fix(ollama): yield during dense stream processing

udaymanish6 · vincentkoc · commit 3a478d76384d · 2026-05-31T17:34:27.000+02:00
diff --git a/extensions/ollama/src/stream.test.ts b/extensions/ollama/src/stream.test.ts
@@ -356,4 +356,90 @@ describe("createOllamaStreamFn thinking events", () => {
       arguments: { path: "/path/to/file", line_start: 1, line_end: 400 },
     });
   });
+
+  it("yields to the event loop while processing dense native stream chunks", async () => {
+    const chunks = [
+      ...Array.from({ length: 65 }, (_value, index) => ({
+        model: "qwen3.5",
+        created_at: `2026-01-01T00:00:${String(index % 60).padStart(2, "0")}Z`,
+        message: { role: "assistant" as const, content: "x" },
+        done: false,
+      })),
+      makeOllamaResponse({ content: "" }),
+    ];
+    const body = makeNdjsonBody(chunks);
+    fetchWithSsrFGuardMock.mockResolvedValue({
+      response: new Response(body, { status: 200 }),
+      release: vi.fn(async () => undefined),
+    });
+
+    const streamFn = createOllamaStreamFn("http://localhost:11434");
+    const stream = streamFn(
+      { api: "ollama", provider: "ollama", id: "qwen3.5", contextWindow: 65536 } as never,
+      { messages: [{ role: "user", content: "test" }] } as never,
+      {},
+    );
+
+    let timerFired = false;
+    const timerPromise = new Promise<void>((resolve) => {
+      setTimeout(() => {
+        timerFired = true;
+        resolve();
+      }, 0);
+    });
+    let yieldedBeforeDone = false;
+    for await (const event of stream as AsyncIterable<{ type: string }>) {
+      if (timerFired && event.type !== "done") {
+        yieldedBeforeDone = true;
+      }
+    }
+    await timerPromise;
+
+    expect(yieldedBeforeDone).toBe(true);
+  });
+
+  it("reports caller aborts during dense native stream processing as aborted", async () => {
+    const chunks = [
+      ...Array.from({ length: 65 }, (_value, index) => ({
+        model: "qwen3.5",
+        created_at: `2026-01-01T00:00:${String(index % 60).padStart(2, "0")}Z`,
+        message: { role: "assistant" as const, content: "x" },
+        done: false,
+      })),
+      makeOllamaResponse({ content: "" }),
+    ];
+    const body = makeNdjsonBody(chunks);
+    fetchWithSsrFGuardMock.mockResolvedValue({
+      response: new Response(body, { status: 200 }),
+      release: vi.fn(async () => undefined),
+    });
+
+    const controller = new AbortController();
+    const streamFn = createOllamaStreamFn("http://localhost:11434");
+    const stream = streamFn(
+      { api: "ollama", provider: "ollama", id: "qwen3.5", contextWindow: 65536 } as never,
+      { messages: [{ role: "user", content: "test" }] } as never,
+      { signal: controller.signal },
+    );
+
+    setTimeout(() => {
+      controller.abort();
+    }, 0);
+
+    const events: Array<{ type: string; reason?: string; error?: { stopReason?: string } }> = [];
+    for await (const event of stream as AsyncIterable<{
+      type: string;
+      reason?: string;
+      error?: { stopReason?: string };
+    }>) {
+      events.push(event);
+    }
+
+    const lastEvent = events.at(-1);
+    expect(lastEvent).toMatchObject({
+      type: "error",
+      reason: "aborted",
+      error: { stopReason: "aborted" },
+    });
+  });
 });
diff --git a/extensions/ollama/src/stream.ts b/extensions/ollama/src/stream.ts
@@ -51,11 +51,49 @@ const log = createSubsystemLogger("ollama-stream");
 
 export const OLLAMA_NATIVE_BASE_URL = OLLAMA_DEFAULT_BASE_URL;
 
+const OLLAMA_STREAM_COOPERATIVE_YIELD_INTERVAL_MS = 12;
+const OLLAMA_STREAM_COOPERATIVE_YIELD_MAX_EVENTS = 64;
 const GARBLED_VISIBLE_TEXT_MODEL_RE = /\b(?:glm|kimi)\b/i;
 const GARBLED_VISIBLE_TEXT_MIN_CHARS = 80;
 const GARBLED_VISIBLE_TEXT_SYMBOL_RE = /[$#%&="'_~`^|\\/*+\-[\]{}()<>:;,.!?]/gu;
 const LETTER_OR_DIGIT_RE = /[\p{L}\p{N}]/gu;
 
+type OllamaStreamCooperativeScheduler = {
+  afterEvent: () => Promise<void>;
+};
+
+function throwIfOllamaStreamAborted(signal?: AbortSignal): void {
+  if (signal?.aborted) {
+    throw new Error("Request was aborted");
+  }
+}
+
+function createOllamaStreamCooperativeScheduler(
+  signal?: AbortSignal,
+): OllamaStreamCooperativeScheduler {
+  let lastYieldedAt = Date.now();
+  let eventsSinceYield = 0;
+  return {
+    async afterEvent() {
+      throwIfOllamaStreamAborted(signal);
+      eventsSinceYield += 1;
+      const now = Date.now();
+      if (
+        eventsSinceYield < OLLAMA_STREAM_COOPERATIVE_YIELD_MAX_EVENTS &&
+        now - lastYieldedAt < OLLAMA_STREAM_COOPERATIVE_YIELD_INTERVAL_MS
+      ) {
+        return;
+      }
+      eventsSinceYield = 0;
+      lastYieldedAt = now;
+      await new Promise<void>((resolve) => {
+        setTimeout(resolve, 0);
+      });
+      throwIfOllamaStreamAborted(signal);
+    },
+  };
+}
+
 function countMatches(text: string, re: RegExp): number {
   re.lastIndex = 0;
   return Array.from(text.matchAll(re)).length;
@@ -523,18 +561,22 @@ function buildStreamAssistantMessage(params: {
 
 function buildStreamErrorAssistantMessage(params: {
   model: StreamModelDescriptor;
+  stopReason: Extract<StopReason, "aborted" | "error">;
   errorMessage: string;
   timestamp?: number;
-}): AssistantMessage & { stopReason: "error"; errorMessage: string } {
+}): AssistantMessage & {
+  stopReason: Extract<StopReason, "aborted" | "error">;
+  errorMessage: string;
+} {
   return {
     ...buildStreamAssistantMessage({
       model: params.model,
       content: [],
-      stopReason: "error",
+      stopReason: params.stopReason,
       usage: buildUsageWithNoCost({}),
       timestamp: params.timestamp,
     }),
-    stopReason: "error",
+    stopReason: params.stopReason,
     errorMessage: params.errorMessage,
   };
 }
@@ -1169,6 +1211,7 @@ function createRawOllamaStreamFn(
           let pendingFinalVisibleContent: string | undefined;
           const modelInfo = { api: model.api, provider: model.provider, id: model.id };
           const visibleContentSanitizer = createOllamaVisibleContentSanitizer(model.id);
+          const cooperativeScheduler = createOllamaStreamCooperativeScheduler(options?.signal);
           let streamStarted = false;
           let thinkingStarted = false;
           let thinkingEnded = false;
@@ -1289,6 +1332,7 @@ function createRawOllamaStreamFn(
           };
 
           for await (const chunk of parseNdjsonStream(reader)) {
+            throwIfOllamaStreamAborted(options?.signal);
             const thinkingDelta = chunk.message?.thinking ?? chunk.message?.reasoning;
             if (thinkingDelta) {
               if (!streamStarted) {
@@ -1341,6 +1385,7 @@ function createRawOllamaStreamFn(
               finalResponse = chunk;
               break;
             }
+            await cooperativeScheduler.afterEvent();
           }
 
           if (!finalResponse) {
@@ -1392,11 +1437,13 @@ function createRawOllamaStreamFn(
           await release();
         }
       } catch (err) {
+        const stopReason = options?.signal?.aborted ? "aborted" : "error";
         stream.push({
           type: "error",
-          reason: "error",
+          reason: stopReason,
           error: buildStreamErrorAssistantMessage({
             model,
+            stopReason,
             errorMessage: formatErrorMessage(err),
           }),
         });
