@dbx-tools/appkit-mastra 0.6.216 → 0.6.217

package/src/agents.ts

@@ -151,8 +151,9 @@ function deriveToolId(description: string): string {
  * });
  * ```
  *
- * Returns the definition unchanged - the wrapper exists only to anchor
- * type inference and to match the AppKit API surface.
+ * Adds the package's default workspace when the definition omits one. That
+ * workspace carries Databricks skill mounts and Databricks Sandbox command
+ * execution. An explicit workspace remains the caller's complete override.
  */
 export function createAgent<T extends MastraAgentDefinition>(def: T): T {
   if (def.workspace) return { ...def };
@@ -250,7 +251,7 @@ export type MastraPlugins = Record<string, MastraPluginToolkitProvider>;
 /** Function form of {@link MastraAgentDefinition.tools}. */
 export type MastraToolsFn = (plugins: MastraPlugins) => MastraTools | Promise<MastraTools>;
 
-/** Function form of {@link MastraAgentDefinition.workspace}. */
+/** Function form of {@link MastraAgentDefinition.workspace}; `undefined` disables it for this agent. */
 export type MastraAgentWorkspaceResolver = () => Workspace | undefined;
 
 /**
@@ -512,8 +513,15 @@ export async function buildAgents(opts: {
   for (const [id, def] of Object.entries(definitions)) {
     const tools = await resolveTools(def.tools, plugins, ambientTools);
     let workspace = resolveAgentWorkspace(def.workspace);
-    if (extraSkillPaths?.length && isDefaultWorkspace(workspace)) {
-      workspace = createWorkspace({ extraSkillPaths });
+    if (
+      (def.workspace === undefined && !workspace) ||
+      ((extraSkillPaths?.length || config.sandbox !== undefined) && isDefaultWorkspace(workspace))
+    ) {
+      workspace = createWorkspace({
+        extraSkillPaths,
+        sandbox:
+          config.sandbox === undefined || config.sandbox === true ? "databricks" : config.sandbox,
+      });
       markDefaultWorkspace(workspace);
     }
     const gated = approvalGatedToolIds(tools);

package/src/config.ts

@@ -26,6 +26,7 @@ import type { MastraAgentDefinition, MastraTools } from "./agents.ts";
 import type { GenieSpacesConfig } from "./genie.ts";
 import { IDENTITY_MODES, type MastraIdentityMode } from "./identity.ts";
 import type { RemoteSkillsOption } from "./remote-skills.ts";
+import type { DatabricksWorkspaceSandboxOptions } from "./sandbox.ts";
 
 /**
  * `RequestContext` key under which {@link MastraServer} stores the
@@ -198,6 +199,13 @@ export interface MastraPluginConfig extends BasePluginConfig {
    * `lakebase` plugin's pool; an object opens a dedicated store.
    */
   memory?: boolean | MastraMemoryConfig;
+  /**
+   * Sandbox for auto-created agent workspaces. Defaults to Databricks Sandbox.
+   * `true` selects Databricks with defaults, an object configures its lifecycle,
+   * and `false` disables workspace command execution. An agent with an explicit
+   * custom `workspace` keeps that workspace and its sandbox.
+   */
+  sandbox?: boolean | "databricks" | "monty" | DatabricksWorkspaceSandboxOptions;
   /**
    * Code-defined agents. Accepts three shapes for convenience:
    *
@@ -630,6 +638,11 @@ export const MASTRA_CONFIG_SCHEMA: ConfigSchema = {
       description:
         "PgVector store for Mastra semantic recall. `true` reuses the `lakebase` plugin's pool, an object opens a dedicated store. Auto-enabled when the `lakebase` plugin is registered.",
     },
+    sandbox: {
+      type: ["boolean", "string", "object"],
+      description:
+        'Command sandbox for auto-created agent workspaces. Defaults to Databricks Sandbox with Monty fallback; false disables command execution, "monty" selects Monty directly, and an object configures Databricks lifecycle/fallback settings.',
+    },
     defaultAgent: {
       type: "string",
       description:

package/src/model.ts

@@ -24,7 +24,7 @@
  */
 
 import { getExecutionContext } from "@databricks/appkit";
-import { classes, resolve } from "@dbx-tools/model";
+import { classes, invoke, resolve } from "@dbx-tools/model";
 import { functionModule, json, log, net } from "@dbx-tools/shared-core";
 import { model } from "@dbx-tools/shared-model";
 import type { MastraModelConfig } from "@mastra/core/llm";
@@ -32,7 +32,7 @@ import type { RequestContext } from "@mastra/core/request-context";
 
 import { MASTRA_USER_KEY, type MastraPluginConfig, type User } from "./config.ts";
 import {
-  rewriteServingBody,
+  rewriteServingRequest,
   rewriteServingResponseBody,
   rewriteServingResponseStream,
 } from "./serving-sanitize.ts";
@@ -131,16 +131,16 @@ export async function buildModel(
   };
 }
 
-/** Path prefix that identifies a Databricks Model Serving REST call. */
-const SERVING_ENDPOINTS_PATH_PREFIX = "/serving-endpoints/";
+/** Chat Completions route whose provider-specific wire shapes are normalized here. */
+const CHAT_COMPLETIONS_PATH = `/${invoke.CHAT_COMPLETIONS_PATH}`;
 
 /**
- * Install a single shared `globalThis.fetch` wrapper for every POST to
- * `/serving-endpoints/...`. The wrapper does two things:
+ * Install a single shared `globalThis.fetch` wrapper for Chat Completions
+ * requests. The wrapper does three things:
  *
- *   1. Rewrites the outgoing `messages` array to repair Mastra/AI SDK
- *      stream-replay quirks that Databricks-hosted Claude rejects (see
- *      {@link rewriteServingBody} in `./serving-sanitize.js`).
+ *   1. Rewrites outgoing JSON from either `init.body` or a `Request` to repair
+ *      provider-specific request constraints (see {@link rewriteServingRequest}
+ *      in `./serving-sanitize.js`).
  *   2. At `LOG_LEVEL=debug`, dumps the (post-sanitize) JSON body so
  *      4xx debugging doesn't have to fight AI SDK's `[Array]`
  *      formatter.
@@ -154,32 +154,30 @@ const SERVING_ENDPOINTS_PATH_PREFIX = "/serving-endpoints/";
  * step.
  */
 const setupFetchInterceptor = functionModule.memoize((): void => {
+  globalThis.fetch = createServingFetchInterceptor(globalThis.fetch.bind(globalThis));
+});
+
+/** Build the serving fetch wrapper; exported for transport-level regression tests. */
+export function createServingFetchInterceptor(original: typeof fetch): typeof fetch {
   const logger = log.logger("mastra/llm");
-  const original = globalThis.fetch.bind(globalThis);
-  globalThis.fetch = (async (input, init) => {
+  return (async (input, init) => {
     const url = net.urlBuilder(input);
-    if (
-      !url ||
-      !url.pathname.startsWith(SERVING_ENDPOINTS_PATH_PREFIX) ||
-      typeof init?.body !== "string"
-    ) {
+    const method = init?.method ?? (input instanceof Request ? input.method : "GET");
+    if (!url || url.pathname !== CHAT_COMPLETIONS_PATH || method.toUpperCase() !== "POST") {
       return original(input, init);
     }
-    const rewritten = rewriteServingBody(init.body);
-    if (rewritten !== init.body) {
-      init = { ...init, body: rewritten };
-    }
-    const parsed = json.parse<unknown>(rewritten);
+    const rewritten = await rewriteServingRequest(input, init);
+    const parsed = json.parse<unknown>(rewritten.body);
     logger.debug(
       "POST",
       parsed === undefined
         ? { url: url.toString(), bodyType: "non-JSON" }
         : { url: url.toString(), body: parsed },
     );
-    const response = await original(input, init);
+    const response = await original(rewritten.input, rewritten.init);
     return repairServingResponse(response);
   }) as typeof globalThis.fetch;
-});
+}
 
 /**
  * Rewrite a serving response whose body needs repair, leaving unsupported

package/src/monty-sandbox.ts

@@ -0,0 +1,406 @@
+/**
+ * Pydantic Monty fallback for Mastra command execution.
+ *
+ * Uses the Node subprocess-worker build of Monty. It executes Python source
+ * with no host filesystem, network, environment, shell, or third-party package
+ * access unless a future caller explicitly adds those capabilities.
+ *
+ * @module
+ */
+
+import { error, functionModule } from "@dbx-tools/shared-core";
+import type {
+  CommandResult,
+  ExecuteCommandOptions,
+  SandboxInfo,
+  WorkspaceSandbox,
+} from "@mastra/core/workspace";
+import type { CheckoutOptions, Monty as MontyPool, MontySession } from "@pydantic/monty/node";
+
+const DEFAULT_COMMAND_TIMEOUT_MS = 30_000;
+const DEFAULT_MAX_MEMORY_BYTES = 10_000_000;
+
+const montyModule = functionModule.memoize(() => import("@pydantic/monty/node"));
+const montyPool = functionModule.memoize(async (): Promise<MontyPool> => {
+  const { Monty } = await montyModule();
+  return Monty.create({
+    minProcesses: 1,
+    maxProcesses: 4,
+    requestTimeout: 35,
+  });
+});
+
+/** Resource controls for the Monty Python fallback. */
+export interface MontySandboxOptions {
+  /** Mastra provider id. */
+  id?: string;
+  /** Display name. */
+  name?: string;
+  /** Default execution timeout. Per-call `options.timeout` wins. */
+  commandTimeoutMs?: number;
+  /** Maximum Monty heap bytes for one checkout. */
+  maxMemoryBytes?: number;
+  /** Type-check each Python snippet before running it. Defaults to false. */
+  typeCheck?: boolean;
+}
+
+/** Python-only, deny-by-default sandbox backed by Monty's Node worker pool. */
+export class MontySandbox implements WorkspaceSandbox {
+  readonly id: string;
+  readonly name: string;
+  readonly provider = "monty";
+  status: WorkspaceSandbox["status"] = "pending";
+  error?: string;
+
+  private readonly commandTimeoutMs: number;
+  private readonly maxMemoryBytes: number;
+  private readonly typeCheck: boolean;
+  private readonly createdAt = new Date();
+  private lastUsedAt?: Date;
+
+  constructor(options: MontySandboxOptions = {}) {
+    this.id = options.id ?? "monty";
+    this.name = options.name ?? "Pydantic Monty";
+    this.commandTimeoutMs = options.commandTimeoutMs ?? DEFAULT_COMMAND_TIMEOUT_MS;
+    this.maxMemoryBytes = options.maxMemoryBytes ?? DEFAULT_MAX_MEMORY_BYTES;
+    this.typeCheck = options.typeCheck ?? false;
+  }
+
+  /** Warm the shared crash-isolated worker pool. */
+  async start(): Promise<void> {
+    if (this.status === "running") return;
+    this.status = "starting";
+    this.error = undefined;
+    try {
+      await montyPool();
+      this.status = "running";
+    } catch (caught) {
+      this.status = "error";
+      this.error = error.errorMessage(caught);
+      throw caught;
+    }
+  }
+
+  /** The process-wide pool stays warm; this instance carries no open session. */
+  async stop(): Promise<void> {
+    this.status = "stopped";
+  }
+
+  /** The process-wide pool stays warm; each command closes its own checkout. */
+  async destroy(): Promise<void> {
+    this.status = "destroyed";
+  }
+
+  /** Monty sessions are ephemeral and do not provide persistent checkpoints. */
+  async snapshot(): Promise<void> {}
+
+  readonly supportsCheckpoints = false;
+
+  /** Execute Python source and map Monty output to Mastra's command result. */
+  async executeCommand(
+    command: string,
+    args: string[] = [],
+    options: ExecuteCommandOptions = {},
+  ): Promise<CommandResult> {
+    const started = performance.now();
+    const code = montyCode(command, args);
+    if (code === undefined) {
+      return failedResult(
+        command,
+        args,
+        started,
+        "Monty fallback accepts Python source directly, or python/python3 with a single -c script.",
+      );
+    }
+    const environment = Object.entries(options.env ?? {}).filter(
+      ([, value]) => value !== undefined,
+    );
+    if (environment.length > 0) {
+      return failedResult(
+        command,
+        args,
+        started,
+        "Monty fallback does not expose host environment variables.",
+      );
+    }
+    if (options.cwd) {
+      return failedResult(
+        command,
+        args,
+        started,
+        "Monty fallback does not expose a host working directory.",
+      );
+    }
+    if (options.abortSignal?.aborted) throw options.abortSignal.reason;
+
+    await this.start();
+    const timeout = options.timeout ?? this.commandTimeoutMs;
+    const pool = await montyPool();
+    const session = await checkoutSession(
+      pool,
+      {
+        limits: {
+          maxDurationSecs: Math.max(0.001, timeout / 1_000),
+          maxMemory: this.maxMemoryBytes,
+        },
+        typeCheck: this.typeCheck,
+      },
+      options.abortSignal,
+    );
+    const output = new RetainedOutput(options);
+    const workerPid = session.workerPid;
+    let aborted = false;
+    const onAbort = (): void => {
+      aborted = true;
+      terminateWorker(workerPid);
+    };
+    options.abortSignal?.addEventListener("abort", onAbort, { once: true });
+    if (options.abortSignal?.aborted) onAbort();
+    try {
+      if (aborted) throw abortReason(options.abortSignal);
+      const value = await session.feedRun(stripPythonFence(code), {
+        printCallback: (stream, text) => output.emit(stream, text),
+      });
+      const returned = formatResult(value);
+      if (returned) {
+        const suffix = returned.endsWith("\n") ? returned : `${returned}\n`;
+        output.emit("stdout", suffix);
+      }
+      this.lastUsedAt = new Date();
+      return {
+        command,
+        args,
+        success: true,
+        exitCode: 0,
+        stdout: output.stdout.value,
+        stderr: output.stderr.value,
+        executionTimeMs: performance.now() - started,
+        ...output.metadata(),
+      };
+    } catch (caught) {
+      if (aborted || options.abortSignal?.aborted) {
+        throw abortReason(options.abortSignal);
+      }
+      const message = error.errorMessage(caught);
+      output.emit("stderr", output.stderr.value ? `\n${message}` : message);
+      const { MontyCrashedError } = await montyModule();
+      const timedOut = caught instanceof MontyCrashedError && caught.timedOut;
+      return {
+        command,
+        args,
+        success: false,
+        exitCode: timedOut ? -1 : 1,
+        stdout: output.stdout.value,
+        stderr: output.stderr.value,
+        executionTimeMs: performance.now() - started,
+        ...(timedOut ? { timedOut: true, killed: true } : {}),
+        ...output.metadata(),
+      };
+    } finally {
+      options.abortSignal?.removeEventListener("abort", onAbort);
+      await session.close();
+    }
+  }
+
+  async isReady(): Promise<boolean> {
+    return this.status === "running";
+  }
+
+  getInfo(): SandboxInfo {
+    return {
+      id: this.id,
+      name: this.name,
+      provider: this.provider,
+      status: this.status,
+      createdAt: this.createdAt,
+      ...(this.lastUsedAt ? { lastUsedAt: this.lastUsedAt } : {}),
+      metadata: {
+        language: "python",
+        hostAccess: false,
+      },
+    };
+  }
+
+  getInstructions(): string {
+    return [
+      "Commands run as Python source in the Pydantic Monty fallback.",
+      "Pass Python code directly rather than shell syntax.",
+      "There is no host filesystem, network, environment, shell, or third-party package access.",
+    ].join(" ");
+  }
+}
+
+function montyCode(command: string, args: readonly string[]): string | undefined {
+  if (args.length === 0) return command;
+  if ((command === "python" || command === "python3") && args.length === 2 && args[0] === "-c") {
+    return args[1];
+  }
+  return undefined;
+}
+
+function stripPythonFence(code: string): string {
+  const trimmed = code.trim();
+  const match = /^```(?:python|py)?\s*\n([\s\S]*?)\n```$/i.exec(trimmed);
+  return match?.[1] ?? code;
+}
+
+function formatResult(value: unknown): string {
+  if (value === undefined || value === null) return "";
+  if (typeof value === "string") return value;
+  try {
+    return JSON.stringify(value, null, 2) ?? String(value);
+  } catch {
+    return String(value);
+  }
+}
+
+class RetainedOutput {
+  readonly stdout: RetainedText;
+  readonly stderr: RetainedText;
+
+  private readonly onStdout: ExecuteCommandOptions["onStdout"];
+  private readonly onStderr: ExecuteCommandOptions["onStderr"];
+
+  constructor(options: ExecuteCommandOptions) {
+    this.stdout = new RetainedText(options.maxRetainedBytes);
+    this.stderr = new RetainedText(options.maxRetainedBytes);
+    this.onStdout = options.onStdout;
+    this.onStderr = options.onStderr;
+  }
+
+  emit(stream: "stdout" | "stderr", text: string): void {
+    if (stream === "stdout") {
+      this.onStdout?.(text);
+      this.stdout.append(text);
+    } else {
+      this.onStderr?.(text);
+      this.stderr.append(text);
+    }
+  }
+
+  metadata(): Pick<
+    CommandResult,
+    "stdoutTruncated" | "stderrTruncated" | "stdoutDroppedBytes" | "stderrDroppedBytes"
+  > {
+    return {
+      ...(this.stdout.droppedBytes > 0
+        ? {
+            stdoutTruncated: true,
+            stdoutDroppedBytes: this.stdout.droppedBytes,
+          }
+        : {}),
+      ...(this.stderr.droppedBytes > 0
+        ? {
+            stderrTruncated: true,
+            stderrDroppedBytes: this.stderr.droppedBytes,
+          }
+        : {}),
+    };
+  }
+}
+
+class RetainedText {
+  value = "";
+  droppedBytes = 0;
+
+  private readonly maxBytes: number;
+  private readonly encoder = new TextEncoder();
+  private readonly decoder = new TextDecoder();
+
+  constructor(maxBytes: number | undefined) {
+    const resolved = maxBytes ?? Number.POSITIVE_INFINITY;
+    if (
+      resolved !== Number.POSITIVE_INFINITY &&
+      (!Number.isSafeInteger(resolved) || resolved < 0)
+    ) {
+      throw new TypeError("maxRetainedBytes must be a non-negative safe integer or Infinity");
+    }
+    this.maxBytes = resolved;
+  }
+
+  append(text: string): void {
+    if (!text) return;
+    if (this.maxBytes === Number.POSITIVE_INFINITY) {
+      this.value += text;
+      return;
+    }
+    const combined = this.encoder.encode(this.value + text);
+    if (combined.length <= this.maxBytes) {
+      this.value += text;
+      return;
+    }
+    let start = combined.length - this.maxBytes;
+    while (start < combined.length && (combined[start]! & 0xc0) === 0x80) start++;
+    const retained = combined.subarray(start);
+    this.droppedBytes += combined.length - retained.length;
+    this.value = this.decoder.decode(retained);
+  }
+}
+
+function terminateWorker(workerPid: number | undefined): void {
+  if (workerPid === undefined) return;
+  try {
+    process.kill(workerPid, "SIGKILL");
+  } catch {
+    // The command can finish between the abort event and the kill request.
+  }
+}
+
+async function checkoutSession(
+  pool: MontyPool,
+  options: CheckoutOptions,
+  signal: AbortSignal | undefined,
+): Promise<MontySession> {
+  const checkout = pool.checkout(options);
+  if (!signal) return checkout;
+  if (signal.aborted) {
+    void checkout.then((session) => session.close()).catch(() => undefined);
+    throw abortReason(signal);
+  }
+
+  return new Promise<MontySession>((resolve, reject) => {
+    let settled = false;
+    const onAbort = (): void => {
+      if (settled) return;
+      settled = true;
+      reject(abortReason(signal));
+      void checkout.then((session) => session.close()).catch(() => undefined);
+    };
+    signal.addEventListener("abort", onAbort, { once: true });
+    checkout.then(
+      (session) => {
+        if (settled) return;
+        settled = true;
+        signal.removeEventListener("abort", onAbort);
+        resolve(session);
+      },
+      (caught: unknown) => {
+        if (settled) return;
+        settled = true;
+        signal.removeEventListener("abort", onAbort);
+        reject(caught);
+      },
+    );
+  });
+}
+
+function abortReason(signal: AbortSignal | undefined): unknown {
+  return signal?.reason ?? new DOMException("The command was aborted", "AbortError");
+}
+
+function failedResult(
+  command: string,
+  args: string[],
+  started: number,
+  stderr: string,
+): CommandResult {
+  return {
+    command,
+    args,
+    success: false,
+    exitCode: 1,
+    stdout: "",
+    stderr,
+    executionTimeMs: performance.now() - started,
+  };
+}