Abort stalled LLM requests

.env.example

@@ -32,6 +32,9 @@ GIT_AUTOCOMMIT=false
 
 PORT=3800
 
+# Maximum total time for one query/mutation model run before it is aborted.
+UNDERSTORY_LLM_TIMEOUT_MS=120000
+
 # Optional: require "Authorization: Bearer <token>" on /mcp and /api.
 # Unset = open (fine for localhost/LAN). Set this before exposing understory anywhere.
 #AUTH_TOKEN=

packages/core/src/agent/agent.ts

@@ -12,9 +12,14 @@ import { buildReadTools, buildWriteTools, formatTree } from "./tools.js";
 import { TraceRecorder, TraceStore } from "./trace.js";
 
 const MAX_STEPS = 12;
+const DEFAULT_TIMEOUT_MS = 120_000;
 
 export interface AgentOptions {
   model?: string;
+  /** Cancel the model run when the caller disconnects or explicitly aborts. */
+  abortSignal?: AbortSignal;
+  /** Total deadline for the model run. Defaults to UNDERSTORY_LLM_TIMEOUT_MS or 120 seconds. */
+  timeoutMs?: number;
 }
 
 export interface QueryResult {
@@ -95,6 +100,17 @@ function errorMessage(err: unknown): string {
   return err instanceof Error ? err.message : String(err);
 }
 
+export function modelAbortSignal(options: AgentOptions, env: NodeJS.ProcessEnv = process.env): AbortSignal {
+  const envTimeout = Number(env.UNDERSTORY_LLM_TIMEOUT_MS);
+  const timeoutMs =
+    options.timeoutMs ??
+    (Number.isFinite(envTimeout) && envTimeout > 0 ? envTimeout : DEFAULT_TIMEOUT_MS);
+  const timeoutSignal = AbortSignal.timeout(timeoutMs);
+  return options.abortSignal
+    ? AbortSignal.any([options.abortSignal, timeoutSignal])
+    : timeoutSignal;
+}
+
 /** Read-only Q&A over the bundle. */
 export async function runQuery(
   kb: KnowledgeBase,
@@ -109,6 +125,7 @@ export async function runQuery(
     modelChain = resolved.modelChain;
     const result = await generateText({
       model: resolved.model,
+      abortSignal: modelAbortSignal(options),
       system: buildSystemPrompt(ctx),
       prompt: question,
       tools: buildReadTools(kb, recorder),
@@ -139,12 +156,21 @@ export async function runMutation(
     modelChain = resolved.modelChain;
     const result = await generateText({
       model: resolved.model,
+      abortSignal: modelAbortSignal(options),
       system: buildSystemPrompt(ctx),
       prompt: instruction,
       tools: { ...buildReadTools(kb, recorder), ...buildWriteTools(kb, filesChanged, recorder) },
       stopWhen: stepCountIs(MAX_STEPS),
       temperature: 0.2,
     });
+    if (filesChanged.size === 0) {
+      const message =
+        `Mutation completed without changing any files. Model response: ` +
+        (result.text.trim() || "(empty response)");
+      const trace = recorder.finalize("mutation", instruction, message, "failed", modelChain);
+      await traceStore(kb).save(trace);
+      return { ok: false, status: "failed", error: message };
+    }
     const trace = recorder.finalize("mutation", instruction, result.text, "success", modelChain);
     await traceStore(kb).save(trace);
     return {
@@ -196,6 +222,7 @@ export async function streamChat(
     modelChain = resolved.modelChain;
     const result = streamText({
       model: resolved.model,
+      abortSignal: modelAbortSignal(options),
       system: buildSystemPrompt(ctx),
       messages,
       tools: { ...buildReadTools(kb, recorder), ...buildWriteTools(kb, filesChanged, recorder) },

packages/core/src/agent/hot-memory.ts

@@ -1,6 +1,6 @@
 import type { KnowledgeBase } from "../okf/index.js";
 import { parseDuration } from "../util/duration.js";
-import type { AgentOptions } from "./agent.js";
+import { modelAbortSignal, type AgentOptions } from "./agent.js";
 
 /**
  * Hot memory: a small working set of recently written concepts and recent
@@ -133,6 +133,12 @@ const defaultGenerate: HotGenerate = async (system, prompt, options) => {
     model = await providers.createModel(options.model ? { ...cfg, model: options.model } : cfg);
   }
   const { generateText } = await import("ai");
-  const result = await generateText({ model, system, prompt, temperature: 0 });
+  const result = await generateText({
+    model,
+    abortSignal: modelAbortSignal(options),
+    system,
+    prompt,
+    temperature: 0,
+  });
   return result.text;
 };

packages/core/test/agent.test.ts

@@ -0,0 +1,82 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { promises as fs } from "node:fs";
+import os from "node:os";
+import path from "node:path";
+
+const { generateText } = vi.hoisted(() => ({ generateText: vi.fn() }));
+
+vi.mock("ai", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("ai")>();
+  return {
+    ...actual,
+    generateText,
+    streamText: vi.fn(),
+    stepCountIs: vi.fn(() => vi.fn()),
+  };
+});
+
+vi.mock("../src/providers/index.js", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("../src/providers/index.js")>();
+  return {
+    ...actual,
+    createModel: vi.fn(async () => ({ modelId: "test-model" })),
+  };
+});
+
+import { modelAbortSignal, runMutation } from "../src/agent/agent.js";
+import { KnowledgeBase } from "../src/okf/index.js";
+
+let root: string;
+
+beforeEach(async () => {
+  root = await fs.mkdtemp(path.join(os.tmpdir(), "ustory-agent-"));
+  generateText.mockReset();
+  process.env.LLM_API_BASE_URL = "http://localhost:1234/v1";
+  process.env.LLM_API_KEY = "test";
+  process.env.LLM_API_FORMAT = "openai";
+  process.env.LLM_MODEL = "test-model";
+});
+
+afterEach(async () => {
+  await fs.rm(root, { recursive: true, force: true });
+  delete process.env.LLM_API_BASE_URL;
+  delete process.env.LLM_API_KEY;
+  delete process.env.LLM_API_FORMAT;
+  delete process.env.LLM_MODEL;
+});
+
+describe("modelAbortSignal", () => {
+  it("aborts when the configured model deadline expires", async () => {
+    const signal = modelAbortSignal({}, { UNDERSTORY_LLM_TIMEOUT_MS: "10" });
+
+    expect(signal.aborted).toBe(false);
+    await new Promise<void>((resolve) => signal.addEventListener("abort", () => resolve(), { once: true }));
+    expect(signal.aborted).toBe(true);
+  });
+
+  it("propagates caller cancellation before the deadline", () => {
+    const caller = new AbortController();
+    const signal = modelAbortSignal({ abortSignal: caller.signal, timeoutMs: 60_000 });
+
+    caller.abort();
+    expect(signal.aborted).toBe(true);
+  });
+});
+
+describe("runMutation", () => {
+  it("fails when the model returns without changing a file", async () => {
+    generateText.mockResolvedValue({ text: "Done", steps: [] });
+    const kb = new KnowledgeBase(root);
+
+    const outcome = await runMutation(kb, "Record a fact");
+
+    expect(outcome).toEqual({
+      ok: false,
+      status: "failed",
+      error: "Mutation completed without changing any files. Model response: Done",
+    });
+    expect(generateText).toHaveBeenCalledWith(
+      expect.objectContaining({ abortSignal: expect.any(AbortSignal) })
+    );
+  });
+});

packages/server/src/mcp/http.ts

@@ -14,12 +14,19 @@ export function mcpRouter(kb: KnowledgeBase): Router {
   const router = express.Router();
 
   const handle = async (req: Request, res: Response) => {
-    const server = await buildMcpServer(kb);
+    const requestAbort = new AbortController();
+    const abortRequest = () => requestAbort.abort();
+    req.once("aborted", abortRequest);
+
+    const server = await buildMcpServer(kb, requestAbort.signal);
     const transport = new StreamableHTTPServerTransport({
       sessionIdGenerator: undefined, // stateless
       enableJsonResponse: true, // one JSON reply per request — no long-lived SSE
     });
     res.on("close", () => {
+      // If the client disconnects while an LLM call is active, stop that work
+      // instead of leaving it occupying the local model indefinitely.
+      abortRequest();
       transport.close();
       server.close();
     });

packages/server/src/mcp/server.ts

@@ -14,7 +14,7 @@ import { buildSeedMemory, seedInstructions } from "./seed.js";
  * fallback; every tool-calling client loads descriptions). Without it the
  * client model has no signal that memory might hold an answer.
  */
-export async function buildMcpServer(kb: KnowledgeBase): Promise<McpServer> {
+export async function buildMcpServer(kb: KnowledgeBase, abortSignal?: AbortSignal): Promise<McpServer> {
   // Seed generation must never prevent the server from starting — a missing
   // or empty bundle root degrades to a minimal seed, not a crash.
   const seed = await buildSeedMemory(kb).catch((err: Error) => {
@@ -40,7 +40,7 @@ export async function buildMcpServer(kb: KnowledgeBase): Promise<McpServer> {
       inputSchema: { question: z.string().describe("The question to answer") },
     },
     async ({ question }) => {
-      const { answer, source } = await runQueryCached(kb, question);
+      const { answer, source } = await runQueryCached(kb, question, { abortSignal });
       const marker = source === "cache" ? "\n\n(cached answer)" : source === "hot" ? "\n\n(hot memory)" : "";
       return {
         content: [{ type: "text", text: `${answer}${marker}` }],
@@ -120,7 +120,7 @@ export async function buildMcpServer(kb: KnowledgeBase): Promise<McpServer> {
         `must use the write tools.\n\n` +
         `KNOWLEDGE TO RECORD:\n${content}` +
         (suggested_path ? `\n\nIf it fits, place new content at ${suggested_path}.` : "");
-      const outcome = await runMutation(kb, instruction);
+      const outcome = await runMutation(kb, instruction, { abortSignal });
       await refreshSeed();
       return mutationOutcomeResponse(outcome);
     }
@@ -137,7 +137,7 @@ export async function buildMcpServer(kb: KnowledgeBase): Promise<McpServer> {
       },
     },
     async ({ instruction }) => {
-      const outcome = await runMutation(kb, instruction);
+      const outcome = await runMutation(kb, instruction, { abortSignal });
       await refreshSeed();
       return mutationOutcomeResponse(outcome);
     }
@@ -219,7 +219,7 @@ export async function buildMcpServer(kb: KnowledgeBase): Promise<McpServer> {
         `or remove the link if the target is gone.\n${brokenList}\n\n` +
         `Follow the enrich / link-both-ways rules. Read concepts before editing.`;
 
-      const outcome = await runMutation(kb, instruction);
+      const outcome = await runMutation(kb, instruction, { abortSignal });
       await refreshSeed();
       if (!outcome.ok) return mutationOutcomeResponse(outcome);
       const { summary, filesChanged } = outcome.result;