diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index f817d85d34..a5b1b6e730 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1991,7 +1991,6 @@ dependencies = [ "codex-windows-sandbox", "core-foundation 0.9.4", "core_test_support", - "crypto_box", "csv", "ctor 0.6.3", "dirs", diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index 886d9c2557..ef6593d09c 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -232,9 +232,7 @@ impl ThreadHistoryBuilder { RolloutItem::EventMsg(event) => self.handle_event(event), RolloutItem::Compacted(payload) => self.handle_compacted(payload), RolloutItem::ResponseItem(item) => self.handle_response_item(item), - RolloutItem::TurnContext(_) - | RolloutItem::SessionMeta(_) - | RolloutItem::SessionState(_) => {} + RolloutItem::TurnContext(_) | RolloutItem::SessionMeta(_) => {} } } diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index d9b2690658..36d20047fd 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -26,7 +26,6 @@ bm25 = { workspace = true } chrono = { workspace = true, features = ["serde"] } clap = { workspace = true, features = ["derive"] } codex-analytics = { workspace = true } -codex-agent-identity = { workspace = true } codex-api = { workspace = true } codex-app-server-protocol = { workspace = true } codex-apply-patch = { workspace = true } @@ -36,7 +35,6 @@ codex-connectors = { workspace = true } codex-config = { workspace = true } codex-core-plugins = { workspace = true } codex-core-skills = { workspace = true } -crypto_box = { workspace = true } codex-exec-server = { workspace = true } codex-features = { workspace = true } codex-feedback = { workspace = true } @@ -44,7 +42,6 @@ codex-login = { workspace = true } codex-mcp = { workspace = true } codex-model-provider-info = { workspace = true } codex-models-manager = { workspace = true } -ed25519-dalek = { workspace = true } codex-shell-command = { workspace = true } codex-execpolicy = { workspace = true } codex-git-utils = { workspace = true } @@ -150,11 +147,13 @@ codex-shell-escalation = { workspace = true } assert_cmd = { workspace = true } assert_matches = { workspace = true } codex-arg0 = { workspace = true } +codex-agent-identity = { workspace = true } codex-otel = { workspace = true } codex-test-binary-support = { workspace = true } codex-utils-cargo-bin = { workspace = true } core_test_support = { workspace = true } ctor = { workspace = true } +ed25519-dalek = { workspace = true } insta = { workspace = true } maplit = { workspace = true } opentelemetry = { workspace = true } diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 8004f6e6ed..d27f51d57d 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -116,7 +116,6 @@ fn keep_forked_rollout_item(item: &RolloutItem) -> bool { | ResponseItem::Compaction { .. } | ResponseItem::Other, ) => false, - RolloutItem::SessionState(_) => false, RolloutItem::Compacted(_) | RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) diff --git a/codex-rs/core/src/agent_identity.rs b/codex-rs/core/src/agent_identity.rs deleted file mode 100644 index eb9b8bee6f..0000000000 --- a/codex-rs/core/src/agent_identity.rs +++ /dev/null @@ -1,768 +0,0 @@ -use std::sync::Arc; -use std::time::Duration; - -use anyhow::Context; -use anyhow::Result; -use chrono::SecondsFormat; -use chrono::Utc; -use codex_agent_identity::AgentBillOfMaterials; -use codex_agent_identity::AgentIdentityKey; -use codex_agent_identity::agent_identity_biscuit_url; -use codex_agent_identity::agent_identity_request_id; -use codex_agent_identity::agent_registration_url; -use codex_agent_identity::build_abom; -use codex_agent_identity::generate_agent_key_material; -use codex_agent_identity::public_key_ssh_from_private_key_pkcs8_base64; -use codex_features::Feature; -use codex_login::AgentIdentityAuthRecord; -use codex_login::AuthManager; -use codex_login::CodexAuth; -use codex_login::default_client::create_client; -use codex_protocol::protocol::SessionSource; -use serde::Deserialize; -use serde::Serialize; -use tokio::sync::Semaphore; -use tracing::debug; -use tracing::info; -use tracing::warn; - -use crate::config::Config; - -mod task_registration; - -pub(crate) use task_registration::RegisteredAgentTask; - -const AGENT_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15); -const AGENT_IDENTITY_BISCUIT_TIMEOUT: Duration = Duration::from_secs(15); - -#[derive(Clone)] -pub(crate) struct AgentIdentityManager { - auth_manager: Arc, - chatgpt_base_url: String, - feature_enabled: bool, - abom: AgentBillOfMaterials, - ensure_lock: Arc, -} - -impl std::fmt::Debug for AgentIdentityManager { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("AgentIdentityManager") - .field("chatgpt_base_url", &self.chatgpt_base_url) - .field("feature_enabled", &self.feature_enabled) - .field("abom", &self.abom) - .finish_non_exhaustive() - } -} - -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] -pub(crate) struct StoredAgentIdentity { - pub(crate) binding_id: String, - pub(crate) chatgpt_account_id: String, - pub(crate) chatgpt_user_id: Option, - pub(crate) agent_runtime_id: String, - pub(crate) private_key_pkcs8_base64: String, - pub(crate) public_key_ssh: String, - pub(crate) registered_at: String, - pub(crate) abom: AgentBillOfMaterials, -} - -#[derive(Debug, Serialize)] -struct RegisterAgentRequest { - abom: AgentBillOfMaterials, - agent_public_key: String, - capabilities: Vec, -} - -#[derive(Debug, Deserialize)] -struct RegisterAgentResponse { - agent_runtime_id: String, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -struct AgentIdentityBinding { - binding_id: String, - chatgpt_account_id: String, - chatgpt_user_id: Option, - access_token: String, -} - -impl AgentIdentityManager { - pub(crate) fn new( - config: &Config, - auth_manager: Arc, - session_source: SessionSource, - ) -> Self { - Self { - auth_manager, - chatgpt_base_url: config.chatgpt_base_url.clone(), - feature_enabled: config.features.enabled(Feature::UseAgentIdentity), - abom: build_abom(session_source), - ensure_lock: Arc::new(Semaphore::new(/*permits*/ 1)), - } - } - - pub(crate) fn is_enabled(&self) -> bool { - self.feature_enabled - } - - pub(crate) async fn ensure_registered_identity(&self) -> Result> { - if !self.feature_enabled { - return Ok(None); - } - - let Some((auth, binding)) = self.current_auth_binding().await else { - return Ok(None); - }; - - self.ensure_registered_identity_for_binding(&auth, &binding) - .await - .map(Some) - } - - async fn ensure_registered_identity_for_binding( - &self, - auth: &CodexAuth, - binding: &AgentIdentityBinding, - ) -> Result { - let _guard = self - .ensure_lock - .acquire() - .await - .map_err(|_| anyhow::anyhow!("agent identity ensure semaphore closed"))?; - - if let Some(stored_identity) = self.load_stored_identity(auth, binding)? { - info!( - agent_runtime_id = %stored_identity.agent_runtime_id, - binding_id = %binding.binding_id, - "reusing stored agent identity" - ); - return Ok(stored_identity); - } - - let stored_identity = self.register_agent_identity(binding).await?; - self.store_identity(auth, &stored_identity)?; - Ok(stored_identity) - } - - pub(crate) async fn task_matches_current_identity(&self, task: &RegisteredAgentTask) -> bool { - if !self.feature_enabled { - return false; - } - - self.current_stored_identity() - .await - .is_some_and(|stored_identity| { - stored_identity.agent_runtime_id == task.agent_runtime_id - }) - } - - async fn current_auth_binding(&self) -> Option<(CodexAuth, AgentIdentityBinding)> { - let Some(auth) = self.auth_manager.auth().await else { - debug!("skipping agent identity flow because no auth is available"); - return None; - }; - - let binding = - AgentIdentityBinding::from_auth(&auth, self.auth_manager.forced_chatgpt_workspace_id()); - if binding.is_none() { - debug!("skipping agent identity flow because ChatGPT auth is unavailable"); - } - binding.map(|binding| (auth, binding)) - } - - async fn current_stored_identity(&self) -> Option { - let (auth, binding) = self.current_auth_binding().await?; - self.load_stored_identity(&auth, &binding).ok().flatten() - } - - async fn register_agent_identity( - &self, - binding: &AgentIdentityBinding, - ) -> Result { - let key_material = generate_agent_key_material()?; - let request_body = RegisterAgentRequest { - abom: self.abom.clone(), - agent_public_key: key_material.public_key_ssh.clone(), - capabilities: Vec::new(), - }; - - let url = agent_registration_url(&self.chatgpt_base_url); - let human_biscuit = self.mint_human_biscuit(binding, "POST", &url).await?; - let client = create_client(); - let response = client - .post(&url) - .header("X-OpenAI-Authorization", human_biscuit) - .json(&request_body) - .timeout(AGENT_REGISTRATION_TIMEOUT) - .send() - .await - .with_context(|| { - format!("failed to send agent identity registration request to {url}") - })?; - - if response.status().is_success() { - let response_body = response - .json::() - .await - .with_context(|| format!("failed to parse agent identity response from {url}"))?; - 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: response_body.agent_runtime_id, - 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(), - }; - info!( - agent_runtime_id = %stored_identity.agent_runtime_id, - binding_id = %binding.binding_id, - "registered agent identity" - ); - return Ok(stored_identity); - } - - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - anyhow::bail!("agent identity registration failed with status {status} from {url}: {body}") - } - - async fn mint_human_biscuit( - &self, - binding: &AgentIdentityBinding, - target_method: &str, - target_url: &str, - ) -> Result { - let url = agent_identity_biscuit_url(&self.chatgpt_base_url); - let request_id = agent_identity_request_id()?; - let client = create_client(); - let response = client - .get(&url) - .bearer_auth(&binding.access_token) - .header("X-Request-Id", request_id.clone()) - .header("X-Original-Method", target_method) - .header("X-Original-Url", target_url) - .timeout(AGENT_IDENTITY_BISCUIT_TIMEOUT) - .send() - .await - .with_context(|| format!("failed to send agent identity biscuit request to {url}"))?; - - if response.status().is_success() { - let human_biscuit = response - .headers() - .get("x-openai-authorization") - .context("agent identity biscuit response did not include x-openai-authorization")? - .to_str() - .context("agent identity biscuit response header was not valid UTF-8")? - .to_string(); - info!( - request_id = %request_id, - "minted human biscuit for agent identity registration" - ); - return Ok(human_biscuit); - } - - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - anyhow::bail!( - "agent identity biscuit minting failed with status {status} from {url}: {body}" - ) - } - - fn load_stored_identity( - &self, - auth: &CodexAuth, - binding: &AgentIdentityBinding, - ) -> Result> { - let Some(record) = auth.get_agent_identity(&binding.chatgpt_account_id) else { - return Ok(None); - }; - - let stored_identity = - match StoredAgentIdentity::from_auth_record(binding, record, self.abom.clone()) { - Ok(stored_identity) => stored_identity, - Err(error) => { - warn!( - binding_id = %binding.binding_id, - error = %error, - "stored agent identity is invalid; deleting cached value" - ); - auth.remove_agent_identity()?; - return Ok(None); - } - }; - - if !stored_identity.matches_binding(binding) { - warn!( - binding_id = %binding.binding_id, - "stored agent identity binding no longer matches current auth; deleting cached value" - ); - auth.remove_agent_identity()?; - return Ok(None); - } - - if let Err(error) = stored_identity.validate_key_material() { - warn!( - agent_runtime_id = %stored_identity.agent_runtime_id, - binding_id = %binding.binding_id, - error = %error, - "stored agent identity key material is invalid; deleting cached value" - ); - auth.remove_agent_identity()?; - return Ok(None); - } - - Ok(Some(stored_identity)) - } - - fn store_identity( - &self, - auth: &CodexAuth, - stored_identity: &StoredAgentIdentity, - ) -> Result<()> { - auth.set_agent_identity(stored_identity.to_auth_record())?; - Ok(()) - } - - #[cfg(test)] - pub(crate) fn new_for_tests( - auth_manager: Arc, - feature_enabled: bool, - chatgpt_base_url: String, - session_source: SessionSource, - ) -> Self { - Self { - auth_manager, - chatgpt_base_url, - feature_enabled, - abom: build_abom(session_source), - ensure_lock: Arc::new(Semaphore::new(/*permits*/ 1)), - } - } - - #[cfg(test)] - pub(crate) async fn seed_generated_identity_for_tests( - &self, - agent_runtime_id: &str, - ) -> Result { - 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 { - fn from_auth_record( - binding: &AgentIdentityBinding, - record: AgentIdentityAuthRecord, - abom: AgentBillOfMaterials, - ) -> Result { - if record.workspace_id != binding.chatgpt_account_id { - anyhow::bail!( - "stored agent identity workspace {:?} does not match current workspace {:?}", - record.workspace_id, - binding.chatgpt_account_id - ); - } - let public_key_ssh = - public_key_ssh_from_private_key_pkcs8_base64(&record.agent_private_key)?; - Ok(Self { - binding_id: binding.binding_id.clone(), - chatgpt_account_id: binding.chatgpt_account_id.clone(), - chatgpt_user_id: record.chatgpt_user_id, - agent_runtime_id: record.agent_runtime_id, - private_key_pkcs8_base64: record.agent_private_key, - public_key_ssh, - registered_at: record.registered_at, - abom, - }) - } - - fn to_auth_record(&self) -> AgentIdentityAuthRecord { - AgentIdentityAuthRecord { - workspace_id: self.chatgpt_account_id.clone(), - chatgpt_user_id: self.chatgpt_user_id.clone(), - agent_runtime_id: self.agent_runtime_id.clone(), - agent_private_key: self.private_key_pkcs8_base64.clone(), - registered_at: self.registered_at.clone(), - background_task_id: None, - } - } - - fn matches_binding(&self, binding: &AgentIdentityBinding) -> bool { - binding.matches_parts( - &self.binding_id, - &self.chatgpt_account_id, - self.chatgpt_user_id.as_deref(), - ) - } - - fn validate_key_material(&self) -> Result<()> { - let derived_public_key = - public_key_ssh_from_private_key_pkcs8_base64(&self.private_key_pkcs8_base64)?; - anyhow::ensure!( - self.public_key_ssh == derived_public_key, - "stored public key does not match the private key" - ); - Ok(()) - } - - pub(crate) fn agent_identity_key(&self) -> AgentIdentityKey<'_> { - AgentIdentityKey { - agent_runtime_id: &self.agent_runtime_id, - private_key_pkcs8_base64: &self.private_key_pkcs8_base64, - } - } -} - -impl AgentIdentityBinding { - fn matches_parts( - &self, - binding_id: &str, - chatgpt_account_id: &str, - chatgpt_user_id: Option<&str>, - ) -> bool { - binding_id == self.binding_id - && chatgpt_account_id == self.chatgpt_account_id - && match self.chatgpt_user_id.as_deref() { - Some(expected_user_id) => chatgpt_user_id == Some(expected_user_id), - None => true, - } - } - - fn from_auth(auth: &CodexAuth, forced_workspace_id: Option) -> Option { - if !auth.is_chatgpt_auth() { - return None; - } - - let token_data = auth.get_token_data().ok()?; - let resolved_account_id = - forced_workspace_id - .filter(|value| !value.is_empty()) - .or(token_data - .account_id - .clone() - .filter(|value| !value.is_empty()))?; - - Some(Self { - binding_id: format!("chatgpt-account-{resolved_account_id}"), - chatgpt_account_id: resolved_account_id, - chatgpt_user_id: token_data - .id_token - .chatgpt_user_id - .filter(|value| !value.is_empty()), - access_token: token_data.access_token, - }) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - use base64::Engine as _; - use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; - use base64::engine::general_purpose::URL_SAFE_NO_PAD; - use codex_app_server_protocol::AuthMode as ApiAuthMode; - use codex_login::AuthCredentialsStoreMode; - use codex_login::AuthDotJson; - use codex_login::save_auth; - use codex_login::token_data::IdTokenInfo; - use codex_login::token_data::TokenData; - use pretty_assertions::assert_eq; - use wiremock::Mock; - use wiremock::MockServer; - use wiremock::ResponseTemplate; - use wiremock::matchers::header; - use wiremock::matchers::method; - use wiremock::matchers::path; - - #[tokio::test] - async fn ensure_registered_identity_skips_when_feature_is_disabled() { - let auth_manager = - AuthManager::from_auth_for_testing(make_chatgpt_auth("account-123", Some("user-123"))); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ false, - "https://chatgpt.com/backend-api/".to_string(), - SessionSource::Cli, - ); - - assert_eq!(manager.ensure_registered_identity().await.unwrap(), None); - } - - #[tokio::test] - async fn ensure_registered_identity_skips_for_api_key_auth() { - let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test-key")); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - "https://chatgpt.com/backend-api/".to_string(), - SessionSource::Cli, - ); - - assert_eq!(manager.ensure_registered_identity().await.unwrap(), None); - } - - #[tokio::test] - async fn ensure_registered_identity_registers_and_reuses_cached_identity() { - let server = MockServer::start().await; - let chatgpt_base_url = server.uri(); - mount_human_biscuit(&server, &chatgpt_base_url).await; - Mock::given(method("POST")) - .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", - }))) - .expect(1) - .mount(&server) - .await; - - let auth_manager = - AuthManager::from_auth_for_testing(make_chatgpt_auth("account-123", Some("user-123"))); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - chatgpt_base_url, - SessionSource::Cli, - ); - - let first = manager - .ensure_registered_identity() - .await - .unwrap() - .expect("identity should be registered"); - let second = manager - .ensure_registered_identity() - .await - .unwrap() - .expect("identity should be reused"); - - 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"); - assert_eq!(first.chatgpt_user_id.as_deref(), Some("user-123")); - } - - #[tokio::test] - async fn ensure_registered_identity_deletes_invalid_cached_identity_and_reregisters() { - let server = MockServer::start().await; - let chatgpt_base_url = server.uri(); - mount_human_biscuit(&server, &chatgpt_base_url).await; - Mock::given(method("POST")) - .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", - }))) - .expect(1) - .mount(&server) - .await; - - let auth = make_chatgpt_auth("account-123", Some("user-123")); - let auth_manager = AuthManager::from_auth_for_testing(auth.clone()); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - chatgpt_base_url, - SessionSource::Cli, - ); - - let binding = - AgentIdentityBinding::from_auth(&auth, /*forced_workspace_id*/ None).expect("binding"); - auth.set_agent_identity(AgentIdentityAuthRecord { - workspace_id: "account-123".to_string(), - chatgpt_user_id: Some("user-123".to_string()), - agent_runtime_id: "agent_invalid".to_string(), - agent_private_key: "not-valid-base64".to_string(), - registered_at: "2026-01-01T00:00:00Z".to_string(), - background_task_id: None, - }) - .expect("seed invalid identity"); - - let stored = manager - .ensure_registered_identity() - .await - .unwrap() - .expect("identity should be registered"); - - 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"); - } - - #[tokio::test] - async fn ensure_registered_identity_deletes_different_user_identity_and_reregisters() { - let server = MockServer::start().await; - let chatgpt_base_url = server.uri(); - mount_human_biscuit(&server, &chatgpt_base_url).await; - Mock::given(method("POST")) - .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_new", - }))) - .expect(1) - .mount(&server) - .await; - - let auth = make_chatgpt_auth("account-123", Some("user-new")); - let stale_key = generate_agent_key_material().expect("key material"); - auth.set_agent_identity(AgentIdentityAuthRecord { - workspace_id: "account-123".to_string(), - chatgpt_user_id: Some("user-old".to_string()), - agent_runtime_id: "agent_old".to_string(), - agent_private_key: stale_key.private_key_pkcs8_base64, - registered_at: "2026-01-01T00:00:00Z".to_string(), - background_task_id: None, - }) - .expect("seed stale identity"); - - let auth_manager = AuthManager::from_auth_for_testing(auth.clone()); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - chatgpt_base_url, - SessionSource::Cli, - ); - - let stored = manager - .ensure_registered_identity() - .await - .unwrap() - .expect("identity should be registered"); - - assert_eq!(stored.agent_runtime_id, "agent_new"); - assert_eq!(stored.chatgpt_user_id.as_deref(), Some("user-new")); - let persisted = auth - .get_agent_identity("account-123") - .expect("stored identity"); - assert_eq!(persisted.agent_runtime_id, "agent_new"); - assert_eq!(persisted.chatgpt_user_id.as_deref(), Some("user-new")); - } - - #[tokio::test] - async fn ensure_registered_identity_uses_chatgpt_base_url() { - let server = MockServer::start().await; - let chatgpt_base_url = format!("{}/backend-api", server.uri()); - mount_human_biscuit(&server, &chatgpt_base_url).await; - Mock::given(method("POST")) - .and(path("/backend-api/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_canonical", - }))) - .expect(1) - .mount(&server) - .await; - - let auth_manager = - AuthManager::from_auth_for_testing(make_chatgpt_auth("account-123", Some("user-123"))); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - chatgpt_base_url, - SessionSource::Cli, - ); - - let stored = manager - .ensure_registered_identity() - .await - .unwrap() - .expect("identity should be registered"); - assert_eq!(stored.agent_runtime_id, "agent_canonical"); - } - - async fn mount_human_biscuit(server: &MockServer, chatgpt_base_url: &str) { - let biscuit_url = agent_identity_biscuit_url(chatgpt_base_url); - let biscuit_path = reqwest::Url::parse(&biscuit_url) - .expect("biscuit URL parses") - .path() - .to_string(); - let target_url = agent_registration_url(chatgpt_base_url); - Mock::given(method("GET")) - .and(path(biscuit_path)) - .and(header("authorization", "Bearer access-token-account-123")) - .and(header("x-original-method", "POST")) - .and(header("x-original-url", target_url)) - .respond_with( - ResponseTemplate::new(200).insert_header("x-openai-authorization", "human-biscuit"), - ) - .expect(1) - .mount(server) - .await; - } - - #[test] - fn encode_ssh_ed25519_public_key_matches_expected_wire_shape() { - let key_material = generate_agent_key_material().expect("key material"); - let (_, encoded_blob) = key_material - .public_key_ssh - .split_once(' ') - .expect("public key contains scheme"); - let decoded = BASE64_STANDARD.decode(encoded_blob).expect("base64"); - - assert_eq!(&decoded[..4], 11u32.to_be_bytes().as_slice()); - assert_eq!(&decoded[4..15], b"ssh-ed25519"); - assert_eq!(&decoded[15..19], 32u32.to_be_bytes().as_slice()); - assert_eq!(decoded.len(), 51); - } - - fn make_chatgpt_auth(account_id: &str, user_id: Option<&str>) -> CodexAuth { - let tempdir = tempfile::tempdir().expect("tempdir"); - let auth_json = AuthDotJson { - auth_mode: Some(ApiAuthMode::Chatgpt), - openai_api_key: None, - tokens: Some(TokenData { - id_token: IdTokenInfo { - email: None, - chatgpt_plan_type: None, - chatgpt_user_id: user_id.map(ToOwned::to_owned), - chatgpt_account_id: Some(account_id.to_string()), - chatgpt_account_is_fedramp: false, - 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(Utc::now()), - agent_identity: None, - }; - save_auth(tempdir.path(), &auth_json, AuthCredentialsStoreMode::File).expect("save auth"); - CodexAuth::from_auth_storage(tempdir.path(), 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") - } -} diff --git a/codex-rs/core/src/agent_identity/task_registration.rs b/codex-rs/core/src/agent_identity/task_registration.rs deleted file mode 100644 index b114cd68cb..0000000000 --- a/codex-rs/core/src/agent_identity/task_registration.rs +++ /dev/null @@ -1,435 +0,0 @@ -use std::time::Duration; - -use anyhow::Context; -use anyhow::Result; -use codex_agent_identity::AgentTaskAuthorizationTarget; -use codex_agent_identity::agent_task_registration_url; -use codex_agent_identity::decrypt_task_id_response; -use codex_agent_identity::sign_task_registration_payload; -use codex_protocol::protocol::SessionAgentTask; -use serde::Deserialize; -use serde::Serialize; -use tracing::info; - -use super::*; - -const AGENT_TASK_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15); - -#[derive(Clone, Debug, PartialEq, Eq)] -pub(crate) struct RegisteredAgentTask { - pub(crate) agent_runtime_id: String, - pub(crate) task_id: String, - pub(crate) registered_at: String, -} - -#[derive(Debug, Serialize)] -struct RegisterTaskRequest { - signature: String, - timestamp: String, -} - -#[derive(Debug, Deserialize)] -struct RegisterTaskResponse { - encrypted_task_id: String, -} - -impl AgentIdentityManager { - pub(crate) async fn register_task(&self) -> Result> { - if !self.feature_enabled { - return Ok(None); - } - - let Some((auth, binding)) = self.current_auth_binding().await else { - return Ok(None); - }; - - self.register_task_for_binding(auth, binding).await - } - - async fn register_task_for_binding( - &self, - auth: CodexAuth, - binding: AgentIdentityBinding, - ) -> Result> { - let stored_identity = self - .ensure_registered_identity_for_binding(&auth, &binding) - .await?; - - let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); - let request_body = RegisterTaskRequest { - signature: sign_task_registration_payload( - stored_identity.agent_identity_key(), - ×tamp, - )?, - timestamp, - }; - - let client = create_client(); - let url = - agent_task_registration_url(&self.chatgpt_base_url, &stored_identity.agent_runtime_id); - let human_biscuit = self.mint_human_biscuit(&binding, "POST", &url).await?; - let response = client - .post(&url) - .header("X-OpenAI-Authorization", human_biscuit) - .json(&request_body) - .timeout(AGENT_TASK_REGISTRATION_TIMEOUT) - .send() - .await - .with_context(|| format!("failed to send agent task registration request to {url}"))?; - - if response.status().is_success() { - let response_body = response - .json::() - .await - .with_context(|| format!("failed to parse agent task response from {url}"))?; - let registered_task = RegisteredAgentTask { - agent_runtime_id: stored_identity.agent_runtime_id.clone(), - task_id: decrypt_task_id_response( - stored_identity.agent_identity_key(), - &response_body.encrypted_task_id, - )?, - registered_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), - }; - info!( - agent_runtime_id = %registered_task.agent_runtime_id, - task_id = %registered_task.task_id, - "registered agent task" - ); - return Ok(Some(registered_task)); - } - - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - anyhow::bail!("agent task registration failed with status {status} from {url}: {body}") - } -} - -impl RegisteredAgentTask { - pub(crate) fn authorization_target(&self) -> AgentTaskAuthorizationTarget<'_> { - AgentTaskAuthorizationTarget { - agent_runtime_id: &self.agent_runtime_id, - task_id: &self.task_id, - } - } - - pub(crate) fn to_session_agent_task(&self) -> SessionAgentTask { - SessionAgentTask { - agent_runtime_id: self.agent_runtime_id.clone(), - task_id: self.task_id.clone(), - registered_at: self.registered_at.clone(), - } - } - - pub(crate) fn from_session_agent_task(task: SessionAgentTask) -> Self { - Self { - agent_runtime_id: task.agent_runtime_id, - task_id: task.task_id, - registered_at: task.registered_at, - } - } -} - -#[cfg(test)] -mod tests { - use base64::Engine as _; - use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; - use base64::engine::general_purpose::URL_SAFE_NO_PAD; - use codex_agent_identity::curve25519_secret_key_from_private_key_pkcs8_base64; - use codex_app_server_protocol::AuthMode as ApiAuthMode; - use codex_login::AuthCredentialsStoreMode; - use codex_login::AuthDotJson; - use codex_login::save_auth; - use codex_login::token_data::IdTokenInfo; - use codex_login::token_data::TokenData; - use pretty_assertions::assert_eq; - use wiremock::Mock; - use wiremock::MockServer; - use wiremock::ResponseTemplate; - use wiremock::matchers::header; - use wiremock::matchers::method; - use wiremock::matchers::path; - - use super::*; - - #[tokio::test] - async fn register_task_skips_when_feature_is_disabled() { - let auth_manager = - AuthManager::from_auth_for_testing(make_chatgpt_auth("account-123", Some("user-123"))); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ false, - "https://chatgpt.com/backend-api/".to_string(), - SessionSource::Cli, - ); - - assert_eq!(manager.register_task().await.unwrap(), None); - } - - #[tokio::test] - async fn register_task_skips_for_api_key_auth() { - let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test-key")); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - "https://chatgpt.com/backend-api/".to_string(), - SessionSource::Cli, - ); - - assert_eq!(manager.register_task().await.unwrap(), None); - } - - #[tokio::test] - async fn register_task_registers_and_decrypts_plaintext_task_id() { - let server = MockServer::start().await; - let chatgpt_base_url = server.uri(); - mount_human_biscuit(&server, &chatgpt_base_url, "agent-123").await; - let auth = make_chatgpt_auth("account-123", Some("user-123")); - let auth_manager = AuthManager::from_auth_for_testing(auth.clone()); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - chatgpt_base_url, - SessionSource::Cli, - ); - let stored_identity = seed_stored_identity(&manager, &auth, "agent-123", "account-123"); - let encrypted_task_id = - encrypt_task_id_for_identity(&stored_identity, "task_123").expect("task ciphertext"); - - Mock::given(method("POST")) - .and(path("/v1/agent/agent-123/task/register")) - .and(header("x-openai-authorization", "human-biscuit")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "encrypted_task_id": encrypted_task_id, - }))) - .expect(1) - .mount(&server) - .await; - - let task = manager - .register_task() - .await - .unwrap() - .expect("task should be registered"); - - assert_eq!( - task, - RegisteredAgentTask { - agent_runtime_id: "agent-123".to_string(), - task_id: "task_123".to_string(), - registered_at: task.registered_at.clone(), - } - ); - } - - #[tokio::test] - async fn register_task_uses_chatgpt_base_url() { - let server = MockServer::start().await; - let chatgpt_base_url = format!("{}/backend-api", server.uri()); - mount_human_biscuit(&server, &chatgpt_base_url, "agent-fallback").await; - let auth = make_chatgpt_auth("account-123", Some("user-123")); - let auth_manager = AuthManager::from_auth_for_testing(auth.clone()); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - chatgpt_base_url, - SessionSource::Cli, - ); - let stored_identity = - seed_stored_identity(&manager, &auth, "agent-fallback", "account-123"); - let encrypted_task_id = encrypt_task_id_for_identity(&stored_identity, "task_fallback") - .expect("task ciphertext"); - - Mock::given(method("POST")) - .and(path("/backend-api/v1/agent/agent-fallback/task/register")) - .and(header("x-openai-authorization", "human-biscuit")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "encrypted_task_id": encrypted_task_id, - }))) - .expect(1) - .mount(&server) - .await; - - let task = manager - .register_task() - .await - .unwrap() - .expect("task should be registered"); - - assert_eq!(task.agent_runtime_id, "agent-fallback"); - assert_eq!(task.task_id, "task_fallback"); - } - - #[tokio::test] - async fn register_task_for_binding_keeps_one_auth_snapshot() { - let server = MockServer::start().await; - let chatgpt_base_url = server.uri(); - mount_human_biscuit(&server, &chatgpt_base_url, "agent-123").await; - let binding_auth = make_chatgpt_auth("account-123", Some("user-123")); - let auth_manager = - AuthManager::from_auth_for_testing(make_chatgpt_auth("account-456", Some("user-456"))); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - chatgpt_base_url, - SessionSource::Cli, - ); - let stored_identity = - seed_stored_identity(&manager, &binding_auth, "agent-123", "account-123"); - let encrypted_task_id = - encrypt_task_id_for_identity(&stored_identity, "task_123").expect("task ciphertext"); - let binding = - AgentIdentityBinding::from_auth(&binding_auth, /*forced_workspace_id*/ None) - .expect("binding"); - - Mock::given(method("POST")) - .and(path("/v1/agent/agent-123/task/register")) - .and(header("x-openai-authorization", "human-biscuit")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "encrypted_task_id": encrypted_task_id, - }))) - .expect(1) - .mount(&server) - .await; - - let task = manager - .register_task_for_binding(binding_auth, binding) - .await - .unwrap() - .expect("task should be registered"); - - assert_eq!( - task, - RegisteredAgentTask { - agent_runtime_id: "agent-123".to_string(), - task_id: "task_123".to_string(), - registered_at: task.registered_at.clone(), - } - ); - } - - #[tokio::test] - async fn task_matches_current_identity_rejects_stale_registered_identity() { - let auth_manager = - AuthManager::from_auth_for_testing(make_chatgpt_auth("account-456", Some("user-456"))); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - "https://chatgpt.com/backend-api/".to_string(), - SessionSource::Cli, - ); - let 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!(!manager.task_matches_current_identity(&task).await); - } - - async fn mount_human_biscuit( - server: &MockServer, - chatgpt_base_url: &str, - agent_runtime_id: &str, - ) { - let biscuit_url = agent_identity_biscuit_url(chatgpt_base_url); - let biscuit_path = reqwest::Url::parse(&biscuit_url) - .expect("biscuit URL parses") - .path() - .to_string(); - let target_url = agent_task_registration_url(chatgpt_base_url, agent_runtime_id); - Mock::given(method("GET")) - .and(path(biscuit_path)) - .and(header("authorization", "Bearer access-token-account-123")) - .and(header("x-original-method", "POST")) - .and(header("x-original-url", target_url)) - .respond_with( - ResponseTemplate::new(200).insert_header("x-openai-authorization", "human-biscuit"), - ) - .expect(1) - .mount(server) - .await; - } - - fn seed_stored_identity( - manager: &AgentIdentityManager, - auth: &CodexAuth, - agent_runtime_id: &str, - account_id: &str, - ) -> StoredAgentIdentity { - let key_material = generate_agent_key_material().expect("key material"); - let binding = - AgentIdentityBinding::from_auth(auth, /*forced_workspace_id*/ None).expect("binding"); - let stored_identity = StoredAgentIdentity { - binding_id: binding.binding_id, - chatgpt_account_id: account_id.to_string(), - chatgpt_user_id: Some("user-123".to_string()), - 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: manager.abom.clone(), - }; - manager - .store_identity(auth, &stored_identity) - .expect("store identity"); - let persisted = auth - .get_agent_identity(account_id) - .expect("persisted identity"); - assert_eq!(persisted.agent_runtime_id, agent_runtime_id); - stored_identity - } - - fn encrypt_task_id_for_identity( - stored_identity: &StoredAgentIdentity, - task_id: &str, - ) -> Result { - let mut rng = crypto_box::aead::OsRng; - let public_key = curve25519_secret_key_from_private_key_pkcs8_base64( - &stored_identity.private_key_pkcs8_base64, - )? - .public_key(); - let ciphertext = public_key - .seal(&mut rng, task_id.as_bytes()) - .map_err(|_| anyhow::anyhow!("failed to encrypt test task id"))?; - Ok(BASE64_STANDARD.encode(ciphertext)) - } - - fn make_chatgpt_auth(account_id: &str, user_id: Option<&str>) -> CodexAuth { - let tempdir = tempfile::tempdir().expect("tempdir"); - let auth_json = AuthDotJson { - auth_mode: Some(ApiAuthMode::Chatgpt), - openai_api_key: None, - tokens: Some(TokenData { - id_token: IdTokenInfo { - email: None, - chatgpt_plan_type: None, - chatgpt_user_id: user_id.map(ToOwned::to_owned), - chatgpt_account_id: Some(account_id.to_string()), - chatgpt_account_is_fedramp: false, - 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(Utc::now()), - agent_identity: None, - }; - save_auth(tempdir.path(), &auth_json, AuthCredentialsStoreMode::File).expect("save auth"); - CodexAuth::from_auth_storage(tempdir.path(), 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") - } -} diff --git a/codex-rs/core/src/arc_monitor.rs b/codex-rs/core/src/arc_monitor.rs index 99767872d5..3340e6c6a4 100644 --- a/codex-rs/core/src/arc_monitor.rs +++ b/codex-rs/core/src/arc_monitor.rs @@ -110,42 +110,32 @@ pub(crate) async fn monitor_action( }, None => None, }; - let (authorization_header_value, account_id) = if let Some(token) = - read_non_empty_env_var(CODEX_ARC_MONITOR_TOKEN) - { - ( - format!("Bearer {token}"), - auth.as_ref().and_then(CodexAuth::get_account_id), - ) - } else if let Some(authorization_header_value) = - match sess.authorization_header_for_current_agent_task().await { - Ok(authorization_header_value) => authorization_header_value, - Err(err) => { - warn!( - error = %err, - "skipping safety monitor because agent assertion authorization is unavailable" - ); + let (authorization_header_value, account_id) = + if let Some(token) = read_non_empty_env_var(CODEX_ARC_MONITOR_TOKEN) { + ( + format!("Bearer {token}"), + auth.as_ref().and_then(CodexAuth::get_account_id), + ) + } else { + let Some(auth) = auth.as_ref() else { return ArcMonitorOutcome::Ok; - } - } - { - (authorization_header_value, None) - } else { - let Some(auth) = auth.as_ref() else { - return ArcMonitorOutcome::Ok; - }; - let token = match auth.get_token() { - Ok(token) => token, - Err(err) => { - warn!( - error = %err, - "skipping safety monitor because auth token is unavailable" - ); + }; + let Some(auth_manager) = turn_context.auth_manager.as_ref() else { return ArcMonitorOutcome::Ok; - } + }; + let Some(authorization_header_value) = auth_manager + .chatgpt_authorization_header_for_auth(auth) + .await + else { + return ArcMonitorOutcome::Ok; + }; + let account_id = if authorization_header_value.starts_with("AgentAssertion ") { + None + } else { + auth.get_account_id() + }; + (authorization_header_value, account_id) }; - (format!("Bearer {token}"), auth.get_account_id()) - }; let url = read_non_empty_env_var(CODEX_ARC_MONITOR_ENDPOINT_OVERRIDE).unwrap_or_else(|| { format!( diff --git a/codex-rs/core/src/arc_monitor_tests.rs b/codex-rs/core/src/arc_monitor_tests.rs index 6bf1aa98f0..5d2315843f 100644 --- a/codex-rs/core/src/arc_monitor_tests.rs +++ b/codex-rs/core/src/arc_monitor_tests.rs @@ -14,13 +14,14 @@ use wiremock::matchers::method; use wiremock::matchers::path; use super::*; -use crate::agent_identity::AgentIdentityManager; -use crate::agent_identity::RegisteredAgentTask; use crate::session::tests::make_session_and_context; use chrono::Utc; +use codex_agent_identity::generate_agent_key_material; +use codex_login::AgentIdentityAuthRecord; use codex_login::AuthCredentialsStoreMode; use codex_login::AuthDotJson; use codex_login::AuthManager; +use codex_login::BackgroundAgentTaskAuthMode; use codex_login::CodexAuth; use codex_login::save_auth; use codex_login::token_data::IdTokenInfo; @@ -31,7 +32,6 @@ use codex_protocol::models::LocalShellExecAction; use codex_protocol::models::LocalShellStatus; use codex_protocol::models::MessagePhase; use codex_protocol::models::ResponseItem; -use codex_protocol::protocol::SessionSource; use tempfile::tempdir; const TEST_ID_TOKEN: &str = concat!( @@ -72,9 +72,10 @@ impl Drop for EnvVarGuard { async fn install_cached_agent_task_auth( session: &mut Session, turn_context: &mut TurnContext, - chatgpt_base_url: String, + _chatgpt_base_url: String, ) { let auth_dir = tempdir().expect("temp auth dir"); + let key_material = generate_agent_key_material().expect("generate test key material"); let auth_json = AuthDotJson { auth_mode: Some(codex_app_server_protocol::AuthMode::Chatgpt), openai_api_key: None, @@ -92,33 +93,26 @@ async fn install_cached_agent_task_auth( account_id: Some("account_id".to_string()), }), last_refresh: Some(Utc::now()), - agent_identity: None, + agent_identity: Some(AgentIdentityAuthRecord { + workspace_id: "account_id".to_string(), + chatgpt_user_id: None, + agent_runtime_id: "agent-123".to_string(), + agent_private_key: key_material.private_key_pkcs8_base64, + registered_at: "2026-04-15T00:00:00Z".to_string(), + background_task_id: Some("task-123".to_string()), + }), }; save_auth(auth_dir.path(), &auth_json, AuthCredentialsStoreMode::File).expect("save test auth"); let auth = CodexAuth::from_auth_storage(auth_dir.path(), AuthCredentialsStoreMode::File) .expect("load test auth") .expect("test auth"); let auth_manager = AuthManager::from_auth_for_testing(auth); - let agent_identity_manager = Arc::new(AgentIdentityManager::new_for_tests( - Arc::clone(&auth_manager), - /*feature_enabled*/ true, - chatgpt_base_url, - SessionSource::Exec, - )); - let stored_identity = agent_identity_manager - .seed_generated_identity_for_tests("agent-123") - .await - .expect("seed test identity"); + auth_manager.set_chatgpt_backend_auth_config( + Some("https://chatgpt.com/backend-api".to_string()), + BackgroundAgentTaskAuthMode::Enabled, + ); session.services.auth_manager = Arc::clone(&auth_manager); - session.services.agent_identity_manager = agent_identity_manager; turn_context.auth_manager = Some(auth_manager); - session - .cache_agent_task_for_tests(RegisteredAgentTask { - agent_runtime_id: stored_identity.agent_runtime_id, - task_id: "task-123".to_string(), - registered_at: "2026-04-15T00:00:00Z".to_string(), - }) - .await; } #[tokio::test] diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 50f9dd295e..762fda76f6 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -31,7 +31,6 @@ use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; -use crate::agent_identity::RegisteredAgentTask; use codex_api::ApiError; use codex_api::AuthProvider; use codex_api::CompactClient as ApiCompactClient; @@ -216,7 +215,6 @@ pub struct ModelClient { pub struct ModelClientSession { client: ModelClient, websocket_session: WebsocketSession, - agent_task: Option, cache_websocket_session_on_drop: bool, /// Turn state for sticky routing. /// @@ -335,25 +333,10 @@ 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(/*agent_task*/ None) - } - - pub(crate) fn new_session_with_agent_task( - &self, - agent_task: Option, - ) -> 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, - agent_task, - cache_websocket_session_on_drop, + websocket_session: self.take_cached_websocket_session(), + cache_websocket_session_on_drop: true, turn_state: Arc::new(OnceLock::new()), } } @@ -436,7 +419,7 @@ impl ModelClient { if prompt.input.is_empty() { return Ok(Vec::new()); } - let client_setup = self.current_client_setup(/*agent_task*/ None).await?; + let client_setup = self.current_client_setup().await?; let transport = ReqwestTransport::new(build_reqwest_client()); let request_telemetry = Self::build_request_telemetry( session_telemetry, @@ -500,7 +483,7 @@ impl ModelClient { ) -> Result { // 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(/*agent_task*/ None).await?; + let client_setup = self.current_client_setup().await?; let mut sideband_headers = extra_headers.clone(); sideband_headers.extend(sideband_websocket_auth_headers( client_setup.api_auth.as_ref(), @@ -535,7 +518,7 @@ impl ModelClient { return Ok(Vec::new()); } - let client_setup = self.current_client_setup(/*agent_task*/ None).await?; + let client_setup = self.current_client_setup().await?; let transport = ReqwestTransport::new(build_reqwest_client()); let request_telemetry = Self::build_request_telemetry( session_telemetry, @@ -685,35 +668,30 @@ 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, - agent_task: Option<&RegisteredAgentTask>, - ) -> Result { + async fn current_client_setup(&self) -> Result { let auth = self.state.provider.auth().await; let api_provider = self.state.provider.api_provider().await?; let auth_manager = self.state.provider.auth_manager(); - let api_auth = match (agent_task, auth_manager.as_ref(), auth.as_ref()) { - (Some(agent_task), Some(auth_manager), Some(auth)) => { + let api_auth = match (auth_manager.as_ref(), auth.as_ref()) { + (Some(auth_manager), Some(auth)) + if self.state.provider.info().requires_openai_auth && auth.is_chatgpt_auth() => + { if let Some(authorization_header_value) = auth_manager - .chatgpt_agent_task_authorization_header_for_auth( - auth, - agent_task.authorization_target(), - ) - .map_err(|err| { - CodexErr::Stream( - format!("failed to build agent assertion authorization: {err}"), - None, - ) - })? + .chatgpt_authorization_header_for_auth(auth) + .await { debug!( - agent_runtime_id = %agent_task.agent_runtime_id, - task_id = %agent_task.task_id, - "using agent assertion authorization for downstream request" + auth_mode = ?auth.api_auth_mode(), + "using auth manager authorization for downstream request" ); + let account_id = if authorization_header_value.starts_with("AgentAssertion ") { + None + } else { + auth.get_account_id() + }; let mut auth_provider = AuthorizationHeaderAuthProvider::new( Some(authorization_header_value), - /*account_id*/ None, + account_id, ); if auth.is_fedramp_account() { auth_provider = auth_provider.with_fedramp_routing_header(); @@ -866,10 +844,6 @@ impl Drop for ModelClientSession { } 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; @@ -1071,15 +1045,11 @@ impl ModelClientSession { return Ok(()); } - 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 client_setup = self.client.current_client_setup().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.as_ref(), @@ -1233,10 +1203,7 @@ impl ModelClientSession { .map(AuthManager::unauthorized_recovery); let mut pending_retry = PendingUnauthorizedRetry::default(); loop { - let client_setup = self - .client - .current_client_setup(self.agent_task.as_ref()) - .await?; + let client_setup = self.client.current_client_setup().await?; let transport = ReqwestTransport::new(build_reqwest_client()); let request_auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), @@ -1325,10 +1292,7 @@ impl ModelClientSession { .map(AuthManager::unauthorized_recovery); let mut pending_retry = PendingUnauthorizedRetry::default(); loop { - let client_setup = self - .client - .current_client_setup(self.agent_task.as_ref()) - .await?; + let client_setup = self.client.current_client_setup().await?; let request_auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), client_setup.api_auth.as_ref(), diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 3336685eb3..db1541ab82 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -1,5 +1,3 @@ -use std::sync::Arc; - use super::AuthRequestTelemetryContext; use super::ModelClient; use super::PendingUnauthorizedRetry; @@ -11,14 +9,14 @@ use super::X_CODEX_WINDOW_ID_HEADER; use super::X_OPENAI_SUBAGENT_HEADER; use crate::Prompt; use crate::ResponseEvent; -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_agent_identity::generate_agent_key_material; use codex_agent_identity::verifying_key_from_private_key_pkcs8_base64; use codex_app_server_protocol::AuthMode; +use codex_login::AgentIdentityAuthRecord; use codex_login::AuthManager; +use codex_login::BackgroundAgentTaskAuthMode; use codex_login::CodexAuth; use codex_model_provider::BearerAuthProvider; use codex_model_provider_info::ModelProviderInfo; @@ -39,7 +37,6 @@ use futures::StreamExt; use pretty_assertions::assert_eq; use serde::Deserialize; 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); @@ -125,32 +122,36 @@ async fn drain_stream_to_completion(stream: &mut crate::ResponseStream) -> anyho Ok(()) } -async fn model_client_with_agent_task( +struct StoredBackgroundTaskAuth { + agent_runtime_id: String, + task_id: String, + private_key_pkcs8_base64: String, +} + +fn model_client_with_stored_background_task( provider: ModelProviderInfo, -) -> ( - TempDir, - ModelClient, - RegisteredAgentTask, - StoredAgentIdentity, -) { - let codex_home = tempfile::tempdir().expect("tempdir"); - 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, - )); - 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(), +) -> (ModelClient, StoredBackgroundTaskAuth) { + let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); + let key_material = generate_agent_key_material().expect("generate test identity key material"); + let stored_auth = StoredBackgroundTaskAuth { + agent_runtime_id: "agent-123".to_string(), task_id: "task-123".to_string(), - registered_at: "2026-03-23T12:00:00Z".to_string(), + private_key_pkcs8_base64: key_material.private_key_pkcs8_base64.clone(), }; + auth.set_agent_identity(AgentIdentityAuthRecord { + workspace_id: "account_id".to_string(), + chatgpt_user_id: None, + agent_runtime_id: stored_auth.agent_runtime_id.clone(), + agent_private_key: key_material.private_key_pkcs8_base64, + registered_at: "2026-03-23T12:00:00Z".to_string(), + background_task_id: Some(stored_auth.task_id.clone()), + }) + .expect("store agent identity"); + let auth_manager = AuthManager::from_auth_for_testing(auth); + auth_manager.set_chatgpt_backend_auth_config( + Some("https://chatgpt.com/backend-api".to_string()), + BackgroundAgentTaskAuthMode::Enabled, + ); let client = ModelClient::new( Some(auth_manager), ThreadId::new(), @@ -162,7 +163,7 @@ async fn model_client_with_agent_task( /*include_timing_metrics*/ false, /*beta_features_header*/ None, ); - (codex_home, client, agent_task, stored_identity) + (client, stored_auth) } #[derive(Debug, Deserialize)] @@ -175,7 +176,7 @@ struct AgentAssertionEnvelope { fn assert_agent_assertion_header( authorization_header: &str, - stored_identity: &StoredAgentIdentity, + stored_auth: &StoredBackgroundTaskAuth, expected_agent_runtime_id: &str, expected_task_id: &str, ) { @@ -198,7 +199,7 @@ fn assert_agent_assertion_header( .expect("base64 signature"), ) .expect("signature bytes"); - verifying_key_from_private_key_pkcs8_base64(&stored_identity.private_key_pkcs8_base64) + verifying_key_from_private_key_pkcs8_base64(&stored_auth.private_key_pkcs8_base64) .expect("verifying key") .verify( format!( @@ -303,7 +304,7 @@ fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() { } #[tokio::test] -async fn responses_http_uses_agent_assertion_when_agent_task_is_present() { +async fn responses_http_uses_agent_assertion_when_stored_background_task_is_present() { core_test_support::skip_if_no_network!(); let server = responses::start_mock_server().await; @@ -317,11 +318,10 @@ async fn responses_http_uses_agent_assertion_when_agent_task_is_present() { .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 (client, stored_auth) = model_client_with_stored_background_task(provider); 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 client_session = client.new_session(); let mut stream = client_session .stream( @@ -345,15 +345,15 @@ async fn responses_http_uses_agent_assertion_when_agent_task_is_present() { .expect("authorization header should be present"); assert_agent_assertion_header( &authorization, - &stored_identity, - &agent_task.agent_runtime_id, - &agent_task.task_id, + &stored_auth, + &stored_auth.agent_runtime_id, + &stored_auth.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() { +async fn websocket_uses_agent_assertion_when_stored_background_task_is_present() { core_test_support::skip_if_no_network!(); let server = responses::start_websocket_server(vec![ @@ -371,29 +371,13 @@ async fn websocket_agent_task_bypasses_cached_bearer_prewarm() { 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 (client, stored_auth) = model_client_with_stored_background_task(provider); 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, - /*effort*/ None, - ReasoningSummary::Auto, - /*service_tier*/ None, - /*turn_metadata_header*/ 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 + let mut session = client.new_session(); + let mut stream = session .stream( &prompt, &model_info, @@ -404,27 +388,23 @@ async fn websocket_agent_task_bypasses_cached_bearer_prewarm() { /*turn_metadata_header*/ None, ) .await - .expect("agent task stream should succeed"); + .expect("stream should succeed"); drain_stream_to_completion(&mut stream) .await - .expect("agent task websocket stream should complete"); + .expect("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] + assert_eq!(handshakes.len(), 1); + let authorization = handshakes[0] .header("authorization") - .expect("agent handshake should include authorization"); + .expect("handshake should include authorization"); assert_agent_assertion_header( - &agent_authorization, - &stored_identity, - &agent_task.agent_runtime_id, - &agent_task.task_id, + &authorization, + &stored_auth, + &stored_auth.agent_runtime_id, + &stored_auth.task_id, ); - assert_eq!(handshakes[1].header("chatgpt-account-id"), None); + assert_eq!(handshakes[0].header("chatgpt-account-id"), None); server.shutdown().await; } diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 1809789a8e..14ee38e6b7 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -5,7 +5,6 @@ // the TUI or the tracing stack). #![deny(clippy::print_stdout, clippy::print_stderr)] -mod agent_identity; mod apply_patch; mod apps; mod arc_monitor; diff --git a/codex-rs/core/src/mcp_openai_file.rs b/codex-rs/core/src/mcp_openai_file.rs index 324eb4ad1a..d80c3241f6 100644 --- a/codex-rs/core/src/mcp_openai_file.rs +++ b/codex-rs/core/src/mcp_openai_file.rs @@ -16,7 +16,6 @@ use codex_api::AuthProvider; use codex_api::upload_local_file; use codex_login::CodexAuth; use codex_model_provider::AuthorizationHeaderAuthProvider; -use codex_model_provider::BearerAuthProvider; use serde_json::Value as JsonValue; pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files( @@ -120,29 +119,25 @@ async fn build_uploaded_local_argument_value( "ChatGPT auth is required to upload local files for Codex Apps tools".to_string(), ); }; - let upload_auth: Box = if let Some(authorization_header_value) = sess - .authorization_header_for_current_agent_task() + let authorization_header_value = sess + .services + .auth_manager + .chatgpt_authorization_header_for_auth(auth) .await - .map_err(|error| format!("failed to build agent assertion authorization: {error}"))? - { - let mut auth_provider = AuthorizationHeaderAuthProvider::new( - Some(authorization_header_value), - /*account_id*/ None, - ); - if auth.is_fedramp_account() { - auth_provider = auth_provider.with_fedramp_routing_header(); - } - Box::new(auth_provider) + .ok_or_else(|| { + "ChatGPT auth is required to upload local files for Codex Apps tools".to_string() + })?; + let account_id = if authorization_header_value.starts_with("AgentAssertion ") { + None } else { - let token_data = auth - .get_token_data() - .map_err(|error| format!("failed to read ChatGPT auth for file upload: {error}"))?; - Box::new(BearerAuthProvider { - token: Some(token_data.access_token), - account_id: token_data.account_id, - is_fedramp_account: auth.is_fedramp_account(), - }) + auth.get_account_id() }; + let mut auth_provider = + AuthorizationHeaderAuthProvider::new(Some(authorization_header_value), account_id); + if auth.is_fedramp_account() { + auth_provider = auth_provider.with_fedramp_routing_header(); + } + let upload_auth: Box = Box::new(auth_provider); let uploaded = upload_local_file( turn_context.config.chatgpt_base_url.trim_end_matches('/'), upload_auth.as_ref(), @@ -168,17 +163,17 @@ async fn build_uploaded_local_argument_value( #[cfg(test)] mod tests { use super::*; - use crate::agent_identity::AgentIdentityManager; - use crate::agent_identity::RegisteredAgentTask; use crate::session::tests::make_session_and_context; use chrono::Utc; + use codex_agent_identity::generate_agent_key_material; + use codex_login::AgentIdentityAuthRecord; use codex_login::AuthCredentialsStoreMode; use codex_login::AuthDotJson; use codex_login::AuthManager; + use codex_login::BackgroundAgentTaskAuthMode; use codex_login::save_auth; use codex_login::token_data::IdTokenInfo; use codex_login::token_data::TokenData; - use codex_protocol::protocol::SessionSource; use codex_utils_absolute_path::AbsolutePathBuf; use pretty_assertions::assert_eq; use std::sync::Arc; @@ -194,9 +189,10 @@ mod tests { async fn install_cached_agent_task_auth( session: &mut Session, turn_context: &mut TurnContext, - chatgpt_base_url: String, + _chatgpt_base_url: String, ) { let auth_dir = tempdir().expect("temp auth dir"); + let key_material = generate_agent_key_material().expect("generate test key material"); let auth_json = AuthDotJson { auth_mode: Some(codex_app_server_protocol::AuthMode::Chatgpt), openai_api_key: None, @@ -214,7 +210,14 @@ mod tests { account_id: Some("account_id".to_string()), }), last_refresh: Some(Utc::now()), - agent_identity: None, + agent_identity: Some(AgentIdentityAuthRecord { + workspace_id: "account_id".to_string(), + chatgpt_user_id: None, + agent_runtime_id: "agent-123".to_string(), + agent_private_key: key_material.private_key_pkcs8_base64, + registered_at: "2026-04-15T00:00:00Z".to_string(), + background_task_id: Some("task-123".to_string()), + }), }; save_auth(auth_dir.path(), &auth_json, AuthCredentialsStoreMode::File) .expect("save test auth"); @@ -222,26 +225,12 @@ mod tests { .expect("load test auth") .expect("test auth"); let auth_manager = AuthManager::from_auth_for_testing(auth); - let agent_identity_manager = Arc::new(AgentIdentityManager::new_for_tests( - Arc::clone(&auth_manager), - /*feature_enabled*/ true, - chatgpt_base_url, - SessionSource::Exec, - )); - let stored_identity = agent_identity_manager - .seed_generated_identity_for_tests("agent-123") - .await - .expect("seed test identity"); + auth_manager.set_chatgpt_backend_auth_config( + Some("https://chatgpt.com/backend-api".to_string()), + BackgroundAgentTaskAuthMode::Enabled, + ); session.services.auth_manager = Arc::clone(&auth_manager); - session.services.agent_identity_manager = agent_identity_manager; turn_context.auth_manager = Some(auth_manager); - session - .cache_agent_task_for_tests(RegisteredAgentTask { - agent_runtime_id: stored_identity.agent_runtime_id, - task_id: "task-123".to_string(), - registered_at: "2026-04-15T00:00:00Z".to_string(), - }) - .await; } #[tokio::test] @@ -618,7 +607,6 @@ mod tests { .await; let (mut session, mut turn_context) = make_session_and_context().await; - let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); let dir = tempdir().expect("temp dir"); let local_path = dir.path().join("file_report.csv"); tokio::fs::write(&local_path, b"hello") @@ -630,6 +618,12 @@ mod tests { config.chatgpt_base_url = format!("{}/backend-api", server.uri()); turn_context.config = Arc::new(config); install_cached_agent_task_auth(&mut session, &mut turn_context, server.uri()).await; + let auth = session + .services + .auth_manager + .auth() + .await + .expect("test auth"); let rewritten = build_uploaded_local_argument_value( &session, diff --git a/codex-rs/core/src/session/agent_task_lifecycle.rs b/codex-rs/core/src/session/agent_task_lifecycle.rs deleted file mode 100644 index 72d0cf0459..0000000000 --- a/codex-rs/core/src/session/agent_task_lifecycle.rs +++ /dev/null @@ -1,218 +0,0 @@ -use crate::agent_identity::RegisteredAgentTask; -use crate::session::session::Session; -use codex_protocol::protocol::RolloutItem; -use codex_protocol::protocol::SessionAgentTask; -use codex_protocol::protocol::SessionStateUpdate; -use tracing::debug; -use tracing::info; -use tracing::warn; - -impl Session { - pub(super) async fn maybe_prewarm_agent_task_registration(&self) { - // Startup task registration is best-effort: regular turns already retry on demand, and - // a prewarm failure should not shut down the session or block unrelated work. - if let Err(error) = self.ensure_agent_task_registered().await { - warn!( - error = %error, - "startup agent task prewarm failed; regular turns will retry registration" - ); - } - } - - fn latest_persisted_agent_task( - rollout_items: &[RolloutItem], - ) -> Option> { - rollout_items.iter().rev().find_map(|item| match item { - RolloutItem::SessionState(update) => Some(update.agent_task.clone()), - _ => None, - }) - } - - pub(super) async fn restore_persisted_agent_task(&self, rollout_items: &[RolloutItem]) { - let Some(agent_task_update) = Self::latest_persisted_agent_task(rollout_items) else { - return; - }; - - match agent_task_update { - Some(agent_task) => { - let registered_task = - RegisteredAgentTask::from_session_agent_task(agent_task.clone()); - if self - .services - .agent_identity_manager - .task_matches_current_identity(®istered_task) - .await - { - let mut state = self.state.lock().await; - state.set_agent_task(agent_task); - } else { - debug!( - agent_runtime_id = %registered_task.agent_runtime_id, - task_id = %registered_task.task_id, - "discarding persisted agent task because it does not match the registered agent identity" - ); - let mut state = self.state.lock().await; - state.clear_agent_task(); - } - } - None => { - let mut state = self.state.lock().await; - state.clear_agent_task(); - } - } - } - - async fn persist_agent_task_update(&self, agent_task: Option<&RegisteredAgentTask>) { - self.persist_rollout_items(&[RolloutItem::SessionState(SessionStateUpdate { - agent_task: agent_task.map(RegisteredAgentTask::to_session_agent_task), - })]) - .await; - } - - async fn clear_cached_agent_task(&self, agent_task: &RegisteredAgentTask) { - let cleared = { - let mut state = self.state.lock().await; - if state.agent_task().as_ref() == Some(&agent_task.to_session_agent_task()) { - state.clear_agent_task(); - true - } else { - false - } - }; - if cleared { - self.persist_agent_task_update(/*agent_task*/ None).await; - } - } - - async fn cache_agent_task(&self, agent_task: RegisteredAgentTask) -> RegisteredAgentTask { - let session_agent_task = agent_task.to_session_agent_task(); - let changed = { - let mut state = self.state.lock().await; - if state.agent_task().as_ref() == Some(&session_agent_task) { - false - } else { - state.set_agent_task(session_agent_task); - true - } - }; - if changed { - self.persist_agent_task_update(Some(&agent_task)).await; - } - agent_task - } - - #[cfg(test)] - pub(crate) async fn cache_agent_task_for_tests(&self, agent_task: RegisteredAgentTask) { - self.cache_agent_task(agent_task).await; - } - - pub(super) async fn cached_agent_task_for_current_identity( - &self, - ) -> Option { - let agent_task = { - let state = self.state.lock().await; - state - .agent_task() - .map(RegisteredAgentTask::from_session_agent_task) - }?; - - if self - .services - .agent_identity_manager - .task_matches_current_identity(&agent_task) - .await - { - debug!( - agent_runtime_id = %agent_task.agent_runtime_id, - task_id = %agent_task.task_id, - "reusing cached agent task" - ); - return Some(agent_task); - } - - debug!( - agent_runtime_id = %agent_task.agent_runtime_id, - task_id = %agent_task.task_id, - "discarding cached agent task because the registered agent identity changed" - ); - self.clear_cached_agent_task(&agent_task).await; - None - } - - pub(crate) async fn authorization_header_for_current_agent_task( - &self, - ) -> anyhow::Result> { - let Some(agent_task) = self.cached_agent_task_for_current_identity().await else { - return Ok(None); - }; - - let Some(auth) = self.services.auth_manager.auth().await else { - return Ok(None); - }; - let authorization_header_value = self - .services - .auth_manager - .chatgpt_agent_task_authorization_header_for_auth( - &auth, - agent_task.authorization_target(), - )?; - if authorization_header_value.is_some() { - debug!( - agent_runtime_id = %agent_task.agent_runtime_id, - task_id = %agent_task.task_id, - "using agent assertion authorization for current task request" - ); - } - Ok(authorization_header_value) - } - - pub(super) async fn ensure_agent_task_registered( - &self, - ) -> anyhow::Result> { - if let Some(agent_task) = self.cached_agent_task_for_current_identity().await { - return Ok(Some(agent_task)); - } - - let _guard = self - .agent_task_registration_lock - .acquire() - .await - .map_err(|_| anyhow::anyhow!("agent task registration semaphore closed"))?; - if let Some(agent_task) = self.cached_agent_task_for_current_identity().await { - return Ok(Some(agent_task)); - } - - for _ in 0..2 { - let Some(agent_task) = self.services.agent_identity_manager.register_task().await? - else { - return Ok(None); - }; - - if !self - .services - .agent_identity_manager - .task_matches_current_identity(&agent_task) - .await - { - debug!( - agent_runtime_id = %agent_task.agent_runtime_id, - task_id = %agent_task.task_id, - "discarding newly registered agent task because the registered agent identity changed" - ); - continue; - } - - let agent_task = self.cache_agent_task(agent_task).await; - - info!( - thread_id = %self.conversation_id, - agent_runtime_id = %agent_task.agent_runtime_id, - task_id = %agent_task.task_id, - "registered agent task for thread" - ); - return Ok(Some(agent_task)); - } - - Ok(None) - } -} diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index f159546575..5bae0b96a6 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -13,7 +13,6 @@ use crate::agent::Mailbox; use crate::agent::MailboxReceiver; use crate::agent::agent_status_from_event; use crate::agent::status::is_final; -use crate::agent_identity::AgentIdentityManager; use crate::apps::render_apps_section; use crate::commit_attribution::commit_message_trailer_instruction; use crate::compact; @@ -164,7 +163,6 @@ use codex_protocol::error::Result as CodexResult; #[cfg(test)] use codex_protocol::exec_output::StreamOutput; -mod agent_task_lifecycle; mod handlers; mod mcp; mod review; @@ -974,58 +972,6 @@ impl Session { }); } - fn start_agent_identity_registration(self: &Arc) { - if !self.services.agent_identity_manager.is_enabled() { - return; - } - - let weak_sess = Arc::downgrade(self); - let mut auth_state_rx = self.services.auth_manager.subscribe_auth_state(); - tokio::spawn(async move { - loop { - let Some(sess) = weak_sess.upgrade() else { - return; - }; - match sess - .services - .agent_identity_manager - .ensure_registered_identity() - .await - { - Ok(Some(_)) => { - sess.maybe_prewarm_agent_task_registration().await; - return; - } - Ok(None) => { - drop(sess); - if auth_state_rx.changed().await.is_err() { - return; - } - } - Err(error) => { - sess.fail_agent_identity_registration(error).await; - return; - } - } - } - }); - } - - async fn fail_agent_identity_registration(self: &Arc, error: anyhow::Error) { - warn!(error = %error, "agent identity registration failed"); - let message = format!( - "Agent identity registration failed while `features.use_agent_identity` is enabled: {error}" - ); - self.send_event_raw(Event { - id: self.next_internal_sub_id(), - msg: EventMsg::Error(ErrorEvent { - message, - codex_error_info: Some(CodexErrorInfo::Other), - }), - }) - .await; - } - pub(crate) fn get_tx_event(&self) -> Sender { self.tx_event.clone() } @@ -1173,7 +1119,6 @@ impl Session { } InitialHistory::Resumed(resumed_history) => { let rollout_items = resumed_history.history; - self.restore_persisted_agent_task(&rollout_items).await; let previous_turn_settings = self .apply_rollout_reconstruction(&turn_context, &rollout_items) .await; diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index 3e407c4cd7..a4c042af0c 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/codex-rs/core/src/session/rollout_reconstruction.rs @@ -207,9 +207,7 @@ impl Session { active_segment.get_or_insert_with(ActiveReplaySegment::default); active_segment.counts_as_user_turn |= is_user_turn_boundary(response_item); } - RolloutItem::EventMsg(_) - | RolloutItem::SessionMeta(_) - | RolloutItem::SessionState(_) => {} + RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) => {} } if base_replacement_history.is_some() @@ -277,7 +275,6 @@ impl Session { history.drop_last_n_user_turns(rollback.num_turns); } RolloutItem::EventMsg(_) - | RolloutItem::SessionState(_) | RolloutItem::TurnContext(_) | RolloutItem::SessionMeta(_) => {} } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index a2080bbeab..ce2aaa58aa 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -26,7 +26,6 @@ pub(crate) struct Session { pub(crate) services: SessionServices, pub(super) js_repl: Arc, pub(super) next_internal_sub_id: AtomicU64, - pub(super) agent_task_registration_lock: Semaphore, } #[derive(Clone)] @@ -632,11 +631,6 @@ 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 @@ -659,7 +653,6 @@ impl Session { hooks, rollout: Mutex::new(rollout_recorder), user_shell: Arc::new(default_shell), - agent_identity_manager: Arc::clone(&agent_identity_manager), shell_snapshot_tx, show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, @@ -722,7 +715,6 @@ impl Session { services, js_repl, next_internal_sub_id: AtomicU64::new(0), - agent_task_registration_lock: Semaphore::new(/*permits*/ 1), }); if let Some(network_policy_decider_session) = network_policy_decider_session { let mut guard = network_policy_decider_session.write().await; @@ -849,7 +841,6 @@ impl Session { // record_initial_history can emit events. We record only after the SessionConfiguredEvent is emitted. sess.record_initial_history(initial_history).await; - sess.start_agent_identity_registration(); { let mut state = sess.state.lock().await; state.set_pending_session_start_source(Some(session_start_source)); diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 34e43f8f47..2a3e35be8e 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -1,6 +1,4 @@ use super::*; -use crate::agent_identity::RegisteredAgentTask; -use crate::agent_identity::StoredAgentIdentity; use crate::config::ConfigBuilder; use crate::config::test_config; use crate::config_loader::ConfigLayerStack; @@ -17,22 +15,9 @@ use crate::shell::default_user_shell; use crate::skills::SkillRenderSideEffects; use crate::skills::render::SkillMetadataBudget; use crate::tools::format_exec_output_str; -use base64::Engine as _; -use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; -use base64::engine::general_purpose::URL_SAFE_NO_PAD; -use chrono::SecondsFormat; -use chrono::Utc; -use codex_agent_identity::AgentBillOfMaterials; -use codex_agent_identity::curve25519_secret_key_from_private_key_pkcs8_base64; use codex_features::Feature; use codex_features::Features; -use codex_login::AgentIdentityAuthRecord; -use codex_login::AuthCredentialsStoreMode; -use codex_login::AuthDotJson; use codex_login::CodexAuth; -use codex_login::save_auth; -use codex_login::token_data::IdTokenInfo; -use codex_login::token_data::TokenData; use codex_model_provider_info::ModelProviderInfo; use codex_models_manager::bundled_models_response; use codex_models_manager::model_info; @@ -131,8 +116,6 @@ use core_test_support::test_codex::test_codex; use core_test_support::test_path_buf; use core_test_support::tracing::install_test_tracing; use core_test_support::wait_for_event; -use ed25519_dalek::SigningKey; -use ed25519_dalek::pkcs8::EncodePrivateKey; use opentelemetry::trace::TraceContextExt; use opentelemetry::trace::TraceId; use opentelemetry_sdk::metrics::InMemoryMetricExporter; @@ -146,12 +129,6 @@ use tokio::sync::Semaphore; use tokio::time::sleep; use tokio::time::timeout; use tracing_opentelemetry::OpenTelemetrySpanExt; -use wiremock::Mock; -use wiremock::MockServer; -use wiremock::ResponseTemplate; -use wiremock::matchers::header; -use wiremock::matchers::method; -use wiremock::matchers::path; use codex_protocol::mcp::CallToolResult as McpCallToolResult; use pretty_assertions::assert_eq; @@ -1290,120 +1267,6 @@ async fn record_initial_history_reconstructs_resumed_transcript() { assert_eq!(expected, history.raw_items()); } -#[tokio::test] -async fn record_initial_history_restores_latest_persisted_agent_task() { - let auth = make_chatgpt_auth("account-123", Some("user-123")); - seed_stored_identity(&auth, "agent-123", "account-123"); - let (session, _turn_context, _rx_event) = make_agent_identity_session_and_context_with_rx( - auth, - "https://chatgpt.com/backend-api".to_string(), - ) - .await; - let expected = RegisteredAgentTask { - agent_runtime_id: "agent-123".to_string(), - task_id: "task-123".to_string(), - registered_at: "2026-03-23T12:00:00Z".to_string(), - }; - let rollout_items = vec![ - RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate { - agent_task: Some(expected.to_session_agent_task()), - }), - RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate { - agent_task: None, - }), - RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate { - agent_task: Some(expected.to_session_agent_task()), - }), - ]; - - session - .record_initial_history(InitialHistory::Resumed(ResumedHistory { - conversation_id: ThreadId::default(), - history: rollout_items, - rollout_path: PathBuf::from("/tmp/resume.jsonl"), - })) - .await; - - assert_eq!( - session.state.lock().await.agent_task(), - Some(expected.to_session_agent_task()) - ); -} - -#[tokio::test] -async fn record_initial_history_discards_persisted_agent_task_for_different_identity() { - let auth = make_chatgpt_auth("account-123", Some("user-123")); - seed_stored_identity(&auth, "agent-123", "account-123"); - let (session, _turn_context, _rx_event) = make_agent_identity_session_and_context_with_rx( - auth, - "https://chatgpt.com/backend-api".to_string(), - ) - .await; - let rollout_items = vec![RolloutItem::SessionState( - codex_protocol::protocol::SessionStateUpdate { - agent_task: Some( - RegisteredAgentTask { - agent_runtime_id: "agent-other".to_string(), - task_id: "task-other".to_string(), - registered_at: "2026-03-23T12:00:00Z".to_string(), - } - .to_session_agent_task(), - ), - }, - )]; - - session - .record_initial_history(InitialHistory::Resumed(ResumedHistory { - conversation_id: ThreadId::default(), - history: rollout_items, - rollout_path: PathBuf::from("/tmp/resume.jsonl"), - })) - .await; - - assert_eq!(session.state.lock().await.agent_task(), None); -} - -#[tokio::test] -async fn record_initial_history_honors_cleared_persisted_agent_task() { - let (session, _turn_context) = make_session_and_context().await; - { - let mut state = session.state.lock().await; - state.set_agent_task( - RegisteredAgentTask { - agent_runtime_id: "agent-fresh".to_string(), - task_id: "task-fresh".to_string(), - registered_at: "2026-03-23T12:01:00Z".to_string(), - } - .to_session_agent_task(), - ); - } - let rollout_items = vec![ - RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate { - agent_task: Some( - RegisteredAgentTask { - agent_runtime_id: "agent-123".to_string(), - task_id: "task-123".to_string(), - registered_at: "2026-03-23T12:00:00Z".to_string(), - } - .to_session_agent_task(), - ), - }), - RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate { - agent_task: None, - }), - ]; - - session - .record_initial_history(InitialHistory::Resumed(ResumedHistory { - conversation_id: ThreadId::default(), - history: rollout_items, - rollout_path: PathBuf::from("/tmp/resume.jsonl"), - })) - .await; - - assert_eq!(session.state.lock().await.agent_task(), None); -} - #[tokio::test] async fn record_initial_history_new_defers_initial_context_until_first_turn() { let (session, _turn_context) = make_session_and_context().await; @@ -3208,11 +3071,6 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { }), rollout: Mutex::new(None), user_shell: Arc::new(default_user_shell()), - agent_identity_manager: Arc::new(crate::agent_identity::AgentIdentityManager::new( - config.as_ref(), - Arc::clone(&auth_manager), - session_configuration.session_source.clone(), - )), shell_snapshot_tx: watch::channel(None).0, show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, @@ -3305,7 +3163,6 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { services, js_repl, next_internal_sub_id: AtomicU64::new(0), - agent_task_registration_lock: Semaphore::new(/*permits*/ 1), }; (session, turn_context) @@ -4178,11 +4035,6 @@ where }), rollout: Mutex::new(None), user_shell: Arc::new(default_user_shell()), - agent_identity_manager: Arc::new(crate::agent_identity::AgentIdentityManager::new( - config.as_ref(), - Arc::clone(&auth_manager), - session_configuration.session_source.clone(), - )), shell_snapshot_tx: watch::channel(None).0, show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, @@ -4275,7 +4127,6 @@ where services, js_repl, next_internal_sub_id: AtomicU64::new(0), - agent_task_registration_lock: Semaphore::new(/*permits*/ 1), }); (session, turn_context, rx_event) @@ -4296,24 +4147,6 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( .await } -async fn make_agent_identity_session_and_context_with_rx( - auth: CodexAuth, - chatgpt_base_url: String, -) -> ( - Arc, - Arc, - async_channel::Receiver, -) { - make_session_and_context_with_auth_and_config_and_rx(auth, Vec::new(), move |config| { - config.chatgpt_base_url = chatgpt_base_url; - config - .features - .enable(Feature::UseAgentIdentity) - .expect("test config should allow use_agent_identity"); - }) - .await -} - // Like make_session_and_context, but returns Arc and the event receiver // so tests can assert on emitted events. pub(crate) async fn make_session_and_context_with_rx() -> ( @@ -4324,240 +4157,6 @@ pub(crate) async fn make_session_and_context_with_rx() -> ( make_session_and_context_with_dynamic_tools_and_rx(Vec::new()).await } -#[tokio::test] -async fn fail_agent_identity_registration_emits_error_without_shutdown() { - let (session, _turn_context, rx_event) = make_session_and_context_with_rx().await; - - session - .fail_agent_identity_registration(anyhow::anyhow!("registration exploded")) - .await; - - let error_event = timeout(Duration::from_secs(1), rx_event.recv()) - .await - .expect("error event should arrive") - .expect("error event should be readable"); - match error_event.msg { - EventMsg::Error(ErrorEvent { - message, - codex_error_info, - }) => { - assert_eq!( - message, - "Agent identity registration failed while `features.use_agent_identity` is enabled: registration exploded".to_string() - ); - assert_eq!(codex_error_info, Some(CodexErrorInfo::Other)); - } - other => panic!("expected error event, got {other:?}"), - } - - assert!(rx_event.try_recv().is_err()); -} - -#[tokio::test] -async fn startup_agent_task_prewarm_caches_registered_task() { - let server = MockServer::start().await; - let chatgpt_base_url = server.uri(); - let auth = make_chatgpt_auth("account-123", Some("user-123")); - let stored_identity = seed_stored_identity(&auth, "agent-123", "account-123"); - let encrypted_task_id = - encrypt_task_id_for_identity(&stored_identity, "task_123").expect("task ciphertext"); - mount_human_biscuit(&server, &chatgpt_base_url, "agent-123").await; - Mock::given(method("POST")) - .and(path("/v1/agent/agent-123/task/register")) - .and(header("x-openai-authorization", "human-biscuit")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "encrypted_task_id": encrypted_task_id, - }))) - .expect(1) - .mount(&server) - .await; - - let (session, _turn_context, rx_event) = - make_agent_identity_session_and_context_with_rx(auth, chatgpt_base_url).await; - - session.maybe_prewarm_agent_task_registration().await; - - let cached_task = session - .state - .lock() - .await - .agent_task() - .expect("task should be cached"); - assert_eq!(cached_task.agent_runtime_id, "agent-123"); - assert_eq!(cached_task.task_id, "task_123"); - assert!(rx_event.try_recv().is_err()); -} - -#[tokio::test] -async fn startup_agent_task_prewarm_failure_does_not_emit_error() { - let server = MockServer::start().await; - let chatgpt_base_url = server.uri(); - let auth = make_chatgpt_auth("account-123", Some("user-123")); - seed_stored_identity(&auth, "agent-123", "account-123"); - mount_human_biscuit(&server, &chatgpt_base_url, "agent-123").await; - Mock::given(method("POST")) - .and(path("/v1/agent/agent-123/task/register")) - .and(header("x-openai-authorization", "human-biscuit")) - .respond_with(ResponseTemplate::new(500)) - .expect(1) - .mount(&server) - .await; - - let (session, _turn_context, rx_event) = - make_agent_identity_session_and_context_with_rx(auth, chatgpt_base_url).await; - - session.maybe_prewarm_agent_task_registration().await; - - assert_eq!(session.state.lock().await.agent_task(), None); - assert!(rx_event.try_recv().is_err()); -} - -#[tokio::test] -async fn cached_agent_task_for_current_identity_clears_stale_task() { - let auth = make_chatgpt_auth("account-123", Some("user-123")); - seed_stored_identity(&auth, "agent-123", "account-123"); - let (session, _turn_context, _rx_event) = make_agent_identity_session_and_context_with_rx( - auth, - "https://chatgpt.com/backend-api".to_string(), - ) - .await; - { - let mut state = session.state.lock().await; - state.set_agent_task( - RegisteredAgentTask { - agent_runtime_id: "agent-old".to_string(), - task_id: "task-old".to_string(), - registered_at: "2026-04-15T00:00:00Z".to_string(), - } - .to_session_agent_task(), - ); - } - - assert_eq!(session.cached_agent_task_for_current_identity().await, None); - assert_eq!(session.state.lock().await.agent_task(), None); -} - -fn seed_stored_identity( - auth: &CodexAuth, - agent_runtime_id: &str, - account_id: &str, -) -> StoredAgentIdentity { - let signing_key = generate_test_signing_key(); - let private_key_pkcs8 = signing_key - .to_pkcs8_der() - .expect("encode test signing key as PKCS#8"); - let stored_identity = StoredAgentIdentity { - binding_id: format!("chatgpt-account-{account_id}"), - chatgpt_account_id: account_id.to_string(), - chatgpt_user_id: Some("user-123".to_string()), - agent_runtime_id: agent_runtime_id.to_string(), - private_key_pkcs8_base64: BASE64_STANDARD.encode(private_key_pkcs8.as_bytes()), - public_key_ssh: "ssh-ed25519 test".to_string(), - registered_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), - abom: AgentBillOfMaterials { - agent_version: env!("CARGO_PKG_VERSION").to_string(), - agent_harness_id: "codex-cli".to_string(), - running_location: format!("{}-{}", SessionSource::Exec, std::env::consts::OS), - }, - }; - - auth.set_agent_identity(AgentIdentityAuthRecord { - workspace_id: account_id.to_string(), - chatgpt_user_id: stored_identity.chatgpt_user_id.clone(), - agent_runtime_id: stored_identity.agent_runtime_id.clone(), - agent_private_key: stored_identity.private_key_pkcs8_base64.clone(), - registered_at: stored_identity.registered_at.clone(), - background_task_id: None, - }) - .expect("store identity"); - - stored_identity -} - -fn encrypt_task_id_for_identity( - stored_identity: &StoredAgentIdentity, - task_id: &str, -) -> anyhow::Result { - let mut rng = crypto_box::aead::OsRng; - let public_key = curve25519_secret_key_from_private_key_pkcs8_base64( - &stored_identity.private_key_pkcs8_base64, - )? - .public_key(); - let ciphertext = public_key - .seal(&mut rng, task_id.as_bytes()) - .map_err(|_| anyhow::anyhow!("failed to encrypt test task id"))?; - Ok(BASE64_STANDARD.encode(ciphertext)) -} - -fn generate_test_signing_key() -> SigningKey { - SigningKey::from_bytes(&[7u8; 32]) -} - -async fn mount_human_biscuit(server: &MockServer, chatgpt_base_url: &str, agent_runtime_id: &str) { - let biscuit_url = format!( - "{}/authenticate_app_v2", - chatgpt_base_url.trim_end_matches('/') - ); - let biscuit_path = reqwest::Url::parse(&biscuit_url) - .expect("biscuit URL parses") - .path() - .to_string(); - let target_url = format!( - "{}/v1/agent/{agent_runtime_id}/task/register", - chatgpt_base_url.trim_end_matches('/') - ); - Mock::given(method("GET")) - .and(path(biscuit_path)) - .and(header("authorization", "Bearer access-token-account-123")) - .and(header("x-original-method", "POST")) - .and(header("x-original-url", target_url)) - .respond_with( - ResponseTemplate::new(200).insert_header("x-openai-authorization", "human-biscuit"), - ) - .expect(1) - .mount(server) - .await; -} - -fn make_chatgpt_auth(account_id: &str, user_id: Option<&str>) -> CodexAuth { - let tempdir = tempfile::tempdir().expect("tempdir"); - let auth_json = AuthDotJson { - auth_mode: Some(codex_app_server_protocol::AuthMode::Chatgpt), - openai_api_key: None, - tokens: Some(TokenData { - id_token: IdTokenInfo { - email: None, - chatgpt_plan_type: None, - chatgpt_user_id: user_id.map(ToOwned::to_owned), - chatgpt_account_id: Some(account_id.to_string()), - chatgpt_account_is_fedramp: false, - 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(Utc::now()), - agent_identity: None, - }; - save_auth(tempdir.path(), &auth_json, AuthCredentialsStoreMode::File).expect("save auth"); - CodexAuth::from_auth_storage(tempdir.path(), 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") -} - #[tokio::test] async fn refresh_mcp_servers_is_deferred_until_next_turn() { let (session, turn_context) = make_session_and_context().await; diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index efbcca6716..adfaa783a9 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -333,24 +333,6 @@ pub(crate) async fn run_turn( })) .await; } - 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.send_event( - turn_context.as_ref(), - EventMsg::Error(ErrorEvent { - message: format!( - "Agent task registration failed. Please try again; Codex will attempt to register the task again on the next turn: {error}" - ), - codex_error_info: Some(CodexErrorInfo::Other), - }), - ) - .await; - return None; - } - }; - if !skill_items.is_empty() { sess.record_conversation_items(&turn_context, &skill_items) .await; @@ -375,16 +357,8 @@ 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 prewarmed_client_session = prewarmed_client_session; - if agent_task.is_some() - && let Some(prewarmed_client_session) = prewarmed_client_session.as_mut() + let mut client_session = if let Some(prewarmed_client_session) = prewarmed_client_session.take() { - 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() diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 5db38f7b72..aae02d61bd 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -4,7 +4,6 @@ use std::sync::Arc; use crate::RolloutRecorder; use crate::SkillsManager; use crate::agent::AgentControl; -use crate::agent_identity::AgentIdentityManager; use crate::client::ModelClient; use crate::config::StartedNetworkProxy; use crate::exec_policy::ExecPolicyManager; @@ -43,7 +42,6 @@ pub(crate) struct SessionServices { pub(crate) hooks: Hooks, pub(crate) rollout: Mutex>, pub(crate) user_shell: Arc, - pub(crate) agent_identity_manager: Arc, pub(crate) shell_snapshot_tx: watch::Sender>>, pub(crate) show_raw_agent_reasoning: bool, pub(crate) exec_policy: Arc, diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 4ac8c4d530..e9fd676332 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -11,7 +11,6 @@ use crate::session::PreviousTurnSettings; use crate::session::session::SessionConfiguration; use crate::session_startup_prewarm::SessionStartupPrewarmHandle; use codex_protocol::protocol::RateLimitSnapshot; -use codex_protocol::protocol::SessionAgentTask; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; use codex_protocol::protocol::TurnContextItem; @@ -31,7 +30,6 @@ pub(crate) struct SessionState { previous_turn_settings: Option, /// Startup prewarmed session prepared during session initialization. pub(crate) startup_prewarm: Option, - pub(crate) agent_task: Option, pub(crate) active_connector_selection: HashSet, pub(crate) pending_session_start_source: Option, granted_permissions: Option, @@ -51,7 +49,6 @@ impl SessionState { mcp_dependency_prompted: HashSet::new(), previous_turn_settings: None, startup_prewarm: None, - agent_task: None, active_connector_selection: HashSet::new(), pending_session_start_source: None, granted_permissions: None, @@ -189,18 +186,6 @@ impl SessionState { self.startup_prewarm.take() } - pub(crate) fn agent_task(&self) -> Option { - self.agent_task.clone() - } - - pub(crate) fn set_agent_task(&mut self, agent_task: SessionAgentTask) { - self.agent_task = Some(agent_task); - } - - pub(crate) fn clear_agent_task(&mut self) { - self.agent_task = None; - } - // Adds connector IDs to the active set and returns the merged selection. pub(crate) fn merge_connector_selection(&mut self, connector_ids: I) -> HashSet where diff --git a/codex-rs/core/src/state/session_tests.rs b/codex-rs/core/src/state/session_tests.rs index 12229ee2f2..5e90cc881d 100644 --- a/codex-rs/core/src/state/session_tests.rs +++ b/codex-rs/core/src/state/session_tests.rs @@ -2,7 +2,6 @@ use super::*; use crate::session::tests::make_session_configuration_for_tests; use codex_protocol::protocol::CreditsSnapshot; use codex_protocol::protocol::RateLimitWindow; -use codex_protocol::protocol::SessionAgentTask; use pretty_assertions::assert_eq; #[tokio::test] @@ -34,37 +33,6 @@ async fn clear_connector_selection_removes_entries() { assert_eq!(state.get_connector_selection(), HashSet::new()); } -#[tokio::test] -async fn set_agent_task_persists_plaintext_task_for_session_reuse() { - let session_configuration = make_session_configuration_for_tests().await; - let mut state = SessionState::new(session_configuration); - let agent_task = SessionAgentTask { - agent_runtime_id: "agent-123".to_string(), - task_id: "task_123".to_string(), - registered_at: "2026-03-23T12:00:00Z".to_string(), - }; - - state.set_agent_task(agent_task.clone()); - - assert_eq!(state.agent_task(), Some(agent_task)); -} - -#[tokio::test] -async fn clear_agent_task_removes_cached_task() { - let session_configuration = make_session_configuration_for_tests().await; - let mut state = SessionState::new(session_configuration); - let agent_task = SessionAgentTask { - agent_runtime_id: "agent_123".to_string(), - task_id: "task_123".to_string(), - registered_at: "2026-03-23T12:00:00Z".to_string(), - }; - - state.set_agent_task(agent_task); - state.clear_agent_task(); - - assert_eq!(state.agent_task(), None); -} - #[tokio::test] async fn set_rate_limits_defaults_limit_id_to_codex_when_missing() { let session_configuration = make_session_configuration_for_tests().await; diff --git a/codex-rs/login/src/agent_identity.rs b/codex-rs/login/src/agent_identity.rs index e85c406399..cd92d201ec 100644 --- a/codex-rs/login/src/agent_identity.rs +++ b/codex-rs/login/src/agent_identity.rs @@ -1,48 +1,27 @@ use std::sync::Arc; -use std::time::Duration; -use anyhow::Context; use anyhow::Result; -use chrono::SecondsFormat; -use chrono::Utc; -use codex_agent_identity::AgentBillOfMaterials; use codex_agent_identity::AgentIdentityKey; -use codex_agent_identity::agent_identity_biscuit_url; -use codex_agent_identity::agent_identity_request_id; -use codex_agent_identity::agent_registration_url; -use codex_agent_identity::agent_task_registration_url; use codex_agent_identity::authorization_header_for_agent_task; -use codex_agent_identity::build_abom; -use codex_agent_identity::decrypt_task_id_response; +#[cfg(test)] use codex_agent_identity::generate_agent_key_material; use codex_agent_identity::normalize_chatgpt_base_url; use codex_agent_identity::public_key_ssh_from_private_key_pkcs8_base64; -use codex_agent_identity::sign_task_registration_payload; use codex_agent_identity::supports_background_agent_task_auth; +#[cfg(test)] use codex_protocol::protocol::SessionSource; -use serde::Deserialize; -use serde::Serialize; -use tokio::sync::Semaphore; use tracing::debug; -use tracing::info; use tracing::warn; use crate::AgentIdentityAuthRecord; use crate::AuthManager; use crate::CodexAuth; -use crate::default_client::create_client; - -const AGENT_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15); -const AGENT_TASK_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15); -const AGENT_IDENTITY_BISCUIT_TIMEOUT: Duration = Duration::from_secs(15); #[derive(Clone)] pub(crate) struct BackgroundAgentTaskManager { auth_manager: Arc, chatgpt_base_url: String, auth_mode: BackgroundAgentTaskAuthMode, - abom: AgentBillOfMaterials, - ensure_lock: Arc, } impl std::fmt::Debug for BackgroundAgentTaskManager { @@ -50,7 +29,6 @@ impl std::fmt::Debug for BackgroundAgentTaskManager { f.debug_struct("BackgroundAgentTaskManager") .field("chatgpt_base_url", &self.chatgpt_base_url) .field("auth_mode", &self.auth_mode) - .field("abom", &self.abom) .finish_non_exhaustive() } } @@ -76,7 +54,7 @@ impl BackgroundAgentTaskAuthMode { } } -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq)] struct StoredAgentIdentity { binding_id: String, chatgpt_account_id: String, @@ -86,30 +64,6 @@ struct StoredAgentIdentity { public_key_ssh: String, registered_at: String, background_task_id: Option, - abom: AgentBillOfMaterials, -} - -#[derive(Debug, Serialize)] -struct RegisterAgentRequest { - abom: AgentBillOfMaterials, - agent_public_key: String, - capabilities: Vec, -} - -#[derive(Debug, Deserialize)] -struct RegisterAgentResponse { - agent_runtime_id: String, -} - -#[derive(Debug, Serialize)] -struct RegisterTaskRequest { - signature: String, - timestamp: String, -} - -#[derive(Debug, Deserialize)] -struct RegisterTaskResponse { - encrypted_task_id: String, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -117,7 +71,6 @@ struct AgentIdentityBinding { binding_id: String, chatgpt_account_id: String, chatgpt_user_id: Option, - access_token: String, } impl BackgroundAgentTaskManager { @@ -125,12 +78,11 @@ impl BackgroundAgentTaskManager { pub(crate) fn new( auth_manager: Arc, chatgpt_base_url: String, - session_source: SessionSource, + _session_source: SessionSource, ) -> Self { Self::new_with_auth_mode( auth_manager, chatgpt_base_url, - session_source, BackgroundAgentTaskAuthMode::Disabled, ) } @@ -138,15 +90,12 @@ impl BackgroundAgentTaskManager { pub(crate) fn new_with_auth_mode( auth_manager: Arc, chatgpt_base_url: String, - session_source: SessionSource, auth_mode: BackgroundAgentTaskAuthMode, ) -> Self { Self { auth_manager, chatgpt_base_url: normalize_chatgpt_base_url(&chatgpt_base_url), auth_mode, - abom: build_abom(session_source), - ensure_lock: Arc::new(Semaphore::new(/*permits*/ 1)), } } @@ -174,29 +123,20 @@ impl BackgroundAgentTaskManager { return Ok(None); }; - let _guard = self - .ensure_lock - .acquire() - .await - .context("background agent task ensure semaphore closed")?; - let mut stored_identity = self - .ensure_registered_identity_for_binding(auth, &binding) - .await?; - let background_task_id = match stored_identity.background_task_id.clone() { - Some(background_task_id) => background_task_id, - _ => { - let background_task_id = self - .register_background_task_for_identity(&binding, &stored_identity) - .await?; - stored_identity.background_task_id = Some(background_task_id.clone()); - self.store_identity(auth, &stored_identity)?; - background_task_id - } + let Some(stored_identity) = self.load_stored_identity(auth, &binding)? else { + return Ok(None); + }; + let Some(background_task_id) = stored_identity.background_task_id.as_ref() else { + debug!( + agent_runtime_id = %stored_identity.agent_runtime_id, + "skipping background agent task auth because stored agent identity has no background task id" + ); + return Ok(None); }; Ok(Some(authorization_header_for_task( &stored_identity, - &background_task_id, + background_task_id, )?)) } @@ -224,167 +164,6 @@ impl BackgroundAgentTaskManager { } } - async fn ensure_registered_identity_for_binding( - &self, - auth: &CodexAuth, - binding: &AgentIdentityBinding, - ) -> Result { - if let Some(stored_identity) = self.load_stored_identity(auth, binding)? { - return Ok(stored_identity); - } - - let stored_identity = self.register_agent_identity(binding).await?; - self.store_identity(auth, &stored_identity)?; - Ok(stored_identity) - } - - async fn register_agent_identity( - &self, - binding: &AgentIdentityBinding, - ) -> Result { - let key_material = generate_agent_key_material()?; - let request_body = RegisterAgentRequest { - abom: self.abom.clone(), - agent_public_key: key_material.public_key_ssh.clone(), - capabilities: Vec::new(), - }; - - let url = agent_registration_url(&self.chatgpt_base_url); - let human_biscuit = self.mint_human_biscuit(binding, "POST", &url).await?; - let client = create_client(); - let response = client - .post(&url) - .header("X-OpenAI-Authorization", human_biscuit) - .json(&request_body) - .timeout(AGENT_REGISTRATION_TIMEOUT) - .send() - .await - .with_context(|| { - format!("failed to send agent identity registration request to {url}") - })?; - - if response.status().is_success() { - let response_body = response - .json::() - .await - .with_context(|| format!("failed to parse agent identity response from {url}"))?; - 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: response_body.agent_runtime_id, - 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), - background_task_id: None, - abom: self.abom.clone(), - }; - info!( - agent_runtime_id = %stored_identity.agent_runtime_id, - binding_id = %binding.binding_id, - "registered background agent identity" - ); - return Ok(stored_identity); - } - - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - anyhow::bail!("agent identity registration failed with status {status} from {url}: {body}") - } - - async fn register_background_task_for_identity( - &self, - binding: &AgentIdentityBinding, - stored_identity: &StoredAgentIdentity, - ) -> Result { - let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); - let request_body = RegisterTaskRequest { - signature: sign_task_registration_payload( - stored_identity.agent_identity_key(), - ×tamp, - )?, - timestamp, - }; - - let client = create_client(); - let url = - agent_task_registration_url(&self.chatgpt_base_url, &stored_identity.agent_runtime_id); - let human_biscuit = self.mint_human_biscuit(binding, "POST", &url).await?; - let response = client - .post(&url) - .header("X-OpenAI-Authorization", human_biscuit) - .json(&request_body) - .timeout(AGENT_TASK_REGISTRATION_TIMEOUT) - .send() - .await - .with_context(|| format!("failed to send background agent task request to {url}"))?; - - if response.status().is_success() { - let response_body = response - .json::() - .await - .with_context(|| format!("failed to parse background task response from {url}"))?; - let background_task_id = decrypt_task_id_response( - stored_identity.agent_identity_key(), - &response_body.encrypted_task_id, - )?; - info!( - agent_runtime_id = %stored_identity.agent_runtime_id, - task_id = %background_task_id, - "registered background agent task" - ); - return Ok(background_task_id); - } - - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - anyhow::bail!( - "background agent task registration failed with status {status} from {url}: {body}" - ) - } - - async fn mint_human_biscuit( - &self, - binding: &AgentIdentityBinding, - target_method: &str, - target_url: &str, - ) -> Result { - let url = agent_identity_biscuit_url(&self.chatgpt_base_url); - let request_id = agent_identity_request_id()?; - let client = create_client(); - let response = client - .get(&url) - .bearer_auth(&binding.access_token) - .header("X-Request-Id", request_id.clone()) - .header("X-Original-Method", target_method) - .header("X-Original-Url", target_url) - .timeout(AGENT_IDENTITY_BISCUIT_TIMEOUT) - .send() - .await - .with_context(|| format!("failed to send agent identity biscuit request to {url}"))?; - - if response.status().is_success() { - let human_biscuit = response - .headers() - .get("x-openai-authorization") - .context("agent identity biscuit response did not include x-openai-authorization")? - .to_str() - .context("agent identity biscuit response header was not valid UTF-8")? - .to_string(); - info!( - request_id = %request_id, - "minted human biscuit for background agent task" - ); - return Ok(human_biscuit); - } - - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - anyhow::bail!( - "agent identity biscuit minting failed with status {status} from {url}: {body}" - ) - } - fn load_stored_identity( &self, auth: &CodexAuth, @@ -394,19 +173,18 @@ impl BackgroundAgentTaskManager { return Ok(None); }; - let stored_identity = - match StoredAgentIdentity::from_auth_record(binding, record, self.abom.clone()) { - Ok(stored_identity) => stored_identity, - Err(error) => { - warn!( - binding_id = %binding.binding_id, - error = %error, - "stored agent identity is invalid; deleting cached value" - ); - auth.remove_agent_identity()?; - return Ok(None); - } - }; + let stored_identity = match StoredAgentIdentity::from_auth_record(binding, record) { + Ok(stored_identity) => stored_identity, + Err(error) => { + warn!( + binding_id = %binding.binding_id, + error = %error, + "stored agent identity is invalid; deleting cached value" + ); + auth.remove_agent_identity()?; + return Ok(None); + } + }; if !stored_identity.matches_binding(binding) { warn!( @@ -430,15 +208,6 @@ impl BackgroundAgentTaskManager { Ok(Some(stored_identity)) } - - fn store_identity( - &self, - auth: &CodexAuth, - stored_identity: &StoredAgentIdentity, - ) -> Result<()> { - auth.set_agent_identity(stored_identity.to_auth_record())?; - Ok(()) - } } pub fn cached_background_agent_task_authorization_header_value( @@ -455,8 +224,7 @@ pub fn cached_background_agent_task_authorization_header_value( let Some(record) = auth.get_agent_identity(&binding.chatgpt_account_id) else { return Ok(None); }; - let stored_identity = - StoredAgentIdentity::from_auth_record(&binding, record, build_abom(SessionSource::Cli))?; + let stored_identity = StoredAgentIdentity::from_auth_record(&binding, record)?; if !stored_identity.matches_binding(&binding) { return Ok(None); } @@ -471,7 +239,6 @@ impl StoredAgentIdentity { fn from_auth_record( binding: &AgentIdentityBinding, record: AgentIdentityAuthRecord, - abom: AgentBillOfMaterials, ) -> Result { if record.workspace_id != binding.chatgpt_account_id { anyhow::bail!( @@ -491,21 +258,9 @@ impl StoredAgentIdentity { public_key_ssh, registered_at: record.registered_at, background_task_id: record.background_task_id, - abom, }) } - fn to_auth_record(&self) -> AgentIdentityAuthRecord { - AgentIdentityAuthRecord { - workspace_id: self.chatgpt_account_id.clone(), - chatgpt_user_id: self.chatgpt_user_id.clone(), - agent_runtime_id: self.agent_runtime_id.clone(), - agent_private_key: self.private_key_pkcs8_base64.clone(), - registered_at: self.registered_at.clone(), - background_task_id: self.background_task_id.clone(), - } - } - fn matches_binding(&self, binding: &AgentIdentityBinding) -> bool { binding.matches_parts( &self.binding_id, @@ -568,7 +323,6 @@ impl AgentIdentityBinding { .id_token .chatgpt_user_id .filter(|value| !value.is_empty()), - access_token: token_data.access_token, }) } } @@ -597,7 +351,6 @@ mod tests { let manager = BackgroundAgentTaskManager::new_with_auth_mode( auth_manager, "https://chatgpt.com/backend-api".to_string(), - SessionSource::Cli, BackgroundAgentTaskAuthMode::Disabled, ); diff --git a/codex-rs/login/src/auth/agent_assertion.rs b/codex-rs/login/src/auth/agent_assertion.rs deleted file mode 100644 index 4f6be2fccd..0000000000 --- a/codex-rs/login/src/auth/agent_assertion.rs +++ /dev/null @@ -1,18 +0,0 @@ -use anyhow::Result; -use codex_agent_identity::AgentIdentityKey; -pub use codex_agent_identity::AgentTaskAuthorizationTarget; - -use super::storage::AgentIdentityAuthRecord; - -pub(super) fn authorization_header_for_agent_task( - record: &AgentIdentityAuthRecord, - target: AgentTaskAuthorizationTarget<'_>, -) -> Result { - codex_agent_identity::authorization_header_for_agent_task( - AgentIdentityKey { - agent_runtime_id: &record.agent_runtime_id, - private_key_pkcs8_base64: &record.agent_private_key, - }, - target, - ) -} diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 29733f13da..4db9f663a7 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -1,4 +1,3 @@ -use anyhow::Context; use async_trait::async_trait; use chrono::Utc; use reqwest::StatusCode; @@ -23,8 +22,6 @@ use codex_app_server_protocol::AuthMode as ApiAuthMode; use codex_protocol::config_types::ForcedLoginMethod; use codex_protocol::config_types::ModelProviderAuthInfo; -use super::agent_assertion; -use super::agent_assertion::AgentTaskAuthorizationTarget; use super::external_bearer::BearerTokenRefresher; use super::revoke::revoke_auth_tokens; pub use crate::auth::storage::AgentIdentityAuthRecord; @@ -43,7 +40,6 @@ use codex_protocol::auth::KnownPlan as InternalKnownPlan; use codex_protocol::auth::PlanType as InternalPlanType; use codex_protocol::auth::RefreshTokenFailedError; use codex_protocol::auth::RefreshTokenFailedReason; -use codex_protocol::protocol::SessionSource; use serde_json::Value; use thiserror::Error; @@ -1528,43 +1524,6 @@ impl AuthManager { .and_then(|guard| guard.clone()) } - pub fn chatgpt_agent_task_authorization_header_for_auth( - &self, - auth: &CodexAuth, - target: AgentTaskAuthorizationTarget<'_>, - ) -> anyhow::Result> { - let Some(record) = self.agent_identity_for_chatgpt_auth(auth)? else { - return Ok(None); - }; - agent_assertion::authorization_header_for_agent_task(&record, target).map(Some) - } - - fn agent_identity_for_chatgpt_auth( - &self, - auth: &CodexAuth, - ) -> anyhow::Result> { - if !auth.is_chatgpt_auth() { - return Ok(None); - } - - let token_data = auth - .get_token_data() - .context("ChatGPT token data is not available")?; - let workspace_id = self - .forced_chatgpt_workspace_id() - .filter(|value| !value.is_empty()) - .or(token_data.account_id.filter(|value| !value.is_empty())); - - let Some(workspace_id) = workspace_id else { - return Ok(None); - }; - let Some(record) = auth.get_agent_identity(&workspace_id) else { - anyhow::bail!("agent identity is not available for workspace {workspace_id}"); - }; - - Ok(Some(record)) - } - pub fn set_chatgpt_backend_auth_config( &self, chatgpt_base_url: Option, @@ -1634,7 +1593,6 @@ impl AuthManager { BackgroundAgentTaskManager::new_with_auth_mode( Arc::clone(self), chatgpt_base_url, - SessionSource::Cli, auth_mode, ) .authorization_header_value_or_bearer(auth) diff --git a/codex-rs/login/src/auth/mod.rs b/codex-rs/login/src/auth/mod.rs index 62ab467d08..b927f9a775 100644 --- a/codex-rs/login/src/auth/mod.rs +++ b/codex-rs/login/src/auth/mod.rs @@ -1,4 +1,3 @@ -mod agent_assertion; pub mod default_client; pub mod error; mod storage; @@ -8,7 +7,6 @@ mod external_bearer; mod manager; mod revoke; -pub use agent_assertion::AgentTaskAuthorizationTarget; pub use error::RefreshTokenFailedError; pub use error::RefreshTokenFailedReason; pub use manager::*; diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index af0e13b9b7..b7d167f593 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -21,7 +21,6 @@ pub use server::run_login_server; pub use agent_identity::BackgroundAgentTaskAuthMode; pub use agent_identity::cached_background_agent_task_authorization_header_value; pub use auth::AgentIdentityAuthRecord; -pub use auth::AgentTaskAuthorizationTarget; pub use auth::AuthConfig; pub use auth::AuthDotJson; pub use auth::AuthManager; diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 475d2941c3..f4dfb00dfa 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2769,26 +2769,6 @@ impl fmt::Display for SubAgentSource { } } -/// Persisted agent-task details that let a resumed thread keep using the same backend task. -/// -/// `agent_runtime_id` is validation metadata for the globally registered agent identity, not a -/// separate session-scoped identity. Resume only restores this task after confirming that runtime -/// id still matches the globally registered identity; otherwise the cached task is discarded and a -/// fresh task can be registered. -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema, TS)] -pub struct SessionAgentTask { - pub agent_runtime_id: String, - pub task_id: String, - pub registered_at: String, -} - -/// Session-scoped state updates that can be appended after the canonical SessionMeta line. -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema, TS, Default)] -pub struct SessionStateUpdate { - #[serde(default)] - pub agent_task: Option, -} - /// SessionMeta contains session-level data that doesn't correspond to a specific turn. /// /// NOTE: There used to be an `instructions` field here, which stored user_instructions, but we @@ -2858,7 +2838,6 @@ pub struct SessionMetaLine { #[serde(tag = "type", content = "payload", rename_all = "snake_case")] pub enum RolloutItem { SessionMeta(SessionMetaLine), - SessionState(SessionStateUpdate), ResponseItem(ResponseItem), Compacted(CompactedItem), TurnContext(TurnContextItem), diff --git a/codex-rs/rollout/src/list.rs b/codex-rs/rollout/src/list.rs index 0df1b9c9f5..bf09d15e86 100644 --- a/codex-rs/rollout/src/list.rs +++ b/codex-rs/rollout/src/list.rs @@ -1082,9 +1082,6 @@ async fn read_head_summary(path: &Path, head_limit: usize) -> io::Result { // Not included in `head`; skip. } - RolloutItem::SessionState(_) => { - // Not included in `head`; skip. - } RolloutItem::EventMsg(ev) => { if let EventMsg::UserMessage(user) = ev { summary.saw_user_event = true; @@ -1137,7 +1134,6 @@ pub async fn read_head_for_summary(path: &Path) -> io::Result {} } diff --git a/codex-rs/rollout/src/metadata.rs b/codex-rs/rollout/src/metadata.rs index 65e347a323..58d55a887d 100644 --- a/codex-rs/rollout/src/metadata.rs +++ b/codex-rs/rollout/src/metadata.rs @@ -70,8 +70,7 @@ pub fn builder_from_items( ) -> Option { if let Some(session_meta) = items.iter().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => Some(meta_line), - RolloutItem::SessionState(_) - | RolloutItem::ResponseItem(_) + RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::EventMsg(_) => None, @@ -125,8 +124,7 @@ pub async fn extract_metadata_from_rollout( metadata, memory_mode: items.iter().rev().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => meta_line.meta.memory_mode.clone(), - RolloutItem::SessionState(_) - | RolloutItem::ResponseItem(_) + RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::EventMsg(_) => None, diff --git a/codex-rs/rollout/src/policy.rs b/codex-rs/rollout/src/policy.rs index 86fc950a24..df41d2c766 100644 --- a/codex-rs/rollout/src/policy.rs +++ b/codex-rs/rollout/src/policy.rs @@ -16,10 +16,9 @@ pub fn is_persisted_response_item(item: &RolloutItem, mode: EventPersistenceMode RolloutItem::ResponseItem(item) => should_persist_response_item(item), RolloutItem::EventMsg(ev) => should_persist_event_msg(ev, mode), // Persist Codex executive markers so we can analyze flows (e.g., compaction, API turns). - RolloutItem::Compacted(_) - | RolloutItem::TurnContext(_) - | RolloutItem::SessionMeta(_) - | RolloutItem::SessionState(_) => true, + RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::SessionMeta(_) => { + true + } } } diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index f6e15a973e..871d53141b 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -707,9 +707,6 @@ impl RolloutRecorder { RolloutItem::Compacted(item) => { items.push(RolloutItem::Compacted(item)); } - RolloutItem::SessionState(update) => { - items.push(RolloutItem::SessionState(update)); - } RolloutItem::TurnContext(item) => { items.push(RolloutItem::TurnContext(item)); } @@ -1566,7 +1563,6 @@ async fn resume_candidate_matches_cwd( && let Some(latest_turn_context_cwd) = items.iter().rev().find_map(|item| match item { RolloutItem::TurnContext(turn_context) => Some(turn_context.cwd.as_path()), RolloutItem::SessionMeta(_) - | RolloutItem::SessionState(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::EventMsg(_) => None, diff --git a/codex-rs/state/src/extract.rs b/codex-rs/state/src/extract.rs index 64d7f5bbb7..6d36f211a4 100644 --- a/codex-rs/state/src/extract.rs +++ b/codex-rs/state/src/extract.rs @@ -19,7 +19,6 @@ pub fn apply_rollout_item( ) { match item { RolloutItem::SessionMeta(meta_line) => apply_session_meta_from_item(metadata, meta_line), - RolloutItem::SessionState(_) => {} RolloutItem::TurnContext(turn_ctx) => apply_turn_context(metadata, turn_ctx), RolloutItem::EventMsg(event) => apply_event_msg(metadata, event), RolloutItem::ResponseItem(item) => apply_response_item(metadata, item), @@ -37,10 +36,9 @@ pub fn rollout_item_affects_thread_metadata(item: &RolloutItem) -> bool { RolloutItem::EventMsg( EventMsg::TokenCount(_) | EventMsg::UserMessage(_) | EventMsg::ThreadNameUpdated(_), ) => true, - RolloutItem::SessionState(_) - | RolloutItem::EventMsg(_) - | RolloutItem::ResponseItem(_) - | RolloutItem::Compacted(_) => false, + RolloutItem::EventMsg(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) => { + false + } } } diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 5592ca54df..eb7494a037 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -954,8 +954,7 @@ SELECT pub(super) fn extract_dynamic_tools(items: &[RolloutItem]) -> Option>> { items.iter().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => Some(meta_line.meta.dynamic_tools.clone()), - RolloutItem::SessionState(_) - | RolloutItem::ResponseItem(_) + RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::EventMsg(_) => None, @@ -965,8 +964,7 @@ pub(super) fn extract_dynamic_tools(items: &[RolloutItem]) -> Option Option { items.iter().rev().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => meta_line.meta.memory_mode.clone(), - RolloutItem::SessionState(_) - | RolloutItem::ResponseItem(_) + RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::EventMsg(_) => None,