@dbx-tools/appkit-mastra 0.3.36 → 0.3.37

package/README.md

@@ -468,7 +468,9 @@ client that talks to these routes.
   defaults, and approval-gated tool inspection.
 - `config` - plugin config types and RequestContext key constants.
 - `model` / `serving` / `servingSanitize` - Mastra model config, request
-  overrides, serving-endpoint config, and request-body cleanup.
+  overrides, serving-endpoint config, and the on-the-wire request/response
+  cleanup that keeps provider-specific payload quirks (Claude's replayed
+  thinking blocks, Gemini's content-parts responses) from failing a turn.
 - `genie` - Genie prompt, space normalization, Genie toolkits, and suggestions.
 - `chart` / `statement` / `writer` - chart cache, statement row fetches, and
   safe writer events.

package/package.json

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

package/src/model.ts

@@ -32,7 +32,7 @@ import type { RequestContext } from "@mastra/core/request-context";
 
 import { MASTRA_USER_KEY, type MastraPluginConfig, type User } from "./config";
 import { MASTRA_MODEL_OVERRIDE_KEY, resolveServingConfig } from "./serving";
-import { rewriteServingBody } from "./serving-sanitize";
+import { rewriteServingBody, rewriteServingResponseBody } from "./serving-sanitize";
 
 type ModelClass = model.ModelClass;
 const { parseModelClass } = classes;
@@ -140,6 +140,10 @@ const SERVING_ENDPOINTS_PATH_PREFIX = "/serving-endpoints/";
  *   2. At `LOG_LEVEL=debug`, dumps the (post-sanitize) JSON body so
  *      4xx debugging doesn't have to fight AI SDK's `[Array]`
  *      formatter.
+ *   3. Repairs the non-streaming JSON response, where Databricks-hosted
+ *      Gemini returns `choices[].message.content` as a parts array that
+ *      the AI SDK's OpenAI schema rejects (see
+ *      {@link rewriteServingResponseBody}).
  *
  * Safe to call from any hot path: {@link functionModule.memoize} ensures
  * the wrapper is installed at most once per process, so subsequent
@@ -169,6 +173,38 @@ const setupFetchInterceptor = functionModule.memoize((): void => {
         ? { url: url.toString(), bodyType: "non-JSON" }
         : { url: url.toString(), body: parsed },
     );
-    return original(input, init);
+    const response = await original(input, init);
+    return repairServingResponse(response);
   }) as typeof globalThis.fetch;
 });
+
+/**
+ * Rewrite a non-streaming serving response whose body needs repair, leaving
+ * everything else byte-identical.
+ *
+ * Streaming turns are passed straight through: an SSE body must stay a live
+ * stream (buffering it to a string would defeat streaming and break
+ * `text/event-stream` parsing), and the delta frames already carry string
+ * content, so they never hit the array-shaped `content` bug. Likewise a
+ * non-JSON or error body is returned untouched, so the caller still sees the
+ * original status and headers.
+ */
+async function repairServingResponse(response: Response): Promise<Response> {
+  const contentType = response.headers.get("content-type") ?? "";
+  if (!contentType.includes("application/json")) return response;
+
+  const body = await response.clone().text();
+  const rewritten = rewriteServingResponseBody(body);
+  if (rewritten === body) return response;
+
+  // `content-length` / `content-encoding` describe the ORIGINAL bytes, so they
+  // are dropped: the rewritten body is a different length and already decoded.
+  const headers = new Headers(response.headers);
+  headers.delete("content-length");
+  headers.delete("content-encoding");
+  return new Response(rewritten, {
+    status: response.status,
+    statusText: response.statusText,
+    headers,
+  });
+}

package/src/serving-sanitize.ts

@@ -1,14 +1,20 @@
 /**
- * Repairs Mastra / AI SDK message replays sent to Databricks Model
- * Serving before they hit the OpenAI-compatible `/chat/completions`
- * route.
+ * Repairs traffic between Mastra / the AI SDK and Databricks Model Serving's
+ * OpenAI-compatible `/chat/completions` route, in both directions.
  *
- * Exists because the transcript Mastra persists is not always a transcript the
- * provider will accept back: Databricks-hosted Claude rejects replayed
- * extended-thinking blocks and reads a trailing assistant message as a prefill
- * request. Both repairs are provider quirks, not schema violations, so they are
- * applied on the wire (see the `globalThis.fetch` wrapper in `model.ts`) rather
- * than by changing what the agent stores or what the UI shows.
+ * Outbound ({@link rewriteServingBody}), because the transcript Mastra
+ * persists is not always a transcript the provider will accept back:
+ * Databricks-hosted Claude rejects replayed extended-thinking blocks and reads
+ * a trailing assistant message as a prefill request.
+ *
+ * Inbound ({@link rewriteServingResponseBody}), because Databricks-hosted
+ * Gemini answers with its native content-parts array where the OpenAI contract
+ * (and therefore the AI SDK's response schema) requires a plain string.
+ *
+ * Every repair here is a provider quirk rather than a schema violation, so all
+ * of them are applied on the wire (see the `globalThis.fetch` wrapper in
+ * `model.ts`) rather than by changing what the agent stores or what the UI
+ * shows.
  *
  * @module
  */
@@ -186,3 +192,56 @@ function isEmptyServingContent(content: ServingChatMessage["content"]): boolean
     return false;
   });
 }
+
+/**
+ * Parse, sanitize, and re-serialize a `/serving-endpoints/...` non-streaming
+ * JSON RESPONSE body. Returns the original string verbatim when the body is
+ * not JSON or no rewrite was needed, mirroring
+ * {@link rewriteServingBody} on the request side.
+ */
+export function rewriteServingResponseBody(body: string): string {
+  const parsed = json.parseRecord(body);
+  if (!parsed) return body;
+  return flattenChoiceMessageContent(parsed) ? JSON.stringify(parsed) : body;
+}
+
+/**
+ * Collapse a structured `choices[].message.content` array to the plain string
+ * the OpenAI Chat Completions contract specifies.
+ *
+ * Databricks-hosted Gemini answers a non-streaming `/chat/completions` call
+ * with the Gemini-native parts shape:
+ *
+ * ```json
+ * "content": [{ "type": "text", "text": "...", "thoughtSignature": "..." }]
+ * ```
+ *
+ * `@ai-sdk/openai-compatible` validates the response against the OpenAI
+ * schema, where `content` is a nullable string, so it throws
+ * `Type validation failed: ... expected string, received array` and the whole
+ * call fails (`AI_APICallError: Invalid JSON response` on an HTTP 200). Mastra
+ * uses `doGenerate` for its side calls, so the visible symptom is
+ * `Error generating title` - every thread keeps its placeholder name.
+ *
+ * Flattening on the wire keeps the repair in one place: the streaming path is
+ * unaffected (deltas already carry string content), and neither the agent's
+ * stored transcript nor the UI has to know the provider emitted parts. Any
+ * non-text part (a `thoughtSignature`-only entry, an inline image) contributes
+ * nothing, matching {@link openaiChat.chatContentToText}, and an all-parts-empty
+ * message flattens to `""` rather than being dropped, so `finish_reason` and
+ * `usage` still round-trip.
+ */
+export function flattenChoiceMessageContent(payload: Record<string, unknown>): boolean {
+  if (!Array.isArray(payload.choices)) return false;
+  let changed = false;
+  for (const choice of payload.choices) {
+    if (!choice || typeof choice !== "object") continue;
+    const message = (choice as { message?: unknown }).message;
+    if (!message || typeof message !== "object") continue;
+    const target = message as { content?: unknown };
+    if (!Array.isArray(target.content)) continue;
+    target.content = openaiChat.chatContentToText(target.content, { types: ["text"] });
+    changed = true;
+  }
+  return changed;
+}

package/test/serving-sanitize.test.ts

@@ -0,0 +1,98 @@
+import assert from "node:assert/strict";
+import { describe, it } from "node:test";
+
+import { rewriteServingResponseBody } from "../src/serving-sanitize";
+
+/**
+ * Trimmed copy of a real Databricks-hosted Gemini reply: `content` is the
+ * Gemini-native parts array, and the text part also carries the signed
+ * `thoughtSignature` the model returns alongside it.
+ */
+const geminiPartsResponse = {
+  model: "gemini-3.1-flash-lite",
+  choices: [
+    {
+      message: {
+        role: "assistant",
+        content: [
+          { type: "text", text: "Fuel Mart Weekly Sales Summary", thoughtSignature: "AY89a18" },
+        ],
+      },
+      index: 0,
+      finish_reason: "stop",
+    },
+  ],
+  usage: { prompt_tokens: 2098, completion_tokens: 5, total_tokens: 2201 },
+  object: "chat.completion",
+};
+
+describe("serving response sanitize", () => {
+  it("flattens Gemini's content parts to the string the AI SDK expects", () => {
+    const result = JSON.parse(rewriteServingResponseBody(JSON.stringify(geminiPartsResponse)));
+    assert.equal(result.choices[0].message.content, "Fuel Mart Weekly Sales Summary");
+  });
+
+  it("preserves every sibling field while rewriting content", () => {
+    const result = JSON.parse(rewriteServingResponseBody(JSON.stringify(geminiPartsResponse)));
+    assert.equal(result.choices[0].finish_reason, "stop");
+    assert.equal(result.choices[0].message.role, "assistant");
+    assert.deepEqual(result.usage, geminiPartsResponse.usage);
+    assert.equal(result.model, "gemini-3.1-flash-lite");
+  });
+
+  it("joins multiple text parts and ignores non-text ones", () => {
+    const body = JSON.stringify({
+      choices: [
+        {
+          message: {
+            content: [
+              { type: "text", text: "Fuel Mart " },
+              { type: "image", image_url: "https://example.invalid/x.png" },
+              { type: "text", text: "Sales" },
+            ],
+          },
+        },
+      ],
+    });
+    const result = JSON.parse(rewriteServingResponseBody(body));
+    assert.equal(result.choices[0].message.content, "Fuel Mart Sales");
+  });
+
+  it("flattens a parts array holding no text to an empty string", () => {
+    const body = JSON.stringify({
+      choices: [{ message: { content: [{ thoughtSignature: "AY89" }] }, finish_reason: "stop" }],
+    });
+    const result = JSON.parse(rewriteServingResponseBody(body));
+    assert.equal(result.choices[0].message.content, "");
+    assert.equal(result.choices[0].finish_reason, "stop");
+  });
+
+  it("returns a compliant OpenAI response byte-identical", () => {
+    const body = JSON.stringify({
+      choices: [{ message: { role: "assistant", content: "already a string" } }],
+    });
+    assert.equal(rewriteServingResponseBody(body), body);
+  });
+
+  it("leaves a non-JSON or choice-less body untouched", () => {
+    assert.equal(rewriteServingResponseBody("not json"), "not json");
+    assert.equal(
+      rewriteServingResponseBody('{"error_code":"BAD_REQUEST"}'),
+      '{"error_code":"BAD_REQUEST"}',
+    );
+  });
+
+  it("repairs every choice when the provider returns more than one", () => {
+    const body = JSON.stringify({
+      choices: [
+        { message: { content: [{ type: "text", text: "first" }] } },
+        { message: { content: [{ type: "text", text: "second" }] } },
+      ],
+    });
+    const result = JSON.parse(rewriteServingResponseBody(body));
+    assert.deepEqual(
+      result.choices.map((c: { message: { content: string } }) => c.message.content),
+      ["first", "second"],
+    );
+  });
+});