@vellumai/credential-executor 0.4.55 → 0.4.56

package/src/managed-main.ts

@@ -19,7 +19,7 @@
 
 import { mkdirSync, unlinkSync } from "node:fs";
 import { createServer as createNetServer, type Socket } from "node:net";
-import { dirname } from "node:path";
+import { dirname, join } from "node:path";
 import { Readable, Writable } from "node:stream";
 
 import { CES_PROTOCOL_VERSION, CesRpcMethod } from "@vellumai/ces-contracts";
@@ -47,10 +47,16 @@ import {
   registerCommandExecutionHandler,
   registerManageSecureCommandToolHandler,
   type RpcHandlerRegistry,
+  type SessionIdRef,
 } from "./server.js";
-import { publishBundle } from "./toolstore/publish.js";
-import type { ManagedSubjectResolverOptions } from "./subjects/managed.js";
-import type { ManagedMaterializerOptions } from "./materializers/managed-platform.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/ces-contracts";
+import { buildLazyGetters, type ApiKeyRef } from "./managed-lazy-getters.js";
+import { MANAGED_LOCAL_STATIC_REJECTION_ERROR } from "./managed-errors.js";
 
 // ---------------------------------------------------------------------------
 // Logging (managed always logs to stderr)
@@ -82,7 +88,7 @@ function ensureDataDirs(): void {
 // Build RPC handler registry (managed mode)
 // ---------------------------------------------------------------------------
 
-function buildHandlers(sessionId: string): RpcHandlerRegistry {
+function buildHandlers(sessionIdRef: SessionIdRef, apiKeyRef: ApiKeyRef): RpcHandlerRegistry {
   // -- Grant stores ----------------------------------------------------------
   const persistentGrantStore = new PersistentGrantStore(
     getCesGrantsDir("managed"),
@@ -97,73 +103,143 @@ function buildHandlers(sessionId: string): RpcHandlerRegistry {
 
   // -- Managed credential options --------------------------------------------
   // In managed mode, credentials are obtained from the platform via its
-  // token-materialization endpoint. The platform URL and API key are provided
-  // through environment variables set by the orchestration layer.
+  // token-materialization endpoint. The platform URL and assistant ID come
+  // from environment variables. The API key may come from the env var OR
+  // from the bootstrap handshake (the assistant forwards it after hatch).
+  // We use a lazy getter so the handshake-provided key takes effect even
+  // though handlers are built before the handshake completes.
   const platformBaseUrl = process.env["PLATFORM_BASE_URL"] ?? "";
-  const assistantApiKey = process.env["ASSISTANT_API_KEY"] ?? "";
-
-  const managedSubjectOptions: ManagedSubjectResolverOptions | undefined =
-    platformBaseUrl && assistantApiKey
-      ? { platformBaseUrl, assistantApiKey }
-      : undefined;
-
-  const managedMaterializerOptions: ManagedMaterializerOptions | undefined =
-    platformBaseUrl && assistantApiKey
-      ? { platformBaseUrl, assistantApiKey }
-      : undefined;
+  const assistantId = process.env["PLATFORM_ASSISTANT_ID"] ?? "";
+
+  const { getAssistantApiKey, getManagedSubjectOptions, getManagedMaterializerOptions } =
+    buildLazyGetters({
+      platformBaseUrl,
+      assistantId,
+      apiKeyRef,
+      envApiKey: process.env["ASSISTANT_API_KEY"] || "",
+    });
 
-  if (!managedSubjectOptions) {
+  if (!platformBaseUrl || !assistantId) {
     warn(
-      "PLATFORM_BASE_URL and/or ASSISTANT_API_KEY not set. " +
-        "Managed credential materialisation will not be available.",
+      "PLATFORM_BASE_URL and/or PLATFORM_ASSISTANT_ID not set. " +
+        "Managed credential materialisation will depend on the handshake-provided API key.",
     );
   }
 
+  // -- Workspace root for command execution cwd ------------------------------
+  const assistantDataMount =
+    process.env["CES_ASSISTANT_DATA_MOUNT"] ?? "/assistant-data-ro";
+  const mountedVellumRoot = join(assistantDataMount, ".vellum");
+
   // -- Build handler registry ------------------------------------------------
+  // NOTE: local_static credential handles are NOT supported in managed mode.
+  // The encrypted key store uses PBKDF2 key derivation from user identity
+  // (username, homedir), but the assistant container runs as root while CES
+  // runs as ces — different derived keys make decryption silently fail.
+  // Managed deployments must use platform_oauth handles exclusively.
+  //
+  // We provide error-returning stubs for localMaterialiser/localSubjectDeps
+  // so the HTTP handler compiles but any local_static request gets a clear
+  // error message rather than a silent decryption failure.
+
+  const localMaterialiserStub = {
+    materialise: async () => ({
+      ok: false as const,
+      error: MANAGED_LOCAL_STATIC_REJECTION_ERROR,
+    }),
+  };
 
-  // In managed mode there is no local secure-key backend. The HTTP handler
-  // uses managed subject resolution and managed materialisation instead.
-  // The localMaterialiser and localSubjectDeps are required by HttpExecutorDeps
-  // but will only be reached for local_static/local_oauth handles (which are
-  // not expected in managed deployments). We stub them to fail closed.
-  const stubLocalMaterialiser = {
-    async materialise() {
-      return {
-        ok: false as const,
-        error: "Local credential materialisation is not available in managed mode.",
-      };
-    },
-    reset() {},
+  const localSubjectDepsStub = {
+    metadataStore: { getById: () => undefined, list: () => [] } as any,
+    oauthConnections: { getById: () => undefined },
   };
 
-  const handlers = buildHandlersWithHttp(
-    {
-      persistentGrantStore,
-      temporaryGrantStore,
-      localMaterialiser: stubLocalMaterialiser as any,
-      localSubjectDeps: {
-        metadataStore: { getByServiceField: () => undefined } as any,
-        oauthConnections: { getById: () => undefined },
-      },
-      managedSubjectOptions,
-      managedMaterializerOptions,
-      sessionId,
-    },
-  );
+  // Use a deps object with getters so the handshake-provided API key
+  // is resolved lazily at RPC call time (after the handshake completes).
+  const httpDeps = {
+    persistentGrantStore,
+    temporaryGrantStore,
+    localMaterialiser: localMaterialiserStub as any,
+    localSubjectDeps: localSubjectDepsStub,
+    get managedSubjectOptions() { return getManagedSubjectOptions(); },
+    get managedMaterializerOptions() { return getManagedMaterializerOptions(); },
+    auditStore,
+    sessionId: sessionIdRef,
+  };
+
+  const handlers = buildHandlersWithHttp(httpDeps);
 
-  // Register run_authenticated_command handler
+  // Register run_authenticated_command handler with managed platform materializer
   registerCommandExecutionHandler(handlers, {
     executorDeps: {
       persistentStore: persistentGrantStore,
       temporaryStore: temporaryGrantStore,
-      materializeCredential: async (_handle) => ({
-        ok: false as const,
-        error:
-          "Command credential materialisation in managed mode is not yet available.",
-      }),
+      materializeCredential: async (handle) => {
+        // Parse handle to determine type
+        const parseResult = parseHandle(handle);
+        if (!parseResult.ok) {
+          return { ok: false as const, error: parseResult.error };
+        }
+
+        switch (parseResult.handle.type) {
+          // -- Local static: NOT supported in managed mode -------------------
+          case HandleType.LocalStatic: {
+            return {
+              ok: false as const,
+              error: MANAGED_LOCAL_STATIC_REJECTION_ERROR,
+            };
+          }
+
+          // -- Platform OAuth: materialise via the platform endpoint ----------
+          case HandleType.PlatformOAuth: {
+            const matOpts = getManagedMaterializerOptions();
+            const subOpts = getManagedSubjectOptions();
+            if (!matOpts || !subOpts) {
+              return {
+                ok: false as const,
+                error:
+                  "PLATFORM_BASE_URL and/or ASSISTANT_API_KEY not set. " +
+                  "Managed credential materialisation is not available.",
+              };
+            }
+
+            const subjectResult = await resolveManagedSubject(
+              handle,
+              subOpts,
+            );
+            if (!subjectResult.ok) {
+              return { ok: false as const, error: subjectResult.error.message };
+            }
+
+            const matResult = await materializeManagedToken(
+              subjectResult.subject,
+              matOpts,
+            );
+            if (!matResult.ok) {
+              return { ok: false as const, error: matResult.error.message };
+            }
+
+            return {
+              ok: true as const,
+              value: matResult.token.accessToken,
+              handleType: HandleType.PlatformOAuth,
+            };
+          }
+
+          default:
+            return {
+              ok: false as const,
+              error: `Handle type "${parseResult.handle.type}" is not supported in managed mode. ` +
+                `Supported types: platform_oauth.`,
+            };
+        }
+      },
+      auditStore,
+      sessionId: sessionIdRef,
       cesMode: "managed",
+      egressHooks: buildCesEgressHooks(),
     },
-    defaultWorkspaceDir: "/workspace",
+    defaultWorkspaceDir: join(mountedVellumRoot, "workspace"),
   });
 
   // Register manage_secure_command_tool handler
@@ -171,15 +247,47 @@ function buildHandlers(sessionId: string): RpcHandlerRegistry {
 
   registerManageSecureCommandToolHandler(handlers, {
     downloadBundle: async (sourceUrl: string) => {
-      const resp = await fetch(sourceUrl);
+      const urlError = validateSourceUrl(sourceUrl);
+      if (urlError) {
+        throw new Error(urlError);
+      }
+      const MAX_BUNDLE_SIZE = 100 * 1024 * 1024; // 100 MB
+      const resp = await fetch(sourceUrl, { signal: AbortSignal.timeout(60_000) });
       if (!resp.ok) {
         throw new Error(`HTTP ${resp.status}: ${resp.statusText}`);
       }
-      return Buffer.from(await resp.arrayBuffer());
+      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})`);
+      }
+      // Stream the body and enforce the size limit on actual bytes received,
+      // since Content-Length can be absent (chunked encoding) or lie.
+      const body = resp.body;
+      if (!body) {
+        throw new Error("Response body is null");
+      }
+      const chunks: Uint8Array[] = [];
+      let totalBytes = 0;
+      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})`);
+        }
+        chunks.push(chunk);
+      }
+      return Buffer.concat(chunks);
     },
     publishBundle: (request) => publishBundle({ ...request, cesMode: "managed" }),
     unregisterTool: (toolName: string) => {
-      return toolRegistry.delete(toolName);
+      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);
+        if (!stillInUse) {
+          deleteBundleFromToolstore(entry.bundleDigest, "managed");
+        }
+      }
+      return removed;
     },
     registerTool: (entry) => {
       toolRegistry.set(entry.toolName, entry);
@@ -194,7 +302,6 @@ function buildHandlers(sessionId: string): RpcHandlerRegistry {
 
   handlers[CesRpcMethod.ListGrants] = createListGrantsHandler({
     persistentGrantStore,
-    sessionId,
   }) as typeof handlers[string];
 
   handlers[CesRpcMethod.RevokeGrant] = createRevokeGrantHandler({
@@ -227,11 +334,16 @@ function startHealthServer(port: number, signal: AbortSignal): ReturnType<typeof
         );
       }
       if (url.pathname === "/readyz") {
-        const ready = rpcConnected;
+        // Always return 200 — pod readiness must not depend on whether the
+        // assistant has connected.  When the CES feature flag is off the
+        // assistant never connects, and a 503 here would block pod
+        // 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({ ready, version: CES_PROTOCOL_VERSION }),
+          JSON.stringify({ ready: true, rpcConnected, version: CES_PROTOCOL_VERSION }),
           {
-            status: ready ? 200 : 503,
+            status: 200,
             headers: { "Content-Type": "application/json" },
           },
         );
@@ -388,9 +500,12 @@ async function main(): Promise<void> {
 
   rpcConnected = true;
 
-  // Build the handler registry with all available RPC implementations
-  const sessionId = `ces-managed-${Date.now()}`;
-  const handlers = buildHandlers(sessionId);
+  // 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).
+  const sessionIdRef: SessionIdRef = { current: `ces-managed-${Date.now()}` };
+  const apiKeyRef: ApiKeyRef = { current: "" };
+  const handlers = buildHandlers(sessionIdRef, apiKeyRef);
 
   const server = new CesRpcServer({
     input: connection.readable,
@@ -405,6 +520,17 @@ async function main(): Promise<void> {
         process.stderr.write(`[ces-managed] ERROR: ${msg} ${args.map(String).join(" ")}\n`),
     },
     signal: controller.signal,
+    onHandshakeComplete: (hsSessionId, hsApiKey) => {
+      sessionIdRef.current = hsSessionId;
+      if (hsApiKey) {
+        apiKeyRef.current = hsApiKey;
+        log(`Received assistant API key via handshake`);
+      }
+    },
+    onApiKeyUpdate: (newKey) => {
+      apiKeyRef.current = newKey;
+      log(`Assistant API key updated via RPC`);
+    },
   });
 
   await server.serve();

package/src/materializers/local-oauth-lookup.ts

@@ -0,0 +1,97 @@
+/**
+ * CES-native read-only OAuth connection lookup for local mode.
+ *
+ * In local mode, CES runs as a child process of the assistant on the same
+ * machine and can read the assistant's SQLite database to look up OAuth
+ * connections.
+ *
+ * This implementation opens the database in read-only mode and queries the
+ * `oauth_connections` table directly using raw SQLite queries. It does not
+ * use Drizzle ORM to avoid importing assistant-internal schema modules.
+ *
+ * The lookup is read-only — CES never modifies OAuth connection records.
+ */
+
+import Database from "bun:sqlite";
+import { existsSync } from "node:fs";
+import { join } from "node:path";
+
+import type { OAuthConnectionRecord } from "@vellumai/credential-storage";
+import { oauthConnectionAccessTokenPath } from "@vellumai/credential-storage";
+import type { OAuthConnectionLookup } from "../subjects/local.js";
+
+// ---------------------------------------------------------------------------
+// Raw SQLite row shape (matches oauth_connections table)
+// ---------------------------------------------------------------------------
+
+interface OAuthConnectionRow {
+  id: string;
+  oauth_app_id: string;
+  provider_key: string;
+  account_info: string | null;
+  granted_scopes: string;
+  expires_at: number | null;
+  has_refresh_token: number;
+  status: string;
+  label: string | null;
+  metadata: string | null;
+  created_at: number;
+  updated_at: number;
+}
+
+// ---------------------------------------------------------------------------
+// Row → OAuthConnectionRecord mapping
+// ---------------------------------------------------------------------------
+
+function rowToRecord(row: OAuthConnectionRow): OAuthConnectionRecord {
+  return {
+    id: row.id,
+    providerKey: row.provider_key,
+    accountInfo: row.account_info,
+    grantedScopes: JSON.parse(row.granted_scopes || "[]"),
+    accessTokenPath: oauthConnectionAccessTokenPath(row.id),
+    hasRefreshToken: row.has_refresh_token === 1,
+    expiresAt: row.expires_at,
+    createdAt: row.created_at,
+    updatedAt: row.updated_at,
+  };
+}
+
+// ---------------------------------------------------------------------------
+// Lookup implementation
+// ---------------------------------------------------------------------------
+
+/**
+ * Create a read-only OAuth connection lookup backed by the assistant's
+ * SQLite database.
+ *
+ * @param vellumRoot - The Vellum root directory (e.g. `~/.vellum`).
+ */
+export function createLocalOAuthLookup(
+  vellumRoot: string,
+): OAuthConnectionLookup {
+  const dbPath = join(vellumRoot, "workspace", "data", "db", "assistant.db");
+
+  return {
+    getById(connectionId: string): OAuthConnectionRecord | undefined {
+      if (!existsSync(dbPath)) return undefined;
+
+      let db: Database | undefined;
+      try {
+        db = new Database(dbPath, { readonly: true });
+        const row = db
+          .query<OAuthConnectionRow, [string, string]>(
+            `SELECT * FROM oauth_connections WHERE id = ? AND status = ? LIMIT 1`,
+          )
+          .get(connectionId, "active");
+
+        if (!row) return undefined;
+        return rowToRecord(row);
+      } catch {
+        return undefined;
+      } finally {
+        db?.close();
+      }
+    },
+  };
+}