From 746f770b10a5d0a566f5f18b68e56377c44ef006 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Wed, 1 Apr 2026 22:59:49 +0200 Subject: [PATCH] test(coding-agent): add session lifecycle characterization suite --- .../agent-session-bash-persistence.test.ts | 242 ++++++++++ .../suite/agent-session-compaction.test.ts | 343 ++++++++++++++ .../agent-session-model-extension.test.ts | 325 +++++++++++++ .../test/suite/agent-session-prompt.test.ts | 323 +++++++++++++ .../test/suite/agent-session-queue.test.ts | 438 +++++++++++++----- .../suite/agent-session-retry-events.test.ts | 358 ++++++++++++++ .../suite/agent-session-runtime-host.test.ts | 321 +++++++++++++ packages/coding-agent/test/suite/harness.ts | 108 +++-- ...2023-queued-slash-command-followup.test.ts | 80 ++++ 9 files changed, 2408 insertions(+), 130 deletions(-) create mode 100644 packages/coding-agent/test/suite/agent-session-bash-persistence.test.ts create mode 100644 packages/coding-agent/test/suite/agent-session-compaction.test.ts create mode 100644 packages/coding-agent/test/suite/agent-session-model-extension.test.ts create mode 100644 packages/coding-agent/test/suite/agent-session-prompt.test.ts create mode 100644 packages/coding-agent/test/suite/agent-session-retry-events.test.ts create mode 100644 packages/coding-agent/test/suite/agent-session-runtime-host.test.ts create mode 100644 packages/coding-agent/test/suite/regressions/2023-queued-slash-command-followup.test.ts diff --git a/packages/coding-agent/test/suite/agent-session-bash-persistence.test.ts b/packages/coding-agent/test/suite/agent-session-bash-persistence.test.ts new file mode 100644 index 00000000..f8903044 --- /dev/null +++ b/packages/coding-agent/test/suite/agent-session-bash-persistence.test.ts @@ -0,0 +1,242 @@ +import { Buffer } from "node:buffer"; +import type { AgentTool } from "@mariozechner/pi-agent-core"; +import { fauxAssistantMessage, fauxToolCall } from "@mariozechner/pi-ai"; +import { Type } from "@sinclair/typebox"; +import { afterEach, describe, expect, it } from "vitest"; +import type { BashOperations } from "../../src/core/tools/bash.js"; +import { createHarness, type Harness } from "./harness.js"; + +function getEntryTypes(harness: Harness): string[] { + return harness.sessionManager.getEntries().map((entry) => entry.type); +} + +describe("AgentSession bash and persistence characterization", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + while (harnesses.length > 0) { + harnesses.pop()?.cleanup(); + } + }); + + it("records bash results immediately while idle", async () => { + const harness = await createHarness(); + harnesses.push(harness); + + harness.session.recordBashResult("echo hi", { + output: "hi", + exitCode: 0, + cancelled: false, + truncated: false, + }); + + expect(harness.session.hasPendingBashMessages).toBe(false); + expect(harness.session.messages[harness.session.messages.length - 1]?.role).toBe("bashExecution"); + expect(getEntryTypes(harness)).toContain("message"); + }); + + it("defers bash results while streaming and flushes them before the next prompt", async () => { + let releaseToolExecution: (() => void) | undefined; + const toolRelease = new Promise((resolve) => { + releaseToolExecution = resolve; + }); + const waitTool: AgentTool = { + name: "wait", + label: "Wait", + description: "Wait for release", + parameters: Type.Object({}), + execute: async () => { + await toolRelease; + return { + content: [{ type: "text", text: "released" }], + details: {}, + }; + }, + }; + const harness = await createHarness({ tools: [waitTool] }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage([fauxToolCall("wait", {})], { stopReason: "toolUse" }), + fauxAssistantMessage("done"), + fauxAssistantMessage("after flush"), + ]); + + const sawToolStart = new Promise((resolve) => { + const unsubscribe = harness.session.subscribe((event) => { + if (event.type === "tool_execution_start") { + unsubscribe(); + resolve(); + } + }); + }); + + const firstPrompt = harness.session.prompt("start"); + await sawToolStart; + harness.session.recordBashResult("echo hi", { + output: "hi", + exitCode: 0, + cancelled: false, + truncated: false, + }); + + expect(harness.session.hasPendingBashMessages).toBe(true); + expect(harness.session.messages.some((message) => message.role === "bashExecution")).toBe(false); + + releaseToolExecution?.(); + await firstPrompt; + + expect(harness.session.hasPendingBashMessages).toBe(true); + expect(harness.session.messages.some((message) => message.role === "bashExecution")).toBe(false); + + await harness.session.prompt("next turn"); + + expect(harness.session.hasPendingBashMessages).toBe(false); + expect(harness.session.messages.some((message) => message.role === "bashExecution")).toBe(true); + expect(getEntryTypes(harness).filter((type) => type === "message").length).toBeGreaterThan(0); + }); + + it("executes bash commands and records the result", async () => { + const harness = await createHarness(); + harnesses.push(harness); + + const result = await harness.session.executeBash("printf 'hello'"); + + expect(result.output).toContain("hello"); + expect(harness.session.messages[harness.session.messages.length - 1]?.role).toBe("bashExecution"); + }); + + it("cancels running bash commands with abortBash", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const operations: BashOperations = { + exec: async (_command, _cwd, options) => { + return await new Promise<{ exitCode: number | null }>((_resolve, reject) => { + options.signal?.addEventListener( + "abort", + () => { + reject(new Error("aborted")); + }, + { once: true }, + ); + }); + }, + }; + + const bashPromise = harness.session.executeBash("sleep", undefined, { operations }); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(harness.session.isBashRunning).toBe(true); + harness.session.abortBash(); + + const result = await bashPromise; + expect(result.cancelled).toBe(true); + expect(harness.session.isBashRunning).toBe(false); + }); + + it("persists user, assistant, toolResult, and custom messages in order", async () => { + const echoTool: AgentTool = { + name: "echo", + label: "Echo", + description: "Echo text back", + parameters: Type.Object({ text: Type.String() }), + execute: async (_toolCallId, params) => { + const text = typeof params === "object" && params !== null && "text" in params ? String(params.text) : ""; + return { content: [{ type: "text", text: `echo:${text}` }], details: { text } }; + }, + }; + const harness = await createHarness({ tools: [echoTool] }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage([fauxToolCall("echo", { text: "hello" })], { stopReason: "toolUse" }), + fauxAssistantMessage("done"), + ]); + + await harness.session.sendCustomMessage({ + customType: "note", + content: "hello", + display: true, + details: { a: 1 }, + }); + await harness.session.prompt("start"); + + const entries = harness.sessionManager.getEntries(); + expect(entries.map((entry) => entry.type)).toEqual([ + "custom_message", + "message", + "message", + "message", + "message", + ]); + expect(harness.session.messages.map((message) => message.role)).toEqual([ + "custom", + "user", + "assistant", + "toolResult", + "assistant", + ]); + }); + + it("does not emit message_end for bash execution messages", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const messageEndRoles: string[] = []; + harness.session.subscribe((event) => { + if (event.type === "message_end") { + messageEndRoles.push(event.message.role); + } + }); + + harness.session.recordBashResult("echo hi", { + output: "hi", + exitCode: 0, + cancelled: false, + truncated: false, + }); + + expect(messageEndRoles).toEqual([]); + }); + + it("persists aborted assistant messages", async () => { + const harness = await createHarness(); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("x".repeat(20_000))]); + + const sawMessageUpdate = new Promise((resolve) => { + const unsubscribe = harness.session.subscribe((event) => { + if (event.type === "message_update") { + unsubscribe(); + resolve(); + } + }); + }); + + const promptPromise = harness.session.prompt("hi"); + await sawMessageUpdate; + await harness.session.abort(); + await promptPromise; + + const lastEntry = harness.sessionManager.getEntries()[harness.sessionManager.getEntries().length - 1]; + expect(lastEntry?.type).toBe("message"); + if (lastEntry?.type === "message") { + expect(lastEntry.message.role).toBe("assistant"); + if (lastEntry.message.role === "assistant") { + expect(lastEntry.message.stopReason).toBe("aborted"); + } + } + }); + + it("records bash output through custom operations", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const operations: BashOperations = { + exec: async (_command, _cwd, options) => { + options.onData(Buffer.from("hello from custom ops")); + return { exitCode: 0 }; + }, + }; + + const result = await harness.session.executeBash("custom", undefined, { operations }); + + expect(result.output).toContain("hello from custom ops"); + expect(harness.session.messages[harness.session.messages.length - 1]?.role).toBe("bashExecution"); + }); +}); diff --git a/packages/coding-agent/test/suite/agent-session-compaction.test.ts b/packages/coding-agent/test/suite/agent-session-compaction.test.ts new file mode 100644 index 00000000..256fb380 --- /dev/null +++ b/packages/coding-agent/test/suite/agent-session-compaction.test.ts @@ -0,0 +1,343 @@ +import { type AssistantMessage, fauxAssistantMessage, type Model } from "@mariozechner/pi-ai"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { createHarness, type Harness } from "./harness.js"; + +type SessionWithCompactionInternals = { + _checkCompaction: (assistantMessage: AssistantMessage, skipAbortedCheck?: boolean) => Promise; + _runAutoCompaction: (reason: "overflow" | "threshold", willRetry: boolean) => Promise; +}; + +function createUsage(totalTokens: number) { + return { + input: totalTokens, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; +} + +function createAssistant( + harness: Harness, + options: { + stopReason?: AssistantMessage["stopReason"]; + errorMessage?: string; + totalTokens?: number; + timestamp?: number; + }, +): AssistantMessage { + const model = harness.getModel(); + return { + ...fauxAssistantMessage("", { + stopReason: options.stopReason, + errorMessage: options.errorMessage, + timestamp: options.timestamp, + }), + api: model.api, + provider: model.provider, + model: model.id, + usage: createUsage(options.totalTokens ?? 0), + }; +} + +describe("AgentSession compaction characterization", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + vi.useRealTimers(); + vi.restoreAllMocks(); + while (harnesses.length > 0) { + harnesses.pop()?.cleanup(); + } + }); + + it("manually compacts using an extension-provided summary", async () => { + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("session_before_compact", async (event) => ({ + compaction: { + summary: "summary from extension", + firstKeptEntryId: event.preparation.firstKeptEntryId, + tokensBefore: event.preparation.tokensBefore, + details: { source: "extension" }, + }, + })); + }, + ], + }); + harnesses.push(harness); + + await harness.session.prompt("one"); + await harness.session.prompt("two"); + + const result = await harness.session.compact(); + const compactionEntries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction"); + + expect(result.summary).toBe("summary from extension"); + expect(compactionEntries).toHaveLength(1); + expect(harness.session.messages[0]?.role).toBe("compactionSummary"); + }); + + it("throws when compacting without a model", async () => { + const harness = await createHarness(); + harnesses.push(harness); + harness.session.agent.state.model = undefined as unknown as Model; + + await expect(harness.session.compact()).rejects.toThrow("No model selected"); + }); + + it("throws when compacting without configured auth", async () => { + const harness = await createHarness({ withConfiguredAuth: false }); + harnesses.push(harness); + + await expect(harness.session.compact()).rejects.toThrow(`No API key found for ${harness.getModel().provider}.`); + }); + + it("cancels in-progress manual compaction when abortCompaction is called", async () => { + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("session_before_compact", async (event) => { + return await new Promise<{ cancel: true }>((resolve) => { + event.signal.addEventListener("abort", () => resolve({ cancel: true }), { once: true }); + }); + }); + }, + ], + }); + harnesses.push(harness); + + await harness.session.prompt("one"); + await harness.session.prompt("two"); + + const compactPromise = harness.session.compact(); + await new Promise((resolve) => setTimeout(resolve, 0)); + harness.session.abortCompaction(); + + await expect(compactPromise).rejects.toThrow("Compaction cancelled"); + }); + + it("resumes after threshold compaction when only agent-level queued messages exist", async () => { + vi.useFakeTimers(); + const harness = await createHarness({ + settings: { compaction: { keepRecentTokens: 1 } }, + extensionFactories: [ + (pi) => { + pi.on("session_before_compact", async (event) => ({ + compaction: { + summary: "auto compacted", + firstKeptEntryId: event.preparation.firstKeptEntryId, + tokensBefore: event.preparation.tokensBefore, + details: {}, + }, + })); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("one"), fauxAssistantMessage("two")]); + await harness.session.prompt("first"); + await harness.session.prompt("second"); + + harness.session.agent.followUp({ + role: "custom", + customType: "test", + content: [{ type: "text", text: "queued custom" }], + display: false, + timestamp: Date.now(), + }); + + const continueSpy = vi.spyOn(harness.session.agent, "continue").mockResolvedValue(); + const sessionInternals = harness.session as unknown as SessionWithCompactionInternals; + + await sessionInternals._runAutoCompaction("threshold", false); + await vi.advanceTimersByTimeAsync(100); + + expect(continueSpy).toHaveBeenCalledTimes(1); + }); + + it("does not retry overflow recovery more than once", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const sessionInternals = harness.session as unknown as SessionWithCompactionInternals; + const overflowMessage = createAssistant(harness, { + stopReason: "error", + errorMessage: "prompt is too long", + timestamp: Date.now(), + }); + const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(); + const compactionErrors: string[] = []; + harness.session.subscribe((event) => { + if (event.type === "compaction_end" && event.errorMessage) { + compactionErrors.push(event.errorMessage); + } + }); + + await sessionInternals._checkCompaction(overflowMessage); + await sessionInternals._checkCompaction({ ...overflowMessage, timestamp: Date.now() + 1 }); + + expect(runAutoCompactionSpy).toHaveBeenCalledTimes(1); + expect(compactionErrors).toContain( + "Context overflow recovery failed after one compact-and-retry attempt. Try reducing context or switching to a larger-context model.", + ); + }); + + it("ignores stale pre-compaction assistant usage on pre-prompt checks", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const sessionInternals = harness.session as unknown as SessionWithCompactionInternals; + const staleTimestamp = Date.now() - 10_000; + const staleAssistant = createAssistant(harness, { + stopReason: "stop", + totalTokens: 610_000, + timestamp: staleTimestamp, + }); + + harness.sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: "before compaction" }], + timestamp: staleTimestamp - 1000, + }); + harness.sessionManager.appendMessage(staleAssistant); + const firstKeptEntryId = harness.sessionManager.getEntries()[0]!.id; + harness.sessionManager.appendCompaction( + "summary", + firstKeptEntryId, + staleAssistant.usage.totalTokens, + undefined, + false, + ); + harness.sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: "after compaction" }], + timestamp: Date.now(), + }); + + const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(); + + await sessionInternals._checkCompaction(staleAssistant, false); + + expect(runAutoCompactionSpy).not.toHaveBeenCalled(); + }); + + it("triggers threshold compaction for error messages using the last successful usage", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const sessionInternals = harness.session as unknown as SessionWithCompactionInternals; + const successfulAssistant = createAssistant(harness, { + stopReason: "stop", + totalTokens: 190_000, + timestamp: Date.now(), + }); + const errorAssistant = createAssistant(harness, { + stopReason: "error", + errorMessage: "529 overloaded", + timestamp: Date.now() + 1000, + }); + harness.session.agent.state.messages = [ + { role: "user", content: [{ type: "text", text: "hello" }], timestamp: Date.now() - 1000 }, + successfulAssistant, + { role: "user", content: [{ type: "text", text: "retry" }], timestamp: Date.now() + 500 }, + errorAssistant, + ]; + + const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(); + + await sessionInternals._checkCompaction(errorAssistant); + + expect(runAutoCompactionSpy).toHaveBeenCalledWith("threshold", false); + }); + + it("does not trigger threshold compaction for error messages when no prior usage exists", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const sessionInternals = harness.session as unknown as SessionWithCompactionInternals; + const errorAssistant = createAssistant(harness, { + stopReason: "error", + errorMessage: "529 overloaded", + timestamp: Date.now(), + }); + harness.session.agent.state.messages = [ + { role: "user", content: [{ type: "text", text: "hello" }], timestamp: Date.now() - 1000 }, + errorAssistant, + ]; + + const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(); + + await sessionInternals._checkCompaction(errorAssistant); + + expect(runAutoCompactionSpy).not.toHaveBeenCalled(); + }); + + it("does not trigger threshold compaction when only kept pre-compaction usage exists", async () => { + const harness = await createHarness(); + harnesses.push(harness); + const sessionInternals = harness.session as unknown as SessionWithCompactionInternals; + const preCompactionTimestamp = Date.now() - 10_000; + const keptAssistant = createAssistant(harness, { + stopReason: "stop", + totalTokens: 190_000, + timestamp: preCompactionTimestamp, + }); + + harness.sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: "before compaction" }], + timestamp: preCompactionTimestamp - 1000, + }); + harness.sessionManager.appendMessage(keptAssistant); + const firstKeptEntryId = harness.sessionManager.getEntries()[0]!.id; + harness.sessionManager.appendCompaction( + "summary", + firstKeptEntryId, + keptAssistant.usage.totalTokens, + undefined, + false, + ); + + const errorAssistant = createAssistant(harness, { + stopReason: "error", + errorMessage: "529 overloaded", + timestamp: Date.now(), + }); + harness.session.agent.state.messages = [ + { role: "user", content: [{ type: "text", text: "kept user" }], timestamp: preCompactionTimestamp - 1000 }, + keptAssistant, + { role: "user", content: [{ type: "text", text: "new prompt" }], timestamp: Date.now() - 500 }, + errorAssistant, + ]; + + const runAutoCompactionSpy = vi.spyOn(sessionInternals, "_runAutoCompaction").mockResolvedValue(); + + await sessionInternals._checkCompaction(errorAssistant); + + expect(runAutoCompactionSpy).not.toHaveBeenCalled(); + }); + + it("does not trigger threshold compaction below the threshold or when disabled", async () => { + const belowThresholdHarness = await createHarness({ + settings: { compaction: { enabled: true, reserveTokens: 1000 } }, + models: [{ id: "faux-1", contextWindow: 200_000 }], + }); + harnesses.push(belowThresholdHarness); + const disabledHarness = await createHarness({ settings: { compaction: { enabled: false } } }); + harnesses.push(disabledHarness); + + const belowThresholdInternals = belowThresholdHarness.session as unknown as SessionWithCompactionInternals; + const disabledInternals = disabledHarness.session as unknown as SessionWithCompactionInternals; + const belowThresholdSpy = vi.spyOn(belowThresholdInternals, "_runAutoCompaction").mockResolvedValue(); + const disabledSpy = vi.spyOn(disabledInternals, "_runAutoCompaction").mockResolvedValue(); + + await belowThresholdInternals._checkCompaction( + createAssistant(belowThresholdHarness, { stopReason: "stop", totalTokens: 1_000, timestamp: Date.now() }), + ); + await disabledInternals._checkCompaction( + createAssistant(disabledHarness, { stopReason: "stop", totalTokens: 1_000_000, timestamp: Date.now() }), + ); + + expect(belowThresholdSpy).not.toHaveBeenCalled(); + expect(disabledSpy).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/coding-agent/test/suite/agent-session-model-extension.test.ts b/packages/coding-agent/test/suite/agent-session-model-extension.test.ts new file mode 100644 index 00000000..1b207959 --- /dev/null +++ b/packages/coding-agent/test/suite/agent-session-model-extension.test.ts @@ -0,0 +1,325 @@ +import type { AgentTool, ThinkingLevel } from "@mariozechner/pi-agent-core"; +import { fauxAssistantMessage, fauxToolCall, type Model } from "@mariozechner/pi-ai"; +import { Type } from "@sinclair/typebox"; +import { afterEach, describe, expect, it } from "vitest"; +import type { ExtensionAPI } from "../../src/index.js"; +import { createHarness, getAssistantTexts, type Harness } from "./harness.js"; + +describe("AgentSession model and extension characterization", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + while (harnesses.length > 0) { + harnesses.pop()?.cleanup(); + } + }); + + it("setModel saves the model and emits model_select", async () => { + const modelEvents: string[] = []; + const harness = await createHarness({ + models: [ + { id: "faux-1", name: "One", reasoning: true }, + { id: "faux-2", name: "Two", reasoning: true }, + ], + extensionFactories: [ + (pi) => { + pi.on("model_select", async (event) => { + modelEvents.push(`${event.previousModel?.id ?? "none"}->${event.model.id}:${event.source}`); + }); + }, + ], + }); + harnesses.push(harness); + const nextModel = harness.getModel("faux-2")!; + + await harness.session.setModel(nextModel); + + expect(harness.session.model?.id).toBe("faux-2"); + expect(modelEvents).toEqual(["faux-1->faux-2:set"]); + expect( + harness.sessionManager + .getEntries() + .filter((entry) => entry.type === "model_change") + .map((entry) => `${entry.provider}/${entry.modelId}`), + ).toEqual([`${nextModel.provider}/${nextModel.id}`]); + }); + + it("cycles through scoped models and preserves the scoped thinking preference", async () => { + const harness = await createHarness({ + models: [ + { id: "faux-1", name: "One", reasoning: true }, + { id: "faux-2", name: "Two", reasoning: false }, + ], + }); + harnesses.push(harness); + const modelOne = harness.getModel("faux-1")!; + const modelTwo = harness.getModel("faux-2")!; + harness.session.setScopedModels([{ model: modelOne, thinkingLevel: "high" }, { model: modelTwo }] as Array<{ + model: Model; + thinkingLevel?: ThinkingLevel; + }>); + harness.session.setThinkingLevel("high"); + + await harness.session.cycleModel(); + expect(harness.session.model?.id).toBe("faux-2"); + expect(harness.session.thinkingLevel).toBe("off"); + + await harness.session.cycleModel(); + expect(harness.session.model?.id).toBe("faux-1"); + expect(harness.session.thinkingLevel).toBe("high"); + }); + + it("clamps thinking levels to model capabilities and cycles available levels", async () => { + const harness = await createHarness({ models: [{ id: "faux-1", reasoning: false }] }); + harnesses.push(harness); + + harness.session.setThinkingLevel("high"); + expect(harness.session.thinkingLevel).toBe("off"); + expect(harness.session.cycleThinkingLevel()).toBeUndefined(); + }); + + it("throws when setModel is called without configured auth", async () => { + const harness = await createHarness({ + models: [ + { id: "faux-1", name: "One", reasoning: true }, + { id: "faux-2", name: "Two", reasoning: true }, + ], + withConfiguredAuth: false, + }); + harnesses.push(harness); + + await expect(harness.session.setModel(harness.getModel("faux-2")!)).rejects.toThrow( + `No API key for ${harness.getModel().provider}/faux-2`, + ); + }); + + it("allows extension tool_call handlers to block tool execution", async () => { + const echoTool: AgentTool = { + name: "echo", + label: "Echo", + description: "Echo text back", + parameters: Type.Object({ text: Type.String() }), + execute: async () => { + throw new Error("tool should have been blocked"); + }, + }; + const harness = await createHarness({ + tools: [echoTool], + extensionFactories: [ + (pi) => { + pi.on("tool_call", async () => ({ block: true, reason: "Blocked by test" })); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage([fauxToolCall("echo", { text: "hello" })], { stopReason: "toolUse" }), + (context) => { + const toolResult = context.messages.find((message) => message.role === "toolResult"); + const errorText = + toolResult?.role === "toolResult" + ? toolResult.content + .filter((part): part is { type: "text"; text: string } => part.type === "text") + .map((part) => part.text) + .join("\n") + : ""; + return fauxAssistantMessage(errorText); + }, + ]); + + await harness.session.prompt("hi"); + + expect(getAssistantTexts(harness)).toContain("Blocked by test"); + expect( + harness.session.messages.find((message) => message.role === "toolResult" && message.isError), + ).toBeDefined(); + }); + + it("allows extension tool_result handlers to modify tool results", async () => { + const echoTool: AgentTool = { + name: "echo", + label: "Echo", + description: "Echo text back", + parameters: Type.Object({ text: Type.String() }), + execute: async (_toolCallId, params) => { + const text = typeof params === "object" && params !== null && "text" in params ? String(params.text) : ""; + return { content: [{ type: "text", text }], details: { text } }; + }, + }; + const harness = await createHarness({ + tools: [echoTool], + extensionFactories: [ + (pi) => { + pi.on("tool_result", async () => ({ + content: [{ type: "text", text: "patched result" }], + details: { patched: true }, + })); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage([fauxToolCall("echo", { text: "hello" })], { stopReason: "toolUse" }), + (context) => { + const toolResult = context.messages.find((message) => message.role === "toolResult"); + const text = + toolResult?.role === "toolResult" + ? toolResult.content + .filter((part): part is { type: "text"; text: string } => part.type === "text") + .map((part) => part.text) + .join("\n") + : ""; + return fauxAssistantMessage(text); + }, + ]); + + await harness.session.prompt("hi"); + + expect(getAssistantTexts(harness)).toContain("patched result"); + expect( + harness.session.messages.find((message) => message.role === "toolResult" && message.details?.patched === true), + ).toBeDefined(); + }); + + it("allows extension context handlers to modify messages before the LLM call", async () => { + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("context", async (event) => ({ + messages: event.messages.map((message) => + message.role === "user" + ? { ...message, content: [{ type: "text", text: "rewritten" }], timestamp: message.timestamp } + : message, + ), + })); + }, + ], + }); + harnesses.push(harness); + let providerUserText = ""; + harness.setResponses([ + (context) => { + const user = context.messages.find((message) => message.role === "user"); + providerUserText = + user && typeof user.content !== "string" + ? user.content + .filter((part): part is { type: "text"; text: string } => part.type === "text") + .map((part) => part.text) + .join("\n") + : ""; + return fauxAssistantMessage("done"); + }, + ]); + + await harness.session.prompt("original"); + + expect(providerUserText).toBe("rewritten"); + const storedUserMessage = harness.session.messages.find((message) => message.role === "user"); + expect(storedUserMessage?.role).toBe("user"); + if (storedUserMessage?.role === "user") { + expect(storedUserMessage.content).toEqual([{ type: "text", text: "original" }]); + } + }); + + it("allows extension input handlers to transform or handle input", async () => { + let extensionApi: ExtensionAPI | undefined; + const transformedHarness = await createHarness({ + extensionFactories: [ + (pi) => { + extensionApi = pi; + pi.on("input", async (event) => { + if (event.text === "ping") { + return { action: "handled" }; + } + return { action: "transform", text: `transformed:${event.text}` }; + }); + }, + ], + }); + harnesses.push(transformedHarness); + let providerUserText = ""; + transformedHarness.setResponses([ + (context) => { + const user = context.messages.find((message) => message.role === "user"); + providerUserText = + user && typeof user.content !== "string" + ? user.content + .filter((part): part is { type: "text"; text: string } => part.type === "text") + .map((part) => part.text) + .join("\n") + : ""; + return fauxAssistantMessage("done"); + }, + ]); + + await transformedHarness.session.prompt("hello"); + await transformedHarness.session.prompt("ping"); + + expect(providerUserText).toBe("transformed:hello"); + expect(transformedHarness.session.messages.filter((message) => message.role === "user")).toHaveLength(1); + expect(extensionApi).toBeDefined(); + }); + + it("allows before_agent_start handlers to inject custom messages and modify the system prompt", async () => { + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("before_agent_start", async (event) => ({ + message: { + customType: "before-start", + content: "injected", + display: true, + details: { injected: true }, + }, + systemPrompt: `${event.systemPrompt}\n\nextra instructions`, + })); + }, + ], + }); + harnesses.push(harness); + let providerSystemPrompt = ""; + let sawInjectedUserMessage = false; + harness.setResponses([ + (context) => { + providerSystemPrompt = context.systemPrompt ?? ""; + sawInjectedUserMessage = context.messages.some( + (message) => + message.role === "user" && + typeof message.content !== "string" && + message.content.some((part) => part.type === "text" && part.text === "injected"), + ); + return fauxAssistantMessage("done"); + }, + ]); + + await harness.session.prompt("hello"); + + expect(providerSystemPrompt).toContain("extra instructions"); + expect(sawInjectedUserMessage).toBe(true); + expect( + harness.session.messages.some((message) => message.role === "custom" && message.customType === "before-start"), + ).toBe(true); + }); + + it("bindExtensions emits session_start and reload emits session_shutdown then session_start", async () => { + const lifecycleEvents: string[] = []; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("session_start", async (event) => { + lifecycleEvents.push(`start:${event.reason}`); + }); + pi.on("session_shutdown", async () => { + lifecycleEvents.push("shutdown"); + }); + }, + ], + }); + harnesses.push(harness); + + await harness.session.bindExtensions({ shutdownHandler: () => {} }); + await harness.session.reload(); + + expect(lifecycleEvents).toEqual(["start:startup", "shutdown", "start:reload"]); + }); +}); diff --git a/packages/coding-agent/test/suite/agent-session-prompt.test.ts b/packages/coding-agent/test/suite/agent-session-prompt.test.ts new file mode 100644 index 00000000..aaeb5896 --- /dev/null +++ b/packages/coding-agent/test/suite/agent-session-prompt.test.ts @@ -0,0 +1,323 @@ +import { mkdirSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type { AgentTool } from "@mariozechner/pi-agent-core"; +import { fauxAssistantMessage, fauxToolCall, type Model } from "@mariozechner/pi-ai"; +import { Type } from "@sinclair/typebox"; +import { afterEach, describe, expect, it } from "vitest"; +import type { PromptTemplate } from "../../src/core/prompt-templates.js"; +import { createSyntheticSourceInfo } from "../../src/core/source-info.js"; +import { createTestResourceLoader } from "../utilities.js"; +import { createHarness, getMessageText, type Harness } from "./harness.js"; + +describe("AgentSession prompt characterization", () => { + const harnesses: Harness[] = []; + const tempDirs: string[] = []; + + afterEach(() => { + while (harnesses.length > 0) { + harnesses.pop()?.cleanup(); + } + while (tempDirs.length > 0) { + const tempDir = tempDirs.pop(); + if (tempDir) { + rmSync(tempDir, { recursive: true, force: true }); + } + } + }); + + it("prompts while idle and records a single text response", async () => { + const harness = await createHarness(); + harnesses.push(harness); + + harness.setResponses([fauxAssistantMessage("hello")]); + + await harness.session.prompt("hi"); + + expect(harness.session.messages.map((message) => message.role)).toEqual(["user", "assistant"]); + expect(getMessageText(harness.session.messages[0]!)).toBe("hi"); + expect(harness.getPendingResponseCount()).toBe(0); + }); + + it("handles a tool call turn and waits for the follow-up LLM response", async () => { + const toolRuns: string[] = []; + const echoTool: AgentTool = { + name: "echo", + label: "Echo", + description: "Echo text back", + parameters: Type.Object({ text: Type.String() }), + execute: async (_toolCallId, params) => { + const text = typeof params === "object" && params !== null && "text" in params ? String(params.text) : ""; + toolRuns.push(text); + return { + content: [{ type: "text", text: `echo:${text}` }], + details: { text }, + }; + }, + }; + const harness = await createHarness({ tools: [echoTool] }); + harnesses.push(harness); + + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("echo", { text: "hello" }), { stopReason: "toolUse" }), + fauxAssistantMessage("done"), + ]); + + await harness.session.prompt("start"); + + expect(toolRuns).toEqual(["hello"]); + expect(harness.session.messages.map((message) => message.role)).toEqual([ + "user", + "assistant", + "toolResult", + "assistant", + ]); + expect(harness.session.messages[2]?.role).toBe("toolResult"); + expect(harness.session.messages[3]?.role).toBe("assistant"); + }); + + it("executes multiple tool calls from one response and continues with a single follow-up response", async () => { + const toolRuns: string[] = []; + const makeTool = (name: string, delayMs: number): AgentTool => ({ + name, + label: name, + description: `${name} tool`, + parameters: Type.Object({ value: Type.String() }), + execute: async (_toolCallId, params) => { + const value = + typeof params === "object" && params !== null && "value" in params ? String(params.value) : ""; + await new Promise((resolve) => setTimeout(resolve, delayMs)); + toolRuns.push(`${name}:${value}`); + return { + content: [{ type: "text", text: `${name}:${value}` }], + details: { value }, + }; + }, + }); + const harness = await createHarness({ tools: [makeTool("slow", 25), makeTool("fast", 0)] }); + harnesses.push(harness); + + harness.setResponses([ + fauxAssistantMessage([fauxToolCall("slow", { value: "a" }), fauxToolCall("fast", { value: "b" })], { + stopReason: "toolUse", + }), + (context) => { + const toolResults = context.messages.filter((message) => message.role === "toolResult"); + return fauxAssistantMessage(`tool results: ${toolResults.length}`); + }, + ]); + + await harness.session.prompt("run tools"); + + expect(toolRuns.sort()).toEqual(["fast:b", "slow:a"]); + expect(harness.session.messages.filter((message) => message.role === "toolResult")).toHaveLength(2); + expect(harness.session.messages[harness.session.messages.length - 1]?.role).toBe("assistant"); + }); + + it("preserves image attachments in the provider context", async () => { + const harness = await createHarness(); + harnesses.push(harness); + let sawImage = false; + + harness.setResponses([ + (context) => { + const user = context.messages.find((message) => message.role === "user"); + sawImage = + user?.role === "user" && + typeof user.content !== "string" && + user.content.some((part) => part.type === "image"); + return fauxAssistantMessage("ok"); + }, + ]); + + await harness.session.prompt("describe", { + images: [ + { + type: "image", + mimeType: "image/png", + data: "ZmFrZQ==", + }, + ], + }); + + expect(sawImage).toBe(true); + }); + + it("expands skill commands before sending the prompt", async () => { + const tempDir = join(tmpdir(), `pi-skill-${Date.now()}-${Math.random().toString(36).slice(2)}`); + mkdirSync(tempDir, { recursive: true }); + tempDirs.push(tempDir); + const skillPath = join(tempDir, "test-skill.md"); + writeFileSync(skillPath, "# Test Skill\n\nUse the skill body."); + + const resourceLoader = { + ...createTestResourceLoader(), + getSkills: () => ({ + skills: [ + { + name: "test", + description: "Test skill", + filePath: skillPath, + disableModelInvocation: false, + baseDir: tempDir, + sourceInfo: createSyntheticSourceInfo(skillPath, { + source: "local", + scope: "project", + origin: "top-level", + baseDir: tempDir, + }), + }, + ], + diagnostics: [], + }), + }; + const harness = await createHarness({ resourceLoader }); + harnesses.push(harness); + let expandedPrompt = ""; + + harness.setResponses([ + (context) => { + const user = context.messages.find((message) => message.role === "user"); + expandedPrompt = user ? getMessageText(user) : ""; + return fauxAssistantMessage("ok"); + }, + ]); + + await harness.session.prompt("/skill:test explain this"); + + expect(expandedPrompt).toContain(' { + const template: PromptTemplate = { + name: "review", + description: "Review template", + content: "Review this code: $1", + filePath: "/virtual/review.md", + sourceInfo: createSyntheticSourceInfo("/virtual/review.md", { + source: "local", + scope: "temporary", + origin: "top-level", + }), + }; + const resourceLoader = { + ...createTestResourceLoader(), + getPrompts: () => ({ prompts: [template], diagnostics: [] }), + }; + const harness = await createHarness({ resourceLoader }); + harnesses.push(harness); + let expandedPrompt = ""; + + harness.setResponses([ + (context) => { + const user = context.messages.find((message) => message.role === "user"); + expandedPrompt = user ? getMessageText(user) : ""; + return fauxAssistantMessage("ok"); + }, + ]); + + await harness.session.prompt("/review src/index.ts"); + + expect(expandedPrompt).toBe("Review this code: src/index.ts"); + }); + + it("dispatches extension commands without consuming a provider response", async () => { + const commandRuns: string[] = []; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.registerCommand("testcmd", { + description: "Test command", + handler: async (args) => { + commandRuns.push(args); + }, + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("should stay queued")]); + + await harness.session.prompt("/testcmd hello world"); + + expect(commandRuns).toEqual(["hello world"]); + expect(harness.session.messages).toEqual([]); + expect(harness.getPendingResponseCount()).toBe(1); + }); + + it("sendUserMessage while idle triggers a turn", async () => { + const harness = await createHarness(); + harnesses.push(harness); + + harness.setResponses([fauxAssistantMessage("response")]); + + await harness.session.sendUserMessage("from extension"); + + expect(harness.session.messages.map((message) => message.role)).toEqual(["user", "assistant"]); + expect(getMessageText(harness.session.messages[0]!)).toBe("from extension"); + }); + + it("throws when prompted during streaming without a streamingBehavior", async () => { + let releaseToolExecution: (() => void) | undefined; + const toolRelease = new Promise((resolve) => { + releaseToolExecution = resolve; + }); + const waitTool: AgentTool = { + name: "wait", + label: "Wait", + description: "Wait for release", + parameters: Type.Object({}), + execute: async () => { + await toolRelease; + return { + content: [{ type: "text", text: "released" }], + details: {}, + }; + }, + }; + const harness = await createHarness({ tools: [waitTool] }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }), + fauxAssistantMessage("done"), + ]); + + const sawToolStart = new Promise((resolve) => { + const unsubscribe = harness.session.subscribe((event) => { + if (event.type === "tool_execution_start") { + unsubscribe(); + resolve(); + } + }); + }); + + const promptPromise = harness.session.prompt("start"); + await sawToolStart; + + await expect(harness.session.prompt("second")).rejects.toThrow( + "Agent is already processing. Specify streamingBehavior ('steer' or 'followUp') to queue the message.", + ); + + releaseToolExecution?.(); + await promptPromise; + }); + + it("throws when prompting without a model", async () => { + const harness = await createHarness(); + harnesses.push(harness); + harness.session.agent.state.model = undefined as unknown as Model; + + await expect(harness.session.prompt("hi")).rejects.toThrow("No model selected."); + }); + + it("throws when prompting without configured auth", async () => { + const harness = await createHarness({ withConfiguredAuth: false }); + harnesses.push(harness); + + await expect(harness.session.prompt("hi")).rejects.toThrow( + `No API key found for ${harness.getModel().provider}.`, + ); + }); +}); diff --git a/packages/coding-agent/test/suite/agent-session-queue.test.ts b/packages/coding-agent/test/suite/agent-session-queue.test.ts index 226227aa..a46cce6f 100644 --- a/packages/coding-agent/test/suite/agent-session-queue.test.ts +++ b/packages/coding-agent/test/suite/agent-session-queue.test.ts @@ -3,32 +3,58 @@ import { fauxAssistantMessage, fauxToolCall } from "@mariozechner/pi-ai"; import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; import { Type } from "@sinclair/typebox"; import { afterEach, describe, expect, it } from "vitest"; -import { createHarness, type Harness } from "./harness.js"; +import { createHarness, getAssistantTexts, getMessageText, getUserTexts, type Harness } from "./harness.js"; -function getUserTexts(harness: Harness): string[] { - return harness.session.messages - .filter((message) => message.role === "user") - .map((message) => { - const content = message.content; - if (typeof content === "string") { - return content; +async function createWaitingHarness( + options: { + tools?: AgentTool[]; + extensionFactories?: Harness["session"]["extensionRunner"] extends never + ? never + : Array<(pi: ExtensionAPI) => void>; + } = {}, +): Promise<{ + harness: Harness; + releaseToolExecution: () => void; + promptPromise: Promise; + waitForToolStart: Promise; +}> { + let releaseToolExecution: (() => void) | undefined; + const toolRelease = new Promise((resolve) => { + releaseToolExecution = resolve; + }); + const waitTool: AgentTool = { + name: "wait", + label: "Wait", + description: "Wait for release", + parameters: Type.Object({}), + execute: async () => { + await toolRelease; + return { + content: [{ type: "text", text: "released" }], + details: {}, + }; + }, + }; + const harness = await createHarness({ + tools: [waitTool, ...(options.tools ?? [])], + extensionFactories: options.extensionFactories, + }); + + const waitForToolStart = new Promise((resolve) => { + const unsubscribe = harness.session.subscribe((event) => { + if (event.type === "tool_execution_start" && event.toolName === "wait") { + unsubscribe(); + resolve(); } - return content - .filter((part): part is { type: "text"; text: string } => part.type === "text") - .map((part) => part.text) - .join("\n"); }); -} + }); -function getAssistantTexts(harness: Harness): string[] { - return harness.session.messages - .filter((message) => message.role === "assistant") - .map((message) => - message.content - .filter((part): part is { type: "text"; text: string } => part.type === "text") - .map((part) => part.text) - .join("\n"), - ); + return { + harness, + releaseToolExecution: () => releaseToolExecution?.(), + promptPromise: harness.session.prompt("start"), + waitForToolStart, + }; } describe("AgentSession queue characterization", () => { @@ -65,128 +91,332 @@ describe("AgentSession queue characterization", () => { it("delivers extension-origin steering messages before the next LLM call", async () => { let extensionApi: ExtensionAPI | undefined; - let releaseToolExecution: (() => void) | undefined; - const toolStarted = new Promise((resolve) => { - releaseToolExecution = resolve; - }); - const waitTool: AgentTool = { - name: "wait", - label: "Wait", - description: "Wait for the test to release execution", - parameters: Type.Object({}), - execute: async () => { - await toolStarted; - return { - content: [{ type: "text", text: "released" }], - details: {}, - }; - }, - }; - const harness = await createHarness({ - tools: [waitTool], + const waiting = await createWaitingHarness({ extensionFactories: [ (pi) => { extensionApi = pi; }, ], }); + const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting; harnesses.push(harness); harness.setResponses([ fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }), (context) => { const sawSteer = context.messages.some( - (message) => - message.role === "user" && - typeof message.content !== "string" && - message.content.some((part) => part.type === "text" && part.text === "steer now"), + (message) => message.role === "user" && getMessageText(message) === "steer now", ); return fauxAssistantMessage(sawSteer ? "saw steer" : "missing steer"); }, ]); - const sawToolStart = new Promise((resolve) => { - const unsubscribe = harness.session.subscribe((event) => { - if (event.type === "tool_execution_start" && event.toolName === "wait") { - unsubscribe(); - resolve(); - } - }); - }); - - const promptPromise = harness.session.prompt("start"); - await sawToolStart; + await waitForToolStart; await new Promise((resolve) => setTimeout(resolve, 0)); extensionApi?.sendUserMessage("steer now", { deliverAs: "steer" }); - releaseToolExecution?.(); + releaseToolExecution(); await promptPromise; expect(getUserTexts(harness)).toEqual(["start", "steer now"]); expect(getAssistantTexts(harness)).toContain("saw steer"); }); - it("treats queued slash-command follow-ups as raw user text instead of dispatching the command", async () => { - let extensionApi: ExtensionAPI | undefined; - const commandRuns: string[] = []; - let releaseToolExecution: (() => void) | undefined; - const toolRelease = new Promise((resolve) => { - releaseToolExecution = resolve; - }); - const waitTool: AgentTool = { - name: "wait", - label: "Wait", - description: "Wait for the test to release execution", - parameters: Type.Object({}), - execute: async () => { - await toolRelease; - return { - content: [{ type: "text", text: "released" }], - details: {}, - }; + it("delivers follow-up messages only after the current run finishes", async () => { + const waiting = await createWaitingHarness(); + const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting; + harnesses.push(harness); + const assistantSeenBeforeFollowUp: string[] = []; + + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }), + (context) => { + assistantSeenBeforeFollowUp.push( + ...context.messages + .filter((message) => message.role === "assistant") + .map((message) => + message.content + .filter((part): part is { type: "text"; text: string } => part.type === "text") + .map((part) => part.text) + .join("\n"), + ), + ); + return fauxAssistantMessage("follow-up response"); }, - }; + ]); + + await waitForToolStart; + await harness.session.followUp("after current run"); + releaseToolExecution(); + await promptPromise; + + expect(getUserTexts(harness)).toEqual(["start", "after current run"]); + expect(assistantSeenBeforeFollowUp).toContain(""); + expect(getAssistantTexts(harness)).toContain("follow-up response"); + }); + + it("delivers multiple steering messages in order in one-at-a-time mode", async () => { + const waiting = await createWaitingHarness(); + const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting; + harnesses.push(harness); + + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }), + fauxAssistantMessage("handled steer 1"), + fauxAssistantMessage("handled steer 2"), + ]); + + await waitForToolStart; + await harness.session.steer("steer 1"); + await harness.session.steer("steer 2"); + releaseToolExecution(); + await promptPromise; + + expect(getUserTexts(harness)).toEqual(["start", "steer 1", "steer 2"]); + expect(getAssistantTexts(harness)).toEqual(["", "handled steer 1", "handled steer 2"]); + }); + + it("delivers multiple follow-up messages in order in one-at-a-time mode", async () => { + const waiting = await createWaitingHarness(); + const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting; + harnesses.push(harness); + + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }), + fauxAssistantMessage("original turn complete"), + fauxAssistantMessage("handled follow-up 1"), + fauxAssistantMessage("handled follow-up 2"), + ]); + + await waitForToolStart; + await harness.session.followUp("follow-up 1"); + await harness.session.followUp("follow-up 2"); + releaseToolExecution(); + await promptPromise; + + expect(getUserTexts(harness)).toEqual(["start", "follow-up 1", "follow-up 2"]); + expect(getAssistantTexts(harness)).toEqual([ + "", + "original turn complete", + "handled follow-up 1", + "handled follow-up 2", + ]); + }); + + it("delivers all steering messages in one batch in all mode", async () => { + const waiting = await createWaitingHarness(); + const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting; + harnesses.push(harness); + harness.session.setSteeringMode("all"); + let batchedUserMessages: string[] = []; + + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }), + (context) => { + batchedUserMessages = context.messages + .filter((message) => message.role === "user") + .map((message) => getMessageText(message)); + return fauxAssistantMessage("batched steer response"); + }, + ]); + + await waitForToolStart; + await harness.session.steer("steer 1"); + await harness.session.steer("steer 2"); + releaseToolExecution(); + await promptPromise; + + expect(batchedUserMessages).toEqual(["start", "steer 1", "steer 2"]); + expect(getAssistantTexts(harness)).toEqual(["", "batched steer response"]); + }); + + it("delivers all follow-up messages in one batch in all mode", async () => { + const waiting = await createWaitingHarness(); + const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting; + harnesses.push(harness); + harness.session.setFollowUpMode("all"); + let batchedUserMessages: string[] = []; + + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }), + fauxAssistantMessage("original turn complete"), + (context) => { + batchedUserMessages = context.messages + .filter((message) => message.role === "user") + .map((message) => getMessageText(message)); + return fauxAssistantMessage("batched follow-up response"); + }, + ]); + + await waitForToolStart; + await harness.session.followUp("follow-up 1"); + await harness.session.followUp("follow-up 2"); + releaseToolExecution(); + await promptPromise; + + expect(batchedUserMessages).toEqual(["start", "follow-up 1", "follow-up 2"]); + expect(getAssistantTexts(harness)).toEqual(["", "original turn complete", "batched follow-up response"]); + }); + + it("queues custom messages with deliverAs steer while streaming", async () => { + const waiting = await createWaitingHarness(); + const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting; + harnesses.push(harness); + let sawCustomMessage = false; + + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }), + (context) => { + sawCustomMessage = context.messages.some( + (message) => + message.role === "user" && + typeof message.content !== "string" && + message.content.some((part) => part.type === "text" && part.text === "steer custom"), + ); + return fauxAssistantMessage("done"); + }, + ]); + + await waitForToolStart; + await harness.session.sendCustomMessage( + { customType: "queue-test", content: "steer custom", display: true, details: { value: 1 } }, + { deliverAs: "steer" }, + ); + releaseToolExecution(); + await promptPromise; + + expect(sawCustomMessage).toBe(true); + expect( + harness.session.messages.some((message) => message.role === "custom" && message.customType === "queue-test"), + ).toBe(true); + }); + + it("queues custom messages with deliverAs followUp while streaming", async () => { + const waiting = await createWaitingHarness(); + const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting; + harnesses.push(harness); + let sawCustomMessage = false; + + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }), + fauxAssistantMessage("original turn complete"), + (context) => { + sawCustomMessage = context.messages.some( + (message) => + message.role === "user" && + typeof message.content !== "string" && + message.content.some((part) => part.type === "text" && part.text === "follow-up custom"), + ); + return fauxAssistantMessage("done"); + }, + ]); + + await waitForToolStart; + await harness.session.sendCustomMessage( + { customType: "queue-test", content: "follow-up custom", display: true, details: { value: 1 } }, + { deliverAs: "followUp" }, + ); + releaseToolExecution(); + await promptPromise; + + expect(sawCustomMessage).toBe(true); + expect( + harness.session.messages.some((message) => message.role === "custom" && message.customType === "queue-test"), + ).toBe(true); + }); + + it("injects nextTurn custom messages into the next prompt", async () => { + const harness = await createHarness(); + harnesses.push(harness); + let sawCustomMessage = false; + + await harness.session.sendCustomMessage( + { customType: "next-turn", content: "carry this", display: true, details: {} }, + { deliverAs: "nextTurn" }, + ); + + harness.setResponses([ + (context) => { + sawCustomMessage = context.messages.some( + (message) => + message.role === "user" && + typeof message.content !== "string" && + message.content.some((part) => part.type === "text" && part.text === "carry this"), + ); + return fauxAssistantMessage("done"); + }, + ]); + + await harness.session.prompt("normal prompt"); + + expect(sawCustomMessage).toBe(true); + expect(harness.session.messages.map((message) => message.role)).toEqual(["user", "custom", "assistant"]); + }); + + it("updates pendingMessageCount and removes queued text before message_start is emitted", async () => { + const waiting = await createWaitingHarness(); + const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting; + harnesses.push(harness); + const countsAtQueuedMessageStart: number[] = []; + + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }), + fauxAssistantMessage("done"), + ]); + + harness.session.subscribe((event) => { + if ( + event.type === "message_start" && + event.message.role === "user" && + getMessageText(event.message) === "queued" + ) { + countsAtQueuedMessageStart.push(harness.session.pendingMessageCount); + } + }); + + await waitForToolStart; + await harness.session.steer("queued"); + expect(harness.session.pendingMessageCount).toBe(1); + releaseToolExecution(); + await promptPromise; + + expect(countsAtQueuedMessageStart).toEqual([0]); + expect(harness.session.pendingMessageCount).toBe(0); + }); + + it("throws when queueing an extension command with steer", async () => { const harness = await createHarness({ - tools: [waitTool], extensionFactories: [ (pi) => { - extensionApi = pi; pi.registerCommand("testcmd", { description: "Test command", - handler: async (args) => { - commandRuns.push(args); - }, + handler: async () => {}, }); }, ], }); harnesses.push(harness); - harness.setResponses([ - fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }), - fauxAssistantMessage("first turn complete"), - fauxAssistantMessage("queued follow-up handled by model"), - ]); + await expect(harness.session.steer("/testcmd queued")).rejects.toThrow( + 'Extension command "/testcmd" cannot be queued. Use prompt() or execute the command when not streaming.', + ); + }); - const sawToolStart = new Promise((resolve) => { - const unsubscribe = harness.session.subscribe((event) => { - if (event.type === "tool_execution_start" && event.toolName === "wait") { - unsubscribe(); - resolve(); - } - }); + it("throws when queueing an extension command with followUp", async () => { + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.registerCommand("testcmd", { + description: "Test command", + handler: async () => {}, + }); + }, + ], }); + harnesses.push(harness); - const promptPromise = harness.session.prompt("start"); - await sawToolStart; - await new Promise((resolve) => setTimeout(resolve, 0)); - - extensionApi?.sendUserMessage("/testcmd queued", { deliverAs: "followUp" }); - releaseToolExecution?.(); - await promptPromise; - - expect(commandRuns).toEqual([]); - expect(getUserTexts(harness)).toEqual(["start", "/testcmd queued"]); - expect(getAssistantTexts(harness)).toContain("queued follow-up handled by model"); + await expect(harness.session.followUp("/testcmd queued")).rejects.toThrow( + 'Extension command "/testcmd" cannot be queued. Use prompt() or execute the command when not streaming.', + ); }); }); diff --git a/packages/coding-agent/test/suite/agent-session-retry-events.test.ts b/packages/coding-agent/test/suite/agent-session-retry-events.test.ts new file mode 100644 index 00000000..7da3c07e --- /dev/null +++ b/packages/coding-agent/test/suite/agent-session-retry-events.test.ts @@ -0,0 +1,358 @@ +import type { AgentTool } from "@mariozechner/pi-agent-core"; +import { fauxAssistantMessage, fauxThinking, fauxToolCall } from "@mariozechner/pi-ai"; +import { Type } from "@sinclair/typebox"; +import { afterEach, describe, expect, it } from "vitest"; +import { createHarness, type Harness } from "./harness.js"; + +function normalizeEventOrder(events: Harness["events"]): string[] { + const normalized: string[] = []; + for (const event of events) { + const label = + event.type === "message_start" || event.type === "message_end" + ? `${event.type}:${event.message.role}` + : event.type === "tool_execution_start" || event.type === "tool_execution_end" + ? `${event.type}:${event.toolName}` + : event.type; + if (label === "message_update" && normalized[normalized.length - 1] === "message_update") { + continue; + } + normalized.push(label); + } + return normalized; +} + +describe("AgentSession retry and event characterization", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + while (harnesses.length > 0) { + harnesses.pop()?.cleanup(); + } + }); + + it("retries after a transient error and succeeds", async () => { + const harness = await createHarness({ settings: { retry: { enabled: true, maxRetries: 3, baseDelayMs: 1 } } }); + harnesses.push(harness); + const retryEvents: string[] = []; + harness.session.subscribe((event) => { + if (event.type === "auto_retry_start") retryEvents.push(`start:${event.attempt}`); + if (event.type === "auto_retry_end") retryEvents.push(`end:${event.success}`); + }); + + harness.setResponses([ + fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }), + fauxAssistantMessage("recovered"), + ]); + + await harness.session.prompt("test"); + + expect(retryEvents).toEqual(["start:1", "end:true"]); + expect(harness.faux.state.callCount).toBe(2); + expect(harness.session.isRetrying).toBe(false); + }); + + it("retries multiple transient failures and succeeds on the final attempt", async () => { + const harness = await createHarness({ settings: { retry: { enabled: true, maxRetries: 3, baseDelayMs: 1 } } }); + harnesses.push(harness); + const retryEvents: string[] = []; + harness.session.subscribe((event) => { + if (event.type === "auto_retry_start") retryEvents.push(`start:${event.attempt}`); + if (event.type === "auto_retry_end") retryEvents.push(`end:${event.success}`); + }); + + harness.setResponses([ + fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }), + fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }), + fauxAssistantMessage("success"), + ]); + + await harness.session.prompt("test"); + + expect(retryEvents).toEqual(["start:1", "start:2", "end:true"]); + expect(harness.faux.state.callCount).toBe(3); + }); + + it("exhausts max retries and emits a failure event", async () => { + const harness = await createHarness({ settings: { retry: { enabled: true, maxRetries: 2, baseDelayMs: 1 } } }); + harnesses.push(harness); + const retryEvents: string[] = []; + harness.session.subscribe((event) => { + if (event.type === "auto_retry_start") retryEvents.push(`start:${event.attempt}`); + if (event.type === "auto_retry_end") retryEvents.push(`end:${event.success}`); + }); + + harness.setResponses([ + fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }), + fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }), + fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }), + ]); + + await harness.session.prompt("test"); + + expect(retryEvents).toEqual(["start:1", "start:2", "end:false"]); + expect(harness.faux.state.callCount).toBe(3); + expect(harness.session.isRetrying).toBe(false); + }); + + it("prompt waits for retry completion even when assistant message_end handling is delayed", async () => { + const harness = await createHarness({ + settings: { retry: { enabled: true, maxRetries: 3, baseDelayMs: 1 } }, + extensionFactories: [ + (pi) => { + pi.on("message_end", async (event) => { + if (event.message.role === "assistant") { + await new Promise((resolve) => setTimeout(resolve, 40)); + } + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }), + fauxAssistantMessage("recovered"), + ]); + + await harness.session.prompt("test"); + + expect(harness.faux.state.callCount).toBe(2); + expect(harness.session.isRetrying).toBe(false); + }); + + it("does not retry when retry is disabled", async () => { + const harness = await createHarness({ settings: { retry: { enabled: false } } }); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" })]); + + await harness.session.prompt("test"); + + expect(harness.faux.state.callCount).toBe(1); + expect(harness.eventsOfType("auto_retry_start")).toEqual([]); + }); + + it("does not retry non-retryable errors", async () => { + const harness = await createHarness({ settings: { retry: { enabled: true, maxRetries: 3, baseDelayMs: 1 } } }); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("", { stopReason: "error", errorMessage: "invalid_api_key" })]); + + await harness.session.prompt("test"); + + expect(harness.faux.state.callCount).toBe(1); + expect(harness.eventsOfType("auto_retry_start")).toEqual([]); + }); + + it("cancels retry sleep when abortRetry is called", async () => { + const harness = await createHarness({ settings: { retry: { enabled: true, maxRetries: 3, baseDelayMs: 100 } } }); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" })]); + + const sawRetryStart = new Promise((resolve) => { + const unsubscribe = harness.session.subscribe((event) => { + if (event.type === "auto_retry_start") { + unsubscribe(); + resolve(); + } + }); + }); + + const promptPromise = harness.session.prompt("test"); + await sawRetryStart; + harness.session.abortRetry(); + await promptPromise; + + expect(harness.session.isRetrying).toBe(false); + expect(harness.eventsOfType("auto_retry_end").map((event) => event.finalError)).toContain("Retry cancelled"); + expect(harness.faux.state.callCount).toBe(1); + }); + + it("waits for the full loop when retry recovery produces tool calls", async () => { + const toolRuns: string[] = []; + const echoTool: AgentTool = { + name: "echo", + label: "Echo", + description: "Echo text back", + parameters: Type.Object({ text: Type.String() }), + execute: async (_toolCallId, params) => { + const text = typeof params === "object" && params !== null && "text" in params ? String(params.text) : ""; + toolRuns.push(text); + return { content: [{ type: "text", text: `echo:${text}` }], details: { text } }; + }, + }; + const harness = await createHarness({ + tools: [echoTool], + settings: { retry: { enabled: true, maxRetries: 3, baseDelayMs: 1 } }, + }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage("", { stopReason: "error", errorMessage: "overloaded_error" }), + fauxAssistantMessage([fauxToolCall("echo", { text: "hello" })], { stopReason: "toolUse" }), + fauxAssistantMessage("final answer"), + ]); + + await harness.session.prompt("test"); + + expect(harness.faux.state.callCount).toBe(3); + expect(toolRuns).toEqual(["hello"]); + expect(harness.session.isStreaming).toBe(false); + await harness.session.prompt("follow-up"); + expect(harness.faux.state.callCount).toBe(4); + }); + + it("emits extension events before public event subscribers", async () => { + const order: string[] = []; + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("message_start", async (event) => { + order.push(`extension:${event.type}:${event.message.role}`); + }); + pi.on("message_end", async (event) => { + order.push(`extension:${event.type}:${event.message.role}`); + }); + }, + ], + }); + harnesses.push(harness); + harness.session.subscribe((event) => { + if (event.type === "message_start" || event.type === "message_end") { + order.push(`public:${event.type}:${event.message.role}`); + } + }); + harness.setResponses([fauxAssistantMessage("done")]); + + await harness.session.prompt("hi"); + + expect(order).toEqual([ + "extension:message_start:user", + "public:message_start:user", + "extension:message_end:user", + "public:message_end:user", + "extension:message_start:assistant", + "public:message_start:assistant", + "extension:message_end:assistant", + "public:message_end:assistant", + ]); + }); + + it("emits the expected event order for a single prompt", async () => { + const harness = await createHarness(); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("hello")]); + + await harness.session.prompt("hi"); + + expect(normalizeEventOrder(harness.events)).toEqual([ + "agent_start", + "turn_start", + "message_start:user", + "message_end:user", + "message_start:assistant", + "message_update", + "message_end:assistant", + "turn_end", + "agent_end", + ]); + }); + + it("emits the expected event order for a tool call turn", async () => { + const toolRuns: string[] = []; + const echoTool: AgentTool = { + name: "echo", + label: "Echo", + description: "Echo text back", + parameters: Type.Object({ text: Type.String() }), + execute: async (_toolCallId, params) => { + const text = typeof params === "object" && params !== null && "text" in params ? String(params.text) : ""; + toolRuns.push(text); + return { content: [{ type: "text", text: `echo:${text}` }], details: { text } }; + }, + }; + const harness = await createHarness({ tools: [echoTool] }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage([fauxToolCall("echo", { text: "hello" })], { stopReason: "toolUse" }), + fauxAssistantMessage("done"), + ]); + + await harness.session.prompt("hi"); + + expect(toolRuns).toEqual(["hello"]); + expect(normalizeEventOrder(harness.events)).toEqual([ + "agent_start", + "turn_start", + "message_start:user", + "message_end:user", + "message_start:assistant", + "message_update", + "message_end:assistant", + "tool_execution_start:echo", + "tool_execution_end:echo", + "message_start:toolResult", + "message_end:toolResult", + "turn_end", + "turn_start", + "message_start:assistant", + "message_update", + "message_end:assistant", + "turn_end", + "agent_end", + ]); + }); + + it("emits streaming deltas for text, thinking, and tool calls in message_update events", async () => { + const harness = await createHarness(); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage( + [fauxThinking("plan"), { type: "text", text: "answer" }, fauxToolCall("echo", { text: "hello" })], + { + stopReason: "toolUse", + }, + ), + ]); + + await harness.session.prompt("hi").catch(() => {}); + + const updateTypes = harness.eventsOfType("message_update").map((event) => event.assistantMessageEvent.type); + expect(updateTypes).toContain("thinking_delta"); + expect(updateTypes).toContain("text_delta"); + expect(updateTypes).toContain("toolcall_delta"); + }); + + it("emits agent_end for error responses", async () => { + const harness = await createHarness(); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("", { stopReason: "error", errorMessage: "broken" })]); + + await harness.session.prompt("hi"); + + expect(harness.events[harness.events.length - 1]?.type).toBe("agent_end"); + }); + + it("emits agent_end for aborted runs and persists the aborted assistant message", async () => { + const harness = await createHarness(); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("x".repeat(20_000))]); + + const sawMessageUpdate = new Promise((resolve) => { + const unsubscribe = harness.session.subscribe((event) => { + if (event.type === "message_update") { + unsubscribe(); + resolve(); + } + }); + }); + + const promptPromise = harness.session.prompt("hi"); + await sawMessageUpdate; + await harness.session.abort(); + await promptPromise; + + expect(harness.events[harness.events.length - 1]?.type).toBe("agent_end"); + const lastMessage = harness.session.messages[harness.session.messages.length - 1]; + expect(lastMessage?.role).toBe("assistant"); + if (lastMessage?.role === "assistant") { + expect(lastMessage.stopReason).toBe("aborted"); + } + }); +}); diff --git a/packages/coding-agent/test/suite/agent-session-runtime-host.test.ts b/packages/coding-agent/test/suite/agent-session-runtime-host.test.ts new file mode 100644 index 00000000..b28628c5 --- /dev/null +++ b/packages/coding-agent/test/suite/agent-session-runtime-host.test.ts @@ -0,0 +1,321 @@ +import { existsSync, mkdirSync, realpathSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { fauxAssistantMessage, registerFauxProvider } from "@mariozechner/pi-ai"; +import { afterEach, describe, expect, it } from "vitest"; +import { + type AgentSessionRuntimeBootstrap, + AgentSessionRuntimeHost, + createAgentSessionRuntime, +} from "../../src/core/agent-session-runtime.js"; +import { AuthStorage } from "../../src/core/auth-storage.js"; +import { SessionManager } from "../../src/core/session-manager.js"; +import type { + ExtensionFactory, + SessionBeforeForkEvent, + SessionBeforeSwitchEvent, + SessionStartEvent, +} from "../../src/index.js"; + +type RecordedSessionEvent = SessionBeforeSwitchEvent | SessionBeforeForkEvent | SessionStartEvent; + +describe("AgentSessionRuntimeHost characterization", () => { + const cleanups: Array<() => Promise | void> = []; + + afterEach(async () => { + while (cleanups.length > 0) { + await cleanups.pop()?.(); + } + process.chdir(tmpdir()); + }); + + async function createRuntimeHost( + extensionFactory: ExtensionFactory, + options?: { cwd?: string; bootstrapModel?: boolean; bootstrapThinkingLevel?: boolean }, + ) { + const tempDir = + options?.cwd ?? join(tmpdir(), `pi-runtime-suite-${Date.now()}-${Math.random().toString(36).slice(2)}`); + mkdirSync(tempDir, { recursive: true }); + + const faux = registerFauxProvider({ + models: [ + { id: "faux-1", reasoning: true }, + { id: "faux-2", reasoning: false }, + ], + }); + faux.setResponses([fauxAssistantMessage("one"), fauxAssistantMessage("two"), fauxAssistantMessage("three")]); + + const authStorage = AuthStorage.inMemory(); + authStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key"); + + const bootstrap: AgentSessionRuntimeBootstrap = { + agentDir: tempDir, + authStorage, + model: options?.bootstrapModel === false ? undefined : faux.getModel(), + thinkingLevel: options?.bootstrapThinkingLevel === false ? undefined : undefined, + resourceLoader: { + extensionFactories: [ + (pi) => { + pi.registerProvider(faux.getModel().provider, { + baseUrl: faux.getModel().baseUrl, + apiKey: "faux-key", + api: faux.api, + models: faux.models.map((registeredModel) => ({ + id: registeredModel.id, + name: registeredModel.name, + api: registeredModel.api, + reasoning: registeredModel.reasoning, + input: registeredModel.input, + cost: registeredModel.cost, + contextWindow: registeredModel.contextWindow, + maxTokens: registeredModel.maxTokens, + })), + }); + extensionFactory(pi); + }, + ], + noSkills: true, + noPromptTemplates: true, + noThemes: true, + }, + }; + const runtime = await createAgentSessionRuntime(bootstrap, { + cwd: tempDir, + sessionManager: SessionManager.create(tempDir), + }); + const runtimeHost = new AgentSessionRuntimeHost(bootstrap, runtime); + await runtimeHost.session.bindExtensions({}); + + cleanups.push(async () => { + await runtimeHost.dispose(); + faux.unregister(); + if (existsSync(tempDir)) { + rmSync(tempDir, { recursive: true, force: true }); + } + }); + + return { runtimeHost, faux, tempDir }; + } + + it("emits session_before_switch and session_start for new and resume flows", async () => { + const events: RecordedSessionEvent[] = []; + const { runtimeHost } = await createRuntimeHost((pi) => { + pi.on("session_before_switch", (event) => { + events.push(event); + }); + pi.on("session_start", (event) => { + events.push(event); + }); + }); + + expect(events).toEqual([{ type: "session_start", reason: "startup" }]); + events.length = 0; + + await runtimeHost.session.prompt("hello"); + const originalSessionFile = runtimeHost.session.sessionFile; + const originalSession = runtimeHost.session; + + const newSessionResult = await runtimeHost.newSession(); + expect(newSessionResult.cancelled).toBe(false); + await runtimeHost.session.bindExtensions({}); + expect(runtimeHost.session).not.toBe(originalSession); + expect(runtimeHost.session.messages).toEqual([]); + expect(events).toEqual([ + { type: "session_before_switch", reason: "new", targetSessionFile: undefined }, + { type: "session_start", reason: "new", previousSessionFile: originalSessionFile }, + ]); + + events.length = 0; + const secondSessionFile = runtimeHost.session.sessionFile; + + const switchResult = await runtimeHost.switchSession(originalSessionFile!); + expect(switchResult.cancelled).toBe(false); + await runtimeHost.session.bindExtensions({}); + expect(events).toEqual([ + { type: "session_before_switch", reason: "resume", targetSessionFile: originalSessionFile }, + { type: "session_start", reason: "resume", previousSessionFile: secondSessionFile }, + ]); + }); + + it("honors session_before_switch cancellation for new and resume", async () => { + const events: RecordedSessionEvent[] = []; + let cancelReason: "new" | "resume" | undefined; + const { runtimeHost } = await createRuntimeHost((pi) => { + pi.on("session_before_switch", (event) => { + events.push(event); + if (event.reason === cancelReason) { + return { cancel: true }; + } + }); + pi.on("session_start", (event) => { + events.push(event); + }); + }); + + await runtimeHost.session.prompt("hello"); + const originalSessionFile = runtimeHost.session.sessionFile; + + cancelReason = "new"; + const newResult = await runtimeHost.newSession(); + expect(newResult.cancelled).toBe(true); + expect(runtimeHost.session.sessionFile).toBe(originalSessionFile); + + events.length = 0; + const otherDir = join(tmpdir(), `pi-runtime-other-${Date.now()}-${Math.random().toString(36).slice(2)}`); + mkdirSync(otherDir, { recursive: true }); + const otherSession = SessionManager.create(otherDir); + otherSession.appendMessage({ role: "user", content: [{ type: "text", text: "other" }], timestamp: Date.now() }); + const otherSessionFile = otherSession.getSessionFile(); + cancelReason = "resume"; + const resumeResult = await runtimeHost.switchSession(otherSessionFile!); + expect(resumeResult.cancelled).toBe(true); + expect(runtimeHost.session.sessionFile).toBe(originalSessionFile); + }); + + it("emits session_before_fork and session_start and honors cancellation", async () => { + const events: RecordedSessionEvent[] = []; + let cancelNextFork = false; + const { runtimeHost } = await createRuntimeHost((pi) => { + pi.on("session_before_fork", (event) => { + events.push(event); + if (cancelNextFork) { + cancelNextFork = false; + return { cancel: true }; + } + }); + pi.on("session_start", (event) => { + events.push(event); + }); + }); + + events.length = 0; + await runtimeHost.session.prompt("hello"); + const userMessage = runtimeHost.session.getUserMessagesForForking()[0]!; + const previousSessionFile = runtimeHost.session.sessionFile; + + const successResult = await runtimeHost.fork(userMessage.entryId); + expect(successResult.cancelled).toBe(false); + expect(successResult.selectedText).toBe("hello"); + await runtimeHost.session.bindExtensions({}); + expect(events).toEqual([ + { type: "session_before_fork", entryId: userMessage.entryId }, + { type: "session_start", reason: "fork", previousSessionFile }, + ]); + + events.length = 0; + cancelNextFork = true; + const cancelResult = await runtimeHost.fork(userMessage.entryId); + expect(cancelResult).toEqual({ cancelled: true }); + expect(events).toEqual([{ type: "session_before_fork", entryId: userMessage.entryId }]); + }); + + it("throws when forking with an invalid entry id", async () => { + const { runtimeHost } = await createRuntimeHost(() => {}); + await expect(runtimeHost.fork("missing-entry")).rejects.toThrow("Invalid entry ID for forking"); + }); + + it("updates process.cwd() on cross-cwd session replacement", async () => { + const firstDir = join(tmpdir(), `pi-runtime-cwd-a-${Date.now()}-${Math.random().toString(36).slice(2)}`); + const secondDir = join(tmpdir(), `pi-runtime-cwd-b-${Date.now()}-${Math.random().toString(36).slice(2)}`); + mkdirSync(firstDir, { recursive: true }); + mkdirSync(secondDir, { recursive: true }); + const { runtimeHost, faux, tempDir } = await createRuntimeHost(() => {}, { cwd: firstDir }); + const otherAuthStorage = AuthStorage.inMemory(); + otherAuthStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key"); + const otherRuntime = await createAgentSessionRuntime( + { + agentDir: tempDir, + authStorage: otherAuthStorage, + resourceLoader: { + extensionFactories: [ + (pi) => { + pi.registerProvider(faux.getModel().provider, { + baseUrl: faux.getModel().baseUrl, + apiKey: "faux-key", + api: faux.api, + models: faux.models.map((registeredModel) => ({ + id: registeredModel.id, + name: registeredModel.name, + api: registeredModel.api, + reasoning: registeredModel.reasoning, + input: registeredModel.input, + cost: registeredModel.cost, + contextWindow: registeredModel.contextWindow, + maxTokens: registeredModel.maxTokens, + })), + }); + }, + ], + noSkills: true, + noPromptTemplates: true, + noThemes: true, + }, + }, + { cwd: secondDir, sessionManager: SessionManager.create(secondDir) }, + ); + cleanups.push(async () => { + otherRuntime.session.dispose(); + }); + await otherRuntime.session.prompt("other"); + const otherSessionFile = otherRuntime.session.sessionFile!; + + await runtimeHost.switchSession(otherSessionFile); + + expect(realpathSync(process.cwd())).toBe(realpathSync(secondDir)); + expect(realpathSync(runtimeHost.session.sessionManager.getCwd())).toBe(realpathSync(secondDir)); + }); + + it("restores model and thinking state from the destination session", async () => { + const { runtimeHost, faux, tempDir } = await createRuntimeHost(() => {}, { + bootstrapModel: false, + bootstrapThinkingLevel: false, + }); + const otherDir = join(tempDir, "other"); + mkdirSync(otherDir, { recursive: true }); + const otherAuthStorage = AuthStorage.inMemory(); + otherAuthStorage.setRuntimeApiKey(faux.getModel().provider, "faux-key"); + const otherRuntime = await createAgentSessionRuntime( + { + agentDir: tempDir, + authStorage: otherAuthStorage, + resourceLoader: { + extensionFactories: [ + (pi) => { + pi.registerProvider(faux.getModel().provider, { + baseUrl: faux.getModel().baseUrl, + apiKey: "faux-key", + api: faux.api, + models: faux.models.map((registeredModel) => ({ + id: registeredModel.id, + name: registeredModel.name, + api: registeredModel.api, + reasoning: registeredModel.reasoning, + input: registeredModel.input, + cost: registeredModel.cost, + contextWindow: registeredModel.contextWindow, + maxTokens: registeredModel.maxTokens, + })), + }); + }, + ], + noSkills: true, + noPromptTemplates: true, + noThemes: true, + }, + }, + { cwd: otherDir, sessionManager: SessionManager.create(otherDir) }, + ); + cleanups.push(async () => { + otherRuntime.session.dispose(); + }); + await otherRuntime.session.setModel(faux.getModel("faux-2")!); + otherRuntime.session.setThinkingLevel("off"); + await otherRuntime.session.prompt("hello"); + const targetSessionFile = otherRuntime.session.sessionFile!; + + await runtimeHost.switchSession(targetSessionFile); + + expect(runtimeHost.session.model?.id).toBe("faux-2"); + expect(runtimeHost.session.thinkingLevel).toBe("off"); + }); +}); diff --git a/packages/coding-agent/test/suite/harness.ts b/packages/coding-agent/test/suite/harness.ts index 9262f013..c1cdb81b 100644 --- a/packages/coding-agent/test/suite/harness.ts +++ b/packages/coding-agent/test/suite/harness.ts @@ -5,12 +5,14 @@ import { existsSync, mkdirSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; -import type { AgentTool } from "@mariozechner/pi-agent-core"; +import type { AgentMessage, AgentTool } from "@mariozechner/pi-agent-core"; import { Agent } from "@mariozechner/pi-agent-core"; import type { FauxModelDefinition, FauxProviderRegistration, FauxResponseStep, Model } from "@mariozechner/pi-ai"; import { registerFauxProvider } from "@mariozechner/pi-ai"; import { AgentSession, type AgentSessionEvent } from "../../src/core/agent-session.js"; import { AuthStorage } from "../../src/core/auth-storage.js"; +import type { ExtensionRunner } from "../../src/core/extensions/index.js"; +import { convertToLlm } from "../../src/core/messages.js"; import { ModelRegistry } from "../../src/core/model-registry.js"; import { SessionManager } from "../../src/core/session-manager.js"; import type { Settings } from "../../src/core/settings-manager.js"; @@ -22,6 +24,37 @@ import { createTestResourceLoader, } from "../utilities.js"; +type MessageTextPart = { type: "text"; text: string }; + +export function getMessageText(message: unknown): string { + if (!message || typeof message !== "object" || !("content" in message)) { + return ""; + } + const content = (message as { content?: string | Array<{ type: string; text?: string }> }).content; + if (content === undefined) { + return ""; + } + if (typeof content === "string") { + return content; + } + return content + .filter((part): part is MessageTextPart => part.type === "text") + .map((part) => part.text) + .join("\n"); +} + +export function getUserTexts(harness: Harness): string[] { + return harness.session.messages + .filter((message) => message.role === "user") + .map((message) => getMessageText(message)); +} + +export function getAssistantTexts(harness: Harness): string[] { + return harness.session.messages + .filter((message) => message.role === "assistant") + .map((message) => getMessageText(message)); +} + export interface HarnessOptions { models?: FauxModelDefinition[]; settings?: Partial; @@ -29,12 +62,14 @@ export interface HarnessOptions { tools?: AgentTool[]; resourceLoader?: ResourceLoader; extensionFactories?: Array; + withConfiguredAuth?: boolean; } export interface Harness { session: AgentSession; sessionManager: SessionManager; settingsManager: SettingsManager; + authStorage: AuthStorage; faux: FauxProviderRegistration; models: [Model, ...Model[]]; getModel(): Model; @@ -62,37 +97,56 @@ export async function createHarness(options: HarnessOptions = {}): Promise [tool.name, tool])) : undefined; - - const agent = new Agent({ - getApiKey: () => "faux-key", - initialState: { - model, - systemPrompt: options.systemPrompt ?? "You are a test assistant.", - tools: [], - }, - }); + const withConfiguredAuth = options.withConfiguredAuth ?? true; + const extensionRunnerRef: { current?: ExtensionRunner } = {}; const sessionManager = SessionManager.inMemory(); const settingsManager = SettingsManager.inMemory(options.settings); const authStorage = AuthStorage.inMemory(); - authStorage.setRuntimeApiKey(model.provider, "faux-key"); + if (withConfiguredAuth) { + authStorage.setRuntimeApiKey(model.provider, "faux-key"); + } const modelRegistry = ModelRegistry.inMemory(authStorage); - modelRegistry.registerProvider(model.provider, { - baseUrl: model.baseUrl, - apiKey: "faux-key", - api: fauxProvider.api, - models: fauxProvider.models.map((registeredModel) => ({ - id: registeredModel.id, - name: registeredModel.name, - api: registeredModel.api, - reasoning: registeredModel.reasoning, - input: registeredModel.input, - cost: registeredModel.cost, - contextWindow: registeredModel.contextWindow, - maxTokens: registeredModel.maxTokens, - baseUrl: registeredModel.baseUrl, - })), + if (withConfiguredAuth) { + modelRegistry.registerProvider(model.provider, { + baseUrl: model.baseUrl, + apiKey: "faux-key", + api: fauxProvider.api, + models: fauxProvider.models.map((registeredModel) => ({ + id: registeredModel.id, + name: registeredModel.name, + api: registeredModel.api, + reasoning: registeredModel.reasoning, + input: registeredModel.input, + cost: registeredModel.cost, + contextWindow: registeredModel.contextWindow, + maxTokens: registeredModel.maxTokens, + baseUrl: registeredModel.baseUrl, + })), + }); + } + + const agent = new Agent({ + getApiKey: () => (withConfiguredAuth ? "faux-key" : undefined), + initialState: { + model, + systemPrompt: options.systemPrompt ?? "You are a test assistant.", + tools: [], + }, + convertToLlm, + onPayload: async (payload) => { + const runner = extensionRunnerRef.current; + if (!runner?.hasHandlers("before_provider_request")) { + return payload; + } + return runner.emitBeforeProviderRequest(payload); + }, + transformContext: async (messages: AgentMessage[]) => { + const runner = extensionRunnerRef.current; + if (!runner) return messages; + return runner.emitContext(messages); + }, }); const extensionsResult = options.extensionFactories ? await createTestExtensionsResult(options.extensionFactories, tempDir) @@ -108,6 +162,7 @@ export async function createHarness(options: HarnessOptions = {}): Promise { + const harnesses: Harness[] = []; + + afterEach(() => { + while (harnesses.length > 0) { + harnesses.pop()?.cleanup(); + } + }); + + it("treats extension-origin queued slash-command follow-ups as raw user text instead of dispatching the command", async () => { + let extensionApi: ExtensionAPI | undefined; + const commandRuns: string[] = []; + let releaseToolExecution: (() => void) | undefined; + const toolRelease = new Promise((resolve) => { + releaseToolExecution = resolve; + }); + const waitTool: AgentTool = { + name: "wait", + label: "Wait", + description: "Wait for the test to release execution", + parameters: Type.Object({}), + execute: async () => { + await toolRelease; + return { + content: [{ type: "text", text: "released" }], + details: {}, + }; + }, + }; + const harness = await createHarness({ + tools: [waitTool], + extensionFactories: [ + (pi) => { + extensionApi = pi; + pi.registerCommand("testcmd", { + description: "Test command", + handler: async (args) => { + commandRuns.push(args); + }, + }); + }, + ], + }); + harnesses.push(harness); + + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }), + fauxAssistantMessage("first turn complete"), + fauxAssistantMessage("queued follow-up handled by model"), + ]); + + const sawToolStart = new Promise((resolve) => { + const unsubscribe = harness.session.subscribe((event) => { + if (event.type === "tool_execution_start" && event.toolName === "wait") { + unsubscribe(); + resolve(); + } + }); + }); + + const promptPromise = harness.session.prompt("start"); + await sawToolStart; + await new Promise((resolve) => setTimeout(resolve, 0)); + + extensionApi?.sendUserMessage("/testcmd queued", { deliverAs: "followUp" }); + releaseToolExecution?.(); + await promptPromise; + + expect(commandRuns).toEqual([]); + expect(getUserTexts(harness)).toEqual(["start", "/testcmd queued"]); + expect(getAssistantTexts(harness)).toContain("queued follow-up handled by model"); + }); +});