fix(subagents): preserve short ids across reloads

extensions/subagents/id-sequence.test.ts

@@ -0,0 +1,130 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+import type { SessionEntry } from "@earendil-works/pi-coding-agent";
+import {
+  restoreSubagentIdCounters,
+  SUBAGENT_ID_WATERMARK_ENTRY_TYPE,
+  subagentIdWatermark,
+} from "./src/id-sequence.ts";
+
+function entry(value: Partial<SessionEntry>) {
+  return {
+    id: crypto.randomUUID(),
+    parentId: null,
+    timestamp: new Date().toISOString(),
+    ...value,
+  } as SessionEntry;
+}
+
+test("watermarks survive reload and retain independent model and btw sequences", () => {
+  const branch = [
+    entry({
+      type: "custom",
+      customType: SUBAGENT_ID_WATERMARK_ENTRY_TYPE,
+      data: subagentIdWatermark("sa-7"),
+    }),
+    entry({
+      type: "custom",
+      customType: SUBAGENT_ID_WATERMARK_ENTRY_TYPE,
+      data: subagentIdWatermark("btw-3"),
+    }),
+  ];
+
+  assert.deepEqual(restoreSubagentIdCounters(branch), {
+    modelCounter: 7,
+    btwCounter: 3,
+  });
+});
+
+test("restore takes the high water when concurrent completions arrive out of order", () => {
+  const branch = ["sa-9", "sa-2", "sa-12"].map((id) =>
+    entry({
+      type: "custom",
+      customType: SUBAGENT_ID_WATERMARK_ENTRY_TYPE,
+      data: subagentIdWatermark(id),
+    }),
+  );
+
+  assert.equal(restoreSubagentIdCounters(branch).modelCounter, 12);
+});
+
+test("legacy tool and result entries migrate without a watermark", () => {
+  const branch = [
+    entry({
+      type: "message",
+      message: {
+        role: "toolResult",
+        toolCallId: "call-1",
+        toolName: "subagent_spawn",
+        content: [{ type: "text", text: "Spawned" }],
+        details: { id: "sa-4" },
+        isError: false,
+        timestamp: Date.now(),
+      },
+    }),
+    entry({
+      type: "custom",
+      customType: "subagent-result",
+      data: { details: { id: "sa-6" } },
+    }),
+    entry({
+      type: "custom",
+      customType: "btw-result",
+      data: { id: "btw-5" },
+    }),
+  ];
+
+  assert.deepEqual(restoreSubagentIdCounters(branch), {
+    modelCounter: 6,
+    btwCounter: 5,
+  });
+});
+
+test("a fork restores only ids visible on its selected branch", () => {
+  const beforeFork = entry({
+    type: "custom",
+    customType: SUBAGENT_ID_WATERMARK_ENTRY_TYPE,
+    data: subagentIdWatermark("sa-2"),
+  });
+  const otherBranch = entry({
+    type: "custom",
+    customType: SUBAGENT_ID_WATERMARK_ENTRY_TYPE,
+    data: subagentIdWatermark("sa-99"),
+  });
+
+  assert.equal(restoreSubagentIdCounters([beforeFork]).modelCounter, 2);
+  assert.equal(
+    restoreSubagentIdCounters([beforeFork, otherBranch]).modelCounter,
+    99,
+  );
+});
+
+test("malformed and unrelated entries cannot advance either sequence", () => {
+  const branch = [
+    entry({
+      type: "custom",
+      customType: SUBAGENT_ID_WATERMARK_ENTRY_TYPE,
+      data: { version: 2, id: "sa-100" },
+    }),
+    entry({
+      type: "custom",
+      customType: SUBAGENT_ID_WATERMARK_ENTRY_TYPE,
+      data: { version: 1, id: "sa-0" },
+    }),
+    entry({
+      type: "custom",
+      customType: "another-extension",
+      data: { id: "sa-50" },
+    }),
+    entry({
+      type: "custom",
+      customType: SUBAGENT_ID_WATERMARK_ENTRY_TYPE,
+      data: { version: 1, id: "sa-3-extra" },
+    }),
+  ];
+
+  assert.deepEqual(restoreSubagentIdCounters(branch), {
+    modelCounter: 0,
+    btwCounter: 0,
+  });
+});

extensions/subagents/index.test.ts

@@ -11,6 +11,8 @@ import type {
 import { PLAN_MODE_CHANNEL } from "../shared/plan-mode-state.ts";
 import subagents, { createSubagentResultDispatcher } from "./index.ts";
 
+const emptySessionManager = { getBranch: () => [] };
+
 test("subagent results render before the hidden wake-up message", () => {
   const events: unknown[] = [];
   const pi = {
@@ -164,6 +166,7 @@ test("session start preserves the complete registered subagent family", () => {
     cwd: process.cwd(),
     hasUI: false,
     isProjectTrusted: () => false,
+    sessionManager: emptySessionManager,
   } as unknown as ExtensionContext);
 
   assert.deepEqual(
@@ -217,6 +220,7 @@ test("the complete subagent family fails closed before the first spawn", async (
     hasUI: false,
     isIdle: () => true,
     isProjectTrusted: () => false,
+    sessionManager: emptySessionManager,
   } as unknown as ExtensionContext;
 
   subagents(pi);
@@ -337,6 +341,7 @@ test("session_start re-registers agent types for its cwd and live trust decision
       cwd,
       hasUI: false,
       isProjectTrusted: () => true,
+      sessionManager: emptySessionManager,
     } as unknown as ExtensionContext);
     assert.ok(spawnTools.length > 1, "session_start re-registers spawn");
     assert.ok(
@@ -357,6 +362,7 @@ test("session_start re-registers agent types for its cwd and live trust decision
       cwd: alternateCwd,
       hasUI: false,
       isProjectTrusted: () => true,
+      sessionManager: emptySessionManager,
     } as unknown as ExtensionContext);
     assert.ok(
       spawnTools
@@ -374,6 +380,7 @@ test("session_start re-registers agent types for its cwd and live trust decision
       cwd,
       hasUI: false,
       isProjectTrusted: () => false,
+      sessionManager: emptySessionManager,
     } as unknown as ExtensionContext);
     assert.equal(
       spawnTools
@@ -431,6 +438,7 @@ test("session_start re-registers agent types for its cwd and live trust decision
       cwd,
       hasUI: false,
       isProjectTrusted: () => true,
+      sessionManager: emptySessionManager,
     } as unknown as ExtensionContext);
     const trustedSpawn = spawnTools.at(-1)?.execute;
     assert.ok(trustedSpawn);

extensions/subagents/index.ts

@@ -116,6 +116,12 @@ import {
   runTool,
   type SubagentRuntime,
 } from "./src/runtime.ts";
+import {
+  restoreSubagentIdCounters,
+  SUBAGENT_ID_WATERMARK_ENTRY_TYPE,
+  subagentIdWatermark,
+  type SubagentIdCounters,
+} from "./src/id-sequence.ts";
 import {
   normalizeSubagentTitle,
   selectSubagentStripEntry,
@@ -301,6 +307,10 @@ function renderSubagentResult(
 export default function (pi: ExtensionAPI) {
   let runtime: SubagentRuntime | undefined;
   let managerPromise: Promise<SubagentManagerShape> | undefined;
+  let restoredIdCounters: SubagentIdCounters = {
+    modelCounter: 0,
+    btwCounter: 0,
+  };
   let sessionContext: ExtensionContext | undefined;
   let ui: ExtensionUIContext | undefined;
   let unsubStatus: (() => void) | undefined;
@@ -329,7 +339,14 @@ export default function (pi: ExtensionAPI) {
       enable: OPENPI_TOOL_SURFACE.subagents.entry,
     });
 
-  const getRuntime = () => (runtime ??= createSubagentRuntime());
+  const getRuntime = () =>
+    (runtime ??= createSubagentRuntime({
+      initialModelCounter: restoredIdCounters.modelCounter,
+      initialBtwCounter: restoredIdCounters.btwCounter,
+    }));
+
+  const persistId = (id: string) =>
+    pi.appendEntry(SUBAGENT_ID_WATERMARK_ENTRY_TYPE, subagentIdWatermark(id));
 
   /** Resolve the manager service once per runtime and wire the extension hooks. */
   const getManager = () => {
@@ -487,6 +504,9 @@ export default function (pi: ExtensionAPI) {
   };
 
   pi.on("session_start", (_event, ctx) => {
+    restoredIdCounters = restoreSubagentIdCounters(
+      ctx.sessionManager.getBranch(),
+    );
     refreshAgentTypes(ctx.cwd, ctx.isProjectTrusted());
     registerStableToolFamily();
     sessionContext = ctx;
@@ -726,6 +746,7 @@ export default function (pi: ExtensionAPI) {
         if (worktree) await reclaimWorktree(cwd, worktree).catch(() => {});
         throw error;
       }
+      persistId(snap.id);
 
       return {
         content: [
@@ -1214,6 +1235,7 @@ export default function (pi: ExtensionAPI) {
       );
       return;
     }
+    persistId(snap.id);
 
     await openSubagentTakeover(ctx, manager.view, snap.id, {
       badge: "by the way",

extensions/subagents/manager.test.ts

@@ -166,6 +166,25 @@ test("spawn origin propagates to ids, snapshots, and settlement", async () => {
   });
 });
 
+test("restored session counters continue both id sequences without reuse", async () => {
+  await withManager(
+    async (manager, runtime) => {
+      const model = await runTool(
+        runtime,
+        manager.spawn("pi", task("model task")),
+      );
+      const btw = await runTool(
+        runtime,
+        manager.spawn("pi", { ...task("side question"), origin: "btw" }),
+      );
+
+      assert.equal(model.id, "sa-42");
+      assert.equal(btw.id, "btw-8");
+    },
+    { initialModelCounter: 41, initialBtwCounter: 7 },
+  );
+});
+
 test("by-the-way sessions run in their own pool, separate from the model pool", async () => {
   await withManager(async (manager, runtime) => {
     // Fill the entire btw pool.

extensions/subagents/src/id-sequence.ts

@@ -0,0 +1,84 @@
+import type { SessionEntry } from "@earendil-works/pi-coding-agent";
+
+export const SUBAGENT_ID_WATERMARK_ENTRY_TYPE = "subagent-id-watermark";
+
+export interface SubagentIdCounters {
+  readonly modelCounter: number;
+  readonly btwCounter: number;
+}
+
+interface IdWatermarkData {
+  readonly version: 1;
+  readonly id: string;
+}
+
+function record(value: unknown): Record<string, unknown> | undefined {
+  return typeof value === "object" && value !== null
+    ? (value as Record<string, unknown>)
+    : undefined;
+}
+
+function idFrom(value: unknown) {
+  const candidate = record(value)?.id;
+  return typeof candidate === "string" ? candidate : undefined;
+}
+
+function idFromEntry(entry: SessionEntry) {
+  if (entry.type === "custom") {
+    const data = record(entry.data);
+    if (
+      entry.customType === SUBAGENT_ID_WATERMARK_ENTRY_TYPE &&
+      data?.version === 1
+    ) {
+      return idFrom(data);
+    }
+    if (
+      entry.customType === "subagent-finished" ||
+      entry.customType === "btw-result"
+    ) {
+      return idFrom(data);
+    }
+    if (entry.customType === "subagent-result") {
+      return idFrom(data?.details) ?? idFrom(data);
+    }
+    return undefined;
+  }
+
+  if (entry.type === "custom_message") {
+    return entry.customType === "subagent-result"
+      ? idFrom(entry.details)
+      : undefined;
+  }
+
+  if (entry.type !== "message") return undefined;
+  const message = record(entry.message);
+  return message?.role === "toolResult" && message.toolName === "subagent_spawn"
+    ? idFrom(message.details)
+    : undefined;
+}
+
+function sequence(id: string | undefined, prefix: "sa" | "btw") {
+  if (!id) return undefined;
+  const match = new RegExp(`^${prefix}-(\\d+)$`, "u").exec(id);
+  if (!match) return undefined;
+  const value = Number(match[1]);
+  return Number.isSafeInteger(value) && value > 0 ? value : undefined;
+}
+
+/** Restore branch-local id high-water marks, including sessions from before watermarks existed. */
+export function restoreSubagentIdCounters(
+  entries: readonly SessionEntry[],
+): SubagentIdCounters {
+  let modelCounter = 0;
+  let btwCounter = 0;
+  for (const entry of entries) {
+    const id = idFromEntry(entry);
+    modelCounter = Math.max(modelCounter, sequence(id, "sa") ?? 0);
+    btwCounter = Math.max(btwCounter, sequence(id, "btw") ?? 0);
+  }
+  return { modelCounter, btwCounter };
+}
+
+export function subagentIdWatermark(id: string): IdWatermarkData {
+  return { version: 1, id };
+}

extensions/subagents/src/manager.ts

@@ -218,8 +218,8 @@ const makeManager = (config: SubagentManagerConfig = {}) =>
     let changeWaiters: Array<() => void> = [];
     const idListeners = new Map<string, Set<() => void>>();
     const cleanups = new Set<Fiber.Fiber<unknown>>();
-    let modelCounter = 0;
-    let btwCounter = 0;
+    let modelCounter = config.initialModelCounter ?? 0;
+    let btwCounter = config.initialBtwCounter ?? 0;
     // Reservations are tracked per pool so the model and user "by the way" asides
     // never contend for the same slots.
     let reservedModel = 0;
@@ -843,6 +843,9 @@ const makeManager = (config: SubagentManagerConfig = {}) =>
 export interface SubagentManagerConfig {
   /** Test-only override for the first-response watchdog timeout. */
   firstResponseTimeoutMs?: number;
+  /** Session-branch high-water marks restored by the extension host. */
+  initialModelCounter?: number;
+  initialBtwCounter?: number;
 }
 
 export const makeSubagentManagerLayer = (config: SubagentManagerConfig = {}) =>