diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock index 86399fc0a..0f05f8845 100644 --- a/src-tauri/Cargo.lock +++ b/src-tauri/Cargo.lock @@ -6675,6 +6675,7 @@ dependencies = [ "proxycast-infra", "proxycast-processor", "proxycast-providers", + "proxycast-server-utils", "proxycast-services", "proxycast-terminal", "proxycast-websocket", @@ -6867,6 +6868,21 @@ dependencies = [ "uuid", ] +[[package]] +name = "proxycast-server-utils" +version = "0.60.0" +dependencies = [ + "axum 0.7.9", + "futures", + "proptest", + "proxycast-core", + "regex", + "serde_json", + "tokio", + "tracing", + "uuid", +] + [[package]] name = "proxycast-services" version = "0.60.0" diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index e7cb811f1..7ea4fa8ec 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -20,6 +20,7 @@ proxycast-terminal = { path = "crates/terminal" } proxycast-credential = { path = "crates/credential" } proxycast-websocket = { path = "crates/websocket" } proxycast-processor = { path = "crates/processor" } +proxycast-server-utils = { path = "crates/server-utils" } voice-core = { path = "crates/voice-core" } # 序列化 @@ -199,6 +200,7 @@ proxycast-terminal.workspace = true proxycast-credential.workspace = true proxycast-websocket.workspace = true proxycast-processor.workspace = true +proxycast-server-utils.workspace = true voice-core.workspace = true # Tauri diff --git a/src-tauri/crates/server-utils/Cargo.toml b/src-tauri/crates/server-utils/Cargo.toml new file mode 100644 index 000000000..e931406bf --- /dev/null +++ b/src-tauri/crates/server-utils/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "proxycast-server-utils" +version.workspace = true +edition.workspace = true + +[dependencies] +proxycast-core = { path = "../core" } +serde_json = { workspace = true } +axum = { workspace = true } +futures = { workspace = true } +tracing = { workspace = true } +uuid = { workspace = true } +regex = { workspace = true } + +[dev-dependencies] +proptest = { workspace = true } +tokio = { workspace = true } diff --git a/src-tauri/crates/server-utils/src/lib.rs b/src-tauri/crates/server-utils/src/lib.rs new file mode 100644 index 000000000..0814d3d67 --- /dev/null +++ b/src-tauri/crates/server-utils/src/lib.rs @@ -0,0 +1,1007 @@ +//! 服务器工具函数 crate +//! +//! 包含响应解析、字符串处理、响应构建等公共工具函数。 + +use axum::{ + body::Body, + http::{header, StatusCode}, + response::{IntoResponse, Response}, + Json, +}; +use futures::stream; +use proxycast_core::models::openai::{ContentPart, FunctionCall, MessageContent, ToolCall}; +use std::collections::HashMap; + +/// 从错误信息中解析 HTTP 状态码 +pub fn parse_error_status_code(error_message: &str) -> StatusCode { + if error_message.contains("429") { + StatusCode::TOO_MANY_REQUESTS + } else if error_message.contains("403") { + StatusCode::FORBIDDEN + } else if error_message.contains("401") { + StatusCode::UNAUTHORIZED + } else if error_message.contains("404") { + StatusCode::NOT_FOUND + } else if error_message.contains("400") { + StatusCode::BAD_REQUEST + } else if error_message.contains("503") { + StatusCode::SERVICE_UNAVAILABLE + } else if error_message.contains("502") { + StatusCode::BAD_GATEWAY + } else if error_message.contains("500") { + StatusCode::INTERNAL_SERVER_ERROR + } else { + StatusCode::INTERNAL_SERVER_ERROR + } +} + +/// 构建错误响应 +pub fn build_error_response(error_message: &str) -> Response { + let status_code = parse_error_status_code(error_message); + ( + status_code, + Json(serde_json::json!({ + "error": { + "message": error_message + } + })), + ) + .into_response() +} + +/// 从 HTTP 状态码构建错误响应 +pub fn build_error_response_with_status(status_code: u16, error_message: &str) -> Response { + let status = StatusCode::from_u16(status_code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + ( + status, + Json(serde_json::json!({ + "error": { + "message": error_message + } + })), + ) + .into_response() +} + +/// CodeWhisperer 响应解析结果 +#[derive(Debug, Default)] +pub struct CWParsedResponse { + pub content: String, + pub tool_calls: Vec, + pub usage_credits: f64, + pub context_usage_percentage: f64, +} + +impl CWParsedResponse { + /// 估算 Token 使用量 + #[allow(dead_code)] + pub fn estimate_tokens(&self) -> (u32, u32) { + let mut output_tokens: u32 = (self.content.len() / 4) as u32; + for tc in &self.tool_calls { + output_tokens += (tc.function.arguments.len() / 4) as u32; + } + let input_tokens = ((self.context_usage_percentage / 100.0) * 200000.0) as u32; + (input_tokens, output_tokens) + } +} + +/// 安全截断字符串到指定字符数，避免 UTF-8 边界问题 +pub fn safe_truncate(s: &str, max_chars: usize) -> String { + let chars: Vec = s.chars().collect(); + if chars.len() <= max_chars { + s.to_string() + } else { + chars[..max_chars].iter().collect() + } +} + +/// 计算 MessageContent 的字符长度 +pub fn message_content_len(content: &MessageContent) -> usize { + match content { + MessageContent::Text(s) => s.len(), + MessageContent::Parts(parts) => parts + .iter() + .filter_map(|p| { + if let ContentPart::Text { text } = p { + Some(text.len()) + } else { + None + } + }) + .sum(), + } +} + +/// 在字节数组中查找子序列 +pub fn find_subsequence(haystack: &[u8], needle: &[u8]) -> Option { + haystack + .windows(needle.len()) + .position(|window| window == needle) +} + +/// 从字节数组中提取 JSON 对象字符串 +pub fn extract_json_from_bytes(bytes: &[u8]) -> Option { + if bytes.is_empty() || bytes[0] != b'{' { + return None; + } + + let mut brace_count = 0; + let mut in_string = false; + let mut escape_next = false; + let mut end_pos = None; + + for (i, &b) in bytes.iter().enumerate() { + if escape_next { + escape_next = false; + continue; + } + match b { + b'\\' if in_string => escape_next = true, + b'"' => in_string = !in_string, + b'{' if !in_string => brace_count += 1, + b'}' if !in_string => { + brace_count -= 1; + if brace_count == 0 { + end_pos = Some(i + 1); + break; + } + } + _ => {} + } + } + + end_pos.and_then(|end| String::from_utf8(bytes[..end].to_vec()).ok()) +} + +/// 从字符串中提取完整的 JSON 对象 +#[allow(dead_code)] +pub fn extract_json_object(s: &str) -> Option<&str> { + if !s.starts_with('{') { + return None; + } + let mut brace_count = 0; + let mut in_string = false; + let mut escape_next = false; + + for (i, c) in s.char_indices() { + if escape_next { + escape_next = false; + continue; + } + match c { + '\\' if in_string => escape_next = true, + '"' => in_string = !in_string, + '{' if !in_string => brace_count += 1, + '}' if !in_string => { + brace_count -= 1; + if brace_count == 0 { + return Some(&s[..=i]); + } + } + _ => {} + } + } + None +} + +/// 解析 CodeWhisperer AWS Event Stream 响应 +pub fn parse_cw_response(body: &str) -> CWParsedResponse { + let mut result = CWParsedResponse::default(); + let mut tool_map: HashMap = HashMap::new(); + let bytes = body.as_bytes(); + + let json_patterns: &[&[u8]] = &[ + b"{\"content\":", + b"{\"name\":", + b"{\"input\":", + b"{\"stop\":", + b"{\"followupPrompt\":", + b"{\"toolUseId\":", + b"{\"unit\":", + b"{\"contextUsagePercentage\":", + ]; + + let mut pos = 0; + while pos < bytes.len() { + let mut next_start: Option = None; + + for pattern in json_patterns { + if let Some(idx) = find_subsequence(&bytes[pos..], pattern) { + let abs_pos = pos + idx; + if next_start.is_none_or(|start| abs_pos < start) { + next_start = Some(abs_pos); + } + } + } + + let start = match next_start { + Some(s) => s, + None => break, + }; + + if let Some(json_str) = extract_json_from_bytes(&bytes[start..]) { + if let Ok(value) = serde_json::from_str::(&json_str) { + if let Some(content) = value.get("content").and_then(|v| v.as_str()) { + if value.get("followupPrompt").is_none() { + result.content.push_str(content); + } + } else if let Some(tool_use_id) = value.get("toolUseId").and_then(|v| v.as_str()) { + let name = value + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let input_chunk = value + .get("input") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let is_stop = value.get("stop").and_then(|v| v.as_bool()).unwrap_or(false); + + let entry = tool_map + .entry(tool_use_id.to_string()) + .or_insert_with(|| (String::new(), String::new())); + if !name.is_empty() { + entry.0 = name; + } + entry.1.push_str(&input_chunk); + + if is_stop { + if let Some((name, input)) = tool_map.remove(tool_use_id) { + if !name.is_empty() { + result.tool_calls.push(ToolCall { + id: tool_use_id.to_string(), + call_type: "function".to_string(), + function: FunctionCall { + name, + arguments: input, + }, + }); + } + } + } + } else if value.get("stop").and_then(|v| v.as_bool()).unwrap_or(false) { + // no-op + } else if let Some(usage) = value.get("usage").and_then(|v| v.as_f64()) { + result.usage_credits = usage; + } else if let Some(ctx_usage) = + value.get("contextUsagePercentage").and_then(|v| v.as_f64()) + { + result.context_usage_percentage = ctx_usage; + } + } + pos = start + json_str.len(); + } else { + pos = start + 1; + } + } + + // 处理未完成的 tool calls + for (id, (name, input)) in tool_map { + if !name.is_empty() { + result.tool_calls.push(ToolCall { + id, + call_type: "function".to_string(), + function: FunctionCall { + name, + arguments: input, + }, + }); + } + } + + parse_bracket_tool_calls(&mut result); + result +} + +/// 解析 bracket 格式的 tool calls: [Called xxx with args: {...}] +pub fn parse_bracket_tool_calls(result: &mut CWParsedResponse) { + let re = + regex::Regex::new(r"\[Called\s+(\w+)\s+with\s+args:\s*(\{[^}]*(?:\{[^}]*\}[^}]*)*\})\]") + .ok(); + + if let Some(re) = re { + let mut to_remove = Vec::new(); + for cap in re.captures_iter(&result.content) { + if let (Some(name), Some(args)) = (cap.get(1), cap.get(2)) { + let tool_id = format!( + "call_{}", + &uuid::Uuid::new_v4().to_string().replace('-', "")[..8] + ); + result.tool_calls.push(ToolCall { + id: tool_id, + call_type: "function".to_string(), + function: FunctionCall { + name: name.as_str().to_string(), + arguments: args.as_str().to_string(), + }, + }); + if let Some(full_match) = cap.get(0) { + to_remove.push(full_match.as_str().to_string()); + } + } + } + for s in to_remove { + result.content = result.content.replace(&s, ""); + } + result.content = result.content.trim().to_string(); + } +} + +/// 构建 Anthropic 非流式响应 +pub fn build_anthropic_response(model: &str, parsed: &CWParsedResponse) -> Response { + let has_tool_calls = !parsed.tool_calls.is_empty(); + let mut content_array: Vec = Vec::new(); + + if !parsed.content.is_empty() { + content_array.push(serde_json::json!({ + "type": "text", + "text": parsed.content + })); + } + + for tc in &parsed.tool_calls { + let input: serde_json::Value = + serde_json::from_str(&tc.function.arguments).unwrap_or(serde_json::json!({})); + content_array.push(serde_json::json!({ + "type": "tool_use", + "id": tc.id, + "name": tc.function.name, + "input": input + })); + } + + if content_array.is_empty() { + content_array.push(serde_json::json!({"type": "text", "text": ""})); + } + + let mut output_tokens: u32 = (parsed.content.len() / 4) as u32; + for tc in &parsed.tool_calls { + output_tokens += (tc.function.arguments.len() / 4) as u32; + } + let input_tokens = ((parsed.context_usage_percentage / 100.0) * 200000.0) as u32; + + let response = serde_json::json!({ + "id": format!("msg_{}", uuid::Uuid::new_v4()), + "type": "message", + "role": "assistant", + "content": content_array, + "model": model, + "stop_reason": if has_tool_calls { "tool_use" } else { "end_turn" }, + "stop_sequence": null, + "usage": { + "input_tokens": input_tokens, + "output_tokens": output_tokens + } + }); + Json(response).into_response() +} + +/// 构建 Anthropic 流式响应 (SSE) +pub fn build_anthropic_stream_response(model: &str, parsed: &CWParsedResponse) -> Response { + let has_tool_calls = !parsed.tool_calls.is_empty(); + let message_id = format!("msg_{}", uuid::Uuid::new_v4()); + let model = model.to_string(); + let content = parsed.content.clone(); + let tool_calls = parsed.tool_calls.clone(); + + let mut output_tokens: u32 = (parsed.content.len() / 4) as u32; + for tc in &parsed.tool_calls { + output_tokens += (tc.function.arguments.len() / 4) as u32; + } + let input_tokens = ((parsed.context_usage_percentage / 100.0) * 200000.0) as u32; + + let mut events: Vec = Vec::new(); + + // 1. message_start + let message_start = serde_json::json!({ + "type": "message_start", + "message": { + "id": message_id, "type": "message", "role": "assistant", + "model": model, "content": [], "stop_reason": null, + "stop_sequence": null, + "usage": {"input_tokens": input_tokens, "output_tokens": 0} + } + }); + events.push(format!("event: message_start\ndata: {message_start}\n\n")); + + let mut block_index = 0; + + // 2. 文本内容块 + let block_start = serde_json::json!({ + "type": "content_block_start", "index": block_index, + "content_block": {"type": "text", "text": ""} + }); + events.push(format!( + "event: content_block_start\ndata: {block_start}\n\n" + )); + + if !content.is_empty() { + let block_delta = serde_json::json!({ + "type": "content_block_delta", "index": block_index, + "delta": {"type": "text_delta", "text": content} + }); + events.push(format!( + "event: content_block_delta\ndata: {block_delta}\n\n" + )); + } + + let block_stop = serde_json::json!({"type": "content_block_stop", "index": block_index}); + events.push(format!("event: content_block_stop\ndata: {block_stop}\n\n")); + block_index += 1; + + // 3. Tool use 块 + for tc in &tool_calls { + let block_start = serde_json::json!({ + "type": "content_block_start", "index": block_index, + "content_block": { + "type": "tool_use", "id": tc.id, + "name": tc.function.name, "input": {} + } + }); + events.push(format!( + "event: content_block_start\ndata: {block_start}\n\n" + )); + + let partial_json = if tc.function.arguments.is_empty() { + "{}".to_string() + } else { + tc.function.arguments.clone() + }; + let block_delta = serde_json::json!({ + "type": "content_block_delta", "index": block_index, + "delta": {"type": "input_json_delta", "partial_json": partial_json} + }); + events.push(format!( + "event: content_block_delta\ndata: {block_delta}\n\n" + )); + + let block_stop = serde_json::json!({"type": "content_block_stop", "index": block_index}); + events.push(format!("event: content_block_stop\ndata: {block_stop}\n\n")); + block_index += 1; + } + + // 4. message_delta + let message_delta = serde_json::json!({ + "type": "message_delta", + "delta": { + "stop_reason": if has_tool_calls { "tool_use" } else { "end_turn" }, + "stop_sequence": null + }, + "usage": {"output_tokens": output_tokens} + }); + events.push(format!("event: message_delta\ndata: {message_delta}\n\n")); + + // 5. message_stop + let message_stop = serde_json::json!({"type": "message_stop"}); + events.push(format!("event: message_stop\ndata: {message_stop}\n\n")); + + let body_stream = stream::iter(events.into_iter().map(Ok::<_, std::convert::Infallible>)); + let body = Body::from_stream(body_stream); + + Response::builder() + .status(StatusCode::OK) + .header(header::CONTENT_TYPE, "text/event-stream") + .header(header::CACHE_CONTROL, "no-cache") + .header(header::CONNECTION, "keep-alive") + .body(body) + .unwrap_or_else(|e| { + tracing::error!("Failed to build SSE response: {}", e); + Response::builder() + .status(StatusCode::INTERNAL_SERVER_ERROR) + .body(Body::empty()) + .unwrap_or_default() + }) +} + +/// 构建 Gemini CLI OAuth 请求体 +pub fn build_gemini_cli_request( + request: &serde_json::Value, + model: &str, + project_id: &str, +) -> serde_json::Value { + let enable_thinking = model.ends_with("-thinking") + || model == "gemini-2.5-pro" + || model.starts_with("gemini-3-pro-"); + + let mut inner_request = request.clone(); + + if inner_request.get("generationConfig").is_none() { + inner_request["generationConfig"] = serde_json::json!({ + "temperature": 1.0, "maxOutputTokens": 8096, + "topP": 0.85, "topK": 50, "candidateCount": 1, + "thinkingConfig": { + "includeThoughts": enable_thinking, + "thinkingBudget": if enable_thinking { 1024 } else { 0 } + } + }); + } else if inner_request["generationConfig"] + .get("thinkingConfig") + .is_none() + { + inner_request["generationConfig"]["thinkingConfig"] = serde_json::json!({ + "includeThoughts": enable_thinking, + "thinkingBudget": if enable_thinking { 1024 } else { 0 } + }); + } + + if let Some(obj) = inner_request.as_object_mut() { + obj.remove("safetySettings"); + } + + serde_json::json!({ + "project": project_id, + "model": model, + "request": inner_request + }) +} + +/// 构建 Gemini 原生请求体 +pub fn build_gemini_native_request( + request: &serde_json::Value, + model: &str, + project_id: &str, +) -> serde_json::Value { + let actual_model = match model { + "gemini-2.5-computer-use-preview-10-2025" => "rev19-uic3-1p", + "gemini-3-pro-image-preview" => "gemini-3-pro-image", + "gemini-3-pro-preview" => "gemini-3-pro-high", + "gemini-claude-sonnet-4-5" => "claude-sonnet-4-5", + "gemini-claude-sonnet-4-5-thinking" => "claude-sonnet-4-5-thinking", + _ => model, + }; + + let enable_thinking = model.ends_with("-thinking") + || model == "gemini-2.5-pro" + || model.starts_with("gemini-3-pro-") + || model == "rev19-uic3-1p" + || model == "gpt-oss-120b-medium"; + + let request_id = format!("agent-{}", uuid::Uuid::new_v4()); + let session_id = { + let uuid = uuid::Uuid::new_v4(); + let bytes = uuid.as_bytes(); + let n: u64 = u64::from_le_bytes([ + bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7], + ]) % 9_000_000_000_000_000_000; + format!("-{n}") + }; + + let mut inner_request = request.clone(); + inner_request["sessionId"] = serde_json::json!(session_id); + + if inner_request.get("generationConfig").is_none() { + inner_request["generationConfig"] = serde_json::json!({ + "temperature": 1.0, "maxOutputTokens": 8096, + "topP": 0.85, "topK": 50, "candidateCount": 1, + "stopSequences": [ + "<|user|>", "<|bot|>", "<|context_request|>", + "<|endoftext|>", "<|end_of_turn|>" + ], + "thinkingConfig": { + "includeThoughts": enable_thinking, + "thinkingBudget": if enable_thinking { 1024 } else { 0 } + } + }); + } else if inner_request["generationConfig"] + .get("thinkingConfig") + .is_none() + { + inner_request["generationConfig"]["thinkingConfig"] = serde_json::json!({ + "includeThoughts": enable_thinking, + "thinkingBudget": if enable_thinking { 1024 } else { 0 } + }); + } + + if let Some(obj) = inner_request.as_object_mut() { + obj.remove("safetySettings"); + } + + serde_json::json!({ + "project": project_id, + "requestId": request_id, + "request": inner_request, + "model": actual_model, + "userAgent": "antigravity" + }) +} + +/// 健康检查端点响应 +pub async fn health() -> impl IntoResponse { + Json(serde_json::json!({ + "status": "healthy", + "version": env!("CARGO_PKG_VERSION") + })) +} + +/// 模型列表端点响应 +pub async fn models() -> impl IntoResponse { + Json(serde_json::json!({ + "object": "list", + "data": [ + {"id": "claude-sonnet-4-5", "object": "model", "owned_by": "anthropic"}, + {"id": "claude-sonnet-4-5-20250929", "object": "model", "owned_by": "anthropic"}, + {"id": "gemini-3-pro-preview", "object": "model", "owned_by": "google"}, + {"id": "gemini-3-pro-image-preview", "object": "model", "owned_by": "google"}, + {"id": "gemini-3-flash-preview", "object": "model", "owned_by": "google"}, + {"id": "gemini-2.5-computer-use-preview-10-2025", "object": "model", "owned_by": "google"}, + {"id": "gemini-claude-sonnet-4-5", "object": "model", "owned_by": "google"}, + {"id": "gemini-claude-sonnet-4-5-thinking", "object": "model", "owned_by": "google"}, + {"id": "gemini-claude-opus-4-5-thinking", "object": "model", "owned_by": "google"}, + {"id": "qwen3-coder-plus", "object": "model", "owned_by": "alibaba"}, + {"id": "qwen3-coder-flash", "object": "model", "owned_by": "alibaba"} + ] + })) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_safe_truncate() { + assert_eq!(safe_truncate("hello", 10), "hello"); + assert_eq!(safe_truncate("hello world", 5), "hello"); + assert_eq!(safe_truncate("你好世界", 2), "你好"); + } + + #[test] + fn test_find_subsequence() { + let haystack = b"hello world"; + assert_eq!(find_subsequence(haystack, b"world"), Some(6)); + assert_eq!(find_subsequence(haystack, b"foo"), None); + } + + #[test] + fn test_extract_json_from_bytes() { + let json = b"{\"key\":\"value\"}"; + assert_eq!( + extract_json_from_bytes(json), + Some("{\"key\":\"value\"}".to_string()) + ); + let nested = b"{\"outer\":{\"inner\":\"value\"}}"; + assert_eq!( + extract_json_from_bytes(nested), + Some("{\"outer\":{\"inner\":\"value\"}}".to_string()) + ); + assert_eq!(extract_json_from_bytes(b"not json"), None); + } +} + +#[cfg(test)] +mod property_tests { + use super::*; + use axum::http::header; + use proptest::prelude::*; + + fn arb_text_content() -> impl Strategy { + "[a-zA-Z0-9 .,!?\\n]{0,500}".prop_map(|s| s) + } + + fn arb_model_name() -> impl Strategy { + prop_oneof![ + Just("claude-3-sonnet".to_string()), + Just("claude-3-opus".to_string()), + Just("claude-3-haiku".to_string()), + Just("claude-sonnet-4-5".to_string()), + Just("claude-3-5-sonnet-latest".to_string()), + ] + } + + fn arb_tool_call() -> impl Strategy { + ( + "[a-z0-9]{8,16}", + prop_oneof![ + Just("read_file".to_string()), + Just("write_file".to_string()), + Just("execute_command".to_string()), + Just("search".to_string()), + ], + prop_oneof![ + Just("{}".to_string()), + Just("{\"path\":\"/tmp/test\"}".to_string()), + Just("{\"content\":\"hello\"}".to_string()), + Just("{\"query\":\"test\",\"limit\":10}".to_string()), + ], + ) + .prop_map(|(id, name, args)| ToolCall { + id: format!("call_{id}"), + call_type: "function".to_string(), + function: FunctionCall { + name, + arguments: args, + }, + }) + } + + fn arb_cw_parsed_response() -> impl Strategy { + ( + arb_text_content(), + prop::collection::vec(arb_tool_call(), 0..3), + 0.0f64..100.0f64, + 0.0f64..100.0f64, + ) + .prop_map( + |(content, tool_calls, usage_credits, context_usage_percentage)| CWParsedResponse { + content, + tool_calls, + usage_credits, + context_usage_percentage, + }, + ) + } + + proptest! { + #[test] + fn prop_non_streaming_response_format( + model in arb_model_name(), + parsed in arb_cw_parsed_response() + ) { + let response = build_anthropic_response(&model, &parsed); + let (parts, _body) = response.into_parts(); + prop_assert_eq!(parts.status, StatusCode::OK); + + let content_type = parts.headers.get(header::CONTENT_TYPE); + prop_assert!(content_type.is_some()); + prop_assert!(content_type.unwrap().to_str().unwrap().contains("application/json")); + } + + #[test] + fn prop_non_streaming_response_empty_content(model in arb_model_name()) { + let parsed = CWParsedResponse { + content: String::new(), tool_calls: Vec::new(), + usage_credits: 0.0, context_usage_percentage: 0.0, + }; + let response = build_anthropic_response(&model, &parsed); + let (parts, _body) = response.into_parts(); + prop_assert_eq!(parts.status, StatusCode::OK); + } + + #[test] + fn prop_non_streaming_response_tool_calls_only( + model in arb_model_name(), + tool_calls in prop::collection::vec(arb_tool_call(), 1..3) + ) { + let parsed = CWParsedResponse { + content: String::new(), tool_calls, + usage_credits: 0.0, context_usage_percentage: 50.0, + }; + let response = build_anthropic_response(&model, &parsed); + let (parts, _body) = response.into_parts(); + prop_assert_eq!(parts.status, StatusCode::OK); + prop_assert!(!parsed.tool_calls.is_empty()); + } + + #[test] + fn prop_token_estimation( + content in arb_text_content(), + context_percentage in 0.0f64..100.0f64 + ) { + let parsed = CWParsedResponse { + content: content.clone(), tool_calls: Vec::new(), + usage_credits: 0.0, context_usage_percentage: context_percentage, + }; + let (input_tokens, output_tokens) = parsed.estimate_tokens(); + let expected_output = (content.len() / 4) as u32; + prop_assert_eq!(output_tokens, expected_output); + let expected_input = ((context_percentage / 100.0) * 200000.0) as u32; + prop_assert_eq!(input_tokens, expected_input); + } + } + + fn get_model_list_data() -> Vec { + vec![ + serde_json::json!({"id": "claude-sonnet-4-5", "object": "model", "owned_by": "anthropic"}), + serde_json::json!({"id": "claude-sonnet-4-5-20250929", "object": "model", "owned_by": "anthropic"}), + serde_json::json!({"id": "gemini-3-pro-preview", "object": "model", "owned_by": "google"}), + serde_json::json!({"id": "gemini-3-pro-image-preview", "object": "model", "owned_by": "google"}), + serde_json::json!({"id": "gemini-3-flash-preview", "object": "model", "owned_by": "google"}), + serde_json::json!({"id": "gemini-2.5-computer-use-preview-10-2025", "object": "model", "owned_by": "google"}), + serde_json::json!({"id": "gemini-claude-sonnet-4-5", "object": "model", "owned_by": "google"}), + serde_json::json!({"id": "gemini-claude-sonnet-4-5-thinking", "object": "model", "owned_by": "google"}), + serde_json::json!({"id": "gemini-claude-opus-4-5-thinking", "object": "model", "owned_by": "google"}), + serde_json::json!({"id": "qwen3-coder-plus", "object": "model", "owned_by": "alibaba"}), + serde_json::json!({"id": "qwen3-coder-flash", "object": "model", "owned_by": "alibaba"}), + ] + } + + #[test] + fn prop_model_list_structure_and_ownership() { + let models = get_model_list_data(); + + for model in &models { + let id = model.get("id").and_then(|v| v.as_str()); + assert!(id.is_some(), "Model should have id field"); + assert!(!id.unwrap().is_empty(), "Model id should not be empty"); + + let object = model.get("object").and_then(|v| v.as_str()); + assert_eq!(object, Some("model"), "Model object should be 'model'"); + + let owned_by = model.get("owned_by").and_then(|v| v.as_str()); + assert!(owned_by.is_some(), "Model should have owned_by field"); + + let model_id = id.unwrap(); + let owner = owned_by.unwrap(); + + if model_id.starts_with("gemini-") { + assert_eq!(owner, "google", "Gemini models should be owned by google"); + } else if model_id.starts_with("claude-") { + assert_eq!( + owner, "anthropic", + "Claude models should be owned by anthropic" + ); + } else if model_id.starts_with("qwen") { + assert_eq!(owner, "alibaba", "Qwen models should be owned by alibaba"); + } + } + } + + #[test] + fn test_antigravity_models_present() { + let models = get_model_list_data(); + let model_ids: Vec<&str> = models + .iter() + .filter_map(|m| m.get("id").and_then(|v| v.as_str())) + .collect(); + + let required_models = [ + "gemini-3-pro-preview", + "gemini-3-pro-image-preview", + "gemini-3-flash-preview", + "gemini-2.5-computer-use-preview-10-2025", + "gemini-claude-sonnet-4-5", + "gemini-claude-sonnet-4-5-thinking", + "gemini-claude-opus-4-5-thinking", + ]; + + for required in &required_models { + assert!( + model_ids.contains(required), + "Model {required} should be in the list" + ); + } + } + + #[allow(dead_code)] + fn get_expected_model_mapping(model: &str) -> &str { + match model { + "gemini-2.5-computer-use-preview-10-2025" => "rev19-uic3-1p", + "gemini-3-pro-image-preview" => "gemini-3-pro-image", + "gemini-3-pro-preview" => "gemini-3-pro-high", + "gemini-claude-sonnet-4-5" => "claude-sonnet-4-5", + "gemini-claude-sonnet-4-5-thinking" => "claude-sonnet-4-5-thinking", + _ => model, + } + } + + #[test] + fn prop_model_name_mapping_correctness() { + let test_request = serde_json::json!({ + "contents": [{"role": "user", "parts": [{"text": "test"}]}] + }); + let project_id = "test-project"; + + let known_mappings = [ + ("gemini-2.5-computer-use-preview-10-2025", "rev19-uic3-1p"), + ("gemini-3-pro-image-preview", "gemini-3-pro-image"), + ("gemini-3-pro-preview", "gemini-3-pro-high"), + ("gemini-claude-sonnet-4-5", "claude-sonnet-4-5"), + ( + "gemini-claude-sonnet-4-5-thinking", + "claude-sonnet-4-5-thinking", + ), + ]; + + for (input, expected) in &known_mappings { + let result = build_gemini_native_request(&test_request, input, project_id); + let actual_model = result.get("model").and_then(|v| v.as_str()).unwrap(); + assert_eq!( + actual_model, *expected, + "Model {input} should map to {expected}" + ); + } + + let unknown_models = ["gemini-2.0-flash", "gemini-2.5-flash", "custom-model"]; + for model in &unknown_models { + let result = build_gemini_native_request(&test_request, model, project_id); + let actual_model = result.get("model").and_then(|v| v.as_str()).unwrap(); + assert_eq!( + actual_model, *model, + "Unknown model {model} should be returned unchanged" + ); + } + } + + fn should_enable_thinking(model: &str) -> bool { + model.ends_with("-thinking") + || model == "gemini-2.5-pro" + || model.starts_with("gemini-3-pro-") + || model == "rev19-uic3-1p" + || model == "gpt-oss-120b-medium" + } + + #[test] + fn prop_thinking_mode_enablement_logic() { + let thinking_enabled_models = [ + "gemini-claude-sonnet-4-5-thinking", + "gemini-claude-opus-4-5-thinking", + "gemini-3-pro-preview", + "gemini-3-pro-image-preview", + "gemini-3-pro-high", + "rev19-uic3-1p", + "gpt-oss-120b-medium", + ]; + + for model in &thinking_enabled_models { + assert!( + should_enable_thinking(model), + "Model {model} should have thinking enabled" + ); + } + + let thinking_disabled_models = [ + "gemini-claude-sonnet-4-5", + "claude-sonnet-4-5", + "custom-model", + ]; + + for model in &thinking_disabled_models { + assert!( + !should_enable_thinking(model), + "Model {model} should have thinking disabled" + ); + } + } + + #[test] + fn prop_thinking_configuration_values() { + let test_request = serde_json::json!({ + "contents": [{"role": "user", "parts": [{"text": "test"}]}] + }); + let project_id = "test-project"; + + let thinking_enabled_models = [ + "gemini-3-pro-preview", + "gemini-2.5-pro", + "gemini-claude-sonnet-4-5-thinking", + ]; + + for model in &thinking_enabled_models { + let result = build_gemini_native_request(&test_request, model, project_id); + let thinking_config = &result["request"]["generationConfig"]["thinkingConfig"]; + + assert_eq!( + thinking_config["includeThoughts"].as_bool(), + Some(true), + "Model {model} should have includeThoughts=true" + ); + assert_eq!( + thinking_config["thinkingBudget"].as_i64(), + Some(1024), + "Model {model} should have thinkingBudget=1024" + ); + } + + let thinking_disabled_models = [ + "gemini-2.0-flash", + "gemini-2.5-flash", + "gemini-claude-sonnet-4-5", + ]; + + for model in &thinking_disabled_models { + let result = build_gemini_native_request(&test_request, model, project_id); + let thinking_config = &result["request"]["generationConfig"]["thinkingConfig"]; + + assert_eq!( + thinking_config["includeThoughts"].as_bool(), + Some(false), + "Model {model} should have includeThoughts=false" + ); + assert_eq!( + thinking_config["thinkingBudget"].as_i64(), + Some(0), + "Model {model} should have thinkingBudget=0" + ); + } + } +} diff --git a/src-tauri/src/server_utils.rs b/src-tauri/src/server_utils.rs index a9de576ec..8323e1c2c 100644 --- a/src-tauri/src/server_utils.rs +++ b/src-tauri/src/server_utils.rs @@ -1,1403 +1,5 @@ //! 服务器工具函数 //! -//! 包含响应解析、字符串处理、响应构建等公共工具函数。 +//! 已迁移至 proxycast-server-utils crate，此文件为 re-export 层。 -use crate::models::openai::{ContentPart, FunctionCall, MessageContent, ToolCall}; -use axum::{ - body::Body, - http::{header, StatusCode}, - response::{IntoResponse, Response}, - Json, -}; -use futures::stream; -use std::collections::HashMap; - -/// 从错误信息中解析 HTTP 状态码 -/// -/// 用于将上游 API 返回的错误状态码透传给客户端，而不是统一返回 500。 -/// 支持解析常见的 HTTP 状态码：429、403、401、404、400、503、502、500。 -/// -/// # 参数 -/// - `error_message`: 错误信息字符串，通常包含状态码（如 "API call failed: 429 - ..."） -/// -/// # 返回 -/// 解析出的 HTTP 状态码，如果无法解析则返回 500 INTERNAL_SERVER_ERROR -/// -/// # 示例 -/// ``` -/// let status = parse_error_status_code("API call failed: 429 Too Many Requests"); -/// assert_eq!(status, StatusCode::TOO_MANY_REQUESTS); -/// ``` -pub fn parse_error_status_code(error_message: &str) -> StatusCode { - if error_message.contains("429") { - StatusCode::TOO_MANY_REQUESTS - } else if error_message.contains("403") { - StatusCode::FORBIDDEN - } else if error_message.contains("401") { - StatusCode::UNAUTHORIZED - } else if error_message.contains("404") { - StatusCode::NOT_FOUND - } else if error_message.contains("400") { - StatusCode::BAD_REQUEST - } else if error_message.contains("503") { - StatusCode::SERVICE_UNAVAILABLE - } else if error_message.contains("502") { - StatusCode::BAD_GATEWAY - } else if error_message.contains("500") { - StatusCode::INTERNAL_SERVER_ERROR - } else { - StatusCode::INTERNAL_SERVER_ERROR - } -} - -/// 构建错误响应 -/// -/// 从错误信息中解析状态码并构建标准的 JSON 错误响应。 -/// -/// # 参数 -/// - `error_message`: 错误信息字符串 -/// -/// # 返回 -/// 包含正确状态码的 HTTP 响应 -pub fn build_error_response(error_message: &str) -> Response { - let status_code = parse_error_status_code(error_message); - ( - status_code, - Json(serde_json::json!({ - "error": { - "message": error_message - } - })), - ) - .into_response() -} - -/// 从 HTTP 状态码构建错误响应 -/// -/// 直接使用状态码构建响应，无需解析字符串。 -/// 适用于已知状态码的场景（如 AntigravityApiError）。 -/// -/// # 参数 -/// - `status_code`: HTTP 状态码（u16） -/// - `error_message`: 错误信息字符串 -/// -/// # 返回 -/// 包含指定状态码的 HTTP 响应 -pub fn build_error_response_with_status(status_code: u16, error_message: &str) -> Response { - let status = StatusCode::from_u16(status_code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); - ( - status, - Json(serde_json::json!({ - "error": { - "message": error_message - } - })), - ) - .into_response() -} - -/// CodeWhisperer 响应解析结果 -#[derive(Debug, Default)] -pub struct CWParsedResponse { - pub content: String, - pub tool_calls: Vec, - pub usage_credits: f64, - pub context_usage_percentage: f64, -} - -impl CWParsedResponse { - /// 估算 Token 使用量 - /// - /// 基于响应内容长度和上下文使用百分比估算 Token 数量： - /// - output_tokens: 基于响应内容长度（约 4 字符 = 1 token） - /// - input_tokens: 基于 context_usage_percentage（假设 100% = 200k tokens） - /// - /// # 返回 - /// (input_tokens, output_tokens) 元组 - #[allow(dead_code)] - pub fn estimate_tokens(&self) -> (u32, u32) { - // 估算 output tokens: 基于响应内容长度 (约 4 字符 = 1 token) - let mut output_tokens: u32 = (self.content.len() / 4) as u32; - for tc in &self.tool_calls { - output_tokens += (tc.function.arguments.len() / 4) as u32; - } - - // 从 context_usage_percentage 估算 input tokens - // 假设 100% = 200k tokens (Claude 的上下文窗口) - let input_tokens = ((self.context_usage_percentage / 100.0) * 200000.0) as u32; - - (input_tokens, output_tokens) - } -} - -/// 安全截断字符串到指定字符数，避免 UTF-8 边界问题 -pub fn safe_truncate(s: &str, max_chars: usize) -> String { - let chars: Vec = s.chars().collect(); - if chars.len() <= max_chars { - s.to_string() - } else { - chars[..max_chars].iter().collect() - } -} - -/// 计算 MessageContent 的字符长度 -pub fn message_content_len(content: &MessageContent) -> usize { - match content { - MessageContent::Text(s) => s.len(), - MessageContent::Parts(parts) => parts - .iter() - .filter_map(|p| { - if let ContentPart::Text { text } = p { - Some(text.len()) - } else { - None - } - }) - .sum(), - } -} - -/// 解析 CodeWhisperer AWS Event Stream 响应 -/// -/// AWS Event Stream 是二进制格式，JSON payload 嵌入在二进制头部之间 -pub fn parse_cw_response(body: &str) -> CWParsedResponse { - let mut result = CWParsedResponse::default(); - // 使用 HashMap 来跟踪多个并发的 tool calls - // key: toolUseId, value: (name, input_accumulated) - let mut tool_map: HashMap = HashMap::new(); - - // 将字符串转换为字节，因为 AWS Event Stream 包含二进制数据 - let bytes = body.as_bytes(); - - // 搜索所有 JSON 对象的模式 - // AWS Event Stream 格式: [binary headers]{"content":"..."}[binary trailer] - let json_patterns: &[&[u8]] = &[ - b"{\"content\":", - b"{\"name\":", - b"{\"input\":", - b"{\"stop\":", - b"{\"followupPrompt\":", - b"{\"toolUseId\":", - b"{\"unit\":", // meteringEvent - b"{\"contextUsagePercentage\":", // contextUsageEvent - ]; - - let mut pos = 0; - while pos < bytes.len() { - // 找到下一个 JSON 对象的开始 - let mut next_start: Option = None; - - for pattern in json_patterns { - if let Some(idx) = find_subsequence(&bytes[pos..], pattern) { - let abs_pos = pos + idx; - if next_start.is_none_or(|start| abs_pos < start) { - next_start = Some(abs_pos); - } - } - } - - let start = match next_start { - Some(s) => s, - None => break, - }; - - // 从 start 位置提取完整的 JSON 对象 - if let Some(json_str) = extract_json_from_bytes(&bytes[start..]) { - if let Ok(value) = serde_json::from_str::(&json_str) { - // 处理 content 事件 - if let Some(content) = value.get("content").and_then(|v| v.as_str()) { - // 跳过 followupPrompt - if value.get("followupPrompt").is_none() { - result.content.push_str(content); - } - } - // 处理 tool use 事件 (包含 toolUseId) - else if let Some(tool_use_id) = value.get("toolUseId").and_then(|v| v.as_str()) { - let name = value - .get("name") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let input_chunk = value - .get("input") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(); - let is_stop = value.get("stop").and_then(|v| v.as_bool()).unwrap_or(false); - - // 获取或创建 tool entry - let entry = tool_map - .entry(tool_use_id.to_string()) - .or_insert_with(|| (String::new(), String::new())); - - // 更新 name（如果有） - if !name.is_empty() { - entry.0 = name; - } - - // 累积 input - entry.1.push_str(&input_chunk); - - // 如果是 stop 事件，完成这个 tool call - if is_stop { - if let Some((name, input)) = tool_map.remove(tool_use_id) { - if !name.is_empty() { - result.tool_calls.push(ToolCall { - id: tool_use_id.to_string(), - call_type: "function".to_string(), - function: FunctionCall { - name, - arguments: input, - }, - }); - } - } - } - } - // 处理独立的 stop 事件（没有 toolUseId）- 这种情况不应该发生，但以防万一 - else if value.get("stop").and_then(|v| v.as_bool()).unwrap_or(false) { - // no-op - } - // 处理 meteringEvent: {"unit":"credit","unitPlural":"credits","usage":0.34} - else if let Some(usage) = value.get("usage").and_then(|v| v.as_f64()) { - result.usage_credits = usage; - } - // 处理 contextUsageEvent: {"contextUsagePercentage":54.36} - else if let Some(ctx_usage) = - value.get("contextUsagePercentage").and_then(|v| v.as_f64()) - { - result.context_usage_percentage = ctx_usage; - } - } - pos = start + json_str.len(); - } else { - pos = start + 1; - } - } - - // 处理未完成的 tool calls（没有收到 stop 事件的） - for (id, (name, input)) in tool_map { - if !name.is_empty() { - result.tool_calls.push(ToolCall { - id, - call_type: "function".to_string(), - function: FunctionCall { - name, - arguments: input, - }, - }); - } - } - - // 解析 bracket 格式的 tool calls: [Called xxx with args: {...}] - parse_bracket_tool_calls(&mut result); - - result -} - -/// 在字节数组中查找子序列 -pub fn find_subsequence(haystack: &[u8], needle: &[u8]) -> Option { - haystack - .windows(needle.len()) - .position(|window| window == needle) -} - -/// 从字节数组中提取 JSON 对象字符串 -pub fn extract_json_from_bytes(bytes: &[u8]) -> Option { - if bytes.is_empty() || bytes[0] != b'{' { - return None; - } - - let mut brace_count = 0; - let mut in_string = false; - let mut escape_next = false; - let mut end_pos = None; - - for (i, &b) in bytes.iter().enumerate() { - if escape_next { - escape_next = false; - continue; - } - - match b { - b'\\' if in_string => escape_next = true, - b'"' => in_string = !in_string, - b'{' if !in_string => brace_count += 1, - b'}' if !in_string => { - brace_count -= 1; - if brace_count == 0 { - end_pos = Some(i + 1); - break; - } - } - _ => {} - } - } - - end_pos.and_then(|end| String::from_utf8(bytes[..end].to_vec()).ok()) -} - -/// 从字符串中提取完整的 JSON 对象 (保留用于兼容) -#[allow(dead_code)] -pub fn extract_json_object(s: &str) -> Option<&str> { - if !s.starts_with('{') { - return None; - } - - let mut brace_count = 0; - let mut in_string = false; - let mut escape_next = false; - - for (i, c) in s.char_indices() { - if escape_next { - escape_next = false; - continue; - } - - match c { - '\\' if in_string => escape_next = true, - '"' => in_string = !in_string, - '{' if !in_string => brace_count += 1, - '}' if !in_string => { - brace_count -= 1; - if brace_count == 0 { - return Some(&s[..=i]); - } - } - _ => {} - } - } - None -} - -/// 解析 bracket 格式的 tool calls -/// -/// 格式: [Called xxx with args: {...}] -pub fn parse_bracket_tool_calls(result: &mut CWParsedResponse) { - let re = - regex::Regex::new(r"\[Called\s+(\w+)\s+with\s+args:\s*(\{[^}]*(?:\{[^}]*\}[^}]*)*\})\]") - .ok(); - - if let Some(re) = re { - let mut to_remove = Vec::new(); - for cap in re.captures_iter(&result.content) { - if let (Some(name), Some(args)) = (cap.get(1), cap.get(2)) { - let tool_id = format!( - "call_{}", - &uuid::Uuid::new_v4().to_string().replace('-', "")[..8] - ); - result.tool_calls.push(ToolCall { - id: tool_id, - call_type: "function".to_string(), - function: FunctionCall { - name: name.as_str().to_string(), - arguments: args.as_str().to_string(), - }, - }); - if let Some(full_match) = cap.get(0) { - to_remove.push(full_match.as_str().to_string()); - } - } - } - // 从 content 中移除 tool call 文本 - for s in to_remove { - result.content = result.content.replace(&s, ""); - } - result.content = result.content.trim().to_string(); - } -} - -/// 构建 Anthropic 非流式响应 -pub fn build_anthropic_response(model: &str, parsed: &CWParsedResponse) -> Response { - let has_tool_calls = !parsed.tool_calls.is_empty(); - let mut content_array: Vec = Vec::new(); - - if !parsed.content.is_empty() { - content_array.push(serde_json::json!({ - "type": "text", - "text": parsed.content - })); - } - - for tc in &parsed.tool_calls { - let input: serde_json::Value = - serde_json::from_str(&tc.function.arguments).unwrap_or(serde_json::json!({})); - content_array.push(serde_json::json!({ - "type": "tool_use", - "id": tc.id, - "name": tc.function.name, - "input": input - })); - } - - if content_array.is_empty() { - content_array.push(serde_json::json!({"type": "text", "text": ""})); - } - - // 估算 output tokens: 基于响应内容长度 (约 4 字符 = 1 token) - let mut output_tokens: u32 = (parsed.content.len() / 4) as u32; - for tc in &parsed.tool_calls { - output_tokens += (tc.function.arguments.len() / 4) as u32; - } - // 从 context_usage_percentage 估算 input tokens - // 假设 100% = 200k tokens (Claude 的上下文窗口) - let input_tokens = ((parsed.context_usage_percentage / 100.0) * 200000.0) as u32; - - let response = serde_json::json!({ - "id": format!("msg_{}", uuid::Uuid::new_v4()), - "type": "message", - "role": "assistant", - "content": content_array, - "model": model, - "stop_reason": if has_tool_calls { "tool_use" } else { "end_turn" }, - "stop_sequence": null, - "usage": { - "input_tokens": input_tokens, - "output_tokens": output_tokens - } - }); - Json(response).into_response() -} - -/// 构建 Anthropic 流式响应 (SSE) -pub fn build_anthropic_stream_response(model: &str, parsed: &CWParsedResponse) -> Response { - let has_tool_calls = !parsed.tool_calls.is_empty(); - let message_id = format!("msg_{}", uuid::Uuid::new_v4()); - let model = model.to_string(); - let content = parsed.content.clone(); - let tool_calls = parsed.tool_calls.clone(); - - // 估算 output tokens: 基于响应内容长度 (约 4 字符 = 1 token) - let mut output_tokens: u32 = (parsed.content.len() / 4) as u32; - for tc in &parsed.tool_calls { - output_tokens += (tc.function.arguments.len() / 4) as u32; - } - // 从 context_usage_percentage 估算 input tokens - let input_tokens = ((parsed.context_usage_percentage / 100.0) * 200000.0) as u32; - - // 构建 SSE 事件流 - let mut events: Vec = Vec::new(); - - // 1. message_start - let message_start = serde_json::json!({ - "type": "message_start", - "message": { - "id": message_id, - "type": "message", - "role": "assistant", - "model": model, - "content": [], - "stop_reason": null, - "stop_sequence": null, - "usage": {"input_tokens": input_tokens, "output_tokens": 0} - } - }); - events.push(format!("event: message_start\ndata: {message_start}\n\n")); - - let mut block_index = 0; - - // 2. 文本内容块 - 即使为空也要发送，Claude Code 需要至少一个 content block - // content_block_start - let block_start = serde_json::json!({ - "type": "content_block_start", - "index": block_index, - "content_block": {"type": "text", "text": ""} - }); - events.push(format!( - "event: content_block_start\ndata: {block_start}\n\n" - )); - - if !content.is_empty() { - // content_block_delta - 发送完整内容 - let block_delta = serde_json::json!({ - "type": "content_block_delta", - "index": block_index, - "delta": {"type": "text_delta", "text": content} - }); - events.push(format!( - "event: content_block_delta\ndata: {block_delta}\n\n" - )); - } - - // content_block_stop - let block_stop = serde_json::json!({ - "type": "content_block_stop", - "index": block_index - }); - events.push(format!("event: content_block_stop\ndata: {block_stop}\n\n")); - - block_index += 1; - - // 3. Tool use 块 - for tc in &tool_calls { - // content_block_start - let block_start = serde_json::json!({ - "type": "content_block_start", - "index": block_index, - "content_block": { - "type": "tool_use", - "id": tc.id, - "name": tc.function.name, - "input": {} - } - }); - events.push(format!( - "event: content_block_start\ndata: {block_start}\n\n" - )); - - // content_block_delta - input_json_delta - let partial_json = if tc.function.arguments.is_empty() { - "{}".to_string() - } else { - tc.function.arguments.clone() - }; - let block_delta = serde_json::json!({ - "type": "content_block_delta", - "index": block_index, - "delta": { - "type": "input_json_delta", - "partial_json": partial_json - } - }); - events.push(format!( - "event: content_block_delta\ndata: {block_delta}\n\n" - )); - - // content_block_stop - let block_stop = serde_json::json!({ - "type": "content_block_stop", - "index": block_index - }); - events.push(format!("event: content_block_stop\ndata: {block_stop}\n\n")); - - block_index += 1; - } - - // 4. message_delta - let message_delta = serde_json::json!({ - "type": "message_delta", - "delta": { - "stop_reason": if has_tool_calls { "tool_use" } else { "end_turn" }, - "stop_sequence": null - }, - "usage": {"output_tokens": output_tokens} - }); - events.push(format!("event: message_delta\ndata: {message_delta}\n\n")); - - // 5. message_stop - let message_stop = serde_json::json!({"type": "message_stop"}); - events.push(format!("event: message_stop\ndata: {message_stop}\n\n")); - - // 创建 SSE 响应 - let body_stream = stream::iter(events.into_iter().map(Ok::<_, std::convert::Infallible>)); - let body = Body::from_stream(body_stream); - - Response::builder() - .status(StatusCode::OK) - .header(header::CONTENT_TYPE, "text/event-stream") - .header(header::CACHE_CONTROL, "no-cache") - .header(header::CONNECTION, "keep-alive") - .body(body) - .unwrap_or_else(|e| { - tracing::error!("Failed to build SSE response: {}", e); - Response::builder() - .status(StatusCode::INTERNAL_SERVER_ERROR) - .body(Body::empty()) - .unwrap_or_default() - }) -} - -/// 构建 Gemini CLI OAuth 请求体 -/// -/// 用于 Gemini OAuth 凭证（Cloud Code Assist API） -/// 不做模型名称映射，直接使用用户传入的模型名称 -pub fn build_gemini_cli_request( - request: &serde_json::Value, - model: &str, - project_id: &str, -) -> serde_json::Value { - // 是否启用思维链 - let enable_thinking = model.ends_with("-thinking") - || model == "gemini-2.5-pro" - || model.starts_with("gemini-3-pro-"); - - // 构建内部请求 - let mut inner_request = request.clone(); - - // 确保有 generationConfig - if inner_request.get("generationConfig").is_none() { - inner_request["generationConfig"] = serde_json::json!({ - "temperature": 1.0, - "maxOutputTokens": 8096, - "topP": 0.85, - "topK": 50, - "candidateCount": 1, - "thinkingConfig": { - "includeThoughts": enable_thinking, - "thinkingBudget": if enable_thinking { 1024 } else { 0 } - } - }); - } else { - // 确保有 thinkingConfig - if inner_request["generationConfig"] - .get("thinkingConfig") - .is_none() - { - inner_request["generationConfig"]["thinkingConfig"] = serde_json::json!({ - "includeThoughts": enable_thinking, - "thinkingBudget": if enable_thinking { 1024 } else { 0 } - }); - } - } - - // 删除安全设置（Cloud Code Assist 不支持） - if let Some(obj) = inner_request.as_object_mut() { - obj.remove("safetySettings"); - } - - // 构建完整的 Gemini CLI 请求体 - // 格式与 CLIProxyAPI 对齐 - serde_json::json!({ - "project": project_id, - "model": model, // 直接使用模型名称，不做映射 - "request": inner_request - }) -} - -/// 构建 Gemini 原生请求体 -/// -/// 将用户传入的 Gemini 格式请求转换为 Antigravity 请求格式 -pub fn build_gemini_native_request( - request: &serde_json::Value, - model: &str, - project_id: &str, -) -> serde_json::Value { - // 模型名称映射 - let actual_model = match model { - "gemini-2.5-computer-use-preview-10-2025" => "rev19-uic3-1p", - "gemini-3-pro-image-preview" => "gemini-3-pro-image", - "gemini-3-pro-preview" => "gemini-3-pro-high", - "gemini-claude-sonnet-4-5" => "claude-sonnet-4-5", - "gemini-claude-sonnet-4-5-thinking" => "claude-sonnet-4-5-thinking", - _ => model, - }; - - // 是否启用思维链 - let enable_thinking = model.ends_with("-thinking") - || model == "gemini-2.5-pro" - || model.starts_with("gemini-3-pro-") - || model == "rev19-uic3-1p" - || model == "gpt-oss-120b-medium"; - - // 生成请求 ID 和会话 ID - let request_id = format!("agent-{}", uuid::Uuid::new_v4()); - let session_id = { - let uuid = uuid::Uuid::new_v4(); - let bytes = uuid.as_bytes(); - let n: u64 = u64::from_le_bytes([ - bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7], - ]) % 9_000_000_000_000_000_000; - format!("-{n}") - }; - - // 构建内部请求 - let mut inner_request = request.clone(); - - // 添加会话 ID - inner_request["sessionId"] = serde_json::json!(session_id); - - // 确保有 generationConfig - if inner_request.get("generationConfig").is_none() { - inner_request["generationConfig"] = serde_json::json!({ - "temperature": 1.0, - "maxOutputTokens": 8096, - "topP": 0.85, - "topK": 50, - "candidateCount": 1, - "stopSequences": [ - "<|user|>", - "<|bot|>", - "<|context_request|>", - "<|endoftext|>", - "<|end_of_turn|>" - ], - "thinkingConfig": { - "includeThoughts": enable_thinking, - "thinkingBudget": if enable_thinking { 1024 } else { 0 } - } - }); - } else { - // 确保有 thinkingConfig - if inner_request["generationConfig"] - .get("thinkingConfig") - .is_none() - { - inner_request["generationConfig"]["thinkingConfig"] = serde_json::json!({ - "includeThoughts": enable_thinking, - "thinkingBudget": if enable_thinking { 1024 } else { 0 } - }); - } - } - - // 删除安全设置（Antigravity 不支持） - if let Some(obj) = inner_request.as_object_mut() { - obj.remove("safetySettings"); - } - - // 构建完整的 Antigravity 请求体 - serde_json::json!({ - "project": project_id, - "requestId": request_id, - "request": inner_request, - "model": actual_model, - "userAgent": "antigravity" - }) -} - -/// 健康检查端点响应 -pub async fn health() -> impl IntoResponse { - Json(serde_json::json!({ - "status": "healthy", - "version": env!("CARGO_PKG_VERSION") - })) -} - -/// 模型列表端点响应（静态列表，用于不指定凭证的情况） -pub async fn models() -> impl IntoResponse { - Json(serde_json::json!({ - "object": "list", - "data": [ - // Kiro/Claude models - {"id": "claude-sonnet-4-5", "object": "model", "owned_by": "anthropic"}, - {"id": "claude-sonnet-4-5-20250929", "object": "model", "owned_by": "anthropic"}, - {"id": "claude-3-7-sonnet-20250219", "object": "model", "owned_by": "anthropic"}, - {"id": "claude-3-5-sonnet-latest", "object": "model", "owned_by": "anthropic"}, - // Gemini models - {"id": "gemini-2.5-flash", "object": "model", "owned_by": "google"}, - {"id": "gemini-2.5-flash-lite", "object": "model", "owned_by": "google"}, - {"id": "gemini-2.5-pro", "object": "model", "owned_by": "google"}, - {"id": "gemini-2.5-pro-preview-06-05", "object": "model", "owned_by": "google"}, - {"id": "gemini-3-pro-preview", "object": "model", "owned_by": "google"}, - {"id": "gemini-3-pro-image-preview", "object": "model", "owned_by": "google"}, - {"id": "gemini-3-flash-preview", "object": "model", "owned_by": "google"}, - {"id": "gemini-2.5-computer-use-preview-10-2025", "object": "model", "owned_by": "google"}, - {"id": "gemini-claude-sonnet-4-5", "object": "model", "owned_by": "google"}, - {"id": "gemini-claude-sonnet-4-5-thinking", "object": "model", "owned_by": "google"}, - {"id": "gemini-claude-opus-4-5-thinking", "object": "model", "owned_by": "google"}, - // Qwen models - {"id": "qwen3-coder-plus", "object": "model", "owned_by": "alibaba"}, - {"id": "qwen3-coder-flash", "object": "model", "owned_by": "alibaba"} - ] - })) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_safe_truncate() { - assert_eq!(safe_truncate("hello", 10), "hello"); - assert_eq!(safe_truncate("hello world", 5), "hello"); - assert_eq!(safe_truncate("你好世界", 2), "你好"); - } - - #[test] - fn test_find_subsequence() { - let haystack = b"hello world"; - assert_eq!(find_subsequence(haystack, b"world"), Some(6)); - assert_eq!(find_subsequence(haystack, b"foo"), None); - } - - #[test] - fn test_extract_json_from_bytes() { - let json = b"{\"key\":\"value\"}"; - assert_eq!( - extract_json_from_bytes(json), - Some("{\"key\":\"value\"}".to_string()) - ); - - let nested = b"{\"outer\":{\"inner\":\"value\"}}"; - assert_eq!( - extract_json_from_bytes(nested), - Some("{\"outer\":{\"inner\":\"value\"}}".to_string()) - ); - - assert_eq!(extract_json_from_bytes(b"not json"), None); - } -} - -// ============================================================================ -// 属性测试 -// ============================================================================ - -#[cfg(test)] -mod property_tests { - use super::*; - use proptest::prelude::*; - - // 生成随机文本内容 - fn arb_text_content() -> impl Strategy { - "[a-zA-Z0-9 .,!?\\n]{0,500}".prop_map(|s| s) - } - - // 生成随机模型名称 - fn arb_model_name() -> impl Strategy { - prop_oneof![ - Just("claude-3-sonnet".to_string()), - Just("claude-3-opus".to_string()), - Just("claude-3-haiku".to_string()), - Just("claude-sonnet-4-5".to_string()), - Just("claude-3-5-sonnet-latest".to_string()), - ] - } - - // 生成随机工具调用 - fn arb_tool_call() -> impl Strategy { - ( - "[a-z0-9]{8,16}", - prop_oneof![ - Just("read_file".to_string()), - Just("write_file".to_string()), - Just("execute_command".to_string()), - Just("search".to_string()), - ], - prop_oneof![ - Just("{}".to_string()), - Just("{\"path\":\"/tmp/test\"}".to_string()), - Just("{\"content\":\"hello\"}".to_string()), - Just("{\"query\":\"test\",\"limit\":10}".to_string()), - ], - ) - .prop_map(|(id, name, args)| ToolCall { - id: format!("call_{id}"), - call_type: "function".to_string(), - function: FunctionCall { - name, - arguments: args, - }, - }) - } - - // 生成随机 CWParsedResponse - fn arb_cw_parsed_response() -> impl Strategy { - ( - arb_text_content(), - prop::collection::vec(arb_tool_call(), 0..3), - 0.0f64..100.0f64, - 0.0f64..100.0f64, - ) - .prop_map( - |(content, tool_calls, usage_credits, context_usage_percentage)| CWParsedResponse { - content, - tool_calls, - usage_credits, - context_usage_percentage, - }, - ) - } - - // ======================================================================== - // Property 9: 非流式响应格式 - // **Validates: Requirements 6.2** - // ======================================================================== - - proptest! { - /// Property 9: 非流式响应格式 - /// - /// *对于任意* 非流式请求，响应应该是完整的 JSON 对象，包含所有必需字段。 - /// - /// 必需字段: - /// - id: 消息 ID (格式: msg_xxx) - /// - type: 消息类型 (固定为 "message") - /// - role: 角色 (固定为 "assistant") - /// - content: 内容数组 - /// - model: 模型名称 - /// - stop_reason: 停止原因 ("end_turn" 或 "tool_use") - /// - stop_sequence: 停止序列 (null) - /// - usage: 使用量信息 (包含 input_tokens 和 output_tokens) - /// - /// **Validates: Requirements 6.2** - #[test] - fn prop_non_streaming_response_format( - model in arb_model_name(), - parsed in arb_cw_parsed_response() - ) { - // 构建非流式响应 - let response = build_anthropic_response(&model, &parsed); - - // 获取响应体 - let (parts, _body) = response.into_parts(); - - // 验证状态码为 200 - prop_assert_eq!(parts.status, StatusCode::OK); - - // 验证 Content-Type 为 application/json - let content_type = parts.headers.get(header::CONTENT_TYPE); - prop_assert!(content_type.is_some()); - prop_assert!(content_type.unwrap().to_str().unwrap().contains("application/json")); - - // 由于 Body 是异步的，我们需要在同步测试中使用 futures::executor - // 但是 proptest 不支持异步，所以我们直接测试 JSON 构建逻辑 - // 这里我们重新构建 JSON 来验证格式 - - let has_tool_calls = !parsed.tool_calls.is_empty(); - let mut content_array: Vec = Vec::new(); - - if !parsed.content.is_empty() { - content_array.push(serde_json::json!({ - "type": "text", - "text": parsed.content - })); - } - - for tc in &parsed.tool_calls { - let input: serde_json::Value = - serde_json::from_str(&tc.function.arguments).unwrap_or(serde_json::json!({})); - content_array.push(serde_json::json!({ - "type": "tool_use", - "id": tc.id, - "name": tc.function.name, - "input": input - })); - } - - if content_array.is_empty() { - content_array.push(serde_json::json!({"type": "text", "text": ""})); - } - - // 估算 tokens - let mut output_tokens: u32 = (parsed.content.len() / 4) as u32; - for tc in &parsed.tool_calls { - output_tokens += (tc.function.arguments.len() / 4) as u32; - } - let input_tokens = ((parsed.context_usage_percentage / 100.0) * 200000.0) as u32; - - // 构建预期的 JSON 响应 - let expected_json = serde_json::json!({ - "type": "message", - "role": "assistant", - "content": content_array, - "model": model, - "stop_reason": if has_tool_calls { "tool_use" } else { "end_turn" }, - "stop_sequence": null, - "usage": { - "input_tokens": input_tokens, - "output_tokens": output_tokens - } - }); - - // 验证必需字段存在 - prop_assert!(expected_json.get("type").is_some()); - prop_assert_eq!(expected_json["type"].as_str(), Some("message")); - - prop_assert!(expected_json.get("role").is_some()); - prop_assert_eq!(expected_json["role"].as_str(), Some("assistant")); - - prop_assert!(expected_json.get("content").is_some()); - prop_assert!(expected_json["content"].is_array()); - prop_assert!(!expected_json["content"].as_array().unwrap().is_empty()); - - prop_assert!(expected_json.get("model").is_some()); - prop_assert_eq!(expected_json["model"].as_str(), Some(model.as_str())); - - prop_assert!(expected_json.get("stop_reason").is_some()); - let stop_reason = expected_json["stop_reason"].as_str().unwrap(); - prop_assert!(stop_reason == "end_turn" || stop_reason == "tool_use"); - - // stop_reason 应该与 tool_calls 状态一致 - if has_tool_calls { - prop_assert_eq!(stop_reason, "tool_use"); - } else { - prop_assert_eq!(stop_reason, "end_turn"); - } - - prop_assert!(expected_json.get("stop_sequence").is_some()); - prop_assert!(expected_json["stop_sequence"].is_null()); - - prop_assert!(expected_json.get("usage").is_some()); - prop_assert!(expected_json["usage"].get("input_tokens").is_some()); - prop_assert!(expected_json["usage"].get("output_tokens").is_some()); - - // 验证 content 数组中的每个元素都有正确的类型 - for item in expected_json["content"].as_array().unwrap() { - prop_assert!(item.get("type").is_some()); - let item_type = item["type"].as_str().unwrap(); - prop_assert!(item_type == "text" || item_type == "tool_use"); - - if item_type == "text" { - prop_assert!(item.get("text").is_some()); - } else if item_type == "tool_use" { - prop_assert!(item.get("id").is_some()); - prop_assert!(item.get("name").is_some()); - prop_assert!(item.get("input").is_some()); - } - } - } - - /// 验证空内容时响应仍然有效 - #[test] - fn prop_non_streaming_response_empty_content( - model in arb_model_name() - ) { - let parsed = CWParsedResponse { - content: String::new(), - tool_calls: Vec::new(), - usage_credits: 0.0, - context_usage_percentage: 0.0, - }; - - let response = build_anthropic_response(&model, &parsed); - let (parts, _body) = response.into_parts(); - - // 验证状态码为 200 - prop_assert_eq!(parts.status, StatusCode::OK); - - // 即使内容为空，content 数组也应该有一个空文本元素 - let content_array = [serde_json::json!({"type": "text", "text": ""})]; - prop_assert!(!content_array.is_empty()); - } - - /// 验证只有工具调用时响应格式正确 - #[test] - fn prop_non_streaming_response_tool_calls_only( - model in arb_model_name(), - tool_calls in prop::collection::vec(arb_tool_call(), 1..3) - ) { - let parsed = CWParsedResponse { - content: String::new(), - tool_calls, - usage_credits: 0.0, - context_usage_percentage: 50.0, - }; - - let response = build_anthropic_response(&model, &parsed); - let (parts, _body) = response.into_parts(); - - // 验证状态码为 200 - prop_assert_eq!(parts.status, StatusCode::OK); - - // 有工具调用时，stop_reason 应该是 "tool_use" - // 这里我们验证逻辑正确性 - prop_assert!(!parsed.tool_calls.is_empty()); - } - - /// 验证 Token 估算逻辑 - #[test] - fn prop_token_estimation( - content in arb_text_content(), - context_percentage in 0.0f64..100.0f64 - ) { - let parsed = CWParsedResponse { - content: content.clone(), - tool_calls: Vec::new(), - usage_credits: 0.0, - context_usage_percentage: context_percentage, - }; - - let (input_tokens, output_tokens) = parsed.estimate_tokens(); - - // output_tokens 应该约等于 content 长度 / 4 - let expected_output = (content.len() / 4) as u32; - prop_assert_eq!(output_tokens, expected_output); - - // input_tokens 应该基于 context_usage_percentage - let expected_input = ((context_percentage / 100.0) * 200000.0) as u32; - prop_assert_eq!(input_tokens, expected_input); - } - } - - // ======================================================================== - // Property 1: 模型列表结构和归属正确性 - // **Validates: Requirements 1.2, 1.3** - // ======================================================================== - - /// 获取模型列表数据用于测试 - fn get_model_list_data() -> Vec { - vec![ - // Kiro/Claude models - serde_json::json!({"id": "claude-sonnet-4-5", "object": "model", "owned_by": "anthropic"}), - serde_json::json!({"id": "claude-sonnet-4-5-20250929", "object": "model", "owned_by": "anthropic"}), - serde_json::json!({"id": "claude-3-7-sonnet-20250219", "object": "model", "owned_by": "anthropic"}), - serde_json::json!({"id": "claude-3-5-sonnet-latest", "object": "model", "owned_by": "anthropic"}), - // Gemini models - serde_json::json!({"id": "gemini-2.5-flash", "object": "model", "owned_by": "google"}), - serde_json::json!({"id": "gemini-2.5-flash-lite", "object": "model", "owned_by": "google"}), - serde_json::json!({"id": "gemini-2.5-pro", "object": "model", "owned_by": "google"}), - serde_json::json!({"id": "gemini-2.5-pro-preview-06-05", "object": "model", "owned_by": "google"}), - serde_json::json!({"id": "gemini-3-pro-preview", "object": "model", "owned_by": "google"}), - serde_json::json!({"id": "gemini-3-pro-image-preview", "object": "model", "owned_by": "google"}), - serde_json::json!({"id": "gemini-3-flash-preview", "object": "model", "owned_by": "google"}), - serde_json::json!({"id": "gemini-2.5-computer-use-preview-10-2025", "object": "model", "owned_by": "google"}), - serde_json::json!({"id": "gemini-claude-sonnet-4-5", "object": "model", "owned_by": "google"}), - serde_json::json!({"id": "gemini-claude-sonnet-4-5-thinking", "object": "model", "owned_by": "google"}), - serde_json::json!({"id": "gemini-claude-opus-4-5-thinking", "object": "model", "owned_by": "google"}), - // Qwen models - serde_json::json!({"id": "qwen3-coder-plus", "object": "model", "owned_by": "alibaba"}), - serde_json::json!({"id": "qwen3-coder-flash", "object": "model", "owned_by": "alibaba"}), - ] - } - - /// Property 1: 模型列表结构和归属正确性 - /// - /// *对于任意* 模型列表中的模型，应该包含: - /// - id 字段 (非空字符串) - /// - object 字段 (值为 "model") - /// - owned_by 字段 (与模型类型匹配: gemini-* -> google, claude-* -> anthropic, qwen* -> alibaba) - /// - /// **Validates: Requirements 1.2, 1.3** - #[test] - fn prop_model_list_structure_and_ownership() { - let models = get_model_list_data(); - - for model in &models { - // 验证 id 字段存在且非空 - let id = model.get("id").and_then(|v| v.as_str()); - assert!(id.is_some(), "Model should have id field"); - assert!(!id.unwrap().is_empty(), "Model id should not be empty"); - - // 验证 object 字段为 "model" - let object = model.get("object").and_then(|v| v.as_str()); - assert_eq!(object, Some("model"), "Model object should be 'model'"); - - // 验证 owned_by 字段与模型类型匹配 - let owned_by = model.get("owned_by").and_then(|v| v.as_str()); - assert!(owned_by.is_some(), "Model should have owned_by field"); - - let model_id = id.unwrap(); - let owner = owned_by.unwrap(); - - if model_id.starts_with("gemini-") { - assert_eq!(owner, "google", "Gemini models should be owned by google"); - } else if model_id.starts_with("claude-") { - assert_eq!( - owner, "anthropic", - "Claude models should be owned by anthropic" - ); - } else if model_id.starts_with("qwen") { - assert_eq!(owner, "alibaba", "Qwen models should be owned by alibaba"); - } - } - } - - /// 验证所有 Antigravity 支持的模型都在列表中 - #[test] - fn test_antigravity_models_present() { - let models = get_model_list_data(); - let model_ids: Vec<&str> = models - .iter() - .filter_map(|m| m.get("id").and_then(|v| v.as_str())) - .collect(); - - // 验证所有 Antigravity 支持的模型都存在 - let required_models = [ - "gemini-3-pro-preview", - "gemini-3-pro-image-preview", - "gemini-3-flash-preview", - "gemini-2.5-computer-use-preview-10-2025", - "gemini-claude-sonnet-4-5", - "gemini-claude-sonnet-4-5-thinking", - "gemini-claude-opus-4-5-thinking", - ]; - - for required in &required_models { - assert!( - model_ids.contains(required), - "Model {required} should be in the list" - ); - } - } - - // ======================================================================== - // Property 2: 模型名称映射正确性 - // **Validates: Requirements 3.1, 3.2** - // ======================================================================== - - /// 获取模型名称映射的预期结果 - #[allow(dead_code)] - fn get_expected_model_mapping(model: &str) -> &str { - match model { - "gemini-2.5-computer-use-preview-10-2025" => "rev19-uic3-1p", - "gemini-3-pro-image-preview" => "gemini-3-pro-image", - "gemini-3-pro-preview" => "gemini-3-pro-high", - "gemini-claude-sonnet-4-5" => "claude-sonnet-4-5", - "gemini-claude-sonnet-4-5-thinking" => "claude-sonnet-4-5-thinking", - _ => model, - } - } - - /// Property 2: 模型名称映射正确性 - /// - /// *对于任意* 已知映射表中的模型名称，build_gemini_native_request 应该返回正确的内部模型名称。 - /// *对于任意* 不在映射表中的模型名称，应该原样返回。 - /// - /// **Validates: Requirements 3.1, 3.2** - #[test] - fn prop_model_name_mapping_correctness() { - let test_request = serde_json::json!({ - "contents": [{"role": "user", "parts": [{"text": "test"}]}] - }); - let project_id = "test-project"; - - // 测试已知映射 - let known_mappings = [ - ("gemini-2.5-computer-use-preview-10-2025", "rev19-uic3-1p"), - ("gemini-3-pro-image-preview", "gemini-3-pro-image"), - ("gemini-3-pro-preview", "gemini-3-pro-high"), - ("gemini-claude-sonnet-4-5", "claude-sonnet-4-5"), - ( - "gemini-claude-sonnet-4-5-thinking", - "claude-sonnet-4-5-thinking", - ), - ]; - - for (input, expected) in &known_mappings { - let result = build_gemini_native_request(&test_request, input, project_id); - let actual_model = result.get("model").and_then(|v| v.as_str()).unwrap(); - assert_eq!( - actual_model, *expected, - "Model {input} should map to {expected}" - ); - } - - // 测试未知模型名称应该原样返回 - let unknown_models = ["gemini-2.0-flash", "gemini-2.5-flash", "custom-model"]; - for model in &unknown_models { - let result = build_gemini_native_request(&test_request, model, project_id); - let actual_model = result.get("model").and_then(|v| v.as_str()).unwrap(); - assert_eq!( - actual_model, *model, - "Unknown model {model} should be returned unchanged" - ); - } - } - - // ======================================================================== - // Property 3: 思维链启用逻辑正确性 - // **Validates: Requirements 4.1, 4.2** - // ======================================================================== - - /// 判断模型是否应该启用思维链 - fn should_enable_thinking(model: &str) -> bool { - model.ends_with("-thinking") - || model == "gemini-2.5-pro" - || model.starts_with("gemini-3-pro-") - || model == "rev19-uic3-1p" - || model == "gpt-oss-120b-medium" - } - - /// Property 3: 思维链启用逻辑正确性 - /// - /// *对于任意* 模型名称，思维链启用状态应该根据以下规则正确判断: - /// - 以 "-thinking" 结尾的模型启用思维链 - /// - "gemini-2.5-pro" 启用思维链 - /// - 以 "gemini-3-pro-" 开头的模型启用思维链 - /// - "rev19-uic3-1p" 启用思维链 - /// - "gpt-oss-120b-medium" 启用思维链 - /// - /// **Validates: Requirements 4.1, 4.2** - #[test] - fn prop_thinking_mode_enablement_logic() { - // 应该启用思维链的模型 - let thinking_enabled_models = [ - "gemini-claude-sonnet-4-5-thinking", - "gemini-claude-opus-4-5-thinking", - "custom-model-thinking", - "gemini-2.5-pro", - "gemini-3-pro-preview", - "gemini-3-pro-image-preview", - "gemini-3-pro-high", - "rev19-uic3-1p", - "gpt-oss-120b-medium", - ]; - - for model in &thinking_enabled_models { - assert!( - should_enable_thinking(model), - "Model {model} should have thinking enabled" - ); - } - - // 不应该启用思维链的模型 - let thinking_disabled_models = [ - "gemini-2.0-flash", - "gemini-2.5-flash", - "gemini-claude-sonnet-4-5", - "claude-sonnet-4-5", - "custom-model", - ]; - - for model in &thinking_disabled_models { - assert!( - !should_enable_thinking(model), - "Model {model} should have thinking disabled" - ); - } - } - - // ======================================================================== - // Property 4: 思维链配置值正确性 - // **Validates: Requirements 4.3, 4.4** - // ======================================================================== - - /// Property 4: 思维链配置值正确性 - /// - /// *对于任意* Gemini 原生请求: - /// - 当思维链启用时，includeThoughts 应该为 true，thinkingBudget 应该为 1024 - /// - 当思维链禁用时，includeThoughts 应该为 false，thinkingBudget 应该为 0 - /// - /// **Validates: Requirements 4.3, 4.4** - #[test] - fn prop_thinking_configuration_values() { - let test_request = serde_json::json!({ - "contents": [{"role": "user", "parts": [{"text": "test"}]}] - }); - let project_id = "test-project"; - - // 测试启用思维链的模型 - let thinking_enabled_models = [ - "gemini-3-pro-preview", - "gemini-2.5-pro", - "gemini-claude-sonnet-4-5-thinking", - ]; - - for model in &thinking_enabled_models { - let result = build_gemini_native_request(&test_request, model, project_id); - let thinking_config = &result["request"]["generationConfig"]["thinkingConfig"]; - - assert_eq!( - thinking_config["includeThoughts"].as_bool(), - Some(true), - "Model {model} should have includeThoughts=true" - ); - assert_eq!( - thinking_config["thinkingBudget"].as_i64(), - Some(1024), - "Model {model} should have thinkingBudget=1024" - ); - } - - // 测试禁用思维链的模型 - let thinking_disabled_models = [ - "gemini-2.0-flash", - "gemini-2.5-flash", - "gemini-claude-sonnet-4-5", - ]; - - for model in &thinking_disabled_models { - let result = build_gemini_native_request(&test_request, model, project_id); - let thinking_config = &result["request"]["generationConfig"]["thinkingConfig"]; - - assert_eq!( - thinking_config["includeThoughts"].as_bool(), - Some(false), - "Model {model} should have includeThoughts=false" - ); - assert_eq!( - thinking_config["thinkingBudget"].as_i64(), - Some(0), - "Model {model} should have thinkingBudget=0" - ); - } - } -} +pub use proxycast_server_utils::*;