From 9e95819952fc6ad3edd3b6b94d67803742d7e0d6 Mon Sep 17 00:00:00 2001 From: Cita Date: Mon, 30 Mar 2026 22:46:31 +0800 Subject: [PATCH] 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 | 42 +++- src/server/proxy-core/executors/types.ts | 179 +++++++++++++++++- src/server/proxy-core/surfaces/chatSurface.ts | 4 +- .../proxy-core/surfaces/geminiSurface.ts | 4 +- .../surfaces/openAiResponsesSurface.ts | 4 +- src/server/routes/proxy/chat.stream.test.ts | 32 ++++ .../proxy/responses.codex-oauth.test.ts | 42 ++++ .../routes/proxy/responsesSseFinal.test.ts | 38 ++++ src/server/routes/proxy/responsesSseFinal.ts | 8 +- 9 files changed, 337 insertions(+), 16 deletions(-) diff --git a/src/server/proxy-core/executors/types.test.ts b/src/server/proxy-core/executors/types.test.ts index c1f3682..7f4eeec 100644 --- a/src/server/proxy-core/executors/types.test.ts +++ b/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); + }); +}); diff --git a/src/server/proxy-core/executors/types.ts b/src/server/proxy-core/executors/types.ts index e0d3115..14008db 100644 --- a/src/server/proxy-core/executors/types.ts +++ b/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 { @@ -136,6 +179,128 @@ export async function readRuntimeResponseText( } } +function asNodeReadableStream( + stream: globalThis.ReadableStream, +): NodeReadableStream { + return stream as unknown as NodeReadableStream; +} + +function asWebReadableStream( + stream: NodeReadableStream, +): globalThis.ReadableStream { + return stream as unknown as globalThis.ReadableStream; +} + +function prependReadableStreamChunks( + initialChunks: Uint8Array[], + sourceReader: ReadableStreamDefaultReader, +): globalThis.ReadableStream { + const pendingChunks = [...initialChunks]; + return new ReadableStream({ + 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, + contentEncoding: string | null, +): Promise> { + 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 | undefined { + const body = response.body as globalThis.ReadableStream | 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> | 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({ + 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 { diff --git a/src/server/proxy-core/surfaces/chatSurface.ts b/src/server/proxy-core/surfaces/chatSurface.ts index b07b942..7193cfe 100644 --- a/src/server/proxy-core/surfaces/chatSurface.ts +++ b/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; diff --git a/src/server/proxy-core/surfaces/geminiSurface.ts b/src/server/proxy-core/surfaces/geminiSurface.ts index 02b3810..05f6374 100644 --- a/src/server/proxy-core/surfaces/geminiSurface.ts +++ b/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; diff --git a/src/server/proxy-core/surfaces/openAiResponsesSurface.ts b/src/server/proxy-core/surfaces/openAiResponsesSurface.ts index e376cac..1a3cb7c 100644 --- a/src/server/proxy-core/surfaces/openAiResponsesSurface.ts +++ b/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; diff --git a/src/server/routes/proxy/chat.stream.test.ts b/src/server/routes/proxy/chat.stream.test.ts index cc3d081..e6e696b 100644 --- a/src/server/routes/proxy/chat.stream.test.ts +++ b/src/server/routes/proxy/chat.stream.test.ts @@ -244,6 +244,38 @@ describe('chat proxy stream behavior', () => { expect(response.body).toContain('data: [DONE]'); }); + it('decodes zstd-compressed native SSE chat streams before converting downstream chunks', async () => { + fetchMock.mockResolvedValue(new Response(zstdCompressSync(Buffer.from([ + 'data: {"id":"chatcmpl-zstd-native","model":"upstream-gpt","choices":[{"delta":{"role":"assistant"},"finish_reason":null}]}\n\n', + 'data: {"id":"chatcmpl-zstd-native","model":"upstream-gpt","choices":[{"delta":{"content":"你好，来自 zstd 原生 SSE"},"finish_reason":null}]}\n\n', + 'data: {"id":"chatcmpl-zstd-native","model":"upstream-gpt","choices":[{"delta":{},"finish_reason":"stop"}]}\n\n', + 'data: [DONE]\n\n', + ].join(''))), { + status: 200, + headers: { + 'content-encoding': 'zstd', + 'content-type': 'text/event-stream; charset=utf-8', + }, + })); + + const response = await app.inject({ + method: 'POST', + url: '/v1/chat/completions', + payload: { + model: 'gpt-4o-mini', + stream: true, + messages: [{ role: 'user', content: 'hi' }], + }, + }); + + expect(response.statusCode).toBe(200); + expect(response.headers['content-type']).toContain('text/event-stream'); + expect(response.body).toContain('"chat.completion.chunk"'); + expect(response.body).toContain('你好，来自 zstd 原生 SSE'); + expect(response.body).not.toContain('(�/�'); + expect(response.body).toContain('data: [DONE]'); + }); + it('returns upstream_error for empty non-stream chat responses when empty-content failure is enabled', async () => { config.proxyEmptyContentFailEnabled = true; diff --git a/src/server/routes/proxy/responses.codex-oauth.test.ts b/src/server/routes/proxy/responses.codex-oauth.test.ts index 7948f25..105bb52 100644 --- a/src/server/routes/proxy/responses.codex-oauth.test.ts +++ b/src/server/routes/proxy/responses.codex-oauth.test.ts @@ -1,3 +1,4 @@ +import { zstdCompressSync } from 'node:zlib'; import Fastify, { type FastifyInstance } from 'fastify'; import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'; import { config } from '../../config.js'; @@ -130,6 +131,17 @@ describe('responses proxy codex oauth refresh', () => { }); }; + const createCompressedSseResponse = (chunks: string[], status = 200) => new Response( + zstdCompressSync(Buffer.from(chunks.join(''))), + { + status, + headers: { + 'content-encoding': 'zstd', + 'content-type': 'text/event-stream; charset=utf-8', + }, + }, + ); + beforeAll(async () => { const { responsesProxyRoute } = await import('./responses.js'); app = Fastify(); @@ -712,6 +724,36 @@ describe('responses proxy codex oauth refresh', () => { }); }); + it('decodes zstd-compressed codex responses SSE before relaying native downstream streams', async () => { + fetchMock.mockResolvedValue(createCompressedSseResponse([ + 'event: response.created\n', + 'data: {"type":"response.created","response":{"id":"resp_codex_zstd_stream","model":"gpt-5.4","created_at":1706000000,"status":"in_progress","output":[]}}\n\n', + 'event: response.output_item.added\n', + 'data: {"type":"response.output_item.added","output_index":0,"item":{"id":"msg_codex_zstd_stream","type":"message","role":"assistant","status":"in_progress","content":[]}}\n\n', + 'event: response.output_text.delta\n', + 'data: {"type":"response.output_text.delta","output_index":0,"item_id":"msg_codex_zstd_stream","delta":"你好，来自 zstd responses SSE"}\n\n', + 'event: response.completed\n', + 'data: {"type":"response.completed","response":{"id":"resp_codex_zstd_stream","model":"gpt-5.4","status":"completed","usage":{"input_tokens":3,"output_tokens":1,"total_tokens":4}}}\n\n', + 'data: [DONE]\n\n', + ])); + + const response = await app.inject({ + method: 'POST', + url: '/v1/responses', + payload: { + model: 'gpt-5.4', + input: 'hello codex', + stream: true, + }, + }); + + expect(response.statusCode).toBe(200); + expect(response.body).toContain('event: response.output_text.delta'); + expect(response.body).toContain('你好，来自 zstd responses SSE'); + expect(response.body).not.toContain('(�/�'); + expect(response.body).toContain('data: [DONE]'); + }); + it('preserves codex-required instructions and store fields across responses compatibility retries', async () => { fetchMock.mockResolvedValue(new Response(JSON.stringify({ error: { message: 'upstream_error', type: 'upstream_error' }, diff --git a/src/server/routes/proxy/responsesSseFinal.test.ts b/src/server/routes/proxy/responsesSseFinal.test.ts index eaba0a6..b443c9a 100644 --- a/src/server/routes/proxy/responsesSseFinal.test.ts +++ b/src/server/routes/proxy/responsesSseFinal.test.ts @@ -1,4 +1,6 @@ +import { zstdCompressSync } from 'node:zlib'; import { describe, expect, it } from 'vitest'; +import { Response } from 'undici'; import { collectResponsesFinalPayloadFromSse } from './responsesSseFinal.js'; @@ -73,6 +75,42 @@ describe('collectResponsesFinalPayloadFromSse', () => { }); }); + it('decodes zstd-compressed responses SSE before aggregating the final payload', async () => { + const upstream = new Response(zstdCompressSync(Buffer.from([ + 'event: response.created', + 'data: {"type":"response.created","response":{"id":"resp_zstd_completed","model":"gpt-5.4","created_at":1706000000,"status":"in_progress","output":[]}}', + '', + 'event: response.output_item.added', + 'data: {"type":"response.output_item.added","output_index":0,"item":{"id":"msg_zstd_completed","type":"message","role":"assistant","status":"in_progress","content":[]}}', + '', + 'event: response.output_text.delta', + 'data: {"type":"response.output_text.delta","output_index":0,"item_id":"msg_zstd_completed","delta":"你好，来自 zstd 聚合 SSE"}', + '', + 'event: response.completed', + 'data: {"type":"response.completed","response":{"id":"resp_zstd_completed","model":"gpt-5.4","status":"completed","output":[],"usage":{"input_tokens":3,"output_tokens":1,"total_tokens":4}}}', + '', + 'data: [DONE]', + '', + ].join('\n'))), { + status: 200, + headers: { + 'content-encoding': 'zstd', + 'content-type': 'text/event-stream; charset=utf-8', + }, + }); + + await expect(collectResponsesFinalPayloadFromSse(upstream, 'gpt-5.4')) + .resolves + .toMatchObject({ + payload: { + id: 'resp_zstd_completed', + status: 'completed', + output_text: '你好，来自 zstd 聚合 SSE', + }, + rawText: expect.stringContaining('你好，来自 zstd 聚合 SSE'), + }); + }); + it('returns response.incomplete payloads instead of treating them as upstream failures', async () => { const upstream = { async text() { diff --git a/src/server/routes/proxy/responsesSseFinal.ts b/src/server/routes/proxy/responsesSseFinal.ts index 4863384..d46466f 100644 --- a/src/server/routes/proxy/responsesSseFinal.ts +++ b/src/server/routes/proxy/responsesSseFinal.ts @@ -1,5 +1,6 @@ import { openAiResponsesTransformer } from '../../transformers/openai/responses/index.js'; import { mergeProxyUsage, parseProxyUsage } from '../../services/proxyUsageParser.js'; +import { readRuntimeResponseText } from '../../proxy-core/executors/types.js'; type ResponsesTerminalStatus = 'completed' | 'incomplete'; @@ -420,8 +421,11 @@ export function collectResponsesFinalPayloadFromSseText( } export async function collectResponsesFinalPayloadFromSse( - upstream: { text(): Promise }, + upstream: { text(): Promise; headers?: { get(name: string): string | null } }, modelName: string, ): Promise<{ payload: Record; rawText: string }> { - return collectResponsesFinalPayloadFromSseText(await upstream.text(), modelName); + const rawText = typeof upstream.headers?.get === 'function' + ? await readRuntimeResponseText(upstream as Parameters[0]) + : await upstream.text(); + return collectResponsesFinalPayloadFromSseText(rawText, modelName); }