Add taskId and checkpointNumber parameters to completePrompt and rela…

src/api/index.ts

@@ -23,11 +23,16 @@ import { HumanRelayHandler } from "./providers/human-relay"
 import { KiloCodeHandler } from "./providers/kilocode"
 
 export interface SingleCompletionHandler {
-	completePrompt(prompt: string): Promise<string>
+	completePrompt(prompt: string, taskId?: string, checkpointNumber?: number): Promise<string>
 }
 
 export interface ApiHandler {
-	createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream
+	createMessage(
+		systemPrompt: string,
+		messages: Anthropic.Messages.MessageParam[],
+		taskId?: string,
+		checkpointNumber?: number,
+	): ApiStream
 	getModel(): { id: string; info: ModelInfo }
 
 	/**

src/api/providers/base-provider.ts

@@ -14,7 +14,12 @@ const TOKEN_FUDGE_FACTOR = 1.5
 export abstract class BaseProvider implements ApiHandler {
 	// Cache the Tiktoken encoder instance since it's stateless
 	private encoder: Tiktoken | null = null
-	abstract createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream
+	abstract createMessage(
+		systemPrompt: string,
+		messages: Anthropic.Messages.MessageParam[],
+		taskId?: string,
+		checkpointNumber?: number,
+	): ApiStream
 	abstract getModel(): { id: string; info: ModelInfo }
 
 	/**

src/api/providers/kilocode.ts

@@ -20,19 +20,62 @@ export class KiloCodeHandler extends BaseProvider implements SingleCompletionHan
 		})
 	}
 
-	async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
+	private getIdempotencyKey(taskId: string, checkpointNumber: number): string {
+		// Create a deterministic idempotency key based on task_id and checkpoint number
+		return `${taskId}-${checkpointNumber}`
+	}
+
+	async *createMessage(
+		systemPrompt: string,
+		messages: Anthropic.Messages.MessageParam[],
+		taskId?: string,
+		checkpointNumber?: number,
+	): ApiStream {
 		let stream: AnthropicStream<Anthropic.Messages.RawMessageStreamEvent>
 		const cacheControl: CacheControlEphemeral = { type: "ephemeral" }
-		let { id: modelId, maxTokens, thinking, temperature, virtualId } = this.getModel()
+		const { id: modelId, maxTokens, thinking, temperature, virtualId } = this.getModel()
 
-		const userMsgIndices = messages.reduce(
-			(acc, msg, index) => (msg.role === "user" ? [...acc, index] : acc),
-			[] as number[],
-		)
+		// Use a for loop instead of reduce with spread to avoid linting error
+		const userMsgIndices: number[] = []
+		for (let i = 0; i < messages.length; i++) {
+			if (messages[i].role === "user") {
+				userMsgIndices.push(i)
+			}
+		}
 
 		const lastUserMsgIndex = userMsgIndices[userMsgIndices.length - 1] ?? -1
 		const secondLastMsgUserIndex = userMsgIndices[userMsgIndices.length - 2] ?? -1
 
+		// Prepare request options with headers
+		const requestOptions: { headers: Record<string, string> } = (() => {
+			const betas: string[] = []
+
+			// Check for models that support prompt caching
+			switch (modelId) {
+				case "claude-3-7-sonnet-20250219":
+				case "claude-3-5-sonnet-20241022":
+				case "claude-3-5-haiku-20241022":
+				case "claude-3-opus-20240229":
+				case "claude-3-haiku-20240307":
+					betas.push("prompt-caching-2024-07-31")
+					break
+			}
+
+			const headers: Record<string, string> = {}
+
+			// Add beta features if any
+			if (betas.length > 0) {
+				headers["anthropic-beta"] = betas.join(",")
+			}
+
+			// Add idempotency key if task_id and checkpoint number are provided
+			if (taskId && checkpointNumber !== undefined) {
+				headers["idempotency-key"] = this.getIdempotencyKey(taskId, checkpointNumber)
+			}
+
+			return { headers }
+		})()
+
 		stream = await this.client.messages.create(
 			{
 				model: modelId,
@@ -62,38 +105,12 @@ export class KiloCodeHandler extends BaseProvider implements SingleCompletionHan
 				// tools: tools,
 				stream: true,
 			},
-			(() => {
-				// prompt caching: https://x.com/alexalbert__/status/1823751995901272068
-				// https://github.com/anthropics/anthropic-sdk-typescript?tab=readme-ov-file#default-headers
-				// https://github.com/anthropics/anthropic-sdk-typescript/commit/c920b77fc67bd839bfeb6716ceab9d7c9bbe7393
-
-				const betas = []
-
-				// // Check for the thinking-128k variant first
-				// if (virtualId === "claude-3-7-sonnet-20250219:thinking") {
-				// 	betas.push("output-128k-2025-02-19")
-				// }
-
-				// Then check for models that support prompt caching
-				switch (modelId) {
-					case "claude-3-7-sonnet-20250219":
-					case "claude-3-5-sonnet-20241022":
-					case "claude-3-5-haiku-20241022":
-					case "claude-3-opus-20240229":
-					case "claude-3-haiku-20240307":
-						betas.push("prompt-caching-2024-07-31")
-						return {
-							headers: { "anthropic-beta": betas.join(",") },
-						}
-					default:
-						return undefined
-				}
-			})(),
+			requestOptions,
 		)
 
 		for await (const chunk of stream) {
 			switch (chunk.type) {
-				case "message_start":
+				case "message_start": {
 					// Tells us cache reads/writes/input/output.
 					const usage = chunk.message.usage
 
@@ -106,6 +123,7 @@ export class KiloCodeHandler extends BaseProvider implements SingleCompletionHan
 					}
 
 					break
+				}
 				case "message_delta":
 					// Tells us stop_reason, stop_sequence, and output tokens
 					// along the way and at the end of the message.
@@ -174,17 +192,30 @@ export class KiloCodeHandler extends BaseProvider implements SingleCompletionHan
 		}
 	}
 
-	async completePrompt(prompt: string) {
-		let { id: modelId, temperature } = this.getModel()
+	async completePrompt(prompt: string, taskId?: string, checkpointNumber?: number) {
+		const { id: modelId, temperature } = this.getModel()
 
-		const message = await this.client.messages.create({
-			model: modelId,
-			max_tokens: ANTHROPIC_DEFAULT_MAX_TOKENS,
-			thinking: undefined,
-			temperature,
-			messages: [{ role: "user", content: prompt }],
-			stream: false,
-		})
+		// Prepare request options with headers
+		const requestOptions: { headers: Record<string, string> } = {
+			headers: {},
+		}
+
+		// Add idempotency key if task_id and checkpoint number are provided
+		if (taskId && checkpointNumber !== undefined) {
+			requestOptions.headers["idempotency-key"] = this.getIdempotencyKey(taskId, checkpointNumber)
+		}
+
+		const message = await this.client.messages.create(
+			{
+				model: modelId,
+				max_tokens: ANTHROPIC_DEFAULT_MAX_TOKENS,
+				thinking: undefined,
+				temperature,
+				messages: [{ role: "user", content: prompt }],
+				stream: false,
+			},
+			requestOptions,
+		)
 
 		const content = message.content.find(({ type }) => type === "text")
 		return content?.type === "text" ? content.text : ""

src/core/Cline.ts

@@ -1162,7 +1162,11 @@ export class Cline {
 			}
 			return { role, content }
 		})
-		const stream = this.api.createMessage(systemPrompt, cleanConversationHistory)
+		// Get the current checkpoint number for idempotency key generation
+		const checkpointNumber = this.clineMessages.filter(({ say }) => say === "checkpoint_saved").length
+
+		// Pass task_id and checkpoint number to the API for idempotency key generation
+		const stream = this.api.createMessage(systemPrompt, cleanConversationHistory, this.taskId, checkpointNumber)
 		const iterator = stream[Symbol.asyncIterator]()
 
 		try {

src/core/webview/ClineProvider.ts

@@ -1507,6 +1507,10 @@ export class ClineProvider implements vscode.WebviewViewProvider {
 										},
 										customSupportPrompts,
 									),
+									this.getCurrentCline()?.taskId,
+									this.getCurrentCline()?.clineMessages.filter(
+										({ say }) => say === "checkpoint_saved",
+									).length,
 								)
 
 								await this.postMessageToWebview({

src/utils/__tests__/enhance-prompt.test.ts

@@ -20,7 +20,7 @@ describe("enhancePrompt", () => {
 
 		// Mock the API handler with a completePrompt method
 		;(buildApiHandler as jest.Mock).mockReturnValue({
-			completePrompt: jest.fn().mockResolvedValue("Enhanced prompt"),
+			completePrompt: jest.fn((prompt, taskId, checkpointNumber) => Promise.resolve("Enhanced prompt")),
 			createMessage: jest.fn(),
 			getModel: jest.fn().mockReturnValue({
 				id: "test-model",
@@ -38,7 +38,7 @@ describe("enhancePrompt", () => {
 
 		expect(result).toBe("Enhanced prompt")
 		const handler = buildApiHandler(mockApiConfig)
-		expect((handler as any).completePrompt).toHaveBeenCalledWith(`Test prompt`)
+		expect((handler as any).completePrompt).toHaveBeenCalledWith(`Test prompt`, undefined, undefined)
 	})
 
 	it("enhances prompt using custom enhancement prompt when provided", async () => {
@@ -60,7 +60,11 @@ describe("enhancePrompt", () => {
 
 		expect(result).toBe("Enhanced prompt")
 		const handler = buildApiHandler(mockApiConfig)
-		expect((handler as any).completePrompt).toHaveBeenCalledWith(`${customEnhancePrompt}\n\nTest prompt`)
+		expect((handler as any).completePrompt).toHaveBeenCalledWith(
+			`${customEnhancePrompt}\n\nTest prompt`,
+			undefined,
+			undefined,
+		)
 	})
 
 	it("throws error for empty prompt input", async () => {
@@ -101,7 +105,7 @@ describe("enhancePrompt", () => {
 
 		// Mock successful enhancement
 		;(buildApiHandler as jest.Mock).mockReturnValue({
-			completePrompt: jest.fn().mockResolvedValue("Enhanced prompt"),
+			completePrompt: jest.fn((prompt, taskId, checkpointNumber) => Promise.resolve("Enhanced prompt")),
 			createMessage: jest.fn(),
 			getModel: jest.fn().mockReturnValue({
 				id: "test-model",
@@ -121,7 +125,7 @@ describe("enhancePrompt", () => {
 
 	it("propagates API errors", async () => {
 		;(buildApiHandler as jest.Mock).mockReturnValue({
-			completePrompt: jest.fn().mockRejectedValue(new Error("API Error")),
+			completePrompt: jest.fn((prompt, taskId, checkpointNumber) => Promise.reject(new Error("API Error"))),
 			createMessage: jest.fn(),
 			getModel: jest.fn().mockReturnValue({
 				id: "test-model",

src/utils/single-completion-handler.ts

@@ -5,7 +5,12 @@ import { buildApiHandler, SingleCompletionHandler } from "../api"
  * Enhances a prompt using the configured API without creating a full Cline instance or task history.
  * This is a lightweight alternative that only uses the API's completion functionality.
  */
-export async function singleCompletionHandler(apiConfiguration: ApiConfiguration, promptText: string): Promise<string> {
+export async function singleCompletionHandler(
+	apiConfiguration: ApiConfiguration,
+	promptText: string,
+	taskId?: string,
+	checkpointNumber?: number,
+): Promise<string> {
 	if (!promptText) {
 		throw new Error("No prompt text provided")
 	}
@@ -20,5 +25,5 @@ export async function singleCompletionHandler(apiConfiguration: ApiConfiguration
 		throw new Error("The selected API provider does not support prompt enhancement")
 	}
 
-	return (handler as SingleCompletionHandler).completePrompt(promptText)
+	return (handler as SingleCompletionHandler).completePrompt(promptText, taskId, checkpointNumber)
 }