Merge branch 'main' into pb/tui-status-coalesce

This commit is contained in:
Mario Zechner 2026-01-01 00:27:54 +01:00
commit ac6f5006a9
216 changed files with 14479 additions and 8725 deletions

View file

@ -10,7 +10,7 @@
import { existsSync, mkdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { Agent, ProviderTransport } from "@mariozechner/pi-agent-core";
import { Agent } from "@mariozechner/pi-agent-core";
import { getModel } from "@mariozechner/pi-ai";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { AgentSession } from "../src/core/agent-session.js";
@ -19,8 +19,7 @@ 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 { codingTools } from "../src/core/tools/index.js";
const API_KEY = process.env.ANTHROPIC_API_KEY || process.env.ANTHROPIC_OAUTH_TOKEN;
import { API_KEY } from "./utilities.js";
describe.skipIf(!API_KEY)("AgentSession branching", () => {
let session: AgentSession;
@ -44,13 +43,8 @@ describe.skipIf(!API_KEY)("AgentSession branching", () => {
function createSession(noSession: boolean = false) {
const model = getModel("anthropic", "claude-sonnet-4-5")!;
const transport = new ProviderTransport({
getApiKey: () => API_KEY,
});
const agent = new Agent({
transport,
getApiKey: () => API_KEY,
initialState: {
model,
systemPrompt: "You are a helpful assistant. Be extremely concise, reply with just a few words.",
@ -89,7 +83,7 @@ describe.skipIf(!API_KEY)("AgentSession branching", () => {
expect(userMessages[0].text).toBe("Say hello");
// Branch from the first message
const result = await session.branch(userMessages[0].entryIndex);
const result = await session.branch(userMessages[0].entryId);
expect(result.selectedText).toBe("Say hello");
expect(result.cancelled).toBe(false);
@ -105,7 +99,7 @@ describe.skipIf(!API_KEY)("AgentSession branching", () => {
createSession(true);
// Verify sessions are disabled
expect(session.sessionFile).toBeNull();
expect(session.sessionFile).toBeUndefined();
// Send one message
await session.prompt("Say hi");
@ -119,15 +113,15 @@ describe.skipIf(!API_KEY)("AgentSession branching", () => {
expect(session.messages.length).toBeGreaterThan(0);
// Branch from the first message
const result = await session.branch(userMessages[0].entryIndex);
const result = await session.branch(userMessages[0].entryId);
expect(result.selectedText).toBe("Say hi");
expect(result.cancelled).toBe(false);
// After branching, conversation should be empty
expect(session.messages.length).toBe(0);
// Session file should still be null (no file created)
expect(session.sessionFile).toBeNull();
// Session file should still be undefined (no file created)
expect(session.sessionFile).toBeUndefined();
});
it("should branch from middle of conversation", async () => {
@ -149,7 +143,7 @@ describe.skipIf(!API_KEY)("AgentSession branching", () => {
// Branch from second message (keeps first message + response)
const secondMessage = userMessages[1];
const result = await session.branch(secondMessage.entryIndex);
const result = await session.branch(secondMessage.entryId);
expect(result.selectedText).toBe("Say two");
// After branching, should have first user message + assistant response

View file

@ -10,7 +10,7 @@
import { existsSync, mkdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { Agent, ProviderTransport } from "@mariozechner/pi-agent-core";
import { Agent } from "@mariozechner/pi-agent-core";
import { getModel } from "@mariozechner/pi-ai";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { AgentSession, type AgentSessionEvent } from "../src/core/agent-session.js";
@ -19,8 +19,7 @@ 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 { codingTools } from "../src/core/tools/index.js";
const API_KEY = process.env.ANTHROPIC_API_KEY || process.env.ANTHROPIC_OAUTH_TOKEN;
import { API_KEY } from "./utilities.js";
describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => {
let session: AgentSession;
@ -46,15 +45,10 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => {
}
});
function createSession() {
function createSession(inMemory = false) {
const model = getModel("anthropic", "claude-sonnet-4-5")!;
const transport = new ProviderTransport({
getApiKey: () => API_KEY,
});
const agent = new Agent({
transport,
getApiKey: () => API_KEY,
initialState: {
model,
systemPrompt: "You are a helpful assistant. Be concise.",
@ -62,8 +56,10 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => {
},
});
sessionManager = SessionManager.create(tempDir);
sessionManager = inMemory ? SessionManager.inMemory() : SessionManager.create(tempDir);
const settingsManager = SettingsManager.create(tempDir, tempDir);
// Use minimal keepRecentTokens so small test conversations have something to summarize
settingsManager.applyOverrides({ compaction: { keepRecentTokens: 1 } });
const authStorage = new AuthStorage(join(tempDir, "auth.json"));
const modelRegistry = new ModelRegistry(authStorage);
@ -105,7 +101,7 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => {
// First message should be the summary (a user message with summary content)
const firstMsg = messages[0];
expect(firstMsg.role).toBe("user");
expect(firstMsg.role).toBe("compactionSummary");
}, 120000);
it("should maintain valid session state after compaction", async () => {
@ -156,64 +152,31 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => {
expect(compaction.type).toBe("compaction");
if (compaction.type === "compaction") {
expect(compaction.summary.length).toBeGreaterThan(0);
// firstKeptEntryIndex can be 0 if all messages fit within keepRecentTokens
// (which is the case for small conversations)
expect(compaction.firstKeptEntryIndex).toBeGreaterThanOrEqual(0);
expect(typeof compaction.firstKeptEntryId).toBe("string");
expect(compaction.tokensBefore).toBeGreaterThan(0);
}
}, 120000);
it("should work with --no-session mode (in-memory only)", async () => {
const model = getModel("anthropic", "claude-sonnet-4-5")!;
createSession(true); // in-memory mode
const transport = new ProviderTransport({
getApiKey: () => API_KEY,
});
// Send prompts
await session.prompt("What is 2+2? Reply with just the number.");
await session.agent.waitForIdle();
const agent = new Agent({
transport,
initialState: {
model,
systemPrompt: "You are a helpful assistant. Be concise.",
tools: codingTools,
},
});
await session.prompt("What is 3+3? Reply with just the number.");
await session.agent.waitForIdle();
// Create in-memory session manager
const noSessionManager = SessionManager.inMemory();
// Compact should work even without file persistence
const result = await session.compact();
const settingsManager = SettingsManager.create(tempDir, tempDir);
const authStorage = new AuthStorage(join(tempDir, "auth.json"));
const modelRegistry = new ModelRegistry(authStorage);
expect(result.summary).toBeDefined();
expect(result.summary.length).toBeGreaterThan(0);
const noSessionSession = new AgentSession({
agent,
sessionManager: noSessionManager,
settingsManager,
modelRegistry,
});
try {
// Send prompts
await noSessionSession.prompt("What is 2+2? Reply with just the number.");
await noSessionSession.agent.waitForIdle();
await noSessionSession.prompt("What is 3+3? Reply with just the number.");
await noSessionSession.agent.waitForIdle();
// Compact should work even without file persistence
const result = await noSessionSession.compact();
expect(result.summary).toBeDefined();
expect(result.summary.length).toBeGreaterThan(0);
// In-memory entries should have the compaction
const entries = noSessionManager.getEntries();
const compactionEntries = entries.filter((e) => e.type === "compaction");
expect(compactionEntries.length).toBe(1);
} finally {
noSessionSession.dispose();
}
// In-memory entries should have the compaction
const entries = sessionManager.getEntries();
const compactionEntries = entries.filter((e) => e.type === "compaction");
expect(compactionEntries.length).toBe(1);
}, 120000);
it("should emit correct events during auto-compaction", async () => {

View file

@ -0,0 +1,318 @@
/**
* E2E tests for AgentSession tree navigation with branch summarization.
*
* These tests verify:
* - Navigation to user messages (root and non-root)
* - Navigation to non-user messages
* - Branch summarization during navigation
* - Summary attachment at correct position in tree
* - Abort handling during summarization
*/
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { API_KEY, createTestSession, type TestSessionContext } from "./utilities.js";
describe.skipIf(!API_KEY)("AgentSession tree navigation e2e", () => {
let ctx: TestSessionContext;
beforeEach(() => {
ctx = createTestSession({
systemPrompt: "You are a helpful assistant. Reply with just a few words.",
settingsOverrides: { compaction: { keepRecentTokens: 1 } },
});
});
afterEach(() => {
ctx.cleanup();
});
it("should navigate to user message and put text in editor", async () => {
const { session } = ctx;
// Build conversation: u1 -> a1 -> u2 -> a2
await session.prompt("First message");
await session.agent.waitForIdle();
await session.prompt("Second message");
await session.agent.waitForIdle();
// Get tree entries
const tree = session.sessionManager.getTree();
expect(tree.length).toBe(1);
// Find the first user entry (u1)
const rootNode = tree[0];
expect(rootNode.entry.type).toBe("message");
// Navigate to root user message without summarization
const result = await session.navigateTree(rootNode.entry.id, { summarize: false });
expect(result.cancelled).toBe(false);
expect(result.editorText).toBe("First message");
// After navigating to root user message, leaf should be null (empty conversation)
expect(session.sessionManager.getLeafId()).toBeNull();
}, 60000);
it("should navigate to non-user message without editor text", async () => {
const { session, sessionManager } = ctx;
// Build conversation
await session.prompt("Hello");
await session.agent.waitForIdle();
// Get the assistant message
const entries = sessionManager.getEntries();
const assistantEntry = entries.find((e) => e.type === "message" && e.message.role === "assistant");
expect(assistantEntry).toBeDefined();
// Navigate to assistant message
const result = await session.navigateTree(assistantEntry!.id, { summarize: false });
expect(result.cancelled).toBe(false);
expect(result.editorText).toBeUndefined();
// Leaf should be the assistant entry
expect(sessionManager.getLeafId()).toBe(assistantEntry!.id);
}, 60000);
it("should create branch summary when navigating with summarize=true", async () => {
const { session, sessionManager } = ctx;
// Build conversation: u1 -> a1 -> u2 -> a2
await session.prompt("What is 2+2?");
await session.agent.waitForIdle();
await session.prompt("What is 3+3?");
await session.agent.waitForIdle();
// Get tree and find first user message
const tree = sessionManager.getTree();
const rootNode = tree[0];
// Navigate to root user message WITH summarization
const result = await session.navigateTree(rootNode.entry.id, { summarize: true });
expect(result.cancelled).toBe(false);
expect(result.editorText).toBe("What is 2+2?");
expect(result.summaryEntry).toBeDefined();
expect(result.summaryEntry?.type).toBe("branch_summary");
expect(result.summaryEntry?.summary).toBeTruthy();
expect(result.summaryEntry?.summary.length).toBeGreaterThan(0);
// Summary should be a root entry (parentId = null) since we navigated to root user
expect(result.summaryEntry?.parentId).toBeNull();
// Leaf should be the summary entry
expect(sessionManager.getLeafId()).toBe(result.summaryEntry?.id);
}, 120000);
it("should attach summary to correct parent when navigating to nested user message", async () => {
const { session, sessionManager } = ctx;
// Build conversation: u1 -> a1 -> u2 -> a2 -> u3 -> a3
await session.prompt("Message one");
await session.agent.waitForIdle();
await session.prompt("Message two");
await session.agent.waitForIdle();
await session.prompt("Message three");
await session.agent.waitForIdle();
// Get the second user message (u2)
const entries = sessionManager.getEntries();
const userEntries = entries.filter((e) => e.type === "message" && e.message.role === "user");
expect(userEntries.length).toBe(3);
const u2 = userEntries[1];
const a1 = entries.find((e) => e.id === u2.parentId); // a1 is parent of u2
// Navigate to u2 with summarization
const result = await session.navigateTree(u2.id, { summarize: true });
expect(result.cancelled).toBe(false);
expect(result.editorText).toBe("Message two");
expect(result.summaryEntry).toBeDefined();
// Summary should be attached to a1 (parent of u2)
// So a1 now has two children: u2 and the summary
expect(result.summaryEntry?.parentId).toBe(a1?.id);
// Verify tree structure
const children = sessionManager.getChildren(a1!.id);
expect(children.length).toBe(2);
const childTypes = children.map((c) => c.type).sort();
expect(childTypes).toContain("branch_summary");
expect(childTypes).toContain("message");
}, 120000);
it("should attach summary to selected node when navigating to assistant message", async () => {
const { session, sessionManager } = ctx;
// Build conversation: u1 -> a1 -> u2 -> a2
await session.prompt("Hello");
await session.agent.waitForIdle();
await session.prompt("Goodbye");
await session.agent.waitForIdle();
// Get the first assistant message (a1)
const entries = sessionManager.getEntries();
const assistantEntries = entries.filter((e) => e.type === "message" && e.message.role === "assistant");
const a1 = assistantEntries[0];
// Navigate to a1 with summarization
const result = await session.navigateTree(a1.id, { summarize: true });
expect(result.cancelled).toBe(false);
expect(result.editorText).toBeUndefined(); // No editor text for assistant messages
expect(result.summaryEntry).toBeDefined();
// Summary should be attached to a1 (the selected node)
expect(result.summaryEntry?.parentId).toBe(a1.id);
// Leaf should be the summary entry
expect(sessionManager.getLeafId()).toBe(result.summaryEntry?.id);
}, 120000);
it("should handle abort during summarization", async () => {
const { session, sessionManager } = ctx;
// Build conversation
await session.prompt("Tell me about something");
await session.agent.waitForIdle();
await session.prompt("Continue");
await session.agent.waitForIdle();
const entriesBefore = sessionManager.getEntries();
const leafBefore = sessionManager.getLeafId();
// Get root user message
const tree = sessionManager.getTree();
const rootNode = tree[0];
// Start navigation with summarization but abort immediately
const navigationPromise = session.navigateTree(rootNode.entry.id, { summarize: true });
// Abort after a short delay (let the LLM call start)
await new Promise((resolve) => setTimeout(resolve, 100));
session.abortBranchSummary();
const result = await navigationPromise;
expect(result.cancelled).toBe(true);
expect(result.aborted).toBe(true);
expect(result.summaryEntry).toBeUndefined();
// Session should be unchanged
const entriesAfter = sessionManager.getEntries();
expect(entriesAfter.length).toBe(entriesBefore.length);
expect(sessionManager.getLeafId()).toBe(leafBefore);
}, 60000);
it("should not create summary when navigating without summarize option", async () => {
const { session, sessionManager } = ctx;
// Build conversation
await session.prompt("First");
await session.agent.waitForIdle();
await session.prompt("Second");
await session.agent.waitForIdle();
const entriesBefore = sessionManager.getEntries().length;
// Navigate without summarization
const tree = sessionManager.getTree();
await session.navigateTree(tree[0].entry.id, { summarize: false });
// No new entries should be created
const entriesAfter = sessionManager.getEntries().length;
expect(entriesAfter).toBe(entriesBefore);
// No branch_summary entries
const summaries = sessionManager.getEntries().filter((e) => e.type === "branch_summary");
expect(summaries.length).toBe(0);
}, 60000);
it("should handle navigation to same position (no-op)", async () => {
const { session, sessionManager } = ctx;
// Build conversation
await session.prompt("Hello");
await session.agent.waitForIdle();
const leafBefore = sessionManager.getLeafId();
expect(leafBefore).toBeTruthy();
const entriesBefore = sessionManager.getEntries().length;
// Navigate to current leaf
const result = await session.navigateTree(leafBefore!, { summarize: false });
expect(result.cancelled).toBe(false);
expect(sessionManager.getLeafId()).toBe(leafBefore);
expect(sessionManager.getEntries().length).toBe(entriesBefore);
}, 60000);
it("should support custom summarization instructions", async () => {
const { session, sessionManager } = ctx;
// Build conversation
await session.prompt("What is TypeScript?");
await session.agent.waitForIdle();
// Navigate with custom instructions
const tree = sessionManager.getTree();
const result = await session.navigateTree(tree[0].entry.id, {
summarize: true,
customInstructions: "Summarize in exactly 3 words.",
});
expect(result.summaryEntry).toBeDefined();
expect(result.summaryEntry?.summary).toBeTruthy();
// Can't reliably test 3 words exactly, but summary should be short
expect(result.summaryEntry?.summary.split(/\s+/).length).toBeLessThan(20);
}, 120000);
});
describe.skipIf(!API_KEY)("AgentSession tree navigation - branch scenarios", () => {
let ctx: TestSessionContext;
beforeEach(() => {
ctx = createTestSession({
systemPrompt: "You are a helpful assistant. Reply with just a few words.",
});
});
afterEach(() => {
ctx.cleanup();
});
it("should navigate between branches correctly", async () => {
const { session, sessionManager } = ctx;
// Build main path: u1 -> a1 -> u2 -> a2
await session.prompt("Main branch start");
await session.agent.waitForIdle();
await session.prompt("Main branch continue");
await session.agent.waitForIdle();
// Get a1 id for branching
const entries = sessionManager.getEntries();
const a1 = entries.find((e) => e.type === "message" && e.message.role === "assistant");
// Create a branch from a1: a1 -> u3 -> a3
sessionManager.branch(a1!.id);
await session.prompt("Branch path");
await session.agent.waitForIdle();
// Now navigate back to u2 (on main branch) with summarization
const userEntries = entries.filter((e) => e.type === "message" && e.message.role === "user");
const u2 = userEntries[1]; // "Main branch continue"
const result = await session.navigateTree(u2.id, { summarize: true });
expect(result.cancelled).toBe(false);
expect(result.editorText).toBe("Main branch continue");
expect(result.summaryEntry).toBeDefined();
// Summary captures the branch we're leaving (the "Branch path" conversation)
expect(result.summaryEntry?.summary.length).toBeGreaterThan(0);
}, 180000);
});

View file

@ -3,46 +3,43 @@
*/
import { describe, expect, it } from "vitest";
import type { HookAPI } from "../src/core/hooks/index.js";
import type { CompactionEntry } from "../src/core/session-manager.js";
import type { HookAPI, SessionBeforeCompactEvent, SessionCompactEvent } from "../src/core/hooks/index.js";
describe("Documentation example", () => {
it("custom compaction example should type-check correctly", () => {
// This is the example from hooks.md - verify it compiles
const exampleHook = (pi: HookAPI) => {
pi.on("session", async (event, _ctx) => {
if (event.reason !== "before_compact") return;
// After narrowing, these should all be accessible
const messages = event.messagesToSummarize;
const messagesToKeep = event.messagesToKeep;
const cutPoint = event.cutPoint;
const tokensBefore = event.tokensBefore;
const model = event.model;
const resolveApiKey = event.resolveApiKey;
pi.on("session_before_compact", async (event: SessionBeforeCompactEvent, ctx) => {
// All these should be accessible on the event
const { preparation, branchEntries } = event;
// sessionManager, modelRegistry, and model come from ctx
const { sessionManager, modelRegistry } = ctx;
const { messagesToSummarize, turnPrefixMessages, tokensBefore, firstKeptEntryId, isSplitTurn } =
preparation;
// Verify types
expect(Array.isArray(messages)).toBe(true);
expect(Array.isArray(messagesToKeep)).toBe(true);
expect(typeof cutPoint.firstKeptEntryIndex).toBe("number");
expect(Array.isArray(messagesToSummarize)).toBe(true);
expect(Array.isArray(turnPrefixMessages)).toBe(true);
expect(typeof isSplitTurn).toBe("boolean");
expect(typeof tokensBefore).toBe("number");
expect(model).toBeDefined();
expect(typeof resolveApiKey).toBe("function");
expect(typeof sessionManager.getEntries).toBe("function");
expect(typeof modelRegistry.getApiKey).toBe("function");
expect(typeof firstKeptEntryId).toBe("string");
expect(Array.isArray(branchEntries)).toBe(true);
const summary = messages
const summary = messagesToSummarize
.filter((m) => m.role === "user")
.map((m) => `- ${typeof m.content === "string" ? m.content.slice(0, 100) : "[complex]"}`)
.join("\n");
const compactionEntry: CompactionEntry = {
type: "compaction",
timestamp: new Date().toISOString(),
summary: `User requests:\n${summary}`,
firstKeptEntryIndex: event.cutPoint.firstKeptEntryIndex,
tokensBefore: event.tokensBefore,
// Hooks return compaction content - SessionManager adds id/parentId
return {
compaction: {
summary: `User requests:\n${summary}`,
firstKeptEntryId,
tokensBefore,
},
};
return { compactionEntry };
});
};
@ -50,19 +47,16 @@ describe("Documentation example", () => {
expect(typeof exampleHook).toBe("function");
});
it("compact event should have correct fields after narrowing", () => {
it("compact event should have correct fields", () => {
const checkCompactEvent = (pi: HookAPI) => {
pi.on("session", async (event, _ctx) => {
if (event.reason !== "compact") return;
// After narrowing, these should all be accessible
pi.on("session_compact", async (event: SessionCompactEvent) => {
// These should all be accessible
const entry = event.compactionEntry;
const tokensBefore = event.tokensBefore;
const fromHook = event.fromHook;
expect(entry.type).toBe("compaction");
expect(typeof entry.summary).toBe("string");
expect(typeof tokensBefore).toBe("number");
expect(typeof entry.tokensBefore).toBe("number");
expect(typeof fromHook).toBe("boolean");
});
};

View file

@ -5,18 +5,24 @@
import { existsSync, mkdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { Agent, ProviderTransport } from "@mariozechner/pi-agent-core";
import { Agent } from "@mariozechner/pi-agent-core";
import { getModel } from "@mariozechner/pi-ai";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { AgentSession } from "../src/core/agent-session.js";
import { AuthStorage } from "../src/core/auth-storage.js";
import { HookRunner, type LoadedHook, type SessionEvent } from "../src/core/hooks/index.js";
import {
HookRunner,
type LoadedHook,
type SessionBeforeCompactEvent,
type SessionCompactEvent,
type SessionEvent,
} from "../src/core/hooks/index.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 { codingTools } from "../src/core/tools/index.js";
const API_KEY = process.env.ANTHROPIC_API_KEY || process.env.ANTHROPIC_OAUTH_TOKEN;
const API_KEY = process.env.ANTHROPIC_OAUTH_TOKEN || process.env.ANTHROPIC_API_KEY;
describe.skipIf(!API_KEY)("Compaction hooks", () => {
let session: AgentSession;
@ -40,19 +46,25 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => {
});
function createHook(
onBeforeCompact?: (event: SessionEvent) => { cancel?: boolean; compactionEntry?: any } | undefined,
onCompact?: (event: SessionEvent) => void,
onBeforeCompact?: (event: SessionBeforeCompactEvent) => { cancel?: boolean; compaction?: any } | undefined,
onCompact?: (event: SessionCompactEvent) => void,
): LoadedHook {
const handlers = new Map<string, ((event: any, ctx: any) => Promise<any>)[]>();
handlers.set("session", [
async (event: SessionEvent) => {
handlers.set("session_before_compact", [
async (event: SessionBeforeCompactEvent) => {
capturedEvents.push(event);
if (event.reason === "before_compact" && onBeforeCompact) {
if (onBeforeCompact) {
return onBeforeCompact(event);
}
if (event.reason === "compact" && onCompact) {
return undefined;
},
]);
handlers.set("session_compact", [
async (event: SessionCompactEvent) => {
capturedEvents.push(event);
if (onCompact) {
onCompact(event);
}
return undefined;
@ -63,19 +75,17 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => {
path: "test-hook",
resolvedPath: "/test/test-hook.ts",
handlers,
setSendHandler: () => {},
messageRenderers: new Map(),
commands: new Map(),
setSendMessageHandler: () => {},
setAppendEntryHandler: () => {},
};
}
function createSession(hooks: LoadedHook[]) {
const model = getModel("anthropic", "claude-sonnet-4-5")!;
const transport = new ProviderTransport({
getApiKey: () => API_KEY,
});
const agent = new Agent({
transport,
getApiKey: () => API_KEY,
initialState: {
model,
systemPrompt: "You are a helpful assistant. Be concise.",
@ -88,17 +98,22 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => {
const authStorage = new AuthStorage(join(tempDir, "auth.json"));
const modelRegistry = new ModelRegistry(authStorage);
hookRunner = new HookRunner(hooks, tempDir);
hookRunner.setUIContext(
{
select: async () => null,
hookRunner = new HookRunner(hooks, tempDir, sessionManager, modelRegistry);
hookRunner.initialize({
getModel: () => session.model,
sendMessageHandler: async () => {},
appendEntryHandler: async () => {},
uiContext: {
select: async () => undefined,
confirm: async () => false,
input: async () => null,
input: async () => undefined,
notify: () => {},
custom: async () => undefined as never,
setEditorText: () => {},
getEditorText: () => "",
},
false,
);
hookRunner.setSessionFile(sessionManager.getSessionFile());
hasUI: false,
});
session = new AgentSession({
agent,
@ -123,30 +138,28 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => {
await session.compact();
const beforeCompactEvents = capturedEvents.filter((e) => e.reason === "before_compact");
const compactEvents = capturedEvents.filter((e) => e.reason === "compact");
const beforeCompactEvents = capturedEvents.filter(
(e): e is SessionBeforeCompactEvent => e.type === "session_before_compact",
);
const compactEvents = capturedEvents.filter((e): e is SessionCompactEvent => e.type === "session_compact");
expect(beforeCompactEvents.length).toBe(1);
expect(compactEvents.length).toBe(1);
const beforeEvent = beforeCompactEvents[0];
if (beforeEvent.reason === "before_compact") {
expect(beforeEvent.cutPoint).toBeDefined();
expect(beforeEvent.cutPoint.firstKeptEntryIndex).toBeGreaterThanOrEqual(0);
expect(beforeEvent.messagesToSummarize).toBeDefined();
expect(beforeEvent.messagesToKeep).toBeDefined();
expect(beforeEvent.tokensBefore).toBeGreaterThanOrEqual(0);
expect(beforeEvent.model).toBeDefined();
expect(beforeEvent.resolveApiKey).toBeDefined();
}
expect(beforeEvent.preparation).toBeDefined();
expect(beforeEvent.preparation.messagesToSummarize).toBeDefined();
expect(beforeEvent.preparation.turnPrefixMessages).toBeDefined();
expect(beforeEvent.preparation.tokensBefore).toBeGreaterThanOrEqual(0);
expect(typeof beforeEvent.preparation.isSplitTurn).toBe("boolean");
expect(beforeEvent.branchEntries).toBeDefined();
// sessionManager, modelRegistry, and model are now on ctx, not event
const afterEvent = compactEvents[0];
if (afterEvent.reason === "compact") {
expect(afterEvent.compactionEntry).toBeDefined();
expect(afterEvent.compactionEntry.summary.length).toBeGreaterThan(0);
expect(afterEvent.tokensBefore).toBeGreaterThanOrEqual(0);
expect(afterEvent.fromHook).toBe(false);
}
expect(afterEvent.compactionEntry).toBeDefined();
expect(afterEvent.compactionEntry.summary.length).toBeGreaterThan(0);
expect(afterEvent.compactionEntry.tokensBefore).toBeGreaterThanOrEqual(0);
expect(afterEvent.fromHook).toBe(false);
}, 120000);
it("should allow hooks to cancel compaction", async () => {
@ -158,22 +171,20 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => {
await expect(session.compact()).rejects.toThrow("Compaction cancelled");
const compactEvents = capturedEvents.filter((e) => e.reason === "compact");
const compactEvents = capturedEvents.filter((e) => e.type === "session_compact");
expect(compactEvents.length).toBe(0);
}, 120000);
it("should allow hooks to provide custom compactionEntry", async () => {
it("should allow hooks to provide custom compaction", async () => {
const customSummary = "Custom summary from hook";
const hook = createHook((event) => {
if (event.reason === "before_compact") {
if (event.type === "session_before_compact") {
return {
compactionEntry: {
type: "compaction" as const,
timestamp: new Date().toISOString(),
compaction: {
summary: customSummary,
firstKeptEntryIndex: event.cutPoint.firstKeptEntryIndex,
tokensBefore: event.tokensBefore,
firstKeptEntryId: event.preparation.firstKeptEntryId,
tokensBefore: event.preparation.tokensBefore,
},
};
}
@ -191,11 +202,11 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => {
expect(result.summary).toBe(customSummary);
const compactEvents = capturedEvents.filter((e) => e.reason === "compact");
const compactEvents = capturedEvents.filter((e) => e.type === "session_compact");
expect(compactEvents.length).toBe(1);
const afterEvent = compactEvents[0];
if (afterEvent.reason === "compact") {
if (afterEvent.type === "session_compact") {
expect(afterEvent.compactionEntry.summary).toBe(customSummary);
expect(afterEvent.fromHook).toBe(true);
}
@ -210,12 +221,14 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => {
await session.compact();
const compactEvents = capturedEvents.filter((e) => e.reason === "compact");
const compactEvents = capturedEvents.filter((e) => e.type === "session_compact");
expect(compactEvents.length).toBe(1);
const afterEvent = compactEvents[0];
if (afterEvent.reason === "compact") {
const hasCompactionEntry = afterEvent.entries.some((e) => e.type === "compaction");
if (afterEvent.type === "session_compact") {
// sessionManager is now on ctx, use session.sessionManager directly
const entries = session.sessionManager.getEntries();
const hasCompactionEntry = entries.some((e: { type: string }) => e.type === "compaction");
expect(hasCompactionEntry).toBe(true);
}
}, 120000);
@ -226,19 +239,28 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => {
resolvedPath: "/test/throwing-hook.ts",
handlers: new Map<string, ((event: any, ctx: any) => Promise<any>)[]>([
[
"session",
"session_before_compact",
[
async (event: SessionEvent) => {
async (event: SessionBeforeCompactEvent) => {
capturedEvents.push(event);
throw new Error("Hook intentionally throws");
},
],
],
[
"session_compact",
[
async (event: SessionCompactEvent) => {
capturedEvents.push(event);
if (event.reason === "before_compact") {
throw new Error("Hook intentionally failed");
}
return undefined;
},
],
],
]),
setSendHandler: () => {},
messageRenderers: new Map(),
commands: new Map(),
setSendMessageHandler: () => {},
setAppendEntryHandler: () => {},
};
createSession([throwingHook]);
@ -251,12 +273,9 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => {
expect(result.summary).toBeDefined();
expect(result.summary.length).toBeGreaterThan(0);
const compactEvents = capturedEvents.filter((e) => e.reason === "compact");
const compactEvents = capturedEvents.filter((e): e is SessionCompactEvent => e.type === "session_compact");
expect(compactEvents.length).toBe(1);
if (compactEvents[0].reason === "compact") {
expect(compactEvents[0].fromHook).toBe(false);
}
expect(compactEvents[0].fromHook).toBe(false);
}, 120000);
it("should call multiple hooks in order", async () => {
@ -267,21 +286,28 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => {
resolvedPath: "/test/hook1.ts",
handlers: new Map<string, ((event: any, ctx: any) => Promise<any>)[]>([
[
"session",
"session_before_compact",
[
async (event: SessionEvent) => {
if (event.reason === "before_compact") {
callOrder.push("hook1-before");
}
if (event.reason === "compact") {
callOrder.push("hook1-after");
}
async () => {
callOrder.push("hook1-before");
return undefined;
},
],
],
[
"session_compact",
[
async () => {
callOrder.push("hook1-after");
return undefined;
},
],
],
]),
setSendHandler: () => {},
messageRenderers: new Map(),
commands: new Map(),
setSendMessageHandler: () => {},
setAppendEntryHandler: () => {},
};
const hook2: LoadedHook = {
@ -289,21 +315,28 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => {
resolvedPath: "/test/hook2.ts",
handlers: new Map<string, ((event: any, ctx: any) => Promise<any>)[]>([
[
"session",
"session_before_compact",
[
async (event: SessionEvent) => {
if (event.reason === "before_compact") {
callOrder.push("hook2-before");
}
if (event.reason === "compact") {
callOrder.push("hook2-after");
}
async () => {
callOrder.push("hook2-before");
return undefined;
},
],
],
[
"session_compact",
[
async () => {
callOrder.push("hook2-after");
return undefined;
},
],
],
]),
setSendHandler: () => {},
messageRenderers: new Map(),
commands: new Map(),
setSendMessageHandler: () => {},
setAppendEntryHandler: () => {},
};
createSession([hook1, hook2]);
@ -317,12 +350,10 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => {
}, 120000);
it("should pass correct data in before_compact event", async () => {
let capturedBeforeEvent: (SessionEvent & { reason: "before_compact" }) | null = null;
let capturedBeforeEvent: SessionBeforeCompactEvent | null = null;
const hook = createHook((event) => {
if (event.reason === "before_compact") {
capturedBeforeEvent = event;
}
capturedBeforeEvent = event;
return undefined;
});
createSession([hook]);
@ -337,35 +368,35 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => {
expect(capturedBeforeEvent).not.toBeNull();
const event = capturedBeforeEvent!;
expect(event.cutPoint).toHaveProperty("firstKeptEntryIndex");
expect(event.cutPoint).toHaveProperty("isSplitTurn");
expect(event.cutPoint).toHaveProperty("turnStartIndex");
expect(typeof event.preparation.isSplitTurn).toBe("boolean");
expect(event.preparation.firstKeptEntryId).toBeDefined();
expect(Array.isArray(event.messagesToSummarize)).toBe(true);
expect(Array.isArray(event.messagesToKeep)).toBe(true);
expect(Array.isArray(event.preparation.messagesToSummarize)).toBe(true);
expect(Array.isArray(event.preparation.turnPrefixMessages)).toBe(true);
expect(typeof event.tokensBefore).toBe("number");
expect(typeof event.preparation.tokensBefore).toBe("number");
expect(event.model).toHaveProperty("provider");
expect(event.model).toHaveProperty("id");
expect(Array.isArray(event.branchEntries)).toBe(true);
expect(typeof event.resolveApiKey).toBe("function");
// sessionManager, modelRegistry, and model are now on ctx, not event
// Verify they're accessible via session
expect(typeof session.sessionManager.getEntries).toBe("function");
expect(typeof session.modelRegistry.getApiKey).toBe("function");
expect(Array.isArray(event.entries)).toBe(true);
expect(event.entries.length).toBeGreaterThan(0);
const entries = session.sessionManager.getEntries();
expect(Array.isArray(entries)).toBe(true);
expect(entries.length).toBeGreaterThan(0);
}, 120000);
it("should use hook compactionEntry even with different firstKeptEntryIndex", async () => {
const customSummary = "Custom summary with modified index";
it("should use hook compaction even with different values", async () => {
const customSummary = "Custom summary with modified values";
const hook = createHook((event) => {
if (event.reason === "before_compact") {
if (event.type === "session_before_compact") {
return {
compactionEntry: {
type: "compaction" as const,
timestamp: new Date().toISOString(),
compaction: {
summary: customSummary,
firstKeptEntryIndex: 0,
firstKeptEntryId: event.preparation.firstKeptEntryId,
tokensBefore: 999,
},
};

View file

@ -1,9 +1,9 @@
import type { AppMessage } from "@mariozechner/pi-agent-core";
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import type { AssistantMessage, Usage } from "@mariozechner/pi-ai";
import { getModel } from "@mariozechner/pi-ai";
import { readFileSync } from "fs";
import { join } from "path";
import { describe, expect, it } from "vitest";
import { beforeEach, describe, expect, it } from "vitest";
import {
type CompactionSettings,
calculateContextTokens,
@ -11,15 +11,18 @@ import {
DEFAULT_COMPACTION_SETTINGS,
findCutPoint,
getLastAssistantUsage,
prepareCompaction,
shouldCompact,
} from "../src/core/compaction.js";
} from "../src/core/compaction/index.js";
import {
buildSessionContext,
type CompactionEntry,
createSummaryMessage,
type ModelChangeEntry,
migrateSessionEntries,
parseSessionEntries,
type SessionEntry,
type SessionMessageEntry,
type ThinkingLevelChangeEntry,
} from "../src/core/session-manager.js";
// ============================================================================
@ -29,7 +32,9 @@ import {
function loadLargeSessionEntries(): SessionEntry[] {
const sessionPath = join(__dirname, "fixtures/large-session.jsonl");
const content = readFileSync(sessionPath, "utf-8");
return parseSessionEntries(content);
const entries = parseSessionEntries(content);
migrateSessionEntries(entries); // Add id/parentId for v1 fixtures
return entries.filter((e): e is SessionEntry => e.type !== "session");
}
function createMockUsage(input: number, output: number, cacheRead = 0, cacheWrite = 0): Usage {
@ -43,7 +48,7 @@ function createMockUsage(input: number, output: number, cacheRead = 0, cacheWrit
};
}
function createUserMessage(text: string): AppMessage {
function createUserMessage(text: string): AgentMessage {
return { role: "user", content: text, timestamp: Date.now() };
}
@ -60,18 +65,72 @@ function createAssistantMessage(text: string, usage?: Usage): AssistantMessage {
};
}
function createMessageEntry(message: AppMessage): SessionMessageEntry {
return { type: "message", timestamp: new Date().toISOString(), message };
let entryCounter = 0;
let lastId: string | null = null;
function resetEntryCounter() {
entryCounter = 0;
lastId = null;
}
function createCompactionEntry(summary: string, firstKeptEntryIndex: number): CompactionEntry {
return {
// Reset counter before each test to get predictable IDs
beforeEach(() => {
resetEntryCounter();
});
function createMessageEntry(message: AgentMessage): SessionMessageEntry {
const id = `test-id-${entryCounter++}`;
const entry: SessionMessageEntry = {
type: "message",
id,
parentId: lastId,
timestamp: new Date().toISOString(),
message,
};
lastId = id;
return entry;
}
function createCompactionEntry(summary: string, firstKeptEntryId: string): CompactionEntry {
const id = `test-id-${entryCounter++}`;
const entry: CompactionEntry = {
type: "compaction",
id,
parentId: lastId,
timestamp: new Date().toISOString(),
summary,
firstKeptEntryIndex,
firstKeptEntryId,
tokensBefore: 10000,
};
lastId = id;
return entry;
}
function createModelChangeEntry(provider: string, modelId: string): ModelChangeEntry {
const id = `test-id-${entryCounter++}`;
const entry: ModelChangeEntry = {
type: "model_change",
id,
parentId: lastId,
timestamp: new Date().toISOString(),
provider,
modelId,
};
lastId = id;
return entry;
}
function createThinkingLevelEntry(thinkingLevel: string): ThinkingLevelChangeEntry {
const id = `test-id-${entryCounter++}`;
const entry: ThinkingLevelChangeEntry = {
type: "thinking_level_change",
id,
parentId: lastId,
timestamp: new Date().toISOString(),
thinkingLevel,
};
lastId = id;
return entry;
}
// ============================================================================
@ -122,9 +181,9 @@ describe("getLastAssistantUsage", () => {
expect(usage!.input).toBe(100);
});
it("should return null if no assistant messages", () => {
it("should return undefined if no assistant messages", () => {
const entries: SessionEntry[] = [createMessageEntry(createUserMessage("Hello"))];
expect(getLastAssistantUsage(entries)).toBeNull();
expect(getLastAssistantUsage(entries)).toBeUndefined();
});
});
@ -213,28 +272,9 @@ describe("findCutPoint", () => {
});
});
describe("createSummaryMessage", () => {
it("should create user message with prefix", () => {
const msg = createSummaryMessage("This is the summary");
expect(msg.role).toBe("user");
if (msg.role === "user") {
expect(msg.content).toContain(
"The conversation history before this point was compacted into the following summary:",
);
expect(msg.content).toContain("This is the summary");
}
});
});
describe("buildSessionContext", () => {
it("should load all messages when no compaction", () => {
const entries: SessionEntry[] = [
{
type: "session",
id: "1",
timestamp: "",
cwd: "",
},
createMessageEntry(createUserMessage("1")),
createMessageEntry(createAssistantMessage("a")),
createMessageEntry(createUserMessage("2")),
@ -248,92 +288,67 @@ describe("buildSessionContext", () => {
});
it("should handle single compaction", () => {
// indices: 0=session, 1=u1, 2=a1, 3=u2, 4=a2, 5=compaction, 6=u3, 7=a3
const entries: SessionEntry[] = [
{
type: "session",
id: "1",
timestamp: "",
cwd: "",
},
createMessageEntry(createUserMessage("1")),
createMessageEntry(createAssistantMessage("a")),
createMessageEntry(createUserMessage("2")),
createMessageEntry(createAssistantMessage("b")),
createCompactionEntry("Summary of 1,a,2,b", 3), // keep from index 3 (u2) onwards
createMessageEntry(createUserMessage("3")),
createMessageEntry(createAssistantMessage("c")),
];
// IDs: u1=test-id-0, a1=test-id-1, u2=test-id-2, a2=test-id-3, compaction=test-id-4, u3=test-id-5, a3=test-id-6
const u1 = createMessageEntry(createUserMessage("1"));
const a1 = createMessageEntry(createAssistantMessage("a"));
const u2 = createMessageEntry(createUserMessage("2"));
const a2 = createMessageEntry(createAssistantMessage("b"));
const compaction = createCompactionEntry("Summary of 1,a,2,b", u2.id); // keep from u2 onwards
const u3 = createMessageEntry(createUserMessage("3"));
const a3 = createMessageEntry(createAssistantMessage("c"));
const entries: SessionEntry[] = [u1, a1, u2, a2, compaction, u3, a3];
const loaded = buildSessionContext(entries);
// summary + kept (u2,a2 from idx 3-4) + after (u3,a3 from idx 6-7) = 5
// summary + kept (u2, a2) + after (u3, a3) = 5
expect(loaded.messages.length).toBe(5);
expect(loaded.messages[0].role).toBe("user");
expect((loaded.messages[0] as any).content).toContain("Summary of 1,a,2,b");
expect(loaded.messages[0].role).toBe("compactionSummary");
expect((loaded.messages[0] as any).summary).toContain("Summary of 1,a,2,b");
});
it("should handle multiple compactions (only latest matters)", () => {
// indices: 0=session, 1=u1, 2=a1, 3=compact1, 4=u2, 5=b, 6=u3, 7=c, 8=compact2, 9=u4, 10=d
const entries: SessionEntry[] = [
{
type: "session",
id: "1",
timestamp: "",
cwd: "",
},
createMessageEntry(createUserMessage("1")),
createMessageEntry(createAssistantMessage("a")),
createCompactionEntry("First summary", 1), // keep from index 1
createMessageEntry(createUserMessage("2")),
createMessageEntry(createAssistantMessage("b")),
createMessageEntry(createUserMessage("3")),
createMessageEntry(createAssistantMessage("c")),
createCompactionEntry("Second summary", 6), // keep from index 6 (u3) onwards
createMessageEntry(createUserMessage("4")),
createMessageEntry(createAssistantMessage("d")),
];
// First batch
const u1 = createMessageEntry(createUserMessage("1"));
const a1 = createMessageEntry(createAssistantMessage("a"));
const compact1 = createCompactionEntry("First summary", u1.id);
// Second batch
const u2 = createMessageEntry(createUserMessage("2"));
const b = createMessageEntry(createAssistantMessage("b"));
const u3 = createMessageEntry(createUserMessage("3"));
const c = createMessageEntry(createAssistantMessage("c"));
const compact2 = createCompactionEntry("Second summary", u3.id); // keep from u3 onwards
// After second compaction
const u4 = createMessageEntry(createUserMessage("4"));
const d = createMessageEntry(createAssistantMessage("d"));
const entries: SessionEntry[] = [u1, a1, compact1, u2, b, u3, c, compact2, u4, d];
const loaded = buildSessionContext(entries);
// summary + kept from idx 6 (u3,c) + after (u4,d) = 5
// summary + kept from u3 (u3, c) + after (u4, d) = 5
expect(loaded.messages.length).toBe(5);
expect((loaded.messages[0] as any).content).toContain("Second summary");
expect((loaded.messages[0] as any).summary).toContain("Second summary");
});
it("should clamp firstKeptEntryIndex to valid range", () => {
// indices: 0=session, 1=u1, 2=a1, 3=compact1, 4=u2, 5=b, 6=compact2
const entries: SessionEntry[] = [
{
type: "session",
id: "1",
timestamp: "",
cwd: "",
},
createMessageEntry(createUserMessage("1")),
createMessageEntry(createAssistantMessage("a")),
createCompactionEntry("First summary", 1),
createMessageEntry(createUserMessage("2")),
createMessageEntry(createAssistantMessage("b")),
createCompactionEntry("Second summary", 0), // index 0 is before compaction1, should still work
];
it("should keep all messages when firstKeptEntryId is first entry", () => {
const u1 = createMessageEntry(createUserMessage("1"));
const a1 = createMessageEntry(createAssistantMessage("a"));
const compact1 = createCompactionEntry("First summary", u1.id); // keep from first entry
const u2 = createMessageEntry(createUserMessage("2"));
const b = createMessageEntry(createAssistantMessage("b"));
const entries: SessionEntry[] = [u1, a1, compact1, u2, b];
const loaded = buildSessionContext(entries);
// Keeps from index 0, but compaction entries are skipped, so u1,a1,u2,b = 4 + summary = 5
// Actually index 0 is session header, so messages are u1,a1,u2,b
expect(loaded.messages.length).toBe(5); // summary + 4 messages
// summary + all messages (u1, a1, u2, b) = 5
expect(loaded.messages.length).toBe(5);
});
it("should track model and thinking level changes", () => {
const entries: SessionEntry[] = [
{
type: "session",
id: "1",
timestamp: "",
cwd: "",
},
createMessageEntry(createUserMessage("1")),
{ type: "model_change", timestamp: "", provider: "openai", modelId: "gpt-4" },
createModelChangeEntry("openai", "gpt-4"),
createMessageEntry(createAssistantMessage("a")),
{ type: "thinking_level_change", timestamp: "", thinkingLevel: "high" },
createThinkingLevelEntry("high"),
];
const loaded = buildSessionContext(entries);
@ -380,27 +395,24 @@ describe("Large session fixture", () => {
// ============================================================================
describe.skipIf(!process.env.ANTHROPIC_OAUTH_TOKEN)("LLM summarization", () => {
it("should generate a compaction event for the large session", async () => {
it("should generate a compaction result for the large session", async () => {
const entries = loadLargeSessionEntries();
const model = getModel("anthropic", "claude-sonnet-4-5")!;
const compactionEvent = await compact(
entries,
model,
DEFAULT_COMPACTION_SETTINGS,
process.env.ANTHROPIC_OAUTH_TOKEN!,
);
const preparation = prepareCompaction(entries, DEFAULT_COMPACTION_SETTINGS);
expect(preparation).toBeDefined();
expect(compactionEvent.type).toBe("compaction");
expect(compactionEvent.summary.length).toBeGreaterThan(100);
expect(compactionEvent.firstKeptEntryIndex).toBeGreaterThan(0);
expect(compactionEvent.tokensBefore).toBeGreaterThan(0);
const compactionResult = await compact(preparation!, model, process.env.ANTHROPIC_OAUTH_TOKEN!);
console.log("Summary length:", compactionEvent.summary.length);
console.log("First kept entry index:", compactionEvent.firstKeptEntryIndex);
console.log("Tokens before:", compactionEvent.tokensBefore);
expect(compactionResult.summary.length).toBeGreaterThan(100);
expect(compactionResult.firstKeptEntryId).toBeTruthy();
expect(compactionResult.tokensBefore).toBeGreaterThan(0);
console.log("Summary length:", compactionResult.summary.length);
console.log("First kept entry ID:", compactionResult.firstKeptEntryId);
console.log("Tokens before:", compactionResult.tokensBefore);
console.log("\n--- SUMMARY ---\n");
console.log(compactionEvent.summary);
console.log(compactionResult.summary);
}, 60000);
it("should produce valid session after compaction", async () => {
@ -408,21 +420,28 @@ describe.skipIf(!process.env.ANTHROPIC_OAUTH_TOKEN)("LLM summarization", () => {
const loaded = buildSessionContext(entries);
const model = getModel("anthropic", "claude-sonnet-4-5")!;
const compactionEvent = await compact(
entries,
model,
DEFAULT_COMPACTION_SETTINGS,
process.env.ANTHROPIC_OAUTH_TOKEN!,
);
const preparation = prepareCompaction(entries, DEFAULT_COMPACTION_SETTINGS);
expect(preparation).toBeDefined();
// Simulate appending compaction to entries
const newEntries = [...entries, compactionEvent];
const compactionResult = await compact(preparation!, model, process.env.ANTHROPIC_OAUTH_TOKEN!);
// Simulate appending compaction to entries by creating a proper entry
const lastEntry = entries[entries.length - 1];
const parentId = lastEntry.id;
const compactionEntry: CompactionEntry = {
type: "compaction",
id: "compaction-test-id",
parentId,
timestamp: new Date().toISOString(),
...compactionResult,
};
const newEntries = [...entries, compactionEntry];
const reloaded = buildSessionContext(newEntries);
// Should have summary + kept messages
expect(reloaded.messages.length).toBeLessThan(loaded.messages.length);
expect(reloaded.messages[0].role).toBe("user");
expect((reloaded.messages[0] as any).content).toContain(compactionEvent.summary);
expect(reloaded.messages[0].role).toBe("compactionSummary");
expect((reloaded.messages[0] as any).summary).toContain(compactionResult.summary);
console.log("Original messages:", loaded.messages.length);
console.log("After compaction:", reloaded.messages.length);

View file

@ -66,21 +66,21 @@ describe("parseModelPattern", () => {
const result = parseModelPattern("claude-sonnet-4-5", allModels);
expect(result.model?.id).toBe("claude-sonnet-4-5");
expect(result.thinkingLevel).toBe("off");
expect(result.warning).toBeNull();
expect(result.warning).toBeUndefined();
});
test("partial match returns best model", () => {
const result = parseModelPattern("sonnet", allModels);
expect(result.model?.id).toBe("claude-sonnet-4-5");
expect(result.thinkingLevel).toBe("off");
expect(result.warning).toBeNull();
expect(result.warning).toBeUndefined();
});
test("no match returns null model", () => {
const result = parseModelPattern("nonexistent", allModels);
expect(result.model).toBeNull();
expect(result.model).toBeUndefined();
expect(result.thinkingLevel).toBe("off");
expect(result.warning).toBeNull();
expect(result.warning).toBeUndefined();
});
});
@ -89,14 +89,14 @@ describe("parseModelPattern", () => {
const result = parseModelPattern("sonnet:high", allModels);
expect(result.model?.id).toBe("claude-sonnet-4-5");
expect(result.thinkingLevel).toBe("high");
expect(result.warning).toBeNull();
expect(result.warning).toBeUndefined();
});
test("gpt-4o:medium returns gpt-4o with medium thinking level", () => {
const result = parseModelPattern("gpt-4o:medium", allModels);
expect(result.model?.id).toBe("gpt-4o");
expect(result.thinkingLevel).toBe("medium");
expect(result.warning).toBeNull();
expect(result.warning).toBeUndefined();
});
test("all valid thinking levels work", () => {
@ -104,7 +104,7 @@ describe("parseModelPattern", () => {
const result = parseModelPattern(`sonnet:${level}`, allModels);
expect(result.model?.id).toBe("claude-sonnet-4-5");
expect(result.thinkingLevel).toBe(level);
expect(result.warning).toBeNull();
expect(result.warning).toBeUndefined();
}
});
});
@ -131,7 +131,7 @@ describe("parseModelPattern", () => {
const result = parseModelPattern("qwen/qwen3-coder:exacto", allModels);
expect(result.model?.id).toBe("qwen/qwen3-coder:exacto");
expect(result.thinkingLevel).toBe("off");
expect(result.warning).toBeNull();
expect(result.warning).toBeUndefined();
});
test("openrouter/qwen/qwen3-coder:exacto matches with provider prefix", () => {
@ -139,14 +139,14 @@ describe("parseModelPattern", () => {
expect(result.model?.id).toBe("qwen/qwen3-coder:exacto");
expect(result.model?.provider).toBe("openrouter");
expect(result.thinkingLevel).toBe("off");
expect(result.warning).toBeNull();
expect(result.warning).toBeUndefined();
});
test("qwen3-coder:exacto:high matches model with high thinking level", () => {
const result = parseModelPattern("qwen/qwen3-coder:exacto:high", allModels);
expect(result.model?.id).toBe("qwen/qwen3-coder:exacto");
expect(result.thinkingLevel).toBe("high");
expect(result.warning).toBeNull();
expect(result.warning).toBeUndefined();
});
test("openrouter/qwen/qwen3-coder:exacto:high matches with provider and thinking level", () => {
@ -154,14 +154,14 @@ describe("parseModelPattern", () => {
expect(result.model?.id).toBe("qwen/qwen3-coder:exacto");
expect(result.model?.provider).toBe("openrouter");
expect(result.thinkingLevel).toBe("high");
expect(result.warning).toBeNull();
expect(result.warning).toBeUndefined();
});
test("gpt-4o:extended matches the extended model", () => {
const result = parseModelPattern("openai/gpt-4o:extended", allModels);
expect(result.model?.id).toBe("openai/gpt-4o:extended");
expect(result.thinkingLevel).toBe("off");
expect(result.warning).toBeNull();
expect(result.warning).toBeUndefined();
});
});

View file

@ -273,7 +273,7 @@ describe.skipIf(!process.env.ANTHROPIC_API_KEY && !process.env.ANTHROPIC_OAUTH_T
// Initially null
let text = await client.getLastAssistantText();
expect(text).toBeNull();
expect(text).toBeUndefined();
// Send prompt
await client.promptAndWait("Reply with just: test123");

View file

@ -0,0 +1,268 @@
import { describe, expect, it } from "vitest";
import {
type BranchSummaryEntry,
buildSessionContext,
type CompactionEntry,
type ModelChangeEntry,
type SessionEntry,
type SessionMessageEntry,
type ThinkingLevelChangeEntry,
} from "../../src/core/session-manager.js";
function msg(id: string, parentId: string | null, role: "user" | "assistant", text: string): SessionMessageEntry {
const base = { type: "message" as const, id, parentId, timestamp: "2025-01-01T00:00:00Z" };
if (role === "user") {
return { ...base, message: { role, content: text, timestamp: 1 } };
}
return {
...base,
message: {
role,
content: [{ type: "text", text }],
api: "anthropic-messages",
provider: "anthropic",
model: "claude-test",
usage: {
input: 1,
output: 1,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 2,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
},
stopReason: "stop",
timestamp: 1,
},
};
}
function compaction(id: string, parentId: string | null, summary: string, firstKeptEntryId: string): CompactionEntry {
return {
type: "compaction",
id,
parentId,
timestamp: "2025-01-01T00:00:00Z",
summary,
firstKeptEntryId,
tokensBefore: 1000,
};
}
function branchSummary(id: string, parentId: string | null, summary: string, fromId: string): BranchSummaryEntry {
return { type: "branch_summary", id, parentId, timestamp: "2025-01-01T00:00:00Z", summary, fromId };
}
function thinkingLevel(id: string, parentId: string | null, level: string): ThinkingLevelChangeEntry {
return { type: "thinking_level_change", id, parentId, timestamp: "2025-01-01T00:00:00Z", thinkingLevel: level };
}
function modelChange(id: string, parentId: string | null, provider: string, modelId: string): ModelChangeEntry {
return { type: "model_change", id, parentId, timestamp: "2025-01-01T00:00:00Z", provider, modelId };
}
describe("buildSessionContext", () => {
describe("trivial cases", () => {
it("empty entries returns empty context", () => {
const ctx = buildSessionContext([]);
expect(ctx.messages).toEqual([]);
expect(ctx.thinkingLevel).toBe("off");
expect(ctx.model).toBeNull();
});
it("single user message", () => {
const entries: SessionEntry[] = [msg("1", null, "user", "hello")];
const ctx = buildSessionContext(entries);
expect(ctx.messages).toHaveLength(1);
expect(ctx.messages[0].role).toBe("user");
});
it("simple conversation", () => {
const entries: SessionEntry[] = [
msg("1", null, "user", "hello"),
msg("2", "1", "assistant", "hi there"),
msg("3", "2", "user", "how are you"),
msg("4", "3", "assistant", "great"),
];
const ctx = buildSessionContext(entries);
expect(ctx.messages).toHaveLength(4);
expect(ctx.messages.map((m) => m.role)).toEqual(["user", "assistant", "user", "assistant"]);
});
it("tracks thinking level changes", () => {
const entries: SessionEntry[] = [
msg("1", null, "user", "hello"),
thinkingLevel("2", "1", "high"),
msg("3", "2", "assistant", "thinking hard"),
];
const ctx = buildSessionContext(entries);
expect(ctx.thinkingLevel).toBe("high");
expect(ctx.messages).toHaveLength(2);
});
it("tracks model from assistant message", () => {
const entries: SessionEntry[] = [msg("1", null, "user", "hello"), msg("2", "1", "assistant", "hi")];
const ctx = buildSessionContext(entries);
expect(ctx.model).toEqual({ provider: "anthropic", modelId: "claude-test" });
});
it("tracks model from model change entry", () => {
const entries: SessionEntry[] = [
msg("1", null, "user", "hello"),
modelChange("2", "1", "openai", "gpt-4"),
msg("3", "2", "assistant", "hi"),
];
const ctx = buildSessionContext(entries);
// Assistant message overwrites model change
expect(ctx.model).toEqual({ provider: "anthropic", modelId: "claude-test" });
});
});
describe("with compaction", () => {
it("includes summary before kept messages", () => {
const entries: SessionEntry[] = [
msg("1", null, "user", "first"),
msg("2", "1", "assistant", "response1"),
msg("3", "2", "user", "second"),
msg("4", "3", "assistant", "response2"),
compaction("5", "4", "Summary of first two turns", "3"),
msg("6", "5", "user", "third"),
msg("7", "6", "assistant", "response3"),
];
const ctx = buildSessionContext(entries);
// Should have: summary + kept (3,4) + after (6,7) = 5 messages
expect(ctx.messages).toHaveLength(5);
expect((ctx.messages[0] as any).summary).toContain("Summary of first two turns");
expect((ctx.messages[1] as any).content).toBe("second");
expect((ctx.messages[2] as any).content[0].text).toBe("response2");
expect((ctx.messages[3] as any).content).toBe("third");
expect((ctx.messages[4] as any).content[0].text).toBe("response3");
});
it("handles compaction keeping from first message", () => {
const entries: SessionEntry[] = [
msg("1", null, "user", "first"),
msg("2", "1", "assistant", "response"),
compaction("3", "2", "Empty summary", "1"),
msg("4", "3", "user", "second"),
];
const ctx = buildSessionContext(entries);
// Summary + all messages (1,2,4)
expect(ctx.messages).toHaveLength(4);
expect((ctx.messages[0] as any).summary).toContain("Empty summary");
});
it("multiple compactions uses latest", () => {
const entries: SessionEntry[] = [
msg("1", null, "user", "a"),
msg("2", "1", "assistant", "b"),
compaction("3", "2", "First summary", "1"),
msg("4", "3", "user", "c"),
msg("5", "4", "assistant", "d"),
compaction("6", "5", "Second summary", "4"),
msg("7", "6", "user", "e"),
];
const ctx = buildSessionContext(entries);
// Should use second summary, keep from 4
expect(ctx.messages).toHaveLength(4);
expect((ctx.messages[0] as any).summary).toContain("Second summary");
});
});
describe("with branches", () => {
it("follows path to specified leaf", () => {
// Tree:
// 1 -> 2 -> 3 (branch A)
// \-> 4 (branch B)
const entries: SessionEntry[] = [
msg("1", null, "user", "start"),
msg("2", "1", "assistant", "response"),
msg("3", "2", "user", "branch A"),
msg("4", "2", "user", "branch B"),
];
const ctxA = buildSessionContext(entries, "3");
expect(ctxA.messages).toHaveLength(3);
expect((ctxA.messages[2] as any).content).toBe("branch A");
const ctxB = buildSessionContext(entries, "4");
expect(ctxB.messages).toHaveLength(3);
expect((ctxB.messages[2] as any).content).toBe("branch B");
});
it("includes branch summary in path", () => {
const entries: SessionEntry[] = [
msg("1", null, "user", "start"),
msg("2", "1", "assistant", "response"),
msg("3", "2", "user", "abandoned path"),
branchSummary("4", "2", "Summary of abandoned work", "3"),
msg("5", "4", "user", "new direction"),
];
const ctx = buildSessionContext(entries, "5");
expect(ctx.messages).toHaveLength(4);
expect((ctx.messages[2] as any).summary).toContain("Summary of abandoned work");
expect((ctx.messages[3] as any).content).toBe("new direction");
});
it("complex tree with multiple branches and compaction", () => {
// Tree:
// 1 -> 2 -> 3 -> 4 -> compaction(5) -> 6 -> 7 (main path)
// \-> 8 -> 9 (abandoned branch)
// \-> branchSummary(10) -> 11 (resumed from 3)
const entries: SessionEntry[] = [
msg("1", null, "user", "start"),
msg("2", "1", "assistant", "r1"),
msg("3", "2", "user", "q2"),
msg("4", "3", "assistant", "r2"),
compaction("5", "4", "Compacted history", "3"),
msg("6", "5", "user", "q3"),
msg("7", "6", "assistant", "r3"),
// Abandoned branch from 3
msg("8", "3", "user", "wrong path"),
msg("9", "8", "assistant", "wrong response"),
// Branch summary resuming from 3
branchSummary("10", "3", "Tried wrong approach", "9"),
msg("11", "10", "user", "better approach"),
];
// Main path to 7: summary + kept(3,4) + after(6,7)
const ctxMain = buildSessionContext(entries, "7");
expect(ctxMain.messages).toHaveLength(5);
expect((ctxMain.messages[0] as any).summary).toContain("Compacted history");
expect((ctxMain.messages[1] as any).content).toBe("q2");
expect((ctxMain.messages[2] as any).content[0].text).toBe("r2");
expect((ctxMain.messages[3] as any).content).toBe("q3");
expect((ctxMain.messages[4] as any).content[0].text).toBe("r3");
// Branch path to 11: 1,2,3 + branch_summary + 11
const ctxBranch = buildSessionContext(entries, "11");
expect(ctxBranch.messages).toHaveLength(5);
expect((ctxBranch.messages[0] as any).content).toBe("start");
expect((ctxBranch.messages[1] as any).content[0].text).toBe("r1");
expect((ctxBranch.messages[2] as any).content).toBe("q2");
expect((ctxBranch.messages[3] as any).summary).toContain("Tried wrong approach");
expect((ctxBranch.messages[4] as any).content).toBe("better approach");
});
});
describe("edge cases", () => {
it("uses last entry when leafId not found", () => {
const entries: SessionEntry[] = [msg("1", null, "user", "hello"), msg("2", "1", "assistant", "hi")];
const ctx = buildSessionContext(entries, "nonexistent");
expect(ctx.messages).toHaveLength(2);
});
it("handles orphaned entries gracefully", () => {
const entries: SessionEntry[] = [
msg("1", null, "user", "hello"),
msg("2", "missing", "assistant", "orphan"), // parent doesn't exist
];
const ctx = buildSessionContext(entries, "2");
// Should only get the orphan since parent chain is broken
expect(ctx.messages).toHaveLength(1);
});
});
});

View file

@ -0,0 +1,127 @@
import { mkdirSync, rmSync, writeFileSync } from "fs";
import { tmpdir } from "os";
import { join } from "path";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { findMostRecentSession, loadEntriesFromFile } from "../../src/core/session-manager.js";
describe("loadEntriesFromFile", () => {
let tempDir: string;
beforeEach(() => {
tempDir = join(tmpdir(), `session-test-${Date.now()}`);
mkdirSync(tempDir, { recursive: true });
});
afterEach(() => {
rmSync(tempDir, { recursive: true, force: true });
});
it("returns empty array for non-existent file", () => {
const entries = loadEntriesFromFile(join(tempDir, "nonexistent.jsonl"));
expect(entries).toEqual([]);
});
it("returns empty array for empty file", () => {
const file = join(tempDir, "empty.jsonl");
writeFileSync(file, "");
expect(loadEntriesFromFile(file)).toEqual([]);
});
it("returns empty array for file without valid session header", () => {
const file = join(tempDir, "no-header.jsonl");
writeFileSync(file, '{"type":"message","id":"1"}\n');
expect(loadEntriesFromFile(file)).toEqual([]);
});
it("returns empty array for malformed JSON", () => {
const file = join(tempDir, "malformed.jsonl");
writeFileSync(file, "not json\n");
expect(loadEntriesFromFile(file)).toEqual([]);
});
it("loads valid session file", () => {
const file = join(tempDir, "valid.jsonl");
writeFileSync(
file,
'{"type":"session","id":"abc","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n' +
'{"type":"message","id":"1","parentId":null,"timestamp":"2025-01-01T00:00:01Z","message":{"role":"user","content":"hi","timestamp":1}}\n',
);
const entries = loadEntriesFromFile(file);
expect(entries).toHaveLength(2);
expect(entries[0].type).toBe("session");
expect(entries[1].type).toBe("message");
});
it("skips malformed lines but keeps valid ones", () => {
const file = join(tempDir, "mixed.jsonl");
writeFileSync(
file,
'{"type":"session","id":"abc","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n' +
"not valid json\n" +
'{"type":"message","id":"1","parentId":null,"timestamp":"2025-01-01T00:00:01Z","message":{"role":"user","content":"hi","timestamp":1}}\n',
);
const entries = loadEntriesFromFile(file);
expect(entries).toHaveLength(2);
});
});
describe("findMostRecentSession", () => {
let tempDir: string;
beforeEach(() => {
tempDir = join(tmpdir(), `session-test-${Date.now()}`);
mkdirSync(tempDir, { recursive: true });
});
afterEach(() => {
rmSync(tempDir, { recursive: true, force: true });
});
it("returns null for empty directory", () => {
expect(findMostRecentSession(tempDir)).toBeNull();
});
it("returns null for non-existent directory", () => {
expect(findMostRecentSession(join(tempDir, "nonexistent"))).toBeNull();
});
it("ignores non-jsonl files", () => {
writeFileSync(join(tempDir, "file.txt"), "hello");
writeFileSync(join(tempDir, "file.json"), "{}");
expect(findMostRecentSession(tempDir)).toBeNull();
});
it("ignores jsonl files without valid session header", () => {
writeFileSync(join(tempDir, "invalid.jsonl"), '{"type":"message"}\n');
expect(findMostRecentSession(tempDir)).toBeNull();
});
it("returns single valid session file", () => {
const file = join(tempDir, "session.jsonl");
writeFileSync(file, '{"type":"session","id":"abc","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n');
expect(findMostRecentSession(tempDir)).toBe(file);
});
it("returns most recently modified session", async () => {
const file1 = join(tempDir, "older.jsonl");
const file2 = join(tempDir, "newer.jsonl");
writeFileSync(file1, '{"type":"session","id":"old","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n');
// Small delay to ensure different mtime
await new Promise((r) => setTimeout(r, 10));
writeFileSync(file2, '{"type":"session","id":"new","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n');
expect(findMostRecentSession(tempDir)).toBe(file2);
});
it("skips invalid files and returns valid one", async () => {
const invalid = join(tempDir, "invalid.jsonl");
const valid = join(tempDir, "valid.jsonl");
writeFileSync(invalid, '{"type":"not-session"}\n');
await new Promise((r) => setTimeout(r, 10));
writeFileSync(valid, '{"type":"session","id":"abc","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n');
expect(findMostRecentSession(tempDir)).toBe(valid);
});
});

View file

@ -0,0 +1,178 @@
import { describe, expect, it } from "vitest";
import { type LabelEntry, SessionManager } from "../../src/core/session-manager.js";
describe("SessionManager labels", () => {
it("sets and gets labels", () => {
const session = SessionManager.inMemory();
const msgId = session.appendMessage({ role: "user", content: "hello", timestamp: 1 });
// No label initially
expect(session.getLabel(msgId)).toBeUndefined();
// Set a label
const labelId = session.appendLabelChange(msgId, "checkpoint");
expect(session.getLabel(msgId)).toBe("checkpoint");
// Label entry should be in entries
const entries = session.getEntries();
const labelEntry = entries.find((e) => e.type === "label") as LabelEntry;
expect(labelEntry).toBeDefined();
expect(labelEntry.id).toBe(labelId);
expect(labelEntry.targetId).toBe(msgId);
expect(labelEntry.label).toBe("checkpoint");
});
it("clears labels with undefined", () => {
const session = SessionManager.inMemory();
const msgId = session.appendMessage({ role: "user", content: "hello", timestamp: 1 });
session.appendLabelChange(msgId, "checkpoint");
expect(session.getLabel(msgId)).toBe("checkpoint");
// Clear the label
session.appendLabelChange(msgId, undefined);
expect(session.getLabel(msgId)).toBeUndefined();
});
it("last label wins", () => {
const session = SessionManager.inMemory();
const msgId = session.appendMessage({ role: "user", content: "hello", timestamp: 1 });
session.appendLabelChange(msgId, "first");
session.appendLabelChange(msgId, "second");
session.appendLabelChange(msgId, "third");
expect(session.getLabel(msgId)).toBe("third");
});
it("labels are included in tree nodes", () => {
const session = SessionManager.inMemory();
const msg1Id = session.appendMessage({ role: "user", content: "hello", timestamp: 1 });
const msg2Id = session.appendMessage({
role: "assistant",
content: [{ type: "text", text: "hi" }],
api: "anthropic-messages",
provider: "anthropic",
model: "test",
usage: {
input: 1,
output: 1,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 2,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
},
stopReason: "stop",
timestamp: 2,
});
session.appendLabelChange(msg1Id, "start");
session.appendLabelChange(msg2Id, "response");
const tree = session.getTree();
// Find the message nodes (skip label entries)
const msg1Node = tree.find((n) => n.entry.id === msg1Id);
expect(msg1Node?.label).toBe("start");
// msg2 is a child of msg1
const msg2Node = msg1Node?.children.find((n) => n.entry.id === msg2Id);
expect(msg2Node?.label).toBe("response");
});
it("labels are preserved in createBranchedSession", () => {
const session = SessionManager.inMemory();
const msg1Id = session.appendMessage({ role: "user", content: "hello", timestamp: 1 });
const msg2Id = session.appendMessage({
role: "assistant",
content: [{ type: "text", text: "hi" }],
api: "anthropic-messages",
provider: "anthropic",
model: "test",
usage: {
input: 1,
output: 1,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 2,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
},
stopReason: "stop",
timestamp: 2,
});
session.appendLabelChange(msg1Id, "important");
session.appendLabelChange(msg2Id, "also-important");
// Branch from msg2 (in-memory mode returns null, but updates internal state)
session.createBranchedSession(msg2Id);
// Labels should be preserved
expect(session.getLabel(msg1Id)).toBe("important");
expect(session.getLabel(msg2Id)).toBe("also-important");
// New label entries should exist
const entries = session.getEntries();
const labelEntries = entries.filter((e) => e.type === "label") as LabelEntry[];
expect(labelEntries).toHaveLength(2);
});
it("labels not on path are not preserved in createBranchedSession", () => {
const session = SessionManager.inMemory();
const msg1Id = session.appendMessage({ role: "user", content: "hello", timestamp: 1 });
const msg2Id = session.appendMessage({
role: "assistant",
content: [{ type: "text", text: "hi" }],
api: "anthropic-messages",
provider: "anthropic",
model: "test",
usage: {
input: 1,
output: 1,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 2,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
},
stopReason: "stop",
timestamp: 2,
});
const msg3Id = session.appendMessage({ role: "user", content: "followup", timestamp: 3 });
// Label all messages
session.appendLabelChange(msg1Id, "first");
session.appendLabelChange(msg2Id, "second");
session.appendLabelChange(msg3Id, "third");
// Branch from msg2 (excludes msg3)
session.createBranchedSession(msg2Id);
// Only labels for msg1 and msg2 should be preserved
expect(session.getLabel(msg1Id)).toBe("first");
expect(session.getLabel(msg2Id)).toBe("second");
expect(session.getLabel(msg3Id)).toBeUndefined();
});
it("labels are not included in buildSessionContext", () => {
const session = SessionManager.inMemory();
const msgId = session.appendMessage({ role: "user", content: "hello", timestamp: 1 });
session.appendLabelChange(msgId, "checkpoint");
const ctx = session.buildSessionContext();
expect(ctx.messages).toHaveLength(1);
expect(ctx.messages[0].role).toBe("user");
});
it("throws when labeling non-existent entry", () => {
const session = SessionManager.inMemory();
expect(() => session.appendLabelChange("non-existent", "label")).toThrow("Entry non-existent not found");
});
});

View file

@ -0,0 +1,78 @@
import { describe, expect, it } from "vitest";
import { type FileEntry, migrateSessionEntries } from "../../src/core/session-manager.js";
describe("migrateSessionEntries", () => {
it("should add id/parentId to v1 entries", () => {
const entries: FileEntry[] = [
{ type: "session", id: "sess-1", timestamp: "2025-01-01T00:00:00Z", cwd: "/tmp" },
{ type: "message", timestamp: "2025-01-01T00:00:01Z", message: { role: "user", content: "hi", timestamp: 1 } },
{
type: "message",
timestamp: "2025-01-01T00:00:02Z",
message: {
role: "assistant",
content: [{ type: "text", text: "hello" }],
api: "test",
provider: "test",
model: "test",
usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0 },
stopReason: "stop",
timestamp: 2,
},
},
] as FileEntry[];
migrateSessionEntries(entries);
// Header should have version set
expect((entries[0] as any).version).toBe(2);
// Entries should have id/parentId
const msg1 = entries[1] as any;
const msg2 = entries[2] as any;
expect(msg1.id).toBeDefined();
expect(msg1.id.length).toBe(8);
expect(msg1.parentId).toBeNull();
expect(msg2.id).toBeDefined();
expect(msg2.id.length).toBe(8);
expect(msg2.parentId).toBe(msg1.id);
});
it("should be idempotent (skip already migrated)", () => {
const entries: FileEntry[] = [
{ type: "session", id: "sess-1", version: 2, timestamp: "2025-01-01T00:00:00Z", cwd: "/tmp" },
{
type: "message",
id: "abc12345",
parentId: null,
timestamp: "2025-01-01T00:00:01Z",
message: { role: "user", content: "hi", timestamp: 1 },
},
{
type: "message",
id: "def67890",
parentId: "abc12345",
timestamp: "2025-01-01T00:00:02Z",
message: {
role: "assistant",
content: [{ type: "text", text: "hello" }],
api: "test",
provider: "test",
model: "test",
usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0 },
stopReason: "stop",
timestamp: 2,
},
},
] as FileEntry[];
migrateSessionEntries(entries);
// IDs should be unchanged
expect((entries[1] as any).id).toBe("abc12345");
expect((entries[2] as any).id).toBe("def67890");
expect((entries[2] as any).parentId).toBe("abc12345");
});
});

View file

@ -0,0 +1,55 @@
import { describe, expect, it } from "vitest";
import { type CustomEntry, SessionManager } from "../../src/core/session-manager.js";
describe("SessionManager.saveCustomEntry", () => {
it("saves custom entries and includes them in tree traversal", () => {
const session = SessionManager.inMemory();
// Save a message
const msgId = session.appendMessage({ role: "user", content: "hello", timestamp: 1 });
// Save a custom entry
const customId = session.appendCustomEntry("my_hook", { foo: "bar" });
// Save another message
const msg2Id = session.appendMessage({
role: "assistant",
content: [{ type: "text", text: "hi" }],
api: "anthropic-messages",
provider: "anthropic",
model: "test",
usage: {
input: 1,
output: 1,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 2,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
},
stopReason: "stop",
timestamp: 2,
});
// Custom entry should be in entries
const entries = session.getEntries();
expect(entries).toHaveLength(3);
const customEntry = entries.find((e) => e.type === "custom") as CustomEntry;
expect(customEntry).toBeDefined();
expect(customEntry.customType).toBe("my_hook");
expect(customEntry.data).toEqual({ foo: "bar" });
expect(customEntry.id).toBe(customId);
expect(customEntry.parentId).toBe(msgId);
// Tree structure should be correct
const path = session.getBranch();
expect(path).toHaveLength(3);
expect(path[0].id).toBe(msgId);
expect(path[1].id).toBe(customId);
expect(path[2].id).toBe(msg2Id);
// buildSessionContext should work (custom entries skipped in messages)
const ctx = session.buildSessionContext();
expect(ctx.messages).toHaveLength(2); // only message entries
});
});

View file

@ -0,0 +1,460 @@
import { describe, expect, it } from "vitest";
import { type CustomEntry, SessionManager } from "../../src/core/session-manager.js";
import { assistantMsg, userMsg } from "../utilities.js";
describe("SessionManager append and tree traversal", () => {
describe("append operations", () => {
it("appendMessage creates entry with correct parentId chain", () => {
const session = SessionManager.inMemory();
const id1 = session.appendMessage(userMsg("first"));
const id2 = session.appendMessage(assistantMsg("second"));
const id3 = session.appendMessage(userMsg("third"));
const entries = session.getEntries();
expect(entries).toHaveLength(3);
expect(entries[0].id).toBe(id1);
expect(entries[0].parentId).toBeNull();
expect(entries[0].type).toBe("message");
expect(entries[1].id).toBe(id2);
expect(entries[1].parentId).toBe(id1);
expect(entries[2].id).toBe(id3);
expect(entries[2].parentId).toBe(id2);
});
it("appendThinkingLevelChange integrates into tree", () => {
const session = SessionManager.inMemory();
const msgId = session.appendMessage(userMsg("hello"));
const thinkingId = session.appendThinkingLevelChange("high");
const _msg2Id = session.appendMessage(assistantMsg("response"));
const entries = session.getEntries();
expect(entries).toHaveLength(3);
const thinkingEntry = entries.find((e) => e.type === "thinking_level_change");
expect(thinkingEntry).toBeDefined();
expect(thinkingEntry!.id).toBe(thinkingId);
expect(thinkingEntry!.parentId).toBe(msgId);
expect(entries[2].parentId).toBe(thinkingId);
});
it("appendModelChange integrates into tree", () => {
const session = SessionManager.inMemory();
const msgId = session.appendMessage(userMsg("hello"));
const modelId = session.appendModelChange("openai", "gpt-4");
const _msg2Id = session.appendMessage(assistantMsg("response"));
const entries = session.getEntries();
const modelEntry = entries.find((e) => e.type === "model_change");
expect(modelEntry).toBeDefined();
expect(modelEntry?.id).toBe(modelId);
expect(modelEntry?.parentId).toBe(msgId);
if (modelEntry?.type === "model_change") {
expect(modelEntry.provider).toBe("openai");
expect(modelEntry.modelId).toBe("gpt-4");
}
expect(entries[2].parentId).toBe(modelId);
});
it("appendCompaction integrates into tree", () => {
const session = SessionManager.inMemory();
const id1 = session.appendMessage(userMsg("1"));
const id2 = session.appendMessage(assistantMsg("2"));
const compactionId = session.appendCompaction("summary", id1, 1000);
const _id3 = session.appendMessage(userMsg("3"));
const entries = session.getEntries();
const compactionEntry = entries.find((e) => e.type === "compaction");
expect(compactionEntry).toBeDefined();
expect(compactionEntry?.id).toBe(compactionId);
expect(compactionEntry?.parentId).toBe(id2);
if (compactionEntry?.type === "compaction") {
expect(compactionEntry.summary).toBe("summary");
expect(compactionEntry.firstKeptEntryId).toBe(id1);
expect(compactionEntry.tokensBefore).toBe(1000);
}
expect(entries[3].parentId).toBe(compactionId);
});
it("appendCustomEntry integrates into tree", () => {
const session = SessionManager.inMemory();
const msgId = session.appendMessage(userMsg("hello"));
const customId = session.appendCustomEntry("my_hook", { key: "value" });
const _msg2Id = session.appendMessage(assistantMsg("response"));
const entries = session.getEntries();
const customEntry = entries.find((e) => e.type === "custom") as CustomEntry;
expect(customEntry).toBeDefined();
expect(customEntry.id).toBe(customId);
expect(customEntry.parentId).toBe(msgId);
expect(customEntry.customType).toBe("my_hook");
expect(customEntry.data).toEqual({ key: "value" });
expect(entries[2].parentId).toBe(customId);
});
it("leaf pointer advances after each append", () => {
const session = SessionManager.inMemory();
expect(session.getLeafId()).toBeNull();
const id1 = session.appendMessage(userMsg("1"));
expect(session.getLeafId()).toBe(id1);
const id2 = session.appendMessage(assistantMsg("2"));
expect(session.getLeafId()).toBe(id2);
const id3 = session.appendThinkingLevelChange("high");
expect(session.getLeafId()).toBe(id3);
});
});
describe("getPath", () => {
it("returns empty array for empty session", () => {
const session = SessionManager.inMemory();
expect(session.getBranch()).toEqual([]);
});
it("returns single entry path", () => {
const session = SessionManager.inMemory();
const id = session.appendMessage(userMsg("hello"));
const path = session.getBranch();
expect(path).toHaveLength(1);
expect(path[0].id).toBe(id);
});
it("returns full path from root to leaf", () => {
const session = SessionManager.inMemory();
const id1 = session.appendMessage(userMsg("1"));
const id2 = session.appendMessage(assistantMsg("2"));
const id3 = session.appendThinkingLevelChange("high");
const id4 = session.appendMessage(userMsg("3"));
const path = session.getBranch();
expect(path).toHaveLength(4);
expect(path.map((e) => e.id)).toEqual([id1, id2, id3, id4]);
});
it("returns path from specified entry to root", () => {
const session = SessionManager.inMemory();
const id1 = session.appendMessage(userMsg("1"));
const id2 = session.appendMessage(assistantMsg("2"));
const _id3 = session.appendMessage(userMsg("3"));
const _id4 = session.appendMessage(assistantMsg("4"));
const path = session.getBranch(id2);
expect(path).toHaveLength(2);
expect(path.map((e) => e.id)).toEqual([id1, id2]);
});
});
describe("getTree", () => {
it("returns empty array for empty session", () => {
const session = SessionManager.inMemory();
expect(session.getTree()).toEqual([]);
});
it("returns single root for linear session", () => {
const session = SessionManager.inMemory();
const id1 = session.appendMessage(userMsg("1"));
const id2 = session.appendMessage(assistantMsg("2"));
const id3 = session.appendMessage(userMsg("3"));
const tree = session.getTree();
expect(tree).toHaveLength(1);
const root = tree[0];
expect(root.entry.id).toBe(id1);
expect(root.children).toHaveLength(1);
expect(root.children[0].entry.id).toBe(id2);
expect(root.children[0].children).toHaveLength(1);
expect(root.children[0].children[0].entry.id).toBe(id3);
expect(root.children[0].children[0].children).toHaveLength(0);
});
it("returns tree with branches after branch", () => {
const session = SessionManager.inMemory();
// Build: 1 -> 2 -> 3
const id1 = session.appendMessage(userMsg("1"));
const id2 = session.appendMessage(assistantMsg("2"));
const id3 = session.appendMessage(userMsg("3"));
// Branch from id2, add new path: 2 -> 4
session.branch(id2);
const id4 = session.appendMessage(userMsg("4-branch"));
const tree = session.getTree();
expect(tree).toHaveLength(1);
const root = tree[0];
expect(root.entry.id).toBe(id1);
expect(root.children).toHaveLength(1);
const node2 = root.children[0];
expect(node2.entry.id).toBe(id2);
expect(node2.children).toHaveLength(2); // id3 and id4 are siblings
const childIds = node2.children.map((c) => c.entry.id).sort();
expect(childIds).toEqual([id3, id4].sort());
});
it("handles multiple branches at same point", () => {
const session = SessionManager.inMemory();
const _id1 = session.appendMessage(userMsg("root"));
const id2 = session.appendMessage(assistantMsg("response"));
// Branch A
session.branch(id2);
const idA = session.appendMessage(userMsg("branch-A"));
// Branch B
session.branch(id2);
const idB = session.appendMessage(userMsg("branch-B"));
// Branch C
session.branch(id2);
const idC = session.appendMessage(userMsg("branch-C"));
const tree = session.getTree();
const node2 = tree[0].children[0];
expect(node2.entry.id).toBe(id2);
expect(node2.children).toHaveLength(3);
const branchIds = node2.children.map((c) => c.entry.id).sort();
expect(branchIds).toEqual([idA, idB, idC].sort());
});
it("handles deep branching", () => {
const session = SessionManager.inMemory();
// Main path: 1 -> 2 -> 3 -> 4
const _id1 = session.appendMessage(userMsg("1"));
const id2 = session.appendMessage(assistantMsg("2"));
const id3 = session.appendMessage(userMsg("3"));
const _id4 = session.appendMessage(assistantMsg("4"));
// Branch from 2: 2 -> 5 -> 6
session.branch(id2);
const id5 = session.appendMessage(userMsg("5"));
const _id6 = session.appendMessage(assistantMsg("6"));
// Branch from 5: 5 -> 7
session.branch(id5);
const _id7 = session.appendMessage(userMsg("7"));
const tree = session.getTree();
// Verify structure
const node2 = tree[0].children[0];
expect(node2.children).toHaveLength(2); // id3 and id5
const node5 = node2.children.find((c) => c.entry.id === id5)!;
expect(node5.children).toHaveLength(2); // id6 and id7
const node3 = node2.children.find((c) => c.entry.id === id3)!;
expect(node3.children).toHaveLength(1); // id4
});
});
describe("branch", () => {
it("moves leaf pointer to specified entry", () => {
const session = SessionManager.inMemory();
const id1 = session.appendMessage(userMsg("1"));
const _id2 = session.appendMessage(assistantMsg("2"));
const id3 = session.appendMessage(userMsg("3"));
expect(session.getLeafId()).toBe(id3);
session.branch(id1);
expect(session.getLeafId()).toBe(id1);
});
it("throws for non-existent entry", () => {
const session = SessionManager.inMemory();
session.appendMessage(userMsg("hello"));
expect(() => session.branch("nonexistent")).toThrow("Entry nonexistent not found");
});
it("new appends become children of branch point", () => {
const session = SessionManager.inMemory();
const id1 = session.appendMessage(userMsg("1"));
const _id2 = session.appendMessage(assistantMsg("2"));
session.branch(id1);
const id3 = session.appendMessage(userMsg("branched"));
const entries = session.getEntries();
const branchedEntry = entries.find((e) => e.id === id3)!;
expect(branchedEntry.parentId).toBe(id1); // sibling of id2
});
});
describe("branchWithSummary", () => {
it("inserts branch summary and advances leaf", () => {
const session = SessionManager.inMemory();
const id1 = session.appendMessage(userMsg("1"));
const _id2 = session.appendMessage(assistantMsg("2"));
const _id3 = session.appendMessage(userMsg("3"));
const summaryId = session.branchWithSummary(id1, "Summary of abandoned work");
expect(session.getLeafId()).toBe(summaryId);
const entries = session.getEntries();
const summaryEntry = entries.find((e) => e.type === "branch_summary");
expect(summaryEntry).toBeDefined();
expect(summaryEntry?.parentId).toBe(id1);
if (summaryEntry?.type === "branch_summary") {
expect(summaryEntry.summary).toBe("Summary of abandoned work");
}
});
it("throws for non-existent entry", () => {
const session = SessionManager.inMemory();
session.appendMessage(userMsg("hello"));
expect(() => session.branchWithSummary("nonexistent", "summary")).toThrow("Entry nonexistent not found");
});
});
describe("getLeafEntry", () => {
it("returns undefined for empty session", () => {
const session = SessionManager.inMemory();
expect(session.getLeafEntry()).toBeUndefined();
});
it("returns current leaf entry", () => {
const session = SessionManager.inMemory();
session.appendMessage(userMsg("1"));
const id2 = session.appendMessage(assistantMsg("2"));
const leaf = session.getLeafEntry();
expect(leaf).toBeDefined();
expect(leaf!.id).toBe(id2);
});
});
describe("getEntry", () => {
it("returns undefined for non-existent id", () => {
const session = SessionManager.inMemory();
expect(session.getEntry("nonexistent")).toBeUndefined();
});
it("returns entry by id", () => {
const session = SessionManager.inMemory();
const id1 = session.appendMessage(userMsg("first"));
const id2 = session.appendMessage(assistantMsg("second"));
const entry1 = session.getEntry(id1);
expect(entry1).toBeDefined();
expect(entry1?.type).toBe("message");
if (entry1?.type === "message" && entry1.message.role === "user") {
expect(entry1.message.content).toBe("first");
}
const entry2 = session.getEntry(id2);
expect(entry2).toBeDefined();
if (entry2?.type === "message" && entry2.message.role === "assistant") {
expect((entry2.message.content as any)[0].text).toBe("second");
}
});
});
describe("buildSessionContext with branches", () => {
it("returns messages from current branch only", () => {
const session = SessionManager.inMemory();
// Main: 1 -> 2 -> 3
session.appendMessage(userMsg("msg1"));
const id2 = session.appendMessage(assistantMsg("msg2"));
session.appendMessage(userMsg("msg3"));
// Branch from 2: 2 -> 4
session.branch(id2);
session.appendMessage(assistantMsg("msg4-branch"));
const ctx = session.buildSessionContext();
expect(ctx.messages).toHaveLength(3); // msg1, msg2, msg4-branch (not msg3)
expect((ctx.messages[0] as any).content).toBe("msg1");
expect((ctx.messages[1] as any).content[0].text).toBe("msg2");
expect((ctx.messages[2] as any).content[0].text).toBe("msg4-branch");
});
});
});
describe("createBranchedSession", () => {
it("throws for non-existent entry", () => {
const session = SessionManager.inMemory();
session.appendMessage(userMsg("hello"));
expect(() => session.createBranchedSession("nonexistent")).toThrow("Entry nonexistent not found");
});
it("creates new session with path to specified leaf (in-memory)", () => {
const session = SessionManager.inMemory();
// Build: 1 -> 2 -> 3 -> 4
const id1 = session.appendMessage(userMsg("1"));
const id2 = session.appendMessage(assistantMsg("2"));
const id3 = session.appendMessage(userMsg("3"));
session.appendMessage(assistantMsg("4"));
// Branch from 3: 3 -> 5
session.branch(id3);
const _id5 = session.appendMessage(userMsg("5"));
// Create branched session from id2 (should only have 1 -> 2)
const result = session.createBranchedSession(id2);
expect(result).toBeUndefined(); // in-memory returns null
// Session should now only have entries 1 and 2
const entries = session.getEntries();
expect(entries).toHaveLength(2);
expect(entries[0].id).toBe(id1);
expect(entries[1].id).toBe(id2);
});
it("extracts correct path from branched tree", () => {
const session = SessionManager.inMemory();
// Build: 1 -> 2 -> 3
const id1 = session.appendMessage(userMsg("1"));
const id2 = session.appendMessage(assistantMsg("2"));
session.appendMessage(userMsg("3"));
// Branch from 2: 2 -> 4 -> 5
session.branch(id2);
const id4 = session.appendMessage(userMsg("4"));
const id5 = session.appendMessage(assistantMsg("5"));
// Create branched session from id5 (should have 1 -> 2 -> 4 -> 5)
session.createBranchedSession(id5);
const entries = session.getEntries();
expect(entries).toHaveLength(4);
expect(entries.map((e) => e.id)).toEqual([id1, id2, id4, id5]);
});
});

View file

@ -0,0 +1,158 @@
/**
* Shared test utilities for coding-agent tests.
*/
import { existsSync, mkdirSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { Agent } from "@mariozechner/pi-agent-core";
import { getModel } from "@mariozechner/pi-ai";
import { AgentSession } from "../src/core/agent-session.js";
import { AuthStorage } from "../src/core/auth-storage.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 { codingTools } from "../src/core/tools/index.js";
/**
* API key for authenticated tests. Tests using this should be wrapped in
* describe.skipIf(!API_KEY)
*/
export const API_KEY = process.env.ANTHROPIC_OAUTH_TOKEN || process.env.ANTHROPIC_API_KEY;
/**
* Create a minimal user message for testing.
*/
export function userMsg(text: string) {
return { role: "user" as const, content: text, timestamp: Date.now() };
}
/**
* Create a minimal assistant message for testing.
*/
export function assistantMsg(text: string) {
return {
role: "assistant" as const,
content: [{ type: "text" as const, text }],
api: "anthropic-messages" as const,
provider: "anthropic",
model: "test",
usage: {
input: 1,
output: 1,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 2,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
},
stopReason: "stop" as const,
timestamp: Date.now(),
};
}
/**
* Options for creating a test session.
*/
export interface TestSessionOptions {
/** Use in-memory session (no file persistence) */
inMemory?: boolean;
/** Custom system prompt */
systemPrompt?: string;
/** Custom settings overrides */
settingsOverrides?: Record<string, unknown>;
}
/**
* Resources returned by createTestSession that need cleanup.
*/
export interface TestSessionContext {
session: AgentSession;
sessionManager: SessionManager;
tempDir: string;
cleanup: () => void;
}
/**
* Create an AgentSession for testing with proper setup and cleanup.
* Use this for e2e tests that need real LLM calls.
*/
export function createTestSession(options: TestSessionOptions = {}): TestSessionContext {
const tempDir = join(tmpdir(), `pi-test-${Date.now()}-${Math.random().toString(36).slice(2)}`);
mkdirSync(tempDir, { recursive: true });
const model = getModel("anthropic", "claude-sonnet-4-5")!;
const agent = new Agent({
getApiKey: () => API_KEY,
initialState: {
model,
systemPrompt: options.systemPrompt ?? "You are a helpful assistant. Be extremely concise.",
tools: codingTools,
},
});
const sessionManager = options.inMemory ? SessionManager.inMemory() : SessionManager.create(tempDir);
const settingsManager = SettingsManager.create(tempDir, tempDir);
if (options.settingsOverrides) {
settingsManager.applyOverrides(options.settingsOverrides);
}
const authStorage = new AuthStorage(join(tempDir, "auth.json"));
const modelRegistry = new ModelRegistry(authStorage, tempDir);
const session = new AgentSession({
agent,
sessionManager,
settingsManager,
modelRegistry,
});
// Must subscribe to enable session persistence
session.subscribe(() => {});
const cleanup = () => {
session.dispose();
if (tempDir && existsSync(tempDir)) {
rmSync(tempDir, { recursive: true });
}
};
return { session, sessionManager, tempDir, cleanup };
}
/**
* Build a session tree for testing using SessionManager.
* Returns the IDs of all created entries.
*
* Example tree structure:
* ```
* u1 -> a1 -> u2 -> a2
* -> u3 -> a3 (branch from a1)
* u4 -> a4 (another root)
* ```
*/
export function buildTestTree(
session: SessionManager,
structure: {
messages: Array<{ role: "user" | "assistant"; text: string; branchFrom?: string }>;
},
): Map<string, string> {
const ids = new Map<string, string>();
for (const msg of structure.messages) {
if (msg.branchFrom) {
const branchFromId = ids.get(msg.branchFrom);
if (!branchFromId) {
throw new Error(`Cannot branch from unknown entry: ${msg.branchFrom}`);
}
session.branch(branchFromId);
}
const id =
msg.role === "user" ? session.appendMessage(userMsg(msg.text)) : session.appendMessage(assistantMsg(msg.text));
ids.set(msg.text, id);
}
return ids;
}