@dbx-tools/appkit-mastra 0.1.4 → 0.1.9

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (46) hide show
  1. package/README.md +304 -503
  2. package/index.ts +46 -36
  3. package/package.json +58 -47
  4. package/src/agents.ts +255 -72
  5. package/src/chart.ts +695 -0
  6. package/src/config.ts +282 -18
  7. package/src/filesystems.ts +1090 -0
  8. package/src/genie.ts +1025 -288
  9. package/src/history.ts +297 -0
  10. package/src/mcp.ts +105 -0
  11. package/src/memory.ts +129 -74
  12. package/src/mlflow.ts +149 -0
  13. package/src/model.ts +80 -408
  14. package/src/observability.ts +144 -0
  15. package/src/pagination.ts +34 -0
  16. package/src/plugin.ts +587 -52
  17. package/src/processors.ts +168 -0
  18. package/src/rest.ts +67 -0
  19. package/src/server.ts +292 -80
  20. package/src/serving-sanitize.ts +167 -0
  21. package/src/serving.ts +27 -224
  22. package/src/statement.ts +89 -0
  23. package/src/storage-schema.ts +41 -0
  24. package/src/summarize.ts +176 -0
  25. package/src/threads.ts +338 -0
  26. package/src/workspaces.ts +346 -0
  27. package/src/writer.ts +44 -0
  28. package/tsconfig.json +41 -0
  29. package/dist/index.d.ts +0 -18
  30. package/dist/index.js +0 -18
  31. package/dist/src/agents.d.ts +0 -306
  32. package/dist/src/agents.js +0 -379
  33. package/dist/src/config.d.ts +0 -170
  34. package/dist/src/config.js +0 -12
  35. package/dist/src/genie.d.ts +0 -109
  36. package/dist/src/genie.js +0 -271
  37. package/dist/src/memory.d.ts +0 -79
  38. package/dist/src/memory.js +0 -197
  39. package/dist/src/model.d.ts +0 -159
  40. package/dist/src/model.js +0 -423
  41. package/dist/src/plugin.d.ts +0 -120
  42. package/dist/src/plugin.js +0 -235
  43. package/dist/src/server.d.ts +0 -42
  44. package/dist/src/server.js +0 -109
  45. package/dist/src/serving.d.ts +0 -156
  46. package/dist/src/serving.js +0 -214
package/README.md CHANGED
@@ -1,36 +1,70 @@
1
- # @dbx-tools/appkit-mastra
2
-
3
- An AppKit plugin that hosts [Mastra](https://mastra.ai) agents inside a
4
- Databricks App with user-scoped workspace auth (OBO), optional
5
- Lakebase-backed memory, and an AI SDK chat route the React client can
6
- consume with `useChat()`.
7
-
8
- The plugin is designed so that wiring it up looks the same as the
9
- AppKit
10
- `[agents](https://developers.databricks.com/docs/appkit/v0/plugins/agents)`
11
- plugin - same `createAgent` / `tool` helpers, same `tools(plugins)`
12
- callback shape, same `ToolkitOptions`. Switching between the two for a
13
- given agent is a one-line import change.
14
-
15
- ## Quick start
16
-
17
- The pattern below is the direct counterpart of AppKit's `agents` plugin
18
- example - swap `agents` for `mastra` and the imports stay structurally
19
- identical:
1
+ # @dbx-tools/node-appkit-mastra
2
+
3
+ AppKit plugin and server-side toolkit for hosting Mastra agents inside a
4
+ Databricks App.
5
+
6
+ Import this package when an AppKit backend needs an agent service with
7
+ Databricks on-behalf-of auth, optional Lakebase-backed memory, Databricks Genie
8
+ tools, model selection, chart/data embeds, MLflow feedback, and MCP exposure.
9
+ The package mounts the standard Mastra agent stream under the AppKit server, so
10
+ clients can use Mastra-compatible chat transports instead of a custom protocol.
11
+
12
+ Key features:
13
+
14
+ - AppKit plugin lifecycle integration: routes, setup, shutdown, sibling plugin
15
+ access, and AppKit request context are handled inside `plugin.mastra()`.
16
+ - Agent composition: define one or more Mastra agents, give each one local tools,
17
+ AppKit plugin toolkits, workspace skills, model defaults, and approval-gated
18
+ tools.
19
+ - Databricks execution model: tool calls run with the active AppKit OBO client
20
+ where available, while storage and background work use service-principal
21
+ connections.
22
+ - Durable conversations: Lakebase-backed Mastra storage provides thread
23
+ history, message persistence, and optional vector memory.
24
+ - Rich data answers: Genie tools, statement fetches, chart preparation, and
25
+ embed markers let an agent answer with text plus delayed chart/table payloads.
26
+ - Operational surfaces: model-list routes, feedback routes, MCP exposure,
27
+ scoped API gating, tracing, and MLflow feedback are bundled with the plugin.
28
+
29
+ ## Why Not Just AppKit Agents?
30
+
31
+ Native AppKit includes a beta Agents plugin with markdown and TypeScript agent
32
+ definitions, AppKit tool-provider integration, streaming chat, thread
33
+ management, cancellation, and HITL approval. Use it when you want the AppKit
34
+ agent model and do not need a separate agent framework.
35
+
36
+ Use this package when you specifically want Mastra inside AppKit:
37
+
38
+ - Mastra's larger plugin/tool ecosystem, MCP support, memory/storage model,
39
+ workflow primitives, and `@mastra/client-js` stream shape.
40
+ - AppKit toolkits as Mastra tools, so Analytics, Files, Genie, and other AppKit
41
+ ToolProvider plugins stay available without rewriting them.
42
+ - Genie as an agent tool that emits typed progress events, result metadata, and
43
+ delayed chart/data markers into the same assistant turn.
44
+ - A paired React client in [`@dbx-tools/ui-mastra`](../../ui/mastra) with model
45
+ picking, thread sidebar, approvals, feedback, exports, and inline embeds.
46
+ - Per-request model override and fuzzy endpoint resolution through
47
+ [`@dbx-tools/node-model`](../model), instead of binding every agent to a fixed
48
+ endpoint name.
49
+
50
+ ## Quick Start
20
51
 
21
52
  ```ts
22
- import { analytics, createApp, files, lakebase, server } from "@databricks/appkit";
23
- import { createAgent, mastra, tool } from "@dbx-tools/appkit-mastra";
53
+ import { analytics, createApp, lakebase, server } from "@databricks/appkit";
54
+ import { agents, genie, plugin } from "@dbx-tools/node-appkit-mastra";
24
55
  import { z } from "zod";
25
56
 
26
- const support = createAgent({
27
- instructions: "You help customers with data and files.",
57
+ const analyst = agents.createAgent({
58
+ name: "analyst",
59
+ instructions: ["You answer questions about workspace data.", genie.GENIE_INSTRUCTIONS].join(
60
+ "\n\n",
61
+ ),
28
62
  tools(plugins) {
29
63
  return {
30
- ...plugins.analytics.toolkit(), // every analytics tool
31
- ...plugins.files.toolkit({ only: ["uploads.read"] }), // filtered subset
32
- get_weather: tool({
33
- description: "Weather",
64
+ ...plugins.analytics.toolkit(),
65
+ ...plugins.genie?.toolkit(),
66
+ get_weather: agents.tool({
67
+ description: "Get a simple weather report.",
34
68
  schema: z.object({ city: z.string() }),
35
69
  execute: async ({ city }) => `Sunny in ${city}`,
36
70
  }),
@@ -42,552 +76,319 @@ await createApp({
42
76
  plugins: [
43
77
  server(),
44
78
  analytics(),
45
- files(),
46
- // Drop `lakebase()` in and `mastra` auto-enables per-agent thread
47
- // storage (`schemaName: "mastra_<agentId>"`) plus a shared
48
- // semantic-recall vector index. Skip it for a stateless agent.
49
79
  lakebase(),
50
- mastra({ agents: support }),
80
+ plugin.mastra({
81
+ agents: { analyst },
82
+ defaultAgent: "analyst",
83
+ genie: { spaces: { sales: "01ef..." } },
84
+ }),
51
85
  ],
52
86
  });
53
87
  ```
54
88
 
55
- `createAgent` is a no-op identity helper that anchors type inference.
56
- `tool` is the AppKit-shaped factory (`{ description, schema, execute }`)
57
- that auto-adapts to Mastra's `createTool` under the hood.
58
-
59
- Memory + storage cascades:
60
-
61
- - **No `lakebase()` registered** ▸ agent is fully stateless. No threads,
62
- no recall. Same as `mastra()` alone.
63
- - `**lakebase()` registered, no `storage` / `memory` config** ▸ both
64
- auto-turn on. Each agent gets its own `PostgresStore` schema; every
65
- agent shares one `PgVector` recall index.
66
- - **Per-agent opt-out** ▸ `createAgent({ ..., memory: false, storage: false })`
67
- for routing / one-shot agents that don't need history.
68
- - **Per-agent override** ▸ pass a `PgVectorConfig` / `PostgresStoreConfig` object on the agent for a private index or a shared external schema.
69
-
70
- See [Memory + storage](#memory--storage) for the full cascade and worked
71
- examples.
72
-
73
- On the React side, never hardcode `/api/mastra/...`. Pull the published
74
- paths from `usePluginClientConfig` and use the `chatUrl` helper. Import
75
- them from the dependency-free `@dbx-tools/appkit-mastra-shared` package
76
- so your browser bundle doesn't pull in `pg`, `fastembed`, or Mastra:
77
-
78
- ```tsx
79
- import { usePluginClientConfig } from "@databricks/appkit-ui/react";
80
- import { chatUrl, type MastraClientConfig } from "@dbx-tools/appkit-mastra-shared";
81
- import { useChat } from "@ai-sdk/react";
82
- import { DefaultChatTransport } from "ai";
83
- import { useMemo } from "react";
84
-
85
- function Chat() {
86
- const config = usePluginClientConfig<MastraClientConfig>("mastra");
87
- const transport = useMemo(
88
- () => new DefaultChatTransport({ api: chatUrl(config) }),
89
- [config],
90
- );
91
- const { messages, sendMessage } = useChat({ transport });
92
- // ...
93
- }
94
- ```
95
-
96
- See [Client wiring](#client-wiring) for the full `MastraClientConfig`
97
- shape and per-agent selection.
98
-
99
- ## Sensible defaults
100
-
101
- The plugin is opinionated about what "no config" should mean. Everything
102
- below can be overridden, but the bare path works:
103
-
104
-
105
- | Scenario | What you get |
106
- | --------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
107
- | `mastra()` | One built-in `default` analyst agent, model from `/serving-endpoints`; memory + storage auto-on if the `lakebase` plugin is registered |
108
- | `mastra({ agents: def })` | Single-agent shorthand - `def` is registered and marked as default |
109
- | `mastra({ agents: [def1, def2] })` | Array shorthand - keys come from each `def.name` (or `agent_<i>`); first one is default |
110
- | `mastra({ agents: { x: def, y: def }})` | Record - keys are the registered ids; first key is the default |
111
- | No `defaultAgent` set | First registered agent wins |
112
- | No `model` on an agent | Falls back to `config.defaultModel`, then `DATABRICKS_SERVING_ENDPOINT_NAME`, then walks `defaultModelFallbacks` (Thinking → Balanced → Fast tiers, Claude ↔ GPT ↔ Gemini interleaved within each, then open-weights) and picks the first endpoint actually present in the workspace |
113
- | No `name` on a definition | Uses the registry key as `Agent.name` |
114
- | No `tools` on an agent | Inherits plugin-level `config.tools` ambient set (if any) |
115
- | No `storage` / `memory` and `lakebase()` registered | Both auto-default to `true`. Pass `false` (or a custom config object) on the plugin or an agent to opt out / override. |
116
- | `storage` / `memory` on an agent | Cascades from plugin: storage namespaces **per-agent** (`schemaName: "mastra_<agentId>"`), vector recall is **shared** across agents on one `PgVector` singleton. |
89
+ Benefits of importing the package:
117
90
 
91
+ - `plugin.mastra()` registers a full AppKit plugin named `mastra`.
92
+ - `agents.createAgent()` keeps agent definitions typed and applies the default
93
+ Databricks workspace/skill mounts.
94
+ - `agents.tool()` lets the same AppKit-shaped tool body work in this Mastra
95
+ plugin.
96
+ - `genie.GENIE_INSTRUCTIONS` and `plugins.genie.toolkit()` give agents a
97
+ Databricks Genie workflow without embedding a second agent.
98
+ - Lakebase registration automatically enables durable thread storage and vector
99
+ memory unless you opt out.
118
100
 
119
- Every field on `MastraAgentDefinition`:
101
+ ## Agent Registration
120
102
 
121
-
122
- | Field | Description |
123
- | -------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
124
- | `name` | Display name. Defaults to the registry key. |
125
- | `description` | Long-form description, surfaced as `Agent.description`. |
126
- | `instructions` | System prompt body. Required. |
127
- | `model` | Per-agent model override. String = `modelId` sugar; otherwise a Mastra `DynamicArgument`. |
128
- | `tools` | Plain record OR `(plugins) => tools` callback (see below). |
129
- | `memory` | `false`, `true`, or a `PgVectorConfig`. Cascades from `config.memory`. **Default: shared singleton `PgVector` across every agent** - object override switches to a dedicated index. |
130
- | `storage` | `false`, `true`, or a `PostgresStoreConfig`. Cascades from `config.storage`. **Default: per-agent namespace** via `schemaName: "mastra_<agentId>"` so threads stay isolated. |
131
-
132
-
133
- Plugin-level fields:
134
-
135
-
136
- | Field | Description |
137
- | ----------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
138
- | `agents` | The registry. Omit for a built-in `default` analyst. |
139
- | `defaultAgent` | Id that `chatRoute` binds to when no `:agentId` is supplied. Defaults to first registered. |
140
- | `defaultModel` | Fallback for any agent that omits `model`. Same shape (string sugar or `DynamicArgument`). |
141
- | `defaultModelFallbacks` | Priority-ordered list walked when no `model` / env / override is set. First entry whose endpoint exists in the workspace wins. Default chains the three `ModelTier`s (Thinking → Balanced → Fast); within each tier providers are interleaved Claude ↔ GPT ↔ Gemini with open-weights appended. Compose your own with `modelsForTier(ModelTier.Fast)` or read straight from `MODEL_CATALOG`. |
142
- | `tools` | Ambient tools merged into every agent (per-agent tools win on collisions). |
143
- | `storage` | `undefined` (default) auto-enables when `lakebase()` is registered; `true` does the same explicitly; `false` opts out; an object opens a dedicated `PostgresStore`. Per-agent default is `schemaName: "mastra_<agentId>"`. |
144
- | `memory` | `undefined` (default) auto-enables when `lakebase()` is registered; `true` does the same explicitly; `false` opts out; an object opens a dedicated `PgVector`. Default behavior: one shared `PgVector` singleton across every agent. |
145
- | `modelFuzzyMatch` | `false` to disable fuzzy snapping of model ids against the workspace's Model Serving catalogue. Defaults to `true`. |
146
- | `modelFuzzyThreshold` | Fuse.js score threshold (`0` exact, `1` anything). Defaults to `0.4`. |
147
- | `modelCacheTtlMs` | TTL for the cached endpoint list, per workspace host. Defaults to 5 minutes. Concurrent callers share one in-flight fetch. |
148
- | `modelOverride` | `false` to disable per-request `X-Mastra-Model` / `?model=` / body overrides. Defaults to `true`. |
149
- | `styleInstructions` | Style guardrails appended to every agent's `instructions` to curb LLM-isms (em dashes, emojis, sycophantic openers, throwaway closers). `undefined` (default) uses the built-in `DEFAULT_STYLE_INSTRUCTIONS`; a string replaces it; `false` disables. |
150
-
151
-
152
- ## The `tools(plugins)` callback
153
-
154
- Each registered agent can supply either a static `tools: { ... }` record
155
- or a `tools(plugins)` callback. The returned record accepts **any**
156
- tool shape Mastra understands:
157
-
158
- - Mastra tools built with `createTool` or `new Tool(...)`
159
- - AppKit-shaped tools built with the `tool()` wrapper (below)
160
- - Vercel AI SDK tools (`tool({...})` from `ai`)
161
- - Provider-defined tools (e.g. `openai.tools.webSearch(...)`)
162
- - Toolkits returned from `plugins.<name>.toolkit(...)`
103
+ `plugin.mastra({ agents })` accepts a single definition, an array, or a record.
104
+ Records are best when clients need stable agent ids:
163
105
 
164
106
  ```ts
165
- tools(plugins) {
166
- return {
167
- // Sibling plugin toolkits.
168
- ...plugins.analytics.toolkit(), // every analytics tool
169
- ...plugins.files.toolkit({ only: ["uploads.read"] }), // filtered subset
170
-
171
- // AppKit-shaped inline tool.
172
- get_weather: tool({
173
- description: "Weather",
174
- schema: z.object({ city: z.string() }),
175
- execute: async ({ city }) => `Sunny in ${city}`,
176
- }),
177
-
178
- // Existing Mastra tool dropped in unchanged.
179
- save_doc: existingMastraTool,
180
- };
181
- }
107
+ plugin.mastra({
108
+ agents: {
109
+ support: agents.createAgent({ instructions: "Answer support questions." }),
110
+ analyst: agents.createAgent({ instructions: "Analyze workspace data." }),
111
+ },
112
+ defaultAgent: "support",
113
+ });
182
114
  ```
183
115
 
184
- `plugins` is a typed `Record<string, { toolkit(opts?): tools }>` matching
185
- AppKit's `Plugins` type. It's backed by a runtime Proxy that
186
- **auto-discovers any registered AppKit `ToolProvider` plugin** -
187
- `analytics`, `files`, `lakebase`, `genie`, plus any third-party plugin
188
- that implements the standard `getAgentTools()` + `executeAgentTool()` +
189
- `toolkit()` interface. Tool calls dispatch through the plugin's
190
- `executeAgentTool`, so OBO auth (`asUser`) and telemetry spans stay
191
- intact.
192
-
193
- `plugins.genie` is special-cased: it swaps the generic AppKit toolkit
194
- (which only emits a single final result chunk per call) for a
195
- streaming-aware tools record built on top of the plugin's
196
- `exports().sendMessage` AsyncGenerator. Each Genie wire event
197
- (`FETCHING_METADATA`, `ASKING_AI`, `EXECUTING_QUERY`, attached SQL,
198
- `query_result` row counts, errors) is normalised into a `GenieProgress`
199
- payload and pushed mid-flight through Mastra's `ctx.writer`, surfacing
200
- as `tool-output` chunks the React client can render as inline status
201
- pills, SQL blocks, and row-count badges while the LLM is still waiting
202
- on the final `tool-result`. Tool ids are stable: `genie` for the
203
- default alias, `genie-<alias>` for additional aliases, and one shared
204
- `genie_get_conversation`. The LLM still receives a single clean final
205
- payload so the streaming UI doesn't leak into the assistant message.
206
-
207
- Plugins that aren't registered (or don't implement the toolkit
208
- interface) resolve to `undefined` at runtime, so guard with `?.` /
209
- `?? {}` when a backing plugin is optional in some environments:
116
+ When no agents are supplied, the plugin registers a built-in `default` analyst so
117
+ the route surface still works for smoke tests. Each agent is streamed through the
118
+ Mastra agent API mounted below the plugin path, typically `/api/mastra`.
119
+
120
+ Use `agents.createTool` when you need Mastra-native fields such as
121
+ `outputSchema`, `suspendSchema`, `requireApproval`, or MCP metadata. Use
122
+ `agents.tool` for the smaller AppKit-compatible shape:
210
123
 
211
124
  ```ts
212
- tools(plugins) {
213
- return {
214
- ...(plugins.analytics?.toolkit() ?? {}),
215
- ...(plugins.genie?.toolkit({ prefix: "g_" }) ?? {}),
216
- };
217
- }
125
+ const approveRefund = agents.createTool({
126
+ id: "approve_refund",
127
+ description: "Approve a refund request.",
128
+ inputSchema: z.object({ orderId: z.string(), amount: z.number() }),
129
+ requireApproval: true,
130
+ execute: async ({ context }) => approve(context.orderId, context.amount),
131
+ });
218
132
  ```
219
133
 
220
- `plugins.<name>.toolkit(opts)` accepts the same `ToolkitOptions` shape
221
- AppKit's own toolkits expose (passed through verbatim):
134
+ ## AppKit Toolkits
222
135
 
223
- - `prefix?: string` - prepended to every key (AppKit default: `${pluginName}.`)
224
- - `only?: string[]` / `except?: string[]` - allow/deny list against the
225
- local tool name
226
- - `rename?: Record<string, string>` - remap individual keys
227
-
228
- ### `tool()` vs `createTool()`
229
-
230
- The `tool()` factory mirrors `@databricks/appkit/beta`'s shape so
231
- sharing tool code between the AppKit `agents` plugin and this one is a
232
- single-line import change:
136
+ The `tools(plugins)` callback receives a dynamic index of registered AppKit
137
+ tool-provider plugins. Each entry exposes `.toolkit(opts)` with AppKit-compatible
138
+ `prefix`, `only`, `except`, and `rename` options.
233
139
 
234
140
  ```ts
235
- import { tool, createTool } from "@dbx-tools/appkit-mastra";
236
-
237
- // AppKit-shaped (description / schema / flat-arg execute).
238
- const weather = tool({
239
- description: "Weather",
240
- schema: z.object({ city: z.string() }),
241
- execute: async ({ city }) => `Sunny in ${city}`,
242
- });
243
-
244
- // Full Mastra `createTool` (id required, inputSchema, advanced fields).
245
- const saveDoc = createTool({
246
- id: "save-doc",
247
- description: "Persist a document",
248
- inputSchema: z.object({ key: z.string(), value: z.any() }),
249
- outputSchema: z.object({ saved: z.boolean() }),
250
- requireApproval: true,
251
- execute: async (input, ctx) => {
252
- await ctx.mastra?.getStorage()?.set(input.key, input.value);
253
- return { saved: true };
141
+ const agent = agents.createAgent({
142
+ instructions: "Use the narrowest tool that answers the question.",
143
+ tools(plugins) {
144
+ return {
145
+ ...plugins.analytics.toolkit({ only: ["query"] }),
146
+ ...plugins.files?.toolkit({ prefix: "files.", except: ["delete"] }),
147
+ };
254
148
  },
255
149
  });
256
150
  ```
257
151
 
258
- When `tool()`'s `id` is omitted it's auto-derived from a slugified
259
- description plus a 6-char SHA-1 prefix - stable across runs so traces
260
- stay readable. Pass an explicit `id` when you want to pin one.
261
-
262
- Reach for `createTool` when you need Mastra-only fields (`outputSchema`,
263
- `suspendSchema`, `requireApproval`, `mcp`, etc.).
264
-
265
- ## Model resolution
266
-
267
- Each agent call resolves a `MastraModelConfig` lazily so concurrent
268
- requests get distinct user identities. There are two paths through
269
- the resolver depending on whether the caller asked for a specific
270
- model.
271
-
272
- ### Explicit ask (override / agent / plugin / env)
273
-
274
- When any of these is set the resolver fuzzy-matches that single id
275
- against the live `/serving-endpoints` list, in priority order:
276
-
277
- 1. Per-request override (`X-Mastra-Model` header, `?model=` query,
278
- or `model` / `modelId` body field; see below)
279
- 2. Per-agent `def.model` (string sugar or `DynamicArgument`)
280
- 3. Plugin-level `config.defaultModel`
281
- 4. `DATABRICKS_SERVING_ENDPOINT_NAME` env var
282
-
283
- The matcher is `fuse.js` extended search with tokens split on
284
- non-word characters and AND-joined. Exact matches win immediately;
285
- loose tokens like `"claude sonnet"` snap to
286
- `databricks-claude-sonnet-4-6`, `"llama 70b"` to
287
- `databricks-meta-llama-3-3-70b-instruct`, `"DBRX"` to
288
- `databricks-dbrx-instruct`, and so on. If no candidate scores below
289
- `modelFuzzyThreshold` (default `0.4`) the input is returned verbatim
290
- and Databricks surfaces the canonical 404.
152
+ Tool calls dispatch back through the owning AppKit plugin, preserving OBO auth
153
+ and AppKit telemetry behavior. Optional plugins should be guarded with `?.` when
154
+ you spread their tools.
291
155
 
292
- ### No explicit ask (tier-aware fallback list)
156
+ ## Memory And Storage
293
157
 
294
- When nothing is set, the resolver walks an opinionated
295
- priority-ordered list and returns **the first id that is actually
296
- present in the workspace's endpoint listing**. This is how a workspace
297
- without Claude Opus still gets a sensible default automatically -
298
- the resolver skips ahead to whichever Sonnet / GPT-5 / Gemini / Llama
299
- variant is wired up.
158
+ The `memory` and `storage` config fields can be `false`, `true`, or a concrete
159
+ Mastra Postgres/PgVector config.
300
160
 
301
- The catalogue is grouped two ways:
302
-
303
- - By **capability tier** via the `ModelTier` enum:
304
- `ModelTier.Thinking` (deepest reasoning), `ModelTier.Balanced`
305
- (cost/latency sweet spot), `ModelTier.Fast` (cheap & quick for
306
- classification / routing / simple summarisation).
307
- - By **provider** within each tier: `claude`, `gpt`, `gemini`,
308
- `openSource`.
309
-
310
- Both views live on `MODEL_CATALOG[tier][provider]`. The walked
311
- `FALLBACK_MODEL_IDS` chains the three tiers in descending power
312
- (Thinking → Balanced → Fast); within each tier providers are
313
- round-robin-zipped (Claude ↔ GPT ↔ Gemini) before the open-weights
314
- tail is appended as the universal floor.
161
+ ```ts
162
+ plugin.mastra({
163
+ agents: analyst,
164
+ storage: true,
165
+ memory: { id: "analytics_memory", tableName: "agent_memory" },
166
+ });
167
+ ```
315
168
 
169
+ With `lakebase()` registered, both default to enabled:
316
170
 
317
- | Tier (most powerful first) | Claude | GPT | Gemini | Open weights |
318
- | -------------------------- | -------------------------------- | ------------------------------------- | ------------------------------- | ---------------------------------------------- |
319
- | `ModelTier.Thinking` | Opus 4.8 4.7 4.6 4.5 4.1 | 5.5 Pro | 3.1 Pro → 3 Pro → 2.5 Pro | Llama 4 Maverick, GPT-OSS 120B, Llama 3.1 405B |
320
- | `ModelTier.Balanced` | Sonnet 4.6 4.5 4 | 5.5 → 5.4 → 5.2 → 5.1 → 5 | 3.5 Flash → 3 Flash → 2.5 Flash | Llama 3.3 70B, Qwen3-Next 80B, Qwen35 122B |
321
- | `ModelTier.Fast` | Haiku 4.5 | 5.4 mini → 5.4 nano → 5 mini → 5 nano | 3.1 Flash Lite | GPT-OSS 20B, Gemma 3 12B, Llama 3.1 8B |
171
+ - storage uses a per-agent schema for durable threads and messages;
172
+ - memory uses a shared vector index for semantic recall;
173
+ - the service-principal pool is created outside any request so OBO user
174
+ identities are not captured in background storage work.
322
175
 
176
+ Without `lakebase()`, agents are stateless unless you provide explicit storage
177
+ and memory configs.
323
178
 
324
- #### Pick a tier-appropriate model for one agent
179
+ ## Workspace Skills
325
180
 
326
- Use `modelForTier(tier)` to grab the top of a tier as a string; the
327
- agent-step resolver fuzzy-matches it against the live catalogue at
328
- call time so it still works when the literal top pick isn't deployed.
181
+ Every `agents.createAgent()` gets a default Mastra `Workspace` from
182
+ `workspaces.createWorkspace()`. It mounts Databricks Workspace files through the
183
+ current OBO user's `WorkspaceClient`, so Mastra can discover Assistant-style
184
+ `SKILL.md` files at request time.
329
185
 
330
186
  ```ts
331
- import { createAgent, ModelTier, modelForTier } from "@dbx-tools/appkit-mastra";
332
-
333
- const classifier = createAgent({
334
- instructions: "Classify this email into one of: billing, support, spam.",
335
- model: modelForTier(ModelTier.Fast),
336
- });
337
-
338
- const planner = createAgent({
339
- instructions: "Plan a multi-step data migration.",
340
- model: modelForTier(ModelTier.Thinking),
187
+ const agent = agents.createAgent({
188
+ instructions: "Use mounted workspace skills when relevant.",
189
+ workspace: workspaces.createWorkspace({
190
+ assistantSkills: true,
191
+ mounts: [
192
+ async () => ({
193
+ mounts: { "/reference": myFilesystem },
194
+ skillPaths: ["/reference/skills"],
195
+ }),
196
+ ],
197
+ }),
341
198
  });
342
199
  ```
343
200
 
344
- #### Bias the plugin-level fallback toward a tier
201
+ Production workspace mounts require a forwarded token with `workspace`,
202
+ `workspace.workspace`, or `all-apis` scope. Development mode skips that gate for
203
+ local iteration.
345
204
 
346
- `modelsForTier(tier)` returns the priority-ordered list for one tier;
347
- pass it to `defaultModelFallbacks` to scope the auto-resolver:
348
-
349
- ```ts
350
- import { mastra, ModelTier, modelsForTier } from "@dbx-tools/appkit-mastra";
205
+ ## Genie Tools
351
206
 
352
- mastra({
353
- // All agents that omit `model` will land on a Fast-tier endpoint.
354
- defaultModelFallbacks: modelsForTier(ModelTier.Fast),
355
- });
356
- ```
207
+ `genie.buildGenieTools()` and `plugins.genie.toolkit()` expose tools for:
357
208
 
358
- #### Pin a custom approved subset
209
+ - asking a configured Genie space;
210
+ - reading space descriptions and serialized space metadata;
211
+ - fetching statement rows by `statement_id`;
212
+ - preparing charts from Genie result sets.
359
213
 
360
- Mix in your own endpoint names (internal fine-tunes, regulated
361
- allowlists, etc) in front of the catalogue:
214
+ The central agent drives those tools directly. Genie events stream through the
215
+ Mastra writer using the shared contract from
216
+ [`@dbx-tools/shared-mastra`](../../shared/mastra), so clients can show thinking,
217
+ SQL, row counts, summaries, chart markers, and data markers as the turn runs.
362
218
 
363
219
  ```ts
364
- mastra({
365
- defaultModelFallbacks: [
366
- "my-org-finetune-v2", // try internal endpoint first
367
- "databricks-claude-sonnet-4-6", // approved fallback
368
- ],
220
+ const agent = agents.createAgent({
221
+ instructions: `${baseInstructions}\n\n${genie.GENIE_INSTRUCTIONS}`,
222
+ tools(plugins) {
223
+ return { ...plugins.genie?.toolkit({ prefix: "" }) };
224
+ },
369
225
  });
370
226
  ```
371
227
 
372
- If the workspace has none of the listed ids, the top fallback is
373
- returned and Databricks surfaces the canonical error.
374
-
375
- The endpoint list is cached per workspace host through AppKit's
376
- built-in `CacheManager` (`CacheManager.getInstanceSync().getOrExecute`),
377
- which is the TypeScript counterpart of Python's `cachetools.TTLCache`
378
- plus `cachetools-async` rolled into one: per-entry TTL (default 5
379
- minutes via `modelCacheTtlMs`), bounded size, in-flight request
380
- coalescing (the manager's internal `inFlightRequests` map shares one
381
- fetch across every concurrent caller), telemetry spans, and optional
382
- Lakebase persistence when the `lakebase` plugin is wired up. No extra
383
- dependency lives in this package; the catalogue piggybacks on whatever
384
- storage backend AppKit picked at boot.
385
-
386
- String values (`"databricks-claude-sonnet-4-6"`) are `modelId` sugar
387
- layered on top of the auto-resolver - workspace URL, provider, and OBO
388
- auth stay default. Pass a `DynamicArgument<MastraModelConfig>` on
389
- `def.model` / `config.defaultModel` when you need full control over
390
- auth, provider, or URL; that path bypasses the fuzzy matcher and
391
- per-request override.
392
-
393
- ### `GET /api/mastra/models`
394
-
395
- The plugin exposes the cached endpoint catalogue at `/models` (mounted
396
- under the plugin prefix, default `/api/mastra`) so clients can populate
397
- model pickers and validate `?model=` choices without a separate
398
- Databricks SDK round-trip:
399
-
400
- ```bash
401
- curl -s http://localhost:8000/api/mastra/models | jq
402
- # {
403
- # "endpoints": [
404
- # { "name": "databricks-claude-sonnet-4-6", "task": "llm/v1/chat", "state": "READY", ... },
405
- # { "name": "databricks-meta-llama-3-3-70b-instruct", ... },
406
- # ...
407
- # ]
408
- # }
409
- ```
228
+ ## Charts And Data Embeds
410
229
 
411
- Same payload from a sibling plugin or script (no HTTP round-trip):
230
+ `chart.prepareChart()` mints a chart id immediately, caches an in-progress
231
+ record, resolves the data in the background, and stores a terminal chart or
232
+ error. `chart.fetchChart()` long-polls that cache for route handlers and custom
233
+ clients.
412
234
 
413
235
  ```ts
414
- import { pluginUtils } from "@dbx-tools/appkit-shared";
415
- import { mastra } from "@dbx-tools/appkit-mastra";
236
+ const { chartId } = await chart.prepareChart({
237
+ config,
238
+ request: {
239
+ title: "Revenue by region",
240
+ chartType: "bar",
241
+ instructions: "Compare total revenue by region.",
242
+ data: rows,
243
+ },
244
+ resolveData: async () => rows,
245
+ });
416
246
 
417
- const m = pluginUtils.require(this.context, mastra).exports();
418
- const endpoints = await m.asUser(req).listModels(); // user-scoped
419
- m.clearModelsCache(); // force the next call to re-fetch
247
+ const resolved = await chart.fetchChart(chartId);
420
248
  ```
421
249
 
422
- ### Per-request model override
423
-
424
- Any in-flight request can pick a different backing endpoint without
425
- redeploying. Sources, checked in priority order:
426
-
427
-
428
- | Source | Example |
429
- | ------------------------- | ------------------------------------------------ |
430
- | `X-Mastra-Model` header | `curl -H 'X-Mastra-Model: claude-haiku' ...` |
431
- | `?model=` query parameter | `POST /api/mastra/route/chat?model=llama-70b` |
432
- | Body `model` or `modelId` | `{ "messages": [...], "model": "claude-haiku" }` |
250
+ Agents can return `[chart:<id>]` and `[data:<statement_id>]` markers in prose.
251
+ The embed route resolves them later, which avoids forcing the language model to
252
+ inline large tables or wait for chart planning before continuing its answer.
433
253
 
254
+ ## Model Selection
434
255
 
435
- The override flows through the same fuzzy matcher as static ids, so
436
- `X-Mastra-Model: claude sonnet` still snaps to
437
- `databricks-claude-sonnet-4-6`. Set `modelOverride: false` on the
438
- plugin config to disable the override path entirely (e.g. for a
439
- multi-tenant deployment where untrusted clients shouldn't pick the
440
- endpoint).
256
+ `model.buildModel()` adapts the generic resolver from
257
+ [`@dbx-tools/node-model`](../model) to Mastra. It resolves the model per request,
258
+ so OBO identity and request-specific overrides stay isolated.
441
259
 
442
- ## Memory + storage
260
+ Model priority is:
443
261
 
444
- Memory and storage are split into two independent knobs and both auto-on
445
- the moment the `lakebase` plugin is registered. Bare `mastra()` next to
446
- `lakebase()` already gets you per-agent threads + shared semantic recall;
447
- zero extra config required.
262
+ 1. request override (`X-Mastra-Model`, `?model=`, body `model` / `modelId`);
263
+ 2. per-agent `model`;
264
+ 3. plugin `defaultModel`;
265
+ 4. `DATABRICKS_SERVING_ENDPOINT_NAME`;
266
+ 5. workspace catalogue ranking and static fallback floor.
448
267
 
268
+ ```ts
269
+ plugin.mastra({
270
+ agents: analyst,
271
+ defaultModel: "claude sonnet",
272
+ modelFuzzyMatch: true,
273
+ modelOverride: true,
274
+ });
275
+ ```
449
276
 
450
- | Knob | Default when `lakebase()` is registered | What it backs |
451
- | --------- | ----------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------- |
452
- | `storage` | **Per-agent** `PostgresStore` namespaced by `schemaName: "mastra_<agentId>"` so threads + messages stay isolated. | Mastra threads, messages, working memory. |
453
- | `memory` | **Shared singleton** `PgVector` across every agent (cross-agent semantic recall on one index). | RAG-style recall over past messages via FastEmbed vectors. |
277
+ Use `serving.extractModelOverride()` and `serving.resolveServingConfig()` when
278
+ building custom routes that should behave like the plugin's `/models` and stream
279
+ routes.
454
280
 
281
+ ## Threads, History, And Suggestions
455
282
 
456
- Override either at the plugin level, the agent level, or both. The agent
457
- value wins when set; otherwise the plugin value cascades.
283
+ When storage is enabled, the plugin provides route helpers and in-process
284
+ functions for conversation management:
458
285
 
459
- ```ts
460
- mastra({
461
- // Plugin defaults. Either field becomes the cascading baseline.
462
- // Omit entirely to inherit "auto-on when lakebase is present".
463
- storage: true, // (default behavior when lakebase is registered)
464
- memory: true, // (default behavior when lakebase is registered)
286
+ - `history.loadHistory()` and `history.clearHistory()` read or clear one thread;
287
+ - `threads.listThreads()`, `threads.renameThread()`, and
288
+ `threads.deleteThread()` operate on the caller's scoped conversations;
289
+ - `genie.collectSpaceSuggestions()` reads starter questions from the configured
290
+ Genie space.
465
291
 
466
- agents: {
467
- analyst: createAgent({
468
- instructions: "...",
469
- // No overrides: inherits the auto-on defaults above.
470
- // - threads stored under schema "mastra_analyst"
471
- // - recalls from the shared vector index
472
- }),
292
+ The plugin resolves the active thread from `x-mastra-thread-id`, `?threadId=`,
293
+ or a per-session fallback cookie. That keeps streaming, history, and clear
294
+ operations aligned around the same conversation id.
473
295
 
474
- router: createAgent({
475
- instructions: "Stateless routing agent.",
476
- // Opt out of both for a fully stateless agent.
477
- storage: false,
478
- memory: false,
479
- }),
296
+ ## Feedback And Observability
480
297
 
481
- legal: createAgent({
482
- instructions: "Compliance-bounded assistant.",
483
- // Private vector index so legal's recall doesn't bleed into
484
- // analyst's. Threads still get their own per-agent schema.
485
- memory: { connectionString: process.env.LEGAL_PG_URL!, /* ... */ },
486
- }),
298
+ `observability.buildObservability()` wires Mastra tracing when OTLP export is
299
+ configured. `mlflow.resolveFeedbackEnabled()` turns MLflow feedback on when both
300
+ trace export and an MLflow experiment are configured, unless the plugin config
301
+ forces a value.
487
302
 
488
- archive: createAgent({
489
- instructions: "Read-only archive viewer.",
490
- // Pin to a specific schema (e.g. shared with another service).
491
- storage: {
492
- schemaName: "shared_history",
493
- pool: archivePool,
494
- },
495
- }),
496
- },
303
+ ```ts
304
+ plugin.mastra({
305
+ agents: analyst,
306
+ feedback: true,
497
307
  });
498
308
  ```
499
309
 
500
- Notes:
310
+ `mlflow.logFeedback()` logs a human assessment against the active MLflow trace.
311
+ The response header name and request/response schemas live in
312
+ [`@dbx-tools/shared-mastra`](../../shared/mastra).
501
313
 
502
- - `PostgresStore` runs `CREATE SCHEMA IF NOT EXISTS` on `init()`, so
503
- per-agent schemas spring into existence the first time an agent saves
504
- a message. No bundle / migration step required.
505
- - Disabling `lakebase()` from your plugin list while leaving `storage` /
506
- `memory` truthy fails fast at setup with a clear "lakebase plugin not
507
- registered" error.
508
- - The `lakebase` plugin is declared as a **required** resource only when
509
- `storage` / `memory` is explicitly truthy at registration time. Auto-on
510
- defaults activate inside `setup:complete`, after lakebase is already
511
- proven to be present.
314
+ ## MCP Exposure
512
315
 
513
- ## Runtime exports
514
-
515
- Other plugins / route handlers can introspect the registry via the
516
- `exports()` surface, modeled on AppKit's:
316
+ `mcp.buildMcpServer()` exposes registered agents as MCP tools by default. The
317
+ AppKit plugin publishes clean aliases under its base path:
517
318
 
518
319
  ```ts
519
- import { pluginUtils } from "@dbx-tools/appkit-shared";
520
- import { mastra } from "@dbx-tools/appkit-mastra";
521
-
522
- const m = pluginUtils.require(this.context, mastra).exports();
523
- m.list(); // ["analyst", "helper"]
524
- m.get("analyst"); // Agent | null
525
- m.getDefault(); // Agent | null
526
- m.getMastra(); // underlying Mastra instance (advanced)
527
- m.listModels(); // Promise<ServingEndpointSummary[]> - cached + OBO when wrapped with asUser(req)
528
- m.clearModelsCache(); // force the next listModels() to re-fetch
529
- ```
530
-
531
- ## Client wiring
532
-
533
- `clientConfig()` publishes the mount paths, default agent id, and the
534
- full registry to `usePluginClientConfig("mastra")` so the React client
535
- never has to hardcode `/api/mastra` or rely on `DEFAULT_AGENT_ID`
536
- constants. A tiny URL helper (`chatUrl`) and the `MastraClientConfig`
537
- type ship from the standalone `@dbx-tools/appkit-mastra-shared`
538
- package; that package is pure (no `pg` / `fastembed` / Mastra
539
- dependencies) so it imports cleanly into Vite / Webpack / esbuild
540
- builds.
541
-
542
- ```tsx
543
- import { usePluginClientConfig } from "@databricks/appkit-ui/react";
544
- import { chatUrl, type MastraClientConfig } from "@dbx-tools/appkit-mastra-shared";
545
- import { useChat } from "@ai-sdk/react";
546
- import { DefaultChatTransport } from "ai";
547
- import { useMemo, useState } from "react";
548
-
549
- function Chat() {
550
- const config = usePluginClientConfig<MastraClientConfig>("mastra");
551
- const [selected, setSelected] = useState<string>();
552
- const api = chatUrl(config, selected); // defaults to config.defaultAgent
553
-
554
- const transport = useMemo(() => new DefaultChatTransport({ api }), [api]);
555
- const { messages, sendMessage } = useChat({ transport });
556
-
557
- return (
558
- <>
559
- <select onChange={(e) => setSelected(e.target.value)}>
560
- {config.agents.map((id) => (
561
- <option key={id} value={id}>
562
- {id}
563
- </option>
564
- ))}
565
- </select>
566
- {/* render messages, etc. */}
567
- </>
568
- );
569
- }
320
+ plugin.mastra({
321
+ agents: analyst,
322
+ mcp: {
323
+ serverId: "analytics",
324
+ name: "Analytics MCP",
325
+ tools: false,
326
+ },
327
+ });
570
328
  ```
571
329
 
572
- `MastraClientConfig` fields (all derived from the server-side plugin
573
- mount, so a custom `mastra({ name: "myMastra" })` rewrites every path):
574
-
575
-
576
- | Field | Example | Description |
577
- | ------------------ | ----------------------------------- | -------------------------------------------------------- |
578
- | `basePath` | `"/api/mastra"` | Plugin mount path. |
579
- | `chatPath` | `"/api/mastra/route/chat"` | Default-agent chat URL. Use `chatUrl(config)` to get it. |
580
- | `chatPathTemplate` | `"/api/mastra/route/chat/:agentId"` | OpenAPI-style template for tools / docs. |
581
- | `modelsPath` | `"/api/mastra/models"` | `GET` cached endpoint catalogue. |
582
- | `defaultAgent` | `"analyst"` | Agent id `chatRoute` binds to when none is supplied. |
583
- | `agents` | `["analyst", "helper"]` | Every registered agent id in order. |
584
-
585
-
586
- `chatUrl(config, agentId?)` returns `config.chatPath` for the default
587
- agent (the registered `chatRoute` mount that omits `:agentId`), and
588
- `${config.chatPath}/${encodeURIComponent(agentId)}` otherwise. Pure
589
- function: no React, no hooks, safe in service workers and SSR.
590
-
591
- ## License
592
-
593
- Apache-2.0
330
+ Use `mcp: false` to disable MCP. Turn on `tools: true` only for ambient tools
331
+ that are safe outside an in-process chat turn.
332
+
333
+ ## API Gate
334
+
335
+ The stock `@mastra/express` app has broad management routes. The plugin's
336
+ default `apiAccess: "scoped"` allows only the chat, read-only metadata,
337
+ plugin-owned `/route/*`, embed, model, suggestion, and MCP surfaces that the
338
+ client needs. Use `apiAccess: "full"` only for a trusted first-party console.
339
+
340
+ `server.isMastraRequestAllowed()` is exported for tests and custom dispatch
341
+ logic that need the same allowlist.
342
+
343
+ ## Configuration Reference
344
+
345
+ The plugin config is intentionally centered on the AppKit lifecycle instead of
346
+ requiring callers to assemble a Mastra server by hand.
347
+
348
+ - `agents` registers a single agent, an array, or a record keyed by stable agent
349
+ ids. Records are best for UIs because the ids become route-visible.
350
+ - `defaultAgent` controls which registered agent handles requests that do not
351
+ name an agent explicitly.
352
+ - `storage` and `memory` accept `true`, `false`, or concrete Mastra Postgres /
353
+ PgVector options. `true` resolves from `lakebase()` when present.
354
+ - `genie.spaces` maps aliases to Genie Space IDs. Those aliases flow into tools,
355
+ suggestions, and chart/data workflows.
356
+ - `defaultModel`, `modelOverride`, and `modelFuzzyMatch` control how loose model
357
+ names are resolved through Databricks Model Serving.
358
+ - `feedback` controls whether MLflow feedback routes are exposed. The automatic
359
+ mode enables feedback when tracing and an MLflow experiment are configured.
360
+ - `mcp` controls whether agents are exposed as MCP tools and how that server is
361
+ named.
362
+ - `apiAccess` chooses the route allowlist. Keep the default scoped mode for
363
+ deployed apps.
364
+
365
+ Use this package when you want an AppKit-native agent runtime. Use the shared
366
+ schemas in [`@dbx-tools/shared-mastra`](../../shared/mastra) when building a
367
+ client that talks to these routes.
368
+
369
+ ## Modules
370
+
371
+ - `plugin` - `MastraPlugin` and `mastra()` AppKit plugin factory.
372
+ - `agents` - `createAgent`, `tool`, `createTool`, agent build helpers, fallback
373
+ defaults, and approval-gated tool inspection.
374
+ - `config` - plugin config types and RequestContext key constants.
375
+ - `model` / `serving` / `servingSanitize` - Mastra model config, request
376
+ overrides, serving-endpoint config, and request-body cleanup.
377
+ - `genie` - Genie prompt, space normalization, Genie toolkits, and suggestions.
378
+ - `chart` / `statement` / `writer` - chart cache, statement row fetches, and
379
+ safe writer events.
380
+ - `history` / `threads` / `pagination` - conversation persistence helpers and
381
+ route handlers.
382
+ - `memory` / `storageSchema` - Lakebase-backed Mastra store/vector setup.
383
+ - `workspaces` / `filesystems` - Mastra workspace creation and Databricks
384
+ Workspace file adapters.
385
+ - `mcp` - MCP server construction.
386
+ - `observability` / `mlflow` - tracing and feedback.
387
+ - `server` / `rest` / `processors` - Express dispatch, Databricks REST helpers,
388
+ stream/result processors.
389
+
390
+ Browser-facing wire types are in
391
+ [`@dbx-tools/shared-mastra`](../../shared/mastra). Genie event contracts are in
392
+ [`@dbx-tools/shared-genie`](../../shared/genie). Model request/result contracts
393
+ are in [`@dbx-tools/shared-model`](../../shared/model). The matching React chat
394
+ surface is [`@dbx-tools/ui-mastra`](../../ui/mastra).