feat: enhanced apply_diff

Author: NaccOll

src/core/assistant-message/AssistantMessageParser.ts

@@ -63,6 +63,9 @@ export class AssistantMessageParser {
 			const char = chunk[i]
 			this.accumulator += char
 			const currentPosition = accumulatorStartLength + i
+			if (this.currentToolUse && toolCallParam?.anthropicContent) {
+				this.currentToolUse.toolUseParam = toolCallParam.anthropicContent
+			}
 
 			// There should not be a param without a tool use.
 			if (this.currentToolUse && this.currentParamName) {
@@ -99,9 +102,6 @@ export class AssistantMessageParser {
 				const currentToolValue = this.accumulator.slice(this.currentToolUseStartIndex)
 				const toolUseClosingTag = `</${this.currentToolUse.name}>`
 				if (currentToolValue.endsWith(toolUseClosingTag)) {
-					if (toolCallParam?.anthropicContent && this.currentToolUse) {
-						this.currentToolUse.toolUseParam = toolCallParam.anthropicContent
-					}
 					// End of a tool use.
 					this.currentToolUse.partial = false
 

src/core/assistant-message/parseAssistantMessage.ts

@@ -21,7 +21,10 @@ export function parseAssistantMessage(
 	for (let i = 0; i < assistantMessage.length; i++) {
 		const char = assistantMessage[i]
 		accumulator += char
-
+		// During streaming, opportunistically attach temporary toolUseParam (if available)
+		if (currentToolUse && toolCallParam?.anthropicContent) {
+			currentToolUse.toolUseParam = toolCallParam.anthropicContent
+		}
 		// There should not be a param without a tool use.
 		if (currentToolUse && currentParamName) {
 			const currentParamValue = accumulator.slice(currentParamValueStartIndex)
@@ -48,9 +51,6 @@ export function parseAssistantMessage(
 			const currentToolValue = accumulator.slice(currentToolUseStartIndex)
 			const toolUseClosingTag = `</${currentToolUse.name}>`
 			if (currentToolValue.endsWith(toolUseClosingTag)) {
-				if (toolCallParam?.anthropicContent && currentToolUse) {
-					currentToolUse.toolUseParam = toolCallParam.anthropicContent
-				}
 				// End of a tool use.
 				currentToolUse.partial = false
 				contentBlocks.push(currentToolUse)

src/core/assistant-message/presentAssistantMessage.ts

@@ -492,7 +492,7 @@ export async function presentAssistantMessage(cline: Task) {
 						)
 					}
 
-					if (isMultiFileApplyDiffEnabled || toolCallEnabled) {
+					if (isMultiFileApplyDiffEnabled) {
 						await checkpointSaveAndMark(cline)
 						await applyDiffTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag)
 					} else {

src/core/prompts/tools/schemas/apply-diff-schema.ts

@@ -2,6 +2,151 @@ import { ToolArgs } from "../types"
 import { BaseToolSchema } from "./base-tool-schema"
 
 export function generateApplyDiffSchema(args: ToolArgs): BaseToolSchema {
+	if (args?.diffStrategy?.getName() === "MultiFileSearchReplace") {
+		return generateMultipleApplyDiffSchema(args)
+	}
+	const schema: BaseToolSchema = {
+		name: "apply_diff",
+		description:
+			"apply PRECISE, TARGETED modifications to an existing file by searching for specific sections of content and replacing them. This tool is for SURGICAL EDITS ONLY - specific changes to existing code",
+		parameters: [
+			{
+				name: "path",
+				type: "string",
+				description: "The relative path to the file that needs to be modified.",
+				required: true,
+			},
+			{
+				name: "diff",
+				type: "array",
+				description:
+					"An array of diff operations to be applied to the file. CRITICAL: For efficiency, include a large surrounding context (3-5 lines above and below) to combine multiple nearby changes into one operation instead of creating separate diffs for each line.",
+				required: true,
+				items: {
+					name: "diffItem",
+					type: "object",
+					description: "A single search-and-replace operation.",
+					required: true,
+					properties: {
+						d1: {
+							name: "start_line",
+							type: "number",
+							description:
+								"The starting line number of the 'search' block. Required when applying multiple diffs to a single file.",
+							required: true,
+						},
+						d2: {
+							name: "search",
+							type: "string",
+							description:
+								"The exact multi-line block of content to search for in the file. MUST match the original file content exactly (including all whitespace, indentation, tabs, and line breaks). Copy the exact text from the original file with perfect whitespace preservation.",
+							required: true,
+						},
+						d3: {
+							name: "replace",
+							type: "string",
+							description:
+								"The new multi-line content that will replace the search block. Preserve the original indentation structure and include all the context lines from the search block, making only the necessary changes to the target lines. This should be the complete replacement for the entire search block.",
+							required: true,
+						},
+					},
+				},
+			},
+		],
+
+		systemPrompt: `## apply_diff
+Description: Request to apply PRECISE, TARGETED modifications to an existing file by searching for specific sections of content and replacing them. This tool is for SURGICAL EDITS ONLY - specific changes to existing code.
+You can perform multiple distinct search and replace operations within a single \`apply_diff\` call by providing multiple SEARCH/REPLACE blocks in the \`diff\` parameter. This is the preferred way to make several targeted changes efficiently.
+The SEARCH section must exactly match existing content including whitespace and indentation.
+If you're not confident in the exact content to search for, use the read_file tool first to get the exact content.
+When applying the diffs, be extra careful to remember to change any closing brackets or other syntax that may be affected by the diff farther down in the file.
+ALWAYS make as many changes in a single 'apply_diff' request as possible using multiple SEARCH/REPLACE blocks
+
+Parameters:
+- path: (required) The path of the file to modify (relative to the current workspace directory ${args.cwd})
+- diff: (required) The search/replace block defining the changes.
+
+Diff format:
+\`\`\`
+<<<<<<< SEARCH
+:start_line: (required) The line number of original content where the search block starts.
+-------
+[exact content to find including whitespace]
+=======
+[new content to replace with]
+>>>>>>> REPLACE
+
+\`\`\`
+
+
+Example:
+
+Original file:
+\`\`\`
+1 | def calculate_total(items):
+2 |     total = 0
+3 |     for item in items:
+4 |         total += item
+5 |     return total
+\`\`\`
+
+Search/Replace content:
+\`\`\`
+<<<<<<< SEARCH
+:start_line:1
+-------
+def calculate_total(items):
+    total = 0
+    for item in items:
+        total += item
+    return total
+=======
+def calculate_total(items):
+    """Calculate total with 10% markup"""
+    return sum(item * 1.1 for item in items)
+>>>>>>> REPLACE
+
+\`\`\`
+
+Search/Replace content with multiple edits:
+\`\`\`
+<<<<<<< SEARCH
+:start_line:1
+-------
+def calculate_total(items):
+    sum = 0
+=======
+def calculate_sum(items):
+    sum = 0
+>>>>>>> REPLACE
+
+<<<<<<< SEARCH
+:start_line:4
+-------
+        total += item
+    return total
+=======
+        sum += item
+    return sum 
+>>>>>>> REPLACE
+\`\`\`
+
+
+Usage:
+<apply_diff>
+<path>File path here</path>
+<diff>
+Your search/replace content here
+You can use multi search/replace block in one diff block, but make sure to include the line numbers for each block.
+Only use a single line of '=======' between search and replacement content, because multiple '=======' will corrupt the file.
+</diff>
+</apply_diff>`,
+	}
+
+	return schema
+}
+
+function generateMultipleApplyDiffSchema(args: ToolArgs): BaseToolSchema {
 	const schema: BaseToolSchema = {
 		name: "apply_diff",
 		description:

src/core/task/Task.ts

@@ -378,7 +378,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 					EXPERIMENT_IDS.MULTI_FILE_APPLY_DIFF,
 				)
 
-				if (isMultiFileApplyDiffEnabled || this.apiConfiguration?.toolCallEnabled === true) {
+				if (isMultiFileApplyDiffEnabled) {
 					this.diffStrategy = new MultiFileSearchReplaceDiffStrategy(this.fuzzyMatchThreshold)
 				}
 			})
@@ -2083,7 +2083,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 						const streamingFailedMessage = this.abort
 							? undefined
 							: (error.message ?? JSON.stringify(serializeError(error), null, 2))
-
+						console.log(error)
 						// Now call abortTask after determining the cancel reason.
 						await this.abortTask()
 						await abortStream(cancelReason, streamingFailedMessage)

src/core/task/tool-call-helper.ts

@@ -161,6 +161,20 @@ export class StreamingToolCallProcessor {
 			originContent: this.accumulatedToolCalls,
 		}
 
+		// Provide a temporary anthropicContent (input) during streaming before final closure
+		const currentState = this.processingStates.get(index)
+		if (currentState && !currentState.functionClosed) {
+			const tmpInput = this.tryBuildTemporaryJson(currentState, toolCall.function.arguments)
+			if (tmpInput != null) {
+				result.anthropicContent = {
+					id: result.toolUserId,
+					name: result.toolName,
+					input: tmpInput,
+					type: "tool_use",
+				}
+			}
+		}
+
 		if (this.processingStates.get(index)?.functionClosed) {
 			let input
 			try {
@@ -484,6 +498,120 @@ export class StreamingToolCallProcessor {
 		return str.substring(pos, pos + 2)
 	}
 
+	/**
+	 * Try to synthesize a temporarily-closable JSON string from the current streaming buffer,
+	 * so we can JSON.parse it and expose a usable anthropicContent.input during partial streaming.
+	 * This function does NOT mutate parser state; it only operates on copies.
+	 *
+	 * Strategy:
+	 * - Safely close an in-flight string (key or value), trimming incomplete escapes/unicode.
+	 * - If a key has just finished (no value yet), inject ": null" to complete the pair.
+	 * - If a primitive token is mid-flight (true/false/null/number), attempt to complete it minimally.
+	 * - Trim trailing commas before closing braces.
+	 * - Close any open objects/arrays according to bracketStack, injecting "null" where a value is expected.
+	 * - Try parse; on failure, add one more "null" if needed and retry.
+	 */
+	private tryBuildTemporaryJson(state: ToolCallProcessingState, rawArgs: string): any | null {
+		let s = rawArgs
+
+		if (!s || s.trim().length === 0) {
+			return null
+		}
+
+		const trimTrailingComma = (str: string): string => str.replace(/,(\s*)$/, "$1")
+
+		const completePrimitiveSuffix = (pb: string): string => {
+			// Complete booleans/null prefixes
+			if (/^(t|tr|tru)$/.test(pb)) return "e" // true
+			if (/^(f|fa|fal|fals)$/.test(pb)) return "e" // false
+			if (/^(n|nu|nul)$/.test(pb)) return "l" // null
+			// Complete numeric partials like "-" or "12."
+			if (/^-?$/.test(pb)) return "0"
+			if (/^-?\d+\.$/.test(pb)) return "0"
+			return ""
+		}
+
+		const stripIncompleteUnicodeAtEnd = (input: string): string => {
+			const uniIndex = input.lastIndexOf("\\u")
+			if (uniIndex !== -1) {
+				const tail = input.slice(uniIndex + 2)
+				if (!/^[0-9a-fA-F]{4}$/.test(tail)) {
+					return input.slice(uniIndex) ? input.slice(0, uniIndex) : input
+				}
+			}
+			return input
+		}
+
+		// 1) Handle in-flight strings
+		if (state.inString) {
+			// Drop dangling backslash to avoid invalid escape at buffer end
+			if (s.endsWith("\\")) s = s.slice(0, -1)
+			// Trim incomplete unicode escape (e.g. \u12)
+			s = stripIncompleteUnicodeAtEnd(s)
+			// Close the string
+			s += `"`
+		} else {
+			// 2) Not inside a string; if a primitive token is partially accumulated, try to complete it minimally
+			if (state.primitiveBuffer && state.primitiveBuffer.length > 0) {
+				const suffix = completePrimitiveSuffix(state.primitiveBuffer)
+				if (suffix) s += suffix
+			}
+		}
+
+		// 4) Before closing brackets, remove trailing commas to avoid JSON syntax errors
+		s = trimTrailingComma(s)
+
+		// 5) Close any open objects/arrays per the current stack
+		if (state.bracketStack.length > 0) {
+			for (let i = state.bracketStack.length - 1; i >= 0; i--) {
+				// Always ensure no trailing comma before we append a closer
+				s = trimTrailingComma(s)
+
+				const b = state.bracketStack[i]
+
+				s += b === "{" ? "}" : "]"
+			}
+		}
+
+		// 6) First parse attempt
+		try {
+			return JSON.parse(s)
+		} catch {
+			// 7) Second attempt: add one more null if still dangling and retry
+			try {
+				let s2 = s
+				const lastNonWs = this.findLastNonWhitespaceChar(s2)
+				if (lastNonWs === ":" || state.parserState === ParserState.EXPECT_VALUE) {
+					s2 += "null"
+				}
+				s2 = trimTrailingComma(s2)
+				return JSON.parse(s2)
+			} catch {
+				// 8) Final fallback: repeatedly trim trailing commas and retry
+				let s3 = s
+				for (let k = 0; k < 3; k++) {
+					const trimmed = trimTrailingComma(s3)
+					if (trimmed === s3) break
+					s3 = trimmed
+					try {
+						return JSON.parse(s3)
+					} catch {
+						// continue
+					}
+				}
+				return null
+			}
+		}
+	}
+
+	private findLastNonWhitespaceChar(str: string): string {
+		for (let i = str.length - 1; i >= 0; i--) {
+			const ch = str[i]
+			if (!/\s/.test(ch)) return ch
+		}
+		return ""
+	}
+
 	private onOpenTag(tag: string, toolName: string): string {
 		return `<${tag}>`
 	}