diff --git a/codex-rs/codex-api/src/auth.rs b/codex-rs/codex-api/src/auth.rs index 2d88dd6e2..b889c359e 100644 --- a/codex-rs/codex-api/src/auth.rs +++ b/codex-rs/codex-api/src/auth.rs @@ -67,6 +67,12 @@ pub type AuthProviderFuture<'a> = /// Shared auth handle passed through API clients. pub type SharedAuthProvider = Arc; +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct AgentIdentityTelemetry { + pub agent_id: String, + pub task_id: String, +} + #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub struct AuthHeaderTelemetry { pub attached: bool, diff --git a/codex-rs/codex-api/src/lib.rs b/codex-rs/codex-api/src/lib.rs index 1e7294be8..72de4d058 100644 --- a/codex-rs/codex-api/src/lib.rs +++ b/codex-rs/codex-api/src/lib.rs @@ -19,6 +19,7 @@ pub use codex_client::ReqwestTransport; pub use codex_client::TransportError; pub use crate::api_bridge::map_api_error; +pub use crate::auth::AgentIdentityTelemetry; pub use crate::auth::AuthError; pub use crate::auth::AuthHeaderTelemetry; pub use crate::auth::AuthProvider; diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index a4b2c491f..ffa8ae8e2 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -30,6 +30,7 @@ use std::sync::OnceLock; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; +use codex_api::AgentIdentityTelemetry; use codex_api::ApiError; use codex_api::AuthProvider; use codex_api::CompactClient as ApiCompactClient; @@ -115,8 +116,11 @@ use crate::responses_metadata::subagent_header_value; use crate::util::emit_feedback_auth_recovery_tags; use codex_feedback::FeedbackRequestTags; use codex_feedback::emit_feedback_request_tags_with_auth_env; +use codex_login::auth::AgentIdentityAuthPolicy; use codex_login::auth_env_telemetry::AuthEnvTelemetry; use codex_login::auth_env_telemetry::collect_auth_env_telemetry; +use codex_model_provider::AgentIdentitySessionFallback; +use codex_model_provider::ProviderAuthScope; use codex_model_provider::SharedModelProvider; use codex_model_provider::create_model_provider; #[cfg(test)] @@ -202,6 +206,7 @@ struct ModelClientState { include_attestation: bool, attestation_provider: Option>, disable_websockets: AtomicBool, + agent_identity_session_fallback: AgentIdentitySessionFallback, cached_websocket_session: StdMutex, } @@ -213,6 +218,7 @@ struct CurrentClientSetup { auth: Option, api_provider: ApiProvider, api_auth: SharedAuthProvider, + agent_identity_telemetry: Option, } #[derive(Clone, Copy)] @@ -240,6 +246,7 @@ impl RequestRouteTelemetry { #[derive(Debug, Clone)] pub struct ModelClient { state: Arc, + agent_identity_policy: AgentIdentityAuthPolicy, prompt_cache_key_override: Option, } @@ -392,6 +399,7 @@ impl ModelClient { /// are passed to [`ModelClientSession::stream`] (and other turn-scoped methods) explicitly. pub fn new( auth_manager: Option>, + agent_identity_policy: AgentIdentityAuthPolicy, thread_id: ThreadId, provider_info: ModelProviderInfo, session_source: SessionSource, @@ -426,8 +434,10 @@ impl ModelClient { include_attestation, attestation_provider, disable_websockets: AtomicBool::new(false), + agent_identity_session_fallback: AgentIdentitySessionFallback::default(), cached_websocket_session: StdMutex::new(WebsocketSession::default()), }), + agent_identity_policy, prompt_cache_key_override: None, } } @@ -528,6 +538,7 @@ impl ModelClient { AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), client_setup.api_auth.as_ref(), + client_setup.agent_identity_telemetry.clone(), PendingUnauthorizedRetry::default(), ), RequestRouteTelemetry::for_endpoint(RESPONSES_COMPACT_ENDPOINT), @@ -660,6 +671,7 @@ impl ModelClient { AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), client_setup.api_auth.as_ref(), + client_setup.agent_identity_telemetry.clone(), PendingUnauthorizedRetry::default(), ), RequestRouteTelemetry::for_endpoint(MEMORIES_SUMMARIZE_ENDPOINT), @@ -909,14 +921,27 @@ impl ModelClient { async fn current_client_setup(&self) -> Result { let auth = self.state.provider.auth().await; let api_provider = self.state.provider.api_provider().await?; - let api_auth = self.state.provider.api_auth().await?; + let resolved_auth = self + .state + .provider + .api_auth_for_scope(ProviderAuthScope { + agent_identity_policy: self.agent_identity_policy, + session_source: self.state.session_source.clone(), + agent_identity_session_fallback: self.state.agent_identity_session_fallback.clone(), + }) + .await?; Ok(CurrentClientSetup { auth, api_provider, - api_auth, + api_auth: resolved_auth.auth, + agent_identity_telemetry: resolved_auth.agent_identity_telemetry, }) } + pub(crate) async fn prewarm_auth(&self) -> Result<()> { + self.current_client_setup().await.map(|_| ()) + } + /// Opens a websocket connection using the same header and telemetry wiring as normal turns. /// /// Both startup prewarm and in-turn `needs_new` reconnects call this path so handshake @@ -934,7 +959,7 @@ impl ModelClient { let headers = self.build_websocket_headers(responses_metadata).await; let websocket_telemetry = ModelClientSession::build_websocket_telemetry( session_telemetry, - auth_context, + auth_context.clone(), request_route_telemetry, self.state.auth_env_telemetry.clone(), ); @@ -976,6 +1001,7 @@ impl ModelClient { response_debug.cf_ray.as_deref(), response_debug.auth_error.as_deref(), response_debug.auth_error_code.as_deref(), + auth_context.agent_identity_telemetry(), ); emit_feedback_request_tags_with_auth_env( &FeedbackRequestTags { @@ -1205,6 +1231,7 @@ impl ModelClientSession { let auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), client_setup.api_auth.as_ref(), + client_setup.agent_identity_telemetry.clone(), PendingUnauthorizedRetry::default(), ); let connection = self @@ -1343,6 +1370,7 @@ impl ModelClientSession { let request_auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), client_setup.api_auth.as_ref(), + client_setup.agent_identity_telemetry.clone(), pending_retry, ); let (request_telemetry, sse_telemetry) = Self::build_streaming_telemetry( @@ -1470,6 +1498,7 @@ impl ModelClientSession { let request_auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), client_setup.api_auth.as_ref(), + client_setup.agent_identity_telemetry.clone(), pending_retry, ); let request = self.client.build_responses_request( @@ -2035,11 +2064,12 @@ impl PendingUnauthorizedRetry { } } -#[derive(Clone, Copy, Debug, Default)] +#[derive(Clone, Debug, Default)] struct AuthRequestTelemetryContext { auth_mode: Option<&'static str>, auth_header_attached: bool, auth_header_name: Option<&'static str>, + agent_identity_telemetry: Option, retry_after_unauthorized: bool, recovery_mode: Option<&'static str>, recovery_phase: Option<&'static str>, @@ -2049,6 +2079,7 @@ impl AuthRequestTelemetryContext { fn new( auth_mode: Option, api_auth: &dyn AuthProvider, + agent_identity_telemetry: Option, retry: PendingUnauthorizedRetry, ) -> Self { let auth_telemetry = auth_header_telemetry(api_auth); @@ -2062,11 +2093,16 @@ impl AuthRequestTelemetryContext { }), auth_header_attached: auth_telemetry.attached, auth_header_name: auth_telemetry.name, + agent_identity_telemetry, retry_after_unauthorized: retry.retry_after_unauthorized, recovery_mode: retry.recovery_mode, recovery_phase: retry.recovery_phase, } } + + fn agent_identity_telemetry(&self) -> Option<&AgentIdentityTelemetry> { + self.agent_identity_telemetry.as_ref() + } } struct WebsocketConnectParams<'a> { @@ -2253,6 +2289,7 @@ impl RequestTelemetry for ApiTelemetry { debug.cf_ray.as_deref(), debug.auth_error.as_deref(), debug.auth_error_code.as_deref(), + self.auth_context.agent_identity_telemetry(), ); emit_feedback_request_tags_with_auth_env( &FeedbackRequestTags { @@ -2307,6 +2344,7 @@ impl WebsocketTelemetry for ApiTelemetry { duration, error_message.as_deref(), connection_reused, + self.auth_context.agent_identity_telemetry(), ); emit_feedback_request_tags_with_auth_env( &FeedbackRequestTags { diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 26b3fa158..dd2b1079f 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -1,6 +1,8 @@ use super::AuthRequestTelemetryContext; +use super::CompactConversationRequestSettings; use super::ModelClient; use super::PendingUnauthorizedRetry; +use super::Prompt; use super::UnauthorizedRecoveryExecution; use super::X_CODEX_INSTALLATION_ID_HEADER; use super::X_CODEX_PARENT_THREAD_ID_HEADER; @@ -13,11 +15,15 @@ use crate::GenerateAttestationFuture; use crate::responses_metadata::CodexResponsesMetadata; use crate::test_support::TestCodexResponsesRequestKind; use crate::test_support::responses_metadata as test_responses_metadata; +use codex_api::AgentIdentityTelemetry; use codex_api::ApiError; use codex_api::ResponseEvent; use codex_api::TransportError; +use codex_login::AuthCredentialsStoreMode; +use codex_login::AuthKeyringBackendKind; use codex_login::AuthManager; use codex_login::CodexAuth; +use codex_login::auth::AgentIdentityAuthPolicy; use codex_model_provider::BearerAuthProvider; use codex_model_provider::SharedModelProvider; use codex_model_provider::create_model_provider; @@ -28,6 +34,7 @@ use codex_model_provider_info::create_oss_provider_with_base_url; use codex_otel::SessionTelemetry; use codex_protocol::ThreadId; use codex_protocol::auth::AuthMode; +use codex_protocol::models::BaseInstructions; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ModelInfo; @@ -35,6 +42,7 @@ use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::InternalSessionSource; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; +use codex_rollout_trace::CompactionTraceContext; use codex_rollout_trace::ExecutionStatus; use codex_rollout_trace::InferenceTraceAttempt; use codex_rollout_trace::InferenceTraceContext; @@ -65,14 +73,27 @@ use tracing_subscriber::layer::Context as LayerContext; use tracing_subscriber::layer::SubscriberExt; use tracing_subscriber::registry::LookupSpan; use tracing_subscriber::util::SubscriberInitExt; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; +const TEST_CHATGPT_ID_TOKEN: &str = "eyJhbGciOiJub25lIiwidHlwIjoiSldUIn0.eyJlbWFpbCI6InVzZXJAZXhhbXBsZS5jb20iLCJlbWFpbF92ZXJpZmllZCI6dHJ1ZSwiaHR0cHM6Ly9hcGkub3BlbmFpLmNvbS9hdXRoIjp7ImNoYXRncHRfdXNlcl9pZCI6InVzZXItMTIzNDUiLCJ1c2VyX2lkIjoidXNlci0xMjM0NSIsImNoYXRncHRfcGxhbl90eXBlIjoicHJvIiwiY2hhdGdwdF9hY2NvdW50X2lkIjoiYWNjb3VudC0xMjMifX0.c2ln"; const TEST_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111"; fn test_model_client(session_source: SessionSource) -> ModelClient { + test_model_client_with_thread_id(ThreadId::new(), session_source) +} + +fn test_model_client_with_thread_id( + thread_id: ThreadId, + session_source: SessionSource, +) -> ModelClient { let provider = create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses); - let thread_id = ThreadId::new(); ModelClient::new( /*auth_manager*/ None, + AgentIdentityAuthPolicy::JwtOnly, thread_id, provider, session_source, @@ -86,6 +107,115 @@ fn test_model_client(session_source: SessionSource) -> ModelClient { ) } +#[tokio::test] +async fn compact_uses_bearer_after_agent_identity_session_fallback() -> anyhow::Result<()> { + let server = MockServer::start().await; + let registration_count = Arc::new(AtomicUsize::new(0)); + let response_count = Arc::clone(®istration_count); + Mock::given(method("POST")) + .and(path("/v1/agent/register")) + .respond_with(move |_request: &wiremock::Request| { + response_count.fetch_add(1, Ordering::SeqCst); + ResponseTemplate::new(/*status*/ 503) + }) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/v1/responses/compact")) + .respond_with(ResponseTemplate::new(/*status*/ 200).set_body_json(json!({ + "output": [] + }))) + .expect(/*requests*/ 1) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let auth_manager = chatgpt_auth_manager(&codex_home, server.uri()).await; + let mut provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None); + provider.base_url = Some(format!("{}/v1", server.uri())); + provider.supports_websockets = false; + let thread_id = ThreadId::new(); + let client = ModelClient::new( + Some(auth_manager), + AgentIdentityAuthPolicy::ChatGptAuth, + thread_id, + provider, + SessionSource::Cli, + "test_originator".to_string(), + /*model_verbosity*/ None, + /*enable_request_compression*/ false, + /*include_timing_metrics*/ false, + /*beta_features_header*/ None, + /*item_ids_enabled*/ false, + /*attestation_provider*/ None, + ); + let prompt = Prompt { + input: vec![ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: "please compact".to_string(), + }], + phase: None, + internal_chat_message_metadata_passthrough: None, + }], + base_instructions: BaseInstructions { + text: "base instructions".to_string(), + }, + ..Default::default() + }; + let responses_metadata = test_responses_metadata_for_client( + &client, + /*turn_id*/ None, + format!("{}:0", client.state.thread_id), + /*parent_thread_id*/ None, + TestCodexResponsesRequestKind::Turn, + ); + + let output = client + .compact_conversation_history( + &prompt, + &test_model_info(), + /*turn_state*/ None, + CompactConversationRequestSettings { + effort: None, + summary: codex_protocol::config_types::ReasoningSummary::None, + service_tier: None, + }, + &test_session_telemetry(), + &CompactionTraceContext::disabled(), + &responses_metadata, + ) + .await?; + + assert!(output.is_empty()); + assert_eq!(registration_count.load(Ordering::SeqCst), 3); + let requests = server + .received_requests() + .await + .expect("server should record requests"); + let compact_request = requests + .iter() + .find(|request| request.url.path() == "/v1/responses/compact") + .expect("compact request should be captured"); + assert_eq!( + compact_request + .headers + .get(http::header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()), + Some("Bearer test-access-token") + ); + assert_eq!( + compact_request + .headers + .get("ChatGPT-Account-ID") + .and_then(|value| value.to_str().ok()), + Some("account-123") + ); + + Ok(()) +} + fn test_model_provider() -> SharedModelProvider { test_model_client(SessionSource::Cli).state.provider.clone() } @@ -169,6 +299,45 @@ fn ultra_reasoning_uses_max_for_requests() { ); } +fn write_chatgpt_auth_json(codex_home: &std::path::Path) { + let auth_json = json!({ + "tokens": { + "id_token": TEST_CHATGPT_ID_TOKEN, + "access_token": "test-access-token", + "refresh_token": "test-refresh-token", + "account_id": "account-123" + }, + "last_refresh": "2099-01-01T00:00:00Z" + }); + std::fs::write( + codex_home.join("auth.json"), + serde_json::to_string_pretty(&auth_json).expect("serialize auth.json"), + ) + .expect("write auth.json"); +} + +async fn chatgpt_auth_manager( + codex_home: &TempDir, + agent_identity_authapi_base_url: String, +) -> Arc { + write_chatgpt_auth_json(codex_home.path()); + let auth_manager = AuthManager::shared( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + /*auth_route_config*/ None, + ) + .await; + let auth = auth_manager.auth().await.expect("auth should load"); + AuthManager::from_auth_for_testing_with_agent_identity_authapi_base_url( + auth, + agent_identity_authapi_base_url, + ) +} + #[derive(Default)] struct TagCollectorVisitor { tags: BTreeMap, @@ -572,6 +741,7 @@ fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() { let auth_context = AuthRequestTelemetryContext::new( Some(AuthMode::Chatgpt), &BearerAuthProvider::for_test(Some("access-token"), Some("workspace-123")), + /*agent_identity_telemetry*/ None, PendingUnauthorizedRetry::from_recovery(UnauthorizedRecoveryExecution { mode: "managed", phase: "refresh_token", @@ -586,6 +756,27 @@ fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() { assert_eq!(auth_context.recovery_phase, Some("refresh_token")); } +#[test] +fn auth_request_telemetry_context_tracks_agent_identity_ids() { + let auth_context = AuthRequestTelemetryContext::new( + Some(AuthMode::Chatgpt), + &BearerAuthProvider::for_test(/*token*/ None, /*account_id*/ None), + Some(AgentIdentityTelemetry { + agent_id: "agent-runtime-context".to_string(), + task_id: "task-run-context".to_string(), + }), + PendingUnauthorizedRetry::default(), + ); + + assert_eq!( + auth_context.agent_identity_telemetry(), + Some(&AgentIdentityTelemetry { + agent_id: "agent-runtime-context".to_string(), + task_id: "task-run-context".to_string(), + }) + ); +} + fn model_client_with_counting_attestation( include_attestation: bool, ) -> (ModelClient, Arc) { @@ -623,6 +814,7 @@ fn model_client_with_counting_attestation( }; let model_client = ModelClient::new( auth_manager, + AgentIdentityAuthPolicy::JwtOnly, ThreadId::new(), provider, SessionSource::Exec, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 216ebd1f9..78abe9df8 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -8,6 +8,7 @@ use crate::shell_snapshot::ShellSnapshot; use crate::skills::SkillError; use crate::state::ActiveTurn; use codex_extension_api::ExtensionDataInit; +use codex_login::auth::AgentIdentityAuthPolicy; use codex_protocol::SessionId; use codex_protocol::config_types::SERVICE_TIER_DEFAULT_REQUEST_VALUE; use codex_protocol::config_types::ServiceTier; @@ -1057,6 +1058,11 @@ impl Session { time_provider, model_client: ModelClient::new( Some(Arc::clone(&auth_manager)), + if config.features.enabled(Feature::UseAgentIdentity) { + AgentIdentityAuthPolicy::ChatGptAuth + } else { + AgentIdentityAuthPolicy::JwtOnly + }, thread_id, session_configuration.provider.clone(), session_configuration.session_source.clone(), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index c6772373d..1030a64e6 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -30,6 +30,7 @@ use core_test_support::test_codex::local_selections; use codex_features::Feature; use codex_login::CodexAuth; +use codex_login::auth::AgentIdentityAuthPolicy; use codex_model_provider_info::ModelProviderInfo; use codex_models_manager::bundled_models_response; use codex_models_manager::model_info; @@ -469,6 +470,7 @@ fn test_model_client_session() -> crate::client::ModelClientSession { .expect("test thread id should be valid"); crate::client::ModelClient::new( /*auth_manager*/ None, + AgentIdentityAuthPolicy::JwtOnly, thread_id, ModelProviderInfo::create_openai_provider(/* base_url */ /*base_url*/ None), codex_protocol::protocol::SessionSource::Exec, @@ -5415,6 +5417,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { time_provider: Arc::new(crate::current_time::SystemTimeProvider), model_client: ModelClient::new( Some(auth_manager.clone()), + AgentIdentityAuthPolicy::JwtOnly, thread_id, session_configuration.provider.clone(), session_configuration.session_source.clone(), @@ -7492,6 +7495,7 @@ where time_provider: Arc::new(crate::current_time::SystemTimeProvider), model_client: ModelClient::new( Some(Arc::clone(&auth_manager)), + AgentIdentityAuthPolicy::JwtOnly, thread_id, session_configuration.provider.clone(), session_configuration.session_source.clone(), diff --git a/codex-rs/core/src/session_startup_prewarm.rs b/codex-rs/core/src/session_startup_prewarm.rs index af9129734..8ce5229a3 100644 --- a/codex-rs/core/src/session_startup_prewarm.rs +++ b/codex-rs/core/src/session_startup_prewarm.rs @@ -183,6 +183,14 @@ impl SessionStartupPrewarmHandle { impl Session { pub(crate) async fn schedule_startup_prewarm(self: &Arc, base_instructions: String) { if !self.services.model_client.responses_websocket_enabled() { + // Without websocket prewarm, resolve auth once so Agent Identity bootstrap can + // register or engage this session's bearer fallback before the first user request. + let model_client = self.services.model_client.clone(); + tokio::spawn(async move { + if let Err(err) = model_client.prewarm_auth().await { + warn!("startup auth prewarm failed: {err:#}"); + } + }); return; } diff --git a/codex-rs/core/tests/responses_headers.rs b/codex-rs/core/tests/responses_headers.rs index 7e213991c..5b7c2c65b 100644 --- a/codex-rs/core/tests/responses_headers.rs +++ b/codex-rs/core/tests/responses_headers.rs @@ -5,6 +5,7 @@ use codex_core::ModelClient; use codex_core::Prompt; use codex_core::ResponseEvent; use codex_login::CodexAuth; +use codex_login::auth::AgentIdentityAuthPolicy; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::WireApi; use codex_otel::SessionTelemetry; @@ -120,6 +121,7 @@ async fn responses_stream_includes_subagent_header_on_review() { let client = ModelClient::new( /*auth_manager*/ None, + AgentIdentityAuthPolicy::JwtOnly, thread_id, provider.clone(), session_source.clone(), @@ -253,6 +255,7 @@ async fn responses_stream_includes_subagent_header_on_other() { let client = ModelClient::new( /*auth_manager*/ None, + AgentIdentityAuthPolicy::JwtOnly, thread_id, provider.clone(), session_source.clone(), @@ -372,6 +375,7 @@ async fn responses_respects_model_info_overrides_from_config() { let client = ModelClient::new( /*auth_manager*/ None, + AgentIdentityAuthPolicy::JwtOnly, thread_id, provider.clone(), session_source.clone(), diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index f934a2434..4df2bca1c 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -12,6 +12,7 @@ use codex_features::Feature; use codex_login::AuthKeyringBackendKind; use codex_login::AuthManager; use codex_login::CodexAuth; +use codex_login::auth::AgentIdentityAuthPolicy; use codex_login::default_client::originator; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::WireApi; @@ -1220,6 +1221,7 @@ async fn send_provider_auth_request(server: &MockServer, auth: ModelProviderAuth Some(AuthManager::from_auth_for_testing(CodexAuth::from_api_key( "unused-api-key", ))), + AgentIdentityAuthPolicy::JwtOnly, thread_id, provider, SessionSource::Exec, @@ -2833,6 +2835,7 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() { let client = ModelClient::new( /*auth_manager*/ None, + AgentIdentityAuthPolicy::JwtOnly, thread_id, provider.clone(), SessionSource::Exec, diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index b1b731a08..ed569805d 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -9,6 +9,7 @@ use codex_core::ResponseEvent; use codex_core::X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER; use codex_features::Feature; use codex_login::CodexAuth; +use codex_login::auth::AgentIdentityAuthPolicy; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::WireApi; use codex_otel::MetricsClient; @@ -2189,6 +2190,7 @@ async fn websocket_harness_with_provider_options( let summary = ReasoningSummary::Auto; let client = ModelClient::new( /*auth_manager*/ None, + AgentIdentityAuthPolicy::JwtOnly, thread_id, provider.clone(), SessionSource::Exec, diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index ececb02b2..d0e69d656 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -5,6 +5,9 @@ use anyhow::Result; use codex_core::compact::SUMMARY_PREFIX; use codex_features::Feature; use codex_login::CodexAuth; +use codex_login::auth::AgentIdentityAuth; +use codex_login::auth::AgentIdentityAuthRecord; +use codex_protocol::account::PlanType as AccountPlanType; use codex_protocol::config_types::ServiceTier; use codex_protocol::dynamic_tools::DynamicToolFunctionSpec; use codex_protocol::dynamic_tools::DynamicToolNamespaceSpec; @@ -133,6 +136,8 @@ const PRETURN_CONTEXT_DIFF_CWD: &str = "/tmp/PRETURN_CONTEXT_DIFF_CWD"; const DUMMY_FUNCTION_NAME: &str = "test_tool"; const TURN_STATE_HEADER: &str = "x-codex-turn-state"; const REMOTE_COMPACT_TURN_COMPLETE_TIMEOUT: Duration = Duration::from_secs(30); +const TEST_AGENT_IDENTITY_PRIVATE_KEY: &str = + "MC4CAQAwBQYDK2VwBCIEIJ7kFBaOujmoz1gvBNEC+BeM2IX87FFB0xmISOZ/XO0c"; fn summary_with_prefix(summary: &str) -> String { format!("{SUMMARY_PREFIX}\n{summary}") @@ -532,6 +537,79 @@ async fn remote_compact_replaces_history_for_followups() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn remote_compact_uses_agent_identity_assertion() -> Result<()> { + skip_if_no_network!(Ok(())); + + let harness = TestCodexHarness::with_builder( + test_codex().with_auth(CodexAuth::AgentIdentity( + AgentIdentityAuth::from_record( + AgentIdentityAuthRecord { + agent_runtime_id: "agent-runtime-compact".to_string(), + agent_private_key: TEST_AGENT_IDENTITY_PRIVATE_KEY.to_string(), + account_id: "account-compact".to_string(), + chatgpt_user_id: "user-compact".to_string(), + email: Some("agent@example.com".to_string()), + plan_type: AccountPlanType::Plus, + chatgpt_account_is_fedramp: false, + task_id: Some("task-compact".to_string()), + }, + "https://auth.openai.com/api/accounts", + /*auth_route_config*/ None, + ) + .await?, + )), + ) + .await?; + let codex = harness.test().codex.clone(); + + let _responses_mock = responses::mount_sse_once( + harness.server(), + responses::sse(vec![ + responses::ev_assistant_message("m1", "REMOTE_REPLY"), + responses::ev_completed("resp-1"), + ]), + ) + .await; + let compact_mock = responses::mount_compact_json_once( + harness.server(), + serde_json::json!({ "output": compacted_summary_only_output("COMPACTED") }), + ) + .await; + + codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "hello remote compact".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + wait_for_turn_complete(&codex).await; + + codex.submit(Op::Compact).await?; + wait_for_turn_complete(&codex).await; + + let compact_request = compact_mock.single_request(); + assert_eq!(compact_request.path(), "/v1/responses/compact"); + assert!( + compact_request + .header("authorization") + .is_some_and(|value| value.starts_with("AgentAssertion ")), + "compact request should use task-scoped AgentAssertion auth" + ); + assert_eq!( + compact_request.header("chatgpt-account-id").as_deref(), + Some("account-compact") + ); + + Ok(()) +} + async fn assert_remote_manual_compact_request_parity( auth: CodexAuth, configured_service_tier: Option, diff --git a/codex-rs/login/src/auth/agent_identity.rs b/codex-rs/login/src/auth/agent_identity.rs index 79ee2e7f5..15cdb1116 100644 --- a/codex-rs/login/src/auth/agent_identity.rs +++ b/codex-rs/login/src/auth/agent_identity.rs @@ -46,7 +46,7 @@ pub(super) fn require_agent_identity_authapi_base_url( }) } -#[derive(Debug, Error)] +#[derive(Clone, Debug, Error)] pub enum AgentIdentityAuthError { #[error( "agent identity bootstrap unavailable after {attempts} attempts during {operation}: {message}" @@ -59,13 +59,14 @@ pub enum AgentIdentityAuthError { } impl AgentIdentityAuthError { - pub fn is_bootstrap_unavailable(error: &std::io::Error) -> bool { - matches!( - error - .get_ref() - .and_then(|source| source.downcast_ref::()), - Some(Self::BootstrapUnavailable { .. }) - ) + pub(super) fn bootstrap_unavailable(error: &std::io::Error) -> Option<&Self> { + match error + .get_ref() + .and_then(|source| source.downcast_ref::()) + { + Some(error @ Self::BootstrapUnavailable { .. }) => Some(error), + None => None, + } } } diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index d52563f02..9509cab45 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -584,7 +584,7 @@ async fn chatgpt_auth_registration_retry_exhaustion_is_fallback_eligible() -> an .await .expect_err("retry exhaustion should return an error"); - assert!(AgentIdentityAuthError::is_bootstrap_unavailable(&err)); + assert!(AgentIdentityAuthError::bootstrap_unavailable(&err).is_some()); assert!( auth.stored_managed_chatgpt_agent_identity_record("account-123") .is_none() @@ -648,7 +648,7 @@ async fn chatgpt_auth_task_registration_retry_exhaustion_is_fallback_eligible() .await .expect_err("task retry exhaustion should return an error"); - assert!(AgentIdentityAuthError::is_bootstrap_unavailable(&err)); + assert!(AgentIdentityAuthError::bootstrap_unavailable(&err).is_some()); record.task_id = None; assert_eq!( auth.stored_managed_chatgpt_agent_identity_record("account-123"), @@ -701,7 +701,7 @@ async fn chatgpt_auth_non_retryable_registration_error_is_hard_failure() -> anyh .await .expect_err("hard registration failure should return an error"); - assert!(!AgentIdentityAuthError::is_bootstrap_unavailable(&err)); + assert!(AgentIdentityAuthError::bootstrap_unavailable(&err).is_none()); assert!( auth.stored_managed_chatgpt_agent_identity_record("account-123") .is_none() @@ -742,7 +742,7 @@ async fn agent_identity_jwt_task_registration_retry_exhaustion_is_strict() -> an .await .expect_err("agent identity jwt task retry exhaustion should fail"); - assert!(!AgentIdentityAuthError::is_bootstrap_unavailable(&err)); + assert!(AgentIdentityAuthError::bootstrap_unavailable(&err).is_none()); Ok(()) } diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 8418c7a2a..844410874 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -15,6 +15,8 @@ use std::sync::Mutex; use std::sync::RwLock; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; +use std::time::Duration; +use std::time::Instant; use tokio::sync::Semaphore; use tokio::sync::watch; use tracing::instrument; @@ -83,6 +85,63 @@ pub enum AgentIdentityAuthPolicy { ChatGptAuth, } +const AGENT_IDENTITY_BOOTSTRAP_FAILURE_COOLDOWN: Duration = Duration::from_secs(60 * 60); + +#[derive(Debug)] +struct CachedAgentIdentityBootstrapFailure { + account_id: String, + authapi_base_url: String, + retry_at: Instant, + error: AgentIdentityAuthError, +} + +#[derive(Debug, Default)] +struct AgentIdentityBootstrapCooldown { + failure: Option, +} + +impl AgentIdentityBootstrapCooldown { + fn error_for( + &mut self, + account_id: &str, + authapi_base_url: &str, + now: Instant, + ) -> Option { + let error = self + .failure + .as_ref() + .filter(|failure| { + failure.account_id == account_id + && failure.authapi_base_url == authapi_base_url + && failure.retry_at > now + }) + .map(|failure| failure.error.clone()); + if error.is_none() { + self.clear(); + } + error + } + + fn record_failure( + &mut self, + account_id: String, + authapi_base_url: String, + error: AgentIdentityAuthError, + now: Instant, + ) { + self.failure = Some(CachedAgentIdentityBootstrapFailure { + account_id, + authapi_base_url, + retry_at: now + AGENT_IDENTITY_BOOTSTRAP_FAILURE_COOLDOWN, + error, + }); + } + + fn clear(&mut self) { + self.failure = None; + } +} + impl PartialEq for CodexAuth { fn eq(&self, other: &Self) -> bool { match (self, other) { @@ -1742,6 +1801,7 @@ pub struct AuthManager { agent_identity_authapi_base_url: Option, refresh_lock: Semaphore, agent_identity_lock: Semaphore, + agent_identity_bootstrap_cooldown: Mutex, external_auth: RwLock>>, auth_route_config: Option, } @@ -1843,6 +1903,7 @@ impl AuthManager { agent_identity_authapi_base_url, refresh_lock: Semaphore::new(/*permits*/ 1), agent_identity_lock: Semaphore::new(/*permits*/ 1), + agent_identity_bootstrap_cooldown: Mutex::default(), external_auth: RwLock::new(None), auth_route_config, } @@ -1868,6 +1929,7 @@ impl AuthManager { agent_identity_authapi_base_url: default_agent_identity_authapi_base_url(), refresh_lock: Semaphore::new(/*permits*/ 1), agent_identity_lock: Semaphore::new(/*permits*/ 1), + agent_identity_bootstrap_cooldown: Mutex::default(), external_auth: RwLock::new(None), auth_route_config: None, }) @@ -1892,6 +1954,7 @@ impl AuthManager { agent_identity_authapi_base_url: default_agent_identity_authapi_base_url(), refresh_lock: Semaphore::new(/*permits*/ 1), agent_identity_lock: Semaphore::new(/*permits*/ 1), + agent_identity_bootstrap_cooldown: Mutex::default(), external_auth: RwLock::new(None), auth_route_config: None, }) @@ -1924,6 +1987,7 @@ impl AuthManager { ), refresh_lock: Semaphore::new(/*permits*/ 1), agent_identity_lock: Semaphore::new(/*permits*/ 1), + agent_identity_bootstrap_cooldown: Mutex::default(), external_auth: RwLock::new(None), auth_route_config: None, }) @@ -1946,6 +2010,7 @@ impl AuthManager { agent_identity_authapi_base_url: default_agent_identity_authapi_base_url(), refresh_lock: Semaphore::new(/*permits*/ 1), agent_identity_lock: Semaphore::new(/*permits*/ 1), + agent_identity_bootstrap_cooldown: Mutex::default(), external_auth: RwLock::new(Some( Arc::new(BearerTokenRefresher::new(config)) as Arc )), @@ -2006,15 +2071,44 @@ impl AuthManager { .acquire() .await .map_err(std::io::Error::other)?; - return auth + let forced_chatgpt_workspace_id = self.forced_chatgpt_workspace_id(); + let cooldown_key = ManagedChatGptAgentIdentityBinding::from_auth( + &auth, + forced_chatgpt_workspace_id.clone(), + ) + .and_then(|binding| { + self.agent_identity_authapi_base_url + .as_ref() + .map(|base_url| (binding.account_id, base_url.clone())) + }); + if let Some((account_id, authapi_base_url)) = cooldown_key.as_ref() + && let Ok(mut cooldown) = self.agent_identity_bootstrap_cooldown.lock() + && let Some(error) = + cooldown.error_for(account_id, authapi_base_url, Instant::now()) + { + tracing::warn!("agent identity bootstrap retry suppressed during shared cooldown"); + return Err(std::io::Error::other(error)); + } + + let result = auth .agent_identity_auth( policy, self.agent_identity_authapi_base_url.as_deref(), - self.forced_chatgpt_workspace_id(), + forced_chatgpt_workspace_id, self.auth_route_config.as_ref(), session_source, ) .await; + if let Ok(mut cooldown) = self.agent_identity_bootstrap_cooldown.lock() { + if let (Err(err), Some((account_id, authapi_base_url))) = (&result, cooldown_key) + && let Some(error) = AgentIdentityAuthError::bootstrap_unavailable(err).cloned() + { + cooldown.record_failure(account_id, authapi_base_url, error, Instant::now()); + } else { + cooldown.clear(); + } + } + return result; } auth.agent_identity_auth( policy, diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs index 3d703e317..0ee0313b5 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -12,6 +12,7 @@ use codex_core::resolve_installation_id; use codex_features::Feature; use codex_login::AuthManager; use codex_login::CodexAuth; +use codex_login::auth::AgentIdentityAuthPolicy; use codex_login::auth_env_telemetry::collect_auth_env_telemetry; use codex_login::default_client::originator; use codex_model_provider::ModelProvider; @@ -250,6 +251,7 @@ impl MemoryStartupContext { let session_id_string = session_id.to_string(); let model_client = ModelClient::new( Some(Arc::clone(&self.auth_manager)), + AgentIdentityAuthPolicy::JwtOnly, self.thread_id, config.model_provider.clone(), session_source.clone(), diff --git a/codex-rs/model-provider/src/auth.rs b/codex-rs/model-provider/src/auth.rs index 555312027..5b667e454 100644 --- a/codex-rs/model-provider/src/auth.rs +++ b/codex-rs/model-provider/src/auth.rs @@ -1,13 +1,20 @@ use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use codex_agent_identity::AgentIdentityKey; use codex_agent_identity::authorization_header_for_agent_task; +use codex_api::AgentIdentityTelemetry; use codex_api::AuthProvider; use codex_api::SharedAuthProvider; use codex_login::AuthManager; use codex_login::CodexAuth; +use codex_login::auth::AgentIdentityAuth; +use codex_login::auth::AgentIdentityAuthError; +use codex_login::auth::AgentIdentityAuthPolicy; use codex_model_provider_info::ModelProviderInfo; use codex_protocol::error::CodexErr; +use codex_protocol::protocol::SessionSource; use http::HeaderMap; use http::HeaderValue; @@ -16,9 +23,62 @@ use crate::bearer_auth_provider::BearerAuthProvider; const BEDROCK_API_KEY_UNSUPPORTED_MESSAGE: &str = "Bedrock API key auth is only supported by the Amazon Bedrock model provider"; +#[derive(Clone, Debug)] +pub struct ProviderAuthScope { + pub agent_identity_policy: AgentIdentityAuthPolicy, + pub session_source: SessionSource, + pub agent_identity_session_fallback: AgentIdentitySessionFallback, +} + +#[derive(Clone, Debug, Default)] +pub struct AgentIdentitySessionFallback { + engaged: Arc, +} + +impl AgentIdentitySessionFallback { + pub fn is_engaged(&self) -> bool { + self.engaged.load(Ordering::Relaxed) + } + + fn engage(&self) -> bool { + !self.engaged.swap(true, Ordering::Relaxed) + } +} + +/// Provider auth resolved for a request, plus metadata describing the effective auth. +#[derive(Clone)] +pub struct ResolvedProviderAuth { + pub auth: SharedAuthProvider, + pub agent_identity_telemetry: Option, +} + +impl ResolvedProviderAuth { + pub(crate) fn new(auth: SharedAuthProvider) -> Self { + Self { + auth, + agent_identity_telemetry: None, + } + } + + fn for_agent_identity(auth: AgentIdentityAuth) -> Self { + let agent_identity_telemetry = agent_identity_telemetry(&auth); + Self { + auth: Arc::new(AgentIdentityAuthProvider { auth }), + agent_identity_telemetry: Some(agent_identity_telemetry), + } + } +} + +pub(crate) fn agent_identity_telemetry(auth: &AgentIdentityAuth) -> AgentIdentityTelemetry { + AgentIdentityTelemetry { + agent_id: auth.record().agent_runtime_id.clone(), + task_id: auth.run_task_id().to_string(), + } +} + #[derive(Clone, Debug)] struct AgentIdentityAuthProvider { - auth: codex_login::auth::AgentIdentityAuth, + auth: AgentIdentityAuth, } impl AuthProvider for AgentIdentityAuthProvider { @@ -95,6 +155,74 @@ pub(crate) fn resolve_provider_auth( }) } +pub(crate) async fn resolve_provider_auth_for_scope( + auth_manager: Option>, + auth: Option<&CodexAuth>, + provider: &ModelProviderInfo, + scope: ProviderAuthScope, +) -> codex_protocol::error::Result { + let ProviderAuthScope { + agent_identity_policy, + session_source, + agent_identity_session_fallback, + } = scope; + if let Some(CodexAuth::AgentIdentity(agent_identity_auth)) = auth { + return Ok(ResolvedProviderAuth::for_agent_identity( + agent_identity_auth.clone(), + )); + } + + if !should_bootstrap_chatgpt_agent_identity(agent_identity_policy, auth) + || agent_identity_session_fallback.is_engaged() + { + return resolve_provider_auth(auth, provider).map(ResolvedProviderAuth::new); + } + + let Some(auth_manager) = auth_manager else { + return resolve_provider_auth(auth, provider).map(ResolvedProviderAuth::new); + }; + + match auth_manager + .agent_identity_auth(agent_identity_policy, session_source) + .await + { + Ok(Some(agent_identity_auth)) => Ok(ResolvedProviderAuth::for_agent_identity( + agent_identity_auth, + )), + Ok(None) => resolve_provider_auth(auth, provider).map(ResolvedProviderAuth::new), + Err(err) => { + if let Some(AgentIdentityAuthError::BootstrapUnavailable { + operation, + attempts, + message, + }) = err + .get_ref() + .and_then(|source| source.downcast_ref::()) + { + let newly_engaged = agent_identity_session_fallback.engage(); + tracing::warn!( + operation, + attempts = *attempts, + error = %message, + newly_engaged, + "agent identity bootstrap unavailable; using ChatGPT bearer auth for this session" + ); + resolve_provider_auth(auth, provider).map(ResolvedProviderAuth::new) + } else { + Err(err.into()) + } + } + } +} + +fn should_bootstrap_chatgpt_agent_identity( + agent_identity_policy: AgentIdentityAuthPolicy, + auth: Option<&CodexAuth>, +) -> bool { + agent_identity_policy == AgentIdentityAuthPolicy::ChatGptAuth + && matches!(auth, Some(CodexAuth::Chatgpt(_))) +} + fn bearer_auth_for_provider( provider: &ModelProviderInfo, ) -> codex_protocol::error::Result> { @@ -129,13 +257,128 @@ pub fn auth_provider_from_auth(auth: &CodexAuth) -> SharedAuthProvider { #[cfg(test)] mod tests { + use codex_agent_identity::generate_agent_key_material; + use codex_login::AuthCredentialsStoreMode; + use codex_login::AuthKeyringBackendKind; + use codex_login::auth::AgentIdentityAuthRecord; use codex_login::auth::BedrockApiKeyAuth; use codex_model_provider_info::WireApi; use codex_model_provider_info::create_oss_provider_with_base_url; + use codex_protocol::account::PlanType; use pretty_assertions::assert_eq; + use serde_json::json; + use std::path::Path; + use std::path::PathBuf; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::method; + use wiremock::matchers::path; use super::*; + static NEXT_CODEX_HOME_ID: AtomicUsize = AtomicUsize::new(0); + const TEST_CHATGPT_ID_TOKEN: &str = "eyJhbGciOiJub25lIiwidHlwIjoiSldUIn0.eyJlbWFpbCI6InVzZXJAZXhhbXBsZS5jb20iLCJlbWFpbF92ZXJpZmllZCI6dHJ1ZSwiaHR0cHM6Ly9hcGkub3BlbmFpLmNvbS9hdXRoIjp7ImNoYXRncHRfdXNlcl9pZCI6InVzZXItMTIzNDUiLCJ1c2VyX2lkIjoidXNlci0xMjM0NSIsImNoYXRncHRfcGxhbl90eXBlIjoicHJvIiwiY2hhdGdwdF9hY2NvdW50X2lkIjoiYWNjb3VudC0xMjMifX0.c2ln"; + + async fn agent_identity_auth(chatgpt_account_is_fedramp: bool) -> AgentIdentityAuth { + let key_material = generate_agent_key_material().expect("generate key material"); + AgentIdentityAuth::from_record( + AgentIdentityAuthRecord { + agent_runtime_id: "agent-runtime-1".to_string(), + agent_private_key: key_material.private_key_pkcs8_base64, + account_id: "account-1".to_string(), + chatgpt_user_id: "user-1".to_string(), + email: Some("agent@example.com".to_string()), + plan_type: PlanType::Plus, + chatgpt_account_is_fedramp, + task_id: Some("task-run-1".to_string()), + }, + "https://auth.openai.com/api/accounts", + /*auth_route_config*/ None, + ) + .await + .expect("agent identity auth record should include task id") + } + + fn provider_auth_scope( + policy: AgentIdentityAuthPolicy, + fallback: AgentIdentitySessionFallback, + ) -> ProviderAuthScope { + ProviderAuthScope { + agent_identity_policy: policy, + session_source: SessionSource::Cli, + agent_identity_session_fallback: fallback, + } + } + + fn test_codex_home() -> PathBuf { + let id = NEXT_CODEX_HOME_ID.fetch_add(1, Ordering::Relaxed); + let path = std::env::temp_dir().join(format!( + "codex-model-provider-agent-identity-{pid}-{id}", + pid = std::process::id() + )); + let _ = std::fs::remove_dir_all(&path); + std::fs::create_dir_all(&path).expect("create temp codex home"); + path + } + + fn write_chatgpt_auth_json(codex_home: &Path) { + let auth_json = json!({ + "tokens": { + "id_token": TEST_CHATGPT_ID_TOKEN, + "access_token": "test-access-token", + "refresh_token": "test-refresh-token", + "account_id": "account-123" + }, + "last_refresh": "2099-01-01T00:00:00Z" + }); + std::fs::write( + codex_home.join("auth.json"), + serde_json::to_string_pretty(&auth_json).expect("serialize auth.json"), + ) + .expect("write auth.json"); + } + + async fn chatgpt_auth_manager( + agent_identity_authapi_base_url: String, + ) -> (PathBuf, Arc, CodexAuth) { + let codex_home = test_codex_home(); + write_chatgpt_auth_json(&codex_home); + let auth_manager = AuthManager::shared( + codex_home.clone(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + /*auth_route_config*/ None, + ) + .await; + let auth = auth_manager.auth().await.expect("auth should load"); + let auth_manager = AuthManager::from_auth_for_testing_with_agent_identity_authapi_base_url( + auth.clone(), + agent_identity_authapi_base_url, + ); + (codex_home, auth_manager, auth) + } + + async fn mount_transient_agent_registration( + server: &MockServer, + status: u16, + registration_count: Arc, + ) { + Mock::given(method("POST")) + .and(path("/v1/agent/register")) + .respond_with(move |_request: &wiremock::Request| { + registration_count.fetch_add(1, Ordering::SeqCst); + ResponseTemplate::new(status) + }) + .mount(server) + .await; + } + #[test] fn unauthenticated_auth_provider_adds_no_headers() { let provider = @@ -161,4 +404,180 @@ mod tests { Ok(_) => panic!("Bedrock API key auth should be rejected"), } } + + #[tokio::test] + async fn first_party_run_scope_uses_agent_assertion_and_exposes_telemetry() { + let auth = CodexAuth::AgentIdentity( + agent_identity_auth(/*chatgpt_account_is_fedramp*/ false).await, + ); + let provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None); + + let auth = resolve_provider_auth_for_scope( + /*auth_manager*/ None, + Some(&auth), + &provider, + provider_auth_scope( + AgentIdentityAuthPolicy::JwtOnly, + AgentIdentitySessionFallback::default(), + ), + ) + .await + .expect("auth should resolve"); + + assert_eq!( + auth.agent_identity_telemetry, + Some(AgentIdentityTelemetry { + agent_id: "agent-runtime-1".to_string(), + task_id: "task-run-1".to_string(), + }) + ); + let headers = auth.auth.to_auth_headers(); + assert!( + headers + .get(http::header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| value.starts_with("AgentAssertion ")) + ); + } + + #[tokio::test] + async fn agent_identity_auth_provider_preserves_account_routing_headers() { + let auth = agent_identity_auth(/*chatgpt_account_is_fedramp*/ true).await; + let provider = auth_provider_from_auth(&CodexAuth::AgentIdentity(auth)); + + let headers = provider.to_auth_headers(); + + assert!( + headers + .get(http::header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| value.starts_with("AgentAssertion ")) + ); + assert_eq!( + headers + .get("ChatGPT-Account-ID") + .and_then(|value| value.to_str().ok()), + Some("account-1") + ); + assert_eq!( + headers + .get("X-OpenAI-Fedramp") + .and_then(|value| value.to_str().ok()), + Some("true") + ); + } + + #[tokio::test] + async fn chatgpt_bootstrap_unavailable_uses_session_bearer_fallback() { + let server = MockServer::start().await; + let registration_count = Arc::new(AtomicUsize::new(0)); + mount_transient_agent_registration( + &server, + /*status*/ 503, + Arc::clone(®istration_count), + ) + .await; + let (_codex_home, auth_manager, auth) = chatgpt_auth_manager(server.uri()).await; + let provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None); + let fallback = AgentIdentitySessionFallback::default(); + + let provider_auth = resolve_provider_auth_for_scope( + Some(auth_manager), + Some(&auth), + &provider, + provider_auth_scope(AgentIdentityAuthPolicy::ChatGptAuth, fallback.clone()), + ) + .await + .expect("fallback should resolve bearer auth"); + + let headers = provider_auth.auth.to_auth_headers(); + assert_eq!( + headers + .get(http::header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()), + Some("Bearer test-access-token") + ); + assert_eq!( + headers + .get("ChatGPT-Account-ID") + .and_then(|value| value.to_str().ok()), + Some("account-123") + ); + assert!(fallback.is_engaged()); + assert_eq!(registration_count.load(Ordering::SeqCst), 3); + } + + #[tokio::test] + async fn chatgpt_session_fallback_skips_later_agent_identity_bootstrap() { + let server = MockServer::start().await; + let registration_count = Arc::new(AtomicUsize::new(0)); + mount_transient_agent_registration( + &server, + /*status*/ 503, + Arc::clone(®istration_count), + ) + .await; + let (_codex_home, auth_manager, auth) = chatgpt_auth_manager(server.uri()).await; + let provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None); + let fallback = AgentIdentitySessionFallback::default(); + + resolve_provider_auth_for_scope( + Some(Arc::clone(&auth_manager)), + Some(&auth), + &provider, + provider_auth_scope(AgentIdentityAuthPolicy::ChatGptAuth, fallback.clone()), + ) + .await + .expect("first fallback should resolve bearer auth"); + resolve_provider_auth_for_scope( + Some(auth_manager), + Some(&auth), + &provider, + provider_auth_scope(AgentIdentityAuthPolicy::ChatGptAuth, fallback), + ) + .await + .expect("second fallback should resolve bearer auth"); + + assert_eq!(registration_count.load(Ordering::SeqCst), 3); + } + + #[tokio::test] + async fn chatgpt_sessions_share_bootstrap_failure_cooldown() { + let server = MockServer::start().await; + let registration_count = Arc::new(AtomicUsize::new(0)); + mount_transient_agent_registration( + &server, + /*status*/ 503, + Arc::clone(®istration_count), + ) + .await; + let (_codex_home, auth_manager, auth) = chatgpt_auth_manager(server.uri()).await; + let provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None); + let first_fallback = AgentIdentitySessionFallback::default(); + let second_fallback = AgentIdentitySessionFallback::default(); + + resolve_provider_auth_for_scope( + Some(Arc::clone(&auth_manager)), + Some(&auth), + &provider, + provider_auth_scope(AgentIdentityAuthPolicy::ChatGptAuth, first_fallback.clone()), + ) + .await + .expect("first session fallback should resolve bearer auth"); + resolve_provider_auth_for_scope( + Some(auth_manager), + Some(&auth), + &provider, + provider_auth_scope( + AgentIdentityAuthPolicy::ChatGptAuth, + second_fallback.clone(), + ), + ) + .await + .expect("second session fallback should resolve bearer auth"); + + assert!(first_fallback.is_engaged()); + assert!(second_fallback.is_engaged()); + assert_eq!(registration_count.load(Ordering::SeqCst), 3); + } } diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index 44ce878e2..15967c4fe 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -4,6 +4,9 @@ mod bearer_auth_provider; mod models_endpoint; mod provider; +pub use auth::AgentIdentitySessionFallback; +pub use auth::ProviderAuthScope; +pub use auth::ResolvedProviderAuth; pub use auth::auth_provider_from_auth; pub use auth::unauthenticated_auth_provider; pub use bearer_auth_provider::BearerAuthProvider; diff --git a/codex-rs/model-provider/src/models_endpoint.rs b/codex-rs/model-provider/src/models_endpoint.rs index ebb51ea56..9a203c26e 100644 --- a/codex-rs/model-provider/src/models_endpoint.rs +++ b/codex-rs/model-provider/src/models_endpoint.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use std::time::Duration; +use codex_api::AgentIdentityTelemetry; use codex_api::ModelsClient; use codex_api::RequestTelemetry; use codex_api::ReqwestTransport; @@ -26,6 +27,7 @@ use codex_response_debug_context::telemetry_transport_error_message; use http::HeaderMap; use tokio::time::timeout; +use crate::auth::agent_identity_telemetry; use crate::auth::resolve_provider_auth; const MODELS_REFRESH_TIMEOUT: Duration = Duration::from_secs(5); @@ -75,10 +77,16 @@ impl OpenAiModelsEndpoint { let api_auth = resolve_provider_auth(auth.as_ref(), &self.provider_info)?; let transport = ReqwestTransport::new(build_reqwest_client()); let auth_telemetry = auth_header_telemetry(api_auth.as_ref()); + let agent_identity_telemetry = if let Some(CodexAuth::AgentIdentity(auth)) = auth.as_ref() { + Some(agent_identity_telemetry(auth)) + } else { + None + }; let request_telemetry: Arc = Arc::new(ModelsRequestTelemetry { auth_mode: auth_mode.map(|mode| TelemetryAuthMode::from(mode).to_string()), auth_header_attached: auth_telemetry.attached, auth_header_name: auth_telemetry.name, + agent_identity_telemetry, auth_env: self.auth_env(), }); let client = ModelsClient::new(transport, api_provider, api_auth) @@ -124,6 +132,7 @@ struct ModelsRequestTelemetry { auth_mode: Option, auth_header_attached: bool, auth_header_name: Option<&'static str>, + agent_identity_telemetry: Option, auth_env: AuthEnvTelemetry, } @@ -164,6 +173,8 @@ impl RequestTelemetry for ModelsRequestTelemetry { auth.error = response_debug.auth_error.as_deref(), auth.error_code = response_debug.auth_error_code.as_deref(), auth.mode = self.auth_mode.as_deref(), + auth.agent_id = self.agent_identity_telemetry.as_ref().map(|metadata| metadata.agent_id.as_str()), + auth.task_id = self.agent_identity_telemetry.as_ref().map(|metadata| metadata.task_id.as_str()), ); tracing::event!( target: "codex_otel.trace_safe", @@ -188,6 +199,8 @@ impl RequestTelemetry for ModelsRequestTelemetry { auth.error = response_debug.auth_error.as_deref(), auth.error_code = response_debug.auth_error_code.as_deref(), auth.mode = self.auth_mode.as_deref(), + auth.agent_id = self.agent_identity_telemetry.as_ref().map(|metadata| metadata.agent_id.as_str()), + auth.task_id = self.agent_identity_telemetry.as_ref().map(|metadata| metadata.task_id.as_str()), ); emit_feedback_request_tags_with_auth_env( &FeedbackRequestTags { diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index 5df5a88ea..6e58f4b74 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -18,8 +18,11 @@ use codex_protocol::error::CodexErr; use codex_protocol::openai_models::ModelsResponse; use crate::amazon_bedrock::AmazonBedrockModelProvider; +use crate::auth::ProviderAuthScope; +use crate::auth::ResolvedProviderAuth; use crate::auth::auth_manager_for_provider; use crate::auth::resolve_provider_auth; +use crate::auth::resolve_provider_auth_for_scope; use crate::models_endpoint::OpenAiModelsEndpoint; /// Optional provider-backed features that Codex may expose at runtime. @@ -175,6 +178,21 @@ pub trait ModelProvider: fmt::Debug + Send + Sync { }) } + /// Returns request credentials, optionally scoped to a Codex session task. + fn api_auth_for_scope( + &self, + scope: ProviderAuthScope, + ) -> ModelProviderFuture<'_, codex_protocol::error::Result> { + Box::pin(async move { + if !provider_uses_first_party_auth_path(self.info()) { + return self.api_auth().await.map(ResolvedProviderAuth::new); + } + let auth = self.auth().await; + resolve_provider_auth_for_scope(self.auth_manager(), auth.as_ref(), self.info(), scope) + .await + }) + } + /// Creates the model manager implementation appropriate for this provider. fn models_manager( &self, @@ -188,6 +206,14 @@ pub type ModelProviderFuture<'a, T> = Pin + Send + 'a /// Shared runtime model provider handle. pub type SharedModelProvider = Arc; +fn provider_uses_first_party_auth_path(provider: &ModelProviderInfo) -> bool { + provider.requires_openai_auth + && provider.env_key.is_none() + && provider.experimental_bearer_token.is_none() + && provider.auth.is_none() + && provider.aws.is_none() +} + /// Creates the default runtime model provider for configured provider metadata. pub fn create_model_provider( provider_info: ModelProviderInfo, @@ -310,14 +336,17 @@ impl ModelProvider for ConfiguredModelProvider { mod tests { use std::num::NonZeroU64; + use codex_login::auth::AgentIdentityAuthPolicy; use codex_login::auth::BedrockApiKeyAuth; use codex_model_provider_info::ModelProviderAwsAuthInfo; use codex_model_provider_info::WireApi; + use codex_model_provider_info::create_oss_provider_with_base_url; use codex_models_manager::manager::RefreshStrategy; use codex_protocol::account::PlanType; use codex_protocol::config_types::ModelProviderAuthInfo; use codex_protocol::openai_models::ModelInfo; use codex_protocol::openai_models::ModelsResponse; + use codex_protocol::protocol::SessionSource; use pretty_assertions::assert_eq; use serde_json::json; use wiremock::Mock; @@ -328,6 +357,7 @@ mod tests { use wiremock::matchers::path; use super::*; + use crate::auth::AgentIdentitySessionFallback; fn provider_info_with_command_auth() -> ModelProviderInfo { ModelProviderInfo { @@ -406,6 +436,25 @@ mod tests { }) } + #[tokio::test] + async fn scoped_auth_ignores_scope_for_non_openai_provider() { + let provider = create_model_provider( + create_oss_provider_with_base_url("http://localhost:11434/v1", WireApi::Responses), + /*auth_manager*/ None, + ); + + let auth = provider + .api_auth_for_scope(ProviderAuthScope { + agent_identity_policy: AgentIdentityAuthPolicy::JwtOnly, + session_source: SessionSource::Cli, + agent_identity_session_fallback: AgentIdentitySessionFallback::default(), + }) + .await + .expect("auth should resolve"); + + assert!(auth.auth.to_auth_headers().is_empty()); + } + #[test] fn configured_provider_uses_default_capabilities() { let provider = create_model_provider( diff --git a/codex-rs/otel/src/events/session_telemetry.rs b/codex-rs/otel/src/events/session_telemetry.rs index 7e2700236..8210c0e5b 100644 --- a/codex-rs/otel/src/events/session_telemetry.rs +++ b/codex-rs/otel/src/events/session_telemetry.rs @@ -32,6 +32,7 @@ use crate::metrics::runtime_metrics::RuntimeMetricsSummary; use crate::metrics::timer::Timer; use crate::provider::OtelProvider; use crate::sanitize_metric_tag_value; +use codex_api::AgentIdentityTelemetry; use codex_api::ApiError; use codex_api::ResponseEvent; use codex_protocol::ThreadId; @@ -516,6 +517,7 @@ impl SessionTelemetry { /*cf_ray*/ None, /*auth_error*/ None, /*auth_error_code*/ None, + /*agent_identity_telemetry*/ None, ); response @@ -538,6 +540,7 @@ impl SessionTelemetry { cf_ray: Option<&str>, auth_error: Option<&str>, auth_error_code: Option<&str>, + agent_identity_telemetry: Option<&AgentIdentityTelemetry>, ) { let success = status.is_some_and(|code| (200..=299).contains(&code)) && error.is_none(); let success_str = if success { "true" } else { "false" }; @@ -578,6 +581,8 @@ impl SessionTelemetry { auth.cf_ray = cf_ray, auth.error = auth_error, auth.error_code = auth_error_code, + auth.agent_id = agent_identity_telemetry.map(|metadata| metadata.agent_id.as_str()), + auth.task_id = agent_identity_telemetry.map(|metadata| metadata.task_id.as_str()), }, log: {}, trace: {}, @@ -601,6 +606,7 @@ impl SessionTelemetry { cf_ray: Option<&str>, auth_error: Option<&str>, auth_error_code: Option<&str>, + agent_identity_telemetry: Option<&AgentIdentityTelemetry>, ) { let success = error.is_none() && status @@ -632,6 +638,8 @@ impl SessionTelemetry { auth.cf_ray = cf_ray, auth.error = auth_error, auth.error_code = auth_error_code, + auth.agent_id = agent_identity_telemetry.map(|metadata| metadata.agent_id.as_str()), + auth.task_id = agent_identity_telemetry.map(|metadata| metadata.task_id.as_str()), }, log: {}, trace: {}, @@ -643,6 +651,7 @@ impl SessionTelemetry { duration: Duration, error: Option<&str>, connection_reused: bool, + agent_identity_telemetry: Option<&AgentIdentityTelemetry>, ) { let success_str = if error.is_none() { "true" } else { "false" }; self.counter( @@ -669,6 +678,8 @@ impl SessionTelemetry { auth.env_provider_key_present = self.metadata.auth_env.provider_env_key_present, auth.env_refresh_token_url_override_present = self.metadata.auth_env.refresh_token_url_override_present, auth.connection_reused = connection_reused, + auth.agent_id = agent_identity_telemetry.map(|metadata| metadata.agent_id.as_str()), + auth.task_id = agent_identity_telemetry.map(|metadata| metadata.task_id.as_str()), }, log: {}, trace: {}, diff --git a/codex-rs/otel/tests/suite/otel_export_routing_policy.rs b/codex-rs/otel/tests/suite/otel_export_routing_policy.rs index 7432aaaa0..17b6f764d 100644 --- a/codex-rs/otel/tests/suite/otel_export_routing_policy.rs +++ b/codex-rs/otel/tests/suite/otel_export_routing_policy.rs @@ -1,3 +1,4 @@ +use codex_api::AgentIdentityTelemetry; use codex_otel::AuthEnvTelemetryMetadata; use codex_otel::OtelProvider; use codex_otel::SessionTelemetry; @@ -511,6 +512,10 @@ fn otel_export_routing_policy_routes_api_request_auth_observability() { SandboxPolicy::DangerFullAccess, Vec::new(), ); + let agent_identity_telemetry = AgentIdentityTelemetry { + agent_id: "agent-runtime-otel".to_string(), + task_id: "task-run-otel".to_string(), + }; manager.record_api_request( /*attempt*/ 1, Some(401), @@ -526,6 +531,7 @@ fn otel_export_routing_policy_routes_api_request_auth_observability() { Some("ray-401"), Some("missing_authorization_header"), Some("token_expired"), + Some(&agent_identity_telemetry), ); }); @@ -599,6 +605,14 @@ fn otel_export_routing_policy_routes_api_request_auth_observability() { .map(String::as_str), Some("true") ); + assert_eq!( + request_log_attrs.get("auth.agent_id").map(String::as_str), + Some("agent-runtime-otel") + ); + assert_eq!( + request_log_attrs.get("auth.task_id").map(String::as_str), + Some("task-run-otel") + ); let spans = span_exporter.get_finished_spans().expect("span export"); let conversation_trace_event = @@ -641,6 +655,14 @@ fn otel_export_routing_policy_routes_api_request_auth_observability() { .map(String::as_str), Some("true") ); + assert_eq!( + request_trace_attrs.get("auth.agent_id").map(String::as_str), + Some("agent-runtime-otel") + ); + assert_eq!( + request_trace_attrs.get("auth.task_id").map(String::as_str), + Some("task-run-otel") + ); } #[test] @@ -685,6 +707,10 @@ fn otel_export_routing_policy_routes_websocket_connect_auth_observability() { .with_auth_env(auth_env_metadata()); let root_span = tracing::info_span!("root"); let _root_guard = root_span.enter(); + let agent_identity_telemetry = AgentIdentityTelemetry { + agent_id: "agent-runtime-ws".to_string(), + task_id: "task-run-ws".to_string(), + }; manager.record_websocket_connect( std::time::Duration::from_millis(17), Some(401), @@ -700,6 +726,7 @@ fn otel_export_routing_policy_routes_websocket_connect_auth_observability() { Some("ray-ws-401"), Some("missing_authorization_header"), Some("token_expired"), + Some(&agent_identity_telemetry), ); }); @@ -741,6 +768,14 @@ fn otel_export_routing_policy_routes_websocket_connect_auth_observability() { .map(String::as_str), Some("configured") ); + assert_eq!( + connect_log_attrs.get("auth.agent_id").map(String::as_str), + Some("agent-runtime-ws") + ); + assert_eq!( + connect_log_attrs.get("auth.task_id").map(String::as_str), + Some("task-run-ws") + ); let spans = span_exporter.get_finished_spans().expect("span export"); let connect_trace_event = @@ -758,6 +793,14 @@ fn otel_export_routing_policy_routes_websocket_connect_auth_observability() { .map(String::as_str), Some("true") ); + assert_eq!( + connect_trace_attrs.get("auth.agent_id").map(String::as_str), + Some("agent-runtime-ws") + ); + assert_eq!( + connect_trace_attrs.get("auth.task_id").map(String::as_str), + Some("task-run-ws") + ); } #[test] @@ -802,10 +845,15 @@ fn otel_export_routing_policy_routes_websocket_request_transport_observability() .with_auth_env(auth_env_metadata()); let root_span = tracing::info_span!("root"); let _root_guard = root_span.enter(); + let agent_identity_telemetry = AgentIdentityTelemetry { + agent_id: "agent-runtime-ws-request".to_string(), + task_id: "task-run-ws-request".to_string(), + }; manager.record_websocket_request( std::time::Duration::from_millis(23), Some("stream error"), /*connection_reused*/ true, + Some(&agent_identity_telemetry), ); }); @@ -831,6 +879,14 @@ fn otel_export_routing_policy_routes_websocket_request_transport_observability() .map(String::as_str), Some("true") ); + assert_eq!( + request_log_attrs.get("auth.agent_id").map(String::as_str), + Some("agent-runtime-ws-request") + ); + assert_eq!( + request_log_attrs.get("auth.task_id").map(String::as_str), + Some("task-run-ws-request") + ); let spans = span_exporter.get_finished_spans().expect("span export"); let request_trace_event = @@ -848,4 +904,12 @@ fn otel_export_routing_policy_routes_websocket_request_transport_observability() .map(String::as_str), Some("true") ); + assert_eq!( + request_trace_attrs.get("auth.agent_id").map(String::as_str), + Some("agent-runtime-ws-request") + ); + assert_eq!( + request_trace_attrs.get("auth.task_id").map(String::as_str), + Some("task-run-ws-request") + ); } diff --git a/codex-rs/otel/tests/suite/runtime_summary.rs b/codex-rs/otel/tests/suite/runtime_summary.rs index 7c869482c..b77e4557a 100644 --- a/codex-rs/otel/tests/suite/runtime_summary.rs +++ b/codex-rs/otel/tests/suite/runtime_summary.rs @@ -61,11 +61,13 @@ fn runtime_metrics_summary_collects_tool_api_and_streaming_metrics() -> Result<( /*cf_ray*/ None, /*auth_error*/ None, /*auth_error_code*/ None, + /*agent_identity_telemetry*/ None, ); manager.record_websocket_request( Duration::from_millis(400), /*error*/ None, /*connection_reused*/ false, + /*agent_identity_telemetry*/ None, ); let sse_response: std::result::Result< Option>>,