use axum::{
extract::State,
http::StatusCode,
response::{
sse::{Event, Sse},
IntoResponse,
},
Json,
};
use futures::stream::StreamExt;
use serde::Deserialize;
use serde_json::Value;
use std::time::{SystemTime, UNIX_EPOCH};
use crate::auth::{get_valid_token_and_project, SharedState};
use crate::mapping::{
map_anthropic_messages_to_gemini, map_anthropic_system_instruction,
map_anthropic_tools_to_gemini, map_model_name, map_openai_messages_to_gemini,
map_openai_tools_to_gemini,
};
use crate::telemetry::send_telemetry_metrics;
#[derive(Debug, Deserialize)]
pub struct ChatCompletionsRequest {
pub messages: Vec<Value>,
pub model: Option<String>,
pub thinking_level: Option<String>,
pub tools: Option<Value>,
}
#[allow(dead_code)]
#[derive(Debug, Deserialize)]
pub struct AnthropicMessagesRequest {
pub model: String,
pub messages: Vec<Value>,
pub system: Option<Value>,
pub max_tokens: Option<u32>,
pub metadata: Option<Value>,
pub stop_sequences: Option<Vec<String>>,
pub stream: Option<bool>,
pub temperature: Option<f32>,
pub tool_choice: Option<Value>,
pub tools: Option<Value>,
pub thinking: Option<Value>,
}
pub struct LineDecoder {
buffer: Vec<u8>,
}
impl LineDecoder {
pub fn new() -> Self {
Self { buffer: Vec::new() }
}
pub fn feed(&mut self, chunk: &[u8]) -> Vec<String> {
self.buffer.extend_from_slice(chunk);
let mut lines = Vec::new();
while let Some(pos) = self.buffer.iter().position(|&b| b == b'\n') {
let line_bytes = self.buffer.drain(..=pos).collect::<Vec<u8>>();
let mut s = String::from_utf8_lossy(&line_bytes).into_owned();
if s.ends_with('\n') {
s.pop();
}
if s.ends_with('\r') {
s.pop();
}
lines.push(s);
}
lines
}
}
/// Translates Gemini `usageMetadata` into OpenAI-shaped usage WITHOUT inventing zeros.
///
/// Token counts Gemini always reports are defaulted to 0 as before. `cachedContentTokenCount`
/// is different: Gemini omits it entirely when no cached content was used AND when it does not
/// measure caching for the request at all. Defaulting that to 0 would present an unmeasured
/// request as a measured cache miss, which downstream cannot distinguish from a real reported
/// zero. So the field is OMITTED when absent and passed through when present - absence stays
/// absent all the way to the consumer.
fn usage_from_gemini(usage_meta: &serde_json::Value) -> serde_json::Value {
let mut usage = serde_json::json!({
"prompt_tokens": usage_meta.get("promptTokenCount").unwrap_or(&serde_json::json!(0)),
"completion_tokens": usage_meta.get("candidatesTokenCount").unwrap_or(&serde_json::json!(0)),
"total_tokens": usage_meta.get("totalTokenCount").unwrap_or(&serde_json::json!(0)),
});
if let Some(cached) = usage_meta.get("cachedContentTokenCount") {
usage["cached_prompt_tokens"] = cached.clone();
}
usage
}
fn hash_messages_incremental_cached(
messages: &[Value],
hash_cache: &mut Vec<crate::auth::CachedSessionHashes>,
) -> (Vec<String>, Vec<(Value, String)>) {
use sha2::{Digest, Sha256};
let mut best_match_index = None;
let mut best_match_len = 0;
for (cache_idx, cached) in hash_cache.iter().enumerate() {
let mut match_len = 0;
for (i, msg) in messages.iter().enumerate() {
if i < cached.hashes.len() && cached.hashes[i].0 == *msg {
match_len += 1;
} else {
break;
}
}
if match_len > best_match_len {
best_match_len = match_len;
best_match_index = Some(cache_idx);
}
}
let mut hashes = Vec::with_capacity(messages.len());
let mut current_hash = [0u8; 32];
let mut new_cache_hashes = Vec::with_capacity(messages.len());
if let Some(cache_idx) = best_match_index {
let cached = &hash_cache[cache_idx];
for (message, (_, hash)) in messages.iter().zip(&cached.hashes).take(best_match_len) {
hashes.push(hash.clone());
new_cache_hashes.push((message.clone(), hash.clone()));
}
if best_match_len > 0 {
let last_hex = &cached.hashes[best_match_len - 1].1;
if let Ok(bytes) = hex::decode(last_hex) {
if bytes.len() == 32 {
current_hash.copy_from_slice(&bytes);
}
}
}
}
for msg in messages.iter().skip(best_match_len) {
let msg_json = serde_json::to_string(msg).unwrap_or_default();
let mut msg_hasher = Sha256::new();
msg_hasher.update(msg_json.as_bytes());
let msg_hash = msg_hasher.finalize();
let mut chain_hasher = Sha256::new();
chain_hasher.update(current_hash);
chain_hasher.update(msg_hash);
let final_hash = chain_hasher.finalize();
current_hash.copy_from_slice(&final_hash);
let hash_str = hex::encode(final_hash);
hashes.push(hash_str.clone());
new_cache_hashes.push((msg.clone(), hash_str));
}
if let Some(cache_idx) = best_match_index {
hash_cache.remove(cache_idx);
}
(hashes, new_cache_hashes)
}
struct ResolvedSession {
trajectory_id: String,
session_id: String,
setup_ms: u64,
t_hashing: std::time::Duration,
t_disk: std::time::Duration,
}
async fn resolve_trajectory_and_session(
state_arc: &SharedState,
messages_arr: &[Value],
start_bridge: std::time::Instant,
) -> ResolvedSession {
let t_start_hashing = std::time::Instant::now();
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::<i64>(&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;
ResolvedSession {
trajectory_id,
session_id,
setup_ms,
t_hashing,
t_disk,
}
}
pub async fn chat_completions(
State(state_arc): State<SharedState>,
Json(body): Json<ChatCompletionsRequest>,
) -> 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_name = body.model.as_deref().unwrap_or("gemini-3.5-flash");
let model_id = map_model_name(model_name, body.thinking_level.as_deref());
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, system_instruction) =
map_openai_messages_to_gemini(&messages_val, &signature_cache);
let t_mapping = t_start_mapping.elapsed();
let mut budget = 0;
if let Some(ref thinking_level) = body.thinking_level {
match thinking_level.to_lowercase().as_str() {
"high" => budget = 4096,
"medium" => budget = 2048,
"low" => budget = 1024,
_ => budget = 2048,
}
} 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 resolved =
resolve_trajectory_and_session(&state_arc, body.messages.as_slice(), start_bridge).await;
let trajectory_id = resolved.trajectory_id;
let session_id = resolved.session_id;
let setup_ms = resolved.setup_ms;
let t_hashing = resolved.t_hashing;
let t_disk = resolved.t_disk;
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": 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_openai_tools_to_gemini(tools) {
payload["request"]["tools"] = mapped_tools;
}
}
if let Some(system_inst) = system_instruction {
payload["request"]["systemInstruction"] = system_inst;
}
tracing::info!(
"Forwarding 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 request_id = format!(
"chatcmpl-{}",
&uuid::Uuid::new_v4().simple().to_string()[..12]
);
let mut bytes_stream = response.bytes_stream();
let mut decoder = LineDecoder::new();
let mut tool_call_index = 0;
let token_clone = token.clone();
let project_clone = project.clone();
let trajectory_id_clone = trajectory_id.clone();
let stream = async_stream::stream! {
let start_time = std::time::Instant::now();
let mut first_msg_latency = None;
let mut trace_id = None;
while let Some(chunk_res) = bytes_stream.next().await {
let chunk = match chunk_res {
Ok(c) => c,
Err(e) => {
yield Err(std::io::Error::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::<Value>(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") {
tracing::info!("USAGE METADATA: {}", usage_meta);
let out = serde_json::json!({
"id": request_id,
"object": "chat.completion.chunk",
"created": SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(),
"model": model_id,
"choices": [],
"usage": usage_from_gemini(usage_meta)
});
yield Ok(Event::default().data(out.to_string()));
}
// Candidates and parts
if let Some(candidates) = response_part["candidates"].as_array() {
if !candidates.is_empty() {
let candidate = &candidates[0];
let mut finish_reason = serde_json::Value::Null;
if let Some(fr) = candidate.get("finishReason").and_then(|v| v.as_str()) {
let mapped = match fr {
"STOP" => {
if tool_call_index > 0 {
"tool_calls"
} else {
"stop"
}
}
"MAX_TOKENS" => "length",
"SAFETY" => "content_filter",
"RECITATION" => "content_filter",
_ => "stop",
};
finish_reason = serde_json::json!(mapped);
}
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");
let in_thinking_phase = is_thought;
if in_thinking_phase {
if !text.is_empty() {
let out = serde_json::json!({
"id": request_id,
"object": "chat.completion.chunk",
"created": SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(),
"model": model_id,
"choices": [{
"index": 0,
"delta": {"reasoning_content": text},
"finish_reason": serde_json::Value::Null
}]
});
yield Ok(Event::default().data(out.to_string()));
}
} else if let Some(f) = func {
if f.is_object() {
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!("call_{}", &uuid::Uuid::new_v4().simple().to_string()[..12])
});
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);
});
}
}
let out = serde_json::json!({
"id": request_id,
"object": "chat.completion.chunk",
"created": SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(),
"model": model_id,
"choices": [{
"index": 0,
"delta": {
"tool_calls": [{
"index": tool_call_index,
"id": tc_id,
"type": "function",
"function": {
"name": name,
"arguments": args_str
}
}]
},
"finish_reason": "tool_calls"
}]
});
tool_call_index += 1;
yield Ok(Event::default().data(out.to_string()));
}
} else {
if !text.is_empty() {
let out = serde_json::json!({
"id": request_id,
"object": "chat.completion.chunk",
"created": SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(),
"model": model_id,
"choices": [{
"index": 0,
"delta": {"content": text},
"finish_reason": serde_json::Value::Null
}]
});
yield Ok(Event::default().data(out.to_string()));
}
}
}
}
if finish_reason != serde_json::Value::Null {
let out = serde_json::json!({
"id": request_id,
"object": "chat.completion.chunk",
"created": SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(),
"model": model_id,
"choices": [{
"index": 0,
"delta": {},
"finish_reason": finish_reason
}]
});
yield Ok(Event::default().data(out.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);
}
});
yield Ok(Event::default().data("[DONE]"));
};
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
}
pub async fn list_models() -> impl IntoResponse {
Json(serde_json::json!({
"data": [
{"id": "Gemini 3.5 Flash (Medium)", "object": "model"},
{"id": "Gemini 3.5 Flash (High)", "object": "model"},
{"id": "Gemini 3.5 Flash (Low)", "object": "model"},
{"id": "Gemini 3.1 Pro (Low)", "object": "model"},
{"id": "Gemini 3.1 Pro (High)", "object": "model"},
{"id": "Claude Sonnet 4.6 (Thinking)", "object": "model"},
{"id": "Claude Opus 4.6 (Thinking)", "object": "model"},
{"id": "GPT-OSS 120B (Medium)", "object": "model"}
]
}))
}
pub async fn anthropic_messages(
State(state_arc): State<SharedState>,
Json(body): Json<AnthropicMessagesRequest>,
) -> 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 resolved =
resolve_trajectory_and_session(&state_arc, body.messages.as_slice(), start_bridge).await;
let trajectory_id = resolved.trajectory_id;
let session_id = resolved.session_id;
let setup_ms = resolved.setup_ms;
let t_hashing = resolved.t_hashing;
let t_disk = resolved.t_disk;
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<String> = 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::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::<Value>(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::<Value>::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::<Value>(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::<Value>(
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
}
}