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.
20 kB · 531 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532import 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 { type AgentStatusResult, 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;}
function assistantError(errorMessage?: string): AgentMessage { return { role: "assistant", content: [], stopReason: "error", ...(errorMessage !== undefined ? { errorMessage } : {}), timestamp: 0, } as unknown as AgentMessage;}
describe("daemon session summarizer", () => { describe("parseAgentStatusResponse", () => { // The recap line is model output, so every tag shape the model actually produces is pinned here. test.each([ [ "parses recap and completion verdict for an idle session", "<recap>Added the API reference page</recap>\n<status>COMPLETED</status>", false, { summary: "Added the API reference page", taskState: "completed" }, ], [ "maps NEEDS_INPUT for idle sessions", "<recap>Asked which database to target</recap>\n<status>NEEDS_INPUT</status>", false, { summary: "Asked which database to target", taskState: "needs_input" }, ], [ "omits the verdict while working and ignores any status tag", "<recap>Refactoring token validation</recap>\n<status>COMPLETED</status>", true, { summary: "Refactoring token validation" }, ], [ "falls back to needs_input on an unrecognized idle verdict", "<recap>Doing something</recap>\n<status>MAYBE</status>", false, { summary: "Doing something", taskState: "needs_input" }, ], [ "falls back to needs_input on a missing idle verdict", "<recap>Doing something</recap>", false, { summary: "Doing something", taskState: "needs_input" }, ], ["ignores narration with no tags at all", "Investigating the failing test.", true, undefined], [ "ignores narration outside the tags", "Recap: . So: <recap>Curating a niche list of Muon optimizer papers</recap>", true, { summary: "Curating a niche list of Muon optimizer papers" }, ], [ "ignores reasoning prose around the tags", "Let me decide. The agent finished editing.\n<recap>Updated the login handler</recap>\n<status>COMPLETED</status>", false, { summary: "Updated the login handler", taskState: "completed" }, ], [ "rejects an echoed prompt template", "<recap>a present-tense clause, at most 12 words, no trailing period</recap>\n<status>COMPLETED</status>", true, undefined, ], ["returns undefined for empty output", "", false, undefined], ["returns undefined when no recap tag is present", "<status>COMPLETED</status>", false, undefined], [ "drops chain-of-thought that falls outside the closing recap tag", "<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>", true, { summary: "Sending SSH auth retry to tcg-autoresearch-rl" }, ], [ "rejects a recap body that is nothing but counting artifacts", "<recap>(1) word(2) count(3) = 3 words</recap>", true, undefined, ], [ "rejects a rambling recap that blows past the word ceiling", "<recap>this is a very long rambling sentence that just keeps going and going well past any reasonable recap length</recap>", true, undefined, ], [ "strips wrapping quotes the model adds around the recap", '<recap>"Wiring the recap line"</recap>\n<status>COMPLETED</status>', false, { summary: "Wiring the recap line", taskState: "completed" }, ], [ "ignores an open recap tag with no close", "<recap>Editing the parser\n<status>NEEDS_INPUT</status>", true, undefined, ], [ "takes the last recap tag when a draft is corrected", "<recap>Draft recap</recap>\n<recap>Final corrected recap</recap>", true, { summary: "Final corrected recap" }, ], [ "takes the last status tag when a draft is corrected", "<recap>Editing the parser</recap>\n<status>NEEDS_INPUT</status>\n<status>COMPLETED</status>", false, { summary: "Editing the parser", taskState: "completed" }, ], [ // The model sometimes emits the unicode lookalikes ‹ › instead of < >. "normalizes unicode angle-bracket lookalikes around the tags", "‹recap›Curating a niche list of Muon optimizer papers‹/recap›", true, { summary: "Curating a niche list of Muon optimizer papers" }, ], ])("%s", (_name, text, isWorking, expected) => { expect(parseAgentStatusResponse(text, isWorking)).toEqual(expected); }); });
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("the recap off switch generates and seeds nothing", async () => { const state = makeState({ messages: [userMessage("hi")], isSessionActive: false, persistedStatus: { summary: "p", basedOnMessageCount: 1 }, }); const gen = vi.fn(async () => ({ summary: "Generated recap" })); const off = new DaemonSessionSummarizer( () => [state], undefined, gen, () => false, ); off.start(); off.seed(state); off.notifyActivity(state); await (off as unknown as { summarize(s: ActiveSessionState): Promise<void> }).summarize(state); off.stop(); expect(gen).not.toHaveBeenCalled(); expect(state.summaryState).toBeUndefined(); });
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(); }); });
describe("errored turn verdicts", () => { const providerError = "400 enable_thinking is not supported for this model";
function erroredTranscript(errorMessage?: string): AgentMessage[] { return [userMessage("write a session marker and verify the file content"), assistantError(errorMessage)]; }
// A journal-backed state: getLatestAgentStatus returns what appends recorded, // falling back to a verdict that predates this run (e.g. written before a // daemon restart). function erroredState(options: { messages: AgentMessage[]; isSessionActive?: boolean; persistedStatus?: AgentStatus; }): { state: ActiveSessionState; appended: AgentStatus[]; } { const appended: AgentStatus[] = []; const state = { activeSessionId: "active-error", summaryState: undefined, runtime: { session: { isSessionActive: options.isSessionActive ?? false, messages: options.messages, modelRegistry: {}, settingsManager: SettingsManager.inMemory(), state: { streamingMessage: undefined }, sessionManager: { appendAgentStatus: (status: AgentStatus) => { appended.push(status); }, getLatestAgentStatus: () => appended.at(-1) ?? options.persistedStatus, getLeafId: () => null, }, }, }, } as unknown as ActiveSessionState; return { state, appended }; }
async function runSummarize( state: ActiveSessionState, generate: () => Promise<AgentStatusResult | undefined>, ): Promise<ReturnType<typeof vi.fn>> { const onStatusChanged = vi.fn(); const summarizer = new DaemonSessionSummarizer(() => [state], onStatusChanged, generate); await (summarizer as unknown as { summarize(state: ActiveSessionState): Promise<void> }).summarize(state); return onStatusChanged; }
test("an errored session persists the real error as its verdict, never a completed one", async () => { const { state, appended } = erroredState({ messages: erroredTranscript(providerError) }); // The classifier would only see the task text and invent completed work. const generate = vi.fn(async () => ({ summary: "Writing session marker and verifying file content", taskState: "completed" as const, }));
const onStatusChanged = await runSummarize(state, generate);
expect(generate).not.toHaveBeenCalled(); expect(state.summaryState).toEqual({ summary: `Model request failed: ${providerError}`, taskState: "error", basedOnMessageCount: 2, }); expect(appended).toEqual([ { summary: `Model request failed: ${providerError}`, taskState: "error", basedOnMessageCount: 2 }, ]); expect(onStatusChanged).toHaveBeenCalledOnce(); });
test("repeated sweeps over an unchanged errored transcript append nothing more", async () => { const { state, appended } = erroredState({ messages: erroredTranscript(providerError) }); const summarizer = new DaemonSessionSummarizer(() => [state], undefined, vi.fn()); const internal = summarizer as unknown as { summarize(state: ActiveSessionState): Promise<void> };
await internal.summarize(state); await internal.summarize(state); await internal.summarize(state);
expect(appended).toHaveLength(1); });
test("a restart-seeded completed verdict for an errored transcript is repaired by the sweep", async () => { // Pre-fix code fabricated a completed verdict for the errored transcript // and the journal kept it; a daemon restart seeds it back before the // first sweep. The unchanged-content fast path must not skip the repair. const persisted: AgentStatus = { summary: "Writing session marker and verifying file content", taskState: "completed", basedOnMessageCount: 2, }; const { state, appended } = erroredState({ messages: erroredTranscript(providerError), persistedStatus: persisted, }); const generate = vi.fn(async () => ({ summary: "unreached", taskState: "completed" as const })); const summarizer = new DaemonSessionSummarizer(() => [state], undefined, generate); // Daemon-mode bind() restores the persisted verdict on restart. summarizer.seed(state); const internal = summarizer as unknown as { summarize(state: ActiveSessionState): Promise<void> };
await internal.summarize(state);
expect(generate).not.toHaveBeenCalled(); expect(state.summaryState).toEqual({ summary: `Model request failed: ${providerError}`, taskState: "error", basedOnMessageCount: 2, }); expect(appended).toEqual([ { summary: `Model request failed: ${providerError}`, taskState: "error", basedOnMessageCount: 2 }, ]);
// Repaired once: later sweeps append nothing more. await internal.summarize(state); await internal.summarize(state); expect(appended).toHaveLength(1); });
test("an errored turn without an error message still settles to the error verdict", async () => { const { state, appended } = erroredState({ messages: erroredTranscript(undefined) });
await runSummarize(state, async () => undefined);
expect(state.summaryState).toMatchObject({ summary: "Model request failed", taskState: "error" }); expect(appended).toHaveLength(1); });
test("a successful final answer after an error earns a normal model verdict", async () => { const { state, appended } = erroredState({ messages: erroredTranscript(providerError) }); const generate = vi.fn(async () => ({ summary: "Wrote the marker file", taskState: "completed" as const })); const summarizer = new DaemonSessionSummarizer(() => [state], undefined, generate); const internal = summarizer as unknown as { summarize(state: ActiveSessionState): Promise<void> };
await internal.summarize(state); // error verdict settles first expect(state.summaryState?.taskState).toBe("error");
// The user retries and the turn succeeds: the classifier may judge again. (state.runtime.session as unknown as { messages: AgentMessage[] }).messages = [ ...erroredTranscript(providerError), userMessage("retry the marker task"), assistantMessage("Wrote .session-marker and verified its content"), ]; await internal.summarize(state);
expect(generate).toHaveBeenCalledTimes(1); expect(state.summaryState).toMatchObject({ summary: "Wrote the marker file", taskState: "completed" }); expect(appended.at(-1)).toMatchObject({ taskState: "completed" }); });
test("a working session is never error-settled, even with an errored trailing assistant message", async () => { // While working (e.g. mid-retry) the model refresh keeps recapping; // verdicts settle only when the session goes idle. const { state, appended } = erroredState({ messages: erroredTranscript(providerError), isSessionActive: true, }); const generate = vi.fn(async () => ({ summary: "Retrying the failed request" }));
await runSummarize(state, generate);
expect(generate).toHaveBeenCalledOnce(); expect(state.summaryState).toEqual({ summary: "Retrying the failed request", taskState: undefined, basedOnMessageCount: 2, }); expect(appended).toHaveLength(0); }); });});