fix(runtime): 共享子执行准入池，并在中断时释放已分配的槽位

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,15 +236,41 @@ export class ChildExecutionAdmission {
   }
 }
 
-const sessionAdmissions = new WeakMap<object, ChildExecutionAdmission>();
+// 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<object, ChildExecutionAdmission> =
+  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(
   sessionOwner: object,
   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);
   }

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);
               }),
             ),

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 });
+  }
+});