use super::{ any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server, to_bytes, Arc, Body, Json, Mutex, Request, 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::provider_catalog::{ StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use sha2::{Digest, Sha256}; const STREAM_IMAGE_TEST_STACK_BYTES: usize = 16 * 1024 * 1024; fn run_stream_image_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_IMAGE_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 image test thread should spawn"); if let Err(payload) = handle.join() { std::panic::resume_unwind(payload); } } #[test] fn gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth_refresh() { run_stream_image_test( "gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth_refresh", gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth_refresh_impl, ); } async fn gateway_executes_codex_image_stream_via_local_decision_gate_after_oauth_refresh_impl() { #[derive(Debug, Clone)] struct SeenExecutionRuntimeStreamRequest { trace_id: String, url: String, authorization: String, headers: serde_json::Value, body: serde_json::Value, plan_stream: bool, } #[derive(Debug, Clone)] struct SeenRefreshRequest { content_type: String, body: 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!(["openai", "codex"])), Some(serde_json::json!(["openai:image"])), Some(serde_json::json!(["gpt-image-2"])), api_key_id.to_string(), Some("default".to_string()), true, false, false, Some(60), Some(5), Some(4_102_444_800_i64), Some(serde_json::json!(["openai", "codex"])), Some(serde_json::json!(["openai:image"])), Some(serde_json::json!(["gpt-image-2"])), ) .expect("auth snapshot should build") } fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow { StoredMinimalCandidateSelectionRow { provider_id: "provider-codex-image-stream-local-1".to_string(), provider_name: "codex".to_string(), provider_type: "codex".to_string(), provider_priority: 10, provider_is_active: true, endpoint_id: "endpoint-codex-image-stream-local-1".to_string(), endpoint_api_format: "openai:image".to_string(), endpoint_api_family: Some("openai".to_string()), endpoint_kind: Some("image".to_string()), endpoint_is_active: true, key_id: "key-codex-image-stream-local-1".to_string(), key_name: "oauth".to_string(), key_auth_type: "oauth".to_string(), key_is_active: true, key_api_formats: Some(vec!["openai:image".to_string()]), key_allowed_models: None, key_capabilities: None, key_internal_priority: 5, key_global_priority_by_format: Some(serde_json::json!({"openai:image": 1})), model_id: "model-codex-image-stream-local-1".to_string(), global_model_id: "global-model-codex-image-stream-local-1".to_string(), global_model_name: "gpt-image-2".to_string(), global_model_mappings: None, global_model_supports_streaming: Some(true), model_provider_model_name: "gpt-image-2".to_string(), model_provider_model_mappings: Some(vec![StoredProviderModelMapping { name: "gpt-image-2".to_string(), priority: 1, api_formats: Some(vec!["openai:image".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-codex-image-stream-local-1".to_string(), "codex".to_string(), Some("https://chatgpt.com".to_string()), "codex".to_string(), ) .expect("provider should build") .with_transport_fields( true, false, false, None, Some(2), None, Some(20.0), None, None, ) } fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint { StoredProviderCatalogEndpoint::new( "endpoint-codex-image-stream-local-1".to_string(), "provider-codex-image-stream-local-1".to_string(), "openai:image".to_string(), Some("openai".to_string()), Some("image".to_string()), true, ) .expect("endpoint should build") .with_transport_fields( "https://chatgpt.com/backend-api/codex".to_string(), None, None, Some(2), None, None, None, None, ) .expect("endpoint transport should build") } fn sample_provider_catalog_key() -> StoredProviderCatalogKey { let encrypted_auth_config = encrypt_python_fernet_plaintext( DEVELOPMENT_ENCRYPTION_KEY, r#"{"provider_type":"codex","refresh_token":"rt-codex-image-stream-local-123"}"#, ) .expect("auth config should encrypt"); StoredProviderCatalogKey::new( "key-codex-image-stream-local-1".to_string(), "provider-codex-image-stream-local-1".to_string(), "oauth".to_string(), "oauth".to_string(), None, true, ) .expect("key should build") .with_transport_fields( Some(serde_json::json!(["openai:image"])), encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "__placeholder__") .expect("placeholder api key should encrypt"), Some(encrypted_auth_config), None, Some(serde_json::json!({"openai:image": 1})), None, None, None, None, ) .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_refresh = Arc::new(Mutex::new(None::)); let seen_refresh_clone = Arc::clone(&seen_refresh); let refresh = Router::new().route( "/oauth/token", any(move |request: Request| { let seen_refresh_inner = Arc::clone(&seen_refresh_clone); async move { let (parts, body) = request.into_parts(); let raw_body = to_bytes(body, usize::MAX).await.expect("body should read"); *seen_refresh_inner.lock().expect("mutex should lock") = Some(SeenRefreshRequest { content_type: parts .headers .get(http::header::CONTENT_TYPE) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), body: String::from_utf8(raw_body.to_vec()) .expect("refresh body should be utf8"), }); Json(json!({ "access_token": "refreshed-codex-image-stream-access-token", "refresh_token": "rt-codex-image-stream-local-456", "token_type": "Bearer", "expires_in": 3600 })) } }), ); 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(), headers: payload.get("headers").cloned().unwrap_or_default(), body: payload .get("body") .and_then(|value| value.get("json_body")) .cloned() .unwrap_or_default(), plan_stream: payload .get("stream") .and_then(|value| value.as_bool()) .unwrap_or(false), }); let frames = concat!( "{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"application/json\"}}}\n", "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"{\\\"created\\\":1776991097,\\\"data\\\":[{\\\"b64_json\\\":\\\"aGVsbG8=\\\",\\\"revised_prompt\\\":\\\"水墨视觉海报\\\"}],\\\"usage\\\":{\\\"input_tokens\\\":11,\\\"output_tokens\\\":22,\\\"total_tokens\\\":33}}\"}}\n", "{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41}}}\n", "{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n" ); let mut response = http::Response::builder() .status(StatusCode::OK) .body(Body::from(frames)) .expect("response should build"); response.headers_mut().insert( http::header::CONTENT_TYPE, http::HeaderValue::from_static("application/x-ndjson"), ); response } }), ); let client_api_key = "sk-client-codex-image-stream-local"; let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key(client_api_key)), sample_auth_snapshot( "key-codex-image-stream-client-123", "user-codex-image-stream-client-123", ), )])); let candidate_selection_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ sample_candidate_row(), ])); let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_provider_catalog_provider()], vec![sample_provider_catalog_endpoint()], vec![sample_provider_catalog_key()], )); let (refresh_url, refresh_handle) = start_server(refresh).await; let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; let oauth_refresh = crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![ Arc::new( crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default() .with_token_url_for_tests("codex", format!("{refresh_url}/oauth/token")), ), ]); 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::new(InMemoryRequestCandidateRepository::default()), DEVELOPMENT_ENCRYPTION_KEY, ), ) .with_oauth_refresh_coordinator_for_tests(oauth_refresh); 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/images/generations")) .header(http::header::CONTENT_TYPE, "application/json") .header( http::header::AUTHORIZATION, format!("Bearer {client_api_key}"), ) .header(TRACE_ID_HEADER, "trace-codex-image-stream-local-123") .body( "{\"model\":\"gpt-image-2\",\"prompt\":\"生成一张水墨视觉海报\",\"background\":\"auto\",\"quality\":\"auto\",\"size\":\"auto\",\"stream\":true,\"response_format\":\"b64_json\"}", ) .send() .await .expect("request should succeed"); let response_status = response.status(); let response_content_type = response .headers() .get(http::header::CONTENT_TYPE) .and_then(|value| value.to_str().ok()) .map(str::to_string); let response_text = response.text().await.expect("body should read"); assert_eq!(response_status, StatusCode::OK, "{response_text}"); assert_eq!( response_content_type.as_deref(), Some("text/event-stream"), "{response_text}" ); assert!(response_text.contains("event: image_generation.completed")); assert!(response_text.contains("\"type\":\"image_generation.completed\"")); assert!(response_text.contains("\"b64_json\":\"aGVsbG8=\"")); assert!(response_text.contains("\"total_tokens\":33")); assert!(!response_text.contains("response.completed")); let seen_refresh_request = seen_refresh .lock() .expect("mutex should lock") .clone() .expect("refresh request should be captured"); assert_eq!( seen_refresh_request.content_type, "application/x-www-form-urlencoded" ); assert!(seen_refresh_request .body .contains("refresh_token=rt-codex-image-stream-local-123")); 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-codex-image-stream-local-123" ); assert_eq!( seen_execution_runtime_request.url, "https://chatgpt.com/backend-api/codex/images/generations" ); assert_eq!( seen_execution_runtime_request.authorization, "Bearer refreshed-codex-image-stream-access-token" ); assert_eq!( seen_execution_runtime_request.body, json!({ "prompt": "生成一张水墨视觉海报", "background": "auto", "model": "gpt-image-2", "quality": "auto", "size": "auto" }) ); assert_eq!( seen_execution_runtime_request.headers["user-agent"], "codex_cli_rs/0.153.3" ); assert_eq!( seen_execution_runtime_request.headers["originator"], "codex_cli_rs" ); for header in ["x-client-request-id", "session-id", "thread-id"] { assert!(seen_execution_runtime_request.headers.get(header).is_none()); } assert!(!seen_execution_runtime_request.plan_stream); gateway_handle.abort(); execution_runtime_handle.abort(); refresh_handle.abort(); } #[test] fn gateway_bridges_codex_image_sync_json_to_streaming_image_sse() { run_stream_image_test( "gateway_bridges_codex_image_sync_json_to_streaming_image_sse", gateway_bridges_codex_image_sync_json_to_streaming_image_sse_impl, ); } async fn gateway_bridges_codex_image_sync_json_to_streaming_image_sse_impl() { #[derive(Debug, Clone)] struct SeenExecutionRuntimeStreamRequest { trace_id: String, request_stream: bool, plan_stream: bool, } 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!(["openai", "codex"])), Some(serde_json::json!(["openai:image"])), Some(serde_json::json!(["gpt-image-2"])), api_key_id.to_string(), Some("default".to_string()), true, false, false, Some(60), Some(5), Some(4_102_444_800_i64), Some(serde_json::json!(["openai", "codex"])), Some(serde_json::json!(["openai:image"])), Some(serde_json::json!(["gpt-image-2"])), ) .expect("auth snapshot should build") } fn sample_candidate_row() -> StoredMinimalCandidateSelectionRow { StoredMinimalCandidateSelectionRow { provider_id: "provider-codex-image-stream-local-1".to_string(), provider_name: "codex".to_string(), provider_type: "codex".to_string(), provider_priority: 10, provider_is_active: true, endpoint_id: "endpoint-codex-image-stream-local-1".to_string(), endpoint_api_format: "openai:image".to_string(), endpoint_api_family: Some("openai".to_string()), endpoint_kind: Some("image".to_string()), endpoint_is_active: true, key_id: "key-codex-image-stream-local-1".to_string(), key_name: "oauth".to_string(), key_auth_type: "oauth".to_string(), key_is_active: true, key_api_formats: Some(vec!["openai:image".to_string()]), key_allowed_models: None, key_capabilities: None, key_internal_priority: 5, key_global_priority_by_format: Some(serde_json::json!({"openai:image": 1})), model_id: "model-codex-image-stream-local-1".to_string(), global_model_id: "global-model-codex-image-stream-local-1".to_string(), global_model_name: "gpt-image-2".to_string(), global_model_mappings: None, global_model_supports_streaming: Some(true), model_provider_model_name: "gpt-image-2".to_string(), model_provider_model_mappings: Some(vec![StoredProviderModelMapping { name: "gpt-image-2".to_string(), priority: 1, api_formats: Some(vec!["openai:image".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-codex-image-stream-local-1".to_string(), "codex".to_string(), Some("https://chatgpt.com".to_string()), "codex".to_string(), ) .expect("provider should build") .with_transport_fields( true, false, false, None, Some(2), None, Some(20.0), None, None, ) } fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint { StoredProviderCatalogEndpoint::new( "endpoint-codex-image-stream-local-1".to_string(), "provider-codex-image-stream-local-1".to_string(), "openai:image".to_string(), Some("openai".to_string()), Some("image".to_string()), true, ) .expect("endpoint should build") .with_transport_fields( "https://chatgpt.com/backend-api/codex".to_string(), None, None, Some(2), None, None, None, None, ) .expect("endpoint transport should build") } fn sample_provider_catalog_key() -> StoredProviderCatalogKey { let encrypted_auth_config = encrypt_python_fernet_plaintext( DEVELOPMENT_ENCRYPTION_KEY, r#"{"provider_type":"codex","refresh_token":"rt-codex-image-stream-local-123"}"#, ) .expect("auth config should encrypt"); StoredProviderCatalogKey::new( "key-codex-image-stream-local-1".to_string(), "provider-codex-image-stream-local-1".to_string(), "oauth".to_string(), "oauth".to_string(), None, true, ) .expect("key should build") .with_transport_fields( Some(serde_json::json!(["openai:image"])), encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "__placeholder__") .expect("placeholder api key should encrypt"), Some(encrypted_auth_config), None, Some(serde_json::json!({"openai:image": 1})), None, None, None, None, ) .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 refresh = Router::new().route( "/oauth/token", any(move |_request: Request| async move { Json(json!({ "access_token": "refreshed-codex-image-stream-access-token", "refresh_token": "rt-codex-image-stream-local-456", "token_type": "Bearer", "expires_in": 3600 })) }), ); 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(), request_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), plan_stream: payload .get("stream") .and_then(|value| value.as_bool()) .unwrap_or(false), }); let frames = concat!( "{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"application/json\"}}}\n", "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"{\\\"created\\\":1776991097,\\\"data\\\":[{\\\"b64_json\\\":\\\"aGVsbG8=\\\"}],\\\"usage\\\":{\\\"total_tokens\\\":100,\\\"input_tokens\\\":50,\\\"output_tokens\\\":50,\\\"input_tokens_details\\\":{\\\"text_tokens\\\":10,\\\"image_tokens\\\":40}}}\"}}\n", "{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41}}}\n", "{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n" ); let mut response = http::Response::builder() .status(StatusCode::OK) .body(Body::from(frames)) .expect("response should build"); response.headers_mut().insert( http::header::CONTENT_TYPE, http::HeaderValue::from_static("application/x-ndjson"), ); response } }), ); let client_api_key = "sk-client-codex-image-stream-local"; let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(hash_api_key(client_api_key)), sample_auth_snapshot( "key-codex-image-stream-client-123", "user-codex-image-stream-client-123", ), )])); let candidate_selection_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ sample_candidate_row(), ])); let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_provider_catalog_provider()], vec![sample_provider_catalog_endpoint()], vec![sample_provider_catalog_key()], )); let (refresh_url, refresh_handle) = start_server(refresh).await; let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; let oauth_refresh = crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![ Arc::new( crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default() .with_token_url_for_tests("codex", format!("{refresh_url}/oauth/token")), ), ]); 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::new(InMemoryRequestCandidateRepository::default()), DEVELOPMENT_ENCRYPTION_KEY, ), ) .with_oauth_refresh_coordinator_for_tests(oauth_refresh); 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/images/generations")) .header(http::header::CONTENT_TYPE, "application/json") .header( http::header::AUTHORIZATION, format!("Bearer {client_api_key}"), ) .header(TRACE_ID_HEADER, "trace-codex-image-stream-json-123") .body("{\"model\":\"gpt-image-2\",\"prompt\":\"生成一张海报\",\"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 response_text = response.text().await.expect("body should read"); assert!(response_text.contains("event: image_generation.completed")); assert!(response_text.contains("\"type\":\"image_generation.completed\"")); assert!(response_text.contains("\"b64_json\":\"aGVsbG8=\"")); assert!(response_text.contains("\"total_tokens\":100")); assert!(!response_text.trim_start().starts_with('{')); assert!(!response_text.contains("\"created\": 1776991097")); 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-codex-image-stream-json-123" ); assert!(!seen_execution_runtime_request.request_stream); assert!(!seen_execution_runtime_request.plan_stream); gateway_handle.abort(); execution_runtime_handle.abort(); refresh_handle.abort(); } #[derive(Debug, Clone)] struct SeenImageBridgeExecutionPlan { trace_id: String, client_api_format: String, provider_api_format: String, url: String, plan_stream: bool, auth_header: String, body_json: serde_json::Value, } fn image_bridge_hash_api_key(value: &str) -> String { let mut hasher = Sha256::new(); hasher.update(value.as_bytes()); format!("{:x}", hasher.finalize()) } fn image_bridge_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, None, Some(serde_json::json!([ "openai:chat", "openai:responses", "openai:image" ])), Some(serde_json::json!(["gpt-image-2"])), api_key_id.to_string(), Some("default".to_string()), true, false, false, Some(60), Some(5), Some(4_102_444_800_i64), None, Some(serde_json::json!([ "openai:chat", "openai:responses", "openai:image" ])), Some(serde_json::json!(["gpt-image-2"])), ) .expect("auth snapshot should build") } fn image_bridge_candidate_row( prefix: &str, provider_name: &str, provider_type: &str, ) -> StoredMinimalCandidateSelectionRow { let key_auth_type = if provider_type == "chatgpt_web" { "bearer" } else { "api_key" }; StoredMinimalCandidateSelectionRow { provider_id: format!("provider-{prefix}"), provider_name: provider_name.to_string(), provider_type: provider_type.to_string(), provider_priority: 10, provider_is_active: true, endpoint_id: format!("endpoint-{prefix}"), endpoint_api_format: "openai:image".to_string(), endpoint_api_family: Some("openai".to_string()), endpoint_kind: Some("image".to_string()), endpoint_is_active: true, key_id: format!("key-{prefix}"), key_name: "prod".to_string(), key_auth_type: key_auth_type.to_string(), key_is_active: true, key_api_formats: Some(vec!["openai:image".to_string()]), key_allowed_models: None, key_capabilities: None, key_internal_priority: 5, key_global_priority_by_format: Some(serde_json::json!({"openai:image": 1})), model_id: format!("model-{prefix}"), global_model_id: format!("global-model-{prefix}"), global_model_name: "gpt-image-2".to_string(), global_model_mappings: None, global_model_supports_streaming: Some(false), model_provider_model_name: "gpt-image-2".to_string(), model_provider_model_mappings: Some(vec![StoredProviderModelMapping { name: "gpt-image-2".to_string(), priority: 1, api_formats: Some(vec!["openai:image".to_string()]), endpoint_ids: None, operations: None, }]), model_supports_streaming: Some(false), model_is_active: true, model_is_available: true, } } fn image_bridge_provider_catalog_provider( prefix: &str, provider_name: &str, provider_type: &str, base_url: &str, ) -> StoredProviderCatalogProvider { StoredProviderCatalogProvider::new( format!("provider-{prefix}"), provider_name.to_string(), Some(base_url.to_string()), provider_type.to_string(), ) .expect("provider should build") .with_transport_fields( true, false, false, None, Some(2), None, Some(20.0), None, None, ) } fn image_bridge_provider_catalog_endpoint( prefix: &str, base_url: &str, ) -> StoredProviderCatalogEndpoint { StoredProviderCatalogEndpoint::new( format!("endpoint-{prefix}"), format!("provider-{prefix}"), "openai:image".to_string(), Some("openai".to_string()), Some("image".to_string()), true, ) .expect("endpoint should build") .with_transport_fields( base_url.to_string(), None, None, Some(2), None, None, None, None, ) .expect("endpoint transport should build") } fn image_bridge_provider_catalog_key( prefix: &str, provider_type: &str, ) -> StoredProviderCatalogKey { let auth_type = if provider_type == "chatgpt_web" { "bearer" } else { "api_key" }; StoredProviderCatalogKey::new( format!("key-{prefix}"), format!("provider-{prefix}"), "prod".to_string(), auth_type.to_string(), None, true, ) .expect("key should build") .with_transport_fields( Some(serde_json::json!(["openai:image"])), encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-upstream-image-bridge") .expect("api key should encrypt"), None, None, Some(serde_json::json!({"openai:image": 1})), None, None, None, None, ) .expect("key transport should build") } async fn start_image_bridge_gateway( prefix: &str, provider_name: &str, provider_type: &str, base_url: &str, execution_runtime_url: String, ) -> ( String, tokio::task::JoinHandle<()>, String, Arc, ) { let client_api_key = format!("sk-client-{prefix}"); let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( Some(image_bridge_hash_api_key(&client_api_key)), image_bridge_auth_snapshot(&format!("api-key-{prefix}"), &format!("user-{prefix}")), )])); let candidate_selection_repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![ image_bridge_candidate_row(prefix, provider_name, provider_type), ])); let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![image_bridge_provider_catalog_provider( prefix, provider_name, provider_type, base_url, )], vec![image_bridge_provider_catalog_endpoint(prefix, base_url)], vec![image_bridge_provider_catalog_key(prefix, provider_type)], )); let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default()); let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url) .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, ), ); let gateway = build_router_with_state(gateway_state); let (gateway_url, gateway_handle) = start_server(gateway).await; ( gateway_url, gateway_handle, client_api_key, request_candidate_repository, ) } fn capture_image_bridge_execution_plan( parts: http::request::Parts, payload: serde_json::Value, ) -> SeenImageBridgeExecutionPlan { SeenImageBridgeExecutionPlan { trace_id: parts .headers .get(TRACE_ID_HEADER) .and_then(|value| value.to_str().ok()) .unwrap_or_default() .to_string(), client_api_format: payload .get("client_api_format") .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), provider_api_format: payload .get("provider_api_format") .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), url: payload .get("url") .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), plan_stream: payload .get("stream") .and_then(|value| value.as_bool()) .unwrap_or(false), auth_header: payload .get("headers") .and_then(|value| value.get("authorization")) .and_then(|value| value.as_str()) .unwrap_or_default() .to_string(), body_json: payload .get("body") .and_then(|value| value.get("json_body")) .cloned() .unwrap_or(serde_json::Value::Null), } } fn image_bridge_execution_runtime( seen_execution_plan: Arc>>, ) -> Router { Router::new().route( "/v1/execute/stream", any(move |request: Request| { let seen_execution_plan_inner = Arc::clone(&seen_execution_plan); 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_plan_inner.lock().expect("mutex should lock") = Some(capture_image_bridge_execution_plan(parts, payload)); 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: response.output_item.done\\ndata: {\\\"type\\\":\\\"response.output_item.done\\\",\\\"output_index\\\":0,\\\"item\\\":{\\\"id\\\":\\\"ig_bridge_123\\\",\\\"type\\\":\\\"image_generation_call\\\",\\\"result\\\":\\\"aGVsbG8=\\\",\\\"output_format\\\":\\\"png\\\"}}\\n\\n\"}}\n", "{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp_bridge_123\\\",\\\"object\\\":\\\"response\\\",\\\"model\\\":\\\"gpt-image-2\\\",\\\"status\\\":\\\"completed\\\",\\\"output\\\":[]}}\\n\\n\"}}\n", "{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n" ); let mut response = http::Response::builder() .status(StatusCode::OK) .body(Body::from(frames)) .expect("response should build"); response.headers_mut().insert( http::header::CONTENT_TYPE, http::HeaderValue::from_static("application/x-ndjson"), ); response } }), ) } #[test] fn gateway_routes_openai_responses_stream_image_intent_to_openai_image_plan_without_streaming_support( ) { run_stream_image_test( "gateway_routes_openai_responses_stream_image_intent_to_openai_image_plan_without_streaming_support", gateway_routes_openai_responses_stream_image_intent_to_openai_image_plan_without_streaming_support_impl, ); } async fn gateway_routes_openai_responses_stream_image_intent_to_openai_image_plan_without_streaming_support_impl( ) { let seen_execution_plan = Arc::new(Mutex::new(None::)); let execution_runtime = image_bridge_execution_runtime(Arc::clone(&seen_execution_plan)); let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; let (gateway_url, gateway_handle, client_api_key, _request_candidate_repository) = start_image_bridge_gateway( "responses-stream-image-bridge", "image-provider", "custom", "https://images.example.com/v1", execution_runtime_url, ) .await; let response = reqwest::Client::new() .post(format!("{gateway_url}/v1/responses")) .header(http::header::CONTENT_TYPE, "application/json") .header(http::header::AUTHORIZATION, format!("Bearer {client_api_key}")) .header(TRACE_ID_HEADER, "trace-responses-stream-image-bridge-123") .body( r#"{"model":"gpt-image-2","input":"Draw a mountain observatory","tools":[{"type":"image_generation","size":"1024x1024"}],"tool_choice":{"type":"image_generation"},"stream":true}"#, ) .send() .await .expect("request should succeed"); let status = response.status(); let response_text = response.text().await.expect("body should read"); assert_eq!(status, StatusCode::OK, "{response_text}"); assert!(response_text.contains("response.output_item.done")); assert!(response_text.contains("image_generation_call")); let seen_plan = seen_execution_plan .lock() .expect("mutex should lock") .clone() .expect("execution plan should be captured"); assert_eq!( seen_plan.trace_id, "trace-responses-stream-image-bridge-123" ); assert_eq!(seen_plan.client_api_format, "openai:responses"); assert_eq!(seen_plan.provider_api_format, "openai:image"); assert_eq!( seen_plan.url, "https://images.example.com/v1/images/generations" ); assert!(!seen_plan.plan_stream); assert_eq!(seen_plan.auth_header, "Bearer sk-upstream-image-bridge"); assert_eq!(seen_plan.body_json["model"], "gpt-image-2"); assert_eq!(seen_plan.body_json["prompt"], "Draw a mountain observatory"); assert_eq!(seen_plan.body_json["size"], "1024x1024"); assert!(seen_plan.body_json.get("stream").is_none()); assert!(seen_plan.body_json.get("input").is_none()); assert!(seen_plan.body_json.get("tools").is_none()); gateway_handle.abort(); execution_runtime_handle.abort(); }