From d766866d3ccf615d66a55f41bdd60114cee040cd Mon Sep 17 00:00:00 2001 From: Cooper Gamble Date: Wed, 3 Jun 2026 11:00:21 +0000 Subject: [PATCH] [codex] cover app-owned HTTP requests [ci changed_files] --- codex-rs/Cargo.lock | 6 + codex-rs/app-server/src/message_processor.rs | 66 +++-- codex-rs/app-server/src/request_processors.rs | 16 + .../request_processors/account_processor.rs | 99 ++++--- .../src/request_processors/apps_processor.rs | 48 ++- .../request_processors/catalog_processor.rs | 41 ++- .../request_processors/config_processor.rs | 29 +- .../src/request_processors/mcp_processor.rs | 27 +- .../src/request_processors/plugins.rs | 276 ++++++++++++------ codex-rs/chatgpt/Cargo.toml | 1 + codex-rs/chatgpt/src/chatgpt_client.rs | 35 ++- codex-rs/chatgpt/src/connectors.rs | 70 +++-- codex-rs/chatgpt/src/get_task.rs | 11 +- codex-rs/chatgpt/src/workspace_settings.rs | 18 +- codex-rs/codex-mcp/Cargo.toml | 1 + codex-rs/codex-mcp/src/connection_manager.rs | 12 +- .../codex-mcp/src/connection_manager_tests.rs | 1 + codex-rs/codex-mcp/src/mcp/mod.rs | 5 + codex-rs/core-plugins/Cargo.toml | 2 + codex-rs/core-plugins/src/manager.rs | 96 +++++- codex-rs/core-plugins/src/remote.rs | 96 ++++-- .../core-plugins/src/remote/http_state.rs | 32 ++ .../src/remote/http_state_tests.rs | 42 +++ codex-rs/core-plugins/src/remote/share.rs | 21 +- .../core-plugins/src/remote/share/tests.rs | 4 +- codex-rs/core-plugins/src/remote_legacy.rs | 35 ++- codex-rs/core/src/agent/control_tests.rs | 20 ++ codex-rs/core/src/client.rs | 42 ++- codex-rs/core/src/connectors.rs | 64 +++- codex-rs/core/src/mcp_openai_file.rs | 11 +- codex-rs/core/src/mcp_tool_call.rs | 38 ++- codex-rs/core/src/mcp_tool_call_tests.rs | 1 + .../core/src/plugins/discoverable_tests.rs | 6 +- codex-rs/core/src/realtime_conversation.rs | 20 +- codex-rs/core/src/session/mcp.rs | 9 +- codex-rs/core/src/session/mod.rs | 6 + codex-rs/core/src/session/session.rs | 2 + codex-rs/core/src/session/turn.rs | 6 +- .../tools/handlers/request_plugin_install.rs | 6 +- codex-rs/core/tests/common/test_codex.rs | 68 +++-- codex-rs/core/tests/suite/client.rs | 129 ++++++++ codex-rs/core/tests/suite/openai_file_mcp.rs | 45 ++- codex-rs/ext/image-generation/Cargo.toml | 1 + codex-rs/ext/image-generation/src/backend.rs | 27 +- .../ext/image-generation/src/extension.rs | 5 + codex-rs/ext/web-search/Cargo.toml | 1 + codex-rs/ext/web-search/src/extension.rs | 5 + codex-rs/ext/web-search/src/tool.rs | 20 +- codex-rs/http-state/src/http_state_tests.rs | 12 + codex-rs/http-state/src/lib.rs | 7 + codex-rs/model-provider/src/provider.rs | 9 + .../rmcp-client/src/http_client_adapter.rs | 123 +++++++- 52 files changed, 1415 insertions(+), 358 deletions(-) create mode 100644 codex-rs/core-plugins/src/remote/http_state.rs create mode 100644 codex-rs/core-plugins/src/remote/http_state_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 8724c30ee5..5c7c584a21 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2230,6 +2230,7 @@ dependencies = [ "codex-core", "codex-core-plugins", "codex-git-utils", + "codex-http-state", "codex-login", "codex-model-provider", "codex-plugin", @@ -2665,12 +2666,14 @@ dependencies = [ "anyhow", "chrono", "codex-analytics", + "codex-api", "codex-app-server-protocol", "codex-config", "codex-core-skills", "codex-exec-server", "codex-git-utils", "codex-hooks", + "codex-http-state", "codex-login", "codex-model-provider", "codex-otel", @@ -3061,6 +3064,7 @@ dependencies = [ "codex-core", "codex-extension-api", "codex-features", + "codex-http-state", "codex-login", "codex-model-provider", "codex-model-provider-info", @@ -3182,6 +3186,7 @@ dependencies = [ "codex-async-utils", "codex-config", "codex-exec-server", + "codex-http-state", "codex-login", "codex-model-provider", "codex-otel", @@ -4215,6 +4220,7 @@ dependencies = [ "codex-core", "codex-extension-api", "codex-features", + "codex-http-state", "codex-login", "codex-model-provider", "codex-model-provider-info", diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index d175d0cb61..c76dec9cfb 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -910,7 +910,11 @@ impl MessageProcessor { } ClientRequest::ExperimentalFeatureEnablementSet { params, .. } => { self.config_processor - .experimental_feature_enablement_set(request_id.clone(), params) + .experimental_feature_enablement_set( + request_id.clone(), + params, + app_server_client_name.as_deref(), + ) .await } ClientRequest::RemoteControlEnable { .. } => self @@ -1143,13 +1147,17 @@ impl MessageProcessor { self.thread_processor.conversation_summary(params).await } ClientRequest::SkillsList { params, .. } => { - self.catalog_processor.skills_list(params).await + self.catalog_processor + .skills_list(params, app_server_client_name.as_deref()) + .await } ClientRequest::SkillsExtraRootsSet { params, .. } => { self.catalog_processor.skills_extra_roots_set(params).await } ClientRequest::HooksList { params, .. } => { - self.catalog_processor.hooks_list(params).await + self.catalog_processor + .hooks_list(params, app_server_client_name.as_deref()) + .await } ClientRequest::MarketplaceAdd { params, .. } => { self.marketplace_processor.marketplace_add(params).await @@ -1161,45 +1169,67 @@ impl MessageProcessor { self.marketplace_processor.marketplace_upgrade(params).await } ClientRequest::PluginList { params, .. } => { - self.plugin_processor.plugin_list(params).await + self.plugin_processor + .plugin_list(params, app_server_client_name.as_deref()) + .await } ClientRequest::PluginInstalled { params, .. } => { - self.plugin_processor.plugin_installed(params).await + self.plugin_processor + .plugin_installed(params, app_server_client_name.as_deref()) + .await } ClientRequest::PluginRead { params, .. } => { - self.plugin_processor.plugin_read(params).await + self.plugin_processor + .plugin_read(params, app_server_client_name.as_deref()) + .await } ClientRequest::PluginSkillRead { params, .. } => { - self.plugin_processor.plugin_skill_read(params).await + self.plugin_processor + .plugin_skill_read(params, app_server_client_name.as_deref()) + .await } ClientRequest::PluginShareSave { params, .. } => { - self.plugin_processor.plugin_share_save(params).await + self.plugin_processor + .plugin_share_save(params, app_server_client_name.as_deref()) + .await } ClientRequest::PluginShareUpdateTargets { params, .. } => { self.plugin_processor - .plugin_share_update_targets(params) + .plugin_share_update_targets(params, app_server_client_name.as_deref()) .await } ClientRequest::PluginShareList { params, .. } => { - self.plugin_processor.plugin_share_list(params).await + self.plugin_processor + .plugin_share_list(params, app_server_client_name.as_deref()) + .await } ClientRequest::PluginShareCheckout { params, .. } => { - self.plugin_processor.plugin_share_checkout(params).await + self.plugin_processor + .plugin_share_checkout(params, app_server_client_name.as_deref()) + .await } ClientRequest::PluginShareDelete { params, .. } => { - self.plugin_processor.plugin_share_delete(params).await + self.plugin_processor + .plugin_share_delete(params, app_server_client_name.as_deref()) + .await } ClientRequest::AppsList { params, .. } => { - self.apps_processor.apps_list(&request_id, params).await + self.apps_processor + .apps_list(&request_id, params, app_server_client_name.as_deref()) + .await } ClientRequest::SkillsConfigWrite { params, .. } => { self.catalog_processor.skills_config_write(params).await } ClientRequest::PluginInstall { params, .. } => { - self.plugin_processor.plugin_install(params).await + self.plugin_processor + .plugin_install(params, app_server_client_name.as_deref()) + .await } ClientRequest::PluginUninstall { params, .. } => { - self.plugin_processor.plugin_uninstall(params).await + self.plugin_processor + .plugin_uninstall(params, app_server_client_name.as_deref()) + .await } ClientRequest::ModelList { params, .. } => { self.catalog_processor @@ -1208,7 +1238,7 @@ impl MessageProcessor { } ClientRequest::ExperimentalFeatureList { params, .. } => { self.catalog_processor - .experimental_feature_list(params) + .experimental_feature_list(params, app_server_client_name.as_deref()) .await } ClientRequest::PermissionProfileList { params, .. } => { @@ -1279,12 +1309,12 @@ impl MessageProcessor { } ClientRequest::McpServerStatusList { params, .. } => { self.mcp_processor - .mcp_server_status_list(&request_id, params) + .mcp_server_status_list(&request_id, params, app_server_client_name.as_deref()) .await } ClientRequest::McpResourceRead { params, .. } => { self.mcp_processor - .mcp_resource_read(&request_id, params) + .mcp_resource_read(&request_id, params, app_server_client_name.as_deref()) .await } ClientRequest::McpServerToolCall { params, .. } => { diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 73ea4e82cb..b71a164ae4 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -533,6 +533,22 @@ pub(crate) use self::thread_summary::summary_to_thread; pub(crate) use self::thread_summary::thread_settings_from_config_snapshot; pub(crate) use self::thread_summary::thread_settings_from_core_snapshot; +fn http_state_context( + config: &Config, + app_server_client_name: Option<&str>, +) -> codex_http_state::HttpStateContext { + codex_http_state::HttpStateContext::new( + config.codex_home.to_path_buf(), + http_state_surface(app_server_client_name), + ) +} + +fn http_state_surface(app_server_client_name: Option<&str>) -> codex_http_state::HttpStateSurface { + codex_http_state::HttpStateSurface::from_app_server_client_name( + app_server_client_name.unwrap_or_default(), + ) +} + pub(crate) fn build_api_turns_from_rollout_items(items: &[RolloutItem]) -> Vec { let mut builder = ThreadHistoryBuilder::new(); for item in items { diff --git a/codex-rs/app-server/src/request_processors/account_processor.rs b/codex-rs/app-server/src/request_processors/account_processor.rs index 98aeda1cea..ea353d80bf 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -183,6 +183,7 @@ impl AccountRequestProcessor { config_manager: &ConfigManager, thread_manager: &Arc, auth: Option, + http_state_surface: Option, ) { match config_manager .load_latest_config(/*fallback_cwd*/ None) @@ -202,6 +203,9 @@ impl AccountRequestProcessor { refresh_config_manager.clone(), ); })), + http_state_surface.map(|surface| { + HttpStateContext::new(config.codex_home.to_path_buf(), surface) + }), ); } Err(err) => { @@ -257,6 +261,7 @@ impl AccountRequestProcessor { access_token, chatgpt_account_id, chatgpt_plan_type, + http_state_surface, ) .await; } @@ -317,8 +322,10 @@ impl AccountRequestProcessor { self.outgoing.send_result(request_id, result).await; if logged_in { - self.send_login_success_notifications(/*login_id*/ None) - .await; + self.send_login_success_notifications( + /*login_id*/ None, /*http_state_surface*/ None, + ) + .await; } } @@ -411,11 +418,7 @@ impl AccountRequestProcessor { }); } - let outgoing_clone = self.outgoing.clone(); - let config_manager = self.config_manager.clone(); - let thread_manager = Arc::clone(&self.thread_manager); - let chatgpt_base_url = self.config.chatgpt_base_url.clone(); - let active_login = self.active_login.clone(); + let processor = self.clone(); let auth_url = server.auth_url.clone(); tokio::spawn(async move { let (success, error_msg) = match tokio::time::timeout( @@ -432,19 +435,17 @@ impl AccountRequestProcessor { } }; - Self::send_chatgpt_login_completion_notifications( - &outgoing_clone, - config_manager, - thread_manager, - chatgpt_base_url, - login_id, - success, - error_msg, - ) - .await; + processor + .send_chatgpt_login_completion_notifications( + login_id, + success, + error_msg, + http_state_surface, + ) + .await; // Clear the active login if it matches this attempt. It may have been replaced or cancelled. - let mut guard = active_login.lock().await; + let mut guard = processor.active_login.lock().await; if guard.as_ref().map(ActiveLogin::login_id) == Some(login_id) { *guard = None; } @@ -494,11 +495,7 @@ impl AccountRequestProcessor { let verification_url = device_code.verification_url.clone(); let user_code = device_code.user_code.clone(); - let outgoing_clone = self.outgoing.clone(); - let config_manager = self.config_manager.clone(); - let thread_manager = Arc::clone(&self.thread_manager); - let chatgpt_base_url = self.config.chatgpt_base_url.clone(); - let active_login = self.active_login.clone(); + let processor = self.clone(); tokio::spawn(async move { let (success, error_msg) = tokio::select! { _ = cancel.cancelled() => { @@ -512,18 +509,16 @@ impl AccountRequestProcessor { } }; - Self::send_chatgpt_login_completion_notifications( - &outgoing_clone, - config_manager, - thread_manager, - chatgpt_base_url, - login_id, - success, - error_msg, - ) - .await; + processor + .send_chatgpt_login_completion_notifications( + login_id, + success, + error_msg, + http_state_surface, + ) + .await; - let mut guard = active_login.lock().await; + let mut guard = processor.active_login.lock().await; if guard.as_ref().map(ActiveLogin::login_id) == Some(login_id) { *guard = None; } @@ -571,6 +566,7 @@ impl AccountRequestProcessor { access_token: String, chatgpt_account_id: String, chatgpt_plan_type: Option, + http_state_surface: HttpStateSurface, ) { let result = self .login_chatgpt_auth_tokens_response(access_token, chatgpt_account_id, chatgpt_plan_type) @@ -579,7 +575,7 @@ impl AccountRequestProcessor { self.outgoing.send_result(request_id, result).await; if logged_in { - self.send_login_success_notifications(/*login_id*/ None) + self.send_login_success_notifications(/*login_id*/ None, Some(http_state_surface)) .await; } } @@ -634,11 +630,16 @@ impl AccountRequestProcessor { Ok(LoginAccountResponse::ChatgptAuthTokens {}) } - async fn send_login_success_notifications(&self, login_id: Option) { + async fn send_login_success_notifications( + &self, + login_id: Option, + http_state_surface: Option, + ) { Self::maybe_refresh_remote_installed_plugins_cache_for_current_config( &self.config_manager, &self.thread_manager, self.auth_manager.auth_cached(), + http_state_surface, ) .await; @@ -661,44 +662,45 @@ impl AccountRequestProcessor { } async fn send_chatgpt_login_completion_notifications( - outgoing: &OutgoingMessageSender, - config_manager: ConfigManager, - thread_manager: Arc, - chatgpt_base_url: String, + &self, login_id: Uuid, success: bool, error_msg: Option, + http_state_surface: HttpStateSurface, ) { let payload_v2 = AccountLoginCompletedNotification { login_id: Some(login_id.to_string()), success, error: error_msg, }; - outgoing + self.outgoing .send_server_notification(ServerNotification::AccountLoginCompleted(payload_v2)) .await; if success { - let auth_manager = thread_manager.auth_manager(); + let auth_manager = self.thread_manager.auth_manager(); auth_manager.reload().await; - config_manager - .replace_cloud_config_bundle_loader(auth_manager.clone(), chatgpt_base_url); - config_manager + self.config_manager.replace_cloud_config_bundle_loader( + auth_manager.clone(), + self.config.chatgpt_base_url.clone(), + ); + self.config_manager .sync_default_client_residency_requirement() .await; let auth = auth_manager.auth_cached(); Self::maybe_refresh_remote_installed_plugins_cache_for_current_config( - &config_manager, - &thread_manager, + &self.config_manager, + &self.thread_manager, auth.clone(), + Some(http_state_surface), ) .await; let payload_v2 = AccountUpdatedNotification { auth_mode: auth.as_ref().map(CodexAuth::api_auth_mode), plan_type: auth.as_ref().and_then(CodexAuth::account_plan_type), }; - outgoing + self.outgoing .send_server_notification(ServerNotification::AccountUpdated(payload_v2)) .await; } @@ -724,6 +726,7 @@ impl AccountRequestProcessor { &self.config_manager, &self.thread_manager, self.auth_manager.auth_cached(), + /*http_state_surface*/ None, ) .await; diff --git a/codex-rs/app-server/src/request_processors/apps_processor.rs b/codex-rs/app-server/src/request_processors/apps_processor.rs index 49a6615f1c..c1a9de4d94 100644 --- a/codex-rs/app-server/src/request_processors/apps_processor.rs +++ b/codex-rs/app-server/src/request_processors/apps_processor.rs @@ -1,4 +1,5 @@ use super::*; +use codex_http_state::HttpStateContext; pub(crate) struct AppsRequestProcessor { auth_manager: Arc, @@ -35,8 +36,9 @@ impl AppsRequestProcessor { &self, request_id: &ConnectionRequestId, params: AppsListParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - self.apps_list_inner(request_id, params) + self.apps_list_inner(request_id, params, app_server_client_name) .await .map(|response| response.map(Into::into)) } @@ -45,6 +47,7 @@ impl AppsRequestProcessor { &self, request_id: &ConnectionRequestId, params: AppsListParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { let thread = if let Some(thread_id) = params.thread_id.as_deref() { let (_, loaded_thread) = self.load_thread(thread_id).await?; @@ -64,7 +67,11 @@ impl AppsRequestProcessor { .set_enabled(Feature::Apps, thread.enabled(Feature::Apps)); } - let auth = self.auth_manager.auth().await; + let auth = self + .auth_manager + .auth_for_surface(http_state_surface(app_server_client_name)) + .await; + let http_state = http_state_context(&config, app_server_client_name); if !config .features .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend)) @@ -76,7 +83,7 @@ impl AppsRequestProcessor { } if !self - .workspace_codex_plugins_enabled(&config, auth.as_ref()) + .workspace_codex_plugins_enabled(&config, auth.as_ref(), &http_state) .await { return Ok(Some(AppsListResponse { @@ -92,7 +99,7 @@ impl AppsRequestProcessor { tokio::spawn(async move { tokio::select! { _ = shutdown_token.cancelled() => {} - _ = Self::apps_list_task(outgoing, request, params, config, environment_manager) => {} + _ = Self::apps_list_task(outgoing, request, params, config, environment_manager, http_state) => {} } }); Ok(None) @@ -108,11 +115,15 @@ impl AppsRequestProcessor { params: AppsListParams, config: Config, environment_manager: Arc, + http_state: HttpStateContext, ) { let retry_params = params.clone(); let retry_config = config.clone(); let retry_environment_manager = Arc::clone(&environment_manager); - let result = Self::apps_list_response(&outgoing, params, config, environment_manager).await; + let retry_http_state = http_state.clone(); + let result = + Self::apps_list_response(&outgoing, params, config, environment_manager, http_state) + .await; let should_retry = result .as_ref() .is_ok_and(|(_, codex_apps_ready)| !codex_apps_ready); @@ -128,6 +139,7 @@ impl AppsRequestProcessor { retry_params, retry_config, retry_environment_manager, + retry_http_state, ) .await { @@ -141,6 +153,7 @@ impl AppsRequestProcessor { params: AppsListParams, config: Config, environment_manager: Arc, + http_state: HttpStateContext, ) -> Result<(AppsListResponse, bool), JSONRPCErrorError> { let AppsListParams { cursor, @@ -157,8 +170,11 @@ impl AppsRequestProcessor { }; let (mut accessible_connectors, mut all_connectors) = tokio::join!( - connectors::list_cached_accessible_connectors_from_mcp_tools(&config), - connectors::list_cached_all_connectors(&config) + connectors::list_cached_accessible_connectors_from_mcp_tools_with_http_state( + &config, + Some(&http_state), + ), + connectors::list_cached_all_connectors_with_http_state(&config, Some(&http_state)) ); let cached_all_connectors = all_connectors.clone(); @@ -166,12 +182,14 @@ impl AppsRequestProcessor { let accessible_config = config.clone(); let accessible_tx = tx.clone(); + let accessible_http_state = http_state.clone(); tokio::spawn(async move { let result = - connectors::list_accessible_connectors_from_mcp_tools_with_environment_manager( + connectors::list_accessible_connectors_from_mcp_tools_with_environment_manager_and_http_state( &accessible_config, force_refetch, Arc::clone(&environment_manager), + Some(accessible_http_state), ) .await .map_err(|err| format!("failed to load accessible apps: {err}")); @@ -180,9 +198,13 @@ impl AppsRequestProcessor { let all_config = config.clone(); tokio::spawn(async move { - let result = connectors::list_all_connectors_with_options(&all_config, force_refetch) - .await - .map_err(|err| format!("failed to list apps: {err}")); + let result = connectors::list_all_connectors_with_options_and_http_state( + &all_config, + force_refetch, + Some(http_state), + ) + .await + .map_err(|err| format!("failed to list apps: {err}")); let _ = tx.send(AppListLoadResult::Directory(result)); }); @@ -303,11 +325,13 @@ impl AppsRequestProcessor { &self, config: &Config, auth: Option<&CodexAuth>, + http_state: &HttpStateContext, ) -> bool { - match workspace_settings::codex_plugins_enabled_for_workspace( + match workspace_settings::codex_plugins_enabled_for_workspace_with_http_state( config, auth, Some(&self.workspace_settings_cache), + Some(http_state.clone()), ) .await { diff --git a/codex-rs/app-server/src/request_processors/catalog_processor.rs b/codex-rs/app-server/src/request_processors/catalog_processor.rs index 5f2d681912..e6aac12cd8 100644 --- a/codex-rs/app-server/src/request_processors/catalog_processor.rs +++ b/codex-rs/app-server/src/request_processors/catalog_processor.rs @@ -1,5 +1,6 @@ use super::*; use codex_config::config_toml::ConfigToml; +use codex_http_state::HttpStateContext; use codex_http_state::HttpStateSurface; use futures::StreamExt; @@ -121,8 +122,9 @@ impl CatalogRequestProcessor { pub(crate) async fn skills_list( &self, params: SkillsListParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - self.skills_list_response(params) + self.skills_list_response(params, app_server_client_name) .await .map(|response| Some(response.into())) } @@ -130,8 +132,9 @@ impl CatalogRequestProcessor { pub(crate) async fn hooks_list( &self, params: HooksListParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - self.hooks_list_response(params) + self.hooks_list_response(params, app_server_client_name) .await .map(|response| Some(response.into())) } @@ -167,8 +170,9 @@ impl CatalogRequestProcessor { pub(crate) async fn experimental_feature_list( &self, params: ExperimentalFeatureListParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - self.experimental_feature_list_response(params) + self.experimental_feature_list_response(params, app_server_client_name) .await .map(|response| Some(response.into())) } @@ -229,11 +233,13 @@ impl CatalogRequestProcessor { &self, config: &Config, auth: Option<&CodexAuth>, + http_state: &HttpStateContext, ) -> bool { - match workspace_settings::codex_plugins_enabled_for_workspace( + match workspace_settings::codex_plugins_enabled_for_workspace_with_http_state( config, auth, Some(&self.workspace_settings_cache), + Some(http_state.clone()), ) .await { @@ -320,6 +326,7 @@ impl CatalogRequestProcessor { async fn experimental_feature_list_response( &self, params: ExperimentalFeatureListParams, + app_server_client_name: Option<&str>, ) -> Result { let ExperimentalFeatureListParams { cursor, @@ -343,9 +350,13 @@ impl CatalogRequestProcessor { } None => self.load_latest_config(/*fallback_cwd*/ None).await?, }; - let auth = self.auth_manager.auth().await; + let http_state = http_state_context(&config, app_server_client_name); + let auth = self + .auth_manager + .auth_for_surface(http_state_surface(app_server_client_name)) + .await; let workspace_codex_plugins_enabled = self - .workspace_codex_plugins_enabled(&config, auth.as_ref()) + .workspace_codex_plugins_enabled(&config, auth.as_ref(), &http_state) .await; let data = FEATURES @@ -509,6 +520,7 @@ impl CatalogRequestProcessor { async fn skills_list_response( &self, params: SkillsListParams, + app_server_client_name: Option<&str>, ) -> Result { let SkillsListParams { cwds, force_reload } = params; let cwds = if cwds.is_empty() { @@ -518,9 +530,13 @@ impl CatalogRequestProcessor { }; let config = self.load_latest_config(/*fallback_cwd*/ None).await?; - let auth = self.auth_manager.auth().await; + let http_state = http_state_context(&config, app_server_client_name); + let auth = self + .auth_manager + .auth_for_surface(http_state_surface(app_server_client_name)) + .await; let workspace_codex_plugins_enabled = self - .workspace_codex_plugins_enabled(&config, auth.as_ref()) + .workspace_codex_plugins_enabled(&config, auth.as_ref(), &http_state) .await; let skills_manager = self.thread_manager.skills_manager(); let plugins_manager = self.thread_manager.plugins_manager(); @@ -615,6 +631,7 @@ impl CatalogRequestProcessor { async fn hooks_list_response( &self, params: HooksListParams, + app_server_client_name: Option<&str>, ) -> Result { let HooksListParams { cwds } = params; let cwds = if cwds.is_empty() { @@ -623,7 +640,10 @@ impl CatalogRequestProcessor { cwds }; - let auth = self.auth_manager.auth().await; + let auth = self + .auth_manager + .auth_for_surface(http_state_surface(app_server_client_name)) + .await; let plugins_manager = self.thread_manager.plugins_manager(); let mut data = Vec::new(); for cwd in cwds { @@ -651,8 +671,9 @@ impl CatalogRequestProcessor { continue; } }; + let http_state = http_state_context(&config, app_server_client_name); let workspace_codex_plugins_enabled = self - .workspace_codex_plugins_enabled(&config, auth.as_ref()) + .workspace_codex_plugins_enabled(&config, auth.as_ref(), &http_state) .await; let plugins_enabled = config.features.enabled(Feature::Plugins) && workspace_codex_plugins_enabled; diff --git a/codex-rs/app-server/src/request_processors/config_processor.rs b/codex-rs/app-server/src/request_processors/config_processor.rs index 3571026193..62c07eac19 100644 --- a/codex-rs/app-server/src/request_processors/config_processor.rs +++ b/codex-rs/app-server/src/request_processors/config_processor.rs @@ -1,5 +1,7 @@ use std::sync::Arc; +use super::http_state_context; +use super::http_state_surface; use crate::config_manager::ConfigManager; use crate::config_manager_service::ConfigManagerError; use crate::error_code::internal_error; @@ -150,6 +152,7 @@ impl ConfigRequestProcessor { &self, request_id: ConnectionRequestId, params: ExperimentalFeatureEnablementSetParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { let should_refresh_apps_list = params.enablement.get("apps").copied() == Some(true); let response = self @@ -162,8 +165,10 @@ impl ConfigRequestProcessor { ) .await; if should_refresh_apps_list { - self.refresh_apps_list_after_experimental_feature_enablement_set() - .await; + self.refresh_apps_list_after_experimental_feature_enablement_set( + app_server_client_name, + ) + .await; } Ok(None) } @@ -195,7 +200,10 @@ impl ConfigRequestProcessor { Ok(response) } - async fn refresh_apps_list_after_experimental_feature_enablement_set(&self) { + async fn refresh_apps_list_after_experimental_feature_enablement_set( + &self, + app_server_client_name: Option<&str>, + ) { let config = match self.load_latest_config(/*fallback_cwd*/ None).await { Ok(config) => config, Err(error) => { @@ -206,7 +214,11 @@ impl ConfigRequestProcessor { return; } }; - let auth = self.auth_manager.auth().await; + let http_state = http_state_context(&config, app_server_client_name); + let auth = self + .auth_manager + .auth_for_surface(http_state_surface(app_server_client_name)) + .await; if !config.features.apps_enabled_for_auth( auth.as_ref() .is_some_and(codex_login::CodexAuth::uses_codex_backend), @@ -218,11 +230,16 @@ impl ConfigRequestProcessor { let environment_manager = self.thread_manager.environment_manager(); tokio::spawn(async move { let (all_connectors_result, accessible_connectors_result) = tokio::join!( - connectors::list_all_connectors_with_options(&config, /*force_refetch*/ true), - connectors::list_accessible_connectors_from_mcp_tools_with_environment_manager( + connectors::list_all_connectors_with_options_and_http_state( + &config, + /*force_refetch*/ true, + Some(http_state.clone()), + ), + connectors::list_accessible_connectors_from_mcp_tools_with_environment_manager_and_http_state( &config, /*force_refetch*/ true, Arc::clone(&environment_manager), + Some(http_state), ), ); let all_connectors = match all_connectors_result { diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index ae62e2e785..2e94a79ffa 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -1,4 +1,5 @@ use super::*; +use codex_http_state::HttpStateContext; const MCP_TOOL_THREAD_ID_META_KEY: &str = "threadId"; @@ -47,8 +48,9 @@ impl McpRequestProcessor { &self, request_id: &ConnectionRequestId, params: ListMcpServerStatusParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - self.list_mcp_server_status(request_id, params) + self.list_mcp_server_status(request_id, params, app_server_client_name) .await .map(|()| None) } @@ -57,8 +59,9 @@ impl McpRequestProcessor { &self, request_id: &ConnectionRequestId, params: McpResourceReadParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - self.read_mcp_resource(request_id, params) + self.read_mcp_resource(request_id, params, app_server_client_name) .await .map(|()| None) } @@ -195,6 +198,7 @@ impl McpRequestProcessor { &self, request_id: &ConnectionRequestId, params: ListMcpServerStatusParams, + app_server_client_name: Option<&str>, ) -> Result<(), JSONRPCErrorError> { let request = request_id.clone(); @@ -213,7 +217,11 @@ impl McpRequestProcessor { let mcp_config = config .to_mcp_config(self.thread_manager.plugins_manager().as_ref()) .await; - let auth = self.auth_manager.auth().await; + let http_state = http_state_context(&config, app_server_client_name); + let auth = self + .auth_manager + .auth_for_surface(http_state_surface(app_server_client_name)) + .await; let environment_manager = self.thread_manager.environment_manager(); // This status path has no turn-selected environment. Use config cwd // as the local stdio fallback; named environment stdio MCPs must @@ -229,6 +237,7 @@ impl McpRequestProcessor { mcp_config, auth, runtime_context, + http_state, ) .await; }); @@ -242,6 +251,7 @@ impl McpRequestProcessor { mcp_config: codex_mcp::McpConfig, auth: Option, runtime_context: McpRuntimeContext, + http_state: HttpStateContext, ) { let result = Self::list_mcp_server_status_response( request_id.request_id.to_string(), @@ -249,6 +259,7 @@ impl McpRequestProcessor { mcp_config, auth, runtime_context, + http_state, ) .await; outgoing.send_result(request_id, result).await; @@ -260,6 +271,7 @@ impl McpRequestProcessor { mcp_config: codex_mcp::McpConfig, auth: Option, runtime_context: McpRuntimeContext, + http_state: HttpStateContext, ) -> Result { let detail = match params.detail.unwrap_or(McpServerStatusDetail::Full) { McpServerStatusDetail::Full => McpSnapshotDetail::Full, @@ -271,6 +283,7 @@ impl McpRequestProcessor { auth.as_ref(), request_id, runtime_context, + Some(http_state), detail, ) .await; @@ -341,6 +354,7 @@ impl McpRequestProcessor { &self, request_id: &ConnectionRequestId, params: McpResourceReadParams, + app_server_client_name: Option<&str>, ) -> Result<(), JSONRPCErrorError> { let outgoing = Arc::clone(&self.outgoing); let McpResourceReadParams { @@ -364,7 +378,11 @@ impl McpRequestProcessor { let mcp_config = config .to_mcp_config(self.thread_manager.plugins_manager().as_ref()) .await; - let auth = self.auth_manager.auth().await; + let http_state = http_state_context(&config, app_server_client_name); + let auth = self + .auth_manager + .auth_for_surface(http_state_surface(app_server_client_name)) + .await; let environment_manager = self.thread_manager.environment_manager(); // This threadless resource-read path has no turn cwd or turn-selected // environment. Use config cwd only as the local stdio fallback; named @@ -378,6 +396,7 @@ impl McpRequestProcessor { &mcp_config, auth.as_ref(), runtime_context, + Some(http_state), &server, &uri, ) diff --git a/codex-rs/app-server/src/request_processors/plugins.rs b/codex-rs/app-server/src/request_processors/plugins.rs index 2335694856..6cebdf3897 100644 --- a/codex-rs/app-server/src/request_processors/plugins.rs +++ b/codex-rs/app-server/src/request_processors/plugins.rs @@ -26,6 +26,14 @@ pub(crate) struct PluginRequestProcessor { workspace_settings_cache: Arc, } +fn remote_plugin_service_config( + config: &Config, + app_server_client_name: Option<&str>, +) -> RemotePluginServiceConfig { + RemotePluginServiceConfig::new(config.chatgpt_base_url.clone()) + .with_http_state(http_state_context(config, app_server_client_name)) +} + fn plugin_skills_to_info( skills: &[codex_core::skills::SkillMetadata], disabled_skill_paths: &HashSet, @@ -341,11 +349,21 @@ impl PluginRequestProcessor { } } + async fn auth_for_app_server_client( + &self, + app_server_client_name: Option<&str>, + ) -> Option { + self.auth_manager + .auth_for_surface(http_state_surface(app_server_client_name)) + .await + } + pub(crate) async fn plugin_list( &self, params: PluginListParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - self.plugin_list_response(params) + self.plugin_list_response(params, app_server_client_name) .await .map(|response| Some(response.into())) } @@ -353,8 +371,9 @@ impl PluginRequestProcessor { pub(crate) async fn plugin_installed( &self, params: PluginInstalledParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - self.plugin_installed_response(params) + self.plugin_installed_response(params, app_server_client_name) .await .map(|response| Some(response.into())) } @@ -362,8 +381,9 @@ impl PluginRequestProcessor { pub(crate) async fn plugin_read( &self, params: PluginReadParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - self.plugin_read_response(params) + self.plugin_read_response(params, app_server_client_name) .await .map(|response| Some(response.into())) } @@ -371,8 +391,9 @@ impl PluginRequestProcessor { pub(crate) async fn plugin_skill_read( &self, params: PluginSkillReadParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - self.plugin_skill_read_response(params) + self.plugin_skill_read_response(params, app_server_client_name) .await .map(|response| Some(response.into())) } @@ -380,8 +401,9 @@ impl PluginRequestProcessor { pub(crate) async fn plugin_share_save( &self, params: PluginShareSaveParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - self.plugin_share_save_response(params) + self.plugin_share_save_response(params, app_server_client_name) .await .map(|response| Some(response.into())) } @@ -389,8 +411,9 @@ impl PluginRequestProcessor { pub(crate) async fn plugin_share_update_targets( &self, params: PluginShareUpdateTargetsParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - self.plugin_share_update_targets_response(params) + self.plugin_share_update_targets_response(params, app_server_client_name) .await .map(|response| Some(response.into())) } @@ -398,8 +421,9 @@ impl PluginRequestProcessor { pub(crate) async fn plugin_share_list( &self, params: PluginShareListParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - self.plugin_share_list_response(params) + self.plugin_share_list_response(params, app_server_client_name) .await .map(|response| Some(response.into())) } @@ -407,8 +431,9 @@ impl PluginRequestProcessor { pub(crate) async fn plugin_share_checkout( &self, params: PluginShareCheckoutParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - self.plugin_share_checkout_response(params) + self.plugin_share_checkout_response(params, app_server_client_name) .await .map(|response| Some(response.into())) } @@ -416,8 +441,9 @@ impl PluginRequestProcessor { pub(crate) async fn plugin_share_delete( &self, params: PluginShareDeleteParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - self.plugin_share_delete_response(params) + self.plugin_share_delete_response(params, app_server_client_name) .await .map(|response| Some(response.into())) } @@ -425,8 +451,9 @@ impl PluginRequestProcessor { pub(crate) async fn plugin_install( &self, params: PluginInstallParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - self.plugin_install_response(params) + self.plugin_install_response(params, app_server_client_name) .await .map(|response| Some(response.into())) } @@ -434,8 +461,9 @@ impl PluginRequestProcessor { pub(crate) async fn plugin_uninstall( &self, params: PluginUninstallParams, + app_server_client_name: Option<&str>, ) -> Result, JSONRPCErrorError> { - self.plugin_uninstall_response(params) + self.plugin_uninstall_response(params, app_server_client_name) .await .map(|response| Some(response.into())) } @@ -491,11 +519,13 @@ impl PluginRequestProcessor { &self, config: &Config, auth: Option<&CodexAuth>, + app_server_client_name: Option<&str>, ) -> bool { - match workspace_settings::codex_plugins_enabled_for_workspace( + match workspace_settings::codex_plugins_enabled_for_workspace_with_http_state( config, auth, Some(&self.workspace_settings_cache), + Some(http_state_context(config, app_server_client_name)), ) .await { @@ -512,6 +542,7 @@ impl PluginRequestProcessor { async fn plugin_list_response( &self, params: PluginListParams, + app_server_client_name: Option<&str>, ) -> Result { let plugins_manager = self.thread_manager.plugins_manager(); let PluginListParams { @@ -534,9 +565,11 @@ impl PluginRequestProcessor { if !config.features.enabled(Feature::Plugins) { return Ok(empty_response()); } - let auth = self.auth_manager.auth().await; + let auth = self + .auth_for_app_server_client(app_server_client_name) + .await; if !self - .workspace_codex_plugins_enabled(&config, auth.as_ref()) + .workspace_codex_plugins_enabled(&config, auth.as_ref(), app_server_client_name) .await { return Ok(empty_response()); @@ -548,6 +581,7 @@ impl PluginRequestProcessor { auth.clone(), &roots, Some(self.effective_plugins_changed_callback()), + Some(http_state_context(&config, app_server_client_name)), ); } let (mut data, marketplace_load_errors) = if include_local { @@ -616,9 +650,8 @@ impl PluginRequestProcessor { // TODO(remote plugins): Remove this once remote plugins are ready and vertical plugins are // served directly from the normal remote catalog. if include_vertical && !config.features.enabled(Feature::RemotePlugin) { - let remote_plugin_service_config = RemotePluginServiceConfig { - chatgpt_base_url: config.chatgpt_base_url.clone(), - }; + let remote_plugin_service_config = + remote_plugin_service_config(&config, app_server_client_name); match codex_core_plugins::remote::fetch_openai_curated_remote_collection_marketplace( &remote_plugin_service_config, auth.as_ref(), @@ -655,9 +688,8 @@ impl PluginRequestProcessor { remote_sources.push(RemoteMarketplaceSource::SharedWithMe); } if !remote_sources.is_empty() { - let remote_plugin_service_config = RemotePluginServiceConfig { - chatgpt_base_url: config.chatgpt_base_url.clone(), - }; + let remote_plugin_service_config = + remote_plugin_service_config(&config, app_server_client_name); match codex_core_plugins::remote::fetch_remote_marketplaces( &remote_plugin_service_config, auth.as_ref(), @@ -707,7 +739,11 @@ impl PluginRequestProcessor { .any(|marketplace| marketplace.name == OPENAI_CURATED_MARKETPLACE_NAME) { match plugins_manager - .featured_plugin_ids_for_config(&plugins_input, auth.as_ref()) + .featured_plugin_ids_for_config_with_http_state( + &plugins_input, + auth.as_ref(), + http_state_context(&config, app_server_client_name), + ) .await { Ok(featured_plugin_ids) => featured_plugin_ids, @@ -733,6 +769,7 @@ impl PluginRequestProcessor { async fn plugin_installed_response( &self, params: PluginInstalledParams, + app_server_client_name: Option<&str>, ) -> Result { let plugins_manager = self.thread_manager.plugins_manager(); let PluginInstalledParams { @@ -753,9 +790,11 @@ impl PluginRequestProcessor { if !config.features.enabled(Feature::Plugins) { return Ok(empty_response()); } - let auth = self.auth_manager.auth().await; + let auth = self + .auth_for_app_server_client(app_server_client_name) + .await; if !self - .workspace_codex_plugins_enabled(&config, auth.as_ref()) + .workspace_codex_plugins_enabled(&config, auth.as_ref(), app_server_client_name) .await { return Ok(empty_response()); @@ -768,6 +807,7 @@ impl PluginRequestProcessor { &plugins_input, auth.clone(), Some(self.effective_plugins_changed_callback()), + Some(http_state_context(&config, app_server_client_name)), ); let (mut data, marketplace_load_errors) = self @@ -783,9 +823,11 @@ impl PluginRequestProcessor { data.extend( self.load_remote_installed_plugins( plugins_manager, + &config, &plugins_input, &remote_installed_plugin_visible_scopes, auth.as_ref(), + app_server_client_name, ) .await, ); @@ -883,9 +925,11 @@ impl PluginRequestProcessor { async fn load_remote_installed_plugins( &self, plugins_manager: Arc, + config: &Config, plugins_input: &codex_core_plugins::PluginsConfigInput, visible_scopes: &[RemotePluginScope], auth: Option<&CodexAuth>, + app_server_client_name: Option<&str>, ) -> Vec { let remote_marketplaces = if let Some(remote_marketplaces) = plugins_manager.build_remote_installed_plugin_marketplaces_from_cache(visible_scopes) @@ -893,11 +937,12 @@ impl PluginRequestProcessor { Ok(remote_marketplaces) } else { plugins_manager - .build_and_cache_remote_installed_plugin_marketplaces( + .build_and_cache_remote_installed_plugin_marketplaces_with_http_state( plugins_input, auth, visible_scopes, Some(self.effective_plugins_changed_callback()), + http_state_context(config, app_server_client_name), ) .await }; @@ -924,6 +969,7 @@ impl PluginRequestProcessor { async fn plugin_read_response( &self, params: PluginReadParams, + app_server_client_name: Option<&str>, ) -> Result { let plugins_manager = self.thread_manager.plugins_manager(); let PluginReadParams { @@ -965,10 +1011,11 @@ impl PluginRequestProcessor { ); let share_context = match share_context { Some(context) => { - let auth = self.auth_manager.auth().await; - let remote_plugin_service_config = RemotePluginServiceConfig { - chatgpt_base_url: config.chatgpt_base_url.clone(), - }; + let auth = self + .auth_for_app_server_client(app_server_client_name) + .await; + let remote_plugin_service_config = + remote_plugin_service_config(&config, app_server_client_name); match codex_core_plugins::remote::fetch_remote_plugin_share_context( &remote_plugin_service_config, auth.as_ref(), @@ -1016,6 +1063,7 @@ impl PluginRequestProcessor { &config, &outcome.plugin.apps, Arc::clone(&environment_manager), + app_server_client_name, ) .await; let visible_skills = outcome @@ -1071,10 +1119,11 @@ impl PluginRequestProcessor { "remote plugin read is not enabled for marketplace {remote_marketplace_name}" ))); } - let auth = self.auth_manager.auth().await; - let remote_plugin_service_config = RemotePluginServiceConfig { - chatgpt_base_url: config.chatgpt_base_url.clone(), - }; + let auth = self + .auth_for_app_server_client(app_server_client_name) + .await; + let remote_plugin_service_config = + remote_plugin_service_config(&config, app_server_client_name); validate_remote_plugin_id(&plugin_name)?; let remote_detail = codex_core_plugins::remote::fetch_remote_plugin_detail( &remote_plugin_service_config, @@ -1097,6 +1146,7 @@ impl PluginRequestProcessor { &config, &plugin_apps, Arc::clone(&environment_manager), + app_server_client_name, ) .await; remote_plugin_detail_to_info(remote_detail, app_summaries) @@ -1109,6 +1159,7 @@ impl PluginRequestProcessor { async fn plugin_skill_read_response( &self, params: PluginSkillReadParams, + app_server_client_name: Option<&str>, ) -> Result { let PluginSkillReadParams { remote_marketplace_name, @@ -1129,10 +1180,11 @@ impl PluginRequestProcessor { )); } - let auth = self.auth_manager.auth().await; - let remote_plugin_service_config = RemotePluginServiceConfig { - chatgpt_base_url: config.chatgpt_base_url.clone(), - }; + let auth = self + .auth_for_app_server_client(app_server_client_name) + .await; + let remote_plugin_service_config = + remote_plugin_service_config(&config, app_server_client_name); let remote_skill_detail = codex_core_plugins::remote::fetch_remote_plugin_skill_detail( &remote_plugin_service_config, auth.as_ref(), @@ -1153,8 +1205,11 @@ impl PluginRequestProcessor { async fn plugin_share_save_response( &self, params: PluginShareSaveParams, + app_server_client_name: Option<&str>, ) -> Result { - let (config, auth) = self.load_plugin_share_config_and_auth().await?; + let (config, auth) = self + .load_plugin_share_config_and_auth(app_server_client_name) + .await?; if !config.features.enabled(Feature::PluginSharing) { return Err(invalid_request("plugin sharing is disabled")); } @@ -1183,9 +1238,8 @@ impl PluginRequestProcessor { validate_client_plugin_share_targets(share_targets)?; } - let remote_plugin_service_config = RemotePluginServiceConfig { - chatgpt_base_url: config.chatgpt_base_url.clone(), - }; + let remote_plugin_service_config = + remote_plugin_service_config(&config, app_server_client_name); let access_policy = codex_core_plugins::remote::RemotePluginShareAccessPolicy { discoverability: discoverability.map(remote_plugin_share_discoverability), share_targets: share_targets.map(remote_plugin_share_targets), @@ -1211,8 +1265,11 @@ impl PluginRequestProcessor { async fn plugin_share_update_targets_response( &self, params: PluginShareUpdateTargetsParams, + app_server_client_name: Option<&str>, ) -> Result { - let (config, auth) = self.load_plugin_share_config_and_auth().await?; + let (config, auth) = self + .load_plugin_share_config_and_auth(app_server_client_name) + .await?; if !config.features.enabled(Feature::PluginSharing) { return Err(invalid_request("plugin sharing is disabled")); } @@ -1226,9 +1283,8 @@ impl PluginRequestProcessor { } validate_client_plugin_share_targets(&share_targets)?; - let remote_plugin_service_config = RemotePluginServiceConfig { - chatgpt_base_url: config.chatgpt_base_url.clone(), - }; + let remote_plugin_service_config = + remote_plugin_service_config(&config, app_server_client_name); let result = codex_core_plugins::remote::update_remote_plugin_share_targets( &remote_plugin_service_config, auth.as_ref(), @@ -1254,11 +1310,13 @@ impl PluginRequestProcessor { async fn plugin_share_list_response( &self, _params: PluginShareListParams, + app_server_client_name: Option<&str>, ) -> Result { - let (config, auth) = self.load_plugin_share_config_and_auth().await?; - let remote_plugin_service_config = RemotePluginServiceConfig { - chatgpt_base_url: config.chatgpt_base_url.clone(), - }; + let (config, auth) = self + .load_plugin_share_config_and_auth(app_server_client_name) + .await?; + let remote_plugin_service_config = + remote_plugin_service_config(&config, app_server_client_name); let data = codex_core_plugins::remote::list_remote_plugin_shares( &remote_plugin_service_config, auth.as_ref(), @@ -1285,8 +1343,11 @@ impl PluginRequestProcessor { async fn plugin_share_checkout_response( &self, params: PluginShareCheckoutParams, + app_server_client_name: Option<&str>, ) -> Result { - let (config, auth) = self.load_plugin_share_config_and_auth().await?; + let (config, auth) = self + .load_plugin_share_config_and_auth(app_server_client_name) + .await?; if !config.features.enabled(Feature::PluginSharing) { return Err(invalid_request("plugin sharing is disabled")); } @@ -1295,9 +1356,8 @@ impl PluginRequestProcessor { return Err(invalid_request("invalid remote plugin id")); } - let remote_plugin_service_config = RemotePluginServiceConfig { - chatgpt_base_url: config.chatgpt_base_url.clone(), - }; + let remote_plugin_service_config = + remote_plugin_service_config(&config, app_server_client_name); let result = codex_core_plugins::remote::checkout_remote_plugin_share( &remote_plugin_service_config, auth.as_ref(), @@ -1321,16 +1381,18 @@ impl PluginRequestProcessor { async fn plugin_share_delete_response( &self, params: PluginShareDeleteParams, + app_server_client_name: Option<&str>, ) -> Result { - let (config, auth) = self.load_plugin_share_config_and_auth().await?; + let (config, auth) = self + .load_plugin_share_config_and_auth(app_server_client_name) + .await?; let PluginShareDeleteParams { remote_plugin_id } = params; if remote_plugin_id.is_empty() || !is_valid_remote_plugin_id(&remote_plugin_id) { return Err(invalid_request("invalid remote plugin id")); } - let remote_plugin_service_config = RemotePluginServiceConfig { - chatgpt_base_url: config.chatgpt_base_url.clone(), - }; + let remote_plugin_service_config = + remote_plugin_service_config(&config, app_server_client_name); codex_core_plugins::remote::delete_remote_plugin_share( &remote_plugin_service_config, auth.as_ref(), @@ -1345,18 +1407,22 @@ impl PluginRequestProcessor { async fn load_plugin_share_config_and_auth( &self, + app_server_client_name: Option<&str>, ) -> Result<(Config, Option), JSONRPCErrorError> { let config = self.load_latest_config(/*fallback_cwd*/ None).await?; if !config.features.enabled(Feature::Plugins) { return Err(invalid_request("plugin sharing is not enabled")); } - let auth = self.auth_manager.auth().await; + let auth = self + .auth_for_app_server_client(app_server_client_name) + .await; Ok((config, auth)) } async fn plugin_install_response( &self, params: PluginInstallParams, + app_server_client_name: Option<&str>, ) -> Result { let PluginInstallParams { marketplace_path, @@ -1367,7 +1433,11 @@ impl PluginRequestProcessor { (Some(marketplace_path), None) => marketplace_path, (None, Some(remote_marketplace_name)) => { return self - .remote_plugin_install_response(remote_marketplace_name, plugin_name) + .remote_plugin_install_response( + remote_marketplace_name, + plugin_name, + app_server_client_name, + ) .await; } (Some(_), Some(_)) | (None, None) => { @@ -1378,10 +1448,12 @@ impl PluginRequestProcessor { }; let config_cwd = marketplace_path.as_path().parent().map(Path::to_path_buf); let config = self.load_latest_config(config_cwd.clone()).await?; - let auth = self.auth_manager.auth().await; + let auth = self + .auth_for_app_server_client(app_server_client_name) + .await; if !self - .workspace_codex_plugins_enabled(&config, auth.as_ref()) + .workspace_codex_plugins_enabled(&config, auth.as_ref(), app_server_client_name) .await { return Err(invalid_request( @@ -1418,13 +1490,13 @@ impl PluginRequestProcessor { } let plugin_apps = load_plugin_apps(result.installed_path.as_path()).await; - let auth = self.auth_manager.auth().await; let apps_needing_auth = self .plugin_apps_needing_auth_for_install( &config, auth.as_ref().is_some_and(CodexAuth::is_chatgpt_auth), &result.plugin_id.as_key(), &plugin_apps, + app_server_client_name, ) .await; @@ -1438,6 +1510,7 @@ impl PluginRequestProcessor { &self, remote_marketplace_name: String, remote_plugin_id: String, + app_server_client_name: Option<&str>, ) -> Result { let config = self.load_latest_config(/*fallback_cwd*/ None).await?; if !config.features.enabled(Feature::Plugins) { @@ -1447,10 +1520,11 @@ impl PluginRequestProcessor { } validate_remote_plugin_id(&remote_plugin_id)?; - let auth = self.auth_manager.auth().await; - let remote_plugin_service_config = RemotePluginServiceConfig { - chatgpt_base_url: config.chatgpt_base_url.clone(), - }; + let auth = self + .auth_for_app_server_client(app_server_client_name) + .await; + let remote_plugin_service_config = + remote_plugin_service_config(&config, app_server_client_name); let remote_detail = codex_core_plugins::remote::fetch_remote_plugin_detail_with_download_urls( &remote_plugin_service_config, @@ -1519,6 +1593,7 @@ impl PluginRequestProcessor { &config.plugins_config_input(), auth.clone(), Some(self.effective_plugins_changed_callback()), + Some(http_state_context(&config, app_server_client_name)), ); let mut plugin_metadata = @@ -1540,6 +1615,7 @@ impl PluginRequestProcessor { auth.as_ref().is_some_and(CodexAuth::is_chatgpt_auth), &result.plugin_id.as_key(), &plugin_apps, + app_server_client_name, ) .await; @@ -1555,18 +1631,25 @@ impl PluginRequestProcessor { is_chatgpt_auth: bool, plugin_id: &str, plugin_apps: &[codex_plugin::AppConnectorId], + app_server_client_name: Option<&str>, ) -> Vec { if plugin_apps.is_empty() || !config.features.apps_enabled_for_auth(is_chatgpt_auth) { return Vec::new(); } let environment_manager = self.thread_manager.environment_manager(); + let http_state = http_state_context(config, app_server_client_name); let (all_connectors_result, accessible_connectors_result) = tokio::join!( - connectors::list_all_connectors_with_options(config, /*force_refetch*/ false), - connectors::list_accessible_connectors_from_mcp_tools_with_environment_manager( + connectors::list_all_connectors_with_options_and_http_state( + config, + /*force_refetch*/ false, + Some(http_state.clone()), + ), + connectors::list_accessible_connectors_from_mcp_tools_with_environment_manager_and_http_state( config, /*force_refetch*/ true, - Arc::clone(&environment_manager) + Arc::clone(&environment_manager), + Some(http_state.clone()), ), ); @@ -1577,7 +1660,7 @@ impl PluginRequestProcessor { plugin = plugin_id, "failed to load app metadata after plugin install: {err:#}" ); - connectors::list_cached_all_connectors(config) + connectors::list_cached_all_connectors_with_http_state(config, Some(&http_state)) .await .unwrap_or_default() } @@ -1591,9 +1674,12 @@ impl PluginRequestProcessor { "failed to load accessible apps after plugin install: {err:#}" ); ( - connectors::list_cached_accessible_connectors_from_mcp_tools(config) - .await - .unwrap_or_default(), + connectors::list_cached_accessible_connectors_from_mcp_tools_with_http_state( + config, + Some(&http_state), + ) + .await + .unwrap_or_default(), false, ) } @@ -1697,6 +1783,7 @@ impl PluginRequestProcessor { async fn plugin_uninstall_response( &self, params: PluginUninstallParams, + app_server_client_name: Option<&str>, ) -> Result { let PluginUninstallParams { plugin_id } = params; if codex_plugin::PluginId::parse(&plugin_id).is_err() @@ -1705,7 +1792,9 @@ impl PluginRequestProcessor { return Err(invalid_request("invalid remote plugin id")); } if is_valid_remote_plugin_id(&plugin_id) { - return self.remote_plugin_uninstall_response(plugin_id).await; + return self + .remote_plugin_uninstall_response(plugin_id, app_server_client_name) + .await; } let plugins_manager = self.thread_manager.plugins_manager(); @@ -1788,6 +1877,7 @@ impl PluginRequestProcessor { async fn remote_plugin_uninstall_response( &self, plugin_id: String, + app_server_client_name: Option<&str>, ) -> Result { let config = self.load_latest_config(/*fallback_cwd*/ None).await?; if !config.features.enabled(Feature::Plugins) { @@ -1795,10 +1885,11 @@ impl PluginRequestProcessor { } validate_remote_plugin_id(&plugin_id)?; - let auth = self.auth_manager.auth().await; - let remote_plugin_service_config = RemotePluginServiceConfig { - chatgpt_base_url: config.chatgpt_base_url.clone(), - }; + let auth = self + .auth_for_app_server_client(app_server_client_name) + .await; + let remote_plugin_service_config = + remote_plugin_service_config(&config, app_server_client_name); let uninstall_result = codex_core_plugins::remote::uninstall_remote_plugin( &remote_plugin_service_config, auth.as_ref(), @@ -1819,6 +1910,7 @@ impl PluginRequestProcessor { &config.plugins_config_input(), auth.clone(), Some(self.effective_plugins_changed_callback()), + Some(http_state_context(&config, app_server_client_name)), ); } @@ -1833,29 +1925,37 @@ async fn load_plugin_app_summaries( config: &Config, plugin_apps: &[codex_plugin::AppConnectorId], environment_manager: Arc, + app_server_client_name: Option<&str>, ) -> Vec { if plugin_apps.is_empty() { return Vec::new(); } - let connectors = - match connectors::list_all_connectors_with_options(config, /*force_refetch*/ false).await { - Ok(connectors) => connectors, - Err(err) => { - warn!("failed to load app metadata for plugin/read: {err:#}"); - connectors::list_cached_all_connectors(config) - .await - .unwrap_or_default() - } - }; + let http_state = http_state_context(config, app_server_client_name); + let connectors = match connectors::list_all_connectors_with_options_and_http_state( + config, + /*force_refetch*/ false, + Some(http_state.clone()), + ) + .await + { + Ok(connectors) => connectors, + Err(err) => { + warn!("failed to load app metadata for plugin/read: {err:#}"); + connectors::list_cached_all_connectors_with_http_state(config, Some(&http_state)) + .await + .unwrap_or_default() + } + }; let plugin_connectors = connectors::connectors_for_plugin_apps(connectors, plugin_apps); let accessible_connectors = - match connectors::list_accessible_connectors_from_mcp_tools_with_environment_manager( + match connectors::list_accessible_connectors_from_mcp_tools_with_environment_manager_and_http_state( config, /*force_refetch*/ false, environment_manager, + Some(http_state), ) .await { diff --git a/codex-rs/chatgpt/Cargo.toml b/codex-rs/chatgpt/Cargo.toml index 6b0e010964..289d8f92a7 100644 --- a/codex-rs/chatgpt/Cargo.toml +++ b/codex-rs/chatgpt/Cargo.toml @@ -15,6 +15,7 @@ codex-connectors = { workspace = true } codex-core = { workspace = true } codex-core-plugins = { workspace = true } codex-git-utils = { workspace = true } +codex-http-state = { workspace = true } codex-login = { workspace = true } codex-model-provider = { workspace = true } codex-plugin = { workspace = true } diff --git a/codex-rs/chatgpt/src/chatgpt_client.rs b/codex-rs/chatgpt/src/chatgpt_client.rs index 372f62e696..41484ef25d 100644 --- a/codex-rs/chatgpt/src/chatgpt_client.rs +++ b/codex-rs/chatgpt/src/chatgpt_client.rs @@ -1,4 +1,5 @@ use codex_core::config::Config; +use codex_http_state::HttpStateContext; use codex_login::AuthManager; use codex_login::default_client::create_client; @@ -9,26 +10,34 @@ use std::time::Duration; const OAI_PRODUCT_SKU_HEADER: &str = "OAI-Product-Sku"; const CODEX_PRODUCT_SKU: &str = "codex"; -/// Make a GET request to the ChatGPT backend API. -pub(crate) async fn chatgpt_get_request( +pub(crate) async fn chatgpt_get_request_with_http_state( config: &Config, path: String, + http_state: HttpStateContext, ) -> anyhow::Result { - chatgpt_get_request_with_timeout(config, path, /*timeout*/ None).await + chatgpt_get_request_with_timeout_and_http_state( + config, + path, + /*timeout*/ None, + Some(http_state), + ) + .await } -pub(crate) async fn chatgpt_get_request_with_timeout( +pub(crate) async fn chatgpt_get_request_with_timeout_and_http_state( config: &Config, path: String, timeout: Option, + http_state: Option, ) -> anyhow::Result { let chatgpt_base_url = &config.chatgpt_base_url; let auth_manager = AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await; - let auth = auth_manager - .auth() - .await - .ok_or_else(|| anyhow::anyhow!("ChatGPT auth not available"))?; + let auth = match http_state.as_ref() { + Some(http_state) => auth_manager.auth_for_surface(http_state.surface()).await, + None => auth_manager.auth().await, + } + .ok_or_else(|| anyhow::anyhow!("ChatGPT auth not available"))?; anyhow::ensure!( auth.uses_codex_backend(), "ChatGPT backend requests require Codex backend auth" @@ -46,9 +55,16 @@ pub(crate) async fn chatgpt_get_request_with_timeout( path.trim_start_matches('/') ); + let auth_provider = codex_model_provider::with_native_integrity_state( + codex_model_provider::auth_provider_from_auth(&auth), + Some(&auth), + http_state, + ); + let mut request_headers = Default::default(); + auth_provider.add_auth_headers_for_url(&url, &mut request_headers); let mut request = client .get(&url) - .headers(codex_model_provider::auth_provider_from_auth(&auth).to_auth_headers()) + .headers(request_headers.clone()) .header(OAI_PRODUCT_SKU_HEADER, CODEX_PRODUCT_SKU) .header("Content-Type", "application/json"); @@ -57,6 +73,7 @@ pub(crate) async fn chatgpt_get_request_with_timeout( } let response = request.send().await.context("Failed to send request")?; + auth_provider.observe_response_headers(&url, &request_headers, response.headers()); if response.status().is_success() { let result: T = response diff --git a/codex-rs/chatgpt/src/connectors.rs b/codex-rs/chatgpt/src/connectors.rs index 63b34e7355..fc75a489a5 100644 --- a/codex-rs/chatgpt/src/connectors.rs +++ b/codex-rs/chatgpt/src/connectors.rs @@ -2,7 +2,7 @@ use std::collections::HashMap; use std::collections::HashSet; use std::time::Duration; -use crate::chatgpt_client::chatgpt_get_request_with_timeout; +use crate::chatgpt_client::chatgpt_get_request_with_timeout_and_http_state; use codex_app_server_protocol::AppInfo; use codex_connectors::ConnectorDirectoryCacheContext; @@ -14,11 +14,14 @@ use codex_connectors::merge::merge_plugin_connectors; use codex_core::config::Config; pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools; pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_environment_manager; +pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_environment_manager_and_http_state; pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_options; pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_options_and_status; pub use codex_core::connectors::list_cached_accessible_connectors_from_mcp_tools; +pub use codex_core::connectors::list_cached_accessible_connectors_from_mcp_tools_with_http_state; pub use codex_core::connectors::with_app_enabled_state; use codex_core_plugins::PluginsManager; +use codex_http_state::HttpStateContext; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::default_client::originator; @@ -26,20 +29,27 @@ use codex_plugin::AppConnectorId; const DIRECTORY_CONNECTORS_TIMEOUT: Duration = Duration::from_secs(60); -async fn apps_enabled(config: &Config) -> bool { +async fn auth(config: &Config, http_state: Option<&HttpStateContext>) -> Option { let auth_manager = AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await; - let auth = auth_manager.auth().await; + match http_state { + Some(http_state) => auth_manager.auth_for_surface(http_state.surface()).await, + None => auth_manager.auth().await, + } +} + +async fn apps_enabled(config: &Config, http_state: Option<&HttpStateContext>) -> bool { + let auth = auth(config, http_state).await; config .features .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend)) } -async fn connector_auth(config: &Config) -> anyhow::Result { - let auth_manager = - AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await; - let auth = auth_manager - .auth() +async fn connector_auth( + config: &Config, + http_state: Option<&HttpStateContext>, +) -> anyhow::Result { + let auth = auth(config, http_state) .await .ok_or_else(|| anyhow::anyhow!("ChatGPT auth not available"))?; anyhow::ensure!( @@ -50,7 +60,7 @@ async fn connector_auth(config: &Config) -> anyhow::Result { } pub async fn list_connectors(config: &Config) -> anyhow::Result> { - if !apps_enabled(config).await { + if !apps_enabled(config, /*http_state*/ None).await { return Ok(Vec::new()); } let (connectors_result, accessible_result) = tokio::join!( @@ -72,11 +82,18 @@ pub async fn list_all_connectors(config: &Config) -> anyhow::Result } pub async fn list_cached_all_connectors(config: &Config) -> Option> { - if !apps_enabled(config).await { + list_cached_all_connectors_with_http_state(config, /*http_state*/ None).await +} + +pub async fn list_cached_all_connectors_with_http_state( + config: &Config, + http_state: Option<&HttpStateContext>, +) -> Option> { + if !apps_enabled(config, http_state).await { return Some(Vec::new()); } - let auth = connector_auth(config).await.ok()?; + let auth = connector_auth(config, http_state).await.ok()?; let cache_context = connector_directory_cache_context(config, &auth); let connectors = codex_connectors::cached_directory_connectors(&cache_context)?; let connectors = merge_plugin_connectors( @@ -96,22 +113,35 @@ pub async fn list_all_connectors_with_options( config: &Config, force_refetch: bool, ) -> anyhow::Result> { - if !apps_enabled(config).await { + list_all_connectors_with_options_and_http_state(config, force_refetch, /*http_state*/ None) + .await +} + +pub async fn list_all_connectors_with_options_and_http_state( + config: &Config, + force_refetch: bool, + http_state: Option, +) -> anyhow::Result> { + if !apps_enabled(config, http_state.as_ref()).await { return Ok(Vec::new()); } - let auth = connector_auth(config).await?; + let auth = connector_auth(config, http_state.as_ref()).await?; let cache_context = connector_directory_cache_context(config, &auth); let connectors = codex_connectors::list_all_connectors_with_options( cache_context, auth.is_workspace_account(), force_refetch, - |path| async move { - chatgpt_get_request_with_timeout::( - config, - path, - Some(DIRECTORY_CONNECTORS_TIMEOUT), - ) - .await + |path| { + let http_state = http_state.clone(); + async move { + chatgpt_get_request_with_timeout_and_http_state::( + config, + path, + Some(DIRECTORY_CONNECTORS_TIMEOUT), + http_state, + ) + .await + } }, ) .await?; diff --git a/codex-rs/chatgpt/src/get_task.rs b/codex-rs/chatgpt/src/get_task.rs index 9301ffc38d..1aaede3388 100644 --- a/codex-rs/chatgpt/src/get_task.rs +++ b/codex-rs/chatgpt/src/get_task.rs @@ -1,7 +1,9 @@ use codex_core::config::Config; +use codex_http_state::HttpStateContext; +use codex_http_state::HttpStateSurface; use serde::Deserialize; -use crate::chatgpt_client::chatgpt_get_request; +use crate::chatgpt_client::chatgpt_get_request_with_http_state; #[derive(Debug, Deserialize)] pub struct GetTaskResponse { @@ -36,5 +38,10 @@ pub struct OutputDiff { pub(crate) async fn get_task(config: &Config, task_id: String) -> anyhow::Result { let path = format!("/wham/tasks/{task_id}"); - chatgpt_get_request(config, path).await + chatgpt_get_request_with_http_state( + config, + path, + HttpStateContext::new(config.codex_home.to_path_buf(), HttpStateSurface::CodexCli), + ) + .await } diff --git a/codex-rs/chatgpt/src/workspace_settings.rs b/codex-rs/chatgpt/src/workspace_settings.rs index a177215518..1383d0b0cf 100644 --- a/codex-rs/chatgpt/src/workspace_settings.rs +++ b/codex-rs/chatgpt/src/workspace_settings.rs @@ -5,10 +5,11 @@ use std::time::Instant; use anyhow::Context; use codex_core::config::Config; +use codex_http_state::HttpStateContext; use codex_login::CodexAuth; use serde::Deserialize; -use crate::chatgpt_client::chatgpt_get_request_with_timeout; +use crate::chatgpt_client::chatgpt_get_request_with_timeout_and_http_state; const WORKSPACE_SETTINGS_TIMEOUT: Duration = Duration::from_secs(10); const WORKSPACE_SETTINGS_CACHE_TTL: Duration = Duration::from_secs(15 * 60); @@ -85,6 +86,18 @@ pub async fn codex_plugins_enabled_for_workspace( config: &Config, auth: Option<&CodexAuth>, cache: Option<&WorkspaceSettingsCache>, +) -> anyhow::Result { + codex_plugins_enabled_for_workspace_with_http_state( + config, auth, cache, /*http_state*/ None, + ) + .await +} + +pub async fn codex_plugins_enabled_for_workspace_with_http_state( + config: &Config, + auth: Option<&CodexAuth>, + cache: Option<&WorkspaceSettingsCache>, + http_state: Option, ) -> anyhow::Result { let Some(auth) = auth else { return Ok(true); @@ -115,10 +128,11 @@ pub async fn codex_plugins_enabled_for_workspace( } let encoded_account_id = encode_path_segment(account_id); - let settings: WorkspaceSettingsResponse = chatgpt_get_request_with_timeout( + let settings: WorkspaceSettingsResponse = chatgpt_get_request_with_timeout_and_http_state( config, format!("/accounts/{encoded_account_id}/settings"), Some(WORKSPACE_SETTINGS_TIMEOUT), + http_state, ) .await?; diff --git a/codex-rs/codex-mcp/Cargo.toml b/codex-rs/codex-mcp/Cargo.toml index d7634f0aae..c35ab3c458 100644 --- a/codex-rs/codex-mcp/Cargo.toml +++ b/codex-rs/codex-mcp/Cargo.toml @@ -19,6 +19,7 @@ codex-async-utils = { workspace = true } codex-api = { workspace = true } codex-config = { workspace = true } codex-exec-server = { workspace = true } +codex-http-state = { workspace = true } codex-login = { workspace = true } codex-model-provider = { workspace = true } codex-otel = { workspace = true } diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index d15a2e1500..27930a2f38 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -43,6 +43,7 @@ use async_channel::Sender; use codex_config::Constrained; use codex_config::McpServerTransportConfig; use codex_config::types::OAuthCredentialsStoreMode; +use codex_http_state::HttpStateContext; use codex_login::CodexAuth; use codex_protocol::mcp::CallToolResult; use codex_protocol::mcp::McpServerInfo; @@ -220,6 +221,7 @@ impl McpConnectionManager { initial_permission_profile: PermissionProfile, runtime_context: McpRuntimeContext, codex_home: PathBuf, + http_state: Option, codex_apps_tools_cache_key: CodexAppsToolsCacheKey, host_owned_codex_apps_enabled: bool, prefix_mcp_tool_names: bool, @@ -239,9 +241,13 @@ impl McpConnectionManager { ); let tool_plugin_provenance = Arc::new(tool_plugin_provenance); let startup_submit_id = submit_id.clone(); - let codex_apps_auth_provider = auth - .filter(|auth| auth.uses_codex_backend()) - .map(codex_model_provider::auth_provider_from_auth); + let codex_apps_auth_provider = auth.filter(|auth| auth.uses_codex_backend()).map(|auth| { + codex_model_provider::with_native_integrity_state( + codex_model_provider::auth_provider_from_auth(auth), + Some(auth), + http_state, + ) + }); let mcp_servers = mcp_servers.clone(); for (server_name, server) in mcp_servers .into_iter() diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 60a0026eeb..0a959ad940 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -1180,6 +1180,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { PathBuf::from("/tmp"), ), codex_home.path().to_path_buf(), + /*http_state*/ None, CodexAppsToolsCacheKey { account_id: None, chatgpt_user_id: None, diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 9633655846..203db3971d 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -22,6 +22,7 @@ use codex_config::McpServerConfig; use codex_config::McpServerTransportConfig; use codex_config::types::AppToolApproval; use codex_config::types::OAuthCredentialsStoreMode; +use codex_http_state::HttpStateContext; use codex_login::CodexAuth; use codex_plugin::PluginCapabilitySummary; use codex_protocol::mcp::McpServerInfo; @@ -264,6 +265,7 @@ pub async fn read_mcp_resource( config: &McpConfig, auth: Option<&CodexAuth>, runtime_context: McpRuntimeContext, + http_state: Option, server: &str, uri: &str, ) -> anyhow::Result { @@ -288,6 +290,7 @@ pub async fn read_mcp_resource( PermissionProfile::default(), runtime_context, config.codex_home.clone(), + http_state, codex_apps_tools_cache_key(auth), host_owned_codex_apps_enabled, config.prefix_mcp_tool_names, @@ -320,6 +323,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( auth: Option<&CodexAuth>, submit_id: String, runtime_context: McpRuntimeContext, + http_state: Option, detail: McpSnapshotDetail, ) -> McpServerStatusSnapshot { let mcp_servers = effective_mcp_servers(config, auth); @@ -358,6 +362,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( PermissionProfile::default(), runtime_context, config.codex_home.clone(), + http_state, codex_apps_tools_cache_key(auth), host_owned_codex_apps_enabled, config.prefix_mcp_tool_names, diff --git a/codex-rs/core-plugins/Cargo.toml b/codex-rs/core-plugins/Cargo.toml index d0096e4455..2499e07ae7 100644 --- a/codex-rs/core-plugins/Cargo.toml +++ b/codex-rs/core-plugins/Cargo.toml @@ -15,12 +15,14 @@ workspace = true [dependencies] anyhow = { workspace = true } codex-analytics = { workspace = true } +codex-api = { workspace = true } codex-app-server-protocol = { workspace = true } codex-config = { workspace = true } codex-core-skills = { workspace = true } codex-exec-server = { workspace = true } codex-git-utils = { workspace = true } codex-hooks = { workspace = true } +codex-http-state = { workspace = true } codex-login = { workspace = true } codex-model-provider = { workspace = true } codex-otel = { workspace = true } diff --git a/codex-rs/core-plugins/src/manager.rs b/codex-rs/core-plugins/src/manager.rs index 8022c6c86d..b35e1fd979 100644 --- a/codex-rs/core-plugins/src/manager.rs +++ b/codex-rs/core-plugins/src/manager.rs @@ -58,6 +58,7 @@ use codex_config::types::PluginConfig; use codex_config::version_for_toml; use codex_core_skills::SkillMetadata; use codex_hooks::plugin_hook_declarations; +use codex_http_state::HttpStateContext; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_plugin::AppConnectorId; @@ -173,8 +174,17 @@ struct ConfiguredMarketplaceUpgradeState { } fn remote_plugin_service_config(config: &PluginsConfigInput) -> RemotePluginServiceConfig { - RemotePluginServiceConfig { - chatgpt_base_url: config.chatgpt_base_url.clone(), + RemotePluginServiceConfig::new(config.chatgpt_base_url.clone()) +} + +fn remote_plugin_service_config_with_http_state( + config: &PluginsConfigInput, + http_state: Option, +) -> RemotePluginServiceConfig { + let service_config = remote_plugin_service_config(config); + match http_state { + Some(http_state) => service_config.with_http_state(http_state), + None => service_config, } } @@ -624,11 +634,40 @@ impl PluginsManager { visible_scopes: &[RemotePluginScope], on_effective_plugins_changed: Option>, ) -> Result, RemotePluginCatalogError> { - let plugins = crate::remote::fetch_remote_installed_plugins( - &remote_plugin_service_config(config), + self.build_and_cache_remote_installed_plugin_marketplaces_with_service_config( + remote_plugin_service_config(config), auth, + visible_scopes, + on_effective_plugins_changed, ) - .await?; + .await + } + + pub async fn build_and_cache_remote_installed_plugin_marketplaces_with_http_state( + &self, + config: &PluginsConfigInput, + auth: Option<&CodexAuth>, + visible_scopes: &[RemotePluginScope], + on_effective_plugins_changed: Option>, + http_state: HttpStateContext, + ) -> Result, RemotePluginCatalogError> { + self.build_and_cache_remote_installed_plugin_marketplaces_with_service_config( + remote_plugin_service_config(config).with_http_state(http_state), + auth, + visible_scopes, + on_effective_plugins_changed, + ) + .await + } + + async fn build_and_cache_remote_installed_plugin_marketplaces_with_service_config( + &self, + service_config: RemotePluginServiceConfig, + auth: Option<&CodexAuth>, + visible_scopes: &[RemotePluginScope], + on_effective_plugins_changed: Option>, + ) -> Result, RemotePluginCatalogError> { + let plugins = crate::remote::fetch_remote_installed_plugins(&service_config, auth).await?; let marketplaces = crate::remote::group_remote_installed_plugins_by_marketplaces(&plugins, visible_scopes); let changed = self.write_remote_installed_plugins_cache(plugins); @@ -671,12 +710,14 @@ impl PluginsManager { config: &PluginsConfigInput, auth: Option, on_effective_plugins_changed: Option>, + http_state: Option, ) { self.maybe_start_remote_installed_plugins_cache_refresh_with_notify( config, auth, RemoteInstalledPluginsCacheRefreshNotify::IfCacheChanged, on_effective_plugins_changed, + http_state, ); } @@ -685,12 +726,14 @@ impl PluginsManager { config: &PluginsConfigInput, auth: Option, on_effective_plugins_changed: Option>, + http_state: Option, ) { self.maybe_start_remote_installed_plugins_cache_refresh_with_notify( config, auth, RemoteInstalledPluginsCacheRefreshNotify::AfterSuccessfulRefresh, on_effective_plugins_changed, + http_state, ); } @@ -700,6 +743,7 @@ impl PluginsManager { auth: Option, notify: RemoteInstalledPluginsCacheRefreshNotify, on_effective_plugins_changed: Option>, + http_state: Option, ) { if !config.plugins_enabled { return; @@ -707,7 +751,7 @@ impl PluginsManager { self.schedule_remote_installed_plugins_cache_refresh( RemoteInstalledPluginsCacheRefreshRequest { - service_config: remote_plugin_service_config(config), + service_config: remote_plugin_service_config_with_http_state(config, http_state), auth, notify, on_effective_plugins_changed, @@ -720,6 +764,7 @@ impl PluginsManager { config: &PluginsConfigInput, auth: Option, on_effective_plugins_changed: Option>, + http_state: Option, ) { if !config.plugins_enabled { return; @@ -728,17 +773,19 @@ impl PluginsManager { let manager = Arc::clone(self); let config_for_refresh = config.clone(); let auth_for_refresh = auth.clone(); + let http_state_for_refresh = http_state.clone(); let on_local_cache_changed = Arc::new(move || { manager.maybe_start_remote_installed_plugins_cache_refresh_after_mutation( &config_for_refresh, auth_for_refresh.clone(), on_effective_plugins_changed.clone(), + http_state_for_refresh.clone(), ); }); crate::remote::maybe_start_remote_installed_plugin_bundle_sync( self.codex_home.clone(), - remote_plugin_service_config(config), + remote_plugin_service_config_with_http_state(config, http_state), auth, Some(on_local_cache_changed), ); @@ -750,17 +797,20 @@ impl PluginsManager { auth: Option, roots: &[AbsolutePathBuf], on_effective_plugins_changed: Option>, + http_state: Option, ) { self.maybe_start_non_curated_plugin_cache_refresh(roots); self.maybe_start_remote_installed_plugins_cache_refresh( config, auth.clone(), on_effective_plugins_changed.clone(), + http_state.clone(), ); self.maybe_start_remote_installed_plugin_bundle_sync( config, auth, on_effective_plugins_changed, + http_state, ); } @@ -816,6 +866,34 @@ impl PluginsManager { &self, config: &PluginsConfigInput, auth: Option<&CodexAuth>, + ) -> Result, RemotePluginFetchError> { + self.featured_plugin_ids_for_config_with_service_config( + config, + auth, + remote_plugin_service_config(config), + ) + .await + } + + pub async fn featured_plugin_ids_for_config_with_http_state( + &self, + config: &PluginsConfigInput, + auth: Option<&CodexAuth>, + http_state: HttpStateContext, + ) -> Result, RemotePluginFetchError> { + self.featured_plugin_ids_for_config_with_service_config( + config, + auth, + remote_plugin_service_config(config).with_http_state(http_state), + ) + .await + } + + async fn featured_plugin_ids_for_config_with_service_config( + &self, + config: &PluginsConfigInput, + auth: Option<&CodexAuth>, + service_config: RemotePluginServiceConfig, ) -> Result, RemotePluginFetchError> { if !config.plugins_enabled { return Ok(Vec::new()); @@ -826,7 +904,7 @@ impl PluginsManager { return Ok(featured_plugin_ids); } let featured_plugin_ids = crate::remote_legacy::fetch_remote_featured_plugin_ids( - &remote_plugin_service_config(config), + &service_config, auth, self.restriction_product, ) @@ -1570,11 +1648,13 @@ impl PluginsManager { &config_for_remote_sync, auth.clone(), on_effective_plugins_changed.clone(), + /*http_state*/ None, ); manager.maybe_start_remote_installed_plugin_bundle_sync( &config_for_remote_sync, auth.clone(), on_effective_plugins_changed, + /*http_state*/ None, ); if config_for_remote_sync.remote_plugin_enabled { match crate::remote::fetch_and_cache_global_remote_plugin_catalog( diff --git a/codex-rs/core-plugins/src/remote.rs b/codex-rs/core-plugins/src/remote.rs index 3c86f6b14b..7be30a8f97 100644 --- a/codex-rs/core-plugins/src/remote.rs +++ b/codex-rs/core-plugins/src/remote.rs @@ -1,16 +1,19 @@ use crate::store::PLUGINS_CACHE_DIR; use crate::store::PluginStore; +use codex_api::SharedAuthProvider; use codex_app_server_protocol::JSONRPCErrorError; use codex_app_server_protocol::PluginAuthPolicy; use codex_app_server_protocol::PluginAvailability; use codex_app_server_protocol::PluginInstallPolicy; use codex_app_server_protocol::PluginInterface; use codex_app_server_protocol::SkillInterface; +use codex_http_state::HttpStateContext; use codex_login::CodexAuth; use codex_login::default_client::build_reqwest_client; use codex_plugin::PluginId; use codex_utils_absolute_path::AbsolutePathBuf; use reqwest::RequestBuilder; +use reqwest::header::HeaderMap; use serde::Deserialize; use serde::Serialize; use serde_json::Value as JsonValue; @@ -24,6 +27,7 @@ use std::time::Duration; use url::Url; mod catalog_cache; +pub(crate) mod http_state; mod remote_installed_plugin_sync; mod share; @@ -91,9 +95,24 @@ const REMOTE_INSTALLED_MARKETPLACE_DISPLAY_ORDER: [(&str, &str); 5] = [ ), ]; -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone)] pub struct RemotePluginServiceConfig { pub chatgpt_base_url: String, + http_state: Option, +} + +impl RemotePluginServiceConfig { + pub fn new(chatgpt_base_url: String) -> Self { + Self { + chatgpt_base_url, + http_state: None, + } + } + + pub fn with_http_state(mut self, http_state: HttpStateContext) -> Self { + self.http_state = Some(http_state); + self + } } #[derive(Debug, Clone, PartialEq)] @@ -871,7 +890,7 @@ pub async fn fetch_remote_plugin_skill_detail( let url = remote_plugin_skill_detail_url(config, plugin_id, skill_name)?; let client = build_reqwest_client(); - let request = authenticated_request(client.get(&url), auth)?; + let request = authenticated_request(config, client.get(&url), auth, &url); let response: RemotePluginSkillDetailResponse = send_and_decode(request, &url).await?; if response.plugin_id != plugin_id { return Err(RemotePluginCatalogError::UnexpectedPluginId { @@ -973,7 +992,7 @@ pub async fn install_remote_plugin( let base_url = config.chatgpt_base_url.trim_end_matches('/'); let url = format!("{base_url}/ps/plugins/{plugin_id}/install"); let client = build_reqwest_client(); - let request = authenticated_request(client.post(&url), auth)?; + let request = authenticated_request(config, client.post(&url), auth, &url); let response: RemotePluginMutationResponse = send_and_decode(request, &url).await?; if response.id != plugin_id { return Err(RemotePluginCatalogError::UnexpectedPluginId { @@ -1009,7 +1028,7 @@ pub async fn uninstall_remote_plugin( let base_url = config.chatgpt_base_url.trim_end_matches('/'); let url = format!("{base_url}/plugins/{plugin_id}/uninstall"); let client = build_reqwest_client(); - let request = authenticated_request(client.post(&url), auth)?; + let request = authenticated_request(config, client.post(&url), auth, &url); let response: RemotePluginMutationResponse = send_and_decode(request, &url).await?; if response.id != plugin_id { return Err(RemotePluginCatalogError::UnexpectedPluginId { @@ -1399,7 +1418,7 @@ async fn get_remote_plugin_list_page( let base_url = config.chatgpt_base_url.trim_end_matches('/'); let url = format!("{base_url}/ps/plugins/list"); let client = build_reqwest_client(); - let mut request = authenticated_request(client.get(&url), auth)?; + let mut request = authenticated_request(config, client.get(&url), auth, &url); request = request.query(&[("scope", scope.api_value())]); request = request.query(&[("limit", REMOTE_PLUGIN_LIST_PAGE_LIMIT)]); if let Some(collection) = collection { @@ -1419,7 +1438,7 @@ async fn get_remote_shared_workspace_plugins_page( let base_url = config.chatgpt_base_url.trim_end_matches('/'); let url = format!("{base_url}/ps/plugins/workspace/shared"); let client = build_reqwest_client(); - let mut request = authenticated_request(client.get(&url), auth)?; + let mut request = authenticated_request(config, client.get(&url), auth, &url); request = request.query(&[("limit", REMOTE_PLUGIN_LIST_PAGE_LIMIT)]); if let Some(page_token) = page_token { request = request.query(&[("pageToken", page_token)]); @@ -1437,7 +1456,7 @@ async fn get_remote_plugin_installed_page( let base_url = config.chatgpt_base_url.trim_end_matches('/'); let url = format!("{base_url}/ps/plugins/installed"); let client = build_reqwest_client(); - let mut request = authenticated_request(client.get(&url), auth)?; + let mut request = authenticated_request(config, client.get(&url), auth, &url); request = request.query(&[("scope", scope.api_value())]); if include_download_urls { request = request.query(&[("includeDownloadUrls", true)]); @@ -1457,7 +1476,7 @@ async fn fetch_plugin_detail( let base_url = config.chatgpt_base_url.trim_end_matches('/'); let url = format!("{base_url}/ps/plugins/{plugin_id}"); let client = build_reqwest_client(); - let mut request = authenticated_request(client.get(&url), auth)?; + let mut request = authenticated_request(config, client.get(&url), auth, &url); if include_download_urls { request = request.query(&[("includeDownloadUrls", true)]); } @@ -1496,25 +1515,62 @@ fn ensure_chatgpt_auth(auth: Option<&CodexAuth>) -> Result<&CodexAuth, RemotePlu } fn authenticated_request( + config: &RemotePluginServiceConfig, request: RequestBuilder, auth: &CodexAuth, -) -> Result { - Ok(request - .timeout(REMOTE_PLUGIN_CATALOG_TIMEOUT) - .headers(codex_model_provider::auth_provider_from_auth(auth).to_auth_headers())) + url: &str, +) -> AuthenticatedRequest { + let (auth_provider, request_headers) = http_state::request_headers(config, auth, url); + AuthenticatedRequest { + request: request + .timeout(REMOTE_PLUGIN_CATALOG_TIMEOUT) + .headers(request_headers.clone()), + auth_provider, + request_headers, + } +} + +struct AuthenticatedRequest { + request: RequestBuilder, + auth_provider: SharedAuthProvider, + request_headers: HeaderMap, +} + +impl AuthenticatedRequest { + fn query(mut self, query: &T) -> Self { + self.request = self.request.query(query); + self + } + + fn json(mut self, body: &T) -> Self { + self.request = self.request.json(body); + self + } + + async fn send(self, url: &str) -> Result { + let response = + self.request + .send() + .await + .map_err(|source| RemotePluginCatalogError::Request { + url: url.to_string(), + source, + })?; + http_state::observe_response_headers( + &self.auth_provider, + url, + &self.request_headers, + response.headers(), + ); + Ok(response) + } } async fn send_and_decode Deserialize<'de>>( - request: RequestBuilder, + request: AuthenticatedRequest, url: &str, ) -> Result { - let response = request - .send() - .await - .map_err(|source| RemotePluginCatalogError::Request { - url: url.to_string(), - source, - })?; + let response = request.send(url).await?; let status = response.status(); let body = response.text().await.unwrap_or_default(); if !status.is_success() { diff --git a/codex-rs/core-plugins/src/remote/http_state.rs b/codex-rs/core-plugins/src/remote/http_state.rs new file mode 100644 index 0000000000..2cabfb156f --- /dev/null +++ b/codex-rs/core-plugins/src/remote/http_state.rs @@ -0,0 +1,32 @@ +use super::RemotePluginServiceConfig; +use codex_api::SharedAuthProvider; +use codex_login::CodexAuth; +use reqwest::header::HeaderMap; + +pub(crate) fn request_headers( + config: &RemotePluginServiceConfig, + auth: &CodexAuth, + url: &str, +) -> (SharedAuthProvider, HeaderMap) { + let auth_provider = codex_model_provider::with_native_integrity_state( + codex_model_provider::auth_provider_from_auth(auth), + Some(auth), + config.http_state.clone(), + ); + let mut request_headers = HeaderMap::new(); + auth_provider.add_auth_headers_for_url(url, &mut request_headers); + (auth_provider, request_headers) +} + +pub(crate) fn observe_response_headers( + auth_provider: &SharedAuthProvider, + url: &str, + request_headers: &HeaderMap, + response_headers: &HeaderMap, +) { + auth_provider.observe_response_headers(url, request_headers, response_headers); +} + +#[cfg(test)] +#[path = "http_state_tests.rs"] +mod tests; diff --git a/codex-rs/core-plugins/src/remote/http_state_tests.rs b/codex-rs/core-plugins/src/remote/http_state_tests.rs new file mode 100644 index 0000000000..f874c2fe43 --- /dev/null +++ b/codex-rs/core-plugins/src/remote/http_state_tests.rs @@ -0,0 +1,42 @@ +use super::*; +use codex_http_state::HttpStateContext; +use codex_http_state::HttpStateSurface; +use pretty_assertions::assert_eq; +use reqwest::header::HeaderValue; +use tempfile::TempDir; + +const REQUEST_URL: &str = "https://chatgpt.com/backend-api/ps/plugins/list"; + +#[test] +fn plugin_request_headers_attach_and_rotate_native_integrity_state() { + let codex_home = TempDir::new().expect("create temp dir"); + let state = HttpStateContext::new(codex_home.path().to_path_buf(), HttpStateSurface::CodexCli); + state + .set_for_surface(HttpStateSurface::CodexCli, "state-before".to_string()) + .expect("seed state"); + let config = RemotePluginServiceConfig::new("https://chatgpt.com/backend-api".to_string()) + .with_http_state(state.clone()); + let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing(); + + let (auth_provider, request_headers) = request_headers(&config, &auth, REQUEST_URL); + assert_eq!( + request_headers.get("X-OAI-IS"), + Some(&HeaderValue::from_static("state-before")) + ); + + let mut response_headers = HeaderMap::new(); + response_headers.insert("X-OAI-IS-Update", HeaderValue::from_static("state-after")); + observe_response_headers( + &auth_provider, + REQUEST_URL, + &request_headers, + &response_headers, + ); + + assert_eq!( + state + .get_for_surface(HttpStateSurface::CodexCli) + .expect("read state"), + Some("state-after".to_string()) + ); +} diff --git a/codex-rs/core-plugins/src/remote/share.rs b/codex-rs/core-plugins/src/remote/share.rs index c400d27c47..c462b7d7fc 100644 --- a/codex-rs/core-plugins/src/remote/share.rs +++ b/codex-rs/core-plugins/src/remote/share.rs @@ -4,7 +4,6 @@ use crate::plugin_bundle_archive::pack_plugin_bundle_tar_gz; use codex_login::CodexAuth; use codex_login::default_client::build_reqwest_client; use codex_utils_absolute_path::AbsolutePathBuf; -use reqwest::RequestBuilder; use reqwest::StatusCode; use serde::Deserialize; use serde::Serialize; @@ -280,7 +279,7 @@ pub async fn delete_remote_plugin_share( let base_url = config.chatgpt_base_url.trim_end_matches('/'); let url = format!("{base_url}/public/plugins/workspace/{remote_plugin_id}"); let client = build_reqwest_client(); - let request = authenticated_request(client.delete(&url), auth)?; + let request = authenticated_request(config, client.delete(&url), auth, &url); send_and_expect_status(request, &url, &[StatusCode::NO_CONTENT]).await?; if let Err(err) = local_paths::remove_plugin_share_local_path(codex_home, remote_plugin_id) { warn!( @@ -313,7 +312,7 @@ pub async fn update_remote_plugin_share_targets( let base_url = config.chatgpt_base_url.trim_end_matches('/'); let url = format!("{base_url}/ps/plugins/{remote_plugin_id}/shares"); let client = build_reqwest_client(); - let request = authenticated_request(client.put(&url), auth)?.json( + let request = authenticated_request(config, client.put(&url), auth, &url).json( &RemotePluginShareUpdateTargetsRequest { discoverability, targets, @@ -379,7 +378,7 @@ async fn get_created_workspace_plugins_page( let base_url = config.chatgpt_base_url.trim_end_matches('/'); let url = format!("{base_url}/ps/plugins/workspace/created"); let client = build_reqwest_client(); - let mut request = authenticated_request(client.get(&url), auth)?; + let mut request = authenticated_request(config, client.get(&url), auth, &url); request = request.query(&[("limit", REMOTE_PLUGIN_LIST_PAGE_LIMIT)]); if let Some(page_token) = page_token { request = request.query(&[("pageToken", page_token)]); @@ -397,7 +396,7 @@ async fn create_workspace_plugin_upload( let base_url = config.chatgpt_base_url.trim_end_matches('/'); let url = format!("{base_url}/public/plugins/workspace/upload-url"); let client = build_reqwest_client(); - let request = authenticated_request(client.post(&url), auth)?.json( + let request = authenticated_request(config, client.post(&url), auth, &url).json( &RemoteWorkspacePluginUploadUrlRequest { filename, mime_type: "application/gzip", @@ -451,7 +450,7 @@ async fn finalize_workspace_plugin_upload( format!("{base_url}/public/plugins/workspace") }; let client = build_reqwest_client(); - let request = authenticated_request(client.post(&url), auth)?.json(&body); + let request = authenticated_request(config, client.post(&url), auth, &url).json(&body); send_and_decode(request, &url).await } @@ -489,17 +488,11 @@ fn archive_plugin_for_upload_with_limit( } async fn send_and_expect_status( - request: RequestBuilder, + request: AuthenticatedRequest, url_for_error: &str, expected_statuses: &[StatusCode], ) -> Result<(), RemotePluginCatalogError> { - let response = request - .send() - .await - .map_err(|source| RemotePluginCatalogError::Request { - url: url_for_error.to_string(), - source, - })?; + let response = request.send(url_for_error).await?; let status = response.status(); let body = response.text().await.unwrap_or_default(); if !expected_statuses.contains(&status) { diff --git a/codex-rs/core-plugins/src/remote/share/tests.rs b/codex-rs/core-plugins/src/remote/share/tests.rs index e33d3c5d51..1f6ee85f26 100644 --- a/codex-rs/core-plugins/src/remote/share/tests.rs +++ b/codex-rs/core-plugins/src/remote/share/tests.rs @@ -23,9 +23,7 @@ use wiremock::matchers::query_param; use wiremock::matchers::query_param_is_missing; fn test_config(server: &MockServer) -> RemotePluginServiceConfig { - RemotePluginServiceConfig { - chatgpt_base_url: format!("{}/backend-api", server.uri()), - } + RemotePluginServiceConfig::new(format!("{}/backend-api", server.uri())) } fn test_auth() -> CodexAuth { diff --git a/codex-rs/core-plugins/src/remote_legacy.rs b/codex-rs/core-plugins/src/remote_legacy.rs index dcf9f79eb8..177d91e0b1 100644 --- a/codex-rs/core-plugins/src/remote_legacy.rs +++ b/codex-rs/core-plugins/src/remote_legacy.rs @@ -1,4 +1,5 @@ use crate::remote::RemotePluginServiceConfig; +use crate::remote::http_state; use codex_login::CodexAuth; use codex_login::default_client::build_reqwest_client; use codex_protocol::protocol::Product; @@ -130,10 +131,11 @@ pub async fn fetch_remote_plugin_status( let base_url = config.chatgpt_base_url.trim_end_matches('/'); let url = format!("{base_url}/plugins/list"); let client = build_reqwest_client(); + let (auth_provider, request_headers) = http_state::request_headers(config, auth, &url); let request = client .get(&url) .timeout(REMOTE_PLUGIN_FETCH_TIMEOUT) - .headers(codex_model_provider::auth_provider_from_auth(auth).to_auth_headers()); + .headers(request_headers.clone()); let response = request .send() @@ -142,6 +144,12 @@ pub async fn fetch_remote_plugin_status( url: url.clone(), source, })?; + http_state::observe_response_headers( + &auth_provider, + &url, + &request_headers, + response.headers(), + ); let status = response.status(); let body = response.text().await.unwrap_or_default(); if !status.is_success() { @@ -170,9 +178,11 @@ pub async fn fetch_remote_featured_plugin_ids( )]) .timeout(REMOTE_FEATURED_PLUGIN_FETCH_TIMEOUT); - if let Some(auth) = auth.filter(|auth| auth.uses_codex_backend()) { - request = - request.headers(codex_model_provider::auth_provider_from_auth(auth).to_auth_headers()); + let auth_context = auth + .filter(|auth| auth.uses_codex_backend()) + .map(|auth| http_state::request_headers(config, auth, &url)); + if let Some((_, request_headers)) = &auth_context { + request = request.headers(request_headers.clone()); } let response = request @@ -182,6 +192,14 @@ pub async fn fetch_remote_featured_plugin_ids( url: url.clone(), source, })?; + if let Some((auth_provider, request_headers)) = &auth_context { + http_state::observe_response_headers( + auth_provider, + &url, + request_headers, + response.headers(), + ); + } let status = response.status(); let body = response.text().await.unwrap_or_default(); if !status.is_success() { @@ -237,10 +255,11 @@ async fn post_remote_plugin_mutation( let auth = ensure_codex_backend_auth(auth)?; let url = remote_plugin_mutation_url(config, plugin_id, action)?; let client = build_reqwest_client(); + let (auth_provider, request_headers) = http_state::request_headers(config, auth, &url); let request = client .post(url.clone()) .timeout(REMOTE_PLUGIN_MUTATION_TIMEOUT) - .headers(codex_model_provider::auth_provider_from_auth(auth).to_auth_headers()); + .headers(request_headers.clone()); let response = request .send() @@ -249,6 +268,12 @@ async fn post_remote_plugin_mutation( url: url.clone(), source, })?; + http_state::observe_response_headers( + &auth_provider, + &url, + &request_headers, + response.headers(), + ); let status = response.status(); let body = response.text().await.unwrap_or_default(); if !status.is_success() { diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 3971434c08..183a5b5001 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -1597,6 +1597,16 @@ async fn spawn_thread_subagent_inherits_parent_http_state_surface() { ) .await .expect("parent client info should update"); + assert_eq!( + parent_thread + .codex + .session + .services + .thread_extension_data + .get::() + .as_deref(), + Some(&HttpStateSurface::CodexDesktop) + ); let child_thread_id = harness .control @@ -1628,6 +1638,16 @@ async fn spawn_thread_subagent_inherits_parent_http_state_surface() { .http_state_surface(), Some(HttpStateSurface::CodexDesktop) ); + assert_eq!( + child_thread + .codex + .session + .services + .thread_extension_data + .get::() + .as_deref(), + Some(&HttpStateSurface::CodexDesktop) + ); } #[tokio::test] diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 253c7ad6e0..f0b96faafe 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -415,6 +415,16 @@ impl ModelClient { self.state.http_state.clone() } + pub(crate) fn http_state_context_for_surface( + &self, + surface: HttpStateSurface, + ) -> Option { + self.state + .http_state + .as_ref() + .map(|state| state.for_surface(surface)) + } + pub(crate) fn set_http_state_surface(&self, surface: HttpStateSurface) { if self .state @@ -869,7 +879,15 @@ impl ModelClient { /// This centralizes setup used by both prewarm and normal request paths so they stay in /// lockstep when auth/provider resolution changes. async fn current_client_setup(&self) -> Result { - let auth = self.state.provider.auth().await; + let auth = match ( + self.state.provider.auth_manager(), + self.http_state_surface(), + ) { + (Some(auth_manager), Some(http_state_surface)) => { + auth_manager.auth_for_surface(http_state_surface).await + } + _ => self.state.provider.auth().await, + }; let api_provider = self.state.provider.api_provider().await?; let api_auth = with_native_integrity_state( self.state.provider.api_auth().await?, @@ -1316,9 +1334,14 @@ impl ModelClientSession { inference_trace: &InferenceTraceContext, ) -> Result { let auth_manager = self.client.state.provider.auth_manager(); - let mut auth_recovery = auth_manager - .as_ref() - .map(AuthManager::unauthorized_recovery); + let http_state_surface = self.client.http_state_surface(); + let mut auth_recovery = auth_manager.as_ref().map(|auth_manager| { + if let Some(http_state_surface) = http_state_surface { + auth_manager.unauthorized_recovery_for_surface(http_state_surface) + } else { + auth_manager.unauthorized_recovery() + } + }); let mut pending_retry = PendingUnauthorizedRetry::default(); loop { let client_setup = self.client.current_client_setup().await?; @@ -1431,9 +1454,14 @@ impl ModelClientSession { inference_trace: &InferenceTraceContext, ) -> Result { let auth_manager = self.client.state.provider.auth_manager(); - let mut auth_recovery = auth_manager - .as_ref() - .map(AuthManager::unauthorized_recovery); + let http_state_surface = self.client.http_state_surface(); + let mut auth_recovery = auth_manager.as_ref().map(|auth_manager| { + if let Some(http_state_surface) = http_state_surface { + auth_manager.unauthorized_recovery_for_surface(http_state_surface) + } else { + auth_manager.unauthorized_recovery() + } + }); let mut pending_retry = PendingUnauthorizedRetry::default(); loop { let client_setup = self.client.current_client_setup().await?; diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 7a8344b3d7..4b6660aa47 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -31,6 +31,7 @@ use codex_config::types::AppsConfigToml; use codex_config::types::ToolSuggestDiscoverableType; use codex_core_plugins::PluginsManager; use codex_features::Feature; +use codex_http_state::HttpStateContext; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::default_client::originator; @@ -87,10 +88,17 @@ pub struct AccessibleConnectorsStatus { pub async fn list_accessible_connectors_from_mcp_tools( config: &Config, +) -> anyhow::Result> { + list_accessible_connectors_from_mcp_tools_with_http_state(config, /*http_state*/ None).await +} + +pub(crate) async fn list_accessible_connectors_from_mcp_tools_with_http_state( + config: &Config, + http_state: Option, ) -> anyhow::Result> { Ok( - list_accessible_connectors_from_mcp_tools_with_options_and_status( - config, /*force_refetch*/ false, + list_accessible_connectors_from_mcp_tools_with_options_and_status_and_http_state( + config, /*force_refetch*/ false, http_state, ) .await? .connectors, @@ -147,10 +155,23 @@ pub(crate) async fn list_tool_suggest_discoverable_tools_with_auth( pub async fn list_cached_accessible_connectors_from_mcp_tools( config: &Config, +) -> Option> { + list_cached_accessible_connectors_from_mcp_tools_with_http_state( + config, /*http_state*/ None, + ) + .await +} + +pub async fn list_cached_accessible_connectors_from_mcp_tools_with_http_state( + config: &Config, + http_state: Option<&HttpStateContext>, ) -> Option> { let auth_manager = AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await; - let auth = auth_manager.auth().await; + let auth = match http_state { + Some(http_state) => auth_manager.auth_for_surface(http_state.surface()).await, + None => auth_manager.auth().await, + }; if !config .features .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend)) @@ -197,6 +218,19 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_options( pub async fn list_accessible_connectors_from_mcp_tools_with_options_and_status( config: &Config, force_refetch: bool, +) -> anyhow::Result { + list_accessible_connectors_from_mcp_tools_with_options_and_status_and_http_state( + config, + force_refetch, + /*http_state*/ None, + ) + .await +} + +async fn list_accessible_connectors_from_mcp_tools_with_options_and_status_and_http_state( + config: &Config, + force_refetch: bool, + http_state: Option, ) -> anyhow::Result { // TODO: Wire callers that already own an EnvironmentManager into // list_accessible_connectors_from_mcp_tools_with_environment_manager instead @@ -208,10 +242,11 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_options_and_status( let environment_manager = EnvironmentManager::from_codex_home(config.codex_home.clone(), Some(local_runtime_paths)) .await?; - list_accessible_connectors_from_mcp_tools_with_environment_manager( + list_accessible_connectors_from_mcp_tools_with_environment_manager_and_http_state( config, force_refetch, Arc::new(environment_manager), + http_state, ) .await } @@ -220,10 +255,28 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager( config: &Config, force_refetch: bool, environment_manager: Arc, +) -> anyhow::Result { + list_accessible_connectors_from_mcp_tools_with_environment_manager_and_http_state( + config, + force_refetch, + environment_manager, + /*http_state*/ None, + ) + .await +} + +pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager_and_http_state( + config: &Config, + force_refetch: bool, + environment_manager: Arc, + http_state: Option, ) -> anyhow::Result { let auth_manager = AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ false).await; - let auth = auth_manager.auth().await; + let auth = match http_state.as_ref() { + Some(http_state) => auth_manager.auth_for_surface(http_state.surface()).await, + None => auth_manager.auth().await, + }; if !config .features .apps_enabled_for_auth(auth.as_ref().is_some_and(CodexAuth::uses_codex_backend)) @@ -282,6 +335,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_environment_manager( // one exists, but do not reintroduce the old hidden-local fallback. McpRuntimeContext::new(environment_manager, config.cwd.to_path_buf()), config.codex_home.to_path_buf(), + http_state, codex_apps_tools_cache_key(auth.as_ref()), host_owned_codex_apps_enabled, mcp_config.prefix_mcp_tool_names, diff --git a/codex-rs/core/src/mcp_openai_file.rs b/codex-rs/core/src/mcp_openai_file.rs index 821cab7de9..f54e1e7210 100644 --- a/codex-rs/core/src/mcp_openai_file.rs +++ b/codex-rs/core/src/mcp_openai_file.rs @@ -33,8 +33,15 @@ pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files( let Some(arguments) = arguments_value.as_object() else { return Ok(Some(arguments_value)); }; - let auth = sess.services.auth_manager.auth().await; - let http_state = sess.services.model_client.http_state_context(); + let auth = sess + .services + .auth_manager + .auth_for_surface(turn_context.http_state_surface) + .await; + let http_state = sess + .services + .model_client + .http_state_context_for_surface(turn_context.http_state_surface); let mut rewritten_arguments = arguments.clone(); for field_name in openai_file_input_params { diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 9c00ba2be6..0b0ba58268 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -35,6 +35,7 @@ use codex_config::types::AppToolApproval; use codex_config::types::ApprovalsReviewer; use codex_features::Feature; use codex_hooks::PermissionRequestDecision; +use codex_http_state::HttpStateContext; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::MCP_TOOL_CODEX_APPS_META_KEY; use codex_mcp::McpPermissionPromptAutoApproveContext; @@ -682,7 +683,11 @@ async fn refresh_codex_apps_after_connector_auth(sess: &Session, turn_context: & match mcp_tools_result { Ok(mcp_tools) => { - let auth = sess.services.auth_manager.auth().await; + let auth = sess + .services + .auth_manager + .auth_for_surface(turn_context.http_state_surface) + .await; connectors::refresh_accessible_connectors_cache_from_mcp_tools( &turn_context.config, auth.as_ref(), @@ -1434,18 +1439,25 @@ pub(crate) async fn lookup_mcp_tool_metadata( .into_iter() .find(|tool_info| tool_info.server_name == server && tool_info.tool.name == tool_name)?; let connector_description = if server == CODEX_APPS_MCP_SERVER_NAME { - let connectors = match connectors::list_cached_accessible_connectors_from_mcp_tools( - turn_context.config.as_ref(), - ) - .await - { - Some(connectors) => Some(connectors), - None => { - connectors::list_accessible_connectors_from_mcp_tools(turn_context.config.as_ref()) - .await - .ok() - } - }; + let http_state = HttpStateContext::new( + turn_context.config.codex_home.to_path_buf(), + turn_context.http_state_surface, + ); + let connectors = + match connectors::list_cached_accessible_connectors_from_mcp_tools_with_http_state( + turn_context.config.as_ref(), + Some(&http_state), + ) + .await + { + Some(connectors) => Some(connectors), + None => connectors::list_accessible_connectors_from_mcp_tools_with_http_state( + turn_context.config.as_ref(), + Some(http_state), + ) + .await + .ok(), + }; connectors.and_then(|connectors| { let connector_id = tool_info.connector_id.as_deref()?; connectors diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index c9259b8f2d..70565c1c63 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1269,6 +1269,7 @@ async fn install_host_owned_codex_apps_manager(session: &Session, turn_context: turn_context.cwd.to_path_buf() }), turn_context.config.codex_home.to_path_buf(), + /*http_state*/ None, codex_mcp::codex_apps_tools_cache_key(auth.as_ref()), /*host_owned_codex_apps_enabled*/ true, turn_context.config.prefix_mcp_tool_names(), diff --git a/codex-rs/core/src/plugins/discoverable_tests.rs b/codex-rs/core/src/plugins/discoverable_tests.rs index 992bf93f45..1a4601449a 100644 --- a/codex-rs/core/src/plugins/discoverable_tests.rs +++ b/codex-rs/core/src/plugins/discoverable_tests.rs @@ -320,11 +320,11 @@ remote_plugin = true let mut config = load_plugins_config(codex_home.path()).await; config.chatgpt_base_url = format!("{}/backend-api", server.uri()); let plugins_manager = PluginsManager::new(config.codex_home.to_path_buf()); + let remote_plugin_service_config = + RemotePluginServiceConfig::new(config.chatgpt_base_url.clone()); fetch_and_cache_global_remote_plugin_catalog( codex_home.path(), - &RemotePluginServiceConfig { - chatgpt_base_url: config.chatgpt_base_url.clone(), - }, + &remote_plugin_service_config, Some(&auth), ) .await diff --git a/codex-rs/core/src/realtime_conversation.rs b/codex-rs/core/src/realtime_conversation.rs index 1040d9e2d7..0887b4b1fc 100644 --- a/codex-rs/core/src/realtime_conversation.rs +++ b/codex-rs/core/src/realtime_conversation.rs @@ -27,6 +27,8 @@ use codex_config::config_toml::RealtimeWsVersion; use codex_login::CodexAuth; use codex_login::default_client::default_headers; use codex_login::read_openai_api_key_from_env; +use codex_model_provider::unauthenticated_auth_provider; +use codex_model_provider::with_native_integrity_state; use codex_model_provider_info::ModelProviderInfo; use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CodexResult; @@ -227,6 +229,7 @@ struct ConversationState { struct RealtimeStart { api_provider: ApiProvider, + websocket_auth: SharedAuthProvider, extra_headers: Option, session_config: RealtimeSessionConfig, model_client: ModelClient, @@ -279,6 +282,7 @@ impl RealtimeConversationManager { async fn start_inner(&self, start: RealtimeStart) -> CodexResult { let RealtimeStart { api_provider, + websocket_auth, extra_headers, session_config, model_client, @@ -307,7 +311,7 @@ impl RealtimeConversationManager { audio_rx, }; - let client = RealtimeWebsocketClient::new(api_provider); + let client = RealtimeWebsocketClient::new(api_provider).with_auth(websocket_auth); let (task, sdp) = if let Some(sdp) = sdp { let call = model_client .create_realtime_call_with_headers( @@ -601,6 +605,7 @@ pub(crate) async fn handle_start( struct PreparedRealtimeConversationStart { api_provider: ApiProvider, + websocket_auth: SharedAuthProvider, extra_headers: Option, requested_realtime_session_id: Option, version: RealtimeWsVersion, @@ -618,7 +623,15 @@ async fn prepare_realtime_start( .model_client .auth_manager() .unwrap_or_else(|| Arc::clone(&sess.services.auth_manager)); - let auth = auth_manager.auth().await; + let auth = match sess.services.model_client.http_state_surface() { + Some(http_state_surface) => auth_manager.auth_for_surface(http_state_surface).await, + None => auth_manager.auth().await, + }; + let websocket_auth = with_native_integrity_state( + unauthenticated_auth_provider(), + auth.as_ref(), + sess.services.model_client.http_state_context(), + ); let config = sess.get_config().await; let transport = params .transport @@ -656,6 +669,7 @@ async fn prepare_realtime_start( }; Ok(PreparedRealtimeConversationStart { api_provider, + websocket_auth, extra_headers, requested_realtime_session_id, version, @@ -777,6 +791,7 @@ async fn handle_start_inner( ) -> CodexResult<()> { let PreparedRealtimeConversationStart { api_provider, + websocket_auth, extra_headers, requested_realtime_session_id, version, @@ -790,6 +805,7 @@ async fn handle_start_inner( }; let start = RealtimeStart { api_provider, + websocket_auth, extra_headers, session_config, model_client: sess.services.model_client.clone(), diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index cc018343eb..467955d8d8 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -309,7 +309,11 @@ impl Session { store_mode: OAuthCredentialsStoreMode, elicitation_reviewer: Option, ) { - let auth = self.services.auth_manager.auth().await; + let auth = self + .services + .auth_manager + .auth_for_surface(turn_context.http_state_surface) + .await; let config = self.get_config().await; let mcp_config = config .to_mcp_config(self.services.plugins_manager.as_ref()) @@ -351,6 +355,9 @@ impl Session { turn_context.permission_profile(), mcp_runtime_context, config.codex_home.to_path_buf(), + self.services + .model_client + .http_state_context_for_surface(turn_context.http_state_surface), codex_apps_tools_cache_key(auth.as_ref()), host_owned_codex_apps_enabled, mcp_config.prefix_mcp_tool_names, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 7f683fb7e4..aacaa5d1b0 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -796,6 +796,12 @@ impl Codex { .services .model_client .set_app_server_client_name(app_server_client_name.as_deref()); + if let Some(http_state_surface) = self.session.services.model_client.http_state_surface() { + self.session + .services + .thread_extension_data + .insert(http_state_surface); + } self.session .update_settings(SessionSettingsUpdate { app_server_client_name, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 77a268b6ee..1c9fafb792 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -973,6 +973,7 @@ impl Session { codex_extension_api::ExtensionData::new(session_id.to_string()); let thread_extension_data = codex_extension_api::ExtensionData::new(thread_id.to_string()); + thread_extension_data.insert(http_state_surface); for contributor in extensions.thread_lifecycle_contributors() { contributor.on_thread_start(codex_extension_api::ThreadStartInput { config: config.as_ref(), @@ -1184,6 +1185,7 @@ impl Session { session_configuration.permission_profile(), mcp_runtime_context, config.codex_home.to_path_buf(), + sess.services.model_client.http_state_context(), codex_apps_tools_cache_key(auth), host_owned_codex_apps_enabled, config.prefix_mcp_tool_names(), diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index f2f0df071e..21a623dc28 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1126,7 +1126,11 @@ pub(crate) async fn built_tools( } else { None }; - let auth = sess.services.auth_manager.auth().await; + let auth = sess + .services + .auth_manager + .auth_for_surface(turn_context.http_state_surface) + .await; let loaded_plugin_app_connector_ids = loaded_plugins .effective_apps() .into_iter() diff --git a/codex-rs/core/src/tools/handlers/request_plugin_install.rs b/codex-rs/core/src/tools/handlers/request_plugin_install.rs index 51032b27c5..f0dce2dd7b 100644 --- a/codex-rs/core/src/tools/handlers/request_plugin_install.rs +++ b/codex-rs/core/src/tools/handlers/request_plugin_install.rs @@ -138,7 +138,11 @@ impl ToolExecutor for RequestPluginInstallHandler { .as_ref() .is_some_and(|response| response.action == ElicitationAction::Accept); - let auth = session.services.auth_manager.auth().await; + let auth = session + .services + .auth_manager + .auth_for_surface(turn.http_state_surface) + .await; let completed = if user_confirmed { verify_request_plugin_install_completed(&session, &turn, &tool, auth.as_ref()).await } else { diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 4fe0609a71..7735212ae7 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -24,6 +24,7 @@ use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::ExecutorFileSystem; use codex_exec_server::RemoveOptions; use codex_extension_api::empty_extension_registry; +use codex_login::AuthManager; use codex_login::CodexAuth; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::built_in_model_providers; @@ -210,6 +211,7 @@ pub fn turn_permission_fields( pub struct TestCodexBuilder { config_mutators: Vec>, auth: CodexAuth, + auth_manager: Option>, pre_build_hooks: Vec>, workspace_setups: Vec>, home: Option>, @@ -232,6 +234,11 @@ impl TestCodexBuilder { self } + pub fn with_auth_manager(mut self, auth_manager: Arc) -> Self { + self.auth_manager = Some(auth_manager); + self + } + pub fn with_model(self, model: &str) -> Self { let new_model = model.to_string(); self.with_config(move |config| { @@ -464,13 +471,16 @@ impl TestCodexBuilder { test_env: TestEnv, environment_manager: Arc, ) -> anyhow::Result { - let auth = self.auth.clone(); + let auth_manager = self + .auth_manager + .clone() + .unwrap_or_else(|| codex_core::test_support::auth_manager_from_auth(self.auth.clone())); let state_db = codex_core::init_state_db(&config).await; let thread_store = thread_store_from_config(&config, state_db.clone()); let installation_id = resolve_installation_id(&config.codex_home).await?; let thread_manager = ThreadManager::new( &config, - codex_core::test_support::auth_manager_from_auth(auth.clone()), + Arc::clone(&auth_manager), SessionSource::Exec, Arc::clone(&environment_manager), empty_extension_registry(), @@ -483,42 +493,39 @@ impl TestCodexBuilder { let thread_manager = Arc::new(thread_manager); let user_shell_override = self.user_shell_override.clone(); - let new_conversation = match (resume_from, user_shell_override) { - (Some(path), Some(user_shell_override)) => { - let auth_manager = codex_core::test_support::auth_manager_from_auth(auth); - Box::pin( + let new_conversation = + match (resume_from, user_shell_override) { + (Some(path), Some(user_shell_override)) => Box::pin( codex_core::test_support::resume_thread_from_rollout_with_user_shell_override( thread_manager.as_ref(), config.clone(), path, - auth_manager, + Arc::clone(&auth_manager), user_shell_override, ), ) - .await? - } - (Some(path), None) => { - let auth_manager = codex_core::test_support::auth_manager_from_auth(auth); - Box::pin(thread_manager.resume_thread_from_rollout( - config.clone(), - path, - auth_manager, - /*parent_trace*/ None, - )) - .await? - } - (None, Some(user_shell_override)) => { - Box::pin( - codex_core::test_support::start_thread_with_user_shell_override( - thread_manager.as_ref(), + .await?, + (Some(path), None) => { + Box::pin(thread_manager.resume_thread_from_rollout( config.clone(), - user_shell_override, - ), - ) - .await? - } - (None, None) => Box::pin(thread_manager.start_thread(config.clone())).await?, - }; + path, + Arc::clone(&auth_manager), + /*parent_trace*/ None, + )) + .await? + } + (None, Some(user_shell_override)) => { + Box::pin( + codex_core::test_support::start_thread_with_user_shell_override( + thread_manager.as_ref(), + config.clone(), + user_shell_override, + ), + ) + .await? + } + (None, None) => Box::pin(thread_manager.start_thread(config.clone())).await?, + }; Ok(TestCodex { home, @@ -1031,6 +1038,7 @@ pub fn test_codex() -> TestCodexBuilder { TestCodexBuilder { config_mutators: vec![], auth: CodexAuth::from_api_key("dummy"), + auth_manager: None, pre_build_hooks: vec![], workspace_setups: vec![], home: None, diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index c954a954fb..57066d689e 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -9,8 +9,11 @@ use codex_core::resolve_installation_id; use codex_core::thread_store_from_config; use codex_extension_api::empty_extension_registry; use codex_features::Feature; +use codex_http_state::HttpStateStore; +use codex_http_state::HttpStateSurface; use codex_login::AuthManager; use codex_login::CodexAuth; +use codex_login::REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR; use codex_login::default_client::originator; use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::WireApi; @@ -58,6 +61,7 @@ use core_test_support::responses::ev_completed_with_tokens; use core_test_support::responses::ev_message_item_added; use core_test_support::responses::ev_output_text_delta; use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_response_sequence; use core_test_support::responses::mount_sse_once; use core_test_support::responses::mount_sse_once_match; use core_test_support::responses::mount_sse_sequence; @@ -71,6 +75,7 @@ use dunce::canonicalize as normalize_path; use futures::StreamExt; use pretty_assertions::assert_eq; use serde_json::json; +use std::ffi::OsString; use std::io::Write; use std::num::NonZeroU64; use std::sync::Arc; @@ -943,6 +948,130 @@ async fn send_provider_auth_request(server: &MockServer, auth: ModelProviderAuth } } +struct EnvGuard { + key: &'static str, + original: Option, +} + +impl EnvGuard { + fn set(key: &'static str, value: String) -> Self { + let original = std::env::var_os(key); + // SAFETY: the test is serialized with the other auth-refresh tests. + unsafe { + std::env::set_var(key, &value); + } + Self { key, original } + } +} + +impl Drop for EnvGuard { + fn drop(&mut self) { + // SAFETY: the guard restores the original environment value before other tests run. + unsafe { + match &self.original { + Some(value) => std::env::set_var(self.key, value), + None => std::env::remove_var(self.key), + } + } + } +} + +#[serial_test::serial(auth_refresh)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn response_401_refresh_preserves_native_http_state_surface() { + skip_if_no_network!(); + + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/oauth/token")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "access_token": "new-access-token", + "refresh_token": "new-refresh-token", + "state": "new-state", + }))) + .expect(1) + .mount(&server) + .await; + let responses_mock = mount_response_sequence( + &server, + vec![ + ResponseTemplate::new(401).set_body_string("unauthorized"), + ResponseTemplate::new(401).set_body_string("unauthorized"), + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_raw( + sse(vec![ev_response_created("resp1"), ev_completed("resp1")]), + "text/event-stream", + ), + ], + ) + .await; + + let home = Arc::new(TempDir::new().unwrap()); + write_auth_json( + home.as_ref(), + /*openai_api_key*/ None, + "pro", + "old-access-token", + Some("acc-123"), + ); + let store = HttpStateStore::new(home.path().to_path_buf()); + store + .set(HttpStateSurface::CodexExec, "old-state".to_string()) + .unwrap(); + let _refresh_endpoint = EnvGuard::set( + REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR, + format!("{}/oauth/token", server.uri()), + ); + let auth_manager = AuthManager::shared( + home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await; + let mut builder = test_codex() + .with_home(home) + .with_auth_manager(auth_manager) + .with_config(|config| { + config.model_catalog = + Some(bundled_models_response().expect("bundled models should parse")); + }); + let test = builder.build(&server).await.unwrap(); + + test.submit_turn("hello").await.unwrap(); + + assert_eq!( + store.get(HttpStateSurface::CodexExec).unwrap(), + Some("new-state".to_string()) + ); + let requests = responses_mock.requests(); + assert_eq!(requests.len(), 3); + assert_eq!( + requests[0].header("authorization").as_deref(), + Some("Bearer old-access-token") + ); + assert_eq!( + requests[1].header("authorization").as_deref(), + Some("Bearer old-access-token") + ); + assert_eq!( + requests[2].header("authorization").as_deref(), + Some("Bearer new-access-token") + ); + let refresh_request = server + .received_requests() + .await + .unwrap_or_default() + .into_iter() + .find(|request| request.url.path() == "/oauth/token") + .expect("refresh request"); + let refresh_body: serde_json::Value = + serde_json::from_slice(&refresh_request.body).expect("refresh JSON body"); + assert_eq!(refresh_body["surface"], "codex_exec"); + server.verify().await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn includes_base_instructions_override_in_request() { skip_if_no_network!(); diff --git a/codex-rs/core/tests/suite/openai_file_mcp.rs b/codex-rs/core/tests/suite/openai_file_mcp.rs index a25cc36a86..f8939fd908 100644 --- a/codex-rs/core/tests/suite/openai_file_mcp.rs +++ b/codex-rs/core/tests/suite/openai_file_mcp.rs @@ -5,6 +5,8 @@ use std::path::Path; use anyhow::Context; use anyhow::Result; +use codex_http_state::HttpStateStore; +use codex_http_state::HttpStateSurface; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; use core_test_support::apps_test_server::AppsTestServer; @@ -33,6 +35,11 @@ use wiremock::matchers::header; use wiremock::matchers::method; use wiremock::matchers::path; +const INITIAL_INTEGRITY_STATE: &str = "state-n"; +const INTEGRITY_STATE_HEADER_NAME: &str = "x-oai-is"; +const INTEGRITY_STATE_UPDATE_HEADER_NAME: &str = "x-oai-is-update"; +const UNTRUSTED_INTEGRITY_STATE_UPDATE: &str = "state-n-plus-one"; + fn write_post_tool_use_hook(home: &Path) -> Result<()> { let script_path = home.join("post_tool_use_hook.py"); let log_path = home.join("post_tool_use_hook_log.jsonl"); @@ -95,10 +102,17 @@ async fn codex_apps_file_params_upload_local_paths_before_mcp_tool_call() -> Res "file_size": 11, "use_case": "codex", }))) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "file_id": "file_123", - "upload_url": format!("{}/upload/file_123", server.uri()), - }))) + .respond_with( + ResponseTemplate::new(200) + .insert_header( + INTEGRITY_STATE_UPDATE_HEADER_NAME, + UNTRUSTED_INTEGRITY_STATE_UPDATE, + ) + .set_body_json(json!({ + "file_id": "file_123", + "upload_url": format!("{}/upload/file_123", server.uri()), + })), + ) .expect(1) .mount(&server) .await; @@ -150,6 +164,12 @@ async fn codex_apps_file_params_upload_local_paths_before_mcp_tool_call() -> Res if let Err(error) = write_post_tool_use_hook(home) { panic!("failed to write apps file post tool use hook fixture: {error}"); } + HttpStateStore::new(home.to_path_buf()) + .set( + HttpStateSurface::CodexExec, + INITIAL_INTEGRITY_STATE.to_string(), + ) + .expect("initial HTTP state should store"); }) .with_config(move |config| { trust_discovered_hooks(config); @@ -218,6 +238,23 @@ async fn codex_apps_file_params_upload_local_paths_before_mcp_tool_call() -> Res "file_size_bytes": 11, }) ); + assert_eq!( + HttpStateStore::new(test.codex_home_path().to_path_buf()) + .get(HttpStateSurface::CodexExec)?, + Some(INITIAL_INTEGRITY_STATE.to_string()) + ); + let file_create_request = server + .received_requests() + .await + .context("capture file upload requests")? + .into_iter() + .find(|request| request.url.path() == "/files") + .context("find file create request")?; + assert!( + !file_create_request + .headers + .contains_key(INTEGRITY_STATE_HEADER_NAME) + ); server.verify().await; Ok(()) diff --git a/codex-rs/ext/image-generation/Cargo.toml b/codex-rs/ext/image-generation/Cargo.toml index 995f1786bc..68a1113336 100644 --- a/codex-rs/ext/image-generation/Cargo.toml +++ b/codex-rs/ext/image-generation/Cargo.toml @@ -18,6 +18,7 @@ codex-api = { workspace = true } codex-core = { workspace = true } codex-extension-api = { workspace = true } codex-features = { workspace = true } +codex-http-state = { workspace = true } codex-login = { workspace = true } codex-model-provider = { workspace = true } codex-model-provider-info = { workspace = true } diff --git a/codex-rs/ext/image-generation/src/backend.rs b/codex-rs/ext/image-generation/src/backend.rs index 9b837772c7..03b66ad087 100644 --- a/codex-rs/ext/image-generation/src/backend.rs +++ b/codex-rs/ext/image-generation/src/backend.rs @@ -3,33 +3,46 @@ use codex_api::ImageGenerationRequest; use codex_api::ImageResponse; use codex_api::ImagesClient; use codex_api::ReqwestTransport; +use codex_http_state::HttpStateContext; use codex_login::default_client::build_reqwest_client; use codex_model_provider::SharedModelProvider; +use codex_model_provider::with_native_integrity_state; use http::HeaderMap; #[derive(Clone)] pub(crate) struct CodexImagesBackend { provider: SharedModelProvider, + http_state: Option, } impl CodexImagesBackend { /// Creates a backend that sends image requests through the active model provider. - pub(crate) fn new(provider: SharedModelProvider) -> Self { - Self { provider } + pub(crate) fn new(provider: SharedModelProvider, http_state: Option) -> Self { + Self { + provider, + http_state, + } } /// Resolves the provider and auth required for the current image API request. async fn client(&self) -> Result, String> { + let auth = match self.http_state.as_ref() { + Some(http_state) => self.provider.auth_for_surface(http_state.surface()).await, + None => self.provider.auth().await, + }; let provider = self .provider .api_provider() .await .map_err(|err| err.to_string())?; - let auth = self - .provider - .api_auth() - .await - .map_err(|err| err.to_string())?; + let auth = with_native_integrity_state( + self.provider + .api_auth() + .await + .map_err(|err| err.to_string())?, + auth.as_ref(), + self.http_state.clone(), + ); Ok(ImagesClient::new( ReqwestTransport::new(build_reqwest_client()), provider, diff --git a/codex-rs/ext/image-generation/src/extension.rs b/codex-rs/ext/image-generation/src/extension.rs index f7a9f2406f..a1d2f543ae 100644 --- a/codex-rs/ext/image-generation/src/extension.rs +++ b/codex-rs/ext/image-generation/src/extension.rs @@ -10,6 +10,8 @@ use codex_extension_api::ToolCall; use codex_extension_api::ToolContributor; use codex_extension_api::ToolExecutor; use codex_features::Feature; +use codex_http_state::HttpStateContext; +use codex_http_state::HttpStateSurface; use codex_login::AuthManager; use codex_model_provider::create_model_provider; use codex_model_provider_info::ModelProviderInfo; @@ -78,6 +80,9 @@ impl ToolContributor for ImageGenerationExtension { vec![Arc::new(ImageGenerationTool::new(CodexImagesBackend::new( create_model_provider(config.provider.clone(), Some(self.auth_manager.clone())), + thread_store.get::().map(|surface| { + HttpStateContext::new(self.auth_manager.codex_home().to_path_buf(), *surface) + }), )))] } } diff --git a/codex-rs/ext/web-search/Cargo.toml b/codex-rs/ext/web-search/Cargo.toml index 0faabf21a8..e007cb1244 100644 --- a/codex-rs/ext/web-search/Cargo.toml +++ b/codex-rs/ext/web-search/Cargo.toml @@ -18,6 +18,7 @@ codex-api = { workspace = true } codex-core = { workspace = true } codex-extension-api = { workspace = true } codex-features = { workspace = true } +codex-http-state = { workspace = true } codex-login = { workspace = true } codex-model-provider = { workspace = true } codex-model-provider-info = { workspace = true } diff --git a/codex-rs/ext/web-search/src/extension.rs b/codex-rs/ext/web-search/src/extension.rs index 30fc48961e..e80847c155 100644 --- a/codex-rs/ext/web-search/src/extension.rs +++ b/codex-rs/ext/web-search/src/extension.rs @@ -14,6 +14,8 @@ use codex_extension_api::ThreadLifecycleContributor; use codex_extension_api::ThreadStartInput; use codex_extension_api::ToolContributor; use codex_features::Feature; +use codex_http_state::HttpStateContext; +use codex_http_state::HttpStateSurface; use codex_login::AuthManager; use codex_model_provider::create_model_provider; use codex_model_provider_info::ModelProviderInfo; @@ -121,6 +123,9 @@ impl ToolContributor for WebSearchExtension { config.provider.clone(), Some(self.auth_manager.clone()), ), + http_state: thread_store.get::().map(|surface| { + HttpStateContext::new(self.auth_manager.codex_home().to_path_buf(), *surface) + }), settings: config.settings.clone(), })] } diff --git a/codex-rs/ext/web-search/src/tool.rs b/codex-rs/ext/web-search/src/tool.rs index 9a09b73ebb..684ddf819d 100644 --- a/codex-rs/ext/web-search/src/tool.rs +++ b/codex-rs/ext/web-search/src/tool.rs @@ -14,8 +14,10 @@ use codex_extension_api::ToolName; use codex_extension_api::ToolOutput; use codex_extension_api::ToolSpec; use codex_extension_api::parse_tool_input_schema_without_compaction; +use codex_http_state::HttpStateContext; use codex_login::default_client::build_reqwest_client; use codex_model_provider::SharedModelProvider; +use codex_model_provider::with_native_integrity_state; use codex_protocol::items::WebSearchItem; use codex_protocol::models::WebSearchAction; use codex_tools::ResponsesApiNamespace; @@ -36,6 +38,7 @@ const WEB_RUN_DESCRIPTION: &str = include_str!("../web_run_description.md"); pub(crate) struct WebSearchTool { pub(crate) session_id: String, pub(crate) provider: SharedModelProvider, + pub(crate) http_state: Option, pub(crate) settings: SearchSettings, } @@ -77,16 +80,23 @@ impl ToolExecutor for WebSearchTool { async fn handle(&self, call: ToolCall) -> Result, FunctionCallError> { let commands = parse_commands(&call)?; let command_action = command_action(&commands); + let auth = match self.http_state.as_ref() { + Some(http_state) => self.provider.auth_for_surface(http_state.surface()).await, + None => self.provider.auth().await, + }; let provider = self .provider .api_provider() .await .map_err(|err| FunctionCallError::Fatal(err.to_string()))?; - let auth = self - .provider - .api_auth() - .await - .map_err(|err| FunctionCallError::Fatal(err.to_string()))?; + let auth = with_native_integrity_state( + self.provider + .api_auth() + .await + .map_err(|err| FunctionCallError::Fatal(err.to_string()))?, + auth.as_ref(), + self.http_state.clone(), + ); let client = SearchClient::new( ReqwestTransport::new(build_reqwest_client()), provider, diff --git a/codex-rs/http-state/src/http_state_tests.rs b/codex-rs/http-state/src/http_state_tests.rs index e54378b90f..9ffd261e42 100644 --- a/codex-rs/http-state/src/http_state_tests.rs +++ b/codex-rs/http-state/src/http_state_tests.rs @@ -77,6 +77,18 @@ fn stores_state_in_one_file_per_surface() { ); } +#[test] +fn derives_an_independent_surface_selection() { + let codex_home = TempDir::new().expect("tempdir"); + let cli = HttpStateContext::new(codex_home.path().to_path_buf(), HttpStateSurface::CodexCli); + let desktop = cli.for_surface(HttpStateSurface::CodexDesktop); + + desktop.set_surface(HttpStateSurface::CodexVscode); + + assert_eq!(cli.surface(), HttpStateSurface::CodexCli); + assert_eq!(desktop.surface(), HttpStateSurface::CodexVscode); +} + #[test] fn compare_and_set_rejects_a_stale_prior_value() { let codex_home = TempDir::new().expect("tempdir"); diff --git a/codex-rs/http-state/src/lib.rs b/codex-rs/http-state/src/lib.rs index 7218a7d44a..93405d1fe6 100644 --- a/codex-rs/http-state/src/lib.rs +++ b/codex-rs/http-state/src/lib.rs @@ -247,6 +247,13 @@ impl HttpStateContext { true } + pub fn for_surface(&self, surface: HttpStateSurface) -> Self { + Self { + store: self.store.clone(), + surface: Arc::new(RwLock::new(surface)), + } + } + pub fn get_for_surface(&self, surface: HttpStateSurface) -> io::Result> { self.store.get(surface) } diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index e2ef3f3650..aac1ed0326 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use codex_api::Provider; use codex_api::SharedAuthProvider; +use codex_http_state::HttpStateSurface; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_model_provider_info::ModelProviderInfo; @@ -112,6 +113,14 @@ pub trait ModelProvider: fmt::Debug + Send + Sync { /// Returns the current provider-scoped auth value, if one is configured. async fn auth(&self) -> Option; + /// Returns the current provider-scoped auth value and attributes any refresh to `surface`. + async fn auth_for_surface(&self, surface: HttpStateSurface) -> Option { + match self.auth_manager() { + Some(auth_manager) => auth_manager.auth_for_surface(surface).await, + None => self.auth().await, + } + } + /// Returns the current app-visible account state for this provider. fn account_state(&self) -> ProviderAccountResult; diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index 6f98f78920..d6ce505e22 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter.rs @@ -27,6 +27,7 @@ use reqwest::header::AUTHORIZATION; use reqwest::header::CONTENT_TYPE; use reqwest::header::HeaderMap; use reqwest::header::HeaderName; +use reqwest::header::HeaderValue; use rmcp::model::ClientJsonRpcMessage; use rmcp::model::JsonRpcMessage; use rmcp::model::ServerJsonRpcMessage; @@ -93,7 +94,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { let has_session_id = session_id.is_some(); let mut headers = self.default_headers.clone(); headers.extend(custom_headers); - self.add_auth_headers(&mut headers); + self.add_auth_headers(&uri, &mut headers); insert_header( &mut headers, ACCEPT, @@ -152,6 +153,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { )); } }; + self.observe_response_headers(&uri, &headers, &response.headers); if response.status == StatusCode::NOT_FOUND.as_u16() && session_id.is_some() { return Err(StreamableHttpError::Client( @@ -216,7 +218,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { ) -> std::result::Result<(), StreamableHttpError> { let mut headers = self.default_headers.clone(); headers.extend(custom_headers); - self.add_auth_headers(&mut headers); + self.add_auth_headers(&uri, &mut headers); if let Some(auth_token) = auth_token { insert_header( &mut headers, @@ -246,6 +248,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { .await .map_err(StreamableHttpClientAdapterError::from) .map_err(StreamableHttpError::Client)?; + self.observe_response_headers(&uri, &headers, &response.headers); if response.status == StatusCode::METHOD_NOT_ALLOWED.as_u16() { return Ok(()); @@ -271,7 +274,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { > { let mut headers = self.default_headers.clone(); headers.extend(custom_headers); - self.add_auth_headers(&mut headers); + self.add_auth_headers(&uri, &mut headers); insert_header( &mut headers, ACCEPT, @@ -315,6 +318,7 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { .await .map_err(StreamableHttpClientAdapterError::from) .map_err(StreamableHttpError::Client)?; + self.observe_response_headers(&uri, &headers, &response.headers); if response.status == StatusCode::METHOD_NOT_ALLOWED.as_u16() { return Err(StreamableHttpError::ServerDoesNotSupportSse); @@ -347,9 +351,24 @@ impl StreamableHttpClient for StreamableHttpClientAdapter { } impl StreamableHttpClientAdapter { - fn add_auth_headers(&self, headers: &mut HeaderMap) { + fn add_auth_headers(&self, request_url: &str, headers: &mut HeaderMap) { if let Some(auth_provider) = &self.auth_provider { - headers.extend(auth_provider.to_auth_headers()); + auth_provider.add_auth_headers_for_url(request_url, headers); + } + } + + fn observe_response_headers( + &self, + request_url: &str, + request_headers: &HeaderMap, + response_headers: &[HttpHeader], + ) { + if let Some(auth_provider) = &self.auth_provider { + auth_provider.observe_response_headers( + request_url, + request_headers, + &header_map(response_headers), + ); } } } @@ -451,6 +470,18 @@ fn protocol_headers(headers: &HeaderMap) -> Vec { .collect() } +fn header_map(headers: &[HttpHeader]) -> HeaderMap { + headers + .iter() + .filter_map(|header| { + Some(( + HeaderName::from_bytes(header.name.as_bytes()).ok()?, + HeaderValue::from_str(&header.value).ok()?, + )) + }) + .collect() +} + fn response_header(headers: &[HttpHeader], name: impl AsRef) -> Option { let name = name.as_ref(); headers @@ -490,3 +521,85 @@ fn sse_stream_from_body( })) .boxed() } + +#[cfg(test)] +mod tests { + use std::sync::Mutex as StdMutex; + + use codex_api::AuthProvider; + use codex_exec_server::ReqwestHttpClient; + use pretty_assertions::assert_eq; + + use super::*; + + #[derive(Debug, Default)] + struct RecordingAuthProvider { + observed_updates: StdMutex>, + } + + impl AuthProvider for RecordingAuthProvider { + fn add_auth_headers(&self, _headers: &mut HeaderMap) {} + + fn add_auth_headers_for_url(&self, _request_url: &str, headers: &mut HeaderMap) { + let _ = headers.insert("x-test-state", HeaderValue::from_static("sent-state")); + } + + fn observe_response_headers( + &self, + request_url: &str, + request_headers: &HeaderMap, + response_headers: &HeaderMap, + ) { + let sent_state = request_headers + .get("x-test-state") + .and_then(|value| value.to_str().ok()); + let update_state = response_headers + .get("x-test-state-update") + .and_then(|value| value.to_str().ok()); + if let (Some(sent_state), Some(update_state)) = (sent_state, update_state) { + self.observed_updates + .lock() + .expect("observed updates lock should not be poisoned") + .push(( + request_url.to_string(), + sent_state.to_string(), + update_state.to_string(), + )); + } + } + } + + #[test] + fn uses_url_scoped_auth_and_observes_forwarded_response_headers() { + let auth = Arc::new(RecordingAuthProvider::default()); + let adapter = StreamableHttpClientAdapter::new( + Arc::new(ReqwestHttpClient), + HeaderMap::new(), + Some(auth.clone()), + ); + let request_url = "https://chatgpt.com/backend-api/wham/apps"; + let mut request_headers = HeaderMap::new(); + + adapter.add_auth_headers(request_url, &mut request_headers); + adapter.observe_response_headers( + request_url, + &request_headers, + &[HttpHeader { + name: "x-test-state-update".to_string(), + value: "next-state".to_string(), + }], + ); + + assert_eq!( + *auth + .observed_updates + .lock() + .expect("observed updates lock should not be poisoned"), + vec![( + request_url.to_string(), + "sent-state".to_string(), + "next-state".to_string(), + )] + ); + } +}