@dbx-tools/appkit-mastra 0.1.12 → 0.1.18

package/dist/src/config.js

@@ -4,9 +4,52 @@
  * Kept in a leaf module so `plugin.ts`, `server.ts`, `model.ts`, and
  * `memory.ts` can import them without creating a cycle.
  */
+import { MASTRA_RESOURCE_ID_KEY, MASTRA_THREAD_ID_KEY, } from "@mastra/core/request-context";
 /**
  * `RequestContext` key under which {@link MastraServer} stores the
  * resolved AppKit user. `model.ts` reads it to mint user-scoped
  * Databricks tokens.
  */
 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 = [
+    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",
+];

package/dist/src/genie.d.ts

@@ -1,131 +1,194 @@
 /**
- * Mastra tool wrappers around the AppKit `genie` plugin's exports.
+ * Genie agent for Mastra.
  *
- * One `sendMessage` tool is registered per configured space alias so
- * the LLM picks the space by tool selection (the description bakes the
- * alias in). `getConversation` is registered once, taking `alias` as a
- * parameter.
+ * Each configured Genie space exposes a single Mastra tool to the
+ * calling agent (`genie` for the `"default"` alias, `genie_<alias>`
+ * otherwise). When invoked, the tool runs end-to-end:
  *
- * All Genie payload types are inferred from the public `genie` factory
- * (`genie().plugin` constructor → `exports()` return type), so any
- * upstream change in `@databricks/appkit` flows in automatically.
+ *   1. Pulls the per-request {@link WorkspaceClient} off
+ *      `ctx.requestContext` (stamped by `MastraServer`) and emits a
+ *      `started` writer event so the host UI can show progress
+ *      immediately, before any LLM round-trip.
+ *   2. Spins up a per-call inner Mastra `Agent` with three tools:
+ *        - `ask_genie`: drives one `genieEventChat` turn, fetches
+ *          the matching statement's rows when the turn ran SQL,
+ *          and forwards every wire event (status, thinking, sql,
+ *          rows) through `ctx.writer` for streaming UI updates.
+ *        - `get_space_description`: cheap title / description /
+ *          warehouse id lookup for grounding.
+ *        - `get_space_serialized`: full `GenieSpace` JSON for
+ *          column-level grounding when the description isn't
+ *          enough.
+ *   3. Runs the inner agent with `structuredOutput` (Mastra's
+ *      two-pass mode + `jsonPromptInjection`) to coerce the
+ *      agent's final answer into a tagged
+ *      `[{type:"text"|"data", ...}]` array. The two-pass design
+ *      avoids Databricks Model Serving's `response_format` +
+ *      `tools` collision; prompt injection sidesteps the
+ *      separate `response_format` + streaming collision in the
+ *      structuring agent.
+ *   4. Charts every `data` item in parallel via
+ *      {@link runChartPlanner}, maps `text` items to the shared
+ *      {@link GenieSummaryItem} `string` variant, and returns the
+ *      hydrated {@link GenieAgentResult}.
  *
- * As Genie streams its long-running events (`FETCHING_METADATA` →
- * `ASKING_AI` → `EXECUTING_QUERY` → `COMPLETED`, plus SQL text and
- * follow-ups in `message_result.attachments`), the tool forwards a
- * normalised {@link GenieProgress} discriminated union out through
- * `ctx.writer` so the client can render an incremental loading pill.
- * Row payloads from `query_result` are intentionally discarded - the
- * LLM never sees rows, and charts come from the separate
- * `render_data` tool when the model decides one is useful.
+ * The legacy AppKit `genie` plugin (`@databricks/appkit`'s `genie`)
+ * is no longer used at runtime. The inner agent talks to Genie
+ * directly via `@dbx-tools/genie` (`genieEventChat`) and the
+ * workspace `statementExecution.getStatement` API. The plugin's
+ * `spaces` config is still honored so existing AppKit-style wiring
+ * keeps working without change.
  */
-import { genie } from "@databricks/appkit";
-import { createTool } from "@mastra/core/tools";
+import { type ChartEvent } from "@dbx-tools/appkit-mastra-shared";
+import { appkitUtils } from "@dbx-tools/shared";
+import type { RequestContext } from "@mastra/core/request-context";
+import type { MastraTools } from "./agents.js";
 import type { MastraPluginConfig } from "./config.js";
-/** Live AppKit `GeniePlugin` instance. */
-export type GeniePluginInstance = InstanceType<ReturnType<typeof genie>["plugin"]>;
-/** Full `exports()` shape of the AppKit `genie` plugin. */
-export type GenieExports = ReturnType<GeniePluginInstance["exports"]>;
+/** Default alias used when a single unnamed Genie space is wired up. */
+export declare const DEFAULT_GENIE_ALIAS = "default";
+/** Per-space Genie agent configuration. */
+export interface GenieSpaceConfig {
+    /** Genie `space_id`. Required; resolves via `client.genie.getSpace`. */
+    spaceId: string;
+    /**
+     * Optional human-readable description appended to the Genie
+     * tool's description so the calling LLM has hints about
+     * *what data* this space covers (e.g. "orders, returns,
+     * fulfillment"). When omitted, only the space's own
+     * `description` (fetched on first use) is shown.
+     */
+    hint?: string;
+}
+/** Map of alias -> space config. Accepts either explicit objects or bare space ids. */
+export type GenieSpacesConfig = Record<string, GenieSpaceConfig | string>;
 /**
- * Stream event yielded by `genie.exports().sendMessage`. Discriminated
- * by `type` (`"message_start" | "status" | "message_result" |
- * "query_result" | "error" | "history_info"`).
+ * Get the chart inventory map for this request, creating it on
+ * first access. Subsequent reads return the same map so callers
+ * mutate in place. The map is request-scoped (collected with the
+ * `RequestContext` at end of request), so there's no per-process
+ * leak.
  */
-export type GenieStreamEvent = ReturnType<GenieExports["sendMessage"]> extends AsyncGenerator<infer E> ? E : never;
-/** Conversation history returned by `genie.exports().getConversation`. */
-export type GenieConversation = Awaited<ReturnType<GenieExports["getConversation"]>>;
+export declare function chartInventoryFromContext(requestContext: RequestContext): Map<string, ChartEvent>;
 /**
- * Normalised progress event surfaced to the UI as a Mastra
- * `tool-output` chunk. Loading pill events (`started`, `status`,
- * `sql`, `suggested`, `error`) are pure UI metadata and never reach
- * the LLM.
- *
- * The `chart` variant is the wire shape emitted by
- * {@link emitChartWithPlanning} (used by both this Genie
- * draining loop and the system-level `render_data` tool). All
- * fields except `chartId` are optional because two events per
- * chartId arrive on the wire: the first carries the rows
- * (`title` + `description?` + `data`); the second, on planner
- * success, carries just the resolved Echarts spec (`option`).
- * The host UI's `<ChartSlot>` merges them by `chartId`.
+ * Options for {@link createGenieTool}. Only carries config that
+ * doesn't vary per request - the per-request {@link WorkspaceClient},
+ * `RequestContext`, writer, and abort signal flow through the
+ * tool's `execute(_, ctx)` and are not captured here.
  */
-export type GenieProgress = {
-    kind: "started";
-    conversationId: string;
-    messageId: string;
+export interface CreateGenieToolOptions {
+    /** Genie space id this tool targets. */
     spaceId: string;
-} | {
-    kind: "status";
-    status: string;
-    label: string;
-} | {
-    kind: "sql";
-    sql: string;
-    title?: string;
-    description?: string;
-    statementId?: string;
-} | {
-    kind: "chart";
-    chartId: string;
-    title?: string;
-    description?: string;
-    data?: Array<Record<string, unknown>>;
-    option?: Record<string, unknown>;
-} | {
-    kind: "text";
-    content: string;
-} | {
-    kind: "suggested";
-    questions: string[];
-} | {
-    kind: "error";
-    error: string;
-};
+    /** Plugin config; resolves the LLM and chart planner agent. */
+    config: MastraPluginConfig;
+    /** Override the registered tool id. Defaults to `"genie"`. */
+    toolId?: string;
+    /** Override the tool description shown to the calling LLM. */
+    toolDescription?: string;
+    /**
+     * Override the inner agent's max tool-loop steps. Defaults to
+     * {@link DEFAULT_MAX_STEPS}.
+     */
+    maxSteps?: number;
+}
 /**
- * Default tool name for a wired Genie alias. The well-known `default`
- * alias collapses to `genie`; everything else gets a `genie_` prefix so
- * multiple spaces stay disambiguated when an agent has more than one
- * wired. Matches the `genie` / `genie_<alias>` naming used elsewhere in
- * dbx-tools AppKit demos.
+ * Build the calling agent's Genie tool. The returned Mastra tool
+ * runs end-to-end on each invocation:
+ *
+ *   1. Pull the per-request `WorkspaceClient` off
+ *      `ctx.requestContext` (stamped by `MastraServer` under
+ *      {@link MASTRA_USER_KEY}) and emit a `started` writer
+ *      event so the host UI shows progress immediately.
+ *   2. Spin up the inner Mastra agent + three tools, fresh per
+ *      call so the row cache stays invocation-scoped.
+ *   3. Run the agent with `structuredOutput` against
+ *      {@link agentSummarySchema}. Mastra's two-pass design keeps
+ *      the inner loop tools-only (no `response_format`), so the
+ *      Databricks Model Serving `response_format`+`tools`
+ *      collision never fires.
+ *   4. Walk the returned `[text|data][]`, map `text` items to
+ *      shared `GenieSummaryItem.string`, and chart every `data`
+ *      item in parallel via {@link runChartPlanner} to a
+ *      `GenieSummaryItem.visualize`. Items referencing a missing
+ *      `statementId` are dropped with a warn log; chart-planner
+ *      failures leave `dataset.chart` unset so the host UI falls
+ *      back to a table.
+ */
+export declare function createGenieTool(opts: CreateGenieToolOptions): import("@mastra/core/tools").Tool<any, any, any, any, import("@mastra/core/tools").ToolExecutionContext<any, any, unknown>, string, unknown>;
+/**
+ * Default tool id for a wired Genie alias. The well-known
+ * `default` alias collapses to `genie`; every other alias gets a
+ * `genie_` prefix so multi-space registrations stay
+ * disambiguated.
  */
 export declare function defaultGenieToolName(alias: string): string;
 /**
- * Build one `sendMessage` tool per configured Genie alias plus a single
- * `getConversation` tool. Returns a record keyed by tool id, ready to
- * spread into an `Agent`'s `tools` map.
+ * Normalize the {@link GenieSpacesConfig} record. Bare-string
+ * entries (`{ default: "01ef..." }`) get wrapped as
+ * `{ spaceId: "01ef..." }`; object entries pass through unchanged.
+ * `undefined` and empty-string values are dropped so callers can
+ * pass `process.env.X` directly (matches AppKit `genie()`'s
+ * defensive treatment of unset env vars).
+ */
+export declare function normalizeGenieSpaces(spaces: GenieSpacesConfig | Record<string, string | GenieSpaceConfig | undefined> | undefined): Record<string, GenieSpaceConfig>;
+/**
+ * Discover Genie space aliases from every supported source and
+ * merge them into a single record. Precedence (highest first):
  *
- * `config` must be the active plugin config; Genie's
- * `query_result` events are routed through
- * {@link emitChartWithPlanning} which uses it to resolve the
- * chart-planner's model.
+ *   1. {@link MastraPluginConfig.genieSpaces} on the `mastra(...)`
+ *      call. Explicit Mastra wiring always wins so users can
+ *      override AppKit's defaults per-agent.
+ *   2. AppKit `genie({ spaces: { ... } })` plugin instance. Lets
+ *      users keep using the existing AppKit config format
+ *      (`genie({ spaces: { sales: "...", ops: "..." } })`)
+ *      without restating the same record on the Mastra plugin.
+ *      Read off the live plugin instance via a structural cast
+ *      since `Plugin.config` is TS-protected (not runtime-private).
+ *   3. `DATABRICKS_GENIE_SPACE_ID` env var (registered under the
+ *      well-known `default` alias). Matches the AppKit `genie()`
+ *      plugin's fallback behavior so a bare `mastra()` + `genie()`
+ *      pair just works.
+ *
+ * Aliases collide cleanly: a higher-precedence source's value
+ * replaces a lower one's wholesale. Sources that contribute zero
+ * aliases (or contribute only `undefined` / empty entries) are
+ * silently ignored.
+ */
+export declare function resolveGenieSpaces(config: MastraPluginConfig, context: appkitUtils.PluginContextLike | undefined): Record<string, GenieSpaceConfig>;
+/**
+ * Build one Mastra tool per configured Genie space. Each tool is
+ * a thin {@link createGenieTool} wrapper with the alias-derived
+ * id and a hint-flavored description so the calling LLM knows
+ * which space covers what data.
+ *
+ * Returns a record keyed by tool id, ready to spread into an
+ * `Agent`'s `tools` map (or surfaced via
+ * `plugins.genie?.toolkit()`).
  */
 export declare function buildGenieTools(opts: {
-    aliases: string[];
-    exports: GenieExports;
+    spaces: GenieSpacesConfig | Record<string, GenieSpaceConfig>;
     config: MastraPluginConfig;
-    signal?: AbortSignal;
-}): Record<string, ReturnType<typeof createTool>>;
+}): MastraTools;
 /**
- * Toolkit provider built from a live AppKit `GeniePlugin` instance.
- * Returned by {@link buildGenieProvider} so that
- * `plugins.genie?.toolkit()` inside an agent's `tools(plugins)` callback
- * resolves to the streaming-aware {@link buildGenieTools} record instead
- * of the AppKit default (which does one blocking call per tool with no
- * mid-flight events).
- *
- * The returned `toolkit()` reads alias names off the plugin's
- * `getAgentTools()` registry (each entry is `${alias}.sendMessage` or
- * `${alias}.getConversation`), then mints one `sendMessage` tool per
- * alias plus a shared `getConversation`. `sendMessage` / `getConversation`
- * are bound back to the plugin instance so they keep their `this`
- * (they are class methods, not free functions).
- *
- * `_opts` is accepted but unused for now - the streaming tools are an
- * all-or-nothing bundle. Wire `only` / `except` / `prefix` / `rename`
- * later if a caller needs them.
+ * Plugin-toolkit adapter so the `plugins.genie?.toolkit()` lookup
+ * inside an agent's `tools(plugins)` callback returns the
+ * Genie agent-backed tools instead of throwing on missing plugin.
+ * Mirrors AppKit's `PluginToolkitProvider` shape.
  */
-export declare function buildGenieProvider(plugin: GeniePluginInstance, opts: {
+export declare function buildGenieToolkitProvider(opts: {
+    spaces: GenieSpacesConfig | Record<string, GenieSpaceConfig>;
     config: MastraPluginConfig;
 }): {
-    toolkit(opts?: unknown): Record<string, ReturnType<typeof createTool>>;
+    toolkit(opts?: unknown): MastraTools;
 };
+/**
+ * Returns `true` when at least one Genie space is reachable
+ * through {@link resolveGenieSpaces} - either via
+ * {@link MastraPluginConfig.genieSpaces}, the AppKit `genie()`
+ * plugin instance, or the `DATABRICKS_GENIE_SPACE_ID` env var.
+ *
+ * Cheap to call from `resolveProvider` to short-circuit `genie`
+ * lookups when nothing is wired, so the `plugins.genie` lookup
+ * still resolves to `undefined` (matching AppKit's
+ * absent-plugin semantics) when neither source is configured.
+ */
+export declare function hasAnyGenieSpaces(config: MastraPluginConfig, context: appkitUtils.PluginContextLike | undefined): boolean;