Compare commits

...

6 Commits

4 changed files with 572 additions and 26 deletions

View File

@@ -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.

View File

@@ -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<string, unknown> {
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, unknown>): string {
const nestedError =
failed.error && typeof failed.error === "object" && !Array.isArray(failed.error)
? (failed.error as Record<string, unknown>)
: 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<typeof buildErrorBody> {
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<ReadableStreamDefaultController<Uint8Array>, "enqueue">;
type SseProcessTarget = Pick<TransformStreamDefaultController<Uint8Array>, "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<string, unknown>) || 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<string, unknown>)
: 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<Uint8Array>({
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: {

View File

@@ -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<typeof createStreamFailureFinalizers>[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<Uint8Array>({
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<Uint8Array>({
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<Uint8Array>({
cancel() {
onCancel();
return new Promise<void>(() => {});
},
});
return wrapZedCompletionStream(
new Response(body, {
status: 200,
headers: { "Content-Type": "application/x-ndjson" },
}),
"x_ai",
TEST_MODEL
);
}
async function resolvesWithin(promise: Promise<unknown>, timeoutMs: number): Promise<void> {
let timeout: ReturnType<typeof setTimeout> | undefined;
try {
await Promise.race([
promise,
new Promise<never>((_, reject) => {
timeout = setTimeout(
() => reject(new Error(`operation exceeded ${timeoutMs}ms`)),
timeoutMs
);
}),
]);
} finally {
if (timeout) clearTimeout(timeout);
}
}
async function waitFor(predicate: () => boolean, timeoutMs: number): Promise<void> {
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<Record<string, unknown>> {
return text
.split(/\r?\n/)
.filter((line) => line.startsWith("data: ") && line.slice(6) !== "[DONE]")
.map((line) => JSON.parse(line.slice(6)) as Record<string, unknown>);
}
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<string, string> = {};
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);
});

View File

@@ -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<FixtureResult> {
// 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);
});