From 1b91e21d49d7741298ffb964e55cc3de4752bd1e Mon Sep 17 00:00:00 2001
From: "Timothy J. Aveni" <tja@berkeley.edu>
Date: Thu, 8 Oct 2026 05:19:28 -0700
Subject: [PATCH] Keep restricted engines reserved through asynchronous
 retirement

---
 restricted-execution/service.mjs      | 10 +++++---
 restricted-execution/service.test.mjs | 36 +++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/restricted-execution/service.mjs b/restricted-execution/service.mjs
index 3db0ade..9edc8bf 100644
--- a/restricted-execution/service.mjs
+++ b/restricted-execution/service.mjs
@@ -239,15 +239,17 @@ export async function serveRestrictedRegistry(config, io = process) {
       }
       const engine = call.engine;
       if (engine) {
-        // Revoke the frame token before returning an engine to the idle pool.
-        // A delayed child request can never acquire the next call's authority.
-        if (engine.current === call) engine.current = null;
         if (!stopped && reason.code === "INVOCATION_ENDED" && !usageFailure && !engine.destroying) {
           const rss = (await engineUsage(engine)).memoryBytes;
           metrics.peakRSSBytes = Math.max(metrics.peakRSSBytes, rss);
           // Bound retained engine heaps as well as lifetime. Hard aggregate
           // physical memory remains the supervising process/cgroup's contract.
-          if (rss <= 128 * 1024 * 1024) {
+          // Keep the lease busy through every await. Another artifact can evict
+          // an idle lease immediately; never publish an engine still retiring
+          // or one whose process died while its final usage read was pending.
+          if (rss <= 128 * 1024 * 1024 && !stopped && !engine.destroying && engine.current === call) {
+            // Revoke the frame token atomically with publishing idle capacity.
+            engine.current = null;
             idleEngines.push(engine);
             engine.idleUntil = Date.now() + limits.engineIdleMs;
             engine.idleTimer = setTimeout(() => void destroyEngine(engine).catch(fail), limits.engineIdleMs);
diff --git a/restricted-execution/service.test.mjs b/restricted-execution/service.test.mjs
index 409f615..5b6dc26 100644
--- a/restricted-execution/service.test.mjs
+++ b/restricted-execution/service.test.mjs
@@ -686,3 +686,39 @@ test("the shared engine budget evicts idle code families and rejects busy oversu
   }
   assert.equal(engineBudget.stats().engines, 0);
 });
+
+test("engine retirement accounting keeps its admission lease busy", async () => {
+  const observations = [];
+  let idle;
+  const engineBudget = new RestrictedEngineBudget(1, async () => ({
+    usage: async () => {
+      observations.push(idle());
+      await new Promise((done) => setImmediate(done));
+      return { cpuMicroseconds: 0, memoryBytes: 1024 };
+    },
+    release: async () => {},
+  }));
+  const acquire = engineBudget.acquire.bind(engineBudget);
+  engineBudget.acquire = (isIdle, evict) => {
+    idle = isIdle;
+    return acquire(isIdle, evict);
+  };
+  const f = await fixture(registry("return request.input;"), { engineBudget });
+  try {
+    f.start("first");
+    assert.equal((await f.receive()).output, "first");
+    assert.ok(observations.length > 0);
+    assert.ok(
+      observations.every((value) => !value),
+      "An engine must not become evictable during asynchronous retirement accounting",
+    );
+    assert.equal(f.stats().idle, 1);
+    assert.equal(idle(), true);
+    f.start("second");
+    assert.equal((await f.receive()).output, "second");
+    assert.equal(f.stats().starts, 1);
+  } finally {
+    await f.close();
+  }
+  assert.equal(engineBudget.stats().engines, 0);
+});