Skip to content

Commit 098ff6d

Browse files
committed
feat: Save large MCP responses to files to prevent context overflow
- Added McpResponseHandler utility class to handle large MCP responses - Automatically saves responses over 50KB (configurable) to temporary files - Provides preview and file path in context instead of full response - Integrated into useMcpToolTool and accessMcpResourceTool - Added mcpResponseSizeThreshold configuration option - Added comprehensive tests for the new functionality Fixes #7042
1 parent 9ffc2b0 commit 098ff6d

File tree

6 files changed

+498
-3
lines changed

6 files changed

+498
-3
lines changed

packages/types/src/global-settings.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,7 @@ export const globalSettingsSchema = z.object({
133133

134134
mcpEnabled: z.boolean().optional(),
135135
enableMcpServerCreation: z.boolean().optional(),
136+
mcpResponseSizeThreshold: z.number().optional(),
136137

137138
remoteControlEnabled: z.boolean().optional(),
138139

src/core/tools/accessMcpResourceTool.ts

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { ClineAskUseMcpServer } from "../../shared/ExtensionMessage"
22
import { ToolUse, RemoveClosingTag, AskApproval, HandleError, PushToolResult } from "../../shared/tools"
33
import { Task } from "../task/Task"
44
import { formatResponse } from "../prompts/responses"
5+
import { defaultMcpResponseHandler } from "../../utils/mcpResponseHandler"
56

67
export async function accessMcpResourceTool(
78
cline: Task,
@@ -57,7 +58,7 @@ export async function accessMcpResourceTool(
5758
await cline.say("mcp_server_request_started")
5859
const resourceResult = await cline.providerRef.deref()?.getMcpHub()?.readResource(server_name, uri)
5960

60-
const resourceResultPretty =
61+
const resourceResultText =
6162
resourceResult?.contents
6263
.map((item) => {
6364
if (item.text) {
@@ -81,6 +82,16 @@ export async function accessMcpResourceTool(
8182
}
8283
})
8384

85+
// Check if response is large and should be saved to file
86+
const processedResponse = await defaultMcpResponseHandler.processResponse(
87+
resourceResultText,
88+
server_name,
89+
uri,
90+
)
91+
92+
// Use the processed content (either original or file reference)
93+
const resourceResultPretty = processedResponse.content
94+
8495
await cline.say("mcp_server_response", resourceResultPretty, images)
8596
pushToolResult(formatResponse.toolResult(resourceResultPretty, images))
8697

src/core/tools/useMcpToolTool.ts

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import { formatResponse } from "../prompts/responses"
44
import { ClineAskUseMcpServer } from "../../shared/ExtensionMessage"
55
import { McpExecutionStatus } from "@roo-code/types"
66
import { t } from "../../i18n"
7+
import { defaultMcpResponseHandler } from "../../utils/mcpResponseHandler"
78

89
interface McpToolParams {
910
server_name?: string
@@ -135,13 +136,18 @@ async function executeToolAndProcessResult(
135136
const outputText = processToolContent(toolResult)
136137

137138
if (outputText) {
139+
// Check if response is large and should be saved to file
140+
const processedResponse = await defaultMcpResponseHandler.processResponse(outputText, serverName, toolName)
141+
138142
await sendExecutionStatus(cline, {
139143
executionId,
140144
status: "output",
141-
response: outputText,
145+
response: processedResponse.savedToFile
146+
? `Response saved to file: ${processedResponse.filePath}`
147+
: outputText,
142148
})
143149

144-
toolResultPretty = (toolResult.isError ? "Error:\n" : "") + outputText
150+
toolResultPretty = (toolResult.isError ? "Error:\n" : "") + processedResponse.content
145151
}
146152

147153
// Send completion status

src/core/webview/ClineProvider.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1701,6 +1701,7 @@ export class ClineProvider
17011701
maxDiagnosticMessages,
17021702
includeTaskHistoryInEnhance,
17031703
remoteControlEnabled,
1704+
mcpResponseSizeThreshold,
17041705
} = await this.getState()
17051706

17061707
const telemetryKey = process.env.POSTHOG_API_KEY
@@ -1829,6 +1830,7 @@ export class ClineProvider
18291830
maxDiagnosticMessages: maxDiagnosticMessages ?? 50,
18301831
includeTaskHistoryInEnhance: includeTaskHistoryInEnhance ?? false,
18311832
remoteControlEnabled: remoteControlEnabled ?? false,
1833+
mcpResponseSizeThreshold: mcpResponseSizeThreshold ?? 50000,
18321834
}
18331835
}
18341836

@@ -2018,6 +2020,8 @@ export class ClineProvider
20182020
includeTaskHistoryInEnhance: stateValues.includeTaskHistoryInEnhance ?? false,
20192021
// Add remoteControlEnabled setting
20202022
remoteControlEnabled: stateValues.remoteControlEnabled ?? false,
2023+
// Add MCP response size threshold setting
2024+
mcpResponseSizeThreshold: stateValues.mcpResponseSizeThreshold ?? 50000,
20212025
}
20222026
}
20232027

Lines changed: 241 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,241 @@
1+
import { describe, it, expect, vi, beforeEach, afterEach, beforeAll } from "vitest"
2+
import * as fs from "fs/promises"
3+
import * as os from "os"
4+
5+
// Mock the modules before importing the module under test
6+
vi.mock("fs/promises", () => ({
7+
mkdir: vi.fn(),
8+
writeFile: vi.fn(),
9+
readdir: vi.fn(),
10+
stat: vi.fn(),
11+
unlink: vi.fn(),
12+
}))
13+
vi.mock("os", () => ({
14+
tmpdir: vi.fn(() => "/mock/tmp/dir"),
15+
}))
16+
vi.mock("./safeWriteJson", () => ({
17+
safeWriteJson: vi.fn(),
18+
}))
19+
20+
// Import after mocks are set up
21+
const { McpResponseHandler } = await import("./mcpResponseHandler")
22+
const { safeWriteJson } = await import("./safeWriteJson")
23+
24+
describe("McpResponseHandler", () => {
25+
let handler: InstanceType<typeof McpResponseHandler>
26+
const mockTmpDir = "/mock/tmp/dir"
27+
const mockResponseDir = `${mockTmpDir}/roo-code-mcp-responses`
28+
29+
beforeEach(() => {
30+
vi.clearAllMocks()
31+
vi.mocked(fs.mkdir).mockResolvedValue(undefined)
32+
vi.mocked(fs.writeFile).mockResolvedValue(undefined)
33+
vi.mocked(safeWriteJson).mockResolvedValue(undefined)
34+
35+
// Mock Date for consistent file naming
36+
vi.spyOn(Date.prototype, "toISOString").mockReturnValue("2024-01-01T00:00:00.000Z")
37+
})
38+
39+
afterEach(() => {
40+
vi.restoreAllMocks()
41+
})
42+
43+
describe("with default threshold", () => {
44+
beforeEach(() => {
45+
handler = new McpResponseHandler()
46+
})
47+
48+
it("should return small responses directly without saving to file", async () => {
49+
const smallResponse = "Small response content"
50+
const result = await handler.processResponse(smallResponse, "testServer", "testTool")
51+
52+
expect(result.savedToFile).toBe(false)
53+
expect(result.content).toBe(smallResponse)
54+
expect(result.filePath).toBeUndefined()
55+
expect(fs.writeFile).not.toHaveBeenCalled()
56+
})
57+
58+
it("should save large responses to file and return preview", async () => {
59+
const largeResponse = "x".repeat(60000) // 60KB response
60+
const result = await handler.processResponse(largeResponse, "testServer", "testTool")
61+
62+
expect(result.savedToFile).toBe(true)
63+
expect(result.filePath).toBeDefined()
64+
expect(result.content).toContain("[MCP Response saved to file due to large size")
65+
expect(result.content).toContain("File:")
66+
expect(result.content).toContain("Preview of response:")
67+
// Check that preview shows limited content
68+
const lines = result.content.split("\n")
69+
const previewStartIndex = lines.findIndex((line: string) => line.includes("Preview of response:"))
70+
expect(previewStartIndex).toBeGreaterThan(-1)
71+
72+
expect(fs.mkdir).toHaveBeenCalledWith(expect.stringContaining("roo-code-mcp-responses"), {
73+
recursive: true,
74+
})
75+
expect(fs.writeFile).toHaveBeenCalledWith(
76+
expect.stringContaining("mcp-response-testServer-testTool"),
77+
largeResponse,
78+
"utf-8",
79+
)
80+
})
81+
82+
it("should handle JSON responses using processStructuredResponse", async () => {
83+
const jsonData = {
84+
data: "x".repeat(50000),
85+
metadata: { count: 100 },
86+
}
87+
const result = await handler.processStructuredResponse(jsonData, "dbServer", "queryTool")
88+
89+
expect(result.savedToFile).toBe(true)
90+
expect(result.content).toContain("Preview of response:")
91+
expect(result.content).toContain('"data"')
92+
expect(result.content).toContain('"metadata"')
93+
94+
expect(safeWriteJson).toHaveBeenCalledWith(
95+
expect.stringContaining("mcp-response-dbServer-queryTool"),
96+
jsonData,
97+
)
98+
})
99+
100+
it("should handle non-JSON responses in preview", async () => {
101+
const textResponse = "Plain text response\n".repeat(3000) // Large plain text
102+
const result = await handler.processResponse(textResponse, "textServer", "textTool")
103+
104+
expect(result.savedToFile).toBe(true)
105+
expect(result.content).toContain("Preview of response:")
106+
expect(result.content).toContain("Plain text response")
107+
})
108+
109+
it("should handle file write errors gracefully", async () => {
110+
vi.mocked(fs.writeFile).mockRejectedValue(new Error("Write failed"))
111+
112+
const largeResponse = "x".repeat(60000)
113+
await expect(handler.processResponse(largeResponse, "testServer", "testTool")).rejects.toThrow(
114+
"Write failed",
115+
)
116+
})
117+
118+
it("should create directory if it doesn't exist", async () => {
119+
const largeResponse = "x".repeat(60000)
120+
await handler.processResponse(largeResponse, "testServer", "testTool")
121+
122+
expect(fs.mkdir).toHaveBeenCalledWith(expect.stringContaining("roo-code-mcp-responses"), {
123+
recursive: true,
124+
})
125+
})
126+
})
127+
128+
describe("with custom threshold", () => {
129+
it("should use custom maxResponseSize when provided", async () => {
130+
handler = new McpResponseHandler({ maxResponseSize: 1000 }) // 1KB threshold
131+
132+
const smallResponse = "x".repeat(500) // 500 bytes - under threshold
133+
const result1 = await handler.processResponse(smallResponse, "server", "tool")
134+
expect(result1.savedToFile).toBe(false)
135+
136+
const largeResponse = "x".repeat(1500) // 1.5KB - over threshold
137+
const result2 = await handler.processResponse(largeResponse, "server", "tool")
138+
expect(result2.savedToFile).toBe(true)
139+
})
140+
})
141+
142+
describe("edge cases", () => {
143+
beforeEach(() => {
144+
handler = new McpResponseHandler()
145+
})
146+
147+
it("should handle empty responses", async () => {
148+
const result = await handler.processResponse("", "server", "tool")
149+
expect(result.savedToFile).toBe(false)
150+
expect(result.content).toBe("")
151+
})
152+
153+
it("should handle responses exactly at threshold", async () => {
154+
const maxResponseSize = 50 * 1024 // 50KB
155+
handler = new McpResponseHandler({ maxResponseSize })
156+
157+
const response = "x".repeat(maxResponseSize)
158+
const result = await handler.processResponse(response, "server", "tool")
159+
expect(result.savedToFile).toBe(false) // Should NOT save when exactly at threshold (<=)
160+
161+
const largerResponse = "x".repeat(maxResponseSize + 1)
162+
const result2 = await handler.processResponse(largerResponse, "server", "tool")
163+
expect(result2.savedToFile).toBe(true) // Should save when over threshold
164+
})
165+
166+
it("should handle special characters in server and tool names", async () => {
167+
const largeResponse = "x".repeat(60000)
168+
const result = await handler.processResponse(
169+
largeResponse,
170+
"server-with-slashes",
171+
"tool-with-dashes_and_underscores",
172+
)
173+
174+
expect(result.savedToFile).toBe(true)
175+
expect(result.filePath).toContain("server-with-slashes")
176+
expect(result.filePath).toContain("tool-with-dashes_and_underscores")
177+
expect(result.content).toContain("[MCP Response saved to file")
178+
})
179+
180+
it("should limit preview to configured number of lines", async () => {
181+
const lines = Array.from({ length: 100 }, (_, i) => `Line ${i + 1}`)
182+
const response = lines.join("\n")
183+
handler = new McpResponseHandler({
184+
maxResponseSize: 100, // Low threshold to trigger save
185+
previewLines: 10, // Only show 10 lines in preview
186+
})
187+
188+
const result = await handler.processResponse(response, "server", "tool")
189+
expect(result.savedToFile).toBe(true)
190+
191+
// Check that preview contains first 10 lines
192+
expect(result.content).toContain("Line 1")
193+
expect(result.content).toContain("Line 10")
194+
expect(result.content).not.toContain("Line 11")
195+
expect(result.content).toContain("... (90 more lines)")
196+
})
197+
198+
describe("cleanupOldFiles", () => {
199+
beforeEach(() => {
200+
handler = new McpResponseHandler()
201+
})
202+
203+
it("should delete old MCP response files", async () => {
204+
const mockFiles = ["mcp-response-old-file.txt", "mcp-response-recent-file.txt", "other-file.txt"]
205+
206+
const oldDate = new Date(Date.now() - 25 * 60 * 60 * 1000) // 25 hours ago
207+
const recentDate = new Date(Date.now() - 1 * 60 * 60 * 1000) // 1 hour ago
208+
209+
const readdir = vi.mocked(fs.readdir)
210+
const stat = vi.mocked(fs.stat)
211+
const unlink = vi.mocked(fs.unlink)
212+
213+
readdir.mockResolvedValue(mockFiles as any)
214+
stat.mockImplementation(async (filePath) => {
215+
const pathStr = String(filePath)
216+
if (pathStr.includes("old-file")) {
217+
return { mtime: oldDate } as any
218+
}
219+
return { mtime: recentDate } as any
220+
})
221+
unlink.mockResolvedValue(undefined)
222+
223+
const deletedCount = await handler.cleanupOldFiles(24)
224+
225+
expect(deletedCount).toBe(1)
226+
expect(unlink).toHaveBeenCalledWith(expect.stringContaining("mcp-response-old-file.txt"))
227+
expect(unlink).not.toHaveBeenCalledWith(expect.stringContaining("recent-file"))
228+
expect(unlink).not.toHaveBeenCalledWith(expect.stringContaining("other-file"))
229+
})
230+
231+
it("should handle missing directory gracefully", async () => {
232+
const readdir = vi.mocked(fs.readdir)
233+
readdir.mockRejectedValue(new Error("ENOENT"))
234+
235+
const deletedCount = await handler.cleanupOldFiles(24)
236+
237+
expect(deletedCount).toBe(0)
238+
})
239+
})
240+
})
241+
})

0 commit comments

Comments
 (0)