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.
18 kB · 459 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import { Agent } from "@earendil-works/pi-agent-core";import { type AssistantMessage, getModel, type Usage } from "@earendil-works/pi-ai";import stripAnsi from "strip-ansi";import { afterEach, beforeAll, describe, expect, it } from "vitest";import { AgentSession } from "../src/core/agent-session.js";import { AuthStorage } from "../src/core/auth-storage.js";import { type ContextTreeNode, loadContextTreeChildrenFromDisk } from "../src/core/context-tree.js";import { ModelRegistry } from "../src/core/model-registry.js";import { SessionManager } from "../src/core/session-manager.js";import { SettingsManager } from "../src/core/settings-manager.js";import { addAssistantUsage, cloneUsage, emptyUsage } from "../src/core/usage.js";import { formatContextTree } from "../src/modes/interactive/components/context-tree-format.js";import { initTheme } from "../src/modes/interactive/theme/theme.js";import { createTestResourceLoader } from "./utilities.js";
const model = getModel("anthropic", "claude-sonnet-4-5")!;
beforeAll(() => { initTheme("dark");});
function createUsage(input: number, output: number, cost: number): Usage { return { input, output, cacheRead: 0, cacheWrite: 0, totalTokens: input + output, cost: { input: cost, output: 0, cacheRead: 0, cacheWrite: 0, total: cost, }, };}
function createAssistantMessage(text: string, usage: Usage): AssistantMessage { return { role: "assistant", content: [{ type: "text", text }], api: model.api, provider: model.provider, model: model.id, usage, stopReason: "stop", timestamp: Date.now(), };}
function createUserMessage(text: string) { return { role: "user" as const, content: text, timestamp: Date.now(), };}
/** * Write a child agent session into `dir` the way RLM child runs persist them: * model change, the spawn prompt as the first user message, and assistant * usage. Returns the assistant entry id for attribution. */function writeChildSession( dir: string, prompt: string, assistantUsage: Usage,): { assistantEntryId: string; sessionManager: SessionManager } { mkdirSync(dir, { recursive: true }); const sessionManager = SessionManager.create(process.cwd(), dir); sessionManager.newSession(); sessionManager.appendModelChange(model.provider, model.id); sessionManager.appendMessage(createUserMessage(prompt)); const assistantEntryId = sessionManager.appendMessage(createAssistantMessage("done", cloneUsage(assistantUsage))); return { assistantEntryId, sessionManager };}
const resolveContextWindow = () => 200000;
let tempDirs: string[] = [];
function makeTempDir(): string { const dir = mkdtempSync(join(tmpdir(), "context-tree-test-")); tempDirs.push(dir); return dir;}
afterEach(() => { for (const dir of tempDirs) { rmSync(dir, { recursive: true, force: true }); } tempDirs = [];});
describe("loadContextTreeChildrenFromDisk", () => { it("returns no nodes for a missing or empty rlm session dir", () => { expect(loadContextTreeChildrenFromDisk(undefined, resolveContextWindow)).toEqual([]); expect(loadContextTreeChildrenFromDisk(join(makeTempDir(), "missing"), resolveContextWindow)).toEqual([]); expect(loadContextTreeChildrenFromDisk(makeTempDir(), resolveContextWindow)).toEqual([]); });
it("builds nodes recursively and separates own usage from attributed child usage", () => { const rlmDir = makeTempDir(); const childDir = join(rlmDir, "sub-aaaa1111"); const grandchildDir = join(childDir, "sub-bbbb2222");
const childUsage = createUsage(1000, 200, 0.1); const grandchildUsage = createUsage(400, 50, 0.04); const child = writeChildSession(childDir, "summarize the auth module", childUsage); writeChildSession(grandchildDir, "read the oauth callbacks", grandchildUsage);
// The child folds the grandchild's usage into its assistant message, // exactly like _attributeRlmChildUsageToParent does. const aggregate = cloneUsage(childUsage); addAssistantUsage(aggregate, grandchildUsage); child.sessionManager.appendChildUsageAttribution(child.assistantEntryId, grandchildUsage, aggregate);
const nodes = loadContextTreeChildrenFromDisk(rlmDir, resolveContextWindow); expect(nodes).toHaveLength(1);
const childNode = nodes[0]; expect(childNode.id).toBe("sub-aaaa1111"); expect(childNode.label).toBe("summarize the auth module"); expect(childNode.status).toBe("done"); expect(childNode.model).toEqual({ provider: model.provider, id: model.id }); // Aggregate includes the grandchild; own does not. expect(childNode.totalUsage.input).toBe(1400); expect(childNode.totalUsage.cost.total).toBeCloseTo(0.14); expect(childNode.ownUsage.input).toBe(1000); expect(childNode.ownUsage.output).toBe(200); expect(childNode.ownUsage.cost.total).toBeCloseTo(0.1);
expect(childNode.children).toHaveLength(1); const grandchildNode = childNode.children[0]; expect(grandchildNode.id).toBe("sub-bbbb2222"); expect(grandchildNode.label).toBe("read the oauth callbacks"); expect(grandchildNode.ownUsage.input).toBe(400); expect(grandchildNode.totalUsage.input).toBe(400); expect(grandchildNode.children).toEqual([]);
// Own usage summed over the tree equals the attributed aggregate. expect(childNode.ownUsage.input + grandchildNode.ownUsage.input).toBe(childNode.totalUsage.input); });
it("reports context usage from the last assistant message and the model context window", () => { const rlmDir = makeTempDir(); writeChildSession(join(rlmDir, "sub-ctx00001"), "check context", createUsage(1500, 500, 0.02));
const nodes = loadContextTreeChildrenFromDisk(rlmDir, resolveContextWindow); expect(nodes[0].contextUsage).toEqual({ tokens: 2000, contextWindow: 200000, percent: 1, });
// Unknown context window: no context usage at all. const unresolved = loadContextTreeChildrenFromDisk(rlmDir, () => undefined); expect(unresolved[0].contextUsage).toBeUndefined(); });
it("reports unknown context usage after a trailing compaction", () => { const rlmDir = makeTempDir(); const child = writeChildSession(join(rlmDir, "sub-comp0001"), "compact me", createUsage(1500, 500, 0.02)); child.sessionManager.appendCompaction("summary", child.assistantEntryId, 2000);
const nodes = loadContextTreeChildrenFromDisk(rlmDir, resolveContextWindow); expect(nodes[0].contextUsage).toEqual({ tokens: null, contextWindow: 200000, percent: null }); });
it("derives error and cancelled status from the last assistant stop reason", () => { const rlmDir = makeTempDir(); const errored = writeChildSession(join(rlmDir, "sub-err00001"), "fail", createUsage(100, 10, 0.01)); errored.sessionManager.appendMessage({ ...createAssistantMessage("boom", createUsage(50, 5, 0.01)), stopReason: "error", }); const aborted = writeChildSession(join(rlmDir, "sub-abr00001"), "stop", createUsage(100, 10, 0.01)); aborted.sessionManager.appendMessage({ ...createAssistantMessage("", createUsage(50, 5, 0.01)), stopReason: "aborted", });
const nodes = loadContextTreeChildrenFromDisk(rlmDir, resolveContextWindow); const byId = new Map(nodes.map((n) => [n.id, n])); expect(byId.get("sub-err00001")?.status).toBe("error"); expect(byId.get("sub-abr00001")?.status).toBe("cancelled"); });
it("sums only the current branch, excluding abandoned forked paths", () => { const rlmDir = makeTempDir(); const childDir = join(rlmDir, "sub-fork0001"); mkdirSync(childDir, { recursive: true });
// Hand-written session with a fork: u1 -> a1 (abandoned) and u1 -> a2 (leaf). const header = { type: "session", version: 3, id: "11111111-2222-7333-8444-555555555555", timestamp: new Date().toISOString(), cwd: process.cwd(), }; const entry = (id: string, parentId: string | null, message: unknown) => ({ type: "message", id, parentId, timestamp: new Date().toISOString(), message, }); const lines = [ header, { type: "model_change", id: "m1", parentId: null, timestamp: new Date().toISOString(), provider: model.provider, modelId: model.id, }, entry("u1", "m1", createUserMessage("forked work")), entry("a1", "u1", createAssistantMessage("abandoned", createUsage(9000, 900, 0.9))), entry("a2", "u1", createAssistantMessage("kept", createUsage(1000, 100, 0.1))), ]; writeFileSync(join(childDir, `${header.id}.jsonl`), `${lines.map((line) => JSON.stringify(line)).join("\n")}\n`);
const nodes = loadContextTreeChildrenFromDisk(rlmDir, resolveContextWindow); expect(nodes).toHaveLength(1); // Only the leaf branch (u1 -> a2) counts; the abandoned a1 path does not. expect(nodes[0].ownUsage.input).toBe(1000); expect(nodes[0].totalUsage.input).toBe(1000); expect(nodes[0].contextUsage?.tokens).toBe(1100); });
it("subtracts attributions targeting branch assistants even when the attribution entry is off-branch", () => { const rlmDir = makeTempDir(); const childDir = join(rlmDir, "sub-attr0001"); mkdirSync(childDir, { recursive: true });
// a1 received a grandchild attribution, then the session forked from a1: // the attribution entry is off-branch but a1 (whose usage was rewritten // to the aggregate) is still on it. const grandchildUsage = createUsage(500, 50, 0.05); const ownAssistantUsage = createUsage(1000, 100, 0.1); const aggregate = cloneUsage(ownAssistantUsage); addAssistantUsage(aggregate, grandchildUsage); const header = { type: "session", version: 3, id: "11111111-2222-7333-8444-666666666666", timestamp: new Date().toISOString(), cwd: process.cwd(), }; const entry = (id: string, parentId: string | null, message: unknown) => ({ type: "message", id, parentId, timestamp: new Date().toISOString(), message, }); const lines = [ header, entry("u1", null, createUserMessage("attributed work")), entry("a1", "u1", createAssistantMessage("answer", cloneUsage(ownAssistantUsage))), { type: "child_usage_attributed", id: "x1", parentId: "a1", timestamp: new Date().toISOString(), targetId: "a1", childUsage: grandchildUsage, aggregateUsage: aggregate, }, entry("u2", "a1", createUserMessage("forked follow-up")), ]; writeFileSync(join(childDir, `${header.id}.jsonl`), `${lines.map((line) => JSON.stringify(line)).join("\n")}\n`);
const nodes = loadContextTreeChildrenFromDisk(rlmDir, resolveContextWindow); expect(nodes).toHaveLength(1); // a1 carries the aggregate (1500) on the branch; own must subtract the // off-branch attribution back out. expect(nodes[0].totalUsage.input).toBe(1500); expect(nodes[0].ownUsage.input).toBe(1000); });
it("includes trailing messages after the last assistant in context usage", () => { const rlmDir = makeTempDir(); const child = writeChildSession(join(rlmDir, "sub-trail001"), "trailing", createUsage(1500, 500, 0.02)); child.sessionManager.appendMessage(createUserMessage("a trailing user message that has not reached the model"));
const nodes = loadContextTreeChildrenFromDisk(rlmDir, resolveContextWindow); const tokens = nodes[0].contextUsage?.tokens; // Last assistant usage is 2000; the trailing message adds an estimate on top. expect(tokens).toBeGreaterThan(2000); });
it("skips children that are already represented live", () => { const rlmDir = makeTempDir(); writeChildSession(join(rlmDir, "sub-live0001"), "live child", createUsage(100, 10, 0.01)); writeChildSession(join(rlmDir, "sub-done0001"), "done child", createUsage(200, 20, 0.02));
const nodes = loadContextTreeChildrenFromDisk(rlmDir, resolveContextWindow, new Set(["sub-live0001"])); expect(nodes.map((node) => node.id)).toEqual(["sub-done0001"]); });});
describe("AgentSession.getContextTree", () => { function createSession() { const settingsManager = SettingsManager.inMemory(); const sessionManager = SessionManager.inMemory(); const authStorage = AuthStorage.inMemory(); authStorage.setRuntimeApiKey("anthropic", "test-key"); const session = new AgentSession({ agent: new Agent({ getApiKey: () => "test-key", initialState: { model, systemPrompt: "You are a helpful assistant.", tools: [], thinkingLevel: "high", }, }), sessionManager, settingsManager, cwd: process.cwd(), modelRegistry: ModelRegistry.inMemory(authStorage), resourceLoader: createTestResourceLoader(), }); return { session, sessionManager }; }
function syncAgentMessages(session: AgentSession, sessionManager: SessionManager): void { session.agent.state.messages = sessionManager.buildSessionContext().messages; }
it("returns a root node whose own usage excludes attributed child usage", () => { const { session, sessionManager } = createSession(); sessionManager.appendMessage(createUserMessage("do the thing")); const assistantEntryId = sessionManager.appendMessage( createAssistantMessage("on it", createUsage(3000, 600, 0.3)), );
const childUsage = createUsage(500, 100, 0.05); const aggregate = createUsage(3500, 700, 0.35); sessionManager.appendChildUsageAttribution(assistantEntryId, childUsage, aggregate); syncAgentMessages(session, sessionManager);
const tree = session.getContextTree(); expect(tree.id).toBe("root"); expect(tree.status).toBe("active"); expect(tree.model).toEqual({ provider: model.provider, id: model.id }); expect(tree.totalUsage.input).toBe(3500); expect(tree.totalUsage.cost.total).toBeCloseTo(0.35); expect(tree.ownUsage.input).toBe(3000); expect(tree.ownUsage.cost.total).toBeCloseTo(0.3); // In-memory session, no rlm dir: completed children cannot be discovered. expect(tree.children).toEqual([]); });
it("keeps pre-compaction spend in the totals after compaction", () => { // Intentional: /context reports cumulative session spend. Compaction // shrinks the model-facing context, but tokens already paid for must not // vanish from the totals (the old /usage undercounted here). const { session, sessionManager } = createSession(); sessionManager.appendMessage(createUserMessage("expensive early work")); sessionManager.appendMessage(createAssistantMessage("done", createUsage(5000, 1000, 0.5))); const keptId = sessionManager.appendMessage(createUserMessage("later work")); sessionManager.appendCompaction("summary of early work", keptId, 6000); sessionManager.appendMessage(createAssistantMessage("after compaction", createUsage(200, 50, 0.02))); syncAgentMessages(session, sessionManager);
const tree = session.getContextTree(); expect(tree.totalUsage.input).toBe(5200); expect(tree.totalUsage.cost.total).toBeCloseTo(0.52); expect(tree.ownUsage.input).toBe(5200); });});
describe("formatContextTree", () => { function node(overrides: Partial<ContextTreeNode>): ContextTreeNode { return { id: "root", label: "main agent", status: "active", ownUsage: emptyUsage(), totalUsage: emptyUsage(), children: [], ...overrides, }; }
it("renders a tree with per-agent rows and own-usage totals", () => { const root = node({ model: { provider: "anthropic", id: "claude-sonnet-4-5" }, ownUsage: createUsage(40000, 5200, 0.84), totalUsage: createUsage(52100, 5950, 1.1), contextUsage: { tokens: 62300, contextWindow: 200000, percent: 31.15 }, children: [ node({ id: "sub-aaaa1111", label: "summarize auth module", status: "done", ownUsage: createUsage(12100, 750, 0.21), totalUsage: createUsage(12100, 750, 0.21), contextUsage: { tokens: 18000, contextWindow: 200000, percent: 9 }, }), node({ id: "sub-bbbb2222", label: "refactor tests", status: "running", ownUsage: createUsage(8000, 700, 0.13), totalUsage: createUsage(8000, 700, 0.13), }), ], });
const output = stripAnsi(formatContextTree(root, 100)); const lines = output.split("\n");
expect(lines[0]).toBe("Context"); expect(output).toContain("Model: anthropic/claude-sonnet-4-5"); expect(output).toContain("● main agent"); expect(output).toContain("├─ ✓ summarize auth module"); expect(output).toContain("└─ ◆ refactor tests"); // Own usage per row: 45200 -> 45k, 12850 -> 13k, 8700 -> 8.7k. expect(output).toContain("45k"); expect(output).toContain("$0.84"); expect(output).toContain("31% (62k/200k)"); // Totals sum own usage across all agents: 66750 -> 67k, $1.18. expect(output).toContain("Total: 67k tokens · $1.18 across 3 agents"); // Detail sections carry the precise numbers /usage used to show. expect(output).toContain("Input: 60,100"); expect(output).toContain("Output: 6,650"); expect(output).toContain("Total: 66,750"); expect(output).toContain("Cost\nTotal: $1.1800"); expect(output).toContain("Context\nCurrent: 62,300 / 200,000 (31.2%)"); });
it("reports unknown context after compaction in the detail section", () => { const root = node({ ownUsage: createUsage(900, 100, 0.01), totalUsage: createUsage(900, 100, 0.01), contextUsage: { tokens: null, contextWindow: 200000, percent: null }, }); const output = stripAnsi(formatContextTree(root, 100)); expect(output).toContain("Current: unknown after compaction"); });
it("renders a lone root without tree glyphs or agent count", () => { const root = node({ ownUsage: createUsage(900, 100, 0.01), totalUsage: createUsage(900, 100, 0.01) }); const output = stripAnsi(formatContextTree(root, 100)); expect(output).toContain("● main agent"); expect(output).not.toContain("├"); expect(output).toContain("Total: 1.0k tokens · $0.01"); expect(output).not.toContain("across"); });});