use std::sync::atomic::{AtomicUsize, Ordering}; use super::usage::{ hash_api_key, sample_local_openai_auth_snapshot, sample_local_openai_candidate_row, sample_local_openai_endpoint, sample_local_openai_key, sample_local_openai_provider, }; use super::{ any, build_router_with_state, build_state_with_execution_runtime_override, start_server, wait_until, AppState, Arc, Body, Bytes, GatewayFallbackMetricKind, GatewayFallbackReason, HeaderValue, Infallible, Request, Response, Router, StatusCode, EXECUTION_PATH_DISTRIBUTED_OVERLOADED, EXECUTION_PATH_HEADER, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS, EXECUTION_PATH_LOCAL_OVERLOADED, }; use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY; use aether_data::repository::auth::InMemoryAuthApiKeySnapshotRepository; use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository; use aether_data::repository::candidates::InMemoryRequestCandidateRepository; use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; use aether_data::repository::usage::InMemoryUsageReadRepository; use aether_runtime_state::{ MemoryRuntimeStateConfig, RuntimeSemaphore, RuntimeSemaphoreConfig, RuntimeState, }; use crate::data::GatewayDataState; const CONCURRENCY_TEST_STACK_BYTES: usize = 16 * 1024 * 1024; fn run_concurrency_test(test_name: &'static str, make_future: F) where F: FnOnce() -> Fut + Send + 'static, Fut: std::future::Future + 'static, { let handle = std::thread::Builder::new() .name(test_name.to_string()) .stack_size(CONCURRENCY_TEST_STACK_BYTES) .spawn(move || { let runtime = tokio::runtime::Builder::new_current_thread() .enable_all() .build() .expect("test runtime should build"); runtime.block_on(make_future()); }) .expect("concurrency test thread should spawn"); if let Err(payload) = handle.join() { std::panic::resume_unwind(payload); } } fn memory_runtime_semaphore(gate: &'static str, limit: usize) -> RuntimeSemaphore { RuntimeState::memory(MemoryRuntimeStateConfig::default()) .semaphore(gate, limit, RuntimeSemaphoreConfig::default()) .expect("memory runtime semaphore should build") } fn sample_decision() -> crate::control::GatewayControlDecision { crate::control::GatewayControlDecision { public_path: "/v1/chat/completions".to_string(), public_query_string: None, route_class: Some("ai_public".to_string()), route_family: Some("openai".to_string()), route_kind: Some("chat".to_string()), request_auth_channel: None, auth_endpoint_signature: None, execution_runtime_candidate: true, auth_context: None, admin_principal: None, local_auth_rejection: None, model_directive_policy: Default::default(), } } fn build_local_openai_gateway_state( execution_runtime_override_base_url: impl Into, ) -> AppState { let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key("sk-client-openai-concurrency")), sample_local_openai_auth_snapshot( "api-key-openai-concurrency-1", "user-openai-concurrency-1", ), )])); let candidate_selection_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ sample_local_openai_candidate_row(), ])); let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_local_openai_provider()], vec![sample_local_openai_endpoint()], vec![sample_local_openai_key()], )); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let usage_repository = Arc::new(InMemoryUsageReadRepository::default()); build_state_with_execution_runtime_override(execution_runtime_override_base_url) .with_data_state_for_tests( GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests( auth_repository, candidate_selection_repository, provider_catalog_repository, request_candidate_repository, usage_repository, DEVELOPMENT_ENCRYPTION_KEY, ), ) } #[test] fn gateway_rejects_second_in_flight_stream_request_with_distributed_overload() { run_concurrency_test( "gateway_rejects_second_in_flight_stream_request_with_distributed_overload", gateway_rejects_second_in_flight_stream_request_with_distributed_overload_impl, ); } async fn gateway_rejects_second_in_flight_stream_request_with_distributed_overload_impl() { let execution_runtime_hits = Arc::new(AtomicUsize::new(0)); let execution_runtime_hits_clone = Arc::clone(&execution_runtime_hits); let execution_runtime = Router::new().route( "/v1/execute/stream", any(move |_request: Request| { let execution_runtime_hits = Arc::clone(&execution_runtime_hits_clone); async move { execution_runtime_hits.fetch_add(1, Ordering::SeqCst); let stream = async_stream::stream! { yield Ok::<_, Infallible>(Bytes::from_static( b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n", )); yield Ok::<_, Infallible>(Bytes::from_static( b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-concurrency-123\\\"}\\n\\n\"}}\n", )); futures_util::future::pending::<()>().await; }; let mut response = Response::builder() .status(StatusCode::OK) .body(Body::from_stream(stream)) .expect("response should build"); response.headers_mut().insert( http::header::CONTENT_TYPE, HeaderValue::from_static("application/x-ndjson"), ); response } }), ); let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; let distributed_gate = memory_runtime_semaphore("gateway_requests_distributed", 1); let gateway_a = build_router_with_state( build_local_openai_gateway_state(execution_runtime_url.clone()) .with_distributed_request_concurrency_gate(distributed_gate.clone()), ); let gateway_b = build_router_with_state( build_local_openai_gateway_state(execution_runtime_url) .with_distributed_request_concurrency_gate(distributed_gate), ); let (gateway_a_url, gateway_a_handle) = start_server(gateway_a).await; let (gateway_b_url, gateway_b_handle) = start_server(gateway_b).await; let client = reqwest::Client::new(); let first_response = client .post(format!("{gateway_a_url}/v1/chat/completions")) .header( http::header::AUTHORIZATION, "Bearer sk-client-openai-concurrency", ) .header(http::header::CONTENT_TYPE, "application/json") .body("{\"model\":\"gpt-5\",\"messages\":[],\"stream\":true}") .send() .await .expect("first request should succeed"); wait_until(500, || execution_runtime_hits.load(Ordering::SeqCst) == 1).await; let second_response = client .post(format!("{gateway_b_url}/v1/chat/completions")) .header( http::header::AUTHORIZATION, "Bearer sk-client-openai-concurrency", ) .header(http::header::CONTENT_TYPE, "application/json") .body("{\"model\":\"gpt-5\",\"messages\":[],\"stream\":true}") .send() .await .expect("second request should complete"); assert_eq!(second_response.status(), StatusCode::SERVICE_UNAVAILABLE); assert_eq!( second_response .headers() .get(EXECUTION_PATH_HEADER) .and_then(|value| value.to_str().ok()), Some(EXECUTION_PATH_DISTRIBUTED_OVERLOADED) ); assert_eq!( second_response .json::() .await .expect("json body should decode")["error"]["details"]["gate"], "gateway_requests_distributed" ); assert_eq!(execution_runtime_hits.load(Ordering::SeqCst), 1); drop(first_response); gateway_a_handle.abort(); gateway_b_handle.abort(); execution_runtime_handle.abort(); } #[test] fn gateway_rejects_second_in_flight_stream_request_with_local_overload() { run_concurrency_test( "gateway_rejects_second_in_flight_stream_request_with_local_overload", gateway_rejects_second_in_flight_stream_request_with_local_overload_impl, ); } async fn gateway_rejects_second_in_flight_stream_request_with_local_overload_impl() { let execution_runtime_hits = Arc::new(AtomicUsize::new(0)); let execution_runtime_hits_clone = Arc::clone(&execution_runtime_hits); let execution_runtime = Router::new().route( "/v1/execute/stream", any(move |_request: Request| { let execution_runtime_hits = Arc::clone(&execution_runtime_hits_clone); async move { execution_runtime_hits.fetch_add(1, Ordering::SeqCst); let stream = async_stream::stream! { yield Ok::<_, Infallible>(Bytes::from_static( b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n", )); yield Ok::<_, Infallible>(Bytes::from_static( b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-concurrency-123\\\"}\\n\\n\"}}\n", )); futures_util::future::pending::<()>().await; }; let mut response = Response::builder() .status(StatusCode::OK) .body(Body::from_stream(stream)) .expect("response should build"); response.headers_mut().insert( http::header::CONTENT_TYPE, HeaderValue::from_static("application/x-ndjson"), ); response } }), ); let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; let gateway = build_router_with_state( build_local_openai_gateway_state(execution_runtime_url).with_request_concurrency_limit(1), ); let (gateway_url, gateway_handle) = start_server(gateway).await; let client = reqwest::Client::new(); let first_response = client .post(format!("{gateway_url}/v1/chat/completions")) .header( http::header::AUTHORIZATION, "Bearer sk-client-openai-concurrency", ) .header(http::header::CONTENT_TYPE, "application/json") .body("{\"model\":\"gpt-5\",\"messages\":[],\"stream\":true}") .send() .await .expect("first request should succeed"); wait_until(500, || execution_runtime_hits.load(Ordering::SeqCst) == 1).await; let second_response = client .post(format!("{gateway_url}/v1/chat/completions")) .header( http::header::AUTHORIZATION, "Bearer sk-client-openai-concurrency", ) .header(http::header::CONTENT_TYPE, "application/json") .body("{\"model\":\"gpt-5\",\"messages\":[],\"stream\":true}") .send() .await .expect("second request should complete"); assert_eq!(second_response.status(), StatusCode::SERVICE_UNAVAILABLE); assert_eq!( second_response .headers() .get(EXECUTION_PATH_HEADER) .and_then(|value| value.to_str().ok()), Some(EXECUTION_PATH_LOCAL_OVERLOADED) ); assert_eq!( second_response .json::() .await .expect("json body should decode")["error"]["type"], "overloaded" ); assert_eq!(execution_runtime_hits.load(Ordering::SeqCst), 1); drop(first_response); gateway_handle.abort(); execution_runtime_handle.abort(); } #[test] fn gateway_exposes_request_concurrency_metrics() { run_concurrency_test( "gateway_exposes_request_concurrency_metrics", gateway_exposes_request_concurrency_metrics_impl, ); } async fn gateway_exposes_request_concurrency_metrics_impl() { let state = AppState::new() .expect("gateway state should build") .with_request_concurrency_limit(3) .with_distributed_request_concurrency_gate(memory_runtime_semaphore( "gateway_requests_distributed", 5, )); assert!(state.prewarm_metric_snapshot().await); let gateway = build_router_with_state(state); let (gateway_url, gateway_handle) = start_server(gateway).await; let response = reqwest::Client::new() .get(format!("{gateway_url}/_gateway/metrics")) .send() .await .expect("request should succeed"); assert_eq!(response.status(), StatusCode::OK); assert_eq!( response .headers() .get(http::header::CONTENT_TYPE) .and_then(|value| value.to_str().ok()), Some("text/plain; version=0.0.4; charset=utf-8") ); let body = response.text().await.expect("body should read"); assert!(body.contains("service_up{service=\"aether-gateway\"} 1")); assert!(body.contains("concurrency_in_flight{gate=\"gateway_requests\"} 0")); assert!(body.contains("concurrency_available_permits{gate=\"gateway_requests\"} 3")); assert!(body.contains("concurrency_in_flight{gate=\"gateway_requests_distributed\"} 0")); assert!(body.contains("concurrency_available_permits{gate=\"gateway_requests_distributed\"} 5")); assert!(body.contains("tunnel_proxy_connections 0")); assert!(body.contains("tunnel_nodes 0")); assert!(body.contains("tunnel_active_streams 0")); assert!(body.contains("gateway_process_cpu_usage_basis_points ")); assert!(body.contains("gateway_process_memory_bytes ")); assert!(body.contains("gateway_process_threads ")); assert!(body.contains("gateway_process_open_fds ")); assert!(body.contains("gateway_process_fd_limit ")); assert!(body.contains("gateway_process_socket_fds ")); assert!(body.contains("gateway_allocator_observability_available ")); assert!(body.contains("gateway_allocator_allocated_bytes ")); assert!(body.contains("gateway_allocator_active_bytes ")); assert!(body.contains("gateway_allocator_resident_bytes ")); assert!(body.contains("gateway_allocator_active_to_allocated_basis_points ")); assert!(body.contains("gateway_network_observability_available ")); assert!(body.contains("gateway_network_received_bytes_total ")); assert!(body.contains("gateway_tcp_state_observability_available ")); assert!(body.contains("gateway_host_tcp_established_connections ")); assert!(body.contains("gateway_process_tcp_established_connections ")); assert!(body.contains("postgres_observability_available{driver=\"postgres\"} 0")); assert!(body.contains("postgres_observability_unavailable{driver=\"postgres\"} 0")); assert!(body.contains("postgres_lock_waiting_connections{driver=\"postgres\"} 0")); assert!(body.contains("postgres_oldest_active_query_age_ms{driver=\"postgres\"} 0")); assert!(body.contains("postgres_oldest_transaction_age_ms{driver=\"postgres\"} 0")); assert!(body.contains("redis_runtime_enabled{backend=\"redis\"} 0")); assert!(body.contains("redis_runtime_health_unavailable{backend=\"redis\"} 0")); assert!(body.contains("redis_runtime_connected_clients{backend=\"redis\"} 0")); assert!(body.contains("redis_runtime_used_memory_bytes{backend=\"redis\"} 0")); assert!(body.contains("usage_runtime_queue_worker_read_batches_total 0")); assert!(body.contains("usage_runtime_queue_worker_read_entries_total 0")); assert!(body.contains("usage_runtime_queue_worker_reclaimed_entries_total 0")); assert!(body.contains("usage_runtime_queue_worker_acked_entries_total 0")); assert!(body.contains("usage_runtime_queue_worker_dead_lettered_entries_total 0")); assert!(body.contains("usage_runtime_queue_worker_process_failures_total 0")); assert!(body.contains("usage_runtime_queue_worker_read_failures_total 0")); assert!(body.contains("usage_runtime_queue_worker_reclaim_failures_total 0")); assert!(body.contains("usage_queue_health_unavailable 0")); assert!( body.contains("usage_queue_enabled{stream=\"usage:events\",group=\"usage_consumers\"} 0") ); assert!(body .contains("usage_queue_configured{stream=\"usage:events\",group=\"usage_consumers\"} 0")); assert!(body.contains("usage_queue_dlq_length{stream=\"usage:events:dlq\"} 0")); assert!(body.contains("usage_counter_health_unavailable 0")); assert!(body.contains("usage_counter_outbox_pending_rows 0")); assert!(body.contains("usage_counter_outbox_oldest_pending_age_seconds 0")); assert!(body.contains("usage_counter_outbox_flush_batches_total 0")); assert!(body.contains("usage_counter_outbox_flush_rows_claimed_total 0")); assert!(body.contains("usage_counter_outbox_flush_failed_batches_total 0")); assert!(body.contains("usage_counter_outbox_cleanup_rows_total 0")); assert!(body.contains("usage_counter_outbox_cleanup_failed_batches_total 0")); assert!(body.contains("gateway_background_tasks_active 0")); assert!(body.contains("gateway_background_tasks_supervised_total 0")); assert!(body.contains("gateway_background_tasks_unexpected_exits_total 0")); assert!(body.contains("gateway_background_tasks_panicked_total 0")); assert!(body.contains("gateway_background_tasks_aborted_total 0")); assert!(body.contains("gateway_tokio_runtime_observability_available 1")); assert!(body.contains("gateway_tokio_runtime_workers ")); assert!(body.contains("gateway_tokio_runtime_alive_tasks ")); assert!(body.contains("gateway_tokio_runtime_global_queue_depth ")); gateway_handle.abort(); } #[test] fn gateway_exposes_fallback_metrics() { run_concurrency_test( "gateway_exposes_fallback_metrics", gateway_exposes_fallback_metrics_impl, ); } async fn gateway_exposes_fallback_metrics_impl() { let state = AppState::new().expect("gateway state should build"); let decision = sample_decision(); state.record_fallback_metric( GatewayFallbackMetricKind::DecisionRemote, Some(&decision), Some("openai_chat_sync"), None, GatewayFallbackReason::LocalDecisionMiss, ); state.record_fallback_metric( GatewayFallbackMetricKind::PlanFallback, Some(&decision), Some("openai_chat_sync"), None, GatewayFallbackReason::RemoteDecisionMiss, ); state.record_fallback_metric( GatewayFallbackMetricKind::LocalExecutionRuntimeMiss, Some(&decision), None, Some(EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS), GatewayFallbackReason::LocalExecutionPathRequired, ); assert!(state.prewarm_metric_snapshot().await); let gateway = build_router_with_state(state); let (gateway_url, gateway_handle) = start_server(gateway).await; let response = reqwest::Client::new() .get(format!("{gateway_url}/_gateway/metrics")) .send() .await .expect("request should succeed"); assert_eq!(response.status(), StatusCode::OK); let body = response.text().await.expect("body should read"); let decision_remote = body .lines() .find(|line| line.starts_with("decision_remote_total{")) .expect("decision_remote_total sample should be rendered"); assert!(decision_remote.contains("route_class=\"ai_public\"")); assert!(decision_remote.contains("route_family=\"openai\"")); assert!(decision_remote.contains("route_kind=\"chat\"")); assert!(decision_remote.contains("plan_kind=\"openai_chat_sync\"")); assert!(decision_remote.contains("execution_path=\"none\"")); assert!(decision_remote.contains("reason=\"local_decision_miss\"")); assert!(decision_remote.ends_with(" 1")); let plan_fallback = body .lines() .find(|line| line.starts_with("plan_fallback_total{")) .expect("plan_fallback_total sample should be rendered"); assert!(plan_fallback.contains("route_class=\"ai_public\"")); assert!(plan_fallback.contains("route_family=\"openai\"")); assert!(plan_fallback.contains("route_kind=\"chat\"")); assert!(plan_fallback.contains("plan_kind=\"openai_chat_sync\"")); assert!(plan_fallback.contains("execution_path=\"none\"")); assert!(plan_fallback.contains("reason=\"remote_decision_miss\"")); assert!(plan_fallback.ends_with(" 1")); let local_execution_runtime_miss = body .lines() .find(|line| line.starts_with("local_execution_runtime_miss_total{")) .expect("local_execution_runtime_miss_total sample should be rendered"); assert!(local_execution_runtime_miss.contains("route_class=\"ai_public\"")); assert!(local_execution_runtime_miss.contains("route_family=\"openai\"")); assert!(local_execution_runtime_miss.contains("route_kind=\"chat\"")); assert!(local_execution_runtime_miss.contains("plan_kind=\"none\"")); assert!(local_execution_runtime_miss.contains(&format!( "execution_path=\"{}\"", EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS ))); assert!(local_execution_runtime_miss.contains("reason=\"local_execution_path_required\"")); assert!(local_execution_runtime_miss.ends_with(" 1")); gateway_handle.abort(); }