Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
36 kB · 1111 lines
TypeScript
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112import { existsSync, mkdirSync, rmSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import { Agent } from "@earendil-works/pi-agent-core";import { type AssistantMessage, type AssistantMessageEvent, EventStream, getModel, type ImageContent, type TextContent,} from "@earendil-works/pi-ai";import { Type } from "typebox";import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";import { createAgentSessionMessage } from "../src/core/agent-messages.js";import { AgentSession } from "../src/core/agent-session.js";import { AuthStorage } from "../src/core/auth-storage.js";import type { CompactionResult } from "../src/core/compaction/index.js";import type { AgentCronJob } from "../src/core/cron-jobs.js";import { ModelRegistry } from "../src/core/model-registry.js";import { SessionManager } from "../src/core/session-manager.js";import { SettingsManager } from "../src/core/settings-manager.js";import type { BuildSystemPromptOptions } from "../src/core/system-prompt.js";import { createTestExtensionsResult, createTestResourceLoader } from "./utilities.js";
class MockAssistantStream extends EventStream<AssistantMessageEvent, AssistantMessage> { constructor() { super( (event) => event.type === "done" || event.type === "error", (event) => { if (event.type === "done") return event.message; if (event.type === "error") return event.error; throw new Error("Unexpected event type"); }, ); }}
function createAssistantMessage(text: string): 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: "stop", timestamp: Date.now(), };}
describe("AgentSession concurrent prompt guard", () => { let session: AgentSession; let tempDir: string;
beforeEach(() => { tempDir = join(tmpdir(), `pi-concurrent-test-${Date.now()}`); mkdirSync(tempDir, { recursive: true }); });
afterEach(async () => { delete (globalThis as typeof globalThis & { testExtensionApi?: unknown }).testExtensionApi; delete (globalThis as typeof globalThis & { testCommandRuns?: unknown }).testCommandRuns; if (session) { session.dispose(); } if (tempDir && existsSync(tempDir)) { rmSync(tempDir, { recursive: true }); } });
function createSession() { const model = getModel("anthropic", "claude-sonnet-4-5")!; let abortSignal: AbortSignal | undefined;
const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: "Test", tools: [], }, streamFn: (_model, _context, options) => { abortSignal = options?.signal; const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ type: "start", partial: createAssistantMessage("") }); const checkAbort = () => { if (abortSignal?.aborted) { stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") }); } else { setTimeout(checkAbort, 5); } }; checkAbort(); }); return stream; }, });
const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); const modelRegistry = ModelRegistry.create(authStorage, tempDir); authStorage.setRuntimeApiKey("anthropic", "test-key");
session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, modelRegistry, resourceLoader: createTestResourceLoader(), });
return session; }
it("awaits asynchronous dispose callbacks during graceful disposal", async () => { createSession(); let releaseCallback: () => void = () => {}; const callbackGate = new Promise<void>((resolve) => { releaseCallback = resolve; }); let callbackStarted = false; session.registerDisposeCallback(async () => { callbackStarted = true; await callbackGate; });
let disposed = false; const disposal = session.disposeAsync().then(() => { disposed = true; }); await new Promise<void>((resolve) => setTimeout(resolve, 0)); expect(callbackStarted).toBe(true); expect(disposed).toBe(false); releaseCallback(); await disposal; expect(disposed).toBe(true); });
it("awaits dispose callbacks when synchronous disposal wins during the refinement drain", async () => { createSession(); let releaseDrain: () => void = () => {}; const drainGate = new Promise<void>((resolve) => { releaseDrain = resolve; }); let drainStarted = false; const internals = session as unknown as { _drainPendingRefinementForDisposal: () => Promise<void>; }; vi.spyOn(internals, "_drainPendingRefinementForDisposal").mockImplementation(async () => { drainStarted = true; await drainGate; });
let releaseCallback: () => void = () => {}; const callbackGate = new Promise<void>((resolve) => { releaseCallback = resolve; }); let callbackStarted = false; session.registerDisposeCallback(async () => { callbackStarted = true; await callbackGate; });
let disposed = false; const disposal = session.disposeAsync().then(() => { disposed = true; }); await vi.waitFor(() => expect(drainStarted).toBe(true)); session.dispose(); expect(callbackStarted).toBe(true);
releaseDrain(); await new Promise<void>((resolve) => setTimeout(resolve, 0)); expect(disposed).toBe(false); releaseCallback(); await disposal; expect(disposed).toBe(true); });
it("forwards kernelSnapshot: false to the kernel provisioner during disposal", async () => { createSession(); const dispose = vi.fn(async () => {}); Reflect.set(session, "_ipythonKernelProvisioner", { dispose });
await session.disposeAsync({ kernelSnapshot: false }); expect(dispose).toHaveBeenCalledWith({ snapshot: false }); });
it("should throw when prompt() called while streaming", async () => { createSession();
const firstPrompt = session.prompt("First message"); await new Promise((resolve) => setTimeout(resolve, 10)); expect(session.isStreaming).toBe(true);
await expect(session.prompt("Second message")).rejects.toThrow( "Agent is already processing. Specify streamingBehavior ('steer' or 'followUp') to queue the message.", );
await session.abort(); await firstPrompt.catch(() => {}); });
it("should allow steer() while streaming", async () => { createSession();
const firstPrompt = session.prompt("First message"); await new Promise((resolve) => setTimeout(resolve, 10));
expect(() => session.steer("Steering message")).not.toThrow(); expect(session.queuedActionCount).toBe(1);
await session.abort(); await firstPrompt.catch(() => {}); });
it("should allow followUp() while streaming", async () => { createSession();
const firstPrompt = session.prompt("First message"); await new Promise((resolve) => setTimeout(resolve, 10));
expect(() => session.followUp("Follow-up message")).not.toThrow(); expect(session.queuedActionCount).toBe(1);
await session.abort(); await firstPrompt.catch(() => {}); });
it("should queue extension-origin steering messages while streaming", async () => { const model = getModel("anthropic", "claude-sonnet-4-5")!; let abortSignal: AbortSignal | undefined; let sawSteeringMessage = false; let lastInputSource: string | undefined; const queueEvents: Array<{ steering: readonly string[]; followUp: readonly string[] }> = [];
const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: "Test", tools: [], }, streamFn: (_model, context, options) => { abortSignal = options?.signal; const stream = new MockAssistantStream(); queueMicrotask(() => { const userTexts = context.messages .filter((message) => message.role === "user") .map((message) => { if (typeof message.content === "string") { return message.content; } return message.content .filter((part): part is TextContent | ImageContent => typeof part === "object" && part !== null) .filter((part): part is TextContent => part.type === "text") .map((part) => part.text) .join("\n"); });
if (userTexts.includes("Steer from extension")) { sawSteeringMessage = true; stream.push({ type: "start", partial: createAssistantMessage("") }); stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Steered") }); return; }
stream.push({ type: "start", partial: createAssistantMessage("") }); const checkAbort = () => { if (abortSignal?.aborted) { stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") }); } else { setTimeout(checkAbort, 5); } }; checkAbort(); }); return stream; }, });
const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); const modelRegistry = ModelRegistry.create(authStorage, tempDir); authStorage.setRuntimeApiKey("anthropic", "test-key");
const extensionsResult = await createTestExtensionsResult([ (pi) => { (globalThis as typeof globalThis & { testExtensionApi?: unknown }).testExtensionApi = pi; }, (pi) => { pi.on("input", async (event) => { lastInputSource = event.source; }); }, ]);
session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, modelRegistry, resourceLoader: createTestResourceLoader({ extensionsResult }), }); session.subscribe((event) => { if (event.type === "session_action_update") { queueEvents.push({ steering: event.actions.steering, followUp: event.actions.followUps }); } });
const firstPrompt = session.prompt("First message"); await new Promise((resolve) => setTimeout(resolve, 10)); expect(session.isStreaming).toBe(true);
const pi = ( globalThis as typeof globalThis & { testExtensionApi?: { sendUserMessage: (content: string, options?: { deliverAs?: "steer" | "followUp" }) => void; }; } ).testExtensionApi; expect(pi).toBeDefined();
pi!.sendUserMessage("Steer from extension", { deliverAs: "steer" }); await new Promise((resolve) => setTimeout(resolve, 25));
expect(session.queuedActionCount).toBe(1); expect(session.getSteeringMessages()).toContain("Steer from extension"); expect(lastInputSource).toBe("extension"); expect(queueEvents.some((event) => event.steering.includes("Steer from extension"))).toBe(true);
await session.abort(); await firstPrompt.catch(() => {});
expect(sawSteeringMessage).toBe(false); expect(session.getSteeringMessages()).toContain("Steer from extension"); expect(session.queuedActionCount).toBe(1);
await session.prompt("After abort");
expect(sawSteeringMessage).toBe(true); expect(session.queuedActionCount).toBe(0); });
it("delivers accepted agent messages without extension input interception", async () => { const model = getModel("anthropic", "claude-sonnet-4-5")!; let inputCalls = 0; let receivedUserText: string | undefined; const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: "Test", tools: [], }, streamFn: (_model, context) => { const stream = new MockAssistantStream(); queueMicrotask(() => { const userMessages = context.messages.filter((message) => message.role === "user"); const user = userMessages.at(-1); if (user && typeof user.content !== "string") { receivedUserText = user.content .filter( (part): part is TextContent => typeof part === "object" && part !== null && part.type === "text", ) .map((part) => part.text) .join("\n"); } stream.push({ type: "start", partial: createAssistantMessage("") }); stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Delivered") }); }); return stream; }, });
const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); const modelRegistry = ModelRegistry.create(authStorage, tempDir); authStorage.setRuntimeApiKey("anthropic", "test-key"); const extensionsResult = await createTestExtensionsResult([ (pi) => { pi.on("input", async () => { inputCalls++; return { action: "handled" }; }); }, ]);
session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, modelRegistry, resourceLoader: createTestResourceLoader({ extensionsResult }), });
await session.acceptAgentMessagePrompt("agent-to-agent payload", { expandPromptTemplates: false }); await session.agent.waitForIdle();
expect(inputCalls).toBe(0); expect(receivedUserText).toBe("agent-to-agent payload"); });
it("queues an agent message behind an active turn", async () => { createSession(); const activeTurn = session.prompt("active turn"); await vi.waitFor(() => expect(session.isStreaming).toBe(true)); let queued: boolean | undefined;
await session.acceptAgentMessagePrompt("message during turn", { queueIfBusy: true, streamingBehavior: "steer", preflightResult: (accepted, didQueue) => { expect(accepted).toBe(true); queued = didQueue; }, });
expect(queued).toBe(true); expect(session.getSessionActionSnapshot().steering).toContain("message during turn"); await session.abort(); await activeTurn.catch(() => undefined); });
it("serializes an agent message behind cron admission", async () => { createSession(); const now = new Date().toISOString(); const job: AgentCronJob = { id: "cron-overlap", status: "active", source: "cron", activeSessionId: "active-1", sessionId: session.sessionId, sessionFile: "/tmp/session.jsonl", cwd: tempDir, prompt: "cron prompt", schedule: { kind: "interval", expression: "every 5m", intervalMs: 300_000 }, createdAt: now, updatedAt: now, runCount: 0, }; const cronTurn = session.promptHeartbeat(job, { streamingBehavior: "followUp", source: "rpc" }); let queued: boolean | undefined;
await session.acceptAgentMessagePrompt("message during cron admission", { queueIfBusy: true, streamingBehavior: "steer", preflightResult: (accepted, didQueue) => { expect(accepted).toBe(true); queued = didQueue; }, });
expect(queued).toBe(true); expect(session.getSessionActionSnapshot().steering).toContain("message during cron admission"); await session.abort(); await cronTurn.catch(() => undefined); });
it("does not admit a cron prompt invalidated while async input handlers ran", async () => { const model = getModel("anthropic", "claude-sonnet-4-5")!; let deliveredUserText: string | undefined; const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: "Test", tools: [], }, streamFn: (_model, context) => { const stream = new MockAssistantStream(); queueMicrotask(() => { const user = context.messages.filter((message) => message.role === "user").at(-1); if (user && typeof user.content !== "string") { deliveredUserText = user.content .filter( (part): part is TextContent => typeof part === "object" && part !== null && part.type === "text", ) .map((part) => part.text) .join("\n"); } stream.push({ type: "start", partial: createAssistantMessage("") }); stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Delivered") }); }); return stream; }, });
const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); const modelRegistry = ModelRegistry.create(authStorage, tempDir); authStorage.setRuntimeApiKey("anthropic", "test-key"); let jobInvalidated = false; const extensionsResult = await createTestExtensionsResult([ (pi) => { pi.on("input", async () => { jobInvalidated = true; return undefined; }); }, ]);
session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, modelRegistry, resourceLoader: createTestResourceLoader({ extensionsResult }), });
// Mirrors the daemon cron path: admissionCommitted re-checks the job and // throws when it was cancelled or updated before the prompt is admitted. await expect( session.promptUntilAccepted("stale cron prompt", { streamingBehavior: "followUp", source: "rpc", admissionCommitted: () => { if (jobInvalidated) throw new Error("Cron job became unrunnable before admission"); }, }), ).rejects.toThrow("Cron job became unrunnable before admission");
await session.agent.waitForIdle(); expect(deliveredUserText).toBeUndefined(); expect(session.isStreaming).toBe(false); });
it("serializes concurrent agent messages", async () => { createSession(); const dispositions = new Map<string, boolean | undefined>(); const accept = (message: string) => session.acceptAgentMessagePrompt(message, { queueIfBusy: true, streamingBehavior: "steer", preflightResult: (accepted, queued) => { expect(accepted).toBe(true); dispositions.set(message, queued); }, });
await Promise.all([accept("first concurrent message"), accept("second concurrent message")]);
expect(dispositions.get("first concurrent message")).toBe(false); expect(dispositions.get("second concurrent message")).toBe(true); expect(session.getSessionActionSnapshot().steering).toContain("second concurrent message"); await session.abort(); });
it("enforces the agent message queue cap inside core admission", async () => { createSession(); const accept = (index: number) => { const message = createAgentSessionMessage({ id: `agentmsg-${index}`, source: "agent_message", message: `message ${index}`, from: { clientId: "test" }, target: { activeSessionId: "target", sessionId: session.sessionId, runtimeKind: "subagent", }, }); return session.queueAgentMessagePrompt(message.content, "steer", message); };
for (let index = 0; index < 20; index++) { await accept(index); } await expect(accept(20)).rejects.toThrow("Target session has too many pending messages"); expect(session.unfinishedActionCount).toBe(20); const snapshot = session.getSessionActionRecoverySnapshot(); session.dispose();
createSession(); await expect(session.restoreSessionActions(snapshot)).resolves.toBe(20); expect(session.unfinishedActionCount).toBe(20); });
it("rejects an agent message when clear wins core admission", async () => { createSession(); const internals = session as unknown as { _acquireSessionActionCommitFence(): Promise<{ release(): void }>; }; const fence = await internals._acquireSessionActionCommitFence(); const message = createAgentSessionMessage({ id: "agentmsg-clear-race", source: "agent_message", message: "clear race", from: { clientId: "test" }, target: { activeSessionId: "target", sessionId: session.sessionId, runtimeKind: "subagent", }, }); const accepted = session.acceptAgentMessagePrompt(message.content, { customMessage: message, queueIfBusy: true, streamingBehavior: "steer", }); await Promise.resolve();
session.clearQueuedAgentMessages(); fence.release();
await expect(accepted).rejects.toThrow("Agent message was cleared before admission"); expect(session.unfinishedActionCount).toBe(0); });
it("accepts agent messages during and after direct compaction", async () => { createSession(); let releaseCompaction: (result: CompactionResult) => void = () => {}; const compactionGate = new Promise<CompactionResult>((resolve) => { releaseCompaction = resolve; }); const internals = session as unknown as { _performCompaction: () => Promise<CompactionResult>; }; vi.spyOn(internals, "_performCompaction").mockImplementation(() => compactionGate); const compaction = session.compact(); await vi.waitFor(() => expect(session.isCompacting).toBe(true)); let queued: boolean | undefined;
await session.acceptAgentMessagePrompt("message during direct compaction", { queueIfBusy: true, streamingBehavior: "steer", preflightResult: (accepted, didQueue) => { expect(accepted).toBe(true); queued = didQueue; }, });
expect(queued).toBe(true); expect(session.getSessionActionSnapshot().steering).toContain("message during direct compaction"); releaseCompaction({ summary: "compacted", firstKeptEntryId: "entry-1", tokensBefore: 1 }); await compaction; let postCompactionQueued: boolean | undefined; await session.acceptAgentMessagePrompt("message after direct compaction", { queueIfBusy: true, streamingBehavior: "steer", preflightResult: (accepted, didQueue) => { expect(accepted).toBe(true); postCompactionQueued = didQueue; }, }); expect(postCompactionQueued).toBe(true); const deliveryCount = (text: string) => session.messages.filter( (message) => message.role === "user" && typeof message.content !== "string" && message.content.some((part) => part.type === "text" && part.text === text), ).length; expect(deliveryCount("message during direct compaction")).toBe(0); expect(deliveryCount("message after direct compaction")).toBe(0);
session.resumeQueuedWork(); await vi.waitFor(() => expect(deliveryCount("message during direct compaction")).toBe(1)); expect(session.getSessionActionSnapshot().steering).toContain("message after direct compaction"); });
it("delivers an accepted agent message once after scheduler-owned compaction", async () => { createSession(); let releaseCompaction: (result: CompactionResult) => void = () => {}; const compactionGate = new Promise<CompactionResult>((resolve) => { releaseCompaction = resolve; }); const internals = session as unknown as { _performCompaction: () => Promise<CompactionResult>; }; vi.spyOn(internals, "_performCompaction").mockImplementation(() => compactionGate); const compaction = session.compact(undefined, { skipAbort: true }); await vi.waitFor(() => expect(session.isCompacting).toBe(true)); let queued: boolean | undefined;
await session.acceptAgentMessagePrompt("message during scheduled compaction", { queueIfBusy: true, streamingBehavior: "steer", preflightResult: (accepted, didQueue) => { expect(accepted).toBe(true); queued = didQueue; }, });
expect(queued).toBe(true); expect(session.getSessionActionSnapshot().steering).toContain("message during scheduled compaction"); releaseCompaction({ summary: "compacted", firstKeptEntryId: "entry-1", tokensBefore: 1 }); await compaction; const deliveryCount = () => session.messages.filter( (message) => message.role === "user" && typeof message.content !== "string" && message.content.some( (part) => part.type === "text" && part.text === "message during scheduled compaction", ), ).length;
await vi.waitFor(() => expect(deliveryCount()).toBe(1)); expect(deliveryCount()).toBe(1); });
it("delivers internal prompts without extension input interception", async () => { const model = getModel("anthropic", "claude-sonnet-4-5")!; let inputCalls = 0; let receivedUserText: string | undefined; const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: "Test", tools: [], }, streamFn: (_model, context) => { const stream = new MockAssistantStream(); queueMicrotask(() => { const userMessages = context.messages.filter((message) => message.role === "user"); const user = userMessages.at(-1); if (user && typeof user.content !== "string") { receivedUserText = user.content .filter( (part): part is TextContent => typeof part === "object" && part !== null && part.type === "text", ) .map((part) => part.text) .join("\n"); } stream.push({ type: "start", partial: createAssistantMessage("") }); stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Delivered") }); }); return stream; }, });
const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); const modelRegistry = ModelRegistry.create(authStorage, tempDir); authStorage.setRuntimeApiKey("anthropic", "test-key"); const extensionsResult = await createTestExtensionsResult([ (pi) => { pi.on("input", async () => { inputCalls++; return { action: "handled" }; }); }, ]);
session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, modelRegistry, resourceLoader: createTestResourceLoader({ extensionsResult }), });
await session.prompt("host gate follow-up", { internalPrompt: true }); await session.agent.waitForIdle();
expect(inputCalls).toBe(0); expect(receivedUserText).toBe("host gate follow-up"); });
it("should allow prompt() after previous completes", async () => { const model = getModel("anthropic", "claude-sonnet-4-5")!; const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: "Test", tools: [], }, streamFn: () => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ type: "start", partial: createAssistantMessage("") }); stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Done") }); }); return stream; }, });
const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); const modelRegistry = ModelRegistry.create(authStorage, tempDir); authStorage.setRuntimeApiKey("anthropic", "test-key");
session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, modelRegistry, resourceLoader: createTestResourceLoader(), });
await session.prompt("First message"); expect(session.isStreaming).toBe(false); await expect(session.prompt("Second message")).resolves.not.toThrow(); });
it("should wait for queued agent events before emitting tool_call", async () => { const model = getModel("anthropic", "claude-sonnet-4-5")!; const tool = { name: "dummy", description: "Dummy tool", label: "dummy", parameters: Type.Object({ q: Type.String() }), execute: async (_toolCallId: string, params: unknown) => { const q = typeof params === "object" && params !== null && "q" in params ? String((params as { q: unknown }).q) : ""; return { content: [{ type: "text" as const, text: `result:${q}` }], details: {}, }; }, };
const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: "Test", tools: [tool], }, streamFn: async (_model, context) => { const stream = new MockAssistantStream(); queueMicrotask(() => { const toolResultCount = context.messages.filter((message) => message.role === "toolResult").length; if (toolResultCount > 0) { const message: AssistantMessage = { role: "assistant", content: [{ type: "text", text: "done" }], api: "anthropic-messages", provider: "anthropic", model: "mock", usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0, totalTokens: 2, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "stop", timestamp: Date.now(), }; stream.push({ type: "start", partial: { ...message, content: [] } }); stream.push({ type: "done", reason: "stop", message }); return; }
const message: AssistantMessage = { role: "assistant", content: [ { type: "toolCall", id: "toolu_1", name: "dummy", arguments: { q: "x" } }, { type: "toolCall", id: "toolu_2", name: "dummy", arguments: { q: "y" } }, ], api: "anthropic-messages", provider: "anthropic", model: "mock", usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0, totalTokens: 2, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "toolUse", timestamp: Date.now(), };
stream.push({ type: "start", partial: { ...message, content: [] } }); stream.push({ type: "done", reason: "toolUse", message }); }); return stream; }, });
const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); const modelRegistry = ModelRegistry.create(authStorage, tempDir); authStorage.setRuntimeApiKey("anthropic", "test-key");
session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, modelRegistry, resourceLoader: createTestResourceLoader(), baseToolsOverride: { dummy: tool }, });
const snapshots: string[][] = []; const sessionWithRunner = session as unknown as { _extensionRunner?: { hasHandlers: (eventType: string) => boolean; emit: (event: { type: string; message?: { role?: string } }) => Promise<void>; emitMessageEnd: (event: { type: string; message?: { role?: string } }) => Promise<undefined>; emitToolCall: (event: { type: string; toolCallId: string }) => Promise<undefined>; emitInput: ( text: string, images: unknown, source: "interactive" | "rpc" | "extension", ) => Promise<{ action: "continue" }>; emitBeforeAgentStart: ( prompt: string, images: unknown, systemPrompt: string, systemPromptOptions: BuildSystemPromptOptions, ) => Promise<undefined>; invalidate: (message?: string) => void; }; }; sessionWithRunner._extensionRunner = { hasHandlers: (eventType) => eventType === "tool_call", emit: async () => {}, emitMessageEnd: async () => undefined, emitToolCall: async () => { snapshots.push( sessionManager .getEntries() .filter((entry) => entry.type === "message") .map((entry) => entry.message.role), ); return undefined; }, emitInput: async () => ({ action: "continue" }), emitBeforeAgentStart: async () => undefined, invalidate: () => {}, };
await session.prompt("hi"); await session.agent.waitForIdle();
expect(snapshots).toEqual([ ["user", "assistant"], ["user", "assistant"], ]); });
it("should persist message_end events in order with slow extension handlers", async () => { const model = getModel("anthropic", "claude-sonnet-4-5")!; const tool = { name: "dummy", description: "Dummy tool", label: "dummy", parameters: Type.Object({ q: Type.String() }), execute: async (_toolCallId: string, params: unknown) => { const q = typeof params === "object" && params !== null && "q" in params ? String((params as { q: unknown }).q) : ""; return { content: [{ type: "text" as const, text: `result:${q}` }], details: {}, }; }, };
const agent = new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: "Test", tools: [tool], }, streamFn: async (_model, context) => { const stream = new MockAssistantStream(); queueMicrotask(() => { const hasToolResult = context.messages.some((message) => message.role === "toolResult");
if (hasToolResult) { const message: AssistantMessage = { role: "assistant", content: [{ type: "text", text: "done" }], api: "anthropic-messages", provider: "anthropic", model: "mock", usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0, totalTokens: 2, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "stop", timestamp: Date.now(), }; stream.push({ type: "start", partial: { ...message, content: [] } }); stream.push({ type: "done", reason: "stop", message }); return; }
const message: AssistantMessage = { role: "assistant", content: [ { type: "text", text: "calling tool" }, { type: "toolCall", id: "toolu_1", name: "dummy", arguments: { q: "x" } }, ], api: "anthropic-messages", provider: "anthropic", model: "mock", usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0, totalTokens: 2, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "toolUse", timestamp: Date.now(), };
stream.push({ type: "start", partial: { ...message, content: [] } }); stream.push({ type: "done", reason: "toolUse", message }); }); return stream; }, });
const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); const modelRegistry = ModelRegistry.create(authStorage, tempDir); authStorage.setRuntimeApiKey("anthropic", "test-key");
session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, modelRegistry, resourceLoader: createTestResourceLoader(), baseToolsOverride: { dummy: tool }, });
const sessionWithRunner = session as unknown as { _extensionRunner?: { hasHandlers: (eventType: string) => boolean; emit: (event: { type: string; message?: { role?: string } }) => Promise<void>; emitMessageEnd: (event: { type: string; message?: { role?: string } }) => Promise<undefined>; emitInput: ( text: string, images: unknown, source: "interactive" | "rpc" | "extension", ) => Promise<{ action: "continue" }>; emitBeforeAgentStart: ( prompt: string, images: unknown, systemPrompt: string, systemPromptOptions: BuildSystemPromptOptions, ) => Promise<undefined>; invalidate: (message?: string) => void; }; }; sessionWithRunner._extensionRunner = { hasHandlers: () => false, emit: async () => {}, emitMessageEnd: async (event) => { if (event.type === "message_end" && event.message?.role === "assistant") { await new Promise((resolve) => setTimeout(resolve, 40)); } return undefined; }, emitInput: async () => ({ action: "continue" }), emitBeforeAgentStart: async () => undefined, invalidate: () => {}, };
await session.prompt("hi"); await session.agent.waitForIdle(); await new Promise((resolve) => setTimeout(resolve, 100));
const messageEntries = sessionManager.getEntries().filter((entry) => entry.type === "message"); expect(messageEntries.map((entry) => entry.message.role)).toEqual([ "user", "assistant", "toolResult", "assistant", ]); });});