Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
159 changes: 159 additions & 0 deletions src/session/optimized-context-store.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -960,3 +960,162 @@ describe("createOptimizedContextStore unpublished rewrite", () => {
expect(turnTexts((await store.load()).turns)).toEqual(["era-2"]);
});
});

describe("createOptimizedContextStore prompt dedupe (CL-9026)", () => {
const PROMPT_FILE = "prompt.jsonl";

function toolResultTurn(body: string): ConversationTurn {
return {
role: "user",
content: [
{
type: "tool_result",
callId: "call-1",
content: [{ type: "text", text: body }],
},
],
timestamp: 1,
};
}

function cloneTurns(turns: ConversationTurn[]): ConversationTurn[] {
return turns.map((t) => JSON.parse(JSON.stringify(t)) as ConversationTurn);
}

test("identical writePrompt writes no prompt segment", async () => {
const dir = tempDir();
const store = await createOptimizedContextStore(dir);
const turns = [turn("a"), turn("b")];
await store.writeTurns([...turns]);
await store.writePrompt([...turns]);
await store.writePrompt(cloneTurns(turns));

expect(fs.existsSync(path.join(dir, PROMPT_FILE))).toBe(false);
expect(await listSegmentFiles(dir, PROMPT_FILE)).toEqual([]);

await store.writeMetadata(EMPTY_CHECKPOINT_METADATA);
await store.commit({ message: "identical prompt is a no-op" });
expect(await gitLsTree(dir)).not.toContain(PROMPT_FILE);

// A fresh instance compares against on-disk live turns and also skips.
const fresh = await createOptimizedContextStore(dir);
await fresh.writePrompt(cloneTurns(turns));
expect(fs.existsSync(path.join(dir, PROMPT_FILE))).toBe(false);
});

test("differing writePrompt still writes the prompt snapshot", async () => {
const dir = tempDir();
const store = await createOptimizedContextStore(dir);
const live = [turn("a"), toolResultTurn("x".repeat(1000))];
await store.writeTurns([...live]);

const shrunk = [turn("a"), toolResultTurn("summarized")];
await store.writePrompt(shrunk);
expect(fs.existsSync(path.join(dir, PROMPT_FILE))).toBe(true);
expect(fs.readFileSync(path.join(dir, PROMPT_FILE), "utf-8")).toBe(
jsonl(shrunk),
);

// An appended ephemeral turn also changes the prompt and still writes.
const ephemeralDir = tempDir();
const ephemeralStore = await createOptimizedContextStore(ephemeralDir);
await ephemeralStore.writeTurns([...live]);
const withEphemeral = [...cloneTurns(live), turn("ephemeral")];
await ephemeralStore.writePrompt(withEphemeral);
expect(fs.readFileSync(path.join(ephemeralDir, PROMPT_FILE), "utf-8")).toBe(
jsonl(withEphemeral),
);
});

test("turns-only store loads live turns and commits no prompt file", async () => {
const dir = tempDir();
const store = await createOptimizedContextStore(dir);
await store.writeTurns([turn("a"), turn("b")]);
await store.writeMetadata(EMPTY_CHECKPOINT_METADATA);
const first = await store.commit({ message: "turns only" });

expect(turnTexts((await store.load()).turns)).toEqual(["a", "b"]);

const fresh = await createOptimizedContextStore(dir);
expect(turnTexts((await fresh.load()).turns)).toEqual(["a", "b"]);
const second = await fresh.commit({ message: "turns-only no-op" });
expect(second.hash).toBe(first.hash);
expect(fs.existsSync(path.join(dir, PROMPT_FILE))).toBe(false);
expect(await gitLsTree(dir)).not.toContain(PROMPT_FILE);
});

test("identical writePrompt after a differing one removes stale prompt segments", async () => {
const dir = tempDir();
const store = await createOptimizedContextStore(dir);
const live = [turn("a"), turn("b")];
await store.writeTurns([...live]);
await store.writePrompt([turn("a")]);
expect(fs.existsSync(path.join(dir, PROMPT_FILE))).toBe(true);

await store.writePrompt(cloneTurns(live));
expect(fs.existsSync(path.join(dir, PROMPT_FILE))).toBe(false);
expect(await listSegmentFiles(dir, PROMPT_FILE)).toEqual([]);

await store.writeMetadata(EMPTY_CHECKPOINT_METADATA);
await store.commit({ message: "stale prompt removed" });
expect(await gitLsTree(dir)).not.toContain(PROMPT_FILE);
expect(turnTexts((await store.load()).turns)).toEqual(["a", "b"]);
});

test("committed then identical writePrompt drops prompt.jsonl from HEAD", async () => {
const dir = tempDir();
const store = await createOptimizedContextStore(dir);
const live = [turn("a"), turn("b")];
await store.writeTurns([...live]);
await store.writePrompt([turn("a")]);
await store.writeMetadata(EMPTY_CHECKPOINT_METADATA);
await store.commit({ message: "differing prompt" });
expect(await gitLsTree(dir)).toContain(PROMPT_FILE);

await store.writePrompt(cloneTurns(live));
await store.writeMetadata(EMPTY_CHECKPOINT_METADATA);
const afterIdentical = await store.commit({ message: "identical prompt" });
expect(await gitLsTree(dir)).not.toContain(PROMPT_FILE);
expect(turnTexts((await store.load()).turns)).toEqual(["a", "b"]);

await store.writeMetadata(EMPTY_CHECKPOINT_METADATA);
const empty = await store.commit({ message: "empty checkpoint" });
expect(empty.hash).toBe(afterIdentical.hash);
});

test("identical writePrompt after multi-segment differing prompt drops all prompt files from HEAD", async () => {
const dir = tempDir();
const store = await createOptimizedContextStore(dir);
const live: ConversationTurn[] = [];
const big = "x".repeat(20_000);
for (let i = 0; i < 18; i++) {
live.push(turn(`${i}-${big}`));
}
await store.writeTurns([...live]);

const differing: ConversationTurn[] = [];
for (let i = 0; i < 18; i++) {
differing.push(turn(`p-${i}-${big}`));
}
await store.writePrompt(differing);
const extraPrompt = segmentFileName(PROMPT_FILE, 1);
expect(fs.existsSync(path.join(dir, PROMPT_FILE))).toBe(true);
expect(fs.existsSync(path.join(dir, extraPrompt))).toBe(true);

await store.writeMetadata(EMPTY_CHECKPOINT_METADATA);
await store.commit({ message: "multi-segment prompt" });
const treeAfterDiffering = await gitLsTree(dir);
expect(treeAfterDiffering).toContain(PROMPT_FILE);
expect(treeAfterDiffering).toContain(extraPrompt);

await store.writePrompt(cloneTurns(live));
await store.writeMetadata(EMPTY_CHECKPOINT_METADATA);
await store.commit({ message: "identical drops prompt segments" });
const tree = await gitLsTree(dir);
expect(tree).not.toContain(PROMPT_FILE);
expect(tree).not.toContain(extraPrompt);
expect(turnTexts((await store.load()).turns)).toEqual(
live.map((t) => (t.content[0] as { text: string }).text),
);
});
});
82 changes: 77 additions & 5 deletions src/session/optimized-context-store.ts
Original file line number Diff line number Diff line change
Expand Up @@ -527,7 +527,7 @@ export async function createSessionStores(
const pendingBlobFilepaths = new Set<string>();
const pendingSegmentPaths = new Set<string>();
let writeTurnsSegmented = createSegmentedJSONLWriter(dir, TURNS_FILE);
const writePromptSegmented = createSegmentedJSONLWriter(dir, PROMPT_FILE);
let writePromptSegmented = createSegmentedJSONLWriter(dir, PROMPT_FILE);
let liveTurnRefs: readonly ConversationTurn[] | null = null;
let unpublishedRewrite: ConversationTurn[] | null = null;

Expand Down Expand Up @@ -606,6 +606,75 @@ export async function createSessionStores(
liveTurnRefs = [...turns];
}

function promptEqualsLiveTurns(
live: readonly ConversationTurn[],
prompt: readonly ConversationTurn[],
): boolean {
if (live.length !== prompt.length) return false;
for (let index = 0; index < prompt.length; index++) {
if (live[index] === prompt[index]) continue;
if (JSON.stringify(live[index]) !== JSON.stringify(prompt[index]))
return false;
}
return true;
}

// The reactor checkpoints the materialized prompt every cycle, but most
// cycles run no transform that changes it — writing an identical snapshot
// next to turns.jsonl doubles disk and re-hash cost for zero information.
// load() never reads prompt.jsonl (base turns + TURNS_FILE extras only),
// so skipping the write leaves resume behavior unchanged; prompts that
// actually differ still write exactly as before.
async function writePromptIfDiffered(
turns: readonly ConversationTurn[],
): Promise<void> {
let live: readonly ConversationTurn[] | null = null;
if (unpublishedRewrite !== null) live = unpublishedRewrite;
else if (liveTurnRefs !== null) live = liveTurnRefs;
else {
// Fresh instance with no writeTurns yet: recover live turns from disk
// the same way writeTurnsLiveOrStage does. Any failure falls through
// to a normal write — an unreadable baseline must not drop the snapshot.
try {
const extraTexts = await readExtraSegmentTexts(dir, TURNS_FILE);
let baseTurns: ConversationTurn[];
try {
baseTurns = (await base.load()).turns;
} catch {
baseTurns = await readBaseTurnsFromDisk(dir);
}
live =
extraTexts.length === 0
? baseTurns
: await loadTurnsWithoutMalformedToolSequence(
baseTurns,
extraTexts,
);
} catch {
live = null;
}
}
if (live === null || !promptEqualsLiveTurns(live, turns)) {
await writeSegmented(writePromptSegmented, turns);
return;
}
// Identical to live turns: converge disk to no prompt segment so a stale
// snapshot from an earlier differing write cannot linger. Removals join
// pendingSegmentPaths so commit stages them out of the tree.
const highest = await highestSegmentIndex(dir, PROMPT_FILE);
let removed = false;
for (let index = 0; index <= highest; index++) {
const name = segmentFileName(PROMPT_FILE, index);
if (!(await pathExists(path.join(dir, name)))) continue;
await fs.promises.unlink(path.join(dir, name));
pendingSegmentPaths.add(name);
removed = true;
}
if (removed) {
writePromptSegmented = createSegmentedJSONLWriter(dir, PROMPT_FILE);
}
}

// Prefer the longest prefix of base + extras whose tool sequence the reactor
// will accept. Orphan tails left by a fresh-writer compaction rewrite are
// dropped and unlinked so the next load does not re-poison the session.
Expand Down Expand Up @@ -718,7 +787,9 @@ export async function createSessionStores(
return [...baseTurns, ...parsedExtras.slice(0, keepExtras).flat()];
},
readBlob: (key, signal) => base.readBlob(key, signal),
writePrompt: (turns) => writeSegmented(writePromptSegmented, turns),
// Skipped identical snapshots fall back to live turns in load(), which
// never reads prompt.jsonl.
writePrompt: (turns) => writePromptIfDiffered(turns),
writeResponse: (turn, signal) => base.writeResponse(turn, signal),
writeManifest: (records, signal) => base.writeManifest(records, signal),
writeTurns: (turns) => writeTurnsLiveOrStage(turns),
Expand Down Expand Up @@ -793,9 +864,10 @@ export async function createSessionStores(
}

const add = extraCommitPaths([...new Set(toAdd)]);
const remove = extraCommitPaths([...new Set(toRemove)]).filter(
(p) => !add.includes(p),
);
// extraCommitPaths strips vendor roots because base.commit() git.adds
// those that still exist. It does not git.remove missing ones, so an
// unlinked prompt.jsonl must stay in `remove`.
const remove = [...new Set(toRemove)].filter((p) => !add.includes(p));
extraPaths = [...new Set([...add, ...remove])];

for (const filepath of add) {
Expand Down
Loading