🤖 fix: truncate oversized MCP images to prevent context overflow

src/cli/run.test.ts

@@ -169,6 +169,65 @@ describe("mux CLI", () => {
       expect(result.exitCode).toBe(1);
       expect(result.output.length).toBeGreaterThan(0);
     });
+
+    test("--help shows --mcp option", async () => {
+      const result = await runCli(["run", "--help"]);
+      expect(result.exitCode).toBe(0);
+      expect(result.stdout).toContain("--mcp");
+      expect(result.stdout).toContain("name=command");
+      expect(result.stdout).toContain("--no-mcp-config");
+    });
+
+    test("--mcp without = shows error", async () => {
+      const result = await runRunDirect(["--mcp", "invalid-format", "test message"]);
+      expect(result.exitCode).toBe(1);
+      expect(result.output).toContain("Invalid --mcp format");
+      expect(result.output).toContain("Expected: name=command");
+    });
+
+    test("--mcp with empty name shows error", async () => {
+      const result = await runRunDirect(["--mcp", "=some-command", "test message"]);
+      expect(result.exitCode).toBe(1);
+      expect(result.output).toContain("Server name is required");
+    });
+
+    test("--mcp with empty command shows error", async () => {
+      const result = await runRunDirect(["--mcp", "myserver=", "test message"]);
+      expect(result.exitCode).toBe(1);
+      expect(result.output).toContain("Command is required");
+    });
+
+    test("--mcp accepts valid name=command format", async () => {
+      // Test with a nonexistent directory to ensure parsing succeeds before failing
+      const result = await runRunDirect([
+        "--dir",
+        "/nonexistent/path/for/mcp/test",
+        "--mcp",
+        "memory=npx -y @modelcontextprotocol/server-memory",
+        "test message",
+      ]);
+      // Should not fail with "Invalid --mcp format" - will fail on directory instead
+      expect(result.output).not.toContain("Invalid --mcp format");
+      // Verify it got past argument parsing to directory validation
+      expect(result.exitCode).toBe(1);
+    });
+
+    test("--mcp can be repeated multiple times", async () => {
+      // Test with a nonexistent directory to ensure parsing succeeds before failing
+      const result = await runRunDirect([
+        "--dir",
+        "/nonexistent/path/for/mcp/test",
+        "--mcp",
+        "server1=command1",
+        "--mcp",
+        "server2=command2 with args",
+        "test message",
+      ]);
+      // Should not fail with "Invalid --mcp format"
+      expect(result.output).not.toContain("Invalid --mcp format");
+      // Verify it got past argument parsing to directory validation
+      expect(result.exitCode).toBe(1);
+    });
   });
 
   describe("mux server", () => {

src/cli/run.ts

@@ -18,6 +18,8 @@ import { InitStateManager } from "@/node/services/initStateManager";
 import { AIService } from "@/node/services/aiService";
 import { AgentSession, type AgentSessionChatEvent } from "@/node/services/agentSession";
 import { BackgroundProcessManager } from "@/node/services/backgroundProcessManager";
+import { MCPConfigService } from "@/node/services/mcpConfigService";
+import { MCPServerManager } from "@/node/services/mcpServerManager";
 import {
   isCaughtUpMessage,
   isStreamAbort,
@@ -153,6 +155,27 @@ function renderUnknown(value: unknown): string {
   }
 }
 
+interface MCPServerEntry {
+  name: string;
+  command: string;
+}
+
+function collectMcpServers(value: string, previous: MCPServerEntry[]): MCPServerEntry[] {
+  const eqIndex = value.indexOf("=");
+  if (eqIndex === -1) {
+    throw new Error(`Invalid --mcp format: "${value}". Expected: name=command`);
+  }
+  const name = value.slice(0, eqIndex).trim();
+  const command = value.slice(eqIndex + 1).trim();
+  if (!name) {
+    throw new Error(`Invalid --mcp format: "${value}". Server name is required`);
+  }
+  if (!command) {
+    throw new Error(`Invalid --mcp format: "${value}". Command is required`);
+  }
+  return [...previous, { name, command }];
+}
+
 const program = new Command();
 
 program
@@ -171,6 +194,8 @@ program
   .option("-q, --quiet", "only output final result")
   .option("--workspace-id <id>", "explicit workspace ID (auto-generated if not provided)")
   .option("--config-root <path>", "mux config directory")
+  .option("--mcp <server>", "MCP server as name=command (can be repeated)", collectMcpServers, [])
+  .option("--no-mcp-config", "ignore .mux/mcp.jsonc, use only --mcp servers")
   .addHelpText(
     "after",
     `
@@ -181,6 +206,8 @@ Examples:
   $ mux run --mode plan "Refactor the auth module"
   $ echo "Add logging" | mux run
   $ mux run --json "List all files" | jq '.type'
+  $ mux run --mcp "memory=npx -y @modelcontextprotocol/server-memory" "Remember this"
+  $ mux run --mcp "chrome=npx chrome-devtools-mcp" --mcp "fs=npx @anthropic/mcp-fs" "Take a screenshot"
 `
   );
 
@@ -199,6 +226,8 @@ interface CLIOptions {
   quiet?: boolean;
   workspaceId?: string;
   configRoot?: string;
+  mcp: MCPServerEntry[];
+  mcpConfig: boolean;
 }
 
 const opts = program.opts<CLIOptions>();
@@ -278,6 +307,18 @@ async function main(): Promise<void> {
   );
   ensureProvidersConfig(config);
 
+  // Initialize MCP support
+  const mcpConfigService = new MCPConfigService();
+  const inlineServers: Record<string, string> = {};
+  for (const entry of opts.mcp) {
+    inlineServers[entry.name] = entry.command;
+  }
+  const mcpServerManager = new MCPServerManager(mcpConfigService, {
+    inlineServers,
+    ignoreConfigFile: !opts.mcpConfig,
+  });
+  aiService.setMCPServerManager(mcpServerManager);
+
   const session = new AgentSession({
     workspaceId,
     config,
@@ -523,6 +564,7 @@ async function main(): Promise<void> {
   } finally {
     unsubscribe();
     session.dispose();
+    mcpServerManager.dispose();
   }
 }
 

src/node/services/mcpResultTransform.test.ts

@@ -0,0 +1,156 @@
+import { describe, it, expect } from "bun:test";
+import { transformMCPResult, MAX_IMAGE_DATA_BYTES } from "./mcpResultTransform";
+
+describe("transformMCPResult", () => {
+  describe("image data overflow handling", () => {
+    it("should pass through small images unchanged", () => {
+      const smallImageData = "a".repeat(1000); // 1KB of base64 data
+      const result = transformMCPResult({
+        content: [
+          { type: "text", text: "Screenshot taken" },
+          { type: "image", data: smallImageData, mimeType: "image/png" },
+        ],
+      });
+
+      expect(result).toEqual({
+        type: "content",
+        value: [
+          { type: "text", text: "Screenshot taken" },
+          { type: "media", data: smallImageData, mediaType: "image/png" },
+        ],
+      });
+    });
+
+    it("should truncate large image data to prevent context overflow", () => {
+      // Create a large base64 string that simulates a big screenshot
+      // A typical screenshot could be 500KB-2MB of base64 data
+      const largeImageData = "x".repeat(MAX_IMAGE_DATA_BYTES + 100_000);
+      const result = transformMCPResult({
+        content: [
+          { type: "text", text: "Screenshot taken" },
+          { type: "image", data: largeImageData, mimeType: "image/png" },
+        ],
+      });
+
+      const transformed = result as {
+        type: "content";
+        value: Array<{ type: string; text?: string; data?: string; mediaType?: string }>;
+      };
+
+      expect(transformed.type).toBe("content");
+      expect(transformed.value).toHaveLength(2);
+      expect(transformed.value[0]).toEqual({ type: "text", text: "Screenshot taken" });
+
+      // The image should be replaced with a text message explaining the truncation
+      const imageResult = transformed.value[1];
+      expect(imageResult.type).toBe("text");
+      expect(imageResult.text).toContain("Image data too large");
+      expect(imageResult.text).toContain(String(largeImageData.length));
+    });
+
+    it("should handle multiple images, truncating only the oversized ones", () => {
+      const smallImageData = "small".repeat(100);
+      const largeImageData = "x".repeat(MAX_IMAGE_DATA_BYTES + 50_000);
+
+      const result = transformMCPResult({
+        content: [
+          { type: "image", data: smallImageData, mimeType: "image/png" },
+          { type: "image", data: largeImageData, mimeType: "image/jpeg" },
+        ],
+      });
+
+      const transformed = result as {
+        type: "content";
+        value: Array<{ type: string; text?: string; data?: string; mediaType?: string }>;
+      };
+
+      expect(transformed.value).toHaveLength(2);
+      // Small image passes through
+      expect(transformed.value[0]).toEqual({
+        type: "media",
+        data: smallImageData,
+        mediaType: "image/png",
+      });
+      // Large image gets truncated with explanation
+      expect(transformed.value[1].type).toBe("text");
+      expect(transformed.value[1].text).toContain("Image data too large");
+    });
+
+    it("should report approximate file size in KB/MB in truncation message", () => {
+      // ~1.5MB of base64 data
+      const largeImageData = "y".repeat(1_500_000);
+      const result = transformMCPResult({
+        content: [{ type: "image", data: largeImageData, mimeType: "image/png" }],
+      });
+
+      const transformed = result as {
+        type: "content";
+        value: Array<{ type: string; text?: string }>;
+      };
+
+      expect(transformed.value[0].type).toBe("text");
+      // Should mention MB since it's over 1MB
+      expect(transformed.value[0].text).toMatch(/\d+(\.\d+)?\s*MB/i);
+    });
+  });
+
+  describe("existing functionality", () => {
+    it("should pass through error results unchanged", () => {
+      const errorResult = {
+        isError: true,
+        content: [{ type: "text" as const, text: "Error!" }],
+      };
+      expect(transformMCPResult(errorResult)).toBe(errorResult);
+    });
+
+    it("should pass through toolResult unchanged", () => {
+      const toolResult = { toolResult: { foo: "bar" } };
+      expect(transformMCPResult(toolResult)).toBe(toolResult);
+    });
+
+    it("should pass through results without content array", () => {
+      const noContent = { something: "else" };
+      expect(transformMCPResult(noContent as never)).toBe(noContent);
+    });
+
+    it("should pass through text-only content without transformation wrapper", () => {
+      const textOnly = {
+        content: [
+          { type: "text" as const, text: "Hello" },
+          { type: "text" as const, text: "World" },
+        ],
+      };
+      // No images = no transformation needed
+      expect(transformMCPResult(textOnly)).toBe(textOnly);
+    });
+
+    it("should convert resource content to text", () => {
+      const result = transformMCPResult({
+        content: [
+          { type: "image", data: "abc", mimeType: "image/png" },
+          { type: "resource", resource: { uri: "file:///test.txt", text: "File content" } },
+        ],
+      });
+
+      const transformed = result as {
+        type: "content";
+        value: Array<{ type: string; text?: string; data?: string }>;
+      };
+
+      expect(transformed.value[1]).toEqual({ type: "text", text: "File content" });
+    });
+
+    it("should default to image/png when mimeType is missing", () => {
+      const result = transformMCPResult({
+        content: [{ type: "image", data: "abc", mimeType: "" }],
+      });
+
+      const transformed = result as {
+        type: "content";
+        value: Array<{ type: string; mediaType?: string }>;
+      };
+
+      expect(transformed.value[0].mediaType).toBe("image/png");
+    });
+  });
+});