Integrate runtime authentication into model requests (#57171)

## What changed

- Add opt-in runtime authentication through `ModelClient::with_runtime_auth`, requiring the provider and runtime manager to share the selected OpenAI credentials.
- Keep runtime identity, credentials, routing, and transport selection consistent for each request and its retries, signing a fresh assertion for each attempt.
- Apply authentication asynchronously to WebSocket handshakes, replace stale authentication headers, and reconnect when the runtime changes.
- Reject account, organization, or project selectors that conflict with the registered runtime, and preserve gateway OAuth composition.

## Testing

Add coverage for dual-header and assertion-only authentication, HTTP retries, runtime replacement and account switches, stale WebSocket headers, selector mismatches, gateway OAuth failures, transport revocation, and Realtime compatibility.

GitOrigin-RevId: 45b9efc2c80481406a0f6f6396a8cdd21d933c12

codex-rs/runtime-agent-instance/src/manager_tests.rs

@@ -1,11 +1,9 @@
 //! Exercise shared bootstrap against HTTP and the existing external credential-provider boundary.
 
 use super::*;
-use crate::RuntimeAuthProvider;
 use crate::RuntimeRole;
 use base64::Engine;
 use codex_api::AuthError;
-use codex_api::AuthProvider;
 use codex_features::RuntimeAuthMode;
 use codex_http_client::DestinationPolicy;
 use codex_http_client::HttpClientFactory;
@@ -132,6 +130,24 @@ async fn concurrent_requests_share_registration_and_replacement_preserves_in_fli
     assert_eq!(first.auth().get_token().unwrap(), "key-a");
     assert!(!format!("{first:?}").contains("key-a"));
     assert_eq!(replacement.auth().get_token().unwrap(), "key-b");
+    assert_eq!(
+        first
+            .clone()
+            .into_auth_provider(RuntimeAuthMode::DualHeader)
+            .resolve_auth_headers()
+            .await
+            .unwrap()["authorization"],
+        "Bearer key-a"
+    );
+    assert_eq!(
+        replacement
+            .clone()
+            .into_auth_provider(RuntimeAuthMode::DualHeader)
+            .resolve_auth_headers()
+            .await
+            .unwrap()["authorization"],
+        "Bearer key-b"
+    );
     let requests = server.received_requests().await.unwrap();
     assert_eq!(requests.len(), 3);
     assert_eq!(requests[0].body, requests[1].body);
@@ -197,24 +213,28 @@ async fn policy_recovery_preserves_pending_registration() {
         },
         format!("{}/agent/register", server.uri()).parse().unwrap(),
     ));
-    let provider = RuntimeAuthProvider::new(manager.clone(), RuntimeAuthMode::DualHeader);
     assert!(matches!(
-        provider.resolve_auth_headers().await,
+        manager.auth().await.map_err(AuthError::from),
         Err(AuthError::Transient(_))
     ));
 
     policy.invalidate();
     assert!(matches!(
-        provider.resolve_auth_headers().await,
+        manager.auth().await.map_err(AuthError::from),
         Err(AuthError::Policy(NetworkPolicyDenied::Unavailable))
     ));
     assert_eq!(server.received_requests().await.unwrap().len(), 1);
 
     assert!(controller.publish(policy.revision(), DestinationPolicy::Unrestricted));
-    let headers = provider.resolve_auth_headers().await.unwrap();
+    let selected = manager.auth().await.unwrap();
+    let headers = selected
+        .clone()
+        .into_auth_provider(RuntimeAuthMode::DualHeader)
+        .resolve_auth_headers()
+        .await
+        .unwrap();
     assert_eq!(headers["authorization"], "Bearer key-a");
     assert!(headers.contains_key("x-openai-agent-assertion"));
-    let selected = manager.auth().await.unwrap();
     let requests = server.received_requests().await.unwrap();
     assert_eq!(requests.len(), 2);
     assert_eq!(requests[0].body, requests[1].body);
@@ -342,13 +362,24 @@ async fn request_auth_replaces_stale_headers_and_signs_each_attempt() {
                 .mount(&server)
                 .await;
             let (manager, _, _) = setup(&server, initial.clone(), initial.clone()).await;
-            let provider = RuntimeAuthProvider::new(manager.clone(), mode);
+            let selected = manager.auth().await.unwrap();
+            let provider = selected.clone().into_auth_provider(mode);
+            assert!(
+                codex_api::auth_header_telemetry(provider.as_ref())
+                    .unwrap()
+                    .attached
+            );
             let mut previous_assertion = None;
             for _ in 0..2 {
                 let mut request = codex_http_client::Request::new(
                     Method::POST,
                     format!("{}/responses", server.uri()),
                 );
+                request
+                    .headers
+                    .insert("openai-project", HeaderValue::from_static("project-a"));
+                let mut sync_headers = request.headers.clone();
+                provider.add_auth_headers(&mut sync_headers).unwrap();
                 for name in [
                     "authorization",
                     "x-openai-agent-assertion",
@@ -359,49 +390,47 @@ async fn request_auth_replaces_stale_headers_and_signs_each_attempt() {
                         .headers
                         .insert(name, HeaderValue::from_static("stale"));
                 }
-                request
-                    .headers
-                    .insert("openai-project", HeaderValue::from_static("project-a"));
-                let mut headers = provider.apply_auth(request).await.unwrap().headers;
-                let mut expected = HeaderMap::new();
-                expected.insert("openai-project", HeaderValue::from_static("project-a"));
-                if initial.is_chatgpt_auth() {
-                    expected.insert(
-                        "chatgpt-account-id",
-                        HeaderValue::from_static("workspace-a"),
-                    );
+                for mut headers in [
+                    sync_headers,
+                    provider.apply_auth(request).await.unwrap().headers,
+                ] {
+                    let mut expected = HeaderMap::new();
+                    expected.insert("openai-project", HeaderValue::from_static("project-a"));
+                    if initial.is_chatgpt_auth() {
+                        expected.insert(
+                            "chatgpt-account-id",
+                            HeaderValue::from_static("workspace-a"),
+                        );
+                    }
+                    let assertion = match mode {
+                        RuntimeAuthMode::DualHeader => {
+                            let bearer = headers.remove(AUTHORIZATION).unwrap();
+                            assert!(bearer.is_sensitive());
+                            assert_eq!(bearer, format!("Bearer {}", initial.get_token().unwrap()));
+                            headers.remove("x-openai-agent-assertion").unwrap()
+                        }
+                        RuntimeAuthMode::AssertionOnly => headers.remove(AUTHORIZATION).unwrap(),
+                    };
+                    assert!(assertion.is_sensitive());
+                    assert_eq!(headers, expected);
+                    assert_ne!(previous_assertion.as_ref(), Some(&assertion));
+                    let assertion_text = assertion.to_str().unwrap();
+                    let assertion_text = match mode {
+                        RuntimeAuthMode::DualHeader => assertion_text,
+                        RuntimeAuthMode::AssertionOnly => {
+                            assertion_text.strip_prefix("AgentAssertion ").unwrap()
+                        }
+                    };
+                    let segments: Vec<_> = assertion_text.split('.').collect();
+                    let claims: Value = serde_json::from_slice(
+                        &base64::engine::general_purpose::URL_SAFE_NO_PAD
+                            .decode(segments[2])
+                            .unwrap(),
+                    )
+                    .unwrap();
+                    assert_eq!(claims["runtime_agent_instance"], selected.runtime().id());
+                    previous_assertion = Some(assertion);
                 }
-                let assertion = match mode {
-                    RuntimeAuthMode::DualHeader => {
-                        let bearer = headers.remove(AUTHORIZATION).unwrap();
-                        assert!(bearer.is_sensitive());
-                        assert_eq!(bearer, format!("Bearer {}", initial.get_token().unwrap()));
-                        headers.remove("x-openai-agent-assertion").unwrap()
-                    }
-                    RuntimeAuthMode::AssertionOnly => headers.remove(AUTHORIZATION).unwrap(),
-                };
-                assert!(assertion.is_sensitive());
-                assert_eq!(headers, expected);
-                assert_ne!(previous_assertion.as_ref(), Some(&assertion));
-                let assertion_text = assertion.to_str().unwrap();
-                let assertion_text = match mode {
-                    RuntimeAuthMode::DualHeader => assertion_text,
-                    RuntimeAuthMode::AssertionOnly => {
-                        assertion_text.strip_prefix("AgentAssertion ").unwrap()
-                    }
-                };
-                let segments: Vec<_> = assertion_text.split('.').collect();
-                let claims: Value = serde_json::from_slice(
-                    &base64::engine::general_purpose::URL_SAFE_NO_PAD
-                        .decode(segments[2])
-                        .unwrap(),
-                )
-                .unwrap();
-                assert_eq!(
-                    claims["runtime_agent_instance"],
-                    manager.auth().await.unwrap().runtime().id()
-                );
-                previous_assertion = Some(assertion);
             }
             assert_eq!(server.received_requests().await.unwrap().len(), 1);
         }
@@ -409,7 +438,7 @@ async fn request_auth_replaces_stale_headers_and_signs_each_attempt() {
 }
 
 #[tokio::test]
-async fn header_auth_uses_refreshed_credentials_without_replacing_the_runtime() {
+async fn new_request_uses_refreshed_credentials_without_replacing_the_runtime() {
     let server = MockServer::start().await;
     Mock::given(method("POST"))
         .respond_with(registered)
@@ -417,26 +446,35 @@ async fn header_auth_uses_refreshed_credentials_without_replacing_the_runtime()
         .await;
     let initial = oauth(Some("login-a"), "old");
     let (manager, source, _) = setup(&server, initial.clone(), initial).await;
-    let provider = RuntimeAuthProvider::new(manager.clone(), RuntimeAuthMode::DualHeader);
-    let first = provider.resolve_auth_headers().await.unwrap();
     let runtime = manager.auth().await.unwrap();
+    let provider = runtime
+        .clone()
+        .into_auth_provider(RuntimeAuthMode::DualHeader);
+    let first = provider.resolve_auth_headers().await.unwrap();
     let refreshed = oauth(Some("login-a"), "new");
     *source.current.write().unwrap() = refreshed.clone();
-    let next = provider.resolve_auth_headers().await.unwrap();
+    let next_runtime = manager.auth().await.unwrap();
+    let next = next_runtime
+        .clone()
+        .into_auth_provider(RuntimeAuthMode::DualHeader)
+        .resolve_auth_headers()
+        .await
+        .unwrap();
     assert_eq!(
         next[AUTHORIZATION],
         format!("Bearer {}", refreshed.get_token().unwrap())
     );
     assert_ne!(first[AUTHORIZATION], next[AUTHORIZATION]);
-    assert!(Arc::ptr_eq(
-        runtime.runtime(),
-        manager.auth().await.unwrap().runtime()
-    ));
+    assert_eq!(
+        provider.resolve_auth_headers().await.unwrap()[AUTHORIZATION],
+        first[AUTHORIZATION]
+    );
+    assert!(Arc::ptr_eq(runtime.runtime(), next_runtime.runtime()));
     assert_eq!(server.received_requests().await.unwrap().len(), 1);
 }
 
 #[tokio::test]
-async fn registration_failure_never_returns_credential_only_auth() {
+async fn registration_failure_reports_retryability() {
     for status in [403, 404, 429, 503] {
         let server = MockServer::start().await;
         Mock::given(method("POST"))
@@ -445,8 +483,7 @@ async fn registration_failure_never_returns_credential_only_auth() {
             .await;
         let initial = CodexAuth::from_api_key("key-a");
         let (manager, _, _) = setup(&server, initial.clone(), initial).await;
-        let provider = RuntimeAuthProvider::new(manager, RuntimeAuthMode::DualHeader);
-        let error = provider.resolve_auth_headers().await.unwrap_err();
+        let error = AuthError::from(manager.auth().await.unwrap_err());
         assert_eq!(
             matches!(error, AuthError::Transient(_)),
             status == 429 || status == 503