@vellumai/cli 0.8.6 → 0.8.7-dev.202606052135.3e62c5a

package/src/commands/client.ts

@@ -1,3 +1,4 @@
+import { spawn } from "node:child_process";
 import { existsSync } from "node:fs";
 import { hostname } from "node:os";
 import path from "node:path";
@@ -16,17 +17,35 @@ import {
   GATEWAY_PORT,
   type Species,
 } from "../lib/constants";
-import { loadGuardianToken } from "../lib/guardian-token";
+import { loadGuardianToken, refreshGuardianToken } from "../lib/guardian-token";
 import { getLocalLanIPv4 } from "../lib/local";
 import {
   CLI_INTERFACE_ID,
   WEB_INTERFACE_ID,
   getClientRegistrationHeaders,
 } from "../lib/client-identity";
+import {
+  getLockfileData,
+  upsertLockfileAssistant,
+  replacePlatformAssistants,
+  runHatch,
+  runRetire,
+  getGuardianAccessToken,
+  parseGatewayUrl,
+  resolveGatewayProxyTarget,
+  readAllowedGatewayPorts,
+  isLoopbackAddr,
+  resolveDevCliInvocation,
+  resolveLockfilePaths,
+  resolveConfigDir,
+  type CliInvocation,
+} from "@vellumai/local-mode";
 import { parseAssistantTargetArg } from "../lib/assistant-target-args.js";
 import {
   fetchOrganizationId,
   fetchPlatformAssistants,
+  getPlatformUrl,
+  getWebUrl,
   readPlatformToken,
 } from "../lib/platform-client";
 import { tuiLog } from "../lib/tui-log";
@@ -64,7 +83,8 @@ function readAssistantName(entry: AssistantEntry | null): string | undefined {
     : undefined;
 }
 
-function parseArgs(): ParsedArgs {
+// Exported for unit testing the arg/auth resolution without launching the TUI.
+export function parseArgs(): ParsedArgs {
   const args = process.argv.slice(3);
 
   const positionalName = parseAssistantTargetArg(args, [
@@ -74,6 +94,8 @@ function parseArgs(): ParsedArgs {
     "-a",
     "--interface",
     "-i",
+    "--token",
+    "-t",
   ]);
   const flagArgs: string[] = [];
   for (let i = 0; i < args.length; i++) {
@@ -87,7 +109,9 @@ function parseArgs(): ParsedArgs {
         arg === "--assistant-id" ||
         arg === "-a" ||
         arg === "--interface" ||
-        arg === "-i") &&
+        arg === "-i" ||
+        arg === "--token" ||
+        arg === "-t") &&
       args[i + 1]
     ) {
       flagArgs.push(arg, args[++i]);
@@ -135,11 +159,31 @@ function parseArgs(): ParsedArgs {
   const cloud = entry?.cloud;
   const species: Species = (entry?.species as Species) ?? "vellum";
 
-  // Platform-hosted assistants use a session token; local assistants use a guardian JWT.
-  const platformToken =
-    cloud === "vellum" ? (readPlatformToken() ?? undefined) : undefined;
-  const bearerToken =
-    cloud === "vellum"
+  // Ephemeral auth: a handed-over token (e.g. from `vellum pair`) used for this
+  // session only. Resolve it BEFORE the credential lookup below so an ephemeral
+  // session never reads (or writes) the local token store.
+  let bearerTokenOverride: string | undefined;
+  for (let i = 0; i < flagArgs.length; i++) {
+    if (
+      (flagArgs[i] === "--token" || flagArgs[i] === "-t") &&
+      flagArgs[i + 1]
+    ) {
+      bearerTokenOverride = flagArgs[i + 1];
+    }
+  }
+
+  // Platform-hosted assistants (cloud "vellum") use a session token; every
+  // other topology — local, docker, and "paired" (a remote assistant paired
+  // from another machine) — uses a bearer guardian JWT. Both are skipped
+  // entirely when --token supplies the credential, so no saved creds are read.
+  const platformToken = bearerTokenOverride
+    ? undefined
+    : cloud === "vellum"
+      ? (readPlatformToken() ?? undefined)
+      : undefined;
+  const bearerToken = bearerTokenOverride
+    ? bearerTokenOverride
+    : cloud === "vellum"
       ? undefined
       : (loadGuardianToken(entry?.assistantId ?? "")?.accessToken ?? undefined);
 
@@ -229,6 +273,9 @@ ${ANSI.bold}ARGUMENTS:${ANSI.reset}
 
 ${ANSI.bold}OPTIONS:${ANSI.reset}
     -u, --url <url>            Runtime URL
+    -t, --token <jwt>          Bearer token to use for this session (e.g. from
+                              'vellum pair'). Overrides the stored token and is
+                              not persisted.
     -a, --assistant-id <id>    Assistant ID
     -i, --interface <id>       Interface identifier: cli (default) or web
     -h, --help                 Show this help message
@@ -242,6 +289,10 @@ ${ANSI.bold}EXAMPLES:${ANSI.reset}
     vellum client vellum-assistant-foo
     vellum client --url http://34.56.78.90:${GATEWAY_PORT}
     vellum client vellum-assistant-foo --url http://localhost:${GATEWAY_PORT}
+
+    # Ephemeral: connect to another machine's assistant with a paired token
+    # (no lockfile entry, nothing persisted):
+    vellum client --url http://10.0.0.196:${GATEWAY_PORT} --token <jwt>
 `);
 }
 
@@ -311,7 +362,291 @@ function findWebDistDir(): string | null {
   return null;
 }
 
+/**
+ * Locate the apps/web source directory for running the Vite dev server.
+ * Only works from a source checkout (not npm-installed).
+ */
+function findWebSourceDir(): string | null {
+  let dir = import.meta.dir;
+  for (let depth = 0; depth < 8; depth++) {
+    const candidate = path.join(dir, "apps", "web", "vite.config.ts");
+    if (existsSync(candidate)) {
+      return path.dirname(candidate);
+    }
+    const parent = path.dirname(dir);
+    if (parent === dir) break;
+    dir = parent;
+  }
+  return null;
+}
+
+const LOCKFILE_PATTERN = /^(?:\/assistant)?\/__local\/lockfile$/;
+const HATCH_PATTERN = /^(?:\/assistant)?\/__local\/hatch$/;
+const RETIRE_PATTERN = /^(?:\/assistant)?\/__local\/retire$/;
+const GUARDIAN_TOKEN_PATTERN =
+  /^(?:\/assistant)?\/__local\/guardian-token\/([^/]+)$/;
+
+function getEnvRecord(): Record<string, string> {
+  const result: Record<string, string> = {};
+  for (const [k, v] of Object.entries(process.env)) {
+    if (v !== undefined) result[k] = v;
+  }
+  return result;
+}
+
+const _localEnv = getEnvRecord();
+const _lockfilePaths = resolveLockfilePaths(_localEnv);
+const _configDir = resolveConfigDir(_localEnv);
+const _baseDir = getBaseDir();
+
+async function handleLocalEndpoints(
+  req: Request,
+  url: URL,
+  server: { requestIP(req: Request): { address: string } | null },
+): Promise<Response | null> {
+  const { pathname } = url;
+  const lockfilePaths = _lockfilePaths;
+  const configDir = _configDir;
+
+  // Check if this is a __local or __gateway route before enforcing loopback.
+  const isLocalRoute =
+    LOCKFILE_PATTERN.test(pathname) ||
+    HATCH_PATTERN.test(pathname) ||
+    RETIRE_PATTERN.test(pathname) ||
+    GUARDIAN_TOKEN_PATTERN.test(pathname) ||
+    parseGatewayUrl(pathname).match;
+
+  if (!isLocalRoute) return null;
+
+  // All __local and __gateway endpoints are restricted to loopback clients.
+  const peer = server.requestIP(req)?.address ?? "";
+  if (!isLoopbackAddr(peer)) {
+    return Response.json({ error: "Forbidden" }, { status: 403 });
+  }
+
+  // Lockfile
+  if (LOCKFILE_PATTERN.test(pathname)) {
+    if (req.method === "GET") {
+      const result = getLockfileData(lockfilePaths);
+      if (result.ok) {
+        return Response.json(result.data);
+      }
+      return new Response(null, { status: result.status });
+    }
+    if (req.method === "POST") {
+      let body: Record<string, unknown>;
+      try {
+        body = (await req.json()) as Record<string, unknown>;
+      } catch {
+        return Response.json(
+          { ok: false, error: "Invalid JSON body" },
+          { status: 400 },
+        );
+      }
+      let result;
+      if (body.syncPlatform && Array.isArray(body.platformAssistants)) {
+        result = replacePlatformAssistants(
+          lockfilePaths,
+          body.platformAssistants as Array<Record<string, unknown>>,
+        );
+      } else {
+        result = upsertLockfileAssistant(
+          lockfilePaths,
+          body.assistant as Record<string, unknown>,
+          body.activeAssistant as string | undefined,
+        );
+      }
+      return Response.json(result, { status: result.ok ? 200 : result.status });
+    }
+    return new Response(null, { status: 405 });
+  }
+
+  // Hatch
+  if (HATCH_PATTERN.test(pathname)) {
+    if (req.method !== "POST") return new Response(null, { status: 405 });
+
+    let species = "vellum";
+    let remote: string | undefined;
+    const contentType = req.headers.get("content-type") ?? "";
+    if (contentType.includes("json")) {
+      try {
+        const body = (await req.json()) as {
+          species?: string;
+          remote?: string;
+        };
+        if (body.species) species = body.species;
+        if (body.remote) remote = body.remote;
+      } catch {
+        return Response.json(
+          { ok: false, error: "Invalid JSON body" },
+          { status: 400 },
+        );
+      }
+    }
+
+    let invocation: CliInvocation;
+    try {
+      invocation = resolveDevCliInvocation(_baseDir);
+    } catch (err) {
+      return Response.json(
+        { ok: false, error: err instanceof Error ? err.message : String(err) },
+        { status: 500 },
+      );
+    }
+
+    const result = await runHatch(
+      invocation,
+      species,
+      remote ? { remote } : undefined,
+    );
+    if (result.ok) {
+      return Response.json({ ok: true, assistantId: result.assistantId });
+    }
+    return Response.json(
+      { ok: false, error: result.error },
+      { status: result.status },
+    );
+  }
+
+  // Retire
+  if (RETIRE_PATTERN.test(pathname)) {
+    if (req.method !== "POST") return new Response(null, { status: 405 });
+
+    let assistantId: string | undefined;
+    try {
+      const body = (await req.json()) as { assistantId?: string };
+      assistantId = body.assistantId;
+    } catch {
+      return Response.json(
+        { ok: false, error: "Invalid JSON body" },
+        { status: 400 },
+      );
+    }
+
+    if (!assistantId) {
+      return Response.json(
+        { ok: false, error: "Missing assistantId" },
+        { status: 400 },
+      );
+    }
+
+    let invocation: CliInvocation;
+    try {
+      invocation = resolveDevCliInvocation(_baseDir);
+    } catch (err) {
+      return Response.json(
+        { ok: false, error: err instanceof Error ? err.message : String(err) },
+        { status: 500 },
+      );
+    }
+
+    const result = await runRetire(invocation, assistantId);
+    if (result.ok) {
+      return Response.json({ ok: true });
+    }
+    return Response.json(
+      { ok: false, error: result.error },
+      { status: result.status },
+    );
+  }
+
+  // Guardian token
+  const guardianMatch = pathname.match(GUARDIAN_TOKEN_PATTERN);
+  if (guardianMatch) {
+    if (req.method !== "GET") return new Response(null, { status: 405 });
+
+    const assistantId = decodeURIComponent(guardianMatch[1]!);
+
+    let invocation: CliInvocation;
+    try {
+      invocation = resolveDevCliInvocation(_baseDir);
+    } catch (err) {
+      return Response.json(
+        { error: err instanceof Error ? err.message : String(err) },
+        { status: 500 },
+      );
+    }
+
+    const result = await getGuardianAccessToken(
+      assistantId,
+      configDir,
+      invocation,
+      true,
+      _localEnv,
+    );
+    if (result.ok) {
+      return Response.json({ accessToken: result.accessToken });
+    }
+    return Response.json({ error: result.error }, { status: result.status });
+  }
+
+  // Gateway proxy — same allowlist decision the web (Vite middleware) and
+  // Electron (`app://` handler) hosts use, so all three can't drift.
+  const gatewayDecision = resolveGatewayProxyTarget(pathname, () =>
+    readAllowedGatewayPorts(lockfilePaths),
+  );
+  if (gatewayDecision.kind === "invalid-port") {
+    return new Response("Port must be between 1024 and 65535", { status: 400 });
+  }
+  if (gatewayDecision.kind === "forbidden-port") {
+    return new Response("Gateway port is not active in lockfile", {
+      status: 403,
+    });
+  }
+  if (gatewayDecision.kind === "forward") {
+    const { target: gatewayTarget } = gatewayDecision;
+    const targetUrl = `http://127.0.0.1:${gatewayTarget.port}${gatewayTarget.path}${url.search}`;
+    const headers = new Headers(req.headers);
+    headers.set("host", `127.0.0.1:${gatewayTarget.port}`);
+
+    try {
+      const hasBody = req.method !== "GET" && req.method !== "HEAD";
+      const proxyRes = await fetch(targetUrl, {
+        method: req.method,
+        headers,
+        body: hasBody ? req.body : undefined,
+        redirect: "manual",
+      });
+      const resHeaders = new Headers(proxyRes.headers);
+      resHeaders.delete("transfer-encoding");
+      return new Response(proxyRes.body, {
+        status: proxyRes.status,
+        statusText: proxyRes.statusText,
+        headers: resHeaders,
+      });
+    } catch {
+      return new Response("Gateway proxy error", { status: 502 });
+    }
+  }
+
+  return null;
+}
+
+function getBaseDir(): string {
+  let dir = import.meta.dir;
+  for (let depth = 0; depth < 8; depth++) {
+    if (existsSync(path.join(dir, "cli", "src", "index.ts"))) {
+      return dir;
+    }
+    const pkgPath = path.join(dir, "package.json");
+    if (existsSync(pkgPath)) {
+      return dir;
+    }
+    const parent = path.dirname(dir);
+    if (parent === dir) break;
+    dir = parent;
+  }
+  return path.resolve(import.meta.dir, "..", "..", "..");
+}
+
 async function runWebInterface(): Promise<void> {
+  // Prefer Vite dev server in source checkouts for full local-mode support
+  // (HMR, __local endpoints, gateway proxy).
+  const webSourceDir = findWebSourceDir();
+  if (webSourceDir) {
+    return runViteDevServer(webSourceDir);
+  }
+
   const distDir = findWebDistDir();
   if (!distDir) {
     console.error(
@@ -323,7 +658,14 @@ async function runWebInterface(): Promise<void> {
     process.exit(1);
   }
 
-  const indexHtml = await Bun.file(path.join(distDir, "index.html")).text();
+  const rawIndexHtml = await Bun.file(path.join(distDir, "index.html")).text();
+  const platformUrl = getPlatformUrl();
+  const webUrl = getWebUrl();
+  const configJson = JSON.stringify({ webUrl, platformUrl });
+  const indexHtml = rawIndexHtml.replace(
+    "</head>",
+    `<script>window.__VELLUM_CONFIG__=${configJson}</script></head>`,
+  );
 
   const server = Bun.serve({
     port: 3000,
@@ -332,10 +674,82 @@ async function runWebInterface(): Promise<void> {
       const url = new URL(req.url);
       const { pathname } = url;
 
-      if (pathname === "/") {
+      if (pathname === "/" || pathname === "/assistant") {
         return Response.redirect(SPA_BASE, 302);
       }
 
+      // Loopback auth: the platform redirects here after login with
+      // ?state=...&session_token=... — forward into the SPA.
+      if (pathname === "/callback") {
+        return Response.redirect(
+          `/account/platform-callback${url.search}`,
+          302,
+        );
+      }
+
+      // Expose environment config to the SPA.
+      if (pathname === "/assistant/__config" || pathname === "/__config") {
+        return new Response(configJson, {
+          headers: { "Content-Type": "application/json" },
+        });
+      }
+
+      // __local endpoints for local-mode (lockfile, hatch, retire, guardian-token, gateway-proxy).
+      const localResponse = await handleLocalEndpoints(req, url, server);
+      if (localResponse) return localResponse;
+
+      // Reverse-proxy platform API requests.
+      if (
+        pathname.startsWith("/v1/") ||
+        pathname.startsWith("/_allauth/") ||
+        pathname.startsWith("/accounts/")
+      ) {
+        const target = new URL(pathname + url.search, platformUrl);
+        const headers = new Headers(req.headers);
+        headers.set("Host", new URL(platformUrl).host);
+        headers.delete("Origin");
+        headers.delete("Referer");
+
+        // Forward the session token — the loopback flow stores it in
+        // the browser cookie jar for localhost, but the platform backend
+        // expects it on its own domain. Set both the Cookie (for Django
+        // session middleware / allauth) and X-Session-Token (for DRF
+        // views that accept header-based auth).
+        const sessionToken = /sessionid=([^;]+)/.exec(
+          req.headers.get("Cookie") ?? "",
+        )?.[1];
+        if (sessionToken) {
+          headers.set(
+            "Cookie",
+            `sessionid=${sessionToken}; __Secure-sessionid=${sessionToken}`,
+          );
+          headers.set("X-Session-Token", sessionToken);
+        }
+
+        try {
+          const hasBody = req.method !== "GET" && req.method !== "HEAD";
+          const body = hasBody ? await req.arrayBuffer() : undefined;
+          const proxyRes = await fetch(target.toString(), {
+            method: req.method,
+            headers,
+            body,
+            redirect: "manual",
+          });
+          const resHeaders = new Headers(proxyRes.headers);
+          resHeaders.delete("transfer-encoding");
+          return new Response(proxyRes.body, {
+            status: proxyRes.status,
+            statusText: proxyRes.statusText,
+            headers: resHeaders,
+          });
+        } catch (err) {
+          return new Response(
+            JSON.stringify({ error: `Platform proxy error: ${err}` }),
+            { status: 502, headers: { "Content-Type": "application/json" } },
+          );
+        }
+      }
+
       if (pathname.startsWith(SPA_BASE)) {
         const relPath = pathname.slice(SPA_BASE.length);
         if (relPath) {
@@ -350,6 +764,13 @@ async function runWebInterface(): Promise<void> {
         });
       }
 
+      // SPA fallback for /account/* routes (login, callback, etc.)
+      if (pathname.startsWith("/account/")) {
+        return new Response(indexHtml, {
+          headers: { "Content-Type": "text/html; charset=utf-8" },
+        });
+      }
+
       return new Response("Not Found", { status: 404 });
     },
   });
@@ -368,6 +789,74 @@ async function runWebInterface(): Promise<void> {
   await new Promise(() => {});
 }
 
+async function runViteDevServer(webSourceDir: string): Promise<void> {
+  const platformUrl = getPlatformUrl();
+
+  const child = spawn("bun", ["run", "dev"], {
+    cwd: webSourceDir,
+    stdio: "inherit",
+    env: {
+      ...process.env,
+      VITE_PLATFORM_MODE: "false",
+      API_PROXY_TARGET: platformUrl,
+      VELLUM_WEB_URL: getWebUrl(),
+      VELLUM_PLATFORM_URL: platformUrl,
+      PORT: "3000",
+    },
+  });
+
+  const shutdown = (): void => {
+    child.kill();
+    process.exit(0);
+  };
+  process.on("SIGINT", shutdown);
+  process.on("SIGTERM", shutdown);
+
+  await new Promise<void>((_, reject) => {
+    child.on("exit", (code) => {
+      if (code !== 0) {
+        reject(new Error(`Vite dev server exited with code ${code}`));
+      }
+    });
+  });
+}
+
+/**
+ * Return a possibly-refreshed bearer token for the TUI's startup auth.
+ *
+ * Only a STORED guardian token is refreshable: platform session auth
+ * (`cloud === "vellum"`) and ephemeral `--token` overrides (whose token won't
+ * match the store) are left untouched, as is a token that's still fresh. When
+ * the stored token has passed its `refreshAfter` (or expiry) and a refresh
+ * token is available, refresh once via the concurrency-safe refreshGuardianToken
+ * and use the rotated access token. Falls back to the existing token if refresh
+ * isn't possible/fails — the session still starts (same as before).
+ */
+export async function resolveFreshBearerToken(
+  runtimeUrl: string,
+  assistantId: string,
+  bearerToken: string | undefined,
+  cloud: string | undefined,
+): Promise<string | undefined> {
+  if (cloud === "vellum" || !bearerToken || !assistantId) return bearerToken;
+
+  const stored = loadGuardianToken(assistantId);
+  // Refresh only the stored token (an ephemeral --token won't match), and only
+  // when a refresh credential is present.
+  if (!stored || stored.accessToken !== bearerToken || !stored.refreshToken) {
+    return bearerToken;
+  }
+
+  // new Date() handles both ISO strings and epoch-ms numbers; Date.parse of an
+  // epoch-ms string would be NaN.
+  const renewAtRaw = stored.refreshAfter || stored.accessTokenExpiresAt;
+  const renewAt = new Date(renewAtRaw).getTime();
+  if (!Number.isFinite(renewAt) || renewAt > Date.now()) return bearerToken;
+
+  const refreshed = await refreshGuardianToken(runtimeUrl, assistantId);
+  return refreshed?.accessToken ?? bearerToken;
+}
+
 export async function client(): Promise<void> {
   const {
     runtimeUrl,
@@ -415,8 +904,19 @@ export async function client(): Promise<void> {
       ...getClientRegistrationHeaders(interfaceId),
     };
   } else {
+    // Proactively refresh a stale STORED guardian token before opening the TUI,
+    // so launching after the access token expired renews transparently rather
+    // than erroring. (Mid-session expiry — the token dying while the TUI is
+    // already open — is a separate follow-up, since the TUI threads a static
+    // auth object through React.)
+    const token = await resolveFreshBearerToken(
+      runtimeUrl,
+      assistantId,
+      bearerToken,
+      cloud,
+    );
     auth = {
-      ...(bearerToken ? { Authorization: `Bearer ${bearerToken}` } : {}),
+      ...(token ? { Authorization: `Bearer ${token}` } : {}),
       ...getClientRegistrationHeaders(interfaceId),
     };
   }