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.
13 kB · 329 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330import type { AgentMessage } from "@earendil-works/pi-agent-core";import { describe, expect, test, vi } from "vitest";import type { AgentStatus } from "../src/core/session-manager.js";import { SettingsManager } from "../src/core/settings-manager.js";import type { ActiveSessionState } from "../src/modes/daemon/active-session-state.js";import { buildStatusContext, DaemonSessionSummarizer, parseAgentStatusResponse,} from "../src/modes/daemon/daemon-session-summarizer.js";
function userMessage(text: string): AgentMessage { return { role: "user", content: [{ type: "text", text }], timestamp: 0 } as unknown as AgentMessage;}
function assistantMessage(text: string, tools: string[] = []): AgentMessage { const content = [{ type: "text", text }, ...tools.map((name) => ({ type: "tool_use", name, id: name, input: {} }))]; return { role: "assistant", content, timestamp: 0 } as unknown as AgentMessage;}
describe("daemon session summarizer", () => { describe("parseAgentStatusResponse", () => { test("parses recap and completion verdict for an idle session", () => { const result = parseAgentStatusResponse( "<recap>Added the API reference page</recap>\n<status>COMPLETED</status>", false, ); expect(result).toEqual({ summary: "Added the API reference page", taskState: "completed" }); });
test("maps NEEDS_INPUT for idle sessions", () => { const result = parseAgentStatusResponse( "<recap>Asked which database to target</recap>\n<status>NEEDS_INPUT</status>", false, ); expect(result?.taskState).toBe("needs_input"); });
test("omits the verdict while working and ignores any status tag", () => { const result = parseAgentStatusResponse( "<recap>Refactoring token validation</recap>\n<status>COMPLETED</status>", true, ); expect(result).toEqual({ summary: "Refactoring token validation" }); });
test("falls back to needs_input on a missing or unrecognized idle verdict", () => { expect( parseAgentStatusResponse("<recap>Doing something</recap>\n<status>MAYBE</status>", false)?.taskState, ).toBe("needs_input"); expect(parseAgentStatusResponse("<recap>Doing something</recap>", false)?.taskState).toBe("needs_input"); });
test("ignores narration outside the tags and never leaks free-form text", () => { // A chatty/reasoning model that narrates instead of using tags yields no recap. expect(parseAgentStatusResponse("Investigating the failing test.", true)).toBeUndefined(); expect( parseAgentStatusResponse( "Recap: . So: <recap>Curating a niche list of Muon optimizer papers</recap>", true, ), ).toEqual({ summary: "Curating a niche list of Muon optimizer papers" }); });
test("ignores reasoning prose around the tags", () => { const text = "Let me decide. The agent finished editing.\n<recap>Updated the login handler</recap>\n<status>COMPLETED</status>"; expect(parseAgentStatusResponse(text, false)).toEqual({ summary: "Updated the login handler", taskState: "completed", }); });
test("rejects an echoed prompt template", () => { const echoed = "<recap>a present-tense clause, at most 12 words, no trailing period</recap>\n<status>COMPLETED</status>"; expect(parseAgentStatusResponse(echoed, true)).toBeUndefined(); });
test("returns undefined when no recap tag is present", () => { expect(parseAgentStatusResponse("", false)).toBeUndefined(); expect(parseAgentStatusResponse("<status>COMPLETED</status>", false)).toBeUndefined(); });
test("drops chain-of-thought that falls outside the closing recap tag", () => { const text = "<recap>Sending SSH auth retry to tcg-autoresearch-rl</recap> That's 5 words? Count: Sending(1) SSH(2) = 6 words.\n<status>NEEDS_INPUT</status>"; expect(parseAgentStatusResponse(text, true)).toEqual({ summary: "Sending SSH auth retry to tcg-autoresearch-rl", }); });
test("rejects a recap body that is nothing but counting artifacts", () => { expect(parseAgentStatusResponse("<recap>(1) word(2) count(3) = 3 words</recap>", true)).toBeUndefined(); });
test("rejects a rambling recap that blows past the word ceiling", () => { const text = "<recap>this is a very long rambling sentence that just keeps going and going well past any reasonable recap length</recap>"; expect(parseAgentStatusResponse(text, true)).toBeUndefined(); });
test("strips wrapping quotes the model adds around the recap", () => { const text = '<recap>"Wiring the recap line"</recap>\n<status>COMPLETED</status>'; expect(parseAgentStatusResponse(text, false)).toEqual({ summary: "Wiring the recap line", taskState: "completed", }); });
test("ignores an open recap tag with no close", () => { expect( parseAgentStatusResponse("<recap>Editing the parser\n<status>NEEDS_INPUT</status>", true), ).toBeUndefined(); });
test("takes the last recap tag when a draft is corrected", () => { const text = "<recap>Draft recap</recap>\n<recap>Final corrected recap</recap>"; expect(parseAgentStatusResponse(text, true)).toEqual({ summary: "Final corrected recap" }); });
test("takes the last status tag when a draft is corrected", () => { const text = "<recap>Editing the parser</recap>\n<status>NEEDS_INPUT</status>\n<status>COMPLETED</status>"; expect(parseAgentStatusResponse(text, false)?.taskState).toBe("completed"); });
test("normalizes unicode angle-bracket lookalikes around the tags", () => { // The model sometimes emits › ‹ instead of > < ; normalize so the tag still parses. const text = "‹recap›Curating a niche list of Muon optimizer papers‹/recap›"; expect(parseAgentStatusResponse(text, true)).toEqual({ summary: "Curating a niche list of Muon optimizer papers", }); }); });
describe("buildStatusContext", () => { test("includes the agent state and the trailing conversation with tool names", () => { const context = buildStatusContext( [userMessage("add a login endpoint"), assistantMessage("Editing the router", ["Edit", "Bash"])], true, ); expect(context).toContain("<agent-state>working</agent-state>"); expect(context).toContain("user: add a login endpoint"); expect(context).toContain("assistant: Editing the router [tools: Edit, Bash]"); });
test("marks idle sessions as finished", () => { expect(buildStatusContext([userMessage("hi")], false)).toContain("idle (finished its turn)"); });
test("only keeps the most recent messages", () => { const messages = Array.from({ length: 20 }, (_, i) => userMessage(`message ${i}`)); const context = buildStatusContext(messages, false); expect(context).toContain("message 19"); expect(context).not.toContain("message 0\n"); }); });
describe("status change notification", () => { function makeState(options: { messages: AgentMessage[]; isSessionActive: boolean; summaryState?: AgentStatus; persistedStatus?: AgentStatus; appendAgentStatus?: (status: AgentStatus) => void; getLeafId?: () => string | null; }): ActiveSessionState { return { activeSessionId: "active-1", summaryState: options.summaryState, runtime: { session: { isSessionActive: options.isSessionActive, messages: options.messages, modelRegistry: {}, settingsManager: SettingsManager.inMemory(), state: { streamingMessage: undefined }, sessionManager: { appendAgentStatus: options.appendAgentStatus ?? (() => {}), getLatestAgentStatus: () => options.persistedStatus, getLeafId: options.getLeafId ?? (() => null), }, }, }, } as unknown as ActiveSessionState; }
async function settle(state: ActiveSessionState, generated: { summary: string; taskState?: "needs_input" }) { const onStatusChanged = vi.fn(); const summarizer = new DaemonSessionSummarizer( () => [state], onStatusChanged, async () => generated, ); await (summarizer as unknown as { summarize(state: ActiveSessionState): Promise<void> }).summarize(state); return onStatusChanged; }
test("an idle settle with unchanged verdict text still notifies: its currency drives the roster", async () => { const previous: AgentStatus = { summary: "Working on it", taskState: "needs_input", basedOnMessageCount: 1 }; const state = makeState({ messages: [userMessage("hi"), userMessage("more")], isSessionActive: false, summaryState: previous, });
const onStatusChanged = await settle(state, { summary: "Working on it", taskState: "needs_input" });
expect(state.summaryState?.basedOnMessageCount).toBe(2); expect(onStatusChanged).toHaveBeenCalledOnce(); });
function failingIdleSetup( stateOptions: Parameters<typeof makeState>[0], generateFn: () => Promise<undefined> = async () => undefined, ) { const state = makeState(stateOptions); const generate = vi.fn(generateFn); const summarizer = new DaemonSessionSummarizer(() => [state], undefined, generate); const internal = summarizer as unknown as { summarize(state: ActiveSessionState): Promise<void>; failedIdleGenerations: Map<string, unknown>; }; return { state, generate, summarizer, internal }; }
// Warmup pins the ceiling itself: repeated failing idle sweeps stop paying // after three attempts and never persist the fabricated in-memory fallback. test.each([ { rearm: "branch navigation moving the leaf at the same length", trigger: (leaf: { id: string }) => { leaf.id = "leaf-b"; }, }, { rearm: "the backoff elapsing so external failures recover", trigger: () => vi.setSystemTime(Date.now() + 31 * 60_000), }, ])("the exhausted idle retry ceiling re-arms on $rearm", async ({ trigger }) => { vi.useFakeTimers(); try { const appendAgentStatus = vi.fn(); const leaf = { id: "leaf-a" }; const { state, generate, internal } = failingIdleSetup({ messages: [userMessage("hi")], isSessionActive: false, appendAgentStatus, getLeafId: () => leaf.id, });
for (let sweep = 0; sweep < 5; sweep++) { await internal.summarize(state); } expect(generate).toHaveBeenCalledTimes(3); expect(appendAgentStatus).not.toHaveBeenCalled(); expect(state.summaryState).toEqual({ summary: "", taskState: "needs_input", basedOnMessageCount: 1 });
trigger(leaf); await internal.summarize(state); expect(generate).toHaveBeenCalledTimes(4); expect(appendAgentStatus).not.toHaveBeenCalled(); } finally { vi.useRealTimers(); } });
test("a forget() during an in-flight idle generation leaves no failure record behind", async () => { let release!: () => void; const gate = new Promise<void>((resolveGate) => { release = resolveGate; }); const { state, summarizer, internal } = failingIdleSetup( { messages: [userMessage("hi")], isSessionActive: false }, async () => { await gate; return undefined; }, );
const pass = internal.summarize(state); summarizer.forget("active-1"); release(); await pass;
expect(internal.failedIdleGenerations.size).toBe(0); });
test("an idle re-settle matching the latest persisted status appends nothing", async () => { const appendAgentStatus = vi.fn(); const persisted: AgentStatus = { summary: "Awaiting review", taskState: "needs_input", basedOnMessageCount: 1, }; const state = makeState({ messages: [userMessage("hi")], isSessionActive: false, persistedStatus: persisted, appendAgentStatus, });
const onStatusChanged = vi.fn(); const summarizer = new DaemonSessionSummarizer( () => [state], onStatusChanged, async () => ({ summary: "Awaiting review", taskState: "needs_input" as const }), ); await (summarizer as unknown as { summarize(state: ActiveSessionState): Promise<void> }).summarize(state);
expect(appendAgentStatus).not.toHaveBeenCalled(); expect(onStatusChanged).toHaveBeenCalledOnce(); expect(state.summaryState).toEqual(persisted); });
test("a working refresh with unchanged text stays quiet", async () => { const previous: AgentStatus = { summary: "Working on it", taskState: "needs_input", basedOnMessageCount: 2 }; const state = makeState({ messages: [userMessage("hi"), userMessage("more")], isSessionActive: true, summaryState: previous, });
const onStatusChanged = await settle(state, { summary: "Working on it" });
expect(onStatusChanged).not.toHaveBeenCalled(); }); });});