fix: unwrap Gemini CLI model test envelopes

This commit is contained in:
Mas0nShi
2026-05-21 17:39:10 +08:00
parent 3e6ce6cf4a
commit 66f21de50e
3 changed files with 158 additions and 8 deletions
@@ -1515,12 +1515,18 @@ fn provider_query_aggregate_standard_stream_sync_response(
fn provider_query_standard_execution_response_body(
provider_api_format: &str,
result: &aether_contracts::ExecutionResult,
report_context: Option<&Value>,
) -> Option<Value> {
let body = provider_query_execution_json_body(result).or_else(|| {
provider_query_decode_execution_body(result).and_then(|body| {
provider_query_aggregate_standard_stream_sync_response(provider_api_format, &body)
})
})?;
let body = report_context
.and_then(|context| {
crate::ai_serving::api::normalize_provider_private_response_value(body.clone(), context)
})
.unwrap_or(body);
if result.status_code < 400
&& provider_query_normalize_api_format_alias(provider_api_format)
== "gemini:generate_content"
@@ -2876,6 +2882,16 @@ async fn provider_query_execute_standard_test_candidate(
}
};
}
let private_report_context =
(crate::provider_transport::is_gemini_cli_provider_transport(&transport)
&& normalized_provider_api_format == "gemini:generate_content")
.then(|| {
json!({
"has_envelope": true,
"envelope_name": crate::provider_transport::GEMINI_CLI_V1INTERNAL_ENVELOPE_NAME,
"provider_api_format": provider_api_format,
})
});
let uses_vertex_query_auth =
crate::provider_transport::uses_vertex_api_key_query_auth(&transport, provider_api_format);
@@ -3080,7 +3096,11 @@ async fn provider_query_execute_standard_test_candidate(
.execute_execution_runtime_sync_plan(Some(trace_id), &plan)
.await?;
let response_body = if result.status_code < 400 {
provider_query_standard_execution_response_body(provider_api_format, &result)
provider_query_standard_execution_response_body(
provider_api_format,
&result,
private_report_context.as_ref(),
)
} else {
result.body.as_ref().and_then(|body| body.json_body.clone())
};
@@ -124,7 +124,7 @@ fn provider_query_execution_json_body_decodes_stream_encoded_json_response() {
Some(body.clone())
);
assert_eq!(
provider_query_standard_execution_response_body("openai:image", &result),
provider_query_standard_execution_response_body("openai:image", &result, None),
Some(body)
);
}
@@ -308,7 +308,7 @@ fn provider_query_standard_test_aggregates_responses_stream_body() {
error: None,
};
let body = provider_query_standard_execution_response_body("openai:responses", &result)
let body = provider_query_standard_execution_response_body("openai:responses", &result, None)
.expect("stream body should aggregate");
assert_eq!(body["model"], json!("gpt-5.4-mini"));
@@ -340,7 +340,7 @@ fn provider_query_standard_test_aggregates_responses_image_generation_call() {
error: None,
};
let body = provider_query_standard_execution_response_body("openai:responses", &result)
let body = provider_query_standard_execution_response_body("openai:responses", &result, None)
.expect("responses image stream body should aggregate");
assert_eq!(body["output"][0]["type"], json!("image_generation_call"));
@@ -502,10 +502,12 @@ fn provider_query_standard_test_rejects_gemini_success_without_visible_output()
error: None,
};
assert!(
provider_query_standard_execution_response_body("gemini:generate_content", &result)
.is_none()
);
assert!(provider_query_standard_execution_response_body(
"gemini:generate_content",
&result,
None
)
.is_none());
}
#[test]
@@ -5230,6 +5230,134 @@ async fn gateway_handles_gemini_cli_test_model_failover_locally() {
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_unwraps_gemini_cli_v1internal_response_for_failover_model_test() {
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |Json(plan): Json<ExecutionPlan>| async move {
assert_eq!(plan.provider_id, "provider-gemini-cli");
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_eq!(
plan.url,
"https://cloudcode-pa.googleapis.com/v1internal:generateContent"
);
assert_eq!(
plan.body.json_body.as_ref().unwrap()["project"],
json!("project-1")
);
assert_eq!(
plan.body.json_body.as_ref().unwrap()["model"],
json!("gemini-3-flash-preview")
);
Json(json!({
"request_id": plan.request_id,
"candidate_id": plan.candidate_id,
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"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"
}
},
"telemetry": {
"elapsed_ms": 23
}
}))
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let mut provider = sample_provider("provider-gemini-cli", "Gemini CLI", 10);
provider.provider_type = "gemini_cli".to_string();
let mut key = sample_key(
"key-gemini-cli",
"provider-gemini-cli",
"gemini:generate_content",
"cached-gemini-cli-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
aether_crypto::encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"gemini_cli","project_id":"project-1"}"#,
)
.expect("auth config should encrypt"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_endpoint(
"endpoint-gemini-cli",
"provider-gemini-cli",
"gemini:generate_content",
"https://cloudcode-pa.googleapis.com",
)],
vec![key],
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
provider_catalog_repository,
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
)),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-query/test-model-failover"
))
.header(GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"provider_id": "provider-gemini-cli",
"failover_models": ["gemini-3-flash-preview"],
"api_format": "gemini:generate_content"
}))
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], json!(true));
assert_eq!(payload["total_attempts"], json!(1));
assert_eq!(
payload["data"]["response"]["candidates"][0]["content"]["parts"][0]["text"],
json!("Gemini CLI v1internal failover response")
);
assert!(payload["data"]["response"].get("response").is_none());
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_provider_query_test_model_failover_with_single_model_name_alias() {
let execution_runtime = Router::new().route(