@dbx-tools/cli 0.9.54 → 0.9.60

package/src/lakebase-proxy/protocol.ts

@@ -0,0 +1,170 @@
+/** Minimal PostgreSQL startup protocol helpers for the loopback proxy. */
+
+import type { Socket } from "node:net";
+
+/** PostgreSQL SSLRequest protocol code. */
+export const SSL_REQUEST = 80_877_103;
+/** PostgreSQL CancelRequest protocol code. */
+export const CANCEL_REQUEST = 80_877_102;
+/** PostgreSQL protocol version 3 code. */
+export const PROTOCOL_3 = 196_608;
+
+/** Parsed PostgreSQL startup packet. */
+export interface StartupMessage {
+  kind: "startup";
+  protocol: number;
+  parameters: Record<string, string>;
+}
+
+/** Parsed PostgreSQL cancellation packet. */
+export interface CancelMessage {
+  kind: "cancel";
+  processId: number;
+  secretKey: number;
+}
+
+/** First supported packet received from a PostgreSQL client. */
+export type InitialMessage = StartupMessage | CancelMessage;
+
+/** Error caused by an invalid or unsupported PostgreSQL startup packet. */
+export class PostgresProtocolError extends Error {}
+
+/** Read and classify the first PostgreSQL frontend message from a socket. */
+export async function readInitialMessage(
+  socket: Socket,
+  timeoutMs: number,
+): Promise<InitialMessage> {
+  const deadline = Date.now() + timeoutMs;
+  for (;;) {
+    const packet = await readPacket(socket, deadline);
+    const code = packet.readInt32BE(4);
+    if (code === SSL_REQUEST) {
+      socket.write("N");
+      continue;
+    }
+    if (code === CANCEL_REQUEST) {
+      if (packet.length !== 16) throw new PostgresProtocolError("Invalid CancelRequest packet");
+      return {
+        kind: "cancel",
+        processId: packet.readInt32BE(8),
+        secretKey: packet.readInt32BE(12),
+      };
+    }
+    if (code !== PROTOCOL_3) {
+      throw new PostgresProtocolError(`Unsupported PostgreSQL protocol version ${code}`);
+    }
+    return { kind: "startup", protocol: code, parameters: parseParameters(packet.subarray(8)) };
+  }
+}
+
+/** Build a PostgreSQL SSL negotiation request packet. */
+export function sslRequest(): Buffer {
+  return requestPacket(SSL_REQUEST);
+}
+
+/** Build a PostgreSQL cancellation request packet. */
+export function cancelRequest(processId: number, secretKey: number): Buffer {
+  return requestPacket(CANCEL_REQUEST, processId, secretKey);
+}
+
+/** Build the backend messages that complete PostgreSQL startup. */
+export function startupComplete(
+  parameters: ReadonlyMap<string, string>,
+  processId: number,
+  secretKey: number,
+  transactionStatus = "I",
+): Buffer {
+  return Buffer.concat([
+    backendMessage("R", int32(0)),
+    ...[...parameters].map(([name, value]) =>
+      backendMessage("S", Buffer.from(`${name}\0${value}\0`)),
+    ),
+    backendMessage("K", Buffer.concat([int32(processId), int32(secretKey)])),
+    backendMessage("Z", Buffer.from(transactionStatus.slice(0, 1) || "I")),
+  ]);
+}
+
+/** Replace client startup identity with the resolved upstream identity. */
+export function upstreamStartupParameters(
+  parameters: Record<string, string>,
+  user: string,
+  database: string,
+): Record<string, string> {
+  return { ...parameters, user, database };
+}
+
+/** Build a fatal PostgreSQL error response. */
+export function fatalError(message: string, sqlstate = "08001"): Buffer {
+  return backendMessage(
+    "E",
+    Buffer.concat([Buffer.from(`SFATAL\0C${sqlstate}\0M${message}\0`), Buffer.from([0])]),
+  );
+}
+
+function requestPacket(code: number, processId?: number, secretKey?: number): Buffer {
+  const values = [code];
+  if (processId !== undefined && secretKey !== undefined) values.push(processId, secretKey);
+  return Buffer.concat([int32(4 + values.length * 4), ...values.map(int32)]);
+}
+
+function backendMessage(code: string, body: Buffer): Buffer {
+  return Buffer.concat([Buffer.from(code), int32(body.length + 4), body]);
+}
+
+function int32(value: number): Buffer {
+  const buffer = Buffer.allocUnsafe(4);
+  buffer.writeInt32BE(value);
+  return buffer;
+}
+
+async function readPacket(socket: Socket, deadline: number): Promise<Buffer> {
+  let buffer = Buffer.alloc(0);
+  while (buffer.length < 4) buffer = Buffer.concat([buffer, await readChunk(socket, deadline)]);
+  const length = buffer.readInt32BE(0);
+  if (length < 8 || length > 1024 * 1024) {
+    throw new PostgresProtocolError("Invalid PostgreSQL startup packet");
+  }
+  while (buffer.length < length) {
+    buffer = Buffer.concat([buffer, await readChunk(socket, deadline)]);
+  }
+  const remainder = buffer.subarray(length);
+  if (remainder.length) socket.unshift(remainder);
+  return buffer.subarray(0, length);
+}
+
+function readChunk(socket: Socket, deadline: number): Promise<Buffer> {
+  const remaining = deadline - Date.now();
+  if (remaining <= 0) {
+    return Promise.reject(new PostgresProtocolError("PostgreSQL startup timed out"));
+  }
+  return new Promise((resolve, reject) => {
+    const timer = setTimeout(
+      () => finish(new PostgresProtocolError("PostgreSQL startup timed out")),
+      remaining,
+    );
+    const onData = (chunk: Buffer) => finish(undefined, chunk);
+    const onError = (error: Error) => finish(error);
+    const onEnd = () =>
+      finish(new PostgresProtocolError("PostgreSQL client closed during startup"));
+    const finish = (error?: Error, chunk?: Buffer) => {
+      clearTimeout(timer);
+      socket.off("data", onData);
+      socket.off("error", onError);
+      socket.off("end", onEnd);
+      if (error) reject(error);
+      else resolve(chunk!);
+    };
+    socket.once("data", onData);
+    socket.once("error", onError);
+    socket.once("end", onEnd);
+  });
+}
+
+function parseParameters(body: Buffer): Record<string, string> {
+  const values = body.toString("utf8").split("\0");
+  const parameters: Record<string, string> = {};
+  for (let index = 0; index + 1 < values.length && values[index]; index += 2) {
+    parameters[values[index]!] = values[index + 1] ?? "";
+  }
+  return parameters;
+}

package/src/lakebase-proxy/proxy.ts

@@ -0,0 +1,256 @@
+/** Pure Node PostgreSQL wire proxy for Databricks Lakebase. */
+
+import { createServer, type Server, type Socket } from "node:net";
+import type { Duplex } from "node:stream";
+import { LakebaseClient, requireAddress } from "@dbx-tools/lakebase";
+import { postgresConnectionOptions } from "@dbx-tools/postgres";
+import { log, stringUtils } from "@dbx-tools/shared-core";
+import { Client } from "pg";
+
+import { CancellationRegistry } from "./cancellation.ts";
+import {
+  resolveLakebaseProxyOptions,
+  type LakebaseProxyOptions,
+  type ResolvedLakebaseProxyOptions,
+} from "./options.ts";
+import {
+  fatalError,
+  PostgresProtocolError,
+  readInitialMessage,
+  startupComplete,
+  upstreamStartupParameters,
+  type CancelMessage,
+} from "./protocol.ts";
+
+const logger = log.logger("lakebase-proxy");
+
+interface PgConnectionInternals {
+  stream: Duplex;
+  on(event: "parameterStatus", listener: (message: ParameterStatus) => void): void;
+}
+
+type PgClientInternals = Client & {
+  connection: Client["connection"] & PgConnectionInternals;
+  processID: number;
+  secretKey: number;
+  _txStatus?: string;
+  getStartupConf(): Record<string, string>;
+};
+
+interface ParameterStatus {
+  parameterName: string;
+  parameterValue: string;
+}
+
+/** Loopback PostgreSQL wire proxy backed by Databricks Lakebase credentials. */
+export class LakebaseProxy {
+  private readonly client: LakebaseClient;
+  private readonly options: ResolvedLakebaseProxyOptions;
+  private readonly cancellations = new CancellationRegistry();
+  private server?: Server;
+  private statsTimer?: ReturnType<typeof setInterval>;
+  private opened = 0;
+  private closed = 0;
+  private failed = 0;
+  private active = 0;
+
+  constructor(options: LakebaseProxyOptions = {}) {
+    this.options = resolveLakebaseProxyOptions(options);
+    this.client = new LakebaseClient({ profile: this.options.profile });
+  }
+
+  async listen(): Promise<{ host: string; port: number }> {
+    if (this.server) throw new Error("Lakebase proxy is already listening");
+    const host = this.options.listen.host;
+    const server = createServer((socket) => void this.handle(socket));
+    this.server = server;
+    this.statsTimer = setInterval(() => this.reportStats(), 60_000);
+    this.statsTimer.unref();
+    await new Promise<void>((resolve, reject) => {
+      server.once("error", reject);
+      server.listen(this.options.listen.port, host, () => {
+        server.off("error", reject);
+        resolve();
+      });
+    });
+    const address = server.address();
+    if (!address || typeof address === "string") {
+      throw new Error("Lakebase listener has no address");
+    }
+    logger.info("Lakebase proxy listening", { host: address.address, port: address.port });
+    return { host: address.address, port: address.port };
+  }
+
+  async close(): Promise<void> {
+    const server = this.server;
+    this.server = undefined;
+    if (this.statsTimer) clearInterval(this.statsTimer);
+    this.statsTimer = undefined;
+    if (!server) return;
+    await new Promise<void>((resolve, reject) =>
+      server.close((error) => (error ? reject(error) : resolve())),
+    );
+  }
+
+  private async handle(local: Socket): Promise<void> {
+    const started = Date.now();
+    let upstream: Duplex | undefined;
+    let cancellation: CancelMessage | undefined;
+    this.opened += 1;
+    this.active += 1;
+    try {
+      local.setNoDelay(true);
+      const initial = await readInitialMessage(local, this.options.startupTimeoutSeconds * 1000);
+      if (initial.kind === "cancel") {
+        await this.cancellations.forward(initial);
+        local.end();
+        return;
+      }
+      const startup = initial;
+      const targetText = startup.parameters.database;
+      if (!targetText) throw new ProxyFailure("startup database is required", "3D000");
+      let target;
+      try {
+        target = requireAddress(targetText);
+      } catch (error) {
+        throw new ProxyFailure(message(error), "3D000");
+      }
+      const startupUser = stringUtils.trimToUndefined(startup.parameters.user);
+      let resolved;
+      try {
+        resolved = await this.client.resolve(target, startupUser);
+      } catch (error) {
+        throw new ProxyFailure(message(error), "3D000");
+      }
+      let password;
+      try {
+        password = await this.client.generateDatabaseCredential(resolved.endpoint, startupUser);
+      } catch (error) {
+        throw new ProxyFailure(message(error), "28000");
+      }
+      let connected;
+      try {
+        connected = await connectUpstream(
+          resolved,
+          password,
+          postgresConnectionOptions(startup.parameters, this.options.postgresRole),
+        );
+      } catch (error) {
+        throw new ProxyFailure(message(error), "08001");
+      }
+      upstream = connected.socket;
+      cancellation = this.cancellations.register({
+        host: resolved.host,
+        port: resolved.port,
+        processId: connected.processId,
+        secretKey: connected.secretKey,
+      });
+      local.write(
+        startupComplete(
+          connected.parameters,
+          cancellation.processId,
+          cancellation.secretKey,
+          connected.transactionStatus,
+        ),
+      );
+      local.pipe(upstream).pipe(local);
+      await Promise.race([closed(local), closed(upstream)]);
+      logger.debug("Lakebase connection closed", {
+        durationMs: Date.now() - started,
+        project: resolved.project,
+        branch: resolved.branch,
+        endpoint: resolved.endpoint,
+      });
+    } catch (error) {
+      this.failed += 1;
+      const failure =
+        error instanceof ProxyFailure
+          ? error
+          : error instanceof PostgresProtocolError
+            ? new ProxyFailure(error.message, "08P01")
+            : new ProxyFailure(message(error));
+      logger.warn("Lakebase connection failed", {
+        durationMs: Date.now() - started,
+        error: failure.message,
+      });
+      if (!local.destroyed) local.end(fatalError(failure.message, failure.sqlstate));
+    } finally {
+      if (cancellation) this.cancellations.remove(cancellation);
+      upstream?.destroy();
+      local.destroy();
+      this.closed += 1;
+      this.active -= 1;
+    }
+  }
+
+  private reportStats(): void {
+    logger.info("Lakebase proxy connections", {
+      opened: this.opened,
+      closed: this.closed,
+      failed: this.failed,
+      active: this.active,
+    });
+  }
+}
+
+async function connectUpstream(
+  resolved: Awaited<ReturnType<LakebaseClient["resolve"]>>,
+  password: string,
+  startup: Record<string, string>,
+): Promise<{
+  socket: Duplex;
+  parameters: Map<string, string>;
+  processId: number;
+  secretKey: number;
+  transactionStatus: string;
+}> {
+  const client = new Client({
+    host: resolved.host,
+    port: resolved.port,
+    database: resolved.database,
+    user: resolved.user,
+    password,
+    ssl: { rejectUnauthorized: true, servername: resolved.host },
+    application_name: startup.application_name,
+    connectionTimeoutMillis: 30_000,
+  }) as unknown as PgClientInternals;
+  const parameters = new Map<string, string>();
+  client.getStartupConf = () =>
+    upstreamStartupParameters(startup, resolved.user, resolved.database);
+  client.on("error", (error) => {
+    logger.debug("Lakebase upstream client closed", { error: error.message });
+  });
+  client.connection.on("parameterStatus", (message) => {
+    parameters.set(message.parameterName, message.parameterValue);
+  });
+  await client.connect();
+  const socket = client.connection.stream;
+  socket.removeAllListeners("data");
+  return {
+    socket,
+    parameters,
+    processId: client.processID,
+    secretKey: client.secretKey,
+    transactionStatus: client._txStatus ?? "I",
+  };
+}
+
+function closed(socket: Duplex): Promise<void> {
+  return new Promise((resolve, reject) => {
+    socket.once("close", () => resolve());
+    socket.once("error", reject);
+  });
+}
+
+function message(error: unknown): string {
+  return error instanceof Error ? error.message : String(error);
+}
+
+class ProxyFailure extends Error {
+  constructor(
+    message: string,
+    readonly sqlstate = "08001",
+  ) {
+    super(message);
+  }
+}

package/src/model-gateway/cli.ts

@@ -0,0 +1,89 @@
+/**
+ * Foreground `dbx model-gateway` command.
+ *
+ * @module
+ */
+
+import { addArgs, parseArgs, serializeArgs } from "@dbx-tools/cli-args";
+import { buildServiceCommand, type CliServiceCliDependencies } from "@dbx-tools/cli-service/cli";
+import { defineService, type CliServiceDefinition } from "@dbx-tools/cli-service/definition";
+import {
+  ModelGatewayCliOptionsSchema,
+  resolveModelGatewayCliOptions,
+  type ModelGatewayCliOptions,
+  type ModelGatewayOptions,
+} from "@dbx-tools/shared-model-gateway/options";
+import { Command } from "commander";
+
+import { PACKAGE_VERSION } from "../../index.ts";
+import { startModelGateway } from "./server.ts";
+
+/** Injectable foreground gateway boundary for CLI tests. */
+export interface ModelGatewayCliDependencies {
+  start(options: ModelGatewayOptions): Promise<void>;
+  /** Optional service CLI boundary for tests and embedding. */
+  readonly service?: CliServiceCliDependencies;
+}
+
+const DEFAULT_DEPENDENCIES: ModelGatewayCliDependencies = {
+  start: startModelGateway,
+};
+
+/** Build the tray-only model-gateway service definition. */
+export function modelGatewayServiceDefinition(
+  options: ModelGatewayOptions = {},
+): CliServiceDefinition {
+  const { runtimeInfo: _, ...resolved } = resolveModelGatewayCliOptions(options);
+  const host = urlHost(resolved.listen.host);
+  return defineService(import.meta.url, {
+    id: "dbx-tools.cli-model-gateway",
+    name: "dbx model gateway",
+    command: {
+      binName: "dbx-model-gateway",
+      arguments: serializeArgs(resolved),
+    },
+    menu: [
+      {
+        type: "url",
+        label: "Models",
+        url: `http://${host}:${resolved.listen.port}/v1/models`,
+      },
+    ],
+  });
+}
+
+/** Build foreground and service model-gateway commands without starting a server. */
+export function buildProgram(
+  name = "dbx model-gateway",
+  dependencies: ModelGatewayCliDependencies = DEFAULT_DEPENDENCIES,
+): Command {
+  const version = PACKAGE_VERSION;
+  const program = addArgs(
+    new Command()
+      .name(name)
+      .description("Run or manage the AppKit Databricks model gateway")
+      .version(version, "-v, --version"),
+    ModelGatewayCliOptionsSchema,
+  ).action(async () => {
+    const options = modelGatewayOptions(program);
+    if (options.runtimeInfo) {
+      process.stdout.write(`${JSON.stringify({ implementation: "typescript-appkit", version })}\n`);
+      return;
+    }
+    const { runtimeInfo: _, ...serverOptions } = options;
+    await dependencies.start(serverOptions);
+  });
+  const serviceCommand = buildServiceCommand(() => {
+    return modelGatewayServiceDefinition(modelGatewayOptions(program));
+  }, dependencies.service);
+  program.addCommand(serviceCommand);
+  return program;
+}
+
+function modelGatewayOptions(command: Command): ModelGatewayCliOptions {
+  return parseArgs(command, ModelGatewayCliOptionsSchema);
+}
+
+function urlHost(host: string): string {
+  return host.includes(":") ? `[${host}]` : host;
+}

package/src/model-gateway/server.ts

@@ -0,0 +1,62 @@
+/**
+ * Foreground model-gateway server owned by the CLI package.
+ *
+ * @module
+ */
+
+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 { workspaceClient } from "@dbx-tools/databricks";
+import {
+  resolveModelGatewayOptions,
+  type ModelGatewayOptions,
+  type ResolvedModelGatewayOptions,
+} from "@dbx-tools/shared-model-gateway/options";
+
+type AppKitServerOptions = NonNullable<Parameters<typeof server>[0]>;
+
+/** Resolve AppKit server settings for the loopback model gateway. */
+export function modelGatewayServerOptions(options: ModelGatewayOptions = {}): AppKitServerOptions {
+  return appKitServerOptions(resolveModelGatewayOptions(options));
+}
+
+function appKitServerOptions(options: ResolvedModelGatewayOptions): AppKitServerOptions {
+  return {
+    bodyLimit: options.bodyLimit,
+    host: options.listen.host,
+    port: options.listen.port,
+  };
+}
+
+/** Start the foreground model gateway and keep its AppKit server active. */
+export async function startModelGateway(options: ModelGatewayOptions = {}): Promise<void> {
+  const resolved = resolveModelGatewayOptions(options);
+  const client = await workspaceClient.createWorkspaceClient({
+    ...(resolved.profile ? { profile: resolved.profile } : {}),
+  });
+  await createApp({
+    client,
+    plugins: [modelGateway(), server(appKitServerOptions(resolved))],
+    onPluginsReady(appkit) {
+      appkit.server.extend((application) => {
+        application.get("/api/healthz", (_request, response) => {
+          sendHealth(response);
+        });
+        application.get("/v1/models", (request, response) => {
+          void appkit.modelGateway.models(request, response);
+        });
+        for (const [path, protocol] of [
+          ["/v1/chat/completions", "openai-chat"],
+          ["/v1/responses", "openai-responses"],
+          ["/v1/messages", "anthropic-messages"],
+          ["/v1/embeddings", "openai-embeddings"],
+        ] as const) {
+          application.post(path, (request, response) => {
+            void appkit.modelGateway.inference(protocol, request, response);
+          });
+        }
+      });
+    },
+  });
+}

package/src/tunnel/app.ts

@@ -0,0 +1,58 @@
+/**
+ * The GATE half of the wrapper: a server-less AppKit app whose only job is to
+ * expose the `authGate` handlers the proxy calls.
+ *
+ * The in-process plugin path (`authGate` in the app's own `plugins`) has a real
+ * HTTP server to mount routes on. A wrapper does not - the app it fronts is a
+ * separate process it must not reach into - so the gate's login routes live on
+ * the proxy instead, and this app exists purely to give the plugin the runtime it
+ * needs: a `CacheManager` for the one-time-code store and signing key, and the
+ * sibling `email` plugin's transport for delivering a code.
+ *
+ * Lazily imported by `cli.ts`, so `dbx tunnel --insecure` (and `install` /
+ * `status`) never load AppKit, the Databricks SDK, or the SMTP stack.
+ *
+ * @module
+ */
+
+import { lakebase, type WorkspaceClient } from "@databricks/appkit";
+import { appkit } from "@dbx-tools/appkit";
+import { authStorage } from "@dbx-tools/auth-gate";
+import { email } from "@dbx-tools/email";
+import { postgresConnectionOptions } from "@dbx-tools/postgres";
+import { object } from "@dbx-tools/shared-core";
+import { authGate, type AuthGateApi, type AuthGateConfig } from "@dbx-tools/tunnel";
+
+/** Boot the gate app and return the handlers the proxy authenticates against. */
+export async function startGateApp(
+  config: AuthGateConfig,
+  client?: WorkspaceClient,
+): Promise<AuthGateApi> {
+  // No `server()` plugin: nothing here listens. `plugins` order matters - `email`
+  // registers its transport before `authGate` looks for a sibling to send with.
+  const plugins = await appkit.createApp({
+    client,
+    plugins: [
+      ...(authStorage.shouldUseLakebase(config)
+        ? [lakebase({ pool: postgresConnectionOptions({}) })]
+        : []),
+      ...(config.sendCode ? [] : [email()]),
+      authGate(config),
+    ],
+  });
+  const api = Reflect.get(plugins, "authGate");
+  if (!isAuthGateApi(api)) throw new Error("the authGate plugin exposed no api");
+  return api;
+}
+
+function isAuthGateApi(value: unknown): value is AuthGateApi {
+  return (
+    object.isRecord(value) &&
+    typeof value.basePath === "string" &&
+    typeof value.passkeysEnabled === "boolean" &&
+    typeof value.handler === "function" &&
+    typeof value.session === "function" &&
+    typeof value.status === "function" &&
+    typeof value.close === "function"
+  );
+}