From 013d3634f7c77f5a23158c252b157daf4fe3161d Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Fri, 9 Oct 2026 17:43:38 +0200 Subject: [PATCH 1/2] feat(shared): add a modular experiments client to the workspace client Add buildExperimentsClient (@databricks/sdk-experiments 0.52.0, pinned) and a lazily built `experiments` accessor on the facade, so the CLI can list MLflow experiments through a typed client instead of raw request(). Co-authored-by: Isaac Signed-off-by: MarioCadenas --- knip.json | 1 + packages/appkit/package.json | 1 + .../src/testing/mock-workspace-client.ts | 1 + packages/shared/package.json | 1 + .../shared/src/workspace-client/client.ts | 11 ++++++ .../shared/src/workspace-client/modular.ts | 18 +++++++++ .../workspace-client/tests/modular.test.ts | 37 +++++++++++++++++++ packages/shared/src/workspace-client/types.ts | 4 ++ pnpm-lock.yaml | 19 ++++++++++ 9 files changed, 93 insertions(+) diff --git a/knip.json b/knip.json index 639990671..6963fc6a9 100644 --- a/knip.json +++ b/knip.json @@ -15,6 +15,7 @@ "@databricks/sdk-experimental", "@databricks/sdk-files", "@databricks/sdk-database", + "@databricks/sdk-experiments", "@databricks/sdk-genie", "@databricks/sdk-uc-functions", "@databricks/sdk-jobs", diff --git a/packages/appkit/package.json b/packages/appkit/package.json index 692478d7e..041761741 100644 --- a/packages/appkit/package.json +++ b/packages/appkit/package.json @@ -76,6 +76,7 @@ "@databricks/sdk-experimental": "0.17.0", "@databricks/sdk-database": "0.51.0", "@databricks/sdk-files": "0.52.0", + "@databricks/sdk-experiments": "0.52.0", "@databricks/sdk-genie": "0.54.0", "@databricks/sdk-uc-functions": "0.51.0", "@databricks/sdk-jobs": "0.57.0", diff --git a/packages/appkit/src/testing/mock-workspace-client.ts b/packages/appkit/src/testing/mock-workspace-client.ts index 68b35ef5e..27a47048a 100644 --- a/packages/appkit/src/testing/mock-workspace-client.ts +++ b/packages/appkit/src/testing/mock-workspace-client.ts @@ -70,6 +70,7 @@ const FACADE_SERVICES = [ "genie", "database", "jobs", + "experiments", "vectorSearch", "tables", "postgres", diff --git a/packages/shared/package.json b/packages/shared/package.json index 47d4a63ca..844833fd2 100644 --- a/packages/shared/package.json +++ b/packages/shared/package.json @@ -53,6 +53,7 @@ "@databricks/sdk-experimental": "0.17.0", "@databricks/sdk-database": "0.51.0", "@databricks/sdk-files": "0.52.0", + "@databricks/sdk-experiments": "0.52.0", "@databricks/sdk-genie": "0.54.0", "@databricks/sdk-uc-functions": "0.51.0", "@databricks/sdk-jobs": "0.57.0", diff --git a/packages/shared/src/workspace-client/client.ts b/packages/shared/src/workspace-client/client.ts index bb61670aa..335a4aeab 100644 --- a/packages/shared/src/workspace-client/client.ts +++ b/packages/shared/src/workspace-client/client.ts @@ -25,6 +25,8 @@ import { type FunctionsClient, type GenieClient, buildJobsClient, + buildExperimentsClient, + type ExperimentsClient, type JobsClient, buildModelServingClient, type ModelServingClient, @@ -58,6 +60,7 @@ export class AppKitWorkspaceClient implements WorkspaceClient { #database?: DatabaseClient; #jobs?: JobsClient; #modelServing?: ModelServingClient; + #experiments?: ExperimentsClient; #files?: FilesClient; #vectorSearch?: VectorSearchClient; #tables?: TablesClient; @@ -110,6 +113,14 @@ export class AppKitWorkspaceClient implements WorkspaceClient { return this.#jobs; } + // Migrated to the modular SDK — built lazily, independent of the legacy client. + get experiments(): ExperimentsClient { + if (!this.#experiments) { + this.#experiments = buildExperimentsClient(this.#opts); + } + return this.#experiments; + } + // Migrated to the modular SDK — built lazily, independent of the legacy client. get vectorSearch(): VectorSearchClient { if (!this.#vectorSearch) { diff --git a/packages/shared/src/workspace-client/modular.ts b/packages/shared/src/workspace-client/modular.ts index 3e5e4bcfb..a8b322f12 100644 --- a/packages/shared/src/workspace-client/modular.ts +++ b/packages/shared/src/workspace-client/modular.ts @@ -53,6 +53,7 @@ import { } from "@databricks/sdk-core/http"; import { type ProfileOptions, resolve } from "@databricks/sdk-core/profiles"; import type { DatabaseClient } from "@databricks/sdk-database/v1"; +import type { ExperimentsClient } from "@databricks/sdk-experiments/v1"; import type { FilesClient } from "@databricks/sdk-files/v2"; import type { GenieClient } from "@databricks/sdk-genie/v1"; import type { JobsClient } from "@databricks/sdk-jobs/v2"; @@ -634,6 +635,17 @@ export function buildModelServingClient( ); } +/** Build a modular MLflow Experiments client from wrapper options. */ +export function buildExperimentsClient( + opts: WorkspaceClientOptions, +): ExperimentsClient { + return lazyClient( + () => + import("@databricks/sdk-experiments/v1").then((m) => m.ExperimentsClient), + mapToClientOptions(opts), + ); +} + /** Build a modular Unity Catalog Tables client from wrapper options. */ export function buildTablesClient(opts: WorkspaceClientOptions): TablesClient { return lazyClient( @@ -702,6 +714,7 @@ export function buildPostgresClient( // ── Client type re-exports (for the facade accessor types) ─────────────── export type { DatabaseClient } from "@databricks/sdk-database/v1"; +export type { ExperimentsClient } from "@databricks/sdk-experiments/v1"; export type { FilesClient } from "@databricks/sdk-files/v2"; export type { DirectoryEntry as FilesDirectoryEntry, @@ -761,6 +774,11 @@ export type { FunctionInfo, GetFunctionRequest, } from "@databricks/sdk-uc-functions/v1"; +export type { + Experiment, + ListExperimentsRequest, + ListExperimentsResponse, +} from "@databricks/sdk-experiments/v1"; export type { BaseRun, CancelRunRequest, diff --git a/packages/shared/src/workspace-client/tests/modular.test.ts b/packages/shared/src/workspace-client/tests/modular.test.ts index 355836239..e8a938375 100644 --- a/packages/shared/src/workspace-client/tests/modular.test.ts +++ b/packages/shared/src/workspace-client/tests/modular.test.ts @@ -42,6 +42,12 @@ vi.mock("@databricks/sdk-warehouses/v1", () => ({ vi.mock("@databricks/sdk-statementexecution/v1", () => ({ StatementExecutionClient: vi.fn().mockImplementation((opts) => ({ opts })), })); +vi.mock("@databricks/sdk-experiments/v1", () => ({ + ExperimentsClient: vi.fn().mockImplementation((opts) => { + ctorOpts.push(opts); + return { opts, getExperiment: async () => ({}) }; + }), +})); vi.mock("@databricks/sdk-genie/v1", () => ({ GenieClient: vi.fn().mockImplementation((opts) => { ctorOpts.push(opts); @@ -157,6 +163,7 @@ import { buildFunctionsClient, buildConnectionsClient, buildDatabaseClient, + buildExperimentsClient, buildGenieClient, buildModelServingClient, buildVolumesClient, @@ -531,6 +538,36 @@ describe("modular mapToClientOptions (via buildWarehousesClient)", () => { expect(m2mOpts).toEqual([]); }); + test("experiments normalizes the env host and resolves the SP via M2M", async () => { + process.env.DATABRICKS_HOST = "x.cloud.databricks.com"; + process.env.DATABRICKS_CLIENT_ID = "sp-client-id"; + process.env.DATABRICKS_CLIENT_SECRET = "sp-secret"; + await buildExperimentsClient({}).getExperiment({} as never); + expect(ctorOpts[0].host).toBe("https://x.cloud.databricks.com"); + expect(m2mOpts).toEqual([ + { + host: "https://x.cloud.databricks.com", + clientId: "sp-client-id", + clientSecret: "sp-secret", + }, + ]); + expect(ctorOpts[0].credentials).toBeDefined(); + }); + + test("experiments (asUser) uses the OBO token as PAT", async () => { + process.env.DATABRICKS_CLIENT_ID = "sp-client-id"; + process.env.DATABRICKS_CLIENT_SECRET = "sp-secret"; + await buildExperimentsClient({ + token: "user-token", + host: "https://x", + }).getExperiment({} as never); + expect(ctorOpts[0].credentials).toMatchObject({ + kind: "pat", + token: "user-token", + }); + expect(m2mOpts).toEqual([]); + }); + test("postgres normalizes the host and passes SP M2M credentials", async () => { process.env.DATABRICKS_HOST = "envhost.cloud.databricks.com"; process.env.DATABRICKS_CLIENT_ID = "sp-client-id"; diff --git a/packages/shared/src/workspace-client/types.ts b/packages/shared/src/workspace-client/types.ts index c4024d86a..95fa14523 100644 --- a/packages/shared/src/workspace-client/types.ts +++ b/packages/shared/src/workspace-client/types.ts @@ -23,6 +23,7 @@ import type { DatabaseClient, JobsClient, ModelServingClient, + ExperimentsClient, FilesClient, TablesClient, VolumesClient, @@ -68,6 +69,9 @@ export interface WorkspaceClient extends WorkspaceAuth { /** Jobs (modular SDK, Jobs API 2.2). */ readonly jobs: JobsClient; + /** MLflow Experiments (modular SDK). */ + readonly experiments: ExperimentsClient; + /** Vector Search (modular SDK; index metadata only, queries use `request()`). */ readonly vectorSearch: VectorSearchClient; diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index b67e4d8ec..64e02c17e 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -293,6 +293,9 @@ importers: '@databricks/sdk-experimental': specifier: 0.17.0 version: 0.17.0 + '@databricks/sdk-experiments': + specifier: 0.52.0 + version: 0.52.0 '@databricks/sdk-files': specifier: 0.52.0 version: 0.52.0 @@ -662,6 +665,9 @@ importers: '@databricks/sdk-experimental': specifier: 0.17.0 version: 0.17.0 + '@databricks/sdk-experiments': + specifier: 0.52.0 + version: 0.52.0 '@databricks/sdk-files': specifier: 0.52.0 version: 0.52.0 @@ -2089,6 +2095,10 @@ packages: resolution: {integrity: sha512-dOJIt4F2nBk6HKObnv7Xbmy/qLYTy2835qhXSuW0Qw1QAXui9plmCet1KqG3yeQcMTyncWGbnhjGdQi8GEGQSA==} engines: {node: '>=22.0', npm: '>=10.0.0'} + '@databricks/sdk-experiments@0.52.0': + resolution: {integrity: sha512-RVpS4fMtfTKI/7WdyeZIdLdy7lBV3liO/W/dLH+zNwga7REbp581zAxZNaomCYJG/y7SrqrtMo77CVDDKH4AKg==} + engines: {node: '>=22.0.0'} + '@databricks/sdk-files@0.52.0': resolution: {integrity: sha512-H3/uAkkl7EI+60/q1FRqtKrG7UdJfWKLug4/87q+P7OXJMF3J8kgECIg1LfU2Ih2Qj42LNhqdjos7LJSC68gyw==} engines: {node: '>=22.0.0'} @@ -14450,6 +14460,15 @@ snapshots: transitivePeerDependencies: - supports-color + '@databricks/sdk-experiments@0.52.0': + dependencies: + '@databricks/sdk-auth': 0.51.0 + '@databricks/sdk-core': 0.51.0 + '@databricks/sdk-options': 0.51.0 + '@js-temporal/polyfill': 0.5.1 + json-bigint: 1.0.0 + zod: 4.3.6 + '@databricks/sdk-files@0.52.0': dependencies: '@databricks/sdk-auth': 0.51.0 From c2f1d6aae39747bde8c4239cb6c4225007b30e67 Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Thu, 15 Oct 2026 18:10:22 +0200 Subject: [PATCH 2/2] fix(appkit): add .js to @mlflow/core deep imports so agent tracing loads under Node ESM @mlflow/core 0.4.0 has no exports map, so Node ESM rejects the extensionless deep imports (ERR_MODULE_NOT_FOUND). UC auto-detect fell back to classic, then the exporter import failed and tracing was disabled behind a misleading "requires the optional @mlflow/core package" warning. vitest adds extensions, so unit tests never caught it. Only warn about the missing package when the bare import fails; a failed deep import now logs its real error. Add a test that resolves the deep imports in a child node process. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- packages/appkit/src/plugins/agents/mlflow.ts | 12 ++-- .../src/plugins/agents/tests/mlflow.test.ts | 65 ++++++++++++++++--- 2 files changed, 64 insertions(+), 13 deletions(-) diff --git a/packages/appkit/src/plugins/agents/mlflow.ts b/packages/appkit/src/plugins/agents/mlflow.ts index ac985241c..19ecbf5a0 100644 --- a/packages/appkit/src/plugins/agents/mlflow.ts +++ b/packages/appkit/src/plugins/agents/mlflow.ts @@ -227,7 +227,7 @@ async function resolveUcLocation( const experiment = await client.getExperiment(experimentId); if (!experiment) return undefined; const { ucLocationFromExperimentTags } = - await import("@mlflow/core/dist/core/destination"); + await import("@mlflow/core/dist/core/destination.js"); return ucLocationFromExperimentTags(experiment.tags) ?? undefined; } catch (err) { logger.warn( @@ -272,14 +272,14 @@ async function buildMlflowSpanProcessor( let processor: SpanProcessor; if (ucLoc) { const { DatabricksUCTableSpanExporter, DatabricksUCTableSpanProcessor } = - await import("@mlflow/core/dist/exporters/uc_table"); + await import("@mlflow/core/dist/exporters/uc_table.js"); processor = new DatabricksUCTableSpanProcessor( new DatabricksUCTableSpanExporter(client), ucLoc, ); } else { const { MlflowSpanExporter, MlflowSpanProcessor } = - await import("@mlflow/core/dist/exporters/mlflow"); + await import("@mlflow/core/dist/exporters/mlflow.js"); processor = new MlflowSpanProcessor(new MlflowSpanExporter(client)); } return { @@ -346,8 +346,10 @@ export async function initAgentTracing(): Promise { const id = experimentId(); if (!id) return; + let coreLoaded = false; try { mlflow = await import("@mlflow/core"); + coreLoaded = true; const host = normalizedDatabricksHost(); initConfig = { trackingUri: process.env.MLFLOW_TRACKING_URI?.trim() || "databricks", @@ -385,7 +387,9 @@ export async function initAgentTracing(): Promise { logger.info("MLflow agent tracing enabled (experiment %s)", id); } } catch (err) { - if (isModuleNotFound(err)) { + // Only the bare package being absent means "not installed"; a failed deep + // import is a real error and must surface as one. + if (!coreLoaded && isModuleNotFound(err)) { logger.warn( "MLflow agent tracing requires the optional `@mlflow/core` package. " + "Install it to enable tracing: npm i @mlflow/core", diff --git a/packages/appkit/src/plugins/agents/tests/mlflow.test.ts b/packages/appkit/src/plugins/agents/tests/mlflow.test.ts index 6747c2b34..b8e19d215 100644 --- a/packages/appkit/src/plugins/agents/tests/mlflow.test.ts +++ b/packages/appkit/src/plugins/agents/tests/mlflow.test.ts @@ -1,3 +1,9 @@ +import { execFileSync } from "node:child_process"; +import { existsSync, readFileSync } from "node:fs"; +import { createRequire } from "node:module"; +import { dirname, join } from "node:path"; +import { fileURLToPath } from "node:url"; + import { SpanKind } from "@opentelemetry/api"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; @@ -57,14 +63,14 @@ function stubSdk( return Promise.resolve(); } } - vi.doMock("@mlflow/core/dist/exporters/mlflow", () => ({ + vi.doMock("@mlflow/core/dist/exporters/mlflow.js", () => ({ MlflowSpanExporter: class {}, MlflowSpanProcessor: StubProcessor, })); // Distinct from the classic stub so a test can tell which processor was built // and assert the resolved UC location actually reaches its constructor. const uc: { built: boolean; location?: unknown } = { built: false }; - vi.doMock("@mlflow/core/dist/exporters/uc_table", () => ({ + vi.doMock("@mlflow/core/dist/exporters/uc_table.js", () => ({ DatabricksUCTableSpanExporter: class {}, DatabricksUCTableSpanProcessor: class { constructor(_exporter: unknown, location: unknown) { @@ -83,7 +89,7 @@ function stubSdk( })); // mlflow's experiment-tag → UC-location parser (deep import), used by the // numeric-id auto-detect path. Default: not a UC experiment (null → classic). - vi.doMock("@mlflow/core/dist/core/destination", () => ({ + vi.doMock("@mlflow/core/dist/core/destination.js", () => ({ ucLocationFromExperimentTags: vi.fn(() => ucFromTags), })); return { sdk, span, setInputs, setOutputs, uc }; @@ -100,9 +106,9 @@ describe("agent tracing (mlflow)", () => { afterEach(() => { vi.doUnmock("@mlflow/core"); - vi.doUnmock("@mlflow/core/dist/exporters/mlflow"); - vi.doUnmock("@mlflow/core/dist/exporters/uc_table"); - vi.doUnmock("@mlflow/core/dist/core/destination"); + vi.doUnmock("@mlflow/core/dist/exporters/mlflow.js"); + vi.doUnmock("@mlflow/core/dist/exporters/uc_table.js"); + vi.doUnmock("@mlflow/core/dist/core/destination.js"); vi.doUnmock("../../../telemetry"); vi.restoreAllMocks(); delete process.env.MLFLOW_EXPERIMENT_ID; @@ -367,6 +373,47 @@ describe("agent tracing (mlflow)", () => { }); }); + // vitest resolves extensionless specifiers, real Node ESM does not (no + // `exports` map in @mlflow/core), so resolve mlflow.ts's deep imports in a + // child `node` process. Regression: extensionless paths broke tracing in prod. + const appkitRoot = join( + dirname(fileURLToPath(import.meta.url)), + "../../../..", + ); + const mlflowInstalled = (() => { + try { + createRequire(join(appkitRoot, "package.json")).resolve("@mlflow/core"); + return true; + } catch { + return false; + } + })(); + test.skipIf(!mlflowInstalled)( + "@mlflow/core deep imports resolve under real Node ESM", + () => { + const source = readFileSync( + join(appkitRoot, "src/plugins/agents/mlflow.ts"), + "utf8", + ); + const specifiers = [ + ...source.matchAll(/import\("(@mlflow\/core\/[^"]+)"\)/g), + ].map((m) => m[1]); + expect(specifiers).toHaveLength(3); + for (const spec of specifiers) { + const url = execFileSync( + process.execPath, + [ + "--input-type=module", + "-e", + `process.stdout.write(import.meta.resolve(${JSON.stringify(spec)}))`, + ], + { cwd: appkitRoot, encoding: "utf8" }, + ); + expect(existsSync(fileURLToPath(url))).toBe(true); + } + }, + ); + // Tripwire: fails loudly if a @mlflow/core version bump moves or renames the // deep-imported exporter/processor classes buildMlflowSpanProcessor relies on, // or moves the public symbols we now import from the entrypoint. Runs against @@ -381,11 +428,11 @@ describe("agent tracing (mlflow)", () => { SpanType, } = mlflow; const { MlflowSpanExporter, MlflowSpanProcessor } = - await import("@mlflow/core/dist/exporters/mlflow"); + await import("@mlflow/core/dist/exporters/mlflow.js"); const { DatabricksUCTableSpanExporter, DatabricksUCTableSpanProcessor } = - await import("@mlflow/core/dist/exporters/uc_table"); + await import("@mlflow/core/dist/exporters/uc_table.js"); const { ucLocationFromExperimentTags } = - await import("@mlflow/core/dist/core/destination"); + await import("@mlflow/core/dist/core/destination.js"); // Public entrypoint symbols (deep imports under mlflow-tracing 0.1.3). expect(typeof createAuthProvider).toBe("function");