diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index e3ea4701..15bdbf80 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -9,6 +9,7 @@ ### Fixed +- Fixed persisted sessions being read and parsed twice when opened, reducing startup latency for large sessions ([#6793](https://github.com/earendil-works/pi/issues/6793)). - Fixed prompt-template defaults for all arguments (`${@:-default}` and `${ARGUMENTS:-default}`) ([#6695](https://github.com/earendil-works/pi/issues/6695)). - Fixed obsolete custom UI, custom tool, and custom editor examples in the extension documentation ([#6735](https://github.com/earendil-works/pi/issues/6735)). - Fixed Kimi Coding sessions to show API-equivalent implied costs with the subscription indicator. diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index eadca30b..f5b75c2c 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -485,6 +485,16 @@ export function getDefaultSessionDir(cwd: string, agentDir: string = getDefaultA } const SESSION_READ_BUFFER_SIZE = 1024 * 1024; +const SESSION_HEADER_READ_BUFFER_SIZE = 4096; +/** Bound synchronous header discovery while allowing large cwd and custom metadata fields. */ +const MAX_SESSION_HEADER_SCAN_BYTES = 1024 * 1024; + +class SessionHeaderScanLimitError extends Error { + constructor(filePath: string) { + super(`Session header exceeds ${MAX_SESSION_HEADER_SCAN_BYTES}-byte scan limit: ${filePath}`); + this.name = "SessionHeaderScanLimitError"; + } +} function parseSessionEntryLine(line: string): FileEntry | null { if (!line.trim()) return null; @@ -541,20 +551,69 @@ export function loadEntriesFromFile(filePath: string): FileEntry[] { return entries; } +/** + * Inspect a physical line while searching for the first parsed session entry. + * Blank and malformed lines are skipped to match loadEntriesFromFile(). + * Returns undefined to keep scanning, null for a parsed non-header entry, or the header. + */ +function parseSessionHeaderCandidate(line: string): SessionHeader | null | undefined { + if (!line.trim()) return undefined; + const entry = parseSessionEntryLine(line); + if (!entry) return undefined; + if (entry.type !== "session" || typeof (entry as { id?: unknown }).id !== "string") return null; + return entry; +} + function readSessionHeader(filePath: string): SessionHeader | null { + const fd = openSync(filePath, "r"); try { - const fd = openSync(filePath, "r"); - const buffer = Buffer.alloc(512); - const bytesRead = readSync(fd, buffer, 0, 512, 0); - closeSync(fd); - const firstLine = buffer.toString("utf8", 0, bytesRead).split("\n")[0]; - if (!firstLine) return null; - const header = JSON.parse(firstLine) as Record; - if (header.type !== "session" || typeof header.id !== "string") { - return null; + const decoder = new StringDecoder("utf8"); + const buffer = Buffer.allocUnsafe(SESSION_HEADER_READ_BUFFER_SIZE); + const lineChunks: string[] = []; + let scannedBytes = 0; + + while (scannedBytes < MAX_SESSION_HEADER_SCAN_BYTES) { + const readLength = Math.min(buffer.length, MAX_SESSION_HEADER_SCAN_BYTES - scannedBytes); + const bytesRead = readSync(fd, buffer, 0, readLength, null); + if (bytesRead === 0) { + lineChunks.push(decoder.end()); + return parseSessionHeaderCandidate(lineChunks.join("")) ?? null; + } + scannedBytes += bytesRead; + + const chunk = decoder.write(buffer.subarray(0, bytesRead)); + let lineStart = 0; + let newlineIndex = chunk.indexOf("\n", lineStart); + while (newlineIndex !== -1) { + lineChunks.push(chunk.slice(lineStart, newlineIndex)); + const header = parseSessionHeaderCandidate(lineChunks.join("")); + if (header !== undefined) return header; + lineChunks.length = 0; + lineStart = newlineIndex + 1; + newlineIndex = chunk.indexOf("\n", lineStart); + } + lineChunks.push(chunk.slice(lineStart)); } - return header as unknown as SessionHeader; + + // Probe for EOF so a final header without a newline is allowed when it ends + // exactly at the scan limit. Any additional byte exceeds the bounded scan. + const probe = Buffer.allocUnsafe(1); + if (readSync(fd, probe, 0, probe.length, null) === 0) { + lineChunks.push(decoder.end()); + return parseSessionHeaderCandidate(lineChunks.join("")) ?? null; + } + throw new SessionHeaderScanLimitError(filePath); + } finally { + closeSync(fd); + } +} + +function readSessionHeaderForDiscovery(filePath: string): SessionHeader | null { + try { + return readSessionHeader(filePath); } catch { + // Discovery is best-effort: unreadable or oversized files are not sessions, + // and one corrupt file must not prevent other sessions from being found. return null; } } @@ -576,7 +635,7 @@ export function findMostRecentSession(sessionDir: string, cwd?: string): string const files = readdirSync(resolvedSessionDir) .filter((f) => f.endsWith(".jsonl")) .map((f) => join(resolvedSessionDir, f)) - .map((path) => ({ path, header: readSessionHeader(path) })) + .map((path) => ({ path, header: readSessionHeaderForDiscovery(path) })) .filter( (file): file is { path: string; header: SessionHeader } => file.header !== null && @@ -587,6 +646,7 @@ export function findMostRecentSession(sessionDir: string, cwd?: string): string return files[0]?.path || null; } catch { + // Directory access and stat races make recent-session discovery unavailable. return null; } } @@ -807,6 +867,7 @@ export class SessionManager { sessionFile: string | undefined, persist: boolean, newSessionOptions?: NewSessionOptions, + preloadedFileEntries?: FileEntry[], ) { this.cwd = resolvePath(cwd); this.sessionDir = normalizePath(sessionDir); @@ -816,7 +877,7 @@ export class SessionManager { } if (sessionFile) { - this.setSessionFile(sessionFile); + this._setSessionFile(sessionFile, preloadedFileEntries); } else { this.newSession(newSessionOptions); } @@ -824,9 +885,13 @@ export class SessionManager { /** Switch to a different session file (used for resume and branching) */ setSessionFile(sessionFile: string): void { + this._setSessionFile(sessionFile); + } + + private _setSessionFile(sessionFile: string, preloadedFileEntries?: FileEntry[]): void { this.sessionFile = resolvePath(sessionFile); if (existsSync(this.sessionFile)) { - this.fileEntries = loadEntriesFromFile(this.sessionFile); + this.fileEntries = preloadedFileEntries ?? loadEntriesFromFile(this.sessionFile); // If file was empty, initialize it with a valid session header. If it was // non-empty but did not parse as a pi session, fail without modifying it. @@ -1451,13 +1516,24 @@ export class SessionManager { */ static open(path: string, sessionDir?: string, cwdOverride?: string): SessionManager { const resolvedPath = resolvePath(path); - // Extract cwd from session header if possible, otherwise use process.cwd() - const entries = loadEntriesFromFile(resolvedPath); - const header = entries.find((e) => e.type === "session") as SessionHeader | undefined; - const cwd = cwdOverride ?? header?.cwd ?? process.cwd(); + let header: SessionHeader | null = null; + let preloadedFileEntries: FileEntry[] | undefined; + if (cwdOverride === undefined && existsSync(resolvedPath)) { + try { + header = readSessionHeader(resolvedPath); + } catch (error) { + if (!(error instanceof SessionHeaderScanLimitError)) throw error; + // The bounded scan is only a discovery optimization. A full load remains + // authoritative for legacy files with very large headers or prefixes. + preloadedFileEntries = loadEntriesFromFile(resolvedPath); + const firstEntry = preloadedFileEntries[0]; + header = firstEntry?.type === "session" ? firstEntry : null; + } + } + const cwd = cwdOverride ?? (header ? getSessionHeaderCwd(header) : undefined) ?? process.cwd(); // If no sessionDir provided, derive from file's parent directory const dir = sessionDir ? normalizePath(sessionDir) : resolve(resolvedPath, ".."); - return new SessionManager(cwd, dir, resolvedPath, true); + return new SessionManager(cwd, dir, resolvedPath, true, undefined, preloadedFileEntries); } /** diff --git a/packages/coding-agent/test/session-manager/file-operations.test.ts b/packages/coding-agent/test/session-manager/file-operations.test.ts index d4ec0d2f..c978b320 100644 --- a/packages/coding-agent/test/session-manager/file-operations.test.ts +++ b/packages/coding-agent/test/session-manager/file-operations.test.ts @@ -5,6 +5,8 @@ import { join } from "path"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { findMostRecentSession, loadEntriesFromFile, SessionManager } from "../../src/core/session-manager.ts"; +const HEADER_SCAN_LIMIT_BYTES = 1024 * 1024; + describe("loadEntriesFromFile", () => { let tempDir: string; @@ -17,6 +19,19 @@ describe("loadEntriesFromFile", () => { rmSync(tempDir, { recursive: true, force: true }); }); + function writeSessionHeader(file: string, cwd: string, id: string, prefix = ""): void { + writeFileSync( + file, + `${prefix}${JSON.stringify({ + type: "session", + version: 3, + id, + timestamp: "2025-01-01T00:00:00Z", + cwd, + })}\n`, + ); + } + it("returns empty array for non-existent file", () => { const entries = loadEntriesFromFile(join(tempDir, "nonexistent.jsonl")); expect(entries).toEqual([]); @@ -65,6 +80,43 @@ describe("loadEntriesFromFile", () => { expect(entries).toHaveLength(2); }); + it.each([ + ["leading blank lines", "\n \n", "leading-blank"], + ["leading malformed lines", "not json\n{broken json\n", "leading-malformed"], + ["a multi-buffer header", "", "a".repeat(8192)], + ])("reads cwd from a session with %s", (_description, prefix, sessionId) => { + const file = join(tempDir, "header.jsonl"); + const storedCwd = join(tempDir, "stored-project"); + writeSessionHeader(file, storedCwd, sessionId, prefix); + + const sessionManager = SessionManager.open(file, tempDir); + expect(sessionManager.getSessionId()).toBe(sessionId); + expect(sessionManager.getCwd()).toBe(storedCwd); + }); + + it("opens compatible sessions beyond the discovery scan limit", () => { + const storedCwd = join(tempDir, "stored-project"); + const overrideCwd = join(tempDir, "override-project"); + const cases = [ + { name: "large-header", id: "a".repeat(HEADER_SCAN_LIMIT_BYTES + 1), prefix: "" }, + { + name: "large-prefix", + id: "large-prefix", + prefix: `${"x".repeat(HEADER_SCAN_LIMIT_BYTES + 1)}\n`, + }, + ]; + + for (const { name, id, prefix } of cases) { + const file = join(tempDir, `${name}.jsonl`); + writeSessionHeader(file, storedCwd, id, prefix); + for (const cwdOverride of [undefined, overrideCwd]) { + const sessionManager = SessionManager.open(file, tempDir, cwdOverride); + expect(sessionManager.getSessionId()).toBe(id); + expect(sessionManager.getCwd()).toBe(cwdOverride ?? storedCwd); + } + } + }); + it("opens session files larger than Node's max string length", () => { const file = join(tempDir, "large.jsonl"); writeFileSync( @@ -155,6 +207,15 @@ describe("findMostRecentSession", () => { expect(findMostRecentSession(tempDir)).toBe(valid); }); + it("skips oversized corrupt files and returns a valid session", () => { + const invalid = join(tempDir, "oversized.jsonl"); + const valid = join(tempDir, "valid.jsonl"); + writeFileSync(invalid, "x".repeat(HEADER_SCAN_LIMIT_BYTES + 1)); + writeFileSync(valid, '{"type":"session","id":"abc","timestamp":"2025-01-01T00:00:00Z","cwd":"/tmp"}\n'); + + expect(findMostRecentSession(tempDir)).toBe(valid); + }); + it("filters most recent session by cwd", async () => { const projectA = join(tempDir, "project-a"); const projectB = join(tempDir, "project-b");