mirror of
https://github.com/earendil-works/pi.git
synced 2026-10-21 06:09:18 +08:00
Armin Ronacher
9e05370b29
This change makes system prompt text and tool changes part of the transcript rather than silently rewriting its starting conditions. This lets Pi record when instructions changed or tools became available, restore that state after resuming or navigating branches, and preserve cached prompt prefixes where the upstream supports it.
499 lines
16 KiB
TypeScript
499 lines
16 KiB
TypeScript
import type { AgentTool } from "@earendil-works/pi-agent-core";
|
|
import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai";
|
|
import type { ExtensionAPI, InputEvent } from "@earendil-works/pi-coding-agent";
|
|
import { Type } from "typebox";
|
|
import { afterEach, describe, expect, it } from "vitest";
|
|
import { createHarness, getAssistantTexts, getMessageText, getUserTexts, type Harness } from "./harness.ts";
|
|
|
|
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 {
|
|
harness,
|
|
releaseToolExecution: () => releaseToolExecution?.(),
|
|
promptPromise: harness.session.prompt("start"),
|
|
waitForToolStart,
|
|
};
|
|
}
|
|
|
|
describe("AgentSession queue characterization", () => {
|
|
const harnesses: Harness[] = [];
|
|
|
|
afterEach(() => {
|
|
while (harnesses.length > 0) {
|
|
harnesses.pop()?.cleanup();
|
|
}
|
|
});
|
|
|
|
it("dispatches extension commands immediately when prompted while idle", 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);
|
|
|
|
await harness.session.prompt("/testcmd hello world");
|
|
|
|
expect(commandRuns).toEqual(["hello world"]);
|
|
expect(harness.getPendingResponseCount()).toBe(0);
|
|
expect(harness.session.messages).toEqual([]);
|
|
});
|
|
|
|
it("delivers extension-origin steering messages before the next LLM call", async () => {
|
|
let extensionApi: ExtensionAPI | undefined;
|
|
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" && getMessageText(message) === "steer now",
|
|
);
|
|
return fauxAssistantMessage(sawSteer ? "saw steer" : "missing steer");
|
|
},
|
|
]);
|
|
|
|
await waitForToolStart;
|
|
await new Promise((resolve) => setTimeout(resolve, 0));
|
|
|
|
extensionApi?.sendUserMessage("steer now", { deliverAs: "steer" });
|
|
releaseToolExecution();
|
|
await promptPromise;
|
|
|
|
expect(getUserTexts(harness)).toEqual(["start", "steer now"]);
|
|
expect(getAssistantTexts(harness)).toContain("saw steer");
|
|
});
|
|
|
|
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");
|
|
});
|
|
|
|
// Regression test for #8718.
|
|
it("runs direct steering and follow-up messages through input handlers", async () => {
|
|
const inputEvents: Array<Pick<InputEvent, "text" | "source" | "streamingBehavior">> = [];
|
|
const waiting = await createWaitingHarness({
|
|
extensionFactories: [
|
|
(pi) => {
|
|
pi.on("input", (event) => {
|
|
inputEvents.push({
|
|
text: event.text,
|
|
source: event.source,
|
|
streamingBehavior: event.streamingBehavior,
|
|
});
|
|
if (event.text.startsWith("handle")) return { action: "handled" };
|
|
return { action: "transform", text: `transformed: ${event.text}` };
|
|
});
|
|
},
|
|
],
|
|
});
|
|
const { harness, waitForToolStart, promptPromise, releaseToolExecution } = waiting;
|
|
harnesses.push(harness);
|
|
harness.setResponses([
|
|
fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }),
|
|
fauxAssistantMessage("steered"),
|
|
fauxAssistantMessage("followed up"),
|
|
]);
|
|
|
|
await waitForToolStart;
|
|
inputEvents.length = 0;
|
|
try {
|
|
await harness.session.steer("steer me", undefined, { source: "rpc" });
|
|
await harness.session.steer("handle steer", undefined, { source: "rpc" });
|
|
await harness.session.followUp("follow me", undefined, { source: "rpc" });
|
|
await harness.session.followUp("handle follow", undefined, { source: "rpc" });
|
|
|
|
expect(inputEvents).toEqual([
|
|
{ text: "steer me", source: "rpc", streamingBehavior: "steer" },
|
|
{ text: "handle steer", source: "rpc", streamingBehavior: "steer" },
|
|
{ text: "follow me", source: "rpc", streamingBehavior: "followUp" },
|
|
{ text: "handle follow", source: "rpc", streamingBehavior: "followUp" },
|
|
]);
|
|
expect(harness.session.getSteeringMessages()).toEqual(["transformed: steer me"]);
|
|
expect(harness.session.getFollowUpMessages()).toEqual(["transformed: follow me"]);
|
|
} finally {
|
|
releaseToolExecution();
|
|
}
|
|
await promptPromise;
|
|
});
|
|
|
|
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([
|
|
"system",
|
|
"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({
|
|
extensionFactories: [
|
|
(pi) => {
|
|
pi.registerCommand("testcmd", {
|
|
description: "Test command",
|
|
handler: async () => {},
|
|
});
|
|
},
|
|
],
|
|
});
|
|
harnesses.push(harness);
|
|
|
|
await expect(harness.session.steer("/testcmd queued")).rejects.toThrow(
|
|
'Extension command "/testcmd" cannot be queued. Use prompt() or execute the command when not streaming.',
|
|
);
|
|
});
|
|
|
|
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);
|
|
|
|
await expect(harness.session.followUp("/testcmd queued")).rejects.toThrow(
|
|
'Extension command "/testcmd" cannot be queued. Use prompt() or execute the command when not streaming.',
|
|
);
|
|
});
|
|
|
|
it("delivers follow-ups queued during agent_end", async () => {
|
|
let sent = false;
|
|
const harness = await createHarness({
|
|
extensionFactories: [
|
|
(pi: ExtensionAPI) => {
|
|
pi.on("agent_end", async () => {
|
|
if (sent) return;
|
|
sent = true;
|
|
pi.sendUserMessage("conflict report", { deliverAs: "followUp" });
|
|
});
|
|
},
|
|
],
|
|
});
|
|
harnesses.push(harness);
|
|
|
|
harness.setResponses([fauxAssistantMessage("reply"), fauxAssistantMessage("follow-up reply")]);
|
|
|
|
await harness.session.prompt("hello");
|
|
await harness.session.agent.waitForIdle();
|
|
|
|
expect(getUserTexts(harness)).toEqual(["hello", "conflict report"]);
|
|
});
|
|
});
|