fix(runtime-host): defer resource drain outside admission

packages/runtime-host/src/__tests__/host-kernel.test.ts

@@ -86,6 +86,7 @@ import {
 import type { RuntimeHostCompositionSource } from '../server/host-composition.js';
 import { createUnavailableDomainOperationHandlers } from '../server/operation-dispatcher.js';
 import { HostChangeFeed } from '../server/host-change-feed.js';
+import { SessionAdmissionGate } from '../server/session-admission-gate.js';
 import { FramedTransport, RuntimeHostTransportError } from '../transport/framed-transport.js';
 import {
   prepareStorageRootControlDirectory,
@@ -946,6 +947,37 @@ describe('non-serving Runtime Host kernel', () => {
     });
   });
 
+  test('requestDrain leaves an active Session admission before beginning composition drain', async () => {
+    await withHostPaths(async (paths) => {
+      const capability = await resolveStorageRoot({ path: paths.root, kind: 'interactive' });
+      const owner = await tryAcquireInteractiveRootOwner(capability);
+      assert.ok(owner);
+      let context: RuntimeHostCompositionContext | undefined;
+      let drainCalls = 0;
+      const host = await RuntimeHostKernel.start({
+        owner,
+        idleGraceMs: 10_000,
+        composition: defineInteractiveRuntimeHostComposition(async (value) => {
+          context = value;
+          return testComposition({
+            beginDrain: () => {
+              drainCalls += 1;
+            },
+          });
+        }),
+      });
+      const admission = new SessionAdmissionGate();
+
+      await admission.run('session', () => {
+        context?.requestDrain();
+        assert.equal(drainCalls, 0);
+      });
+
+      assert.equal(drainCalls, 1);
+      await host.closed;
+    });
+  });
+
   test('execution settlement can exclude environment resources without releasing Host ownership', async () => {
     await withHostPaths(async (paths) => {
       const capability = await resolveStorageRoot({ path: paths.root, kind: 'interactive' });

packages/runtime-host/src/__tests__/runtime-resource-coordinator.test.ts

@@ -34,7 +34,10 @@ import {
   HostRuntimeResourceCoordinator,
   type HostRuntimeResourceCoordinatorInput,
 } from '../server/runtime-resource-coordinator.js';
-import { SessionAdmissionGate } from '../server/session-admission-gate.js';
+import {
+  runAfterCurrentSessionAdmission,
+  SessionAdmissionGate,
+} from '../server/session-admission-gate.js';
 
 const SESSION_ID = 'session-1';
 const RUNTIME_REF = 'maka://runtime/background-tasks/shell-1';
@@ -227,7 +230,8 @@ describe('Host Runtime Resource coordinator', () => {
     assert.equal(!revoked.ok && revoked.error.code, 'not_found');
   });
 
-  test('drains for canonical state failure but keeps projection failure scoped to its query', async () => {
+  test('drains for canonical state failure but keeps projection failure scoped to its query', async (t) => {
+    t.mock.method(console, 'error', () => {});
     const harness = createHarness();
     harness.updates = [
       resourceUpdate(0, {
@@ -257,6 +261,66 @@ describe('Host Runtime Resource coordinator', () => {
     assert.equal(harness.terminateCount, 0);
   });
 
+  test('requests a canonical state drain only after leaving Session admission', async (t) => {
+    t.mock.method(console, 'error', () => {});
+    let drainAdmission: Promise<void> | undefined;
+    let harness!: ReturnType<typeof createHarness>;
+    harness = createHarness({
+      requestDrain: () => {
+        runAfterCurrentSessionAdmission(() => {
+          harness.drainCount += 1;
+          drainAdmission = harness.sessionAdmission.run(SESSION_ID, async () => {});
+          void drainAdmission.catch(() => {});
+        });
+      },
+    });
+    harness.stateReadFailure = new Error('canonical state unavailable');
+
+    const result = await harness.coordinator.handlers['runtime.resource.query'](
+      { kind: 'list_start', sessionId: SESSION_ID },
+      connection('connection-1'),
+    );
+
+    assert.equal(result.ok, false);
+    assert.equal(!result.ok && result.error.code, 'internal_failure');
+    assert.equal(harness.drainCount, 1);
+    assert.ok(drainAdmission, 'the canonical read failure requests a drain');
+    await assert.doesNotReject(drainAdmission);
+  });
+
+  test('logs a bounded redacted canonical state failure before draining', async (t) => {
+    const logs: string[] = [];
+    let drainCount = 0;
+    let logCountAtDrain = 0;
+    t.mock.method(console, 'error', (...args: unknown[]) => {
+      logs.push(args.map(String).join(' '));
+    });
+    const harness = createHarness({
+      requestDrain: () => {
+        drainCount += 1;
+        logCountAtDrain = logs.length;
+      },
+    });
+    harness.stateReadFailure = new Error(
+      `canonical state unavailable api_key=sk-secretvalue123 ${'x'.repeat(16 * 1024)}`,
+    );
+
+    const result = await harness.coordinator.handlers['runtime.resource.query'](
+      { kind: 'list_start', sessionId: SESSION_ID },
+      connection('connection-1'),
+    );
+
+    assert.equal(result.ok, false);
+    assert.equal(!result.ok && result.error.code, 'internal_failure');
+    assert.equal(drainCount, 1);
+    assert.equal(logCountAtDrain, 1);
+    assert.equal(logs.length, 1);
+    assert.match(logs[0] ?? '', /canonical state unavailable/);
+    assert.match(logs[0] ?? '', /\[redacted\]/i);
+    assert.doesNotMatch(logs[0] ?? '', /sk-secretvalue123/);
+    assert.ok(Buffer.byteLength(logs[0] ?? '', 'utf8') < 9 * 1024);
+  });
+
   test('fences PTY control by connection and retains only exact sequence retries', async () => {
     const harness = createHarness();
     const firstConnection = connection('connection-1');
@@ -399,6 +463,7 @@ describe('Host Runtime Resource coordinator', () => {
     assert.equal(started.ok, false);
     assert.ok(harness.lastBackgroundInput);
     assert.equal(harness.stopCount, 1);
+    assert.equal(harness.drainCount, 1);
     harness.finishBackground({ successful: false });
   });
 
@@ -702,12 +767,29 @@ describe('Host Runtime Resource coordinator', () => {
     assert.equal(!missing.ok && missing.error.code, 'not_found');
     assert.equal(harness.stopCount, 0);
   });
+
+  test('drains when the admitted mutable Session read fails', async () => {
+    const harness = createHarness();
+    harness.sessionReadFailureAt = 2;
+
+    const started = await harness.coordinator.handlers['runtime.resource.start'](
+      { sessionId: SESSION_ID, launchId: 'session-read-failure' },
+      connection('connection-1'),
+    );
+
+    assert.equal(started.ok, false);
+    assert.equal(!started.ok && started.error.code, 'internal_failure');
+    assert.equal(harness.drainCount, 1);
+    assert.equal(harness.lastBackgroundInput, undefined);
+  });
 });
 
 function createHarness(
-  options: Pick<
-    HostRuntimeResourceCoordinatorInput,
-    'resolveShell' | 'sessionAccessAuthority'
+  options: Partial<
+    Pick<
+      HostRuntimeResourceCoordinatorInput,
+      'requestDrain' | 'resolveShell' | 'sessionAccessAuthority'
+    >
   > = {},
 ) {
   let backgroundCompletion: ShellRunBashInput['onCompletion'];
@@ -716,6 +798,8 @@ function createHarness(
   const state = {
     updates: [resourceUpdate(0)],
     sessionState: 'active' as 'active' | 'archived' | 'missing',
+    sessionReadCount: 0,
+    sessionReadFailureAt: undefined as number | undefined,
     writeCount: 0,
     stopCount: 0,
     terminateCount: 0,
@@ -829,6 +913,10 @@ function createHarness(
     },
     sessionHeaders: {
       readHeader: async (sessionId) => {
+        state.sessionReadCount += 1;
+        if (state.sessionReadCount === state.sessionReadFailureAt) {
+          throw new Error('Session state unavailable');
+        }
         if (state.sessionState === 'missing') throw new SessionNotFoundError(sessionId);
         return {
           cwd: '/workspace',

packages/runtime-host/src/__tests__/session-admission-gate.test.ts

@@ -20,7 +20,11 @@
 import { deferred } from '@maka/core/test-only/async-primitives';
 import assert from 'node:assert/strict';
 import { test } from 'node:test';
-import { SessionAdmissionGate } from '../server/session-admission-gate.js';
+import {
+  runAfterCurrentSessionAdmission,
+  SessionAdmissionGate,
+  type SessionAdmissionLease,
+} from '../server/session-admission-gate.js';
 
 test('serializes operations for one Session', async () => {
   const gate = new SessionAdmissionGate();
@@ -148,6 +152,72 @@ test('rejects accidental admission re-entry instead of deadlocking', async () =>
   });
 });
 
+test('runs outside-admission work synchronously when no admission is active', () => {
+  const gate = new SessionAdmissionGate();
+  let ran = false;
+
+  runAfterCurrentSessionAdmission(() => {
+    ran = true;
+  });
+
+  assert.equal(ran, true);
+});
+
+test('runs outside-admission work after release and before the next queued admission', async () => {
+  const gate = new SessionAdmissionGate();
+  const entered = deferred();
+  const release = deferred();
+  const order: string[] = [];
+
+  const active = gate.run('session', async () => {
+    order.push('active:start');
+    runAfterCurrentSessionAdmission(() => {
+      order.push('after-release');
+    });
+    entered.resolve();
+    await release.promise;
+    order.push('active:end');
+  });
+  await entered.promise;
+  const queued = gate.run('session', () => {
+    order.push('queued');
+  });
+
+  assert.deepEqual(order, ['active:start']);
+  release.resolve();
+  await Promise.all([active, queued]);
+  assert.deepEqual(order, ['active:start', 'active:end', 'after-release', 'queued']);
+});
+
+test('tracks admitted work started outside the owning async chain until release', async () => {
+  const gate = new SessionAdmissionGate();
+  const leaseReady = deferred<SessionAdmissionLease>();
+  const release = deferred();
+  const order: string[] = [];
+
+  const active = gate.run('session', async (lease) => {
+    order.push('active:start');
+    leaseReady.resolve(lease);
+    await release.promise;
+    order.push('active:end');
+  });
+  const lease = await leaseReady.promise;
+  await gate.runAdmitted('session', lease, () => {
+    order.push('admitted');
+    runAfterCurrentSessionAdmission(() => {
+      order.push('after-release');
+    });
+  });
+  const queued = gate.run('session', () => {
+    order.push('queued');
+  });
+
+  assert.deepEqual(order, ['active:start', 'admitted']);
+  release.resolve();
+  await Promise.all([active, queued]);
+  assert.deepEqual(order, ['active:start', 'admitted', 'active:end', 'after-release', 'queued']);
+});
+
 test('work detached from an admission takes admissions of its own', async () => {
   const gate = new SessionAdmissionGate();
   const release = deferred();
@@ -172,3 +242,20 @@ test('work detached from an admission takes admissions of its own', async () =>
   await detached;
   assert.deepEqual(order, ['active:start', 'active:end', 'detached:admitted']);
 });
+
+test('treats detached work as outside the current admission', async () => {
+  const gate = new SessionAdmissionGate();
+  const order: string[] = [];
+
+  await gate.run('session', async () => {
+    order.push('active:start');
+    await gate.detach(async () => {
+      runAfterCurrentSessionAdmission(() => {
+        order.push('detached:outside');
+      });
+    });
+    order.push('active:end');
+  });
+
+  assert.deepEqual(order, ['active:start', 'detached:outside', 'active:end']);
+});

packages/runtime-host/src/server/failure-diagnostic.ts

@@ -0,0 +1,27 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { truncateUtf8 } from '@maka/core/diagnostic-log';
+import { redactSecrets } from '@maka/core/redaction';
+
+export function boundedFailureDiagnostic(error: unknown): string {
+  const details =
+    error instanceof Error ? error.stack || `${error.name}: ${error.message}` : String(error);
+  return truncateUtf8(redactSecrets(details), 8 * 1024, '\n<diagnostic truncated>');
+}

packages/runtime-host/src/server/host-kernel.ts

@@ -94,6 +94,7 @@ import {
 import { HostResidencyRegistry } from './host-residency-registry.js';
 import type { PeerMeshNode } from '../peer-mesh/node.js';
 import { createPeerMeshOperationHandlers } from './peer-mesh-authority.js';
+import { runAfterCurrentSessionAdmission } from './session-admission-gate.js';
 
 const DEFAULT_IDLE_GRACE_MS = 30_000;
 const DEFAULT_HANDSHAKE_TIMEOUT_MS = 5_000;
@@ -342,9 +343,11 @@ export class RuntimeHostKernel {
       this.#cancelIdle();
       this.#cancelInitialConnectionDeadline();
       this.#armShutdownDeadline();
-      this.#beginCompositionDrain();
     }
-    this.#commitRequestedShutdownIfQuiescent();
+    runAfterCurrentSessionAdmission(() => {
+      this.#beginCompositionDrain();
+      this.#commitRequestedShutdownIfQuiescent();
+    });
   }
 
   async #start(): Promise<void> {

packages/runtime-host/src/server/operation-dispatcher.ts

@@ -17,8 +17,6 @@
  * under the License.
  */
 
-import { truncateUtf8 } from '@maka/core/diagnostic-log';
-import { redactSecrets } from '@maka/core/redaction';
 import type { RootTurnAdmissionAuthorization } from '@maka/storage/execution-stores';
 import {
   HOST_OPERATION_SPECS,
@@ -72,6 +70,7 @@ import { USAGE_PRICING_OPERATION_SPECS } from '../protocol/usage-pricing.js';
 import { WEB_SEARCH_OPERATION_SPECS } from '../protocol/web-search.js';
 import { WORKHUB_COORDINATION_OPERATION_SPECS } from '../protocol/workhub-coordination.js';
 import { PLUGIN_PLATFORM_OPERATION_SPECS } from '../protocol/plugin-platform.js';
+import { boundedFailureDiagnostic } from './failure-diagnostic.js';
 import { createPeerMeshOperationHandlers } from './peer-mesh-authority.js';
 import type { RuntimeHostConnectionAuthority } from './connection-authority.js';
 
@@ -365,7 +364,7 @@ async function dispatchTypedOperation<K extends OperationKey>(
     outcome = decodeOperationOutcome(request.operation, await handler(request.input, context));
   } catch (error) {
     console.error(
-      `[runtime-host] unexpected ${request.operation} failure: ${boundedUnexpectedFailure(error)}`,
+      `[runtime-host] unexpected ${request.operation} failure: ${boundedFailureDiagnostic(error)}`,
     );
     return operationFailureResponse(
       request as RequestFrame,
@@ -387,9 +386,3 @@ async function dispatchTypedOperation<K extends OperationKey>(
         error: outcome.error,
       };
 }
-
-function boundedUnexpectedFailure(error: unknown): string {
-  const details =
-    error instanceof Error ? error.stack || `${error.name}: ${error.message}` : String(error);
-  return truncateUtf8(redactSecrets(details), 8 * 1024, '\n<diagnostic truncated>');
-}