diff --git a/codex-rs/codex-api/src/api_bridge.rs b/codex-rs/codex-api/src/api_bridge.rs index 0ad2b13979..740f7cb783 100644 --- a/codex-rs/codex-api/src/api_bridge.rs +++ b/codex-rs/codex-api/src/api_bridge.rs @@ -178,13 +178,33 @@ struct UsageErrorBody { pub struct CoreAuthProvider { pub token: Option, pub account_id: Option, + authorization_header_override: Option, } impl CoreAuthProvider { + pub fn from_bearer_token(token: Option, account_id: Option) -> Self { + Self { + token, + account_id, + authorization_header_override: None, + } + } + + pub fn from_authorization_header_value( + authorization_header_value: Option, + account_id: Option, + ) -> Self { + Self { + token: None, + account_id, + authorization_header_override: authorization_header_value, + } + } + pub fn auth_header_attached(&self) -> bool { - self.token + self.authorization_header_value() .as_ref() - .is_some_and(|token| http::HeaderValue::from_str(&format!("Bearer {token}")).is_ok()) + .is_some_and(|value| http::HeaderValue::from_str(value).is_ok()) } pub fn auth_header_name(&self) -> Option<&'static str> { @@ -195,8 +215,20 @@ impl CoreAuthProvider { Self { token: token.map(str::to_string), account_id: account_id.map(str::to_string), + authorization_header_override: None, } } + + #[cfg(test)] + pub fn for_test_authorization_header( + authorization_header_value: Option<&str>, + account_id: Option<&str>, + ) -> Self { + Self::from_authorization_header_value( + authorization_header_value.map(str::to_string), + account_id.map(str::to_string), + ) + } } impl ApiAuthProvider for CoreAuthProvider { @@ -204,6 +236,12 @@ impl ApiAuthProvider for CoreAuthProvider { self.token.clone() } + fn authorization_header_value(&self) -> Option { + self.authorization_header_override + .clone() + .or_else(|| self.bearer_token().map(|token| format!("Bearer {token}"))) + } + fn account_id(&self) -> Option { self.account_id.clone() } diff --git a/codex-rs/codex-api/src/api_bridge_tests.rs b/codex-rs/codex-api/src/api_bridge_tests.rs index 71d3889915..c5274946a3 100644 --- a/codex-rs/codex-api/src/api_bridge_tests.rs +++ b/codex-rs/codex-api/src/api_bridge_tests.rs @@ -133,11 +133,22 @@ fn map_api_error_extracts_identity_auth_details_from_headers() { #[test] fn core_auth_provider_reports_when_auth_header_will_attach() { - let auth = CoreAuthProvider { - token: Some("access-token".to_string()), - account_id: None, - }; + let auth = CoreAuthProvider::from_bearer_token(Some("access-token".to_string()), None); assert!(auth.auth_header_attached()); assert_eq!(auth.auth_header_name(), Some("authorization")); } + +#[test] +fn core_auth_provider_supports_non_bearer_authorization_headers() { + let auth = + CoreAuthProvider::for_test_authorization_header(Some("AgentAssertion opaque-token"), None); + + assert!(auth.auth_header_attached()); + assert_eq!(auth.auth_header_name(), Some("authorization")); + assert_eq!(auth.bearer_token(), None); + assert_eq!( + auth.authorization_header_value(), + Some("AgentAssertion opaque-token".to_string()) + ); +} diff --git a/codex-rs/codex-api/src/auth.rs b/codex-rs/codex-api/src/auth.rs index f649062db1..4f27264a97 100644 --- a/codex-rs/codex-api/src/auth.rs +++ b/codex-rs/codex-api/src/auth.rs @@ -9,14 +9,17 @@ use http::HeaderValue; /// reach this interface. pub trait AuthProvider: Send + Sync { fn bearer_token(&self) -> Option; + fn authorization_header_value(&self) -> Option { + self.bearer_token().map(|token| format!("Bearer {token}")) + } fn account_id(&self) -> Option { None } } pub(crate) fn add_auth_headers_to_header_map(auth: &A, headers: &mut HeaderMap) { - if let Some(token) = auth.bearer_token() - && let Ok(header) = HeaderValue::from_str(&format!("Bearer {token}")) + if let Some(authorization) = auth.authorization_header_value() + && let Ok(header) = HeaderValue::from_str(&authorization) { let _ = headers.insert(http::header::AUTHORIZATION, header); } diff --git a/codex-rs/core/src/agent_identity.rs b/codex-rs/core/src/agent_identity.rs index 71c759e59e..515bbc7fef 100644 --- a/codex-rs/core/src/agent_identity.rs +++ b/codex-rs/core/src/agent_identity.rs @@ -27,12 +27,15 @@ use tracing::debug; use tracing::info; use tracing::warn; +use crate::config::Config; + +mod assertion; mod task_registration; +#[cfg(test)] +pub(crate) use assertion::AgentAssertionEnvelope; pub(crate) use task_registration::RegisteredAgentTask; -use crate::config::Config; - const AGENT_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15); const AGENT_IDENTITY_BISCUIT_TIMEOUT: Duration = Duration::from_secs(15); @@ -334,7 +337,7 @@ impl AgentIdentityManager { } #[cfg(test)] - fn new_for_tests( + pub(crate) fn new_for_tests( auth_manager: Arc, feature_enabled: bool, chatgpt_base_url: String, @@ -348,6 +351,30 @@ impl AgentIdentityManager { ensure_lock: Arc::new(Mutex::new(())), } } + + #[cfg(test)] + pub(crate) async fn seed_generated_identity_for_tests( + &self, + agent_runtime_id: &str, + ) -> Result { + let (auth, binding) = self + .current_auth_binding() + .await + .context("test agent identity requires ChatGPT auth")?; + let key_material = generate_agent_key_material()?; + let stored_identity = StoredAgentIdentity { + binding_id: binding.binding_id.clone(), + chatgpt_account_id: binding.chatgpt_account_id.clone(), + chatgpt_user_id: binding.chatgpt_user_id.clone(), + agent_runtime_id: agent_runtime_id.to_string(), + private_key_pkcs8_base64: key_material.private_key_pkcs8_base64, + public_key_ssh: key_material.public_key_ssh, + registered_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + abom: self.abom.clone(), + }; + self.store_identity(&auth, &stored_identity)?; + Ok(stored_identity) + } } impl StoredAgentIdentity { @@ -577,7 +604,7 @@ mod tests { .and(path("/v1/agent/register")) .and(header("x-openai-authorization", "human-biscuit")) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "agent_runtime_id": "agent_123", + "agent_runtime_id": "agent-123", }))) .expect(1) .mount(&server) @@ -603,7 +630,7 @@ mod tests { .unwrap() .expect("identity should be reused"); - assert_eq!(first.agent_runtime_id, "agent_123"); + assert_eq!(first.agent_runtime_id, "agent-123"); assert_eq!(first, second); assert_eq!(first.abom.agent_harness_id, "codex-cli"); assert_eq!(first.chatgpt_account_id, "account-123"); @@ -619,7 +646,7 @@ mod tests { .and(path("/v1/agent/register")) .and(header("x-openai-authorization", "human-biscuit")) .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "agent_runtime_id": "agent_456", + "agent_runtime_id": "agent-456", }))) .expect(1) .mount(&server) @@ -650,11 +677,11 @@ mod tests { .unwrap() .expect("identity should be registered"); - assert_eq!(stored.agent_runtime_id, "agent_456"); + assert_eq!(stored.agent_runtime_id, "agent-456"); let persisted = auth .get_agent_identity(&binding.chatgpt_account_id) .expect("stored identity"); - assert_eq!(persisted.agent_runtime_id, "agent_456"); + assert_eq!(persisted.agent_runtime_id, "agent-456"); } #[tokio::test] diff --git a/codex-rs/core/src/agent_identity/assertion.rs b/codex-rs/core/src/agent_identity/assertion.rs new file mode 100644 index 0000000000..5943ae5346 --- /dev/null +++ b/codex-rs/core/src/agent_identity/assertion.rs @@ -0,0 +1,182 @@ +use std::collections::BTreeMap; + +use anyhow::Context; +use anyhow::Result; +use base64::Engine as _; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use ed25519_dalek::Signer as _; +use serde::Deserialize; +use serde::Serialize; +use tracing::debug; + +use super::*; + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +pub(crate) struct AgentAssertionEnvelope { + pub(crate) agent_runtime_id: String, + pub(crate) task_id: String, + pub(crate) timestamp: String, + pub(crate) signature: String, +} + +impl AgentIdentityManager { + pub(crate) async fn authorization_header_for_task( + &self, + agent_task: &RegisteredAgentTask, + ) -> Result> { + if !self.feature_enabled { + return Ok(None); + } + + let Some(stored_identity) = self.ensure_registered_identity().await? else { + return Ok(None); + }; + anyhow::ensure!( + stored_identity.agent_runtime_id == agent_task.agent_runtime_id, + "agent task runtime {} does not match stored agent identity {}", + agent_task.agent_runtime_id, + stored_identity.agent_runtime_id + ); + + let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); + let envelope = AgentAssertionEnvelope { + agent_runtime_id: agent_task.agent_runtime_id.clone(), + task_id: agent_task.task_id.clone(), + timestamp: timestamp.clone(), + signature: sign_agent_assertion_payload(&stored_identity, agent_task, ×tamp)?, + }; + let serialized_assertion = serialize_agent_assertion(&envelope)?; + debug!( + agent_runtime_id = %envelope.agent_runtime_id, + task_id = %envelope.task_id, + "attaching agent assertion authorization to downstream request" + ); + Ok(Some(format!("AgentAssertion {serialized_assertion}"))) + } +} + +fn sign_agent_assertion_payload( + stored_identity: &StoredAgentIdentity, + agent_task: &RegisteredAgentTask, + timestamp: &str, +) -> Result { + let signing_key = stored_identity.signing_key()?; + let payload = format!( + "{}:{}:{timestamp}", + agent_task.agent_runtime_id, agent_task.task_id + ); + Ok(BASE64_STANDARD.encode(signing_key.sign(payload.as_bytes()).to_bytes())) +} + +fn serialize_agent_assertion(envelope: &AgentAssertionEnvelope) -> Result { + let payload = serde_json::to_vec(&BTreeMap::from([ + ("agent_runtime_id", envelope.agent_runtime_id.as_str()), + ("signature", envelope.signature.as_str()), + ("task_id", envelope.task_id.as_str()), + ("timestamp", envelope.timestamp.as_str()), + ])) + .context("failed to serialize agent assertion envelope")?; + Ok(URL_SAFE_NO_PAD.encode(payload)) +} + +#[cfg(test)] +mod tests { + use base64::engine::general_purpose::URL_SAFE_NO_PAD; + use ed25519_dalek::Signature; + use ed25519_dalek::Verifier as _; + use pretty_assertions::assert_eq; + + use super::*; + + #[tokio::test] + async fn authorization_header_for_task_skips_when_feature_is_disabled() { + let auth_manager = + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ false, + "https://chatgpt.com/backend-api/".to_string(), + SessionSource::Cli, + ); + let agent_task = RegisteredAgentTask { + binding_id: "chatgpt-account-account_id".to_string(), + chatgpt_account_id: "account_id".to_string(), + chatgpt_user_id: None, + agent_runtime_id: "agent-123".to_string(), + task_id: "task-123".to_string(), + registered_at: "2026-03-23T12:00:00Z".to_string(), + }; + + assert_eq!( + manager + .authorization_header_for_task(&agent_task) + .await + .unwrap(), + None + ); + } + + #[tokio::test] + async fn authorization_header_for_task_serializes_signed_agent_assertion() { + let auth_manager = + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + "https://chatgpt.com/backend-api/".to_string(), + SessionSource::Cli, + ); + let stored_identity = manager + .seed_generated_identity_for_tests("agent-123") + .await + .expect("seed test identity"); + let agent_task = RegisteredAgentTask { + binding_id: "chatgpt-account-account_id".to_string(), + chatgpt_account_id: "account_id".to_string(), + chatgpt_user_id: None, + agent_runtime_id: "agent-123".to_string(), + task_id: "task-123".to_string(), + registered_at: "2026-03-23T12:00:00Z".to_string(), + }; + + let header = manager + .authorization_header_for_task(&agent_task) + .await + .expect("build agent assertion") + .expect("header should exist"); + let token = header + .strip_prefix("AgentAssertion ") + .expect("agent assertion scheme"); + let payload = URL_SAFE_NO_PAD + .decode(token) + .expect("valid base64url payload"); + let envelope: AgentAssertionEnvelope = + serde_json::from_slice(&payload).expect("valid assertion envelope"); + + assert_eq!( + envelope, + AgentAssertionEnvelope { + agent_runtime_id: "agent-123".to_string(), + task_id: "task-123".to_string(), + timestamp: envelope.timestamp.clone(), + signature: envelope.signature.clone(), + } + ); + let signature_bytes = BASE64_STANDARD + .decode(&envelope.signature) + .expect("valid base64 signature"); + let signature = Signature::from_slice(&signature_bytes).expect("valid signature bytes"); + let signing_key = stored_identity.signing_key().expect("signing key"); + signing_key + .verifying_key() + .verify( + format!( + "{}:{}:{}", + envelope.agent_runtime_id, envelope.task_id, envelope.timestamp + ) + .as_bytes(), + &signature, + ) + .expect("signature should verify"); + } +} diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 16f743943a..0eb6c03138 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -31,6 +31,8 @@ use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; +use crate::agent_identity::AgentIdentityManager; +use crate::agent_identity::RegisteredAgentTask; use codex_api::ApiError; use codex_api::CompactClient as ApiCompactClient; use codex_api::CompactionInput as ApiCompactionInput; @@ -92,6 +94,7 @@ use tokio::sync::oneshot; use tokio::sync::oneshot::error::TryRecvError; use tokio_tungstenite::tungstenite::Error; use tokio_tungstenite::tungstenite::Message; +use tracing::debug; use tracing::instrument; use tracing::trace; use tracing::warn; @@ -144,6 +147,7 @@ pub(crate) const WEBSOCKET_CONNECT_TIMEOUT: Duration = #[derive(Debug)] struct ModelClientState { auth_manager: Option>, + agent_identity_manager: Option>, conversation_id: ThreadId, window_generation: AtomicU64, installation_id: String, @@ -211,6 +215,8 @@ pub struct ModelClient { pub struct ModelClientSession { client: ModelClient, websocket_session: WebsocketSession, + agent_task: Option, + cache_websocket_session_on_drop: bool, /// Turn state for sticky routing. /// /// This is an `OnceLock` that stores the turn state value received from the server @@ -306,6 +312,33 @@ impl ModelClient { enable_request_compression: bool, include_timing_metrics: bool, beta_features_header: Option, + ) -> Self { + Self::new_with_agent_identity_manager( + auth_manager, + /*agent_identity_manager*/ None, + conversation_id, + installation_id, + provider, + session_source, + model_verbosity, + enable_request_compression, + include_timing_metrics, + beta_features_header, + ) + } + + #[allow(clippy::too_many_arguments)] + pub(crate) fn new_with_agent_identity_manager( + auth_manager: Option>, + agent_identity_manager: Option>, + conversation_id: ThreadId, + installation_id: String, + provider: ModelProviderInfo, + session_source: SessionSource, + model_verbosity: Option, + enable_request_compression: bool, + include_timing_metrics: bool, + beta_features_header: Option, ) -> Self { let auth_manager = auth_manager_for_provider(auth_manager, &provider); let codex_api_key_env_enabled = auth_manager @@ -315,6 +348,7 @@ impl ModelClient { Self { state: Arc::new(ModelClientState { auth_manager, + agent_identity_manager, conversation_id, window_generation: AtomicU64::new(0), installation_id, @@ -336,9 +370,25 @@ impl ModelClient { /// This constructor does not perform network I/O itself; the session opens a websocket lazily /// when the first stream request is issued. pub fn new_session(&self) -> ModelClientSession { + self.new_session_with_agent_task(None) + } + + pub(crate) fn new_session_with_agent_task( + &self, + agent_task: Option, + ) -> ModelClientSession { + let cache_websocket_session_on_drop = agent_task.is_none(); + let websocket_session = if agent_task.is_some() { + drop(self.take_cached_websocket_session()); + WebsocketSession::default() + } else { + self.take_cached_websocket_session() + }; ModelClientSession { client: self.clone(), - websocket_session: self.take_cached_websocket_session(), + websocket_session, + agent_task, + cache_websocket_session_on_drop, turn_state: Arc::new(OnceLock::new()), } } @@ -421,7 +471,7 @@ impl ModelClient { if prompt.input.is_empty() { return Ok(Vec::new()); } - let client_setup = self.current_client_setup().await?; + let client_setup = self.current_client_setup(None).await?; let transport = ReqwestTransport::new(build_reqwest_client()); let request_telemetry = Self::build_request_telemetry( session_telemetry, @@ -485,7 +535,7 @@ impl ModelClient { ) -> Result { // Create the media call over HTTP first, then retain matching auth so realtime can attach // the server-side control WebSocket to the call id from that HTTP response. - let client_setup = self.current_client_setup().await?; + let client_setup = self.current_client_setup(/*agent_task*/ None).await?; let mut sideband_headers = extra_headers.clone(); sideband_headers.extend(sideband_websocket_auth_headers(&client_setup.api_auth)); let transport = ReqwestTransport::new(build_reqwest_client()); @@ -518,7 +568,7 @@ impl ModelClient { return Ok(Vec::new()); } - let client_setup = self.current_client_setup().await?; + let client_setup = self.current_client_setup(None).await?; let transport = ReqwestTransport::new(build_reqwest_client()); let request_telemetry = Self::build_request_telemetry( session_telemetry, @@ -659,7 +709,10 @@ impl ModelClient { /// /// This centralizes setup used by both prewarm and normal request paths so they stay in /// lockstep when auth/provider resolution changes. - async fn current_client_setup(&self) -> Result { + async fn current_client_setup( + &self, + agent_task: Option<&RegisteredAgentTask>, + ) -> Result { let auth = match self.state.auth_manager.as_ref() { Some(manager) => manager.auth().await, None => None, @@ -668,7 +721,33 @@ impl ModelClient { .state .provider .to_api_provider(auth.as_ref().map(CodexAuth::auth_mode))?; - let api_auth = auth_provider_from_auth(auth.clone(), &self.state.provider)?; + let api_auth = match (agent_task, self.state.agent_identity_manager.as_ref()) { + (Some(agent_task), Some(agent_identity_manager)) => { + if let Some(authorization_header_value) = agent_identity_manager + .authorization_header_for_task(agent_task) + .await + .map_err(|err| { + CodexErr::Stream( + format!("failed to build agent assertion authorization: {err}"), + None, + ) + })? + { + debug!( + agent_runtime_id = %agent_task.agent_runtime_id, + task_id = %agent_task.task_id, + "using agent assertion authorization for downstream request" + ); + CoreAuthProvider::from_authorization_header_value( + Some(authorization_header_value), + None, + ) + } else { + auth_provider_from_auth(auth.clone(), &self.state.provider)? + } + } + _ => auth_provider_from_auth(auth.clone(), &self.state.provider)?, + }; Ok(CurrentClientSetup { auth, api_provider, @@ -802,12 +881,18 @@ impl ModelClient { impl Drop for ModelClientSession { fn drop(&mut self) { let websocket_session = std::mem::take(&mut self.websocket_session); - self.client - .store_cached_websocket_session(websocket_session); + if self.cache_websocket_session_on_drop { + self.client + .store_cached_websocket_session(websocket_session); + } } } impl ModelClientSession { + pub(crate) fn disable_cached_websocket_session_on_drop(&mut self) { + self.cache_websocket_session_on_drop = false; + } + pub(crate) fn reset_websocket_session(&mut self) { self.websocket_session.connection = None; self.websocket_session.last_request = None; @@ -1009,11 +1094,15 @@ impl ModelClientSession { return Ok(()); } - let client_setup = self.client.current_client_setup().await.map_err(|err| { - ApiError::Stream(format!( - "failed to build websocket prewarm client setup: {err}" - )) - })?; + let client_setup = self + .client + .current_client_setup(self.agent_task.as_ref()) + .await + .map_err(|err| { + ApiError::Stream(format!( + "failed to build websocket prewarm client setup: {err}" + )) + })?; let auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), &client_setup.api_auth, @@ -1167,7 +1256,10 @@ impl ModelClientSession { .map(AuthManager::unauthorized_recovery); let mut pending_retry = PendingUnauthorizedRetry::default(); loop { - let client_setup = self.client.current_client_setup().await?; + let client_setup = self + .client + .current_client_setup(self.agent_task.as_ref()) + .await?; let transport = ReqwestTransport::new(build_reqwest_client()); let request_auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), @@ -1256,7 +1348,10 @@ impl ModelClientSession { .map(AuthManager::unauthorized_recovery); let mut pending_retry = PendingUnauthorizedRetry::default(); loop { - let client_setup = self.client.current_client_setup().await?; + let client_setup = self + .client + .current_client_setup(self.agent_task.as_ref()) + .await?; let request_auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), &client_setup.api_auth, diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 2fd5f04f95..9bf0aaa94d 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use super::AuthRequestTelemetryContext; use super::ModelClient; use super::PendingUnauthorizedRetry; @@ -7,17 +9,39 @@ use super::X_CODEX_PARENT_THREAD_ID_HEADER; use super::X_CODEX_TURN_METADATA_HEADER; use super::X_CODEX_WINDOW_ID_HEADER; use super::X_OPENAI_SUBAGENT_HEADER; +use crate::Prompt; +use crate::ResponseEvent; +use crate::agent_identity::AgentAssertionEnvelope; +use crate::agent_identity::AgentIdentityManager; +use crate::agent_identity::RegisteredAgentTask; +use crate::agent_identity::StoredAgentIdentity; +use base64::Engine as _; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; use codex_api::CoreAuthProvider; use codex_app_server_protocol::AuthMode; +use codex_keyring_store::tests::MockKeyringStore; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::WireApi; use codex_model_provider_info::create_oss_provider_with_base_url; use codex_otel::SessionTelemetry; use codex_protocol::ThreadId; +use codex_protocol::config_types::ReasoningSummary; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ModelInfo; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; +use codex_secrets::SecretsBackendKind; +use codex_secrets::SecretsManager; +use core_test_support::responses; +use ed25519_dalek::Signature; +use ed25519_dalek::Verifier as _; +use futures::StreamExt; use pretty_assertions::assert_eq; use serde_json::json; +use tempfile::TempDir; fn test_model_client(session_source: SessionSource) -> ModelClient { let provider = create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses); @@ -79,6 +103,118 @@ fn test_session_telemetry() -> SessionTelemetry { ) } +fn test_prompt(text: &str) -> Prompt { + Prompt { + input: vec![ResponseItem::Message { + id: None, + role: "user".into(), + content: vec![ContentItem::InputText { + text: text.to_string(), + }], + end_turn: None, + phase: None, + }], + ..Prompt::default() + } +} + +async fn drain_stream_to_completion(stream: &mut crate::ResponseStream) -> anyhow::Result<()> { + while let Some(event) = stream.next().await { + if matches!(event?, ResponseEvent::Completed { .. }) { + break; + } + } + Ok(()) +} + +async fn model_client_with_agent_task( + provider: ModelProviderInfo, +) -> ( + TempDir, + ModelClient, + RegisteredAgentTask, + StoredAgentIdentity, +) { + let codex_home = tempfile::tempdir().expect("tempdir"); + let keyring_store = Arc::new(MockKeyringStore::default()); + let secrets_manager = SecretsManager::new_with_keyring_store( + codex_home.path().to_path_buf(), + SecretsBackendKind::Local, + keyring_store, + ); + let auth_manager = + AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); + let agent_identity_manager = Arc::new(AgentIdentityManager::new_for_tests( + Arc::clone(&auth_manager), + /*feature_enabled*/ true, + "https://chatgpt.com/backend-api/".to_string(), + SessionSource::Cli, + secrets_manager, + )); + let stored_identity = agent_identity_manager + .seed_generated_identity_for_tests("agent-123") + .await + .expect("seed test identity"); + let agent_task = RegisteredAgentTask { + agent_runtime_id: stored_identity.agent_runtime_id.clone(), + task_id: "task-123".to_string(), + registered_at: "2026-03-23T12:00:00Z".to_string(), + }; + let client = ModelClient::new_with_agent_identity_manager( + Some(auth_manager), + Some(agent_identity_manager), + ThreadId::new(), + /*installation_id*/ "11111111-1111-4111-8111-111111111111".to_string(), + provider, + SessionSource::Cli, + None, + false, + false, + None, + ); + (codex_home, client, agent_task, stored_identity) +} + +fn assert_agent_assertion_header( + authorization_header: &str, + stored_identity: &StoredAgentIdentity, + expected_agent_runtime_id: &str, + expected_task_id: &str, +) { + let token = authorization_header + .strip_prefix("AgentAssertion ") + .expect("agent assertion authorization scheme"); + let envelope: AgentAssertionEnvelope = serde_json::from_slice( + &URL_SAFE_NO_PAD + .decode(token) + .expect("base64url-encoded agent assertion"), + ) + .expect("valid agent assertion envelope"); + + assert_eq!(envelope.agent_runtime_id, expected_agent_runtime_id); + assert_eq!(envelope.task_id, expected_task_id); + + let signature = Signature::from_slice( + &base64::engine::general_purpose::STANDARD + .decode(&envelope.signature) + .expect("base64 signature"), + ) + .expect("signature bytes"); + stored_identity + .signing_key() + .expect("signing key") + .verifying_key() + .verify( + format!( + "{}:{}:{}", + envelope.agent_runtime_id, envelope.task_id, envelope.timestamp + ) + .as_bytes(), + &signature, + ) + .expect("signature should verify"); +} + #[test] fn build_subagent_headers_sets_other_subagent_label() { let client = test_model_client(SessionSource::SubAgent(SubAgentSource::Other( @@ -169,3 +305,130 @@ fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() { assert_eq!(auth_context.recovery_mode, Some("managed")); assert_eq!(auth_context.recovery_phase, Some("refresh_token")); } + +#[tokio::test] +async fn responses_http_uses_agent_assertion_when_agent_task_is_present() { + core_test_support::skip_if_no_network!(); + + let server = responses::start_mock_server().await; + let request_recorder = responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_completed("resp-1"), + ]), + ) + .await; + let provider = + create_oss_provider_with_base_url(&format!("{}/v1", server.uri()), WireApi::Responses); + let (_codex_home, client, agent_task, stored_identity) = + model_client_with_agent_task(provider).await; + let model_info = test_model_info(); + let session_telemetry = test_session_telemetry(); + let mut client_session = client.new_session_with_agent_task(Some(agent_task.clone())); + + let mut stream = client_session + .stream( + &test_prompt("hello"), + &model_info, + &session_telemetry, + None, + ReasoningSummary::Auto, + None, + None, + ) + .await + .expect("stream request should succeed"); + drain_stream_to_completion(&mut stream) + .await + .expect("stream should complete"); + + let request = request_recorder.single_request(); + let authorization = request + .header("authorization") + .expect("authorization header should be present"); + assert_agent_assertion_header( + &authorization, + &stored_identity, + &agent_task.agent_runtime_id, + &agent_task.task_id, + ); + assert_eq!(request.header("chatgpt-account-id"), None); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn websocket_agent_task_bypasses_cached_bearer_prewarm() { + core_test_support::skip_if_no_network!(); + + let server = responses::start_websocket_server(vec![ + vec![vec![ + responses::ev_response_created("resp-prewarm"), + responses::ev_completed("resp-prewarm"), + ]], + vec![vec![ + responses::ev_response_created("resp-1"), + responses::ev_completed("resp-1"), + ]], + ]) + .await; + let mut provider = + create_oss_provider_with_base_url(&format!("{}/v1", server.uri()), WireApi::Responses); + provider.supports_websockets = true; + provider.websocket_connect_timeout_ms = Some(5_000); + let (_codex_home, client, agent_task, stored_identity) = + model_client_with_agent_task(provider).await; + let model_info = test_model_info(); + let session_telemetry = test_session_telemetry(); + let prompt = test_prompt("hello"); + + let mut prewarm_session = client.new_session(); + prewarm_session + .prewarm_websocket( + &prompt, + &model_info, + &session_telemetry, + None, + ReasoningSummary::Auto, + None, + None, + ) + .await + .expect("bearer prewarm should succeed"); + drop(prewarm_session); + + let mut agent_task_session = client.new_session_with_agent_task(Some(agent_task.clone())); + let mut stream = agent_task_session + .stream( + &prompt, + &model_info, + &session_telemetry, + None, + ReasoningSummary::Auto, + None, + None, + ) + .await + .expect("agent task stream should succeed"); + drain_stream_to_completion(&mut stream) + .await + .expect("agent task websocket stream should complete"); + + let handshakes = server.handshakes(); + assert_eq!(handshakes.len(), 2); + assert_eq!( + handshakes[0].header("authorization"), + Some("Bearer Access Token".to_string()) + ); + let agent_authorization = handshakes[1] + .header("authorization") + .expect("agent handshake should include authorization"); + assert_agent_assertion_header( + &agent_authorization, + &stored_identity, + &agent_task.agent_runtime_id, + &agent_task.task_id, + ); + assert_eq!(handshakes[1].header("chatgpt-account-id"), None); + + server.shutdown().await; +} diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index a2398849df..38be995a38 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -2185,6 +2185,11 @@ impl Session { config.analytics_enabled, ) }); + let agent_identity_manager = Arc::new(AgentIdentityManager::new( + config.as_ref(), + Arc::clone(&auth_manager), + session_configuration.session_source.clone(), + )); let services = SessionServices { // Initialize the MCP connection manager with an uninitialized // instance. It will be replaced with one created via @@ -2207,11 +2212,7 @@ impl Session { hooks, rollout: Mutex::new(rollout_recorder), user_shell: Arc::new(default_shell), - agent_identity_manager: Arc::new(AgentIdentityManager::new( - config.as_ref(), - Arc::clone(&auth_manager), - session_configuration.session_source.clone(), - )), + agent_identity_manager: Arc::clone(&agent_identity_manager), shell_snapshot_tx, show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, @@ -2228,8 +2229,9 @@ impl Session { network_proxy, network_approval: Arc::clone(&network_approval), state_db: state_db_ctx.clone(), - model_client: ModelClient::new( + model_client: ModelClient::new_with_agent_identity_manager( Some(Arc::clone(&auth_manager)), + Some(agent_identity_manager), conversation_id, installation_id, session_configuration.provider.clone(), @@ -6373,11 +6375,14 @@ pub(crate) async fn run_turn( })) .await; } - if let Err(error) = sess.ensure_agent_task_registered().await { - warn!(error = %error, "agent task registration failed"); - sess.fail_agent_identity_registration(error).await; - return None; - } + let agent_task = match sess.ensure_agent_task_registered().await { + Ok(agent_task) => agent_task, + Err(error) => { + warn!(error = %error, "agent task registration failed"); + sess.fail_agent_identity_registration(error).await; + return None; + } + }; if !skill_items.is_empty() { sess.record_conversation_items(&turn_context, &skill_items) @@ -6400,8 +6405,21 @@ pub(crate) async fn run_turn( // `ModelClientSession` is turn-scoped and caches WebSocket + sticky routing state, so we reuse // one instance across retries within this turn. - let mut client_session = - prewarmed_client_session.unwrap_or_else(|| sess.services.model_client.new_session()); + let mut prewarmed_client_session = prewarmed_client_session; + if agent_task.is_some() + && let Some(prewarmed_client_session) = prewarmed_client_session.as_mut() + { + prewarmed_client_session.disable_cached_websocket_session_on_drop(); + } + let mut client_session = if let Some(agent_task) = agent_task { + sess.services + .model_client + .new_session_with_agent_task(Some(agent_task)) + } else if let Some(prewarmed_client_session) = prewarmed_client_session.take() { + prewarmed_client_session + } else { + sess.services.model_client.new_session() + }; // Pending input is drained into history before building the next model request. // However, we defer that drain until after sampling in two cases: // 1. At the start of a turn, so the fresh user prompt in `input` gets sampled first. diff --git a/codex-rs/core/src/mcp_openai_file.rs b/codex-rs/core/src/mcp_openai_file.rs index 587dd4b770..8ee0700f2e 100644 --- a/codex-rs/core/src/mcp_openai_file.rs +++ b/codex-rs/core/src/mcp_openai_file.rs @@ -112,10 +112,8 @@ async fn build_uploaded_local_argument_value( let token_data = auth .get_token_data() .map_err(|error| format!("failed to read ChatGPT auth for file upload: {error}"))?; - let upload_auth = CoreAuthProvider { - token: Some(token_data.access_token), - account_id: token_data.account_id, - }; + let upload_auth = + CoreAuthProvider::from_bearer_token(Some(token_data.access_token), token_data.account_id); let uploaded = upload_local_file( turn_context.config.chatgpt_base_url.trim_end_matches('/'), &upload_auth, diff --git a/codex-rs/login/src/api_bridge.rs b/codex-rs/login/src/api_bridge.rs index d8b9dbb77c..49dd66eb05 100644 --- a/codex-rs/login/src/api_bridge.rs +++ b/codex-rs/login/src/api_bridge.rs @@ -8,29 +8,20 @@ pub fn auth_provider_from_auth( provider: &ModelProviderInfo, ) -> codex_protocol::error::Result { if let Some(api_key) = provider.api_key()? { - return Ok(CoreAuthProvider { - token: Some(api_key), - account_id: None, - }); + return Ok(CoreAuthProvider::from_bearer_token(Some(api_key), None)); } if let Some(token) = provider.experimental_bearer_token.clone() { - return Ok(CoreAuthProvider { - token: Some(token), - account_id: None, - }); + return Ok(CoreAuthProvider::from_bearer_token(Some(token), None)); } if let Some(auth) = auth { let token = auth.get_token()?; - Ok(CoreAuthProvider { - token: Some(token), - account_id: auth.get_account_id(), - }) + Ok(CoreAuthProvider::from_bearer_token( + Some(token), + auth.get_account_id(), + )) } else { - Ok(CoreAuthProvider { - token: None, - account_id: None, - }) + Ok(CoreAuthProvider::from_bearer_token(None, None)) } }