@vellumai/credential-executor 0.8.4 → 0.8.6

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

@@ -18,10 +18,10 @@
  */
 
 import { afterEach, describe, expect, test } from "bun:test";
-import { mkdirSync, mkdtempSync, rmSync } from "node:fs";
+import { mkdirSync, mkdtempSync, rmSync, unlinkSync } from "node:fs";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
-import { createConnection, type Socket } from "node:net";
+import { createConnection, createServer as createNetServer, type Socket } from "node:net";
 import { Readable, Writable } from "node:stream";
 
 import {
@@ -37,13 +37,12 @@ import {
 } from "@vellumai/service-contracts/credential-rpc";
 
 import { PersistentGrantStore } from "../grants/persistent-store.js";
-import { TemporaryGrantStore } from "../grants/temporary-store.js";
 import { AuditStore } from "../audit/store.js";
 import {
   createListGrantsHandler,
   createListAuditRecordsHandler,
 } from "../grants/rpc-handlers.js";
-import { CesRpcServer, type RpcHandlerRegistry, type SessionIdRef } from "../server.js";
+import { CesRpcServer, type RpcHandlerRegistry, type ServeEndReason, type SessionIdRef } from "../server.js";
 import { createLocalSecureKeyBackend } from "../materializers/local-secure-key-backend.js";
 
 // ---------------------------------------------------------------------------
@@ -159,12 +158,11 @@ function acceptOneConnection(socketPath: string, signal: AbortSignal): Promise<{
   socket: Socket;
 }> {
   return new Promise((resolve, reject) => {
-    const { createServer: createNetServer } = require("node:net");
     const netServer = createNetServer();
 
     const cleanup = () => {
       netServer.close();
-      try { require("node:fs").unlinkSync(socketPath); } catch { /* ok */ }
+      try { unlinkSync(socketPath); } catch { /* ok */ }
     };
 
     if (signal.aborted) {
@@ -188,7 +186,7 @@ function acceptOneConnection(socketPath: string, signal: AbortSignal): Promise<{
 
     netServer.on("connection", (sock: Socket) => {
       netServer.close();
-      try { require("node:fs").unlinkSync(socketPath); } catch { /* ok */ }
+      try { unlinkSync(socketPath); } catch { /* ok */ }
 
       const readable = new Readable({ read() {} });
       const writable = new Writable({
@@ -356,9 +354,10 @@ describe("managed CES integration (real Unix socket)", () => {
     // -- Pick a free port for health server ------------------------------------
     // Use port 0 trick: bind, read the port, close, then use it.
     const healthPort = await new Promise<number>((resolve) => {
-      const srv = require("node:net").createServer();
+      const srv = createNetServer();
       srv.listen(0, () => {
-        const port = srv.address().port;
+        const address = srv.address();
+        const port = typeof address === "object" && address ? address.port : 0;
         srv.close(() => resolve(port));
       });
     });
@@ -681,7 +680,7 @@ describe("credential CRUD RPC", () => {
    */
   async function setupCredentialRpc(): Promise<{
     clientSock: Socket;
-    servePromise: Promise<void>;
+    servePromise: Promise<ServeEndReason>;
   }> {
     savedEnv = saveEnv();
     tmpDir = mkdtempSync(join(tmpdir(), "ces-cred-integ-"));

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-materializers.test.ts

@@ -19,13 +19,11 @@ import { platformOAuthHandle } from "@vellumai/service-contracts/credential-rpc"
 import {
   type ManagedSubject,
   resolveManagedSubject,
-  SubjectResolutionError,
   type PlatformCatalogEntry,
   type ResolvedSubject,
 } from "../subjects/managed.js";
 import {
   materializeManagedToken,
-  MaterializationError,
   type ManagedMaterializerOptions,
 } from "../materializers/managed-platform.js";
 
@@ -82,7 +80,7 @@ function createMockFetch(handlers: {
     error?: Error;
   };
 }): typeof globalThis.fetch {
-  return (async (input: RequestInfo | URL, init?: RequestInit) => {
+  return (async (input: RequestInfo | URL, _init?: RequestInit) => {
     const url = typeof input === "string" ? input : input.toString();
 
     if (url.includes("/oauth/managed/catalog")) {

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/__tests__/toolstore.test.ts

@@ -4,6 +4,7 @@ import {
   rmSync,
   existsSync,
   readFileSync,
+  readdirSync,
   writeFileSync,
   symlinkSync,
 } from "node:fs";
@@ -395,7 +396,6 @@ describe("publishBundle — digest mismatch rejection", () => {
 
     // Also check that no staging directories were left behind
     if (existsSync(toolstoreDir)) {
-      const { readdirSync } = require("node:fs");
       const entries = readdirSync(toolstoreDir) as string[];
       const stagingDirs = entries.filter((e: string) =>
         e.startsWith(".staging-"),

package/src/__tests__/transport.test.ts

@@ -160,7 +160,7 @@ describe("health probes", () => {
     expect(src).toMatch(/\/healthz/);
     expect(src).toMatch(/\/readyz/);
     // Health server uses Bun.serve on a dedicated port, not the socket
-    expect(src).toMatch(/startHealthServer\(healthPort/);
+    expect(src).toMatch(/startHealthServer\(\s*healthPort/);
   });
 
   test("getHealthPort defaults to 8090", () => {
@@ -221,7 +221,7 @@ describe("local entrypoint transport isolation", () => {
 describe("CesRpcServer", () => {
   test("completes handshake with correct protocol version", async () => {
     const { server, handshake, input } = createTestServer();
-    const servePromise = server.serve();
+    const _servePromise = server.serve();
 
     const ack = await handshake();
     expect(ack.type).toBe("handshake_ack");
@@ -237,7 +237,7 @@ describe("CesRpcServer", () => {
 
   test("rejects handshake with wrong protocol version", async () => {
     const { server, send, collectOutputLines, input } = createTestServer();
-    const servePromise = server.serve();
+    const _servePromise = server.serve();
 
     send({
       type: "handshake_request",
@@ -259,7 +259,7 @@ describe("CesRpcServer", () => {
 
   test("rejects RPC before handshake", async () => {
     const { server, send, collectOutputLines, input } = createTestServer();
-    const servePromise = server.serve();
+    const _servePromise = server.serve();
 
     send({
       type: "rpc",
@@ -283,13 +283,13 @@ describe("CesRpcServer", () => {
 
   test("dispatches RPC to registered handler", async () => {
     const handlers: RpcHandlerRegistry = {
-      list_grants: async (req: unknown) => {
+      list_grants: async (_req: unknown) => {
         return { grants: [] };
       },
     };
 
     const { server, handshake, rpc, input } = createTestServer(handlers);
-    const servePromise = server.serve();
+    const _servePromise = server.serve();
 
     await handshake();
 
@@ -304,7 +304,7 @@ describe("CesRpcServer", () => {
 
   test("returns METHOD_NOT_FOUND for unknown methods", async () => {
     const { server, handshake, rpc, input } = createTestServer();
-    const servePromise = server.serve();
+    const _servePromise = server.serve();
 
     await handshake();
 
@@ -326,7 +326,7 @@ describe("CesRpcServer", () => {
     };
 
     const { server, handshake, rpc, input } = createTestServer(handlers);
-    const servePromise = server.serve();
+    const _servePromise = server.serve();
 
     await handshake();
 

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