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
parent 24a8dec151
commit aaa593ac31
1 changed files with 212 additions and 0 deletions
@@ -1159,3 +1159,215 @@ pub async fn execute_image_edit(

    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");
+    }
+}