diff --git a/src/chat-history.ts b/src/chat-history.ts index 36f98f8..8f27b2f 100644 --- a/src/chat-history.ts +++ b/src/chat-history.ts @@ -99,6 +99,7 @@ export async function summarizeHistory(agentDir: string, branch: string): Promis prompt, dir: agentDir, maxTurns: 1, + timeoutMs: 30_000, replaceBuiltinTools: true, tools: [], systemPrompt: "You are a concise summarizer. Output only the summary, nothing else.", diff --git a/src/sdk-types.ts b/src/sdk-types.ts index 5ab79ca..be049ac 100644 --- a/src/sdk-types.ts +++ b/src/sdk-types.ts @@ -145,6 +145,8 @@ export interface QueryOptions { hooks?: GCHooks; maxTurns?: number; abortController?: AbortController; + /** Maximum duration of each model turn before the query is aborted; this is not a whole-query timeout. */ + timeoutMs?: number; sessionId?: string; /** MCP servers to connect to. Merged with manifest `mcp_servers` (these win on key collision). */ mcpServers?: Record; diff --git a/src/sdk.ts b/src/sdk.ts index 55b941c..893391a 100644 --- a/src/sdk.ts +++ b/src/sdk.ts @@ -92,6 +92,10 @@ export function query(options: QueryOptions): Query { const collectedMessages: GCMessage[] = []; const ac = options.abortController ?? new AbortController(); const costTracker = new CostTracker(); + let removeAbortForwarder: (() => void) | undefined; + const abortQuery = () => { + ac.abort(); + }; // These are set once the agent is loaded (async init below) let _sessionId = options.sessionId ?? ""; @@ -128,6 +132,11 @@ export function query(options: QueryOptions): Query { // Async initialization + run const runPromise = (async () => { try { + if (options.timeoutMs !== undefined && + (!Number.isFinite(options.timeoutMs) || options.timeoutMs <= 0)) { + throw new Error("timeoutMs must be a finite number greater than zero"); + } + // Validate mutually exclusive options if (options.repo && options.sandbox) { throw new Error("repo and sandbox options are mutually exclusive"); @@ -315,6 +324,21 @@ export function query(options: QueryOptions): Query { ...modelOptions, }, }); + const forwardAbort = () => agent.abort(); + ac.signal.addEventListener("abort", forwardAbort, { once: true }); + removeAbortForwarder = () => ac.signal.removeEventListener("abort", forwardAbort); + + const promptWithTimeout = async (prompt: string) => { + if (ac.signal.aborted) return; + const timer = options.timeoutMs === undefined + ? undefined + : setTimeout(abortQuery, options.timeoutMs); + try { + await otelContext.with(_session.ctx, () => agent.prompt(prompt)); + } finally { + if (timer !== undefined) clearTimeout(timer); + } + }; // 9. Subscribe to events and map to GCMessage agent.subscribe((event: AgentEvent) => { @@ -494,10 +518,42 @@ export function query(options: QueryOptions): Query { } }); + const emitAbortedResponse = () => { + pushMsg({ + type: "assistant", + content: "", + model: (loaded.model as any).id ?? "unknown", + provider: (loaded.model as any).provider ?? "unknown", + stopReason: "aborted", + usage: { + inputTokens: 0, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + totalTokens: 0, + costUsd: 0, + }, + }); + }; + const finishAbortedQuery = () => { + emitAbortedResponse(); + channel.finish(); + }; + // 10. Send prompt — run inside the session span's context so that // gen_ai.chat and gitagent.tool.execute spans become children of // gitagent.agent.session. if (typeof options.prompt === "string") { + if (ac.signal.aborted) { + pushMsg({ + type: "system", + subtype: "session_start", + content: `Agent ${loaded.manifest.name} initialized`, + metadata: { sessionId: _sessionId }, + }); + finishAbortedQuery(); + return; + } // Fire pre_query hook before sending to LLM if (hooksConfig?.hooks.pre_query) { const result = await runHooks(hooksConfig.hooks.pre_query, loaded.agentDir, { @@ -515,12 +571,18 @@ export function query(options: QueryOptions): Query { return; } } - await otelContext.with(_session.ctx, () => - agent.prompt(options.prompt as string), - ); + await promptWithTimeout(options.prompt as string); + if (ac.signal.aborted) { + finishAbortedQuery(); + return; + } } else { // Multi-turn: iterate the async iterable for await (const userMsg of options.prompt) { + if (ac.signal.aborted) { + finishAbortedQuery(); + return; + } pushMsg({ type: "user", content: userMsg.content }); // Fire pre_query hook for each turn if (hooksConfig?.hooks.pre_query) { @@ -539,9 +601,11 @@ export function query(options: QueryOptions): Query { return; } } - await otelContext.with(_session.ctx, () => - agent.prompt(userMsg.content), - ); + await promptWithTimeout(userMsg.content); + if (ac.signal.aborted) { + finishAbortedQuery(); + return; + } } } @@ -558,6 +622,8 @@ export function query(options: QueryOptions): Query { // Ensure channel finishes even if no agent_end event channel.finish(); } finally { + removeAbortForwarder?.(); + removeAbortForwarder = undefined; // Tear down MCP servers on every exit path — success, hook-block // early-return, abort, and error (this finally runs before the // .catch() handler below). cleanup() is idempotent. @@ -604,7 +670,7 @@ export function query(options: QueryOptions): Query { // Build the Query object (AsyncGenerator + helpers) const generator: Query = { abort() { - ac.abort(); + abortQuery(); }, steer(_message: string) { diff --git a/test/sdk.test.ts b/test/sdk.test.ts index a71fcc6..1e4f0cd 100644 --- a/test/sdk.test.ts +++ b/test/sdk.test.ts @@ -1,5 +1,8 @@ import { describe, it, before } from "node:test"; import assert from "node:assert/strict"; +import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { join } from "node:path"; +import { tmpdir } from "node:os"; // Dynamic imports since the project is ESM let query: typeof import("../dist/exports.js").query; @@ -204,6 +207,65 @@ describe("wrapToolWithProgrammaticHooks()", () => { // ── query() error handling ───────────────────────────────────────────── describe("query()", () => { + it("reports invalid timeout values without loading an agent", async () => { + const messages: any[] = []; + for await (const msg of query({ + prompt: "hello", + dir: "/nonexistent/path", + timeoutMs: 0, + })) { + messages.push(msg); + } + + const errorMsg = messages.find((m) => m.type === "system" && m.subtype === "error"); + assert.match(errorMsg?.content ?? "", /timeoutMs must be a finite number greater than zero/); + }); + + it("honors an already-aborted controller before starting the model", async () => { + const agentDir = mkdtempSync(join(tmpdir(), "gitagent-abort-")); + try { + writeFileSync(join(agentDir, "agent.yaml"), `spec_version: "1.0"\nname: test\nversion: "1.0.0"\ndescription: test\nmodel:\n preferred: openai:gpt-4o-mini\n fallback: []\ntools: []\nruntime:\n max_turns: 1\n`); + const controller = new AbortController(); + controller.abort(); + const messages: any[] = []; + const drain = (async () => { + for await (const msg of query({ + prompt: "hello", + dir: agentDir, + abortController: controller, + timeoutMs: 100, + })) { + messages.push(msg); + } + })(); + let deadlineTimer: ReturnType; + const deadline = new Promise((_, reject) => { + deadlineTimer = setTimeout( + () => reject(new Error("aborted query did not settle")), + 2000, + ); + }); + try { + await Promise.race([drain, deadline]); + } finally { + clearTimeout(deadlineTimer!); + } + assert.ok(messages.some((m) => m.type === "system")); + const assistantMessages = messages.filter((m) => m.type === "assistant"); + assert.equal(messages.some((m) => m.type === "delta"), false); + assert.equal(assistantMessages.length, 1); + assert.equal(assistantMessages[0].stopReason, "aborted"); + assert.equal(assistantMessages[0].usage.totalTokens, 0); + assert.equal( + assistantMessages[0].content, + "", + "a pre-aborted query must not produce model content", + ); + } finally { + rmSync(agentDir, { recursive: true, force: true }); + } + }); + it("emits error system message when agent dir is invalid", async () => { const messages: any[] = []; for await (const msg of query({