fix(proxy): stream Server-Sent Events instead of buffering them for capture

CHANGELOG.md

@@ -4,6 +4,12 @@ Gatekeeper is a standalone credential-injecting TLS-intercepting proxy. It trans
 
 Gatekeeper is pre-1.0. The configuration schema and credential source interface may change between minor versions.
 
+## v0.14.1 — 2026-06-23
+
+### Fixed
+
+- **Streamed responses are no longer buffered during capture** — the response-body log sampler did a blocking read of up to `MaxBodySize` (8 KB) before forwarding, so any incrementally-produced text response (Server-Sent Events, `application/x-ndjson`, chunked JSON, …) had its status line and every chunk — including the keepalive pings an upstream sends during a long time-to-first-token — withheld from the client until 8 KB accumulated, or until the stream ended for responses under 8 KB. The client received nothing, tripped its first-byte timeout, and retried in a loop; most visibly on large or cache-cold LLM streaming requests (e.g. `/v1/messages`), which worked when sent directly (un-proxied). The sampler now captures text responses lazily via a non-blocking tee instead of a blocking read-ahead, so a slow or streamed response (SSE, ndjson, chunked JSON) is forwarded immediately rather than withheld. The canonical log line for a text response is now written when the body completes rather than when its headers arrive; non-text and non-streaming responses are unchanged. This does **not** apply to hosts with an `llm-gateway` Keep engine: response-policy evaluation reads the full body before forwarding, so streams on those hosts are still buffered end-to-end ([#37](https://github.com/majorcontext/gatekeeper/pull/37))
+
 ## v0.14.0 — 2026-06-23
 
 ### Added

proxy/intercept_test.go

@@ -15,10 +15,31 @@ import (
 	"sync"
 	"sync/atomic"
 	"testing"
+	"time"
 
 	keeplib "github.com/majorcontext/keep"
 )
 
+// captureLog installs a logger that records logged requests and returns a
+// function that waits for the next one. The canonical log line for a text
+// response is written when its body is closed (capture streams lazily), so
+// tests must wait for it rather than read a shared variable synchronously.
+func captureLog(t *testing.T, p *Proxy) func() RequestLogData {
+	t.Helper()
+	ch := make(chan RequestLogData, 8)
+	p.SetLogger(func(d RequestLogData) { ch <- d })
+	return func() RequestLogData {
+		t.Helper()
+		select {
+		case d := <-ch:
+			return d
+		case <-time.After(5 * time.Second):
+			t.Fatal("timed out waiting for canonical log line")
+			return RequestLogData{}
+		}
+	}
+}
+
 // interceptTestSetup creates a proxy with TLS interception enabled and an HTTPS
 // backend server. The proxy is configured to trust the backend's TLS cert and
 // the returned client trusts the proxy's interception CA.
@@ -110,17 +131,16 @@ func TestIntercept_CredentialInjectionCanonicalLog(t *testing.T) {
 
 	setup.Proxy.SetCredentialWithGrant(setup.BackendHost, "Authorization", "Bearer granted-token", "my-grant")
 
-	var logged RequestLogData
-	setup.Proxy.SetLogger(func(data RequestLogData) {
-		logged = data
-	})
+	waitLog := captureLog(t, setup.Proxy)
 
 	resp, err := setup.Client.Get(setup.Backend.URL + "/resource")
 	if err != nil {
 		t.Fatalf("request: %v", err)
 	}
+	io.ReadAll(resp.Body)
 	resp.Body.Close()
 
+	logged := waitLog()
 	if !logged.AuthInjected {
 		t.Error("expected AuthInjected=true")
 	}
@@ -249,17 +269,16 @@ func TestIntercept_CanonicalLogFields(t *testing.T) {
 
 	setup.Proxy.SetCredentialWithGrant(setup.BackendHost, "Authorization", "Bearer tok", "test-grant")
 
-	var logged RequestLogData
-	setup.Proxy.SetLogger(func(data RequestLogData) {
-		logged = data
-	})
+	waitLog := captureLog(t, setup.Proxy)
 
 	resp, err := setup.Client.Get(setup.Backend.URL + "/some/path")
 	if err != nil {
 		t.Fatalf("request: %v", err)
 	}
+	io.ReadAll(resp.Body)
 	resp.Body.Close()
 
+	logged := waitLog()
 	if logged.Method != "GET" {
 		t.Errorf("Method = %q, want GET", logged.Method)
 	}
@@ -625,10 +644,7 @@ func TestIntercept_CaptureHeaders_StrippedBeforeForwarding(t *testing.T) {
 
 	setup.Proxy.SetCaptureHeaders([]string{"X-Workspace-Slug", "X-Request-Source"})
 
-	var logged RequestLogData
-	setup.Proxy.SetLogger(func(data RequestLogData) {
-		logged = data
-	})
+	waitLog := captureLog(t, setup.Proxy)
 
 	req, _ := http.NewRequest("GET", setup.Backend.URL+"/test", nil)
 	req.Header.Set("X-Workspace-Slug", "sneaky-plum")
@@ -641,6 +657,7 @@ func TestIntercept_CaptureHeaders_StrippedBeforeForwarding(t *testing.T) {
 	}
 	defer resp.Body.Close()
 	io.ReadAll(resp.Body)
+	logged := waitLog()
 
 	// Verify capture headers were stripped before forwarding.
 	if receivedHeaders.Get("X-Workspace-Slug") != "" {
@@ -769,10 +786,7 @@ func TestIntercept_UserID_ContextResolver(t *testing.T) {
 		return nil, false
 	})
 
-	var logged RequestLogData
-	setup.Proxy.SetLogger(func(data RequestLogData) {
-		logged = data
-	})
+	waitLog := captureLog(t, setup.Proxy)
 
 	// Rebuild the client with proxy auth credentials (user:token).
 	proxyURL := mustParseURL(setup.ProxyServer.URL)
@@ -793,6 +807,7 @@ func TestIntercept_UserID_ContextResolver(t *testing.T) {
 	}
 	defer resp.Body.Close()
 	io.ReadAll(resp.Body)
+	logged := waitLog()
 
 	if logged.UserID != "alice" {
 		t.Errorf("UserID = %q, want %q (CONNECT path)", logged.UserID, "alice")
@@ -951,8 +966,7 @@ func TestIntercept_LLMPolicy_DeniedLogged(t *testing.T) {
 		}, true
 	})
 
-	var logged RequestLogData
-	p.SetLogger(func(data RequestLogData) { logged = data })
+	waitLog := captureLog(t, p)
 
 	proxyServer := httptest.NewServer(p)
 	t.Cleanup(proxyServer.Close)
@@ -979,6 +993,7 @@ func TestIntercept_LLMPolicy_DeniedLogged(t *testing.T) {
 	}
 	io.ReadAll(resp.Body)
 	resp.Body.Close()
+	logged := waitLog()
 
 	if resp.StatusCode != http.StatusBadRequest {
 		t.Errorf("status = %d, want 400", resp.StatusCode)

proxy/proxy.go

@@ -372,6 +372,12 @@ type readCloserWrapper struct {
 // the captured data and a new ReadCloser that streams the full content.
 // For small bodies (<=MaxBodySize), the body is fully buffered.
 // For large bodies, only MaxBodySize is buffered; the rest streams through.
+//
+// captureBody does a blocking read-ahead, so it must only be used on bodies that
+// are already fully available — i.e. request bodies. For response bodies, which
+// may be produced incrementally (Server-Sent Events, ndjson, chunked JSON), use
+// capturingBody: a read-ahead there would withhold the response from the client
+// until MaxBodySize accumulated, starving streamed responses.
 func captureBody(body io.ReadCloser, contentType string) ([]byte, io.ReadCloser) {
 	if body == nil {
 		return nil, nil
@@ -407,6 +413,54 @@ func captureBody(body io.ReadCloser, contentType string) ([]byte, io.ReadCloser)
 	return captured, &readCloserWrapper{Reader: fullBody, Closer: body}
 }
 
+// capturingBody wraps a response body, copying up to limit bytes into an
+// in-memory buffer as the body is read, then invoking onClose (if non-nil)
+// exactly once when the body is closed. Capture happens lazily as the downstream
+// consumer reads, so — unlike captureBody's read-ahead — it never blocks the
+// forwarding of a streamed response. Used for response-body logging on the
+// streaming proxy paths.
+//
+// Not safe for concurrent Read/Close; the proxy paths read the body to
+// completion and then close it, which is sequential.
+type capturingBody struct {
+	rc      io.ReadCloser
+	buf     bytes.Buffer
+	limit   int
+	onClose func(captured []byte)
+	closed  bool
+}
+
+func newCapturingBody(rc io.ReadCloser, limit int, onClose func(captured []byte)) *capturingBody {
+	return &capturingBody{rc: rc, limit: limit, onClose: onClose}
+}
+
+func (c *capturingBody) Read(p []byte) (int, error) {
+	n, err := c.rc.Read(p)
+	if n > 0 {
+		if room := c.limit - c.buf.Len(); room > 0 {
+			if room > n {
+				room = n
+			}
+			c.buf.Write(p[:room])
+		}
+	}
+	return n, err
+}
+
+// Captured returns the bytes buffered so far (up to limit).
+func (c *capturingBody) Captured() []byte { return c.buf.Bytes() }
+
+func (c *capturingBody) Close() error {
+	err := c.rc.Close()
+	if !c.closed {
+		c.closed = true
+		if c.onClose != nil {
+			c.onClose(c.buf.Bytes())
+		}
+	}
+	return err
+}
+
 // FilterHeaders creates a copy of headers with sensitive values filtered.
 // injectedHeaders is a set of lower-cased header names whose values should be
 // redacted (credential headers the proxy injected).
@@ -1832,52 +1886,56 @@ func (p *Proxy) handleHTTP(w http.ResponseWriter, r *http.Request) {
 
 	// Forward request
 	resp, err := httpTransport.RoundTrip(outReq)
-	duration := time.Since(start)
-
-	// Capture response body and headers
-	var respBody []byte
-	var respHeaders http.Header
-	statusCode := http.StatusBadGateway
-	var responseSize int64 = -1
-	if resp != nil {
-		respHeaders = resp.Header.Clone()
-		respBody, resp.Body = captureBody(resp.Body, resp.Header.Get("Content-Type"))
-		statusCode = resp.StatusCode
-		responseSize = resp.ContentLength
-	}
 
-	p.logRequest(r, RequestLogData{
+	// Fields shared by the error and success log lines.
+	logData := RequestLogData{
 		Method:          r.Method,
 		URL:             r.URL.String(),
 		Host:            host,
 		Path:            r.URL.Path,
 		RequestType:     "http",
-		StatusCode:      statusCode,
-		Duration:        duration,
-		Err:             err,
 		RequestHeaders:  originalReqHeaders,
-		ResponseHeaders: respHeaders,
 		RequestBody:     reqBody,
-		ResponseBody:    respBody,
 		RequestSize:     r.ContentLength,
-		ResponseSize:    responseSize,
+		ResponseSize:    -1,
 		InjectedHeaders: credResult.InjectedHeaders,
 		Grants:          credResult.Grants,
-	})
+	}
 
 	if err != nil {
+		logData.StatusCode = http.StatusBadGateway
+		logData.Duration = time.Since(start)
+		logData.Err = err
+		p.logRequest(r, logData)
 		http.Error(w, "moat proxy: upstream request failed", http.StatusBadGateway)
 		return
 	}
 	defer resp.Body.Close()
 
+	logData.StatusCode = resp.StatusCode
+	logData.ResponseHeaders = resp.Header.Clone()
+	logData.ResponseSize = resp.ContentLength
+
 	for key, values := range resp.Header {
 		for _, value := range values {
 			w.Header().Add(key, value)
 		}
 	}
 	w.WriteHeader(resp.StatusCode)
-	_, _ = io.Copy(w, resp.Body)
+
+	// Stream to the client, capturing a bounded sample for logging without a
+	// blocking read-ahead (which would starve a streamed response). The log line
+	// is written after the body completes. Non-text bodies carry no useful
+	// sample, so they stream straight through.
+	if isTextContentType(resp.Header.Get("Content-Type")) {
+		cb := newCapturingBody(resp.Body, MaxBodySize, nil)
+		_, _ = io.Copy(w, cb)
+		logData.ResponseBody = cb.Captured()
+	} else {
+		_, _ = io.Copy(w, resp.Body)
+	}
+	logData.Duration = time.Since(start)
+	p.logRequest(r, logData)
 }
 
 func (p *Proxy) handleConnect(w http.ResponseWriter, r *http.Request) {
@@ -2334,37 +2392,56 @@ func (p *Proxy) handleConnectWithInterception(w http.ResponseWriter, r *http.Req
 			if logURL == "" {
 				logURL = req.URL.String()
 			}
-			var respBody []byte
-			respBody, resp.Body = captureBody(resp.Body, resp.Header.Get("Content-Type"))
-
 			// Use pre-injection headers so credential values don't appear in logs.
 			preHeaders, _ := req.Context().Value(interceptPreInjHeadersKey{}).(http.Header)
 			if preHeaders == nil {
 				preHeaders = req.Header.Clone()
 			}
 			reqBody, _ := req.Context().Value(interceptReqBodyKey{}).([]byte)
 
-			p.logRequest(r, RequestLogData{
+			logData := RequestLogData{
 				RequestID:       req.Header.Get("X-Request-Id"),
 				Method:          req.Method,
 				URL:             logURL,
 				Host:            host,
 				Path:            req.URL.Path,
 				RequestType:     "connect",
 				StatusCode:      resp.StatusCode,
-				Duration:        time.Since(reqStartFromContext(req.Context())),
 				RequestHeaders:  preHeaders,
 				ResponseHeaders: resp.Header.Clone(),
 				RequestBody:     reqBody,
-				ResponseBody:    respBody,
 				RequestSize:     req.ContentLength,
 				ResponseSize:    resp.ContentLength,
 				AuthInjected:    len(credResult.InjectedHeaders) > 0,
 				InjectedHeaders: credResult.InjectedHeaders,
 				Grants:          credResult.Grants,
 				Denied:          llmDenied,
 				DenyReason:      llmDenyReason,
-			})
+			}
+			reqStart := reqStartFromContext(req.Context())
+
+			// For text responses (SSE, ndjson, chunked JSON, …) defer the canonical
+			// log line until the body is read and closed, capturing a bounded sample
+			// as it streams. A blocking read-ahead here would withhold the status
+			// line and every chunk from the client until the sample filled, starving
+			// streamed responses and tripping the client's first-byte timeout.
+			//
+			// The deferred line is emitted exactly once: httputil.ReverseProxy
+			// closes the response body unconditionally (success and copy-error
+			// paths), so onClose always runs. Responses with no useful streamed
+			// sample are logged synchronously instead: non-text bodies, and protocol
+			// upgrades (101), whose body is the hijacked connection rather than a
+			// readable stream.
+			if isTextContentType(resp.Header.Get("Content-Type")) && resp.StatusCode != http.StatusSwitchingProtocols {
+				resp.Body = newCapturingBody(resp.Body, MaxBodySize, func(captured []byte) {
+					logData.ResponseBody = captured
+					logData.Duration = time.Since(reqStart)
+					p.logRequest(r, logData)
+				})
+			} else {
+				logData.Duration = time.Since(reqStart)
+				p.logRequest(r, logData)
+			}
 
 			return nil
 		},