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.
19 kB · 503 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504import type { AgentMessage } from "@earendil-works/pi-agent-core";import type { Api, Model } from "@earendil-works/pi-ai";import { completeSimple } from "@earendil-works/pi-ai";import type { ModelRegistry } from "../../core/model-registry.js";import { completeWithProviderRetry, type ProviderRetryPolicy, providerRetryPolicy } from "../../core/provider-retry.js";import type { AgentStatus, AgentTaskState } from "../../core/session-manager.js";import type { ActiveSessionState } from "./active-session-state.js";
const SWEEP_INTERVAL_MS = 25_000;// Collapse a tool-use loop's rapid turn_end bursts into one summarization.const SETTLE_DEBOUNCE_MS = 2_000;// Idle generations stop retrying (and paying) on unchanged content until the// backoff elapses, so transient outages and late credentials still recover.const IDLE_GENERATION_ATTEMPT_LIMIT = 3;const IDLE_GENERATION_RETRY_BACKOFF_MS = 30 * 60_000;
const SUMMARY_MODEL_PROVIDER = "prime-inference";const SUMMARY_MODEL_ID = "qwen/qwen3-30b-a3b-instruct-2507";
const SUMMARY_CONTEXT_MESSAGES = 8;const SUMMARY_MAX_CHARS_PER_MESSAGE = 600;// Generous so a chatty model still closes the tags before truncation.const SUMMARY_MAX_TOKENS = 400;
export const AGENT_STATUS_SYSTEM_PROMPT = `You generate a status line for an AI coding agent dashboard. You are given the recent conversation between a user and the agent, plus whether the agent is currently working or idle.
Output ONLY these two tags, nothing before, between, or after. Do not think out loud, explain, or count words.<recap>a present-tense clause, at most 12 words, saying what the agent is doing or just did, no trailing period</recap><status>one of NEEDS_INPUT, COMPLETED</status>
STATUS meaning:- COMPLETED: the agent finished its turn AND the user's request is fully done with nothing left.- NEEDS_INPUT: the agent finished its turn but the task is not fully done — it asked a question, hit a blocker, or needs more prompting.When you are unsure between COMPLETED and NEEDS_INPUT, choose NEEDS_INPUT.
Example:<recap>Refactoring the auth middleware and updating its tests</recap><status>NEEDS_INPUT</status>`;
export interface AgentStatusResult { summary: string; taskState?: AgentTaskState;}
/** Resolve the cheap summary model, or undefined when it has no configured auth. */export function resolveSummaryModel(registry: ModelRegistry): Model<Api> | undefined { const model = registry.find(SUMMARY_MODEL_PROVIDER, SUMMARY_MODEL_ID); if (model && registry.hasConfiguredAuth(model)) { return model; } return undefined;}
function messageText(content: unknown): { text: string; tools: string[] } { if (typeof content === "string") { return { text: content, tools: [] }; } if (!Array.isArray(content)) { return { text: "", tools: [] }; } const parts: string[] = []; const tools: string[] = []; for (const block of content) { if (typeof block !== "object" || block === null) { continue; } const type = (block as { type?: unknown }).type; if (type === "text" && typeof (block as { text?: unknown }).text === "string") { parts.push((block as { text: string }).text); } else if (type === "tool_use" || type === "toolUse") { const name = (block as { name?: unknown }).name; if (typeof name === "string") { tools.push(name); } } } return { text: parts.join("\n"), tools };}
function clamp(text: string, max: number): string { const normalized = text.replace(/\s+/g, " ").trim(); return normalized.length > max ? `${normalized.slice(0, max)}…` : normalized;}
/** Serialize the trailing messages into a compact prompt body (tool calls by name only). */export function buildStatusContext(messages: readonly AgentMessage[], isWorking: boolean): string { const recent = messages.slice(-SUMMARY_CONTEXT_MESSAGES); const lines: string[] = []; for (const message of recent) { const role = message.role; if (role !== "user" && role !== "assistant" && role !== "toolResult" && role !== "custom") { continue; } const { text, tools } = messageText(message.content); const body = clamp(text, SUMMARY_MAX_CHARS_PER_MESSAGE); const toolNote = tools.length > 0 ? `[tools: ${[...new Set(tools)].join(", ")}]` : ""; const rendered = [body, toolNote].filter((part) => part.length > 0).join(" "); if (rendered) { lines.push(`${role}: ${rendered}`); } } const state = isWorking ? "working" : "idle (finished its turn)"; return `<agent-state>${state}</agent-state>\n<conversation>\n${lines.join("\n")}\n</conversation>`;}
// Cuts a word-counting trailer the model sometimes appends, e.g.// `Sending X. That's 5 words? Count: X(1)... = 6 words.`. Kept to structural// counting markers so plain words ("Waiting for CI") survive.const REASONING_TRAILER = /\s*(?:["”]\s*)?(?:\bthat['’]?s\s+\d+\s*words?\b|\bcount\s*:|\(\d+\)|=\s*\d+\s*words?\b).*/i;const COUNTING_ARTIFACT = /\(\d+\)|=\s*\d+\s*words?\b/i;const MAX_RECAP_WORDS = 16;
function cleanRecap(raw: string): string | undefined { const value = raw .trim() .replace(REASONING_TRAILER, "") .replace(/^["“']+|["”']+$/g, "") .replace(/[.\s]+$/, "") .trim(); if (!value || value.startsWith("<") || /present-tense|12 words/i.test(value)) { return undefined; } if (COUNTING_ARTIFACT.test(value) || value.split(/\s+/).length > MAX_RECAP_WORDS) { return undefined; } return value;}
/** Take the content of the last `<recap>` and `<status>` tags; idle verdicts default to needs_input. */export function parseAgentStatusResponse(text: string, isWorking: boolean): AgentStatusResult | undefined { // Normalize unicode angle-bracket lookalikes (‹ › ＜ ＞) so a tag written with them still parses. const cleaned = text.replace(/[‹＜]/g, "<").replace(/[›＞]/g, ">");
const recapMatch = [...cleaned.matchAll(/<recap>([\s\S]*?)<\/recap>/gi)].at(-1); const summary = recapMatch ? cleanRecap(recapMatch[1]!) : undefined; if (!summary) { return undefined; } if (isWorking) { return { summary }; } const statusMatch = [...cleaned.matchAll(/<status>\s*([a-z_]+)\s*<\/status>/gi)].at(-1); const status = statusMatch ? statusMatch[1]!.toUpperCase() : undefined; const taskState: AgentTaskState = status === "COMPLETED" ? "completed" : "needs_input"; return { summary, taskState };}
export interface GenerateAgentStatusParams { registry: ModelRegistry; messages: readonly AgentMessage[]; isWorking: boolean; retryPolicy?: ProviderRetryPolicy; signal?: AbortSignal;}
/** One cheap model call for a fresh status, or undefined if unavailable/empty/failed. */export async function generateAgentStatus(params: GenerateAgentStatusParams): Promise<AgentStatusResult | undefined> { const { registry, messages, isWorking, retryPolicy, signal } = params; if (messages.length === 0) { return undefined; } const model = resolveSummaryModel(registry); if (!model) { return undefined; } const auth = await registry.getApiKeyAndHeaders(model); if (!auth.ok || !auth.apiKey) { return undefined; } try { // One failed attempt would settle an idle session to a stale needs_input verdict. const response = await completeWithProviderRetry( () => completeSimple( model, { systemPrompt: AGENT_STATUS_SYSTEM_PROMPT, messages: [ { role: "user" as const, content: [{ type: "text" as const, text: buildStatusContext(messages, isWorking) }], timestamp: Date.now(), }, ], }, { maxTokens: SUMMARY_MAX_TOKENS, apiKey: auth.apiKey, headers: auth.headers, signal }, ), { policy: retryPolicy, signal }, ); if (response.stopReason === "error") { return undefined; } const textContent = response.content .filter((c): c is { type: "text"; text: string } => c.type === "text") .map((c) => c.text) .join("\n"); return parseAgentStatusResponse(textContent, isWorking); } catch { return undefined; }}
function isSessionWorking(state: ActiveSessionState): boolean { const session = state.runtime.session; return session.isSessionActive;}
// Recap prefix for a turn that errored; the transcript's own error text follows// it so the persisted verdict reports the real last event, never invented work.const ERROR_RECAP_PREFIX = "Model request failed";// Generous; the agents view truncates recaps further for display.const ERROR_RECAP_MAX_CHARS = 160;
/** * Recap for an idle session whose last turn errored, or undefined when the last * turn ended normally (or never produced an assistant message). A turn that * errored produced no final answer, so its verdict must come from the * transcript's error — the classifier would only see the task text and invent * work that never happened. */function terminalTurnError(messages: readonly AgentMessage[]): string | undefined { for (let i = messages.length - 1; i >= 0; i--) { const message = messages[i]; if (message.role !== "assistant") { continue; } const { stopReason, errorMessage } = message as { stopReason?: unknown; errorMessage?: unknown }; if (stopReason !== "error") { return undefined; } const detail = typeof errorMessage === "string" ? errorMessage.trim() : ""; return detail ? `${ERROR_RECAP_PREFIX}: ${clamp(detail, ERROR_RECAP_MAX_CHARS)}` : ERROR_RECAP_PREFIX; } return undefined;}
/** * Background status summarization for daemon-hosted sessions, top-level and * subagents alike. A periodic sweep refreshes working sessions; debounced * turn-end activity drives the idle verdict. Status lives in memory; settled * idle verdicts are persisted. */export class DaemonSessionSummarizer { private interval: ReturnType<typeof setInterval> | undefined; private readonly debounceTimers = new Map<string, ReturnType<typeof setTimeout>>(); // Controller per in-flight summary so a closing session can abort its write. private readonly inFlight = new Map<string, AbortController>(); // Sessions requested while one was running; get one more pass on completion. private readonly rerunRequested = new Set<string>(); // Failed idle generations per session, keyed to the settled content they saw. private readonly failedIdleGenerations = new Map< string, { contentKey: string; attempts: number; lastFailureAt: number } >();
constructor( private readonly listSessions: () => readonly ActiveSessionState[], private readonly onStatusChanged?: (state: ActiveSessionState) => void, // Injectable for tests. private readonly generate: ( params: GenerateAgentStatusParams, ) => Promise<AgentStatusResult | undefined> = generateAgentStatus, // Off switch (`sessionRecap.enabled: false`): a disabled summarizer never // generates, seeds, or persists a status, so no recap line is produced. private readonly enabled: () => boolean = () => true, ) {}
start(): void { if (this.interval || !this.enabled()) { return; } this.interval = setInterval(() => { for (const state of this.listSessions()) { void this.summarize(state); } }, SWEEP_INTERVAL_MS); this.interval.unref?.(); }
stop(): void { if (this.interval) { clearInterval(this.interval); this.interval = undefined; } for (const timer of this.debounceTimers.values()) { clearTimeout(timer); } this.debounceTimers.clear(); for (const controller of this.inFlight.values()) { controller.abort(); } this.rerunRequested.clear(); this.failedIdleGenerations.clear(); }
/** Drop any pending work for a session that is closing. */ forget(activeSessionId: string): void { const timer = this.debounceTimers.get(activeSessionId); if (timer) { clearTimeout(timer); this.debounceTimers.delete(activeSessionId); } this.inFlight.get(activeSessionId)?.abort(); this.rerunRequested.delete(activeSessionId); this.failedIdleGenerations.delete(activeSessionId); }
/** Seed in-memory status from the persisted entry when a session is added. */ seed(state: ActiveSessionState): void { if (state.summaryState || !this.enabled()) { return; } const persisted = state.runtime.session.sessionManager.getLatestAgentStatus(); if (persisted) { state.summaryState = persisted; } }
/** Called when a session finishes a turn; debounce until the agent settles. */ notifyActivity(state: ActiveSessionState): void { if (!this.enabled()) { return; } const id = state.activeSessionId; const existing = this.debounceTimers.get(id); if (existing) { clearTimeout(existing); } const timer = setTimeout(() => { this.debounceTimers.delete(id); void this.summarize(state); }, SETTLE_DEBOUNCE_MS); timer.unref?.(); this.debounceTimers.set(id, timer); }
private async summarize(state: ActiveSessionState): Promise<void> { if (!this.enabled()) { return; } const id = state.activeSessionId; if (this.inFlight.has(id)) { this.rerunRequested.add(id); // run once more after the current pass return; } const session = state.runtime.session; const messages = session.messages; if (messages.length === 0) { return; } const messageCount = messages.length; const isWorking = isSessionWorking(state); const previous = state.summaryState; // Idle sessions with a current verdict need no refresh — except a // transcript whose terminal turn errored (owesErrorVerdict below) — // while working sessions always refresh so the recap keeps up with the // in-progress turn. const contentUnchanged = previous?.basedOnMessageCount === messageCount; const owesIdleVerdict = !isWorking && previous?.taskState === undefined; // A blank recap means the model call hasn't succeeded yet (e.g. the // needs_input fallback fired on a transient failure); keep retrying until a // real summary lands so the recap isn't left permanently empty. const owesSummary = !isWorking && !previous?.summary; // A turn that errored produced no final answer. Only a real final answer may // earn a completed verdict, so skip the classifier entirely — it sees the // task text with no evidence of the failure and would invent work — and // settle the verdict from the transcript's actual last event instead. // The pass below is synchronous, so no stale-state discard is needed. const turnError = !isWorking ? terminalTurnError(messages) : undefined; // A daemon restart seeds the latest persisted verdict, which pre-fix code // may have fabricated as `completed` for a transcript whose terminal turn // errored at this same message count. Only the error branch below repairs // it, so exempt terminal errors from the unchanged-content fast return. const owesErrorVerdict = turnError !== undefined && previous?.taskState !== "error"; if (contentUnchanged && !isWorking && !owesIdleVerdict && !owesSummary && !owesErrorVerdict) { return; } if (turnError !== undefined) { this.failedIdleGenerations.delete(id); const status: AgentStatus = { summary: turnError, taskState: "error", basedOnMessageCount: messageCount, }; this.commitStatus(state, session, status, { isWorking, previous, persist: true }); return; } // The leaf entry id is the branch-tip identity (appends, edits, and branch // navigation all move it; counts and timestamps collide across siblings). const contentKey = `${session.sessionManager.getLeafId() ?? "root"}:${messageCount}`; const failed = this.failedIdleGenerations.get(id); if ( !isWorking && failed?.contentKey === contentKey && failed.attempts >= IDLE_GENERATION_ATTEMPT_LIMIT && Date.now() - failed.lastFailureAt < IDLE_GENERATION_RETRY_BACKOFF_MS ) { return; } // Include the in-progress message so a long streaming turn gets a live recap. const streaming = isWorking ? session.state.streamingMessage : undefined; const contextMessages = streaming ? [...messages, streaming] : messages;
const controller = new AbortController(); this.inFlight.set(id, controller); try { const generated = await this.generate({ registry: session.modelRegistry, messages: contextMessages, isWorking, retryPolicy: providerRetryPolicy(session.settingsManager), signal: controller.signal, }); if (generated) { this.failedIdleGenerations.delete(id); } else if (!isWorking && !controller.signal.aborted) { // The aborted check keeps a racing forget() from repopulating the map. this.failedIdleGenerations.set(id, { contentKey, attempts: failed?.contentKey === contentKey ? failed.attempts + 1 : 1, lastFailureAt: Date.now(), }); } // A failed classification on an idle session would spin at "working" // forever (the activity axis holds unjudged idle sessions there), so // settle it to needs_input. const result = generated ?? (!isWorking && (owesIdleVerdict || owesSummary) ? { summary: previous?.summary ?? "", taskState: "needs_input" as const } : undefined); if (!result) { return; } // Discard if the session closed, was swapped, or moved to a new turn // during the async call — never write a verdict for stale state. if ( controller.signal.aborted || state.runtime.session !== session || isSessionWorking(state) !== isWorking || session.messages.length !== messageCount ) { return; } // A working refresh carries no verdict; keep the prior one at the same // message count so a still-valid needs_input isn't dropped. const taskState = result.taskState ?? (previous?.basedOnMessageCount === messageCount ? previous?.taskState : undefined); // Only a real classification persists; the needs_input fallback above // must not grow the journal. this.commitStatus( state, session, { summary: result.summary, taskState, basedOnMessageCount: messageCount }, { isWorking, previous, persist: generated !== undefined }, ); } finally { this.inFlight.delete(id); // Re-debounce a request that arrived mid-pass instead of dropping it. if (this.rerunRequested.delete(id)) { this.notifyActivity(state); } } }
/** * Publish an in-memory status and persist a settled idle verdict that differs * from the latest persisted entry: sweeps must not grow the journal. `persist` * is true only for real verdicts — model classifications and * transcript-derived error verdicts — never for the needs_input fallback. */ private commitStatus( state: ActiveSessionState, session: ActiveSessionState["runtime"]["session"], status: AgentStatus, { isWorking, previous, persist }: { isWorking: boolean; previous: AgentStatus | undefined; persist: boolean }, ): void { // An idle settle refreshes the verdict's currency, which drives the roster's // activity axis: it must publish even when the verdict text is unchanged. const changed = previous?.summary !== status.summary || previous?.taskState !== status.taskState || (!isWorking && previous?.basedOnMessageCount !== status.basedOnMessageCount); state.summaryState = status; if (!isWorking && persist) { const persisted = session.sessionManager.getLatestAgentStatus(); if ( persisted?.summary !== status.summary || persisted.taskState !== status.taskState || persisted.basedOnMessageCount !== status.basedOnMessageCount ) { try { session.sessionManager.appendAgentStatus(status); } catch { // best-effort; in-memory status still shows } } } if (changed) { this.onStatusChanged?.(state); } }}