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.
16 kB · 427 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428import 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;}
/** * 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, ) {}
start(): void { if (this.interval) { 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) { 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 { 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> { 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; 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; if (contentUnchanged && !isWorking && !owesIdleVerdict && !owesSummary) { 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); const status: AgentStatus = { summary: result.summary, taskState, basedOnMessageCount: messageCount, }; // 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; // Persist only settled idle verdicts from real classifications that // differ from the latest persisted entry: idle sweeps must not grow the journal. if (!isWorking && generated) { 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); } } 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); } } }}