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.
10 kB · 321 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322import { existsSync, readdirSync, statSync } from "node:fs";import { basename, join } from "node:path";import type { AssistantMessage, Usage } from "@earendil-works/pi-ai";import type { RlmChildAgentStatus } from "./agent-session.js";import { calculateContextTokens, estimateContextTokens } from "./compaction/index.js";import type { ContextUsage } from "./extensions/index.js";import { buildSessionContext, type FileEntry, loadEntriesFromFile, type SessionEntry } from "./session-manager.js";import { addAssistantUsage, cloneUsage, emptyUsage, subtractAssistantUsage } from "./usage.js";
/** Resolves a model's context window so disk-only nodes can report utilization. */export type ContextWindowResolver = (provider: string, modelId: string) => number | undefined;
/** * One agent in the context overview: the main session or an RLM (sub-)agent. * `ownUsage` excludes descendants; `totalUsage` includes completed descendants, matching /usage. */export interface ContextTreeNode { /** "root" for the session itself; sub-xxxx for an RLM child. */ id: string; label: string; status: "active" | RlmChildAgentStatus; model?: { provider: string; id: string }; ownUsage: Usage; totalUsage: Usage; contextUsage?: ContextUsage; children: ContextTreeNode[];}
function isAssistantEntry(entry: SessionEntry): entry is SessionEntry & { type: "message"; message: AssistantMessage;} { return entry.type === "message" && entry.message.role === "assistant";}
function readUserMessageText(content: unknown): string { if (typeof content === "string") { return content; } if (!Array.isArray(content)) { return ""; } return content .filter( (block): block is { type: "text"; text: string } => typeof block === "object" && block !== null && (block as { type?: unknown }).type === "text" && typeof (block as { text?: unknown }).text === "string", ) .map((block) => block.text) .join("\n");}
function compactLabel(text: string, maxLength = 80): string { const compact = text.replace(/\s+/g, " ").trim(); if (compact.length <= maxLength) { return compact; } return `${compact.slice(0, Math.max(0, maxLength - 3)).trimEnd()}...`;}
/** * Usage totals for one agent: `totalUsage` sums the branch's assistant usage * (attributed aggregates, so descendants are included), `ownUsage` removes the * attributions targeting those assistants. Attribution entries are matched by * target across ALL entries, not just the branch: attributions rewrite the * target assistant's usage no matter which branch they were appended on, so a * fork that keeps the assistant but drops the attribution entry must still * subtract it. * * Totals are deliberately cumulative across compactions: compaction shrinks * the model-facing context, not what the session has spent, so assistants * dropped from the resolved context still count here. */export function computeOwnAndTotalUsage( branch: SessionEntry[], allEntries: SessionEntry[],): { ownUsage: Usage; totalUsage: Usage } { const totalUsage = emptyUsage(); const branchAssistantIds = new Set<string>(); for (const entry of branch) { if (isAssistantEntry(entry)) { branchAssistantIds.add(entry.id); addAssistantUsage(totalUsage, entry.message.usage); } else if ((entry.type === "compaction" || entry.type === "branch_summary") && entry.usage) { addAssistantUsage(totalUsage, entry.usage); } } const ownUsage = cloneUsage(totalUsage); for (const entry of allEntries) { if (entry.type === "child_usage_attributed" && branchAssistantIds.has(entry.targetId)) { subtractAssistantUsage(ownUsage, entry.childUsage); } } return { ownUsage, totalUsage };}
/** * Current context utilization from persisted entries, mirroring * AgentSession.getContextUsage(): unknown right after a compaction until the * next assistant response, otherwise the last assistant usage plus an * estimate for trailing messages (tool results, queued user input) that have * not hit the model yet. */function computeContextUsageFromEntries( allEntries: SessionEntry[], branch: SessionEntry[], contextWindow: number | undefined,): ContextUsage | undefined { if (!contextWindow || contextWindow <= 0) { return undefined; }
let latestCompactionIndex = -1; for (let i = branch.length - 1; i >= 0; i--) { if (branch[i].type === "compaction") { latestCompactionIndex = i; break; } }
if (latestCompactionIndex >= 0) { let hasPostCompactionUsage = false; for (let i = branch.length - 1; i > latestCompactionIndex; i--) { const entry = branch[i]; if (!isAssistantEntry(entry)) { continue; } const assistant = entry.message; if (assistant.stopReason === "aborted" || assistant.stopReason === "error") { continue; } if (calculateContextTokens(assistant.usage) > 0) { hasPostCompactionUsage = true; } break; } if (!hasPostCompactionUsage) { return { tokens: null, contextWindow, percent: null }; } }
const estimate = estimateContextTokens(buildSessionContext(allEntries).messages); if (estimate.tokens <= 0) { return undefined; } return { tokens: estimate.tokens, contextWindow, percent: (estimate.tokens / contextWindow) * 100 };}
function sessionEntriesFromFile(file: string): SessionEntry[] { return loadEntriesFromFile(file).filter((entry: FileEntry): entry is SessionEntry => entry.type !== "session");}
/** * Entries on the current branch, root to leaf, mirroring * SessionManager.getBranch(): the leaf is the last appended entry and the * branch is its parentId chain. Keeps forked/abandoned paths out of usage * sums so disk nodes match what a live session would report. */function branchEntries(entries: SessionEntry[]): SessionEntry[] { if (entries.length === 0) { return []; } const byId = new Map(entries.map((entry) => [entry.id, entry])); const branch: SessionEntry[] = []; const seen = new Set<string>(); let current: SessionEntry | undefined = entries[entries.length - 1]; while (current && !seen.has(current.id)) { seen.add(current.id); branch.push(current); current = current.parentId ? byId.get(current.parentId) : undefined; } return branch.reverse();}
/** * Terminal status for a persisted child, inferred from how its last assistant * turn ended: errored and aborted runs should not render as successful. */function statusFromBranch(entries: SessionEntry[]): "done" | "error" | "cancelled" { for (let i = entries.length - 1; i >= 0; i--) { const entry = entries[i]; if (!isAssistantEntry(entry)) { continue; } if (entry.message.stopReason === "error") { return "error"; } if (entry.message.stopReason === "aborted") { return "cancelled"; } return "done"; } return "done";}
function findSessionFile(dir: string): string | undefined { let newest: { path: string; mtime: number } | undefined; for (const name of readdirSync(dir)) { if (!name.endsWith(".jsonl")) { continue; } const path = join(dir, name); try { const mtime = statSync(path).mtime.getTime(); if (!newest || mtime > newest.mtime) { newest = { path, mtime }; } } catch { // Skip unreadable files. } } return newest?.path;}
function listChildSessionDirs(rlmSessionDir: string): string[] { let names: string[]; try { names = readdirSync(rlmSessionDir); } catch { return []; } return names .filter((name) => name.startsWith("sub-")) .map((name) => join(rlmSessionDir, name)) .filter((path) => { try { return statSync(path).isDirectory(); } catch { return false; } }) .sort((a, b) => { try { return statSync(a).mtime.getTime() - statSync(b).mtime.getTime(); } catch { return 0; } });}
/** * Build a context node for a completed RLM child from its persisted session * dir (sub-xxxx/). Children that already attributed grandchild usage carry the * aggregate on their assistant messages (applyChildUsageAttributions), so own * usage is recovered by subtracting the attribution entries. Returns undefined * when the dir holds no readable session. */export function loadContextTreeChildFromDisk( childSessionDir: string, resolveContextWindow: ContextWindowResolver,): ContextTreeNode | undefined { const sessionFile = findSessionFile(childSessionDir); if (!sessionFile) { return undefined; } const allEntries = sessionEntriesFromFile(sessionFile); const branch = branchEntries(allEntries); if (branch.length === 0) { return undefined; }
const { ownUsage, totalUsage } = computeOwnAndTotalUsage(branch, allEntries);
let model: { provider: string; id: string } | undefined; for (const entry of branch) { if (entry.type === "model_change") { model = { provider: entry.provider, id: entry.modelId }; } }
let label = ""; for (const entry of branch) { if (entry.type === "message" && entry.message.role === "user") { label = compactLabel(readUserMessageText(entry.message.content)); if (label) { break; } } }
const contextWindow = model ? resolveContextWindow(model.provider, model.id) : undefined;
return { id: basename(childSessionDir), label: label || "child agent", status: statusFromBranch(branch), model, ownUsage, totalUsage, contextUsage: computeContextUsageFromEntries(allEntries, branch, contextWindow), children: loadContextTreeChildrenFromDisk(childSessionDir, resolveContextWindow), };}
/** * Build context nodes for all persisted RLM children under an RLM session * dir, recursing into nested sub-* dirs for grandchildren. `skipIds` * excludes children that are already represented live. */export function loadContextTreeChildrenFromDisk( rlmSessionDir: string | undefined, resolveContextWindow: ContextWindowResolver, skipIds?: ReadonlySet<string>,): ContextTreeNode[] { if (!rlmSessionDir || !existsSync(rlmSessionDir)) { return []; } const nodes: ContextTreeNode[] = []; for (const childDir of listChildSessionDirs(rlmSessionDir)) { if (skipIds?.has(basename(childDir))) { continue; } const node = loadContextTreeChildFromDisk(childDir, resolveContextWindow); if (node) { nodes.push(node); } } return nodes;}