fix(coding-agent): avoid duplicate session reads

fixes #6793
This commit is contained in:
Armin Ronacher
2026-07-19 20:18:49 +02:00
parent 3da591ab74
commit f1c587dde3
3 changed files with 156 additions and 18 deletions
+1
View File
@@ -9,6 +9,7 @@
### Fixed ### 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 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 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. - Fixed Kimi Coding sessions to show API-equivalent implied costs with the subscription indicator.
@@ -485,6 +485,16 @@ export function getDefaultSessionDir(cwd: string, agentDir: string = getDefaultA
} }
const SESSION_READ_BUFFER_SIZE = 1024 * 1024; 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 { function parseSessionEntryLine(line: string): FileEntry | null {
if (!line.trim()) return null; if (!line.trim()) return null;
@@ -541,20 +551,69 @@ export function loadEntriesFromFile(filePath: string): FileEntry[] {
return entries; 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 { function readSessionHeader(filePath: string): SessionHeader | null {
const fd = openSync(filePath, "r");
try { try {
const fd = openSync(filePath, "r"); const decoder = new StringDecoder("utf8");
const buffer = Buffer.alloc(512); const buffer = Buffer.allocUnsafe(SESSION_HEADER_READ_BUFFER_SIZE);
const bytesRead = readSync(fd, buffer, 0, 512, 0); const lineChunks: string[] = [];
closeSync(fd); let scannedBytes = 0;
const firstLine = buffer.toString("utf8", 0, bytesRead).split("\n")[0];
if (!firstLine) return null; while (scannedBytes < MAX_SESSION_HEADER_SCAN_BYTES) {
const header = JSON.parse(firstLine) as Record<string, unknown>; const readLength = Math.min(buffer.length, MAX_SESSION_HEADER_SCAN_BYTES - scannedBytes);
if (header.type !== "session" || typeof header.id !== "string") { const bytesRead = readSync(fd, buffer, 0, readLength, null);
return 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 { } 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; return null;
} }
} }
@@ -576,7 +635,7 @@ export function findMostRecentSession(sessionDir: string, cwd?: string): string
const files = readdirSync(resolvedSessionDir) const files = readdirSync(resolvedSessionDir)
.filter((f) => f.endsWith(".jsonl")) .filter((f) => f.endsWith(".jsonl"))
.map((f) => join(resolvedSessionDir, f)) .map((f) => join(resolvedSessionDir, f))
.map((path) => ({ path, header: readSessionHeader(path) })) .map((path) => ({ path, header: readSessionHeaderForDiscovery(path) }))
.filter( .filter(
(file): file is { path: string; header: SessionHeader } => (file): file is { path: string; header: SessionHeader } =>
file.header !== null && file.header !== null &&
@@ -587,6 +646,7 @@ export function findMostRecentSession(sessionDir: string, cwd?: string): string
return files[0]?.path || null; return files[0]?.path || null;
} catch { } catch {
// Directory access and stat races make recent-session discovery unavailable.
return null; return null;
} }
} }
@@ -807,6 +867,7 @@ export class SessionManager {
sessionFile: string | undefined, sessionFile: string | undefined,
persist: boolean, persist: boolean,
newSessionOptions?: NewSessionOptions, newSessionOptions?: NewSessionOptions,
preloadedFileEntries?: FileEntry[],
) { ) {
this.cwd = resolvePath(cwd); this.cwd = resolvePath(cwd);
this.sessionDir = normalizePath(sessionDir); this.sessionDir = normalizePath(sessionDir);
@@ -816,7 +877,7 @@ export class SessionManager {
} }
if (sessionFile) { if (sessionFile) {
this.setSessionFile(sessionFile); this._setSessionFile(sessionFile, preloadedFileEntries);
} else { } else {
this.newSession(newSessionOptions); this.newSession(newSessionOptions);
} }
@@ -824,9 +885,13 @@ export class SessionManager {
/** Switch to a different session file (used for resume and branching) */ /** Switch to a different session file (used for resume and branching) */
setSessionFile(sessionFile: string): void { setSessionFile(sessionFile: string): void {
this._setSessionFile(sessionFile);
}
private _setSessionFile(sessionFile: string, preloadedFileEntries?: FileEntry[]): void {
this.sessionFile = resolvePath(sessionFile); this.sessionFile = resolvePath(sessionFile);
if (existsSync(this.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 // 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. // 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 { static open(path: string, sessionDir?: string, cwdOverride?: string): SessionManager {
const resolvedPath = resolvePath(path); const resolvedPath = resolvePath(path);
// Extract cwd from session header if possible, otherwise use process.cwd() let header: SessionHeader | null = null;
const entries = loadEntriesFromFile(resolvedPath); let preloadedFileEntries: FileEntry[] | undefined;
const header = entries.find((e) => e.type === "session") as SessionHeader | undefined; if (cwdOverride === undefined && existsSync(resolvedPath)) {
const cwd = cwdOverride ?? header?.cwd ?? process.cwd(); 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 // If no sessionDir provided, derive from file's parent directory
const dir = sessionDir ? normalizePath(sessionDir) : resolve(resolvedPath, ".."); const dir = sessionDir ? normalizePath(sessionDir) : resolve(resolvedPath, "..");
return new SessionManager(cwd, dir, resolvedPath, true); return new SessionManager(cwd, dir, resolvedPath, true, undefined, preloadedFileEntries);
} }
/** /**
@@ -5,6 +5,8 @@ import { join } from "path";
import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { afterEach, beforeEach, describe, expect, it } from "vitest";
import { findMostRecentSession, loadEntriesFromFile, SessionManager } from "../../src/core/session-manager.ts"; import { findMostRecentSession, loadEntriesFromFile, SessionManager } from "../../src/core/session-manager.ts";
const HEADER_SCAN_LIMIT_BYTES = 1024 * 1024;
describe("loadEntriesFromFile", () => { describe("loadEntriesFromFile", () => {
let tempDir: string; let tempDir: string;
@@ -17,6 +19,19 @@ describe("loadEntriesFromFile", () => {
rmSync(tempDir, { recursive: true, force: true }); 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", () => { it("returns empty array for non-existent file", () => {
const entries = loadEntriesFromFile(join(tempDir, "nonexistent.jsonl")); const entries = loadEntriesFromFile(join(tempDir, "nonexistent.jsonl"));
expect(entries).toEqual([]); expect(entries).toEqual([]);
@@ -65,6 +80,43 @@ describe("loadEntriesFromFile", () => {
expect(entries).toHaveLength(2); 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", () => { it("opens session files larger than Node's max string length", () => {
const file = join(tempDir, "large.jsonl"); const file = join(tempDir, "large.jsonl");
writeFileSync( writeFileSync(
@@ -155,6 +207,15 @@ describe("findMostRecentSession", () => {
expect(findMostRecentSession(tempDir)).toBe(valid); 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 () => { it("filters most recent session by cwd", async () => {
const projectA = join(tempDir, "project-a"); const projectA = join(tempDir, "project-a");
const projectB = join(tempDir, "project-b"); const projectB = join(tempDir, "project-b");