@@ -2,43 +2,40 @@ package middleware
22
33import (
44 "bytes"
5+ "github.com/emirpasic/gods/v2/queues/circularbuffer"
56 "io"
67 "net/http"
8+ "sort"
79 "sync"
10+ "time"
811
912 "github.com/labstack/echo/v4"
1013 "github.com/mudler/LocalAI/core/application"
1114 "github.com/rs/zerolog/log"
1215)
1316
17+ type APIExchangeRequest struct {
18+ Method string `json:"method"`
19+ Path string `json:"path"`
20+ Headers * http.Header `json:"headers"`
21+ Body * []byte `json:"body"`
22+ }
23+
24+ type APIExchangeResponse struct {
25+ Status int `json:"status"`
26+ Headers * http.Header `json:"headers"`
27+ Body * []byte `json:"body"`
28+ }
29+
1430type APIExchange struct {
15- Request struct {
16- Method string
17- Path string
18- Headers http.Header
19- Body []byte
20- }
21- Response struct {
22- Status int
23- Headers http.Header
24- Body []byte
25- }
31+ Timestamp time.Time `json:"timestamp"`
32+ Request APIExchangeRequest `json:"request"`
33+ Response APIExchangeResponse `json:"response"`
2634}
2735
28- var apiLogs [] APIExchange
36+ var traceBuffer * circularbuffer. Queue [ APIExchange ]
2937var mu sync.Mutex
30- var logChan = make (chan APIExchange , 100 ) // Buffered channel for serialization
31-
32- func init () {
33- go func () {
34- for exchange := range logChan {
35- mu .Lock ()
36- apiLogs = append (apiLogs , exchange )
37- mu .Unlock ()
38- log .Debug ().Msgf ("Logged exchange: %s %s - Status: %d" , exchange .Request .Method , exchange .Request .Path , exchange .Response .Status )
39- }
40- }()
41- }
38+ var logChan = make (chan APIExchange , 100 )
4239
4340type bodyWriter struct {
4441 http.ResponseWriter
@@ -58,13 +55,24 @@ func (w *bodyWriter) Flush() {
5855
5956// TraceMiddleware intercepts and logs JSON API requests and responses
6057func TraceMiddleware (app * application.Application ) echo.MiddlewareFunc {
58+ if app .ApplicationConfig ().EnableTracing && traceBuffer == nil {
59+ traceBuffer = circularbuffer.New [APIExchange ](app .ApplicationConfig ().TracingMaxItems )
60+
61+ go func () {
62+ for exchange := range logChan {
63+ mu .Lock ()
64+ traceBuffer .Enqueue (exchange )
65+ mu .Unlock ()
66+ }
67+ }()
68+ }
69+
6170 return func (next echo.HandlerFunc ) echo.HandlerFunc {
6271 return func (c echo.Context ) error {
6372 if ! app .ApplicationConfig ().EnableTracing {
6473 return next (c )
6574 }
6675
67- // Only log if Content-Type is application/json
6876 if c .Request ().Header .Get ("Content-Type" ) != "application/json" {
6977 return next (c )
7078 }
@@ -78,6 +86,8 @@ func TraceMiddleware(app *application.Application) echo.MiddlewareFunc {
7886 // Restore the body for downstream handlers
7987 c .Request ().Body = io .NopCloser (bytes .NewBuffer (body ))
8088
89+ startTime := time .Now ()
90+
8191 // Wrap response writer to capture body
8292 resBody := new (bytes.Buffer )
8393 mw := & bodyWriter {
@@ -93,51 +103,54 @@ func TraceMiddleware(app *application.Application) echo.MiddlewareFunc {
93103 }
94104
95105 // Create exchange log
106+ requestHeaders := c .Request ().Header .Clone ()
107+ requestBody := make ([]byte , len (body ))
108+ copy (requestBody , body )
109+ responseHeaders := c .Response ().Header ().Clone ()
110+ responseBody := make ([]byte , resBody .Len ())
111+ copy (responseBody , resBody .Bytes ())
96112 exchange := APIExchange {
97- Request : struct {
98- Method string
99- Path string
100- Headers http.Header
101- Body []byte
102- }{
113+ Timestamp : startTime ,
114+ Request : APIExchangeRequest {
103115 Method : c .Request ().Method ,
104116 Path : c .Path (),
105- Headers : c . Request (). Header . Clone () ,
106- Body : body ,
117+ Headers : & requestHeaders ,
118+ Body : & requestBody ,
107119 },
108- Response : struct {
109- Status int
110- Headers http.Header
111- Body []byte
112- }{
120+ Response : APIExchangeResponse {
113121 Status : c .Response ().Status ,
114- Headers : c . Response (). Header (). Clone () ,
115- Body : resBody . Bytes () ,
122+ Headers : & responseHeaders ,
123+ Body : & responseBody ,
116124 },
117125 }
118126
119- // Send to channel (non-blocking)
120127 select {
121128 case logChan <- exchange :
122129 default :
123- log .Warn ().Msg ("API log channel full, dropping log " )
130+ log .Warn ().Msg ("Trace channel full, dropping trace " )
124131 }
125132
126133 return nil
127134 }
128135 }
129136}
130137
131- // GetAPILogs returns a copy of the logged API exchanges for display
132- func GetAPILogs () []APIExchange {
138+ // GetTraces returns a copy of the logged API exchanges for display
139+ func GetTraces () []APIExchange {
133140 mu .Lock ()
134- defer mu .Unlock ()
135- return append ([]APIExchange {}, apiLogs ... )
141+ traces := traceBuffer .Values ()
142+ mu .Unlock ()
143+
144+ sort .Slice (traces , func (i , j int ) bool {
145+ return traces [i ].Timestamp .Before (traces [j ].Timestamp )
146+ })
147+
148+ return traces
136149}
137150
138- // ClearAPILogs clears the in-memory logs
139- func ClearAPILogs () {
151+ // ClearTraces clears the in-memory logs
152+ func ClearTraces () {
140153 mu .Lock ()
141- apiLogs = nil
154+ traceBuffer . Clear ()
142155 mu .Unlock ()
143156}
0 commit comments