@dbx-tools/appkit-mastra 0.9.153 → 0.9.157

package/src/mlflow.ts

@@ -26,9 +26,13 @@
  */
 
 import { ConfigurationError, createWorkspaceClient } from "@databricks/appkit";
-import { appkit } from "@dbx-tools/appkit";
-import { asyncUtils, environmentUtils, errorUtils, log, object } from "@dbx-tools/shared-core";
-import { feedback } from "@dbx-tools/shared-mastra";
+import { type WorkspaceClientLike } from "@dbx-tools/appkit";
+import { log, sleep, isDatabricksAppEnv, toError, isRecord } from "@dbx-tools/shared-core";
+import {
+  type MlflowFeedbackRequest,
+  DEFAULT_FEEDBACK_NAME,
+  DEFAULT_COMMENT_NAME,
+} from "@dbx-tools/shared-mastra";
 import { TraceLocationType, TraceMetadataKey } from "@mlflow/core";
 import { TraceInfo } from "@mlflow/core/dist/core/entities/trace_info.js";
 import { TraceState } from "@mlflow/core/dist/core/entities/trace_state.js";
@@ -118,7 +122,7 @@ export class AgentTraceSpanProcessor implements SpanProcessor {
 }
 
 /** Workspace client carried on an AppKit execution context. */
-type WorkspaceClient = appkit.WorkspaceClientLike;
+type WorkspaceClient = WorkspaceClientLike;
 
 /** Experiment tag Databricks stores as `catalog.schema.table_prefix`. */
 const TRACE_DESTINATION_PATH_TAG = "mlflow.experiment.databricksTraceDestinationPath";
@@ -157,14 +161,14 @@ function experimentTagMap(raw: unknown): Map<string, string> {
   const tags = new Map<string, string>();
   if (Array.isArray(raw)) {
     for (const tag of raw) {
-      if (!object.isRecord(tag) || typeof tag.key !== "string" || typeof tag.value !== "string") {
+      if (!isRecord(tag) || typeof tag.key !== "string" || typeof tag.value !== "string") {
         continue;
       }
       tags.set(tag.key, tag.value);
     }
     return tags;
   }
-  if (!object.isRecord(raw)) return tags;
+  if (!isRecord(raw)) return tags;
   for (const [key, value] of Object.entries(raw)) {
     if (typeof value === "string") tags.set(key, value);
   }
@@ -327,7 +331,7 @@ function configuredExperiment(): { id?: string; name?: string } {
 
 /** Tracking URI for direct local MLflow tracing, with CLI-profile auth injected. */
 export function directMlflowTrackingUri(): string | undefined {
-  if (environmentUtils.isDatabricksAppEnv()) return undefined;
+  if (isDatabricksAppEnv()) return undefined;
   const configured = process.env.MLFLOW_TRACKING_URI?.trim();
   if (configured && configured !== "databricks") return configured;
   const profile = process.env.DATABRICKS_CONFIG_PROFILE?.trim();
@@ -360,16 +364,14 @@ export function directMlflowTraceLocation():
 /** Whether the local direct-to-experiment MLflow SDK has enough configuration. */
 export function directMlflowTracingConfigured(): boolean {
   return Boolean(
-    !environmentUtils.isDatabricksAppEnv() &&
-    process.env.MLFLOW_EXPERIMENT_ID?.trim() &&
-    directMlflowTrackingUri(),
+    !isDatabricksAppEnv() && process.env.MLFLOW_EXPERIMENT_ID?.trim() && directMlflowTrackingUri(),
   );
 }
 
 /** Whether a Databricks App can promote OTLP spans into full MLflow trace info. */
 export function appMlflowTraceInfoConfigured(): boolean {
   return Boolean(
-    environmentUtils.isDatabricksAppEnv() &&
+    isDatabricksAppEnv() &&
     (process.env.MLFLOW_EXPERIMENT_ID?.trim() || process.env.MLFLOW_EXPERIMENT_NAME?.trim()) &&
     (process.env.OTEL_EXPORTER_OTLP_ENDPOINT?.trim() ||
       process.env.OTEL_EXPORTER_OTLP_TRACES_ENDPOINT?.trim()),
@@ -395,7 +397,7 @@ export async function initializeAppMlflowTraceInfo(client?: WorkspaceClient): Pr
     return true;
   } catch (err) {
     logger.warn("Databricks Apps MLflow trace-info promotion disabled", {
-      error: errorUtils.toError(err),
+      error: toError(err),
     });
     // Startup uses the service principal. Permit one later request-scoped
     // client to resolve experiment metadata when that principal cannot read it.
@@ -468,7 +470,7 @@ export function persistAppMlflowTraceInfo(params: AppMlflowTraceInfo): void {
     .catch((err) => {
       logger.warn("Databricks Apps MLflow trace-info promotion failed", {
         traceId: params.traceId,
-        error: errorUtils.toError(err),
+        error: toError(err),
       });
     })
     .finally(() => pendingAppMlflowTraceInfo.delete(pending));
@@ -511,7 +513,7 @@ async function detectedDirectMlflowTraceLocation(
   } catch (err) {
     logger.warn("direct MLflow UC trace-location detection failed", {
       experimentId,
-      error: errorUtils.toError(err),
+      error: toError(err),
     });
     return undefined;
   }
@@ -575,7 +577,7 @@ export async function initializeDirectMlflowTracing(): Promise<boolean> {
     return true;
   } catch (err) {
     logger.warn("direct MLflow tracing disabled", {
-      error: errorUtils.toError(err),
+      error: toError(err),
     });
     return false;
   }
@@ -712,7 +714,7 @@ async function resolveExperimentId(client: WorkspaceClient): Promise<string | un
 }
 
 /** Parameters for {@link logFeedback}. */
-export type LogFeedbackParams = feedback.MlflowFeedbackRequest & {
+export type LogFeedbackParams = MlflowFeedbackRequest & {
   /** Identity the feedback is attributed to (user email / id). */
   sourceId?: string;
 };
@@ -731,9 +733,7 @@ export async function logFeedback(
   // A comment with no thumbs value is logged as text feedback; a value
   // (with an optional comment as the rationale) is the thumbs path.
   const hasValue = params.value !== undefined;
-  const name =
-    params.name?.trim() ||
-    (hasValue ? feedback.DEFAULT_FEEDBACK_NAME : feedback.DEFAULT_COMMENT_NAME);
+  const name = params.name?.trim() || (hasValue ? DEFAULT_FEEDBACK_NAME : DEFAULT_COMMENT_NAME);
   const value = hasValue ? params.value : params.comment;
   const assessmentTraceId = mlflowAssessmentTraceId(
     params.traceId,
@@ -762,7 +762,7 @@ export async function logFeedback(
       logger.warn("feedback request failed", {
         traceId: params.traceId,
         assessmentTraceId,
-        error: errorUtils.toError(err),
+        error: toError(err),
       });
       return undefined;
     }
@@ -776,7 +776,7 @@ export async function logFeedback(
     // Trace export is async; a fresh trace may not exist yet. Retry a
     // few times with a short backoff before giving up softly.
     if (res.status === 404 && attempt < NOT_FOUND_RETRIES) {
-      await asyncUtils.sleep(NOT_FOUND_BACKOFF_MS * (attempt + 1));
+      await sleep(NOT_FOUND_BACKOFF_MS * (attempt + 1));
       continue;
     }
     logger.warn("feedback not recorded", {

package/src/model.ts

@@ -23,14 +23,14 @@
  */
 
 import { createWorkspaceClient, type WorkspaceClient } from "@databricks/appkit";
-import { appkit } from "@dbx-tools/appkit";
-import { classes, policy, resolve } from "@dbx-tools/model";
+import { tryGetExecutionContext } from "@dbx-tools/appkit";
+import { classes, policy, resolve, selection } from "@dbx-tools/model";
 import {
   createDatabricksLanguageModel,
   type DatabricksLanguageModel,
-} from "@dbx-tools/model-protocol/provider";
+} from "@dbx-tools/model-protocol";
 import { log } from "@dbx-tools/shared-core";
-import { model, type ServingEndpointSummary } from "@dbx-tools/shared-model";
+import { contracts, type ServingEndpointSummary } from "@dbx-tools/shared-model";
 import { RequestContext } from "@mastra/core/request-context";
 
 import { MASTRA_USER_KEY, type MastraPluginConfig, type User } from "./config.ts";
@@ -41,7 +41,7 @@ import {
 } from "./serving.ts";
 import { recordActiveTraceAuth, recordActiveTraceModel } from "./telemetry.ts";
 
-type ModelClass = model.ModelClass;
+type ModelClass = contracts.ModelClass;
 const { parseModelClass } = classes;
 const { selectModel } = resolve;
 
@@ -89,7 +89,7 @@ function resolveModelWorkspaceContext(
   requestContext?: RequestContext,
 ): ResolvedModelWorkspaceClient {
   const user = requestContext?.get(MASTRA_USER_KEY) as User | undefined;
-  const executionContext = user?.executionContext ?? appkit.tryGetExecutionContext();
+  const executionContext = user?.executionContext ?? tryGetExecutionContext();
   if (executionContext?.client.config) {
     return {
       client: executionContext.client,
@@ -137,7 +137,7 @@ export function resolveDefaultModelId(
   overrides: BuildModelOverrides = {},
 ): string {
   const requested = overrides.modelId ?? process.env.DATABRICKS_SERVING_ENDPOINT_NAME;
-  return resolve.resolveModel(endpoints, selectionInput(config, requested, overrides.modelClass))
+  return selection.resolveModel(endpoints, selectionInput(config, requested, overrides.modelClass))
     .modelId;
 }
 

package/src/monty-sandbox.ts

@@ -10,8 +10,8 @@
  */
 
 import { availableParallelism } from "node:os";
-import { configUtils } from "@dbx-tools/core";
-import { errorUtils, functionUtils, log } from "@dbx-tools/shared-core";
+import { boolean } from "@dbx-tools/core";
+import { log, toError, errorMessage, memoize } from "@dbx-tools/shared-core";
 import {
   MastraSandbox,
   type CommandResult,
@@ -45,9 +45,9 @@ export function montyMaxProcesses(parallelism: number = availableParallelism()):
   return Math.max(4, Math.floor(parallelism) * 2);
 }
 
-const montyContext = functionUtils.memoize(async (): Promise<MontyContext> => {
+const montyContext = memoize(async (): Promise<MontyContext> => {
   const errors: Error[] = [];
-  const targets = configUtils.boolean(undefined, "MONTY_FORCE_WASM")
+  const targets = boolean(undefined, "MONTY_FORCE_WASM")
     ? [WASM_MONTY_TARGET]
     : [NATIVE_MONTY_TARGET, WASM_MONTY_TARGET];
   for (const target of targets) {
@@ -66,9 +66,9 @@ const montyContext = functionUtils.memoize(async (): Promise<MontyContext> => {
         pool,
       };
     } catch (error) {
-      errors.push(errorUtils.toError(error));
+      errors.push(toError(error));
       logger.debug(`${target.name} unavailable`, {
-        error: errorUtils.toError(error),
+        error: toError(error),
       });
     }
   }
@@ -208,7 +208,7 @@ export class MontySandbox extends MastraSandbox {
       if (aborted || options.abortSignal?.aborted) {
         throw abortReason(options.abortSignal);
       }
-      const message = errorUtils.errorMessage(caught);
+      const message = errorMessage(caught);
       output.emit("stderr", output.stderr.value ? `\n${message}` : message);
       const { crashedError } = await montyContext();
       const timedOut = caught instanceof crashedError && caught.timedOut;

package/src/observability.ts

@@ -17,10 +17,11 @@
  *    ambient OTel context, so when Mastra is invoked from inside an
  *    AppKit HTTP span the trace stays connected.
  *
- * Net effect: Mastra spans get exactly the treatment AppKit's
- * `agents` plugin gets. No custom OTLP pipeline lives in this
- * package; the OTLP endpoint, headers, and resource attributes are
- * driven by the standard OTel env vars
+ * Net effect: deployed Mastra spans get exactly the treatment AppKit's
+ * `agents` plugin gets. Local direct MLflow tracing uses its own global
+ * provider because AppKit has no OTLP processor to start there. No custom
+ * OTLP pipeline lives in this package; the OTLP endpoint, headers, and resource
+ * attributes are driven by the standard OTel env vars
  * (`OTEL_EXPORTER_OTLP_ENDPOINT`, `OTEL_EXPORTER_OTLP_HEADERS`,
  * `OTEL_SERVICE_NAME`, `OTEL_RESOURCE_ATTRIBUTES`, ...) and consumed
  * by AppKit's `TelemetryManager`. Set those once and both AppKit and
@@ -92,7 +93,7 @@ export interface BuildObservabilityOptions {
    * Supports dot notation for nested values per the Mastra docs.
    */
   requestContextKeys?: readonly string[];
-  /** AppKit-owned tracer provider used by the Mastra OTel bridge. */
+  /** AppKit-owned provider used unless local direct MLflow tracing owns the global provider. */
   tracerProvider?: TracerProvider;
 }
 
@@ -140,6 +141,7 @@ export async function buildObservability(
   const directMlflowConfigured =
     options?.enabled === false || otlpConfigured ? false : await initializeDirectMlflowTracing();
   if (directMlflowConfigured) configureOtelPropagation();
+  const bridgeTracerProvider = directMlflowConfigured ? undefined : options?.tracerProvider;
   const tracingConfigured = otlpConfigured || directMlflowConfigured;
 
   if (options?.enabled === false || (options?.enabled !== true && !tracingConfigured)) {
@@ -190,7 +192,7 @@ export async function buildObservability(
       serviceName: {
         serviceName,
         bridge: new OtelBridge(
-          options?.tracerProvider ? { tracerProvider: options.tracerProvider } : undefined,
+          bridgeTracerProvider ? { tracerProvider: bridgeTracerProvider } : undefined,
         ),
         requestContextKeys,
       },