diff --git a/internal/server/web/proxy/x_custom.go b/internal/server/web/proxy/x_custom.go index c0132f6..a61df56 100644 --- a/internal/server/web/proxy/x_custom.go +++ b/internal/server/web/proxy/x_custom.go @@ -1,6 +1,7 @@ package proxy import ( + "bytes" "context" "errors" "fmt" @@ -9,12 +10,65 @@ import ( "github.com/bricks-cloud/bricksllm/internal/telemetry" "github.com/bricks-cloud/bricksllm/internal/util" "github.com/gin-gonic/gin" + "io" + "mime" "net/http" "net/http/httputil" "net/url" "strings" ) +// xCustomMaxCapturedResponseBytes caps how much of a response body gets +// buffered for later processing (c.Set). Responses larger than this are +// still proxied in full, they're just not captured. +const xCustomMaxCapturedResponseBytes = 5 * 1024 * 1024 // 5MB + +// xCustomCapturingBody wraps a response body so its bytes keep flowing to +// the client exactly as they arrive (no buffering delay, streaming stays +// real-time), while also being copied into an in-memory buffer for later +// use. onClose runs once the upstream body has been fully read/closed. +type xCustomCapturingBody struct { + io.ReadCloser + buf bytes.Buffer + exceeded bool + // failed is set when the upstream body read ends in an error other than + // io.EOF (aborted/truncated response), so we don't capture partial data + // as if it were the complete response. + failed bool + onClose func(data []byte, exceeded bool) +} + +func (b *xCustomCapturingBody) Read(p []byte) (int, error) { + n, err := b.ReadCloser.Read(p) + if n > 0 && !b.exceeded { + if b.buf.Len()+n > xCustomMaxCapturedResponseBytes { + b.exceeded = true + // drop the reference so the already-buffered bytes can be + // garbage collected instead of being held for the rest of + // a possibly long-lived stream. + b.buf = bytes.Buffer{} + } else { + b.buf.Write(p[:n]) + } + } + + if err != nil && err != io.EOF { + b.failed = true + } + + return n, err +} + +func (b *xCustomCapturingBody) Close() error { + err := b.ReadCloser.Close() + + if b.onClose != nil && !b.failed { + b.onClose(b.buf.Bytes(), b.exceeded) + } + + return err +} + func getXCustomHandler(prod bool) gin.HandlerFunc { return func(c *gin.Context) { log := util.GetLogFromCtx(c) @@ -67,6 +121,43 @@ func getXCustomHandler(prod bool) gin.HandlerFunc { r.SetURL(target) r.Out.URL.Path, r.Out.URL.RawPath = target.Path, target.RawPath r.Out.WithContext(ctx) + + // Let the transport negotiate and transparently decompress + // the upstream response itself; otherwise a forwarded + // client Accept-Encoding disables that and ModifyResponse + // would capture raw compressed bytes instead of text. + r.Out.Header.Del("Accept-Encoding") + }, + ModifyResponse: func(res *http.Response) error { + if res.Body == nil || res.StatusCode == http.StatusSwitchingProtocols { + // A 101 response's Body is an io.ReadWriteCloser used + // for bidirectional upgrade proxying (e.g. WebSocket); + // wrapping it would strip that and break the upgrade. + return nil + } + + mediaType, _, _ := mime.ParseMediaType(res.Header.Get("Content-Type")) + isStreaming := mediaType == "text/event-stream" + + res.Body = &xCustomCapturingBody{ + ReadCloser: res.Body, + onClose: func(data []byte, exceeded bool) { + if exceeded || len(data) == 0 { + return + } + + if isStreaming { + c.Set("content", string(data)) + c.Set("streaming_response", data) + c.Set("stream", true) + return + } + + c.Set("response", data) + }, + } + + return nil }, } proxy.ServeHTTP(c.Writer, c.Request)