/** * wait-queued.test.ts — get_subagent_result(wait: true) lifecycle behavior. * * Queued records have no promise yet (it's created when the queue starts * them), so the old `status === "running" && record.promise` condition * skipped the wait entirely and returned "still running" — forcing the * caller into a poll loop against the concurrency queue. * * Wiring test through the REAL extension: spawn background agents until one * queues, call the real tool with wait:true, drain the queue, and assert the * call returns the final result. */ import { describe, expect, it, vi } from "vitest"; vi.mock("../src/agent-runner.js", async () => { const actual = await vi.importActual("../src/agent-runner.js"); return { ...actual, runAgent: vi.fn() }; }); import { runAgent } from "../src/agent-runner.js"; import subagentsExtension from "../src/index.js"; function makePi() { const tools = new Map(); const lifecycle = new Map(); const pi = { registerMessageRenderer: vi.fn(), registerTool: vi.fn((t: any) => tools.set(t.name, t)), registerCommand: vi.fn(), on: vi.fn((event: string, handler: any) => lifecycle.set(event, handler)), events: { emit: vi.fn(), on: vi.fn(() => vi.fn()), }, appendEntry: vi.fn(), sendMessage: vi.fn(), } as any; return { pi, tools, lifecycle }; } function ctx() { return { hasUI: false, ui: { setStatus: vi.fn(), setWidget: vi.fn(), notify: vi.fn() }, cwd: process.cwd(), model: undefined, modelRegistry: { find: vi.fn(), getAvailable: vi.fn(() => []) }, sessionManager: { getSessionId: vi.fn(() => "s1"), getBranch: vi.fn(() => []) }, getSystemPrompt: vi.fn(() => "parent"), } as any; } const textOf = (r: any): string => r.content[0].text; const flush = async () => { await new Promise((r) => setImmediate(r)); await new Promise((r) => setImmediate(r)); }; /** runAgent mock where each call blocks until we resolve it manually. */ function deferredRuns() { const resolvers: Array<(v: any) => void> = []; vi.mocked(runAgent).mockImplementation( () => new Promise((resolve) => { resolvers.push(() => resolve({ responseText: "THE-RESULT-PAYLOAD", session: { dispose: vi.fn() } as any, aborted: false, steered: false, }), ); }) as any, ); return resolvers; } async function spawnBackground(tools: Map): Promise<{ id: string; queued: boolean }> { const r = await tools.get("Agent").execute( "tc-spawn", { prompt: "go", description: "queued-wait test agent", subagent_type: "general-purpose", run_in_background: true }, undefined, undefined, ctx(), ); const id = /Agent ID: (\S+)/.exec(textOf(r))![1]; return { id, queued: textOf(r).includes("queued in background") }; } describe("get_subagent_result wait:true on a queued agent", () => { it("waits through queue start and returns the result (no 'still running')", async () => { const { pi, tools, lifecycle } = makePi(); subagentsExtension(pi); const resolvers = deferredRuns(); // Spawn until one lands in the queue (concurrency limit is config-dependent). let queuedId: string | undefined; for (let i = 0; i < 20 && !queuedId; i++) { const { id, queued } = await spawnBackground(tools); if (queued) queuedId = id; } expect(queuedId, "expected to hit the concurrency limit within 20 spawns").toBeDefined(); // wait:true on the QUEUED agent — must not return "still running". const waitPromise = tools .get("get_subagent_result") .execute("tc-wait", { agent_id: queuedId, wait: true }, undefined, undefined, ctx()); // Drain: resolve running agents until the queued one starts and finishes. let settled = false; void waitPromise.then(() => { settled = true; }); for (let i = 0; i < 40 && !settled; i++) { while (resolvers.length > 0) resolvers.shift()!(); await flush(); await new Promise((r) => setTimeout(r, 100)); // outlive one 250ms poll tick } const result = await waitPromise; expect(textOf(result)).toContain("THE-RESULT-PAYLOAD"); expect(textOf(result)).not.toContain("still running"); await new Promise((r) => setTimeout(r, 350)); expect(JSON.stringify(pi.sendMessage.mock.calls)).not.toContain(queuedId); await lifecycle.get("session_shutdown")?.(); }, 20_000); it("aborts a running result wait without aborting or consuming the child", async () => { const { pi, tools, lifecycle } = makePi(); subagentsExtension(pi); let resolveRun: (() => void) | undefined; let childSignal: AbortSignal | undefined; vi.mocked(runAgent).mockImplementation( (_ctx, _type, _prompt, options) => new Promise((resolve) => { childSignal = options.signal; resolveRun = () => resolve({ responseText: "THE-RESULT-PAYLOAD", session: { dispose: vi.fn() } as any, aborted: false, steered: false, }); }), ); const { id } = await spawnBackground(tools); const controller = new AbortController(); const removeListener = vi.spyOn(controller.signal, "removeEventListener"); const waitOutcome = tools .get("get_subagent_result") .execute("tc-wait-abort", { agent_id: id, wait: true }, controller.signal, undefined, ctx()) .then( () => "resolved", (error: unknown) => error instanceof Error ? error.name : String(error), ); controller.abort(); const outcome = await Promise.race([ waitOutcome, new Promise((resolve) => setTimeout(() => resolve("timed-out"), 100)), ]); const childWasAborted = childSignal?.aborted; resolveRun?.(); await flush(); await waitOutcome; await new Promise((r) => setTimeout(r, 350)); const completedResult = await tools .get("get_subagent_result") .execute("tc-result", { agent_id: id }, undefined, undefined, ctx()); await lifecycle.get("session_shutdown")?.(); expect(outcome).toBe("AbortError"); expect(childWasAborted).toBe(false); expect(removeListener).toHaveBeenCalledTimes(1); expect(pi.sendMessage).toHaveBeenCalledTimes(1); expect(textOf(completedResult)).toContain("THE-RESULT-PAYLOAD"); }); it("aborts a queued result wait before the agent starts", async () => { const { pi, tools, lifecycle } = makePi(); subagentsExtension(pi); const resolvers = deferredRuns(); let queuedId: string | undefined; for (let i = 0; i < 20 && !queuedId; i++) { const { id, queued } = await spawnBackground(tools); if (queued) queuedId = id; } expect(queuedId, "expected to hit the concurrency limit within 20 spawns").toBeDefined(); const controller = new AbortController(); const waitOutcome = tools .get("get_subagent_result") .execute("tc-queued-abort", { agent_id: queuedId, wait: true }, controller.signal, undefined, ctx()) .then( () => "resolved", (error: unknown) => error instanceof Error ? error.name : String(error), ); controller.abort(); const outcome = await Promise.race([ waitOutcome, new Promise((resolve) => setTimeout(() => resolve("timed-out"), 100)), ]); let completedResult: any; for (let i = 0; i < 40 && !completedResult; i++) { while (resolvers.length > 0) resolvers.shift()!(); await flush(); const result = await tools .get("get_subagent_result") .execute("tc-queued-result", { agent_id: queuedId }, undefined, undefined, ctx()); if (textOf(result).includes("THE-RESULT-PAYLOAD")) completedResult = result; await new Promise((r) => setTimeout(r, 25)); } await waitOutcome; await lifecycle.get("session_shutdown")?.(); expect(outcome).toBe("AbortError"); expect(textOf(completedResult)).toContain("THE-RESULT-PAYLOAD"); }); });