@dbx-tools/cli 0.9.60 → 0.9.64

package/src/genie-code/options.ts

@@ -0,0 +1,14 @@
+/**
+ * CLI projection of the shared managed Genie Code options.
+ *
+ * @module
+ */
+
+export {
+  GenieCodeOptionsSchema as GenieCodeCliOptionsSchema,
+  resolveGenieCodeOptions as resolveGenieCodeCliOptions,
+} from "@dbx-tools/shared-genie-code/options";
+export type {
+  GenieCodeOptions as GenieCodeCliOptions,
+  ResolvedGenieCodeOptions as ResolvedGenieCodeCliOptions,
+} from "@dbx-tools/shared-genie-code/options";

package/src/genie-code/runner.ts

@@ -0,0 +1,106 @@
+/**
+ * Internal readiness gate and interactive Genie Code child process.
+ *
+ * @module
+ */
+
+import { AppKitChildProcess } from "@dbx-tools/appkit/child-process";
+import { asyncUtils, log } from "@dbx-tools/shared-core";
+import { GenieCodeRunnerOptionsSchema } from "@dbx-tools/shared-genie-code/options";
+
+/** Private environment field carrying serialized runner options. */
+const RUNNER_OPTIONS_ENV = "DBX_TOOLS_GENIE_RUNNER_OPTIONS";
+const logger = log.logger("genie-code:runner");
+
+/** Parse the serialized runner boundary from its private environment field. */
+export function genieCodeRunnerOptions(environment: NodeJS.ProcessEnv = process.env) {
+  const serialized = environment[RUNNER_OPTIONS_ENV];
+  if (!serialized) throw new Error(`${RUNNER_OPTIONS_ENV} is required`);
+  return GenieCodeRunnerOptionsSchema.parse(JSON.parse(serialized));
+}
+
+/** Wait until the guarded model-gateway health route accepts the runner token. */
+export async function waitForGenieCodeGateway(
+  healthUrl: string,
+  bearerToken: string,
+  timeoutMs: number,
+): Promise<void> {
+  for await (const ready of asyncUtils.poll(
+    async ({ signal }) => {
+      try {
+        const response = await fetch(healthUrl, {
+          headers: { authorization: `Bearer ${bearerToken}` },
+          signal,
+        });
+        return response.ok;
+      } catch {
+        return false;
+      }
+    },
+    {
+      intervalMs: 200,
+      timeoutMs,
+      predicate: (ready) => !ready,
+    },
+  )) {
+    if (ready) return;
+  }
+}
+
+/** Run Genie Code with inherited terminal streams after its gateway is ready. */
+export async function runGenieCodeChild(
+  environment: NodeJS.ProcessEnv = process.env,
+): Promise<void> {
+  const options = genieCodeRunnerOptions(environment);
+  logger.info("waiting for model gateway", {
+    gateway: options.gatewayHealthUrl,
+    timeoutMs: options.startupTimeoutMs,
+  });
+  await waitForGenieCodeGateway(
+    options.gatewayHealthUrl,
+    options.bearerToken,
+    options.startupTimeoutMs,
+  );
+  logger.info("model gateway ready; starting Genie Code", {
+    executable: options.executable,
+    home: options.home,
+  });
+  const child = new AppKitChildProcess(
+    [
+      options.executable,
+      [...options.arguments],
+      {
+        cwd: process.cwd(),
+        env: {
+          ...environment,
+          GENIE_HOME: options.home,
+        },
+        stdin: "inherit",
+        stdout: "inherit",
+        stderr: "inherit",
+      },
+    ],
+    {
+      gracefulTimeoutMs: 5_000,
+      forceTimeoutMs: 1_000,
+    },
+  );
+  const running = child.start();
+  const shutdown = () => {
+    void child.shutdown();
+  };
+  for (const signal of ["SIGINT", "SIGTERM", "SIGHUP"] as const) {
+    process.once(signal, shutdown);
+  }
+  try {
+    const result = await running;
+    if (result.exitCode !== 0) process.exitCode = result.exitCode;
+  } finally {
+    for (const signal of ["SIGINT", "SIGTERM", "SIGHUP"] as const) {
+      process.off(signal, shutdown);
+    }
+    await child.shutdown();
+  }
+}
+
+export { RUNNER_OPTIONS_ENV };

package/src/genie-code/runtime.ts

@@ -0,0 +1,231 @@
+/**
+ * Supervise Genie Code and its authenticated local model gateway.
+ *
+ * @module
+ */
+
+import { randomBytes } from "node:crypto";
+import { rm } from "node:fs/promises";
+import { delimiter, dirname } from "node:path";
+import { DatabricksModelRegistry, models } from "@dbx-tools/appkit-model-gateway";
+import * as auth from "@dbx-tools/auth";
+import { resolveServicePackageBin } from "@dbx-tools/cli-service/definition";
+import type { BinContext } from "@dbx-tools/core/bin";
+import { workspaceClient } from "@dbx-tools/databricks";
+import { log } from "@dbx-tools/shared-core";
+import {
+  resolveGenieCodeOptions,
+  type GenieCodeOptions,
+  type GenieCodeRunnerOptions,
+} from "@dbx-tools/shared-genie-code/options";
+import type { ModelClass } from "@dbx-tools/shared-model/contracts";
+import concurrently from "concurrently";
+import getPort from "get-port";
+
+import { writeGenieCodeConfig, type GenieCodeHome } from "./config.ts";
+import { installGenieCode } from "./install.ts";
+import { RUNNER_OPTIONS_ENV } from "./runner.ts";
+
+const logger = log.logger("genie-code");
+
+/** Injectable owners used to prepare and supervise managed Genie Code. */
+export interface GenieCodeRuntimeDependencies {
+  createAuthClient: typeof auth.client.createAuthClient;
+  install: typeof installGenieCode;
+  port(options: { host: string; port?: number }): Promise<number>;
+  resolveBin(packageReference: string, binName: string): string;
+  resolveCodexModel(
+    model: string | undefined,
+    profile: string,
+    modelClass?: ModelClass,
+  ): Promise<string>;
+  supervise: typeof concurrently;
+  token(): string;
+  writeConfig: typeof writeGenieCodeConfig;
+}
+
+const DEFAULT_DEPENDENCIES: GenieCodeRuntimeDependencies = {
+  createAuthClient: auth.client.createAuthClient,
+  install: installGenieCode,
+  port: ({ host, port }) => getPort({ host, ...(port ? { port } : {}) }),
+  resolveBin: resolveServicePackageBin,
+  resolveCodexModel,
+  supervise: concurrently,
+  token: () => randomBytes(32).toString("base64url"),
+  writeConfig: writeGenieCodeConfig,
+};
+
+/** Resolve a Genie `--model` value to the Codex catalogue slug the gateway publishes. */
+async function resolveCodexModel(
+  model: string | undefined,
+  profile: string,
+  modelClass?: ModelClass,
+): Promise<string> {
+  const client = await workspaceClient.createWorkspaceClient({ profile });
+  const target = await new DatabricksModelRegistry({ client }).resolve(model, {
+    requiresTools: true,
+    ...(modelClass ? { modelClass } : {}),
+  });
+  if (!target) {
+    throw new Error(
+      model ? `Genie Code model not found: ${model}` : "No tool-capable chat model is available",
+    );
+  }
+  return models.codexModelSlug(target);
+}
+
+/** Inputs accepted by managed Genie Code preparation and execution. */
+export interface RunGenieCodeOptions {
+  cwd?: string;
+  dependencies?: Partial<GenieCodeRuntimeDependencies>;
+  genieArgs?: readonly string[];
+  options?: GenieCodeOptions;
+}
+
+/** Resolved installation, auth, listener, and home for one invocation. */
+export interface PreparedGenieCodeRuntime {
+  bearerToken: string;
+  gatewayBaseUrl: string;
+  gatewayHealthUrl: string;
+  gatewayListen: string;
+  gatewayPort: number;
+  home: GenieCodeHome;
+  installation: BinContext;
+  model: string;
+  profile: string;
+}
+
+/** Resolve auth, installation, listener, and configuration before spawning children. */
+export async function prepareGenieCodeRuntime(
+  input: RunGenieCodeOptions = {},
+): Promise<PreparedGenieCodeRuntime> {
+  const dependencies = { ...DEFAULT_DEPENDENCIES, ...input.dependencies };
+  const options = resolveGenieCodeOptions(input.options);
+  logger.info("resolving Databricks authentication", {
+    requestedProfile: options.profile,
+  });
+  const authClient = await dependencies.createAuthClient(
+    options.profile ? { profile: options.profile } : {},
+  );
+  const profile = authClient.profile;
+  if (!profile) {
+    throw new Error("Genie Code requires a configured Databricks profile");
+  }
+  logger.info("preparing Genie Code runtime", {
+    profile,
+    model: options.model,
+    modelClass: options.modelClass,
+  });
+  const installation = await dependencies.install();
+  const host = options.gatewayListen.host;
+  const requestedPort = options.gatewayListen.port;
+  const gatewayPort = await dependencies.port({
+    host,
+    ...(requestedPort > 0 ? { port: requestedPort } : {}),
+  });
+  if (requestedPort > 0 && gatewayPort !== requestedPort) {
+    throw new Error(`Genie Code gateway port ${requestedPort} is already in use`);
+  }
+  const model = await dependencies.resolveCodexModel(options.model, profile, options.modelClass);
+  logger.info("resolved Genie Code model", {
+    requested: options.model,
+    modelClass: options.modelClass,
+    model,
+  });
+  const urlHost = host.includes(":") ? `[${host}]` : host;
+  const gatewayBaseUrl = `http://${urlHost}:${gatewayPort}/v1`;
+  const gatewayHealthUrl = `http://${urlHost}:${gatewayPort}/api/healthz`;
+  const gatewayListen = `tcp://${urlHost}:${gatewayPort}`;
+  const bearerToken = dependencies.token();
+  const home = await dependencies.writeConfig({
+    bearerToken,
+    gatewayBaseUrl,
+    model,
+    profile,
+    projectDirectory: input.cwd ?? process.cwd(),
+  });
+  return {
+    bearerToken,
+    gatewayBaseUrl,
+    gatewayHealthUrl,
+    gatewayListen,
+    gatewayPort,
+    home,
+    installation,
+    model,
+    profile,
+  };
+}
+
+function shellExecutable(path: string): string {
+  if (process.platform === "win32") return `"${path.replaceAll('"', '""')}"`;
+  return `'${path.replaceAll("'", "'\\''")}'`;
+}
+
+function childPath(...paths: string[]): string {
+  return [...new Set([...paths.map(dirname), process.env.PATH].filter(Boolean))].join(delimiter);
+}
+
+/** Run Genie Code and the model gateway until either process exits. */
+export async function runGenieCode(input: RunGenieCodeOptions = {}): Promise<void> {
+  const dependencies = { ...DEFAULT_DEPENDENCIES, ...input.dependencies };
+  const prepared = await prepareGenieCodeRuntime({
+    ...input,
+    dependencies,
+  });
+  const gatewayBin = dependencies.resolveBin(import.meta.url, "dbx-model-gateway");
+  const runnerBin = dependencies.resolveBin(import.meta.url, "dbx-genie");
+  const path = childPath(gatewayBin, runnerBin);
+  const runnerOptions: GenieCodeRunnerOptions = {
+    executable: prepared.installation.path,
+    arguments: ["-p", prepared.home.overlayName, ...(input.genieArgs ?? [])],
+    home: prepared.home.home,
+    gatewayHealthUrl: prepared.gatewayHealthUrl,
+    bearerToken: prepared.bearerToken,
+    startupTimeoutMs: 60_000,
+  };
+  const { result } = dependencies.supervise(
+    [
+      {
+        command: shellExecutable(gatewayBin),
+        name: "gateway",
+        env: {
+          ...process.env,
+          PATH: path,
+          LISTEN: prepared.gatewayListen,
+          DATABRICKS_CONFIG_PROFILE: prepared.profile,
+          DBX_TOOLS_MODEL_GATEWAY_BEARER_TOKEN: prepared.bearerToken,
+        },
+      },
+      {
+        command: shellExecutable(runnerBin),
+        name: "genie",
+        env: {
+          ...process.env,
+          PATH: path,
+          DATABRICKS_CONFIG_PROFILE: prepared.profile,
+          [RUNNER_OPTIONS_ENV]: JSON.stringify(runnerOptions),
+        },
+      },
+    ],
+    {
+      cwd: input.cwd ?? process.cwd(),
+      raw: true,
+      killOthersOn: ["failure", "success"],
+      killSignal: "SIGTERM",
+      killTimeout: 17_000,
+      successCondition: "command-genie",
+    },
+  );
+  logger.info("started Genie Code and model gateway", {
+    profile: prepared.profile,
+    model: prepared.model,
+    gateway: prepared.gatewayBaseUrl,
+    home: prepared.home.home,
+  });
+  try {
+    await result;
+  } finally {
+    await rm(prepared.home.overlayPath, { force: true });
+  }
+}

package/src/graphiti/cli.ts

@@ -13,12 +13,12 @@ import { defineService, type CliServiceDefinition } from "@dbx-tools/cli-service
 import { runGraphiti, type GraphitiRuntimeOptions } from "@dbx-tools/graphiti/runtime";
 import { Command } from "commander";
 
-import { PACKAGE_VERSION } from "../../index.ts";
 import {
   GraphitiCliOptionsSchema,
   GraphitiOptionsSchema,
   type GraphitiOptions,
 } from "./options.ts";
+import { PACKAGE_VERSION } from "../../index.ts";
 
 /** Injectable runtime and service lifecycle boundaries for CLI callers. */
 export interface GraphitiCliDependencies {

package/src/lakebase-proxy/cli.ts

@@ -6,7 +6,6 @@ import { defineService, type CliServiceDefinition } from "@dbx-tools/cli-service
 import { connectionUrl } from "@dbx-tools/lakebase";
 import { Command } from "commander";
 
-import { PACKAGE_VERSION } from "../../index.ts";
 import {
   LakebaseProxyOptionsSchema,
   LakebaseProxyUrlOptionsSchema,
@@ -14,6 +13,7 @@ import {
   type LakebaseProxyOptions,
 } from "./options.ts";
 import { LakebaseProxy } from "./proxy.ts";
+import { PACKAGE_VERSION } from "../../index.ts";
 
 /** Injectable service boundary for Lakebase proxy CLI composition and tests. */
 export interface LakebaseProxyCliDependencies {

package/src/model-gateway/cli.ts

@@ -15,8 +15,8 @@ import {
 } from "@dbx-tools/shared-model-gateway/options";
 import { Command } from "commander";
 
-import { PACKAGE_VERSION } from "../../index.ts";
 import { startModelGateway } from "./server.ts";
+import { PACKAGE_VERSION } from "../../index.ts";
 
 /** Injectable foreground gateway boundary for CLI tests. */
 export interface ModelGatewayCliDependencies {
@@ -33,7 +33,11 @@ const DEFAULT_DEPENDENCIES: ModelGatewayCliDependencies = {
 export function modelGatewayServiceDefinition(
   options: ModelGatewayOptions = {},
 ): CliServiceDefinition {
-  const { runtimeInfo: _, ...resolved } = resolveModelGatewayCliOptions(options);
+  const {
+    runtimeInfo: _,
+    bearerToken: _bearerToken,
+    ...resolved
+  } = resolveModelGatewayCliOptions(options);
   const host = urlHost(resolved.listen.host);
   return defineService(import.meta.url, {
     id: "dbx-tools.cli-model-gateway",

package/src/model-gateway/server.ts

@@ -4,15 +4,17 @@
  * @module
  */
 
+import { createHash, timingSafeEqual } from "node:crypto";
 import { createApp, server } from "@databricks/appkit";
-import { modelGateway } from "@dbx-tools/appkit-model-gateway/plugin";
-import { sendHealth } from "@dbx-tools/appkit-model-gateway/routes";
+import { ModelGateway } from "@dbx-tools/appkit-model-gateway";
+import { sendHealth, sendInference, sendModels } from "@dbx-tools/appkit-model-gateway/routes";
 import { workspaceClient } from "@dbx-tools/databricks";
 import {
   resolveModelGatewayOptions,
   type ModelGatewayOptions,
   type ResolvedModelGatewayOptions,
 } from "@dbx-tools/shared-model-gateway/options";
+import type express from "express";
 
 type AppKitServerOptions = NonNullable<Parameters<typeof server>[0]>;
 
@@ -29,22 +31,55 @@ function appKitServerOptions(options: ResolvedModelGatewayOptions): AppKitServer
   };
 }
 
+/** Require one bearer token without retaining or logging request credentials. */
+export function gatewayBearerMiddleware(token: string): express.RequestHandler {
+  const expected = createHash("sha256").update(token).digest();
+  return (request, response, next) => {
+    const authorization = request.header("authorization");
+    const match = /^Bearer\s+(.+)$/i.exec(authorization ?? "");
+    const actual = createHash("sha256")
+      .update(match?.[1] ?? "")
+      .digest();
+    if (match && timingSafeEqual(actual, expected)) {
+      next();
+      return;
+    }
+    response
+      .status(401)
+      .setHeader("www-authenticate", "Bearer")
+      .json({
+        error: {
+          message: "Invalid or missing model-gateway bearer token.",
+          type: "authentication_error",
+          code: "invalid_bearer_token",
+        },
+      });
+  };
+}
+
 /** Start the foreground model gateway and keep its AppKit server active. */
 export async function startModelGateway(options: ModelGatewayOptions = {}): Promise<void> {
   const resolved = resolveModelGatewayOptions(options);
+  const gateway = new ModelGateway({
+    ...(resolved.model ? { model: resolved.model } : {}),
+    ...(resolved.modelClass ? { modelClass: resolved.modelClass } : {}),
+  });
   const client = await workspaceClient.createWorkspaceClient({
     ...(resolved.profile ? { profile: resolved.profile } : {}),
   });
   await createApp({
     client,
-    plugins: [modelGateway(), server(appKitServerOptions(resolved))],
+    plugins: [server(appKitServerOptions(resolved))],
     onPluginsReady(appkit) {
       appkit.server.extend((application) => {
+        if (resolved.bearerToken) {
+          application.use(gatewayBearerMiddleware(resolved.bearerToken));
+        }
         application.get("/api/healthz", (_request, response) => {
           sendHealth(response);
         });
         application.get("/v1/models", (request, response) => {
-          void appkit.modelGateway.models(request, response);
+          void sendModels(gateway, request, response);
         });
         for (const [path, protocol] of [
           ["/v1/chat/completions", "openai-chat"],
@@ -53,7 +88,7 @@ export async function startModelGateway(options: ModelGatewayOptions = {}): Prom
           ["/v1/embeddings", "openai-embeddings"],
         ] as const) {
           application.post(path, (request, response) => {
-            void appkit.modelGateway.inference(protocol, request, response);
+            void sendInference(gateway, protocol, request, response);
           });
         }
       });