Merge pull request #3261 from mcowger/mcowger/smallToolFixes

fix(tools): improve tool calling consistency

Author: chrarnoldus

.changeset/grumpy-snails-walk.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Improve native tool calling consistency

src/core/prompts/sections/markdown-formatting.ts

@@ -1,7 +1,11 @@
-export function markdownFormattingSection(): string {
+import { ToolUseStyle } from "../../../../packages/types/src/kilocode/native-function-calling"
+
+export function markdownFormattingSection(
+	toolUseStyle: ToolUseStyle, // kilocode_change
+): string {
 	return `====
 
 MARKDOWN RULES
 
-ALL responses MUST show ANY \`language construct\` OR filename reference as clickable, exactly as [\`filename OR language.declaration()\`](relative/file/path.ext:line); line is required for \`syntax\` and optional for filename links. This applies to ALL markdown responses and ALSO those in <attempt_completion>`
+ALL responses MUST show ANY \`language construct\` OR filename reference as clickable, exactly as [\`filename OR language.declaration()\`](relative/file/path.ext:line); line is required for \`syntax\` and optional for filename links. This applies to ALL markdown responses and ALSO those in ${toolUseStyle === "json" ? "attempt_completion" : "<attempt_completion>" /*kilocode_change*/}`
 }

src/core/prompts/sections/modes.ts

@@ -38,7 +38,7 @@ ${allModes
 
 	// kilocode_change: toolUseStyle
 	modesContent += `
-If the user asks you to create or edit a new mode for this project, you should read the instructions${toolUseStyle !== "json" ? " by using the fetch_instructions tool, like this:\n<fetch_instructions>\n<task>create_mode</task>\n</fetch_instructions>" : "."}`
+If the user asks you to create or edit a new mode for this project, you should read the instructions by using the fetch_instructions tool${toolUseStyle !== "json" ? ", like this:\n<fetch_instructions>\n<task>create_mode</task>\n</fetch_instructions>" : "."}`
 
 	return modesContent
 }

src/core/prompts/sections/rules.ts

@@ -5,6 +5,7 @@ import { CodeIndexManager } from "../../../services/code-index/manager"
 import { getFastApplyEditingInstructions } from "../tools/edit-file"
 import { type ClineProviderState } from "../../webview/ClineProvider"
 import { getFastApplyModelType, isFastApplyAvailable } from "../../tools/editFileTool"
+import { ToolUseStyle } from "../../../../packages/types/src/kilocode/native-function-calling"
 // kilocode_change end
 
 function getEditingInstructions(diffStrategy?: DiffStrategy): string {
@@ -57,6 +58,7 @@ export function getRulesSection(
 	diffStrategy?: DiffStrategy,
 	codeIndexManager?: CodeIndexManager,
 	clineProviderState?: ClineProviderState, // kilocode_change
+	toolUseStyle?: ToolUseStyle, // kilocode_change
 ): string {
 	const isCodebaseSearchAvailable =
 		codeIndexManager &&
@@ -70,12 +72,15 @@ export function getRulesSection(
 
 	const kiloCodeUseMorph = isFastApplyAvailable(clineProviderState)
 
-	return `====
+	let rulesContent = ""
+	rulesContent += `====
 
 RULES
 
 - The project base directory is: ${cwd.toPosix()}
-- All file paths must be relative to this directory. However, commands may change directories in terminals, so respect working directory specified by the response to <execute_command>.
+`
+	rulesContent += `- All file paths must be relative to this directory. However, commands may change directories in terminals, so respect working directory specified by the response to ${toolUseStyle === "json" ? '"execute_command"' : "<execute_command>" /*kilocode_change*/}.`
+	rulesContent += `
 - You cannot \`cd\` into a different directory to complete a task. You are stuck operating from '${cwd.toPosix()}', so be sure to pass in the correct 'path' parameter when using tools that require a path.
 - Do not use the ~ character or $HOME to refer to the home directory.
 - Before using the execute_command tool, you must first think about the SYSTEM INFORMATION context provided to understand the user's environment and tailor your commands to ensure they are compatible with their system. You must also consider if the command you need to run should be executed in a specific directory outside of the current working directory '${cwd.toPosix()}', and if so prepend with \`cd\`'ing into that directory && then executing the command (as one command since you are stuck operating from '${cwd.toPosix()}'). For example, if you needed to run \`npm install\` in a project outside of '${cwd.toPosix()}', you would need to prepend with a \`cd\` i.e. pseudocode for this would be \`cd (path to project) && (command, in this case npm install)\`.
@@ -106,4 +111,5 @@ ${kiloCodeUseMorph ? getFastApplyEditingInstructions(getFastApplyModelType(cline
 			? " Then if you want to test your work, you might use browser_action to launch the site, wait for the user's response confirming the site was launched along with a screenshot, then perhaps e.g., click a button to test functionality if needed, wait for the user's response confirming the button was clicked along with a screenshot of the new state, before finally closing the browser."
 			: ""
 	}`
+	return rulesContent
 }

src/core/prompts/sections/tool-use.ts

@@ -1,14 +1,17 @@
 import { ToolUseStyle } from "../../../../packages/types/src" // kilocode_change
 
-const basicInfo = `====
+export function getSharedToolUseSection(
+	toolUseStyle?: ToolUseStyle, // kilocode_change
+): string {
+	return `====
 
 TOOL USE
 
-You have access to a set of tools that are executed upon the user's approval. You must use exactly one tool per message, and every assistant message must include a tool call. You use tools step-by-step to accomplish a given task, with each tool use informed by the result of the previous tool use.
-`
+You have access to a set of tools that are executed upon the user's approval. You ${toolUseStyle === "json" ? "MUST USE" : "must use" /*kilocode_change*/} exactly one tool per message, and ${toolUseStyle === "json" ? "EVERY" : "every" /*kilocode_change*/} assistant message ${toolUseStyle === "json" ? "MUST" : "must" /*kilocode_change*/} include a tool call. You use tools step-by-step to accomplish a given task, with each tool use informed by the result of the previous tool use.${
+		toolUseStyle === "json" // kilocode_change
+			? ""
+			: `
 
-// kilocode_change: xml-specifix instruction split off
-const xmlInfo = `
 # Tool Use Formatting
 
 Tool uses are formatted using XML-style tags. The tool name itself becomes the XML tag name. Each parameter is enclosed within its own set of tags. Here's the structure:
@@ -20,11 +23,5 @@ Tool uses are formatted using XML-style tags. The tool name itself becomes the X
 </actual_tool_name>
 
 Always use the actual tool name as the XML tag name for proper parsing and execution.`
-
-export function getSharedToolUseSection(toolUseStyle?: ToolUseStyle): string {
-	if (toolUseStyle === "json") {
-		return `${basicInfo}`
-	} else {
-		return `${basicInfo}${xmlInfo}`
-	}
+	}`
 }

src/core/prompts/system.ts

@@ -99,7 +99,7 @@ async function generatePrompt(
 
 	const basePrompt = `${roleDefinition}
 
-${markdownFormattingSection()}
+${markdownFormattingSection(toolUseStyle ?? "xml" /*kilocode_change*/)}
 
 ${getSharedToolUseSection(toolUseStyle /*kilocode_change*/)}
 
@@ -132,7 +132,7 @@ ${getCapabilitiesSection(cwd, supportsComputerUse, shouldIncludeMcp ? mcpHub : u
 
 ${modesSection}
 
-${getRulesSection(cwd, supportsComputerUse, effectiveDiffStrategy, codeIndexManager, clineProviderState /* kilocode_change */)}
+${getRulesSection(cwd, supportsComputerUse, effectiveDiffStrategy, codeIndexManager, clineProviderState, toolUseStyle /* kilocode_change */)}
 
 ${getSystemInfoSection(cwd)}
 

src/core/prompts/tools/native-tools/__tests__/getAllowedJSONToolsForMode.spec.ts

@@ -0,0 +1,70 @@
+import { describe, it, expect } from "vitest"
+import { getAllowedJSONToolsForMode } from "../getAllowedJSONToolsForMode"
+import { Mode } from "../../../../../shared/modes"
+import { ClineProviderState } from "../../../../webview/ClineProvider"
+import { apply_diff_multi_file, apply_diff_single_file } from "../apply_diff"
+
+describe("getAllowedJSONToolsForMode", () => {
+	const mockCodeIndexManager = {
+		isFeatureEnabled: true,
+		isFeatureConfigured: true,
+		isInitialized: true,
+	} as any
+
+	const baseProviderState: Partial<ClineProviderState> = {
+		apiConfiguration: {
+			diffEnabled: true,
+		},
+		experiments: {},
+	}
+
+	it("should use single file diff when multiFileApplyDiff experiment is disabled", () => {
+		const providerState: Partial<ClineProviderState> = {
+			...baseProviderState,
+			apiConfiguration: {
+				diffEnabled: true,
+			},
+			experiments: {
+				multiFileApplyDiff: false,
+			},
+		}
+
+		const tools = getAllowedJSONToolsForMode(
+			"code" as Mode,
+			mockCodeIndexManager,
+			providerState as ClineProviderState,
+			true,
+			false,
+		)
+
+		const applyDiffTool = tools.find((tool) => "function" in tool && tool.function.name === "apply_diff")
+		expect(applyDiffTool).toBeDefined()
+		expect(applyDiffTool).toEqual(apply_diff_single_file)
+		expect(applyDiffTool).not.toEqual(apply_diff_multi_file)
+	})
+
+	it("should use multi file diff when multiFileApplyDiff experiment is enabled", () => {
+		const providerState: Partial<ClineProviderState> = {
+			...baseProviderState,
+			apiConfiguration: {
+				diffEnabled: true,
+			},
+			experiments: {
+				multiFileApplyDiff: true,
+			},
+		}
+
+		const tools = getAllowedJSONToolsForMode(
+			"code" as Mode,
+			mockCodeIndexManager,
+			providerState as ClineProviderState,
+			true,
+			false,
+		)
+
+		const applyDiffTool = tools.find((tool) => "function" in tool && tool.function.name === "apply_diff")
+		expect(applyDiffTool).toBeDefined()
+		expect(applyDiffTool).toEqual(apply_diff_multi_file)
+		expect(applyDiffTool).not.toEqual(apply_diff_single_file)
+	})
+})

src/core/prompts/tools/native-tools/getAllowedJSONToolsForMode.ts

@@ -10,9 +10,10 @@ import { apply_diff_multi_file, apply_diff_single_file } from "./apply_diff"
 
 export function getAllowedJSONToolsForMode(
 	mode: Mode,
-	codeIndexManager?: CodeIndexManager,
-	clineProviderState?: ClineProviderState,
-	supportsImages?: boolean,
+	codeIndexManager: CodeIndexManager | undefined,
+	clineProviderState: ClineProviderState | undefined,
+	diffEnabled: boolean,
+	supportsImages: boolean,
 ): OpenAI.Chat.ChatCompletionTool[] {
 	const config = getModeConfig(mode, clineProviderState?.customModes)
 
@@ -79,27 +80,32 @@ export function getAllowedJSONToolsForMode(
 	}
 
 	// Create a map of tool names to native tool definitions for quick lookup
-	const nativeToolsMap = new Map<string, OpenAI.Chat.ChatCompletionTool>()
-	nativeTools.forEach((tool) => {
-		nativeToolsMap.set(tool.function.name, tool)
-	})
+	// Exclude apply_diff tools as they are handled specially below
+	const allowedTools: OpenAI.Chat.ChatCompletionTool[] = []
 
-	if (clineProviderState?.apiConfiguration.diffEnabled) {
-		if (clineProviderState?.experiments.multiFileApplyDiff) {
-			nativeToolsMap.set("apply_diff", apply_diff_multi_file)
-		} else {
-			nativeToolsMap.set("apply_diff", apply_diff_single_file)
+	let isApplyDiffToolAllowedForMode = false
+	for (const nativeTool of nativeTools) {
+		const toolName = nativeTool.function.name
+
+		// If the tool is in the allowed set, add it.
+		if (tools.has(toolName)) {
+			if (toolName === "apply_diff") {
+				isApplyDiffToolAllowedForMode = true
+			} else {
+				allowedTools.push(nativeTool)
+			}
 		}
 	}
 
-	// Map allowed tools to their native definitions
-	const allowedTools: OpenAI.Chat.ChatCompletionTool[] = []
-	tools.forEach((toolName) => {
-		const nativeTool = nativeToolsMap.get(toolName)
-		if (nativeTool) {
-			allowedTools.push(nativeTool)
+	// Handle the "apply_diff" logic separately because the same tool has different
+	// implementations depending on whether multi-file diffs are enabled, but the same name is used.
+	if (isApplyDiffToolAllowedForMode && diffEnabled) {
+		if (clineProviderState?.experiments.multiFileApplyDiff) {
+			allowedTools.push(apply_diff_multi_file)
+		} else {
+			allowedTools.push(apply_diff_single_file)
 		}
-	})
+	}
 
 	return allowedTools
 }

src/core/task/Task.ts

@@ -2941,7 +2941,8 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 					mode,
 					undefined, // codeIndexManager is private, not accessible here
 					providerState,
-					this.api?.getModel()?.info?.supportsImages,
+					this.diffEnabled,
+					this.api?.getModel()?.info?.supportsImages ?? false,
 				)
 
 				metadata.allowedTools = allowedTools