feat: add run_sandbox_command and get_task_run_status MCP tools

CLAUDE.md

@@ -266,6 +266,8 @@ const { accountId, orgId, authToken } = authResult;
 
 ### MCP Tools
 
+**CRITICAL: Never manually extract `accountId` from `extra.authInfo` (e.g. `authInfo?.extra?.accountId`).** Always use `resolveAccountId()` — it handles validation, org-key overrides, and access control in one place.
+
 ```typescript
 import { resolveAccountId } from "@/lib/mcp/resolveAccountId";
 import type { McpAuthInfo } from "@/lib/mcp/verifyApiKey";
@@ -275,6 +277,14 @@ const { accountId, error } = await resolveAccountId({
   authInfo,
   accountIdOverride: undefined,
 });
+
+if (error) {
+  return getToolResultError(error);
+}
+
+if (!accountId) {
+  return getToolResultError("Failed to resolve account ID");
+}
 ```
 
 This ensures:

lib/mcp/tools/index.ts

@@ -19,6 +19,7 @@ import { registerSendEmailTool } from "./registerSendEmailTool";
 import { registerAllArtistTools } from "./artists";
 import { registerAllChatsTools } from "./chats";
 import { registerAllPulseTools } from "./pulse";
+import { registerAllSandboxTools } from "./sandbox";
 
 /**
  * Registers all MCP tools on the server.
@@ -37,6 +38,7 @@ export const registerAllTools = (server: McpServer): void => {
   registerAllFileTools(server);
   registerAllImageTools(server);
   registerAllPulseTools(server);
+  registerAllSandboxTools(server);
   registerAllSearchTools(server);
   registerAllSora2Tools(server);
   registerAllSpotifyTools(server);

lib/mcp/tools/sandbox/__tests__/registerRunSandboxCommandTool.test.ts

@@ -0,0 +1,199 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import type { RequestHandlerExtra } from "@modelcontextprotocol/sdk/shared/protocol.js";
+import type { ServerRequest, ServerNotification } from "@modelcontextprotocol/sdk/types.js";
+
+import { registerRunSandboxCommandTool } from "../registerRunSandboxCommandTool";
+
+const mockProcessCreateSandbox = vi.fn();
+const mockResolveAccountId = vi.fn();
+
+vi.mock("@/lib/sandbox/processCreateSandbox", () => ({
+  processCreateSandbox: (...args: unknown[]) => mockProcessCreateSandbox(...args),
+}));
+
+vi.mock("@/lib/mcp/resolveAccountId", () => ({
+  resolveAccountId: (...args: unknown[]) => mockResolveAccountId(...args),
+}));
+
+type ServerRequestHandlerExtra = RequestHandlerExtra<ServerRequest, ServerNotification>;
+
+/**
+ * Creates a mock extra object with optional authInfo.
+ *
+ * @param authInfo
+ * @param authInfo.accountId
+ * @param authInfo.orgId
+ */
+function createMockExtra(authInfo?: {
+  accountId?: string;
+  orgId?: string | null;
+}): ServerRequestHandlerExtra {
+  return {
+    authInfo: authInfo
+      ? {
+          token: "test-token",
+          scopes: ["mcp:tools"],
+          clientId: authInfo.accountId,
+          extra: {
+            accountId: authInfo.accountId,
+            orgId: authInfo.orgId ?? null,
+          },
+        }
+      : undefined,
+  } as unknown as ServerRequestHandlerExtra;
+}
+
+describe("registerRunSandboxCommandTool", () => {
+  let mockServer: McpServer;
+  let registeredHandler: (args: unknown, extra: ServerRequestHandlerExtra) => Promise<unknown>;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+
+    mockServer = {
+      registerTool: vi.fn((name, config, handler) => {
+        registeredHandler = handler;
+      }),
+    } as unknown as McpServer;
+
+    registerRunSandboxCommandTool(mockServer);
+  });
+
+  it("registers the run_sandbox_command tool", () => {
+    expect(mockServer.registerTool).toHaveBeenCalledWith(
+      "run_sandbox_command",
+      expect.objectContaining({
+        description: expect.any(String),
+      }),
+      expect.any(Function),
+    );
+  });
+
+  it("returns error when resolveAccountId returns an error", async () => {
+    mockResolveAccountId.mockResolvedValue({
+      accountId: null,
+      error: "Authentication required. Provide an API key via Authorization: Bearer header, or provide account_id from the system prompt context.",
+    });
+
+    const result = await registeredHandler(
+      { command: "ls" },
+      createMockExtra(),
+    );
+
+    expect(result).toEqual({
+      content: [
+        {
+          type: "text",
+          text: expect.stringContaining("Authentication required"),
+        },
+      ],
+    });
+  });
+
+  it("returns error when resolveAccountId returns null accountId without error", async () => {
+    mockResolveAccountId.mockResolvedValue({
+      accountId: null,
+      error: null,
+    });
+
+    const result = await registeredHandler(
+      { command: "ls" },
+      createMockExtra(),
+    );
+
+    expect(result).toEqual({
+      content: [
+        {
+          type: "text",
+          text: expect.stringContaining("Failed to resolve account ID"),
+        },
+      ],
+    });
+  });
+
+  it("calls processCreateSandbox with command and returns success", async () => {
+    mockResolveAccountId.mockResolvedValue({
+      accountId: "acc_123",
+      error: null,
+    });
+    mockProcessCreateSandbox.mockResolvedValue({
+      sandboxId: "sbx_123",
+      sandboxStatus: "running",
+      timeout: 600000,
+      createdAt: "2024-01-01T00:00:00.000Z",
+      runId: "run_abc123",
+    });
+
+    const result = await registeredHandler(
+      { command: "npm install", args: ["express"], cwd: "/app" },
+      createMockExtra({ accountId: "acc_123" }),
+    );
+
+    expect(mockProcessCreateSandbox).toHaveBeenCalledWith({
+      accountId: "acc_123",
+      command: "npm install",
+      args: ["express"],
+      cwd: "/app",
+    });
+    expect(result).toEqual({
+      content: [
+        {
+          type: "text",
+          text: expect.stringContaining('"sandboxId":"sbx_123"'),
+        },
+      ],
+    });
+    expect(result).toEqual({
+      content: [
+        {
+          type: "text",
+          text: expect.stringContaining('"runId":"run_abc123"'),
+        },
+      ],
+    });
+  });
+
+  it("passes authInfo to resolveAccountId", async () => {
+    mockResolveAccountId.mockResolvedValue({
+      accountId: "acc_123",
+      error: null,
+    });
+    mockProcessCreateSandbox.mockResolvedValue({
+      sandboxId: "sbx_123",
+      sandboxStatus: "running",
+      timeout: 600000,
+      createdAt: "2024-01-01T00:00:00.000Z",
+    });
+
+    const extra = createMockExtra({ accountId: "acc_123" });
+    await registeredHandler({ command: "ls" }, extra);
+
+    expect(mockResolveAccountId).toHaveBeenCalledWith({
+      authInfo: extra.authInfo,
+      accountIdOverride: undefined,
+    });
+  });
+
+  it("returns error when processCreateSandbox throws", async () => {
+    mockResolveAccountId.mockResolvedValue({
+      accountId: "acc_123",
+      error: null,
+    });
+    mockProcessCreateSandbox.mockRejectedValue(new Error("Sandbox creation failed"));
+
+    const result = await registeredHandler(
+      { command: "ls" },
+      createMockExtra({ accountId: "acc_123" }),
+    );
+
+    expect(result).toEqual({
+      content: [
+        {
+          type: "text",
+          text: expect.stringContaining("Sandbox creation failed"),
+        },
+      ],
+    });
+  });
+});

lib/mcp/tools/sandbox/index.ts

@@ -0,0 +1,11 @@
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import { registerRunSandboxCommandTool } from "./registerRunSandboxCommandTool";
+
+/**
+ * Registers all sandbox-related MCP tools on the server.
+ *
+ * @param server - The MCP server instance to register tools on.
+ */
+export const registerAllSandboxTools = (server: McpServer): void => {
+  registerRunSandboxCommandTool(server);
+};

lib/mcp/tools/sandbox/registerRunSandboxCommandTool.ts

@@ -0,0 +1,61 @@
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import type { RequestHandlerExtra } from "@modelcontextprotocol/sdk/shared/protocol.js";
+import type { ServerRequest, ServerNotification } from "@modelcontextprotocol/sdk/types.js";
+import { z } from "zod";
+import type { McpAuthInfo } from "@/lib/mcp/verifyApiKey";
+import { resolveAccountId } from "@/lib/mcp/resolveAccountId";
+import { getToolResultSuccess } from "@/lib/mcp/getToolResultSuccess";
+import { getToolResultError } from "@/lib/mcp/getToolResultError";
+import { processCreateSandbox } from "@/lib/sandbox/processCreateSandbox";
+
+const runSandboxCommandSchema = z.object({
+  command: z.string().describe("The command to run in the sandbox."),
+  args: z.array(z.string()).optional().describe("Arguments for the command."),
+  cwd: z.string().optional().describe("Working directory for the command."),
+});
+
+/**
+ * Registers the "run_sandbox_command" tool on the MCP server.
+ * Creates a sandbox and runs a command in it.
+ *
+ * @param server - The MCP server instance to register the tool on.
+ */
+export function registerRunSandboxCommandTool(server: McpServer): void {
+  server.registerTool(
+    "run_sandbox_command",
+    {
+      description:
+        "Create a sandbox and run a command in it. Returns the sandbox ID and a run ID to track progress.",
+      inputSchema: runSandboxCommandSchema,
+    },
+    async (args, extra: RequestHandlerExtra<ServerRequest, ServerNotification>) => {
+      const authInfo = extra.authInfo as McpAuthInfo | undefined;
+      const { accountId, error } = await resolveAccountId({
+        authInfo,
+        accountIdOverride: undefined,
+      });
+
+      if (error) {
+        return getToolResultError(error);
+      }
+
+      if (!accountId) {
+        return getToolResultError("Failed to resolve account ID");
+      }
+
+      try {
+        const result = await processCreateSandbox({
+          accountId,
+          command: args.command,
+          args: args.args,
+          cwd: args.cwd,
+        });
+
+        return getToolResultSuccess(result);
+      } catch (error) {
+        const message = error instanceof Error ? error.message : "Failed to create sandbox";
+        return getToolResultError(message);
+      }
+    },
+  );
+}