fix: decode zstd-compressed SSE proxy surfaces (#335)

* Fix zstd-compressed SSE decoding in proxy surfaces

* Fix stream/web type bridge for zstd SSE

* Guard zstd stream readers against decoded bodies

* Handle stacked zstd stream encodings

src/server/proxy-core/executors/types.test.ts

@@ -1,7 +1,7 @@
 import { gzipSync, zstdCompressSync } from 'node:zlib';
 import { Response } from 'undici';
 import { describe, expect, it } from 'vitest';
-import { materializeErrorResponse, readRuntimeResponseText } from './types.js';
+import { getRuntimeResponseReader, materializeErrorResponse, readRuntimeResponseText } from './types.js';
 
 describe('readRuntimeResponseText', () => {
   it('decompresses zstd responses before reading the body text', async () => {
@@ -54,3 +54,43 @@ describe('materializeErrorResponse', () => {
     expect(materialized.headers.get('content-type')).toBe('application/json; charset=utf-8');
   });
 });
+
+describe('getRuntimeResponseReader', () => {
+  it('falls back to the original stream when the body is already decompressed despite a zstd header', async () => {
+    const payload = 'data: {"ok":true}\n\n';
+    const response = new Response(payload, {
+      status: 200,
+      headers: {
+        'content-encoding': 'zstd',
+        'content-type': 'text/event-stream; charset=utf-8',
+      },
+    });
+
+    const reader = getRuntimeResponseReader(response);
+
+    expect(reader).toBeDefined();
+    const firstChunk = await reader?.read();
+
+    expect(firstChunk?.done).toBe(false);
+    expect(Buffer.from(firstChunk?.value ?? []).toString('utf8')).toBe(payload);
+  });
+
+  it('decompresses stacked streaming encodings when zstd is not the outermost layer', async () => {
+    const payload = 'data: {"ok":true,"kind":"stacked"}\n\n';
+    const response = new Response(gzipSync(zstdCompressSync(Buffer.from(payload))), {
+      status: 200,
+      headers: {
+        'content-encoding': 'zstd, gzip',
+        'content-type': 'text/event-stream; charset=utf-8',
+      },
+    });
+
+    const reader = getRuntimeResponseReader(response);
+
+    expect(reader).toBeDefined();
+    const firstChunk = await reader?.read();
+
+    expect(firstChunk?.done).toBe(false);
+    expect(Buffer.from(firstChunk?.value ?? []).toString('utf8')).toBe(payload);
+  });
+});

src/server/proxy-core/executors/types.ts

@@ -1,9 +1,15 @@
 import {
   brotliDecompressSync,
+  createBrotliDecompress,
+  createGunzip,
+  createInflate,
+  createZstdDecompress,
   gunzipSync,
   inflateSync,
   zstdDecompressSync,
 } from 'node:zlib';
+import { Readable } from 'node:stream';
+import type { ReadableStream as NodeReadableStream } from 'node:stream/web';
 import {
   Response,
   fetch,
@@ -70,10 +76,20 @@ export async function performFetch(
 }
 
 function hasZstdContentEncoding(contentEncoding: string | null): boolean {
-  if (!contentEncoding) return false;
+  return getContentEncodings(contentEncoding).some((encoding) => encoding === 'zstd');
+}
+
+function getContentEncodings(contentEncoding: string | null): string[] {
+  if (!contentEncoding) return [];
   return contentEncoding
     .split(',')
-    .some((encoding) => encoding.trim().toLowerCase() === 'zstd');
+    .map((encoding) => encoding.trim().toLowerCase())
+    .filter(Boolean);
+}
+
+function getOutermostContentEncoding(contentEncoding: string | null): string | null {
+  const encodings = getContentEncodings(contentEncoding);
+  return encodings.length > 0 ? encodings[encodings.length - 1] : null;
 }
 
 function looksLikeZstdFrame(buffer: Buffer): boolean {
@@ -88,11 +104,7 @@ function decodeRuntimeResponseBuffer(buffer: Buffer, contentEncoding: string | n
   if (!contentEncoding) return buffer;
 
   let decoded = buffer;
-  const encodings = contentEncoding
-    .split(',')
-    .map((encoding) => encoding.trim().toLowerCase())
-    .filter(Boolean)
-    .reverse();
+  const encodings = getContentEncodings(contentEncoding).reverse();
 
   for (const encoding of encodings) {
     if (encoding === 'zstd') {
@@ -116,6 +128,37 @@ function decodeRuntimeResponseBuffer(buffer: Buffer, contentEncoding: string | n
   return decoded;
 }
 
+function decodeRuntimeResponseStream(
+  stream: Readable,
+  contentEncoding: string | null,
+): Readable {
+  if (!contentEncoding) return stream;
+
+  let decoded = stream;
+  const encodings = getContentEncodings(contentEncoding).reverse();
+
+  for (const encoding of encodings) {
+    if (encoding === 'zstd') {
+      decoded = decoded.pipe(createZstdDecompress()) as Readable;
+      continue;
+    }
+    if (encoding === 'br') {
+      decoded = decoded.pipe(createBrotliDecompress()) as Readable;
+      continue;
+    }
+    if (encoding === 'gzip' || encoding === 'x-gzip') {
+      decoded = decoded.pipe(createGunzip()) as Readable;
+      continue;
+    }
+    if (encoding === 'deflate') {
+      decoded = decoded.pipe(createInflate()) as Readable;
+      continue;
+    }
+  }
+
+  return decoded;
+}
+
 export async function readRuntimeResponseText(
   response: RuntimeResponse,
 ): Promise<string> {
@@ -136,6 +179,128 @@ export async function readRuntimeResponseText(
   }
 }
 
+function asNodeReadableStream(
+  stream: globalThis.ReadableStream<Uint8Array>,
+): NodeReadableStream<any> {
+  return stream as unknown as NodeReadableStream<any>;
+}
+
+function asWebReadableStream(
+  stream: NodeReadableStream<any>,
+): globalThis.ReadableStream<Uint8Array> {
+  return stream as unknown as globalThis.ReadableStream<Uint8Array>;
+}
+
+function prependReadableStreamChunks(
+  initialChunks: Uint8Array[],
+  sourceReader: ReadableStreamDefaultReader<Uint8Array>,
+): globalThis.ReadableStream<Uint8Array> {
+  const pendingChunks = [...initialChunks];
+  return new ReadableStream<Uint8Array>({
+    async pull(controller) {
+      const nextPendingChunk = pendingChunks.shift();
+      if (nextPendingChunk) {
+        controller.enqueue(nextPendingChunk);
+        return;
+      }
+
+      const nextChunk = await sourceReader.read();
+      if (nextChunk.done) {
+        controller.close();
+        return;
+      }
+
+      controller.enqueue(nextChunk.value);
+    },
+    cancel(reason) {
+      return sourceReader.cancel(reason);
+    },
+  });
+}
+
+async function resolveRuntimeResponseReader(
+  sourceReader: ReadableStreamDefaultReader<Uint8Array>,
+  contentEncoding: string | null,
+): Promise<ReadableStreamDefaultReader<Uint8Array>> {
+  const initialChunks: Uint8Array[] = [];
+  let probeBuffer = Buffer.alloc(0);
+
+  while (probeBuffer.length < 4) {
+    const nextChunk = await sourceReader.read();
+    if (nextChunk.done) {
+      break;
+    }
+    if (!nextChunk.value || nextChunk.value.byteLength === 0) {
+      continue;
+    }
+
+    initialChunks.push(nextChunk.value);
+    probeBuffer = Buffer.concat([probeBuffer, Buffer.from(nextChunk.value)]);
+  }
+
+  if (initialChunks.length === 0) {
+    return sourceReader;
+  }
+
+  const reconstructedBody = prependReadableStreamChunks(initialChunks, sourceReader);
+  const outermostEncoding = getOutermostContentEncoding(contentEncoding);
+  if (outermostEncoding === 'zstd' && !looksLikeZstdFrame(probeBuffer)) {
+    return reconstructedBody.getReader();
+  }
+
+  const decoded = decodeRuntimeResponseStream(
+    Readable.fromWeb(asNodeReadableStream(reconstructedBody)),
+    contentEncoding,
+  );
+  return asWebReadableStream(Readable.toWeb(decoded)).getReader();
+}
+
+export function getRuntimeResponseReader(
+  response: RuntimeResponse,
+): ReadableStreamDefaultReader<Uint8Array> | undefined {
+  const body = response.body as globalThis.ReadableStream<Uint8Array> | null | undefined;
+  if (!body) return undefined;
+
+  const contentEncoding = typeof response.headers?.get === 'function'
+    ? response.headers.get('content-encoding')
+    : null;
+  if (!hasZstdContentEncoding(contentEncoding)) {
+    return body.getReader();
+  }
+
+  const sourceReader = body.getReader();
+  let resolvedReaderPromise: Promise<ReadableStreamDefaultReader<Uint8Array>> | null = null;
+  const ensureResolvedReader = () => {
+    if (!resolvedReaderPromise) {
+      resolvedReaderPromise = resolveRuntimeResponseReader(sourceReader, contentEncoding);
+    }
+    return resolvedReaderPromise;
+  };
+
+  // Keep the public API synchronous while delaying the zstd probe until the first read.
+  return new ReadableStream<Uint8Array>({
+    async pull(controller) {
+      const reader = await ensureResolvedReader();
+      const nextChunk = await reader.read();
+      if (nextChunk.done) {
+        controller.close();
+        return;
+      }
+
+      controller.enqueue(nextChunk.value);
+    },
+    async cancel(reason) {
+      if (!resolvedReaderPromise) {
+        await sourceReader.cancel(reason);
+        return;
+      }
+
+      const reader = await resolvedReaderPromise.catch(() => sourceReader);
+      await reader.cancel(reason);
+    },
+  }).getReader();
+}
+
 export async function materializeErrorResponse(
   response: RuntimeResponse,
 ): Promise<RuntimeResponse> {

src/server/proxy-core/surfaces/chatSurface.ts

@@ -44,7 +44,7 @@ import {
   unwrapGeminiCliPayload,
 } from '../../routes/proxy/geminiCliCompat.js';
 import { summarizeConversationFileInputsInOpenAiBody } from '../capabilities/conversationFileCapabilities.js';
-import { readRuntimeResponseText } from '../executors/types.js';
+import { getRuntimeResponseReader, readRuntimeResponseText } from '../executors/types.js';
 import { detectDownstreamClientContext } from '../../routes/proxy/downstreamClientContext.js';
 import { canRetryProxyChannel, getProxyMaxChannelRetries } from '../../services/proxyChannelRetry.js';
 import {
@@ -665,7 +665,7 @@ export async function handleChatSurfaceRequest(
           });
         } else {
           startSseResponse();
-          const upstreamReader = upstream.body?.getReader();
+          const upstreamReader = getRuntimeResponseReader(upstream);
           const baseReader = String(selected.site.platform || '').trim().toLowerCase() === 'gemini-cli' && upstreamReader
             ? createGeminiCliStreamReader(upstreamReader)
             : upstreamReader;

src/server/proxy-core/surfaces/geminiSurface.ts

@@ -39,7 +39,7 @@ import { dispatchRuntimeRequest } from '../../routes/proxy/runtimeExecutor.js';
 import { detectDownstreamClientContext, type DownstreamClientContext } from '../../routes/proxy/downstreamClientContext.js';
 import { insertProxyLog } from '../../services/proxyLogStore.js';
 import { summarizeConversationFileInputsInOpenAiBody } from '../capabilities/conversationFileCapabilities.js';
-import { readRuntimeResponseText } from '../executors/types.js';
+import { getRuntimeResponseReader, readRuntimeResponseText } from '../executors/types.js';
 import { canRetryProxyChannel, getProxyMaxChannelRetries } from '../../services/proxyChannelRetry.js';
 import {
   buildSurfaceProxyDebugResponseHeaders,
@@ -757,7 +757,7 @@ export async function geminiProxyRoute(app: FastifyInstance) {
           }
 
           if (geminiGenerateContentTransformer.stream.isSseContentType(contentType)) {
-            const upstreamReader = upstream.body?.getReader();
+            const upstreamReader = getRuntimeResponseReader(upstream);
             const reader = isInternalGemini && !isGeminiCliDownstream && upstreamReader
               ? createGeminiCliStreamReader(upstreamReader)
               : upstreamReader;

src/server/proxy-core/surfaces/openAiResponsesSurface.ts

@@ -37,7 +37,7 @@ import {
   unwrapGeminiCliPayload,
 } from '../../routes/proxy/geminiCliCompat.js';
 import { isCodexResponsesSurface } from '../cliProfiles/codexProfile.js';
-import { readRuntimeResponseText } from '../executors/types.js';
+import { getRuntimeResponseReader, readRuntimeResponseText } from '../executors/types.js';
 import { runCodexHttpSessionTask } from '../runtime/codexHttpSessionQueue.js';
 import {
   summarizeConversationFileInputsInOpenAiBody,
@@ -874,7 +874,7 @@ export async function handleOpenAiResponsesSurfaceRequest(
             replayReader = createSingleChunkStreamReader(rawText);
           }
 
-          const upstreamReader = replayReader ?? upstream.body?.getReader();
+          const upstreamReader = replayReader ?? getRuntimeResponseReader(upstream);
           const baseReader = String(selected.site.platform || '').trim().toLowerCase() === 'gemini-cli' && upstreamReader
             ? createGeminiCliStreamReader(upstreamReader)
             : upstreamReader;