fix(session): preserve title slots in entry visitor

This commit is contained in:
Aleksandr Khaustov
2026-08-15 16:57:08 +04:00
parent 2076dc3960
commit cbbaceac36
3 changed files with 46 additions and 16 deletions
@@ -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<SessionTitleUpdate | undefined> {
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);
});
@@ -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 () => {
@@ -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();
}
});