use super::{ any, build_router_with_state, build_state_with_execution_runtime_override, json, next_non_keepalive_chunk, start_server, strip_sse_keepalive_comments, to_bytes, Arc, Body, Bytes, HeaderName, HeaderValue, Json, Mutex, Request, Response, Router, StatusCode, TRACE_ID_HEADER, }; use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY}; use aether_data::repository::auth::{ InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot, }; use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository; use aether_data::repository::candidates::InMemoryRequestCandidateRepository; use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; use aether_data_contracts::repository::candidate_selection::{ StoredMinimalCandidateSelectionRow, StoredProviderModelMapping, }; use aether_data_contracts::repository::candidates::{ RequestCandidateReadRepository, RequestCandidateStatus, }; use aether_data_contracts::repository::provider_catalog::{ StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use sha2::{Digest, Sha256}; const STREAM_PROVIDER_TEST_STACK_BYTES: usize = 16 * 1024 * 1024; fn run_stream_provider_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(STREAM_PROVIDER_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("stream provider test thread should spawn"); if let Err(payload) = handle.join() { std::panic::resume_unwind(payload); } } #[test] fn gateway_executes_kiro_claude_cli_stream_via_local_provider_catalog_candidate() { run_stream_provider_test( "gateway_executes_kiro_claude_cli_stream_via_local_provider_catalog_candidate", gateway_executes_kiro_claude_cli_stream_via_local_provider_catalog_candidate_impl, ); } async fn gateway_executes_kiro_claude_cli_stream_via_local_provider_catalog_candidate_impl() { use base64::Engine as _; #[derive(Debug, Clone)] struct SeenExecutionRuntimeStreamRequest { trace_id: String, url: String, authorization: String, accept: String, host: String, endpoint_tag: String, mapped_model: String, current_content: String, debug_tag: String, proxy_node_id: String, transport_profile_id: String, } fn crc32(data: &[u8]) -> u32 { let mut crc = 0xffff_ffffu32; for &byte in data { crc ^= byte as u32; for _ in 0..8 { let mask = if crc & 1 == 1 { 0xedb8_8320 } else { 0 }; crc = (crc >> 1) ^ mask; } } !crc } fn encode_string_header(name: &str, value: &str) -> Vec { let mut out = Vec::new(); out.push(name.len() as u8); out.extend_from_slice(name.as_bytes()); out.push(7); out.extend_from_slice(&(value.len() as u16).to_be_bytes()); out.extend_from_slice(value.as_bytes()); out } fn encode_event_frame( message_type: &str, event_type: Option<&str>, payload: serde_json::Value, ) -> Vec { let mut headers = encode_string_header(":message-type", message_type); if let Some(event_type) = event_type { headers.extend_from_slice(&encode_string_header(":event-type", event_type)); } let payload = serde_json::to_vec(&payload).expect("payload should encode"); let total_len = 12 + headers.len() + payload.len() + 4; let mut out = Vec::with_capacity(total_len); out.extend_from_slice(&(total_len as u32).to_be_bytes()); out.extend_from_slice(&(headers.len() as u32).to_be_bytes()); let prelude_crc = crc32(&out[..8]); out.extend_from_slice(&prelude_crc.to_be_bytes()); out.extend_from_slice(&headers); out.extend_from_slice(&payload); let message_crc = crc32(&out); out.extend_from_slice(&message_crc.to_be_bytes()); out } 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) -> 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!(["claude", "kiro"])), Some(serde_json::json!(["claude:messages"])), Some(serde_json::json!(["claude-sonnet-4"])), 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!(["claude", "kiro"])), Some(serde_json::json!(["claude:messages"])), Some(serde_json::json!(["claude-sonnet-4"])), ) .expect("auth snapshot should build") } fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow { StoredMinimalCandidateSelectionRow { provider_id: "provider-kiro-cli-local-stream-1".to_string(), provider_name: "kiro".to_string(), provider_type: "kiro".to_string(), provider_priority: 10, provider_is_active: true, endpoint_id: "endpoint-kiro-cli-local-stream-1".to_string(), endpoint_api_format: "claude:messages".to_string(), endpoint_api_family: Some("claude".to_string()), endpoint_kind: Some("cli".to_string()), endpoint_is_active: true, key_id: "key-kiro-cli-local-stream-1".to_string(), key_name: "prod".to_string(), key_auth_type: "bearer".to_string(), key_is_active: true, key_api_formats: Some(vec!["claude:messages".to_string()]), key_allowed_models: None, key_capabilities: None, key_internal_priority: 5, key_global_priority_by_format: Some(serde_json::json!({"claude:messages": 1})), model_id: "model-kiro-cli-local-stream-1".to_string(), global_model_id: "global-model-kiro-cli-local-stream-1".to_string(), global_model_name: "claude-sonnet-4".to_string(), global_model_mappings: None, global_model_supports_streaming: Some(true), model_provider_model_name: "claude-sonnet-4-upstream".to_string(), model_provider_model_mappings: Some(vec![StoredProviderModelMapping { name: "claude-sonnet-4-upstream".to_string(), priority: 1, api_formats: Some(vec!["claude:messages".to_string()]), endpoint_ids: None, operations: None, }]), model_supports_streaming: Some(true), model_is_active: true, model_is_available: true, } } fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider { StoredProviderCatalogProvider::new( "provider-kiro-cli-local-stream-1".to_string(), "kiro".to_string(), Some("https://example.com".to_string()), "kiro".to_string(), ) .expect("provider should build") .with_transport_fields( true, false, false, None, Some(2), Some(serde_json::json!({"url":"http://provider-proxy.internal:8080"})), Some(20.0), None, None, ) } fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint { StoredProviderCatalogEndpoint::new( "endpoint-kiro-cli-local-stream-1".to_string(), "provider-kiro-cli-local-stream-1".to_string(), "claude:messages".to_string(), Some("claude".to_string()), Some("cli".to_string()), true, ) .expect("endpoint should build") .with_transport_fields( "https://kiro.{region}.example?tenant=demo".to_string(), Some(serde_json::json!([ {"action":"set","key":"accept","value":"text/plain"}, {"action":"set","key":"x-endpoint-tag","value":"kiro-cli-local-stream"} ])), Some(serde_json::json!([ {"action":"set","path":"debugTag","value":"kiro-local-stream"} ])), Some(2), None, None, None, None, ) .expect("endpoint transport should build") } fn sample_provider_catalog_key() -> StoredProviderCatalogKey { let auth_config = serde_json::json!({ "provider_type": "kiro", "access_token": "cached-kiro-access-token", "expires_at": 4102444800_u64, "refresh_token": "rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr", "machine_id": "123e4567-e89b-12d3-a456-426614174000", "api_region": "us-east-1", "kiro_version": "0.8.0", "system_version": "darwin#24.6.0", "node_version": "22.21.1" }); StoredProviderCatalogKey::new( "key-kiro-cli-local-stream-1".to_string(), "provider-kiro-cli-local-stream-1".to_string(), "prod".to_string(), "bearer".to_string(), None, true, ) .expect("key should build") .with_transport_fields( Some(serde_json::json!(["claude:messages"])), encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "__placeholder__") .expect("api key should encrypt"), Some( encrypt_python_fernet_plaintext( DEVELOPMENT_ENCRYPTION_KEY, auth_config.to_string().as_str(), ) .expect("auth config should encrypt"), ), None, Some(serde_json::json!({"claude:messages": 1})), None, None, Some( serde_json::json!({"enabled": true, "node_id":"proxy-node-kiro-cli-local-stream"}), ), Some(serde_json::json!({"transport_profile":"chrome_136"})), ) .expect("key transport should build") } let seen_execution_runtime = Arc::new(Mutex::new(None::)); let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime); let seen_report = Arc::new(Mutex::new(false)); let seen_report_clone = Arc::clone(&seen_report); let decision_hits = Arc::new(Mutex::new(0usize)); let decision_hits_clone = Arc::clone(&decision_hits); let plan_hits = Arc::new(Mutex::new(0usize)); let plan_hits_clone = Arc::clone(&plan_hits); let public_hits = Arc::new(Mutex::new(0usize)); let public_hits_clone = Arc::clone(&public_hits); let upstream = Router::new() .route( "/api/internal/gateway/resolve", any(|_request: Request| async move { Json(json!({ "action": "proxy_public", "route_class": "ai_public", "route_family": "claude", "route_kind": "messages", "request_auth_channel": "bearer_like", "auth_endpoint_signature": "claude:messages", "execution_runtime_candidate": true, "auth_context": { "user_id": "user-kiro-cli-local-stream-123", "api_key_id": "key-kiro-cli-local-stream-123", "access_allowed": true }, "public_path": "/v1/messages" })) }), ) .route( "/api/internal/gateway/decision-stream", any(move |_request: Request| { let decision_hits_inner = Arc::clone(&decision_hits_clone); async move { *decision_hits_inner.lock().expect("mutex should lock") += 1; Json(json!({"action": "proxy_public"})) } }), ) .route( "/api/internal/gateway/plan-stream", any(move |_request: Request| { let plan_hits_inner = Arc::clone(&plan_hits_clone); async move { *plan_hits_inner.lock().expect("mutex should lock") += 1; Json(json!({"action": "proxy_public"})) } }), ) .route( "/api/internal/gateway/report-stream", any(move |request: Request| { let seen_report_inner = Arc::clone(&seen_report_clone); async move { let (_parts, body) = request.into_parts(); let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read"); *seen_report_inner.lock().expect("mutex should lock") = true; Json(json!({"ok": true})) } }), ) .route( "/v1/messages", any(move |_request: Request| { let public_hits_inner = Arc::clone(&public_hits_clone); async move { *public_hits_inner.lock().expect("mutex should lock") += 1; (StatusCode::IM_A_TEAPOT, Body::from("public-route-hit")) } }), ); let execution_runtime = Router::new().route( "/v1/execute/stream", any(move |request: Request| { let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone); async move { let (parts, body) = request.into_parts(); let raw_body = to_bytes(body, usize::MAX).await.expect("body should read"); let payload: serde_json::Value = serde_json::from_slice(&raw_body).expect("execution runtime payload should parse"); *seen_execution_runtime_inner.lock().expect("mutex should lock") = Some(SeenExecutionRuntimeStreamRequest { trace_id: parts .headers .get(TRACE_ID_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), url: payload .get("url") .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), authorization: payload .get("headers") .and_then(|value| value.get("authorization")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), accept: payload .get("headers") .and_then(|value| value.get("accept")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), host: payload .get("headers") .and_then(|value| value.get("host")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), endpoint_tag: payload .get("headers") .and_then(|value| value.get("x-endpoint-tag")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), mapped_model: payload .get("body") .and_then(|value| value.get("json_body")) .and_then(|value| value.get("conversationState")) .and_then(|value| value.get("currentMessage")) .and_then(|value| value.get("userInputMessage")) .and_then(|value| value.get("modelId")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), current_content: payload .get("body") .and_then(|value| value.get("json_body")) .and_then(|value| value.get("conversationState")) .and_then(|value| value.get("currentMessage")) .and_then(|value| value.get("userInputMessage")) .and_then(|value| value.get("content")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), debug_tag: payload .get("body") .and_then(|value| value.get("json_body")) .and_then(|value| value.get("debugTag")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), proxy_node_id: payload .get("proxy") .and_then(|value| value.get("node_id")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), transport_profile_id: payload .get("transport_profile").and_then(|value| value.get("profile_id")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), }); let kiro_frames = [ encode_event_frame( "event", Some("assistantResponseEvent"), json!({"content": "Hello from Kiro stream"}), ), encode_event_frame( "event", Some("contextUsageEvent"), json!({"contextUsagePercentage": 1.0}), ), ] .concat(); let frames = format!( concat!( "{{\"type\":\"headers\",\"payload\":{{\"kind\":\"headers\",\"status_code\":200,\"headers\":{{\"content-type\":\"application/vnd.amazon.eventstream\"}}}}}}\n", "{{\"type\":\"data\",\"payload\":{{\"kind\":\"data\",\"chunk_b64\":\"{}\"}}}}\n", "{{\"type\":\"telemetry\",\"payload\":{{\"kind\":\"telemetry\",\"telemetry\":{{\"elapsed_ms\":27,\"upstream_bytes\":64}}}}}}\n", "{{\"type\":\"eof\",\"payload\":{{\"kind\":\"eof\"}}}}\n" ), base64::engine::general_purpose::STANDARD.encode(kiro_frames) ); let mut response = Response::builder() .status(StatusCode::OK) .body(Body::from(frames)) .expect("response should build"); response.headers_mut().insert( http::header::CONTENT_TYPE, HeaderValue::from_static("application/x-ndjson"), ); response } }), ); let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key("sk-client-kiro-cli-local-stream")), sample_auth_snapshot( "key-kiro-cli-local-stream-123", "user-kiro-cli-local-stream-123", ), )])); let candidate_selection_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ sample_candidate_row(), ])); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_provider_catalog_provider()], vec![sample_provider_catalog_endpoint()], vec![sample_provider_catalog_key()], )); let (upstream_url, upstream_handle) = start_server(upstream).await; let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone()) .with_data_state_for_tests( crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests( auth_repository, candidate_selection_repository, provider_catalog_repository, Arc::clone(&request_candidate_repository), DEVELOPMENT_ENCRYPTION_KEY, ) .attach_proxy_node_repository_for_tests( crate::tests::ai_execute::ai_execute_proxy_node_repository([ "proxy-node-kiro-cli-local-stream", ]), ), ); let gateway = build_router_with_state(gateway_state); let (gateway_url, gateway_handle) = start_server(gateway).await; let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/messages")) .header(http::header::CONTENT_TYPE, "application/json") .header(http::header::USER_AGENT, "Claude-Code/2.1.0") .header( http::header::AUTHORIZATION, "Bearer sk-client-kiro-cli-local-stream", ) .header(TRACE_ID_HEADER, "trace-kiro-cli-local-stream-123") .body("{\"model\":\"claude-sonnet-4\",\"messages\":[{\"role\":\"user\",\"content\":\"hello stream\"}],\"stream\":true}") .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/event-stream") ); let body = response.text().await.expect("body should read"); assert!(body.contains("Hello from Kiro stream")); assert!(body.contains("event: message_start")); let seen_execution_runtime_request = seen_execution_runtime .lock() .expect("mutex should lock") .clone() .expect("execution runtime stream should be captured"); assert_eq!( seen_execution_runtime_request.trace_id, "trace-kiro-cli-local-stream-123" ); assert_eq!( seen_execution_runtime_request.url, "https://kiro.us-east-1.example/generateAssistantResponse?tenant=demo" ); assert_eq!( seen_execution_runtime_request.authorization, "Bearer cached-kiro-access-token" ); assert_eq!( seen_execution_runtime_request.accept, "application/vnd.amazon.eventstream" ); assert_eq!( seen_execution_runtime_request.host, "q.us-east-1.amazonaws.com" ); assert_eq!( seen_execution_runtime_request.endpoint_tag, "kiro-cli-local-stream" ); assert_eq!( seen_execution_runtime_request.mapped_model, "claude-sonnet-4-upstream" ); assert_eq!( seen_execution_runtime_request.current_content, "hello stream" ); assert_eq!( seen_execution_runtime_request.debug_tag, "kiro-local-stream" ); assert_eq!( seen_execution_runtime_request.proxy_node_id, "proxy-node-kiro-cli-local-stream" ); assert_eq!( seen_execution_runtime_request.transport_profile_id, "chrome_136" ); let stored_candidates = request_candidate_repository .list_by_request_id("trace-kiro-cli-local-stream-123") .await .expect("request candidate trace should read"); assert_eq!(stored_candidates.len(), 1); assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success); tokio::time::sleep(std::time::Duration::from_millis(100)).await; assert!( !*seen_report.lock().expect("mutex should lock"), "report-stream should stay local when request candidate persistence is available" ); assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0); assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0); assert_eq!(*public_hits.lock().expect("mutex should lock"), 0); gateway_handle.abort(); execution_runtime_handle.abort(); upstream_handle.abort(); } #[test] fn gateway_executes_claude_cli_stream_via_local_decision_gate_without_waiting_for_same_format_prefetch( ) { run_stream_provider_test( "gateway_executes_claude_cli_stream_via_local_decision_gate_without_waiting_for_same_format_prefetch", gateway_executes_claude_cli_stream_via_local_decision_gate_without_waiting_for_same_format_prefetch_impl, ); } async fn gateway_executes_claude_cli_stream_via_local_decision_gate_without_waiting_for_same_format_prefetch_impl( ) { #[derive(Debug, Clone)] struct SeenExecutionRuntimeStreamRequest { trace_id: String, url: String, model: String, stream: bool, accept: String, authorization: String, endpoint_tag: String, metadata_mode: String, metadata_source: String, proxy_node_id: String, transport_profile_id: String, } 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) -> 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!(["claude"])), Some(serde_json::json!(["claude:messages"])), Some(serde_json::json!(["claude-code"])), 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!(["claude"])), Some(serde_json::json!(["claude:messages"])), Some(serde_json::json!(["claude-code"])), ) .expect("auth snapshot should build") } fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow { StoredMinimalCandidateSelectionRow { provider_id: "provider-claude-cli-local-1".to_string(), provider_name: "claude".to_string(), provider_type: "custom".to_string(), provider_priority: 10, provider_is_active: true, endpoint_id: "endpoint-claude-cli-local-1".to_string(), endpoint_api_format: "claude:messages".to_string(), endpoint_api_family: Some("claude".to_string()), endpoint_kind: Some("cli".to_string()), endpoint_is_active: true, key_id: "key-claude-cli-local-1".to_string(), key_name: "prod".to_string(), key_auth_type: "bearer".to_string(), key_is_active: true, key_api_formats: Some(vec!["claude:messages".to_string()]), key_allowed_models: None, key_capabilities: None, key_internal_priority: 5, key_global_priority_by_format: Some(serde_json::json!({"claude:messages": 1})), model_id: "model-claude-cli-local-1".to_string(), global_model_id: "global-model-claude-cli-local-1".to_string(), global_model_name: "claude-code".to_string(), global_model_mappings: None, global_model_supports_streaming: Some(true), model_provider_model_name: "claude-code-upstream".to_string(), model_provider_model_mappings: Some(vec![StoredProviderModelMapping { name: "claude-code-upstream".to_string(), priority: 1, api_formats: Some(vec!["claude:messages".to_string()]), endpoint_ids: None, operations: None, }]), model_supports_streaming: Some(true), model_is_active: true, model_is_available: true, } } fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider { StoredProviderCatalogProvider::new( "provider-claude-cli-local-1".to_string(), "claude".to_string(), Some("https://example.com".to_string()), "custom".to_string(), ) .expect("provider should build") .with_transport_fields( true, false, false, None, Some(2), Some(serde_json::json!({"url":"http://provider-proxy.internal:8080"})), Some(20.0), None, None, ) } fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint { StoredProviderCatalogEndpoint::new( "endpoint-claude-cli-local-1".to_string(), "provider-claude-cli-local-1".to_string(), "claude:messages".to_string(), Some("claude".to_string()), Some("cli".to_string()), true, ) .expect("endpoint should build") .with_transport_fields( "https://api.anthropic.example".to_string(), Some(serde_json::json!([ {"action":"set","key":"x-endpoint-tag","value":"claude-cli-local"} ])), Some(serde_json::json!([ {"action":"set","path":"metadata.mode","value":"safe"}, {"action":"rename","from":"metadata.client","to":"metadata.source"} ])), Some(2), Some("/custom/v1/messages".to_string()), None, None, None, ) .expect("endpoint transport should build") } fn sample_provider_catalog_key() -> StoredProviderCatalogKey { StoredProviderCatalogKey::new( "key-claude-cli-local-1".to_string(), "provider-claude-cli-local-1".to_string(), "prod".to_string(), "bearer".to_string(), None, true, ) .expect("key should build") .with_transport_fields( Some(serde_json::json!(["claude:messages"])), encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-claude-cli") .expect("api key should encrypt"), None, None, Some(serde_json::json!({"claude:messages": 1})), None, None, Some(serde_json::json!({"enabled": true, "node_id":"proxy-node-claude-cli-local"})), Some(serde_json::json!({"transport_profile":"chrome_136"})), ) .expect("key transport should build") } let seen_execution_runtime = Arc::new(Mutex::new(None::)); let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime); let seen_report = Arc::new(Mutex::new(false)); let seen_report_clone = Arc::clone(&seen_report); let decision_hits = Arc::new(Mutex::new(0usize)); let decision_hits_clone = Arc::clone(&decision_hits); let plan_hits = Arc::new(Mutex::new(0usize)); let plan_hits_clone = Arc::clone(&plan_hits); let public_hits = Arc::new(Mutex::new(0usize)); let public_hits_clone = Arc::clone(&public_hits); let upstream = Router::new() .route( "/api/internal/gateway/decision-stream", any(move |_request: Request| { let decision_hits_inner = Arc::clone(&decision_hits_clone); async move { *decision_hits_inner.lock().expect("mutex should lock") += 1; Json(json!({"action": "proxy_public"})) } }), ) .route( "/api/internal/gateway/plan-stream", any(move |_request: Request| { let plan_hits_inner = Arc::clone(&plan_hits_clone); async move { *plan_hits_inner.lock().expect("mutex should lock") += 1; Json(json!({"action": "proxy_public"})) } }), ) .route( "/api/internal/gateway/report-stream", any(move |request: Request| { let seen_report_inner = Arc::clone(&seen_report_clone); async move { let (_parts, body) = request.into_parts(); let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read"); *seen_report_inner.lock().expect("mutex should lock") = true; Json(json!({"ok": true})) } }), ) .route( "/v1/messages", any(move |_request: Request| { let public_hits_inner = Arc::clone(&public_hits_clone); async move { *public_hits_inner.lock().expect("mutex should lock") += 1; (StatusCode::IM_A_TEAPOT, Body::from("public-route-hit")) } }), ); let execution_runtime = Router::new().route( "/v1/execute/stream", any(move |request: Request| { let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone); async move { let (parts, body) = request.into_parts(); let raw_body = to_bytes(body, usize::MAX).await.expect("body should read"); let payload: serde_json::Value = serde_json::from_slice(&raw_body).expect("execution runtime payload should parse"); *seen_execution_runtime_inner.lock().expect("mutex should lock") = Some(SeenExecutionRuntimeStreamRequest { trace_id: parts .headers .get(TRACE_ID_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), url: payload .get("url") .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), model: payload .get("body") .and_then(|value| value.get("json_body")) .and_then(|value| value.get("model")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), stream: payload .get("body") .and_then(|value| value.get("json_body")) .and_then(|value| value.get("stream")) .and_then(|value| value.as_bool()) .unwrap_or(false), accept: payload .get("headers") .and_then(|value| value.get("accept")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), authorization: payload .get("headers") .and_then(|value| value.get("authorization")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), endpoint_tag: payload .get("headers") .and_then(|value| value.get("x-endpoint-tag")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), metadata_mode: payload .get("body") .and_then(|value| value.get("json_body")) .and_then(|value| value.get("metadata")) .and_then(|value| value.get("mode")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), metadata_source: payload .get("body") .and_then(|value| value.get("json_body")) .and_then(|value| value.get("metadata")) .and_then(|value| value.get("source")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), proxy_node_id: payload .get("proxy") .and_then(|value| value.get("node_id")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), transport_profile_id: payload .get("transport_profile").and_then(|value| value.get("profile_id")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), }); let body_stream = async_stream::stream! { yield Ok::(Bytes::from_static( b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n" )); yield Ok::(Bytes::from_static( b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: message_start\\ndata: {\\\"type\\\":\\\"message_start\\\"}\\n\\n\"}}\n" )); tokio::time::sleep(std::time::Duration::from_millis(250)).await; yield Ok::(Bytes::from_static( b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: message_stop\\ndata: {\\\"type\\\":\\\"message_stop\\\"}\\n\\n\"}}\n" )); yield Ok::(Bytes::from_static( b"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":31,\"ttfb_ms\":11,\"upstream_bytes\":37}}}\n" )); yield Ok::(Bytes::from_static( b"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n" )); }; let mut response = Response::builder() .status(StatusCode::OK) .body(Body::from_stream(body_stream)) .expect("response should build"); response.headers_mut().insert( http::header::CONTENT_TYPE, HeaderValue::from_static("application/x-ndjson"), ); response } }), ); let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key("sk-client-claude-cli-local")), sample_auth_snapshot("api-key-claude-cli-local-1", "user-claude-cli-local-1"), )])); let candidate_selection_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ sample_candidate_row(), ])); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_provider_catalog_provider()], vec![sample_provider_catalog_endpoint()], vec![sample_provider_catalog_key()], )); let (upstream_url, upstream_handle) = start_server(upstream).await; let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone()) .with_data_state_for_tests( crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests( auth_repository, candidate_selection_repository, provider_catalog_repository, Arc::clone(&request_candidate_repository), DEVELOPMENT_ENCRYPTION_KEY, ) .attach_proxy_node_repository_for_tests( crate::tests::ai_execute::ai_execute_proxy_node_repository([ "proxy-node-claude-cli-local", ]), ), ); let gateway = build_router_with_state(gateway_state); let (gateway_url, gateway_handle) = start_server(gateway).await; let mut response = reqwest::Client::new() .post(format!("{gateway_url}/v1/messages")) .header(http::header::CONTENT_TYPE, "application/json") .header(http::header::USER_AGENT, "Claude-Code/2.1.0") .header( http::header::AUTHORIZATION, "Bearer sk-client-claude-cli-local", ) .header(TRACE_ID_HEADER, "trace-claude-cli-local-stream-123") .body( "{\"model\":\"claude-code\",\"messages\":[],\"stream\":true,\"metadata\":{\"client\":\"desktop-claude-cli\"}}", ) .send() .await .expect("request should succeed"); assert_eq!(response.status(), StatusCode::OK); assert_eq!( tokio::time::timeout( std::time::Duration::from_millis(100), next_non_keepalive_chunk(&mut response), ) .await .expect("same-format passthrough should yield first chunk before eof"), Bytes::from_static(b"event: message_start\ndata: {\"type\":\"message_start\"}\n\n") ); assert_eq!( response.text().await.expect("remaining body should read"), "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n" ); let seen_execution_runtime_request = seen_execution_runtime .lock() .expect("mutex should lock") .clone() .expect("execution runtime stream should be captured"); assert_eq!( seen_execution_runtime_request.trace_id, "trace-claude-cli-local-stream-123" ); assert_eq!( seen_execution_runtime_request.url, "https://api.anthropic.example/custom/v1/messages" ); assert_eq!(seen_execution_runtime_request.model, "claude-code-upstream"); assert!(seen_execution_runtime_request.stream); assert_eq!(seen_execution_runtime_request.accept, "text/event-stream"); assert_eq!( seen_execution_runtime_request.authorization, "Bearer sk-upstream-claude-cli" ); assert_eq!( seen_execution_runtime_request.endpoint_tag, "claude-cli-local" ); assert_eq!(seen_execution_runtime_request.metadata_mode, "safe"); assert_eq!( seen_execution_runtime_request.metadata_source, "desktop-claude-cli" ); assert_eq!( seen_execution_runtime_request.proxy_node_id, "proxy-node-claude-cli-local" ); assert_eq!( seen_execution_runtime_request.transport_profile_id, "chrome_136" ); let stored_candidates = request_candidate_repository .list_by_request_id("trace-claude-cli-local-stream-123") .await .expect("request candidate trace should read"); assert_eq!(stored_candidates.len(), 1); assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success); tokio::time::sleep(std::time::Duration::from_millis(100)).await; assert!( !*seen_report.lock().expect("mutex should lock"), "report-stream should stay local when request candidate persistence is available" ); assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0); assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0); assert_eq!(*public_hits.lock().expect("mutex should lock"), 0); gateway_handle.abort(); execution_runtime_handle.abort(); upstream_handle.abort(); } #[test] fn gateway_executes_claude_code_cli_stream_via_local_decision_gate_with_local_stream_decision() { run_stream_provider_test( "gateway_executes_claude_code_cli_stream_via_local_decision_gate_with_local_stream_decision", gateway_executes_claude_code_cli_stream_via_local_decision_gate_with_local_stream_decision_impl, ); } async fn gateway_executes_claude_code_cli_stream_via_local_decision_gate_with_local_stream_decision_impl( ) { #[derive(Debug, Clone)] struct SeenExecutionRuntimeStreamRequest { trace_id: String, url: String, model: String, stream: bool, accept: String, authorization: String, anthropic_version: String, anthropic_beta: String, x_app: String, x_stainless_helper_method: String, user_agent: String, endpoint_tag: String, metadata_mode: String, metadata_source: String, assistant_content: serde_json::Value, proxy_node_id: String, transport_profile_id: String, } 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) -> 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!(["claude", "claude_code"])), Some(serde_json::json!(["claude:messages"])), Some(serde_json::json!(["claude-code"])), 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!(["claude", "claude_code"])), Some(serde_json::json!(["claude:messages"])), Some(serde_json::json!(["claude-code"])), ) .expect("auth snapshot should build") } fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow { StoredMinimalCandidateSelectionRow { provider_id: "provider-claude-code-cli-local-1".to_string(), provider_name: "claude_code".to_string(), provider_type: "claude_code".to_string(), provider_priority: 10, provider_is_active: true, endpoint_id: "endpoint-claude-code-cli-local-1".to_string(), endpoint_api_format: "claude:messages".to_string(), endpoint_api_family: Some("claude".to_string()), endpoint_kind: Some("cli".to_string()), endpoint_is_active: true, key_id: "key-claude-code-cli-local-1".to_string(), key_name: "prod".to_string(), key_auth_type: "oauth".to_string(), key_is_active: true, key_api_formats: Some(vec!["claude:messages".to_string()]), key_allowed_models: None, key_capabilities: None, key_internal_priority: 5, key_global_priority_by_format: Some(serde_json::json!({"claude:messages": 1})), model_id: "model-claude-code-cli-local-1".to_string(), global_model_id: "global-model-claude-code-cli-local-1".to_string(), global_model_name: "claude-code".to_string(), global_model_mappings: None, global_model_supports_streaming: Some(true), model_provider_model_name: "claude-code-upstream".to_string(), model_provider_model_mappings: Some(vec![StoredProviderModelMapping { name: "claude-code-upstream".to_string(), priority: 1, api_formats: Some(vec!["claude:messages".to_string()]), endpoint_ids: None, operations: None, }]), model_supports_streaming: Some(true), model_is_active: true, model_is_available: true, } } fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider { StoredProviderCatalogProvider::new( "provider-claude-code-cli-local-1".to_string(), "claude_code".to_string(), Some("https://example.com".to_string()), "claude_code".to_string(), ) .expect("provider should build") .with_transport_fields( true, false, false, None, Some(2), Some(serde_json::json!({"url":"http://provider-proxy.internal:8080"})), Some(20.0), None, Some(serde_json::json!({ "claude_code_advanced": { "cli_only_enabled": false } })), ) } fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint { StoredProviderCatalogEndpoint::new( "endpoint-claude-code-cli-local-1".to_string(), "provider-claude-code-cli-local-1".to_string(), "claude:messages".to_string(), Some("claude".to_string()), Some("cli".to_string()), true, ) .expect("endpoint should build") .with_transport_fields( "https://api.anthropic.example/v1/messages".to_string(), Some(serde_json::json!([ {"action":"set","key":"x-endpoint-tag","value":"claude-code-cli-local"} ])), Some(serde_json::json!([ {"action":"set","path":"metadata.mode","value":"safe"}, {"action":"rename","from":"metadata.client","to":"metadata.source"} ])), Some(2), None, None, None, None, ) .expect("endpoint transport should build") } fn sample_provider_catalog_key() -> StoredProviderCatalogKey { StoredProviderCatalogKey::new( "key-claude-code-cli-local-1".to_string(), "provider-claude-code-cli-local-1".to_string(), "prod".to_string(), "oauth".to_string(), None, true, ) .expect("key should build") .with_transport_fields( Some(serde_json::json!(["claude:messages"])), encrypt_python_fernet_plaintext( DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-claude-code-oauth", ) .expect("api key should encrypt"), None, None, Some(serde_json::json!({"claude:messages": 1})), None, None, Some( serde_json::json!({"enabled": true, "node_id":"proxy-node-claude-code-cli-local"}), ), Some(serde_json::json!({ "transport_profile": { "profile_id": "claude_code_nodejs", "header_fingerprint": { "user_agent":"Claude-Code/9.9" } } })), ) .expect("key transport should build") } let seen_execution_runtime = Arc::new(Mutex::new(None::)); let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime); let seen_report = Arc::new(Mutex::new(false)); let seen_report_clone = Arc::clone(&seen_report); let decision_hits = Arc::new(Mutex::new(0usize)); let decision_hits_clone = Arc::clone(&decision_hits); let plan_hits = Arc::new(Mutex::new(0usize)); let plan_hits_clone = Arc::clone(&plan_hits); let public_hits = Arc::new(Mutex::new(0usize)); let public_hits_clone = Arc::clone(&public_hits); let upstream = Router::new() .route( "/api/internal/gateway/decision-stream", any(move |_request: Request| { let decision_hits_inner = Arc::clone(&decision_hits_clone); async move { *decision_hits_inner.lock().expect("mutex should lock") += 1; Json(json!({"action": "proxy_public"})) } }), ) .route( "/api/internal/gateway/plan-stream", any(move |_request: Request| { let plan_hits_inner = Arc::clone(&plan_hits_clone); async move { *plan_hits_inner.lock().expect("mutex should lock") += 1; Json(json!({"action": "proxy_public"})) } }), ) .route( "/api/internal/gateway/report-stream", any(move |request: Request| { let seen_report_inner = Arc::clone(&seen_report_clone); async move { let (_parts, body) = request.into_parts(); let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read"); *seen_report_inner.lock().expect("mutex should lock") = true; Json(json!({"ok": true})) } }), ) .route( "/v1/messages", any(move |_request: Request| { let public_hits_inner = Arc::clone(&public_hits_clone); async move { *public_hits_inner.lock().expect("mutex should lock") += 1; (StatusCode::IM_A_TEAPOT, Body::from("public-route-hit")) } }), ); let execution_runtime = Router::new().route( "/v1/execute/stream", any(move |request: Request| { let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone); async move { let (parts, body) = request.into_parts(); let raw_body = to_bytes(body, usize::MAX).await.expect("body should read"); let payload: serde_json::Value = serde_json::from_slice(&raw_body).expect("execution runtime payload should parse"); *seen_execution_runtime_inner.lock().expect("mutex should lock") = Some(SeenExecutionRuntimeStreamRequest { trace_id: parts .headers .get(TRACE_ID_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), url: payload .get("url") .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), model: payload .get("body") .and_then(|value| value.get("json_body")) .and_then(|value| value.get("model")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), stream: payload .get("body") .and_then(|value| value.get("json_body")) .and_then(|value| value.get("stream")) .and_then(|value| value.as_bool()) .unwrap_or(false), accept: payload .get("headers") .and_then(|value| value.get("accept")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), authorization: payload .get("headers") .and_then(|value| value.get("authorization")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), anthropic_version: payload .get("headers") .and_then(|value| value.get("anthropic-version")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), anthropic_beta: payload .get("headers") .and_then(|value| value.get("anthropic-beta")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), x_app: payload .get("headers") .and_then(|value| value.get("x-app")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), x_stainless_helper_method: payload .get("headers") .and_then(|value| value.get("x-stainless-helper-method")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), user_agent: payload .get("headers") .and_then(|value| value.get("user-agent")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), endpoint_tag: payload .get("headers") .and_then(|value| value.get("x-endpoint-tag")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), metadata_mode: payload .get("body") .and_then(|value| value.get("json_body")) .and_then(|value| value.get("metadata")) .and_then(|value| value.get("mode")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), metadata_source: payload .get("body") .and_then(|value| value.get("json_body")) .and_then(|value| value.get("metadata")) .and_then(|value| value.get("source")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), assistant_content: payload .get("body") .and_then(|value| value.get("json_body")) .and_then(|value| value.get("messages")) .and_then(|value| value.get(0)) .and_then(|value| value.get("content")) .cloned() .unwrap_or(serde_json::Value::Null), proxy_node_id: payload .get("proxy") .and_then(|value| value.get("node_id")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), transport_profile_id: payload .get("transport_profile").and_then(|value| value.get("profile_id")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), }); let frames = concat!( "{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n", "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: message_start\\ndata: {\\\"type\\\":\\\"message_start\\\"}\\n\\n\"}}\n", "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: message_stop\\ndata: {\\\"type\\\":\\\"message_stop\\\"}\\n\\n\"}}\n", "{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":31,\"ttfb_ms\":11,\"upstream_bytes\":37}}}\n", "{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n" ); let mut response = Response::builder() .status(StatusCode::OK) .body(Body::from(frames)) .expect("response should build"); response.headers_mut().insert( http::header::CONTENT_TYPE, HeaderValue::from_static("application/x-ndjson"), ); response } }), ); let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key("sk-client-claude-code-cli-local")), sample_auth_snapshot( "api-key-claude-code-cli-local-1", "user-claude-code-cli-local-1", ), )])); let candidate_selection_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ sample_candidate_row(), ])); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_provider_catalog_provider()], vec![sample_provider_catalog_endpoint()], vec![sample_provider_catalog_key()], )); let (upstream_url, upstream_handle) = start_server(upstream).await; let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone()) .with_data_state_for_tests( crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests( auth_repository, candidate_selection_repository, provider_catalog_repository, Arc::clone(&request_candidate_repository), DEVELOPMENT_ENCRYPTION_KEY, ) .attach_proxy_node_repository_for_tests( crate::tests::ai_execute::ai_execute_proxy_node_repository([ "proxy-node-claude-code-cli-local", ]), ), ); let gateway = build_router_with_state(gateway_state); let (gateway_url, gateway_handle) = start_server(gateway).await; let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/messages")) .header(http::header::CONTENT_TYPE, "application/json") .header(http::header::USER_AGENT, "Claude-Code/2.1.0") .header( http::header::AUTHORIZATION, "Bearer sk-client-claude-code-cli-local", ) .header("anthropic-beta", "context-1m-2025-08-07,custom-beta") .header(TRACE_ID_HEADER, "trace-claude-code-cli-local-stream-123") .body( serde_json::json!({ "model":"claude-code", "stream":true, "thinking":{"type":"enabled"}, "messages":[{ "role":"assistant", "content":[ {"type":"thinking","thinking":"keep","signature":"sig_valid"}, {"type":"thinking","thinking":"drop-empty-signature","signature":""}, {"type":"redacted_thinking","data":"keep-redacted","signature":"sig_redacted"}, {"type":"redacted_thinking","data":"drop-no-signature"}, {"type":"text","text":"ok"} ] }], "metadata":{"client":"desktop-claude-code-cli"} }) .to_string(), ) .send() .await .expect("request should succeed"); assert_eq!(response.status(), StatusCode::OK); assert_eq!( strip_sse_keepalive_comments(&response.text().await.expect("body should read")), concat!( "event: message_start\ndata: {\"type\":\"message_start\"}\n\n", "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", ) ); let seen_execution_runtime_request = seen_execution_runtime .lock() .expect("mutex should lock") .clone() .expect("execution runtime stream should be captured"); assert_eq!( seen_execution_runtime_request.trace_id, "trace-claude-code-cli-local-stream-123" ); assert_eq!( seen_execution_runtime_request.url, "https://api.anthropic.example/v1/messages" ); assert_eq!(seen_execution_runtime_request.model, "claude-code-upstream"); assert!(seen_execution_runtime_request.stream); assert_eq!(seen_execution_runtime_request.accept, "text/event-stream"); assert_eq!( seen_execution_runtime_request.authorization, "Bearer sk-upstream-claude-code-oauth" ); assert_eq!( seen_execution_runtime_request.anthropic_version, "2023-06-01" ); assert_eq!( seen_execution_runtime_request.anthropic_beta, "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,prompt-caching-scope-2026-01-05,effort-2025-11-24,context-management-2025-06-27,extended-cache-ttl-2025-04-11,context-1m-2025-08-07,custom-beta" ); assert_eq!(seen_execution_runtime_request.x_app, "cli"); assert_eq!( seen_execution_runtime_request.x_stainless_helper_method, "stream" ); assert_eq!( seen_execution_runtime_request.user_agent, "claude-cli/2.1.284 (external, cli)" ); assert_eq!( seen_execution_runtime_request.endpoint_tag, "claude-code-cli-local" ); assert_eq!(seen_execution_runtime_request.metadata_mode, "safe"); assert_eq!( seen_execution_runtime_request.metadata_source, "desktop-claude-code-cli" ); assert_eq!( seen_execution_runtime_request.assistant_content, json!([ {"type":"thinking","thinking":"keep","signature":"sig_valid"}, {"type":"redacted_thinking","data":"keep-redacted","signature":"sig_redacted"}, {"type":"text","text":"ok"} ]) ); assert_eq!( seen_execution_runtime_request.proxy_node_id, "proxy-node-claude-code-cli-local" ); assert_eq!( seen_execution_runtime_request.transport_profile_id, "claude_code_nodejs" ); let stored_candidates = request_candidate_repository .list_by_request_id("trace-claude-code-cli-local-stream-123") .await .expect("request candidate trace should read"); assert_eq!(stored_candidates.len(), 1); assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success); tokio::time::sleep(std::time::Duration::from_millis(100)).await; assert!( !*seen_report.lock().expect("mutex should lock"), "report-stream should stay local when request candidate persistence is available" ); assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0); assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0); assert_eq!(*public_hits.lock().expect("mutex should lock"), 0); gateway_handle.abort(); execution_runtime_handle.abort(); upstream_handle.abort(); } #[test] fn gateway_executes_claude_chat_stream_via_local_decision_gate_with_local_stream_decision() { run_stream_provider_test( "gateway_executes_claude_chat_stream_via_local_decision_gate_with_local_stream_decision", gateway_executes_claude_chat_stream_via_local_decision_gate_with_local_stream_decision_impl, ); } async fn gateway_executes_claude_chat_stream_via_local_decision_gate_with_local_stream_decision_impl( ) { #[derive(Debug, Clone)] struct SeenExecutionRuntimeStreamRequest { trace_id: String, url: String, model: String, stream: bool, accept: String, auth_header_value: String, endpoint_tag: String, metadata_mode: String, metadata_source: String, proxy_node_id: String, transport_profile_id: String, } 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) -> 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!(["claude"])), Some(serde_json::json!(["claude:messages"])), Some(serde_json::json!(["claude-sonnet-4-5"])), 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!(["claude"])), Some(serde_json::json!(["claude:messages"])), Some(serde_json::json!(["claude-sonnet-4-5"])), ) .expect("auth snapshot should build") } fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow { StoredMinimalCandidateSelectionRow { provider_id: "provider-claude-chat-local-stream-1".to_string(), provider_name: "claude".to_string(), provider_type: "custom".to_string(), provider_priority: 10, provider_is_active: true, endpoint_id: "endpoint-claude-chat-local-stream-1".to_string(), endpoint_api_format: "claude:messages".to_string(), endpoint_api_family: Some("claude".to_string()), endpoint_kind: Some("chat".to_string()), endpoint_is_active: true, key_id: "key-claude-chat-local-stream-1".to_string(), key_name: "prod".to_string(), key_auth_type: "api_key".to_string(), key_is_active: true, key_api_formats: Some(vec!["claude:messages".to_string()]), key_allowed_models: None, key_capabilities: None, key_internal_priority: 5, key_global_priority_by_format: Some(serde_json::json!({"claude:messages": 1})), model_id: "model-claude-chat-local-stream-1".to_string(), global_model_id: "global-model-claude-chat-local-stream-1".to_string(), global_model_name: "claude-sonnet-4-5".to_string(), global_model_mappings: None, global_model_supports_streaming: Some(true), model_provider_model_name: "claude-sonnet-4-5-upstream".to_string(), model_provider_model_mappings: Some(vec![StoredProviderModelMapping { name: "claude-sonnet-4-5-upstream".to_string(), priority: 1, api_formats: Some(vec!["claude:messages".to_string()]), endpoint_ids: None, operations: None, }]), model_supports_streaming: Some(true), model_is_active: true, model_is_available: true, } } fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider { StoredProviderCatalogProvider::new( "provider-claude-chat-local-stream-1".to_string(), "claude".to_string(), Some("https://example.com".to_string()), "custom".to_string(), ) .expect("provider should build") .with_transport_fields( true, false, false, None, Some(2), Some(serde_json::json!({"url":"http://provider-proxy.internal:8080"})), Some(20.0), None, None, ) } fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint { StoredProviderCatalogEndpoint::new( "endpoint-claude-chat-local-stream-1".to_string(), "provider-claude-chat-local-stream-1".to_string(), "claude:messages".to_string(), Some("claude".to_string()), Some("chat".to_string()), true, ) .expect("endpoint should build") .with_transport_fields( "https://api.anthropic.example".to_string(), Some(serde_json::json!([ {"action":"set","key":"x-endpoint-tag","value":"claude-chat-local-stream"} ])), Some(serde_json::json!([ {"action":"set","path":"metadata.mode","value":"safe"}, {"action":"rename","from":"metadata.client","to":"metadata.source"} ])), Some(2), Some("/custom/v1/messages".to_string()), None, None, None, ) .expect("endpoint transport should build") } fn sample_provider_catalog_key() -> StoredProviderCatalogKey { StoredProviderCatalogKey::new( "key-claude-chat-local-stream-1".to_string(), "provider-claude-chat-local-stream-1".to_string(), "prod".to_string(), "api_key".to_string(), None, true, ) .expect("key should build") .with_transport_fields( Some(serde_json::json!(["claude:messages"])), encrypt_python_fernet_plaintext( DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-claude-chat-stream", ) .expect("api key should encrypt"), None, None, Some(serde_json::json!({"claude:messages": 1})), None, None, Some(serde_json::json!({"enabled": true, "node_id":"proxy-node-claude-chat-stream"})), Some(serde_json::json!({"transport_profile":"chrome_136"})), ) .expect("key transport should build") } let seen_execution_runtime = Arc::new(Mutex::new(None::)); let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime); let seen_report = Arc::new(Mutex::new(false)); let seen_report_clone = Arc::clone(&seen_report); let decision_hits = Arc::new(Mutex::new(0usize)); let decision_hits_clone = Arc::clone(&decision_hits); let plan_hits = Arc::new(Mutex::new(0usize)); let plan_hits_clone = Arc::clone(&plan_hits); let public_hits = Arc::new(Mutex::new(0usize)); let public_hits_clone = Arc::clone(&public_hits); let upstream = Router::new() .route( "/api/internal/gateway/decision-stream", any(move |_request: Request| { let decision_hits_inner = Arc::clone(&decision_hits_clone); async move { *decision_hits_inner.lock().expect("mutex should lock") += 1; Json(json!({"action": "proxy_public"})) } }), ) .route( "/api/internal/gateway/plan-stream", any(move |_request: Request| { let plan_hits_inner = Arc::clone(&plan_hits_clone); async move { *plan_hits_inner.lock().expect("mutex should lock") += 1; Json(json!({"action": "proxy_public"})) } }), ) .route( "/api/internal/gateway/report-stream", any(move |request: Request| { let seen_report_inner = Arc::clone(&seen_report_clone); async move { let (_parts, body) = request.into_parts(); let _raw_body = to_bytes(body, usize::MAX).await.expect("body should read"); *seen_report_inner.lock().expect("mutex should lock") = true; Json(json!({"ok": true})) } }), ) .route( "/v1/messages", any(move |_request: Request| { let public_hits_inner = Arc::clone(&public_hits_clone); async move { *public_hits_inner.lock().expect("mutex should lock") += 1; (StatusCode::IM_A_TEAPOT, Body::from("public-route-hit")) } }), ); let execution_runtime = Router::new().route( "/v1/execute/stream", any(move |request: Request| { let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone); async move { let (parts, body) = request.into_parts(); let raw_body = to_bytes(body, usize::MAX).await.expect("body should read"); let payload: serde_json::Value = serde_json::from_slice(&raw_body).expect("execution runtime payload should parse"); *seen_execution_runtime_inner.lock().expect("mutex should lock") = Some(SeenExecutionRuntimeStreamRequest { trace_id: parts .headers .get(TRACE_ID_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), url: payload .get("url") .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), model: payload .get("body") .and_then(|value| value.get("json_body")) .and_then(|value| value.get("model")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), stream: payload .get("body") .and_then(|value| value.get("json_body")) .and_then(|value| value.get("stream")) .and_then(|value| value.as_bool()) .unwrap_or(false), accept: payload .get("headers") .and_then(|value| value.get("accept")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), auth_header_value: payload .get("headers") .and_then(|value| value.get("x-api-key")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), endpoint_tag: payload .get("headers") .and_then(|value| value.get("x-endpoint-tag")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), metadata_mode: payload .get("body") .and_then(|value| value.get("json_body")) .and_then(|value| value.get("metadata")) .and_then(|value| value.get("mode")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), metadata_source: payload .get("body") .and_then(|value| value.get("json_body")) .and_then(|value| value.get("metadata")) .and_then(|value| value.get("source")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), proxy_node_id: payload .get("proxy") .and_then(|value| value.get("node_id")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), transport_profile_id: payload .get("transport_profile").and_then(|value| value.get("profile_id")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), }); let frames = concat!( "{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n", "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: message_start\\ndata: {\\\"type\\\":\\\"message_start\\\"}\\n\\n\"}}\n", "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: message_stop\\ndata: {\\\"type\\\":\\\"message_stop\\\"}\\n\\n\"}}\n", "{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":31,\"ttfb_ms\":11,\"upstream_bytes\":37}}}\n", "{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n" ); let mut response = Response::builder() .status(StatusCode::OK) .body(Body::from(frames)) .expect("response should build"); response.headers_mut().insert( http::header::CONTENT_TYPE, HeaderValue::from_static("application/x-ndjson"), ); response } }), ); let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key("sk-client-claude-chat-stream-local")), sample_auth_snapshot( "api-key-claude-chat-local-stream-1", "user-claude-chat-local-stream-1", ), )])); let candidate_selection_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ sample_candidate_row(), ])); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_provider_catalog_provider()], vec![sample_provider_catalog_endpoint()], vec![sample_provider_catalog_key()], )); let (upstream_url, upstream_handle) = start_server(upstream).await; let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone()) .with_data_state_for_tests( crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests( auth_repository, candidate_selection_repository, provider_catalog_repository, Arc::clone(&request_candidate_repository), DEVELOPMENT_ENCRYPTION_KEY, ) .attach_proxy_node_repository_for_tests( crate::tests::ai_execute::ai_execute_proxy_node_repository([ "proxy-node-claude-chat-stream", ]), ), ); let gateway = build_router_with_state(gateway_state); let (gateway_url, gateway_handle) = start_server(gateway).await; let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/messages")) .header(http::header::CONTENT_TYPE, "application/json") .header("x-api-key", "sk-client-claude-chat-stream-local") .header(TRACE_ID_HEADER, "trace-claude-chat-local-stream-123") .body( "{\"model\":\"claude-sonnet-4-5\",\"messages\":[],\"stream\":true,\"metadata\":{\"client\":\"desktop-claude-stream\"}}", ) .send() .await .expect("request should succeed"); assert_eq!(response.status(), StatusCode::OK); assert_eq!( strip_sse_keepalive_comments(&response.text().await.expect("body should read")), concat!( "event: message_start\ndata: {\"type\":\"message_start\"}\n\n", "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n", ) ); let seen_execution_runtime_request = seen_execution_runtime .lock() .expect("mutex should lock") .clone() .expect("execution runtime stream should be captured"); assert_eq!( seen_execution_runtime_request.trace_id, "trace-claude-chat-local-stream-123" ); assert_eq!( seen_execution_runtime_request.url, "https://api.anthropic.example/custom/v1/messages" ); assert_eq!( seen_execution_runtime_request.model, "claude-sonnet-4-5-upstream" ); assert!(seen_execution_runtime_request.stream); assert_eq!(seen_execution_runtime_request.accept, "text/event-stream"); assert_eq!( seen_execution_runtime_request.auth_header_value, "sk-upstream-claude-chat-stream" ); assert_eq!( seen_execution_runtime_request.endpoint_tag, "claude-chat-local-stream" ); assert_eq!(seen_execution_runtime_request.metadata_mode, "safe"); assert_eq!( seen_execution_runtime_request.metadata_source, "desktop-claude-stream" ); assert_eq!( seen_execution_runtime_request.proxy_node_id, "proxy-node-claude-chat-stream" ); assert_eq!( seen_execution_runtime_request.transport_profile_id, "chrome_136" ); let stored_candidates = request_candidate_repository .list_by_request_id("trace-claude-chat-local-stream-123") .await .expect("request candidate trace should read"); assert_eq!(stored_candidates.len(), 1); assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success); tokio::time::sleep(std::time::Duration::from_millis(100)).await; assert!( !*seen_report.lock().expect("mutex should lock"), "report-stream should stay local when request candidate persistence is available" ); assert_eq!(*decision_hits.lock().expect("mutex should lock"), 0); assert_eq!(*plan_hits.lock().expect("mutex should lock"), 0); assert_eq!(*public_hits.lock().expect("mutex should lock"), 0); gateway_handle.abort(); execution_runtime_handle.abort(); upstream_handle.abort(); }