Use AgentAssertion downstream behind use_agent_identity

codex-rs/codex-api/src/api_bridge.rs

@@ -178,13 +178,33 @@ struct UsageErrorBody {
 pub struct CoreAuthProvider {
     pub token: Option<String>,
     pub account_id: Option<String>,
+    authorization_header_override: Option<String>,
 }
 
 impl CoreAuthProvider {
+    pub fn from_bearer_token(token: Option<String>, account_id: Option<String>) -> Self {
+        Self {
+            token,
+            account_id,
+            authorization_header_override: None,
+        }
+    }
+
+    pub fn from_authorization_header_value(
+        authorization_header_value: Option<String>,
+        account_id: Option<String>,
+    ) -> Self {
+        Self {
+            token: None,
+            account_id,
+            authorization_header_override: authorization_header_value,
+        }
+    }
+
     pub fn auth_header_attached(&self) -> bool {
-        self.token
+        self.authorization_header_value()
             .as_ref()
-            .is_some_and(|token| http::HeaderValue::from_str(&format!("Bearer {token}")).is_ok())
+            .is_some_and(|value| http::HeaderValue::from_str(value).is_ok())
     }
 
     pub fn auth_header_name(&self) -> Option<&'static str> {
@@ -195,8 +215,20 @@ impl CoreAuthProvider {
         Self {
             token: token.map(str::to_string),
             account_id: account_id.map(str::to_string),
+            authorization_header_override: None,
         }
     }
+
+    #[cfg(test)]
+    pub fn for_test_authorization_header(
+        authorization_header_value: Option<&str>,
+        account_id: Option<&str>,
+    ) -> Self {
+        Self::from_authorization_header_value(
+            authorization_header_value.map(str::to_string),
+            account_id.map(str::to_string),
+        )
+    }
 }
 
 impl ApiAuthProvider for CoreAuthProvider {
@@ -204,6 +236,12 @@ impl ApiAuthProvider for CoreAuthProvider {
         self.token.clone()
     }
 
+    fn authorization_header_value(&self) -> Option<String> {
+        self.authorization_header_override
+            .clone()
+            .or_else(|| self.bearer_token().map(|token| format!("Bearer {token}")))
+    }
+
     fn account_id(&self) -> Option<String> {
         self.account_id.clone()
     }

codex-rs/codex-api/src/api_bridge_tests.rs

@@ -133,11 +133,22 @@ fn map_api_error_extracts_identity_auth_details_from_headers() {
 
 #[test]
 fn core_auth_provider_reports_when_auth_header_will_attach() {
-    let auth = CoreAuthProvider {
-        token: Some("access-token".to_string()),
-        account_id: None,
-    };
+    let auth = CoreAuthProvider::from_bearer_token(Some("access-token".to_string()), None);
 
     assert!(auth.auth_header_attached());
     assert_eq!(auth.auth_header_name(), Some("authorization"));
 }
+
+#[test]
+fn core_auth_provider_supports_non_bearer_authorization_headers() {
+    let auth =
+        CoreAuthProvider::for_test_authorization_header(Some("AgentAssertion opaque-token"), None);
+
+    assert!(auth.auth_header_attached());
+    assert_eq!(auth.auth_header_name(), Some("authorization"));
+    assert_eq!(auth.bearer_token(), None);
+    assert_eq!(
+        auth.authorization_header_value(),
+        Some("AgentAssertion opaque-token".to_string())
+    );
+}

codex-rs/codex-api/src/auth.rs

@@ -9,14 +9,17 @@ use http::HeaderValue;
 /// reach this interface.
 pub trait AuthProvider: Send + Sync {
     fn bearer_token(&self) -> Option<String>;
+    fn authorization_header_value(&self) -> Option<String> {
+        self.bearer_token().map(|token| format!("Bearer {token}"))
+    }
     fn account_id(&self) -> Option<String> {
         None
     }
 }
 
 pub(crate) fn add_auth_headers_to_header_map<A: AuthProvider>(auth: &A, headers: &mut HeaderMap) {
-    if let Some(token) = auth.bearer_token()
-        && let Ok(header) = HeaderValue::from_str(&format!("Bearer {token}"))
+    if let Some(authorization) = auth.authorization_header_value()
+        && let Ok(header) = HeaderValue::from_str(&authorization)
     {
         let _ = headers.insert(http::header::AUTHORIZATION, header);
     }

codex-rs/core/src/agent_identity.rs

@@ -27,12 +27,15 @@ use tracing::debug;
 use tracing::info;
 use tracing::warn;
 
+use crate::config::Config;
+
+mod assertion;
 mod task_registration;
 
+#[cfg(test)]
+pub(crate) use assertion::AgentAssertionEnvelope;
 pub(crate) use task_registration::RegisteredAgentTask;
 
-use crate::config::Config;
-
 const AGENT_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15);
 const AGENT_IDENTITY_BISCUIT_TIMEOUT: Duration = Duration::from_secs(15);
 
@@ -334,7 +337,7 @@ impl AgentIdentityManager {
     }
 
     #[cfg(test)]
-    fn new_for_tests(
+    pub(crate) fn new_for_tests(
         auth_manager: Arc<AuthManager>,
         feature_enabled: bool,
         chatgpt_base_url: String,
@@ -348,6 +351,30 @@ impl AgentIdentityManager {
             ensure_lock: Arc::new(Mutex::new(())),
         }
     }
+
+    #[cfg(test)]
+    pub(crate) async fn seed_generated_identity_for_tests(
+        &self,
+        agent_runtime_id: &str,
+    ) -> Result<StoredAgentIdentity> {
+        let (auth, binding) = self
+            .current_auth_binding()
+            .await
+            .context("test agent identity requires ChatGPT auth")?;
+        let key_material = generate_agent_key_material()?;
+        let stored_identity = StoredAgentIdentity {
+            binding_id: binding.binding_id.clone(),
+            chatgpt_account_id: binding.chatgpt_account_id.clone(),
+            chatgpt_user_id: binding.chatgpt_user_id.clone(),
+            agent_runtime_id: agent_runtime_id.to_string(),
+            private_key_pkcs8_base64: key_material.private_key_pkcs8_base64,
+            public_key_ssh: key_material.public_key_ssh,
+            registered_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true),
+            abom: self.abom.clone(),
+        };
+        self.store_identity(&auth, &stored_identity)?;
+        Ok(stored_identity)
+    }
 }
 
 impl StoredAgentIdentity {
@@ -577,7 +604,7 @@ mod tests {
             .and(path("/v1/agent/register"))
             .and(header("x-openai-authorization", "human-biscuit"))
             .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
-                "agent_runtime_id": "agent_123",
+                "agent_runtime_id": "agent-123",
             })))
             .expect(1)
             .mount(&server)
@@ -603,7 +630,7 @@ mod tests {
             .unwrap()
             .expect("identity should be reused");
 
-        assert_eq!(first.agent_runtime_id, "agent_123");
+        assert_eq!(first.agent_runtime_id, "agent-123");
         assert_eq!(first, second);
         assert_eq!(first.abom.agent_harness_id, "codex-cli");
         assert_eq!(first.chatgpt_account_id, "account-123");
@@ -619,7 +646,7 @@ mod tests {
             .and(path("/v1/agent/register"))
             .and(header("x-openai-authorization", "human-biscuit"))
             .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
-                "agent_runtime_id": "agent_456",
+                "agent_runtime_id": "agent-456",
             })))
             .expect(1)
             .mount(&server)
@@ -650,11 +677,11 @@ mod tests {
             .unwrap()
             .expect("identity should be registered");
 
-        assert_eq!(stored.agent_runtime_id, "agent_456");
+        assert_eq!(stored.agent_runtime_id, "agent-456");
         let persisted = auth
             .get_agent_identity(&binding.chatgpt_account_id)
             .expect("stored identity");
-        assert_eq!(persisted.agent_runtime_id, "agent_456");
+        assert_eq!(persisted.agent_runtime_id, "agent-456");
     }
 
     #[tokio::test]

codex-rs/core/src/agent_identity/assertion.rs

@@ -0,0 +1,182 @@
+use std::collections::BTreeMap;
+
+use anyhow::Context;
+use anyhow::Result;
+use base64::Engine as _;
+use base64::engine::general_purpose::URL_SAFE_NO_PAD;
+use ed25519_dalek::Signer as _;
+use serde::Deserialize;
+use serde::Serialize;
+use tracing::debug;
+
+use super::*;
+
+#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
+pub(crate) struct AgentAssertionEnvelope {
+    pub(crate) agent_runtime_id: String,
+    pub(crate) task_id: String,
+    pub(crate) timestamp: String,
+    pub(crate) signature: String,
+}
+
+impl AgentIdentityManager {
+    pub(crate) async fn authorization_header_for_task(
+        &self,
+        agent_task: &RegisteredAgentTask,
+    ) -> Result<Option<String>> {
+        if !self.feature_enabled {
+            return Ok(None);
+        }
+
+        let Some(stored_identity) = self.ensure_registered_identity().await? else {
+            return Ok(None);
+        };
+        anyhow::ensure!(
+            stored_identity.agent_runtime_id == agent_task.agent_runtime_id,
+            "agent task runtime {} does not match stored agent identity {}",
+            agent_task.agent_runtime_id,
+            stored_identity.agent_runtime_id
+        );
+
+        let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true);
+        let envelope = AgentAssertionEnvelope {
+            agent_runtime_id: agent_task.agent_runtime_id.clone(),
+            task_id: agent_task.task_id.clone(),
+            timestamp: timestamp.clone(),
+            signature: sign_agent_assertion_payload(&stored_identity, agent_task, &timestamp)?,
+        };
+        let serialized_assertion = serialize_agent_assertion(&envelope)?;
+        debug!(
+            agent_runtime_id = %envelope.agent_runtime_id,
+            task_id = %envelope.task_id,
+            "attaching agent assertion authorization to downstream request"
+        );
+        Ok(Some(format!("AgentAssertion {serialized_assertion}")))
+    }
+}
+
+fn sign_agent_assertion_payload(
+    stored_identity: &StoredAgentIdentity,
+    agent_task: &RegisteredAgentTask,
+    timestamp: &str,
+) -> Result<String> {
+    let signing_key = stored_identity.signing_key()?;
+    let payload = format!(
+        "{}:{}:{timestamp}",
+        agent_task.agent_runtime_id, agent_task.task_id
+    );
+    Ok(BASE64_STANDARD.encode(signing_key.sign(payload.as_bytes()).to_bytes()))
+}
+
+fn serialize_agent_assertion(envelope: &AgentAssertionEnvelope) -> Result<String> {
+    let payload = serde_json::to_vec(&BTreeMap::from([
+        ("agent_runtime_id", envelope.agent_runtime_id.as_str()),
+        ("signature", envelope.signature.as_str()),
+        ("task_id", envelope.task_id.as_str()),
+        ("timestamp", envelope.timestamp.as_str()),
+    ]))
+    .context("failed to serialize agent assertion envelope")?;
+    Ok(URL_SAFE_NO_PAD.encode(payload))
+}
+
+#[cfg(test)]
+mod tests {
+    use base64::engine::general_purpose::URL_SAFE_NO_PAD;
+    use ed25519_dalek::Signature;
+    use ed25519_dalek::Verifier as _;
+    use pretty_assertions::assert_eq;
+
+    use super::*;
+
+    #[tokio::test]
+    async fn authorization_header_for_task_skips_when_feature_is_disabled() {
+        let auth_manager =
+            AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
+        let manager = AgentIdentityManager::new_for_tests(
+            auth_manager,
+            /*feature_enabled*/ false,
+            "https://chatgpt.com/backend-api/".to_string(),
+            SessionSource::Cli,
+        );
+        let agent_task = RegisteredAgentTask {
+            binding_id: "chatgpt-account-account_id".to_string(),
+            chatgpt_account_id: "account_id".to_string(),
+            chatgpt_user_id: None,
+            agent_runtime_id: "agent-123".to_string(),
+            task_id: "task-123".to_string(),
+            registered_at: "2026-03-23T12:00:00Z".to_string(),
+        };
+
+        assert_eq!(
+            manager
+                .authorization_header_for_task(&agent_task)
+                .await
+                .unwrap(),
+            None
+        );
+    }
+
+    #[tokio::test]
+    async fn authorization_header_for_task_serializes_signed_agent_assertion() {
+        let auth_manager =
+            AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
+        let manager = AgentIdentityManager::new_for_tests(
+            auth_manager,
+            /*feature_enabled*/ true,
+            "https://chatgpt.com/backend-api/".to_string(),
+            SessionSource::Cli,
+        );
+        let stored_identity = manager
+            .seed_generated_identity_for_tests("agent-123")
+            .await
+            .expect("seed test identity");
+        let agent_task = RegisteredAgentTask {
+            binding_id: "chatgpt-account-account_id".to_string(),
+            chatgpt_account_id: "account_id".to_string(),
+            chatgpt_user_id: None,
+            agent_runtime_id: "agent-123".to_string(),
+            task_id: "task-123".to_string(),
+            registered_at: "2026-03-23T12:00:00Z".to_string(),
+        };
+
+        let header = manager
+            .authorization_header_for_task(&agent_task)
+            .await
+            .expect("build agent assertion")
+            .expect("header should exist");
+        let token = header
+            .strip_prefix("AgentAssertion ")
+            .expect("agent assertion scheme");
+        let payload = URL_SAFE_NO_PAD
+            .decode(token)
+            .expect("valid base64url payload");
+        let envelope: AgentAssertionEnvelope =
+            serde_json::from_slice(&payload).expect("valid assertion envelope");
+
+        assert_eq!(
+            envelope,
+            AgentAssertionEnvelope {
+                agent_runtime_id: "agent-123".to_string(),
+                task_id: "task-123".to_string(),
+                timestamp: envelope.timestamp.clone(),
+                signature: envelope.signature.clone(),
+            }
+        );
+        let signature_bytes = BASE64_STANDARD
+            .decode(&envelope.signature)
+            .expect("valid base64 signature");
+        let signature = Signature::from_slice(&signature_bytes).expect("valid signature bytes");
+        let signing_key = stored_identity.signing_key().expect("signing key");
+        signing_key
+            .verifying_key()
+            .verify(
+                format!(
+                    "{}:{}:{}",
+                    envelope.agent_runtime_id, envelope.task_id, envelope.timestamp
+                )
+                .as_bytes(),
+                &signature,
+            )
+            .expect("signature should verify");
+    }
+}

codex-rs/core/src/client.rs

@@ -31,6 +31,8 @@ use std::sync::atomic::AtomicBool;
 use std::sync::atomic::AtomicU64;
 use std::sync::atomic::Ordering;
 
+use crate::agent_identity::AgentIdentityManager;
+use crate::agent_identity::RegisteredAgentTask;
 use codex_api::ApiError;
 use codex_api::CompactClient as ApiCompactClient;
 use codex_api::CompactionInput as ApiCompactionInput;
@@ -92,6 +94,7 @@ use tokio::sync::oneshot;
 use tokio::sync::oneshot::error::TryRecvError;
 use tokio_tungstenite::tungstenite::Error;
 use tokio_tungstenite::tungstenite::Message;
+use tracing::debug;
 use tracing::instrument;
 use tracing::trace;
 use tracing::warn;
@@ -144,6 +147,7 @@ pub(crate) const WEBSOCKET_CONNECT_TIMEOUT: Duration =
 #[derive(Debug)]
 struct ModelClientState {
     auth_manager: Option<Arc<AuthManager>>,
+    agent_identity_manager: Option<Arc<AgentIdentityManager>>,
     conversation_id: ThreadId,
     window_generation: AtomicU64,
     installation_id: String,
@@ -211,6 +215,8 @@ pub struct ModelClient {
 pub struct ModelClientSession {
     client: ModelClient,
     websocket_session: WebsocketSession,
+    agent_task: Option<RegisteredAgentTask>,
+    cache_websocket_session_on_drop: bool,
     /// Turn state for sticky routing.
     ///
     /// This is an `OnceLock` that stores the turn state value received from the server
@@ -306,6 +312,33 @@ impl ModelClient {
         enable_request_compression: bool,
         include_timing_metrics: bool,
         beta_features_header: Option<String>,
+    ) -> Self {
+        Self::new_with_agent_identity_manager(
+            auth_manager,
+            /*agent_identity_manager*/ None,
+            conversation_id,
+            installation_id,
+            provider,
+            session_source,
+            model_verbosity,
+            enable_request_compression,
+            include_timing_metrics,
+            beta_features_header,
+        )
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub(crate) fn new_with_agent_identity_manager(
+        auth_manager: Option<Arc<AuthManager>>,
+        agent_identity_manager: Option<Arc<AgentIdentityManager>>,
+        conversation_id: ThreadId,
+        installation_id: String,
+        provider: ModelProviderInfo,
+        session_source: SessionSource,
+        model_verbosity: Option<VerbosityConfig>,
+        enable_request_compression: bool,
+        include_timing_metrics: bool,
+        beta_features_header: Option<String>,
     ) -> Self {
         let auth_manager = auth_manager_for_provider(auth_manager, &provider);
         let codex_api_key_env_enabled = auth_manager
@@ -315,6 +348,7 @@ impl ModelClient {
         Self {
             state: Arc::new(ModelClientState {
                 auth_manager,
+                agent_identity_manager,
                 conversation_id,
                 window_generation: AtomicU64::new(0),
                 installation_id,
@@ -336,9 +370,25 @@ 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 {
+        self.new_session_with_agent_task(None)
+    }
+
+    pub(crate) fn new_session_with_agent_task(
+        &self,
+        agent_task: Option<RegisteredAgentTask>,
+    ) -> ModelClientSession {
+        let cache_websocket_session_on_drop = agent_task.is_none();
+        let websocket_session = if agent_task.is_some() {
+            drop(self.take_cached_websocket_session());
+            WebsocketSession::default()
+        } else {
+            self.take_cached_websocket_session()
+        };
         ModelClientSession {
             client: self.clone(),
-            websocket_session: self.take_cached_websocket_session(),
+            websocket_session,
+            agent_task,
+            cache_websocket_session_on_drop,
             turn_state: Arc::new(OnceLock::new()),
         }
     }
@@ -421,7 +471,7 @@ impl ModelClient {
         if prompt.input.is_empty() {
             return Ok(Vec::new());
         }
-        let client_setup = self.current_client_setup().await?;
+        let client_setup = self.current_client_setup(None).await?;
         let transport = ReqwestTransport::new(build_reqwest_client());
         let request_telemetry = Self::build_request_telemetry(
             session_telemetry,
@@ -485,7 +535,7 @@ impl ModelClient {
     ) -> Result<RealtimeWebrtcCallStart> {
         // Create the media call over HTTP first, then retain matching auth so realtime can attach
         // the server-side control WebSocket to the call id from that HTTP response.
-        let client_setup = self.current_client_setup().await?;
+        let client_setup = self.current_client_setup(/*agent_task*/ None).await?;
         let mut sideband_headers = extra_headers.clone();
         sideband_headers.extend(sideband_websocket_auth_headers(&client_setup.api_auth));
         let transport = ReqwestTransport::new(build_reqwest_client());
@@ -518,7 +568,7 @@ impl ModelClient {
             return Ok(Vec::new());
         }
 
-        let client_setup = self.current_client_setup().await?;
+        let client_setup = self.current_client_setup(None).await?;
         let transport = ReqwestTransport::new(build_reqwest_client());
         let request_telemetry = Self::build_request_telemetry(
             session_telemetry,
@@ -659,7 +709,10 @@ impl ModelClient {
     ///
     /// This centralizes setup used by both prewarm and normal request paths so they stay in
     /// lockstep when auth/provider resolution changes.
-    async fn current_client_setup(&self) -> Result<CurrentClientSetup> {
+    async fn current_client_setup(
+        &self,
+        agent_task: Option<&RegisteredAgentTask>,
+    ) -> Result<CurrentClientSetup> {
         let auth = match self.state.auth_manager.as_ref() {
             Some(manager) => manager.auth().await,
             None => None,
@@ -668,7 +721,33 @@ impl ModelClient {
             .state
             .provider
             .to_api_provider(auth.as_ref().map(CodexAuth::auth_mode))?;
-        let api_auth = auth_provider_from_auth(auth.clone(), &self.state.provider)?;
+        let api_auth = match (agent_task, self.state.agent_identity_manager.as_ref()) {
+            (Some(agent_task), Some(agent_identity_manager)) => {
+                if let Some(authorization_header_value) = agent_identity_manager
+                    .authorization_header_for_task(agent_task)
+                    .await
+                    .map_err(|err| {
+                        CodexErr::Stream(
+                            format!("failed to build agent assertion authorization: {err}"),
+                            None,
+                        )
+                    })?
+                {
+                    debug!(
+                        agent_runtime_id = %agent_task.agent_runtime_id,
+                        task_id = %agent_task.task_id,
+                        "using agent assertion authorization for downstream request"
+                    );
+                    CoreAuthProvider::from_authorization_header_value(
+                        Some(authorization_header_value),
+                        None,
+                    )
+                } else {
+                    auth_provider_from_auth(auth.clone(), &self.state.provider)?
+                }
+            }
+            _ => auth_provider_from_auth(auth.clone(), &self.state.provider)?,
+        };
         Ok(CurrentClientSetup {
             auth,
             api_provider,
@@ -802,12 +881,18 @@ impl ModelClient {
 impl Drop for ModelClientSession {
     fn drop(&mut self) {
         let websocket_session = std::mem::take(&mut self.websocket_session);
-        self.client
-            .store_cached_websocket_session(websocket_session);
+        if self.cache_websocket_session_on_drop {
+            self.client
+                .store_cached_websocket_session(websocket_session);
+        }
     }
 }
 
 impl ModelClientSession {
+    pub(crate) fn disable_cached_websocket_session_on_drop(&mut self) {
+        self.cache_websocket_session_on_drop = false;
+    }
+
     pub(crate) fn reset_websocket_session(&mut self) {
         self.websocket_session.connection = None;
         self.websocket_session.last_request = None;
@@ -1009,11 +1094,15 @@ impl ModelClientSession {
             return Ok(());
         }
 
-        let client_setup = self.client.current_client_setup().await.map_err(|err| {
-            ApiError::Stream(format!(
-                "failed to build websocket prewarm client setup: {err}"
-            ))
-        })?;
+        let client_setup = self
+            .client
+            .current_client_setup(self.agent_task.as_ref())
+            .await
+            .map_err(|err| {
+                ApiError::Stream(format!(
+                    "failed to build websocket prewarm client setup: {err}"
+                ))
+            })?;
         let auth_context = AuthRequestTelemetryContext::new(
             client_setup.auth.as_ref().map(CodexAuth::auth_mode),
             &client_setup.api_auth,
@@ -1167,7 +1256,10 @@ impl ModelClientSession {
             .map(AuthManager::unauthorized_recovery);
         let mut pending_retry = PendingUnauthorizedRetry::default();
         loop {
-            let client_setup = self.client.current_client_setup().await?;
+            let client_setup = self
+                .client
+                .current_client_setup(self.agent_task.as_ref())
+                .await?;
             let transport = ReqwestTransport::new(build_reqwest_client());
             let request_auth_context = AuthRequestTelemetryContext::new(
                 client_setup.auth.as_ref().map(CodexAuth::auth_mode),
@@ -1256,7 +1348,10 @@ impl ModelClientSession {
             .map(AuthManager::unauthorized_recovery);
         let mut pending_retry = PendingUnauthorizedRetry::default();
         loop {
-            let client_setup = self.client.current_client_setup().await?;
+            let client_setup = self
+                .client
+                .current_client_setup(self.agent_task.as_ref())
+                .await?;
             let request_auth_context = AuthRequestTelemetryContext::new(
                 client_setup.auth.as_ref().map(CodexAuth::auth_mode),
                 &client_setup.api_auth,

codex-rs/core/src/client_tests.rs

@@ -1,3 +1,5 @@
+use std::sync::Arc;
+
 use super::AuthRequestTelemetryContext;
 use super::ModelClient;
 use super::PendingUnauthorizedRetry;
@@ -7,17 +9,39 @@ use super::X_CODEX_PARENT_THREAD_ID_HEADER;
 use super::X_CODEX_TURN_METADATA_HEADER;
 use super::X_CODEX_WINDOW_ID_HEADER;
 use super::X_OPENAI_SUBAGENT_HEADER;
+use crate::Prompt;
+use crate::ResponseEvent;
+use crate::agent_identity::AgentAssertionEnvelope;
+use crate::agent_identity::AgentIdentityManager;
+use crate::agent_identity::RegisteredAgentTask;
+use crate::agent_identity::StoredAgentIdentity;
+use base64::Engine as _;
+use base64::engine::general_purpose::URL_SAFE_NO_PAD;
 use codex_api::CoreAuthProvider;
 use codex_app_server_protocol::AuthMode;
+use codex_keyring_store::tests::MockKeyringStore;
+use codex_login::AuthManager;
+use codex_login::CodexAuth;
+use codex_model_provider_info::ModelProviderInfo;
 use codex_model_provider_info::WireApi;
 use codex_model_provider_info::create_oss_provider_with_base_url;
 use codex_otel::SessionTelemetry;
 use codex_protocol::ThreadId;
+use codex_protocol::config_types::ReasoningSummary;
+use codex_protocol::models::ContentItem;
+use codex_protocol::models::ResponseItem;
 use codex_protocol::openai_models::ModelInfo;
 use codex_protocol::protocol::SessionSource;
 use codex_protocol::protocol::SubAgentSource;
+use codex_secrets::SecretsBackendKind;
+use codex_secrets::SecretsManager;
+use core_test_support::responses;
+use ed25519_dalek::Signature;
+use ed25519_dalek::Verifier as _;
+use futures::StreamExt;
 use pretty_assertions::assert_eq;
 use serde_json::json;
+use tempfile::TempDir;
 
 fn test_model_client(session_source: SessionSource) -> ModelClient {
     let provider = create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses);
@@ -79,6 +103,118 @@ fn test_session_telemetry() -> SessionTelemetry {
     )
 }
 
+fn test_prompt(text: &str) -> Prompt {
+    Prompt {
+        input: vec![ResponseItem::Message {
+            id: None,
+            role: "user".into(),
+            content: vec![ContentItem::InputText {
+                text: text.to_string(),
+            }],
+            end_turn: None,
+            phase: None,
+        }],
+        ..Prompt::default()
+    }
+}
+
+async fn drain_stream_to_completion(stream: &mut crate::ResponseStream) -> anyhow::Result<()> {
+    while let Some(event) = stream.next().await {
+        if matches!(event?, ResponseEvent::Completed { .. }) {
+            break;
+        }
+    }
+    Ok(())
+}
+
+async fn model_client_with_agent_task(
+    provider: ModelProviderInfo,
+) -> (
+    TempDir,
+    ModelClient,
+    RegisteredAgentTask,
+    StoredAgentIdentity,
+) {
+    let codex_home = tempfile::tempdir().expect("tempdir");
+    let keyring_store = Arc::new(MockKeyringStore::default());
+    let secrets_manager = SecretsManager::new_with_keyring_store(
+        codex_home.path().to_path_buf(),
+        SecretsBackendKind::Local,
+        keyring_store,
+    );
+    let auth_manager =
+        AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing());
+    let agent_identity_manager = Arc::new(AgentIdentityManager::new_for_tests(
+        Arc::clone(&auth_manager),
+        /*feature_enabled*/ true,
+        "https://chatgpt.com/backend-api/".to_string(),
+        SessionSource::Cli,
+        secrets_manager,
+    ));
+    let stored_identity = agent_identity_manager
+        .seed_generated_identity_for_tests("agent-123")
+        .await
+        .expect("seed test identity");
+    let agent_task = RegisteredAgentTask {
+        agent_runtime_id: stored_identity.agent_runtime_id.clone(),
+        task_id: "task-123".to_string(),
+        registered_at: "2026-03-23T12:00:00Z".to_string(),
+    };
+    let client = ModelClient::new_with_agent_identity_manager(
+        Some(auth_manager),
+        Some(agent_identity_manager),
+        ThreadId::new(),
+        /*installation_id*/ "11111111-1111-4111-8111-111111111111".to_string(),
+        provider,
+        SessionSource::Cli,
+        None,
+        false,
+        false,
+        None,
+    );
+    (codex_home, client, agent_task, stored_identity)
+}
+
+fn assert_agent_assertion_header(
+    authorization_header: &str,
+    stored_identity: &StoredAgentIdentity,
+    expected_agent_runtime_id: &str,
+    expected_task_id: &str,
+) {
+    let token = authorization_header
+        .strip_prefix("AgentAssertion ")
+        .expect("agent assertion authorization scheme");
+    let envelope: AgentAssertionEnvelope = serde_json::from_slice(
+        &URL_SAFE_NO_PAD
+            .decode(token)
+            .expect("base64url-encoded agent assertion"),
+    )
+    .expect("valid agent assertion envelope");
+
+    assert_eq!(envelope.agent_runtime_id, expected_agent_runtime_id);
+    assert_eq!(envelope.task_id, expected_task_id);
+
+    let signature = Signature::from_slice(
+        &base64::engine::general_purpose::STANDARD
+            .decode(&envelope.signature)
+            .expect("base64 signature"),
+    )
+    .expect("signature bytes");
+    stored_identity
+        .signing_key()
+        .expect("signing key")
+        .verifying_key()
+        .verify(
+            format!(
+                "{}:{}:{}",
+                envelope.agent_runtime_id, envelope.task_id, envelope.timestamp
+            )
+            .as_bytes(),
+            &signature,
+        )
+        .expect("signature should verify");
+}
+
 #[test]
 fn build_subagent_headers_sets_other_subagent_label() {
     let client = test_model_client(SessionSource::SubAgent(SubAgentSource::Other(
@@ -169,3 +305,130 @@ fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() {
     assert_eq!(auth_context.recovery_mode, Some("managed"));
     assert_eq!(auth_context.recovery_phase, Some("refresh_token"));
 }
+
+#[tokio::test]
+async fn responses_http_uses_agent_assertion_when_agent_task_is_present() {
+    core_test_support::skip_if_no_network!();
+
+    let server = responses::start_mock_server().await;
+    let request_recorder = responses::mount_sse_once(
+        &server,
+        responses::sse(vec![
+            responses::ev_response_created("resp-1"),
+            responses::ev_completed("resp-1"),
+        ]),
+    )
+    .await;
+    let provider =
+        create_oss_provider_with_base_url(&format!("{}/v1", server.uri()), WireApi::Responses);
+    let (_codex_home, client, agent_task, stored_identity) =
+        model_client_with_agent_task(provider).await;
+    let model_info = test_model_info();
+    let session_telemetry = test_session_telemetry();
+    let mut client_session = client.new_session_with_agent_task(Some(agent_task.clone()));
+
+    let mut stream = client_session
+        .stream(
+            &test_prompt("hello"),
+            &model_info,
+            &session_telemetry,
+            None,
+            ReasoningSummary::Auto,
+            None,
+            None,
+        )
+        .await
+        .expect("stream request should succeed");
+    drain_stream_to_completion(&mut stream)
+        .await
+        .expect("stream should complete");
+
+    let request = request_recorder.single_request();
+    let authorization = request
+        .header("authorization")
+        .expect("authorization header should be present");
+    assert_agent_assertion_header(
+        &authorization,
+        &stored_identity,
+        &agent_task.agent_runtime_id,
+        &agent_task.task_id,
+    );
+    assert_eq!(request.header("chatgpt-account-id"), None);
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn websocket_agent_task_bypasses_cached_bearer_prewarm() {
+    core_test_support::skip_if_no_network!();
+
+    let server = responses::start_websocket_server(vec![
+        vec![vec![
+            responses::ev_response_created("resp-prewarm"),
+            responses::ev_completed("resp-prewarm"),
+        ]],
+        vec![vec![
+            responses::ev_response_created("resp-1"),
+            responses::ev_completed("resp-1"),
+        ]],
+    ])
+    .await;
+    let mut provider =
+        create_oss_provider_with_base_url(&format!("{}/v1", server.uri()), WireApi::Responses);
+    provider.supports_websockets = true;
+    provider.websocket_connect_timeout_ms = Some(5_000);
+    let (_codex_home, client, agent_task, stored_identity) =
+        model_client_with_agent_task(provider).await;
+    let model_info = test_model_info();
+    let session_telemetry = test_session_telemetry();
+    let prompt = test_prompt("hello");
+
+    let mut prewarm_session = client.new_session();
+    prewarm_session
+        .prewarm_websocket(
+            &prompt,
+            &model_info,
+            &session_telemetry,
+            None,
+            ReasoningSummary::Auto,
+            None,
+            None,
+        )
+        .await
+        .expect("bearer prewarm should succeed");
+    drop(prewarm_session);
+
+    let mut agent_task_session = client.new_session_with_agent_task(Some(agent_task.clone()));
+    let mut stream = agent_task_session
+        .stream(
+            &prompt,
+            &model_info,
+            &session_telemetry,
+            None,
+            ReasoningSummary::Auto,
+            None,
+            None,
+        )
+        .await
+        .expect("agent task stream should succeed");
+    drain_stream_to_completion(&mut stream)
+        .await
+        .expect("agent task websocket stream should complete");
+
+    let handshakes = server.handshakes();
+    assert_eq!(handshakes.len(), 2);
+    assert_eq!(
+        handshakes[0].header("authorization"),
+        Some("Bearer Access Token".to_string())
+    );
+    let agent_authorization = handshakes[1]
+        .header("authorization")
+        .expect("agent handshake should include authorization");
+    assert_agent_assertion_header(
+        &agent_authorization,
+        &stored_identity,
+        &agent_task.agent_runtime_id,
+        &agent_task.task_id,
+    );
+    assert_eq!(handshakes[1].header("chatgpt-account-id"), None);
+
+    server.shutdown().await;
+}

codex-rs/core/src/codex.rs

@@ -2185,6 +2185,11 @@ impl Session {
                 config.analytics_enabled,
             )
         });
+        let agent_identity_manager = Arc::new(AgentIdentityManager::new(
+            config.as_ref(),
+            Arc::clone(&auth_manager),
+            session_configuration.session_source.clone(),
+        ));
         let services = SessionServices {
             // Initialize the MCP connection manager with an uninitialized
             // instance. It will be replaced with one created via
@@ -2207,11 +2212,7 @@ impl Session {
             hooks,
             rollout: Mutex::new(rollout_recorder),
             user_shell: Arc::new(default_shell),
-            agent_identity_manager: Arc::new(AgentIdentityManager::new(
-                config.as_ref(),
-                Arc::clone(&auth_manager),
-                session_configuration.session_source.clone(),
-            )),
+            agent_identity_manager: Arc::clone(&agent_identity_manager),
             shell_snapshot_tx,
             show_raw_agent_reasoning: config.show_raw_agent_reasoning,
             exec_policy,
@@ -2228,8 +2229,9 @@ impl Session {
             network_proxy,
             network_approval: Arc::clone(&network_approval),
             state_db: state_db_ctx.clone(),
-            model_client: ModelClient::new(
+            model_client: ModelClient::new_with_agent_identity_manager(
                 Some(Arc::clone(&auth_manager)),
+                Some(agent_identity_manager),
                 conversation_id,
                 installation_id,
                 session_configuration.provider.clone(),
@@ -6373,11 +6375,14 @@ pub(crate) async fn run_turn(
         }))
         .await;
     }
-    if let Err(error) = sess.ensure_agent_task_registered().await {
-        warn!(error = %error, "agent task registration failed");
-        sess.fail_agent_identity_registration(error).await;
-        return None;
-    }
+    let agent_task = match sess.ensure_agent_task_registered().await {
+        Ok(agent_task) => agent_task,
+        Err(error) => {
+            warn!(error = %error, "agent task registration failed");
+            sess.fail_agent_identity_registration(error).await;
+            return None;
+        }
+    };
 
     if !skill_items.is_empty() {
         sess.record_conversation_items(&turn_context, &skill_items)
@@ -6400,8 +6405,21 @@ pub(crate) async fn run_turn(
 
     // `ModelClientSession` is turn-scoped and caches WebSocket + sticky routing state, so we reuse
     // one instance across retries within this turn.
-    let mut client_session =
-        prewarmed_client_session.unwrap_or_else(|| sess.services.model_client.new_session());
+    let mut prewarmed_client_session = prewarmed_client_session;
+    if agent_task.is_some()
+        && let Some(prewarmed_client_session) = prewarmed_client_session.as_mut()
+    {
+        prewarmed_client_session.disable_cached_websocket_session_on_drop();
+    }
+    let mut client_session = if let Some(agent_task) = agent_task {
+        sess.services
+            .model_client
+            .new_session_with_agent_task(Some(agent_task))
+    } else if let Some(prewarmed_client_session) = prewarmed_client_session.take() {
+        prewarmed_client_session
+    } else {
+        sess.services.model_client.new_session()
+    };
     // Pending input is drained into history before building the next model request.
     // However, we defer that drain until after sampling in two cases:
     // 1. At the start of a turn, so the fresh user prompt in `input` gets sampled first.

codex-rs/core/src/mcp_openai_file.rs

@@ -112,10 +112,8 @@ async fn build_uploaded_local_argument_value(
     let token_data = auth
         .get_token_data()
         .map_err(|error| format!("failed to read ChatGPT auth for file upload: {error}"))?;
-    let upload_auth = CoreAuthProvider {
-        token: Some(token_data.access_token),
-        account_id: token_data.account_id,
-    };
+    let upload_auth =
+        CoreAuthProvider::from_bearer_token(Some(token_data.access_token), token_data.account_id);
     let uploaded = upload_local_file(
         turn_context.config.chatgpt_base_url.trim_end_matches('/'),
         &upload_auth,

codex-rs/login/src/api_bridge.rs

@@ -8,29 +8,20 @@ pub fn auth_provider_from_auth(
     provider: &ModelProviderInfo,
 ) -> codex_protocol::error::Result<CoreAuthProvider> {
     if let Some(api_key) = provider.api_key()? {
-        return Ok(CoreAuthProvider {
-            token: Some(api_key),
-            account_id: None,
-        });
+        return Ok(CoreAuthProvider::from_bearer_token(Some(api_key), None));
     }
 
     if let Some(token) = provider.experimental_bearer_token.clone() {
-        return Ok(CoreAuthProvider {
-            token: Some(token),
-            account_id: None,
-        });
+        return Ok(CoreAuthProvider::from_bearer_token(Some(token), None));
     }
 
     if let Some(auth) = auth {
         let token = auth.get_token()?;
-        Ok(CoreAuthProvider {
-            token: Some(token),
-            account_id: auth.get_account_id(),
-        })
+        Ok(CoreAuthProvider::from_bearer_token(
+            Some(token),
+            auth.get_account_id(),
+        ))
     } else {
-        Ok(CoreAuthProvider {
-            token: None,
-            account_id: None,
-        })
+        Ok(CoreAuthProvider::from_bearer_token(None, None))
     }
 }