fix: enforce exact subsession recovery links

This commit is contained in:
Federico Jaramillo Martinez
2026-06-25 09:27:10 +02:00
parent 417b04a23f
commit 5550c60950
2 changed files with 127 additions and 67 deletions
@@ -1291,6 +1291,60 @@ describe("PiSessionService", () => {
} }
}); });
it("does not relink a child marker when the current child file header no longer records the parent", async () => {
const tempDir = await mkdtemp(join(tmpdir(), "pi-web-subsession-stale-child-header-"));
const parentFile = join(tempDir, "parent.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(childFile, `${JSON.stringify({ type: "session", version: 3, id: "child-1", timestamp: "2026-01-01T00:00:00.000Z", cwd: "/workspace-feature" })}\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 child = fakeRuntime("child-1", { sessionFile: childFile, 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: childFile, parentSessionPath: parentFile }]),
listAll: () => Promise.resolve([]),
open,
},
archiveStore: {
...emptyArchiveStore(),
get: (sessionId) => Promise.resolve(sessionId === "child-1" ? { sessionId: "child-1", cwd: "/workspace-feature", archivedAt: "2026-01-01T00:00:00.000Z", parentSessionPath: parentFile } : undefined),
},
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);
expect(open).not.toHaveBeenCalledWith(parentFile);
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");
+70 -64
View File
@@ -468,14 +468,15 @@ export class PiSessionService {
if (!decision.allowed) throw spawnTargetError(decision); if (!decision.allowed) throw spawnTargetError(decision);
const created = await this.start(decision.cwd, input.parentSessionFile); const created = await this.start(decision.cwd, input.parentSessionFile);
const parentSessionFile = nonEmptyString(input.parentSessionFile); const parentSessionFile = nonEmptyString(input.parentSessionFile);
const link = { const link: TrackedSubsessionLink = {
parentSessionId: input.parentSessionId,
childSessionId: created.id, childSessionId: created.id,
...(created.path === "" ? {} : { childSessionFile: created.path }), ...(created.path === "" ? {} : { childSessionFile: created.path }),
...(parentSessionFile === undefined ? {} : { parentSessionFile }), ...(parentSessionFile === undefined ? {} : { parentSessionFile }),
cwd: decision.cwd, cwd: decision.cwd,
}; };
this.registerSubsession(input.parentSessionId, link); this.registerVerifiedSubsession(link);
this.persistSubsessionLink(input.parentSessionId, link); this.persistSubsessionLink(link);
this.persistSubsessionChildMarker(input.parentSessionId, created.id); this.persistSubsessionChildMarker(input.parentSessionId, created.id);
await this.prompt(created.id, input.prompt); await this.prompt(created.id, input.prompt);
this.logger.info( this.logger.info(
@@ -527,8 +528,8 @@ export class PiSessionService {
return this.getOrOpenTrackedSubsession(sessionId); return this.getOrOpenTrackedSubsession(sessionId);
} }
private registerSubsession(parentSessionId: string, link: Omit<TrackedSubsessionLink, "parentSessionId">): void { private registerVerifiedSubsession(link: TrackedSubsessionLink): void {
const childSessionId = link.childSessionId; const { childSessionId, parentSessionId } = link;
const previousParentId = this.subsessionParents.get(childSessionId); const previousParentId = this.subsessionParents.get(childSessionId);
if (previousParentId !== undefined && previousParentId !== parentSessionId) { if (previousParentId !== undefined && previousParentId !== parentSessionId) {
const previousChildren = this.subsessionChildren.get(previousParentId); const previousChildren = this.subsessionChildren.get(previousParentId);
@@ -541,8 +542,7 @@ export class PiSessionService {
children.add(childSessionId); children.add(childSessionId);
this.subsessionChildren.set(parentSessionId, children); this.subsessionChildren.set(parentSessionId, children);
const previous = this.subsessionLinks.get(childSessionId); this.subsessionLinks.set(childSessionId, link);
this.subsessionLinks.set(childSessionId, mergeSubsessionLink(previous, { ...link, parentSessionId }));
if (!this.subsessionNotifyArmed.has(childSessionId)) this.subsessionNotifyArmed.set(childSessionId, false); if (!this.subsessionNotifyArmed.has(childSessionId)) this.subsessionNotifyArmed.set(childSessionId, false);
} }
@@ -557,15 +557,15 @@ export class PiSessionService {
if (children?.size === 0) this.subsessionChildren.delete(parentSessionId); if (children?.size === 0) this.subsessionChildren.delete(parentSessionId);
} }
private persistSubsessionLink(parentSessionId: string, link: Omit<TrackedSubsessionLink, "parentSessionId">): void { private persistSubsessionLink(link: TrackedSubsessionLink): void {
const parent = this.active.get(parentSessionId)?.runtime.session; const parent = this.active.get(link.parentSessionId)?.runtime.session;
if (parent === undefined) return; if (parent === undefined) return;
if (parent.sessionManager.appendCustomEntry === undefined) return; if (parent.sessionManager.appendCustomEntry === undefined) return;
try { try {
parent.sessionManager.appendCustomEntry(SUBSESSION_LINK_CUSTOM_TYPE, persistedParentSubsessionLinkData(parentSessionId, link)); parent.sessionManager.appendCustomEntry(SUBSESSION_LINK_CUSTOM_TYPE, persistedParentSubsessionLinkData(link));
} catch (error: unknown) { } catch (error: unknown) {
this.logger.info( this.logger.info(
{ parentSessionId, sessionId: link.childSessionId, error: error instanceof Error ? error.message : String(error) }, { parentSessionId: link.parentSessionId, sessionId: link.childSessionId, error: error instanceof Error ? error.message : String(error) },
"failed to persist subsession link", "failed to persist subsession link",
); );
} }
@@ -596,60 +596,81 @@ export class PiSessionService {
} }
private async registerPersistedSubsessionLinks(parentSessionId: string, parent: PiAgentSession, parentSessionFile: string | undefined): Promise<void> { private async registerPersistedSubsessionLinks(parentSessionId: string, parent: PiAgentSession, parentSessionFile: string | undefined): Promise<void> {
// Parent custom links are the authoritative recovery record: verify the
// exact live child file/header or an exact archived child before tracking.
const entries = parent.sessionManager.getEntries?.() ?? parent.sessionManager.getBranch(); const entries = parent.sessionManager.getEntries?.() ?? parent.sessionManager.getBranch();
for (const entry of entries) { for (const entry of entries) {
const link = parsePersistedParentSubsessionLink(entry); const link = parsePersistedParentSubsessionLink(entry);
if (link === undefined) continue; if (link === undefined) continue;
if (link.spawnedBySessionId !== parentSessionId) continue; const verified = await this.verifiedSubsessionLinkFromParentLink(parentSessionId, parentSessionFile, link);
if (!await this.persistedSubsessionLinkMatchesParent(parentSessionFile, link)) continue; if (verified === undefined) continue;
this.registerSubsession(parentSessionId, trackedSubsessionLinkFromParentLink(link, parentSessionFile)); this.registerVerifiedSubsession(verified);
} }
} }
private async persistedSubsessionLinkMatchesParent(parentSessionFile: string | undefined, link: PersistedParentSubsessionLink): Promise<boolean> { private async verifiedSubsessionLinkFromParentLink(parentSessionId: string, parentSessionFile: string | undefined, link: PersistedParentSubsessionLink): Promise<TrackedSubsessionLink | undefined> {
if (parentSessionFile === undefined) return false; if (parentSessionFile === undefined) return undefined;
if (link.spawnedSessionFile !== undefined) { if (link.spawnedBySessionId !== parentSessionId) return undefined;
const header = await readSessionHeaderSummary(link.spawnedSessionFile); if (!(await this.parentLinkHasValidChildTarget(parentSessionFile, link))) return undefined;
if (header?.id === link.spawnedSessionId) { return trackedSubsessionLinkFromParentLink(parentSessionId, link, parentSessionFile);
return header.parentSession !== undefined && sessionPathsEqual(header.parentSession, parentSessionFile);
}
} }
private async parentLinkHasValidChildTarget(parentSessionFile: string, link: PersistedParentSubsessionLink): Promise<boolean> {
if (link.spawnedSessionFile !== undefined && (await sessionFileHeaderMatches(link.spawnedSessionFile, { sessionId: link.spawnedSessionId, parentSessionFile }))) return true;
return this.archivedSubsessionLinkMatchesParent(parentSessionFile, link);
}
private async archivedSubsessionLinkMatchesParent(parentSessionFile: string, link: PersistedParentSubsessionLink): Promise<boolean> {
const archived = await this.getArchivedExact(link.spawnedSessionId); const archived = await this.getArchivedExact(link.spawnedSessionId);
return archived?.parentSessionPath !== undefined && sessionPathsEqual(archived.parentSessionPath, parentSessionFile); if (archived?.parentSessionPath === undefined) return false;
if (!sessionPathsEqual(archived.parentSessionPath, parentSessionFile)) return false;
if (archived.originalPath !== undefined && link.spawnedSessionFile !== undefined && !sessionPathsEqual(archived.originalPath, link.spawnedSessionFile)) return false;
return true;
} }
private async recoverSubsessionTrackingForOpenedSession(session: PiAgentSession): Promise<void> { private async recoverSubsessionTrackingForOpenedSession(session: PiAgentSession): Promise<void> {
const link = await this.verifiedSubsessionLinkFromOpenedChild(session);
if (link === undefined) return;
this.registerVerifiedSubsession(link);
}
private async verifiedSubsessionLinkFromOpenedChild(session: PiAgentSession): Promise<TrackedSubsessionLink | undefined> {
// Child markers are only hints; the current child header and reciprocal
// parent custom link must agree on the exact ids and files before relinking.
const entries = session.sessionManager.getEntries?.() ?? session.sessionManager.getBranch(); const entries = session.sessionManager.getEntries?.() ?? session.sessionManager.getBranch();
let marker: PersistedChildSubsessionLink | undefined; let marker: PersistedChildSubsessionLink | undefined;
for (const entry of entries) { for (const entry of entries) {
const parsed = parsePersistedChildSubsessionLink(entry); const parsed = parsePersistedChildSubsessionLink(entry);
if (parsed?.spawnedSessionId === session.sessionId) marker = parsed; if (parsed?.spawnedSessionId === session.sessionId) marker = parsed;
} }
if (marker === undefined) return; if (marker === undefined) return undefined;
const parentSessionFile = await parentSessionFileForSession(session);
if (parentSessionFile === undefined) return;
const parentHeader = await readSessionHeaderSummary(parentSessionFile);
if (parentHeader?.id !== marker.spawnedBySessionId) return;
const childSessionFile = nonEmptyString(session.sessionFile); const childSessionFile = nonEmptyString(session.sessionFile);
if (childSessionFile === undefined) return; if (childSessionFile === undefined) return undefined;
const hasReciprocalLink = await this.parentHasReciprocalSubsessionLink(parentSessionFile, marker.spawnedBySessionId, session.sessionId, childSessionFile); const childHeader = await readSessionHeaderSummary(childSessionFile);
if (!hasReciprocalLink) return; if (childHeader?.id !== session.sessionId) return undefined;
this.registerSubsession(marker.spawnedBySessionId, { const parentSessionFile = nonEmptyString(childHeader.parentSession);
if (parentSessionFile === undefined) return undefined;
const parentHeader = await readSessionHeaderSummary(parentSessionFile);
if (parentHeader?.id !== marker.spawnedBySessionId) return undefined;
const parentLink = this.findReciprocalParentSubsessionLink(parentSessionFile, marker.spawnedBySessionId, session.sessionId, childSessionFile);
if (parentLink === undefined) return undefined;
return {
parentSessionId: marker.spawnedBySessionId,
childSessionId: session.sessionId, childSessionId: session.sessionId,
childSessionFile, childSessionFile,
parentSessionFile, parentSessionFile,
cwd: session.sessionManager.getCwd(), cwd: parentLink.cwd ?? session.sessionManager.getCwd(),
}); };
} }
private async parentHasReciprocalSubsessionLink(parentSessionFile: string, parentSessionId: string, childSessionId: string, childSessionFile: string): Promise<boolean> { private findReciprocalParentSubsessionLink(parentSessionFile: string, parentSessionId: string, childSessionId: string, childSessionFile: string): PersistedParentSubsessionLink | undefined {
let parentManager: PiSessionManager; let parentManager: PiSessionManager;
try { try {
parentManager = this.sessionManager.open(parentSessionFile); parentManager = this.sessionManager.open(parentSessionFile);
} catch { } catch {
return false; return undefined;
} }
const entries = parentManager.getEntries?.() ?? parentManager.getBranch(); const entries = parentManager.getEntries?.() ?? parentManager.getBranch();
for (const entry of entries) { for (const entry of entries) {
@@ -657,9 +678,9 @@ export class PiSessionService {
if (link === undefined) continue; if (link === undefined) continue;
if (link.spawnedBySessionId !== parentSessionId || link.spawnedSessionId !== childSessionId) continue; if (link.spawnedBySessionId !== parentSessionId || link.spawnedSessionId !== childSessionId) continue;
if (link.spawnedSessionFile === undefined || !sessionPathsEqual(link.spawnedSessionFile, childSessionFile)) continue; if (link.spawnedSessionFile === undefined || !sessionPathsEqual(link.spawnedSessionFile, childSessionFile)) continue;
if (await this.persistedSubsessionLinkMatchesParent(parentSessionFile, link)) return true; return link;
} }
return false; return undefined;
} }
private async getOrOpenTrackedSubsession(sessionId: string): Promise<PiAgentSession> { private async getOrOpenTrackedSubsession(sessionId: string): Promise<PiAgentSession> {
@@ -671,12 +692,10 @@ export class PiSessionService {
const link = this.subsessionLinks.get(sessionId); const link = this.subsessionLinks.get(sessionId);
if (link?.childSessionFile !== undefined) { if (link?.childSessionFile !== undefined) {
const header = await readSessionHeaderSummary(link.childSessionFile); if (!(await sessionFileHeaderMatches(link.childSessionFile, { sessionId, parentSessionFile: link.parentSessionFile }))) throw new Error("Session not found");
if (header?.id === sessionId) {
const sessionManager = this.sessionManager.open(link.childSessionFile); const sessionManager = this.sessionManager.open(link.childSessionFile);
return (await this.create(sessionManager, link.cwd ?? sessionManager.getCwd())).runtime.session; return (await this.create(sessionManager, link.cwd ?? sessionManager.getCwd())).runtime.session;
} }
}
const listed = link?.cwd === undefined const listed = link?.cwd === undefined
? (await this.sessionManager.listAll?.() ?? []).find((session) => session.id === sessionId) ? (await this.sessionManager.listAll?.() ?? []).find((session) => session.id === sessionId)
@@ -693,7 +712,7 @@ export class PiSessionService {
const archived = await this.getArchivedExact(childSessionId); const archived = await this.getArchivedExact(childSessionId);
if (archived !== undefined) return { cwd: archived.cwd, status: "archived" }; if (archived !== undefined) return { cwd: archived.cwd, status: "archived" };
const link = this.subsessionLinks.get(childSessionId); const link = this.subsessionLinks.get(childSessionId);
if (link?.childSessionFile !== undefined && (await readSessionHeaderSummary(link.childSessionFile))?.id === childSessionId) { if (link?.childSessionFile !== undefined && (await sessionFileHeaderMatches(link.childSessionFile, { sessionId: childSessionId, parentSessionFile: link.parentSessionFile }))) {
return { cwd: link.cwd ?? "", status: "idle" }; return { cwd: link.cwd ?? "", status: "idle" };
} }
if (link?.cwd !== undefined) return { cwd: link.cwd, status: "unknown" }; if (link?.cwd !== undefined) return { cwd: link.cwd, status: "unknown" };
@@ -1627,32 +1646,20 @@ function isDefined<T>(value: T | undefined): value is T {
return value !== undefined; return value !== undefined;
} }
function mergeSubsessionLink(previous: TrackedSubsessionLink | undefined, next: TrackedSubsessionLink): TrackedSubsessionLink { function trackedSubsessionLinkFromParentLink(parentSessionId: string, link: PersistedParentSubsessionLink, parentSessionFile: string): TrackedSubsessionLink {
return {
parentSessionId: next.parentSessionId,
childSessionId: next.childSessionId,
...(previous?.childSessionFile === undefined ? {} : { childSessionFile: previous.childSessionFile }),
...(previous?.parentSessionFile === undefined ? {} : { parentSessionFile: previous.parentSessionFile }),
...(previous?.cwd === undefined ? {} : { cwd: previous.cwd }),
...(next.childSessionFile === undefined ? {} : { childSessionFile: next.childSessionFile }),
...(next.parentSessionFile === undefined ? {} : { parentSessionFile: next.parentSessionFile }),
...(next.cwd === undefined ? {} : { cwd: next.cwd }),
};
}
function trackedSubsessionLinkFromParentLink(link: PersistedParentSubsessionLink, parentSessionFile: string | undefined): Omit<TrackedSubsessionLink, "parentSessionId"> {
return { return {
parentSessionId,
childSessionId: link.spawnedSessionId, childSessionId: link.spawnedSessionId,
...(link.spawnedSessionFile === undefined ? {} : { childSessionFile: link.spawnedSessionFile }), ...(link.spawnedSessionFile === undefined ? {} : { childSessionFile: link.spawnedSessionFile }),
...(parentSessionFile === undefined ? {} : { parentSessionFile }), parentSessionFile,
...(link.cwd === undefined ? {} : { cwd: link.cwd }), ...(link.cwd === undefined ? {} : { cwd: link.cwd }),
}; };
} }
function persistedParentSubsessionLinkData(parentSessionId: string, link: Omit<TrackedSubsessionLink, "parentSessionId">): Record<string, unknown> { function persistedParentSubsessionLinkData(link: TrackedSubsessionLink): Record<string, unknown> {
return { return {
version: 1, version: 1,
spawnedBySessionId: parentSessionId, spawnedBySessionId: link.parentSessionId,
spawnedSessionId: link.childSessionId, spawnedSessionId: link.childSessionId,
...(link.childSessionFile === undefined ? {} : { spawnedSessionFile: link.childSessionFile }), ...(link.childSessionFile === undefined ? {} : { spawnedSessionFile: link.childSessionFile }),
...(link.cwd === undefined ? {} : { cwd: link.cwd }), ...(link.cwd === undefined ? {} : { cwd: link.cwd }),
@@ -1726,12 +1733,11 @@ async function readSessionHeaderSummary(sessionFile: string): Promise<SessionHea
} }
} }
async function parentSessionFileForSession(session: PiAgentSession): Promise<string | undefined> { async function sessionFileHeaderMatches(sessionFile: string, expected: { sessionId: string; parentSessionFile?: string | undefined }): Promise<boolean> {
const headerParentSession = nonEmptyString(session.sessionManager.getHeader?.()?.parentSession); const header = await readSessionHeaderSummary(sessionFile);
if (headerParentSession !== undefined) return headerParentSession; if (header?.id !== expected.sessionId) return false;
const sessionFile = nonEmptyString(session.sessionFile); if (expected.parentSessionFile === undefined) return true;
if (sessionFile === undefined) return undefined; return header.parentSession !== undefined && sessionPathsEqual(header.parentSession, expected.parentSessionFile);
return (await readSessionHeaderSummary(sessionFile))?.parentSession;
} }
async function clearParentSession(sessionFile: string): Promise<void> { async function clearParentSession(sessionFile: string): Promise<void> {