feat: 添加远程 count_tokens API 支持并完善相关配置逻辑

src/anthropic/token.rs

@@ -7,7 +7,34 @@
 //! - 西文字符：每个计 1 个字符单位
 //! - 4 个字符单位 = 1 token（四舍五入）
 
-use crate::anthropic::types::{Message, SystemMessage, Tool};
+use std::sync::OnceLock;
+use crate::anthropic::types::{CountTokensRequest, CountTokensResponse, Message, SystemMessage, Tool};
+
+/// Count Tokens API 配置
+#[derive(Clone, Default)]
+pub struct CountTokensConfig {
+    /// 外部 count_tokens API 地址
+    pub api_url: Option<String>,
+    /// count_tokens API 密钥
+    pub api_key: Option<String>,
+    /// count_tokens API 认证类型（"x-api-key" 或 "bearer"）
+    pub auth_type: String,
+}
+
+/// 全局配置存储
+static COUNT_TOKENS_CONFIG: OnceLock<CountTokensConfig> = OnceLock::new();
+
+/// 初始化 count_tokens 配置
+///
+/// 应在应用启动时调用一次
+pub fn init_config(config: CountTokensConfig) {
+    let _ = COUNT_TOKENS_CONFIG.set(config);
+}
+
+/// 获取配置
+fn get_config() -> Option<&'static CountTokensConfig> {
+    COUNT_TOKENS_CONFIG.get()
+}
 
 
 /// 判断字符是否为非西文字符
@@ -71,7 +98,83 @@ pub fn count_tokens(text: &str) -> u64 {
 
 
 /// 估算请求的输入 tokens
+///
+/// 优先调用远程 API，失败时回退到本地计算
 pub(crate) fn count_all_tokens(system: Option<Vec<SystemMessage>>, messages: Vec<Message>, tools: Option<Vec<Tool>>) -> u64 {
+    // 检查是否配置了远程 API
+    if let Some(config) = get_config() {
+        if let Some(api_url) = &config.api_url {
+            // 尝试调用远程 API
+            let result = tokio::task::block_in_place(|| {
+                tokio::runtime::Handle::current().block_on(
+                    call_remote_count_tokens(api_url, config, &system, &messages, &tools)
+                )
+            });
+
+            match result {
+                Ok(tokens) => {
+                    tracing::debug!("远程 count_tokens API 返回: {}", tokens);
+                    return tokens;
+                }
+                Err(e) => {
+                    tracing::warn!("远程 count_tokens API 调用失败，回退到本地计算: {}", e);
+                }
+            }
+        }
+    }
+
+    // 本地计算
+    count_all_tokens_local(system, messages, tools)
+}
+
+/// 调用远程 count_tokens API
+async fn call_remote_count_tokens(
+    api_url: &str,
+    config: &CountTokensConfig,
+    system: &Option<Vec<SystemMessage>>,
+    messages: &Vec<Message>,
+    tools: &Option<Vec<Tool>>,
+) -> Result<u64, Box<dyn std::error::Error + Send + Sync>> {
+    let client = reqwest::Client::new();
+
+    // 构建请求体
+    let request = CountTokensRequest {
+        model: "claude-sonnet-4-5-20250929".to_string(), // 模型名称用于 token 计算
+        messages: messages.clone(),
+        system: system.clone(),
+        tools: tools.clone(),
+    };
+
+    // 构建请求
+    let mut req_builder = client.post(api_url);
+
+    // 设置认证头
+    if let Some(api_key) = &config.api_key {
+        if config.auth_type == "bearer" {
+            req_builder = req_builder.header("Authorization", format!("Bearer {}", api_key));
+        } else {
+            req_builder = req_builder.header("x-api-key", api_key);
+        }
+    }
+
+    // 发送请求
+    let response = req_builder
+        .header("Content-Type", "application/json")
+        .json(&request)
+        .timeout(std::time::Duration::from_secs(5))
+        .send()
+        .await?;
+
+    if !response.status().is_success() {
+        return Err(format!("API 返回错误状态: {}", response.status()).into());
+    }
+
+    let result: CountTokensResponse = response.json().await?;
+    Ok(result.input_tokens as u64)
+}
+
+/// 本地计算请求的输入 tokens
+fn count_all_tokens_local(system: Option<Vec<SystemMessage>>, messages: Vec<Message>, tools: Option<Vec<Tool>>) -> u64 {
     let mut total = 0;
 
     // 系统消息

src/anthropic/types.rs

@@ -98,7 +98,7 @@ pub struct MessagesRequest {
 }
 
 /// 消息
-#[derive(Debug, Deserialize)]
+#[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct Message {
     pub role: String,
     /// 可以是 string 或 ContentBlock 数组
@@ -106,13 +106,13 @@ pub struct Message {
 }
 
 /// 系统消息
-#[derive(Debug, Deserialize)]
+#[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct SystemMessage {
-    pub text: String
+    pub text: String,
 }
 
 /// 工具定义
-#[derive(Debug, Deserialize)]
+#[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct Tool {
     pub name: String,
     pub description: String,
@@ -154,16 +154,18 @@ pub struct ImageSource {
 // === Count Tokens 端点类型 ===
 
 /// Token 计数请求
-#[derive(Debug, Deserialize)]
+#[derive(Debug, Serialize, Deserialize)]
 pub struct CountTokensRequest {
     pub model: String,
     pub messages: Vec<Message>,
+    #[serde(skip_serializing_if = "Option::is_none")]
     pub system: Option<Vec<SystemMessage>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
     pub tools: Option<Vec<Tool>>,
 }
 
 /// Token 计数响应
-#[derive(Debug, Serialize)]
+#[derive(Debug, Serialize, Deserialize)]
 pub struct CountTokensResponse {
     pub input_tokens: i32,
 }

src/main.rs

@@ -48,6 +48,13 @@ async fn main() {
     let token_manager = TokenManager::new(config.clone(), credentials.clone());
     let kiro_provider = KiroProvider::new(token_manager);
 
+    // 初始化 count_tokens 配置
+    anthropic::token::init_config(anthropic::token::CountTokensConfig {
+        api_url: config.count_tokens_api_url.clone(),
+        api_key: config.count_tokens_api_key.clone(),
+        auth_type: config.count_tokens_auth_type.clone(),
+    });
+
     // 构建路由（从凭据获取 profile_arn）
     let app = anthropic::create_router_with_provider(&api_key, Some(kiro_provider), credentials.profile_arn.clone());