mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
refactor gateway orchestration and failover effects
This commit is contained in:
@@ -177,13 +177,11 @@ fn usage_runtime_paths_depend_on_shared_crates_not_app_runtime_shims() {
|
||||
"is_local_ai_sync_report_kind",
|
||||
"is_local_ai_stream_report_kind",
|
||||
"sync_report_represents_failure",
|
||||
"extract_gemini_file_mapping_entries",
|
||||
"gemini_file_mapping_cache_key",
|
||||
"normalize_gemini_file_name",
|
||||
"report_request_id",
|
||||
"should_handle_local_sync_report",
|
||||
"should_handle_local_stream_report",
|
||||
"GEMINI_FILE_MAPPING_TTL_SECONDS",
|
||||
"apply_local_report_effect",
|
||||
"LocalReportEffect",
|
||||
] {
|
||||
assert!(
|
||||
usage_reporting_mod.contains(pattern),
|
||||
@@ -206,12 +204,39 @@ fn usage_runtime_paths_depend_on_shared_crates_not_app_runtime_shims() {
|
||||
"fn should_handle_local_sync_report(",
|
||||
"fn should_handle_local_stream_report(",
|
||||
"\"openai_video_delete_sync_success\" && payload.status_code == 404",
|
||||
"sync_codex_quota_from_response_headers(",
|
||||
"apply_local_gemini_file_mapping_report_effect(",
|
||||
"pub(crate) async fn store_local_gemini_file_mapping(",
|
||||
] {
|
||||
assert!(
|
||||
!usage_reporting_mod.contains(pattern),
|
||||
"usage/reporting/mod.rs should not own local report classification logic {pattern}"
|
||||
);
|
||||
}
|
||||
|
||||
let report_effects =
|
||||
read_workspace_file("apps/aether-gateway/src/orchestration/report_effects.rs");
|
||||
assert!(
|
||||
report_effects.contains("aether_usage_runtime"),
|
||||
"orchestration/report_effects.rs should depend on aether_usage_runtime"
|
||||
);
|
||||
for pattern in [
|
||||
"extract_gemini_file_mapping_entries",
|
||||
"gemini_file_mapping_cache_key",
|
||||
"normalize_gemini_file_name",
|
||||
"report_request_id",
|
||||
"GEMINI_FILE_MAPPING_TTL_SECONDS",
|
||||
"sync_codex_quota_from_response_headers",
|
||||
"store_local_gemini_file_mapping",
|
||||
"delete_local_gemini_file_mapping",
|
||||
"GatewaySyncReportRequest",
|
||||
"GatewayStreamReportRequest",
|
||||
] {
|
||||
assert!(
|
||||
report_effects.contains(pattern),
|
||||
"orchestration/report_effects.rs should own local report effect detail {pattern}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -2768,6 +2768,49 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
|
||||
}),
|
||||
);
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct SeenExecutionRuntimeRequest {
|
||||
url: String,
|
||||
authorization: String,
|
||||
}
|
||||
|
||||
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeRequest>));
|
||||
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||
let execution_runtime = Router::new().route(
|
||||
"/v1/execute/sync",
|
||||
any(move |request: Request| {
|
||||
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
|
||||
async move {
|
||||
let plan: aether_contracts::ExecutionPlan = serde_json::from_slice(
|
||||
&to_bytes(request.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should read"),
|
||||
)
|
||||
.expect("plan should parse");
|
||||
*seen_execution_runtime_inner
|
||||
.lock()
|
||||
.expect("mutex should lock") = Some(SeenExecutionRuntimeRequest {
|
||||
url: plan.url.clone(),
|
||||
authorization: plan
|
||||
.headers
|
||||
.get("authorization")
|
||||
.cloned()
|
||||
.unwrap_or_default(),
|
||||
});
|
||||
let result = aether_contracts::ExecutionResult {
|
||||
request_id: plan.request_id,
|
||||
candidate_id: None,
|
||||
status_code: 401,
|
||||
headers: std::collections::BTreeMap::new(),
|
||||
body: None,
|
||||
telemetry: None,
|
||||
error: None,
|
||||
};
|
||||
(StatusCode::OK, Json(result))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let mut provider = sample_provider("provider-codex", "codex", 10);
|
||||
provider.provider_type = "codex".to_string();
|
||||
let endpoint = sample_endpoint(
|
||||
@@ -2832,6 +2875,7 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
|
||||
|
||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||
let (token_url, token_handle) = start_server(token_server).await;
|
||||
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||
let oauth_refresh =
|
||||
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
|
||||
Arc::new(
|
||||
@@ -2840,8 +2884,7 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
|
||||
),
|
||||
]);
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
build_state_with_execution_runtime_override(execution_runtime_url)
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository.clone(),
|
||||
@@ -2889,6 +2932,19 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
|
||||
}
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
|
||||
let seen_execution_runtime_request = seen_execution_runtime
|
||||
.lock()
|
||||
.expect("mutex should lock")
|
||||
.clone()
|
||||
.expect("execution runtime request should be captured");
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.url,
|
||||
"https://chatgpt.com/backend-api/wham/usage"
|
||||
);
|
||||
assert_eq!(
|
||||
seen_execution_runtime_request.authorization,
|
||||
"Bearer refreshed-codex-access-token"
|
||||
);
|
||||
|
||||
let stored_key = provider_catalog_repository
|
||||
.list_keys_by_ids(&["key-codex-oauth-refresh".to_string()])
|
||||
@@ -2974,6 +3030,7 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
|
||||
);
|
||||
|
||||
gateway_handle.abort();
|
||||
execution_runtime_handle.abort();
|
||||
token_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user