From e79affa257ea5076d05dfac6b9b91f91c3d13fe5 Mon Sep 17 00:00:00 2001
From: wolfy <adam@shockvpn.com>
Date: Sun, 31 May 2026 18:10:51 -0500
Subject: [PATCH] fix(agent): stream Ollama responses to prevent prefill
 timeout

---
 agents/ollama.ts | 70 +++++++++++++++++++++++++++---------------------
 mcp/readFile.ts  |  2 +-
 2 files changed, 41 insertions(+), 31 deletions(-)

diff --git a/agents/ollama.ts b/agents/ollama.ts
index 3a40b39..b25dae0 100644
--- a/agents/ollama.ts
+++ b/agents/ollama.ts
@@ -152,39 +152,50 @@ async function runOllamaLoop(ctx: AgentRunContext): Promise<AgentResult> {
     iterations++;
     log.info(`» Ollama turn ${iterations}/${MAX_ITERATIONS}...`);
 
-    let response: Awaited<ReturnType<typeof ollama.chat>>;
+    // Stream the response so the activity monitor sees tokens arriving during
+    // long prefill. Without streaming, a 50k+ token context can take 200-300s
+    // of silent prefill before the first token, tripping the activity timeout.
+    let assistantMessage: Message;
+    let promptTokens: number | undefined;
+    let evalTokens: number | undefined;
     try {
-      response = await retry(
-        () =>
-          ollama.chat({
-            model,
-            messages,
-            tools,
-            keep_alive: -1,
-            think: false,
-            options: { num_ctx: 262144, temperature: 0.1 },
-          }),
-        {
-          delaysMs: [3_000, 8_000],
-          shouldRetry: (err) => {
-            const msg = err instanceof Error ? err.message : String(err);
-            return /unexpected EOF|XML syntax error|ECONNRESET|ETIMEDOUT|fetch failed/i.test(
-              msg,
-            );
-          },
-          label: `Ollama turn ${iterations}`,
-        },
-      );
-    } catch (err) {
-      await unloadModel(ollama, model);
+      const stream = await ollama.chat({
+        model,
+        messages,
+        tools,
+        keep_alive: -1,
+        think: false,
+        options: { num_ctx: 262144, temperature: 0.1 },
+        stream: true,
+      });
 
-      const lastError = err instanceof Error ? err.message : String(err);
-      log.error(`Ollama error: ${lastError}`);
-      return { success: false, error: `Ollama request failed: ${lastError}` };
+      let content = "";
+      let streamToolCalls: ToolCall[] | undefined;
+      let chunkCount = 0;
+      for await (const chunk of stream) {
+        content += chunk.message.content ?? "";
+        if (chunk.message.tool_calls?.length) streamToolCalls = chunk.message.tool_calls;
+        if (chunk.prompt_eval_count !== undefined) promptTokens = chunk.prompt_eval_count;
+        if (chunk.eval_count !== undefined) evalTokens = chunk.eval_count;
+        chunkCount++;
+        // Log a heartbeat every 50 chunks so the activity monitor stays alive
+        if (chunkCount % 50 === 0) log.debug(`  streaming… (${chunkCount} chunks)`);
+      }
+      assistantMessage = { role: "assistant", content, tool_calls: streamToolCalls };
+    } catch (err) {
+      const msg = err instanceof Error ? err.message : String(err);
+      const isRetryable = /unexpected EOF|XML syntax error|ECONNRESET|ETIMEDOUT|fetch failed/i.test(msg);
+      if (!isRetryable) {
+        await unloadModel(ollama, model);
+        log.error(`Ollama error: ${msg}`);
+        return { success: false, error: `Ollama request failed: ${msg}` };
+      }
+      // Retryable: wait and loop — the outer while will retry on the next iteration
+      log.warning(`» stream error (retryable): ${msg} — retrying in 5s`);
+      await new Promise((r) => setTimeout(r, 5_000));
+      continue;
     }
 
-    const promptTokens = response.prompt_eval_count;
-    const evalTokens = response.eval_count;
     if (promptTokens !== undefined) {
       const total = promptTokens + (evalTokens ?? 0);
       const pct = Math.round((total / 262144) * 100);
@@ -196,7 +207,6 @@ async function runOllamaLoop(ctx: AgentRunContext): Promise<AgentResult> {
       }
     }
 
-    const assistantMessage = response.message;
     messages.push(assistantMessage);
 
     const toolCalls: ToolCall[] | undefined = assistantMessage.tool_calls;
diff --git a/mcp/readFile.ts b/mcp/readFile.ts
index 7e45dd0..012b158 100644
--- a/mcp/readFile.ts
+++ b/mcp/readFile.ts
@@ -4,7 +4,7 @@ import type { ToolContext } from "./server.ts";
 import { execute, tool } from "./shared.ts";
 
 /** Hard cap on returned content to avoid flooding the model's context window. */
-const MAX_CHARS = 20000;
+const MAX_CHARS = 12000;
 
 export const ReadFileParams = type({
   path: type.string.describe("absolute path to the file to read"),