fix(ollama): yield during dense stream processing

udaymanish6 · udaymanish6 · commit e40e6a473100 · 2026-05-25T17:03:18.000-04:00
diff --git a/extensions/ollama/src/stream.test.ts b/extensions/ollama/src/stream.test.ts
@@ -249,4 +249,45 @@ describe("createOllamaStreamFn thinking events", () => {
       auditContext: "ollama-stream.chat",
     });
   });
+
+  it("yields to the event loop while processing dense native stream chunks", async () => {
+    const chunks = [
+      ...Array.from({ length: 200 }, (_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);
+  });
 });
diff --git a/extensions/ollama/src/stream.ts b/extensions/ollama/src/stream.ts
@@ -50,11 +50,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) => {
+        setImmediate(resolve);
+      });
+      throwIfOllamaStreamAborted(signal);
+    },
+  };
+}
+
 function countMatches(text: string, re: RegExp): number {
   re.lastIndex = 0;
   return Array.from(text.matchAll(re)).length;
@@ -1155,6 +1193,7 @@ export function createOllamaStreamFn(
           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;
@@ -1275,6 +1314,7 @@ export function createOllamaStreamFn(
           };
 
           for await (const chunk of parseNdjsonStream(reader)) {
+            throwIfOllamaStreamAborted(options?.signal);
             const thinkingDelta = chunk.message?.thinking ?? chunk.message?.reasoning;
             if (thinkingDelta) {
               if (!streamStarted) {
@@ -1327,6 +1367,7 @@ export function createOllamaStreamFn(
               finalResponse = chunk;
               break;
             }
+            await cooperativeScheduler.afterEvent();
           }
 
           if (!finalResponse) {
