diff --git a/packages/junior/src/chat/tools/web/network.ts b/packages/junior/src/chat/tools/web/network.ts index f9801369..f6694bdd 100644 --- a/packages/junior/src/chat/tools/web/network.ts +++ b/packages/junior/src/chat/tools/web/network.ts @@ -265,13 +265,18 @@ export async function withTimeout( task: Promise, timeoutMs: number, label: string, + options?: { onTimeout?: () => void }, ): Promise { let timer: ReturnType | undefined; const timeoutPromise = new Promise((_, reject) => { - timer = setTimeout( - () => reject(new Error(`${label} timed out`)), - timeoutMs, - ); + timer = setTimeout(() => { + reject(new Error(`${label} timed out`)); + try { + options?.onTimeout?.(); + } catch { + // Timeout semantics must not depend on cleanup success. + } + }, timeoutMs); }); try { diff --git a/packages/junior/src/chat/tools/web/search.ts b/packages/junior/src/chat/tools/web/search.ts index 689b9524..b0a2a61f 100644 --- a/packages/junior/src/chat/tools/web/search.ts +++ b/packages/junior/src/chat/tools/web/search.ts @@ -97,6 +97,7 @@ export function createWebSearchTool() { // search tool on a search-tuned model and only allow an explicit // override. const model = process.env.AI_WEB_SEARCH_MODEL ?? DEFAULT_SEARCH_MODEL; + const controller = new AbortController(); try { // AI SDK Gateway reads AI_GATEWAY_API_KEY or ambient Vercel OIDC @@ -113,9 +114,11 @@ export function createWebSearchTool() { }), }, toolChoice: { type: "tool", toolName: SEARCH_TOOL_NAME }, + abortSignal: controller.signal, }), SEARCH_TIMEOUT_MS, "webSearch", + { onTimeout: () => controller.abort() }, ); const results = parseSearchResults(response.toolResults, maxResults); diff --git a/packages/junior/tests/unit/web/web-search.test.ts b/packages/junior/tests/unit/web/web-search.test.ts index 51537469..da55b3fb 100644 --- a/packages/junior/tests/unit/web/web-search.test.ts +++ b/packages/junior/tests/unit/web/web-search.test.ts @@ -74,6 +74,7 @@ describe("createWebSearchTool", () => { model: { model: "xai/grok-4-fast-reasoning" }, prompt: "vercel ai gateway", toolChoice: { type: "tool", toolName: "parallelSearch" }, + abortSignal: expect.any(AbortSignal), }), ); expect(result).toEqual({ @@ -160,6 +161,93 @@ describe("createWebSearchTool", () => { vi.useRealTimers(); }); + it("aborts the generateText call on timeout", async () => { + vi.useFakeTimers(); + let capturedSignal: AbortSignal | undefined; + vi.mocked(generateText).mockImplementation(((opts: { + abortSignal?: AbortSignal; + }) => { + capturedSignal = opts.abortSignal; + return new Promise(() => { + // Intentionally unresolved to trigger tool timeout. + }); + }) as never); + + const tool = createWebSearchTool(); + if (typeof tool.execute !== "function") { + throw new Error("webSearch execute function missing"); + } + + const pending = tool.execute({ query: "slow query" }, {} as never); + expect(capturedSignal?.aborted).toBe(false); + await vi.advanceTimersByTimeAsync(60_000); + await pending; + expect(capturedSignal?.aborted).toBe(true); + vi.useRealTimers(); + }); + + it("does not abort signal on successful search", async () => { + let capturedSignal: AbortSignal | undefined; + vi.mocked(generateText).mockImplementation(((opts: { + abortSignal?: AbortSignal; + }) => { + capturedSignal = opts.abortSignal; + return Promise.resolve({ toolResults: [] }); + }) as never); + + const tool = createWebSearchTool(); + if (typeof tool.execute !== "function") { + throw new Error("webSearch execute function missing"); + } + + await tool.execute({ query: "fast query" }, {} as never); + expect(capturedSignal?.aborted).toBe(false); + }); + + it("still reports timeout even if abort signal cleanup throws", async () => { + vi.useFakeTimers(); + const brokenController = new AbortController(); + const originalAbort = brokenController.abort.bind(brokenController); + brokenController.abort = () => { + originalAbort(); + throw new Error("abort listener blew up"); + }; + + // Patch AbortController to return our broken one + const originalAC = globalThis.AbortController; + globalThis.AbortController = class extends originalAC { + constructor() { + super(); + return brokenController as unknown as AbortController; + } + } as typeof AbortController; + + vi.mocked(generateText).mockImplementation( + () => + new Promise(() => { + // Intentionally unresolved to trigger tool timeout. + }) as never, + ); + + const tool = createWebSearchTool(); + if (typeof tool.execute !== "function") { + throw new Error("webSearch execute function missing"); + } + + const pending = tool.execute({ query: "boom query" }, {} as never); + await vi.advanceTimersByTimeAsync(60_000); + const result = await pending; + + globalThis.AbortController = originalAC; + + expect(result).toMatchObject({ + ok: false, + timeout: true, + error: "web search failed: webSearch timed out", + }); + vi.useRealTimers(); + }); + it("marks authentication failures as non-retryable", async () => { vi.mocked(generateText).mockRejectedValueOnce( new Error(