From ae93d0f921b33cc26009c3c0e55a97c4b58be12b Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Fri, 25 Sep 2026 16:43:03 +0200 Subject: [PATCH 1/5] feat(appkit): add DatabricksAdapter.fromAiGateway for gateway chat completions Route agents through the Databricks AI Gateway Chat Completions endpoint (/ai-gateway/mlflow/v1/chat/completions), naming the model in the request body (e.g. system.ai.claude-opus-5-5) instead of in the URL. Reuses the existing Chat Completions tool loop and SSE parsing unchanged. Additive and non-breaking: the serving-endpoint path (fromModelServing, fromServingEndpoint, bare-string model) never sets a body model, so its request is byte-identical to before. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- packages/appkit/src/agents/databricks.ts | 111 ++++++++++++++++++ .../src/agents/tests/databricks.test.ts | 52 ++++++++ .../appkit/src/connectors/serving/client.ts | 23 ++++ 3 files changed, 186 insertions(+) diff --git a/packages/appkit/src/agents/databricks.ts b/packages/appkit/src/agents/databricks.ts index 097f4fdba..1bcd43f65 100644 --- a/packages/appkit/src/agents/databricks.ts +++ b/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>; 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 { + 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[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) { diff --git a/packages/appkit/src/agents/tests/databricks.test.ts b/packages/appkit/src/agents/tests/databricks.test.ts index 4fff9c9a0..e0e9a740a 100644 --- a/packages/appkit/src/agents/tests/databricks.test.ts +++ b/packages/appkit/src/agents/tests/databricks.test.ts @@ -1211,6 +1211,58 @@ 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("parseTextToolCalls", () => { test("parses Llama JSON format", () => { const text = diff --git a/packages/appkit/src/connectors/serving/client.ts b/packages/appkit/src/connectors/serving/client.ts index de9d0465c..60d3cdae0 100644 --- a/packages/appkit/src/connectors/serving/client.ts +++ b/packages/appkit/src/connectors/serving/client.ts @@ -120,3 +120,26 @@ export async function stream( signal, ); } + +/** + * Returns the raw SSE byte stream from the Databricks AI Gateway Chat + * Completions endpoint. Thin wrapper over {@link streamPath} that hard-codes + * the gateway path and forces `stream: true`. + * + * Unlike a serving endpoint, the target model is named in the request body + * (`body.model`, e.g. `"system.ai.claude-opus-5-5"`) — the gateway is a single + * fixed path that routes by the body's `model`, so the caller sets it there. + */ +export async function streamAiGateway( + client: ApiClientLike, + body: Record, + signal?: AbortSignal, +): Promise> { + const { stream: _stream, ...cleanBody } = body; + return streamPath( + client, + "/ai-gateway/mlflow/v1/chat/completions", + { ...cleanBody, stream: true }, + signal, + ); +} From af71144a8fc1cc5959063f27d01564abef43017b Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Mon, 28 Sep 2026 11:52:56 +0200 Subject: [PATCH 2/5] feat(appkit): auto-route system.* model names to the AI Gateway Add adapterFromModelString(), the single decision point both string-model resolution sites (agents plugin + standalone runAgent) now share: a `system.*` UC model name (e.g. system.ai.claude-opus-5-5) resolves to the AI Gateway; any other string (serving-endpoint names like databricks-claude-sonnet-4-5, or custom endpoints) resolves to Model Serving. Non-breaking: serving-endpoint names are [a-zA-Z0-9_-] (no dots), so the `system.` prefix never matches an existing endpoint name. To force a `system.*` name onto a serving endpoint, pass a pre-built adapter as the agent's model. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- packages/appkit/src/agents/databricks.ts | 36 ++++++++++++++++ .../src/agents/tests/databricks.test.ts | 43 +++++++++++++++++++ packages/appkit/src/core/agent/run-agent.ts | 4 +- packages/appkit/src/plugins/agents/agents.ts | 5 ++- 4 files changed, 84 insertions(+), 4 deletions(-) diff --git a/packages/appkit/src/agents/databricks.ts b/packages/appkit/src/agents/databricks.ts index 1bcd43f65..d0d18f3bb 100644 --- a/packages/appkit/src/agents/databricks.ts +++ b/packages/appkit/src/agents/databricks.ts @@ -954,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 { + return model.startsWith("system.") + ? DatabricksAdapter.fromAiGateway({ model, ...options }) + : DatabricksAdapter.fromModelServing(model, options); +} + // --------------------------------------------------------------------------- // Text-based tool call parsing (fallback) // --------------------------------------------------------------------------- diff --git a/packages/appkit/src/agents/tests/databricks.test.ts b/packages/appkit/src/agents/tests/databricks.test.ts index e0e9a740a..722d943e7 100644 --- a/packages/appkit/src/agents/tests/databricks.test.ts +++ b/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, @@ -1263,6 +1264,48 @@ describe("DatabricksAdapter.fromAiGateway", () => { }); }); +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 = diff --git a/packages/appkit/src/core/agent/run-agent.ts b/packages/appkit/src/core/agent/run-agent.ts index 57a659686..0dd79d1d5 100644 --- a/packages/appkit/src/core/agent/run-agent.ts +++ b/packages/appkit/src/core/agent/run-agent.ts @@ -270,8 +270,8 @@ async function resolveAdapter(def: AgentDefinition): Promise { return DatabricksAdapter.fromModelServing(); } if (typeof model === "string") { - const { DatabricksAdapter } = await import("../../agents/databricks"); - return DatabricksAdapter.fromModelServing(model); + const { adapterFromModelString } = await import("../../agents/databricks"); + return adapterFromModelString(model); } return await model; } diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index adcf2e077..b9eba7bb5 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -573,8 +573,9 @@ export class AgentsPlugin extends Plugin implements ToolProvider { } } if (typeof source === "string") { - const { DatabricksAdapter } = await import("../../agents/databricks"); - return DatabricksAdapter.fromModelServing(source, adapterOptions); + const { adapterFromModelString } = + await import("../../agents/databricks"); + return adapterFromModelString(source, adapterOptions); } return await source; } From 44602de95244f53570ee1b297e880a89b749ddda Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Mon, 28 Sep 2026 13:53:28 +0200 Subject: [PATCH 3/5] fix(appkit): route DATABRICKS_SERVING_ENDPOINT_NAME default through gateway routing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The no-model default path resolved DATABRICKS_SERVING_ENDPOINT_NAME inside fromModelServing(), always building a serving-endpoint URL — so a `system.*` value (e.g. system.ai.claude-opus-4-6) 404'd instead of reaching the AI Gateway. Read the env default into the resolver's `source` so it flows through adapterFromModelString() like an explicit string model, in both the agents plugin and standalone runAgent. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- packages/appkit/src/core/agent/run-agent.ts | 15 ++++++++---- packages/appkit/src/plugins/agents/agents.ts | 24 +++++++++----------- 2 files changed, 21 insertions(+), 18 deletions(-) diff --git a/packages/appkit/src/core/agent/run-agent.ts b/packages/appkit/src/core/agent/run-agent.ts index 0dd79d1d5..abb14782b 100644 --- a/packages/appkit/src/core/agent/run-agent.ts +++ b/packages/appkit/src/core/agent/run-agent.ts @@ -264,16 +264,21 @@ async function initStandalonePlugins( } async function resolveAdapter(def: AgentDefinition): Promise { - const { model } = def; - if (!model) { + // Explicit model wins; otherwise fall back to the + // DATABRICKS_SERVING_ENDPOINT_NAME env default. A string from either source + // routes by name (system.* → AI Gateway, else Model Serving). + const source = def.model ?? process.env.DATABRICKS_SERVING_ENDPOINT_NAME; + if (!source) { + // No model and no env default: let fromModelServing() surface the + // canonical "set DATABRICKS_SERVING_ENDPOINT_NAME" error. const { DatabricksAdapter } = await import("../../agents/databricks"); return DatabricksAdapter.fromModelServing(); } - if (typeof model === "string") { + if (typeof source === "string") { const { adapterFromModelString } = await import("../../agents/databricks"); - return adapterFromModelString(model); + return adapterFromModelString(source); } - return await model; + return await source; } function normalizeMessages( diff --git a/packages/appkit/src/plugins/agents/agents.ts b/packages/appkit/src/plugins/agents/agents.ts index b9eba7bb5..db06f80eb 100644 --- a/packages/appkit/src/plugins/agents/agents.ts +++ b/packages/appkit/src/plugins/agents/agents.ts @@ -544,7 +544,14 @@ export class AgentsPlugin extends Plugin implements ToolProvider { def: AgentDefinition, name: string, ): Promise { - const source = def.model ?? this.config.defaultModel; + // Explicit model (adapter or string) wins; otherwise fall back to the + // DATABRICKS_SERVING_ENDPOINT_NAME env default. A string from any source + // routes by name in `adapterFromModelString` (system.* → AI Gateway, + // everything else → Model Serving). + const source = + def.model ?? + this.config.defaultModel ?? + process.env.DATABRICKS_SERVING_ENDPOINT_NAME; // Per-agent adapter knobs from `AgentDefinition` / markdown frontmatter. // Only applied when AppKit builds the adapter itself (string or omitted // model). Users who pass a pre-built `AgentAdapter` own these settings. @@ -559,18 +566,9 @@ export class AgentsPlugin extends Plugin implements ToolProvider { adapterOptions.generationParams = def.generationParams; if (!source) { - const { DatabricksAdapter } = await import("../../agents/databricks"); - try { - return await DatabricksAdapter.fromModelServing( - undefined, - adapterOptions, - ); - } catch (err) { - throw new Error( - `Agent '${name}' has no model configured and no DATABRICKS_SERVING_ENDPOINT_NAME default available`, - { cause: err instanceof Error ? err : undefined }, - ); - } + throw new Error( + `Agent '${name}' has no model configured and no DATABRICKS_SERVING_ENDPOINT_NAME default available`, + ); } if (typeof source === "string") { const { adapterFromModelString } = From 86b27f93193cafe318e7ce534c1c222be3a42b74 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Mon, 28 Sep 2026 14:28:15 +0200 Subject: [PATCH 4/5] chore(playground): use a system.ai gateway model for the autocomplete demo Points the autocomplete agent's `endpoint` at system.ai.gemini-3-5-flash-lite so the reference app exercises the new AI Gateway routing (a system.* string resolves to /ai-gateway/mlflow/v1/chat/completions). Co-authored-by: Isaac Signed-off-by: MarioCadenas --- apps/dev-playground/server/agents/autocomplete/agent.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/dev-playground/server/agents/autocomplete/agent.md b/apps/dev-playground/server/agents/autocomplete/agent.md index 0b8270f07..740296f86 100644 --- a/apps/dev-playground/server/agents/autocomplete/agent.md +++ b/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 --- From 978c25ef6df029435a2123d14a254b23702d70b1 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Mon, 28 Sep 2026 14:28:28 +0200 Subject: [PATCH 5/5] docs(appkit): document DatabricksAdapter.fromAiGateway Regenerated API reference for the new fromAiGateway factory. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../api/appkit/Class.DatabricksAdapter.md | 52 +++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/docs/docs/api/appkit/Class.DatabricksAdapter.md b/docs/docs/api/appkit/Class.DatabricksAdapter.md index 66772a20a..7efc5e890 100644 --- a/docs/docs/api/appkit/Class.DatabricksAdapter.md +++ b/docs/docs/api/appkit/Class.DatabricksAdapter.md @@ -88,6 +88,58 @@ run(input: AgentInput, context: AgentRunContext): AsyncGenerator; +``` + +Creates a DatabricksAdapter that talks to the Databricks AI Gateway +Chat Completions endpoint (`/ai-gateway/mlflow/v1/chat/completions`). + +Unlike [fromModelServing](#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. + +#### Parameters + +| Parameter | Type | +| ------ | ------ | +| `options` | `AiGatewayOptions` | + +#### Returns + +`Promise`\<`DatabricksAdapter`\> + +#### 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, + }), + }, + }), + ], +}); +``` + +*** + ### fromModelServing() ```ts