diff --git a/.env.example b/.env.example index 25009dd..1edf1b1 100644 --- a/.env.example +++ b/.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 " 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 3af435e..bc55ed5 100644 --- a/packages/core/src/agent/agent.ts +++ b/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) }, diff --git a/packages/core/src/agent/hot-memory.ts b/packages/core/src/agent/hot-memory.ts index 34596c8..bc9d1eb 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/test/agent.test.ts b/packages/core/test/agent.test.ts new file mode 100644 index 0000000..477fdb4 --- /dev/null +++ b/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(); + 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, 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((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) }) + ); + }); +}); diff --git a/packages/server/src/mcp/http.ts b/packages/server/src/mcp/http.ts index 1394e21..bc4bcf0 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 9ea0f75..d5718d8 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;