Merge PR #8645: perf(session): stream persisted session init probes (@lemonleks)

2026-08-16 02:13:37 +02:00
parent f901273918 9c695cf98a
commit 0128fe323f
5 changed files with 189 additions and 39 deletions
@@ -2,9 +2,11 @@ import { afterEach, describe, expect, it } from "bun:test";
 import * as path from "node:path";
 import { getBundledModel } from "@oh-my-pi/pi-catalog/models";
 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 +18,15 @@ 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 };
+	}
+	override async readText(): Promise<string> {
+		throw new Error("Large sessions must stream");
+	}
+}
+
 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 +78,36 @@ describe("SessionManager.peekSessionInit", () => {
 		expect(peek?.init?.restrictToolNames).toBe(true);
 	});

+	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();
+		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, new LargeFileSessionStorage());
+		expect(peek?.cwd).toBe(manager.getCwd());
+		expect(peek?.init?.systemPrompt).toBe("second");
+	});
+
+	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 +120,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();