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,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);
 
@@ -327,7 +330,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,
@@ -341,6 +344,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 {
@@ -570,7 +597,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)
@@ -596,7 +623,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");
@@ -612,7 +639,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)
@@ -643,11 +670,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]