From 2076dc3960b4ef822824c71f89e3c259748629c2 Mon Sep 17 00:00:00 2001 From: Aleksandr Khaustov <6837579+lemonleks@users.noreply.github.com> Date: Sat, 15 Aug 2026 16:33:44 +0400 Subject: [PATCH 1/3] perf(session): stream persisted session init probes --- packages/coding-agent/CHANGELOG.md | 4 + .../src/session/session-loader.ts | 75 ++++++++++++---- .../src/session/session-manager.ts | 35 +++++--- .../test/session-loader-stream.test.ts | 32 +++++++ .../test/session/peek-session-init.test.ts | 88 ++++++++++++++++++- 5 files changed, 203 insertions(+), 31 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 5597d222b..e6daed58f 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Changed + +- Reduced peak memory for persisted subagent revival probes by streaming large file-backed session journals instead of loading the complete journal ([#8117](https://github.com/can1357/oh-my-pi/issues/8117)). + ## [17.3.4] - 2026-08-14 ### Changed diff --git a/packages/coding-agent/src/session/session-loader.ts b/packages/coding-agent/src/session/session-loader.ts index 2c2f0fab2..5b7c5fa08 100644 --- a/packages/coding-agent/src/session/session-loader.ts +++ b/packages/coding-agent/src/session/session-loader.ts @@ -42,10 +42,14 @@ function splitTitleSlot(content: string): { body: string; slot: SessionTitleUpda return { body: content.slice(newlineIndex + 1), slot }; } +function isValidSessionHeader(entry: FileEntry | undefined): entry is SessionHeader { + 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] as SessionHeader; - if (header.type !== "session" || typeof header.id !== "string") return entries; + const header = entries[0]; + if (!isValidSessionHeader(header)) return entries; if (slot.title && slot.title.length > 0) { header.title = slot.title; } else { @@ -235,32 +239,71 @@ export function parseSessionEntries(content: string): FileEntry[] { return parseSessionContent(content).entries; } +function shouldStreamEntries(storage: SessionStorage, size: number): boolean { + return storage instanceof FileSessionStorage && size >= STREAM_LOAD_THRESHOLD_BYTES; +} + +async function loadEntriesWithKnownSize(filePath: string, storage: SessionStorage, size: number): Promise { + const loaded = shouldStreamEntries(storage, size) + ? await loadEntriesFromFileStream(filePath) + : parseSessionContent(await storage.readText(filePath)); + const { entries } = loaded; + return isValidSessionHeader(entries[0]) ? entries : []; +} + /** Exported for testing */ export async function loadEntriesFromFile( filePath: string, storage: SessionStorage = new FileSessionStorage(), ): Promise { - let loaded: { entries: FileEntry[]; titleSlot: SessionTitleUpdate | undefined }; try { - const stat = storage.statSync(filePath); - loaded = - storage instanceof FileSessionStorage && stat.size >= STREAM_LOAD_THRESHOLD_BYTES - ? await loadEntriesFromFileStream(filePath) - : parseSessionContent(await storage.readText(filePath)); + return await loadEntriesWithKnownSize(filePath, storage, storage.statSync(filePath).size); } catch (err) { if (isEnoent(err)) return []; throw err; } - const { entries } = loaded; +} - // Validate session header - if (entries.length === 0) return entries; - const header = entries[0] as SessionHeader; - if (header.type !== "session" || typeof header.id !== "string") { - return []; +/** + * Visit session entries, using bounded streaming for large file-backed journals. + * Small files and non-file backends keep the existing full-load path. + */ +export async function visitEntriesFromFile( + filePath: string, + visit: (entry: FileEntry) => void | boolean, + storage: SessionStorage = new FileSessionStorage(), +): Promise { + let visitorThrew = false; + const callVisitor = (entry: FileEntry): void | boolean => { + try { + return visit(entry); + } catch (err) { + visitorThrew = true; + throw err; + } + }; + try { + const size = storage.statSync(filePath).size; + if (shouldStreamEntries(storage, size)) { + let sawFirstEntry = false; + await visitEntriesFromFileStream(filePath, entry => { + if (!sawFirstEntry) { + sawFirstEntry = true; + if (!isValidSessionHeader(entry)) return false; + } + return callVisitor(entry); + }); + return; + } + + for (const entry of await loadEntriesWithKnownSize(filePath, storage, size)) { + if (callVisitor(entry) === false) return; + } + } catch (err) { + if (visitorThrew) throw err; + if (isEnoent(err)) return; + throw err; } - - return entries; } /** diff --git a/packages/coding-agent/src/session/session-manager.ts b/packages/coding-agent/src/session/session-manager.ts index 37a50703d..47aa1f1c2 100644 --- a/packages/coding-agent/src/session/session-manager.ts +++ b/packages/coding-agent/src/session/session-manager.ts @@ -59,7 +59,12 @@ import { type UsageStatistics, } from "./session-entries"; import { findMostRecentSession, listAllSessions, listSessions, type SessionInfo } from "./session-listing"; -import { loadEntriesFromFile, readTitleSlotFromFile, resolveBlobRefsInEntries } from "./session-loader"; +import { + loadEntriesFromFile, + readTitleSlotFromFile, + resolveBlobRefsInEntries, + visitEntriesFromFile, +} from "./session-loader"; import { generateId, migrateToCurrentVersion } from "./session-migrations"; import { computeDefaultSessionDir, @@ -2638,15 +2643,7 @@ export class SessionManager { advisor?: string; } | null; } | null> { - let loaded: FileEntry[]; - try { - loaded = await loadEntriesFromFile(filePath, storage); - } catch { - return null; - } - // A missing/empty file has no usable session — nothing to revive from. - if (loaded.length === 0) return null; - const header = loaded.find(entry => entry.type === "session") as SessionHeader | undefined; + let header: SessionHeader | undefined; let init: { systemPrompt: string; task: string; @@ -2661,8 +2658,11 @@ export class SessionManager { readSummarize?: boolean; advisor?: string; } | null = null; - for (let index = loaded.length - 1; index >= 0; index--) { - const entry = loaded[index]; + const visit = (entry: FileEntry): void => { + if (entry.type === "session") { + header ??= entry; + return; + } if (entry.type === "session_init") { init = { systemPrompt: entry.systemPrompt, @@ -2678,10 +2678,17 @@ export class SessionManager { spawns: entry.spawns, advisor: entry.advisor, }; - break; } + }; + + try { + await visitEntriesFromFile(filePath, visit, storage); + } catch { + return null; } - return { cwd: header?.cwd ?? getProjectDir(), init }; + // A missing, empty, or invalid file has no usable session. + if (!header) return null; + return { cwd: header.cwd ?? getProjectDir(), init }; } /** Continue the most recent session, or create a new one if none exists. */ diff --git a/packages/coding-agent/test/session-loader-stream.test.ts b/packages/coding-agent/test/session-loader-stream.test.ts index a0900f904..6fb60ea9a 100644 --- a/packages/coding-agent/test/session-loader-stream.test.ts +++ b/packages/coding-agent/test/session-loader-stream.test.ts @@ -93,6 +93,27 @@ describe("loadEntriesFromFileStream (Bun.JSONL parity)", () => { expect(titleSlot?.title).toBe("Visitor"); expect(entryIds(visited)).toEqual(["s1", "m1", "m2"]); }); + + it("visits a large journal before reading its tail", async () => { + const largeText = "x".repeat(1024 * 1024); + const lines = [JSON.stringify(HEADER)]; + for (let index = 1; index <= 9; index++) { + lines.push(JSON.stringify(msg(`m${index}`, index === 1 ? "s1" : `m${index - 1}`, largeText))); + } + const file = await writeTemp(`${lines.join("\n")}\n`); + expect(fs.statSync(file).size).toBeGreaterThan(8 * 1024 * 1024); + + let visited = 0; + await sessionLoader.visitEntriesFromFile(file, () => { + visited++; + 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); + }); + it("does not revisit entries before a malformed line spanning stream chunks", async () => { const content = [ JSON.stringify(HEADER), @@ -143,6 +164,17 @@ describe("loadEntriesFromFileStream (Bun.JSONL parity)", () => { ).rejects.toBe(failure); }); + it("propagates ENOENT errors through the routed visitor", async () => { + const file = await writeTemp(`${JSON.stringify(HEADER)}\n`); + const failure = Object.assign(new Error("visitor failed"), { code: "ENOENT" }); + + await expect( + sessionLoader.visitEntriesFromFile(file, () => { + throw failure; + }), + ).rejects.toBe(failure); + }); + it("matches parseSessionContent on title slot + valid + malformed + blank lines", async () => { const slotLine = serializeTitleSlot({ title: "Hello world", source: "user", updatedAt: ISO }); // title slot | header | valid | blank | malformed | valid | malformed-no-newline-at-EOF 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 2f2e0e652..8446f2ede 100644 --- a/packages/coding-agent/test/session/peek-session-init.test.ts +++ b/packages/coding-agent/test/session/peek-session-init.test.ts @@ -1,10 +1,13 @@ -import { afterEach, describe, expect, it } from "bun:test"; +import { afterEach, describe, expect, it, spyOn } from "bun:test"; import * as path from "node:path"; import { getBundledModel } from "@oh-my-pi/pi-catalog/models"; +import * as sessionLoader from "@oh-my-pi/pi-coding-agent/session/session-loader"; import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; +import { FileSessionStorage, MemorySessionStorage } from "@oh-my-pi/pi-coding-agent/session/session-storage"; import { TempDir } from "@oh-my-pi/pi-utils"; const tempDirs: TempDir[] = []; +const LARGE_SESSION_BYTES = 9 * 1024 * 1024; function makeTempDir(prefix: string): string { const dir = TempDir.createSync(prefix); @@ -16,6 +19,12 @@ afterEach(async () => { await Promise.all(tempDirs.splice(0).map(dir => dir.remove())); }); +class LargeFileSessionStorage extends FileSessionStorage { + override statSync(filePath: string) { + return { ...super.statSync(filePath), size: LARGE_SESSION_BYTES }; + } +} + function assistantMessage(text: string) { const model = getBundledModel("anthropic", "claude-sonnet-4-5"); if (!model) throw new Error("Expected built-in anthropic model to exist"); @@ -67,6 +76,46 @@ describe("SessionManager.peekSessionInit", () => { expect(peek?.init?.restrictToolNames).toBe(true); }); + it("returns metadata through the large-file stream path", 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.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(visitEntries).toHaveBeenCalledTimes(1); + expect(loadEntries).not.toHaveBeenCalled(); + } finally { + visitEntries.mockRestore(); + loadEntries.mockRestore(); + } + }); + + it("preserves non-file storage behavior", async () => { + const cwd = makeTempDir("@pi-peek-memory-"); + const storage = new MemorySessionStorage(); + const manager = SessionManager.create(cwd, path.join(cwd, "sessions"), storage); + const sessionFile = manager.getSessionFile(); + if (!sessionFile) throw new Error("Expected a persisted session file path"); + manager.appendSessionInit({ systemPrompt: "first", task: "task", tools: ["read"], spawns: "" }); + manager.appendSessionInit({ systemPrompt: "second", task: "task", tools: ["read"], spawns: "" }); + manager.appendMessage(assistantMessage("journal tail")); + + const peek = await SessionManager.peekSessionInit(sessionFile, storage); + expect(peek?.cwd).toBe(manager.getCwd()); + expect(peek?.init?.systemPrompt).toBe("second"); + }); + it("returns init: null for a session file with no session_init (a main/legacy session)", async () => { const cwd = makeTempDir("@pi-peek-legacy-"); const manager = SessionManager.create(cwd, path.join(cwd, "sessions")); @@ -79,6 +128,43 @@ describe("SessionManager.peekSessionInit", () => { expect(peek?.init).toBeNull(); }); + it("returns null when the first entry is not a session header", async () => { + const file = path.join(makeTempDir("@pi-peek-invalid-header-"), "invalid.jsonl"); + const content = [ + { + type: "session_init", + id: "invalid-first", + parentId: null, + timestamp: "2026-08-15T00:00:00.000Z", + systemPrompt: "invalid", + task: "task", + tools: [], + }, + { + type: "session", + version: 3, + id: "late-header", + timestamp: "2026-08-15T00:00:00.000Z", + cwd: "/wrong", + }, + { + type: "session_init", + id: "late-init", + parentId: "late-header", + timestamp: "2026-08-15T00:00:00.000Z", + systemPrompt: "late", + task: "task", + tools: [], + }, + ] + .map(entry => JSON.stringify(entry)) + .join("\n"); + await Bun.write(file, `${content}\n`); + + expect(await SessionManager.peekSessionInit(file)).toBeNull(); + expect(await SessionManager.peekSessionInit(file, new LargeFileSessionStorage())).toBeNull(); + }); + it("returns null for a file that cannot be read", async () => { const peek = await SessionManager.peekSessionInit(path.join(makeTempDir("@pi-peek-missing-"), "nope.jsonl")); expect(peek).toBeNull(); From cbbaceac36b522d3455eaf25cec9bc7d1d6672a7 Mon Sep 17 00:00:00 2001 From: Aleksandr Khaustov <6837579+lemonleks@users.noreply.github.com> Date: Sat, 15 Aug 2026 16:57:08 +0400 Subject: [PATCH 2/3] fix(session): preserve title slots in entry visitor --- .../src/session/session-loader.ts | 19 +++++++---- .../test/session-loader-stream.test.ts | 33 +++++++++++++++++-- .../test/session/peek-session-init.test.ts | 10 +++--- 3 files changed, 46 insertions(+), 16 deletions(-) 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(); } }); From 9c695cf98a160da08136130826b89f4b18f18664 Mon Sep 17 00:00:00 2001 From: Aleksandr Khaustov <6837579+lemonleks@users.noreply.github.com> Date: Sat, 15 Aug 2026 23:54:30 +0400 Subject: [PATCH 3/3] refactor(session): simplify streamed entry visits --- .../src/session/session-loader.ts | 76 ++++++++----------- .../test/session-loader-stream.test.ts | 29 +------ .../test/session/peek-session-init.test.ts | 22 ++---- 3 files changed, 39 insertions(+), 88 deletions(-) diff --git a/packages/coding-agent/src/session/session-loader.ts b/packages/coding-agent/src/session/session-loader.ts index efba21feb..168913bbc 100644 --- a/packages/coding-agent/src/session/session-loader.ts +++ b/packages/coding-agent/src/session/session-loader.ts @@ -46,27 +46,20 @@ function isValidSessionHeader(entry: FileEntry | undefined): entry is SessionHea return entry?.type === "session" && typeof entry.id === "string"; } -function applyTitleSlot(header: SessionHeader, slot: SessionTitleUpdate): void { +function applyTitleSlot(entry: FileEntry | undefined, slot: SessionTitleUpdate | undefined): void { + if (!slot || !isValidSessionHeader(entry)) return; if (slot.title && slot.title.length > 0) { - header.title = slot.title; + entry.title = slot.title; } else { - delete header.title; + delete entry.title; } if (slot.source) { - header.titleSource = slot.source; + entry.titleSource = slot.source; } else { - delete header.titleSource; + delete entry.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; -} - /** Parse session JSONL while stripping and folding the optional fixed title slot. */ export function parseSessionContent(content: string): { entries: FileEntry[]; @@ -74,17 +67,19 @@ export function parseSessionContent(content: string): { } { const { body, slot } = splitTitleSlot(content); const entries = parseJsonlLenient(body) as FileEntry[]; - return { entries: foldTitleSlot(entries, slot), titleSlot: slot }; + applyTitleSlot(entries[0], slot); + return { entries, titleSlot: slot }; } /** Parse session JSONL and visit each entry without retaining prior entries. */ export async function visitEntriesFromFileStream( filePath: string, - visit: (entry: FileEntry, titleSlot: SessionTitleUpdate | undefined) => void | boolean, + visit: (entry: FileEntry) => void | boolean, options: VisitEntriesFromFileStreamOptions = {}, ): Promise { let titleSlot: SessionTitleUpdate | undefined; let sawFirstLine = false; + let sawFirstEntry = false; let bytesSinceYield = 0; let entriesSinceYield = 0; let recordsSeen = 0; @@ -129,8 +124,13 @@ export async function visitEntriesFromFileStream( stopped = true; break; } + const entry = value as FileEntry; + if (!sawFirstEntry) { + sawFirstEntry = true; + applyTitleSlot(entry, titleSlot); + } try { - if (visit(value as FileEntry, titleSlot) === false) { + if (visit(entry) === false) { stopped = true; break; } @@ -219,7 +219,7 @@ export async function loadEntriesFromFileStream(filePath: string): Promise<{ const titleSlot = await visitEntriesFromFileStream(filePath, entry => { entries.push(entry); }); - return { entries: foldTitleSlot(entries, titleSlot), titleSlot }; + return { entries, titleSlot }; } /** Read only the fixed-size head window to detect a physical title slot. */ @@ -277,37 +277,21 @@ export async function visitEntriesFromFile( visit: (entry: FileEntry) => void | boolean, storage: SessionStorage = new FileSessionStorage(), ): Promise { - let visitorThrew = false; - const callVisitor = (entry: FileEntry): void | boolean => { - try { + const size = storage.statSync(filePath).size; + if (shouldStreamEntries(storage, size)) { + let sawFirstEntry = false; + await visitEntriesFromFileStream(filePath, entry => { + if (!sawFirstEntry) { + sawFirstEntry = true; + if (!isValidSessionHeader(entry)) return false; + } return visit(entry); - } catch (err) { - visitorThrew = true; - throw err; - } - }; - try { - const size = storage.statSync(filePath).size; - if (shouldStreamEntries(storage, size)) { - let sawFirstEntry = false; - await visitEntriesFromFileStream(filePath, (entry, titleSlot) => { - if (!sawFirstEntry) { - sawFirstEntry = true; - if (!isValidSessionHeader(entry)) return false; - if (titleSlot) applyTitleSlot(entry, titleSlot); - } - return callVisitor(entry); - }); - return; - } + }); + return; + } - for (const entry of await loadEntriesWithKnownSize(filePath, storage, size)) { - if (callVisitor(entry) === false) return; - } - } catch (err) { - if (visitorThrew) throw err; - if (isEnoent(err)) return; - throw err; + for (const entry of await loadEntriesWithKnownSize(filePath, storage, size)) { + if (visit(entry) === false) return; } } diff --git a/packages/coding-agent/test/session-loader-stream.test.ts b/packages/coding-agent/test/session-loader-stream.test.ts index 79a2a2480..9878d55b8 100644 --- a/packages/coding-agent/test/session-loader-stream.test.ts +++ b/packages/coding-agent/test/session-loader-stream.test.ts @@ -4,7 +4,6 @@ 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 @@ -17,13 +16,6 @@ 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", @@ -101,6 +93,7 @@ describe("loadEntriesFromFileStream (Bun.JSONL parity)", () => { expect(titleSlot?.title).toBe("Visitor"); expect(entryIds(visited)).toEqual(["s1", "m1", "m2"]); + expect(visited[0]).toMatchObject({ title: "Visitor", titleSource: "user" }); }); it("visits a large journal before reading its tail", async () => { @@ -182,26 +175,6 @@ describe("loadEntriesFromFileStream (Bun.JSONL parity)", () => { ).rejects.toBe(failure); }); - it("propagates ENOENT errors through the routed visitor", async () => { - const file = await writeTemp(`${JSON.stringify(HEADER)}\n`); - const failure = Object.assign(new Error("visitor failed"), { code: "ENOENT" }); - - await expect( - sessionLoader.visitEntriesFromFile(file, () => { - 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 () => { const slotLine = serializeTitleSlot({ title: "Hello world", source: "user", updatedAt: ISO }); // title slot | header | valid | blank | malformed | valid | malformed-no-newline-at-EOF 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 152aab972..119457741 100644 --- a/packages/coding-agent/test/session/peek-session-init.test.ts +++ b/packages/coding-agent/test/session/peek-session-init.test.ts @@ -1,7 +1,6 @@ -import { afterEach, describe, expect, it, spyOn } from "bun:test"; +import { afterEach, describe, expect, it } from "bun:test"; import * as path from "node:path"; import { getBundledModel } from "@oh-my-pi/pi-catalog/models"; -import * as sessionLoader from "@oh-my-pi/pi-coding-agent/session/session-loader"; import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; import { FileSessionStorage, MemorySessionStorage } from "@oh-my-pi/pi-coding-agent/session/session-storage"; import { TempDir } from "@oh-my-pi/pi-utils"; @@ -23,6 +22,9 @@ class LargeFileSessionStorage extends FileSessionStorage { override statSync(filePath: string) { return { ...super.statSync(filePath), size: LARGE_SESSION_BYTES }; } + override async readText(): Promise { + throw new Error("Large sessions must stream"); + } } function assistantMessage(text: string) { @@ -76,7 +78,7 @@ describe("SessionManager.peekSessionInit", () => { expect(peek?.init?.restrictToolNames).toBe(true); }); - it("routes the latest large-file metadata through the entry visitor", async () => { + it("streams large file-backed sessions without a full read", async () => { const cwd = makeTempDir("@pi-peek-stream-"); const manager = SessionManager.create(cwd, path.join(cwd, "sessions")); const sessionFile = manager.getSessionFile(); @@ -86,17 +88,9 @@ describe("SessionManager.peekSessionInit", () => { manager.appendSessionInit({ systemPrompt: "second", task: "task", tools: ["read"], spawns: "" }); manager.appendMessage(assistantMessage("journal tail")); - const storage = new LargeFileSessionStorage(); - const visitEntries = spyOn(sessionLoader, "visitEntriesFromFile"); - try { - const peek = await SessionManager.peekSessionInit(sessionFile, storage); - - expect(peek?.cwd).toBe(manager.getCwd()); - expect(peek?.init?.systemPrompt).toBe("second"); - expect(visitEntries).toHaveBeenCalledTimes(1); - } finally { - visitEntries.mockRestore(); - } + const peek = await SessionManager.peekSessionInit(sessionFile, new LargeFileSessionStorage()); + expect(peek?.cwd).toBe(manager.getCwd()); + expect(peek?.init?.systemPrompt).toBe("second"); }); it("preserves non-file storage behavior", async () => {