@dbx-tools/appkit-mastra 0.1.13 → 0.1.14

package/src/config.ts

@@ -5,11 +5,14 @@
  * `memory.ts` can import them without creating a cycle.
  */
 
-import type { BasePluginConfig, getExecutionContext } from "@databricks/appkit";
+import type { BasePluginConfig } from "@databricks/appkit";
 import type { AgentConfig } from "@mastra/core/agent";
+import { MASTRA_RESOURCE_ID_KEY, MASTRA_THREAD_ID_KEY } from "@mastra/core/request-context";
 import type { PgVectorConfig, PostgresStoreConfig } from "@mastra/pg";
 
-import type { MastraAgentDefinition, MastraTools } from "./agents.js";
+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
@@ -18,10 +21,63 @@ import type { MastraAgentDefinition, MastraTools } from "./agents.js";
  */
 export const MASTRA_USER_KEY = "mastra__user";
 
+/**
+ * `RequestContext` keys for AppKit user metadata stamped by
+ * {@link MastraServer}. Surfaced as trace metadata via
+ * {@link TRACE_REQUEST_CONTEXT_KEYS} so traces are filterable by who
+ * issued the request without leaking the full user object.
+ */
+export const MASTRA_USER_NAME_KEY = "mastra__userName";
+export const MASTRA_USER_EMAIL_KEY = "mastra__userEmail";
+
+/**
+ * `RequestContext` key for the per-HTTP-request id stamped by
+ * {@link MastraServer}. Reads `X-Request-Id` from the incoming
+ * headers when present (so an upstream load balancer / API gateway
+ * can keep its trace correlation), falls back to a freshly minted
+ * UUID. Echoed back on the response and surfaced on every span via
+ * {@link TRACE_REQUEST_CONTEXT_KEYS} so logs and traces share a
+ * join key.
+ */
+export const MASTRA_REQUEST_ID_KEY = "mastra__requestId";
+
+/**
+ * `RequestContext` key for OAuth scopes parsed from the forwarded
+ * access token by {@link MastraServer.configureRequestContextScopes}.
+ * Workspace mounts that touch Databricks workspace files require
+ * `workspace` or `all-apis` in this list.
+ */
+export const MASTRA_SCOPES_KEY = "mastra__scopes";
+
+/**
+ * Canonical list of `RequestContext` keys we want Mastra to extract
+ * as metadata on every observability span (agent runs, model calls,
+ * tool invocations, workflow steps).
+ *
+ * Mirrors {@link https://mastra.ai/docs/observability/tracing/overview#automatic-metadata-from-requestcontext}:
+ * passed verbatim into `Observability.configs[*].requestContextKeys`,
+ * so any key listed here is read from `RequestContext` at trace
+ * start and attached as scalar span metadata. Keep the set to plain
+ * scalars - never include {@link MASTRA_USER_KEY} (it carries the
+ * full AppKit execution context with a `WorkspaceClient` reference).
+ *
+ * Order is purely cosmetic; Mastra de-dupes internally.
+ */
+export const TRACE_REQUEST_CONTEXT_KEYS: readonly string[] = [
+  MASTRA_RESOURCE_ID_KEY,
+  MASTRA_THREAD_ID_KEY,
+  MASTRA_REQUEST_ID_KEY,
+  MASTRA_USER_NAME_KEY,
+  MASTRA_USER_EMAIL_KEY,
+  // Model override key is owned by `serving.ts`; spelled inline here
+  // so this module stays leaf-level (no cycles with `serving.ts`).
+  "mastra__model_override",
+];
+
 /** AppKit execution context plus the canonical user id. */
 export interface User {
   id: string;
-  executionContext: ReturnType<typeof getExecutionContext>;
+  executionContext: appkit.ExecutionContextLike;
 }
 
 /** PgVector config with an optional Mastra store id. */
@@ -29,6 +85,45 @@ export type MastraMemoryConfig = PgVectorConfig & {
   id?: string;
 };
 
+/**
+ * Fine-grained control for the optional MCP server exposure
+ * ({@link MastraPluginConfig.mcp}). Every field is optional; the object
+ * form only needs to set what differs from the defaults.
+ */
+export interface MastraMcpConfig {
+  /**
+   * Server id used in the route path (`/mcp/<serverId>/...`) and as the
+   * MCP registry id. Defaults to the plugin's registered name.
+   */
+  serverId?: string;
+  /** Display name advertised over MCP. Defaults to `"<displayName> MCP"`. */
+  name?: string;
+  /** Semantic version advertised over MCP. Defaults to `"1.0.0"`. */
+  version?: string;
+  /** Optional human-readable description advertised over MCP. */
+  description?: string;
+  /**
+   * Expose every registered agent as an `ask_<agentId>` MCP tool.
+   * Defaults to `true` - this is the "leverage the Mastra agents over
+   * MCP" behavior most callers want.
+   */
+  agents?: boolean;
+  /**
+   * Also expose the plugin's ambient tools (the built-in `render_data`
+   * plus anything in `config.tools`) as MCP tools. Defaults to `false`:
+   * the ambient tools assume an in-process chat turn (they publish
+   * writer events the chat UI consumes), so they aren't useful to a
+   * standalone MCP client. Turn this on only when those tools are safe
+   * to call out-of-band.
+   */
+  tools?: boolean;
+  /**
+   * Extra tools to expose over MCP beyond the agent / ambient sets.
+   * Use this for tools written specifically for MCP consumers.
+   */
+  extraTools?: MastraTools;
+}
+
 /** Configuration accepted by the Mastra AppKit plugin. */
 export interface MastraPluginConfig extends BasePluginConfig {
   /** Mastra OpenAI-compatible provider id. Defaults to `"databricks"`. */
@@ -89,10 +184,7 @@ export interface MastraPluginConfig extends BasePluginConfig {
    * });
    * ```
    */
-  agents?:
-    | Record<string, MastraAgentDefinition>
-    | MastraAgentDefinition
-    | MastraAgentDefinition[];
+  agents?: Record<string, MastraAgentDefinition> | MastraAgentDefinition | MastraAgentDefinition[];
   /**
    * Ambient tools spread into every registered agent's tools record;
    * per-agent tools win on key collision. Use for a small shared
@@ -100,7 +192,8 @@ export interface MastraPluginConfig extends BasePluginConfig {
    */
   tools?: MastraTools;
   /**
-   * Agent id used by `chatRoute` when the client doesn't specify one.
+   * Agent id used when the client doesn't specify one (the bare,
+   * un-suffixed history / suggestions routes resolve to it).
    * Defaults to the first key in `agents` (or `"default"` when
    * `agents` is omitted). Must match an id in `agents` when both are
    * set; a mismatch throws at setup with the available candidates.
@@ -148,17 +241,20 @@ export interface MastraPluginConfig extends BasePluginConfig {
    */
   modelOverride?: boolean;
   /**
-   * Priority-ordered list of endpoint names tried when no agent /
-   * plugin / env / request-override model id is set. The resolver
-   * picks the first id that is actually present in the workspace's
-   * `/serving-endpoints` listing - this is what lets a workspace
-   * without Claude Opus still get a sensible default automatically.
+   * Priority-ordered list of endpoint names tried *first* when no
+   * agent / plugin / env / request-override model id is set, ahead of
+   * the dynamic score-classified catalogue. The resolver picks the
+   * first id that is actually present in the workspace's
+   * `/serving-endpoints` listing.
    *
-   * Defaults to the built-in list in `model.ts` (`FALLBACK_MODEL_IDS`):
-   * Claude (Opus -> Sonnet -> Haiku), then OpenAI GPT-5 family, then
-   * open weights (Llama 4, Llama 3.3, GPT-OSS, Qwen, Llama 3.1).
-   * Override here to pin a regulated workspace to an approved subset
-   * or to add custom endpoints in front of the public catalogue.
+   * When unset, resolution is driven by the live Foundation Model API
+   * `quality` / `speed` / `cost` scores: endpoints are classified into
+   * chat classes (`classifyEndpoints`) and walked best-first
+   * (ChatThinking -> ChatBalanced -> ChatFast), with the small built-in
+   * `FALLBACK_MODEL_IDS` list as the floor when the catalogue can't be
+   * read. Set this to
+   * pin a regulated workspace to an approved subset, or to put custom
+   * endpoints in front of the auto-classified catalogue.
    */
   defaultModelFallbacks?: readonly string[];
   /**
@@ -166,7 +262,7 @@ export interface MastraPluginConfig extends BasePluginConfig {
    * processor that strips `chartId` fields from prior assistant
    * tool-invocation results before they reach the model. This
    * prevents the model from reusing turn-scoped chartIds it sees
-   * in memory recall (which would leave `[[chart:<id>]]` markers
+   * in memory recall (which would leave `[chart:<id>]` markers
    * pointing at writer events that no longer exist).
    *
    * Set to `false` to opt out - useful if a non-default agent
@@ -189,4 +285,159 @@ export interface MastraPluginConfig extends BasePluginConfig {
    * and the style block leans on the model's recency bias.
    */
   styleInstructions?: string | false;
+  /**
+   * Genie spaces this plugin's agents can delegate to. One Mastra
+   * tool is registered per alias (`genie` for the well-known
+   * `default` alias, `genie_<alias>` otherwise). Each tool spins
+   * up a per-question Genie sub-agent that runs Databricks
+   * "agent mode" against the space, broadcasts wire events to the
+   * UI, fetches statement rows for non-empty results, and returns
+   * a `(string | data | chart)[]` summary the host UI renders
+   * inline.
+   *
+   * Entries accept either a full {@link GenieSpaceConfig} object
+   * or a bare `space_id` string when no extras are needed:
+   *
+   * ```ts
+   * mastra({
+   *   genieSpaces: {
+   *     default: "01ef0d3c0e1b1f4a8d2c3e4f5a6b7c8d",
+   *     forecasts: { spaceId: "01ef...", hint: "weekly demand forecasts" },
+   *   },
+   * });
+   * ```
+   *
+   * Reach the spaces from an agent's `tools(plugins)` callback via
+   * `plugins.genie?.toolkit()`; the resulting tools accept
+   * `{ content, conversationId? }` and return a hydrated summary.
+   *
+   * **Fallback discovery** (highest precedence first): if this
+   * field is omitted, the Genie agent also picks up spaces from
+   * (1) the AppKit `genie({ spaces: { ... } })` plugin instance
+   * when registered, and (2) the `DATABRICKS_GENIE_SPACE_ID`
+   * env var (registered under the `default` alias). This keeps
+   * existing AppKit deployments working without restating the
+   * spaces config in two places.
+   */
+  genieSpaces?: GenieSpacesConfig;
+  /**
+   * TTL for the in-memory Genie space metadata cache, in
+   * milliseconds. Defaults to 5 minutes. The Genie agent calls
+   * `client.genie.getSpace(...)` on every cold-start to get the
+   * title / description / warehouse id; cached responses skip the
+   * round-trip and concurrent callers coalesce on a single
+   * in-flight fetch. Drop to a smaller value when analysts are
+   * actively editing space metadata and you want changes visible
+   * within seconds; raise it to amortise the round-trip when
+   * space metadata is effectively frozen.
+   *
+   * Backed by AppKit's `CacheManager`, so the cache participates
+   * in telemetry spans (`cache.getOrExecute`) and benefits from
+   * Lakebase persistence when the `lakebase` plugin is wired up.
+   */
+  genieSpaceCacheTtlMs?: number;
+  /**
+   * Maximum LLM steps each agent gets per turn. One step = one
+   * round-trip to the underlying model (a tool call consumes a
+   * step, the final-text reply consumes one too). Applies to
+   * every agent registered through {@link MastraPluginConfig.agents}
+   * - per-agent overrides aren't surfaced yet because the same
+   * ceiling has been sufficient across every workload we've run.
+   *
+   * Defaults to {@link DEFAULT_AGENT_MAX_STEPS} (25), sized to fit
+   * a decomposed Genie turn (grounding + several `ask_genie` calls
+   * + `prepare_chart` per dataset + the final-text reply) with
+   * headroom for the model to chain a couple of follow-ups before
+   * answering. Mastra's own `agent.generate` default of 5 would
+   * cut multi-step orchestration off after 2-3 tool calls, so
+   * explicitly raising the ceiling here is what lets the
+   * agent-mode loop play out.
+   *
+   * Lower when an unusually slow or expensive model makes long
+   * turns unaffordable; raise for exploratory workloads that need
+   * to drill deep into a dataset within a single turn.
+   */
+  agentMaxSteps?: number;
+  /**
+   * Wire Mastra spans into AppKit's global OTel pipeline via
+   * `@mastra/otel-bridge`.
+   *
+   * - `undefined` (default, auto): on only when
+   *   `OTEL_EXPORTER_OTLP_ENDPOINT` or
+   *   `OTEL_EXPORTER_OTLP_TRACES_ENDPOINT` is set. When unset, the
+   *   bridge is skipped so Mastra does not log
+   *   `[OtelBridge] No OTEL span found` on the noop tracer.
+   * - `true`: force on even without an OTLP endpoint.
+   * - `false`: force off.
+   */
+  observability?: boolean;
+  /**
+   * Log user feedback (thumbs up/down + freeform comments) to MLflow as
+   * trace assessments, and surface the feedback controls in the chat UI.
+   *
+   * - `undefined` (default, auto): enabled only when MLflow tracing is
+   *   wired - an OTLP exporter endpoint is set and an MLflow experiment
+   *   is named (the same signals the observability pipeline needs to
+   *   ship traces to MLflow). Otherwise off, since there'd be no trace
+   *   to attach feedback to.
+   * - `true`: force on. Feedback controls show and writes are attempted
+   *   regardless of env detection (use when the env is configured in a
+   *   way the auto-probe doesn't recognize).
+   * - `false`: force off. No trace-id header, no feedback route, no UI.
+   *
+   * Feedback attaches to a turn's MLflow trace via the OpenTelemetry
+   * trace id the server stamps on each response; see `mlflow.ts`.
+   */
+  feedback?: boolean;
+  /**
+   * Expose the plugin's agents (and optionally tools) as a Mastra MCP
+   * server so external MCP clients - Claude Desktop, Cursor, the Mastra
+   * playground, or another agent - can call them over the standard MCP
+   * transports. Enabled by default (agents only): wrapping the
+   * already-registered agents costs nothing extra, so the endpoint is on
+   * out of the box; only the ambient tools (which assume an in-process
+   * chat turn) stay off unless explicitly opted in.
+   *
+   * - `undefined` (default) / `true`: expose every registered agent as
+   *   an `ask_<agentId>` MCP tool under a server whose id is the plugin
+   *   name.
+   * - `false`: no MCP endpoints.
+   * - {@link MastraMcpConfig}: fine-grained control over the server id,
+   *   advertised metadata, and which agents / tools are exposed.
+   *
+   * When enabled, the stock Mastra MCP routes mount under the plugin's
+   * base path (no bespoke route is added - the server is handed to the
+   * `Mastra` instance via `mcpServers`, which `@mastra/express` serves):
+   *
+   * - Streamable HTTP: `POST /api/<plugin>/mcp/<serverId>/mcp`
+   * - SSE (legacy):    `GET  /api/<plugin>/mcp/<serverId>/sse`
+   *                    `POST /api/<plugin>/mcp/<serverId>/messages`
+   *
+   * Requests run under the same AppKit OBO scope as the chat routes, so
+   * an agent invoked over MCP resolves its model and tools as the
+   * calling user.
+   */
+  mcp?: boolean | MastraMcpConfig;
+  /**
+   * How much of the stock `@mastra/express` management API is reachable
+   * through the plugin mount. `@mastra/express` registers its full route
+   * table (agent inference plus admin / mutating routes: direct tool
+   * execution, workflow control, raw memory read/write, telemetry, logs,
+   * scores). AppKit already authenticates every request as the OBO user,
+   * but nothing there restricts *which* of those operations the browser
+   * client may invoke.
+   *
+   * - `"scoped"` (default): only the routes the chat client legitimately
+   *   needs are dispatched to Mastra - agent inference
+   *   (`stream` / `generate` / `network`), read-only agent metadata, this
+   *   plugin's own OBO- and resource-scoped `/route/*` routes (history /
+   *   threads), and, when {@link mcp} is enabled, the MCP transport.
+   *   Everything else (tool execution, workflow control, raw memory,
+   *   telemetry, logs, scores, and other mutations) is rejected with
+   *   `403` before it reaches Mastra.
+   * - `"full"`: dispatch the entire stock Mastra API. Use only for a
+   *   trusted first-party console that genuinely needs the management
+   *   surface.
+   */
+  apiAccess?: "scoped" | "full";
 }