diff --git a/CLAUDE.md b/CLAUDE.md new file mode 120000 index 0000000..47dc3e3 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +AGENTS.md \ No newline at end of file diff --git a/antigravity-bridge/src/auth.rs b/antigravity-bridge/src/auth.rs index 679d964..5a0e120 100644 --- a/antigravity-bridge/src/auth.rs +++ b/antigravity-bridge/src/auth.rs @@ -4,6 +4,7 @@ use std::time::{SystemTime, UNIX_EPOCH}; pub const BINARY_PATH: &str = "/home/mayer/.local/bin/agy"; +#[allow(dead_code)] #[derive(Clone)] pub struct CachedSessionHashes { pub trajectory_id: String, diff --git a/antigravity-bridge/src/main.rs b/antigravity-bridge/src/main.rs index 4ed9cc4..45c9459 100644 --- a/antigravity-bridge/src/main.rs +++ b/antigravity-bridge/src/main.rs @@ -14,7 +14,7 @@ use auth::{ extract_credentials_candidates, get_binary_sha256, load_expected_sha256, load_refresh_token, verify_and_bind_credentials, BridgeState, BINARY_PATH, }; -use routes::{chat_completions, list_models}; +use routes::{chat_completions, list_models, anthropic_messages}; #[tokio::main] async fn main() -> anyhow::Result<()> { @@ -61,6 +61,8 @@ async fn main() -> anyhow::Result<()> { let app = Router::new() .route("/v1/chat/completions", post(chat_completions)) + .route("/v1/messages", post(anthropic_messages)) + .route("/messages", post(anthropic_messages)) .route("/v1/models", get(list_models)) .route("/models", get(list_models)) .layer(DefaultBodyLimit::max(15 * 1024 * 1024)) diff --git a/antigravity-bridge/src/mapping.rs b/antigravity-bridge/src/mapping.rs index fc91b5f..1a9807b 100644 --- a/antigravity-bridge/src/mapping.rs +++ b/antigravity-bridge/src/mapping.rs @@ -1,25 +1,64 @@ use serde_json::Value; use std::collections::HashMap; -pub fn convert_schema_types(mut schema: Value) -> Value { - if let Some(obj) = schema.as_object_mut() { - if let Some(t_val) = obj.get_mut("type") { - if let Some(t_str) = t_val.as_str() { - *t_val = serde_json::json!(t_str.to_uppercase()); - } +pub fn convert_schema_types(schema: Value) -> Value { + let Some(obj) = schema.as_object() else { + return schema; + }; + + let mut cleaned = serde_json::Map::new(); + + // 1. Map type and convert to uppercase if string + if let Some(t_val) = obj.get("type") { + if let Some(t_str) = t_val.as_str() { + cleaned.insert("type".to_string(), serde_json::json!(t_str.to_uppercase())); + } else { + cleaned.insert("type".to_string(), t_val.clone()); } - if let Some(props) = obj.get_mut("properties") { - if let Some(props_obj) = props.as_object_mut() { - for (_, val) in props_obj.iter_mut() { - *val = convert_schema_types(val.clone()); - } + } + + // 2. Keep format + if let Some(format_val) = obj.get("format") { + cleaned.insert("format".to_string(), format_val.clone()); + } + + // 3. Keep description + if let Some(desc_val) = obj.get("description") { + cleaned.insert("description".to_string(), desc_val.clone()); + } + + // 4. Keep nullable + if let Some(null_val) = obj.get("nullable") { + cleaned.insert("nullable".to_string(), null_val.clone()); + } + + // 5. Keep enum + if let Some(enum_val) = obj.get("enum") { + cleaned.insert("enum".to_string(), enum_val.clone()); + } + + // 6. Keep required + if let Some(req_val) = obj.get("required") { + cleaned.insert("required".to_string(), req_val.clone()); + } + + // 7. Recurse properties + if let Some(props_val) = obj.get("properties") { + if let Some(props_obj) = props_val.as_object() { + let mut cleaned_props = serde_json::Map::new(); + for (k, v) in props_obj { + cleaned_props.insert(k.clone(), convert_schema_types(v.clone())); } + cleaned.insert("properties".to_string(), Value::Object(cleaned_props)); } - if let Some(items) = obj.get_mut("items") { - *items = convert_schema_types(items.clone()); - } } - schema + + // 8. Recurse items + if let Some(items_val) = obj.get("items") { + cleaned.insert("items".to_string(), convert_schema_types(items_val.clone())); + } + + Value::Object(cleaned) } pub fn map_openai_tools_to_gemini(openai_tools: &Value) -> Option { @@ -64,16 +103,16 @@ pub fn map_model_name(model: &str, thinking_level: Option<&str>) -> String { ("gemini-3.5-flash-low", "gemini-3.5-flash-extra-low"), ("gemini-3.1-pro-low", "gemini-3.1-pro-low"), ("gemini-3.1-pro-high", "gemini-pro-agent"), - ("claude-sonnet-4-6", "claude-sonnet-4-6"), - ("claude-opus-4-6-thinking", "claude-opus-4-6-thinking"), + ("claude-sonnet-4-6", "gemini-3.5-flash-low"), + ("claude-opus-4-6-thinking", "gemini-3-flash-agent"), ("gpt-oss-120b-medium", "gpt-oss-120b-medium"), ("gemini 3.5 flash (medium)", "gemini-3.5-flash-low"), ("gemini 3.5 flash (high)", "gemini-3-flash-agent"), ("gemini 3.5 flash (low)", "gemini-3.5-flash-extra-low"), ("gemini 3.1 pro (low)", "gemini-3.1-pro-low"), ("gemini 3.1 pro (high)", "gemini-pro-agent"), - ("claude sonnet 4.6 (thinking)", "claude-sonnet-4-6"), - ("claude opus 4.6 (thinking)", "claude-opus-4-6-thinking"), + ("claude sonnet 4.6 (thinking)", "gemini-3.5-flash-low"), + ("claude opus 4.6 (thinking)", "gemini-3-flash-agent"), ("gpt-oss 120b (medium)", "gpt-oss-120b-medium"), ]; @@ -126,9 +165,11 @@ pub fn map_model_name(model: &str, thinking_level: Option<&str>) -> String { "gemini-3.1-pro-low".to_string() } } else if model_lower.contains("sonnet") { - "claude-sonnet-4-6".to_string() + "gemini-3.5-flash-low".to_string() + } else if model_lower.contains("haiku") { + "gemini-3.5-flash-extra-low".to_string() } else if model_lower.contains("opus") { - "claude-opus-4-6-thinking".to_string() + "gemini-3-flash-agent".to_string() } else if model_lower.contains("gpt-oss") { "gpt-oss-120b-medium".to_string() } else { @@ -536,6 +577,259 @@ fn parse_data_url(url: &str) -> Option<(String, String)> { Some((mime.to_string(), data.to_string())) } +pub fn map_anthropic_messages_to_gemini( + messages_val: &Value, + signature_cache: &HashMap, +) -> (Vec, Option) { + let mut contents = Vec::new(); + let mut tool_name_by_id = HashMap::new(); + + // First scan all messages to extract tool names from tool_use blocks + if let Some(arr) = messages_val.as_array() { + for msg in arr { + let content_val = &msg["content"]; + if let Some(content_arr) = content_val.as_array() { + for block in content_arr { + if block["type"].as_str() == Some("tool_use") { + if let (Some(id), Some(name)) = (block["id"].as_str(), block["name"].as_str()) { + tool_name_by_id.insert(id.to_string(), name.to_string()); + } + } + } + } + } + } + + if let Some(arr) = messages_val.as_array() { + for msg in arr { + let role = msg["role"].as_str().unwrap_or("user"); + let gemini_role = match role { + "assistant" => "model", + _ => "user", + }; + + let mut parts = Vec::new(); + let content_val = &msg["content"]; + + if content_val.is_string() { + let content_str = content_val.as_str().unwrap_or(""); + if !content_str.is_empty() { + parse_text_with_media(content_str, &mut parts); + } + } else if let Some(content_arr) = content_val.as_array() { + for block in content_arr { + let block_type = block["type"].as_str().unwrap_or(""); + match block_type { + "text" => { + if let Some(text) = block["text"].as_str() { + if !text.is_empty() { + parts.push(serde_json::json!({ "text": text })); + } + } + } + "image" | "document" => { + if let Some(source) = block["source"].as_object() { + if source.get("type").and_then(|t| t.as_str()) == Some("base64") { + if let (Some(media_type), Some(data)) = ( + source.get("media_type").and_then(|t| t.as_str()), + source.get("data").and_then(|d| d.as_str()), + ) { + let cleaned_data: String = data + .chars() + .filter(|c| !c.is_whitespace()) + .collect(); + use base64::{engine::general_purpose, Engine as _}; + let final_mime = if let Ok(decoded) = + general_purpose::STANDARD.decode(cleaned_data.as_bytes()) + { + sniff_or_validate_mime(media_type, &decoded) + .unwrap_or_else(|| media_type.to_string()) + } else { + media_type.to_string() + }; + parts.push(serde_json::json!({ + "inlineData": { + "mimeType": final_mime, + "data": cleaned_data + } + })); + } + } + } + } + "tool_use" => { + let func_name = block["name"].as_str().unwrap_or(""); + let args_val = block["input"].clone(); + let tc_id = block["id"].as_str().unwrap_or(""); + let signature = signature_cache.get(tc_id).map(|s| s.as_str()).unwrap_or(""); + + parts.push(serde_json::json!({ + "functionCall": { + "name": func_name, + "args": args_val, + "id": tc_id + }, + "thoughtSignature": signature + })); + } + "tool_result" => { + let tc_id = block["tool_use_id"].as_str().unwrap_or(""); + let fn_name = tool_name_by_id + .get(tc_id) + .cloned() + .unwrap_or_else(|| "unknown_tool".to_string()); + + let block_content = &block["content"]; + let (response_obj, nested_parts) = if block_content.is_string() { + let content_str = block_content.as_str().unwrap_or(""); + let response_val = match serde_json::from_str::(content_str) { + Ok(val) => { + if val.is_object() { + val + } else { + serde_json::json!({ "result": val }) + } + } + Err(_) => { + serde_json::json!({ "result": content_str }) + } + }; + (response_val, vec![]) + } else if block_content.is_object() { + (block_content.clone(), vec![]) + } else if let Some(nested_arr) = block_content.as_array() { + let mut text_accum = String::new(); + let mut media_parts = Vec::new(); + for nested_block in nested_arr { + if nested_block["type"].as_str() == Some("text") { + if let Some(txt) = nested_block["text"].as_str() { + text_accum.push_str(txt); + } + } else if nested_block["type"].as_str() == Some("image") || nested_block["type"].as_str() == Some("document") { + if let Some(source) = nested_block["source"].as_object() { + if source.get("type").and_then(|t| t.as_str()) == Some("base64") { + if let (Some(media_type), Some(data)) = ( + source.get("media_type").and_then(|t| t.as_str()), + source.get("data").and_then(|d| d.as_str()), + ) { + let cleaned_data: String = data + .chars() + .filter(|c| !c.is_whitespace()) + .collect(); + use base64::{engine::general_purpose, Engine as _}; + let final_mime = if let Ok(decoded) = + general_purpose::STANDARD.decode(cleaned_data.as_bytes()) + { + sniff_or_validate_mime(media_type, &decoded) + .unwrap_or_else(|| media_type.to_string()) + } else { + media_type.to_string() + }; + media_parts.push(serde_json::json!({ + "inlineData": { + "mimeType": final_mime, + "data": cleaned_data + } + })); + } + } + } + } + } + let response_val = serde_json::json!({ "result": text_accum }); + (response_val, media_parts) + } else { + (serde_json::json!({ "result": block_content.to_string() }), vec![]) + }; + + parts.push(serde_json::json!({ + "functionResponse": { + "name": fn_name, + "response": response_obj, + "id": tc_id + } + })); + for p in nested_parts { + parts.push(p); + } + } + _ => {} + } + } + } + + if !parts.is_empty() { + contents.push(serde_json::json!({ + "role": gemini_role, + "parts": parts + })); + } + } + } + + (contents, None) +} + +pub fn map_anthropic_tools_to_gemini(anthropic_tools: &Value) -> Option { + let Some(tools_arr) = anthropic_tools.as_array() else { + return None; + }; + + let mut declarations = Vec::new(); + for tool in tools_arr { + let mut decl = serde_json::json!({ + "name": tool["name"].as_str().unwrap_or(""), + "description": tool["description"].as_str().unwrap_or(""), + }); + + if let Some(schema) = tool.get("input_schema") { + let mapped_params = convert_schema_types(schema.clone()); + decl.as_object_mut() + .unwrap() + .insert("parameters".to_string(), mapped_params); + } + declarations.push(decl); + } + + if !declarations.is_empty() { + Some(serde_json::json!([{ + "functionDeclarations": declarations + }])) + } else { + None + } +} + +pub fn map_anthropic_system_instruction(system: &Value) -> Option { + if system.is_string() { + let text = system.as_str().unwrap_or(""); + if !text.is_empty() { + return Some(serde_json::json!({ + "role": "user", + "parts": [{"text": text}] + })); + } + } else if let Some(arr) = system.as_array() { + let mut parts = Vec::new(); + for item in arr { + if item["type"].as_str() == Some("text") { + if let Some(text) = item["text"].as_str() { + if !text.is_empty() { + parts.push(serde_json::json!({ "text": text })); + } + } + } + } + if !parts.is_empty() { + return Some(serde_json::json!({ + "role": "user", + "parts": parts + })); + } + } + None +} + #[cfg(test)] mod tests { use super::*; @@ -649,4 +943,124 @@ mod tests { assert_eq!(parts[1]["inlineData"]["mimeType"], "image/png"); assert_eq!(parts[1]["inlineData"]["data"], "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFWM0AAAADUlEQVR4nGNgAAAAAgABSK+kcQAAAABJRU5ErkJggg=="); } + + #[test] + fn test_map_anthropic_tools() { + let anthropic_tools = serde_json::json!([ + { + "name": "get_weather", + "description": "Get weather info", + "input_schema": { + "$schema": "http://json-schema.org/draft-07/schema#", + "type": "object", + "properties": { + "location": { + "type": "string", + "const": "San Francisco", + "exclusiveMinimum": 10 + } + }, + "required": ["location"], + "additionalProperties": false + } + } + ]); + let mapped = map_anthropic_tools_to_gemini(&anthropic_tools).unwrap(); + let tool_decl = &mapped[0]["functionDeclarations"][0]; + assert_eq!(tool_decl["name"], "get_weather"); + assert_eq!(tool_decl["description"], "Get weather info"); + + let params = &tool_decl["parameters"]; + assert_eq!(params["type"], "OBJECT"); + assert!(params.get("$schema").is_none()); + assert!(params.get("additionalProperties").is_none()); + + let location = ¶ms["properties"]["location"]; + assert_eq!(location["type"], "STRING"); + assert!(location.get("const").is_none()); + assert!(location.get("exclusiveMinimum").is_none()); + } + + #[test] + fn test_map_anthropic_messages() { + let messages = serde_json::json!([ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Check this image" + }, + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFWM0AAAADUlEQVR4nGNgAAAAAgABSK+kcQAAAABJRU5ErkJggg==" + } + } + ] + }, + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Thinking..." + }, + { + "type": "tool_use", + "id": "toolu_1", + "name": "get_weather", + "input": {"location": "San Francisco"} + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_1", + "content": "Sunny" + } + ] + } + ]); + + let (contents, _) = map_anthropic_messages_to_gemini(&messages, &HashMap::new()); + assert_eq!(contents.len(), 3); + + assert_eq!(contents[0]["role"], "user"); + let parts0 = contents[0]["parts"].as_array().unwrap(); + assert_eq!(parts0[0]["text"], "Check this image"); + assert_eq!(parts0[1]["inlineData"]["mimeType"], "image/png"); + + assert_eq!(contents[1]["role"], "model"); + let parts1 = contents[1]["parts"].as_array().unwrap(); + assert_eq!(parts1[0]["text"], "Thinking..."); + assert_eq!(parts1[1]["functionCall"]["name"], "get_weather"); + assert_eq!(parts1[1]["functionCall"]["args"]["location"], "San Francisco"); + + assert_eq!(contents[2]["role"], "user"); + let parts2 = contents[2]["parts"].as_array().unwrap(); + assert_eq!(parts2[0]["functionResponse"]["name"], "get_weather"); + assert_eq!(parts2[0]["functionResponse"]["response"]["result"], "Sunny"); + } + + #[test] + fn test_map_anthropic_system_instruction() { + let system_str = serde_json::json!("You are helpful"); + let mapped_str = map_anthropic_system_instruction(&system_str).unwrap(); + assert_eq!(mapped_str["role"], "user"); + assert_eq!(mapped_str["parts"][0]["text"], "You are helpful"); + + let system_arr = serde_json::json!([ + {"type": "text", "text": "Be concise"}, + {"type": "other", "text": "ignored"} + ]); + let mapped_arr = map_anthropic_system_instruction(&system_arr).unwrap(); + assert_eq!(mapped_arr["parts"].as_array().unwrap().len(), 1); + assert_eq!(mapped_arr["parts"][0]["text"], "Be concise"); + } } diff --git a/antigravity-bridge/src/routes.rs b/antigravity-bridge/src/routes.rs index 0c5ebb2..8d05487 100644 --- a/antigravity-bridge/src/routes.rs +++ b/antigravity-bridge/src/routes.rs @@ -13,7 +13,11 @@ use serde_json::Value; use std::time::{SystemTime, UNIX_EPOCH}; use crate::auth::{get_valid_token_and_project, SharedState}; -use crate::mapping::{map_model_name, map_openai_messages_to_gemini, map_openai_tools_to_gemini}; +use crate::mapping::{ + map_model_name, map_openai_messages_to_gemini, map_openai_tools_to_gemini, + map_anthropic_messages_to_gemini, map_anthropic_tools_to_gemini, + map_anthropic_system_instruction, +}; use crate::telemetry::send_telemetry_metrics; #[derive(Debug, Deserialize)] @@ -24,6 +28,22 @@ pub struct ChatCompletionsRequest { pub tools: Option, } +#[allow(dead_code)] +#[derive(Debug, Deserialize)] +pub struct AnthropicMessagesRequest { + pub model: String, + pub messages: Vec, + pub system: Option, + pub max_tokens: Option, + pub metadata: Option, + pub stop_sequences: Option>, + pub stream: Option, + pub temperature: Option, + pub tool_choice: Option, + pub tools: Option, + pub thinking: Option, +} + pub struct LineDecoder { buffer: Vec, } @@ -531,3 +551,664 @@ pub async fn list_models() -> impl IntoResponse { ] })) } + +pub async fn anthropic_messages( + State(state_arc): State, + Json(body): Json, +) -> impl IntoResponse { + let start_bridge = std::time::Instant::now(); + let (token, project) = match get_valid_token_and_project(&state_arc).await { + Ok(vals) => vals, + Err(e) => { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ "error": format!("Auth failed: {}", e) })), + ) + .into_response(); + } + }; + let t_auth = start_bridge.elapsed(); + + let model_id = map_model_name(&body.model, None); + + let t_start_mapping = std::time::Instant::now(); + let signature_cache = { + let state = state_arc.read().await; + state.thought_signature_cache.clone() + }; + let messages_val = serde_json::json!(body.messages); + let (contents, _) = map_anthropic_messages_to_gemini(&messages_val, &signature_cache); + let system_instruction = body.system.as_ref().and_then(map_anthropic_system_instruction); + let t_mapping = t_start_mapping.elapsed(); + + let mut budget = 0; + if let Some(ref thinking) = body.thinking { + if thinking["type"].as_str() == Some("enabled") { + budget = thinking["budget_tokens"].as_u64().unwrap_or(2048) as i32; + } + } else { + let model_lower = model_id.to_lowercase(); + if model_lower.contains("high") || model_lower.contains("agent") { + budget = 4096; + } else if model_lower.contains("low") { + budget = 2048; + } + } + + let t_start_hashing = std::time::Instant::now(); + let messages_arr = body.messages.as_slice(); + let (prefix_hashes, new_cache_hashes) = { + let mut state = state_arc.write().await; + hash_messages_incremental_cached(messages_arr, &mut state.hash_cache) + }; + + let mut traj_id = None; + let mut sess_id = None; + + { + let state = state_arc.read().await; + if !prefix_hashes.is_empty() { + // Search all prefixes from longest (excluding full array) down to 1 + for prefix_hash in prefix_hashes.iter().take(prefix_hashes.len() - 1).rev() { + if let Some((t_id, s_id)) = state.sessions_cache.get(prefix_hash) { + traj_id = Some(t_id.clone()); + sess_id = Some(s_id.clone()); + break; + } + } + } + } + + let trajectory_id = traj_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + let session_id = + sess_id.unwrap_or_else(|| rand::Rng::gen::(&mut rand::thread_rng()).to_string()); + let t_hashing = t_start_hashing.elapsed(); + + let t_start_disk = std::time::Instant::now(); + { + let mut state = state_arc.write().await; + + state.hash_cache.insert( + 0, + crate::auth::CachedSessionHashes { + trajectory_id: trajectory_id.clone(), + hashes: new_cache_hashes, + }, + ); + state.hash_cache.truncate(10); + + if let Some(current_hash) = prefix_hashes.last() { + state.sessions_cache.insert( + current_hash.clone(), + (trajectory_id.clone(), session_id.clone()), + ); + let sessions_clone = state.sessions_cache.clone(); + tokio::spawn(async move { + crate::auth::save_sessions(&sessions_clone); + }); + } + } + let t_disk = t_start_disk.elapsed(); + + let total_setup = start_bridge.elapsed(); + let setup_ms = total_setup.as_millis() as u64; + tracing::info!( + "bridge setup finished in {}ms (auth: {}ms, mapping: {}ms, hashing: {}ms, disk_io: {}ms)", + setup_ms, + t_auth.as_millis(), + t_mapping.as_millis(), + t_hashing.as_millis(), + t_disk.as_millis() + ); + + let mut payload = serde_json::json!({ + "project": project, + "requestId": format!("agent/{}/{}/{}/1", uuid::Uuid::new_v4(), SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_millis(), uuid::Uuid::new_v4()), + "request": { + "contents": contents, + "generationConfig": { + "maxOutputTokens": body.max_tokens.unwrap_or(8192), + "thinkingConfig": { + "includeThoughts": budget > 0, + "thinkingBudget": budget + } + }, + "labels": { + "last_step_index": "1", + "model_enum": "MODEL_PLACEHOLDER_M132", + "trajectory_id": trajectory_id.clone(), + "used_claude": "false", + "used_claude_conservative": "false" + }, + "sessionId": session_id + }, + "model": model_id, + "userAgent": "antigravity", + "requestType": "chat", + "enabledCreditTypes": ["GOOGLE_ONE_AI"] + }); + + if let Some(tools) = body.tools.as_ref() { + if let Some(mapped_tools) = map_anthropic_tools_to_gemini(tools) { + payload["request"]["tools"] = mapped_tools; + } + } + + if let Some(system_inst) = system_instruction { + payload["request"]["systemInstruction"] = system_inst; + } + + tracing::info!( + "Forwarding Anthropic streamGenerateContent request for model: {} (thinkingBudget={})...", + model_id, + budget + ); + + let client = reqwest::Client::new(); + let url = "https://daily-cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse"; + + let response_res = client + .post(url) + .header("Authorization", format!("Bearer {}", token)) + .header("Content-Type", "application/json") + .header("User-Agent", "antigravity/cli/1.0.10 linux/amd64") + .header("Accept", "text/event-stream") + .header("Cache-Control", "no-cache") + .json(&payload) + .send() + .await; + + let response = match response_res { + Ok(r) => r, + Err(e) => { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ "error": format!("API request failed: {}", e) })), + ) + .into_response(); + } + }; + + if !response.status().is_success() { + let status = response.status(); + let err_text = response.text().await.unwrap_or_default(); + tracing::error!( + "daily-cloudcode API returned status {}: {}", + status, + err_text + ); + return ( + status, + Json(serde_json::json!({ "error": format!("API returned {}: {}", status, err_text) })), + ) + .into_response(); + } + + let message_id = format!( + "msg_01{}", + &uuid::Uuid::new_v4().simple().to_string()[..20] + ); + let mut bytes_stream = response.bytes_stream(); + let mut decoder = LineDecoder::new(); + + let is_streaming = body.stream.unwrap_or(false); + + if is_streaming { + let token_clone = token.clone(); + let project_clone = project.clone(); + let trajectory_id_clone = trajectory_id.clone(); + let model_id_clone = model_id.clone(); + let state_arc_clone = state_arc.clone(); + + let stream = async_stream::stream! { + let start_time = std::time::Instant::now(); + let mut first_msg_latency = None; + let mut trace_id = None; + let mut current_block_index: usize = 0; + let mut current_block_type: Option = None; + let mut has_tool_calls = false; + let mut completion_tokens = 0; + let mut prompt_tokens = 0; + + // Emit message_start first + let msg_start_val = serde_json::json!({ + "type": "message_start", + "message": { + "id": message_id.clone(), + "type": "message", + "role": "assistant", + "content": [], + "model": model_id_clone.clone(), + "stop_reason": serde_json::Value::Null, + "stop_sequence": serde_json::Value::Null, + "usage": { + "input_tokens": 0, + "output_tokens": 0 + } + } + }); + yield Ok(Event::default().event("message_start").data(msg_start_val.to_string())); + + while let Some(chunk_res) = bytes_stream.next().await { + let chunk = match chunk_res { + Ok(c) => c, + Err(e) => { + yield Err(std::io::Error::new(std::io::ErrorKind::Other, e)); + break; + } + }; + + let lines = decoder.feed(&chunk); + for line in lines { + let trimmed = line.trim(); + if !trimmed.starts_with("data: ") { + continue; + } + let data_str = &trimmed["data: ".len()..]; + if data_str == "[DONE]" { + break; + } + + if let Ok(chunk_val) = serde_json::from_str::(data_str) { + if first_msg_latency.is_none() { + first_msg_latency = Some(start_time.elapsed()); + } + if let Some(tid) = chunk_val["traceId"].as_str() { + trace_id = Some(tid.to_string()); + } + + let response_part = &chunk_val["response"]; + + // Usage metadata + if let Some(usage_meta) = response_part.get("usageMetadata") { + if let Some(candidates_tokens) = usage_meta.get("candidatesTokenCount").and_then(|v| v.as_u64()) { + completion_tokens = candidates_tokens; + } + if let Some(p_tokens) = usage_meta.get("promptTokenCount").and_then(|v| v.as_u64()) { + prompt_tokens = p_tokens; + } + } + + // Candidates and parts + if let Some(candidates) = response_part["candidates"].as_array() { + if !candidates.is_empty() { + let candidate = &candidates[0]; + if let Some(parts) = candidate["content"]["parts"].as_array() { + for part in parts { + let text = part["text"].as_str().unwrap_or(""); + let is_thought = part["thought"].as_bool().unwrap_or(false); + let func = part.get("functionCall"); + + if is_thought { + if !text.is_empty() { + if current_block_type.as_deref() != Some("thinking") { + if current_block_type.is_some() { + let stop_val = serde_json::json!({ + "type": "content_block_stop", + "index": current_block_index + }); + yield Ok(Event::default().event("content_block_stop").data(stop_val.to_string())); + current_block_index += 1; + } + current_block_type = Some("thinking".to_string()); + let start_val = serde_json::json!({ + "type": "content_block_start", + "index": current_block_index, + "content_block": { + "type": "thinking", + "thinking": "" + } + }); + yield Ok(Event::default().event("content_block_start").data(start_val.to_string())); + } + let delta_val = serde_json::json!({ + "type": "content_block_delta", + "index": current_block_index, + "delta": { + "type": "thinking_delta", + "thinking": text + } + }); + yield Ok(Event::default().event("content_block_delta").data(delta_val.to_string())); + } + } else if let Some(f) = func { + if f.is_object() { + has_tool_calls = true; + if current_block_type.is_some() { + let stop_val = serde_json::json!({ + "type": "content_block_stop", + "index": current_block_index + }); + yield Ok(Event::default().event("content_block_stop").data(stop_val.to_string())); + current_block_index += 1; + } + + let name = f["name"].as_str().unwrap_or(""); + let args_val = &f["args"]; + let args_str = if args_val.is_object() { + args_val.to_string() + } else { + args_val.as_str().unwrap_or("{}").to_string() + }; + + let tc_id = f["id"].as_str() + .map(|s| s.to_string()) + .unwrap_or_else(|| { + format!("toolu_01{}", &uuid::Uuid::new_v4().simple().to_string()[..20]) + }); + + let thought_sig = part.get("thoughtSignature") + .or_else(|| part.get("thought_signature")); + if let Some(ts_val) = thought_sig { + if let Some(ts_str) = ts_val.as_str() { + let mut state = state_arc_clone.write().await; + state.thought_signature_cache.insert(tc_id.clone(), ts_str.to_string()); + let sigs_clone = state.thought_signature_cache.clone(); + tokio::spawn(async move { + crate::auth::save_signatures(&sigs_clone); + }); + } + } + + // Start tool block + let start_val = serde_json::json!({ + "type": "content_block_start", + "index": current_block_index, + "content_block": { + "type": "tool_use", + "id": tc_id, + "name": name, + "input": {} + } + }); + yield Ok(Event::default().event("content_block_start").data(start_val.to_string())); + + // Stream the delta input + let delta_val = serde_json::json!({ + "type": "content_block_delta", + "index": current_block_index, + "delta": { + "type": "input_json_delta", + "partial_json": args_str + } + }); + yield Ok(Event::default().event("content_block_delta").data(delta_val.to_string())); + + current_block_type = Some("tool_use".to_string()); + } + } else { + if !text.is_empty() { + if current_block_type.as_deref() != Some("text") { + if current_block_type.is_some() { + let stop_val = serde_json::json!({ + "type": "content_block_stop", + "index": current_block_index + }); + yield Ok(Event::default().event("content_block_stop").data(stop_val.to_string())); + current_block_index += 1; + } + current_block_type = Some("text".to_string()); + let start_val = serde_json::json!({ + "type": "content_block_start", + "index": current_block_index, + "content_block": { + "type": "text", + "text": "" + } + }); + yield Ok(Event::default().event("content_block_start").data(start_val.to_string())); + } + let delta_val = serde_json::json!({ + "type": "content_block_delta", + "index": current_block_index, + "delta": { + "type": "text_delta", + "text": text + } + }); + yield Ok(Event::default().event("content_block_delta").data(delta_val.to_string())); + } + } + } + } + } + } + } + } + } + + if current_block_type.is_some() { + let stop_val = serde_json::json!({ + "type": "content_block_stop", + "index": current_block_index + }); + yield Ok(Event::default().event("content_block_stop").data(stop_val.to_string())); + } + + let total_latency = start_time.elapsed(); + let first_latency = first_msg_latency.unwrap_or(total_latency); + + let token_clone2 = token_clone.clone(); + let project_clone2 = project_clone.clone(); + let trajectory_id_clone2 = trajectory_id_clone.clone(); + let trace_id_clone = trace_id.clone(); + + tokio::spawn(async move { + if let Err(e) = send_telemetry_metrics( + &token_clone2, + &project_clone2, + &trajectory_id_clone2, + trace_id_clone.as_deref(), + first_latency, + total_latency, + ).await { + tracing::error!("Failed to record telemetry metrics: {:?}", e); + } + }); + + let stop_reason = if has_tool_calls { "tool_use" } else { "end_turn" }; + let msg_delta_val = serde_json::json!({ + "type": "message_delta", + "delta": { + "stop_reason": stop_reason, + "stop_sequence": serde_json::Value::Null + }, + "usage": { + "output_tokens": completion_tokens + } + }); + yield Ok(Event::default().event("message_delta").data(msg_delta_val.to_string())); + + let msg_stop_val = serde_json::json!({ + "type": "message_stop" + }); + yield Ok(Event::default().event("message_stop").data(msg_stop_val.to_string())); + + tracing::info!("Anthropic stream finished. Usage: prompt_tokens={}, completion_tokens={}", prompt_tokens, completion_tokens); + }; + + let mut response = Sse::new(stream).into_response(); + response.headers_mut().insert( + "x-bridge-setup-ms", + axum::http::HeaderValue::from_str(&setup_ms.to_string()).unwrap(), + ); + response + } else { + let mut trace_id = None; + let mut content_blocks = Vec::::new(); + let mut completion_tokens = 0; + let mut prompt_tokens = 0; + let mut stop_reason = "end_turn"; + + let start_time = std::time::Instant::now(); + let mut first_msg_latency = None; + + while let Some(chunk_res) = bytes_stream.next().await { + let chunk = match chunk_res { + Ok(c) => c, + Err(e) => { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ "error": format!("Chunk read error: {}", e) })), + ).into_response(); + } + }; + + let lines = decoder.feed(&chunk); + for line in lines { + let trimmed = line.trim(); + if !trimmed.starts_with("data: ") { + continue; + } + let data_str = &trimmed["data: ".len()..]; + if data_str == "[DONE]" { + break; + } + + if let Ok(chunk_val) = serde_json::from_str::(data_str) { + if first_msg_latency.is_none() { + first_msg_latency = Some(start_time.elapsed()); + } + if let Some(tid) = chunk_val["traceId"].as_str() { + trace_id = Some(tid.to_string()); + } + + let response_part = &chunk_val["response"]; + + // Usage metadata + if let Some(usage_meta) = response_part.get("usageMetadata") { + if let Some(candidates_tokens) = usage_meta.get("candidatesTokenCount").and_then(|v| v.as_u64()) { + completion_tokens = candidates_tokens; + } + if let Some(p_tokens) = usage_meta.get("promptTokenCount").and_then(|v| v.as_u64()) { + prompt_tokens = p_tokens; + } + } + + // Candidates and parts + if let Some(candidates) = response_part["candidates"].as_array() { + if !candidates.is_empty() { + let candidate = &candidates[0]; + if let Some(parts) = candidate["content"]["parts"].as_array() { + for part in parts { + let text = part["text"].as_str().unwrap_or(""); + let is_thought = part["thought"].as_bool().unwrap_or(false); + let func = part.get("functionCall"); + + if is_thought { + if !text.is_empty() { + if let Some(last_block) = content_blocks.last_mut() { + if last_block["type"].as_str() == Some("thinking") { + let existing = last_block["thinking"].as_str().unwrap_or(""); + last_block["thinking"] = serde_json::json!(format!("{}{}", existing, text)); + continue; + } + } + content_blocks.push(serde_json::json!({ + "type": "thinking", + "thinking": text + })); + } + } else if let Some(f) = func { + if f.is_object() { + stop_reason = "tool_use"; + let name = f["name"].as_str().unwrap_or(""); + let args_val = &f["args"]; + let args_obj = if args_val.is_object() { + args_val.clone() + } else { + serde_json::from_str::(args_val.as_str().unwrap_or("{}")).unwrap_or_else(|_| serde_json::json!({})) + }; + + let tc_id = f["id"].as_str() + .map(|s| s.to_string()) + .unwrap_or_else(|| { + format!("toolu_01{}", &uuid::Uuid::new_v4().simple().to_string()[..20]) + }); + + let thought_sig = part.get("thoughtSignature") + .or_else(|| part.get("thought_signature")); + if let Some(ts_val) = thought_sig { + if let Some(ts_str) = ts_val.as_str() { + let mut state = state_arc.write().await; + state.thought_signature_cache.insert(tc_id.clone(), ts_str.to_string()); + let sigs_clone = state.thought_signature_cache.clone(); + tokio::spawn(async move { + crate::auth::save_signatures(&sigs_clone); + }); + } + } + + content_blocks.push(serde_json::json!({ + "type": "tool_use", + "id": tc_id, + "name": name, + "input": args_obj + })); + } + } else { + if !text.is_empty() { + if let Some(last_block) = content_blocks.last_mut() { + if last_block["type"].as_str() == Some("text") { + let existing = last_block["text"].as_str().unwrap_or(""); + last_block["text"] = serde_json::json!(format!("{}{}", existing, text)); + continue; + } + } + content_blocks.push(serde_json::json!({ + "type": "text", + "text": text + })); + } + } + } + } + } + } + } + } + } + + let total_latency = start_time.elapsed(); + let first_latency = first_msg_latency.unwrap_or(total_latency); + + let token_clone2 = token.clone(); + let project_clone2 = project.clone(); + let trajectory_id_clone2 = trajectory_id.clone(); + let trace_id_clone = trace_id.clone(); + + tokio::spawn(async move { + if let Err(e) = send_telemetry_metrics( + &token_clone2, + &project_clone2, + &trajectory_id_clone2, + trace_id_clone.as_deref(), + first_latency, + total_latency, + ).await { + tracing::error!("Failed to record telemetry metrics: {:?}", e); + } + }); + + let response_payload = serde_json::json!({ + "id": message_id, + "type": "message", + "role": "assistant", + "content": content_blocks, + "model": model_id, + "stop_reason": stop_reason, + "stop_sequence": serde_json::Value::Null, + "usage": { + "input_tokens": prompt_tokens, + "output_tokens": completion_tokens + } + }); + + let mut response = Json(response_payload).into_response(); + response.headers_mut().insert( + "x-bridge-setup-ms", + axum::http::HeaderValue::from_str(&setup_ms.to_string()).unwrap(), + ); + response + } +}