diff --git a/packages/coding-agent/src/session/session-loader.ts b/packages/coding-agent/src/session/session-loader.ts index 5b7c5fa08..efba21feb 100644 --- a/packages/coding-agent/src/session/session-loader.ts +++ b/packages/coding-agent/src/session/session-loader.ts @@ -46,10 +46,7 @@ function isValidSessionHeader(entry: FileEntry | undefined): entry is SessionHea return entry?.type === "session" && typeof entry.id === "string"; } -function foldTitleSlot(entries: FileEntry[], slot: SessionTitleUpdate | undefined): FileEntry[] { - if (!slot || entries.length === 0) return entries; - const header = entries[0]; - if (!isValidSessionHeader(header)) return entries; +function applyTitleSlot(header: SessionHeader, slot: SessionTitleUpdate): void { if (slot.title && slot.title.length > 0) { header.title = slot.title; } else { @@ -60,6 +57,13 @@ function foldTitleSlot(entries: FileEntry[], slot: SessionTitleUpdate | undefine } else { delete header.titleSource; } +} + +function foldTitleSlot(entries: FileEntry[], slot: SessionTitleUpdate | undefined): FileEntry[] { + if (!slot || entries.length === 0) return entries; + const header = entries[0]; + if (!isValidSessionHeader(header)) return entries; + applyTitleSlot(header, slot); return entries; } @@ -76,7 +80,7 @@ export function parseSessionContent(content: string): { /** Parse session JSONL and visit each entry without retaining prior entries. */ export async function visitEntriesFromFileStream( filePath: string, - visit: (entry: FileEntry) => void | boolean, + visit: (entry: FileEntry, titleSlot: SessionTitleUpdate | undefined) => void | boolean, options: VisitEntriesFromFileStreamOptions = {}, ): Promise { let titleSlot: SessionTitleUpdate | undefined; @@ -126,7 +130,7 @@ export async function visitEntriesFromFileStream( break; } try { - if (visit(value as FileEntry) === false) { + if (visit(value as FileEntry, titleSlot) === false) { stopped = true; break; } @@ -286,10 +290,11 @@ export async function visitEntriesFromFile( const size = storage.statSync(filePath).size; if (shouldStreamEntries(storage, size)) { let sawFirstEntry = false; - await visitEntriesFromFileStream(filePath, entry => { + await visitEntriesFromFileStream(filePath, (entry, titleSlot) => { if (!sawFirstEntry) { sawFirstEntry = true; if (!isValidSessionHeader(entry)) return false; + if (titleSlot) applyTitleSlot(entry, titleSlot); } return callVisitor(entry); }); diff --git a/packages/coding-agent/test/session-loader-stream.test.ts b/packages/coding-agent/test/session-loader-stream.test.ts index 6fb60ea9a..79a2a2480 100644 --- a/packages/coding-agent/test/session-loader-stream.test.ts +++ b/packages/coding-agent/test/session-loader-stream.test.ts @@ -4,6 +4,7 @@ import * as os from "node:os"; import * as path from "node:path"; import type { FileEntry } from "@oh-my-pi/pi-coding-agent/session/session-entries"; import * as sessionLoader from "@oh-my-pi/pi-coding-agent/session/session-loader"; +import { FileSessionStorage } from "@oh-my-pi/pi-coding-agent/session/session-storage"; import { serializeTitleSlot } from "@oh-my-pi/pi-coding-agent/session/session-title-slot"; // Parity contract for the ≥8MiB streaming loader (now Bun.JSONL-based): it must @@ -16,6 +17,14 @@ import { serializeTitleSlot } from "@oh-my-pi/pi-coding-agent/session/session-ti const ISO = "2026-06-29T12:00:00.000Z"; const HEADER = { type: "session", version: 3, id: "s1", timestamp: ISO, cwd: "/tmp" }; +const LARGE_SESSION_BYTES = 9 * 1024 * 1024; + +class LargeFileSessionStorage extends FileSessionStorage { + override statSync(filePath: string) { + return { ...super.statSync(filePath), size: LARGE_SESSION_BYTES }; + } +} + const msg = (id: string, parentId: string, text: string) => ({ type: "message", id, @@ -96,7 +105,8 @@ describe("loadEntriesFromFileStream (Bun.JSONL parity)", () => { it("visits a large journal before reading its tail", async () => { const largeText = "x".repeat(1024 * 1024); - const lines = [JSON.stringify(HEADER)]; + const slotLine = serializeTitleSlot({ title: "Visitor", source: "user", updatedAt: ISO }); + const lines = [slotLine, JSON.stringify({ ...HEADER, title: "stale", titleSource: "generated" })]; for (let index = 1; index <= 9; index++) { lines.push(JSON.stringify(msg(`m${index}`, index === 1 ? "s1" : `m${index - 1}`, largeText))); } @@ -104,14 +114,22 @@ describe("loadEntriesFromFileStream (Bun.JSONL parity)", () => { expect(fs.statSync(file).size).toBeGreaterThan(8 * 1024 * 1024); let visited = 0; - await sessionLoader.visitEntriesFromFile(file, () => { + let headerTitle: string | undefined; + let headerTitleSource: string | undefined; + await sessionLoader.visitEntriesFromFile(file, entry => { visited++; + if (entry.type === "session") { + headerTitle = entry.title; + headerTitleSource = entry.titleSource; + } if (visited === 1) fs.truncateSync(file, 0); }); // A collecting load reads the tail before the first callback and would // still visit every in-memory entry after the file is truncated. - expect(visited).toBe(1); + expect(visited).toBeLessThan(10); + expect(headerTitle).toBe("Visitor"); + expect(headerTitleSource).toBe("user"); }); it("does not revisit entries before a malformed line spanning stream chunks", async () => { @@ -173,6 +191,15 @@ describe("loadEntriesFromFileStream (Bun.JSONL parity)", () => { throw failure; }), ).rejects.toBe(failure); + await expect( + sessionLoader.visitEntriesFromFile( + file, + () => { + throw failure; + }, + new LargeFileSessionStorage(), + ), + ).rejects.toBe(failure); }); it("matches parseSessionContent on title slot + valid + malformed + blank lines", async () => { diff --git a/packages/coding-agent/test/session/peek-session-init.test.ts b/packages/coding-agent/test/session/peek-session-init.test.ts index 8446f2ede..152aab972 100644 --- a/packages/coding-agent/test/session/peek-session-init.test.ts +++ b/packages/coding-agent/test/session/peek-session-init.test.ts @@ -76,28 +76,26 @@ describe("SessionManager.peekSessionInit", () => { expect(peek?.init?.restrictToolNames).toBe(true); }); - it("returns metadata through the large-file stream path", async () => { + it("routes the latest large-file metadata through the entry visitor", async () => { const cwd = makeTempDir("@pi-peek-stream-"); const manager = SessionManager.create(cwd, path.join(cwd, "sessions")); const sessionFile = manager.getSessionFile(); if (!sessionFile) throw new Error("Expected a persisted session file path"); - manager.appendSessionInit({ systemPrompt: "streamed", task: "task", tools: ["read"], spawns: "" }); + manager.appendSessionInit({ systemPrompt: "first", task: "task", tools: ["read"], spawns: "" }); + manager.appendSessionInit({ systemPrompt: "second", task: "task", tools: ["read"], spawns: "" }); manager.appendMessage(assistantMessage("journal tail")); const storage = new LargeFileSessionStorage(); const visitEntries = spyOn(sessionLoader, "visitEntriesFromFile"); - const loadEntries = spyOn(sessionLoader, "loadEntriesFromFile"); try { const peek = await SessionManager.peekSessionInit(sessionFile, storage); expect(peek?.cwd).toBe(manager.getCwd()); - expect(peek?.init?.systemPrompt).toBe("streamed"); + expect(peek?.init?.systemPrompt).toBe("second"); expect(visitEntries).toHaveBeenCalledTimes(1); - expect(loadEntries).not.toHaveBeenCalled(); } finally { visitEntries.mockRestore(); - loadEntries.mockRestore(); } });