From 4bc68b45e25dfe42bf4f2bc632d2fde665cdd868 Mon Sep 17 00:00:00 2001 From: roboomp Date: Tue, 14 Jul 2026 17:58:14 +0000 Subject: [PATCH 1/4] fix(coding-agent): persisted vibe sessions across restarts Vibe worker roster lived only in a process-local Map, so a resumed parent session started with an empty registry and vibe_send failed with "Unknown vibe session". Persist a versioned, parent-scoped lifecycle journal (spawn/turn/tombstone events), rehydrate validated idle workers through the persisted-subagent reviver on resume, and gate the flow with generation/CAS protection so stale finalizers cannot clobber a replacement worker. Killed transcripts stay readable but non-revivable; mode-exit commits tombstones atomically with the mode change and rolls back cleanly on storage failure. Ported from @mastertyko's fork branch fix/vibe-session-persistence. Fixes #5303 --- packages/coding-agent/CHANGELOG.md | 4 + packages/coding-agent/src/collab/host.ts | 5 +- .../src/modes/components/agent-hub.ts | 25 +- .../modes/controllers/command-controller.ts | 2 +- .../src/modes/interactive-mode.ts | 72 +- .../src/registry/agent-lifecycle.ts | 140 +- .../src/registry/agent-registry.ts | 43 +- packages/coding-agent/src/sdk.ts | 63 +- .../coding-agent/src/session/agent-session.ts | 30 + .../src/session/indexed-session-storage.ts | 11 +- .../src/session/redis-session-storage.ts | 62 +- .../src/session/session-manager.ts | 291 ++- .../src/slash-commands/builtin-registry.ts | 2 +- packages/coding-agent/src/task/executor.ts | 51 +- .../coding-agent/src/task/output-manager.ts | 28 +- .../coding-agent/src/task/persisted-revive.ts | 7 +- packages/coding-agent/src/tools/index.ts | 3 + packages/coding-agent/src/tools/vibe.ts | 3 +- packages/coding-agent/src/vibe/runtime.ts | 897 +++++++++- .../test/agent-hub-activate.test.ts | 62 + .../test/interactive-mode-vibe-toggle.test.ts | 137 +- .../test/registry/agent-lifecycle.test.ts | 47 + .../test/sdk-session-isolation.test.ts | 131 +- ...ession-manager-atomic-rewrite-race.test.ts | 276 ++- .../redis-session-storage-manager.test.ts | 32 + .../session/redis-session-storage.test.ts | 54 + .../test/task/executor-wall-clock.test.ts | 65 +- .../test/task/output-manager.test.ts | 26 +- .../test/vibe/vibe-runtime.test.ts | 1579 ++++++++++++++++- 29 files changed, 3919 insertions(+), 229 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 10f7883ed..4676cc148 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed persisted vibe workers disappearing or being replaced across graceful restarts, session switches, failed mode exits, and late cancelled initialization: resumable conversations now restore safely, mode exit atomically commits worker tombstones with the mode change and rolls back cleanly on storage failure, explicit kills tear workers down monotonically while repairing uncertain append tails, killed transcripts remain readable but non-revivable, and stale initializers cannot overwrite newer same-ID workers ([#5303](https://github.com/can1357/oh-my-pi/issues/5303) by [@mastertyko](https://github.com/mastertyko)). + ## [16.4.8] - 2026-07-12 ### Fixed diff --git a/packages/coding-agent/src/collab/host.ts b/packages/coding-agent/src/collab/host.ts index 74a021a80..5d9d103e7 100644 --- a/packages/coding-agent/src/collab/host.ts +++ b/packages/coding-agent/src/collab/host.ts @@ -614,10 +614,11 @@ export class CollabHost { case "kill": { const kill = async () => { const ref = AgentRegistry.global().get(agentId); - if (ref && ref.status === "running" && ref.session) { + if (!ref) return; + if (ref.status === "running" && ref.session) { await ref.session.abort({ reason: USER_INTERRUPT_LABEL }); } - await AgentLifecycleManager.global().release(agentId); + await AgentLifecycleManager.global().release(agentId, ref); }; kill().catch(fail); break; diff --git a/packages/coding-agent/src/modes/components/agent-hub.ts b/packages/coding-agent/src/modes/components/agent-hub.ts index 1c8413605..469584a06 100644 --- a/packages/coding-agent/src/modes/components/agent-hub.ts +++ b/packages/coding-agent/src/modes/components/agent-hub.ts @@ -25,7 +25,9 @@ import { IrcBus } from "../../irc/bus"; import { AgentLifecycleManager } from "../../registry/agent-lifecycle"; import { type AgentRef, AgentRegistry, type AgentStatus, MAIN_AGENT_ID } from "../../registry/agent-registry"; import { USER_INTERRUPT_LABEL } from "../../session/messages"; +import { SessionManager } from "../../session/session-manager"; import { replaceTabs, TRUNCATE_LENGTHS, truncateToWidth } from "../../tools/render-utils"; +import { persistedVibeChildIds } from "../../vibe/runtime"; import type { ObservableSession, SessionObserverRegistry } from "../session-observer-registry"; import { theme } from "../theme/theme"; import { matchesSelectDown, matchesSelectUp } from "../utils/keybinding-matchers"; @@ -69,19 +71,35 @@ function statusBadge(status: AgentStatus): string { } } +async function readPersistedVibeChildIds(sessionFile: string): Promise> { + 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(); + } +} + async function registerPersistedSubagents( registry: AgentRegistry, sessionFile: string | null | undefined, ): Promise { 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, ): Promise { let entries: fs.Dirent[]; try { @@ -123,6 +141,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, @@ -134,7 +153,7 @@ async function registerPersistedSubagentsFromDir( status: "parked", }); } - await registerPersistedSubagentsFromDir(registry, path.join(dir, id), id); + await registerPersistedSubagentsFromDir(registry, path.join(dir, id), id, vibeOwnedIds); } } @@ -607,7 +626,7 @@ export class AgentHubOverlayComponent extends Container { if (ref.status === "running" && ref.session) { await ref.session.abort({ reason: USER_INTERRUPT_LABEL }); } - await this.#lifecycle().release(ref.id); + await this.#lifecycle().release(ref.id, ref); } catch (error) { logger.warn("Agent hub: kill failed", { id: ref.id, error: String(error) }); this.#notice = error instanceof Error ? error.message : String(error); diff --git a/packages/coding-agent/src/modes/controllers/command-controller.ts b/packages/coding-agent/src/modes/controllers/command-controller.ts index ef6363d89..1657ccdf4 100644 --- a/packages/coding-agent/src/modes/controllers/command-controller.ts +++ b/packages/coding-agent/src/modes/controllers/command-controller.ts @@ -992,7 +992,7 @@ export class CommandController { } try { - await this.ctx.sessionManager.moveTo(resolvedPath); + await this.ctx.session.moveSession(resolvedPath); } catch (err) { this.ctx.showError(`Move failed: ${err instanceof Error ? err.message : String(err)}`); return; diff --git a/packages/coding-agent/src/modes/interactive-mode.ts b/packages/coding-agent/src/modes/interactive-mode.ts index c53bea59b..ab09d0ca6 100644 --- a/packages/coding-agent/src/modes/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive-mode.ts @@ -92,7 +92,7 @@ import planModeApprovedPrompt from "../prompts/system/plan-mode-approved.md" wit import planModeCompactInstructionsPrompt from "../prompts/system/plan-mode-compact-instructions.md" with { type: "text", }; -import { type AgentRegistry, MAIN_AGENT_ID } from "../registry/agent-registry"; +import type { AgentRegistry } from "../registry/agent-registry"; import { type AgentSession, type AgentSessionEvent, @@ -126,7 +126,7 @@ import { getEditorCommand, openInEditor } from "../utils/external-editor"; import { getSessionAccentAnsi, getSessionAccentHex } from "../utils/session-color"; import { messageHasDisplayableThinking } from "../utils/thinking-display"; import { popTerminalTitle, pushTerminalTitle, setSessionTerminalTitle } from "../utils/title-generator"; -import { VibeSessionRegistry } from "../vibe/runtime"; +import { type VibeOwnerScope, type VibeParentSession, VibeSessionRegistry } from "../vibe/runtime"; import type { AssistantMessageComponent } from "./components/assistant-message"; import type { BashExecutionComponent } from "./components/bash-execution"; import { ChatBlock, type ChatBlockHost } from "./components/chat-block"; @@ -532,6 +532,8 @@ export class InteractiveMode implements InteractiveModeContext { #planModePreviousTools: string[] | undefined; #goalModePreviousTools: string[] | undefined; #vibeModePreviousTools: string[] | undefined; + #vibeModeOwnerScope: VibeOwnerScope | undefined; + #vibeScopeSuspendedForSwitch = false; #goalContinuationTimer: NodeJS.Timeout | undefined; #goalTurnHadToolCalls = false; #goalContinuationTurnInFlight = false; @@ -975,6 +977,7 @@ export class InteractiveMode implements InteractiveModeContext { await this.initHooksAndCustomTools(); // Restore mode from session (e.g. plan mode on resume) + this.session.setSessionBeforeSwitchReconciler?.(() => this.#quiesceVibeForSessionSwitch()); this.session.setSessionSwitchReconciler?.(() => this.#reconcileModeFromSession({ preserveActiveGoal: true })); await this.#reconcileModeFromSession(); @@ -1955,6 +1958,24 @@ export class InteractiveMode implements InteractiveModeContext { this.ui.requestRender(); } + #vibeParentSession(): VibeParentSession { + return { + getAgentId: () => this.session.getAgentId() ?? null, + getSessionId: () => this.sessionManager.getSessionId(), + getSessionFile: () => this.sessionManager.getSessionFile() ?? null, + sessionManager: this.sessionManager, + asyncJobManager: this.session.asyncJobManager, + settings: this.session.settings, + }; + } + + async #quiesceVibeForSessionSwitch(): Promise { + const ownerScope = this.#vibeModeOwnerScope; + if (!this.vibeModeEnabled || !ownerScope) return; + await VibeSessionRegistry.global().suspendScope(ownerScope, this.session.asyncJobManager); + this.#vibeScopeSuspendedForSwitch = true; + } + #updateGoalModeStatus(): void { const status = this.goalModeEnabled || this.goalModePaused @@ -2092,7 +2113,10 @@ export class InteractiveMode implements InteractiveModeContext { } } - async #clearTransientModeState(): Promise { + async #clearTransientModeState(options?: { + preserveVibe?: boolean; + vibeScopeAlreadySuspended?: boolean; + }): Promise { if (this.planModeEnabled || this.planModePaused) { if (this.#planModePreviousTools !== undefined) { await this.session.setActiveToolsByName(this.#planModePreviousTools); @@ -2124,23 +2148,35 @@ export class InteractiveMode implements InteractiveModeContext { this.#updateGoalModeStatus(); } - if (this.vibeModeEnabled) { + if (this.vibeModeEnabled && !options?.preserveVibe) { + const ownerScope = this.#vibeModeOwnerScope; await this.session.deactivateVibeTools(this.#vibeModePreviousTools ?? []); this.session.setVibeModeState(undefined); this.vibeModeEnabled = false; this.#vibeModePreviousTools = undefined; - await VibeSessionRegistry.global().killAll( - this.session.getAgentId() ?? MAIN_AGENT_ID, - this.session.asyncJobManager, - ); + this.#vibeModeOwnerScope = undefined; + if (ownerScope && !options?.vibeScopeAlreadySuspended) { + await VibeSessionRegistry.global().suspendScope(ownerScope, this.session.asyncJobManager); + } this.#updateVibeModeStatus(); } } /** Reconcile mode state from session entries on resume/switch. */ async #reconcileModeFromSession(options?: { preserveActiveGoal?: boolean }): Promise { - await this.#clearTransientModeState(); + const vibeScopeAlreadySuspended = this.#vibeScopeSuspendedForSwitch; + this.#vibeScopeSuspendedForSwitch = false; const sessionContext = this.sessionManager.buildSessionContext(); + const vibeSession = this.#vibeParentSession(); + const targetVibeScope = VibeSessionRegistry.global().ownerScope(vibeSession); + const preserveVibe = + this.vibeModeEnabled && + sessionContext.mode === "vibe" && + this.#vibeModeOwnerScope?.ownerId === targetVibeScope.ownerId && + this.#vibeModeOwnerScope.parentSessionId === targetVibeScope.parentSessionId && + this.#vibeModeOwnerScope.parentSessionFile === targetVibeScope.parentSessionFile; + await this.#clearTransientModeState({ preserveVibe, vibeScopeAlreadySuspended }); + await VibeSessionRegistry.global().rehydrate(vibeSession); const goalEnabled = this.session.settings.get("goal.enabled"); if (!goalEnabled && (sessionContext.mode === "goal" || sessionContext.mode === "goal_paused")) { this.session.goalRuntime.clearAccounting(); @@ -2175,7 +2211,7 @@ export class InteractiveMode implements InteractiveModeContext { } this.session.goalRuntime.clearAccounting(); if (sessionContext.mode === "vibe") { - await this.#enterVibeMode(); + if (!preserveVibe) await this.#enterVibeMode({ persistModeChange: false }); return; } if (!this.session.settings.get("plan.enabled")) { @@ -2959,7 +2995,7 @@ export class InteractiveMode implements InteractiveModeContext { } } - async #enterVibeMode(): Promise { + async #enterVibeMode(options?: { persistModeChange?: boolean }): Promise { if (this.vibeModeEnabled) { return; } @@ -2972,9 +3008,13 @@ export class InteractiveMode implements InteractiveModeContext { return; } + const vibeRegistry = VibeSessionRegistry.global(); + const ownerScope = vibeRegistry.ownerScope(this.#vibeParentSession()); + vibeRegistry.activateScope(ownerScope); const previousTools = this.session.getActiveToolNames(); await this.session.activateVibeTools(["read"]); this.#vibeModePreviousTools = previousTools; + this.#vibeModeOwnerScope = ownerScope; this.vibeModeEnabled = true; // Suppress cache-miss marker on the next turn: vibe mode changes the // injected context, which predictably invalidates the cache. @@ -2984,7 +3024,7 @@ export class InteractiveMode implements InteractiveModeContext { await this.session.sendVibeModeContext({ deliverAs: "steer" }); } this.#updateVibeModeStatus(); - this.sessionManager.appendModeChange("vibe"); + if (options?.persistModeChange !== false) this.sessionManager.appendModeChange("vibe"); this.showStatus("Vibe mode enabled. You direct fast/good worker sessions; toolset is read + vibe tools."); } @@ -2992,17 +3032,15 @@ export class InteractiveMode implements InteractiveModeContext { if (!this.vibeModeEnabled) { return; } + const ownerScope = this.#vibeModeOwnerScope; + const killed = await VibeSessionRegistry.global().killAll(this.#vibeParentSession(), ownerScope); await this.session.deactivateVibeTools(this.#vibeModePreviousTools ?? []); this.session.setVibeModeState(undefined); this.vibeModeEnabled = false; this.#vibeModePreviousTools = undefined; + this.#vibeModeOwnerScope = undefined; this.lastAssistantUsage = undefined; - const killed = await VibeSessionRegistry.global().killAll( - this.session.getAgentId() ?? MAIN_AGENT_ID, - this.session.asyncJobManager, - ); this.#updateVibeModeStatus(); - this.sessionManager.appendModeChange("none"); this.showStatus( killed > 0 ? `Vibe mode disabled. Killed ${killed} worker session${killed === 1 ? "" : "s"}.` diff --git a/packages/coding-agent/src/registry/agent-lifecycle.ts b/packages/coding-agent/src/registry/agent-lifecycle.ts index 91678e2a7..f234c5de9 100644 --- a/packages/coding-agent/src/registry/agent-lifecycle.ts +++ b/packages/coding-agent/src/registry/agent-lifecycle.ts @@ -12,9 +12,15 @@ 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; +export type AgentReviver = (expected: AgentRef) => Promise; /** * Builds a reviver for a `parked` ref restored from disk (Agent Hub scan, @@ -33,11 +39,17 @@ export interface AdoptOptions { } interface AdoptedAgent { + ref: AgentRef; idleTtlMs: number; revive?: AgentReviver; timer?: NodeJS.Timeout; } +interface RevivingAgent { + ref: AgentRef; + promise: Promise; +} + export class AgentLifecycleManager { static #global: AgentLifecycleManager | undefined; @@ -67,10 +79,10 @@ export class AgentLifecycleManager { readonly #registry: AgentRegistry; readonly #adopted = new Map(); - /** Ids whose session is being disposed by {@link park} right now. */ - readonly #parking = new Set(); - /** In-flight revives, so concurrent {@link ensureLive} calls coalesce. */ - readonly #revivals = new Map>(); + /** Exact refs whose session is being disposed by {@link park} right now. */ + readonly #parking = new Map(); + /** In-flight revives, bound to the parked ref that initiated them. */ + readonly #revivals = new Map(); #unsubscribe: (() => void) | undefined; #persistedReviverFactory: PersistedSubagentReviverFactory | undefined; /** TTL applied when a cold-revived ref is adopted on demand. */ @@ -96,27 +108,32 @@ 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). */ - 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); + has(id: string, expected?: AgentRefExpectation): boolean { + const adopted = this.#adopted.get(id); + return Boolean( + adopted && (expected === undefined || adopted.ref === expected || adopted.ref.session === expected), + ); } /** True while {@link park} is disposing this agent's session (lets dispose hooks distinguish park from teardown). */ - isParking(id: string): boolean { - return this.#parking.has(id); + isParking(id: string, expected?: AgentRefExpectation): boolean { + const ref = this.#parking.get(id); + return Boolean(ref && (expected === undefined || ref === expected || ref.session === expected)); } /** @@ -125,24 +142,24 @@ export class AgentLifecycleManager { */ async park(id: string): Promise { const adopted = this.#adopted.get(id); - if (!adopted) return; const ref = this.#registry.get(id); - if (!ref?.session) return; + if (!adopted || adopted.ref !== ref || !ref.session) return; + const liveSession = ref.session; if (adopted.timer) { clearTimeout(adopted.timer); adopted.timer = undefined; } - this.#parking.add(id); + this.#parking.set(id, ref); try { try { - await ref.session.dispose(); + await liveSession.dispose(); } catch (error) { logger.warn("AgentLifecycleManager.park: session dispose failed", { id, error: String(error) }); } - this.#registry.detachSession(id); - this.#registry.setStatus(id, "parked"); + this.#registry.detachSession(id, ref); + this.#registry.setStatus(id, "parked", ref); } finally { - this.#parking.delete(id); + if (this.#parking.get(id) === ref) this.#parking.delete(id); } } @@ -160,13 +177,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); } } @@ -178,63 +196,97 @@ export class AgentLifecycleManager { * when the agent is not revivable or no reviver can be produced. */ async #resolveAndRevive(id: string, ref: AgentRef): Promise { - 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 { + async release(id: string, expected?: AgentRefExpectation): Promise { const adopted = this.#adopted.get(id); - clearTimeout(adopted?.timer); - this.#adopted.delete(id); - const ref = this.#registry.get(id); - if (ref?.session) { + 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 liveSession = ref.session; + if (liveSession) { try { - await ref.session.dispose(); + await liveSession.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). */ async dispose(): Promise { this.#unsubscribe?.(); this.#unsubscribe = undefined; - const ids = [...this.#adopted.keys()]; - await Promise.all(ids.map(id => this.release(id))); + const adopted = [...this.#adopted.values()]; + await Promise.all(adopted.map(entry => this.release(entry.ref.id, entry.ref))); this.#revivals.clear(); this.#parking.clear(); this.#persistedReviverFactory = undefined; } - async #revive(id: string, revive: AgentReviver, sessionFile: string | null): Promise { - const session = await revive(); - this.#registry.attachSession(id, session, sessionFile); - // Emits status_changed → "idle", which re-arms the TTL timer below. - this.#registry.setStatus(id, "idle"); + 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 (!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; + if (!this.#registry.setStatus(id, "idle", liveRef)) { + await session.dispose(); + throw new Error(`Agent "${id}" changed before its persisted session became idle.`); + } return session; } @@ -251,7 +303,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); diff --git a/packages/coding-agent/src/registry/agent-registry.ts b/packages/coding-agent/src/registry/agent-registry.ts index 84849bc9a..69972d38c 100644 --- a/packages/coding-agent/src/registry/agent-registry.ts +++ b/packages/coding-agent/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(); readonly #listeners = new Set(); + #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 { diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index d03b414a4..0a94fc53a 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/packages/coding-agent/src/sdk.ts @@ -99,7 +99,7 @@ import type { MnemopiSessionState } from "./mnemopi/state"; import asyncResultTemplate from "./prompts/tools/async-result.md" with { type: "text" }; import lateDiagnosticTemplate from "./prompts/tools/lsp-late-diagnostic.md" with { type: "text" }; import { AgentLifecycleManager } from "./registry/agent-lifecycle"; -import { AgentRegistry, MAIN_AGENT_ID } from "./registry/agent-registry"; +import { type AgentRef, AgentRegistry, MAIN_AGENT_ID } from "./registry/agent-registry"; import { collectEnvSecrets, deobfuscateSessionContext, @@ -199,6 +199,7 @@ import { ttsTool } from "./tools/tts"; import { resolveActiveRepoContext } from "./utils/active-repo-context"; import { EventBus } from "./utils/event-bus"; import { buildNamedToolChoice } from "./utils/tool-choice"; +import { VibeSessionRegistry } from "./vibe/runtime"; import { buildWorkspaceTree, type WorkspaceTree } from "./workspace-tree"; type AsyncResultEntry = { @@ -522,6 +523,13 @@ export interface CreateAgentSessionOptions { agentDisplayName?: string; /** Optional shared agent registry for IRC routing. Default: AgentRegistry.global(). */ agentRegistry?: AgentRegistry; + /** + * Registry generation authorized for this creation. `null` requires the id + * to be absent; an AgentRef allows a parked revival to reuse only that ref. + * Undefined preserves legacy unconditional registration for external SDK callers. + * @internal + */ + expectedAgentRef?: AgentRef | null; /** Parent task ID prefix for nested artifact naming (e.g., "Extensions") */ parentTaskPrefix?: string; /** @@ -1525,15 +1533,18 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} const resolvedAgentDisplayName = options.agentDisplayName ?? ((options.taskDepth ?? 0) > 0 || options.parentTaskPrefix ? "sub" : "main"); const agentKind = (options.taskDepth ?? 0) > 0 || options.parentTaskPrefix ? ("sub" as const) : ("main" as const); + let registeredAgentRef: AgentRef | undefined; /** * Forget the agent ref on teardown — unless the agent is being parked (or is * already parked). Parking disposes the session but keeps the ref addressable * (history://, revive); only process teardown / explicit kill unregisters. */ const unregisterUnlessParked = (): void => { - if (agentRegistry.get(resolvedAgentId)?.status === "parked") return; - if (AgentLifecycleManager.global().isParking(resolvedAgentId)) return; - agentRegistry.unregister(resolvedAgentId); + const ref = registeredAgentRef; + if (!ref || agentRegistry.get(resolvedAgentId) !== ref) return; + if (ref.status === "parked") return; + if (AgentLifecycleManager.global().isParking(resolvedAgentId, ref)) return; + agentRegistry.unregister(resolvedAgentId, ref); }; const evalKernelOwnerId = `agent-session:${Snowflake.next()}`; @@ -1577,6 +1588,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} requireYieldTool: options.requireYieldTool, taskDepth: options.taskDepth ?? 0, getSessionFile: () => sessionManager.getSessionFile() ?? null, + sessionManager, getEvalKernelOwnerId: () => evalKernelOwnerId, getEvalSessionId: () => session?.getEvalSessionId() ?? options.parentEvalSessionId ?? defaultEvalSessionId(toolSession), @@ -2590,16 +2602,25 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} // so that subagents launched in the same parallel batch can see each other in // their initial `# IRC Peers` block (rendered inside `rebuildSystemPrompt`). // The session reference is attached after construction below. - agentRegistry.register({ + const registrationInput = { id: resolvedAgentId, displayName: resolvedAgentDisplayName, kind: agentKind, parentId: options.parentAgentId, session: null, sessionFile: sessionManager.getSessionFile() ?? null, - status: "running", - }); - hasRegistered = true; + status: "running" as const, + }; + registeredAgentRef = + options.expectedAgentRef === undefined + ? agentRegistry.register(registrationInput) + : agentRegistry.registerIfAvailable(registrationInput, options.expectedAgentRef); + if (!registeredAgentRef) { + throw new Error(`Agent "${resolvedAgentId}" is already owned by another session generation.`); + } + // A reused parked ref remains parked until the new AgentSession is fully + // constructed and attached. Startup failure therefore leaves it revivable. + hasRegistered = options.expectedAgentRef === undefined || options.expectedAgentRef === null; setActiveToolNames(initialToolNames); const { systemPrompt } = await logger.time( @@ -2949,7 +2970,19 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} // Attach the live session to the pre-registered ref so peers can route IRC // messages here. Refresh sessionFile in case it was unavailable at pre-register // time. The dispose wrapper below unregisters on teardown (unless parked). - agentRegistry.attachSession(resolvedAgentId, session, sessionManager.getSessionFile() ?? null); + if ( + !registeredAgentRef || + !agentRegistry.attachSession( + resolvedAgentId, + session, + sessionManager.getSessionFile() ?? null, + registeredAgentRef, + ) || + !agentRegistry.setStatus(resolvedAgentId, "running", registeredAgentRef) + ) { + throw new Error(`Agent "${resolvedAgentId}" was replaced during session initialization.`); + } + hasRegistered = true; { const originalDispose = session.dispose.bind(session); session.dispose = async () => { @@ -2963,6 +2996,17 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} // adopted subagent sessions, revivers. Tear it down while shared // resources (kernels, MCP, LSP) are still live. Subagent disposal // must NOT touch the global lifecycle. + const vibeRegistry = VibeSessionRegistry.global(); + const vibeParentSession = { + getAgentId: () => resolvedAgentId, + getSessionId: () => sessionManager.getSessionId(), + getSessionFile: () => sessionManager.getSessionFile() ?? null, + sessionManager, + asyncJobManager: scopedAsyncJobManager, + settings, + getActiveModelString, + }; + await vibeRegistry.suspendScope(vibeRegistry.ownerScope(vibeParentSession), scopedAsyncJobManager); await AgentLifecycleManager.global().dispose(); } await originalDispose(); @@ -3158,6 +3202,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} try { if (hasSession) { await session.dispose(); + if (hasRegistered) unregisterUnlessParked(); } else { if (hasRegistered) unregisterUnlessParked(); if (asyncJobManager) { diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 31d6b1383..cc0cc4a8f 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -3092,6 +3092,12 @@ export class AgentSession { this.#standingResolveHandler = handler ?? undefined; } + #sessionBeforeSwitchReconciler: (() => Promise) | undefined; + + setSessionBeforeSwitchReconciler(reconciler: (() => Promise) | null): void { + this.#sessionBeforeSwitchReconciler = reconciler ?? undefined; + } + #sessionSwitchReconciler: (() => Promise) | undefined; setSessionSwitchReconciler(reconciler: (() => Promise) | null): void { @@ -7161,6 +7167,12 @@ export class AgentSession { this.#vibeModeState = state; } + #assertVibeSessionTransitionAllowed(action: string): void { + if (this.#vibeModeState?.enabled) { + throw new Error(`Cannot ${action} while vibe mode is active. Exit vibe mode first.`); + } + } + get goalRuntime(): GoalRuntime { return this.#goalRuntime; } @@ -8858,6 +8870,7 @@ export class AgentSession { * @returns true if completed, false if cancelled by hook */ async newSession(options?: NewSessionOptions): Promise { + this.#assertVibeSessionTransitionAllowed("start a new session"); const previousSessionFile = this.sessionFile; const nextDiscoverySessionToolNames = this.#mcpDiscoveryEnabled ? [ @@ -8964,6 +8977,7 @@ export class AgentSession { * @returns true if completed, false if cancelled by hook or not persisting */ async fork(): Promise { + this.#assertVibeSessionTransitionAllowed("fork the session"); const previousSessionFile = this.sessionFile; // Emit session_before_switch event with reason "fork" (can be cancelled) @@ -9026,6 +9040,12 @@ export class AgentSession { return true; } + /** Move the active session and artifacts after enforcing mode transition invariants. */ + async moveSession(newCwd: string, targetSessionDir?: string): Promise { + this.#assertVibeSessionTransitionAllowed("move the session"); + await this.sessionManager.moveTo(newCwd, targetSessionDir); + } + // ========================================================================= // Model Management // ========================================================================= @@ -10221,6 +10241,7 @@ export class AgentSession { * @returns The handoff document text, or undefined if cancelled/failed */ async handoff(customInstructions?: string, options?: SessionHandoffOptions): Promise { + this.#assertVibeSessionTransitionAllowed("handoff to a new session"); const entries = this.sessionManager.getBranch(); const messageCount = entries.filter(e => e.type === "message").length; @@ -15019,6 +15040,7 @@ export class AgentSession { this.#disconnectFromAgent(); await this.abort({ goalReason: "internal" }); + await this.#sessionBeforeSwitchReconciler?.(); // Flush pending writes before switching so restore snapshots reflect committed state. await this.sessionManager.flush(); @@ -15244,6 +15266,14 @@ export class AgentSession { this.#syncTodoPhasesFromBranch(); this.#resetAllAdvisorRuntimes(); this.#reconnectToAgent(); + try { + await this.#sessionSwitchReconciler?.(); + } catch (reconcileError) { + logger.warn("Failed to reconcile session mode after switch rollback", { + targetSessionFile: sessionPath, + error: String(reconcileError), + }); + } if (restoreMcpError) { throw restoreMcpError; } diff --git a/packages/coding-agent/src/session/indexed-session-storage.ts b/packages/coding-agent/src/session/indexed-session-storage.ts index b620765b8..4786f4229 100644 --- a/packages/coding-agent/src/session/indexed-session-storage.ts +++ b/packages/coding-agent/src/session/indexed-session-storage.ts @@ -270,8 +270,15 @@ export class IndexedSessionStorage implements SessionStorage { { trackDrain: false }, ); } catch (err) { - this.#restoreIndex(path, previous); - throw toError(err); + const error = toError(err); + try { + if ((await this.#backend.readFull(path)) === content) return; + } catch { + // Preserve the original write failure; verification was unavailable. + } + const current = this.#index.get(path); + if (current?.mtimeMs === mtimeMs) this.#restoreIndex(path, previous); + throw error; } } diff --git a/packages/coding-agent/src/session/redis-session-storage.ts b/packages/coding-agent/src/session/redis-session-storage.ts index 54f6761bf..e5200f1d4 100644 --- a/packages/coding-agent/src/session/redis-session-storage.ts +++ b/packages/coding-agent/src/session/redis-session-storage.ts @@ -13,6 +13,7 @@ import type { SessionTitleUpdate } from "./session-title-slot"; * without dragging the entire Bun typings into this module. */ export interface RedisSessionStorageClient { + send(command: string, args: string[]): Promise; get(key: string): Promise; getrange(key: string, start: number, end: number): Promise; strlen(key: string): Promise; @@ -45,6 +46,26 @@ export interface RedisSessionStorageOptions { const DEFAULT_PREFIX = "omp:sessions:"; const DEFAULT_SCAN_COUNT = 500; +const WRITE_FULL_SCRIPT = `-- OMP_WRITE_FULL +redis.call("SET", KEYS[1], ARGV[1]) +redis.call("HSET", KEYS[2], ARGV[2], ARGV[3]) +if ARGV[4] == "1" then + redis.call("HSET", KEYS[3], ARGV[2], ARGV[5]) +else + redis.call("HDEL", KEYS[3], ARGV[2]) +end +return 1`; + +const APPEND_SCRIPT = `-- OMP_APPEND +local size = redis.call("APPEND", KEYS[1], ARGV[1]) +redis.call("HSET", KEYS[2], ARGV[2], ARGV[3]) +return size`; + +const UPDATE_TITLE_SCRIPT = `-- OMP_UPDATE_TITLE +redis.call("HSET", KEYS[1], ARGV[1], ARGV[2]) +redis.call("HSET", KEYS[2], ARGV[1], ARGV[3]) +return 1`; + function encodeTitleMeta(title: SessionTitleUpdate): string { return JSON.stringify(title); } @@ -155,23 +176,42 @@ class RedisSessionStorageBackend implements SessionStorageBackend { } async writeFull(path: string, content: string, mtimeMs: number, title?: SessionTitleUpdate): Promise { - await this.#client.set(this.#fileKey(path), content); - await this.#client.hset(this.#metaKey(), path, String(mtimeMs)); - if (title) { - await this.#client.hset(this.#titleMetaKey(), path, encodeTitleMeta(title)); - } else { - await this.#client.hdel(this.#titleMetaKey(), path); - } + await this.#client.send("EVAL", [ + WRITE_FULL_SCRIPT, + "3", + this.#fileKey(path), + this.#metaKey(), + this.#titleMetaKey(), + content, + path, + String(mtimeMs), + title ? "1" : "0", + title ? encodeTitleMeta(title) : "", + ]); } async append(path: string, line: string, mtimeMs: number): Promise { - await this.#client.append(this.#fileKey(path), line); - await this.#client.hset(this.#metaKey(), path, String(mtimeMs)); + await this.#client.send("EVAL", [ + APPEND_SCRIPT, + "2", + this.#fileKey(path), + this.#metaKey(), + line, + path, + String(mtimeMs), + ]); } async updateSessionTitle(path: string, title: SessionTitleUpdate, mtimeMs: number): Promise { - await this.#client.hset(this.#metaKey(), path, String(mtimeMs)); - await this.#client.hset(this.#titleMetaKey(), path, encodeTitleMeta(title)); + await this.#client.send("EVAL", [ + UPDATE_TITLE_SCRIPT, + "2", + this.#metaKey(), + this.#titleMetaKey(), + path, + String(mtimeMs), + encodeTitleMeta(title), + ]); } async truncate(path: string, mtimeMs: number): Promise { diff --git a/packages/coding-agent/src/session/session-manager.ts b/packages/coding-agent/src/session/session-manager.ts index be61be344..cd134c63f 100644 --- a/packages/coding-agent/src/session/session-manager.ts +++ b/packages/coding-agent/src/session/session-manager.ts @@ -365,6 +365,34 @@ interface DiskQueueOptions { epoch?: number; } +interface AtomicEntryBatch { + collecting: boolean; + entryIds: Set; + deferredNotifications: SessionEntry[]; + preBatchLeafId: string | null; + externalLeafChanged: boolean; + externalLeafId: string | null; +} + +/** + * The storage may have published a write that rejected, and an authoritative + * repair could not be proven durable. Callers must fail closed until recovery. + */ +export class SessionPersistenceIndeterminateError extends AggregateError { + readonly operationError: Error; + readonly recoveryErrors: readonly Error[]; + + constructor(operationError: Error, recoveryErrors: readonly Error[]) { + super( + [operationError, ...recoveryErrors], + `Session persistence is indeterminate after "${operationError.message}" and authoritative repair failed.`, + ); + this.name = "SessionPersistenceIndeterminateError"; + this.operationError = operationError; + this.recoveryErrors = [...recoveryErrors]; + } +} + /** * Stores and navigates an append-only conversation journal. * @@ -425,6 +453,10 @@ export class SessionManager { #diskTail: Promise = Promise.resolve(); #diskFailure: Error | undefined; #diskFailureLogged = false; + /** FIFO reservation for atomic batches and authoritative recovery. */ + #atomicPersistenceTail: Promise = Promise.resolve(); + /** Observer notifications withheld until their entries are proven durable. */ + #pendingDurabilityNotifications: SessionEntry[] = []; /** Bumped on every sync rewrite / chain reset so stale queued tasks become no-ops. */ #diskEpoch = 0; /** @@ -438,6 +470,8 @@ export class SessionManager { #atomicRewriteFenceEpoch: number | null = null; /** Set by synchronous appends that land while an atomic replacement is active. */ #atomicRewriteDirty = false; + /** Atomic entry batch currently staged for a full-file commit. */ + #atomicEntryBatch: AtomicEntryBatch | undefined; #artifactManager: ArtifactManager | null = null; #artifactManagerSessionFile: string | null = null; @@ -500,6 +534,18 @@ export class SessionManager { return reported; } + async #withAtomicPersistenceLock(operation: () => Promise): Promise { + const predecessor = this.#atomicPersistenceTail; + const turn = Promise.withResolvers(); + this.#atomicPersistenceTail = predecessor.catch(() => undefined).then(() => turn.promise); + await predecessor.catch(() => undefined); + try { + return await operation(); + } finally { + turn.resolve(); + } + } + async #drainAndCloseWriter(): Promise { try { await this.#scheduleDiskWork( @@ -527,6 +573,117 @@ export class SessionManager { await writer.close(); } + #latchIndeterminate(operationError: Error, recoveryErrors: readonly Error[]): SessionPersistenceIndeterminateError { + const error = new SessionPersistenceIndeterminateError(operationError, recoveryErrors); + this.#diskFailure = error; + if (!this.#diskFailureLogged) { + this.#diskFailureLogged = true; + logger.error("Session persistence became indeterminate.", { + sessionFile: this.#sessionFile, + error: error.message, + }); + } + return error; + } + + #notifyDurableEntries(entries: readonly SessionEntry[] = []): void { + const notifications = [...this.#pendingDurabilityNotifications, ...entries]; + this.#pendingDurabilityNotifications = []; + const seen = new Set(); + for (const entry of notifications) { + if (seen.has(entry.id)) continue; + seen.add(entry.id); + this.#notifyEntryAppended(entry); + } + } + + async #authoritativelyRewriteCurrentStateLocked(operationError: Error): Promise { + if (!this.#persist || !this.#sessionFile) return; + const previousDiskTail = this.#diskTail; + const writer = this.#writer; + this.#diskEpoch++; + const epoch = this.#diskEpoch; + this.#writer = undefined; + this.#diskTail = Promise.resolve(); + this.#forceFileCreation = true; + this.#fileIsCurrent = false; + this.#rewriteRequired = true; + this.#atomicRewriteFenceEpoch = epoch; + if (!this.#diskFailure) this.#diskFailure = operationError; + try { + await previousDiskTail.catch(() => undefined); + let closeError: Error | undefined; + if (writer) { + try { + await writer.close(); + } catch (error) { + closeError = toError(error); + } + } + let drainError: Error | undefined; + try { + await this.#storage.drain(); + } catch (error) { + drainError = toError(error); + } + if (writer?.isOpen()) { + throw this.#latchIndeterminate(operationError, [ + closeError ?? new Error("Failed to close session writer before authoritative repair."), + ...(drainError ? [drainError] : []), + ]); + } + + do { + this.#atomicRewriteDirty = false; + const sessionFile = this.#sessionFile; + if (!sessionFile) { + throw this.#latchIndeterminate(operationError, [ + new Error("Session file disappeared during authoritative repair."), + ]); + } + const body = this.#fileBody(); + try { + await this.#storage.writeTextAtomic(sessionFile, body, { + commitGuard: () => this.#diskEpoch === epoch, + }); + } catch (error) { + const recoveryErrors = [toError(error)]; + try { + await this.#storage.drain(); + } catch (drainFailure) { + recoveryErrors.push(toError(drainFailure)); + } + let actual: string; + try { + actual = await this.#storage.readText(sessionFile); + } catch (readFailure) { + recoveryErrors.push(toError(readFailure)); + throw this.#latchIndeterminate(operationError, recoveryErrors); + } + if (actual !== body) { + recoveryErrors.push(new Error("Authoritative session repair did not match durable storage.")); + throw this.#latchIndeterminate(operationError, recoveryErrors); + } + } + if (this.#diskEpoch !== epoch) { + throw this.#latchIndeterminate(operationError, [ + new Error("Authoritative session repair was superseded before verification."), + ]); + } + } while (this.#atomicRewriteDirty); + + this.#fileIsCurrent = true; + this.#rewriteRequired = false; + this.#hasTitleSlot = true; + this.#clearDiskError(); + } catch (error) { + if (error instanceof SessionPersistenceIndeterminateError) throw error; + throw this.#latchIndeterminate(operationError, [toError(error)]); + } finally { + if (this.#atomicRewriteFenceEpoch === epoch) this.#atomicRewriteFenceEpoch = null; + } + } + #appendWriter(): SessionStorageWriter { if (!this.#sessionFile) throw new Error("Cannot open a session writer before a session file exists"); @@ -671,6 +828,12 @@ export class SessionManager { #appendToSessionFile(entry: SessionEntry): void { if (!this.#persist || !this.#sessionFile) return; + if (this.#atomicEntryBatch) { + this.#fileIsCurrent = false; + this.#rewriteRequired = true; + this.#atomicRewriteDirty = true; + return; + } if (this.#diskFailure) throw this.#diskFailure; // Lazy gate: a brand-new session is not written until it has an assistant @@ -834,11 +997,44 @@ export class SessionManager { }; } + #setLeaf(id: string | null): void { + this.#index.setLeaf(id); + const batch = this.#atomicEntryBatch; + if (batch && !batch.collecting) { + batch.externalLeafChanged = true; + batch.externalLeafId = id; + } + } + #recordEntry(entry: SessionEntry): void { this.#entries.push(entry); this.#index.insert(entry); + const batch = this.#atomicEntryBatch; + if (batch?.collecting) batch.entryIds.add(entry.id); + if (batch && !batch.collecting) { + batch.externalLeafChanged = true; + batch.externalLeafId = entry.id; + } this.#appendToSessionFile(entry); - this.#notifyEntryAppended(entry); + if (batch) batch.deferredNotifications.push(entry); + else this.#notifyEntryAppended(entry); + } + + #rollbackAtomicEntryBatch(batch: AtomicEntryBatch): void { + const retainedAncestor = (id: string | null): string | null => { + const seen = new Set(); + while (id && batch.entryIds.has(id) && !seen.has(id)) { + seen.add(id); + id = this.#index.get(id)?.parentId ?? null; + } + return id; + }; + const retained = this.#entries.filter(entry => !batch.entryIds.has(entry.id)); + for (const entry of retained) entry.parentId = retainedAncestor(entry.parentId); + const restoredLeaf = retainedAncestor(batch.externalLeafChanged ? batch.externalLeafId : batch.preBatchLeafId); + this.#entries = retained; + this.#index.rebuild(retained); + this.#index.setLeaf(restoredLeaf && this.#index.has(restoredLeaf) ? restoredLeaf : null); } #draftPath(): string | null { @@ -1181,6 +1377,92 @@ export class SessionManager { await this.#rewriteAtomically(); } + /** + * Stage a synchronous group of entry appends and publish the resulting full + * journal with one atomic replace. A failed publish removes only the staged + * entries, preserves/reparents entries appended concurrently, restores the + * prior durable file view, and clears the failed writer latch for retry. + * + * The callback MUST be synchronous. + */ + appendEntriesAtomically(append: () => T): Promise { + return this.#withAtomicPersistenceLock(() => this.#appendEntriesAtomicallyLocked(append)); + } + + async #appendEntriesAtomicallyLocked(append: () => T): Promise { + if (!this.#persist || !this.#sessionFile) return append(); + if (this.#atomicEntryBatch) throw new Error("Atomic persistence lock ownership was violated."); + try { + await this.ensureOnDisk(); + await this.flush(); + } catch (error) { + const operationError = toError(error); + await this.#authoritativelyRewriteCurrentStateLocked(operationError); + this.#notifyDurableEntries(); + throw error; + } + + const batch: AtomicEntryBatch = { + collecting: true, + entryIds: new Set(), + deferredNotifications: [], + preBatchLeafId: this.#index.leafId(), + externalLeafChanged: false, + externalLeafId: null, + }; + this.#atomicEntryBatch = batch; + let result!: T; + try { + try { + result = append(); + } finally { + batch.collecting = false; + } + await this.#rewriteAtomically(); + if (!this.#fileIsCurrent || this.#rewriteRequired) { + throw new Error("Atomic session batch was superseded before commit."); + } + this.#atomicEntryBatch = undefined; + this.#notifyDurableEntries(batch.deferredNotifications); + return result; + } catch (error) { + batch.collecting = false; + const operationError = toError(error); + this.#rollbackAtomicEntryBatch(batch); + try { + await this.#authoritativelyRewriteCurrentStateLocked(operationError); + } catch (repairError) { + const retainedNotifications = batch.deferredNotifications.filter(entry => !batch.entryIds.has(entry.id)); + this.#pendingDurabilityNotifications.push(...retainedNotifications); + this.#atomicEntryBatch = undefined; + this.#fileIsCurrent = false; + this.#rewriteRequired = true; + if (repairError instanceof SessionPersistenceIndeterminateError) throw repairError; + throw this.#latchIndeterminate(operationError, [toError(repairError)]); + } + const retainedNotifications = batch.deferredNotifications.filter(entry => !batch.entryIds.has(entry.id)); + this.#atomicEntryBatch = undefined; + this.#notifyDurableEntries(retainedNotifications); + throw error; + } + } + + /** + * Replace an uncertain append tail with the authoritative in-memory journal. + * Callers must only use this for monotonic recovery where every retained + * entry remains intended (for example, an explicit terminal tombstone). + */ + recoverPersistenceFromCurrentState(): Promise { + return this.#withAtomicPersistenceLock(async () => { + if (!this.#persist || !this.#sessionFile) return; + if (this.#atomicEntryBatch) throw new Error("Atomic persistence lock ownership was violated."); + const operationError = + this.#diskFailure ?? new Error("Authoritative session persistence recovery was requested."); + await this.#authoritativelyRewriteCurrentStateLocked(operationError); + this.#notifyDurableEntries(); + }); + } + /** Flush pending writes. Call before switching sessions or on shutdown. */ async flush(): Promise { if (!this.#persist || !this.#sessionFile) return; @@ -1201,6 +1483,7 @@ export class SessionManager { */ flushSync(): void { if (!this.#persist || !this.#sessionFile) return; + if (this.#atomicEntryBatch) throw new Error("Cannot synchronously flush during an atomic session batch."); if (this.#diskFailure) throw this.#diskFailure; if (this.#fileIsCurrent && !this.#rewriteRequired) { const writerError = this.#writer?.getError(); @@ -1726,19 +2009,19 @@ export class SessionManager { */ branch(branchFromId: string): void { if (!this.#index.has(branchFromId)) throw new Error(`Entry ${branchFromId} not found`); - this.#index.setLeaf(branchFromId); + this.#setLeaf(branchFromId); } /** Reset the leaf to null so the next append creates a new root entry. */ resetLeaf(): void { - this.#index.setLeaf(null); + this.#setLeaf(null); } /** Like branch(), but also records a branch_summary of the abandoned path. */ branchWithSummary(branchFromId: string | null, summary: string, details?: unknown, fromExtension?: boolean): string { if (branchFromId !== null && !this.#index.has(branchFromId)) throw new Error(`Entry ${branchFromId} not found`); - this.#index.setLeaf(branchFromId); + this.#setLeaf(branchFromId); const entry: BranchSummaryEntry = { type: "branch_summary", id: generateId(this.#index), diff --git a/packages/coding-agent/src/slash-commands/builtin-registry.ts b/packages/coding-agent/src/slash-commands/builtin-registry.ts index e5a5711b9..431b767c9 100644 --- a/packages/coding-agent/src/slash-commands/builtin-registry.ts +++ b/packages/coding-agent/src/slash-commands/builtin-registry.ts @@ -1675,7 +1675,7 @@ const BUILTIN_SLASH_COMMAND_REGISTRY: ReadonlyArray = [ return usage(`Directory does not exist: ${resolvedPath}`, runtime); } try { - await runtime.sessionManager.moveTo(resolvedPath); + await runtime.session.moveSession(resolvedPath); } catch (err) { return usage(`Move failed: ${errorMessage(err)}`, runtime); } diff --git a/packages/coding-agent/src/task/executor.ts b/packages/coding-agent/src/task/executor.ts index 62744a633..7c9ae9a61 100644 --- a/packages/coding-agent/src/task/executor.ts +++ b/packages/coding-agent/src/task/executor.ts @@ -33,7 +33,7 @@ import type { MCPManager } from "../mcp/manager"; import type { MnemopiSessionState } from "../mnemopi/state"; import subagentSystemPromptTemplate from "../prompts/system/subagent-system-prompt.md" with { type: "text" }; import submitReminderTemplate from "../prompts/system/subagent-yield-reminder.md" with { type: "text" }; -import { AgentLifecycleManager } from "../registry/agent-lifecycle"; +import { AgentLifecycleManager, type AgentReviver } from "../registry/agent-lifecycle"; import { AgentRegistry } from "../registry/agent-registry"; import { type CreateAgentSessionOptions, createAgentSession, discoverAuthStorage } from "../sdk"; import type { AgentSession, AgentSessionEvent } from "../session/agent-session"; @@ -1945,9 +1945,11 @@ export async function finalizeSubagentLifecycle(args: { keepAlive: boolean; isolated: boolean; agentIdleTtlMs: number; - reviveSession: (() => Promise) | null; + reviveSession: AgentReviver | null; }): Promise { const registry = AgentRegistry.global(); + const ref = registry.get(args.id); + const ownsRef = Boolean(ref && ref.session === args.session); const disposeSession = async (): Promise => { try { await untilAborted(AbortSignal.timeout(5000), () => args.session.dispose()); @@ -1962,7 +1964,7 @@ export async function finalizeSubagentLifecycle(args: { const resumableAbort = args.abortKind === "budget" && args.keepAlive && !args.isolated && args.reviveSession !== null; if (args.aborted && !resumableAbort) { - registry.setStatus(args.id, "aborted"); + if (ref && ownsRef) registry.setStatus(args.id, "aborted", ref); await disposeSession(); return; } @@ -1970,7 +1972,7 @@ export async function finalizeSubagentLifecycle(args: { if (!args.keepAlive) { // One-shot helper: dispose and unregister. No IRC, no revival. await disposeSession(); - registry.unregister(args.id); + if (ref && ownsRef) registry.unregister(args.id, ref); return; } @@ -1980,19 +1982,26 @@ export async function finalizeSubagentLifecycle(args: { // transcript stays reachable (history://), but ensureLive will throw. // Status must flip to "parked" before dispose so the sdk dispose // wrapper skips unregister. - registry.setStatus(args.id, "parked"); + if (ref && ownsRef) registry.setStatus(args.id, "parked", ref); await disposeSession(); - registry.detachSession(args.id); + if (ref && ownsRef) registry.detachSession(args.id, ref); return; } // Keep-alive: finished and failed subagents both stay interrogable. // The lifecycle manager owns idle-TTL parking + revival from here on. - registry.setStatus(args.id, "idle"); - AgentLifecycleManager.global().adopt(args.id, { - idleTtlMs: args.agentIdleTtlMs, - revive: args.reviveSession ?? undefined, - }); + if (!ref || !ownsRef || !registry.setStatus(args.id, "idle", ref)) { + await disposeSession(); + return; + } + AgentLifecycleManager.global().adopt( + args.id, + { + idleTtlMs: args.agentIdleTtlMs, + revive: args.reviveSession ?? undefined, + }, + ref, + ); } /** Options for {@link runSubagentFollowUpTurn}. */ @@ -2240,7 +2249,7 @@ export async function runSubprocess(options: ExecutorOptions): Promise void) | null = null; - let reviveSession: (() => Promise) | null = null; + let reviveSession: AgentReviver | null = null; // Adopted (kept-alive) subagents flip registry status from session events on // later turns: revive/wake → running, turn drained → idle. The subscription // intentionally survives this run; a disposed session emits nothing, so it @@ -2248,9 +2257,9 @@ export async function runSubprocess(options: ExecutorOptions): Promise { target.subscribe(event => { if (event.type === "agent_start") { - AgentRegistry.global().setStatus(id, "running"); + AgentRegistry.global().setStatus(id, "running", target); } else if (event.type === "agent_end") { - AgentRegistry.global().setStatus(id, "idle"); + AgentRegistry.global().setStatus(id, "idle", target); } }); }; @@ -2426,7 +2435,10 @@ export async function runSubprocess(options: ExecutorOptions): Promise ({ + const buildSubagentSessionOptions = ( + sessionManagerForRun: SessionManager, + expectedAgentRef: CreateAgentSessionOptions["expectedAgentRef"], + ): CreateAgentSessionOptions => ({ cwd: worktree ?? cwd, authStorage, modelRegistry, @@ -2474,6 +2486,7 @@ export async function runSubprocess(options: ExecutorOptions): Promise { + reviveSession = async expectedAgentRef => { const reopened = await SessionManager.open(sessionFile, undefined, undefined, { suppressBreadcrumb: true, }); if (options.parentArtifactManager) { reopened.adoptArtifactManager(options.parentArtifactManager); } - const { session: revived } = await createAgentSession(buildSubagentSessionOptions(reopened)); + const { session: revived } = await createAgentSession( + buildSubagentSessionOptions(reopened, expectedAgentRef), + ); installRegistryStatusSync(revived); return revived; }; diff --git a/packages/coding-agent/src/task/output-manager.ts b/packages/coding-agent/src/task/output-manager.ts index a4954d91f..d675c3d82 100644 --- a/packages/coding-agent/src/task/output-manager.ts +++ b/packages/coding-agent/src/task/output-manager.ts @@ -18,10 +18,11 @@ import { ADVISOR_TRANSCRIPT_STEM } from "../advisor/transcript-recorder"; * * The first allocation of a given name keeps the name as-is; subsequent * allocations of the same name get a `-2`, `-3`, … suffix. On resume, scans - * existing output files so previously written outputs are never overwritten. + * existing output and child-session files so prior state is never overwritten. */ export class AgentOutputManager { #initialized = false; + #initializing: Promise | undefined; /** Final ids already handed out, relative to this manager's scope. */ readonly #taken = new Set(); readonly #getArtifactsDir: () => string | null; @@ -42,8 +43,12 @@ export class AgentOutputManager { */ async #ensureInitialized(): Promise { if (this.#initialized) return; + this.#initializing ??= this.#seedFromDisk(); + await this.#initializing; this.#initialized = true; + } + async #seedFromDisk(): Promise { const dir = this.#getArtifactsDir(); if (!dir) return; @@ -56,8 +61,9 @@ export class AgentOutputManager { const prefix = this.#parentPrefix ? `${this.#parentPrefix}.` : ""; for (const file of files) { - if (!file.endsWith(".md")) continue; - let rest = file.slice(0, -3); // drop ".md" + const extension = file.endsWith(".jsonl") ? ".jsonl" : file.endsWith(".md") ? ".md" : undefined; + if (!extension) continue; + let rest = file.slice(0, -extension.length); if (prefix) { if (!rest.startsWith(prefix)) continue; rest = rest.slice(prefix.length); @@ -80,6 +86,22 @@ export class AgentOutputManager { return this.#parentPrefix ? `${this.#parentPrefix}.${candidate}` : candidate; } + /** Reserve final IDs discovered outside the output directory scan. */ + async reserve(ids: Iterable): Promise { + await this.#ensureInitialized(); + const prefix = this.#parentPrefix ? `${this.#parentPrefix}.` : ""; + for (const id of ids) { + let rest = id; + if (prefix) { + if (!rest.startsWith(prefix)) continue; + rest = rest.slice(prefix.length); + } + const dot = rest.indexOf("."); + const segment = dot === -1 ? rest : rest.slice(0, dot); + if (segment) this.#taken.add(segment); + } + } + /** * Allocate a unique ID. * diff --git a/packages/coding-agent/src/task/persisted-revive.ts b/packages/coding-agent/src/task/persisted-revive.ts index 16f2edf23..a41c91991 100644 --- a/packages/coding-agent/src/task/persisted-revive.ts +++ b/packages/coding-agent/src/task/persisted-revive.ts @@ -71,7 +71,7 @@ export function createPersistedSubagentReviverFactory( taskDepth++; parentId = registry.get(parentId)?.parentId; } - return async () => { + return async expectedRef => { // Re-open fresh on every revive: park closes the writer, so this takes // the single-writer lock cleanly and restores the full message history. const reopened = await SessionManager.open(sessionFile, undefined, undefined, { @@ -96,6 +96,7 @@ export function createPersistedSubagentReviverFactory( agentDisplayName: ref.displayName, parentTaskPrefix: ref.id, parentAgentId: ref.parentId, + expectedAgentRef: expectedRef, taskDepth, toolNames: init.tools, outputSchema: init.outputSchema, @@ -119,8 +120,8 @@ export function createPersistedSubagentReviverFactory( // Without it the idle-TTL timer never clears on a turn and the lifecycle // could park the agent mid-run. session.subscribe(event => { - if (event.type === "agent_start") registry.setStatus(ref.id, "running"); - else if (event.type === "agent_end") registry.setStatus(ref.id, "idle"); + if (event.type === "agent_start") registry.setStatus(ref.id, "running", session); + else if (event.type === "agent_end") registry.setStatus(ref.id, "idle", session); }); return session; }; diff --git a/packages/coding-agent/src/tools/index.ts b/packages/coding-agent/src/tools/index.ts index 287c22ff6..b7d5202a9 100644 --- a/packages/coding-agent/src/tools/index.ts +++ b/packages/coding-agent/src/tools/index.ts @@ -25,6 +25,7 @@ import type { ArtifactManager } from "../session/artifacts"; import type { ClientBridge } from "../session/client-bridge"; import type { CustomMessage } from "../session/messages"; import type { UsageStatistics } from "../session/session-entries"; +import type { SessionManager } from "../session/session-manager"; import type { ToolChoiceQueue } from "../session/tool-choice-queue"; import { TaskTool } from "../task"; import type { AgentOutputManager } from "../task/output-manager"; @@ -209,6 +210,8 @@ export interface ToolSession { getEvalSessionId?: () => string | null; /** Get session file */ getSessionFile: () => string | null; + /** Parent session journal used by tools that persist runtime lifecycle state. */ + sessionManager?: Pick; /** Get eval kernel owner ID for session-scoped retained-kernel cleanup. */ getEvalKernelOwnerId?: () => string | null; /** Reject new eval work once session disposal has started. */ diff --git a/packages/coding-agent/src/tools/vibe.ts b/packages/coding-agent/src/tools/vibe.ts index 223ff965d..edc513724 100644 --- a/packages/coding-agent/src/tools/vibe.ts +++ b/packages/coding-agent/src/tools/vibe.ts @@ -23,7 +23,6 @@ import vibeListDescription from "../prompts/tools/vibe-list.md" with { type: "te import vibeSendDescription from "../prompts/tools/vibe-send.md" with { type: "text" }; import vibeSpawnDescription from "../prompts/tools/vibe-spawn.md" with { type: "text" }; import vibeWaitDescription from "../prompts/tools/vibe-wait.md" with { type: "text" }; -import { MAIN_AGENT_ID } from "../registry/agent-registry"; import { oneLineLabel } from "../task/types"; import { renderStatusLine } from "../tui"; import { @@ -93,7 +92,7 @@ export interface VibeToolDetails { } function screensOf(session: ToolSession, ids?: string[]): VibeScreenSnapshot[] { - return VibeSessionRegistry.global().screens(session.getAgentId?.() ?? MAIN_AGENT_ID, ids); + return VibeSessionRegistry.global().screens(session, ids); } function textResult(text: string, details: VibeToolDetails): AgentToolResult { diff --git a/packages/coding-agent/src/vibe/runtime.ts b/packages/coding-agent/src/vibe/runtime.ts index d7e1215aa..a040d5b95 100644 --- a/packages/coding-agent/src/vibe/runtime.ts +++ b/packages/coding-agent/src/vibe/runtime.ts @@ -24,7 +24,9 @@ import { registerArtifactsDir } from "../internal-urls/registry-helpers"; import { MCPManager } from "../mcp/manager"; import vibeTurnResultTemplate from "../prompts/tools/vibe-turn-result.md" with { type: "text" }; import { AgentLifecycleManager } from "../registry/agent-lifecycle"; -import { AgentRegistry, MAIN_AGENT_ID } from "../registry/agent-registry"; +import { type AgentRef, AgentRegistry, MAIN_AGENT_ID } from "../registry/agent-registry"; +import type { SessionEntry } from "../session/session-entries"; +import { SessionManager, SessionPersistenceIndeterminateError } from "../session/session-manager"; import { getBundledAgent } from "../task/agents"; import { type ExecutorOptions, runSubagentFollowUpTurn, runSubprocess } from "../task/executor"; import { generateTaskName } from "../task/name-generator"; @@ -68,6 +70,80 @@ const DEFAULT_WAIT_TIMEOUT_MS = 30_000; /** Response text cap inside a delivered turn result; full output stays at agent://. */ const RESPONSE_PREVIEW_MAX = 6000; +const VIBE_LIFECYCLE_CUSTOM_TYPE = "vibe-session-lifecycle"; +const VIBE_LIFECYCLE_VERSION = 1; + +export interface VibeOwnerScope { + ownerId: string; + parentSessionId: string; + parentSessionFile: string | null; +} + +export interface VibeParentSession { + getAgentId?: () => string | null; + getSessionId?: () => string | null; + getSessionFile: () => string | null; + sessionManager?: ToolSession["sessionManager"] & + Partial< + Pick + >; + asyncJobManager?: AsyncJobManager; + settings: ToolSession["settings"]; + getActiveModelString?: () => string | undefined; + getModelString?: () => string | undefined; +} + +type VibeTombstoneReason = "explicit-kill" | "mode-exit" | "spawn-failed" | "unrecoverable"; + +interface VibeLifecycleBase { + version: typeof VIBE_LIFECYCLE_VERSION; + id: string; + ownerId: string; + parentSessionId: string; +} + +interface VibeSpawnLifecycleEvent extends VibeLifecycleBase { + action: "spawn"; + cli: VibeCli; + agent: string; + childSessionFile: string; + createdAt: number; +} + +interface VibeTurnLifecycleEvent extends VibeLifecycleBase { + action: "turn-started" | "turn-settled"; + turn: number; +} + +interface VibeTombstoneLifecycleEvent extends VibeLifecycleBase { + action: "tombstone"; + reason: VibeTombstoneReason; +} + +interface VibeTombstoneRevocationEvent extends VibeLifecycleBase { + action: "tombstone-revoked"; + reason: "mode-exit"; +} + +type VibeLifecycleEvent = + | VibeSpawnLifecycleEvent + | VibeTurnLifecycleEvent + | VibeTombstoneLifecycleEvent + | VibeTombstoneRevocationEvent; + +interface VibeRestoreCandidate { + spawn: VibeSpawnLifecycleEvent; + turnCount: number; + lastActivityAt: number; + inFlight: boolean; + tombstoneReason?: VibeTombstoneReason; +} + +interface ResolvedVibeWorker { + agent: AgentDefinition; + modelOverride?: string | string[]; +} + interface VibeTurn { jobId: string; message: string; @@ -82,6 +158,9 @@ interface VibeRecord { id: string; cli: VibeCli; ownerId: string; + parentSessionId: string; + parentSessionFile: string | null; + childSessionFile?: string; agent: AgentDefinition; modelOverride?: string | string[]; state: VibeSessionState; @@ -106,6 +185,10 @@ interface VibeRecord { queue: string[]; turnCount: number; killed: boolean; + /** True while a parent switch is detaching this process-local record without terminating it. */ + suspended: boolean; + /** True only after a terminal lifecycle event has durably flushed. */ + terminalPersisted: boolean; } /** @@ -171,6 +254,87 @@ function firstLine(text: string, max = 100): string { return oneLineLabel(text, max); } +function scopeKey(scope: VibeOwnerScope, id: string): string { + return `${scope.parentSessionId}\0${scope.parentSessionFile ?? ""}\0${scope.ownerId}\0${id}`; +} + +function matchesScope(record: VibeRecord, scope: VibeOwnerScope): boolean { + return ( + record.ownerId === scope.ownerId && + record.parentSessionId === scope.parentSessionId && + record.parentSessionFile === scope.parentSessionFile + ); +} + +function objectRecord(value: unknown): Record | undefined { + if (typeof value !== "object" || value === null || Array.isArray(value)) return undefined; + return value as Record; +} + +function parseLifecycleEvent(value: unknown): VibeLifecycleEvent | undefined { + const data = objectRecord(value); + if (!data || data.version !== VIBE_LIFECYCLE_VERSION) return undefined; + if (typeof data.id !== "string" || !data.id) return undefined; + if (typeof data.ownerId !== "string" || !data.ownerId) return undefined; + if (typeof data.parentSessionId !== "string" || !data.parentSessionId) return undefined; + const base: VibeLifecycleBase = { + version: VIBE_LIFECYCLE_VERSION, + id: data.id, + ownerId: data.ownerId, + parentSessionId: data.parentSessionId, + }; + if (data.action === "spawn") { + const cli = data.cli === "fast" || data.cli === "good" ? data.cli : undefined; + if (!cli || typeof data.agent !== "string" || typeof data.childSessionFile !== "string") return undefined; + if (typeof data.createdAt !== "number" || !Number.isFinite(data.createdAt)) return undefined; + return { + ...base, + action: "spawn", + cli, + agent: data.agent, + childSessionFile: data.childSessionFile, + createdAt: data.createdAt, + }; + } + if (data.action === "turn-started" || data.action === "turn-settled") { + if (typeof data.turn !== "number" || !Number.isInteger(data.turn) || data.turn < 1) return undefined; + return { ...base, action: data.action, turn: data.turn }; + } + if (data.action === "tombstone") { + const reason = data.reason; + if ( + reason !== "explicit-kill" && + reason !== "mode-exit" && + reason !== "spawn-failed" && + reason !== "unrecoverable" + ) { + return undefined; + } + return { ...base, action: "tombstone", reason }; + } + if (data.action === "tombstone-revoked" && data.reason === "mode-exit") { + return { ...base, action: "tombstone-revoked", reason: "mode-exit" }; + } + return undefined; +} + +/** Child ids claimed by any valid Vibe spawn event, independent of current parent scope. */ +export function persistedVibeChildIds(entries: Iterable): Set { + const ids = new Set(); + for (const entry of entries) { + if (entry.type !== "custom" || entry.customType !== VIBE_LIFECYCLE_CUSTOM_TYPE) continue; + const event = parseLifecycleEvent(entry.data); + if ( + event?.action === "spawn" && + /^[A-Za-z0-9_-]+$/.test(event.id) && + event.childSessionFile === `${event.id}.jsonl` + ) { + ids.add(event.id); + } + } + return ids; +} + /** Merge the monitor's rolling `recentTools` window (newest first) into the per-turn trace (oldest first). */ function mergeTrace(turn: VibeTurn, progress: AgentProgress): void { turn.toolCount = progress.toolCount; @@ -188,9 +352,9 @@ function mergeTrace(turn: VibeTurn, progress: AgentProgress): void { export class VibeTurnError extends Error {} /** - * Process-global registry of vibe worker sessions, scoped per owner agent id - * (same convention as AsyncJobManager owner filters). The interactive mode - * kills an owner's sessions on vibe-mode exit via {@link killAll}. + * Process-global registry of vibe worker sessions, scoped by both owner agent + * id and stable parent session id. Persisted lifecycle events rebuild idle + * records after a process restart; live turn jobs remain process-local. */ export class VibeSessionRegistry { static #global: VibeSessionRegistry | undefined; @@ -208,6 +372,176 @@ export class VibeSessionRegistry { } readonly #records = new Map(); + readonly #terminationTails = new Map>(); + readonly #terminatedScopes = new Set(); + + ownerScope(session: VibeParentSession): VibeOwnerScope { + const parentSessionId = session.getSessionId?.(); + if (!parentSessionId) { + throw new ToolError("Vibe sessions require a stable parent session id."); + } + const parentSessionFile = session.getSessionFile(); + return { + ownerId: session.getAgentId?.() ?? MAIN_AGENT_ID, + parentSessionId, + parentSessionFile: parentSessionFile ? path.resolve(parentSessionFile) : null, + }; + } + + /** Re-open spawn admission after an explicit Vibe-mode entry. */ + activateScope(scope: VibeOwnerScope): void { + this.#terminatedScopes.delete(scopeKey(scope, "")); + } + + async #withTerminationLock(scope: VibeOwnerScope, operation: () => Promise): Promise { + const key = scopeKey(scope, ""); + const predecessor = this.#terminationTails.get(key) ?? Promise.resolve(); + const released = Promise.withResolvers(); + const tail = predecessor.then(() => released.promise); + this.#terminationTails.set(key, tail); + await predecessor; + try { + return await operation(); + } finally { + released.resolve(); + if (this.#terminationTails.get(key) === tail) this.#terminationTails.delete(key); + } + } + + #resolveWorker(session: VibeParentSession, cli: VibeCli): ResolvedVibeWorker { + const agentName = VIBE_CLI_AGENT[cli]; + const agent = getBundledAgent(agentName); + if (!agent) { + throw new ToolError(`Bundled agent "${agentName}" for vibe cli "${cli}" is unavailable.`); + } + const agentModelOverrides = session.settings.get("task.agentModelOverrides"); + return { + agent, + modelOverride: resolveAgentModelPatterns({ + settingsOverride: agentModelOverrides[agentName], + agentModel: agent.model, + settings: session.settings, + activeModelPattern: session.getActiveModelString?.(), + fallbackModelPattern: session.getModelString?.(), + }), + }; + } + + async #appendLifecycleEvent( + session: VibeParentSession, + event: VibeLifecycleEvent, + expectedParentSessionFile: string | null, + ): Promise { + if (!expectedParentSessionFile || !session.sessionManager) return false; + const matchesCurrentScope = (): boolean => { + const currentSessionFile = session.getSessionFile(); + return ( + session.getSessionId?.() === event.parentSessionId && + (session.getAgentId?.() ?? MAIN_AGENT_ID) === event.ownerId && + currentSessionFile !== null && + path.resolve(currentSessionFile) === expectedParentSessionFile + ); + }; + if (!matchesCurrentScope()) return false; + await session.sessionManager.ensureOnDisk(); + if (!matchesCurrentScope()) return false; + session.sessionManager.appendCustomEntry(VIBE_LIFECYCLE_CUSTOM_TYPE, event); + await session.sessionManager.flush(); + return true; + } + + #eventBase(record: VibeRecord): VibeLifecycleBase { + return { + version: VIBE_LIFECYCLE_VERSION, + id: record.id, + ownerId: record.ownerId, + parentSessionId: record.parentSessionId, + }; + } + + async #appendTombstone( + session: VibeParentSession, + record: VibeRecord, + reason: VibeTombstoneReason, + ): Promise { + return this.#appendLifecycleEvent( + session, + { + ...this.#eventBase(record), + action: "tombstone", + reason, + }, + record.parentSessionFile, + ); + } + + #hasInMemoryTombstone(session: VibeParentSession, record: VibeRecord): boolean { + let terminalReason: VibeTombstoneReason | undefined; + for (const entry of session.sessionManager?.getEntries() ?? []) { + if (entry.type !== "custom" || entry.customType !== VIBE_LIFECYCLE_CUSTOM_TYPE) continue; + const event = parseLifecycleEvent(entry.data); + if ( + !event || + event.id !== record.id || + event.ownerId !== record.ownerId || + event.parentSessionId !== record.parentSessionId + ) { + continue; + } + if (event.action === "tombstone") terminalReason = event.reason; + else if (event.action === "tombstone-revoked" && terminalReason === "mode-exit") terminalReason = undefined; + } + return terminalReason !== undefined; + } + + async #persistModeExit( + session: VibeParentSession, + scope: VibeOwnerScope, + records: readonly VibeRecord[], + ): Promise { + const pending = records.filter(record => !record.terminalPersisted); + const sessionManager = session.sessionManager; + if (!sessionManager) { + if (pending.some(record => record.childSessionFile)) { + throw new ToolError("Vibe mode exit cannot persist worker tombstones without the parent session manager."); + } + for (const record of pending) record.terminalPersisted = true; + return; + } + const currentScope = this.ownerScope(session); + if ( + currentScope.ownerId !== scope.ownerId || + currentScope.parentSessionId !== scope.parentSessionId || + currentScope.parentSessionFile !== scope.parentSessionFile + ) { + throw new ToolError("Vibe parent session changed before mode exit could be persisted."); + } + const parentSessionFile = currentScope.parentSessionFile; + const persistedPending = pending.filter(record => record.childSessionFile !== undefined); + for (const record of persistedPending) { + if ( + !parentSessionFile || + path.resolve(parentSessionFile.slice(0, -6), `${record.id}.jsonl`) !== record.childSessionFile + ) { + throw new ToolError(`Vibe session "${record.id}" changed parent scope before termination.`); + } + } + const appendEntriesAtomically = sessionManager.appendEntriesAtomically; + if (!appendEntriesAtomically) { + throw new ToolError("Vibe mode exit requires atomic parent-session persistence."); + } + await appendEntriesAtomically.call(sessionManager, () => { + for (const record of persistedPending) { + sessionManager.appendCustomEntry(VIBE_LIFECYCLE_CUSTOM_TYPE, { + ...this.#eventBase(record), + action: "tombstone", + reason: "mode-exit", + }); + } + sessionManager.appendModeChange?.("none"); + }); + for (const record of pending) record.terminalPersisted = true; + } #manager(session: ToolSession): AsyncJobManager { const manager = session.asyncJobManager; @@ -217,10 +551,10 @@ export class VibeSessionRegistry { return manager; } - #record(owner: string, id: string): VibeRecord { - const record = this.#records.get(id.trim()); - if (!record || record.ownerId !== owner) { - const roster = this.listIds(owner); + #record(scope: VibeOwnerScope, id: string): VibeRecord { + const record = this.#records.get(scopeKey(scope, id.trim())); + if (!record || !matchesScope(record, scope)) { + const roster = this.#listIds(scope); throw new ToolError( `Unknown vibe session "${id}".${roster.length > 0 ? ` Active sessions: ${roster.join(", ")}` : " No sessions — spawn one with vibe_spawn."}`, ); @@ -228,25 +562,37 @@ export class VibeSessionRegistry { return record; } - listIds(owner: string): string[] { + #registeredAgent(record: VibeRecord): AgentRef | undefined { + const ref = AgentRegistry.global().get(record.id); + if (ref?.kind !== "sub" || ref.parentId !== record.ownerId) return undefined; + if (record.childSessionFile && ref.sessionFile !== record.childSessionFile) return undefined; + return ref; + } + + #listIds(scope: VibeOwnerScope): string[] { const ids: string[] = []; for (const record of this.#records.values()) { - if (record.ownerId === owner && record.state !== "dead") ids.push(record.id); + if (matchesScope(record, scope) && record.state !== "dead") ids.push(record.id); } return ids; } + listIds(session: ToolSession): string[] { + return this.#listIds(this.ownerScope(session)); + } + /** * Live screen snapshots for rich rendering (the "TV wall"): one entry per * session in creation order, carrying the in-flight turn's trace, current * tool, and streamed text tail. All strings are one-line sanitized here so * renderers can print them verbatim. */ - screens(owner: string, ids?: string[]): VibeScreenSnapshot[] { + screens(session: ToolSession, ids?: string[]): VibeScreenSnapshot[] { + const scope = this.ownerScope(session); const wanted = ids?.length ? new Set(ids.map(id => id.trim())) : undefined; const records: VibeRecord[] = []; for (const record of this.#records.values()) { - if (record.ownerId !== owner) continue; + if (!matchesScope(record, scope)) continue; if (wanted && !wanted.has(record.id)) continue; records.push(record); } @@ -275,51 +621,272 @@ export class VibeSessionRegistry { })); } - /** Spawn a persistent worker session and start its first turn in the background. */ - async spawn(session: ToolSession, args: { cli: VibeCli; name?: string; prompt: string }): Promise { - const owner = session.getAgentId?.() ?? MAIN_AGENT_ID; - const manager = this.#manager(session); - const agentName = VIBE_CLI_AGENT[args.cli]; - const agent = getBundledAgent(agentName); - if (!agent) { - throw new ToolError(`Bundled agent "${agentName}" for vibe cli "${args.cli}" is unavailable.`); + #persistedIds(session: VibeParentSession, scope: VibeOwnerScope): Set { + const ids = new Set(); + for (const entry of session.sessionManager?.getEntries() ?? []) { + if (entry.type !== "custom" || entry.customType !== VIBE_LIFECYCLE_CUSTOM_TYPE) continue; + const event = parseLifecycleEvent(entry.data); + if (event?.ownerId === scope.ownerId && event.parentSessionId === scope.parentSessionId) ids.add(event.id); + } + for (const record of this.#records.values()) { + if (matchesScope(record, scope)) ids.add(record.id); + } + return ids; + } + + async #resolvePersistedChild( + parentSessionFile: string, + spawn: VibeSpawnLifecycleEvent, + options?: { requireAgentMatch?: boolean }, + ): Promise { + if (options?.requireAgentMatch !== false && spawn.agent !== VIBE_CLI_AGENT[spawn.cli]) return undefined; + if (!/^[A-Za-z0-9_-]+$/.test(spawn.id) || spawn.childSessionFile !== `${spawn.id}.jsonl`) return undefined; + const artifactsDir = path.resolve(parentSessionFile.slice(0, -6)); + const childSessionFile = path.resolve(artifactsDir, spawn.childSessionFile); + const relative = path.relative(artifactsDir, childSessionFile); + if (!relative || path.isAbsolute(relative) || relative.startsWith(`..${path.sep}`) || relative === "..") { + return undefined; + } + try { + const persisted = await SessionManager.peekSessionInit(childSessionFile); + return persisted?.init ? childSessionFile : undefined; + } catch { + return undefined; + } + } + + async #markTerminalRef( + id: string, + ownerId: string, + childSessionFile: string, + expected?: AgentRef | null, + ): Promise { + const registry = AgentRegistry.global(); + const existing = registry.get(id); + if (expected !== undefined && existing !== undefined && existing !== expected) return; + if ( + existing && + (existing.kind !== "sub" || existing.parentId !== ownerId || existing.sessionFile !== childSessionFile) + ) { + return; + } + if (existing?.status === "aborted" && !existing.session) return; + if (existing && !registry.setStatus(id, "aborted", existing)) return; + if (existing && AgentLifecycleManager.global().has(id, existing)) { + await AgentLifecycleManager.global().release(id, existing); + } else if (existing?.session) { + await existing.session.dispose(); + } + const current = registry.get(id); + if (current && current !== existing) return; + if (current) registry.unregister(id, current); + registry.register({ + id, + displayName: id, + kind: "sub", + parentId: ownerId, + session: null, + sessionFile: childSessionFile, + status: "aborted", + }); + } + + /** Reconcile resumable and terminal workers from the persisted parent journal. */ + async rehydrate(session: VibeParentSession): Promise { + const sessionFile = session.getSessionFile(); + const sessionManager = session.sessionManager; + if (!sessionFile || !sessionManager) return 0; + const scope = this.ownerScope(session); + const allSpawns = new Map(); + const terminalIntents = new Map(); + for (const entry of sessionManager.getEntries()) { + if (entry.type !== "custom" || entry.customType !== VIBE_LIFECYCLE_CUSTOM_TYPE) continue; + const event = parseLifecycleEvent(entry.data); + if (!event || event.ownerId !== scope.ownerId || event.parentSessionId !== scope.parentSessionId) continue; + if (event.action === "spawn") allSpawns.set(event.id, event); + else if (event.action === "tombstone") terminalIntents.set(event.id, event.reason); + else if (event.action === "tombstone-revoked" && terminalIntents.get(event.id) === "mode-exit") { + terminalIntents.delete(event.id); + } } - const agentModelOverrides = session.settings.get("task.agentModelOverrides"); - const modelOverride = resolveAgentModelPatterns({ - settingsOverride: agentModelOverrides[agentName], - agentModel: agent.model, - settings: session.settings, - activeModelPattern: session.getActiveModelString?.(), - fallbackModelPattern: session.getModelString?.(), - }); + const candidates = new Map(); + for (const entry of sessionManager.getBranch()) { + if (entry.type !== "custom" || entry.customType !== VIBE_LIFECYCLE_CUSTOM_TYPE) continue; + const event = parseLifecycleEvent(entry.data); + if (!event || event.ownerId !== scope.ownerId || event.parentSessionId !== scope.parentSessionId) continue; + const eventTime = Date.parse(entry.timestamp); + if (event.action === "spawn") { + candidates.set(event.id, { + spawn: event, + turnCount: 0, + lastActivityAt: Number.isFinite(eventTime) ? eventTime : event.createdAt, + inFlight: false, + }); + continue; + } + const candidate = candidates.get(event.id); + if (!candidate) continue; + candidate.lastActivityAt = Number.isFinite(eventTime) ? eventTime : candidate.lastActivityAt; + if (event.action === "turn-started" && event.turn >= candidate.turnCount) { + candidate.turnCount = event.turn; + candidate.inFlight = true; + } else if (event.action === "turn-settled" && event.turn >= candidate.turnCount) { + candidate.turnCount = event.turn; + candidate.inFlight = false; + } else if (event.action === "tombstone") { + candidate.tombstoneReason = event.reason; + } else if (event.action === "tombstone-revoked" && candidate.tombstoneReason === "mode-exit") { + candidate.tombstoneReason = undefined; + } + } + for (const id of terminalIntents.keys()) { + const spawn = allSpawns.get(id); + if (!spawn) continue; + const childSessionFile = await this.#resolvePersistedChild(sessionFile, spawn, { requireAgentMatch: false }); + if (!childSessionFile) continue; + await this.#markTerminalRef(id, scope.ownerId, childSessionFile); + this.#records.delete(scopeKey(scope, id)); + } + + let restored = 0; + for (const candidate of candidates.values()) { + const { spawn } = candidate; + if (candidate.tombstoneReason || terminalIntents.has(spawn.id) || candidate.turnCount < 1) continue; + const childSessionFile = await this.#resolvePersistedChild(sessionFile, spawn); + if (!childSessionFile) continue; + const key = scopeKey(scope, spawn.id); + if (this.#records.has(key)) continue; + const existing = AgentRegistry.global().get(spawn.id); + const existingIsResumable = + existing?.kind === "sub" && + existing.parentId === scope.ownerId && + existing.sessionFile === childSessionFile && + (existing.status === "idle" || existing.status === "parked"); + const blockedByCollision = Boolean(existing && !existingIsResumable); + const { agent, modelOverride } = this.#resolveWorker(session, spawn.cli); + if (!existing) { + AgentRegistry.global().register({ + id: spawn.id, + displayName: spawn.id, + kind: "sub", + parentId: scope.ownerId, + session: null, + sessionFile: childSessionFile, + status: "parked", + }); + } + this.#records.set(key, { + id: spawn.id, + cli: spawn.cli, + ownerId: scope.ownerId, + parentSessionId: scope.parentSessionId, + parentSessionFile: scope.parentSessionFile, + childSessionFile, + agent, + modelOverride, + state: "idle", + createdAt: spawn.createdAt, + lastActivityAt: candidate.lastActivityAt, + lastActivity: blockedByCollision + ? "blocked by an agent id collision" + : candidate.inFlight + ? `turn ${candidate.turnCount} interrupted by process restart` + : undefined, + queue: [], + turnCount: candidate.turnCount, + killed: false, + suspended: false, + terminalPersisted: false, + }); + restored++; + } + return restored; + } + + /** Spawn a persistent worker session and start its first turn in the background. */ + async spawn(session: ToolSession, args: { cli: VibeCli; name?: string; prompt: string }): Promise { + const scope = this.ownerScope(session); + return this.#withTerminationLock(scope, () => this.#spawnLocked(session, scope, args)); + } + + async #spawnLocked( + session: ToolSession, + scope: VibeOwnerScope, + args: { cli: VibeCli; name?: string; prompt: string }, + ): Promise { + if (this.#terminatedScopes.has(scopeKey(scope, ""))) { + throw new ToolError("Vibe mode has exited; enter Vibe mode again before spawning a worker."); + } + const manager = this.#manager(session); + const { agent, modelOverride } = this.#resolveWorker(session, args.cli); if (!session.agentOutputManager) { session.agentOutputManager = new AgentOutputManager(session.getArtifactsDir ?? (() => null)); } + const reservedIds = this.#persistedIds(session, scope); + for (const ref of AgentRegistry.global().list()) reservedIds.add(ref.id); + await session.agentOutputManager.reserve(reservedIds); const requestedName = args.name?.replace(/[^A-Za-z0-9_-]+/g, "").slice(0, 48); const id = await session.agentOutputManager.allocate(requestedName || generateTaskName()); - + const parentSessionFile = scope.parentSessionFile; + const childSessionName = `${id}.jsonl`; + const childSessionFile = parentSessionFile + ? path.resolve(parentSessionFile.slice(0, -6), childSessionName) + : undefined; + const createdAt = Date.now(); const record: VibeRecord = { id, cli: args.cli, - ownerId: owner, + ownerId: scope.ownerId, + parentSessionId: scope.parentSessionId, + parentSessionFile, + childSessionFile, agent, modelOverride, state: "starting", - createdAt: Date.now(), - lastActivityAt: Date.now(), + createdAt, + lastActivityAt: createdAt, queue: [], turnCount: 0, killed: false, + suspended: false, + terminalPersisted: false, }; - this.#records.set(id, record); - + const key = scopeKey(scope, id); + this.#records.set(key, record); + let spawnPersisted = false; try { + if (childSessionFile) { + spawnPersisted = await this.#appendLifecycleEvent( + session, + { + ...this.#eventBase(record), + action: "spawn", + cli: args.cli, + agent: agent.name, + childSessionFile: childSessionName, + createdAt, + }, + record.parentSessionFile, + ); + if (!spawnPersisted) throw new ToolError("Vibe parent session changed before the worker could start."); + } const jobId = this.#registerTurnJob(session, manager, record, args.prompt, { first: true }); return { id, jobId }; } catch (error) { - this.#records.delete(id); + record.killed = true; + record.state = "dead"; + record.lastActivityAt = Date.now(); + record.lastActivity = "spawn failed"; + if (childSessionFile) { + // A rejected terminal write leaves this dead record in the map so mode exit can retry it. + record.terminalPersisted = await this.#appendTombstone(session, record, "spawn-failed"); + if (!record.terminalPersisted) { + throw new ToolError("Vibe parent session changed before spawn failure could be persisted."); + } + } + this.#records.delete(key); throw error; } } @@ -330,16 +897,20 @@ export class VibeSessionRegistry { * background turn immediately. */ async send(session: ToolSession, args: { session: string; message: string }): Promise { - const owner = session.getAgentId?.() ?? MAIN_AGENT_ID; - const record = this.#record(owner, args.session); + const scope = this.ownerScope(session); + const record = this.#record(scope, args.session); if (record.state === "dead") { throw new ToolError(`Vibe session "${record.id}" is dead. Spawn a new one with vibe_spawn.`); } const message = args.message.trim(); if (!message) throw new ToolError("Message must not be empty."); + const registered = this.#registeredAgent(record); + if (AgentRegistry.global().get(record.id) && !registered) { + throw new ToolError(`Vibe session "${record.id}" no longer resolves to this parent session.`); + } if (record.turn) { - const live = AgentRegistry.global().get(record.id)?.session; + const live = registered?.session; if (live?.isStreaming) { await live.steer(message); record.lastActivityAt = Date.now(); @@ -350,6 +921,10 @@ export class VibeSessionRegistry { return { id: record.id, mode: "queued" }; } + if (!registered || (registered.status !== "idle" && registered.status !== "parked")) { + throw new ToolError(`Vibe session "${record.id}" no longer resolves to this parent session.`); + } + const manager = this.#manager(session); const jobId = this.#registerTurnJob(session, manager, record, message, { first: false }); return { id: record.id, mode: "turn", jobId }; @@ -365,14 +940,14 @@ export class VibeSessionRegistry { session: ToolSession, args: { sessions?: string[]; timeoutMs?: number; signal?: AbortSignal }, ): Promise { - const owner = session.getAgentId?.() ?? MAIN_AGENT_ID; + const scope = this.ownerScope(session); const manager = this.#manager(session); // Named sessions are watched regardless of state (a just-settled turn is // reported from its retained job); the no-args form watches every // session with a turn actually in flight. const watched = args.sessions?.length - ? args.sessions.map(id => this.#record(owner, id)) - : [...this.#records.values()].filter(record => record.ownerId === owner && record.turn !== undefined); + ? args.sessions.map(id => this.#record(scope, id)) + : [...this.#records.values()].filter(record => matchesScope(record, scope) && record.turn !== undefined); // Snapshot each watched turn's job at entry: #finishTurn installs a // queued follow-up turn inside the settling job's callback (before that @@ -440,41 +1015,179 @@ export class VibeSessionRegistry { return { settled, stillRunning, timedOut: waited && settled.length === 0 }; } - /** Terminate a worker: cancel its in-flight turn and dispose + unregister its session. */ - async kill(session: ToolSession, id: string): Promise { - const owner = session.getAgentId?.() ?? MAIN_AGENT_ID; - const record = this.#record(owner, id); - return this.#killRecord(record, session.asyncJobManager); - } - - /** Kill every session belonging to `owner` (vibe-mode exit / teardown). Returns the number killed. */ - async killAll(owner: string, manager?: AsyncJobManager): Promise { - let killed = 0; - for (const record of this.#records.values()) { - if (record.ownerId !== owner || record.state === "dead") continue; - await this.#killRecord(record, manager); - killed++; + /** Detach one parent's process-local workers without tombstoning their persisted conversations. */ + async suspendScope(scope: VibeOwnerScope, manager?: AsyncJobManager): Promise { + const records = [...this.#records.values()].filter(record => matchesScope(record, scope)); + const teardown = records.map(record => ({ + record, + ref: this.#registeredAgent(record), + job: record.turn && manager ? manager.getJob(record.turn.jobId) : undefined, + })); + for (const { record } of teardown) { + record.suspended = true; + record.queue.length = 0; + record.state = "dead"; + record.lastActivityAt = Date.now(); + record.lastActivity = "suspended for parent-session switch"; + this.#records.delete(scopeKey(scope, record.id)); + if (record.turn && manager) manager.cancel(record.turn.jobId, { ownerId: record.ownerId }); } - return killed; + for (const { record, ref } of teardown) { + if (!ref) continue; + try { + await AgentLifecycleManager.global().release(record.id, ref); + } catch (error) { + logger.warn("vibe: failed to detach worker session", { + id: record.id, + error: error instanceof Error ? error.message : String(error), + }); + } + } + await Promise.allSettled(teardown.flatMap(entry => (entry.job ? [entry.job.promise] : []))); + return records.length; } - async #killRecord(record: VibeRecord, manager: AsyncJobManager | undefined): Promise { + /** Terminate one worker; a tombstone failure still tears it down before reconciliation and error delivery. */ + async kill(session: ToolSession, id: string): Promise { + const scope = this.ownerScope(session); + return this.#withTerminationLock(scope, () => { + const record = this.#record(scope, id); + return this.#killRecord(record, session.asyncJobManager, session, "explicit-kill"); + }); + } + + /** Kill every live session in one parent scope after durably recording the complete mode-exit intent. */ + async killAll(session: VibeParentSession, ownerScope?: VibeOwnerScope): Promise { + const scope = ownerScope ?? this.ownerScope(session); + return this.#withTerminationLock(scope, async () => { + try { + const killed = await this.#killAllLocked(session, scope); + this.#terminatedScopes.add(scopeKey(scope, "")); + return killed; + } catch (error) { + if (error instanceof SessionPersistenceIndeterminateError) { + this.#terminatedScopes.add(scopeKey(scope, "")); + } + throw error; + } + }); + } + + async #killAllLocked(session: VibeParentSession, scope: VibeOwnerScope): Promise { + const records = [...this.#records.values()].filter( + record => matchesScope(record, scope) && !(record.state === "dead" && record.terminalPersisted), + ); + if (records.length === 0) { + const entries = session.sessionManager?.getBranch() ?? []; + for (let index = entries.length - 1; index >= 0; index--) { + const entry = entries[index]; + if (entry.type !== "mode_change") continue; + if (entry.mode === "none") return 0; + break; + } + } + try { + await this.#persistModeExit(session, scope, records); + } catch (error) { + if (error instanceof SessionPersistenceIndeterminateError) { + for (const record of records) { + await this.#killRecord(record, session.asyncJobManager, session, "mode-exit", false); + } + } + throw error; + } + for (const record of records) { + await this.#killRecord(record, session.asyncJobManager, session, "mode-exit"); + } + return records.length; + } + + async #killRecord( + record: VibeRecord, + manager: AsyncJobManager | undefined, + session: VibeParentSession, + reason: VibeTombstoneReason, + persistTerminal = true, + ): Promise { + const registered = this.#registeredAgent(record); + const settlingJobs = new Set(); + if (record.turn && manager) { + const job = manager.getJob(record.turn.jobId); + if (job) settlingJobs.add(job); + } + let persistenceError: unknown; + if (persistTerminal && !record.terminalPersisted) { + try { + if (record.killed) { + const recover = session.sessionManager?.recoverPersistenceFromCurrentState; + if (!recover) throw new ToolError("Vibe tombstone recovery requires parent-session persistence."); + await recover.call(session.sessionManager); + } + if (!this.#hasInMemoryTombstone(session, record) && record.childSessionFile) { + if (!(await this.#appendTombstone(session, record, reason))) { + throw new ToolError(`Vibe session "${record.id}" changed parent scope before termination.`); + } + } + record.terminalPersisted = true; + } catch (error) { + persistenceError = error; + } + } record.killed = true; record.queue.length = 0; let cancelledTurn = false; if (record.turn && manager) { + const job = manager.getJob(record.turn.jobId); + if (job) settlingJobs.add(job); cancelledTurn = manager.cancel(record.turn.jobId, { ownerId: record.ownerId }); } record.state = "dead"; record.lastActivityAt = Date.now(); record.lastActivity = "killed"; - try { - await AgentLifecycleManager.global().release(record.id); - } catch (error) { - logger.warn("vibe: failed to release worker session", { - id: record.id, - error: error instanceof Error ? error.message : String(error), - }); + if (registered) { + try { + await AgentLifecycleManager.global().release(record.id, registered); + } catch (error) { + logger.warn("vibe: failed to release worker session", { + id: record.id, + error: error instanceof Error ? error.message : String(error), + }); + } + } + await Promise.allSettled([...settlingJobs].map(job => job.promise)); + const terminalRef = registered ?? this.#registeredAgent(record) ?? null; + if (record.childSessionFile) { + try { + const persisted = await SessionManager.peekSessionInit(record.childSessionFile); + if (persisted?.init) + await this.#markTerminalRef(record.id, record.ownerId, record.childSessionFile, terminalRef); + } catch { + // A turn killed before session initialization has no readable history to retain. + } + } + if (persistenceError) { + let finalPersistenceError = persistenceError; + const recover = session.sessionManager?.recoverPersistenceFromCurrentState; + if (recover) { + try { + await recover.call(session.sessionManager); + if (!this.#hasInMemoryTombstone(session, record) && record.childSessionFile) { + if (!(await this.#appendTombstone(session, record, reason))) { + throw new ToolError(`Vibe session "${record.id}" changed parent scope before termination.`); + } + } + record.terminalPersisted = true; + } catch (recoveryError) { + if (recoveryError instanceof SessionPersistenceIndeterminateError) { + finalPersistenceError = recoveryError; + } + logger.warn("vibe: failed to reconcile explicit tombstone persistence", { + id: record.id, + error: recoveryError instanceof Error ? recoveryError.message : String(recoveryError), + }); + } + } + throw finalPersistenceError; } return { id: record.id, cancelledTurn }; } @@ -580,6 +1293,18 @@ export class VibeSessionRegistry { record.turnCount = turnIndex; record.lastActivityAt = Date.now(); try { + const turnStartedPersisted = await this.#appendLifecycleEvent( + session, + { + ...this.#eventBase(record), + action: "turn-started", + turn: turnIndex, + }, + record.parentSessionFile, + ); + if (record.childSessionFile && !turnStartedPersisted) { + throw new ToolError(`Vibe session "${record.id}" changed parent scope before its turn started.`); + } const result = options.first ? await runSubprocess(await this.#buildSpawnOptions(session, record, message, signal, onProgress)) : await runSubagentFollowUpTurn({ @@ -592,10 +1317,10 @@ export class VibeSessionRegistry { eventBus: session.eventBus, artifactsDir: session.getSessionFile()?.slice(0, -6), }); - return this.#settleTurn(session, manager, record, turn, ownJobId, turnIndex, result); + return await this.#settleTurn(session, manager, record, turn, ownJobId, turnIndex, result); } catch (error) { if (error instanceof VibeTurnError) throw error; - this.#finishTurn(session, manager, record, ownJobId); + await this.#finishTurn(session, manager, record, ownJobId); const reason = error instanceof Error ? error.message : String(error); record.lastActivity = firstLine(`turn failed: ${reason}`); throw new VibeTurnError( @@ -611,19 +1336,41 @@ export class VibeSessionRegistry { } /** Post-turn bookkeeping shared by success and failure paths: clear the in-flight turn, flush the queue. */ - #finishTurn(session: ToolSession, manager: AsyncJobManager, record: VibeRecord, settledJobId: string): void { + async #finishTurn( + session: ToolSession, + manager: AsyncJobManager, + record: VibeRecord, + settledJobId: string, + ): Promise { record.lastJobId = settledJobId; record.turn = undefined; record.live = undefined; record.lastActivityAt = Date.now(); - if (record.killed) { + if (record.killed || record.suspended) { record.state = "dead"; return; } - // A spawn that failed before its session ever registered leaves nothing - // to continue — mark the record dead so sends fail with clear guidance. - record.state = AgentRegistry.global().get(record.id) ? "idle" : "dead"; - if (record.state === "dead" || record.queue.length === 0) return; + // Only an idle/parked ref with this parent's exact child file is resumable. + const registered = this.#registeredAgent(record); + record.state = registered && (registered.status === "idle" || registered.status === "parked") ? "idle" : "dead"; + if (record.state === "dead") { + record.terminalPersisted = await this.#appendTombstone(session, record, "unrecoverable"); + return; + } + const settledPersisted = await this.#appendLifecycleEvent( + session, + { + ...this.#eventBase(record), + action: "turn-settled", + turn: record.turnCount, + }, + record.parentSessionFile, + ); + if (record.childSessionFile && !settledPersisted) { + record.state = "dead"; + return; + } + if (record.queue.length === 0) return; const nextMessage = record.queue.splice(0, record.queue.length).join("\n\n"); try { this.#registerTurnJob(session, manager, record, nextMessage, { first: false }); @@ -638,7 +1385,7 @@ export class VibeSessionRegistry { } /** Format a settled turn into the self-delivering result text (activity trace + response). */ - #settleTurn( + async #settleTurn( session: ToolSession, manager: AsyncJobManager, record: VibeRecord, @@ -646,8 +1393,8 @@ export class VibeSessionRegistry { settledJobId: string, turnIndex: number, result: SingleResult, - ): string { - this.#finishTurn(session, manager, record, settledJobId); + ): Promise { + await this.#finishTurn(session, manager, record, settledJobId); const failed = result.exitCode !== 0 || result.aborted === true; const status = result.aborted ? "aborted" : failed ? "failed" : "completed"; record.lastActivity = firstLine( diff --git a/packages/coding-agent/test/agent-hub-activate.test.ts b/packages/coding-agent/test/agent-hub-activate.test.ts index 4cc4033c7..50ecba602 100644 --- a/packages/coding-agent/test/agent-hub-activate.test.ts +++ b/packages/coding-agent/test/agent-hub-activate.test.ts @@ -4,6 +4,7 @@ * focus failure keeps the hub open and surfaces the error as a notice. */ import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "bun:test"; +import * as fs from "node:fs/promises"; import * as path from "node:path"; import { resetSettingsForTest, Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { IrcBus } from "@oh-my-pi/pi-coding-agent/irc/bus"; @@ -14,6 +15,7 @@ import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; import type { InteractiveModeContext } from "@oh-my-pi/pi-coding-agent/modes/types"; import { AgentRegistry } from "@oh-my-pi/pi-coding-agent/registry/agent-registry"; import type { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; +import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; import { TempDir } from "@oh-my-pi/pi-utils"; const AGENT_ID = "Worker"; @@ -115,6 +117,66 @@ describe("Agent hub Enter activation", () => { hub.dispose(); }); + it("does not generically revive active or tombstoned Vibe children copied by a post-exit fork", async () => { + using tempDir = TempDir.createSync("@omp-agent-hub-vibe-fork-"); + const manager = SessionManager.create(tempDir.path(), tempDir.path()); + manager.appendModeChange("vibe"); + const parentSessionId = manager.getSessionId(); + for (const id of ["ActiveVibe", "KilledVibe"]) { + manager.appendCustomEntry("vibe-session-lifecycle", { + version: 1, + action: "spawn", + id, + ownerId: "Main", + parentSessionId, + cli: "fast", + agent: "sonic", + childSessionFile: `${id}.jsonl`, + createdAt: Date.now(), + }); + } + manager.appendCustomEntry("vibe-session-lifecycle", { + version: 1, + action: "tombstone", + id: "KilledVibe", + ownerId: "Main", + parentSessionId, + reason: "mode-exit", + }); + manager.appendModeChange("none"); + await manager.ensureOnDisk(); + await manager.flush(); + const sourceSessionFile = manager.getSessionFile(); + if (!sourceSessionFile) throw new Error("Expected source session file"); + const sourceArtifacts = sourceSessionFile.slice(0, -6); + await fs.mkdir(sourceArtifacts, { recursive: true }); + for (const id of ["ActiveVibe", "KilledVibe", "OrdinaryTask"]) { + await fs.writeFile(path.join(sourceArtifacts, `${id}.jsonl`), "persisted child"); + } + const fork = await manager.fork(); + if (!fork) throw new Error("Expected persisted fork"); + await fs.cp(sourceArtifacts, fork.newSessionFile.slice(0, -6), { recursive: true }); + await manager.close(); + + const agents = new AgentRegistry(); + const hub = new AgentHubOverlayComponent({ + observers: new SessionObserverRegistry(), + hubKeys: [], + onDone: () => {}, + requestRender: () => {}, + registry: agents, + irc: new IrcBus(agents), + focusAgent: async () => {}, + sessionFile: fork.newSessionFile, + }); + await hub.persistedSubagentsReady; + + expect(agents.get("ActiveVibe")).toBeUndefined(); + expect(agents.get("KilledVibe")).toBeUndefined(); + expect(agents.get("OrdinaryTask")?.status).toBe("parked"); + hub.dispose(); + }); + it("selector controller restores focus to the editor after Enter focuses an agent", async () => { const agents = new AgentRegistry(); agents.register({ diff --git a/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts b/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts index 84cb48cd6..1a1b20ac9 100644 --- a/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts +++ b/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts @@ -16,8 +16,10 @@ import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage"; import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; +import { FileSessionStorage, type WriteTextAtomicOptions } from "@oh-my-pi/pi-coding-agent/session/session-storage"; import { VIBE_TOOL_NAMES } from "@oh-my-pi/pi-coding-agent/tools/vibe"; import { EventBus } from "@oh-my-pi/pi-coding-agent/utils/event-bus"; +import { VibeSessionRegistry } from "@oh-my-pi/pi-coding-agent/vibe/runtime"; import { TempDir } from "@oh-my-pi/pi-utils"; import { type } from "arktype"; @@ -34,11 +36,56 @@ function stubTool(name: string): AgentTool { }; } +function vibeModeEntryCount(manager: SessionManager): number { + return manager.getEntries().filter(entry => entry.type === "mode_change" && entry.mode === "vibe").length; +} + +class ExitFaultStorage extends FileSessionStorage { + failNextAtomicWrite = false; + #readGate: + | { + filePath: string; + started: ReturnType>; + release: ReturnType>; + } + | undefined; + + gateNextRead(filePath: string): { started: Promise; release: () => void } { + const started = Promise.withResolvers(); + const release = Promise.withResolvers(); + this.#readGate = { filePath, started, release }; + return { started: started.promise, release: release.resolve }; + } + + override async readTextSlices( + filePath: string, + prefixBytes: number, + suffixBytes: number, + ): Promise<[string, string]> { + const gate = this.#readGate; + if (gate?.filePath === filePath) { + this.#readGate = undefined; + gate.started.resolve(); + await gate.release.promise; + } + return super.readTextSlices(filePath, prefixBytes, suffixBytes); + } + + override async writeTextAtomic(filePath: string, content: string, options?: WriteTextAtomicOptions): Promise { + if (this.failNextAtomicWrite) { + this.failNextAtomicWrite = false; + throw Object.assign(new Error("journal atomic publish failed"), { code: "ENOSPC" }); + } + await super.writeTextAtomic(filePath, content, options); + } +} + describe("InteractiveMode vibe mode toggle", () => { let tempDir: TempDir; let authStorage: AuthStorage; let session: AgentSession; let mode: InteractiveMode; + let storage: ExitFaultStorage; beforeAll(async () => { await initTheme(); @@ -46,6 +93,7 @@ describe("InteractiveMode vibe mode toggle", () => { beforeEach(async () => { resetSettingsForTest(); + VibeSessionRegistry.resetGlobalForTests(); tempDir = TempDir.createSync("@pi-vibe-toggle-"); await Settings.init({ inMemory: true, cwd: tempDir.path() }); authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth.db")); @@ -55,6 +103,7 @@ describe("InteractiveMode vibe mode toggle", () => { const registryTools = [stubTool("read")]; + storage = new ExitFaultStorage(); session = new AgentSession({ agent: new Agent({ initialState: { @@ -64,7 +113,7 @@ describe("InteractiveMode vibe mode toggle", () => { messages: [], }, }), - sessionManager: SessionManager.create(tempDir.path(), tempDir.path()), + sessionManager: SessionManager.create(tempDir.path(), tempDir.path(), storage), settings: Settings.isolated({}), modelRegistry, toolRegistry: new Map(registryTools.map(tool => [tool.name, tool])), @@ -76,6 +125,7 @@ describe("InteractiveMode vibe mode toggle", () => { afterEach(async () => { mode?.stop(); await session?.dispose(); + VibeSessionRegistry.resetGlobalForTests(); authStorage?.close(); tempDir?.removeSync(); vi.restoreAllMocks(); @@ -103,4 +153,89 @@ describe("InteractiveMode vibe mode toggle", () => { expect(session.getActiveToolNames()).toEqual([]); expect(session.getAllToolNames()).toEqual(["read"]); }); + + it("preserves workers and mode metadata on a same-session reload", async () => { + await mode.init({ suppressWelcomeIntro: true }); + await mode.handleVibeModeCommand(); + await session.sessionManager.ensureOnDisk(); + const sessionFile = session.sessionFile; + if (!sessionFile) throw new Error("Expected persisted session file"); + const registry = VibeSessionRegistry.global(); + const suspend = vi.spyOn(registry, "suspendScope"); + const terminate = vi.spyOn(registry, "killAll"); + + const readGate = storage.gateNextRead(sessionFile); + const switching = session.switchSession(sessionFile); + await readGate.started; + const suspendCallsBeforeRead = suspend.mock.calls.length; + readGate.release(); + expect(suspendCallsBeforeRead).toBe(1); + expect(await switching).toBe(true); + + expect(mode.vibeModeEnabled).toBe(true); + expect(suspend).toHaveBeenCalledTimes(1); + expect(terminate).not.toHaveBeenCalled(); + expect(vibeModeEntryCount(session.sessionManager)).toBe(1); + }); + + it("suspends the old scope without tombstones when switching to another vibe parent", async () => { + await mode.init({ suppressWelcomeIntro: true }); + await mode.handleVibeModeCommand(); + await session.sessionManager.ensureOnDisk(); + const originalSessionId = session.sessionManager.getSessionId(); + const targetManager = SessionManager.create(tempDir.path(), tempDir.path()); + targetManager.appendModeChange("vibe"); + await targetManager.ensureOnDisk(); + const targetFile = targetManager.getSessionFile(); + if (!targetFile) throw new Error("Expected target session file"); + await targetManager.close(); + const registry = VibeSessionRegistry.global(); + const suspend = vi.spyOn(registry, "suspendScope"); + const terminate = vi.spyOn(registry, "killAll"); + + expect(await session.switchSession(targetFile)).toBe(true); + + expect(mode.vibeModeEnabled).toBe(true); + expect(suspend).toHaveBeenCalledTimes(1); + expect(suspend.mock.calls[0]?.[0]).toMatchObject({ parentSessionId: originalSessionId }); + expect(terminate).not.toHaveBeenCalled(); + expect(vibeModeEntryCount(session.sessionManager)).toBe(1); + }); + + it("rejects new, drop, fork, and move transitions at the AgentSession boundary while vibe is active", async () => { + await mode.init({ suppressWelcomeIntro: true }); + await mode.handleVibeModeCommand(); + await session.sessionManager.ensureOnDisk(); + const sessionFile = session.sessionFile; + if (!sessionFile) throw new Error("Expected persisted session file"); + + await expect(session.newSession()).rejects.toThrow("Exit vibe mode first"); + await expect(session.newSession({ drop: true })).rejects.toThrow("Exit vibe mode first"); + await expect(session.fork()).rejects.toThrow("Exit vibe mode first"); + await expect(session.moveSession(path.join(tempDir.path(), "other-project"))).rejects.toThrow( + "Exit vibe mode first", + ); + expect(session.sessionFile).toBe(sessionFile); + expect(session.sessionManager.getCwd()).toBe(tempDir.path()); + expect(mode.vibeModeEnabled).toBe(true); + }); + + it("keeps vibe mode and tools active after a real storage failure, then allows a retry", async () => { + await mode.init({ suppressWelcomeIntro: true }); + await mode.handleVibeModeCommand(); + const activeTools = session.getActiveToolNames(); + storage.failNextAtomicWrite = true; + + const exitError = await mode.handleVibeModeCommand().catch(error => error); + expect(exitError).toBeInstanceOf(Error); + expect(String(exitError)).toContain("journal atomic publish failed"); + + expect(mode.vibeModeEnabled).toBe(true); + expect(session.getVibeModeState()).toEqual({ enabled: true }); + expect(session.getActiveToolNames()).toEqual(activeTools); + expect(vibeModeEntryCount(session.sessionManager)).toBe(1); + + await mode.handleVibeModeCommand(); + expect(mode.vibeModeEnabled).toBe(false); + }); }); diff --git a/packages/coding-agent/test/registry/agent-lifecycle.test.ts b/packages/coding-agent/test/registry/agent-lifecycle.test.ts index 577ab4467..3284629c4 100644 --- a/packages/coding-agent/test/registry/agent-lifecycle.test.ts +++ b/packages/coding-agent/test/registry/agent-lifecycle.test.ts @@ -56,6 +56,27 @@ describe("AgentLifecycleManager", () => { return registry.register({ id, displayName: "task", kind: "sub", session, sessionFile, status: "idle" }); } + it("registerIfAvailable never replaces a collision and reuses only the exact expected ref", () => { + const parked = registerIdleSub("generation-Sub", null); + registry.setStatus("generation-Sub", "parked", parked); + const next = { + id: "generation-Sub", + displayName: "replacement", + kind: "sub" as const, + session: null, + status: "running" as const, + }; + + expect(registry.registerIfAvailable(next, null)).toBeUndefined(); + expect(registry.get("generation-Sub")).toBe(parked); + expect(registry.registerIfAvailable(next, parked)).toBe(parked); + expect(registry.get("generation-Sub")).toBe(parked); + + registry.unregister("generation-Sub", parked); + expect(registry.registerIfAvailable(next, parked)).toBeUndefined(); + expect(registry.get("generation-Sub")).toBeUndefined(); + }); + it("adopt arms the TTL: an idle agent is parked — session disposed, ref + sessionFile retained", async () => { vi.useFakeTimers(); const stub = makeSessionStub(); @@ -250,6 +271,32 @@ describe("AgentLifecycleManager", () => { expect(registry.get("6-Sub")).toBeUndefined(); }); + it("a delayed release cannot remove or mutate a replacement ref with the same id", async () => { + const gate = deferred(); + const oldSession = makeSessionStub(() => gate.promise); + const oldRef = registerIdleSub("cas-Sub", oldSession.session); + lifecycle.adopt("cas-Sub", { idleTtlMs: 0 }, oldRef); + const releasing = lifecycle.release("cas-Sub", oldRef); + await flushAsync(); + expect(oldSession.disposeCalls()).toBe(1); + + const replacementSession = makeSessionStub(); + const replacement = registerIdleSub("cas-Sub", replacementSession.session, "/tmp/replacement.jsonl"); + lifecycle.adopt("cas-Sub", { idleTtlMs: 0 }, replacement); + expect(registry.setStatus("cas-Sub", "aborted", oldRef)).toBe(false); + expect(registry.detachSession("cas-Sub", oldRef)).toBe(false); + expect(registry.unregister("cas-Sub", oldRef)).toBe(false); + + gate.resolve(); + await releasing; + + expect(registry.get("cas-Sub")).toBe(replacement); + expect(replacement.status).toBe("idle"); + expect(replacement.session).toBe(replacementSession.session); + expect(replacementSession.disposeCalls()).toBe(0); + expect(lifecycle.has("cas-Sub", replacement)).toBe(true); + }); + it("adopt(Main) is a no-op: Main is never adopted or parked", async () => { vi.useFakeTimers(); const stub = makeSessionStub(); diff --git a/packages/coding-agent/test/sdk-session-isolation.test.ts b/packages/coding-agent/test/sdk-session-isolation.test.ts index dba4d4b65..60a834f2d 100644 --- a/packages/coding-agent/test/sdk-session-isolation.test.ts +++ b/packages/coding-agent/test/sdk-session-isolation.test.ts @@ -1,4 +1,4 @@ -import { afterAll, afterEach, beforeAll, describe, expect, it } from "bun:test"; +import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "bun:test"; import * as fs from "node:fs"; import * as os from "node:os"; import * as path from "node:path"; @@ -7,10 +7,13 @@ import { getBundledModel } from "@oh-my-pi/pi-catalog/models"; import type { Rule } from "@oh-my-pi/pi-coding-agent/capability/rule"; import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import { AgentLifecycleManager } from "@oh-my-pi/pi-coding-agent/registry/agent-lifecycle"; +import { AgentRegistry } from "@oh-my-pi/pi-coding-agent/registry/agent-registry"; import { createAgentSession } from "@oh-my-pi/pi-coding-agent/sdk"; import { SecretObfuscator } from "@oh-my-pi/pi-coding-agent/secrets"; import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage"; import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; +import { VibeSessionRegistry } from "@oh-my-pi/pi-coding-agent/vibe/runtime"; import { getSessionsDir, removeSyncWithRetries, Snowflake } from "@oh-my-pi/pi-utils"; function createTtsrRule(name: string): Rule { @@ -77,6 +80,7 @@ describe("createAgentSession session storage isolation", () => { }); afterEach(async () => { + vi.restoreAllMocks(); for (const tempDir of tempDirs.splice(0)) { removeSyncWithRetries(tempDir); } @@ -115,6 +119,131 @@ describe("createAgentSession session storage isolation", () => { await session.dispose(); } }); + it("does not replace a newer registry generation when creation expected the id to be absent", async () => { + const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), `pi-sdk-generation-cas-${Snowflake.next()}-`)); + tempDirs.push(tempDir); + const cwd = path.join(tempDir, "project"); + fs.mkdirSync(cwd, { recursive: true }); + const registry = new AgentRegistry(); + const replacement = registry.register({ + id: "shared-worker", + displayName: "replacement B", + kind: "sub", + parentId: "Main", + session: null, + status: "idle", + }); + + await expect( + createAgentSession({ + cwd, + agentDir: path.join(tempDir, "agent"), + modelRegistry: sharedModelRegistry, + settings: Settings.isolated(), + disableExtensionDiscovery: true, + skills: [], + contextFiles: [], + promptTemplates: [], + slashCommands: [], + enableMCP: false, + enableLsp: false, + agentRegistry: registry, + agentId: "shared-worker", + agentDisplayName: "late A", + parentTaskPrefix: "shared-worker", + parentAgentId: "Main", + taskDepth: 1, + expectedAgentRef: null, + }), + ).rejects.toThrow("already owned by another session generation"); + expect(registry.get("shared-worker")).toBe(replacement); + expect(replacement).toMatchObject({ status: "idle", session: null }); + }); + + it("reuses the exact parked ref authorized for revival", async () => { + const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), `pi-sdk-generation-revive-${Snowflake.next()}-`)); + tempDirs.push(tempDir); + const cwd = path.join(tempDir, "project"); + fs.mkdirSync(cwd, { recursive: true }); + const sessionManager = SessionManager.create(cwd, tempDir); + await sessionManager.ensureOnDisk(); + const sessionFile = sessionManager.getSessionFile(); + if (!sessionFile) throw new Error("Expected persisted worker session file"); + const registry = new AgentRegistry(); + const parked = registry.register({ + id: "revived-worker", + displayName: "revived worker", + kind: "sub", + parentId: "Main", + session: null, + sessionFile, + status: "parked", + }); + + const { session } = await createAgentSession({ + cwd, + modelRegistry: sharedModelRegistry, + settings: Settings.isolated(), + disableExtensionDiscovery: true, + skills: [], + contextFiles: [], + promptTemplates: [], + slashCommands: [], + enableMCP: false, + enableLsp: false, + sessionManager, + agentRegistry: registry, + agentId: "revived-worker", + agentDisplayName: "revived worker", + parentTaskPrefix: "revived-worker", + parentAgentId: "Main", + taskDepth: 1, + expectedAgentRef: parked, + }); + try { + expect(registry.get("revived-worker")).toBe(parked); + expect(parked).toMatchObject({ status: "running", session, sessionFile }); + } finally { + await session.dispose(); + } + expect(registry.get("revived-worker")).toBeUndefined(); + }); + + it("suspends the exact Vibe owner scope before global lifecycle teardown", async () => { + VibeSessionRegistry.resetGlobalForTests(); + const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), `pi-sdk-vibe-dispose-${Snowflake.next()}-`)); + tempDirs.push(tempDir); + const cwd = path.join(tempDir, "project"); + fs.mkdirSync(cwd, { recursive: true }); + const { session } = await createAgentSession({ + cwd, + agentDir: path.join(tempDir, "agent"), + modelRegistry: sharedModelRegistry, + settings: Settings.isolated(), + disableExtensionDiscovery: true, + skills: [], + contextFiles: [], + promptTemplates: [], + slashCommands: [], + enableMCP: false, + enableLsp: false, + }); + const vibeRegistry = VibeSessionRegistry.global(); + const suspend = vi.spyOn(vibeRegistry, "suspendScope"); + const lifecycleDispose = vi.spyOn(AgentLifecycleManager.global(), "dispose"); + const parentSessionId = session.sessionManager.getSessionId(); + const parentSessionFile = session.sessionManager.getSessionFile(); + if (!parentSessionFile) throw new Error("Expected persisted parent session file"); + + await session.dispose(); + + expect(suspend).toHaveBeenCalledWith( + { ownerId: "Main", parentSessionId, parentSessionFile }, + session.asyncJobManager, + ); + expect(suspend.mock.invocationCallOrder[0]).toBeLessThan(lifecycleDispose.mock.invocationCallOrder[0]); + }); + it("wires the discovered TTSR manager into the created session", async () => { const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), `pi-sdk-ttsr-${Snowflake.next()}-`)); tempDirs.push(tempDir); diff --git a/packages/coding-agent/test/session-manager-atomic-rewrite-race.test.ts b/packages/coding-agent/test/session-manager-atomic-rewrite-race.test.ts index 41b8dc695..8048e2a73 100644 --- a/packages/coding-agent/test/session-manager-atomic-rewrite-race.test.ts +++ b/packages/coding-agent/test/session-manager-atomic-rewrite-race.test.ts @@ -1,6 +1,13 @@ import { describe, expect, it } from "bun:test"; import { getBundledModel } from "@oh-my-pi/pi-catalog/models"; -import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; +import { + IndexedSessionStorage, + type SessionStorageBackend, +} from "@oh-my-pi/pi-coding-agent/session/indexed-session-storage"; +import { + SessionManager, + SessionPersistenceIndeterminateError, +} from "@oh-my-pi/pi-coding-agent/session/session-manager"; import { MemorySessionStorage, type SessionStorageWriter, @@ -679,3 +686,270 @@ describe("SessionManager fence handoff across superseded rewrites", () => { expect(storage.detachedLines).toEqual([]); }); }); + +interface AtomicFailureHandle { + started: Promise; + release: () => void; +} + +class GatedAtomicFailureStorage extends MemorySessionStorage { + #nextFailure: + | { + error: Error; + started: ReturnType>; + release: ReturnType>; + } + | undefined; + + failNextAtomicWrite(error: Error): AtomicFailureHandle { + if (this.#nextFailure) throw new Error("Atomic failure already armed"); + const started = Promise.withResolvers(); + const release = Promise.withResolvers(); + this.#nextFailure = { error, started, release }; + return { started: started.promise, release: release.resolve }; + } + + override async writeTextAtomic(path: string, content: string, options?: WriteTextAtomicOptions): Promise { + const failure = this.#nextFailure; + if (!failure) { + await super.writeTextAtomic(path, content, options); + return; + } + this.#nextFailure = undefined; + failure.started.resolve(); + await failure.release.promise; + throw failure.error; + } +} + +class ScriptedAtomicFailureStorage extends MemorySessionStorage { + readonly behaviors: Array<{ commit: boolean; error: Error }> = []; + + override async writeTextAtomic(path: string, content: string, options?: WriteTextAtomicOptions): Promise { + const behavior = this.behaviors.shift(); + if (!behavior) { + await super.writeTextAtomic(path, content, options); + return; + } + if (behavior.commit) await super.writeTextAtomic(path, content, options); + throw behavior.error; + } +} + +describe("SessionManager atomic entry batches", () => { + it("restores the exact active branch when an atomic batch publish fails", async () => { + const storage = new GatedAtomicFailureStorage(); + const manager = SessionManager.create("/cwd", "/sessions", storage); + const rootId = manager.appendCustomEntry("root"); + manager.appendCustomEntry("abandoned-tail"); + await manager.ensureOnDisk(); + const sessionFile = manager.getSessionFile(); + if (!sessionFile) throw new Error("Expected session file"); + manager.branch(rootId); + const before = await storage.readText(sessionFile); + const failure = storage.failNextAtomicWrite(new Error("batch publish failed")); + const commit = manager.appendEntriesAtomically(() => manager.appendCustomEntry("staged-terminal")); + await failure.started; + failure.release(); + + await expect(commit).rejects.toThrow("batch publish failed"); + expect(manager.getBranch().map(entry => entry.id)).toEqual([rootId]); + expect( + manager.getEntries().some(entry => entry.type === "custom" && entry.customType === "staged-terminal"), + ).toBe(false); + expect(await storage.readText(sessionFile)).toBe(before); + + await manager.appendEntriesAtomically(() => manager.appendCustomEntry("committed-terminal")); + expect(manager.getBranch().at(-1)).toMatchObject({ type: "custom", customType: "committed-terminal" }); + await manager.close(); + }); + + it("reparents and durably preserves a concurrent append when the staged batch rolls back", async () => { + const storage = new GatedAtomicFailureStorage(); + const manager = SessionManager.create("/cwd", "/sessions", storage); + const rootId = manager.appendCustomEntry("root"); + await manager.ensureOnDisk(); + const sessionFile = manager.getSessionFile(); + if (!sessionFile) throw new Error("Expected session file"); + const notifiedIds: string[] = []; + manager.onEntryAppended = entry => notifiedIds.push(entry.id); + const failure = storage.failNextAtomicWrite(new Error("batch publish failed")); + let stagedId = ""; + const commit = manager.appendEntriesAtomically(() => { + stagedId = manager.appendCustomEntry("staged-terminal"); + }); + await failure.started; + const concurrentId = manager.appendCustomEntry("concurrent-survivor"); + failure.release(); + + await expect(commit).rejects.toThrow("batch publish failed"); + expect(manager.getEntries().some(entry => entry.id === stagedId)).toBe(false); + expect(manager.getEntries().find(entry => entry.id === concurrentId)?.parentId).toBe(rootId); + expect(manager.getBranch().at(-1)?.id).toBe(concurrentId); + expect(notifiedIds).toEqual([concurrentId]); + + const content = await storage.readText(sessionFile); + expect(content).not.toContain('"customType":"staged-terminal"'); + expect(content).toContain('"customType":"concurrent-survivor"'); + await manager.close(); + }); + + it("repairs authoritative rollback after commit-then-throw, including a rejecting repair acknowledgement", async () => { + const storage = new ScriptedAtomicFailureStorage(); + const manager = SessionManager.create("/cwd", "/sessions", storage); + manager.appendCustomEntry("root"); + await manager.ensureOnDisk(); + const sessionFile = manager.getSessionFile(); + if (!sessionFile) throw new Error("Expected session file"); + const before = await storage.readText(sessionFile); + storage.behaviors.push( + { commit: true, error: new Error("batch committed but acknowledgement failed") }, + { commit: true, error: new Error("repair committed but acknowledgement failed") }, + ); + + await expect(manager.appendEntriesAtomically(() => manager.appendCustomEntry("staged-terminal"))).rejects.toThrow( + "batch committed but acknowledgement failed", + ); + + expect(await storage.readText(sessionFile)).toBe(before); + expect( + manager.getEntries().some(entry => entry.type === "custom" && entry.customType === "staged-terminal"), + ).toBe(false); + await manager.appendEntriesAtomically(() => manager.appendCustomEntry("retry-terminal")); + expect(await storage.readText(sessionFile)).toContain('"customType":"retry-terminal"'); + await manager.close(); + }); + + it("reserves concurrent atomic batches FIFO before their first await", async () => { + const storage = new GatedAtomicFailureStorage(); + const manager = SessionManager.create("/cwd", "/sessions", storage); + manager.appendCustomEntry("root"); + await manager.ensureOnDisk(); + const sessionFile = manager.getSessionFile(); + if (!sessionFile) throw new Error("Expected session file"); + const failure = storage.failNextAtomicWrite(new Error("first batch failed")); + const first = manager.appendEntriesAtomically(() => manager.appendCustomEntry("batch-a")); + await failure.started; + let secondCallbackRan = false; + const second = manager.appendEntriesAtomically(() => { + secondCallbackRan = true; + return manager.appendCustomEntry("batch-b"); + }); + await Promise.resolve(); + expect(secondCallbackRan).toBe(false); + failure.release(); + + await expect(first).rejects.toThrow("first batch failed"); + await second; + expect(secondCallbackRan).toBe(true); + const content = await storage.readText(sessionFile); + expect(content).not.toContain('"customType":"batch-a"'); + expect(content).toContain('"customType":"batch-b"'); + await manager.close(); + }); + + it("latches a typed indeterminate error when rollback repair cannot be verified", async () => { + const storage = new ScriptedAtomicFailureStorage(); + const manager = SessionManager.create("/cwd", "/sessions", storage); + manager.appendCustomEntry("root"); + await manager.ensureOnDisk(); + const sessionFile = manager.getSessionFile(); + if (!sessionFile) throw new Error("Expected session file"); + storage.behaviors.push( + { commit: true, error: new Error("batch committed but acknowledgement failed") }, + { commit: false, error: new Error("authoritative repair failed before commit") }, + ); + + const failure = await manager + .appendEntriesAtomically(() => manager.appendCustomEntry("possibly-durable-terminal")) + .catch(error => error); + + expect(failure).toBeInstanceOf(SessionPersistenceIndeterminateError); + expect((failure as SessionPersistenceIndeterminateError).errors).toHaveLength(3); + expect(await storage.readText(sessionFile)).toContain('"customType":"possibly-durable-terminal"'); + expect( + manager + .getEntries() + .some(entry => entry.type === "custom" && entry.customType === "possibly-durable-terminal"), + ).toBe(false); + await expect(manager.flush()).rejects.toBeInstanceOf(SessionPersistenceIndeterminateError); + + await manager.recoverPersistenceFromCurrentState(); + expect(await storage.readText(sessionFile)).not.toContain('"customType":"possibly-durable-terminal"'); + await manager.close(); + }); +}); + +class CommitThenThrowIndexedBackend implements SessionStorageBackend { + content: string | null = null; + readonly atomicWriteStarted = Promise.withResolvers(); + readonly releaseAtomicWrite = Promise.withResolvers(); + readonly newerWriteFinished = Promise.withResolvers(); + #writeCount = 0; + #atomicWriteFailed = false; + + async init(): Promise {} + + async loadIndex(): Promise<[]> { + return []; + } + + async readFull(): Promise { + if (this.#atomicWriteFailed) await this.newerWriteFinished.promise; + return this.content; + } + + async readSlices(): Promise<[string, string]> { + return [this.content ?? "", this.content ?? ""]; + } + + async writeFull(_path: string, content: string): Promise { + this.#writeCount++; + if (this.#writeCount === 2) { + this.atomicWriteStarted.resolve(); + await this.releaseAtomicWrite.promise; + this.content = content; + this.#atomicWriteFailed = true; + throw new Error("atomic write committed but acknowledgement failed"); + } + this.content = content; + if (this.#writeCount === 3) this.newerWriteFinished.resolve(); + } + + async append(_path: string, line: string): Promise { + this.content = (this.content ?? "") + line; + } + + async updateSessionTitle(): Promise {} + + async truncate(): Promise { + this.content = ""; + } + + async remove(): Promise { + this.content = null; + } + + async move(): Promise {} +} + +describe("IndexedSessionStorage atomic readback", () => { + it("preserves a newer synchronous takeover when failed-write readback sees its body", async () => { + const backend = new CommitThenThrowIndexedBackend(); + const storage = new IndexedSessionStorage(backend); + const sessionPath = "/sessions/current.jsonl"; + const newerBody = "newer-body-that-is-long"; + await storage.initialize(); + await storage.writeText(sessionPath, "old"); + + const failure = storage.writeTextAtomic(sessionPath, "atomic-body").catch(error => error); + await backend.atomicWriteStarted.promise; + storage.writeTextSync(sessionPath, newerBody); + backend.releaseAtomicWrite.resolve(); + + expect(await failure).toMatchObject({ message: "atomic write committed but acknowledgement failed" }); + await storage.drain(); + expect(await storage.readText(sessionPath)).toBe(newerBody); + expect(storage.statSync(sessionPath).size).toBe(Buffer.byteLength(newerBody)); + }); +}); diff --git a/packages/coding-agent/test/session/redis-session-storage-manager.test.ts b/packages/coding-agent/test/session/redis-session-storage-manager.test.ts index 241d0657e..f74c67055 100644 --- a/packages/coding-agent/test/session/redis-session-storage-manager.test.ts +++ b/packages/coding-agent/test/session/redis-session-storage-manager.test.ts @@ -37,6 +37,38 @@ function createFakeRedis(): FakeRedis { return { strings, hashes, + async send(command, args) { + if (command !== "EVAL") throw new Error(`Unsupported Redis command: ${command}`); + const script = args[0] ?? ""; + const keyCount = Number(args[1] ?? "0"); + const keys = args.slice(2, 2 + keyCount); + const argv = args.slice(2 + keyCount); + if (script.includes("OMP_WRITE_FULL")) { + const [fileKey, metaKey, titleKey] = keys; + const [content, filePath, mtimeMs, hasTitle, title] = argv; + strings.set(fileKey, content); + getHash(metaKey).set(filePath, mtimeMs); + if (hasTitle === "1") getHash(titleKey).set(filePath, title); + else getHash(titleKey).delete(filePath); + return 1; + } + if (script.includes("OMP_APPEND")) { + const [fileKey, metaKey] = keys; + const [line, filePath, mtimeMs] = argv; + const next = (strings.get(fileKey) ?? "") + line; + strings.set(fileKey, next); + getHash(metaKey).set(filePath, mtimeMs); + return Buffer.byteLength(next, "utf-8"); + } + if (script.includes("OMP_UPDATE_TITLE")) { + const [metaKey, titleKey] = keys; + const [filePath, mtimeMs, title] = argv; + getHash(metaKey).set(filePath, mtimeMs); + getHash(titleKey).set(filePath, title); + return 1; + } + throw new Error("Unsupported Redis script"); + }, async get(key) { return strings.has(key) ? (strings.get(key) as string) : null; }, diff --git a/packages/coding-agent/test/session/redis-session-storage.test.ts b/packages/coding-agent/test/session/redis-session-storage.test.ts index b151f42e5..18247f778 100644 --- a/packages/coding-agent/test/session/redis-session-storage.test.ts +++ b/packages/coding-agent/test/session/redis-session-storage.test.ts @@ -68,6 +68,45 @@ function createFakeRedis(): FakeRedis { queue.push(error); failures.set(method, queue); }, + async send(command, args) { + record("send", [command, args]); + checkFailure("send"); + if (command !== "EVAL") throw new Error(`Unsupported Redis command: ${command}`); + const script = args[0] ?? ""; + const keyCount = Number(args[1] ?? "0"); + const keys = args.slice(2, 2 + keyCount); + const argv = args.slice(2 + keyCount); + if (script.includes("OMP_WRITE_FULL")) { + checkFailure("set"); + checkFailure("hset"); + const [fileKey, metaKey, titleKey] = keys; + const [content, filePath, mtimeMs, hasTitle, title] = argv; + strings.set(fileKey, content); + getHash(metaKey).set(filePath, mtimeMs); + if (hasTitle === "1") getHash(titleKey).set(filePath, title); + else getHash(titleKey).delete(filePath); + return 1; + } + if (script.includes("OMP_APPEND")) { + checkFailure("append"); + checkFailure("hset"); + const [fileKey, metaKey] = keys; + const [line, filePath, mtimeMs] = argv; + const next = (strings.get(fileKey) ?? "") + line; + strings.set(fileKey, next); + getHash(metaKey).set(filePath, mtimeMs); + return Buffer.byteLength(next, "utf-8"); + } + if (script.includes("OMP_UPDATE_TITLE")) { + checkFailure("hset"); + const [metaKey, titleKey] = keys; + const [filePath, mtimeMs, title] = argv; + getHash(metaKey).set(filePath, mtimeMs); + getHash(titleKey).set(filePath, title); + return 1; + } + throw new Error("Unsupported Redis script"); + }, async get(key) { record("get", [key]); checkFailure("get"); @@ -185,6 +224,21 @@ describe("RedisSessionStorage", () => { expect(typeof stat.mtimeMs).toBe("number"); }); + it("commits file content and metadata through one atomic Redis script", async () => { + const storage = await RedisSessionStorage.create({ client: redis }); + const sessionPath = "/sessions/p/atomic.jsonl"; + await storage.writeText(sessionPath, "old\n"); + const oldMtime = redis.hashes.get("omp:sessions:meta")?.get(sessionPath); + redis.failNext("send", new Error("EVAL transport failed")); + + await expect(storage.writeTextAtomic(sessionPath, "new\n")).rejects.toThrow("EVAL transport failed"); + + expect(redis.strings.get(`omp:sessions:file:${sessionPath}`)).toBe("old\n"); + expect(redis.hashes.get("omp:sessions:meta")?.get(sessionPath)).toBe(oldMtime); + expect(storage.statSync(sessionPath).size).toBe(4); + expect(redis.calls.some(call => call.method === "send" && call.args[0] === "EVAL")).toBe(true); + }); + it("create() warms the metadata index with STRLEN and never GETs full content", async () => { redis.strings.set("omp:sessions:file:/sessions/p/huge.jsonl", "0123456789"); redis.hashes.set("omp:sessions:meta", new Map([["/sessions/p/huge.jsonl", String(Date.now())]])); diff --git a/packages/coding-agent/test/task/executor-wall-clock.test.ts b/packages/coding-agent/test/task/executor-wall-clock.test.ts index 0157249ba..dd54fd79c 100644 --- a/packages/coding-agent/test/task/executor-wall-clock.test.ts +++ b/packages/coding-agent/test/task/executor-wall-clock.test.ts @@ -2,7 +2,8 @@ import { afterEach, describe, expect, it, vi } from "bun:test"; import type { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import type { LoadExtensionsResult } from "@oh-my-pi/pi-coding-agent/extensibility/extensions/types"; -import type { CreateAgentSessionResult } from "@oh-my-pi/pi-coding-agent/sdk"; +import { AgentRegistry } from "@oh-my-pi/pi-coding-agent/registry/agent-registry"; +import type { CreateAgentSessionOptions, CreateAgentSessionResult } from "@oh-my-pi/pi-coding-agent/sdk"; import * as sdkModule from "@oh-my-pi/pi-coding-agent/sdk"; import type { AgentSession, AgentSessionEvent, PromptOptions } from "@oh-my-pi/pi-coding-agent/session/agent-session"; import { runSubprocess } from "@oh-my-pi/pi-coding-agent/task/executor"; @@ -71,6 +72,7 @@ function mockCreateAgentSession(session: AgentSession) { describe("runSubprocess wall clock (task.maxRuntimeMs)", () => { afterEach(() => { vi.restoreAllMocks(); + AgentRegistry.resetGlobalForTests(); }); const baseAgent: AgentDefinition = { @@ -197,6 +199,67 @@ describe("runSubprocess wall clock (task.maxRuntimeMs)", () => { expect(promptCalls).toBe(0); }); + it("a cancelled late initializer cannot replace a newer same-id worker", async () => { + AgentRegistry.resetGlobalForTests(); + const registry = AgentRegistry.global(); + const creationGate = Promise.withResolvers(); + const creationStarted = Promise.withResolvers(); + const lateDisposed = Promise.withResolvers(); + const lateSession = { + dispose: async () => lateDisposed.resolve(), + } as unknown as AgentSession; + let lateInstall = registry.get("late-generation"); + vi.spyOn(sdkModule, "createAgentSession").mockImplementation(async (options = {}) => { + creationStarted.resolve(options); + await creationGate.promise; + lateInstall = registry.registerIfAvailable( + { + id: "late-generation", + displayName: "late A", + kind: "sub", + parentId: "Main", + session: null, + status: "running", + }, + options.expectedAgentRef ?? null, + ); + return { + session: lateSession, + extensionsResult: {} as unknown as LoadExtensionsResult, + setToolUIContext: () => {}, + eventBus: new EventBus(), + } satisfies CreateAgentSessionResult; + }); + const abortController = new AbortController(); + const run = runSubprocess({ + ...baseOptions, + id: "late-generation", + settings: Settings.isolated({ "task.maxRuntimeMs": 0 }), + signal: abortController.signal, + }); + const creationOptions = await creationStarted.promise; + expect(creationOptions.expectedAgentRef).toBeNull(); + abortController.abort(); + const cancelled = await run; + expect(cancelled.aborted).toBe(true); + + const replacementSession = { dispose: async () => {} } as unknown as AgentSession; + const replacement = registry.register({ + id: "late-generation", + displayName: "replacement B", + kind: "sub", + parentId: "Main", + session: replacementSession, + status: "idle", + }); + creationGate.resolve(); + await lateDisposed.promise; + + expect(lateInstall).toBeUndefined(); + expect(registry.get("late-generation")).toBe(replacement); + expect(replacement).toMatchObject({ status: "idle", session: replacementSession }); + }); + it("a late successful yield does not flip a timed-out run to success", async () => { // A hung subagent emits a successful `yield` event during teardown (after // the timer has already aborted). Without the fix, `hasYield=true` would diff --git a/packages/coding-agent/test/task/output-manager.test.ts b/packages/coding-agent/test/task/output-manager.test.ts index de33b13ad..426665304 100644 --- a/packages/coding-agent/test/task/output-manager.test.ts +++ b/packages/coding-agent/test/task/output-manager.test.ts @@ -5,8 +5,8 @@ import { TempDir } from "@oh-my-pi/pi-utils"; // Contract: subagent output ids are the requested name, used verbatim the first // time and suffixed (`-2`, `-3`, …) only when the same name recurs. A parent -// prefix nests ids under it. On resume the manager scans existing `.md` outputs -// so it never reuses a name that would clobber a previously written output. +// prefix nests ids under it. On resume the manager scans existing output and +// child-session files so it never reuses a name that would clobber prior state. describe("AgentOutputManager", () => { it("uses the requested name verbatim and suffixes only on repeat", async () => { @@ -42,14 +42,24 @@ describe("AgentOutputManager", () => { const dir = tmp.path(); await Bun.write(path.join(dir, "Anna.md"), "prior"); await Bun.write(path.join(dir, "Anna-2.md"), "prior"); + await Bun.write(path.join(dir, "Bob.jsonl"), "persisted child session"); // Unrelated tool artifacts (numeric `.log` ids) must not be mistaken for names. await Bun.write(path.join(dir, "7.bash.log"), "noise"); const mgr = new AgentOutputManager(() => dir); expect(await mgr.allocate("Anna")).toBe("Anna-3"); - // A name with no file on disk is still pristine. - expect(await mgr.allocate("Bob")).toBe("Bob"); + // A child JSONL without a result markdown still reserves its worker id. + expect(await mgr.allocate("Bob")).toBe("Bob-2"); + }); + + it("awaits one disk scan before concurrent allocations", async () => { + using tmp = TempDir.createSync("@omp-output-manager-"); + const dir = tmp.path(); + await Bun.write(path.join(dir, "Anna.jsonl"), "persisted child session"); + const mgr = new AgentOutputManager(() => dir); + + expect(await Promise.all([mgr.allocate("Anna"), mgr.allocate("Anna")])).toEqual(["Anna-2", "Anna-3"]); }); it("only counts files within its own prefix scope on resume", async () => { @@ -66,6 +76,14 @@ describe("AgentOutputManager", () => { expect(await mgr.allocate("Dave")).toBe("Anna.Dave"); }); + it("reserves lifecycle-known ids that no longer have files on disk", async () => { + const mgr = new AgentOutputManager(() => null); + await mgr.reserve(["Gone", "Other"]); + + expect(await mgr.allocate("Gone")).toBe("Gone-2"); + expect(await mgr.allocate("Fresh")).toBe("Fresh"); + }); + it("reserves the advisor transcript stem so a task can't clobber __advisor.jsonl", async () => { const mgr = new AgentOutputManager(() => null); // A subagent allocated `__advisor` would write `__advisor.jsonl`, colliding with diff --git a/packages/coding-agent/test/vibe/vibe-runtime.test.ts b/packages/coding-agent/test/vibe/vibe-runtime.test.ts index a2e42c607..3eeb044e2 100644 --- a/packages/coding-agent/test/vibe/vibe-runtime.test.ts +++ b/packages/coding-agent/test/vibe/vibe-runtime.test.ts @@ -16,26 +16,178 @@ * acknowledges its delivery so the result is not delivered twice. * 5. `kill` cancels the in-flight turn job and releases the worker session. */ + import { afterEach, beforeEach, describe, expect, it, vi } from "bun:test"; +import * as fs from "node:fs/promises"; +import * as os from "node:os"; +import * as path from "node:path"; import { AsyncJobManager } from "@oh-my-pi/pi-coding-agent/async/job-manager"; import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { AgentLifecycleManager } from "@oh-my-pi/pi-coding-agent/registry/agent-lifecycle"; import { AgentRegistry } from "@oh-my-pi/pi-coding-agent/registry/agent-registry"; import type { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; +import { + SessionManager, + SessionPersistenceIndeterminateError, +} from "@oh-my-pi/pi-coding-agent/session/session-manager"; +import { + FileSessionStorage, + type SessionStorage, + type SessionStorageWriter, + type WriteTextAtomicOptions, +} from "@oh-my-pi/pi-coding-agent/session/session-storage"; import * as executorModule from "@oh-my-pi/pi-coding-agent/task/executor"; import type { AgentProgress, SingleResult } from "@oh-my-pi/pi-coding-agent/task/types"; import type { ToolSession } from "@oh-my-pi/pi-coding-agent/tools"; import { VibeSessionRegistry } from "@oh-my-pi/pi-coding-agent/vibe/runtime"; -function createSession(options: { manager?: AsyncJobManager } = {}): ToolSession { +const PERSISTED_WORKER_SYSTEM_PROMPT = "Persisted vibe worker"; +const PERSISTED_WORKER_TOOLS = ["read", "yield"]; +const INITIAL_VIBE_TASK = "Complete the first persisted turn."; +const FOLLOW_UP_VIBE_TASK = "Continue from the persisted conversation."; +const RESTORED_VIBE_RESPONSE = "Continued from persisted context."; + +async function fileExists(file: string): Promise { + try { + return (await fs.stat(file)).isFile(); + } catch { + return false; + } +} + +type AtomicWriteHook = (commit: () => Promise) => Promise; + +class FaultInjectingSessionStorage extends FileSessionStorage { + readonly atomicWriteHooks: AtomicWriteHook[] = []; + atomicWriteAttempts = 0; + failedWriterClosed = false; + #appendFault: { error: Error; prefixBytes: number } | undefined; + + failNextAppendWithPrefix(error: Error, prefixBytes: number): void { + this.#appendFault = { error, prefixBytes }; + } + + override async writeTextAtomic(filePath: string, content: string, options?: WriteTextAtomicOptions): Promise { + this.atomicWriteAttempts++; + const hook = this.atomicWriteHooks.shift(); + const commit = () => super.writeTextAtomic(filePath, content, options); + if (hook) await hook(commit); + else await commit(); + } + + override openWriter( + filePath: string, + options?: { flags?: "a" | "w"; onError?: (error: Error) => void }, + ): SessionStorageWriter { + const inner = super.openWriter(filePath, options); + let injectedError: Error | undefined; + let pendingFault: Promise = Promise.resolve(); + let faulted = false; + return { + append: line => { + const fault = this.#appendFault; + if (!fault) return inner.append(line); + this.#appendFault = undefined; + faulted = true; + const prefix = Buffer.from(line).subarray(0, fault.prefixBytes); + const operation = fs.appendFile(filePath, prefix).then(() => { + injectedError = fault.error; + options?.onError?.(fault.error); + throw fault.error; + }); + pendingFault = operation.catch(() => {}); + return operation; + }, + flush: async () => { + await pendingFault; + if (injectedError) throw injectedError; + await inner.flush(); + }, + isOpen: () => inner.isOpen(), + close: async () => { + await pendingFault; + await inner.close(); + if (faulted) this.failedWriterClosed = true; + }, + getError: () => injectedError ?? inner.getError(), + }; + } +} + +class SwitchGatedSessionStorage extends FaultInjectingSessionStorage { + #readGate: + | { + filePath: string; + started: ReturnType>; + release: ReturnType>; + } + | undefined; + + gateNextRead(filePath: string): { started: Promise; release: () => void } { + const started = Promise.withResolvers(); + const release = Promise.withResolvers(); + this.#readGate = { filePath, started, release }; + return { started: started.promise, release: release.resolve }; + } + + override async readTextSlices( + filePath: string, + prefixBytes: number, + suffixBytes: number, + ): Promise<[string, string]> { + const gate = this.#readGate; + if (gate?.filePath === filePath) { + this.#readGate = undefined; + gate.started.resolve(); + await gate.release.promise; + } + return super.readTextSlices(filePath, prefixBytes, suffixBytes); + } +} + +interface TestSessionOptions { + manager?: AsyncJobManager; + sessionManager?: SessionManager; + ownerId?: string; + parentSessionId?: string; +} + +function createSession(options: TestSessionOptions = {}): ToolSession { + const sessionManager = options.sessionManager; return { - cwd: "/tmp", + cwd: sessionManager?.getCwd() ?? "/tmp", hasUI: false, settings: Settings.isolated({}), - getSessionFile: () => null, + getSessionFile: () => sessionManager?.getSessionFile() ?? null, + getSessionId: () => options.parentSessionId ?? sessionManager?.getSessionId() ?? "vibe-test-parent", + getAgentId: () => options.ownerId ?? "Main", + getArtifactsDir: () => sessionManager?.getArtifactsDir() ?? null, getSessionSpawns: () => "*", + sessionManager, asyncJobManager: options.manager, - } as unknown as ToolSession; + }; +} + +interface PersistWorkerOptions { + cwd: string; + artifactsDir: string; + id: string; + task: string; +} + +async function persistWorkerSession(options: PersistWorkerOptions): Promise { + const childSessionFile = path.join(options.artifactsDir, `${options.id}.jsonl`); + const childManager = SessionManager.create(options.cwd, options.artifactsDir); + await childManager.setSessionFile(childSessionFile); + childManager.appendSessionInit({ + systemPrompt: PERSISTED_WORKER_SYSTEM_PROMPT, + task: options.task, + tools: [...PERSISTED_WORKER_TOOLS], + spawns: "", + }); + await childManager.flush(); + await childManager.close(); + return childSessionFile; } function makeResult(id: string, overrides: Partial = {}): SingleResult { @@ -66,6 +218,10 @@ function deferred(): Deferred { return { promise, resolve }; } +async function flushMicrotasks(): Promise { + for (let index = 0; index < 5; index++) await Promise.resolve(); +} + async function pollUntil(predicate: () => boolean, timeoutMs = 2000): Promise { const start = Date.now(); while (!predicate()) { @@ -80,7 +236,7 @@ async function pollUntil(predicate: () => boolean, timeoutMs = 2000): Promise void | Promise } = {}) { const listeners = new Set<(event: unknown) => void>(); const prompts: string[] = []; const steers: string[] = []; @@ -118,6 +274,7 @@ function createFakeWorkerSession(options: { streaming?: boolean } = {}) { async abort(): Promise {}, async dispose(): Promise { disposed = true; + await options.onDispose?.(); }, }; return { @@ -171,6 +328,18 @@ function progressSnapshot(id: string, overrides: Partial = {}): A describe("vibe session registry", () => { const managers: AsyncJobManager[] = []; + const persistedManagers: SessionManager[] = []; + const tempRoots: string[] = []; + + async function createPersistedParent(storage?: SessionStorage): Promise { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "omp-vibe-resume-")); + tempRoots.push(root); + const cwd = path.join(root, "workspace"); + await fs.mkdir(cwd, { recursive: true }); + const manager = SessionManager.create(cwd, path.join(root, "sessions"), storage); + persistedManagers.push(manager); + return manager; + } function createManager(): AsyncJobManager { const manager = new AsyncJobManager({ onJobComplete: () => {} }); @@ -178,6 +347,74 @@ describe("vibe session registry", () => { return manager; } + function installPersistedSpawnMock(): void { + vi.spyOn(executorModule, "runSubprocess").mockImplementation(async options => { + const artifactsDir = options.artifactsDir; + if (!artifactsDir) throw new Error("Persisted vibe test requires an artifacts directory"); + const parentSessionFile = options.sessionFile; + if (!parentSessionFile) throw new Error("Persisted vibe test requires a parent session file"); + const snapshot = await SessionManager.open(parentSessionFile, undefined, undefined, { + suppressBreadcrumb: true, + }); + try { + const actions = snapshot.getEntries().flatMap(entry => { + if (entry.type !== "custom" || typeof entry.data !== "object" || entry.data === null) return []; + const data = entry.data as Record; + return data.id === options.id && typeof data.action === "string" ? [data.action] : []; + }); + expect(actions).toEqual(["spawn", "turn-started"]); + } finally { + await snapshot.close(); + } + const childSessionFile = await persistWorkerSession({ + cwd: options.cwd, + artifactsDir, + id: options.id, + task: options.task, + }); + const worker = createFakeWorkerSession(); + AgentRegistry.global().register({ + id: options.id, + displayName: options.id, + kind: "sub", + parentId: options.parentAgentId ?? "Main", + session: worker.session, + sessionFile: childSessionFile, + status: "running", + }); + AgentRegistry.global().setStatus(options.id, "idle"); + AgentLifecycleManager.global().adopt(options.id, { + idleTtlMs: 0, + revive: async () => worker.session, + }); + return makeResult(options.id, { output: "Persisted first turn." }); + }); + } + + function installPersistedReviver(capture: { sessionFile?: string; prompts?: string[] }): void { + AgentLifecycleManager.global().setPersistedSubagentReviverFactory(async ref => { + if (!ref.sessionFile) return undefined; + const persisted = await SessionManager.peekSessionInit(ref.sessionFile); + if (!persisted?.init) return undefined; + const worker = createFakeWorkerSession(); + worker.prompts.push(persisted.init.task); + worker.setScript({ + events: yieldTurnEvents({ report: RESTORED_VIBE_RESPONSE }), + responseText: RESTORED_VIBE_RESPONSE, + }); + capture.sessionFile = ref.sessionFile; + capture.prompts = worker.prompts; + return async () => worker.session; + }, 0); + } + + async function simulateProcessBoundary(): Promise { + await AgentLifecycleManager.global().dispose(); + VibeSessionRegistry.resetGlobalForTests(); + AgentLifecycleManager.resetGlobalForTests(); + AgentRegistry.resetGlobalForTests(); + } + beforeEach(() => { AgentRegistry.resetGlobalForTests(); AgentLifecycleManager.resetGlobalForTests(); @@ -189,6 +426,12 @@ describe("vibe session registry", () => { for (const manager of managers.splice(0)) { await manager.dispose({ timeoutMs: 1000 }); } + for (const manager of persistedManagers.splice(0)) { + await manager.close(); + } + for (const root of tempRoots.splice(0)) { + await fs.rm(root, { recursive: true, force: true }); + } VibeSessionRegistry.resetGlobalForTests(); AgentLifecycleManager.resetGlobalForTests(); AgentRegistry.resetGlobalForTests(); @@ -231,7 +474,7 @@ describe("vibe session registry", () => { // Ack is immediate: the job is still running behind the gate. const job = manager.getJob(jobId)!; expect(job.status).toBe("running"); - expect(registry.screens("Main")[0]?.cli).toBe("fast"); + expect(registry.screens(session)[0]?.cli).toBe("fast"); gate.resolve(); await job.promise; @@ -245,11 +488,46 @@ describe("vibe session registry", () => { expect(text.indexOf("read(src/foo.ts)")).toBeLessThan(text.indexOf("bash(bun test)")); expect(text).toContain("Implemented the widget."); // Session survives the turn, addressable for follow-ups. - const entry = registry.screens("Main")[0]!; + const entry = registry.screens(session)[0]!; expect(entry.state).toBe("idle"); expect(entry.turns).toBe(1); }); + it("retains a failed spawn when its tombstone cannot flush and retries it on mode exit", async () => { + const parentManager = await createPersistedParent(); + parentManager.appendModeChange("vibe"); + const jobs = createManager(); + const session = createSession({ manager: jobs, sessionManager: parentManager }); + const originalFlush = parentManager.flush.bind(parentManager); + const flush = vi + .spyOn(parentManager, "flush") + .mockImplementationOnce(originalFlush) + .mockRejectedValueOnce(new Error("tombstone write failed")); + const register = vi.spyOn(jobs, "register").mockImplementation(() => { + throw new Error("job registry failed"); + }); + + await expect( + VibeSessionRegistry.global().spawn(session, { + cli: "fast", + name: "failed-spawn", + prompt: INITIAL_VIBE_TASK, + }), + ).rejects.toThrow("tombstone write failed"); + + expect(VibeSessionRegistry.global().screens(session)[0]).toMatchObject({ + id: "failed-spawn", + state: "dead", + lastActivity: "spawn failed", + }); + expect(parentManager.buildSessionContext().mode).toBe("vibe"); + flush.mockRestore(); + register.mockRestore(); + expect(await VibeSessionRegistry.global().killAll(session)).toBe(1); + expect(VibeSessionRegistry.global().listIds(session)).toEqual([]); + expect(parentManager.buildSessionContext().mode).toBe("none"); + }); + it("send steers a streaming mid-turn worker and queues for a non-steerable one", async () => { const gate = deferred(); const fake = createFakeWorkerSession({ streaming: true }); @@ -287,7 +565,7 @@ describe("vibe session registry", () => { fake.setStreaming(false); const queued = await registry.send(session, { session: "Good", message: "Then write tests." }); expect(queued.mode).toBe("queued"); - expect(registry.screens("Main")[0]?.queued).toBe(1); + expect(registry.screens(session)[0]?.queued).toBe(1); // Settling the turn drains the queue into an automatic follow-up turn. gate.resolve(); @@ -340,7 +618,1099 @@ describe("vibe session registry", () => { expect(text).toContain('turn="2"'); expect(text).toContain("edit(src/foo.ts)"); expect(text).toContain("Renamed everything."); - expect(registry.screens("Main")[0]?.turns).toBe(2); + expect(registry.screens(session)[0]?.turns).toBe(2); + }); + + it("rehydrates an idle worker after a process boundary and continues turn two with prior context", async () => { + installPersistedSpawnMock(); + const parentManager = await createPersistedParent(); + parentManager.appendModeChange("vibe"); + const parentSessionFile = parentManager.getSessionFile(); + if (!parentSessionFile) throw new Error("Persisted parent session file was not created"); + expect(await fileExists(parentSessionFile)).toBe(false); + const firstManager = createManager(); + const firstSession = createSession({ manager: firstManager, sessionManager: parentManager }); + const firstRegistry = VibeSessionRegistry.global(); + const spawned = await firstRegistry.spawn(firstSession, { + cli: "fast", + name: "push-fixes", + prompt: INITIAL_VIBE_TASK, + }); + expect(await fileExists(parentSessionFile)).toBe(true); + await firstManager.getJob(spawned.jobId)!.promise; + expect(firstRegistry.screens(firstSession)[0]).toMatchObject({ id: "push-fixes", state: "idle", turns: 1 }); + + const wrongOwner = createSession({ manager: firstManager, sessionManager: parentManager, ownerId: "Other" }); + await expect( + firstRegistry.send(wrongOwner, { session: "push-fixes", message: FOLLOW_UP_VIBE_TASK }), + ).rejects.toThrow('Unknown vibe session "push-fixes"'); + const wrongParent = createSession({ + manager: firstManager, + sessionManager: parentManager, + parentSessionId: "different-parent", + }); + await expect( + firstRegistry.send(wrongParent, { session: "push-fixes", message: FOLLOW_UP_VIBE_TASK }), + ).rejects.toThrow('Unknown vibe session "push-fixes"'); + + await parentManager.flush(); + await parentManager.close(); + await simulateProcessBoundary(); + const resumedManager = await SessionManager.open(parentSessionFile, undefined, undefined, { + suppressBreadcrumb: true, + }); + persistedManagers.push(resumedManager); + const resumedJobs = createManager(); + const resumedSession = createSession({ manager: resumedJobs, sessionManager: resumedManager }); + const revived: { sessionFile?: string; prompts?: string[] } = {}; + installPersistedReviver(revived); + + const resumedRegistry = VibeSessionRegistry.global(); + expect(await resumedRegistry.rehydrate(resumedSession)).toBe(1); + expect(AgentRegistry.global().get("push-fixes")?.status).toBe("parked"); + expect(resumedRegistry.screens(resumedSession)[0]).toMatchObject({ + id: "push-fixes", + state: "idle", + turns: 1, + }); + + const outcome = await resumedRegistry.send(resumedSession, { + session: "push-fixes", + message: FOLLOW_UP_VIBE_TASK, + }); + expect(outcome.mode).toBe("turn"); + const turnJob = resumedJobs.getJob(outcome.jobId!)!; + await turnJob.promise; + expect(revived.sessionFile).toBe(path.join(parentSessionFile.slice(0, -6), "push-fixes.jsonl")); + expect(revived.prompts).toEqual([INITIAL_VIBE_TASK, FOLLOW_UP_VIBE_TASK]); + expect(turnJob.resultText).toContain('turn="2"'); + expect(turnJob.resultText).toContain(RESTORED_VIBE_RESPONSE); + }); + + it("suspends an idle persisted worker for same-process disposal and resumes turn two", async () => { + installPersistedSpawnMock(); + const parentManager = await createPersistedParent(); + parentManager.appendModeChange("vibe"); + const parentSessionFile = parentManager.getSessionFile(); + if (!parentSessionFile) throw new Error("Persisted parent session file was not created"); + const firstJobs = createManager(); + const firstSession = createSession({ manager: firstJobs, sessionManager: parentManager }); + const registry = VibeSessionRegistry.global(); + const spawned = await registry.spawn(firstSession, { + cli: "fast", + name: "idle-dispose", + prompt: INITIAL_VIBE_TASK, + }); + await firstJobs.getJob(spawned.jobId)!.promise; + expect(await registry.suspendScope(registry.ownerScope(firstSession), firstJobs)).toBe(1); + expect(AgentRegistry.global().get("idle-dispose")).toBeUndefined(); + expect( + parentManager.getEntries().some(entry => { + if (entry.type !== "custom" || typeof entry.data !== "object" || entry.data === null) return false; + return (entry.data as Record).action === "tombstone"; + }), + ).toBe(false); + await parentManager.close(); + const reopened = await SessionManager.open(parentSessionFile, undefined, undefined, { suppressBreadcrumb: true }); + persistedManagers.push(reopened); + const resumedJobs = createManager(); + const resumedSession = createSession({ manager: resumedJobs, sessionManager: reopened }); + const revived: { sessionFile?: string; prompts?: string[] } = {}; + installPersistedReviver(revived); + expect(await registry.rehydrate(resumedSession)).toBe(1); + const outcome = await registry.send(resumedSession, { + session: "idle-dispose", + message: FOLLOW_UP_VIBE_TASK, + }); + const turnJob = resumedJobs.getJob(outcome.jobId!)!; + await turnJob.promise; + expect(revived.prompts).toEqual([INITIAL_VIBE_TASK, FOLLOW_UP_VIBE_TASK]); + expect(turnJob.resultText).toContain('turn="2"'); + }); + + it("suspends a blocked in-flight worker for fresh-process disposal without tombstoning it", async () => { + vi.spyOn(executorModule, "runSubprocess").mockImplementation(async options => { + const artifactsDir = options.artifactsDir; + if (!artifactsDir) throw new Error("Persisted vibe test requires an artifacts directory"); + const childSessionFile = await persistWorkerSession({ + cwd: options.cwd, + artifactsDir, + id: options.id, + task: options.task, + }); + const worker = createFakeWorkerSession(); + AgentRegistry.global().register({ + id: options.id, + displayName: options.id, + kind: "sub", + parentId: options.parentAgentId ?? "Main", + session: worker.session, + sessionFile: childSessionFile, + status: "running", + }); + const signal = options.signal; + if (!signal) throw new Error("Persisted blocked worker requires a cancellation signal"); + await new Promise(resolve => { + if (signal.aborted) resolve(); + else signal.addEventListener("abort", () => resolve(), { once: true }); + }); + return makeResult(options.id, { output: "Interrupted by disposal.", aborted: true }); + }); + const parentManager = await createPersistedParent(); + parentManager.appendModeChange("vibe"); + const parentSessionFile = parentManager.getSessionFile(); + if (!parentSessionFile) throw new Error("Persisted parent session file was not created"); + const firstJobs = createManager(); + const firstSession = createSession({ manager: firstJobs, sessionManager: parentManager }); + const registry = VibeSessionRegistry.global(); + await registry.spawn(firstSession, { + cli: "fast", + name: "running-dispose", + prompt: INITIAL_VIBE_TASK, + }); + await pollUntil(() => AgentRegistry.global().get("running-dispose")?.status === "running"); + expect(await registry.suspendScope(registry.ownerScope(firstSession), firstJobs)).toBe(1); + expect( + parentManager.getEntries().some(entry => { + if (entry.type !== "custom" || typeof entry.data !== "object" || entry.data === null) return false; + return (entry.data as Record).action === "tombstone"; + }), + ).toBe(false); + await parentManager.close(); + await simulateProcessBoundary(); + + const reopened = await SessionManager.open(parentSessionFile, undefined, undefined, { suppressBreadcrumb: true }); + persistedManagers.push(reopened); + const resumedJobs = createManager(); + const resumedSession = createSession({ manager: resumedJobs, sessionManager: reopened }); + const revived: { sessionFile?: string; prompts?: string[] } = {}; + installPersistedReviver(revived); + const freshRegistry = VibeSessionRegistry.global(); + expect(await freshRegistry.rehydrate(resumedSession)).toBe(1); + expect(freshRegistry.screens(resumedSession)[0]).toMatchObject({ + id: "running-dispose", + state: "idle", + turns: 1, + lastActivity: "turn 1 interrupted by process restart", + }); + const outcome = await freshRegistry.send(resumedSession, { + session: "running-dispose", + message: FOLLOW_UP_VIBE_TASK, + }); + const turnJob = resumedJobs.getJob(outcome.jobId!)!; + await turnJob.promise; + expect(revived.prompts).toEqual([INITIAL_VIBE_TASK, FOLLOW_UP_VIBE_TASK]); + expect(turnJob.resultText).toContain('turn="2"'); + }); + + it("restores an interrupted turn as idle without replay and continues only after send", async () => { + installPersistedSpawnMock(); + const parentManager = await createPersistedParent(); + parentManager.appendModeChange("vibe"); + const parentSessionFile = parentManager.getSessionFile(); + if (!parentSessionFile) throw new Error("Persisted parent session file was not created"); + const firstJobs = createManager(); + const firstSession = createSession({ manager: firstJobs, sessionManager: parentManager }); + const spawned = await VibeSessionRegistry.global().spawn(firstSession, { + cli: "fast", + name: "interrupted", + prompt: INITIAL_VIBE_TASK, + }); + await firstJobs.getJob(spawned.jobId)!.promise; + const turnStarted = parentManager.getEntries().find(entry => { + if (entry.type !== "custom" || typeof entry.data !== "object" || entry.data === null) return false; + const data = entry.data as Record; + return data.id === "interrupted" && data.action === "turn-started"; + }); + if (!turnStarted) throw new Error("Expected a persisted turn-started lifecycle event"); + parentManager.branch(turnStarted.id); + parentManager.appendModeChange("vibe"); + await parentManager.flush(); + await parentManager.close(); + await simulateProcessBoundary(); + + const resumedManager = await SessionManager.open(parentSessionFile, undefined, undefined, { + suppressBreadcrumb: true, + }); + persistedManagers.push(resumedManager); + const resumedJobs = createManager(); + const resumedSession = createSession({ manager: resumedJobs, sessionManager: resumedManager }); + const revived: { sessionFile?: string; prompts?: string[] } = {}; + installPersistedReviver(revived); + const registry = VibeSessionRegistry.global(); + expect(await registry.rehydrate(resumedSession)).toBe(1); + expect(registry.screens(resumedSession)[0]).toMatchObject({ + id: "interrupted", + state: "idle", + turns: 1, + lastActivity: "turn 1 interrupted by process restart", + }); + expect(revived).toEqual({}); + + const outcome = await registry.send(resumedSession, { + session: "interrupted", + message: FOLLOW_UP_VIBE_TASK, + }); + const turnJob = resumedJobs.getJob(outcome.jobId!)!; + await turnJob.promise; + expect(revived.prompts).toEqual([INITIAL_VIBE_TASK, FOLLOW_UP_VIBE_TASK]); + expect(turnJob.resultText).toContain('turn="2"'); + }); + + it("never lets an old lifecycle append cross into a target session during its load window", async () => { + installPersistedSpawnMock(); + const storage = new SwitchGatedSessionStorage(); + const sourceManager = await createPersistedParent(storage); + sourceManager.appendModeChange("vibe"); + const sourceSessionId = sourceManager.getSessionId(); + const sourceSessionFile = sourceManager.getSessionFile(); + if (!sourceSessionFile) throw new Error("Persisted source session file was not created"); + const jobs = createManager(); + const session = createSession({ manager: jobs, sessionManager: sourceManager }); + const registry = VibeSessionRegistry.global(); + const spawned = await registry.spawn(session, { + cli: "fast", + name: "switch-race", + prompt: INITIAL_VIBE_TASK, + }); + await jobs.getJob(spawned.jobId)!.promise; + + const targetManager = await createPersistedParent(storage); + targetManager.appendModeChange("goal"); + targetManager.appendCustomEntry("target-only"); + await targetManager.ensureOnDisk(); + const targetSessionFile = targetManager.getSessionFile(); + if (!targetSessionFile) throw new Error("Persisted target session file was not created"); + await targetManager.close(); + + const lifecycleAppendStarted = Promise.withResolvers(); + const releaseLifecycleAppend = Promise.withResolvers(); + const originalEnsureOnDisk = sourceManager.ensureOnDisk.bind(sourceManager); + let gateLifecycleAppend = true; + const ensureOnDisk = vi.spyOn(sourceManager, "ensureOnDisk").mockImplementation(async () => { + if (gateLifecycleAppend) { + gateLifecycleAppend = false; + lifecycleAppendStarted.resolve(); + await releaseLifecycleAppend.promise; + } + await originalEnsureOnDisk(); + }); + const followUp = await registry.send(session, { + session: "switch-race", + message: FOLLOW_UP_VIBE_TASK, + }); + const turnJob = jobs.getJob(followUp.jobId!)!; + await lifecycleAppendStarted.promise; + + const targetRead = storage.gateNextRead(targetSessionFile); + const switching = sourceManager.setSessionFile(targetSessionFile); + await targetRead.started; + releaseLifecycleAppend.resolve(); + await turnJob.promise.catch(() => undefined); + targetRead.release(); + await switching; + ensureOnDisk.mockRestore(); + + expect(turnJob.status).toBe("failed"); + expect(sourceManager.buildSessionContext().mode).toBe("goal"); + expect( + sourceManager + .getBranch() + .map(entry => + entry.type === "mode_change" + ? `mode:${entry.mode}` + : `custom:${entry.type === "custom" ? entry.customType : entry.type}`, + ), + ).toEqual(["mode:goal", "custom:target-only"]); + expect( + sourceManager.getEntries().some(entry => { + if (entry.type !== "custom" || entry.customType !== "vibe-session-lifecycle") return false; + if (typeof entry.data !== "object" || entry.data === null) return false; + return (entry.data as Record).parentSessionId === sourceSessionId; + }), + ).toBe(false); + expect(await fs.readFile(targetSessionFile, "utf8")).not.toContain(sourceSessionId); + }); + + it("never sends through or releases a registry ref with the wrong child session file", async () => { + installPersistedSpawnMock(); + const parentManager = await createPersistedParent(); + parentManager.appendModeChange("vibe"); + const parentSessionFile = parentManager.getSessionFile(); + if (!parentSessionFile) throw new Error("Persisted parent session file was not created"); + const jobs = createManager(); + const session = createSession({ manager: jobs, sessionManager: parentManager }); + const registry = VibeSessionRegistry.global(); + const spawned = await registry.spawn(session, { + cli: "fast", + name: "exact-ref", + prompt: INITIAL_VIBE_TASK, + }); + await jobs.getJob(spawned.jobId)!.promise; + await AgentLifecycleManager.global().release("exact-ref"); + const mismatchedFile = path.join(parentSessionFile.slice(0, -6), "another-parent.jsonl"); + AgentRegistry.global().register({ + id: "exact-ref", + displayName: "exact-ref", + kind: "sub", + parentId: "Main", + session: null, + sessionFile: mismatchedFile, + status: "parked", + }); + + await expect(registry.send(session, { session: "exact-ref", message: FOLLOW_UP_VIBE_TASK })).rejects.toThrow( + "no longer resolves to this parent session", + ); + await registry.kill(session, "exact-ref"); + expect(AgentRegistry.global().get("exact-ref")).toMatchObject({ + status: "parked", + sessionFile: mismatchedFile, + }); + }); + + it("suspends a parent without tombstones so the same worker id can move between isolated scopes", async () => { + installPersistedSpawnMock(); + const parentA = await createPersistedParent(); + const parentB = await createPersistedParent(); + parentA.appendModeChange("vibe"); + parentB.appendModeChange("vibe"); + const jobsA = createManager(); + const jobsB = createManager(); + const sessionA = createSession({ manager: jobsA, sessionManager: parentA }); + const sessionB = createSession({ manager: jobsB, sessionManager: parentB }); + const registry = VibeSessionRegistry.global(); + const workerA = await registry.spawn(sessionA, { + cli: "fast", + name: "shared-name", + prompt: INITIAL_VIBE_TASK, + }); + await jobsA.getJob(workerA.jobId)!.promise; + expect(await registry.suspendScope(registry.ownerScope(sessionA), jobsA)).toBe(1); + expect(AgentRegistry.global().get("shared-name")).toBeUndefined(); + + const workerB = await registry.spawn(sessionB, { + cli: "fast", + name: "shared-name", + prompt: INITIAL_VIBE_TASK, + }); + expect(workerB.id).toBe("shared-name"); + await jobsB.getJob(workerB.jobId)!.promise; + expect(await registry.suspendScope(registry.ownerScope(sessionB), jobsB)).toBe(1); + const revived: { sessionFile?: string; prompts?: string[] } = {}; + installPersistedReviver(revived); + expect(await registry.rehydrate(sessionA)).toBe(1); + expect(registry.listIds(sessionA)).toEqual(["shared-name"]); + await expect(registry.send(sessionB, { session: "shared-name", message: FOLLOW_UP_VIBE_TASK })).rejects.toThrow( + 'Unknown vibe session "shared-name"', + ); + + const outcome = await registry.send(sessionA, { + session: "shared-name", + message: FOLLOW_UP_VIBE_TASK, + }); + await jobsA.getJob(outcome.jobId!)!.promise; + expect(revived.prompts).toEqual([INITIAL_VIBE_TASK, FOLLOW_UP_VIBE_TASK]); + }); + + it("does not let late teardown from a suspended parent mutate a new same-id worker", async () => { + const parentA = await createPersistedParent(); + const parentB = await createPersistedParent(); + parentA.appendModeChange("vibe"); + parentB.appendModeChange("vibe"); + const parentAFile = parentA.getSessionFile(); + if (!parentAFile) throw new Error("Persisted parent A session file was not created"); + let workerA: ReturnType | undefined; + vi.spyOn(executorModule, "runSubprocess").mockImplementation(async options => { + const artifactsDir = options.artifactsDir; + if (!artifactsDir) throw new Error("Persisted vibe test requires an artifacts directory"); + const childSessionFile = await persistWorkerSession({ + cwd: options.cwd, + artifactsDir, + id: options.id, + task: options.task, + }); + const worker = createFakeWorkerSession(); + if (options.sessionFile === parentAFile) workerA = worker; + AgentRegistry.global().register({ + id: options.id, + displayName: options.id, + kind: "sub", + parentId: options.parentAgentId ?? "Main", + session: worker.session, + sessionFile: childSessionFile, + status: "running", + }); + if (options.sessionFile === parentAFile) { + const signal = options.signal; + if (!signal) throw new Error("Persisted blocked worker requires a cancellation signal"); + await new Promise(resolve => { + if (signal.aborted) resolve(); + else signal.addEventListener("abort", () => resolve(), { once: true }); + }); + return makeResult(options.id, { output: "Parent A suspended.", aborted: true }); + } + AgentRegistry.global().setStatus(options.id, "idle"); + return makeResult(options.id, { output: "Parent B finished." }); + }); + const jobsA = createManager(); + const jobsB = createManager(); + const sessionA = createSession({ manager: jobsA, sessionManager: parentA }); + const sessionB = createSession({ manager: jobsB, sessionManager: parentB }); + const registry = VibeSessionRegistry.global(); + await registry.spawn(sessionA, { cli: "fast", name: "reused", prompt: INITIAL_VIBE_TASK }); + await pollUntil(() => workerA !== undefined); + const oldRef = AgentRegistry.global().get("reused"); + if (!oldRef) throw new Error("Expected parent A worker ref"); + expect(await registry.suspendScope(registry.ownerScope(sessionA), jobsA)).toBe(1); + + const second = await registry.spawn(sessionB, { cli: "fast", name: "reused", prompt: INITIAL_VIBE_TASK }); + await jobsB.getJob(second.jobId)!.promise; + const replacement = AgentRegistry.global().get("reused"); + expect(replacement).toMatchObject({ status: "idle" }); + expect(replacement).not.toBe(oldRef); + + await executorModule.finalizeSubagentLifecycle({ + id: "reused", + session: workerA!.session, + aborted: true, + keepAlive: true, + isolated: false, + agentIdleTtlMs: 0, + reviveSession: null, + }); + await AgentLifecycleManager.global().release("reused", oldRef); + expect(AgentRegistry.global().get("reused")).toBe(replacement); + expect(registry.screens(sessionB)[0]).toMatchObject({ id: "reused", state: "idle" }); + }); + + it("rejects escaped child paths and JSONL files without persisted worker init", async () => { + installPersistedSpawnMock(); + const parentManager = await createPersistedParent(); + parentManager.appendModeChange("vibe"); + const parentSessionFile = parentManager.getSessionFile(); + if (!parentSessionFile) throw new Error("Persisted parent session file was not created"); + const jobs = createManager(); + const session = createSession({ manager: jobs, sessionManager: parentManager }); + const spawned = await VibeSessionRegistry.global().spawn(session, { + cli: "fast", + name: "valid-worker", + prompt: INITIAL_VIBE_TASK, + }); + await jobs.getJob(spawned.jobId)!.promise; + + const lifecycleEntries = parentManager.getEntries().filter(entry => { + if (entry.type !== "custom" || typeof entry.data !== "object" || entry.data === null) return false; + return (entry.data as Record).id === "valid-worker"; + }); + if (lifecycleEntries.length < 3) throw new Error("Expected complete persisted worker lifecycle events"); + const artifactsDir = parentSessionFile.slice(0, -6); + const invalidManager = SessionManager.create(parentManager.getCwd(), artifactsDir); + persistedManagers.push(invalidManager); + await invalidManager.ensureOnDisk(); + const invalidSource = invalidManager.getSessionFile(); + if (!invalidSource) throw new Error("Expected an invalid child fixture path"); + await invalidManager.close(); + await fs.rename(invalidSource, path.join(artifactsDir, "invalid-init.jsonl")); + + for (const { id, childSessionFile } of [ + { id: "escaped", childSessionFile: "../escaped.jsonl" }, + { id: "invalid-init", childSessionFile: "invalid-init.jsonl" }, + ]) { + for (const entry of lifecycleEntries) { + if (entry.type !== "custom" || typeof entry.data !== "object" || entry.data === null) continue; + const data: Record = { ...(entry.data as Record), id }; + if (data.action === "spawn") data.childSessionFile = childSessionFile; + parentManager.appendCustomEntry(entry.customType, data); + } + } + await parentManager.flush(); + await parentManager.close(); + await simulateProcessBoundary(); + + const resumedManager = await SessionManager.open(parentSessionFile, undefined, undefined, { + suppressBreadcrumb: true, + }); + persistedManagers.push(resumedManager); + const resumedSession = createSession({ manager: createManager(), sessionManager: resumedManager }); + const registry = VibeSessionRegistry.global(); + expect(await registry.rehydrate(resumedSession)).toBe(1); + expect(registry.listIds(resumedSession)).toEqual(["valid-worker"]); + expect(AgentRegistry.global().get("escaped")).toBeUndefined(); + expect(AgentRegistry.global().get("invalid-init")).toBeUndefined(); + }); + + it("reserves orphan JSONL and lifecycle-known ids before allocating a worker name", async () => { + installPersistedSpawnMock(); + const parentManager = await createPersistedParent(); + parentManager.appendModeChange("vibe"); + const parentSessionFile = parentManager.getSessionFile(); + if (!parentSessionFile) throw new Error("Persisted parent session file was not created"); + const artifactsDir = parentSessionFile.slice(0, -6); + await fs.mkdir(artifactsDir, { recursive: true }); + await fs.writeFile(path.join(artifactsDir, "orphan.jsonl"), "orphaned transcript"); + const firstJobs = createManager(); + const firstSession = createSession({ manager: firstJobs, sessionManager: parentManager }); + const registry = VibeSessionRegistry.global(); + const orphanCollision = await registry.spawn(firstSession, { + cli: "fast", + name: "orphan", + prompt: INITIAL_VIBE_TASK, + }); + expect(orphanCollision.id).toBe("orphan-2"); + await firstJobs.getJob(orphanCollision.jobId)!.promise; + + const metadataOnly = await registry.spawn(firstSession, { + cli: "fast", + name: "metadata-only", + prompt: INITIAL_VIBE_TASK, + }); + await firstJobs.getJob(metadataOnly.jobId)!.promise; + await registry.kill(firstSession, metadataOnly.id); + await fs.rm(path.join(artifactsDir, "metadata-only.jsonl"), { force: true }); + await simulateProcessBoundary(); + + const resumedJobs = createManager(); + const resumedSession = createSession({ manager: resumedJobs, sessionManager: parentManager }); + const metadataCollision = await VibeSessionRegistry.global().spawn(resumedSession, { + cli: "fast", + name: "metadata-only", + prompt: INITIAL_VIBE_TASK, + }); + expect(metadataCollision.id).toBe("metadata-only-2"); + await resumedJobs.getJob(metadataCollision.jobId)!.promise; + }); + + it("retains a cold candidate blocked by another parent id collision so mode exit can tombstone it", async () => { + installPersistedSpawnMock(); + const parentManager = await createPersistedParent(); + parentManager.appendModeChange("vibe"); + const parentSessionFile = parentManager.getSessionFile(); + if (!parentSessionFile) throw new Error("Persisted parent session file was not created"); + const jobs = createManager(); + const session = createSession({ manager: jobs, sessionManager: parentManager }); + const registry = VibeSessionRegistry.global(); + const spawned = await registry.spawn(session, { + cli: "fast", + name: "collision", + prompt: INITIAL_VIBE_TASK, + }); + await jobs.getJob(spawned.jobId)!.promise; + await registry.suspendScope(registry.ownerScope(session), jobs); + const otherWorker = createFakeWorkerSession(); + const otherSessionFile = path.join(path.dirname(parentSessionFile), "other-parent", "collision.jsonl"); + const otherRef = AgentRegistry.global().register({ + id: "collision", + displayName: "collision", + kind: "sub", + parentId: "Main", + session: otherWorker.session, + sessionFile: otherSessionFile, + status: "idle", + }); + + expect(await registry.rehydrate(session)).toBe(1); + expect(registry.screens(session)[0]?.lastActivity).toBe("blocked by an agent id collision"); + expect(await registry.killAll(session)).toBe(1); + expect(AgentRegistry.global().get("collision")).toBe(otherRef); + expect(otherWorker.isDisposed()).toBe(false); + expect( + parentManager.getEntries().some(entry => { + if (entry.type !== "custom" || typeof entry.data !== "object" || entry.data === null) return false; + const data = entry.data as Record; + return data.id === "collision" && data.action === "tombstone"; + }), + ).toBe(true); + + AgentRegistry.global().unregister("collision", otherRef); + await parentManager.close(); + await simulateProcessBoundary(); + const reopened = await SessionManager.open(parentSessionFile, undefined, undefined, { suppressBreadcrumb: true }); + persistedManagers.push(reopened); + const resumedSession = createSession({ manager: createManager(), sessionManager: reopened }); + expect(await VibeSessionRegistry.global().rehydrate(resumedSession)).toBe(0); + expect(AgentRegistry.global().get("collision")).toMatchObject({ status: "aborted", session: null }); + }); + + it("rejects a spawn queued behind mode exit and leaves no live or untombstoned worker", async () => { + const storage = new FaultInjectingSessionStorage(); + const parentManager = await createPersistedParent(storage); + parentManager.appendModeChange("vibe"); + await parentManager.ensureOnDisk(); + const jobs = createManager(); + const session = createSession({ manager: jobs, sessionManager: parentManager }); + const registry = VibeSessionRegistry.global(); + const exitWriteStarted = Promise.withResolvers(); + const releaseExitWrite = Promise.withResolvers(); + storage.atomicWriteHooks.push(async commit => { + exitWriteStarted.resolve(); + await releaseExitWrite.promise; + await commit(); + }); + + const exiting = registry.killAll(session); + await exitWriteStarted.promise; + let spawnSettled = false; + const lateSpawn = registry + .spawn(session, { cli: "fast", name: "late-after-exit", prompt: INITIAL_VIBE_TASK }) + .then( + () => ({ error: undefined }), + error => ({ error }), + ) + .finally(() => { + spawnSettled = true; + }); + await flushMicrotasks(); + expect(spawnSettled).toBe(false); + + releaseExitWrite.resolve(); + expect(await exiting).toBe(0); + const { error } = await lateSpawn; + expect(error).toBeInstanceOf(Error); + expect(String(error)).toContain("Vibe mode has exited"); + expect(registry.listIds(session)).toEqual([]); + expect(AgentRegistry.global().get("late-after-exit")).toBeUndefined(); + expect(parentManager.buildSessionContext().mode).toBe("none"); + expect( + parentManager.getEntries().some(entry => { + if (entry.type !== "custom" || typeof entry.data !== "object" || entry.data === null) return false; + const data = entry.data as Record; + return data.id === "late-after-exit" && data.action === "spawn"; + }), + ).toBe(false); + }); + + it("keeps mode and workers live after a real atomic mode-exit failure, then retries cleanly", async () => { + installPersistedSpawnMock(); + const storage = new FaultInjectingSessionStorage(); + const parentManager = await createPersistedParent(storage); + parentManager.appendModeChange("vibe"); + const parentSessionFile = parentManager.getSessionFile(); + if (!parentSessionFile) throw new Error("Persisted parent session file was not created"); + const jobs = createManager(); + const session = createSession({ manager: jobs, sessionManager: parentManager }); + const registry = VibeSessionRegistry.global(); + const first = await registry.spawn(session, { + cli: "fast", + name: "retry-exit-one", + prompt: INITIAL_VIBE_TASK, + }); + const second = await registry.spawn(session, { + cli: "good", + name: "retry-exit-two", + prompt: INITIAL_VIBE_TASK, + }); + await Promise.all([jobs.getJob(first.jobId)!.promise, jobs.getJob(second.jobId)!.promise]); + const firstRef = AgentRegistry.global().get("retry-exit-one"); + const secondRef = AgentRegistry.global().get("retry-exit-two"); + if (!firstRef || !secondRef) throw new Error("Expected both live worker refs"); + await parentManager.flush(); + const beforeBytes = await fs.readFile(parentSessionFile); + const beforeSize = (await fs.stat(parentSessionFile)).size; + storage.atomicWriteHooks.push(async () => { + throw Object.assign(new Error("atomic publish unavailable"), { code: "ENOSPC" }); + }); + + await expect(registry.killAll(session)).rejects.toThrow("atomic publish unavailable"); + + expect((await fs.stat(parentSessionFile)).size).toBe(beforeSize); + expect(await fs.readFile(parentSessionFile)).toEqual(beforeBytes); + expect(registry.screens(session).map(screen => [screen.id, screen.state])).toEqual([ + ["retry-exit-one", "idle"], + ["retry-exit-two", "idle"], + ]); + expect(AgentRegistry.global().get("retry-exit-one")?.status).toBe("idle"); + expect(AgentRegistry.global().get("retry-exit-two")?.status).toBe("idle"); + expect(parentManager.buildSessionContext().mode).toBe("vibe"); + const failedExitActions = parentManager.getEntries().flatMap(entry => { + if (entry.type !== "custom" || typeof entry.data !== "object" || entry.data === null) return []; + const data = entry.data as Record; + return typeof data.id === "string" && data.id.startsWith("retry-exit-") && data.action === "tombstone" + ? [data.action] + : []; + }); + expect(failedExitActions).toEqual([]); + + VibeSessionRegistry.resetGlobalForTests(); + const reloadedRegistry = VibeSessionRegistry.global(); + expect(await reloadedRegistry.rehydrate(session)).toBe(2); + expect(reloadedRegistry.listIds(session)).toEqual(["retry-exit-one", "retry-exit-two"]); + expect(AgentRegistry.global().get("retry-exit-one")).toBe(firstRef); + expect(AgentRegistry.global().get("retry-exit-two")).toBe(secondRef); + + expect(await reloadedRegistry.killAll(session)).toBe(2); + expect(reloadedRegistry.listIds(session)).toEqual([]); + expect(AgentRegistry.global().get("retry-exit-one")?.status).toBe("aborted"); + expect(AgentRegistry.global().get("retry-exit-two")?.status).toBe("aborted"); + expect(parentManager.buildSessionContext().mode).toBe("none"); + + await parentManager.close(); + await simulateProcessBoundary(); + const reopened = await SessionManager.open(parentSessionFile, undefined, undefined, { suppressBreadcrumb: true }); + persistedManagers.push(reopened); + const resumedSession = createSession({ manager: createManager(), sessionManager: reopened }); + expect(reopened.buildSessionContext().mode).toBe("none"); + expect(await VibeSessionRegistry.global().rehydrate(resumedSession)).toBe(0); + expect(AgentRegistry.global().get("retry-exit-one")).toMatchObject({ status: "aborted", session: null }); + expect(AgentRegistry.global().get("retry-exit-two")).toMatchObject({ status: "aborted", session: null }); + }); + + it("fail-closes workers when mode-exit rollback durability is indeterminate", async () => { + installPersistedSpawnMock(); + const storage = new FaultInjectingSessionStorage(); + const parentManager = await createPersistedParent(storage); + parentManager.appendModeChange("vibe"); + const parentSessionFile = parentManager.getSessionFile(); + if (!parentSessionFile) throw new Error("Persisted parent session file was not created"); + const jobs = createManager(); + const session = createSession({ manager: jobs, sessionManager: parentManager }); + const registry = VibeSessionRegistry.global(); + const spawned = await registry.spawn(session, { + cli: "fast", + name: "indeterminate-exit", + prompt: INITIAL_VIBE_TASK, + }); + await jobs.getJob(spawned.jobId)!.promise; + storage.atomicWriteHooks.push( + async commit => { + await commit(); + throw new Error("mode exit committed but acknowledgement failed"); + }, + async () => { + throw new Error("authoritative rollback publish failed"); + }, + ); + + const failure = await registry.killAll(session).catch(error => error); + + expect(failure).toBeInstanceOf(SessionPersistenceIndeterminateError); + expect(registry.listIds(session)).toEqual([]); + expect(AgentRegistry.global().get("indeterminate-exit")).toMatchObject({ status: "aborted", session: null }); + expect(parentManager.buildSessionContext().mode).toBe("vibe"); + const durableText = await fs.readFile(parentSessionFile, "utf8"); + expect(durableText).toContain('"reason":"mode-exit"'); + expect(durableText).toContain('"mode":"none"'); + + await parentManager.recoverPersistenceFromCurrentState(); + expect(await registry.killAll(session)).toBe(1); + expect(parentManager.buildSessionContext().mode).toBe("none"); + }); + + it("tears down and repairs a partial writer failure while persisting an explicit tombstone", async () => { + installPersistedSpawnMock(); + const storage = new FaultInjectingSessionStorage(); + const parentManager = await createPersistedParent(storage); + parentManager.appendModeChange("vibe"); + const parentSessionFile = parentManager.getSessionFile(); + if (!parentSessionFile) throw new Error("Persisted parent session file was not created"); + const jobs = createManager(); + const session = createSession({ manager: jobs, sessionManager: parentManager }); + const registry = VibeSessionRegistry.global(); + const spawned = await registry.spawn(session, { + cli: "fast", + name: "explicit-io-failure", + prompt: INITIAL_VIBE_TASK, + }); + await jobs.getJob(spawned.jobId)!.promise; + const liveRef = AgentRegistry.global().get("explicit-io-failure"); + if (!liveRef?.session) throw new Error("Expected a live worker ref"); + const dispose = vi.spyOn(liveRef.session, "dispose"); + await parentManager.flush(); + const beforeSize = (await fs.stat(parentSessionFile)).size; + storage.failNextAppendWithPrefix(Object.assign(new Error("explicit tombstone ENOSPC"), { code: "ENOSPC" }), 23); + + await expect(registry.kill(session, "explicit-io-failure")).rejects.toThrow("explicit tombstone ENOSPC"); + + expect(dispose).toHaveBeenCalled(); + expect(storage.failedWriterClosed).toBe(true); + expect(registry.listIds(session)).toEqual([]); + expect(AgentRegistry.global().get("explicit-io-failure")).toMatchObject({ status: "aborted", session: null }); + await parentManager.flush(); + expect((await fs.stat(parentSessionFile)).size).toBeGreaterThan(beforeSize); + const repairedLines = (await fs.readFile(parentSessionFile, "utf8")).trimEnd().split("\n"); + expect(repairedLines.every(line => Boolean(JSON.parse(line)))).toBe(true); + expect( + parentManager.getEntries().some(entry => { + if (entry.type !== "custom" || typeof entry.data !== "object" || entry.data === null) return false; + const data = entry.data as Record; + return data.id === "explicit-io-failure" && data.action === "tombstone" && data.reason === "explicit-kill"; + }), + ).toBe(true); + + await parentManager.close(); + await simulateProcessBoundary(); + const reopened = await SessionManager.open(parentSessionFile, undefined, undefined, { suppressBreadcrumb: true }); + persistedManagers.push(reopened); + const resumedSession = createSession({ manager: createManager(), sessionManager: reopened }); + expect(await VibeSessionRegistry.global().rehydrate(resumedSession)).toBe(0); + expect(AgentRegistry.global().get("explicit-io-failure")).toMatchObject({ status: "aborted", session: null }); + }); + + it("serializes explicit kill ahead of a failing mode exit so rollback cannot erase it", async () => { + installPersistedSpawnMock(); + const parentManager = await createPersistedParent(); + parentManager.appendModeChange("vibe"); + const parentSessionFile = parentManager.getSessionFile(); + if (!parentSessionFile) throw new Error("Persisted parent session file was not created"); + const jobs = createManager(); + const session = createSession({ manager: jobs, sessionManager: parentManager }); + const registry = VibeSessionRegistry.global(); + const explicit = await registry.spawn(session, { + cli: "fast", + name: "overlap-explicit", + prompt: INITIAL_VIBE_TASK, + }); + const survivor = await registry.spawn(session, { + cli: "good", + name: "overlap-survivor", + prompt: INITIAL_VIBE_TASK, + }); + await Promise.all([jobs.getJob(explicit.jobId)!.promise, jobs.getJob(survivor.jobId)!.promise]); + + const originalFlush = parentManager.flush.bind(parentManager); + const explicitFlushStarted = Promise.withResolvers(); + const releaseExplicitFlush = Promise.withResolvers(); + const explicitTeardownStarted = Promise.withResolvers(); + const releaseExplicitTeardown = Promise.withResolvers(); + const lifecycle = AgentLifecycleManager.global(); + const originalRelease = lifecycle.release.bind(lifecycle); + const releaseSpy = vi.spyOn(lifecycle, "release").mockImplementation(async (id, expected) => { + if (id === "overlap-explicit") { + explicitTeardownStarted.resolve(); + await releaseExplicitTeardown.promise; + } + return originalRelease(id, expected); + }); + let flushCalls = 0; + const flush = vi.spyOn(parentManager, "flush").mockImplementation(async () => { + flushCalls++; + if (flushCalls === 1) { + explicitFlushStarted.resolve(); + await releaseExplicitFlush.promise; + await originalFlush(); + return; + } + if (flushCalls === 2) throw new Error("overlapping mode exit failed"); + await originalFlush(); + }); + + const explicitKill = registry.kill(session, "overlap-explicit"); + await explicitFlushStarted.promise; + const failedExit = registry.killAll(session).catch(error => error); + await flushMicrotasks(); + expect(flushCalls).toBe(1); + releaseExplicitFlush.resolve(); + await explicitTeardownStarted.promise; + await flushMicrotasks(); + expect(flushCalls).toBe(1); + releaseExplicitTeardown.resolve(); + expect((await explicitKill).id).toBe("overlap-explicit"); + const exitError = await failedExit; + expect(exitError).toBeInstanceOf(Error); + expect(String(exitError)).toContain("overlapping mode exit failed"); + expect(flushCalls).toBe(2); + expect(parentManager.buildSessionContext().mode).toBe("vibe"); + expect(AgentRegistry.global().get("overlap-explicit")).toMatchObject({ status: "aborted", session: null }); + expect(AgentRegistry.global().get("overlap-survivor")?.status).toBe("idle"); + + flush.mockRestore(); + releaseSpy.mockRestore(); + await parentManager.close(); + await simulateProcessBoundary(); + const reopened = await SessionManager.open(parentSessionFile, undefined, undefined, { suppressBreadcrumb: true }); + persistedManagers.push(reopened); + const resumedSession = createSession({ manager: createManager(), sessionManager: reopened }); + expect(await VibeSessionRegistry.global().rehydrate(resumedSession)).toBe(1); + expect(AgentRegistry.global().get("overlap-explicit")).toMatchObject({ status: "aborted", session: null }); + expect(AgentRegistry.global().get("overlap-survivor")).toMatchObject({ status: "parked", session: null }); + }); + + it("serializes overlapping mode exits so a later success cannot be revoked by an earlier failure", async () => { + installPersistedSpawnMock(); + const parentManager = await createPersistedParent(); + parentManager.appendModeChange("vibe"); + const parentSessionFile = parentManager.getSessionFile(); + if (!parentSessionFile) throw new Error("Persisted parent session file was not created"); + const jobs = createManager(); + const session = createSession({ manager: jobs, sessionManager: parentManager }); + const registry = VibeSessionRegistry.global(); + const spawned = await registry.spawn(session, { + cli: "fast", + name: "overlap-mode-exit", + prompt: INITIAL_VIBE_TASK, + }); + await jobs.getJob(spawned.jobId)!.promise; + + const originalFlush = parentManager.flush.bind(parentManager); + const firstExitFlushStarted = Promise.withResolvers(); + const releaseFirstExitFlush = Promise.withResolvers(); + let flushCalls = 0; + const flush = vi.spyOn(parentManager, "flush").mockImplementation(async () => { + flushCalls++; + if (flushCalls === 1) { + firstExitFlushStarted.resolve(); + await releaseFirstExitFlush.promise; + throw new Error("first mode exit failed"); + } + await originalFlush(); + }); + + const firstExit = registry.killAll(session).catch(error => error); + await firstExitFlushStarted.promise; + const successfulExit = registry.killAll(session); + await flushMicrotasks(); + expect(flushCalls).toBe(1); + releaseFirstExitFlush.resolve(); + const firstError = await firstExit; + expect(firstError).toBeInstanceOf(Error); + expect(String(firstError)).toContain("first mode exit failed"); + expect(await successfulExit).toBe(1); + expect(flushCalls).toBe(2); + expect(parentManager.buildSessionContext().mode).toBe("none"); + expect(registry.listIds(session)).toEqual([]); + expect(AgentRegistry.global().get("overlap-mode-exit")).toMatchObject({ status: "aborted", session: null }); + expect(await registry.killAll(session)).toBe(0); + expect(flushCalls).toBe(2); + + flush.mockRestore(); + await parentManager.close(); + await simulateProcessBoundary(); + const reopened = await SessionManager.open(parentSessionFile, undefined, undefined, { suppressBreadcrumb: true }); + persistedManagers.push(reopened); + const resumedSession = createSession({ manager: createManager(), sessionManager: reopened }); + expect(await VibeSessionRegistry.global().rehydrate(resumedSession)).toBe(0); + expect(AgentRegistry.global().get("overlap-mode-exit")).toMatchObject({ status: "aborted", session: null }); + }); + + it("persists mode none when re-exiting a rewound pre-exit vibe branch", async () => { + installPersistedSpawnMock(); + const parentManager = await createPersistedParent(); + parentManager.appendModeChange("vibe"); + const parentSessionFile = parentManager.getSessionFile(); + if (!parentSessionFile) throw new Error("Persisted parent session file was not created"); + const jobs = createManager(); + const session = createSession({ manager: jobs, sessionManager: parentManager }); + const registry = VibeSessionRegistry.global(); + const spawned = await registry.spawn(session, { + cli: "fast", + name: "rewound-mode-exit", + prompt: INITIAL_VIBE_TASK, + }); + await jobs.getJob(spawned.jobId)!.promise; + const preExitLeaf = parentManager.getLeafId(); + if (!preExitLeaf) throw new Error("Expected a pre-exit Vibe branch leaf"); + + expect(await registry.killAll(session)).toBe(1); + expect(parentManager.buildSessionContext().mode).toBe("none"); + parentManager.branch(preExitLeaf); + expect(parentManager.buildSessionContext().mode).toBe("vibe"); + expect(await registry.rehydrate(session)).toBe(0); + expect(registry.listIds(session)).toEqual([]); + + expect(await registry.killAll(session)).toBe(0); + expect(parentManager.buildSessionContext().mode).toBe("none"); + + await parentManager.close(); + await simulateProcessBoundary(); + const reopened = await SessionManager.open(parentSessionFile, undefined, undefined, { suppressBreadcrumb: true }); + persistedManagers.push(reopened); + expect(reopened.buildSessionContext().mode).toBe("none"); + }); + + it("does not rehydrate workers tombstoned by explicit kill or deliberate mode exit", async () => { + installPersistedSpawnMock(); + const parentManager = await createPersistedParent(); + parentManager.appendModeChange("vibe"); + const parentSessionFile = parentManager.getSessionFile(); + if (!parentSessionFile) throw new Error("Persisted parent session file was not created"); + const firstManager = createManager(); + const firstSession = createSession({ manager: firstManager, sessionManager: parentManager }); + const registry = VibeSessionRegistry.global(); + const explicitlyKilled = await registry.spawn(firstSession, { + cli: "fast", + name: "explicitly-killed", + prompt: INITIAL_VIBE_TASK, + }); + const modeExited = await registry.spawn(firstSession, { + cli: "good", + name: "mode-exited", + prompt: INITIAL_VIBE_TASK, + }); + await Promise.all([ + firstManager.getJob(explicitlyKilled.jobId)!.promise, + firstManager.getJob(modeExited.jobId)!.promise, + ]); + const preKillLeaf = parentManager.getLeafId(); + if (!preKillLeaf) throw new Error("Expected a persisted lifecycle leaf before kill"); + expect((await registry.kill(firstSession, "explicitly-killed")).cancelledTurn).toBe(false); + parentManager.appendCustomEntry("vibe-session-lifecycle", { + version: 1, + action: "tombstone-revoked", + id: "explicitly-killed", + ownerId: "Main", + parentSessionId: parentManager.getSessionId(), + reason: "mode-exit", + }); + await parentManager.flush(); + expect(await registry.killAll(firstSession)).toBe(1); + const killedSnapshot = await SessionManager.open(parentSessionFile, undefined, undefined, { + suppressBreadcrumb: true, + }); + try { + const tombstonedIds = killedSnapshot.getEntries().flatMap(entry => { + if (entry.type !== "custom" || typeof entry.data !== "object" || entry.data === null) return []; + const data = entry.data as Record; + return data.action === "tombstone" && typeof data.id === "string" ? [data.id] : []; + }); + expect(tombstonedIds.toSorted()).toEqual(["explicitly-killed", "mode-exited"]); + } finally { + await killedSnapshot.close(); + } + parentManager.branch(preKillLeaf); + + parentManager.appendModeChange("none"); + await parentManager.flush(); + await parentManager.close(); + await simulateProcessBoundary(); + const artifactsDir = parentSessionFile.slice(0, -6); + for (const id of ["explicitly-killed", "mode-exited"]) { + AgentRegistry.global().register({ + id, + displayName: id, + kind: "sub", + parentId: "Main", + session: null, + sessionFile: path.join(artifactsDir, `${id}.jsonl`), + status: "parked", + }); + } + const resumedManager = await SessionManager.open(parentSessionFile, undefined, undefined, { + suppressBreadcrumb: true, + }); + persistedManagers.push(resumedManager); + expect(resumedManager.buildSessionContext().mode).toBe("none"); + const resumedSession = createSession({ manager: createManager(), sessionManager: resumedManager }); + const resumedRegistry = VibeSessionRegistry.global(); + expect(await resumedRegistry.rehydrate(resumedSession)).toBe(0); + expect(resumedRegistry.listIds(resumedSession)).toEqual([]); + const explicitRef = AgentRegistry.global().get("explicitly-killed"); + const modeExitRef = AgentRegistry.global().get("mode-exited"); + expect(explicitRef).toMatchObject({ status: "aborted", session: null }); + expect(modeExitRef).toMatchObject({ status: "aborted", session: null }); + if (!explicitRef?.sessionFile || !modeExitRef?.sessionFile) { + throw new Error("Tombstoned workers must retain readable transcript references"); + } + expect((await SessionManager.peekSessionInit(explicitRef.sessionFile))?.init?.task).toBe(INITIAL_VIBE_TASK); + expect((await SessionManager.peekSessionInit(modeExitRef.sessionFile))?.init?.task).toBe(INITIAL_VIBE_TASK); + await expect(AgentLifecycleManager.global().ensureLive("explicitly-killed")).rejects.toThrow( + "aborted and cannot be revived", + ); + await expect( + resumedRegistry.send(resumedSession, { session: "explicitly-killed", message: FOLLOW_UP_VIBE_TASK }), + ).rejects.toThrow('Unknown vibe session "explicitly-killed"'); + await expect( + resumedRegistry.send(resumedSession, { session: "mode-exited", message: FOLLOW_UP_VIBE_TASK }), + ).rejects.toThrow('Unknown vibe session "mode-exited"'); }); it("runSubagentFollowUpTurn continues the same live session and finalizes trace + yield response", async () => { @@ -487,15 +1857,185 @@ describe("vibe session registry", () => { const { jobId } = await registry.spawn(session, { cli: "fast", name: "Doomed", prompt: "Never mind." }); await pollUntil(() => AgentRegistry.global().get("Doomed") !== undefined); - const outcome = await registry.kill(session, "Doomed"); + const killPromise = registry.kill(session, "Doomed"); + await pollUntil(() => manager.getJob(jobId)?.status === "cancelled"); + gate.resolve(); + const outcome = await killPromise; expect(outcome.cancelledTurn).toBe(true); expect(manager.getJob(jobId)!.status).toBe("cancelled"); expect(fake.isDisposed()).toBe(true); expect(AgentRegistry.global().get("Doomed")).toBeUndefined(); - expect(registry.screens("Main")[0]?.state).toBe("dead"); + expect(registry.screens(session)[0]?.state).toBe("dead"); await expect(registry.send(session, { session: "Doomed", message: "hello?" })).rejects.toThrow("dead"); + }); - gate.resolve(); + it("keeps a persisted in-flight kill terminal when the old executor finalizes late", async () => { + let worker: ReturnType | undefined; + vi.spyOn(executorModule, "runSubprocess").mockImplementation(async options => { + const artifactsDir = options.artifactsDir; + if (!artifactsDir) throw new Error("Persisted vibe test requires an artifacts directory"); + const childSessionFile = await persistWorkerSession({ + cwd: options.cwd, + artifactsDir, + id: options.id, + task: options.task, + }); + worker = createFakeWorkerSession(); + AgentRegistry.global().register({ + id: options.id, + displayName: options.id, + kind: "sub", + parentId: options.parentAgentId ?? "Main", + session: worker.session, + sessionFile: childSessionFile, + status: "running", + }); + const signal = options.signal; + if (!signal) throw new Error("Persisted blocked worker requires a cancellation signal"); + await new Promise(resolve => { + if (signal.aborted) resolve(); + else signal.addEventListener("abort", () => resolve(), { once: true }); + }); + return makeResult(options.id, { output: "Killed during work.", aborted: true }); + }); + const parentManager = await createPersistedParent(); + parentManager.appendModeChange("vibe"); + const jobs = createManager(); + const session = createSession({ manager: jobs, sessionManager: parentManager }); + const registry = VibeSessionRegistry.global(); + await registry.spawn(session, { cli: "fast", name: "persisted-kill", prompt: "Keep working." }); + await pollUntil(() => worker !== undefined); + + expect((await registry.kill(session, "persisted-kill")).cancelledTurn).toBe(true); + const terminal = AgentRegistry.global().get("persisted-kill"); + expect(terminal).toMatchObject({ status: "aborted", session: null }); + + await executorModule.finalizeSubagentLifecycle({ + id: "persisted-kill", + session: worker!.session, + aborted: true, + keepAlive: true, + isolated: false, + agentIdleTtlMs: 0, + reviveSession: null, + }); + expect(AgentRegistry.global().get("persisted-kill")).toBe(terminal); + expect(AgentRegistry.global().get("persisted-kill")).toMatchObject({ status: "aborted", session: null }); + await expect(AgentLifecycleManager.global().ensureLive("persisted-kill")).rejects.toThrow("cannot be revived"); + }); + + it("does not terminalize a same-path replacement installed while the killed worker disposes", async () => { + let worker: ReturnType | undefined; + let replacement: ReturnType | undefined; + let replacementWorker: ReturnType | undefined; + vi.spyOn(executorModule, "runSubprocess").mockImplementation(async options => { + const artifactsDir = options.artifactsDir; + if (!artifactsDir) throw new Error("Persisted vibe test requires an artifacts directory"); + const childSessionFile = await persistWorkerSession({ + cwd: options.cwd, + artifactsDir, + id: options.id, + task: options.task, + }); + replacementWorker = createFakeWorkerSession(); + worker = createFakeWorkerSession({ + onDispose: () => { + replacement = AgentRegistry.global().register({ + id: options.id, + displayName: options.id, + kind: "sub", + parentId: options.parentAgentId ?? "Main", + session: replacementWorker!.session, + sessionFile: childSessionFile, + status: "idle", + }); + }, + }); + AgentRegistry.global().register({ + id: options.id, + displayName: options.id, + kind: "sub", + parentId: options.parentAgentId ?? "Main", + session: worker.session, + sessionFile: childSessionFile, + status: "running", + }); + const signal = options.signal; + if (!signal) throw new Error("Persisted blocked worker requires a cancellation signal"); + await new Promise(resolve => { + if (signal.aborted) resolve(); + else signal.addEventListener("abort", () => resolve(), { once: true }); + }); + return makeResult(options.id, { output: "Old worker killed.", aborted: true }); + }); + const parentManager = await createPersistedParent(); + parentManager.appendModeChange("vibe"); + const jobs = createManager(); + const session = createSession({ manager: jobs, sessionManager: parentManager }); + const registry = VibeSessionRegistry.global(); + await registry.spawn(session, { cli: "fast", name: "same-path-replacement", prompt: "Keep working." }); + await pollUntil(() => worker !== undefined); + + expect((await registry.kill(session, "same-path-replacement")).cancelledTurn).toBe(true); + expect(replacement).toBeDefined(); + expect(AgentRegistry.global().get("same-path-replacement")).toBe(replacement); + expect(replacement).toMatchObject({ status: "idle", session: replacementWorker!.session }); + expect(replacementWorker!.isDisposed()).toBe(false); + }); + + it("persists a kill issued before child initialization and terminalizes the worker if it registers late", async () => { + vi.spyOn(executorModule, "runSubprocess").mockImplementation(async options => { + const artifactsDir = options.artifactsDir; + if (!artifactsDir) throw new Error("Persisted vibe test requires an artifacts directory"); + const signal = options.signal; + if (!signal) throw new Error("Pre-initialization worker requires a cancellation signal"); + await new Promise(resolve => { + if (signal.aborted) resolve(); + else signal.addEventListener("abort", () => resolve(), { once: true }); + }); + const childSessionFile = await persistWorkerSession({ + cwd: options.cwd, + artifactsDir, + id: options.id, + task: options.task, + }); + AgentRegistry.global().register({ + id: options.id, + displayName: options.id, + kind: "sub", + parentId: options.parentAgentId ?? "Main", + session: createFakeWorkerSession().session, + sessionFile: childSessionFile, + status: "idle", + }); + return makeResult(options.id, { output: "Killed before initialization.", aborted: true }); + }); + const parentManager = await createPersistedParent(); + parentManager.appendModeChange("vibe"); + const parentSessionFile = parentManager.getSessionFile(); + if (!parentSessionFile) throw new Error("Persisted parent session file was not created"); + const jobs = createManager(); + const session = createSession({ manager: jobs, sessionManager: parentManager }); + const registry = VibeSessionRegistry.global(); + await registry.spawn(session, { cli: "fast", name: "pre-init-kill", prompt: "Start later." }); + + expect((await registry.kill(session, "pre-init-kill")).cancelledTurn).toBe(true); + expect(AgentRegistry.global().get("pre-init-kill")).toMatchObject({ status: "aborted", session: null }); + expect( + parentManager.getEntries().some(entry => { + if (entry.type !== "custom" || typeof entry.data !== "object" || entry.data === null) return false; + const data = entry.data as Record; + return data.id === "pre-init-kill" && data.action === "tombstone"; + }), + ).toBe(true); + + await parentManager.close(); + await simulateProcessBoundary(); + const reopened = await SessionManager.open(parentSessionFile, undefined, undefined, { suppressBreadcrumb: true }); + persistedManagers.push(reopened); + const resumedSession = createSession({ manager: createManager(), sessionManager: reopened }); + expect(await VibeSessionRegistry.global().rehydrate(resumedSession)).toBe(0); + expect(AgentRegistry.global().get("pre-init-kill")).toMatchObject({ status: "aborted", session: null }); }); it("killAll terminates every session for the owner (mode-exit path)", async () => { @@ -518,16 +2058,19 @@ describe("vibe session registry", () => { const manager = createManager(); const session = createSession({ manager }); const registry = VibeSessionRegistry.global(); - await registry.spawn(session, { cli: "fast", name: "One", prompt: "A." }); - await registry.spawn(session, { cli: "good", name: "Two", prompt: "B." }); + const one = await registry.spawn(session, { cli: "fast", name: "One", prompt: "A." }); + const two = await registry.spawn(session, { cli: "good", name: "Two", prompt: "B." }); await pollUntil(() => gates.size === 2); - const killed = await registry.killAll("Main", manager); + const killPromise = registry.killAll(session); + await pollUntil(() => manager.getJob(one.jobId)?.status === "cancelled"); + gates.get("One")!.resolve(); + await pollUntil(() => manager.getJob(two.jobId)?.status === "cancelled"); + gates.get("Two")!.resolve(); + const killed = await killPromise; expect(killed).toBe(2); - expect(registry.listIds("Main")).toEqual([]); + expect(registry.listIds(session)).toEqual([]); expect(AgentRegistry.global().get("One")).toBeUndefined(); expect(AgentRegistry.global().get("Two")).toBeUndefined(); - - for (const gate of gates.values()) gate.resolve(); }); }); From febf3eb73dc49b6672b7b2c3ae43f7c8c77be9df Mon Sep 17 00:00:00 2001 From: roboomp Date: Tue, 14 Jul 2026 23:42:18 +0000 Subject: [PATCH 2/4] fix(coding-agent): bounded cancelled vibe teardown vibe_kill and Vibe-mode exit / session-switch suspension awaited every cancelled turn's job promise unconditionally. When a provider or tool ignores the abort signal the promise never settles, so teardown hung indefinitely. Give cancelled jobs a 250ms unref'd settlement grace, then log and detach any still-pending job while preserving cancelled state. Ported from @RensTillmann's fix (65446c4). Fixes #5303 --- packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/src/vibe/runtime.ts | 24 +++- .../test/vibe/vibe-runtime.test.ts | 111 +++++++++++++++--- 3 files changed, 120 insertions(+), 16 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 4676cc148..4257621c2 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -5,6 +5,7 @@ ### Fixed - Fixed persisted vibe workers disappearing or being replaced across graceful restarts, session switches, failed mode exits, and late cancelled initialization: resumable conversations now restore safely, mode exit atomically commits worker tombstones with the mode change and rolls back cleanly on storage failure, explicit kills tear workers down monotonically while repairing uncertain append tails, killed transcripts remain readable but non-revivable, and stale initializers cannot overwrite newer same-ID workers ([#5303](https://github.com/can1357/oh-my-pi/issues/5303) by [@mastertyko](https://github.com/mastertyko)). +- Bounded vibe teardown so `vibe_kill` and Vibe-mode exit / session-switch suspension no longer hang when a cancelled turn's provider or tool ignores the abort signal: cancelled jobs get a short unref'd settlement grace, then detach while staying marked cancelled ([#5303](https://github.com/can1357/oh-my-pi/issues/5303) by [@RensTillmann](https://github.com/RensTillmann)). ## [16.4.8] - 2026-07-12 diff --git a/packages/coding-agent/src/vibe/runtime.ts b/packages/coding-agent/src/vibe/runtime.ts index a040d5b95..56214dfa7 100644 --- a/packages/coding-agent/src/vibe/runtime.ts +++ b/packages/coding-agent/src/vibe/runtime.ts @@ -69,6 +69,8 @@ const TRACE_LINE_MAX = 120; const DEFAULT_WAIT_TIMEOUT_MS = 30_000; /** Response text cap inside a delivered turn result; full output stays at agent://. */ const RESPONSE_PREVIEW_MAX = 6000; +/** Grace period for abort-aware turns before teardown detaches a stuck provider/tool call. */ +const CANCELLED_TURN_SETTLE_GRACE_MS = 250; const VIBE_LIFECYCLE_CUSTOM_TYPE = "vibe-session-lifecycle"; const VIBE_LIFECYCLE_VERSION = 1; @@ -351,6 +353,24 @@ function mergeTrace(turn: VibeTurn, progress: AgentProgress): void { /** Thrown from a turn job body so the job manager marks the job failed while carrying the formatted result. */ export class VibeTurnError extends Error {} +async function awaitCancelledTurnJobs(jobs: ReadonlySet): Promise { + if (jobs.size === 0) return; + const settled = Promise.allSettled([...jobs].map(job => job.promise)).then(() => true); + const timeout = Promise.withResolvers(); + const timer = setTimeout(() => timeout.resolve(false), CANCELLED_TURN_SETTLE_GRACE_MS); + timer.unref(); + try { + if (!(await Promise.race([settled, timeout.promise]))) { + logger.warn("vibe: detached cancelled turn that did not settle within teardown grace period", { + jobCount: jobs.size, + graceMs: CANCELLED_TURN_SETTLE_GRACE_MS, + }); + } + } finally { + clearTimeout(timer); + } +} + /** * Process-global registry of vibe worker sessions, scoped by both owner agent * id and stable parent session id. Persisted lifecycle events rebuild idle @@ -1043,7 +1063,7 @@ export class VibeSessionRegistry { }); } } - await Promise.allSettled(teardown.flatMap(entry => (entry.job ? [entry.job.promise] : []))); + await awaitCancelledTurnJobs(new Set(teardown.flatMap(entry => (entry.job ? [entry.job] : [])))); return records.length; } @@ -1154,7 +1174,7 @@ export class VibeSessionRegistry { }); } } - await Promise.allSettled([...settlingJobs].map(job => job.promise)); + await awaitCancelledTurnJobs(settlingJobs); const terminalRef = registered ?? this.#registeredAgent(record) ?? null; if (record.childSessionFile) { try { diff --git a/packages/coding-agent/test/vibe/vibe-runtime.test.ts b/packages/coding-agent/test/vibe/vibe-runtime.test.ts index 3eeb044e2..e745acb6b 100644 --- a/packages/coding-agent/test/vibe/vibe-runtime.test.ts +++ b/packages/coding-agent/test/vibe/vibe-runtime.test.ts @@ -803,6 +803,61 @@ describe("vibe session registry", () => { expect(turnJob.resultText).toContain('turn="2"'); }); + it("bounds parent-session suspension when a cancelled turn ignores abort and settles late", async () => { + const gate = deferred(); + const started = deferred(); + const disposed = deferred(); + const fake = createFakeWorkerSession({ onDispose: disposed.resolve }); + vi.spyOn(executorModule, "runSubprocess").mockImplementation(async options => { + AgentRegistry.global().register({ + id: options.id, + displayName: options.id, + kind: "sub", + parentId: "Main", + session: fake.session, + status: "running", + }); + started.resolve(); + await gate.promise; + return makeResult(options.id); + }); + + const manager = createManager(); + const session = createSession({ manager }); + const registry = VibeSessionRegistry.global(); + const { jobId } = await registry.spawn(session, { + cli: "fast", + name: "IgnoresSuspendAbort", + prompt: "Keep working through a parent-session switch.", + }); + await started.promise; + + vi.useFakeTimers(); + try { + const suspension = registry.suspendScope(registry.ownerScope(session), manager); + await disposed.promise; + await flushMicrotasks(); + expect(vi.getTimerCount()).toBeGreaterThan(0); + vi.advanceTimersByTime(250); + + expect(await suspension).toBe(1); + expect(manager.getJob(jobId)!.status).toBe("cancelled"); + expect(fake.isDisposed()).toBe(true); + expect(AgentRegistry.global().get("IgnoresSuspendAbort")).toBeUndefined(); + expect(registry.listIds(session)).toEqual([]); + + gate.resolve(); + await manager.getJob(jobId)!.promise; + expect(manager.getJob(jobId)!.status).toBe("cancelled"); + expect(fake.isDisposed()).toBe(true); + expect(AgentRegistry.global().get("IgnoresSuspendAbort")).toBeUndefined(); + expect(registry.listIds(session)).toEqual([]); + } finally { + gate.resolve(); + vi.useRealTimers(); + } + }); + it("restores an interrupted turn as idle without replay and continues only after send", async () => { installPersistedSpawnMock(); const parentManager = await createPersistedParent(); @@ -1835,9 +1890,11 @@ describe("vibe session registry", () => { await manager.getJob("Fast-t2")!.promise; }); - it("kill cancels the in-flight turn and releases the worker session", async () => { + it("bounds kill teardown when a cancelled turn ignores abort and settles late", async () => { const gate = deferred(); - const fake = createFakeWorkerSession(); + const started = deferred(); + const disposed = deferred(); + const fake = createFakeWorkerSession({ onDispose: disposed.resolve }); vi.spyOn(executorModule, "runSubprocess").mockImplementation(async options => { AgentRegistry.global().register({ id: options.id, @@ -1847,6 +1904,7 @@ describe("vibe session registry", () => { session: fake.session, status: "running", }); + started.resolve(); await gate.promise; return makeResult(options.id); }); @@ -1854,19 +1912,44 @@ describe("vibe session registry", () => { const manager = createManager(); const session = createSession({ manager }); const registry = VibeSessionRegistry.global(); - const { jobId } = await registry.spawn(session, { cli: "fast", name: "Doomed", prompt: "Never mind." }); - await pollUntil(() => AgentRegistry.global().get("Doomed") !== undefined); + const { jobId } = await registry.spawn(session, { + cli: "fast", + name: "IgnoresKillAbort", + prompt: "Keep working through explicit termination.", + }); + await started.promise; - const killPromise = registry.kill(session, "Doomed"); - await pollUntil(() => manager.getJob(jobId)?.status === "cancelled"); - gate.resolve(); - const outcome = await killPromise; - expect(outcome.cancelledTurn).toBe(true); - expect(manager.getJob(jobId)!.status).toBe("cancelled"); - expect(fake.isDisposed()).toBe(true); - expect(AgentRegistry.global().get("Doomed")).toBeUndefined(); - expect(registry.screens(session)[0]?.state).toBe("dead"); - await expect(registry.send(session, { session: "Doomed", message: "hello?" })).rejects.toThrow("dead"); + vi.useFakeTimers(); + try { + const kill = registry.kill(session, "IgnoresKillAbort"); + await disposed.promise; + await flushMicrotasks(); + expect(vi.getTimerCount()).toBeGreaterThan(0); + vi.advanceTimersByTime(250); + + const outcome = await kill; + expect(outcome.cancelledTurn).toBe(true); + expect(manager.getJob(jobId)!.status).toBe("cancelled"); + expect(fake.isDisposed()).toBe(true); + expect(AgentRegistry.global().get("IgnoresKillAbort")).toBeUndefined(); + expect(registry.screens(session)[0]?.state).toBe("dead"); + await expect(registry.send(session, { session: "IgnoresKillAbort", message: "hello?" })).rejects.toThrow( + "dead", + ); + + gate.resolve(); + await manager.getJob(jobId)!.promise; + expect(manager.getJob(jobId)!.status).toBe("cancelled"); + expect(fake.isDisposed()).toBe(true); + expect(AgentRegistry.global().get("IgnoresKillAbort")).toBeUndefined(); + expect(registry.screens(session)[0]?.state).toBe("dead"); + await expect(registry.send(session, { session: "IgnoresKillAbort", message: "still there?" })).rejects.toThrow( + "dead", + ); + } finally { + gate.resolve(); + vi.useRealTimers(); + } }); it("keeps a persisted in-flight kill terminal when the old executor finalizes late", async () => { From 709066b7a1e068fbeb0367badd2d48baf93b639e Mon Sep 17 00:00:00 2001 From: roboomp Date: Tue, 14 Jul 2026 23:58:30 +0000 Subject: [PATCH 3/4] fix(coding-agent): preserve target tools when switching out of vibe MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #clearTransientModeState's vibe branch runs only from #reconcileModeFromSession, after switchSession has already loaded and restored the target session's active tools. It called deactivateVibeTools(#vibeModePreviousTools), which SETS the active tool set to that snapshot — but the snapshot belongs to the SOURCE session, so switching to a non-vibe (or different vibe) target clobbered the target's tools. Add AgentSession.removeVibeToolsPreservingActive(), which strips only the transient vibe tools and keeps the current active set, and use it on the switch path. The live-exit path (#exitVibeMode) still restores its own snapshot. Fixes #5303 --- packages/coding-agent/CHANGELOG.md | 1 + .../src/modes/interactive-mode.ts | 7 ++++- .../coding-agent/src/session/agent-session.ts | 20 +++++++++++++- .../test/interactive-mode-vibe-toggle.test.ts | 27 +++++++++++++++++++ 4 files changed, 53 insertions(+), 2 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 4257621c2..21f7879b0 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -6,6 +6,7 @@ - Fixed persisted vibe workers disappearing or being replaced across graceful restarts, session switches, failed mode exits, and late cancelled initialization: resumable conversations now restore safely, mode exit atomically commits worker tombstones with the mode change and rolls back cleanly on storage failure, explicit kills tear workers down monotonically while repairing uncertain append tails, killed transcripts remain readable but non-revivable, and stale initializers cannot overwrite newer same-ID workers ([#5303](https://github.com/can1357/oh-my-pi/issues/5303) by [@mastertyko](https://github.com/mastertyko)). - Bounded vibe teardown so `vibe_kill` and Vibe-mode exit / session-switch suspension no longer hang when a cancelled turn's provider or tool ignores the abort signal: cancelled jobs get a short unref'd settlement grace, then detach while staying marked cancelled ([#5303](https://github.com/can1357/oh-my-pi/issues/5303) by [@RensTillmann](https://github.com/RensTillmann)). +- Fixed switching out of a Vibe-mode session clobbering the target session's active tools: mode reconciliation now strips only the transient vibe tools and preserves the freshly loaded target's tool set, instead of re-applying the source session's stale pre-vibe snapshot ([#5303](https://github.com/can1357/oh-my-pi/issues/5303)). ## [16.4.8] - 2026-07-12 diff --git a/packages/coding-agent/src/modes/interactive-mode.ts b/packages/coding-agent/src/modes/interactive-mode.ts index ab09d0ca6..ebba1e4b7 100644 --- a/packages/coding-agent/src/modes/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive-mode.ts @@ -2150,7 +2150,12 @@ export class InteractiveMode implements InteractiveModeContext { if (this.vibeModeEnabled && !options?.preserveVibe) { const ownerScope = this.#vibeModeOwnerScope; - await this.session.deactivateVibeTools(this.#vibeModePreviousTools ?? []); + // This runs only from #reconcileModeFromSession, i.e. after switchSession + // already loaded and restored the target session's active tools. The + // #vibeModePreviousTools snapshot belongs to the SOURCE session, so + // applying it here would clobber the target's tools — strip only the + // transient vibe tools and keep the target's active set intact. + await this.session.removeVibeToolsPreservingActive(); this.session.setVibeModeState(undefined); this.vibeModeEnabled = false; this.#vibeModePreviousTools = undefined; diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index cc0cc4a8f..f6e71fb5c 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -6191,13 +6191,31 @@ export class AgentSession { /** Removes tools installed by {@link activateVibeTools} and activates `nextToolNames`. */ async deactivateVibeTools(nextToolNames: string[]): Promise { + this.#uninstallVibeTools(); + await this.#applyActiveToolsByName(nextToolNames); + } + + /** + * Removes the ephemeral vibe tools while keeping whatever active tool set the + * session currently holds (minus those vibe tools). Unlike + * {@link deactivateVibeTools}, this never restores a caller-held pre-vibe + * snapshot — use it on the session-switch path, where that snapshot belongs to + * the source session and would clobber the freshly loaded target's tools. + */ + async removeVibeToolsPreservingActive(): Promise { + const removed = new Set(this.#installedVibeToolNames); + this.#uninstallVibeTools(); + const nextActive = this.getActiveToolNames().filter(name => !removed.has(name)); + await this.#applyActiveToolsByName(nextActive); + } + + #uninstallVibeTools(): void { for (const name of this.#installedVibeToolNames) { this.#toolRegistry.delete(name); this.#builtInToolNames.delete(name); this.#selectedDiscoveredToolNames.delete(name); } this.#installedVibeToolNames.clear(); - await this.#applyActiveToolsByName(nextToolNames); } #getEditModeSession() { diff --git a/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts b/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts index 1a1b20ac9..eddb22822 100644 --- a/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts +++ b/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts @@ -202,6 +202,33 @@ describe("InteractiveMode vibe mode toggle", () => { expect(vibeModeEntryCount(session.sessionManager)).toBe(1); }); + it("does not clobber the target's active tools with the source snapshot when switching out of vibe", async () => { + await mode.init({ suppressWelcomeIntro: true }); + // Pre-vibe snapshot on the source session is empty; entering vibe activates + // read + the vibe tools. + await mode.handleVibeModeCommand(); + expect(mode.vibeModeEnabled).toBe(true); + expect(session.getActiveToolNames()).toContain("read"); + + // Target is a distinct, non-vibe session. + const targetManager = SessionManager.create(tempDir.path(), tempDir.path()); + targetManager.appendModeChange("none"); + await targetManager.ensureOnDisk(); + const targetFile = targetManager.getSessionFile(); + if (!targetFile) throw new Error("Expected target session file"); + await targetManager.close(); + + expect(await session.switchSession(targetFile)).toBe(true); + + expect(mode.vibeModeEnabled).toBe(false); + // The transient vibe tools are gone, but the genuinely-active `read` tool + // must survive — the source's empty pre-vibe snapshot must not wipe it. + expect(session.getActiveToolNames()).toEqual(["read"]); + for (const name of VIBE_TOOL_NAMES) { + expect(session.getActiveToolNames()).not.toContain(name); + } + }); + it("rejects new, drop, fork, and move transitions at the AgentSession boundary while vibe is active", async () => { await mode.init({ suppressWelcomeIntro: true }); await mode.handleVibeModeCommand(); From 122ab2a4efc851ce816cc2edbe30e6d314b50deb Mon Sep 17 00:00:00 2001 From: roboomp Date: Wed, 15 Jul 2026 00:19:36 +0000 Subject: [PATCH 4/4] fix(coding-agent): resolve restored vibe workers against session model #vibeParentSession() omitted getActiveModelString, so on resume/switch #reconcileModeFromSession -> VibeSessionRegistry.rehydrate() -> #resolveWorker() fell back to the settings default model. A reopened or switched-to session whose active model came from the journal or a non-default launch therefore resolved the good (pi/task) worker on the wrong model. Wire getActiveModelString from the session's active model, matching the spawn-path ToolSession and sdk's dispose-path helper. Fixes #5303 --- packages/coding-agent/CHANGELOG.md | 1 + .../src/modes/interactive-mode.ts | 6 +++++ .../test/interactive-mode-vibe-toggle.test.ts | 25 +++++++++++++++++++ 3 files changed, 32 insertions(+) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 21f7879b0..3cf9422c2 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -7,6 +7,7 @@ - Fixed persisted vibe workers disappearing or being replaced across graceful restarts, session switches, failed mode exits, and late cancelled initialization: resumable conversations now restore safely, mode exit atomically commits worker tombstones with the mode change and rolls back cleanly on storage failure, explicit kills tear workers down monotonically while repairing uncertain append tails, killed transcripts remain readable but non-revivable, and stale initializers cannot overwrite newer same-ID workers ([#5303](https://github.com/can1357/oh-my-pi/issues/5303) by [@mastertyko](https://github.com/mastertyko)). - Bounded vibe teardown so `vibe_kill` and Vibe-mode exit / session-switch suspension no longer hang when a cancelled turn's provider or tool ignores the abort signal: cancelled jobs get a short unref'd settlement grace, then detach while staying marked cancelled ([#5303](https://github.com/can1357/oh-my-pi/issues/5303) by [@RensTillmann](https://github.com/RensTillmann)). - Fixed switching out of a Vibe-mode session clobbering the target session's active tools: mode reconciliation now strips only the transient vibe tools and preserves the freshly loaded target's tool set, instead of re-applying the source session's stale pre-vibe snapshot ([#5303](https://github.com/can1357/oh-my-pi/issues/5303)). +- Fixed restored Vibe workers resolving against the settings default model instead of the reopened session's active model: `#reconcileModeFromSession` now passes the session's active model string into rehydration so the `good` (`pi/task`) worker's inherited model tracks the resumed or switched-to session ([#5303](https://github.com/can1357/oh-my-pi/issues/5303)). ## [16.4.8] - 2026-07-12 diff --git a/packages/coding-agent/src/modes/interactive-mode.ts b/packages/coding-agent/src/modes/interactive-mode.ts index ebba1e4b7..ef1dde6d3 100644 --- a/packages/coding-agent/src/modes/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive-mode.ts @@ -56,6 +56,7 @@ import { reset as resetCapabilities } from "../capability"; import type { CollabGuestLink } from "../collab/guest"; import type { CollabHost } from "../collab/host"; import { KeybindingsManager } from "../config/keybindings"; +import { formatModelString } from "../config/model-resolver"; import { applyProviderGlobalsFromSettings } from "../config/provider-globals"; import { isSettingsInitialized, onStatusLineSessionAccentChanged, Settings, settings } from "../config/settings"; import { clearClaudePluginRootsCache } from "../discovery/helpers"; @@ -1966,6 +1967,11 @@ export class InteractiveMode implements InteractiveModeContext { sessionManager: this.sessionManager, asyncJobManager: this.session.asyncJobManager, settings: this.session.settings, + // Resolve restored/switched-to workers against this session's active model + // (same as the spawn-path ToolSession), not the settings default. This is + // the primary fallback in resolveAgentModelPatterns, so the `good` worker's + // pi/task inheritance tracks the reopened session's model. + getActiveModelString: () => (this.session.model ? formatModelString(this.session.model) : undefined), }; } diff --git a/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts b/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts index eddb22822..f096b1ee0 100644 --- a/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts +++ b/packages/coding-agent/test/interactive-mode-vibe-toggle.test.ts @@ -178,6 +178,31 @@ describe("InteractiveMode vibe mode toggle", () => { expect(vibeModeEntryCount(session.sessionManager)).toBe(1); }); + it("passes the session's active model into vibe rehydration on resume", async () => { + await mode.init({ suppressWelcomeIntro: true }); + await mode.handleVibeModeCommand(); + await session.sessionManager.ensureOnDisk(); + const sessionFile = session.sessionFile; + if (!sessionFile) throw new Error("Expected persisted session file"); + const expectedModel = session.model; + if (!expectedModel) throw new Error("Expected an active session model"); + const registry = VibeSessionRegistry.global(); + let rehydrateCalled = false; + let activeModelDuringRehydrate: string | undefined; + vi.spyOn(registry, "rehydrate").mockImplementation(async parent => { + rehydrateCalled = true; + activeModelDuringRehydrate = parent.getActiveModelString?.(); + return 0; + }); + + expect(await session.switchSession(sessionFile)).toBe(true); + + // Rehydration must resolve workers against the reopened session's active + // model (so the `good`/pi/task worker tracks it), not the settings default. + expect(rehydrateCalled).toBe(true); + expect(activeModelDuringRehydrate).toBe(`${expectedModel.provider}/${expectedModel.id}`); + }); + it("suspends the old scope without tombstones when switching to another vibe parent", async () => { await mode.init({ suppressWelcomeIntro: true }); await mode.handleVibeModeCommand();