@oh-my-pi/pi-coding-agent 17.0.8 → 17.1.0

package/src/registry/agent-lifecycle.ts

@@ -14,13 +14,23 @@
  * - ensureLive during an in-flight park either cancels the park (session still
  *   live) or waits for detach+park and then revives.
  * - Concurrent ensureLive/park operations coalesce per id.
+ *
+ * Every adoption, park, and revival is bound to the exact {@link AgentRef} it
+ * started from, so stale async work (a late finalizer, a cancelled initializer,
+ * a superseded revive) can never clobber a newer same-id ref.
  */
 
 import { logger } from "@oh-my-pi/pi-utils";
 import type { AgentSession } from "../session/agent-session";
-import { type AgentRef, AgentRegistry, MAIN_AGENT_ID, type RegistryEvent } from "./agent-registry";
+import {
+	type AgentRef,
+	type AgentRefExpectation,
+	AgentRegistry,
+	MAIN_AGENT_ID,
+	type RegistryEvent,
+} from "./agent-registry";
 
-export type AgentReviver = () => Promise<AgentSession>;
+export type AgentReviver = (expected: AgentRef) => Promise<AgentSession>;
 
 /**
  * Builds a reviver for a `parked` ref restored from disk (Agent Hub scan,
@@ -39,12 +49,15 @@ export interface AdoptOptions {
 }
 
 interface AdoptedAgent {
+	ref: AgentRef;
 	idleTtlMs: number;
 	revive?: AgentReviver;
 	timer?: NodeJS.Timeout;
 }
 
 interface ParkInFlight {
+	/** The exact ref this park was started for. */
+	ref: AgentRef;
 	/** Resolves when the park attempt finishes (success, cancel, or dispose error). */
 	promise: Promise<void>;
 	/** Cancel before the session is detached. Returns true if cancel took effect. */
@@ -55,6 +68,11 @@ interface ParkInFlight {
 	detached: boolean;
 }
 
+interface RevivingAgent {
+	ref: AgentRef;
+	promise: Promise<AgentSession>;
+}
+
 export class AgentLifecycleManager {
 	static #global: AgentLifecycleManager | undefined;
 
@@ -85,12 +103,13 @@ export class AgentLifecycleManager {
 	readonly #registry: AgentRegistry;
 	readonly #adopted = new Map<string, AdoptedAgent>();
 	/**
-	 * In-flight park attempts. A park is cancelable until the live session is
-	 * detached; after detach, ensureLive waits for the park and revives.
+	 * In-flight park attempts, each bound to the ref it started from. A park is
+	 * cancelable until the live session is detached; after detach, ensureLive
+	 * waits for the park and revives.
 	 */
 	readonly #parks = new Map<string, ParkInFlight>();
-	/** In-flight revives, so concurrent {@link ensureLive} calls coalesce. */
-	readonly #revivals = new Map<string, Promise<AgentSession>>();
+	/** In-flight revives, bound to the parked ref that initiated them, so concurrent {@link ensureLive} calls coalesce. */
+	readonly #revivals = new Map<string, RevivingAgent>();
 	#unsubscribe: (() => void) | undefined;
 	#persistedReviverFactory: PersistedSubagentReviverFactory | undefined;
 	/** TTL applied when a cold-revived ref is adopted on demand. */
@@ -115,23 +134,29 @@ export class AgentLifecycleManager {
 	/**
 	 * Take ownership of a finished subagent. Caller has already set registry
 	 * status to "idle". Arms the TTL timer (idleTtlMs <= 0 adopts without one).
+	 * When `expected` is given, the adoption is refused if the id no longer
+	 * resolves to that ref (or that ref's session).
 	 */
-	adopt(id: string, opts: AdoptOptions): void {
+	adopt(id: string, opts: AdoptOptions, expected?: AgentRefExpectation): void {
 		if (id === MAIN_AGENT_ID) return;
-		if (!this.#registry.get(id)) {
-			logger.warn("AgentLifecycleManager.adopt: unknown agent id", { id });
+		const ref = this.#registry.get(id);
+		if (!ref || (expected !== undefined && ref !== expected && ref.session !== expected)) {
+			logger.warn("AgentLifecycleManager.adopt: unknown or replaced agent id", { id });
 			return;
 		}
 		const existing = this.#adopted.get(id);
 		clearTimeout(existing?.timer);
-		const adopted: AdoptedAgent = { idleTtlMs: opts.idleTtlMs, revive: opts.revive };
+		const adopted: AdoptedAgent = { ref, idleTtlMs: opts.idleTtlMs, revive: opts.revive };
 		this.#adopted.set(id, adopted);
 		this.#armTimer(id, adopted);
 	}
 
-	/** True if the id is adopted (parked or live). */
-	has(id: string): boolean {
-		return this.#adopted.has(id);
+	/** True if the id is adopted (parked or live) — and, when `expected` is given, still bound to that ref. */
+	has(id: string, expected?: AgentRefExpectation): boolean {
+		const adopted = this.#adopted.get(id);
+		return Boolean(
+			adopted && (expected === undefined || adopted.ref === expected || adopted.ref.session === expected),
+		);
 	}
 
 	/**
@@ -147,11 +172,14 @@ export class AgentLifecycleManager {
 	/**
 	 * True while {@link park} is disposing this agent's session (lets dispose
 	 * hooks distinguish park from teardown). False once the park is cancelled
-	 * by ensureLive or after detach+dispose completes.
+	 * by ensureLive or after detach+dispose completes. When `expected` is
+	 * given, only a park bound to that ref (or its session) counts.
 	 */
-	isParking(id: string): boolean {
+	isParking(id: string, expected?: AgentRefExpectation): boolean {
 		const park = this.#parks.get(id);
-		return Boolean(park && !park.cancelled);
+		return Boolean(
+			park && !park.cancelled && (expected === undefined || park.ref === expected || park.ref.session === expected),
+		);
 	}
 
 	/**
@@ -170,7 +198,8 @@ export class AgentLifecycleManager {
 		const adopted = this.#adopted.get(id);
 		if (!adopted) return;
 		const ref = this.#registry.get(id);
-		const session = ref?.session;
+		if (!ref || adopted.ref !== ref) return;
+		const session = ref.session;
 		if (!session) return;
 
 		if (adopted.timer) {
@@ -180,6 +209,7 @@ export class AgentLifecycleManager {
 
 		let cancelled = false;
 		const park: ParkInFlight = {
+			ref,
 			promise: undefined as unknown as Promise<void>,
 			cancel: () => {
 				// Cancel only before detach — once detached the old session is already
@@ -200,16 +230,16 @@ export class AgentLifecycleManager {
 				await Promise.resolve();
 				if (cancelled) return;
 
-				// Re-check liveness: release/unregister may have raced us.
+				// Re-check liveness: release/unregister/replace may have raced us.
 				const live = this.#registry.get(id);
-				if (!live?.session || live.session !== session) return;
-				if (!this.#adopted.has(id)) return;
+				if (live !== ref || !live.session || live.session !== session) return;
+				if (this.#adopted.get(id)?.ref !== ref) return;
 
 				// Commit: detach + parked *before* dispose so callers never see a
 				// dying session via ref.session / idle status.
 				park.detached = true;
-				this.#registry.detachSession(id);
-				this.#registry.setStatus(id, "parked");
+				this.#registry.detachSession(id, ref);
+				this.#registry.setStatus(id, "parked", ref);
 
 				try {
 					await session.dispose();
@@ -238,16 +268,16 @@ export class AgentLifecycleManager {
 	async ensureLive(id: string): Promise<AgentSession> {
 		const park = this.#parks.get(id);
 		if (park) {
-			const ref = this.#registry.get(id);
+			const parked = this.#registry.get(id);
 			// Cancel if the live session is still attached — keep it instead of
 			// thrashing dispose + revive.
-			if (ref?.session && !park.detached && park.cancel()) {
+			if (parked?.session && !park.detached && park.cancel()) {
 				await park.promise;
 				const kept = this.#registry.get(id)?.session;
 				if (kept) {
 					// Park cleared the idle timer; re-arm so TTL park still works.
 					const adopted = this.#adopted.get(id);
-					if (adopted && ref.status === "idle") this.#armTimer(id, adopted);
+					if (adopted && adopted.ref === parked && parked.status === "idle") this.#armTimer(id, adopted);
 					return kept;
 				}
 			} else {
@@ -265,13 +295,14 @@ export class AgentLifecycleManager {
 		}
 		if (ref.session) return ref.session;
 		const inflight = this.#revivals.get(id);
-		if (inflight) return inflight;
+		if (inflight?.ref === ref) return inflight.promise;
 		const revival = this.#resolveAndRevive(id, ref);
-		this.#revivals.set(id, revival);
+		const pending: RevivingAgent = { ref, promise: revival };
+		this.#revivals.set(id, pending);
 		try {
 			return await revival;
 		} finally {
-			this.#revivals.delete(id);
+			if (this.#revivals.get(id) === pending) this.#revivals.delete(id);
 		}
 	}
 
@@ -283,53 +314,72 @@ export class AgentLifecycleManager {
 	 * when the agent is not revivable or no reviver can be produced.
 	 */
 	async #resolveAndRevive(id: string, ref: AgentRef): Promise<AgentSession> {
-		let revive = this.#adopted.get(id)?.revive;
+		let adoption = this.#adopted.get(id);
+		let revive = adoption?.ref === ref ? adoption.revive : undefined;
 		let coldAdopted = false;
 		if (!revive && ref.status === "parked" && ref.sessionFile && this.#persistedReviverFactory) {
 			revive = await this.#persistedReviverFactory(ref);
 			if (revive) {
-				this.#adopted.set(id, { idleTtlMs: this.#persistedReviveTtlMs, revive });
+				adoption = { ref, idleTtlMs: this.#persistedReviveTtlMs, revive };
+				this.#adopted.set(id, adoption);
 				coldAdopted = true;
 			}
 		}
-		if (ref.status !== "parked" || !revive) {
+		if (this.#registry.get(id) !== ref) {
+			throw new Error(`Agent "${id}" changed while its persisted session was being prepared.`);
+		}
+		if (ref.status !== "parked" || !revive || !adoption) {
 			throw new Error(
 				`Agent "${id}" is ${ref.status} and cannot be revived${revive ? "" : " (no reviver registered)"}. Its transcript remains readable at history://${id}.`,
 			);
 		}
 		try {
-			return await this.#revive(id, revive, ref.sessionFile);
+			return await this.#revive(id, revive, ref, adoption);
 		} catch (error) {
 			// A failed cold revive (stale ctx, missing cwd, bad MCP) must not leave a
 			// poisoned reviver stuck in #adopted — drop it so a later ensureLive
 			// rebuilds via the factory (which may have fresher context by then).
-			if (coldAdopted) this.#adopted.delete(id);
+			if (coldAdopted && this.#adopted.get(id) === adoption) this.#adopted.delete(id);
 			throw error;
 		}
 	}
 
-	/** Hard removal: dispose if live, unregister from registry, drop timers. */
-	async release(id: string): Promise<void> {
+	/**
+	 * Hard removal: dispose if live, unregister from registry, drop timers.
+	 * When `expected` is given, only a ref matching it is released; a stale
+	 * release can never take down a newer same-id ref. Returns true when a
+	 * matching ref was released.
+	 */
+	async release(id: string, expected?: AgentRefExpectation): Promise<boolean> {
 		const adopted = this.#adopted.get(id);
-		clearTimeout(adopted?.timer);
-		this.#adopted.delete(id);
+		const current = this.#registry.get(id);
+		const currentMatches =
+			current && (expected === undefined || current === expected || current.session === expected);
+		const adoptedMatches =
+			adopted && (expected === undefined || adopted.ref === expected || adopted.ref.session === expected);
+		const ref = currentMatches ? current : adoptedMatches ? adopted.ref : undefined;
+		if (!ref) return false;
+		if (adopted?.ref === ref) {
+			clearTimeout(adopted.timer);
+			this.#adopted.delete(id);
+		}
 
 		const park = this.#parks.get(id);
-		if (park) {
+		if (park && park.ref === ref) {
 			// Prefer cancel when the session is still live so release owns dispose.
 			if (!park.detached) park.cancel();
 			await park.promise;
 		}
 
-		const ref = this.#registry.get(id);
-		if (ref?.session) {
+		if (this.#registry.get(id) === ref && ref.session) {
 			try {
 				await ref.session.dispose();
 			} catch (error) {
 				logger.warn("AgentLifecycleManager.release: session dispose failed", { id, error: String(error) });
 			}
 		}
-		this.#registry.unregister(id);
+		this.#registry.unregister(id, ref);
+		return true;
 	}
 
 	/** Teardown everything (process exit / main session dispose). */
@@ -343,11 +393,31 @@ export class AgentLifecycleManager {
 		this.#persistedReviverFactory = undefined;
 	}
 
-	async #revive(id: string, revive: AgentReviver, sessionFile: string | null): Promise<AgentSession> {
-		const session = await revive();
-		this.#registry.attachSession(id, session, sessionFile);
+	async #revive(id: string, revive: AgentReviver, ref: AgentRef, adopted: AdoptedAgent): Promise<AgentSession> {
+		const session = await revive(ref);
+		let liveRef = this.#registry.get(id);
+		if (liveRef === ref) {
+			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.session !== session ||
+			liveRef.kind !== ref.kind ||
+			liveRef.parentId !== ref.parentId ||
+			liveRef.sessionFile !== ref.sessionFile
+		) {
+			await session.dispose();
+			throw new Error(`Agent "${id}" was replaced while its persisted session was reviving.`);
+		}
+		adopted.ref = liveRef;
 		// Emits status_changed → "idle", which re-arms the TTL timer below.
-		this.#registry.setStatus(id, "idle");
+		if (!this.#registry.setStatus(id, "idle", liveRef)) {
+			await session.dispose();
+			throw new Error(`Agent "${id}" changed before its persisted session became idle.`);
+		}
 		return session;
 	}
 
@@ -364,7 +434,7 @@ export class AgentLifecycleManager {
 
 	#onRegistryEvent(event: RegistryEvent): void {
 		const adopted = this.#adopted.get(event.ref.id);
-		if (!adopted) return;
+		if (!adopted || adopted.ref !== event.ref) return;
 		if (event.type === "removed") {
 			clearTimeout(adopted.timer);
 			this.#adopted.delete(event.ref.id);

package/src/registry/agent-registry.ts

@@ -45,6 +45,8 @@ export interface AgentRef {
 	activity?: string;
 }
 
+export type AgentRefExpectation = AgentRef | AgentSession;
+
 export type RegistryEvent =
 	| { type: "registered"; ref: AgentRef }
 	| { type: "status_changed"; ref: AgentRef }
@@ -80,6 +82,10 @@ export class AgentRegistry {
 	readonly #refs = new Map<string, AgentRef>();
 	readonly #listeners = new Set<RegistryListener>();
 
+	#matchesExpected(ref: AgentRef, expected?: AgentRefExpectation): boolean {
+		return expected === undefined || ref === expected || ref.session === expected;
+	}
+
 	register(input: RegisterInput): AgentRef {
 		const now = Date.now();
 		const ref: AgentRef = {
@@ -98,15 +104,28 @@ export class AgentRegistry {
 		return ref;
 	}
 
-	setStatus(id: string, status: AgentStatus): void {
+	/**
+	 * 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.
+	 */
+	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;
+	}
+
+	setStatus(id: string, status: AgentStatus, expected?: AgentRefExpectation): boolean {
 		const ref = this.#refs.get(id);
-		if (!ref || ref.status === status) return;
+		if (!ref || !this.#matchesExpected(ref, expected)) return false;
+		if (ref.status === status) return true;
 		ref.status = status;
 		// Activity describes current work; it is meaningless once the agent
 		// leaves `running`, so drop it to avoid showing stale work in rosters.
 		if (status !== "running") ref.activity = undefined;
 		ref.lastActivity = Date.now();
 		this.#emit({ type: "status_changed", ref });
+		return true;
 	}
 
 	/**
@@ -133,25 +152,33 @@ export class AgentRegistry {
 		ref.activity = gist;
 	}
 
-	attachSession(id: string, session: AgentSession, sessionFile?: string | null): void {
+	attachSession(
+		id: string,
+		session: AgentSession,
+		sessionFile?: string | null,
+		expected?: AgentRefExpectation,
+	): boolean {
 		const ref = this.#refs.get(id);
-		if (!ref) return;
+		if (!ref || !this.#matchesExpected(ref, expected)) return false;
 		ref.session = session;
 		if (sessionFile !== undefined) ref.sessionFile = sessionFile;
 		ref.lastActivity = Date.now();
+		return true;
 	}
 
-	detachSession(id: string): void {
+	detachSession(id: string, expected?: AgentRefExpectation): boolean {
 		const ref = this.#refs.get(id);
-		if (!ref) return;
+		if (!ref || !this.#matchesExpected(ref, expected)) return false;
 		ref.session = null;
+		return true;
 	}
 
-	unregister(id: string): void {
+	unregister(id: string, expected?: AgentRefExpectation): boolean {
 		const ref = this.#refs.get(id);
-		if (!ref) return;
+		if (!ref || !this.#matchesExpected(ref, expected)) return false;
 		this.#refs.delete(id);
 		this.#emit({ type: "removed", ref });
+		return true;
 	}
 
 	get(id: string): AgentRef | undefined {

package/src/registry/persisted-agents.ts

@@ -1,22 +1,45 @@
 import * as fs from "node:fs";
 import * as path from "node:path";
 import { ADVISOR_TRANSCRIPT_FILENAME, isAdvisorTranscriptName } from "../advisor/transcript-recorder";
+import { SessionManager } from "../session/session-manager";
+import { persistedVibeChildIds } from "../vibe/runtime";
 import { type AgentRegistry, MAIN_AGENT_ID } from "./agent-registry";
 
+/**
+ * Child ids owned by the Vibe roster persisted in this session file. Vibe
+ * workers are revived through the Vibe registry's own journal, so the generic
+ * persisted-subagent scan must not register them as plain `sub` refs.
+ */
+async function readPersistedVibeChildIds(sessionFile: string): Promise<Set<string>> {
+	let sessionManager: SessionManager;
+	try {
+		sessionManager = await SessionManager.open(sessionFile, undefined, undefined, { suppressBreadcrumb: true });
+	} catch {
+		return new Set();
+	}
+	try {
+		return persistedVibeChildIds(sessionManager.getEntries());
+	} finally {
+		await sessionManager.close();
+	}
+}
+
 /** Register persisted subagent and advisor transcripts as parked registry refs. */
 export async function registerPersistedSubagents(
 	registry: AgentRegistry,
 	sessionFile: string | null | undefined,
 ): Promise<void> {
 	if (!sessionFile?.endsWith(".jsonl")) return;
+	const vibeOwnedIds = await readPersistedVibeChildIds(sessionFile);
 	const root = sessionFile.slice(0, -6);
-	await registerPersistedSubagentsFromDir(registry, root, undefined);
+	await registerPersistedSubagentsFromDir(registry, root, undefined, vibeOwnedIds);
 }
 
 async function registerPersistedSubagentsFromDir(
 	registry: AgentRegistry,
 	dir: string,
 	parentId: string | undefined,
+	vibeOwnedIds: ReadonlySet<string>,
 ): Promise<void> {
 	let entries: fs.Dirent[];
 	try {
@@ -58,6 +81,7 @@ async function registerPersistedSubagentsFromDir(
 			continue;
 		}
 		const id = entry.name.slice(0, -6);
+		if (vibeOwnedIds.has(id) && registry.get(id)?.sessionFile !== sessionFile) continue;
 		if (!registry.get(id)) {
 			registry.register({
 				id,
@@ -69,6 +93,6 @@ async function registerPersistedSubagentsFromDir(
 				status: "parked",
 			});
 		}
-		await registerPersistedSubagentsFromDir(registry, path.join(dir, id), id);
+		await registerPersistedSubagentsFromDir(registry, path.join(dir, id), id, vibeOwnedIds);
 	}
 }