diff --git a/src/api/gemini.rs b/src/api/gemini.rs index 80110d2..31d96fb 100644 --- a/src/api/gemini.rs +++ b/src/api/gemini.rs @@ -1,9 +1,13 @@ -//! Gemini-native endpoint (/v1/gemini) — zero-translation tool call passthrough. +//! Gemini-native endpoint — zero-translation tool call passthrough. +//! +//! Routes: +//! POST /v1/gemini (custom, stream via body flag) +//! POST /v1beta/models/{model}:generateContent (sync, official Gemini API path) +//! POST /v1beta/models/{model}:streamGenerateContent (stream, official Gemini API path) //! //! Accepts tools in Gemini `functionDeclarations` format directly, //! returns `functionCall` in Gemini format directly. //! No OpenAI ↔ Gemini format conversion. -//! Supports both sync (`stream: false`) and streaming (`stream: true`) modes. use axum::{ extract::State, @@ -12,7 +16,7 @@ use axum::{ }; use rand::Rng; use std::sync::Arc; -use tracing::{info, warn}; +use tracing::{debug, info, warn}; use super::models::{lookup_model, DEFAULT_MODEL, MODELS}; use super::polling::{ @@ -115,9 +119,51 @@ async fn build_usage_metadata( } } -pub(crate) async fn handle_gemini( +/// POST /v1beta/*path — handles both :generateContent and :streamGenerateContent +/// +/// Parses paths like: +/// /v1beta/models/gemini-3-flash:generateContent +/// /v1beta/models/gemini-3-flash:streamGenerateContent +pub(crate) async fn handle_gemini_v1beta( State(state): State>, - Json(body): Json, + axum::extract::Path(path): axum::extract::Path, + Json(mut body): Json, +) -> axum::response::Response { + // Expected path: "models/{model}:{action}" + let path = path.trim_start_matches('/'); + let model_action = path.strip_prefix("models/").unwrap_or(path); + + // Split on the LAST colon — model name may contain hyphens/numbers but not colons + if let Some(colon_pos) = model_action.rfind(':') { + let model = &model_action[..colon_pos]; + let action = &model_action[colon_pos + 1..]; + + body.model = Some(model.to_string()); + match action { + "generateContent" => body.stream = false, + "streamGenerateContent" => body.stream = true, + _ => { + return err_response( + StatusCode::BAD_REQUEST, + format!("Unknown action: {action}. Use :generateContent or :streamGenerateContent"), + "invalid_request_error", + ); + } + } + } else { + return err_response( + StatusCode::BAD_REQUEST, + format!("Invalid path: /v1beta/{path}. Expected /v1beta/models/{{model}}:generateContent"), + "invalid_request_error", + ); + } + + handle_gemini_inner(state, body).await +} + +async fn handle_gemini_inner( + state: Arc, + body: GeminiRequest, ) -> axum::response::Response { info!( "POST /v1/gemini model={} stream={}", @@ -209,7 +255,11 @@ pub(crate) async fn handle_gemini( count = tools.len(), "Stored Gemini-native tools for MITM injection" ); + } else { + state.mitm_store.clear_tools().await; } + } else { + state.mitm_store.clear_tools().await; } if let Some(ref config) = body.tool_config { state.mitm_store.set_tool_config(config.clone()).await; @@ -317,6 +367,8 @@ pub(crate) async fn handle_gemini( // Send message state.mitm_store.set_active_cascade(&cascade_id).await; + // Store real user text for MITM injection — LS gets a dummy prompt + state.mitm_store.set_pending_user_text(user_text.clone()).await; // Store image for MITM injection (LS doesn't forward images to Google API) if let Some(ref img) = image { use base64::Engine; @@ -328,9 +380,24 @@ pub(crate) async fn handle_gemini( }) .await; } + + // Pre-flight: install channel BEFORE send_message so the MITM proxy + // can grab it when the LS fires its API call. + let has_custom_tools = state.mitm_store.get_tools().await.is_some(); + let mitm_rx = if has_custom_tools { + state.mitm_store.clear_response_async().await; + state.mitm_store.clear_upstream_error().await; + let _ = state.mitm_store.take_any_function_calls().await; + let (tx, rx) = tokio::sync::mpsc::channel(64); + state.mitm_store.set_channel(tx).await; + Some(rx) + } else { + None + }; + match state .backend - .send_message_with_image(&cascade_id, &user_text, model.model_enum, image.as_ref()) + .send_message_with_image(&cascade_id, ".", model.model_enum, image.as_ref()) .await { Ok((200, _)) => { @@ -341,6 +408,7 @@ pub(crate) async fn handle_gemini( }); } Ok((status, _)) => { + state.mitm_store.drop_channel().await; return err_response( StatusCode::BAD_GATEWAY, format!("Antigravity returned {status}"), @@ -348,6 +416,7 @@ pub(crate) async fn handle_gemini( ); } Err(e) => { + state.mitm_store.drop_channel().await; return err_response( StatusCode::BAD_GATEWAY, format!("Send message failed: {e}"), @@ -360,9 +429,9 @@ pub(crate) async fn handle_gemini( let model_name = model_name.to_string(); let timeout = body.timeout; if body.stream { - gemini_stream(state, model_name, cascade_id, timeout).await + gemini_stream(state, model_name, cascade_id, timeout, mitm_rx).await } else { - gemini_sync(state, model_name, cascade_id, timeout).await + gemini_sync(state, model_name, cascade_id, timeout, mitm_rx).await } } @@ -373,26 +442,31 @@ async fn gemini_sync( model_name: String, cascade_id: String, timeout: u64, + mitm_rx: Option>, ) -> axum::response::Response { - let has_custom_tools = state.mitm_store.get_tools().await.is_some(); + // Clear stale response and upstream errors (only if no pre-installed channel) + if mitm_rx.is_none() { + state.mitm_store.clear_response_async().await; + state.mitm_store.clear_upstream_error().await; + } - // Clear stale response and upstream errors - state.mitm_store.clear_response_async().await; - state.mitm_store.clear_upstream_error().await; - - // ── MITM bypass: when tools active, poll MitmStore directly ── - if has_custom_tools { + // ── MITM bypass: channel-based pipeline when custom tools active ── + if let Some(mut rx) = mitm_rx { let start = std::time::Instant::now(); - while start.elapsed().as_secs() < timeout { - // Check for upstream errors from MITM (Google API errors) - if let Some(err) = state.mitm_store.take_upstream_error().await { - return upstream_err_response(&err); - } - // Check for function calls - let captured = state.mitm_store.take_any_function_calls().await; - if let Some(ref calls) = captured { - if !calls.is_empty() { + let mut acc_text = String::new(); + let mut acc_thinking: Option = None; + + while let Some(event) = tokio::time::timeout( + std::time::Duration::from_secs(timeout.saturating_sub(start.elapsed().as_secs())), + rx.recv(), + ).await.ok().flatten() { + use crate::mitm::store::MitmEvent; + match event { + MitmEvent::ThinkingDelta(t) => { acc_thinking = Some(t); } + MitmEvent::TextDelta(t) => { acc_text = t; } + MitmEvent::Usage(_) | MitmEvent::Grounding(_) => {} + MitmEvent::FunctionCall(calls) => { let parts: Vec = calls .iter() .map(|fc| { @@ -404,7 +478,7 @@ async fn gemini_sync( }) }) .collect(); - + state.mitm_store.drop_channel().await; return Json(serde_json::json!({ "candidates": [{ "content": { @@ -417,49 +491,53 @@ async fn gemini_sync( })) .into_response(); } - } - - // Check for completed text response - if state.mitm_store.is_response_complete() { - let text = state - .mitm_store - .take_response_text() - .await - .unwrap_or_default(); - let thinking = state.mitm_store.take_thinking_text().await; - - // Guard against stale response_complete with no data - if text.is_empty() && thinking.is_none() { - warn!("Gemini sync bypass: stale response_complete — clearing"); - state.mitm_store.clear_response_async().await; - tokio::time::sleep(tokio::time::Duration::from_millis(200)).await; - continue; + MitmEvent::ResponseComplete => { + if acc_text.is_empty() && acc_thinking.is_none() { + // Empty response — continue waiting + continue; + } + if acc_text.is_empty() && acc_thinking.is_some() { + // Thinking-only — LS needs to make a follow-up request. + // Reinstall channel and unblock gate. + let (new_tx, new_rx) = tokio::sync::mpsc::channel(64); + state.mitm_store.set_channel(new_tx).await; + state.mitm_store.clear_request_in_flight(); + let _ = state.mitm_store.take_any_function_calls().await; + rx = new_rx; + debug!( + "Gemini sync: thinking-only — new channel for follow-up, thinking_len={}", + acc_thinking.as_ref().map(|t| t.len()).unwrap_or(0) + ); + continue; + } + let mut parts: Vec = Vec::new(); + if let Some(ref t) = acc_thinking { + parts.push(serde_json::json!({"text": t, "thought": true})); + } + parts.push(serde_json::json!({"text": acc_text})); + state.mitm_store.drop_channel().await; + return Json(serde_json::json!({ + "candidates": [{ + "content": { + "parts": parts, + "role": "model", + }, + "finishReason": "STOP", + }], + "modelVersion": model_name, + "usageMetadata": build_usage_metadata(&state.mitm_store, &cascade_id).await, + })) + .into_response(); } - - let mut parts: Vec = Vec::new(); - if let Some(ref t) = thinking { - parts.push(serde_json::json!({"text": t, "thought": true})); + MitmEvent::UpstreamError(err) => { + state.mitm_store.drop_channel().await; + return upstream_err_response(&err); } - parts.push(serde_json::json!({"text": text})); - - return Json(serde_json::json!({ - "candidates": [{ - "content": { - "parts": parts, - "role": "model", - }, - "finishReason": "STOP", - }], - "modelVersion": model_name, - "usageMetadata": build_usage_metadata(&state.mitm_store, &cascade_id).await, - })) - .into_response(); } - - tokio::time::sleep(tokio::time::Duration::from_millis(200)).await; } - // Timeout — return proper error with status code + // Timeout + state.mitm_store.drop_channel().await; return ( axum::http::StatusCode::GATEWAY_TIMEOUT, Json(serde_json::json!({ @@ -543,131 +621,158 @@ async fn gemini_stream( model_name: String, cascade_id: String, timeout: u64, + mitm_rx: Option>, ) -> axum::response::Response { let stream = async_stream::stream! { let start = std::time::Instant::now(); let mut last_text = String::new(); let mut last_thinking = String::new(); - let has_custom_tools = state.mitm_store.get_tools().await.is_some(); - // Clear stale response - state.mitm_store.clear_response_async().await; + // Clear stale response (only if no pre-installed channel) + if mitm_rx.is_none() { + state.mitm_store.clear_response_async().await; + } + + // ── MITM bypass path (custom tools active) ── + // Channel-based pipeline: read events directly from MITM proxy. + // Channel is pre-installed before send_message to avoid race conditions. + if let Some(mut rx) = mitm_rx { + let mut did_unblock_ls = false; + + 'channel_loop: while let Some(event) = tokio::time::timeout( + std::time::Duration::from_secs(timeout.saturating_sub(start.elapsed().as_secs())), + rx.recv(), + ).await.ok().flatten() { + use crate::mitm::store::MitmEvent; + match event { + MitmEvent::ThinkingDelta(full_thinking) => { + if full_thinking.len() > last_thinking.len() { + let delta = full_thinking[last_thinking.len()..].to_string(); + last_thinking = full_thinking; + + yield Ok::<_, std::convert::Infallible>(Event::default().data(serde_json::to_string(&serde_json::json!({ + "candidates": [{ + "content": { + "parts": [{"text": delta, "thought": true}], + "role": "model", + }, + }], + "modelVersion": model_name, + })).unwrap_or_default())); + } + } + MitmEvent::TextDelta(full_text) => { + if full_text.len() > last_text.len() { + let delta = full_text[last_text.len()..].to_string(); + + yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({ + "candidates": [{ + "content": { + "parts": [{"text": delta}], + "role": "model", + }, + }], + "modelVersion": model_name, + })).unwrap_or_default())); + last_text = full_text; + } + } + MitmEvent::FunctionCall(calls) => { + let parts: Vec = calls + .iter() + .map(|fc| { + serde_json::json!({ + "functionCall": { + "name": fc.name, + "args": fc.args, + } + }) + }) + .collect(); + + let usage_meta = build_usage_metadata(&state.mitm_store, &cascade_id).await; + yield Ok::<_, std::convert::Infallible>(Event::default().data(serde_json::to_string(&serde_json::json!({ + "candidates": [{ + "content": { + "parts": parts, + "role": "model", + }, + "finishReason": "STOP", + }], + "usageMetadata": usage_meta, + "modelVersion": model_name, + })).unwrap_or_default())); + yield Ok(Event::default().data("[DONE]")); + state.mitm_store.drop_channel().await; + return; + } + MitmEvent::ResponseComplete => { + if !last_text.is_empty() { + // Final chunk with finishReason + usageMetadata + let usage_meta = build_usage_metadata(&state.mitm_store, &cascade_id).await; + yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({ + "candidates": [{ + "content": { + "parts": [{"text": ""}], + "role": "model", + }, + "finishReason": "STOP", + }], + "usageMetadata": usage_meta, + "modelVersion": model_name, + })).unwrap_or_default())); + yield Ok(Event::default().data("[DONE]")); + state.mitm_store.drop_channel().await; + return; + } else if !last_thinking.is_empty() && !did_unblock_ls { + // Thinking-only response — LS needs follow-up API calls. + // Create a new channel and unblock the gate. + did_unblock_ls = true; + let (new_tx, new_rx) = tokio::sync::mpsc::channel(64); + state.mitm_store.set_channel(new_tx).await; + state.mitm_store.clear_request_in_flight(); + let _ = state.mitm_store.take_any_function_calls().await; + rx = new_rx; + debug!( + "Gemini stream: thinking-only — new channel for follow-up, thinking_len={}", + last_thinking.len() + ); + continue 'channel_loop; + } + // Empty response — continue waiting + } + MitmEvent::UpstreamError(err) => { + let error_msg = super::util::upstream_error_message(&err); + let error_type = super::util::upstream_error_type(&err); + yield Ok::<_, std::convert::Infallible>(Event::default().data(serde_json::to_string(&serde_json::json!({ + "error": { + "message": error_msg, + "type": error_type, + "code": err.status, + } + })).unwrap())); + yield Ok(Event::default().data("[DONE]")); + state.mitm_store.drop_channel().await; + return; + } + MitmEvent::Usage(_) | MitmEvent::Grounding(_) => {} + } + } + + // Timeout or channel closed + state.mitm_store.drop_channel().await; + yield Ok::<_, std::convert::Infallible>(Event::default().data(serde_json::to_string(&serde_json::json!({ + "error": { + "message": format!("Timeout: no response from Google API after {timeout}s"), + "type": "upstream_error", + "code": 504, + } + })).unwrap())); + yield Ok(Event::default().data("[DONE]")); + return; + } while start.elapsed().as_secs() < timeout { - // Check for upstream errors from MITM (Google API errors) - if let Some(err) = state.mitm_store.take_upstream_error().await { - let error_msg = super::util::upstream_error_message(&err); - let error_type = super::util::upstream_error_type(&err); - yield Ok::<_, std::convert::Infallible>(Event::default().data(serde_json::to_string(&serde_json::json!({ - "error": { - "message": error_msg, - "type": error_type, - "code": err.status, - } - })).unwrap())); - yield Ok(Event::default().data("[DONE]".to_string())); - break; - } - - // ── Check for MITM-captured function calls FIRST ── - let captured = state.mitm_store.take_any_function_calls().await; - if let Some(ref calls) = captured { - if !calls.is_empty() { - let parts: Vec = calls - .iter() - .map(|fc| { - serde_json::json!({ - "functionCall": { - "name": fc.name, - "args": fc.args, - } - }) - }) - .collect(); - - let usage_meta = build_usage_metadata(&state.mitm_store, &cascade_id).await; - yield Ok::<_, std::convert::Infallible>(Event::default().data(serde_json::to_string(&serde_json::json!({ - "candidates": [{ - "content": { - "parts": parts, - "role": "model", - }, - "finishReason": "STOP", - }], - "usageMetadata": usage_meta, - "modelVersion": model_name, - })).unwrap_or_default())); - yield Ok(Event::default().data("[DONE]")); - return; - } - } - - // ── MITM bypass path (custom tools active) ── - if has_custom_tools { - // Stream thinking - if let Some(tc) = state.mitm_store.peek_thinking_text().await { - if tc.len() > last_thinking.len() { - let delta = tc[last_thinking.len()..].to_string(); - last_thinking = tc; - - yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({ - "candidates": [{ - "content": { - "parts": [{"text": delta, "thought": true}], - "role": "model", - }, - }], - "modelVersion": model_name, - })).unwrap_or_default())); - } - } - - // Stream text - if let Some(text) = state.mitm_store.peek_response_text().await { - if !text.is_empty() && text.len() > last_text.len() { - let delta = &text[last_text.len()..]; - - yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({ - "candidates": [{ - "content": { - "parts": [{"text": delta}], - "role": "model", - }, - }], - "modelVersion": model_name, - })).unwrap_or_default())); - last_text = text; - } - } - - // Check completion - let complete = state.mitm_store.is_response_complete(); - if complete && !last_text.is_empty() { - // Final chunk with finishReason + usageMetadata - let usage_meta = build_usage_metadata(&state.mitm_store, &cascade_id).await; - yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({ - "candidates": [{ - "content": { - "parts": [{"text": ""}], - "role": "model", - }, - "finishReason": "STOP", - }], - "usageMetadata": usage_meta, - "modelVersion": model_name, - })).unwrap_or_default())); - yield Ok(Event::default().data("[DONE]")); - return; - } else if complete && last_text.is_empty() && last_thinking.is_empty() { - // Stale state — clear and retry - warn!("Gemini stream bypass: stale response_complete — clearing"); - state.mitm_store.clear_response_async().await; - } - - let poll_ms: u64 = rand::thread_rng().gen_range(150..300); - tokio::time::sleep(tokio::time::Duration::from_millis(poll_ms)).await; - continue; - } // ── Normal LS path (no custom tools) ── if let Ok((status, data)) = state.backend.get_steps(&cascade_id).await { diff --git a/src/api/mod.rs b/src/api/mod.rs index 29a0f68..ae7007f 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -43,7 +43,10 @@ pub fn router(state: Arc) -> Router { "/v1/chat/completions", post(completions::handle_completions), ) - .route("/v1/gemini", post(gemini::handle_gemini)) + .route( + "/v1beta/{*path}", + post(gemini::handle_gemini_v1beta), + ) .route("/v1/models", get(handle_models)) .route("/v1/sessions", get(handle_list_sessions)) .route("/v1/sessions/{id}", delete(handle_delete_session)) @@ -67,7 +70,8 @@ async fn handle_root() -> Json { "endpoints": [ "/v1/chat/completions", "/v1/responses", - "/v1/gemini", + "/v1beta/models/{model}:generateContent", + "/v1beta/models/{model}:streamGenerateContent", "/v1/models", "/v1/sessions", diff --git a/src/main.rs b/src/main.rs index ca017dd..666c9ab 100644 --- a/src/main.rs +++ b/src/main.rs @@ -412,7 +412,8 @@ fn print_banner( println!(" \x1b[1mroutes\x1b[0m"); println!(" \x1b[33m POST\x1b[0m /v1/responses"); println!(" \x1b[33m POST\x1b[0m /v1/chat/completions"); - println!(" \x1b[33m POST\x1b[0m /v1/gemini"); + println!(" \x1b[33m POST\x1b[0m /v1beta/models/:model:generateContent"); + println!(" \x1b[33m POST\x1b[0m /v1beta/models/:model:streamGenerateContent"); println!(" \x1b[32m GET \x1b[0m /v1/models"); println!(" \x1b[32m GET \x1b[0m /v1/sessions"); println!(" \x1b[31m DEL \x1b[0m /v1/sessions/:id");