mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
192 lines
8.1 KiB
TypeScript
192 lines
8.1 KiB
TypeScript
import { existsSync, readFileSync, writeFileSync } from "node:fs";
|
|
import { describe, expect, it } from "vitest";
|
|
import { NodeExecutionEnv } from "../../src/harness/env/nodejs.ts";
|
|
import { createJsonlSessionStore, loadJsonlSessionMetadata } from "../../src/harness/session/jsonl-store.ts";
|
|
import { createInMemorySessionStore } from "../../src/harness/session/memory-store.ts";
|
|
import { createSessionRepository } from "../../src/harness/session/repository.ts";
|
|
import type { Session } from "../../src/harness/session/session.ts";
|
|
import { createAssistantMessage, createTempDir, createUserMessage } from "./session-test-utils.ts";
|
|
|
|
async function appendUsageEntries(session: Session) {
|
|
const assistant = createAssistantMessage("reply");
|
|
if (assistant.role !== "assistant") throw new Error("Expected assistant test message");
|
|
assistant.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 },
|
|
};
|
|
await session.appendMessage(assistant);
|
|
await session.appendCompaction("summary", undefined, 1234, undefined, undefined, {
|
|
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 session.moveTo(await session.getLeafId(), {
|
|
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 },
|
|
},
|
|
});
|
|
}
|
|
|
|
function createInMemoryRepository() {
|
|
return createSessionRepository({ store: createInMemorySessionStore() });
|
|
}
|
|
|
|
function createJsonlRepository(options: Parameters<typeof createJsonlSessionStore>[0]) {
|
|
return createSessionRepository({ store: createJsonlSessionStore(options) });
|
|
}
|
|
|
|
describe("Session aggregate", () => {
|
|
it("owns leaf navigation, labels, names, stats, and branch traversal", async () => {
|
|
const repo = createInMemoryRepository();
|
|
const session = await repo.create({ id: "session-1" });
|
|
const root = await session.appendMessage(createUserMessage("root"));
|
|
const child = await session.appendMessage(createAssistantMessage("child"));
|
|
await session.appendLabel(root, "checkpoint");
|
|
await session.appendSessionName(" review\nname ");
|
|
|
|
expect(await session.getLeafId()).not.toBeNull();
|
|
expect(await session.getLabel(root)).toBe("checkpoint");
|
|
expect(await session.getSessionName()).toBe("review name");
|
|
expect((await session.getBranch(child)).map((entry) => entry.id)).toEqual([root, child]);
|
|
|
|
await session.moveTo(root);
|
|
expect(await session.getLeafId()).toBe(root);
|
|
expect((await session.getEntries()).at(-1)).toMatchObject({ type: "leaf", targetId: root });
|
|
await expect(session.moveTo("missing")).rejects.toThrow("Entry missing not found");
|
|
});
|
|
|
|
it("serializes concurrent appends into one parent chain", async () => {
|
|
const session = await createInMemoryRepository().create({});
|
|
await Promise.all(
|
|
Array.from({ length: 20 }, (_, index) => session.appendMessage(createUserMessage(`message ${index}`))),
|
|
);
|
|
const entries = await session.getEntries();
|
|
expect(entries).toHaveLength(20);
|
|
for (let index = 0; index < entries.length; index++) {
|
|
expect(entries[index]!.parentId).toBe(index === 0 ? null : entries[index - 1]!.id);
|
|
}
|
|
});
|
|
|
|
it("includes assistant and summary usage in statistics", async () => {
|
|
const session = await createInMemoryRepository().create({});
|
|
await appendUsageEntries(session);
|
|
expect(await session.getSessionStats()).toEqual({
|
|
messageCount: 1,
|
|
cachedTokens: 40,
|
|
uncachedTokens: 68,
|
|
totalTokens: 136,
|
|
costTotal: 1.36,
|
|
});
|
|
});
|
|
|
|
it("stops branch traversal at retained-tail compaction", async () => {
|
|
const session = await createInMemoryRepository().create({});
|
|
await session.appendMessage(createUserMessage("root"));
|
|
const child = await session.appendMessage(createAssistantMessage("child"));
|
|
const compaction = await session.appendCompaction("summary", child, 1234, undefined, undefined, undefined, [
|
|
createAssistantMessage("child"),
|
|
]);
|
|
const tail = await session.appendMessage(createUserMessage("after"));
|
|
expect((await session.getBranch(tail)).map((entry) => entry.id)).toEqual([compaction, tail]);
|
|
});
|
|
});
|
|
|
|
describe("JsonlSessionStore", () => {
|
|
it("writes headers and entries and reopens the aggregate", async () => {
|
|
const root = createTempDir();
|
|
const repo = createJsonlRepository({ fs: new NodeExecutionEnv({ cwd: root }), sessionsRoot: root });
|
|
const session = await repo.create({ cwd: root, id: "session-1", metadata: { profile: "reviewer" } });
|
|
const metadata = await session.getMetadata();
|
|
const entryId = await session.appendMessage(createUserMessage("one"));
|
|
const reopened = await repo.open(metadata);
|
|
|
|
expect(existsSync(metadata.path)).toBe(true);
|
|
expect((await reopened.getEntries()).map((entry) => entry.id)).toEqual([entryId]);
|
|
expect((await loadJsonlSessionMetadata(new NodeExecutionEnv({ cwd: root }), metadata.path)).metadata).toEqual({
|
|
profile: "reviewer",
|
|
});
|
|
const lines = readFileSync(metadata.path, "utf8").trim().split("\n");
|
|
expect(JSON.parse(lines[0]!)).toMatchObject({ type: "session", version: 3, id: "session-1" });
|
|
expect(JSON.parse(lines[1]!)).toMatchObject({ id: entryId, type: "message" });
|
|
});
|
|
|
|
it("fails loudly for malformed headers and entries", async () => {
|
|
const root = createTempDir();
|
|
const env = new NodeExecutionEnv({ cwd: root });
|
|
const repo = createJsonlRepository({ fs: env, sessionsRoot: root });
|
|
const session = await repo.create({ cwd: root, id: "session-1" });
|
|
const metadata = await session.getMetadata();
|
|
writeFileSync(metadata.path, "not json\n");
|
|
await expect(repo.open(metadata)).rejects.toMatchObject({ code: "invalid_session" });
|
|
|
|
const header = { type: "session", version: 3, id: "session-1", timestamp: metadata.createdAt, cwd: root };
|
|
writeFileSync(metadata.path, `${JSON.stringify(header)}\nnot json\n`);
|
|
await expect(repo.open(metadata)).rejects.toMatchObject({ code: "invalid_entry" });
|
|
});
|
|
|
|
it("enforces entry uniqueness and does not recreate deleted files", async () => {
|
|
const root = createTempDir();
|
|
const store = createJsonlSessionStore({ fs: new NodeExecutionEnv({ cwd: root }), sessionsRoot: root });
|
|
const snapshot = await store.sessions.create({ cwd: root, id: "session-1" });
|
|
const entry = {
|
|
type: "message" as const,
|
|
id: "entry-1",
|
|
parentId: null,
|
|
timestamp: "2026-01-01T00:00:00.000Z",
|
|
message: createUserMessage("one"),
|
|
};
|
|
await store.entries.append(snapshot, entry);
|
|
await expect(store.entries.append(snapshot, entry)).rejects.toThrow("Entry entry-1 already exists");
|
|
await store.sessions.delete(snapshot);
|
|
await expect(store.entries.append(snapshot, { ...entry, id: "entry-2" })).rejects.toThrow("Session not found");
|
|
expect(existsSync(snapshot.path)).toBe(false);
|
|
});
|
|
|
|
it("scopes entry uniqueness to the session path", async () => {
|
|
const root = createTempDir();
|
|
const store = createJsonlSessionStore({ fs: new NodeExecutionEnv({ cwd: root }), sessionsRoot: root });
|
|
const first = await store.sessions.create({ cwd: "/tmp/first", id: "shared-session-id" });
|
|
const second = await store.sessions.create({ cwd: "/tmp/second", id: "shared-session-id" });
|
|
const entry = {
|
|
type: "message" as const,
|
|
id: "shared-entry-id",
|
|
parentId: null,
|
|
timestamp: "2026-01-01T00:00:00.000Z",
|
|
message: createUserMessage("one"),
|
|
};
|
|
|
|
await store.entries.append(first, entry);
|
|
await expect(store.entries.append(second, entry)).resolves.toBeUndefined();
|
|
});
|
|
|
|
it("rejects non-object header metadata", async () => {
|
|
const root = createTempDir();
|
|
const env = new NodeExecutionEnv({ cwd: root });
|
|
const repo = createJsonlRepository({ fs: env, sessionsRoot: root });
|
|
const session = await repo.create({ cwd: root, id: "session-1" });
|
|
const metadata = await session.getMetadata();
|
|
const header = {
|
|
type: "session",
|
|
version: 3,
|
|
id: metadata.id,
|
|
timestamp: metadata.createdAt,
|
|
cwd: root,
|
|
metadata: "profile",
|
|
};
|
|
writeFileSync(metadata.path, `${JSON.stringify(header)}\n`);
|
|
await expect(repo.open(metadata)).rejects.toThrow("session header metadata must be an object");
|
|
});
|
|
});
|