@dbx-tools/appkit-mastra 0.9.69 → 0.9.71

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.
@@ -0,0 +1,252 @@
1
+ /**
2
+ * OpenAPI operation loading and JSON Schema extraction.
3
+ *
4
+ * @module
5
+ */
6
+ import { readFile } from "node:fs/promises";
7
+ import { object } from "@dbx-tools/shared-core";
8
+ import { urlBuilder, type UrlBuilder } from "@dbx-tools/shared-core/net";
9
+ import type { JSONSchema7 } from "json-schema";
10
+ import { z } from "zod";
11
+
12
+ const HTTP_METHODS = [
13
+ "GET",
14
+ "POST",
15
+ "PUT",
16
+ "PATCH",
17
+ "DELETE",
18
+ "OPTIONS",
19
+ "HEAD",
20
+ "TRACE",
21
+ ] as const;
22
+
23
+ const JsonObjectSchema = z
24
+ .record(z.string(), z.unknown())
25
+ .describe("A JSON object embedded in an OpenAPI document.");
26
+
27
+ const OpenApiMediaTypeSchema = z
28
+ .object({
29
+ schema: JsonObjectSchema.describe("The JSON Schema for this media type."),
30
+ })
31
+ .describe("One OpenAPI media type.");
32
+
33
+ const OpenApiRequestBodySchema = z
34
+ .object({
35
+ content: z
36
+ .record(z.string(), OpenApiMediaTypeSchema)
37
+ .describe("Request media types keyed by content type."),
38
+ })
39
+ .describe("An OpenAPI request body.");
40
+
41
+ const OpenApiResponseSchema = z
42
+ .object({
43
+ content: z
44
+ .record(z.string(), OpenApiMediaTypeSchema)
45
+ .optional()
46
+ .describe("Response media types keyed by content type."),
47
+ })
48
+ .describe("One OpenAPI response.");
49
+
50
+ const OpenApiOperationSchema = z
51
+ .object({
52
+ operationId: z.string().describe("The stable operation identifier."),
53
+ summary: z.string().optional().describe("The short operation summary."),
54
+ description: z.string().optional().describe("The full operation description."),
55
+ requestBody: OpenApiRequestBodySchema.optional().describe(
56
+ "The optional operation request body.",
57
+ ),
58
+ responses: z
59
+ .record(z.string(), OpenApiResponseSchema)
60
+ .describe("Operation responses keyed by HTTP status."),
61
+ })
62
+ .describe("The OpenAPI fields required to define a tool.");
63
+
64
+ const OpenApiPathSchema = z
65
+ .object({
66
+ get: OpenApiOperationSchema.optional(),
67
+ post: OpenApiOperationSchema.optional(),
68
+ put: OpenApiOperationSchema.optional(),
69
+ patch: OpenApiOperationSchema.optional(),
70
+ delete: OpenApiOperationSchema.optional(),
71
+ options: OpenApiOperationSchema.optional(),
72
+ head: OpenApiOperationSchema.optional(),
73
+ trace: OpenApiOperationSchema.optional(),
74
+ })
75
+ .describe("Operations available at one OpenAPI path.");
76
+
77
+ const OpenApiDocumentSchema = z
78
+ .object({
79
+ paths: z
80
+ .record(z.string(), OpenApiPathSchema)
81
+ .refine((paths) => Object.keys(paths).length > 0, "OpenAPI paths must not be empty")
82
+ .describe("OpenAPI path items keyed by URL path."),
83
+ components: JsonObjectSchema.optional().describe("Reusable OpenAPI components."),
84
+ servers: z
85
+ .array(
86
+ z.object({
87
+ url: z.string().describe("Server base URL."),
88
+ }),
89
+ )
90
+ .optional()
91
+ .describe("OpenAPI server definitions."),
92
+ })
93
+ .describe("An OpenAPI document used to define HTTP tools.");
94
+
95
+ type OpenApiDocument = z.infer<typeof OpenApiDocumentSchema>;
96
+ type OpenApiOperation = z.infer<typeof OpenApiOperationSchema>;
97
+
98
+ /** HTTP method supported by an {@link OpenApiTool}. */
99
+ export type OpenApiHttpMethod = (typeof HTTP_METHODS)[number];
100
+
101
+ /** One executable HTTP operation derived from an OpenAPI document. */
102
+ export interface OpenApiTool {
103
+ /** Stable operation identifier. */
104
+ id: string;
105
+ /** Model-facing operation description. */
106
+ description: string;
107
+ /** Resolved operation URL, or the OpenAPI path for a local document without a server URL. */
108
+ url: string;
109
+ /** HTTP method used by the operation. */
110
+ method: OpenApiHttpMethod;
111
+ /** JSON Schema describing operation input. */
112
+ inputSchema: JSONSchema7;
113
+ /** JSON Schema describing a successful operation result. */
114
+ outputSchema: JSONSchema7;
115
+ }
116
+
117
+ /** Request options used while loading an OpenAPI document. */
118
+ export interface OpenApiToolsOptions {
119
+ /** Cancels an HTTP request or file read. */
120
+ signal?: AbortSignal;
121
+ /** Headers sent when `source` is an HTTP(S) URL. */
122
+ headers?: RequestInit["headers"];
123
+ }
124
+
125
+ /** Load an HTTP(S) URL or local file and return its OpenAPI operations as tool schemas. */
126
+ export async function openApiTools(
127
+ source: string,
128
+ options: OpenApiToolsOptions = {},
129
+ ): Promise<OpenApiTool[]> {
130
+ const sourceUrl = httpUrl(source);
131
+ const document = OpenApiDocumentSchema.parse(
132
+ sourceUrl
133
+ ? await fetchDocument(sourceUrl, options)
134
+ : await readDocument(source, options.signal),
135
+ );
136
+ return Object.entries(document.paths).flatMap(([path, item]) =>
137
+ HTTP_METHODS.flatMap((method) => {
138
+ const operation = item[method.toLowerCase() as Lowercase<OpenApiHttpMethod>];
139
+ return operation ? [operationTool(document, sourceUrl, path, method, operation)] : [];
140
+ }),
141
+ );
142
+ }
143
+
144
+ function httpUrl(source: string): UrlBuilder | undefined {
145
+ const url = urlBuilder(source);
146
+ return url?.scheme === "http" || url?.scheme === "https" ? url : undefined;
147
+ }
148
+
149
+ async function fetchDocument(url: URL, options: OpenApiToolsOptions): Promise<unknown> {
150
+ const response = await fetch(url, {
151
+ signal: options.signal,
152
+ headers: options.headers,
153
+ });
154
+ if (!response.ok) {
155
+ throw new Error(`OpenAPI request failed with HTTP ${response.status}: ${url}`);
156
+ }
157
+ return response.json();
158
+ }
159
+
160
+ async function readDocument(path: string, signal?: AbortSignal): Promise<unknown> {
161
+ return JSON.parse(await readFile(path, { encoding: "utf8", signal }));
162
+ }
163
+
164
+ function operationTool(
165
+ document: OpenApiDocument,
166
+ sourceUrl: URL | undefined,
167
+ path: string,
168
+ method: OpenApiHttpMethod,
169
+ operation: OpenApiOperation,
170
+ ): OpenApiTool {
171
+ const requestSchema = jsonSchema(operation.requestBody?.content) ?? { type: "object" };
172
+ const responseSchema = successfulResponseSchema(operation);
173
+ return {
174
+ id: operation.operationId,
175
+ description:
176
+ operation.description?.trim() || operation.summary?.trim() || operation.operationId,
177
+ url: operationUrl(document, sourceUrl, path),
178
+ method,
179
+ inputSchema: resolveReferences(document, requestSchema) as JSONSchema7,
180
+ outputSchema: resolveReferences(document, responseSchema) as JSONSchema7,
181
+ };
182
+ }
183
+
184
+ function successfulResponseSchema(operation: OpenApiOperation): Record<string, unknown> {
185
+ for (const [status, response] of Object.entries(operation.responses)) {
186
+ const schema = jsonSchema(response.content);
187
+ if (/^2(?:\d{2}|XX)$/i.test(status) && schema) return schema;
188
+ }
189
+ return {};
190
+ }
191
+
192
+ function jsonSchema(
193
+ content: Record<string, z.infer<typeof OpenApiMediaTypeSchema>> | undefined,
194
+ ): Record<string, unknown> | undefined {
195
+ const mediaType = Object.entries(content ?? {}).find(([type]) =>
196
+ /^application\/(?:[^+;]+\+)?json(?:;|$)/i.test(type),
197
+ );
198
+ return mediaType?.[1].schema;
199
+ }
200
+
201
+ function operationUrl(document: OpenApiDocument, sourceUrl: URL | undefined, path: string): string {
202
+ const server = document.servers?.[0]?.url;
203
+ if (server) {
204
+ const base = sourceUrl ? urlBuilder(new URL(server, sourceUrl)) : httpUrl(server);
205
+ return base
206
+ ? base.withPathAppend(path).toString()
207
+ : `${server.replace(/\/$/, "")}/${path.replace(/^\//, "")}`;
208
+ }
209
+ return sourceUrl ? urlBuilder(sourceUrl)?.withPathReplace(path).toString() ?? path : path;
210
+ }
211
+
212
+ function resolveReferences(
213
+ document: OpenApiDocument,
214
+ value: unknown,
215
+ active: ReadonlySet<string> = new Set(),
216
+ ): unknown {
217
+ if (Array.isArray(value)) return value.map((item) => resolveReferences(document, item, active));
218
+ if (!object.isRecord(value)) return value;
219
+ const reference = typeof value.$ref === "string" ? value.$ref : undefined;
220
+ if (reference) {
221
+ if (active.has(reference)) throw new Error(`Cyclic OpenAPI reference: ${reference}`);
222
+ const target = referenceTarget(document, reference);
223
+ const next = new Set(active).add(reference);
224
+ const { $ref: _ignored, ...siblings } = value;
225
+ return {
226
+ ...(resolveReferences(document, target, next) as Record<string, unknown>),
227
+ ...(resolveReferences(document, siblings, next) as Record<string, unknown>),
228
+ };
229
+ }
230
+ return Object.fromEntries(
231
+ Object.entries(value).map(([key, child]) => [key, resolveReferences(document, child, active)]),
232
+ );
233
+ }
234
+
235
+ function referenceTarget(document: OpenApiDocument, reference: string): unknown {
236
+ if (!reference.startsWith("#/")) throw new Error(`Unsupported OpenAPI reference: ${reference}`);
237
+ let current: unknown = document;
238
+ for (const rawSegment of reference.slice(2).split("/")) {
239
+ const segment = rawSegment.replaceAll("~1", "/").replaceAll("~0", "~");
240
+ if (!object.isRecord(current) || !(segment in current)) {
241
+ throw new Error(`Unresolved OpenAPI reference: ${reference}`);
242
+ }
243
+ current = current[segment];
244
+ }
245
+ return current;
246
+ }
247
+
248
+ if (import.meta.main) {
249
+ const source = process.argv[2];
250
+ if (!source) throw new Error("Usage: bun openapi-tool.ts <openapi-url-or-file>");
251
+ console.log(JSON.stringify(await openApiTools(source), null, 2));
252
+ }
package/src/server.ts CHANGED
@@ -29,13 +29,13 @@ import {
29
29
  } from "./config.ts";
30
30
  import { requestUserEmail, requestUserId } from "./identity.ts";
31
31
  import { resolveFeedbackEnabled } from "./mlflow.ts";
32
- import { recordActiveTraceAuth, recordActiveTraceUser } from "./telemetry.ts";
33
-
34
32
  import {
35
33
  extractModelOverride,
36
34
  MASTRA_MODEL_OVERRIDE_KEY,
35
+ MASTRA_RESOLVED_MODEL_KEY,
37
36
  resolveServingConfig,
38
37
  } from "./serving.ts";
38
+ import { recordActiveTraceAuth, recordActiveTraceUser } from "./telemetry.ts";
39
39
  /**
40
40
  * OpenTelemetry's sentinel for "no valid trace" - 32 zero hex chars.
41
41
  * `trace.getActiveSpan()` returns a non-recording span with this id
@@ -52,6 +52,7 @@ const TRUSTED_REQUEST_CONTEXT_KEYS = [
52
52
  MASTRA_REQUEST_ID_KEY,
53
53
  MASTRA_SCOPES_KEY,
54
54
  MASTRA_MODEL_OVERRIDE_KEY,
55
+ MASTRA_RESOLVED_MODEL_KEY,
55
56
  ] as const;
56
57
 
57
58
  /** Remove browser-supplied values for fields owned by trusted server middleware. */
@@ -225,9 +226,7 @@ export class MastraServer extends ExpressMastraServer {
225
226
  });
226
227
  const user = requestContext.get(MASTRA_USER_KEY) as User | undefined;
227
228
  if (user) {
228
- recordActiveTraceAuth(
229
- "isUserContext" in user.executionContext ? "obo" : "service-principal",
230
- );
229
+ recordActiveTraceAuth("isUserContext" in user.executionContext ? "obo" : "service-principal");
231
230
  }
232
231
  const traceUser = [
233
232
  requestContext.get(MASTRA_USER_EMAIL_KEY),
package/src/serving.ts CHANGED
@@ -24,6 +24,9 @@ import type { MastraPluginConfig } from "./config.ts";
24
24
  */
25
25
  export const MASTRA_MODEL_OVERRIDE_KEY = "mastra__model_override";
26
26
 
27
+ /** RequestContext key carrying the endpoint selected for the active agent step. */
28
+ export const MASTRA_RESOLVED_MODEL_KEY = "mastra__resolved_model";
29
+
27
30
  /**
28
31
  * Minimal Express-ish request shape used by {@link extractModelOverride}.
29
32
  * Keeps this module independent of `express` so the helper can be
package/src/telemetry.ts CHANGED
@@ -36,11 +36,11 @@ import {
36
36
  updateActiveMlflowTraceTags,
37
37
  } from "./mlflow.ts";
38
38
  import {
39
+ APPKIT_AGENT_TRACE_ROOT_ATTR,
39
40
  MLFLOW_AGENT_TAG,
40
41
  MLFLOW_AGENT_TAG_ATTR,
41
42
  MLFLOW_GENIE_TAG,
42
43
  MLFLOW_GENIE_TAG_ATTR,
43
- MLFLOW_MODEL_TAG,
44
44
  MLFLOW_LOCAL_TAG,
45
45
  MLFLOW_OBO_AUTH_TAG,
46
46
  MLFLOW_SP_AUTH_TAG,
@@ -50,10 +50,9 @@ import {
50
50
  export {
51
51
  MLFLOW_AGENT_TAG,
52
52
  MLFLOW_AGENT_TAG_ATTR,
53
+ APPKIT_AGENT_TRACE_ROOT_ATTR,
53
54
  MLFLOW_GENIE_TAG,
54
55
  MLFLOW_GENIE_TAG_ATTR,
55
- MLFLOW_MODEL_TAG,
56
- MLFLOW_MODEL_TAG_ATTR,
57
56
  MLFLOW_LOCAL_TAG,
58
57
  MLFLOW_LOCAL_TAG_ATTR,
59
58
  MLFLOW_OBO_AUTH_TAG,
@@ -130,8 +129,8 @@ interface TraceTarget {
130
129
  }
131
130
 
132
131
  interface ChatTraceState {
133
- readonly models: Set<string>;
134
132
  readonly tags: Map<string, string>;
133
+ firstModel?: string;
135
134
  user?: string;
136
135
  }
137
136
 
@@ -187,21 +186,14 @@ function activeChatTraceState(): ChatTraceState | undefined {
187
186
  return context.active().getValue(CHAT_TRACE_STATE_KEY) as ChatTraceState | undefined;
188
187
  }
189
188
 
190
- function setChatTraceTag(
191
- state: ChatTraceState,
192
- name: string,
193
- value: ChatTraceTagValue,
194
- ): void {
189
+ function setChatTraceTag(state: ChatTraceState, name: string, value: ChatTraceTagValue): void {
195
190
  const tag = stringUtils.trimToNull(name);
196
191
  if (!tag) return;
197
192
  state.tags.set(tag, String(value));
198
193
  }
199
194
 
200
195
  /** Inject one MLflow trace tag into the active chat request. */
201
- export function recordActiveTraceTag(
202
- name: string,
203
- value: ChatTraceTagValue = true,
204
- ): void {
196
+ export function recordActiveTraceTag(name: string, value: ChatTraceTagValue = true): void {
205
197
  const state = activeChatTraceState();
206
198
  if (state) setChatTraceTag(state, name, value);
207
199
  }
@@ -217,11 +209,17 @@ export function recordActiveTraceTags(
217
209
  }
218
210
  }
219
211
 
220
- /** Record a selected model on the active chat trace, preserving first-use order. */
212
+ /** Record the first selected model as the active chat trace's `agent` tag. */
221
213
  export function recordActiveTraceModel(modelId: string): void {
222
214
  const value = stringUtils.trimToNull(modelId);
223
215
  if (!value) return;
224
- activeChatTraceState()?.models.add(value);
216
+ const state = activeChatTraceState();
217
+ if (!state || state.firstModel) return;
218
+ state.firstModel = value;
219
+ setChatTraceTag(state, MLFLOW_AGENT_TAG, value);
220
+ const span = trace.getActiveSpan();
221
+ if (span?.isRecording()) span.setAttribute(MLFLOW_AGENT_TAG_ATTR, value);
222
+ updateActiveMlflowTraceTags({ [MLFLOW_AGENT_TAG]: value });
225
223
  }
226
224
 
227
225
  /** Record the best resolved user identifier on the active chat trace. */
@@ -237,20 +235,6 @@ export function recordActiveTraceAuth(auth: "obo" | "service-principal"): void {
237
235
  recordActiveTraceTag(auth === "obo" ? MLFLOW_OBO_AUTH_TAG : MLFLOW_SP_AUTH_TAG);
238
236
  }
239
237
 
240
- /** Serialize model tags as a scalar or an ordered JSON array for multiple values. */
241
- export function modelTraceTag(models: Iterable<string>): string | undefined {
242
- const values: string[] = [];
243
- const seen = new Set<string>();
244
- for (const model of models) {
245
- const value = stringUtils.trimToNull(model);
246
- if (!value || seen.has(value)) continue;
247
- seen.add(value);
248
- values.push(value);
249
- }
250
- if (values.length === 0) return undefined;
251
- return values.length === 1 ? values[0] : JSON.stringify(values);
252
- }
253
-
254
238
  /** Publish the turn's MLflow trace id before the handler writes the body. */
255
239
  function stampMlflowTraceId(res: express.Response, span: Span): void {
256
240
  if (res.headersSent) return;
@@ -425,20 +409,18 @@ class AssistantResponseCollector {
425
409
  function resolveTraceTarget(): TraceTarget | undefined {
426
410
  const activeContext = context.active();
427
411
  if (directMlflowTracingActive()) {
428
- const span = trace
429
- .getTracer(CHAT_TURN_TRACER)
430
- .startSpan(
431
- "mastra.chat_turn",
432
- {
433
- kind: SpanKind.SERVER,
434
- attributes: {
435
- [MLFLOW_AGENT_TAG_ATTR]: "true",
436
- [MLFLOW_SPAN_TYPE_ATTR]: MLFLOW_SPAN_TYPE_AGENT,
437
- [GEN_AI_OPERATION_NAME_ATTR]: "invoke_agent",
438
- },
412
+ const span = trace.getTracer(CHAT_TURN_TRACER).startSpan(
413
+ "mastra.chat_turn",
414
+ {
415
+ kind: SpanKind.SERVER,
416
+ attributes: {
417
+ [APPKIT_AGENT_TRACE_ROOT_ATTR]: true,
418
+ [MLFLOW_SPAN_TYPE_ATTR]: MLFLOW_SPAN_TYPE_AGENT,
419
+ [GEN_AI_OPERATION_NAME_ATTR]: "invoke_agent",
439
420
  },
440
- trace.deleteSpan(activeContext),
441
- );
421
+ },
422
+ trace.deleteSpan(activeContext),
423
+ );
442
424
  if (!span.isRecording()) {
443
425
  span.end();
444
426
  return undefined;
@@ -454,20 +436,18 @@ function resolveTraceTarget(): TraceTarget | undefined {
454
436
  const active = trace.getActiveSpan();
455
437
  if (active?.isRecording()) return { span: active, owned: false };
456
438
 
457
- const span = trace
458
- .getTracer(CHAT_TURN_TRACER)
459
- .startSpan(
460
- "mastra.chat_turn",
461
- {
462
- kind: SpanKind.SERVER,
463
- attributes: {
464
- [MLFLOW_AGENT_TAG_ATTR]: "true",
465
- [MLFLOW_SPAN_TYPE_ATTR]: MLFLOW_SPAN_TYPE_AGENT,
466
- [GEN_AI_OPERATION_NAME_ATTR]: "invoke_agent",
467
- },
439
+ const span = trace.getTracer(CHAT_TURN_TRACER).startSpan(
440
+ "mastra.chat_turn",
441
+ {
442
+ kind: SpanKind.SERVER,
443
+ attributes: {
444
+ [APPKIT_AGENT_TRACE_ROOT_ATTR]: true,
445
+ [MLFLOW_SPAN_TYPE_ATTR]: MLFLOW_SPAN_TYPE_AGENT,
446
+ [GEN_AI_OPERATION_NAME_ATTR]: "invoke_agent",
468
447
  },
469
- activeContext,
470
- );
448
+ },
449
+ activeContext,
450
+ );
471
451
  if (!span.isRecording()) {
472
452
  span.end();
473
453
  return undefined;
@@ -506,9 +486,8 @@ export function chatTurnTelemetryMiddleware(
506
486
  next();
507
487
  return;
508
488
  }
509
- const traceState: ChatTraceState = { models: new Set(), tags: new Map() };
489
+ const traceState: ChatTraceState = { tags: new Map() };
510
490
  const traceRequestTime = Date.now();
511
- setChatTraceTag(traceState, MLFLOW_AGENT_TAG, true);
512
491
  if (!appkit.isDatabricksAppEnv()) {
513
492
  setChatTraceTag(traceState, MLFLOW_LOCAL_TAG, true);
514
493
  }
@@ -525,7 +504,6 @@ export function chatTurnTelemetryMiddleware(
525
504
  traceState.user = stampMlflowActor(target.span, req);
526
505
  target.span.setAttribute(MLFLOW_SPAN_TYPE_ATTR, MLFLOW_SPAN_TYPE_AGENT);
527
506
  target.span.setAttribute(GEN_AI_OPERATION_NAME_ATTR, "invoke_agent");
528
- target.span.setAttribute(MLFLOW_AGENT_TAG_ATTR, "true");
529
507
  const identity =
530
508
  typeof options.identity === "function" ? options.identity(req) : options.identity;
531
509
  if (identity) {
@@ -544,10 +522,7 @@ export function chatTurnTelemetryMiddleware(
544
522
  const fullInput = serialized(messages);
545
523
  requestPreview = textOnlyChatInput(messages) ?? fullInput;
546
524
  target.span.setAttribute(CHAT_MESSAGES_ATTR, fullInput);
547
- target.span.setAttribute(
548
- MLFLOW_SPAN_INPUTS_ATTR,
549
- requestPreview.slice(0, TRACE_IO_LIMIT),
550
- );
525
+ target.span.setAttribute(MLFLOW_SPAN_INPUTS_ATTR, requestPreview.slice(0, TRACE_IO_LIMIT));
551
526
  }
552
527
 
553
528
  const collector = new AssistantResponseCollector();
@@ -584,8 +559,6 @@ export function chatTurnTelemetryMiddleware(
584
559
  const finalizeTrace = (): void => {
585
560
  if (traceFinalized) return;
586
561
  traceFinalized = true;
587
- const model = modelTraceTag(traceState.models);
588
- if (model) setChatTraceTag(traceState, MLFLOW_MODEL_TAG, model);
589
562
  const tags = Object.fromEntries(traceState.tags);
590
563
  for (const [name, value] of traceState.tags) {
591
564
  target.span.setAttribute(`${MLFLOW_TRACE_TAG_PREFIX}${name}`, value);
@@ -651,7 +624,6 @@ export function chatTurnTelemetryMiddleware(
651
624
  .setSpan(context.active(), target.span)
652
625
  .setValue(CHAT_TRACE_STATE_KEY, traceState);
653
626
  context.with(chatContext, () => {
654
- updateActiveMlflowTraceTags({ [MLFLOW_AGENT_TAG]: "true" });
655
627
  next();
656
628
  });
657
629
  }
@@ -7,9 +7,6 @@ export const MLFLOW_AGENT_TAG = "agent";
7
7
  /** Experiment-UI tag marking a turn that called Genie. */
8
8
  export const MLFLOW_GENIE_TAG = "genie";
9
9
 
10
- /** Experiment-UI tag recording the model or ordered models used by a turn. */
11
- export const MLFLOW_MODEL_TAG = "model";
12
-
13
10
  /** Experiment-UI tag marking a turn that invoked OBO authentication. */
14
11
  export const MLFLOW_OBO_AUTH_TAG = "obo_auth";
15
12
 
@@ -25,8 +22,8 @@ export const MLFLOW_AGENT_TAG_ATTR = `${MLFLOW_TRACE_TAG_PREFIX}${MLFLOW_AGENT_T
25
22
  /** Attribute key for the `genie` trace tag. */
26
23
  export const MLFLOW_GENIE_TAG_ATTR = `${MLFLOW_TRACE_TAG_PREFIX}${MLFLOW_GENIE_TAG}`;
27
24
 
28
- /** Attribute key for the `model` trace tag. */
29
- export const MLFLOW_MODEL_TAG_ATTR = `${MLFLOW_TRACE_TAG_PREFIX}${MLFLOW_MODEL_TAG}`;
25
+ /** Internal root marker used by direct MLflow span filtering. */
26
+ export const APPKIT_AGENT_TRACE_ROOT_ATTR = "appkit.mastra.trace.agent";
30
27
 
31
28
  /** Attribute key for the `obo_auth` trace tag. */
32
29
  export const MLFLOW_OBO_AUTH_TAG_ATTR = `${MLFLOW_TRACE_TAG_PREFIX}${MLFLOW_OBO_AUTH_TAG}`;
package/src/workspaces.ts CHANGED
@@ -23,6 +23,7 @@ import { ConfigurationError, createWorkspaceClient } from "@databricks/appkit";
23
23
  import type { WorkspaceClient } from "@databricks/appkit";
24
24
  import { DatabricksFileSystem, workspaceClient } from "@dbx-tools/databricks";
25
25
  import { errorUtils, log, object, stringUtils, token } from "@dbx-tools/shared-core";
26
+ import { posixPath } from "@dbx-tools/shared-fs";
26
27
  import type { RequestContext } from "@mastra/core/request-context";
27
28
  import {
28
29
  CompositeFilesystem,
@@ -676,9 +677,83 @@ function buildWorkspaceSkillsResolver(
676
677
  extraSkillPaths: string[] = [],
677
678
  ): SkillsResolver {
678
679
  return async (context: SkillsContext) => {
679
- const { skillPaths } = await resolveWorkspaceContribution(resolvers, context);
680
- const merged = [...(skillPaths ?? []), ...extraSkillPaths];
681
- logger.debug("skills:resolved", { skillPaths, extraSkillPaths });
680
+ const contribution = await resolveWorkspaceContribution(resolvers, context);
681
+ const skillPaths = await uniqueSkillPaths(contribution);
682
+ const merged = [...skillPaths, ...extraSkillPaths];
683
+ logger.debug("skills:resolved", {
684
+ configuredSkillPaths: contribution.skillPaths,
685
+ skillPaths,
686
+ extraSkillPaths,
687
+ });
682
688
  return merged;
683
689
  };
684
690
  }
691
+
692
+ /** Expand mounted roots and keep the first concrete skill for each directory name. */
693
+ async function uniqueSkillPaths(contribution: WorkspaceMountContribution): Promise<string[]> {
694
+ const selected = new Map<string, string>();
695
+ const unresolved: string[] = [];
696
+ for (const root of contribution.skillPaths ?? []) {
697
+ const mount = mountedFilesystem(root, contribution.mounts);
698
+ if (!mount) {
699
+ unresolved.push(root);
700
+ continue;
701
+ }
702
+ let candidates: string[];
703
+ try {
704
+ candidates = await concreteSkillPaths(root, mount.path, mount.filesystem);
705
+ } catch (error) {
706
+ logger.debug("skills:root-expansion-failed", {
707
+ root,
708
+ error: errorUtils.errorMessage(error),
709
+ });
710
+ unresolved.push(root);
711
+ continue;
712
+ }
713
+ if (candidates.length === 0) {
714
+ unresolved.push(root);
715
+ continue;
716
+ }
717
+ for (const candidate of candidates) {
718
+ const name = posixPath.basename(candidate);
719
+ const existing = selected.get(name);
720
+ if (existing) {
721
+ logger.debug("skills:duplicate-skipped", { name, path: candidate, selected: existing });
722
+ } else {
723
+ selected.set(name, candidate);
724
+ }
725
+ }
726
+ }
727
+ return [...selected.values(), ...unresolved];
728
+ }
729
+
730
+ /** Find the most-specific mounted filesystem containing one skill path. */
731
+ function mountedFilesystem(
732
+ skillPath: string,
733
+ mounts: Readonly<Record<string, WorkspaceFilesystem>>,
734
+ ): { filesystem: WorkspaceFilesystem; path: string } | undefined {
735
+ return Object.entries(mounts)
736
+ .filter(([mount]) => skillPath === mount || skillPath.startsWith(`${mount}/`))
737
+ .sort(([left], [right]) => right.length - left.length)
738
+ .map(([path, filesystem]) => ({ filesystem, path }))[0];
739
+ }
740
+
741
+ /** Return concrete child skill directories from one mounted skill root. */
742
+ async function concreteSkillPaths(
743
+ root: string,
744
+ mount: string,
745
+ filesystem: WorkspaceFilesystem,
746
+ ): Promise<string[]> {
747
+ const relativeRoot = root === mount ? "." : root.slice(mount.length + 1);
748
+ if (await filesystem.exists(posixPath.join(relativeRoot, "SKILL.md"))) return [root];
749
+ const entries = await filesystem.readdir(relativeRoot);
750
+ const candidates: string[] = [];
751
+ for (const entry of entries) {
752
+ if (entry.type !== "directory") continue;
753
+ const relative = posixPath.join(relativeRoot, entry.name);
754
+ if (await filesystem.exists(posixPath.join(relative, "SKILL.md"))) {
755
+ candidates.push(posixPath.join(root, entry.name));
756
+ }
757
+ }
758
+ return candidates;
759
+ }