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.
11 kB · 254 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255import * as ai from "@earendil-works/pi-ai";import { afterEach, describe, expect, test, vi } from "vitest";import { SettingsManager } from "../src/core/settings-manager.js";import type { ActiveSessionState } from "../src/modes/daemon/active-session-state.js";import { DaemonSessionSummarizer } from "../src/modes/daemon/daemon-session-summarizer.js";
// The debounce the summarizer waits for after a turn settles (kept in sync with// SETTLE_DEBOUNCE_MS in the module).const SETTLE_MS = 2000;
function makeState( opts: { working?: boolean; messages?: number; kind?: "top-level" | "subagent"; persisted?: unknown } = {},): ActiveSessionState { const appended: unknown[] = []; const state = { activeSessionId: "a1", summaryState: undefined, runtime: { metadata: { kind: opts.kind ?? "top-level" }, session: { isStreaming: opts.working ?? false, isCompacting: false, isSessionActive: opts.working ?? false, sessionActions: { queuedCount: 0, steering: [], followUps: [] }, messages: Array.from({ length: opts.messages ?? 2 }, () => ({ role: "user", content: "hi" })), state: { streamingMessage: undefined }, modelRegistry: {}, settingsManager: SettingsManager.inMemory(), sessionManager: { appendAgentStatus: (s: unknown) => appended.push(s), getLatestAgentStatus: () => opts.persisted, getLeafId: () => null, }, }, }, } as unknown as ActiveSessionState; (state as unknown as { appendedStatuses: unknown[] }).appendedStatuses = appended; return state;}
describe("DaemonSessionSummarizer lifecycle", () => { afterEach(() => { vi.useRealTimers(); vi.restoreAllMocks(); });
test.each([ { enabled: false, maxRetries: 3 }, { enabled: true, maxRetries: 0 }, ])("honors session retry settings %j on a transient summary failure", async (retry) => { vi.useFakeTimers(); const model = ai.getModel("openai", "gpt-4o-mini"); const complete = vi.spyOn(ai, "completeSimple").mockResolvedValue({ role: "assistant", content: [], api: model.api, provider: model.provider, model: model.id, usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "error", errorMessage: "503 Service Unavailable", diagnostics: [{ type: "provider_stream_failure", timestamp: 0, details: { kind: "server_error" } }], timestamp: 0, }); const state = makeState(); Object.assign(state.runtime.session, { settingsManager: SettingsManager.inMemory({ retry }) }); Object.assign(state.runtime.session.modelRegistry, { find: () => model, hasConfiguredAuth: () => true, getApiKeyAndHeaders: async () => ({ ok: true, apiKey: "test-key" }), }); const summarizer = new DaemonSessionSummarizer(() => [state]);
summarizer.notifyActivity(state); await vi.runAllTimersAsync();
expect(complete).toHaveBeenCalledOnce(); expect(state.summaryState).toMatchObject({ summary: "", taskState: "needs_input", basedOnMessageCount: 2 }); });
test("runs the model call after the settle debounce and records the verdict", async () => { vi.useFakeTimers(); const generate = vi.fn().mockResolvedValue({ summary: "Added the health endpoint", taskState: "completed" }); const onStatusChanged = vi.fn(); const summarizer = new DaemonSessionSummarizer(() => [], onStatusChanged, generate); const state = makeState({ working: false });
summarizer.notifyActivity(state); // No model call until the agent has settled for the debounce window. await vi.advanceTimersByTimeAsync(SETTLE_MS - 500); expect(generate).not.toHaveBeenCalled();
await vi.advanceTimersByTimeAsync(600); expect(generate).toHaveBeenCalledOnce(); expect(state.summaryState).toMatchObject({ summary: "Added the health endpoint", taskState: "completed" }); expect(onStatusChanged).toHaveBeenCalled(); });
test("a failing model on an idle session settles to a needs_input fallback verdict", async () => { vi.useFakeTimers(); const generate = vi.fn().mockResolvedValue(undefined); // 404 / 401 / timeout / unparseable const summarizer = new DaemonSessionSummarizer(() => [], undefined, generate); const state = makeState({ working: false });
summarizer.notifyActivity(state); await vi.advanceTimersByTimeAsync(SETTLE_MS + 500); expect(generate).toHaveBeenCalledOnce(); // The activity axis holds an unjudged idle session at "working"; the fallback // settles it to needs_input so it doesn't spin forever. expect(state.summaryState).toMatchObject({ taskState: "needs_input", basedOnMessageCount: 2 }); });
test("retries after a needs_input fallback until a real summary lands", async () => { vi.useFakeTimers(); const generate = vi .fn() .mockResolvedValueOnce(undefined) // transient failure → blank needs_input fallback .mockResolvedValue({ summary: "Reviewed the diff", taskState: "completed" }); const summarizer = new DaemonSessionSummarizer(() => [], undefined, generate); const state = makeState({ working: false });
summarizer.notifyActivity(state); await vi.advanceTimersByTimeAsync(SETTLE_MS + 500); expect(state.summaryState).toMatchObject({ summary: "", taskState: "needs_input" });
summarizer.notifyActivity(state); await vi.advanceTimersByTimeAsync(SETTLE_MS + 500); expect(generate).toHaveBeenCalledTimes(2); expect(state.summaryState).toMatchObject({ summary: "Reviewed the diff", taskState: "completed" }); });
test("refreshes a working session even when the message count is unchanged", async () => { vi.useFakeTimers(); const generate = vi.fn().mockResolvedValue({ summary: "Editing the router" }); const summarizer = new DaemonSessionSummarizer(() => [], undefined, generate); const state = makeState({ working: true }); // A status already exists for the current message count, so an idle session // would be skipped — but a working one must still refresh its recap. state.summaryState = { summary: "Editing the router", taskState: undefined, basedOnMessageCount: 2 };
summarizer.notifyActivity(state); await vi.advanceTimersByTimeAsync(SETTLE_MS + 500); expect(generate).toHaveBeenCalledOnce(); expect(generate.mock.calls[0]?.[0]).toMatchObject({ isWorking: true }); });
test("a working refresh preserves a still-valid verdict at the same message count", async () => { vi.useFakeTimers(); const generate = vi.fn().mockResolvedValue({ summary: "Editing the router" }); // working → no verdict const summarizer = new DaemonSessionSummarizer(() => [], undefined, generate); const state = makeState({ working: true }); state.summaryState = { summary: "Asked which db", taskState: "needs_input", basedOnMessageCount: 2 };
summarizer.notifyActivity(state); await vi.advanceTimersByTimeAsync(SETTLE_MS + 500); expect(state.summaryState).toMatchObject({ summary: "Editing the router", taskState: "needs_input" }); });
test("discards a verdict when the session is forgotten (closed) mid-call", async () => { vi.useFakeTimers(); const state = makeState({ working: false }); let summarizer!: DaemonSessionSummarizer; const generate = vi.fn().mockImplementation(async () => { summarizer.forget(state.activeSessionId); // session closes while the model runs return { summary: "Result after close", taskState: "completed" }; }); summarizer = new DaemonSessionSummarizer(() => [], undefined, generate);
summarizer.notifyActivity(state); await vi.advanceTimersByTimeAsync(SETTLE_MS + 500); expect(generate).toHaveBeenCalledOnce(); expect(state.summaryState).toBeUndefined(); });
test("discards a verdict when a new turn arrives during the model call", async () => { vi.useFakeTimers(); const state = makeState({ working: false }); const generate = vi.fn().mockImplementation(async () => { (state.runtime.session.messages as unknown[]).push({ role: "user", content: "another task" }); return { summary: "Stale summary for the old turn", taskState: "completed" }; }); const summarizer = new DaemonSessionSummarizer(() => [], undefined, generate);
summarizer.notifyActivity(state); await vi.advanceTimersByTimeAsync(SETTLE_MS + 500); expect(generate).toHaveBeenCalledOnce(); expect(state.summaryState).toBeUndefined(); });
test("summarizes a subagent and persists its settled idle verdict", async () => { vi.useFakeTimers(); const generate = vi.fn().mockResolvedValue({ summary: "Auditing the migration scripts", taskState: "completed" }); const summarizer = new DaemonSessionSummarizer(() => [], undefined, generate); const state = makeState({ working: false, kind: "subagent" });
summarizer.notifyActivity(state); await vi.advanceTimersByTimeAsync(SETTLE_MS + 500);
expect(generate).toHaveBeenCalledOnce(); expect(state.summaryState).toMatchObject({ summary: "Auditing the migration scripts", taskState: "completed" }); expect((state as unknown as { appendedStatuses: unknown[] }).appendedStatuses).toHaveLength(1); });
test("seeds a subagent's persisted recap into memory", () => { const summarizer = new DaemonSessionSummarizer(() => [], undefined, vi.fn()); const persisted = { summary: "Reviewing the diff", taskState: "needs_input", basedOnMessageCount: 3 }; const state = makeState({ kind: "subagent", persisted });
summarizer.seed(state); expect(state.summaryState).toEqual(persisted); });
test("a session that errored before any work settles from the transcript, never the classifier", async () => { vi.useFakeTimers(); // Every model request errored (e.g. provider 400s): no work was done, so a // completed verdict would be fabricated. The real last event is the error. const generate = vi.fn(); const onStatusChanged = vi.fn(); const summarizer = new DaemonSessionSummarizer(() => [], onStatusChanged, generate); const state = makeState({ working: false }); Object.assign(state.runtime.session, { messages: [ { role: "user", content: "write a session marker and verify the file content" }, { role: "assistant", content: [], stopReason: "error", errorMessage: "400 enable_thinking not supported" }, ], });
summarizer.notifyActivity(state); await vi.advanceTimersByTimeAsync(SETTLE_MS + 500);
// Zero work done: the classifier must not be paid to invent a verdict. expect(generate).not.toHaveBeenCalled(); expect(state.summaryState).toMatchObject({ summary: "Model request failed: 400 enable_thinking not supported", taskState: "error", basedOnMessageCount: 2, }); expect((state as unknown as { appendedStatuses: unknown[] }).appendedStatuses).toEqual([ { summary: "Model request failed: 400 enable_thinking not supported", taskState: "error", basedOnMessageCount: 2, }, ]); expect(onStatusChanged).toHaveBeenCalled(); });});