From 36c5d19437d669a68977e72acc1626588d1722b6 Mon Sep 17 00:00:00 2001 From: amosblomqvist Date: Sat, 19 Sep 2026 17:26:57 +0300 Subject: [PATCH] =?UTF-8?q?fix(tools):=20the=20abort=20cascade=20=E2=80=94?= =?UTF-8?q?=20kill=20running=20tool=20calls,=20honest=20activity=20labels?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause of the live-test bugs (2026-09-19): pi passes every run's AbortSignal to tool execute, but toPiTool dropped it — agent.abort() only landed at the next step boundary, so a long local grep held its turn hostage until the dev server's memory restart killed the run mid-tool, freezing the record at its write-ahead state (the 'went back in time' rollback). - lh-tool: execute ctx gains signal (parent + timeout combined); every call races a 30s system-level timeout with a 250ms cooperative grace — abort kills a stuck tool but lets a poll-the-signal tool deliver its partial results. Pre-aborted signals never fire listeners: checked upfront. toolTimeoutMs is the test seam (pdfModelOverride idiom). - fs-shared: walkFiles returns {timedOut, aborted} instead of throwing — abort yields partial results. - grep: the scan phase was unbounded (the walk's 10s deadline only bounded collection) — same budget now checked per file; find/grep check the signal and return partials with a notice. - pending/chat: streamingToolActivity names the ACTUAL running tool — a part sits input-streaming through its whole execution, so the old shimmer called minutes-long greps 'Writing a question…'. Tests: toPiTool guard (forward, abort, timeout, grace, question bypass, ToolInputError passthrough), walk mid-abort, find/grep pre-abort + mid-run abort, activity labels. Gates: 617/617. --- components/chat.tsx | 18 ++--- src/agent/lh-tool.ts | 89 ++++++++++++++++++++++- src/agent/tools/find.ts | 11 ++- src/agent/tools/fs-shared.ts | 15 +++- src/agent/tools/grep.ts | 30 +++++++- src/core/pending.ts | 38 ++++++++++ tests/fs-tools.test.ts | 59 ++++++++++++++- tests/stream-activity.test.ts | 59 +++++++++++++++ tests/tool-abort.test.ts | 132 ++++++++++++++++++++++++++++++++++ 9 files changed, 430 insertions(+), 21 deletions(-) create mode 100644 tests/stream-activity.test.ts create mode 100644 tests/tool-abort.test.ts diff --git a/components/chat.tsx b/components/chat.tsx index 8942db2..dbb403f 100644 --- a/components/chat.tsx +++ b/components/chat.tsx @@ -2,7 +2,7 @@ import { createContext, Fragment, useCallback, useEffect, useLayoutEffect, useMemo, useRef, useState } from "react"; import { useChat } from "@ai-sdk/react"; -import { isToolUIPart, lastAssistantMessageIsCompleteWithToolCalls } from "ai"; +import { lastAssistantMessageIsCompleteWithToolCalls } from "ai"; import { createChatTransport } from "@/components/chat-transport"; import { useStickToBottomContext } from "use-stick-to-bottom"; import { @@ -34,7 +34,7 @@ import { QuestionDock } from "@/components/question-dock"; import { SidebarToggle } from "@/components/sidebar"; import type { ArtifactPartProps, JobRecord, Phase3, QuestionAnswerData, TeacherUIMessage, TurnUsage } from "@/src/core"; import { PhaseChip, PhaseDivider, planNodeCount } from "@/components/phase-divider"; -import { findPendingToolCall, phaseAt3, phaseTimeline } from "@/src/core"; +import { findPendingToolCall, phaseAt3, phaseTimeline, streamingToolActivity } from "@/src/core"; /** One width for everything: transcript, cards, and the bottom blob share this column. */ /** @@ -262,12 +262,12 @@ export function Chat({ // session's current phase. const phases = useMemo(() => phaseTimeline(messages), [messages]); - // A tool call still streaming its input: the teacher is writing a question. - // Shown as a transcript shimmer, exactly like the thinking indicator. - const questionStreaming = - !pending && - messages.at(-1)?.role === "assistant" && - (messages.at(-1)?.parts.some((p) => isToolUIPart(p) && p.state === "input-streaming") ?? false); + // A tool call still streaming its input: the teacher is mid-activity. The + // label names the ACTUAL tool (a long grep used to shimmer as “Writing a + // question…” for minutes — a part stays input-streaming through its whole + // execution, see streamingToolActivity); shown as the transcript shimmer, + // exactly like the thinking indicator. + const streamingActivity = !pending ? streamingToolActivity(messages) : null; const [text, setText] = useState(""); // Voice dictation (STT wave, 2026-09-17): finals append at the end of @@ -941,7 +941,7 @@ export function Chat({ ))} {(status === "submitted" || turnInFlight) && Thinking…} - {questionStreaming && Writing a question…} + {streamingActivity && {streamingActivity}} {error && (
Something went wrong: {error.message} diff --git a/src/agent/lh-tool.ts b/src/agent/lh-tool.ts index 88671c5..bf46a5c 100644 --- a/src/agent/lh-tool.ts +++ b/src/agent/lh-tool.ts @@ -31,8 +31,14 @@ export interface LhTool { /** * Server-side execution. Absent + clientQuestion = the client answers. * Throws on validation failure — the loop converts to an error result. + * + * ctx.signal (2026-09-19, the abort-cascade fix): the run's abort signal + * COMBINED with the system-level per-call timeout — aborting the run + * kills a running tool, the timeout is the net so no tool can hold a + * turn hostage (the 2.7-minute grep lesson). Tools with loops check it + * and return partial results; fast single-op tools may ignore it. */ - execute?: (input: TInput, ctx: { toolCallId: string }) => TResult | Promise; + execute?: (input: TInput, ctx: { toolCallId: string; signal?: AbortSignal }) => TResult | Promise; /** Client-answered question tool (quiz/ask): a call ends the turn. */ clientQuestion?: boolean; } @@ -81,6 +87,83 @@ export class ToolInputError extends Error { } } +/** + * System-level per-call timeout (2026-09-19, the abort-cascade fix): pi + * passes every run's AbortSignal to tool execute, but the bridge used to + * drop it — so agent.abort() only landed at the NEXT step boundary and a + * long-running local tool (the grep that outlived its turn by minutes) + * held the turn hostage until the dev server's memory restart killed the + * run mid-tool, freezing the record at its write-ahead state (the "went + * back in time" rollback). The guard now races every execute against (a) + * the run's signal and (b) this timeout, rejecting with pi's own wording + * ("Operation aborted") on abort and a clear timeout message otherwise — + * pi turns the throw into an error tool result the model can see. + */ +const TOOL_TIMEOUT_MS = 30_000; +/** Test seam (runner.ts's pdfModelOverride idiom). */ +export const toolTimeoutMs = { current: TOOL_TIMEOUT_MS }; + +/** How long an aborted/timed-out call may take to settle on its own before + * the hard reject lands — cooperative tools poll the signal and return + * partial results within a tick or two (the fs tools check per file). */ +const TOOL_ABORT_GRACE_MS = 250; + +function runWithAbortGuard( + toolName: string, + run: (signal: AbortSignal) => T | Promise, + parent: AbortSignal | undefined, +): Promise { + const controller = new AbortController(); + const forwardAbort = () => controller.abort(); + if (parent) { + if (parent.aborted) controller.abort(); + else parent.addEventListener("abort", forwardAbort, { once: true }); + } + const timeoutMs = toolTimeoutMs.current; + const timer = setTimeout(() => controller.abort(), timeoutMs); + const humanTimeout = timeoutMs >= 1_000 ? `${Math.round(timeoutMs / 1000)}s` : `${timeoutMs}ms`; + let graceTimer: ReturnType | undefined; + const cleanup = () => { + clearTimeout(timer); + if (graceTimer) clearTimeout(graceTimer); + parent?.removeEventListener("abort", forwardAbort); + }; + return new Promise((resolve, reject) => { + const settle = (fn: () => void) => { + controller.signal.removeEventListener("abort", onAbort); + if (graceTimer) clearTimeout(graceTimer); + fn(); + }; + const onAbort = () => { + // Cooperative tools poll the signal and settle with partial + // results within a tick or two (the fs tools check per file) — + // give them a brief grace to deliver those before the hard + // reject, so abort KILLS a stuck tool without racing away a + // graceful one's partial result. + graceTimer = setTimeout( + () => + reject( + new Error( + parent?.aborted + ? "Operation aborted" + : `Tool "${toolName}" timed out after ${humanTimeout} — the result was discarded`, + ), + ), + TOOL_ABORT_GRACE_MS, + ); + }; + controller.signal.addEventListener("abort", onAbort, { once: true }); + if (controller.signal.aborted) onAbort(); // an already-aborted signal never fires listeners + Promise.resolve() + .then(() => run(controller.signal)) + .then( + (v) => settle(() => resolve(v)), + (e) => settle(() => reject(e)), + ); + }) + .finally(cleanup); +} + /** Convert one lh tool to a pi AgentTool. */ export function toPiTool(def: LhTool): AgentTool { return { @@ -97,7 +180,7 @@ export function toPiTool(def: LhTool): AgentTool { // dock's subject can never flip mid-turn between two parallel // questions. Question executes are instant; no latency is lost. executionMode: def.clientQuestion ? "sequential" : "parallel", - execute: async (toolCallId, args) => { + execute: async (toolCallId, args, signal) => { // The authoring check: structural JSON Schema validation already // ran in the loop's prepare step; this re-run adds the refinements // (superRefine rules) the JSON Schema can't carry. @@ -111,7 +194,7 @@ export function toPiTool(def: LhTool): AgentTool { details: null, }; } - const result = await def.execute(parsed.data, { toolCallId }); + const result = await runWithAbortGuard(def.name, (toolSignal) => def.execute!(parsed.data, { toolCallId, signal: toolSignal }), signal); // Rich image result: model-facing content parts carry the images // (pi's ToolResultMessage.content is (TextContent | ImageContent)[]). if (isImageResult(result)) { diff --git a/src/agent/tools/find.ts b/src/agent/tools/find.ts index 8df5df1..9d4c018 100644 --- a/src/agent/tools/find.ts +++ b/src/agent/tools/find.ts @@ -36,7 +36,8 @@ export const findTool = lhTool({ name: "find", description: `Search for files by glob pattern. Returns matching file paths relative to the search directory. Ignores node_modules, .git, and Library directories. Output is truncated to ${DEFAULT_LIMIT} results or ${formatSize(MAX_OUTPUT_BYTES)} (whichever is hit first).`, schema: findInputSchema, - async execute(input) { + async execute(input, ctx) { + if (ctx?.signal?.aborted) return "Operation aborted"; const searchPath = resolveSourcePath(input.path); let stat; try { @@ -49,7 +50,7 @@ export const findTool = lhTool({ const effectiveLimit = input.limit ?? DEFAULT_LIMIT; const matches: string[] = []; let limitReached = false; - const { timedOut } = await walkFiles( + const { timedOut, aborted } = await walkFiles( searchPath, (filePath) => { const rel = path.relative(searchPath, filePath).split(path.sep).join("/"); @@ -61,8 +62,14 @@ export const findTool = lhTool({ matches.push(rel); return false; }, + ctx?.signal, ); + if (aborted) { + return matches.length > 0 + ? `${matches.join("\n")}\n\n[search aborted — partial results]` + : "Search aborted before any results."; + } if (matches.length === 0 && !limitReached) { return timedOut ? `Search timed out after 10s with no matches yet. Narrow the search to a subdirectory.` diff --git a/src/agent/tools/fs-shared.ts b/src/agent/tools/fs-shared.ts index 3f3015b..f0396d2 100644 --- a/src/agent/tools/fs-shared.ts +++ b/src/agent/tools/fs-shared.ts @@ -136,16 +136,21 @@ export async function walkFiles( root: string, visit: (filePath: string) => boolean | Promise, signal?: AbortSignal, -): Promise<{ timedOut: boolean }> { +): Promise<{ timedOut: boolean; aborted: boolean }> { const deadline = Date.now() + WALK_TIMEOUT_MS; let timedOut = false; + let aborted = false; + if (signal?.aborted) return { timedOut: false, aborted: true }; const walkDir = async (dir: string): Promise => { if (Date.now() > deadline) { timedOut = true; return true; // stop } - if (signal?.aborted) throw new Error("Operation aborted"); + if (signal?.aborted) { + aborted = true; + return true; // stop — abort yields partial results, not a throw + } let entries: import("node:fs").Dirent[]; try { entries = await fs.readdir(dir, { withFileTypes: true }); @@ -157,6 +162,10 @@ export async function walkFiles( timedOut = true; return true; } + if (signal?.aborted) { + aborted = true; + return true; + } const full = path.join(dir, entry.name); if (entry.isDirectory()) { if (EXCLUDED_DIR_NAMES.has(entry.name)) continue; @@ -170,7 +179,7 @@ export async function walkFiles( }; await walkDir(root); - return { timedOut }; + return { timedOut, aborted }; } /** True if the first 8KB chunk contains a NUL byte — pi/rg's binary skip. */ diff --git a/src/agent/tools/grep.ts b/src/agent/tools/grep.ts index c493a26..ec8312c 100644 --- a/src/agent/tools/grep.ts +++ b/src/agent/tools/grep.ts @@ -12,6 +12,7 @@ import { resolveSourcePath, truncateGrepLine, walkFiles, + WALK_TIMEOUT_MS, } from "./fs-shared.ts"; /** @@ -45,7 +46,8 @@ export const grepTool = lhTool({ name: "grep", description: `Search file contents for a pattern. Returns matching lines with file paths and line numbers. Ignores node_modules, .git, and Library directories. Output is truncated to ${DEFAULT_LIMIT} matches or ${formatSize(MAX_OUTPUT_BYTES)} (whichever is hit first). Long lines are truncated to ${GREP_MAX_LINE_LENGTH} chars.`, schema: grepInputSchema, - async execute(input) { + async execute(input, ctx) { + if (ctx?.signal?.aborted) return "Operation aborted"; const searchPath = resolveSourcePath(input.path); let stat; @@ -66,6 +68,7 @@ export const grepTool = lhTool({ // Files to scan: one explicit file, or the walked tree (glob-filtered). const files: string[] = []; let timedOut = false; + let aborted = false; if (!isDirectory) { files.push(searchPath); } else { @@ -79,16 +82,32 @@ export const grepTool = lhTool({ files.push(filePath); return false; }, + ctx?.signal, ); timedOut = result.timedOut; + aborted = result.aborted; } const outLines: string[] = []; let matchCount = 0; let limitReached = false; let linesTruncated = false; + // The scan phase is bounded too (2026-09-19): the walk's 10s deadline + // bounds COLLECTION, but reading every collected file was unbounded — + // the grep that outlived its turn by minutes. Same budget, checked per + // file; abort lands mid-scan with partial results. + const scanDeadline = Date.now() + WALK_TIMEOUT_MS; + let scanTimedOut = false; for (const filePath of files) { + if (ctx?.signal?.aborted) { + aborted = true; + break; + } + if (Date.now() > scanDeadline) { + scanTimedOut = true; + break; + } if (matchCount >= effectiveLimit) { limitReached = true; break; @@ -133,8 +152,13 @@ export const grepTool = lhTool({ if (limitReached) break; } + if (aborted) { + return outLines.length > 0 + ? `${capOutput(outLines).content}\n\n[search aborted — partial results]` + : "Search aborted before any results."; + } if (outLines.length === 0) { - return timedOut + return timedOut || scanTimedOut ? `Search timed out after 10s with no matches yet. Narrow the search to a subdirectory.` : "No matches found"; } @@ -143,7 +167,7 @@ export const grepTool = lhTool({ if (limitReached || matchCount >= effectiveLimit) notices.push(`${effectiveLimit} matches limit reached`); if (linesTruncated) notices.push(`Some lines truncated to ${GREP_MAX_LINE_LENGTH} chars. Use read tool to see full lines`); if (truncated) notices.push(`${formatSize(MAX_OUTPUT_BYTES)} limit reached`); - if (timedOut) notices.push("search timed out after 10s — partial results, narrow the directory"); + if (timedOut || scanTimedOut) notices.push("search timed out after 10s — partial results, narrow the directory"); return notices.length > 0 ? `${content}\n\n[${notices.join(". ")}]` : content; }, }); diff --git a/src/core/pending.ts b/src/core/pending.ts index 32e1fd2..618e821 100644 --- a/src/core/pending.ts +++ b/src/core/pending.ts @@ -83,4 +83,42 @@ export function findPendingToolCall(messages: UIMessage[]): PendingToolCall | nu return findPending(messages, QUESTION_TOOLS); } +// ─── Turn activity label (2026-09-19, the “Writing a question…” fix) ────── + +/** + * The transcript's turn-activity shimmer text: what the teacher is + * doing RIGHT NOW. A tool part sits in `input-streaming` for its whole + * execution (the bridge defers `tool-input-available` to + * tool_execution_end — the invalid-input gate for the question dock), so + * keying the old “Writing a question…” shimmer on ANY streaming tool + * part mislabeled minutes-long local greps as question-writing. The label + * now names the actual activity; chat.tsx gates it on `!pending` and + * renders it as the quiet transcript shimmer. + */ +const ACTIVITY_LABELS: Record = { + read: "Reading a file…", + ls: "Listing a directory…", + find: "Finding files…", + grep: "Searching file contents…", + research: "Dispatching research…", + demo: "Building a demo…", + visual: "Making a visual…", + pdf: "Reading a PDF…", + plan_init: "Drawing up the plan…", + probe_done: "Concluding the assessment…", +}; + +export function streamingToolActivity(messages: UIMessage[]): string | null { + const last = messages.at(-1); + if (!last || last.role !== "assistant") return null; + for (const part of last.parts) { + if (isToolUIPart(part) && part.state === "input-streaming") { + const tool = part.type.slice("tool-".length); + if ((QUESTION_TOOLS as readonly string[]).includes(tool)) return "Writing a question…"; + return ACTIVITY_LABELS[tool] ?? "Working…"; + } + } + return null; +} + diff --git a/tests/fs-tools.test.ts b/tests/fs-tools.test.ts index eb2d56d..85bfc56 100644 --- a/tests/fs-tools.test.ts +++ b/tests/fs-tools.test.ts @@ -4,7 +4,7 @@ import * as os from "node:os"; import * as path from "node:path"; import { describe, it } from "node:test"; -import { expandBraces, globToRegex, matchGlob } from "../src/agent/tools/fs-shared.ts"; +import { expandBraces, globToRegex, matchGlob, walkFiles } from "../src/agent/tools/fs-shared.ts"; import { findTool } from "../src/agent/tools/find.ts"; import { grepTool } from "../src/agent/tools/grep.ts"; import { lsTool } from "../src/agent/tools/ls.ts"; @@ -186,3 +186,60 @@ describe("read (home-anchored re-port)", () => { assert.match(out, /Eigenvectors keep their direction/); }); }); + +describe("abort plumbing (the 2026-09-19 cascade fix)", () => { + it("walkFiles stops promptly when the signal aborts mid-walk", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "lh-fs-abort-")); + for (let i = 0; i < 20; i++) await fs.writeFile(path.join(root, `f${i}.txt`), "needle\n"); + const ac = new AbortController(); + let visited = 0; + const result = await walkFiles( + root, + () => { + visited++; + if (visited === 3) ac.abort(); // deterministic: abort from inside the walk + return false; + }, + ac.signal, + ); + assert.equal(result.aborted, true); + assert.equal(result.timedOut, false); + assert.ok(visited < 20, "the walk stopped early"); + }); + + it("find honors a pre-aborted signal without walking", async () => { + FIXTURE = await setup(); + const ac = new AbortController(); + ac.abort(); + const out = (await findTool.execute!({ pattern: "*.md" }, { toolCallId: "c", signal: ac.signal })) as string; + assert.equal(out, "Operation aborted"); + }); + + it("grep honors a pre-aborted signal without scanning", async () => { + const ac = new AbortController(); + ac.abort(); + const out = (await grepTool.execute!({ pattern: "eigenvector" }, { toolCallId: "c", signal: ac.signal })) as string; + assert.equal(out, "Operation aborted"); + }); + + it("grep aborts mid-run with partial results, not a hang", async () => { + // The exact shape of the 2.7-minute lesson: a tree big enough that + // the walk+scan outlast an abort mid-flight. Whichever phase the + // abort lands in, the tool returns promptly with an abort notice — + // never "No matches found", never a hang. + const root = await fs.mkdtemp(path.join(os.tmpdir(), "lh-fs-midabort-")); + await fs.mkdir(path.join(root, "deep"), { recursive: true }); + for (let i = 0; i < 200; i++) { + await fs.writeFile(path.join(root, i % 2 ? `f${i}.txt` : `deep/f${i}.txt`), "needle line\n"); + } + const ac = new AbortController(); + setTimeout(() => ac.abort(), 1); + const started = Date.now(); + const out = (await grepTool.execute!( + { pattern: "needle", path: root }, + { toolCallId: "c", signal: ac.signal }, + )) as string; + assert.match(out, /aborted/); + assert.ok(Date.now() - started < 5_000, "returned promptly"); + }); +}); diff --git a/tests/stream-activity.test.ts b/tests/stream-activity.test.ts new file mode 100644 index 0000000..5938d47 --- /dev/null +++ b/tests/stream-activity.test.ts @@ -0,0 +1,59 @@ +/** + * streamingToolActivity (2026-09-19, the "Writing a question…" fix): a + * tool part stays input-streaming for its WHOLE execution (the bridge + * defers input-available to tool_execution_end), so the transcript + * activity line must name the ACTUAL running tool — the old code shimmered + * "Writing a question…" over minutes-long greps. + */ +import assert from "node:assert/strict"; +import { describe, it } from "node:test"; +import type { UIMessage } from "ai"; +import { streamingToolActivity } from "../src/core/pending.ts"; + +function msg(role: string, parts: unknown[]): UIMessage { + return { id: "m1", role, parts } as unknown as UIMessage; +} + +function toolPart(tool: string, state: string) { + return { type: `tool-${tool}`, toolCallId: "c1", state, input: {} }; +} + +describe("streamingToolActivity (the turn-activity label)", () => { + it("question tools keep their question shimmer", () => { + assert.equal(streamingToolActivity([msg("assistant", [toolPart("quiz", "input-streaming")])]), "Writing a question…"); + assert.equal( + streamingToolActivity([msg("assistant", [toolPart("ask_user_question", "input-streaming")])]), + "Writing a question…", + ); + }); + + it("local tools name their actual activity", () => { + assert.equal(streamingToolActivity([msg("assistant", [toolPart("grep", "input-streaming")])]), "Searching file contents…"); + assert.equal(streamingToolActivity([msg("assistant", [toolPart("read", "input-streaming")])]), "Reading a file…"); + assert.equal(streamingToolActivity([msg("assistant", [toolPart("find", "input-streaming")])]), "Finding files…"); + assert.equal(streamingToolActivity([msg("assistant", [toolPart("ls", "input-streaming")])]), "Listing a directory…"); + }); + + it("job dispatch tools get their own labels; unknown tools stay generic", () => { + assert.equal(streamingToolActivity([msg("assistant", [toolPart("pdf", "input-streaming")])]), "Reading a PDF…"); + assert.equal( + streamingToolActivity([msg("assistant", [toolPart("mystery_tool", "input-streaming")])]), + "Working…", + ); + }); + + it("no activity once the call settles or the turn ends", () => { + // output-available = done executing → no shimmer (the transcript + // shows the settled card instead). + assert.equal(streamingToolActivity([msg("assistant", [toolPart("grep", "output-available")])]), null); + // the last message being the user's = nothing running. + assert.equal(streamingToolActivity([msg("user", [{ type: "text", text: "hi" }])]), null); + assert.equal(streamingToolActivity([]), null); + }); + + it("the SECOND streaming part wins only if the first settles — parts scan forward", () => { + // A settled call followed by a live one: the live one is the activity. + const parts = [toolPart("read", "output-available"), toolPart("grep", "input-streaming")]; + assert.equal(streamingToolActivity([msg("assistant", parts)]), "Searching file contents…"); + }); +}); diff --git a/tests/tool-abort.test.ts b/tests/tool-abort.test.ts new file mode 100644 index 0000000..bf16e66 --- /dev/null +++ b/tests/tool-abort.test.ts @@ -0,0 +1,132 @@ +/** + * The abort-cascade fix (2026-09-19): pi passes every run's AbortSignal to + * tool execute — the bridge used to drop it, so agent.abort() only landed + * at the NEXT step boundary and a long-running local tool (the grep that + * outlived its turn by minutes) held the turn hostage until the dev + * server's memory restart killed the run mid-tool, freezing the record at + * its write-ahead state. toPiTool now threads the signal and races every + * execute against the system-level timeout; the fs tools honor it with + * partial results (fs-tools.test.ts covers their side). + */ +import assert from "node:assert/strict"; +import { afterEach, describe, it } from "node:test"; +import { z } from "zod"; +import { lhTool, toPiTool, toolTimeoutMs, ToolInputError } from "../src/agent/lh-tool.ts"; + +const DEFAULT_TIMEOUT_MS = 30_000; + +afterEach(() => { + toolTimeoutMs.current = DEFAULT_TIMEOUT_MS; +}); + +/** A tool that captures ctx.signal, then either hangs or resolves `value`. */ +function probeTool(captured: { signal?: AbortSignal }, opts: { hang?: boolean; value?: string } = {}) { + return lhTool({ + name: "probe", + description: "test probe", + schema: z.object({}), + execute: (_input, ctx) => { + captured.signal = ctx?.signal; + if (opts.hang) return new Promise(() => {}); // never settles, no timer + return opts.value ?? "ok"; + }, + }); +} + +describe("toPiTool abort + timeout guard", () => { + it("forwards a live abort signal to the tool's ctx", async () => { + // Even with no parent signal, the tool still gets the guard's combined + // signal (parent+timeout) — a stable reference tools can poll. + const captured: { signal?: AbortSignal } = {}; + const pi = toPiTool(probeTool(captured)); + await pi.execute("c1", {}, undefined as never); + assert.ok(captured.signal, "the tool received a signal"); + assert.equal(captured.signal!.aborted, false); + // The forwarded signal is the guard's COMBINED one — aborting the + // parent flips it, which is what long-running tools poll. + const ac = new AbortController(); + const captured2: { signal?: AbortSignal } = {}; + const pi2 = toPiTool(probeTool(captured2, { value: "done" })); + const pending = pi2.execute("c1b", {}, ac.signal); + ac.abort(); + await pending; + assert.equal(captured2.signal!.aborted, true); + }); + + it("a signal-wired tool observes the run's abort", async () => { + const ac = new AbortController(); + const captured: { signal?: AbortSignal } = {}; + const pi = toPiTool(probeTool(captured, { hang: true })); + const pending = pi.execute("c2", {}, ac.signal).catch((e: Error) => e.message); + ac.abort(); + assert.match((await pending) as string, /Operation aborted/); + assert.equal(captured.signal!.aborted, true); + }); + + it("rejects a hanging tool at the system-level timeout", async () => { + toolTimeoutMs.current = 50; + const pi = toPiTool(probeTool({}, { hang: true })); + await assert.rejects(pi.execute("c3", {}, undefined as never), /timed out/); + }); + + it("a pre-aborted run rejects immediately with pi's wording", async () => { + const ac = new AbortController(); + ac.abort(); + const pi = toPiTool(probeTool({}, { hang: true })); + await assert.rejects(pi.execute("c4", {}, ac.signal), /Operation aborted/); + }); + + it("a tool that honors the signal settles before the timeout wins", async () => { + // The race resolves with the tool's own (partial) result — the fs + // tools' contract: abort yields partial results, not an error throw. + toolTimeoutMs.current = 2_000; + const tool = lhTool({ + name: "partial", + description: "test", + schema: z.object({}), + execute: (_input, ctx) => + new Promise((resolve) => { + // The fs-tools contract: poll upfront (a listener on an + // already-aborted signal never fires), else listen. + const signal = ctx!.signal!; + if (signal.aborted) return resolve("partial results"); + signal.addEventListener("abort", () => resolve("partial results"), { once: true }); + }), + }); + const ac = new AbortController(); + const pi = toPiTool(tool); + const pending = pi.execute("c5", {}, ac.signal); + ac.abort(); + const result = (await pending) as { content: Array<{ text?: string }> }; + assert.equal(result.content[0]!.text, "partial results"); + }); + + it("question tools bypass the guard — the placeholder is instant", async () => { + const ac = new AbortController(); + ac.abort(); + const tool = lhTool({ + name: "quiz", + description: "test", + schema: z.object({}), + clientQuestion: true, + }); + const pi = toPiTool(tool); + const result = (await pi.execute("c6", {}, ac.signal)) as { content: Array<{ text?: string }> }; + assert.match(result.content[0]!.text!, /Question posed/); + }); + + it("validation errors pass through the guard unchanged", async () => { + const tool = lhTool({ + name: "strict", + description: "test", + schema: z.object({ n: z.number() }), + execute: () => "never", + }); + const pi = toPiTool(tool); + await assert.rejects(pi.execute("c7", { n: "not a number" }, undefined as never), (e: unknown) => { + assert.ok(e instanceof ToolInputError); + assert.equal(e.toolName, "strict"); + return true; + }); + }); +});