@dbx-tools/appkit-mastra 0.1.112 → 0.3.2

package/src/memory.ts

@@ -0,0 +1,300 @@
+/**
+ * Lakebase-backed Mastra memory wiring.
+ *
+ * Provides a {@link MemoryBuilder} that mints one `Memory` per agent
+ * with two independent knobs:
+ *
+ * - **Storage** (threads / messages via `PostgresStore`): defaults to
+ *   **per-agent** namespacing via {@link agentStorageSchemaName} so
+ *   conversation history stays isolated between agents in the same
+ *   database. `PostgresStore` auto-creates the schema with
+ *   `CREATE SCHEMA IF NOT EXISTS` on init.
+ * - **Memory** (semantic recall via `PgVector`): defaults to a single
+ *   **shared** instance across every agent. Cross-agent recall on one
+ *   index is almost always what users want; opt into per-agent recall
+ *   by passing a {@link MastraMemoryConfigOverride} on the agent.
+ *
+ * Additionally, {@link MemoryBuilder.instanceStorage} returns a
+ * **Mastra-instance-level** `PostgresStore` (schema `mastra_instance`)
+ * used for workflow snapshots - the persistence layer
+ * `agent.resumeStream()` reads from when waking a suspended
+ * `requireApproval` tool call. Per-agent stores are not enough for
+ * this: workflow runs are scoped to the Mastra instance, not an
+ * individual agent's `Memory`.
+ *
+ * Plugin-level `config.storage` / `config.memory` act as the baseline
+ * (auto-defaulted to `true` in `plugin.ts` when the `lakebase` plugin
+ * is registered); per-agent settings cascade on top of that.
+ */
+
+import { getUsernameWithApiLookup } from "@databricks/appkit";
+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");
+
+/**
+ * Build a dedicated **service-principal** Lakebase pool for Mastra
+ * memory from the lakebase plugin's resolved SP pg config.
+ *
+ * The plugin's `exports().pool` is a `RoutingPool` that switches to
+ * the per-user (OBO) pool whenever a query runs inside an `asUser`
+ * scope - exactly the context the mastra plugin establishes around
+ * every chat turn. Memory (threads / messages + semantic recall) must
+ * instead always act as the app service principal: it owns the
+ * auto-created `mastra_*` schemas (a per-user role usually can't
+ * `CREATE SCHEMA`) and is shared across users, so it cannot inherit a
+ * request's OBO identity.
+ *
+ * `pgConfig` must be the plugin's `exports().getPgConfig()` evaluated
+ * **outside** any `asUser` scope (i.e. during setup), so it carries
+ * the SP connection target, OAuth token-refresh `password` callback,
+ * and any `lakebase({ pool })` tuning overrides - all of which this
+ * pool inherits. See the call site in `plugin.ts`.
+ */
+export async function createServicePrincipalPool(pgConfig: PoolConfig): Promise<Pool> {
+  // `getPgConfig()` resolves the SP username synchronously from
+  // `PGUSER` / `DATABRICKS_CLIENT_ID`; fall back to the async API
+  // lookup (e.g. local dev authenticating via PAT) so the pool always
+  // has an identity to connect with.
+  const user = pgConfig.user ?? (await getUsernameWithApiLookup());
+  return new Pool({ ...pgConfig, user });
+}
+
+/** Effective per-knob setting after the plugin/agent cascade. */
+type StorageSetting = MastraAgentDefinition["storage"];
+type MemorySetting = MastraAgentDefinition["memory"];
+
+/**
+ * True when any plugin-level or per-agent setting could need the
+ * Lakebase pool. Used by `plugin.ts` to gate creation of the
+ * service-principal pool and the {@link MemoryBuilder} that consumes
+ * it; when false neither is built.
+ */
+export function needsLakebase(config: MastraPluginConfig): boolean {
+  if (settingNeedsSharedPool(config.storage)) return true;
+  if (settingNeedsSharedPool(config.memory)) return true;
+  const defs = collectAgentDefinitions(config);
+  return defs.some((d) => settingNeedsSharedPool(d.storage) || settingNeedsSharedPool(d.memory));
+}
+
+/**
+ * Construct a per-agent {@link Memory} factory bound to the supplied
+ * service-principal pool (see {@link createServicePrincipalPool}).
+ * Caches the shared `PgVector` singleton (built on first need) so each
+ * agent build is O(1) after the first.
+ */
+export function createMemoryBuilder(
+  config: MastraPluginConfig,
+  servicePrincipalPool: Pool,
+): MemoryBuilder {
+  return new MemoryBuilder(config, servicePrincipalPool);
+}
+
+/**
+ * Builds one `Memory` per agent against a shared service-principal
+ * Lakebase pool. Per-instance state keeps the shared `PgVector` alive
+ * across calls so registering N agents stays cheap.
+ */
+export class MemoryBuilder {
+  private sharedVector: PgVector | undefined;
+
+  constructor(
+    private readonly config: MastraPluginConfig,
+    private readonly servicePrincipalPool: Pool,
+  ) {}
+
+  /**
+   * Build the Mastra-instance-level storage used for workflow
+   * snapshots. Returns `undefined` when plugin-level `storage` is
+   * disabled, in which case `agent.resumeStream()` (and therefore
+   * the `requireApproval` flow) will not be available.
+   *
+   * The store lives in a dedicated `mastra_instance` schema so it
+   * never collides with per-agent {@link agentStorageSchemaName} namespaces.
+   * Workflow snapshots are not per-agent state; they belong to the
+   * `Mastra` instance that owns the workflow execution.
+   */
+  instanceStorage(): PostgresStore | undefined {
+    const setting = this.config.storage;
+    if (!setting) return undefined;
+    if (typeof setting === "object") {
+      return new PostgresStore(
+        withId(setting, "mastra-store__instance") as ConstructorParameters<typeof PostgresStore>[0],
+      );
+    }
+    return new PostgresStore({
+      id: "mastra-store__instance",
+      schemaName: "mastra_instance",
+      pool: this.servicePrincipalPool,
+    });
+  }
+
+  /**
+   * Build a `Memory` for `agentId` after the plugin/agent cascade.
+   * Returns `undefined` when the agent has neither storage nor a
+   * vector store enabled - Mastra accepts a missing `memory` field
+   * and treats the agent as stateless.
+   */
+  forAgent(agentId: string, def: MastraAgentDefinition): Memory | undefined {
+    const storageSetting = def.storage ?? this.config.storage;
+    const memorySetting = def.memory ?? this.config.memory;
+
+    const storage = this.buildStorage(agentId, storageSetting);
+    const vector = this.buildVector(memorySetting);
+    if (!storage && !vector) {
+      logger.debug("agent:stateless", { agentId });
+      return undefined;
+    }
+
+    logger.debug("agent:configured", {
+      agentId,
+      storage: storage !== undefined,
+      vector: vector !== undefined,
+      vectorMode:
+        vector === undefined ? "off" : typeof memorySetting === "object" ? "dedicated" : "shared",
+    });
+
+    return new Memory({
+      ...(storage ? { storage } : {}),
+      ...(vector ? { vector, embedder: fastembed } : {}),
+      options: {
+        lastMessages: 10,
+        ...(vector ? { semanticRecall: { topK: 3, messageRange: 2 } } : {}),
+        // Auto-name each thread from its opening turn so the
+        // conversation list the UI renders shows meaningful titles
+        // instead of raw ids. Titling runs on the small / fast chat
+        // tier (see `summarize.ts`) rather than the agent's primary
+        // model, so naming a thread never spends the heavyweight model.
+        // Only meaningful when storage is on; harmless otherwise.
+        ...(storage
+          ? {
+              generateTitle: {
+                model: summaryModel(this.config),
+                instructions: TITLE_INSTRUCTIONS,
+              },
+            }
+          : {}),
+      },
+    });
+  }
+
+  private buildStorage(agentId: string, setting: StorageSetting): PostgresStore | undefined {
+    if (!setting) return undefined;
+    if (typeof setting === "boolean") {
+      return new PostgresStore({
+        id: `mastra-store__${agentId}`,
+        schemaName: agentStorageSchemaName(agentId),
+        pool: this.servicePrincipalPool,
+      });
+    }
+    // Cast: `withId` guarantees `id` is set, but the distributive
+    // Omit + `id?: string` shape doesn't structurally narrow to the
+    // discriminated union members. Runtime shape is identical.
+    return new PostgresStore(
+      withId(setting, `mastra-store__${agentId}`) as ConstructorParameters<typeof PostgresStore>[0],
+    );
+  }
+
+  /**
+   * Resolve the agent's vector store. Cascade:
+   *
+   * - falsy: no vector.
+   * - `boolean` / `undefined-inheriting-true`: return the shared
+   *   singleton (built lazily on first call). All agents that
+   *   default-enable memory write into and recall from one index.
+   * - object: build a dedicated `PgVector` for this agent.
+   */
+  private buildVector(setting: MemorySetting): PgVector | undefined {
+    if (!setting) return undefined;
+    if (typeof setting === "boolean") return this.getSharedVector();
+    return buildPgVector(setting);
+  }
+
+  private getSharedVector(): PgVector {
+    if (!this.sharedVector) {
+      this.sharedVector = buildSharedPgVector(this.servicePrincipalPool);
+    }
+    return this.sharedVector;
+  }
+}
+
+/**
+ * Build the shared `PgVector` that backs the default
+ * `def.memory === true` case across every agent.
+ *
+ * `PgVector`'s constructor accepts only connection-style configs
+ * (`HostConfig` / `ConnectionStringConfig` / `ClientConfig`); there is
+ * no `{ pool }` shorthand the way `PostgresStore` has one. Worse, the
+ * constructor synchronously kicks off a `cacheWarmupPromise` IIFE that
+ * calls `this.pool.connect()` before returning, so we can't cleanly
+ * hand it an inert config and patch the pool afterwards.
+ *
+ * The trick: pass illegal-but-validation-passing placeholders so the
+ * warmup's `net.connect()` rejects synchronously with `RangeError`
+ * (Node validates `0 <= port < 65536`). The IIFE's `catch {}` swallows
+ * it, no DNS lookup or TCP attempt happens, and we then swap
+ * `pgVector.pool` to the lakebase pool. Every subsequent `PgVector`
+ * method reads `this.pool` at call time, so all real I/O goes through
+ * the lakebase pool from then on. The placeholder pool is `.end()`'d
+ * so its socket book-keeping is released.
+ */
+function buildSharedPgVector(pool: Pool): PgVector {
+  const vector = new PgVector({
+    id: `pg${randomUUID()}`,
+    // Keep the recall index out of `public`: on a Lakebase database the app
+    // service principal has no CREATE on `public` (PG15+ locks it down), so a
+    // default-schema PgVector fails on CREATE INDEX with "permission denied for
+    // schema public". PgVector runs CREATE SCHEMA IF NOT EXISTS for a named
+    // schema, so the SP creates and owns `mastra_instance` (the same
+    // instance-level schema the workflow-snapshot store uses) and the index
+    // lands there. Table names don't collide with the storage tables.
+    schemaName: "mastra_instance",
+    host: "-1",
+    port: -1,
+    database: "_",
+    user: "_",
+    password: "_",
+  });
+  const placeholder = vector.pool;
+  vector.pool = pool;
+  void placeholder.end().catch(() => undefined);
+  return vector;
+}
+
+/** Per-agent dedicated `PgVector` (rare; opt-in via object override). */
+function buildPgVector(setting: MastraMemoryConfigOverride): PgVector {
+  return new PgVector(
+    withId(setting, `pg-vector__${randomUUID()}`) as ConstructorParameters<typeof PgVector>[0],
+  );
+}
+
+/** True when this setting requires the shared Lakebase pool. */
+function settingNeedsSharedPool(setting: StorageSetting | MemorySetting | undefined): boolean {
+  return setting === true;
+}
+
+/** Walk the three shapes of `config.agents` into a flat list. */
+function collectAgentDefinitions(config: MastraPluginConfig): MastraAgentDefinition[] {
+  const agents = config.agents;
+  if (!agents) return [];
+  if (Array.isArray(agents)) return agents;
+  if (typeof (agents as MastraAgentDefinition).instructions === "string") {
+    return [agents as MastraAgentDefinition];
+  }
+  return Object.values(agents as Record<string, MastraAgentDefinition>);
+}
+
+/** Fill in a default `id` when the caller didn't supply one. */
+function withId<T extends { id?: string }>(value: T, fallback: string): T {
+  return value.id ? value : { ...value, id: fallback };
+}

package/src/mlflow.ts

@@ -0,0 +1,149 @@
+/**
+ * MLflow user-feedback logging: detect whether MLflow tracing is wired
+ * for this deployment, and log a thumbs / comment as a trace
+ * *assessment* via the Databricks MLflow REST API.
+ *
+ * Feedback attaches to a trace, and the plugin's spans reach MLflow
+ * through the same OTel pipeline as every other AppKit span (see
+ * `observability.ts`). MLflow derives its trace id from the OpenTelemetry
+ * trace id (`tr-<hex(otelTraceId)>`), so the server stamps the active
+ * trace id on each turn's response and the client sends it back here.
+ *
+ * There is no MLflow JS SDK, so this posts to the assessments REST
+ * endpoint directly using the OBO-scoped workspace client (the feedback
+ * is thus attributed to the signed-in user). Trace export is
+ * asynchronous, so the just-finished trace may not exist in MLflow yet
+ * when the user reacts; the log call retries briefly on "not found"
+ * before giving up softly.
+ */
+
+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");
+
+/** Workspace client carried on an AppKit execution context. */
+type WorkspaceClient = appkit.WorkspaceClientLike;
+
+/** Assessments REST path for a trace. `3.0` is the current MLflow API version. */
+const assessmentsPath = (traceId: string): string =>
+  `/api/3.0/mlflow/traces/${encodeURIComponent(traceId)}/assessments`;
+
+/** Number of times to retry a "trace not found" response before giving up. */
+const NOT_FOUND_RETRIES = 3;
+/** Base backoff between "trace not found" retries, in ms (grows linearly). */
+const NOT_FOUND_BACKOFF_MS = 1200;
+
+/**
+ * Whether MLflow feedback logging is available for this deployment.
+ *
+ * Enabled when an OTLP exporter endpoint is configured (traces are
+ * actually shipped somewhere) AND an MLflow experiment is named - the
+ * two signals that the OTLP backend is MLflow and traces will
+ * materialize there. Both are standard env vars, so no plugin config is
+ * required; a deployment opts in simply by wiring MLflow tracing.
+ */
+export function mlflowEnabled(): boolean {
+  const hasExporter = Boolean(
+    process.env.OTEL_EXPORTER_OTLP_ENDPOINT?.trim() ||
+    process.env.OTEL_EXPORTER_OTLP_TRACES_ENDPOINT?.trim(),
+  );
+  const hasExperiment = Boolean(
+    process.env.MLFLOW_EXPERIMENT_ID?.trim() || process.env.MLFLOW_EXPERIMENT_NAME?.trim(),
+  );
+  return hasExporter && hasExperiment;
+}
+
+/**
+ * Resolve whether user feedback is enabled from the plugin's optional
+ * `config.feedback` override: the explicit boolean wins, otherwise fall
+ * back to auto-detecting MLflow tracing ({@link mlflowEnabled}). Shared
+ * by the plugin's client-config gate and the server's trace-id header so
+ * the two never disagree.
+ */
+export function resolveFeedbackEnabled(explicit: boolean | undefined): boolean {
+  return explicit ?? mlflowEnabled();
+}
+
+/** Parameters for {@link logFeedback}. */
+export interface LogFeedbackParams {
+  /** MLflow trace id the assessment attaches to (`tr-<hex>`). */
+  traceId: string;
+  /** Assessment name; defaults per whether a value or a comment is sent. */
+  name?: string;
+  /** Thumbs / rating / label value. Omit for a comment-only submission. */
+  value?: boolean | number | string;
+  /** Freeform comment: the rationale alongside a value, or the value itself when none. */
+  comment?: string;
+  /** Identity the feedback is attributed to (user email / id). */
+  sourceId?: string;
+}
+
+/**
+ * Log a HUMAN feedback assessment to a trace. Returns the created
+ * assessment id on success, or `undefined` when the trace can't be
+ * found (even after retrying for export lag) or the request otherwise
+ * fails - callers surface that as a soft "not recorded" rather than an
+ * error, keeping the chat usable.
+ */
+export async function logFeedback(
+  client: WorkspaceClient,
+  params: LogFeedbackParams,
+): Promise<string | undefined> {
+  // A comment with no thumbs value is logged as text feedback; a value
+  // (with an optional comment as the rationale) is the thumbs path.
+  const hasValue = params.value !== undefined;
+  const name =
+    params.name?.trim() ||
+    (hasValue ? feedback.DEFAULT_FEEDBACK_NAME : feedback.DEFAULT_COMMENT_NAME);
+  const value = hasValue ? params.value : params.comment;
+  const assessment: Record<string, unknown> = {
+    trace_id: params.traceId,
+    assessment_name: name,
+    source: {
+      source_type: "HUMAN",
+      source_id: params.sourceId?.trim() || "user",
+    },
+    feedback: { value },
+    ...(hasValue && params.comment?.trim() ? { rationale: params.comment } : {}),
+  };
+  const body = { assessment };
+
+  for (let attempt = 0; attempt <= NOT_FOUND_RETRIES; attempt++) {
+    let res: Response;
+    try {
+      res = await databricksFetch(client, assessmentsPath(params.traceId), {
+        method: "POST",
+        body,
+      });
+    } catch (err) {
+      logger.warn("feedback request failed", {
+        traceId: params.traceId,
+        error: error.errorMessage(err),
+      });
+      return undefined;
+    }
+    if (res.ok) {
+      const parsed = await readResponseJson(res);
+      const assessmentId =
+        (parsed as { assessment?: { assessment_id?: unknown } })?.assessment?.assessment_id ??
+        (parsed as { assessment_id?: unknown })?.assessment_id;
+      return typeof assessmentId === "string" ? assessmentId : "";
+    }
+    // Trace export is async; a fresh trace may not exist yet. Retry a
+    // few times with a short backoff before giving up softly.
+    if (res.status === 404 && attempt < NOT_FOUND_RETRIES) {
+      await async.sleep(NOT_FOUND_BACKOFF_MS * (attempt + 1));
+      continue;
+    }
+    logger.warn("feedback not recorded", {
+      traceId: params.traceId,
+      status: res.status,
+      body: await readResponseText(res),
+    });
+    return undefined;
+  }
+  return undefined;
+}

package/src/model.ts

@@ -0,0 +1,163 @@
+/**
+ * Databricks Model Serving resolver for Mastra agents.
+ *
+ * Each agent step calls {@link buildModel} with the active
+ * `RequestContext`. The user stamped by `MastraServer` carries an
+ * AppKit `WorkspaceClient`; we ask it for the workspace host and a
+ * fresh bearer header, then point Mastra's OpenAI-compatible provider
+ * at `/serving-endpoints` on that host.
+ *
+ * This module only adds the Mastra-specific glue. The actual model
+ * selection - listing the workspace catalogue and resolving an
+ * explicit name / class / fallback chain to a real endpoint id - lives
+ * in `@dbx-tools/model` ({@link selectModel}) so non-Mastra consumers
+ * (e.g. a job that just needs a model name) can reuse it. Here we
+ * assemble the explicit ask from Mastra's request context (the
+ * per-request override under {@link MASTRA_MODEL_OVERRIDE_KEY}, the
+ * agent / plugin `modelId`, or `DATABRICKS_SERVING_ENDPOINT_NAME`),
+ * pass the plugin's fuzzy / class / fallback knobs through, and wrap
+ * the resolved id in the OpenAI-compatible provider config Mastra
+ * expects. Catalogue fetches fail loud: network / auth errors
+ * propagate so callers see the real SDK message.
+ */
+
+import { getExecutionContext } from "@databricks/appkit";
+import { classes, resolve } from "@dbx-tools/model";
+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";
+
+type ModelClass = model.ModelClass;
+const { parseModelClass } = classes;
+const { selectModel } = resolve;
+
+/** Optional overrides accepted by {@link buildModel}. */
+export interface BuildModelOverrides {
+  /**
+   * Static model id from the agent / plugin config (string sugar on
+   * `def.model` or `config.defaultModel`). Loses to the per-request
+   * override but wins over env / class / fallback.
+   */
+  modelId?: string;
+  /**
+   * Chat capability class to resolve when no explicit model id is
+   * supplied. Used by internal agents (e.g. the chart planner asks for
+   * {@link model.ModelClass.ChatFast}) to express intent without pinning an
+   * endpoint name; the live catalogue is classified and the top
+   * available model in the class is chosen, falling back to the
+   * class's static list when the workspace has none.
+   */
+  modelClass?: ModelClass;
+}
+
+/**
+ * Resolve a `MastraModelConfig` for the current agent step. Runs
+ * while `agent.stream` is inside the `asUser(req)` scope so tokens
+ * are user-scoped; outside an active user context the workspace
+ * client falls back to the service principal.
+ */
+export async function buildModel(
+  config: MastraPluginConfig,
+  requestContext: RequestContext,
+  overrides: BuildModelOverrides = {},
+): Promise<MastraModelConfig> {
+  void setupFetchInterceptor();
+  // The chat path stamps the AppKit user on the request context via
+  // `MastraServer`. The MCP transport routes don't thread that context
+  // into tool execution, so fall back to the ambient execution context
+  // (the active OBO scope, or the service principal) when it's absent.
+  const user = requestContext.get(MASTRA_USER_KEY) as User | undefined;
+  const executionContext = user?.executionContext ?? getExecutionContext();
+  const clientConfig = executionContext.client.config;
+  const host = (await clientConfig.getHost()).toString();
+  const headers = new Headers();
+  await clientConfig.authenticate(headers);
+  // The OpenAI Node SDK appends paths like `/chat/completions` to whatever
+  // URL we hand it. Drop the trailing slash so the resulting URL stays
+  // well-formed (`/serving-endpoints/chat/completions`).
+  const url = new URL("/serving-endpoints", host).toString().replace(/\/$/, "");
+
+  const logger = log.logger(config);
+  const serving = resolveServingConfig(config);
+  const override = serving.allowOverride
+    ? (requestContext.get(MASTRA_MODEL_OVERRIDE_KEY) as string | undefined)
+    : undefined;
+
+  // The override / agent default / env value can be either a concrete
+  // endpoint name or a model class slug ("chat-thinking" /
+  // "chat-balanced" / "chat-fast"). A class slug becomes a class intent
+  // (let the live catalogue pick the best model in that band); anything
+  // else is an explicit name fuzzy-matched against the catalogue. An
+  // internal `overrides.modelClass` (e.g. the chart planner) is the
+  // floor when nothing was requested.
+  const requested = override ?? overrides.modelId ?? process.env.DATABRICKS_SERVING_ENDPOINT_NAME;
+  const requestedClass = requested !== undefined ? parseModelClass(requested) : null;
+  const explicit = requestedClass === null ? requested : undefined;
+  const modelClass = requestedClass ?? overrides.modelClass;
+
+  const { modelId, source } = await selectModel(executionContext.client, host, {
+    ...(explicit !== undefined ? { explicit } : {}),
+    fuzzy: serving.fuzzy,
+    threshold: serving.threshold,
+    ...(modelClass !== undefined ? { modelClass } : {}),
+    fallbacks: serving.fallbacks,
+    ttlMs: serving.ttlMs,
+  });
+  logger.debug("model selected", { modelId, source, requested });
+
+  return {
+    providerId: config.providerId ?? "databricks",
+    modelId,
+    url,
+    headers: Object.fromEntries(headers.entries()),
+  };
+}
+
+/** Path prefix that identifies a Databricks Model Serving REST call. */
+const SERVING_ENDPOINTS_PATH_PREFIX = "/serving-endpoints/";
+
+/**
+ * Install a single shared `globalThis.fetch` wrapper for every POST to
+ * `/serving-endpoints/...`. The wrapper does two things:
+ *
+ *   1. Rewrites the outgoing `messages` array to repair Mastra/AI SDK
+ *      stream-replay quirks that Databricks-hosted Claude rejects (see
+ *      {@link rewriteServingBody} in `./serving-sanitize.js`).
+ *   2. At `LOG_LEVEL=debug`, dumps the (post-sanitize) JSON body so
+ *      4xx debugging doesn't have to fight AI SDK's `[Array]`
+ *      formatter.
+ *
+ * Safe to call from any hot path: {@link functionModule.memoize} ensures
+ * the wrapper is installed at most once per process, so subsequent
+ * calls are a no-op even when {@link buildModel} fires on every agent
+ * step.
+ */
+const setupFetchInterceptor = functionModule.memoize((): void => {
+  const logger = log.logger("mastra/llm");
+  const original = globalThis.fetch.bind(globalThis);
+  globalThis.fetch = (async (input, init) => {
+    const url = net.urlBuilder(input);
+    if (
+      !url ||
+      !url.pathname.startsWith(SERVING_ENDPOINTS_PATH_PREFIX) ||
+      typeof init?.body !== "string"
+    ) {
+      return original(input, init);
+    }
+    const rewritten = rewriteServingBody(init.body);
+    if (rewritten !== init.body) {
+      init = { ...init, body: rewritten };
+    }
+    try {
+      logger.debug("POST", { url: url.toString(), body: JSON.parse(rewritten) });
+    } catch {
+      logger.debug("POST", { url: url.toString(), bodyType: "non-JSON" });
+    }
+    return original(input, init);
+  }) as typeof globalThis.fetch;
+});