@vellumai/credential-executor 0.8.4 → 0.8.6

package/src/managed-main.ts

@@ -6,11 +6,15 @@
  *
  * 1. Ensures the CES-private data directories exist.
  * 2. Binds a bootstrap Unix socket on the shared bootstrap volume.
- * 3. Accepts exactly **one** assistant runtime connection.
+ * 3. Accepts a single assistant runtime connection.
  * 4. Unlinks the socket path immediately after the connection is accepted,
- *    preventing any second process from connecting.
+ *    preventing any second process from connecting while the session is live.
  * 5. Serves RPC on the accepted stream only.
- * 6. Simultaneously serves health probes (`/healthz`, `/readyz`) on a
+ * 6. When that session ends (the assistant disconnects or its container is
+ *    restarted), re-binds the socket and awaits a reconnection. CES is a
+ *    long-lived sidecar — it outlives any single assistant session and only
+ *    shuts down on SIGTERM/SIGINT. At most one connection is ever active.
+ * 7. Simultaneously serves health probes (`/healthz`, `/readyz`) on a
  *    dedicated HTTP port for Kubernetes liveness/readiness checks.
  *
  * The managed entrypoint never opens a generic TCP or HTTP command API.
@@ -22,7 +26,10 @@ import { createServer as createNetServer, type Socket } from "node:net";
 import { dirname, join } from "node:path";
 import { Readable, Writable } from "node:stream";
 
-import { CES_PROTOCOL_VERSION, CesRpcMethod } from "@vellumai/service-contracts/credential-rpc";
+import {
+  CES_PROTOCOL_VERSION,
+  CesRpcMethod,
+} from "@vellumai/service-contracts/credential-rpc";
 
 import { AuditStore } from "./audit/store.js";
 import { PersistentGrantStore } from "./grants/persistent-store.js";
@@ -49,19 +56,36 @@ import {
   registerCommandExecutionHandler,
   registerManageSecureCommandToolHandler,
   type RpcHandlerRegistry,
+  type ServeEndReason,
   type SessionIdRef,
 } from "./server.js";
-import { deleteBundleFromToolstore, publishBundle } from "./toolstore/publish.js";
+import {
+  deleteBundleFromToolstore,
+  publishBundle,
+} from "./toolstore/publish.js";
 import { validateSourceUrl } from "./toolstore/manifest.js";
 import { buildCesEgressHooks } from "./commands/egress-hooks.js";
 import { resolveManagedSubject } from "./subjects/managed.js";
 import { materializeManagedToken } from "./materializers/managed-platform.js";
-import { HandleType, parseHandle } from "@vellumai/service-contracts/credential-rpc";
-import { buildLazyGetters, type ApiKeyRef, type AssistantIdRef } from "./managed-lazy-getters.js";
+import {
+  HandleType,
+  parseHandle,
+} from "@vellumai/service-contracts/credential-rpc";
+import {
+  applyManagedCredentialRefs,
+  buildLazyGetters,
+  type ApiKeyRef,
+  type AssistantIdRef,
+} from "./managed-lazy-getters.js";
 import { MANAGED_LOCAL_STATIC_REJECTION_ERROR } from "./managed-errors.js";
 import type { SecureKeyBackend } from "@vellumai/credential-storage";
 import { createLocalSecureKeyBackend } from "./materializers/local-secure-key-backend.js";
-import { handleCredentialRoute, type CredentialRouteDeps } from "./http/credential-routes.js";
+import type { LocalMaterialiser } from "./materializers/local.js";
+import type { LocalSubjectResolverDeps } from "./subjects/local.js";
+import {
+  handleCredentialRoute,
+  type CredentialRouteDeps,
+} from "./http/credential-routes.js";
 import { handleLogExportRoute } from "./http/log-export-routes.js";
 import { CES_MIGRATIONS } from "./migrations/registry.js";
 import { runCesMigrations } from "./migrations/runner.js";
@@ -95,7 +119,12 @@ function ensureDataDirs(): void {
 // Build RPC handler registry (managed mode)
 // ---------------------------------------------------------------------------
 
-function buildHandlers(sessionIdRef: SessionIdRef, apiKeyRef: ApiKeyRef, assistantIdRef: AssistantIdRef, secureKeyBackend: SecureKeyBackend): RpcHandlerRegistry {
+function buildHandlers(
+  sessionIdRef: SessionIdRef,
+  apiKeyRef: ApiKeyRef,
+  assistantIdRef: AssistantIdRef,
+  secureKeyBackend: SecureKeyBackend,
+): { handlers: RpcHandlerRegistry; temporaryGrantStore: TemporaryGrantStore } {
   // -- Grant stores ----------------------------------------------------------
   const persistentGrantStore = new PersistentGrantStore(
     getCesGrantsDir("managed"),
@@ -117,7 +146,7 @@ function buildHandlers(sessionIdRef: SessionIdRef, apiKeyRef: ApiKeyRef, assista
   // though handlers are built before the handshake completes.
   const platformBaseUrl = process.env["VELLUM_PLATFORM_URL"] ?? "";
 
-  const { getAssistantApiKey, getManagedSubjectOptions, getManagedMaterializerOptions } =
+  const { getManagedSubjectOptions, getManagedMaterializerOptions } =
     buildLazyGetters({
       platformBaseUrl,
       assistantIdRef,
@@ -135,11 +164,13 @@ function buildHandlers(sessionIdRef: SessionIdRef, apiKeyRef: ApiKeyRef, assista
   // -- Workspace root for command execution cwd ------------------------------
   // Use VELLUM_WORKSPACE_DIR when set, otherwise fall back to the legacy
   // path derived from the assistant data mount.
-  const defaultWorkspaceDir = process.env["VELLUM_WORKSPACE_DIR"] ?? (() => {
-    const assistantDataMount =
-      process.env["CES_ASSISTANT_DATA_MOUNT"] ?? "/assistant-data-ro";
-    return join(join(assistantDataMount, ".vellum"), "workspace");
-  })();
+  const defaultWorkspaceDir =
+    process.env["VELLUM_WORKSPACE_DIR"] ??
+    (() => {
+      const assistantDataMount =
+        process.env["CES_ASSISTANT_DATA_MOUNT"] ?? "/assistant-data-ro";
+      return join(join(assistantDataMount, ".vellum"), "workspace");
+    })();
 
   // -- Build handler registry ------------------------------------------------
   // NOTE: local_static credential handles are NOT supported in managed mode.
@@ -161,8 +192,11 @@ function buildHandlers(sessionIdRef: SessionIdRef, apiKeyRef: ApiKeyRef, assista
     }),
   };
 
-  const localSubjectDepsStub = {
-    metadataStore: { getById: () => undefined, list: () => [] } as any,
+  const localSubjectDepsStub: LocalSubjectResolverDeps = {
+    metadataStore: {
+      getById: () => undefined,
+      list: () => [],
+    } as unknown as LocalSubjectResolverDeps["metadataStore"],
     oauthConnections: { getById: () => undefined },
   };
 
@@ -171,10 +205,14 @@ function buildHandlers(sessionIdRef: SessionIdRef, apiKeyRef: ApiKeyRef, assista
   const httpDeps = {
     persistentGrantStore,
     temporaryGrantStore,
-    localMaterialiser: localMaterialiserStub as any,
+    localMaterialiser: localMaterialiserStub as unknown as LocalMaterialiser,
     localSubjectDeps: localSubjectDepsStub,
-    get managedSubjectOptions() { return getManagedSubjectOptions(); },
-    get managedMaterializerOptions() { return getManagedMaterializerOptions(); },
+    get managedSubjectOptions() {
+      return getManagedSubjectOptions();
+    },
+    get managedMaterializerOptions() {
+      return getManagedMaterializerOptions();
+    },
     auditStore,
     sessionId: sessionIdRef,
   };
@@ -215,10 +253,7 @@ function buildHandlers(sessionIdRef: SessionIdRef, apiKeyRef: ApiKeyRef, assista
               };
             }
 
-            const subjectResult = await resolveManagedSubject(
-              handle,
-              subOpts,
-            );
+            const subjectResult = await resolveManagedSubject(handle, subOpts);
             if (!subjectResult.ok) {
               return { ok: false as const, error: subjectResult.error.message };
             }
@@ -241,7 +276,8 @@ function buildHandlers(sessionIdRef: SessionIdRef, apiKeyRef: ApiKeyRef, assista
           default:
             return {
               ok: false as const,
-              error: `Handle type "${parseResult.handle.type}" is not supported in managed mode. ` +
+              error:
+                `Handle type "${parseResult.handle.type}" is not supported in managed mode. ` +
                 `Supported types: platform_oauth.`,
             };
         }
@@ -255,7 +291,15 @@ function buildHandlers(sessionIdRef: SessionIdRef, apiKeyRef: ApiKeyRef, assista
   });
 
   // Register manage_secure_command_tool handler
-  const toolRegistry = new Map<string, { toolName: string; credentialHandle: string; description: string; bundleDigest: string }>();
+  const toolRegistry = new Map<
+    string,
+    {
+      toolName: string;
+      credentialHandle: string;
+      description: string;
+      bundleDigest: string;
+    }
+  >();
 
   registerManageSecureCommandToolHandler(handlers, {
     downloadBundle: async (sourceUrl: string) => {
@@ -264,13 +308,17 @@ function buildHandlers(sessionIdRef: SessionIdRef, apiKeyRef: ApiKeyRef, assista
         throw new Error(urlError);
       }
       const MAX_BUNDLE_SIZE = 100 * 1024 * 1024; // 100 MB
-      const resp = await fetch(sourceUrl, { signal: AbortSignal.timeout(60_000) });
+      const resp = await fetch(sourceUrl, {
+        signal: AbortSignal.timeout(60_000),
+      });
       if (!resp.ok) {
         throw new Error(`HTTP ${resp.status}: ${resp.statusText}`);
       }
       const contentLength = resp.headers.get("content-length");
       if (contentLength && parseInt(contentLength, 10) > MAX_BUNDLE_SIZE) {
-        throw new Error(`Bundle too large: ${contentLength} bytes (max ${MAX_BUNDLE_SIZE})`);
+        throw new Error(
+          `Bundle too large: ${contentLength} bytes (max ${MAX_BUNDLE_SIZE})`,
+        );
       }
       // Stream the body and enforce the size limit on actual bytes received,
       // since Content-Length can be absent (chunked encoding) or lie.
@@ -283,18 +331,23 @@ function buildHandlers(sessionIdRef: SessionIdRef, apiKeyRef: ApiKeyRef, assista
       for await (const chunk of body) {
         totalBytes += chunk.byteLength;
         if (totalBytes > MAX_BUNDLE_SIZE) {
-          throw new Error(`Bundle too large: received >${MAX_BUNDLE_SIZE} bytes (max ${MAX_BUNDLE_SIZE})`);
+          throw new Error(
+            `Bundle too large: received >${MAX_BUNDLE_SIZE} bytes (max ${MAX_BUNDLE_SIZE})`,
+          );
         }
         chunks.push(chunk);
       }
       return Buffer.concat(chunks);
     },
-    publishBundle: (request) => publishBundle({ ...request, cesMode: "managed" }),
+    publishBundle: (request) =>
+      publishBundle({ ...request, cesMode: "managed" }),
     unregisterTool: (toolName: string) => {
       const entry = toolRegistry.get(toolName);
       const removed = toolRegistry.delete(toolName);
       if (removed && entry?.bundleDigest) {
-        const stillInUse = Array.from(toolRegistry.values()).some(t => t.bundleDigest === entry.bundleDigest);
+        const stillInUse = Array.from(toolRegistry.values()).some(
+          (t) => t.bundleDigest === entry.bundleDigest,
+        );
         if (!stillInUse) {
           deleteBundleFromToolstore(entry.bundleDigest, "managed");
         }
@@ -310,52 +363,59 @@ function buildHandlers(sessionIdRef: SessionIdRef, apiKeyRef: ApiKeyRef, assista
   handlers[CesRpcMethod.RecordGrant] = createRecordGrantHandler({
     persistentGrantStore,
     temporaryGrantStore,
-  }) as typeof handlers[string];
+  }) as (typeof handlers)[string];
 
   handlers[CesRpcMethod.ListGrants] = createListGrantsHandler({
     persistentGrantStore,
-  }) as typeof handlers[string];
+  }) as (typeof handlers)[string];
 
   handlers[CesRpcMethod.RevokeGrant] = createRevokeGrantHandler({
     persistentGrantStore,
-  }) as typeof handlers[string];
+  }) as (typeof handlers)[string];
 
   // Register audit record handler
   handlers[CesRpcMethod.ListAuditRecords] = createListAuditRecordsHandler({
     auditStore,
-  }) as typeof handlers[string];
+  }) as (typeof handlers)[string];
 
   // Register credential CRUD handlers
   handlers[CesRpcMethod.GetCredential] = (async (req: { account: string }) => {
     const value = await secureKeyBackend.get(req.account);
     return { found: value !== undefined, value };
-  }) as typeof handlers[string];
+  }) as (typeof handlers)[string];
 
-  handlers[CesRpcMethod.SetCredential] = (async (req: { account: string; value: string }) => {
+  handlers[CesRpcMethod.SetCredential] = (async (req: {
+    account: string;
+    value: string;
+  }) => {
     const ok = await secureKeyBackend.set(req.account, req.value);
     return { ok };
-  }) as typeof handlers[string];
+  }) as (typeof handlers)[string];
 
-  handlers[CesRpcMethod.DeleteCredential] = (async (req: { account: string }) => {
+  handlers[CesRpcMethod.DeleteCredential] = (async (req: {
+    account: string;
+  }) => {
     const result = await secureKeyBackend.delete(req.account);
     return { result };
-  }) as typeof handlers[string];
+  }) as (typeof handlers)[string];
 
   handlers[CesRpcMethod.ListCredentials] = (async () => {
     const accounts = await secureKeyBackend.list();
     return { accounts };
-  }) as typeof handlers[string];
+  }) as (typeof handlers)[string];
 
-  handlers[CesRpcMethod.BulkSetCredentials] = (async (req: { credentials: Array<{ account: string; value: string }> }) => {
+  handlers[CesRpcMethod.BulkSetCredentials] = (async (req: {
+    credentials: Array<{ account: string; value: string }>;
+  }) => {
     const results = [];
     for (const { account, value } of req.credentials) {
       const ok = await secureKeyBackend.set(account, value);
       results.push({ account, ok });
     }
     return { results };
-  }) as typeof handlers[string];
+  }) as (typeof handlers)[string];
 
-  return handlers;
+  return { handlers, temporaryGrantStore };
 }
 
 // ---------------------------------------------------------------------------
@@ -374,10 +434,9 @@ function startHealthServer(
     async fetch(req) {
       const url = new URL(req.url);
       if (url.pathname === "/healthz") {
-        return new Response(
-          JSON.stringify({ status: "ok" }),
-          { headers: { "Content-Type": "application/json" } },
-        );
+        return new Response(JSON.stringify({ status: "ok" }), {
+          headers: { "Content-Type": "application/json" },
+        });
       }
       if (url.pathname === "/readyz") {
         // Always return 200 — pod readiness must not depend on whether the
@@ -386,44 +445,56 @@ function startHealthServer(
         // scheduling during dark-launch.  The sidecar can't do useful work
         // without a connection anyway, so readiness is purely about the
         // process being up and able to accept a future connection.
-        return new Response(
-          JSON.stringify({ status: "ok", rpcConnected }),
-          {
-            status: 200,
-            headers: { "Content-Type": "application/json" },
-          },
-        );
+        return new Response(JSON.stringify({ status: "ok", rpcConnected }), {
+          status: 200,
+          headers: { "Content-Type": "application/json" },
+        });
       }
 
       // Credential CRUD routes (only if service token is configured)
       if (credentialDeps) {
-        const credentialResponse = await handleCredentialRoute(req, credentialDeps);
+        const credentialResponse = await handleCredentialRoute(
+          req,
+          credentialDeps,
+        );
         if (credentialResponse) return credentialResponse;
       }
 
       // Log export route
-      const logExportResponse = await handleLogExportRoute(req, getCesLogDir("managed"));
+      const logExportResponse = await handleLogExportRoute(
+        req,
+        getCesLogDir("managed"),
+      );
       if (logExportResponse) return logExportResponse;
 
       return new Response("Not Found", { status: 404 });
     },
   });
 
-  signal.addEventListener("abort", () => {
-    server.stop(true);
-  }, { once: true });
+  signal.addEventListener(
+    "abort",
+    () => {
+      server.stop(true);
+    },
+    { once: true },
+  );
 
   return server;
 }
 
 // ---------------------------------------------------------------------------
-// Bootstrap socket server (accepts exactly one connection)
+// Bootstrap socket server (accepts one connection at a time)
 // ---------------------------------------------------------------------------
 
 /**
- * Listen on a Unix socket, accept exactly one connection, unlink the
- * socket path, and return readable/writable streams for the accepted
- * connection.
+ * Listen on a Unix socket, accept one connection, unlink the socket path,
+ * and return readable/writable streams for the accepted connection.
+ *
+ * The socket is unlinked while a connection is active so no second process
+ * can connect concurrently (only one assistant ever talks to CES at a time).
+ * When that session ends, the caller re-invokes this function to re-bind the
+ * socket and accept the assistant's reconnection — CES outlives any single
+ * assistant session (see `main()`).
  */
 function acceptOneConnection(
   socketPath: string,
@@ -456,12 +527,17 @@ function acceptOneConnection(
       return;
     }
 
-    signal.addEventListener("abort", () => {
+    // Remove this listener once the promise settles. Because CES re-binds
+    // the socket after each session ends, a long-lived AbortSignal would
+    // otherwise accumulate one dangling listener per reconnection.
+    const onAbort = () => {
       cleanup();
       reject(new Error("Aborted while waiting for connection"));
-    }, { once: true });
+    };
+    signal.addEventListener("abort", onAbort, { once: true });
 
     netServer.on("error", (err) => {
+      signal.removeEventListener("abort", onAbort);
       cleanup();
       reject(err);
     });
@@ -471,8 +547,10 @@ function acceptOneConnection(
     });
 
     netServer.on("connection", (socket: Socket) => {
-      // Accept exactly one connection, then close the listener and
-      // unlink the socket path so no other process can connect.
+      // Accept the connection, then close the listener and unlink the
+      // socket path so no other process can connect while this session
+      // is active.
+      signal.removeEventListener("abort", onAbort);
       log.info("Assistant connected via bootstrap socket");
       netServer.close();
       try {
@@ -480,7 +558,7 @@ function acceptOneConnection(
       } catch {
         // Already unlinked
       }
-      log.info("Bootstrap socket unlinked (single-connection enforced)");
+      log.info("Bootstrap socket unlinked (single active connection enforced)");
 
       const readable = new Readable({
         read() {
@@ -529,13 +607,22 @@ async function main(): Promise<void> {
 
   const controller = new AbortController();
 
-  // Graceful shutdown
-  const shutdown = () => {
-    log.info("Shutting down...");
-    controller.abort();
-  };
-  process.on("SIGTERM", shutdown);
-  process.on("SIGINT", shutdown);
+  // Graceful shutdown — pass the signal as the abort reason so consumers
+  // of controller.signal can inspect signal.reason for triage.
+  process.on("SIGTERM", () => {
+    log.warn(
+      { signal: "SIGTERM", pid: process.pid, uptime: process.uptime() },
+      "Received SIGTERM — shutting down",
+    );
+    controller.abort("SIGTERM");
+  });
+  process.on("SIGINT", () => {
+    log.warn(
+      { signal: "SIGINT", pid: process.pid, uptime: process.uptime() },
+      "Received SIGINT — shutting down",
+    );
+    controller.abort("SIGINT");
+  });
 
   // Create the secure key backend unconditionally — it's needed by both
   // HTTP credential routes (when CES_SERVICE_TOKEN is set) and RPC
@@ -546,7 +633,11 @@ async function main(): Promise<void> {
   const secureKeyBackend = createLocalSecureKeyBackend(vellumRoot);
 
   // Run one-time credential store migrations before accepting connections.
-  await runCesMigrations(getCesDataRoot("managed"), secureKeyBackend, CES_MIGRATIONS);
+  await runCesMigrations(
+    getCesDataRoot("managed"),
+    secureKeyBackend,
+    CES_MIGRATIONS,
+  );
   log.info("CES managed startup: migrations complete");
 
   // Set up credential CRUD routes if a service token is configured.
@@ -565,73 +656,164 @@ async function main(): Promise<void> {
     );
   }
 
-  // Start health server on dedicated port
+  // Start health server on dedicated port. The returned handle isn't
+  // needed because the server lifetime is bound to controller.signal,
+  // which fires on shutdown and triggers Bun.serve's stop().
   const healthPort = getHealthPort();
-  const healthServer = startHealthServer(healthPort, controller.signal, credentialDeps);
+  startHealthServer(healthPort, controller.signal, credentialDeps);
   log.info(`Health server listening on port ${healthPort}`);
 
-  // Wait for exactly one assistant connection on the bootstrap socket
-  const socketPath = getBootstrapSocketPath();
-  log.info(`Waiting for assistant connection on ${socketPath}...`);
-
-  let connection: Awaited<ReturnType<typeof acceptOneConnection>>;
-  try {
-    connection = await acceptOneConnection(socketPath, controller.signal);
-  } catch (err) {
-    if (controller.signal.aborted) {
-      log.info("Shutdown before assistant connected.");
-      return;
-    }
-    throw err;
-  }
-
-  rpcConnected = true;
-
-  // Build the handler registry with all available RPC implementations.
-  // Use mutable refs so the handshake-provided session ID and API key
-  // are available to handlers at call time (after the handshake completes).
+  // Build the handler registry once, up front, and reuse it across every
+  // assistant session. All CES state lives behind these handlers — file-backed
+  // grant/audit stores plus the in-memory temporary-grant store and the
+  // secure-command tool registry — and must be process-scoped so it survives
+  // an assistant reconnection. In particular, the tool registry mirrors the
+  // persistent toolstore on disk; rebuilding it per session would let a later
+  // `unregister` miss a tool registered in an earlier session and orphan its
+  // bundle.
+  //
+  // The in-memory temporary-grant store is the exception: `allow_once` /
+  // `allow_10m` grants are keyed by proposal hash only (not session), so they
+  // would otherwise leak ephemeral approvals across sessions. It is cleared at
+  // the end of every session below so a reconnecting assistant must re-prompt.
+  //
+  // The mutable refs carry the handshake-provided session ID, API key, and
+  // assistant ID; handlers read them at call time, so updating the refs when
+  // each session's handshake completes is all that's needed per connection.
   const sessionIdRef: SessionIdRef = { current: `ces-managed-${Date.now()}` };
   const apiKeyRef: ApiKeyRef = { current: "" };
   const assistantIdRef: AssistantIdRef = { current: "" };
-  const handlers = buildHandlers(sessionIdRef, apiKeyRef, assistantIdRef, secureKeyBackend);
+  const { handlers, temporaryGrantStore } = buildHandlers(
+    sessionIdRef,
+    apiKeyRef,
+    assistantIdRef,
+    secureKeyBackend,
+  );
 
+  // Serve loop. CES is a long-lived sidecar that must outlive any single
+  // assistant session: the assistant container can crash and be restarted
+  // independently of the CES container (Kubernetes restarts containers, not
+  // the whole pod), so when the RPC stream ends we re-bind the bootstrap
+  // socket and wait for the assistant to reconnect rather than tearing the
+  // sidecar down. The loop only exits on a shutdown signal (SIGTERM/SIGINT),
+  // which aborts the controller.
   const rpcLog = getLogger("rpc");
-  const server = new CesRpcServer({
-    input: connection.readable,
-    output: connection.writable,
-    handlers,
-    logger: {
-      log: (msg: string, ...args: unknown[]) => rpcLog.info({ args }, msg),
-      warn: (msg: string, ...args: unknown[]) => rpcLog.warn({ args }, msg),
-      error: (msg: string, ...args: unknown[]) => rpcLog.error({ args }, msg),
-    },
-    signal: controller.signal,
-    onHandshakeComplete: (hsSessionId, hsApiKey, hsAssistantId) => {
-      sessionIdRef.current = hsSessionId;
-      if (hsApiKey) {
-        apiKeyRef.current = hsApiKey;
-        log.info("Received assistant API key via handshake");
-      }
-      if (hsAssistantId) {
-        assistantIdRef.current = hsAssistantId;
-        log.info("Received assistant ID via handshake");
-      }
-    },
-    onApiKeyUpdate: (newKey, newAssistantId) => {
-      apiKeyRef.current = newKey;
-      log.info("Assistant API key updated via RPC");
-      if (newAssistantId) {
-        assistantIdRef.current = newAssistantId;
-        log.info("Assistant ID updated via RPC");
+  const socketPath = getBootstrapSocketPath();
+
+  while (!controller.signal.aborted) {
+    log.info(`Waiting for assistant connection on ${socketPath}...`);
+
+    let connection: Awaited<ReturnType<typeof acceptOneConnection>>;
+    try {
+      connection = await acceptOneConnection(socketPath, controller.signal);
+    } catch (err) {
+      if (controller.signal.aborted) {
+        log.info("Shutdown before assistant connected.");
+        return;
       }
-    },
-  });
+      throw err;
+    }
+
+    rpcConnected = true;
 
-  await server.serve();
+    const server = new CesRpcServer({
+      input: connection.readable,
+      output: connection.writable,
+      handlers,
+      logger: {
+        log: (msg: string, ...args: unknown[]) => rpcLog.info({ args }, msg),
+        warn: (msg: string, ...args: unknown[]) => rpcLog.warn({ args }, msg),
+        error: (msg: string, ...args: unknown[]) => rpcLog.error({ args }, msg),
+      },
+      signal: controller.signal,
+      onHandshakeComplete: (hsSessionId, hsApiKey, hsAssistantId) => {
+        sessionIdRef.current = hsSessionId;
+        // Overwrite the credential refs on every handshake. The handler
+        // registry persists across reconnects, so a new session that omits
+        // the API key / assistant ID must fail closed (falling back to the
+        // env key, or no key) rather than reusing the previous session's
+        // credentials.
+        applyManagedCredentialRefs(
+          apiKeyRef,
+          assistantIdRef,
+          hsApiKey,
+          hsAssistantId,
+        );
+        if (hsApiKey) {
+          log.info("Received assistant API key via handshake");
+        }
+        if (hsAssistantId) {
+          log.info("Received assistant ID via handshake");
+        }
+      },
+      onApiKeyUpdate: (newKey, newAssistantId) => {
+        // Overwrite both refs on every credential update, for the same
+        // fail-closed reason as the handshake: the assistant sources the
+        // assistant ID from the same place it sources the key, so an update
+        // that omits the ID means it has none — CES must clear the stale ID
+        // rather than keep materializing for the previous session's assistant.
+        applyManagedCredentialRefs(
+          apiKeyRef,
+          assistantIdRef,
+          newKey,
+          newAssistantId,
+        );
+        log.info("Assistant API key updated via RPC");
+        if (newAssistantId) {
+          log.info("Assistant ID updated via RPC");
+        }
+      },
+    });
 
-  rpcConnected = false;
-  log.info("RPC session ended. Shutting down...");
-  controller.abort();
+    // `serve()` resolves on a clean stream end or signal abort, and rejects
+    // when the transport stream errors — which is precisely what a hard
+    // disconnect (connection reset when the assistant container crashes)
+    // looks like. Both cases must keep the sidecar up; only a shutdown
+    // signal should tear it down. So treat a serve() rejection the same as
+    // a session end and fall through to await reconnection.
+    let endReason: ServeEndReason | "transport_error";
+    try {
+      endReason = await server.serve();
+    } catch (err) {
+      server.close();
+      endReason = "transport_error";
+      log.warn(
+        { err, uptime: process.uptime(), pid: process.pid },
+        "RPC transport errored — treating as session end",
+      );
+    }
+
+    rpcConnected = false;
+
+    // Drop all ephemeral approvals when the session ends. `allow_once` /
+    // `allow_10m` grants are keyed by proposal hash only, so reusing the
+    // store across a reconnect would let a pre-disconnect approval be
+    // consumed by a later session without re-prompting. Clearing here
+    // restores the prior behavior, where the process exited on stream end
+    // and these grants never survived.
+    temporaryGrantStore.clear();
+
+    // A signal-driven end means the process is shutting down; exit the loop.
+    // Any other end reason (the assistant disconnected, its stream closed,
+    // or the transport errored) means we keep the sidecar up and await a
+    // reconnection.
+    if (
+      controller.signal.aborted ||
+      endReason === "signal_aborted" ||
+      endReason === "signal_aborted_before_start"
+    ) {
+      log.info(
+        { reason: endReason, uptime: process.uptime(), pid: process.pid },
+        "RPC session ended due to shutdown — exiting serve loop",
+      );
+      break;
+    }
+
+    log.warn(
+      { reason: endReason, uptime: process.uptime(), pid: process.pid },
+      "RPC session ended (assistant disconnected) — awaiting reconnection",
+    );
+  }
 }
 
 main().catch((err) => {