3 Commits

Author SHA1 Message Date
boos 4d093548de 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).
2026-06-22 08:39:38 +08:00
boos aaa593ac31 feat(proxy): add Gemini native upstream executors
Add Gemini-shaped upstream calls that mirror AQBot's native passthrough
and BoosAPI's /v1/responses forwarding: the model lives in the URL path
(/v1beta/models/{model}:{op}) and the key in ?key=, so the body is sent
verbatim with no model rewrite.

- gemini_base_url / gemini_operation_url normalize the v1beta root and
  build {base}/v1beta/models/{model}:{op}?key=...&{extra}, stripping a
  leading models/ prefix and handling the no-key case.
- execute_gemini_operation runs non-streaming ops (generateContent /
  countTokens) and returns upstream JSON verbatim.
- execute_gemini_stream_response opens streamGenerateContent for SSE
  passthrough, enforcing alt=sse via merge_alt_sse.

Unit tests cover URL building and alt=sse merging.
2026-06-22 08:39:26 +08:00
boos 24a8dec151 fix: define DEFAULT_NEW_USER_MAX_COST_USD and repair test AppConfig initializers
Pre-existing breakage on this branch (introduced by the billing commits),
unrelated to the Gemini work but blocking both `cargo build` and
`cargo test`:

- admin/users.rs imported routes::auth::DEFAULT_NEW_USER_MAX_COST_USD,
  which was never defined. Add it to auth.rs as $1.0 (the billing design's
  "new users default to $1" starter allowance).
- AppConfig gained `payment` and `pixelle_video` fields but several test
  helpers still used struct literals missing them. Initialize both with
  Default::default() at each affected test site.
2026-06-22 08:39:17 +08:00
11 changed files with 743 additions and 0 deletions
+212
View File
@@ -1159,3 +1159,215 @@ pub async fn execute_image_edit(
Ok(json) Ok(json)
} }
// ── Gemini native passthrough ───────────────────────────────────────────
//
// Gemini's native protocol carries the model in the URL path
// (`/v1beta/models/{model}:{op}`) and the API key in the `?key=` query, so the
// request body needs no rewrite — we forward it verbatim. This mirrors AQBot's
// `handle_native_request`, which deliberately skips `rewrite_body_model` for
// Gemini, and BoosAPI's own `/v1/responses` passthrough.
/// Normalize an upstream base into a Gemini `v1beta` root. Accepts bases that
/// already end in `/v1beta` (or contain `/v1beta/`) and leaves them as-is;
/// otherwise appends `/v1beta`. Trailing slashes are trimmed.
fn gemini_base_url(base_url: &str) -> String {
let trimmed = base_url.trim().trim_end_matches('/');
if trimmed.ends_with("/v1beta") || trimmed.contains("/v1beta/") {
trimmed.to_string()
} else {
format!("{trimmed}/v1beta")
}
}
/// Build `{base}/v1beta/models/{model}:{op}?key={api_key}[&{extra_query}]`.
/// `extra_query` carries client query params already filtered (e.g. `alt=sse`),
/// without a leading `?`/`&`.
fn gemini_operation_url(
base_url: &str,
model: &str,
operation: &str,
api_key: &str,
extra_query: Option<&str>,
) -> String {
let base = gemini_base_url(base_url);
let model = model.trim().trim_start_matches('/');
let model = model.strip_prefix("models/").unwrap_or(model);
let mut url = format!("{base}/models/{model}:{operation}");
let mut sep = '?';
if !api_key.is_empty() {
url.push(sep);
url.push_str("key=");
url.push_str(api_key);
sep = '&';
}
if let Some(query) = extra_query.filter(|q| !q.is_empty()) {
url.push(sep);
url.push_str(query);
}
url
}
/// Execute a non-streaming Gemini operation (`generateContent` / `countTokens`)
/// and return the upstream JSON verbatim.
pub async fn execute_gemini_operation(
state: AppState,
route: UpstreamRoute,
operation: &str,
body: Value,
extra_query: Option<&str>,
) -> AppResult<Value> {
let client = state.http.clone();
let upstream_url =
gemini_operation_url(&route.base_url, &route.model, operation, &route.api_key, extra_query);
let resp = client
.post(&upstream_url)
.header("Content-Type", "application/json")
.json(&body)
.send()
.await
.map_err(|e| AppError::UpstreamError(format!("Upstream request failed: {e}")))?;
let status = resp.status();
let text = resp
.text()
.await
.map_err(|e| AppError::UpstreamError(format!("Failed to read upstream response: {e}")))?;
if !status.is_success() {
return Err(AppError::UpstreamError(format!(
"Upstream returned {status}: {text}"
)));
}
parse_upstream_json(&text, &upstream_url, status)
}
/// Execute a streaming Gemini operation (`streamGenerateContent`) and return the
/// live upstream response for SSE passthrough. `alt=sse` is enforced in the
/// query so the upstream emits `data:`-framed events the accumulator can parse.
pub async fn execute_gemini_stream_response(
state: AppState,
route: UpstreamRoute,
operation: &str,
body: Value,
extra_query: Option<&str>,
) -> AppResult<reqwest::Response> {
let client = state.http.clone();
let query = merge_alt_sse(extra_query);
let upstream_url = gemini_operation_url(
&route.base_url,
&route.model,
operation,
&route.api_key,
Some(&query),
);
let resp = client
.post(&upstream_url)
.header("Content-Type", "application/json")
.json(&body)
.send()
.await
.map_err(|e| AppError::UpstreamError(format!("Upstream request failed: {e}")))?;
let status = resp.status();
if !status.is_success() {
let text = resp
.text()
.await
.unwrap_or_else(|_| "unknown error".to_string());
return Err(AppError::UpstreamError(format!(
"Upstream returned {status}: {text}"
)));
}
Ok(resp)
}
/// Ensure `alt=sse` is present in the extra query (Gemini streams plain JSON
/// arrays without it). Preserves any other client-supplied params.
fn merge_alt_sse(extra_query: Option<&str>) -> String {
match extra_query.filter(|q| !q.is_empty()) {
Some(q) if q.split('&').any(|p| p == "alt=sse" || p.starts_with("alt=")) => q.to_string(),
Some(q) => format!("{q}&alt=sse"),
None => "alt=sse".to_string(),
}
}
#[cfg(test)]
mod gemini_tests {
use super::*;
#[test]
fn base_url_appends_v1beta_when_missing() {
assert_eq!(
gemini_base_url("https://generativelanguage.googleapis.com"),
"https://generativelanguage.googleapis.com/v1beta"
);
}
#[test]
fn base_url_preserves_existing_v1beta() {
assert_eq!(
gemini_base_url("https://generativelanguage.googleapis.com/v1beta/"),
"https://generativelanguage.googleapis.com/v1beta"
);
}
#[test]
fn operation_url_embeds_model_op_and_key() {
let url = gemini_operation_url(
"https://generativelanguage.googleapis.com",
"gemini-2.0-flash",
"generateContent",
"secret",
None,
);
assert_eq!(
url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key=secret"
);
}
#[test]
fn operation_url_strips_models_prefix_and_appends_query() {
let url = gemini_operation_url(
"https://generativelanguage.googleapis.com/v1beta",
"models/gemini-1.5-pro",
"streamGenerateContent",
"k",
Some("alt=sse"),
);
assert_eq!(
url,
"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-pro:streamGenerateContent?key=k&alt=sse"
);
}
#[test]
fn operation_url_without_key_uses_question_mark_for_query() {
let url = gemini_operation_url(
"https://example.test/v1beta",
"gemini-x",
"countTokens",
"",
Some("foo=bar"),
);
assert_eq!(
url,
"https://example.test/v1beta/models/gemini-x:countTokens?foo=bar"
);
}
#[test]
fn merge_alt_sse_variants() {
assert_eq!(merge_alt_sse(None), "alt=sse");
assert_eq!(merge_alt_sse(Some("")), "alt=sse");
assert_eq!(merge_alt_sse(Some("foo=bar")), "foo=bar&alt=sse");
// already present → unchanged
assert_eq!(merge_alt_sse(Some("alt=sse")), "alt=sse");
assert_eq!(merge_alt_sse(Some("alt=json")), "alt=json");
}
}
+2
View File
@@ -2840,6 +2840,8 @@ mod tests {
template_proxy_url: None, template_proxy_url: None,
}, },
comfyui: Default::default(), comfyui: Default::default(),
payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
BuiltinKeys { BuiltinKeys {
+2
View File
@@ -412,6 +412,8 @@ mod tests {
template_proxy_url: None, template_proxy_url: None,
}, },
comfyui: Default::default(), comfyui: Default::default(),
payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
BuiltinKeys { BuiltinKeys {
+6
View File
@@ -19,6 +19,11 @@ use serde_json::{json, Value};
use std::net::{IpAddr, SocketAddr}; use std::net::{IpAddr, SocketAddr};
use tracing::{info, warn}; use tracing::{info, warn};
/// Default spending cap (USD) granted to a newly created user's auto-generated
/// API key. Per the billing design (新用户默认 $1 额度), new users start with a
/// $1 allowance on their default downstream key.
pub const DEFAULT_NEW_USER_MAX_COST_USD: f64 = 1.0;
// ── Request / Response types ───────────────────────────────────────── // ── Request / Response types ─────────────────────────────────────────
#[derive(Debug, Deserialize)] #[derive(Debug, Deserialize)]
@@ -1019,6 +1024,7 @@ mod tests {
}, },
comfyui: Default::default(), comfyui: Default::default(),
payment: Default::default(), payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
BuiltinKeys { BuiltinKeys {
+2
View File
@@ -127,6 +127,8 @@ mod tests {
template_proxy_url: None, template_proxy_url: None,
}, },
comfyui: ComfyUIConfig::default(), comfyui: ComfyUIConfig::default(),
payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
BuiltinKeys { BuiltinKeys {
+2
View File
@@ -1305,6 +1305,8 @@ mod tests {
template_proxy_url: None, template_proxy_url: None,
}, },
comfyui: Default::default(), comfyui: Default::default(),
payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
BuiltinKeys { BuiltinKeys {
+2
View File
@@ -536,6 +536,8 @@ mod tests {
template_proxy_url: None, template_proxy_url: None,
}, },
comfyui: Default::default(), comfyui: Default::default(),
payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
BuiltinKeys { BuiltinKeys {
+2
View File
@@ -2071,6 +2071,8 @@ mod tests {
template_proxy_url: None, template_proxy_url: None,
}, },
comfyui, comfyui,
payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
BuiltinKeys { BuiltinKeys {
+2
View File
@@ -2550,6 +2550,8 @@ mod tests {
template_proxy_url: None, template_proxy_url: None,
}, },
comfyui: ComfyUIConfig::default(), comfyui: ComfyUIConfig::default(),
payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
BuiltinKeys { BuiltinKeys {
+509
View File
@@ -73,6 +73,16 @@ pub fn routes() -> Router<AppState> {
.route("/v1/videos", post(videos_handler)) .route("/v1/videos", post(videos_handler))
.route("/v1/videos/{id}", get(video_task_handler)) .route("/v1/videos/{id}", get(video_task_handler))
.route("/v1/models", get(list_models)) .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 shim aliases ──────────────────────────────────────
// ComfyUI's OpenAI nodes hit /proxy/openai/... paths (see // ComfyUI's OpenAI nodes hit /proxy/openai/... paths (see
// comfy_api_nodes/nodes_openai.py). Mounting them here as aliases // 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 /// POST /v1/chat/completions
async fn chat_completions( async fn chat_completions(
State(state): State<AppState>, 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)) (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 { fn build_conductor(state: &AppState) -> ProxyConductor {
let db = state.db.clone(); let db = state.db.clone();
let select_db = 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)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
@@ -2121,6 +2545,8 @@ mod tests {
template_proxy_url: None, template_proxy_url: None,
}, },
comfyui: Default::default(), comfyui: Default::default(),
payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
crate::services::builtin_keys::BuiltinKeys { crate::services::builtin_keys::BuiltinKeys {
@@ -2155,4 +2581,87 @@ mod tests {
.unwrap(); .unwrap();
assert!(fail_count > 0); 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())
);
}
} }
+2
View File
@@ -696,6 +696,8 @@ mod tests {
template_proxy_url: None, template_proxy_url: None,
}, },
comfyui: ComfyUIConfig::default(), comfyui: ComfyUIConfig::default(),
payment: Default::default(),
pixelle_video: Default::default(),
}, },
pool, pool,
BuiltinKeys { BuiltinKeys {