From c2a8a7702d789cab7c00d84fdedb6639185e3235 Mon Sep 17 00:00:00 2001 From: Edward Frazer Date: Tue, 21 Apr 2026 09:28:37 -0700 Subject: [PATCH 1/9] Revert "[codex] Use background agent task auth for backend calls (#18094)" This reverts commit 904c751a4068e5a4f7005ef165f5d3afa756f420. --- codex-rs/Cargo.lock | 1 - codex-rs/analytics/src/client.rs | 22 +- .../app-server/src/codex_message_processor.rs | 222 +++-- codex-rs/app-server/src/lib.rs | 24 +- codex-rs/app-server/src/transport/mod.rs | 3 +- .../src/transport/remote_control/enroll.rs | 18 +- .../src/transport/remote_control/mod.rs | 41 +- .../src/transport/remote_control/websocket.rs | 140 +-- codex-rs/backend-client/src/client.rs | 24 +- codex-rs/chatgpt/src/chatgpt_client.rs | 23 +- codex-rs/cloud-requirements/src/lib.rs | 33 +- codex-rs/cloud-tasks-client/src/http.rs | 10 - codex-rs/cloud-tasks/src/lib.rs | 52 +- codex-rs/cloud-tasks/src/util.rs | 48 +- codex-rs/codex-mcp/src/lib.rs | 4 - codex-rs/codex-mcp/src/mcp/mod.rs | 92 +- codex-rs/core-skills/src/remote.rs | 31 +- codex-rs/core/src/agent_identity.rs | 3 - codex-rs/core/src/config/mod.rs | 13 +- codex-rs/core/src/connectors.rs | 53 +- codex-rs/core/src/mcp.rs | 20 +- codex-rs/core/src/session/handlers.rs | 14 +- codex-rs/core/src/session/mcp.rs | 15 +- codex-rs/core/src/session/mod.rs | 2 +- codex-rs/core/src/session/session.rs | 14 +- codex-rs/core/src/session/tests.rs | 1 - codex-rs/login/Cargo.toml | 1 - codex-rs/login/src/agent_identity.rs | 834 ------------------ codex-rs/login/src/auth/agent_assertion.rs | 1 - codex-rs/login/src/auth/auth_tests.rs | 1 - codex-rs/login/src/auth/manager.rs | 165 +--- codex-rs/login/src/auth/storage.rs | 2 - codex-rs/login/src/auth/storage_tests.rs | 1 - codex-rs/login/src/lib.rs | 3 - codex-rs/models-manager/src/manager.rs | 23 +- 35 files changed, 248 insertions(+), 1706 deletions(-) delete mode 100644 codex-rs/login/src/agent_identity.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 25c12ac154c3..0f87b3ec2ebd 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2416,7 +2416,6 @@ dependencies = [ "codex-terminal-detection", "codex-utils-template", "core_test_support", - "crypto_box", "ed25519-dalek", "keyring", "once_cell", diff --git a/codex-rs/analytics/src/client.rs b/codex-rs/analytics/src/client.rs index 25ab40953f99..1a4b5defe9cd 100644 --- a/codex-rs/analytics/src/client.rs +++ b/codex-rs/analytics/src/client.rs @@ -297,7 +297,7 @@ impl AnalyticsEventsClient { } async fn send_track_events( - auth_manager: &Arc, + auth_manager: &AuthManager, base_url: &str, events: Vec, ) { @@ -310,11 +310,9 @@ async fn send_track_events( if !auth.is_chatgpt_auth() { return; } - let Some(authorization_header_value) = auth_manager - .chatgpt_authorization_header_for_auth(&auth) - .await - else { - return; + let access_token = match auth.get_token() { + Ok(token) => token, + Err(_) => return, }; let Some(account_id) = auth.get_account_id() else { return; @@ -324,17 +322,15 @@ async fn send_track_events( let url = format!("{base_url}/codex/analytics-events/events"); let payload = TrackEventsRequest { events }; - let mut request = create_client() + let response = create_client() .post(&url) .timeout(ANALYTICS_EVENTS_TIMEOUT) - .header("authorization", authorization_header_value) + .bearer_auth(&access_token) .header("chatgpt-account-id", &account_id) .header("Content-Type", "application/json") - .json(&payload); - if auth.is_fedramp_account() { - request = request.header("X-OpenAI-Fedramp", "true"); - } - let response = request.send().await; + .json(&payload) + .send() + .await; match response { Ok(response) if response.status().is_success() => {} diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 3b92d4199f58..038d23862bbb 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -281,10 +281,10 @@ use codex_login::run_login_server; use codex_mcp::McpRuntimeEnvironment; use codex_mcp::McpServerStatusSnapshot; use codex_mcp::McpSnapshotDetail; -use codex_mcp::collect_mcp_server_status_snapshot_with_detail_and_authorization_header; +use codex_mcp::collect_mcp_server_status_snapshot_with_detail; use codex_mcp::discover_supported_scopes; -use codex_mcp::effective_mcp_servers_with_authorization_header; use codex_mcp::read_mcp_resource as read_mcp_resource_without_thread; +use codex_mcp::effective_mcp_servers; use codex_mcp::resolve_oauth_scopes; use codex_models_manager::collaboration_mode_presets::CollaborationModesConfig; use codex_protocol::ThreadId; @@ -1970,28 +1970,12 @@ impl CodexMessageProcessor { }); } - let authorization_header_value = self - .auth_manager - .chatgpt_authorization_header_for_auth(&auth) - .await; - let mut client = BackendClient::new(self.config.chatgpt_base_url.clone()) - .map(|client| { - client.with_user_agent(codex_login::default_client::get_codex_user_agent()) - }) + let client = BackendClient::from_auth(self.config.chatgpt_base_url.clone(), &auth) .map_err(|err| JSONRPCErrorError { code: INTERNAL_ERROR_CODE, message: format!("failed to construct backend client: {err}"), data: None, })?; - if let Some(authorization_header_value) = authorization_header_value { - client = client.with_authorization_header_value(authorization_header_value); - } - if let Some(account_id) = auth.get_account_id() { - client = client.with_chatgpt_account_id(account_id); - } - if auth.is_fedramp_account() { - client = client.with_fedramp_routing_header(); - } let snapshots = client .get_rate_limits_many() @@ -5692,8 +5676,7 @@ impl CodexMessageProcessor { let mcp_config = config .to_mcp_config(self.thread_manager.plugins_manager().as_ref()) .await; - let auth_manager = Arc::clone(&self.auth_manager); - let auth = auth_manager.auth().await; + let auth = self.auth_manager.auth().await; let runtime_environment = match self.thread_manager.environment_manager().current().await { Ok(Some(environment)) => { // Status listing has no turn cwd. This fallback is used only @@ -5718,111 +5701,120 @@ impl CodexMessageProcessor { }; tokio::spawn(async move { - let detail = match params.detail.unwrap_or(McpServerStatusDetail::Full) { - McpServerStatusDetail::Full => McpSnapshotDetail::Full, - McpServerStatusDetail::ToolsAndAuthOnly => McpSnapshotDetail::ToolsAndAuthOnly, - }; - - let background_authorization_header_value = if let Some(auth) = auth.as_ref() { - auth_manager - .chatgpt_authorization_header_for_auth(auth) - .await - } else { - None - }; - let snapshot = collect_mcp_server_status_snapshot_with_detail_and_authorization_header( - &mcp_config, - auth.as_ref(), - request.request_id.to_string(), + Self::list_mcp_server_status_task( + outgoing, + request, + params, + config, + mcp_config, + auth, runtime_environment, - detail, - background_authorization_header_value.as_deref(), ) .await; + }); + } - let effective_servers = effective_mcp_servers_with_authorization_header( - &mcp_config, - auth.as_ref(), - background_authorization_header_value.as_deref(), - ); - let McpServerStatusSnapshot { - tools_by_server, - resources, - resource_templates, - auth_statuses, - } = snapshot; - - let mut server_names: Vec = config - .mcp_servers - .keys() - .cloned() - // Include built-in/plugin MCP servers that are present in the - // effective runtime config even when they are not user-declared in - // `config.mcp_servers`. - .chain(effective_servers.keys().cloned()) - .chain(auth_statuses.keys().cloned()) - .chain(resources.keys().cloned()) - .chain(resource_templates.keys().cloned()) - .collect(); - server_names.sort(); - server_names.dedup(); - - let total = server_names.len(); - let limit = params.limit.unwrap_or(total as u32).max(1) as usize; - let effective_limit = limit.min(total); - let start = match params.cursor { - Some(cursor) => match cursor.parse::() { - Ok(idx) => idx, - Err(_) => { - let error = JSONRPCErrorError { - code: INVALID_REQUEST_ERROR_CODE, - message: format!("invalid cursor: {cursor}"), - data: None, - }; - outgoing.send_error(request, error).await; - return; - } - }, - None => 0, - }; + async fn list_mcp_server_status_task( + outgoing: Arc, + request_id: ConnectionRequestId, + params: ListMcpServerStatusParams, + config: Config, + mcp_config: codex_mcp::McpConfig, + auth: Option, + runtime_environment: McpRuntimeEnvironment, + ) { + let detail = match params.detail.unwrap_or(McpServerStatusDetail::Full) { + McpServerStatusDetail::Full => McpSnapshotDetail::Full, + McpServerStatusDetail::ToolsAndAuthOnly => McpSnapshotDetail::ToolsAndAuthOnly, + }; - if start > total { - let error = JSONRPCErrorError { - code: INVALID_REQUEST_ERROR_CODE, - message: format!("cursor {start} exceeds total MCP servers {total}"), - data: None, - }; - outgoing.send_error(request, error).await; - return; - } + let snapshot = collect_mcp_server_status_snapshot_with_detail( + &mcp_config, + auth.as_ref(), + request_id.request_id.to_string(), + runtime_environment, + detail, + ) + .await; - let end = start.saturating_add(effective_limit).min(total); + let effective_servers = effective_mcp_servers(&mcp_config, auth.as_ref()); + let McpServerStatusSnapshot { + tools_by_server, + resources, + resource_templates, + auth_statuses, + } = snapshot; + + let mut server_names: Vec = config + .mcp_servers + .keys() + .cloned() + // Include built-in/plugin MCP servers that are present in the + // effective runtime config even when they are not user-declared in + // `config.mcp_servers`. + .chain(effective_servers.keys().cloned()) + .chain(auth_statuses.keys().cloned()) + .chain(resources.keys().cloned()) + .chain(resource_templates.keys().cloned()) + .collect(); + server_names.sort(); + server_names.dedup(); - let data: Vec = server_names[start..end] - .iter() - .map(|name| McpServerStatus { - name: name.clone(), - tools: tools_by_server.get(name).cloned().unwrap_or_default(), - resources: resources.get(name).cloned().unwrap_or_default(), - resource_templates: resource_templates.get(name).cloned().unwrap_or_default(), - auth_status: auth_statuses - .get(name) - .cloned() - .unwrap_or(CoreMcpAuthStatus::Unsupported) - .into(), - }) - .collect(); + let total = server_names.len(); + let limit = params.limit.unwrap_or(total as u32).max(1) as usize; + let effective_limit = limit.min(total); + let start = match params.cursor { + Some(cursor) => match cursor.parse::() { + Ok(idx) => idx, + Err(_) => { + let error = JSONRPCErrorError { + code: INVALID_REQUEST_ERROR_CODE, + message: format!("invalid cursor: {cursor}"), + data: None, + }; + outgoing.send_error(request_id, error).await; + return; + } + }, + None => 0, + }; - let next_cursor = if end < total { - Some(end.to_string()) - } else { - None + if start > total { + let error = JSONRPCErrorError { + code: INVALID_REQUEST_ERROR_CODE, + message: format!("cursor {start} exceeds total MCP servers {total}"), + data: None, }; + outgoing.send_error(request_id, error).await; + return; + } - let response = ListMcpServerStatusResponse { data, next_cursor }; + let end = start.saturating_add(effective_limit).min(total); - outgoing.send_response(request, response).await; - }); + let data: Vec = server_names[start..end] + .iter() + .map(|name| McpServerStatus { + name: name.clone(), + tools: tools_by_server.get(name).cloned().unwrap_or_default(), + resources: resources.get(name).cloned().unwrap_or_default(), + resource_templates: resource_templates.get(name).cloned().unwrap_or_default(), + auth_status: auth_statuses + .get(name) + .cloned() + .unwrap_or(CoreMcpAuthStatus::Unsupported) + .into(), + }) + .collect(); + + let next_cursor = if end < total { + Some(end.to_string()) + } else { + None + }; + + let response = ListMcpServerStatusResponse { data, next_cursor }; + + outgoing.send_response(request_id, response).await; } async fn read_mcp_resource( diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index d4573b267a83..2e4f04e16ea8 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -28,11 +28,10 @@ use crate::outgoing_message::QueuedOutgoingMessage; use crate::transport::CHANNEL_CAPACITY; use crate::transport::ConnectionState; use crate::transport::OutboundConnectionState; -use crate::transport::RemoteControlStartOptions; use crate::transport::TransportEvent; use crate::transport::auth::policy_from_settings; use crate::transport::route_outgoing_envelope; -use crate::transport::start_remote_control_with_options; +use crate::transport::start_remote_control; use crate::transport::start_stdio_connection; use crate::transport::start_websocket_acceptor; use codex_analytics::AppServerRpcTransport; @@ -581,17 +580,16 @@ pub async fn run_main_with_transport( )); } - let (remote_control_accept_handle, remote_control_handle) = - start_remote_control_with_options(RemoteControlStartOptions { - remote_control_url: config.chatgpt_base_url.clone(), - state_db: state_db.clone(), - auth_manager: auth_manager.clone(), - transport_event_tx: transport_event_tx.clone(), - shutdown_token: transport_shutdown_token.clone(), - app_server_client_name_rx, - initial_enabled: remote_control_enabled, - }) - .await?; + let (remote_control_accept_handle, remote_control_handle) = start_remote_control( + config.chatgpt_base_url.clone(), + state_db.clone(), + auth_manager.clone(), + transport_event_tx.clone(), + transport_shutdown_token.clone(), + app_server_client_name_rx, + remote_control_enabled, + ) + .await?; transport_accept_handles.push(remote_control_accept_handle); let outbound_handle = tokio::spawn(async move { diff --git a/codex-rs/app-server/src/transport/mod.rs b/codex-rs/app-server/src/transport/mod.rs index a07bd44188ac..712a92d2cb2a 100644 --- a/codex-rs/app-server/src/transport/mod.rs +++ b/codex-rs/app-server/src/transport/mod.rs @@ -34,8 +34,7 @@ mod stdio; mod websocket; pub(crate) use remote_control::RemoteControlHandle; -pub(crate) use remote_control::RemoteControlStartOptions; -pub(crate) use remote_control::start_remote_control_with_options; +pub(crate) use remote_control::start_remote_control; pub(crate) use stdio::start_stdio_connection; pub(crate) use websocket::start_websocket_acceptor; diff --git a/codex-rs/app-server/src/transport/remote_control/enroll.rs b/codex-rs/app-server/src/transport/remote_control/enroll.rs index 4d9929818a28..dbe18c8355db 100644 --- a/codex-rs/app-server/src/transport/remote_control/enroll.rs +++ b/codex-rs/app-server/src/transport/remote_control/enroll.rs @@ -17,7 +17,6 @@ const REMOTE_CONTROL_RESPONSE_BODY_MAX_BYTES: usize = 4096; const REQUEST_ID_HEADER: &str = "x-request-id"; const OAI_REQUEST_ID_HEADER: &str = "x-oai-request-id"; const CF_RAY_HEADER: &str = "cf-ray"; -const REMOTE_CONTROL_FEDRAMP_HEADER: &str = "X-OpenAI-Fedramp"; pub(super) const REMOTE_CONTROL_ACCOUNT_ID_HEADER: &str = "chatgpt-account-id"; #[derive(Debug, Clone, PartialEq, Eq)] @@ -30,9 +29,8 @@ pub(super) struct RemoteControlEnrollment { #[derive(Debug, Clone, PartialEq, Eq)] pub(super) struct RemoteControlConnectionAuth { - pub(super) authorization_header_value: String, + pub(super) bearer_token: String, pub(super) account_id: String, - pub(super) is_fedramp_account: bool, } pub(super) async fn load_persisted_remote_control_enrollment( @@ -201,15 +199,12 @@ pub(super) async fn enroll_remote_control_server( app_server_version: env!("CARGO_PKG_VERSION"), }; let client = build_reqwest_client(); - let mut http_request = client + let http_request = client .post(enroll_url) .timeout(REMOTE_CONTROL_ENROLL_TIMEOUT) - .header("authorization", &auth.authorization_header_value) - .header(REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id); - if auth.is_fedramp_account { - http_request = http_request.header(REMOTE_CONTROL_FEDRAMP_HEADER, "true"); - } - let http_request = http_request.json(&request); + .bearer_auth(&auth.bearer_token) + .header(REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id) + .json(&request); let response = http_request.send().await.map_err(|err| { io::Error::other(format!( @@ -450,9 +445,8 @@ mod tests { let err = enroll_remote_control_server( &remote_control_target, &RemoteControlConnectionAuth { - authorization_header_value: "Bearer Access Token".to_string(), + bearer_token: "Access Token".to_string(), account_id: "account_id".to_string(), - is_fedramp_account: false, }, ) .await diff --git a/codex-rs/app-server/src/transport/remote_control/mod.rs b/codex-rs/app-server/src/transport/remote_control/mod.rs index 8764928ea91a..c014c7a2c902 100644 --- a/codex-rs/app-server/src/transport/remote_control/mod.rs +++ b/codex-rs/app-server/src/transport/remote_control/mod.rs @@ -4,7 +4,6 @@ mod protocol; mod websocket; use crate::transport::remote_control::websocket::RemoteControlWebsocket; -use crate::transport::remote_control::websocket::RemoteControlWebsocketOptions; pub use self::protocol::ClientId; use self::protocol::ServerEvent; @@ -45,17 +44,6 @@ impl RemoteControlHandle { } } -pub(crate) struct RemoteControlStartOptions { - pub(crate) remote_control_url: String, - pub(crate) state_db: Option>, - pub(crate) auth_manager: Arc, - pub(crate) transport_event_tx: mpsc::Sender, - pub(crate) shutdown_token: CancellationToken, - pub(crate) app_server_client_name_rx: Option>, - pub(crate) initial_enabled: bool, -} - -#[cfg(test)] pub(crate) async fn start_remote_control( remote_control_url: String, state_db: Option>, @@ -65,38 +53,15 @@ pub(crate) async fn start_remote_control( app_server_client_name_rx: Option>, initial_enabled: bool, ) -> io::Result<(JoinHandle<()>, RemoteControlHandle)> { - start_remote_control_with_options(RemoteControlStartOptions { - remote_control_url, - state_db, - auth_manager, - transport_event_tx, - shutdown_token, - app_server_client_name_rx, - initial_enabled, - }) - .await -} - -pub(crate) async fn start_remote_control_with_options( - options: RemoteControlStartOptions, -) -> io::Result<(JoinHandle<()>, RemoteControlHandle)> { - let RemoteControlStartOptions { - remote_control_url, - state_db, - auth_manager, - transport_event_tx, - shutdown_token, - app_server_client_name_rx, - initial_enabled, - } = options; let remote_control_target = if initial_enabled { Some(normalize_remote_control_url(&remote_control_url)?) } else { None }; + let (enabled_tx, enabled_rx) = watch::channel(initial_enabled); let join_handle = tokio::spawn(async move { - RemoteControlWebsocket::from_options(RemoteControlWebsocketOptions { + RemoteControlWebsocket::new( remote_control_url, remote_control_target, state_db, @@ -104,7 +69,7 @@ pub(crate) async fn start_remote_control_with_options( transport_event_tx, shutdown_token, enabled_rx, - }) + ) .run(app_server_client_name_rx) .await; }); diff --git a/codex-rs/app-server/src/transport/remote_control/websocket.rs b/codex-rs/app-server/src/transport/remote_control/websocket.rs index dee177d30e07..523c93f3beed 100644 --- a/codex-rs/app-server/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server/src/transport/remote_control/websocket.rs @@ -49,7 +49,6 @@ use tracing::warn; pub(super) const REMOTE_CONTROL_PROTOCOL_VERSION: &str = "2"; pub(super) const REMOTE_CONTROL_ACCOUNT_ID_HEADER: &str = "chatgpt-account-id"; -const REMOTE_CONTROL_FEDRAMP_HEADER: &str = "X-OpenAI-Fedramp"; const REMOTE_CONTROL_SUBSCRIBE_CURSOR_HEADER: &str = "x-codex-subscribe-cursor"; const REMOTE_CONTROL_WEBSOCKET_PING_INTERVAL: std::time::Duration = std::time::Duration::from_secs(10); @@ -129,16 +128,6 @@ pub(crate) struct RemoteControlWebsocket { enabled_rx: watch::Receiver, } -pub(crate) struct RemoteControlWebsocketOptions { - pub(crate) remote_control_url: String, - pub(crate) remote_control_target: Option, - pub(crate) state_db: Option>, - pub(crate) auth_manager: Arc, - pub(crate) transport_event_tx: mpsc::Sender, - pub(crate) shutdown_token: CancellationToken, - pub(crate) enabled_rx: watch::Receiver, -} - enum ConnectOutcome { Connected(Box>>), Disabled, @@ -146,7 +135,6 @@ enum ConnectOutcome { } impl RemoteControlWebsocket { - #[cfg(test)] pub(crate) fn new( remote_control_url: String, remote_control_target: Option, @@ -156,27 +144,6 @@ impl RemoteControlWebsocket { shutdown_token: CancellationToken, enabled_rx: watch::Receiver, ) -> Self { - Self::from_options(RemoteControlWebsocketOptions { - remote_control_url, - remote_control_target, - state_db, - auth_manager, - transport_event_tx, - shutdown_token, - enabled_rx, - }) - } - - pub(crate) fn from_options(options: RemoteControlWebsocketOptions) -> Self { - let RemoteControlWebsocketOptions { - remote_control_url, - remote_control_target, - state_db, - auth_manager, - transport_event_tx, - shutdown_token, - enabled_rx, - } = options; let shutdown_token = shutdown_token.child_token(); let (server_event_tx, server_event_rx) = mpsc::channel(super::CHANNEL_CAPACITY); let client_tracker = @@ -308,16 +275,14 @@ impl RemoteControlWebsocket { } return ConnectOutcome::Disabled; } - connect_result = connect_remote_control_websocket_with_options( - ConnectRemoteControlWebsocketOptions { - remote_control_target: &remote_control_target, - state_db: self.state_db.as_deref(), - auth_manager: &self.auth_manager, - auth_recovery: &mut self.auth_recovery, - enrollment: &mut self.enrollment, - subscribe_cursor: subscribe_cursor.as_deref(), - app_server_client_name, - }, + connect_result = connect_remote_control_websocket( + &remote_control_target, + self.state_db.as_deref(), + &self.auth_manager, + &mut self.auth_recovery, + &mut self.enrollment, + subscribe_cursor.as_deref(), + app_server_client_name, ) => connect_result, }; @@ -715,11 +680,12 @@ fn build_remote_control_websocket_request( "x-codex-protocol-version", REMOTE_CONTROL_PROTOCOL_VERSION, )?; - set_remote_control_header(headers, "authorization", &auth.authorization_header_value)?; + set_remote_control_header( + headers, + "authorization", + &format!("Bearer {}", auth.bearer_token), + )?; set_remote_control_header(headers, REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id)?; - if auth.is_fedramp_account { - set_remote_control_header(headers, REMOTE_CONTROL_FEDRAMP_HEADER, "true")?; - } if let Some(subscribe_cursor) = subscribe_cursor { set_remote_control_header( headers, @@ -764,19 +730,8 @@ pub(crate) async fn load_remote_control_auth( )); } - let authorization_header_value = auth_manager - .chatgpt_authorization_header_for_auth(&auth) - .await - .ok_or_else(|| { - io::Error::new( - ErrorKind::PermissionDenied, - "remote control requires ChatGPT authentication", - ) - })?; - Ok(RemoteControlConnectionAuth { - authorization_header_value, - is_fedramp_account: auth.is_fedramp_account(), + bearer_token: auth.get_token().map_err(io::Error::other)?, account_id: auth.get_account_id().ok_or_else(|| { io::Error::new( ErrorKind::WouldBlock, @@ -786,7 +741,6 @@ pub(crate) async fn load_remote_control_auth( }) } -#[cfg(test)] pub(super) async fn connect_remote_control_websocket( remote_control_target: &RemoteControlTarget, state_db: Option<&StateRuntime>, @@ -799,44 +753,6 @@ pub(super) async fn connect_remote_control_websocket( WebSocketStream>, tungstenite::http::Response<()>, )> { - connect_remote_control_websocket_with_options(ConnectRemoteControlWebsocketOptions { - remote_control_target, - state_db, - auth_manager, - auth_recovery, - enrollment, - subscribe_cursor, - app_server_client_name, - }) - .await -} - -struct ConnectRemoteControlWebsocketOptions<'a> { - remote_control_target: &'a RemoteControlTarget, - state_db: Option<&'a StateRuntime>, - auth_manager: &'a Arc, - auth_recovery: &'a mut UnauthorizedRecovery, - enrollment: &'a mut Option, - subscribe_cursor: Option<&'a str>, - app_server_client_name: Option<&'a str>, -} - -async fn connect_remote_control_websocket_with_options( - options: ConnectRemoteControlWebsocketOptions<'_>, -) -> io::Result<( - WebSocketStream>, - tungstenite::http::Response<()>, -)> { - let ConnectRemoteControlWebsocketOptions { - remote_control_target, - state_db, - auth_manager, - auth_recovery, - enrollment, - subscribe_cursor, - app_server_client_name, - } = options; - ensure_rustls_crypto_provider(); let auth = load_remote_control_auth(auth_manager).await?; @@ -1099,34 +1015,6 @@ mod tests { } } - #[test] - fn build_remote_control_websocket_request_includes_fedramp_header() { - let request = build_remote_control_websocket_request( - "ws://127.0.0.1/backend-api/wham/remote/control/server", - &RemoteControlEnrollment { - account_id: "account_id".to_string(), - environment_id: "env_test".to_string(), - server_id: "srv_e_test".to_string(), - server_name: "test-server".to_string(), - }, - &RemoteControlConnectionAuth { - authorization_header_value: "AgentAssertion assertion".to_string(), - account_id: "account_id".to_string(), - is_fedramp_account: true, - }, - /*subscribe_cursor*/ None, - ) - .expect("request should build"); - - assert_eq!( - request - .headers() - .get(REMOTE_CONTROL_FEDRAMP_HEADER) - .and_then(|value| value.to_str().ok()), - Some("true") - ); - } - #[tokio::test] async fn connect_remote_control_websocket_includes_http_error_details() { let listener = TcpListener::bind("127.0.0.1:0") diff --git a/codex-rs/backend-client/src/client.rs b/codex-rs/backend-client/src/client.rs index e6ea0253b896..8f84ef28f4b0 100644 --- a/codex-rs/backend-client/src/client.rs +++ b/codex-rs/backend-client/src/client.rs @@ -116,7 +116,7 @@ impl PathStyle { pub struct Client { base_url: String, http: reqwest::Client, - authorization_header_value: Option, + bearer_token: Option, user_agent: Option, chatgpt_account_id: Option, chatgpt_account_is_fedramp: bool, @@ -142,7 +142,7 @@ impl Client { Ok(Self { base_url, http, - authorization_header_value: None, + bearer_token: None, user_agent: None, chatgpt_account_id: None, chatgpt_account_is_fedramp: false, @@ -165,12 +165,7 @@ impl Client { } pub fn with_bearer_token(mut self, token: impl Into) -> Self { - self.authorization_header_value = Some(format!("Bearer {}", token.into())); - self - } - - pub fn with_authorization_header_value(mut self, value: impl Into) -> Self { - self.authorization_header_value = Some(value.into()); + self.bearer_token = Some(token.into()); self } @@ -203,10 +198,11 @@ impl Client { } else { h.insert(USER_AGENT, HeaderValue::from_static("codex-cli")); } - if let Some(value) = &self.authorization_header_value - && let Ok(hv) = HeaderValue::from_str(value) - { - h.insert(AUTHORIZATION, hv); + if let Some(token) = &self.bearer_token { + let value = format!("Bearer {token}"); + if let Ok(hv) = HeaderValue::from_str(&value) { + h.insert(AUTHORIZATION, hv); + } } if let Some(acc) = &self.chatgpt_account_id && let Ok(name) = HeaderName::from_bytes(b"ChatGPT-Account-Id") @@ -820,7 +816,7 @@ mod tests { let codex_client = Client { base_url: "https://example.test".to_string(), http: reqwest::Client::new(), - authorization_header_value: None, + bearer_token: None, user_agent: None, chatgpt_account_id: None, chatgpt_account_is_fedramp: false, @@ -834,7 +830,7 @@ mod tests { let chatgpt_client = Client { base_url: "https://chatgpt.com/backend-api".to_string(), http: reqwest::Client::new(), - authorization_header_value: None, + bearer_token: None, user_agent: None, chatgpt_account_id: None, chatgpt_account_is_fedramp: false, diff --git a/codex-rs/chatgpt/src/chatgpt_client.rs b/codex-rs/chatgpt/src/chatgpt_client.rs index 6bdb16347e14..fa3a63dadbb4 100644 --- a/codex-rs/chatgpt/src/chatgpt_client.rs +++ b/codex-rs/chatgpt/src/chatgpt_client.rs @@ -1,5 +1,4 @@ use codex_core::config::Config; -use codex_login::AuthManager; use codex_login::default_client::create_client; use crate::chatgpt_token::get_chatgpt_token_data; @@ -32,32 +31,16 @@ pub(crate) async fn chatgpt_get_request_with_timeout( let token = get_chatgpt_token_data().ok_or_else(|| anyhow::anyhow!("ChatGPT token not available"))?; - let auth_manager = - AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false); - let auth = auth_manager.auth().await; - let is_fedramp_account = auth - .as_ref() - .is_some_and(codex_login::CodexAuth::is_fedramp_account); - let authorization_header_value = match auth.as_ref() { - Some(auth) if auth.is_chatgpt_auth() => auth_manager - .chatgpt_authorization_header_for_auth(auth) - .await - .unwrap_or_else(|| format!("Bearer {}", token.access_token)), - _ => format!("Bearer {}", token.access_token), - }; let account_id = token.account_id.ok_or_else(|| { anyhow::anyhow!("ChatGPT account ID not available, please re-run `codex login`") - })?; + }); let mut request = client .get(&url) - .header("authorization", authorization_header_value) - .header("chatgpt-account-id", account_id) + .bearer_auth(&token.access_token) + .header("chatgpt-account-id", account_id?) .header("Content-Type", "application/json"); - if is_fedramp_account { - request = request.header("X-OpenAI-Fedramp", "true"); - } if let Some(timeout) = timeout { request = request.timeout(timeout); diff --git a/codex-rs/cloud-requirements/src/lib.rs b/codex-rs/cloud-requirements/src/lib.rs index 23be6f410da2..914958d2c72c 100644 --- a/codex-rs/cloud-requirements/src/lib.rs +++ b/codex-rs/cloud-requirements/src/lib.rs @@ -195,15 +195,11 @@ trait RequirementsFetcher: Send + Sync { struct BackendRequirementsFetcher { base_url: String, - auth_manager: Arc, } impl BackendRequirementsFetcher { - fn new(auth_manager: Arc, base_url: String) -> Self { - Self { - base_url, - auth_manager, - } + fn new(base_url: String) -> Self { + Self { base_url } } } @@ -213,14 +209,7 @@ impl RequirementsFetcher for BackendRequirementsFetcher { &self, auth: &CodexAuth, ) -> Result, FetchAttemptError> { - let authorization_header_value = self - .auth_manager - .chatgpt_authorization_header_for_auth(auth) - .await; - let mut client = BackendClient::new(self.base_url.clone()) - .map(|client| { - client.with_user_agent(codex_login::default_client::get_codex_user_agent()) - }) + let client = BackendClient::from_auth(self.base_url.clone(), auth) .inspect_err(|err| { tracing::warn!( error = %err, @@ -228,15 +217,6 @@ impl RequirementsFetcher for BackendRequirementsFetcher { ); }) .map_err(|_| FetchAttemptError::Retryable(RetryableFailureKind::BackendClientInit))?; - if let Some(authorization_header_value) = authorization_header_value { - client = client.with_authorization_header_value(authorization_header_value); - } - if let Some(account_id) = auth.get_account_id() { - client = client.with_chatgpt_account_id(account_id); - } - if auth.is_fedramp_account() { - client = client.with_fedramp_routing_header(); - } let response = client .get_config_requirements_file() @@ -713,11 +693,8 @@ pub fn cloud_requirements_loader( codex_home: PathBuf, ) -> CloudRequirementsLoader { let service = CloudRequirementsService::new( - auth_manager.clone(), - Arc::new(BackendRequirementsFetcher::new( - auth_manager, - chatgpt_base_url, - )), + auth_manager, + Arc::new(BackendRequirementsFetcher::new(chatgpt_base_url)), codex_home, CLOUD_REQUIREMENTS_TIMEOUT, ); diff --git a/codex-rs/cloud-tasks-client/src/http.rs b/codex-rs/cloud-tasks-client/src/http.rs index 3ada7023655e..4ea098022737 100644 --- a/codex-rs/cloud-tasks-client/src/http.rs +++ b/codex-rs/cloud-tasks-client/src/http.rs @@ -37,11 +37,6 @@ impl HttpClient { self } - pub fn with_authorization_header_value(mut self, value: impl Into) -> Self { - self.backend = self.backend.clone().with_authorization_header_value(value); - self - } - pub fn with_user_agent(mut self, ua: impl Into) -> Self { self.backend = self.backend.clone().with_user_agent(ua); self @@ -52,11 +47,6 @@ impl HttpClient { self } - pub fn with_fedramp_routing_header(mut self) -> Self { - self.backend = self.backend.clone().with_fedramp_routing_header(); - self - } - fn tasks_api(&self) -> api::Tasks<'_> { api::Tasks::new(self) } diff --git a/codex-rs/cloud-tasks/src/lib.rs b/codex-rs/cloud-tasks/src/lib.rs index 155587003cd6..7006d52b921d 100644 --- a/codex-rs/cloud-tasks/src/lib.rs +++ b/codex-rs/cloud-tasks/src/lib.rs @@ -68,45 +68,43 @@ async fn init_backend(user_agent_suffix: &str) -> anyhow::Result }; append_error_log(format!("startup: base_url={base_url} path_style={style}")); - let Some(auth_manager) = util::load_auth_manager(Some(base_url.clone())).await else { - eprintln!( - "Not signed in. Please run 'codex login' to sign in with ChatGPT, then re-run 'codex cloud'." - ); - std::process::exit(1); + let auth_manager = util::load_auth_manager().await; + let auth = match auth_manager.as_ref() { + Some(manager) => manager.auth().await, + None => None, }; - let Some(auth) = auth_manager.auth().await else { - eprintln!( - "Not signed in. Please run 'codex login' to sign in with ChatGPT, then re-run 'codex cloud'." - ); - std::process::exit(1); + let auth = match auth { + Some(auth) => auth, + None => { + eprintln!( + "Not signed in. Please run 'codex login' to sign in with ChatGPT, then re-run 'codex cloud'." + ); + std::process::exit(1); + } }; if let Some(acc) = auth.get_account_id() { append_error_log(format!("auth: mode=ChatGPT account_id={acc}")); } - let authorization_header_value = auth_manager - .chatgpt_authorization_header_for_auth(&auth) - .await; - let Some(authorization_header_value) = authorization_header_value else { - eprintln!( - "Not signed in. Please run 'codex login' to sign in with ChatGPT, then re-run 'codex cloud'." - ); - std::process::exit(1); + let token = match auth.get_token() { + Ok(t) if !t.is_empty() => t, + _ => { + eprintln!( + "Not signed in. Please run 'codex login' to sign in with ChatGPT, then re-run 'codex cloud'." + ); + std::process::exit(1); + } }; - http = http.with_authorization_header_value(authorization_header_value); - if let Some(acc) = auth.get_account_id().or_else(|| { - auth.get_token() - .ok() - .and_then(|token| util::extract_chatgpt_account_id(&token)) - }) { + http = http.with_bearer_token(token.clone()); + if let Some(acc) = auth + .get_account_id() + .or_else(|| util::extract_chatgpt_account_id(&token)) + { append_error_log(format!("auth: set ChatGPT-Account-Id header: {acc}")); http = http.with_chatgpt_account_id(acc); } - if auth.is_fedramp_account() { - http = http.with_fedramp_routing_header(); - } Ok(BackendContext { backend: Arc::new(http), diff --git a/codex-rs/cloud-tasks/src/util.rs b/codex-rs/cloud-tasks/src/util.rs index 693ff7839c72..cbaed17beaf3 100644 --- a/codex-rs/cloud-tasks/src/util.rs +++ b/codex-rs/cloud-tasks/src/util.rs @@ -3,7 +3,6 @@ use chrono::DateTime; use chrono::Local; use chrono::Utc; use reqwest::header::HeaderMap; -use std::sync::Arc; use codex_core::config::Config; use codex_login::AuthManager; @@ -60,18 +59,18 @@ pub fn extract_chatgpt_account_id(token: &str) -> Option { .map(str::to_string) } -pub async fn load_auth_manager(chatgpt_base_url: Option) -> Option> { +pub async fn load_auth_manager() -> Option { // TODO: pass in cli overrides once cloud tasks properly support them. let config = Config::load_with_cli_overrides(Vec::new()).await.ok()?; - let auth_manager = - AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false); - if let Some(chatgpt_base_url) = chatgpt_base_url { - auth_manager.set_chatgpt_backend_base_url(Some(chatgpt_base_url)); - } - Some(auth_manager) + Some(AuthManager::new( + config.codex_home.to_path_buf(), + /*enable_codex_api_key_env*/ false, + config.cli_auth_credentials_store_mode, + )) } -/// Build headers for ChatGPT-backed requests. +/// Build headers for ChatGPT-backed requests: `User-Agent`, optional `Authorization`, +/// and optional `ChatGPT-Account-Id`. pub async fn build_chatgpt_headers() -> HeaderMap { use reqwest::header::AUTHORIZATION; use reqwest::header::HeaderName; @@ -85,34 +84,23 @@ pub async fn build_chatgpt_headers() -> HeaderMap { USER_AGENT, HeaderValue::from_str(&ua).unwrap_or(HeaderValue::from_static("codex-cli")), ); - let base_url = normalize_base_url( - &std::env::var("CODEX_CLOUD_TASKS_BASE_URL") - .unwrap_or_else(|_| "https://chatgpt.com/backend-api".to_string()), - ); - if let Some(auth_manager) = load_auth_manager(Some(base_url)).await - && let Some(auth) = auth_manager.auth().await + if let Some(am) = load_auth_manager().await + && let Some(auth) = am.auth().await + && let Ok(tok) = auth.get_token() + && !tok.is_empty() { - if let Some(authorization_header_value) = auth_manager - .chatgpt_authorization_header_for_auth(&auth) - .await - && let Ok(hv) = HeaderValue::from_str(&authorization_header_value) - { + let v = format!("Bearer {tok}"); + if let Ok(hv) = HeaderValue::from_str(&v) { headers.insert(AUTHORIZATION, hv); } - if let Some(acc) = auth.get_account_id().or_else(|| { - auth.get_token() - .ok() - .and_then(|token| extract_chatgpt_account_id(&token)) - }) && let Ok(name) = HeaderName::from_bytes(b"ChatGPT-Account-Id") + if let Some(acc) = auth + .get_account_id() + .or_else(|| extract_chatgpt_account_id(&tok)) + && let Ok(name) = HeaderName::from_bytes(b"ChatGPT-Account-Id") && let Ok(hv) = HeaderValue::from_str(&acc) { headers.insert(name, hv); } - if auth.is_fedramp_account() - && let Ok(name) = HeaderName::from_bytes(b"X-OpenAI-Fedramp") - { - headers.insert(name, HeaderValue::from_static("true")); - } } headers } diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 760fc64cfd0f..8b77086e5e17 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -16,18 +16,15 @@ pub use mcp::ToolPluginProvenance; pub use mcp::canonical_mcp_server_key; pub use mcp::collect_mcp_server_status_snapshot; pub use mcp::collect_mcp_server_status_snapshot_with_detail; -pub use mcp::collect_mcp_server_status_snapshot_with_detail_and_authorization_header; pub use mcp::collect_mcp_snapshot; pub use mcp::collect_mcp_snapshot_from_manager; pub use mcp::collect_mcp_snapshot_from_manager_with_detail; pub use mcp::collect_mcp_snapshot_with_detail; -pub use mcp::collect_mcp_snapshot_with_detail_and_authorization_header; pub use mcp::collect_missing_mcp_dependencies; pub use mcp::compute_auth_statuses; pub use mcp::configured_mcp_servers; pub use mcp::discover_supported_scopes; pub use mcp::effective_mcp_servers; -pub use mcp::effective_mcp_servers_with_authorization_header; pub use mcp::group_tools_by_server; pub use mcp::mcp_permission_prompt_is_auto_approved; pub use mcp::oauth_login_support; @@ -38,7 +35,6 @@ pub use mcp::should_retry_without_scopes; pub use mcp::split_qualified_tool_name; pub use mcp::tool_plugin_provenance; pub use mcp::with_codex_apps_mcp; -pub use mcp::with_codex_apps_mcp_with_authorization_header; pub use mcp_connection_manager::CodexAppsToolsCacheKey; pub use mcp_connection_manager::DEFAULT_STARTUP_TIMEOUT; pub use mcp_connection_manager::MCP_SANDBOX_STATE_META_CAPABILITY; diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 3118c7c171aa..97053cbe53e8 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -215,25 +215,14 @@ fn codex_apps_mcp_bearer_token(auth: Option<&CodexAuth>) -> Option { } } -fn codex_apps_mcp_http_headers( - auth: Option<&CodexAuth>, - authorization_header_value: Option<&str>, -) -> Option> { +fn codex_apps_mcp_http_headers(auth: Option<&CodexAuth>) -> Option> { let mut headers = HashMap::new(); - if let Some(authorization_header_value) = authorization_header_value { - headers.insert( - "Authorization".to_string(), - authorization_header_value.to_string(), - ); - } else if let Some(token) = codex_apps_mcp_bearer_token(auth) { + if let Some(token) = codex_apps_mcp_bearer_token(auth) { headers.insert("Authorization".to_string(), format!("Bearer {token}")); } if let Some(account_id) = auth.and_then(CodexAuth::get_account_id) { headers.insert("ChatGPT-Account-ID".to_string(), account_id); } - if auth.is_some_and(CodexAuth::is_fedramp_account) { - headers.insert("X-OpenAI-Fedramp".to_string(), "true".to_string()); - } if headers.is_empty() { None } else { @@ -267,16 +256,12 @@ pub(crate) fn codex_apps_mcp_url(config: &McpConfig) -> String { codex_apps_mcp_url_for_base_url(&config.chatgpt_base_url) } -fn codex_apps_mcp_server_config( - config: &McpConfig, - auth: Option<&CodexAuth>, - authorization_header_value: Option<&str>, -) -> McpServerConfig { +fn codex_apps_mcp_server_config(config: &McpConfig, auth: Option<&CodexAuth>) -> McpServerConfig { let bearer_token_env_var = codex_apps_mcp_bearer_token_env_var(); let http_headers = if bearer_token_env_var.is_some() { None } else { - codex_apps_mcp_http_headers(auth, authorization_header_value) + codex_apps_mcp_http_headers(auth) }; let url = codex_apps_mcp_url(config); @@ -304,25 +289,14 @@ fn codex_apps_mcp_server_config( } pub fn with_codex_apps_mcp( - servers: HashMap, - auth: Option<&CodexAuth>, - config: &McpConfig, -) -> HashMap { - with_codex_apps_mcp_with_authorization_header( - servers, auth, config, /*authorization_header_value*/ None, - ) -} - -pub fn with_codex_apps_mcp_with_authorization_header( mut servers: HashMap, auth: Option<&CodexAuth>, config: &McpConfig, - authorization_header_value: Option<&str>, ) -> HashMap { if config.apps_enabled && auth.is_some_and(CodexAuth::is_chatgpt_auth) { servers.insert( CODEX_APPS_MCP_SERVER_NAME.to_string(), - codex_apps_mcp_server_config(config, auth, authorization_header_value), + codex_apps_mcp_server_config(config, auth), ); } else { servers.remove(CODEX_APPS_MCP_SERVER_NAME); @@ -337,19 +311,9 @@ pub fn configured_mcp_servers(config: &McpConfig) -> HashMap, -) -> HashMap { - effective_mcp_servers_with_authorization_header( - config, auth, /*authorization_header_value*/ None, - ) -} - -pub fn effective_mcp_servers_with_authorization_header( - config: &McpConfig, - auth: Option<&CodexAuth>, - authorization_header_value: Option<&str>, ) -> HashMap { let servers = configured_mcp_servers(config); - with_codex_apps_mcp_with_authorization_header(servers, auth, config, authorization_header_value) + with_codex_apps_mcp(servers, auth, config) } pub fn tool_plugin_provenance(config: &McpConfig) -> ToolPluginProvenance { @@ -420,27 +384,7 @@ pub async fn collect_mcp_snapshot_with_detail( runtime_environment: McpRuntimeEnvironment, detail: McpSnapshotDetail, ) -> McpListToolsResponseEvent { - collect_mcp_snapshot_with_detail_and_authorization_header( - config, - auth, - submit_id, - runtime_environment, - detail, - /*authorization_header_value*/ None, - ) - .await -} - -pub async fn collect_mcp_snapshot_with_detail_and_authorization_header( - config: &McpConfig, - auth: Option<&CodexAuth>, - submit_id: String, - runtime_environment: McpRuntimeEnvironment, - detail: McpSnapshotDetail, - authorization_header_value: Option<&str>, -) -> McpListToolsResponseEvent { - let mcp_servers = - effective_mcp_servers_with_authorization_header(config, auth, authorization_header_value); + let mcp_servers = effective_mcp_servers(config, auth); let tool_plugin_provenance = tool_plugin_provenance(config); if mcp_servers.is_empty() { return McpListToolsResponseEvent { @@ -515,27 +459,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( runtime_environment: McpRuntimeEnvironment, detail: McpSnapshotDetail, ) -> McpServerStatusSnapshot { - collect_mcp_server_status_snapshot_with_detail_and_authorization_header( - config, - auth, - submit_id, - runtime_environment, - detail, - /*authorization_header_value*/ None, - ) - .await -} - -pub async fn collect_mcp_server_status_snapshot_with_detail_and_authorization_header( - config: &McpConfig, - auth: Option<&CodexAuth>, - submit_id: String, - runtime_environment: McpRuntimeEnvironment, - detail: McpSnapshotDetail, - authorization_header_value: Option<&str>, -) -> McpServerStatusSnapshot { - let mcp_servers = - effective_mcp_servers_with_authorization_header(config, auth, authorization_header_value); + let mcp_servers = effective_mcp_servers(config, auth); let tool_plugin_provenance = tool_plugin_provenance(config); if mcp_servers.is_empty() { return McpServerStatusSnapshot { diff --git a/codex-rs/core-skills/src/remote.rs b/codex-rs/core-skills/src/remote.rs index c85c78ff8d50..2dc620b864d8 100644 --- a/codex-rs/core-skills/src/remote.rs +++ b/codex-rs/core-skills/src/remote.rs @@ -6,9 +6,7 @@ use std::path::Path; use std::path::PathBuf; use std::time::Duration; -use codex_login::BackgroundAgentTaskAuthMode; use codex_login::CodexAuth; -use codex_login::cached_background_agent_task_authorization_header_value; use codex_login::default_client::build_reqwest_client; const REMOTE_SKILLS_API_TIMEOUT: Duration = Duration::from_secs(30); @@ -114,15 +112,13 @@ pub async fn list_remote_skills( .get(&url) .timeout(REMOTE_SKILLS_API_TIMEOUT) .query(&query_params); - let authorization_header_value = authorization_header_value_for_auth(auth) + let token = auth + .get_token() .context("Failed to read auth token for remote skills")?; - request = request.header("authorization", authorization_header_value); + request = request.bearer_auth(token); if let Some(account_id) = auth.get_account_id() { request = request.header("chatgpt-account-id", account_id); } - if auth.is_fedramp_account() { - request = request.header("X-OpenAI-Fedramp", "true"); - } let response = request .send() .await @@ -161,15 +157,13 @@ pub async fn export_remote_skill( let url = format!("{base_url}/hazelnuts/{skill_id}/export"); let mut request = client.get(&url).timeout(REMOTE_SKILLS_API_TIMEOUT); - let authorization_header_value = authorization_header_value_for_auth(auth) + let token = auth + .get_token() .context("Failed to read auth token for remote skills")?; - request = request.header("authorization", authorization_header_value); + request = request.bearer_auth(token); if let Some(account_id) = auth.get_account_id() { request = request.header("chatgpt-account-id", account_id); } - if auth.is_fedramp_account() { - request = request.header("X-OpenAI-Fedramp", "true"); - } let response = request .send() @@ -207,19 +201,6 @@ pub async fn export_remote_skill( }) } -fn authorization_header_value_for_auth(auth: &CodexAuth) -> std::io::Result { - if let Ok(Some(authorization_header_value)) = - cached_background_agent_task_authorization_header_value( - auth, - BackgroundAgentTaskAuthMode::Disabled, - ) - { - Ok(authorization_header_value) - } else { - auth.get_token().map(|token| format!("Bearer {token}")) - } -} - fn safe_join(base: &Path, name: &str) -> Result { let path = Path::new(name); for component in path.components() { diff --git a/codex-rs/core/src/agent_identity.rs b/codex-rs/core/src/agent_identity.rs index 3a4c6787f21a..5bd15de76c9c 100644 --- a/codex-rs/core/src/agent_identity.rs +++ b/codex-rs/core/src/agent_identity.rs @@ -412,7 +412,6 @@ impl StoredAgentIdentity { agent_runtime_id: self.agent_runtime_id.clone(), agent_private_key: self.private_key_pkcs8_base64.clone(), registered_at: self.registered_at.clone(), - background_task_id: None, } } @@ -673,7 +672,6 @@ mod tests { agent_runtime_id: "agent_invalid".to_string(), agent_private_key: "not-valid-base64".to_string(), registered_at: "2026-01-01T00:00:00Z".to_string(), - background_task_id: None, }) .expect("seed invalid identity"); @@ -713,7 +711,6 @@ mod tests { agent_runtime_id: "agent_old".to_string(), agent_private_key: stale_key.private_key_pkcs8_base64, registered_at: "2026-01-01T00:00:00Z".to_string(), - background_task_id: None, }) .expect("seed stale identity"); diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 7ae710f83cbc..995c2008c54a 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -51,7 +51,7 @@ use codex_config::types::UriBasedFileOpener; use codex_config::types::WindowsSandboxModeToml; use codex_exec_server::ExecutorFileSystem; use codex_exec_server::LOCAL_FS; -pub use codex_features::Feature; +use codex_features::Feature; use codex_features::FeatureConfigSource; use codex_features::FeatureOverrides; use codex_features::FeatureToml; @@ -60,7 +60,6 @@ use codex_features::FeaturesToml; use codex_features::MultiAgentV2ConfigToml; use codex_git_utils::resolve_root_git_project_for_trust; use codex_login::AuthManagerConfig; -use codex_login::BackgroundAgentTaskAuthMode; use codex_mcp::McpConfig; use codex_model_provider_info::LEGACY_OLLAMA_CHAT_PROVIDER_ID; use codex_model_provider_info::ModelProviderInfo; @@ -635,16 +634,6 @@ impl AuthManagerConfig for Config { fn forced_chatgpt_workspace_id(&self) -> Option { self.forced_chatgpt_workspace_id.clone() } - - fn chatgpt_base_url(&self) -> Option { - Some(self.chatgpt_base_url.clone()) - } - - fn background_agent_task_auth_mode(&self) -> BackgroundAgentTaskAuthMode { - BackgroundAgentTaskAuthMode::from_feature_enabled( - self.features.enabled(Feature::UseAgentIdentity), - ) - } } #[derive(Clone)] diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 965521ae3d9e..933ce0ac764e 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -43,7 +43,7 @@ use codex_mcp::ToolInfo; use codex_mcp::ToolPluginProvenance; use codex_mcp::codex_apps_tools_cache_key; use codex_mcp::compute_auth_statuses; -use codex_mcp::with_codex_apps_mcp_with_authorization_header; +use codex_mcp::with_codex_apps_mcp; const CONNECTORS_READY_TIMEOUT_ON_EMPTY_TOOLS: Duration = Duration::from_secs(30); const DIRECTORY_CONNECTORS_TIMEOUT: Duration = Duration::from_secs(60); @@ -220,20 +220,8 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_options_and_status( }); } - let background_authorization_header_value = if let Some(auth) = auth.as_ref() { - auth_manager - .chatgpt_authorization_header_for_auth(auth) - .await - } else { - None - }; let mcp_config = config.to_mcp_config(plugins_manager.as_ref()).await; - let mcp_servers = with_codex_apps_mcp_with_authorization_header( - HashMap::new(), - auth.as_ref(), - &mcp_config, - background_authorization_header_value.as_deref(), - ); + let mcp_servers = with_codex_apps_mcp(HashMap::new(), auth.as_ref(), &mcp_config); if mcp_servers.is_empty() { return Ok(AccessibleConnectorsStatus { connectors: Vec::new(), @@ -435,18 +423,6 @@ async fn list_directory_connectors_for_tool_suggest_with_auth( }; let access_token = token_data.access_token.clone(); let account_id = account_id.to_string(); - let is_fedramp_account = token_data.id_token.is_fedramp_account(); - let authorization_header_value = { - let auth_manager = - AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false); - match auth { - Some(auth) if auth.is_chatgpt_auth() => auth_manager - .chatgpt_authorization_header_for_auth(auth) - .await - .unwrap_or_else(|| format!("Bearer {access_token}")), - _ => format!("Bearer {access_token}"), - } - }; let is_workspace_account = token_data.id_token.is_workspace_account(); let cache_key = AllConnectorsCacheKey::new( config.chatgpt_base_url.clone(), @@ -460,15 +436,14 @@ async fn list_directory_connectors_for_tool_suggest_with_auth( is_workspace_account, /*force_refetch*/ false, |path| { - let authorization_header_value = authorization_header_value.clone(); + let access_token = access_token.clone(); let account_id = account_id.clone(); async move { - chatgpt_get_request_with_authorization_header::( + chatgpt_get_request_with_token::( config, path, - authorization_header_value.as_str(), + access_token.as_str(), account_id.as_str(), - is_fedramp_account, ) .await } @@ -477,25 +452,23 @@ async fn list_directory_connectors_for_tool_suggest_with_auth( .await } -async fn chatgpt_get_request_with_authorization_header( +async fn chatgpt_get_request_with_token( config: &Config, path: String, - authorization_header_value: &str, + access_token: &str, account_id: &str, - is_fedramp_account: bool, ) -> anyhow::Result { let client = create_client(); let url = format!("{}{}", config.chatgpt_base_url, path); - let mut request = client + let response = client .get(&url) - .header("authorization", authorization_header_value) + .bearer_auth(access_token) .header("chatgpt-account-id", account_id) .header("Content-Type", "application/json") - .timeout(DIRECTORY_CONNECTORS_TIMEOUT); - if is_fedramp_account { - request = request.header("X-OpenAI-Fedramp", "true"); - } - let response = request.send().await.context("failed to send request")?; + .timeout(DIRECTORY_CONNECTORS_TIMEOUT) + .send() + .await + .context("failed to send request")?; if response.status().is_success() { response diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 715a1300d8b4..0d4c26991d40 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -7,7 +7,7 @@ use codex_config::McpServerConfig; use codex_login::CodexAuth; use codex_mcp::ToolPluginProvenance; use codex_mcp::configured_mcp_servers; -use codex_mcp::effective_mcp_servers_with_authorization_header; +use codex_mcp::effective_mcp_servers; use codex_mcp::tool_plugin_provenance as collect_tool_plugin_provenance; #[derive(Clone)] @@ -29,25 +29,9 @@ impl McpManager { &self, config: &Config, auth: Option<&CodexAuth>, - ) -> HashMap { - self.effective_servers_with_authorization_header( - config, auth, /*authorization_header_value*/ None, - ) - .await - } - - pub async fn effective_servers_with_authorization_header( - &self, - config: &Config, - auth: Option<&CodexAuth>, - authorization_header_value: Option<&str>, ) -> HashMap { let mcp_config = config.to_mcp_config(self.plugins_manager.as_ref()).await; - effective_mcp_servers_with_authorization_header( - &mcp_config, - auth, - authorization_header_value, - ) + effective_mcp_servers(&mcp_config, auth) } pub async fn tool_plugin_provenance(&self, config: &Config) -> ToolPluginProvenance { diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 71efd2332650..958b98b2f168 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -475,22 +475,10 @@ pub async fn reload_user_config(sess: &Arc) { pub async fn list_mcp_tools(sess: &Session, config: &Arc, sub_id: String) { let mcp_connection_manager = sess.services.mcp_connection_manager.read().await; let auth = sess.services.auth_manager.auth().await; - let background_authorization_header_value = if let Some(auth) = auth.as_ref() { - sess.services - .auth_manager - .chatgpt_authorization_header_for_auth(auth) - .await - } else { - None - }; let mcp_servers = sess .services .mcp_manager - .effective_servers_with_authorization_header( - config, - auth.as_ref(), - background_authorization_header_value.as_deref(), - ) + .effective_servers(config, auth.as_ref()) .await; let snapshot = collect_mcp_snapshot_from_manager( &mcp_connection_manager, diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index be7504d9e79f..6fcc5336754a 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -218,20 +218,7 @@ impl Session { .mcp_manager .tool_plugin_provenance(config.as_ref()) .await; - let background_authorization_header_value = if let Some(auth) = auth.as_ref() { - self.services - .auth_manager - .chatgpt_authorization_header_for_auth(auth) - .await - } else { - None - }; - let mcp_servers = with_codex_apps_mcp_with_authorization_header( - mcp_servers, - auth.as_ref(), - &mcp_config, - background_authorization_header_value.as_deref(), - ); + let mcp_servers = with_codex_apps_mcp(mcp_servers, auth.as_ref(), &mcp_config); let auth_statuses = compute_auth_statuses(mcp_servers.iter(), store_mode).await; { let mut guard = self.services.mcp_startup_cancellation_token.lock().await; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 4b919800cd93..87ecd14c9fe8 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -302,7 +302,7 @@ use crate::unified_exec::UnifiedExecProcessManager; use crate::windows_sandbox::WindowsSandboxLevelExt; use codex_git_utils::get_git_repo_root; use codex_mcp::compute_auth_statuses; -use codex_mcp::with_codex_apps_mcp_with_authorization_header; +use codex_mcp::with_codex_apps_mcp; use codex_otel::SessionTelemetry; use codex_otel::THREAD_STARTED_METRIC; use codex_otel::TelemetryAuthMode; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 16e86e3aeac8..b96f56abaa51 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -337,20 +337,8 @@ impl Session { let mcp_manager_for_mcp = Arc::clone(&mcp_manager); let auth_and_mcp_fut = async move { let auth = auth_manager_clone.auth().await; - let authorization_header_value = match auth.as_ref() { - Some(auth) => { - auth_manager_clone - .chatgpt_authorization_header_for_auth(auth) - .await - } - None => None, - }; let mcp_servers = mcp_manager_for_mcp - .effective_servers_with_authorization_header( - &config_for_mcp, - auth.as_ref(), - authorization_header_value.as_deref(), - ) + .effective_servers(&config_for_mcp, auth.as_ref()) .await; let auth_statuses = compute_auth_statuses( mcp_servers.iter(), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index f1a11bfcb604..6669115113d1 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4603,7 +4603,6 @@ fn seed_stored_identity( agent_runtime_id: stored_identity.agent_runtime_id.clone(), agent_private_key: stored_identity.private_key_pkcs8_base64.clone(), registered_at: stored_identity.registered_at.clone(), - background_task_id: None, }) .expect("store identity"); diff --git a/codex-rs/login/Cargo.toml b/codex-rs/login/Cargo.toml index c56c891e1d68..16d183602485 100644 --- a/codex-rs/login/Cargo.toml +++ b/codex-rs/login/Cargo.toml @@ -21,7 +21,6 @@ codex-otel = { workspace = true } codex-protocol = { workspace = true } codex-terminal-detection = { workspace = true } codex-utils-template = { workspace = true } -crypto_box = { workspace = true } ed25519-dalek = { workspace = true } once_cell = { workspace = true } os_info = { workspace = true } diff --git a/codex-rs/login/src/agent_identity.rs b/codex-rs/login/src/agent_identity.rs deleted file mode 100644 index 864c45ad0fa3..000000000000 --- a/codex-rs/login/src/agent_identity.rs +++ /dev/null @@ -1,834 +0,0 @@ -use std::collections::BTreeMap; -use std::sync::Arc; -use std::time::Duration; - -use anyhow::Context; -use anyhow::Result; -use base64::Engine as _; -use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; -use base64::engine::general_purpose::URL_SAFE_NO_PAD; -use chrono::SecondsFormat; -use chrono::Utc; -use codex_protocol::protocol::SessionSource; -use crypto_box::SecretKey as Curve25519SecretKey; -use ed25519_dalek::Signer as _; -use ed25519_dalek::SigningKey; -use ed25519_dalek::VerifyingKey; -use ed25519_dalek::pkcs8::DecodePrivateKey; -use ed25519_dalek::pkcs8::EncodePrivateKey; -use rand::TryRngCore; -use rand::rngs::OsRng; -use serde::Deserialize; -use serde::Serialize; -use sha2::Digest as _; -use sha2::Sha512; -use tokio::sync::Semaphore; -use tracing::debug; -use tracing::info; -use tracing::warn; - -use crate::AgentIdentityAuthRecord; -use crate::AuthManager; -use crate::CodexAuth; -use crate::default_client::create_client; - -const AGENT_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15); -const AGENT_TASK_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15); -const AGENT_IDENTITY_BISCUIT_TIMEOUT: Duration = Duration::from_secs(15); - -#[derive(Clone)] -pub(crate) struct BackgroundAgentTaskManager { - auth_manager: Arc, - chatgpt_base_url: String, - auth_mode: BackgroundAgentTaskAuthMode, - abom: AgentBillOfMaterials, - ensure_lock: Arc, -} - -impl std::fmt::Debug for BackgroundAgentTaskManager { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("BackgroundAgentTaskManager") - .field("chatgpt_base_url", &self.chatgpt_base_url) - .field("auth_mode", &self.auth_mode) - .field("abom", &self.abom) - .finish_non_exhaustive() - } -} - -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] -pub enum BackgroundAgentTaskAuthMode { - Enabled, - #[default] - Disabled, -} - -impl BackgroundAgentTaskAuthMode { - pub fn from_feature_enabled(enabled: bool) -> Self { - if enabled { - Self::Enabled - } else { - Self::Disabled - } - } - - fn is_enabled(self) -> bool { - matches!(self, Self::Enabled) - } -} - -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] -struct StoredAgentIdentity { - binding_id: String, - chatgpt_account_id: String, - chatgpt_user_id: Option, - agent_runtime_id: String, - private_key_pkcs8_base64: String, - public_key_ssh: String, - registered_at: String, - background_task_id: Option, - abom: AgentBillOfMaterials, -} - -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] -struct AgentBillOfMaterials { - agent_version: String, - agent_harness_id: String, - running_location: String, -} - -#[derive(Debug, Serialize)] -struct RegisterAgentRequest { - abom: AgentBillOfMaterials, - agent_public_key: String, - capabilities: Vec, -} - -#[derive(Debug, Deserialize)] -struct RegisterAgentResponse { - agent_runtime_id: String, -} - -#[derive(Debug, Serialize)] -struct RegisterTaskRequest { - signature: String, - timestamp: String, -} - -#[derive(Debug, Deserialize)] -struct RegisterTaskResponse { - encrypted_task_id: String, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -struct AgentIdentityBinding { - binding_id: String, - chatgpt_account_id: String, - chatgpt_user_id: Option, - access_token: String, -} - -struct GeneratedAgentKeyMaterial { - private_key_pkcs8_base64: String, - public_key_ssh: String, -} - -impl BackgroundAgentTaskManager { - #[cfg(test)] - pub(crate) fn new( - auth_manager: Arc, - chatgpt_base_url: String, - session_source: SessionSource, - ) -> Self { - Self::new_with_auth_mode( - auth_manager, - chatgpt_base_url, - session_source, - BackgroundAgentTaskAuthMode::Disabled, - ) - } - - pub(crate) fn new_with_auth_mode( - auth_manager: Arc, - chatgpt_base_url: String, - session_source: SessionSource, - auth_mode: BackgroundAgentTaskAuthMode, - ) -> Self { - Self { - auth_manager, - chatgpt_base_url: normalize_chatgpt_base_url(&chatgpt_base_url), - auth_mode, - abom: build_abom(session_source), - ensure_lock: Arc::new(Semaphore::new(/*permits*/ 1)), - } - } - - pub(crate) async fn authorization_header_value_for_auth( - &self, - auth: &CodexAuth, - ) -> Result> { - if !self.auth_mode.is_enabled() { - debug!("skipping background agent task auth because agent identity is disabled"); - return Ok(None); - } - - if !supports_background_agent_task_auth(&self.chatgpt_base_url) { - debug!( - chatgpt_base_url = %self.chatgpt_base_url, - "skipping background agent task auth for unsupported backend host" - ); - return Ok(None); - } - - let Some(binding) = - AgentIdentityBinding::from_auth(auth, self.auth_manager.forced_chatgpt_workspace_id()) - else { - debug!("skipping background agent task auth because ChatGPT auth is unavailable"); - return Ok(None); - }; - - let _guard = self - .ensure_lock - .acquire() - .await - .context("background agent task ensure semaphore closed")?; - let mut stored_identity = self - .ensure_registered_identity_for_binding(auth, &binding) - .await?; - let background_task_id = match stored_identity.background_task_id.clone() { - Some(background_task_id) => background_task_id, - _ => { - let background_task_id = self - .register_background_task_for_identity(&binding, &stored_identity) - .await?; - stored_identity.background_task_id = Some(background_task_id.clone()); - self.store_identity(auth, &stored_identity)?; - background_task_id - } - }; - - Ok(Some(authorization_header_for_task( - &stored_identity, - &background_task_id, - )?)) - } - - pub(crate) async fn authorization_header_value_or_bearer( - &self, - auth: &CodexAuth, - ) -> Option { - match self.authorization_header_value_for_auth(auth).await { - Ok(Some(authorization_header_value)) => Some(authorization_header_value), - Ok(None) => auth - .get_token() - .ok() - .filter(|token| !token.is_empty()) - .map(|token| format!("Bearer {token}")), - Err(error) => { - warn!( - error = %error, - "falling back to bearer authorization because background agent task auth failed" - ); - auth.get_token() - .ok() - .filter(|token| !token.is_empty()) - .map(|token| format!("Bearer {token}")) - } - } - } - - async fn ensure_registered_identity_for_binding( - &self, - auth: &CodexAuth, - binding: &AgentIdentityBinding, - ) -> Result { - if let Some(stored_identity) = self.load_stored_identity(auth, binding)? { - return Ok(stored_identity); - } - - let stored_identity = self.register_agent_identity(binding).await?; - self.store_identity(auth, &stored_identity)?; - Ok(stored_identity) - } - - async fn register_agent_identity( - &self, - binding: &AgentIdentityBinding, - ) -> Result { - let key_material = generate_agent_key_material()?; - let request_body = RegisterAgentRequest { - abom: self.abom.clone(), - agent_public_key: key_material.public_key_ssh.clone(), - capabilities: Vec::new(), - }; - - let url = agent_registration_url(&self.chatgpt_base_url); - let human_biscuit = self.mint_human_biscuit(binding, "POST", &url).await?; - let client = create_client(); - let response = client - .post(&url) - .header("X-OpenAI-Authorization", human_biscuit) - .json(&request_body) - .timeout(AGENT_REGISTRATION_TIMEOUT) - .send() - .await - .with_context(|| { - format!("failed to send agent identity registration request to {url}") - })?; - - if response.status().is_success() { - let response_body = response - .json::() - .await - .with_context(|| format!("failed to parse agent identity response from {url}"))?; - let stored_identity = StoredAgentIdentity { - binding_id: binding.binding_id.clone(), - chatgpt_account_id: binding.chatgpt_account_id.clone(), - chatgpt_user_id: binding.chatgpt_user_id.clone(), - agent_runtime_id: response_body.agent_runtime_id, - private_key_pkcs8_base64: key_material.private_key_pkcs8_base64, - public_key_ssh: key_material.public_key_ssh, - registered_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), - background_task_id: None, - abom: self.abom.clone(), - }; - info!( - agent_runtime_id = %stored_identity.agent_runtime_id, - binding_id = %binding.binding_id, - "registered background agent identity" - ); - return Ok(stored_identity); - } - - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - anyhow::bail!("agent identity registration failed with status {status} from {url}: {body}") - } - - async fn register_background_task_for_identity( - &self, - binding: &AgentIdentityBinding, - stored_identity: &StoredAgentIdentity, - ) -> Result { - let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); - let request_body = RegisterTaskRequest { - signature: sign_task_registration_payload(stored_identity, ×tamp)?, - timestamp, - }; - - let client = create_client(); - let url = - agent_task_registration_url(&self.chatgpt_base_url, &stored_identity.agent_runtime_id); - let human_biscuit = self.mint_human_biscuit(binding, "POST", &url).await?; - let response = client - .post(&url) - .header("X-OpenAI-Authorization", human_biscuit) - .json(&request_body) - .timeout(AGENT_TASK_REGISTRATION_TIMEOUT) - .send() - .await - .with_context(|| format!("failed to send background agent task request to {url}"))?; - - if response.status().is_success() { - let response_body = response - .json::() - .await - .with_context(|| format!("failed to parse background task response from {url}"))?; - let background_task_id = - decrypt_task_id_response(stored_identity, &response_body.encrypted_task_id)?; - info!( - agent_runtime_id = %stored_identity.agent_runtime_id, - task_id = %background_task_id, - "registered background agent task" - ); - return Ok(background_task_id); - } - - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - anyhow::bail!( - "background agent task registration failed with status {status} from {url}: {body}" - ) - } - - async fn mint_human_biscuit( - &self, - binding: &AgentIdentityBinding, - target_method: &str, - target_url: &str, - ) -> Result { - let url = agent_identity_biscuit_url(&self.chatgpt_base_url); - let request_id = agent_identity_request_id()?; - let client = create_client(); - let response = client - .get(&url) - .bearer_auth(&binding.access_token) - .header("X-Request-Id", request_id.clone()) - .header("X-Original-Method", target_method) - .header("X-Original-Url", target_url) - .timeout(AGENT_IDENTITY_BISCUIT_TIMEOUT) - .send() - .await - .with_context(|| format!("failed to send agent identity biscuit request to {url}"))?; - - if response.status().is_success() { - let human_biscuit = response - .headers() - .get("x-openai-authorization") - .context("agent identity biscuit response did not include x-openai-authorization")? - .to_str() - .context("agent identity biscuit response header was not valid UTF-8")? - .to_string(); - info!( - request_id = %request_id, - "minted human biscuit for background agent task" - ); - return Ok(human_biscuit); - } - - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - anyhow::bail!( - "agent identity biscuit minting failed with status {status} from {url}: {body}" - ) - } - - fn load_stored_identity( - &self, - auth: &CodexAuth, - binding: &AgentIdentityBinding, - ) -> Result> { - let Some(record) = auth.get_agent_identity(&binding.chatgpt_account_id) else { - return Ok(None); - }; - - let stored_identity = - match StoredAgentIdentity::from_auth_record(binding, record, self.abom.clone()) { - Ok(stored_identity) => stored_identity, - Err(error) => { - warn!( - binding_id = %binding.binding_id, - error = %error, - "stored agent identity is invalid; deleting cached value" - ); - auth.remove_agent_identity()?; - return Ok(None); - } - }; - - if !stored_identity.matches_binding(binding) { - warn!( - binding_id = %binding.binding_id, - "stored agent identity binding no longer matches current auth; deleting cached value" - ); - auth.remove_agent_identity()?; - return Ok(None); - } - - if let Err(error) = stored_identity.validate_key_material() { - warn!( - agent_runtime_id = %stored_identity.agent_runtime_id, - binding_id = %binding.binding_id, - error = %error, - "stored agent identity key material is invalid; deleting cached value" - ); - auth.remove_agent_identity()?; - return Ok(None); - } - - Ok(Some(stored_identity)) - } - - fn store_identity( - &self, - auth: &CodexAuth, - stored_identity: &StoredAgentIdentity, - ) -> Result<()> { - auth.set_agent_identity(stored_identity.to_auth_record())?; - Ok(()) - } -} - -pub fn cached_background_agent_task_authorization_header_value( - auth: &CodexAuth, - auth_mode: BackgroundAgentTaskAuthMode, -) -> Result> { - if !auth_mode.is_enabled() { - return Ok(None); - } - - let Some(binding) = AgentIdentityBinding::from_auth(auth, /*forced_workspace_id*/ None) else { - return Ok(None); - }; - let Some(record) = auth.get_agent_identity(&binding.chatgpt_account_id) else { - return Ok(None); - }; - let stored_identity = - StoredAgentIdentity::from_auth_record(&binding, record, build_abom(SessionSource::Cli))?; - if !stored_identity.matches_binding(&binding) { - return Ok(None); - } - stored_identity.validate_key_material()?; - let Some(background_task_id) = stored_identity.background_task_id.as_ref() else { - return Ok(None); - }; - authorization_header_for_task(&stored_identity, background_task_id).map(Some) -} - -impl StoredAgentIdentity { - fn from_auth_record( - binding: &AgentIdentityBinding, - record: AgentIdentityAuthRecord, - abom: AgentBillOfMaterials, - ) -> Result { - if record.workspace_id != binding.chatgpt_account_id { - anyhow::bail!( - "stored agent identity workspace {:?} does not match current workspace {:?}", - record.workspace_id, - binding.chatgpt_account_id - ); - } - let signing_key = signing_key_from_private_key_pkcs8_base64(&record.agent_private_key)?; - Ok(Self { - binding_id: binding.binding_id.clone(), - chatgpt_account_id: binding.chatgpt_account_id.clone(), - chatgpt_user_id: record.chatgpt_user_id, - agent_runtime_id: record.agent_runtime_id.clone(), - private_key_pkcs8_base64: record.agent_private_key, - public_key_ssh: encode_ssh_ed25519_public_key(&signing_key.verifying_key()), - registered_at: record.registered_at, - background_task_id: record.background_task_id, - abom, - }) - } - - fn to_auth_record(&self) -> AgentIdentityAuthRecord { - AgentIdentityAuthRecord { - workspace_id: self.chatgpt_account_id.clone(), - chatgpt_user_id: self.chatgpt_user_id.clone(), - agent_runtime_id: self.agent_runtime_id.clone(), - agent_private_key: self.private_key_pkcs8_base64.clone(), - registered_at: self.registered_at.clone(), - background_task_id: self.background_task_id.clone(), - } - } - - fn matches_binding(&self, binding: &AgentIdentityBinding) -> bool { - binding.matches_parts( - &self.binding_id, - &self.chatgpt_account_id, - self.chatgpt_user_id.as_deref(), - ) - } - - fn validate_key_material(&self) -> Result<()> { - let signing_key = self.signing_key()?; - let derived_public_key = encode_ssh_ed25519_public_key(&signing_key.verifying_key()); - anyhow::ensure!( - self.public_key_ssh == derived_public_key, - "stored public key does not match the private key" - ); - Ok(()) - } - - fn signing_key(&self) -> Result { - signing_key_from_private_key_pkcs8_base64(&self.private_key_pkcs8_base64) - } -} - -impl AgentIdentityBinding { - fn matches_parts( - &self, - binding_id: &str, - chatgpt_account_id: &str, - chatgpt_user_id: Option<&str>, - ) -> bool { - binding_id == self.binding_id - && chatgpt_account_id == self.chatgpt_account_id - && match self.chatgpt_user_id.as_deref() { - Some(expected_user_id) => chatgpt_user_id == Some(expected_user_id), - None => true, - } - } - - fn from_auth(auth: &CodexAuth, forced_workspace_id: Option) -> Option { - if !auth.is_chatgpt_auth() { - return None; - } - - let token_data = auth.get_token_data().ok()?; - let resolved_account_id = - forced_workspace_id - .filter(|value| !value.is_empty()) - .or(token_data - .account_id - .clone() - .filter(|value| !value.is_empty()))?; - - Some(Self { - binding_id: format!("chatgpt-account-{resolved_account_id}"), - chatgpt_account_id: resolved_account_id, - chatgpt_user_id: token_data - .id_token - .chatgpt_user_id - .filter(|value| !value.is_empty()), - access_token: token_data.access_token, - }) - } -} - -fn authorization_header_for_task( - stored_identity: &StoredAgentIdentity, - background_task_id: &str, -) -> Result { - let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); - let signature = sign_agent_assertion_payload(stored_identity, background_task_id, ×tamp)?; - let payload = serde_json::to_vec(&BTreeMap::from([ - ( - "agent_runtime_id", - stored_identity.agent_runtime_id.as_str(), - ), - ("signature", signature.as_str()), - ("task_id", background_task_id), - ("timestamp", timestamp.as_str()), - ])) - .context("failed to serialize agent assertion envelope")?; - Ok(format!( - "AgentAssertion {}", - URL_SAFE_NO_PAD.encode(payload) - )) -} - -fn sign_agent_assertion_payload( - stored_identity: &StoredAgentIdentity, - background_task_id: &str, - timestamp: &str, -) -> Result { - let signing_key = stored_identity.signing_key()?; - let payload = format!( - "{}:{background_task_id}:{timestamp}", - stored_identity.agent_runtime_id - ); - Ok(BASE64_STANDARD.encode(signing_key.sign(payload.as_bytes()).to_bytes())) -} - -fn sign_task_registration_payload( - stored_identity: &StoredAgentIdentity, - timestamp: &str, -) -> Result { - let signing_key = stored_identity.signing_key()?; - let payload = format!("{}:{timestamp}", stored_identity.agent_runtime_id); - Ok(BASE64_STANDARD.encode(signing_key.sign(payload.as_bytes()).to_bytes())) -} - -fn decrypt_task_id_response( - stored_identity: &StoredAgentIdentity, - encrypted_task_id: &str, -) -> Result { - let signing_key = stored_identity.signing_key()?; - let ciphertext = BASE64_STANDARD - .decode(encrypted_task_id) - .context("encrypted task id is not valid base64")?; - let plaintext = curve25519_secret_key_from_signing_key(&signing_key) - .unseal(&ciphertext) - .map_err(|_| anyhow::anyhow!("failed to decrypt encrypted task id"))?; - String::from_utf8(plaintext).context("decrypted task id is not valid UTF-8") -} - -fn curve25519_secret_key_from_signing_key(signing_key: &SigningKey) -> Curve25519SecretKey { - let digest = Sha512::digest(signing_key.to_bytes()); - let mut secret_key = [0u8; 32]; - secret_key.copy_from_slice(&digest[..32]); - secret_key[0] &= 248; - secret_key[31] &= 127; - secret_key[31] |= 64; - Curve25519SecretKey::from(secret_key) -} - -fn build_abom(session_source: SessionSource) -> AgentBillOfMaterials { - AgentBillOfMaterials { - agent_version: env!("CARGO_PKG_VERSION").to_string(), - agent_harness_id: match &session_source { - SessionSource::VSCode => "codex-app".to_string(), - SessionSource::Cli - | SessionSource::Exec - | SessionSource::Mcp - | SessionSource::Custom(_) - | SessionSource::SubAgent(_) - | SessionSource::Unknown => "codex-cli".to_string(), - }, - running_location: format!("{}-{}", session_source, std::env::consts::OS), - } -} - -fn generate_agent_key_material() -> Result { - let mut secret_key_bytes = [0u8; 32]; - OsRng - .try_fill_bytes(&mut secret_key_bytes) - .context("failed to generate agent identity private key bytes")?; - let signing_key = SigningKey::from_bytes(&secret_key_bytes); - let private_key_pkcs8 = signing_key - .to_pkcs8_der() - .context("failed to encode agent identity private key as PKCS#8")?; - - Ok(GeneratedAgentKeyMaterial { - private_key_pkcs8_base64: BASE64_STANDARD.encode(private_key_pkcs8.as_bytes()), - public_key_ssh: encode_ssh_ed25519_public_key(&signing_key.verifying_key()), - }) -} - -fn encode_ssh_ed25519_public_key(verifying_key: &VerifyingKey) -> String { - let mut blob = Vec::with_capacity(4 + 11 + 4 + 32); - append_ssh_string(&mut blob, b"ssh-ed25519"); - append_ssh_string(&mut blob, verifying_key.as_bytes()); - format!("ssh-ed25519 {}", BASE64_STANDARD.encode(blob)) -} - -fn append_ssh_string(buf: &mut Vec, value: &[u8]) { - buf.extend_from_slice(&(value.len() as u32).to_be_bytes()); - buf.extend_from_slice(value); -} - -fn signing_key_from_private_key_pkcs8_base64(private_key_pkcs8_base64: &str) -> Result { - let private_key = BASE64_STANDARD - .decode(private_key_pkcs8_base64) - .context("stored agent identity private key is not valid base64")?; - SigningKey::from_pkcs8_der(&private_key) - .context("stored agent identity private key is not valid PKCS#8") -} - -fn agent_registration_url(chatgpt_base_url: &str) -> String { - let trimmed = chatgpt_base_url.trim_end_matches('/'); - format!("{trimmed}/v1/agent/register") -} - -fn agent_task_registration_url(chatgpt_base_url: &str, agent_runtime_id: &str) -> String { - let trimmed = chatgpt_base_url.trim_end_matches('/'); - format!("{trimmed}/v1/agent/{agent_runtime_id}/task/register") -} - -fn agent_identity_biscuit_url(chatgpt_base_url: &str) -> String { - let trimmed = chatgpt_base_url.trim_end_matches('/'); - format!("{trimmed}/authenticate_app_v2") -} - -fn agent_identity_request_id() -> Result { - let mut request_id_bytes = [0u8; 16]; - OsRng - .try_fill_bytes(&mut request_id_bytes) - .context("failed to generate agent identity request id")?; - Ok(format!( - "codex-agent-identity-{}", - URL_SAFE_NO_PAD.encode(request_id_bytes) - )) -} - -fn normalize_chatgpt_base_url(chatgpt_base_url: &str) -> String { - let mut base_url = chatgpt_base_url.trim_end_matches('/').to_string(); - for suffix in [ - "/wham/remote/control/server/enroll", - "/wham/remote/control/server", - ] { - if let Some(stripped) = base_url.strip_suffix(suffix) { - base_url = stripped.to_string(); - break; - } - } - if (base_url.starts_with("https://chatgpt.com") - || base_url.starts_with("https://chat.openai.com")) - && !base_url.contains("/backend-api") - { - base_url = format!("{base_url}/backend-api"); - } - if let Some(stripped) = base_url.strip_suffix("/codex") { - stripped.to_string() - } else { - base_url - } -} - -fn supports_background_agent_task_auth(chatgpt_base_url: &str) -> bool { - let Ok(url) = url::Url::parse(chatgpt_base_url) else { - return false; - }; - let Some(host) = url.host_str() else { - return false; - }; - host == "chatgpt.com" - || host == "chat.openai.com" - || host == "chatgpt-staging.com" - || host.ends_with(".chatgpt.com") - || host.ends_with(".chatgpt-staging.com") -} - -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn disabled_background_agent_task_auth_returns_none_for_supported_host() { - let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); - let auth_manager = AuthManager::from_auth_for_testing(auth.clone()); - let manager = BackgroundAgentTaskManager::new_with_auth_mode( - auth_manager, - "https://chatgpt.com/backend-api".to_string(), - SessionSource::Cli, - BackgroundAgentTaskAuthMode::Disabled, - ); - - let authorization_header_value = manager - .authorization_header_value_for_auth(&auth) - .await - .expect("disabled manager should not fail"); - - assert_eq!(None, authorization_header_value); - } - - #[tokio::test] - async fn default_background_agent_task_auth_returns_none_for_supported_host() { - let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); - let auth_manager = AuthManager::from_auth_for_testing(auth.clone()); - let manager = BackgroundAgentTaskManager::new( - auth_manager, - "https://chatgpt.com/backend-api".to_string(), - SessionSource::Cli, - ); - - let authorization_header_value = manager - .authorization_header_value_for_auth(&auth) - .await - .expect("default manager should not fail"); - - assert_eq!(None, authorization_header_value); - } - - #[test] - fn cached_background_agent_task_auth_honors_disabled_mode() { - let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); - let key_material = generate_agent_key_material().expect("generate key material"); - auth.set_agent_identity(AgentIdentityAuthRecord { - workspace_id: "account_id".to_string(), - chatgpt_user_id: None, - agent_runtime_id: "agent_123".to_string(), - agent_private_key: key_material.private_key_pkcs8_base64, - registered_at: "2026-04-13T12:00:00Z".to_string(), - background_task_id: Some("task_123".to_string()), - }) - .expect("set agent identity"); - - let disabled_authorization_header_value = - cached_background_agent_task_authorization_header_value( - &auth, - BackgroundAgentTaskAuthMode::Disabled, - ) - .expect("disabled cached auth should not fail"); - let enabled_authorization_header_value = - cached_background_agent_task_authorization_header_value( - &auth, - BackgroundAgentTaskAuthMode::Enabled, - ) - .expect("enabled cached auth should not fail"); - - assert_eq!(None, disabled_authorization_header_value); - assert!(enabled_authorization_header_value.is_some()); - } -} diff --git a/codex-rs/login/src/auth/agent_assertion.rs b/codex-rs/login/src/auth/agent_assertion.rs index ce24c57d08aa..e4a1731f6538 100644 --- a/codex-rs/login/src/auth/agent_assertion.rs +++ b/codex-rs/login/src/auth/agent_assertion.rs @@ -167,7 +167,6 @@ mod tests { agent_runtime_id: agent_runtime_id.to_string(), agent_private_key: BASE64_STANDARD.encode(private_key.as_bytes()), registered_at: "2026-03-23T12:00:00Z".to_string(), - background_task_id: None, } } } diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 326adea6e438..fed497307077 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -211,7 +211,6 @@ fn chatgpt_auth_persists_agent_identity_for_workspace() { agent_runtime_id: "agent_123".to_string(), agent_private_key: "pkcs8-base64".to_string(), registered_at: "2026-04-13T12:00:00Z".to_string(), - background_task_id: None, }; auth.set_agent_identity(record.clone()) diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 29733f13da9d..9ead5ce8b27c 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -1,4 +1,3 @@ -use anyhow::Context; use async_trait::async_trait; use chrono::Utc; use reqwest::StatusCode; @@ -23,8 +22,6 @@ use codex_app_server_protocol::AuthMode as ApiAuthMode; use codex_protocol::config_types::ForcedLoginMethod; use codex_protocol::config_types::ModelProviderAuthInfo; -use super::agent_assertion; -use super::agent_assertion::AgentTaskAuthorizationTarget; use super::external_bearer::BearerTokenRefresher; use super::revoke::revoke_auth_tokens; pub use crate::auth::storage::AgentIdentityAuthRecord; @@ -43,12 +40,9 @@ use codex_protocol::auth::KnownPlan as InternalKnownPlan; use codex_protocol::auth::PlanType as InternalPlanType; use codex_protocol::auth::RefreshTokenFailedError; use codex_protocol::auth::RefreshTokenFailedReason; -use codex_protocol::protocol::SessionSource; use serde_json::Value; use thiserror::Error; -use crate::agent_identity::BackgroundAgentTaskAuthMode; -use crate::agent_identity::BackgroundAgentTaskManager; /// Authentication mechanism used by the current user. #[derive(Debug, Clone)] pub enum CodexAuth { @@ -378,7 +372,7 @@ impl CodexAuth { .filter(|identity| identity.workspace_id == workspace_id) } - pub fn set_agent_identity(&self, mut record: AgentIdentityAuthRecord) -> std::io::Result<()> { + pub fn set_agent_identity(&self, record: AgentIdentityAuthRecord) -> std::io::Result<()> { let (state, storage) = match self { Self::Chatgpt(auth) => (&auth.state, &auth.storage), Self::ChatgptAuthTokens(auth) => (&auth.state, &auth.storage), @@ -391,13 +385,6 @@ impl CodexAuth { let mut auth = guard .clone() .ok_or_else(|| std::io::Error::other("auth data is not available"))?; - if record.background_task_id.is_none() - && let Some(existing) = auth.agent_identity.as_ref() - && existing.workspace_id == record.workspace_id - && existing.agent_runtime_id == record.agent_runtime_id - { - record.background_task_id = existing.background_task_id.clone(); - } auth.agent_identity = Some(record); storage.save(&auth)?; *guard = Some(auth); @@ -1198,8 +1185,6 @@ pub struct AuthManager { enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, forced_chatgpt_workspace_id: RwLock>, - chatgpt_base_url: RwLock>, - background_agent_task_auth_mode: RwLock, refresh_lock: Semaphore, external_auth: RwLock>>, auth_state_tx: watch::Sender<()>, @@ -1220,16 +1205,6 @@ pub trait AuthManagerConfig { /// Returns the workspace ID that ChatGPT auth should be restricted to, if any. fn forced_chatgpt_workspace_id(&self) -> Option; - - /// Returns the ChatGPT backend base URL used for first-party backend authorization. - fn chatgpt_base_url(&self) -> Option { - None - } - - /// Returns whether default ChatGPT backend authorization may use background AgentAssertion. - fn background_agent_task_auth_mode(&self) -> BackgroundAgentTaskAuthMode { - BackgroundAgentTaskAuthMode::Disabled - } } impl Debug for AuthManager { @@ -1246,11 +1221,6 @@ impl Debug for AuthManager { "forced_chatgpt_workspace_id", &self.forced_chatgpt_workspace_id, ) - .field("chatgpt_base_url", &self.chatgpt_base_url) - .field( - "background_agent_task_auth_mode", - &self.background_agent_task_auth_mode, - ) .field("has_external_auth", &self.has_external_auth()) .finish_non_exhaustive() } @@ -1283,8 +1253,6 @@ impl AuthManager { enable_codex_api_key_env, auth_credentials_store_mode, forced_chatgpt_workspace_id: RwLock::new(None), - chatgpt_base_url: RwLock::new(None), - background_agent_task_auth_mode: RwLock::new(BackgroundAgentTaskAuthMode::Disabled), refresh_lock: Semaphore::new(/*permits*/ 1), external_auth: RwLock::new(None), auth_state_tx, @@ -1305,8 +1273,6 @@ impl AuthManager { enable_codex_api_key_env: false, auth_credentials_store_mode: AuthCredentialsStoreMode::File, forced_chatgpt_workspace_id: RwLock::new(None), - chatgpt_base_url: RwLock::new(None), - background_agent_task_auth_mode: RwLock::new(BackgroundAgentTaskAuthMode::Disabled), refresh_lock: Semaphore::new(/*permits*/ 1), external_auth: RwLock::new(None), auth_state_tx, @@ -1326,8 +1292,6 @@ impl AuthManager { enable_codex_api_key_env: false, auth_credentials_store_mode: AuthCredentialsStoreMode::File, forced_chatgpt_workspace_id: RwLock::new(None), - chatgpt_base_url: RwLock::new(None), - background_agent_task_auth_mode: RwLock::new(BackgroundAgentTaskAuthMode::Disabled), refresh_lock: Semaphore::new(/*permits*/ 1), external_auth: RwLock::new(None), auth_state_tx, @@ -1345,8 +1309,6 @@ impl AuthManager { enable_codex_api_key_env: false, auth_credentials_store_mode: AuthCredentialsStoreMode::File, forced_chatgpt_workspace_id: RwLock::new(None), - chatgpt_base_url: RwLock::new(None), - background_agent_task_auth_mode: RwLock::new(BackgroundAgentTaskAuthMode::Disabled), refresh_lock: Semaphore::new(/*permits*/ 1), external_auth: RwLock::new(Some( Arc::new(BearerTokenRefresher::new(config)) as Arc @@ -1528,127 +1490,6 @@ impl AuthManager { .and_then(|guard| guard.clone()) } - pub fn chatgpt_agent_task_authorization_header_for_auth( - &self, - auth: &CodexAuth, - target: AgentTaskAuthorizationTarget<'_>, - ) -> anyhow::Result> { - let Some(record) = self.agent_identity_for_chatgpt_auth(auth)? else { - return Ok(None); - }; - agent_assertion::authorization_header_for_agent_task(&record, target).map(Some) - } - - fn agent_identity_for_chatgpt_auth( - &self, - auth: &CodexAuth, - ) -> anyhow::Result> { - if !auth.is_chatgpt_auth() { - return Ok(None); - } - - let token_data = auth - .get_token_data() - .context("ChatGPT token data is not available")?; - let workspace_id = self - .forced_chatgpt_workspace_id() - .filter(|value| !value.is_empty()) - .or(token_data.account_id.filter(|value| !value.is_empty())); - - let Some(workspace_id) = workspace_id else { - return Ok(None); - }; - let Some(record) = auth.get_agent_identity(&workspace_id) else { - anyhow::bail!("agent identity is not available for workspace {workspace_id}"); - }; - - Ok(Some(record)) - } - - pub fn set_chatgpt_backend_auth_config( - &self, - chatgpt_base_url: Option, - background_agent_task_auth_mode: BackgroundAgentTaskAuthMode, - ) { - let mut changed = false; - if let Ok(mut guard) = self.chatgpt_base_url.write() - && *guard != chatgpt_base_url - { - *guard = chatgpt_base_url; - changed = true; - } - if let Ok(mut guard) = self.background_agent_task_auth_mode.write() - && *guard != background_agent_task_auth_mode - { - *guard = background_agent_task_auth_mode; - changed = true; - } - if changed { - self.auth_state_tx.send_replace(()); - } - } - - pub fn set_chatgpt_backend_base_url(&self, chatgpt_base_url: Option) { - let (_, auth_mode) = self.chatgpt_backend_auth_config(); - self.set_chatgpt_backend_auth_config(chatgpt_base_url, auth_mode); - } - - fn chatgpt_backend_auth_config(&self) -> (Option, BackgroundAgentTaskAuthMode) { - let chatgpt_base_url = self - .chatgpt_base_url - .read() - .ok() - .and_then(|guard| guard.clone()); - let auth_mode = self - .background_agent_task_auth_mode - .read() - .ok() - .map(|guard| *guard) - .unwrap_or_default(); - (chatgpt_base_url, auth_mode) - } - - /// Returns the default authorization header for ChatGPT backend requests. - /// - /// This uses background AgentAssertion when configured and available, otherwise it falls back - /// to the ChatGPT bearer token. Low-level bootstrap calls that must never use AgentAssertion - /// should use [`Self::chatgpt_bearer_authorization_header_for_auth`] instead. - pub async fn chatgpt_authorization_header(self: &Arc) -> Option { - let auth = self.auth().await?; - self.chatgpt_authorization_header_for_auth(&auth).await - } - - pub async fn chatgpt_authorization_header_for_auth( - self: &Arc, - auth: &CodexAuth, - ) -> Option { - if !auth.is_chatgpt_auth() { - return None; - } - - let (chatgpt_base_url, auth_mode) = self.chatgpt_backend_auth_config(); - let Some(chatgpt_base_url) = chatgpt_base_url else { - return Self::chatgpt_bearer_authorization_header_for_auth(auth); - }; - - BackgroundAgentTaskManager::new_with_auth_mode( - Arc::clone(self), - chatgpt_base_url, - SessionSource::Cli, - auth_mode, - ) - .authorization_header_value_or_bearer(auth) - .await - } - - pub fn chatgpt_bearer_token_for_auth(auth: &CodexAuth) -> Option { - auth.get_token().ok().filter(|token| !token.is_empty()) - } - - pub fn chatgpt_bearer_authorization_header_for_auth(auth: &CodexAuth) -> Option { - Self::chatgpt_bearer_token_for_auth(auth).map(|token| format!("Bearer {token}")) - } - pub fn subscribe_auth_state(&self) -> watch::Receiver<()> { self.auth_state_tx.subscribe() } @@ -1691,10 +1532,6 @@ impl AuthManager { config.cli_auth_credentials_store_mode(), ); auth_manager.set_forced_chatgpt_workspace_id(config.forced_chatgpt_workspace_id()); - auth_manager.set_chatgpt_backend_auth_config( - config.chatgpt_base_url(), - config.background_agent_task_auth_mode(), - ); auth_manager } diff --git a/codex-rs/login/src/auth/storage.rs b/codex-rs/login/src/auth/storage.rs index 5b752ee6725d..f0d7a3169c1d 100644 --- a/codex-rs/login/src/auth/storage.rs +++ b/codex-rs/login/src/auth/storage.rs @@ -52,8 +52,6 @@ pub struct AgentIdentityAuthRecord { pub agent_runtime_id: String, pub agent_private_key: String, pub registered_at: String, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub background_task_id: Option, } pub(super) fn get_auth_file(codex_home: &Path) -> PathBuf { diff --git a/codex-rs/login/src/auth/storage_tests.rs b/codex-rs/login/src/auth/storage_tests.rs index e00f09d46408..2e1cc8502977 100644 --- a/codex-rs/login/src/auth/storage_tests.rs +++ b/codex-rs/login/src/auth/storage_tests.rs @@ -69,7 +69,6 @@ async fn file_storage_persists_agent_identity() -> anyhow::Result<()> { agent_runtime_id: "agent_123".to_string(), agent_private_key: "pkcs8-base64".to_string(), registered_at: "2026-04-13T12:00:00Z".to_string(), - background_task_id: None, }), }; diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index af0e13b9b704..046d878e87b4 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -1,4 +1,3 @@ -pub mod agent_identity; pub mod auth; pub mod auth_env_telemetry; pub mod token_data; @@ -18,8 +17,6 @@ pub use server::ServerOptions; pub use server::ShutdownHandle; pub use server::run_login_server; -pub use agent_identity::BackgroundAgentTaskAuthMode; -pub use agent_identity::cached_background_agent_task_authorization_header_value; pub use auth::AgentIdentityAuthRecord; pub use auth::AgentTaskAuthorizationTarget; pub use auth::AuthConfig; diff --git a/codex-rs/models-manager/src/manager.rs b/codex-rs/models-manager/src/manager.rs index 6d5bdb9f7346..c029960a7039 100644 --- a/codex-rs/models-manager/src/manager.rs +++ b/codex-rs/models-manager/src/manager.rs @@ -17,7 +17,6 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::collect_auth_env_telemetry; use codex_login::default_client::build_reqwest_client; -use codex_model_provider::AuthorizationHeaderAuthProvider; use codex_model_provider::SharedModelProvider; use codex_model_provider::create_model_provider; use codex_model_provider_info::ModelProviderInfo; @@ -454,23 +453,7 @@ impl ModelsManager { let auth = self.provider.auth().await; let auth_mode = auth.as_ref().map(CodexAuth::auth_mode); let api_provider = self.provider.api_provider().await?; - let mut api_auth = self.provider.api_auth().await?; - if let Some(auth_manager) = auth_manager.as_ref() - && let Some(auth) = auth.as_ref().filter(|auth| auth.is_chatgpt_auth()) - && provider_uses_codex_login_auth(self.provider.info()) - && let Some(authorization_header_value) = auth_manager - .chatgpt_authorization_header_for_auth(auth) - .await - { - let mut auth_provider = AuthorizationHeaderAuthProvider::new( - Some(authorization_header_value), - auth.get_account_id(), - ); - if auth.is_fedramp_account() { - auth_provider = auth_provider.with_fedramp_routing_header(); - } - api_auth = Arc::new(auth_provider); - } + let api_auth = self.provider.api_auth().await?; let auth_env = collect_auth_env_telemetry(self.provider.info(), codex_api_key_env_enabled); let transport = ReqwestTransport::new(build_reqwest_client()); let auth_telemetry = auth_header_telemetry(api_auth.as_ref()); @@ -618,10 +601,6 @@ impl ModelsManager { } } -fn provider_uses_codex_login_auth(provider: &ModelProviderInfo) -> bool { - provider.env_key.is_none() && provider.experimental_bearer_token.is_none() -} - #[cfg(test)] #[path = "manager_tests.rs"] mod tests; From dbbe810b1fd53c90383ceb0ea14e880c6f21e1d6 Mon Sep 17 00:00:00 2001 From: Edward Frazer Date: Tue, 21 Apr 2026 09:29:10 -0700 Subject: [PATCH 2/9] Revert "[codex] Use AgentAssertion downstream behind use_agent_identity (#17980)" This reverts commit b44d2851cf0e7b728d4a848d4a4f68da3483dfd6. --- codex-rs/Cargo.lock | 1 - codex-rs/codex-api/src/files.rs | 61 ---- codex-rs/core/src/agent_identity.rs | 40 +-- .../src/agent_identity/task_registration.rs | 8 - codex-rs/core/src/arc_monitor.rs | 32 +-- codex-rs/core/src/arc_monitor_tests.rs | 146 ---------- codex-rs/core/src/client.rs | 103 +------ codex-rs/core/src/client_tests.rs | 260 ------------------ codex-rs/core/src/mcp_openai_file.rs | 211 +------------- .../core/src/session/agent_task_lifecycle.rs | 32 --- codex-rs/core/src/session/session.rs | 11 +- codex-rs/core/src/session/turn.rs | 48 ++-- codex-rs/login/Cargo.toml | 3 +- codex-rs/login/src/auth/agent_assertion.rs | 172 ------------ codex-rs/login/src/auth/mod.rs | 2 - codex-rs/login/src/lib.rs | 1 - .../src/bearer_auth_provider.rs | 99 ------- codex-rs/model-provider/src/lib.rs | 1 - 18 files changed, 65 insertions(+), 1166 deletions(-) delete mode 100644 codex-rs/login/src/auth/agent_assertion.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 0f87b3ec2ebd..859e54b926cf 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2416,7 +2416,6 @@ dependencies = [ "codex-terminal-detection", "codex-utils-template", "core_test_support", - "ed25519-dalek", "keyring", "once_cell", "os_info", diff --git a/codex-rs/codex-api/src/files.rs b/codex-rs/codex-api/src/files.rs index a938242934ca..d1e2840066d8 100644 --- a/codex-rs/codex-api/src/files.rs +++ b/codex-rs/codex-api/src/files.rs @@ -308,18 +308,6 @@ mod tests { ChatGptTestAuth } - #[derive(Clone, Copy)] - struct AgentAssertionTestAuth; - - impl AuthProvider for AgentAssertionTestAuth { - fn add_auth_headers(&self, headers: &mut reqwest::header::HeaderMap) { - headers.insert( - reqwest::header::AUTHORIZATION, - HeaderValue::from_static("AgentAssertion test-assertion"), - ); - } - } - fn base_url_for(server: &MockServer) -> String { format!("{}/backend-api", server.uri()) } @@ -329,7 +317,6 @@ mod tests { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/backend-api/files")) - .and(header("authorization", "Bearer token")) .and(header("chatgpt-account-id", "account_id")) .and(body_json(serde_json::json!({ "file_name": "hello.txt", @@ -390,52 +377,4 @@ mod tests { assert_eq!(uploaded.mime_type, Some("text/plain".to_string())); assert_eq!(finalize_attempts.load(Ordering::SeqCst), 2); } - - #[tokio::test] - async fn upload_local_file_uses_authorization_header_value() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/backend-api/files")) - .and(header("authorization", "AgentAssertion test-assertion")) - .and(body_json(serde_json::json!({ - "file_name": "hello.txt", - "file_size": 5, - "use_case": "codex", - }))) - .respond_with( - ResponseTemplate::new(200) - .set_body_json(serde_json::json!({"file_id": "file_123", "upload_url": format!("{}/upload/file_123", server.uri())})), - ) - .mount(&server) - .await; - Mock::given(method("PUT")) - .and(path("/upload/file_123")) - .and(header("content-length", "5")) - .respond_with(ResponseTemplate::new(200)) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/backend-api/files/file_123/uploaded")) - .and(header("authorization", "AgentAssertion test-assertion")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "status": "success", - "download_url": format!("{}/download/file_123", server.uri()), - "file_name": "hello.txt", - "mime_type": "text/plain", - "file_size_bytes": 5 - }))) - .mount(&server) - .await; - - let base_url = base_url_for(&server); - let dir = TempDir::new().expect("temp dir"); - let path = dir.path().join("hello.txt"); - tokio::fs::write(&path, b"hello").await.expect("write file"); - - let uploaded = upload_local_file(&base_url, &AgentAssertionTestAuth, &path) - .await - .expect("upload succeeds"); - - assert_eq!(uploaded.file_id, "file_123"); - } } diff --git a/codex-rs/core/src/agent_identity.rs b/codex-rs/core/src/agent_identity.rs index 5bd15de76c9c..e7a087461eb4 100644 --- a/codex-rs/core/src/agent_identity.rs +++ b/codex-rs/core/src/agent_identity.rs @@ -27,12 +27,12 @@ use tracing::debug; use tracing::info; use tracing::warn; -use crate::config::Config; - mod task_registration; 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); @@ -339,7 +339,7 @@ impl AgentIdentityManager { } #[cfg(test)] - pub(crate) fn new_for_tests( + fn new_for_tests( auth_manager: Arc, feature_enabled: bool, chatgpt_base_url: String, @@ -353,30 +353,6 @@ impl AgentIdentityManager { ensure_lock: Arc::new(Semaphore::new(/*permits*/ 1)), } } - - #[cfg(test)] - pub(crate) async fn seed_generated_identity_for_tests( - &self, - agent_runtime_id: &str, - ) -> Result { - let (auth, binding) = self - .current_auth_binding() - .await - .context("test agent identity requires ChatGPT auth")?; - let key_material = generate_agent_key_material()?; - let stored_identity = StoredAgentIdentity { - binding_id: binding.binding_id.clone(), - chatgpt_account_id: binding.chatgpt_account_id.clone(), - chatgpt_user_id: binding.chatgpt_user_id, - agent_runtime_id: agent_runtime_id.to_string(), - private_key_pkcs8_base64: key_material.private_key_pkcs8_base64, - public_key_ssh: key_material.public_key_ssh, - registered_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), - abom: self.abom.clone(), - }; - self.store_identity(&auth, &stored_identity)?; - Ok(stored_identity) - } } impl StoredAgentIdentity { @@ -607,7 +583,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) @@ -633,7 +609,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"); @@ -649,7 +625,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) @@ -681,11 +657,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/task_registration.rs b/codex-rs/core/src/agent_identity/task_registration.rs index d00f7d25e4a6..53bb272cbc00 100644 --- a/codex-rs/core/src/agent_identity/task_registration.rs +++ b/codex-rs/core/src/agent_identity/task_registration.rs @@ -2,7 +2,6 @@ use std::time::Duration; use anyhow::Context; use anyhow::Result; -use codex_login::AgentTaskAuthorizationTarget; use codex_protocol::protocol::SessionAgentTask; use crypto_box::SecretKey as Curve25519SecretKey; use ed25519_dalek::Signer as _; @@ -103,13 +102,6 @@ impl AgentIdentityManager { } impl RegisteredAgentTask { - pub(crate) fn authorization_target(&self) -> AgentTaskAuthorizationTarget<'_> { - AgentTaskAuthorizationTarget { - agent_runtime_id: &self.agent_runtime_id, - task_id: &self.task_id, - } - } - pub(crate) fn to_session_agent_task(&self) -> SessionAgentTask { SessionAgentTask { agent_runtime_id: self.agent_runtime_id.clone(), diff --git a/codex-rs/core/src/arc_monitor.rs b/codex-rs/core/src/arc_monitor.rs index 99767872d5e3..ecd7f3966628 100644 --- a/codex-rs/core/src/arc_monitor.rs +++ b/codex-rs/core/src/arc_monitor.rs @@ -13,7 +13,6 @@ use codex_login::CodexAuth; use codex_login::default_client::build_reqwest_client; use codex_protocol::models::MessagePhase; use codex_protocol::models::ResponseItem; -use reqwest::header::AUTHORIZATION; const ARC_MONITOR_TIMEOUT: Duration = Duration::from_secs(30); const CODEX_ARC_MONITOR_ENDPOINT_OVERRIDE: &str = "CODEX_ARC_MONITOR_ENDPOINT_OVERRIDE"; @@ -110,31 +109,13 @@ pub(crate) async fn monitor_action( }, None => None, }; - let (authorization_header_value, account_id) = if let Some(token) = - read_non_empty_env_var(CODEX_ARC_MONITOR_TOKEN) - { - ( - format!("Bearer {token}"), - auth.as_ref().and_then(CodexAuth::get_account_id), - ) - } else if let Some(authorization_header_value) = - match sess.authorization_header_for_current_agent_task().await { - Ok(authorization_header_value) => authorization_header_value, - Err(err) => { - warn!( - error = %err, - "skipping safety monitor because agent assertion authorization is unavailable" - ); - return ArcMonitorOutcome::Ok; - } - } - { - (authorization_header_value, None) + let token = if let Some(token) = read_non_empty_env_var(CODEX_ARC_MONITOR_TOKEN) { + token } else { let Some(auth) = auth.as_ref() else { return ArcMonitorOutcome::Ok; }; - let token = match auth.get_token() { + match auth.get_token() { Ok(token) => token, Err(err) => { warn!( @@ -143,8 +124,7 @@ pub(crate) async fn monitor_action( ); return ArcMonitorOutcome::Ok; } - }; - (format!("Bearer {token}"), auth.get_account_id()) + } }; let url = read_non_empty_env_var(CODEX_ARC_MONITOR_ENDPOINT_OVERRIDE).unwrap_or_else(|| { @@ -167,8 +147,8 @@ pub(crate) async fn monitor_action( .post(&url) .timeout(ARC_MONITOR_TIMEOUT) .json(&body) - .header(AUTHORIZATION, authorization_header_value); - if let Some(account_id) = account_id { + .bearer_auth(token); + if let Some(account_id) = auth.as_ref().and_then(CodexAuth::get_account_id) { request = request.header("chatgpt-account-id", account_id); } diff --git a/codex-rs/core/src/arc_monitor_tests.rs b/codex-rs/core/src/arc_monitor_tests.rs index 5fd691b93052..1cb29ce08cfc 100644 --- a/codex-rs/core/src/arc_monitor_tests.rs +++ b/codex-rs/core/src/arc_monitor_tests.rs @@ -10,38 +10,18 @@ use wiremock::MockServer; use wiremock::ResponseTemplate; use wiremock::matchers::body_json; use wiremock::matchers::header; -use wiremock::matchers::header_regex; use wiremock::matchers::method; use wiremock::matchers::path; use super::*; -use crate::agent_identity::AgentIdentityManager; -use crate::agent_identity::RegisteredAgentTask; use crate::context::ContextualUserFragment; use crate::session::tests::make_session_and_context; -use chrono::Utc; -use codex_login::AuthCredentialsStoreMode; -use codex_login::AuthDotJson; -use codex_login::AuthManager; -use codex_login::CodexAuth; -use codex_login::save_auth; -use codex_login::token_data::IdTokenInfo; -use codex_login::token_data::TokenData; use codex_protocol::models::ContentItem; use codex_protocol::models::LocalShellAction; use codex_protocol::models::LocalShellExecAction; use codex_protocol::models::LocalShellStatus; use codex_protocol::models::MessagePhase; use codex_protocol::models::ResponseItem; -use codex_protocol::protocol::SessionSource; -use tempfile::tempdir; - -const TEST_ID_TOKEN: &str = concat!( - "eyJhbGciOiJub25lIiwidHlwIjoiSldUIn0.", - "eyJodHRwczovL2FwaS5vcGVuYWkuY29tL2F1dGgiOnsiY2hhdGdwdF91c2VyX2lk", - "IjpudWxsLCJjaGF0Z3B0X2FjY291bnRfaWQiOiJhY2NvdW50X2lkIn19.", - "c2ln", -); struct EnvVarGuard { key: &'static str, @@ -71,58 +51,6 @@ impl Drop for EnvVarGuard { } } -async fn install_cached_agent_task_auth( - session: &mut Session, - turn_context: &mut TurnContext, - chatgpt_base_url: String, -) { - let auth_dir = tempdir().expect("temp auth dir"); - let auth_json = AuthDotJson { - auth_mode: Some(codex_app_server_protocol::AuthMode::Chatgpt), - openai_api_key: None, - tokens: Some(TokenData { - id_token: IdTokenInfo { - email: None, - chatgpt_plan_type: None, - chatgpt_user_id: None, - chatgpt_account_id: Some("account_id".to_string()), - chatgpt_account_is_fedramp: false, - raw_jwt: TEST_ID_TOKEN.to_string(), - }, - access_token: "Access Token".to_string(), - refresh_token: "test".to_string(), - account_id: Some("account_id".to_string()), - }), - last_refresh: Some(Utc::now()), - agent_identity: None, - }; - save_auth(auth_dir.path(), &auth_json, AuthCredentialsStoreMode::File).expect("save test auth"); - let auth = CodexAuth::from_auth_storage(auth_dir.path(), AuthCredentialsStoreMode::File) - .expect("load test auth") - .expect("test auth"); - let auth_manager = AuthManager::from_auth_for_testing(auth); - let agent_identity_manager = Arc::new(AgentIdentityManager::new_for_tests( - Arc::clone(&auth_manager), - /*feature_enabled*/ true, - chatgpt_base_url, - SessionSource::Exec, - )); - let stored_identity = agent_identity_manager - .seed_generated_identity_for_tests("agent-123") - .await - .expect("seed test identity"); - session.services.auth_manager = Arc::clone(&auth_manager); - session.services.agent_identity_manager = agent_identity_manager; - turn_context.auth_manager = Some(auth_manager); - session - .cache_agent_task_for_tests(RegisteredAgentTask { - agent_runtime_id: stored_identity.agent_runtime_id, - task_id: "task-123".to_string(), - registered_at: "2026-04-15T00:00:00Z".to_string(), - }) - .await; -} - #[tokio::test] async fn build_arc_monitor_request_includes_relevant_history_and_null_policies() { let (session, mut turn_context) = make_session_and_context().await; @@ -326,80 +254,6 @@ async fn build_arc_monitor_request_includes_relevant_history_and_null_policies() ); } -#[tokio::test] -#[serial(arc_monitor_env)] -async fn monitor_action_uses_agent_assertion_for_cached_task() { - let server = MockServer::start().await; - let (mut session, mut turn_context) = make_session_and_context().await; - install_cached_agent_task_auth(&mut session, &mut turn_context, server.uri()).await; - - let mut config = (*turn_context.config).clone(); - config.chatgpt_base_url = server.uri(); - turn_context.config = Arc::new(config); - - session - .record_into_history( - &[ResponseItem::Message { - id: None, - role: "user".to_string(), - content: vec![ContentItem::InputText { - text: "please run the tool".to_string(), - }], - end_turn: None, - phase: None, - }], - &turn_context, - ) - .await; - - Mock::given(method("POST")) - .and(path("/codex/safety/arc")) - .and(header_regex("authorization", r"^AgentAssertion .+")) - .and(body_json(serde_json::json!({ - "metadata": { - "codex_thread_id": session.conversation_id.to_string(), - "codex_turn_id": turn_context.sub_id.clone(), - "conversation_id": session.conversation_id.to_string(), - "protection_client_callsite": "normal", - }, - "messages": [{ - "role": "user", - "content": [{ - "type": "input_text", - "text": "please run the tool", - }], - }], - "policies": { - "developer": null, - "user": null, - }, - "action": { - "tool": "mcp_tool_call", - }, - }))) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "outcome": "ok", - "short_reason": "", - "rationale": "", - "risk_score": 1, - "risk_level": "low", - "evidence": [], - }))) - .expect(1) - .mount(&server) - .await; - - let outcome = monitor_action( - &session, - &turn_context, - serde_json::json!({ "tool": "mcp_tool_call" }), - "normal", - ) - .await; - - assert_eq!(outcome, ArcMonitorOutcome::Ok); -} - #[tokio::test] #[serial(arc_monitor_env)] async fn monitor_action_posts_expected_arc_request() { diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 50f9dd295ec4..601b8106df8f 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -31,7 +31,6 @@ use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; -use crate::agent_identity::RegisteredAgentTask; use codex_api::ApiError; use codex_api::AuthProvider; use codex_api::CompactClient as ApiCompactClient; @@ -96,7 +95,6 @@ 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; @@ -111,7 +109,6 @@ use codex_feedback::FeedbackRequestTags; use codex_feedback::emit_feedback_request_tags_with_auth_env; use codex_login::auth_env_telemetry::AuthEnvTelemetry; use codex_login::auth_env_telemetry::collect_auth_env_telemetry; -use codex_model_provider::AuthorizationHeaderAuthProvider; use codex_model_provider::SharedModelProvider; use codex_model_provider::create_model_provider; #[cfg(test)] @@ -216,8 +213,6 @@ pub struct ModelClient { pub struct ModelClientSession { client: ModelClient, websocket_session: WebsocketSession, - agent_task: Option, - cache_websocket_session_on_drop: bool, /// Turn state for sticky routing. /// /// This is an `OnceLock` that stores the turn state value received from the server @@ -335,25 +330,9 @@ impl ModelClient { /// This constructor does not perform network I/O itself; the session opens a websocket lazily /// when the first stream request is issued. pub fn new_session(&self) -> ModelClientSession { - self.new_session_with_agent_task(/*agent_task*/ None) - } - - pub(crate) fn new_session_with_agent_task( - &self, - agent_task: Option, - ) -> ModelClientSession { - let cache_websocket_session_on_drop = agent_task.is_none(); - let websocket_session = if agent_task.is_some() { - drop(self.take_cached_websocket_session()); - WebsocketSession::default() - } else { - self.take_cached_websocket_session() - }; ModelClientSession { client: self.clone(), - websocket_session, - agent_task, - cache_websocket_session_on_drop, + websocket_session: self.take_cached_websocket_session(), turn_state: Arc::new(OnceLock::new()), } } @@ -436,7 +415,7 @@ impl ModelClient { if prompt.input.is_empty() { return Ok(Vec::new()); } - let client_setup = self.current_client_setup(/*agent_task*/ None).await?; + let client_setup = self.current_client_setup().await?; let transport = ReqwestTransport::new(build_reqwest_client()); let request_telemetry = Self::build_request_telemetry( session_telemetry, @@ -500,7 +479,7 @@ impl ModelClient { ) -> Result { // Create the media call over HTTP first, then retain matching auth so realtime can attach // the server-side control WebSocket to the call id from that HTTP response. - let client_setup = self.current_client_setup(/*agent_task*/ None).await?; + let client_setup = self.current_client_setup().await?; let mut sideband_headers = extra_headers.clone(); sideband_headers.extend(sideband_websocket_auth_headers( client_setup.api_auth.as_ref(), @@ -535,7 +514,7 @@ impl ModelClient { return Ok(Vec::new()); } - let client_setup = self.current_client_setup(/*agent_task*/ None).await?; + let client_setup = self.current_client_setup().await?; let transport = ReqwestTransport::new(build_reqwest_client()); let request_telemetry = Self::build_request_telemetry( session_telemetry, @@ -685,46 +664,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, - agent_task: Option<&RegisteredAgentTask>, - ) -> Result { + async fn current_client_setup(&self) -> Result { let auth = self.state.provider.auth().await; let api_provider = self.state.provider.api_provider().await?; - let auth_manager = self.state.provider.auth_manager(); - let api_auth = match (agent_task, auth_manager.as_ref(), auth.as_ref()) { - (Some(agent_task), Some(auth_manager), Some(auth)) => { - if let Some(authorization_header_value) = auth_manager - .chatgpt_agent_task_authorization_header_for_auth( - auth, - agent_task.authorization_target(), - ) - .map_err(|err| { - CodexErr::Stream( - format!("failed to build agent assertion authorization: {err}"), - None, - ) - })? - { - debug!( - agent_runtime_id = %agent_task.agent_runtime_id, - task_id = %agent_task.task_id, - "using agent assertion authorization for downstream request" - ); - let mut auth_provider = AuthorizationHeaderAuthProvider::new( - Some(authorization_header_value), - /*account_id*/ None, - ); - if auth.is_fedramp_account() { - auth_provider = auth_provider.with_fedramp_routing_header(); - } - Arc::new(auth_provider) - } else { - self.state.provider.api_auth().await? - } - } - _ => self.state.provider.api_auth().await?, - }; + let api_auth = self.state.provider.api_auth().await?; Ok(CurrentClientSetup { auth, api_provider, @@ -858,18 +801,12 @@ impl ModelClient { impl Drop for ModelClientSession { fn drop(&mut self) { let websocket_session = std::mem::take(&mut self.websocket_session); - if self.cache_websocket_session_on_drop { - self.client - .store_cached_websocket_session(websocket_session); - } + 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; @@ -1071,15 +1008,11 @@ impl ModelClientSession { return Ok(()); } - let client_setup = self - .client - .current_client_setup(self.agent_task.as_ref()) - .await - .map_err(|err| { - ApiError::Stream(format!( - "failed to build websocket prewarm client setup: {err}" - )) - })?; + let client_setup = self.client.current_client_setup().await.map_err(|err| { + ApiError::Stream(format!( + "failed to build websocket prewarm client setup: {err}" + )) + })?; let auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), client_setup.api_auth.as_ref(), @@ -1233,10 +1166,7 @@ impl ModelClientSession { .map(AuthManager::unauthorized_recovery); let mut pending_retry = PendingUnauthorizedRetry::default(); loop { - let client_setup = self - .client - .current_client_setup(self.agent_task.as_ref()) - .await?; + let client_setup = self.client.current_client_setup().await?; let transport = ReqwestTransport::new(build_reqwest_client()); let request_auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), @@ -1325,10 +1255,7 @@ impl ModelClientSession { .map(AuthManager::unauthorized_recovery); let mut pending_retry = PendingUnauthorizedRetry::default(); loop { - let client_setup = self - .client - .current_client_setup(self.agent_task.as_ref()) - .await?; + let client_setup = self.client.current_client_setup().await?; let request_auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), client_setup.api_auth.as_ref(), diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 357a892264ef..f4575b26a0b2 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -1,5 +1,3 @@ -use std::sync::Arc; - use super::AuthRequestTelemetryContext; use super::ModelClient; use super::PendingUnauthorizedRetry; @@ -9,36 +7,17 @@ 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::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_app_server_protocol::AuthMode; -use codex_login::AuthManager; -use codex_login::CodexAuth; use codex_model_provider::BearerAuthProvider; -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 core_test_support::responses; -use ed25519_dalek::Signature; -use ed25519_dalek::Verifier as _; -use futures::StreamExt; use pretty_assertions::assert_eq; -use serde::Deserialize; use serde_json::json; -use tempfile::TempDir; fn test_model_client(session_source: SessionSource) -> ModelClient { let provider = create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses); @@ -100,118 +79,6 @@ 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 auth_manager = - AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()); - let agent_identity_manager = Arc::new(AgentIdentityManager::new_for_tests( - Arc::clone(&auth_manager), - /*feature_enabled*/ true, - "https://chatgpt.com/backend-api/".to_string(), - SessionSource::Cli, - )); - let stored_identity = agent_identity_manager - .seed_generated_identity_for_tests("agent-123") - .await - .expect("seed test identity"); - let agent_task = RegisteredAgentTask { - agent_runtime_id: stored_identity.agent_runtime_id.clone(), - task_id: "task-123".to_string(), - registered_at: "2026-03-23T12:00:00Z".to_string(), - }; - let client = ModelClient::new( - Some(auth_manager), - ThreadId::new(), - /*installation_id*/ "11111111-1111-4111-8111-111111111111".to_string(), - provider, - SessionSource::Cli, - /*model_verbosity*/ None, - /*enable_request_compression*/ false, - /*include_timing_metrics*/ false, - /*beta_features_header*/ None, - ); - (codex_home, client, agent_task, stored_identity) -} - -#[derive(Debug, Deserialize)] -struct AgentAssertionEnvelope { - agent_runtime_id: String, - task_id: String, - timestamp: String, - signature: String, -} - -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( @@ -302,130 +169,3 @@ 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, - /*effort*/ None, - ReasoningSummary::Auto, - /*service_tier*/ None, - /*turn_metadata_header*/ 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, - /*effort*/ None, - ReasoningSummary::Auto, - /*service_tier*/ None, - /*turn_metadata_header*/ None, - ) - .await - .expect("bearer prewarm should succeed"); - drop(prewarm_session); - - let mut agent_task_session = client.new_session_with_agent_task(Some(agent_task.clone())); - let mut stream = agent_task_session - .stream( - &prompt, - &model_info, - &session_telemetry, - /*effort*/ None, - ReasoningSummary::Auto, - /*service_tier*/ None, - /*turn_metadata_header*/ 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/mcp_openai_file.rs b/codex-rs/core/src/mcp_openai_file.rs index 324eb4ad1a4e..d6e6d1f9c072 100644 --- a/codex-rs/core/src/mcp_openai_file.rs +++ b/codex-rs/core/src/mcp_openai_file.rs @@ -12,10 +12,8 @@ use crate::session::session::Session; use crate::session::turn_context::TurnContext; -use codex_api::AuthProvider; use codex_api::upload_local_file; use codex_login::CodexAuth; -use codex_model_provider::AuthorizationHeaderAuthProvider; use codex_model_provider::BearerAuthProvider; use serde_json::Value as JsonValue; @@ -42,14 +40,9 @@ pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files( let Some(value) = arguments.get(field_name) else { continue; }; - let Some(uploaded_value) = rewrite_argument_value_for_openai_files( - sess, - turn_context, - auth.as_ref(), - field_name, - value, - ) - .await? + let Some(uploaded_value) = + rewrite_argument_value_for_openai_files(turn_context, auth.as_ref(), field_name, value) + .await? else { continue; }; @@ -64,7 +57,6 @@ pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files( } async fn rewrite_argument_value_for_openai_files( - sess: &Session, turn_context: &TurnContext, auth: Option<&CodexAuth>, field_name: &str, @@ -73,7 +65,6 @@ async fn rewrite_argument_value_for_openai_files( match value { JsonValue::String(path_or_file_ref) => { let rewritten = build_uploaded_local_argument_value( - sess, turn_context, auth, field_name, @@ -90,7 +81,6 @@ async fn rewrite_argument_value_for_openai_files( return Ok(None); }; let rewritten = build_uploaded_local_argument_value( - sess, turn_context, auth, field_name, @@ -107,7 +97,6 @@ async fn rewrite_argument_value_for_openai_files( } async fn build_uploaded_local_argument_value( - sess: &Session, turn_context: &TurnContext, auth: Option<&CodexAuth>, field_name: &str, @@ -120,32 +109,17 @@ async fn build_uploaded_local_argument_value( "ChatGPT auth is required to upload local files for Codex Apps tools".to_string(), ); }; - let upload_auth: Box = if let Some(authorization_header_value) = sess - .authorization_header_for_current_agent_task() - .await - .map_err(|error| format!("failed to build agent assertion authorization: {error}"))? - { - let mut auth_provider = AuthorizationHeaderAuthProvider::new( - Some(authorization_header_value), - /*account_id*/ None, - ); - if auth.is_fedramp_account() { - auth_provider = auth_provider.with_fedramp_routing_header(); - } - Box::new(auth_provider) - } else { - let token_data = auth - .get_token_data() - .map_err(|error| format!("failed to read ChatGPT auth for file upload: {error}"))?; - Box::new(BearerAuthProvider { - token: Some(token_data.access_token), - account_id: token_data.account_id, - is_fedramp_account: auth.is_fedramp_account(), - }) + let token_data = auth + .get_token_data() + .map_err(|error| format!("failed to read ChatGPT auth for file upload: {error}"))?; + let upload_auth = BearerAuthProvider { + token: Some(token_data.access_token), + account_id: token_data.account_id, + is_fedramp_account: auth.is_fedramp_account(), }; let uploaded = upload_local_file( turn_context.config.chatgpt_base_url.trim_end_matches('/'), - upload_auth.as_ref(), + &upload_auth, &resolved_path, ) .await @@ -168,82 +142,12 @@ async fn build_uploaded_local_argument_value( #[cfg(test)] mod tests { use super::*; - use crate::agent_identity::AgentIdentityManager; - use crate::agent_identity::RegisteredAgentTask; use crate::session::tests::make_session_and_context; - use chrono::Utc; - use codex_login::AuthCredentialsStoreMode; - use codex_login::AuthDotJson; - use codex_login::AuthManager; - use codex_login::save_auth; - use codex_login::token_data::IdTokenInfo; - use codex_login::token_data::TokenData; - use codex_protocol::protocol::SessionSource; use codex_utils_absolute_path::AbsolutePathBuf; use pretty_assertions::assert_eq; use std::sync::Arc; use tempfile::tempdir; - const TEST_ID_TOKEN: &str = concat!( - "eyJhbGciOiJub25lIiwidHlwIjoiSldUIn0.", - "eyJodHRwczovL2FwaS5vcGVuYWkuY29tL2F1dGgiOnsiY2hhdGdwdF91c2VyX2lk", - "IjpudWxsLCJjaGF0Z3B0X2FjY291bnRfaWQiOiJhY2NvdW50X2lkIn19.", - "c2ln", - ); - - async fn install_cached_agent_task_auth( - session: &mut Session, - turn_context: &mut TurnContext, - chatgpt_base_url: String, - ) { - let auth_dir = tempdir().expect("temp auth dir"); - let auth_json = AuthDotJson { - auth_mode: Some(codex_app_server_protocol::AuthMode::Chatgpt), - openai_api_key: None, - tokens: Some(TokenData { - id_token: IdTokenInfo { - email: None, - chatgpt_plan_type: None, - chatgpt_user_id: None, - chatgpt_account_id: Some("account_id".to_string()), - chatgpt_account_is_fedramp: false, - raw_jwt: TEST_ID_TOKEN.to_string(), - }, - access_token: "Access Token".to_string(), - refresh_token: "test".to_string(), - account_id: Some("account_id".to_string()), - }), - last_refresh: Some(Utc::now()), - agent_identity: None, - }; - save_auth(auth_dir.path(), &auth_json, AuthCredentialsStoreMode::File) - .expect("save test auth"); - let auth = CodexAuth::from_auth_storage(auth_dir.path(), AuthCredentialsStoreMode::File) - .expect("load test auth") - .expect("test auth"); - let auth_manager = AuthManager::from_auth_for_testing(auth); - let agent_identity_manager = Arc::new(AgentIdentityManager::new_for_tests( - Arc::clone(&auth_manager), - /*feature_enabled*/ true, - chatgpt_base_url, - SessionSource::Exec, - )); - let stored_identity = agent_identity_manager - .seed_generated_identity_for_tests("agent-123") - .await - .expect("seed test identity"); - session.services.auth_manager = Arc::clone(&auth_manager); - session.services.agent_identity_manager = agent_identity_manager; - turn_context.auth_manager = Some(auth_manager); - session - .cache_agent_task_for_tests(RegisteredAgentTask { - agent_runtime_id: stored_identity.agent_runtime_id, - task_id: "task-123".to_string(), - registered_at: "2026-04-15T00:00:00Z".to_string(), - }) - .await; - } - #[tokio::test] async fn openai_file_argument_rewrite_requires_declared_file_params() { let (session, turn_context) = make_session_and_context().await; @@ -308,7 +212,7 @@ mod tests { .mount(&server) .await; - let (session, mut turn_context) = make_session_and_context().await; + let (_, mut turn_context) = make_session_and_context().await; let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); let dir = tempdir().expect("temp dir"); let local_path = dir.path().join("file_report.csv"); @@ -322,7 +226,6 @@ mod tests { turn_context.config = Arc::new(config); let rewritten = build_uploaded_local_argument_value( - &session, &turn_context, Some(&auth), "file", @@ -390,7 +293,7 @@ mod tests { .mount(&server) .await; - let (session, mut turn_context) = make_session_and_context().await; + let (_, mut turn_context) = make_session_and_context().await; let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); let dir = tempdir().expect("temp dir"); let local_path = dir.path().join("file_report.csv"); @@ -403,7 +306,6 @@ mod tests { config.chatgpt_base_url = format!("{}/backend-api", server.uri()); turn_context.config = Arc::new(config); let rewritten = rewrite_argument_value_for_openai_files( - &session, &turn_context, Some(&auth), "file", @@ -503,7 +405,7 @@ mod tests { .mount(&server) .await; - let (session, mut turn_context) = make_session_and_context().await; + let (_, mut turn_context) = make_session_and_context().await; let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); let dir = tempdir().expect("temp dir"); tokio::fs::write(dir.path().join("one.csv"), b"one") @@ -518,7 +420,6 @@ mod tests { config.chatgpt_base_url = format!("{}/backend-api", server.uri()); turn_context.config = Arc::new(config); let rewritten = rewrite_argument_value_for_openai_files( - &session, &turn_context, Some(&auth), "files", @@ -570,88 +471,4 @@ mod tests { assert!(error.contains("failed to upload")); assert!(error.contains("file")); } - - #[tokio::test] - async fn build_uploaded_local_argument_value_uses_agent_assertion_for_cached_task() { - use wiremock::Mock; - use wiremock::MockServer; - use wiremock::ResponseTemplate; - use wiremock::matchers::body_json; - use wiremock::matchers::header_regex; - use wiremock::matchers::method; - use wiremock::matchers::path; - - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/backend-api/files")) - .and(header_regex("authorization", r"^AgentAssertion .+")) - .and(body_json(serde_json::json!({ - "file_name": "file_report.csv", - "file_size": 5, - "use_case": "codex", - }))) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "file_id": "file_123", - "upload_url": format!("{}/upload/file_123", server.uri()), - }))) - .expect(1) - .mount(&server) - .await; - Mock::given(method("PUT")) - .and(path("/upload/file_123")) - .respond_with(ResponseTemplate::new(200)) - .expect(1) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/backend-api/files/file_123/uploaded")) - .and(header_regex("authorization", r"^AgentAssertion .+")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "status": "success", - "download_url": format!("{}/download/file_123", server.uri()), - "file_name": "file_report.csv", - "mime_type": "text/csv", - "file_size_bytes": 5, - }))) - .expect(1) - .mount(&server) - .await; - - let (mut session, mut turn_context) = make_session_and_context().await; - let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); - let dir = tempdir().expect("temp dir"); - let local_path = dir.path().join("file_report.csv"); - tokio::fs::write(&local_path, b"hello") - .await - .expect("write local file"); - turn_context.cwd = AbsolutePathBuf::try_from(dir.path()).expect("absolute path"); - - let mut config = (*turn_context.config).clone(); - config.chatgpt_base_url = format!("{}/backend-api", server.uri()); - turn_context.config = Arc::new(config); - install_cached_agent_task_auth(&mut session, &mut turn_context, server.uri()).await; - - let rewritten = build_uploaded_local_argument_value( - &session, - &turn_context, - Some(&auth), - "file", - /*index*/ None, - "file_report.csv", - ) - .await - .expect("rewrite should upload the local file"); - - assert_eq!( - rewritten, - serde_json::json!({ - "download_url": format!("{}/download/file_123", server.uri()), - "file_id": "file_123", - "mime_type": "text/csv", - "file_name": "file_report.csv", - "uri": "sediment://file_123", - "file_size_bytes": 5, - }) - ); - } } diff --git a/codex-rs/core/src/session/agent_task_lifecycle.rs b/codex-rs/core/src/session/agent_task_lifecycle.rs index 72d0cf0459c6..f5c1c4c6e16d 100644 --- a/codex-rs/core/src/session/agent_task_lifecycle.rs +++ b/codex-rs/core/src/session/agent_task_lifecycle.rs @@ -101,11 +101,6 @@ impl Session { agent_task } - #[cfg(test)] - pub(crate) async fn cache_agent_task_for_tests(&self, agent_task: RegisteredAgentTask) { - self.cache_agent_task(agent_task).await; - } - pub(super) async fn cached_agent_task_for_current_identity( &self, ) -> Option { @@ -139,33 +134,6 @@ impl Session { None } - pub(crate) async fn authorization_header_for_current_agent_task( - &self, - ) -> anyhow::Result> { - let Some(agent_task) = self.cached_agent_task_for_current_identity().await else { - return Ok(None); - }; - - let Some(auth) = self.services.auth_manager.auth().await else { - return Ok(None); - }; - let authorization_header_value = self - .services - .auth_manager - .chatgpt_agent_task_authorization_header_for_auth( - &auth, - agent_task.authorization_target(), - )?; - if authorization_header_value.is_some() { - debug!( - agent_runtime_id = %agent_task.agent_runtime_id, - task_id = %agent_task.task_id, - "using agent assertion authorization for current task request" - ); - } - Ok(authorization_header_value) - } - pub(super) async fn ensure_agent_task_registered( &self, ) -> anyhow::Result> { diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index b96f56abaa51..d8d0aef4db5e 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -624,11 +624,6 @@ impl Session { config.analytics_enabled, ) }); - let agent_identity_manager = Arc::new(AgentIdentityManager::new( - config.as_ref(), - Arc::clone(&auth_manager), - session_configuration.session_source.clone(), - )); let services = SessionServices { // Initialize the MCP connection manager with an uninitialized // instance. It will be replaced with one created via @@ -651,7 +646,11 @@ impl Session { hooks, rollout: Mutex::new(rollout_recorder), user_shell: Arc::new(default_shell), - agent_identity_manager: Arc::clone(&agent_identity_manager), + agent_identity_manager: Arc::new(AgentIdentityManager::new( + config.as_ref(), + Arc::clone(&auth_manager), + session_configuration.session_source.clone(), + )), shell_snapshot_tx, show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 7b27c8080a9f..7f156c97858d 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -345,23 +345,20 @@ pub(crate) async fn run_turn( })) .await; } - let agent_task = match sess.ensure_agent_task_registered().await { - Ok(agent_task) => agent_task, - Err(error) => { - warn!(error = %error, "agent task registration failed"); - sess.send_event( - turn_context.as_ref(), - EventMsg::Error(ErrorEvent { - message: format!( - "Agent task registration failed. Please try again; Codex will attempt to register the task again on the next turn: {error}" - ), - codex_error_info: Some(CodexErrorInfo::Other), - }), - ) - .await; - return None; - } - }; + if let Err(error) = sess.ensure_agent_task_registered().await { + warn!(error = %error, "agent task registration failed"); + sess.send_event( + turn_context.as_ref(), + EventMsg::Error(ErrorEvent { + message: format!( + "Agent task registration failed. Please try again; Codex will attempt to register the task again on the next turn: {error}" + ), + codex_error_info: Some(CodexErrorInfo::Other), + }), + ) + .await; + return None; + } if !skill_items.is_empty() { sess.record_conversation_items(&turn_context, &skill_items) @@ -386,21 +383,8 @@ pub(crate) async fn run_turn( // `ModelClientSession` is turn-scoped and caches WebSocket + sticky routing state, so we reuse // one instance across retries within this turn. - let mut prewarmed_client_session = prewarmed_client_session; - if agent_task.is_some() - && let Some(prewarmed_client_session) = prewarmed_client_session.as_mut() - { - 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() - }; + let mut client_session = + prewarmed_client_session.unwrap_or_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/login/Cargo.toml b/codex-rs/login/Cargo.toml index 16d183602485..d5303ea54cf3 100644 --- a/codex-rs/login/Cargo.toml +++ b/codex-rs/login/Cargo.toml @@ -8,7 +8,6 @@ license.workspace = true workspace = true [dependencies] -anyhow = { workspace = true } async-trait = { workspace = true } base64 = { workspace = true } chrono = { workspace = true, features = ["serde"] } @@ -21,7 +20,6 @@ codex-otel = { workspace = true } codex-protocol = { workspace = true } codex-terminal-detection = { workspace = true } codex-utils-template = { workspace = true } -ed25519-dalek = { workspace = true } once_cell = { workspace = true } os_info = { workspace = true } rand = { workspace = true } @@ -44,6 +42,7 @@ urlencoding = { workspace = true } webbrowser = { workspace = true } [dev-dependencies] +anyhow = { workspace = true } core_test_support = { workspace = true } keyring = { workspace = true } pretty_assertions = { workspace = true } diff --git a/codex-rs/login/src/auth/agent_assertion.rs b/codex-rs/login/src/auth/agent_assertion.rs deleted file mode 100644 index e4a1731f6538..000000000000 --- a/codex-rs/login/src/auth/agent_assertion.rs +++ /dev/null @@ -1,172 +0,0 @@ -use std::collections::BTreeMap; - -use anyhow::Context; -use anyhow::Result; -use base64::Engine as _; -use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; -use base64::engine::general_purpose::URL_SAFE_NO_PAD; -use chrono::SecondsFormat; -use chrono::Utc; -use ed25519_dalek::Signer as _; -use ed25519_dalek::SigningKey; -use ed25519_dalek::pkcs8::DecodePrivateKey; -use serde::Deserialize; -use serde::Serialize; - -use super::storage::AgentIdentityAuthRecord; - -/// Task binding to use when constructing a task-scoped AgentAssertion. -/// -/// The caller owns the task lifecycle. `AuthManager` only uses this target to -/// sign an authorization header with the stored agent identity key material. -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct AgentTaskAuthorizationTarget<'a> { - pub agent_runtime_id: &'a str, - pub task_id: &'a str, -} - -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] -struct AgentAssertionEnvelope { - agent_runtime_id: String, - task_id: String, - timestamp: String, - signature: String, -} - -pub(super) fn authorization_header_for_agent_task( - record: &AgentIdentityAuthRecord, - target: AgentTaskAuthorizationTarget<'_>, -) -> Result { - anyhow::ensure!( - record.agent_runtime_id == target.agent_runtime_id, - "agent task runtime {} does not match stored agent identity {}", - target.agent_runtime_id, - record.agent_runtime_id - ); - - let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); - let envelope = AgentAssertionEnvelope { - agent_runtime_id: target.agent_runtime_id.to_string(), - task_id: target.task_id.to_string(), - timestamp: timestamp.clone(), - signature: sign_agent_assertion_payload(record, target, ×tamp)?, - }; - let serialized_assertion = serialize_agent_assertion(&envelope)?; - Ok(format!("AgentAssertion {serialized_assertion}")) -} - -fn sign_agent_assertion_payload( - record: &AgentIdentityAuthRecord, - target: AgentTaskAuthorizationTarget<'_>, - timestamp: &str, -) -> Result { - let signing_key = signing_key_from_agent_private_key(&record.agent_private_key)?; - let payload = format!("{}:{}:{timestamp}", target.agent_runtime_id, target.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)) -} - -fn signing_key_from_agent_private_key(agent_private_key: &str) -> Result { - let private_key = BASE64_STANDARD - .decode(agent_private_key) - .context("stored agent identity private key is not valid base64")?; - SigningKey::from_pkcs8_der(&private_key) - .context("stored agent identity private key is not valid PKCS#8") -} - -#[cfg(test)] -mod tests { - use ed25519_dalek::Signature; - use ed25519_dalek::Verifier as _; - use ed25519_dalek::pkcs8::EncodePrivateKey; - use pretty_assertions::assert_eq; - - use super::*; - - #[test] - fn authorization_header_for_agent_task_serializes_signed_agent_assertion() { - let record = test_agent_identity_record("agent-123"); - let target = AgentTaskAuthorizationTarget { - agent_runtime_id: "agent-123", - task_id: "task-123", - }; - - let header = authorization_header_for_agent_task(&record, target) - .expect("build agent assertion header"); - 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"); - signing_key_from_agent_private_key(&record.agent_private_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 authorization_header_for_agent_task_rejects_mismatched_runtime() { - let record = test_agent_identity_record("agent-123"); - let target = AgentTaskAuthorizationTarget { - agent_runtime_id: "agent-456", - task_id: "task-123", - }; - - let error = authorization_header_for_agent_task(&record, target) - .expect_err("runtime mismatch should fail"); - - assert_eq!( - error.to_string(), - "agent task runtime agent-456 does not match stored agent identity agent-123" - ); - } - - fn test_agent_identity_record(agent_runtime_id: &str) -> AgentIdentityAuthRecord { - let signing_key = SigningKey::from_bytes(&[7u8; 32]); - let private_key = signing_key - .to_pkcs8_der() - .expect("encode test key material"); - AgentIdentityAuthRecord { - workspace_id: "account-123".to_string(), - chatgpt_user_id: Some("user-123".to_string()), - agent_runtime_id: agent_runtime_id.to_string(), - agent_private_key: BASE64_STANDARD.encode(private_key.as_bytes()), - registered_at: "2026-03-23T12:00:00Z".to_string(), - } - } -} diff --git a/codex-rs/login/src/auth/mod.rs b/codex-rs/login/src/auth/mod.rs index 62ab467d08fa..b927f9a77520 100644 --- a/codex-rs/login/src/auth/mod.rs +++ b/codex-rs/login/src/auth/mod.rs @@ -1,4 +1,3 @@ -mod agent_assertion; pub mod default_client; pub mod error; mod storage; @@ -8,7 +7,6 @@ mod external_bearer; mod manager; mod revoke; -pub use agent_assertion::AgentTaskAuthorizationTarget; pub use error::RefreshTokenFailedError; pub use error::RefreshTokenFailedReason; pub use manager::*; diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 046d878e87b4..d819b0946d33 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -18,7 +18,6 @@ pub use server::ShutdownHandle; pub use server::run_login_server; pub use auth::AgentIdentityAuthRecord; -pub use auth::AgentTaskAuthorizationTarget; pub use auth::AuthConfig; pub use auth::AuthDotJson; pub use auth::AuthManager; diff --git a/codex-rs/model-provider/src/bearer_auth_provider.rs b/codex-rs/model-provider/src/bearer_auth_provider.rs index 970574c752c8..5a24ca6f78da 100644 --- a/codex-rs/model-provider/src/bearer_auth_provider.rs +++ b/codex-rs/model-provider/src/bearer_auth_provider.rs @@ -38,55 +38,6 @@ impl AuthProvider for BearerAuthProvider { } } -/// Auth provider for callers that already resolved the complete Authorization header value. -#[derive(Clone, Default)] -pub struct AuthorizationHeaderAuthProvider { - pub authorization_header_value: Option, - pub account_id: Option, - pub is_fedramp_account: bool, -} - -impl AuthorizationHeaderAuthProvider { - pub fn new(authorization_header_value: Option, account_id: Option) -> Self { - Self { - authorization_header_value, - account_id, - is_fedramp_account: false, - } - } - - pub fn for_test(authorization_header_value: Option<&str>, account_id: Option<&str>) -> Self { - Self { - authorization_header_value: authorization_header_value.map(str::to_string), - account_id: account_id.map(str::to_string), - is_fedramp_account: false, - } - } - - pub fn with_fedramp_routing_header(mut self) -> Self { - self.is_fedramp_account = true; - self - } -} - -impl AuthProvider for AuthorizationHeaderAuthProvider { - fn add_auth_headers(&self, headers: &mut HeaderMap) { - if let Some(authorization_header_value) = self.authorization_header_value.as_ref() - && let Ok(header) = HeaderValue::from_str(authorization_header_value) - { - let _ = headers.insert(http::header::AUTHORIZATION, header); - } - if let Some(account_id) = self.account_id.as_ref() - && let Ok(header) = HeaderValue::from_str(account_id) - { - let _ = headers.insert("ChatGPT-Account-ID", header); - } - if self.is_fedramp_account { - let _ = headers.insert("X-OpenAI-Fedramp", HeaderValue::from_static("true")); - } - } -} - #[cfg(test)] mod tests { use super::*; @@ -148,54 +99,4 @@ mod tests { Some("true") ); } - - #[test] - fn authorization_header_auth_provider_supports_non_bearer_authorization_headers() { - let auth = AuthorizationHeaderAuthProvider::for_test( - Some("AgentAssertion opaque-token"), - Some("workspace-123"), - ); - let mut headers = HeaderMap::new(); - - auth.add_auth_headers(&mut headers); - - assert_eq!( - headers - .get(http::header::AUTHORIZATION) - .and_then(|value| value.to_str().ok()), - Some("AgentAssertion opaque-token") - ); - assert_eq!( - headers - .get("ChatGPT-Account-ID") - .and_then(|value| value.to_str().ok()), - Some("workspace-123") - ); - assert_eq!( - codex_api::auth_header_telemetry(&auth), - codex_api::AuthHeaderTelemetry { - attached: true, - name: Some("authorization"), - } - ); - } - - #[test] - fn authorization_header_auth_provider_adds_fedramp_routing_header_when_enabled() { - let auth = AuthorizationHeaderAuthProvider::for_test( - Some("AgentAssertion opaque-token"), - Some("workspace-123"), - ) - .with_fedramp_routing_header(); - let mut headers = HeaderMap::new(); - - auth.add_auth_headers(&mut headers); - - assert_eq!( - headers - .get("X-OpenAI-Fedramp") - .and_then(|value| value.to_str().ok()), - Some("true") - ); - } } diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index 1874f37e31a2..f240c47db0b4 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -2,7 +2,6 @@ mod auth; mod bearer_auth_provider; mod provider; -pub use bearer_auth_provider::AuthorizationHeaderAuthProvider; pub use bearer_auth_provider::BearerAuthProvider; pub use bearer_auth_provider::BearerAuthProvider as CoreAuthProvider; pub use provider::ModelProvider; From 292a4af9e28fe6615910a37c88264a791c2feb80 Mon Sep 17 00:00:00 2001 From: Edward Frazer Date: Tue, 21 Apr 2026 09:30:52 -0700 Subject: [PATCH 3/9] Revert "Persist and prewarm agent tasks per thread (#17978)" This reverts commit e5b52a3caa7a050c4df81570f5698d74793ef942. --- .../src/protocol/thread_history.rs | 4 +- .../app-server/tests/suite/v2/initialize.rs | 4 +- codex-rs/core/src/agent/control.rs | 1 - codex-rs/core/src/agent_identity.rs | 13 +- .../src/agent_identity/task_registration.rs | 42 +- .../core/src/session/agent_task_lifecycle.rs | 186 -------- codex-rs/core/src/session/mod.rs | 91 +++- .../src/session/rollout_reconstruction.rs | 5 +- codex-rs/core/src/session/session.rs | 4 +- codex-rs/core/src/session/tests.rs | 402 +----------------- codex-rs/core/src/state/session.rs | 8 +- codex-rs/core/src/state/session_tests.rs | 14 +- codex-rs/core/tests/suite/otel.rs | 2 - codex-rs/protocol/src/protocol.rs | 21 - codex-rs/rollout/src/list.rs | 4 - codex-rs/rollout/src/metadata.rs | 6 +- codex-rs/rollout/src/policy.rs | 7 +- codex-rs/rollout/src/recorder.rs | 4 - codex-rs/state/src/extract.rs | 8 +- codex-rs/state/src/runtime/threads.rs | 6 +- 20 files changed, 148 insertions(+), 684 deletions(-) delete mode 100644 codex-rs/core/src/session/agent_task_lifecycle.rs diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index bd3911d0d25b..f7973a45de81 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -232,9 +232,7 @@ impl ThreadHistoryBuilder { RolloutItem::EventMsg(event) => self.handle_event(event), RolloutItem::Compacted(payload) => self.handle_compacted(payload), RolloutItem::ResponseItem(item) => self.handle_response_item(item), - RolloutItem::TurnContext(_) - | RolloutItem::SessionMeta(_) - | RolloutItem::SessionState(_) => {} + RolloutItem::TurnContext(_) | RolloutItem::SessionMeta(_) => {} } } diff --git a/codex-rs/app-server/tests/suite/v2/initialize.rs b/codex-rs/app-server/tests/suite/v2/initialize.rs index 4fe5f2fbe751..165160468f78 100644 --- a/codex-rs/app-server/tests/suite/v2/initialize.rs +++ b/codex-rs/app-server/tests/suite/v2/initialize.rs @@ -24,9 +24,7 @@ use std::time::Duration; use tempfile::TempDir; use tokio::time::timeout; -// This covers debug app-server process startup before the first JSON-RPC response, -// not expected steady-state request latency. -const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(20); +const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); #[tokio::test] async fn initialize_uses_client_info_name_as_originator() -> Result<()> { diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 3c8b2cf3e57a..1c30d613e008 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -116,7 +116,6 @@ fn keep_forked_rollout_item(item: &RolloutItem) -> bool { | ResponseItem::Compaction { .. } | ResponseItem::Other, ) => false, - RolloutItem::SessionState(_) => false, // A forked child gets its own runtime config, including spawned-agent // instructions, so it must establish a fresh context diff baseline. RolloutItem::TurnContext(_) => false, diff --git a/codex-rs/core/src/agent_identity.rs b/codex-rs/core/src/agent_identity.rs index e7a087461eb4..8edcb7bb10d1 100644 --- a/codex-rs/core/src/agent_identity.rs +++ b/codex-rs/core/src/agent_identity.rs @@ -157,16 +157,14 @@ impl AgentIdentityManager { Ok(stored_identity) } - pub(crate) async fn task_matches_current_identity(&self, task: &RegisteredAgentTask) -> bool { + pub(crate) async fn task_matches_current_binding(&self, task: &RegisteredAgentTask) -> bool { if !self.feature_enabled { return false; } - self.current_stored_identity() + self.current_auth_binding() .await - .is_some_and(|stored_identity| { - stored_identity.agent_runtime_id == task.agent_runtime_id - }) + .is_some_and(|(_, binding)| task.matches_binding(&binding)) } async fn current_auth_binding(&self) -> Option<(CodexAuth, AgentIdentityBinding)> { @@ -183,11 +181,6 @@ impl AgentIdentityManager { binding.map(|binding| (auth, binding)) } - async fn current_stored_identity(&self) -> Option { - let (auth, binding) = self.current_auth_binding().await?; - self.load_stored_identity(&auth, &binding).ok().flatten() - } - async fn register_agent_identity( &self, binding: &AgentIdentityBinding, diff --git a/codex-rs/core/src/agent_identity/task_registration.rs b/codex-rs/core/src/agent_identity/task_registration.rs index 53bb272cbc00..4fc5d51282d3 100644 --- a/codex-rs/core/src/agent_identity/task_registration.rs +++ b/codex-rs/core/src/agent_identity/task_registration.rs @@ -2,7 +2,6 @@ use std::time::Duration; use anyhow::Context; use anyhow::Result; -use codex_protocol::protocol::SessionAgentTask; use crypto_box::SecretKey as Curve25519SecretKey; use ed25519_dalek::Signer as _; use serde::Deserialize; @@ -17,6 +16,9 @@ const AGENT_TASK_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15); #[derive(Clone, Debug, PartialEq, Eq)] pub(crate) struct RegisteredAgentTask { + pub(crate) binding_id: String, + pub(crate) chatgpt_account_id: String, + pub(crate) chatgpt_user_id: Option, pub(crate) agent_runtime_id: String, pub(crate) task_id: String, pub(crate) registered_at: String, @@ -80,6 +82,9 @@ impl AgentIdentityManager { .await .with_context(|| format!("failed to parse agent task response from {url}"))?; let registered_task = RegisteredAgentTask { + binding_id: stored_identity.binding_id.clone(), + chatgpt_account_id: stored_identity.chatgpt_account_id.clone(), + chatgpt_user_id: stored_identity.chatgpt_user_id.clone(), agent_runtime_id: stored_identity.agent_runtime_id.clone(), task_id: decrypt_task_id_response( &stored_identity, @@ -102,20 +107,18 @@ impl AgentIdentityManager { } impl RegisteredAgentTask { - pub(crate) fn to_session_agent_task(&self) -> SessionAgentTask { - SessionAgentTask { - agent_runtime_id: self.agent_runtime_id.clone(), - task_id: self.task_id.clone(), - registered_at: self.registered_at.clone(), - } + pub(super) fn matches_binding(&self, binding: &AgentIdentityBinding) -> bool { + binding.matches_parts( + &self.binding_id, + &self.chatgpt_account_id, + self.chatgpt_user_id.as_deref(), + ) } - pub(crate) fn from_session_agent_task(task: SessionAgentTask) -> Self { - Self { - agent_runtime_id: task.agent_runtime_id, - task_id: task.task_id, - registered_at: task.registered_at, - } + pub(crate) fn has_same_binding(&self, other: &Self) -> bool { + self.binding_id == other.binding_id + && self.chatgpt_account_id == other.chatgpt_account_id + && self.chatgpt_user_id == other.chatgpt_user_id } } @@ -239,6 +242,9 @@ mod tests { assert_eq!( task, RegisteredAgentTask { + binding_id: "chatgpt-account-account-123".to_string(), + chatgpt_account_id: "account-123".to_string(), + chatgpt_user_id: Some("user-123".to_string()), agent_runtime_id: "agent-123".to_string(), task_id: "task_123".to_string(), registered_at: task.registered_at.clone(), @@ -325,6 +331,9 @@ mod tests { assert_eq!( task, RegisteredAgentTask { + binding_id: "chatgpt-account-account-123".to_string(), + chatgpt_account_id: "account-123".to_string(), + chatgpt_user_id: Some("user-123".to_string()), agent_runtime_id: "agent-123".to_string(), task_id: "task_123".to_string(), registered_at: task.registered_at.clone(), @@ -333,7 +342,7 @@ mod tests { } #[tokio::test] - async fn task_matches_current_identity_rejects_stale_registered_identity() { + async fn task_matches_current_binding_rejects_stale_auth_binding() { let auth_manager = AuthManager::from_auth_for_testing(make_chatgpt_auth("account-456", Some("user-456"))); let manager = AgentIdentityManager::new_for_tests( @@ -343,12 +352,15 @@ mod tests { SessionSource::Cli, ); let task = RegisteredAgentTask { + binding_id: "chatgpt-account-account-123".to_string(), + chatgpt_account_id: "account-123".to_string(), + chatgpt_user_id: Some("user-123".to_string()), agent_runtime_id: "agent-123".to_string(), task_id: "task_123".to_string(), registered_at: "2026-03-23T12:00:00Z".to_string(), }; - assert!(!manager.task_matches_current_identity(&task).await); + assert!(!manager.task_matches_current_binding(&task).await); } async fn mount_human_biscuit( diff --git a/codex-rs/core/src/session/agent_task_lifecycle.rs b/codex-rs/core/src/session/agent_task_lifecycle.rs deleted file mode 100644 index f5c1c4c6e16d..000000000000 --- a/codex-rs/core/src/session/agent_task_lifecycle.rs +++ /dev/null @@ -1,186 +0,0 @@ -use crate::agent_identity::RegisteredAgentTask; -use crate::session::session::Session; -use codex_protocol::protocol::RolloutItem; -use codex_protocol::protocol::SessionAgentTask; -use codex_protocol::protocol::SessionStateUpdate; -use tracing::debug; -use tracing::info; -use tracing::warn; - -impl Session { - pub(super) async fn maybe_prewarm_agent_task_registration(&self) { - // Startup task registration is best-effort: regular turns already retry on demand, and - // a prewarm failure should not shut down the session or block unrelated work. - if let Err(error) = self.ensure_agent_task_registered().await { - warn!( - error = %error, - "startup agent task prewarm failed; regular turns will retry registration" - ); - } - } - - fn latest_persisted_agent_task( - rollout_items: &[RolloutItem], - ) -> Option> { - rollout_items.iter().rev().find_map(|item| match item { - RolloutItem::SessionState(update) => Some(update.agent_task.clone()), - _ => None, - }) - } - - pub(super) async fn restore_persisted_agent_task(&self, rollout_items: &[RolloutItem]) { - let Some(agent_task_update) = Self::latest_persisted_agent_task(rollout_items) else { - return; - }; - - match agent_task_update { - Some(agent_task) => { - let registered_task = - RegisteredAgentTask::from_session_agent_task(agent_task.clone()); - if self - .services - .agent_identity_manager - .task_matches_current_identity(®istered_task) - .await - { - let mut state = self.state.lock().await; - state.set_agent_task(agent_task); - } else { - debug!( - agent_runtime_id = %registered_task.agent_runtime_id, - task_id = %registered_task.task_id, - "discarding persisted agent task because it does not match the registered agent identity" - ); - let mut state = self.state.lock().await; - state.clear_agent_task(); - } - } - None => { - let mut state = self.state.lock().await; - state.clear_agent_task(); - } - } - } - - async fn persist_agent_task_update(&self, agent_task: Option<&RegisteredAgentTask>) { - self.persist_rollout_items(&[RolloutItem::SessionState(SessionStateUpdate { - agent_task: agent_task.map(RegisteredAgentTask::to_session_agent_task), - })]) - .await; - } - - async fn clear_cached_agent_task(&self, agent_task: &RegisteredAgentTask) { - let cleared = { - let mut state = self.state.lock().await; - if state.agent_task().as_ref() == Some(&agent_task.to_session_agent_task()) { - state.clear_agent_task(); - true - } else { - false - } - }; - if cleared { - self.persist_agent_task_update(/*agent_task*/ None).await; - } - } - - async fn cache_agent_task(&self, agent_task: RegisteredAgentTask) -> RegisteredAgentTask { - let session_agent_task = agent_task.to_session_agent_task(); - let changed = { - let mut state = self.state.lock().await; - if state.agent_task().as_ref() == Some(&session_agent_task) { - false - } else { - state.set_agent_task(session_agent_task); - true - } - }; - if changed { - self.persist_agent_task_update(Some(&agent_task)).await; - } - agent_task - } - - pub(super) async fn cached_agent_task_for_current_identity( - &self, - ) -> Option { - let agent_task = { - let state = self.state.lock().await; - state - .agent_task() - .map(RegisteredAgentTask::from_session_agent_task) - }?; - - if self - .services - .agent_identity_manager - .task_matches_current_identity(&agent_task) - .await - { - debug!( - agent_runtime_id = %agent_task.agent_runtime_id, - task_id = %agent_task.task_id, - "reusing cached agent task" - ); - return Some(agent_task); - } - - debug!( - agent_runtime_id = %agent_task.agent_runtime_id, - task_id = %agent_task.task_id, - "discarding cached agent task because the registered agent identity changed" - ); - self.clear_cached_agent_task(&agent_task).await; - None - } - - pub(super) async fn ensure_agent_task_registered( - &self, - ) -> anyhow::Result> { - if let Some(agent_task) = self.cached_agent_task_for_current_identity().await { - return Ok(Some(agent_task)); - } - - let _guard = self - .agent_task_registration_lock - .acquire() - .await - .map_err(|_| anyhow::anyhow!("agent task registration semaphore closed"))?; - if let Some(agent_task) = self.cached_agent_task_for_current_identity().await { - return Ok(Some(agent_task)); - } - - for _ in 0..2 { - let Some(agent_task) = self.services.agent_identity_manager.register_task().await? - else { - return Ok(None); - }; - - if !self - .services - .agent_identity_manager - .task_matches_current_identity(&agent_task) - .await - { - debug!( - agent_runtime_id = %agent_task.agent_runtime_id, - task_id = %agent_task.task_id, - "discarding newly registered agent task because the registered agent identity changed" - ); - continue; - } - - let agent_task = self.cache_agent_task(agent_task).await; - - info!( - thread_id = %self.conversation_id, - agent_runtime_id = %agent_task.agent_runtime_id, - task_id = %agent_task.task_id, - "registered agent task for thread" - ); - return Ok(Some(agent_task)); - } - - Ok(None) - } -} diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 87ecd14c9fe8..77b868f19f16 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -173,7 +173,6 @@ use codex_protocol::error::Result as CodexResult; #[cfg(test)] use codex_protocol::exec_output::StreamOutput; -mod agent_task_lifecycle; mod handlers; mod mcp; mod review; @@ -1000,10 +999,7 @@ impl Session { .ensure_registered_identity() .await { - Ok(Some(_)) => { - sess.maybe_prewarm_agent_task_registration().await; - return; - } + Ok(Some(_)) => return, Ok(None) => { drop(sess); if auth_state_rx.changed().await.is_err() { @@ -1034,6 +1030,90 @@ impl Session { .await; } + async fn cached_agent_task_for_current_binding(&self) -> Option { + let agent_task = { + let state = self.state.lock().await; + state.agent_task() + }?; + + if self + .services + .agent_identity_manager + .task_matches_current_binding(&agent_task) + .await + { + debug!( + agent_runtime_id = %agent_task.agent_runtime_id, + task_id = %agent_task.task_id, + "reusing cached agent task" + ); + return Some(agent_task); + } + + debug!( + agent_runtime_id = %agent_task.agent_runtime_id, + task_id = %agent_task.task_id, + "discarding cached agent task because auth binding changed" + ); + let mut state = self.state.lock().await; + if state.agent_task().as_ref() == Some(&agent_task) { + state.clear_agent_task(); + } + None + } + + async fn ensure_agent_task_registered(&self) -> anyhow::Result> { + if let Some(agent_task) = self.cached_agent_task_for_current_binding().await { + return Ok(Some(agent_task)); + } + + for _ in 0..2 { + let Some(agent_task) = self.services.agent_identity_manager.register_task().await? + else { + return Ok(None); + }; + + if !self + .services + .agent_identity_manager + .task_matches_current_binding(&agent_task) + .await + { + debug!( + agent_runtime_id = %agent_task.agent_runtime_id, + task_id = %agent_task.task_id, + "discarding newly registered agent task because auth binding changed" + ); + continue; + } + + { + let mut state = self.state.lock().await; + if let Some(existing_agent_task) = state.agent_task() { + if existing_agent_task.has_same_binding(&agent_task) { + return Ok(Some(existing_agent_task)); + } + debug!( + agent_runtime_id = %existing_agent_task.agent_runtime_id, + task_id = %existing_agent_task.task_id, + "replacing cached agent task because auth binding changed" + ); + } + state.set_agent_task(agent_task.clone()); + } + + info!( + thread_id = %self.conversation_id, + agent_runtime_id = %agent_task.agent_runtime_id, + task_id = %agent_task.task_id, + "registered agent task for thread" + ); + return Ok(Some(agent_task)); + } + + Ok(None) + } + pub(crate) fn get_tx_event(&self) -> Sender { self.tx_event.clone() } @@ -1181,7 +1261,6 @@ impl Session { } InitialHistory::Resumed(resumed_history) => { let rollout_items = resumed_history.history; - self.restore_persisted_agent_task(&rollout_items).await; let previous_turn_settings = self .apply_rollout_reconstruction(&turn_context, &rollout_items) .await; diff --git a/codex-rs/core/src/session/rollout_reconstruction.rs b/codex-rs/core/src/session/rollout_reconstruction.rs index 3e407c4cd797..a4c042af0c83 100644 --- a/codex-rs/core/src/session/rollout_reconstruction.rs +++ b/codex-rs/core/src/session/rollout_reconstruction.rs @@ -207,9 +207,7 @@ impl Session { active_segment.get_or_insert_with(ActiveReplaySegment::default); active_segment.counts_as_user_turn |= is_user_turn_boundary(response_item); } - RolloutItem::EventMsg(_) - | RolloutItem::SessionMeta(_) - | RolloutItem::SessionState(_) => {} + RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) => {} } if base_replacement_history.is_some() @@ -277,7 +275,6 @@ impl Session { history.drop_last_n_user_turns(rollback.num_turns); } RolloutItem::EventMsg(_) - | RolloutItem::SessionState(_) | RolloutItem::TurnContext(_) | RolloutItem::SessionMeta(_) => {} } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index d8d0aef4db5e..015feea0c6f8 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -26,7 +26,6 @@ pub(crate) struct Session { pub(crate) services: SessionServices, pub(super) js_repl: Arc, pub(super) next_internal_sub_id: AtomicU64, - pub(super) agent_task_registration_lock: Semaphore, } #[derive(Clone)] @@ -713,7 +712,6 @@ impl Session { services, js_repl, next_internal_sub_id: AtomicU64::new(0), - agent_task_registration_lock: Semaphore::new(/*permits*/ 1), }); if let Some(network_policy_decider_session) = network_policy_decider_session { let mut guard = network_policy_decider_session.write().await; @@ -754,6 +752,7 @@ impl Session { // Start the watcher after SessionConfigured so it cannot emit earlier events. sess.start_skills_watcher_listener(); + sess.start_agent_identity_registration(); let mut required_mcp_servers: Vec = mcp_servers .iter() .filter(|(_, server)| server.enabled && server.required) @@ -840,7 +839,6 @@ impl Session { // record_initial_history can emit events. We record only after the SessionConfiguredEvent is emitted. sess.record_initial_history(initial_history).await; - sess.start_agent_identity_registration(); { let mut state = sess.state.lock().await; state.set_pending_session_start_source(Some(session_start_source)); diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 6669115113d1..2ff4be2ed8a1 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -1,6 +1,4 @@ use super::*; -use crate::agent_identity::RegisteredAgentTask; -use crate::agent_identity::StoredAgentIdentity; use crate::config::ConfigBuilder; use crate::config::test_config; use crate::config_loader::ConfigLayerStack; @@ -19,20 +17,10 @@ use crate::shell::default_user_shell; use crate::skills::SkillRenderSideEffects; use crate::skills::render::SkillMetadataBudget; use crate::tools::format_exec_output_str; -use base64::Engine as _; -use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; -use base64::engine::general_purpose::URL_SAFE_NO_PAD; -use chrono::SecondsFormat; -use chrono::Utc; + use codex_features::Feature; use codex_features::Features; -use codex_login::AgentIdentityAuthRecord; -use codex_login::AuthCredentialsStoreMode; -use codex_login::AuthDotJson; use codex_login::CodexAuth; -use codex_login::save_auth; -use codex_login::token_data::IdTokenInfo; -use codex_login::token_data::TokenData; use codex_model_provider_info::ModelProviderInfo; use codex_models_manager::bundled_models_response; use codex_models_manager::model_info; @@ -131,9 +119,6 @@ use core_test_support::test_codex::test_codex; use core_test_support::test_path_buf; use core_test_support::tracing::install_test_tracing; use core_test_support::wait_for_event; -use crypto_box::SecretKey as Curve25519SecretKey; -use ed25519_dalek::SigningKey; -use ed25519_dalek::pkcs8::EncodePrivateKey; use opentelemetry::trace::TraceContextExt; use opentelemetry::trace::TraceId; use opentelemetry_sdk::metrics::InMemoryMetricExporter; @@ -141,20 +126,12 @@ use opentelemetry_sdk::metrics::data::AggregatedMetrics; use opentelemetry_sdk::metrics::data::Metric; use opentelemetry_sdk::metrics::data::MetricData; use opentelemetry_sdk::metrics::data::ResourceMetrics; -use sha2::Digest as _; -use sha2::Sha512; use std::path::Path; use std::time::Duration; use tokio::sync::Semaphore; use tokio::time::sleep; use tokio::time::timeout; use tracing_opentelemetry::OpenTelemetrySpanExt; -use wiremock::Mock; -use wiremock::MockServer; -use wiremock::ResponseTemplate; -use wiremock::matchers::header; -use wiremock::matchers::method; -use wiremock::matchers::path; use codex_protocol::mcp::CallToolResult as McpCallToolResult; use pretty_assertions::assert_eq; @@ -1293,120 +1270,6 @@ async fn record_initial_history_reconstructs_resumed_transcript() { assert_eq!(expected, history.raw_items()); } -#[tokio::test] -async fn record_initial_history_restores_latest_persisted_agent_task() { - let auth = make_chatgpt_auth("account-123", Some("user-123")); - seed_stored_identity(&auth, "agent-123", "account-123"); - let (session, _turn_context, _rx_event) = make_agent_identity_session_and_context_with_rx( - auth, - "https://chatgpt.com/backend-api".to_string(), - ) - .await; - let expected = RegisteredAgentTask { - agent_runtime_id: "agent-123".to_string(), - task_id: "task-123".to_string(), - registered_at: "2026-03-23T12:00:00Z".to_string(), - }; - let rollout_items = vec![ - RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate { - agent_task: Some(expected.to_session_agent_task()), - }), - RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate { - agent_task: None, - }), - RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate { - agent_task: Some(expected.to_session_agent_task()), - }), - ]; - - session - .record_initial_history(InitialHistory::Resumed(ResumedHistory { - conversation_id: ThreadId::default(), - history: rollout_items, - rollout_path: PathBuf::from("/tmp/resume.jsonl"), - })) - .await; - - assert_eq!( - session.state.lock().await.agent_task(), - Some(expected.to_session_agent_task()) - ); -} - -#[tokio::test] -async fn record_initial_history_discards_persisted_agent_task_for_different_identity() { - let auth = make_chatgpt_auth("account-123", Some("user-123")); - seed_stored_identity(&auth, "agent-123", "account-123"); - let (session, _turn_context, _rx_event) = make_agent_identity_session_and_context_with_rx( - auth, - "https://chatgpt.com/backend-api".to_string(), - ) - .await; - let rollout_items = vec![RolloutItem::SessionState( - codex_protocol::protocol::SessionStateUpdate { - agent_task: Some( - RegisteredAgentTask { - agent_runtime_id: "agent-other".to_string(), - task_id: "task-other".to_string(), - registered_at: "2026-03-23T12:00:00Z".to_string(), - } - .to_session_agent_task(), - ), - }, - )]; - - session - .record_initial_history(InitialHistory::Resumed(ResumedHistory { - conversation_id: ThreadId::default(), - history: rollout_items, - rollout_path: PathBuf::from("/tmp/resume.jsonl"), - })) - .await; - - assert_eq!(session.state.lock().await.agent_task(), None); -} - -#[tokio::test] -async fn record_initial_history_honors_cleared_persisted_agent_task() { - let (session, _turn_context) = make_session_and_context().await; - { - let mut state = session.state.lock().await; - state.set_agent_task( - RegisteredAgentTask { - agent_runtime_id: "agent-fresh".to_string(), - task_id: "task-fresh".to_string(), - registered_at: "2026-03-23T12:01:00Z".to_string(), - } - .to_session_agent_task(), - ); - } - let rollout_items = vec![ - RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate { - agent_task: Some( - RegisteredAgentTask { - agent_runtime_id: "agent-123".to_string(), - task_id: "task-123".to_string(), - registered_at: "2026-03-23T12:00:00Z".to_string(), - } - .to_session_agent_task(), - ), - }), - RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate { - agent_task: None, - }), - ]; - - session - .record_initial_history(InitialHistory::Resumed(ResumedHistory { - conversation_id: ThreadId::default(), - history: rollout_items, - rollout_path: PathBuf::from("/tmp/resume.jsonl"), - })) - .await; - - assert_eq!(session.state.lock().await.agent_task(), None); -} - #[tokio::test] async fn record_initial_history_new_defers_initial_context_until_first_turn() { let (session, _turn_context) = make_session_and_context().await; @@ -3308,7 +3171,6 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { services, js_repl, next_internal_sub_id: AtomicU64::new(0), - agent_task_registration_lock: Semaphore::new(/*permits*/ 1), }; (session, turn_context) @@ -4191,25 +4053,19 @@ async fn shutdown_and_wait_shuts_down_tracked_ephemeral_guardian_review() { .expect("ephemeral guardian review should receive a shutdown op"); } -async fn make_session_and_context_with_auth_and_config_and_rx( - auth: CodexAuth, +pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( dynamic_tools: Vec, - configure_config: F, ) -> ( Arc, Arc, async_channel::Receiver, -) -where - F: FnOnce(&mut Config), -{ +) { let (tx_event, rx_event) = async_channel::unbounded(); let codex_home = tempfile::tempdir().expect("create temp dir"); - let mut config = build_test_config(codex_home.path()).await; - configure_config(&mut config); + let config = build_test_config(codex_home.path()).await; let config = Arc::new(config); let conversation_id = ThreadId::default(); - let auth_manager = AuthManager::from_auth_for_testing(auth); + let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("Test API Key")); let models_manager = Arc::new(ModelsManager::new( config.codex_home.to_path_buf(), auth_manager.clone(), @@ -4411,45 +4267,11 @@ where services, js_repl, next_internal_sub_id: AtomicU64::new(0), - agent_task_registration_lock: Semaphore::new(/*permits*/ 1), }); (session, turn_context, rx_event) } -pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( - dynamic_tools: Vec, -) -> ( - Arc, - Arc, - async_channel::Receiver, -) { - make_session_and_context_with_auth_and_config_and_rx( - CodexAuth::from_api_key("Test API Key"), - dynamic_tools, - |_config| {}, - ) - .await -} - -async fn make_agent_identity_session_and_context_with_rx( - auth: CodexAuth, - chatgpt_base_url: String, -) -> ( - Arc, - Arc, - async_channel::Receiver, -) { - make_session_and_context_with_auth_and_config_and_rx(auth, Vec::new(), move |config| { - config.chatgpt_base_url = chatgpt_base_url; - config - .features - .enable(Feature::UseAgentIdentity) - .expect("test config should allow use_agent_identity"); - }) - .await -} - // Like make_session_and_context, but returns Arc and the event receiver // so tests can assert on emitted events. pub(crate) async fn make_session_and_context_with_rx() -> ( @@ -4489,220 +4311,6 @@ async fn fail_agent_identity_registration_emits_error_without_shutdown() { assert!(rx_event.try_recv().is_err()); } -#[tokio::test] -async fn startup_agent_task_prewarm_caches_registered_task() { - let server = MockServer::start().await; - let chatgpt_base_url = server.uri(); - let auth = make_chatgpt_auth("account-123", Some("user-123")); - let stored_identity = seed_stored_identity(&auth, "agent-123", "account-123"); - let encrypted_task_id = - encrypt_task_id_for_identity(&stored_identity, "task_123").expect("task ciphertext"); - mount_human_biscuit(&server, &chatgpt_base_url, "agent-123").await; - Mock::given(method("POST")) - .and(path("/v1/agent/agent-123/task/register")) - .and(header("x-openai-authorization", "human-biscuit")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "encrypted_task_id": encrypted_task_id, - }))) - .expect(1) - .mount(&server) - .await; - - let (session, _turn_context, rx_event) = - make_agent_identity_session_and_context_with_rx(auth, chatgpt_base_url).await; - - session.maybe_prewarm_agent_task_registration().await; - - let cached_task = session - .state - .lock() - .await - .agent_task() - .expect("task should be cached"); - assert_eq!(cached_task.agent_runtime_id, "agent-123"); - assert_eq!(cached_task.task_id, "task_123"); - assert!(rx_event.try_recv().is_err()); -} - -#[tokio::test] -async fn startup_agent_task_prewarm_failure_does_not_emit_error() { - let server = MockServer::start().await; - let chatgpt_base_url = server.uri(); - let auth = make_chatgpt_auth("account-123", Some("user-123")); - seed_stored_identity(&auth, "agent-123", "account-123"); - mount_human_biscuit(&server, &chatgpt_base_url, "agent-123").await; - Mock::given(method("POST")) - .and(path("/v1/agent/agent-123/task/register")) - .and(header("x-openai-authorization", "human-biscuit")) - .respond_with(ResponseTemplate::new(500)) - .expect(1) - .mount(&server) - .await; - - let (session, _turn_context, rx_event) = - make_agent_identity_session_and_context_with_rx(auth, chatgpt_base_url).await; - - session.maybe_prewarm_agent_task_registration().await; - - assert_eq!(session.state.lock().await.agent_task(), None); - assert!(rx_event.try_recv().is_err()); -} - -#[tokio::test] -async fn cached_agent_task_for_current_identity_clears_stale_task() { - let auth = make_chatgpt_auth("account-123", Some("user-123")); - seed_stored_identity(&auth, "agent-123", "account-123"); - let (session, _turn_context, _rx_event) = make_agent_identity_session_and_context_with_rx( - auth, - "https://chatgpt.com/backend-api".to_string(), - ) - .await; - { - let mut state = session.state.lock().await; - state.set_agent_task( - RegisteredAgentTask { - agent_runtime_id: "agent-old".to_string(), - task_id: "task-old".to_string(), - registered_at: "2026-04-15T00:00:00Z".to_string(), - } - .to_session_agent_task(), - ); - } - - assert_eq!(session.cached_agent_task_for_current_identity().await, None); - assert_eq!(session.state.lock().await.agent_task(), None); -} - -fn seed_stored_identity( - auth: &CodexAuth, - agent_runtime_id: &str, - account_id: &str, -) -> StoredAgentIdentity { - let signing_key = generate_test_signing_key(); - let private_key_pkcs8 = signing_key - .to_pkcs8_der() - .expect("encode test signing key as PKCS#8"); - let stored_identity = StoredAgentIdentity { - binding_id: format!("chatgpt-account-{account_id}"), - chatgpt_account_id: account_id.to_string(), - chatgpt_user_id: Some("user-123".to_string()), - agent_runtime_id: agent_runtime_id.to_string(), - private_key_pkcs8_base64: BASE64_STANDARD.encode(private_key_pkcs8.as_bytes()), - public_key_ssh: "ssh-ed25519 test".to_string(), - registered_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), - abom: crate::agent_identity::AgentBillOfMaterials { - agent_version: env!("CARGO_PKG_VERSION").to_string(), - agent_harness_id: "codex-cli".to_string(), - running_location: format!("{}-{}", SessionSource::Exec, std::env::consts::OS), - }, - }; - - auth.set_agent_identity(AgentIdentityAuthRecord { - workspace_id: account_id.to_string(), - chatgpt_user_id: stored_identity.chatgpt_user_id.clone(), - agent_runtime_id: stored_identity.agent_runtime_id.clone(), - agent_private_key: stored_identity.private_key_pkcs8_base64.clone(), - registered_at: stored_identity.registered_at.clone(), - }) - .expect("store identity"); - - stored_identity -} - -fn encrypt_task_id_for_identity( - stored_identity: &StoredAgentIdentity, - task_id: &str, -) -> anyhow::Result { - let signing_key = stored_identity.signing_key()?; - let mut rng = crypto_box::aead::OsRng; - let public_key = curve25519_secret_key_from_signing_key_for_tests(&signing_key).public_key(); - let ciphertext = public_key - .seal(&mut rng, task_id.as_bytes()) - .map_err(|_| anyhow::anyhow!("failed to encrypt test task id"))?; - Ok(BASE64_STANDARD.encode(ciphertext)) -} - -fn curve25519_secret_key_from_signing_key_for_tests( - signing_key: &SigningKey, -) -> Curve25519SecretKey { - let digest = Sha512::digest(signing_key.to_bytes()); - let mut secret_key = [0u8; 32]; - secret_key.copy_from_slice(&digest[..32]); - secret_key[0] &= 248; - secret_key[31] &= 127; - secret_key[31] |= 64; - Curve25519SecretKey::from(secret_key) -} - -fn generate_test_signing_key() -> SigningKey { - SigningKey::from_bytes(&[7u8; 32]) -} - -async fn mount_human_biscuit(server: &MockServer, chatgpt_base_url: &str, agent_runtime_id: &str) { - let biscuit_url = format!( - "{}/authenticate_app_v2", - chatgpt_base_url.trim_end_matches('/') - ); - let biscuit_path = reqwest::Url::parse(&biscuit_url) - .expect("biscuit URL parses") - .path() - .to_string(); - let target_url = format!( - "{}/v1/agent/{agent_runtime_id}/task/register", - chatgpt_base_url.trim_end_matches('/') - ); - Mock::given(method("GET")) - .and(path(biscuit_path)) - .and(header("authorization", "Bearer access-token-account-123")) - .and(header("x-original-method", "POST")) - .and(header("x-original-url", target_url)) - .respond_with( - ResponseTemplate::new(200).insert_header("x-openai-authorization", "human-biscuit"), - ) - .expect(1) - .mount(server) - .await; -} - -fn make_chatgpt_auth(account_id: &str, user_id: Option<&str>) -> CodexAuth { - let tempdir = tempfile::tempdir().expect("tempdir"); - let auth_json = AuthDotJson { - auth_mode: Some(codex_app_server_protocol::AuthMode::Chatgpt), - openai_api_key: None, - tokens: Some(TokenData { - id_token: IdTokenInfo { - email: None, - chatgpt_plan_type: None, - chatgpt_user_id: user_id.map(ToOwned::to_owned), - chatgpt_account_id: Some(account_id.to_string()), - chatgpt_account_is_fedramp: false, - raw_jwt: fake_id_token(account_id, user_id), - }, - access_token: format!("access-token-{account_id}"), - refresh_token: "refresh-token".to_string(), - account_id: Some(account_id.to_string()), - }), - last_refresh: Some(Utc::now()), - agent_identity: None, - }; - save_auth(tempdir.path(), &auth_json, AuthCredentialsStoreMode::File).expect("save auth"); - CodexAuth::from_auth_storage(tempdir.path(), AuthCredentialsStoreMode::File) - .expect("load auth") - .expect("auth") -} - -fn fake_id_token(account_id: &str, user_id: Option<&str>) -> String { - let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#); - let payload = serde_json::json!({ - "https://api.openai.com/auth": { - "chatgpt_user_id": user_id, - "chatgpt_account_id": account_id, - } - }); - let payload = URL_SAFE_NO_PAD.encode(payload.to_string()); - format!("{header}.{payload}.signature") -} - #[tokio::test] async fn refresh_mcp_servers_is_deferred_until_next_turn() { let (session, turn_context) = make_session_and_context().await; diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 4ac8c4d5303a..bf6e06d8a93b 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -6,12 +6,12 @@ use codex_sandboxing::policy_transforms::merge_permission_profiles; use std::collections::HashMap; use std::collections::HashSet; +use crate::agent_identity::RegisteredAgentTask; use crate::context_manager::ContextManager; use crate::session::PreviousTurnSettings; use crate::session::session::SessionConfiguration; use crate::session_startup_prewarm::SessionStartupPrewarmHandle; use codex_protocol::protocol::RateLimitSnapshot; -use codex_protocol::protocol::SessionAgentTask; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; use codex_protocol::protocol::TurnContextItem; @@ -31,7 +31,7 @@ pub(crate) struct SessionState { previous_turn_settings: Option, /// Startup prewarmed session prepared during session initialization. pub(crate) startup_prewarm: Option, - pub(crate) agent_task: Option, + pub(crate) agent_task: Option, pub(crate) active_connector_selection: HashSet, pub(crate) pending_session_start_source: Option, granted_permissions: Option, @@ -189,11 +189,11 @@ impl SessionState { self.startup_prewarm.take() } - pub(crate) fn agent_task(&self) -> Option { + pub(crate) fn agent_task(&self) -> Option { self.agent_task.clone() } - pub(crate) fn set_agent_task(&mut self, agent_task: SessionAgentTask) { + pub(crate) fn set_agent_task(&mut self, agent_task: RegisteredAgentTask) { self.agent_task = Some(agent_task); } diff --git a/codex-rs/core/src/state/session_tests.rs b/codex-rs/core/src/state/session_tests.rs index 12229ee2f270..927c1fbb70fd 100644 --- a/codex-rs/core/src/state/session_tests.rs +++ b/codex-rs/core/src/state/session_tests.rs @@ -1,8 +1,8 @@ use super::*; +use crate::agent_identity::RegisteredAgentTask; use crate::session::tests::make_session_configuration_for_tests; use codex_protocol::protocol::CreditsSnapshot; use codex_protocol::protocol::RateLimitWindow; -use codex_protocol::protocol::SessionAgentTask; use pretty_assertions::assert_eq; #[tokio::test] @@ -38,8 +38,11 @@ async fn clear_connector_selection_removes_entries() { async fn set_agent_task_persists_plaintext_task_for_session_reuse() { let session_configuration = make_session_configuration_for_tests().await; let mut state = SessionState::new(session_configuration); - let agent_task = SessionAgentTask { - agent_runtime_id: "agent-123".to_string(), + let agent_task = RegisteredAgentTask { + binding_id: "chatgpt-account-account-123".to_string(), + chatgpt_account_id: "account-123".to_string(), + chatgpt_user_id: Some("user-123".to_string()), + agent_runtime_id: "agent_123".to_string(), task_id: "task_123".to_string(), registered_at: "2026-03-23T12:00:00Z".to_string(), }; @@ -53,7 +56,10 @@ async fn set_agent_task_persists_plaintext_task_for_session_reuse() { async fn clear_agent_task_removes_cached_task() { let session_configuration = make_session_configuration_for_tests().await; let mut state = SessionState::new(session_configuration); - let agent_task = SessionAgentTask { + let agent_task = RegisteredAgentTask { + binding_id: "chatgpt-account-account-123".to_string(), + chatgpt_account_id: "account-123".to_string(), + chatgpt_user_id: Some("user-123".to_string()), agent_runtime_id: "agent_123".to_string(), task_id: "task_123".to_string(), registered_at: "2026-03-23T12:00:00Z".to_string(), diff --git a/codex-rs/core/tests/suite/otel.rs b/codex-rs/core/tests/suite/otel.rs index f93945e78fff..d771ed19aced 100644 --- a/codex-rs/core/tests/suite/otel.rs +++ b/codex-rs/core/tests/suite/otel.rs @@ -228,7 +228,6 @@ async fn process_sse_records_failed_event_when_stream_closes_without_completed() .features .disable(Feature::GhostCommit) .expect("test config should allow feature update"); - config.model_provider.stream_max_retries = Some(0); }) .build(&server) .await @@ -407,7 +406,6 @@ async fn process_sse_failed_event_logs_missing_error() { .features .disable(Feature::GhostCommit) .expect("test config should allow feature update"); - config.model_provider.stream_max_retries = Some(0); }) .build(&server) .await diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index b9a7a1395f52..f1c6046c7b53 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2779,26 +2779,6 @@ impl fmt::Display for SubAgentSource { } } -/// Persisted agent-task details that let a resumed thread keep using the same backend task. -/// -/// `agent_runtime_id` is validation metadata for the globally registered agent identity, not a -/// separate session-scoped identity. Resume only restores this task after confirming that runtime -/// id still matches the globally registered identity; otherwise the cached task is discarded and a -/// fresh task can be registered. -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema, TS)] -pub struct SessionAgentTask { - pub agent_runtime_id: String, - pub task_id: String, - pub registered_at: String, -} - -/// Session-scoped state updates that can be appended after the canonical SessionMeta line. -#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema, TS, Default)] -pub struct SessionStateUpdate { - #[serde(default)] - pub agent_task: Option, -} - /// SessionMeta contains session-level data that doesn't correspond to a specific turn. /// /// NOTE: There used to be an `instructions` field here, which stored user_instructions, but we @@ -2868,7 +2848,6 @@ pub struct SessionMetaLine { #[serde(tag = "type", content = "payload", rename_all = "snake_case")] pub enum RolloutItem { SessionMeta(SessionMetaLine), - SessionState(SessionStateUpdate), ResponseItem(ResponseItem), Compacted(CompactedItem), TurnContext(TurnContextItem), diff --git a/codex-rs/rollout/src/list.rs b/codex-rs/rollout/src/list.rs index 0df1b9c9f548..bf09d15e8677 100644 --- a/codex-rs/rollout/src/list.rs +++ b/codex-rs/rollout/src/list.rs @@ -1082,9 +1082,6 @@ async fn read_head_summary(path: &Path, head_limit: usize) -> io::Result { // Not included in `head`; skip. } - RolloutItem::SessionState(_) => { - // Not included in `head`; skip. - } RolloutItem::EventMsg(ev) => { if let EventMsg::UserMessage(user) = ev { summary.saw_user_event = true; @@ -1137,7 +1134,6 @@ pub async fn read_head_for_summary(path: &Path) -> io::Result {} } diff --git a/codex-rs/rollout/src/metadata.rs b/codex-rs/rollout/src/metadata.rs index 65e347a3232e..58d55a887df3 100644 --- a/codex-rs/rollout/src/metadata.rs +++ b/codex-rs/rollout/src/metadata.rs @@ -70,8 +70,7 @@ pub fn builder_from_items( ) -> Option { if let Some(session_meta) = items.iter().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => Some(meta_line), - RolloutItem::SessionState(_) - | RolloutItem::ResponseItem(_) + RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::EventMsg(_) => None, @@ -125,8 +124,7 @@ pub async fn extract_metadata_from_rollout( metadata, memory_mode: items.iter().rev().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => meta_line.meta.memory_mode.clone(), - RolloutItem::SessionState(_) - | RolloutItem::ResponseItem(_) + RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::EventMsg(_) => None, diff --git a/codex-rs/rollout/src/policy.rs b/codex-rs/rollout/src/policy.rs index 86fc950a24b9..df41d2c76676 100644 --- a/codex-rs/rollout/src/policy.rs +++ b/codex-rs/rollout/src/policy.rs @@ -16,10 +16,9 @@ pub fn is_persisted_response_item(item: &RolloutItem, mode: EventPersistenceMode RolloutItem::ResponseItem(item) => should_persist_response_item(item), RolloutItem::EventMsg(ev) => should_persist_event_msg(ev, mode), // Persist Codex executive markers so we can analyze flows (e.g., compaction, API turns). - RolloutItem::Compacted(_) - | RolloutItem::TurnContext(_) - | RolloutItem::SessionMeta(_) - | RolloutItem::SessionState(_) => true, + RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::SessionMeta(_) => { + true + } } } diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index f6e15a973eb9..871d53141bed 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -707,9 +707,6 @@ impl RolloutRecorder { RolloutItem::Compacted(item) => { items.push(RolloutItem::Compacted(item)); } - RolloutItem::SessionState(update) => { - items.push(RolloutItem::SessionState(update)); - } RolloutItem::TurnContext(item) => { items.push(RolloutItem::TurnContext(item)); } @@ -1566,7 +1563,6 @@ async fn resume_candidate_matches_cwd( && let Some(latest_turn_context_cwd) = items.iter().rev().find_map(|item| match item { RolloutItem::TurnContext(turn_context) => Some(turn_context.cwd.as_path()), RolloutItem::SessionMeta(_) - | RolloutItem::SessionState(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::EventMsg(_) => None, diff --git a/codex-rs/state/src/extract.rs b/codex-rs/state/src/extract.rs index 64d7f5bbb71a..6d36f211a4c4 100644 --- a/codex-rs/state/src/extract.rs +++ b/codex-rs/state/src/extract.rs @@ -19,7 +19,6 @@ pub fn apply_rollout_item( ) { match item { RolloutItem::SessionMeta(meta_line) => apply_session_meta_from_item(metadata, meta_line), - RolloutItem::SessionState(_) => {} RolloutItem::TurnContext(turn_ctx) => apply_turn_context(metadata, turn_ctx), RolloutItem::EventMsg(event) => apply_event_msg(metadata, event), RolloutItem::ResponseItem(item) => apply_response_item(metadata, item), @@ -37,10 +36,9 @@ pub fn rollout_item_affects_thread_metadata(item: &RolloutItem) -> bool { RolloutItem::EventMsg( EventMsg::TokenCount(_) | EventMsg::UserMessage(_) | EventMsg::ThreadNameUpdated(_), ) => true, - RolloutItem::SessionState(_) - | RolloutItem::EventMsg(_) - | RolloutItem::ResponseItem(_) - | RolloutItem::Compacted(_) => false, + RolloutItem::EventMsg(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) => { + false + } } } diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index ec094a7339c0..70fea73ff4a2 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -968,8 +968,7 @@ SELECT pub(super) fn extract_dynamic_tools(items: &[RolloutItem]) -> Option>> { items.iter().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => Some(meta_line.meta.dynamic_tools.clone()), - RolloutItem::SessionState(_) - | RolloutItem::ResponseItem(_) + RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::EventMsg(_) => None, @@ -979,8 +978,7 @@ pub(super) fn extract_dynamic_tools(items: &[RolloutItem]) -> Option Option { items.iter().rev().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => meta_line.meta.memory_mode.clone(), - RolloutItem::SessionState(_) - | RolloutItem::ResponseItem(_) + RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::EventMsg(_) => None, From 59727a8901730c50a15913a7b9847b43af897d66 Mon Sep 17 00:00:00 2001 From: Edward Frazer Date: Tue, 21 Apr 2026 09:32:10 -0700 Subject: [PATCH 4/9] Revert "Register agent tasks behind use_agent_identity (#17387)" This reverts commit 55c3de75cba1a65088ff02b91527b94a1b60a69b. --- MODULE.bazel.lock | 3 - codex-rs/Cargo.lock | 43 -- codex-rs/Cargo.toml | 1 - codex-rs/app-server/tests/suite/auth.rs | 8 +- .../app-server/tests/suite/v2/plugin_list.rs | 3 +- .../tests/suite/v2/realtime_conversation.rs | 5 +- codex-rs/core/Cargo.toml | 2 - codex-rs/core/src/agent_identity.rs | 81 +-- .../src/agent_identity/task_registration.rs | 470 ------------------ codex-rs/core/src/session/tests.rs | 13 +- codex-rs/core/src/state/session.rs | 15 - codex-rs/core/src/state/session_tests.rs | 38 -- .../core/tests/suite/realtime_conversation.rs | 63 +-- 13 files changed, 41 insertions(+), 704 deletions(-) delete mode 100644 codex-rs/core/src/agent_identity/task_registration.rs diff --git a/MODULE.bazel.lock b/MODULE.bazel.lock index c107884e1dfa..24fc9dbbd93b 100644 --- a/MODULE.bazel.lock +++ b/MODULE.bazel.lock @@ -681,7 +681,6 @@ "bitflags_1.3.2": "{\"dependencies\":[{\"name\":\"compiler_builtins\",\"optional\":true,\"req\":\"^0.1.2\"},{\"name\":\"core\",\"optional\":true,\"package\":\"rustc-std-workspace-core\",\"req\":\"^1.0.0\"},{\"kind\":\"dev\",\"name\":\"rustversion\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde_derive\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"trybuild\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"walkdir\",\"req\":\"^2.3\"}],\"features\":{\"default\":[],\"example_generated\":[],\"rustc-dep-of-std\":[\"core\",\"compiler_builtins\"]}}", "bitflags_2.10.0": "{\"dependencies\":[{\"name\":\"arbitrary\",\"optional\":true,\"req\":\"^1.0\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"arbitrary\",\"req\":\"^1.0\"},{\"name\":\"bytemuck\",\"optional\":true,\"req\":\"^1.12\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"bytemuck\",\"req\":\"^1.12.2\"},{\"kind\":\"dev\",\"name\":\"rustversion\",\"req\":\"^1.0\"},{\"default_features\":false,\"name\":\"serde_core\",\"optional\":true,\"req\":\"^1.0.228\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde_lib\",\"package\":\"serde\",\"req\":\"^1.0.103\"},{\"kind\":\"dev\",\"name\":\"serde_test\",\"req\":\"^1.0.19\"},{\"kind\":\"dev\",\"name\":\"trybuild\",\"req\":\"^1.0.18\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"zerocopy\",\"req\":\"^0.8\"}],\"features\":{\"example_generated\":[],\"serde\":[\"serde_core\"],\"std\":[]}}", "bitflags_2.11.0": "{\"dependencies\":[{\"name\":\"arbitrary\",\"optional\":true,\"req\":\"^1.0\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"arbitrary\",\"req\":\"^1.0\"},{\"name\":\"bytemuck\",\"optional\":true,\"req\":\"^1.12\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"bytemuck\",\"req\":\"^1.12.2\"},{\"kind\":\"dev\",\"name\":\"rustversion\",\"req\":\"^1.0\"},{\"default_features\":false,\"name\":\"serde_core\",\"optional\":true,\"req\":\"^1.0.228\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde_lib\",\"package\":\"serde\",\"req\":\"^1.0.103\"},{\"kind\":\"dev\",\"name\":\"serde_test\",\"req\":\"^1.0.19\"},{\"kind\":\"dev\",\"name\":\"trybuild\",\"req\":\"^1.0.18\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"zerocopy\",\"req\":\"^0.8\"}],\"features\":{\"example_generated\":[],\"serde\":[\"serde_core\"],\"std\":[]}}", - "blake2_0.10.6": "{\"dependencies\":[{\"features\":[\"mac\"],\"name\":\"digest\",\"req\":\"^0.10.3\"},{\"features\":[\"dev\"],\"kind\":\"dev\",\"name\":\"digest\",\"req\":\"^0.10.3\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.2.2\"}],\"features\":{\"default\":[\"std\"],\"reset\":[],\"simd\":[],\"simd_asm\":[\"simd_opt\"],\"simd_opt\":[\"simd\"],\"size_opt\":[],\"std\":[\"digest/std\"]}}", "block-buffer_0.10.4": "{\"dependencies\":[{\"name\":\"generic-array\",\"req\":\"^0.14\"}],\"features\":{}}", "block-padding_0.3.3": "{\"dependencies\":[{\"name\":\"generic-array\",\"req\":\"^0.14\"}],\"features\":{\"std\":[]}}", "block2_0.6.2": "{\"dependencies\":[{\"default_features\":false,\"features\":[\"std\"],\"name\":\"objc2\",\"req\":\">=0.6.2, <0.8.0\"}],\"features\":{\"alloc\":[],\"compiler-rt\":[\"objc2/unstable-compiler-rt\"],\"default\":[\"std\"],\"gnustep-1-7\":[\"objc2/gnustep-1-7\"],\"gnustep-1-8\":[\"gnustep-1-7\",\"objc2/gnustep-1-8\"],\"gnustep-1-9\":[\"gnustep-1-8\",\"objc2/gnustep-1-9\"],\"gnustep-2-0\":[\"gnustep-1-9\",\"objc2/gnustep-2-0\"],\"gnustep-2-1\":[\"gnustep-2-0\",\"objc2/gnustep-2-1\"],\"std\":[\"alloc\"],\"unstable-coerce-pointee\":[],\"unstable-objfw\":[],\"unstable-private\":[],\"unstable-winobjc\":[\"gnustep-1-8\"]}}", @@ -777,8 +776,6 @@ "crunchy_0.2.4": "{\"dependencies\":[],\"features\":{\"default\":[\"limit_128\"],\"limit_1024\":[],\"limit_128\":[],\"limit_2048\":[],\"limit_256\":[],\"limit_512\":[],\"limit_64\":[],\"std\":[]}}", "crypto-bigint_0.5.5": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"bincode\",\"req\":\"^1\"},{\"features\":[\"html_reports\"],\"kind\":\"dev\",\"name\":\"criterion\",\"req\":\"^0.5\"},{\"default_features\":false,\"name\":\"der\",\"optional\":true,\"req\":\"^0.7\"},{\"name\":\"generic-array\",\"optional\":true,\"req\":\"^0.14\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.4\"},{\"kind\":\"dev\",\"name\":\"num-bigint\",\"req\":\"^0.4\"},{\"kind\":\"dev\",\"name\":\"num-integer\",\"req\":\"^0.1\"},{\"kind\":\"dev\",\"name\":\"num-traits\",\"req\":\"^0.2\"},{\"kind\":\"dev\",\"name\":\"proptest\",\"req\":\"^1\"},{\"kind\":\"dev\",\"name\":\"rand_chacha\",\"req\":\"^0.3\"},{\"name\":\"rand_core\",\"optional\":true,\"req\":\"^0.6.4\"},{\"features\":[\"std\"],\"kind\":\"dev\",\"name\":\"rand_core\",\"req\":\"^0.6\"},{\"default_features\":false,\"name\":\"rlp\",\"optional\":true,\"req\":\"^0.5\"},{\"default_features\":false,\"name\":\"serdect\",\"optional\":true,\"req\":\"^0.2\"},{\"default_features\":false,\"name\":\"subtle\",\"req\":\"^2.4\"},{\"default_features\":false,\"name\":\"zeroize\",\"optional\":true,\"req\":\"^1\"}],\"features\":{\"alloc\":[\"serdect?/alloc\"],\"default\":[\"rand\"],\"extra-sizes\":[],\"rand\":[\"rand_core/std\"],\"serde\":[\"dep:serdect\"]}}", "crypto-common_0.1.7": "{\"dependencies\":[{\"features\":[\"more_lengths\"],\"name\":\"generic-array\",\"req\":\"=0.14.7\"},{\"name\":\"rand_core\",\"optional\":true,\"req\":\"^0.6\"},{\"name\":\"typenum\",\"req\":\"^1.14\"}],\"features\":{\"getrandom\":[\"rand_core/getrandom\"],\"std\":[]}}", - "crypto_box_0.9.1": "{\"dependencies\":[{\"default_features\":false,\"name\":\"aead\",\"req\":\"^0.5.2\"},{\"kind\":\"dev\",\"name\":\"bincode\",\"req\":\"^1\"},{\"default_features\":false,\"name\":\"blake2\",\"optional\":true,\"req\":\"^0.10\"},{\"name\":\"chacha20\",\"optional\":true,\"req\":\"^0.9\"},{\"default_features\":false,\"name\":\"crypto_secretbox\",\"req\":\"^0.1.1\"},{\"default_features\":false,\"features\":[\"zeroize\"],\"name\":\"curve25519-dalek\",\"req\":\"^4\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.4\"},{\"kind\":\"dev\",\"name\":\"rand\",\"req\":\"^0.8\"},{\"kind\":\"dev\",\"name\":\"rmp-serde\",\"req\":\"^1\"},{\"name\":\"salsa20\",\"optional\":true,\"req\":\"^0.10\"},{\"default_features\":false,\"name\":\"serdect\",\"optional\":true,\"req\":\"^0.2\"},{\"default_features\":false,\"name\":\"subtle\",\"req\":\"^2\"},{\"default_features\":false,\"name\":\"zeroize\",\"req\":\"^1\"}],\"features\":{\"alloc\":[\"aead/alloc\"],\"chacha20\":[\"dep:chacha20\",\"crypto_secretbox/chacha20\"],\"default\":[\"alloc\",\"getrandom\",\"salsa20\"],\"getrandom\":[\"aead/getrandom\",\"rand_core\"],\"heapless\":[\"aead/heapless\"],\"rand_core\":[\"aead/rand_core\"],\"salsa20\":[\"dep:salsa20\",\"crypto_secretbox/salsa20\"],\"seal\":[\"dep:blake2\",\"alloc\"],\"serde\":[\"dep:serdect\"],\"std\":[\"aead/std\"]}}", - "crypto_secretbox_0.1.1": "{\"dependencies\":[{\"default_features\":false,\"name\":\"aead\",\"req\":\"^0.5\"},{\"features\":[\"zeroize\"],\"name\":\"chacha20\",\"optional\":true,\"req\":\"^0.9\"},{\"default_features\":false,\"name\":\"cipher\",\"req\":\"^0.4\"},{\"default_features\":false,\"features\":[\"zeroize\"],\"name\":\"generic-array\",\"req\":\"^0.14.7\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.4\"},{\"name\":\"poly1305\",\"req\":\"^0.8\"},{\"features\":[\"zeroize\"],\"name\":\"salsa20\",\"optional\":true,\"req\":\"^0.10\"},{\"default_features\":false,\"name\":\"subtle\",\"req\":\"^2\"},{\"default_features\":false,\"name\":\"zeroize\",\"req\":\"^1\"}],\"features\":{\"alloc\":[\"aead/alloc\"],\"default\":[\"alloc\",\"getrandom\",\"salsa20\"],\"getrandom\":[\"aead/getrandom\",\"rand_core\"],\"heapless\":[\"aead/heapless\"],\"rand_core\":[\"aead/rand_core\"],\"std\":[\"aead/std\",\"alloc\"],\"stream\":[\"aead/stream\"]}}", "csv-core_0.1.13": "{\"dependencies\":[{\"default_features\":false,\"kind\":\"dev\",\"name\":\"arrayvec\",\"req\":\"^0.5\"},{\"default_features\":false,\"name\":\"memchr\",\"req\":\"^2\"}],\"features\":{\"default\":[],\"libc\":[\"memchr/libc\"]}}", "csv_1.4.0": "{\"dependencies\":[{\"default_features\":false,\"features\":[\"alloc\",\"serde\"],\"kind\":\"dev\",\"name\":\"bstr\",\"req\":\"^1.7.0\"},{\"name\":\"csv-core\",\"req\":\"^0.1.11\"},{\"name\":\"itoa\",\"req\":\"^1\"},{\"name\":\"ryu\",\"req\":\"^1\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1.0.221\"},{\"name\":\"serde_core\",\"req\":\"^1.0.221\"}],\"features\":{}}", "ctor-proc-macro_0.0.7": "{\"dependencies\":[],\"features\":{\"default\":[]}}", diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 859e54b926cf..41ddc808dde6 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -946,15 +946,6 @@ dependencies = [ "serde_core", ] -[[package]] -name = "blake2" -version = "0.10.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46502ad458c9a52b69d4d4d32775c788b7a1b85e8bc9d482d92250fc0e3f8efe" -dependencies = [ - "digest", -] - [[package]] name = "block-buffer" version = "0.10.4" @@ -1995,7 +1986,6 @@ dependencies = [ "codex-windows-sandbox", "core-foundation 0.9.4", "core_test_support", - "crypto_box", "csv", "ctor 0.6.3", "dirs", @@ -2026,7 +2016,6 @@ dependencies = [ "serde_json", "serial_test", "sha1", - "sha2", "shlex", "similar", "tempfile", @@ -3750,40 +3739,9 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" dependencies = [ "generic-array", - "rand_core 0.6.4", "typenum", ] -[[package]] -name = "crypto_box" -version = "0.9.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "16182b4f39a82ec8a6851155cc4c0cda3065bb1db33651726a29e1951de0f009" -dependencies = [ - "aead", - "blake2", - "crypto_secretbox", - "curve25519-dalek", - "salsa20", - "subtle", - "zeroize", -] - -[[package]] -name = "crypto_secretbox" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9d6cf87adf719ddf43a805e92c6870a531aedda35ff640442cbaf8674e141e1" -dependencies = [ - "aead", - "cipher", - "generic-array", - "poly1305", - "salsa20", - "subtle", - "zeroize", -] - [[package]] name = "csv" version = "1.4.0" @@ -5164,7 +5122,6 @@ checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" dependencies = [ "typenum", "version_check", - "zeroize", ] [[package]] diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 236412051177..04fd2fa06c42 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -228,7 +228,6 @@ color-eyre = "0.6.3" constant_time_eq = "0.3.1" crossbeam-channel = "0.5.15" crossterm = "0.28.1" -crypto_box = { version = "0.9.1", features = ["seal"] } csv = "1.3.1" ctor = "0.6.3" deno_core_icudata = "0.77.0" diff --git a/codex-rs/app-server/tests/suite/auth.rs b/codex-rs/app-server/tests/suite/auth.rs index b35a6fcba152..1e608710126d 100644 --- a/codex-rs/app-server/tests/suite/auth.rs +++ b/codex-rs/app-server/tests/suite/auth.rs @@ -351,8 +351,6 @@ async fn get_auth_status_omits_token_after_proactive_refresh_failure() -> Result )?; let server = MockServer::start().await; - // App-server startup may proactively read stale auth before this test sends - // getAuthStatus; require the refresh path without depending on that race. Mock::given(method("POST")) .and(path("/oauth/token")) .respond_with(ResponseTemplate::new(401).set_body_json(serde_json::json!({ @@ -360,7 +358,7 @@ async fn get_auth_status_omits_token_after_proactive_refresh_failure() -> Result "code": "refresh_token_reused" } }))) - .expect(1..=2) + .expect(2) .mount(&server) .await; @@ -420,8 +418,6 @@ async fn get_auth_status_returns_token_after_proactive_refresh_recovery() -> Res )?; let server = MockServer::start().await; - // App-server startup may proactively read stale auth before this test sends - // getAuthStatus; require the refresh path without depending on that race. Mock::given(method("POST")) .and(path("/oauth/token")) .respond_with(ResponseTemplate::new(401).set_body_json(serde_json::json!({ @@ -429,7 +425,7 @@ async fn get_auth_status_returns_token_after_proactive_refresh_recovery() -> Res "code": "refresh_token_reused" } }))) - .expect(1..=2) + .expect(2) .mount(&server) .await; diff --git a/codex-rs/app-server/tests/suite/v2/plugin_list.rs b/codex-rs/app-server/tests/suite/v2/plugin_list.rs index bf69df3c470c..9718761e3359 100644 --- a/codex-rs/app-server/tests/suite/v2/plugin_list.rs +++ b/codex-rs/app-server/tests/suite/v2/plugin_list.rs @@ -30,8 +30,7 @@ use wiremock::matchers::method; use wiremock::matchers::path; use wiremock::matchers::query_param; -// These tests start full app-server processes and can also run plugin startup warmers. -const DEFAULT_TIMEOUT: Duration = Duration::from_secs(60); +const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10); const TEST_CURATED_PLUGIN_SHA: &str = "0123456789abcdef0123456789abcdef01234567"; const STARTUP_REMOTE_PLUGIN_SYNC_MARKER_FILE: &str = ".tmp/app-server-remote-plugin-sync-v1"; const ALTERNATE_MARKETPLACE_RELATIVE_PATH: &str = ".claude-plugin/marketplace.json"; diff --git a/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs b/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs index dfc3fea31820..793efbe1e182 100644 --- a/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs +++ b/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs @@ -72,7 +72,6 @@ use wiremock::matchers::path; use wiremock::matchers::path_regex; const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10); -const DELEGATED_SHELL_TOOL_TIMEOUT_MS: u64 = 30_000; const STARTUP_CONTEXT_HEADER: &str = "Startup context from Codex."; const V2_STEERING_ACKNOWLEDGEMENT: &str = "This was sent to steer the previous background agent task."; @@ -1782,9 +1781,7 @@ async fn webrtc_v2_tool_call_delegated_turn_can_execute_shell_tool() -> Result<( create_shell_command_sse_response( realtime_tool_ok_command(), /*workdir*/ None, - // Windows CI can spend several seconds starting the nested PowerShell command. This - // test verifies delegated shell-tool plumbing, not timeout enforcement. - Some(DELEGATED_SHELL_TOOL_TIMEOUT_MS), + Some(5000), "shell_call", )?, create_final_assistant_message_sse_response("shell tool finished")?, diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index d1cca0ac32dc..5b003aa1d6b5 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -35,7 +35,6 @@ codex-connectors = { workspace = true } codex-config = { workspace = true } codex-core-plugins = { workspace = true } codex-core-skills = { workspace = true } -crypto_box = { workspace = true } codex-exec-server = { workspace = true } codex-features = { workspace = true } codex-feedback = { workspace = true } @@ -100,7 +99,6 @@ rmcp = { workspace = true, default-features = false, features = [ serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } sha1 = { workspace = true } -sha2 = { workspace = true } shlex = { workspace = true } similar = { workspace = true } tempfile = { workspace = true } diff --git a/codex-rs/core/src/agent_identity.rs b/codex-rs/core/src/agent_identity.rs index 8edcb7bb10d1..7e86ba22ae4a 100644 --- a/codex-rs/core/src/agent_identity.rs +++ b/codex-rs/core/src/agent_identity.rs @@ -27,10 +27,6 @@ use tracing::debug; use tracing::info; use tracing::warn; -mod task_registration; - -pub(crate) use task_registration::RegisteredAgentTask; - use crate::config::Config; const AGENT_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15); @@ -123,62 +119,36 @@ impl AgentIdentityManager { return Ok(None); } - let Some((auth, binding)) = self.current_auth_binding().await else { + let Some(auth) = self.auth_manager.auth().await else { + debug!("skipping agent identity registration because no auth is available"); return Ok(None); }; - self.ensure_registered_identity_for_binding(&auth, &binding) - .await - .map(Some) - } + let Some(binding) = + AgentIdentityBinding::from_auth(&auth, self.auth_manager.forced_chatgpt_workspace_id()) + else { + debug!("skipping agent identity registration because ChatGPT auth is unavailable"); + return Ok(None); + }; - async fn ensure_registered_identity_for_binding( - &self, - auth: &CodexAuth, - binding: &AgentIdentityBinding, - ) -> Result { let _guard = self .ensure_lock .acquire() .await .map_err(|_| anyhow::anyhow!("agent identity ensure semaphore closed"))?; - if let Some(stored_identity) = self.load_stored_identity(auth, binding)? { + if let Some(stored_identity) = self.load_stored_identity(&auth, &binding)? { info!( agent_runtime_id = %stored_identity.agent_runtime_id, binding_id = %binding.binding_id, "reusing stored agent identity" ); - return Ok(stored_identity); + return Ok(Some(stored_identity)); } - let stored_identity = self.register_agent_identity(binding).await?; - self.store_identity(auth, &stored_identity)?; - Ok(stored_identity) - } - - pub(crate) async fn task_matches_current_binding(&self, task: &RegisteredAgentTask) -> bool { - if !self.feature_enabled { - return false; - } - - self.current_auth_binding() - .await - .is_some_and(|(_, binding)| task.matches_binding(&binding)) - } - - async fn current_auth_binding(&self) -> Option<(CodexAuth, AgentIdentityBinding)> { - let Some(auth) = self.auth_manager.auth().await else { - debug!("skipping agent identity flow because no auth is available"); - return None; - }; - - let binding = - AgentIdentityBinding::from_auth(&auth, self.auth_manager.forced_chatgpt_workspace_id()); - if binding.is_none() { - debug!("skipping agent identity flow because ChatGPT auth is unavailable"); - } - binding.map(|binding| (auth, binding)) + let stored_identity = self.register_agent_identity(&binding).await?; + self.store_identity(&auth, &stored_identity)?; + Ok(Some(stored_identity)) } async fn register_agent_identity( @@ -385,11 +355,12 @@ impl StoredAgentIdentity { } fn matches_binding(&self, binding: &AgentIdentityBinding) -> bool { - binding.matches_parts( - &self.binding_id, - &self.chatgpt_account_id, - self.chatgpt_user_id.as_deref(), - ) + self.binding_id == binding.binding_id + && self.chatgpt_account_id == binding.chatgpt_account_id + && match binding.chatgpt_user_id.as_deref() { + Some(chatgpt_user_id) => self.chatgpt_user_id.as_deref() == Some(chatgpt_user_id), + None => true, + } } fn validate_key_material(&self) -> Result<()> { @@ -408,20 +379,6 @@ impl StoredAgentIdentity { } impl AgentIdentityBinding { - fn matches_parts( - &self, - binding_id: &str, - chatgpt_account_id: &str, - chatgpt_user_id: Option<&str>, - ) -> bool { - binding_id == self.binding_id - && chatgpt_account_id == self.chatgpt_account_id - && match self.chatgpt_user_id.as_deref() { - Some(expected_user_id) => chatgpt_user_id == Some(expected_user_id), - None => true, - } - } - fn from_auth(auth: &CodexAuth, forced_workspace_id: Option) -> Option { if !auth.is_chatgpt_auth() { return None; diff --git a/codex-rs/core/src/agent_identity/task_registration.rs b/codex-rs/core/src/agent_identity/task_registration.rs deleted file mode 100644 index 4fc5d51282d3..000000000000 --- a/codex-rs/core/src/agent_identity/task_registration.rs +++ /dev/null @@ -1,470 +0,0 @@ -use std::time::Duration; - -use anyhow::Context; -use anyhow::Result; -use crypto_box::SecretKey as Curve25519SecretKey; -use ed25519_dalek::Signer as _; -use serde::Deserialize; -use serde::Serialize; -use sha2::Digest as _; -use sha2::Sha512; -use tracing::info; - -use super::*; - -const AGENT_TASK_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15); - -#[derive(Clone, Debug, PartialEq, Eq)] -pub(crate) struct RegisteredAgentTask { - pub(crate) binding_id: String, - pub(crate) chatgpt_account_id: String, - pub(crate) chatgpt_user_id: Option, - pub(crate) agent_runtime_id: String, - pub(crate) task_id: String, - pub(crate) registered_at: String, -} - -#[derive(Debug, Serialize)] -struct RegisterTaskRequest { - signature: String, - timestamp: String, -} - -#[derive(Debug, Deserialize)] -struct RegisterTaskResponse { - encrypted_task_id: String, -} - -impl AgentIdentityManager { - pub(crate) async fn register_task(&self) -> Result> { - if !self.feature_enabled { - return Ok(None); - } - - let Some((auth, binding)) = self.current_auth_binding().await else { - return Ok(None); - }; - - self.register_task_for_binding(auth, binding).await - } - - async fn register_task_for_binding( - &self, - auth: CodexAuth, - binding: AgentIdentityBinding, - ) -> Result> { - let stored_identity = self - .ensure_registered_identity_for_binding(&auth, &binding) - .await?; - - let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); - let request_body = RegisterTaskRequest { - signature: sign_task_registration_payload(&stored_identity, ×tamp)?, - timestamp, - }; - - let client = create_client(); - let url = - agent_task_registration_url(&self.chatgpt_base_url, &stored_identity.agent_runtime_id); - let human_biscuit = self.mint_human_biscuit(&binding, "POST", &url).await?; - let response = client - .post(&url) - .header("X-OpenAI-Authorization", human_biscuit) - .json(&request_body) - .timeout(AGENT_TASK_REGISTRATION_TIMEOUT) - .send() - .await - .with_context(|| format!("failed to send agent task registration request to {url}"))?; - - if response.status().is_success() { - let response_body = response - .json::() - .await - .with_context(|| format!("failed to parse agent task response from {url}"))?; - let registered_task = RegisteredAgentTask { - binding_id: stored_identity.binding_id.clone(), - chatgpt_account_id: stored_identity.chatgpt_account_id.clone(), - chatgpt_user_id: stored_identity.chatgpt_user_id.clone(), - agent_runtime_id: stored_identity.agent_runtime_id.clone(), - task_id: decrypt_task_id_response( - &stored_identity, - &response_body.encrypted_task_id, - )?, - registered_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), - }; - info!( - agent_runtime_id = %registered_task.agent_runtime_id, - task_id = %registered_task.task_id, - "registered agent task" - ); - return Ok(Some(registered_task)); - } - - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - anyhow::bail!("agent task registration failed with status {status} from {url}: {body}") - } -} - -impl RegisteredAgentTask { - pub(super) fn matches_binding(&self, binding: &AgentIdentityBinding) -> bool { - binding.matches_parts( - &self.binding_id, - &self.chatgpt_account_id, - self.chatgpt_user_id.as_deref(), - ) - } - - pub(crate) fn has_same_binding(&self, other: &Self) -> bool { - self.binding_id == other.binding_id - && self.chatgpt_account_id == other.chatgpt_account_id - && self.chatgpt_user_id == other.chatgpt_user_id - } -} - -fn sign_task_registration_payload( - stored_identity: &StoredAgentIdentity, - timestamp: &str, -) -> Result { - let signing_key = stored_identity.signing_key()?; - let payload = format!("{}:{timestamp}", stored_identity.agent_runtime_id); - Ok(BASE64_STANDARD.encode(signing_key.sign(payload.as_bytes()).to_bytes())) -} - -fn decrypt_task_id_response( - stored_identity: &StoredAgentIdentity, - encrypted_task_id: &str, -) -> Result { - let signing_key = stored_identity.signing_key()?; - let ciphertext = BASE64_STANDARD - .decode(encrypted_task_id) - .context("encrypted task id is not valid base64")?; - let plaintext = curve25519_secret_key_from_signing_key(&signing_key) - .unseal(&ciphertext) - .map_err(|_| anyhow::anyhow!("failed to decrypt encrypted task id"))?; - String::from_utf8(plaintext).context("decrypted task id is not valid UTF-8") -} - -fn curve25519_secret_key_from_signing_key(signing_key: &SigningKey) -> Curve25519SecretKey { - let digest = Sha512::digest(signing_key.to_bytes()); - let mut secret_key = [0u8; 32]; - secret_key.copy_from_slice(&digest[..32]); - secret_key[0] &= 248; - secret_key[31] &= 127; - secret_key[31] |= 64; - Curve25519SecretKey::from(secret_key) -} - -fn agent_task_registration_url(chatgpt_base_url: &str, agent_runtime_id: &str) -> String { - let trimmed = chatgpt_base_url.trim_end_matches('/'); - format!("{trimmed}/v1/agent/{agent_runtime_id}/task/register") -} - -#[cfg(test)] -mod tests { - use base64::engine::general_purpose::URL_SAFE_NO_PAD; - use codex_app_server_protocol::AuthMode as ApiAuthMode; - use codex_login::AuthCredentialsStoreMode; - use codex_login::AuthDotJson; - use codex_login::save_auth; - use codex_login::token_data::IdTokenInfo; - use codex_login::token_data::TokenData; - use pretty_assertions::assert_eq; - use wiremock::Mock; - use wiremock::MockServer; - use wiremock::ResponseTemplate; - use wiremock::matchers::header; - use wiremock::matchers::method; - use wiremock::matchers::path; - - use super::*; - - #[tokio::test] - async fn register_task_skips_when_feature_is_disabled() { - let auth_manager = - AuthManager::from_auth_for_testing(make_chatgpt_auth("account-123", Some("user-123"))); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ false, - "https://chatgpt.com/backend-api/".to_string(), - SessionSource::Cli, - ); - - assert_eq!(manager.register_task().await.unwrap(), None); - } - - #[tokio::test] - async fn register_task_skips_for_api_key_auth() { - let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test-key")); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - "https://chatgpt.com/backend-api/".to_string(), - SessionSource::Cli, - ); - - assert_eq!(manager.register_task().await.unwrap(), None); - } - - #[tokio::test] - async fn register_task_registers_and_decrypts_plaintext_task_id() { - let server = MockServer::start().await; - let chatgpt_base_url = server.uri(); - mount_human_biscuit(&server, &chatgpt_base_url, "agent-123").await; - let auth = make_chatgpt_auth("account-123", Some("user-123")); - let auth_manager = AuthManager::from_auth_for_testing(auth.clone()); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - chatgpt_base_url, - SessionSource::Cli, - ); - let stored_identity = seed_stored_identity(&manager, &auth, "agent-123", "account-123"); - let encrypted_task_id = - encrypt_task_id_for_identity(&stored_identity, "task_123").expect("task ciphertext"); - - Mock::given(method("POST")) - .and(path("/v1/agent/agent-123/task/register")) - .and(header("x-openai-authorization", "human-biscuit")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "encrypted_task_id": encrypted_task_id, - }))) - .expect(1) - .mount(&server) - .await; - - let task = manager - .register_task() - .await - .unwrap() - .expect("task should be registered"); - - assert_eq!( - task, - RegisteredAgentTask { - binding_id: "chatgpt-account-account-123".to_string(), - chatgpt_account_id: "account-123".to_string(), - chatgpt_user_id: Some("user-123".to_string()), - agent_runtime_id: "agent-123".to_string(), - task_id: "task_123".to_string(), - registered_at: task.registered_at.clone(), - } - ); - } - - #[tokio::test] - async fn register_task_uses_chatgpt_base_url() { - let server = MockServer::start().await; - let chatgpt_base_url = format!("{}/backend-api", server.uri()); - mount_human_biscuit(&server, &chatgpt_base_url, "agent-fallback").await; - let auth = make_chatgpt_auth("account-123", Some("user-123")); - let auth_manager = AuthManager::from_auth_for_testing(auth.clone()); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - chatgpt_base_url, - SessionSource::Cli, - ); - let stored_identity = - seed_stored_identity(&manager, &auth, "agent-fallback", "account-123"); - let encrypted_task_id = encrypt_task_id_for_identity(&stored_identity, "task_fallback") - .expect("task ciphertext"); - - Mock::given(method("POST")) - .and(path("/backend-api/v1/agent/agent-fallback/task/register")) - .and(header("x-openai-authorization", "human-biscuit")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "encrypted_task_id": encrypted_task_id, - }))) - .expect(1) - .mount(&server) - .await; - - let task = manager - .register_task() - .await - .unwrap() - .expect("task should be registered"); - - assert_eq!(task.agent_runtime_id, "agent-fallback"); - assert_eq!(task.task_id, "task_fallback"); - } - - #[tokio::test] - async fn register_task_for_binding_keeps_one_auth_snapshot() { - let server = MockServer::start().await; - let chatgpt_base_url = server.uri(); - mount_human_biscuit(&server, &chatgpt_base_url, "agent-123").await; - let binding_auth = make_chatgpt_auth("account-123", Some("user-123")); - let auth_manager = - AuthManager::from_auth_for_testing(make_chatgpt_auth("account-456", Some("user-456"))); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - chatgpt_base_url, - SessionSource::Cli, - ); - let stored_identity = - seed_stored_identity(&manager, &binding_auth, "agent-123", "account-123"); - let encrypted_task_id = - encrypt_task_id_for_identity(&stored_identity, "task_123").expect("task ciphertext"); - let binding = - AgentIdentityBinding::from_auth(&binding_auth, /*forced_workspace_id*/ None) - .expect("binding"); - - Mock::given(method("POST")) - .and(path("/v1/agent/agent-123/task/register")) - .and(header("x-openai-authorization", "human-biscuit")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "encrypted_task_id": encrypted_task_id, - }))) - .expect(1) - .mount(&server) - .await; - - let task = manager - .register_task_for_binding(binding_auth, binding) - .await - .unwrap() - .expect("task should be registered"); - - assert_eq!( - task, - RegisteredAgentTask { - binding_id: "chatgpt-account-account-123".to_string(), - chatgpt_account_id: "account-123".to_string(), - chatgpt_user_id: Some("user-123".to_string()), - agent_runtime_id: "agent-123".to_string(), - task_id: "task_123".to_string(), - registered_at: task.registered_at.clone(), - } - ); - } - - #[tokio::test] - async fn task_matches_current_binding_rejects_stale_auth_binding() { - let auth_manager = - AuthManager::from_auth_for_testing(make_chatgpt_auth("account-456", Some("user-456"))); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - "https://chatgpt.com/backend-api/".to_string(), - SessionSource::Cli, - ); - let task = RegisteredAgentTask { - binding_id: "chatgpt-account-account-123".to_string(), - chatgpt_account_id: "account-123".to_string(), - chatgpt_user_id: Some("user-123".to_string()), - agent_runtime_id: "agent-123".to_string(), - task_id: "task_123".to_string(), - registered_at: "2026-03-23T12:00:00Z".to_string(), - }; - - assert!(!manager.task_matches_current_binding(&task).await); - } - - async fn mount_human_biscuit( - server: &MockServer, - chatgpt_base_url: &str, - agent_runtime_id: &str, - ) { - let biscuit_url = agent_identity_biscuit_url(chatgpt_base_url); - let biscuit_path = reqwest::Url::parse(&biscuit_url) - .expect("biscuit URL parses") - .path() - .to_string(); - let target_url = agent_task_registration_url(chatgpt_base_url, agent_runtime_id); - Mock::given(method("GET")) - .and(path(biscuit_path)) - .and(header("authorization", "Bearer access-token-account-123")) - .and(header("x-original-method", "POST")) - .and(header("x-original-url", target_url)) - .respond_with( - ResponseTemplate::new(200).insert_header("x-openai-authorization", "human-biscuit"), - ) - .expect(1) - .mount(server) - .await; - } - - fn seed_stored_identity( - manager: &AgentIdentityManager, - auth: &CodexAuth, - agent_runtime_id: &str, - account_id: &str, - ) -> StoredAgentIdentity { - let key_material = generate_agent_key_material().expect("key material"); - let binding = - AgentIdentityBinding::from_auth(auth, /*forced_workspace_id*/ None).expect("binding"); - let stored_identity = StoredAgentIdentity { - binding_id: binding.binding_id, - chatgpt_account_id: account_id.to_string(), - chatgpt_user_id: Some("user-123".to_string()), - agent_runtime_id: agent_runtime_id.to_string(), - private_key_pkcs8_base64: key_material.private_key_pkcs8_base64, - public_key_ssh: key_material.public_key_ssh, - registered_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), - abom: manager.abom.clone(), - }; - manager - .store_identity(auth, &stored_identity) - .expect("store identity"); - let persisted = auth - .get_agent_identity(account_id) - .expect("persisted identity"); - assert_eq!(persisted.agent_runtime_id, agent_runtime_id); - stored_identity - } - - fn encrypt_task_id_for_identity( - stored_identity: &StoredAgentIdentity, - task_id: &str, - ) -> Result { - let mut rng = crypto_box::aead::OsRng; - let public_key = - curve25519_secret_key_from_signing_key(&stored_identity.signing_key()?).public_key(); - let ciphertext = public_key - .seal(&mut rng, task_id.as_bytes()) - .map_err(|_| anyhow::anyhow!("failed to encrypt test task id"))?; - Ok(BASE64_STANDARD.encode(ciphertext)) - } - - fn make_chatgpt_auth(account_id: &str, user_id: Option<&str>) -> CodexAuth { - let tempdir = tempfile::tempdir().expect("tempdir"); - let auth_json = AuthDotJson { - auth_mode: Some(ApiAuthMode::Chatgpt), - openai_api_key: None, - tokens: Some(TokenData { - id_token: IdTokenInfo { - email: None, - chatgpt_plan_type: None, - chatgpt_user_id: user_id.map(ToOwned::to_owned), - chatgpt_account_id: Some(account_id.to_string()), - chatgpt_account_is_fedramp: false, - raw_jwt: fake_id_token(account_id, user_id), - }, - access_token: format!("access-token-{account_id}"), - refresh_token: "refresh-token".to_string(), - account_id: Some(account_id.to_string()), - }), - last_refresh: Some(Utc::now()), - agent_identity: None, - }; - save_auth(tempdir.path(), &auth_json, AuthCredentialsStoreMode::File).expect("save auth"); - CodexAuth::from_auth_storage(tempdir.path(), AuthCredentialsStoreMode::File) - .expect("load auth") - .expect("auth") - } - - fn fake_id_token(account_id: &str, user_id: Option<&str>) -> String { - let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#); - let payload = serde_json::json!({ - "https://api.openai.com/auth": { - "chatgpt_user_id": user_id, - "chatgpt_account_id": account_id, - } - }); - let payload = URL_SAFE_NO_PAD.encode(payload.to_string()); - format!("{header}.{payload}.signature") - } -} diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 2ff4be2ed8a1..1f987738a85c 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4283,7 +4283,7 @@ pub(crate) async fn make_session_and_context_with_rx() -> ( } #[tokio::test] -async fn fail_agent_identity_registration_emits_error_without_shutdown() { +async fn fail_agent_identity_registration_emits_error_and_shutdown() { let (session, _turn_context, rx_event) = make_session_and_context_with_rx().await; session @@ -4301,14 +4301,21 @@ async fn fail_agent_identity_registration_emits_error_without_shutdown() { }) => { assert_eq!( message, - "Agent identity registration failed while `features.use_agent_identity` is enabled: registration exploded".to_string() + "Agent identity registration failed. Codex cannot continue while `features.use_agent_identity` is enabled: registration exploded".to_string() ); assert_eq!(codex_error_info, Some(CodexErrorInfo::Other)); } other => panic!("expected error event, got {other:?}"), } - assert!(rx_event.try_recv().is_err()); + let shutdown_event = timeout(Duration::from_secs(1), rx_event.recv()) + .await + .expect("shutdown event should arrive") + .expect("shutdown event should be readable"); + match shutdown_event.msg { + EventMsg::ShutdownComplete => {} + other => panic!("expected shutdown event, got {other:?}"), + } } #[tokio::test] diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index bf6e06d8a93b..e9fd67633299 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -6,7 +6,6 @@ use codex_sandboxing::policy_transforms::merge_permission_profiles; use std::collections::HashMap; use std::collections::HashSet; -use crate::agent_identity::RegisteredAgentTask; use crate::context_manager::ContextManager; use crate::session::PreviousTurnSettings; use crate::session::session::SessionConfiguration; @@ -31,7 +30,6 @@ pub(crate) struct SessionState { previous_turn_settings: Option, /// Startup prewarmed session prepared during session initialization. pub(crate) startup_prewarm: Option, - pub(crate) agent_task: Option, pub(crate) active_connector_selection: HashSet, pub(crate) pending_session_start_source: Option, granted_permissions: Option, @@ -51,7 +49,6 @@ impl SessionState { mcp_dependency_prompted: HashSet::new(), previous_turn_settings: None, startup_prewarm: None, - agent_task: None, active_connector_selection: HashSet::new(), pending_session_start_source: None, granted_permissions: None, @@ -189,18 +186,6 @@ impl SessionState { self.startup_prewarm.take() } - pub(crate) fn agent_task(&self) -> Option { - self.agent_task.clone() - } - - pub(crate) fn set_agent_task(&mut self, agent_task: RegisteredAgentTask) { - self.agent_task = Some(agent_task); - } - - pub(crate) fn clear_agent_task(&mut self) { - self.agent_task = None; - } - // Adds connector IDs to the active set and returns the merged selection. pub(crate) fn merge_connector_selection(&mut self, connector_ids: I) -> HashSet where diff --git a/codex-rs/core/src/state/session_tests.rs b/codex-rs/core/src/state/session_tests.rs index 927c1fbb70fd..5e90cc881dd2 100644 --- a/codex-rs/core/src/state/session_tests.rs +++ b/codex-rs/core/src/state/session_tests.rs @@ -1,5 +1,4 @@ use super::*; -use crate::agent_identity::RegisteredAgentTask; use crate::session::tests::make_session_configuration_for_tests; use codex_protocol::protocol::CreditsSnapshot; use codex_protocol::protocol::RateLimitWindow; @@ -34,43 +33,6 @@ async fn clear_connector_selection_removes_entries() { assert_eq!(state.get_connector_selection(), HashSet::new()); } -#[tokio::test] -async fn set_agent_task_persists_plaintext_task_for_session_reuse() { - let session_configuration = make_session_configuration_for_tests().await; - let mut state = SessionState::new(session_configuration); - let agent_task = RegisteredAgentTask { - binding_id: "chatgpt-account-account-123".to_string(), - chatgpt_account_id: "account-123".to_string(), - chatgpt_user_id: Some("user-123".to_string()), - agent_runtime_id: "agent_123".to_string(), - task_id: "task_123".to_string(), - registered_at: "2026-03-23T12:00:00Z".to_string(), - }; - - state.set_agent_task(agent_task.clone()); - - assert_eq!(state.agent_task(), Some(agent_task)); -} - -#[tokio::test] -async fn clear_agent_task_removes_cached_task() { - let session_configuration = make_session_configuration_for_tests().await; - let mut state = SessionState::new(session_configuration); - let agent_task = RegisteredAgentTask { - binding_id: "chatgpt-account-account-123".to_string(), - chatgpt_account_id: "account-123".to_string(), - chatgpt_user_id: Some("user-123".to_string()), - agent_runtime_id: "agent_123".to_string(), - task_id: "task_123".to_string(), - registered_at: "2026-03-23T12:00:00Z".to_string(), - }; - - state.set_agent_task(agent_task); - state.clear_agent_task(); - - assert_eq!(state.agent_task(), None); -} - #[tokio::test] async fn set_rate_limits_defaults_limit_id_to_codex_when_missing() { let session_configuration = make_session_configuration_for_tests().await; diff --git a/codex-rs/core/tests/suite/realtime_conversation.rs b/codex-rs/core/tests/suite/realtime_conversation.rs index 60167a15586a..e044a216f81f 100644 --- a/codex-rs/core/tests/suite/realtime_conversation.rs +++ b/codex-rs/core/tests/suite/realtime_conversation.rs @@ -15,7 +15,6 @@ use codex_protocol::protocol::ConversationStartTransport; use codex_protocol::protocol::ConversationTextParams; use codex_protocol::protocol::ErrorEvent; use codex_protocol::protocol::EventMsg; -use codex_protocol::protocol::GitInfo; use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::Op; use codex_protocol::protocol::RealtimeAudioFrame; @@ -26,11 +25,7 @@ use codex_protocol::protocol::RealtimeNoopRequested; use codex_protocol::protocol::RealtimeOutputModality; use codex_protocol::protocol::RealtimeVoice; use codex_protocol::protocol::RolloutItem; -use codex_protocol::protocol::RolloutLine; -use codex_protocol::protocol::SessionMeta; -use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; -use codex_protocol::protocol::UserMessageEvent; use codex_protocol::user_input::UserInput; use codex_utils_output_truncation::approx_token_count; use core_test_support::responses; @@ -71,7 +66,6 @@ const MEMORY_PROMPT_PHRASE: &str = "You have access to a memory folder with guidance from prior runs."; const REALTIME_CONVERSATION_TEST_SUBPROCESS_ENV_VAR: &str = "CODEX_REALTIME_CONVERSATION_TEST_SUBPROCESS"; -const WEBSOCKET_REQUEST_TIMEOUT: Duration = Duration::from_secs(30); #[derive(Debug, Clone)] struct RealtimeCallRequestCapture { @@ -148,7 +142,7 @@ async fn wait_for_matching_websocket_request( where F: Fn(&core_test_support::responses::WebSocketRequest) -> bool, { - let deadline = tokio::time::Instant::now() + WEBSOCKET_REQUEST_TIMEOUT; + let deadline = tokio::time::Instant::now() + Duration::from_secs(10); loop { if let Some(request) = server .connections() @@ -208,18 +202,16 @@ async fn seed_recent_thread( let db = test.codex.state_db().context("state db enabled")?; let thread_id = ThreadId::new(); let updated_at = Utc::now(); - let rollout_dir = test + let rollout_path = test .codex_home_path() - .join("sessions") - .join(updated_at.format("%Y/%m/%d").to_string()); - fs::create_dir_all(&rollout_dir)?; - let rollout_path = rollout_dir.join(format!( - "rollout-{}-{thread_id}.jsonl", - updated_at.format("%Y-%m-%dT%H-%M-%S") - )); + .join(format!("rollout-{thread_id}.jsonl")); + // This helper seeds SQLite metadata directly. Local listing drops stale metadata rows whose + // rollout path no longer exists, so create the placeholder path that the test metadata points + // at without exercising rollout writing in this realtime-context test. + std::fs::write(&rollout_path, "")?; let mut metadata_builder = codex_state::ThreadMetadataBuilder::new( thread_id, - rollout_path.clone(), + rollout_path, updated_at, SessionSource::Cli, ); @@ -229,45 +221,6 @@ async fn seed_recent_thread( let mut metadata = metadata_builder.build("test-provider"); metadata.title = title.to_string(); metadata.first_user_message = Some(first_user_message.to_string()); - - let timestamp = updated_at.to_rfc3339(); - let session_meta = RolloutLine { - timestamp: timestamp.clone(), - item: RolloutItem::SessionMeta(SessionMetaLine { - meta: SessionMeta { - id: thread_id, - timestamp: timestamp.clone(), - cwd: metadata.cwd.clone(), - originator: "cli".to_string(), - cli_version: "0.0.0".to_string(), - source: SessionSource::Cli, - model_provider: Some("test-provider".to_string()), - ..Default::default() - }, - git: Some(GitInfo { - commit_hash: None, - branch: metadata.git_branch.clone(), - repository_url: None, - }), - }), - }; - let user_message = RolloutLine { - timestamp, - item: RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent { - message: first_user_message.to_string(), - images: None, - local_images: Vec::new(), - text_elements: Vec::new(), - })), - }; - fs::write( - &rollout_path, - format!( - "{}\n{}\n", - serde_json::to_string(&session_meta)?, - serde_json::to_string(&user_message)? - ), - )?; db.upsert_thread(&metadata).await?; Ok(()) From 68f0c6bf42c5fb37a080449503ddd0588656b8d7 Mon Sep 17 00:00:00 2001 From: Edward Frazer Date: Tue, 21 Apr 2026 09:33:41 -0700 Subject: [PATCH 5/9] Revert "Register agent identities behind use_agent_identity (#17386)" This reverts commit 8e784bba2fb795ec9248bbfd9a26fdad1b809dc6. --- codex-rs/Cargo.toml | 1 - .../src/transport/remote_control/tests.rs | 1 - .../src/transport/remote_control/websocket.rs | 1 - .../app-server/tests/common/auth_fixtures.rs | 1 - .../app-server/tests/suite/v2/app_list.rs | 1 - codex-rs/core/BUILD.bazel | 1 - codex-rs/core/Cargo.toml | 2 - codex-rs/core/src/agent_identity.rs | 771 ------------------ codex-rs/core/src/lib.rs | 1 - codex-rs/core/src/session/tests.rs | 47 -- codex-rs/core/src/state/service.rs | 2 - .../suite/responses_api_proxy_headers.rs | 2 - codex-rs/login/src/auth/auth_tests.rs | 136 --- codex-rs/login/src/auth/manager.rs | 93 +-- codex-rs/login/src/auth/storage.rs | 13 - codex-rs/login/src/auth/storage_tests.rs | 31 - codex-rs/login/src/lib.rs | 2 +- codex-rs/login/src/server.rs | 1 - codex-rs/login/tests/suite/auth_refresh.rs | 22 - codex-rs/tui/src/local_chatgpt_auth.rs | 2 - 20 files changed, 10 insertions(+), 1121 deletions(-) delete mode 100644 codex-rs/core/src/agent_identity.rs diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 04fd2fa06c42..5da0b4a8575e 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -237,7 +237,6 @@ dirs = "6" dns-lookup = "3.0.1" dotenvy = "0.15.7" dunce = "1.0.4" -ed25519-dalek = { version = "2.2.0", features = ["pkcs8"] } encoding_rs = "0.8.35" env-flags = "0.1.1" env_logger = "0.11.9" diff --git a/codex-rs/app-server/src/transport/remote_control/tests.rs b/codex-rs/app-server/src/transport/remote_control/tests.rs index d702a99c2cf3..eb49116167eb 100644 --- a/codex-rs/app-server/src/transport/remote_control/tests.rs +++ b/codex-rs/app-server/src/transport/remote_control/tests.rs @@ -97,7 +97,6 @@ fn remote_control_auth_dot_json(account_id: Option<&str>) -> AuthDotJson { account_id: account_id.map(str::to_string), }), last_refresh: Some(chrono::Utc::now()), - agent_identity: None, } } diff --git a/codex-rs/app-server/src/transport/remote_control/websocket.rs b/codex-rs/app-server/src/transport/remote_control/websocket.rs index 523c93f3beed..b68c98e733ca 100644 --- a/codex-rs/app-server/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server/src/transport/remote_control/websocket.rs @@ -1011,7 +1011,6 @@ mod tests { account_id: Some("account_id".to_string()), }), last_refresh: Some(Utc::now()), - agent_identity: None, } } diff --git a/codex-rs/app-server/tests/common/auth_fixtures.rs b/codex-rs/app-server/tests/common/auth_fixtures.rs index 86f0fb456ddb..99334f07706f 100644 --- a/codex-rs/app-server/tests/common/auth_fixtures.rs +++ b/codex-rs/app-server/tests/common/auth_fixtures.rs @@ -163,7 +163,6 @@ pub fn write_chatgpt_auth( openai_api_key: None, tokens: Some(tokens), last_refresh, - agent_identity: None, }; save_auth(codex_home, &auth, cli_auth_credentials_store_mode).context("write auth.json") diff --git a/codex-rs/app-server/tests/suite/v2/app_list.rs b/codex-rs/app-server/tests/suite/v2/app_list.rs index 78a915d178bf..7aa803452805 100644 --- a/codex-rs/app-server/tests/suite/v2/app_list.rs +++ b/codex-rs/app-server/tests/suite/v2/app_list.rs @@ -119,7 +119,6 @@ async fn list_apps_returns_empty_with_api_key_auth() -> Result<()> { openai_api_key: Some("test-api-key".to_string()), tokens: None, last_refresh: None, - agent_identity: None, }, AuthCredentialsStoreMode::File, )?; diff --git a/codex-rs/core/BUILD.bazel b/codex-rs/core/BUILD.bazel index dd52bce43d3a..df5f4da1fabb 100644 --- a/codex-rs/core/BUILD.bazel +++ b/codex-rs/core/BUILD.bazel @@ -57,7 +57,6 @@ codex_rust_crate( "//codex-rs/linux-sandbox:codex-linux-sandbox", "//codex-rs/rmcp-client:test_stdio_server", "//codex-rs/rmcp-client:test_streamable_http_server", - "//codex-rs/responses-api-proxy:codex-responses-api-proxy", "//codex-rs/cli:codex", ], ) diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 5b003aa1d6b5..86d6f8b0f3ec 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -42,7 +42,6 @@ codex-login = { workspace = true } codex-mcp = { workspace = true } codex-model-provider-info = { workspace = true } codex-models-manager = { workspace = true } -ed25519-dalek = { workspace = true } codex-shell-command = { workspace = true } codex-execpolicy = { workspace = true } codex-git-utils = { workspace = true } @@ -145,7 +144,6 @@ codex-shell-escalation = { workspace = true } [dev-dependencies] assert_cmd = { workspace = true } assert_matches = { workspace = true } -codex-arg0 = { workspace = true } codex-otel = { workspace = true } codex-test-binary-support = { workspace = true } codex-utils-cargo-bin = { workspace = true } diff --git a/codex-rs/core/src/agent_identity.rs b/codex-rs/core/src/agent_identity.rs deleted file mode 100644 index 7e86ba22ae4a..000000000000 --- a/codex-rs/core/src/agent_identity.rs +++ /dev/null @@ -1,771 +0,0 @@ -use std::sync::Arc; -use std::time::Duration; - -use anyhow::Context; -use anyhow::Result; -use base64::Engine as _; -use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; -use base64::engine::general_purpose::URL_SAFE_NO_PAD; -use chrono::SecondsFormat; -use chrono::Utc; -use codex_features::Feature; -use codex_login::AgentIdentityAuthRecord; -use codex_login::AuthManager; -use codex_login::CodexAuth; -use codex_login::default_client::create_client; -use codex_protocol::protocol::SessionSource; -use ed25519_dalek::SigningKey; -use ed25519_dalek::VerifyingKey; -use ed25519_dalek::pkcs8::DecodePrivateKey; -use ed25519_dalek::pkcs8::EncodePrivateKey; -use rand::TryRngCore; -use rand::rngs::OsRng; -use serde::Deserialize; -use serde::Serialize; -use tokio::sync::Semaphore; -use tracing::debug; -use tracing::info; -use tracing::warn; - -use crate::config::Config; - -const AGENT_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15); -const AGENT_IDENTITY_BISCUIT_TIMEOUT: Duration = Duration::from_secs(15); - -#[derive(Clone)] -pub(crate) struct AgentIdentityManager { - auth_manager: Arc, - chatgpt_base_url: String, - feature_enabled: bool, - abom: AgentBillOfMaterials, - ensure_lock: Arc, -} - -impl std::fmt::Debug for AgentIdentityManager { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("AgentIdentityManager") - .field("chatgpt_base_url", &self.chatgpt_base_url) - .field("feature_enabled", &self.feature_enabled) - .field("abom", &self.abom) - .finish_non_exhaustive() - } -} - -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] -pub(crate) struct StoredAgentIdentity { - pub(crate) binding_id: String, - pub(crate) chatgpt_account_id: String, - pub(crate) chatgpt_user_id: Option, - pub(crate) agent_runtime_id: String, - pub(crate) private_key_pkcs8_base64: String, - pub(crate) public_key_ssh: String, - pub(crate) registered_at: String, - pub(crate) abom: AgentBillOfMaterials, -} - -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] -pub(crate) struct AgentBillOfMaterials { - pub(crate) agent_version: String, - pub(crate) agent_harness_id: String, - pub(crate) running_location: String, -} - -#[derive(Debug, Serialize)] -struct RegisterAgentRequest { - abom: AgentBillOfMaterials, - agent_public_key: String, - capabilities: Vec, -} - -#[derive(Debug, Deserialize)] -struct RegisterAgentResponse { - agent_runtime_id: String, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -struct AgentIdentityBinding { - binding_id: String, - chatgpt_account_id: String, - chatgpt_user_id: Option, - access_token: String, -} - -struct GeneratedAgentKeyMaterial { - private_key_pkcs8_base64: String, - public_key_ssh: String, -} - -impl AgentIdentityManager { - pub(crate) fn new( - config: &Config, - auth_manager: Arc, - session_source: SessionSource, - ) -> Self { - Self { - auth_manager, - chatgpt_base_url: config.chatgpt_base_url.clone(), - feature_enabled: config.features.enabled(Feature::UseAgentIdentity), - abom: build_abom(session_source), - ensure_lock: Arc::new(Semaphore::new(/*permits*/ 1)), - } - } - - pub(crate) fn is_enabled(&self) -> bool { - self.feature_enabled - } - - pub(crate) async fn ensure_registered_identity(&self) -> Result> { - if !self.feature_enabled { - return Ok(None); - } - - let Some(auth) = self.auth_manager.auth().await else { - debug!("skipping agent identity registration because no auth is available"); - return Ok(None); - }; - - let Some(binding) = - AgentIdentityBinding::from_auth(&auth, self.auth_manager.forced_chatgpt_workspace_id()) - else { - debug!("skipping agent identity registration because ChatGPT auth is unavailable"); - return Ok(None); - }; - - let _guard = self - .ensure_lock - .acquire() - .await - .map_err(|_| anyhow::anyhow!("agent identity ensure semaphore closed"))?; - - if let Some(stored_identity) = self.load_stored_identity(&auth, &binding)? { - info!( - agent_runtime_id = %stored_identity.agent_runtime_id, - binding_id = %binding.binding_id, - "reusing stored agent identity" - ); - return Ok(Some(stored_identity)); - } - - let stored_identity = self.register_agent_identity(&binding).await?; - self.store_identity(&auth, &stored_identity)?; - Ok(Some(stored_identity)) - } - - async fn register_agent_identity( - &self, - binding: &AgentIdentityBinding, - ) -> Result { - let key_material = generate_agent_key_material()?; - let request_body = RegisterAgentRequest { - abom: self.abom.clone(), - agent_public_key: key_material.public_key_ssh.clone(), - capabilities: Vec::new(), - }; - - let url = agent_registration_url(&self.chatgpt_base_url); - let human_biscuit = self.mint_human_biscuit(binding, "POST", &url).await?; - let client = create_client(); - let response = client - .post(&url) - .header("X-OpenAI-Authorization", human_biscuit) - .json(&request_body) - .timeout(AGENT_REGISTRATION_TIMEOUT) - .send() - .await - .with_context(|| { - format!("failed to send agent identity registration request to {url}") - })?; - - if response.status().is_success() { - let response_body = response - .json::() - .await - .with_context(|| format!("failed to parse agent identity response from {url}"))?; - let stored_identity = StoredAgentIdentity { - binding_id: binding.binding_id.clone(), - chatgpt_account_id: binding.chatgpt_account_id.clone(), - chatgpt_user_id: binding.chatgpt_user_id.clone(), - agent_runtime_id: response_body.agent_runtime_id, - private_key_pkcs8_base64: key_material.private_key_pkcs8_base64, - public_key_ssh: key_material.public_key_ssh, - registered_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), - abom: self.abom.clone(), - }; - info!( - agent_runtime_id = %stored_identity.agent_runtime_id, - binding_id = %binding.binding_id, - "registered agent identity" - ); - return Ok(stored_identity); - } - - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - anyhow::bail!("agent identity registration failed with status {status} from {url}: {body}") - } - - async fn mint_human_biscuit( - &self, - binding: &AgentIdentityBinding, - target_method: &str, - target_url: &str, - ) -> Result { - let url = agent_identity_biscuit_url(&self.chatgpt_base_url); - let request_id = agent_identity_request_id()?; - let client = create_client(); - let response = client - .get(&url) - .bearer_auth(&binding.access_token) - .header("X-Request-Id", request_id.clone()) - .header("X-Original-Method", target_method) - .header("X-Original-Url", target_url) - .timeout(AGENT_IDENTITY_BISCUIT_TIMEOUT) - .send() - .await - .with_context(|| format!("failed to send agent identity biscuit request to {url}"))?; - - if response.status().is_success() { - let human_biscuit = response - .headers() - .get("x-openai-authorization") - .context("agent identity biscuit response did not include x-openai-authorization")? - .to_str() - .context("agent identity biscuit response header was not valid UTF-8")? - .to_string(); - info!( - request_id = %request_id, - "minted human biscuit for agent identity registration" - ); - return Ok(human_biscuit); - } - - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - anyhow::bail!( - "agent identity biscuit minting failed with status {status} from {url}: {body}" - ) - } - - fn load_stored_identity( - &self, - auth: &CodexAuth, - binding: &AgentIdentityBinding, - ) -> Result> { - let Some(record) = auth.get_agent_identity(&binding.chatgpt_account_id) else { - return Ok(None); - }; - - let stored_identity = - match StoredAgentIdentity::from_auth_record(binding, record, self.abom.clone()) { - Ok(stored_identity) => stored_identity, - Err(error) => { - warn!( - binding_id = %binding.binding_id, - error = %error, - "stored agent identity is invalid; deleting cached value" - ); - auth.remove_agent_identity()?; - return Ok(None); - } - }; - - if !stored_identity.matches_binding(binding) { - warn!( - binding_id = %binding.binding_id, - "stored agent identity binding no longer matches current auth; deleting cached value" - ); - auth.remove_agent_identity()?; - return Ok(None); - } - - if let Err(error) = stored_identity.validate_key_material() { - warn!( - agent_runtime_id = %stored_identity.agent_runtime_id, - binding_id = %binding.binding_id, - error = %error, - "stored agent identity key material is invalid; deleting cached value" - ); - auth.remove_agent_identity()?; - return Ok(None); - } - - Ok(Some(stored_identity)) - } - - fn store_identity( - &self, - auth: &CodexAuth, - stored_identity: &StoredAgentIdentity, - ) -> Result<()> { - auth.set_agent_identity(stored_identity.to_auth_record())?; - Ok(()) - } - - #[cfg(test)] - fn new_for_tests( - auth_manager: Arc, - feature_enabled: bool, - chatgpt_base_url: String, - session_source: SessionSource, - ) -> Self { - Self { - auth_manager, - chatgpt_base_url, - feature_enabled, - abom: build_abom(session_source), - ensure_lock: Arc::new(Semaphore::new(/*permits*/ 1)), - } - } -} - -impl StoredAgentIdentity { - fn from_auth_record( - binding: &AgentIdentityBinding, - record: AgentIdentityAuthRecord, - abom: AgentBillOfMaterials, - ) -> Result { - if record.workspace_id != binding.chatgpt_account_id { - anyhow::bail!( - "stored agent identity workspace {:?} does not match current workspace {:?}", - record.workspace_id, - binding.chatgpt_account_id - ); - } - let signing_key = signing_key_from_private_key_pkcs8_base64(&record.agent_private_key)?; - Ok(Self { - binding_id: binding.binding_id.clone(), - chatgpt_account_id: binding.chatgpt_account_id.clone(), - chatgpt_user_id: record.chatgpt_user_id, - agent_runtime_id: record.agent_runtime_id, - private_key_pkcs8_base64: record.agent_private_key, - public_key_ssh: encode_ssh_ed25519_public_key(&signing_key.verifying_key()), - registered_at: record.registered_at, - abom, - }) - } - - fn to_auth_record(&self) -> AgentIdentityAuthRecord { - AgentIdentityAuthRecord { - workspace_id: self.chatgpt_account_id.clone(), - chatgpt_user_id: self.chatgpt_user_id.clone(), - agent_runtime_id: self.agent_runtime_id.clone(), - agent_private_key: self.private_key_pkcs8_base64.clone(), - registered_at: self.registered_at.clone(), - } - } - - fn matches_binding(&self, binding: &AgentIdentityBinding) -> bool { - self.binding_id == binding.binding_id - && self.chatgpt_account_id == binding.chatgpt_account_id - && match binding.chatgpt_user_id.as_deref() { - Some(chatgpt_user_id) => self.chatgpt_user_id.as_deref() == Some(chatgpt_user_id), - None => true, - } - } - - fn validate_key_material(&self) -> Result<()> { - let signing_key = self.signing_key()?; - let derived_public_key = encode_ssh_ed25519_public_key(&signing_key.verifying_key()); - anyhow::ensure!( - self.public_key_ssh == derived_public_key, - "stored public key does not match the private key" - ); - Ok(()) - } - - pub(crate) fn signing_key(&self) -> Result { - signing_key_from_private_key_pkcs8_base64(&self.private_key_pkcs8_base64) - } -} - -impl AgentIdentityBinding { - fn from_auth(auth: &CodexAuth, forced_workspace_id: Option) -> Option { - if !auth.is_chatgpt_auth() { - return None; - } - - let token_data = auth.get_token_data().ok()?; - let resolved_account_id = - forced_workspace_id - .filter(|value| !value.is_empty()) - .or(token_data - .account_id - .clone() - .filter(|value| !value.is_empty()))?; - - Some(Self { - binding_id: format!("chatgpt-account-{resolved_account_id}"), - chatgpt_account_id: resolved_account_id, - chatgpt_user_id: token_data - .id_token - .chatgpt_user_id - .filter(|value| !value.is_empty()), - access_token: token_data.access_token, - }) - } -} - -fn build_abom(session_source: SessionSource) -> AgentBillOfMaterials { - AgentBillOfMaterials { - agent_version: env!("CARGO_PKG_VERSION").to_string(), - agent_harness_id: match &session_source { - SessionSource::VSCode => "codex-app".to_string(), - SessionSource::Cli - | SessionSource::Exec - | SessionSource::Mcp - | SessionSource::Custom(_) - | SessionSource::SubAgent(_) - | SessionSource::Unknown => "codex-cli".to_string(), - }, - running_location: format!("{}-{}", session_source, std::env::consts::OS), - } -} - -fn generate_agent_key_material() -> Result { - let mut secret_key_bytes = [0u8; 32]; - OsRng - .try_fill_bytes(&mut secret_key_bytes) - .context("failed to generate agent identity private key bytes")?; - let signing_key = SigningKey::from_bytes(&secret_key_bytes); - let private_key_pkcs8 = signing_key - .to_pkcs8_der() - .context("failed to encode agent identity private key as PKCS#8")?; - - Ok(GeneratedAgentKeyMaterial { - private_key_pkcs8_base64: BASE64_STANDARD.encode(private_key_pkcs8.as_bytes()), - public_key_ssh: encode_ssh_ed25519_public_key(&signing_key.verifying_key()), - }) -} - -fn encode_ssh_ed25519_public_key(verifying_key: &VerifyingKey) -> String { - let mut blob = Vec::with_capacity(4 + 11 + 4 + 32); - append_ssh_string(&mut blob, b"ssh-ed25519"); - append_ssh_string(&mut blob, verifying_key.as_bytes()); - format!("ssh-ed25519 {}", BASE64_STANDARD.encode(blob)) -} - -fn append_ssh_string(buf: &mut Vec, value: &[u8]) { - buf.extend_from_slice(&(value.len() as u32).to_be_bytes()); - buf.extend_from_slice(value); -} - -fn agent_registration_url(chatgpt_base_url: &str) -> String { - let trimmed = chatgpt_base_url.trim_end_matches('/'); - format!("{trimmed}/v1/agent/register") -} - -fn signing_key_from_private_key_pkcs8_base64(private_key_pkcs8_base64: &str) -> Result { - let private_key = BASE64_STANDARD - .decode(private_key_pkcs8_base64) - .context("stored agent identity private key is not valid base64")?; - SigningKey::from_pkcs8_der(&private_key) - .context("stored agent identity private key is not valid PKCS#8") -} - -fn agent_identity_biscuit_url(chatgpt_base_url: &str) -> String { - let trimmed = chatgpt_base_url.trim_end_matches('/'); - format!("{trimmed}/authenticate_app_v2") -} - -fn agent_identity_request_id() -> Result { - let mut request_id_bytes = [0u8; 16]; - OsRng - .try_fill_bytes(&mut request_id_bytes) - .context("failed to generate agent identity request id")?; - Ok(format!( - "codex-agent-identity-{}", - URL_SAFE_NO_PAD.encode(request_id_bytes) - )) -} - -#[cfg(test)] -mod tests { - use super::*; - - use base64::engine::general_purpose::URL_SAFE_NO_PAD; - use codex_app_server_protocol::AuthMode as ApiAuthMode; - use codex_login::AuthCredentialsStoreMode; - use codex_login::AuthDotJson; - use codex_login::save_auth; - use codex_login::token_data::IdTokenInfo; - use codex_login::token_data::TokenData; - use pretty_assertions::assert_eq; - use wiremock::Mock; - use wiremock::MockServer; - use wiremock::ResponseTemplate; - use wiremock::matchers::header; - use wiremock::matchers::method; - use wiremock::matchers::path; - - #[tokio::test] - async fn ensure_registered_identity_skips_when_feature_is_disabled() { - let auth_manager = - AuthManager::from_auth_for_testing(make_chatgpt_auth("account-123", Some("user-123"))); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ false, - "https://chatgpt.com/backend-api/".to_string(), - SessionSource::Cli, - ); - - assert_eq!(manager.ensure_registered_identity().await.unwrap(), None); - } - - #[tokio::test] - async fn ensure_registered_identity_skips_for_api_key_auth() { - let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test-key")); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - "https://chatgpt.com/backend-api/".to_string(), - SessionSource::Cli, - ); - - assert_eq!(manager.ensure_registered_identity().await.unwrap(), None); - } - - #[tokio::test] - async fn ensure_registered_identity_registers_and_reuses_cached_identity() { - let server = MockServer::start().await; - let chatgpt_base_url = server.uri(); - mount_human_biscuit(&server, &chatgpt_base_url).await; - Mock::given(method("POST")) - .and(path("/v1/agent/register")) - .and(header("x-openai-authorization", "human-biscuit")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "agent_runtime_id": "agent_123", - }))) - .expect(1) - .mount(&server) - .await; - - let auth_manager = - AuthManager::from_auth_for_testing(make_chatgpt_auth("account-123", Some("user-123"))); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - chatgpt_base_url, - SessionSource::Cli, - ); - - let first = manager - .ensure_registered_identity() - .await - .unwrap() - .expect("identity should be registered"); - let second = manager - .ensure_registered_identity() - .await - .unwrap() - .expect("identity should be reused"); - - assert_eq!(first.agent_runtime_id, "agent_123"); - assert_eq!(first, second); - assert_eq!(first.abom.agent_harness_id, "codex-cli"); - assert_eq!(first.chatgpt_account_id, "account-123"); - assert_eq!(first.chatgpt_user_id.as_deref(), Some("user-123")); - } - - #[tokio::test] - async fn ensure_registered_identity_deletes_invalid_cached_identity_and_reregisters() { - let server = MockServer::start().await; - let chatgpt_base_url = server.uri(); - mount_human_biscuit(&server, &chatgpt_base_url).await; - Mock::given(method("POST")) - .and(path("/v1/agent/register")) - .and(header("x-openai-authorization", "human-biscuit")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "agent_runtime_id": "agent_456", - }))) - .expect(1) - .mount(&server) - .await; - - let auth = make_chatgpt_auth("account-123", Some("user-123")); - let auth_manager = AuthManager::from_auth_for_testing(auth.clone()); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - chatgpt_base_url, - SessionSource::Cli, - ); - - let binding = - AgentIdentityBinding::from_auth(&auth, /*forced_workspace_id*/ None).expect("binding"); - auth.set_agent_identity(AgentIdentityAuthRecord { - workspace_id: "account-123".to_string(), - chatgpt_user_id: Some("user-123".to_string()), - agent_runtime_id: "agent_invalid".to_string(), - agent_private_key: "not-valid-base64".to_string(), - registered_at: "2026-01-01T00:00:00Z".to_string(), - }) - .expect("seed invalid identity"); - - let stored = manager - .ensure_registered_identity() - .await - .unwrap() - .expect("identity should be registered"); - - assert_eq!(stored.agent_runtime_id, "agent_456"); - let persisted = auth - .get_agent_identity(&binding.chatgpt_account_id) - .expect("stored identity"); - assert_eq!(persisted.agent_runtime_id, "agent_456"); - } - - #[tokio::test] - async fn ensure_registered_identity_deletes_different_user_identity_and_reregisters() { - let server = MockServer::start().await; - let chatgpt_base_url = server.uri(); - mount_human_biscuit(&server, &chatgpt_base_url).await; - Mock::given(method("POST")) - .and(path("/v1/agent/register")) - .and(header("x-openai-authorization", "human-biscuit")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "agent_runtime_id": "agent_new", - }))) - .expect(1) - .mount(&server) - .await; - - let auth = make_chatgpt_auth("account-123", Some("user-new")); - let stale_key = generate_agent_key_material().expect("key material"); - auth.set_agent_identity(AgentIdentityAuthRecord { - workspace_id: "account-123".to_string(), - chatgpt_user_id: Some("user-old".to_string()), - agent_runtime_id: "agent_old".to_string(), - agent_private_key: stale_key.private_key_pkcs8_base64, - registered_at: "2026-01-01T00:00:00Z".to_string(), - }) - .expect("seed stale identity"); - - let auth_manager = AuthManager::from_auth_for_testing(auth.clone()); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - chatgpt_base_url, - SessionSource::Cli, - ); - - let stored = manager - .ensure_registered_identity() - .await - .unwrap() - .expect("identity should be registered"); - - assert_eq!(stored.agent_runtime_id, "agent_new"); - assert_eq!(stored.chatgpt_user_id.as_deref(), Some("user-new")); - let persisted = auth - .get_agent_identity("account-123") - .expect("stored identity"); - assert_eq!(persisted.agent_runtime_id, "agent_new"); - assert_eq!(persisted.chatgpt_user_id.as_deref(), Some("user-new")); - } - - #[tokio::test] - async fn ensure_registered_identity_uses_chatgpt_base_url() { - let server = MockServer::start().await; - let chatgpt_base_url = format!("{}/backend-api", server.uri()); - mount_human_biscuit(&server, &chatgpt_base_url).await; - Mock::given(method("POST")) - .and(path("/backend-api/v1/agent/register")) - .and(header("x-openai-authorization", "human-biscuit")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "agent_runtime_id": "agent_canonical", - }))) - .expect(1) - .mount(&server) - .await; - - let auth_manager = - AuthManager::from_auth_for_testing(make_chatgpt_auth("account-123", Some("user-123"))); - let manager = AgentIdentityManager::new_for_tests( - auth_manager, - /*feature_enabled*/ true, - chatgpt_base_url, - SessionSource::Cli, - ); - - let stored = manager - .ensure_registered_identity() - .await - .unwrap() - .expect("identity should be registered"); - assert_eq!(stored.agent_runtime_id, "agent_canonical"); - } - - async fn mount_human_biscuit(server: &MockServer, chatgpt_base_url: &str) { - let biscuit_url = agent_identity_biscuit_url(chatgpt_base_url); - let biscuit_path = reqwest::Url::parse(&biscuit_url) - .expect("biscuit URL parses") - .path() - .to_string(); - let target_url = agent_registration_url(chatgpt_base_url); - Mock::given(method("GET")) - .and(path(biscuit_path)) - .and(header("authorization", "Bearer access-token-account-123")) - .and(header("x-original-method", "POST")) - .and(header("x-original-url", target_url)) - .respond_with( - ResponseTemplate::new(200).insert_header("x-openai-authorization", "human-biscuit"), - ) - .expect(1) - .mount(server) - .await; - } - - #[test] - fn encode_ssh_ed25519_public_key_matches_expected_wire_shape() { - let key_material = generate_agent_key_material().expect("key material"); - let (_, encoded_blob) = key_material - .public_key_ssh - .split_once(' ') - .expect("public key contains scheme"); - let decoded = BASE64_STANDARD.decode(encoded_blob).expect("base64"); - - assert_eq!(&decoded[..4], 11u32.to_be_bytes().as_slice()); - assert_eq!(&decoded[4..15], b"ssh-ed25519"); - assert_eq!(&decoded[15..19], 32u32.to_be_bytes().as_slice()); - assert_eq!(decoded.len(), 51); - } - - fn make_chatgpt_auth(account_id: &str, user_id: Option<&str>) -> CodexAuth { - let tempdir = tempfile::tempdir().expect("tempdir"); - let auth_json = AuthDotJson { - auth_mode: Some(ApiAuthMode::Chatgpt), - openai_api_key: None, - tokens: Some(TokenData { - id_token: IdTokenInfo { - email: None, - chatgpt_plan_type: None, - chatgpt_user_id: user_id.map(ToOwned::to_owned), - chatgpt_account_id: Some(account_id.to_string()), - chatgpt_account_is_fedramp: false, - raw_jwt: fake_id_token(account_id, user_id), - }, - access_token: format!("access-token-{account_id}"), - refresh_token: "refresh-token".to_string(), - account_id: Some(account_id.to_string()), - }), - last_refresh: Some(Utc::now()), - agent_identity: None, - }; - save_auth(tempdir.path(), &auth_json, AuthCredentialsStoreMode::File).expect("save auth"); - CodexAuth::from_auth_storage(tempdir.path(), AuthCredentialsStoreMode::File) - .expect("load auth") - .expect("auth") - } - - fn fake_id_token(account_id: &str, user_id: Option<&str>) -> String { - let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#); - let payload = serde_json::json!({ - "https://api.openai.com/auth": { - "chatgpt_user_id": user_id, - "chatgpt_account_id": account_id, - } - }); - let payload = URL_SAFE_NO_PAD.encode(payload.to_string()); - format!("{header}.{payload}.signature") - } -} diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index cf61c5faf93f..373e57737be0 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -5,7 +5,6 @@ // the TUI or the tracing stack). #![deny(clippy::print_stdout, clippy::print_stderr)] -mod agent_identity; mod apply_patch; mod apps; mod arc_monitor; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 1f987738a85c..1bc55cbbce79 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -130,7 +130,6 @@ use std::path::Path; use std::time::Duration; use tokio::sync::Semaphore; use tokio::time::sleep; -use tokio::time::timeout; use tracing_opentelemetry::OpenTelemetrySpanExt; use codex_protocol::mcp::CallToolResult as McpCallToolResult; @@ -3074,11 +3073,6 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { }), rollout: Mutex::new(None), user_shell: Arc::new(default_user_shell()), - agent_identity_manager: Arc::new(crate::agent_identity::AgentIdentityManager::new( - config.as_ref(), - Arc::clone(&auth_manager), - session_configuration.session_source.clone(), - )), shell_snapshot_tx: watch::channel(None).0, show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, @@ -4170,11 +4164,6 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( }), rollout: Mutex::new(None), user_shell: Arc::new(default_user_shell()), - agent_identity_manager: Arc::new(crate::agent_identity::AgentIdentityManager::new( - config.as_ref(), - Arc::clone(&auth_manager), - session_configuration.session_source.clone(), - )), shell_snapshot_tx: watch::channel(None).0, show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, @@ -4282,42 +4271,6 @@ pub(crate) async fn make_session_and_context_with_rx() -> ( make_session_and_context_with_dynamic_tools_and_rx(Vec::new()).await } -#[tokio::test] -async fn fail_agent_identity_registration_emits_error_and_shutdown() { - let (session, _turn_context, rx_event) = make_session_and_context_with_rx().await; - - session - .fail_agent_identity_registration(anyhow::anyhow!("registration exploded")) - .await; - - let error_event = timeout(Duration::from_secs(1), rx_event.recv()) - .await - .expect("error event should arrive") - .expect("error event should be readable"); - match error_event.msg { - EventMsg::Error(ErrorEvent { - message, - codex_error_info, - }) => { - assert_eq!( - message, - "Agent identity registration failed. Codex cannot continue while `features.use_agent_identity` is enabled: registration exploded".to_string() - ); - assert_eq!(codex_error_info, Some(CodexErrorInfo::Other)); - } - other => panic!("expected error event, got {other:?}"), - } - - let shutdown_event = timeout(Duration::from_secs(1), rx_event.recv()) - .await - .expect("shutdown event should arrive") - .expect("shutdown event should be readable"); - match shutdown_event.msg { - EventMsg::ShutdownComplete => {} - other => panic!("expected shutdown event, got {other:?}"), - } -} - #[tokio::test] async fn refresh_mcp_servers_is_deferred_until_next_turn() { let (session, turn_context) = make_session_and_context().await; diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 5db38f7b72a0..aae02d61bdc6 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -4,7 +4,6 @@ use std::sync::Arc; use crate::RolloutRecorder; use crate::SkillsManager; use crate::agent::AgentControl; -use crate::agent_identity::AgentIdentityManager; use crate::client::ModelClient; use crate::config::StartedNetworkProxy; use crate::exec_policy::ExecPolicyManager; @@ -43,7 +42,6 @@ pub(crate) struct SessionServices { pub(crate) hooks: Hooks, pub(crate) rollout: Mutex>, pub(crate) user_shell: Arc, - pub(crate) agent_identity_manager: Arc, pub(crate) shell_snapshot_tx: watch::Sender>>, pub(crate) show_raw_agent_reasoning: bool, pub(crate) exec_policy: Arc, diff --git a/codex-rs/core/tests/suite/responses_api_proxy_headers.rs b/codex-rs/core/tests/suite/responses_api_proxy_headers.rs index 67f6e86ab0e2..13d3b305d3d1 100644 --- a/codex-rs/core/tests/suite/responses_api_proxy_headers.rs +++ b/codex-rs/core/tests/suite/responses_api_proxy_headers.rs @@ -19,7 +19,6 @@ use core_test_support::responses::mount_sse_once_match; use core_test_support::responses::sse; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; -use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::test_codex; use pretty_assertions::assert_eq; use serde_json::json; @@ -30,7 +29,6 @@ const CHILD_PROMPT: &str = "child: say done"; const SPAWN_CALL_ID: &str = "spawn-call-1"; const REQUEST_POLL_INTERVAL: Duration = Duration::from_millis(/*millis*/ 20); const TURN_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 60); - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn responses_api_parent_and_subagent_requests_include_identity_headers() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index fed497307077..c008b0fc86b6 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -16,8 +16,6 @@ use serde_json::json; use std::sync::Arc; use tempfile::TempDir; use tempfile::tempdir; -use tokio::time::Duration; -use tokio::time::timeout; #[tokio::test] async fn refresh_without_id_token() { @@ -138,7 +136,6 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() { account_id: None, }), last_refresh: Some(last_refresh), - agent_identity: None, }, auth_dot_json ); @@ -176,7 +173,6 @@ fn logout_removes_auth_file() -> Result<(), std::io::Error> { openai_api_key: Some("sk-test-key".to_string()), tokens: None, last_refresh: None, - agent_identity: None, }; super::save_auth(dir.path(), &auth_dot_json, AuthCredentialsStoreMode::File)?; let auth_file = get_auth_file(dir.path()); @@ -186,77 +182,6 @@ fn logout_removes_auth_file() -> Result<(), std::io::Error> { Ok(()) } -#[test] -fn chatgpt_auth_persists_agent_identity_for_workspace() { - let codex_home = tempdir().unwrap(); - write_auth_file( - AuthFileParams { - openai_api_key: None, - chatgpt_plan_type: Some("pro".to_string()), - chatgpt_account_id: Some("account-123".to_string()), - }, - codex_home.path(), - ) - .expect("failed to write auth file"); - let auth = super::load_auth( - codex_home.path(), - /*enable_codex_api_key_env*/ false, - AuthCredentialsStoreMode::File, - ) - .expect("load auth") - .expect("auth available"); - let record = AgentIdentityAuthRecord { - workspace_id: "account-123".to_string(), - chatgpt_user_id: Some("user-123".to_string()), - agent_runtime_id: "agent_123".to_string(), - agent_private_key: "pkcs8-base64".to_string(), - registered_at: "2026-04-13T12:00:00Z".to_string(), - }; - - auth.set_agent_identity(record.clone()) - .expect("set agent identity"); - - assert_eq!(auth.get_agent_identity("account-123"), Some(record.clone())); - assert_eq!(auth.get_agent_identity("other-account"), None); - let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); - let persisted = storage - .load() - .expect("load auth") - .expect("auth should exist"); - assert_eq!(persisted.agent_identity, Some(record)); - - assert!(auth.remove_agent_identity().expect("remove agent identity")); - assert_eq!(auth.get_agent_identity("account-123"), None); -} - -#[test] -fn dummy_chatgpt_auth_does_not_create_cwd_auth_json_when_identity_is_set() { - let cwd_auth = std::env::current_dir() - .expect("current dir") - .join("auth.json"); - let had_auth_json_before_test = cwd_auth.exists(); - let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); - let record = AgentIdentityAuthRecord { - workspace_id: "account_id".to_string(), - chatgpt_user_id: None, - agent_runtime_id: "agent_123".to_string(), - agent_private_key: "pkcs8-base64".to_string(), - registered_at: "2026-04-13T12:00:00Z".to_string(), - background_task_id: None, - }; - - auth.set_agent_identity(record.clone()) - .expect("set agent identity"); - - assert_eq!(auth.get_agent_identity("account_id"), Some(record)); - if !had_auth_json_before_test { - assert!( - !cwd_auth.exists(), - "dummy ChatGPT auth must not write auth.json in the test process cwd" - ); - } -} - #[test] fn unauthorized_recovery_reports_mode_and_step_names() { let dir = tempdir().unwrap(); @@ -550,67 +475,6 @@ exit 1 } } -#[tokio::test] -async fn auth_manager_notifies_when_auth_state_changes() { - let dir = tempdir().unwrap(); - let manager = AuthManager::shared( - dir.path().to_path_buf(), - /*enable_codex_api_key_env*/ false, - AuthCredentialsStoreMode::File, - ); - let mut auth_state_rx = manager.subscribe_auth_state(); - - save_auth( - dir.path(), - &AuthDotJson { - auth_mode: Some(ApiAuthMode::ApiKey), - openai_api_key: Some("sk-test-key".to_string()), - tokens: None, - last_refresh: None, - agent_identity: None, - }, - AuthCredentialsStoreMode::File, - ) - .expect("save auth"); - - assert!( - manager.reload(), - "reload should report a changed auth state" - ); - timeout(Duration::from_secs(1), auth_state_rx.changed()) - .await - .expect("auth change notification should arrive") - .expect("auth state watch should remain open"); - - save_auth( - dir.path(), - &AuthDotJson { - auth_mode: Some(ApiAuthMode::ApiKey), - openai_api_key: Some("sk-updated-key".to_string()), - tokens: None, - last_refresh: None, - agent_identity: None, - }, - AuthCredentialsStoreMode::File, - ) - .expect("save updated auth"); - - assert!( - !manager.reload(), - "reload remains mode-stable even when the underlying credentials change" - ); - timeout(Duration::from_secs(1), auth_state_rx.changed()) - .await - .expect("auth reload notification should still arrive") - .expect("auth state watch should remain open"); - - manager.set_forced_chatgpt_workspace_id(Some("workspace-123".to_string())); - timeout(Duration::from_secs(1), auth_state_rx.changed()) - .await - .expect("workspace change notification should arrive") - .expect("auth state watch should remain open"); -} - struct AuthFileParams { openai_api_key: Option, chatgpt_plan_type: Option, diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 9ead5ce8b27c..e6bdb64a6be3 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -24,7 +24,6 @@ use codex_protocol::config_types::ModelProviderAuthInfo; use super::external_bearer::BearerTokenRefresher; use super::revoke::revoke_auth_tokens; -pub use crate::auth::storage::AgentIdentityAuthRecord; pub use crate::auth::storage::AuthDotJson; use crate::auth::storage::AuthStorageBackend; use crate::auth::storage::create_auth_storage; @@ -65,7 +64,6 @@ pub struct ChatgptAuth { #[derive(Debug, Clone)] pub struct ChatgptAuthTokens { state: ChatgptAuthState, - storage: Arc, } #[derive(Debug, Clone)] @@ -213,13 +211,14 @@ impl CodexAuth { client, }; - let storage = create_auth_storage(codex_home.to_path_buf(), storage_mode); match auth_mode { - ApiAuthMode::Chatgpt => Ok(Self::Chatgpt(ChatgptAuth { state, storage })), - ApiAuthMode::ChatgptAuthTokens => Ok(Self::ChatgptAuthTokens(ChatgptAuthTokens { - state, - storage, - })), + ApiAuthMode::Chatgpt => { + let storage = create_auth_storage(codex_home.to_path_buf(), storage_mode); + Ok(Self::Chatgpt(ChatgptAuth { state, storage })) + } + ApiAuthMode::ChatgptAuthTokens => { + Ok(Self::ChatgptAuthTokens(ChatgptAuthTokens { state })) + } ApiAuthMode::ApiKey => unreachable!("api key mode is handled above"), } } @@ -366,52 +365,6 @@ impl CodexAuth { self.get_current_auth_json().and_then(|t| t.tokens) } - pub fn get_agent_identity(&self, workspace_id: &str) -> Option { - self.get_current_auth_json() - .and_then(|auth| auth.agent_identity) - .filter(|identity| identity.workspace_id == workspace_id) - } - - pub fn set_agent_identity(&self, record: AgentIdentityAuthRecord) -> std::io::Result<()> { - let (state, storage) = match self { - Self::Chatgpt(auth) => (&auth.state, &auth.storage), - Self::ChatgptAuthTokens(auth) => (&auth.state, &auth.storage), - Self::ApiKey(_) => return Ok(()), - }; - let mut guard = state - .auth_dot_json - .lock() - .map_err(|_| std::io::Error::other("failed to lock auth state"))?; - let mut auth = guard - .clone() - .ok_or_else(|| std::io::Error::other("auth data is not available"))?; - auth.agent_identity = Some(record); - storage.save(&auth)?; - *guard = Some(auth); - Ok(()) - } - - pub fn remove_agent_identity(&self) -> std::io::Result { - let (state, storage) = match self { - Self::Chatgpt(auth) => (&auth.state, &auth.storage), - Self::ChatgptAuthTokens(auth) => (&auth.state, &auth.storage), - Self::ApiKey(_) => return Ok(false), - }; - let mut guard = state - .auth_dot_json - .lock() - .map_err(|_| std::io::Error::other("failed to lock auth state"))?; - let Some(mut auth) = guard.clone() else { - return Ok(false); - }; - let removed = auth.agent_identity.take().is_some(); - if removed { - storage.save(&auth)?; - *guard = Some(auth); - } - Ok(removed) - } - /// Consider this private to integration tests. pub fn create_dummy_chatgpt_auth_for_testing() -> Self { let auth_dot_json = AuthDotJson { @@ -424,7 +377,6 @@ impl CodexAuth { account_id: Some("account_id".to_string()), }), last_refresh: Some(Utc::now()), - agent_identity: None, }; let client = create_client(); @@ -517,7 +469,6 @@ pub fn login_with_api_key( openai_api_key: Some(api_key.to_string()), tokens: None, last_refresh: None, - agent_identity: None, }; save_auth(codex_home, &auth_dot_json, auth_credentials_store_mode) } @@ -891,7 +842,6 @@ impl AuthDotJson { openai_api_key: None, tokens: Some(tokens), last_refresh: Some(Utc::now()), - agent_identity: None, }) } @@ -1187,7 +1137,6 @@ pub struct AuthManager { forced_chatgpt_workspace_id: RwLock>, refresh_lock: Semaphore, external_auth: RwLock>>, - auth_state_tx: watch::Sender<()>, } /// Configuration view required to construct a shared [`AuthManager`]. @@ -1236,7 +1185,6 @@ impl AuthManager { enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, ) -> Self { - let (auth_state_tx, _) = watch::channel(()); let managed_auth = load_auth( &codex_home, enable_codex_api_key_env, @@ -1255,13 +1203,11 @@ impl AuthManager { forced_chatgpt_workspace_id: RwLock::new(None), refresh_lock: Semaphore::new(/*permits*/ 1), external_auth: RwLock::new(None), - auth_state_tx, } } /// Create an AuthManager with a specific CodexAuth, for testing only. pub fn from_auth_for_testing(auth: CodexAuth) -> Arc { - let (auth_state_tx, _) = watch::channel(()); let cached = CachedAuth { auth: Some(auth), permanent_refresh_failure: None, @@ -1275,13 +1221,11 @@ impl AuthManager { forced_chatgpt_workspace_id: RwLock::new(None), refresh_lock: Semaphore::new(/*permits*/ 1), external_auth: RwLock::new(None), - auth_state_tx, }) } /// Create an AuthManager with a specific CodexAuth and codex home, for testing only. pub fn from_auth_for_testing_with_home(auth: CodexAuth, codex_home: PathBuf) -> Arc { - let (auth_state_tx, _) = watch::channel(()); let cached = CachedAuth { auth: Some(auth), permanent_refresh_failure: None, @@ -1294,12 +1238,10 @@ impl AuthManager { forced_chatgpt_workspace_id: RwLock::new(None), refresh_lock: Semaphore::new(/*permits*/ 1), external_auth: RwLock::new(None), - auth_state_tx, }) } pub fn external_bearer_only(config: ModelProviderAuthInfo) -> Arc { - let (auth_state_tx, _) = watch::channel(()); Arc::new(Self { codex_home: PathBuf::from("non-existent"), inner: RwLock::new(CachedAuth { @@ -1313,7 +1255,6 @@ impl AuthManager { external_auth: RwLock::new(Some( Arc::new(BearerTokenRefresher::new(config)) as Arc )), - auth_state_tx, }) } @@ -1453,7 +1394,6 @@ impl AuthManager { } tracing::info!("Reloaded auth, changed: {changed}"); guard.auth = new_auth; - self.auth_state_tx.send_replace(()); changed } else { false @@ -1463,23 +1403,18 @@ impl AuthManager { pub fn set_external_auth(&self, external_auth: Arc) { if let Ok(mut guard) = self.external_auth.write() { *guard = Some(external_auth); - self.auth_state_tx.send_replace(()); } } pub fn clear_external_auth(&self) { if let Ok(mut guard) = self.external_auth.write() { *guard = None; - self.auth_state_tx.send_replace(()); } } pub fn set_forced_chatgpt_workspace_id(&self, workspace_id: Option) { - if let Ok(mut guard) = self.forced_chatgpt_workspace_id.write() - && *guard != workspace_id - { + if let Ok(mut guard) = self.forced_chatgpt_workspace_id.write() { *guard = workspace_id; - self.auth_state_tx.send_replace(()); } } @@ -1490,10 +1425,6 @@ impl AuthManager { .and_then(|guard| guard.clone()) } - pub fn subscribe_auth_state(&self) -> watch::Receiver<()> { - self.auth_state_tx.subscribe() - } - pub fn has_external_auth(&self) -> bool { self.external_auth().is_some() } @@ -1760,14 +1691,8 @@ impl AuthManager { ), ))); } - let mut auth_dot_json = + let auth_dot_json = AuthDotJson::from_external_tokens(&refreshed).map_err(RefreshTokenError::Transient)?; - if let Some(previous_auth) = self - .auth_cached() - .and_then(|auth| auth.get_current_auth_json()) - { - auth_dot_json.agent_identity = previous_auth.agent_identity; - } save_auth( &self.codex_home, &auth_dot_json, diff --git a/codex-rs/login/src/auth/storage.rs b/codex-rs/login/src/auth/storage.rs index f0d7a3169c1d..97e801415cab 100644 --- a/codex-rs/login/src/auth/storage.rs +++ b/codex-rs/login/src/auth/storage.rs @@ -39,19 +39,6 @@ pub struct AuthDotJson { #[serde(default, skip_serializing_if = "Option::is_none")] pub last_refresh: Option>, - - #[serde(default, skip_serializing_if = "Option::is_none")] - pub agent_identity: Option, -} - -#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)] -pub struct AgentIdentityAuthRecord { - pub workspace_id: String, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub chatgpt_user_id: Option, - pub agent_runtime_id: String, - pub agent_private_key: String, - pub registered_at: String, } pub(super) fn get_auth_file(codex_home: &Path) -> PathBuf { diff --git a/codex-rs/login/src/auth/storage_tests.rs b/codex-rs/login/src/auth/storage_tests.rs index 2e1cc8502977..4bf72c11b94d 100644 --- a/codex-rs/login/src/auth/storage_tests.rs +++ b/codex-rs/login/src/auth/storage_tests.rs @@ -18,7 +18,6 @@ async fn file_storage_load_returns_auth_dot_json() -> anyhow::Result<()> { openai_api_key: Some("test-key".to_string()), tokens: None, last_refresh: Some(Utc::now()), - agent_identity: None, }; storage @@ -39,7 +38,6 @@ async fn file_storage_save_persists_auth_dot_json() -> anyhow::Result<()> { openai_api_key: Some("test-key".to_string()), tokens: None, last_refresh: Some(Utc::now()), - agent_identity: None, }; let file = get_auth_file(codex_home.path()); @@ -54,30 +52,6 @@ async fn file_storage_save_persists_auth_dot_json() -> anyhow::Result<()> { Ok(()) } -#[tokio::test] -async fn file_storage_persists_agent_identity() -> anyhow::Result<()> { - let codex_home = tempdir()?; - let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); - let auth_dot_json = AuthDotJson { - auth_mode: Some(AuthMode::Chatgpt), - openai_api_key: None, - tokens: None, - last_refresh: Some(Utc::now()), - agent_identity: Some(AgentIdentityAuthRecord { - workspace_id: "account-123".to_string(), - chatgpt_user_id: Some("user-123".to_string()), - agent_runtime_id: "agent_123".to_string(), - agent_private_key: "pkcs8-base64".to_string(), - registered_at: "2026-04-13T12:00:00Z".to_string(), - }), - }; - - storage.save(&auth_dot_json)?; - - assert_eq!(storage.load()?, Some(auth_dot_json)); - Ok(()) -} - #[test] fn file_storage_delete_removes_auth_file() -> anyhow::Result<()> { let dir = tempdir()?; @@ -86,7 +60,6 @@ fn file_storage_delete_removes_auth_file() -> anyhow::Result<()> { openai_api_key: Some("sk-test-key".to_string()), tokens: None, last_refresh: None, - agent_identity: None, }; let storage = create_auth_storage(dir.path().to_path_buf(), AuthCredentialsStoreMode::File); storage.save(&auth_dot_json)?; @@ -110,7 +83,6 @@ fn ephemeral_storage_save_load_delete_is_in_memory_only() -> anyhow::Result<()> openai_api_key: Some("sk-ephemeral".to_string()), tokens: None, last_refresh: Some(Utc::now()), - agent_identity: None, }; storage.save(&auth_dot_json)?; @@ -209,7 +181,6 @@ fn auth_with_prefix(prefix: &str) -> AuthDotJson { account_id: Some(format!("{prefix}-account-id")), }), last_refresh: None, - agent_identity: None, } } @@ -226,7 +197,6 @@ fn keyring_auth_storage_load_returns_deserialized_auth() -> anyhow::Result<()> { openai_api_key: Some("sk-test".to_string()), tokens: None, last_refresh: None, - agent_identity: None, }; seed_keyring_with_auth( &mock_keyring, @@ -269,7 +239,6 @@ fn keyring_auth_storage_save_persists_and_removes_fallback_file() -> anyhow::Res account_id: Some("account".to_string()), }), last_refresh: Some(Utc::now()), - agent_identity: None, }; storage.save(&auth)?; diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index d819b0946d33..16bee1e4da0a 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -17,7 +17,7 @@ pub use server::ServerOptions; pub use server::ShutdownHandle; pub use server::run_login_server; -pub use auth::AgentIdentityAuthRecord; +pub use api_bridge::auth_provider_from_auth; pub use auth::AuthConfig; pub use auth::AuthDotJson; pub use auth::AuthManager; diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index 0c7b81018432..169a8a3091bc 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -781,7 +781,6 @@ pub(crate) async fn persist_tokens_async( openai_api_key: api_key, tokens: Some(tokens), last_refresh: Some(Utc::now()), - agent_identity: None, }; save_auth(&codex_home, &auth, auth_credentials_store_mode) }) diff --git a/codex-rs/login/tests/suite/auth_refresh.rs b/codex-rs/login/tests/suite/auth_refresh.rs index d754c9589d96..bf9e03bc263d 100644 --- a/codex-rs/login/tests/suite/auth_refresh.rs +++ b/codex-rs/login/tests/suite/auth_refresh.rs @@ -54,7 +54,6 @@ async fn refresh_token_succeeds_updates_storage() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), - agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -118,7 +117,6 @@ async fn refresh_token_refreshes_when_auth_is_unchanged() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), - agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -173,7 +171,6 @@ async fn refresh_token_skips_refresh_when_auth_changed() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens), last_refresh: Some(initial_last_refresh), - agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -183,7 +180,6 @@ async fn refresh_token_skips_refresh_when_auth_changed() -> Result<()> { openai_api_key: None, tokens: Some(disk_tokens.clone()), last_refresh: Some(initial_last_refresh), - agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -238,7 +234,6 @@ async fn refresh_token_errors_on_account_mismatch() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), - agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -249,7 +244,6 @@ async fn refresh_token_errors_on_account_mismatch() -> Result<()> { openai_api_key: None, tokens: Some(disk_tokens), last_refresh: Some(initial_last_refresh), - agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -308,7 +302,6 @@ async fn returns_fresh_tokens_as_is() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(stale_refresh), - agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -356,7 +349,6 @@ async fn refreshes_token_when_access_token_is_expired() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(fresh_refresh), - agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -406,7 +398,6 @@ async fn auth_reloads_disk_auth_when_cached_auth_is_stale() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens), last_refresh: Some(stale_refresh), - agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -417,7 +408,6 @@ async fn auth_reloads_disk_auth_when_cached_auth_is_stale() -> Result<()> { openai_api_key: None, tokens: Some(disk_tokens.clone()), last_refresh: Some(fresh_refresh), - agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -469,7 +459,6 @@ async fn auth_reloads_disk_auth_without_calling_expired_refresh_token() -> Resul openai_api_key: None, tokens: Some(initial_tokens), last_refresh: Some(stale_refresh), - agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -480,7 +469,6 @@ async fn auth_reloads_disk_auth_without_calling_expired_refresh_token() -> Resul openai_api_key: None, tokens: Some(disk_tokens.clone()), last_refresh: Some(fresh_refresh), - agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -530,7 +518,6 @@ async fn refresh_token_returns_permanent_error_for_expired_refresh_token() -> Re openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), - agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -583,7 +570,6 @@ async fn refresh_token_does_not_retry_after_permanent_failure() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), - agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -650,7 +636,6 @@ async fn refresh_token_reloads_changed_auth_after_permanent_failure() -> Result< openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), - agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -672,7 +657,6 @@ async fn refresh_token_reloads_changed_auth_after_permanent_failure() -> Result< openai_api_key: None, tokens: Some(disk_tokens.clone()), last_refresh: Some(fresh_refresh), - agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -731,7 +715,6 @@ async fn refresh_token_returns_transient_error_on_server_failure() -> Result<()> openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), - agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -784,7 +767,6 @@ async fn unauthorized_recovery_reloads_then_refreshes_tokens() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), - agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -794,7 +776,6 @@ async fn unauthorized_recovery_reloads_then_refreshes_tokens() -> Result<()> { openai_api_key: None, tokens: Some(disk_tokens.clone()), last_refresh: Some(initial_last_refresh), - agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -878,7 +859,6 @@ async fn unauthorized_recovery_errors_on_account_mismatch() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), - agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -889,7 +869,6 @@ async fn unauthorized_recovery_errors_on_account_mismatch() -> Result<()> { openai_api_key: None, tokens: Some(disk_tokens), last_refresh: Some(initial_last_refresh), - agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -947,7 +926,6 @@ async fn unauthorized_recovery_requires_chatgpt_auth() -> Result<()> { openai_api_key: Some("sk-test".to_string()), tokens: None, last_refresh: None, - agent_identity: None, }; ctx.write_auth(&auth)?; diff --git a/codex-rs/tui/src/local_chatgpt_auth.rs b/codex-rs/tui/src/local_chatgpt_auth.rs index 1f84b289a78c..e888c0387c36 100644 --- a/codex-rs/tui/src/local_chatgpt_auth.rs +++ b/codex-rs/tui/src/local_chatgpt_auth.rs @@ -108,7 +108,6 @@ mod tests { account_id: Some("workspace-1".to_string()), }), last_refresh: Some(Utc::now()), - agent_identity: None, }; save_auth(codex_home, &auth, AuthCredentialsStoreMode::File) .expect("chatgpt auth should save"); @@ -155,7 +154,6 @@ mod tests { openai_api_key: Some("sk-test".to_string()), tokens: None, last_refresh: None, - agent_identity: None, }, AuthCredentialsStoreMode::File, ) From 63491004827fd2ddc565aba995878f8afd9be954 Mon Sep 17 00:00:00 2001 From: Edward Frazer Date: Tue, 21 Apr 2026 09:34:14 -0700 Subject: [PATCH 6/9] Revert "Add use_agent_identity feature flag (#17385)" This reverts commit 39cc85310fbb1c4d04034e596cd7420090875799. --- codex-rs/app-server/tests/suite/v2/account.rs | 4 +--- codex-rs/core/config.schema.json | 6 ------ codex-rs/exec-server/src/server/handler/tests.rs | 3 +-- codex-rs/features/src/lib.rs | 8 -------- codex-rs/features/src/tests.rs | 6 ------ 5 files changed, 2 insertions(+), 25 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/account.rs b/codex-rs/app-server/tests/suite/v2/account.rs index f7981226351a..3c88bcb7a430 100644 --- a/codex-rs/app-server/tests/suite/v2/account.rs +++ b/codex-rs/app-server/tests/suite/v2/account.rs @@ -45,9 +45,7 @@ use wiremock::ResponseTemplate; use wiremock::matchers::method; use wiremock::matchers::path; -// Account tests spin up fresh app-server processes repeatedly, which can take -// longer on slower Bazel macOS runners once the suite is already warm. -const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(20); +const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); const LOGIN_ISSUER_ENV_VAR: &str = "CODEX_APP_SERVER_LOGIN_ISSUER"; // Helper to create a minimal config.toml for the app server diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index b0faea2b56d5..3f973599176c 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -521,9 +521,6 @@ "unified_exec": { "type": "boolean" }, - "use_agent_identity": { - "type": "boolean" - }, "use_legacy_landlock": { "type": "boolean" }, @@ -2520,9 +2517,6 @@ "unified_exec": { "type": "boolean" }, - "use_agent_identity": { - "type": "boolean" - }, "use_legacy_landlock": { "type": "boolean" }, diff --git a/codex-rs/exec-server/src/server/handler/tests.rs b/codex-rs/exec-server/src/server/handler/tests.rs index 7f2cdba629d8..9d1e4c470cf3 100644 --- a/codex-rs/exec-server/src/server/handler/tests.rs +++ b/codex-rs/exec-server/src/server/handler/tests.rs @@ -292,8 +292,7 @@ async fn output_and_exit_are_retained_after_notification_receiver_closes() { process_id.as_str(), shell_argv( "sleep 0.05; printf 'first\\n'; sleep 0.05; printf 'second\\n'", - // `cmd.exe` retains the space before `&&` in `echo first && ...`. - "(echo first) && ping -n 2 127.0.0.1 >NUL && (echo second)", + "echo first&& ping -n 2 127.0.0.1 >NUL&& echo second", ), )) .await diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 046063a554e4..a5a2ef7a304c 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -199,8 +199,6 @@ pub enum Feature { ResponsesWebsockets, /// Legacy rollout flag for Responses API WebSocket transport v2 experiments. ResponsesWebsocketsV2, - /// Use the agent identity registration flow for ChatGPT-authenticated sessions. - UseAgentIdentity, /// Enable workspace dependency support. WorkspaceDependencies, } @@ -988,12 +986,6 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::Removed, default_enabled: false, }, - FeatureSpec { - id: Feature::UseAgentIdentity, - key: "use_agent_identity", - stage: Stage::UnderDevelopment, - default_enabled: false, - }, FeatureSpec { id: Feature::WorkspaceDependencies, key: "workspace_dependencies", diff --git a/codex-rs/features/src/tests.rs b/codex-rs/features/src/tests.rs index 199ef9522a48..3f865e316b5b 100644 --- a/codex-rs/features/src/tests.rs +++ b/codex-rs/features/src/tests.rs @@ -225,12 +225,6 @@ fn remote_control_is_under_development() { assert_eq!(Feature::RemoteControl.default_enabled(), false); } -#[test] -fn use_agent_identity_is_under_development() { - assert_eq!(Feature::UseAgentIdentity.stage(), Stage::UnderDevelopment); - assert_eq!(Feature::UseAgentIdentity.default_enabled(), false); -} - #[test] fn workspace_dependencies_is_stable_and_enabled_by_default() { assert_eq!(Feature::WorkspaceDependencies.stage(), Stage::Stable); From 096892f3dda57cae9b58c26f1a67d024290070b6 Mon Sep 17 00:00:00 2001 From: Edward Frazer Date: Tue, 21 Apr 2026 09:35:27 -0700 Subject: [PATCH 7/9] fix: finish agent identity revert cleanup --- .../app-server/src/codex_message_processor.rs | 2 +- codex-rs/core/src/session/mod.rs | 134 ------------------ codex-rs/core/src/session/session.rs | 6 - codex-rs/core/src/session/turn.rs | 15 -- codex-rs/login/src/auth/manager.rs | 1 - codex-rs/login/src/lib.rs | 1 - codex-rs/login/tests/suite/logout.rs | 1 - 7 files changed, 1 insertion(+), 159 deletions(-) diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 038d23862bbb..2b299675ac56 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -283,8 +283,8 @@ use codex_mcp::McpServerStatusSnapshot; use codex_mcp::McpSnapshotDetail; use codex_mcp::collect_mcp_server_status_snapshot_with_detail; use codex_mcp::discover_supported_scopes; -use codex_mcp::read_mcp_resource as read_mcp_resource_without_thread; use codex_mcp::effective_mcp_servers; +use codex_mcp::read_mcp_resource as read_mcp_resource_without_thread; use codex_mcp::resolve_oauth_scopes; use codex_models_manager::collaboration_mode_presets::CollaborationModesConfig; use codex_protocol::ThreadId; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 77b868f19f16..80084f197930 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -14,7 +14,6 @@ use crate::agent::Mailbox; use crate::agent::MailboxReceiver; use crate::agent::agent_status_from_event; use crate::agent::status::is_final; -use crate::agent_identity::AgentIdentityManager; use crate::build_available_skills; use crate::commit_attribution::commit_message_trailer_instruction; use crate::compact; @@ -981,139 +980,6 @@ impl Session { }); } - fn start_agent_identity_registration(self: &Arc) { - if !self.services.agent_identity_manager.is_enabled() { - return; - } - - let weak_sess = Arc::downgrade(self); - let mut auth_state_rx = self.services.auth_manager.subscribe_auth_state(); - tokio::spawn(async move { - loop { - let Some(sess) = weak_sess.upgrade() else { - return; - }; - match sess - .services - .agent_identity_manager - .ensure_registered_identity() - .await - { - Ok(Some(_)) => return, - Ok(None) => { - drop(sess); - if auth_state_rx.changed().await.is_err() { - return; - } - } - Err(error) => { - sess.fail_agent_identity_registration(error).await; - return; - } - } - } - }); - } - - async fn fail_agent_identity_registration(self: &Arc, error: anyhow::Error) { - warn!(error = %error, "agent identity registration failed"); - let message = format!( - "Agent identity registration failed while `features.use_agent_identity` is enabled: {error}" - ); - self.send_event_raw(Event { - id: self.next_internal_sub_id(), - msg: EventMsg::Error(ErrorEvent { - message, - codex_error_info: Some(CodexErrorInfo::Other), - }), - }) - .await; - } - - async fn cached_agent_task_for_current_binding(&self) -> Option { - let agent_task = { - let state = self.state.lock().await; - state.agent_task() - }?; - - if self - .services - .agent_identity_manager - .task_matches_current_binding(&agent_task) - .await - { - debug!( - agent_runtime_id = %agent_task.agent_runtime_id, - task_id = %agent_task.task_id, - "reusing cached agent task" - ); - return Some(agent_task); - } - - debug!( - agent_runtime_id = %agent_task.agent_runtime_id, - task_id = %agent_task.task_id, - "discarding cached agent task because auth binding changed" - ); - let mut state = self.state.lock().await; - if state.agent_task().as_ref() == Some(&agent_task) { - state.clear_agent_task(); - } - None - } - - async fn ensure_agent_task_registered(&self) -> anyhow::Result> { - if let Some(agent_task) = self.cached_agent_task_for_current_binding().await { - return Ok(Some(agent_task)); - } - - for _ in 0..2 { - let Some(agent_task) = self.services.agent_identity_manager.register_task().await? - else { - return Ok(None); - }; - - if !self - .services - .agent_identity_manager - .task_matches_current_binding(&agent_task) - .await - { - debug!( - agent_runtime_id = %agent_task.agent_runtime_id, - task_id = %agent_task.task_id, - "discarding newly registered agent task because auth binding changed" - ); - continue; - } - - { - let mut state = self.state.lock().await; - if let Some(existing_agent_task) = state.agent_task() { - if existing_agent_task.has_same_binding(&agent_task) { - return Ok(Some(existing_agent_task)); - } - debug!( - agent_runtime_id = %existing_agent_task.agent_runtime_id, - task_id = %existing_agent_task.task_id, - "replacing cached agent task because auth binding changed" - ); - } - state.set_agent_task(agent_task.clone()); - } - - info!( - thread_id = %self.conversation_id, - agent_runtime_id = %agent_task.agent_runtime_id, - task_id = %agent_task.task_id, - "registered agent task for thread" - ); - return Ok(Some(agent_task)); - } - - Ok(None) - } - pub(crate) fn get_tx_event(&self) -> Sender { self.tx_event.clone() } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 015feea0c6f8..640e59b5b712 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -645,11 +645,6 @@ 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(), - )), shell_snapshot_tx, show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, @@ -752,7 +747,6 @@ impl Session { // Start the watcher after SessionConfigured so it cannot emit earlier events. sess.start_skills_watcher_listener(); - sess.start_agent_identity_registration(); let mut required_mcp_servers: Vec = mcp_servers .iter() .filter(|(_, server)| server.enabled && server.required) diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 7f156c97858d..98f572243a51 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -345,21 +345,6 @@ 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.send_event( - turn_context.as_ref(), - EventMsg::Error(ErrorEvent { - message: format!( - "Agent task registration failed. Please try again; Codex will attempt to register the task again on the next turn: {error}" - ), - codex_error_info: Some(CodexErrorInfo::Other), - }), - ) - .await; - return None; - } - if !skill_items.is_empty() { sess.record_conversation_items(&turn_context, &skill_items) .await; diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index e6bdb64a6be3..449591da9300 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -15,7 +15,6 @@ use std::sync::RwLock; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use tokio::sync::Semaphore; -use tokio::sync::watch; use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::AuthMode as ApiAuthMode; diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 16bee1e4da0a..d69a77a97d92 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -17,7 +17,6 @@ pub use server::ServerOptions; pub use server::ShutdownHandle; pub use server::run_login_server; -pub use api_bridge::auth_provider_from_auth; pub use auth::AuthConfig; pub use auth::AuthDotJson; pub use auth::AuthManager; diff --git a/codex-rs/login/tests/suite/logout.rs b/codex-rs/login/tests/suite/logout.rs index 59de9dabda92..a53006602ea5 100644 --- a/codex-rs/login/tests/suite/logout.rs +++ b/codex-rs/login/tests/suite/logout.rs @@ -192,7 +192,6 @@ fn chatgpt_auth_with_refresh_token(refresh_token: &str) -> AuthDotJson { account_id: Some("account-id".to_string()), }), last_refresh: None, - agent_identity: None, } } From aaeb20da43cdc829b64b9d75b12f248a17e1985d Mon Sep 17 00:00:00 2001 From: Edward Frazer Date: Tue, 21 Apr 2026 13:46:59 -0700 Subject: [PATCH 8/9] fix: add missing core test imports --- codex-rs/core/src/session/tests.rs | 1 + codex-rs/core/src/session/tests/guardian_tests.rs | 1 + codex-rs/core/tests/suite/responses_api_proxy_headers.rs | 1 + 3 files changed, 3 insertions(+) diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 1bc55cbbce79..ee72769bda6a 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -130,6 +130,7 @@ use std::path::Path; use std::time::Duration; use tokio::sync::Semaphore; use tokio::time::sleep; +use tokio::time::timeout; use tracing_opentelemetry::OpenTelemetrySpanExt; use codex_protocol::mcp::CallToolResult as McpCallToolResult; diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 9b1172d5788a..73a0de285a0f 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -48,6 +48,7 @@ use std::fs; use std::sync::Arc; use std::time::Duration; use tempfile::tempdir; +use tokio::time::timeout; use tokio_util::sync::CancellationToken; fn expect_text_output(output: &FunctionToolOutput) -> String { diff --git a/codex-rs/core/tests/suite/responses_api_proxy_headers.rs b/codex-rs/core/tests/suite/responses_api_proxy_headers.rs index 13d3b305d3d1..7519822a36b1 100644 --- a/codex-rs/core/tests/suite/responses_api_proxy_headers.rs +++ b/codex-rs/core/tests/suite/responses_api_proxy_headers.rs @@ -19,6 +19,7 @@ use core_test_support::responses::mount_sse_once_match; use core_test_support::responses::sse; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; +use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::test_codex; use pretty_assertions::assert_eq; use serde_json::json; From c787e22b7b5ff618f1ed925522f482eb37398557 Mon Sep 17 00:00:00 2001 From: Edward Frazer Date: Tue, 21 Apr 2026 13:47:43 -0700 Subject: [PATCH 9/9] fix: refresh locks after agent identity revert --- MODULE.bazel.lock | 2 -- codex-rs/Cargo.lock | 28 +--------------------------- 2 files changed, 1 insertion(+), 29 deletions(-) diff --git a/MODULE.bazel.lock b/MODULE.bazel.lock index 24fc9dbbd93b..5ebcb32edf04 100644 --- a/MODULE.bazel.lock +++ b/MODULE.bazel.lock @@ -846,8 +846,6 @@ "dylint_testing_5.0.0": "{\"dependencies\":[{\"name\":\"anyhow\",\"req\":\"^1.0\"},{\"name\":\"cargo_metadata\",\"req\":\"^0.23\"},{\"name\":\"compiletest_rs\",\"req\":\"^0.11\"},{\"name\":\"dylint\",\"req\":\"=5.0.0\"},{\"name\":\"dylint_internal\",\"req\":\"=5.0.0\"},{\"name\":\"env_logger\",\"req\":\"^0.11\"},{\"name\":\"once_cell\",\"req\":\"^1.21\"},{\"name\":\"regex\",\"req\":\"^1.11\"},{\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"name\":\"tempfile\",\"req\":\"^3.23\"}],\"features\":{\"default\":[],\"deny_warnings\":[]}}", "dyn-clone_1.0.20": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"rustversion\",\"req\":\"^1.0\"},{\"features\":[\"diff\"],\"kind\":\"dev\",\"name\":\"trybuild\",\"req\":\"^1.0.66\"}],\"features\":{}}", "ecdsa_0.16.9": "{\"dependencies\":[{\"name\":\"der\",\"optional\":true,\"req\":\"^0.7\"},{\"default_features\":false,\"features\":[\"oid\"],\"name\":\"digest\",\"optional\":true,\"req\":\"^0.10.7\"},{\"default_features\":false,\"features\":[\"digest\",\"sec1\"],\"name\":\"elliptic-curve\",\"req\":\"^0.13.6\"},{\"default_features\":false,\"features\":[\"dev\"],\"kind\":\"dev\",\"name\":\"elliptic-curve\",\"req\":\"^0.13\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.4\"},{\"name\":\"rfc6979\",\"optional\":true,\"req\":\"^0.4\"},{\"default_features\":false,\"features\":[\"alloc\"],\"name\":\"serdect\",\"optional\":true,\"req\":\"^0.2\"},{\"default_features\":false,\"features\":[\"oid\"],\"name\":\"sha2\",\"optional\":true,\"req\":\"^0.10\"},{\"default_features\":false,\"kind\":\"dev\",\"name\":\"sha2\",\"req\":\"^0.10\"},{\"default_features\":false,\"features\":[\"rand_core\"],\"name\":\"signature\",\"req\":\"^2.0, <2.3\"},{\"default_features\":false,\"name\":\"spki\",\"optional\":true,\"req\":\"^0.7.2\"}],\"features\":{\"alloc\":[\"elliptic-curve/alloc\",\"signature/alloc\",\"spki/alloc\"],\"arithmetic\":[\"elliptic-curve/arithmetic\"],\"default\":[\"digest\"],\"dev\":[\"arithmetic\",\"digest\",\"elliptic-curve/dev\",\"hazmat\"],\"digest\":[\"dep:digest\",\"signature/digest\"],\"hazmat\":[],\"pem\":[\"elliptic-curve/pem\",\"pkcs8\"],\"pkcs8\":[\"digest\",\"elliptic-curve/pkcs8\",\"der\"],\"serde\":[\"elliptic-curve/serde\",\"serdect\"],\"signing\":[\"arithmetic\",\"digest\",\"hazmat\",\"rfc6979\"],\"std\":[\"alloc\",\"elliptic-curve/std\",\"signature/std\"],\"verifying\":[\"arithmetic\",\"digest\",\"hazmat\"]}}", - "ed25519-dalek_2.2.0": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"bincode\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"blake2\",\"req\":\"^0.10\"},{\"features\":[\"html_reports\"],\"kind\":\"dev\",\"name\":\"criterion\",\"req\":\"^0.5\"},{\"default_features\":false,\"features\":[\"digest\"],\"name\":\"curve25519-dalek\",\"req\":\"^4\"},{\"default_features\":false,\"features\":[\"digest\",\"rand_core\"],\"kind\":\"dev\",\"name\":\"curve25519-dalek\",\"req\":\"^4\"},{\"default_features\":false,\"name\":\"ed25519\",\"req\":\">=2.2, <2.3\"},{\"kind\":\"dev\",\"name\":\"hex\",\"req\":\"^0.4\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.4\"},{\"default_features\":false,\"name\":\"merlin\",\"optional\":true,\"req\":\"^3\"},{\"kind\":\"dev\",\"name\":\"rand\",\"req\":\"^0.8\"},{\"default_features\":false,\"name\":\"rand_core\",\"optional\":true,\"req\":\"^0.6.4\"},{\"default_features\":false,\"kind\":\"dev\",\"name\":\"rand_core\",\"req\":\"^0.6.4\"},{\"default_features\":false,\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"default_features\":false,\"name\":\"sha2\",\"req\":\"^0.10\"},{\"kind\":\"dev\",\"name\":\"sha3\",\"req\":\"^0.10\"},{\"default_features\":false,\"name\":\"signature\",\"optional\":true,\"req\":\">=2.0, <2.3\"},{\"default_features\":false,\"name\":\"subtle\",\"req\":\"^2.3.0\"},{\"kind\":\"dev\",\"name\":\"toml\",\"req\":\"^0.7\"},{\"default_features\":false,\"features\":[\"static_secrets\"],\"kind\":\"dev\",\"name\":\"x25519-dalek\",\"req\":\"^2\"},{\"default_features\":false,\"name\":\"zeroize\",\"optional\":true,\"req\":\"^1.5\"}],\"features\":{\"alloc\":[\"curve25519-dalek/alloc\",\"ed25519/alloc\",\"serde?/alloc\",\"zeroize/alloc\"],\"asm\":[\"sha2/asm\"],\"batch\":[\"alloc\",\"merlin\",\"rand_core\"],\"default\":[\"fast\",\"std\",\"zeroize\"],\"digest\":[\"signature/digest\"],\"fast\":[\"curve25519-dalek/precomputed-tables\"],\"hazmat\":[],\"legacy_compatibility\":[\"curve25519-dalek/legacy_compatibility\"],\"pem\":[\"alloc\",\"ed25519/pem\",\"pkcs8\"],\"pkcs8\":[\"ed25519/pkcs8\"],\"rand_core\":[\"dep:rand_core\"],\"serde\":[\"dep:serde\",\"ed25519/serde\"],\"std\":[\"alloc\",\"ed25519/std\",\"serde?/std\",\"sha2/std\"],\"zeroize\":[\"dep:zeroize\",\"curve25519-dalek/zeroize\"]}}", - "ed25519_2.2.3": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"bincode\",\"req\":\"^1\"},{\"features\":[\"rand_core\"],\"kind\":\"dev\",\"name\":\"ed25519-dalek\",\"req\":\"^2\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.4\"},{\"name\":\"pkcs8\",\"optional\":true,\"req\":\"^0.10\"},{\"features\":[\"std\"],\"kind\":\"dev\",\"name\":\"rand_core\",\"req\":\"^0.6\"},{\"default_features\":false,\"features\":[\"signature\"],\"kind\":\"dev\",\"name\":\"ring-compat\",\"req\":\"^0.8\"},{\"default_features\":false,\"name\":\"serde\",\"optional\":true,\"req\":\"^1\"},{\"name\":\"serde_bytes\",\"optional\":true,\"req\":\"^0.11\"},{\"default_features\":false,\"name\":\"signature\",\"req\":\"^2\"},{\"default_features\":false,\"name\":\"zeroize\",\"optional\":true,\"req\":\"^1\"}],\"features\":{\"alloc\":[\"pkcs8?/alloc\"],\"default\":[\"std\"],\"pem\":[\"alloc\",\"pkcs8/pem\"],\"serde_bytes\":[\"serde\",\"dep:serde_bytes\"],\"std\":[\"pkcs8?/std\",\"signature/std\"]}}", "either_1.15.0": "{\"dependencies\":[{\"default_features\":false,\"features\":[\"alloc\",\"derive\"],\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0.95\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0.0\"}],\"features\":{\"default\":[\"std\"],\"std\":[],\"use_std\":[\"std\"]}}", "elliptic-curve_0.13.8": "{\"dependencies\":[{\"name\":\"base16ct\",\"req\":\"^0.2\"},{\"default_features\":false,\"features\":[\"alloc\"],\"name\":\"base64ct\",\"optional\":true,\"req\":\"^1\"},{\"default_features\":false,\"features\":[\"rand_core\",\"generic-array\",\"zeroize\"],\"name\":\"crypto-bigint\",\"req\":\"^0.5\"},{\"name\":\"digest\",\"optional\":true,\"req\":\"^0.10\"},{\"default_features\":false,\"name\":\"ff\",\"optional\":true,\"req\":\"^0.13\"},{\"default_features\":false,\"features\":[\"zeroize\"],\"name\":\"generic-array\",\"req\":\"^0.14.6\"},{\"default_features\":false,\"name\":\"group\",\"optional\":true,\"req\":\"^0.13\"},{\"name\":\"hex-literal\",\"optional\":true,\"req\":\"^0.4\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.4\"},{\"default_features\":false,\"name\":\"hkdf\",\"optional\":true,\"req\":\"^0.12.1\"},{\"features\":[\"alloc\"],\"name\":\"pem-rfc7468\",\"optional\":true,\"req\":\"^0.7\"},{\"default_features\":false,\"name\":\"pkcs8\",\"optional\":true,\"req\":\"^0.10.2\"},{\"default_features\":false,\"name\":\"rand_core\",\"req\":\"^0.6.4\"},{\"features\":[\"subtle\",\"zeroize\"],\"name\":\"sec1\",\"optional\":true,\"req\":\"^0.7.1\"},{\"default_features\":false,\"features\":[\"alloc\"],\"name\":\"serde_json\",\"optional\":true,\"req\":\"^1.0.47\"},{\"default_features\":false,\"features\":[\"alloc\"],\"name\":\"serdect\",\"optional\":true,\"req\":\"^0.2\"},{\"kind\":\"dev\",\"name\":\"sha2\",\"req\":\"^0.10\"},{\"kind\":\"dev\",\"name\":\"sha3\",\"req\":\"^0.10\"},{\"default_features\":false,\"name\":\"subtle\",\"req\":\"^2\"},{\"default_features\":false,\"name\":\"tap\",\"optional\":true,\"req\":\"^1.0.1\"},{\"default_features\":false,\"name\":\"zeroize\",\"req\":\"^1.7\"}],\"features\":{\"alloc\":[\"base16ct/alloc\",\"ff?/alloc\",\"group?/alloc\",\"pkcs8?/alloc\",\"sec1?/alloc\",\"zeroize/alloc\"],\"arithmetic\":[\"group\"],\"bits\":[\"arithmetic\",\"ff/bits\",\"dep:tap\"],\"default\":[\"arithmetic\"],\"dev\":[\"arithmetic\",\"dep:hex-literal\",\"pem\",\"pkcs8\"],\"ecdh\":[\"arithmetic\",\"digest\",\"dep:hkdf\"],\"group\":[\"dep:group\",\"ff\"],\"hash2curve\":[\"arithmetic\",\"digest\"],\"hazmat\":[],\"jwk\":[\"dep:base64ct\",\"dep:serde_json\",\"alloc\",\"serde\",\"zeroize/alloc\"],\"pem\":[\"dep:pem-rfc7468\",\"alloc\",\"arithmetic\",\"pkcs8\",\"sec1/pem\"],\"pkcs8\":[\"dep:pkcs8\",\"sec1\"],\"serde\":[\"dep:serdect\",\"alloc\",\"pkcs8\",\"sec1/serde\"],\"std\":[\"alloc\",\"rand_core/std\",\"pkcs8?/std\",\"sec1?/std\"],\"voprf\":[\"digest\"]}}", "ena_0.14.3": "{\"dependencies\":[{\"name\":\"dogged\",\"optional\":true,\"req\":\"^0.2.0\"},{\"name\":\"log\",\"req\":\"^0.4\"}],\"features\":{\"bench\":[],\"persistent\":[\"dogged\"]}}", diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 41ddc808dde6..1b565dbd20a7 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1936,7 +1936,6 @@ dependencies = [ "codex-api", "codex-app-server-protocol", "codex-apply-patch", - "codex-arg0", "codex-async-utils", "codex-code-mode", "codex-config", @@ -1990,7 +1989,6 @@ dependencies = [ "ctor 0.6.3", "dirs", "dunce", - "ed25519-dalek", "env-flags", "eventsource-stream", "futures", @@ -3798,7 +3796,6 @@ dependencies = [ "cfg-if", "cpufeatures", "curve25519-dalek-derive", - "digest", "fiat-crypto", "rustc_version", "subtle", @@ -4419,30 +4416,6 @@ dependencies = [ "spki", ] -[[package]] -name = "ed25519" -version = "2.2.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "115531babc129696a58c64a4fef0a8bf9e9698629fb97e9e40767d235cfbcd53" -dependencies = [ - "pkcs8", - "signature", -] - -[[package]] -name = "ed25519-dalek" -version = "2.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70e796c081cee67dc755e1a36a0a172b897fab85fc3f6bc48307991f64e4eca9" -dependencies = [ - "curve25519-dalek", - "ed25519", - "serde", - "sha2", - "subtle", - "zeroize", -] - [[package]] name = "either" version = "1.15.0" @@ -5122,6 +5095,7 @@ checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" dependencies = [ "typenum", "version_check", + "zeroize", ] [[package]]