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.
8.1 kB · 247 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248import type { AssistantMessage, ImageContent, ThinkingContent, ToolCall, ToolResultMessage, UserMessage,} from "@earendil-works/pi-ai";import { describe, expect, it } from "vitest";import type { SessionEntry, SessionTreeFlatNode } from "../src/core/session-manager.js";import { buildSessionTreeFromFlatNodes, formatTreeToolCall, projectSessionTreeEntry, projectSessionTreeFlatNodes, TREE_MESSAGE_TEXT_LIMIT,} from "../src/modes/agent-connection/session-tree-projection.js";
const userMessage = (id: string, parentId: string | null, content: UserMessage["content"]): SessionEntry => ({ type: "message", id, parentId, timestamp: "2026-01-01T00:00:00.000Z", message: { role: "user", content, timestamp: 1 } as UserMessage,});
const assistantMessage = (id: string, parentId: string | null, content: AssistantMessage["content"]): SessionEntry => ({ type: "message", id, parentId, timestamp: "2026-01-01T00:00:01.000Z", message: { role: "assistant", content, api: "anthropic-messages", provider: "anthropic", model: "claude-sonnet-4", usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "toolUse", timestamp: 2, } as AssistantMessage,});
describe("session tree projection", () => { it("keeps the display line and drops image payloads", () => { const projected = projectSessionTreeEntry( userMessage("user-1", null, [ { type: "text", text: "hello" }, { type: "image", data: "aGVsbG8=", mimeType: "image/png" } satisfies ImageContent, ]), );
expect(projected).toEqual({ type: "message", id: "user-1", parentId: null, timestamp: "2026-01-01T00:00:00.000Z", role: "user", text: "hello", hasText: true, }); expect(JSON.stringify(projected)).not.toContain("aGVsbG8="); });
it("clips message text to the budget the tree already searched within", () => { const long = "x".repeat(TREE_MESSAGE_TEXT_LIMIT + 500); const projected = projectSessionTreeEntry(userMessage("user-1", null, long)); if (projected.type !== "message") throw new Error("expected a message entry"); expect(projected.text).toHaveLength(TREE_MESSAGE_TEXT_LIMIT); });
it("drops thinking blocks and keeps tool calls as one-line descriptions", () => { const projected = projectSessionTreeEntry( assistantMessage("asst-1", "user-1", [ { type: "thinking", thinking: "secret reasoning" } satisfies ThinkingContent, { type: "text", text: "working" }, { type: "toolCall", id: "call-1", name: "bash", arguments: { command: "npm run check" }, } satisfies ToolCall, ]), );
expect(projected).toEqual({ type: "message", id: "asst-1", parentId: "user-1", timestamp: "2026-01-01T00:00:01.000Z", role: "assistant", text: "working", hasText: true, stopReason: "toolUse", toolCalls: [{ id: "call-1", display: "[bash: npm run check]" }], }); expect(JSON.stringify(projected)).not.toContain("secret reasoning"); });
it("keeps the search text of a tool result but drops its structured details", () => { const entry: SessionEntry = { type: "message", id: "result-1", parentId: "asst-1", timestamp: "2026-01-01T00:00:02.000Z", message: { role: "toolResult", toolCallId: "call-1", toolName: "bash", content: [{ type: "text", text: "checks passed" }], details: { rawOutput: "x".repeat(5000) }, isError: false, timestamp: 3, } as ToolResultMessage, };
const projected = projectSessionTreeEntry(entry); expect(projected).toEqual({ type: "message", id: "result-1", parentId: "asst-1", timestamp: "2026-01-01T00:00:02.000Z", role: "toolResult", text: "checks passed", hasText: false, toolCallId: "call-1", toolName: "bash", }); expect(JSON.stringify(projected).length).toBeLessThan(JSON.stringify(entry).length / 10); });
it("keeps custom message and branch summary text whole so search results do not change", () => { const body = `head ${"y".repeat(5000)} tail`; const custom: SessionEntry = { type: "custom_message", id: "custom-1", parentId: null, timestamp: "2026-01-01T00:00:00.000Z", customType: "harness_digest", content: body, display: true, }; const summary: SessionEntry = { type: "branch_summary", id: "summary-1", parentId: "custom-1", timestamp: "2026-01-01T00:00:01.000Z", fromId: "custom-1", summary: body, };
expect(projectSessionTreeEntry(custom)).toMatchObject({ type: "custom_message", text: body }); expect(projectSessionTreeEntry(summary)).toMatchObject({ type: "branch_summary", text: body }); });
it("projects bookkeeping entries to their structure alone", () => { const states: SessionEntry[] = [ { type: "session_state", id: "s-1", parentId: null, timestamp: "t", state: { status: "active" } }, { type: "agent_status", id: "s-2", parentId: "s-1", timestamp: "t", status: { summary: "done", basedOnMessageCount: 1 }, }, { type: "git_state", id: "s-3", parentId: "s-2", timestamp: "t", git: { branch: "main" } }, ];
expect(states.map((entry) => projectSessionTreeEntry(entry))).toEqual([ { type: "session_state", id: "s-1", parentId: null, timestamp: "t" }, { type: "agent_status", id: "s-2", parentId: "s-1", timestamp: "t" }, { type: "git_state", id: "s-3", parentId: "s-2", timestamp: "t" }, ]); });
it("formats tool call previews the way the tree line does", () => { expect(formatTreeToolCall("edit", { path: "/home/dev/proj/file.ts" }, "/home/dev")).toBe( "[edit: ~/proj/file.ts]", ); expect(formatTreeToolCall("bash", { command: "a\nb" }, undefined)).toBe("[bash: a b]"); expect(formatTreeToolCall("bash", { command: "x".repeat(80) }, undefined)).toBe(`[bash: ${"x".repeat(50)}...]`); expect(formatTreeToolCall("custom", { a: 1 }, undefined)).toBe('[custom: {"a":1}]'); });
it("nests projected nodes by parent and orders siblings by timestamp", () => { const flatNodes: SessionTreeFlatNode[] = [ { entry: userMessage("late", "root", "second"), label: "later" }, { entry: userMessage("root", null, "first") }, { entry: userMessage("early", "root", "first child") }, ]; // Sibling order follows timestamps, so give the two children distinct ones. const timed = flatNodes.map((flatNode, index) => ({ ...flatNode, entry: { ...flatNode.entry, timestamp: `2026-01-01T00:00:0${index}.000Z` } as SessionEntry, }));
const tree = buildSessionTreeFromFlatNodes(projectSessionTreeFlatNodes(timed)); expect(tree).toHaveLength(1); expect(tree[0]!.entry.id).toBe("root"); expect(tree[0]!.children.map((child) => child.entry.id)).toEqual(["late", "early"]); expect(tree[0]!.children[0]!.label).toBe("later"); });
it("shrinks a transcript-shaped tree by orders of magnitude", () => { const entries: SessionEntry[] = []; let parentId: string | null = null; for (let index = 0; index < 40; index++) { const assistant = assistantMessage(`asst-${index}`, parentId, [ { type: "thinking", thinking: "t".repeat(2000) } satisfies ThinkingContent, { type: "text", text: `step ${index}` }, { type: "toolCall", id: `call-${index}`, name: "bash", arguments: { command: `cmd ${index}` }, } satisfies ToolCall, ]); const result: SessionEntry = { type: "message", id: `result-${index}`, parentId: `asst-${index}`, timestamp: "2026-01-01T00:00:00.000Z", message: { role: "toolResult", toolCallId: `call-${index}`, toolName: "bash", content: [{ type: "text", text: "o".repeat(8000) }], isError: false, timestamp: 4, } as ToolResultMessage, }; entries.push(assistant, result); parentId = `result-${index}`; }
const full = JSON.stringify(entries.map((entry) => ({ entry }))).length; const projected = JSON.stringify(projectSessionTreeFlatNodes(entries.map((entry) => ({ entry })))).length; expect(projected * 10).toBeLessThan(full); });});