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
2026-11-13 21:30:47 +00:00 · 2026-11-06 04:33:14 +00:00
parent df38cb3b2a
commit faeb379646
37 changed files with 1875 additions and 477 deletions
@@ -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<CodexAuth>,
    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<RuntimeAuthManager>, Arc<Credentials>, Arc<AuthManager>) {
    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;
+}