fix(sse): fail Zed streams without false success

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.

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<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: {