mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
fix(gateway): harden Gemini endpoint routing
This commit is contained in:
@@ -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(
|
||||
|
||||
@@ -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!(
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user