From 0915f5a00e5e63fa57c3fba9a3cc8090233f3a79 Mon Sep 17 00:00:00 2001 From: Roo Code Date: Tue, 6 Jan 2026 20:41:30 +0000 Subject: [PATCH] feat: implement dynamic context management for MCP responses - Add mcpResponseHandler helper to dynamically check if MCP responses fit within available context budget - Calculate available budget based on model context window, current token usage, and reserved output tokens - Save oversized responses to .roo/tmp/mcp-responses/ directory with a preview in context - Integrate handler into UseMcpToolTool and accessMcpResourceTool - Add comprehensive tests for the new functionality Closes #7042 --- src/core/tools/UseMcpToolTool.ts | 16 +- .../tools/__tests__/useMcpToolTool.spec.ts | 10 + src/core/tools/accessMcpResourceTool.ts | 16 +- .../__tests__/mcpResponseHandler.spec.ts | 255 ++++++++++++++++++ src/core/tools/helpers/mcpResponseHandler.ts | 223 +++++++++++++++ 5 files changed, 518 insertions(+), 2 deletions(-) create mode 100644 src/core/tools/helpers/__tests__/mcpResponseHandler.spec.ts create mode 100644 src/core/tools/helpers/mcpResponseHandler.ts diff --git a/src/core/tools/UseMcpToolTool.ts b/src/core/tools/UseMcpToolTool.ts index b44b96054e..a2d7558559 100644 --- a/src/core/tools/UseMcpToolTool.ts +++ b/src/core/tools/UseMcpToolTool.ts @@ -5,6 +5,7 @@ import { McpExecutionStatus } from "@roo-code/types" import { t } from "../../i18n" import { BaseTool, ToolCallbacks } from "./BaseTool" import type { ToolUse } from "../../shared/tools" +import { handleMcpResponse } from "./helpers/mcpResponseHandler" interface UseMcpToolParams { server_name: string @@ -340,7 +341,20 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> { } await task.say("mcp_server_response", toolResultPretty) - pushToolResult(formatResponse.toolResult(toolResultPretty)) + + // Handle potentially oversized MCP responses by checking against available context budget + const mcpResponseResult = await handleMcpResponse(task, toolResultPretty, { + fileNamePrefix: `mcp-tool-${serverName}-${toolName}`, + }) + + if (mcpResponseResult.savedToFile) { + console.log( + `[UseMcpToolTool] MCP response saved to file: ${mcpResponseResult.filePath} ` + + `(${mcpResponseResult.originalTokenCount} tokens -> ${mcpResponseResult.returnedTokenCount} tokens)`, + ) + } + + pushToolResult(formatResponse.toolResult(mcpResponseResult.content)) } } diff --git a/src/core/tools/__tests__/useMcpToolTool.spec.ts b/src/core/tools/__tests__/useMcpToolTool.spec.ts index 130047ae15..0fa0335091 100644 --- a/src/core/tools/__tests__/useMcpToolTool.spec.ts +++ b/src/core/tools/__tests__/useMcpToolTool.spec.ts @@ -4,6 +4,16 @@ import { useMcpToolTool } from "../UseMcpToolTool" import { Task } from "../../task/Task" import { ToolUse } from "../../../shared/tools" +// Mock the mcpResponseHandler to pass through content unchanged in tests +vi.mock("../helpers/mcpResponseHandler", () => ({ + handleMcpResponse: vi.fn().mockImplementation(async (_task, response) => ({ + content: response, + savedToFile: false, + originalTokenCount: 100, + returnedTokenCount: 100, + })), +})) + // Mock dependencies vi.mock("../../prompts/responses", () => ({ formatResponse: { diff --git a/src/core/tools/accessMcpResourceTool.ts b/src/core/tools/accessMcpResourceTool.ts index ff6705ef6c..8df0c5a9c7 100644 --- a/src/core/tools/accessMcpResourceTool.ts +++ b/src/core/tools/accessMcpResourceTool.ts @@ -3,6 +3,7 @@ import type { ToolUse } from "../../shared/tools" import { Task } from "../task/Task" import { formatResponse } from "../prompts/responses" import { BaseTool, ToolCallbacks } from "./BaseTool" +import { handleMcpResponse } from "./helpers/mcpResponseHandler" interface AccessMcpResourceParams { server_name: string @@ -82,7 +83,20 @@ export class AccessMcpResourceTool extends BaseTool<"access_mcp_resource"> { }) await task.say("mcp_server_response", resourceResultPretty, images) - pushToolResult(formatResponse.toolResult(resourceResultPretty, images)) + + // Handle potentially oversized MCP responses by checking against available context budget + const mcpResponseResult = await handleMcpResponse(task, resourceResultPretty, { + fileNamePrefix: `mcp-resource-${server_name}`, + }) + + if (mcpResponseResult.savedToFile) { + console.log( + `[AccessMcpResourceTool] MCP response saved to file: ${mcpResponseResult.filePath} ` + + `(${mcpResponseResult.originalTokenCount} tokens -> ${mcpResponseResult.returnedTokenCount} tokens)`, + ) + } + + pushToolResult(formatResponse.toolResult(mcpResponseResult.content, images)) } catch (error) { await handleError("accessing MCP resource", error instanceof Error ? error : new Error(String(error))) } diff --git a/src/core/tools/helpers/__tests__/mcpResponseHandler.spec.ts b/src/core/tools/helpers/__tests__/mcpResponseHandler.spec.ts new file mode 100644 index 0000000000..a7f4809b43 --- /dev/null +++ b/src/core/tools/helpers/__tests__/mcpResponseHandler.spec.ts @@ -0,0 +1,255 @@ +import * as fs from "fs/promises" +import * as path from "path" +import { Task } from "../../../task/Task" +import { + handleMcpResponse, + getAvailableMcpResponseBudget, + MCP_RESPONSE_BUDGET_PERCENT, + MCP_RESPONSE_DIR, + DEFAULT_PREVIEW_LINES, +} from "../mcpResponseHandler" + +// Mock dependencies +vi.mock("fs/promises") +vi.mock("../../../../utils/countTokens", () => ({ + countTokens: vi.fn().mockResolvedValue(100), +})) + +const mockFs = vi.mocked(fs) + +describe("mcpResponseHandler", () => { + let mockTask: Partial + + beforeEach(() => { + vi.clearAllMocks() + + // Setup mock task with typical values + mockTask = { + cwd: "/test/workspace", + api: { + getModel: vi.fn().mockReturnValue({ + id: "claude-sonnet-4-20250514", + info: { + contextWindow: 200000, + maxTokens: 8192, + supportsPromptCache: true, + }, + }), + } as any, + getTokenUsage: vi.fn().mockReturnValue({ + contextTokens: 50000, + cacheReadTokens: 0, + cacheWriteTokens: 0, + }), + apiConfiguration: {}, + } + + // Mock fs operations + mockFs.mkdir.mockResolvedValue(undefined) + mockFs.writeFile.mockResolvedValue(undefined) + }) + + describe("handleMcpResponse", () => { + it("should return response directly when it fits within context budget", async () => { + const smallResponse = "This is a small MCP response" + + const result = await handleMcpResponse(mockTask as Task, smallResponse) + + expect(result.savedToFile).toBe(false) + expect(result.content).toBe(smallResponse) + expect(result.filePath).toBeUndefined() + expect(mockFs.writeFile).not.toHaveBeenCalled() + }) + + it("should save response to file when it exceeds context budget", async () => { + // Mock countTokens to return a very large number for the first call (response) + // and small number for preview + const { countTokens } = await import("../../../../utils/countTokens") + vi.mocked(countTokens) + .mockResolvedValueOnce(500000) // Original response tokens - exceeds budget + .mockResolvedValueOnce(100) // Preview tokens + + const largeResponse = "A".repeat(1000000) // Very large response + + const result = await handleMcpResponse(mockTask as Task, largeResponse) + + expect(result.savedToFile).toBe(true) + expect(result.filePath).toBeDefined() + expect(result.filePath).toContain(MCP_RESPONSE_DIR.replace(".roo/", "")) + expect(result.content).toContain("[MCP Response Saved to File]") + expect(result.content).toContain("read_file tool") + expect(mockFs.mkdir).toHaveBeenCalled() + expect(mockFs.writeFile).toHaveBeenCalled() + }) + + it("should generate preview with correct number of lines", async () => { + const { countTokens } = await import("../../../../utils/countTokens") + vi.mocked(countTokens) + .mockResolvedValueOnce(500000) // Original exceeds budget + .mockResolvedValueOnce(50) // Preview tokens + + // Create response with many lines + const lines = Array.from({ length: 100 }, (_, i) => `Line ${i + 1}: Content`) + const multilineResponse = lines.join("\n") + + const result = await handleMcpResponse(mockTask as Task, multilineResponse, { + previewLines: 20, + }) + + expect(result.savedToFile).toBe(true) + // Preview should contain only first 20 lines + const previewSection = result.content.split("---")[1] + const previewLineCount = previewSection.trim().split("\n").length + expect(previewLineCount).toBeLessThanOrEqual(20) + }) + + it("should use custom file name prefix when provided", async () => { + const { countTokens } = await import("../../../../utils/countTokens") + vi.mocked(countTokens) + .mockResolvedValueOnce(500000) // Exceeds budget + .mockResolvedValueOnce(50) + + const response = "Large response content" + + const result = await handleMcpResponse(mockTask as Task, response, { + fileNamePrefix: "custom-prefix", + }) + + expect(result.savedToFile).toBe(true) + // Check that writeFile was called with a path containing the custom prefix + const writeFileCall = mockFs.writeFile.mock.calls[0] + expect(writeFileCall[0]).toContain("custom-prefix") + }) + + it("should handle token counting errors gracefully", async () => { + const { countTokens } = await import("../../../../utils/countTokens") + vi.mocked(countTokens).mockRejectedValue(new Error("Token counting failed")) + + const response = "Test response" + + // Should not throw, should fall back to character-based estimation + const result = await handleMcpResponse(mockTask as Task, response) + + expect(result).toBeDefined() + expect(typeof result.originalTokenCount).toBe("number") + }) + + it("should include token count information in result", async () => { + const { countTokens } = await import("../../../../utils/countTokens") + vi.mocked(countTokens).mockResolvedValue(1500) + + const response = "Test response with moderate size" + + const result = await handleMcpResponse(mockTask as Task, response) + + expect(result.originalTokenCount).toBe(1500) + expect(result.returnedTokenCount).toBe(1500) + }) + + it("should save to correct directory structure", async () => { + const { countTokens } = await import("../../../../utils/countTokens") + vi.mocked(countTokens) + .mockResolvedValueOnce(500000) + .mockResolvedValueOnce(50) + + const response = "Large response" + + await handleMcpResponse(mockTask as Task, response) + + // Check mkdir was called with correct path + expect(mockFs.mkdir).toHaveBeenCalledWith( + path.join("/test/workspace", MCP_RESPONSE_DIR), + { recursive: true }, + ) + }) + + it("should handle zero available budget gracefully", async () => { + // Set up task with very high current token usage + mockTask.getTokenUsage = vi.fn().mockReturnValue({ + contextTokens: 195000, // Nearly at context window limit + cacheReadTokens: 0, + cacheWriteTokens: 0, + }) + + const { countTokens } = await import("../../../../utils/countTokens") + vi.mocked(countTokens) + .mockResolvedValueOnce(100) + .mockResolvedValueOnce(50) + + const response = "Small response" + + const result = await handleMcpResponse(mockTask as Task, response) + + // Even small response should be saved to file when budget is exhausted + expect(result.savedToFile).toBe(true) + }) + }) + + describe("getAvailableMcpResponseBudget", () => { + it("should calculate correct budget based on context window and usage", () => { + const budget = getAvailableMcpResponseBudget(mockTask as Task) + + // contextWindow: 200000 + // maxOutputTokens: ~8192 (from model info) + // contextTokens: 50000 + // remaining = 200000 - 8192 - 50000 = 141808 + // budget = 141808 * 0.5 = 70904 + expect(budget).toBeGreaterThan(0) + expect(budget).toBeLessThan(200000) + }) + + it("should return smaller budget when context is nearly full", () => { + mockTask.getTokenUsage = vi.fn().mockReturnValue({ + contextTokens: 180000, + cacheReadTokens: 0, + cacheWriteTokens: 0, + }) + + const budget = getAvailableMcpResponseBudget(mockTask as Task) + + // remaining = 200000 - 8192 - 180000 = 11808 + // budget = 11808 * 0.5 = 5904 + expect(budget).toBeLessThan(10000) + }) + + it("should handle different model context windows", () => { + // Test with smaller context window model + mockTask.api = { + getModel: vi.fn().mockReturnValue({ + id: "small-model", + info: { + contextWindow: 32000, + maxTokens: 4096, + supportsPromptCache: false, + }, + }), + } as any + + mockTask.getTokenUsage = vi.fn().mockReturnValue({ + contextTokens: 10000, + cacheReadTokens: 0, + cacheWriteTokens: 0, + }) + + const budget = getAvailableMcpResponseBudget(mockTask as Task) + + // remaining = 32000 - 4096 - 10000 = 17904 + // budget = 17904 * 0.5 = 8952 + expect(budget).toBeLessThan(10000) + }) + }) + + describe("constants", () => { + it("should have correct budget percentage", () => { + expect(MCP_RESPONSE_BUDGET_PERCENT).toBe(0.5) + }) + + it("should have correct default preview lines", () => { + expect(DEFAULT_PREVIEW_LINES).toBe(50) + }) + + it("should have correct response directory", () => { + expect(MCP_RESPONSE_DIR).toBe(".roo/tmp/mcp-responses") + }) + }) +}) diff --git a/src/core/tools/helpers/mcpResponseHandler.ts b/src/core/tools/helpers/mcpResponseHandler.ts new file mode 100644 index 0000000000..942f769de7 --- /dev/null +++ b/src/core/tools/helpers/mcpResponseHandler.ts @@ -0,0 +1,223 @@ +import * as fs from "fs/promises" +import * as path from "path" +import { Anthropic } from "@anthropic-ai/sdk" +import { ANTHROPIC_DEFAULT_MAX_TOKENS } from "@roo-code/types" + +import { Task } from "../../task/Task" +import { countTokens } from "../../../utils/countTokens" +import { getModelMaxOutputTokens } from "../../../shared/api" + +/** + * Percentage of available context to use as budget for MCP responses. + * Uses the same percentage as file reading for consistency. + */ +export const MCP_RESPONSE_BUDGET_PERCENT = 0.5 + +/** + * Default number of preview lines to show when response is saved to file. + */ +export const DEFAULT_PREVIEW_LINES = 50 + +/** + * Directory name for storing oversized MCP responses within .roo folder. + */ +export const MCP_RESPONSE_DIR = ".roo/tmp/mcp-responses" + +export interface McpResponseHandlerResult { + /** The content to include in the tool result (either full response or preview with file path) */ + content: string + /** Whether the response was saved to a file */ + savedToFile: boolean + /** Path to the saved file (if savedToFile is true) */ + filePath?: string + /** Token count of the original response */ + originalTokenCount: number + /** Token count of the content being returned */ + returnedTokenCount: number +} + +export interface McpResponseHandlerOptions { + /** Number of preview lines to show when response is saved to file (default: 50) */ + previewLines?: number + /** Custom file name prefix for saved files */ + fileNamePrefix?: string +} + +/** + * Handles MCP responses by checking if they fit within the available context budget. + * If the response is too large, it saves it to a file and returns a preview. + * + * This implements dynamic context management for MCP responses, similar to how + * ReadFileTool handles large files. + * + * @param task - The current Task instance + * @param response - The MCP response content + * @param options - Optional configuration + * @returns Result containing the content to use and metadata about the operation + */ +export async function handleMcpResponse( + task: Task, + response: string, + options: McpResponseHandlerOptions = {}, +): Promise { + const { previewLines = DEFAULT_PREVIEW_LINES, fileNamePrefix = "mcp-response" } = options + + // Get model info and calculate available context budget + const { id: modelId, info: modelInfo } = task.api.getModel() + const { contextTokens } = task.getTokenUsage() + const contextWindow = modelInfo.contextWindow + + const maxOutputTokens = + getModelMaxOutputTokens({ + modelId, + model: modelInfo, + settings: task.apiConfiguration, + }) ?? ANTHROPIC_DEFAULT_MAX_TOKENS + + // Calculate available token budget for MCP response + const remainingTokens = contextWindow - maxOutputTokens - (contextTokens || 0) + const mcpResponseBudget = Math.floor(remainingTokens * MCP_RESPONSE_BUDGET_PERCENT) + + // Count tokens in the response + const contentBlocks: Anthropic.Messages.ContentBlockParam[] = [{ type: "text", text: response }] + let responseTokens: number + try { + responseTokens = await countTokens(contentBlocks, { useWorker: false }) + } catch { + // Fallback: conservative estimate (2 chars per token) + responseTokens = Math.ceil(response.length / 2) + } + + // If response fits within budget, return it directly + if (responseTokens <= mcpResponseBudget && mcpResponseBudget > 0) { + return { + content: response, + savedToFile: false, + originalTokenCount: responseTokens, + returnedTokenCount: responseTokens, + } + } + + // Response is too large - save to file and return preview + const filePath = await saveResponseToFile(task.cwd, response, fileNamePrefix) + const preview = generatePreview(response, previewLines) + + // Count tokens in preview content + let previewTokens: number + try { + const previewBlocks: Anthropic.Messages.ContentBlockParam[] = [{ type: "text", text: preview }] + previewTokens = await countTokens(previewBlocks, { useWorker: false }) + } catch { + previewTokens = Math.ceil(preview.length / 2) + } + + // Build the response with file reference and preview + const relativePath = path.relative(task.cwd, filePath) + const resultContent = formatOversizedResponse(relativePath, response.length, responseTokens, preview) + + return { + content: resultContent, + savedToFile: true, + filePath: relativePath, + originalTokenCount: responseTokens, + returnedTokenCount: previewTokens, + } +} + +/** + * Saves the MCP response to a file in the .roo/tmp/mcp-responses directory. + * + * @param cwd - Current working directory + * @param content - Content to save + * @param prefix - File name prefix + * @returns Absolute path to the saved file + */ +async function saveResponseToFile(cwd: string, content: string, prefix: string): Promise { + const responseDir = path.join(cwd, MCP_RESPONSE_DIR) + + // Ensure directory exists + await fs.mkdir(responseDir, { recursive: true }) + + // Generate unique filename with timestamp + const timestamp = Date.now() + const randomSuffix = Math.random().toString(36).substring(2, 8) + const fileName = `${prefix}-${timestamp}-${randomSuffix}.txt` + const filePath = path.join(responseDir, fileName) + + // Write content to file + await fs.writeFile(filePath, content, "utf-8") + + return filePath +} + +/** + * Generates a preview of the response content. + * + * @param content - Full response content + * @param maxLines - Maximum number of lines to include in preview + * @returns Preview string + */ +function generatePreview(content: string, maxLines: number): string { + const lines = content.split("\n") + if (lines.length <= maxLines) { + return content + } + + const previewLines = lines.slice(0, maxLines) + return previewLines.join("\n") +} + +/** + * Formats the response when it's been saved to a file. + * + * @param filePath - Relative path to saved file + * @param contentLength - Length of original content in characters + * @param tokenCount - Token count of original content + * @param preview - Preview of the content + * @returns Formatted response string + */ +function formatOversizedResponse( + filePath: string, + contentLength: number, + tokenCount: number, + preview: string, +): string { + return `[MCP Response Saved to File] +The MCP response was too large to include in context (${tokenCount.toLocaleString()} tokens, ${contentLength.toLocaleString()} characters). +The full response has been saved to: ${filePath} + +You can read the complete response using the read_file tool with the path above. + +Preview (first ${DEFAULT_PREVIEW_LINES} lines): +--- +${preview} +--- + +Suggested actions: +- Use read_file with line_range to read specific sections of the file +- Process the data using bash/python scripts if needed for analysis +- Extract specific information by reading relevant portions` +} + +/** + * Calculates the available context budget for MCP responses. + * Useful for pre-checks or logging. + * + * @param task - The current Task instance + * @returns Available token budget for MCP responses + */ +export function getAvailableMcpResponseBudget(task: Task): number { + const { id: modelId, info: modelInfo } = task.api.getModel() + const { contextTokens } = task.getTokenUsage() + const contextWindow = modelInfo.contextWindow + + const maxOutputTokens = + getModelMaxOutputTokens({ + modelId, + model: modelInfo, + settings: task.apiConfiguration, + }) ?? ANTHROPIC_DEFAULT_MAX_TOKENS + + const remainingTokens = contextWindow - maxOutputTokens - (contextTokens || 0) + return Math.floor(remainingTokens * MCP_RESPONSE_BUDGET_PERCENT) +}