diff --git a/packages/appkit/src/agents/databricks.ts b/packages/appkit/src/agents/databricks.ts index 3e5d9317a..803723fc2 100644 --- a/packages/appkit/src/agents/databricks.ts +++ b/packages/appkit/src/agents/databricks.ts @@ -10,6 +10,7 @@ import { type StreamBody, stream as servingStream, streamAiGateway, + type WorkspaceRequestClientLike, } from "../connectors/serving/client"; import { APPKIT_USER_AGENT, getClientOptions } from "../context/client-options"; import { createWorkspaceClient } from "../workspace-client"; @@ -149,7 +150,7 @@ interface RawFetchAdapterOptions { * Preferred options: caller provides the transport function directly. * The `fromServingEndpoint` / `fromModelServing` factories use this to route * through `connectors/serving/stream`, which centralises URL encoding, auth - * via the SDK's `apiClient.request`, and any future retries/telemetry. + * via the workspace client's `request()`, and any future retries/telemetry. */ interface StreamBodyAdapterOptions { streamBody: StreamBody; @@ -174,17 +175,19 @@ function isStreamBodyOptions( } /** - * Duck-typed subset of the Databricks SDK `WorkspaceClient`. Callers of - * `fromServingEndpoint` and `fromModelServing` pass a real `WorkspaceClient`, - * but we only need the `apiClient.request` surface — so we declare the minimal - * interface rather than importing the SDK type directly. This keeps the adapter - * free of a hard compile-time dependency on `@databricks/sdk-experimental`. + * Duck-typed subset of the AppKit workspace client. Callers of + * `fromServingEndpoint` and `fromModelServing` pass `createWorkspaceClient()`, + * but we only need its raw `request()` seam, so we declare the minimal shape + * rather than importing the client type. A legacy SDK client exposing only + * `apiClient.request` is still accepted (deprecated). */ -export interface WorkspaceClientLike { - apiClient: { - request(options: Record): Promise; - }; -} +type WorkspaceClientLike = + | WorkspaceRequestClientLike + | { + apiClient: { + request(options: Record): Promise; + }; + }; /** * A fixed client, or a provider resolved on every model call. A provider lets @@ -403,7 +406,7 @@ export class DatabricksAdapter implements AgentAdapter { * Creates a DatabricksAdapter for a Databricks Model Serving endpoint. * * Routes through the shared `connectors/serving/stream` helper, which - * delegates to the SDK's `apiClient.request({ raw: true })`. That gives the + * delegates to the workspace client's raw `request()`. That gives the * adapter centralised URL encoding + authentication with the rest of the * serving surface — no bespoke `fetch()` + `authenticate()` plumbing. */ @@ -423,15 +426,7 @@ export class DatabricksAdapter implements AgentAdapter { const resolveClient = clientResolver(workspaceClient); const adapter = new DatabricksAdapter({ streamBody: (body, signal) => - // Cast through the structural shape: the connector types - // `workspaceClient` as the SDK's concrete `WorkspaceClient`, but we - // only need `apiClient.request`. - servingStream( - resolveClient() as unknown as Parameters[0], - endpointName, - body, - signal, - ), + servingStream(resolveClient(), endpointName, body, signal), maxSteps, maxTokens, generationParams, @@ -481,7 +476,7 @@ export class DatabricksAdapter implements AgentAdapter { if (!workspaceClient) { workspaceClient = createWorkspaceClient({ clientOptions: getClientOptions(), - }) as unknown as WorkspaceClientLike; + }); } return DatabricksAdapter.fromServingEndpoint({ @@ -503,7 +498,7 @@ export class DatabricksAdapter implements AgentAdapter { * 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 + * and transport reuse the workspace client's `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. * @@ -546,21 +541,14 @@ export class DatabricksAdapter implements AgentAdapter { const resolveClient = clientResolver( workspaceClient ?? - (createWorkspaceClient({ + createWorkspaceClient({ clientOptions: getClientOptions(), - }) as unknown as WorkspaceClientLike), + }), ); const adapter = 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( - resolveClient() as unknown as Parameters[0], - body, - signal, - ), + streamAiGateway(resolveClient(), body, signal), model, maxSteps, maxTokens, diff --git a/packages/appkit/src/agents/supervisor-api.ts b/packages/appkit/src/agents/supervisor-api.ts index b635b92f4..e5b4dbd4b 100644 --- a/packages/appkit/src/agents/supervisor-api.ts +++ b/packages/appkit/src/agents/supervisor-api.ts @@ -11,6 +11,7 @@ import { type ApiClientLike, type StreamBody, streamPath, + type WorkspaceRequestClientLike, } from "../connectors/serving/client"; import { createLogger } from "../logging/logger"; import { readSseEvents } from "../stream"; @@ -90,16 +91,17 @@ function summariseErrorPayload(payload: unknown): string { /** * Structural shape of a Databricks SDK client used by {@link fromSupervisorApi}. - * Only what we need: `apiClient.request` for streaming and - * `config.ensureResolved` to materialise the host/credentials. + * Only what we need: `request` for streaming and `getHost` to materialise + * the host/credentials up front. A legacy SDK client (`apiClient.request` + + * `config.ensureResolved`) is still accepted (deprecated). * * Exported because {@link SupervisorApiAdapterOptions.workspaceClient} (a * public type) references it — callers passing their own client can name * the shape they need to satisfy. */ -export interface WorkspaceClientLike extends ApiClientLike { - config: { ensureResolved(): Promise }; -} +export type WorkspaceClientLike = + | (WorkspaceRequestClientLike & { getHost(): Promise }) + | (ApiClientLike & { config: { ensureResolved(): Promise } }); // --------------------------------------------------------------------------- // Supervisor API tool surface (wire format) @@ -346,8 +348,8 @@ interface SupervisorApiAdapterCtorOptions { * Authentication is handled via the Databricks SDK credential chain — the * same mechanism used by `DatabricksAdapter.fromModelServing`. The transport * is injected via {@link SupervisorApiAdapterCtorOptions.streamBody}; the - * {@link fromSupervisorApi} factory wires it through the SDK's - * `apiClient.request({ raw: true })`. + * {@link fromSupervisorApi} factory wires it through the workspace client's + * raw `request()`. * * Set `DEBUG=appkit:agents:supervisor-api` to log the outbound request * shape (model, instructions length, input shape, tool count) and to be @@ -877,22 +879,14 @@ function mapEvent( export async function fromSupervisorApi( options: SupervisorApiAdapterOptions, ): Promise { - let client = options.workspaceClient; - if (!client) { - // The wrapper's client provides everything `WorkspaceClientLike` needs - // (`apiClient.request` + `config.ensureResolved`) but its - // `apiClient.request` signature is narrower than our structural - // `Record` shape, so a direct assignment doesn't type. - // The cast bridges the structural gap — same pattern the serving - // connector uses for `ApiClientLike`. - client = createWorkspaceClient() as unknown as WorkspaceClientLike; - } + const resolved: WorkspaceClientLike = + options.workspaceClient ?? createWorkspaceClient(); - await client.config.ensureResolved(); + // Resolve host + credentials now so misconfiguration fails at construction, + // not on the first stream. + if ("request" in resolved) await resolved.getHost(); + else await resolved.config.ensureResolved(); - // Capture the resolved client so the closure doesn't depend on the outer - // `let` binding being reassigned later. - const resolved = client; return new SupervisorApiAdapter({ streamBody: (body, signal) => streamPath(resolved, "/ai-gateway/mlflow/v1/responses", body, signal), diff --git a/packages/appkit/src/agents/tests/databricks.test.ts b/packages/appkit/src/agents/tests/databricks.test.ts index a511ef6d4..e4e2ff833 100644 --- a/packages/appkit/src/agents/tests/databricks.test.ts +++ b/packages/appkit/src/agents/tests/databricks.test.ts @@ -1094,7 +1094,7 @@ describe("DatabricksAdapter.fromServingEndpoint", () => { test.each([ [ "fromServingEndpoint", - (workspaceClient: () => { apiClient: unknown }) => + (workspaceClient: () => { request: unknown }) => DatabricksAdapter.fromServingEndpoint({ workspaceClient: workspaceClient as never, endpointName: "my-model", @@ -1102,7 +1102,7 @@ describe("DatabricksAdapter.fromServingEndpoint", () => { ], [ "fromAiGateway", - (workspaceClient: () => { apiClient: unknown }) => + (workspaceClient: () => { request: unknown }) => DatabricksAdapter.fromAiGateway({ workspaceClient: workspaceClient as never, model: "system.ai.claude", @@ -1111,11 +1111,12 @@ describe("DatabricksAdapter.fromServingEndpoint", () => { ])("%s resolves a client provider on every run", async (_name, build) => { const clients = ["alice", "bob"].map((user) => ({ user, - apiClient: { - request: vi.fn(async () => ({ - contents: createReadableStream([textDelta(user), sseChunk("[DONE]")]), - })), - }, + request: vi.fn( + async () => + new Response( + createReadableStream([textDelta(user), sseChunk("[DONE]")]), + ), + ), })); let next = 0; const adapter = await build(() => clients[next++]); @@ -1127,19 +1128,23 @@ describe("DatabricksAdapter.fromServingEndpoint", () => { // drain } } - expect(clients[0].apiClient.request).toHaveBeenCalledTimes(1); - expect(clients[1].apiClient.request).toHaveBeenCalledTimes(1); + expect(clients[0].request).toHaveBeenCalledTimes(1); + expect(clients[1].request).toHaveBeenCalledTimes(1); }); - test("routes tool-free chat through apiClient.request with a streaming payload", async () => { - const apiClient = { - request: vi.fn().mockResolvedValue({ - contents: createReadableStream([textDelta("Hi"), sseChunk("[DONE]")]), - }), + test("routes tool-free chat through request() with a streaming payload", async () => { + const client = { + request: vi + .fn() + .mockResolvedValue( + new Response( + createReadableStream([textDelta("Hi"), sseChunk("[DONE]")]), + ), + ), }; const adapter = await DatabricksAdapter.fromServingEndpoint({ - workspaceClient: { apiClient }, + workspaceClient: client, endpointName: "my-model", }); @@ -1150,25 +1155,28 @@ describe("DatabricksAdapter.fromServingEndpoint", () => { // drain } - expect(apiClient.request).toHaveBeenCalledTimes(1); - const [requestArgs] = apiClient.request.mock.calls[0]; + expect(client.request).toHaveBeenCalledTimes(1); + const [requestArgs] = client.request.mock.calls[0]; expect(requestArgs.path).toBe("/serving-endpoints/my-model/invocations"); expect(requestArgs.method).toBe("POST"); - expect(requestArgs.raw).toBe(true); - expect(requestArgs.payload.stream).toBe(true); + expect(JSON.parse(requestArgs.body).stream).toBe(true); // Auth + url encoding are the connector's (and the SDK's) concerns — the // adapter no longer reaches into the workspace config. }); test("URL-encodes endpoint names with special characters", async () => { - const apiClient = { - request: vi.fn().mockResolvedValue({ - contents: createReadableStream([textDelta("Hi"), sseChunk("[DONE]")]), - }), + const client = { + request: vi + .fn() + .mockResolvedValue( + new Response( + createReadableStream([textDelta("Hi"), sseChunk("[DONE]")]), + ), + ), }; const adapter = await DatabricksAdapter.fromServingEndpoint({ - workspaceClient: { apiClient }, + workspaceClient: client, endpointName: "my model/with spaces", }); @@ -1179,7 +1187,7 @@ describe("DatabricksAdapter.fromServingEndpoint", () => { // drain } - const [requestArgs] = apiClient.request.mock.calls[0]; + const [requestArgs] = client.request.mock.calls[0]; expect(requestArgs.path).toBe( "/serving-endpoints/my%20model%2Fwith%20spaces/invocations", ); @@ -1206,7 +1214,7 @@ describe("DatabricksAdapter.fromModelServing", () => { return { ...actual, createWorkspaceClient: vi.fn().mockImplementation(() => ({ - apiClient: { request: vi.fn() }, + request: vi.fn(), })), }; }); @@ -1226,14 +1234,18 @@ describe("DatabricksAdapter.fromModelServing", () => { test("explicit endpoint name takes precedence over env var", async () => { process.env.DATABRICKS_SERVING_ENDPOINT_NAME = "env-model"; - const apiClient = { - request: vi.fn().mockResolvedValue({ - contents: createReadableStream([textDelta("Hi"), sseChunk("[DONE]")]), - }), + const client = { + request: vi + .fn() + .mockResolvedValue( + new Response( + createReadableStream([textDelta("Hi"), sseChunk("[DONE]")]), + ), + ), }; const adapter = await DatabricksAdapter.fromModelServing("explicit-model", { - workspaceClient: { apiClient }, + workspaceClient: client, }); expect(adapter).toBeInstanceOf(DatabricksAdapter); @@ -1245,7 +1257,7 @@ describe("DatabricksAdapter.fromModelServing", () => { // drain } - const [requestArgs] = apiClient.request.mock.calls[0]; + const [requestArgs] = client.request.mock.calls[0]; expect(requestArgs.path).toBe( "/serving-endpoints/explicit-model/invocations", ); @@ -1254,15 +1266,19 @@ 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 client = { + request: vi + .fn() + .mockResolvedValue( + new Response( + createReadableStream([textDelta("Hi"), sseChunk("[DONE]")]), + ), + ), }; const adapter = await DatabricksAdapter.fromAiGateway({ model: "system.ai.claude-opus-5-5", - workspaceClient: { apiClient }, + workspaceClient: client, }); for await (const _ of adapter.run( @@ -1272,23 +1288,28 @@ describe("DatabricksAdapter.fromAiGateway", () => { // drain } - const [requestArgs] = apiClient.request.mock.calls[0]; + const [requestArgs] = client.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); + expect(JSON.parse(requestArgs.body).model).toBe( + "system.ai.claude-opus-5-5", + ); + expect(JSON.parse(requestArgs.body).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 client = { + request: vi + .fn() + .mockResolvedValue( + new Response( + createReadableStream([textDelta("Hi"), sseChunk("[DONE]")]), + ), + ), }; const adapter = await DatabricksAdapter.fromServingEndpoint({ - workspaceClient: { apiClient }, + workspaceClient: client, endpointName: "my-model", }); @@ -1299,8 +1320,8 @@ describe("DatabricksAdapter.fromAiGateway", () => { // drain } - const [requestArgs] = apiClient.request.mock.calls[0]; - expect(requestArgs.payload.model).toBeUndefined(); + const [requestArgs] = client.request.mock.calls[0]; + expect(JSON.parse(requestArgs.body).model).toBeUndefined(); }); }); diff --git a/packages/appkit/src/agents/tests/supervisor-api.test.ts b/packages/appkit/src/agents/tests/supervisor-api.test.ts index c9d632825..c7cd02429 100644 --- a/packages/appkit/src/agents/tests/supervisor-api.test.ts +++ b/packages/appkit/src/agents/tests/supervisor-api.test.ts @@ -1065,7 +1065,17 @@ describe("SupervisorApiAdapter", () => { }); describe("fromSupervisorApi", () => { - test("calls ensureResolved on the supplied workspace client", async () => { + test("resolves the host on the supplied workspace client", async () => { + const getHost = vi.fn(async () => "https://example.databricks.com"); + const adapter = await fromSupervisorApi({ + model: "databricks-claude-sonnet-4", + workspaceClient: { getHost, request: vi.fn() }, + }); + expect(getHost).toHaveBeenCalledTimes(1); + expect(adapter).toBeInstanceOf(SupervisorApiAdapter); + }); + + test("still calls ensureResolved on a legacy SDK client", async () => { const ensureResolved = vi.fn(async () => {}); const adapter = await fromSupervisorApi({ model: "databricks-claude-sonnet-4", @@ -1078,7 +1088,7 @@ describe("fromSupervisorApi", () => { expect(adapter).toBeInstanceOf(SupervisorApiAdapter); }); - test("routes streaming through apiClient.request with the SA path", async () => { + test("routes streaming through request() with the SA path", async () => { const encoder = new TextEncoder(); const contents = new ReadableStream({ start(controller) { @@ -1086,14 +1096,11 @@ describe("fromSupervisorApi", () => { controller.close(); }, }); - const request = vi.fn().mockResolvedValue({ contents }); + const request = vi.fn().mockResolvedValue(new Response(contents)); const adapter = await fromSupervisorApi({ model: "databricks-claude-sonnet-4", - workspaceClient: { - config: { ensureResolved: vi.fn(async () => {}) }, - apiClient: { request }, - }, + workspaceClient: { getHost: vi.fn(async () => "h"), request }, }); await collect(adapter.run(createInput(), { executeTool: vi.fn() })); @@ -1102,13 +1109,13 @@ describe("fromSupervisorApi", () => { const [requestArgs] = request.mock.calls[0]; expect(requestArgs.path).toBe("/ai-gateway/mlflow/v1/responses"); expect(requestArgs.method).toBe("POST"); - expect(requestArgs.raw).toBe(true); - expect(requestArgs.payload).toMatchObject({ + const payload = JSON.parse(requestArgs.body); + expect(payload).toMatchObject({ model: "databricks-claude-sonnet-4", input: "Hello", stream: true, }); - expect(requestArgs.payload).not.toHaveProperty("tools"); + expect(payload).not.toHaveProperty("tools"); }); }); diff --git a/packages/appkit/src/connectors/serving/client.ts b/packages/appkit/src/connectors/serving/client.ts index 3a9b44a80..9c46f29de 100644 --- a/packages/appkit/src/connectors/serving/client.ts +++ b/packages/appkit/src/connectors/serving/client.ts @@ -1,19 +1,30 @@ import { createLogger } from "../../logging/logger"; -import type { - serving, - WorkspaceClient, - WorkspaceRequest, -} from "../../workspace-client"; +import type { serving, WorkspaceClient } from "../../workspace-client"; import { contextFromAbortSignal } from "../context"; const logger = createLogger("connectors:serving"); /** - * Structural shape of a Databricks SDK client we need for the low-level - * request call. Lets `streamPath` be reused by adapters that don't want a - * hard dependency on the concrete `WorkspaceClient` type. AppKit's own client - * provides `request` (modular transport); a caller-supplied legacy SDK client - * only has `apiClient.request`, which stays supported. + * Structural shape of the AppKit workspace client's raw-request seam + * (`createWorkspaceClient().request`): sends through the modular transport + * (AppKit User-Agent + auth) and returns the unread fetch `Response`. + */ +export interface WorkspaceRequestClientLike { + request(req: { + method: string; + path: string; + headers?: Record; + body?: string; + signal?: AbortSignal; + }): Promise; +} + +/** + * Structural shape of a legacy Databricks SDK client's low-level + * `apiClient.request` call. + * + * @deprecated Pass an AppKit workspace client ({@link WorkspaceRequestClientLike}). + * Still accepted so callers passing a raw legacy SDK client keep working. */ export interface ApiClientLike { apiClient: { @@ -22,7 +33,6 @@ export interface ApiClientLike { context?: unknown, ): Promise; }; - request?(req: WorkspaceRequest): Promise; } // The legacy SDK's `servingEndpoints.query` copied only these fields into the @@ -45,6 +55,9 @@ const QUERY_FIELDS = [ "usage_context", ]; +/** A client {@link streamPath} can send through. */ +type StreamClientLike = WorkspaceRequestClientLike | ApiClientLike; + /** * Transport shim shared by the agent adapters: given a request body, returns * the raw SSE byte stream from a serving / AI-gateway endpoint. Injected at @@ -108,11 +121,9 @@ export async function invoke( * POSTs `body` as JSON to an arbitrary workspace API path and returns the raw * SSE byte stream. No parsing is performed — bytes are passed through as-is. * - * Uses the client's `request` (modular transport) when available, else the - * legacy SDK's `apiClient.request({ raw: true })`, so callers inherit URL - * resolution and the SDK credential chain (PAT/OAuth/OIDC). - * - * When `signal` is provided it aborts the outbound HTTP request. + * Uses the workspace client's `request()` so callers inherit URL resolution, + * the credential chain (PAT/OAuth/OIDC), and the AppKit User-Agent. A non-2xx + * status throws `ApiError`. `signal` aborts the outbound HTTP request. * * @internal * @@ -124,14 +135,14 @@ export async function invoke( * `beta.ts` or any other entry point. */ export async function streamPath( - client: ApiClientLike, + client: StreamClientLike, path: string, body: Record, signal?: AbortSignal, ): Promise> { logger.debug("Streaming from path %s", path); - if (client.request) { + if ("request" in client) { const response = await client.request({ method: "POST", path, @@ -177,14 +188,14 @@ export async function streamPath( * `stream: true` in the payload. */ export async function stream( - client: WorkspaceClient, + client: StreamClientLike, endpointName: string, body: Record, signal?: AbortSignal, ): Promise> { const { stream: _stream, ...cleanBody } = body; return streamPath( - client as unknown as ApiClientLike, + client, `/serving-endpoints/${encodeURIComponent(endpointName)}/invocations`, { ...cleanBody, stream: true }, signal, @@ -201,7 +212,7 @@ export async function stream( * fixed path that routes by the body's `model`, so the caller sets it there. */ export async function streamAiGateway( - client: ApiClientLike, + client: StreamClientLike, body: Record, signal?: AbortSignal, ): Promise> { diff --git a/packages/appkit/src/core/agent/run-agent.ts b/packages/appkit/src/core/agent/run-agent.ts index 847dd4f1a..ed7290191 100644 --- a/packages/appkit/src/core/agent/run-agent.ts +++ b/packages/appkit/src/core/agent/run-agent.ts @@ -10,7 +10,6 @@ import type { ToolProvider, } from "shared"; -import type { WorkspaceClientLike } from "../../agents/databricks"; import { isSupervisorTool, SUPERVISOR_EXTENSION_KEY, @@ -357,8 +356,7 @@ async function resolveAdapter( // On behalf of the user the model client is the caller's, per call. return onBehalfOfUser ? adapterFromModelString(source, { - workspaceClient: () => - getWorkspaceClient() as unknown as WorkspaceClientLike, + workspaceClient: () => getWorkspaceClient(), }) : adapterFromModelString(source); }