src/api/providers/claude-code.ts - main - outbreak/kilocode

src/api/providers/claude-code.ts
196 строк6 KB
Kevin van Dijk
Merge branch 'main' into roo-v3.34.7
03 дек 2025, 21:59
03 дек 2025, 21:594acb0bc
import type { Anthropic } from "@anthropic-ai/sdk"
import {
	claudeCodeDefaultModelId,
	type ClaudeCodeModelId,
	claudeCodeModels,
	type ModelInfo,
	getClaudeCodeModelId,
} from "@roo-code/types"
import { type ApiHandler, ApiHandlerCreateMessageMetadata } from ".."
import { ApiStreamUsageChunk, type ApiStream } from "../transform/stream"
import { runClaudeCode } from "../../integrations/claude-code/run"
import { filterMessagesForClaudeCode } from "../../integrations/claude-code/message-filter"
import { t } from "../../i18n"
import { ApiHandlerOptions } from "../../shared/api"
import { countTokens } from "../../utils/countTokens"

export class ClaudeCodeHandler implements ApiHandler {
	private options: ApiHandlerOptions

	constructor(options: ApiHandlerOptions) {
		this.options = options
	}

	async *createMessage(
		systemPrompt: string,
		messages: Anthropic.Messages.MessageParam[],
		_metadata?: ApiHandlerCreateMessageMetadata,
	): ApiStream {
		// Filter out image blocks since Claude Code doesn't support them
		const filteredMessages = filterMessagesForClaudeCode(messages)

		const useVertex = process.env.CLAUDE_CODE_USE_VERTEX === "1"
		const model = this.getModel()

		// Validate that the model ID is a valid ClaudeCodeModelId
		const modelId = model.id in claudeCodeModels ? (model.id as ClaudeCodeModelId) : claudeCodeDefaultModelId

		const claudeProcess = runClaudeCode({
			systemPrompt,
			messages: filteredMessages,
			path: this.options.claudeCodePath,
			modelId: getClaudeCodeModelId(modelId, useVertex),
			maxOutputTokens: this.options.claudeCodeMaxOutputTokens,
		})

		// Usage is included with assistant messages,
		// but cost is included in the result chunk
		const usage: ApiStreamUsageChunk = {
			type: "usage",
			inputTokens: 0,
			outputTokens: 0,
			cacheReadTokens: 0,
			cacheWriteTokens: 0,
		}

		let isPaidUsage = true

		for await (const chunk of claudeProcess) {
			if (typeof chunk === "string") {
				yield {
					type: "text",
					text: chunk,
				}

				continue
			}

			if (chunk.type === "system" && chunk.subtype === "init") {
				// Based on my tests, subscription usage sets the `apiKeySource` to "none"
				isPaidUsage = chunk.apiKeySource !== "none"
				continue
			}

			if (chunk.type === "assistant" && "message" in chunk) {
				const message = chunk.message

				if (message.stop_reason !== null) {
					const content = "text" in message.content[0] ? message.content[0] : undefined

					const isError = content && content.text.startsWith(`API Error`)
					if (isError) {
						// Error messages are formatted as: `API Error: <<status code>> <<json>>`
						const errorMessageStart = content.text.indexOf("{")
						const errorMessage = content.text.slice(errorMessageStart)

						const error = this.attemptParse(errorMessage)
						if (!error) {
							throw new Error(content.text)
						}

						if (error.error.message.includes("Invalid model name")) {
							throw new Error(
								content.text + `\n\n${t("common:errors.claudeCode.apiKeyModelPlanMismatch")}`,
							)
						}

						throw new Error(errorMessage)
					}
				}

				let toolReminder = false // kilocode_change
				for (const content of message.content) {
					switch (content.type) {
						case "text":
							yield {
								type: "text",
								text: content.text,
							}
							break
						case "thinking":
							yield {
								type: "reasoning",
								text: content.thinking || "",
							}
							break
						case "redacted_thinking":
							yield {
								type: "reasoning",
								text: "[Redacted thinking block]",
							}
							break
						case "tool_use":
							console.error(`tool_use is not supported yet. Received: ${JSON.stringify(content)}`)
							// kilocode_change start
							if (!toolReminder) {
								yield {
									type: "text",
									text: "Preparing to use tools in XML format... ",
								}
								toolReminder = true
							}
							// kilocode_change end
							break
					}
				}

				// Accumulate usage across streaming chunks
				usage.inputTokens += message.usage.input_tokens
				usage.outputTokens += message.usage.output_tokens
				usage.cacheReadTokens = (usage.cacheReadTokens || 0) + (message.usage.cache_read_input_tokens || 0)
				usage.cacheWriteTokens =
					(usage.cacheWriteTokens || 0) + (message.usage.cache_creation_input_tokens || 0)

				continue
			}

			if (chunk.type === "result" && "result" in chunk) {
				usage.totalCost = isPaidUsage ? chunk.total_cost_usd : 0

				yield usage
			}
		}
	}

	getModel(): { id: string; info: ModelInfo } {
		const modelId = this.options.apiModelId
		if (modelId && modelId in claudeCodeModels) {
			const id = modelId as ClaudeCodeModelId
			const modelInfo: ModelInfo = { ...claudeCodeModels[id] }

			// Override maxTokens with the configured value if provided
			if (this.options.claudeCodeMaxOutputTokens !== undefined) {
				modelInfo.maxTokens = this.options.claudeCodeMaxOutputTokens
			}

			return { id, info: modelInfo }
		}

		const defaultModelInfo: ModelInfo = { ...claudeCodeModels[claudeCodeDefaultModelId] }

		// Override maxTokens with the configured value if provided
		if (this.options.claudeCodeMaxOutputTokens !== undefined) {
			defaultModelInfo.maxTokens = this.options.claudeCodeMaxOutputTokens
		}

		return {
			id: claudeCodeDefaultModelId,
			info: defaultModelInfo,
		}
	}

	async countTokens(content: Anthropic.Messages.ContentBlockParam[]): Promise<number> {
		if (content.length === 0) {
			return 0
		}
		return countTokens(content, { useWorker: true })
	}

	private attemptParse(str: string) {
		try {
			return JSON.parse(str)
		} catch (err) {
			return null
		}
	}
}