Use AgentAssertion downstream behind use_agent_identity

Author: adrian-openai

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,15 +215,33 @@ 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 {
     fn bearer_token(&self) -> Option<String> {
         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,27 @@ 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()),
+        /*account_id*/ 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"),
+        /*account_id*/ 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);
 
@@ -335,7 +338,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,
@@ -349,6 +352,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,
+            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 {
@@ -579,7 +606,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)
@@ -605,7 +632,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");
@@ -621,7 +648,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)
@@ -653,11 +680,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,176 @@
+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 {
+            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 {
+            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");
+    }
+}