/** Tests for AgentSession.mutateQueuedMessage: (lane, index, expectedText) queue edits. */
import { existsSync, mkdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { Agent } from "@earendil-works/pi-agent-core";
import {
type AssistantMessage,
type AssistantMessageEvent,
EventStream,
getModel,
type ImageContent,
} from "@earendil-works/pi-ai";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { AgentSession } from "../src/core/agent-session.js";
import { AuthStorage } from "../src/core/auth-storage.js";
import { ModelRegistry } from "../src/core/model-registry.js";
import type { QueuedMessageMutation } from "../src/core/session-action-store.js";
import { SessionManager } from "../src/core/session-manager.js";
import { SettingsManager } from "../src/core/settings-manager.js";
import { createTestResourceLoader } from "./utilities.js";
class MockAssistantStream extends EventStream<AssistantMessageEvent, AssistantMessage> {
constructor() {
super(
(event) => event.type === "done" || event.type === "error",
(event) => {
if (event.type === "done") return event.message;
if (event.type === "error") return event.error;
throw new Error("Unexpected event type");
},
);
}
}
const zeroUsage = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 };
function createAssistantMessage(text: string): AssistantMessage {
return {
role: "assistant",
content: [{ type: "text", text }],
api: "anthropic-messages",
provider: "anthropic",
model: "mock",
usage: { ...zeroUsage, totalTokens: 0, cost: zeroUsage },
stopReason: "stop",
timestamp: Date.now(),
};
}
describe("AgentSession queue mutation", () => {
let session: AgentSession;
let tempDir: string;
beforeEach(() => {
tempDir = join(tmpdir(), `pi-queue-mutation-test-${Date.now()}-${Math.random()}`);
mkdirSync(tempDir, { recursive: true });
});
afterEach(async () => {
if (session) session.dispose();
if (tempDir && existsSync(tempDir)) rmSync(tempDir, { recursive: true });
});
function createSession() {
const model = getModel("anthropic", "claude-sonnet-4-5")!;
const agent = new Agent({
getApiKey: () => "test-key",
initialState: { model, systemPrompt: "Test", tools: [] },
streamFn: (_model, _context, options) => {
const stream = new MockAssistantStream();
queueMicrotask(() => stream.push({ type: "start", partial: createAssistantMessage("") }));
options?.signal?.addEventListener("abort", () =>
stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") }),
);
return stream;
},
});
const authStorage = AuthStorage.create(join(tempDir, "auth.json"));
authStorage.setRuntimeApiKey("anthropic", "test-key");
session = new AgentSession({
agent,
sessionManager: SessionManager.inMemory(),
settingsManager: SettingsManager.create(tempDir, tempDir),
cwd: tempDir,
modelRegistry: ModelRegistry.create(authStorage, tempDir),
resourceLoader: createTestResourceLoader(),
});
return session;
}
/** Start a never-finishing turn so later submissions stay queued; returns the prompt promise. */
async function blockSession(): Promise<{ running: Promise<void> }> {
const running = session.prompt("running");
await new Promise((resolve) => setTimeout(resolve, 10));
return { running };
}
const mutate = (l: "steering" | "followUp", i: number, text: string, m: QueuedMessageMutation) =>
session.mutateQueuedMessage(l, i, text, m);
const queue = () => {
const snapshot = session.getSessionActionSnapshot();
return { steering: [...snapshot.steering], followUp: [...snapshot.followUps] };
};
const payloadOf = (text: string) =>
session.getSessionActionRecoverySnapshot().actions.find((action) => action.payload.text === text)?.payload;
const primaryContent = (text: string) => {
const payload = payloadOf(text);
if (payload?.kind !== "turn") throw new Error(`expected turn payload for ${text}`);
return payload.records[0]?.message.content;
};
it("addresses items in snapshot projection order and rejects stale addresses", async () => {
createSession();
const { running } = await blockSession();
await session.steer("s1");
await session.steer("s2");
await session.followUp("f1");
expect(queue()).toEqual({ steering: ["s1", "s2"], followUp: ["f1"] });
expect(session.queuedActionCount).toBe(3);
expect(mutate("steering", 0, "wrong", { type: "delete" })).toBe("rejected");
expect(mutate("steering", 3, "s1", { type: "delete" })).toBe("rejected");
expect(mutate("followUp", 0, "s1", { type: "delete" })).toBe("rejected"); // right text, wrong lane
expect(mutate("steering", 0, "s1", { type: "delete" })).toBe("applied");
expect(queue()).toEqual({ steering: ["s2"], followUp: ["f1"] });
expect(session.queuedActionCount).toBe(2);
await session.abort();
await running.catch(() => {});
});
it("moves within a lane, rejects boundary moves, and flips lanes to the target lane end", async () => {
createSession();
const { running } = await blockSession();
await session.steer("s1");
await session.steer("s2");
await session.followUp("f1");
expect(mutate("steering", 0, "s1", { type: "move", direction: 1 })).toBe("applied");
expect(queue().steering).toEqual(["s2", "s1"]);
expect(mutate("steering", 1, "s1", { type: "move", direction: 1 })).toBe("rejected");
expect(mutate("steering", 0, "s2", { type: "move", direction: -1 })).toBe("rejected");
expect(mutate("steering", 0, "s2", { type: "replace", text: "s2 flipped", lane: "followUp" })).toBe("applied");
expect(queue()).toEqual({ steering: ["s1"], followUp: ["f1", "s2 flipped"] });
expect(mutate("followUp", 1, "s2 flipped", { type: "replace", text: "s2 flipped", lane: "steering" })).toBe(
"applied",
);
expect(queue()).toEqual({ steering: ["s1", "s2 flipped"], followUp: ["f1"] });
await session.abort();
await running.catch(() => {});
});
it("replaces turn text in place across the image tri-state and rebuilds content and records", async () => {
createSession();
const { running } = await blockSession();
const original: ImageContent = { type: "image", data: "original", mimeType: "image/png" };
await session.prompt("edit me [image #1]", {
streamingBehavior: "steer",
images: [original],
content: [{ type: "text", text: "edit me" }, { type: "text", text: " [image #1]" }, original],
});
const edit = (expected: string, text: string, images?: ImageContent[]) =>
mutate("steering", 0, expected, { type: "replace", text, images, lane: "steering" });
// images undefined: preserve images, swap the text block, keep non-text blocks.
expect(edit("edit me [image #1]", "edited")).toBe("applied");
expect(payloadOf("edited")).toMatchObject({
images: [original],
content: [{ type: "text", text: "edited" }, original],
});
expect(primaryContent("edited")).toEqual([{ type: "text", text: "edited" }, original]);
// explicit images: content rebuilt from text + clones.
const replacement: ImageContent = { type: "image", data: "replacement", mimeType: "image/png" };
expect(edit("edited", "re-imaged", [replacement])).toBe("applied");
const reimaged = payloadOf("re-imaged");
expect(reimaged).toMatchObject({
images: [replacement],
content: [{ type: "text", text: "re-imaged" }, replacement],
});
expect(primaryContent("re-imaged")).toEqual([{ type: "text", text: "re-imaged" }, replacement]);
// empty images: clear to a bare text block.
expect(edit("re-imaged", "plain", [])).toBe("applied");
const cleared = payloadOf("plain");
expect(cleared).not.toHaveProperty("images");
expect(cleared?.kind === "turn" ? cleared.content : undefined).toEqual([{ type: "text", text: "plain" }]);
expect(primaryContent("plain")).toEqual([{ type: "text", text: "plain" }]);
await session.abort();
await running.catch(() => {});
});
it("validates session command text and protects accepted agent messages", async () => {
createSession();
const { running } = await blockSession();
await session.prompt("/compact focus on tools", { streamingBehavior: "steer" });
await session.acceptAgentMessagePrompt("agent item", { streamingBehavior: "followUp" });
expect(queue()).toEqual({ steering: ["/compact focus on tools"], followUp: ["agent item"] });
expect(
mutate("steering", 0, "/compact focus on tools", { type: "replace", text: "plain", lane: "steering" }),
).toBe("invalid");
expect(
mutate("steering", 0, "/compact focus on tools", {
type: "replace",
text: "/compact revised focus",
lane: "followUp",
}),
).toBe("applied");
expect(queue().steering).toEqual([]);
expect(payloadOf("/compact revised focus")).toMatchObject({
kind: "session_command",
command: { name: "compact", args: "revised focus" },
});
expect(mutate("followUp", 0, "agent item", { type: "replace", text: "hijacked", lane: "followUp" })).toBe(
"rejected",
);
expect(mutate("followUp", 0, "agent item", { type: "delete" })).toBe("applied");
expect(queue().followUp).toEqual(["/compact revised focus"]); // flipped command appended to the lane end
await session.abort();
await running.catch(() => {});
});
it("resumes suspended queued work after a successful replace", async () => {
createSession();
const { running } = await blockSession();
await session.followUp("f1");
await session.abort();
await running.catch(() => {});
expect(session.isQueuedWorkSuspended).toBe(true);
expect(queue().followUp).toEqual(["f1"]); // queue survives abort
expect(mutate("followUp", 0, "f1", { type: "replace", text: "f1 edited", lane: "followUp" })).toBe("applied");
expect(session.isQueuedWorkSuspended).toBe(false);
await session.abort();
});
it("resumes suspended queued work after a successful delete", async () => {
createSession();
const { running } = await blockSession();
await session.followUp("kept");
await session.followUp("deleted");
await session.abort();
await running.catch(() => {});
expect(session.isQueuedWorkSuspended).toBe(true);
expect(mutate("followUp", 1, "deleted", { type: "delete" })).toBe("applied");
expect(session.isQueuedWorkSuspended).toBe(false);
expect(queue().followUp).toEqual(["kept"]);
await session.abort();
});
it("settles the completion of a deleted promptAndWait turn", async () => {
createSession();
const { running } = await blockSession();
const waiting = session.promptAndWait("waited", { streamingBehavior: "steer" });
await new Promise((resolve) => setTimeout(resolve, 10));
expect(queue().steering).toEqual(["waited"]);
expect(mutate("steering", 0, "waited", { type: "delete" })).toBe("applied");
await expect(waiting).rejects.toThrow("Queued prompt was deleted before delivery.");
await session.abort();
await running.catch(() => {});
});
it("rejects replacing an injected custom-message turn but allows deleting it", async () => {
createSession();
const { running } = await blockSession();
const injected = (
session as unknown as {
_promptInjectedMessage(
text: string,
message: unknown,
options?: { streamingBehavior?: "steer" },
): Promise<void>;
}
)._promptInjectedMessage(
"injected",
{ role: "custom", customType: "test", content: "injected", display: true, timestamp: Date.now() },
{ streamingBehavior: "steer" },
);
await new Promise((resolve) => setTimeout(resolve, 10));
const [expected] = queue().steering;
expect(expected).toBeDefined();
expect(mutate("steering", 0, expected!, { type: "replace", text: "rewritten", lane: "steering" })).toBe(
"rejected",
);
expect(mutate("steering", 0, expected!, { type: "delete" })).toBe("applied");
await session.abort();
await Promise.allSettled([running, injected]);
});
it("updates the wake policy when flipping lanes", async () => {
createSession();
const { running } = await blockSession();
await session.followUp("f1");
expect(mutate("followUp", 0, "f1", { type: "replace", text: "f1", lane: "steering" })).toBe("applied");
const recovered = session.getSessionActionRecoverySnapshot().actions.find((a) => a.payload.text === "f1");
expect(recovered?.delivery).toBe("next_turn_boundary");
expect(recovered?.wake).toBe("on_lower_boundary");
await session.abort();
await running.catch(() => {});
});
});