fix(gateway): harden Gemini endpoint routing

This commit is contained in:
MMEXA
2026-05-18 00:53:34 +00:00
parent 81ff375bfd
commit b004a02e4a
20 changed files with 1348 additions and 83 deletions

View File

@@ -18,6 +18,7 @@ use crate::ai_serving::{
};
use crate::clock::current_unix_ms;
use crate::execution_runtime;
use crate::handlers::admin::provider::write::provider::reconcile_admin_fixed_provider_template_endpoints;
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::handlers::shared::provider_pool::{
admin_provider_pool_config_from_config_value, read_admin_provider_pool_runtime_state,
@@ -555,7 +556,6 @@ fn provider_query_build_test_request_body_with_model_policy(
"content": provider_query_extract_message(payload)
.unwrap_or_else(|| DEFAULT_PROVIDER_QUERY_TEST_MESSAGE.to_string())
}],
"max_tokens": 30,
"temperature": 0.7,
"stream": true,
})
@@ -1022,6 +1022,8 @@ async fn provider_query_build_kiro_test_candidates(
payload: &Value,
requested_model_override: Option<&str>,
) -> Result<Vec<ProviderQueryTestCandidate>, Response<Body>> {
provider_query_reconcile_fixed_provider_endpoints_for_test_model(state, provider).await?;
let provider_ids = vec![provider.id.clone()];
let endpoints = state
.app()
@@ -1227,6 +1229,33 @@ async fn provider_query_build_kiro_test_candidates(
Ok(candidates)
}
async fn provider_query_reconcile_fixed_provider_endpoints_for_test_model(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
) -> Result<(), Response<Body>> {
if state
.fixed_provider_template(&provider.provider_type)
.is_none()
|| !state.has_provider_catalog_data_writer()
{
return Ok(());
}
reconcile_admin_fixed_provider_template_endpoints(state, provider)
.await
.map_err(|err| {
warn!(
provider_id = %provider.id,
provider_type = %provider.provider_type,
error = ?err,
"admin provider-query test-model: failed to reconcile fixed provider endpoints"
);
build_admin_provider_query_bad_request_response(
ADMIN_PROVIDER_QUERY_NO_ACTIVE_API_KEY_DETAIL,
)
})
}
fn provider_query_decode_execution_body(
result: &aether_contracts::ExecutionResult,
) -> Option<Vec<u8>> {
@@ -1256,7 +1285,7 @@ fn provider_query_standard_execution_response_body(
provider_api_format: &str,
result: &aether_contracts::ExecutionResult,
) -> Option<Value> {
result
let body = result
.body
.as_ref()
.and_then(|body| body.json_body.clone())
@@ -1264,7 +1293,15 @@ fn provider_query_standard_execution_response_body(
provider_query_decode_execution_body(result).and_then(|body| {
provider_query_aggregate_standard_stream_sync_response(provider_api_format, &body)
})
})
})?;
if result.status_code < 400
&& provider_query_normalize_api_format_alias(provider_api_format)
== "gemini:generate_content"
&& aether_ai_formats::formats::gemini::generate_content::response::from_raw(&body).is_none()
{
return None;
}
Some(body)
}
fn provider_query_extract_error_message(

View File

@@ -28,6 +28,17 @@ fn provider_query_test_request_body_defaults_missing_model() {
assert_eq!(body["model"], json!("fallback-model"));
}
#[test]
fn provider_query_default_test_request_body_does_not_set_max_tokens() {
let body = provider_query_build_test_request_body(&json!({}), "fallback-model");
assert_eq!(body["model"], json!("fallback-model"));
assert!(
body.get("max_tokens").is_none(),
"admin model test must not silently force a low max_tokens value"
);
}
#[test]
fn provider_query_failover_request_body_overrides_custom_model() {
let payload = json!({
@@ -168,6 +179,40 @@ fn provider_query_standard_test_aggregates_responses_stream_body() {
assert_eq!(body["output"][0]["content"][0]["text"], json!("Hello"));
}
#[test]
fn provider_query_standard_test_rejects_gemini_success_without_visible_output() {
let result = aether_contracts::ExecutionResult {
request_id: "provider-test".to_string(),
candidate_id: Some("candidate-0".to_string()),
status_code: 200,
headers: BTreeMap::new(),
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"candidates": [{
"content": {"role": "model"},
"finishReason": "MAX_TOKENS"
}],
"usageMetadata": {
"promptTokenCount": 8,
"candidatesTokenCount": 1,
"thoughtsTokenCount": 25,
"totalTokenCount": 34
},
"modelVersion": "gemini-3-flash-preview",
"responseId": "resp-empty"
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
};
assert!(
provider_query_standard_execution_response_body("gemini:generate_content", &result)
.is_none()
);
}
#[test]
fn provider_query_test_adapter_routes_fixed_provider_endpoint_types() {
assert_eq!(

View File

@@ -179,9 +179,6 @@ pub(super) async fn maybe_build_local_test_connection_route_response(
"role": "user",
"parts": [{"text": "Health check"}],
}],
"generationConfig": {
"maxOutputTokens": 5,
},
}),
_ => return None,
};