@dbx-tools/appkit-mastra 0.3.21 → 0.3.23

package/package.json

@@ -28,22 +28,22 @@
     "express": "^5.1.0",
     "pg": "^8.22.0",
     "zod": "^4.3.6",
-    "@dbx-tools/core": "0.3.21",
-    "@dbx-tools/appkit": "0.3.21",
-    "@dbx-tools/databricks": "0.3.21",
-    "@dbx-tools/shared-genie": "0.3.21",
-    "@dbx-tools/model": "0.3.21",
-    "@dbx-tools/shared-core": "0.3.21",
-    "@dbx-tools/genie": "0.3.21",
-    "@dbx-tools/shared-mastra": "0.3.21",
-    "@dbx-tools/shared-model": "0.3.21"
+    "@dbx-tools/appkit": "0.3.23",
+    "@dbx-tools/genie": "0.3.23",
+    "@dbx-tools/databricks": "0.3.23",
+    "@dbx-tools/core": "0.3.23",
+    "@dbx-tools/model": "0.3.23",
+    "@dbx-tools/shared-core": "0.3.23",
+    "@dbx-tools/shared-mastra": "0.3.23",
+    "@dbx-tools/shared-genie": "0.3.23",
+    "@dbx-tools/shared-model": "0.3.23"
   },
   "main": "index.ts",
   "license": "UNLICENSED",
   "publishConfig": {
     "access": "public"
   },
-  "version": "0.3.21",
+  "version": "0.3.23",
   "types": "index.ts",
   "type": "module",
   "exports": {

package/src/agents.ts

@@ -15,6 +15,9 @@
  * @module
  */
 
+import { plugin } from "@dbx-tools/appkit";
+import { fallback } from "@dbx-tools/model";
+import { log, object, string } from "@dbx-tools/shared-core";
 import type { AgentConfig, ToolsInput } from "@mastra/core/agent";
 import { Agent } from "@mastra/core/agent";
 import type { Tool } from "@mastra/core/tools";
@@ -27,12 +30,9 @@ import type { MastraPluginConfig } from "./config";
 import { buildGenieToolkitProvider, resolveGenieSpaces } from "./genie";
 import type { MemoryBuilder } from "./memory";
 import { buildModel } from "./model";
-import { fallback } from "@dbx-tools/model";
 import { ResultProcessor, stripStaleChartsProcessor } from "./processors";
 import { buildSummarizeTool } from "./summarize";
 import { createWorkspace } from "./workspaces";
-import { log, object, string } from "@dbx-tools/shared-core";
-import { plugin } from "@dbx-tools/appkit";
 
 /**
  * Tool record accepted by every Mastra `Agent.tools` field and by the

package/src/chart.ts

@@ -28,7 +28,9 @@
  */
 
 import { CacheManager } from "@databricks/appkit";
+import { async, error, hash, log, string, type BrandContext } from "@dbx-tools/shared-core";
 import { wire, type Chart, type ChartResult } from "@dbx-tools/shared-mastra";
+import { model } from "@dbx-tools/shared-model";
 import { Agent } from "@mastra/core/agent";
 import type { RequestContext } from "@mastra/core/request-context";
 import { createTool } from "@mastra/core/tools";
@@ -36,8 +38,6 @@ import { z } from "zod";
 
 import type { MastraPluginConfig } from "./config";
 import { buildModel } from "./model";
-import { model } from "@dbx-tools/shared-model";
-import { async, error, hash, log, string, type BrandContext } from "@dbx-tools/shared-core";
 
 const logger = log.logger("mastra/chart");
 

package/src/config.ts

@@ -8,6 +8,7 @@
  */
 
 import type { BasePluginConfig } from "@databricks/appkit";
+import { appkit } from "@dbx-tools/appkit";
 import type { BrandContext } from "@dbx-tools/shared-core";
 import type { AgentConfig } from "@mastra/core/agent";
 import { MASTRA_RESOURCE_ID_KEY, MASTRA_THREAD_ID_KEY } from "@mastra/core/request-context";
@@ -15,7 +16,6 @@ import type { PgVectorConfig, PostgresStoreConfig } from "@mastra/pg";
 
 import type { MastraAgentDefinition, MastraTools } from "./agents";
 import type { GenieSpacesConfig } from "./genie";
-import { appkit } from "@dbx-tools/appkit";
 
 /**
  * `RequestContext` key under which {@link MastraServer} stores the

package/src/filesystems.ts

@@ -13,8 +13,10 @@
  * @module
  */
 
+import { posix as path } from "node:path";
 import { getExecutionContext } from "@databricks/appkit";
 import { WorkspaceClient } from "@databricks/sdk-experimental";
+import { error, functionModule, hash, log } from "@dbx-tools/shared-core";
 import type {
   CopyOptions,
   FileContent,
@@ -26,8 +28,6 @@ import type {
   ReadOptions,
   RemoveOptions,
   WriteOptions,
-} from "@mastra/core/workspace";
-import {
   DirectoryNotEmptyError,
   DirectoryNotFoundError,
   FileExistsError,
@@ -39,8 +39,6 @@ import {
   WorkspaceReadOnlyError,
   type MastraFilesystemOptions,
 } from "@mastra/core/workspace";
-import { posix as path } from "node:path";
-import { error, functionModule, hash, log } from "@dbx-tools/shared-core";
 
 /* ------------------------------ constants ------------------------------ */
 
@@ -764,7 +762,7 @@ export class DatabricksWorkspaceFilesystem extends MastraFilesystem {
       const entries = await this.listAbsoluteDirectory(absolutePath);
       const filtered = this.filterEntries(entries, options);
       if (!options?.recursive) return filtered;
-      return this.readDirectoryRecursive(absolutePath, options, 0);
+      return await this.readDirectoryRecursive(absolutePath, options, 0);
     } catch (err) {
       this.rethrow(err, inputPath);
     }

package/src/genie.ts

@@ -38,11 +38,11 @@
 
 import { CacheManager, genie } from "@databricks/appkit";
 import { WorkspaceClient } from "@databricks/sdk-experimental";
-import { wire, type MastraWriter, type StartedEvent } from "@dbx-tools/shared-mastra";
+import { plugin } from "@dbx-tools/appkit";
 import { chat, space as genieSpace } from "@dbx-tools/genie";
-import { genieModel, type GenieMessage } from "@dbx-tools/shared-genie";
 import { error, log, string } from "@dbx-tools/shared-core";
-import { plugin } from "@dbx-tools/appkit";
+import { genieModel, type GenieMessage } from "@dbx-tools/shared-genie";
+import { wire, type MastraWriter, type StartedEvent } from "@dbx-tools/shared-mastra";
 import type { RequestContext } from "@mastra/core/request-context";
 import { MASTRA_THREAD_ID_KEY } from "@mastra/core/request-context";
 import { createTool } from "@mastra/core/tools";
@@ -50,8 +50,7 @@ import { z } from "zod";
 
 import type { MastraTools } from "./agents";
 import { chartPlannerRequestSchema, prepareChart } from "./chart";
-import type { MastraPluginConfig } from "./config";
-import { MASTRA_USER_KEY, type User } from "./config";
+import type { MastraPluginConfig, MASTRA_USER_KEY, type User } from "./config";
 import { fetchStatementData } from "./statement";
 import { safeWrite } from "./writer";
 
@@ -918,7 +917,7 @@ export function resolveGenieSpaces(
   const merged: Record<string, GenieSpaceConfig> = {};
 
   // Source 3 (lowest precedence): env var.
-  const envSpaceId = process.env["DATABRICKS_GENIE_SPACE_ID"];
+  const envSpaceId = process.env.DATABRICKS_GENIE_SPACE_ID;
   if (envSpaceId) {
     merged[DEFAULT_GENIE_ALIAS] = { spaceId: envSpaceId };
   }

package/src/history.ts

@@ -18,6 +18,7 @@
  * @module
  */
 
+import { error, log } from "@dbx-tools/shared-core";
 import type {
   MastraClearHistoryResponse,
   MastraHistoryResponse,
@@ -31,7 +32,6 @@ import type { ContextWithMastra } from "@mastra/core/server";
 import { registerApiRoute } from "@mastra/core/server";
 
 import { clampPerPage, parseIntParam } from "./pagination";
-import { error, log } from "@dbx-tools/shared-core";
 
 const logger = log.logger("mastra/history");
 

package/src/memory.ts

@@ -29,18 +29,18 @@
  * @module
  */
 
+import { randomUUID } from "node:crypto";
 import { getUsernameWithApiLookup } from "@databricks/appkit";
+import { log } from "@dbx-tools/shared-core";
 import { fastembed } from "@mastra/fastembed";
 import { Memory } from "@mastra/memory";
 import { PgVector, PostgresStore } from "@mastra/pg";
-import { randomUUID } from "node:crypto";
 import { Pool, type PoolConfig } from "pg";
 
 import type { MastraAgentDefinition, MastraMemoryConfigOverride } from "./agents";
 import type { MastraPluginConfig } from "./config";
 import { agentStorageSchemaName } from "./storage-schema";
 import { summaryModel, TITLE_INSTRUCTIONS } from "./summarize";
-import { log } from "@dbx-tools/shared-core";
 
 const logger = log.logger("mastra/memory");
 

package/src/mlflow.ts

@@ -19,10 +19,10 @@
  * @module
  */
 
+import { appkit } from "@dbx-tools/appkit";
+import { async, error, log } from "@dbx-tools/shared-core";
 import { feedback } from "@dbx-tools/shared-mastra";
 import { databricksFetch, readResponseJson, readResponseText } from "./rest";
-import { async, error, log } from "@dbx-tools/shared-core";
-import { appkit } from "@dbx-tools/appkit";
 
 const logger = log.logger("mastra/mlflow");
 

package/src/model.ts

@@ -25,14 +25,14 @@
 
 import { getExecutionContext } from "@databricks/appkit";
 import { classes, resolve } from "@dbx-tools/model";
+import { functionModule, log, net } from "@dbx-tools/shared-core";
 import { model } from "@dbx-tools/shared-model";
 import type { MastraModelConfig } from "@mastra/core/llm";
 import type { RequestContext } from "@mastra/core/request-context";
 
 import { MASTRA_USER_KEY, type MastraPluginConfig, type User } from "./config";
-import { rewriteServingBody } from "./serving-sanitize";
 import { MASTRA_MODEL_OVERRIDE_KEY, resolveServingConfig } from "./serving";
-import { functionModule, log, net } from "@dbx-tools/shared-core";
+import { rewriteServingBody } from "./serving-sanitize";
 
 type ModelClass = model.ModelClass;
 const { parseModelClass } = classes;

package/src/observability.ts

@@ -33,12 +33,12 @@
  * @module
  */
 
+import { project } from "@dbx-tools/core";
+import { log } from "@dbx-tools/shared-core";
 import { Observability } from "@mastra/observability";
 import { OtelBridge } from "@mastra/otel-bridge";
 
 import { TRACE_REQUEST_CONTEXT_KEYS } from "./config";
-import { log } from "@dbx-tools/shared-core";
-import { project } from "@dbx-tools/core";
 
 const logger = log.logger("mastra/observability");
 

package/src/plugin.ts

@@ -51,11 +51,9 @@ import {
   type PluginManifest,
   type ResourceRequirement,
 } from "@databricks/appkit";
-import type { Agent } from "@mastra/core/agent";
-import { Mastra } from "@mastra/core/mastra";
-import express from "express";
-import type { Pool } from "pg";
-
+import { plugin } from "@dbx-tools/appkit";
+import { serving as nodeServing } from "@dbx-tools/model";
+import { error, log, string } from "@dbx-tools/shared-core";
 import {
   feedback,
   routes,
@@ -63,8 +61,12 @@ import {
   type MastraFeedbackRequest,
   type StatementData,
 } from "@dbx-tools/shared-mastra";
-import { serving as nodeServing } from "@dbx-tools/model";
 import { display, type ServingEndpointSummary } from "@dbx-tools/shared-model";
+import type { Agent } from "@mastra/core/agent";
+import { Mastra } from "@mastra/core/mastra";
+import express from "express";
+import type { Pool } from "pg";
+
 import { buildAgents, FALLBACK_AGENT_ID, type BuiltAgents } from "./agents";
 import { fetchChart } from "./chart";
 import type { MastraPluginConfig } from "./config";
@@ -78,8 +80,6 @@ import { attachRoutePatchMiddleware, isMastraRequestAllowed, MastraServer } from
 import { resolveServingConfig } from "./serving";
 import { fetchStatementData, STATEMENT_ROW_CAP } from "./statement";
 import { threadsRoute } from "./threads";
-import { error, log, string } from "@dbx-tools/shared-core";
-import { plugin } from "@dbx-tools/appkit";
 
 const GENIE_MANIFEST = plugin.data(genie).plugin.manifest;
 const LAKEBASE_MANIFEST = plugin.data(lakebase).plugin.manifest;
@@ -330,7 +330,7 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     // Reads only in-memory build state, so it's synchronous and needs no OBO
     // scoping. Registered before the catch-all, same as `/models`.
     const handleDefaultModel = (req: express.Request, res: express.Response): void => {
-      const requested = string.firstNonEmpty(req.params["agentId"]);
+      const requested = string.firstNonEmpty(req.params.agentId);
       const agentId = requested ?? this.built?.defaultAgentId ?? FALLBACK_AGENT_ID;
       const raw = this.built?.defaultModels[agentId];
       // `"<dynamic>"` (a call-time function) has no fixed id to advertise.
@@ -378,14 +378,14 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     // clean 404; thrown errors bubble through `next(err)`.
     const embedResolvers: Record<string, EmbedResolver> = {
       chart: (req, id, signal) => {
-        const timeoutMs = parseTimeoutMs(req.query["timeoutMs"]);
+        const timeoutMs = parseTimeoutMs(req.query.timeoutMs);
         return fetchChart(id, {
           ...(timeoutMs !== undefined ? { timeoutMs } : {}),
           signal,
         });
       },
       data: (req, id, signal) => {
-        const limit = parseStatementLimit(req.query["limit"]);
+        const limit = parseStatementLimit(req.query.limit);
         return this.userScopedSelf(req).fetchStatement(id, {
           ...(limit !== undefined ? { limit } : {}),
           signal,
@@ -394,8 +394,8 @@ export class MastraPlugin extends Plugin<MastraPluginConfig> {
     };
 
     router.get(`${routes.MASTRA_ROUTES.embed}/:type/:id`, (req, res, next) => {
-      const type = req.params["type"] ?? "";
-      const id = req.params["id"];
+      const type = req.params.type ?? "";
+      const id = req.params.id;
       const resolve = embedResolvers[type];
       if (!resolve) {
         res.status(404).json({ error: `unsupported embed type: ${type}` });

package/src/server.ts

@@ -7,7 +7,9 @@
  * @module
  */
 
+import { randomUUID } from "node:crypto";
 import { getExecutionContext } from "@databricks/appkit";
+import { http, log, object, string, token } from "@dbx-tools/shared-core";
 import { feedback, thread } from "@dbx-tools/shared-mastra";
 import {
   MASTRA_RESOURCE_ID_KEY,
@@ -17,9 +19,6 @@ import {
 import { MastraServer as MastraServerExpress } from "@mastra/express";
 import { trace } from "@opentelemetry/api";
 import type express from "express";
-import { randomUUID } from "node:crypto";
-
-import { resolveFeedbackEnabled } from "./mlflow";
 
 import {
   MASTRA_REQUEST_ID_KEY,
@@ -30,7 +29,8 @@ import {
   type MastraPluginConfig,
   type User,
 } from "./config";
-import { http, log, object, string, token } from "@dbx-tools/shared-core";
+import { resolveFeedbackEnabled } from "./mlflow";
+
 import { extractModelOverride, MASTRA_MODEL_OVERRIDE_KEY, resolveServingConfig } from "./serving";
 /**
  * OpenTelemetry's sentinel for "no valid trace" - 32 zero hex chars.
@@ -154,10 +154,7 @@ export class MastraServer extends MastraServerExpress {
   configureRequestContextScopes(req: express.Request, requestContext: RequestContext) {
     const scopes = [
       ...object
-        .sequence(
-          token.getAccessTokenScopes(req),
-          token.getAccessTokenScopes(req, "authorization"),
-        )
+        .sequence(token.getAccessTokenScopes(req), token.getAccessTokenScopes(req, "authorization"))
         .distinct(),
     ];
     requestContext.set(MASTRA_SCOPES_KEY, scopes);

package/src/serving-sanitize.ts

@@ -1,4 +1,5 @@
 import { string } from "@dbx-tools/shared-core";
+import { type ChatMessage, type ChatRole, openaiChat } from "@dbx-tools/shared-model";
 /**
  * Repairs Mastra / AI SDK message replays sent to Databricks Model
  * Serving before they hit the OpenAI-compatible `/chat/completions`
@@ -6,43 +7,46 @@ import { string } from "@dbx-tools/shared-core";
  */
 
 /**
- * OpenAI-flavoured chat message shape we need to mutate. We do not
- * import the OpenAI / AI SDK types because both packages keep these
- * fields under internal namespaces; the wire payload is the contract
- * here and it's stable enough to inline.
+ * A chat message as it arrives on the serving wire, plus the extended-thinking
+ * fields Databricks-hosted Claude adds. The OpenAI-standard part of the shape
+ * is {@link ChatMessage} from `@dbx-tools/shared-model`; only the provider
+ * extensions this module strips are declared here.
  */
-export interface ServingChatMessage {
-  role: "system" | "user" | "assistant" | "tool" | "reasoning";
-  content?: string | ServingContentPart[];
-  tool_calls?: Array<{ id: string; type: string; function: unknown }>;
-  tool_call_id?: string;
+export interface ServingChatMessage extends ChatMessage {
+  /** Narrowed to the roles this repair pass reasons about, plus Claude's `reasoning` turn. */
+  role: ChatRole | "reasoning";
   reasoning?: unknown;
   reasoning_content?: unknown;
 }
 
-type ServingContentPart = {
-  type?: string;
-  text?: string;
-  [key: string]: unknown;
-};
-
 const REASONING_PART_TYPES = new Set(["reasoning", "thinking", "redacted_thinking"]);
 
 /**
  * Parse, sanitize, and re-serialize a `/serving-endpoints/...` POST
  * body. Returns the original string verbatim when the body is not
- * JSON, has no `messages`, or no rewrite was needed.
+ * JSON or no rewrite was needed.
  */
 export function rewriteServingBody(body: string): string {
-  let parsed: { messages?: unknown };
+  let parsed: Record<string, unknown>;
   try {
-    parsed = JSON.parse(body);
+    parsed = JSON.parse(body) as Record<string, unknown>;
   } catch {
     return body;
   }
-  if (!Array.isArray(parsed.messages)) return body;
-  const messages = parsed.messages as ServingChatMessage[];
-  const changed = stripReasoningFromServingMessages(messages) || repairAssistantPrefill(messages);
+
+  // Runs regardless of `messages`: Databricks refuses to parse a body carrying
+  // an unknown top-level field, so this failure is not specific to a transcript.
+  let changed = openaiChat.stripUnsupportedChatFields(parsed).length > 0;
+
+  if (Array.isArray(parsed.messages)) {
+    const messages = parsed.messages as ServingChatMessage[];
+    // Evaluated eagerly, not short-circuited: one transcript can need both
+    // reasoning blocks stripped AND a trailing assistant prefill folded back.
+    const stripped = stripReasoningFromServingMessages(messages);
+    const repaired = repairAssistantPrefill(messages);
+    changed = changed || stripped || repaired;
+  }
+
   return changed ? JSON.stringify(parsed) : body;
 }
 
@@ -145,15 +149,19 @@ export function repairAssistantPrefill(messages: ServingChatMessage[]): boolean
   return true;
 }
 
+/**
+ * Flatten content to the prose a human would read: text parts only, separated
+ * by blank lines, since the result is spliced into another turn's message body.
+ */
 function textFromServingContent(content: ServingChatMessage["content"]): string {
-  if (typeof content === "string") return content;
-  if (!Array.isArray(content)) return "";
-  return content
-    .filter((part) => part?.type === "text" && typeof part.text === "string")
-    .map((part) => part.text as string)
-    .join("\n\n");
+  return openaiChat.chatContentToText(content, { separator: "\n\n", types: ["text"] });
 }
 
+/**
+ * Whether a turn carries nothing worth replaying. Deliberately stricter than
+ * "has no text": a part of any other type (an image, a cache marker) counts as
+ * content, so a message holding one is kept even though it has no prose.
+ */
 function isEmptyServingContent(content: ServingChatMessage["content"]): boolean {
   if (content === undefined) return true;
   if (typeof content === "string") return content.trim().length === 0;

package/src/serving.ts

@@ -11,11 +11,11 @@
  * @module
  */
 
-import { override } from "@dbx-tools/shared-mastra";
 import { serving as nodeServing } from "@dbx-tools/model";
+import { string } from "@dbx-tools/shared-core";
+import { override } from "@dbx-tools/shared-mastra";
 
 import type { MastraPluginConfig } from "./config";
-import { string } from "@dbx-tools/shared-core";
 
 /**
  * `RequestContext` key under which {@link MastraServer} stores the

package/src/statement.ts

@@ -20,8 +20,8 @@
  */
 
 import { WorkspaceClient } from "@databricks/sdk-experimental";
-import type { GenieDatasetData } from "@dbx-tools/shared-mastra";
 import { databricks } from "@dbx-tools/appkit";
+import type { GenieDatasetData } from "@dbx-tools/shared-mastra";
 
 /**
  * Hard server-side cap on rows returned by the

package/src/summarize.ts

@@ -19,6 +19,7 @@
  * @module
  */
 
+import { string } from "@dbx-tools/shared-core";
 import { model } from "@dbx-tools/shared-model";
 import { Agent } from "@mastra/core/agent";
 import type { MastraModelConfig } from "@mastra/core/llm";
@@ -28,7 +29,6 @@ import { z } from "zod";
 
 import type { MastraPluginConfig } from "./config";
 import { buildModel } from "./model";
-import { string } from "@dbx-tools/shared-core";
 
 /** Fast / small chat tier used for both titling and summaries. */
 const SUMMARY_MODEL_CLASS = model.ModelClass.ChatFast;

package/src/threads.ts

@@ -21,6 +21,7 @@
  * @module
  */
 
+import { log } from "@dbx-tools/shared-core";
 import {
   wire,
   type MastraDeleteThreadResponse,
@@ -35,7 +36,6 @@ import type { ContextWithMastra } from "@mastra/core/server";
 import { registerApiRoute } from "@mastra/core/server";
 
 import { clampPerPage, parseIntParam } from "./pagination";
-import { log } from "@dbx-tools/shared-core";
 
 const logger = log.logger("mastra/threads");
 

package/src/workspaces.ts

@@ -11,6 +11,7 @@
  */
 
 import type { WorkspaceClient } from "@databricks/sdk-experimental";
+import { log, string, token } from "@dbx-tools/shared-core";
 import type { RequestContext } from "@mastra/core/request-context";
 import {
   CompositeFilesystem,
@@ -22,7 +23,6 @@ import {
 
 import { MASTRA_SCOPES_KEY, MASTRA_USER_EMAIL_KEY, MASTRA_USER_KEY, type User } from "./config";
 import { DatabricksWorkspaceFilesystem, emptyFilesystem } from "./filesystems";
-import { log, string, token } from "@dbx-tools/shared-core";
 
 /* ------------------------------ constants ------------------------------ */
 

package/src/writer.ts

@@ -12,8 +12,8 @@
  * @module
  */
 
-import type { MastraWriter } from "@dbx-tools/shared-mastra";
 import { error, log } from "@dbx-tools/shared-core";
+import type { MastraWriter } from "@dbx-tools/shared-mastra";
 
 /**
  * Best-effort `writer.write`. No-op when `writer` is undefined;

package/test/tsconfig.json

@@ -0,0 +1,14 @@
+// ~~ Generated by projen. To modify, edit .projenrc.js and run "pnpm exec projen".
+{
+  "extends": "../tsconfig.json",
+  "compilerOptions": {
+    "noEmit": true,
+    "rootDir": ".."
+  },
+  "include": [
+    "**/*.ts"
+  ],
+  "exclude": [
+    "node_modules"
+  ]
+}