Cristina Poncela Cubeiro
9e7582aa03
This PR: - Adds retainedTail to compaction entries in the new agent harness so we don't have to walk up the tree for the 2000 tokens before compaction, - Changes getPathToRoot to getPathToRootOrCompaction to only load until last compaction, as unnecessary to access all nodes where it is called, - Adds a SQLite storage backend, in a separate packages/session-backend-sqlite, with a migration system and schemas as per on-site discussions: sessions to match session header messages (except for metadata, which I couldn't understand what it's used for or where it gets written, so I omitted it), session_entries for shared entry types as columns plus payload as a json for what remains, session_sequences to represent the append-only, serialized nature of the jsonl files, branch_entries to attribute nodes to branches (relationship one-to-many), and session_materialized with the session info (see /session in TUI) to act as a "cache" or quick-access for costs, message count, token info, labels, session name, and model-thinking-level config (e.g. for fast resume). - This is compatible with the new agent harness Session abstraction.
505 lines
16 KiB
TypeScript
505 lines
16 KiB
TypeScript
import { existsSync, readFileSync, writeFileSync } from "node:fs";
|
|
import { join } from "node:path";
|
|
import { describe, expect, it } from "vitest";
|
|
import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts";
|
|
import { JsonlSessionStorage, loadJsonlSessionMetadata } from "../../src/harness/session/jsonl-storage.ts";
|
|
import { InMemorySessionStorage } from "../../src/harness/session/memory-storage.ts";
|
|
import {
|
|
type BranchSummaryEntry,
|
|
type CompactionEntry,
|
|
type MessageEntry,
|
|
ok,
|
|
type SessionMetadata,
|
|
} from "../../src/harness/types.ts";
|
|
import { createAssistantMessage, createTempDir, createUserMessage } from "./session-test-utils.ts";
|
|
|
|
describe("InMemorySessionStorage", () => {
|
|
it("returns configured session metadata", async () => {
|
|
const metadata: SessionMetadata = { id: "session-1", createdAt: "2026-01-01T00:00:00.000Z" };
|
|
const storage = new InMemorySessionStorage({ metadata });
|
|
expect(await storage.getMetadata()).toEqual(metadata);
|
|
});
|
|
|
|
it("copies initial entries and persists leaf changes", async () => {
|
|
const entry: MessageEntry = {
|
|
type: "message",
|
|
id: "entry-1",
|
|
parentId: null,
|
|
timestamp: "2026-01-01T00:00:00.000Z",
|
|
message: createUserMessage("one"),
|
|
};
|
|
const initialEntries = [entry];
|
|
const storage = new InMemorySessionStorage({ entries: initialEntries });
|
|
initialEntries.push({ ...entry, id: "entry-2" });
|
|
expect((await storage.getEntries()).map((storedEntry) => storedEntry.id)).toEqual(["entry-1"]);
|
|
expect(await storage.getLeafId()).toBe("entry-1");
|
|
await storage.setLeafId(null);
|
|
expect(await storage.getLeafId()).toBeNull();
|
|
expect((await storage.getEntries()).at(-1)).toMatchObject({ type: "leaf", targetId: null });
|
|
});
|
|
|
|
it("rejects invalid leaf ids", async () => {
|
|
const storage = new InMemorySessionStorage();
|
|
await expect(storage.setLeafId("missing")).rejects.toThrow("Entry missing not found");
|
|
});
|
|
|
|
it("finds entries by type", async () => {
|
|
const entry: MessageEntry = {
|
|
type: "message",
|
|
id: "entry-1",
|
|
parentId: null,
|
|
timestamp: "2026-01-01T00:00:00.000Z",
|
|
message: createUserMessage("one"),
|
|
};
|
|
const storage = new InMemorySessionStorage({ entries: [entry] });
|
|
expect((await storage.findEntries("message")).map((found) => found.id)).toEqual(["entry-1"]);
|
|
expect(await storage.findEntries("session_info")).toEqual([]);
|
|
});
|
|
|
|
it("maintains label lookup", async () => {
|
|
const entry: MessageEntry = {
|
|
type: "message",
|
|
id: "entry-1",
|
|
parentId: null,
|
|
timestamp: "2026-01-01T00:00:00.000Z",
|
|
message: createUserMessage("one"),
|
|
};
|
|
const storage = new InMemorySessionStorage({ entries: [entry] });
|
|
expect(await storage.getLabel("entry-1")).toBeUndefined();
|
|
await storage.appendEntry({
|
|
type: "label",
|
|
id: "label-1",
|
|
parentId: "entry-1",
|
|
timestamp: "2026-01-01T00:00:01.000Z",
|
|
targetId: "entry-1",
|
|
label: "checkpoint",
|
|
});
|
|
expect(await storage.getLabel("entry-1")).toBe("checkpoint");
|
|
await storage.appendEntry({
|
|
type: "label",
|
|
id: "label-2",
|
|
parentId: "label-1",
|
|
timestamp: "2026-01-01T00:00:02.000Z",
|
|
targetId: "entry-1",
|
|
label: undefined,
|
|
});
|
|
expect(await storage.getLabel("entry-1")).toBeUndefined();
|
|
});
|
|
|
|
it("includes summary-entry usage in session stats", async () => {
|
|
const assistant: MessageEntry = {
|
|
type: "message",
|
|
id: "assistant",
|
|
parentId: null,
|
|
timestamp: "2026-01-01T00:00:00.000Z",
|
|
message: {
|
|
role: "assistant",
|
|
content: [{ type: "text", text: "reply" }],
|
|
api: "anthropic-messages",
|
|
provider: "anthropic",
|
|
model: "claude-sonnet-4-5",
|
|
usage: {
|
|
input: 10,
|
|
output: 20,
|
|
cacheRead: 30,
|
|
cacheWrite: 40,
|
|
totalTokens: 100,
|
|
cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 },
|
|
},
|
|
stopReason: "stop",
|
|
timestamp: 0,
|
|
},
|
|
};
|
|
const compaction: CompactionEntry = {
|
|
type: "compaction",
|
|
id: "compaction",
|
|
parentId: "assistant",
|
|
timestamp: "2026-01-01T00:00:01.000Z",
|
|
summary: "summary",
|
|
firstKeptEntryId: "assistant",
|
|
tokensBefore: 1234,
|
|
usage: {
|
|
input: 1,
|
|
output: 2,
|
|
cacheRead: 3,
|
|
cacheWrite: 4,
|
|
totalTokens: 10,
|
|
cost: { input: 0.01, output: 0.02, cacheRead: 0.03, cacheWrite: 0.04, total: 0.1 },
|
|
},
|
|
};
|
|
const branchSummary: BranchSummaryEntry = {
|
|
type: "branch_summary",
|
|
id: "branch-summary",
|
|
parentId: "compaction",
|
|
timestamp: "2026-01-01T00:00:02.000Z",
|
|
fromId: "assistant",
|
|
summary: "branch",
|
|
usage: {
|
|
input: 5,
|
|
output: 6,
|
|
cacheRead: 7,
|
|
cacheWrite: 8,
|
|
totalTokens: 26,
|
|
cost: { input: 0.05, output: 0.06, cacheRead: 0.07, cacheWrite: 0.08, total: 0.26 },
|
|
},
|
|
};
|
|
const storage = new InMemorySessionStorage({ entries: [assistant, compaction, branchSummary] });
|
|
expect(await storage.getSessionStats()).toEqual({
|
|
messageCount: 1,
|
|
cachedTokens: 40,
|
|
uncachedTokens: 68,
|
|
totalTokens: 136,
|
|
costTotal: 1.36,
|
|
});
|
|
});
|
|
|
|
it("walks paths to root or retained-tail compaction", async () => {
|
|
const root: MessageEntry = {
|
|
type: "message",
|
|
id: "root",
|
|
parentId: null,
|
|
timestamp: "2026-01-01T00:00:00.000Z",
|
|
message: createUserMessage("root"),
|
|
};
|
|
const child: MessageEntry = {
|
|
...root,
|
|
id: "child",
|
|
parentId: "root",
|
|
message: createAssistantMessage("child"),
|
|
};
|
|
const compaction: CompactionEntry = {
|
|
type: "compaction",
|
|
id: "compaction",
|
|
parentId: "child",
|
|
timestamp: "2026-01-01T00:00:01.000Z",
|
|
summary: "summary",
|
|
firstKeptEntryId: "child",
|
|
tokensBefore: 1234,
|
|
retainedTail: [createAssistantMessage("child")],
|
|
};
|
|
const afterCompaction: MessageEntry = {
|
|
...root,
|
|
id: "after-compaction",
|
|
parentId: "compaction",
|
|
message: createUserMessage("after"),
|
|
};
|
|
const storage = new InMemorySessionStorage({ entries: [root, child, compaction, afterCompaction] });
|
|
expect((await storage.getPathToRootOrCompaction("child")).map((entry) => entry.id)).toEqual(["root", "child"]);
|
|
expect((await storage.getPathToRootOrCompaction("after-compaction")).map((entry) => entry.id)).toEqual([
|
|
"compaction",
|
|
"after-compaction",
|
|
]);
|
|
expect(await storage.getPathToRootOrCompaction(null)).toEqual([]);
|
|
});
|
|
});
|
|
|
|
describe("JsonlSessionStorage", () => {
|
|
it("throws for missing files when opening", async () => {
|
|
const dir = createTempDir();
|
|
const env = new NodeExecutionEnv({ cwd: dir });
|
|
const filePath = join(dir, "session.jsonl");
|
|
await expect(JsonlSessionStorage.open(env, filePath)).rejects.toMatchObject({ code: "not_found" });
|
|
});
|
|
|
|
it("writes the header on create", async () => {
|
|
const dir = createTempDir();
|
|
const env = new NodeExecutionEnv({ cwd: dir });
|
|
const filePath = join(dir, "session.jsonl");
|
|
const storage = await JsonlSessionStorage.create(env, filePath, { cwd: dir, sessionId: "session-1" });
|
|
expect(existsSync(filePath)).toBe(true);
|
|
expect(readFileSync(filePath, "utf8").trim().split("\n")).toHaveLength(1);
|
|
expect(await storage.getLeafId()).toBeNull();
|
|
expect(await storage.getEntries()).toEqual([]);
|
|
await storage.appendEntry({
|
|
type: "message",
|
|
id: "user-1",
|
|
parentId: null,
|
|
timestamp: "2026-01-01T00:00:00.000Z",
|
|
message: createUserMessage("one"),
|
|
});
|
|
const lines = readFileSync(filePath, "utf8").trim().split("\n");
|
|
expect(JSON.parse(lines[0]!).type).toBe("session");
|
|
expect(JSON.parse(lines[1]!).id).toBe("user-1");
|
|
expect(lines).toHaveLength(2);
|
|
});
|
|
|
|
it("throws for malformed session headers", async () => {
|
|
const dir = createTempDir();
|
|
const env = new NodeExecutionEnv({ cwd: dir });
|
|
const filePath = join(dir, "session.jsonl");
|
|
writeFileSync(filePath, "not json\n");
|
|
await expect(JsonlSessionStorage.open(env, filePath)).rejects.toThrow("first line is not a valid session header");
|
|
});
|
|
|
|
it("throws for malformed entry lines", async () => {
|
|
const dir = createTempDir();
|
|
const env = new NodeExecutionEnv({ cwd: dir });
|
|
const filePath = join(dir, "session.jsonl");
|
|
const header = {
|
|
type: "session",
|
|
version: 3,
|
|
id: "session-1",
|
|
timestamp: "2026-01-01T00:00:00.000Z",
|
|
cwd: dir,
|
|
};
|
|
const entry: MessageEntry = {
|
|
type: "message",
|
|
id: "entry-1",
|
|
parentId: null,
|
|
timestamp: "2026-01-01T00:00:00.000Z",
|
|
message: createUserMessage("one"),
|
|
};
|
|
writeFileSync(filePath, `${JSON.stringify(header)}\nnot json\n${JSON.stringify(entry)}\n`);
|
|
await expect(JsonlSessionStorage.open(env, filePath)).rejects.toMatchObject({ code: "invalid_entry" });
|
|
});
|
|
|
|
it("creates and reads session metadata from the header", async () => {
|
|
const dir = createTempDir();
|
|
const env = new NodeExecutionEnv({ cwd: dir });
|
|
const filePath = join(dir, "session.jsonl");
|
|
const storage = await JsonlSessionStorage.create(env, filePath, {
|
|
cwd: dir,
|
|
sessionId: "session-1",
|
|
parentSessionPath: "/tmp/parent.jsonl",
|
|
});
|
|
const metadata = await storage.getMetadata();
|
|
expect(metadata).toMatchObject({
|
|
id: "session-1",
|
|
cwd: dir,
|
|
path: filePath,
|
|
parentSessionPath: "/tmp/parent.jsonl",
|
|
});
|
|
await storage.appendEntry({
|
|
type: "message",
|
|
id: "user-1",
|
|
parentId: null,
|
|
timestamp: "2026-01-01T00:00:00.000Z",
|
|
message: createUserMessage("one"),
|
|
});
|
|
expect(await loadJsonlSessionMetadata(env, filePath)).toEqual(metadata);
|
|
});
|
|
|
|
it("round-trips custom header metadata", async () => {
|
|
const dir = createTempDir();
|
|
const env = new NodeExecutionEnv({ cwd: dir });
|
|
const filePath = join(dir, "session.jsonl");
|
|
const storage = await JsonlSessionStorage.create(env, filePath, {
|
|
cwd: dir,
|
|
sessionId: "session-1",
|
|
metadata: { profile: "reviewer" },
|
|
});
|
|
expect((await storage.getMetadata()).metadata).toEqual({ profile: "reviewer" });
|
|
const loaded = await JsonlSessionStorage.open(env, filePath);
|
|
expect((await loaded.getMetadata()).metadata).toEqual({ profile: "reviewer" });
|
|
expect((await loadJsonlSessionMetadata(env, filePath)).metadata).toEqual({ profile: "reviewer" });
|
|
});
|
|
|
|
it("omits header metadata when not provided", async () => {
|
|
const dir = createTempDir();
|
|
const env = new NodeExecutionEnv({ cwd: dir });
|
|
const filePath = join(dir, "session.jsonl");
|
|
await JsonlSessionStorage.create(env, filePath, { cwd: dir, sessionId: "session-1" });
|
|
expect(JSON.parse(readFileSync(filePath, "utf8").trim())).not.toHaveProperty("metadata");
|
|
expect((await loadJsonlSessionMetadata(env, filePath)).metadata).toBeUndefined();
|
|
});
|
|
|
|
it("throws for non-object header metadata", async () => {
|
|
const dir = createTempDir();
|
|
const env = new NodeExecutionEnv({ cwd: dir });
|
|
const filePath = join(dir, "session.jsonl");
|
|
const header = {
|
|
type: "session",
|
|
version: 3,
|
|
id: "session-1",
|
|
timestamp: "2026-01-01T00:00:00.000Z",
|
|
cwd: dir,
|
|
metadata: "profile",
|
|
};
|
|
writeFileSync(filePath, `${JSON.stringify(header)}\n`);
|
|
await expect(JsonlSessionStorage.open(env, filePath)).rejects.toThrow(
|
|
"session header metadata must be an object",
|
|
);
|
|
});
|
|
|
|
it("loads existing entries and reconstructs leaf", async () => {
|
|
const dir = createTempDir();
|
|
const env = new NodeExecutionEnv({ cwd: dir });
|
|
const filePath = join(dir, "session.jsonl");
|
|
const storage = await JsonlSessionStorage.create(env, filePath, { cwd: dir, sessionId: "session-1" });
|
|
const root: MessageEntry = {
|
|
type: "message",
|
|
id: "root",
|
|
parentId: null,
|
|
timestamp: "2026-01-01T00:00:00.000Z",
|
|
message: createUserMessage("root"),
|
|
};
|
|
const child: MessageEntry = {
|
|
...root,
|
|
id: "child",
|
|
parentId: "root",
|
|
message: createAssistantMessage("child"),
|
|
};
|
|
await storage.appendEntry(root);
|
|
await storage.appendEntry(child);
|
|
const loaded = await JsonlSessionStorage.open(env, filePath);
|
|
expect(await loaded.getLeafId()).toBe("child");
|
|
expect((await loaded.getEntries()).map((entry) => entry.id)).toEqual(["root", "child"]);
|
|
await loaded.setLeafId("root");
|
|
const reloaded = await JsonlSessionStorage.open(env, filePath);
|
|
expect(await reloaded.getLeafId()).toBe("root");
|
|
expect((await reloaded.getEntries()).at(-1)).toMatchObject({ type: "leaf", targetId: "root" });
|
|
expect((await loaded.getPathToRootOrCompaction("child")).map((entry) => entry.id)).toEqual(["root", "child"]);
|
|
});
|
|
|
|
it("finds entries by type", async () => {
|
|
const dir = createTempDir();
|
|
const env = new NodeExecutionEnv({ cwd: dir });
|
|
const filePath = join(dir, "session.jsonl");
|
|
const storage = await JsonlSessionStorage.create(env, filePath, { cwd: dir, sessionId: "session-1" });
|
|
await storage.appendEntry({
|
|
type: "message",
|
|
id: "entry-1",
|
|
parentId: null,
|
|
timestamp: "2026-01-01T00:00:00.000Z",
|
|
message: createUserMessage("one"),
|
|
});
|
|
expect((await storage.findEntries("message")).map((found) => found.id)).toEqual(["entry-1"]);
|
|
expect(await storage.findEntries("session_info")).toEqual([]);
|
|
});
|
|
|
|
it("maintains label lookup", async () => {
|
|
const dir = createTempDir();
|
|
const env = new NodeExecutionEnv({ cwd: dir });
|
|
const filePath = join(dir, "session.jsonl");
|
|
const storage = await JsonlSessionStorage.create(env, filePath, { cwd: dir, sessionId: "session-1" });
|
|
await storage.appendEntry({
|
|
type: "message",
|
|
id: "entry-1",
|
|
parentId: null,
|
|
timestamp: "2026-01-01T00:00:00.000Z",
|
|
message: createUserMessage("one"),
|
|
});
|
|
expect(await storage.getLabel("entry-1")).toBeUndefined();
|
|
await storage.appendEntry({
|
|
type: "label",
|
|
id: "label-1",
|
|
parentId: "entry-1",
|
|
timestamp: "2026-01-01T00:00:01.000Z",
|
|
targetId: "entry-1",
|
|
label: "checkpoint",
|
|
});
|
|
expect(await storage.getLabel("entry-1")).toBe("checkpoint");
|
|
await storage.appendEntry({
|
|
type: "label",
|
|
id: "label-2",
|
|
parentId: "label-1",
|
|
timestamp: "2026-01-01T00:00:02.000Z",
|
|
targetId: "entry-1",
|
|
label: undefined,
|
|
});
|
|
expect(await storage.getLabel("entry-1")).toBeUndefined();
|
|
const loaded = await JsonlSessionStorage.open(env, filePath);
|
|
expect(await loaded.getLabel("entry-1")).toBeUndefined();
|
|
});
|
|
|
|
it("includes summary-entry usage in session stats", async () => {
|
|
const dir = createTempDir();
|
|
const env = new NodeExecutionEnv({ cwd: dir });
|
|
const filePath = join(dir, "session.jsonl");
|
|
const storage = await JsonlSessionStorage.create(env, filePath, { cwd: dir, sessionId: "session-1" });
|
|
await storage.appendEntry({
|
|
type: "message",
|
|
id: "assistant",
|
|
parentId: null,
|
|
timestamp: "2026-01-01T00:00:00.000Z",
|
|
message: {
|
|
role: "assistant",
|
|
content: [{ type: "text", text: "reply" }],
|
|
api: "anthropic-messages",
|
|
provider: "anthropic",
|
|
model: "claude-sonnet-4-5",
|
|
usage: {
|
|
input: 10,
|
|
output: 20,
|
|
cacheRead: 30,
|
|
cacheWrite: 40,
|
|
totalTokens: 100,
|
|
cost: { input: 0.1, output: 0.2, cacheRead: 0.3, cacheWrite: 0.4, total: 1 },
|
|
},
|
|
stopReason: "stop",
|
|
timestamp: 0,
|
|
},
|
|
});
|
|
await storage.appendEntry({
|
|
type: "compaction",
|
|
id: "compaction",
|
|
parentId: "assistant",
|
|
timestamp: "2026-01-01T00:00:01.000Z",
|
|
summary: "summary",
|
|
firstKeptEntryId: "assistant",
|
|
tokensBefore: 1234,
|
|
usage: {
|
|
input: 1,
|
|
output: 2,
|
|
cacheRead: 3,
|
|
cacheWrite: 4,
|
|
totalTokens: 10,
|
|
cost: { input: 0.01, output: 0.02, cacheRead: 0.03, cacheWrite: 0.04, total: 0.1 },
|
|
},
|
|
});
|
|
await storage.appendEntry({
|
|
type: "branch_summary",
|
|
id: "branch-summary",
|
|
parentId: "compaction",
|
|
timestamp: "2026-01-01T00:00:02.000Z",
|
|
fromId: "assistant",
|
|
summary: "branch",
|
|
usage: {
|
|
input: 5,
|
|
output: 6,
|
|
cacheRead: 7,
|
|
cacheWrite: 8,
|
|
totalTokens: 26,
|
|
cost: { input: 0.05, output: 0.06, cacheRead: 0.07, cacheWrite: 0.08, total: 0.26 },
|
|
},
|
|
});
|
|
expect(await storage.getSessionStats()).toEqual({
|
|
messageCount: 1,
|
|
cachedTokens: 40,
|
|
uncachedTokens: 68,
|
|
totalTokens: 136,
|
|
costTotal: 1.36,
|
|
});
|
|
});
|
|
|
|
it("reads session metadata through the line-reading filesystem operation", async () => {
|
|
const dir = createTempDir();
|
|
const filePath = join(dir, "session.jsonl");
|
|
const header = {
|
|
type: "session",
|
|
version: 3,
|
|
id: "session-1",
|
|
timestamp: "2026-01-01T00:00:00.000Z",
|
|
cwd: dir,
|
|
};
|
|
const metadata = await loadJsonlSessionMetadata(
|
|
{
|
|
readTextLines: async () => ok([JSON.stringify(header)]),
|
|
readTextFile: async () => {
|
|
throw new Error("readTextFile should not be called for metadata");
|
|
},
|
|
writeFile: async () => ok(undefined),
|
|
appendFile: async () => ok(undefined),
|
|
},
|
|
filePath,
|
|
);
|
|
expect(metadata).toEqual({
|
|
id: "session-1",
|
|
createdAt: "2026-01-01T00:00:00.000Z",
|
|
cwd: dir,
|
|
path: filePath,
|
|
parentSessionPath: undefined,
|
|
});
|
|
});
|
|
});
|