Roo-Code/src/core/task/validateToolResultIds.ts

import { Anthropic } from "@anthropic-ai/sdk"
import { TelemetryService } from "@roo-code/telemetry"
import { findLastIndex } from "../../shared/array"

/**
 * Custom error class for tool result ID mismatches.
 * Used for structured error tracking via PostHog.
 */
export class ToolResultIdMismatchError extends Error {
	constructor(
		message: string,
		public readonly toolResultIds: string[],
		public readonly toolUseIds: string[],
	) {
		super(message)
		this.name = "ToolResultIdMismatchError"
	}
}

/**
 * Custom error class for missing tool results.
 * Used for structured error tracking via PostHog when tool_use blocks
 * don't have corresponding tool_result blocks.
 */
export class MissingToolResultError extends Error {
	constructor(
		message: string,
		public readonly missingToolUseIds: string[],
		public readonly existingToolResultIds: string[],
	) {
		super(message)
		this.name = "MissingToolResultError"
	}
}

/**
 * Validates and fixes tool_result IDs in a user message against the previous assistant message.
 *
 * This is a centralized validation that catches all tool_use/tool_result issues
 * before messages are added to the API conversation history. It handles scenarios like:
 * - Race conditions during streaming
 * - Message editing scenarios
 * - Resume/delegation scenarios
 * - Missing tool_result blocks for tool_use calls
 *
 * @param userMessage - The user message being added to history
 * @param apiConversationHistory - The conversation history to find the previous assistant message from
 * @returns The validated user message with corrected tool_use_ids and any missing tool_results added
 */
export function validateAndFixToolResultIds(
	userMessage: Anthropic.MessageParam,
	apiConversationHistory: Anthropic.MessageParam[],
): Anthropic.MessageParam {
	// Only process user messages with array content
	if (userMessage.role !== "user" || !Array.isArray(userMessage.content)) {
		return userMessage
	}

	// Find the previous assistant message from conversation history
	const prevAssistantIdx = findLastIndex(apiConversationHistory, (msg) => msg.role === "assistant")
	if (prevAssistantIdx === -1) {
		return userMessage
	}

	const previousAssistantMessage = apiConversationHistory[prevAssistantIdx]

	// Get tool_use blocks from the assistant message
	const assistantContent = previousAssistantMessage.content
	if (!Array.isArray(assistantContent)) {
		return userMessage
	}

	const toolUseBlocks = assistantContent.filter((block): block is Anthropic.ToolUseBlock => block.type === "tool_use")

	// No tool_use blocks to match against - no validation needed
	if (toolUseBlocks.length === 0) {
		return userMessage
	}

	// Find tool_result blocks in the user message
	let toolResults = userMessage.content.filter(
		(block): block is Anthropic.ToolResultBlockParam => block.type === "tool_result",
	)

	// Deduplicate tool_result blocks to prevent API protocol violations (GitHub #10465)
	// This serves as a safety net for any potential race conditions that could generate
	// duplicate tool_results with the same tool_use_id. The root cause (approval feedback
	// creating duplicate results) has been fixed in presentAssistantMessage.ts, but this
	// deduplication remains as a defensive measure for unknown edge cases.
	const seenToolResultIds = new Set<string>()
	const deduplicatedContent = userMessage.content.filter((block) => {
		if (block.type !== "tool_result") {
			return true
		}
		if (seenToolResultIds.has(block.tool_use_id)) {
			return false // Duplicate - filter out
		}
		seenToolResultIds.add(block.tool_use_id)
		return true
	})

	userMessage = {
		...userMessage,
		content: deduplicatedContent,
	}

	toolResults = deduplicatedContent.filter(
		(block): block is Anthropic.ToolResultBlockParam => block.type === "tool_result",
	)

	// Build a set of valid tool_use IDs
	const validToolUseIds = new Set(toolUseBlocks.map((block) => block.id))

	// Build a set of existing tool_result IDs
	const existingToolResultIds = new Set(toolResults.map((r) => r.tool_use_id))

	// Check for missing tool_results (tool_use IDs that don't have corresponding tool_results)
	const missingToolUseIds = toolUseBlocks
		.filter((toolUse) => !existingToolResultIds.has(toolUse.id))
		.map((toolUse) => toolUse.id)

	// Check if any tool_result has an invalid ID
	const hasInvalidIds = toolResults.some((result) => !validToolUseIds.has(result.tool_use_id))

	// If no missing tool_results and no invalid IDs, check if reordering is needed
	if (missingToolUseIds.length === 0 && !hasInvalidIds) {
		// Reorder tool_result blocks to match tool_use order (required by Anthropic API)
		const reordered = reorderToolResults(
			userMessage.content as Anthropic.Messages.ContentBlockParam[],
			toolUseBlocks,
		)
		if (reordered) {
			return { ...userMessage, content: reordered }
		}
		return userMessage
	}

	// We have issues - need to fix them
	const toolResultIdList = toolResults.map((r) => r.tool_use_id)
	const toolUseIdList = toolUseBlocks.map((b) => b.id)

	// Report missing tool_results to PostHog error tracking
	if (missingToolUseIds.length > 0 && TelemetryService.hasInstance()) {
		TelemetryService.instance.captureException(
			new MissingToolResultError(
				`Detected missing tool_result blocks. Missing tool_use IDs: [${missingToolUseIds.join(", ")}], existing tool_result IDs: [${toolResultIdList.join(", ")}]`,
				missingToolUseIds,
				toolResultIdList,
			),
			{
				missingToolUseIds,
				existingToolResultIds: toolResultIdList,
				toolUseCount: toolUseBlocks.length,
				toolResultCount: toolResults.length,
			},
		)
	}

	// Report ID mismatches to PostHog error tracking
	if (hasInvalidIds && TelemetryService.hasInstance()) {
		TelemetryService.instance.captureException(
			new ToolResultIdMismatchError(
				`Detected tool_result ID mismatch. tool_result IDs: [${toolResultIdList.join(", ")}], tool_use IDs: [${toolUseIdList.join(", ")}]`,
				toolResultIdList,
				toolUseIdList,
			),
			{
				toolResultIds: toolResultIdList,
				toolUseIds: toolUseIdList,
				toolResultCount: toolResults.length,
				toolUseCount: toolUseBlocks.length,
			},
		)
	}

	// Match tool_results to tool_uses by position and fix incorrect IDs
	const usedToolUseIds = new Set<string>()
	const contentArray = userMessage.content as Anthropic.Messages.ContentBlockParam[]

	const correctedContent = contentArray
		.map((block: Anthropic.Messages.ContentBlockParam) => {
			if (block.type !== "tool_result") {
				return block
			}

			// If the ID is already valid and not yet used, keep it
			if (validToolUseIds.has(block.tool_use_id) && !usedToolUseIds.has(block.tool_use_id)) {
				usedToolUseIds.add(block.tool_use_id)
				return block
			}

			// Find which tool_result index this block is by comparing references.
			// This correctly handles duplicate tool_use_ids - we find the actual block's
			// position among all tool_results, not the first block with a matching ID.
			const toolResultIndex = toolResults.indexOf(block as Anthropic.ToolResultBlockParam)

			// Try to match by position - only fix if there's a corresponding tool_use
			if (toolResultIndex !== -1 && toolResultIndex < toolUseBlocks.length) {
				const correctId = toolUseBlocks[toolResultIndex].id
				// Only use this ID if it hasn't been used yet
				if (!usedToolUseIds.has(correctId)) {
					usedToolUseIds.add(correctId)
					return {
						...block,
						tool_use_id: correctId,
					}
				}
			}

			// No corresponding tool_use for this tool_result, or the ID is already used
			return null
		})
		.filter((block): block is NonNullable<typeof block> => block !== null)

	// Add missing tool_result blocks for any tool_use that doesn't have one
	const coveredToolUseIds = new Set(
		correctedContent
			.filter(
				(b: Anthropic.Messages.ContentBlockParam): b is Anthropic.ToolResultBlockParam =>
					b.type === "tool_result",
			)
			.map((r: Anthropic.ToolResultBlockParam) => r.tool_use_id),
	)

	const stillMissingToolUseIds = toolUseBlocks.filter((toolUse) => !coveredToolUseIds.has(toolUse.id))

	// Build final content: add missing tool_results at the beginning if any
	const missingToolResults: Anthropic.ToolResultBlockParam[] = stillMissingToolUseIds.map((toolUse) => ({
		type: "tool_result" as const,
		tool_use_id: toolUse.id,
		content: "Tool execution was interrupted before completion.",
	}))

	// Combine missing tool_results with corrected content
	const combinedContent =
		missingToolResults.length > 0 ? [...missingToolResults, ...correctedContent] : correctedContent

	// Reorder tool_result blocks to match the tool_use order (required by Anthropic API).
	// This handles the case where tool results were appended in completion order rather
	// than the original tool_use order.
	const finalContent = reorderToolResults(combinedContent, toolUseBlocks) ?? combinedContent

	return {
		...userMessage,
		content: finalContent,
	}
}

/**
 * Reorders tool_result blocks within a content array to match the order of
 * their corresponding tool_use blocks from the assistant message.
 *
 * Non-tool-result blocks (text, image, etc.) remain in their original
 * positions relative to the tool_result blocks -- only tool_results are
 * reordered among themselves.
 *
 * Returns `null` if the tool_results are already in the correct order
 * (no reordering needed).
 */
function reorderToolResults(
	content: Anthropic.Messages.ContentBlockParam[],
	toolUseBlocks: Anthropic.ToolUseBlock[],
): Anthropic.Messages.ContentBlockParam[] | null {
	if (toolUseBlocks.length === 0) {
		return null
	}

	// Build an order map: tool_use_id -> position index
	const orderMap = new Map<string, number>()
	toolUseBlocks.forEach((block, index) => {
		orderMap.set(block.id, index)
	})

	// Separate tool_result blocks from non-tool-result blocks, preserving indices
	const toolResultEntries: { index: number; block: Anthropic.ToolResultBlockParam }[] = []
	const nonToolResultEntries: { index: number; block: Anthropic.Messages.ContentBlockParam }[] = []

	content.forEach((block, index) => {
		if (block.type === "tool_result") {
			toolResultEntries.push({ index, block: block as Anthropic.ToolResultBlockParam })
		} else {
			nonToolResultEntries.push({ index, block })
		}
	})

	if (toolResultEntries.length <= 1) {
		return null // Nothing to reorder
	}

	// Sort tool_result blocks by their corresponding tool_use order
	const sortedToolResults = [...toolResultEntries].sort((a, b) => {
		const orderA = orderMap.get(a.block.tool_use_id) ?? Number.MAX_SAFE_INTEGER
		const orderB = orderMap.get(b.block.tool_use_id) ?? Number.MAX_SAFE_INTEGER
		return orderA - orderB
	})

	// Check if already in correct order
	const alreadyOrdered = sortedToolResults.every((entry, i) => entry === toolResultEntries[i])
	if (alreadyOrdered) {
		return null
	}

	// Reconstruct the array: place sorted tool_results into the original
	// tool_result positions, keeping non-tool-result blocks where they were.
	const result: Anthropic.Messages.ContentBlockParam[] = new Array(content.length)

	// First, place non-tool-result blocks back at their original indices
	for (const entry of nonToolResultEntries) {
		result[entry.index] = entry.block
	}

	// Then, place sorted tool_results into the slots that were originally
	// occupied by tool_result blocks (preserving relative position of non-tool blocks)
	const toolResultSlots = toolResultEntries.map((e) => e.index)
	sortedToolResults.forEach((entry, i) => {
		result[toolResultSlots[i]] = entry.block
	})

	return result
}