Implement stateful per-chunk streaming compression

gateway/gateway-runtime/policy-engine/internal/kernel/decompression.go

@@ -168,8 +168,94 @@ func (sd *streamDecompressor) Close() {
 	}
 }
 
+// streamWriteFlushCloser is the common interface implemented by *gzip.Writer and
+// *brotli.Writer: incremental writes, a mid-stream flush that keeps the stream open,
+// and a final close that emits the trailer.
+type streamWriteFlushCloser interface {
+	io.WriteCloser
+	Flush() error
+}
+
+// streamCompressor provides stateful, per-chunk streaming compression that produces
+// a SINGLE continuous compressed stream across all chunks.
+//
+// This is the compression counterpart of streamDecompressor and must be used instead
+// of recompressBody for streaming bodies. recompressBody creates a fresh writer and
+// Close()es it on every call, so each chunk becomes an independent, self-contained
+// gzip/brotli member. A concatenation of independent members is not what the
+// Content-Encoding header promises (one stream for the whole body), and downstream
+// HTTP decoders (e.g. the Anthropic/Claude Code client) decode only the first member,
+// see its end-of-stream trailer, and treat the response as finished — dropping every
+// subsequent chunk. streamCompressor instead keeps one writer alive for the life of
+// the stream, flushing (Z_SYNC_FLUSH) after each chunk and closing only at EndOfStream.
+type streamCompressor struct {
+	buf      *bytes.Buffer
+	writer   streamWriteFlushCloser
+	encoding string
+	closed   bool
+}
+
+// newStreamCompressor returns a streamCompressor for the given Content-Encoding.
+// For unknown encodings the writer is nil and FeedChunk passes bytes through unchanged.
+func newStreamCompressor(encoding string) *streamCompressor {
+	buf := &bytes.Buffer{}
+	switch encoding {
+	case "gzip":
+		return &streamCompressor{buf: buf, writer: gzip.NewWriter(buf), encoding: encoding}
+	case "br":
+		return &streamCompressor{buf: buf, writer: brotli.NewWriter(buf), encoding: encoding}
+	default:
+		return &streamCompressor{buf: buf, encoding: encoding}
+	}
+}
+
+// FeedChunk compresses a chunk of the decompressed body and returns the compressed
+// bytes produced so far. On intermediate chunks the writer is flushed (keeping the
+// stream open); on endOfStream the writer is closed, emitting the final block and
+// trailer. The returned bytes belong to the caller (a fresh copy).
+func (sc *streamCompressor) FeedChunk(chunk []byte, endOfStream bool) ([]byte, error) {
+	// Passthrough for unknown encodings.
+	if sc.writer == nil {
+		return chunk, nil
+	}
+	if sc.closed {
+		return nil, fmt.Errorf("stream compressor: write after close")
+	}
+	if len(chunk) > 0 {
+		if _, err := sc.writer.Write(chunk); err != nil {
+			return nil, fmt.Errorf("stream compressor write: %w", err)
+		}
+	}
+	if endOfStream {
+		if err := sc.writer.Close(); err != nil {
+			return nil, fmt.Errorf("stream compressor close: %w", err)
+		}
+		sc.closed = true
+	} else {
+		if err := sc.writer.Flush(); err != nil {
+			return nil, fmt.Errorf("stream compressor flush: %w", err)
+		}
+	}
+	out := make([]byte, sc.buf.Len())
+	copy(out, sc.buf.Bytes())
+	sc.buf.Reset()
+	return out, nil
+}
+
+// Close releases the compressor's writer on error paths where endOfStream will never
+// arrive. Safe to call multiple times.
+func (sc *streamCompressor) Close() {
+	if sc.writer != nil && !sc.closed {
+		_ = sc.writer.Close()
+		sc.closed = true
+	}
+}
+
 // recompressBody re-compresses body bytes using the original Content-Encoding.
 // Used to restore compression after policies have processed the decompressed body.
+// NOTE: This produces a complete standalone stream and is only correct for
+// non-streaming (fully buffered) bodies. For streaming bodies use streamCompressor,
+// which keeps a single stream open across chunks.
 // Supported encodings: "gzip", "br" (Brotli). Unknown encodings are returned as-is.
 func recompressBody(body []byte, encoding string) ([]byte, error) {
 	switch encoding {

gateway/gateway-runtime/policy-engine/internal/kernel/decompression_test.go

@@ -21,6 +21,7 @@ package kernel
 import (
 	"bytes"
 	"compress/gzip"
+	"io"
 	"testing"
 	"time"
 
@@ -280,3 +281,147 @@ func TestStreamDecompressor_RoundTrip(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, original, final)
 }
+
+// =============================================================================
+// streamCompressor Tests
+// =============================================================================
+
+// gzipHeaderCount counts the number of gzip member headers (magic bytes 1f 8b 08)
+// in a byte slice. A correct single continuous stream has exactly one.
+func gzipHeaderCount(data []byte) int {
+	return bytes.Count(data, []byte{0x1f, 0x8b, 0x08})
+}
+
+// singleMemberGunzip decodes only the FIRST gzip member, mimicking downstream HTTP
+// decoders (such as the Claude Code client) that stop at the first member's trailer.
+func singleMemberGunzip(t *testing.T, data []byte) []byte {
+	t.Helper()
+	zr, err := gzip.NewReader(bytes.NewReader(data))
+	require.NoError(t, err)
+	zr.Multistream(false) // do NOT transparently continue into later members
+	out, err := io.ReadAll(zr)
+	require.NoError(t, err)
+	return out
+}
+
+// TestStreamCompressor_Gzip_MultipleChunks is the regression test for the analytics
+// streaming bug: multiple chunks must be re-compressed into ONE continuous gzip stream.
+func TestStreamCompressor_Gzip_MultipleChunks(t *testing.T) {
+	chunks := [][]byte{
+		[]byte("event: message_start\ndata: {\"type\":\"message_start\"}\n\n"),
+		[]byte("event: content_block_delta\ndata: {\"delta\":\"Hi\"}\n\n"),
+		[]byte("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"),
+	}
+	var full []byte
+	for _, c := range chunks {
+		full = append(full, c...)
+	}
+
+	sc := newStreamCompressor("gzip")
+	var compressed []byte
+	for i, c := range chunks {
+		endOfStream := i == len(chunks)-1
+		out, err := sc.FeedChunk(c, endOfStream)
+		require.NoError(t, err)
+		compressed = append(compressed, out...)
+	}
+
+	// The whole body must be a SINGLE gzip member — this is the property the bug violated.
+	assert.Equal(t, 1, gzipHeaderCount(compressed),
+		"streaming compression must emit exactly one gzip member, not one per chunk")
+
+	// A single-member decoder (like the downstream client) must recover ALL chunks,
+	// not just the first. This is what failed for Claude Code before the fix.
+	assert.Equal(t, full, singleMemberGunzip(t, compressed))
+
+	// And a normal (multistream) decode must also yield the full body.
+	final, err := decompressBody(compressed, "gzip")
+	require.NoError(t, err)
+	assert.Equal(t, full, final)
+}
+
+// TestRecompressBody_PerChunk_ProducesMultipleMembers documents the OLD, buggy
+// behaviour that streamCompressor replaces: re-compressing each chunk independently
+// produces N gzip members, and a single-member decoder recovers only the first chunk.
+func TestRecompressBody_PerChunk_ProducesMultipleMembers(t *testing.T) {
+	chunks := [][]byte{
+		[]byte("event: message_start\n\n"),
+		[]byte("event: content_block_delta\n\n"),
+		[]byte("event: message_stop\n\n"),
+	}
+
+	var buggy []byte
+	for _, c := range chunks {
+		out, err := recompressBody(c, "gzip")
+		require.NoError(t, err)
+		buggy = append(buggy, out...)
+	}
+
+	// Per-chunk recompress yields one member per chunk...
+	assert.Equal(t, len(chunks), gzipHeaderCount(buggy))
+	// ...and a single-member decoder sees only the first chunk — the dropped-stream bug.
+	assert.Equal(t, chunks[0], singleMemberGunzip(t, buggy))
+}
+
+// TestStreamCompressor_Brotli_MultipleChunks verifies the brotli path round-trips
+// across multiple chunks into a single continuous stream.
+func TestStreamCompressor_Brotli_MultipleChunks(t *testing.T) {
+	chunks := [][]byte{
+		[]byte("first-chunk-payload"),
+		[]byte("second-chunk-payload"),
+		[]byte("third-chunk-payload"),
+	}
+	var full []byte
+	for _, c := range chunks {
+		full = append(full, c...)
+	}
+
+	sc := newStreamCompressor("br")
+	var compressed []byte
+	for i, c := range chunks {
+		out, err := sc.FeedChunk(c, i == len(chunks)-1)
+		require.NoError(t, err)
+		compressed = append(compressed, out...)
+	}
+
+	final, err := decompressBody(compressed, "br")
+	require.NoError(t, err)
+	assert.Equal(t, full, final)
+}
+
+// TestStreamCompressor_Gzip_EmptyFinalChunk mirrors the real Envoy flow where the
+// EndOfStream chunk carries zero bytes: the trailer must still be emitted so the
+// stream is well-formed.
+func TestStreamCompressor_Gzip_EmptyFinalChunk(t *testing.T) {
+	sc := newStreamCompressor("gzip")
+
+	out1, err := sc.FeedChunk([]byte("payload-one"), false)
+	require.NoError(t, err)
+	out2, err := sc.FeedChunk([]byte("payload-two"), false)
+	require.NoError(t, err)
+	// Final chunk with no data, only EndOfStream — as Envoy delivers it.
+	out3, err := sc.FeedChunk(nil, true)
+	require.NoError(t, err)
+
+	compressed := append(append(append([]byte{}, out1...), out2...), out3...)
+	assert.Equal(t, 1, gzipHeaderCount(compressed))
+	assert.Equal(t, []byte("payload-onepayload-two"), singleMemberGunzip(t, compressed))
+}
+
+// TestStreamCompressor_UnknownEncoding_Passthrough verifies unknown encodings pass
+// bytes through unchanged.
+func TestStreamCompressor_UnknownEncoding_Passthrough(t *testing.T) {
+	sc := newStreamCompressor("identity")
+	out, err := sc.FeedChunk([]byte("raw-bytes"), true)
+	require.NoError(t, err)
+	assert.Equal(t, []byte("raw-bytes"), out)
+}
+
+// TestStreamCompressor_WriteAfterClose returns an error rather than corrupting output.
+func TestStreamCompressor_WriteAfterClose(t *testing.T) {
+	sc := newStreamCompressor("gzip")
+	_, err := sc.FeedChunk([]byte("data"), true)
+	require.NoError(t, err)
+	_, err = sc.FeedChunk([]byte("more"), false)
+	require.Error(t, err)
+}

gateway/gateway-runtime/policy-engine/internal/kernel/execution_context.go

@@ -118,6 +118,10 @@ type PolicyExecutionContext struct {
 	// requestStreamDecomp performs per-chunk decompression for compressed streaming
 	// request bodies. Nil when the request is not Content-Encoded.
 	requestStreamDecomp *streamDecompressor
+	// requestStreamComp performs per-chunk re-compression for compressed streaming
+	// request bodies, keeping a single continuous stream open across chunks. Nil until
+	// the first compressed chunk is forwarded.
+	requestStreamComp *streamCompressor
 
 	// isStreamingResponse is set to true during response headers processing when
 	// streaming indicators are detected AND the policy chain supports streaming.
@@ -127,6 +131,10 @@ type PolicyExecutionContext struct {
 	// responseStreamDecomp performs per-chunk decompression for compressed streaming
 	// response bodies. Nil when the response is not Content-Encoded.
 	responseStreamDecomp *streamDecompressor
+	// responseStreamComp performs per-chunk re-compression for compressed streaming
+	// response bodies, keeping a single continuous stream open across chunks. Nil until
+	// the first compressed chunk is forwarded.
+	responseStreamComp *streamCompressor
 	// streamTerminated is set when a policy returns TerminateStream=true. Any
 	// subsequent upstream chunks that Envoy delivers after we have already sent
 	// EndOfStream downstream are silently suppressed — the downstream connection

gateway/gateway-runtime/policy-engine/internal/kernel/translator.go

@@ -1514,17 +1514,31 @@ func TranslateStreamingRequestChunkAction(result *executor.StreamingRequestExecu
 	// Re-compress the output if the original request was Content-Encoded.
 	// The upstream receives the Content-Encoding header as-is, so the body
 	// bytes must match the encoding the upstream expects.
+	//
+	// Use a stateful streamCompressor (kept for the life of the stream), NOT
+	// recompressBody: per-chunk recompressBody emits an independent closed member per
+	// chunk, so the upstream decoder would stop after the first one. See
+	// TranslateStreamingResponseChunkAction for the full rationale.
 	if execCtx.requestContentEncoding != "" {
-		recompressed, err := recompressBody(outputBody, execCtx.requestContentEncoding)
+		if execCtx.requestStreamComp == nil {
+			execCtx.requestStreamComp = newStreamCompressor(execCtx.requestContentEncoding)
+		}
+		recompressed, err := execCtx.requestStreamComp.FeedChunk(outputBody, originalChunk.EndOfStream)
 		if err != nil {
-			slog.Warn("[streaming] failed to re-compress request body; sending uncompressed — Content-Encoding mismatch",
+			// The Content-Encoding header has already been forwarded to the upstream in
+			// streaming mode and cannot be retracted. Sending the body uncompressed under
+			// a committed Content-Encoding would yield a stream the upstream cannot decode.
+			// Fail the ext_proc stream so Envoy aborts the request rather than delivering
+			// corrupt bytes. The forwarded Content-Encoding state is left intact.
+			slog.Error("[streaming] failed to re-compress request body; terminating stream to avoid Content-Encoding mismatch",
 				"encoding", execCtx.requestContentEncoding,
 				"error", err,
 			)
-			execCtx.requestContentEncoding = ""
-		} else {
-			outputBody = recompressed
+			execCtx.requestStreamComp.Close()
+			return nil, fmt.Errorf("failed to re-compress streaming request body (encoding %q): %w",
+				execCtx.requestContentEncoding, err)
 		}
+		outputBody = recompressed
 	}
 
 	analyticsData := make(map[string]any)
@@ -1585,21 +1599,40 @@ func TranslateStreamingResponseChunkAction(result *executor.StreamingResponseExe
 		outputBody = originalChunk.Chunk
 	}
 
+	// If a policy terminated the stream early (e.g. guardrail intervention), force
+	// EndOfStream so Envoy closes the connection cleanly after delivering the final chunk.
+	endOfStream := originalChunk.EndOfStream || result.StreamTerminated
+
 	// Re-compress the output if the original response was Content-Encoded.
 	// Response headers (including Content-Encoding) are already committed downstream
 	// in streaming mode and cannot be changed — the body must match the encoding
 	// the client expects.
+	//
+	// Use a stateful streamCompressor (created once and kept for the life of the
+	// stream), NOT recompressBody: recompressBody would emit an independent, fully
+	// closed compressed member per chunk, and downstream clients decode only the first
+	// member and treat the stream as finished — dropping every subsequent chunk.
 	if execCtx.responseContentEncoding != "" {
-		recompressed, err := recompressBody(outputBody, execCtx.responseContentEncoding)
+		if execCtx.responseStreamComp == nil {
+			execCtx.responseStreamComp = newStreamCompressor(execCtx.responseContentEncoding)
+		}
+		recompressed, err := execCtx.responseStreamComp.FeedChunk(outputBody, endOfStream)
 		if err != nil {
-			slog.Warn("[streaming] failed to re-compress response body; sending uncompressed — Content-Encoding mismatch",
+			// The Content-Encoding header has already been committed to the downstream
+			// client in streaming mode and cannot be changed. Sending the body
+			// uncompressed under a committed Content-Encoding would yield a stream the
+			// client cannot decode. Fail the ext_proc stream so Envoy aborts the response
+			// rather than delivering corrupt bytes. The committed Content-Encoding state
+			// is left intact.
+			slog.Error("[streaming] failed to re-compress response body; terminating stream to avoid Content-Encoding mismatch",
 				"encoding", execCtx.responseContentEncoding,
 				"error", err,
 			)
-			execCtx.responseContentEncoding = ""
-		} else {
-			outputBody = recompressed
+			execCtx.responseStreamComp.Close()
+			return nil, fmt.Errorf("failed to re-compress streaming response body (encoding %q): %w",
+				execCtx.responseContentEncoding, err)
 		}
+		outputBody = recompressed
 	}
 
 	analyticsData := make(map[string]any)
@@ -1634,9 +1667,6 @@ func TranslateStreamingResponseChunkAction(result *executor.StreamingResponseExe
 		mergeDynamicMetadata(execCtx.dynamicMetadata, dm)
 	}
 
-	// If a policy terminated the stream early (e.g. guardrail intervention), force
-	// EndOfStream so Envoy closes the connection cleanly after delivering the final chunk.
-	endOfStream := originalChunk.EndOfStream || result.StreamTerminated
 	if result.StreamTerminated {
 		slog.Info("[streaming] stream terminated by policy; forcing EndOfStream on final chunk")
 	}