use std::time::{Duration, SystemTime, UNIX_EPOCH}; use super::{ any, build_router, build_router_with_state, json, start_server, to_bytes, AppState, Arc, Body, HeaderValue, Json, Mutex, Request, Response, Router, StatusCode, DEPENDENCY_REASON_HEADER, EXECUTION_PATH_HEADER, EXECUTION_PATH_LOCAL_EXECUTION_LOOP_DETECTED, EXECUTION_PATH_LOCAL_ROUTE_NOT_FOUND, EXECUTION_RUNTIME_LOOP_GUARD_HEADER, EXECUTION_RUNTIME_LOOP_GUARD_VALUE, FORWARDED_FOR_HEADER, GATEWAY_HEADER, TRACE_ID_HEADER, TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, TRUSTED_AUTH_API_KEY_ID_HEADER, TRUSTED_AUTH_USER_ID_HEADER, TUNNEL_AFFINITY_FORWARDED_BY_HEADER, TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER, }; use aether_data::repository::auth::{ InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot, }; use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; use aether_data_contracts::repository::provider_catalog::{ StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use sha2::{Digest, Sha256}; fn hash_api_key(value: &str) -> String { let mut hasher = Sha256::new(); hasher.update(value.as_bytes()); format!("{:x}", hasher.finalize()) } fn sample_auth_snapshot( api_key_id: &str, user_id: &str, allowed_model: &str, ) -> StoredAuthApiKeySnapshot { StoredAuthApiKeySnapshot::new( user_id.to_string(), "alice".to_string(), Some("alice@example.com".to_string()), "user".to_string(), "local".to_string(), true, false, Some(serde_json::json!(["openai"])), Some(serde_json::json!(["openai:chat"])), Some(serde_json::json!([allowed_model])), api_key_id.to_string(), Some("default".to_string()), true, false, false, Some(60), Some(5), Some(4_102_444_800), Some(serde_json::json!(["openai"])), Some(serde_json::json!(["openai:chat"])), Some(serde_json::json!([allowed_model])), ) .expect("auth snapshot should build") } fn sample_cli_auth_snapshot( api_key_id: &str, user_id: &str, allowed_model: &str, ) -> StoredAuthApiKeySnapshot { StoredAuthApiKeySnapshot::new( user_id.to_string(), "alice".to_string(), Some("alice@example.com".to_string()), "user".to_string(), "local".to_string(), true, false, Some(serde_json::json!(["openai"])), Some(serde_json::json!(["openai:responses"])), Some(serde_json::json!([allowed_model])), api_key_id.to_string(), Some("default".to_string()), true, false, false, Some(60), Some(5), Some(4_102_444_800), Some(serde_json::json!(["openai"])), Some(serde_json::json!(["openai:responses"])), Some(serde_json::json!([allowed_model])), ) .expect("auth snapshot should build") } fn sample_provider(provider_id: &str) -> StoredProviderCatalogProvider { sample_provider_with_request_timeout(provider_id, None) } fn sample_provider_with_request_timeout( provider_id: &str, request_timeout_secs: Option, ) -> StoredProviderCatalogProvider { StoredProviderCatalogProvider::new( provider_id.to_string(), provider_id.to_string(), Some("https://provider.example".to_string()), "custom".to_string(), ) .expect("provider should build") .with_transport_fields( true, false, false, None, None, None, request_timeout_secs, None, None, ) } fn sample_endpoint(endpoint_id: &str, provider_id: &str) -> StoredProviderCatalogEndpoint { StoredProviderCatalogEndpoint::new( endpoint_id.to_string(), provider_id.to_string(), "openai:chat".to_string(), Some("openai".to_string()), Some("chat".to_string()), true, ) .expect("endpoint should build") .with_transport_fields( "https://api.provider.example".to_string(), None, None, None, None, None, None, None, ) .expect("endpoint transport should build") } fn sample_codex_endpoint(endpoint_id: &str, provider_id: &str) -> StoredProviderCatalogEndpoint { StoredProviderCatalogEndpoint::new( endpoint_id.to_string(), provider_id.to_string(), "openai:responses".to_string(), Some("openai".to_string()), Some("cli".to_string()), true, ) .expect("endpoint should build") .with_transport_fields( "https://chatgpt.com/backend-api/codex".to_string(), None, None, None, None, None, None, None, ) .expect("endpoint transport should build") } fn sample_key(key_id: &str, provider_id: &str, node_id: &str) -> StoredProviderCatalogKey { StoredProviderCatalogKey::new( key_id.to_string(), provider_id.to_string(), "default".to_string(), "api_key".to_string(), None, true, ) .expect("key should build") .with_transport_fields( Some(json!(["openai:chat"])), "plain-upstream-key".to_string(), None, None, Some(json!({"openai:chat": 1})), None, None, Some(json!({ "enabled": true, "mode": "tunnel", "node_id": node_id, })), None, ) .expect("key transport should build") } fn sample_codex_key(key_id: &str, provider_id: &str, node_id: &str) -> StoredProviderCatalogKey { StoredProviderCatalogKey::new( key_id.to_string(), provider_id.to_string(), "default".to_string(), "api_key".to_string(), None, true, ) .expect("key should build") .with_transport_fields( Some(json!(["openai:responses"])), "plain-upstream-key".to_string(), None, None, Some(json!({"openai:responses": 1})), None, None, Some(json!({ "enabled": true, "mode": "tunnel", "node_id": node_id, })), None, ) .expect("key transport should build") } fn tunnel_attachment_key(node_id: &str) -> String { format!("tunnel.attachments.{node_id}") } fn current_unix_secs() -> u64 { SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() .as_secs() } #[tokio::test] async fn gateway_rejects_unknown_path_locally_and_generates_trace_id() { let upstream_hits = Arc::new(Mutex::new(0usize)); let upstream_hits_clone = Arc::clone(&upstream_hits); let upstream = Router::new().route( "/{*path}", any(move |_request: Request| { let upstream_hits_inner = Arc::clone(&upstream_hits_clone); async move { *upstream_hits_inner.lock().expect("mutex should lock") += 1; (StatusCode::CREATED, Body::from("unexpected upstream hit")) } }), ); let (upstream_url, upstream_handle) = start_server(upstream).await; let gateway = build_router().expect("gateway should build"); let (gateway_url, gateway_handle) = start_server(gateway).await; let client = reqwest::Client::new(); let response = client .post(format!("{gateway_url}/does/not/exist?stream=true")) .header(http::header::HOST, "api.example.com") .header(DEPENDENCY_REASON_HEADER, "forged") .body("{\"hello\":\"world\"}") .send() .await .expect("request should succeed"); assert_eq!(response.status(), StatusCode::NOT_FOUND); assert_eq!( response .headers() .get(GATEWAY_HEADER) .and_then(|value| value.to_str().ok()), Some("rust-phase3b") ); assert_eq!( response .headers() .get(EXECUTION_PATH_HEADER) .and_then(|value| value.to_str().ok()), Some(EXECUTION_PATH_LOCAL_ROUTE_NOT_FOUND) ); assert_eq!( response .headers() .get(DEPENDENCY_REASON_HEADER) .and_then(|value| value.to_str().ok()), None ); let response_trace_id = response .headers() .get(TRACE_ID_HEADER) .and_then(|value| value.to_str().ok()) .expect("response trace id should exist") .to_string(); let payload: serde_json::Value = response.json().await.expect("body should parse"); assert_eq!(payload["error"]["type"], "http_error"); assert_eq!(payload["error"]["message"], "Route not found"); assert!(!response_trace_id.is_empty()); assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0); gateway_handle.abort(); upstream_handle.abort(); } #[tokio::test] async fn gateway_preserves_existing_trace_id_on_unknown_local_not_found() { let upstream_hits = Arc::new(Mutex::new(0usize)); let upstream_hits_clone = Arc::clone(&upstream_hits); let upstream = Router::new().route( "/{*path}", any(move |_request: Request| { let upstream_hits_inner = Arc::clone(&upstream_hits_clone); async move { *upstream_hits_inner.lock().expect("mutex should lock") += 1; (StatusCode::OK, Body::from("unexpected upstream hit")) } }), ); let (upstream_url, upstream_handle) = start_server(upstream).await; let gateway = build_router().expect("gateway should build"); let (gateway_url, gateway_handle) = start_server(gateway).await; let response = reqwest::Client::new() .get(format!("{gateway_url}/streaming-proxy")) .header(TRACE_ID_HEADER, "trace-123") .send() .await .expect("request should succeed"); assert_eq!(response.status(), StatusCode::NOT_FOUND); assert_eq!( response .headers() .get(TRACE_ID_HEADER) .and_then(|value| value.to_str().ok()), Some("trace-123") ); assert_eq!( response .headers() .get(EXECUTION_PATH_HEADER) .and_then(|value| value.to_str().ok()), Some(EXECUTION_PATH_LOCAL_ROUTE_NOT_FOUND) ); let payload: serde_json::Value = response.json().await.expect("body should parse"); assert_eq!(payload["error"]["type"], "http_error"); assert_eq!(payload["error"]["message"], "Route not found"); assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0); gateway_handle.abort(); upstream_handle.abort(); } #[tokio::test] async fn gateway_rejects_execution_runtime_loop_guarded_ai_request() { let gateway = build_router().expect("gateway should build"); let (gateway_url, gateway_handle) = start_server(gateway).await; let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/responses")) .header(TRACE_ID_HEADER, "trace-loop-guard-123") .header( EXECUTION_RUNTIME_LOOP_GUARD_HEADER, EXECUTION_RUNTIME_LOOP_GUARD_VALUE, ) .header(http::header::CONTENT_TYPE, "application/json") .body(r#"{"model":"gpt-5.4","input":"hello"}"#) .send() .await .expect("request should succeed"); assert_eq!(response.status(), StatusCode::LOOP_DETECTED); assert_eq!( response .headers() .get(EXECUTION_PATH_HEADER) .and_then(|value| value.to_str().ok()), Some(EXECUTION_PATH_LOCAL_EXECUTION_LOOP_DETECTED) ); let payload: serde_json::Value = response.json().await.expect("body should parse"); assert_eq!(payload["error"]["type"], "http_error"); assert_eq!( payload["error"]["message"], "Gateway detected an execution runtime request loop back into the local frontdoor" ); gateway_handle.abort(); } #[tokio::test] async fn gateway_shapes_execution_loop_rejections_for_claude_routes() { let gateway = build_router().expect("gateway should build"); let (gateway_url, gateway_handle) = start_server(gateway).await; for path in ["/v1/messages", "/v1/messages/count_tokens"] { let response = reqwest::Client::new() .post(format!("{gateway_url}{path}")) .header( EXECUTION_RUNTIME_LOOP_GUARD_HEADER, EXECUTION_RUNTIME_LOOP_GUARD_VALUE, ) .header(http::header::CONTENT_TYPE, "application/json") .body(r#"{"model":"claude-sonnet-4","messages":[]}"#) .send() .await .expect("request should succeed"); assert_eq!(response.status(), StatusCode::LOOP_DETECTED, "path: {path}"); assert_eq!( response .headers() .get(EXECUTION_PATH_HEADER) .and_then(|value| value.to_str().ok()), Some(EXECUTION_PATH_LOCAL_EXECUTION_LOOP_DETECTED), "path: {path}" ); let payload: serde_json::Value = response.json().await.expect("body should parse"); assert_eq!(payload["type"], "error", "path: {path}"); assert_eq!(payload["error"]["type"], "api_error", "path: {path}"); assert_eq!( payload["error"]["message"], "Gateway detected an execution runtime request loop back into the local frontdoor", "path: {path}" ); } gateway_handle.abort(); } #[tokio::test] async fn gateway_shapes_wrong_method_rejections_for_claude_routes() { let gateway = build_router().expect("gateway should build"); let (gateway_url, gateway_handle) = start_server(gateway).await; for path in ["/v1/messages", "/v1/messages/count_tokens"] { let response = reqwest::Client::new() .get(format!("{gateway_url}{path}")) .send() .await .expect("request should succeed"); assert_eq!( response.status(), StatusCode::METHOD_NOT_ALLOWED, "path: {path}" ); assert_eq!( response .headers() .get(http::header::ALLOW) .and_then(|value| value.to_str().ok()), Some("POST"), "path: {path}" ); let payload: serde_json::Value = response.json().await.expect("body should parse"); assert_eq!(payload["type"], "error", "path: {path}"); assert_eq!( payload["error"]["type"], "invalid_request_error", "path: {path}" ); assert_eq!( payload["error"]["message"], "Method not allowed", "path: {path}" ); } gateway_handle.abort(); } #[tokio::test] async fn gateway_rejects_execution_runtime_via_guarded_ai_request() { let gateway = build_router().expect("gateway should build"); let (gateway_url, gateway_handle) = start_server(gateway).await; let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/messages")) .header(TRACE_ID_HEADER, "trace-loop-via-123") .header("via", "1.1 aether-execution-runtime") .header(http::header::CONTENT_TYPE, "application/json") .body(r#"{"model":"claude-sonnet-4","messages":[{"role":"user","content":"hello"}]}"#) .send() .await .expect("request should succeed"); assert_eq!(response.status(), StatusCode::LOOP_DETECTED); assert_eq!( response .headers() .get(EXECUTION_PATH_HEADER) .and_then(|value| value.to_str().ok()), Some(EXECUTION_PATH_LOCAL_EXECUTION_LOOP_DETECTED) ); gateway_handle.abort(); } #[tokio::test] async fn gateway_forwards_public_request_to_remote_tunnel_owner_before_fallback_probe() { #[derive(Debug, Clone)] struct SeenOwnerRequest { path: String, body: String, trace_id: String, gateway_marker: String, authorization: String, trusted_user_id: String, trusted_api_key_id: String, trusted_access_allowed: String, forwarded_for: String, forwarded_by: String, owner_instance_id: String, } let fallback_probe_hits = Arc::new(Mutex::new(0usize)); let fallback_probe_hits_clone = Arc::clone(&fallback_probe_hits); let fallback_probe = Router::new().route( "/{*path}", any(move |_request: Request| { let fallback_probe_hits_inner = Arc::clone(&fallback_probe_hits_clone); async move { *fallback_probe_hits_inner.lock().expect("mutex should lock") += 1; ( StatusCode::OK, Body::from("fallback-probe-should-not-be-hit"), ) } }), ); let seen_owner = Arc::new(Mutex::new(None::)); let seen_owner_clone = Arc::clone(&seen_owner); let owner = Router::new().route( "/v1/chat/completions", any(move |request: Request| { let seen_owner_inner = Arc::clone(&seen_owner_clone); async move { let (parts, body) = request.into_parts(); let raw_body = to_bytes(body, usize::MAX).await.expect("body should read"); tokio::time::sleep(Duration::from_millis(40)).await; *seen_owner_inner.lock().expect("mutex should lock") = Some(SeenOwnerRequest { path: parts .uri .path_and_query() .map(|value| value.as_str()) .unwrap_or("/") .to_string(), body: String::from_utf8(raw_body.to_vec()).expect("utf-8 body"), trace_id: parts .headers .get(TRACE_ID_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), gateway_marker: parts .headers .get(GATEWAY_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), authorization: parts .headers .get(http::header::AUTHORIZATION) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), trusted_user_id: parts .headers .get(TRUSTED_AUTH_USER_ID_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), trusted_api_key_id: parts .headers .get(TRUSTED_AUTH_API_KEY_ID_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), trusted_access_allowed: parts .headers .get(TRUSTED_AUTH_ACCESS_ALLOWED_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), forwarded_for: parts .headers .get(FORWARDED_FOR_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), forwarded_by: parts .headers .get(TUNNEL_AFFINITY_FORWARDED_BY_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), owner_instance_id: parts .headers .get(TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), }); ( StatusCode::OK, [(GATEWAY_HEADER, "gateway-b-owner")], Body::from("owner-gateway-response"), ) } }), ); let (_unused_fallback_probe_url, fallback_probe_handle) = start_server(fallback_probe).await; let (owner_url, owner_handle) = start_server(owner).await; let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_provider_with_request_timeout( "provider-owner", Some(0.1), )], vec![sample_endpoint("endpoint-owner", "provider-owner")], vec![sample_key("key-owner", "provider-owner", "node-owner")], )); let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key("sk-client-openai-affinity")), sample_auth_snapshot("api-key-affinity-1", "user-affinity-1", "gpt-4.1"), )])); let observed_at_unix_secs = current_unix_secs(); let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests( provider_catalog_repository, "development-key", ) .with_auth_api_key_reader(auth_repository) .with_system_config_values_for_tests(vec![( tunnel_attachment_key("node-owner"), serde_json::to_value(crate::tunnel::TunnelAttachmentRecord { gateway_instance_id: "gateway-b".to_string(), relay_base_url: owner_url.clone(), conn_count: 1, observed_at_unix_secs, }) .expect("attachment should serialize"), )]); let mut state = AppState::new().expect("gateway state should build"); state = state .with_data_state_for_tests(data_state) .with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a:8080")); let short_timeout_client = reqwest::Client::builder() .timeout(Duration::from_millis(10)) .build() .expect("test client should build"); state.client = short_timeout_client.clone(); state.owner_forward_client = short_timeout_client; state.remember_scheduler_affinity_target( "scheduler_affinity:api-key-affinity-1:openai:chat:gpt-4.1", crate::cache::SchedulerAffinityTarget { provider_id: "provider-owner".to_string(), endpoint_id: "endpoint-owner".to_string(), key_id: "key-owner".to_string(), }, Duration::from_secs(300), 100, ); let gateway = build_router_with_state(state); let (gateway_url, gateway_handle) = start_server(gateway).await; let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/chat/completions?stream=false")) .header(http::header::CONTENT_TYPE, "application/json") .header( http::header::AUTHORIZATION, "Bearer sk-client-openai-affinity", ) .header(TRACE_ID_HEADER, "trace-tunnel-affinity-forward-1") .body("{\"model\":\"gpt-4.1\",\"messages\":[]}") .send() .await .expect("request should succeed"); assert_eq!(response.status(), StatusCode::OK); assert_eq!( response .headers() .get(GATEWAY_HEADER) .and_then(|value| value.to_str().ok()), Some("gateway-b-owner") ); assert_eq!( response .headers() .get(EXECUTION_PATH_HEADER) .and_then(|value| value.to_str().ok()), Some("tunnel_affinity_forward") ); assert_eq!( response .headers() .get(TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER) .and_then(|value| value.to_str().ok()), Some("gateway-b") ); assert_eq!( response.text().await.expect("body should read"), "owner-gateway-response" ); assert_eq!(*fallback_probe_hits.lock().expect("mutex should lock"), 0); let owner_request = seen_owner .lock() .expect("mutex should lock") .clone() .expect("owner request should be captured"); assert_eq!(owner_request.path, "/v1/chat/completions?stream=false"); assert_eq!( owner_request.body, "{\"model\":\"gpt-4.1\",\"messages\":[]}" ); assert_eq!(owner_request.trace_id, "trace-tunnel-affinity-forward-1"); assert_eq!(owner_request.gateway_marker, "rust-phase3b-affinity"); assert_eq!(owner_request.authorization, ""); assert_eq!(owner_request.trusted_user_id, "user-affinity-1"); assert_eq!(owner_request.trusted_api_key_id, "api-key-affinity-1"); assert_eq!(owner_request.trusted_access_allowed, "true"); assert_eq!(owner_request.forwarded_for, "127.0.0.1"); assert_eq!(owner_request.forwarded_by, "gateway-a"); assert_eq!(owner_request.owner_instance_id, "gateway-b"); gateway_handle.abort(); owner_handle.abort(); fallback_probe_handle.abort(); } #[tokio::test] async fn gateway_aggregates_sync_sse_from_remote_tunnel_owner_before_returning_to_client() { #[derive(Debug, Clone)] struct SeenOwnerRequest { path: String, body: String, trace_id: String, gateway_marker: String, trusted_user_id: String, trusted_api_key_id: String, trusted_access_allowed: String, forwarded_by: String, owner_instance_id: String, } let seen_owner = Arc::new(Mutex::new(None::)); let seen_owner_clone = Arc::clone(&seen_owner); let owner = Router::new().route( "/v1/responses", any(move |request: Request| { let seen_owner_inner = Arc::clone(&seen_owner_clone); async move { let (parts, body) = request.into_parts(); let raw_body = to_bytes(body, usize::MAX).await.expect("body should read"); *seen_owner_inner.lock().expect("mutex should lock") = Some(SeenOwnerRequest { path: parts .uri .path_and_query() .map(|value| value.as_str()) .unwrap_or("/") .to_string(), body: String::from_utf8(raw_body.to_vec()).expect("utf-8 body"), trace_id: parts .headers .get(TRACE_ID_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), gateway_marker: parts .headers .get(GATEWAY_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), trusted_user_id: parts .headers .get(TRUSTED_AUTH_USER_ID_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), trusted_api_key_id: parts .headers .get(TRUSTED_AUTH_API_KEY_ID_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), trusted_access_allowed: parts .headers .get(TRUSTED_AUTH_ACCESS_ALLOWED_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), forwarded_by: parts .headers .get(TUNNEL_AFFINITY_FORWARDED_BY_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), owner_instance_id: parts .headers .get(TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), }); let mut response = Response::builder() .status(StatusCode::OK) .body(Body::from(concat!( "event: response.created\n", "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp-codex-affinity-123\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"in_progress\",\"output\":[]}}\n\n", "event: response.output_text.delta\n", "data: {\"type\":\"response.output_text.delta\",\"output_index\":0,\"content_index\":0,\"delta\":\"Hello from Codex\"}\n\n", "event: response.completed\n", "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp-codex-affinity-123\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"total_tokens\":3}}}\n\n" ))) .expect("response should build"); response.headers_mut().insert( http::header::CONTENT_TYPE, HeaderValue::from_static("text/event-stream"), ); response.headers_mut().insert( http::header::CACHE_CONTROL, HeaderValue::from_static("no-cache"), ); response.headers_mut().insert( http::header::HeaderName::from_static(GATEWAY_HEADER), HeaderValue::from_static("gateway-b-owner"), ); response } }), ); let (owner_url, owner_handle) = start_server(owner).await; let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key("sk-client-openai-cli-affinity")), sample_cli_auth_snapshot("api-key-affinity-cli-1", "user-affinity-cli-1", "gpt-5.4"), )])); let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_provider("provider-cli-owner")], vec![sample_codex_endpoint( "endpoint-cli-owner", "provider-cli-owner", )], vec![sample_codex_key( "key-cli-owner", "provider-cli-owner", "node-cli-owner", )], )); let observed_at_unix_secs = current_unix_secs(); let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests( provider_catalog_repository, "development-key", ) .with_auth_api_key_reader(auth_repository) .with_system_config_values_for_tests(vec![( tunnel_attachment_key("node-cli-owner"), serde_json::to_value(crate::tunnel::TunnelAttachmentRecord { gateway_instance_id: "gateway-b".to_string(), relay_base_url: owner_url.clone(), conn_count: 1, observed_at_unix_secs, }) .expect("attachment should serialize"), )]); let mut state = AppState::new().expect("gateway state should build"); state = state .with_data_state_for_tests(data_state) .with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a:8080")); state.remember_scheduler_affinity_target( "scheduler_affinity:api-key-affinity-cli-1:openai:responses:gpt-5.4", crate::cache::SchedulerAffinityTarget { provider_id: "provider-cli-owner".to_string(), endpoint_id: "endpoint-cli-owner".to_string(), key_id: "key-cli-owner".to_string(), }, Duration::from_secs(300), 100, ); let gateway = build_router_with_state(state); let (gateway_url, gateway_handle) = start_server(gateway).await; let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/responses")) .header(http::header::CONTENT_TYPE, "application/json") .header( http::header::AUTHORIZATION, "Bearer sk-client-openai-cli-affinity", ) .header(TRACE_ID_HEADER, "trace-tunnel-affinity-cli-sync-1") .json(&json!({ "model": "gpt-5.4", "input": "hello", "stream": false })) .send() .await .expect("request should succeed"); assert_eq!(response.status(), StatusCode::OK); assert_eq!( response .headers() .get(GATEWAY_HEADER) .and_then(|value| value.to_str().ok()), Some("gateway-b-owner") ); assert_eq!( response .headers() .get(EXECUTION_PATH_HEADER) .and_then(|value| value.to_str().ok()), Some("tunnel_affinity_forward") ); assert_eq!( response .headers() .get(TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER) .and_then(|value| value.to_str().ok()), Some("gateway-b") ); assert!(response .headers() .get(http::header::CONTENT_TYPE) .and_then(|value| value.to_str().ok()) .is_some_and(|value| value.contains("application/json"))); let body: serde_json::Value = response.json().await.expect("json body should parse"); assert_eq!(body["id"], "resp-codex-affinity-123"); assert_eq!(body["object"], "response"); assert_eq!(body["status"], "completed"); assert_eq!(body["output"][0]["content"][0]["text"], "Hello from Codex"); assert_eq!(body["usage"]["total_tokens"], 3); let owner_request = seen_owner .lock() .expect("mutex should lock") .clone() .expect("owner request should be captured"); assert_eq!(owner_request.path, "/v1/responses"); let owner_body: serde_json::Value = serde_json::from_str(&owner_request.body).expect("owner body should parse"); assert_eq!(owner_body["model"], "gpt-5.4"); assert_eq!(owner_body["stream"], false); assert_eq!(owner_request.trace_id, "trace-tunnel-affinity-cli-sync-1"); assert_eq!(owner_request.gateway_marker, "rust-phase3b-affinity"); assert_eq!(owner_request.trusted_user_id, "user-affinity-cli-1"); assert_eq!(owner_request.trusted_api_key_id, "api-key-affinity-cli-1"); assert_eq!(owner_request.trusted_access_allowed, "true"); assert_eq!(owner_request.forwarded_by, "gateway-a"); assert_eq!(owner_request.owner_instance_id, "gateway-b"); gateway_handle.abort(); owner_handle.abort(); } #[tokio::test] async fn gateway_streamifies_sync_json_from_remote_tunnel_owner_before_returning_to_client() { #[derive(Debug, Clone)] struct SeenOwnerRequest { path: String, body: String, trace_id: String, gateway_marker: String, trusted_user_id: String, trusted_api_key_id: String, trusted_access_allowed: String, forwarded_by: String, owner_instance_id: String, } let seen_owner = Arc::new(Mutex::new(None::)); let seen_owner_clone = Arc::clone(&seen_owner); let owner = Router::new().route( "/v1/responses", any(move |request: Request| { let seen_owner_inner = Arc::clone(&seen_owner_clone); async move { let (parts, body) = request.into_parts(); let raw_body = to_bytes(body, usize::MAX).await.expect("body should read"); *seen_owner_inner.lock().expect("mutex should lock") = Some(SeenOwnerRequest { path: parts .uri .path_and_query() .map(|value| value.as_str()) .unwrap_or("/") .to_string(), body: String::from_utf8(raw_body.to_vec()).expect("utf-8 body"), trace_id: parts .headers .get(TRACE_ID_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), gateway_marker: parts .headers .get(GATEWAY_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), trusted_user_id: parts .headers .get(TRUSTED_AUTH_USER_ID_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), trusted_api_key_id: parts .headers .get(TRUSTED_AUTH_API_KEY_ID_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), trusted_access_allowed: parts .headers .get(TRUSTED_AUTH_ACCESS_ALLOWED_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), forwarded_by: parts .headers .get(TUNNEL_AFFINITY_FORWARDED_BY_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), owner_instance_id: parts .headers .get(TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), }); let encoded_response = serde_json::to_vec(&json!({ "id": "resp-codex-affinity-stream-123", "object": "response", "model": "gpt-5.4", "status": "completed", "output": [{ "type": "message", "id": "msg-codex-affinity-stream-123", "role": "assistant", "content": [{ "type": "output_text", "text": "Hello from affinity sync json", "annotations": [] }] }], "usage": { "input_tokens": 1, "output_tokens": 2, "total_tokens": 3 } })) .expect("body should encode"); let split_at = encoded_response.len() / 2; let first = axum::body::Bytes::copy_from_slice(&encoded_response[..split_at]); let second = axum::body::Bytes::copy_from_slice(&encoded_response[split_at..]); let response_body = Body::from_stream(async_stream::stream! { yield Ok::<_, std::io::Error>(first); tokio::time::sleep(Duration::from_millis(40)).await; yield Ok::<_, std::io::Error>(second); }); let mut response = Response::builder() .status(StatusCode::OK) .body(response_body) .expect("response should build"); response.headers_mut().insert( http::header::CONTENT_TYPE, HeaderValue::from_static("application/json"), ); response.headers_mut().insert( http::header::HeaderName::from_static(GATEWAY_HEADER), HeaderValue::from_static("gateway-b-owner"), ); response } }), ); let (owner_url, owner_handle) = start_server(owner).await; let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key("sk-client-openai-cli-affinity")), sample_cli_auth_snapshot("api-key-affinity-cli-1", "user-affinity-cli-1", "gpt-5.4"), )])); let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_provider("provider-cli-owner")], vec![sample_codex_endpoint( "endpoint-cli-owner", "provider-cli-owner", )], vec![sample_codex_key( "key-cli-owner", "provider-cli-owner", "node-cli-owner", )], )); let observed_at_unix_secs = current_unix_secs(); let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests( provider_catalog_repository, "development-key", ) .with_auth_api_key_reader(auth_repository) .with_system_config_values_for_tests(vec![( tunnel_attachment_key("node-cli-owner"), serde_json::to_value(crate::tunnel::TunnelAttachmentRecord { gateway_instance_id: "gateway-b".to_string(), relay_base_url: owner_url.clone(), conn_count: 1, observed_at_unix_secs, }) .expect("attachment should serialize"), )]); let mut state = AppState::new().expect("gateway state should build"); state = state .with_data_state_for_tests(data_state) .with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a:8080")); state.client = reqwest::Client::builder() .timeout(Duration::from_millis(10)) .build() .expect("short shared client should build"); state.remember_scheduler_affinity_target( "scheduler_affinity:api-key-affinity-cli-1:openai:responses:gpt-5.4", crate::cache::SchedulerAffinityTarget { provider_id: "provider-cli-owner".to_string(), endpoint_id: "endpoint-cli-owner".to_string(), key_id: "key-cli-owner".to_string(), }, Duration::from_secs(300), 100, ); let gateway = build_router_with_state(state); let (gateway_url, gateway_handle) = start_server(gateway).await; let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/responses")) .header(http::header::CONTENT_TYPE, "application/json") .header( http::header::AUTHORIZATION, "Bearer sk-client-openai-cli-affinity", ) .header(TRACE_ID_HEADER, "trace-tunnel-affinity-cli-stream-1") .json(&json!({ "model": "gpt-5.4", "input": "hello", "stream": true })) .send() .await .expect("request should succeed"); assert_eq!(response.status(), StatusCode::OK); assert_eq!( response .headers() .get(GATEWAY_HEADER) .and_then(|value| value.to_str().ok()), Some("gateway-b-owner") ); assert_eq!( response .headers() .get(EXECUTION_PATH_HEADER) .and_then(|value| value.to_str().ok()), Some("tunnel_affinity_forward") ); assert_eq!( response .headers() .get(TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER) .and_then(|value| value.to_str().ok()), Some("gateway-b") ); assert_eq!( response .headers() .get(http::header::CONTENT_TYPE) .and_then(|value| value.to_str().ok()), Some("text/event-stream") ); let body = response.text().await.expect("body should read"); assert!(body.contains("event: response.output_text.delta")); assert!(body.contains("Hello from affinity sync json")); assert!(body.contains("event: response.completed")); let owner_request = seen_owner .lock() .expect("mutex should lock") .clone() .expect("owner request should be captured"); assert_eq!(owner_request.path, "/v1/responses"); let owner_body: serde_json::Value = serde_json::from_str(&owner_request.body).expect("owner body should parse"); assert_eq!(owner_body["model"], "gpt-5.4"); assert_eq!(owner_body["stream"], true); assert_eq!(owner_request.trace_id, "trace-tunnel-affinity-cli-stream-1"); assert_eq!(owner_request.gateway_marker, "rust-phase3b-affinity"); assert_eq!(owner_request.trusted_user_id, "user-affinity-cli-1"); assert_eq!(owner_request.trusted_api_key_id, "api-key-affinity-cli-1"); assert_eq!(owner_request.trusted_access_allowed, "true"); assert_eq!(owner_request.forwarded_by, "gateway-a"); assert_eq!(owner_request.owner_instance_id, "gateway-b"); gateway_handle.abort(); owner_handle.abort(); }