feat: implement apply_code tool for two-stage code generation (#6159)

packages/types/src/provider-settings.ts

@@ -63,6 +63,7 @@ export const DEFAULT_CONSECUTIVE_MISTAKE_LIMIT = 3
 const baseProviderSettingsSchema = z.object({
 	includeMaxTokens: z.boolean().optional(),
 	diffEnabled: z.boolean().optional(),
+	applyEnabled: z.boolean().optional(),
 	todoListEnabled: z.boolean().optional(),
 	fuzzyMatchThreshold: z.number().optional(),
 	modelTemperature: z.number().nullish(),

packages/types/src/tool.ts

@@ -19,6 +19,7 @@ export const toolNames = [
 	"read_file",
 	"write_to_file",
 	"apply_diff",
+	"apply_code",
 	"insert_content",
 	"search_and_replace",
 	"search_files",

src/core/assistant-message/presentAssistantMessage.ts

@@ -34,6 +34,7 @@ import { Task } from "../task/Task"
 import { codebaseSearchTool } from "../tools/codebaseSearchTool"
 import { experiments, EXPERIMENT_IDS } from "../../shared/experiments"
 import { applyDiffToolLegacy } from "../tools/applyDiffTool"
+import { applyCodeTool } from "../tools/applyCodeTool"
 
 /**
  * Processes and presents assistant message content to the user interface.
@@ -214,6 +215,8 @@ export async function presentAssistantMessage(cline: Task) {
 						const modeName = getModeBySlug(mode, customModes)?.name ?? mode
 						return `[${block.name} in ${modeName} mode: '${message}']`
 					}
+					case "apply_code":
+						return `[${block.name} for '${block.params.path}' with instruction: '${block.params.instruction}']`
 				}
 			}
 
@@ -522,6 +525,9 @@ export async function presentAssistantMessage(cline: Task) {
 						askFinishSubTaskApproval,
 					)
 					break
+				case "apply_code":
+					await applyCodeTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag)
+					break
 			}
 
 			break

src/core/prompts/tools/apply-code.ts

@@ -0,0 +1,34 @@
+import { ToolArgs } from "./types"
+
+export function getApplyCodeDescription(args: ToolArgs): string {
+	return `## apply_code
+Description: Request to apply code changes using a two-stage approach for improved reliability. This tool first generates code based on your instruction, then creates an accurate diff to integrate it into the existing file. This approach separates creative code generation from technical diff creation, resulting in more reliable code modifications.
+
+Parameters:
+- path: (required) The path of the file to modify (relative to the current workspace directory ${args.cwd})
+- instruction: (required) Clear instruction describing what code changes to make
+
+Usage:
+<apply_code>
+<path>File path here</path>
+<instruction>Your instruction for code changes</instruction>
+</apply_code>
+
+Example: Adding a new function to an existing file
+<apply_code>
+<path>src/utils.ts</path>
+<instruction>Add a function called calculateAverage that takes an array of numbers and returns their average</instruction>
+</apply_code>
+
+Example: Modifying existing code
+<apply_code>
+<path>src/api/handler.ts</path>
+<instruction>Update the error handling in the fetchData function to include retry logic with exponential backoff</instruction>
+</apply_code>
+
+Benefits over apply_diff:
+- More reliable: Separates code generation from diff creation
+- Cleaner context: Each stage has focused, minimal context
+- Better success rate: Reduces failures due to inaccurate diffs
+- Natural instructions: Use plain language instead of crafting diffs`
+}

src/core/prompts/tools/index.ts

@@ -23,6 +23,7 @@ import { getSwitchModeDescription } from "./switch-mode"
 import { getNewTaskDescription } from "./new-task"
 import { getCodebaseSearchDescription } from "./codebase-search"
 import { getUpdateTodoListDescription } from "./update-todo-list"
+import { getApplyCodeDescription } from "./apply-code"
 import { CodeIndexManager } from "../../../services/code-index/manager"
 
 // Map of tool names to their description functions
@@ -46,6 +47,7 @@ const toolDescriptionMap: Record<string, (args: ToolArgs) => string | undefined>
 	search_and_replace: (args) => getSearchAndReplaceDescription(args),
 	apply_diff: (args) =>
 		args.diffStrategy ? args.diffStrategy.getToolDescription({ cwd: args.cwd, toolOptions: args.toolOptions }) : "",
+	apply_code: (args) => getApplyCodeDescription(args),
 	update_todo_list: (args) => getUpdateTodoListDescription(args),
 }
 
@@ -114,6 +116,11 @@ export function getToolDescriptionsForMode(
 		tools.delete("update_todo_list")
 	}
 
+	// Conditionally exclude apply_code if disabled in settings
+	if (settings?.applyEnabled === false) {
+		tools.delete("apply_code")
+	}
+
 	// Map tool descriptions for allowed tools
 	const descriptions = Array.from(tools).map((toolName) => {
 		const descriptionFn = toolDescriptionMap[toolName]
@@ -148,4 +155,5 @@ export {
 	getInsertContentDescription,
 	getSearchAndReplaceDescription,
 	getCodebaseSearchDescription,
+	getApplyCodeDescription,
 }