b3nw/zerogravity
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add official Gemini v1beta API routes

Browse Source 
Replace /v1/gemini with proper Gemini API paths:
- POST /v1beta/models/{model}:generateContent (sync)
- POST /v1beta/models/{model}:streamGenerateContent (streaming)

Model is extracted from URL path. Uses axum wildcard
catch-all since colons in path segments are not supported.
This commit is contained in:
Nikketryhard
2026-02-16 21:46:52 -06:00
parent eb4c846b24
commit 34799fa2a9
3 changed files with 293 additions and 183 deletions
Download Patch File Download Diff File Expand all files Collapse all files

335
src/api/gemini.rs
View File

@@ -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<Arc<AppState>>,
Json(body): Json<GeminiRequest>,
axum::extract::Path(path): axum::extract::Path<String>,
Json(mut body): Json<GeminiRequest>,
) -> 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<AppState>,
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<tokio::sync::mpsc::Receiver<crate::mitm::store::MitmEvent>>,
) -> axum::response::Response {
let has_custom_tools = state.mitm_store.get_tools().await.is_some();
// Clear stale response and upstream errors
// 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;
// ── MITM bypass: when tools active, poll MitmStore directly ──
if has_custom_tools {
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() {
// ── MITM bypass: channel-based pipeline when custom tools active ──
if let Some(mut rx) = mitm_rx {
let start = std::time::Instant::now();
let mut acc_text = String::new();
let mut acc_thinking: Option<String> = 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<serde_json::Value> = 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,31 +491,31 @@ 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;
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<serde_json::Value> = Vec::new();
if let Some(ref t) = thinking {
if let Some(ref t) = acc_thinking {
parts.push(serde_json::json!({"text": t, "thought": true}));
}
parts.push(serde_json::json!({"text": text}));
parts.push(serde_json::json!({"text": acc_text}));
state.mitm_store.drop_channel().await;
return Json(serde_json::json!({
"candidates": [{
"content": {
@@ -455,11 +529,15 @@ async fn gemini_sync(
}))
.into_response();
}
tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;
MitmEvent::UpstreamError(err) => {
state.mitm_store.drop_channel().await;
return upstream_err_response(&err);
}
}
}
// 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,36 +621,63 @@ async fn gemini_stream(
model_name: String,
cascade_id: String,
timeout: u64,
mitm_rx: Option<tokio::sync::mpsc::Receiver<crate::mitm::store::MitmEvent>>,
) -> 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
// 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;
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,
"candidates": [{
"content": {
"parts": [{"text": delta, "thought": true}],
"role": "model",
},
}],
"modelVersion": model_name,
})).unwrap_or_default()));
}
})).unwrap()));
yield Ok(Event::default().data("[DONE]".to_string()));
break;
}
MitmEvent::TextDelta(full_text) => {
if full_text.len() > last_text.len() {
let delta = full_text[last_text.len()..].to_string();
// ── 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() {
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<serde_json::Value> = calls
.iter()
.map(|fc| {
@@ -598,51 +703,11 @@ async fn gemini_stream(
"modelVersion": model_name,
})).unwrap_or_default()));
yield Ok(Event::default().data("[DONE]"));
state.mitm_store.drop_channel().await;
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() {
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!({
@@ -657,17 +722,57 @@ async fn gemini_stream(
"modelVersion": model_name,
})).unwrap_or_default()));
yield Ok(Event::default().data("[DONE]"));
state.mitm_store.drop_channel().await;
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;
} 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(_) => {}
}
}
let poll_ms: u64 = rand::thread_rng().gen_range(150..300);
tokio::time::sleep(tokio::time::Duration::from_millis(poll_ms)).await;
continue;
// 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 {
// ── Normal LS path (no custom tools) ──
if let Ok((status, data)) = state.backend.get_steps(&cascade_id).await {

8
src/api/mod.rs
View File

@@ -43,7 +43,10 @@ pub fn router(state: Arc<AppState>) -> 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<serde_json::Value> {
"endpoints": [
"/v1/chat/completions",
"/v1/responses",
"/v1/gemini",
"/v1beta/models/{model}:generateContent",
"/v1beta/models/{model}:streamGenerateContent",
"/v1/models",
"/v1/sessions",

3
src/main.rs
View File

@@ -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");