diff --git a/.env.example b/.env.example index ed809972..b73dc6b4 100644 --- a/.env.example +++ b/.env.example @@ -35,6 +35,11 @@ GIT_AUTOCOMMIT=false PORT=3800 +# Optional total deadline for one query/mutation model run (e.g. 300s, 10m, 2h). +# Unset = no deadline: local models routinely take many minutes per run. +# A disconnecting client always cancels its in-flight model work regardless. +#LLM_TIMEOUT=30m + # Optional: require "Authorization: Bearer " on /mcp and /api. # Unset = open (fine for localhost/LAN). Set this before exposing understory anywhere. #AUTH_TOKEN= diff --git a/packages/core/src/agent/agent.ts b/packages/core/src/agent/agent.ts index c993c5e2..15d2767c 100644 --- a/packages/core/src/agent/agent.ts +++ b/packages/core/src/agent/agent.ts @@ -10,11 +10,16 @@ import { withFallback } from "../providers/fallback.js"; import { buildSystemPrompt } from "./system-prompt.js"; import { buildReadTools, buildWriteTools, formatTree } from "./tools.js"; import { TraceRecorder, TraceStore, type TraceUsage } from "./trace.js"; +import { parseDuration } from "../util/duration.js"; const MAX_STEPS = 12; 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 LLM_TIMEOUT; no deadline when neither is set. */ + timeoutMs?: number; } export interface QueryResult { @@ -95,6 +100,46 @@ function errorMessage(err: unknown): string { return err instanceof Error ? err.message : String(err); } +/** + * The configured total deadline for one model run, from LLM_TIMEOUT + * ("10m", "2h", ...). There is no deadline unless one is configured: healthy + * runs on local hardware routinely take many minutes, so any default would + * abort working setups. An unparseable value means no deadline plus a + * warning, never a silent fallback. + */ +export function resolveLlmTimeout(env: NodeJS.ProcessEnv = process.env): { + ms: number | null; + warning?: string; +} { + const raw = env.LLM_TIMEOUT; + if (raw === undefined || raw.trim() === "") return { ms: null }; + const ms = parseDuration(raw); + if (ms === null || ms <= 0) { + return { + ms: null, + warning: `LLM_TIMEOUT="${raw}" is not a positive duration (use e.g. 300s, 10m, 2h) — no deadline applied`, + }; + } + return { ms }; +} + +/** + * Abort signal for a model run: fires when the caller cancels (always + * honored) or when the configured deadline expires (opt-in). Undefined when + * neither applies. + */ +export function modelAbortSignal( + options: AgentOptions, + env: NodeJS.ProcessEnv = process.env +): AbortSignal | undefined { + const timeoutMs = options.timeoutMs ?? resolveLlmTimeout(env).ms; + const signals: AbortSignal[] = []; + if (options.abortSignal) signals.push(options.abortSignal); + if (timeoutMs != null) signals.push(AbortSignal.timeout(timeoutMs)); + if (signals.length === 0) return undefined; + return signals.length === 1 ? signals[0] : AbortSignal.any(signals); +} + /** Sum token usage across the run's steps (issue #15). Undefined when the provider reports none. */ function sumStepsUsage( steps: ReadonlyArray<{ usage?: { inputTokens?: number; outputTokens?: number } }> @@ -126,6 +171,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), @@ -156,12 +202,22 @@ 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, }); + // A run that wrote nothing AND said nothing is a stalled/flaked model. + // Writing nothing while explaining why (knowledge already present, + // nothing genuinely related to link) is a legitimate outcome. + if (filesChanged.size === 0 && result.text.trim() === "") { + const message = "Mutation returned an empty response without changing any files."; + const trace = recorder.finalize("mutation", instruction, message, "failed", modelChain, sumStepsUsage(result.steps)); + await traceStore(kb).save(trace); + return { ok: false, status: "failed", error: message }; + } const trace = recorder.finalize("mutation", instruction, result.text, "success", modelChain, sumStepsUsage(result.steps)); await traceStore(kb).save(trace); return { @@ -213,6 +269,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) }, diff --git a/packages/core/src/agent/hot-memory.ts b/packages/core/src/agent/hot-memory.ts index 34596c80..bc9d1ebc 100644 --- a/packages/core/src/agent/hot-memory.ts +++ b/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; }; diff --git a/packages/core/src/agent/index.ts b/packages/core/src/agent/index.ts index bd0f6cd6..6fb90ab7 100644 --- a/packages/core/src/agent/index.ts +++ b/packages/core/src/agent/index.ts @@ -1,4 +1,4 @@ -export { runQuery, runMutation, streamChat } from "./agent.js"; +export { runQuery, runMutation, streamChat, resolveLlmTimeout } from "./agent.js"; export type { AgentOptions, QueryResult, MutationResult, MutationOutcome } from "./agent.js"; export { buildSystemPrompt } from "./system-prompt.js"; export { buildReadTools, buildWriteTools, formatTree } from "./tools.js"; diff --git a/packages/core/test/agent.test.ts b/packages/core/test/agent.test.ts new file mode 100644 index 00000000..8edc8cb6 --- /dev/null +++ b/packages/core/test/agent.test.ts @@ -0,0 +1,127 @@ +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(); + return { + ...actual, + generateText, + streamText: vi.fn(), + stepCountIs: vi.fn(() => vi.fn()), + }; +}); + +vi.mock("../src/providers/index.js", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + createModel: vi.fn(async () => ({ modelId: "test-model" })), + }; +}); + +import { modelAbortSignal, resolveLlmTimeout, 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("resolveLlmTimeout", () => { + it("applies no deadline unless LLM_TIMEOUT is set", () => { + expect(resolveLlmTimeout({})).toEqual({ ms: null }); + expect(resolveLlmTimeout({ LLM_TIMEOUT: "" })).toEqual({ ms: null }); + }); + + it("parses duration strings", () => { + expect(resolveLlmTimeout({ LLM_TIMEOUT: "10m" })).toEqual({ ms: 600_000 }); + expect(resolveLlmTimeout({ LLM_TIMEOUT: "2h" })).toEqual({ ms: 7_200_000 }); + }); + + it("warns and applies no deadline for unparseable or non-positive values", () => { + for (const raw of ["soon", "120000", "0s", "NaN"]) { + const resolved = resolveLlmTimeout({ LLM_TIMEOUT: raw }); + expect(resolved.ms).toBeNull(); + expect(resolved.warning).toContain(raw); + } + }); +}); + +describe("modelAbortSignal", () => { + it("returns no signal when there is neither a caller signal nor a deadline", () => { + expect(modelAbortSignal({}, {})).toBeUndefined(); + }); + + it("aborts when the configured model deadline expires", async () => { + const signal = modelAbortSignal({}, { LLM_TIMEOUT: "10ms" })!; + + expect(signal.aborted).toBe(false); + await new Promise((resolve) => signal.addEventListener("abort", () => resolve(), { once: true })); + expect(signal.aborted).toBe(true); + }); + + it("propagates caller cancellation with no deadline configured", () => { + const caller = new AbortController(); + const signal = modelAbortSignal({ abortSignal: caller.signal }, {})!; + + caller.abort(); + 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 an empty response without changing a file", async () => { + generateText.mockResolvedValue({ text: " ", steps: [] }); + const kb = new KnowledgeBase(root); + + const outcome = await runMutation(kb, "Record a fact", { abortSignal: new AbortController().signal }); + + expect(outcome).toEqual({ + ok: false, + status: "failed", + error: "Mutation returned an empty response without changing any files.", + }); + expect(generateText).toHaveBeenCalledWith( + expect.objectContaining({ abortSignal: expect.any(AbortSignal) }) + ); + }); + + it("succeeds when the model explains why nothing needed writing", async () => { + generateText.mockResolvedValue({ text: "Already recorded in /facts/a.md; nothing to change.", steps: [] }); + const kb = new KnowledgeBase(root); + + const outcome = await runMutation(kb, "Record a fact"); + + expect(outcome.ok).toBe(true); + if (outcome.ok) { + expect(outcome.result.filesChanged).toEqual([]); + expect(outcome.result.summary).toContain("Already recorded"); + } + }); +}); diff --git a/packages/server/src/index.ts b/packages/server/src/index.ts index 0faadd36..445ede72 100644 --- a/packages/server/src/index.ts +++ b/packages/server/src/index.ts @@ -3,7 +3,7 @@ import { existsSync } from "node:fs"; import { fileURLToPath } from "node:url"; import express from "express"; import cors from "cors"; -import { KnowledgeBase, resolveFallbackConfig, resolveModelConfig } from "@understory/core"; +import { KnowledgeBase, resolveFallbackConfig, resolveLlmTimeout, resolveModelConfig } from "@understory/core"; import { mcpRouter } from "./mcp/http.js"; import { browseRouter } from "./api/browse.js"; import { chatRouter } from "./api/chat.js"; @@ -52,6 +52,10 @@ try { process.exit(1); } +const llmTimeout = resolveLlmTimeout(); +if (llmTimeout.warning) console.error(`[understory] ${llmTimeout.warning}`); +console.log(`[understory] llm timeout: ${process.env.LLM_TIMEOUT && llmTimeout.ms != null ? process.env.LLM_TIMEOUT.trim() : "none"}`); + // Reflect the request origin; expose Mcp-Session-Id so browser MCP clients can // read it back off the initialize response. app.use( diff --git a/packages/server/src/mcp/http.ts b/packages/server/src/mcp/http.ts index 1394e217..bc4bcf0d 100644 --- a/packages/server/src/mcp/http.ts +++ b/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(); }); diff --git a/packages/server/src/mcp/server.ts b/packages/server/src/mcp/server.ts index 9ea0f753..d5718d88 100644 --- a/packages/server/src/mcp/server.ts +++ b/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 { +export async function buildMcpServer(kb: KnowledgeBase, abortSignal?: AbortSignal): Promise { // 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 { 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 { `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 { }, }, 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 { `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; diff --git a/packages/server/src/mcp/stdio.ts b/packages/server/src/mcp/stdio.ts index 8888eac9..25eab91b 100644 --- a/packages/server/src/mcp/stdio.ts +++ b/packages/server/src/mcp/stdio.ts @@ -5,7 +5,7 @@ * -e LLM_PROVIDER=openrouter -- node /packages/server/dist/mcp/stdio.js */ import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; -import { KnowledgeBase, resolveFallbackConfig, resolveModelConfig } from "@understory/core"; +import { KnowledgeBase, resolveFallbackConfig, resolveLlmTimeout, resolveModelConfig } from "@understory/core"; import { buildMcpServer } from "./server.js"; const bundleRoot = process.env.BUNDLE_ROOT; @@ -34,6 +34,10 @@ try { process.exit(1); } +const llmTimeout = resolveLlmTimeout(); +if (llmTimeout.warning) console.error(`[understory] ${llmTimeout.warning}`); +console.error(`[understory] llm timeout: ${process.env.LLM_TIMEOUT && llmTimeout.ms != null ? process.env.LLM_TIMEOUT.trim() : "none"}`); + const kb = new KnowledgeBase(bundleRoot, { gitAutocommit: process.env.GIT_AUTOCOMMIT === "true", });