diff --git a/codex-rs/app-server/src/request_processors/account_processor.rs b/codex-rs/app-server/src/request_processors/account_processor.rs index 178edfc5ae..ba8672181c 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -1553,7 +1553,10 @@ pub(crate) async fn logout_configured_account( ) -> Result { // Another process may have changed the persisted workspace. Reload both // account authority and its policy before selecting a grant to remove. - auth_manager.reload().await; + auth_manager + .reload_with_cached_auth() + .await + .map_err(|err| internal_error(format!("logout failed: {err}")))?; config_manager.update_managed_config_loader(Arc::clone(auth_manager), fallback_config); let config = config_manager .load_latest_config(/*fallback_cwd*/ None) diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index 339438a5fd..b5c39b109a 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -386,6 +386,69 @@ async fn logout_account_removes_auth_and_notifies(enterprise_cleanup_failure: bo Ok(()) } +#[tokio::test] +async fn logout_account_requires_a_successful_auth_reload() -> Result<()> { + let oauth_server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/revoke")) + .and(wiremock::matchers::body_partial_json( + json!({"token": "refresh-b"}), + )) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&oauth_server) + .await; + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?; + write_chatgpt_auth( + codex_home.path(), + ChatGptAuthFixture::new("access-a").refresh_token("refresh-a"), + AuthCredentialsStoreMode::File, + )?; + let refresh_url = format!("{}/oauth/token", oauth_server.uri()); + let mut mcp = TestAppServer::builder() + .with_codex_home(codex_home.path()) + .without_auto_env() + .with_env_overrides(&[ + ("OPENAI_API_KEY", None), + (REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR, Some(&refresh_url)), + ]) + .build_initialized_with_timeout(DEFAULT_READ_TIMEOUT) + .await?; + + let auth_file = codex_home.path().join("auth.json"); + let unreadable = "{"; + std::fs::write(&auth_file, unreadable)?; + let id = mcp.send_logout_account_request().await?; + let error = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(id)), + ) + .await??; + assert_eq!(error.error.code, -32603); + assert_eq!(std::fs::read_to_string(&auth_file)?, unreadable); + assert!( + oauth_server + .received_requests() + .await + .unwrap() + .iter() + .all(|request| request.url.path() != "/oauth/revoke"), + ); + + // Once storage is readable, logout must revoke the newly loaded account. + write_chatgpt_auth( + codex_home.path(), + ChatGptAuthFixture::new("access-b").refresh_token("refresh-b"), + AuthCredentialsStoreMode::File, + )?; + let id = mcp.send_logout_account_request().await?; + let _: LogoutAccountResponse = timeout(DEFAULT_READ_TIMEOUT, mcp.read_response(id)).await??; + assert!(!auth_file.exists()); + oauth_server.verify().await; + Ok(()) +} + #[tokio::test] async fn logout_account_succeeds_when_config_reload_fails() -> Result<()> { let codex_home = TempDir::new()?; diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 47c780a3a7..ae7dfc7a33 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -25,6 +25,7 @@ //! WebSocket prewarm is treated as the first websocket connection attempt for a turn. If it //! fails, normal stream retry/fallback logic handles recovery on the same turn. +use codex_login::AuthOwner; use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex as StdMutex; @@ -368,7 +369,7 @@ struct WebsocketSession { parent_routing_affinity: Option, runtime_id: Option, /// Owner of the cached state, including before a connection is opened. - auth_owner_generation: Option, + auth_owner: Option, connection_key: Option, last_request: Option, last_response_rx: Option>, @@ -524,7 +525,7 @@ impl WebsocketSession { .continuation_reset_reason .or_else(|| self.last_request.as_ref().and(reason)); *self = Self { - auth_owner_generation: self.auth_owner_generation, + auth_owner: self.auth_owner.clone(), continuation_reset_reason, last_connected_auth_revision: self.last_connected_auth_revision, ..Default::default() @@ -754,12 +755,12 @@ impl ModelClient { /// This constructor does not perform network I/O itself; the session opens a websocket lazily /// when the first stream request is issued. pub fn new_session(&self) -> ModelClientSession { - let auth_owner_generation = self.auth_owner_generation(); + let auth_owner = self.auth_owner(); let mut websocket_session = self.take_cached_websocket_session(); - if websocket_session.auth_owner_generation != auth_owner_generation { + if websocket_session.auth_owner != auth_owner { // Drop the old owner's cache before this turn can establish fresh routing state. websocket_session.reset(Some("other")); - websocket_session.auth_owner_generation = auth_owner_generation; + websocket_session.auth_owner = auth_owner; } ModelClientSession { client: self.clone(), @@ -772,13 +773,8 @@ impl ModelClient { self.state.provider.auth_manager() } - pub(crate) fn auth_owner_generation(&self) -> Option { - self.auth_manager().map(|manager| { - manager - .auth_change_state_receiver() - .borrow() - .owner_generation - }) + pub(crate) fn auth_owner(&self) -> Option { + codex_model_provider::request_auth_owner(self.state.provider.as_ref()) } fn take_cached_websocket_session(&self) -> WebsocketSession { @@ -1545,9 +1541,11 @@ impl<'a> ResponsesHttpExecutor<'a> { } impl ModelClientSession { - pub(crate) fn turn_state_with_auth_owner(&self) -> (Option, Option) { - let owner = self.client.auth_owner_generation(); - let turn_state = (self.websocket_session.auth_owner_generation == owner) + pub(crate) fn turn_state_with_auth_owner(&self) -> (Option, Option) { + let owner = self.websocket_session.auth_owner.clone(); + let turn_state = owner + .as_ref() + .is_none_or(AuthOwner::is_current) .then(|| self.turn_state.get().cloned()) .flatten(); (turn_state, owner) @@ -1671,7 +1669,7 @@ impl ModelClientSession { auth_revision: client_setup.auth_revision, api_auth: client_setup.api_auth, http_client_factory: client_setup.http_client_factory, - auth_owner_generation: client_setup.auth_owner_generation, + auth_owner: client_setup.auth_owner, responses_metadata: &websocket_metadata, auth_context, request_route_telemetry: RequestRouteTelemetry::for_endpoint("/responses"), @@ -1731,7 +1729,7 @@ impl ModelClientSession { auth_revision, api_auth, http_client_factory, - auth_owner_generation, + auth_owner, responses_metadata, auth_context, request_route_telemetry, @@ -1759,8 +1757,8 @@ impl ModelClientSession { }; let needs_new = self.websocket_session.connection.is_none() || reset_reason.is_some(); // Resolving an external auth provider can change ownership during client setup. - let owner_changed = self.websocket_session.auth_owner_generation != auth_owner_generation - || self.client.auth_owner_generation() != auth_owner_generation; + let owner_changed = self.websocket_session.auth_owner != auth_owner + || auth_owner.as_ref().is_some_and(|owner| !owner.is_current()); if owner_changed { self.turn_state = Arc::new(OnceLock::new()); } @@ -1795,7 +1793,7 @@ impl ModelClientSession { self.websocket_session.parent_routing_affinity = parent_routing_affinity; self.websocket_session.responses_headers = responses_headers.clone(); self.websocket_session.runtime_id = runtime_id; - self.websocket_session.auth_owner_generation = auth_owner_generation; + self.websocket_session.auth_owner = auth_owner; self.websocket_session.connection_key = Some(connection_key); self.websocket_session .set_connection_reused(/*connection_reused*/ false); @@ -1845,7 +1843,7 @@ impl ResponsesHttpExecutor<'_> { pub(crate) async fn stream_responses_http( &self, turn_state: &Arc>, - expected_auth_owner: Option, + expected_auth_owner: Option, event_sender: Option<&Sender>, prompt: &Prompt, model_info: &ModelInfo, @@ -1867,9 +1865,8 @@ impl ResponsesHttpExecutor<'_> { .client .current_client_setup(ClientRouting::Workspace) .await?; - if expected_auth_owner.is_some_and(|owner| { - Some(owner) != client_setup.auth_owner_generation - || Some(owner) != self.client.auth_owner_generation() + if expected_auth_owner.as_ref().is_some_and(|owner| { + Some(owner) != client_setup.auth_owner.as_ref() || !owner.is_current() }) { return Err(CodexErr::TurnAborted); } @@ -2141,7 +2138,7 @@ impl ModelClientSession { auth_revision: client_setup.auth_revision, api_auth: client_setup.api_auth, http_client_factory: client_setup.http_client_factory, - auth_owner_generation: client_setup.auth_owner_generation, + auth_owner: client_setup.auth_owner, runtime_id: client_setup.runtime_id, responses_metadata: &websocket_metadata, auth_context: request_auth_context, @@ -2889,7 +2886,7 @@ struct WebsocketConnectParams<'a> { auth_revision: Option, api_auth: SharedAuthProvider, http_client_factory: HttpClientFactory, - auth_owner_generation: Option, + auth_owner: Option, responses_metadata: &'a CodexResponsesMetadata, auth_context: AuthRequestTelemetryContext, request_route_telemetry: RequestRouteTelemetry, diff --git a/codex-rs/core/src/client_runtime_auth_tests.rs b/codex-rs/core/src/client_runtime_auth_tests.rs index e1f5b32da5..982f29ea45 100644 --- a/codex-rs/core/src/client_runtime_auth_tests.rs +++ b/codex-rs/core/src/client_runtime_auth_tests.rs @@ -218,3 +218,82 @@ async fn realtime_keeps_provider_auth_until_runtime_integration(mode: RuntimeAut assert!(!requests[0].headers.contains_key("x-openai-agent-assertion")); assert!(registration.received_requests().await.unwrap().is_empty()); } + +#[test_case::test_case(RuntimeAuthMode::DualHeader; "dual_header")] +#[test_case::test_case(RuntimeAuthMode::AssertionOnly; "assertion_only")] +#[tokio::test] +async fn request_owner_tracks_the_selected_explicit_key(mode: RuntimeAuthMode) { + let server = MockServer::start().await; + Mock::given(path("/register")) + .respond_with(|request: &wiremock::Request| ResponseTemplate::new(200).set_body_json(json!({ + "runtime_agent_instance_id": request.headers["idempotency-key"].to_str().unwrap(), + "agent_host_id": "host-a", + "actor": {"kind": "platform", "user_id": "user-a", "organization_id": "org-a", "project_id": null}, + }))) + .expect(2) + .mount(&server).await; + let credentials = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("saved-key")); + let context = codex_model_provider::ProviderAuthContext { + runtime_auth: Some(( + Arc::new(RuntimeAuthManager::new( + credentials.clone(), + AgentHostAuthorization::InstallationId { + installation_id: uuid::Uuid::nil(), + }, + RuntimeRegistrationMetadata { + harness: None, + harness_version: None, + role: RuntimeRole::Orchestrator, + }, + format!("{}/register", server.uri()).parse().unwrap(), + )), + mode, + )), + credentials, + }; + let mut owners = Vec::new(); + for key in ["key-a", "key-b", "key-a"] { + let mut info = ModelProviderInfo::create_openai_provider(/*base_url*/ None); + info.experimental_bearer_token = Some(key.into()); + let mut client = test_model_client(SessionSource::Exec); + Arc::get_mut(&mut client.state).unwrap().provider = context.model_provider(info); + let owner = client.auth_owner(); + let setup = client + .current_client_setup(ClientRouting::ConfiguredProvider) + .await + .unwrap(); + assert!(owner.is_some()); + assert_eq!(setup.auth_owner, owner); + assert_eq!(setup.auth.unwrap().api_key(), Some(key)); + let mut headers = http::HeaderMap::new(); + setup.api_auth.add_auth_headers(&mut headers).unwrap(); + match mode { + RuntimeAuthMode::DualHeader => { + assert_eq!( + headers[http::header::AUTHORIZATION], + format!("Bearer {key}") + ); + assert!(headers.contains_key("x-openai-agent-assertion")); + } + RuntimeAuthMode::AssertionOnly => { + assert!( + headers[http::header::AUTHORIZATION] + .to_str() + .unwrap() + .starts_with("AgentAssertion ") + ); + assert!(!headers.contains_key("x-openai-agent-assertion")); + } + } + owners.push(owner); + } + assert_ne!( + owners[0], owners[1], + "different keys cannot share turn or connection ownership" + ); + assert_eq!( + owners[0], owners[2], + "reselecting a key preserves its owner" + ); + server.verify().await; +} diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 7e5ba85741..cfa1f589ed 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -400,7 +400,7 @@ async fn client_setup_accepts_command_credential_refresh() { refreshed_revision ), ); - assert_ne!(setup.auth_owner_generation, client.auth_owner_generation()); + assert_ne!(setup.auth_owner, client.auth_owner()); } } @@ -454,8 +454,11 @@ async fn client_setup_rebuilds_chatgpt_refresh_but_rejects_account_switches() { ); assert_eq!(setup.auth.unwrap().get_token().unwrap(), refreshed_token); assert_eq!( - (setup.auth_revision, setup.auth_owner_generation), - (Some(*manager.auth_change_receiver().borrow()), Some(0)) + (setup.auth_revision, setup.auth_owner), + ( + Some(*manager.auth_change_receiver().borrow()), + Some(manager.auth_owner()) + ) ); } else { assert_eq!( diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index dd9b289780..5656fb486b 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -1,3 +1,4 @@ +use codex_login::AuthOwner; use std::sync::Arc; use std::sync::OnceLock; @@ -102,7 +103,7 @@ enum CompactionRequestKind<'a> { enum CompactionRequestTransport<'a> { Foreground(&'a mut ModelClientSession), - Background(Option, Option), + Background(Option, Option), } /// The caller decides whether a declined prepared result needs a foreground request. @@ -297,7 +298,7 @@ async fn run_remote_compaction_request_v2( ) -> CodexResult { let turn_context = &step_context.turn; let client_session = match transport { - CompactionRequestTransport::Background(turn_state, auth_owner_generation) => { + CompactionRequestTransport::Background(turn_state, auth_owner) => { let settings = &step_context.settings; let cancellation = crate::session::async_compaction::shutdown_cancellation(sess).await; let effort = sess @@ -311,7 +312,7 @@ async fn run_remote_compaction_request_v2( biased; _ = cancellation.cancelled() => return Err(CodexErr::TurnAborted), stream = http_executor.stream_responses_http( - &turn_state, *auth_owner_generation, /*event_sender*/ None, prompt, &settings.model_info, &step_context.session_telemetry, effort, + &turn_state, auth_owner.clone(), /*event_sender*/ None, prompt, &settings.model_info, &step_context.session_telemetry, effort, settings.reasoning_summary, settings.service_tier.clone(), responses_metadata, &trace, ) => stream?, diff --git a/codex-rs/core/src/compact_remote_v2_attempt.rs b/codex-rs/core/src/compact_remote_v2_attempt.rs index 23c6dda1cc..30fad7f13a 100644 --- a/codex-rs/core/src/compact_remote_v2_attempt.rs +++ b/codex-rs/core/src/compact_remote_v2_attempt.rs @@ -1,6 +1,7 @@ //! Builds and sends foreground and speculative remote V2 compaction attempts. use crate::context::UserGoalUpdate; +use codex_login::AuthOwner; use codex_protocol::ResponseItemId; use codex_protocol::openai_models::CatalogModel; use std::collections::BTreeSet; @@ -165,10 +166,7 @@ pub(super) async fn run_remote_compact_v2_attempt( input.base_instructions, input.input, input.metadata, - CompactionRequestTransport::Background( - input.turn_state, - input.prefix.auth_owner_generation, - ), + CompactionRequestTransport::Background(input.turn_state, input.prefix.auth_owner), Some(input.window), ), }; @@ -283,7 +281,7 @@ pub(crate) async fn capture_background_input( mut snapshot: ContextManager, base_instructions: BaseInstructions, turn_state: Option, - auth_owner_generation: Option, + auth_owner: Option, window: (String, u64, codex_protocol::ContextWindowId), ) -> BackgroundCompactionInput { let base_instructions = if step_context.settings.model_info.use_responses_lite @@ -301,7 +299,7 @@ pub(crate) async fn capture_background_input( let prefix = CompactionPrefix { generation: snapshot.history_version(), len: snapshot.annotated_items().len(), - auth_owner_generation, + auth_owner, }; trim_function_call_history_to_fit_context_window( &mut snapshot, @@ -344,7 +342,7 @@ pub(crate) async fn run_background_compaction( sess.services.image_store.as_ref(), ) => {}, } - let prefix = prepared_input.prefix; + let prefix = prepared_input.prefix.clone(); let metadata = CompactionTurnMetadata::new( CompactionTrigger::Auto, CompactionReason::ContextLimit, diff --git a/codex-rs/core/src/compact_remote_v2_installation_tests.rs b/codex-rs/core/src/compact_remote_v2_installation_tests.rs index 48878fa955..6f8caa23b7 100644 --- a/codex-rs/core/src/compact_remote_v2_installation_tests.rs +++ b/codex-rs/core/src/compact_remote_v2_installation_tests.rs @@ -96,13 +96,19 @@ async fn prepared_prefix_rejection_respects_caller_intent( sess.replace_history(vec![original.clone()], /*reference_context_item*/ None) .await; let history = sess.clone_history().await; - let auth = sess.services.model_client.auth_owner_generation(); + let auth = sess.services.model_client.auth_owner(); let prefix = CompactionPrefix { generation: history.history_version(), len: history.annotated_items().len(), - auth_owner_generation: match change { + auth_owner: match change { HistoryChange::Replacement => auth, - HistoryChange::Account => Some(auth.unwrap_or_default().wrapping_add(1)), + HistoryChange::Account => Some( + codex_login::AuthManager::from_auth_for_testing_with_home( + CodexAuth::from_api_key("other-account"), + std::path::PathBuf::from("unused"), + ) + .auth_owner(), + ), }, }; let roster = Arc::new(FailingRosterRead::default()); @@ -119,7 +125,7 @@ async fn prepared_prefix_rejection_respects_caller_intent( cutoff_world_state.host_roster_reader = Some(roster.clone()); } let candidate = BackgroundCompactionCandidate { - prefix, + prefix: prefix.clone(), model_info: Arc::clone(&step.settings.model_info), cutoff_world_state: Arc::new(cutoff_world_state), attempt: RemoteCompactV2Attempt { diff --git a/codex-rs/core/src/session/async_compaction.rs b/codex-rs/core/src/session/async_compaction.rs index cc263a4560..98a28dba87 100644 --- a/codex-rs/core/src/session/async_compaction.rs +++ b/codex-rs/core/src/session/async_compaction.rs @@ -2,6 +2,7 @@ //! Discarded requests still charge completed usage; session shutdown cancels remaining network work. //! Snapshot capture and launch registration hold session state; prompt preparation does not. +use codex_login::AuthOwner; use std::sync::Arc; use crate::compact_remote_v2::BackgroundCompactionCandidate; @@ -26,7 +27,7 @@ struct CompactionAttempt { history_version: u64, model_slug: String, compaction_hash: Option, - auth_owner_generation: Option, + auth_owner: Option, } #[derive(Default)] @@ -57,7 +58,7 @@ pub(super) async fn maybe_start( phase: CompactionPhase, turn_cancellation: &CancellationToken, token_status: &ContextWindowTokenStatus, - (turn_state, auth_owner_generation): (Option, Option), + (turn_state, auth_owner): (Option, Option), ) { let turn_context = &step_context.turn; // Let the foreground establish the backend's model pin before starting a parallel request. @@ -89,14 +90,14 @@ pub(super) async fn maybe_start( history_version: state.history.history_version(), model_slug: model_info.slug.clone(), compaction_hash: model_info.comp_hash.clone(), - auth_owner_generation, + auth_owner: auth_owner.clone(), }; // Stop blocks new launches; already-running compaction remains session-owned. if turn_cancellation.is_cancelled() || state.shutting_down || state.async_compaction.last_attempt.as_ref() == Some(&attempt) || state.history.annotated_items().is_empty() - || sess.services.model_client.auth_owner_generation() != auth_owner_generation + || sess.services.model_client.auth_owner() != auth_owner { return; } @@ -108,7 +109,7 @@ pub(super) async fn maybe_start( state.history.clone(), base_instructions, turn_state, - auth_owner_generation, + auth_owner, sess.current_window_from_state(&state), ); let tasks = state.async_compaction.tasks.clone(); @@ -187,7 +188,7 @@ pub(crate) async fn wait_for_pending( if attempt.history_version != state.history.history_version() || attempt.model_slug != model.slug || attempt.compaction_hash != model.comp_hash - || attempt.auth_owner_generation != sess.services.model_client.auth_owner_generation() + || attempt.auth_owner != sess.services.model_client.auth_owner() { return None; } diff --git a/codex-rs/core/src/session/async_compaction_tests.rs b/codex-rs/core/src/session/async_compaction_tests.rs index f85b360d23..90bae90d13 100644 --- a/codex-rs/core/src/session/async_compaction_tests.rs +++ b/codex-rs/core/src/session/async_compaction_tests.rs @@ -120,7 +120,7 @@ async fn cancelled_turn_skips_launch_but_preserves_existing_work(trigger: Compac let world_state = Arc::new(WorldState::default()); let cancelled = CancellationToken::new(); cancelled.cancel(); - let auth = sess.services.model_client.auth_owner_generation(); + let auth = sess.services.model_client.auth_owner(); maybe_start( &sess, &step, @@ -128,7 +128,7 @@ async fn cancelled_turn_skips_launch_but_preserves_existing_work(trigger: Compac phase, &cancelled, &status, - (Some("routing".to_string()), auth), + (Some("routing".to_string()), auth.clone()), ) .await; assert!( @@ -149,7 +149,7 @@ async fn cancelled_turn_skips_launch_but_preserves_existing_work(trigger: Compac phase, &turn_cancellation, &status, - (Some("routing".to_string()), auth), + (Some("routing".to_string()), auth.clone()), ) .await; turn_cancellation.cancel(); diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 13ba7bdd54..fb42967c46 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3,6 +3,7 @@ pub(crate) mod startup; use crate::PromptCacheKey; use crate::WithTurnExtensionData; use crate::runtime_auth::RuntimeAuthSlot; +use codex_login::AuthOwner; use std::collections::BTreeMap; use std::collections::HashMap; use std::collections::HashSet; @@ -477,11 +478,11 @@ pub(crate) enum ForkPersistence { } /// Normal appends keep this generation; compaction and rollback start a new one. -#[derive(Clone, Copy)] +#[derive(Clone)] pub(crate) struct CompactionPrefix { pub(crate) generation: u64, pub(crate) len: usize, - pub(crate) auth_owner_generation: Option, + pub(crate) auth_owner: Option, } pub(crate) enum CompactedHistoryInstallation<'a> { @@ -4778,8 +4779,7 @@ impl Session { // Validate immediately before slicing and replacing the captured prefix. if state.history.history_version() != prefix.generation || state.history.annotated_items().len() < prefix.len - || self.services.model_client.auth_owner_generation() - != prefix.auth_owner_generation + || self.services.model_client.auth_owner() != prefix.auth_owner { return false; } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs index fc10f402ba..4f53dfddc7 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs @@ -10,6 +10,7 @@ pub(super) use telemetry::SampledScore; use super::request::PreparedRequest; use super::score::PendingScore; +use codex_login::AuthOwner; use connection_pool::ConnectionPool; use std::sync::Arc; @@ -162,17 +163,13 @@ impl LunaSampler { self.connections.prewarm().await; } - fn auth_owner_generation(&self) -> Option { - self.config - .provider - .auth_manager() - .filter(|_| self.config.provider.info().auth.is_none()) - .map(|manager| { - manager - .auth_change_state_receiver() - .borrow() - .owner_generation - }) + fn auth_owner(&self) -> Option { + let provider = &self.config.provider; + if provider.info().auth.is_none() { + codex_model_provider::request_auth_owner(provider.as_ref()) + } else { + None + } } // Normalize new evidence while preserving the retained prefix and its item IDs. @@ -228,7 +225,7 @@ impl LunaSampler { &self, request: &LunaSamplingRequest, input: Vec, - auth_owner_generation: Option, + auth_owner: Option, ) -> execution::SamplingExecution { let api_request = ResponsesApiRequest { model: MODEL.to_owned(), @@ -253,7 +250,7 @@ impl LunaSampler { access_programs: None, }; execution::SamplingExecution { - auth_owner_generation, + auth_owner, cyber_access_program: codex_guardian_reviewer::guardian_cyber_access_program( request .cyber_access_program @@ -284,7 +281,7 @@ impl LunaSampler { input: review_request, sampling, } = sample; - let auth_owner_generation = self.auth_owner_generation(); + let auth_owner = self.auth_owner(); let mut completion = execution::RetainedCompletion { ready: Some(ready), output: Some(Vec::new()), @@ -301,7 +298,7 @@ impl LunaSampler { .into_iter() .map(|envelope| (envelope.item, envelope.metadata)) .unzip(); - let mut execution = self.execution(&sampling, input, auth_owner_generation); + let mut execution = self.execution(&sampling, input, auth_owner); execution.request.include = vec!["reasoning.encrypted_content".to_owned()]; let mut telemetry = telemetry::SamplingTelemetry::new( self.config.metrics.clone(), diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs index b89fb46c7c..2470a1c017 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool.rs @@ -28,6 +28,7 @@ use codex_http_client::ClientRouteClass; use codex_http_client::HttpClientFactory; use codex_http_client::NetworkPolicy; use codex_login::AuthChangeState; +use codex_login::AuthOwner; use codex_login::CodexAuth; use codex_login::default_client::ClientRedirectPolicy; use codex_login::default_client::add_originator_header; @@ -92,6 +93,7 @@ pub(super) struct PooledConnection { auth_changes: Option>, key: ResponsesConnectionKey, runtime_id: Option, + auth_owner: Option, _permit: OwnedSemaphorePermit, } @@ -99,6 +101,7 @@ struct ClientSetup { provider: Provider, http_client_factory: HttpClientFactory, runtime_id: Option, + auth_owner: Option, request_kind: RequestMode, auth: SharedAuthProvider, redirect_policy: ClientRedirectPolicy, @@ -114,6 +117,7 @@ enum Connection { pub(super) struct ConnectionLease { pub(super) thread_id: String, pub(super) request_kind: RequestMode, + pub(super) auth_owner: Option, connection: Connection, pool: Arc, _permit: OwnedSemaphorePermit, @@ -227,6 +231,7 @@ impl ConnectionPool { auth, http_client_factory, runtime_id, + auth_owner, request_kind, redirect_policy, key, @@ -243,6 +248,7 @@ impl ConnectionPool { Some(connection) if connection.key == key && connection.runtime_id == runtime_id + && connection.auth_owner == auth_owner && connection .auth_changes .as_ref() @@ -356,6 +362,7 @@ impl ConnectionPool { } } Ok(ConnectionLease { + auth_owner, thread_id, request_kind, connection, @@ -377,11 +384,7 @@ impl ConnectionPool { ) .await .map_err(LunaSamplerError::Provider)?; - let mut auth_changes = self - .config - .provider - .auth_manager() - .map(|manager| manager.auth_change_state_receiver()); + let mut auth_changes = setup.auth_owner.as_ref().map(AuthOwner::changes); if auth_changes .as_mut() .map(|auth| auth.borrow_and_update().generation) @@ -399,6 +402,7 @@ impl ConnectionPool { redirect_policy: setup.redirect_policy, http_client_factory: setup.http_client_factory, runtime_id: setup.runtime_id, + auth_owner: setup.auth_owner, request_kind, key, auth_changes, @@ -463,6 +467,7 @@ impl ConnectionPool { auth, http_client_factory, runtime_id, + auth_owner, request_kind, redirect_policy: _, key, @@ -545,6 +550,7 @@ impl ConnectionPool { auth_changes, key, runtime_id, + auth_owner, _permit: permit, }) } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_cancellation_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_cancellation_tests.rs index d2b39209de..a69de4927f 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_cancellation_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_cancellation_tests.rs @@ -153,7 +153,7 @@ async fn cancellation_preserves_unused_sockets_and_backpressured_completion() -> .execution( &sample_request("stale"), Vec::new(), - /*auth_owner_generation*/ None, + /*auth_owner*/ None, ) .request; // Cancel before this current-thread runtime can dispatch, then reuse its only socket. diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs index 4c93e1e050..d5e87ba2d2 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs @@ -20,6 +20,7 @@ use codex_api::ResponsesApiRequest; use codex_api::TransportError; use codex_core::cyber_access_program; use codex_extension_api::ExtensionMetrics; +use codex_login::AuthOwner; use codex_login::UnauthorizedRecovery; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; @@ -37,7 +38,7 @@ const RESPONSES_LITE_METADATA_KEY: &str = const TURN_METADATA_KEY: &str = "x-codex-turn-metadata"; pub(super) struct SamplingExecution { - pub(super) auth_owner_generation: Option, + pub(super) auth_owner: Option, pub(super) cyber_access_program: Option, pub(super) config: Arc, pub(super) connections: Arc, @@ -184,34 +185,24 @@ impl SamplingExecution { completion: &mut RetainedCompletion, telemetry: &mut SamplingTelemetry, ) -> Result { - let auth_changes = self - .config - .provider - .auth_manager() - .filter(|_| self.config.provider.info().auth.is_none()) - .map(|manager| manager.auth_change_state_receiver()); - let owner_generation = self.auth_owner_generation; + let owner = self.auth_owner.clone(); let account_changed_error = || { LunaSamplerError::Provider( std::io::Error::other("account changed during Luna classification").into(), ) }; let ensure_account_owner = || { - if auth_changes - .as_ref() - .map(|changes| changes.borrow().owner_generation) - != owner_generation - { + if owner.as_ref().is_some_and(|owner| !owner.is_current()) { return Err(account_changed_error()); } Ok(()) }; - let mut owner_changes = auth_changes.clone(); - let owner_changed = async move { - let Some(changes) = owner_changes.as_mut() else { + let owner_changed = async { + let Some(owner) = owner.as_ref() else { return std::future::pending::<()>().await; }; - while Some(changes.borrow_and_update().owner_generation) == owner_generation { + let mut changes = owner.changes(); + while owner.is_current() { if changes.changed().await.is_err() { break; } @@ -246,6 +237,12 @@ impl SamplingExecution { return Err(error); } }; + if owner + .as_ref() + .is_some_and(|owner| lease.auth_owner.as_ref() != Some(owner)) + { + return Err(account_changed_error()); + } // Omitted programs need no auth lookup, which can refresh credentials. self.request.access_programs = match self.cyber_access_program { None => None, diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 3b286c6b48..9a3df938f4 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -1461,6 +1461,51 @@ async fn refresh_failure_is_scoped_to_the_matching_auth_snapshot() { assert_eq!(manager.refresh_failure_for_auth(&updated_auth), None); } +#[tokio::test] +async fn provider_credentials_follow_the_primary_network_policy() { + let script = ProviderAuthScript::new(&["provider-key", "rotated-provider-key"]).unwrap(); + let controller = codex_http_client::NetworkPolicyController::default(); + let policy = controller.policy(); + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) + .with_network_policy(policy.clone()); + let mut primary = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("login-key")); + Arc::get_mut(&mut primary).unwrap().auth_route_config = + AuthRouteConfig::from_http_client_factory(factory); + assert!(controller.publish( + policy.revision(), + codex_http_client::DestinationPolicy::Restricted { + allowed_hosts: ["allowed.example".to_string()].into_iter().collect(), + }, + )); + + let provider = + AuthManager::external_bearer_with_runtime(script.auth_config(), primary.runtime_config()); + let transport = provider.http_client_factory(); + let allowed = "https://allowed.example/".parse().unwrap(); + let blocked = "https://blocked.example/".parse().unwrap(); + let permit = transport.network_policy().acquire(&allowed).unwrap(); + assert_eq!( + provider.auth().await.unwrap().api_key(), + Some("provider-key"), + ); + permit.check().unwrap(); + let first_owner = provider.auth_owner(); + + provider.refresh_token().await.unwrap(); + assert_eq!( + provider.auth_cached().unwrap().api_key(), + Some("rotated-provider-key"), + ); + assert!(!first_owner.is_current()); + permit.check().unwrap(); + assert!(transport.network_policy().acquire(&allowed).is_ok()); + assert!(transport.network_policy().acquire(&blocked).is_err()); + + primary.set_cached_auth(Some(CodexAuth::from_api_key("another-login-key"))); + assert!(permit.check().is_err()); + assert!(transport.network_policy().acquire(&allowed).is_err()); +} + #[tokio::test] async fn external_bearer_only_auth_manager_uses_cached_provider_token() { let script = ProviderAuthScript::new(&["provider-token", "next-token"]).unwrap(); @@ -1825,6 +1870,32 @@ async fn external_auth_does_not_reuse_disallowed_credentials_after_a_transient_f assert_eq!(manager.auth().await, Some(recovered)); } +#[tokio::test] +#[serial_test::serial(codex_auth_env)] +async fn failed_storage_reload_preserves_cached_identity() -> anyhow::Result<()> { + let _access_token = remove_access_token_env_var(); + let home = tempdir()?; + let auth = CodexAuth::from_api_key("stored-key"); + let manager = + AuthManager::from_auth_for_testing_with_home(auth.clone(), home.path().to_path_buf()); + let initial = manager + .auth_change_state_for(&auth) + .expect("the cached credential has a source"); + std::fs::create_dir(get_auth_file(home.path()))?; + + assert!( + manager.reload_with_cached_auth().await.is_err(), + "the fixture must fail the storage read" + ); + assert!( + !manager.reload().await, + "failed storage reads must not publish a logout" + ); + assert_eq!(manager.auth_cached(), Some(auth.clone())); + assert_eq!(manager.auth_change_state_for(&auth), Some(initial)); + Ok(()) +} + #[tokio::test] async fn replacing_external_auth_clears_permanent_failure() { let manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("seed")); diff --git a/codex-rs/login/src/auth/change_state.rs b/codex-rs/login/src/auth/change_state.rs index 6fbf5cc717..e2aae121d5 100644 --- a/codex-rs/login/src/auth/change_state.rs +++ b/codex-rs/login/src/auth/change_state.rs @@ -1,6 +1,49 @@ //! Tracks credential changes separately from ownership changes, before notifications coalesce. +use super::AuthManager; use super::CodexAuth; +use std::sync::Arc; + +/// An account owner within a particular credential source, retained by in-flight work. +#[derive(Clone)] +pub struct AuthOwner { + manager: Arc, + generation: u64, +} + +impl AuthOwner { + pub(super) fn new(manager: Arc, generation: u64) -> Self { + Self { + manager, + generation, + } + } + + /// Watch credential changes from the source that owns this work. + pub fn changes(&self) -> tokio::sync::watch::Receiver { + self.manager.auth_change_state_receiver() + } + + pub fn is_current(&self) -> bool { + self.changes().borrow().owner_generation == self.generation + } +} + +impl PartialEq for AuthOwner { + fn eq(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.manager, &other.manager) && self.generation == other.generation + } +} + +impl Eq for AuthOwner {} + +impl std::fmt::Debug for AuthOwner { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AuthOwner") + .field("generation", &self.generation) + .finish_non_exhaustive() + } +} /// Opaque revisions local to one auth manager. Consumers must reset on reconnect. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] @@ -8,7 +51,6 @@ pub struct AuthChangeState { /// Advances on cached credential changes and explicit invalidation of pending reads. pub generation: u64, /// Advances on login, logout, or changes of user, workspace, service account, or auth mode. - /// A verified legacy-to-OAuth credential migration preserves the same owner. /// Credential changes with incomplete owner identity also advance this revision. /// Losing login metadata alone preserves the account owner. pub owner_generation: u64, diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 79388f6485..7a3a2e68aa 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -55,6 +55,7 @@ use super::agent_identity::require_agent_identity_authapi_base_url; use super::agent_identity::verified_record_from_jwt; use super::api_key::is_fedramp_api_key; use super::change_state::AuthChangeState; +use super::change_state::AuthOwner; use super::change_state::same_login; use super::change_state::same_owner; use super::external_bearer::BearerTokenRefresher; @@ -313,6 +314,18 @@ pub trait ExternalAuth: Send + Sync { } } +struct FixedApiKeyAuth(CodexAuth); + +impl ExternalAuth for FixedApiKeyAuth { + fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { Ok(self.0.clone()) }) + } + + fn refresh(&self, _context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { + self.resolve() + } +} + pub type ExternalAuthFuture<'a, T> = Pin> + Send + 'a>>; fn permanent_external_auth_error(message: impl Into) -> RefreshTokenError { @@ -2240,6 +2253,8 @@ pub struct AuthManager { external_auth: RwLock>>, workload_identity_selected: bool, auth_route_config: AuthRouteConfig, + // Provider credentials follow the primary account's policy without owning its lifecycle. + owns_network_policy: bool, allowed_auth_modes: RwLock>>, } @@ -2404,6 +2419,7 @@ impl AuthManager { external_auth: RwLock::new(None), workload_identity_selected: false, auth_route_config, + owns_network_policy: true, allowed_auth_modes: RwLock::new(allowed_auth_modes), } } @@ -2415,6 +2431,47 @@ impl AuthManager { /// Create an AuthManager with an optional CodexAuth, for testing only. pub(crate) fn from_optional_auth_for_testing(auth: Option) -> Arc { + Self::from_cached_auth( + auth, + /*external_auth*/ None, + AuthRuntimeConfig { + codex_home: PathBuf::from("non-existent"), + auth_route_config: crate::test_support::transport_default_auth_route_config(), + }, + /*owns_network_policy*/ true, + ) + } + + /// An explicitly selected API key with no ambient auth-storage or refresh fallback. + pub fn from_api_key(api_key: &str, runtime: AuthRuntimeConfig) -> Arc { + let auth = CodexAuth::from_api_key(api_key); + Self::from_cached_auth( + Some(auth.clone()), + Some(Arc::new(FixedApiKeyAuth(auth))), + runtime, + /*owns_network_policy*/ false, + ) + } + + /// Give a command-backed credential source the caller's storage and network context. + pub fn external_bearer_with_runtime( + config: ModelProviderAuthInfo, + runtime: AuthRuntimeConfig, + ) -> Arc { + Self::from_cached_auth( + /*auth*/ None, + Some(Arc::new(BearerTokenRefresher::new(config, HashMap::new()))), + runtime, + /*owns_network_policy*/ false, + ) + } + + fn from_cached_auth( + auth: Option, + external_auth: Option>, + runtime: AuthRuntimeConfig, + owns_network_policy: bool, + ) -> Arc { let cached = CachedAuth { auth, permanent_refresh_failure: None, @@ -2422,13 +2479,17 @@ impl AuthManager { let (auth_change_tx, _auth_change_rx) = watch::channel(0); Arc::new(Self { - codex_home: PathBuf::from("non-existent"), + codex_home: runtime.codex_home, inner: RwLock::new(cached), auth_change_tx, auth_change_state_tx: watch::channel(AuthChangeState::default()).0, workspace_routing_resolver: OnceLock::new(), enable_codex_api_key_env: false, - auth_credentials_store_mode: AuthCredentialsStoreMode::File, + auth_credentials_store_mode: if external_auth.is_some() { + AuthCredentialsStoreMode::Ephemeral + } else { + AuthCredentialsStoreMode::File + }, keyring_backend_kind: AuthKeyringBackendKind::default(), forced_login_method: None, forced_chatgpt_workspace_id: RwLock::new(None), @@ -2439,10 +2500,11 @@ impl AuthManager { unsettled_chatgpt_refresh: Mutex::new(None), agent_identity_lock: Semaphore::new(/*permits*/ 1), agent_identity_bootstrap_cooldown: Mutex::default(), - external_auth: RwLock::new(None), + external_auth: RwLock::new(external_auth), workload_identity_selected: false, allowed_auth_modes: RwLock::new(None), - auth_route_config: crate::test_support::transport_default_auth_route_config(), + auth_route_config: runtime.auth_route_config, + owns_network_policy, }) } @@ -2475,6 +2537,7 @@ impl AuthManager { workload_identity_selected: false, allowed_auth_modes: RwLock::new(None), auth_route_config: crate::test_support::transport_default_auth_route_config(), + owns_network_policy: true, }) } @@ -2515,6 +2578,7 @@ impl AuthManager { workload_identity_selected: false, allowed_auth_modes: RwLock::new(None), auth_route_config: crate::test_support::transport_default_auth_route_config(), + owns_network_policy: true, }) } @@ -2553,6 +2617,7 @@ impl AuthManager { external_auth: RwLock::new(Some(Arc::new(BearerTokenRefresher::new(config, env)))), workload_identity_selected: false, allowed_auth_modes: RwLock::new(None), + owns_network_policy: false, // External bearer auth refreshes by running the provider's command and never makes // auth-owned HTTP requests, so this route is intentionally inert. auth_route_config: AuthRouteConfig::from_http_client_factory(HttpClientFactory::new( @@ -2578,6 +2643,20 @@ impl AuthManager { Some(*self.auth_change_state_tx.borrow()) } + /// Identify this credential source's current account owner without refreshing auth. + pub fn auth_owner(self: &Arc) -> AuthOwner { + AuthOwner::new( + Arc::clone(self), + self.auth_change_state_tx.borrow().owner_generation, + ) + } + + /// Retain the owner of a selected credential, rejecting a changed source. + pub fn auth_owner_for(self: &Arc, expected: &CodexAuth) -> Option { + let state = self.auth_change_state_for(expected)?; + Some(AuthOwner::new(Arc::clone(self), state.owner_generation)) + } + /// Subscribes to cached auth changes that can affect request recovery. pub fn auth_change_receiver(&self) -> watch::Receiver { self.auth_change_tx.subscribe() @@ -2621,7 +2700,7 @@ impl AuthManager { #[instrument(level = "trace", skip_all)] pub async fn auth(&self) -> Option { if self.has_external_auth() { - self.reload().await; + self.reload_inner().await.ok()?; return self.auth_cached(); } @@ -2638,7 +2717,7 @@ impl AuthManager { /// Refreshes auth, then captures credentials and their account-bound factory together. /// The auth read lock prevents an identity change between the two snapshots. pub async fn auth_with_http_client_factory(&self) -> Option<(CodexAuth, HttpClientFactory)> { - self.auth().await; + self.auth().await?; let cached = self.inner.read().ok()?; Some((cached.auth.clone()?, self.http_client_factory())) } @@ -2839,15 +2918,19 @@ impl AuthManager { /// Reloads auth from the active source. Returns whether the auth value changed. pub async fn reload(&self) -> bool { + self.reload_inner().await.unwrap_or(false) + } + + async fn reload_inner(&self) -> Result { tracing::info!("Reloading auth"); loop { let generation = self.auth_change_state_tx.borrow().generation; - let new_auth = self.load_auth(/*cached_auth*/ None).await.ok().flatten(); + let new_auth = self.load_auth(/*cached_auth*/ None).await?; if let Some(publication) = self.publish_reloaded_auth(generation, new_auth) { - return publication.changed; + return Ok(publication.changed); } if self.inner.is_poisoned() { - return false; + return Ok(false); } } } @@ -2992,7 +3075,7 @@ impl AuthManager { async fn load_auth( &self, cached_auth: Option<&CodexAuth>, - ) -> std::io::Result> { + ) -> Result, RefreshTokenError> { if let Some(external_auth) = self.external_auth_provider() { let cached_auth = self.auth_cached(); if cached_auth @@ -3001,8 +3084,8 @@ impl AuthManager { { return Ok(cached_auth); } - let auth = match self.resolve_external_auth(external_auth.as_ref()).await { - Ok(auth) => Some(auth), + return match self.resolve_external_auth(external_auth.as_ref()).await { + Ok(auth) => Ok(Some(auth)), Err(err) => { tracing::error!("Failed to resolve external auth: {err}"); match err { @@ -3010,26 +3093,27 @@ impl AuthManager { if let Some(auth) = cached_auth.as_ref() { self.record_permanent_refresh_failure_if_unchanged(auth, &error); } - cached_auth + Ok(cached_auth) } // A temporary provider outage is not a logout. Preserve the identity // only while its cached credentials still satisfy the current policy. - RefreshTokenError::Transient(_) => cached_auth.filter(|auth| { + RefreshTokenError::Transient(_) => Ok(cached_auth.filter(|auth| { validate_auth_restrictions( Some(&self.allowed_login_methods()), self.effective_chatgpt_workspaces().as_deref(), auth, ) .is_ok() - }), - RefreshTokenError::Policy(_) => cached_auth, + })), + RefreshTokenError::Policy(_) => Ok(cached_auth), } } }; - return Ok(auth); } - self.try_load_managed_auth(cached_auth).await + self.try_load_managed_auth(cached_auth) + .await + .map_err(RefreshTokenError::Transient) } async fn try_load_managed_auth( @@ -3107,7 +3191,7 @@ impl AuthManager { let source_changed = auth_changed_for_refresh && !super::source::same_source(previous, new_auth.as_ref()); let login_changed = auth_changed_for_refresh && !same_login(previous, new_auth.as_ref()); - if owner_changed { + if owner_changed && self.owns_network_policy { self.auth_route_config .application_network_policy() .invalidate(); diff --git a/codex-rs/login/src/auth/mod.rs b/codex-rs/login/src/auth/mod.rs index e629be9e6a..7ea499ff3c 100644 --- a/codex-rs/login/src/auth/mod.rs +++ b/codex-rs/login/src/auth/mod.rs @@ -24,6 +24,7 @@ pub use bedrock_access_keys::login_with_bedrock_access_keys; pub use bedrock_api_key::BedrockApiKeyAuth; pub use bedrock_api_key::login_with_bedrock_api_key; pub use change_state::AuthChangeState; +pub use change_state::AuthOwner; pub use change_state::RequestAuthChange; pub use error::RefreshTokenFailedError; pub use error::RefreshTokenFailedReason; diff --git a/codex-rs/login/src/auth/source.rs b/codex-rs/login/src/auth/source.rs index 3e25bdb0e1..b50fdbfd70 100644 --- a/codex-rs/login/src/auth/source.rs +++ b/codex-rs/login/src/auth/source.rs @@ -22,6 +22,11 @@ pub(super) fn same_source(previous: Option<&CodexAuth>, current: Option<&CodexAu if previous_subject.is_some() || current_subject.is_some() { return previous_subject == current_subject; } + if let (CodexAuth::AgentIdentity(previous), CodexAuth::AgentIdentity(current)) = + (previous, current) + { + return previous.record() == current.record(); + } let (Ok(previous_token), Ok(current_token)) = (previous.get_token(), current.get_token()) else { return false; diff --git a/codex-rs/login/src/auth/source_tests.rs b/codex-rs/login/src/auth/source_tests.rs index 76ab4bcd2b..f2950bbfd5 100644 --- a/codex-rs/login/src/auth/source_tests.rs +++ b/codex-rs/login/src/auth/source_tests.rs @@ -123,6 +123,99 @@ fn api_key_replacement_and_logout_retire_the_previous_source() { ); } +#[tokio::test] +async fn explicit_api_key_never_reloads_an_ambient_credential() -> anyhow::Result<()> { + let home = tempfile::tempdir()?; + login_with_api_key( + home.path(), + "ambient-key", + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::Direct, + )?; + let runtime = AuthRuntimeConfig { + codex_home: home.path().to_path_buf(), + auth_route_config: crate::test_support::transport_default_auth_route_config(), + }; + let manager = AuthManager::from_api_key("selected-key", runtime.clone()); + for _ in 0..2 { + manager.reload().await; + assert_eq!( + manager.auth().await.expect("selected key").api_key(), + Some("selected-key") + ); + assert_eq!(manager.runtime_config(), runtime); + } + let stored = create_auth_storage( + home.path().to_path_buf(), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::Direct, + ) + .load()? + .expect("ambient login remains stored"); + assert_eq!(stored.openai_api_key.as_deref(), Some("ambient-key")); + Ok(()) +} + +struct RecoverableExternalAuth { + current: RwLock>, +} + +impl ExternalAuth for RecoverableExternalAuth { + fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { + Box::pin(async { + self.current + .read() + .unwrap() + .clone() + .ok_or_else(|| std::io::Error::other("credential source unavailable")) + }) + } + + fn refresh(&self, _context: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { + self.resolve() + } +} + +#[tokio::test] +async fn external_source_outage_keeps_cached_credentials_and_identity() { + let initial_auth = oauth("login-a", "token-1"); + let manager = AuthManager::from_auth_for_testing(initial_auth.clone()); + let source = Arc::new(RecoverableExternalAuth { + current: RwLock::new(Some(initial_auth.clone())), + }); + manager.set_external_auth(source.clone()).await.unwrap(); + let initial = manager.auth_change_state_for(&initial_auth).unwrap(); + *source.current.write().unwrap() = None; + assert_eq!(manager.auth().await, Some(initial_auth.clone())); + let (cached_auth, _) = manager.auth_with_http_client_factory().await.unwrap(); + assert_eq!(cached_auth, initial_auth); + assert!(!manager.reload().await); + assert_eq!(manager.auth_change_state_for(&initial_auth), Some(initial)); + + let refreshed = oauth("login-a", "token-2"); + *source.current.write().unwrap() = Some(refreshed.clone()); + assert_eq!(manager.auth().await, Some(refreshed)); + assert_eq!( + initial.source_generation, + manager + .auth_change_state_for(&initial_auth) + .unwrap() + .source_generation + ); + + manager.clear_external_auth(); + assert!(manager.auth_change_state_for(&initial_auth).is_none()); + *source.current.write().unwrap() = Some(initial_auth.clone()); + manager.set_external_auth(source).await.unwrap(); + assert_ne!( + initial.source_generation, + manager + .auth_change_state_for(&initial_auth) + .unwrap() + .source_generation + ); +} + #[tokio::test] #[serial_test::serial(codex_auth_env)] async fn legacy_identity_persistence_does_not_reload_credentials() -> anyhow::Result<()> { diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 293037f561..ab98e892df 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -46,6 +46,7 @@ pub use auth::AuthKeyringBackendKind; pub use auth::AuthManager; pub use auth::AuthManagerConfig; pub use auth::AuthManagerInitializationError; +pub use auth::AuthOwner; pub use auth::AuthRuntimeConfig; pub use auth::CLIENT_ID; pub use auth::CLIENT_ID_OVERRIDE_ENV_VAR; diff --git a/codex-rs/model-provider/src/auth.rs b/codex-rs/model-provider/src/auth.rs index 95dca9a74f..8fd61b3fa6 100644 --- a/codex-rs/model-provider/src/auth.rs +++ b/codex-rs/model-provider/src/auth.rs @@ -318,16 +318,22 @@ pub(crate) fn auth_manager_for_provider( ) -> Option> { match provider.auth.clone() { Some(config) => { - let mut env = HashMap::new(); - if provider.is_gcp() - && let Some(credential_file) = &provider.credential_file - { - env.insert( - "GOOGLE_APPLICATION_CREDENTIALS".to_string(), - credential_file.as_os_str().to_os_string(), - ); + if provider.is_gcp() { + let mut env = HashMap::new(); + if let Some(credential_file) = &provider.credential_file { + env.insert( + "GOOGLE_APPLICATION_CREDENTIALS".to_string(), + credential_file.as_os_str().to_os_string(), + ); + } + Some(AuthManager::external_bearer_only_with_env(config, env)) + } else { + Some(match auth_manager { + Some(manager) => crate::shared_state::process_shared_state() + .command_credentials(&config, &manager.runtime_config()), + None => AuthManager::external_bearer_only(config), + }) } - Some(AuthManager::external_bearer_only_with_env(config, env)) } None => auth_manager, } @@ -428,7 +434,7 @@ fn should_bootstrap_chatgpt_agent_identity( && matches!(auth, Some(CodexAuth::Chatgpt(_))) } -fn bearer_auth_for_provider( +pub(crate) fn bearer_auth_for_provider( provider: &ModelProviderInfo, ) -> codex_protocol::error::Result> { if let Some(api_key) = provider.api_key()? { diff --git a/codex-rs/model-provider/src/client_auth.rs b/codex-rs/model-provider/src/client_auth.rs index 3e229cb4a4..28f67aac0f 100644 --- a/codex-rs/model-provider/src/client_auth.rs +++ b/codex-rs/model-provider/src/client_auth.rs @@ -7,6 +7,7 @@ use codex_features::RuntimeAuthMode; use codex_http_client::HttpClientFactory; use codex_login::CodexAuth; use codex_protocol::error::Result as CodexResult; +use codex_runtime_agent_instance::RuntimeAuthBinding; use codex_runtime_agent_instance::RuntimeAuthManager; use http::HeaderMap; use std::sync::Arc; @@ -22,7 +23,7 @@ pub struct ClientAuth { enum ClientAuthSource { Provider(SharedAuthProvider), Runtime { - manager: Arc, + manager: Arc, mode: RuntimeAuthMode, source_generation: u64, }, @@ -46,8 +47,8 @@ impl ClientAuth { let resolved = crate::auth_provider_for_request(auth, runtime, headers).await?; let factory = resolved.bind_http_client_factory(factory)?; let source = match (runtime, resolved.runtime) { - (Some((manager, mode)), Some(selected)) => ClientAuthSource::Runtime { - manager: Arc::clone(manager), + (Some((_, mode)), Some(selected)) => ClientAuthSource::Runtime { + manager: Arc::clone(selected.binding()), mode: *mode, source_generation: selected.change_state().source_generation, }, diff --git a/codex-rs/model-provider/src/gcp/gcp_tests.rs b/codex-rs/model-provider/src/gcp/gcp_tests.rs index 08a9c2d325..7b8757b085 100644 --- a/codex-rs/model-provider/src/gcp/gcp_tests.rs +++ b/codex-rs/model-provider/src/gcp/gcp_tests.rs @@ -204,8 +204,9 @@ async fn credential_files_scope_token_caches_and_refreshes() { first_info.credential_file = Some(first_file.clone()); let mut second_info = first_info.clone(); second_info.credential_file = Some(second_file); - let first = GcpModelProvider::new(first_info, /*auth_manager*/ None); - let second = GcpModelProvider::new(second_info, /*auth_manager*/ None); + let primary = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("openai-key")); + let first = GcpModelProvider::new(first_info, Some(Arc::clone(&primary))); + let second = GcpModelProvider::new(second_info, Some(primary)); assert_eq!( (first.auth().await, second.auth().await), diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index 9a84789afd..ec8213c966 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -62,9 +62,11 @@ pub use provider::resolve_api_provider; pub use request_preparation::prepare_provider_request; pub use request_setup::ResolvedProviderRequest; +pub use request_setup::request_auth_owner; pub use request_setup::resolve_request; pub use runtime_auth::runtime_auth_for_url; pub(crate) use runtime_auth::select_runtime_auth_for_source; +pub use runtime_auth::select_runtime_provider; pub use runtime_auth::validate_runtime_destination; -pub use runtime_auth::validate_runtime_provider; +pub use runtime_auth::validate_runtime_environment; pub use runtime_auth::validate_runtime_selectors; diff --git a/codex-rs/model-provider/src/models_endpoint_refresh_tests.rs b/codex-rs/model-provider/src/models_endpoint_refresh_tests.rs index 7518892843..da34c1d88d 100644 --- a/codex-rs/model-provider/src/models_endpoint_refresh_tests.rs +++ b/codex-rs/model-provider/src/models_endpoint_refresh_tests.rs @@ -199,12 +199,16 @@ async fn check_timeout_preserves_refreshed_tokens(caller: Caller) { .auth_change_state_receiver() .borrow() .source_generation; - let runtime = &runtime_auth.as_ref().unwrap().0; + let runtime = runtime_auth + .as_ref() + .unwrap() + .0 + .bind(Arc::clone(&credentials)); assert!( timeout( MODELS_REFRESH_TIMEOUT, crate::select_runtime_auth_for_source( - runtime, + &runtime, generation, &http::HeaderMap::new() ) diff --git a/codex-rs/model-provider/src/models_endpoint_runtime_auth_tests.rs b/codex-rs/model-provider/src/models_endpoint_runtime_auth_tests.rs index 3dee3965dd..0f95816401 100644 --- a/codex-rs/model-provider/src/models_endpoint_runtime_auth_tests.rs +++ b/codex-rs/model-provider/src/models_endpoint_runtime_auth_tests.rs @@ -205,6 +205,87 @@ async fn model_requests_keep_the_selected_project_while_reusing_one_runtime() { } } +#[tokio::test] +async fn runtime_registration_environment_matches_inference_and_catalog() { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + for (registration_base, destination, matching_destination) in [ + ( + "https://chatgpt.com", + Some("https://api.openai-staging.com/v1"), + "https://eu.api.openai.com/v1", + ), + ( + "https://chatgpt-staging.com", + Some("https://api.openai.com/v1"), + "https://api.openai-staging.com/v1", + ), + ( + "https://chatgpt-staging.com", + None, + "https://chatgpt-staging.com/backend-api/codex", + ), + ] { + let home = tempfile::tempdir().unwrap(); + // Keep external network access disabled even if destination validation regresses. + let controller = codex_http_client::NetworkPolicyController::default(); + let factory = HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault) + .with_network_policy(controller.policy()); + let credentials = AuthManager::from_api_key( + "ambient-key", + codex_login::AuthRuntimeConfig { + codex_home: home.path().to_path_buf(), + auth_route_config: codex_login::AuthRouteConfig::from_http_client_factory( + factory.clone(), + ), + }, + ); + let runtime = Arc::new(RuntimeAuthManager::for_chatgpt( + Arc::clone(&credentials), + "00000000-0000-0000-0000-000000000000".into(), + registration_base.into(), + RuntimeRegistrationMetadata { + harness: None, + harness_version: None, + role: RuntimeRole::Orchestrator, + }, + )); + crate::validate_runtime_environment(&runtime, Some(matching_destination)).unwrap(); + let info = ModelProviderInfo { + experimental_bearer_token: Some("provider-key".into()), + env_http_headers: None, + ..ModelProviderInfo::create_openai_provider(destination.map(str::to_owned)) + }; + let provider = crate::provider::create_model_provider_with_runtime( + info.clone(), + Some(Arc::clone(&credentials)), + Some((Arc::clone(&runtime), mode)), + ); + let error = provider.request_auth(/*scope*/ None).await.err().unwrap(); + assert!(matches!( + error.details(), + codex_protocol::error::CodexErrorDetails::InvalidRequest(_) + )); + + let mut catalog_info = info; + if let Some(destination) = destination { + catalog_info.base_url = Some(matching_destination.into()); + catalog_info.model_catalog_url = Some(format!("{destination}/models").into()); + } + let catalog = OpenAiModelsEndpoint::new( + catalog_info, + Some(credentials), + /*gateway_auth_manager*/ None, + ) + .with_runtime_auth(Some((runtime, mode))); + let error = catalog.list_models("0.0.0", factory).await.unwrap_err(); + assert!(matches!( + error.details(), + codex_protocol::error::CodexErrorDetails::InvalidRequest(_) + )); + } + } +} + #[tokio::test] async fn unsupported_openai_destinations_do_not_fall_back_to_bearer_auth() { for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { diff --git a/codex-rs/model-provider/src/request_setup.rs b/codex-rs/model-provider/src/request_setup.rs index 407b89f97d..c09f331927 100644 --- a/codex-rs/model-provider/src/request_setup.rs +++ b/codex-rs/model-provider/src/request_setup.rs @@ -14,6 +14,7 @@ use codex_features::RuntimeAuthMode; use codex_http_client::ClientRouteClass; use codex_http_client::HttpClientFactory; use codex_login::AuthManager; +use codex_login::AuthOwner; use codex_login::CodexAuth; use codex_login::default_client::ClientRedirectPolicy; use codex_login::default_client::create_client_for_route; @@ -40,7 +41,9 @@ pub(crate) async fn resolve_request_auth( async move { if let Some((runtime, mode)) = runtime { crate::validate_runtime_destination(destination.as_deref())?; - crate::validate_runtime_provider(&info, manager.as_ref(), &runtime)?; + crate::validate_runtime_environment(&runtime, destination.as_deref())?; + let runtime = + crate::select_runtime_provider(&info, manager.as_ref(), &runtime)?; let headers = info.build_header_map()?; let selected = crate::runtime_auth::select_runtime_auth(&runtime, &headers) .await @@ -74,9 +77,23 @@ pub(crate) async fn resolve_request_auth( .await? } +/// Capture the selected source's owner before starting work such as a turn or classification. +/// Request preparation returns the owner it actually used so callers can detect a source change. +pub fn request_auth_owner(provider: &dyn ModelProvider) -> Option { + let credentials = provider.auth_manager(); + match provider.runtime_auth() { + Some((runtime, _)) => { + crate::select_runtime_provider(provider.info(), credentials.as_ref(), runtime) + .ok() + .map(|binding| binding.auth_owner()) + } + None => credentials.map(|manager| manager.auth_owner()), + } +} + pub struct ResolvedProviderRequest { pub auth: Option, - pub auth_owner_generation: Option, + pub auth_owner: Option, pub auth_revision: Option, pub api_provider: Provider, pub redirect_policy: ClientRedirectPolicy, @@ -123,6 +140,9 @@ pub async fn resolve_request( crate::prepare_provider_request(provider, check_auth, |state| { let scope = scope.clone(); let http_client_factory = http_client_factory.clone(); + let owner = state + .and_then(|_| provider.auth_manager()) + .map(|manager| manager.auth_owner()); async move { let (api_provider, resolved_auth) = provider.request_provider_and_auth(scope).await?; let auth = match &resolved_auth.credential { @@ -152,9 +172,11 @@ pub async fn resolve_request( .map(|selected| selected.runtime().id().to_owned()); Ok(ResolvedProviderRequest { auth, - auth_owner_generation: runtime_state - .map(|state| state.owner_generation) - .or(state.map(|state| state.owner_generation)), + auth_owner: resolved_auth + .runtime + .as_ref() + .map(|selected| selected.auth_owner().clone()) + .or(owner), auth_revision: runtime_state .map(|state| state.generation) .or(state.map(|state| state.generation)), diff --git a/codex-rs/model-provider/src/runtime_auth.rs b/codex-rs/model-provider/src/runtime_auth.rs index e02fb78b4c..62272bf204 100644 --- a/codex-rs/model-provider/src/runtime_auth.rs +++ b/codex-rs/model-provider/src/runtime_auth.rs @@ -14,6 +14,7 @@ use codex_protocol::error::Result as CodexResult; use codex_runtime_agent_instance::ActorInfo; use codex_runtime_agent_instance::PlatformResourceSelection; use codex_runtime_agent_instance::RuntimeAuth; +use codex_runtime_agent_instance::RuntimeAuthBinding; use codex_runtime_agent_instance::RuntimeAuthError; use codex_runtime_agent_instance::RuntimeAuthManager; use http::HeaderMap; @@ -21,7 +22,7 @@ use std::sync::Arc; /// Select the runtime for the resource scope carried by the model request. pub(crate) async fn select_runtime_auth( - manager: &RuntimeAuthManager, + manager: &Arc, headers: &HeaderMap, ) -> Result { manager @@ -32,7 +33,7 @@ pub(crate) async fn select_runtime_auth( /// Select fresh request auth without allowing a reusable client's source to change. pub(crate) async fn select_runtime_auth_for_source( - manager: &Arc, + manager: &Arc, source_generation: u64, headers: &HeaderMap, ) -> Result { @@ -74,24 +75,29 @@ pub(crate) fn resource_selection( }) } -/// A runtime may not replace an independently selected provider credential. -pub fn validate_runtime_provider( +/// Use the credential selected by provider configuration, never an ambient actor's proof. +pub fn select_runtime_provider( info: &ModelProviderInfo, credentials: Option<&Arc>, - runtime: &RuntimeAuthManager, -) -> CodexResult<()> { - if !info.requires_openai_auth - || info.env_key.is_some() - || info.experimental_bearer_token.is_some() - || info.auth.is_some() - || info.aws.is_some() - || !credentials.is_some_and(|auth| runtime.uses_credentials(auth)) - { - return Err(CodexErr::UnsupportedOperation( + runtime: &Arc, +) -> CodexResult> { + let unsupported = || { + CodexErr::UnsupportedOperation( "runtime auth requires the model provider's selected OpenAI credential".to_string(), - )); + ) + }; + if info.aws.is_some() { + return Err(unsupported()); } - Ok(()) + let credentials = credentials.ok_or_else(unsupported)?; + if let Some(bearer) = crate::auth::bearer_auth_for_provider(info)? { + let token = bearer.token.ok_or_else(unsupported)?; + return Ok(runtime.bind_api_key(&token, &credentials.runtime_config())); + } + if !info.requires_openai_auth && info.auth.is_none() { + return Err(unsupported()); + } + Ok(runtime.bind(Arc::clone(credentials))) } /// An unrecognized OpenAI origin must not silently fall back to legacy auth. @@ -105,6 +111,48 @@ pub fn validate_runtime_destination(destination: Option<&str>) -> CodexResult<() Ok(()) } +/// Do not send a provider's credentials to registration in another OpenAI environment. +pub fn validate_runtime_environment( + runtime: &RuntimeAuthManager, + destination: Option<&str>, +) -> CodexResult<()> { + let registration_url = runtime + .registration_url() + .map_err(|err| CodexErr::InvalidRequest(err.to_string()))?; + let registration_is_staging = match registration_url + .host_str() + .map(|host| host.trim_end_matches('.')) + { + Some("auth.openai.com") => false, + Some("auth.api.openai.org") => true, + // Explicit registration targets also support local test servers. + _ => return Ok(()), + }; + let destination_is_staging = match destination { + Some(destination) => { + let destination = url::Url::parse(destination) + .map_err(|_| CodexErr::InvalidRequest("invalid runtime auth destination".into()))?; + let host = destination + .host_str() + .unwrap_or_default() + .trim_end_matches('.'); + matches!( + host, + "api.openai.org" | "api.openai-staging.com" | "chatgpt-staging.com" + ) || host.ends_with(".chatgpt-staging.com") + } + // Both implicit model-provider defaults use production. + None => false, + }; + if registration_is_staging != destination_is_staging { + return Err(CodexErr::InvalidRequest( + "Requests must use the same production or staging environment as runtime registration" + .into(), + )); + } + Ok(()) +} + /// Select runtime auth for supported OpenAI URLs; reject unsupported OpenAI URLs. pub fn runtime_auth_for_url<'a>( runtime: Option<&'a (Arc, RuntimeAuthMode)>, @@ -113,7 +161,11 @@ pub fn runtime_auth_for_url<'a>( if runtime.is_some() { validate_runtime_destination(Some(url))?; } - Ok(runtime.filter(|_| classify_openai_api_url(url) == OpenAiDestination::Supported)) + let runtime = runtime.filter(|_| classify_openai_api_url(url) == OpenAiDestination::Supported); + if let Some((manager, _)) = runtime { + validate_runtime_environment(manager, Some(url))?; + } + Ok(runtime) } /// Account and organization selectors must agree with the registered actor. @@ -143,3 +195,7 @@ pub fn validate_runtime_selectors(headers: &HeaderMap, actor: &ActorInfo) -> Cod } } } + +#[cfg(test)] +#[path = "runtime_auth_tests.rs"] +mod tests; diff --git a/codex-rs/model-provider/src/runtime_auth_tests.rs b/codex-rs/model-provider/src/runtime_auth_tests.rs new file mode 100644 index 0000000000..65bca05e98 --- /dev/null +++ b/codex-rs/model-provider/src/runtime_auth_tests.rs @@ -0,0 +1,243 @@ +use super::*; +use crate::ProviderAuthContext; +use codex_features::RuntimeAuthMode; +use codex_login::CodexAuth; +use codex_login::auth::BedrockApiKeyAuth; +use codex_protocol::config_types::ModelProviderAuthInfo; +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::num::NonZeroU64; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; + +async fn context( + server: &MockServer, + credentials: Arc, + mode: RuntimeAuthMode, +) -> ProviderAuthContext { + Mock::given(method("POST")) + .respond_with(|request: &wiremock::Request| ResponseTemplate::new(200).set_body_json(json!({ + "runtime_agent_instance_id": request.headers["idempotency-key"].to_str().unwrap(), + "agent_host_id": "host-a", + "actor": {"kind": "platform", "user_id": "user-a", "organization_id": "org-a", "project_id": null}, + }))) + .mount(server).await; + let runtime = RuntimeAuthManager::new( + Arc::clone(&credentials), + AgentHostAuthorization::InstallationId { + installation_id: "00000000-0000-0000-0000-000000000000".parse().unwrap(), + }, + RuntimeRegistrationMetadata { + harness: None, + harness_version: None, + role: RuntimeRole::Orchestrator, + }, + format!("{}/agent/register", server.uri()).parse().unwrap(), + ); + ProviderAuthContext { + credentials, + runtime_auth: Some((Arc::new(runtime), mode)), + } +} + +#[tokio::test] +async fn explicit_keys_and_commands_use_their_own_shared_runtime() { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + let home = tempfile::tempdir().unwrap(); + let credentials = AuthManager::from_auth_for_testing_with_home( + CodexAuth::from_api_key("ambient-key"), + home.path().to_path_buf(), + ); + let context = context(&server, credentials, mode).await; + let literal = ModelProviderInfo { + experimental_bearer_token: Some("literal-key".into()), + requires_openai_auth: false, + ..ModelProviderInfo::create_openai_provider(/*base_url*/ None) + }; + #[cfg(unix)] + let (command, args) = ("sh", vec!["-c".into(), "printf command-key".into()]); + #[cfg(windows)] + let (command, args) = ( + "cmd.exe", + vec![ + "/D".into(), + "/Q".into(), + "/C".into(), + "echo command-key".into(), + ], + ); + let commanded = ModelProviderInfo { + auth: Some(ModelProviderAuthInfo { + command: command.into(), + args, + timeout_ms: NonZeroU64::new(5000).unwrap(), + refresh_interval_ms: 300_000, + cwd: home.path().to_path_buf().try_into().unwrap(), + }), + requires_openai_auth: false, + ..ModelProviderInfo::create_openai_provider(/*base_url*/ None) + }; + for (info, key) in [(literal.clone(), "literal-key"), (commanded, "command-key")] { + let mut runtime_id = None; + for _ in 0..2 { + let provider = context.model_provider(info.clone()); + let selected = provider.request_auth(/*scope*/ None).await.unwrap(); + let runtime = selected.runtime.as_ref().unwrap(); + assert_eq!(runtime.auth().api_key(), Some(key)); + match &runtime_id { + Some(id) => assert_eq!(runtime.runtime().id(), id), + None => runtime_id = Some(runtime.runtime().id().to_string()), + } + let headers = selected + .auth + .apply_auth_headers(Default::default()) + .await + .unwrap(); + match mode { + RuntimeAuthMode::DualHeader => { + assert_eq!(headers["authorization"], format!("Bearer {key}")); + assert!(headers.contains_key("x-openai-agent-assertion")); + } + RuntimeAuthMode::AssertionOnly => { + assert!( + headers["authorization"] + .to_str() + .unwrap() + .starts_with("AgentAssertion ") + ); + assert!(!headers.contains_key("x-openai-agent-assertion")); + } + } + } + } + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 2); + for (request, key) in requests.iter().zip(["literal-key", "command-key"]) { + assert_eq!( + request.body_json::().unwrap()["actor_auth"]["credential"], + key + ); + } + let custom = context.model_provider(ModelProviderInfo { + base_url: Some("https://provider.example/v1".into()), + ..literal + }); + let auth = custom.request_auth(/*scope*/ None).await.unwrap(); + assert!(auth.runtime.is_none()); + let headers = auth + .auth + .apply_auth_headers(Default::default()) + .await + .unwrap(); + assert_eq!(headers["authorization"], "Bearer literal-key"); + assert!(!headers.contains_key("x-openai-agent-assertion")); + assert_eq!(server.received_requests().await.unwrap().len(), 2); + } +} + +#[tokio::test] +async fn explicit_provider_key_keeps_the_requested_project_scope() { + for mode in [RuntimeAuthMode::DualHeader, RuntimeAuthMode::AssertionOnly] { + let server = MockServer::start().await; + let credentials = + AuthManager::from_auth_for_testing(CodexAuth::from_api_key("ambient-key")); + let context = context(&server, credentials, mode).await; + server.reset().await; + Mock::given(method("POST")) + .and(wiremock::matchers::body_partial_json(json!({ + "actor_auth": { + "type": "platform_api_key", + "credential": "selected-key", + "project_id": "proj-selected", + }, + }))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "runtime_agent_instance_id": "rai-selected", + "agent_host_id": "host-a", + "actor": { + "kind": "platform", + "user_id": "user-a", + "organization_id": "org-a", + "project_id": "proj-selected", + }, + }))) + .mount(&server) + .await; + let provider = context.model_provider(ModelProviderInfo { + experimental_bearer_token: Some("selected-key".into()), + http_headers: Some(std::collections::HashMap::from([( + OPENAI_PROJECT_HEADER.into(), + "proj-selected".into(), + )])), + requires_openai_auth: false, + ..ModelProviderInfo::create_openai_provider(/*base_url*/ None) + }); + let selected = provider.request_auth(/*scope*/ None).await.unwrap(); + let runtime = selected.runtime.as_ref().unwrap(); + assert_eq!(runtime.auth().api_key(), Some("selected-key")); + assert!(matches!( + runtime.runtime().actor().info(), + ActorInfo::Platform { project_id: Some(project), .. } if project == "proj-selected" + )); + } +} + +#[tokio::test] +async fn bedrock_auth_stays_independent_when_openai_runtime_auth_is_enabled() { + let server = MockServer::start().await; + let credentials = + AuthManager::from_auth_for_testing(CodexAuth::BedrockApiKey(BedrockApiKeyAuth { + api_key: "bedrock-key".into(), + region: "us-east-1".into(), + })); + let context = context(&server, credentials, RuntimeAuthMode::DualHeader).await; + let bedrock = context.model_provider(ModelProviderInfo::create_amazon_bedrock_provider( + /*aws*/ None, + )); + let headers = bedrock + .api_auth() + .await + .unwrap() + .apply_auth_headers(Default::default()) + .await + .unwrap(); + assert_eq!(headers["authorization"], "Bearer bedrock-key"); + assert!(!headers.contains_key("x-openai-agent-assertion")); + assert!(server.received_requests().await.unwrap().is_empty()); + + let openai = context.model_provider(ModelProviderInfo { + experimental_bearer_token: Some("openai-key".into()), + ..ModelProviderInfo::create_openai_provider(/*base_url*/ None) + }); + let headers = openai + .api_auth() + .await + .unwrap() + .apply_auth_headers(Default::default()) + .await + .unwrap(); + assert_eq!(headers["authorization"], "Bearer openai-key"); + assert!(headers.contains_key("x-openai-agent-assertion")); + let headers = bedrock + .api_auth() + .await + .unwrap() + .apply_auth_headers(Default::default()) + .await + .unwrap(); + assert_eq!(headers["authorization"], "Bearer bedrock-key"); + assert!(!headers.contains_key("x-openai-agent-assertion")); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!( + requests[0].body_json::().unwrap()["actor_auth"]["credential"], + "openai-key" + ); +} diff --git a/codex-rs/model-provider/src/shared_state.rs b/codex-rs/model-provider/src/shared_state.rs index 24c2116b15..45689e8694 100644 --- a/codex-rs/model-provider/src/shared_state.rs +++ b/codex-rs/model-provider/src/shared_state.rs @@ -13,6 +13,7 @@ use codex_model_provider_info::CredentialExportConfig; use codex_model_provider_info::GatewayOAuthSettings; use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_model_provider_info::ModelProviderInfo; +use codex_protocol::config_types::ModelProviderAuthInfo; use crate::amazon_bedrock::AwsAuthRecovery; use crate::amazon_bedrock::AwsCredentialExport; @@ -20,6 +21,7 @@ use crate::amazon_bedrock::AwsCredentialExport; /// Provider-owned runtime state shared across independently configured sessions. #[derive(Debug, Default)] pub(crate) struct ModelProviderSharedState { + command_credentials: Mutex)>>, aws_credential_exports: Mutex)>>, aws_auth_recoveries: Mutex)>>, gateway_exports: Mutex>, @@ -45,6 +47,28 @@ pub(crate) fn process_shared_state() -> &'static ModelProviderSharedState { } impl ModelProviderSharedState { + pub(crate) fn command_credentials( + &self, + config: &ModelProviderAuthInfo, + runtime: &AuthRuntimeConfig, + ) -> Arc { + let mut owners = self + .command_credentials + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + owners.retain(|(_, _, owner)| owner.strong_count() != 0); + if let Some(owner) = owners + .iter() + .find(|(source, context, _)| source == config && context == runtime) + .and_then(|(_, _, owner)| owner.upgrade()) + { + return owner; + } + let owner = AuthManager::external_bearer_with_runtime(config.clone(), runtime.clone()); + owners.push((config.clone(), runtime.clone(), Arc::downgrade(&owner))); + owner + } + pub(crate) fn gateway_oauth_session( &self, config: &GatewayOAuthSettings, diff --git a/codex-rs/runtime-agent-instance/src/auth.rs b/codex-rs/runtime-agent-instance/src/auth.rs index 5cd88e3da9..a59952e526 100644 --- a/codex-rs/runtime-agent-instance/src/auth.rs +++ b/codex-rs/runtime-agent-instance/src/auth.rs @@ -3,6 +3,7 @@ use crate::ActorInfo; use crate::RegistrationError; use crate::RuntimeAgentInstance; +use crate::RuntimeAuthBinding; use crate::RuntimeAuthError; use codex_api::AuthError; use codex_api::AuthHeaderTelemetry; @@ -12,6 +13,7 @@ use codex_features::RuntimeAuthMode; use codex_http_client::CHATGPT_ACCOUNT_ID_HEADER; use codex_http_client::HttpClientFactory; use codex_login::AuthChangeState; +use codex_login::AuthOwner; use codex_login::CodexAuth; use http::HeaderMap; use http::HeaderValue; @@ -28,6 +30,8 @@ pub struct RuntimeAuth { pub(crate) auth: CodexAuth, pub(crate) http_client_factory: HttpClientFactory, pub(crate) change_state: AuthChangeState, + pub(crate) owner: AuthOwner, + pub(crate) binding: Arc, } impl std::fmt::Debug for RuntimeAuth { @@ -63,6 +67,14 @@ impl RuntimeAuth { self.change_state } + pub fn auth_owner(&self) -> &AuthOwner { + &self.owner + } + + pub fn binding(&self) -> &Arc { + &self.binding + } + /// Preserve the caller's transport settings and restrictions while binding requests to /// this credential's account. Independent managed policies must not replace each other. pub fn bind_http_client_factory( diff --git a/codex-rs/runtime-agent-instance/src/binding.rs b/codex-rs/runtime-agent-instance/src/binding.rs new file mode 100644 index 0000000000..5e4e93e66e --- /dev/null +++ b/codex-rs/runtime-agent-instance/src/binding.rs @@ -0,0 +1,408 @@ +//! Register and refresh a runtime within one selected credential source. + +use crate::AgentHostAuthorization; +use crate::AuthenticatedActorAuthorization; +use crate::PlatformResourceSelection; +use crate::RegistrationError; +use crate::RuntimeAgentInstance; +use crate::RuntimeAuth; +use crate::RuntimeRegistration; +use crate::RuntimeRegistrationClient; +use crate::RuntimeRegistrationMetadata; +use codex_agent_identity::ChatGptEnvironment; +use codex_http_client::HttpClientFactory; +use codex_http_client::NetworkPolicyDenied; +use codex_login::AuthChangeState; +use codex_login::AuthManager; +use codex_login::AuthOwner; +use codex_login::CodexAuth; +use codex_login::RefreshTokenError; +use http::StatusCode; +use std::sync::Arc; +use tokio::sync::Mutex; +use url::Url; + +/// A selected credential source and its registration state, retained across request preparation. +pub struct RuntimeAuthBinding { + pub(crate) credentials: Arc, + pub(crate) registration: RegistrationTarget, + pub(crate) metadata: RuntimeRegistrationMetadata, + state: Mutex>, +} + +#[derive(Clone)] +pub(crate) enum RegistrationTarget { + Explicit { + host: AgentHostAuthorization, + url: Url, + }, + ChatGpt { + installation_id: String, + base_url: String, + }, +} + +impl RegistrationTarget { + fn host(&self) -> Result { + match self { + Self::Explicit { host, .. } => Ok(host.clone()), + Self::ChatGpt { + installation_id, .. + } => Ok(AgentHostAuthorization::InstallationId { + installation_id: installation_id + .parse() + .map_err(|_| RegistrationError::InvalidRequest)?, + }), + } + } + + fn url(&self) -> Result { + match self { + Self::Explicit { url, .. } => Ok(url.clone()), + Self::ChatGpt { base_url, .. } => { + let url = match ChatGptEnvironment::from_chatgpt_base_url(base_url) + .map_err(|_| RegistrationError::InvalidEndpoint)? + { + ChatGptEnvironment::Production => crate::DEFAULT_REGISTRATION_URL, + ChatGptEnvironment::Staging => "https://auth.api.openai.org/agent/register", + }; + url.parse().map_err(|_| RegistrationError::InvalidEndpoint) + } + } + } +} + +struct CachedRuntime { + source_generation: u64, + selection: PlatformResourceSelection, + state: RuntimeState, +} + +enum RuntimeState { + Pending { registration: RuntimeRegistration }, + Ready { runtime: Arc }, + Denied { error: RegistrationError }, +} + +impl std::fmt::Debug for RuntimeAuthBinding { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RuntimeAuthBinding").finish_non_exhaustive() + } +} + +impl RuntimeAuthBinding { + pub(crate) fn new( + credentials: Arc, + registration: RegistrationTarget, + metadata: RuntimeRegistrationMetadata, + ) -> Self { + Self { + credentials, + registration, + metadata, + state: Mutex::new(None), + } + } + + pub(crate) fn uses_credentials(&self, credentials: &Arc) -> bool { + Arc::ptr_eq(&self.credentials, credentials) + } + + pub fn auth_owner(&self) -> AuthOwner { + self.credentials.auth_owner() + } + + pub fn registration_url(&self) -> Result { + self.registration.url() + } + + pub async fn auth(self: &Arc) -> Result { + self.auth_with_selection(&PlatformResourceSelection::default()) + .await + } + + /// Select auth and a runtime for the requested organization and project. + pub async fn auth_with_selection( + self: &Arc, + selection: &PlatformResourceSelection, + ) -> Result { + self.select_auth(selection, /*expected_source*/ None).await + } + + #[expect( + clippy::await_holding_invalid_type, + reason = "runtime initialization serializes registration and refresh for one credential source" + )] + async fn select_auth( + self: &Arc, + selection: &PlatformResourceSelection, + expected_source: Option, + ) -> Result { + // This initializer's lock can span IO. No AuthManager lock is held across registration. + let mut state = self.state.lock().await; + self.check_source(expected_source)?; + let (auth, http_client_factory) = self + .credentials + .auth_with_http_client_factory() + .await + .ok_or(RuntimeAuthError::Unavailable)?; + let mut recovery = self.credentials.unauthorized_recovery(); + let mut result = self + .auth_once( + &mut state, + auth, + http_client_factory, + selection, + expected_source, + ) + .await; + while matches!( + &result, + Err(RuntimeAuthError::Registration( + RegistrationError::Rejected { + status: StatusCode::UNAUTHORIZED, + .. + } + )) + ) && recovery.has_next() + { + self.check_source(expected_source)?; + recovery.next().await.map_err(|error| match error { + RefreshTokenError::Policy(denied) => RuntimeAuthError::Policy(denied), + RefreshTokenError::Permanent(_) | RefreshTokenError::Transient(_) => { + RuntimeAuthError::RefreshFailed + } + })?; + let (auth, http_client_factory) = self + .credentials + .auth_with_http_client_factory() + .await + .ok_or(RuntimeAuthError::Unavailable)?; + result = self + .auth_once( + &mut state, + auth, + http_client_factory, + selection, + expected_source, + ) + .await; + } + result + } + + async fn auth_once( + self: &Arc, + state: &mut Option, + auth: CodexAuth, + http_client_factory: HttpClientFactory, + selection: &PlatformResourceSelection, + expected_source: Option, + ) -> Result { + let actor_auth = actor_authorization(&auth, selection)?; + let http_client_factory = http_client_factory.with_sensitive_header_protection(); + let change_state = self.auth_state(&auth)?; + self.check_source(expected_source)?; + // Registered runtimes can serve different projects. Pending and denied + // registrations retain their original scope and idempotency key. + if state.as_ref().is_some_and(|cached| { + cached.source_generation != change_state.source_generation + || cached.selection.organization_id != selection.organization_id + || (!matches!(cached.state, RuntimeState::Ready { .. }) + && cached.selection != *selection) + }) { + *state = None; + } + let cached = match state { + Some(cached) => cached, + None => state.insert(CachedRuntime { + source_generation: change_state.source_generation, + selection: selection.clone(), + state: RuntimeState::Pending { + registration: RuntimeRegistration::new( + actor_auth.clone(), + self.registration.host()?, + self.metadata.clone(), + ) + .map_err(|_| RuntimeAuthError::KeyGeneration)?, + }, + }), + }; + let runtime_state = &mut cached.state; + let registration = match runtime_state { + RuntimeState::Ready { runtime } => { + return self.selected_auth( + runtime.clone(), + auth, + http_client_factory, + change_state, + ); + } + RuntimeState::Denied { error } => return Err(error.clone().into()), + RuntimeState::Pending { registration } => { + registration.request.actor_auth = actor_auth; + registration + } + }; + let client = + RuntimeRegistrationClient::new(http_client_factory.clone(), self.registration.url()?)?; + let result = client.register(registration).await; + if self.auth_state(&auth)?.source_generation != change_state.source_generation { + // A late response must not publish an old identity after a credential switch. + *state = None; + return Err(RuntimeAuthError::Changed); + } + match result { + Ok(info) => { + let runtime = Arc::new(registration.finish(info)); + *runtime_state = RuntimeState::Ready { + runtime: runtime.clone(), + }; + self.selected_auth(runtime, auth, http_client_factory, change_state) + } + Err( + error @ RegistrationError::Rejected { + status: StatusCode::FORBIDDEN, + .. + }, + ) => { + *runtime_state = RuntimeState::Denied { + error: error.clone(), + }; + Err(error.into()) + } + Err(error) => Err(error.into()), + } + } + + fn selected_auth( + self: &Arc, + runtime: Arc, + auth: CodexAuth, + http_client_factory: HttpClientFactory, + change_state: AuthChangeState, + ) -> Result { + let owner = self + .credentials + .auth_owner_for(&auth) + .ok_or(RuntimeAuthError::Changed)?; + Ok(RuntimeAuth { + runtime, + auth, + http_client_factory, + change_state, + owner, + binding: Arc::clone(self), + }) + } + + /// Resolve auth for work whose routing, cache key, or payload uses a captured credential. + /// A same-source refresh is allowed; a different source must not take over that work. + /// This binds credential identity, not when the caller captured it. Restoring the exact + /// same credential can select a new runtime after logout; a different source cannot. + pub async fn auth_for( + self: &Arc, + expected: &CodexAuth, + selection: &PlatformResourceSelection, + ) -> Result { + let source = self + .credentials + .auth_change_state_for(expected) + .ok_or(RuntimeAuthError::Changed)?; + self.auth_for_source(source.source_generation, selection) + .await + } + + /// Select another request within a reusable client's original credential source. + pub async fn auth_for_source( + self: &Arc, + source_generation: u64, + selection: &PlatformResourceSelection, + ) -> Result { + self.select_auth(selection, Some(source_generation)).await + } + + fn check_source(&self, expected: Option) -> Result<(), RuntimeAuthError> { + if expected.is_some_and(|expected| { + self.credentials + .auth_change_state_receiver() + .borrow() + .source_generation + != expected + }) { + return Err(RuntimeAuthError::Changed); + } + Ok(()) + } + + fn auth_state(&self, auth: &CodexAuth) -> Result { + let state = self + .credentials + .auth_change_state_for(auth) + .ok_or(RuntimeAuthError::Changed)?; + if self.credentials.refresh_failure_for_auth(auth).is_some() { + return Err(RuntimeAuthError::RefreshFailed); + } + Ok(state) + } +} + +fn actor_authorization( + auth: &CodexAuth, + selection: &PlatformResourceSelection, +) -> Result { + if !matches!(auth, CodexAuth::ApiKey(_)) && selection != &PlatformResourceSelection::default() { + return Err(RuntimeAuthError::UnsupportedResourceSelection); + } + let credential = auth + .get_token() + .map_err(|_| RuntimeAuthError::Unsupported)? + .into(); + Ok(match auth { + CodexAuth::ApiKey(_) => AuthenticatedActorAuthorization::PlatformApiKey { + credential, + selection: selection.clone(), + }, + CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) => { + AuthenticatedActorAuthorization::Oauth { credential } + } + CodexAuth::PersonalAccessToken(_) => { + AuthenticatedActorAuthorization::PersonalAccessToken { credential } + } + CodexAuth::Headers(_) + | CodexAuth::AgentIdentity(_) + | CodexAuth::BedrockApiKey(_) + | CodexAuth::BedrockAccessKeys(_) => return Err(RuntimeAuthError::Unsupported), + }) +} + +#[derive(Debug, thiserror::Error)] +pub enum RuntimeAuthError { + #[error("runtime registration requires an actor credential")] + Unavailable, + #[error( + "this credential cannot register a runtime; sign in again with a supported OpenAI credential" + )] + Unsupported, + #[error("organization and project selection requires an OpenAI API key")] + UnsupportedResourceSelection, + #[error("credentials changed during runtime registration; try again")] + Changed, + #[error("runtime registration could not refresh actor credentials")] + RefreshFailed, + #[error("failed to generate runtime signing key")] + KeyGeneration, + #[error(transparent)] + Policy(#[from] NetworkPolicyDenied), + #[error(transparent)] + Registration(RegistrationError), +} + +impl From for RuntimeAuthError { + fn from(error: RegistrationError) -> Self { + match error { + RegistrationError::Policy(denied) => Self::Policy(denied), + error => Self::Registration(error), + } + } +} diff --git a/codex-rs/runtime-agent-instance/src/lib.rs b/codex-rs/runtime-agent-instance/src/lib.rs index f28529e367..97a39bfc5a 100644 --- a/codex-rs/runtime-agent-instance/src/lib.rs +++ b/codex-rs/runtime-agent-instance/src/lib.rs @@ -2,6 +2,7 @@ mod actor; mod auth; +mod binding; mod bootstrap; mod manager; mod registration; @@ -20,7 +21,8 @@ pub use bootstrap::RuntimeRole; pub use codex_http_client::is_openai_api_url; pub use auth::RuntimeAuth; -pub use manager::RuntimeAuthError; +pub use binding::RuntimeAuthBinding; +pub use binding::RuntimeAuthError; pub use manager::RuntimeAuthManager; pub use registration::DEFAULT_REGISTRATION_URL; pub use registration::RegistrationError; diff --git a/codex-rs/runtime-agent-instance/src/manager.rs b/codex-rs/runtime-agent-instance/src/manager.rs index e1bb6f204f..fcdde1a933 100644 --- a/codex-rs/runtime-agent-instance/src/manager.rs +++ b/codex-rs/runtime-agent-instance/src/manager.rs @@ -1,85 +1,28 @@ -//! Reuse one active runtime for the selected credential source and organization. +//! Bind credential sources to shared runtime registrations. use crate::AgentHostAuthorization; -use crate::AuthenticatedActorAuthorization; use crate::PlatformResourceSelection; use crate::RegistrationError; -use crate::RuntimeAgentInstance; use crate::RuntimeAuth; -use crate::RuntimeRegistration; -use crate::RuntimeRegistrationClient; +use crate::RuntimeAuthBinding; +use crate::RuntimeAuthError; use crate::RuntimeRegistrationMetadata; -use codex_agent_identity::ChatGptEnvironment; -use codex_http_client::HttpClientFactory; -use codex_http_client::NetworkPolicyDenied; -use codex_login::AuthChangeState; +use crate::binding::RegistrationTarget; use codex_login::AuthManager; +use codex_login::AuthRuntimeConfig; use codex_login::CodexAuth; -use codex_login::RefreshTokenError; -use http::StatusCode; use std::sync::Arc; -use tokio::sync::Mutex; use url::Url; -/// Owned alongside an AuthManager, not created for every thread or request. +/// Shares registration state across requests using the same credential source and context. pub struct RuntimeAuthManager { - credentials: Arc, - registration: RegistrationTarget, - metadata: RuntimeRegistrationMetadata, - state: Mutex>, + login: Arc, + registrations: std::sync::Mutex>, } -enum RegistrationTarget { - Explicit { - host: AgentHostAuthorization, - url: Url, - }, - ChatGpt { - installation_id: String, - base_url: String, - }, -} - -impl RegistrationTarget { - fn host(&self) -> Result { - match self { - Self::Explicit { host, .. } => Ok(host.clone()), - Self::ChatGpt { - installation_id, .. - } => Ok(AgentHostAuthorization::InstallationId { - installation_id: installation_id - .parse() - .map_err(|_| RegistrationError::InvalidRequest)?, - }), - } - } - - fn url(&self) -> Result { - match self { - Self::Explicit { url, .. } => Ok(url.clone()), - Self::ChatGpt { base_url, .. } => { - let url = match ChatGptEnvironment::from_chatgpt_base_url(base_url) - .map_err(|_| RegistrationError::InvalidEndpoint)? - { - ChatGptEnvironment::Production => crate::DEFAULT_REGISTRATION_URL, - ChatGptEnvironment::Staging => "https://auth.api.openai.org/agent/register", - }; - url.parse().map_err(|_| RegistrationError::InvalidEndpoint) - } - } - } -} - -struct CachedRuntime { - source_generation: u64, - selection: PlatformResourceSelection, - state: RuntimeState, -} - -enum RuntimeState { - Pending { registration: RuntimeRegistration }, - Ready { runtime: Arc }, - Denied { error: RegistrationError }, +enum CredentialRegistration { + Provider(Arc), + ApiKey(Arc), } impl std::fmt::Debug for RuntimeAuthManager { @@ -89,9 +32,8 @@ impl std::fmt::Debug for RuntimeAuthManager { } impl RuntimeAuthManager { - /// Check that a consumer uses this owner's credential provider, not an ambient fallback. pub fn uses_credentials(&self, credentials: &Arc) -> bool { - Arc::ptr_eq(&self.credentials, credentials) + self.login.uses_credentials(credentials) } pub fn new( @@ -100,15 +42,14 @@ impl RuntimeAuthManager { metadata: RuntimeRegistrationMetadata, registration_url: Url, ) -> Self { - Self { + Self::from_registration( credentials, - registration: RegistrationTarget::Explicit { + RegistrationTarget::Explicit { host, url: registration_url, }, metadata, - state: Mutex::new(None), - } + ) } /// Resolve the ChatGPT registration origin only when a first-party request needs a runtime. @@ -118,287 +59,119 @@ impl RuntimeAuthManager { base_url: String, metadata: RuntimeRegistrationMetadata, ) -> Self { - Self { + Self::from_registration( credentials, - registration: RegistrationTarget::ChatGpt { + RegistrationTarget::ChatGpt { installation_id, base_url, }, metadata, - state: Mutex::new(None), + ) + } + + fn from_registration( + credentials: Arc, + registration: RegistrationTarget, + metadata: RuntimeRegistrationMetadata, + ) -> Self { + Self { + login: Arc::new(RuntimeAuthBinding::new(credentials, registration, metadata)), + registrations: std::sync::Mutex::new(Vec::new()), } } - pub async fn auth(&self) -> Result { - self.auth_with_selection(&PlatformResourceSelection::default()) - .await + pub fn registration_url(&self) -> Result { + self.login.registration_url() + } + + pub async fn auth(&self) -> Result { + self.login.auth().await } - /// Select auth and a runtime for the requested organization and project. pub async fn auth_with_selection( &self, selection: &PlatformResourceSelection, ) -> Result { - self.select_auth(selection, /*expected_source*/ None).await + self.login.auth_with_selection(selection).await } - #[expect( - clippy::await_holding_invalid_type, - reason = "runtime initialization serializes registration and refresh for one credential source" - )] - async fn select_auth( - &self, - selection: &PlatformResourceSelection, - expected_source: Option, - ) -> Result { - // This initializer's lock can span IO. No AuthManager lock is held across registration. - let mut state = self.state.lock().await; - self.check_source(expected_source)?; - let (auth, http_client_factory) = self - .credentials - .auth_with_http_client_factory() - .await - .ok_or(RuntimeAuthError::Unavailable)?; - let mut recovery = self.credentials.unauthorized_recovery(); - let mut result = self - .auth_once( - &mut state, - auth, - http_client_factory, - selection, - expected_source, - ) - .await; - while matches!( - &result, - Err(RuntimeAuthError::Registration( - RegistrationError::Rejected { - status: StatusCode::UNAUTHORIZED, - .. - } - )) - ) && recovery.has_next() - { - self.check_source(expected_source)?; - recovery.next().await.map_err(|error| match error { - RefreshTokenError::Policy(denied) => RuntimeAuthError::Policy(denied), - RefreshTokenError::Permanent(_) | RefreshTokenError::Transient(_) => { - RuntimeAuthError::RefreshFailed - } - })?; - let (auth, http_client_factory) = self - .credentials - .auth_with_http_client_factory() - .await - .ok_or(RuntimeAuthError::Unavailable)?; - result = self - .auth_once( - &mut state, - auth, - http_client_factory, - selection, - expected_source, - ) - .await; - } - result - } - - async fn auth_once( - &self, - state: &mut Option, - auth: CodexAuth, - http_client_factory: HttpClientFactory, - selection: &PlatformResourceSelection, - expected_source: Option, - ) -> Result { - let actor_auth = actor_authorization(&auth, selection)?; - let http_client_factory = http_client_factory.with_sensitive_header_protection(); - let change_state = self.auth_state(&auth)?; - self.check_source(expected_source)?; - // Registered runtimes can serve different projects. Pending and denied - // registrations retain their original scope and idempotency key. - if state.as_ref().is_some_and(|cached| { - cached.source_generation != change_state.source_generation - || cached.selection.organization_id != selection.organization_id - || (!matches!(cached.state, RuntimeState::Ready { .. }) - && cached.selection != *selection) - }) { - *state = None; - } - let cached = match state { - Some(cached) => cached, - None => state.insert(CachedRuntime { - source_generation: change_state.source_generation, - selection: selection.clone(), - state: RuntimeState::Pending { - registration: RuntimeRegistration::new( - actor_auth.clone(), - self.registration.host()?, - self.metadata.clone(), - ) - .map_err(|_| RuntimeAuthError::KeyGeneration)?, - }, - }), - }; - let runtime_state = &mut cached.state; - let registration = match runtime_state { - RuntimeState::Ready { runtime } => { - return Ok(RuntimeAuth { - runtime: runtime.clone(), - auth, - http_client_factory, - change_state, - }); - } - RuntimeState::Denied { error } => return Err(error.clone().into()), - RuntimeState::Pending { registration } => { - registration.request.actor_auth = actor_auth; - registration - } - }; - let client = - RuntimeRegistrationClient::new(http_client_factory.clone(), self.registration.url()?)?; - let result = client.register(registration).await; - if self.auth_state(&auth)?.source_generation != change_state.source_generation { - // A late response must not publish an old identity after a credential switch. - *state = None; - return Err(RuntimeAuthError::Changed); - } - match result { - Ok(info) => { - let runtime = Arc::new(registration.finish(info)); - *runtime_state = RuntimeState::Ready { - runtime: runtime.clone(), - }; - Ok(RuntimeAuth { - runtime, - auth, - http_client_factory, - change_state, - }) - } - Err( - error @ RegistrationError::Rejected { - status: StatusCode::FORBIDDEN, - .. - }, - ) => { - *runtime_state = RuntimeState::Denied { - error: error.clone(), - }; - Err(error.into()) - } - Err(error) => Err(error.into()), - } - } - - /// Resolve auth for work whose routing, cache key, or payload uses a captured credential. - /// A same-source refresh is allowed; a different source must not take over that work. - /// This binds credential identity, not when the caller captured it. Restoring the exact - /// same credential can select a new runtime after logout; a different source cannot. pub async fn auth_for( &self, expected: &CodexAuth, selection: &PlatformResourceSelection, ) -> Result { - let source = self - .credentials - .auth_change_state_for(expected) - .ok_or(RuntimeAuthError::Changed)?; - self.auth_for_source(source.source_generation, selection) - .await + self.login.auth_for(expected, selection).await } - /// Select another request within a reusable client's original credential source. pub async fn auth_for_source( &self, source_generation: u64, selection: &PlatformResourceSelection, ) -> Result { - self.select_auth(selection, Some(source_generation)).await + self.login + .auth_for_source(source_generation, selection) + .await } - fn check_source(&self, expected: Option) -> Result<(), RuntimeAuthError> { - if expected.is_some_and(|expected| { - self.credentials - .auth_change_state_receiver() - .borrow() - .source_generation - != expected + /// Bind a login or command-backed credential without changing its refresh lifecycle. + pub fn bind(&self, credentials: Arc) -> Arc { + if self.login.uses_credentials(&credentials) { + return Arc::clone(&self.login); + } + let mut registrations = self + .registrations + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(binding) = registrations.iter().find_map(|entry| match entry { + CredentialRegistration::Provider(binding) if binding.uses_credentials(&credentials) => { + Some(binding) + } + _ => None, }) { - return Err(RuntimeAuthError::Changed); + return Arc::clone(binding); } - Ok(()) + let binding = Arc::new(RuntimeAuthBinding::new( + credentials, + self.login.registration.clone(), + self.login.metadata.clone(), + )); + registrations.push(CredentialRegistration::Provider(Arc::clone(&binding))); + binding } - fn auth_state(&self, auth: &CodexAuth) -> Result { - let state = self - .credentials - .auth_change_state_for(auth) - .ok_or(RuntimeAuthError::Changed)?; - if self.credentials.refresh_failure_for_auth(auth).is_some() { - return Err(RuntimeAuthError::RefreshFailed); - } - Ok(state) - } -} - -fn actor_authorization( - auth: &CodexAuth, - selection: &PlatformResourceSelection, -) -> Result { - if !matches!(auth, CodexAuth::ApiKey(_)) && selection != &PlatformResourceSelection::default() { - return Err(RuntimeAuthError::UnsupportedResourceSelection); - } - let credential = auth - .get_token() - .map_err(|_| RuntimeAuthError::Unsupported)? - .into(); - Ok(match auth { - CodexAuth::ApiKey(_) => AuthenticatedActorAuthorization::PlatformApiKey { - credential, - selection: selection.clone(), - }, - CodexAuth::Chatgpt(_) | CodexAuth::ChatgptAuthTokens(_) => { - AuthenticatedActorAuthorization::Oauth { credential } - } - CodexAuth::PersonalAccessToken(_) => { - AuthenticatedActorAuthorization::PersonalAccessToken { credential } - } - CodexAuth::Headers(_) - | CodexAuth::AgentIdentity(_) - | CodexAuth::BedrockApiKey(_) - | CodexAuth::BedrockAccessKeys(_) => return Err(RuntimeAuthError::Unsupported), - }) -} - -#[derive(Debug, thiserror::Error)] -pub enum RuntimeAuthError { - #[error("runtime registration requires an actor credential")] - Unavailable, - #[error("this credential cannot register a runtime; use a supported OpenAI sign-in")] - Unsupported, - #[error("organization and project selection requires an OpenAI API key")] - UnsupportedResourceSelection, - #[error("credentials changed during runtime registration; try again")] - Changed, - #[error("runtime registration could not refresh actor credentials")] - RefreshFailed, - #[error("failed to generate runtime signing key")] - KeyGeneration, - #[error(transparent)] - Policy(#[from] NetworkPolicyDenied), - #[error(transparent)] - Registration(RegistrationError), -} - -impl From for RuntimeAuthError { - fn from(error: RegistrationError) -> Self { - match error { - RegistrationError::Policy(denied) => Self::Policy(denied), - error => Self::Registration(error), + /// An explicit key is independent of login storage and command output, even when bytes match. + pub fn bind_api_key( + &self, + api_key: &str, + runtime: &AuthRuntimeConfig, + ) -> Arc { + let mut registrations = self + .registrations + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(binding) = registrations.iter().find_map(|entry| match entry { + CredentialRegistration::ApiKey(binding) + if binding + .credentials + .auth_cached() + .as_ref() + .and_then(CodexAuth::api_key) + == Some(api_key) + && binding.credentials.runtime_config() == *runtime => + { + Some(binding) + } + _ => None, + }) { + return Arc::clone(binding); } + let binding = Arc::new(RuntimeAuthBinding::new( + AuthManager::from_api_key(api_key, runtime.clone()), + self.login.registration.clone(), + self.login.metadata.clone(), + )); + registrations.push(CredentialRegistration::ApiKey(Arc::clone(&binding))); + binding } } diff --git a/codex-rs/runtime-agent-instance/src/manager_tests.rs b/codex-rs/runtime-agent-instance/src/manager_tests.rs index 6e4eee9b09..4608dbc2c8 100644 --- a/codex-rs/runtime-agent-instance/src/manager_tests.rs +++ b/codex-rs/runtime-agent-instance/src/manager_tests.rs @@ -8,6 +8,7 @@ use codex_features::RuntimeAuthMode; use codex_http_client::DestinationPolicy; use codex_http_client::HttpClientFactory; use codex_http_client::NetworkPolicyController; +use codex_http_client::NetworkPolicyDenied; use codex_http_client::OutboundProxyPolicy; use codex_login::AuthCredentialsStoreMode; use codex_login::AuthKeyringBackendKind; @@ -34,13 +35,19 @@ use wiremock::matchers::body_partial_json; use wiremock::matchers::method; struct Credentials { + available: std::sync::atomic::AtomicBool, current: RwLock, refreshed: CodexAuth, } impl ExternalAuth for Credentials { fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> { - Box::pin(async { Ok(self.current.read().unwrap().clone()) }) + Box::pin(async { + if !self.available.load(std::sync::atomic::Ordering::SeqCst) { + return Err(std::io::Error::other("credential source unavailable")); + } + Ok(self.current.read().unwrap().clone()) + }) } fn refresh(&self, _: ExternalAuthRefreshContext) -> ExternalAuthFuture<'_, CodexAuth> { @@ -58,6 +65,7 @@ async fn setup( ) -> (Arc, Arc, Arc) { let credentials = AuthManager::from_auth_for_testing(initial.clone()); let source = Arc::new(Credentials { + available: std::sync::atomic::AtomicBool::new(true), current: RwLock::new(initial), refreshed, }); @@ -106,6 +114,97 @@ fn oauth(session: Option<&str>, token: &str) -> CodexAuth { .unwrap() } +#[tokio::test] +async fn legacy_login_requires_sign_in_for_runtime_auth() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(registered) + .mount(&server) + .await; + let home = tempfile::tempdir().unwrap(); + let claims = json!({ + "exp": 4102444800_u64, + "session_id": "session-a", + "iss": "https://auth.openai.com", + "https://api.openai.com/auth": { + "chatgpt_user_id": "user-a", "chatgpt_account_id": "workspace-a", + }, + }); + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(claims.to_string()); + let token = format!("e30.{payload}.c2ln"); + let stored: codex_login::AuthDotJson = serde_json::from_value(json!({ + "auth_mode": "agentIdentity", + "agent_identity": { + "agent_runtime_id": "legacy-runtime", + "agent_private_key": codex_agent_identity::generate_agent_key_material() + .unwrap().private_key_pkcs8_base64, + "account_id": "workspace-a", "chatgpt_user_id": "user-a", "plan_type": "pro", + "chatgpt_account_is_fedramp": false, "task_id": "legacy-task", + }, + "tokens": { + "id_token": token, "access_token": token, + "refresh_token": "refresh", "account_id": "workspace-a", + }, + "last_refresh": "2099-01-01T00:00:00Z", + })) + .unwrap(); + codex_login::save_auth( + home.path(), + &stored, + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::Direct, + ) + .unwrap(); + let credentials = AuthManager::shared( + home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::Direct, + codex_login::test_support::transport_default_auth_route_config(), + ) + .await; + let legacy = credentials.auth_cached().unwrap(); + assert!(matches!(&legacy, CodexAuth::AgentIdentity(_))); + let runtime = RuntimeAuthManager::new( + credentials.clone(), + AgentHostAuthorization::InstallationId { + installation_id: Uuid::nil(), + }, + RuntimeRegistrationMetadata { + harness: Some("codex".into()), + harness_version: None, + role: RuntimeRole::OrchestratorAndExecutor, + }, + format!("{}/agent/register", server.uri()).parse().unwrap(), + ); + assert!(matches!( + runtime.auth().await.unwrap_err(), + RuntimeAuthError::Unsupported + )); + assert!(matches!( + runtime + .auth_for(&legacy, &PlatformResourceSelection::default()) + .await + .unwrap_err(), + RuntimeAuthError::Unsupported, + )); + + codex_login::login_with_api_key( + home.path(), + "replacement-key", + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::Direct, + ) + .unwrap(); + credentials.reload().await; + assert_eq!( + runtime.auth().await.unwrap().auth().api_key(), + Some("replacement-key") + ); +} + #[tokio::test] async fn concurrent_requests_share_registration_and_replacement_preserves_in_flight_auth() { let server = MockServer::start().await; @@ -276,6 +375,7 @@ async fn policy_recovery_preserves_pending_registration() { let auth = CodexAuth::from_api_key("key-a"); credentials .set_external_auth(Arc::new(Credentials { + available: std::sync::atomic::AtomicBool::new(true), current: RwLock::new(auth.clone()), refreshed: auth, })) @@ -708,3 +808,218 @@ async fn file_upload_runtime_auth_signs_control_requests_but_not_blob_uploads() assert_eq!(blobs[0].body, b"hello"); } } + +#[tokio::test] +async fn source_outage_preserves_registered_pending_and_denied_runtimes() { + for status in [200, 403, 503] { + let server = MockServer::start().await; + let initial = oauth(Some("login-a"), "initial"); + let initial_token = initial.get_token().unwrap(); + let attempts = std::sync::atomic::AtomicUsize::new(0); + Mock::given(method("POST")) + .respond_with(move |request: &Request| { + let attempt = attempts.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let body: Value = request.body_json().unwrap(); + if (status == 403 && body["actor_auth"]["credential"] == initial_token) + || (status == 503 && attempt == 0) + { + ResponseTemplate::new(status) + } else { + registered(request) + } + }) + .mount(&server) + .await; + let (manager, source, credentials) = setup(&server, initial.clone(), initial).await; + let first = manager.auth().await; + assert_eq!(first.is_ok(), status == 200); + let initial_source = *credentials.auth_change_state_receiver().borrow(); + source + .available + .store(false, std::sync::atomic::Ordering::SeqCst); + let cached = manager.auth().await; + assert_eq!(cached.is_ok(), status != 403); + assert_eq!( + initial_source, + *credentials.auth_change_state_receiver().borrow() + ); + let attempts_before_replacement = if status == 503 { 2 } else { 1 }; + assert_eq!( + server.received_requests().await.unwrap().len(), + attempts_before_replacement + ); + source + .available + .store(true, std::sync::atomic::Ordering::SeqCst); + let recovered = manager.auth().await; + assert_eq!(recovered.is_ok(), status != 403); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), attempts_before_replacement); + if let Ok(cached) = cached { + assert_eq!( + cached.runtime().id(), + recovered.as_ref().unwrap().runtime().id() + ); + } + if status == 200 { + assert_eq!( + first.unwrap().runtime().id(), + recovered.unwrap().runtime().id() + ); + } + if status == 503 { + assert_eq!( + requests[0].headers["idempotency-key"], + requests[1].headers["idempotency-key"] + ); + assert_eq!(requests[0].body, requests[1].body); + } + // Actual source replacement still retires the old key, including a denied registration. + *source.current.write().unwrap() = oauth(Some("login-b"), "replacement"); + assert!(manager.auth().await.is_ok()); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), attempts_before_replacement + 1); + assert_ne!( + requests[0].headers["idempotency-key"], + requests.last().unwrap().headers["idempotency-key"] + ); + } +} + +#[tokio::test] +async fn explicit_keys_reuse_their_owner_without_borrowing_a_command_source() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(registered) + .mount(&server) + .await; + let (root, _, _) = setup( + &server, + CodexAuth::from_api_key("ambient"), + CodexAuth::from_api_key("ambient"), + ) + .await; + let command_credentials = + AuthManager::from_auth_for_testing(CodexAuth::from_api_key("shared-key")); + let command = Arc::new(Credentials { + available: std::sync::atomic::AtomicBool::new(true), + current: RwLock::new(CodexAuth::from_api_key("shared-key")), + refreshed: CodexAuth::from_api_key("shared-key"), + }); + command_credentials + .set_external_auth(command.clone()) + .await + .unwrap(); + let context = command_credentials.runtime_config(); + let command_runtime = root.bind(Arc::clone(&command_credentials)); + let original_command = command_runtime.auth().await.unwrap(); + command + .available + .store(false, std::sync::atomic::Ordering::SeqCst); + + let explicit = root + .bind_api_key("shared-key", &context) + .auth() + .await + .unwrap(); + assert_eq!(explicit.auth().api_key(), Some("shared-key")); + assert_ne!(explicit.runtime().id(), original_command.runtime().id()); + let explicit_id = explicit.runtime().id().to_string(); + drop(explicit); + let reused = root + .bind_api_key("shared-key", &context) + .auth() + .await + .unwrap(); + assert_eq!(reused.runtime().id(), explicit_id); + assert_eq!( + command_runtime.auth().await.unwrap().runtime().id(), + original_command.runtime().id() + ); + + command + .available + .store(true, std::sync::atomic::Ordering::SeqCst); + let recovered = root.bind(command_credentials).auth().await.unwrap(); + assert_eq!(recovered.runtime().id(), original_command.runtime().id()); + *command.current.write().unwrap() = CodexAuth::from_api_key("replacement"); + assert_ne!( + command_runtime.auth().await.unwrap().runtime().id(), + original_command.runtime().id() + ); + assert_eq!( + root.bind_api_key("shared-key", &context) + .auth() + .await + .unwrap() + .runtime() + .id(), + explicit_id + ); + assert_ne!( + root.bind_api_key("other-key", &context) + .auth() + .await + .unwrap() + .runtime() + .id(), + explicit_id + ); + assert_eq!(server.received_requests().await.unwrap().len(), 4); +} + +#[tokio::test] +async fn explicit_key_stays_independent_when_matching_saved_login_changes() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(registered) + .expect(2) + .mount(&server) + .await; + let home = tempfile::tempdir().unwrap(); + let credentials = AuthManager::from_auth_for_testing_with_home( + CodexAuth::from_api_key("shared-key"), + home.path().to_path_buf(), + ); + let manager = Arc::new(RuntimeAuthManager::new( + credentials.clone(), + AgentHostAuthorization::InstallationId { + installation_id: Uuid::nil(), + }, + RuntimeRegistrationMetadata { + harness: None, + harness_version: None, + role: RuntimeRole::Orchestrator, + }, + format!("{}/agent/register", server.uri()).parse().unwrap(), + )); + let explicit = manager.bind_api_key("shared-key", &credentials.runtime_config()); + let initial = explicit.auth().await.unwrap(); + codex_login::login_with_api_key( + home.path(), + "replacement", + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::Direct, + ) + .unwrap(); + credentials.reload().await; + + let selected = explicit.auth().await.unwrap(); + assert_eq!(selected.auth().api_key(), Some("shared-key")); + assert_eq!(selected.runtime().id(), initial.runtime().id()); + assert_eq!( + manager.auth().await.unwrap().auth().api_key(), + Some("replacement") + ); + assert_eq!( + manager + .bind_api_key("shared-key", &credentials.runtime_config()) + .auth() + .await + .unwrap() + .runtime() + .id(), + initial.runtime().id(), + ); + server.verify().await; +}