fix: route Gemini CLI generateContent through stream

This commit is contained in:
Mas0nShi
2026-05-22 09:30:55 +08:00
parent 66f21de50e
commit 9533bd7043
3 changed files with 38 additions and 31 deletions
@@ -1273,7 +1273,7 @@ async fn gateway_executes_gemini_cli_sync_via_local_decision_gate_after_oauth_re
);
assert_eq!(
seen_execution_runtime_request.url,
"https://cloudcode-pa.googleapis.com/v1internal:generateContent"
"https://cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse"
);
assert!(seen_execution_runtime_request.has_model_field);
assert_eq!(
@@ -4773,9 +4773,10 @@ async fn gateway_handles_gemini_cli_test_model_with_oauth_header_fallback() {
assert_eq!(plan.endpoint_id, "endpoint-gemini-cli");
assert_eq!(plan.key_id, "key-gemini-cli");
assert_eq!(plan.provider_api_format, "gemini:generate_content");
assert!(plan.stream);
assert_eq!(
plan.url,
"https://cloudcode-pa.googleapis.com/v1internal:generateContent"
"https://cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse"
);
assert_eq!(
plan.body.json_body.as_ref().unwrap()["project"],
@@ -4929,8 +4930,9 @@ async fn gateway_hydrates_gemini_cli_project_id_from_load_code_assist_for_test_m
assert_eq!(
plan.url,
"https://cloudcode-pa.googleapis.com/v1internal:generateContent"
"https://cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse"
);
assert!(plan.stream);
assert_eq!(
plan.body.json_body.as_ref().unwrap()["project"],
json!("project-from-load-code-assist")
@@ -5030,7 +5032,8 @@ async fn gateway_hydrates_gemini_cli_project_id_from_load_code_assist_for_test_m
*seen_urls.lock().expect("mutex should lock"),
vec![
"https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist".to_string(),
"https://cloudcode-pa.googleapis.com/v1internal:generateContent".to_string(),
"https://cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse"
.to_string(),
]
);
let reloaded = provider_catalog_repository
@@ -5239,9 +5242,10 @@ async fn gateway_unwraps_gemini_cli_v1internal_response_for_failover_model_test(
assert_eq!(plan.endpoint_id, "endpoint-gemini-cli");
assert_eq!(plan.key_id, "key-gemini-cli");
assert_eq!(plan.provider_api_format, "gemini:generate_content");
assert!(plan.stream);
assert_eq!(
plan.url,
"https://cloudcode-pa.googleapis.com/v1internal:generateContent"
"https://cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse"
);
assert_eq!(
plan.body.json_body.as_ref().unwrap()["project"],
@@ -5256,32 +5260,15 @@ async fn gateway_unwraps_gemini_cli_v1internal_response_for_failover_model_test(
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
"content-type": "text/event-stream"
},
"body": {
"json_body": {
"response": {
"candidates": [{
"content": {
"parts": [{
"text": "Gemini CLI v1internal failover response"
}],
"role": "model"
},
"finishReason": "STOP",
"index": 0
}],
"modelVersion": "gemini-3-flash-preview",
"usageMetadata": {
"promptTokenCount": 2,
"candidatesTokenCount": 5,
"totalTokenCount": 7
}
},
"remainingCredits": 123,
"consumedCredits": 1,
"traceId": "trace-gemini-cli-1"
}
"body_bytes_b64": base64::engine::general_purpose::STANDARD.encode(
concat!(
"data: {\"response\":{\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"Gemini CLI v1internal failover response\"}],\"role\":\"model\"},\"finishReason\":\"STOP\",\"index\":0}],\"modelVersion\":\"gemini-3-flash-preview\",\"usageMetadata\":{\"promptTokenCount\":2,\"candidatesTokenCount\":5,\"totalTokenCount\":7}},\"remainingCredits\":123,\"consumedCredits\":1,\"traceId\":\"trace-gemini-cli-1\"}\n\n"
)
.as_bytes()
)
},
"telemetry": {
"elapsed_ms": 23
@@ -34,8 +34,12 @@ pub fn force_upstream_streaming_for_provider(
provider_type: &str,
provider_api_format: &str,
) -> bool {
provider_type.trim().eq_ignore_ascii_case("codex")
&& aether_ai_formats::is_openai_responses_format(provider_api_format)
let provider_type = provider_type.trim();
(provider_type.eq_ignore_ascii_case("codex")
&& aether_ai_formats::is_openai_responses_format(provider_api_format))
|| (provider_type.eq_ignore_ascii_case("gemini_cli")
&& aether_ai_formats::normalize_api_format_alias(provider_api_format)
== "gemini:generate_content")
}
pub(crate) fn parse_upstream_stream_policy(
@@ -196,6 +200,22 @@ mod tests {
));
}
#[test]
fn forces_streaming_for_gemini_cli_generate_content() {
assert!(force_upstream_streaming_for_provider(
"gemini_cli",
"gemini:generate_content"
));
assert!(force_upstream_streaming_for_provider(
" GEMINI_CLI ",
"gemini:generate_content"
));
assert!(!force_upstream_streaming_for_provider(
"google",
"gemini:generate_content"
));
}
#[test]
fn does_not_force_streaming_for_compact_or_other_provider_types() {
assert!(!force_upstream_streaming_for_provider(