From 388aee4b26226778701a7fa749354c44f74b4394 Mon Sep 17 00:00:00 2001 From: cooper-oai Date: Thu, 5 Nov 2026 04:58:48 +0000 Subject: [PATCH] Use runtime authentication across backend request paths (#59572) ## What changed - Share the harness runtime identity with analytics delivery and turn-cost queries, and resolve request authentication together with its account-bound transport policy. - Add runtime authentication to ChatGPT request helpers and `codex cloud`, including environment discovery. Resolve environment credentials for each request and honor CLI configuration overrides when loading cloud authentication. - Preserve custom-backend authentication, validate runtime destinations, and keep cloud requests from following redirects. - Forward only organization and project selectors from model-provider headers to API-key analytics requests. ## Testing Add regression coverage for runtime request signing, credential refresh, account and resource selection, redirect credential handling, custom destinations, and cloud credential-store overrides. GitOrigin-RevId: 75f33961ee6482881e9ff9817e99d977f31b65b8 --- codex-rs/Cargo.lock | 21 ++ codex-rs/analytics/Cargo.toml | 7 +- codex-rs/analytics/src/client.rs | 157 +++++++-- .../analytics/src/client_product_tests.rs | 15 +- .../src/client_runtime_auth_tests.rs | 300 ++++++++++++++++++ codex-rs/analytics/src/client_tests.rs | 27 +- codex-rs/app-server/Cargo.toml | 1 + codex-rs/app-server/src/analytics_utils.rs | 18 +- codex-rs/app-server/src/in_process.rs | 13 +- codex-rs/app-server/src/lib.rs | 12 +- codex-rs/app-server/src/message_processor.rs | 15 +- .../src/message_processor_tracing_tests.rs | 8 +- ...ssage_processor_workspace_routing_tests.rs | 1 + codex-rs/app-server/src/turn_cost_worker.rs | 50 ++- .../src/turn_cost_worker_chatgpt.rs | 8 +- .../turn_cost_worker_runtime_auth_tests.rs | 124 ++++++++ .../app-server/src/turn_cost_worker_tests.rs | 15 +- .../src/user_verification_test_support.rs | 1 + codex-rs/backend-client/src/client.rs | 35 +- .../src/client/runtime_auth_tests.rs | 84 +++++ codex-rs/chatgpt/Cargo.toml | 6 + codex-rs/chatgpt/src/chatgpt_client.rs | 87 ++++- .../src/chatgpt_client_runtime_auth_tests.rs | 212 +++++++++++++ codex-rs/chatgpt/src/connectors.rs | 10 + codex-rs/cli/tests/cloud_auth.rs | 24 +- codex-rs/cloud-tasks-client/src/http.rs | 4 + codex-rs/cloud-tasks/Cargo.toml | 9 + .../cloud-tasks/src/backend_auth_tests.rs | 277 ++++++++++++++++ codex-rs/cloud-tasks/src/env_detect.rs | 36 ++- .../src/env_detect_runtime_auth_tests.rs | 130 ++++++++ codex-rs/cloud-tasks/src/env_detect_tests.rs | 11 +- codex-rs/cloud-tasks/src/lib.rs | 184 +++++++---- codex-rs/cloud-tasks/src/util.rs | 29 +- codex-rs/core/src/session/session.rs | 20 +- codex-rs/core/src/thread_manager.rs | 5 + codex-rs/model-provider/src/provider.rs | 25 -- codex-rs/model-provider/src/provider_tests.rs | 21 +- .../src/request_preparation_tests.rs | 4 +- .../model-provider/src/shared_state_tests.rs | 6 +- 39 files changed, 1821 insertions(+), 191 deletions(-) create mode 100644 codex-rs/analytics/src/client_runtime_auth_tests.rs create mode 100644 codex-rs/app-server/src/turn_cost_worker_runtime_auth_tests.rs create mode 100644 codex-rs/chatgpt/src/chatgpt_client_runtime_auth_tests.rs create mode 100644 codex-rs/cloud-tasks/src/backend_auth_tests.rs create mode 100644 codex-rs/cloud-tasks/src/env_detect_runtime_auth_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 347cf9d64b..3ad82efd49 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2010,15 +2010,19 @@ name = "codex-analytics" version = "0.0.0" dependencies = [ "codex-app-server-protocol", + "codex-features", "codex-git-utils", "codex-http-client", "codex-login", "codex-model-provider", + "codex-model-provider-info", "codex-plugin", "codex-protocol", + "codex-runtime-agent-instance", "codex-state", "codex-utils-absolute-path", "codex-utils-path-uri", + "http 1.4.0", "os_info", "pretty_assertions", "serde", @@ -2026,6 +2030,7 @@ dependencies = [ "sha1 0.10.6", "tokio", "tracing", + "wiremock", ] [[package]] @@ -2138,6 +2143,7 @@ dependencies = [ "codex-queue-extension", "codex-rmcp-client", "codex-rollout", + "codex-runtime-agent-instance", "codex-sandboxing", "codex-shell-command", "codex-skills", @@ -2512,14 +2518,18 @@ name = "codex-chatgpt" version = "0.0.0" dependencies = [ "anyhow", + "base64 0.22.1", "clap", "codex-connectors", "codex-core", + "codex-features", "codex-git-utils", "codex-http-client", "codex-login", "codex-model-provider", "codex-plugin", + "codex-protocol", + "codex-runtime-agent-instance", "codex-utils-cargo-bin", "codex-utils-cli", "pretty_assertions", @@ -2527,6 +2537,8 @@ dependencies = [ "serde_json", "tempfile", "tokio", + "uuid", + "wiremock", ] [[package]] @@ -2706,17 +2718,23 @@ name = "codex-cloud-tasks" version = "0.0.0" dependencies = [ "anyhow", + "base64 0.22.1", "chrono", "clap", + "codex-agent-identity", + "codex-backend-client", "codex-cloud-tasks-client", "codex-cloud-tasks-mock-client", "codex-core", + "codex-features", "codex-git-utils", "codex-http-client", "codex-login", "codex-model-provider", + "codex-runtime-agent-instance", "codex-tui", "codex-utils-cli", + "core_test_support", "crossterm", "http 1.4.0", "insta", @@ -2726,12 +2744,15 @@ dependencies = [ "serde", "serde_json", "supports-color 3.0.2", + "tempfile", "tokio", "tokio-stream", "tracing", "tracing-subscriber", "unicode-segmentation", "unicode-width 0.2.1", + "uuid", + "wiremock", ] [[package]] diff --git a/codex-rs/analytics/Cargo.toml b/codex-rs/analytics/Cargo.toml index 61e6f7210c..ad134107e6 100644 --- a/codex-rs/analytics/Cargo.toml +++ b/codex-rs/analytics/Cargo.toml @@ -13,15 +13,18 @@ path = "src/lib.rs" workspace = true [dependencies] -codex-http-client = { workspace = true } +codex-features = { workspace = true } codex-app-server-protocol = { workspace = true } codex-git-utils = { workspace = true } +codex-http-client = { workspace = true } +codex-runtime-agent-instance = { workspace = true } codex-login = { workspace = true } codex-model-provider = { workspace = true } codex-plugin = { workspace = true } codex-protocol = { workspace = true } codex-state = { workspace = true } codex-utils-path-uri = { workspace = true } +http = { workspace = true } os_info = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } @@ -33,5 +36,7 @@ tokio = { workspace = true, features = [ tracing = { workspace = true, features = ["log"] } [dev-dependencies] +codex-model-provider-info = { workspace = true } +wiremock = { workspace = true } codex-utils-absolute-path = { workspace = true } pretty_assertions = { workspace = true } diff --git a/codex-rs/analytics/src/client.rs b/codex-rs/analytics/src/client.rs index 0c54d0167a..967b45e4ad 100644 --- a/codex-rs/analytics/src/client.rs +++ b/codex-rs/analytics/src/client.rs @@ -63,6 +63,9 @@ use codex_app_server_protocol::TurnItemsView; use codex_app_server_protocol::TurnStartedNotification; use codex_app_server_protocol::TurnStatus; use codex_app_server_protocol::item_event_to_server_notification; +use codex_features::RuntimeAuthMode; +use codex_http_client::OPENAI_ORGANIZATION_HEADER; +use codex_http_client::OPENAI_PROJECT_HEADER; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_model_provider::ChatgptBaseUrl; @@ -75,6 +78,8 @@ use codex_protocol::items::TurnItem; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::request_permissions::RequestPermissionsResponse; +use codex_runtime_agent_instance::RuntimeAuthSlot; +use http::HeaderMap; use std::collections::HashSet; use std::path::PathBuf; use std::sync::Arc; @@ -118,17 +123,31 @@ pub struct AnalyticsEventsClient { enum AnalyticsEventsDestination { Http { base_url: ChatgptBaseUrl, + api_key_headers: HeaderMap, }, #[cfg(debug_assertions)] - CaptureFile { - path: PathBuf, - }, + CaptureFile { path: PathBuf }, } impl AnalyticsEventsDestination { - fn from_base_url(base_url: ChatgptBaseUrl) -> Self { + fn from_base_url(base_url: ChatgptBaseUrl, headers: HeaderMap) -> Self { let capture_file = analytics_capture_file_from_env(); - Self::from_base_url_and_capture_file(base_url, capture_file) + let mut destination = Self::from_base_url_and_capture_file(base_url, capture_file); + match &mut destination { + Self::Http { + api_key_headers, .. + } => { + // Model-provider secrets and gateway headers do not belong on analytics requests. + for name in [OPENAI_ORGANIZATION_HEADER, OPENAI_PROJECT_HEADER] { + if let Some(value) = headers.get(name) { + api_key_headers.insert(name, value.clone()); + } + } + } + #[cfg(debug_assertions)] + Self::CaptureFile { .. } => {} + } + destination } fn from_base_url_and_capture_file( @@ -153,12 +172,15 @@ impl AnalyticsEventsDestination { #[cfg(not(debug_assertions))] let _ = capture_file; - Self::Http { base_url } + Self::Http { + base_url, + api_key_headers: HeaderMap::new(), + } } fn url_for_auth(&self, auth: &CodexAuth) -> Option { match self { - Self::Http { base_url } => { + Self::Http { base_url, .. } => { let base_url = base_url.resolve(auth).trim_end_matches('/'); Some(format!("{base_url}/codex/analytics-events/events")) } @@ -181,7 +203,11 @@ fn analytics_capture_file_from_env() -> Option { } impl AnalyticsEventsQueue { - fn new(auth_manager: Arc, destination: AnalyticsEventsDestination) -> Self { + fn new( + auth_manager: Arc, + destination: AnalyticsEventsDestination, + runtime_auth: Option<(RuntimeAuthSlot, RuntimeAuthMode)>, + ) -> Self { let (sender, mut receiver) = mpsc::channel(ANALYTICS_EVENTS_QUEUE_SIZE); let product_generation = Arc::new(AtomicU64::new(0)); let worker_product_generation = Arc::clone(&product_generation); @@ -218,23 +244,42 @@ impl AnalyticsEventsQueue { { reducer.flush_thread(evicted, &mut events); } - send_track_events(&auth_manager, &destination, events, &thread_products) - .await; + send_track_events( + &auth_manager, + &destination, + events, + &thread_products, + runtime_auth.as_ref(), + ) + .await; thread_products.register(thread_id, product_sku); continue; } AnalyticsEventsQueueMessage::Flush(done_tx) => { let mut events = Vec::new(); reducer.flush(&mut events); - send_track_events(&auth_manager, &destination, events, &thread_products) - .await; + send_track_events( + &auth_manager, + &destination, + events, + &thread_products, + runtime_auth.as_ref(), + ) + .await; let _ = done_tx.send(()); continue; } }; let mut events = Vec::new(); reducer.ingest(input, &mut events).await; - send_track_events(&auth_manager, &destination, events, &thread_products).await; + send_track_events( + &auth_manager, + &destination, + events, + &thread_products, + runtime_auth.as_ref(), + ) + .await; } }); Self { @@ -328,10 +373,36 @@ impl AnalyticsEventsClient { base_url: impl Into, analytics_enabled: Option, ) -> Self { - let destination = AnalyticsEventsDestination::from_base_url(base_url.into()); + Self::new_with_runtime( + auth_manager, + base_url, + analytics_enabled, + /*runtime_auth*/ None, + /*api_key_headers*/ HeaderMap::new(), + ) + } + + /// Share runtime identity for OpenAI delivery; file capture never authenticates. + pub fn new_with_runtime( + auth_manager: Arc, + base_url: impl Into, + analytics_enabled: Option, + runtime_auth: Option<(RuntimeAuthSlot, RuntimeAuthMode)>, + api_key_headers: HeaderMap, + ) -> Self { + let base_url = base_url.into(); + let runtime_auth = runtime_auth.filter(|_| match &base_url { + ChatgptBaseUrl::Default => true, + ChatgptBaseUrl::Configured(url) => { + codex_http_client::classify_openai_api_url(url) + != codex_http_client::OpenAiDestination::Other + } + }); + let destination = AnalyticsEventsDestination::from_base_url(base_url, api_key_headers); Self { - queue: (analytics_enabled != Some(false)) - .then(|| AnalyticsEventsQueue::new(Arc::clone(&auth_manager), destination)), + queue: (analytics_enabled != Some(false)).then(|| { + AnalyticsEventsQueue::new(Arc::clone(&auth_manager), destination, runtime_auth) + }), } } @@ -980,6 +1051,7 @@ async fn send_track_events( destination: &AnalyticsEventsDestination, mut events: Vec, thread_products: &ThreadProducts, + runtime_auth: Option<&(RuntimeAuthSlot, RuntimeAuthMode)>, ) { if events.is_empty() { return; @@ -1008,6 +1080,7 @@ async fn send_track_events( http_client_factory.clone(), ); let product_sku = product_sku.map(str::to_string); + let runtime_auth = runtime_auth.cloned(); requests.spawn(async move { for events in track_event_request_batches(events) { send_track_events_request( @@ -1016,6 +1089,7 @@ async fn send_track_events( events, &factory, product_sku.as_deref(), + runtime_auth.as_ref(), ) .await; } @@ -1057,6 +1131,7 @@ async fn send_track_events_request( events: Vec, http_client_factory: &codex_http_client::HttpClientFactory, product_sku: Option<&str>, + runtime_auth: Option<&(RuntimeAuthSlot, RuntimeAuthMode)>, ) { if events.is_empty() { return; @@ -1072,21 +1147,55 @@ async fn send_track_events_request( let Some(url) = destination.url_for_auth(auth) else { return; }; - let headers = match codex_model_provider::auth_provider_from_auth(auth).to_auth_headers() { - Ok(headers) => headers, + if runtime_auth.is_some() + && let Err(error) = codex_model_provider::validate_runtime_destination(Some(&url)) + { + tracing::warn!("failed to authenticate events request: {error}"); + return; + } + let runtime = match runtime_auth { + Some((slot, mode)) => { + let Some(manager) = slot.get() else { + tracing::warn!("runtime authentication has not been initialized for events"); + return; + }; + Some((Arc::clone(manager), *mode)) + } + None => None, + }; + let headers = match destination { + AnalyticsEventsDestination::Http { + api_key_headers, .. + } if auth.is_api_key_auth() => api_key_headers.clone(), + _ => HeaderMap::new(), + }; + let resolved = async { + let resolved = + codex_model_provider::auth_provider_for_request(auth, runtime.as_ref(), &headers) + .await?; + let factory = resolved.bind_http_client_factory(http_client_factory.clone())?; + resolved + .auth + .apply_auth_headers(headers) + .await + .map(|headers| (headers, factory)) + } + .await; + let (headers, factory) = match resolved { + Ok(resolved) => resolved, Err(error) => { tracing::warn!("failed to authenticate events request: {error}"); return; } }; - let client = match codex_login::default_client::create_client_for_route_async( - http_client_factory.clone(), + let client = codex_login::default_client::create_client_for_route_async( + factory, url.clone(), codex_http_client::ClientRouteClass::Api, codex_login::default_client::ClientRedirectPolicy::Default, ) - .await - { + .await; + let client = match client { Ok(client) => client, Err(error) => { tracing::warn!(%error, "failed to build events client"); @@ -1141,3 +1250,7 @@ fn capture_track_events_request( #[cfg(test)] #[path = "client_tests.rs"] mod tests; + +#[cfg(test)] +#[path = "client_runtime_auth_tests.rs"] +mod runtime_auth_tests; diff --git a/codex-rs/analytics/src/client_product_tests.rs b/codex-rs/analytics/src/client_product_tests.rs index 39c50f73ce..97b3ac6ae0 100644 --- a/codex-rs/analytics/src/client_product_tests.rs +++ b/codex-rs/analytics/src/client_product_tests.rs @@ -90,6 +90,7 @@ async fn capture_requests( let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let destination = AnalyticsEventsDestination::Http { base_url: format!("http://{}", listener.local_addr().unwrap()).into(), + api_key_headers: Default::default(), }; let (paused_tx, paused_rx) = tokio::sync::oneshot::channel(); let (resume_tx, resume_rx) = tokio::sync::oneshot::channel(); @@ -168,7 +169,13 @@ async fn product_streams_send_concurrently_and_preserve_event_order() { let events = (0..32).map(event).collect(); tokio::time::timeout( std::time::Duration::from_secs(15), - send_track_events(&manager, &destination, events, &products), + send_track_events( + &manager, + &destination, + events, + &products, + /*runtime_auth*/ None, + ), ) .await .unwrap(); @@ -207,7 +214,11 @@ async fn buffered_tool_events_preserve_attribution_or_drop_it_on_queue_overflow( codex_login::CodexAuth::create_dummy_chatgpt_auth_for_testing(), ); let client = AnalyticsEventsClient { - queue: Some(AnalyticsEventsQueue::new(manager, destination)), + queue: Some(AnalyticsEventsQueue::new( + manager, + destination, + /*runtime_auth*/ None, + )), }; let thread_id = ThreadId::new(); client.update_thread_product_sku(thread_id, ThreadProductUpdate::Set("aeon".to_string())); diff --git a/codex-rs/analytics/src/client_runtime_auth_tests.rs b/codex-rs/analytics/src/client_runtime_auth_tests.rs new file mode 100644 index 0000000000..cf2f25f869 --- /dev/null +++ b/codex-rs/analytics/src/client_runtime_auth_tests.rs @@ -0,0 +1,300 @@ +//! Runtime signing and destination isolation at the analytics delivery boundary. + +use super::*; +use crate::facts::PluginMeasurementRow; +use crate::facts::PluginMeasurementsInput; +use codex_features::RuntimeAuthMode; +use codex_runtime_agent_instance::AgentHostAuthorization; +use codex_runtime_agent_instance::RuntimeAuthManager; +use codex_runtime_agent_instance::RuntimeAuthSlot; +use codex_runtime_agent_instance::RuntimeRegistrationMetadata; +use codex_runtime_agent_instance::RuntimeRole; +use pretty_assertions::assert_eq; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::path; + +async fn emit(client: &AnalyticsEventsClient, execution_id: &str) { + client.track_plugin_measurements(PluginMeasurementsInput { + thread_id: "thread-a".into(), + turn_id: "turn-a".into(), + item_id: "item-a".into(), + originator: "codex_cli_rs".into(), + model_slug: None, + reasoning_effort: None, + plugin_id: "plugin-a".into(), + execution_id: execution_id.into(), + operation: "security_scan".into(), + rows: vec![PluginMeasurementRow { + measurement_name: "findings".into(), + number_value: 1.0, + dimensions: Default::default(), + }], + }); + client.flush().await; +} + +#[tokio::test] +async fn runtime_delivery_signs_each_batch_and_strips_cross_origin_auth() { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + let redirected = MockServer::start().await; + Mock::given(path("/")) + .respond_with(ResponseTemplate::new(200)) + .expect(2) + .mount(&redirected) + .await; + Mock::given(path("/agent/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "runtime_agent_instance_id": "runtime-analytics", "agent_host_id": "host-a", + "actor": {"kind": "platform", "user_id": "user-a", "organization_id": "org-a", "project_id": null}, + }))) + .expect(1).mount(&server).await; + Mock::given(path("/codex/analytics-events/events")) + .respond_with(ResponseTemplate::new(307).insert_header("location", redirected.uri())) + .expect(2) + .mount(&server) + .await; + let credentials = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("key-a")); + let runtime = ( + Arc::new(RuntimeAuthManager::new( + Arc::clone(&credentials), + AgentHostAuthorization::InstallationId { + installation_id: "00000000-0000-0000-0000-000000000000".parse().unwrap(), + }, + RuntimeRegistrationMetadata { + harness: Some("codex".into()), + harness_version: None, + role: RuntimeRole::Orchestrator, + }, + format!("{}/agent/register", server.uri()).parse().unwrap(), + )), + mode, + ); + // Inject the selected context at the delivery boundary for the loopback fixture. + let client = AnalyticsEventsClient { + queue: Some(AnalyticsEventsQueue::new( + credentials, + AnalyticsEventsDestination::Http { + base_url: server.uri().into(), + api_key_headers: Default::default(), + }, + Some((Arc::new(std::sync::OnceLock::from(runtime.0)), runtime.1)), + )), + }; + emit(&client, "first").await; + emit(&client, "second").await; + let forwarded = redirected.received_requests().await.unwrap(); + assert_eq!(forwarded.len(), 2); + for request in &forwarded { + assert!(!request.headers.contains_key("authorization")); + assert!(!request.headers.contains_key("x-openai-agent-assertion")); + } + redirected.verify().await; + let requests = server.received_requests().await.unwrap(); + let proofs: Vec<_> = requests + .iter() + .filter(|r| r.url.path() == "/codex/analytics-events/events") + .map(|request| { + let authorization = request.headers["authorization"].to_str().unwrap(); + match mode { + RuntimeAuthMode::DualHeader => { + assert_eq!(authorization, "Bearer key-a"); + request.headers["x-openai-agent-assertion"] + .to_str() + .unwrap() + .to_string() + } + RuntimeAuthMode::AssertionOnly => { + assert!(!request.headers.contains_key("x-openai-agent-assertion")); + authorization + .strip_prefix("AgentAssertion ") + .unwrap() + .to_string() + } + } + }) + .collect(); + assert_eq!(proofs.len(), 2); + assert_ne!(proofs[0], proofs[1]); + server.verify().await; + } +} + +#[tokio::test] +async fn custom_analytics_destination_does_not_initialize_runtime_auth() { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + Mock::given(path("/codex/analytics-events/events")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + let client = AnalyticsEventsClient::new_with_runtime( + AuthManager::from_auth_for_testing(CodexAuth::from_api_key("key-a")), + server.uri(), + Some(true), + Some((RuntimeAuthSlot::default(), mode)), + /*api_key_headers*/ Default::default(), + ); + emit(&client, "custom").await; + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].headers["authorization"], "Bearer key-a"); + assert!(!requests[0].headers.contains_key("x-openai-agent-assertion")); + server.verify().await; + } +} + +#[tokio::test] +async fn analytics_between_model_calls_keeps_the_selected_api_resources() { + use codex_http_client::ClientRouteClass; + use codex_login::default_client::ClientRedirectPolicy; + use codex_model_provider_info::ModelProviderInfo; + use std::collections::HashMap; + + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + Mock::given(path("/agent/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "runtime_agent_instance_id": "runtime-selected", "agent_host_id": "host-a", + "actor": {"kind": "platform", "user_id": "user-a", "organization_id": "org-selected", "project_id": null}, + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(path("/v1/responses")) + .respond_with( + ResponseTemplate::new(200).set_body_json(serde_json::json!({"output": []})), + ) + .expect(2) + .mount(&server) + .await; + Mock::given(path("/codex/analytics-events/events")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + let credentials = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("key-a")); + let owner = Arc::new(RuntimeAuthManager::new( + Arc::clone(&credentials), + AgentHostAuthorization::InstallationId { + installation_id: "00000000-0000-0000-0000-000000000000".parse().unwrap(), + }, + RuntimeRegistrationMetadata { + harness: Some("codex".into()), + harness_version: None, + role: RuntimeRole::Orchestrator, + }, + format!("{}/agent/register", server.uri()).parse().unwrap(), + )); + let mut info = + ModelProviderInfo::create_openai_provider(Some(format!("{}/v1", server.uri()))); + info.env_http_headers = None; + info.http_headers = Some(HashMap::from([ + ("OpenAI-Organization".into(), "org-selected".into()), + ("OpenAI-Project".into(), "proj-selected".into()), + ("X-Model-Secret".into(), "model-only".into()), + ])); + let destination = AnalyticsEventsDestination::from_base_url( + server.uri().into(), + info.build_header_map().unwrap(), + ); + let analytics = AnalyticsEventsClient { + queue: Some(AnalyticsEventsQueue::new( + Arc::clone(&credentials), + destination, + Some(( + Arc::new(std::sync::OnceLock::from(Arc::clone(&owner))), + mode, + )), + )), + }; + let provider = codex_model_provider::test_support::runtime_model_provider( + info, + Arc::clone(&credentials), + (owner, mode), + ); + for attempt in 0..2 { + let request = codex_model_provider::resolve_request( + provider.as_ref(), + /*scope*/ None, + /*routing*/ None, + credentials.http_client_factory(), + ) + .await + .unwrap(); + let url = request.api_provider.url_for_path("responses"); + let headers = request + .api_auth + .apply_auth_headers(request.api_provider.headers) + .await + .unwrap(); + let client = codex_login::default_client::create_client_for_route_async( + request.http_client_factory, + url.clone(), + ClientRouteClass::Api, + ClientRedirectPolicy::Default, + ) + .await + .unwrap(); + assert!( + client + .post(&url) + .headers(headers) + .json(&serde_json::json!({"input": "hello"})) + .send() + .await + .unwrap() + .status() + .is_success() + ); + if attempt == 0 { + emit(&analytics, "between-model-calls").await; + } + } + let requests = server.received_requests().await.unwrap(); + let registrations: Vec<_> = requests + .iter() + .filter(|request| request.url.path() == "/agent/register") + .collect(); + assert_eq!( + registrations.len(), + 1, + "analytics must reuse the model registration" + ); + let registration: serde_json::Value = registrations[0].body_json().unwrap(); + assert_eq!( + registration["actor_auth"]["organization_id"], + "org-selected" + ); + assert_eq!(registration["actor_auth"]["project_id"], "proj-selected"); + for request in requests + .iter() + .filter(|request| request.url.path() != "/agent/register") + { + assert_eq!(request.headers["openai-organization"], "org-selected"); + assert_eq!(request.headers["openai-project"], "proj-selected"); + match mode { + RuntimeAuthMode::DualHeader => { + assert_eq!(request.headers["authorization"], "Bearer key-a"); + assert!(request.headers.contains_key("x-openai-agent-assertion")); + } + RuntimeAuthMode::AssertionOnly => { + assert!( + request.headers["authorization"] + .to_str() + .unwrap() + .starts_with("AgentAssertion ") + ); + assert!(!request.headers.contains_key("x-openai-agent-assertion")); + } + } + if request.url.path() == "/codex/analytics-events/events" { + assert!(!request.headers.contains_key("x-model-secret")); + } + } + server.verify().await; + } +} diff --git a/codex-rs/analytics/src/client_tests.rs b/codex-rs/analytics/src/client_tests.rs index d7d8ff637e..6819823d9e 100644 --- a/codex-rs/analytics/src/client_tests.rs +++ b/codex-rs/analytics/src/client_tests.rs @@ -123,6 +123,7 @@ impl AnalyticsEventsClient { queue: Some(AnalyticsEventsQueue::new( auth_manager, AnalyticsEventsDestination::CaptureFile { path }, + /*runtime_auth*/ None, )), } } @@ -358,7 +359,13 @@ fn analytics_destination_uses_http_without_capture_file() { base_url.clone(), /*capture_file*/ None, ); - assert_eq!(destination, AnalyticsEventsDestination::Http { base_url }); + assert_eq!( + destination, + AnalyticsEventsDestination::Http { + base_url, + api_key_headers: Default::default() + } + ); for (api_key, expected_url) in ["sk-test", "sk-proj-fedramp-test", "sk-test"] .into_iter() @@ -384,7 +391,8 @@ fn analytics_destination_ignores_capture_file_in_release() { assert_eq!( destination, AnalyticsEventsDestination::Http { - base_url: "https://chatgpt.com/backend-api/".into() + base_url: "https://chatgpt.com/backend-api/".into(), + api_key_headers: Default::default(), } ); } @@ -401,12 +409,15 @@ async fn capture_file_writes_exact_serialized_request() { let auth = codex_login::CodexAuth::create_dummy_chatgpt_auth_for_testing(); let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); + // Capture-only delivery must not initialize or depend on runtime authentication. + let runtime = (Default::default(), Default::default()); send_track_events_request( &auth, &destination, vec![event], &factory, /*product_sku*/ None, + Some(&runtime), ) .await; @@ -442,6 +453,7 @@ async fn capture_file_writes_final_batches_as_separate_lines() { batch, &factory, /*product_sku*/ None, + /*runtime_auth*/ None, ) .await; } @@ -520,6 +532,7 @@ async fn api_key_auth_sends_only_plugin_events_to_codex_backend() { plugin_measurement("plugin-measurement", "sample@test"), ], &ThreadProducts::default(), + /*runtime_auth*/ None, ) .await; @@ -1196,6 +1209,7 @@ async fn account_switch_between_batches_does_not_send_old_credentials() { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let destination = AnalyticsEventsDestination::Http { base_url: format!("http://{}", listener.local_addr().unwrap()).into(), + api_key_headers: Default::default(), }; let requests = Arc::new(Mutex::new(Vec::new())); let captured = Arc::clone(&requests); @@ -1246,7 +1260,14 @@ async fn account_switch_between_batches_does_not_send_old_credentials() { sample_accepted_line_fingerprint_event("one"), sample_accepted_line_fingerprint_event("two"), ]; - send_track_events(&manager, &destination, events, &ThreadProducts::default()).await; + send_track_events( + &manager, + &destination, + events, + &ThreadProducts::default(), + /*runtime_auth*/ None, + ) + .await; server.abort(); let _ = server.await; let requests = requests.lock().unwrap(); diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index 628e2f47ef..959373d773 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -74,6 +74,7 @@ codex-backend-client = { workspace = true } codex-file-search = { workspace = true } codex-chatgpt = { workspace = true } codex-login = { workspace = true } +codex-runtime-agent-instance = { workspace = true } codex-image-generation-extension = { workspace = true } codex-memories-extension = { workspace = true } codex-web-search-extension = { workspace = true } diff --git a/codex-rs/app-server/src/analytics_utils.rs b/codex-rs/app-server/src/analytics_utils.rs index 2766fb759c..9ff55f7dd3 100644 --- a/codex-rs/app-server/src/analytics_utils.rs +++ b/codex-rs/app-server/src/analytics_utils.rs @@ -3,14 +3,30 @@ use std::sync::Arc; use codex_analytics::AnalyticsEventsClient; use codex_core::config::Config; use codex_login::AuthManager; +use codex_runtime_agent_instance::RuntimeAuthSlot; pub(crate) fn analytics_events_client_from_config( auth_manager: Arc, config: &Config, + runtime_auth_slot: Option, ) -> AnalyticsEventsClient { - AnalyticsEventsClient::new( + let api_key_headers = + if codex_model_provider::provider_uses_first_party_auth_path(&config.model_provider) { + match config.model_provider.build_header_map() { + Ok(headers) => headers, + Err(error) => { + tracing::warn!(%error, "failed to resolve analytics resource selectors"); + return AnalyticsEventsClient::disabled(); + } + } + } else { + Default::default() + }; + AnalyticsEventsClient::new_with_runtime( auth_manager, config.chatgpt_backend_url(), config.analytics_enabled, + runtime_auth_slot.map(|slot| (slot, config.runtime_auth_mode)), + api_key_headers, ) } diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 4691c05134..3e520ea358 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -460,8 +460,16 @@ async fn start_uninitialized(mut args: InProcessStartArgs) -> IoResult(channel_capacity); - let analytics_events_client = - analytics_events_client_from_config(Arc::clone(&auth_manager), args.config.as_ref()); + let runtime_auth_slot = args + .config + .features + .enabled(codex_features::Feature::RuntimeAgentInstances) + .then(codex_runtime_agent_instance::RuntimeAuthSlot::default); + let analytics_events_client = analytics_events_client_from_config( + Arc::clone(&auth_manager), + args.config.as_ref(), + runtime_auth_slot.clone(), + ); let analytics_events_flush_client = analytics_events_client.clone(); // A client can be waiting for process/kill while nobody is polling next_event. let (process_reply_tx, mut process_reply_rx) = mpsc::channel(channel_capacity); @@ -528,6 +536,7 @@ async fn start_uninitialized(mut args: InProcessStartArgs) -> IoResult, pub(crate) session_source: SessionSource, pub(crate) auth_manager: Arc, + pub(crate) runtime_auth_slot: Option, pub(crate) user_verification: Arc, pub(crate) installation_id: String, pub(crate) code_mode_session_provider: Option>, @@ -304,6 +306,7 @@ impl MessageProcessor { startup_warnings_tx, session_source, auth_manager, + runtime_auth_slot, user_verification, installation_id, code_mode_session_provider, @@ -425,6 +428,11 @@ impl MessageProcessor { None => manager, } }); + let runtime_auth = thread_manager.runtime_auth_manager(); + if let (Some(slot), Some(runtime)) = (runtime_auth_slot, &runtime_auth) { + // Analytics starts before the harness. Bind its requests to the harness registration. + slot.get_or_init(|| Arc::clone(runtime)); + } let model_catalog = Arc::new(crate::model_catalog::ModelCatalog::new( config_manager.clone(), startup_cwd, @@ -432,8 +440,11 @@ impl MessageProcessor { Arc::clone(&thread_manager), )); let models_refresh_worker = crate::models_refresh_worker::spawn(&model_catalog); - let turn_cost_worker = - TurnCostWorker::spawn(Arc::clone(&config), Arc::clone(&auth_manager)); + let turn_cost_worker = TurnCostWorker::spawn( + Arc::clone(&config), + Arc::clone(&auth_manager), + runtime_auth.map(|runtime| (runtime, config.runtime_auth_mode)), + ); thread_manager .plugins_manager() .set_analytics_events_client(analytics_events_client.clone()); diff --git a/codex-rs/app-server/src/message_processor_tracing_tests.rs b/codex-rs/app-server/src/message_processor_tracing_tests.rs index aeff494f09..9f85ed0b87 100644 --- a/codex-rs/app-server/src/message_processor_tracing_tests.rs +++ b/codex-rs/app-server/src/message_processor_tracing_tests.rs @@ -254,8 +254,11 @@ pub(super) async fn build_test_processor( Arg0DispatchPaths::default(), Arc::new(codex_config::NoopThreadConfigLoader), ); - let analytics_events_client = - analytics_events_client_from_config(Arc::clone(&auth_manager), config.as_ref()); + let analytics_events_client = analytics_events_client_from_config( + Arc::clone(&auth_manager), + config.as_ref(), + /*runtime_auth_slot*/ None, + ); let outgoing = Arc::new(OutgoingMessageSender::new( outgoing_tx, analytics_events_client.clone(), @@ -279,6 +282,7 @@ pub(super) async fn build_test_processor( &auth_manager, ))), auth_manager, + runtime_auth_slot: None, installation_id: "11111111-1111-4111-8111-111111111111".to_string(), code_mode_session_provider: None, rpc_transport: AppServerRpcTransport::Stdio, diff --git a/codex-rs/app-server/src/message_processor_workspace_routing_tests.rs b/codex-rs/app-server/src/message_processor_workspace_routing_tests.rs index e3f911f6db..6452bb14e4 100644 --- a/codex-rs/app-server/src/message_processor_workspace_routing_tests.rs +++ b/codex-rs/app-server/src/message_processor_workspace_routing_tests.rs @@ -82,6 +82,7 @@ async fn background_routing_call_sites_skip_config(api_key: Option<&str>) -> any &auth_manager, ))), auth_manager, + runtime_auth_slot: None, installation_id: "11111111-1111-4111-8111-111111111111".to_string(), code_mode_session_provider: None, rpc_transport: AppServerRpcTransport::Stdio, diff --git a/codex-rs/app-server/src/turn_cost_worker.rs b/codex-rs/app-server/src/turn_cost_worker.rs index 197f17f95e..cdbaeae92f 100644 --- a/codex-rs/app-server/src/turn_cost_worker.rs +++ b/codex-rs/app-server/src/turn_cost_worker.rs @@ -4,14 +4,15 @@ use codex_backend_client::Client as BackendClient; use codex_backend_client::RequestError; use codex_config::types::OtelExporterKind; use codex_core::config::Config; +use codex_features::RuntimeAuthMode; use codex_login::AuthManager; use codex_model_provider::SharedModelProvider; -use codex_model_provider::create_model_provider; use codex_otel::SessionTelemetry; use codex_protocol::ThreadId; use codex_protocol::auth::AuthMode; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; +use codex_runtime_agent_instance::RuntimeAuthManager; use std::collections::HashMap; use std::collections::HashSet; use std::sync::Arc; @@ -81,6 +82,7 @@ struct WorkerRuntime { config: Arc, backend: TurnCostBackend, turns: HashMap, + runtime_auth: Option<(Arc, RuntimeAuthMode)>, } #[derive(Clone)] @@ -98,7 +100,11 @@ enum BackendAvailability { } impl TurnCostWorker { - pub(crate) fn spawn(config: Arc, auth_manager: Arc) -> Option { + pub(crate) fn spawn( + config: Arc, + auth_manager: Arc, + runtime_auth: Option<(Arc, RuntimeAuthMode)>, + ) -> Option { let has_otel_log_exporter = matches!( config.otel.exporter, OtelExporterKind::OtlpHttp { .. } | OtelExporterKind::OtlpGrpc { .. } @@ -116,10 +122,13 @@ impl TurnCostWorker { let backend = if is_openai { TurnCostBackend::OpenAi(Arc::clone(&auth_manager)) } else { - TurnCostBackend::ModelProvider(create_model_provider( - config.model_provider.clone(), - Some(Arc::clone(&auth_manager)), - )) + TurnCostBackend::ModelProvider( + codex_model_provider::ProviderAuthContext { + credentials: Arc::clone(&auth_manager), + runtime_auth: runtime_auth.clone(), + } + .model_provider(config.model_provider.clone()), + ) }; let (sender, receiver) = mpsc::channel(OBSERVATION_CHANNEL_CAPACITY); let shutdown = CancellationToken::new(); @@ -127,6 +136,7 @@ impl TurnCostWorker { config: Arc::clone(&config), backend: backend.clone(), turns: HashMap::new(), + runtime_auth, }; let worker_shutdown = shutdown.clone(); let task = tokio::spawn(async move { @@ -423,11 +433,15 @@ impl WorkerRuntime { .model_provider .to_api_provider(Some(AuthMode::ApiKey)) .map_err(|error| RequestError::Other(error.into()))?; - let client = BackendClient::from_auth( + let client = BackendClient::from_auth_with_runtime( self.config.chatgpt_backend_url().resolve(&auth), &auth, http_client_factory, - ); + self.runtime_auth.as_ref(), + &provider.headers, + ) + .await + .map_err(RequestError::Other)?; client .query_api_key_turn_costs(turn_ids, &provider.headers) .await @@ -448,14 +462,20 @@ impl WorkerRuntime { return Ok(None); } } - let (provider, auth) = model_provider - .api_provider_and_auth() - .await - .map_err(|error| RequestError::Other(error.into()))?; + let resolved = codex_model_provider::resolve_request( + model_provider.as_ref(), + /*scope*/ None, + /*routing*/ None, + http_client_factory, + ) + .await + .map_err(|error| RequestError::Other(error.into()))?; + let provider = resolved.api_provider; let endpoint = provider.url_for_path("analytics/codex/turn-costs"); - let client = BackendClient::new(provider.base_url.clone(), http_client_factory) - .with_auth_provider(auth) - .map_err(RequestError::Other)?; + let client = + BackendClient::new(provider.base_url.clone(), resolved.http_client_factory) + .with_auth_provider(resolved.api_auth) + .map_err(RequestError::Other)?; client .query_api_key_turn_costs_at(&endpoint, turn_ids, &provider.headers) .await diff --git a/codex-rs/app-server/src/turn_cost_worker_chatgpt.rs b/codex-rs/app-server/src/turn_cost_worker_chatgpt.rs index a9bd6834b2..b79fe31401 100644 --- a/codex-rs/app-server/src/turn_cost_worker_chatgpt.rs +++ b/codex-rs/app-server/src/turn_cost_worker_chatgpt.rs @@ -30,11 +30,15 @@ impl WorkerRuntime { .or_default() .push(turn_id.clone()); } - let client = BackendClient::from_auth( + let client = BackendClient::from_auth_with_runtime( self.config.chatgpt_base_url.clone(), auth, http_client_factory, - ); + self.runtime_auth.as_ref(), + /*headers*/ &Default::default(), + ) + .await + .map_err(RequestError::Other)?; let costs = client.query_chatgpt_turn_costs(&threads).await?; let mut settled_costs = Vec::new(); for thread in costs { diff --git a/codex-rs/app-server/src/turn_cost_worker_runtime_auth_tests.rs b/codex-rs/app-server/src/turn_cost_worker_runtime_auth_tests.rs new file mode 100644 index 0000000000..b95c9e8c0e --- /dev/null +++ b/codex-rs/app-server/src/turn_cost_worker_runtime_auth_tests.rs @@ -0,0 +1,124 @@ +//! Exercise runtime-authenticated turn-cost requests through OpenAI and provider aliases. + +use super::*; +use codex_features::RuntimeAuthMode; +use codex_runtime_agent_instance::AgentHostAuthorization; +use codex_runtime_agent_instance::RuntimeAuthManager; +use codex_runtime_agent_instance::RuntimeRegistrationMetadata; +use codex_runtime_agent_instance::RuntimeRole; +use pretty_assertions::assert_eq; +use test_case::test_case; +use wiremock::matchers::body_partial_json; + +#[test_case(false; "default OpenAI provider")] +#[test_case(true; "OpenAI alias")] +#[tokio::test] +async fn turn_costs_register_with_requested_organization(openai_alias: bool) { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + let registration = if openai_alias { + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "runtime_agent_instance_id": "runtime-telemetry", "agent_host_id": "host-a", + "actor": {"kind": "platform", "user_id": "user-a", "organization_id": "org-a", "project_id": null}, + })) + } else { + // Inspect registration without sending a request to the real OpenAI destination. + ResponseTemplate::new(403) + }; + Mock::given(path("/agent/register")) + .and(body_partial_json(serde_json::json!({ + "actor_auth": {"organization_id": "org-a"}, + }))) + .respond_with(registration) + .expect(1) + .mount(&server) + .await; + Mock::given(path(TURN_COST_PATH)) + .and(header("openai-organization", "org-a")) + .respond_with( + ResponseTemplate::new(200).set_body_json(serde_json::json!({"turns": []})), + ) + .expect(if openai_alias { 2 } else { 0 }) + .mount(&server) + .await; + let credentials = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("key-a")); + let mut worker = test_runtime(&server, credentials.clone()).await; + let owner = Arc::new(RuntimeAuthManager::new( + credentials.clone(), + AgentHostAuthorization::InstallationId { + installation_id: uuid::Uuid::nil(), + }, + RuntimeRegistrationMetadata { + harness: Some("codex".into()), + harness_version: None, + role: RuntimeRole::Orchestrator, + }, + format!("{}/agent/register", server.uri()).parse().unwrap(), + )); + let mut info = codex_model_provider_info::ModelProviderInfo::create_openai_provider(Some( + format!("{}/v1", server.uri()), + )); + info.name = "OpenAI alias".to_string(); + info.http_headers = Some(HashMap::from([( + "OpenAI-Organization".to_string(), + "org-a".into(), + )])); + if openai_alias { + worker.backend = TurnCostBackend::ModelProvider( + codex_model_provider::test_support::runtime_model_provider( + info, + credentials, + (owner, mode), + ), + ); + } else { + let config = Arc::make_mut(&mut worker.config); + config.model_provider = info; + config.chatgpt_base_url = "https://chatgpt.com".to_string(); + worker.runtime_auth = Some((owner, mode)); + worker + .query_turn_costs(&[]) + .await + .expect_err("registration denial must stop the cost request"); + server.verify().await; + continue; + } + for _ in 0..2 { + assert!( + worker + .query_turn_costs(&[]) + .await + .unwrap() + .unwrap() + .is_empty() + ); + } + let requests = server.received_requests().await.unwrap(); + let proofs: Vec<_> = requests + .iter() + .filter(|r| r.url.path() != "/agent/register") + .map(|r| { + let authorization = r.headers["authorization"].to_str().unwrap(); + match mode { + RuntimeAuthMode::DualHeader => { + assert_eq!(authorization, "Bearer key-a"); + r.headers["x-openai-agent-assertion"] + .to_str() + .unwrap() + .to_owned() + } + RuntimeAuthMode::AssertionOnly => { + assert!(!r.headers.contains_key("x-openai-agent-assertion")); + authorization + .strip_prefix("AgentAssertion ") + .unwrap() + .to_owned() + } + } + }) + .collect(); + assert_eq!(proofs.len(), 2); + assert_eq!(proofs.iter().collect::>().len(), 2); + server.verify().await; + } +} diff --git a/codex-rs/app-server/src/turn_cost_worker_tests.rs b/codex-rs/app-server/src/turn_cost_worker_tests.rs index e6a5f81ae1..6a5b4a0631 100644 --- a/codex-rs/app-server/src/turn_cost_worker_tests.rs +++ b/codex-rs/app-server/src/turn_cost_worker_tests.rs @@ -5,6 +5,7 @@ use codex_login::AuthCredentialsStoreMode; use codex_login::AuthKeyringBackendKind; use codex_login::CodexAuth; use codex_login::login_with_api_key; +use codex_model_provider::create_model_provider; use codex_model_provider_info::ModelProviderInfo; use codex_otel::TelemetryAuthMode; use codex_protocol::protocol::SessionSource; @@ -14,6 +15,7 @@ use tempfile::TempDir; use test_case::test_case; use tokio::time::timeout; use tracing::instrument::WithSubscriber; +use tracing_subscriber::layer::SubscriberExt; use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; @@ -52,7 +54,7 @@ async fn worker_starts_with_otlp_metrics_exporter_without_log_exporter() { }; let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("sk-test")); - let worker = TurnCostWorker::spawn(Arc::new(config), auth_manager) + let worker = TurnCostWorker::spawn(Arc::new(config), auth_manager, /*runtime_auth*/ None) .expect("OTLP metrics exporter should enable turn-cost collection"); wait_for_request_count(&server, /*expected*/ 1).await; worker.shutdown(); @@ -207,6 +209,7 @@ async fn custom_provider_auth_failure_retries_without_auth_changes() { config: Arc::new(config), backend, turns: HashMap::new(), + runtime_auth: None, }; assert_eq!( runtime.probe_backend().await, @@ -255,6 +258,7 @@ async fn custom_provider_does_not_send_chatgpt_auth_for_turn_costs() { config: Arc::new(config), backend, turns: HashMap::new(), + runtime_auth: None, }; assert_eq!(runtime.probe_backend().await, BackendAvailability::Disabled); @@ -406,6 +410,7 @@ async fn test_runtime(server: &MockServer, auth_manager: Arc) -> Wo config: Arc::new(config), backend, turns: HashMap::new(), + runtime_auth: None, } } @@ -477,6 +482,11 @@ async fn turn_cost_failures_do_not_log_response_bodies(status: u16) { .with_max_level(tracing::Level::DEBUG) .with_writer(Arc::new(capture.reopen().expect("log writer"))) .finish(); + // A parallel test can first register this callsite without a subscriber. + // Keep two dispatchers alive so it cannot globally disable our log capture. + let _parallel_dispatch = tracing::Dispatch::new( + tracing_subscriber::registry().with(tracing_subscriber::filter::LevelFilter::OFF), + ); async { assert_eq!( runtime.probe_backend().await, @@ -492,3 +502,6 @@ async fn turn_cost_failures_do_not_log_response_bodies(status: u16) { assert!(!logs.contains("private-workspace-response-marker")); server.verify().await; } + +#[path = "turn_cost_worker_runtime_auth_tests.rs"] +mod runtime_auth; diff --git a/codex-rs/app-server/src/user_verification_test_support.rs b/codex-rs/app-server/src/user_verification_test_support.rs index 1a392d102c..ab2c292f18 100644 --- a/codex-rs/app-server/src/user_verification_test_support.rs +++ b/codex-rs/app-server/src/user_verification_test_support.rs @@ -170,6 +170,7 @@ impl Harness { session_source: SessionSource::Cli, auth_manager: Arc::clone(&auth), user_verification: Arc::clone(&service), + runtime_auth_slot: None, installation_id: "11111111-1111-4111-8111-111111111111".into(), code_mode_session_provider: None, rpc_transport: match origin { diff --git a/codex-rs/backend-client/src/client.rs b/codex-rs/backend-client/src/client.rs index a411ab9a20..05a2493a9f 100644 --- a/codex-rs/backend-client/src/client.rs +++ b/codex-rs/backend-client/src/client.rs @@ -283,8 +283,36 @@ impl Client { ) -> Result { let base_url = base_url.into(); let runtime = codex_model_provider::runtime_auth_for_url(runtime, &base_url)?; - Self::from_auth_with_selected_runtime(base_url, auth, http_client_factory, runtime, headers) - .await + Self::from_auth_with_selected_runtime( + base_url, + auth, + http_client_factory, + runtime, + headers, + Self::new, + ) + .await + } + + /// Build an authenticated client without following redirects. + pub async fn from_auth_with_runtime_without_redirects( + base_url: impl Into, + auth: &CodexAuth, + http_client_factory: HttpClientFactory, + runtime: Option<&(Arc, RuntimeAuthMode)>, + headers: &HeaderMap, + ) -> Result { + let base_url = base_url.into(); + let runtime = codex_model_provider::runtime_auth_for_url(runtime, &base_url)?; + Self::from_auth_with_selected_runtime( + base_url, + auth, + http_client_factory, + runtime, + headers, + Self::new_without_redirects, + ) + .await } async fn from_auth_with_selected_runtime( @@ -293,10 +321,11 @@ impl Client { http_client_factory: HttpClientFactory, runtime: Option<&(Arc, RuntimeAuthMode)>, headers: &HeaderMap, + new_client: fn(String, HttpClientFactory) -> Self, ) -> Result { let (auth, http_client_factory) = ClientAuth::for_auth(auth, runtime, http_client_factory, headers).await?; - let mut client = Self::new(base_url, http_client_factory); + let mut client = new_client(base_url, http_client_factory); client.auth = auth; Ok(client.with_user_agent(get_codex_user_agent())) } diff --git a/codex-rs/backend-client/src/client/runtime_auth_tests.rs b/codex-rs/backend-client/src/client/runtime_auth_tests.rs index 4cbe176dc2..a3e4c5ac3d 100644 --- a/codex-rs/backend-client/src/client/runtime_auth_tests.rs +++ b/codex-rs/backend-client/src/client/runtime_auth_tests.rs @@ -95,6 +95,7 @@ async fn account_runtime_auth_signs_each_request_and_strips_credentials_on_redir factory, Some(&runtime), &HeaderMap::new(), + Client::new, ) .await .unwrap(); @@ -223,6 +224,7 @@ async fn client_registers_in_requested_organization_when_default_is_denied() { HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), Some(&runtime), &headers, + Client::new, ) .await .unwrap(); @@ -249,6 +251,85 @@ async fn client_registers_in_requested_organization_when_default_is_denied() { } } +#[tokio::test] +async fn authenticated_clients_without_redirects_return_the_redirect_response() { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + let destination = MockServer::start().await; + Mock::given(path("/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "runtime_agent_instance_id": "runtime-a", "agent_host_id": "host-a", + "actor": {"kind": "platform", "user_id": "user-a", "organization_id": "org-a", "project_id": null}, + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(307).insert_header("location", destination.uri())) + .expect(2) + .mount(&server) + .await; + let auth = CodexAuth::from_api_key("key-a"); + let runtime = runtime_provider( + &server, + AuthManager::from_auth_for_testing(auth.clone()), + mode, + ); + let custom_backend = Client::from_auth_with_runtime_without_redirects( + server.uri(), + &auth, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + Some(&runtime), + &HeaderMap::new(), + ) + .await + .unwrap(); + assert_eq!( + custom_backend + .get_accounts_check() + .await + .unwrap_err() + .status(), + Some(StatusCode::TEMPORARY_REDIRECT), + ); + let client = Client::from_auth_with_selected_runtime( + server.uri(), + &auth, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + Some(&runtime), + &HeaderMap::new(), + Client::new_without_redirects, + ) + .await + .unwrap(); + + let error = client.get_accounts_check().await.unwrap_err(); + assert_eq!(error.status(), Some(StatusCode::TEMPORARY_REDIRECT)); + assert!(destination.received_requests().await.unwrap().is_empty()); + let requests = server.received_requests().await.unwrap(); + let requests: Vec<_> = requests + .iter() + .filter(|request| request.method == Method::GET) + .collect(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[0].headers["authorization"], "Bearer key-a"); + assert!(!requests[0].headers.contains_key("x-openai-agent-assertion")); + let request = requests[1]; + let authorization = request.headers["authorization"].to_str().unwrap(); + match mode { + RuntimeAuthMode::DualHeader => { + assert_eq!(authorization, "Bearer key-a"); + assert!(request.headers.contains_key("x-openai-agent-assertion")); + } + RuntimeAuthMode::AssertionOnly => { + assert!(authorization.starts_with("AgentAssertion ")); + assert!(!request.headers.contains_key("x-openai-agent-assertion")); + } + } + server.verify().await; + } +} + #[tokio::test] async fn unsupported_openai_backend_rejects_runtime_auth_before_registration() { for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { @@ -372,6 +453,7 @@ async fn reusable_runtime_client_keeps_the_selected_accounts_transport_policy() factory.clone(), Some(&runtime), &HeaderMap::new(), + Client::new, ) .await .unwrap(); @@ -398,6 +480,7 @@ async fn reusable_runtime_client_keeps_the_selected_accounts_transport_policy() factory, Some(&runtime), &HeaderMap::new(), + Client::new, ) .await .unwrap(); @@ -500,6 +583,7 @@ async fn reusable_runtime_client_refreshes_credentials_without_accepting_a_new_s HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), Some(&runtime), &HeaderMap::new(), + Client::new, ) .await .unwrap(); diff --git a/codex-rs/chatgpt/Cargo.toml b/codex-rs/chatgpt/Cargo.toml index 04b9cb7df6..252b380813 100644 --- a/codex-rs/chatgpt/Cargo.toml +++ b/codex-rs/chatgpt/Cargo.toml @@ -12,6 +12,8 @@ anyhow = { workspace = true } clap = { workspace = true, features = ["derive"] } codex-connectors = { workspace = true } codex-core = { workspace = true } +codex-features = { workspace = true } +codex-runtime-agent-instance = { workspace = true } codex-git-utils = { workspace = true } codex-http-client = { workspace = true } codex-login = { workspace = true } @@ -22,6 +24,10 @@ serde = { workspace = true, features = ["derive"] } tokio = { workspace = true, features = ["full"] } [dev-dependencies] +codex-protocol = { workspace = true } +base64 = { workspace = true } +uuid = { workspace = true } +wiremock = { workspace = true } codex-utils-cargo-bin = { workspace = true } pretty_assertions = { workspace = true } serde_json = { workspace = true } diff --git a/codex-rs/chatgpt/src/chatgpt_client.rs b/codex-rs/chatgpt/src/chatgpt_client.rs index ec6efb67fc..fd43eabff0 100644 --- a/codex-rs/chatgpt/src/chatgpt_client.rs +++ b/codex-rs/chatgpt/src/chatgpt_client.rs @@ -1,4 +1,6 @@ use codex_core::config::Config; +use codex_features::Feature; +use codex_features::RuntimeAuthMode; use codex_http_client::HttpClient; use codex_http_client::HttpClientFactory; use codex_login::AuthManager; @@ -6,6 +8,8 @@ use codex_login::CodexAuth; use codex_login::default_client::RESIDENCY_HEADER_NAME; use codex_login::default_client::create_client_with_chatgpt_cookies; use codex_login::default_client::default_headers; +use codex_runtime_agent_instance::RuntimeAuthManager; +use std::sync::Arc; use anyhow::Context; use serde::Serialize; @@ -55,27 +59,43 @@ pub(crate) async fn chatgpt_get_request( config: &Config, path: String, ) -> anyhow::Result { - chatgpt_get_request_with_timeout(config, path, /*timeout*/ None).await -} - -pub(crate) async fn chatgpt_get_request_with_timeout( - config: &Config, - path: String, - timeout: Option, -) -> anyhow::Result { - // Bind before loading credentials: the temporary auth manager does not own account changes. - let policy = config + // This helper is used by the standalone apply command. Harness callers pass their owner below. + // The temporary credential manager does not own account changes. + let mut captured_config = config.clone(); + captured_config.application_network_policy = config .application_network_policy .clone() .for_current_account(); - let client = chatgpt_client(config.http_client_factory().with_network_policy(policy)); - let chatgpt_base_url = &config.chatgpt_base_url; + let config = &captured_config; + if config.features.enabled(Feature::RuntimeAgentInstances) { + codex_model_provider::validate_runtime_destination(Some(&config.chatgpt_base_url))?; + } let auth_manager = AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await?; + let runtime = if config.features.enabled(Feature::RuntimeAgentInstances) + && codex_runtime_agent_instance::is_openai_api_url(&config.chatgpt_base_url) + { + let installation_id = codex_core::resolve_installation_id(&config.codex_home).await?; + codex_core::build_runtime_auth_manager(config, auth_manager.clone(), &installation_id) + .map(|runtime| (runtime, config.runtime_auth_mode)) + } else { + None + }; let auth = auth_manager .auth() .await .ok_or_else(|| anyhow::anyhow!("ChatGPT auth not available"))?; + chatgpt_get_request_with_timeout(config, &auth, runtime.as_ref(), path, /*timeout*/ None).await +} + +pub(crate) async fn chatgpt_get_request_with_timeout( + config: &Config, + auth: &CodexAuth, + runtime: Option<&(Arc, RuntimeAuthMode)>, + path: String, + timeout: Option, +) -> anyhow::Result { + let chatgpt_base_url = &config.chatgpt_base_url; anyhow::ensure!( auth.uses_codex_backend(), "ChatGPT backend requests require Codex backend auth" @@ -90,10 +110,16 @@ pub(crate) async fn chatgpt_get_request_with_timeout( path.trim_start_matches('/') ); + let (client, request_auth) = prepare_client_auth(config, auth, runtime, &url).await?; let mut request = client .get(&url) .headers(default_headers()) - .headers(codex_model_provider::auth_provider_from_auth(&auth).to_auth_headers()?) + .headers( + request_auth + .auth + .apply_auth_headers(Default::default()) + .await?, + ) .header(OAI_PRODUCT_SKU_HEADER, CODEX_PRODUCT_SKU) .header("Content-Type", "application/json"); if let Some(timeout) = timeout { @@ -128,6 +154,7 @@ pub(crate) async fn chatgpt_post_request_with_timeout< body: &TRequest, timeout: Duration, product_sku: &str, + runtime: Option<&(Arc, RuntimeAuthMode)>, ) -> anyhow::Result { anyhow::ensure!( auth.uses_codex_backend(), @@ -143,11 +170,16 @@ pub(crate) async fn chatgpt_post_request_with_timeout< config.chatgpt_base_url.trim_end_matches('/'), path.trim_start_matches('/') ); - let client = chatgpt_client(config.http_client_factory()); + let (client, request_auth) = prepare_client_auth(config, auth, runtime, &url).await?; let response = client .post(&url) .headers(default_headers()) - .headers(codex_model_provider::auth_provider_from_auth(auth).to_auth_headers()?) + .headers( + request_auth + .auth + .apply_auth_headers(Default::default()) + .await?, + ) .header(OAI_PRODUCT_SKU_HEADER, product_sku) .header("Content-Type", "application/json") .timeout(timeout) @@ -167,3 +199,28 @@ pub(crate) async fn chatgpt_post_request_with_timeout< anyhow::bail!("Request failed with status {status}: {body}") } } + +/// Keep request credentials and the transport bound to the same account snapshot. +async fn prepare_client_auth( + config: &Config, + auth: &CodexAuth, + runtime: Option<&(Arc, RuntimeAuthMode)>, + url: &str, +) -> anyhow::Result<(HttpClient, codex_model_provider::ResolvedProviderAuth)> { + if runtime.is_some() { + codex_model_provider::validate_runtime_destination(Some(url))?; + } + let resolved = codex_model_provider::auth_provider_for_request( + auth, + runtime, + /*headers*/ &Default::default(), + ) + .await?; + let factory = resolved.bind_http_client_factory(config.http_client_factory())?; + let client = chatgpt_client(factory); + Ok((client, resolved)) +} + +#[cfg(test)] +#[path = "chatgpt_client_runtime_auth_tests.rs"] +mod runtime_auth_tests; diff --git a/codex-rs/chatgpt/src/chatgpt_client_runtime_auth_tests.rs b/codex-rs/chatgpt/src/chatgpt_client_runtime_auth_tests.rs new file mode 100644 index 0000000000..d14d53828a --- /dev/null +++ b/codex-rs/chatgpt/src/chatgpt_client_runtime_auth_tests.rs @@ -0,0 +1,212 @@ +//! Test captured ChatGPT credentials and runtime proofs on helper requests. + +use super::*; +use base64::Engine; +use codex_features::RuntimeAuthMode; +use codex_runtime_agent_instance::AgentHostAuthorization; +use codex_runtime_agent_instance::RuntimeAuthManager; +use codex_runtime_agent_instance::RuntimeRegistrationMetadata; +use codex_runtime_agent_instance::RuntimeRole; +use pretty_assertions::assert_eq; +use serde_json::json; +use std::collections::HashSet; +use std::sync::Arc; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::path; + +#[tokio::test] +async fn helpers_use_captured_auth_and_strip_cross_origin_proofs() { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + let redirected = MockServer::start().await; + Mock::given(path("/")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"ok": true}))) + .expect(1) + .mount(&redirected) + .await; + Mock::given(path("/agent/register")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "runtime_agent_instance_id": "runtime-chatgpt", "agent_host_id": "host-a", + "actor": {"kind": "chatgpt", "user_id": "user-a", "account_id": "workspace-a"}, + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(path("/data")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"ok": true}))) + .expect(2) + .mount(&server) + .await; + Mock::given(path("/redirect")) + .respond_with(ResponseTemplate::new(307).insert_header("location", redirected.uri())) + .expect(1) + .mount(&server) + .await; + let home = tempfile::TempDir::new().unwrap(); + let mut config = codex_core::config::ConfigBuilder::default() + .codex_home(home.path().to_path_buf()) + .build() + .await + .unwrap(); + config.chatgpt_base_url = server.uri(); + let token = actor_token(); + let auth = CodexAuth::from_external_chatgpt_tokens( + &token, + "workspace-a", + /*chatgpt_plan_type*/ None, + ) + .unwrap(); + let runtime = ( + Arc::new(RuntimeAuthManager::new( + AuthManager::from_auth_for_testing(auth.clone()), + AgentHostAuthorization::InstallationId { + installation_id: uuid::Uuid::nil(), + }, + RuntimeRegistrationMetadata { + harness: Some("codex".into()), + harness_version: None, + role: RuntimeRole::Orchestrator, + }, + format!("{}/agent/register", server.uri()).parse().unwrap(), + )), + mode, + ); + // No saved auth exists: these helpers must use the supplied snapshot, not reload auth. + let get: serde_json::Value = chatgpt_get_request_with_timeout( + &config, + &auth, + Some(&runtime), + "data".into(), + /*timeout*/ None, + ) + .await + .unwrap(); + let post: serde_json::Value = chatgpt_post_request_with_timeout( + &config, + &auth, + "data".into(), + &json!({}), + Duration::from_secs(10), + "codex", + Some(&runtime), + ) + .await + .unwrap(); + assert_eq!(get, json!({"ok": true})); + assert_eq!(post, get); + let redirected_result = chatgpt_get_request_with_timeout::( + &config, + &auth, + Some(&runtime), + "redirect".into(), + /*timeout*/ None, + ) + .await + .unwrap(); + assert_eq!(redirected_result, get); + let forwarded = redirected.received_requests().await.unwrap(); + assert_eq!(forwarded.len(), 1); + assert!(!forwarded[0].headers.contains_key("authorization")); + assert!( + !forwarded[0] + .headers + .contains_key("x-openai-agent-assertion") + ); + redirected.verify().await; + let requests = server.received_requests().await.unwrap(); + let proofs: HashSet<_> = requests + .iter() + .filter(|r| r.url.path() != "/agent/register") + .map(|r| { + assert_eq!(r.headers["chatgpt-account-id"], "workspace-a"); + let authorization = r.headers["authorization"].to_str().unwrap(); + match mode { + RuntimeAuthMode::DualHeader => { + assert_eq!(authorization, format!("Bearer {token}")); + r.headers["x-openai-agent-assertion"].to_str().unwrap() + } + RuntimeAuthMode::AssertionOnly => { + assert!(!r.headers.contains_key("x-openai-agent-assertion")); + authorization.strip_prefix("AgentAssertion ").unwrap() + } + } + }) + .collect(); + assert_eq!(proofs.len(), 3); + server.verify().await; + } +} + +fn actor_token() -> String { + let claims = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(json!({ + "iss": "https://auth.openai.com/", "session_id": "login-a", + "https://api.openai.com/auth": {"user_id": "user-a", "chatgpt_account_id": "workspace-a"}, + }).to_string()); + format!("header.{claims}.signature") +} + +#[tokio::test] +async fn standalone_apply_keeps_custom_backend_auth_when_runtime_feature_is_enabled() { + let server = MockServer::start().await; + let token = actor_token(); + Mock::given(path("/task")) + .and(wiremock::matchers::header( + "authorization", + format!("Bearer {token}"), + )) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"ok": true}))) + .expect(1) + .mount(&server) + .await; + let home = tempfile::TempDir::new().unwrap(); + codex_login::auth::login_with_chatgpt_auth_tokens( + home.path(), + &token, + "workspace-a", + /*chatgpt_plan_type*/ None, + ) + .unwrap(); + let mut config = codex_core::config::ConfigBuilder::default() + .codex_home(home.path().to_path_buf()) + .build() + .await + .unwrap(); + config.chatgpt_base_url = server.uri(); + config + .features + .enable(Feature::RuntimeAgentInstances) + .unwrap(); + let response: serde_json::Value = chatgpt_get_request(&config, "task".into()).await.unwrap(); + assert_eq!(response, json!({"ok": true})); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 1); + assert!(!requests[0].headers.contains_key("x-openai-agent-assertion")); + server.verify().await; +} + +#[tokio::test] +async fn standalone_apply_rejects_unsupported_openai_destination() { + let home = tempfile::TempDir::new().unwrap(); + let mut config = codex_core::config::ConfigBuilder::default() + .codex_home(home.path().to_path_buf()) + .build() + .await + .unwrap(); + config.chatgpt_base_url = "https://future.openai.com/backend-api".into(); + config + .features + .enable(Feature::RuntimeAgentInstances) + .unwrap(); + + let error = chatgpt_get_request::(&config, "task".into()) + .await + .unwrap_err(); + assert!(matches!( + error + .downcast_ref::() + .map(codex_protocol::error::CodexErr::details), + Some(codex_protocol::error::CodexErrorDetails::UnsupportedOperation(_)) + )); +} diff --git a/codex-rs/chatgpt/src/connectors.rs b/codex-rs/chatgpt/src/connectors.rs index 51a5fbe609..d0c5d82907 100644 --- a/codex-rs/chatgpt/src/connectors.rs +++ b/codex-rs/chatgpt/src/connectors.rs @@ -103,8 +103,15 @@ pub async fn list_all_connectors_with_options( if !apps_enabled(config).await? { return Ok(Vec::new()); } + let mut captured_config = config.clone(); + captured_config.application_network_policy = config + .application_network_policy + .clone() + .for_current_account(); + let config = &captured_config; let auth = connector_auth(config).await?; let cache_context = connector_directory_cache_context(config, &auth); + let auth = &auth; let connectors = codex_connectors::list_all_connectors_with_options( cache_context, auth.is_workspace_account(), @@ -112,6 +119,8 @@ pub async fn list_all_connectors_with_options( |path| async move { chatgpt_get_request_with_timeout::( config, + auth, + /*runtime*/ None, path, Some(DIRECTORY_CONNECTORS_TIMEOUT), ) @@ -179,6 +188,7 @@ pub async fn read_connector_metadata( }, CONNECTOR_METADATA_TIMEOUT, product_sku, + /*runtime*/ None, ) .await?; let mut requested_ids = missing_ids.iter().cloned().collect::>(); diff --git a/codex-rs/cli/tests/cloud_auth.rs b/codex-rs/cli/tests/cloud_auth.rs index 01e7eb87e3..a1c46efeea 100644 --- a/codex-rs/cli/tests/cloud_auth.rs +++ b/codex-rs/cli/tests/cloud_auth.rs @@ -11,7 +11,7 @@ use wiremock::matchers::header; use wiremock::matchers::path; #[tokio::test] -async fn cloud_list_only_allows_trusted_credential_destinations() -> Result<()> { +async fn cloud_list_respects_the_credential_store_and_trusted_destinations() -> Result<()> { let server = MockServer::start().await; let codex_home = TempDir::new()?; std::fs::write( @@ -66,6 +66,28 @@ async fn cloud_list_only_allows_trusted_credential_destinations() -> Result<()> insta::assert_snapshot!(String::from_utf8(output)?, @"Error: CODEX_CLOUD_TASKS_BASE_URL must use a trusted HTTPS origin on port 443, without user information, a query, or a fragment; custom backends cannot use saved ChatGPT credentials"); assert!(server.received_requests().await.unwrap().is_empty()); + // An explicit ephemeral store must exclude the saved login for Cloud itself. + command()? + .env( + "CODEX_CLOUD_TASKS_BASE_URL", + "https://chatgpt-staging.com/backend-api", + ) + .env("HTTPS_PROXY", server.uri()) + .env("https_proxy", server.uri()) + .args([ + "-c", + "cli_auth_credentials_store='ephemeral'", + "cloud", + "list", + "--limit", + "1", + "--json", + ]) + .assert() + .failure() + .stderr(contains("Not signed in. Please run 'codex login'")); + assert!(server.received_requests().await.unwrap().is_empty()); + // Staging must reach explicit PAT authentication. Reject the synthetic PAT locally so // this test never sends a request to the real staging backend. let auth_server = MockServer::start().await; diff --git a/codex-rs/cloud-tasks-client/src/http.rs b/codex-rs/cloud-tasks-client/src/http.rs index c70ae596e8..b6ea46bbf9 100644 --- a/codex-rs/cloud-tasks-client/src/http.rs +++ b/codex-rs/cloud-tasks-client/src/http.rs @@ -37,6 +37,10 @@ impl HttpClient { Self { base_url, backend } } + pub fn from_backend(base_url: String, backend: backend::Client) -> Self { + Self { base_url, backend } + } + pub fn with_user_agent(mut self, ua: impl Into) -> Self { self.backend = self.backend.clone().with_user_agent(ua); self diff --git a/codex-rs/cloud-tasks/Cargo.toml b/codex-rs/cloud-tasks/Cargo.toml index 12bd6fdf7f..26564894e8 100644 --- a/codex-rs/cloud-tasks/Cargo.toml +++ b/codex-rs/cloud-tasks/Cargo.toml @@ -13,6 +13,7 @@ doctest = false workspace = true [dependencies] +codex-backend-client = { workspace = true } anyhow = { workspace = true } chrono = { workspace = true, features = ["serde"] } clap = { workspace = true, features = ["derive"] } @@ -21,6 +22,7 @@ codex-cloud-tasks-client = { workspace = true } # TODO: codex-cloud-tasks-mock-client should be in dev-dependencies. codex-cloud-tasks-mock-client = { workspace = true } codex-core = { workspace = true } +codex-features = { workspace = true } codex-git-utils = { workspace = true } codex-login = { path = "../login" } codex-model-provider = { workspace = true } @@ -41,5 +43,12 @@ unicode-segmentation = { workspace = true } unicode-width = { workspace = true } [dev-dependencies] +codex-agent-identity = { workspace = true } +core_test_support = { workspace = true } +base64 = { workspace = true } +tempfile = { workspace = true } +codex-runtime-agent-instance = { workspace = true } +uuid = { workspace = true } +wiremock = { workspace = true } insta = { workspace = true } pretty_assertions = { workspace = true } diff --git a/codex-rs/cloud-tasks/src/backend_auth_tests.rs b/codex-rs/cloud-tasks/src/backend_auth_tests.rs new file mode 100644 index 0000000000..1350158782 --- /dev/null +++ b/codex-rs/cloud-tasks/src/backend_auth_tests.rs @@ -0,0 +1,277 @@ +//! Cover credential refresh between Cloud startup and environment discovery. + +#![cfg(test)] + +use base64::Engine; +use chrono::Utc; +use codex_utils_cli::CliConfigOverrides; +use pretty_assertions::assert_eq; +use serde_json::json; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::header; +use wiremock::matchers::method; +use wiremock::matchers::path; + +fn access_token(expires_at: i64) -> String { + let claims = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode( + json!({ + "exp": expires_at, + "iss": "https://auth.openai.com/", + "sub": "user-a", + "session_id": "session-a", + "https://api.openai.com/auth": { + "chatgpt_account_id": "account-a", + "chatgpt_user_id": "user-a", + "chatgpt_plan_type": "plus", + }, + }) + .to_string(), + ); + format!("header.{claims}.signature") +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn environment_discovery_refreshes_after_cloud_startup() { + const SERVER: &str = "CODEX_CLOUD_REFRESH_TEST_SERVER"; + if let Ok(server) = std::env::var(SERVER) { + let overrides = CliConfigOverrides { + raw_overrides: vec![ + "cli_auth_credentials_store=\"file\"".into(), + "features.runtime_agent_instances=false".into(), + "features.use_agent_identity=false".into(), + ], + }; + let context = super::init_backend("cloud-refresh-test", &overrides) + .await + .expect("Cloud startup"); + for _ in 0..2 { + let environments = crate::env_detect::list_environments( + &context.environment_http, + &server, + &crate::util::build_chatgpt_headers(), + ) + .await + .expect("environment discovery after token refresh"); + assert_eq!(environments.len(), 1); + assert_eq!(environments[0].id, "env-a"); + } + return; + } + + let server = MockServer::start().await; + let home = tempfile::tempdir().expect("isolated Cloud credentials"); + let expired = access_token(Utc::now().timestamp() - 60); + // The startup token is near expiry, so discovery must refresh it through AuthManager. + let startup = access_token(Utc::now().timestamp() + 240); + let current = access_token(Utc::now().timestamp() + 3_600); + std::fs::write( + home.path().join("auth.json"), + json!({ + "tokens": { + "id_token": expired, "access_token": expired, + "refresh_token": "refresh-initial", "account_id": "account-a", + }, + "last_refresh": Utc::now(), + }) + .to_string(), + ) + .expect("write initial credentials"); + let refreshed = current.clone(); + Mock::given(method("POST")) + .and(path("/token")) + .respond_with(move |request: &wiremock::Request| { + let body: serde_json::Value = request.body_json().unwrap(); + let token = match body["refresh_token"].as_str().unwrap() { + "refresh-initial" => &startup, + "refresh-startup" => &refreshed, + other => panic!("unexpected test refresh token: {other}"), + }; + ResponseTemplate::new(200).set_body_json(json!({ + "id_token": token, "access_token": token, + "refresh_token": "refresh-startup", + })) + }) + .expect(2) + .mount(&server) + .await; + Mock::given(path("/api/codex/environments")) + .respond_with(ResponseTemplate::new(401)) + .with_priority(10) + .expect(0) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/api/codex/environments")) + .and(header("authorization", format!("Bearer {current}"))) + .and(header("chatgpt-account-id", "account-a")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!([ + {"id": "env-a", "label": "Environment A"}, + ]))) + .expect(2) + .mount(&server) + .await; + + let mut command = std::process::Command::new(std::env::current_exe().unwrap()); + command + .args([ + "--exact", + "backend_auth_tests::environment_discovery_refreshes_after_cloud_startup", + "--nocapture", + ]) + .current_dir(home.path()) + .env(SERVER, server.uri()) + .env("CODEX_HOME", home.path()) + .env( + "CODEX_REFRESH_TOKEN_URL_OVERRIDE", + format!("{}/token", server.uri()), + ) + .env( + "CODEX_CLOUD_TASKS_BASE_URL", + "https://chatgpt.com/backend-api", + ) + .env_remove("CODEX_CLOUD_TASKS_MODE") + .env_remove("CODEX_ACCESS_TOKEN") + .env_remove("CODEX_API_KEY") + .env_remove("OPENAI_API_KEY"); + let output = tokio::task::spawn_blocking(move || command.output()) + .await + .expect("Cloud test child task") + .expect("Cloud test child process"); + assert!( + output.status.success(), + "Cloud startup regression failed:\n{}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr), + ); + server.verify().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cloud_list_preserves_existing_agent_identity_during_runtime_rollout() -> anyhow::Result<()> +{ + use core_test_support::runtime_auth::PROBE_SERVER; + use core_test_support::runtime_auth::run_harness_through_proxy; + + const TEST: &str = + "backend_auth_tests::cloud_list_preserves_existing_agent_identity_during_runtime_rollout"; + const MODE: &str = "CODEX_CLOUD_PROGRAMMATIC_TEST_MODE"; + if let Ok(mode) = std::env::var(MODE) { + let enabled = mode != "off"; + let mode = if enabled { + mode.as_str() + } else { + "dual_header" + }; + let overrides = CliConfigOverrides { + raw_overrides: vec![ + "cli_auth_credentials_store=\"file\"".into(), + format!("features.runtime_agent_instances={{ enabled={enabled}, mode={mode:?} }}"), + "features.use_agent_identity=false".into(), + ], + }; + let context = super::init_backend("cloud-programmatic-test", &overrides).await?; + let page = codex_cloud_tasks_client::CloudBackend::list_tasks( + &*context.backend, + /*env*/ None, + /*limit*/ Some(10), + /*cursor*/ None, + ) + .await?; + assert!(page.tasks.is_empty()); + let environments = crate::env_detect::list_environments( + &context.environment_http, + "https://chatgpt.com", + &crate::util::build_chatgpt_headers(), + ) + .await?; + assert_eq!(environments.len(), 1); + assert_eq!(environments[0].id, "env-a"); + return Ok(()); + } + if std::env::var(PROBE_SERVER).is_ok() { + let home = tempfile::tempdir()?; + let key = codex_agent_identity::generate_agent_key_material()?; + std::fs::write( + home.path().join("auth.json"), + json!({ + "auth_mode": "agentIdentity", + "agent_identity": { + "agent_runtime_id": "existing-agent", + "agent_private_key": key.private_key_pkcs8_base64, + "account_id": "account-a", + "chatgpt_user_id": "user-a", + "plan_type": "plus", + "chatgpt_account_is_fedramp": false, + "task_id": "existing-task", + }, + }) + .to_string(), + )?; + // Keep each credential/configuration case in its own process, retaining the proxy + // from the outer fixture without mutating environment variables in a running test. + for mode in ["off", "dual_header", "assertion_only"] { + let mut command = std::process::Command::new(std::env::current_exe()?); + command + .args(["--exact", TEST, "--nocapture"]) + .current_dir(home.path()) + .env(MODE, mode) + .env("CODEX_HOME", home.path()) + .env( + "CODEX_CLOUD_TASKS_BASE_URL", + "https://chatgpt.com/backend-api", + ) + .env_remove("CODEX_CLOUD_TASKS_MODE") + .env_remove("CODEX_ACCESS_TOKEN") + .env_remove("CODEX_API_KEY") + .env_remove("OPENAI_API_KEY"); + let output = tokio::task::spawn_blocking(move || command.output()).await??; + assert!( + output.status.success(), + "{mode}: {}\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + assert!(String::from_utf8_lossy(&output.stdout).contains("1 passed")); + } + return Ok(()); + } + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/backend-api/wham/tasks/list")) + .and(header("chatgpt-account-id", "account-a")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"items": [], "cursor": null})), + ) + .expect(3) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/api/codex/environments")) + .and(header("chatgpt-account-id", "account-a")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(json!([{ "id": "env-a", "label": "Environment A" }])), + ) + .expect(3) + .mount(&server) + .await; + run_harness_through_proxy(&server, TEST, /*realtime_target*/ None).await?; + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 6); + for request in requests { + let assertion = request.headers["authorization"] + .to_str()? + .strip_prefix("AgentAssertion ") + .unwrap(); + let envelope: serde_json::Value = serde_json::from_slice( + &base64::engine::general_purpose::URL_SAFE_NO_PAD.decode(assertion)?, + )?; + assert_eq!(envelope["agent_runtime_id"], "existing-agent"); + assert_eq!(envelope["task_id"], "existing-task"); + assert!(!request.headers.contains_key("x-openai-agent-assertion")); + } + server.verify().await; + Ok(()) +} diff --git a/codex-rs/cloud-tasks/src/env_detect.rs b/codex-rs/cloud-tasks/src/env_detect.rs index 31fbc99e2f..d2535e0a9f 100644 --- a/codex-rs/cloud-tasks/src/env_detect.rs +++ b/codex-rs/cloud-tasks/src/env_detect.rs @@ -1,5 +1,6 @@ use codex_git_utils::SanitizedGitUrl; use codex_http_client::RouteAwareClientPool; +use codex_model_provider::ClientAuth; use http::StatusCode; use http::header::CONTENT_TYPE; use http::header::HeaderMap; @@ -7,6 +8,12 @@ use std::collections::HashMap; use tracing::info; use tracing::warn; +#[derive(Clone)] +pub(crate) struct EnvironmentClient { + pub(crate) http: RouteAwareClientPool, + pub(crate) auth: Option, +} + #[derive(Debug, Clone, serde::Deserialize)] struct CodeEnvironment { id: String, @@ -24,8 +31,8 @@ pub struct AutodetectSelection { pub label: Option, } -pub async fn autodetect_environment_id( - http: &RouteAwareClientPool, +pub(crate) async fn autodetect_environment_id( + http: &EnvironmentClient, base_url: &str, headers: &HeaderMap, desired_label: Option, @@ -195,12 +202,19 @@ trait EnvironmentHttp: Send + Sync { ) -> impl std::future::Future> + Send + 'a; } -impl EnvironmentHttp for RouteAwareClientPool { +impl EnvironmentHttp for EnvironmentClient { async fn get(&self, url: &str, headers: &HeaderMap) -> anyhow::Result { - let response = RouteAwareClientPool::get(self, url) - .headers(headers.clone()) - .send() - .await?; + let headers = match &self.auth { + Some(auth) => { + auth.resolve(headers) + .await? + .auth + .apply_auth_headers(headers.clone()) + .await? + } + None => headers.clone(), + }; + let response = self.http.get(url).headers(headers).send().await?; Ok(EnvironmentResponse { status: response.status(), content_type: response @@ -306,8 +320,8 @@ fn parse_owner_repo(url: &str) -> Option<(String, String)> { /// List environments for the current repo(s) with a fallback to the global list. /// Returns a de-duplicated, sorted set suitable for the TUI modal. -pub async fn list_environments( - http: &RouteAwareClientPool, +pub(crate) async fn list_environments( + http: &EnvironmentClient, base_url: &str, headers: &HeaderMap, ) -> anyhow::Result> { @@ -426,3 +440,7 @@ async fn list_environments_with_origins( #[cfg(test)] #[path = "env_detect_tests.rs"] mod tests; + +#[cfg(test)] +#[path = "env_detect_runtime_auth_tests.rs"] +mod runtime_auth_tests; diff --git a/codex-rs/cloud-tasks/src/env_detect_runtime_auth_tests.rs b/codex-rs/cloud-tasks/src/env_detect_runtime_auth_tests.rs new file mode 100644 index 0000000000..9cb8d92b2c --- /dev/null +++ b/codex-rs/cloud-tasks/src/env_detect_runtime_auth_tests.rs @@ -0,0 +1,130 @@ +//! Environment discovery must resolve runtime authentication for each HTTP request. + +use super::*; +use codex_features::RuntimeAuthMode; +use codex_http_client::ClientRouteClass; +use codex_http_client::HttpClientFactory; +use codex_http_client::OutboundProxyPolicy; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_runtime_agent_instance::AgentHostAuthorization; +use codex_runtime_agent_instance::RuntimeAuthManager; +use codex_runtime_agent_instance::RuntimeRegistrationMetadata; +use codex_runtime_agent_instance::RuntimeRole; +use pretty_assertions::assert_eq; +use std::collections::HashSet; +use std::sync::Arc; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::path; + +struct ChangingCredentials(std::sync::RwLock); + +impl codex_login::ExternalAuth for ChangingCredentials { + fn resolve(&self) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Ok(self.0.read().unwrap().clone()) }) + } + + fn refresh( + &self, + _: codex_login::ExternalAuthRefreshContext, + ) -> codex_login::ExternalAuthFuture<'_, CodexAuth> { + self.resolve() + } +} + +#[tokio::test] +async fn discovery_signs_each_get_and_rejects_a_changed_credential_source() { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + Mock::given(path("/agent/register")).respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "runtime_agent_instance_id": "runtime-cloud", "agent_host_id": "host-a", + "actor": {"kind": "platform", "user_id": "user-a", "organization_id": "org-a", "project_id": null}, + }))).mount(&server).await; + Mock::given(path("/api/codex/environments")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!([]))) + .expect(2) + .mount(&server) + .await; + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); + let auth = CodexAuth::from_api_key("key-a"); + let credentials = AuthManager::from_auth_for_testing(auth.clone()); + let source = Arc::new(ChangingCredentials(std::sync::RwLock::new(auth.clone()))); + credentials.set_external_auth(source.clone()).await.unwrap(); + let runtime = ( + Arc::new(RuntimeAuthManager::new( + credentials, + AgentHostAuthorization::InstallationId { + installation_id: uuid::Uuid::nil(), + }, + RuntimeRegistrationMetadata { + harness: Some("codex".into()), + harness_version: None, + role: RuntimeRole::Orchestrator, + }, + format!("{}/agent/register", server.uri()).parse().unwrap(), + )), + mode, + ); + let (auth, factory) = ClientAuth::for_auth( + &auth, + Some(&runtime), + factory, + /*headers*/ &Default::default(), + ) + .await + .unwrap(); + let client = EnvironmentClient { + http: RouteAwareClientPool::new_without_redirects_or_request_logging( + factory, + ClientRouteClass::Api, + ), + auth: Some(auth), + }; + for _ in 0..2 { + assert!( + list_environments_with_origins(&client, &server.uri(), &HeaderMap::new(), &[]) + .await + .unwrap() + .is_empty() + ); + } + assert_eq!( + server + .received_requests() + .await + .unwrap() + .iter() + .filter(|request| request.url.path() == "/agent/register") + .count(), + 1, + ); + *source.0.write().unwrap() = CodexAuth::from_api_key("key-b"); + assert!( + list_environments_with_origins(&client, &server.uri(), &HeaderMap::new(), &[]) + .await + .is_err() + ); + let requests = server.received_requests().await.unwrap(); + let proofs: HashSet<_> = requests + .iter() + .filter(|r| r.url.path() != "/agent/register") + .map(|r| { + let authorization = r.headers["authorization"].to_str().unwrap(); + match mode { + RuntimeAuthMode::DualHeader => { + assert_eq!(authorization, "Bearer key-a"); + r.headers["x-openai-agent-assertion"].to_str().unwrap() + } + RuntimeAuthMode::AssertionOnly => { + assert!(!r.headers.contains_key("x-openai-agent-assertion")); + authorization.strip_prefix("AgentAssertion ").unwrap() + } + } + }) + .collect(); + assert_eq!(proofs.len(), 2); + server.verify().await; + } +} diff --git a/codex-rs/cloud-tasks/src/env_detect_tests.rs b/codex-rs/cloud-tasks/src/env_detect_tests.rs index 3299ff87f3..2b14833d19 100644 --- a/codex-rs/cloud-tasks/src/env_detect_tests.rs +++ b/codex-rs/cloud-tasks/src/env_detect_tests.rs @@ -77,10 +77,13 @@ async fn production_http_forwards_headers_and_decodes_response() { .expect("environment HTTP response should write"); String::from_utf8(request).expect("environment HTTP request should be UTF-8") }); - let http = RouteAwareClientPool::new_without_request_logging( - HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), - ClientRouteClass::Api, - ); + let http = EnvironmentClient { + http: RouteAwareClientPool::new_without_request_logging( + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ClientRouteClass::Api, + ), + auth: None, + }; let base_url = format!("http://{address}"); let headers = HeaderMap::from_iter([(AUTHORIZATION, HeaderValue::from_static("Bearer real-token"))]); diff --git a/codex-rs/cloud-tasks/src/lib.rs b/codex-rs/cloud-tasks/src/lib.rs index c3f43f3e19..864e9e067c 100644 --- a/codex-rs/cloud-tasks/src/lib.rs +++ b/codex-rs/cloud-tasks/src/lib.rs @@ -1,4 +1,6 @@ mod app; +#[cfg(test)] +mod backend_auth_tests; mod cli; pub(crate) mod env_detect; mod new_task; @@ -13,10 +15,12 @@ use codex_cloud_tasks_client::TaskStatus; use codex_git_utils::current_branch_name; use codex_git_utils::default_branch_name; use codex_http_client::ClientRouteClass; -use codex_http_client::HttpClientFactory; -use codex_http_client::OutboundProxyPolicy; use codex_http_client::RouteAwareClientPool; +use codex_login::CodexAuth; use codex_login::default_client::get_codex_user_agent; +use codex_model_provider::ClientAuth; +use codex_model_provider::auth_provider_from_auth_manager; +use codex_utils_cli::CliConfigOverrides; use owo_colors::OwoColorize; use owo_colors::Stream; use std::cmp::Ordering; @@ -42,10 +46,13 @@ struct ApplyJob { struct BackendContext { backend: Arc, base_url: String, - environment_http: RouteAwareClientPool, + environment_http: env_detect::EnvironmentClient, } -async fn init_backend(user_agent_suffix: &str) -> anyhow::Result { +async fn init_backend( + user_agent_suffix: &str, + config_overrides: &CliConfigOverrides, +) -> anyhow::Result { #[cfg(debug_assertions)] let use_mock = matches!( std::env::var("CODEX_CLOUD_TASKS_MODE").ok().as_deref(), @@ -57,27 +64,73 @@ async fn init_backend(user_agent_suffix: &str) -> anyhow::Result set_user_agent_suffix(user_agent_suffix); + let ua = get_codex_user_agent(); + let (auth_manager, http_client_factory, config) = + util::load_auth_manager(Some(base_url.clone()), config_overrides).await?; + let auth = match auth_manager.as_ref() { + Some(manager) => manager.auth().await, + None => None, + }; + // Existing agent assertions and supplied headers cannot register a new runtime. + let runtime = if let (Some(config), Some(auth_manager)) = (config, auth_manager.as_ref()) + && config + .features + .enabled(codex_features::Feature::RuntimeAgentInstances) + && matches!( + auth.as_ref(), + Some( + CodexAuth::Chatgpt(_) + | CodexAuth::ChatgptAuthTokens(_) + | CodexAuth::PersonalAccessToken(_) + ) + ) { + // Use the validated command destination for registration as well as actor refresh. + let mut config = config; + config.chatgpt_base_url = base_url.clone(); + let installation_id = codex_core::resolve_installation_id(&config.codex_home).await?; + codex_core::build_runtime_auth_manager(&config, Arc::clone(auth_manager), &installation_id) + .map(|runtime| (runtime, config.runtime_auth_mode)) + } else { + None + }; + let (environment_auth, http_client_factory) = + match auth.as_ref().filter(|auth| auth.uses_codex_backend()) { + Some(auth) => { + let (mut client_auth, factory) = ClientAuth::for_auth( + auth, + runtime.as_ref(), + http_client_factory, + /*headers*/ &Default::default(), + ) + .await?; + if runtime.is_none() + && let Some(auth_manager) = auth_manager.as_ref() + { + client_auth.set_provider(auth_provider_from_auth_manager( + Arc::clone(auth_manager), + auth, + ))?; + } + (Some(client_auth), factory) + } + None => (None, http_client_factory), + }; + let environment_http = env_detect::EnvironmentClient { + http: RouteAwareClientPool::new_without_redirects_or_request_logging( + http_client_factory.clone(), + ClientRouteClass::Api, + ), + auth: environment_auth, + }; #[cfg(debug_assertions)] if use_mock { - let http_client_factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); return Ok(BackendContext { backend: Arc::new(codex_cloud_tasks_mock_client::MockClient), base_url, - environment_http: RouteAwareClientPool::new_without_redirects_or_request_logging( - http_client_factory, - ClientRouteClass::Api, - ), + environment_http, }); } - let ua = get_codex_user_agent(); - let (auth_manager, http_client_factory) = util::load_auth_manager(Some(base_url.clone())).await; - let environment_http = RouteAwareClientPool::new_without_redirects_or_request_logging( - http_client_factory.clone(), - ClientRouteClass::Api, - ); - let mut http = codex_cloud_tasks_client::HttpClient::new(base_url.clone(), http_client_factory) - .with_user_agent(ua); let style = if base_url.contains("/backend-api") { "wham" } else { @@ -85,10 +138,6 @@ async fn init_backend(user_agent_suffix: &str) -> anyhow::Result }; append_error_log(format!("startup: base_url={base_url} path_style={style}")); - let auth = match auth_manager.as_ref() { - Some(manager) => manager.auth().await, - None => None, - }; let auth = match auth { Some(auth) => auth, None => { @@ -110,8 +159,18 @@ async fn init_backend(user_agent_suffix: &str) -> anyhow::Result std::process::exit(1); } - let auth_provider = codex_model_provider::auth_provider_from_auth(&auth); - http = http.with_auth_provider(auth_provider)?; + let backend = codex_backend_client::Client::from_auth_with_runtime_without_redirects( + base_url.clone(), + &auth, + http_client_factory, + runtime.as_ref(), + /*headers*/ &Default::default(), + ) + .await?; + let http = codex_cloud_tasks_client::HttpClient::from_backend( + base_url.clone(), + backend.with_user_agent(ua), + ); if let Some(acc) = auth.get_account_id() { append_error_log(format!("auth: set ChatGPT-Account-Id header: {acc}")); } @@ -175,14 +234,16 @@ async fn resolve_git_ref_with_git_info( } } -async fn run_exec_command(args: crate::cli::ExecCommand) -> anyhow::Result<()> { +async fn run_exec_command( + args: crate::cli::ExecCommand, + ctx: BackendContext, +) -> anyhow::Result<()> { let crate::cli::ExecCommand { query, environment, branch, attempts, } = args; - let ctx = init_backend("codex_cloud_tasks_exec").await?; let prompt = resolve_query_input(query)?; let env_id = resolve_environment_id(&ctx, &environment).await?; let git_ref = resolve_git_ref(branch.as_ref()).await; @@ -206,7 +267,7 @@ async fn resolve_environment_id(ctx: &BackendContext, requested: &str) -> anyhow return Err(anyhow!("environment id must not be empty")); } let normalized = util::normalize_base_url(&ctx.base_url); - let headers = util::build_chatgpt_headers().await?; + let headers = util::build_chatgpt_headers(); let environments = crate::env_detect::list_environments(&ctx.environment_http, &normalized, &headers).await?; if environments.is_empty() { @@ -512,8 +573,10 @@ fn format_task_list_lines( lines } -async fn run_status_command(args: crate::cli::StatusCommand) -> anyhow::Result<()> { - let ctx = init_backend("codex_cloud_tasks_status").await?; +async fn run_status_command( + args: crate::cli::StatusCommand, + ctx: BackendContext, +) -> anyhow::Result<()> { let task_id = parse_task_id(&args.task_id)?; let summary = codex_cloud_tasks_client::CloudBackend::get_task_summary(&*ctx.backend, task_id).await?; @@ -528,8 +591,10 @@ async fn run_status_command(args: crate::cli::StatusCommand) -> anyhow::Result<( Ok(()) } -async fn run_list_command(args: crate::cli::ListCommand) -> anyhow::Result<()> { - let ctx = init_backend("codex_cloud_tasks_list").await?; +async fn run_list_command( + args: crate::cli::ListCommand, + ctx: BackendContext, +) -> anyhow::Result<()> { let env_filter = if let Some(env) = args.environment { Some(resolve_environment_id(&ctx, &env).await?) } else { @@ -595,8 +660,10 @@ async fn run_list_command(args: crate::cli::ListCommand) -> anyhow::Result<()> { Ok(()) } -async fn run_diff_command(args: crate::cli::DiffCommand) -> anyhow::Result<()> { - let ctx = init_backend("codex_cloud_tasks_diff").await?; +async fn run_diff_command( + args: crate::cli::DiffCommand, + ctx: BackendContext, +) -> anyhow::Result<()> { let task_id = parse_task_id(&args.task_id)?; let attempts = collect_attempt_diffs(&*ctx.backend, &task_id).await?; let selected = select_attempt(&attempts, args.attempt)?; @@ -604,8 +671,10 @@ async fn run_diff_command(args: crate::cli::DiffCommand) -> anyhow::Result<()> { Ok(()) } -async fn run_apply_command(args: crate::cli::ApplyCommand) -> anyhow::Result<()> { - let ctx = init_backend("codex_cloud_tasks_apply").await?; +async fn run_apply_command( + args: crate::cli::ApplyCommand, + ctx: BackendContext, +) -> anyhow::Result<()> { let task_id = parse_task_id(&args.task_id)?; let attempts = collect_attempt_diffs(&*ctx.backend, &task_id).await?; let selected = select_attempt(&attempts, args.attempt)?; @@ -751,16 +820,31 @@ fn spawn_apply( /// Entry point for the `codex cloud` subcommand. pub async fn run_main(cli: Cli, _codex_linux_sandbox_exe: Option) -> anyhow::Result<()> { + // Keep the Cloud future out of the CLI dispatcher frame while the TUI is running. + Box::pin(run_cloud(cli)).await +} + +async fn run_cloud(cli: Cli) -> anyhow::Result<()> { + let backend = |suffix| init_backend(suffix, &cli.config_overrides); if let Some(command) = cli.command { return match command { - crate::cli::Command::Exec(args) => run_exec_command(args).await, - crate::cli::Command::Status(args) => run_status_command(args).await, - crate::cli::Command::List(args) => run_list_command(args).await, - crate::cli::Command::Apply(args) => run_apply_command(args).await, - crate::cli::Command::Diff(args) => run_diff_command(args).await, + crate::cli::Command::Exec(args) => { + run_exec_command(args, backend("codex_cloud_tasks_exec").await?).await + } + crate::cli::Command::Status(args) => { + run_status_command(args, backend("codex_cloud_tasks_status").await?).await + } + crate::cli::Command::List(args) => { + run_list_command(args, backend("codex_cloud_tasks_list").await?).await + } + crate::cli::Command::Apply(args) => { + run_apply_command(args, backend("codex_cloud_tasks_apply").await?).await + } + crate::cli::Command::Diff(args) => { + run_diff_command(args, backend("codex_cloud_tasks_diff").await?).await + } }; } - let Cli { .. } = cli; // Very minimal logging setup; mirrors other crates' pattern. let default_level = "error"; @@ -779,7 +863,7 @@ pub async fn run_main(cli: Cli, _codex_linux_sandbox_exe: Option) -> an backend, base_url, environment_http, - } = init_backend("codex_cloud_tasks_tui").await?; + } = backend("codex_cloud_tasks_tui").await?; // Terminal setup use crossterm::ExecutableCommand; @@ -869,14 +953,8 @@ pub async fn run_main(cli: Cli, _codex_linux_sandbox_exe: Option) -> an let environment_http = environment_http.clone(); tokio::spawn(async move { let base_url = util::normalize_base_url(&base_url); - // Build headers: UA + ChatGPT auth if available - let headers = match util::build_chatgpt_headers().await { - Ok(headers) => headers, - Err(error) => { - let _ = tx.send(app::AppEvent::EnvironmentAutodetected(Err(error))); - return; - } - }; + // Authentication is resolved by the environment client for each request. + let headers = util::build_chatgpt_headers(); // Run autodetect. If it fails, we keep using "All". let res = crate::env_detect::autodetect_environment_id( @@ -2032,14 +2110,12 @@ pub async fn run_main(cli: Cli, _codex_linux_sandbox_exe: Option) -> an fn spawn_environment_load( tx: UnboundedSender, base_url: String, - http: RouteAwareClientPool, + http: env_detect::EnvironmentClient, ) { tokio::spawn(async move { let base_url = util::normalize_base_url(&base_url); - let result = match util::build_chatgpt_headers().await { - Ok(headers) => crate::env_detect::list_environments(&http, &base_url, &headers).await, - Err(error) => Err(error), - }; + let headers = util::build_chatgpt_headers(); + let result = crate::env_detect::list_environments(&http, &base_url, &headers).await; let _ = tx.send(app::AppEvent::EnvironmentsLoaded(result)); }); } diff --git a/codex-rs/cloud-tasks/src/util.rs b/codex-rs/cloud-tasks/src/util.rs index f5f770b197..da1a4b64c5 100644 --- a/codex-rs/cloud-tasks/src/util.rs +++ b/codex-rs/cloud-tasks/src/util.rs @@ -7,6 +7,7 @@ use codex_core::config::Config; use codex_http_client::HttpClientFactory; use codex_http_client::OutboundProxyPolicy; use codex_login::AuthManager; +use codex_utils_cli::CliConfigOverrides; use std::sync::Arc; pub fn set_user_agent_suffix(suffix: &str) { @@ -78,16 +79,19 @@ pub(crate) fn validate_chatgpt_base_url(input: &str) -> anyhow::Result { pub async fn load_auth_manager( chatgpt_base_url: Option, -) -> (Option>, HttpClientFactory) { - // TODO: pass in cli overrides once cloud tasks properly support them. - let config = match Config::load_with_cli_overrides(Vec::new()).await { + config_overrides: &CliConfigOverrides, +) -> anyhow::Result<(Option>, HttpClientFactory, Option)> { + let overrides = config_overrides + .parse_overrides() + .map_err(anyhow::Error::msg)?; + let config = match Config::load_with_cli_overrides(overrides).await { Ok(config) => config, Err(error) => { append_error_log(format!( "failed to load auth config; using transport-default proxy handling: {error}" )); let http_client_factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault); - return (None, http_client_factory); + return Ok((None, http_client_factory, None)); } }; let http_client_factory = config.http_client_factory(); @@ -102,15 +106,14 @@ pub async fn load_auth_manager( Ok(auth_manager) => auth_manager, Err(error) => { append_error_log(format!("failed to load auth: {error}")); - return (None, http_client_factory); + return Ok((None, http_client_factory, None)); } }; - (Some(auth_manager), http_client_factory) + Ok((Some(auth_manager), http_client_factory, Some(config))) } -/// Build headers for ChatGPT-backed requests: `User-Agent`, optional `Authorization`, -/// and optional `ChatGPT-Account-Id`. -pub async fn build_chatgpt_headers() -> anyhow::Result { +/// Non-authentication headers shared by environment requests. The client resolves auth per GET. +pub fn build_chatgpt_headers() -> HeaderMap { use http::header::HeaderValue; use http::header::USER_AGENT; @@ -121,13 +124,7 @@ pub async fn build_chatgpt_headers() -> anyhow::Result { USER_AGENT, HeaderValue::from_str(&ua).unwrap_or(HeaderValue::from_static("codex-cli")), ); - if let Some(am) = load_auth_manager(/*chatgpt_base_url*/ None).await.0 - && let Some(auth) = am.auth().await - && auth.uses_codex_backend() - { - headers.extend(codex_model_provider::auth_provider_from_auth(&auth).to_auth_headers()?); - } - Ok(headers) + headers } /// Construct a browser-friendly task URL for the given backend base URL. diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 3cc3efbd39..d00faa1663 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1765,10 +1765,28 @@ impl Session { AnalyticsEventsClient::disabled() } else { analytics_events_client.unwrap_or_else(|| { - AnalyticsEventsClient::new( + let api_key_headers = if codex_model_provider::provider_uses_first_party_auth_path( + &config.model_provider, + ) { + match config.model_provider.build_header_map() { + Ok(headers) => headers, + Err(error) => { + tracing::warn!(%error, "failed to resolve analytics resource selectors"); + return AnalyticsEventsClient::disabled(); + } + } + } else { + Default::default() + }; + AnalyticsEventsClient::new_with_runtime( Arc::clone(&auth_manager), config.chatgpt_backend_url(), config.analytics_enabled, + config + .features + .enabled(Feature::RuntimeAgentInstances) + .then(|| (Arc::clone(&runtime_auth), config.runtime_auth_mode)), + api_key_headers, ) }) }; diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 12562fa3a0..5002100b71 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -886,6 +886,11 @@ impl ThreadManager { self.state.auth_manager.clone() } + /// Return the harness owner initialized at startup, if runtime auth was enabled. + pub fn runtime_auth_manager(&self) -> Option> { + self.state.runtime_auth.get().cloned() + } + /// Share the harness runtime with an API client using this effective ChatGPT configuration. pub fn runtime_auth_for_config( &self, diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index dbdb4ef011..a0d1e53773 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -227,15 +227,6 @@ pub trait ModelProvider: fmt::Debug + Send + Sync { }) } - /// Returns the endpoint and authentication for one API client. - /// Providers with credential-dependent routing should resolve both from the same auth read. - fn api_provider_and_auth( - &self, - ) -> ModelProviderFuture<'_, codex_protocol::error::Result<(Provider, SharedAuthProvider)>> - { - Box::pin(async move { Ok((self.api_provider().await?, self.api_auth().await?)) }) - } - /// Resolves routing for Responses HTTP, compaction, and WebSocket handshakes. fn responses_api_provider<'a>( &'a self, @@ -619,22 +610,6 @@ impl ModelProvider for ConfiguredModelProvider { }) } - fn api_provider_and_auth( - &self, - ) -> ModelProviderFuture<'_, codex_protocol::error::Result<(Provider, SharedAuthProvider)>> - { - Box::pin(async move { - let auth = self.auth().await; - let provider = self.api_provider_for_auth(auth.as_ref())?; - let primary = resolve_provider_auth(auth.as_ref(), &self.info)?; - let auth = self - .compose_auth(ResolvedProviderAuth::new(primary)) - .await? - .auth; - Ok((provider, auth)) - }) - } - fn compose_auth( &self, resolved: ResolvedProviderAuth, diff --git a/codex-rs/model-provider/src/provider_tests.rs b/codex-rs/model-provider/src/provider_tests.rs index aa74446aaf..97c8680bfa 100644 --- a/codex-rs/model-provider/src/provider_tests.rs +++ b/codex-rs/model-provider/src/provider_tests.rs @@ -267,7 +267,12 @@ async fn configured_provider_resolves_default_and_explicit_endpoints() { Some("https://example.test/v1".to_string()) ); assert_eq!( - provider.api_provider_and_auth().await.unwrap().0.base_url, + provider + .request_provider_and_auth(/*scope*/ None) + .await + .unwrap() + .0 + .base_url, "https://example.test/v1" ); @@ -296,9 +301,12 @@ async fn configured_provider_resolves_default_and_explicit_endpoints() { ) .unwrap(); auth.reload().await; - let (api_provider, api_auth) = provider.api_provider_and_auth().await.unwrap(); + let (api_provider, resolved) = provider + .request_provider_and_auth(/*scope*/ None) + .await + .unwrap(); assert_eq!(api_provider.base_url, expected_base_url); - let headers = api_auth.to_auth_headers().unwrap(); + let headers = resolved.auth.to_auth_headers().unwrap(); assert_eq!( headers[http::header::AUTHORIZATION].to_str().unwrap(), format!("Bearer {api_key}") @@ -318,9 +326,12 @@ async fn configured_provider_resolves_default_and_explicit_endpoints() { ("commercial-key", "https://api.openai.com/v1", None), ] { let reads_before = rotating_auth.0.load(Ordering::SeqCst); - let (api_provider, api_auth) = provider.api_provider_and_auth().await.unwrap(); + let (api_provider, resolved) = provider + .request_provider_and_auth(/*scope*/ None) + .await + .unwrap(); assert_eq!(rotating_auth.0.load(Ordering::SeqCst), reads_before + 1); - let headers = api_auth.to_auth_headers().unwrap(); + let headers = resolved.auth.to_auth_headers().unwrap(); let expected_authorization = format!("Bearer {api_key}"); assert_eq!( ( diff --git a/codex-rs/model-provider/src/request_preparation_tests.rs b/codex-rs/model-provider/src/request_preparation_tests.rs index 6fe63f8185..a02249242a 100644 --- a/codex-rs/model-provider/src/request_preparation_tests.rs +++ b/codex-rs/model-provider/src/request_preparation_tests.rs @@ -33,7 +33,7 @@ async fn explicit_provider_token_survives_a_workspace_switch_during_setup() { |_| { attempts += 1; async { - let (_, auth) = provider.api_provider_and_auth().await?; + let (_, resolved) = provider.request_provider_and_auth(/*scope*/ None).await?; codex_login::auth::login_with_chatgpt_auth_tokens( home.path(), &token, @@ -41,7 +41,7 @@ async fn explicit_provider_token_survives_a_workspace_switch_during_setup() { /*chatgpt_plan_type*/ None, )?; manager.reload().await; - Ok(auth) + Ok(resolved.auth) } }, ) diff --git a/codex-rs/model-provider/src/shared_state_tests.rs b/codex-rs/model-provider/src/shared_state_tests.rs index 3979a9988f..d687bac987 100644 --- a/codex-rs/model-provider/src/shared_state_tests.rs +++ b/codex-rs/model-provider/src/shared_state_tests.rs @@ -107,7 +107,11 @@ async fn gateway_credentials_accompany_primary_auth_in_models_and_responses() { models .raw_model_catalog(RefreshStrategy::Online, factory.clone()) .await; - let (api, auth) = provider.api_provider_and_auth().await.unwrap(); + let (api, resolved) = provider + .request_provider_and_auth(/*scope*/ None) + .await + .unwrap(); + let auth = resolved.auth; let scoped = provider .api_auth_for_scope(ProviderAuthScope { agent_identity_policy: AgentIdentityAuthPolicy::JwtOnly,