From 75d619565dcb5452687851bebaf24f0312830fbf Mon Sep 17 00:00:00 2001 From: Ammar Date: Mon, 8 Dec 2025 14:56:47 -0600 Subject: [PATCH 1/2] fix: truncate oversized MCP images to prevent context overflow Large base64-encoded images from MCP servers (like chrome-devtools-mcp screenshots) were overflowing the model's context window. - Extract mcpResultTransform module with 256KB limit for image data - Oversized images replaced with descriptive text message - Add MCP server support to `mux run` CLI command - Add comprehensive tests for image truncation logic _Generated with `mux`_ --- src/cli/run.ts | 8 + src/node/services/mcpResultTransform.test.ts | 156 +++++++++++++++++++ src/node/services/mcpResultTransform.ts | 120 ++++++++++++++ src/node/services/mcpServerManager.ts | 88 +---------- 4 files changed, 285 insertions(+), 87 deletions(-) create mode 100644 src/node/services/mcpResultTransform.test.ts create mode 100644 src/node/services/mcpResultTransform.ts diff --git a/src/cli/run.ts b/src/cli/run.ts index 0cb440d6b2..fa400e2f64 100644 --- a/src/cli/run.ts +++ b/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, @@ -278,6 +280,11 @@ async function main(): Promise { ); ensureProvidersConfig(config); + // Initialize MCP support + const mcpConfigService = new MCPConfigService(); + const mcpServerManager = new MCPServerManager(mcpConfigService); + aiService.setMCPServerManager(mcpServerManager); + const session = new AgentSession({ workspaceId, config, @@ -523,6 +530,7 @@ async function main(): Promise { } finally { unsubscribe(); session.dispose(); + mcpServerManager.dispose(); } } diff --git a/src/node/services/mcpResultTransform.test.ts b/src/node/services/mcpResultTransform.test.ts new file mode 100644 index 0000000000..f169843a42 --- /dev/null +++ b/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"); + }); + }); +}); diff --git a/src/node/services/mcpResultTransform.ts b/src/node/services/mcpResultTransform.ts new file mode 100644 index 0000000000..8afc044d80 --- /dev/null +++ b/src/node/services/mcpResultTransform.ts @@ -0,0 +1,120 @@ +import { log } from "@/node/services/log"; + +/** + * Maximum size of base64 image data in bytes before truncation. + * Large images can overflow the model's context window. 256KB of base64 + * represents roughly 192KB of actual image data, which is sufficient for + * screenshots while preventing context overflow. + */ +export const MAX_IMAGE_DATA_BYTES = 256 * 1024; // 256KB of base64 data + +/** + * MCP CallToolResult content types (from @ai-sdk/mcp) + */ +interface MCPTextContent { + type: "text"; + text: string; +} + +interface MCPImageContent { + type: "image"; + data: string; // base64 + mimeType: string; +} + +interface MCPResourceContent { + type: "resource"; + resource: { uri: string; text?: string; blob?: string; mimeType?: string }; +} + +type MCPContent = MCPTextContent | MCPImageContent | MCPResourceContent; + +export interface MCPCallToolResult { + content?: MCPContent[]; + isError?: boolean; + toolResult?: unknown; +} + +/** + * AI SDK LanguageModelV2ToolResultOutput content types + */ +type AISDKContentPart = + | { type: "text"; text: string } + | { type: "media"; data: string; mediaType: string }; + +/** + * Format byte size as human-readable string (KB or MB) + */ +function formatBytes(bytes: number): string { + if (bytes >= 1_000_000) { + return `${(bytes / 1_000_000).toFixed(1)} MB`; + } + return `${Math.round(bytes / 1000)} KB`; +} + +/** + * Transform MCP tool result to AI SDK format. + * Converts MCP's "image" content type to AI SDK's "media" type. + * Truncates large images to prevent context overflow. + */ +export function transformMCPResult(result: MCPCallToolResult): unknown { + // If it's an error or has toolResult, pass through as-is + if (result.isError || result.toolResult !== undefined) { + return result; + } + + // If no content array, pass through + if (!result.content || !Array.isArray(result.content)) { + return result; + } + + // Check if any content is an image + const hasImage = result.content.some((c) => c.type === "image"); + if (!hasImage) { + return result; + } + + // Debug: log what we received from MCP + log.debug("[MCP] transformMCPResult input", { + contentTypes: result.content.map((c) => c.type), + imageItems: result.content + .filter((c): c is MCPImageContent => c.type === "image") + .map((c) => ({ type: c.type, mimeType: c.mimeType, dataLen: c.data?.length })), + }); + + // Transform to AI SDK content format + const transformedContent: AISDKContentPart[] = result.content.map((item) => { + if (item.type === "text") { + return { type: "text" as const, text: item.text }; + } + if (item.type === "image") { + const imageItem = item; + // Check if image data exceeds the limit + const dataLength = imageItem.data?.length ?? 0; + if (dataLength > MAX_IMAGE_DATA_BYTES) { + log.warn("[MCP] Image data too large, truncating", { + mimeType: imageItem.mimeType, + dataLength, + maxAllowed: MAX_IMAGE_DATA_BYTES, + }); + return { + type: "text" as const, + text: `[Image data too large to include in context: ${formatBytes(dataLength)} (${dataLength} bytes). The image was captured but cannot be displayed inline. Consider using a smaller viewport or requesting a specific region.]`, + }; + } + // Ensure mediaType is present - default to image/png if missing + const mediaType = imageItem.mimeType || "image/png"; + log.debug("[MCP] Transforming image content", { mimeType: imageItem.mimeType, mediaType }); + return { type: "media" as const, data: imageItem.data, mediaType }; + } + // For resource type, convert to text representation + if (item.type === "resource") { + const text = item.resource.text ?? item.resource.uri; + return { type: "text" as const, text }; + } + // Fallback: stringify unknown content + return { type: "text" as const, text: JSON.stringify(item) }; + }); + + return { type: "content", value: transformedContent }; +} diff --git a/src/node/services/mcpServerManager.ts b/src/node/services/mcpServerManager.ts index 99f635c5b9..c90eed7c0e 100644 --- a/src/node/services/mcpServerManager.ts +++ b/src/node/services/mcpServerManager.ts @@ -6,98 +6,12 @@ import type { MCPServerMap, MCPTestResult } from "@/common/types/mcp"; import type { Runtime } from "@/node/runtime/Runtime"; import type { MCPConfigService } from "@/node/services/mcpConfigService"; import { createRuntime } from "@/node/runtime/runtimeFactory"; +import { transformMCPResult, type MCPCallToolResult } from "@/node/services/mcpResultTransform"; const TEST_TIMEOUT_MS = 10_000; const IDLE_TIMEOUT_MS = 10 * 60 * 1000; // 10 minutes const IDLE_CHECK_INTERVAL_MS = 60 * 1000; // Check every minute -/** - * MCP CallToolResult content types (from @ai-sdk/mcp) - */ -interface MCPTextContent { - type: "text"; - text: string; -} - -interface MCPImageContent { - type: "image"; - data: string; // base64 - mimeType: string; -} - -interface MCPResourceContent { - type: "resource"; - resource: { uri: string; text?: string; blob?: string; mimeType?: string }; -} - -type MCPContent = MCPTextContent | MCPImageContent | MCPResourceContent; - -interface MCPCallToolResult { - content?: MCPContent[]; - isError?: boolean; - toolResult?: unknown; -} - -/** - * AI SDK LanguageModelV2ToolResultOutput content types - */ -type AISDKContentPart = - | { type: "text"; text: string } - | { type: "media"; data: string; mediaType: string }; - -/** - * Transform MCP tool result to AI SDK format. - * Converts MCP's "image" content type to AI SDK's "media" type. - */ -function transformMCPResult(result: MCPCallToolResult): unknown { - // If it's an error or has toolResult, pass through as-is - if (result.isError || result.toolResult !== undefined) { - return result; - } - - // If no content array, pass through - if (!result.content || !Array.isArray(result.content)) { - return result; - } - - // Check if any content is an image - const hasImage = result.content.some((c) => c.type === "image"); - if (!hasImage) { - return result; - } - - // Debug: log what we received from MCP - log.debug("[MCP] transformMCPResult input", { - contentTypes: result.content.map((c) => c.type), - imageItems: result.content - .filter((c): c is MCPImageContent => c.type === "image") - .map((c) => ({ type: c.type, mimeType: c.mimeType, dataLen: c.data?.length })), - }); - - // Transform to AI SDK content format - const transformedContent: AISDKContentPart[] = result.content.map((item) => { - if (item.type === "text") { - return { type: "text" as const, text: item.text }; - } - if (item.type === "image") { - const imageItem = item; - // Ensure mediaType is present - default to image/png if missing - const mediaType = imageItem.mimeType || "image/png"; - log.debug("[MCP] Transforming image content", { mimeType: imageItem.mimeType, mediaType }); - return { type: "media" as const, data: imageItem.data, mediaType }; - } - // For resource type, convert to text representation - if (item.type === "resource") { - const text = item.resource.text ?? item.resource.uri; - return { type: "text" as const, text }; - } - // Fallback: stringify unknown content - return { type: "text" as const, text: JSON.stringify(item) }; - }); - - return { type: "content", value: transformedContent }; -} - /** * Wrap MCP tools to transform their results to AI SDK format. * This ensures image content is properly converted to media type. From 3fdc8ed1f9a33b3fb09100ffa00b2f99153049ad Mon Sep 17 00:00:00 2001 From: Ammar Date: Mon, 8 Dec 2025 15:27:58 -0600 Subject: [PATCH 2/2] feat: add --mcp option for inline MCP server specification - Add --mcp flag: mux run --mcp "name=command" (can be repeated) - Add --no-mcp-config flag to ignore .mux/mcp.jsonc file - MCPServerManager accepts inlineServers and ignoreConfigFile options - Inline servers override config file servers with the same name - Add CLI tests for --mcp argument parsing and validation _Generated with `mux`_ --- src/cli/run.test.ts | 59 +++++++++++++++++++++++++++ src/cli/run.ts | 36 +++++++++++++++- src/node/services/mcpServerManager.ts | 36 ++++++++++++++-- 3 files changed, 127 insertions(+), 4 deletions(-) diff --git a/src/cli/run.test.ts b/src/cli/run.test.ts index 58746c6c54..516b1bd0c6 100644 --- a/src/cli/run.test.ts +++ b/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", () => { diff --git a/src/cli/run.ts b/src/cli/run.ts index fa400e2f64..b53bba92fe 100644 --- a/src/cli/run.ts +++ b/src/cli/run.ts @@ -155,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 @@ -173,6 +194,8 @@ program .option("-q, --quiet", "only output final result") .option("--workspace-id ", "explicit workspace ID (auto-generated if not provided)") .option("--config-root ", "mux config directory") + .option("--mcp ", "MCP server as name=command (can be repeated)", collectMcpServers, []) + .option("--no-mcp-config", "ignore .mux/mcp.jsonc, use only --mcp servers") .addHelpText( "after", ` @@ -183,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" ` ); @@ -201,6 +226,8 @@ interface CLIOptions { quiet?: boolean; workspaceId?: string; configRoot?: string; + mcp: MCPServerEntry[]; + mcpConfig: boolean; } const opts = program.opts(); @@ -282,7 +309,14 @@ async function main(): Promise { // Initialize MCP support const mcpConfigService = new MCPConfigService(); - const mcpServerManager = new MCPServerManager(mcpConfigService); + const inlineServers: Record = {}; + 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({ diff --git a/src/node/services/mcpServerManager.ts b/src/node/services/mcpServerManager.ts index c90eed7c0e..c3300745fb 100644 --- a/src/node/services/mcpServerManager.ts +++ b/src/node/services/mcpServerManager.ts @@ -100,12 +100,30 @@ interface WorkspaceServers { lastActivity: number; } +export interface MCPServerManagerOptions { + /** Inline servers to use (merged with config file servers by default) */ + inlineServers?: MCPServerMap; + /** If true, ignore config file servers and use only inline servers */ + ignoreConfigFile?: boolean; +} + export class MCPServerManager { private readonly workspaceServers = new Map(); private readonly idleCheckInterval: ReturnType; + private inlineServers: MCPServerMap = {}; + private ignoreConfigFile = false; - constructor(private readonly configService: MCPConfigService) { + constructor( + private readonly configService: MCPConfigService, + options?: MCPServerManagerOptions + ) { this.idleCheckInterval = setInterval(() => this.cleanupIdleServers(), IDLE_CHECK_INTERVAL_MS); + if (options?.inlineServers) { + this.inlineServers = options.inlineServers; + } + if (options?.ignoreConfigFile) { + this.ignoreConfigFile = options.ignoreConfigFile; + } } /** @@ -130,12 +148,24 @@ export class MCPServerManager { } } + /** + * Get merged servers: config file servers (unless ignoreConfigFile) + inline servers. + * Inline servers take precedence over config file servers with the same name. + */ + private async getMergedServers(projectPath: string): Promise { + const configServers = this.ignoreConfigFile + ? {} + : await this.configService.listServers(projectPath); + // Inline servers override config file servers + return { ...configServers, ...this.inlineServers }; + } + /** * List configured MCP servers for a project (name -> command). * Used to show server info in the system prompt. */ async listServers(projectPath: string): Promise { - return this.configService.listServers(projectPath); + return this.getMergedServers(projectPath); } async getToolsForWorkspace(options: { @@ -145,7 +175,7 @@ export class MCPServerManager { workspacePath: string; }): Promise> { const { workspaceId, projectPath, runtime, workspacePath } = options; - const servers = await this.configService.listServers(projectPath); + const servers = await this.getMergedServers(projectPath); const signature = JSON.stringify(servers); const serverNames = Object.keys(servers);