diff --git a/extensions/subagents/id-sequence.test.ts b/extensions/subagents/id-sequence.test.ts new file mode 100644 index 00000000..4b78aea1 --- /dev/null +++ b/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) { + 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, + }); +}); diff --git a/extensions/subagents/index.test.ts b/extensions/subagents/index.test.ts index 5fd1f9be..9d037509 100644 --- a/extensions/subagents/index.test.ts +++ b/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); diff --git a/extensions/subagents/index.ts b/extensions/subagents/index.ts index d595392b..8f714248 100644 --- a/extensions/subagents/index.ts +++ b/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 | 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", diff --git a/extensions/subagents/manager.test.ts b/extensions/subagents/manager.test.ts index 8dd45c3c..1cfd7ddb 100644 --- a/extensions/subagents/manager.test.ts +++ b/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. diff --git a/extensions/subagents/src/id-sequence.ts b/extensions/subagents/src/id-sequence.ts new file mode 100644 index 00000000..d9ff1556 --- /dev/null +++ b/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 | undefined { + return typeof value === "object" && value !== null + ? (value as Record) + : 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 }; +} diff --git a/extensions/subagents/src/manager.ts b/extensions/subagents/src/manager.ts index 763a3f7e..816344f6 100644 --- a/extensions/subagents/src/manager.ts +++ b/extensions/subagents/src/manager.ts @@ -218,8 +218,8 @@ const makeManager = (config: SubagentManagerConfig = {}) => let changeWaiters: Array<() => void> = []; const idListeners = new Map void>>(); const cleanups = new Set>(); - 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 = {}) => diff --git a/extensions/subagents/src/runtime.ts b/extensions/subagents/src/runtime.ts index 3ecdbcd5..09b4c852 100644 --- a/extensions/subagents/src/runtime.ts +++ b/extensions/subagents/src/runtime.ts @@ -18,12 +18,15 @@ const BackendRegistryLive = Layer.sync(BackendRegistry, () => { ); }); -import { SubagentManagerLive } from "./manager.ts"; +import { + makeSubagentManagerLayer, + type SubagentManagerConfig, +} from "./manager.ts"; -const AppLayer = SubagentManagerLive.pipe(Layer.provide(BackendRegistryLive)); - -export function createSubagentRuntime() { - return ManagedRuntime.make(AppLayer); +export function createSubagentRuntime(config: SubagentManagerConfig = {}) { + return ManagedRuntime.make( + makeSubagentManagerLayer(config).pipe(Layer.provide(BackendRegistryLive)), + ); } export type SubagentRuntime = ReturnType;