Revert "refactor: 由于模型事件中实现了精准上下文统计, 移除 count_tokens 强制 CC 使用会话测试"

This reverts commit aa2a16e9c76091a0a648d6fa738e2f667f28e003.

README.md

@@ -17,6 +17,7 @@
 |------|------|-------------|
 | `/v1/models` | GET | 获取可用模型列表    |
 | `/v1/messages` | POST | 创建消息（对话）    |
+| `/v1/messages/count_tokens` | POST | 估算 Token 数量 |
 
 ## 快速开始
 
@@ -39,7 +40,10 @@ cargo build --release
    "kiroVersion": "0.8.0",
    "machineId": "如果你需要自定义机器码请将64位机器码填到这里",  // 不是标准格式会自动忽略, 自动生成
    "systemVersion": "darwin#24.6.0",
-   "nodeVersion": "22.21.1"
+   "nodeVersion": "22.21.1",
+   "countTokensApiUrl": "https://api.example.com/v1/messages/count_tokens",  // 可选，外部 count_tokens API 地址
+   "countTokensApiKey": "sk-your-count-tokens-api-key",  // 可选，外部 API 密钥
+   "countTokensAuthType": "x-api-key"  // 可选，认证类型：x-api-key 或 bearer
 }
 ```
 
@@ -99,6 +103,9 @@ curl http://127.0.0.1:8990/v1/messages \
 | `machineId` | string | - | 自定义机器码（64位十六进制）不定义则自动生成 |
 | `systemVersion` | string | 随机 | 系统版本标识                  |
 | `nodeVersion` | string | `22.21.1` | Node.js 版本标识            |
+| `countTokensApiUrl` | string | - | 外部 count_tokens API 地址（可选） |
+| `countTokensApiKey` | string | - | 外部 count_tokens API 密钥（可选） |
+| `countTokensAuthType` | string | `x-api-key` | 外部 API 认证类型：`x-api-key` 或 `bearer` |
 
 ### credentials.json
 

config.example.json

@@ -6,5 +6,8 @@
   "kiroVersion": "0.8.0",
   "machineId": "如果你需要自定义机器码请将64位机器码填到这里",
   "systemVersion": "darwin#24.6.0",
-  "nodeVersion": "22.21.1"
+  "nodeVersion": "22.21.1",
+  "countTokensApiUrl": "https://api.example.com/v1/messages/count_tokens",
+  "countTokensApiKey": "sk-your-count-tokens-api-key",
+  "countTokensAuthType": "x-api-key"
 }

src/anthropic/handlers.rs

@@ -23,7 +23,9 @@ use crate::kiro::parser::decoder::EventStreamDecoder;
 use super::converter::{convert_request, ConversionError};
 use super::middleware::AppState;
 use super::stream::{SseEvent, StreamContext};
-use super::types::{ErrorResponse, MessagesRequest, Model, ModelsResponse};
+use super::types::{
+    CountTokensRequest, CountTokensResponse, ErrorResponse, MessagesRequest, Model, ModelsResponse,
+};
 
 /// GET /v1/models
 ///
@@ -468,3 +470,19 @@ async fn handle_non_stream_request(
 
 
 
+/// POST /v1/messages/count_tokens
+///
+/// 计算消息的 token 数量
+pub async fn count_tokens(JsonExtractor(payload): JsonExtractor<CountTokensRequest>) -> impl IntoResponse {
+    tracing::info!(
+        model = %payload.model,
+        message_count = %payload.messages.len(),
+        "Received POST /v1/messages/count_tokens request"
+    );
+
+    let total_tokens = token::count_all_tokens(payload.model, payload.system, payload.messages, payload.tools) as i32;
+
+    Json(CountTokensResponse {
+        input_tokens: total_tokens.max(1) as i32,
+    })
+}

src/anthropic/mod.rs

@@ -5,6 +5,7 @@
 //! # 支持的端点
 //! - `GET /v1/models` - 获取可用模型列表
 //! - `POST /v1/messages` - 创建消息（对话）
+//! - `POST /v1/messages/count_tokens` - 计算 token 数量
 //!
 //! # 使用示例
 //! ```rust,ignore

src/anthropic/router.rs

@@ -9,7 +9,7 @@ use axum::{
 use crate::kiro::provider::KiroProvider;
 
 use super::{
-    handlers::{get_models, post_messages},
+    handlers::{count_tokens, get_models, post_messages},
     middleware::{auth_middleware, cors_layer, AppState},
 };
 
@@ -18,6 +18,7 @@ use super::{
 /// # 端点
 /// - `GET /v1/models` - 获取可用模型列表
 /// - `POST /v1/messages` - 创建消息（对话）
+/// - `POST /v1/messages/count_tokens` - 计算 token 数量
 ///
 /// # 认证
 /// 所有 `/v1` 路径需要 API Key 认证，支持：
@@ -46,6 +47,7 @@ pub fn create_router_with_provider(
     let v1_routes = Router::new()
         .route("/models", get(get_models))
         .route("/messages", post(post_messages))
+        .route("/messages/count_tokens", post(count_tokens))
         .layer(middleware::from_fn_with_state(
             state.clone(),
             auth_middleware,

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,84 @@ pub fn count_tokens(text: &str) -> u64 {
 
 
 /// 估算请求的输入 tokens
-pub(crate) fn count_all_tokens(_model: String, system: Option<Vec<SystemMessage>>, messages: Vec<Message>, tools: Option<Vec<Tool>>) -> u64 {
+///
+/// 优先调用远程 API，失败时回退到本地计算
+pub(crate) fn count_all_tokens(model: String, 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, model, &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,
+    model: String,
+    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: model, // 模型名称用于 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

@@ -150,3 +150,22 @@ pub struct ImageSource {
     pub media_type: String,
     pub data: String,
 }
+
+// === Count Tokens 端点类型 ===
+
+/// Token 计数请求
+#[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, 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());
 
@@ -58,6 +65,7 @@ async fn main() {
     tracing::info!("可用 API:");
     tracing::info!("  GET  /v1/models");
     tracing::info!("  POST /v1/messages");
+    tracing::info!("  POST /v1/messages/count_tokens");
 
     let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
     axum::serve(listener, app).await.unwrap();

src/model/config.rs

@@ -29,6 +29,18 @@ pub struct Config {
 
     #[serde(default = "default_node_version")]
     pub node_version: String,
+
+    /// 外部 count_tokens API 地址（可选）
+    #[serde(default)]
+    pub count_tokens_api_url: Option<String>,
+
+    /// count_tokens API 密钥（可选）
+    #[serde(default)]
+    pub count_tokens_api_key: Option<String>,
+
+    /// count_tokens API 认证类型（可选，"x-api-key" 或 "bearer"，默认 "x-api-key"）
+    #[serde(default = "default_count_tokens_auth_type")]
+    pub count_tokens_auth_type: String,
 }
 
 fn default_host() -> String {
@@ -56,6 +68,10 @@ fn default_node_version() -> String {
     "22.21.1".to_string()
 }
 
+fn default_count_tokens_auth_type() -> String {
+    "x-api-key".to_string()
+}
+
 impl Default for Config {
     fn default() -> Self {
         Self {
@@ -67,6 +83,9 @@ impl Default for Config {
             api_key: None,
             system_version: default_system_version(),
             node_version: default_node_version(),
+            count_tokens_api_url: None,
+            count_tokens_api_key: None,
+            count_tokens_auth_type: default_count_tokens_auth_type(),
         }
     }
 }