Bind runtime auth to the provider's selected credential source (#59823)

## Why

Runtime authentication must use the credential selected by the model provider. Explicit API keys and command credentials need independent ownership so cached connections and in-flight work cannot cross credential sources.

## What changed

- Register and reuse runtimes per credential source, supporting explicit provider keys and command credentials without falling back to saved login credentials.
- Track `AuthOwner` by credential manager and owner generation across WebSocket state, compaction, and scoring.
- Preserve cached identity and runtime registration state during credential-source outages or failed storage reloads. Require a successful auth reload before logout revokes credentials.
- Reject requests whose production or staging environment differs from runtime registration, including model catalog requests.

## Testing

Add regression coverage for credential selection and runtime reuse, owner isolation, source outages, failed reloads and logout recovery, and environment mismatches.

GitOrigin-RevId: 7bab55e0ea502593393df10046309fd8502edbd0

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<AuthManager>,
-    registration: RegistrationTarget,
-    metadata: RuntimeRegistrationMetadata,
-    state: Mutex<Option<CachedRuntime>>,
+    login: Arc<RuntimeAuthBinding>,
+    registrations: std::sync::Mutex<Vec<CredentialRegistration>>,
 }
 
-enum RegistrationTarget {
-    Explicit {
-        host: AgentHostAuthorization,
-        url: Url,
-    },
-    ChatGpt {
-        installation_id: String,
-        base_url: String,
-    },
-}
-
-impl RegistrationTarget {
-    fn host(&self) -> Result<AgentHostAuthorization, RegistrationError> {
-        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<Url, RegistrationError> {
-        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<RuntimeAgentInstance> },
-    Denied { error: RegistrationError },
+enum CredentialRegistration {
+    Provider(Arc<RuntimeAuthBinding>),
+    ApiKey(Arc<RuntimeAuthBinding>),
 }
 
 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<AuthManager>) -> 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<AuthManager>,
+        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<RuntimeAuth, RuntimeAuthError> {
-        self.auth_with_selection(&PlatformResourceSelection::default())
-            .await
+    pub fn registration_url(&self) -> Result<Url, RegistrationError> {
+        self.login.registration_url()
+    }
+
+    pub async fn auth(&self) -> Result<RuntimeAuth, RuntimeAuthError> {
+        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<RuntimeAuth, RuntimeAuthError> {
-        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<u64>,
-    ) -> Result<RuntimeAuth, RuntimeAuthError> {
-        // 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<CachedRuntime>,
-        auth: CodexAuth,
-        http_client_factory: HttpClientFactory,
-        selection: &PlatformResourceSelection,
-        expected_source: Option<u64>,
-    ) -> Result<RuntimeAuth, RuntimeAuthError> {
-        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<RuntimeAuth, RuntimeAuthError> {
-        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<RuntimeAuth, RuntimeAuthError> {
-        self.select_auth(selection, Some(source_generation)).await
+        self.login
+            .auth_for_source(source_generation, selection)
+            .await
     }
 
-    fn check_source(&self, expected: Option<u64>) -> 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<AuthManager>) -> Arc<RuntimeAuthBinding> {
+        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<AuthChangeState, RuntimeAuthError> {
-        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<AuthenticatedActorAuthorization, RuntimeAuthError> {
-    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<RegistrationError> 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<RuntimeAuthBinding> {
+        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
     }
 }