shumengya/sproutclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

test(coding-agent): add session lifecycle characterization suite

Browse Source
This commit is contained in:
Mario Zechner
2026-04-01 22:59:49 +02:00
parent 957f21cea8
commit 746f770b10
9 changed files with 2408 additions and 130 deletions
Download Patch File Download Diff File Expand all files Collapse all files

242
packages/coding-agent/test/suite/agent-session-bash-persistence.test.ts Normal file
View File

@@ -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<void>((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<void>((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<void>((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");
});
});

343
packages/coding-agent/test/suite/agent-session-compaction.test.ts Normal file
View File

@@ -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<void>;
_runAutoCompaction: (reason: "overflow" | "threshold", willRetry: boolean) => Promise<void>;
};
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<any>;
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();
});
});

325
packages/coding-agent/test/suite/agent-session-model-extension.test.ts Normal file
View File

@@ -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<string>;
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"]);
});
});

323
packages/coding-agent/test/suite/agent-session-prompt.test.ts Normal file
View File

@@ -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('<skill name="test" location="');
expect(expandedPrompt).toContain("Use the skill body.");
expect(expandedPrompt).toContain("explain this");
});
it("expands prompt templates before sending the prompt", async () => {
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<void>((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<void>((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<any>;
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}.`,
);
});
});

438
packages/coding-agent/test/suite/agent-session-queue.test.ts
View File

@@ -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<void>;
waitForToolStart: Promise<void>;
}> {
let releaseToolExecution: (() => void) | undefined;
const toolRelease = new Promise<void>((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<void>((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<void>((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<void>((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<void>((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<void>((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.',
);
});
});

358
packages/coding-agent/test/suite/agent-session-retry-events.test.ts Normal file
View File

@@ -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<void>((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<void>((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");
}
});
});

321
packages/coding-agent/test/suite/agent-session-runtime-host.test.ts Normal file
View File

@@ -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> | 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");
});
});

108
packages/coding-agent/test/suite/harness.ts
View File

@@ -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<Settings>;
@@ -29,12 +62,14 @@ export interface HarnessOptions {
tools?: AgentTool[];
resourceLoader?: ResourceLoader;
extensionFactories?: Array<ExtensionFactory | CreateTestExtensionsResultInput>;
withConfiguredAuth?: boolean;
}
export interface Harness {
session: AgentSession;
sessionManager: SessionManager;
settingsManager: SettingsManager;
authStorage: AuthStorage;
faux: FauxProviderRegistration;
models: [Model<string>, ...Model<string>[]];
getModel(): Model<string>;
@@ -62,37 +97,56 @@ export async function createHarness(options: HarnessOptions = {}): Promise<Harne
fauxProvider.setResponses([]);
const model = fauxProvider.getModel();
const toolMap = options.tools ? Object.fromEntries(options.tools.map((tool) => [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<Harne
modelRegistry,
resourceLoader,
baseToolsOverride: toolMap,
extensionRunnerRef,
});
const events: AgentSessionEvent[] = [];
@@ -119,6 +174,7 @@ export async function createHarness(options: HarnessOptions = {}): Promise<Harne
session,
sessionManager,
settingsManager,
authStorage,
faux: fauxProvider,
models: fauxProvider.models,
getModel: fauxProvider.getModel,

80
packages/coding-agent/test/suite/regressions/2023-queued-slash-command-followup.test.ts Normal file
View File

@@ -0,0 +1,80 @@
import type { AgentTool } from "@mariozechner/pi-agent-core";
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, getAssistantTexts, getUserTexts, type Harness } from "../harness.js";
describe("issue #2023 queued slash-command follow-up", () => {
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<void>((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<void>((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");
});
});