feat(proxy): wire Gemini native v1beta endpoints through the conductor

Implement the four Gemini native inbound endpoints so Gemini CLI / Google
GenAI SDKs can point their base URL at the gateway:

- GET  /v1beta/models — directory listing in Gemini shape, built from
  configured routes (like /v1/models); no upstream call, no billing.
- POST /v1beta/models/{model}:generateContent — non-streaming, routed
  through ProxyConductor with failover + usage_logs + billing.
- POST /v1beta/models/{model}:streamGenerateContent — SSE passthrough via
  open_stream_with_failover + sse_passthrough_response.
- POST /v1beta/models/{model}:countTokens — metadata call, routed for a
  channel but not billed.

Auth accepts the gateway key from ?key= or x-goog-api-key (Gemini
convention) and reuses the downstream-key authorization + limits.

Billing closes the loop: gemini_extract_usage reads usageMetadata
(promptTokenCount / candidatesTokenCount) for non-streaming, and
gemini_stream_usage_from_line feeds the shared StreamUsageAccumulator so
streamed usageMetadata is billed too.

Unit tests cover model:action parsing, usage extraction, the stream usage
parser, and query forwarding (key/alt dropped).
This commit is contained in:
boos
2026-06-22 08:39:38 +08:00
parent aaa593ac31
commit 4d093548de
+509
View File
@@ -73,6 +73,16 @@ pub fn routes() -> Router<AppState> {
.route("/v1/videos", post(videos_handler))
.route("/v1/videos/{id}", get(video_task_handler))
.route("/v1/models", get(list_models))
// ── Gemini native inbound (v1beta) ────────────────────────────
// Lets Gemini CLI / Google GenAI SDKs point their base URL straight
// at this gateway. The model lives in the path and the key in the
// `?key=` query (or x-goog-api-key header), so requests pass through
// to a Gemini-shaped upstream without body rewriting.
.route("/v1beta/models", get(gemini_list_models))
.route(
"/v1beta/models/{model_action}",
post(gemini_model_operation),
)
// ── ComfyUI shim aliases ──────────────────────────────────────
// ComfyUI's OpenAI nodes hit /proxy/openai/... paths (see
// comfy_api_nodes/nodes_openai.py). Mounting them here as aliases
@@ -157,6 +167,209 @@ async fn responses_handler(
}
}
// ── Gemini native inbound handlers (v1beta) ────────────────────────────
//
// Gemini clients (Gemini CLI / Google GenAI SDK) carry the model in the URL
// path (`/v1beta/models/{model}:{op}`) and the gateway key in `?key=` or the
// `x-goog-api-key` header. We authenticate against our downstream keys, route
// the Gemini model through the same ProxyConductor as the OpenAI endpoints, and
// forward the body verbatim to a Gemini-shaped upstream (no body rewrite).
/// Gemini operations carried in the `{model}:{op}` path segment.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum GeminiOperation {
GenerateContent,
StreamGenerateContent,
CountTokens,
}
impl GeminiOperation {
fn as_str(self) -> &'static str {
match self {
GeminiOperation::GenerateContent => "generateContent",
GeminiOperation::StreamGenerateContent => "streamGenerateContent",
GeminiOperation::CountTokens => "countTokens",
}
}
}
/// Parse the `{model}:{operation}` path segment (e.g.
/// `gemini-2.5-pro:streamGenerateContent` or `models/gemini-x:countTokens`).
/// Splits on the last `:` so model names that contain none are rejected.
fn parse_gemini_model_action(model_action: &str) -> Option<(String, GeminiOperation)> {
let (model, action) = model_action.rsplit_once(':')?;
let operation = match action {
"generateContent" => GeminiOperation::GenerateContent,
"streamGenerateContent" => GeminiOperation::StreamGenerateContent,
"countTokens" => GeminiOperation::CountTokens,
_ => return None,
};
if model.trim().is_empty() {
return None;
}
Some((model.to_string(), operation))
}
/// Authenticate a Gemini-native request: the key arrives in `?key=` or the
/// `x-goog-api-key` header (Gemini convention), not `Authorization: Bearer`.
/// Reuses the same downstream-key authorization + limits as `check_auth`.
async fn check_gemini_auth(
state: &AppState,
headers: &HeaderMap,
query: &std::collections::HashMap<String, String>,
) -> Result<Option<i64>, AppError> {
let key = query
.get("key")
.map(|s| s.as_str())
.or_else(|| headers.get("x-goog-api-key").and_then(|v| v.to_str().ok()))
.map(str::trim)
.filter(|s| !s.is_empty())
.ok_or_else(|| {
AppError::Unauthorized("Missing API key (use ?key= or x-goog-api-key)".into())
})?;
let proxy_token = state.runtime_proxy_token();
let result = authorize_downstream_token(key, &state.db, Some(&proxy_token))
.await
.map_err(AppError::Unauthorized)?;
match result {
DownstreamAuthResult::Managed { key } => {
check_downstream_key_limits(&key)?;
Ok(Some(key.id))
}
DownstreamAuthResult::Global { .. } => Ok(None),
}
}
/// Rebuild the upstream query string from client params, dropping the inbound
/// gateway `key` (the upstream key is injected by the executor) and `alt` (the
/// streaming path enforces its own `alt=sse`). Returns `None` when nothing
/// remains. Values are forwarded as-is: real Gemini clients only put `key`/`alt`
/// in the query and carry everything else (generationConfig, etc.) in the body.
fn gemini_forward_query(
query: &std::collections::HashMap<String, String>,
drop_alt: bool,
) -> Option<String> {
let mut parts: Vec<(&String, &String)> = query
.iter()
.filter(|(k, _)| k.as_str() != "key" && !(drop_alt && k.as_str() == "alt"))
.collect();
if parts.is_empty() {
return None;
}
// Stable order so the URL is deterministic (and testable).
parts.sort_by(|a, b| a.0.cmp(b.0));
let encoded = parts
.iter()
.map(|(k, v)| format!("{k}={v}"))
.collect::<Vec<_>>()
.join("&");
Some(encoded)
}
/// GET /v1beta/models — list models in Gemini's native shape.
/// Like `/v1/models`, this is a directory listing from configured routes; no
/// upstream call, no billing.
async fn gemini_list_models(
State(state): State<AppState>,
headers: HeaderMap,
Query(query): Query<std::collections::HashMap<String, String>>,
) -> AppResult<Json<Value>> {
info!("GET /v1beta/models");
let _ = check_gemini_auth(&state, &headers, &query).await?;
let rows: Vec<(String, String)> = sqlx::query_as(
"SELECT DISTINCT
tr.model_pattern,
COALESCE(NULLIF(TRIM(tr.model_mapping), ''), tr.model_pattern)
FROM token_routes tr
JOIN route_channels rc ON rc.route_id = tr.id
JOIN accounts a ON a.id = rc.account_id
JOIN sites s ON s.id = a.site_id
LEFT JOIN account_tokens at ON at.id = rc.token_id AND at.enabled = 1
WHERE tr.enabled = 1
AND rc.enabled = 1
AND a.status = 'active'
AND s.status = 'active'
AND TRIM(COALESCE(s.url, '')) <> ''
AND (rc.cooldown_until IS NULL OR rc.cooldown_until <= datetime('now'))
AND (rc.token_id IS NULL OR at.id IS NOT NULL)
ORDER BY tr.id",
)
.fetch_all(&*state.db)
.await?;
let models: Vec<Value> = rows
.into_iter()
.map(|(_pattern, model)| {
let name = if model.starts_with("models/") {
model.clone()
} else {
format!("models/{model}")
};
json!({
"name": name,
"supportedGenerationMethods": [
"generateContent",
"streamGenerateContent",
"countTokens",
],
})
})
.collect();
Ok(Json(json!({ "models": models })))
}
/// POST /v1beta/models/{model}:{op} — dispatch to the matching Gemini op.
async fn gemini_model_operation(
State(state): State<AppState>,
axum::extract::Path(model_action): axum::extract::Path<String>,
Query(query): Query<std::collections::HashMap<String, String>>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> Result<Response, AppError> {
let Some((model, operation)) = parse_gemini_model_action(&model_action) else {
return Err(AppError::BadRequest(format!(
"Unsupported Gemini model action '{model_action}' \
(expected {{model}}:generateContent|streamGenerateContent|countTokens)"
)));
};
info!("POST /v1beta/models/{model}:{}", operation.as_str());
let key_id = check_gemini_auth(&state, &headers, &query).await?;
match operation {
GeminiOperation::GenerateContent => {
let extra_query = gemini_forward_query(&query, false);
let result =
execute_gemini_with_conductor(state, key_id, model, body, extra_query).await?;
Ok(Json(result).into_response())
}
GeminiOperation::StreamGenerateContent => {
let extra_query = gemini_forward_query(&query, true);
execute_gemini_stream_with_conductor(state, key_id, model, body, extra_query).await
}
GeminiOperation::CountTokens => {
// countTokens is a metadata call (no generation) — route for a
// channel but skip billing, mirroring the lighter helper endpoints.
let extra_query = gemini_forward_query(&query, false);
let route = resolve_route(&state, &model).await?;
let result = executors::execute_gemini_operation(
state,
route,
"countTokens",
body,
extra_query.as_deref(),
)
.await?;
Ok(Json(result).into_response())
}
}
}
/// POST /v1/chat/completions
async fn chat_completions(
State(state): State<AppState>,
@@ -1334,6 +1547,38 @@ fn responses_stream_usage_from_line(line: &str) -> Option<(i64, i64)> {
(prompt != 0 || completion != 0).then_some((prompt, completion))
}
/// Extract `(prompt, completion)` tokens from a Gemini response body's
/// `usageMetadata` (`promptTokenCount` / `candidatesTokenCount`). Returns
/// `(0, 0)` when the field is absent — mirrors `extract_usage` for OpenAI.
fn gemini_extract_usage(response: &Value) -> (i64, i64) {
let usage = response.get("usageMetadata");
let prompt = usage
.and_then(|u| u.get("promptTokenCount"))
.and_then(|v| v.as_i64())
.unwrap_or(0);
let completion = usage
.and_then(|u| u.get("candidatesTokenCount"))
.and_then(|v| v.as_i64())
.unwrap_or(0);
(prompt, completion)
}
/// Per-line usage parser for Gemini `streamGenerateContent?alt=sse` output.
/// Gemini emits `data:`-framed JSON chunks; the final chunk carries the
/// cumulative `usageMetadata`. Feeds the shared `StreamUsageAccumulator`.
fn gemini_stream_usage_from_line(line: &str) -> Option<(i64, i64)> {
let payload = line
.strip_prefix("data: ")
.or_else(|| line.strip_prefix("data:"))?
.trim();
if payload.is_empty() || payload == "[DONE]" {
return None;
}
let value: Value = serde_json::from_str(payload).ok()?;
let (prompt, completion) = gemini_extract_usage(&value);
(prompt != 0 || completion != 0).then_some((prompt, completion))
}
fn build_conductor(state: &AppState) -> ProxyConductor {
let db = state.db.clone();
let select_db = db.clone();
@@ -1881,6 +2126,185 @@ async fn execute_responses_stream_with_conductor(
}
}
/// Non-streaming Gemini op (`generateContent`) through the ProxyConductor.
/// Mirrors `execute_responses_with_conductor` but builds a Gemini upstream URL
/// (model in path, key in query) and forwards the body verbatim — Gemini bodies
/// carry no `model` field, so there is no body rewrite.
async fn execute_gemini_with_conductor(
state: AppState,
key_id: Option<i64>,
model: String,
body: Value,
extra_query: Option<String>,
) -> AppResult<Value> {
let started_at = Instant::now();
let request_log_body = body.clone();
let conductor = build_conductor(&state);
let billing_state = state.clone();
let billing_model = model.clone();
let result = conductor
.execute(ExecuteInput {
requested_model: model,
downstream_policy: None,
attempt_fn: Box::new(move |ctx| {
let state = state.clone();
let body = body.clone();
let extra_query = extra_query.clone();
async move {
match executors::execute_gemini_operation(
state,
selected_channel_to_route(ctx.selected),
"generateContent",
body,
extra_query.as_deref(),
)
.await
{
Ok(response) => AttemptResult::Success {
response,
latency_ms: None,
cost: None,
},
Err(err) => AttemptResult::Failure {
action: AttemptFailureAction::Failover,
status: upstream_status_from_error(&err),
raw_error_text: Some(err.to_string()),
},
}
}
.boxed()
}),
on_terminal_failure: None,
})
.await;
let (selected, response, attempts) = match result {
ExecuteResult::Success {
selected,
response,
attempts,
} => (selected, response, attempts),
ExecuteResult::Failure {
reason,
selected,
status,
raw_error_text,
attempts,
} => {
record_proxy_log_failure(
&billing_state,
key_id,
selected.as_ref(),
&billing_model,
status,
Some(elapsed_ms_since(started_at)),
raw_error_text.as_deref(),
attempts,
false,
proxy_log_content_from_request(&request_log_body),
)
.await;
return Err(AppError::UpstreamError(format!(
"Proxy conductor failed after {attempts} attempts ({reason:?}, status: {:?}): {}",
status,
raw_error_text.unwrap_or_else(|| "no upstream error text".to_string())
)));
}
};
let (prompt, completion) = gemini_extract_usage(&response);
let cost = estimate_request_cost(&selected, prompt, completion);
record_proxy_log_success(
&billing_state,
key_id,
&selected,
&billing_model,
Some(elapsed_ms_since(started_at)),
prompt,
completion,
cost,
attempts,
false,
proxy_log_content_from_request_response(&request_log_body, Some(&response)),
)
.await;
record_billing(
&billing_state,
key_id,
&billing_model,
prompt,
completion,
cost,
)
.await;
Ok(response)
}
/// Streaming Gemini op (`streamGenerateContent`) through the ProxyConductor.
/// Mirrors `execute_responses_stream_with_conductor`, opening the upstream with
/// failover and SSE-passing it through with the Gemini usage parser so streamed
/// `usageMetadata` is billed.
async fn execute_gemini_stream_with_conductor(
state: AppState,
key_id: Option<i64>,
model: String,
body: Value,
extra_query: Option<String>,
) -> Result<Response, AppError> {
let started_at = Instant::now();
let request_log_body = body.clone();
let body_for_attempt = body.clone();
let opened = open_stream_with_failover(&state, &model, move |state, selected| {
let body = body_for_attempt.clone();
let extra_query = extra_query.clone();
async move {
executors::execute_gemini_stream_response(
state,
selected_channel_to_route(selected),
"streamGenerateContent",
body,
extra_query.as_deref(),
)
.await
}
})
.await;
match opened {
Ok(opened) => sse_passthrough_response(
state,
opened,
key_id,
model,
started_at,
proxy_log_content_from_request(&request_log_body),
gemini_stream_usage_from_line,
),
Err(failure) => {
record_proxy_log_failure(
&state,
key_id,
failure.selected.as_ref(),
&model,
failure.status,
Some(elapsed_ms_since(started_at)),
failure.raw_error_text.as_deref(),
failure.attempts,
true,
proxy_log_content_from_request(&request_log_body),
)
.await;
Err(AppError::UpstreamError(format!(
"Proxy conductor failed after {} attempts (stream open, status: {:?}): {}",
failure.attempts,
failure.status,
failure
.raw_error_text
.unwrap_or_else(|| "no upstream error text".to_string())
)))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
@@ -2121,6 +2545,8 @@ mod tests {
template_proxy_url: None,
},
comfyui: Default::default(),
payment: Default::default(),
pixelle_video: Default::default(),
},
pool,
crate::services::builtin_keys::BuiltinKeys {
@@ -2155,4 +2581,87 @@ mod tests {
.unwrap();
assert!(fail_count > 0);
}
// ── Gemini native inbound unit tests ───────────────────────────────
#[test]
fn parse_gemini_model_action_known_ops() {
let (m, op) = parse_gemini_model_action("gemini-2.5-pro:generateContent").unwrap();
assert_eq!(m, "gemini-2.5-pro");
assert_eq!(op, GeminiOperation::GenerateContent);
let (m, op) = parse_gemini_model_action("gemini-2.5-pro:streamGenerateContent").unwrap();
assert_eq!(m, "gemini-2.5-pro");
assert_eq!(op, GeminiOperation::StreamGenerateContent);
// model names may carry the `models/` prefix
let (m, op) = parse_gemini_model_action("models/gemini-1.5-pro:countTokens").unwrap();
assert_eq!(m, "models/gemini-1.5-pro");
assert_eq!(op, GeminiOperation::CountTokens);
}
#[test]
fn parse_gemini_model_action_rejects_bad_input() {
// no operation suffix
assert!(parse_gemini_model_action("gemini-2.5-pro").is_none());
// unknown operation
assert!(parse_gemini_model_action("gemini-2.5-pro:embedContent").is_none());
// empty model
assert!(parse_gemini_model_action(":generateContent").is_none());
}
#[test]
fn gemini_extract_usage_reads_usage_metadata() {
let body = json!({
"candidates": [],
"usageMetadata": {
"promptTokenCount": 11,
"candidatesTokenCount": 22,
"totalTokenCount": 33
}
});
assert_eq!(gemini_extract_usage(&body), (11, 22));
// absent usageMetadata → zeros
assert_eq!(gemini_extract_usage(&json!({})), (0, 0));
}
#[test]
fn gemini_stream_usage_from_line_parses_data_frame() {
let line = r#"data: {"candidates":[{"content":{"parts":[]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":6,"totalTokenCount":10}}"#;
assert_eq!(gemini_stream_usage_from_line(line), Some((4, 6)));
}
#[test]
fn gemini_stream_usage_from_line_ignores_non_usage_and_done() {
// a chunk with no usageMetadata yields None (don't overwrite accumulator)
let no_usage = r#"data: {"candidates":[{"content":{"parts":[{"text":"hi"}]}}]}"#;
assert_eq!(gemini_stream_usage_from_line(no_usage), None);
// sentinels / blanks / non-data lines
assert_eq!(gemini_stream_usage_from_line("data: [DONE]"), None);
assert_eq!(gemini_stream_usage_from_line("data:"), None);
assert_eq!(gemini_stream_usage_from_line(": comment"), None);
}
#[test]
fn gemini_forward_query_drops_key_and_alt() {
let mut q = std::collections::HashMap::new();
q.insert("key".to_string(), "secret".to_string());
q.insert("alt".to_string(), "sse".to_string());
// stream path drops both key and alt → nothing left
assert_eq!(gemini_forward_query(&q, true), None);
// non-stream keeps alt but always drops key
assert_eq!(gemini_forward_query(&q, false), Some("alt=sse".to_string()));
}
#[test]
fn gemini_forward_query_preserves_other_params_sorted() {
let mut q = std::collections::HashMap::new();
q.insert("key".to_string(), "secret".to_string());
q.insert("pageSize".to_string(), "50".to_string());
q.insert("filter".to_string(), "x".to_string());
assert_eq!(
gemini_forward_query(&q, true),
Some("filter=x&pageSize=50".to_string())
);
}
}