Custom tools UI (#10244)

Co-authored-by: Matt Rubens <mrubens@users.noreply.github.com>
Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

Author: 3 people

packages/types/src/experiment.ts

@@ -13,6 +13,7 @@ export const experimentIds = [
 	"imageGeneration",
 	"runSlashCommand",
 	"multipleNativeToolCalls",
+	"customTools",
 ] as const
 
 export const experimentIdsSchema = z.enum(experimentIds)
@@ -30,6 +31,7 @@ export const experimentsSchema = z.object({
 	imageGeneration: z.boolean().optional(),
 	runSlashCommand: z.boolean().optional(),
 	multipleNativeToolCalls: z.boolean().optional(),
+	customTools: z.boolean().optional(),
 })
 
 export type Experiments = z.infer<typeof experimentsSchema>

src/core/assistant-message/presentAssistantMessage.ts

@@ -1081,17 +1081,15 @@ export async function presentAssistantMessage(cline: Task) {
 						break
 					}
 
-					const customTool = customToolRegistry.get(block.name)
+					const customTool = stateExperiments?.customTools ? customToolRegistry.get(block.name) : undefined
 
 					if (customTool) {
 						try {
-							console.log(`executing customTool -> ${JSON.stringify(customTool, null, 2)}`)
 							let customToolArgs
 
 							if (customTool.parameters) {
 								try {
 									customToolArgs = customTool.parameters.parse(block.nativeArgs || block.params || {})
-									console.log(`customToolArgs -> ${JSON.stringify(customToolArgs, null, 2)}`)
 								} catch (parseParamsError) {
 									const message = `Custom tool "${block.name}" argument validation failed: ${parseParamsError.message}`
 									console.error(message)
@@ -1102,13 +1100,15 @@ export async function presentAssistantMessage(cline: Task) {
 								}
 							}
 
-							console.log(`${customTool.name}.execute() -> ${JSON.stringify(customToolArgs, null, 2)}`)
-
 							const result = await customTool.execute(customToolArgs, {
 								mode: mode ?? defaultModeSlug,
 								task: cline,
 							})
 
+							console.log(
+								`${customTool.name}.execute(): ${JSON.stringify(customToolArgs)} -> ${JSON.stringify(result)}`,
+							)
+
 							pushToolResult(result)
 							cline.consecutiveMistakeCount = 0
 						} catch (executionError: any) {

src/core/prompts/system.ts

@@ -124,7 +124,7 @@ async function generatePrompt(
 
 	let customToolsSection = ""
 
-	if (!isNativeProtocol(effectiveProtocol)) {
+	if (experiments?.customTools && !isNativeProtocol(effectiveProtocol)) {
 		const customTools = customToolRegistry.getAllSerialized()
 
 		if (customTools.length > 0) {

src/core/task/build-tools.ts

@@ -79,13 +79,16 @@ export async function buildNativeToolsArray(options: BuildToolsOptions): Promise
 	const mcpTools = getMcpServerTools(mcpHub)
 	const filteredMcpTools = filterMcpToolsForMode(mcpTools, mode, customModes, experiments)
 
-	// Add custom tools if they are available.
-	await customToolRegistry.loadFromDirectoryIfStale(path.join(cwd, ".roo", "tools"))
-	const customTools = customToolRegistry.getAllSerialized()
+	// Add custom tools if they are available and the experiment is enabled.
 	let nativeCustomTools: OpenAI.Chat.ChatCompletionFunctionTool[] = []
 
-	if (customTools.length > 0) {
-		nativeCustomTools = customTools.map(formatNative)
+	if (experiments?.customTools) {
+		await customToolRegistry.loadFromDirectoryIfStale(path.join(cwd, ".roo", "tools"))
+		const customTools = customToolRegistry.getAllSerialized()
+
+		if (customTools.length > 0) {
+			nativeCustomTools = customTools.map(formatNative)
+		}
 	}
 
 	return [...filteredNativeTools, ...filteredMcpTools, ...nativeCustomTools]

src/core/tools/validateToolUse.ts

@@ -11,13 +11,13 @@ import { TOOL_GROUPS, ALWAYS_AVAILABLE_TOOLS } from "../../shared/tools"
  * Note: This does NOT check if the tool is allowed for a specific mode,
  * only that the tool actually exists.
  */
-export function isValidToolName(toolName: string): toolName is ToolName {
+export function isValidToolName(toolName: string, experiments?: Record<string, boolean>): toolName is ToolName {
 	// Check if it's a valid static tool
 	if ((validToolNames as readonly string[]).includes(toolName)) {
 		return true
 	}
 
-	if (customToolRegistry.has(toolName)) {
+	if (experiments?.customTools && customToolRegistry.has(toolName)) {
 		return true
 	}
 
@@ -40,7 +40,7 @@ export function validateToolUse(
 ): void {
 	// First, check if the tool name is actually a valid/known tool
 	// This catches completely invalid tool names like "edit_file" that don't exist
-	if (!isValidToolName(toolName)) {
+	if (!isValidToolName(toolName, experiments)) {
 		throw new Error(
 			`Unknown tool "${toolName}". This tool does not exist. Please use one of the available tools: ${validToolNames.join(", ")}.`,
 		)
@@ -94,7 +94,7 @@ export function isToolAllowedForMode(
 
 	// For now, allow all custom tools in any mode.
 	// As a follow-up we should expand the custom tool definition to include mode restrictions.
-	if (customToolRegistry.has(tool)) {
+	if (experiments?.customTools && customToolRegistry.has(tool)) {
 		return true
 	}
 

src/core/webview/webviewMessageHandler.ts

@@ -16,6 +16,7 @@ import {
 	Experiments,
 	ExperimentId,
 } from "@roo-code/types"
+import { customToolRegistry } from "@roo-code/core"
 import { CloudService } from "@roo-code/cloud"
 import { TelemetryService } from "@roo-code/telemetry"
 
@@ -1725,6 +1726,24 @@ export const webviewMessageHandler = async (
 			}
 			break
 		}
+		case "refreshCustomTools": {
+			try {
+				await customToolRegistry.loadFromDirectory(path.join(getCurrentCwd(), ".roo", "tools"))
+
+				await provider.postMessageToWebview({
+					type: "customToolsResult",
+					tools: customToolRegistry.getAllSerialized(),
+				})
+			} catch (error) {
+				await provider.postMessageToWebview({
+					type: "customToolsResult",
+					tools: [],
+					error: error instanceof Error ? error.message : String(error),
+				})
+			}
+
+			break
+		}
 		case "saveApiConfiguration":
 			if (message.text && message.apiConfiguration) {
 				try {

src/shared/ExtensionMessage.ts

@@ -14,6 +14,7 @@ import type {
 	OrganizationAllowList,
 	ShareVisibility,
 	QueuedMessage,
+	SerializedCustomToolDefinition,
 } from "@roo-code/types"
 
 import { GitCommit } from "../utils/git"
@@ -133,6 +134,7 @@ export interface ExtensionMessage {
 		| "browserSessionUpdate"
 		| "browserSessionNavigate"
 		| "claudeCodeRateLimits"
+		| "customToolsResult"
 	text?: string
 	payload?: any // Add a generic payload for now, can refine later
 	// Checkpoint warning message
@@ -218,6 +220,7 @@ export interface ExtensionMessage {
 	browserSessionMessages?: ClineMessage[] // For browser session panel updates
 	isBrowserSessionActive?: boolean // For browser session panel updates
 	stepIndex?: number // For browserSessionNavigate: the target step index to display
+	tools?: SerializedCustomToolDefinition[] // For customToolsResult
 }
 
 export type ExtensionState = Pick<

src/shared/WebviewMessage.ts

@@ -180,6 +180,7 @@ export interface WebviewMessage {
 		| "openDebugUiHistory"
 		| "downloadErrorDiagnostics"
 		| "requestClaudeCodeRateLimits"
+		| "refreshCustomTools"
 	text?: string
 	editedMessageContent?: string
 	tab?: "settings" | "history" | "mcp" | "modes" | "chat" | "marketplace" | "cloud"

src/shared/__tests__/experiments.spec.ts

@@ -32,6 +32,7 @@ describe("experiments", () => {
 				imageGeneration: false,
 				runSlashCommand: false,
 				multipleNativeToolCalls: false,
+				customTools: false,
 			}
 			expect(Experiments.isEnabled(experiments, EXPERIMENT_IDS.POWER_STEERING)).toBe(false)
 		})
@@ -44,6 +45,7 @@ describe("experiments", () => {
 				imageGeneration: false,
 				runSlashCommand: false,
 				multipleNativeToolCalls: false,
+				customTools: false,
 			}
 			expect(Experiments.isEnabled(experiments, EXPERIMENT_IDS.POWER_STEERING)).toBe(true)
 		})
@@ -56,6 +58,7 @@ describe("experiments", () => {
 				imageGeneration: false,
 				runSlashCommand: false,
 				multipleNativeToolCalls: false,
+				customTools: false,
 			}
 			expect(Experiments.isEnabled(experiments, EXPERIMENT_IDS.POWER_STEERING)).toBe(false)
 		})

src/shared/experiments.ts

@@ -7,6 +7,7 @@ export const EXPERIMENT_IDS = {
 	IMAGE_GENERATION: "imageGeneration",
 	RUN_SLASH_COMMAND: "runSlashCommand",
 	MULTIPLE_NATIVE_TOOL_CALLS: "multipleNativeToolCalls",
+	CUSTOM_TOOLS: "customTools",
 } as const satisfies Record<string, ExperimentId>
 
 type _AssertExperimentIds = AssertEqual<Equals<ExperimentId, Values<typeof EXPERIMENT_IDS>>>
@@ -24,6 +25,7 @@ export const experimentConfigsMap: Record<ExperimentKey, ExperimentConfig> = {
 	IMAGE_GENERATION: { enabled: false },
 	RUN_SLASH_COMMAND: { enabled: false },
 	MULTIPLE_NATIVE_TOOL_CALLS: { enabled: false },
+	CUSTOM_TOOLS: { enabled: false },
 }
 
 export const experimentDefault = Object.fromEntries(