@dbx-tools/appkit-mastra 0.3.21 → 0.3.22

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.22",
+    "@dbx-tools/databricks": "0.3.22",
+    "@dbx-tools/genie": "0.3.22",
+    "@dbx-tools/model": "0.3.22",
+    "@dbx-tools/shared-core": "0.3.22",
+    "@dbx-tools/shared-mastra": "0.3.22",
+    "@dbx-tools/shared-genie": "0.3.22",
+    "@dbx-tools/core": "0.3.22",
+    "@dbx-tools/shared-model": "0.3.22"
   },
   "main": "index.ts",
   "license": "UNLICENSED",
   "publishConfig": {
     "access": "public"
   },
-  "version": "0.3.21",
+  "version": "0.3.22",
   "types": "index.ts",
   "type": "module",
   "exports": {

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;