fix: harden subsession recovery validation

This commit is contained in:
Federico Jaramillo Martinez
2026-06-25 00:30:16 +02:00
parent a99696bd09
commit 417b04a23f
2 changed files with 151 additions and 9 deletions
+118 -1
View File
@@ -1136,7 +1136,9 @@ describe("PiSessionService", () => {
getHeader: () => ({ parentSession: parentFile }), getHeader: () => ({ parentSession: parentFile }),
getEntries: () => [{ type: "custom", customType: "pi-web.subsession.spawned", data: { version: 1, spawnedBySessionId: "parent-1", spawnedSessionId: "child-1" } }], getEntries: () => [{ type: "custom", customType: "pi-web.subsession.spawned", data: { version: 1, spawnedBySessionId: "parent-1", spawnedSessionId: "child-1" } }],
}); });
const parentManager = fakeSessionManager("/workspace"); const parentManager = fakeSessionManager("/workspace", {
getEntries: () => [{ type: "custom", customType: "pi-web.subsession.link", data: { version: 1, spawnedBySessionId: "parent-1", spawnedSessionId: "child-1", spawnedSessionFile: childFile, cwd: "/workspace-feature" } }],
});
const child = fakeRuntime("child-1", { sessionFile: childFile, sessionManager: childManager }); const child = fakeRuntime("child-1", { sessionFile: childFile, sessionManager: childManager });
const parent = fakeRuntime("parent-1", { sessionFile: parentFile, sessionManager: parentManager }); const parent = fakeRuntime("parent-1", { sessionFile: parentFile, sessionManager: parentManager });
const runtimes = [child.runtime, parent.runtime]; const runtimes = [child.runtime, parent.runtime];
@@ -1174,6 +1176,121 @@ describe("PiSessionService", () => {
} }
}); });
it("notifies the validated parent file instead of an active prefix-matched parent id", async () => {
const tempDir = await mkdtemp(join(tmpdir(), "pi-web-subsession-prefix-parent-"));
const parentFile = join(tempDir, "parent.jsonl");
const forkParentFile = join(tempDir, "parent-fork.jsonl");
const childFile = join(tempDir, "child.jsonl");
await writeFile(parentFile, `${JSON.stringify({ type: "session", version: 3, id: "parent-1", timestamp: "2026-01-01T00:00:00.000Z", cwd: "/workspace" })}\n`, "utf8");
await writeFile(forkParentFile, `${JSON.stringify({ type: "session", version: 3, id: "parent-1-fork", timestamp: "2026-01-01T00:00:00.000Z", cwd: "/workspace" })}\n`, "utf8");
await writeFile(childFile, `${JSON.stringify({ type: "session", version: 3, id: "child-1", timestamp: "2026-01-01T00:00:00.000Z", cwd: "/workspace-feature", parentSession: parentFile })}\n`, "utf8");
try {
const childManager = fakeSessionManager("/workspace-feature", {
getHeader: () => ({ parentSession: parentFile }),
getEntries: () => [{ type: "custom", customType: "pi-web.subsession.spawned", data: { version: 1, spawnedBySessionId: "parent-1", spawnedSessionId: "child-1" } }],
});
const parentManager = fakeSessionManager("/workspace", {
getEntries: () => [{ type: "custom", customType: "pi-web.subsession.link", data: { version: 1, spawnedBySessionId: "parent-1", spawnedSessionId: "child-1", spawnedSessionFile: childFile, cwd: "/workspace-feature" } }],
});
const forkManager = fakeSessionManager("/workspace");
const fork = fakeRuntime("parent-1-fork", { sessionFile: forkParentFile, sessionManager: forkManager });
const child = fakeRuntime("child-1", { sessionFile: childFile, sessionManager: childManager });
const parent = fakeRuntime("parent-1", { sessionFile: parentFile, sessionManager: parentManager });
const runtimes = [fork.runtime, child.runtime, parent.runtime];
let index = 0;
const open = vi.fn((path: string) => {
if (path === parentFile) return parentManager;
if (path === forkParentFile) return forkManager;
return childManager;
});
const service = new PiSessionService(new CapturingSessionEventHub(), {
createAgentRuntime: () => {
const runtime = runtimes[index] ?? parent.runtime;
index += 1;
return Promise.resolve(runtime);
},
sessionManager: {
create: () => forkManager,
list: (cwd: string) => Promise.resolve(cwd === "/workspace"
? [{ ...sessionRecord("parent-1-fork", "/workspace"), path: forkParentFile }]
: [{ ...sessionRecord("child-1", "/workspace-feature"), path: childFile, parentSessionPath: parentFile }]),
listAll: () => Promise.resolve([]),
open,
},
archiveStore: emptyArchiveStore(),
heartbeatIntervalMs: 60_000,
});
await service.status(sessionRef("parent-1-fork", "/workspace"));
await service.status(sessionRef("child-1", "/workspace-feature"));
child.session.isStreaming = true;
child.emit({ type: "agent_start" });
child.session.isStreaming = false;
child.emit({ type: "agent_end" });
await new Promise((resolve) => setTimeout(resolve, 20));
expect(fork.calls.sendCustomMessage).toHaveLength(0);
expect(parent.calls.sendCustomMessage).toHaveLength(1);
expect(open).toHaveBeenCalledWith(parentFile);
await service.dispose();
} finally {
await rm(tempDir, { recursive: true, force: true });
}
});
it("does not relink a copied child with the original session id unless the parent link names the current child file", async () => {
const tempDir = await mkdtemp(join(tmpdir(), "pi-web-subsession-copied-child-"));
const parentFile = join(tempDir, "parent.jsonl");
const originalChildFile = join(tempDir, "original-child.jsonl");
const copiedChildFile = join(tempDir, "copied-child.jsonl");
await writeFile(parentFile, `${JSON.stringify({ type: "session", version: 3, id: "parent-1", timestamp: "2026-01-01T00:00:00.000Z", cwd: "/workspace" })}\n`, "utf8");
await writeFile(originalChildFile, `${JSON.stringify({ type: "session", version: 3, id: "child-1", timestamp: "2026-01-01T00:00:00.000Z", cwd: "/workspace-feature", parentSession: parentFile })}\n`, "utf8");
await writeFile(copiedChildFile, `${JSON.stringify({ type: "session", version: 3, id: "child-1", timestamp: "2026-01-01T00:00:00.000Z", cwd: "/workspace-feature", parentSession: parentFile })}\n`, "utf8");
try {
const childManager = fakeSessionManager("/workspace-feature", {
getHeader: () => ({ parentSession: parentFile }),
getEntries: () => [{ type: "custom", customType: "pi-web.subsession.spawned", data: { version: 1, spawnedBySessionId: "parent-1", spawnedSessionId: "child-1" } }],
});
const parentManager = fakeSessionManager("/workspace", {
getEntries: () => [{ type: "custom", customType: "pi-web.subsession.link", data: { version: 1, spawnedBySessionId: "parent-1", spawnedSessionId: "child-1", spawnedSessionFile: originalChildFile, cwd: "/workspace-feature" } }],
});
const child = fakeRuntime("child-1", { sessionFile: copiedChildFile, sessionManager: childManager });
const parent = fakeRuntime("parent-1", { sessionFile: parentFile, sessionManager: parentManager });
const runtimes = [child.runtime, parent.runtime];
let index = 0;
const open = vi.fn((path: string) => path === parentFile ? parentManager : childManager);
const service = new PiSessionService(new CapturingSessionEventHub(), {
createAgentRuntime: () => {
const runtime = runtimes[index] ?? parent.runtime;
index += 1;
return Promise.resolve(runtime);
},
sessionManager: {
create: () => childManager,
list: () => Promise.resolve([{ ...sessionRecord("child-1", "/workspace-feature"), path: copiedChildFile, parentSessionPath: parentFile }]),
listAll: () => Promise.resolve([]),
open,
},
archiveStore: emptyArchiveStore(),
heartbeatIntervalMs: 60_000,
});
await service.status(sessionRef("child-1", "/workspace-feature"));
child.session.isStreaming = true;
child.emit({ type: "agent_start" });
child.session.isStreaming = false;
child.emit({ type: "agent_end" });
await new Promise((resolve) => setTimeout(resolve, 20));
expect(parent.calls.sendCustomMessage).toHaveLength(0);
await service.dispose();
} finally {
await rm(tempDir, { recursive: true, force: true });
}
});
it("does not relink a child marker when the child header points at a different parent id", async () => { it("does not relink a child marker when the child header points at a different parent id", async () => {
const tempDir = await mkdtemp(join(tmpdir(), "pi-web-subsession-wrong-parent-")); const tempDir = await mkdtemp(join(tmpdir(), "pi-web-subsession-wrong-parent-"));
const mismatchedParentFile = join(tempDir, "other-parent.jsonl"); const mismatchedParentFile = join(tempDir, "other-parent.jsonl");
+32 -7
View File
@@ -633,14 +633,35 @@ export class PiSessionService {
const parentHeader = await readSessionHeaderSummary(parentSessionFile); const parentHeader = await readSessionHeaderSummary(parentSessionFile);
if (parentHeader?.id !== marker.spawnedBySessionId) return; if (parentHeader?.id !== marker.spawnedBySessionId) return;
const childSessionFile = nonEmptyString(session.sessionFile); const childSessionFile = nonEmptyString(session.sessionFile);
if (childSessionFile === undefined) return;
const hasReciprocalLink = await this.parentHasReciprocalSubsessionLink(parentSessionFile, marker.spawnedBySessionId, session.sessionId, childSessionFile);
if (!hasReciprocalLink) return;
this.registerSubsession(marker.spawnedBySessionId, { this.registerSubsession(marker.spawnedBySessionId, {
childSessionId: session.sessionId, childSessionId: session.sessionId,
...(childSessionFile === undefined ? {} : { childSessionFile }), childSessionFile,
parentSessionFile, parentSessionFile,
cwd: session.sessionManager.getCwd(), cwd: session.sessionManager.getCwd(),
}); });
} }
private async parentHasReciprocalSubsessionLink(parentSessionFile: string, parentSessionId: string, childSessionId: string, childSessionFile: string): Promise<boolean> {
let parentManager: PiSessionManager;
try {
parentManager = this.sessionManager.open(parentSessionFile);
} catch {
return false;
}
const entries = parentManager.getEntries?.() ?? parentManager.getBranch();
for (const entry of entries) {
const link = parsePersistedParentSubsessionLink(entry);
if (link === undefined) continue;
if (link.spawnedBySessionId !== parentSessionId || link.spawnedSessionId !== childSessionId) continue;
if (link.spawnedSessionFile === undefined || !sessionPathsEqual(link.spawnedSessionFile, childSessionFile)) continue;
if (await this.persistedSubsessionLinkMatchesParent(parentSessionFile, link)) return true;
}
return false;
}
private async getOrOpenTrackedSubsession(sessionId: string): Promise<PiAgentSession> { private async getOrOpenTrackedSubsession(sessionId: string): Promise<PiAgentSession> {
const active = this.active.get(sessionId); const active = this.active.get(sessionId);
if (active !== undefined) return active.runtime.session; if (active !== undefined) return active.runtime.session;
@@ -657,7 +678,11 @@ export class PiSessionService {
} }
} }
return this.getOrOpen(sessionId); const listed = link?.cwd === undefined
? (await this.sessionManager.listAll?.() ?? []).find((session) => session.id === sessionId)
: (await this.sessionManager.list(link.cwd)).find((session) => session.id === sessionId);
if (listed === undefined) throw new Error("Session not found");
return (await this.create(this.sessionManager.open(listed.path), listed.cwd)).runtime.session;
} }
private async subsessionSummaryFields(childSessionId: string): Promise<{ cwd: string; status: SubsessionStatus }> { private async subsessionSummaryFields(childSessionId: string): Promise<{ cwd: string; status: SubsessionStatus }> {
@@ -706,18 +731,18 @@ export class PiSessionService {
} }
private async getOrOpenParentForSubsession(parentSessionId: string, childSessionId: string): Promise<PiAgentSession> { private async getOrOpenParentForSubsession(parentSessionId: string, childSessionId: string): Promise<PiAgentSession> {
const active = this.activeForLookup(parentSessionId); const active = this.active.get(parentSessionId);
if (active !== undefined) return active.runtime.session; if (active !== undefined) return active.runtime.session;
const parentSessionFile = this.subsessionLinks.get(childSessionId)?.parentSessionFile; const parentSessionFile = this.subsessionLinks.get(childSessionId)?.parentSessionFile;
if (parentSessionFile !== undefined && (await readSessionHeaderSummary(parentSessionFile))?.id === parentSessionId) { if (parentSessionFile === undefined) throw new Error(`Parent session ${parentSessionId} is not available for subsession notification`);
if ((await readSessionHeaderSummary(parentSessionFile))?.id !== parentSessionId) {
throw new Error(`Parent session ${parentSessionId} is not available for subsession notification`);
}
const sessionManager = this.sessionManager.open(parentSessionFile); const sessionManager = this.sessionManager.open(parentSessionFile);
return (await this.create(sessionManager, sessionManager.getCwd())).runtime.session; return (await this.create(sessionManager, sessionManager.getCwd())).runtime.session;
} }
return this.getOrOpen(parentSessionId);
}
/** /**
* Deliver a subsession-completion notice to the parent as a system-authored * Deliver a subsession-completion notice to the parent as a system-authored
* custom message rather than a user message, so it is not attributed to the * custom message rather than a user message, so it is not attributed to the