Skip to content

Commit a582f9b

Browse files
jrjr
committed
Handle custom tool use similarly to MCP tools for ipc schema purposes
The tool name field is an enum, but custom tool names are dynamic. Follow the MCP tool pattern of just tracking all the custom tools in a single bucket.
1 parent 7980cd3 commit a582f9b

File tree

4 files changed

+358
-2
lines changed

4 files changed

+358
-2
lines changed

‎packages/types/src/tool.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ export const toolNames = [
3737
"update_todo_list",
3838
"run_slash_command",
3939
"generate_image",
40+
"custom_tool",
4041
] as const
4142

4243
export const toolNamesSchema = z.enum(toolNames)

‎src/core/assistant-message/__tests__/presentAssistantMessage-custom-tool.spec.ts‎

Lines changed: 349 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,349 @@
1+
// npx vitest src/core/assistant-message/__tests__/presentAssistantMessage-custom-tool.spec.ts
2+
3+
import { describe, it, expect, beforeEach, vi } from "vitest"
4+
import { presentAssistantMessage } from "../presentAssistantMessage"
5+
6+
// Mock dependencies
7+
vi.mock("../../task/Task")
8+
vi.mock("../../tools/validateToolUse", () => ({
9+
validateToolUse: vi.fn(),
10+
}))
11+
12+
// Mock custom tool registry - must be done inline without external variable references
13+
vi.mock("@roo-code/core", () => ({
14+
customToolRegistry: {
15+
has: vi.fn(),
16+
get: vi.fn(),
17+
},
18+
}))
19+
20+
vi.mock("@roo-code/telemetry", () => ({
21+
TelemetryService: {
22+
instance: {
23+
captureToolUsage: vi.fn(),
24+
captureConsecutiveMistakeError: vi.fn(),
25+
},
26+
},
27+
}))
28+
29+
import { TelemetryService } from "@roo-code/telemetry"
30+
import { customToolRegistry } from "@roo-code/core"
31+
32+
describe("presentAssistantMessage - Custom Tool Recording", () => {
33+
let mockTask: any
34+
35+
beforeEach(() => {
36+
// Reset all mocks
37+
vi.clearAllMocks()
38+
39+
// Create a mock Task with minimal properties needed for testing
40+
mockTask = {
41+
taskId: "test-task-id",
42+
instanceId: "test-instance",
43+
abort: false,
44+
presentAssistantMessageLocked: false,
45+
presentAssistantMessageHasPendingUpdates: false,
46+
currentStreamingContentIndex: 0,
47+
assistantMessageContent: [],
48+
userMessageContent: [],
49+
didCompleteReadingStream: false,
50+
didRejectTool: false,
51+
didAlreadyUseTool: false,
52+
diffEnabled: false,
53+
consecutiveMistakeCount: 0,
54+
clineMessages: [],
55+
api: {
56+
getModel: () => ({ id: "test-model", info: {} }),
57+
},
58+
browserSession: {
59+
closeBrowser: vi.fn().mockResolvedValue(undefined),
60+
},
61+
recordToolUsage: vi.fn(),
62+
recordToolError: vi.fn(),
63+
toolRepetitionDetector: {
64+
check: vi.fn().mockReturnValue({ allowExecution: true }),
65+
},
66+
providerRef: {
67+
deref: () => ({
68+
getState: vi.fn().mockResolvedValue({
69+
mode: "code",
70+
customModes: [],
71+
experiments: {
72+
customTools: true, // Enable by default
73+
},
74+
}),
75+
}),
76+
},
77+
say: vi.fn().mockResolvedValue(undefined),
78+
ask: vi.fn().mockResolvedValue({ response: "yesButtonClicked" }),
79+
}
80+
})
81+
82+
describe("Custom tool usage recording", () => {
83+
it("should record custom tool usage as 'custom_tool' when experiment is enabled", async () => {
84+
const toolCallId = "tool_call_custom_123"
85+
mockTask.assistantMessageContent = [
86+
{
87+
type: "tool_use",
88+
id: toolCallId,
89+
name: "my_custom_tool",
90+
params: { value: "test" },
91+
partial: false,
92+
},
93+
]
94+
95+
// Mock customToolRegistry to recognize this as a custom tool
96+
vi.mocked(customToolRegistry.has).mockReturnValue(true)
97+
vi.mocked(customToolRegistry.get).mockReturnValue({
98+
name: "my_custom_tool",
99+
description: "A custom tool",
100+
execute: vi.fn().mockResolvedValue("Custom tool result"),
101+
})
102+
103+
await presentAssistantMessage(mockTask)
104+
105+
// Should record as "custom_tool", not "my_custom_tool"
106+
expect(mockTask.recordToolUsage).toHaveBeenCalledWith("custom_tool")
107+
expect(TelemetryService.instance.captureToolUsage).toHaveBeenCalledWith(
108+
mockTask.taskId,
109+
"custom_tool",
110+
"native",
111+
)
112+
})
113+
114+
it("should record custom tool usage as 'custom_tool' in XML protocol", async () => {
115+
mockTask.assistantMessageContent = [
116+
{
117+
type: "tool_use",
118+
// No ID = XML protocol
119+
name: "my_custom_tool",
120+
params: { value: "test" },
121+
partial: false,
122+
},
123+
]
124+
125+
vi.mocked(customToolRegistry.has).mockReturnValue(true)
126+
vi.mocked(customToolRegistry.get).mockReturnValue({
127+
name: "my_custom_tool",
128+
description: "A custom tool",
129+
execute: vi.fn().mockResolvedValue("Custom tool result"),
130+
})
131+
132+
await presentAssistantMessage(mockTask)
133+
134+
expect(mockTask.recordToolUsage).toHaveBeenCalledWith("custom_tool")
135+
expect(TelemetryService.instance.captureToolUsage).toHaveBeenCalledWith(
136+
mockTask.taskId,
137+
"custom_tool",
138+
"xml",
139+
)
140+
})
141+
})
142+
143+
describe("Custom tool error recording", () => {
144+
it("should record custom tool error as 'custom_tool'", async () => {
145+
const toolCallId = "tool_call_custom_error_123"
146+
mockTask.assistantMessageContent = [
147+
{
148+
type: "tool_use",
149+
id: toolCallId,
150+
name: "failing_custom_tool",
151+
params: {},
152+
partial: false,
153+
},
154+
]
155+
156+
// Mock customToolRegistry with a tool that throws an error
157+
vi.mocked(customToolRegistry.has).mockReturnValue(true)
158+
vi.mocked(customToolRegistry.get).mockReturnValue({
159+
name: "failing_custom_tool",
160+
description: "A failing custom tool",
161+
execute: vi.fn().mockRejectedValue(new Error("Custom tool execution failed")),
162+
})
163+
164+
await presentAssistantMessage(mockTask)
165+
166+
// Should record error as "custom_tool", not "failing_custom_tool"
167+
expect(mockTask.recordToolError).toHaveBeenCalledWith("custom_tool", "Custom tool execution failed")
168+
expect(mockTask.consecutiveMistakeCount).toBe(1)
169+
})
170+
})
171+
172+
describe("Regular tool recording", () => {
173+
it("should record regular tool usage with actual tool name", async () => {
174+
const toolCallId = "tool_call_read_file_123"
175+
mockTask.assistantMessageContent = [
176+
{
177+
type: "tool_use",
178+
id: toolCallId,
179+
name: "read_file",
180+
params: { path: "test.txt" },
181+
partial: false,
182+
},
183+
]
184+
185+
// read_file is not a custom tool
186+
vi.mocked(customToolRegistry.has).mockReturnValue(false)
187+
188+
await presentAssistantMessage(mockTask)
189+
190+
// Should record as "read_file", not "custom_tool"
191+
expect(mockTask.recordToolUsage).toHaveBeenCalledWith("read_file")
192+
expect(TelemetryService.instance.captureToolUsage).toHaveBeenCalledWith(
193+
mockTask.taskId,
194+
"read_file",
195+
"native",
196+
)
197+
})
198+
199+
it("should record MCP tool usage as 'use_mcp_tool' (not custom_tool)", async () => {
200+
const toolCallId = "tool_call_mcp_123"
201+
mockTask.assistantMessageContent = [
202+
{
203+
type: "tool_use",
204+
id: toolCallId,
205+
name: "use_mcp_tool",
206+
params: {
207+
server_name: "test-server",
208+
tool_name: "test-tool",
209+
arguments: "{}",
210+
},
211+
partial: false,
212+
},
213+
]
214+
215+
vi.mocked(customToolRegistry.has).mockReturnValue(false)
216+
217+
// Mock MCP hub for use_mcp_tool
218+
mockTask.providerRef = {
219+
deref: () => ({
220+
getState: vi.fn().mockResolvedValue({
221+
mode: "code",
222+
customModes: [],
223+
experiments: {
224+
customTools: true,
225+
},
226+
}),
227+
getMcpHub: () => ({
228+
findServerNameBySanitizedName: () => "test-server",
229+
executeToolCall: vi.fn().mockResolvedValue({ content: [{ type: "text", text: "result" }] }),
230+
}),
231+
}),
232+
}
233+
234+
await presentAssistantMessage(mockTask)
235+
236+
// Should record as "use_mcp_tool", not "custom_tool"
237+
expect(mockTask.recordToolUsage).toHaveBeenCalledWith("use_mcp_tool")
238+
expect(TelemetryService.instance.captureToolUsage).toHaveBeenCalledWith(
239+
mockTask.taskId,
240+
"use_mcp_tool",
241+
"native",
242+
)
243+
})
244+
})
245+
246+
describe("Custom tool experiment gate", () => {
247+
it("should treat custom tool as unknown when experiment is disabled", async () => {
248+
const toolCallId = "tool_call_disabled_123"
249+
mockTask.assistantMessageContent = [
250+
{
251+
type: "tool_use",
252+
id: toolCallId,
253+
name: "my_custom_tool",
254+
params: {},
255+
partial: false,
256+
},
257+
]
258+
259+
// Mock provider state with customTools experiment DISABLED
260+
mockTask.providerRef = {
261+
deref: () => ({
262+
getState: vi.fn().mockResolvedValue({
263+
mode: "code",
264+
customModes: [],
265+
experiments: {
266+
customTools: false, // Disabled
267+
},
268+
}),
269+
}),
270+
}
271+
272+
// Even if registry recognizes it, experiment gate should prevent execution
273+
vi.mocked(customToolRegistry.has).mockReturnValue(true)
274+
vi.mocked(customToolRegistry.get).mockReturnValue({
275+
name: "my_custom_tool",
276+
description: "A custom tool",
277+
execute: vi.fn().mockResolvedValue("Should not execute"),
278+
})
279+
280+
await presentAssistantMessage(mockTask)
281+
282+
// Should be treated as unknown tool (not executed)
283+
expect(mockTask.say).toHaveBeenCalledWith("error", "unknownToolError")
284+
expect(mockTask.consecutiveMistakeCount).toBe(1)
285+
286+
// Custom tool should NOT have been executed
287+
const getMock = vi.mocked(customToolRegistry.get)
288+
if (getMock.mock.results.length > 0) {
289+
const customTool = getMock.mock.results[0].value
290+
if (customTool) {
291+
expect(customTool.execute).not.toHaveBeenCalled()
292+
}
293+
}
294+
})
295+
296+
it("should not call customToolRegistry.has() when experiment is disabled", async () => {
297+
mockTask.assistantMessageContent = [
298+
{
299+
type: "tool_use",
300+
id: "tool_call_123",
301+
name: "some_tool",
302+
params: {},
303+
partial: false,
304+
},
305+
]
306+
307+
// Disable experiment
308+
mockTask.providerRef = {
309+
deref: () => ({
310+
getState: vi.fn().mockResolvedValue({
311+
mode: "code",
312+
customModes: [],
313+
experiments: {
314+
customTools: false,
315+
},
316+
}),
317+
}),
318+
}
319+
320+
await presentAssistantMessage(mockTask)
321+
322+
// When experiment is off, shouldn't even check the registry
323+
// (Code checks stateExperiments?.customTools before calling has())
324+
expect(customToolRegistry.has).not.toHaveBeenCalled()
325+
})
326+
})
327+
328+
describe("Partial blocks", () => {
329+
it("should not record usage for partial custom tool blocks", async () => {
330+
mockTask.assistantMessageContent = [
331+
{
332+
type: "tool_use",
333+
id: "tool_call_partial_123",
334+
name: "my_custom_tool",
335+
params: { value: "test" },
336+
partial: true, // Still streaming
337+
},
338+
]
339+
340+
vi.mocked(customToolRegistry.has).mockReturnValue(true)
341+
342+
await presentAssistantMessage(mockTask)
343+
344+
// Should not record usage for partial blocks
345+
expect(mockTask.recordToolUsage).not.toHaveBeenCalled()
346+
expect(TelemetryService.instance.captureToolUsage).not.toHaveBeenCalled()
347+
})
348+
})
349+
})

‎src/core/assistant-message/presentAssistantMessage.ts‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -695,8 +695,11 @@ export async function presentAssistantMessage(cline: Task) {
695695
}
696696

697697
if (!block.partial) {
698-
cline.recordToolUsage(block.name)
699-
TelemetryService.instance.captureToolUsage(cline.taskId, block.name, toolProtocol)
698+
// Check if this is a custom tool - if so, record as "custom_tool" (like MCP tools)
699+
const isCustomTool = stateExperiments?.customTools && customToolRegistry.has(block.name)
700+
const recordName = isCustomTool ? "custom_tool" : block.name
701+
cline.recordToolUsage(recordName)
702+
TelemetryService.instance.captureToolUsage(cline.taskId, recordName, toolProtocol)
700703
}
701704

702705
// Validate tool use before execution - ONLY for complete (non-partial) blocks.
@@ -1091,6 +1094,8 @@ export async function presentAssistantMessage(cline: Task) {
10911094
cline.consecutiveMistakeCount = 0
10921095
} catch (executionError: any) {
10931096
cline.consecutiveMistakeCount++
1097+
// Record custom tool error with static name
1098+
cline.recordToolError("custom_tool", executionError.message)
10941099
await handleError(`executing custom tool "${block.name}"`, executionError)
10951100
}
10961101

‎src/shared/tools.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -266,6 +266,7 @@ export const TOOL_DISPLAY_NAMES: Record<ToolName, string> = {
266266
update_todo_list: "update todo list",
267267
run_slash_command: "run slash command",
268268
generate_image: "generate images",
269+
custom_tool: "use custom tools",
269270
} as const
270271

271272
// Define available tool groups.

0 commit comments

Comments
 (0)