diff --git a/extensions/shared/child-execution-admission.ts b/extensions/shared/child-execution-admission.ts index 48534834b..c92ad51b9 100644 --- a/extensions/shared/child-execution-admission.ts +++ b/extensions/shared/child-execution-admission.ts @@ -5,7 +5,7 @@ * settled Direct children may remain available for a later `send`, but they do * not retain a permit. The parent `SessionManager` is the owner identity, so a * second top-level Pi Session gets an independent pool even though extensions - * share this module instance in one process. + * share one Session map in this process. */ export const MAX_CHILD_EXECUTION_QUEUE = 256; @@ -89,6 +89,10 @@ export class ChildExecutionAdmission { this.maxQueued = options.maxQueued ?? MAX_CHILD_EXECUTION_QUEUE; } + get isShutdown() { + return this.closed; + } + snapshot(): ChildExecutionAdmissionSnapshot { const heldByOrigin = counts(); for (const origin of this.held.values()) heldByOrigin[origin]++; @@ -232,7 +236,30 @@ export class ChildExecutionAdmission { } } -const sessionAdmissions = new WeakMap(); +// Pi loads every extension with `moduleCache: false`, so Workflows and +// Subagents each evaluate their own copy of this module. Symbol.for keeps those +// same-process copies on one versioned Session map; a module-level WeakMap +// would silently give each extension its own pool. +const SESSION_ADMISSIONS_KEY = Symbol.for( + "@tt-a1i/openpi/child-execution-admission/v1", +); +const existingAdmissions: unknown = Reflect.get( + globalThis, + SESSION_ADMISSIONS_KEY, +); +if ( + existingAdmissions !== undefined && + !(existingAdmissions instanceof WeakMap) +) { + throw new Error("Incompatible OpenPI child execution admission registry"); +} +const sessionAdmissions: WeakMap = + existingAdmissions ?? new WeakMap(); +if (existingAdmissions === undefined) { + Object.defineProperty(globalThis, SESSION_ADMISSIONS_KEY, { + value: sessionAdmissions, + }); +} /** Return the one admission owner for this top-level Pi Session. */ export function sessionChildExecutionAdmission( @@ -240,7 +267,10 @@ export function sessionChildExecutionAdmission( options: ChildExecutionAdmissionOptions, ) { let admission = sessionAdmissions.get(sessionOwner); - if (!admission) { + // Pi's reload emits session_shutdown and then session_start for the same + // SessionManager. A shut-down pool rejects every acquire, so the reloaded + // extensions start a fresh one, as separate module copies did before. + if (!admission || admission.isShutdown) { admission = new ChildExecutionAdmission(options); sessionAdmissions.set(sessionOwner, admission); } diff --git a/extensions/subagents/src/manager.ts b/extensions/subagents/src/manager.ts index aa87099d5..9a0b18b55 100644 --- a/extensions/subagents/src/manager.ts +++ b/extensions/subagents/src/manager.ts @@ -548,6 +548,7 @@ const makeManager = (config: SubagentManagerConfig = {}) => Effect.gen(function* () { const origin: SubagentOrigin = task.origin ?? "model"; let admissionLease: ChildExecutionLease | undefined; + let admissionAbandoned = false; let poolReservationHeld = false; let leaseTransferred = false; const doSpawn = Effect.gen(function* () { @@ -574,12 +575,19 @@ const makeManager = (config: SubagentManagerConfig = {}) => }, ); if (config.admission) { - admissionLease = yield* Effect.tryPromise({ + const admission = config.admission; + yield* Effect.tryPromise({ + // Take the lease in the promise continuation, not after the + // fiber resumes: an interrupt can land after admission granted + // the permit but before this fiber runs again, and the + // finalizer below must still find (or already have released) it. try: (signal) => - config.admission!.acquire( - origin === "btw" ? "btw" : "direct", - signal, - ), + admission + .acquire(origin === "btw" ? "btw" : "direct", signal) + .then((lease) => { + if (admissionAbandoned) lease.release(); + else admissionLease = lease; + }), catch: (error) => new SpawnError({ message: @@ -672,6 +680,7 @@ const makeManager = (config: SubagentManagerConfig = {}) => return yield* doSpawn.pipe( Effect.ensuring( Effect.sync(() => { + admissionAbandoned = true; if (poolReservationHeld) { if (origin === "btw") reservedBtw--; else reservedModel--; @@ -795,6 +804,12 @@ const makeManager = (config: SubagentManagerConfig = {}) => message: `Subagent "${id}" is no longer tracked.`, }); } + if (entry.admissionAbort) { + return new SendError({ + message: + "A subagent restart is waiting for a shared child execution slot.", + }); + } if (entry.admissionUncertain) { return new SendError({ message: @@ -818,22 +833,40 @@ const makeManager = (config: SubagentManagerConfig = {}) => // both pass the check in that window. Cleared by RunStarted/settle, // or here when the backend rejects the send. entry.restarting = true; + let admissionAbandoned = false; + let restartLease: ChildExecutionLease | undefined; + const admissionAbort = config.admission + ? new AbortController() + : undefined; + entry.admissionAbort = admissionAbort; return Effect.gen(function* () { - if (config.admission) { - const admissionAbort = new AbortController(); - entry.admissionAbort = admissionAbort; - const lease = yield* Effect.tryPromise({ + if (config.admission && admissionAbort) { + const admission = config.admission; + yield* Effect.tryPromise({ + // Same interrupt window as spawn: hand the lease to the entry + // from the continuation so onError below can always release it. try: (signal) => { const abortForEffect = () => admissionAbort.abort(signal.reason); signal.addEventListener("abort", abortForEffect, { once: true, }); - return config - .admission!.acquire( + return admission + .acquire( origin === "btw" ? "btw" : "direct", admissionAbort.signal, ) + .then((lease) => { + if (admissionAbandoned || admissionAbort.signal.aborted) { + lease.release(); + throw ( + admissionAbort.signal.reason ?? + new Error("Subagent restart was aborted") + ); + } + restartLease = lease; + entry.admissionLease = lease; + }) .finally(() => signal.removeEventListener("abort", abortForEffect), ); @@ -846,18 +879,36 @@ const makeManager = (config: SubagentManagerConfig = {}) => : "Could not acquire a shared child execution slot.", }), }); + // Explicit cancel can land after the promise continuation but + // before this fiber resumes. Never dispatch that stale restart. + if (admissionAbort.signal.aborted) { + return yield* new SendError({ + message: "Subagent restart was aborted", + }); + } entry.admissionAbort = undefined; - entry.admissionLease = lease; } return yield* entry.session.send(text); }).pipe( Effect.onError(() => Effect.sync(() => { - entry.restarting = false; - entry.admissionAbort?.abort(); - entry.admissionAbort = undefined; - entry.admissionLease?.release(); - entry.admissionLease = undefined; + admissionAbandoned = true; + // A cancelled attempt can reject after another send has + // already reserved this entry. Only clear our own attempt. + if ( + !admissionAbort || + entry.admissionAbort === admissionAbort || + (entry.admissionAbort === undefined && + restartLease !== undefined && + entry.admissionLease === restartLease) + ) { + entry.restarting = false; + entry.admissionAbort = undefined; + } + admissionAbort?.abort(); + if (entry.admissionLease === restartLease) + entry.admissionLease = undefined; + restartLease?.release(); notify(entry.snapshot.id); }), ), diff --git a/tests/extensions/shared/child-execution-admission.test.ts b/tests/extensions/shared/child-execution-admission.test.ts index ea9b2db34..f89bb8cf8 100644 --- a/tests/extensions/shared/child-execution-admission.test.ts +++ b/tests/extensions/shared/child-execution-admission.test.ts @@ -1,5 +1,10 @@ import assert from "node:assert/strict"; +import { mkdtemp, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import * as path from "node:path"; import test from "node:test"; +import { pathToFileURL } from "node:url"; +import { discoverAndLoadExtensions } from "@earendil-works/pi-coding-agent"; import { ChildExecutionAdmission, ChildExecutionAdmissionAbortedError, @@ -200,3 +205,70 @@ test("FAIL: a grant racing cancellation leaves no leaked lease", async () => { lease.release(); assert.equal(admission.snapshot().held, 0); }); + +test("FAIL: Pi-loaded extensions share one admission per Session and a fresh one after reload", async () => { + // Pi evaluates each extension with `moduleCache: false`, so Workflows and + // Subagents import separate copies of the admission module. Load two + // fixtures through Pi's real loader instead of one ESM test process. + const directory = await mkdtemp(path.join(tmpdir(), "pi-admission-")); + const sink = Symbol.for("@tt-a1i/openpi/test/admission-loader-sink"); + const moduleUrl = pathToFileURL( + path.resolve("extensions/shared/child-execution-admission.ts"), + ).href; + const fixture = (origin: string) => + [ + `import { sessionChildExecutionAdmission } from ${JSON.stringify(moduleUrl)};`, + "export default function () {", + ` globalThis[Symbol.for(${JSON.stringify(sink.description)})].${origin} = (owner) =>`, + " sessionChildExecutionAdmission(owner, { maxActive: 1 });", + "}", + ].join("\n"); + const factories: Record< + "workflow" | "direct", + (owner: object) => ChildExecutionAdmission + > = Object.create(null); + Reflect.set(globalThis, sink, factories); + try { + await writeFile(path.join(directory, "workflows.ts"), fixture("workflow")); + await writeFile(path.join(directory, "subagents.ts"), fixture("direct")); + const loaded = await discoverAndLoadExtensions( + [ + path.join(directory, "workflows.ts"), + path.join(directory, "subagents.ts"), + ], + directory, + path.join(directory, "agent"), + ); + assert.deepEqual(loaded.errors, []); + + const session = {}; + const workflow = factories.workflow(session); + const direct = factories.direct(session); + assert.equal(workflow, direct); + + const directLease = await direct.acquire("direct"); + const queued = workflow.acquire("workflow"); + assert.deepEqual(workflow.snapshot().heldByOrigin, { + workflow: 0, + direct: 1, + btw: 0, + }); + assert.equal(workflow.snapshot().queued, 1); + directLease.release(); + (await queued).release(); + + const otherSession = factories.direct({}); + assert.notEqual(otherSession, workflow); + + // Pi's reload shuts the old extension runtime down and then starts the + // new one with the same SessionManager. + workflow.shutdown(); + const reloaded = factories.direct(session); + assert.notEqual(reloaded, workflow); + assert.equal(factories.workflow(session), reloaded); + (await reloaded.acquire("direct")).release(); + } finally { + Reflect.deleteProperty(globalThis, sink); + await rm(directory, { recursive: true, force: true }); + } +}); diff --git a/tests/extensions/subagents/manager.test.ts b/tests/extensions/subagents/manager.test.ts index c2d9c4402..4e3ab4432 100644 --- a/tests/extensions/subagents/manager.test.ts +++ b/tests/extensions/subagents/manager.test.ts @@ -19,6 +19,7 @@ import type { SpawnTask, SubagentStatus, } from "../../../extensions/subagents/src/domain.ts"; +import { SendError } from "../../../extensions/subagents/src/domain.ts"; import { makeSubagentManagerLayer, MAX_RUNNING, @@ -354,6 +355,356 @@ test("releases local pool reservations when shared admission rejects queued spaw } }); +test("FAIL: a Direct spawn aborted as admission grants its slot releases the permit", async () => { + const admission = new ChildExecutionAdmission({ maxActive: 1 }); + await withManager( + async (manager, runtime) => { + const blocker = await admission.acquire("workflow"); + const abort = new AbortController(); + const queued = runTool( + runtime, + manager.spawn("pi", task("granted then aborted")), + { signal: abort.signal }, + ); + await new Promise((resolve) => setTimeout(resolve, 0)); + assert.equal(admission.snapshot().queuedByOrigin.direct, 1); + // The Workflow child settles and the user presses Esc in the same turn: + // admission grants the queued spawn before its fiber can resume. + blocker.release(); + abort.abort(new Error("Esc")); + await assert.rejects(queued); + await new Promise((resolve) => setTimeout(resolve, 20)); + assert.equal(manager.view.list().length, 0); + assert.equal(admission.snapshot().held, 0); + }, + { admission }, + ); +}); + +test("FAIL: one Esc aborting two parallel Direct spawns leaves no held permit", async () => { + const base = stubBackend(); + // Real Pi child startup takes long enough for one turn abort to interrupt + // the first spawn mid-startup while the second is still queued. + const slow: SubagentBackend = { + ...base, + spawn: (spawnTask) => + Effect.andThen(Effect.sleep("2 seconds"), base.spawn(spawnTask)), + }; + const admission = new ChildExecutionAdmission({ maxActive: 1 }); + await withManager( + async (manager, runtime) => { + const turn = new AbortController(); + const spawns = ["A", "B"].map((name) => + runTool(runtime, manager.spawn("pi", task(name)), { + signal: turn.signal, + }).catch((error: unknown) => error), + ); + await new Promise((resolve) => setTimeout(resolve, 20)); + assert.equal(admission.snapshot().held, 1); + assert.equal(admission.snapshot().queued, 1); + turn.abort(new Error("Esc")); + await Promise.all(spawns); + await new Promise((resolve) => setTimeout(resolve, 20)); + assert.equal(manager.view.list().length, 0); + assert.equal(admission.snapshot().held, 0); + assert.equal(admission.snapshot().queued, 0); + }, + { admission }, + [slow], + ); +}); + +test("FAIL: a Direct restart aborted as admission grants its slot releases the permit", async () => { + const admission = new ChildExecutionAdmission({ maxActive: 1 }); + await withManager( + async (manager, runtime) => { + const snap = await runTool( + runtime, + manager.spawn("pi", task("settle before restart")), + ); + await runTool(runtime, manager.waitFor([snap.id])); + assert.equal(admission.snapshot().held, 0); + + const blocker = await admission.acquire("workflow"); + const abort = new AbortController(); + const restart = runTool(runtime, manager.send(snap.id, "again"), { + signal: abort.signal, + }); + await new Promise((resolve) => setTimeout(resolve, 0)); + assert.equal(admission.snapshot().queuedByOrigin.direct, 1); + blocker.release(); + abort.abort(new Error("Esc")); + await assert.rejects(restart); + await new Promise((resolve) => setTimeout(resolve, 20)); + assert.equal(admission.snapshot().held, 0); + }, + { admission }, + ); +}); + +test("FAIL: explicit cancel at a restart grant prevents the backend send", async () => { + const admission = new ChildExecutionAdmission({ maxActive: 1 }); + const sent: string[] = []; + const base = stubBackend(); + const backend: SubagentBackend = { + ...base, + spawn: (spawnTask) => + base.spawn(spawnTask).pipe( + Effect.map((session) => ({ + ...session, + send: (text) => { + sent.push(text); + return session.send(text); + }, + })), + ), + }; + await withManager( + async (manager, runtime) => { + const snap = await runTool(runtime, manager.spawn("pi", task("First"))); + await runTool(runtime, manager.waitFor([snap.id])); + const firstGeneration = snap.runGeneration; + const firstText = snap.finalText; + const blocker = await admission.acquire("workflow"); + const restart = runTool(runtime, manager.send(snap.id, "cancelled")); + const rejected = assert.rejects(restart, /aborted/); + await new Promise((resolve) => setTimeout(resolve, 0)); + assert.equal(admission.snapshot().queued, 1); + let cancelled = false; + const unsubscribe = admission.subscribe(() => { + if (cancelled || admission.snapshot().heldByOrigin.direct !== 1) return; + cancelled = runtime.runSync(manager.cancel([snap.id]))[0]!.cancelled; + }); + try { + blocker.release(); + await rejected; + await runTool(runtime, manager.waitFor([snap.id])); + assert.equal(cancelled, true); + assert.deepEqual(sent, []); + assert.equal(snap.runGeneration, firstGeneration); + assert.equal(snap.finalText, firstText); + assert.equal(admission.snapshot().held, 0); + assert.equal(admission.snapshot().queued, 0); + } finally { + unsubscribe(); + } + }, + { admission }, + [backend], + ); +}); + +test("FAIL: cancelled restart cleanup cannot overwrite an immediate replacement", async () => { + const admission = new ChildExecutionAdmission({ maxActive: 1 }); + await withManager( + async (manager, runtime) => { + const snap = await runTool(runtime, manager.spawn("pi", task("First"))); + await runTool(runtime, manager.waitFor([snap.id])); + const firstGeneration = snap.runGeneration!; + const blocker = await admission.acquire("workflow"); + const restart = runTool(runtime, manager.send(snap.id, "cancelled")); + const rejected = assert.rejects(restart, /aborted/); + await new Promise((resolve) => setTimeout(resolve, 0)); + assert.equal(admission.snapshot().queued, 1); + let cancelled = false; + let replacement: Promise | undefined; + const unsubscribe = admission.subscribe(() => { + if (cancelled || admission.snapshot().heldByOrigin.direct !== 1) return; + cancelled = runtime.runSync(manager.cancel([snap.id]))[0]!.cancelled; + replacement = runTool(runtime, manager.send(snap.id, "replacement")); + }); + try { + blocker.release(); + await rejected; + assert.equal(cancelled, true); + assert.ok(replacement); + await replacement; + await runTool(runtime, manager.waitFor([snap.id])); + assert.equal(snap.runGeneration, firstGeneration + 1); + assert.match(snap.finalText, /replacement/); + assert.doesNotMatch(snap.finalText, /cancelled/); + assert.equal(admission.snapshot().held, 0); + assert.equal(admission.snapshot().queued, 0); + } finally { + unsubscribe(); + await replacement?.catch(() => {}); + } + }, + { admission }, + ); +}); + +test("FAIL: a second send cannot bypass a queued restart's admission", async () => { + const admission = new ChildExecutionAdmission({ maxActive: 1 }); + const sent: string[] = []; + const base = stubBackend(); + const backend: SubagentBackend = { + ...base, + spawn: (spawnTask) => + base.spawn(spawnTask).pipe( + Effect.map((session) => ({ + ...session, + send: (text) => { + sent.push(text); + return session.send(text); + }, + })), + ), + }; + await withManager( + async (manager, runtime) => { + const snap = await runTool(runtime, manager.spawn("pi", task("First"))); + await runTool(runtime, manager.waitFor([snap.id])); + const firstGeneration = snap.runGeneration; + const blocker = await admission.acquire("workflow"); + const restart = runTool(runtime, manager.send(snap.id, "admitted")); + await new Promise((resolve) => setTimeout(resolve, 0)); + const queued = admission.snapshot(); + assert.equal(queued.queuedByOrigin.direct, 1); + try { + await assert.rejects( + runTool(runtime, manager.send(snap.id, "bypass")), + /restart is waiting for a shared child execution slot/, + ); + assert.deepEqual(sent, []); + assert.equal(snap.runGeneration, firstGeneration); + assert.deepEqual(admission.snapshot(), queued); + } finally { + blocker.release(); + await restart; + } + await runTool(runtime, manager.waitFor([snap.id])); + assert.deepEqual(sent, ["admitted"]); + assert.equal(admission.snapshot().held, 0); + assert.equal(admission.snapshot().queued, 0); + }, + { admission }, + [backend], + ); +}); + +test("FAIL: cleanup after lease storage preserves a replacement's cancellation identity", async () => { + const admission = new ChildExecutionAdmission({ maxActive: 1 }); + await withManager( + async (manager, runtime) => { + const snap = await runTool(runtime, manager.spawn("pi", task("First"))); + await runTool(runtime, manager.waitFor([snap.id])); + const firstGeneration = snap.runGeneration; + const firstText = snap.finalText; + const blocker = await admission.acquire("workflow"); + const acquire = admission.acquire.bind(admission); + let intercept = true; + let firstCancelled = false; + let replacementCancelled = false; + let replacementGranted = false; + let replacement: Promise | undefined; + admission.acquire = (origin, signal) => { + const pending = acquire(origin, signal); + if (origin === "direct" && intercept) { + intercept = false; + void pending.then(() => + // Let the manager's continuation store the old lease, then + // cancel before the Effect fiber resumes and dispatches it. + queueMicrotask(() => { + firstCancelled = runtime.runSync(manager.cancel([snap.id]))[0]! + .cancelled; + replacement = runTool( + runtime, + manager.send(snap.id, "replacement"), + ); + void replacement.catch(() => {}); + }), + ); + } + return pending; + }; + const unsubscribe = admission.subscribe(() => { + const state = admission.snapshot(); + if ( + !firstCancelled || + replacementGranted || + state.heldByOrigin.direct !== 1 || + state.queued !== 0 + ) + return; + replacementGranted = true; + replacementCancelled = runtime.runSync(manager.cancel([snap.id]))[0]! + .cancelled; + }); + const restart = runTool(runtime, manager.send(snap.id, "cancelled")); + const rejected = assert.rejects(restart, /aborted/); + await new Promise((resolve) => setTimeout(resolve, 0)); + assert.equal(admission.snapshot().queued, 1); + try { + blocker.release(); + await rejected; + assert.ok(replacement); + assert.equal(replacementCancelled, true); + await assert.rejects(replacement, /aborted/); + await runTool(runtime, manager.waitFor([snap.id])); + assert.equal(snap.runGeneration, firstGeneration); + assert.equal(snap.finalText, firstText); + assert.equal(admission.snapshot().held, 0); + assert.equal(admission.snapshot().queued, 0); + } finally { + unsubscribe(); + await replacement?.catch(() => {}); + } + }, + { admission }, + ); +}); + +test("FAIL: a backend-rejected restart releases its lease and retry still needs admission", async () => { + const admission = new ChildExecutionAdmission({ maxActive: 1 }); + const base = stubBackend(); + let rejectNextSend = true; + const backend: SubagentBackend = { + ...base, + spawn: (spawnTask) => + base.spawn(spawnTask).pipe( + Effect.map((session) => ({ + ...session, + send: (text) => + Effect.suspend(() => { + if (rejectNextSend) { + rejectNextSend = false; + return new SendError({ message: "backend rejected restart" }); + } + return session.send(text); + }), + })), + ), + }; + await withManager( + async (manager, runtime) => { + const snap = await runTool(runtime, manager.spawn("pi", task("First"))); + await runTool(runtime, manager.waitFor([snap.id])); + await assert.rejects( + runTool(runtime, manager.send(snap.id, "rejected")), + /backend rejected restart/, + ); + assert.equal(admission.snapshot().held, 0); + + const blocker = await admission.acquire("workflow"); + const retry = runTool(runtime, manager.send(snap.id, "retry")); + try { + await new Promise((resolve) => setTimeout(resolve, 0)); + assert.equal(admission.snapshot().queuedByOrigin.direct, 1); + } finally { + blocker.release(); + await retry; + } + await runTool(runtime, manager.waitFor([snap.id])); + assert.match(snap.finalText, /retry/); + assert.equal(admission.snapshot().held, 0); + assert.equal(admission.snapshot().queued, 0); + }, + { admission }, + [backend], + ); +}); + test("FAIL: an uncertain Direct stop retains its shared slot and cannot restart", async () => { const admission = new ChildExecutionAdmission({ maxActive: 1 }); await withManager(