@ian-pascoe/pi-minimal-subagents 0.1.0 → 0.2.0

package/src/minimal-subagents-extension.ts

@@ -5,14 +5,25 @@ import {
   getAgentDir,
   SessionManager,
   SettingsManager,
+  type AgentSettledEvent,
   type ExtensionAPI,
   type ExtensionContext,
+  type ExtensionFactory,
+  type MessageEndEvent,
+  type SessionBeforeForkEvent,
   type SessionEntry,
+  type SessionShutdownEvent,
+  type SessionStartEvent,
+  type SessionTreeEvent,
 } from "@earendil-works/pi-coding-agent";
 import { MinimalSubagentsCoordinator } from "./minimal-subagents-coordinator.js";
 import { resolveMinimalSubagentsSettings } from "./minimal-subagents-config.js";
 import { snapshotCommittedContext } from "./minimal-subagents-context.js";
-import { rememberForkSnapshot, takeForkSnapshot } from "./minimal-subagents-fork-lifecycle.js";
+import {
+  isForkDestinationForSource,
+  rememberForkSnapshot,
+  takeForkSnapshot,
+} from "./minimal-subagents-fork-lifecycle.js";
 import {
   buildEligibleModelIds,
   COORDINATOR_TOOL_NAMES,
@@ -22,9 +33,14 @@ import {
   CHILD_IDENTITY_ENTRY_TYPE,
   REGISTRY_ENTRY_TYPE,
   replayRegistryEntries,
+  type RegistryReplayDiagnostic,
 } from "./minimal-subagents-registry.js";
 import { createCoordinatorToolSchemas } from "./minimal-subagents-tool-schemas.js";
-import { findDeliveryEvidence, PiAgentSessionFactory } from "./minimal-subagents-sessions.js";
+import {
+  findDeliveryEvidence,
+  PiAgentSessionFactory,
+  type PiAgentSessionFactoryOptions,
+} from "./minimal-subagents-sessions.js";
 import { shutdownMinimalSubagentsSession } from "./minimal-subagents-shutdown.js";
 import { createCoordinatorToolDefinitions } from "./minimal-subagents-tools.js";
 import {
@@ -33,9 +49,11 @@ import {
 } from "./minimal-subagents-rendering.js";
 import { MinimalSubagentsUiController } from "./minimal-subagents-ui.js";
 import type {
+  AgentSessionFactory,
   CallerSnapshot,
   CoordinatorNotification,
   ForkSnapshot,
+  PersistedAgent,
   RegistrySnapshot,
   RootConversationEndpoint,
 } from "./minimal-subagents-types.js";
@@ -43,7 +61,7 @@ import type {
 const EXTENSION_ENTRYPOINT = fileURLToPath(new URL("./index.ts", import.meta.url));
 
 function currentConversationMessages(context: ExtensionContext): AgentMessage[] {
-  const entries = context.sessionManager.getEntries() as SessionEntry[];
+  const entries = context.sessionManager.getEntries();
   const messages = buildSessionContext(entries, context.sessionManager.getLeafId()).messages;
   return snapshotCommittedContext(messages, !context.isIdle());
 }
@@ -68,7 +86,7 @@ function createRootConversationEndpoint(
   context: ExtensionContext,
 ): RootConversationEndpoint {
   return {
-    async steerCoordinatorMessage(message) {
+    async queueCoordinatorMessage(message) {
       pi.sendMessage(
         {
           customType: message.customType,
@@ -82,8 +100,14 @@ function createRootConversationEndpoint(
         },
       );
     },
-    hasDeliveryEvidence: (sourceAgentId, sourceTurnId) =>
-      findDeliveryEvidence(context.sessionManager.getEntries(), sourceAgentId, sourceTurnId),
+    hasDeliveryEvidence: (sourceAgentId, sourceTurnId, deliveryId) =>
+      findDeliveryEvidence(
+        context.sessionManager.getBranch(),
+        sourceAgentId,
+        sourceTurnId,
+        deliveryId,
+      ),
+    isIdle: () => context.isIdle(),
   };
 }
 
@@ -111,10 +135,158 @@ function hasHistoricalChildIdentity(entries: readonly SessionEntry[]): boolean {
   );
 }
 
-function replayPreviousRoot(previousSessionFile: string): RegistrySnapshot {
-  const previousSession = SessionManager.open(previousSessionFile);
-  const previousRootSessionId = previousSession.getSessionId();
-  return replayRegistryEntries(previousSession.getEntries(), previousRootSessionId);
+function reportInvalidRegistryRecords(
+  context: ExtensionContext,
+): (diagnostics: RegistryReplayDiagnostic[]) => void {
+  return (diagnostics) => {
+    const summaries = diagnostics
+      .slice(0, 3)
+      .map(({ entry_index: entryIndex, code }) => `entry ${entryIndex}: ${code}`)
+      .join(", ");
+    context.ui.notify(
+      `Minimal subagents Registry ignored ${diagnostics.length} invalid active-branch record${diagnostics.length === 1 ? "" : "s"} (${summaries}${diagnostics.length > 3 ? ", …" : ""}).`,
+      "warning",
+    );
+  };
+}
+
+function interruptSelectedForkBranch(snapshot: RegistrySnapshot): RegistrySnapshot {
+  const interrupted = structuredClone(snapshot);
+  for (const agent of interrupted.agents) {
+    if (!agent.active_turn_id) continue;
+    agent.latest_result = {
+      agent_id: agent.agent_id,
+      turn_id: agent.active_turn_id,
+      status: "interrupted",
+      output: "",
+      error: "Turn interrupted because its source branch was forked",
+    };
+    agent.active_turn_id = undefined;
+    agent.active_turn_started_at = undefined;
+  }
+  return interrupted;
+}
+
+function orderForkAgentsParentFirst(agents: readonly PersistedAgent[]): PersistedAgent[] {
+  const byId = new Map(agents.map((agent) => [agent.agent_id, agent]));
+  const depth = (agent: PersistedAgent) => {
+    let value = 0;
+    let parentId = agent.parent_id;
+    const visited = new Set<string>();
+    while (parentId !== "root" && !visited.has(parentId)) {
+      visited.add(parentId);
+      value++;
+      parentId = byId.get(parentId)?.parent_id ?? "root";
+    }
+    return value;
+  };
+  return [...agents].sort((left, right) => depth(left) - depth(right));
+}
+
+function unavailableForkAgent(agent: PersistedAgent, error: string): PersistedAgent {
+  return {
+    ...structuredClone(agent),
+    session_file: undefined,
+    session_id: undefined,
+    session_leaf_id: undefined,
+    clone_error: error,
+    active_turn_id: undefined,
+    active_turn_started_at: undefined,
+    availability: "unavailable",
+    missing_dependencies: [error],
+    unavailable_reason: error,
+  };
+}
+
+async function cloneSelectedForkSessions(
+  sessionFactory: AgentSessionFactory,
+  snapshot: RegistrySnapshot,
+  sourceRootSessionFile: string,
+  sourceRootSessionId: string,
+  context: ExtensionContext,
+): Promise<ForkSnapshot> {
+  const agents: PersistedAgent[] = [];
+  const failedSubtrees = new Set<string>();
+  for (const agent of orderForkAgentsParentFirst(snapshot.agents)) {
+    const failedAncestor = [...failedSubtrees].find(
+      (agentId) => agent.agent_id === agentId || agent.agent_id.startsWith(`${agentId}.`),
+    );
+    if (failedAncestor) {
+      agents.push(unavailableForkAgent(agent, `Ancestor clone failed: ${failedAncestor}`));
+      continue;
+    }
+    try {
+      const clone = await sessionFactory.cloneForkSourceSession(agent, sourceRootSessionId);
+      if (!clone.sessionLeafId) {
+        throw new Error(
+          `Minimal subagents fork recovery: no selected session leaf for ${agent.agent_id}`,
+        );
+      }
+      agents.push({
+        ...structuredClone(agent),
+        session_file: clone.sessionFile,
+        session_id: clone.sessionId,
+        session_leaf_id: clone.sessionLeafId,
+      });
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      failedSubtrees.add(agent.agent_id);
+      agents.push(unavailableForkAgent(agent, message));
+      context.ui.notify(`Fork recovery clone failed for ${agent.agent_id}: ${message}`, "error");
+    }
+  }
+  return {
+    ...structuredClone(snapshot),
+    source_root_session_file: sourceRootSessionFile,
+    source_root_session_id: sourceRootSessionId,
+    agents,
+  };
+}
+
+async function bindForkSnapshotToDestination(
+  sessionFactory: AgentSessionFactory,
+  snapshot: ForkSnapshot,
+  context: ExtensionContext,
+): Promise<ForkSnapshot> {
+  const agents: PersistedAgent[] = [];
+  const failedSubtrees = new Set<string>();
+  for (const agent of orderForkAgentsParentFirst(snapshot.agents)) {
+    const failedAncestor = [...failedSubtrees].find(
+      (agentId) => agent.agent_id === agentId || agent.agent_id.startsWith(`${agentId}.`),
+    );
+    if (failedAncestor || !agent.session_file || !agent.session_id) {
+      agents.push(
+        unavailableForkAgent(
+          agent,
+          agent.clone_error ?? `Ancestor ownership failed: ${failedAncestor ?? agent.agent_id}`,
+        ),
+      );
+      continue;
+    }
+    try {
+      const owned = await sessionFactory.adoptForkSessionOwnership(
+        agent,
+        snapshot.source_root_session_id,
+      );
+      if (!owned.sessionLeafId) {
+        throw new Error(
+          `Minimal subagents fork ownership: no selected session leaf for ${agent.agent_id}`,
+        );
+      }
+      agents.push({
+        ...structuredClone(agent),
+        session_file: owned.sessionFile,
+        session_id: owned.sessionId,
+        session_leaf_id: owned.sessionLeafId,
+      });
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      failedSubtrees.add(agent.agent_id);
+      agents.push(unavailableForkAgent(agent, message));
+      context.ui.notify(`Fork ownership failed for ${agent.agent_id}: ${message}`, "error");
+    }
+  }
+  return { ...structuredClone(snapshot), agents };
 }
 
 async function waitForRootSessionIdle(context: ExtensionContext): Promise<void> {
@@ -123,18 +295,49 @@ async function waitForRootSessionIdle(context: ExtensionContext): Promise<void>
   }
 }
 
-/** Register the six root coordinator tools and bind root-owned persistent subagent lifecycle hooks. */
-export default function minimalSubagentsExtension(pi: ExtensionAPI) {
-  let coordinator: MinimalSubagentsCoordinator | undefined;
-  let uiController: MinimalSubagentsUiController | undefined;
-  let preparedFork: ForkSnapshot | undefined;
+/** SDK and runtime construction effects required by the Minimal Subagents lifecycle controller. */
+export interface MinimalSubagentsLifecycleEffects {
+  /** Resolve Pi's current agent directory without coupling lifecycle tests to process configuration. */
+  getAgentDirectory(): string;
+  /** Construct the Pi child-session adapter after root settings and model scope are resolved. */
+  createSessionFactory(options: PiAgentSessionFactoryOptions): AgentSessionFactory;
+}
+
+const productionLifecycleEffects: MinimalSubagentsLifecycleEffects = {
+  getAgentDirectory: getAgentDir,
+  createSessionFactory: (options) => new PiAgentSessionFactory(options),
+};
 
-  pi.registerMessageRenderer("minimal-subagents.message", renderMinimalSubagentsMessage);
-  pi.registerMessageRenderer("minimal-subagents.result", renderMinimalSubagentsResult);
+/** Own coordinator, UI, and prepared-fork state for one root Pi session lifecycle. */
+export class MinimalSubagentsLifecycleController {
+  private coordinator: MinimalSubagentsCoordinator | undefined;
+  private uiController: MinimalSubagentsUiController | undefined;
+  private preparedFork:
+    | { sourceSessionFile: string; selectedBranchSnapshot: RegistrySnapshot }
+    | undefined;
+
+  /** Bind one controller to one Pi extension instance and its runtime construction effects. */
+  constructor(
+    private readonly pi: ExtensionAPI,
+    private readonly effects: MinimalSubagentsLifecycleEffects,
+  ) {}
+
+  /** Register renderers and the six Pi lifecycle event handlers owned by Minimal Subagents. */
+  register(): void {
+    this.pi.registerMessageRenderer("minimal-subagents.message", renderMinimalSubagentsMessage);
+    this.pi.registerMessageRenderer("minimal-subagents.result", renderMinimalSubagentsResult);
+
+    this.pi.on("session_start", (event, context) => this.startSession(event, context));
+    this.pi.on("session_before_fork", (event, context) => this.prepareSessionFork(event, context));
+    this.pi.on("session_tree", (event, context) => this.restoreSessionTree(event, context));
+    this.pi.on("message_end", (event, context) => this.reconcileMessageDelivery(event, context));
+    this.pi.on("agent_settled", (event, context) => this.releaseSettledRecipient(event, context));
+    this.pi.on("session_shutdown", (event, context) => this.shutdownSession(event, context));
+  }
 
-  pi.on("session_start", async (event, context) => {
+  private async startSession(event: SessionStartEvent, context: ExtensionContext): Promise<void> {
     const rootSessionId = context.sessionManager.getSessionId();
-    const agentDir = getAgentDir();
+    const agentDir = this.effects.getAgentDirectory();
     const settingsManager = SettingsManager.create(context.cwd, agentDir, {
       projectTrusted: context.isProjectTrusted(),
     });
@@ -164,10 +367,18 @@ export default function minimalSubagentsExtension(pi: ExtensionAPI) {
     ) {
       models.push(context.model);
     }
-    const availableToolNames = excludeCoordinatorTools(pi.getAllTools().map((tool) => tool.name));
+    const availableToolNames = excludeCoordinatorTools(
+      this.pi.getAllTools().map((tool) => tool.name),
+    );
     const schemas = createCoordinatorToolSchemas(eligibleModelIds);
-    let activeCoordinator!: MinimalSubagentsCoordinator;
-    const sessionFactory = new PiAgentSessionFactory({
+    let activeCoordinator: MinimalSubagentsCoordinator | undefined;
+    const requireActiveCoordinator = (): MinimalSubagentsCoordinator => {
+      if (!activeCoordinator) {
+        throw new Error("Minimal subagents lifecycle: coordinator is not initialized");
+      }
+      return activeCoordinator;
+    };
+    const sessionFactory = this.effects.createSessionFactory({
       cwd: context.cwd,
       agentDir,
       sessionDir: context.sessionManager.getSessionDir(),
@@ -179,54 +390,94 @@ export default function minimalSubagentsExtension(pi: ExtensionAPI) {
       availableToolNames,
       projectTrusted: context.isProjectTrusted(),
       maxSubagentDepth: minimalSubagentsConfig.maxSubagentDepth,
-      onChildSessionActivity: () => activeCoordinator.scheduleDeliveryReconciliation(),
-      getCoordinatorTools: (callerId) =>
-        createCoordinatorToolDefinitions({
-          coordinator: activeCoordinator,
+      onChildSessionActivity: () => activeCoordinator?.scheduleDeliveryReconciliation(),
+      getCoordinatorTools: (callerId) => {
+        const coordinator = requireActiveCoordinator();
+        return createCoordinatorToolDefinitions({
+          coordinator,
           callerId,
-          allowFanoutTools: activeCoordinator.canAgentSpawn(callerId),
+          allowFanoutTools: coordinator.canAgentSpawn(callerId),
           modelRoles: minimalSubagentsConfig.modelRoles,
           schemas,
           captureCaller: (childContext) =>
-            activeCoordinator.snapshotChildCaller(
+            coordinator.snapshotChildCaller(
               callerId,
               childContext.sessionManager.getLeafId() ?? callerId,
             ),
           onAttention: (message) => context.ui.notify(message, "error"),
-        }),
+        });
+      },
     });
     activeCoordinator = new MinimalSubagentsCoordinator({
       sessions: sessionFactory,
-      root: createRootConversationEndpoint(pi, context),
+      root: createRootConversationEndpoint(this.pi, context),
       maxSubagentDepth: minimalSubagentsConfig.maxSubagentDepth,
       registry: {
         rootSessionId,
-        append: (registryEvent) => pi.appendEntry(REGISTRY_ENTRY_TYPE, registryEvent),
+        append: (registryEvent) => this.pi.appendEntry(REGISTRY_ENTRY_TYPE, registryEvent),
       },
       notify: (notification) => {
-        uiController?.refresh();
+        this.uiController?.refresh();
         if (shouldSurfaceNotification(notification)) {
           context.ui.notify(notification.message, notificationLevel(notification));
         }
       },
     });
-    coordinator = activeCoordinator;
+    this.coordinator = activeCoordinator;
 
     let snapshot: RegistrySnapshot;
     if (event.reason === "fork" && event.previousSessionFile) {
+      const previousSession = SessionManager.open(event.previousSessionFile);
+      const sourceRootSessionId = previousSession.getSessionId();
       let forkSnapshot = takeForkSnapshot(event.previousSessionFile);
+      if (forkSnapshot?.source_root_session_id !== sourceRootSessionId) {
+        if (forkSnapshot) {
+          context.ui.notify(
+            "Minimal subagents fork handoff rejected because its source root identity did not match.",
+            "error",
+          );
+        }
+        forkSnapshot = undefined;
+      }
       if (!forkSnapshot) {
-        await activeCoordinator.restore(replayPreviousRoot(event.previousSessionFile));
-        forkSnapshot = await activeCoordinator.prepareFork(event.previousSessionFile);
+        if (
+          isForkDestinationForSource(context.sessionManager.getHeader(), event.previousSessionFile)
+        ) {
+          const selectedSnapshot = interruptSelectedForkBranch(
+            replayRegistryEntries(
+              context.sessionManager.getBranch(),
+              sourceRootSessionId,
+              reportInvalidRegistryRecords(context),
+            ),
+          );
+          forkSnapshot = await cloneSelectedForkSessions(
+            sessionFactory,
+            selectedSnapshot,
+            event.previousSessionFile,
+            sourceRootSessionId,
+            context,
+          );
+        } else {
+          context.ui.notify(
+            "Minimal subagents fork recovery skipped because the destination selected branch could not be proven from parentSession provenance.",
+            "warning",
+          );
+        }
       }
-      snapshot = forkSnapshot;
+      snapshot = forkSnapshot
+        ? await bindForkSnapshotToDestination(sessionFactory, forkSnapshot, context)
+        : { agents: [], tombstones: [], deliveries: [] };
     } else {
-      snapshot = replayRegistryEntries(context.sessionManager.getEntries(), rootSessionId);
+      snapshot = replayRegistryEntries(
+        context.sessionManager.getBranch(),
+        rootSessionId,
+        reportInvalidRegistryRecords(context),
+      );
     }
     await activeCoordinator.restore(snapshot);
     activeCoordinator.writeCheckpoint();
-    uiController = new MinimalSubagentsUiController(activeCoordinator, context);
-    uiController.refresh();
+    this.uiController = new MinimalSubagentsUiController(activeCoordinator, context);
+    this.uiController.refresh();
 
     const rootTools = createCoordinatorToolDefinitions({
       coordinator: activeCoordinator,
@@ -234,46 +485,113 @@ export default function minimalSubagentsExtension(pi: ExtensionAPI) {
       allowFanoutTools: true,
       modelRoles: minimalSubagentsConfig.modelRoles,
       schemas,
-      captureCaller: (toolContext) => rootCallerSnapshot(pi, toolContext),
-      onActivity: () => uiController?.refresh(),
+      captureCaller: (toolContext) => rootCallerSnapshot(this.pi, toolContext),
+      onActivity: () => this.uiController?.refresh(),
       onAttention: (message) => context.ui.notify(message, "error"),
     });
-    for (const tool of rootTools) pi.registerTool(tool);
-    pi.setActiveTools([...new Set([...pi.getActiveTools(), ...COORDINATOR_TOOL_NAMES])]);
+    for (const tool of rootTools) this.pi.registerTool(tool);
+    this.pi.setActiveTools([...new Set([...this.pi.getActiveTools(), ...COORDINATOR_TOOL_NAMES])]);
 
-    if (hasHistoricalChildIdentity(context.sessionManager.getEntries() as SessionEntry[])) {
+    if (hasHistoricalChildIdentity(context.sessionManager.getBranch())) {
       context.ui.notify(
         "Opened a former subagent session directly. It is now an independent root; former descendants and parent messaging were not restored. Concurrent ownership by its original root is unsupported.",
         "warning",
       );
     }
-  });
+  }
 
-  pi.on("session_before_fork", async (_event, context) => {
-    const sessionFile = context.sessionManager.getSessionFile();
-    if (!coordinator || !sessionFile) return;
-    preparedFork = await coordinator.prepareFork(sessionFile);
-    rememberForkSnapshot(preparedFork);
-  });
+  private async prepareSessionFork(
+    event: SessionBeforeForkEvent,
+    context: ExtensionContext,
+  ): Promise<void> {
+    const sourceSessionFile = context.sessionManager.getSessionFile();
+    if (!sourceSessionFile) {
+      this.preparedFork = undefined;
+      return;
+    }
+    const selectedEntry = context.sessionManager.getEntry(event.entryId);
+    const selectedLeafId =
+      event.position === "before" ? (selectedEntry?.parentId ?? undefined) : event.entryId;
+    const selectedBranch =
+      event.position === "before" && selectedEntry?.parentId === null
+        ? []
+        : context.sessionManager.getBranch(selectedLeafId);
+    this.preparedFork = {
+      sourceSessionFile,
+      selectedBranchSnapshot: interruptSelectedForkBranch(
+        replayRegistryEntries(
+          selectedBranch,
+          context.sessionManager.getSessionId(),
+          reportInvalidRegistryRecords(context),
+        ),
+      ),
+    };
+  }
 
-  pi.on("message_end", async (event) => {
-    if (!coordinator) return;
+  private async restoreSessionTree(
+    _event: SessionTreeEvent,
+    context: ExtensionContext,
+  ): Promise<void> {
+    if (!this.coordinator) return;
+    const snapshot = replayRegistryEntries(
+      context.sessionManager.getBranch(),
+      context.sessionManager.getSessionId(),
+      reportInvalidRegistryRecords(context),
+    );
+    await this.coordinator.restore(snapshot);
+    this.coordinator.writeCheckpoint();
+    this.uiController?.refresh();
+  }
+
+  private async reconcileMessageDelivery(
+    event: MessageEndEvent,
+    _context: ExtensionContext,
+  ): Promise<void> {
+    if (!this.coordinator) return;
     if (event.message.role === "toolResult" || event.message.role === "custom") {
-      await coordinator.reconcileDeliveries();
-      uiController?.refresh();
+      await this.coordinator.reconcileDeliveries();
+      this.uiController?.refresh();
     }
-  });
+  }
 
-  pi.on("session_shutdown", async (event, context) => {
-    if (coordinator) {
-      await shutdownMinimalSubagentsSession(event.reason, coordinator, {
+  private releaseSettledRecipient(_event: AgentSettledEvent, context: ExtensionContext): void {
+    if (!this.coordinator) return;
+    if (context.isIdle()) this.coordinator.markRecipientIdle("root");
+    this.uiController?.refresh();
+  }
+
+  private async shutdownSession(
+    event: SessionShutdownEvent,
+    context: ExtensionContext,
+  ): Promise<void> {
+    if (this.coordinator) {
+      if (event.reason === "fork" && this.preparedFork) {
+        await this.coordinator.restore(this.preparedFork.selectedBranchSnapshot);
+        rememberForkSnapshot(
+          await this.coordinator.prepareFork(this.preparedFork.sourceSessionFile),
+        );
+      }
+      await shutdownMinimalSubagentsSession(event.reason, this.coordinator, {
         isRootIdle: () => context.isIdle(),
         waitForRootIdle: () => waitForRootSessionIdle(context),
       });
     }
-    uiController?.dispose();
-    uiController = undefined;
-    coordinator = undefined;
-    preparedFork = undefined;
-  });
+    this.uiController?.dispose();
+    this.uiController = undefined;
+    this.coordinator = undefined;
+    this.preparedFork = undefined;
+  }
+}
+
+/** Compose the Minimal Subagents extension with production or faithful SDK runtime effects. */
+export function createMinimalSubagentsExtension(
+  effects: MinimalSubagentsLifecycleEffects = productionLifecycleEffects,
+): ExtensionFactory {
+  return (pi) => {
+    new MinimalSubagentsLifecycleController(pi, effects).register();
+  };
 }
+
+const minimalSubagentsExtension = createMinimalSubagentsExtension();
+
+export default minimalSubagentsExtension;

package/src/minimal-subagents-fork-lifecycle.ts

@@ -2,16 +2,14 @@ import { existsSync, realpathSync } from "node:fs";
 import { resolve } from "node:path";
 import type { ForkSnapshot } from "./minimal-subagents-types.js";
 
-const FORK_SNAPSHOT_SYMBOL = Symbol.for("minimal-subagents.pending-fork-snapshots.v1");
-
-type GlobalWithForkSnapshots = typeof globalThis & {
-  [FORK_SNAPSHOT_SYMBOL]?: Map<string, ForkSnapshot>;
-};
+declare global {
+  // eslint-disable-next-line no-var -- A process-global handoff must be visible to replacement extension instances.
+  var minimalSubagentsForkSnapshots: Map<string, ForkSnapshot> | undefined;
+}
 
 function forkSnapshotStore(): Map<string, ForkSnapshot> {
-  const processGlobal = globalThis as GlobalWithForkSnapshots;
-  processGlobal[FORK_SNAPSHOT_SYMBOL] ??= new Map();
-  return processGlobal[FORK_SNAPSHOT_SYMBOL];
+  globalThis.minimalSubagentsForkSnapshots ??= new Map();
+  return globalThis.minimalSubagentsForkSnapshots;
 }
 
 function canonicalSessionFile(sessionFile: string): string {
@@ -19,12 +17,24 @@ function canonicalSessionFile(sessionFile: string): string {
   return existsSync(absolutePath) ? realpathSync(absolutePath) : absolutePath;
 }
 
+/** Prove a process-loss fork destination was derived from the expected canonical source file. */
+export function isForkDestinationForSource(
+  destinationHeader: { parentSession?: string } | null,
+  previousSessionFile: string,
+): boolean {
+  return (
+    destinationHeader?.parentSession !== undefined &&
+    canonicalSessionFile(destinationHeader.parentSession) ===
+      canonicalSessionFile(previousSessionFile)
+  );
+}
+
 /** Retain a complete pre-fork hierarchy across Pi extension-instance replacement. */
 export function rememberForkSnapshot(snapshot: ForkSnapshot): void {
-  forkSnapshotStore().set(
-    canonicalSessionFile(snapshot.source_root_session_file),
-    structuredClone(snapshot),
-  );
+  const canonicalSourceFile = canonicalSessionFile(snapshot.source_root_session_file);
+  const retained = structuredClone(snapshot);
+  retained.source_root_session_file = canonicalSourceFile;
+  forkSnapshotStore().set(canonicalSourceFile, retained);
 }
 
 /** Consume the pre-fork hierarchy once when the destination root session starts. */

package/src/minimal-subagents-message-envelope.ts

@@ -0,0 +1,20 @@
+import type { CoordinatorMessage } from "./minimal-subagents-types.js";
+
+/** Add the stable source identity that the receiving model must see. */
+export function addCoordinatorMessageEnvelope(message: CoordinatorMessage): CoordinatorMessage {
+  const kind = message.customType === "minimal-subagents.result" ? "result" : "message";
+  const status = message.details.status ? ` | status=${message.details.status}` : "";
+  const envelope = `[Subagent ${kind} | agent=${message.details.source_agent_id} | turn=${message.details.source_turn_id}${status}]`;
+  return {
+    ...message,
+    content: message.content.length > 0 ? `${envelope}\n${message.content}` : envelope,
+  };
+}
+
+/** Remove the model-only envelope from the styled TUI message body. */
+export function stripCoordinatorMessageEnvelope(content: string): string {
+  return content.replace(
+    /^\[Subagent (?:message|result) \| agent=[^|\]\n]+ \| turn=[^|\]\n]+(?: \| status=[^|\]\n]+)?\]\n?/,
+    "",
+  );
+}