@dbx-tools/appkit-mastra 0.3.44 → 0.4.1

package/lib/src/plugin.d.ts

@@ -0,0 +1,352 @@
+/**
+ * AppKit plugin that builds one or more Mastra `Agent` instances and
+ * mounts the `@mastra/express` server. Clients drive the conversation
+ * over the standard Mastra agent stream (`@mastra/client-js`'s
+ * `getAgent(id).stream()`), so there's no bespoke chat transport to
+ * keep in sync.
+ *
+ * - Agents: registered through `config.agents` at plugin creation
+ *   ({@link MastraAgentDefinition}). Each entry's `tools` field accepts
+ *   either a plain record or a `(plugins) => tools` callback that gets
+ *   a typed sibling-plugin index ({@link MastraPlugins}). Omit
+ *   `config.agents` to get a single built-in `default` analyst.
+ * - Model: each agent call resolves a `MastraModelConfig` via
+ *   {@link buildModel} from `./model.js`. Per-agent `model` overrides
+ *   (`AgentConfig["model"]` or a `modelId` string) flow through
+ *   {@link buildAgents}.
+ * - Memory / storage: per-agent, built by {@link createMemoryBuilder}
+ *   from `./memory.js`. Both auto-default to `true` when the
+ *   `lakebase` plugin is registered (unless the caller passed
+ *   `false` or a custom config). Storage namespaces per agent via
+ *   {@link agentStorageSchemaName} per agent; the vector store is a single
+ *   shared singleton across every agent.
+ * - Server: the Express subapp wiring lives in `./server.js`.
+ * - HTTP: AppKit mounts this plugin under `/api/mastra`. Alongside the
+ *   Mastra agent routes, the plugin registers `/route/history`
+ *   (load + clear a thread's messages), `/route/threads` (list the
+ *   caller's conversations + delete one), `/models`, `/default-model`,
+ *   `/suggestions`, `/route/feedback` (log a thumbs / comment to MLflow
+ *   when feedback is enabled), and the generic `/embed/:type/:id`
+ *   resolver for inline chart / data markers. Each is registered through
+ *   AppKit's `route()` so it lands in the plugin's endpoint map, and
+ *   every outbound call one makes runs through `execute()` (see
+ *   `./defaults.js`). The stock `@mastra/express` surface is gated
+ *   by `config.apiAccess` (default `"scoped"`): only agent inference,
+ *   read-only agent metadata, the `/route/*` routes, and (when enabled)
+ *   MCP are dispatched to Mastra; admin / mutating / bulk-export routes
+ *   are refused with `403`. See {@link isMastraRequestAllowed}.
+ * - MCP: opt in with `config.mcp` to expose the agents (and optionally
+ *   tools) as a Mastra `MCPServer`. It is registered on the `Mastra`
+ *   instance via `mcpServers`, so `@mastra/express` serves the stock
+ *   MCP transport routes (`/mcp/<serverId>/...`) under the mount. See
+ *   `./mcp.js`.
+ *
+ * @module
+ */
+import { Plugin, type IAppRouter, type PluginManifest, type ResourceRequirement } from "@databricks/appkit";
+import { type ServingEndpointSummary } from "@dbx-tools/shared-model";
+import type { Agent } from "@mastra/core/agent";
+import { Mastra } from "@mastra/core/mastra";
+import type { RequestContext } from "@mastra/core/request-context";
+import { type MastraPluginConfig } from "./config.js";
+import { MastraServer } from "./server.js";
+/**
+ * AppKit plugin (registered name: `mastra`) that hosts Mastra agents
+ * with optional Lakebase-backed memory and AI SDK chat routes under
+ * the plugin mount (typically `/api/mastra`).
+ *
+ * @example Register the plugin
+ * ```ts
+ * import { createApp, lakebase } from "@databricks/appkit";
+ * import { createAgent, mastra } from "@dbx-tools/appkit-mastra";
+ *
+ * // `lakebase` first: registering it auto-enables Mastra storage + memory.
+ * const app = await createApp({
+ *   plugins: [
+ *     lakebase(),
+ *     mastra({
+ *       genieSpaces: { default: process.env.DATABRICKS_GENIE_SPACE_ID! },
+ *       agents: createAgent({
+ *         name: "analyst",
+ *         instructions: "You answer questions about revenue and returns.",
+ *         tools(plugins) {
+ *           return { ...plugins.genie?.toolkit() };
+ *         },
+ *       }),
+ *     }),
+ *   ],
+ * });
+ * ```
+ *
+ * @example Read the agents back off the AppKit instance
+ * ```ts
+ * const agentIds = app.mastra.list();
+ * const endpoints = await app.mastra.listModels();
+ * ```
+ */
+export declare class MastraPlugin extends Plugin<MastraPluginConfig> {
+    static manifest: PluginManifest<"mastra">;
+    /**
+     * Tighten resource requirements based on which features are enabled.
+     * AppKit calls this at registration time (config-aware) so disabled
+     * features don't surface their resource asks to the host app.
+     */
+    static getResourceRequirements(config: MastraPluginConfig): ResourceRequirement[];
+    private logger;
+    private built;
+    private mastra;
+    private mastraApp;
+    private mastraServer;
+    /**
+     * The optional MCP server exposing this plugin's agents / tools, or
+     * `null` when `config.mcp` is disabled (the default). Built in
+     * {@link buildAgentAndServer} and registered on the Mastra instance.
+     */
+    private mcp;
+    /**
+     * Dedicated service-principal Lakebase pool backing Mastra memory /
+     * storage. Built once in {@link buildAgentAndServer} (outside any
+     * `asUser` scope, so it never inherits a request's OBO identity) and
+     * drained in {@link shutdown}. `null` until setup runs or when
+     * Lakebase isn't needed.
+     */
+    private servicePrincipalPool;
+    setup(): Promise<void>;
+    /**
+     * When the `lakebase` plugin is registered, auto-enable `storage`
+     * and `memory` unless the caller opted out explicitly (`false` or a
+     * custom config object). Run after `setup:complete` so the lookup
+     * is reliable: any plugin that registers itself synchronously is
+     * already in the registry by the time this fires.
+     */
+    private applyLakebaseAutoDefaults;
+    /**
+     * Drain the memory service-principal pool. Idempotent: the handle is
+     * cleared before the drain starts, so a second call is a no-op and a
+     * later `setup()` rebuilds the pool. Bounded by
+     * {@link POOL_DRAIN_TIMEOUT_MS} to stay well inside the 15s graceful
+     * shutdown budget.
+     */
+    shutdown(): Promise<void>;
+    /**
+     * Abort in-flight work. AppKit's graceful shutdown calls this hook
+     * synchronously and never awaits {@link shutdown}, so the pool drain is
+     * started here too; it cannot be awaited from a `void` hook, and
+     * {@link shutdown} is idempotent so the duplicate call is free.
+     */
+    abortActiveOperations(): void;
+    exports(): {
+        /**
+         * Ids of every registered agent in registration order. Matches
+         * AppKit `agents.list()` so callers can iterate the registry the
+         * same way under both plugins.
+         */
+        list: () => string[];
+        /**
+         * Look up a registered agent by id. Returns `null` (not
+         * undefined) when unknown so call sites can early-return without
+         * a separate `in` check.
+         */
+        get: (id: string) => Agent | null;
+        /**
+         * The agent the client converses with when it doesn't name one.
+         * Resolves to `config.defaultAgent`, the first registered id, or
+         * the built-in `default` fallback.
+         */
+        getDefault: () => Agent | null;
+        /** Underlying Mastra instance for advanced use (custom routes etc.). */
+        getMastra: () => Mastra<Record<string, Agent<any, import("@mastra/core/agent").ToolsInput, undefined, unknown, import("@mastra/core/agent").AgentEditorConfig | undefined>>, Record<string, import("@mastra/core/workflows").AnyWorkflow>, Record<string, import("@mastra/core/vector").MastraVector<any>>, Record<string, import("@mastra/core/tts").MastraTTS>, import("@mastra/core/logger").IMastraLogger, Record<string, import("@mastra/core/mcp").MCPServerBase<any>>, Record<string, import("@mastra/core/evals").MastraScorer<any, any, any, any>>, Record<string, import("@mastra/core/tools").ToolAction<any, any, any, any, any, any, unknown>>, Record<string, import("@mastra/core/processors").Processor<any, unknown>>, Record<string, import("@mastra/core/memory").MastraMemory>, Record<string, import("@mastra/core/channels").ChannelProvider>> | null;
+        /**
+         * Build the `RequestContext` an agent turn driven from OUTSIDE this
+         * plugin's routes needs (a Teams activity on another plugin's endpoint, a
+         * scheduled job).
+         *
+         * Mastra's user-scoped tools - `ask_genie` most visibly - read the AppKit
+         * user off the request context, which only the HTTP middleware stamps.
+         * Calling `agent.generate` with a raw prompt therefore answers "the data
+         * source is unreachable" where the chat routes answer with real data.
+         * Passing this as `requestContext` closes that gap, so an out-of-band turn
+         * has exactly the capabilities a chat turn has.
+         *
+         * Must be called INSIDE an `asUser(req)` scope to inherit the caller's
+         * OBO identity; outside one it resolves to the service principal.
+         */
+        createRequestContext: (options?: {
+            threadId?: string;
+            resourceId?: string;
+            requestId?: string;
+        }) => Promise<RequestContext>;
+        /**
+         * MCP endpoint info when `config.mcp` is enabled, else `null`.
+         * Streamable HTTP is `http`; the SSE pair is the legacy transport.
+         *
+         * Each path is given twice: mount-relative (`httpPath`, `ssePath`,
+         * `messagePath`) for anything that already knows where the plugin is
+         * mounted, and absolute (`http`, `sse`, `messages`) for MCP clients,
+         * which take a single URL and cannot compose one. The absolute form is
+         * built from {@link basePath}, so it honors a `config.name` override but
+         * still assumes AppKit's default `/api/<name>` mount.
+         */
+        getMcp: () => {
+            serverId: string;
+            httpPath: string;
+            ssePath: string;
+            messagePath: string;
+            http: string;
+            sse: string;
+            messages: string;
+        } | null;
+        /** Express subapp Mastra is mounted on; mostly for tests. */
+        getMastraServer: () => MastraServer | null;
+        /**
+         * Fetch the workspace's Model Serving endpoints (cached). Same
+         * payload the `GET /models` route returns; surfaced here so
+         * other plugins / scripts can introspect the catalogue without
+         * an HTTP round-trip. AppKit wraps this with `asUser(req)` for
+         * OBO scoping automatically. Throws when the listing fails, since
+         * there is no status code to hand back on this surface.
+         */
+        listModels: () => Promise<ServingEndpointSummary[]>;
+        /**
+         * Force-evict cached endpoint listings via AppKit's
+         * `CacheManager`. Useful in tests or right after an admin
+         * deploys a new endpoint and doesn't want to wait for the TTL.
+         * Returns the underlying `CacheManager.delete`/`clear` promise.
+         */
+        clearModelsCache: (host?: string) => Promise<void>;
+    };
+    /**
+     * Absolute mount this plugin's routes answer on. AppKit mounts every plugin
+     * at `/api/<plugin.name>` and `this.name` honors a `config.name` override,
+     * so the per-route segments (`routes.MASTRA_ROUTES`) hang off this.
+     */
+    private get basePath();
+    clientConfig(): Record<string, unknown>;
+    /**
+     * Whether user feedback can be logged to MLflow. Delegates to
+     * {@link resolveFeedbackEnabled} so the client-config flag and the
+     * feedback route share the same gate as the server's trace-id header.
+     */
+    private feedbackEnabled;
+    injectRoutes(router: IAppRouter): void;
+    /**
+     * Map a failed {@link ExecutionResult} onto the response: the status the
+     * execution pipeline resolved, paired with `clientMessage` rather than the
+     * result's own text, which can carry upstream detail. The full message is
+     * logged under `<operation>:error`.
+     */
+    private sendFailure;
+    /** 404 body for a request naming an agent that isn't registered. */
+    private unknownAgentMessage;
+    /**
+     * Invoke the Mastra express sub-app. Exists as a method (instead of reading
+     * `this.mastraApp` through the `asUser(req)` proxy at the call site) so the
+     * proxy binds this plain method - whose `.bind` is `Function.prototype.bind`
+     * - rather than the express app, whose `.bind` is the HTTP BIND route
+     * registrar (see the note in `injectRoutes`). Runs inside the user scope so
+     * `getExecutionContext()` returns the OBO client for the agent/model
+     * resolvers.
+     */
+    private dispatchMastra;
+    /**
+     * Map a clean, mount-relative MCP alias path to the underlying
+     * `@mastra/express` route. Returns `null` when MCP is off or the path
+     * isn't an alias. Collapses the stock `/mcp/<serverId>/<transport>`
+     * layout (serverId defaults to the plugin name) down to `/mcp`,
+     * `/sse`, and `/messages`.
+     */
+    private mcpRouteAlias;
+    /**
+     * Implementation backing the `/suggestions` route. Runs inside the
+     * AppKit user-context proxy so `getExecutionContext()` returns the
+     * OBO-scoped client. Resolves the plugin's Genie spaces and merges
+     * their curated `sample_questions` (see {@link collectSpaceSuggestions}).
+     * Returns `[]` when no Genie space is configured so the client
+     * shows a bare empty state instead of built-in example prompts.
+     */
+    private fetchSuggestions;
+    /**
+     * Implementation backing the `/route/feedback` route. Runs inside the
+     * AppKit user-context proxy so `getExecutionContext()` returns the
+     * OBO-scoped client and the assessment is attributed to the signed-in
+     * user (their email / id as the assessment source). Returns the
+     * created assessment id on success, or `undefined` on a soft failure
+     * (see {@link logFeedback} in `./mlflow.js`).
+     */
+    private logFeedback;
+    /**
+     * Implementation backing the `data` embed resolver
+     * (`GET /embed/data/:id`). Runs inside the AppKit user-context proxy so
+     * `getExecutionContext()` returns the OBO-scoped workspace
+     * client, then reuses the same `fetchStatementData` pipeline
+     * the `get_statement` tool runs so the LLM and the UI see the
+     * exact same shape for the same statement.
+     *
+     * Resolves to `undefined` data for upstream 404s so the route can map them
+     * to a clean HTTP 404; any other failure comes back as a non-`ok` result
+     * carrying the status to answer with.
+     */
+    private fetchStatement;
+    /**
+     * Implementation backing the `chart` embed resolver
+     * (`GET /embed/chart/:id`). Runs inside the AppKit user-context proxy so the
+     * lookup is namespaced to the calling identity: a chart minted for another
+     * user is indistinguishable from an unknown id, and the route answers 404.
+     */
+    private fetchChartEntry;
+    /**
+     * Implementation backing the `/models` route. Runs inside the AppKit
+     * user-context proxy so `getExecutionContext()` returns the OBO-scoped
+     * client and the catalogue reflects what the caller can invoke.
+     */
+    private listModelsResult;
+    /**
+     * Implementation backing the `listModels` export. The programmatic surface
+     * carries no status code, so a failed execution becomes a throw with a
+     * stable message; the routes branch on the {@link ExecutionResult} instead.
+     */
+    private listModels;
+    private buildAgentAndServer;
+}
+/**
+ * Register the Mastra plugin on an AppKit app. Mounts the agents, the
+ * `/route/*` chat surface, and (unless disabled) the MCP transport under
+ * `/api/mastra`.
+ *
+ * @example Minimal app
+ * ```ts
+ * import { createApp } from "@databricks/appkit";
+ * import { mastra } from "@dbx-tools/appkit-mastra";
+ *
+ * // No `agents`: a single built-in `default` analyst is registered, so the
+ * // chat endpoint works out of the box.
+ * const app = await createApp({ plugins: [mastra()] });
+ * ```
+ *
+ * @example Two agents, a pinned model, and Genie
+ * ```ts
+ * import { createApp, lakebase } from "@databricks/appkit";
+ * import { createAgent, mastra } from "@dbx-tools/appkit-mastra";
+ *
+ * const app = await createApp({
+ *   plugins: [
+ *     lakebase(),
+ *     mastra({
+ *       defaultModel: "databricks-claude-sonnet-4-6",
+ *       defaultAgent: "analyst",
+ *       genieSpaces: { sales: { spaceId: "01ef...", hint: "orders, returns" } },
+ *       agents: {
+ *         analyst: createAgent({
+ *           instructions: "You answer questions about sales.",
+ *           tools(plugins) {
+ *             return { ...plugins.genie?.toolkit() };
+ *           },
+ *         }),
+ *         helper: createAgent({ instructions: "You explain the analyst's answers." }),
+ *       },
+ *     }),
+ *   ],
+ * });
+ * ```
+ */
+export declare const mastra: import("@databricks/appkit").ToPlugin<typeof MastraPlugin, MastraPluginConfig, "mastra">;