Skip to content

Commit a9c1dad

Browse files
samhvw8daniel-lxs
andauthored
feat(mcp): improve mcp with McpExecution, make UX like terminal execute (RooCodeInc#3485)
* feat: add MCP viewer and enhanced UI - Add McpExecutionStatus type with discriminated union for execution states - Implement McpExecution component with live status updates and collapsible output - Replace inline MCP UI in ChatRow with dedicated McpExecution component - Add comprehensive test coverage for useMcpToolTool - Enhance CodeAccordian to support custom headers - Improve combineCommandSequences to handle MCP execution flows * refactor: merge loops in combineCommandSequences for better efficiency - Reduced from 3 separate loops to 1 main processing loop - Improved time complexity from O(n²) to O(n) - Reduced total passes through array from 4 to 2 - Better memory access patterns and cache utilization * feat: enhance JSON parsing and rendering in McpExecution component * feat: add mock implementations for react-markdown and remark-gfm * fix: remove unreachable return statement in ChatRowContent component * feat(i18n): add execution status messages and error handling for invalid JSON arguments * test: add parameter validation tests for missing server_name and tool_name * Update McpExecution component to show server name in header and tool name in approval section --------- Co-authored-by: Daniel Riccio <ricciodaniel98@gmail.com>
1 parent 1b1e5a2 commit a9c1dad

47 files changed

Lines changed: 1241 additions & 267 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

packages/types/src/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ export * from "./experiment.js"
77
export * from "./global-settings.js"
88
export * from "./history.js"
99
export * from "./ipc.js"
10+
export * from "./mcp.js"
1011
export * from "./message.js"
1112
export * from "./mode.js"
1213
export * from "./model.js"

packages/types/src/mcp.ts

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
import { z } from "zod"
2+
3+
/**
4+
* McpExecutionStatus
5+
*/
6+
7+
export const mcpExecutionStatusSchema = z.discriminatedUnion("status", [
8+
z.object({
9+
executionId: z.string(),
10+
status: z.literal("started"),
11+
serverName: z.string(),
12+
toolName: z.string(),
13+
}),
14+
z.object({
15+
executionId: z.string(),
16+
status: z.literal("output"),
17+
response: z.string(),
18+
}),
19+
z.object({
20+
executionId: z.string(),
21+
status: z.literal("completed"),
22+
response: z.string().optional(),
23+
}),
24+
z.object({
25+
executionId: z.string(),
26+
status: z.literal("error"),
27+
error: z.string().optional(),
28+
}),
29+
])
30+
31+
export type McpExecutionStatus = z.infer<typeof mcpExecutionStatusSchema>
Lines changed: 269 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,269 @@
1+
import { useMcpToolTool } from "../useMcpToolTool"
2+
import { Task } from "../../task/Task"
3+
import { ToolUse } from "../../../shared/tools"
4+
import { formatResponse } from "../../prompts/responses"
5+
6+
// Mock dependencies
7+
jest.mock("../../prompts/responses", () => ({
8+
formatResponse: {
9+
toolResult: jest.fn((result: string) => `Tool result: ${result}`),
10+
toolError: jest.fn((error: string) => `Tool error: ${error}`),
11+
invalidMcpToolArgumentError: jest.fn((server: string, tool: string) => `Invalid args for ${server}:${tool}`),
12+
},
13+
}))
14+
15+
jest.mock("../../../i18n", () => ({
16+
t: jest.fn((key: string, params?: any) => {
17+
if (key === "mcp:errors.invalidJsonArgument" && params?.toolName) {
18+
return `Roo tried to use ${params.toolName} with an invalid JSON argument. Retrying...`
19+
}
20+
return key
21+
}),
22+
}))
23+
24+
describe("useMcpToolTool", () => {
25+
let mockTask: Partial<Task>
26+
let mockAskApproval: jest.Mock
27+
let mockHandleError: jest.Mock
28+
let mockPushToolResult: jest.Mock
29+
let mockRemoveClosingTag: jest.Mock
30+
let mockProviderRef: any
31+
32+
beforeEach(() => {
33+
mockAskApproval = jest.fn()
34+
mockHandleError = jest.fn()
35+
mockPushToolResult = jest.fn()
36+
mockRemoveClosingTag = jest.fn((tag: string, value?: string) => value || "")
37+
38+
mockProviderRef = {
39+
deref: jest.fn().mockReturnValue({
40+
getMcpHub: jest.fn().mockReturnValue({
41+
callTool: jest.fn(),
42+
}),
43+
postMessageToWebview: jest.fn(),
44+
}),
45+
}
46+
47+
mockTask = {
48+
consecutiveMistakeCount: 0,
49+
recordToolError: jest.fn(),
50+
sayAndCreateMissingParamError: jest.fn(),
51+
say: jest.fn(),
52+
ask: jest.fn(),
53+
lastMessageTs: 123456789,
54+
providerRef: mockProviderRef,
55+
}
56+
})
57+
58+
describe("parameter validation", () => {
59+
it("should handle missing server_name", async () => {
60+
const block: ToolUse = {
61+
type: "tool_use",
62+
name: "use_mcp_tool",
63+
params: {
64+
tool_name: "test_tool",
65+
arguments: "{}",
66+
},
67+
partial: false,
68+
}
69+
70+
mockTask.sayAndCreateMissingParamError = jest.fn().mockResolvedValue("Missing server_name error")
71+
72+
await useMcpToolTool(
73+
mockTask as Task,
74+
block,
75+
mockAskApproval,
76+
mockHandleError,
77+
mockPushToolResult,
78+
mockRemoveClosingTag,
79+
)
80+
81+
expect(mockTask.consecutiveMistakeCount).toBe(1)
82+
expect(mockTask.recordToolError).toHaveBeenCalledWith("use_mcp_tool")
83+
expect(mockTask.sayAndCreateMissingParamError).toHaveBeenCalledWith("use_mcp_tool", "server_name")
84+
expect(mockPushToolResult).toHaveBeenCalledWith("Missing server_name error")
85+
})
86+
87+
it("should handle missing tool_name", async () => {
88+
const block: ToolUse = {
89+
type: "tool_use",
90+
name: "use_mcp_tool",
91+
params: {
92+
server_name: "test_server",
93+
arguments: "{}",
94+
},
95+
partial: false,
96+
}
97+
98+
mockTask.sayAndCreateMissingParamError = jest.fn().mockResolvedValue("Missing tool_name error")
99+
100+
await useMcpToolTool(
101+
mockTask as Task,
102+
block,
103+
mockAskApproval,
104+
mockHandleError,
105+
mockPushToolResult,
106+
mockRemoveClosingTag,
107+
)
108+
109+
expect(mockTask.consecutiveMistakeCount).toBe(1)
110+
expect(mockTask.recordToolError).toHaveBeenCalledWith("use_mcp_tool")
111+
expect(mockTask.sayAndCreateMissingParamError).toHaveBeenCalledWith("use_mcp_tool", "tool_name")
112+
expect(mockPushToolResult).toHaveBeenCalledWith("Missing tool_name error")
113+
})
114+
115+
it("should handle invalid JSON arguments", async () => {
116+
const block: ToolUse = {
117+
type: "tool_use",
118+
name: "use_mcp_tool",
119+
params: {
120+
server_name: "test_server",
121+
tool_name: "test_tool",
122+
arguments: "invalid json",
123+
},
124+
partial: false,
125+
}
126+
127+
await useMcpToolTool(
128+
mockTask as Task,
129+
block,
130+
mockAskApproval,
131+
mockHandleError,
132+
mockPushToolResult,
133+
mockRemoveClosingTag,
134+
)
135+
136+
expect(mockTask.consecutiveMistakeCount).toBe(1)
137+
expect(mockTask.recordToolError).toHaveBeenCalledWith("use_mcp_tool")
138+
expect(mockTask.say).toHaveBeenCalledWith("error", expect.stringContaining("invalid JSON argument"))
139+
expect(mockPushToolResult).toHaveBeenCalledWith("Tool error: Invalid args for test_server:test_tool")
140+
})
141+
})
142+
143+
describe("partial requests", () => {
144+
it("should handle partial requests", async () => {
145+
const block: ToolUse = {
146+
type: "tool_use",
147+
name: "use_mcp_tool",
148+
params: {
149+
server_name: "test_server",
150+
tool_name: "test_tool",
151+
arguments: "{}",
152+
},
153+
partial: true,
154+
}
155+
156+
mockTask.ask = jest.fn().mockResolvedValue(true)
157+
158+
await useMcpToolTool(
159+
mockTask as Task,
160+
block,
161+
mockAskApproval,
162+
mockHandleError,
163+
mockPushToolResult,
164+
mockRemoveClosingTag,
165+
)
166+
167+
expect(mockTask.ask).toHaveBeenCalledWith("use_mcp_server", expect.stringContaining("use_mcp_tool"), true)
168+
})
169+
})
170+
171+
describe("successful execution", () => {
172+
it("should execute tool successfully with valid parameters", async () => {
173+
const block: ToolUse = {
174+
type: "tool_use",
175+
name: "use_mcp_tool",
176+
params: {
177+
server_name: "test_server",
178+
tool_name: "test_tool",
179+
arguments: '{"param": "value"}',
180+
},
181+
partial: false,
182+
}
183+
184+
mockAskApproval.mockResolvedValue(true)
185+
186+
const mockToolResult = {
187+
content: [{ type: "text", text: "Tool executed successfully" }],
188+
isError: false,
189+
}
190+
191+
mockProviderRef.deref.mockReturnValue({
192+
getMcpHub: () => ({
193+
callTool: jest.fn().mockResolvedValue(mockToolResult),
194+
}),
195+
postMessageToWebview: jest.fn(),
196+
})
197+
198+
await useMcpToolTool(
199+
mockTask as Task,
200+
block,
201+
mockAskApproval,
202+
mockHandleError,
203+
mockPushToolResult,
204+
mockRemoveClosingTag,
205+
)
206+
207+
expect(mockTask.consecutiveMistakeCount).toBe(0)
208+
expect(mockAskApproval).toHaveBeenCalled()
209+
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_request_started")
210+
expect(mockTask.say).toHaveBeenCalledWith("mcp_server_response", "Tool executed successfully")
211+
expect(mockPushToolResult).toHaveBeenCalledWith("Tool result: Tool executed successfully")
212+
})
213+
214+
it("should handle user rejection", async () => {
215+
const block: ToolUse = {
216+
type: "tool_use",
217+
name: "use_mcp_tool",
218+
params: {
219+
server_name: "test_server",
220+
tool_name: "test_tool",
221+
arguments: "{}",
222+
},
223+
partial: false,
224+
}
225+
226+
mockAskApproval.mockResolvedValue(false)
227+
228+
await useMcpToolTool(
229+
mockTask as Task,
230+
block,
231+
mockAskApproval,
232+
mockHandleError,
233+
mockPushToolResult,
234+
mockRemoveClosingTag,
235+
)
236+
237+
expect(mockTask.say).not.toHaveBeenCalledWith("mcp_server_request_started")
238+
expect(mockPushToolResult).not.toHaveBeenCalled()
239+
})
240+
})
241+
242+
describe("error handling", () => {
243+
it("should handle unexpected errors", async () => {
244+
const block: ToolUse = {
245+
type: "tool_use",
246+
name: "use_mcp_tool",
247+
params: {
248+
server_name: "test_server",
249+
tool_name: "test_tool",
250+
},
251+
partial: false,
252+
}
253+
254+
const error = new Error("Unexpected error")
255+
mockAskApproval.mockRejectedValue(error)
256+
257+
await useMcpToolTool(
258+
mockTask as Task,
259+
block,
260+
mockAskApproval,
261+
mockHandleError,
262+
mockPushToolResult,
263+
mockRemoveClosingTag,
264+
)
265+
266+
expect(mockHandleError).toHaveBeenCalledWith("executing MCP tool", error)
267+
})
268+
})
269+
})

0 commit comments

Comments
 (0)