@dbx-tools/appkit-mastra 0.1.13 → 0.1.14

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 (59) hide show
  1. package/README.md +304 -645
  2. package/index.ts +46 -38
  3. package/package.json +58 -45
  4. package/src/agents.ts +224 -66
  5. package/src/chart.ts +531 -429
  6. package/src/config.ts +270 -19
  7. package/src/filesystems.ts +1090 -0
  8. package/src/genie.ts +1000 -660
  9. package/src/history.ts +166 -79
  10. package/src/mcp.ts +105 -0
  11. package/src/memory.ts +94 -92
  12. package/src/mlflow.ts +149 -0
  13. package/src/model.ts +75 -408
  14. package/src/observability.ts +121 -69
  15. package/src/pagination.ts +34 -0
  16. package/src/plugin.ts +552 -67
  17. package/src/processors.ts +168 -0
  18. package/src/rest.ts +67 -0
  19. package/src/server.ts +232 -45
  20. package/src/serving-sanitize.ts +167 -0
  21. package/src/serving.ts +27 -243
  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 -20
  30. package/dist/index.js +0 -20
  31. package/dist/src/agents.d.ts +0 -306
  32. package/dist/src/agents.js +0 -403
  33. package/dist/src/chart.d.ts +0 -170
  34. package/dist/src/chart.js +0 -491
  35. package/dist/src/config.d.ts +0 -183
  36. package/dist/src/config.js +0 -12
  37. package/dist/src/genie.d.ts +0 -131
  38. package/dist/src/genie.js +0 -630
  39. package/dist/src/history.d.ts +0 -67
  40. package/dist/src/history.js +0 -172
  41. package/dist/src/memory.d.ts +0 -100
  42. package/dist/src/memory.js +0 -242
  43. package/dist/src/model.d.ts +0 -159
  44. package/dist/src/model.js +0 -427
  45. package/dist/src/observability.d.ts +0 -33
  46. package/dist/src/observability.js +0 -71
  47. package/dist/src/plugin.d.ts +0 -130
  48. package/dist/src/plugin.js +0 -283
  49. package/dist/src/processors/strip-stale-charts.d.ts +0 -29
  50. package/dist/src/processors/strip-stale-charts.js +0 -96
  51. package/dist/src/server.d.ts +0 -46
  52. package/dist/src/server.js +0 -123
  53. package/dist/src/serving.d.ts +0 -156
  54. package/dist/src/serving.js +0 -231
  55. package/dist/src/tools/email.d.ts +0 -74
  56. package/dist/src/tools/email.js +0 -122
  57. package/dist/tsconfig.build.tsbuildinfo +0 -1
  58. package/src/processors/strip-stale-charts.ts +0 -105
  59. package/src/tools/email.ts +0 -147
@@ -0,0 +1,167 @@
1
+ import { string } from "@dbx-tools/shared-core";
2
+ /**
3
+ * Repairs Mastra / AI SDK message replays sent to Databricks Model
4
+ * Serving before they hit the OpenAI-compatible `/chat/completions`
5
+ * route.
6
+ */
7
+
8
+ /**
9
+ * OpenAI-flavoured chat message shape we need to mutate. We do not
10
+ * import the OpenAI / AI SDK types because both packages keep these
11
+ * fields under internal namespaces; the wire payload is the contract
12
+ * here and it's stable enough to inline.
13
+ */
14
+ export interface ServingChatMessage {
15
+ role: "system" | "user" | "assistant" | "tool" | "reasoning";
16
+ content?: string | ServingContentPart[];
17
+ tool_calls?: Array<{ id: string; type: string; function: unknown }>;
18
+ tool_call_id?: string;
19
+ reasoning?: unknown;
20
+ reasoning_content?: unknown;
21
+ }
22
+
23
+ type ServingContentPart = {
24
+ type?: string;
25
+ text?: string;
26
+ [key: string]: unknown;
27
+ };
28
+
29
+ const REASONING_PART_TYPES = new Set(["reasoning", "thinking", "redacted_thinking"]);
30
+
31
+ /**
32
+ * Parse, sanitize, and re-serialize a `/serving-endpoints/...` POST
33
+ * body. Returns the original string verbatim when the body is not
34
+ * JSON, has no `messages`, or no rewrite was needed.
35
+ */
36
+ export function rewriteServingBody(body: string): string {
37
+ let parsed: { messages?: unknown };
38
+ try {
39
+ parsed = JSON.parse(body);
40
+ } catch {
41
+ return body;
42
+ }
43
+ if (!Array.isArray(parsed.messages)) return body;
44
+ const messages = parsed.messages as ServingChatMessage[];
45
+ const changed = stripReasoningFromServingMessages(messages) || repairAssistantPrefill(messages);
46
+ return changed ? JSON.stringify(parsed) : body;
47
+ }
48
+
49
+ /**
50
+ * Drop extended-thinking / reasoning blocks from a replayed transcript.
51
+ *
52
+ * Hybrid Claude endpoints (e.g. Sonnet 4.5+) may emit `reasoning` content
53
+ * parts on the first turn. Mastra persists and replays them on the next
54
+ * agent step, but Databricks-hosted Claude rejects those blocks on
55
+ * multi-turn tool continuations. The UI already captured reasoning for
56
+ * display; stripping here keeps provider replay compatible without
57
+ * changing what users see in the chat bubble.
58
+ */
59
+ export function stripReasoningFromServingMessages(messages: ServingChatMessage[]): boolean {
60
+ let changed = false;
61
+ for (let i = messages.length - 1; i >= 0; i--) {
62
+ const msg = messages[i];
63
+ if (!msg) continue;
64
+ if (msg.role === "reasoning") {
65
+ messages.splice(i, 1);
66
+ changed = true;
67
+ continue;
68
+ }
69
+ if (msg.reasoning !== undefined) {
70
+ delete msg.reasoning;
71
+ changed = true;
72
+ }
73
+ if (msg.reasoning_content !== undefined) {
74
+ delete msg.reasoning_content;
75
+ changed = true;
76
+ }
77
+ if (!Array.isArray(msg.content)) continue;
78
+ const filtered = msg.content.filter((part) => {
79
+ const type = part?.type;
80
+ if (typeof type === "string" && REASONING_PART_TYPES.has(type)) {
81
+ changed = true;
82
+ return false;
83
+ }
84
+ return true;
85
+ });
86
+ if (filtered.length !== msg.content.length) {
87
+ msg.content = filtered;
88
+ }
89
+ const hasToolCalls = Array.isArray(msg.tool_calls) && msg.tool_calls.length > 0;
90
+ if (msg.role === "assistant" && !hasToolCalls && isEmptyServingContent(msg.content)) {
91
+ messages.splice(i, 1);
92
+ changed = true;
93
+ }
94
+ }
95
+ return changed;
96
+ }
97
+
98
+ /**
99
+ * Repair a Mastra/AI SDK message replay that Databricks-hosted Claude
100
+ * rejects with `"This model does not support assistant message
101
+ * prefill. The conversation must end with a user message."`.
102
+ *
103
+ * The bug pattern: when an assistant turn streams text *and* a
104
+ * `tool_call`, the AI SDK persists them as two separate assistant
105
+ * entries (text-only and tool-call-only). On the next agent step the
106
+ * tool-call entry is replayed *before* the tool result and the
107
+ * text entry is replayed *after* it, so the conversation ends with a
108
+ * trailing assistant text message. Anthropic interprets that as a
109
+ * prefill request and rejects it on Databricks (the upstream Bedrock
110
+ * route disallows prefill).
111
+ *
112
+ * Fix: when the last message is an assistant text with no `tool_calls`
113
+ * and the chain immediately before it is `assistant(tool_calls=...)`
114
+ * followed only by `tool(...)` results, fold the trailing text back
115
+ * into the `content` of that opening assistant and drop the duplicate.
116
+ */
117
+ export function repairAssistantPrefill(messages: ServingChatMessage[]): boolean {
118
+ if (messages.length < 2) return false;
119
+ const last = messages[messages.length - 1];
120
+ if (!last || last.role !== "assistant" || (last.tool_calls && last.tool_calls.length > 0)) {
121
+ return false;
122
+ }
123
+
124
+ let i = messages.length - 2;
125
+ while (i >= 0 && messages[i]?.role === "tool") i--;
126
+ if (i < 0) return false;
127
+ const opener = messages[i];
128
+ if (
129
+ !opener ||
130
+ opener.role !== "assistant" ||
131
+ !opener.tool_calls ||
132
+ opener.tool_calls.length === 0
133
+ ) {
134
+ return false;
135
+ }
136
+
137
+ const merged = [
138
+ string.trimToNull(textFromServingContent(opener.content)),
139
+ string.trimToNull(textFromServingContent(last.content)),
140
+ ]
141
+ .filter((s): s is string => s !== null)
142
+ .join("\n\n");
143
+ opener.content = merged;
144
+ messages.pop();
145
+ return true;
146
+ }
147
+
148
+ function textFromServingContent(content: ServingChatMessage["content"]): string {
149
+ if (typeof content === "string") return content;
150
+ if (!Array.isArray(content)) return "";
151
+ return content
152
+ .filter((part) => part?.type === "text" && typeof part.text === "string")
153
+ .map((part) => part.text as string)
154
+ .join("\n\n");
155
+ }
156
+
157
+ function isEmptyServingContent(content: ServingChatMessage["content"]): boolean {
158
+ if (content === undefined) return true;
159
+ if (typeof content === "string") return content.trim().length === 0;
160
+ if (!Array.isArray(content)) return true;
161
+ return content.every((part) => {
162
+ if (part?.type === "text") {
163
+ return typeof part.text !== "string" || part.text.trim().length === 0;
164
+ }
165
+ return false;
166
+ });
167
+ }
package/src/serving.ts CHANGED
@@ -1,43 +1,19 @@
1
1
  /**
2
- * Dynamic model resolution against Databricks Model Serving.
2
+ * Mastra-specific glue over the generic `@dbx-tools/model` toolkit.
3
3
  *
4
- * Three concerns live here:
5
- *
6
- * 1. **Listing** - {@link listServingEndpoints} pulls the workspace's
7
- * `/serving-endpoints` via the SDK and caches the result per host
8
- * with a TTL. Concurrent callers share one in-flight promise (the
9
- * same coalescing pattern as Python's `cachetools-async`).
10
- * 2. **Fuzzy matching** - {@link resolveModelId} runs the user's input
11
- * through `fuse.js` extended search so loose tokens like
12
- * `"claude sonnet"` snap to `databricks-claude-sonnet-4-6` even
13
- * when typed without the full endpoint name.
14
- * 3. **Per-request override** - {@link extractModelOverride} pulls a
15
- * model name from the `X-Mastra-Model` header, `?model=` query
16
- * string, or `model` body field so the same agent can be exercised
17
- * against different endpoints without redeploying.
18
- *
19
- * `model.ts` glues these together inside the per-step model resolver;
20
- * `plugin.ts` exposes the cached list at `GET /models`.
4
+ * The live `/serving-endpoints` catalogue access, fuzzy name
5
+ * resolution, and class/fallback selection all live in
6
+ * `@dbx-tools/model`; this module only adds what is specific to the
7
+ * Mastra plugin: pulling a per-request model override off an HTTP
8
+ * request (header / query / body) and projecting the plugin config
9
+ * onto the knobs `buildModel` and the `/models` route share.
21
10
  */
22
11
 
23
- import { CacheManager, type getExecutionContext } from "@databricks/appkit";
24
- import { logUtils, stringUtils } from "@dbx-tools/appkit-shared";
25
- import Fuse from "fuse.js";
26
-
27
- import type { ServingEndpointSummary } from "@dbx-tools/appkit-mastra-shared";
28
- import type { MastraPluginConfig } from "./config.js";
12
+ import { override } from "@dbx-tools/shared-mastra";
13
+ import { serving as nodeServing } from "@dbx-tools/model";
29
14
 
30
- export type { ServingEndpointSummary };
31
-
32
- const log = logUtils.logger("mastra/serving");
33
-
34
- /**
35
- * Structural type for the Databricks workspace client. Derived from
36
- * AppKit's `ExecutionContext` so this module doesn't take a direct
37
- * dependency on `@databricks/sdk-experimental`; the dep flows in
38
- * transitively through `@databricks/appkit`.
39
- */
40
- type WorkspaceClientLike = ReturnType<typeof getExecutionContext>["client"];
15
+ import type { MastraPluginConfig } from "./config";
16
+ import { string } from "@dbx-tools/shared-core";
41
17
 
42
18
  /**
43
19
  * `RequestContext` key under which {@link MastraServer} stores the
@@ -46,196 +22,6 @@ type WorkspaceClientLike = ReturnType<typeof getExecutionContext>["client"];
46
22
  */
47
23
  export const MASTRA_MODEL_OVERRIDE_KEY = "mastra__model_override";
48
24
 
49
- /** HTTP header inspected for a per-request model override. */
50
- export const MODEL_OVERRIDE_HEADER = "x-mastra-model";
51
-
52
- /** Query string parameter inspected for a per-request model override. */
53
- export const MODEL_OVERRIDE_QUERY = "model";
54
-
55
- /** Body fields (in priority order) inspected for a per-request model override. */
56
- export const MODEL_OVERRIDE_BODY_FIELDS = ["model", "modelId"] as const;
57
-
58
- /** Default TTL for the in-memory endpoint cache. Matches the Databricks SDK's session lifetime budget. */
59
- const DEFAULT_TTL_MS = 5 * 60 * 1000;
60
-
61
- /** Default Fuse.js score threshold below which a fuzzy match is accepted. */
62
- const DEFAULT_FUZZY_THRESHOLD = 0.4;
63
-
64
- /** Cache key parts under which endpoint listings are stored. */
65
- const CACHE_KEY_NAMESPACE = "mastra:serving-endpoints";
66
-
67
- /**
68
- * Stable `userKey` arg for AppKit's `CacheManager.getOrExecute`.
69
- * Endpoint visibility is effectively workspace-scoped (we cache by
70
- * host in the key parts), so a single shared key lets every user of
71
- * the same workspace share one cached fetch and coalesce on the
72
- * in-flight promise. Permissions can differ in theory, but the
73
- * Foundation Model API catalogue is the same view for every caller.
74
- */
75
- const SHARED_USER_KEY = "mastra-shared";
76
-
77
- /**
78
- * List Model Serving endpoints for the workspace owning `client`,
79
- * routed through AppKit's `CacheManager`. The manager gives us
80
- * everything `cachetools.TTLCache` provides plus what
81
- * `cachetools-async` adds on top: per-entry TTL, in-flight request
82
- * coalescing (concurrent callers share one fetch via the manager's
83
- * internal `inFlightRequests` map), bounded size, telemetry spans
84
- * (`cache.getOrExecute`), and optional Lakebase persistence so the
85
- * catalogue survives restarts when the lakebase plugin is wired up.
86
- *
87
- * Returns plain {@link ServingEndpointSummary} objects (a stable
88
- * subset of the SDK type) so cache hits never expose stale SDK
89
- * internals. Errors from `CacheManager` or the SDK fetch propagate
90
- * to the caller - we don't swallow them so users see the real
91
- * auth / network issue.
92
- *
93
- * @param host - Workspace host used as the cache key. Pass the value
94
- * resolved from `client.config.getHost()` so multi-host apps share
95
- * one entry per workspace.
96
- * @param opts.ttlMs - Override the default TTL just for this call.
97
- * Forwarded to `CacheManager` as seconds.
98
- */
99
- export async function listServingEndpoints(
100
- client: WorkspaceClientLike,
101
- host: string,
102
- opts: { ttlMs?: number } = {},
103
- ): Promise<ServingEndpointSummary[]> {
104
- const ttlSec = Math.max(1, Math.round((opts.ttlMs ?? DEFAULT_TTL_MS) / 1000));
105
- return CacheManager.getInstanceSync().getOrExecute(
106
- [CACHE_KEY_NAMESPACE, host],
107
- () => fetchEndpoints(client),
108
- SHARED_USER_KEY,
109
- { ttl: ttlSec },
110
- );
111
- }
112
-
113
- async function fetchEndpoints(
114
- client: WorkspaceClientLike,
115
- ): Promise<ServingEndpointSummary[]> {
116
- const startedAt = Date.now();
117
- const out: ServingEndpointSummary[] = [];
118
- for await (const ep of client.servingEndpoints.list()) {
119
- if (!ep.name) continue;
120
- out.push({
121
- name: ep.name,
122
- ...(ep.task !== undefined ? { task: ep.task } : {}),
123
- ...(ep.state?.ready !== undefined ? { state: String(ep.state.ready) } : {}),
124
- ...(ep.description !== undefined ? { description: ep.description } : {}),
125
- });
126
- }
127
- log.debug("listed", { count: out.length, elapsedMs: Date.now() - startedAt });
128
- return out;
129
- }
130
-
131
- /**
132
- * Force-evict cached endpoint listings via AppKit's `CacheManager`.
133
- * With a `host` deletes that one workspace's entry; without one
134
- * clears every cache entry on the manager (since `CacheManager`
135
- * doesn't expose a namespace-scoped clear, this is the brute-force
136
- * path - fine for tests, avoid in steady-state code).
137
- */
138
- export async function clearServingEndpointsCache(host?: string): Promise<void> {
139
- const cache = CacheManager.getInstanceSync();
140
- if (host) {
141
- const key = cache.generateKey([CACHE_KEY_NAMESPACE, host], SHARED_USER_KEY);
142
- await cache.delete(key);
143
- } else {
144
- await cache.clear();
145
- }
146
- }
147
-
148
- /**
149
- * Result of fuzzy-resolving a user-supplied model name against the
150
- * live endpoint list. `score` is Fuse.js's distance (`0` is exact,
151
- * `1` is no match); `matched` is `false` when the score exceeds the
152
- * configured threshold so callers can fall back to the original
153
- * input (Databricks will then return a clean 404).
154
- */
155
- export interface ResolvedModel {
156
- modelId: string;
157
- matched: boolean;
158
- score?: number;
159
- }
160
-
161
- /** Options accepted by {@link resolveModelId}. */
162
- export interface ResolveModelOptions {
163
- /** Fuse.js threshold (0 = exact, 1 = anything). Default `0.4`. */
164
- threshold?: number;
165
- }
166
-
167
- /**
168
- * Snap a user-supplied model name to the closest configured serving
169
- * endpoint:
170
- *
171
- * 1. Exact name match wins immediately (no fuzzy needed).
172
- * 2. Otherwise the input is tokenized (dashes / underscores / spaces
173
- * become separators) and fed through Fuse.js extended search,
174
- * which AND-s each token with fuzzy matching enabled. This is the
175
- * "tokenized fuzzy match" the user reaches for when they type
176
- * `"claude sonnet"` instead of the full endpoint name.
177
- * 3. If the best Fuse score is above `threshold`, return the input
178
- * unchanged and let the upstream call surface the 404. This keeps
179
- * deliberate model ids (e.g. brand new endpoints) from being
180
- * silently rewritten to a similar-looking neighbour.
181
- *
182
- * Pass an empty endpoint list to short-circuit fuzzy matching - the
183
- * input is returned verbatim. This is what {@link buildModel} does
184
- * when the workspace client can't be reached at resolve time.
185
- */
186
- export function resolveModelId(
187
- input: string,
188
- endpoints: readonly ServingEndpointSummary[],
189
- opts: ResolveModelOptions = {},
190
- ): ResolvedModel {
191
- if (endpoints.length === 0) {
192
- log.debug("resolve:no-endpoints", { input });
193
- return { modelId: input, matched: false };
194
- }
195
- for (const ep of endpoints) {
196
- if (ep.name === input) {
197
- log.debug("resolve:exact", { input });
198
- return { modelId: ep.name, matched: true, score: 0 };
199
- }
200
- }
201
- const threshold = opts.threshold ?? DEFAULT_FUZZY_THRESHOLD;
202
- const fuse = new Fuse(endpoints, {
203
- keys: ["name"],
204
- threshold,
205
- ignoreLocation: true,
206
- includeScore: true,
207
- useExtendedSearch: true,
208
- isCaseSensitive: false,
209
- });
210
- // Fuse 7.3 has no built-in tokenize hook; in extended search,
211
- // space-separated tokens are AND-ed with fuzzy matching enabled. We
212
- // lean on the shared tokenizer so the splitting rules stay
213
- // consistent with the rest of the toolkit.
214
- const query = Array.from(
215
- stringUtils.tokenizeWithOptions({ lowerCase: true, camelCase: false }, input),
216
- ).join(" ");
217
- if (!query) {
218
- log.debug("resolve:empty-tokens", { input });
219
- return { modelId: input, matched: false };
220
- }
221
- const results = fuse.search(query);
222
- const best = results[0];
223
- if (best?.item.name && (best.score ?? 0) <= threshold) {
224
- log.debug("resolve:fuzzy-match", {
225
- input,
226
- modelId: best.item.name,
227
- score: best.score,
228
- });
229
- return { modelId: best.item.name, matched: true, score: best.score };
230
- }
231
- log.debug("resolve:no-match", {
232
- input,
233
- bestScore: best?.score,
234
- threshold,
235
- });
236
- return { modelId: input, matched: false };
237
- }
238
-
239
25
  /**
240
26
  * Minimal Express-ish request shape used by {@link extractModelOverride}.
241
27
  * Keeps this module independent of `express` so the helper can be
@@ -264,19 +50,20 @@ export interface ModelOverrideRequest {
264
50
  export function extractModelOverride(req: ModelOverrideRequest): string | null {
265
51
  const headers = req.headers;
266
52
  if (headers) {
267
- const headerVal = stringUtils.firstNonEmpty(
268
- headers[MODEL_OVERRIDE_HEADER] ?? headers[MODEL_OVERRIDE_HEADER.toLowerCase()],
53
+ const headerVal = string.firstNonEmpty(
54
+ headers[override.MODEL_OVERRIDE_HEADER] ??
55
+ headers[override.MODEL_OVERRIDE_HEADER.toLowerCase()],
269
56
  );
270
57
  if (headerVal) return headerVal;
271
58
  }
272
59
  if (req.query) {
273
- const queryVal = stringUtils.firstNonEmpty(req.query[MODEL_OVERRIDE_QUERY]);
60
+ const queryVal = string.firstNonEmpty(req.query[override.MODEL_OVERRIDE_QUERY]);
274
61
  if (queryVal) return queryVal;
275
62
  }
276
63
  if (req.body && typeof req.body === "object") {
277
64
  const record = req.body as Record<string, unknown>;
278
- for (const field of MODEL_OVERRIDE_BODY_FIELDS) {
279
- const bodyVal = stringUtils.firstNonEmpty(record[field]);
65
+ for (const field of override.MODEL_OVERRIDE_BODY_FIELDS) {
66
+ const bodyVal = string.firstNonEmpty(record[field]);
280
67
  if (bodyVal) return bodyVal;
281
68
  }
282
69
  }
@@ -285,18 +72,15 @@ export function extractModelOverride(req: ModelOverrideRequest): string | null {
285
72
 
286
73
  /**
287
74
  * Read the fuzzy-resolution config knobs off the plugin config with
288
- * defaults applied. Kept here so `buildModel` and the `/models` route
289
- * agree on what "enabled" means.
75
+ * `@dbx-tools/model` defaults applied. Kept here so `buildModel` and
76
+ * the `/models` route agree on what "enabled" means.
290
77
  *
291
- * `fallbacks` is the priority-ordered list `pickModelId` walks when
292
- * nothing explicit is set; defaults live in `model.ts`
293
- * (`FALLBACK_MODEL_IDS`) and are passed in by callers to avoid a
294
- * circular import between `serving.ts` and `model.ts`.
78
+ * `fallbacks` is the priority-ordered list `resolveModel` walks first
79
+ * when nothing explicit is set; defaults to an empty list so the
80
+ * generic resolver falls through to the live catalogue and its own
81
+ * `FALLBACK_MODEL_IDS` floor.
295
82
  */
296
- export function resolveServingConfig(
297
- config: MastraPluginConfig,
298
- defaultFallbacks: readonly string[] = [],
299
- ): {
83
+ export function resolveServingConfig(config: MastraPluginConfig): {
300
84
  ttlMs: number;
301
85
  threshold: number;
302
86
  fuzzy: boolean;
@@ -304,10 +88,10 @@ export function resolveServingConfig(
304
88
  fallbacks: readonly string[];
305
89
  } {
306
90
  return {
307
- ttlMs: config.modelCacheTtlMs ?? DEFAULT_TTL_MS,
308
- threshold: config.modelFuzzyThreshold ?? DEFAULT_FUZZY_THRESHOLD,
91
+ ttlMs: config.modelCacheTtlMs ?? nodeServing.DEFAULT_MODEL_CACHE_TTL_MS,
92
+ threshold: config.modelFuzzyThreshold ?? nodeServing.DEFAULT_FUZZY_THRESHOLD,
309
93
  fuzzy: config.modelFuzzyMatch !== false,
310
94
  allowOverride: config.modelOverride !== false,
311
- fallbacks: config.defaultModelFallbacks ?? defaultFallbacks,
95
+ fallbacks: config.defaultModelFallbacks ?? [],
312
96
  };
313
97
  }
@@ -0,0 +1,89 @@
1
+ /**
2
+ * Databricks Statement Execution helpers for the Mastra plugin.
3
+ *
4
+ * Wraps `client.statementExecution.getStatement` with the shape and
5
+ * size handling the plugin's tools and the `/embed/data/:id` route
6
+ * both need: a low-level fetch that returns a raw
7
+ * `{columns, rows, rowCount}` shape and coerces numeric strings to
8
+ * numbers so downstream charts and aggregations don't have to, plus a
9
+ * hard row cap callers clamp `limit` to so a runaway result set can't
10
+ * hose a response. Upstream 404s are detected via
11
+ * `error.errorContext(err).notFound` at the call sites.
12
+ *
13
+ * Not Genie-specific: a Databricks `statement_id` is workspace
14
+ * scoped and lives in the Statement Execution API regardless of
15
+ * which producer (Genie, a tool, a notebook, etc.) submitted the
16
+ * query. Co-located here so consumers can fetch / cap / handle
17
+ * 404s without reaching into the Genie tool module.
18
+ */
19
+
20
+ import { WorkspaceClient } from "@databricks/sdk-experimental";
21
+ import type { GenieDatasetData } from "@dbx-tools/shared-mastra";
22
+ import { databricks } from "@dbx-tools/appkit";
23
+
24
+ /**
25
+ * Hard server-side cap on rows returned by the
26
+ * `/embed/data/:id` route. Sized to keep responses small
27
+ * enough for inline tables to render snappily; the route surfaces
28
+ * a `truncated` flag whenever the upstream `rowCount` exceeds
29
+ * this so end users know they're seeing a sample.
30
+ */
31
+ export const STATEMENT_ROW_CAP = 500;
32
+
33
+ /**
34
+ * Best-effort numeric coercion for the Statement Execution API's
35
+ * all-strings cells. Leaves non-numeric strings (and explicit
36
+ * `null`s) intact; everything else flows through `Number`.
37
+ */
38
+ function coerceCell(cell: string | null): unknown {
39
+ if (cell === null) return null;
40
+ if (/^-?\d+(\.\d+)?$/.test(cell)) {
41
+ const n = Number(cell);
42
+ if (Number.isFinite(n)) return n;
43
+ }
44
+ return cell;
45
+ }
46
+
47
+ /**
48
+ * Fetch a single statement's rows via the Statement Execution API
49
+ * and reshape into the shared {@link GenieDatasetData} shape
50
+ * (column array + row records).
51
+ *
52
+ * Optional `limit` slices the returned `rows` client-side so the
53
+ * agent can scan a small sample without paging the full result
54
+ * set into context. `rowCount` always reflects the upstream total
55
+ * so callers know when the slice truncated.
56
+ *
57
+ * Exported because every consumer in the plugin (the
58
+ * `get_statement` tool, the `prepare_chart` dataset resolver, and
59
+ * the `/embed/data/:id` route) needs the exact same
60
+ * fetch + coercion pipeline so LLM-side `get_statement` output
61
+ * and UI-side `[data:<id>]` rendering stay shape-identical for
62
+ * the same `statement_id`.
63
+ */
64
+ export async function fetchStatementData(
65
+ client: WorkspaceClient,
66
+ statementId: string,
67
+ options?: { limit?: number; signal?: AbortSignal },
68
+ ): Promise<GenieDatasetData> {
69
+ const ctx = options?.signal ? databricks.toContext(options.signal) : undefined;
70
+ const r = await client.statementExecution.getStatement({ statement_id: statementId }, ctx);
71
+ const columns = (r.manifest?.schema?.columns ?? []).map((c) => c.name ?? "");
72
+ const dataArray = (r.result?.data_array ?? []) as Array<Array<string | null>>;
73
+ const sliced =
74
+ options?.limit !== undefined && options.limit >= 0
75
+ ? dataArray.slice(0, options.limit)
76
+ : dataArray;
77
+ const rows = sliced.map((row) => {
78
+ const obj: Record<string, unknown> = {};
79
+ columns.forEach((col, i) => {
80
+ obj[col] = coerceCell(row[i] ?? null);
81
+ });
82
+ return obj;
83
+ });
84
+ return {
85
+ columns,
86
+ rows,
87
+ rowCount: r.manifest?.total_row_count ?? dataArray.length,
88
+ };
89
+ }
@@ -0,0 +1,41 @@
1
+ import { string } from "@dbx-tools/shared-core";
2
+ /**
3
+ * Derive a Postgres-safe schema name for per-agent Mastra storage.
4
+ *
5
+ * Agent ids are often kebab-case route segments (`data-mesh-book-assistant`)
6
+ * but {@link PostgresStore} validates `schemaName` with Mastra's
7
+ * `parseSqlIdentifier` (letter/underscore start, `[A-Za-z0-9_]` only, max 63).
8
+ */
9
+
10
+ const SCHEMA_PREFIX = "mastra_";
11
+ const MAX_PG_IDENTIFIER_LEN = 63;
12
+
13
+ /**
14
+ * Default Lakebase schema for one agent's thread/message store:
15
+ * `mastra_<sanitized-agent-id>`.
16
+ */
17
+ export function agentStorageSchemaName(agentId: string): string {
18
+ const maxSlugLen = MAX_PG_IDENTIFIER_LEN - SCHEMA_PREFIX.length;
19
+ const slug = string.toIdentifierWithOptions(
20
+ {
21
+ delimiter: "_",
22
+ maxLength: maxSlugLen,
23
+ truncateStrategy: "hash",
24
+ truncateHashLength: 6,
25
+ },
26
+ agentId,
27
+ );
28
+ const body =
29
+ slug ||
30
+ string.toIdentifierWithOptions(
31
+ {
32
+ delimiter: "_",
33
+ maxLength: maxSlugLen,
34
+ truncateStrategy: "hash",
35
+ truncateHashLength: 6,
36
+ },
37
+ "agent",
38
+ agentId,
39
+ );
40
+ return `${SCHEMA_PREFIX}${body}`;
41
+ }