Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
12 kB · 295 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296/** 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(() => {}); });});