/** * Tests for AgentSession concurrent prompt guard. */ import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; import * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; import { scheduler } from "node:timers/promises"; import { Agent, AgentBusyError, type AgentMessage, type AgentTool } from "@oh-my-pi/pi-agent-core"; import type { AssistantMessage, Message, ToolCall } from "@oh-my-pi/pi-ai"; import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock"; import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream"; import { getBundledModel } from "@oh-my-pi/pi-catalog/models"; import { AsyncJobManager } from "@oh-my-pi/pi-coding-agent/async"; import type { Rule } from "@oh-my-pi/pi-coding-agent/capability/rule"; import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; import { type SettingPath, Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { TtsrManager } from "@oh-my-pi/pi-coding-agent/export/ttsr"; import type { ExtensionRunner } from "@oh-my-pi/pi-coding-agent/extensibility/extensions"; import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage"; import { convertToLlm } from "@oh-my-pi/pi-coding-agent/session/messages"; import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; import { removeSyncWithRetries, Snowflake } from "@oh-my-pi/pi-utils"; import { type } from "arktype"; import { createAssistantMessage } from "./helpers/agent-session-setup"; // Mock stream that mimics AssistantMessageEventStream // AgentSession schedules its TTSR retry and context-promotion continuations // through `scheduler.wait(delayMs, { signal })` (node:timers/promises), with // blind 50ms/100ms "settle" delays. Tests that drive a continuation to // completion would otherwise pay that wall-clock time on every run. This spy // collapses the blind delay to a single macrotask hop (`scheduler.wait(0)`) // while preserving the real abort-signal semantics, so the continuation still // fires only after the aborted/overflowed turn has been recorded. Each test // that opts in must run inside a block whose afterEach restores mocks. const originalSchedulerWait = scheduler.wait.bind(scheduler); function collapseSchedulerSettleDelays(): void { vi.spyOn(scheduler, "wait").mockImplementation((_delayMs, options) => originalSchedulerWait(0, options)); } describe("AgentSession concurrent prompt guard", () => { let session: AgentSession; let tempDir: string; const authStorages: AuthStorage[] = []; beforeEach(() => { // Collapse scheduler settle delays so the post-abort auto-continue and // dispose teardown are deterministic instead of racing the wall clock. collapseSchedulerSettleDelays(); tempDir = path.join(os.tmpdir(), `pi-concurrent-test-${Snowflake.next()}`); fs.mkdirSync(tempDir, { recursive: true }); }); afterEach(async () => { if (session) { await session.dispose(); } for (const authStorage of authStorages.splice(0)) { authStorage.close(); } if (tempDir && fs.existsSync(tempDir)) { removeSyncWithRetries(tempDir); } vi.restoreAllMocks(); AsyncJobManager.resetForTests(); }); async function createSession(settingsOverrides?: Partial>) { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; let abortSignal: AbortSignal | undefined; // Use a stream function that responds to abort const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [], }, streamFn: (_model, _context, options) => { abortSignal = options?.signal; const stream = new AssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "start", partial: createAssistantMessage("") }); if (abortSignal) { abortSignal.addEventListener( "abort", () => { stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") }); }, { once: true }, ); } }); return stream; }, }); const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(settingsOverrides); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, }); return session; } async function waitFor(predicate: () => boolean, timeoutMs = 500): Promise { const deadline = Date.now() + timeoutMs; while (Date.now() < deadline) { if (predicate()) return; await Bun.sleep(1); } throw new Error("Timed out waiting for condition"); } it("should throw when prompt() called while streaming", async () => { await createSession(); // Start first prompt (don't await, it will block until abort) const firstPrompt = session.prompt("First message"); await waitFor(() => session.isStreaming); // Second prompt should reject await expect(session.prompt("Second message")).rejects.toBeInstanceOf(AgentBusyError); // Cleanup await session.abort(); await firstPrompt.catch(() => {}); // Ignore abort error }); it("should allow steer() while streaming", async () => { await createSession(); // Start first prompt const firstPrompt = session.prompt("First message"); await waitFor(() => session.isStreaming); // steer should work while streaming await session.steer("Steer while streaming"); expect(session.queuedMessageCount).toBe(1); // Cleanup session.agent.clearAllQueues(); await session.abort(); await firstPrompt.catch(() => {}); }); it("should allow followUp() while streaming", async () => { await createSession(); // Start first prompt const firstPrompt = session.prompt("First message"); await waitFor(() => session.isStreaming); // followUp should work while streaming await session.followUp("Follow-up while streaming"); expect(session.queuedMessageCount).toBe(1); // Cleanup session.agent.clearAllQueues(); await session.abort(); await firstPrompt.catch(() => {}); }); it("queues sendUserMessage as steer while streaming without AgentBusyError", async () => { await createSession(); const firstPrompt = session.prompt("First message"); await waitFor(() => session.isStreaming); // The first agent loop may dequeue a steer before the assertion runs, so // observe agent.steer itself rather than the residual queue length. const steered: AgentMessage[] = []; const originalSteer = session.agent.steer.bind(session.agent); session.agent.steer = (message: AgentMessage) => { steered.push(message); originalSteer(message); }; // Extension path: no deliverAs while busy must queue, not throw. await expect(session.sendUserMessage("hello from extension")).resolves.toBeUndefined(); expect(steered).toHaveLength(1); const queued = steered[0]; expect(queued?.role).toBe("user"); if (queued?.role === "user") { expect(queued.content).toEqual([{ type: "text", text: "hello from extension" }]); expect(queued.steering).toBe(true); } session.agent.clearAllQueues(); await session.abort(); await firstPrompt.catch(() => {}); }); it("sendUserMessage without deliverAs preserves prompt-flow keyword notices while streaming", async () => { await createSession({ "magicKeywords.enabled": true, "magicKeywords.ultrathink": true }); const firstPrompt = session.prompt("First message"); await waitFor(() => session.isStreaming); try { await session.sendUserMessage("ultrathink fix via extension"); const queuedShape = session.agent .peekSteeringQueue() .map(message => (message.role === "custom" ? message.customType : message.role)); expect(queuedShape).toEqual(["ultrathink-notice", "user"]); expect(session.getQueuedMessages()).toEqual({ steering: ["ultrathink fix via extension"], followUp: [], }); } finally { session.agent.clearAllQueues(); await session.abort(); await firstPrompt.catch(() => {}); } }); it("sendUserMessage without deliverAs starts a normal prompt when idle", async () => { await createSession(); let rejected: unknown; let settled = false; const turn = session .sendUserMessage("Idle extension message") .catch(error => { rejected = error; }) .finally(() => { settled = true; }); try { await waitFor(() => session.isStreaming || settled); if (rejected) throw rejected; expect(session.isStreaming).toBe(true); expect(settled).toBe(false); expect(session.getQueuedMessages()).toEqual({ steering: [], followUp: [] }); } finally { await session.abort(); await turn; } }); it("delivers hidden nextTurn stop reactions through the next LLM call without exposing them in the visible queue", async () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; let firstStream: AssistantMessageEventStream | undefined; const callMessages: Message[][] = []; const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [], }, convertToLlm, streamFn: (_model, context) => { callMessages.push([...context.messages]); const stream = new AssistantMessageEventStream(); queueMicrotask(() => { stream.push({ type: "start", partial: createAssistantMessage("") }); if (callMessages.length > 1) { stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Resumed") }); return; } }); firstStream = stream; return stream; }, }); const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, }); const firstPrompt = session.prompt("First message"); await waitFor(() => session.isStreaming && firstStream !== undefined && callMessages.length === 1); await session.sendCustomMessage( { customType: "autoresearch-resume", content: "Hidden stop reaction", display: false, attribution: "agent", }, { deliverAs: "nextTurn", triggerTurn: true }, ); expect(session.queuedMessageCount).toBe(0); expect(session.getQueuedMessages()).toEqual({ steering: [], followUp: [] }); firstStream?.push({ type: "done", reason: "stop", message: createAssistantMessage("Done") }); await firstPrompt; await session.waitForIdle(); expect(callMessages).toHaveLength(2); expect( callMessages[1]?.some(message => { if (typeof message.content === "string") { return message.content.includes("Hidden stop reaction"); } return message.content.some( content => content.type === "text" && content.text.includes("Hidden stop reaction"), ); }), ).toBe(true); }); it("continues a main session from session_stop feedback before settling", async () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; const mock = createMockModel({ handler: () => ({ content: ["Done"] }), }); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [] }, streamFn: mock.stream, convertToLlm, }); const stopEvents: Array<{ stop_hook_active: boolean; session_id: string; turn_id: number; last_assistant_message?: AgentMessage; }> = []; const eventOrder: string[] = []; const extensionRunner = { emit: vi.fn(event => { eventOrder.push(event.type); return Promise.resolve(undefined); }), emitBeforeAgentStart: vi.fn().mockResolvedValue(undefined), hasHandlers: vi.fn((eventType: string) => eventType === "session_stop"), emitSessionStop: vi.fn(event => { eventOrder.push("session_stop"); stopEvents.push(event); if (stopEvents.length === 1) { return Promise.resolve({ continue: true, additionalContext: "Mission incomplete; continue." }); } return Promise.resolve(undefined); }), } as unknown as ExtensionRunner; const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, extensionRunner }); await session.prompt("First message"); await session.waitForIdle(); const callMessages = mock.calls.map(call => call.context.messages); expect(callMessages).toHaveLength(2); expect( callMessages[1]?.some(message => typeof message.content === "string" ? message.content.includes("Mission incomplete; continue.") : message.content.some( content => content.type === "text" && content.text.includes("Mission incomplete; continue."), ), ), ).toBe(true); expect(eventOrder.filter(type => type === "session_stop" || type === "agent_end")).toEqual([ "session_stop", "agent_end", "session_stop", "agent_end", ]); expect(stopEvents.map(event => event.stop_hook_active)).toEqual([false, true]); expect(stopEvents.map(event => event.turn_id)).toEqual([0, 0]); expect(stopEvents[0]?.session_id).toBe(session.sessionId); expect(stopEvents[0]?.last_assistant_message?.role).toBe("assistant"); }); it("uses non-empty session_stop reason when additional context is empty", async () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; const mock = createMockModel({ handler: () => ({ content: ["Done"] }), }); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [] }, streamFn: mock.stream, convertToLlm, }); let stopCount = 0; const extensionRunner = { emit: vi.fn().mockResolvedValue(undefined), emitBeforeAgentStart: vi.fn().mockResolvedValue(undefined), hasHandlers: vi.fn((eventType: string) => eventType === "session_stop"), emitSessionStop: vi.fn(() => { stopCount++; if (stopCount === 1) { return Promise.resolve({ continue: true, additionalContext: "", reason: "Continue from fallback reason.", }); } return Promise.resolve(undefined); }), } as unknown as ExtensionRunner; const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, extensionRunner }); await session.prompt("First message"); await session.waitForIdle(); expect(mock.calls).toHaveLength(2); expect( mock.calls[1]?.context.messages.some(message => typeof message.content === "string" ? message.content.includes("Continue from fallback reason.") : message.content.some( content => content.type === "text" && content.text.includes("Continue from fallback reason."), ), ), ).toBe(true); }); it("does not continue session_stop feedback after aborting a slow hook", async () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; const mock = createMockModel({ handler: () => ({ content: ["Done"] }), }); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [] }, streamFn: mock.stream, convertToLlm, }); const stopHook = Promise.withResolvers<{ continue: true; additionalContext: string }>(); const emitSessionStop = vi.fn(() => (emitSessionStop.mock.calls.length === 1 ? stopHook.promise : undefined)); const extensionRunner = { emit: vi.fn().mockResolvedValue(undefined), emitBeforeAgentStart: vi.fn().mockResolvedValue(undefined), hasHandlers: vi.fn((eventType: string) => eventType === "session_stop"), emitSessionStop, } as unknown as ExtensionRunner; const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, extensionRunner }); const promptPromise = session.prompt("First message"); await waitFor(() => emitSessionStop.mock.calls.length === 1); const abortPromise = session.abort(); stopHook.resolve({ continue: true, additionalContext: "Should not run after abort." }); await abortPromise; await promptPromise; await session.waitForIdle(); expect(mock.calls).toHaveLength(1); expect(session.queuedMessageCount).toBe(0); await session.prompt("Second message"); await session.waitForIdle(); expect(mock.calls).toHaveLength(2); expect( mock.calls[1]?.context.messages.some(message => typeof message.content === "string" ? message.content.includes("Should not run after abort.") : message.content.some( content => content.type === "text" && content.text.includes("Should not run after abort."), ), ), ).toBe(false); }); it("caps consecutive session_stop continuations at eight", async () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; const mock = createMockModel({ handler: () => ({ content: ["Pass"] }), }); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [] }, streamFn: mock.stream, convertToLlm, }); const extensionRunner = { emit: vi.fn().mockResolvedValue(undefined), emitBeforeAgentStart: vi.fn().mockResolvedValue(undefined), hasHandlers: vi.fn((eventType: string) => eventType === "session_stop"), emitSessionStop: vi.fn(() => Promise.resolve({ decision: "block" as const, reason: "Run another pass." })), } as unknown as ExtensionRunner; const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, extensionRunner }); await session.prompt("First message"); await session.waitForIdle(); expect(mock.calls).toHaveLength(9); expect(extensionRunner.emitSessionStop).toHaveBeenCalledTimes(9); }); it("emits session_stop only after empty-stop recovery reaches a final stop", async () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; const mock = createMockModel({ responses: [{ content: [""] }, { content: ["Recovered"] }], }); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [] }, streamFn: mock.stream, convertToLlm, }); const extensionRunner = { emit: vi.fn().mockResolvedValue(undefined), emitBeforeAgentStart: vi.fn().mockResolvedValue(undefined), hasHandlers: vi.fn((eventType: string) => eventType === "session_stop"), emitSessionStop: vi.fn().mockResolvedValue(undefined), } as unknown as ExtensionRunner; const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, extensionRunner }); await session.prompt("First message"); await session.waitForIdle(); expect(mock.calls).toHaveLength(2); expect(extensionRunner.emitSessionStop).toHaveBeenCalledTimes(1); }); it("emits session_stop after empty-stop retry cap settles", async () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; const mock = createMockModel({ responses: [{ content: [""] }, { content: [""] }, { content: [""] }, { content: [""] }], }); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [] }, streamFn: mock.stream, convertToLlm, }); const extensionRunner = { emit: vi.fn().mockResolvedValue(undefined), emitBeforeAgentStart: vi.fn().mockResolvedValue(undefined), hasHandlers: vi.fn((eventType: string) => eventType === "session_stop"), emitSessionStop: vi.fn().mockResolvedValue(undefined), } as unknown as ExtensionRunner; const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, extensionRunner }); await session.prompt("First message"); await session.waitForIdle(); expect(mock.calls).toHaveLength(4); expect(extensionRunner.emitSessionStop).toHaveBeenCalledTimes(1); }); it("continues session_stop feedback in ACP sessions with deferred client turns", async () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; const mock = createMockModel({ handler: () => ({ content: ["Done"] }), }); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [] }, streamFn: mock.stream, convertToLlm, }); let stopCount = 0; const extensionRunner = { emit: vi.fn().mockResolvedValue(undefined), emitBeforeAgentStart: vi.fn().mockResolvedValue(undefined), hasHandlers: vi.fn((eventType: string) => eventType === "session_stop"), emitSessionStop: vi.fn(() => { stopCount++; if (stopCount === 1) { return Promise.resolve({ continue: true, additionalContext: "ACP stop continuation." }); } return Promise.resolve(undefined); }), } as unknown as ExtensionRunner; const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, extensionRunner }); session.setClientBridge({ capabilities: {}, deferAgentInitiatedTurns: true, }); await session.prompt("First message"); await session.waitForIdle(); expect(mock.calls).toHaveLength(2); expect( mock.calls[1]?.context.messages.some(message => typeof message.content === "string" ? message.content.includes("ACP stop continuation.") : message.content.some( content => content.type === "text" && content.text.includes("ACP stop continuation."), ), ), ).toBe(true); }); it("does not emit session_stop for subagent sessions", async () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; const mock = createMockModel({ handler: () => ({ content: ["Subagent done"] }), }); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [] }, streamFn: mock.stream, convertToLlm, }); const extensionRunner = { emit: vi.fn().mockResolvedValue(undefined), emitBeforeAgentStart: vi.fn().mockResolvedValue(undefined), hasHandlers: vi.fn((eventType: string) => eventType === "session_stop"), emitSessionStop: vi.fn().mockResolvedValue(undefined), } as unknown as ExtensionRunner; const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, extensionRunner, agentKind: "sub", }); await session.prompt("Subagent message"); await session.waitForIdle(); expect(mock.calls).toHaveLength(1); expect(extensionRunner.emit).toHaveBeenCalledWith({ type: "agent_end", messages: expect.any(Array) }); expect(extensionRunner.emitSessionStop).not.toHaveBeenCalled(); }); it("should allow prompt() after previous completes", async () => { // Create session with a stream that completes immediately const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; const mock = createMockModel({ handler: () => ({ content: ["Done"] }) }); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [], }, streamFn: mock.stream, }); const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, }); // First prompt completes await session.prompt("First message"); // Should not be streaming anymore expect(session.isStreaming).toBe(false); // Second prompt should work await expect(session.prompt("Second message")).resolves.toBe(true); }); it("queues extension follow-up user messages on an idle session without starting a turn", async () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; const mock = createMockModel({ handler: () => ({ content: ["Done"] }) }); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [], }, streamFn: mock.stream, }); const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-idle-followup.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models-idle-followup.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, }); await session.sendUserMessage("hello from session_start", { deliverAs: "followUp" }); expect(mock.calls).toHaveLength(0); expect(session.queuedMessageCount).toBe(1); }); // Regression: a subscriber that fires the next prompt synchronously from the // agent_end listener (the shape every wire transport ends up in — rpc-mode // stdout subscriber, ACP bridge, Cursor exec) must not collide with the // outgoing turn's still-unwinding in-flight bookkeeping. Before the wire-level // agent_end was deferred until #promptInFlightCount drops to 0, the // subscriber observed agent_end while Session.isStreaming was still true (the // agent's own `isStreaming` had flipped, but #promptWithMessage's finally had // not yet decremented the prompt-in-flight counter), and the next prompt // threw AgentBusyError. Surfaced as `RpcCommandError: prompt: Agent is // already processing` from omp-rpc clients (robomp triage reminder path). it("subscriber may prompt() synchronously from agent_end without AgentBusyError", async () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; const mock = createMockModel({ handler: () => ({ content: ["Done"] }) }); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [] }, streamFn: mock.stream, }); const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry }); const observedIsStreamingAtAgentEnd: boolean[] = []; const reentrantPromptResults: Array<"resolved" | { error: string }> = []; let reentrantPrompted = false; session.subscribe(event => { if (event.type !== "agent_end") return; observedIsStreamingAtAgentEnd.push(session.isStreaming); if (reentrantPrompted) return; reentrantPrompted = true; void session .prompt("Second message") .then(() => reentrantPromptResults.push("resolved")) .catch((err: Error) => reentrantPromptResults.push({ error: err.message })); }); await session.prompt("First message"); await waitFor(() => reentrantPromptResults.length > 0, 2000); await session.waitForIdle(); expect(observedIsStreamingAtAgentEnd).not.toContain(true); expect(reentrantPromptResults).toEqual(["resolved"]); }); it("queues idle ACP client-triggered custom messages instead of starting an ownerless turn", async () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; const mock = createMockModel({ handler: () => ({ content: ["Done"] }) }); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [], }, convertToLlm, streamFn: mock.stream, }); const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-acp-idle.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models-acp-idle.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, }); session.setClientBridge({ capabilities: {}, deferAgentInitiatedTurns: true, }); await session.prompt("First message"); expect(session.isStreaming).toBe(false); const callsAfterFirstPrompt = mock.calls.length; await session.sendCustomMessage( { customType: "async-result", content: "Background result", display: true, attribution: "agent", }, { deliverAs: "followUp", triggerTurn: true }, ); expect(mock.calls).toHaveLength(callsAfterFirstPrompt); expect(session.isStreaming).toBe(false); await session.prompt("Next user prompt"); await session.dispose(); session = undefined as unknown as AgentSession; expect(mock.calls).toHaveLength(callsAfterFirstPrompt + 1); expect( mock.calls.at(-1)?.context.messages.some(message => { if (typeof message.content === "string") { return message.content.includes("Background result"); } return message.content.some( content => content.type === "text" && content.text.includes("Background result"), ); }), ).toBe(true); }); it("runs drained ACP async completions as owned follow-up turns despite deferred client turns", async () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; const mock = createMockModel({ handler: () => ({ content: ["Done"] }) }); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [], }, convertToLlm, streamFn: mock.stream, }); const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-acp-async.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models-acp-async.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); const ownerId = "acp-session-a"; const deliveryGate = Promise.withResolvers(); let deliveryStarted = false; const asyncJobManager = new AsyncJobManager({ maxRunningJobs: 2, retentionMs: 1_000, onJobComplete: async () => { deliveryStarted = true; await deliveryGate.promise; await session.sendCustomMessage( { customType: "async-result", content: "Background result", display: true, attribution: "agent", }, { deliverAs: "followUp", triggerTurn: true }, ); }, }); AsyncJobManager.setInstance(asyncJobManager); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, agentId: ownerId, ownedAsyncJobManager: asyncJobManager, }); session.setClientBridge({ capabilities: {}, deferAgentInitiatedTurns: true, }); await session.prompt("First message"); expect(session.isStreaming).toBe(false); const callsAfterFirstPrompt = mock.calls.length; try { asyncJobManager.register("bash", "owned job", async () => "Background result", { id: "owned-job", ownerId, }); await waitFor(() => deliveryStarted); const drainedPromise = session.drainAsyncJobDeliveriesForAcp({ timeoutMs: 1_000 }); await waitFor(() => asyncJobManager.getDeliveryState({ ownerId }).delivering); deliveryGate.resolve(); await expect(drainedPromise).resolves.toBe(true); await session.waitForIdle(); expect(mock.calls).toHaveLength(callsAfterFirstPrompt + 1); expect( mock.calls.at(-1)?.context.messages.some(message => { if (typeof message.content === "string") { return message.content.includes("Background result"); } return message.content.some( content => content.type === "text" && content.text.includes("Background result"), ); }), ).toBe(true); } finally { deliveryGate.resolve(); } }); it("scopes ACP async job snapshots and drains to the owning session id", async () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-acp-scope.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models-acp-scope.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); const settings = Settings.isolated(); const deliveryGate = Promise.withResolvers(); const delivered: string[] = []; const started = new Set(); const asyncJobManager = new AsyncJobManager({ maxRunningJobs: 3, retentionMs: 1_000, onJobComplete: async jobId => { started.add(jobId); if (jobId === "job-a") { await deliveryGate.promise; } delivered.push(jobId); }, }); AsyncJobManager.setInstance(asyncJobManager); const agentA = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [] }, streamFn: createMockModel({ handler: () => ({ content: ["Done"] }) }).stream, }); const agentB = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [] }, streamFn: createMockModel({ handler: () => ({ content: ["Done"] }) }).stream, }); const sessionB = new AgentSession({ agent: agentB, sessionManager: SessionManager.inMemory(), settings, modelRegistry, agentId: "acp-session-b", asyncJobManager, }); session = new AgentSession({ agent: agentA, sessionManager: SessionManager.inMemory(), settings, modelRegistry, agentId: "acp-session-a", ownedAsyncJobManager: asyncJobManager, }); try { asyncJobManager.register("bash", "A", async () => "A", { id: "job-a", ownerId: "acp-session-a" }); await waitFor(() => started.has("job-a")); asyncJobManager.register("bash", "B", async () => "B", { id: "job-b", ownerId: "acp-session-b" }); await waitFor(() => asyncJobManager.getDeliveryState({ ownerId: "acp-session-b" }).queued > 0); expect(sessionB.getAsyncJobSnapshot()?.delivery.pendingJobIds).not.toContain("job-a"); await expect(sessionB.drainAsyncJobDeliveriesForAcp({ timeoutMs: 1_000 })).resolves.toBe(true); expect(delivered).toEqual(["job-b"]); } finally { deliveryGate.resolve(); await sessionB.dispose(); } }); }); describe("AgentSession TTSR resume gate", () => { let session: AgentSession; let tempDir: string; const authStorages: AuthStorage[] = []; beforeEach(() => { tempDir = path.join(os.tmpdir(), `pi-ttsr-gate-test-${Snowflake.next()}`); fs.mkdirSync(tempDir, { recursive: true }); }); afterEach(async () => { if (session) { await session.dispose(); } for (const authStorage of authStorages.splice(0)) { authStorage.close(); } if (tempDir && fs.existsSync(tempDir)) { removeSyncWithRetries(tempDir); } vi.restoreAllMocks(); }); async function waitFor(predicate: () => boolean, timeoutMs = 500): Promise { const deadline = Date.now() + timeoutMs; while (Date.now() < deadline) { if (predicate()) return; await Bun.sleep(1); } throw new Error("Timed out waiting for condition"); } const testRule: Rule = { name: "no-unwrap", path: "/tmp/no-unwrap.md", content: "Do not use .unwrap()", condition: ["\\.unwrap\\("], _source: { provider: "test", providerName: "test", path: "/tmp/no-unwrap.md", level: "project" }, }; function makeMsg(text: string, stopReason: "stop" | "aborted" = "stop"): AssistantMessage { return { role: "assistant", content: [{ type: "text", text }], api: "anthropic-messages", provider: "anthropic", model: "mock", usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason, timestamp: Date.now(), }; } function pushContinuationStream(stream: AssistantMessageEventStream, onComplete: () => void): void { queueMicrotask(() => { const partial = makeMsg(""); stream.push({ type: "start", partial }); onComplete(); stream.push({ type: "done", reason: "stop", message: makeMsg('Fixed: let val = result.expect("msg")'), }); }); } function pushAbortableTtsrStream(stream: AssistantMessageEventStream, signal: AbortSignal | undefined): void { queueMicrotask(() => { const partial = makeMsg(""); stream.push({ type: "start", partial }); stream.push({ type: "text_delta", contentIndex: 0, delta: "let val = result.unwrap(", partial: makeMsg("let val = result.unwrap("), }); if (signal) { signal.addEventListener( "abort", () => { stream.push({ type: "error", reason: "aborted", error: makeMsg("let val = result.unwrap(", "aborted"), }); }, { once: true }, ); } }); } it("prompt() blocks until TTSR interrupt continuation completes", async () => { collapseSchedulerSettleDelays(); const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; let streamCallCount = 0; let continuationCompleted = false; const ttsrManager = new TtsrManager({ enabled: true, contextMode: "discard", interruptMode: "always", repeatMode: "once", repeatGap: 10, }); ttsrManager.addRule(testRule); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [] }, streamFn: (_model, _context, options) => { streamCallCount++; const stream = new AssistantMessageEventStream(); const signal = options?.signal; if (streamCallCount === 1) { // First stream: emit text that triggers TTSR, then respond to abort pushAbortableTtsrStream(stream, signal); } else { // Continuation stream: complete normally after a delay pushContinuationStream(stream, () => { continuationCompleted = true; }); } return stream; }, }); const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-int.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, ttsrManager, }); // prompt() must block until the TTSR continuation completes await session.prompt("Write some Rust code"); // By the time prompt() returns, the continuation must have finished expect(continuationCompleted).toBe(true); expect(streamCallCount).toBeGreaterThanOrEqual(2); expect(session.isStreaming).toBe(false); }); it("labels aborted tool placeholders with the TTSR rule reason", async () => { collapseSchedulerSettleDelays(); const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; let streamCallCount = 0; const ttsrManager = new TtsrManager({ enabled: true, contextMode: "discard", interruptMode: "always", repeatMode: "once", repeatGap: 10, }); ttsrManager.addRule(testRule); const toolCallContent: ToolCall = { type: "toolCall", id: "call_ttsr_abort_reason", name: "mock_edit", arguments: { snippet: "let val = result.unwrap(" }, }; const makeToolCallMsg = (stopReason: "toolUse" | "aborted" = "toolUse"): AssistantMessage => ({ role: "assistant", content: [toolCallContent], api: "anthropic-messages", provider: "anthropic", model: "mock", usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason, timestamp: Date.now(), }); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [] }, streamFn: (_model, _context, options) => { streamCallCount++; const stream = new AssistantMessageEventStream(); const signal = options?.signal; if (streamCallCount === 1) { queueMicrotask(() => { const partial = makeToolCallMsg(); if (signal) { signal.addEventListener( "abort", () => { stream.push({ type: "error", reason: "aborted", error: makeToolCallMsg("aborted"), }); }, { once: true }, ); } stream.push({ type: "start", partial }); stream.push({ type: "toolcall_start", contentIndex: 0, partial }); stream.push({ type: "toolcall_delta", contentIndex: 0, delta: 'let val = result.unwrap("oops")', partial, }); // The TTSR abort placeholder is only minted for tool calls that reached // `toolcall_end`: the agent loop drops incomplete tool calls from an // aborted turn (partial args are unsafe to replay). Complete the call // before the rule-driven abort fires so the labeled placeholder survives. stream.push({ type: "toolcall_end", contentIndex: 0, toolCall: toolCallContent, partial }); }); } else { pushContinuationStream(stream, () => {}); } return stream; }, }); const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-abort-reason.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, ttsrManager }); await session.prompt("Write some Rust code"); const toolResult = sessionManager .getEntries() .find( entry => entry.type === "message" && entry.message.role === "toolResult" && entry.message.toolCallId === toolCallContent.id, ); expect(toolResult?.type).toBe("message"); const text = toolResult?.type === "message" && toolResult.message.role === "toolResult" ? (toolResult.message.content.find((part): part is { type: "text"; text: string } => part.type === "text") ?.text ?? "") : ""; expect(text).toContain("Tool execution was aborted: TTSR matched rule: no-unwrap"); expect(text).not.toContain("Request was aborted"); }); it("labels only the matching aborted tool placeholder with the TTSR rule reason", async () => { collapseSchedulerSettleDelays(); const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; let streamCallCount = 0; const ttsrManager = new TtsrManager({ enabled: true, contextMode: "discard", interruptMode: "always", repeatMode: "once", repeatGap: 10, }); ttsrManager.addRule(testRule); const readToolCallContent: ToolCall = { type: "toolCall", id: "call_innocent_read", name: "read", arguments: { path: "history://Eval1WithSkill" }, }; const matchedToolCallContent: ToolCall = { type: "toolCall", id: "call_ttsr_abort_reason", name: "mock_edit", arguments: { snippet: "let val = result.unwrap(" }, }; const makeToolCallMsg = (stopReason: "toolUse" | "aborted" = "toolUse"): AssistantMessage => ({ role: "assistant", content: [readToolCallContent, matchedToolCallContent], api: "anthropic-messages", provider: "anthropic", model: "mock", usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason, timestamp: Date.now(), }); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [] }, streamFn: (_model, _context, options) => { streamCallCount++; const stream = new AssistantMessageEventStream(); const signal = options?.signal; if (streamCallCount === 1) { queueMicrotask(() => { const partial = makeToolCallMsg(); if (signal) { signal.addEventListener( "abort", () => { stream.push({ type: "error", reason: "aborted", error: makeToolCallMsg("aborted"), }); }, { once: true }, ); } stream.push({ type: "start", partial }); stream.push({ type: "toolcall_start", contentIndex: 1, partial }); stream.push({ type: "toolcall_delta", contentIndex: 1, delta: 'let val = result.unwrap("oops")', partial, }); // The abort placeholder is only minted for tool calls that reached // `toolcall_end`: the agent loop drops incomplete tool calls from an // aborted turn (partial args are unsafe to replay). Complete the // innocent read before the rule-driven abort fires so its placeholder // survives and can carry the neutral sibling label. stream.push({ type: "toolcall_end", contentIndex: 0, toolCall: readToolCallContent, partial }); }); } else { pushContinuationStream(stream, () => {}); } return stream; }, }); const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-abort-reason.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, ttsrManager }); await session.prompt("Write some Rust code"); const toolResults = sessionManager .getEntries() .filter(entry => entry.type === "message" && entry.message.role === "toolResult") .map(entry => (entry.type === "message" && entry.message.role === "toolResult" ? entry.message : undefined)) .filter(message => message !== undefined); const toolResultText = (toolCallId: string): string => toolResults .find(message => message.toolCallId === toolCallId) ?.content.find((part): part is { type: "text"; text: string } => part.type === "text")?.text ?? ""; const readText = toolResultText(readToolCallContent.id); expect(readText).toContain("Tool execution was aborted: TTSR interrupt on another tool call"); expect(readText).not.toContain("TTSR matched rule: no-unwrap"); // The matching call never reached `toolcall_end`, so the loop drops it from // the aborted turn (partial args are unsafe to replay) and no placeholder is // minted. The rule label for a completed matching call is covered by the // single-call test above. expect(toolResultText(matchedToolCallContent.id)).toBe(""); }); it("relativizes the rule file path in the TTSR interrupt injection (no absolute leak)", async () => { collapseSchedulerSettleDelays(); const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; let streamCallCount = 0; const sessionManager = SessionManager.inMemory(); const cwd = sessionManager.getCwd(); const ruleAbsPath = path.join(cwd, ".omp", "rules", "no-unwrap.md"); const expectedRel = path.relative(cwd, ruleAbsPath); const rule: Rule = { name: "no-unwrap", path: ruleAbsPath, content: "Do not use .unwrap()", condition: ["\\.unwrap\\("], _source: { provider: "test", providerName: "test", path: ruleAbsPath, level: "project" }, }; const ttsrManager = new TtsrManager({ enabled: true, contextMode: "discard", interruptMode: "always", repeatMode: "once", repeatGap: 10, }); ttsrManager.addRule(rule); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [] }, streamFn: (_model, _context, options) => { streamCallCount++; const stream = new AssistantMessageEventStream(); if (streamCallCount === 1) { pushAbortableTtsrStream(stream, options?.signal); } else { pushContinuationStream(stream, () => {}); } return stream; }, }); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-rel.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, ttsrManager }); await session.prompt("Write some Rust code"); const injection = sessionManager .getEntries() .find(e => e.type === "custom_message" && e.customType === "ttsr-injection"); expect(injection?.type).toBe("custom_message"); const content = injection?.type === "custom_message" ? injection.content : undefined; expect(typeof content).toBe("string"); const text = content as string; // The rendered interrupt the model receives references the rule by a // project-relative path, never the absolute home path. expect(text).toContain('reason="rule_violation"'); expect(text).toContain(`path="${expectedRel}"`); expect(text).not.toContain(ruleAbsPath); }); it("prompt() blocks until TTSR deferred continuation completes", async () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; let streamCallCount = 0; let continuationCompleted = false; // interruptMode: "never" -> TTSR match queues deferred injection instead of aborting const ttsrManager = new TtsrManager({ enabled: true, contextMode: "discard", interruptMode: "never", repeatMode: "once", repeatGap: 10, }); ttsrManager.addRule(testRule); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [] }, streamFn: (_model, _context, _options) => { streamCallCount++; const stream = new AssistantMessageEventStream(); if (streamCallCount === 1) { // First stream: emit matching text and complete normally queueMicrotask(() => { const partial = makeMsg(""); stream.push({ type: "start", partial }); stream.push({ type: "text_delta", contentIndex: 0, delta: "let val = result.unwrap(", partial: makeMsg("let val = result.unwrap("), }); // Complete normally (no abort) -- deferred path stream.push({ type: "done", reason: "stop", message: makeMsg("let val = result.unwrap()"), }); }); } else { // Continuation stream after deferred TTSR injection pushContinuationStream(stream, () => { continuationCompleted = true; }); } return stream; }, }); const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-def.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, ttsrManager, }); // prompt() must block until the deferred TTSR continuation completes await session.prompt("Write some Rust code"); // By the time prompt() returns, the deferred continuation must have finished expect(continuationCompleted).toBe(true); expect(streamCallCount).toBeGreaterThanOrEqual(2); expect(session.isStreaming).toBe(false); }); it("prompt() returns immediately when session is aborted during TTSR wait", async () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; const ttsrManager = new TtsrManager({ enabled: true, contextMode: "discard", interruptMode: "always", repeatMode: "once", repeatGap: 10, }); ttsrManager.addRule(testRule); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [] }, streamFn: (_model, _context, options) => { const stream = new AssistantMessageEventStream(); const signal = options?.signal; queueMicrotask(() => { const partial = makeMsg(""); stream.push({ type: "start", partial }); stream.push({ type: "text_delta", contentIndex: 0, delta: "result.unwrap(", partial: makeMsg("result.unwrap("), }); if (signal) { signal.addEventListener( "abort", () => { stream.push({ type: "error", reason: "aborted", error: makeMsg("result.unwrap(", "aborted"), }); }, { once: true }, ); } }); return stream; }, }); const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-abt.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, ttsrManager, }); // Start prompt (will trigger TTSR and create resume gate) const promptPromise = session.prompt("Write some Rust code"); await waitFor(() => session.isStreaming); // Abort session — prompt() should unblock await session.abort(); await promptPromise; expect(session.isStreaming).toBe(false); }); it("prompt() waits for TTSR continuation with tool calls to finish", async () => { collapseSchedulerSettleDelays(); const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; let streamCallCount = 0; let toolExecutionFinished = false; let allTurnsCompleted = false; const ttsrManager = new TtsrManager({ enabled: true, contextMode: "discard", interruptMode: "always", repeatMode: "once", repeatGap: 10, }); ttsrManager.addRule(testRule); const mockTool: AgentTool = { name: "mock_edit", label: "Mock Edit", description: "A mock edit tool", parameters: type({}), execute: async () => { toolExecutionFinished = true; return { content: [{ type: "text" as const, text: "edit applied" }] }; }, }; const toolCallContent: ToolCall = { type: "toolCall", id: "call_test_001", name: "mock_edit", arguments: {}, }; function makeToolCallMsg(): AssistantMessage { return { role: "assistant", content: [toolCallContent], api: "anthropic-messages", provider: "anthropic", model: "mock", usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "toolUse", timestamp: Date.now(), }; } const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [mockTool] }, streamFn: (_model, _context, options) => { streamCallCount++; const stream = new AssistantMessageEventStream(); const signal = options?.signal; if (streamCallCount === 1) { // First stream: emit text that triggers TTSR, then respond to abort pushAbortableTtsrStream(stream, signal); } else if (streamCallCount === 2) { // Continuation: return assistant message with a tool call queueMicrotask(() => { const msg = makeToolCallMsg(); stream.push({ type: "start", partial: msg }); stream.push({ type: "done", reason: "toolUse", message: msg }); }); } else { // After tool execution: return final response queueMicrotask(() => { allTurnsCompleted = true; const msg = makeMsg('Fixed: let val = result.expect("msg")'); stream.push({ type: "start", partial: msg }); stream.push({ type: "done", reason: "stop", message: msg }); }); } return stream; }, }); const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-tool.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, ttsrManager, }); // prompt() must block until the TTSR continuation (including tool execution) completes. // Before the fix, prompt() returned after the continuation's first assistant message_end, // while the agent was still executing tool calls in the background. await session.prompt("Write some Rust code"); // By the time prompt() returns, ALL turns must have completed expect(toolExecutionFinished).toBe(true); expect(allTurnsCompleted).toBe(true); expect(streamCallCount).toBeGreaterThanOrEqual(3); expect(session.isStreaming).toBe(false); }); it("interruptMode never folds tool-match reminder into the toolResult instead of driving an extra turn", async () => { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; let streamCallCount = 0; let toolExecuted = false; const ttsrManager = new TtsrManager({ enabled: true, contextMode: "discard", interruptMode: "never", repeatMode: "once", repeatGap: 10, }); ttsrManager.addRule(testRule); const mockTool: AgentTool = { name: "mock_edit", label: "Mock Edit", description: "A mock edit tool", parameters: type({ snippet: "string?" }), execute: async () => { toolExecuted = true; return { content: [{ type: "text" as const, text: "edit applied" }] }; }, }; const toolCallContent: ToolCall = { type: "toolCall", id: "call_never_001", name: "mock_edit", arguments: { snippet: "let val = result.unwrap()" }, }; const makeToolCallMsg = (): AssistantMessage => ({ role: "assistant", content: [toolCallContent], api: "anthropic-messages", provider: "anthropic", model: "mock", usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "toolUse", timestamp: Date.now(), }); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [mockTool] }, streamFn: () => { streamCallCount++; const stream = new AssistantMessageEventStream(); if (streamCallCount === 1) { // Emit a tool call whose argument delta matches the TTSR rule. queueMicrotask(() => { const partial = makeToolCallMsg(); stream.push({ type: "start", partial }); stream.push({ type: "toolcall_start", contentIndex: 0, partial }); stream.push({ type: "toolcall_delta", contentIndex: 0, delta: 'let val = result.unwrap("oops")', partial, }); stream.push({ type: "toolcall_end", contentIndex: 0, toolCall: toolCallContent, partial }); stream.push({ type: "done", reason: "toolUse", message: partial }); }); } else { // Continuation after tool result; finish cleanly. queueMicrotask(() => { const done = makeMsg("ok"); stream.push({ type: "start", partial: done }); stream.push({ type: "done", reason: "stop", message: done }); }); } return stream; }, }); const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-never-tool.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, ttsrManager, }); await session.prompt("Write some Rust code"); // Tool ran (no interrupt) and the loop didn't spawn an extra follow-up turn for injection. expect(toolExecuted).toBe(true); expect(streamCallCount).toBe(2); // The matched tool's result must carry the in-band reminder. const toolResult = agent.state.messages.find( (m): m is Extract => m.role === "toolResult" && m.toolCallId === toolCallContent.id, ); expect(toolResult).toBeDefined(); const text = Array.isArray(toolResult?.content) ? toolResult.content .filter((c): c is { type: "text"; text: string } => c.type === "text") .map(c => c.text) .join("\n") : ""; expect(text).toContain(" { const model = getBundledModel("anthropic", "claude-sonnet-4-5")!; let streamCallCount = 0; let executedCount = 0; const ttsrManager = new TtsrManager({ enabled: true, contextMode: "discard", interruptMode: "never", repeatMode: "once", repeatGap: 10, }); ttsrManager.addRule(testRule); const mockTool: AgentTool = { name: "mock_edit", label: "Mock Edit", description: "A mock edit tool", parameters: type({ snippet: "string?" }), execute: async () => { executedCount++; return { content: [{ type: "text" as const, text: "edit applied" }] }; }, }; const toolCallA: ToolCall = { type: "toolCall", id: "call_dup_A", name: "mock_edit", arguments: { snippet: "a.unwrap()" }, }; const toolCallB: ToolCall = { type: "toolCall", id: "call_dup_B", name: "mock_edit", arguments: { snippet: "b.unwrap()" }, }; const toolCallC: ToolCall = { type: "toolCall", id: "call_dup_C", name: "mock_edit", arguments: { snippet: "c.unwrap()" }, }; const makeBatchMsg = (): AssistantMessage => ({ role: "assistant", content: [toolCallA, toolCallB, toolCallC], api: "anthropic-messages", provider: "anthropic", model: "mock", usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "toolUse", timestamp: Date.now(), }); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: ["Test"], tools: [mockTool] }, streamFn: () => { streamCallCount++; const stream = new AssistantMessageEventStream(); if (streamCallCount === 1) { queueMicrotask(() => { const partial = makeBatchMsg(); stream.push({ type: "start", partial }); const calls: ToolCall[] = [toolCallA, toolCallB, toolCallC]; for (let i = 0; i < calls.length; i++) { const call = calls[i]!; stream.push({ type: "toolcall_start", contentIndex: i, partial }); stream.push({ type: "toolcall_delta", contentIndex: i, delta: `let val = result.unwrap("oops-${call.id}")`, partial, }); stream.push({ type: "toolcall_end", contentIndex: i, toolCall: call, partial }); } stream.push({ type: "done", reason: "toolUse", message: partial }); }); } else { queueMicrotask(() => { const done = makeMsg("ok"); stream.push({ type: "start", partial: done }); stream.push({ type: "done", reason: "stop", message: done }); }); } return stream; }, }); const sessionManager = SessionManager.inMemory(); const settings = Settings.isolated(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-dup.db")); authStorages.push(authStorage); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); authStorage.setRuntimeApiKey("anthropic", "test-key"); session = new AgentSession({ agent, sessionManager, settings, modelRegistry, ttsrManager, }); await session.prompt("Write some Rust code"); expect(executedCount).toBe(3); const toolResults = agent.state.messages.filter( (m): m is Extract => m.role === "toolResult", ); expect(toolResults).toHaveLength(3); const withReminder = toolResults.filter(r => Array.isArray(r.content) ? r.content.some(c => c.type === "text" && c.text.includes(" { collapseSchedulerSettleDelays(); const authStorage = await AuthStorage.create(path.join(tempDir, "testauth-promo.db")); authStorages.push(authStorage); authStorage.setRuntimeApiKey("openai-codex", "test-key"); const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml")); const sparkModel = modelRegistry.find("openai-codex", "gpt-5.3-codex-spark"); const codexModel = modelRegistry.find("openai-codex", "gpt-5.5"); if (!sparkModel || !codexModel) { throw new Error("Expected codex spark and codex models to exist"); } let streamCallCount = 0; let continuationCompleted = false; const makeOverflowMessage = (): AssistantMessage => ({ role: "assistant", content: [{ type: "text", text: "" }], api: sparkModel.api, provider: sparkModel.provider, model: sparkModel.id, usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "error", errorMessage: "context_length_exceeded: Your input exceeds the context window of this model.", timestamp: Date.now(), }); const makeSuccessMessage = (): AssistantMessage => ({ role: "assistant", content: [{ type: "text", text: "Recovered after promotion" }], api: codexModel.api, provider: codexModel.provider, model: codexModel.id, usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "stop", timestamp: Date.now(), }); const agent = new Agent({ getApiKey: () => "test-key", initialState: { model: sparkModel, systemPrompt: ["Test"], tools: [] }, streamFn: () => { streamCallCount++; const stream = new AssistantMessageEventStream(); if (streamCallCount === 1) { queueMicrotask(() => { const message = makeOverflowMessage(); stream.push({ type: "start", partial: message }); stream.push({ type: "error", reason: "error", error: message }); }); } else { queueMicrotask(() => { continuationCompleted = true; const message = makeSuccessMessage(); stream.push({ type: "start", partial: message }); stream.push({ type: "done", reason: "stop", message }); }); } return stream; }, }); const extensionRunner = { emit: vi.fn().mockResolvedValue(undefined), emitBeforeAgentStart: vi.fn().mockResolvedValue(undefined), hasHandlers: vi.fn((eventType: string) => eventType === "session_stop"), emitSessionStop: vi.fn().mockResolvedValue(undefined), } as unknown as ExtensionRunner; session = new AgentSession({ agent, sessionManager: SessionManager.inMemory(), settings: Settings.isolated({ "compaction.enabled": false, "contextPromotion.enabled": true }), modelRegistry, extensionRunner, }); await session.prompt("Handle overflow"); expect(continuationCompleted).toBe(true); expect(streamCallCount).toBeGreaterThanOrEqual(2); expect(session.model?.id).toBe(codexModel.id); expect(session.isStreaming).toBe(false); expect(extensionRunner.emitSessionStop).toHaveBeenCalledTimes(1); }); });