move pi-mono into companion-cloud as apps/companion-os

- Copy all pi-mono source into apps/companion-os/
- Update Dockerfile to COPY pre-built binary instead of downloading from GitHub Releases
- Update deploy-staging.yml to build pi from source (bun compile) before Docker build
- Add apps/companion-os/** to path triggers
- No more cross-repo dispatch needed

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Harivansh Rathi 2026-03-07 09:22:50 -08:00
commit 0250f72976
579 changed files with 206942 additions and 0 deletions

View file

@ -0,0 +1,342 @@
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,224 @@
import { mkdirSync, readFileSync, rmSync, writeFileSync } from "fs";
import { tmpdir } from "os";
import { join } from "path";
import { afterEach, beforeEach, describe, expect, it } from "vitest";
import {
findMostRecentSession,
loadEntriesFromFile,
SessionManager,
} 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);
});
});
describe("SessionManager.setSessionFile with corrupted files", () => {
let tempDir: string;
beforeEach(() => {
tempDir = join(tmpdir(), `session-test-${Date.now()}`);
mkdirSync(tempDir, { recursive: true });
});
afterEach(() => {
rmSync(tempDir, { recursive: true, force: true });
});
it("truncates and rewrites empty file with valid header", () => {
const emptyFile = join(tempDir, "empty.jsonl");
writeFileSync(emptyFile, "");
const sm = SessionManager.open(emptyFile, tempDir);
// Should have created a new session with valid header
expect(sm.getSessionId()).toBeTruthy();
expect(sm.getHeader()).toBeTruthy();
expect(sm.getHeader()?.type).toBe("session");
// File should now contain a valid header
const content = readFileSync(emptyFile, "utf-8");
const lines = content.trim().split("\n").filter(Boolean);
expect(lines.length).toBe(1);
const header = JSON.parse(lines[0]);
expect(header.type).toBe("session");
expect(header.id).toBe(sm.getSessionId());
});
it("truncates and rewrites file without valid header", () => {
const noHeaderFile = join(tempDir, "no-header.jsonl");
// File with messages but no session header (corrupted state)
writeFileSync(
noHeaderFile,
'{"type":"message","id":"abc","parentId":"orphaned","timestamp":"2025-01-01T00:00:00Z","message":{"role":"assistant","content":"test"}}\n',
);
const sm = SessionManager.open(noHeaderFile, tempDir);
// Should have created a new session with valid header
expect(sm.getSessionId()).toBeTruthy();
expect(sm.getHeader()).toBeTruthy();
expect(sm.getHeader()?.type).toBe("session");
// File should now contain only a valid header (old content truncated)
const content = readFileSync(noHeaderFile, "utf-8");
const lines = content.trim().split("\n").filter(Boolean);
expect(lines.length).toBe(1);
const header = JSON.parse(lines[0]);
expect(header.type).toBe("session");
expect(header.id).toBe(sm.getSessionId());
});
it("preserves explicit session file path when recovering from corrupted file", () => {
const explicitPath = join(tempDir, "my-session.jsonl");
writeFileSync(explicitPath, "");
const sm = SessionManager.open(explicitPath, tempDir);
// The session file path should be preserved
expect(sm.getSessionFile()).toBe(explicitPath);
});
it("subsequent loads of recovered file work correctly", () => {
const corruptedFile = join(tempDir, "corrupted.jsonl");
writeFileSync(corruptedFile, "garbage content\n");
// First open recovers the file
const sm1 = SessionManager.open(corruptedFile, tempDir);
const sessionId = sm1.getSessionId();
// Second open should load the recovered file successfully
const sm2 = SessionManager.open(corruptedFile, tempDir);
expect(sm2.getSessionId()).toBe(sessionId);
expect(sm2.getHeader()?.type).toBe("session");
});
});

View file

@ -0,0 +1,217 @@
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,96 @@
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 (v3 is current after hookMessage->custom migration)
expect((entries[0] as any).version).toBe(3);
// 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,62 @@
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_data", { 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_data");
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,549 @@
import { existsSync, mkdirSync, readFileSync, rmSync } from "fs";
import { tmpdir } from "os";
import { join } from "path";
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_data", { 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_data");
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]);
});
it("does not duplicate entries when forking from first user message", () => {
const tempDir = join(tmpdir(), `session-fork-dedup-${Date.now()}`);
mkdirSync(tempDir, { recursive: true });
try {
// Create a persisted session with a couple of turns
const session = SessionManager.create(tempDir, tempDir);
const id1 = session.appendMessage(userMsg("first question"));
session.appendMessage(assistantMsg("first answer"));
session.appendMessage(userMsg("second question"));
session.appendMessage(assistantMsg("second answer"));
// Fork from the very first user message (no assistant in the branched path)
const newFile = session.createBranchedSession(id1);
expect(newFile).toBeDefined();
// The branched path has no assistant, so the file should not exist yet
// (deferred to _persist on first assistant, matching newSession() contract)
expect(existsSync(newFile!)).toBe(false);
// Simulate extension adding entry before assistant (like preset on turn_start)
session.appendCustomEntry("preset-state", { name: "plan" });
// Now the assistant responds
session.appendMessage(assistantMsg("new answer"));
// File should now exist with exactly one header and no duplicate IDs
expect(existsSync(newFile!)).toBe(true);
const content = readFileSync(newFile!, "utf-8");
const lines = content.trim().split("\n").filter(Boolean);
const records = lines.map((line) => JSON.parse(line));
expect(records.filter((r) => r.type === "session")).toHaveLength(1);
const entryIds = records
.filter((r) => r.type !== "session")
.map((r) => r.id)
.filter((id): id is string => typeof id === "string");
expect(new Set(entryIds).size).toBe(entryIds.length);
} finally {
rmSync(tempDir, { recursive: true, force: true });
}
});
it("writes file immediately when forking from a point with assistant messages", () => {
const tempDir = join(tmpdir(), `session-fork-with-assistant-${Date.now()}`);
mkdirSync(tempDir, { recursive: true });
try {
const session = SessionManager.create(tempDir, tempDir);
session.appendMessage(userMsg("first question"));
const id2 = session.appendMessage(assistantMsg("first answer"));
session.appendMessage(userMsg("second question"));
session.appendMessage(assistantMsg("second answer"));
// Fork including the assistant message
const newFile = session.createBranchedSession(id2);
expect(newFile).toBeDefined();
// Path includes an assistant, so file should be written immediately
expect(existsSync(newFile!)).toBe(true);
const content = readFileSync(newFile!, "utf-8");
const lines = content.trim().split("\n").filter(Boolean);
const records = lines.map((line) => JSON.parse(line));
expect(records.filter((r) => r.type === "session")).toHaveLength(1);
} finally {
rmSync(tempDir, { recursive: true, force: true });
}
});
});