@dbx-tools/appkit-mastra 0.1.12 → 0.1.18

package/src/config.ts

@@ -7,9 +7,14 @@
 
 import type { BasePluginConfig, getExecutionContext } 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 { GenieSpacesConfig } from "./genie.js";
 
 /**
  * `RequestContext` key under which {@link MastraServer} stores the
@@ -18,6 +23,51 @@ 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";
+
+/**
+ * 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;
@@ -189,4 +239,74 @@ 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 the Genie agent gets per turn.
+   * Each step is one round-trip to the underlying model
+   * (`get_space_description`, each `ask_genie` call, and the
+   * mandatory `submit_summary` closer all consume one step).
+   *
+   * Defaults to 16, which fits ~10 `ask_genie` calls plus
+   * grounding plus the summary closer - matches the
+   * decomposition pattern the Genie agent instructions prescribe
+   * ("2-6 ask_genie calls for any non-trivial question").
+   * Mastra's own `agent.generate` default of 5 would cut the
+   * Genie agent off after 2-3 Genie calls, so explicitly raising
+   * the ceiling here is what lets the agent-mode loop play out.
+   *
+   * Lower this when an unusually slow or expensive model makes
+   * long turns unaffordable; raise it for exploratory workloads
+   * that need to drill into a dataset.
+   */
+  genieAgentMaxSteps?: number;
 }