@dbx-tools/appkit-mastra 0.1.4 → 0.1.9

package/src/agents.ts

@@ -10,21 +10,27 @@
  *
  * When no agents are registered the plugin falls back to a single
  * built-in analyst so the bare `mastra()` call still mounts a working
- * `chatRoute` agent for demos.
+ * streamable agent for demos.
  */
 
-import { genie } from "@databricks/appkit";
-import { logUtils, pluginUtils, stringUtils } from "@dbx-tools/appkit-shared";
-import { Agent } from "@mastra/core/agent";
 import type { AgentConfig, ToolsInput } from "@mastra/core/agent";
-import { createTool } from "@mastra/core/tools";
+import { Agent } from "@mastra/core/agent";
 import type { Tool } from "@mastra/core/tools";
+import { createTool } from "@mastra/core/tools";
+import type { Workspace } from "@mastra/core/workspace";
 import type { PgVectorConfig, PostgresStoreConfig } from "@mastra/pg";
 
-import type { MastraPluginConfig } from "./config.js";
-import { buildGenieProvider } from "./genie.js";
-import type { MemoryBuilder } from "./memory.js";
-import { buildModel } from "./model.js";
+import { buildRenderDataTool } from "./chart";
+import type { MastraPluginConfig } from "./config";
+import { buildGenieToolkitProvider, resolveGenieSpaces } from "./genie";
+import type { MemoryBuilder } from "./memory";
+import { buildModel } from "./model";
+import { fallback } from "@dbx-tools/model";
+import { ResultProcessor, stripStaleChartsProcessor } from "./processors";
+import { buildSummarizeTool } from "./summarize";
+import { createWorkspace } from "./workspaces";
+import { log, string } from "@dbx-tools/shared-core";
+import { plugin } from "@dbx-tools/appkit";
 
 /**
  * Tool record accepted by every Mastra `Agent.tools` field and by the
@@ -120,12 +126,12 @@ export interface AppKitToolOptions {
 
 /**
  * Build a deterministic Mastra tool id from a description.
- * Delegates to {@link stringUtils.toUniqueSlug}: slug + always-on
- * SHA-1 suffix so two tools with the same leading words don't
- * collide in traces. Stable across runs.
+ * Delegates to {@link string.toUniqueSlug}: slug + always-on
+ * 6-char FNV-1a base-32 suffix so two tools with the same leading
+ * words don't collide in traces. Stable across runs.
  */
 function deriveToolId(description: string): string {
-  return stringUtils.toUniqueSlug(description, { fallbackPrefix: "tool" });
+  return string.toUniqueSlug(description, { fallbackPrefix: "tool" });
 }
 
 /**
@@ -144,7 +150,8 @@ function deriveToolId(description: string): string {
  * type inference and to match the AppKit API surface.
  */
 export function createAgent<T extends MastraAgentDefinition>(def: T): T {
-  return def;
+  const workspace = def.workspace ?? createWorkspace();
+  return { ...def, workspace };
 }
 
 /**
@@ -217,15 +224,16 @@ export interface MastraPluginToolkitProvider {
 export type MastraPlugins = Record<string, MastraPluginToolkitProvider>;
 
 /** Function form of {@link MastraAgentDefinition.tools}. */
-export type MastraToolsFn = (
-  plugins: MastraPlugins,
-) => MastraTools | Promise<MastraTools>;
+export type MastraToolsFn = (plugins: MastraPlugins) => MastraTools | Promise<MastraTools>;
+
+/** Function form of {@link MastraAgentDefinition.workspace}. */
+export type MastraAgentWorkspaceResolver = () => Workspace | undefined;
 
 /**
  * A code-defined Mastra agent. Mirrors the shape AppKit's `agents`
  * plugin uses for `AgentDefinition`. The registry key under
- * `config.agents` is what `chatRoute` matches on; `name` is purely
- * informational (defaults to the key).
+ * `config.agents` is the `agentId` the client streams against; `name`
+ * is purely informational (defaults to the key).
  */
 export interface MastraAgentDefinition {
   /** Display name used as `Agent.name`. Defaults to the registry key. */
@@ -274,7 +282,7 @@ export interface MastraAgentDefinition {
    *
    * - `undefined` (default): inherit `config.storage`. When that's
    *   enabled, the agent gets its **own per-agent `PostgresStore`**
-   *   keyed by `schemaName: "mastra_<agentId>"` so threads and
+   *   keyed by {@link agentStorageSchemaName} so threads and
    *   messages stay isolated between agents in the same database.
    * - `false`: disable storage for this agent only (purely in-memory).
    * - `true`: enable with the per-agent default schema.
@@ -282,6 +290,12 @@ export interface MastraAgentDefinition {
    *   `PostgresStore` config (custom schema, connection, etc.).
    */
   storage?: boolean | MastraStorageConfigOverride;
+  /**
+   * Mastra {@link Workspace} for this agent (filesystem, sandbox, or other
+   * providers). Assistant skills from Databricks workspace paths are wired
+   * via {@link createWorkspace}.
+   */
+  workspace?: Workspace | MastraAgentWorkspaceResolver;
 }
 
 /**
@@ -290,19 +304,16 @@ export interface MastraAgentDefinition {
  * strip `id`. The built-in `Omit` collapses unions to one shape with
  * common fields only, which loses the connection-style discriminants.
  */
-type DistributiveOmit<T, K extends keyof never> = T extends unknown
-  ? Omit<T, K>
-  : never;
+type DistributiveOmit<T, K extends keyof never> = T extends unknown ? Omit<T, K> : never;
 
 /**
  * `PostgresStoreConfig` minus `id` - per-agent overrides accept any
  * Mastra-supported storage shape. `id` is filled in automatically
  * from the agent registry key so traces stay stable.
  */
-export type MastraStorageConfigOverride = DistributiveOmit<
-  PostgresStoreConfig,
-  "id"
-> & { id?: string };
+export type MastraStorageConfigOverride = DistributiveOmit<PostgresStoreConfig, "id"> & {
+  id?: string;
+};
 
 /**
  * `PgVectorConfig` minus `id` - per-agent overrides accept any
@@ -317,11 +328,28 @@ export type MastraMemoryConfigOverride = DistributiveOmit<PgVectorConfig, "id">
 export interface BuiltAgents {
   agents: Record<string, Agent>;
   defaultAgentId: string;
+  /**
+   * Ambient tools shared across every agent (the built-in system tools
+   * spread with `config.tools`). Surfaced so the optional MCP server
+   * can re-expose them when {@link MastraMcpConfig.tools} is enabled.
+   */
+  ambientTools: MastraTools;
 }
 
 /** Fallback agent id used when `config.agents` is omitted entirely. */
 export const FALLBACK_AGENT_ID = "default";
 
+/**
+ * Default per-turn step ceiling applied to every registered agent
+ * when {@link MastraPluginConfig.agentMaxSteps} is unset. 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 - well above Mastra's own `agent.generate` default of
+ * 5, which would cut multi-step orchestration off mid-loop.
+ */
+export const DEFAULT_AGENT_MAX_STEPS = 25;
+
 const FALLBACK_AGENT_INSTRUCTIONS = `You are a data analyst. The user will ask questions about
 business metrics and may share personal preferences you should remember across turns.
 
@@ -344,16 +372,21 @@ Rules:
  * Override globally via {@link MastraPluginConfig.styleInstructions}
  * (pass `false` to disable entirely, or a string to replace).
  */
-export const DEFAULT_STYLE_INSTRUCTIONS = `Output style:
-
-- Plain prose. Use hyphens (-) only. Never use em dashes (—) or en dashes (–).
-- Never use emojis.
-- Skip openers like "Great question", "Absolutely", "I'd be happy to help".
-- Skip closers like "Let me know if you have any questions".
-- Skip self-disclaimers ("I should mention", "It's important to note").
-- Answer directly. No preamble before the actual answer.
-- Use lists and headers only when they clarify a multi-part answer; not for short replies.
-- Quote numbers, code, identifiers, and tool output verbatim. Never paraphrase them.`;
+export const DEFAULT_STYLE_INSTRUCTIONS = [
+  "Output style:",
+  "",
+  "Use markdown formatting, including headings, lists, and code blocks.",
+  "Avoid lists and headers for short replies.",
+  "Plain prose.",
+  "Use hyphens (-) only. Never use em dashes or en dashes.",
+  "Never use emojis.",
+  "Skip openers like 'Great question', 'Absolutely', and 'I'd be happy to help'.",
+  "Skip closers like 'Let me know if you have any questions'.",
+  "Skip self-disclaimers like 'I should mention' and 'It's important to note'.",
+  "Answer directly.",
+  "Do not include a preamble before the actual answer.",
+  "Use lists and headers only when they clarify a multi-part answer.",
+].join("\n");
 
 /**
  * Resolve the style block to append to every agent's instructions.
@@ -364,6 +397,7 @@ function resolveStyleInstructions(config: MastraPluginConfig): string | null {
   if (typeof config.styleInstructions === "string") {
     return config.styleInstructions;
   }
+
   return DEFAULT_STYLE_INSTRUCTIONS;
 }
 
@@ -371,14 +405,21 @@ function resolveStyleInstructions(config: MastraPluginConfig): string | null {
  * Join an agent's bespoke instructions with the resolved style block.
  * Returns the bespoke text unchanged when the style block is disabled.
  */
-function composeInstructions(
-  agentInstructions: string,
-  style: string | null,
-): string {
+function composeInstructions(agentInstructions: string, style: string | null): string {
   if (!style) return agentInstructions;
   return `${agentInstructions.trimEnd()}\n\n${style}`;
 }
 
+/**
+ * Generated description for an agent whose definition omits one. Kept
+ * generic (the agent's own name is the only reliable signal at build
+ * time) so it reads sensibly as the `ask_<id>` MCP tool's description
+ * and in agent metadata.
+ */
+function defaultAgentDescription(name: string): string {
+  return `Conversational AI agent "${name}".`;
+}
+
 /**
  * Resolve every entry in `config.agents` into a Mastra `Agent`
  * instance. When `config.agents` is omitted the plugin registers a
@@ -394,31 +435,78 @@ function composeInstructions(
  */
 export async function buildAgents(opts: {
   config: MastraPluginConfig;
-  context: pluginUtils.PluginContextLike | undefined;
+  context: plugin.PluginContextLike | undefined;
   memoryBuilder?: MemoryBuilder;
-  log: logUtils.Logger;
+  log: log.Logger;
 }): Promise<BuiltAgents> {
   const { config, context, memoryBuilder, log } = opts;
   const definitions = resolveDefinitions(config);
   const ids = Object.keys(definitions);
   const defaultAgentId = config.defaultAgent ?? ids[0] ?? FALLBACK_AGENT_ID;
 
-  const plugins = buildPluginsMap(context);
-  const ambientTools = config.tools ?? {};
+  const plugins = buildPluginsMap(config, context);
+  // System-default ambient tools every agent gets out of the box:
+  // `render_data` for inline visualizations and `summarize` for
+  // offloading text condensing to the fast / small chat tier. The
+  // user can shadow either by including a same-named tool in their own
+  // `config.tools` or per-agent `tools`. Order in {@link resolveTools}
+  // is `system -> user-ambient -> per-agent`, last write wins.
+  const systemTools: MastraTools = {
+    render_data: buildRenderDataTool(config),
+    summarize: buildSummarizeTool(config),
+  };
+  const ambientTools = {
+    ...systemTools,
+    ...(config.tools ?? {}),
+  };
   const style = resolveStyleInstructions(config);
+  // Default-on protection against the model copying turn-scoped
+  // chartIds from prior assistant tool results into the new
+  // turn's `[chart:<id>]` markers. Opt out per-plugin via
+  // `config.stripStaleCharts: false`.
+  const outputProcessors = [new ResultProcessor()];
+  const inputProcessors = [
+    ...(config.stripStaleCharts === false ? [] : [stripStaleChartsProcessor]),
+  ];
   const agents: Record<string, Agent> = {};
+  const approvalGatedByAgent: Array<{ agentId: string; toolIds: string[] }> = [];
 
   for (const [id, def] of Object.entries(definitions)) {
     const tools = await resolveTools(def.tools, plugins, ambientTools);
+    const workspace = resolveAgentWorkspace(def.workspace);
+    const gated = approvalGatedToolIds(tools);
+    if (gated.length > 0) approvalGatedByAgent.push({ agentId: id, toolIds: gated });
     const memory = memoryBuilder?.forAgent(id, def);
     agents[id] = new Agent({
       id,
       name: def.name ?? id,
-      ...(def.description !== undefined ? { description: def.description } : {}),
+      // Always carry a non-empty description: it's surfaced to MCP
+      // clients (which reject a description-less agent) and in agent
+      // metadata. Fall back to a generated one when the definition
+      // omits it.
+      description: def.description?.trim() || defaultAgentDescription(def.name ?? id),
       instructions: composeInstructions(def.instructions, style),
       model: resolveModel(config, def.model),
+      defaultOptions: {
+        maxSteps: config.agentMaxSteps ?? DEFAULT_AGENT_MAX_STEPS,
+      },
       tools,
       ...(memory ? { memory } : {}),
+      ...(workspace ? { workspace } : {}),
+      inputProcessors,
+      outputProcessors,
+    });
+    // Surface the effective default model per agent so operators can
+    // see at a glance which endpoint each agent points at without
+    // having to fire a request and inspect a trace. The value is the
+    // *static* default; per-request overrides (header / query /
+    // body) and the workspace-catalogue fuzzy match still apply at
+    // call time.
+    log.info("agent registered", {
+      id,
+      name: def.name ?? id,
+      defaultModel: describeAgentDefaultModel(config, def),
+      tools: Object.keys(tools),
     });
   }
 
@@ -428,8 +516,88 @@ export async function buildAgents(opts: {
     );
   }
 
-  log.info("agents registered", { ids, defaultAgentId });
-  return { agents, defaultAgentId };
+  assertApprovalGatedToolsHaveStorage(approvalGatedByAgent, memoryBuilder);
+
+  log.info("agents ready", { ids, defaultAgentId });
+  return { agents, defaultAgentId, ambientTools };
+}
+
+/**
+ * Tool ids on `tools` that are approval-gated (`requireApproval: true`).
+ * Keys are used as a fallback when a tool omits an explicit `id`.
+ */
+export function approvalGatedToolIds(tools: MastraTools): string[] {
+  if (!tools || typeof tools !== "object") return [];
+  const ids: string[] = [];
+  for (const [key, tool] of Object.entries(tools)) {
+    if (!isApprovalGatedTool(tool)) continue;
+    ids.push(resolveToolId(tool, key));
+  }
+  return ids;
+}
+
+/** True when a Mastra / AI SDK tool pauses for human approval before execute. */
+function isApprovalGatedTool(tool: unknown): boolean {
+  if (!tool || typeof tool !== "object") return false;
+  return (tool as { requireApproval?: boolean }).requireApproval === true;
+}
+
+function resolveToolId(tool: unknown, fallbackKey: string): string {
+  if (tool && typeof tool === "object" && typeof (tool as { id?: string }).id === "string") {
+    return (tool as { id: string }).id;
+  }
+  return fallbackKey;
+}
+
+/**
+ * Approval-gated tools suspend the agent loop until a human approves.
+ * Mastra persists those suspended runs in Mastra-instance-level storage
+ * (`memoryBuilder.instanceStorage()`), not per-agent memory - so boot
+ * must fail fast when such a tool is registered without storage.
+ */
+function assertApprovalGatedToolsHaveStorage(
+  gated: Array<{ agentId: string; toolIds: string[] }>,
+  memoryBuilder: MemoryBuilder | undefined,
+): void {
+  if (gated.length === 0) return;
+  if (memoryBuilder?.instanceStorage()) return;
+
+  const detail = gated
+    .map(({ agentId, toolIds }) => `${agentId}: ${toolIds.join(", ")}`)
+    .join("; ");
+  throw new Error(
+    "mastra: approval-gated tools require plugin storage (PostgresStore) to persist suspended runs. " +
+      `Affected agents/tools: ${detail}. ` +
+      "Register lakebase() before mastra() so storage auto-enables, or pass storage: true explicitly.",
+  );
+}
+
+/**
+ * Best-effort description of the *static* default model an agent will
+ * resolve to at call time. Walks the same precedence ladder as
+ * {@link resolveModel} / {@link buildModel}:
+ *
+ *   1. Per-agent `def.model` (string sugar -> the literal id;
+ *      function / `DynamicArgument` -> `"<dynamic>"` because the
+ *      resolver decides at call time).
+ *   2. Plugin-level `config.defaultModel` (same rules).
+ *   3. `DATABRICKS_SERVING_ENDPOINT_NAME` env var.
+ *   4. First entry of `config.defaultModelFallbacks ?? fallback.FALLBACK_MODEL_IDS`.
+ *
+ * Used for the startup `agent registered` log so operators can see
+ * which endpoint each agent points at by default. Per-request
+ * overrides (`X-Mastra-Model` etc.) and the workspace-catalogue
+ * fuzzy match are still applied at runtime.
+ */
+function describeAgentDefaultModel(config: MastraPluginConfig, def: MastraAgentDefinition): string {
+  const effective = def.model ?? config.defaultModel;
+  if (typeof effective === "string") return effective;
+  if (effective !== undefined) return "<dynamic>";
+  return (
+    process.env.DATABRICKS_SERVING_ENDPOINT_NAME ??
+    config.defaultModelFallbacks?.[0] ??
+    fallback.FALLBACK_MODEL_IDS[0]!
+  );
 }
 
 /**
@@ -446,9 +614,7 @@ export async function buildAgents(opts: {
  * Omitted or empty inputs fall back to a single built-in analyst so
  * the bare `mastra()` call still mounts a working chat route.
  */
-function resolveDefinitions(
-  config: MastraPluginConfig,
-): Record<string, MastraAgentDefinition> {
+function resolveDefinitions(config: MastraPluginConfig): Record<string, MastraAgentDefinition> {
   const input = config.agents;
   if (!input) return fallbackDefinitions();
 
@@ -484,7 +650,7 @@ function resolveDefinitions(
 /** Derive a registry id from a definition's `name`, with a fallback. */
 function deriveAgentKey(def: MastraAgentDefinition, index?: number): string {
   if (def.name) {
-    const slug = stringUtils.toIdentifier(def.name);
+    const slug = string.toIdentifier(def.name);
     if (slug) return slug;
   }
   return index === undefined ? FALLBACK_AGENT_ID : `agent_${index}`;
@@ -548,6 +714,13 @@ async function resolveTools(
   return { ...ambientTools, ...resolved };
 }
 
+function resolveAgentWorkspace(
+  workspace: MastraAgentDefinition["workspace"],
+): Workspace | undefined {
+  if (!workspace) return undefined;
+  return typeof workspace === "function" ? workspace() : workspace;
+}
+
 /**
  * Build the {@link MastraPlugins} runtime proxy handed to
  * `tools(plugins)` callbacks.
@@ -569,14 +742,15 @@ async function resolveTools(
  * time instead of staring at a spinner for the full Genie round-trip.
  */
 function buildPluginsMap(
-  context: pluginUtils.PluginContextLike | undefined,
+  config: MastraPluginConfig,
+  context: plugin.PluginContextLike | undefined,
 ): MastraPlugins {
   const cache = new Map<string, MastraPluginToolkitProvider | null>();
   return new Proxy({} as MastraPlugins, {
     get(_target, propName) {
       if (typeof propName !== "string") return undefined;
       if (cache.has(propName)) return cache.get(propName) ?? undefined;
-      const provider = resolveProvider(context, propName);
+      const provider = resolveProvider(config, context, propName);
       cache.set(propName, provider);
       return provider ?? undefined;
     },
@@ -585,18 +759,35 @@ function buildPluginsMap(
 
 /**
  * Pick the right {@link MastraPluginToolkitProvider} for a sibling
- * plugin lookup. Returns the streaming-aware Genie adapter when the
- * caller asks for `genie`; falls back to the generic AppKit
- * `ToolProvider` adapter for every other plugin name.
+ * plugin lookup. Returns the Genie agent-backed adapter when
+ * the caller asks for `genie` AND at least one space is reachable
+ * via {@link resolveGenieSpaces} (the explicit
+ * `config.genieSpaces`, the registered AppKit `genie()` plugin's
+ * `spaces` config, or the `DATABRICKS_GENIE_SPACE_ID` env var).
+ * Falls back to the generic AppKit `ToolProvider` adapter for
+ * every other plugin name. `config` is threaded through so the
+ * Genie agent inherits the same model resolver / fallback
+ * ladder the calling agents use.
+ *
+ * The Genie agent talks to Genie directly via `@dbx-tools/genie`
+ * (`genieEventChat`) and the workspace
+ * `statementExecution.getStatement` API. AppKit's stock `genie`
+ * plugin is honored only for its resource manifest and `spaces`
+ * config so existing `app.yaml` configs and `genie({ spaces })`
+ * wiring keep working without change.
  */
 function resolveProvider(
-  context: pluginUtils.PluginContextLike | undefined,
+  config: MastraPluginConfig,
+  context: plugin.PluginContextLike | undefined,
   propName: string,
 ): MastraPluginToolkitProvider | null {
   if (propName === "genie") {
-    const geniePlugin = pluginUtils.instance(context, genie);
-    if (!geniePlugin) return null;
-    return buildGenieProvider(geniePlugin) as MastraPluginToolkitProvider;
+    const spaces = resolveGenieSpaces(config, context);
+    if (Object.keys(spaces).length === 0) return null;
+    return buildGenieToolkitProvider({
+      spaces,
+      config,
+    }) as MastraPluginToolkitProvider;
   }
   const plugin = context?.getPlugins().get(propName);
   return adaptPluginToolkit(plugin);
@@ -609,11 +800,7 @@ function resolveProvider(
  */
 interface AppKitToolkitProvider {
   toolkit?: (opts?: ToolkitOptions) => Record<string, AppKitToolkitEntry>;
-  executeAgentTool?: (
-    name: string,
-    args: unknown,
-    signal?: AbortSignal,
-  ) => Promise<unknown>;
+  executeAgentTool?: (name: string, args: unknown, signal?: AbortSignal) => Promise<unknown>;
 }
 
 /** Single entry returned by an AppKit plugin's `.toolkit(opts)` call. */
@@ -658,17 +845,13 @@ function adaptPluginToolkit(plugin: unknown): MastraPluginToolkitProvider | null
  * Schema parameters pass through unchanged - Mastra's `PublicSchema`
  * accepts `JSONSchema7` directly via `@mastra/schema-compat`.
  */
-function toolkitEntryToMastraTool(
-  entry: AppKitToolkitEntry,
-  plugin: AppKitToolkitProvider,
-): Tool {
+function toolkitEntryToMastraTool(entry: AppKitToolkitEntry, plugin: AppKitToolkitProvider): Tool {
   return createTool({
     id: `${entry.pluginName}__${entry.localName}`,
     description: entry.def.description,
     ...(entry.def.parameters ? { inputSchema: entry.def.parameters as never } : {}),
     execute: async (input: unknown, context: unknown) => {
-      const signal = (context as { abortSignal?: AbortSignal } | undefined)
-        ?.abortSignal;
+      const signal = (context as { abortSignal?: AbortSignal } | undefined)?.abortSignal;
       return plugin.executeAgentTool!(entry.localName, input, signal);
     },
   });