@vellumai/credential-executor 0.8.5 → 0.8.7

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@vellumai/credential-executor",
-  "version": "0.8.5",
+  "version": "0.8.7",
   "license": "MIT",
   "type": "module",
   "exports": {

package/src/__tests__/managed-lazy-getters.test.ts

@@ -9,11 +9,76 @@
 import { describe, expect, test } from "bun:test";
 
 import {
+  applyManagedCredentialRefs,
   buildLazyGetters,
   type ApiKeyRef,
   type AssistantIdRef,
 } from "../managed-lazy-getters.js";
 
+// ---------------------------------------------------------------------------
+// applyManagedCredentialRefs — fail-closed overwrite across sessions
+// ---------------------------------------------------------------------------
+
+describe("applyManagedCredentialRefs", () => {
+  test("overwrites both refs with the provided values", () => {
+    const apiKeyRef: ApiKeyRef = { current: "old-key" };
+    const assistantIdRef: AssistantIdRef = { current: "ast_old" };
+
+    applyManagedCredentialRefs(apiKeyRef, assistantIdRef, "new-key", "ast_new");
+
+    expect(apiKeyRef.current).toBe("new-key");
+    expect(assistantIdRef.current).toBe("ast_new");
+  });
+
+  test("clears a stale assistant ID when the new value is omitted", () => {
+    // A new session (or an API-key-only update) that does not carry an
+    // assistant ID must not inherit the previous session's ID.
+    const apiKeyRef: ApiKeyRef = { current: "prev-key" };
+    const assistantIdRef: AssistantIdRef = { current: "ast_prev" };
+
+    applyManagedCredentialRefs(
+      apiKeyRef,
+      assistantIdRef,
+      "rotated-key",
+      undefined,
+    );
+
+    expect(apiKeyRef.current).toBe("rotated-key");
+    expect(assistantIdRef.current).toBe("");
+  });
+
+  test("clears a stale API key when the new value is omitted", () => {
+    const apiKeyRef: ApiKeyRef = { current: "prev-key" };
+    const assistantIdRef: AssistantIdRef = { current: "ast_prev" };
+
+    applyManagedCredentialRefs(apiKeyRef, assistantIdRef, undefined, undefined);
+
+    expect(apiKeyRef.current).toBe("");
+    expect(assistantIdRef.current).toBe("");
+  });
+
+  test("lazy getters fail closed after a reconnect that omits the ID", () => {
+    const apiKeyRef: ApiKeyRef = { current: "session1-key" };
+    const assistantIdRef: AssistantIdRef = { current: "ast_session1" };
+    const { getManagedMaterializerOptions } = buildLazyGetters({
+      platformBaseUrl: "https://api.vellum.ai",
+      assistantIdRef,
+      apiKeyRef,
+    });
+
+    // A reconnecting session provides a key but no assistant ID.
+    applyManagedCredentialRefs(
+      apiKeyRef,
+      assistantIdRef,
+      "session2-key",
+      undefined,
+    );
+
+    // Materialization must not proceed with the prior session's assistant ID.
+    expect(getManagedMaterializerOptions()).toBeUndefined();
+  });
+});
+
 // ---------------------------------------------------------------------------
 // Before API key arrives
 // ---------------------------------------------------------------------------

package/src/__tests__/managed-reconnect.test.ts

@@ -0,0 +1,244 @@
+/**
+ * Managed CES reconnection test (real entrypoint subprocess).
+ *
+ * Spawns the actual `managed-main.ts` entrypoint and verifies that the CES
+ * sidecar survives the assistant disconnecting and accepts a reconnection,
+ * rather than shutting down when the RPC stream ends.
+ *
+ * This guards the core invariant that CES runs independently of whether the
+ * assistant is actively connected: the assistant container can crash and be
+ * restarted (Kubernetes restarts containers, not the whole pod), and the
+ * restarted assistant must be able to reconnect to a still-running CES.
+ */
+
+import { afterEach, describe, expect, test } from "bun:test";
+import { createConnection, createServer, type Socket } from "node:net";
+import { existsSync, mkdirSync, mkdtempSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join, resolve } from "node:path";
+
+import {
+  CES_PROTOCOL_VERSION,
+  type HandshakeAck,
+} from "@vellumai/service-contracts/credential-rpc";
+
+import type { Subprocess } from "bun";
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+/** Sleep for the given number of milliseconds. */
+function delay(ms: number): Promise<void> {
+  return new Promise((r) => setTimeout(r, ms));
+}
+
+/** Pick a currently-free TCP port by binding to port 0 and reading it back. */
+function pickFreePort(): Promise<number> {
+  return new Promise((resolve, reject) => {
+    const srv = createServer();
+    srv.listen(0, () => {
+      const address = srv.address();
+      const port = typeof address === "object" && address ? address.port : 0;
+      srv.close(() => (port ? resolve(port) : reject(new Error("no port"))));
+    });
+    srv.on("error", reject);
+  });
+}
+
+/** Poll the health endpoint until it responds OK or the deadline passes. */
+async function waitForHealth(port: number, timeoutMs = 15_000): Promise<void> {
+  const deadline = Date.now() + timeoutMs;
+  while (Date.now() < deadline) {
+    try {
+      const resp = await fetch(`http://127.0.0.1:${port}/healthz`);
+      if (resp.ok) return;
+    } catch {
+      // not up yet
+    }
+    await delay(100);
+  }
+  throw new Error(`CES health endpoint did not come up within ${timeoutMs}ms`);
+}
+
+/** Read the `rpcConnected` field from /readyz. */
+async function readyzRpcConnected(port: number): Promise<boolean> {
+  const resp = await fetch(`http://127.0.0.1:${port}/readyz`);
+  const body = (await resp.json()) as { rpcConnected?: boolean };
+  return body.rpcConnected === true;
+}
+
+/** Wait until the socket path exists (CES has bound the bootstrap socket). */
+async function waitForSocket(socketPath: string, timeoutMs = 10_000): Promise<void> {
+  const deadline = Date.now() + timeoutMs;
+  while (Date.now() < deadline) {
+    if (existsSync(socketPath)) return;
+    await delay(50);
+  }
+  throw new Error(`Bootstrap socket did not appear within ${timeoutMs}ms`);
+}
+
+/** Connect to the bootstrap socket, retrying past the listen/connect race. */
+function connectToSocket(
+  socketPath: string,
+  { maxRetries = 40, baseDelayMs = 25 } = {},
+): Promise<Socket> {
+  return new Promise((resolveConn, reject) => {
+    let attempt = 0;
+    const tryConnect = () => {
+      const sock = createConnection(socketPath, () => {
+        sock.removeAllListeners("error");
+        resolveConn(sock);
+      });
+      sock.on("error", (err: NodeJS.ErrnoException) => {
+        sock.destroy();
+        attempt++;
+        if (
+          attempt < maxRetries &&
+          (err.code === "ENOENT" || err.code === "ECONNREFUSED")
+        ) {
+          setTimeout(tryConnect, baseDelayMs);
+        } else {
+          reject(err);
+        }
+      });
+    };
+    tryConnect();
+  });
+}
+
+/** Send a handshake and resolve the resulting ack. */
+function handshake(sock: Socket, sessionId: string): Promise<HandshakeAck> {
+  return new Promise((resolveAck, reject) => {
+    let buffer = "";
+    const timer = setTimeout(() => {
+      sock.removeAllListeners("data");
+      reject(new Error("Timed out waiting for handshake ack"));
+    }, 5_000);
+
+    sock.on("data", (chunk: Buffer) => {
+      buffer += chunk.toString("utf-8");
+      const idx = buffer.indexOf("\n");
+      if (idx === -1) return;
+      const line = buffer.slice(0, idx).trim();
+      clearTimeout(timer);
+      sock.removeAllListeners("data");
+      try {
+        resolveAck(JSON.parse(line) as HandshakeAck);
+      } catch (err) {
+        reject(err as Error);
+      }
+    });
+    sock.on("error", (err) => {
+      clearTimeout(timer);
+      reject(err);
+    });
+
+    sock.write(
+      JSON.stringify({
+        type: "handshake_request",
+        protocolVersion: CES_PROTOCOL_VERSION,
+        sessionId,
+      }) + "\n",
+    );
+  });
+}
+
+// ---------------------------------------------------------------------------
+// Test
+// ---------------------------------------------------------------------------
+
+let tmpDir: string | undefined;
+let proc: Subprocess | undefined;
+
+afterEach(async () => {
+  if (proc) {
+    proc.kill("SIGTERM");
+    await Promise.race([proc.exited, delay(3_000)]);
+    proc = undefined;
+  }
+  if (tmpDir) {
+    try {
+      rmSync(tmpDir, { recursive: true, force: true });
+    } catch {
+      /* ok */
+    }
+    tmpDir = undefined;
+  }
+});
+
+describe("managed CES reconnection (real entrypoint)", () => {
+  test("survives an assistant disconnect and accepts a reconnection", async () => {
+    tmpDir = mkdtempSync(join(tmpdir(), "ces-reconnect-"));
+    const dataDir = join(tmpDir, "ces-data");
+    const socketDir = join(tmpDir, "bootstrap");
+    const socketPath = join(socketDir, "ces.sock");
+    const assistantDataMount = join(tmpDir, "assistant-data-ro");
+    mkdirSync(dataDir, { recursive: true });
+    mkdirSync(socketDir, { recursive: true });
+    mkdirSync(join(assistantDataMount, ".vellum"), { recursive: true });
+
+    const healthPort = await pickFreePort();
+    const managedMain = resolve(__dirname, "..", "managed-main.ts");
+
+    proc = Bun.spawn({
+      cmd: [process.execPath, managedMain],
+      env: {
+        ...process.env,
+        CES_MODE: "managed",
+        CES_DATA_DIR: dataDir,
+        CES_BOOTSTRAP_SOCKET: socketPath,
+        CES_HEALTH_PORT: String(healthPort),
+        CES_ASSISTANT_DATA_MOUNT: assistantDataMount,
+      },
+      stdout: "ignore",
+      stderr: "ignore",
+    });
+
+    // Sidecar comes up and binds the bootstrap socket.
+    await waitForHealth(healthPort);
+    await waitForSocket(socketPath);
+    expect(await readyzRpcConnected(healthPort)).toBe(false);
+
+    // First assistant session connects and completes a handshake.
+    const first = await connectToSocket(socketPath);
+    const ack1 = await handshake(first, "session-1");
+    expect(ack1.accepted).toBe(true);
+
+    // Give /readyz a moment to flip, then confirm CES sees the connection.
+    await delay(200);
+    expect(await readyzRpcConnected(healthPort)).toBe(true);
+
+    // Simulate the assistant pod crashing: drop the connection hard.
+    first.destroy();
+
+    // CES must NOT exit. It should stay healthy, flip rpcConnected back to
+    // false, and re-bind the bootstrap socket to await a reconnection.
+    await waitForSocket(socketPath);
+    expect(proc.killed).toBe(false);
+    const resp = await fetch(`http://127.0.0.1:${healthPort}/healthz`);
+    expect(resp.ok).toBe(true);
+
+    // Wait for the new session's rpcConnected to clear before reconnecting.
+    const cleared = await (async () => {
+      const deadline = Date.now() + 5_000;
+      while (Date.now() < deadline) {
+        if (!(await readyzRpcConnected(healthPort))) return true;
+        await delay(100);
+      }
+      return false;
+    })();
+    expect(cleared).toBe(true);
+
+    // The restarted assistant reconnects and handshakes successfully.
+    const second = await connectToSocket(socketPath);
+    const ack2 = await handshake(second, "session-2");
+    expect(ack2.accepted).toBe(true);
+    expect(ack2.sessionId).toBe("session-2");
+
+    await delay(200);
+    expect(await readyzRpcConnected(healthPort)).toBe(true);
+
+    second.destroy();
+  }, 30_000);
+});

package/src/managed-lazy-getters.ts

@@ -32,6 +32,27 @@ export interface AssistantIdRef {
   current: string;
 }
 
+/**
+ * Overwrite the session-scoped managed credential refs.
+ *
+ * The managed handler registry is long-lived — it persists across assistant
+ * reconnects — so every handshake and every `update_managed_credential` must
+ * fully overwrite these refs, including clearing them when a value is absent.
+ * Otherwise a new or reprovisioned session could keep materializing platform
+ * credentials with the previous session's API key or assistant ID. Absent
+ * values fall back to "" (fail closed): the lazy getters then return no
+ * materialization options, and `getAssistantApiKey` falls back to the env key.
+ */
+export function applyManagedCredentialRefs(
+  apiKeyRef: ApiKeyRef,
+  assistantIdRef: AssistantIdRef,
+  apiKey: string | undefined,
+  assistantId: string | undefined,
+): void {
+  apiKeyRef.current = apiKey ?? "";
+  assistantIdRef.current = assistantId ?? "";
+}
+
 export interface LazyGetterOptions {
   platformBaseUrl: string;
   assistantIdRef: AssistantIdRef;
@@ -55,10 +76,11 @@ export interface LazyGetters {
 export function buildLazyGetters(opts: LazyGetterOptions): LazyGetters {
   const { platformBaseUrl, assistantIdRef, apiKeyRef, envApiKey } = opts;
 
-  const getAssistantApiKey = (): string =>
-    apiKeyRef.current || envApiKey || "";
+  const getAssistantApiKey = (): string => apiKeyRef.current || envApiKey || "";
 
-  const getManagedSubjectOptions = (): ManagedSubjectResolverOptions | undefined => {
+  const getManagedSubjectOptions = ():
+    | ManagedSubjectResolverOptions
+    | undefined => {
     const key = getAssistantApiKey();
     const id = assistantIdRef.current;
     return platformBaseUrl && key && id
@@ -66,7 +88,9 @@ export function buildLazyGetters(opts: LazyGetterOptions): LazyGetters {
       : undefined;
   };
 
-  const getManagedMaterializerOptions = (): ManagedMaterializerOptions | undefined => {
+  const getManagedMaterializerOptions = ():
+    | ManagedMaterializerOptions
+    | undefined => {
     const key = getAssistantApiKey();
     const id = assistantIdRef.current;
     return platformBaseUrl && key && id

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.
@@ -52,6 +56,7 @@ import {
   registerCommandExecutionHandler,
   registerManageSecureCommandToolHandler,
   type RpcHandlerRegistry,
+  type ServeEndReason,
   type SessionIdRef,
 } from "./server.js";
 import {
@@ -67,6 +72,7 @@ import {
   parseHandle,
 } from "@vellumai/service-contracts/credential-rpc";
 import {
+  applyManagedCredentialRefs,
   buildLazyGetters,
   type ApiKeyRef,
   type AssistantIdRef,
@@ -118,7 +124,7 @@ function buildHandlers(
   apiKeyRef: ApiKeyRef,
   assistantIdRef: AssistantIdRef,
   secureKeyBackend: SecureKeyBackend,
-): RpcHandlerRegistry {
+): { handlers: RpcHandlerRegistry; temporaryGrantStore: TemporaryGrantStore } {
   // -- Grant stores ----------------------------------------------------------
   const persistentGrantStore = new PersistentGrantStore(
     getCesGrantsDir("managed"),
@@ -140,15 +146,13 @@ function buildHandlers(
   // though handlers are built before the handshake completes.
   const platformBaseUrl = process.env["VELLUM_PLATFORM_URL"] ?? "";
 
-  const {
-    getManagedSubjectOptions,
-    getManagedMaterializerOptions,
-  } = buildLazyGetters({
-    platformBaseUrl,
-    assistantIdRef,
-    apiKeyRef,
-    envApiKey: process.env["ASSISTANT_API_KEY"] || "",
-  });
+  const { getManagedSubjectOptions, getManagedMaterializerOptions } =
+    buildLazyGetters({
+      platformBaseUrl,
+      assistantIdRef,
+      apiKeyRef,
+      envApiKey: process.env["ASSISTANT_API_KEY"] || "",
+    });
 
   if (!platformBaseUrl) {
     log.warn(
@@ -189,7 +193,10 @@ function buildHandlers(
   };
 
   const localSubjectDepsStub: LocalSubjectResolverDeps = {
-    metadataStore: { getById: () => undefined, list: () => [] } as unknown as LocalSubjectResolverDeps["metadataStore"],
+    metadataStore: {
+      getById: () => undefined,
+      list: () => [],
+    } as unknown as LocalSubjectResolverDeps["metadataStore"],
     oauthConnections: { getById: () => undefined },
   };
 
@@ -408,7 +415,7 @@ function buildHandlers(
     return { results };
   }) as (typeof handlers)[string];
 
-  return handlers;
+  return { handlers, temporaryGrantStore };
 }
 
 // ---------------------------------------------------------------------------
@@ -476,13 +483,18 @@ function startHealthServer(
 }
 
 // ---------------------------------------------------------------------------
-// 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,
@@ -515,16 +527,17 @@ function acceptOneConnection(
       return;
     }
 
-    signal.addEventListener(
-      "abort",
-      () => {
-        cleanup();
-        reject(new Error("Aborted while waiting for connection"));
-      },
-      { once: true },
-    );
+    // 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"));
+    };
+    signal.addEventListener("abort", onAbort, { once: true });
 
     netServer.on("error", (err) => {
+      signal.removeEventListener("abort", onAbort);
       cleanup();
       reject(err);
     });
@@ -534,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 {
@@ -543,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() {
@@ -648,76 +663,157 @@ async function main(): Promise<void> {
   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(
+  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;
+    }
 
-  const endReason = await server.serve();
+    rpcConnected = true;
 
-  rpcConnected = false;
-  log.warn(
-    { reason: endReason, uptime: process.uptime(), pid: process.pid },
-    "RPC session ended — shutting down",
-  );
-  controller.abort("rpc_session_ended");
+    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");
+        }
+      },
+    });
+
+    // `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) => {