feat(shared): add a modular Experiments client to the workspace client

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",

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",

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<void> {
   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<void> {
       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",

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");

packages/appkit/src/testing/mock-workspace-client.ts

@@ -70,6 +70,7 @@ const FACADE_SERVICES = [
   "genie",
   "database",
   "jobs",
+  "experiments",
   "vectorSearch",
   "tables",
   "postgres",

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",

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) {

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,

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";

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;