Replace pointer-comparison stale check with epoch-based invalidation in spawn renderer

Author: ofriw

agenticoding.test.ts

@@ -969,6 +969,7 @@ test("nested spawn live action tracks tool execution events", () => {
 	const childSpawnTool = createChildTools(new MockPi() as any, state, "medium", 0).find(t => t.name === "spawn")!;
 	const { session, emit } = createSubscribableSession([]);
 	state.childSessions.set("tool-call-1", session);
+	state.liveChildSessions.set("tool-call-1", session);
 
 	// Mock console.warn to suppress any expected-but-harmless warnings
 	// (e.g., streaming component errors in headless test env).
@@ -1012,6 +1013,7 @@ test("nested spawn handleEvent recovers from malformed events", () => {
 	const childSpawnTool = createChildTools(new MockPi() as any, state, "medium", 0).find(t => t.name === "spawn")!;
 	const { session, emit } = createSubscribableSession([]);
 	state.childSessions.set("tool-call-1", session);
+	state.liveChildSessions.set("tool-call-1", session);
 
 	const warnings: any[] = [];
 	const originalWarn = console.warn;
@@ -1044,6 +1046,7 @@ test("nested spawn message_end with aborted stopReason clears pending tools", ()
 	const childSpawnTool = createChildTools(new MockPi() as any, state, "medium", 0).find(t => t.name === "spawn")!;
 	const { session, emit } = createSubscribableSession([]);
 	state.childSessions.set("tool-call-1", session);
+	state.liveChildSessions.set("tool-call-1", session);
 
 	const component = childSpawnTool.renderResult(
 		{ content: [{ type: "text", text: "ignored" }], details: { depth: 1, model: "m", thinking: "low", truncated: false } },
@@ -1066,6 +1069,7 @@ test("nested spawn dispose stops event processing", () => {
 	const childSpawnTool = createChildTools(new MockPi() as any, state, "medium", 0).find(t => t.name === "spawn")!;
 	const { session, emit } = createSubscribableSession([]);
 	state.childSessions.set("tool-call-1", session);
+	state.liveChildSessions.set("tool-call-1", session);
 
 	const component = childSpawnTool.renderResult(
 		{ content: [{ type: "text", text: "ignored" }], details: { depth: 1, model: "m", thinking: "low", truncated: false } },
@@ -1670,6 +1674,7 @@ test("nested spawn rapid events collapse to last state", () => {
 	const childSpawnTool = createChildTools(new MockPi() as any, state, "medium", 0).find(t => t.name === "spawn")!;
 	const { session, emit } = createSubscribableSession([]);
 	state.childSessions.set("tool-call-1", session);
+	state.liveChildSessions.set("tool-call-1", session);
 
 	const component = childSpawnTool.renderResult(
 		{ content: [], details: { depth: 1, model: "m", thinking: "low", truncated: false } },
@@ -1697,30 +1702,107 @@ test("nested spawn rapid events collapse to last state", () => {
 	assert.ok(finalLines.some((l: string) => l.includes("✓")));
 });
 
-test("nested spawn isDetachedFromLiveState drops events when session is replaced", () => {
+test("nested spawn drops late events after live registry deletion", () => {
 	const state = createState();
 	const childSpawnTool = createChildTools(new MockPi() as any, state, "medium", 0).find(t => t.name === "spawn")!;
 	const { session, emit } = createSubscribableSession([]);
 	state.childSessions.set("tool-call-1", session);
 	state.liveChildSessions.set("tool-call-1", session);
+	let invalidateCalls = 0;
 
 	const component = childSpawnTool.renderResult(
 		{ content: [{ type: "text", text: "initial" }], details: { depth: 1, model: "m", thinking: "low", truncated: false } },
 		{ expanded: false },
 		theme,
-		createRenderContext(),
+		createRenderContext({ invalidate: () => { invalidateCalls++; } }),
+	) as any;
+	const before = component.render(120);
+
+	state.liveChildSessions.delete("tool-call-1");
+	emit({ type: "message_start", message: { role: "assistant", content: [] } });
+
+	const after = component.render(120);
+	assert.equal(invalidateCalls, 0, "completed-session deletion should stop rerenders from late events");
+	assert.deepEqual(after, before, "completed-session deletion should freeze the rendered state");
+});
+
+test("nested spawn drops events after resetState bumps child epoch", () => {
+	const state = createState();
+	const childSpawnTool = createChildTools(new MockPi() as any, state, "medium", 0).find(t => t.name === "spawn")!;
+	const { session, emit } = createSubscribableSession([]);
+	state.childSessions.set("tool-call-1", session);
+	state.liveChildSessions.set("tool-call-1", session);
+	let invalidateCalls = 0;
+
+	const component = childSpawnTool.renderResult(
+		{ content: [{ type: "text", text: "initial" }], details: { depth: 1, model: "m", thinking: "low", truncated: false } },
+		{ expanded: false },
+		theme,
+		createRenderContext({ invalidate: () => { invalidateCalls++; } }),
+	) as any;
+	const before = component.render(120);
+
+	resetState(state);
+	emit({ type: "message_start", message: { role: "assistant", content: [] } });
+
+	const after = component.render(120);
+	assert.equal(invalidateCalls, 0, "stale events should not request rerender after reset");
+	assert.deepEqual(after, before, "stale events should not change rendered state after reset");
+});
+
+test("nested spawn drops events when session is replaced in live state", () => {
+	const state = createState();
+	const childSpawnTool = createChildTools(new MockPi() as any, state, "medium", 0).find(t => t.name === "spawn")!;
+	const { session, emit } = createSubscribableSession([]);
+	state.childSessions.set("tool-call-1", session);
+	state.liveChildSessions.set("tool-call-1", session);
+	let invalidateCalls = 0;
+
+	const component = childSpawnTool.renderResult(
+		{ content: [{ type: "text", text: "initial" }], details: { depth: 1, model: "m", thinking: "low", truncated: false } },
+		{ expanded: false },
+		theme,
+		createRenderContext({ invalidate: () => { invalidateCalls++; } }),
 	) as any;
+	const before = component.render(120);
 
-	// Replace the session in liveChildSessions with a different object
 	const replacementSession = createSubscribableSession([]).session;
 	state.liveChildSessions.set("tool-call-1", replacementSession);
+	emit({ type: "message_start", message: { role: "assistant", content: [] } });
+
+	const after = component.render(120);
+	assert.equal(invalidateCalls, 0, "replaced sessions should not request rerender");
+	assert.deepEqual(after, before, "replaced sessions should not change rendered state");
+});
 
-	// Emit a message_start event — should be silently dropped
+test("nested spawn completed-session deletion stays stale even if the toolCallId is later reused", () => {
+	const state = createState();
+	const childSpawnTool = createChildTools(new MockPi() as any, state, "medium", 0).find(t => t.name === "spawn")!;
+	const { session, emit } = createSubscribableSession([]);
+	state.childSessions.set("tool-call-1", session);
+	state.liveChildSessions.set("tool-call-1", session);
+	let invalidateCalls = 0;
+
+	const component = childSpawnTool.renderResult(
+		{ content: [{ type: "text", text: "initial" }], details: { depth: 1, model: "m", thinking: "low", truncated: false } },
+		{ expanded: false },
+		theme,
+		createRenderContext({ invalidate: () => { invalidateCalls++; } }),
+	) as any;
+	const before = component.render(120);
+
+	state.liveChildSessions.delete("tool-call-1");
 	emit({ type: "message_start", message: { role: "assistant", content: [] } });
+	const afterDeletion = component.render(120);
+	assert.equal(invalidateCalls, 0, "completed-session deletion should immediately stale the old session");
+	assert.deepEqual(afterDeletion, before, "completed-session deletion should freeze the rendered state before reuse");
 
-	const lines = component.render(120);
-	// Should NOT contain "thinking" because the event was dropped
-	assert.ok(lines.every((l: string) => !l.includes("thinking")), `expected no thinking after detach, got: ${lines.join("\n")}`);
+	state.liveChildSessions.set("tool-call-1", createSubscribableSession([]).session);
+	emit({ type: "message_update", message: { role: "assistant", content: [{ type: "text", text: "should be dropped" }] } });
+	const afterReuse = component.render(120);
+	assert.equal(invalidateCalls, 0, "toolCallId reuse should not revive a completed stale session");
+	assert.deepEqual(afterReuse, before, "toolCallId reuse should keep the old rendered state frozen");
+	assert.ok(afterReuse.every((l: string) => !l.includes("should be dropped")), "toolCallId reuse should not admit stale text updates");
 });
 
 test("concurrent spawn executions produce independent results", async () => {
@@ -1792,6 +1874,7 @@ test("nested spawn render cache preserves stable output for identical params", (
 	const childSpawnTool = createChildTools(new MockPi() as any, state, "medium", 0).find(t => t.name === "spawn")!;
 	const { session } = createSubscribableSession([]);
 	state.childSessions.set("tool-call-1", session);
+	state.liveChildSessions.set("tool-call-1", session);
 
 	const component = childSpawnTool.renderResult(
 		{ content: [{ type: "text", text: "hello" }], details: { depth: 1, model: "m", thinking: "low", truncated: false } },
@@ -1944,6 +2027,7 @@ test("handleEvent gracefully degrades with null message events", () => {
 	const childSpawnTool = createChildTools(new MockPi() as any, state, "medium", 0).find(t => t.name === "spawn")!;
 	const { session, emit } = createSubscribableSession([]);
 	state.childSessions.set("tool-call-1", session);
+	state.liveChildSessions.set("tool-call-1", session);
 
 	const component = childSpawnTool.renderResult(
 		{ content: [], details: { depth: 1, model: "m", thinking: "low", truncated: false } },

spawn/index.ts

@@ -335,8 +335,8 @@ export async function executeSpawn(
 		await abortAndInvalidate();
 	}
 
-	// liveChildSessions must be set first — renderSpawnResult checks it to decide
-	// whether to pass the live registry to attachSession for stale detection.
+	// liveChildSessions must be set before childSessions so the renderer can
+	// attach with a fully-published live ownership record.
 	state.liveChildSessions.set(toolCallId, session);
 	state.childSessions.set(toolCallId, session);
 

spawn/renderer.ts

@@ -178,7 +178,8 @@ class NestedAgentSessionComponent extends Container {
 	private details?: SpawnResultDetails;
 	private nestTheme?: Theme;
 	private ownedToolCallId?: string;
-	private liveChildSessions?: Map<string, AgentSession>;
+	private state?: AgenticodingState;
+	private attachedChildSessionEpoch?: number;
 	private liveOutcome: SpawnOutcome = "running";
 	// States: "⏳ initializing…" → "💭 thinking…" → "[tool] …/preview" or live text → terminal outcome
 	private lastAction = "";
@@ -226,11 +227,16 @@ class NestedAgentSessionComponent extends Container {
 		if (changed) this.clearRenderCache();
 	}
 
-	attachSession(toolCallId: string, session: AgentSession, liveChildSessions?: Map<string, AgentSession>): void {
+	attachSession(
+		toolCallId: string,
+		session: AgentSession,
+		state: AgenticodingState,
+	): void {
 		if (
 			this.session === session
 			&& this.ownedToolCallId === toolCallId
-			&& this.liveChildSessions === liveChildSessions
+			&& this.state === state
+			&& this.attachedChildSessionEpoch === state.childSessionEpoch
 		) {
 			return;
 		}
@@ -239,7 +245,8 @@ class NestedAgentSessionComponent extends Container {
 		this.unsubscribe = undefined;
 		this.session = session;
 		this.ownedToolCallId = toolCallId;
-		this.liveChildSessions = liveChildSessions;
+		this.state = state;
+		this.attachedChildSessionEpoch = state.childSessionEpoch;
 		this.liveOutcome = this.details?.outcome ?? "running";
 		this.toolNames.clear();
 		this.toolComponentFailures.clear();
@@ -270,18 +277,34 @@ class NestedAgentSessionComponent extends Container {
 	}
 
 	/**
-	 * Returns true when the session held by this component has been replaced
-	 * in the liveChildSessions map (e.g., after a resetState). When stale, the
-	 * component silently drops all events to avoid operating on a different
-	 * session than what it was attached to.
+	 * Returns the ownership invalidation reason for the attached session.
+	 *
+	 * Three stale paths:
+	 *   1. resetState() bumped childSessionEpoch after attach, invalidating all
+	 *      prior child sessions even if their objects still exist.
+	 *   2. state.liveChildSessions no longer contains this toolCallId because the
+	 *      child completed and cleared its live ownership.
+	 *   3. state.liveChildSessions now points this toolCallId at a different
+	 *      session, meaning a newer child claimed the slot.
 	 */
+	private getStaleSessionReason(): "epoch" | "completion" | "replacement" | undefined {
+		if (!this.session || !this.ownedToolCallId) {
+			return undefined;
+		}
+		if (this.state && this.attachedChildSessionEpoch !== this.state.childSessionEpoch) {
+			return "epoch";
+		}
+		const liveChildSessions = this.state?.liveChildSessions;
+		if (!liveChildSessions?.has(this.ownedToolCallId)) {
+			return "completion";
+		}
+		return liveChildSessions.get(this.ownedToolCallId) !== this.session
+			? "replacement"
+			: undefined;
+	}
+
 	private isStaleSession(): boolean {
-		return !!(
-			this.session
-			&& this.ownedToolCallId
-			&& this.liveChildSessions
-			&& this.liveChildSessions.get(this.ownedToolCallId) !== this.session
-		);
+		return this.getStaleSessionReason() !== undefined;
 	}
 
 	dispose(): void {
@@ -291,7 +314,7 @@ class NestedAgentSessionComponent extends Container {
 		// dispose, the nulled-out fields prevent double-abort.
 		const session = this.session;
 		const ownedToolCallId = this.ownedToolCallId;
-		const liveChildSessions = this.liveChildSessions;
+		const liveChildSessions = this.state?.liveChildSessions;
 		this.clearRenderCache();
 		this.details = undefined;
 		this.nestTheme = undefined;
@@ -300,7 +323,8 @@ class NestedAgentSessionComponent extends Container {
 		this.toolComponentFailures.clear();
 		this.session = undefined;
 		this.ownedToolCallId = undefined;
-		this.liveChildSessions = undefined;
+		this.state = undefined;
+		this.attachedChildSessionEpoch = undefined;
 		if (session && ownedToolCallId && liveChildSessions?.get(ownedToolCallId) === session) {
 			session.abort().catch(e => console.error("[spawn] abort failed:", ownedToolCallId, e));
 			liveChildSessions.delete(ownedToolCallId);
@@ -776,10 +800,7 @@ function renderSpawnResult(
 	}
 	const child = state.childSessions.get(context.toolCallId);
 	if (child) {
-		const liveChildSessions = state.liveChildSessions.get(context.toolCallId) === child
-			? state.liveChildSessions
-			: undefined;
-		component.attachSession(context.toolCallId, child, liveChildSessions);
+		component.attachSession(context.toolCallId, child, state);
 		state.childSessions.delete(context.toolCallId);
 		return component;
 	}

state.ts

@@ -37,10 +37,11 @@ export interface AgenticodingState {
 	/**
 	 * All live child agent sessions keyed by toolCallId, including claimed ones.
 	 * Reset/teardown aborts this registry so claimed children cannot outlive /new or UI disposal.
+	 * Completed children remove themselves from this registry before returning.
 	 *
 	 * INVARIANT: This Map is never replaced — only cleared via .clear().
-	 * NestedAgentSessionComponent holds a direct reference and depends on it
-	 * staying valid. If you change this, update attachSession in spawn/renderer.ts.
+	 * Spawn renderer ownership checks read this registry after attach, so its
+	 * identity must stay stable across resets, completion cleanup, and disposal.
 	 */
 	liveChildSessions: Map<string, AgentSession>;
 
@@ -65,9 +66,9 @@ export function createState(): AgenticodingState {
 		liveChildSessions,
 		childSessionEpoch: 0,
 	};
-	// Prevent replacement — NestedAgentSessionComponent holds direct references
-	// to both maps and depends on reference stability. Only .clear() and .delete()
-	// are valid — assigning a new Map would silently break session lifecycle.
+	// Prevent replacement — spawn lifecycle code and renderer ownership checks
+	// depend on stable map identity. Only .clear() and .delete() are valid —
+	// assigning a new Map would silently break child-session invalidation.
 	Object.defineProperty(state, 'childSessions', {
 		get: () => childSessions,
 		set: () => { throw new Error('childSessions cannot be replaced — use .clear() instead'); },