diff --git a/changelog.d/fixes/zed-hosted-stream-error-boundary.md b/changelog.d/fixes/zed-hosted-stream-error-boundary.md new file mode 100644 index 0000000000..e15742f092 --- /dev/null +++ b/changelog.d/fixes/zed-hosted-stream-error-boundary.md @@ -0,0 +1 @@ +- **fix(providers):** Zed Hosted streaming failures now trigger fallback before content and end partial streams with a sanitized structured error instead of fake assistant text and a normal-success stop. diff --git a/open-sse/executors/zed-hosted.ts b/open-sse/executors/zed-hosted.ts index ef1ae4ade6..ba66706f10 100644 --- a/open-sse/executors/zed-hosted.ts +++ b/open-sse/executors/zed-hosted.ts @@ -44,6 +44,8 @@ import { zedLlmFetch, type ZedCredentials, } from "../shared/zedAuth.ts"; +import { buildErrorBody } from "../utils/error.ts"; +import { hasUsefulStreamContent } from "../utils/streamReadiness.ts"; import { resolveSuppressThinkClose, THINKING_MARKER_HEADER } from "../utils/thinkCloseMarker.ts"; // Wire values for the `provider` field of POST /completions. These are NOT @@ -122,37 +124,72 @@ function convertProviderEvent( return event; } -function createErrorChunk(model: string, message: string): Record { - return { - id: `chatcmpl-zed-error-${Date.now()}`, - object: "chat.completion.chunk", - created: Math.floor(Date.now() / 1000), - model, - choices: [{ index: 0, delta: { content: `[Zed error] ${message}` }, finish_reason: "stop" }], - }; +const MAX_ZED_FAILURE_MESSAGE_LENGTH = 512; +const MAX_PENDING_ZED_OUTPUT_LENGTH = 64 * 1024; +const ZED_STREAM_FAILURE_PUBLIC_MESSAGE = "Zed upstream stream failed"; + +function boundedFailureText(value: unknown): string | null { + if (typeof value !== "string" && typeof value !== "number") return null; + const text = String(value).trim(); + return text ? text.slice(0, MAX_ZED_FAILURE_MESSAGE_LENGTH) : null; +} + +function extractZedFailureMessage(failed: Record): string { + const nestedError = + failed.error && typeof failed.error === "object" && !Array.isArray(failed.error) + ? (failed.error as Record) + : null; + const candidates = [ + failed.message, + nestedError?.message, + typeof failed.error === "object" ? undefined : failed.error, + failed.code, + nestedError?.code, + ]; + for (const candidate of candidates) { + const text = boundedFailureText(candidate); + if (text) return text; + } + return "request failed"; +} + +function createErrorChunk(message: string): ReturnType { + return buildErrorBody(502, `Zed stream failed: ${message}`, undefined, { + type: "upstream_error", + code: "ZED_STREAM_FAILED", + }); } /** - * The single controller capability these SSE helpers use. They only ever enqueue — - * never `close()`, never read `desiredSize` — so typing them by that one method lets - * the same code serve both stream kinds. The wider + * The controller capabilities these SSE helpers use. Normal frames only enqueue; + * terminal failures also terminate so they do not depend on the upstream socket + * eventually reaching EOF. Narrow controller types keep the helpers honest. The wider * `ReadableStreamDefaultController` annotation rejected every call site, because the * helpers are driven from a TransformStream and `TransformStreamDefaultController` * has no `close()`. */ type SseEnqueueTarget = Pick, "enqueue">; +type SseProcessTarget = Pick, "enqueue" | "terminate">; + +function serializeSseObject(chunk: unknown): string { + if (!chunk) return ""; + let serialized = ""; + const items = Array.isArray(chunk) ? chunk : [chunk]; + for (const item of items) { + if (!item) continue; + serialized += `data: ${JSON.stringify(item)}\n\n`; + } + return serialized; +} function enqueueSseObject( controller: SseEnqueueTarget, encoder: TextEncoder, chunk: unknown ): void { - if (!chunk) return; - const items = Array.isArray(chunk) ? chunk : [chunk]; - for (const item of items) { - if (!item) continue; - controller.enqueue(encoder.encode(`data: ${JSON.stringify(item)}\n\n`)); - } + const serialized = serializeSseObject(chunk); + if (!serialized) return; + controller.enqueue(encoder.encode(serialized)); } type ZedLine = { done?: true; status?: unknown; event?: unknown } | null; @@ -226,16 +263,47 @@ function wrapZedCompletionStream( } let buffer = ""; let done = false; + let providerOutputForwarded = false; + let pendingProviderOutput = ""; + let pendingFailure: (Error & { statusCode: number }) | null = null; + + const forwardProviderOutput = (controller: SseEnqueueTarget, chunk: unknown) => { + const serialized = serializeSseObject(chunk); + if (!serialized) return; + if (providerOutputForwarded) { + controller.enqueue(encoder.encode(serialized)); + return; + } + + // A role/bootstrap-only chunk makes ensureStreamReadiness release the response before any + // model output exists. If the next chunk is status.failed, downstream read-ahead can discard + // the first real content while propagating the error. Hold structural frames until the first + // substantive text/reasoning/tool delta, then release them atomically with that output. + const outputWithBootstrap = pendingProviderOutput + serialized; + if (!hasUsefulStreamContent(outputWithBootstrap)) { + pendingProviderOutput = + outputWithBootstrap.length <= MAX_PENDING_ZED_OUTPUT_LENGTH + ? outputWithBootstrap + : serialized.length <= MAX_PENDING_ZED_OUTPUT_LENGTH + ? serialized + : ""; + return; + } + controller.enqueue(encoder.encode(outputWithBootstrap)); + pendingProviderOutput = ""; + providerOutputForwarded = true; + }; const finish = (controller: SseEnqueueTarget) => { if (done) return; const finalChunk = convertProviderEvent(provider, null, state); - enqueueSseObject(controller, encoder, finalChunk); - controller.enqueue(encoder.encode("data: [DONE]\n\n")); + const finalOutput = `${pendingProviderOutput}${serializeSseObject(finalChunk)}data: [DONE]\n\n`; + pendingProviderOutput = ""; + controller.enqueue(encoder.encode(finalOutput)); done = true; }; - const processLine = (line: string, controller: SseEnqueueTarget) => { + const processLine = (line: string, controller: SseProcessTarget) => { if (done) return; const payload = unwrapZedLine(line); if (!payload) return; @@ -246,17 +314,29 @@ function wrapZedCompletionStream( if (payload.status) { const status = normalizeStatus(payload.status); if (status?.type === "failed" || status?.failed) { - const failed = (status.failed as Record) || status; - const message = String(failed.message || failed.error || failed.code || "request failed"); - enqueueSseObject(controller, encoder, createErrorChunk(model, message)); - finish(controller); + const failed = + status.failed && typeof status.failed === "object" && !Array.isArray(status.failed) + ? (status.failed as Record) + : status; + if (providerOutputForwarded) { + pendingFailure = Object.assign(new Error(ZED_STREAM_FAILURE_PUBLIC_MESSAGE), { + statusCode: 502, + }); + done = true; + controller.terminate(); + return; + } + pendingProviderOutput = ""; + enqueueSseObject(controller, encoder, createErrorChunk(extractZedFailureMessage(failed))); + done = true; + controller.terminate(); } else if (status?.type === "stream_ended" || status === ("stream_ended" as unknown)) { finish(controller); } return; } const converted = convertProviderEvent(provider, payload.event, state); - enqueueSseObject(controller, encoder, converted); + forwardProviderOutput(controller, converted); }; const transformed = response.body.pipeThrough( @@ -281,7 +361,47 @@ function wrapZedCompletionStream( }) ); - return new Response(transformed, { + // `TransformStreamDefaultController.error()` discards already-enqueued output. A failed + // status can share one upstream network chunk with the last content delta, so erroring the + // transform immediately would erase that partial answer. Drain the transformed chunks through + // a backpressure-aware reader first, then reject the next read with the fixed public error. + // The normal chat pipeline turns that rejection into its client-format terminal frame and + // records the 502 through the existing failure finalizers. + const transformedReader = transformed.getReader(); + let guardedStreamCancelled = false; + const cancelTransformedReader = (reason: unknown) => { + if (guardedStreamCancelled) return; + guardedStreamCancelled = true; + // Client cancellation must settle independently of an upstream body whose cancel hook hangs. + // Request cancellation once, but do not await provider cleanup on the client-facing boundary. + void transformedReader.cancel(reason).catch(() => { + console.debug("[ZED] upstream stream cancellation rejected"); + }); + }; + const guardedStream = new ReadableStream({ + async pull(controller) { + try { + const next = await transformedReader.read(); + if (guardedStreamCancelled) return; + if (!next.done) { + controller.enqueue(next.value); + return; + } + if (pendingFailure) { + controller.error(pendingFailure); + return; + } + controller.close(); + } catch (error) { + if (!guardedStreamCancelled) controller.error(error); + } + }, + cancel(reason) { + cancelTransformedReader(reason); + }, + }); + + return new Response(guardedStream, { status: response.status, statusText: response.statusText, headers: { diff --git a/tests/fixtures/zed-hosted-stream-error-boundary-child.ts b/tests/fixtures/zed-hosted-stream-error-boundary-child.ts new file mode 100644 index 0000000000..c5c8f8352d --- /dev/null +++ b/tests/fixtures/zed-hosted-stream-error-boundary-child.ts @@ -0,0 +1,341 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +// This file is executed only by the process-isolated unit-test wrapper. State +// mutations and repository imports must remain here, never in the parent test. +const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-zed-stream-data-")); +const TEST_PLUGINS_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-zed-stream-plugins-")); +const originalDataDir = process.env.DATA_DIR; +const originalPluginsDir = process.env.OMNIROUTE_PLUGINS_DIR; +const originalFetch = globalThis.fetch; +let networkCalls = 0; + +process.env.DATA_DIR = TEST_DATA_DIR; +process.env.OMNIROUTE_PLUGINS_DIR = TEST_PLUGINS_DIR; +globalThis.fetch = async () => { + networkCalls += 1; + throw new Error("Unexpected network access in Zed stream boundary test"); +}; + +const core = await import("../../src/lib/db/core.ts"); +const loggerResource = await import("../../src/shared/utils/loggerResource.ts"); +const { __test__ } = await import("../../open-sse/executors/zed-hosted.ts"); +const { FORMATS } = await import("../../open-sse/translator/formats.ts"); +const { assembleStreamingPipeline } = + await import("../../open-sse/handlers/chatCore/streamingPipeline.ts"); +const { createPassthroughStreamWithLogger } = await import("../../open-sse/utils/stream.ts"); +const { createStreamFailureFinalizers } = + await import("../../open-sse/utils/streamFailureFinalization.ts"); +const { createStreamController } = await import("../../open-sse/utils/streamHandler.ts"); +const { ensureStreamReadiness } = await import("../../open-sse/utils/streamReadiness.ts"); +const { wrapZedCompletionStream } = __test__; + +type StreamCompletionEvent = Parameters< + Parameters[0]["onStreamComplete"] +>[0]; + +const RAW_FAILURE = "Bearer TOP_SECRET /srv/omniroute/zed-handler.ts:42 api_key=zed-secret"; +const TEST_MODEL = "grok-test-zed-stream-boundary"; +const TEST_CONNECTION_ID = "zed-stream-boundary-partial-connection"; + +function failedStatusLine(): string { + return JSON.stringify({ status: { failed: { message: RAW_FAILURE } } }); +} + +function nestedFailedStatusLine(): string { + return JSON.stringify({ + status: { + type: "failed", + error: { message: `${RAW_FAILURE} ${"x".repeat(2_000)}` }, + }, + }); +} + +function wrapOpenNdjson(lines: unknown[]): Response { + const encoder = new TextEncoder(); + const body = lines + .map((line) => (typeof line === "string" ? line : JSON.stringify(line))) + .join("\n"); + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(`${body}\n`)); + // Keep the upstream open: status.failed must terminate the wrapped stream itself. + }, + cancel() {}, + }); + return wrapZedCompletionStream( + new Response(stream, { + status: 200, + headers: { "Content-Type": "application/x-ndjson" }, + }), + "x_ai", + TEST_MODEL + ); +} + +function wrapOpenFailedNdjson(): Response { + const encoder = new TextEncoder(); + const body = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(`${failedStatusLine()}\n`)); + // Deliberately stay open: status.failed is terminal by itself and must not + // depend on the upstream socket eventually reaching EOF. + }, + cancel() {}, + }); + return wrapZedCompletionStream( + new Response(body, { + status: 200, + headers: { "Content-Type": "application/x-ndjson" }, + }), + "x_ai", + TEST_MODEL + ); +} + +function wrapStalledNdjson(onCancel: () => void): Response { + const body = new ReadableStream({ + cancel() { + onCancel(); + return new Promise(() => {}); + }, + }); + return wrapZedCompletionStream( + new Response(body, { + status: 200, + headers: { "Content-Type": "application/x-ndjson" }, + }), + "x_ai", + TEST_MODEL + ); +} + +async function resolvesWithin(promise: Promise, timeoutMs: number): Promise { + let timeout: ReturnType | undefined; + try { + await Promise.race([ + promise, + new Promise((_, reject) => { + timeout = setTimeout( + () => reject(new Error(`operation exceeded ${timeoutMs}ms`)), + timeoutMs + ); + }), + ]); + } finally { + if (timeout) clearTimeout(timeout); + } +} + +async function waitFor(predicate: () => boolean, timeoutMs: number): Promise { + const deadline = Date.now() + timeoutMs; + while (!predicate() && Date.now() < deadline) { + await new Promise((resolve) => setTimeout(resolve, 1)); + } + assert.equal(predicate(), true, `condition was not met within ${timeoutMs}ms`); +} + +function parseSsePayloads(text: string): Array> { + return text + .split(/\r?\n/) + .filter((line) => line.startsWith("data: ") && line.slice(6) !== "[DONE]") + .map((line) => JSON.parse(line.slice(6)) as Record); +} + +function assertNoSensitiveFailureText(text: string): void { + assert.doesNotMatch(text, /TOP_SECRET|zed-secret|\/srv\/omniroute\/zed-handler\.ts/); +} + +test.after(async () => { + core.resetDbInstance(); + await loggerResource.closeSharedLoggerResource(); + globalThis.fetch = originalFetch; + if (originalDataDir === undefined) delete process.env.DATA_DIR; + else process.env.DATA_DIR = originalDataDir; + if (originalPluginsDir === undefined) delete process.env.OMNIROUTE_PLUGINS_DIR; + else process.env.OMNIROUTE_PLUGINS_DIR = originalPluginsDir; + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 }); + fs.rmSync(TEST_PLUGINS_DIR, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 }); +}); + +test("zed-hosted pre-content status.failed becomes a sanitized 502 readiness failure", async () => { + const readiness = await ensureStreamReadiness(wrapOpenFailedNdjson(), { + timeoutMs: 100, + provider: "zed-hosted", + model: TEST_MODEL, + }); + + assert.equal(readiness.ok, false, "the structured error must remain eligible for fallback"); + if (readiness.ok) assert.fail("pre-content Zed failure must not make the stream ready"); + assert.equal(readiness.response.status, 502); + assert.equal(readiness.code, "STREAM_EARLY_EOF"); + + const bodyText = await readiness.response.text(); + const body = JSON.parse(bodyText) as { + error: { message: string; type: string; code: string }; + upstream_details?: { error?: { message?: string } }; + }; + assert.equal(body.error.type, "stream_early_eof"); + assert.equal(body.error.code, "STREAM_EARLY_EOF"); + assert.match(body.upstream_details?.error?.message ?? "", /Zed stream failed/i); + assertNoSensitiveFailureText(bodyText); + assert.equal(networkCalls, 0); +}); + +test("zed-hosted partial failure reaches stream finalization and persistence as 502", async () => { + const roleChunk = { + event: { + id: "chatcmpl-zed-partial", + object: "chat.completion.chunk", + choices: [{ index: 0, delta: { role: "assistant" }, finish_reason: null }], + }, + }; + const contentChunk = { + event: { + id: "chatcmpl-zed-partial", + object: "chat.completion.chunk", + choices: [{ index: 0, delta: { content: "partial answer" }, finish_reason: null }], + }, + }; + const readiness = await ensureStreamReadiness( + wrapOpenNdjson([ + roleChunk, + contentChunk, + nestedFailedStatusLine(), + { event: { ignored: "after failure" } }, + ]), + { + timeoutMs: 100, + provider: "zed-hosted", + model: TEST_MODEL, + } + ); + + assert.equal(readiness.ok, true, "partial model output must remain deliverable"); + const completionEvents: StreamCompletionEvent[] = []; + const persistedFailures: Array<{ + connectionId: string; + model: string; + status: number; + code?: string; + }> = []; + const streamFailures: Array<{ status: number; message: string; code?: string; type?: string }> = + []; + const pipelineErrors: Array<{ message: string; statusCode: number }> = []; + let streamCompletionRecorded = false; + let failureCompletionRecorded = false; + + const recordCompletion = (payload: StreamCompletionEvent): void => { + if (streamCompletionRecorded) return; + streamCompletionRecorded = true; + if (payload.status !== 200) failureCompletionRecorded = true; + completionEvents.push(payload); + }; + const finalizers = createStreamFailureFinalizers({ + isFailureCompletionRecorded: () => failureCompletionRecorded, + isStreamCompletionRecorded: () => streamCompletionRecorded, + onStreamComplete: recordCompletion, + persistFailureUsage: (status, code) => + persistedFailures.push({ + connectionId: TEST_CONNECTION_ID, + model: TEST_MODEL, + status, + code, + }), + onStreamFailure: (failure) => streamFailures.push(failure), + }); + const streamController = createStreamController({ + onError: (event) => { + pipelineErrors.push({ message: event.message, statusCode: event.statusCode }); + return finalizers.onPipelineStreamError(event); + }, + provider: "zed-hosted", + model: TEST_MODEL, + connectionId: TEST_CONNECTION_ID, + clientResponseFormat: FORMATS.OPENAI, + }); + const transformStream = createPassthroughStreamWithLogger( + "zed-hosted", + null, + null, + TEST_MODEL, + TEST_CONNECTION_ID, + { messages: [{ role: "user", content: "test" }] }, + recordCompletion, + null, + finalizers.handleStreamFailure, + FORMATS.OPENAI + ); + const responseHeaders: Record = {}; + const finalStream = assembleStreamingPipeline({ + providerResponse: readiness.response, + transformStream, + streamController, + createPiiTransform: null, + clientRawRequestHeaders: null, + clientResponseFormat: FORMATS.OPENAI, + echoModel: null, + responseHeaders, + }); + const text = await new Response(finalStream, { headers: responseHeaders }).text(); + const payloads = parseSsePayloads(text); + const errorPayload = payloads.find((payload) => "error" in payload) as + { error: { message: string; type: string; code: string } } | undefined; + + assert.match(text, /partial answer/); + assert.ok(errorPayload, "the stream handler must emit its format-safe terminal error"); + assert.equal(errorPayload.error.type, "server_error"); + assert.equal(errorPayload.error.code, "server_error"); + assert.equal(errorPayload.error.message, "Zed upstream stream failed"); + assert.match(text, /"finish_reason":"error"/); + assert.doesNotMatch(text, /\[Zed error\]|"finish_reason":"stop"|response\.failed/); + assert.doesNotMatch(text, /"ignored":"after failure"/); + assertNoSensitiveFailureText(text); + + assert.equal(completionEvents.length, 1, "the failure must finalize exactly once"); + assert.equal(completionEvents[0].status, 502); + assert.equal(completionEvents[0].error, "Zed upstream stream failed"); + assert.equal(completionEvents[0].errorCode, "stream_pipeline_error"); + assert.deepEqual(persistedFailures, [ + { + connectionId: TEST_CONNECTION_ID, + model: TEST_MODEL, + status: 502, + code: "stream_pipeline_error", + }, + ]); + assert.deepEqual(streamFailures, [ + { + status: 502, + message: "Zed upstream stream failed", + code: "stream_pipeline_error", + type: "stream_error", + }, + ]); + assert.deepEqual(pipelineErrors, [{ message: "Zed upstream stream failed", statusCode: 502 }]); + assertNoSensitiveFailureText(JSON.stringify(completionEvents)); + assert.equal(networkCalls, 0); +}); + +test("zed-hosted client cancellation does not await a stalled upstream cancel hook", async () => { + let upstreamCancelCalls = 0; + const response = wrapStalledNdjson(() => { + upstreamCancelCalls += 1; + }); + assert.ok(response.body); + + const reader = response.body.getReader(); + const pendingRead = reader.read(); + await resolvesWithin(reader.cancel("client disconnected"), 100); + const readResult = await pendingRead; + assert.equal(readResult.done, true); + await waitFor(() => upstreamCancelCalls === 1, 100); + + await resolvesWithin(reader.cancel("duplicate cancel"), 100); + await new Promise((resolve) => setTimeout(resolve, 10)); + assert.equal(upstreamCancelCalls, 1, "the upstream cancel hook must be requested exactly once"); + assert.equal(networkCalls, 0); +}); diff --git a/tests/unit/zed-hosted-stream-error-boundary.test.ts b/tests/unit/zed-hosted-stream-error-boundary.test.ts new file mode 100644 index 0000000000..1565c1c1bd --- /dev/null +++ b/tests/unit/zed-hosted-stream-error-boundary.test.ts @@ -0,0 +1,84 @@ +import assert from "node:assert/strict"; +import { spawn } from "node:child_process"; +import { fileURLToPath } from "node:url"; +import test from "node:test"; + +const repoRoot = fileURLToPath(new URL("../../", import.meta.url)); +const fixturePath = fileURLToPath( + new URL("../fixtures/zed-hosted-stream-error-boundary-child.ts", import.meta.url) +); + +type FixtureResult = { + code: number | null; + signal: NodeJS.Signals | null; + stdout: string; + stderr: string; +}; + +function runFixture(): Promise { + // Keep the parent process pristine: the fast unit suite can run files with + // --test-isolation=none, so all stateful imports and mutations live in the child. + const childEnv: NodeJS.ProcessEnv = { + PATH: process.env.PATH, + NODE_PATH: process.env.NODE_PATH, + LANG: process.env.LANG, + LC_ALL: process.env.LC_ALL, + TZ: process.env.TZ, + TMPDIR: process.env.TMPDIR, + NODE_ENV: "test", + API_KEY_SECRET: "zed-boundary-test-only-secret-with-32-plus-characters", + DISABLE_SQLITE_AUTO_BACKUP: "true", + NO_COLOR: "1", + }; + // Inheriting this marker makes Node silently skip the nested --test run. + delete childEnv.NODE_TEST_CONTEXT; + + return new Promise((resolve, reject) => { + const child = spawn(process.execPath, ["--import", "tsx/esm", "--test", fixturePath], { + cwd: repoRoot, + env: childEnv, + stdio: ["ignore", "pipe", "pipe"], + }); + let stdout = ""; + let stderr = ""; + let timedOut = false; + + child.stdout.setEncoding("utf8"); + child.stderr.setEncoding("utf8"); + child.stdout.on("data", (chunk: string) => { + stdout += chunk; + }); + child.stderr.on("data", (chunk: string) => { + stderr += chunk; + }); + + const timeout = setTimeout(() => { + timedOut = true; + child.kill("SIGKILL"); + }, 120_000); + + child.once("error", (error) => { + clearTimeout(timeout); + reject(error); + }); + child.once("close", (code, signal) => { + clearTimeout(timeout); + if (timedOut) { + reject(new Error("Zed stream error boundary fixture timed out after 120 seconds")); + return; + } + resolve({ code, signal, stdout, stderr }); + }); + }); +} + +test("Zed stream error boundary passes in a process-isolated runtime", async () => { + const result = await runFixture(); + const output = `${result.stdout}\n${result.stderr}`; + + assert.equal(result.signal, null, output.slice(-12_000)); + assert.equal(result.code, 0, output.slice(-12_000)); + assert.match(output, /(?:^|\s)tests\s+3(?:\s|$)/m); + assert.match(output, /(?:^|\s)pass\s+3(?:\s|$)/m); + assert.match(output, /(?:^|\s)fail\s+0(?:\s|$)/m); +});