Skip to content

Commit 2749c6e

Browse files
mrubensmrubens
committed
Filter tools
1 parent a840ec4 commit 2749c6e

File tree

3 files changed

+278
-26
lines changed

3 files changed

+278
-26
lines changed

‎src/core/prompts/__tests__/system-native-tools.spec.ts‎

Lines changed: 218 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,218 @@
1+
import { describe, it, expect, vi, beforeEach } from "vitest"
2+
import { SYSTEM_PROMPT } from "../system"
3+
import { CodeIndexManager } from "../../../services/code-index/manager"
4+
import type { SystemPromptSettings } from "../types"
5+
6+
vi.mock("../../../services/code-index/manager")
7+
vi.mock("../../../utils/storage", () => ({
8+
getSettingsDirectoryPath: vi.fn().mockResolvedValue("/test/settings"),
9+
}))
10+
vi.mock("../../../utils/globalContext", () => ({
11+
ensureSettingsDirectoryExists: vi.fn().mockResolvedValue("/test/settings"),
12+
}))
13+
14+
describe("SYSTEM_PROMPT with native tools", () => {
15+
const mockContext = {
16+
extensionUri: { fsPath: "/test/path" },
17+
globalStorageUri: { fsPath: "/test/global-storage" },
18+
globalState: {
19+
get: vi.fn(),
20+
update: vi.fn(),
21+
},
22+
workspaceState: {
23+
get: vi.fn(),
24+
update: vi.fn(),
25+
},
26+
} as any
27+
28+
const defaultSettings: SystemPromptSettings = {
29+
maxConcurrentFileReads: 5,
30+
todoListEnabled: true,
31+
useAgentRules: true,
32+
newTaskRequireTodos: true,
33+
}
34+
35+
beforeEach(() => {
36+
vi.clearAllMocks()
37+
})
38+
39+
it("should filter out update_todo_list when todoListEnabled is false", async () => {
40+
const mockCodeIndexManager = {
41+
isFeatureEnabled: true,
42+
isFeatureConfigured: true,
43+
isInitialized: true,
44+
}
45+
vi.mocked(CodeIndexManager.getInstance).mockReturnValue(mockCodeIndexManager as any)
46+
47+
const result = await SYSTEM_PROMPT(
48+
mockContext,
49+
"/test/cwd",
50+
false,
51+
undefined,
52+
undefined,
53+
undefined,
54+
"code",
55+
undefined,
56+
undefined,
57+
undefined,
58+
true,
59+
undefined,
60+
false,
61+
undefined,
62+
undefined,
63+
false,
64+
{ ...defaultSettings, todoListEnabled: false },
65+
undefined,
66+
undefined,
67+
true, // useNativeTools
68+
)
69+
70+
expect(result.tools).toBeDefined()
71+
const toolNames = result.tools?.map((t) => t.name) || []
72+
expect(toolNames).not.toContain("update_todo_list")
73+
})
74+
75+
it("should include update_todo_list when todoListEnabled is true", async () => {
76+
const mockCodeIndexManager = {
77+
isFeatureEnabled: true,
78+
isFeatureConfigured: true,
79+
isInitialized: true,
80+
}
81+
vi.mocked(CodeIndexManager.getInstance).mockReturnValue(mockCodeIndexManager as any)
82+
83+
const result = await SYSTEM_PROMPT(
84+
mockContext,
85+
"/test/cwd",
86+
false,
87+
undefined,
88+
undefined,
89+
undefined,
90+
"code",
91+
undefined,
92+
undefined,
93+
undefined,
94+
true,
95+
undefined,
96+
false,
97+
undefined,
98+
undefined,
99+
false,
100+
{ ...defaultSettings, todoListEnabled: true },
101+
undefined,
102+
undefined,
103+
true, // useNativeTools
104+
)
105+
106+
expect(result.tools).toBeDefined()
107+
const toolNames = result.tools?.map((t) => t.name) || []
108+
expect(toolNames).toContain("update_todo_list")
109+
})
110+
111+
it("should filter out codebase_search when feature is not configured", async () => {
112+
const mockCodeIndexManager = {
113+
isFeatureEnabled: false,
114+
isFeatureConfigured: false,
115+
isInitialized: false,
116+
}
117+
vi.mocked(CodeIndexManager.getInstance).mockReturnValue(mockCodeIndexManager as any)
118+
119+
const result = await SYSTEM_PROMPT(
120+
mockContext,
121+
"/test/cwd",
122+
false,
123+
undefined,
124+
undefined,
125+
undefined,
126+
"code",
127+
undefined,
128+
undefined,
129+
undefined,
130+
true,
131+
undefined,
132+
false,
133+
undefined,
134+
undefined,
135+
false,
136+
undefined,
137+
undefined,
138+
undefined,
139+
true, // useNativeTools
140+
)
141+
142+
expect(result.tools).toBeDefined()
143+
const toolNames = result.tools?.map((t) => t.name) || []
144+
expect(toolNames).not.toContain("codebase_search")
145+
})
146+
147+
it("should filter out generate_image when experiment is disabled", async () => {
148+
const mockCodeIndexManager = {
149+
isFeatureEnabled: true,
150+
isFeatureConfigured: true,
151+
isInitialized: true,
152+
}
153+
vi.mocked(CodeIndexManager.getInstance).mockReturnValue(mockCodeIndexManager as any)
154+
155+
const result = await SYSTEM_PROMPT(
156+
mockContext,
157+
"/test/cwd",
158+
false,
159+
undefined,
160+
undefined,
161+
undefined,
162+
"code",
163+
undefined,
164+
undefined,
165+
undefined,
166+
true,
167+
{ imageGeneration: false },
168+
false,
169+
undefined,
170+
undefined,
171+
false,
172+
undefined,
173+
undefined,
174+
undefined,
175+
true, // useNativeTools
176+
)
177+
178+
expect(result.tools).toBeDefined()
179+
const toolNames = result.tools?.map((t) => t.name) || []
180+
expect(toolNames).not.toContain("generate_image")
181+
})
182+
183+
it("should filter out run_slash_command when experiment is disabled", async () => {
184+
const mockCodeIndexManager = {
185+
isFeatureEnabled: true,
186+
isFeatureConfigured: true,
187+
isInitialized: true,
188+
}
189+
vi.mocked(CodeIndexManager.getInstance).mockReturnValue(mockCodeIndexManager as any)
190+
191+
const result = await SYSTEM_PROMPT(
192+
mockContext,
193+
"/test/cwd",
194+
false,
195+
undefined,
196+
undefined,
197+
undefined,
198+
"code",
199+
undefined,
200+
undefined,
201+
undefined,
202+
true,
203+
{ runSlashCommand: false },
204+
false,
205+
undefined,
206+
undefined,
207+
false,
208+
undefined,
209+
undefined,
210+
undefined,
211+
true, // useNativeTools
212+
)
213+
214+
expect(result.tools).toBeDefined()
215+
const toolNames = result.tools?.map((t) => t.name) || []
216+
expect(toolNames).not.toContain("run_slash_command")
217+
})
218+
})

‎src/core/prompts/system.ts‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ import { getToolSpecs } from "./tool-specs"
1717

1818
import { PromptVariables, loadSystemPromptFile } from "./sections/custom-system-prompt"
1919

20-
import { getToolDescriptionsForMode } from "./tools"
20+
import { getToolDescriptionsForMode, filterToolsByAvailability } from "./tools"
2121
import {
2222
getRulesSection,
2323
getSystemInfoSection,
@@ -252,8 +252,12 @@ ${customInstructions}`,
252252
// Remove duplicates
253253
const uniqueToolNames = [...new Set(toolNames)]
254254

255+
// Apply centralized filtering to ensure consistency with XML tool mode
256+
const codeIndexManager = CodeIndexManager.getInstance(context, cwd)
257+
const filteredTools = filterToolsByAvailability(uniqueToolNames, codeIndexManager, settings, experiments)
258+
255259
// Get tool specifications
256-
const tools = getToolSpecs(uniqueToolNames)
260+
const tools = getToolSpecs(filteredTools)
257261

258262
return {
259263
systemPrompt: basePrompt,

‎src/core/prompts/tools/index.ts‎

Lines changed: 54 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import type { ToolName, ModeConfig } from "@roo-code/types"
33
import { TOOL_GROUPS, ALWAYS_AVAILABLE_TOOLS, DiffStrategy } from "../../../shared/tools"
44
import { McpHub } from "../../../services/mcp/McpHub"
55
import { Mode, getModeConfig, isToolAllowedForMode, getGroupName } from "../../../shared/modes"
6+
import { CodeIndexManager } from "../../../services/code-index/manager"
67

78
import { ToolArgs } from "./types"
89
import { getExecuteCommandDescription } from "./execute-command"
@@ -27,7 +28,50 @@ import { getCodebaseSearchDescription } from "./codebase-search"
2728
import { getUpdateTodoListDescription } from "./update-todo-list"
2829
import { getRunSlashCommandDescription } from "./run-slash-command"
2930
import { getGenerateImageDescription } from "./generate-image"
30-
import { CodeIndexManager } from "../../../services/code-index/manager"
31+
32+
/**
33+
* Filters a list of tool names based on feature flags, settings, and experiments.
34+
* This ensures consistent tool availability across XML and native tool modes.
35+
*/
36+
export function filterToolsByAvailability(
37+
tools: ToolName[],
38+
codeIndexManager: CodeIndexManager | undefined,
39+
settings?: Record<string, any>,
40+
experiments?: Record<string, boolean>,
41+
): ToolName[] {
42+
return tools.filter((tool) => {
43+
// Conditionally exclude codebase_search if feature is disabled or not configured
44+
if (tool === "codebase_search") {
45+
if (
46+
!codeIndexManager ||
47+
!(
48+
codeIndexManager.isFeatureEnabled &&
49+
codeIndexManager.isFeatureConfigured &&
50+
codeIndexManager.isInitialized
51+
)
52+
) {
53+
return false
54+
}
55+
}
56+
57+
// Conditionally exclude update_todo_list if disabled in settings
58+
if (tool === "update_todo_list" && settings?.todoListEnabled === false) {
59+
return false
60+
}
61+
62+
// Conditionally exclude generate_image if experiment is not enabled
63+
if (tool === "generate_image" && !experiments?.imageGeneration) {
64+
return false
65+
}
66+
67+
// Conditionally exclude run_slash_command if experiment is not enabled
68+
if (tool === "run_slash_command" && !experiments?.runSlashCommand) {
69+
return false
70+
}
71+
72+
return true
73+
})
74+
}
3175

3276
// Map of tool names to their description functions
3377
const toolDescriptionMap: Record<string, (args: ToolArgs) => string | undefined> = {
@@ -120,31 +164,17 @@ export function getToolDescriptionsForMode(
120164
// Add always available tools
121165
ALWAYS_AVAILABLE_TOOLS.forEach((tool) => tools.add(tool))
122166

123-
// Conditionally exclude codebase_search if feature is disabled or not configured
124-
if (
125-
!codeIndexManager ||
126-
!(codeIndexManager.isFeatureEnabled && codeIndexManager.isFeatureConfigured && codeIndexManager.isInitialized)
127-
) {
128-
tools.delete("codebase_search")
129-
}
130-
131-
// Conditionally exclude update_todo_list if disabled in settings
132-
if (settings?.todoListEnabled === false) {
133-
tools.delete("update_todo_list")
134-
}
135-
136-
// Conditionally exclude generate_image if experiment is not enabled
137-
if (!experiments?.imageGeneration) {
138-
tools.delete("generate_image")
139-
}
140-
141-
// Conditionally exclude run_slash_command if experiment is not enabled
142-
if (!experiments?.runSlashCommand) {
143-
tools.delete("run_slash_command")
144-
}
167+
// Apply consistent filtering across all tool modes
168+
const filteredTools = filterToolsByAvailability(
169+
Array.from(tools) as ToolName[],
170+
codeIndexManager,
171+
settings,
172+
experiments,
173+
)
174+
const filteredToolsSet = new Set(filteredTools)
145175

146176
// Map tool descriptions for allowed tools
147-
const descriptions = Array.from(tools).map((toolName) => {
177+
const descriptions = Array.from(filteredToolsSet).map((toolName) => {
148178
const descriptionFn = toolDescriptionMap[toolName]
149179
if (!descriptionFn) {
150180
return undefined

0 commit comments

Comments
 (0)