feat(appkit): introduce workspace-client wrapper over the Databricks SDK

biome.json

@@ -41,6 +41,22 @@
       },
       "security": {
         "noDangerouslySetInnerHtml": "off"
+      },
+      "style": {
+        "noRestrictedImports": {
+          "level": "error",
+          "options": {
+            "patterns": [
+              {
+                "group": [
+                  "@databricks/sdk-experimental",
+                  "@databricks/sdk-experimental/**"
+                ],
+                "message": "Import the Databricks SDK only through the wrapper in packages/appkit/src/workspace-client. Add a re-export there if you need a new symbol."
+              }
+            ]
+          }
+        }
       }
     }
   },
@@ -56,5 +72,17 @@
         "organizeImports": "on"
       }
     }
-  }
+  },
+  "overrides": [
+    {
+      "includes": ["packages/appkit/src/workspace-client/**"],
+      "linter": {
+        "rules": {
+          "style": {
+            "noRestrictedImports": "off"
+          }
+        }
+      }
+    }
+  ]
 }

packages/appkit/src/agents/databricks.ts

@@ -7,6 +7,7 @@ import type {
 } from "shared";
 import { stream as servingStream } from "../connectors/serving/client";
 import { APPKIT_USER_AGENT, getClientOptions } from "../context/client-options";
+import { createWorkspaceClient } from "../workspace-client";
 
 /** Default cap for a single incomplete SSE line tail (DoS guard). */
 const DEFAULT_MAX_SSE_LINE_CHARS = 1024 * 1024;
@@ -203,12 +204,11 @@ interface DeltaToolCall {
  *
  * @example Using the factory (recommended)
  * ```ts
- * import { createApp, createAgent, agents } from "@databricks/appkit";
+ * import { createApp, createAgent, agents, createWorkspaceClient } from "@databricks/appkit";
  * import { DatabricksAdapter } from "@databricks/appkit/beta";
- * import { WorkspaceClient } from "@databricks/sdk-experimental";
  *
  * const adapter = DatabricksAdapter.fromServingEndpoint({
- *   workspaceClient: new WorkspaceClient({}),
+ *   workspaceClient: createWorkspaceClient(),
  *   endpointName: "my-endpoint",
  * });
  *
@@ -358,11 +358,9 @@ export class DatabricksAdapter implements AgentAdapter {
     let workspaceClient: WorkspaceClientLike | undefined =
       options?.workspaceClient;
     if (!workspaceClient) {
-      const sdk = await import("@databricks/sdk-experimental");
-      workspaceClient = new sdk.WorkspaceClient(
-        {},
-        getClientOptions(),
-      ) as unknown as WorkspaceClientLike;
+      workspaceClient = createWorkspaceClient({
+        clientOptions: getClientOptions(),
+      }) as unknown as WorkspaceClientLike;
     }
 
     return DatabricksAdapter.fromServingEndpoint({

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

@@ -944,11 +944,16 @@ describe("DatabricksAdapter.fromModelServing", () => {
   test("reads endpoint from DATABRICKS_SERVING_ENDPOINT_NAME env var", async () => {
     process.env.DATABRICKS_SERVING_ENDPOINT_NAME = "my-model";
 
-    vi.mock("@databricks/sdk-experimental", () => ({
-      WorkspaceClient: vi.fn().mockImplementation(() => ({
-        apiClient: { request: vi.fn() },
-      })),
-    }));
+    vi.mock("../../workspace-client", async (importOriginal) => {
+      const actual =
+        await importOriginal<typeof import("../../workspace-client")>();
+      return {
+        ...actual,
+        createWorkspaceClient: vi.fn().mockImplementation(() => ({
+          apiClient: { request: vi.fn() },
+        })),
+      };
+    });
 
     const adapter = await DatabricksAdapter.fromModelServing();
     expect(adapter).toBeInstanceOf(DatabricksAdapter);

packages/appkit/src/cache/index.ts

@@ -1,5 +1,4 @@
 import { createHash } from "node:crypto";
-import { ApiError, WorkspaceClient } from "@databricks/sdk-experimental";
 import type { CacheConfig, CacheEntry, CacheStorage } from "shared";
 import { createLakebasePool } from "../connectors/lakebase";
 import { getClientOptions } from "../context/client-options";
@@ -8,6 +7,7 @@ import { createLogger } from "../logging/logger";
 import type { Counter, TelemetryProvider } from "../telemetry";
 import { SpanStatusCode, TelemetryManager } from "../telemetry";
 import { deepMerge } from "../utils";
+import { ApiError, createWorkspaceClient } from "../workspace-client";
 import { cacheDefaults } from "./defaults";
 import { InMemoryStorage, PersistentStorage } from "./storage";
 
@@ -171,7 +171,9 @@ export class CacheManager {
 
     // try to use lakebase storage
     try {
-      const workspaceClient = new WorkspaceClient({}, getClientOptions());
+      const workspaceClient = createWorkspaceClient({
+        clientOptions: getClientOptions(),
+      }).toLegacyWorkspaceClient();
       const pool = createLakebasePool({ workspaceClient });
       const persistentStorage = new PersistentStorage(config, pool);