diff --git a/packages/coding-agent/src/registry/agent-lifecycle.ts b/packages/coding-agent/src/registry/agent-lifecycle.ts index 7fb86946a..3864c60f1 100644 --- a/packages/coding-agent/src/registry/agent-lifecycle.ts +++ b/packages/coding-agent/src/registry/agent-lifecycle.ts @@ -415,21 +415,27 @@ export class AgentLifecycleManager { async #revive(id: string, revive: AgentReviver, ref: AgentRef, adopted: AdoptedAgent): Promise { const session = await revive(ref); let liveRef = this.#registry.get(id); - if (liveRef === ref) { + if (liveRef === ref && ref.status === "parked" && !ref.session) { + // A simple reviver returned a session without claiming the parked ref; + // attach it here while the exact ref is still revivable. if (!this.#registry.attachSession(id, session, ref.sessionFile, ref)) { await session.dispose(); throw new Error(`Agent "${id}" changed before its persisted session could attach.`); } liveRef = ref; } else if ( - !liveRef || + liveRef !== ref || + liveRef.status !== "running" || liveRef.session !== session || liveRef.kind !== ref.kind || liveRef.parentId !== ref.parentId || liveRef.sessionFile !== ref.sessionFile ) { + // createAgentSession may have already claimed this exact parked ref and + // attached the returned session. Any other state — especially an + // `aborted` tombstone set while revive() was in flight — is stale. await session.dispose(); - throw new Error(`Agent "${id}" was replaced while its persisted session was reviving.`); + throw new Error(`Agent "${id}" was replaced or became terminal while its persisted session was reviving.`); } adopted.ref = liveRef; // Emits status_changed → "idle", which re-arms the TTL timer below. diff --git a/packages/coding-agent/src/registry/agent-registry.ts b/packages/coding-agent/src/registry/agent-registry.ts index fe85908cb..7baf1dd05 100644 --- a/packages/coding-agent/src/registry/agent-registry.ts +++ b/packages/coding-agent/src/registry/agent-registry.ts @@ -105,19 +105,23 @@ export class AgentRegistry { } /** - * Register a new id only when it is absent, or reuse the exact ref a parked - * revival was authorized to revive. A missing expected ref is a failed CAS: - * callers must never claim an id after its prior generation disappeared. + * Register a new id only when it is absent, or reuse the exact detached + * `parked` ref a revival was authorized to revive. A missing, replaced, or + * terminal expected ref is a failed CAS: delayed revivers must never claim an + * id after its prior generation disappeared or was hard-killed. */ registerIfAvailable(input: RegisterInput, expected: AgentRef | null): AgentRef | undefined { const current = this.#refs.get(input.id); if (expected === null) return current ? undefined : this.register(input); - return current === expected ? current : undefined; + return current === expected && current.status === "parked" && !current.session ? current : undefined; } setStatus(id: string, status: AgentStatus, expected?: AgentRefExpectation): boolean { const ref = this.#refs.get(id); if (!ref || !this.#matchesExpected(ref, expected)) return false; + // `aborted` is terminal: delayed progress/revival work from the killed + // generation must never transition the tombstone back to a live status. + if (ref.status === "aborted") return status === "aborted"; if (ref.status === status) return true; ref.status = status; // Activity describes current work; it is meaningless once the agent @@ -159,7 +163,10 @@ export class AgentRegistry { expected?: AgentRefExpectation, ): boolean { const ref = this.#refs.get(id); - if (!ref || !this.#matchesExpected(ref, expected)) return false; + // Never attach a late-created session to a hard-killed tombstone. This + // closes the race between a parked reviver claiming the ref and finishing + // createAgentSession after an explicit kill. + if (!ref || ref.status === "aborted" || !this.#matchesExpected(ref, expected)) return false; ref.session = session; if (sessionFile !== undefined) ref.sessionFile = sessionFile; ref.lastActivity = Date.now(); diff --git a/packages/coding-agent/test/registry/agent-lifecycle.test.ts b/packages/coding-agent/test/registry/agent-lifecycle.test.ts index b9923aa02..6c3356f66 100644 --- a/packages/coding-agent/test/registry/agent-lifecycle.test.ts +++ b/packages/coding-agent/test/registry/agent-lifecycle.test.ts @@ -75,6 +75,13 @@ describe("AgentLifecycleManager", () => { expect(registry.registerIfAvailable(next, parked)).toBe(parked); expect(registry.get("generation-Sub")).toBe(parked); + registry.setStatus("generation-Sub", "aborted", parked); + expect(registry.registerIfAvailable(next, parked)).toBeUndefined(); + const staleSession = makeSessionStub().session; + expect(registry.attachSession("generation-Sub", staleSession, undefined, parked)).toBe(false); + expect(registry.setStatus("generation-Sub", "idle", parked)).toBe(false); + expect(registry.get("generation-Sub")).toMatchObject({ status: "aborted", session: null }); + registry.unregister("generation-Sub", parked); expect(registry.registerIfAvailable(next, parked)).toBeUndefined(); expect(registry.get("generation-Sub")).toBeUndefined(); @@ -169,6 +176,39 @@ describe("AgentLifecycleManager", () => { expect(b).toBe(revived.session); }); + it("tombstoning a parked agent during revive prevents the stale session from attaching", async () => { + const gate = deferred(); + const revived = makeSessionStub(); + const ref = registry.register({ + id: "Revive-Killed", + displayName: "task", + kind: "sub", + session: null, + sessionFile: "/tmp/Revive-Killed.jsonl", + status: "parked", + }); + lifecycle.adopt( + "Revive-Killed", + { + idleTtlMs: 0, + revive: async () => { + await gate.promise; + return revived.session; + }, + }, + ref, + ); + + const revival = lifecycle.ensureLive("Revive-Killed"); + expect(await lifecycle.release("Revive-Killed", ref, { tombstone: true })).toBe(true); + expect(registry.get("Revive-Killed")).toMatchObject({ status: "aborted", session: null }); + + gate.resolve(); + await expect(revival).rejects.toThrow(/became terminal/); + expect(revived.disposeCalls()).toBe(1); + expect(registry.get("Revive-Killed")).toMatchObject({ status: "aborted", session: null }); + }); + it("ensureLive on an unknown id throws and points at history://", async () => { await expect(lifecycle.ensureLive("9-Ghost")).rejects.toThrow(/history:\/\/9-Ghost/); });