feat(appkit): add DatabricksAdapter.fromAiGateway for gateway chat completions

apps/dev-playground/server/agents/autocomplete/agent.md

@@ -1,5 +1,5 @@
 ---
-endpoint: databricks-gemini-3-1-flash-lite
+endpoint:  system.ai.gemini-3-5-flash-lite
 maxSteps: 1
 ephemeral: true
 ---

packages/appkit/src/agents/databricks.ts

@@ -9,6 +9,7 @@ import type {
 import {
   type StreamBody,
   stream as servingStream,
+  streamAiGateway,
 } from "../connectors/serving/client";
 import { APPKIT_USER_AGENT, getClientOptions } from "../context/client-options";
 import { createWorkspaceClient } from "../workspace-client";
@@ -128,6 +129,12 @@ interface RawFetchAdapterOptions {
   authenticate: () => Promise<Record<string, string>>;
   maxSteps?: number;
   maxTokens?: number;
+  /**
+   * When set, sent in the request body as `model`. Required by the AI Gateway,
+   * which routes by the body's `model` (not the URL). Serving-endpoint paths
+   * name the model in the URL and leave this unset.
+   */
+  model?: string;
   /** Optional generation params forwarded to the serving request body. */
   generationParams?: GenerationParams;
   /** Max length of one SSE line (including an incomplete tail in the buffer). */
@@ -148,6 +155,8 @@ interface StreamBodyAdapterOptions {
   streamBody: StreamBody;
   maxSteps?: number;
   maxTokens?: number;
+  /** See {@link RawFetchAdapterOptions.model}. Set by {@link DatabricksAdapter.fromAiGateway}. */
+  model?: string;
   generationParams?: GenerationParams;
   maxSseLineChars?: number;
   maxStreamTextChars?: number;
@@ -198,6 +207,27 @@ interface ModelServingOptions {
   maxToolArgumentsChars?: number;
 }
 
+interface AiGatewayOptions {
+  /**
+   * Model to name in the request body, e.g. `"system.ai.claude-opus-5-5"`.
+   * The gateway routes by this value; there is no per-endpoint URL.
+   */
+  model: string;
+  /**
+   * Pre-built WorkspaceClient (or structural equivalent). When omitted, one
+   * is created from the ambient client options (SDK credential chain). It is
+   * captured once and reused across requests — do not pass a per-request OBO
+   * client (it would leak the first request's identity into later ones).
+   */
+  workspaceClient?: WorkspaceClientLike;
+  maxSteps?: number;
+  maxTokens?: number;
+  generationParams?: GenerationParams;
+  maxSseLineChars?: number;
+  maxStreamTextChars?: number;
+  maxToolArgumentsChars?: number;
+}
+
 interface OpenAIMessage {
   role: "system" | "user" | "assistant" | "tool";
   content: string | null;
@@ -286,6 +316,8 @@ export class DatabricksAdapter implements AgentAdapter {
   private streamBody: StreamBody;
   private maxSteps: number;
   private maxTokens: number;
+  /** UC model name sent in the request body (AI Gateway); unset for serving endpoints. */
+  private model?: string;
   private generationParams: GenerationParams;
   private maxSseLineChars: number;
   private maxStreamTextChars: number;
@@ -294,6 +326,7 @@ export class DatabricksAdapter implements AgentAdapter {
   constructor(options: DatabricksAdapterOptions) {
     this.maxSteps = options.maxSteps ?? 10;
     this.maxTokens = options.maxTokens ?? 4096;
+    this.model = options.model;
     this.generationParams = options.generationParams ?? {};
     this.maxSseLineChars =
       options.maxSseLineChars ?? DEFAULT_MAX_SSE_LINE_CHARS;
@@ -427,6 +460,80 @@ export class DatabricksAdapter implements AgentAdapter {
     });
   }
 
+  /**
+   * Creates a DatabricksAdapter that talks to the Databricks AI Gateway
+   * Chat Completions endpoint (`/ai-gateway/mlflow/v1/chat/completions`).
+   *
+   * Unlike {@link fromModelServing}, the target model is named in the request
+   * body (`model`, e.g. `"system.ai.claude-opus-5-5"`) rather than in the URL:
+   * the gateway is a single fixed path that routes by the body's `model`. Auth
+   * and transport reuse the SDK's `apiClient.request`, same as the serving
+   * path, so no bespoke `fetch()` + token handling. The request/response wire
+   * format and tool-calling loop are identical to the serving path.
+   *
+   * @example
+   * ```ts
+   * import { createApp, createAgent } from "@databricks/appkit";
+   * import { agents, DatabricksAdapter } from "@databricks/appkit/beta";
+   *
+   * const adapter = await DatabricksAdapter.fromAiGateway({
+   *   model: "system.ai.claude-opus-5-5",
+   * });
+   *
+   * await createApp({
+   *   plugins: [
+   *     agents({
+   *       agents: {
+   *         assistant: createAgent({
+   *           instructions: "You are a helpful assistant.",
+   *           model: adapter,
+   *         }),
+   *       },
+   *     }),
+   *   ],
+   * });
+   * ```
+   */
+  static async fromAiGateway(
+    options: AiGatewayOptions,
+  ): Promise<DatabricksAdapter> {
+    const {
+      model,
+      workspaceClient,
+      maxSteps,
+      maxTokens,
+      generationParams,
+      maxSseLineChars,
+      maxStreamTextChars,
+      maxToolArgumentsChars,
+    } = options;
+
+    const client =
+      workspaceClient ??
+      (createWorkspaceClient({
+        clientOptions: getClientOptions(),
+      }) as unknown as WorkspaceClientLike);
+
+    return new DatabricksAdapter({
+      streamBody: (body, signal) =>
+        // Same structural cast as `fromServingEndpoint`: the connector types
+        // the client as the SDK's `WorkspaceClient`, but we only need
+        // `apiClient.request`.
+        streamAiGateway(
+          client as unknown as Parameters<typeof streamAiGateway>[0],
+          body,
+          signal,
+        ),
+      model,
+      maxSteps,
+      maxTokens,
+      generationParams,
+      maxSseLineChars,
+      maxStreamTextChars,
+      maxToolArgumentsChars,
+    });
+  }
+
   /**
    * Discoverability shim for the Supervisor API adapter. Returns an
    * {@link AgentAdapter} (a `SupervisorApiAdapter` at runtime), NOT a
@@ -574,6 +681,10 @@ export class DatabricksAdapter implements AgentAdapter {
       max_tokens: this.maxTokens,
     };
 
+    // AI Gateway routes by the body's `model`; serving endpoints name it in
+    // the URL and leave this unset.
+    if (this.model) body.model = this.model;
+
     applyGenerationParams(body, this.generationParams);
 
     if (tools.length > 0) {
@@ -843,6 +954,42 @@ export class DatabricksAdapter implements AgentAdapter {
   }
 }
 
+/**
+ * Adapter knobs an agent definition contributes to string-model resolution.
+ * A subset of {@link AiGatewayOptions} / {@link ModelServingOptions} — the two
+ * factories accept the same fields, so one object routes to either.
+ */
+type ModelStringOptions = Pick<
+  AiGatewayOptions,
+  "maxSteps" | "maxTokens" | "generationParams"
+>;
+
+/**
+ * Resolves a string `model` to a {@link DatabricksAdapter}, routing by name:
+ *
+ * - UC model names (`system.*`, e.g. `"system.ai.claude-opus-5-5"`) → AI
+ *   Gateway ({@link DatabricksAdapter.fromAiGateway}); the model is named in
+ *   the request body.
+ * - Everything else — serving-endpoint names like
+ *   `"databricks-claude-sonnet-4-5"` or a custom endpoint → Model Serving
+ *   ({@link DatabricksAdapter.fromModelServing}); the name goes in the URL.
+ *
+ * Serving-endpoint names are `[a-zA-Z0-9_-]` (no dots), so the `system.`
+ * prefix cleanly separates the two namespaces and no existing endpoint name
+ * changes routing. This is the single decision point shared by the agents
+ * plugin and standalone `runAgent`, so the two never drift. To force a
+ * `system.*` name onto a serving endpoint instead, pass a pre-built adapter
+ * (`DatabricksAdapter.fromServingEndpoint(...)`) as the agent's `model`.
+ */
+export function adapterFromModelString(
+  model: string,
+  options?: ModelStringOptions,
+): Promise<DatabricksAdapter> {
+  return model.startsWith("system.")
+    ? DatabricksAdapter.fromAiGateway({ model, ...options })
+    : DatabricksAdapter.fromModelServing(model, options);
+}
+
 // ---------------------------------------------------------------------------
 // Text-based tool call parsing (fallback)
 // ---------------------------------------------------------------------------

packages/appkit/src/agents/tests/databricks.test.ts

@@ -2,6 +2,7 @@ import type { AgentEvent, AgentToolDefinition, Message } from "shared";
 import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
 
 import {
+  adapterFromModelString,
   DatabricksAdapter,
   type GenerationParams,
   parseTextToolCalls,
@@ -1211,6 +1212,100 @@ describe("DatabricksAdapter.fromModelServing", () => {
   });
 });
 
+describe("DatabricksAdapter.fromAiGateway", () => {
+  test("routes to the gateway path with `model` in the request body", async () => {
+    const apiClient = {
+      request: vi.fn().mockResolvedValue({
+        contents: createReadableStream([textDelta("Hi"), sseChunk("[DONE]")]),
+      }),
+    };
+
+    const adapter = await DatabricksAdapter.fromAiGateway({
+      model: "system.ai.claude-opus-5-5",
+      workspaceClient: { apiClient },
+    });
+
+    for await (const _ of adapter.run(
+      { messages: createTestMessages(), tools: [], threadId: "t1" },
+      { executeTool: vi.fn() },
+    )) {
+      // drain
+    }
+
+    const [requestArgs] = apiClient.request.mock.calls[0];
+    expect(requestArgs.path).toBe("/ai-gateway/mlflow/v1/chat/completions");
+    expect(requestArgs.method).toBe("POST");
+    expect(requestArgs.raw).toBe(true);
+    expect(requestArgs.payload.model).toBe("system.ai.claude-opus-5-5");
+    expect(requestArgs.payload.stream).toBe(true);
+  });
+
+  test("serving-endpoint path leaves `model` out of the body (non-breaking)", async () => {
+    const apiClient = {
+      request: vi.fn().mockResolvedValue({
+        contents: createReadableStream([textDelta("Hi"), sseChunk("[DONE]")]),
+      }),
+    };
+
+    const adapter = await DatabricksAdapter.fromServingEndpoint({
+      workspaceClient: { apiClient },
+      endpointName: "my-model",
+    });
+
+    for await (const _ of adapter.run(
+      { messages: createTestMessages(), tools: [], threadId: "t1" },
+      { executeTool: vi.fn() },
+    )) {
+      // drain
+    }
+
+    const [requestArgs] = apiClient.request.mock.calls[0];
+    expect(requestArgs.payload.model).toBeUndefined();
+  });
+});
+
+describe("adapterFromModelString", () => {
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  test("routes `system.*` model names to the AI Gateway", async () => {
+    const gateway = vi
+      .spyOn(DatabricksAdapter, "fromAiGateway")
+      .mockResolvedValue({} as unknown as DatabricksAdapter);
+    const serving = vi
+      .spyOn(DatabricksAdapter, "fromModelServing")
+      .mockResolvedValue({} as unknown as DatabricksAdapter);
+
+    await adapterFromModelString("system.ai.claude-opus-5-5", {
+      maxTokens: 128,
+    });
+
+    expect(gateway).toHaveBeenCalledWith({
+      model: "system.ai.claude-opus-5-5",
+      maxTokens: 128,
+    });
+    expect(serving).not.toHaveBeenCalled();
+  });
+
+  test("routes non-`system` names to Model Serving", async () => {
+    const gateway = vi
+      .spyOn(DatabricksAdapter, "fromAiGateway")
+      .mockResolvedValue({} as unknown as DatabricksAdapter);
+    const serving = vi
+      .spyOn(DatabricksAdapter, "fromModelServing")
+      .mockResolvedValue({} as unknown as DatabricksAdapter);
+
+    await adapterFromModelString("databricks-claude-sonnet-4-5");
+
+    expect(serving).toHaveBeenCalledWith(
+      "databricks-claude-sonnet-4-5",
+      undefined,
+    );
+    expect(gateway).not.toHaveBeenCalled();
+  });
+});
+
 describe("parseTextToolCalls", () => {
   test("parses Llama JSON format", () => {
     const text =