mirror of
https://github.com/mudler/LocalAI.git
synced 2026-05-30 11:36:31 -04:00
fix(traces): cap captured body size to keep admin Traces UI responsive (#9946)
The trace middleware buffered the full request and response bodies for every JSON exchange. With a chatty agent-pool RAG workload, /embeddings responses (large vector arrays) accumulated to tens of MB in the in-memory buffer; the admin Traces page would then download and parse 40+ MB on every load and on every 5s auto-refresh, locking the UI in a loading state. Add LOCALAI_TRACING_MAX_BODY_BYTES (default 64 KiB) that caps each captured body. The full payload still flows through to the real client; only the trace copy is bounded. Exchanges record body_truncated and original body_bytes so the dashboard can show that truncation happened. The cap is configurable via env, CLI, and runtime_settings.json. Also unblock recovery: the Traces page now keeps the Clear button enabled while loading, since "buffer too large to render" is exactly when the user needs to clear it. Assisted-by: Claude:claude-opus-4-7 Signed-off-by: Ettore Di Giacinto <mudler@localai.io> Co-authored-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
@@ -17,16 +17,20 @@ import (
|
||||
)
|
||||
|
||||
type APIExchangeRequest struct {
|
||||
Method string `json:"method"`
|
||||
Path string `json:"path"`
|
||||
Headers *http.Header `json:"headers"`
|
||||
Body *[]byte `json:"body"`
|
||||
Method string `json:"method"`
|
||||
Path string `json:"path"`
|
||||
Headers *http.Header `json:"headers"`
|
||||
Body *[]byte `json:"body"`
|
||||
BodyTruncated bool `json:"body_truncated,omitempty"`
|
||||
BodyBytes int `json:"body_bytes,omitempty"` // original size before truncation
|
||||
}
|
||||
|
||||
type APIExchangeResponse struct {
|
||||
Status int `json:"status"`
|
||||
Headers *http.Header `json:"headers"`
|
||||
Body *[]byte `json:"body"`
|
||||
Status int `json:"status"`
|
||||
Headers *http.Header `json:"headers"`
|
||||
Body *[]byte `json:"body"`
|
||||
BodyTruncated bool `json:"body_truncated,omitempty"`
|
||||
BodyBytes int `json:"body_bytes,omitempty"` // original size before truncation
|
||||
}
|
||||
|
||||
type APIExchange struct {
|
||||
@@ -66,11 +70,29 @@ var doInitializeTracing = sync.OnceFunc(func() {
|
||||
|
||||
type bodyWriter struct {
|
||||
http.ResponseWriter
|
||||
body *bytes.Buffer
|
||||
body *bytes.Buffer
|
||||
maxBytes int // 0 = unlimited capture
|
||||
truncated bool
|
||||
totalBytes int // bytes the upstream handler wrote, even past the cap
|
||||
}
|
||||
|
||||
func (w *bodyWriter) Write(b []byte) (int, error) {
|
||||
w.body.Write(b)
|
||||
// Capture into the trace buffer up to maxBytes, then drop the overflow
|
||||
// so a chatty endpoint can't grow the buffer without bound. The full
|
||||
// payload still flows through to the real client below.
|
||||
w.totalBytes += len(b)
|
||||
if w.maxBytes <= 0 {
|
||||
w.body.Write(b)
|
||||
} else if remain := w.maxBytes - w.body.Len(); remain > 0 {
|
||||
if remain >= len(b) {
|
||||
w.body.Write(b)
|
||||
} else {
|
||||
w.body.Write(b[:remain])
|
||||
w.truncated = true
|
||||
}
|
||||
} else {
|
||||
w.truncated = true
|
||||
}
|
||||
return w.ResponseWriter.Write(b)
|
||||
}
|
||||
|
||||
@@ -80,6 +102,20 @@ func (w *bodyWriter) Flush() {
|
||||
}
|
||||
}
|
||||
|
||||
// truncateForTrace returns a defensive copy of body capped at maxBytes,
|
||||
// and a flag indicating whether the cap forced truncation. maxBytes <= 0
|
||||
// disables the cap.
|
||||
func truncateForTrace(body []byte, maxBytes int) ([]byte, bool) {
|
||||
if maxBytes <= 0 || len(body) <= maxBytes {
|
||||
out := make([]byte, len(body))
|
||||
copy(out, body)
|
||||
return out, false
|
||||
}
|
||||
out := make([]byte, maxBytes)
|
||||
copy(out, body[:maxBytes])
|
||||
return out, true
|
||||
}
|
||||
|
||||
func initializeTracing(maxItems int) {
|
||||
tracingMaxItems = maxItems
|
||||
doInitializeTracing()
|
||||
@@ -134,11 +170,18 @@ func TraceMiddleware(app *application.Application) echo.MiddlewareFunc {
|
||||
|
||||
startTime := time.Now()
|
||||
|
||||
// Cap captured payload size. Without this, /embeddings and
|
||||
// streaming /chat/completions blow the in-memory buffer into the
|
||||
// tens of MB, which then locks the admin Traces UI fetching the
|
||||
// JSON dump faster than the 5s auto-refresh.
|
||||
maxBodyBytes := app.ApplicationConfig().TracingMaxBodyBytes
|
||||
|
||||
// Wrap response writer to capture body
|
||||
resBody := new(bytes.Buffer)
|
||||
mw := &bodyWriter{
|
||||
ResponseWriter: c.Response().Writer,
|
||||
body: resBody,
|
||||
maxBytes: maxBodyBytes,
|
||||
}
|
||||
c.Response().Writer = mw
|
||||
|
||||
@@ -159,8 +202,7 @@ func TraceMiddleware(app *application.Application) echo.MiddlewareFunc {
|
||||
// via any heap-dump-style introspection, and tokens shouldn't
|
||||
// outlive the request that carried them.
|
||||
requestHeaders := redactSensitiveHeaders(c.Request().Header)
|
||||
requestBody := make([]byte, len(body))
|
||||
copy(requestBody, body)
|
||||
requestBody, requestTruncated := truncateForTrace(body, maxBodyBytes)
|
||||
responseHeaders := redactSensitiveHeaders(c.Response().Header())
|
||||
responseBody := make([]byte, resBody.Len())
|
||||
copy(responseBody, resBody.Bytes())
|
||||
@@ -168,15 +210,19 @@ func TraceMiddleware(app *application.Application) echo.MiddlewareFunc {
|
||||
Timestamp: startTime,
|
||||
Duration: time.Since(startTime),
|
||||
Request: APIExchangeRequest{
|
||||
Method: c.Request().Method,
|
||||
Path: c.Path(),
|
||||
Headers: &requestHeaders,
|
||||
Body: &requestBody,
|
||||
Method: c.Request().Method,
|
||||
Path: c.Path(),
|
||||
Headers: &requestHeaders,
|
||||
Body: &requestBody,
|
||||
BodyTruncated: requestTruncated,
|
||||
BodyBytes: len(body),
|
||||
},
|
||||
Response: APIExchangeResponse{
|
||||
Status: status,
|
||||
Headers: &responseHeaders,
|
||||
Body: &responseBody,
|
||||
Status: status,
|
||||
Headers: &responseHeaders,
|
||||
Body: &responseBody,
|
||||
BodyTruncated: mw.truncated,
|
||||
BodyBytes: mw.totalBytes,
|
||||
},
|
||||
}
|
||||
if handlerErr != nil {
|
||||
|
||||
116
core/http/middleware/trace_body_cap_test.go
Normal file
116
core/http/middleware/trace_body_cap_test.go
Normal file
@@ -0,0 +1,116 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// The trace middleware copies request and response bodies into an in-memory
|
||||
// buffer that backs the admin /api/traces endpoint. With no upper bound a
|
||||
// chatty workload (embeddings, large completions) trivially produces a
|
||||
// multi-MB response that locks the Traces UI in a loading state — fetching
|
||||
// and parsing the payload outruns the 5-second auto-refresh. These specs
|
||||
// pin the capping contract so future refactors keep both the cap and the
|
||||
// passthrough to the real client intact.
|
||||
|
||||
var _ = Describe("bodyWriter capping", func() {
|
||||
It("captures the full body when maxBytes is 0 (unlimited)", func() {
|
||||
downstream := httptest.NewRecorder()
|
||||
buf := &bytes.Buffer{}
|
||||
bw := &bodyWriter{ResponseWriter: downstream, body: buf, maxBytes: 0}
|
||||
|
||||
payload := []byte(strings.Repeat("x", 4096))
|
||||
n, err := bw.Write(payload)
|
||||
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(n).To(Equal(len(payload)))
|
||||
Expect(buf.Len()).To(Equal(len(payload)))
|
||||
Expect(downstream.Body.Len()).To(Equal(len(payload)))
|
||||
Expect(bw.truncated).To(BeFalse())
|
||||
})
|
||||
|
||||
It("stops appending to the trace buffer once maxBytes is reached but still forwards to the client", func() {
|
||||
downstream := httptest.NewRecorder()
|
||||
buf := &bytes.Buffer{}
|
||||
bw := &bodyWriter{ResponseWriter: downstream, body: buf, maxBytes: 100}
|
||||
|
||||
payload := []byte(strings.Repeat("a", 250))
|
||||
n, err := bw.Write(payload)
|
||||
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(n).To(Equal(len(payload)), "Write must return the full byte count so callers see no short write")
|
||||
Expect(buf.Len()).To(Equal(100), "trace buffer should hold exactly maxBytes")
|
||||
Expect(downstream.Body.Len()).To(Equal(len(payload)), "client must still receive every byte")
|
||||
Expect(bw.truncated).To(BeTrue())
|
||||
})
|
||||
|
||||
It("handles a write that straddles the cap by keeping only the leading slice", func() {
|
||||
downstream := httptest.NewRecorder()
|
||||
buf := &bytes.Buffer{}
|
||||
bw := &bodyWriter{ResponseWriter: downstream, body: buf, maxBytes: 10}
|
||||
|
||||
_, err := bw.Write([]byte("12345"))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(bw.truncated).To(BeFalse())
|
||||
|
||||
_, err = bw.Write([]byte("67890ABCDE"))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
Expect(buf.String()).To(Equal("1234567890"))
|
||||
Expect(downstream.Body.String()).To(Equal("1234567890ABCDE"))
|
||||
Expect(bw.truncated).To(BeTrue())
|
||||
})
|
||||
|
||||
It("ignores further writes after the cap was already hit", func() {
|
||||
downstream := httptest.NewRecorder()
|
||||
buf := &bytes.Buffer{}
|
||||
bw := &bodyWriter{ResponseWriter: downstream, body: buf, maxBytes: 4}
|
||||
|
||||
_, _ = bw.Write([]byte("AAAA"))
|
||||
_, _ = bw.Write([]byte("BBBB"))
|
||||
_, _ = bw.Write([]byte("CCCC"))
|
||||
|
||||
Expect(buf.String()).To(Equal("AAAA"))
|
||||
Expect(downstream.Body.String()).To(Equal("AAAABBBBCCCC"))
|
||||
Expect(bw.truncated).To(BeTrue())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("truncateForTrace", func() {
|
||||
It("returns the input unchanged when below the cap", func() {
|
||||
in := []byte("hello")
|
||||
out, truncated := truncateForTrace(in, 1024)
|
||||
Expect(truncated).To(BeFalse())
|
||||
Expect(out).To(Equal(in))
|
||||
})
|
||||
|
||||
It("truncates when the input exceeds the cap and signals truncation", func() {
|
||||
in := []byte(strings.Repeat("z", 200))
|
||||
out, truncated := truncateForTrace(in, 64)
|
||||
Expect(truncated).To(BeTrue())
|
||||
Expect(out).To(HaveLen(64))
|
||||
Expect(string(out)).To(Equal(strings.Repeat("z", 64)))
|
||||
})
|
||||
|
||||
It("treats maxBytes <= 0 as unlimited (back-compat with current default)", func() {
|
||||
in := []byte(strings.Repeat("q", 10_000))
|
||||
out, truncated := truncateForTrace(in, 0)
|
||||
Expect(truncated).To(BeFalse())
|
||||
Expect(out).To(HaveLen(len(in)))
|
||||
})
|
||||
|
||||
It("does not retain the caller's backing array (defensive copy)", func() {
|
||||
in := []byte("abcdefghij")
|
||||
out, truncated := truncateForTrace(in, 4)
|
||||
Expect(truncated).To(BeTrue())
|
||||
Expect(string(out)).To(Equal("abcd"))
|
||||
|
||||
// Mutating the source must not corrupt the trace copy.
|
||||
in[0] = 'Z'
|
||||
Expect(string(out)).To(Equal("abcd"))
|
||||
})
|
||||
})
|
||||
@@ -406,7 +406,15 @@ export default function Traces() {
|
||||
<button className="btn btn-secondary btn-sm" onClick={fetchTraces}><i className="fas fa-rotate" /> Refresh</button>
|
||||
<button className="btn btn-secondary btn-sm" onClick={handleExport} disabled={traces.length === 0}><i className="fas fa-download" /> Export</button>
|
||||
<div style={{ flex: 1 }} />
|
||||
<button className="btn btn-danger btn-sm" onClick={handleClear} disabled={traces.length === 0}><i className="fas fa-trash" /> Clear</button>
|
||||
<button
|
||||
className="btn btn-danger btn-sm"
|
||||
onClick={handleClear}
|
||||
/* Stay enabled while loading: a massive in-memory trace buffer is
|
||||
precisely the case where the user can't see the table yet and
|
||||
needs Clear to recover. Clearing an already-empty server-side
|
||||
buffer is a harmless no-op. */
|
||||
disabled={!loading && traces.length === 0}
|
||||
><i className="fas fa-trash" /> Clear</button>
|
||||
</div>
|
||||
|
||||
{settings && (() => {
|
||||
|
||||
Reference in New Issue
Block a user