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

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 codex_home = tempfile::tempdir().expect("tempdir");
        let auth = make_chatgpt_auth(codex_home.path(), "account-123", Some("user-123"));
        let auth_manager = AuthManager::from_auth_for_testing(auth);
        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-123".to_string(),
            chatgpt_account_id: "account-123".to_string(),
            chatgpt_user_id: Some("user-123".to_string()),
            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 codex_home = tempfile::tempdir().expect("tempdir");
        let auth = make_chatgpt_auth(codex_home.path(), "account-123", Some("user-123"));
        let auth_manager = AuthManager::from_auth_for_testing(auth);
        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-123".to_string(),
            chatgpt_account_id: "account-123".to_string(),
            chatgpt_user_id: Some("user-123".to_string()),
            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");
    }

    fn make_chatgpt_auth(
        codex_home: &std::path::Path,
        account_id: &str,
        user_id: Option<&str>,
    ) -> CodexAuth {
        let auth_json = codex_login::AuthDotJson {
            auth_mode: Some(codex_app_server_protocol::AuthMode::Chatgpt),
            openai_api_key: None,
            tokens: Some(codex_login::token_data::TokenData {
                id_token: codex_login::token_data::IdTokenInfo {
                    email: None,
                    chatgpt_plan_type: None,
                    chatgpt_user_id: user_id.map(ToOwned::to_owned),
                    chatgpt_account_id: Some(account_id.to_string()),
                    raw_jwt: fake_id_token(account_id, user_id),
                },
                access_token: format!("access-token-{account_id}"),
                refresh_token: "refresh-token".to_string(),
                account_id: Some(account_id.to_string()),
            }),
            last_refresh: Some(chrono::Utc::now()),
            agent_identity: None,
        };
        codex_login::save_auth(
            codex_home,
            &auth_json,
            codex_login::AuthCredentialsStoreMode::File,
        )
        .expect("save auth");
        CodexAuth::from_auth_storage(codex_home, codex_login::AuthCredentialsStoreMode::File)
            .expect("load auth")
            .expect("auth")
    }

    fn fake_id_token(account_id: &str, user_id: Option<&str>) -> String {
        let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#);
        let payload = serde_json::json!({
            "https://api.openai.com/auth": {
                "chatgpt_user_id": user_id,
                "chatgpt_account_id": account_id,
            }
        });
        let payload = URL_SAFE_NO_PAD.encode(payload.to_string());
        format!("{header}.{payload}.signature")
    }
}