From a0d5fd772e5dbc1be57c5a689cc4db1eabe5035d Mon Sep 17 00:00:00 2001 From: canvrno-oai Date: Wed, 24 Jun 2026 18:48:11 -0700 Subject: [PATCH 1/7] TUI Plugin Sharing 5 - polish remote plugin catalog rows (#26705) This is the final plugin sharing PR in the 5-PR stack. It applies the remaining TUI polish for remote plugin catalog rows and tabs: admin-disabled plugins now read as blocked/view-only instead of looking toggleable, admin-installed/default-installed plugins count and sort like installed plugins, plugin search matches richer metadata, and an empty successful `Shared with me` section stays hidden. - Admin-disabled rows use a blocked marker, show `Disabled`, and keep Enter-only detail behavior without a toggle hint. - Admin-installed/default-installed plugins show as installed in counts, ordering, tabs, and detail copy. - Plugin search now matches descriptions and keywords in addition to existing row metadata. - Successful-empty `Shared with me` tabs are hidden, while loading, error, workspace-empty, and real shared-plugin states remain visible. - Updates coverage in `plugins_popup_snapshot_shows_all_marketplaces_and_sorts_installed_then_name`, `plugins_popup_admin_disabled_installed_plugin_has_no_toggle_hint`, `plugins_popup_search_matches_plugin_descriptions`, and `plugins_popup_remote_section_fallback_states_snapshot`. - Updates snapshots `plugins_popup_curated_marketplace` and `plugins_popup_empty_shared_section_hidden`. image image --- .../src/bottom_pane/list_selection_view.rs | 98 +++++---- codex-rs/tui/src/bottom_pane/mod.rs | 2 + codex-rs/tui/src/chatwidget/plugin_catalog.rs | 97 +++++++-- ..._tests__plugin_detail_popup_installed.snap | 18 +- ...lugins_popup_admin_disabled_installed.snap | 14 ++ ...ts__plugins_popup_curated_marketplace.snap | 10 +- ...ins_popup_empty_shared_section_hidden.snap | 14 ++ ...ins_popup_newly_installed_marketplace.snap | 4 +- ..._tests__plugins_popup_search_filtered.snap | 4 +- .../chatwidget/tests/popups_and_settings.rs | 200 ++++++++++++++---- codex-rs/tui/src/custom_terminal.rs | 11 +- codex-rs/tui/src/tui/test_support.rs | 2 +- 12 files changed, 362 insertions(+), 112 deletions(-) create mode 100644 codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_admin_disabled_installed.snap create mode 100644 codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_empty_shared_section_hidden.snap diff --git a/codex-rs/tui/src/bottom_pane/list_selection_view.rs b/codex-rs/tui/src/bottom_pane/list_selection_view.rs index 7af981d7de..4455a6b670 100644 --- a/codex-rs/tui/src/bottom_pane/list_selection_view.rs +++ b/codex-rs/tui/src/bottom_pane/list_selection_view.rs @@ -281,6 +281,63 @@ pub(crate) struct ListSelectionView { keymap: ListKeymap, } +const SELECTION_TOGGLE_ON_PREFIX: &str = "[*] "; +const SELECTION_TOGGLE_OFF_PREFIX: &str = "[ ] "; +pub(crate) const SELECTION_TOGGLE_UNAVAILABLE_PREFIX: &str = "[-] "; +pub(crate) const SELECTION_TOGGLE_BLOCKED_PREFIX: &str = "[!] "; + +fn selection_toggle_prefix(toggle: &SelectionToggle) -> &'static str { + if toggle.is_on { + SELECTION_TOGGLE_ON_PREFIX + } else { + SELECTION_TOGGLE_OFF_PREFIX + } +} + +fn selection_item_toggle_prefix(item: &SelectionItem) -> Option<&'static str> { + item.toggle + .as_ref() + .map(selection_toggle_prefix) + .or(item.toggle_placeholder) +} + +impl ListSelectionView { + fn selected_item_has_toggle(&self) -> bool { + self.selected_actual_idx() + .and_then(|actual_idx| self.active_items().get(actual_idx)) + .is_some_and(|item| item.toggle.is_some() && Self::item_is_enabled(item)) + } + + fn selected_item_has_toggle_placeholder(&self) -> bool { + self.selected_actual_idx() + .and_then(|actual_idx| self.active_items().get(actual_idx)) + .is_some_and(|item| { + item.toggle.is_none() + && item.toggle_placeholder.is_some() + && Self::item_is_enabled(item) + }) + } + + fn toggle_selected(&mut self) { + let Some(actual_idx) = self.selected_actual_idx() else { + return; + }; + let app_event_tx = self.app_event_tx.clone(); + let Some(item) = self.active_items_mut().get_mut(actual_idx) else { + return; + }; + if !Self::item_is_enabled(item) { + return; + } + let Some(toggle) = item.toggle.as_mut() else { + return; + }; + + toggle.is_on = !toggle.is_on; + (toggle.action)(toggle.is_on, &app_event_tx); + } +} + impl ListSelectionView { /// Create a selection popup view with filtering, scrolling, and callbacks wired. /// @@ -533,10 +590,8 @@ impl ListSelectionView { let wrap_prefix_width = UnicodeWidthStr::width(wrap_prefix.as_str()); let mut name_prefix_spans = Vec::new(); name_prefix_spans.push(wrap_prefix.into()); - if let Some(toggle) = &item.toggle { - name_prefix_spans.push(if toggle.is_on { "[*] " } else { "[ ] " }.into()); - } else if let Some(placeholder) = item.toggle_placeholder { - name_prefix_spans.push(placeholder.into()); + if let Some(toggle_prefix) = selection_item_toggle_prefix(item) { + name_prefix_spans.push(toggle_prefix.into()); } name_prefix_spans.extend(item.name_prefix_spans.clone()); let description = is_selected @@ -611,22 +666,6 @@ impl ListSelectionView { item.disabled_reason.is_none() && !item.is_disabled } - fn selected_item_has_toggle(&self) -> bool { - self.selected_actual_idx() - .and_then(|actual_idx| self.active_items().get(actual_idx)) - .is_some_and(|item| item.toggle.is_some() && Self::item_is_enabled(item)) - } - - fn selected_item_has_toggle_placeholder(&self) -> bool { - self.selected_actual_idx() - .and_then(|actual_idx| self.active_items().get(actual_idx)) - .is_some_and(|item| { - item.toggle.is_none() - && item.toggle_placeholder.is_some() - && Self::item_is_enabled(item) - }) - } - fn actual_idx_for_enabled_number(&self, number: usize) -> Option { if number == 0 { return None; @@ -640,25 +679,6 @@ impl ListSelectionView { .map(|(idx, _)| idx) } - fn toggle_selected(&mut self) { - let Some(actual_idx) = self.selected_actual_idx() else { - return; - }; - let app_event_tx = self.app_event_tx.clone(); - let Some(item) = self.active_items_mut().get_mut(actual_idx) else { - return; - }; - if !Self::item_is_enabled(item) { - return; - } - let Some(toggle) = item.toggle.as_mut() else { - return; - }; - - toggle.is_on = !toggle.is_on; - (toggle.action)(toggle.is_on, &app_event_tx); - } - fn move_up(&mut self) { let before = self.selected_actual_idx(); let len = self.visible_len(); diff --git a/codex-rs/tui/src/bottom_pane/mod.rs b/codex-rs/tui/src/bottom_pane/mod.rs index 7fffdeb98e..ad30dedf5a 100644 --- a/codex-rs/tui/src/bottom_pane/mod.rs +++ b/codex-rs/tui/src/bottom_pane/mod.rs @@ -190,6 +190,8 @@ use crate::status_indicator_widget::StatusDetailsCapitalization; use crate::status_indicator_widget::StatusIndicatorWidget; pub(crate) use experimental_features_view::ExperimentalFeatureItem; pub(crate) use experimental_features_view::ExperimentalFeaturesView; +pub(crate) use list_selection_view::SELECTION_TOGGLE_BLOCKED_PREFIX; +pub(crate) use list_selection_view::SELECTION_TOGGLE_UNAVAILABLE_PREFIX; pub(crate) use list_selection_view::SelectionAction; pub(crate) use list_selection_view::SelectionItem; diff --git a/codex-rs/tui/src/chatwidget/plugin_catalog.rs b/codex-rs/tui/src/chatwidget/plugin_catalog.rs index a19ede3ab5..cb92ea7d31 100644 --- a/codex-rs/tui/src/chatwidget/plugin_catalog.rs +++ b/codex-rs/tui/src/chatwidget/plugin_catalog.rs @@ -12,6 +12,8 @@ use crate::app_event::AppEvent; use crate::app_event::PluginLocation; use crate::app_event::PluginRemoteSectionError; use crate::bottom_pane::ColumnWidthMode; +use crate::bottom_pane::SELECTION_TOGGLE_BLOCKED_PREFIX; +use crate::bottom_pane::SELECTION_TOGGLE_UNAVAILABLE_PREFIX; use crate::bottom_pane::SelectionAction; use crate::bottom_pane::SelectionItem; use crate::bottom_pane::SelectionRowDisplay; @@ -67,6 +69,8 @@ const PERSONAL_MARKETPLACE_RELATIVE_PATH: &str = ".agents/plugins/marketplace.js const REMOTE_LOADING_TAB_ID_PREFIX: &str = "remote-loading:"; const REMOTE_EMPTY_TAB_ID_PREFIX: &str = "remote-empty:"; const REMOTE_ERROR_TAB_ID_PREFIX: &str = "remote-error:"; +const OPENAI_CURATED_LOADING_DESCRIPTION: &str = + "This updates when OpenAI Curated plugins finish loading."; const WORKSPACE_SECTION_TAB_ORDER: u8 = 0; const SHARED_WITH_ME_SECTION_TAB_ORDER: u8 = 1; const SHARED_WITH_ME_LINK_SECTION_TAB_ORDER: u8 = 2; @@ -78,6 +82,7 @@ struct PreferredLocalPluginSource { marketplace_path: AbsolutePathBuf, plugin_name: String, installed: bool, + install_policy: PluginInstallPolicy, } #[derive(Debug, Clone, Copy)] @@ -153,7 +158,9 @@ struct RemoteMarketplaceSection { id: &'static str, label: &'static str, loading_tab_id: &'static str, + loading_item_description: &'static str, marketplace_names: &'static [&'static str], + show_empty_tab: bool, empty_item_name: &'static str, empty_item_description: &'static str, tab_order: u8, @@ -164,7 +171,9 @@ const REMOTE_MARKETPLACE_SECTIONS: [RemoteMarketplaceSection; 2] = [ id: "workspace", label: "Workspace", loading_tab_id: "workspace-loading", + loading_item_description: "This updates when workspace plugins finish loading.", marketplace_names: &[REMOTE_WORKSPACE_MARKETPLACE_NAME], + show_empty_tab: true, empty_item_name: "No workspace plugins available", empty_item_description: "No workspace directory plugins are available.", tab_order: WORKSPACE_SECTION_TAB_ORDER, @@ -173,11 +182,13 @@ const REMOTE_MARKETPLACE_SECTIONS: [RemoteMarketplaceSection; 2] = [ id: "shared-with-me", label: "Shared with me", loading_tab_id: "shared-with-me-loading", + loading_item_description: "This updates when shared plugins finish loading.", marketplace_names: &[ REMOTE_WORKSPACE_SHARED_WITH_ME_MARKETPLACE_NAME, REMOTE_WORKSPACE_SHARED_WITH_ME_PRIVATE_MARKETPLACE_NAME, REMOTE_WORKSPACE_SHARED_WITH_ME_UNLISTED_MARKETPLACE_NAME, ], + show_empty_tab: false, empty_item_name: "No shared plugins available", empty_item_description: "No plugins have been shared with you.", tab_order: SHARED_WITH_ME_SECTION_TAB_ORDER, @@ -200,10 +211,16 @@ impl RemoteMarketplaceSection { } let tab = if remote_sections_loading { - remote_section_loading_tab(self.loading_tab_id, self.label) + remote_section_loading_tab( + self.loading_tab_id, + self.label, + self.loading_item_description, + ) } else if remote_sections_loaded { if let Some(section_error) = plugin_remote_section_error(section_errors, self.id) { remote_section_error_tab(section_error) + } else if !self.show_empty_tab { + return None; } else { remote_section_empty_tab( self.id, @@ -799,7 +816,7 @@ impl ChatWidget { if curated_loading && !curated_has_entries { ( "Loading OpenAI Curated plugins...", - "This section updates when app-server returns it.", + OPENAI_CURATED_LOADING_DESCRIPTION, ) } else if let Some(section_error) = by_openai_section_error && !curated_has_entries @@ -819,7 +836,10 @@ impl ChatWidget { curated_empty_description, ); if curated_loading && curated_has_entries { - curated_items.push(remote_section_loading_item("OpenAI Curated")); + curated_items.push(remote_section_loading_item( + "OpenAI Curated", + OPENAI_CURATED_LOADING_DESCRIPTION, + )); } if let Some(section_error) = by_openai_section_error && curated_has_entries @@ -1024,7 +1044,16 @@ impl ChatWidget { }]; if plugin.summary.installed { - if let Some(plugin_id) = plugin_uninstall_id(&plugin.summary) { + if plugin.summary.install_policy == PluginInstallPolicy::InstalledByDefault { + items.push(SelectionItem { + name: "Installed by admin".to_string(), + description: Some( + "This plugin is installed by your workspace admin.".to_string(), + ), + is_disabled: true, + ..Default::default() + }); + } else if let Some(plugin_id) = plugin_uninstall_id(&plugin.summary) { let uninstall_cwd = self.config.cwd.to_path_buf(); let plugin_display_name = display_name; items.push(SelectionItem { @@ -1164,7 +1193,9 @@ impl ChatWidget { plugin_detail_request_for_entry(marketplace, plugin, preferred_local_sources); let can_view_details = plugin_detail_request.is_some(); let disabled_by_admin = plugin.availability == PluginAvailability::DisabledByAdmin; - let can_toggle_plugin = plugin.installed && !disabled_by_admin; + let can_toggle_plugin = plugin.installed + && plugin.install_policy != PluginInstallPolicy::InstalledByDefault + && !disabled_by_admin; let selected_status_label = format!("{status_label: &'static str { if plugin.availability == PluginAvailability::DisabledByAdmin { return "Disabled by admin"; } + if plugin.install_policy == PluginInstallPolicy::InstalledByDefault { + return if plugin.installed { + "Installed by admin" + } else { + "Enabled by Admin" + }; + } if plugin.installed { if plugin.enabled { "Installed" @@ -1370,7 +1429,7 @@ fn plugin_detail_status_label(plugin: &PluginSummary) -> &'static str { match plugin.install_policy { PluginInstallPolicy::NotAvailable => "Not installable", PluginInstallPolicy::Available => "Can be installed", - PluginInstallPolicy::InstalledByDefault => "Available by default", + PluginInstallPolicy::InstalledByDefault => "Installed by admin", } } } @@ -1643,10 +1702,10 @@ fn is_personal_marketplace_path(marketplace_path: &AbsolutePathBuf) -> bool { .is_some_and(|personal_path| personal_path.as_path() == marketplace_path.as_path()) } -fn remote_section_loading_item(label: &str) -> SelectionItem { +fn remote_section_loading_item(label: &str, description: &str) -> SelectionItem { SelectionItem { name: format!("Loading {label} plugins..."), - description: Some("This section updates when app-server returns it.".to_string()), + description: Some(description.to_string()), is_disabled: true, ..Default::default() } @@ -1670,7 +1729,7 @@ fn plugin_remote_section_error<'a>( .find(|section_error| section_error.section_id == section_id) } -fn remote_section_loading_tab(id: &str, label: &str) -> SelectionTab { +fn remote_section_loading_tab(id: &str, label: &str, item_description: &str) -> SelectionTab { SelectionTab { id: format!("{REMOTE_LOADING_TAB_ID_PREFIX}{id}"), label: label.to_string(), @@ -1678,7 +1737,7 @@ fn remote_section_loading_tab(id: &str, label: &str) -> SelectionTab { format!("Loading {label} plugins."), "Local plugin functionality is already available.".to_string(), ), - items: vec![remote_section_loading_item(label)], + items: vec![remote_section_loading_item(label, item_description)], } } @@ -1816,7 +1875,10 @@ fn plugin_brief_description_without_marketplace( fn plugin_status_label(plugin: &PluginSummary) -> &'static str { if plugin.availability == PluginAvailability::DisabledByAdmin { - return "Disabled by admin"; + return "Disabled"; + } + if !plugin.installed && plugin.install_policy == PluginInstallPolicy::InstalledByDefault { + return "Admin assigned"; } if plugin.installed { if plugin.enabled { @@ -1828,7 +1890,7 @@ fn plugin_status_label(plugin: &PluginSummary) -> &'static str { match plugin.install_policy { PluginInstallPolicy::NotAvailable => "Not installable", PluginInstallPolicy::Available => "Available", - PluginInstallPolicy::InstalledByDefault => "Available by default", + PluginInstallPolicy::InstalledByDefault => "Installed", } } } @@ -1863,6 +1925,7 @@ fn plugin_detail_request_for_entry( && let Some(remote_plugin_id) = plugin_remote_identity(plugin) && let Some(preferred_source) = preferred_local_sources.get(remote_plugin_id) && preferred_source.installed == plugin.installed + && preferred_source.install_policy == plugin.install_policy { return Some(( PluginLocation::Local { diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugin_detail_popup_installed.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugin_detail_popup_installed.snap index 2b69864cc5..bf1a4165b2 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugin_detail_popup_installed.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugin_detail_popup_installed.snap @@ -3,16 +3,16 @@ source: tui/src/chatwidget/tests/popups_and_settings.rs expression: strip_osc8_for_snapshot(&popup) --- Plugins - Figma · Installed · ChatGPT Marketplace + Figma · Installed by admin · ChatGPT Marketplace Turn Figma files into implementation context. -› 1. Back to plugins Return to the plugin list. - 2. Uninstall plugin Remove this plugin now. - Source Local - Auth Auth on install - Skills design-review, extract-copy - Hooks PreToolUse (1), Stop (2) - Apps Figma, Slack - MCP Servers figma-mcp, docs-mcp +› 1. Back to plugins Return to the plugin list. + Installed by admin This plugin is installed by your workspace admin. + Source Local + Auth Auth on install + Skills design-review, extract-copy + Hooks PreToolUse (1), Stop (2) + Apps Figma, Slack + MCP Servers figma-mcp, docs-mcp Press esc to close. diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_admin_disabled_installed.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_admin_disabled_installed.snap new file mode 100644 index 0000000000..b8f109a9fc --- /dev/null +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_admin_disabled_installed.snap @@ -0,0 +1,14 @@ +--- +source: tui/src/chatwidget/tests/popups_and_settings.rs +expression: popup +--- + Plugins + Browse plugins from available marketplaces. + Installed 1 of 1 available plugins. + + [All Plugins] Installed (1) OpenAI Curated Workspace Add Marketplace + + Type to search plugins +› [!] Admin Blocked Disabled Press Enter to view plugin details. + + space enable/disable · ←/→ select marketplace · enter view details · esc close diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_curated_marketplace.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_curated_marketplace.snap index 2ec91937fb..fbf5104c37 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_curated_marketplace.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_curated_marketplace.snap @@ -6,12 +6,12 @@ expression: popup Browse plugins from available marketplaces. Installed 1 of 4 available plugins. - [All Plugins] Installed (1) OpenAI Curated Workspace Shared with me Repo Marketplace Add Marketplace + [All Plugins] Installed (1) OpenAI Curated Workspace Repo Marketplace Add Marketplace Type to search plugins -› [ ] Alpha Sync Disabled Space to enable; Enter view details. - [-] Bravo Search Available · OpenAI Curated · Search docs and tickets. - [-] Hidden Repo Plugin Available · Repo Marketplace · Should not be shown in /plugins. - [-] Starter Available by default · OpenAI Curated · Included by default. +› [ ] Alpha Sync Disabled Space to enable; Enter view details. + [-] Bravo Search Available · OpenAI Curated · Search docs and tickets. + [-] Hidden Repo Plugin Available · Repo Marketplace · Should not be shown in /plugins. + [-] Starter Admin assigned · OpenAI Curated · Included by default. space enable/disable · ←/→ select marketplace · enter view details · esc close diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_empty_shared_section_hidden.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_empty_shared_section_hidden.snap new file mode 100644 index 0000000000..146d02a5a2 --- /dev/null +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_empty_shared_section_hidden.snap @@ -0,0 +1,14 @@ +--- +source: tui/src/chatwidget/tests/popups_and_settings.rs +expression: loaded_popup +--- + Plugins + Workspace. + This section loaded successfully. + + All Plugins Installed (0) OpenAI Curated [Workspace] Add Marketplace + + Type to search plugins +› No workspace plugins available No workspace directory plugins are available. + + space enable/disable · ←/→ select marketplace · enter view details · esc close diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_newly_installed_marketplace.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_newly_installed_marketplace.snap index 803ff9fbe6..60ebee6f9a 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_newly_installed_marketplace.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_newly_installed_marketplace.snap @@ -1,13 +1,13 @@ --- source: tui/src/chatwidget/tests/popups_and_settings.rs +assertion_line: 473 expression: popup --- Plugins Debug Marketplace installed successfully. Select the plugins you want to use and press Enter to install or view details. - All Plugins Installed (0) OpenAI Curated Workspace Shared with me [Debug Marketplace] - Add Marketplace + All Plugins Installed (0) OpenAI Curated Workspace [Debug Marketplace] Add Marketplace Type to search plugins › [-] Debug Plugin Available Press Enter to install or view plugin details. diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_search_filtered.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_search_filtered.snap index bfb0affb21..3a2e455f6d 100644 --- a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_search_filtered.snap +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugins_popup_search_filtered.snap @@ -1,13 +1,13 @@ --- source: tui/src/chatwidget/tests/popups_and_settings.rs -assertion_line: 767 +assertion_line: 1684 expression: popup --- Plugins Browse plugins from available marketplaces. Installed 0 of 3 available plugins. - [All Plugins] Installed (0) OpenAI Curated Workspace Shared with me Add Marketplace + [All Plugins] Installed (0) OpenAI Curated Workspace Add Marketplace sla › [-] Slack Available Press Enter to install or view plugin details. diff --git a/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs b/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs index a8053dc5a3..7d162d41e4 100644 --- a/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs +++ b/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs @@ -186,7 +186,7 @@ async fn plugins_popup_snapshot_shows_all_marketplaces_and_sorts_installed_then_ Some("Starter"), Some("Included by default."), /*installed*/ false, - /*enabled*/ true, + /*enabled*/ false, PluginInstallPolicy::InstalledByDefault, ), ]), @@ -665,10 +665,9 @@ async fn plugin_detail_popup_snapshot_labels_personal_marketplace_as_local() { } #[tokio::test] -async fn plugin_detail_popup_hides_disclosure_for_installed_plugins() { +async fn plugin_detail_popup_distinguishes_admin_installed_from_enabled() { let (mut chat, _rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; chat.set_feature_enabled(Feature::Plugins, /*enabled*/ true); - let summary = plugins_test_summary( "plugin-figma", "figma", @@ -676,7 +675,7 @@ async fn plugin_detail_popup_hides_disclosure_for_installed_plugins() { Some("Design handoff."), /*installed*/ true, /*enabled*/ true, - PluginInstallPolicy::Available, + PluginInstallPolicy::InstalledByDefault, ); let response = plugins_test_response(vec![plugins_test_curated_marketplace(vec![ summary.clone(), @@ -684,32 +683,54 @@ async fn plugin_detail_popup_hides_disclosure_for_installed_plugins() { let cwd = chat.config.cwd.clone(); chat.on_plugins_loaded(cwd.to_path_buf(), Ok(response)); chat.add_plugins_output(); + let mut plugin = plugins_test_detail( + summary, + Some("Turn Figma files into implementation context."), + &["design-review", "extract-copy"], + &[ + (codex_app_server_protocol::HookEventName::PreToolUse, 1), + (codex_app_server_protocol::HookEventName::Stop, 2), + ], + &["Figma", "Slack"], + &["figma-mcp", "docs-mcp"], + ); chat.on_plugin_detail_loaded( cwd.to_path_buf(), Ok(PluginReadResponse { - plugin: plugins_test_detail( - summary, - Some("Turn Figma files into implementation context."), - &["design-review", "extract-copy"], - &[ - (codex_app_server_protocol::HookEventName::PreToolUse, 1), - (codex_app_server_protocol::HookEventName::Stop, 2), - ], - &["Figma", "Slack"], - &["figma-mcp", "docs-mcp"], - ), + plugin: plugin.clone(), }), ); let popup = render_bottom_popup(&chat, /*width*/ 100); assert!( - !popup.contains("Data shared with this app is subject to the app's"), - "expected installed plugin details to hide the disclosure line, got:\n{popup}" + popup.contains("Installed by admin") + && !popup.contains("Uninstall plugin") + && !popup.contains("Data shared with this app is subject to the app's"), + "expected admin-installed plugin details to block uninstall and hide the disclosure line, got:\n{popup}" ); assert_chatwidget_snapshot!( "plugin_detail_popup_installed", strip_osc8_for_snapshot(&popup) ); + + plugin.summary.installed = false; + plugin.summary.enabled = false; + chat.on_plugin_detail_loaded(cwd.to_path_buf(), Ok(PluginReadResponse { plugin })); + let popup = render_bottom_popup(&chat, /*width*/ 100); + assert!( + popup.contains("Enabled by Admin") + && popup.contains("Install plugin") + && !popup.contains("Installed by admin"), + "expected an unmaterialized default plugin to show its admin assignment and remain installable, got:\n{popup}" + ); + insta::assert_snapshot!( + popup + .lines() + .find(|line| line.contains("Figma ·")) + .expect("expected plugin detail header") + .trim(), + @"Figma · Enabled by Admin · ChatGPT Marketplace" + ); } #[tokio::test] @@ -769,17 +790,20 @@ async fn plugins_popup_remote_row_opens_remote_detail() { } #[tokio::test] -async fn plugin_detail_remote_install_uses_remote_location() { +async fn plugin_detail_unmaterialized_default_uses_remote_install_path() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; chat.set_feature_enabled(Feature::Plugins, /*enabled*/ true); - let summary = plugins_test_remote_summary( - "plugins~Plugin_linear", - "linear", - Some("Linear"), - Some("Issue tracking."), - /*installed*/ false, - ); + let summary = PluginSummary { + install_policy: PluginInstallPolicy::InstalledByDefault, + ..plugins_test_remote_summary( + "plugins~Plugin_linear", + "linear", + Some("Linear"), + Some("Issue tracking."), + /*installed*/ false, + ) + }; let cwd = chat.config.cwd.clone(); chat.on_plugins_loaded( cwd.to_path_buf(), @@ -1075,9 +1099,12 @@ async fn plugins_popup_admin_disabled_installed_plugin_has_no_toggle_hint() { ); let popup = render_bottom_popup(&chat, /*width*/ 120); + assert_chatwidget_snapshot!("plugins_popup_admin_disabled_installed", popup); assert!( - popup.contains("Disabled by admin") + popup.contains("[!] Admin Blocked") + && popup.contains("Disabled") && popup.contains("Press Enter to view plugin details.") + && !popup.contains("Disabled by admin") && !popup.contains("Space to disable"), "expected admin-disabled installed row to omit toggle hint, got:\n{popup}" ); @@ -1170,9 +1197,16 @@ async fn plugins_popup_remote_section_fallback_states_snapshot() { let curated_loading_popup = select_tab_containing(&mut chat, "Loading OpenAI Curated plugins..."); let workspace_loading_popup = select_tab_containing(&mut chat, "Loading Workspace plugins."); + let shared_loading_popup = select_tab_containing(&mut chat, "Loading Shared with me plugins."); + let _ = select_tab_containing(&mut chat, "Loading Workspace plugins."); chat.on_plugin_remote_sections_loaded(cwd.to_path_buf(), Vec::new(), Vec::new()); - let shared_empty_popup = select_tab_containing(&mut chat, "Shared with me."); + let loaded_popup = render_bottom_popup(&chat, /*width*/ 100); + assert_chatwidget_snapshot!("plugins_popup_empty_shared_section_hidden", loaded_popup); + assert!( + !loaded_popup.contains("Shared with me"), + "expected empty shared section to stay hidden, got:\n{loaded_popup}" + ); chat.on_plugin_remote_sections_loaded( cwd.to_path_buf(), @@ -1203,20 +1237,20 @@ async fn plugins_popup_remote_section_fallback_states_snapshot() { [ remote_section_state(&curated_loading_popup), remote_section_state(&workspace_loading_popup), - remote_section_state(&shared_empty_popup), + remote_section_state(&shared_loading_popup), remote_section_state(&workspace_error_popup), remote_section_state(&remote_curated_empty_popup), ] .join("\n\n"), @r###" OpenAI Curated marketplace. - Loading OpenAI Curated plugins... This section updates when app-server returns it. + Loading OpenAI Curated plugins... This updates when OpenAI Curated plugins finish loading. Loading Workspace plugins. - Loading Workspace plugins... This section updates when app-server returns it. + Loading Workspace plugins... This updates when workspace plugins finish loading. - Shared with me. - No shared plugins available No plugins have been shared with you. + Loading Shared with me plugins. + Loading Shared with me plugins... This updates when shared plugins finish loading. Workspace unavailable. Workspace unavailable Sign in to ChatGPT to load workspace plugins. @@ -1228,13 +1262,13 @@ async fn plugins_popup_remote_section_fallback_states_snapshot() { } #[tokio::test] -async fn plugins_popup_installed_remote_row_keeps_remote_detail_when_local_share_is_uninstalled() { +async fn plugins_popup_remote_detail_tracks_physical_and_policy_install_state() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; chat.set_feature_enabled(Feature::Plugins, /*enabled*/ true); let remote_plugin_id = "plugins~Plugin_docs"; let remote_marketplace_name = "workspace-shared-with-me-private"; - let local_summary = PluginSummary { + let mut local_summary = PluginSummary { share_context: Some(PluginShareContext { remote_plugin_id: remote_plugin_id.to_string(), remote_version: None, @@ -1251,13 +1285,13 @@ async fn plugins_popup_installed_remote_row_keeps_remote_detail_when_local_share Some("Local editable docs plugin."), /*installed*/ false, /*enabled*/ true, - PluginInstallPolicy::Available, + PluginInstallPolicy::InstalledByDefault, ) }; let popup = render_loaded_plugins_popup( &mut chat, plugins_test_response(vec![ - plugins_test_curated_marketplace(vec![local_summary]), + plugins_test_curated_marketplace(vec![local_summary.clone()]), PluginMarketplaceEntry { name: remote_marketplace_name.to_string(), path: None, @@ -1315,6 +1349,62 @@ async fn plugins_popup_installed_remote_row_keeps_remote_detail_when_local_share } other => panic!("expected FetchPluginDetail event, got {other:?}"), } + + local_summary.install_policy = PluginInstallPolicy::Available; + let mut remote_summary = plugins_test_remote_summary( + remote_plugin_id, + "docs", + Some("Docs"), + Some("Shared docs plugin."), + /*installed*/ false, + ); + remote_summary.install_policy = PluginInstallPolicy::InstalledByDefault; + let popup = render_loaded_plugins_popup( + &mut chat, + plugins_test_response(vec![ + plugins_test_curated_marketplace(vec![local_summary]), + PluginMarketplaceEntry { + name: remote_marketplace_name.to_string(), + path: None, + interface: Some(MarketplaceInterface { + display_name: Some("Shared with me".to_string()), + }), + plugins: vec![remote_summary], + }, + ]), + ); + assert!( + popup.contains("Installed 0 of 1 available plugins.") && popup.contains("Admin assigned"), + "expected the unmaterialized admin-assigned remote duplicate to win without counting as installed, got:\n{popup}" + ); + + chat.handle_key_event(KeyEvent::from(KeyCode::Right)); + let installed_popup = render_bottom_popup(&chat, /*width*/ 100); + assert!( + installed_popup.contains("Showing 0 installed plugins.") + && installed_popup.contains("No installed plugins") + && !installed_popup.contains("Docs"), + "expected the unmaterialized admin-assigned plugin to stay out of the Installed tab, got:\n{installed_popup}" + ); + chat.handle_key_event(KeyEvent::from(KeyCode::Left)); + + while rx.try_recv().is_ok() {} + chat.handle_key_event(KeyEvent::from(KeyCode::Enter)); + assert!(matches!( + rx.try_recv(), + Ok(AppEvent::OpenPluginDetailLoading { .. }) + )); + match rx.try_recv() { + Ok(AppEvent::FetchPluginDetail { params, .. }) => { + assert_eq!(params.marketplace_path, None); + assert_eq!( + params.remote_marketplace_name, + Some(remote_marketplace_name.to_string()) + ); + assert_eq!(params.plugin_name, remote_plugin_id); + } + other => panic!("expected FetchPluginDetail event, got {other:?}"), + } } #[tokio::test] @@ -1685,6 +1775,44 @@ async fn plugins_popup_search_filters_visible_rows_snapshot() { ); } +#[tokio::test] +async fn plugins_popup_search_matches_plugin_descriptions() { + let (mut chat, _rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + chat.set_feature_enabled(Feature::Plugins, /*enabled*/ true); + + render_loaded_plugins_popup( + &mut chat, + plugins_test_response(vec![plugins_test_curated_marketplace(vec![ + plugins_test_summary( + "plugin-calendar", + "calendar", + Some("Calendar"), + Some("Schedule management."), + /*installed*/ false, + /*enabled*/ true, + PluginInstallPolicy::Available, + ), + plugins_test_summary( + "plugin-drive", + "drive", + Some("Drive"), + Some("Document access."), + /*installed*/ false, + /*enabled*/ true, + PluginInstallPolicy::Available, + ), + ])]), + ); + + type_plugins_search_query(&mut chat, "document"); + + let popup = render_bottom_popup(&chat, /*width*/ 100); + assert!( + popup.contains("Drive") && !popup.contains("Calendar"), + "expected plugin search to match descriptions, got:\n{popup}" + ); +} + #[tokio::test] async fn plugins_popup_installed_tab_filters_rows_and_clears_search() { let (mut chat, _rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; diff --git a/codex-rs/tui/src/custom_terminal.rs b/codex-rs/tui/src/custom_terminal.rs index 4800103915..fed3e5456f 100644 --- a/codex-rs/tui/src/custom_terminal.rs +++ b/codex-rs/tui/src/custom_terminal.rs @@ -223,7 +223,7 @@ where )) } - pub(crate) fn with_screen_size_and_cursor_position( + fn with_screen_size_and_cursor_position( backend: B, screen_size: Size, cursor_pos: Position, @@ -245,6 +245,15 @@ where } } + #[cfg(test)] + pub(crate) fn with_screen_size_and_cursor_position_for_test( + backend: B, + screen_size: Size, + cursor_pos: Position, + ) -> Self { + Self::with_screen_size_and_cursor_position(backend, screen_size, cursor_pos) + } + /// Get a Frame object which provides a consistent view into the terminal state for rendering. pub fn get_frame(&mut self) -> Frame<'_> { Frame { diff --git a/codex-rs/tui/src/tui/test_support.rs b/codex-rs/tui/src/tui/test_support.rs index 5c1f30eec0..c3f9f6bc3b 100644 --- a/codex-rs/tui/src/tui/test_support.rs +++ b/codex-rs/tui/src/tui/test_support.rs @@ -11,7 +11,7 @@ use crate::custom_terminal::Terminal; pub(crate) fn make_test_tui() -> io::Result { let backend = CrosstermBackend::new(stdout()); - let terminal = Terminal::with_screen_size_and_cursor_position( + let terminal = Terminal::with_screen_size_and_cursor_position_for_test( backend, Size { width: 80, From 4c0706e24a290f89c40d99b92912367c63fe273d Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Wed, 24 Jun 2026 19:21:28 -0700 Subject: [PATCH 2/7] Allow ChatGPT-hosted MCP servers to use session auth (#29733) ## Why ChatGPT session authentication was inferred from the reserved Codex Apps server name. That couples credential routing to Codex Apps-specific behavior and prevents other MCP endpoints hosted by ChatGPT from explicitly using the current session. The opt-in also needs a clear security boundary: an arbitrary MCP configuration must not be able to redirect ChatGPT credentials to another origin. ## What changed - Add `use_chatgpt_auth` to HTTP MCP server configuration, defaulting to `false`. - Honor the setting only when the parsed server URL has the same HTTP(S) origin as the configured `chatgpt_base_url`; otherwise remove the capability before startup. - Resolve bearer tokens and static or environment-backed authorization headers before selecting authentication, with configured authorization taking precedence over ChatGPT session auth. - Enable the setting for the built-in Codex Apps and hosted plugin runtime endpoints while keeping Codex Apps caching and tool normalization scoped to the reserved server. - Persist the setting through MCP config rewrite paths and expose it in the generated config schema. - Load the current login state for `codex mcp list` so reported auth status matches runtime behavior. ## Verification Core integration coverage exercises the complete streamable HTTP MCP startup path and verifies that: - a same-origin opted-in server receives the current ChatGPT access token; - an explicitly configured authorization header takes precedence; - a different-origin server completes MCP initialization and tool listing without receiving any ChatGPT authorization header. --- codex-rs/cli/src/mcp_cmd.rs | 9 +- codex-rs/codex-mcp/src/catalog_tests.rs | 1 + codex-rs/codex-mcp/src/connection_manager.rs | 83 ++++------ .../codex-mcp/src/connection_manager_tests.rs | 4 + codex-rs/codex-mcp/src/mcp/auth.rs | 6 +- codex-rs/codex-mcp/src/mcp/mod.rs | 33 +++- codex-rs/codex-mcp/src/mcp/mod_tests.rs | 2 + codex-rs/codex-mcp/src/plugin_config_tests.rs | 2 + codex-rs/codex-mcp/src/runtime.rs | 2 + codex-rs/config/src/mcp_edit.rs | 3 + codex-rs/config/src/mcp_edit_tests.rs | 2 + codex-rs/config/src/mcp_types.rs | 11 ++ codex-rs/config/src/mcp_types_tests.rs | 1 + codex-rs/core-plugins/src/manager_tests.rs | 5 + codex-rs/core/config.schema.json | 4 + codex-rs/core/src/config/config_tests.rs | 17 ++ .../core/src/config/edit/document_helpers.rs | 3 + codex-rs/core/src/config/edit_tests.rs | 7 + codex-rs/core/src/mcp_skill_dependencies.rs | 2 + codex-rs/core/tests/suite/code_mode.rs | 1 + codex-rs/core/tests/suite/hooks_mcp.rs | 1 + .../core/tests/suite/mcp_refresh_cleanup.rs | 1 + codex-rs/core/tests/suite/rmcp_client.rs | 154 ++++++++++++++++++ codex-rs/core/tests/suite/search_tool.rs | 3 + codex-rs/core/tests/suite/sqlite_state.rs | 1 + codex-rs/core/tests/suite/token_budget.rs | 1 + codex-rs/core/tests/suite/truncation.rs | 3 + .../mcp/src/executor_plugin/provider_tests.rs | 2 + codex-rs/rmcp-client/src/rmcp_client.rs | 8 +- 29 files changed, 313 insertions(+), 59 deletions(-) diff --git a/codex-rs/cli/src/mcp_cmd.rs b/codex-rs/cli/src/mcp_cmd.rs index 7baee7d349..e874ca6ab9 100644 --- a/codex-rs/cli/src/mcp_cmd.rs +++ b/codex-rs/cli/src/mcp_cmd.rs @@ -18,6 +18,7 @@ use codex_core::config::edit::ConfigEditsBuilder; use codex_core::config::find_codex_home; use codex_core::config::load_global_mcp_servers; use codex_core_plugins::PluginsManager; +use codex_login::AuthManager; use codex_mcp::McpOAuthLoginSupport; use codex_mcp::ResolvedMcpOAuthScopes; use codex_mcp::compute_auth_statuses; @@ -344,6 +345,7 @@ async fn run_add(config_overrides: &CliConfigOverrides, add_args: AddArgs) -> Re }; let new_entry = McpServerConfig { + use_chatgpt_auth: false, transport: transport.clone(), environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, @@ -543,8 +545,11 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> let mcp_manager = McpManager::new(Arc::new(PluginsManager::new( config.codex_home.to_path_buf(), ))); + let auth_manager = + AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ true).await; + let auth = auth_manager.auth().await; let mcp_servers = mcp_manager.configured_servers(&config).await; - let effective_mcp_servers = mcp_manager.effective_servers(&config, /*auth*/ None).await; + let effective_mcp_servers = mcp_manager.effective_servers(&config, auth.as_ref()).await; let mut entries: Vec<_> = mcp_servers.iter().collect(); entries.sort_by_key(|(name, _)| *name); @@ -552,7 +557,7 @@ async fn run_list(config_overrides: &CliConfigOverrides, list_args: ListArgs) -> effective_mcp_servers.iter(), config.mcp_oauth_credentials_store_mode, config.auth_keyring_backend_kind(), - /*auth*/ None, + auth.as_ref(), ) .await; diff --git a/codex-rs/codex-mcp/src/catalog_tests.rs b/codex-rs/codex-mcp/src/catalog_tests.rs index 696876cc40..09a061441b 100644 --- a/codex-rs/codex-mcp/src/catalog_tests.rs +++ b/codex-rs/codex-mcp/src/catalog_tests.rs @@ -17,6 +17,7 @@ use super::ResolvedMcpCatalog; fn server(url: &str) -> McpServerConfig { McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: url.to_string(), bearer_token_env_var: None, diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 33631393cc..0747173470 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -154,7 +154,7 @@ impl McpConnectionManager { ); let tool_plugin_provenance = Arc::new(tool_plugin_provenance); let startup_submit_id = submit_id.clone(); - let codex_apps_auth_provider = auth + let chatgpt_auth_provider = auth .filter(|auth| auth.uses_codex_backend()) .map(codex_model_provider::auth_provider_from_auth); let mcp_servers = mcp_servers.clone(); @@ -173,17 +173,13 @@ impl McpConnectionManager { }, ) .await; - let (codex_apps_tools_cache_context, runtime_auth_provider) = - if server_name == CODEX_APPS_MCP_SERVER_NAME { - codex_apps_cache_context_and_auth_provider( - &server, - &codex_home, - &codex_apps_tools_cache_key, - codex_apps_auth_provider.clone(), - ) - } else { - regular_mcp_cache_context_and_auth_provider() - }; + let codex_apps_tools_cache_context = if server_name == CODEX_APPS_MCP_SERVER_NAME { + codex_apps_tools_cache_context(&codex_home, &codex_apps_tools_cache_key) + } else { + regular_mcp_tools_cache_context() + }; + let runtime_auth_provider = + chatgpt_auth_provider_for_server(&server, chatgpt_auth_provider.clone()); let async_managed_client = AsyncManagedClient::new( server_name.clone(), server, @@ -852,48 +848,35 @@ impl Drop for McpConnectionManager { } } -/// Creates the host-owned state used only by the Codex Apps server. -/// -/// The tools cache is scoped to the authenticated user. Runtime authentication is supplied only -/// when the server is not already configured to read a bearer token from the environment. -fn codex_apps_cache_context_and_auth_provider( - server: &EffectiveMcpServer, +/// Creates the per-user tools cache context used only by the Codex Apps server. +fn codex_apps_tools_cache_context( codex_home: &Path, codex_apps_tools_cache_key: &CodexAppsToolsCacheKey, - codex_apps_auth_provider: Option, -) -> ( - Option, - Option, -) { - let uses_env_bearer_token = - server - .configured_config() - .is_some_and(|config| match &config.transport { - McpServerTransportConfig::StreamableHttp { - bearer_token_env_var, - .. - } => bearer_token_env_var.is_some(), - McpServerTransportConfig::Stdio { .. } => false, - }); - ( - Some(CodexAppsToolsCacheContext { - codex_home: codex_home.to_path_buf(), - user_key: codex_apps_tools_cache_key.clone(), - }), - if uses_env_bearer_token { - None - } else { - codex_apps_auth_provider - }, - ) +) -> Option { + Some(CodexAppsToolsCacheContext { + codex_home: codex_home.to_path_buf(), + user_key: codex_apps_tools_cache_key.clone(), + }) } -/// Keeps regular MCP servers isolated from the host-owned Codex Apps cache and auth provider. -fn regular_mcp_cache_context_and_auth_provider() -> ( - Option, - Option, -) { - (None, None) +/// Keeps regular MCP servers isolated from the Codex Apps tools cache. +fn regular_mcp_tools_cache_context() -> Option { + None +} + +/// Makes ChatGPT authentication available to servers that explicitly opt in. +/// The HTTP transport applies it only when no configured authorization resolves. +fn chatgpt_auth_provider_for_server( + server: &EffectiveMcpServer, + chatgpt_auth_provider: Option, +) -> Option { + if !server + .configured_config() + .is_some_and(|config| config.use_chatgpt_auth) + { + return None; + } + chatgpt_auth_provider } async fn emit_update( diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index f4dfc1fd6d..46fa7fae6c 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -1272,6 +1272,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { ( "stdio".to_string(), EffectiveMcpServer::configured(McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "echo".to_string(), args: Vec::new(), @@ -1298,6 +1299,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { ( "http".to_string(), EffectiveMcpServer::configured(McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "http://127.0.0.1:1".to_string(), bearer_token_env_var: None, @@ -1405,6 +1407,7 @@ fn mcp_init_error_display_prompts_for_github_pat() { let server_name = "github"; let entry = McpAuthStatusEntry { config: Some(McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "https://api.githubcopilot.com/mcp/".to_string(), bearer_token_env_var: None, @@ -1458,6 +1461,7 @@ fn mcp_init_error_display_reports_generic_errors() { let server_name = "custom"; let entry = McpAuthStatusEntry { config: Some(McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com".to_string(), bearer_token_env_var: Some("TOKEN".to_string()), diff --git a/codex-rs/codex-mcp/src/mcp/auth.rs b/codex-rs/codex-mcp/src/mcp/auth.rs index 78d55fb789..182c6b6c3e 100644 --- a/codex-rs/codex-mcp/src/mcp/auth.rs +++ b/codex-rs/codex-mcp/src/mcp/auth.rs @@ -16,8 +16,6 @@ use tracing::warn; use crate::server::EffectiveMcpServer; -use super::CODEX_APPS_MCP_SERVER_NAME; - #[derive(Debug, Clone)] pub struct McpOAuthLoginConfig { pub url: String, @@ -141,7 +139,9 @@ where let futures = servers.into_iter().map(|(name, server)| { let name = name.clone(); let config = server.configured_config().cloned(); - let has_runtime_auth = name == CODEX_APPS_MCP_SERVER_NAME + let has_runtime_auth = config + .as_ref() + .is_some_and(|config| config.use_chatgpt_auth) && auth.is_some_and(CodexAuth::uses_codex_backend) && config.as_ref().is_some_and(|config| { matches!( diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 93ccf1a9e5..d2a6a9f2d6 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -257,9 +257,26 @@ pub fn effective_mcp_servers_from_configured( config: &McpConfig, auth: Option<&CodexAuth>, ) -> HashMap { + let chatgpt_origin = url::Url::parse(&config.chatgpt_base_url) + .ok() + .filter(|url| matches!(url.scheme(), "http" | "https")) + .map(|url| url.origin()); let mut servers = configured_servers .into_iter() - .map(|(name, server)| (name, EffectiveMcpServer::configured(server))) + .map(|(name, mut server)| { + if server.use_chatgpt_auth { + let server_origin = match &server.transport { + McpServerTransportConfig::StreamableHttp { url, .. } => url::Url::parse(url) + .ok() + .filter(|url| matches!(url.scheme(), "http" | "https")) + .map(|url| url.origin()), + McpServerTransportConfig::Stdio { .. } => None, + }; + server.use_chatgpt_auth = + server_origin.is_some() && server_origin.as_ref() == chatgpt_origin.as_ref(); + } + (name, EffectiveMcpServer::configured(server)) + }) .collect::>(); if !host_owned_codex_apps_enabled(config, auth) { servers.remove(CODEX_APPS_MCP_SERVER_NAME); @@ -457,6 +474,7 @@ pub fn codex_apps_mcp_server_config( mcp_server_config_for_url( codex_apps_mcp_url_for_base_url(chatgpt_base_url), apps_mcp_product_sku, + /*use_chatgpt_auth*/ true, ) } @@ -471,10 +489,18 @@ pub fn hosted_plugin_runtime_mcp_server_config( } else { format!("{base_url}/api/codex") }; - mcp_server_config_for_url(format!("{base_url}/ps/mcp"), apps_mcp_product_sku) + mcp_server_config_for_url( + format!("{base_url}/ps/mcp"), + apps_mcp_product_sku, + /*use_chatgpt_auth*/ true, + ) } -fn mcp_server_config_for_url(url: String, apps_mcp_product_sku: Option<&str>) -> McpServerConfig { +fn mcp_server_config_for_url( + url: String, + apps_mcp_product_sku: Option<&str>, + use_chatgpt_auth: bool, +) -> McpServerConfig { let http_headers = apps_mcp_product_sku.map(|product_sku| { HashMap::from([("X-OpenAI-Product-Sku".to_string(), product_sku.to_string())]) }); @@ -486,6 +512,7 @@ fn mcp_server_config_for_url(url: String, apps_mcp_product_sku: Option<&str>) -> http_headers, env_http_headers: None, }, + use_chatgpt_auth, environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, required: false, diff --git a/codex-rs/codex-mcp/src/mcp/mod_tests.rs b/codex-rs/codex-mcp/src/mcp/mod_tests.rs index 3a159e111d..389bf1adcb 100644 --- a/codex-rs/codex-mcp/src/mcp/mod_tests.rs +++ b/codex-rs/codex-mcp/src/mcp/mod_tests.rs @@ -296,6 +296,7 @@ async fn effective_mcp_servers_preserve_runtime_servers() { catalog.register(McpServerRegistration::from_config( "sample".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "https://user.example/mcp".to_string(), bearer_token_env_var: None, @@ -321,6 +322,7 @@ async fn effective_mcp_servers_preserve_runtime_servers() { catalog.register(McpServerRegistration::from_config( "docs".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "https://docs.example/mcp".to_string(), bearer_token_env_var: None, diff --git a/codex-rs/codex-mcp/src/plugin_config_tests.rs b/codex-rs/codex-mcp/src/plugin_config_tests.rs index d95cd5095e..fc495d7ee8 100644 --- a/codex-rs/codex-mcp/src/plugin_config_tests.rs +++ b/codex-rs/codex-mcp/src/plugin_config_tests.rs @@ -32,6 +32,7 @@ fn stdio_server( env_vars: Vec, ) -> McpServerConfig { McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: command.to_string(), args: Vec::new(), @@ -66,6 +67,7 @@ fn declared_placement_preserves_local_plugin_normalization() { Vec::new(), ); let expected_http = McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), bearer_token_env_var: None, diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index ce8d4156ab..be18660e2b 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -104,6 +104,7 @@ mod tests { fn stdio_server(environment_id: &str) -> McpServerConfig { McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "echo".to_string(), args: Vec::new(), @@ -130,6 +131,7 @@ mod tests { fn http_server(environment_id: &str) -> McpServerConfig { McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "http://127.0.0.1:1".to_string(), bearer_token_env_var: None, diff --git a/codex-rs/config/src/mcp_edit.rs b/codex-rs/config/src/mcp_edit.rs index 1d896e5f31..bc5d4d3316 100644 --- a/codex-rs/config/src/mcp_edit.rs +++ b/codex-rs/config/src/mcp_edit.rs @@ -172,6 +172,9 @@ fn serialize_mcp_server(config: &McpServerConfig) -> TomlItem { } } + if config.use_chatgpt_auth { + entry["use_chatgpt_auth"] = value(true); + } if !config.enabled { entry["enabled"] = value(false); } diff --git a/codex-rs/config/src/mcp_edit_tests.rs b/codex-rs/config/src/mcp_edit_tests.rs index 030b0a0233..3c2801b269 100644 --- a/codex-rs/config/src/mcp_edit_tests.rs +++ b/codex-rs/config/src/mcp_edit_tests.rs @@ -16,6 +16,7 @@ async fn replace_mcp_servers_serializes_per_tool_approval_overrides() -> anyhow: let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "docs-server".to_string(), args: Vec::new(), @@ -95,6 +96,7 @@ async fn replace_mcp_servers_serializes_oauth_client_id() -> anyhow::Result<()> let servers = BTreeMap::from([( "maas_outlook".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), bearer_token_env_var: None, diff --git a/codex-rs/config/src/mcp_types.rs b/codex-rs/config/src/mcp_types.rs index 42142212c5..08bf2fef2f 100644 --- a/codex-rs/config/src/mcp_types.rs +++ b/codex-rs/config/src/mcp_types.rs @@ -131,6 +131,12 @@ pub struct McpServerConfig { #[serde(flatten)] pub transport: McpServerTransportConfig, + /// When `true`, request authentication with the current ChatGPT session. Codex honors this + /// only when the server URL has the same origin as the configured ChatGPT base URL and no + /// configured bearer token or authorization header resolves. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub use_chatgpt_auth: bool, + /// Effective environment id for where Codex should start this MCP server. pub environment_id: String, @@ -239,6 +245,8 @@ pub struct RawMcpServerConfig { #[serde(default)] pub environment_id: Option, #[serde(default)] + pub use_chatgpt_auth: Option, + #[serde(default)] pub startup_timeout_sec: Option, #[serde(default)] pub startup_timeout_ms: Option, @@ -286,6 +294,7 @@ impl TryFrom for McpServerConfig { bearer_token, bearer_token_env_var, environment_id, + use_chatgpt_auth, startup_timeout_sec, startup_timeout_ms, tool_timeout_sec, @@ -329,6 +338,7 @@ impl TryFrom for McpServerConfig { throw_if_set("stdio", "env_http_headers", env_http_headers.as_ref())?; throw_if_set("stdio", "oauth", oauth.as_ref())?; throw_if_set("stdio", "oauth_resource", oauth_resource.as_ref())?; + throw_if_set("stdio", "use_chatgpt_auth", use_chatgpt_auth.as_ref())?; let env_vars = env_vars.unwrap_or_default(); for env_var in &env_vars { env_var.validate_source()?; @@ -361,6 +371,7 @@ impl TryFrom for McpServerConfig { Ok(Self { transport, + use_chatgpt_auth: use_chatgpt_auth.unwrap_or_default(), environment_id, startup_timeout_sec, tool_timeout_sec, diff --git a/codex-rs/config/src/mcp_types_tests.rs b/codex-rs/config/src/mcp_types_tests.rs index d3ec9d77e6..cab51f14c7 100644 --- a/codex-rs/config/src/mcp_types_tests.rs +++ b/codex-rs/config/src/mcp_types_tests.rs @@ -426,6 +426,7 @@ fn deserialize_ignores_unknown_server_fields() { assert_eq!( cfg, McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "echo".to_string(), args: vec![], diff --git a/codex-rs/core-plugins/src/manager_tests.rs b/codex-rs/core-plugins/src/manager_tests.rs index 72d357b90e..e4441e15d0 100644 --- a/codex-rs/core-plugins/src/manager_tests.rs +++ b/codex-rs/core-plugins/src/manager_tests.rs @@ -677,6 +677,7 @@ async fn load_plugins_loads_default_skills_and_mcp_servers() { mcp_servers: HashMap::from([( "sample".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "https://sample.example/mcp".to_string(), bearer_token_env_var: None, @@ -768,6 +769,7 @@ enabled = true HashMap::from([( "counter".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "https://sample.example/counter/mcp".to_string(), bearer_token_env_var: None, @@ -1662,6 +1664,7 @@ async fn load_plugins_uses_manifest_configured_component_paths() { HashMap::from([( "custom".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "https://custom.example/mcp".to_string(), bearer_token_env_var: None, @@ -1842,6 +1845,7 @@ async fn load_plugins_ignores_manifest_component_paths_without_dot_slash() { HashMap::from([( "default".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "https://default.example/mcp".to_string(), bearer_token_env_var: None, @@ -2092,6 +2096,7 @@ fn capability_index_filters_inactive_and_zero_capability_plugins() { let connector = |id: &str| AppConnectorId(id.to_string()); let app = |name: &str, connector_id: &str| app_declaration(name, connector_id); let http_server = |url: &str| McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: url.to_string(), bearer_token_env_var: None, diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 25ae84cfbd..f8cfa6f726 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -2555,6 +2555,10 @@ }, "url": { "type": "string" + }, + "use_chatgpt_auth": { + "default": null, + "type": "boolean" } }, "type": "object" diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 4d3a7bcc63..e0c2f16638 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -114,6 +114,7 @@ use tempfile::TempDir; fn stdio_mcp(command: &str) -> McpServerConfig { McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: command.to_string(), args: Vec::new(), @@ -140,6 +141,7 @@ fn stdio_mcp(command: &str) -> McpServerConfig { fn http_mcp(url: &str) -> McpServerConfig { McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: url.to_string(), bearer_token_env_var: None, @@ -5566,6 +5568,7 @@ async fn replace_mcp_servers_round_trips_entries() -> anyhow::Result<()> { servers.insert( "docs".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "echo".to_string(), args: vec!["hello".to_string()], @@ -5923,6 +5926,7 @@ async fn replace_mcp_servers_serializes_env_sorted() -> anyhow::Result<()> { let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "docs-server".to_string(), args: vec!["--verbose".to_string()], @@ -6002,6 +6006,7 @@ async fn replace_mcp_servers_serializes_env_vars() -> anyhow::Result<()> { let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "docs-server".to_string(), args: Vec::new(), @@ -6057,6 +6062,7 @@ async fn replace_mcp_servers_serializes_sourced_env_vars() -> anyhow::Result<()> let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "docs-server".to_string(), args: Vec::new(), @@ -6115,6 +6121,7 @@ async fn replace_mcp_servers_serializes_cwd() -> anyhow::Result<()> { let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "docs-server".to_string(), args: Vec::new(), @@ -6170,6 +6177,7 @@ async fn replace_mcp_servers_streamable_http_serializes_bearer_token() -> anyhow let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), bearer_token_env_var: Some("MCP_TOKEN".to_string()), @@ -6237,6 +6245,7 @@ async fn replace_mcp_servers_streamable_http_serializes_custom_headers() -> anyh let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), bearer_token_env_var: Some("MCP_TOKEN".to_string()), @@ -6319,6 +6328,7 @@ async fn replace_mcp_servers_streamable_http_removes_optional_sections() -> anyh let mut servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), bearer_token_env_var: Some("MCP_TOKEN".to_string()), @@ -6357,6 +6367,7 @@ async fn replace_mcp_servers_streamable_http_removes_optional_sections() -> anyh servers.insert( "docs".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), bearer_token_env_var: None, @@ -6424,6 +6435,7 @@ async fn replace_mcp_servers_streamable_http_isolates_headers_between_servers() ( "docs".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), bearer_token_env_var: Some("MCP_TOKEN".to_string()), @@ -6452,6 +6464,7 @@ async fn replace_mcp_servers_streamable_http_isolates_headers_between_servers() ( "logs".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "logs-server".to_string(), args: vec!["--follow".to_string()], @@ -6540,6 +6553,7 @@ async fn replace_mcp_servers_serializes_disabled_flag() -> anyhow::Result<()> { let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "docs-server".to_string(), args: Vec::new(), @@ -6590,6 +6604,7 @@ async fn replace_mcp_servers_serializes_required_flag() -> anyhow::Result<()> { let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "docs-server".to_string(), args: Vec::new(), @@ -6640,6 +6655,7 @@ async fn replace_mcp_servers_serializes_tool_filters() -> anyhow::Result<()> { let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "docs-server".to_string(), args: Vec::new(), @@ -6695,6 +6711,7 @@ async fn replace_mcp_servers_streamable_http_serializes_oauth_resource() -> anyh let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), bearer_token_env_var: None, diff --git a/codex-rs/core/src/config/edit/document_helpers.rs b/codex-rs/core/src/config/edit/document_helpers.rs index 3aed9a49ba..5f9445ad1c 100644 --- a/codex-rs/core/src/config/edit/document_helpers.rs +++ b/codex-rs/core/src/config/edit/document_helpers.rs @@ -95,6 +95,9 @@ fn serialize_mcp_server_table(config: &McpServerConfig) -> TomlTable { } } + if config.use_chatgpt_auth { + entry["use_chatgpt_auth"] = value(true); + } if !config.enabled { entry["enabled"] = value(false); } diff --git a/codex-rs/core/src/config/edit_tests.rs b/codex-rs/core/src/config/edit_tests.rs index dce192831b..9c6cfe50cc 100644 --- a/codex-rs/core/src/config/edit_tests.rs +++ b/codex-rs/core/src/config/edit_tests.rs @@ -873,6 +873,7 @@ fn blocking_replace_mcp_servers_round_trips() { servers.insert( "stdio".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "cmd".to_string(), args: vec!["--flag".to_string()], @@ -907,6 +908,7 @@ fn blocking_replace_mcp_servers_round_trips() { servers.insert( "http".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com".to_string(), bearer_token_env_var: Some("TOKEN".to_string()), @@ -981,6 +983,7 @@ fn blocking_replace_mcp_servers_serializes_tool_approval_overrides() { servers.insert( "docs".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "docs-server".to_string(), args: Vec::new(), @@ -1041,6 +1044,7 @@ foo = { command = "cmd" } servers.insert( "foo".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "cmd".to_string(), args: Vec::new(), @@ -1091,6 +1095,7 @@ foo = { command = "cmd" } # keep me servers.insert( "foo".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "cmd".to_string(), args: Vec::new(), @@ -1140,6 +1145,7 @@ foo = { command = "cmd", args = ["--flag"] } # keep me servers.insert( "foo".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "cmd".to_string(), args: Vec::new(), @@ -1190,6 +1196,7 @@ foo = { command = "cmd" } servers.insert( "foo".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "cmd".to_string(), args: Vec::new(), diff --git a/codex-rs/core/src/mcp_skill_dependencies.rs b/codex-rs/core/src/mcp_skill_dependencies.rs index 32e95b80ee..0574fd14f5 100644 --- a/codex-rs/core/src/mcp_skill_dependencies.rs +++ b/codex-rs/core/src/mcp_skill_dependencies.rs @@ -357,6 +357,7 @@ fn mcp_dependency_to_server_config( .as_ref() .ok_or_else(|| "missing url for streamable_http dependency".to_string())?; return Ok(McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::StreamableHttp { url: url.clone(), bearer_token_env_var: None, @@ -386,6 +387,7 @@ fn mcp_dependency_to_server_config( .as_ref() .ok_or_else(|| "missing command for stdio dependency".to_string())?; return Ok(McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: command.clone(), args: Vec::new(), diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index d7d0b9be74..4db905e6b6 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -395,6 +395,7 @@ async fn run_code_mode_turn_with_rmcp_config( servers.insert( "rmcp".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, args: Vec::new(), diff --git a/codex-rs/core/tests/suite/hooks_mcp.rs b/codex-rs/core/tests/suite/hooks_mcp.rs index db35577b0a..b3d13035cb 100644 --- a/codex-rs/core/tests/suite/hooks_mcp.rs +++ b/codex-rs/core/tests/suite/hooks_mcp.rs @@ -184,6 +184,7 @@ fn insert_rmcp_test_server(config: &mut Config, command: String, approval_mode: servers.insert( RMCP_SERVER.to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command, args: Vec::new(), diff --git a/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs b/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs index a8e06bb35b..8c523cb477 100644 --- a/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs +++ b/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs @@ -30,6 +30,7 @@ async fn refresh_shuts_down_superseded_mcp_stdio_server() -> anyhow::Result<()> servers.insert( "refresh_cleanup".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command, args: Vec::new(), diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index 6fdf130658..ab53c553f2 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -47,6 +47,7 @@ use codex_protocol::protocol::Op; use codex_protocol::user_input::UserInput; use codex_utils_cargo_bin::cargo_bin; use codex_utils_path_uri::PathUri; +use core_test_support::apps_test_server::AppsTestServer; use core_test_support::assert_regex_match; use core_test_support::is_remote_test_environment; use core_test_support::responses; @@ -267,6 +268,7 @@ fn copy_binary_to_remote_env( struct TestMcpServerOptions { environment_id: String, + use_chatgpt_auth: bool, supports_parallel_tool_calls: bool, tool_timeout_sec: Option, } @@ -275,6 +277,7 @@ impl Default for TestMcpServerOptions { fn default() -> Self { Self { environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), + use_chatgpt_auth: false, supports_parallel_tool_calls: false, tool_timeout_sec: None, } @@ -315,6 +318,7 @@ fn insert_mcp_server( server_name.to_string(), McpServerConfig { transport, + use_chatgpt_auth: options.use_chatgpt_auth, environment_id: options.environment_id, enabled: true, required: false, @@ -1227,6 +1231,7 @@ async fn stdio_mcp_parallel_tool_calls_opt_in_runs_concurrently() -> anyhow::Res stdio_transport(rmcp_test_server_bin, /*env*/ None, Vec::new()), TestMcpServerOptions { environment_id: remote_aware_environment_id(), + use_chatgpt_auth: false, supports_parallel_tool_calls: true, tool_timeout_sec: Some(Duration::from_secs(2)), }, @@ -2298,6 +2303,155 @@ async fn streamable_http_tool_call_round_trip() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_chatgpt_auth_respects_configured_authorization() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let Some(chatgpt_auth_server) = + start_streamable_http_test_server("chatgpt-auth", Some("Access Token")).await? + else { + return Ok(()); + }; + let chatgpt_auth_url = chatgpt_auth_server.url().to_string(); + let chatgpt_base_url = chatgpt_auth_url + .strip_suffix("/mcp") + .expect("test MCP URL should end in /mcp") + .to_string(); + + let chatgpt_auth_fixture = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(move |config| { + config.chatgpt_base_url = chatgpt_base_url; + insert_mcp_server( + config, + "chatgpt_auth", + McpServerTransportConfig::StreamableHttp { + url: chatgpt_auth_url, + bearer_token_env_var: None, + http_headers: None, + env_http_headers: None, + }, + TestMcpServerOptions { + environment_id: remote_aware_environment_id(), + use_chatgpt_auth: true, + ..Default::default() + }, + ); + }) + .build_with_remote_env(&server) + .await?; + wait_for_mcp_server(&chatgpt_auth_fixture.codex, "chatgpt_auth").await?; + drop(chatgpt_auth_fixture); + chatgpt_auth_server.shutdown().await; + + let Some(configured_auth_server) = + start_streamable_http_test_server("configured-auth", Some("configured-token")).await? + else { + return Ok(()); + }; + let configured_auth_url = configured_auth_server.url().to_string(); + let configured_auth_base_url = configured_auth_url + .strip_suffix("/mcp") + .expect("test MCP URL should end in /mcp") + .to_string(); + + let configured_auth_fixture = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(move |config| { + config.chatgpt_base_url = configured_auth_base_url; + insert_mcp_server( + config, + "configured_auth", + McpServerTransportConfig::StreamableHttp { + url: configured_auth_url, + bearer_token_env_var: None, + http_headers: Some(HashMap::from([( + "Authorization".to_string(), + "Bearer configured-token".to_string(), + )])), + env_http_headers: None, + }, + TestMcpServerOptions { + environment_id: remote_aware_environment_id(), + use_chatgpt_auth: true, + ..Default::default() + }, + ); + }) + .build_with_remote_env(&server) + .await?; + + wait_for_mcp_server(&configured_auth_fixture.codex, "configured_auth").await?; + drop(configured_auth_fixture); + configured_auth_server.shutdown().await; + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn streamable_http_chatgpt_auth_is_not_sent_to_another_origin() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let untrusted_server = MockServer::start().await; + let untrusted_apps = AppsTestServer::mount(&untrusted_server).await?; + let untrusted_mcp_url = format!("{}/api/codex/apps", untrusted_apps.chatgpt_base_url); + let trusted_chatgpt_base_url = server.uri(); + + let fixture = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(move |config| { + config.chatgpt_base_url = trusted_chatgpt_base_url; + insert_mcp_server( + config, + "untrusted_origin", + McpServerTransportConfig::StreamableHttp { + url: untrusted_mcp_url, + bearer_token_env_var: None, + http_headers: None, + env_http_headers: None, + }, + TestMcpServerOptions { + use_chatgpt_auth: true, + ..Default::default() + }, + ); + }) + .build(&server) + .await?; + + wait_for_mcp_server(&fixture.codex, "untrusted_origin").await?; + let observed_requests = untrusted_server + .received_requests() + .await + .expect("mock server should capture MCP startup requests") + .into_iter() + .filter(|request| request.url.path() == "/api/codex/apps") + .filter_map(|request| { + let body: Value = serde_json::from_slice(&request.body).ok()?; + let method = body.get("method")?.as_str()?.to_string(); + let authorization = request + .headers + .get("authorization") + .and_then(|value| value.to_str().ok()) + .map(str::to_string); + Some((method, authorization)) + }) + .collect::>(); + + assert_eq!( + observed_requests, + vec![ + ("initialize".to_string(), None), + ("notifications/initialized".to_string(), None), + ("tools/list".to_string(), None), + ], + ); + + Ok(()) +} + /// This test writes to a fallback credentials file in CODEX_HOME. /// Ideally, we wouldn't need to serialize the test but it's much more cumbersome to wire CODEX_HOME through the code. #[test] diff --git a/codex-rs/core/tests/suite/search_tool.rs b/codex-rs/core/tests/suite/search_tool.rs index af67827927..e0361d4770 100644 --- a/codex-rs/core/tests/suite/search_tool.rs +++ b/codex-rs/core/tests/suite/search_tool.rs @@ -1090,6 +1090,7 @@ async fn tool_search_indexes_only_enabled_non_app_mcp_tools() -> Result<()> { servers.insert( "rmcp".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, args: Vec::new(), @@ -1216,6 +1217,7 @@ async fn tool_search_surfaced_mcp_tool_errors_are_returned_to_model() -> Result< servers.insert( "rmcp".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, args: Vec::new(), @@ -1364,6 +1366,7 @@ async fn tool_search_uses_non_app_mcp_server_instructions_as_namespace_descripti servers.insert( "rmcp".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, args: Vec::new(), diff --git a/codex-rs/core/tests/suite/sqlite_state.rs b/codex-rs/core/tests/suite/sqlite_state.rs index ddfc8ef210..d9aefecf7f 100644 --- a/codex-rs/core/tests/suite/sqlite_state.rs +++ b/codex-rs/core/tests/suite/sqlite_state.rs @@ -657,6 +657,7 @@ async fn mcp_call_marks_thread_memory_mode_polluted_when_configured() -> Result< servers.insert( server_name.to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, args: Vec::new(), diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index 18d8008507..a55c182181 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -276,6 +276,7 @@ async fn token_budget_context_injects_plain_thread_hint_text() -> Result<()> { servers.insert( "notes".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, args: Vec::new(), diff --git a/codex-rs/core/tests/suite/truncation.rs b/codex-rs/core/tests/suite/truncation.rs index c7e296ec0c..7080050977 100644 --- a/codex-rs/core/tests/suite/truncation.rs +++ b/codex-rs/core/tests/suite/truncation.rs @@ -383,6 +383,7 @@ async fn mcp_tool_call_output_exceeds_limit_truncated_for_model() -> Result<()> servers.insert( server_name.to_string(), codex_config::types::McpServerConfig { + use_chatgpt_auth: false, transport: codex_config::types::McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, args: Vec::new(), @@ -481,6 +482,7 @@ async fn mcp_image_output_preserves_image_and_no_text_summary() -> Result<()> { servers.insert( server_name.to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, args: Vec::new(), @@ -773,6 +775,7 @@ async fn mcp_tool_call_output_not_truncated_with_custom_limit() -> Result<()> { servers.insert( server_name.to_string(), codex_config::types::McpServerConfig { + use_chatgpt_auth: false, transport: codex_config::types::McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, args: Vec::new(), diff --git a/codex-rs/ext/mcp/src/executor_plugin/provider_tests.rs b/codex-rs/ext/mcp/src/executor_plugin/provider_tests.rs index a007acaac7..cbc87b8182 100644 --- a/codex-rs/ext/mcp/src/executor_plugin/provider_tests.rs +++ b/codex-rs/ext/mcp/src/executor_plugin/provider_tests.rs @@ -166,6 +166,7 @@ async fn reads_declared_config_only_through_executor_file_system() { vec![( "demo".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "demo-mcp".to_string(), args: Vec::new(), @@ -221,6 +222,7 @@ async fn reads_manifest_object_config_without_executor_file_system_access() { vec![( "counter".to_string(), McpServerConfig { + use_chatgpt_auth: false, transport: McpServerTransportConfig::Stdio { command: "counter-mcp".to_string(), args: Vec::new(), diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index 721d801e11..93abf7d114 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -787,6 +787,12 @@ impl RmcpClient { } => { let default_headers = build_default_headers(http_headers.clone(), env_http_headers.clone())?; + let auth_provider = + if bearer_token.is_some() || default_headers.contains_key(AUTHORIZATION) { + None + } else { + auth_provider.clone() + }; let initial_oauth_tokens = if bearer_token.is_none() && auth_provider.is_none() @@ -861,7 +867,7 @@ impl RmcpClient { StreamableHttpClientAdapter::new( Arc::clone(http_client), default_headers, - auth_provider.clone(), + auth_provider, ), http_config, ); From 3e51b46eba036364a735f450c2bb4d7c35d48d2b Mon Sep 17 00:00:00 2001 From: sayan-oai Date: Wed, 24 Jun 2026 19:26:55 -0700 Subject: [PATCH 3/7] [1/3] core: make world state snapshots serializable (#29833) ## Why `WorldState` currently keeps its diff baseline as live Rust objects keyed by process-local `TypeId`. That baseline cannot be written to a rollout or restored after resume, so Codex reconstructs an approximation from `TurnContextItem`. This is the first change in the WorldState persistence stack. It gives every section a stable persisted identity and a compact serializable comparison snapshot without changing rollout behavior yet. ## What changed - Require each `WorldStateSection` to define a stable ID and serializable snapshot type. - Reject duplicate section IDs when constructing `WorldState`. - Persist a dedicated environment comparison snapshot using model-visible strings instead of runtime path types. - Store only `WorldStateSnapshot` in `ContextManager`, removing the parallel live-object baseline. - Render diffs by restoring each section's typed snapshot; invalid snapshots fall back to a full section render. - Omit null object fields for future RFC 7386 patches while preserving null values inside arrays. Follow-up PRs will record full snapshots and merge patches, then restore the baseline during resume, fork, and rollback. ## Test plan - WorldState snapshot tests cover stable IDs, duplicate rejection, null omission, and array preservation. - Environment tests cover persistence-safe snapshot values and existing diff rendering. - ContextManager baseline deduplication and session context-update persistence tests. Related: #29249 --- .../src/context/world_state/environment.rs | 70 ++++++++-- .../context/world_state/environment_tests.rs | 42 +++++- codex-rs/core/src/context/world_state/mod.rs | 130 +++++++++++++----- .../context/world_state/world_state_tests.rs | 120 ++++++++++++++++ codex-rs/core/src/context_manager/history.rs | 14 +- .../core/src/context_manager/history_tests.rs | 9 +- codex-rs/core/src/session/mod.rs | 19 +-- codex-rs/core/src/session/tests.rs | 13 +- 8 files changed, 347 insertions(+), 70 deletions(-) create mode 100644 codex-rs/core/src/context/world_state/world_state_tests.rs diff --git a/codex-rs/core/src/context/world_state/environment.rs b/codex-rs/core/src/context/world_state/environment.rs index d6e3284667..8056a846fe 100644 --- a/codex-rs/core/src/context/world_state/environment.rs +++ b/codex-rs/core/src/context/world_state/environment.rs @@ -10,6 +10,8 @@ use codex_protocol::protocol::TurnContextItem; use codex_protocol::protocol::TurnContextNetworkItem; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; +use serde::Deserialize; +use serde::Serialize; use std::collections::BTreeMap; /// Environment values visible to the model. @@ -91,17 +93,49 @@ impl EnvironmentsState { } impl WorldStateSection for EnvironmentsState { - fn render_diff(&self, previous: Option<&Self>) -> Option> { - let empty = Self::default(); + const ID: &'static str = "environments"; + type Snapshot = EnvironmentsSnapshot; + + fn snapshot(&self) -> Self::Snapshot { + EnvironmentsSnapshot { + environments: self + .environments + .iter() + .map(|(id, environment)| { + ( + id.clone(), + EnvironmentSnapshot { + cwd: environment.cwd.inferred_native_path_string(), + status: environment.status, + shell: environment.shell.clone(), + }, + ) + }) + .collect(), + current_date: self.current_date.clone(), + timezone: self.timezone.clone(), + network: self.network.as_ref().map(NetworkContext::render), + filesystem: self.filesystem.as_ref().map(FileSystemContext::render), + subagents: self.subagents.clone(), + } + } + + fn render_diff( + &self, + previous: Option<&Self::Snapshot>, + ) -> Option> { + let current = self.snapshot(); + let empty = EnvironmentsSnapshot::default(); let previous = previous.unwrap_or(&empty); - let turn_context_values_changed = self.current_date != previous.current_date - || self.timezone != previous.timezone - || self.network != previous.network - || self.filesystem != previous.filesystem; + let turn_context_values_changed = current.current_date != previous.current_date + || current.timezone != previous.timezone + || current.network != previous.network + || current.filesystem != previous.filesystem; let mut updates = self .environments .iter() - .filter(|(id, environment)| { + .filter(|(id, _)| { + let environment = ¤t.environments[*id]; previous .environments .get(*id) @@ -269,7 +303,24 @@ struct EnvironmentState { shell: Option, } -impl EnvironmentState { +#[derive(Default, Deserialize, Serialize)] +pub(crate) struct EnvironmentsSnapshot { + environments: BTreeMap, + current_date: Option, + timezone: Option, + network: Option, + filesystem: Option, + subagents: Option, +} + +#[derive(Deserialize, Serialize)] +struct EnvironmentSnapshot { + cwd: String, + status: EnvironmentStatus, + shell: Option, +} + +impl EnvironmentSnapshot { fn has_same_diff_value(&self, other: &Self) -> bool { self.cwd == other.cwd && self.status == other.status @@ -281,7 +332,8 @@ impl EnvironmentState { } } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] +#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] enum EnvironmentStatus { Starting, Available, diff --git a/codex-rs/core/src/context/world_state/environment_tests.rs b/codex-rs/core/src/context/world_state/environment_tests.rs index 961ee5ed73..e96a6c3091 100644 --- a/codex-rs/core/src/context/world_state/environment_tests.rs +++ b/codex-rs/core/src/context/world_state/environment_tests.rs @@ -8,6 +8,7 @@ use codex_protocol::models::PermissionProfile; use codex_protocol::models::ResponseItem; use codex_protocol::permissions::NetworkSandboxPolicy; use pretty_assertions::assert_eq; +use serde_json::json; #[test] fn renders_full_environment_state() -> Result<()> { @@ -95,7 +96,7 @@ fn renders_only_changed_environments() -> Result<()> { "#, )], - render_fragments(current.render_diff(&previous)), + render_fragments(current.render_diff(&previous.snapshot())), ); Ok(()) } @@ -151,7 +152,42 @@ fn persisted_turn_context_values_render_a_diff() -> Result<()> { "#, )], - render_fragments(current.render_diff(&previous)), + render_fragments(current.render_diff(&previous.snapshot())), + ); + Ok(()) +} + +#[test] +fn persisted_snapshot_uses_model_visible_path_and_context_values() -> Result<()> { + let mut world_state = WorldState::default(); + world_state.add_section(EnvironmentsState { + environments: [( + "remote".to_string(), + available("file:///C:/windows", "powershell")?, + )] + .into_iter() + .collect(), + filesystem: Some(FileSystemContext::from_permission_profile( + &PermissionProfile::Disabled, + &[], + )), + ..Default::default() + }); + + assert_eq!( + serde_json::to_value(world_state.snapshot())?, + json!({ + "environments": { + "environments": { + "remote": { + "cwd": "C:\\windows", + "status": "available", + "shell": "powershell" + } + }, + "filesystem": "" + } + }), ); Ok(()) } @@ -180,6 +216,7 @@ fn single_environment_diff_ignores_unknown_shell() -> Result<()> { .collect(), ..Default::default() }; + let previous = WorldStateSection::snapshot(&previous); assert_eq!( None, @@ -199,6 +236,7 @@ fn removed_legacy_environment_renders_unavailable() -> Result<()> { .collect(), ..Default::default() }; + let previous = WorldStateSection::snapshot(&previous); assert_eq!( Some(user_message( diff --git a/codex-rs/core/src/context/world_state/mod.rs b/codex-rs/core/src/context/world_state/mod.rs index 06f750ef86..e290e9fadd 100644 --- a/codex-rs/core/src/context/world_state/mod.rs +++ b/codex-rs/core/src/context/world_state/mod.rs @@ -2,50 +2,82 @@ mod environment; use crate::context::ContextualUserFragment; use indexmap::IndexMap; -use std::any::Any; -use std::any::TypeId; +use serde::Serialize; +use serde::de::DeserializeOwned; +use serde_json::Value; +use std::collections::BTreeMap; use std::fmt; pub(crate) use environment::EnvironmentsState; trait ErasedWorldStateSection: Send + Sync { - fn as_any(&self) -> &dyn Any; + fn snapshot(&self) -> Option; - fn render_diff(&self, previous: Option<&dyn Any>) -> Option>; + fn render_diff(&self, previous: Option<&Value>) -> Option>; } impl ErasedWorldStateSection for S { - fn as_any(&self) -> &dyn Any { - self + fn snapshot(&self) -> Option { + let mut snapshot = match serde_json::to_value(WorldStateSection::snapshot(self)) { + Ok(snapshot) => snapshot, + Err(err) => { + tracing::error!( + section_id = S::ID, + %err, + "failed to serialize world-state section snapshot" + ); + return None; + } + }; + remove_null_object_fields(&mut snapshot); + Some(snapshot) } - fn render_diff(&self, previous: Option<&dyn Any>) -> Option> { - let previous = match previous { - Some(previous) => { - let Some(previous) = previous.downcast_ref::() else { - unreachable!("world-state section type must match its type ID"); - }; - Some(previous) - } - None => None, - }; - WorldStateSection::render_diff(self, previous) + fn render_diff(&self, previous: Option<&Value>) -> Option> { + let previous = previous.and_then(|previous| { + serde_json::from_value::(previous.clone()) + .inspect_err(|err| { + tracing::warn!( + section_id = S::ID, + %err, + "failed to restore world-state section snapshot" + ); + }) + .ok() + }); + WorldStateSection::render_diff(self, previous.as_ref()) } } /// A typed portion of the state visible to the model. /// /// Implementations own how their current state is rendered relative to an -/// earlier value of the same section type. A missing previous value requests -/// the section's complete current representation. -pub(crate) trait WorldStateSection: Any + Send + Sync { - fn render_diff(&self, previous: Option<&Self>) -> Option>; +/// earlier snapshot of the same section. `ID` is persisted in rollouts and +/// must remain stable. `Snapshot` should contain only the comparison data +/// needed to decide what the model must be told next. +pub(crate) trait WorldStateSection: Send + Sync + 'static { + const ID: &'static str; + type Snapshot: DeserializeOwned + Serialize; + + fn snapshot(&self) -> Self::Snapshot; + + fn render_diff( + &self, + previous: Option<&Self::Snapshot>, + ) -> Option>; } -/// A snapshot of the model-visible world with one section per concrete type. +/// Live model-visible state, keyed by the same stable section IDs used in rollouts. #[derive(Default)] pub(crate) struct WorldState { - sections: IndexMap>, + sections: IndexMap<&'static str, Box>, +} + +/// Compact comparison state for each model-visible world-state section. +#[derive(Clone, Debug, Default, PartialEq, Serialize, serde::Deserialize)] +#[serde(transparent)] +pub(crate) struct WorldStateSnapshot { + sections: BTreeMap, } impl fmt::Debug for WorldState { @@ -58,23 +90,55 @@ impl fmt::Debug for WorldState { impl WorldState { pub(crate) fn add_section(&mut self, section: S) { - self.sections.insert(TypeId::of::(), Box::new(section)); + let id = S::ID; + assert!( + !self.sections.contains_key(id), + "duplicate world-state section ID: {id}" + ); + self.sections.insert(id, Box::new(section)); + } + + pub(crate) fn snapshot(&self) -> WorldStateSnapshot { + WorldStateSnapshot { + sections: self + .sections + .iter() + .filter_map(|(id, section)| { + section + .snapshot() + .map(|snapshot| ((*id).to_string(), snapshot)) + }) + .collect(), + } } pub(crate) fn render_full(&self) -> Vec> { - self.render_diff(&Self::default()) + self.render_diff(&WorldStateSnapshot::default()) } - pub(crate) fn render_diff(&self, previous: &Self) -> Vec> { + pub(crate) fn render_diff( + &self, + previous: &WorldStateSnapshot, + ) -> Vec> { self.sections .iter() - .filter_map(|(type_id, section)| { - let previous = previous - .sections - .get(type_id) - .map(|section| section.as_any()); - section.render_diff(previous) - }) + .filter_map(|(id, section)| section.render_diff(previous.sections.get(*id))) .collect() } } + +fn remove_null_object_fields(value: &mut Value) { + // RFC 7386 reserves object-valued nulls for deletion, but arrays are replaced whole. + match value { + Value::Object(values) => { + values.retain(|_, value| !value.is_null()); + values.values_mut().for_each(remove_null_object_fields); + } + Value::Array(_) => {} + Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {} + } +} + +#[cfg(test)] +#[path = "world_state_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/context/world_state/world_state_tests.rs b/codex-rs/core/src/context/world_state/world_state_tests.rs new file mode 100644 index 0000000000..77736b6c26 --- /dev/null +++ b/codex-rs/core/src/context/world_state/world_state_tests.rs @@ -0,0 +1,120 @@ +use super::*; +use pretty_assertions::assert_eq; +use serde::Deserialize; +use serde::Serialize; +use serde_json::json; + +#[derive(Clone, Deserialize, Serialize)] +struct TestSection { + value: String, + optional: Option, + array: Vec, +} + +impl WorldStateSection for TestSection { + const ID: &'static str = "test"; + type Snapshot = Self; + + fn snapshot(&self) -> Self::Snapshot { + self.clone() + } + + fn render_diff( + &self, + previous: Option<&Self::Snapshot>, + ) -> Option> { + let previous = previous?; + (self.value != previous.value) + .then(|| Box::new(TestFragment(self.value.clone())) as Box) + } +} + +struct TestFragment(String); + +impl ContextualUserFragment for TestFragment { + fn role(&self) -> &'static str { + "user" + } + + fn markers(&self) -> (&'static str, &'static str) { + Self::type_markers() + } + + fn type_markers() -> (&'static str, &'static str) { + ("", "") + } + + fn body(&self) -> String { + self.0.clone() + } +} + +struct DuplicateTestSection; + +impl WorldStateSection for DuplicateTestSection { + const ID: &'static str = "test"; + type Snapshot = (); + + fn snapshot(&self) -> Self::Snapshot {} + + fn render_diff( + &self, + _previous: Option<&Self::Snapshot>, + ) -> Option> { + None + } +} + +#[test] +fn snapshot_uses_stable_section_ids_and_omits_null_fields() { + let mut world_state = WorldState::default(); + world_state.add_section(TestSection { + value: "current".to_string(), + optional: None, + array: vec![json!({"value": null})], + }); + + assert_eq!( + serde_json::to_value(world_state.snapshot()).expect("serialize world-state snapshot"), + json!({"test": {"value": "current", "array": [{"value": null}]}}) + ); +} + +#[test] +fn render_diff_restores_the_typed_section_snapshot() { + let mut previous = WorldState::default(); + previous.add_section(TestSection { + value: "before".to_string(), + optional: None, + array: Vec::new(), + }); + let mut current = WorldState::default(); + current.add_section(TestSection { + value: "after".to_string(), + optional: None, + array: Vec::new(), + }); + + let rendered = current.render_diff(&previous.snapshot()); + + assert_eq!( + vec!["after"], + rendered + .into_iter() + .map(|fragment| fragment.body()) + .collect::>() + ); +} + +#[test] +#[should_panic(expected = "duplicate world-state section ID: test")] +fn duplicate_section_ids_are_rejected() { + let mut world_state = WorldState::default(); + world_state.add_section(TestSection { + value: "current".to_string(), + optional: None, + array: Vec::new(), + }); + + world_state.add_section(DuplicateTestSection); +} diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index daa9e3cd40..0fe72f652e 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -1,5 +1,6 @@ use crate::context::ContextualUserFragment; use crate::context::world_state::WorldState; +use crate::context::world_state::WorldStateSnapshot; use crate::context_manager::normalize; use crate::event_mapping::has_non_contextual_dev_message_content; use crate::event_mapping::is_contextual_dev_message_content; @@ -29,7 +30,6 @@ use codex_utils_output_truncation::truncate_function_output_items_with_policy; use codex_utils_output_truncation::truncate_text; use std::num::NonZeroUsize; use std::ops::Deref; -use std::sync::Arc; use std::sync::LazyLock; /// Transcript of thread history @@ -52,7 +52,7 @@ pub(crate) struct ContextManager { /// whose non-diff fragments no longer exist in the surviving history. reference_context_item: Option, /// World state most recently appended to model-visible history. - world_state_baseline: Option>, + world_state_baseline: Option, } impl ContextManager { @@ -86,18 +86,18 @@ impl ContextManager { pub(crate) fn update_world_state( &mut self, - world_state: Arc, + world_state: &WorldState, ) -> Vec> { - let fragments = self.world_state_baseline.as_deref().map_or_else( + let fragments = self.world_state_baseline.as_ref().map_or_else( || world_state.render_full(), |previous| world_state.render_diff(previous), ); - self.world_state_baseline = Some(world_state); + self.world_state_baseline = Some(world_state.snapshot()); fragments } - pub(crate) fn set_world_state_baseline(&mut self, world_state: Arc) { - self.world_state_baseline = Some(world_state); + pub(crate) fn set_world_state_baseline(&mut self, snapshot: WorldStateSnapshot) { + self.world_state_baseline = Some(snapshot); } pub(crate) fn set_token_usage_full(&mut self, context_window: i64) { diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index 7e7f7c3773..937003efa9 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -32,7 +32,6 @@ use image::Luma; use image::Rgba; use pretty_assertions::assert_eq; use regex_lite::Regex; -use std::sync::Arc; const EXEC_FORMAT_MAX_BYTES: usize = 10_000; const EXEC_FORMAT_MAX_TOKENS: usize = 2_500; @@ -83,16 +82,16 @@ fn world_state_baseline_deduplicates_until_history_is_replaced() { state.add_section(EnvironmentsState::from_turn_context_item( &reference_context_item(), )); - Arc::new(state) + state }; let mut history = ContextManager::new(); - assert_eq!(1, history.update_world_state(world_state()).len()); - assert!(history.update_world_state(world_state()).is_empty()); + assert_eq!(1, history.update_world_state(&world_state()).len()); + assert!(history.update_world_state(&world_state()).is_empty()); history.replace(Vec::new()); - assert_eq!(1, history.update_world_state(world_state()).len()); + assert_eq!(1, history.update_world_state(&world_state()).len()); } fn user_msg(text: &str) -> ResponseItem { diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 8a9055aa22..38a8939b72 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1402,14 +1402,13 @@ impl Session { prepare_response_items(&mut history); let world_state_baseline = reference_context_item .as_ref() - .map(build_world_state_from_turn_context_item); + .map(build_world_state_from_turn_context_item) + .map(|world_state| world_state.snapshot()); { let mut state = self.state.lock().await; state.replace_history(history, reference_context_item); if let Some(world_state) = world_state_baseline { - state - .history - .set_world_state_baseline(Arc::new(world_state)); + state.history.set_world_state_baseline(world_state); } let fallback_ids = state.auto_compact_window_ids(); let window_id = window_id.unwrap_or(fallback_ids.window_id); @@ -2796,7 +2795,7 @@ impl Session { .await, ); let items = crate::context_manager::updates::merge_contextual_fragments( - world_state.render_diff(previous_world_state.as_ref()), + world_state.render_diff(&previous_world_state.snapshot()), ); if !items.is_empty() { self.record_conversation_items(turn_context, &items).await; @@ -2807,7 +2806,7 @@ impl Session { .lock() .await .history - .set_world_state_baseline(Arc::clone(&world_state)); + .set_world_state_baseline(world_state.snapshot()); world_state } @@ -2949,7 +2948,9 @@ impl Session { let mut state = self.state.lock().await; state.replace_history(items, reference_context_item.clone()); if let Some(world_state) = world_state_baseline { - state.history.set_world_state_baseline(world_state); + state + .history + .set_world_state_baseline(world_state.snapshot()); } } @@ -3532,7 +3533,7 @@ impl Session { .lock() .await .history - .set_world_state_baseline(Arc::clone(&world_state)); + .set_world_state_baseline(world_state.snapshot()); context_items } else { // Steady-state path: append only built-in context diffs here; turn-scoped extension @@ -3543,7 +3544,7 @@ impl Session { let world_state_items = { let mut state = self.state.lock().await; crate::context_manager::updates::merge_contextual_fragments( - state.history.update_world_state(Arc::clone(&world_state)), + state.history.update_world_state(world_state.as_ref()), ) }; context_items.extend(world_state_items); diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 983ed20094..da3b9e2954 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -8099,11 +8099,13 @@ async fn record_context_updates_includes_turn_context_fragments_on_steady_state_ }); let mut previous_context_item = turn_context.to_turn_context_item(); previous_context_item.turn_id = Some("previous-turn-id".to_string()); - let world_state = Arc::new(build_world_state_from_turn_context(&session, &turn_context).await); + let world_state = build_world_state_from_turn_context(&session, &turn_context).await; { let mut state = session.state.lock().await; state.set_reference_context_item(Some(previous_context_item)); - state.history.set_world_state_baseline(world_state); + state + .history + .set_world_state_baseline(world_state.snapshot()); } session @@ -8773,12 +8775,13 @@ async fn record_context_updates_and_set_reference_context_item_persists_baseline .with_model(next_model.to_string(), &session.services.models_manager) .await; let previous_context_item = previous_context.to_turn_context_item(); - let world_state = - Arc::new(build_world_state_from_turn_context(&session, &previous_context).await); + let world_state = build_world_state_from_turn_context(&session, &previous_context).await; { let mut state = session.state.lock().await; state.set_reference_context_item(Some(previous_context_item.clone())); - state.history.set_world_state_baseline(world_state); + state + .history + .set_world_state_baseline(world_state.snapshot()); } let rollout_path = attach_thread_persistence(&mut session).await; From 6801941cfea0a09721041daa24f47f70a5fc4aae Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Wed, 24 Jun 2026 19:50:50 -0700 Subject: [PATCH 4/7] TUI support for buffer experience (#29919) --- codex-rs/tui/src/app/event_dispatch.rs | 73 +++++++ codex-rs/tui/src/app/thread_routing.rs | 37 +++- codex-rs/tui/src/app_event.rs | 8 + codex-rs/tui/src/chatwidget.rs | 3 + codex-rs/tui/src/chatwidget/constructor.rs | 1 + codex-rs/tui/src/chatwidget/protocol.rs | 4 +- codex-rs/tui/src/chatwidget/replay.rs | 1 + .../tui/src/chatwidget/safety_buffering.rs | 196 +++++++++++++++++ codex-rs/tui/src/chatwidget/session_flow.rs | 1 + ..._tests__safety_buffering_retry_prompt.snap | 13 ++ ...safety_buffering_status_without_retry.snap | 11 + codex-rs/tui/src/chatwidget/streaming.rs | 6 + codex-rs/tui/src/chatwidget/tests.rs | 1 + .../tui/src/chatwidget/tests/app_server.rs | 206 ++++++++++++++++++ codex-rs/tui/src/chatwidget/turn_runtime.rs | 17 ++ codex-rs/tui/src/history_cell/notices.rs | 31 +++ ...s__safety_access_block_event_snapshot.snap | 7 + codex-rs/tui/src/history_cell/tests.rs | 7 + 18 files changed, 620 insertions(+), 3 deletions(-) create mode 100644 codex-rs/tui/src/chatwidget/safety_buffering.rs create mode 100644 codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__safety_buffering_retry_prompt.snap create mode 100644 codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__safety_buffering_status_without_retry.snap create mode 100644 codex-rs/tui/src/history_cell/snapshots/codex_tui__history_cell__tests__safety_access_block_event_snapshot.snap diff --git a/codex-rs/tui/src/app/event_dispatch.rs b/codex-rs/tui/src/app/event_dispatch.rs index 280c556af5..afe269195a 100644 --- a/codex-rs/tui/src/app/event_dispatch.rs +++ b/codex-rs/tui/src/app/event_dispatch.rs @@ -337,6 +337,79 @@ impl App { self.chat_widget.prepare_local_op_submission(&op); self.submit_active_thread_op(app_server, op).await?; } + AppEvent::RetrySafetyBufferedTurn { + thread_id, + turn_id, + model, + mut turn, + } => { + if self.active_thread_id != Some(thread_id) + || self.chat_widget.thread_id() != Some(thread_id) + { + return Ok(AppRunControl::Continue); + } + if !self.chat_widget.can_retry_safety_buffered_turn(&turn_id) { + self.app_event_tx.send(AppEvent::UpdateModel(model)); + self.app_event_tx.send(AppEvent::UpdateReasoningEffort(Some( + ReasoningEffortConfig::Low, + ))); + return Ok(AppRunControl::Continue); + } + + let AppCommand::UserTurn { + model: turn_model, + effort, + collaboration_mode, + .. + } = &mut turn + else { + self.chat_widget.add_error_message( + "Failed to retry with a faster model: original turn is unavailable." + .to_string(), + ); + return Ok(AppRunControl::Continue); + }; + *turn_model = model.clone(); + *effort = Some(ReasoningEffortConfig::Low); + *collaboration_mode = collaboration_mode.as_ref().map(|mode| { + mode.with_updates( + Some(model), + Some(Some(ReasoningEffortConfig::Low)), + /*developer_instructions*/ None, + ) + }); + + if let Err(err) = app_server.turn_interrupt(thread_id, turn_id).await { + self.chat_widget + .add_error_message(format!("Failed to retry with a faster model: {err}")); + return Ok(AppRunControl::Continue); + } + let rollback_response = + match app_server.thread_rollback(thread_id, /*num_turns*/ 1).await { + Ok(response) => response, + Err(err) => { + self.chat_widget.add_error_message(format!( + "Failed to retry with a faster model: {err}" + )); + return Ok(AppRunControl::Continue); + } + }; + + self.chat_widget.prepare_safety_buffering_retry(); + self.handle_thread_rollback_response_with_origin( + thread_id, + /*num_turns*/ 1, + &rollback_response, + super::thread_routing::ThreadRollbackOrigin::SafetyBufferingRetry, + ) + .await; + + if let Err(err) = self.submit_thread_op(app_server, thread_id, turn).await { + self.chat_widget.fail_safety_buffering_retry(); + self.chat_widget + .add_error_message(format!("Failed to retry with a faster model: {err}")); + } + } AppEvent::RestoreCancelledTurn(prompt) => { self.apply_cancelled_turn_edit(prompt); } diff --git a/codex-rs/tui/src/app/thread_routing.rs b/codex-rs/tui/src/app/thread_routing.rs index 76cf56ef06..b149fd53f3 100644 --- a/codex-rs/tui/src/app/thread_routing.rs +++ b/codex-rs/tui/src/app/thread_routing.rs @@ -7,6 +7,12 @@ use super::*; use crate::session_resume::read_session_model; +#[derive(Clone, Copy)] +pub(super) enum ThreadRollbackOrigin { + Backtrack, + SafetyBufferingRetry, +} + impl App { pub(super) async fn shutdown_current_thread(&mut self, app_server: &mut AppServerSession) { if let Some(thread_id) = self.chat_widget.thread_id() { @@ -643,7 +649,7 @@ impl App { .as_ref() .map(|profile| &profile.permission_profile), ); - app_server + let response = app_server .turn_start( thread_id, items.to_vec(), @@ -661,6 +667,12 @@ impl App { final_output_json_schema.clone(), ) .await?; + if self.active_thread_id == Some(thread_id) + && self.chat_widget.thread_id() == Some(thread_id) + { + self.chat_widget + .record_safety_buffering_turn(response.turn.id, op); + } } Ok(true) } @@ -1395,6 +1407,22 @@ impl App { thread_id: ThreadId, num_turns: u32, response: &ThreadRollbackResponse, + ) { + self.handle_thread_rollback_response_with_origin( + thread_id, + num_turns, + response, + ThreadRollbackOrigin::Backtrack, + ) + .await; + } + + pub(super) async fn handle_thread_rollback_response_with_origin( + &mut self, + thread_id: ThreadId, + num_turns: u32, + response: &ThreadRollbackResponse, + origin: ThreadRollbackOrigin, ) { if let Some(channel) = self.thread_event_channels.get(&thread_id) { let mut store = channel.store.lock().await; @@ -1421,7 +1449,12 @@ impl App { self.clear_active_thread().await; } } - self.handle_backtrack_rollback_succeeded(num_turns); + match origin { + ThreadRollbackOrigin::Backtrack => self.handle_backtrack_rollback_succeeded(num_turns), + ThreadRollbackOrigin::SafetyBufferingRetry => { + self.apply_non_pending_thread_rollback(num_turns); + } + } } pub(super) fn handle_thread_event_now(&mut self, event: ThreadBufferedEvent) { diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index 1c5c949ebc..23e42ce8fe 100644 --- a/codex-rs/tui/src/app_event.rs +++ b/codex-rs/tui/src/app_event.rs @@ -161,6 +161,14 @@ pub(crate) enum AppEvent { op: AppCommand, }, + /// Interrupt, roll back, and retry a safety-buffered turn with the server-selected model. + RetrySafetyBufferedTurn { + thread_id: ThreadId, + turn_id: String, + model: String, + turn: AppCommand, + }, + /// Deliver a synthetic history lookup response to a specific thread channel. ThreadHistoryEntryResponse { thread_id: ThreadId, diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 75d319bae6..42903d2e7f 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -395,10 +395,12 @@ mod review_popups; use self::review::ReviewState; #[cfg(test)] pub(crate) use self::review_popups::show_review_commit_picker_with_entries; +mod safety_buffering; mod service_tiers; mod settings; mod settings_popups; mod side; +use self::safety_buffering::SafetyBufferingState; mod status_state; mod windows_sandbox_prompts; use self::status_state::StatusIndicatorState; @@ -584,6 +586,7 @@ pub(crate) struct ChatWidget { last_unified_wait: Option, unified_exec_wait_streak: Option, turn_lifecycle: TurnLifecycleState, + safety_buffering: SafetyBufferingState, task_complete_pending: bool, unified_exec_processes: Vec, /// Tracks per-server MCP startup state while startup is in progress. diff --git a/codex-rs/tui/src/chatwidget/constructor.rs b/codex-rs/tui/src/chatwidget/constructor.rs index 414a693bde..7fb6fe54e9 100644 --- a/codex-rs/tui/src/chatwidget/constructor.rs +++ b/codex-rs/tui/src/chatwidget/constructor.rs @@ -154,6 +154,7 @@ impl ChatWidget { last_unified_wait: None, unified_exec_wait_streak: None, turn_lifecycle: TurnLifecycleState::new(prevent_idle_sleep), + safety_buffering: SafetyBufferingState::default(), task_complete_pending: false, unified_exec_processes: Vec::new(), mcp_startup_status: None, diff --git a/codex-rs/tui/src/chatwidget/protocol.rs b/codex-rs/tui/src/chatwidget/protocol.rs index 085255a1cf..0c5b0ed3ee 100644 --- a/codex-rs/tui/src/chatwidget/protocol.rs +++ b/codex-rs/tui/src/chatwidget/protocol.rs @@ -147,6 +147,9 @@ impl ChatWidget { ServerNotification::ModelVerification(notification) => { self.on_app_server_model_verification(¬ification.verifications) } + ServerNotification::ModelSafetyBufferingUpdated(notification) => { + self.on_model_safety_buffering_updated(notification, replay_kind) + } ServerNotification::Warning(notification) => self.on_warning(notification.message), ServerNotification::GuardianWarning(notification) => { self.on_warning(notification.message) @@ -208,7 +211,6 @@ impl ChatWidget { | ServerNotification::ExternalAgentConfigImportProgress(_) | ServerNotification::ExternalAgentConfigImportCompleted(_) | ServerNotification::FsChanged(_) - | ServerNotification::ModelSafetyBufferingUpdated(_) | ServerNotification::TurnModerationMetadata(_) | ServerNotification::FuzzyFileSearchSessionUpdated(_) | ServerNotification::FuzzyFileSearchSessionCompleted(_) diff --git a/codex-rs/tui/src/chatwidget/replay.rs b/codex-rs/tui/src/chatwidget/replay.rs index a96ab614ca..3ab5b3486b 100644 --- a/codex-rs/tui/src/chatwidget/replay.rs +++ b/codex-rs/tui/src/chatwidget/replay.rs @@ -24,6 +24,7 @@ impl ChatWidget { duration_ms, } = turn; if matches!(status, TurnStatus::InProgress) { + self.turn_lifecycle.last_turn_id = Some(turn_id.clone()); self.last_non_retry_error = None; self.on_task_started(); } diff --git a/codex-rs/tui/src/chatwidget/safety_buffering.rs b/codex-rs/tui/src/chatwidget/safety_buffering.rs new file mode 100644 index 0000000000..389048fb0b --- /dev/null +++ b/codex-rs/tui/src/chatwidget/safety_buffering.rs @@ -0,0 +1,196 @@ +//! Safety-buffering status and retry UI for active turns. + +use super::*; +use codex_app_server_protocol::ModelSafetyBufferingUpdatedNotification; + +const SAFETY_BUFFERING_PROMPT_VIEW_ID: &str = "safety-buffering-prompt"; + +const SAFETY_BUFFERING_MESSAGE_WITH_RETRY: &str = "This request requires additional safety checks, which can take extra time. You can retry with a faster model for a quicker response, though it may be less capable of handling complex requests."; +const SAFETY_BUFFERING_MESSAGE_WITHOUT_RETRY: &str = + "This request requires additional safety checks, which can take extra time."; + +#[derive(Debug)] +struct ActiveSafetyBuffering { + turn_id: String, + retry_prompt_shown: bool, + agent_message_started: bool, +} + +#[derive(Debug, Default)] +pub(super) struct SafetyBufferingState { + submitted_turn: Option<(String, AppCommand)>, + active: Option, +} + +impl ChatWidget { + pub(crate) fn record_safety_buffering_turn(&mut self, turn_id: String, turn: &AppCommand) { + self.safety_buffering.submitted_turn = Some((turn_id, turn.clone())); + } + + pub(super) fn reset_safety_buffering_for_turn_start(&mut self) { + self.bottom_pane + .dismiss_view_by_id(SAFETY_BUFFERING_PROMPT_VIEW_ID); + self.safety_buffering.active = None; + } + + pub(crate) fn clear_safety_buffering(&mut self) { + self.bottom_pane + .dismiss_view_by_id(SAFETY_BUFFERING_PROMPT_VIEW_ID); + self.safety_buffering = SafetyBufferingState::default(); + } + + pub(super) fn mark_safety_buffering_agent_message_started(&mut self) { + if let Some(active) = self.safety_buffering.active.as_mut() { + active.agent_message_started = true; + } + } + + pub(super) fn safety_buffering_is_waiting(&self) -> bool { + self.safety_buffering + .active + .as_ref() + .is_some_and(|active| !active.agent_message_started) + } + + pub(crate) fn can_retry_safety_buffered_turn(&self, turn_id: &str) -> bool { + self.turn_lifecycle.agent_turn_running + && self + .safety_buffering + .active + .as_ref() + .is_some_and(|active| active.turn_id == turn_id && !active.agent_message_started) + } + + pub(crate) fn prepare_safety_buffering_retry(&mut self) { + let cancel_edit = std::mem::take(&mut self.cancel_edit); + self.last_rendered_user_message_display = None; + self.finalize_turn(); + self.cancel_edit = cancel_edit; + self.input_queue.user_turn_pending_start = true; + } + + pub(crate) fn fail_safety_buffering_retry(&mut self) { + self.input_queue.user_turn_pending_start = false; + self.clear_safety_buffering(); + let prompt = self.cancel_edit.prompt.take(); + self.clear_cancel_edit(); + if let Some(prompt) = prompt { + self.restore_user_message_to_composer(prompt); + } + } + + pub(super) fn on_model_safety_buffering_updated( + &mut self, + notification: ModelSafetyBufferingUpdatedNotification, + replay_kind: Option, + ) { + let ModelSafetyBufferingUpdatedNotification { + turn_id, + show_buffering_ui, + faster_model, + .. + } = notification; + if matches!(replay_kind, Some(ReplayKind::ResumeInitialMessages)) + || !self.turn_lifecycle.agent_turn_running + || self.turn_lifecycle.last_turn_id.as_deref() != Some(turn_id.as_str()) + { + return; + } + if !show_buffering_ui { + if self + .safety_buffering + .active + .as_ref() + .is_some_and(|active| active.turn_id == turn_id) + { + self.bottom_pane + .dismiss_view_by_id(SAFETY_BUFFERING_PROMPT_VIEW_ID); + self.safety_buffering.active = None; + self.restore_reasoning_status_header(); + } + return; + } + + let retry_turn = self + .safety_buffering + .submitted_turn + .as_ref() + .filter(|(submitted_turn_id, _)| replay_kind.is_none() && submitted_turn_id == &turn_id) + .map(|(_, turn)| turn.clone()); + let thread_id = self.thread_id; + let can_offer_retry = faster_model.is_some() && retry_turn.is_some() && thread_id.is_some(); + if !can_offer_retry { + self.bottom_pane + .dismiss_view_by_id(SAFETY_BUFFERING_PROMPT_VIEW_ID); + } + let previous_active = self + .safety_buffering + .active + .as_ref() + .filter(|active| active.turn_id == turn_id); + let retry_prompt_shown = previous_active.is_some_and(|active| active.retry_prompt_shown); + let should_show_retry_prompt = can_offer_retry && !retry_prompt_shown; + let agent_message_started = + previous_active.is_some_and(|active| active.agent_message_started); + self.safety_buffering.active = Some(ActiveSafetyBuffering { + turn_id: turn_id.clone(), + retry_prompt_shown: retry_prompt_shown || should_show_retry_prompt, + agent_message_started, + }); + + let message = if can_offer_retry { + SAFETY_BUFFERING_MESSAGE_WITH_RETRY + } else { + SAFETY_BUFFERING_MESSAGE_WITHOUT_RETRY + }; + self.bottom_pane.ensure_status_indicator(); + self.set_status( + "Working".to_string(), + Some(message.to_string()), + StatusDetailsCapitalization::Preserve, + /*details_max_lines*/ 6, + ); + + let (Some(faster_model), Some(turn), Some(thread_id)) = + (faster_model, retry_turn, thread_id) + else { + return; + }; + if !should_show_retry_prompt { + return; + } + + let header = ColumnRenderable::with(vec![ + Box::new(Line::from("Additional safety checks").bold()) as Box, + Box::new( + Paragraph::new(Line::from(SAFETY_BUFFERING_MESSAGE_WITH_RETRY).dim()) + .wrap(Wrap { trim: false }), + ), + ]); + self.bottom_pane.show_selection_view(SelectionViewParams { + view_id: Some(SAFETY_BUFFERING_PROMPT_VIEW_ID), + header: Box::new(header), + items: vec![ + SelectionItem { + name: "Retry with a faster model".to_string(), + actions: vec![Box::new(move |tx| { + tx.send(AppEvent::RetrySafetyBufferedTurn { + thread_id, + turn_id: turn_id.clone(), + model: faster_model.clone(), + turn: turn.clone(), + }); + })], + dismiss_on_select: true, + ..Default::default() + }, + SelectionItem { + name: "Keep waiting".to_string(), + dismiss_on_select: true, + ..Default::default() + }, + ], + ..Default::default() + }); + } +} diff --git a/codex-rs/tui/src/chatwidget/session_flow.rs b/codex-rs/tui/src/chatwidget/session_flow.rs index 004f64c726..177700d4f4 100644 --- a/codex-rs/tui/src/chatwidget/session_flow.rs +++ b/codex-rs/tui/src/chatwidget/session_flow.rs @@ -27,6 +27,7 @@ impl ChatWidget { } self.refresh_plan_mode_nudge(); self.turn_lifecycle.reset_thread(); + self.clear_safety_buffering(); self.thread_name = session.thread_name.clone(); self.current_goal_status_indicator = None; self.current_goal_status = None; diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__safety_buffering_retry_prompt.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__safety_buffering_retry_prompt.snap new file mode 100644 index 0000000000..60fd4d38c3 --- /dev/null +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__safety_buffering_retry_prompt.snap @@ -0,0 +1,13 @@ +--- +source: tui/src/chatwidget/tests/app_server.rs +expression: popup +--- + Additional safety checks + This request requires additional safety checks, which can take extra time. + You can retry with a faster model for a quicker response, though it may be + less capable of handling complex requests. + +› 1. Retry with a faster model + 2. Keep waiting + + Press enter to confirm or esc to go back diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__safety_buffering_status_without_retry.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__safety_buffering_status_without_retry.snap new file mode 100644 index 0000000000..a719f937cd --- /dev/null +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__safety_buffering_status_without_retry.snap @@ -0,0 +1,11 @@ +--- +source: tui/src/chatwidget/tests/app_server.rs +expression: "render_bottom_popup(&chat, 80)" +--- +• Working (0s • esc to interrupt) + └ This request requires additional safety checks, which can take extra time. + + +› Ask Codex to do anything + + gpt-5.5 default · /tmp/project diff --git a/codex-rs/tui/src/chatwidget/streaming.rs b/codex-rs/tui/src/chatwidget/streaming.rs index 897094b1f6..f8b0124203 100644 --- a/codex-rs/tui/src/chatwidget/streaming.rs +++ b/codex-rs/tui/src/chatwidget/streaming.rs @@ -203,6 +203,11 @@ impl ChatWidget { // (between **/**) as the chunk header. Show this header as status. self.reasoning_buffer.push_str(&delta); + if self.safety_buffering_is_waiting() { + self.request_redraw(); + return; + } + if self.unified_exec_wait_streak.is_some() { // Unified exec waiting should take precedence over reasoning-derived status headers. self.request_redraw(); @@ -383,6 +388,7 @@ impl ChatWidget { pub(super) fn handle_streaming_delta(&mut self, delta: String) { if !delta.is_empty() { self.record_visible_turn_activity(); + self.mark_safety_buffering_agent_message_started(); } if self.stream_controller.is_none() { // Before starting an agent stream, flush any active exec cell group. diff --git a/codex-rs/tui/src/chatwidget/tests.rs b/codex-rs/tui/src/chatwidget/tests.rs index 1af814eb76..24423d3737 100644 --- a/codex-rs/tui/src/chatwidget/tests.rs +++ b/codex-rs/tui/src/chatwidget/tests.rs @@ -77,6 +77,7 @@ pub(super) use codex_app_server_protocol::MarketplaceUpgradeResponse; pub(super) use codex_app_server_protocol::McpServerStartupState; pub(super) use codex_app_server_protocol::McpServerStatusDetail; pub(super) use codex_app_server_protocol::McpServerStatusUpdatedNotification; +pub(super) use codex_app_server_protocol::ModelSafetyBufferingUpdatedNotification; pub(super) use codex_app_server_protocol::ModelVerification as AppServerModelVerification; pub(super) use codex_app_server_protocol::ModelVerificationNotification; pub(super) use codex_app_server_protocol::NonSteerableTurnKind; diff --git a/codex-rs/tui/src/chatwidget/tests/app_server.rs b/codex-rs/tui/src/chatwidget/tests/app_server.rs index a10c2cca93..2e6d7f8f7f 100644 --- a/codex-rs/tui/src/chatwidget/tests/app_server.rs +++ b/codex-rs/tui/src/chatwidget/tests/app_server.rs @@ -61,6 +61,194 @@ fn configured_thread_session(thread_id: ThreadId) -> crate::session_state::Threa } } +fn start_safety_buffering_test_turn( + chat: &mut ChatWidget, + op_rx: &mut tokio::sync::mpsc::UnboundedReceiver, +) -> (ThreadId, &'static str, Op) { + let thread_id = ThreadId::new(); + let turn_id = "turn-safety-buffering"; + chat.thread_id = Some(thread_id); + chat.submit_user_message(UserMessage::from("Explain the request")); + let turn = next_submit_op(op_rx); + assert_matches!(&turn, Op::UserTurn { .. }); + chat.record_safety_buffering_turn(turn_id.to_string(), &turn); + chat.handle_server_notification( + ServerNotification::TurnStarted(TurnStartedNotification { + thread_id: thread_id.to_string(), + turn: AppServerTurn { + id: turn_id.to_string(), + items_view: codex_app_server_protocol::TurnItemsView::Full, + items: Vec::new(), + status: AppServerTurnStatus::InProgress, + error: None, + started_at: Some(0), + completed_at: None, + duration_ms: None, + }, + }), + /*replay_kind*/ None, + ); + (thread_id, turn_id, turn) +} + +fn safety_buffering_notification( + thread_id: ThreadId, + turn_id: &str, + faster_model: Option<&str>, +) -> ModelSafetyBufferingUpdatedNotification { + ModelSafetyBufferingUpdatedNotification { + thread_id: thread_id.to_string(), + turn_id: turn_id.to_string(), + model: "current-model".to_string(), + use_cases: Vec::new(), + reasons: Vec::new(), + show_buffering_ui: true, + faster_model: faster_model.map(str::to_string), + } +} + +#[tokio::test] +async fn safety_buffering_offers_one_retry_with_app_wording() { + let (mut chat, mut rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + let (thread_id, turn_id, _) = start_safety_buffering_test_turn(&mut chat, &mut op_rx); + + let notification = safety_buffering_notification(thread_id, turn_id, Some("faster-model")); + chat.handle_server_notification( + ServerNotification::ModelSafetyBufferingUpdated(notification.clone()), + /*replay_kind*/ None, + ); + chat.handle_server_notification( + ServerNotification::ModelSafetyBufferingUpdated(notification), + /*replay_kind*/ None, + ); + + let popup = render_bottom_popup(&chat, /*width*/ 80); + assert_chatwidget_snapshot!("safety_buffering_retry_prompt", popup); + + chat.handle_key_event(KeyEvent::new(KeyCode::Enter, KeyModifiers::NONE)); + let (event_thread_id, event_turn_id, model, turn) = loop { + match rx.try_recv() { + Ok(AppEvent::RetrySafetyBufferedTurn { + thread_id, + turn_id, + model, + turn, + }) => break (thread_id, turn_id, model, turn), + Ok(_) => continue, + Err(err) => panic!("expected safety-buffering retry event: {err}"), + } + }; + assert_eq!(event_thread_id, thread_id); + assert_eq!(event_turn_id, turn_id); + assert_eq!(model, "faster-model"); + assert_matches!(turn, Op::UserTurn { .. }); + assert!(!render_bottom_popup(&chat, /*width*/ 80).contains("Additional safety checks")); +} + +#[tokio::test] +async fn safety_buffering_stops_retrying_after_agent_message_starts() { + let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + let (thread_id, turn_id, _) = start_safety_buffering_test_turn(&mut chat, &mut op_rx); + chat.handle_server_notification( + ServerNotification::ModelSafetyBufferingUpdated(safety_buffering_notification( + thread_id, + turn_id, + Some("faster-model"), + )), + /*replay_kind*/ None, + ); + assert!(chat.can_retry_safety_buffered_turn(turn_id)); + + chat.on_agent_message_delta("Visible response".to_string()); + + assert!(!chat.can_retry_safety_buffered_turn(turn_id)); +} + +#[tokio::test] +async fn safety_buffering_without_retry_shows_short_app_message() { + let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + let (thread_id, turn_id, turn) = start_safety_buffering_test_turn(&mut chat, &mut op_rx); + + chat.handle_server_notification( + ServerNotification::ModelSafetyBufferingUpdated(safety_buffering_notification( + thread_id, turn_id, /*faster_model*/ None, + )), + /*replay_kind*/ None, + ); + + let render_popup = |chat: &ChatWidget| { + normalize_snapshot_paths(render_bottom_popup(chat, /*width*/ 80)) + }; + let popup = render_popup(&chat); + assert_chatwidget_snapshot!("safety_buffering_status_without_retry", popup,); + + let notification = safety_buffering_notification(thread_id, turn_id, Some("faster-model")); + chat.record_safety_buffering_turn("other-turn".to_string(), &turn); + chat.handle_server_notification( + ServerNotification::ModelSafetyBufferingUpdated(notification.clone()), + /*replay_kind*/ None, + ); + assert_eq!(render_popup(&chat), popup); + + chat.record_safety_buffering_turn(turn_id.to_string(), &turn); + chat.handle_server_notification( + ServerNotification::ModelSafetyBufferingUpdated(notification), + Some(ReplayKind::ThreadSnapshot), + ); + assert_eq!(render_popup(&chat), popup); +} + +#[tokio::test] +async fn safety_buffering_ignores_hidden_stale_and_historical_updates() { + let (mut chat, _rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + let (thread_id, turn_id, _) = start_safety_buffering_test_turn(&mut chat, &mut op_rx); + + let mut hidden = safety_buffering_notification(thread_id, turn_id, Some("faster-model")); + hidden.show_buffering_ui = false; + chat.handle_server_notification( + ServerNotification::ModelSafetyBufferingUpdated(hidden), + /*replay_kind*/ None, + ); + chat.handle_server_notification( + ServerNotification::ModelSafetyBufferingUpdated(safety_buffering_notification( + thread_id, + "stale-turn", + Some("faster-model"), + )), + /*replay_kind*/ None, + ); + chat.handle_server_notification( + ServerNotification::ModelSafetyBufferingUpdated(safety_buffering_notification( + thread_id, + turn_id, + Some("faster-model"), + )), + Some(ReplayKind::ResumeInitialMessages), + ); + assert!(!render_bottom_popup(&chat, /*width*/ 80).contains("Additional safety checks")); + + let mut hidden = safety_buffering_notification(thread_id, turn_id, Some("faster-model")); + chat.handle_server_notification( + ServerNotification::ModelSafetyBufferingUpdated(hidden.clone()), + /*replay_kind*/ None, + ); + assert!(render_bottom_popup(&chat, /*width*/ 80).contains("Additional safety checks")); + hidden.show_buffering_ui = false; + chat.handle_server_notification( + ServerNotification::ModelSafetyBufferingUpdated(hidden), + /*replay_kind*/ None, + ); + + assert_eq!( + chat.bottom_pane + .status_widget() + .expect("status indicator should be visible") + .details(), + None + ); + assert!(!render_bottom_popup(&chat, /*width*/ 80).contains("Additional safety checks")); +} + #[tokio::test] async fn invalid_url_elicitation_is_declined() { let (mut chat, _app_event_tx, mut rx, _op_rx) = make_chatwidget_manual_with_sender().await; @@ -969,6 +1157,24 @@ async fn live_app_server_cyber_policy_error_renders_dedicated_notice() { assert!(!chat.bottom_pane.is_task_running()); } +#[tokio::test] +async fn app_server_safety_access_errors_render_dedicated_notice() { + let message = "Invalid prompt: we've limited access to this content for safety reasons."; + for message in [ + message.to_string(), + json!({ "error": { "message": message } }).to_string(), + ] { + let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + chat.handle_non_retry_error(message, /*codex_error_info*/ None); + + let cells = drain_insert_history(&mut rx); + assert_eq!(cells.len(), 1); + let rendered = lines_to_single_string(&cells[0]); + assert!(rendered.contains("This content can't be shown")); + assert!(!rendered.contains("Invalid prompt:")); + } +} + #[tokio::test] async fn live_app_server_model_verification_renders_warning() { let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; diff --git a/codex-rs/tui/src/chatwidget/turn_runtime.rs b/codex-rs/tui/src/chatwidget/turn_runtime.rs index df3c03ed5a..8af43550e3 100644 --- a/codex-rs/tui/src/chatwidget/turn_runtime.rs +++ b/codex-rs/tui/src/chatwidget/turn_runtime.rs @@ -5,6 +5,9 @@ use super::*; +const SAFETY_ACCESS_BLOCK_PREFIX: &str = + "Invalid prompt: we've limited access to this content for safety reasons."; + impl ChatWidget { /// Synchronize the bottom-pane "task running" indicator with the current lifecycles. /// @@ -48,6 +51,7 @@ impl ChatWidget { pub(super) fn on_task_started(&mut self) { self.input_queue.user_turn_pending_start = false; + self.reset_safety_buffering_for_turn_start(); self.turn_lifecycle.start(Instant::now()); self.transcript.reset_turn_flags(); self.adaptive_chunking.reset(); @@ -297,6 +301,7 @@ impl ChatWidget { /// This does not clear MCP startup tracking, because MCP startup can overlap with turn cleanup /// and should continue to drive the bottom-pane running indicator while it is in progress. pub(super) fn finalize_turn(&mut self) { + self.clear_safety_buffering(); // Drop preview-only stream tail content on any termination path before // failed-cell finalization, so transient tail cells are never persisted. self.clear_active_stream_tail(); @@ -424,6 +429,18 @@ impl ChatWidget { .is_some_and(is_app_server_cyber_policy_error) { self.on_cyber_policy_error(); + } else if message.starts_with(SAFETY_ACCESS_BLOCK_PREFIX) + || serde_json::from_str::(&message).is_ok_and(|response| { + response["error"]["message"] + .as_str() + .is_some_and(|message| message.starts_with(SAFETY_ACCESS_BLOCK_PREFIX)) + }) + { + self.input_queue.submit_pending_steers_after_interrupt = false; + self.finalize_turn(); + self.add_to_history(history_cell::new_safety_access_block_event()); + self.request_redraw(); + self.maybe_send_next_queued_input(); } else if let Some(info) = codex_error_info .as_ref() .and_then(app_server_rate_limit_error_kind) diff --git a/codex-rs/tui/src/history_cell/notices.rs b/codex-rs/tui/src/history_cell/notices.rs index d9dff4f5a0..c4cc79cf70 100644 --- a/codex-rs/tui/src/history_cell/notices.rs +++ b/codex-rs/tui/src/history_cell/notices.rs @@ -85,6 +85,37 @@ pub(crate) fn new_warning_event(message: String) -> PrefixedWrappedHistoryCell { PrefixedWrappedHistoryCell::new(message.yellow(), "⚠ ".yellow(), " ") } +#[derive(Debug)] +pub(crate) struct SafetyAccessBlockCell; + +const SAFETY_ACCESS_BLOCK_TITLE: &str = "This content can't be shown"; +const SAFETY_ACCESS_BLOCK_BODY: &str = "We've limited access to this content for safety reasons. This type of information may be used to benefit or to harm people."; + +pub(crate) fn new_safety_access_block_event() -> SafetyAccessBlockCell { + SafetyAccessBlockCell +} + +impl HistoryCell for SafetyAccessBlockCell { + fn display_lines(&self, width: u16) -> Vec> { + let mut lines = vec![vec!["ⓘ ".cyan(), SAFETY_ACCESS_BLOCK_TITLE.bold()].into()]; + let body = Line::from(vec![" ".into(), SAFETY_ACCESS_BLOCK_BODY.dim()]); + let wrap_width = width.saturating_sub(2).max(1) as usize; + let wrapped = adaptive_wrap_line( + &body, + RtOptions::new(wrap_width).subsequent_indent(" ".into()), + ); + push_owned_lines(&wrapped, &mut lines); + lines + } + + fn raw_lines(&self) -> Vec> { + vec![ + Line::from(SAFETY_ACCESS_BLOCK_TITLE), + Line::from(SAFETY_ACCESS_BLOCK_BODY), + ] + } +} + const TRUSTED_ACCESS_FOR_CYBER_URL: &str = "https://chatgpt.com/cyber"; #[derive(Debug)] diff --git a/codex-rs/tui/src/history_cell/snapshots/codex_tui__history_cell__tests__safety_access_block_event_snapshot.snap b/codex-rs/tui/src/history_cell/snapshots/codex_tui__history_cell__tests__safety_access_block_event_snapshot.snap new file mode 100644 index 0000000000..f512697888 --- /dev/null +++ b/codex-rs/tui/src/history_cell/snapshots/codex_tui__history_cell__tests__safety_access_block_event_snapshot.snap @@ -0,0 +1,7 @@ +--- +source: tui/src/history_cell/tests.rs +expression: rendered +--- +ⓘ This content can't be shown + We've limited access to this content for safety reasons. This type of + information may be used to benefit or to harm people. diff --git a/codex-rs/tui/src/history_cell/tests.rs b/codex-rs/tui/src/history_cell/tests.rs index 99f7a3a93b..42dadb8535 100644 --- a/codex-rs/tui/src/history_cell/tests.rs +++ b/codex-rs/tui/src/history_cell/tests.rs @@ -699,6 +699,13 @@ fn cyber_policy_error_event_snapshot() { insta::assert_snapshot!(rendered); } +#[test] +fn safety_access_block_event_snapshot() { + let cell = new_safety_access_block_event(); + let rendered = render_lines(&cell.display_lines(/*width*/ 80)).join("\n"); + insta::assert_snapshot!(rendered); +} + #[test] fn cyber_policy_error_event_narrow_snapshot() { let cell = new_cyber_policy_error_event(); From f8937b7d862e3d507fe22918e5daa35c7da993e9 Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Wed, 24 Jun 2026 19:51:51 -0700 Subject: [PATCH 5/7] Represent MCP authentication with an enum (#29924) ## Why MCP authentication has distinct OAuth and ChatGPT-session flows. Representing that choice as `use_chatgpt_auth` makes one flow implicit and allows the configuration model to express the distinction only through a boolean. ChatGPT credential forwarding also needs a first-party trust boundary. A configurable `chatgpt_base_url` controls routing, but must not grant an MCP server permission to receive session credentials. This change builds on #29733, where the boolean was introduced. ## What changed - Replace `use_chatgpt_auth` with an `auth` field backed by the exhaustive `McpServerAuth` enum. - Support `auth = "oauth"` and `auth = "chatgpt"`, with OAuth remaining the default. - Trust only the origin derived from the existing hardcoded `CHATGPT_CODEX_BASE_URL` when granting ChatGPT auth to an MCP server. - Keep configured bearer tokens and authorization headers ahead of the selected authentication flow. - Update config writers, schema output, fixtures, and integration-test setup to use the enum. ## Verification Integration coverage exercises the complete streamable HTTP startup path in two independent configurations: - A directly constructed MCP configuration verifies that matching an overridden `chatgpt_base_url` does not grant ChatGPT auth. - A persisted `config.toml` containing an attacker-controlled `chatgpt_base_url` and `auth = "chatgpt"` verifies the same boundary through normal config parsing. Both tests complete MCP initialization and tool listing and assert that the full captured request sequence contains no authorization headers. Separate integration coverage verifies that configured authorization takes precedence over ChatGPT auth. --- codex-rs/cli/src/mcp_cmd.rs | 2 +- codex-rs/codex-mcp/src/catalog_tests.rs | 2 +- codex-rs/codex-mcp/src/connection_manager.rs | 3 +- .../codex-mcp/src/connection_manager_tests.rs | 8 +- codex-rs/codex-mcp/src/mcp/auth.rs | 3 +- codex-rs/codex-mcp/src/mcp/mod.rs | 39 +++--- codex-rs/codex-mcp/src/mcp/mod_tests.rs | 4 +- codex-rs/codex-mcp/src/plugin_config_tests.rs | 4 +- codex-rs/codex-mcp/src/runtime.rs | 4 +- codex-rs/config/src/lib.rs | 1 + codex-rs/config/src/mcp_edit.rs | 5 +- codex-rs/config/src/mcp_edit_tests.rs | 4 +- codex-rs/config/src/mcp_types.rs | 42 ++++-- codex-rs/config/src/mcp_types_tests.rs | 2 +- codex-rs/config/src/types.rs | 1 + codex-rs/core-plugins/src/manager_tests.rs | 10 +- codex-rs/core/config.schema.json | 31 ++++- codex-rs/core/src/config/config_tests.rs | 34 ++--- .../core/src/config/edit/document_helpers.rs | 5 +- codex-rs/core/src/config/edit_tests.rs | 14 +- codex-rs/core/src/mcp_skill_dependencies.rs | 4 +- codex-rs/core/tests/suite/code_mode.rs | 2 +- codex-rs/core/tests/suite/hooks_mcp.rs | 2 +- .../core/tests/suite/mcp_refresh_cleanup.rs | 2 +- codex-rs/core/tests/suite/rmcp_client.rs | 126 ++++++++++-------- codex-rs/core/tests/suite/search_tool.rs | 6 +- codex-rs/core/tests/suite/sqlite_state.rs | 2 +- codex-rs/core/tests/suite/token_budget.rs | 2 +- codex-rs/core/tests/suite/truncation.rs | 6 +- .../mcp/src/executor_plugin/provider_tests.rs | 4 +- codex-rs/model-provider/src/lib.rs | 1 + 31 files changed, 228 insertions(+), 147 deletions(-) diff --git a/codex-rs/cli/src/mcp_cmd.rs b/codex-rs/cli/src/mcp_cmd.rs index e874ca6ab9..15f0b0bb4d 100644 --- a/codex-rs/cli/src/mcp_cmd.rs +++ b/codex-rs/cli/src/mcp_cmd.rs @@ -345,7 +345,7 @@ async fn run_add(config_overrides: &CliConfigOverrides, add_args: AddArgs) -> Re }; let new_entry = McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: transport.clone(), environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, diff --git a/codex-rs/codex-mcp/src/catalog_tests.rs b/codex-rs/codex-mcp/src/catalog_tests.rs index 09a061441b..0b2ffef50d 100644 --- a/codex-rs/codex-mcp/src/catalog_tests.rs +++ b/codex-rs/codex-mcp/src/catalog_tests.rs @@ -17,7 +17,7 @@ use super::ResolvedMcpCatalog; fn server(url: &str) -> McpServerConfig { McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: url.to_string(), bearer_token_env_var: None, diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 0747173470..55c72b0467 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -43,6 +43,7 @@ use anyhow::anyhow; use async_channel::Sender; use codex_api::SharedAuthProvider; use codex_config::Constrained; +use codex_config::McpServerAuth; use codex_config::McpServerTransportConfig; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; @@ -872,7 +873,7 @@ fn chatgpt_auth_provider_for_server( ) -> Option { if !server .configured_config() - .is_some_and(|config| config.use_chatgpt_auth) + .is_some_and(|config| matches!(&config.auth, McpServerAuth::ChatGpt)) { return None; } diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 46fa7fae6c..d66474c7bb 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -1272,7 +1272,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { ( "stdio".to_string(), EffectiveMcpServer::configured(McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "echo".to_string(), args: Vec::new(), @@ -1299,7 +1299,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { ( "http".to_string(), EffectiveMcpServer::configured(McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "http://127.0.0.1:1".to_string(), bearer_token_env_var: None, @@ -1407,7 +1407,7 @@ fn mcp_init_error_display_prompts_for_github_pat() { let server_name = "github"; let entry = McpAuthStatusEntry { config: Some(McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "https://api.githubcopilot.com/mcp/".to_string(), bearer_token_env_var: None, @@ -1461,7 +1461,7 @@ fn mcp_init_error_display_reports_generic_errors() { let server_name = "custom"; let entry = McpAuthStatusEntry { config: Some(McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com".to_string(), bearer_token_env_var: Some("TOKEN".to_string()), diff --git a/codex-rs/codex-mcp/src/mcp/auth.rs b/codex-rs/codex-mcp/src/mcp/auth.rs index 182c6b6c3e..d19766af1f 100644 --- a/codex-rs/codex-mcp/src/mcp/auth.rs +++ b/codex-rs/codex-mcp/src/mcp/auth.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use anyhow::Result; +use codex_config::McpServerAuth; use codex_config::McpServerConfig; use codex_config::McpServerTransportConfig; use codex_config::types::AuthKeyringBackendKind; @@ -141,7 +142,7 @@ where let config = server.configured_config().cloned(); let has_runtime_auth = config .as_ref() - .is_some_and(|config| config.use_chatgpt_auth) + .is_some_and(|config| matches!(&config.auth, McpServerAuth::ChatGpt)) && auth.is_some_and(CodexAuth::uses_codex_backend) && config.as_ref().is_some_and(|config| { matches!( diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index d2a6a9f2d6..f2c8f1d020 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -19,6 +19,7 @@ use std::time::Duration; use async_channel::unbounded; use codex_config::Constrained; +use codex_config::McpServerAuth; use codex_config::McpServerConfig; use codex_config::McpServerTransportConfig; use codex_config::types::AppToolApproval; @@ -26,6 +27,7 @@ use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; use codex_connectors::ConnectorSnapshot; use codex_login::CodexAuth; +use codex_model_provider::CHATGPT_CODEX_BASE_URL; use codex_protocol::mcp::McpServerInfo; use codex_protocol::mcp::Resource; use codex_protocol::mcp::ResourceTemplate; @@ -257,23 +259,28 @@ pub fn effective_mcp_servers_from_configured( config: &McpConfig, auth: Option<&CodexAuth>, ) -> HashMap { - let chatgpt_origin = url::Url::parse(&config.chatgpt_base_url) + let chatgpt_origin = url::Url::parse(CHATGPT_CODEX_BASE_URL) .ok() - .filter(|url| matches!(url.scheme(), "http" | "https")) .map(|url| url.origin()); let mut servers = configured_servers .into_iter() .map(|(name, mut server)| { - if server.use_chatgpt_auth { - let server_origin = match &server.transport { - McpServerTransportConfig::StreamableHttp { url, .. } => url::Url::parse(url) - .ok() - .filter(|url| matches!(url.scheme(), "http" | "https")) - .map(|url| url.origin()), - McpServerTransportConfig::Stdio { .. } => None, - }; - server.use_chatgpt_auth = - server_origin.is_some() && server_origin.as_ref() == chatgpt_origin.as_ref(); + match server.auth.clone() { + McpServerAuth::ChatGpt => { + let server_origin = match &server.transport { + McpServerTransportConfig::StreamableHttp { url, .. } => { + url::Url::parse(url) + .ok() + .filter(|url| matches!(url.scheme(), "http" | "https")) + .map(|url| url.origin()) + } + McpServerTransportConfig::Stdio { .. } => None, + }; + if server_origin.as_ref() != chatgpt_origin.as_ref() { + server.auth = McpServerAuth::OAuth; + } + } + McpServerAuth::OAuth => {} } (name, EffectiveMcpServer::configured(server)) }) @@ -474,7 +481,7 @@ pub fn codex_apps_mcp_server_config( mcp_server_config_for_url( codex_apps_mcp_url_for_base_url(chatgpt_base_url), apps_mcp_product_sku, - /*use_chatgpt_auth*/ true, + McpServerAuth::ChatGpt, ) } @@ -492,14 +499,14 @@ pub fn hosted_plugin_runtime_mcp_server_config( mcp_server_config_for_url( format!("{base_url}/ps/mcp"), apps_mcp_product_sku, - /*use_chatgpt_auth*/ true, + McpServerAuth::ChatGpt, ) } fn mcp_server_config_for_url( url: String, apps_mcp_product_sku: Option<&str>, - use_chatgpt_auth: bool, + auth_mode: McpServerAuth, ) -> McpServerConfig { let http_headers = apps_mcp_product_sku.map(|product_sku| { HashMap::from([("X-OpenAI-Product-Sku".to_string(), product_sku.to_string())]) @@ -512,7 +519,7 @@ fn mcp_server_config_for_url( http_headers, env_http_headers: None, }, - use_chatgpt_auth, + auth: auth_mode, environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), enabled: true, required: false, diff --git a/codex-rs/codex-mcp/src/mcp/mod_tests.rs b/codex-rs/codex-mcp/src/mcp/mod_tests.rs index 389bf1adcb..7d166cbab8 100644 --- a/codex-rs/codex-mcp/src/mcp/mod_tests.rs +++ b/codex-rs/codex-mcp/src/mcp/mod_tests.rs @@ -296,7 +296,7 @@ async fn effective_mcp_servers_preserve_runtime_servers() { catalog.register(McpServerRegistration::from_config( "sample".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "https://user.example/mcp".to_string(), bearer_token_env_var: None, @@ -322,7 +322,7 @@ async fn effective_mcp_servers_preserve_runtime_servers() { catalog.register(McpServerRegistration::from_config( "docs".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "https://docs.example/mcp".to_string(), bearer_token_env_var: None, diff --git a/codex-rs/codex-mcp/src/plugin_config_tests.rs b/codex-rs/codex-mcp/src/plugin_config_tests.rs index fc495d7ee8..576bd29618 100644 --- a/codex-rs/codex-mcp/src/plugin_config_tests.rs +++ b/codex-rs/codex-mcp/src/plugin_config_tests.rs @@ -32,7 +32,7 @@ fn stdio_server( env_vars: Vec, ) -> McpServerConfig { McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: command.to_string(), args: Vec::new(), @@ -67,7 +67,7 @@ fn declared_placement_preserves_local_plugin_normalization() { Vec::new(), ); let expected_http = McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), bearer_token_env_var: None, diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index be18660e2b..be9d50d3f4 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -104,7 +104,7 @@ mod tests { fn stdio_server(environment_id: &str) -> McpServerConfig { McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "echo".to_string(), args: Vec::new(), @@ -131,7 +131,7 @@ mod tests { fn http_server(environment_id: &str) -> McpServerConfig { McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "http://127.0.0.1:1".to_string(), bearer_token_env_var: None, diff --git a/codex-rs/config/src/lib.rs b/codex-rs/config/src/lib.rs index 1cbbbaf0a9..afda2cf9f3 100644 --- a/codex-rs/config/src/lib.rs +++ b/codex-rs/config/src/lib.rs @@ -115,6 +115,7 @@ pub use mcp_edit::ConfigEditsBuilder; pub use mcp_edit::load_global_mcp_servers; pub use mcp_types::AppToolApproval; pub use mcp_types::DEFAULT_MCP_SERVER_ENVIRONMENT_ID; +pub use mcp_types::McpServerAuth; pub use mcp_types::McpServerConfig; pub use mcp_types::McpServerDisabledReason; pub use mcp_types::McpServerEnvVar; diff --git a/codex-rs/config/src/mcp_edit.rs b/codex-rs/config/src/mcp_edit.rs index bc5d4d3316..16b3904917 100644 --- a/codex-rs/config/src/mcp_edit.rs +++ b/codex-rs/config/src/mcp_edit.rs @@ -13,6 +13,7 @@ use toml_edit::value; use crate::AppToolApproval; use crate::CONFIG_TOML_FILE; +use crate::McpServerAuth; use crate::McpServerConfig; use crate::McpServerEnvVar; use crate::McpServerTransportConfig; @@ -172,8 +173,8 @@ fn serialize_mcp_server(config: &McpServerConfig) -> TomlItem { } } - if config.use_chatgpt_auth { - entry["use_chatgpt_auth"] = value(true); + if matches!(&config.auth, McpServerAuth::ChatGpt) { + entry["auth"] = value("chatgpt"); } if !config.enabled { entry["enabled"] = value(false); diff --git a/codex-rs/config/src/mcp_edit_tests.rs b/codex-rs/config/src/mcp_edit_tests.rs index 3c2801b269..10ef23401b 100644 --- a/codex-rs/config/src/mcp_edit_tests.rs +++ b/codex-rs/config/src/mcp_edit_tests.rs @@ -16,7 +16,7 @@ async fn replace_mcp_servers_serializes_per_tool_approval_overrides() -> anyhow: let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "docs-server".to_string(), args: Vec::new(), @@ -96,7 +96,7 @@ async fn replace_mcp_servers_serializes_oauth_client_id() -> anyhow::Result<()> let servers = BTreeMap::from([( "maas_outlook".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), bearer_token_env_var: None, diff --git a/codex-rs/config/src/mcp_types.rs b/codex-rs/config/src/mcp_types.rs index 08bf2fef2f..550de84a0c 100644 --- a/codex-rs/config/src/mcp_types.rs +++ b/codex-rs/config/src/mcp_types.rs @@ -126,16 +126,40 @@ pub struct McpServerOAuthConfig { pub client_id: Option, } +/// Authentication flow Codex attempts after resolving an HTTP MCP server's +/// configured bearer token and authorization headers, which always take +/// precedence. ChatGPT authentication falls back to stored OAuth credentials +/// when its session provider is unavailable; both modes ultimately fall back +/// to an unauthenticated connection. +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum McpServerAuth { + /// Use stored MCP OAuth credentials when available. Starting an OAuth login + /// is a separate operation. + #[default] + #[serde(rename = "oauth")] + OAuth, + /// Use the current ChatGPT session for servers on the trusted first-party + /// ChatGPT origin. If no ChatGPT session provider is available, startup can + /// still fall back to stored OAuth credentials. + #[serde(rename = "chatgpt")] + ChatGpt, +} + +impl McpServerAuth { + fn is_default(&self) -> bool { + self == &Self::default() + } +} + #[derive(Serialize, Debug, Clone, PartialEq)] pub struct McpServerConfig { #[serde(flatten)] pub transport: McpServerTransportConfig, - /// When `true`, request authentication with the current ChatGPT session. Codex honors this - /// only when the server URL has the same origin as the configured ChatGPT base URL and no - /// configured bearer token or authorization header resolves. - #[serde(default, skip_serializing_if = "std::ops::Not::not")] - pub use_chatgpt_auth: bool, + /// Authentication flow to use when no configured authorization resolves. + #[serde(default, skip_serializing_if = "McpServerAuth::is_default")] + pub auth: McpServerAuth, /// Effective environment id for where Codex should start this MCP server. pub environment_id: String, @@ -245,7 +269,7 @@ pub struct RawMcpServerConfig { #[serde(default)] pub environment_id: Option, #[serde(default)] - pub use_chatgpt_auth: Option, + pub auth: Option, #[serde(default)] pub startup_timeout_sec: Option, #[serde(default)] @@ -294,7 +318,7 @@ impl TryFrom for McpServerConfig { bearer_token, bearer_token_env_var, environment_id, - use_chatgpt_auth, + auth, startup_timeout_sec, startup_timeout_ms, tool_timeout_sec, @@ -338,7 +362,7 @@ impl TryFrom for McpServerConfig { throw_if_set("stdio", "env_http_headers", env_http_headers.as_ref())?; throw_if_set("stdio", "oauth", oauth.as_ref())?; throw_if_set("stdio", "oauth_resource", oauth_resource.as_ref())?; - throw_if_set("stdio", "use_chatgpt_auth", use_chatgpt_auth.as_ref())?; + throw_if_set("stdio", "auth", auth.as_ref())?; let env_vars = env_vars.unwrap_or_default(); for env_var in &env_vars { env_var.validate_source()?; @@ -371,7 +395,7 @@ impl TryFrom for McpServerConfig { Ok(Self { transport, - use_chatgpt_auth: use_chatgpt_auth.unwrap_or_default(), + auth: auth.unwrap_or_default(), environment_id, startup_timeout_sec, tool_timeout_sec, diff --git a/codex-rs/config/src/mcp_types_tests.rs b/codex-rs/config/src/mcp_types_tests.rs index cab51f14c7..598a532971 100644 --- a/codex-rs/config/src/mcp_types_tests.rs +++ b/codex-rs/config/src/mcp_types_tests.rs @@ -426,7 +426,7 @@ fn deserialize_ignores_unknown_server_fields() { assert_eq!( cfg, McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "echo".to_string(), args: vec![], diff --git a/codex-rs/config/src/types.rs b/codex-rs/config/src/types.rs index d25ab64c66..5303fccd43 100644 --- a/codex-rs/config/src/types.rs +++ b/codex-rs/config/src/types.rs @@ -4,6 +4,7 @@ // definitions that do not contain business logic. pub use crate::mcp_types::AppToolApproval; +pub use crate::mcp_types::McpServerAuth; pub use crate::mcp_types::McpServerConfig; pub use crate::mcp_types::McpServerDisabledReason; pub use crate::mcp_types::McpServerEnvVar; diff --git a/codex-rs/core-plugins/src/manager_tests.rs b/codex-rs/core-plugins/src/manager_tests.rs index e4441e15d0..cefca287ad 100644 --- a/codex-rs/core-plugins/src/manager_tests.rs +++ b/codex-rs/core-plugins/src/manager_tests.rs @@ -677,7 +677,7 @@ async fn load_plugins_loads_default_skills_and_mcp_servers() { mcp_servers: HashMap::from([( "sample".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "https://sample.example/mcp".to_string(), bearer_token_env_var: None, @@ -769,7 +769,7 @@ enabled = true HashMap::from([( "counter".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "https://sample.example/counter/mcp".to_string(), bearer_token_env_var: None, @@ -1664,7 +1664,7 @@ async fn load_plugins_uses_manifest_configured_component_paths() { HashMap::from([( "custom".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "https://custom.example/mcp".to_string(), bearer_token_env_var: None, @@ -1845,7 +1845,7 @@ async fn load_plugins_ignores_manifest_component_paths_without_dot_slash() { HashMap::from([( "default".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "https://default.example/mcp".to_string(), bearer_token_env_var: None, @@ -2096,7 +2096,7 @@ fn capability_index_filters_inactive_and_zero_capability_plugins() { let connector = |id: &str| AppConnectorId(id.to_string()); let app = |name: &str, connector_id: &str| app_declaration(name, connector_id); let http_server = |url: &str| McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: url.to_string(), bearer_token_env_var: None, diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index f8cfa6f726..842f4d5cc5 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -1384,6 +1384,25 @@ }, "type": "object" }, + "McpServerAuth": { + "description": "Authentication flow Codex attempts after resolving an HTTP MCP server's configured bearer token and authorization headers, which always take precedence. ChatGPT authentication falls back to stored OAuth credentials when its session provider is unavailable; both modes ultimately fall back to an unauthenticated connection.", + "oneOf": [ + { + "description": "Use stored MCP OAuth credentials when available. Starting an OAuth login is a separate operation.", + "enum": [ + "oauth" + ], + "type": "string" + }, + { + "description": "Use the current ChatGPT session for servers on the trusted first-party ChatGPT origin. If no ChatGPT session provider is available, startup can still fall back to stored OAuth credentials.", + "enum": [ + "chatgpt" + ], + "type": "string" + } + ] + }, "McpServerEnvVar": { "anyOf": [ { @@ -2427,6 +2446,14 @@ }, "type": "array" }, + "auth": { + "allOf": [ + { + "$ref": "#/definitions/McpServerAuth" + } + ], + "default": null + }, "bearer_token_env_var": { "type": "string" }, @@ -2555,10 +2582,6 @@ }, "url": { "type": "string" - }, - "use_chatgpt_auth": { - "default": null, - "type": "boolean" } }, "type": "object" diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index e0c2f16638..a178dfae04 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -114,7 +114,7 @@ use tempfile::TempDir; fn stdio_mcp(command: &str) -> McpServerConfig { McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: command.to_string(), args: Vec::new(), @@ -141,7 +141,7 @@ fn stdio_mcp(command: &str) -> McpServerConfig { fn http_mcp(url: &str) -> McpServerConfig { McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: url.to_string(), bearer_token_env_var: None, @@ -5568,7 +5568,7 @@ async fn replace_mcp_servers_round_trips_entries() -> anyhow::Result<()> { servers.insert( "docs".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "echo".to_string(), args: vec!["hello".to_string()], @@ -5926,7 +5926,7 @@ async fn replace_mcp_servers_serializes_env_sorted() -> anyhow::Result<()> { let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "docs-server".to_string(), args: vec!["--verbose".to_string()], @@ -6006,7 +6006,7 @@ async fn replace_mcp_servers_serializes_env_vars() -> anyhow::Result<()> { let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "docs-server".to_string(), args: Vec::new(), @@ -6062,7 +6062,7 @@ async fn replace_mcp_servers_serializes_sourced_env_vars() -> anyhow::Result<()> let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "docs-server".to_string(), args: Vec::new(), @@ -6121,7 +6121,7 @@ async fn replace_mcp_servers_serializes_cwd() -> anyhow::Result<()> { let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "docs-server".to_string(), args: Vec::new(), @@ -6177,7 +6177,7 @@ async fn replace_mcp_servers_streamable_http_serializes_bearer_token() -> anyhow let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), bearer_token_env_var: Some("MCP_TOKEN".to_string()), @@ -6245,7 +6245,7 @@ async fn replace_mcp_servers_streamable_http_serializes_custom_headers() -> anyh let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), bearer_token_env_var: Some("MCP_TOKEN".to_string()), @@ -6328,7 +6328,7 @@ async fn replace_mcp_servers_streamable_http_removes_optional_sections() -> anyh let mut servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), bearer_token_env_var: Some("MCP_TOKEN".to_string()), @@ -6367,7 +6367,7 @@ async fn replace_mcp_servers_streamable_http_removes_optional_sections() -> anyh servers.insert( "docs".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), bearer_token_env_var: None, @@ -6435,7 +6435,7 @@ async fn replace_mcp_servers_streamable_http_isolates_headers_between_servers() ( "docs".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), bearer_token_env_var: Some("MCP_TOKEN".to_string()), @@ -6464,7 +6464,7 @@ async fn replace_mcp_servers_streamable_http_isolates_headers_between_servers() ( "logs".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "logs-server".to_string(), args: vec!["--follow".to_string()], @@ -6553,7 +6553,7 @@ async fn replace_mcp_servers_serializes_disabled_flag() -> anyhow::Result<()> { let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "docs-server".to_string(), args: Vec::new(), @@ -6604,7 +6604,7 @@ async fn replace_mcp_servers_serializes_required_flag() -> anyhow::Result<()> { let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "docs-server".to_string(), args: Vec::new(), @@ -6655,7 +6655,7 @@ async fn replace_mcp_servers_serializes_tool_filters() -> anyhow::Result<()> { let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "docs-server".to_string(), args: Vec::new(), @@ -6711,7 +6711,7 @@ async fn replace_mcp_servers_streamable_http_serializes_oauth_resource() -> anyh let servers = BTreeMap::from([( "docs".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com/mcp".to_string(), bearer_token_env_var: None, diff --git a/codex-rs/core/src/config/edit/document_helpers.rs b/codex-rs/core/src/config/edit/document_helpers.rs index 5f9445ad1c..308eb9922f 100644 --- a/codex-rs/core/src/config/edit/document_helpers.rs +++ b/codex-rs/core/src/config/edit/document_helpers.rs @@ -1,4 +1,5 @@ use codex_config::types::AppToolApproval; +use codex_config::types::McpServerAuth; use codex_config::types::McpServerConfig; use codex_config::types::McpServerEnvVar; use codex_config::types::McpServerToolConfig; @@ -95,8 +96,8 @@ fn serialize_mcp_server_table(config: &McpServerConfig) -> TomlTable { } } - if config.use_chatgpt_auth { - entry["use_chatgpt_auth"] = value(true); + if matches!(&config.auth, McpServerAuth::ChatGpt) { + entry["auth"] = value("chatgpt"); } if !config.enabled { entry["enabled"] = value(false); diff --git a/codex-rs/core/src/config/edit_tests.rs b/codex-rs/core/src/config/edit_tests.rs index 9c6cfe50cc..d47f128f94 100644 --- a/codex-rs/core/src/config/edit_tests.rs +++ b/codex-rs/core/src/config/edit_tests.rs @@ -873,7 +873,7 @@ fn blocking_replace_mcp_servers_round_trips() { servers.insert( "stdio".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "cmd".to_string(), args: vec!["--flag".to_string()], @@ -908,7 +908,7 @@ fn blocking_replace_mcp_servers_round_trips() { servers.insert( "http".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: "https://example.com".to_string(), bearer_token_env_var: Some("TOKEN".to_string()), @@ -983,7 +983,7 @@ fn blocking_replace_mcp_servers_serializes_tool_approval_overrides() { servers.insert( "docs".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "docs-server".to_string(), args: Vec::new(), @@ -1044,7 +1044,7 @@ foo = { command = "cmd" } servers.insert( "foo".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "cmd".to_string(), args: Vec::new(), @@ -1095,7 +1095,7 @@ foo = { command = "cmd" } # keep me servers.insert( "foo".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "cmd".to_string(), args: Vec::new(), @@ -1145,7 +1145,7 @@ foo = { command = "cmd", args = ["--flag"] } # keep me servers.insert( "foo".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "cmd".to_string(), args: Vec::new(), @@ -1196,7 +1196,7 @@ foo = { command = "cmd" } servers.insert( "foo".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "cmd".to_string(), args: Vec::new(), diff --git a/codex-rs/core/src/mcp_skill_dependencies.rs b/codex-rs/core/src/mcp_skill_dependencies.rs index 0574fd14f5..9b2b85aa2d 100644 --- a/codex-rs/core/src/mcp_skill_dependencies.rs +++ b/codex-rs/core/src/mcp_skill_dependencies.rs @@ -357,7 +357,7 @@ fn mcp_dependency_to_server_config( .as_ref() .ok_or_else(|| "missing url for streamable_http dependency".to_string())?; return Ok(McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::StreamableHttp { url: url.clone(), bearer_token_env_var: None, @@ -387,7 +387,7 @@ fn mcp_dependency_to_server_config( .as_ref() .ok_or_else(|| "missing command for stdio dependency".to_string())?; return Ok(McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: command.clone(), args: Vec::new(), diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 4db905e6b6..92b6c9d473 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -395,7 +395,7 @@ async fn run_code_mode_turn_with_rmcp_config( servers.insert( "rmcp".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, args: Vec::new(), diff --git a/codex-rs/core/tests/suite/hooks_mcp.rs b/codex-rs/core/tests/suite/hooks_mcp.rs index b3d13035cb..44c91a0d48 100644 --- a/codex-rs/core/tests/suite/hooks_mcp.rs +++ b/codex-rs/core/tests/suite/hooks_mcp.rs @@ -184,7 +184,7 @@ fn insert_rmcp_test_server(config: &mut Config, command: String, approval_mode: servers.insert( RMCP_SERVER.to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command, args: Vec::new(), diff --git a/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs b/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs index 8c523cb477..7df18727b9 100644 --- a/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs +++ b/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs @@ -30,7 +30,7 @@ async fn refresh_shuts_down_superseded_mcp_stdio_server() -> anyhow::Result<()> servers.insert( "refresh_cleanup".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command, args: Vec::new(), diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index ab53c553f2..a4d56f8cca 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -17,6 +17,7 @@ use std::time::Duration; use std::time::SystemTime; use std::time::UNIX_EPOCH; +use codex_config::types::McpServerAuth; use codex_config::types::McpServerConfig; use codex_config::types::McpServerEnvVar; use codex_config::types::McpServerTransportConfig; @@ -268,7 +269,7 @@ fn copy_binary_to_remote_env( struct TestMcpServerOptions { environment_id: String, - use_chatgpt_auth: bool, + auth: McpServerAuth, supports_parallel_tool_calls: bool, tool_timeout_sec: Option, } @@ -277,7 +278,7 @@ impl Default for TestMcpServerOptions { fn default() -> Self { Self { environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), - use_chatgpt_auth: false, + auth: McpServerAuth::default(), supports_parallel_tool_calls: false, tool_timeout_sec: None, } @@ -318,7 +319,7 @@ fn insert_mcp_server( server_name.to_string(), McpServerConfig { transport, - use_chatgpt_auth: options.use_chatgpt_auth, + auth: options.auth, environment_id: options.environment_id, enabled: true, required: false, @@ -1231,7 +1232,7 @@ async fn stdio_mcp_parallel_tool_calls_opt_in_runs_concurrently() -> anyhow::Res stdio_transport(rmcp_test_server_bin, /*env*/ None, Vec::new()), TestMcpServerOptions { environment_id: remote_aware_environment_id(), - use_chatgpt_auth: false, + auth: Default::default(), supports_parallel_tool_calls: true, tool_timeout_sec: Some(Duration::from_secs(2)), }, @@ -2304,62 +2305,20 @@ async fn streamable_http_tool_call_round_trip() -> anyhow::Result<()> { } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn streamable_http_chatgpt_auth_respects_configured_authorization() -> anyhow::Result<()> { +async fn streamable_http_configured_auth_precedes_chatgpt_auth() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; - let Some(chatgpt_auth_server) = - start_streamable_http_test_server("chatgpt-auth", Some("Access Token")).await? - else { - return Ok(()); - }; - let chatgpt_auth_url = chatgpt_auth_server.url().to_string(); - let chatgpt_base_url = chatgpt_auth_url - .strip_suffix("/mcp") - .expect("test MCP URL should end in /mcp") - .to_string(); - - let chatgpt_auth_fixture = test_codex() - .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) - .with_config(move |config| { - config.chatgpt_base_url = chatgpt_base_url; - insert_mcp_server( - config, - "chatgpt_auth", - McpServerTransportConfig::StreamableHttp { - url: chatgpt_auth_url, - bearer_token_env_var: None, - http_headers: None, - env_http_headers: None, - }, - TestMcpServerOptions { - environment_id: remote_aware_environment_id(), - use_chatgpt_auth: true, - ..Default::default() - }, - ); - }) - .build_with_remote_env(&server) - .await?; - wait_for_mcp_server(&chatgpt_auth_fixture.codex, "chatgpt_auth").await?; - drop(chatgpt_auth_fixture); - chatgpt_auth_server.shutdown().await; - let Some(configured_auth_server) = start_streamable_http_test_server("configured-auth", Some("configured-token")).await? else { return Ok(()); }; let configured_auth_url = configured_auth_server.url().to_string(); - let configured_auth_base_url = configured_auth_url - .strip_suffix("/mcp") - .expect("test MCP URL should end in /mcp") - .to_string(); let configured_auth_fixture = test_codex() .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) .with_config(move |config| { - config.chatgpt_base_url = configured_auth_base_url; insert_mcp_server( config, "configured_auth", @@ -2374,12 +2333,12 @@ async fn streamable_http_chatgpt_auth_respects_configured_authorization() -> any }, TestMcpServerOptions { environment_id: remote_aware_environment_id(), - use_chatgpt_auth: true, + auth: McpServerAuth::ChatGpt, ..Default::default() }, ); }) - .build_with_remote_env(&server) + .build_with_auto_env(&server) .await?; wait_for_mcp_server(&configured_auth_fixture.codex, "configured_auth").await?; @@ -2390,19 +2349,19 @@ async fn streamable_http_chatgpt_auth_respects_configured_authorization() -> any } #[tokio::test(flavor = "multi_thread", worker_threads = 1)] -async fn streamable_http_chatgpt_auth_is_not_sent_to_another_origin() -> anyhow::Result<()> { +async fn streamable_http_chatgpt_auth_is_not_sent_to_configured_origin() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; let untrusted_server = MockServer::start().await; let untrusted_apps = AppsTestServer::mount(&untrusted_server).await?; let untrusted_mcp_url = format!("{}/api/codex/apps", untrusted_apps.chatgpt_base_url); - let trusted_chatgpt_base_url = server.uri(); + let untrusted_chatgpt_base_url = untrusted_apps.chatgpt_base_url; let fixture = test_codex() .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) .with_config(move |config| { - config.chatgpt_base_url = trusted_chatgpt_base_url; + config.chatgpt_base_url = untrusted_chatgpt_base_url; insert_mcp_server( config, "untrusted_origin", @@ -2413,7 +2372,7 @@ async fn streamable_http_chatgpt_auth_is_not_sent_to_another_origin() -> anyhow: env_http_headers: None, }, TestMcpServerOptions { - use_chatgpt_auth: true, + auth: McpServerAuth::ChatGpt, ..Default::default() }, ); @@ -2452,6 +2411,67 @@ async fn streamable_http_chatgpt_auth_is_not_sent_to_another_origin() -> anyhow: Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn configured_chatgpt_base_url_does_not_grant_mcp_chatgpt_auth() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let untrusted_server = MockServer::start().await; + let untrusted_apps = AppsTestServer::mount(&untrusted_server).await?; + let untrusted_mcp_url = format!("{}/api/codex/apps", untrusted_apps.chatgpt_base_url); + let untrusted_chatgpt_base_url = untrusted_apps.chatgpt_base_url; + + let fixture = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_pre_build_hook(move |codex_home| { + fs::write( + codex_home.join("config.toml"), + format!( + r#" +chatgpt_base_url = "{untrusted_chatgpt_base_url}" + +[mcp_servers.untrusted_origin] +url = "{untrusted_mcp_url}" +auth = "chatgpt" +"#, + ), + ) + .expect("write attacker-controlled MCP config"); + }) + .build(&server) + .await?; + + wait_for_mcp_server(&fixture.codex, "untrusted_origin").await?; + let observed_requests = untrusted_server + .received_requests() + .await + .expect("mock server should capture MCP startup requests") + .into_iter() + .filter(|request| request.url.path() == "/api/codex/apps") + .filter_map(|request| { + let body: Value = serde_json::from_slice(&request.body).ok()?; + let method = body.get("method")?.as_str()?.to_string(); + let authorization = request + .headers + .get("authorization") + .and_then(|value| value.to_str().ok()) + .map(str::to_string); + Some((method, authorization)) + }) + .collect::>(); + + assert_eq!( + observed_requests, + vec![ + ("initialize".to_string(), None), + ("notifications/initialized".to_string(), None), + ("tools/list".to_string(), None), + ], + ); + + Ok(()) +} + /// This test writes to a fallback credentials file in CODEX_HOME. /// Ideally, we wouldn't need to serialize the test but it's much more cumbersome to wire CODEX_HOME through the code. #[test] diff --git a/codex-rs/core/tests/suite/search_tool.rs b/codex-rs/core/tests/suite/search_tool.rs index e0361d4770..edfc53fdc6 100644 --- a/codex-rs/core/tests/suite/search_tool.rs +++ b/codex-rs/core/tests/suite/search_tool.rs @@ -1090,7 +1090,7 @@ async fn tool_search_indexes_only_enabled_non_app_mcp_tools() -> Result<()> { servers.insert( "rmcp".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, args: Vec::new(), @@ -1217,7 +1217,7 @@ async fn tool_search_surfaced_mcp_tool_errors_are_returned_to_model() -> Result< servers.insert( "rmcp".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, args: Vec::new(), @@ -1366,7 +1366,7 @@ async fn tool_search_uses_non_app_mcp_server_instructions_as_namespace_descripti servers.insert( "rmcp".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, args: Vec::new(), diff --git a/codex-rs/core/tests/suite/sqlite_state.rs b/codex-rs/core/tests/suite/sqlite_state.rs index d9aefecf7f..5eea5a09e7 100644 --- a/codex-rs/core/tests/suite/sqlite_state.rs +++ b/codex-rs/core/tests/suite/sqlite_state.rs @@ -657,7 +657,7 @@ async fn mcp_call_marks_thread_memory_mode_polluted_when_configured() -> Result< servers.insert( server_name.to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, args: Vec::new(), diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index a55c182181..e536d6b888 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -276,7 +276,7 @@ async fn token_budget_context_injects_plain_thread_hint_text() -> Result<()> { servers.insert( "notes".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, args: Vec::new(), diff --git a/codex-rs/core/tests/suite/truncation.rs b/codex-rs/core/tests/suite/truncation.rs index 7080050977..c4b574abf5 100644 --- a/codex-rs/core/tests/suite/truncation.rs +++ b/codex-rs/core/tests/suite/truncation.rs @@ -383,7 +383,7 @@ async fn mcp_tool_call_output_exceeds_limit_truncated_for_model() -> Result<()> servers.insert( server_name.to_string(), codex_config::types::McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: codex_config::types::McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, args: Vec::new(), @@ -482,7 +482,7 @@ async fn mcp_image_output_preserves_image_and_no_text_summary() -> Result<()> { servers.insert( server_name.to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, args: Vec::new(), @@ -775,7 +775,7 @@ async fn mcp_tool_call_output_not_truncated_with_custom_limit() -> Result<()> { servers.insert( server_name.to_string(), codex_config::types::McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: codex_config::types::McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, args: Vec::new(), diff --git a/codex-rs/ext/mcp/src/executor_plugin/provider_tests.rs b/codex-rs/ext/mcp/src/executor_plugin/provider_tests.rs index cbc87b8182..22e5d1bec5 100644 --- a/codex-rs/ext/mcp/src/executor_plugin/provider_tests.rs +++ b/codex-rs/ext/mcp/src/executor_plugin/provider_tests.rs @@ -166,7 +166,7 @@ async fn reads_declared_config_only_through_executor_file_system() { vec![( "demo".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "demo-mcp".to_string(), args: Vec::new(), @@ -222,7 +222,7 @@ async fn reads_manifest_object_config_without_executor_file_system_access() { vec![( "counter".to_string(), McpServerConfig { - use_chatgpt_auth: false, + auth: Default::default(), transport: McpServerTransportConfig::Stdio { command: "counter-mcp".to_string(), args: Vec::new(), diff --git a/codex-rs/model-provider/src/lib.rs b/codex-rs/model-provider/src/lib.rs index f5efe1f575..44ce878e21 100644 --- a/codex-rs/model-provider/src/lib.rs +++ b/codex-rs/model-provider/src/lib.rs @@ -8,6 +8,7 @@ pub use auth::auth_provider_from_auth; pub use auth::unauthenticated_auth_provider; pub use bearer_auth_provider::BearerAuthProvider; pub use bearer_auth_provider::BearerAuthProvider as CoreAuthProvider; +pub use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; pub use codex_protocol::account::ProviderAccount; pub use provider::ModelProvider; pub use provider::ModelProviderFuture; From b3e1c33776466d6b393e5b8700c322cd9897b929 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Wed, 24 Jun 2026 20:03:22 -0700 Subject: [PATCH 6/7] code-mode: define process host wire protocol (#29804) ## Why The process-owned code mode implementation needs an explicit, bounded wire contract before either side depends on it. Keeping framing and message semantics in `codex-code-mode-protocol` gives the client and sidecar one shared source of truth and makes compatibility failures detectable during connection setup. ## What changed - adds a versioned client/host handshake with required and optional capabilities - defines operation requests and responses for session lifecycle and cell control - defines reverse delegate request, response, cancellation, and cell-closure messages - adds a four-byte little-endian length-prefixed JSON codec with a hard frame cap - rejects malformed frames, unknown fields, invalid identifiers, and unsupported protocol states - locks the wire representation down with explicit JSON round-trip tests ## Testing - `just test -p codex-code-mode-protocol` ## Stack Part 1 of 6. Followed by [#29805](https://github.com/openai/codex/pull/29805). --- codex-rs/code-mode-protocol/src/host/codec.rs | 101 +++ .../src/host/codec_tests.rs | 104 +++ .../code-mode-protocol/src/host/host_tests.rs | 720 +++++++++++++++--- .../code-mode-protocol/src/host/message.rs | 127 ++- codex-rs/code-mode-protocol/src/host/mod.rs | 37 +- .../code-mode-protocol/src/host/payload.rs | 412 ++++++++++ codex-rs/code-mode-protocol/src/host/types.rs | 22 + 7 files changed, 1424 insertions(+), 99 deletions(-) create mode 100644 codex-rs/code-mode-protocol/src/host/codec.rs create mode 100644 codex-rs/code-mode-protocol/src/host/codec_tests.rs create mode 100644 codex-rs/code-mode-protocol/src/host/payload.rs diff --git a/codex-rs/code-mode-protocol/src/host/codec.rs b/codex-rs/code-mode-protocol/src/host/codec.rs new file mode 100644 index 0000000000..6bf83d367d --- /dev/null +++ b/codex-rs/code-mode-protocol/src/host/codec.rs @@ -0,0 +1,101 @@ +use std::io; +use std::mem::size_of; + +use serde::Serialize; +use serde::de::DeserializeOwned; +use tokio::io::AsyncRead; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWrite; +use tokio::io::AsyncWriteExt; + +/// Maximum JSON payload size accepted for one IPC frame. +pub const MAX_FRAME_BYTES: usize = 64 * 1024 * 1024; + +/// Decodes JSON messages prefixed by a four-byte little-endian payload length. +pub struct FramedReader { + reader: R, +} + +impl FramedReader +where + R: AsyncRead + Unpin, +{ + pub fn new(reader: R) -> Self { + Self { reader } + } + + /// Reads the next frame, returning `None` only for EOF at a frame boundary. + pub async fn read(&mut self) -> io::Result> + where + T: DeserializeOwned, + { + let mut length_bytes = [0_u8; size_of::()]; + if self.reader.read(&mut length_bytes[..1]).await? == 0 { + return Ok(None); + } + self.reader.read_exact(&mut length_bytes[1..]).await?; + + let length = u32::from_le_bytes(length_bytes) as usize; + if length > MAX_FRAME_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("code-mode IPC frame length {length} exceeds {MAX_FRAME_BYTES} bytes"), + )); + } + + let mut payload = vec![0; length]; + self.reader.read_exact(&mut payload).await?; + serde_json::from_slice(&payload).map(Some).map_err(|err| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("failed to decode code-mode IPC frame: {err}"), + ) + }) + } +} + +/// Encodes JSON messages with a four-byte little-endian payload length. +pub struct FramedWriter { + writer: W, +} + +impl FramedWriter +where + W: AsyncWrite + Unpin, +{ + pub fn new(writer: W) -> Self { + Self { writer } + } + + /// Writes and flushes one complete frame. + pub async fn write(&mut self, message: &T) -> io::Result<()> + where + T: Serialize, + { + let payload = serde_json::to_vec(message).map_err(|err| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("failed to encode code-mode IPC frame: {err}"), + ) + })?; + if payload.len() > MAX_FRAME_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "code-mode IPC frame length {} exceeds {MAX_FRAME_BYTES} bytes", + payload.len() + ), + )); + } + let length = u32::try_from(payload.len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidData, + "code-mode IPC frame length exceeds u32", + ) + })?; + + self.writer.write_all(&length.to_le_bytes()).await?; + self.writer.write_all(&payload).await?; + self.writer.flush().await + } +} diff --git a/codex-rs/code-mode-protocol/src/host/codec_tests.rs b/codex-rs/code-mode-protocol/src/host/codec_tests.rs new file mode 100644 index 0000000000..996bd0679a --- /dev/null +++ b/codex-rs/code-mode-protocol/src/host/codec_tests.rs @@ -0,0 +1,104 @@ +use pretty_assertions::assert_eq; +use serde_json::json; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; + +use super::FramedReader; +use super::FramedWriter; +use super::MAX_FRAME_BYTES; + +#[tokio::test] +async fn frame_wire_format_is_little_endian_length_prefixed_json() { + let (writer, mut reader) = tokio::io::duplex(/*max_buf_size*/ 128); + let write = tokio::spawn(async move { + FramedWriter::new(writer) + .write(&json!({"value": 1})) + .await + .expect("write frame"); + }); + + let mut bytes = Vec::new(); + reader.read_to_end(&mut bytes).await.expect("read bytes"); + write.await.expect("writer task"); + + let payload = br#"{"value":1}"#; + let mut expected = (payload.len() as u32).to_le_bytes().to_vec(); + expected.extend_from_slice(payload); + assert_eq!(bytes, expected); +} + +#[tokio::test] +async fn fragmented_frame_round_trips() { + let value = json!({"type": "session/open", "sessionId": "session-1"}); + let payload = serde_json::to_vec(&value).expect("serialize"); + let mut bytes = (payload.len() as u32).to_le_bytes().to_vec(); + bytes.extend(payload); + + let (mut writer, reader) = tokio::io::duplex(/*max_buf_size*/ 128); + let write = tokio::spawn(async move { + for byte in bytes { + writer.write_all(&[byte]).await.expect("write byte"); + tokio::task::yield_now().await; + } + }); + + assert_eq!( + FramedReader::new(reader) + .read::() + .await + .expect("read frame"), + Some(value) + ); + write.await.expect("writer task"); +} + +#[tokio::test] +async fn eof_is_clean_only_at_a_frame_boundary() { + let (writer, reader) = tokio::io::duplex(/*max_buf_size*/ 16); + drop(writer); + assert_eq!( + FramedReader::new(reader) + .read::() + .await + .expect("clean eof"), + None + ); + + let (mut writer, reader) = tokio::io::duplex(/*max_buf_size*/ 16); + writer + .write_all(&[1, 0]) + .await + .expect("write partial header"); + drop(writer); + let err = FramedReader::new(reader) + .read::() + .await + .expect_err("truncated header"); + assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof); +} + +#[tokio::test] +async fn oversized_and_malformed_frames_are_rejected() { + let (mut writer, reader) = tokio::io::duplex(/*max_buf_size*/ 16); + writer + .write_all(&((MAX_FRAME_BYTES as u32) + 1).to_le_bytes()) + .await + .expect("write oversized header"); + let err = FramedReader::new(reader) + .read::() + .await + .expect_err("oversized frame"); + assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); + + let (mut writer, reader) = tokio::io::duplex(/*max_buf_size*/ 16); + writer + .write_all(&(1_u32).to_le_bytes()) + .await + .expect("write length"); + writer.write_all(b"{").await.expect("write malformed json"); + let err = FramedReader::new(reader) + .read::() + .await + .expect_err("malformed frame"); + assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); +} diff --git a/codex-rs/code-mode-protocol/src/host/host_tests.rs b/codex-rs/code-mode-protocol/src/host/host_tests.rs index 2d64993440..3aa2a18307 100644 --- a/codex-rs/code-mode-protocol/src/host/host_tests.rs +++ b/codex-rs/code-mode-protocol/src/host/host_tests.rs @@ -1,21 +1,57 @@ +use std::fmt::Debug; + use pretty_assertions::assert_eq; +use serde::Serialize; +use serde::de::DeserializeOwned; +use serde_json::Value; use serde_json::json; use super::Capability; use super::CapabilitySet; use super::ClientHello; use super::ClientToHost; +use super::DelegateRequest; +use super::DelegateRequestId; +use super::DelegateResponse; use super::HandshakeRejectReason; use super::HostHello; +use super::HostRequest; +use super::HostResponse; use super::HostToClient; use super::ProtocolVersion; +use super::RequestId; use super::SessionId; use super::SupportedProtocolVersions; +use super::WireCellId; +use super::WireContentItem; +use super::WireExecuteRequest; +use super::WireImageDetail; +use super::WireNestedToolCall; +use super::WireResult; +use super::WireRuntimeResponse; +use super::WireToolDefinition; +use super::WireToolKind; +use super::WireToolName; +use super::WireWaitOutcome; +use super::WireWaitRequest; +use crate::ExecuteRequest; fn session_id() -> SessionId { SessionId::new("session-1").expect("valid session ID") } +fn cell_id(value: &str) -> WireCellId { + WireCellId::new(value) +} + +fn request_id(value: i64) -> RequestId { + RequestId::new(value) +} + +fn delegate_request_id(value: i64) -> DelegateRequestId { + DelegateRequestId::new(value) +} + fn capability(value: &str) -> Capability { Capability::new(value).expect("valid capability") } @@ -25,116 +61,539 @@ fn supported_versions() -> SupportedProtocolVersions { .expect("nonempty unique protocol versions") } -#[test] -fn handshake_wire_contract_is_explicit_and_round_trips() { - let client_hello = ClientToHost::ClientHello( - ClientHello::new( - supported_versions(), - CapabilitySet::try_new([capability("required")]).expect("valid required set"), - CapabilitySet::try_new([capability("optional")]).expect("valid optional set"), - ) - .expect("disjoint capabilities"), - ); - let client_hello_json = json!({ - "type": "connection/hello", - "supportedVersions": [1], - "requiredCapabilities": ["required"], - "optionalCapabilities": ["optional"], - }); +fn assert_wire_round_trip(message: T, encoded: Value) +where + T: Debug + DeserializeOwned + PartialEq + Serialize, +{ + assert_eq!(serde_json::to_value(&message).expect("serialize"), encoded); assert_eq!( - serde_json::to_value(&client_hello).expect("serialize"), - client_hello_json - ); - assert_eq!( - serde_json::from_value::(client_hello_json).expect("deserialize"), - client_hello - ); - - let host_hello = HostToClient::HostHello(HostHello::new( - ProtocolVersion::V1, - CapabilitySet::try_new([capability("required")]).expect("valid capabilities"), - )); - let host_hello_json = json!({ - "type": "connection/ready", - "selectedVersion": 1, - "capabilities": ["required"], - }); - assert_eq!( - serde_json::to_value(&host_hello).expect("serialize"), - host_hello_json - ); - assert_eq!( - serde_json::from_value::(host_hello_json).expect("deserialize"), - host_hello - ); - - let rejected = HostToClient::HandshakeRejected { - reason: HandshakeRejectReason::NoCompatibleVersion { - supported_versions: supported_versions(), - }, - }; - let rejected_json = json!({ - "type": "connection/rejected", - "reason": { - "type": "noCompatibleVersion", - "supportedVersions": [1], - }, - }); - assert_eq!( - serde_json::to_value(&rejected).expect("serialize"), - rejected_json - ); - assert_eq!( - serde_json::from_value::(rejected_json).expect("deserialize"), - rejected + serde_json::from_value::(encoded).expect("deserialize"), + message ); } +fn execute_request() -> WireExecuteRequest { + WireExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: vec![ + WireToolDefinition { + name: "function_tool".to_string(), + tool_name: WireToolName { + name: "function_tool".to_string(), + namespace: None, + }, + description: "function tool".to_string(), + kind: WireToolKind::Function, + input_schema: Some(json!({ "type": "object" })), + output_schema: None, + }, + WireToolDefinition { + name: "freeform_tool".to_string(), + tool_name: WireToolName { + name: "freeform_tool".to_string(), + namespace: Some("mcp__sample__".to_string()), + }, + description: "freeform tool".to_string(), + kind: WireToolKind::Freeform, + input_schema: None, + output_schema: Some(json!({ "type": "string" })), + }, + ], + source: "text('hello');".to_string(), + yield_time_ms: Some(25), + max_output_tokens: Some(100), + } +} + +fn content_items() -> Vec { + vec![ + WireContentItem::InputText { + text: "hello".to_string(), + }, + WireContentItem::InputImage { + image_url: "data:image/png;base64,none".to_string(), + detail: None, + }, + WireContentItem::InputImage { + image_url: "data:image/png;base64,auto".to_string(), + detail: Some(WireImageDetail::Auto), + }, + WireContentItem::InputImage { + image_url: "data:image/png;base64,low".to_string(), + detail: Some(WireImageDetail::Low), + }, + WireContentItem::InputImage { + image_url: "data:image/png;base64,high".to_string(), + detail: Some(WireImageDetail::High), + }, + WireContentItem::InputImage { + image_url: "data:image/png;base64,original".to_string(), + detail: Some(WireImageDetail::Original), + }, + ] +} + +fn content_items_json() -> Value { + json!([ + { "type": "input_text", "text": "hello" }, + { "type": "input_image", "image_url": "data:image/png;base64,none" }, + { + "type": "input_image", + "image_url": "data:image/png;base64,auto", + "detail": "auto", + }, + { + "type": "input_image", + "image_url": "data:image/png;base64,low", + "detail": "low", + }, + { + "type": "input_image", + "image_url": "data:image/png;base64,high", + "detail": "high", + }, + { + "type": "input_image", + "image_url": "data:image/png;base64,original", + "detail": "original", + }, + ]) +} + #[test] -fn session_lifecycle_wire_contract_is_explicit_and_round_trips() { - let client_messages = [ +fn handshake_v1_variants_are_pinned() { + assert_wire_round_trip( + ClientToHost::ClientHello( + ClientHello::new( + supported_versions(), + CapabilitySet::try_new([capability("required")]).expect("valid required set"), + CapabilitySet::try_new([capability("optional")]).expect("valid optional set"), + ) + .expect("disjoint capabilities"), + ), + json!({ + "type": "connection/hello", + "supportedVersions": [1], + "requiredCapabilities": ["required"], + "optionalCapabilities": ["optional"], + }), + ); + assert_wire_round_trip( + HostToClient::HostHello(HostHello::new( + ProtocolVersion::V1, + CapabilitySet::try_new([capability("required")]).expect("valid capabilities"), + )), + json!({ + "type": "connection/ready", + "selectedVersion": 1, + "capabilities": ["required"], + }), + ); + for (reason, encoded) in [ ( - ClientToHost::OpenSession { - session_id: session_id(), + HandshakeRejectReason::NoCompatibleVersion { + supported_versions: supported_versions(), }, - json!({ "type": "session/open", "sessionId": "session-1" }), + json!({ + "type": "connection/rejected", + "reason": { + "type": "noCompatibleVersion", + "supportedVersions": [1], + }, + }), ), ( - ClientToHost::CloseSession { + HandshakeRejectReason::MissingRequiredCapability { + capability: capability("required"), + }, + json!({ + "type": "connection/rejected", + "reason": { + "type": "missingRequiredCapability", + "capability": "required", + }, + }), + ), + ( + HandshakeRejectReason::InvalidHello { + message: "invalid hello".to_string(), + }, + json!({ + "type": "connection/rejected", + "reason": { + "type": "invalidHello", + "message": "invalid hello", + }, + }), + ), + ] { + assert_wire_round_trip(HostToClient::HandshakeRejected { reason }, encoded); + } +} + +#[test] +fn client_to_host_v1_variants_are_pinned() { + let execute_request = execute_request(); + for (id, request, encoded_request) in [ + ( + request_id(/*value*/ 1), + HostRequest::OpenSession { session_id: session_id(), }, - json!({ "type": "session/close", "sessionId": "session-1" }), + json!({ "method": "session/open", "sessionId": "session-1" }), ), - ]; - for (message, encoded) in client_messages { - assert_eq!(serde_json::to_value(&message).expect("serialize"), encoded); - assert_eq!( - serde_json::from_value::(encoded).expect("deserialize"), - message + ( + request_id(/*value*/ 2), + HostRequest::Execute { + session_id: session_id(), + request: execute_request, + }, + json!({ + "method": "session/execute", + "sessionId": "session-1", + "request": { + "tool_call_id": "call-1", + "enabled_tools": [ + { + "name": "function_tool", + "tool_name": { "name": "function_tool", "namespace": null }, + "description": "function tool", + "kind": "function", + "input_schema": { "type": "object" }, + "output_schema": null, + }, + { + "name": "freeform_tool", + "tool_name": { + "name": "freeform_tool", + "namespace": "mcp__sample__", + }, + "description": "freeform tool", + "kind": "freeform", + "input_schema": null, + "output_schema": { "type": "string" }, + }, + ], + "source": "text('hello');", + "yield_time_ms": 25, + "max_output_tokens": 100, + }, + }), + ), + ( + request_id(/*value*/ 3), + HostRequest::Wait { + session_id: session_id(), + request: WireWaitRequest { + cell_id: cell_id("cell-1"), + yield_time_ms: 50, + }, + }, + json!({ + "method": "session/wait", + "sessionId": "session-1", + "request": { "cell_id": "cell-1", "yield_time_ms": 50 }, + }), + ), + ( + request_id(/*value*/ 4), + HostRequest::Terminate { + session_id: session_id(), + cell_id: cell_id("cell-1"), + }, + json!({ + "method": "session/terminate", + "sessionId": "session-1", + "cellId": "cell-1", + }), + ), + ( + request_id(/*value*/ 5), + HostRequest::ShutdownSession { + session_id: session_id(), + }, + json!({ "method": "session/shutdown", "sessionId": "session-1" }), + ), + ] { + assert_wire_round_trip( + ClientToHost::Request { id, request }, + json!({ + "type": "operation/request", + "id": id, + "request": encoded_request, + }), ); } - let host_messages = [ + for (id, result, encoded_result) in [ ( - HostToClient::SessionReady { + delegate_request_id(/*value*/ 6), + WireResult::Ok { + value: DelegateResponse::ToolResult { + result: json!({ "answer": 42 }), + }, + }, + json!({ + "status": "ok", + "value": { "type": "tool/result", "result": { "answer": 42 } }, + }), + ), + ( + delegate_request_id(/*value*/ 7), + WireResult::Ok { + value: DelegateResponse::NotificationDelivered, + }, + json!({ + "status": "ok", + "value": { "type": "notification/delivered" }, + }), + ), + ( + delegate_request_id(/*value*/ 8), + WireResult::Err { + message: "delegate failed".to_string(), + }, + json!({ "status": "error", "message": "delegate failed" }), + ), + ] { + assert_wire_round_trip( + ClientToHost::DelegateResponse { id, result }, + json!({ + "type": "delegate/response", + "id": id, + "result": encoded_result, + }), + ); + } +} + +#[test] +fn host_to_client_v1_variants_are_pinned() { + for (id, response, encoded_response) in [ + ( + request_id(/*value*/ 1), + HostResponse::SessionReady { session_id: session_id(), }, json!({ "type": "session/ready", "sessionId": "session-1" }), ), ( - HostToClient::SessionClosed { + request_id(/*value*/ 2), + HostResponse::ExecutionStarted { + cell_id: cell_id("cell-1"), + }, + json!({ "type": "execution/started", "cellId": "cell-1" }), + ), + ( + request_id(/*value*/ 3), + HostResponse::WaitCompleted { + outcome: WireWaitOutcome::LiveCell(WireRuntimeResponse::Yielded { + cell_id: cell_id("cell-1"), + content_items: content_items(), + }), + }, + json!({ + "type": "wait/completed", + "outcome": { + "LiveCell": { + "Yielded": { + "cell_id": "cell-1", + "content_items": content_items_json(), + }, + }, + }, + }), + ), + ( + request_id(/*value*/ 4), + HostResponse::WaitCompleted { + outcome: WireWaitOutcome::MissingCell(WireRuntimeResponse::Result { + cell_id: cell_id("missing-cell"), + content_items: Vec::new(), + error_text: Some("cell not found".to_string()), + }), + }, + json!({ + "type": "wait/completed", + "outcome": { + "MissingCell": { + "Result": { + "cell_id": "missing-cell", + "content_items": [], + "error_text": "cell not found", + }, + }, + }, + }), + ), + ( + request_id(/*value*/ 5), + HostResponse::SessionClosed { session_id: session_id(), }, json!({ "type": "session/closed", "sessionId": "session-1" }), ), - ]; - for (message, encoded) in host_messages { - assert_eq!(serde_json::to_value(&message).expect("serialize"), encoded); - assert_eq!( - serde_json::from_value::(encoded).expect("deserialize"), - message + ] { + assert_wire_round_trip( + HostToClient::Response { + id, + result: WireResult::Ok { value: response }, + }, + json!({ + "type": "operation/response", + "id": id, + "result": { "status": "ok", "value": encoded_response }, + }), ); } + assert_wire_round_trip( + HostToClient::Response { + id: request_id(/*value*/ 6), + result: WireResult::Err { + message: "operation failed".to_string(), + }, + }, + json!({ + "type": "operation/response", + "id": 6, + "result": { "status": "error", "message": "operation failed" }, + }), + ); + + assert_wire_round_trip( + HostToClient::InitialResponse { + id: request_id(/*value*/ 7), + result: WireResult::Ok { + value: WireRuntimeResponse::Terminated { + cell_id: cell_id("cell-1"), + content_items: Vec::new(), + }, + }, + }, + json!({ + "type": "execute/initialResponse", + "id": 7, + "result": { + "status": "ok", + "value": { + "Terminated": { "cell_id": "cell-1", "content_items": [] }, + }, + }, + }), + ); + assert_wire_round_trip( + HostToClient::InitialResponse { + id: request_id(/*value*/ 8), + result: WireResult::Err { + message: "execution failed".to_string(), + }, + }, + json!({ + "type": "execute/initialResponse", + "id": 8, + "result": { "status": "error", "message": "execution failed" }, + }), + ); + + assert_wire_round_trip( + HostToClient::DelegateRequest { + id: delegate_request_id(/*value*/ 9), + session_id: session_id(), + request: DelegateRequest::InvokeTool { + invocation: WireNestedToolCall { + cell_id: cell_id("cell-1"), + runtime_tool_call_id: "runtime-call-1".to_string(), + tool_name: WireToolName { + name: "freeform_tool".to_string(), + namespace: Some("mcp__sample__".to_string()), + }, + tool_kind: WireToolKind::Freeform, + input: Some(json!({ "value": 1 })), + }, + }, + }, + json!({ + "type": "delegate/request", + "id": 9, + "sessionId": "session-1", + "request": { + "type": "tool/invoke", + "invocation": { + "cell_id": "cell-1", + "runtime_tool_call_id": "runtime-call-1", + "tool_name": { + "name": "freeform_tool", + "namespace": "mcp__sample__", + }, + "tool_kind": "freeform", + "input": { "value": 1 }, + }, + }, + }), + ); + assert_wire_round_trip( + HostToClient::DelegateRequest { + id: delegate_request_id(/*value*/ 10), + session_id: session_id(), + request: DelegateRequest::Notify { + call_id: "call-1".to_string(), + cell_id: cell_id("cell-1"), + text: "important".to_string(), + }, + }, + json!({ + "type": "delegate/request", + "id": 10, + "sessionId": "session-1", + "request": { + "type": "notification/send", + "callId": "call-1", + "cellId": "cell-1", + "text": "important", + }, + }), + ); + assert_wire_round_trip( + HostToClient::CancelDelegateRequest { + id: delegate_request_id(/*value*/ 11), + }, + json!({ "type": "delegate/cancel", "id": 11 }), + ); + assert_wire_round_trip( + HostToClient::CellClosed { + session_id: session_id(), + cell_id: cell_id("cell-1"), + }, + json!({ + "type": "cell/closed", + "sessionId": "session-1", + "cellId": "cell-1", + }), + ); +} + +#[test] +fn execute_request_integer_bounds_are_enforced() { + let wire_request = execute_request(); + let domain_request = ExecuteRequest::try_from(wire_request.clone()) + .expect("valid wire request converts to the domain"); + assert_eq!( + WireExecuteRequest::try_from(domain_request.clone()) + .expect("valid domain request converts to the wire"), + wire_request + ); + + let too_large = ExecuteRequest { + max_output_tokens: Some(usize::try_from(i32::MAX).expect("i32::MAX fits usize") + 1), + ..domain_request + }; + assert!(WireExecuteRequest::try_from(too_large).is_err()); + + let negative = WireExecuteRequest { + max_output_tokens: Some(-1), + ..wire_request + }; + assert!(ExecuteRequest::try_from(negative).is_err()); } #[test] @@ -168,7 +627,11 @@ fn invalid_protocol_states_cannot_be_constructed_or_decoded() { ); for invalid in [ - json!({ "type": "session/open", "sessionId": "" }), + json!({ + "type": "operation/request", + "id": 1, + "request": { "method": "session/open", "sessionId": "" }, + }), json!({ "type": "connection/hello", "supportedVersions": [], @@ -187,19 +650,100 @@ fn invalid_protocol_states_cannot_be_constructed_or_decoded() { } #[test] -fn unknown_fields_are_rejected() { +fn every_nested_v1_object_rejects_unknown_fields() { assert!( serde_json::from_value::(json!({ - "type": "session/open", + "type": "operation/request", + "id": 1, + "request": { "method": "session/open", "sessionId": "session-1" }, + "unexpected": true, + })) + .is_err() + ); + assert!( + serde_json::from_value::(json!({ + "method": "session/open", "sessionId": "session-1", "unexpected": true, })) .is_err() ); + assert!( + serde_json::from_value::(json!({ + "tool_call_id": "call-1", + "enabled_tools": [], + "source": "text('hello');", + "yield_time_ms": null, + "max_output_tokens": null, + "unexpected": true, + })) + .is_err() + ); + assert!( + serde_json::from_value::(json!({ + "name": "tool", + "tool_name": { "name": "tool", "namespace": null }, + "description": "tool", + "kind": "function", + "input_schema": null, + "output_schema": null, + "unexpected": true, + })) + .is_err() + ); + assert!( + serde_json::from_value::(json!({ + "name": "tool", + "namespace": null, + "unexpected": true, + })) + .is_err() + ); + assert!( + serde_json::from_value::(json!({ + "cell_id": "cell-1", + "yield_time_ms": 50, + "unexpected": true, + })) + .is_err() + ); + assert!( + serde_json::from_value::(json!({ + "Yielded": { + "cell_id": "cell-1", + "content_items": [], + "unexpected": true, + }, + })) + .is_err() + ); + assert!( + serde_json::from_value::(json!({ + "type": "input_text", + "text": "hello", + "unexpected": true, + })) + .is_err() + ); + assert!( + serde_json::from_value::(json!({ + "cell_id": "cell-1", + "runtime_tool_call_id": "runtime-call-1", + "tool_name": { "name": "tool", "namespace": null }, + "tool_kind": "function", + "input": null, + "unexpected": true, + })) + .is_err() + ); assert!( serde_json::from_value::(json!({ - "type": "session/ready", - "sessionId": "session-1", + "type": "operation/response", + "id": 1, + "result": { + "status": "ok", + "value": { "type": "session/ready", "sessionId": "session-1" }, + }, "unexpected": true, })) .is_err() diff --git a/codex-rs/code-mode-protocol/src/host/message.rs b/codex-rs/code-mode-protocol/src/host/message.rs index ed30ac8f17..78f0961d51 100644 --- a/codex-rs/code-mode-protocol/src/host/message.rs +++ b/codex-rs/code-mode-protocol/src/host/message.rs @@ -2,13 +2,22 @@ use std::fmt; use serde::Deserialize; use serde::Serialize; +use serde_json::Value as JsonValue; use super::Capability; use super::CapabilitySet; +use super::DelegateRequestId; use super::HandshakeRejectReason; use super::ProtocolVersion; +use super::RequestId; use super::SessionId; use super::SupportedProtocolVersions; +use super::WireCellId; +use super::WireExecuteRequest; +use super::WireNestedToolCall; +use super::WireRuntimeResponse; +use super::WireWaitOutcome; +use super::WireWaitRequest; #[derive(Clone, Debug, PartialEq, Serialize)] #[serde(rename_all = "camelCase")] @@ -116,27 +125,133 @@ impl HostHello { } /// Messages sent from a client to the code-mode host. -#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[derive(Debug, Deserialize, PartialEq, Serialize)] #[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")] pub enum ClientToHost { #[serde(rename = "connection/hello")] ClientHello(ClientHello), - #[serde(rename = "session/open")] - OpenSession { session_id: SessionId }, - #[serde(rename = "session/close")] - CloseSession { session_id: SessionId }, + #[serde(rename = "operation/request")] + Request { id: RequestId, request: HostRequest }, + #[serde(rename = "delegate/response")] + DelegateResponse { + id: DelegateRequestId, + result: WireResult, + }, } /// Messages sent from the code-mode host to a client. -#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[derive(Debug, Deserialize, PartialEq, Serialize)] #[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")] pub enum HostToClient { #[serde(rename = "connection/ready")] HostHello(HostHello), #[serde(rename = "connection/rejected")] HandshakeRejected { reason: HandshakeRejectReason }, + #[serde(rename = "operation/response")] + Response { + id: RequestId, + result: WireResult, + }, + #[serde(rename = "execute/initialResponse")] + InitialResponse { + id: RequestId, + result: WireResult, + }, + #[serde(rename = "delegate/request")] + DelegateRequest { + id: DelegateRequestId, + session_id: SessionId, + request: DelegateRequest, + }, + #[serde(rename = "delegate/cancel")] + CancelDelegateRequest { id: DelegateRequestId }, + #[serde(rename = "cell/closed")] + CellClosed { + session_id: SessionId, + cell_id: WireCellId, + }, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields, tag = "method", rename_all_fields = "camelCase")] +pub enum HostRequest { + #[serde(rename = "session/open")] + OpenSession { session_id: SessionId }, + #[serde(rename = "session/execute")] + Execute { + session_id: SessionId, + request: WireExecuteRequest, + }, + #[serde(rename = "session/wait")] + Wait { + session_id: SessionId, + request: WireWaitRequest, + }, + #[serde(rename = "session/terminate")] + Terminate { + session_id: SessionId, + cell_id: WireCellId, + }, + #[serde(rename = "session/shutdown")] + ShutdownSession { session_id: SessionId }, +} + +#[derive(Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")] +pub enum HostResponse { #[serde(rename = "session/ready")] SessionReady { session_id: SessionId }, + #[serde(rename = "execution/started")] + ExecutionStarted { cell_id: WireCellId }, + #[serde(rename = "wait/completed")] + WaitCompleted { outcome: WireWaitOutcome }, #[serde(rename = "session/closed")] SessionClosed { session_id: SessionId }, } + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")] +pub enum DelegateRequest { + #[serde(rename = "tool/invoke")] + InvokeTool { invocation: WireNestedToolCall }, + #[serde(rename = "notification/send")] + Notify { + call_id: String, + cell_id: WireCellId, + text: String, + }, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")] +pub enum DelegateResponse { + #[serde(rename = "tool/result")] + ToolResult { result: JsonValue }, + #[serde(rename = "notification/delivered")] + NotificationDelivered, +} + +#[derive(Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields, tag = "status", rename_all_fields = "camelCase")] +pub enum WireResult { + #[serde(rename = "ok")] + Ok { value: T }, + #[serde(rename = "error")] + Err { message: String }, +} + +impl WireResult { + pub fn from_result(result: Result) -> Self { + match result { + Ok(value) => Self::Ok { value }, + Err(message) => Self::Err { message }, + } + } + + pub fn into_result(self) -> Result { + match self { + Self::Ok { value } => Ok(value), + Self::Err { message } => Err(message), + } + } +} diff --git a/codex-rs/code-mode-protocol/src/host/mod.rs b/codex-rs/code-mode-protocol/src/host/mod.rs index f705f80e36..7e78768ab6 100644 --- a/codex-rs/code-mode-protocol/src/host/mod.rs +++ b/codex-rs/code-mode-protocol/src/host/mod.rs @@ -1,29 +1,56 @@ -//! Transport-neutral messages for the callback-only code-mode host boundary. +//! Messages and local IPC framing for the code-mode host boundary. //! -//! Protocol version 1 relies on ordered framing and connection-scoped -//! fail-stop behavior rather than message sequence numbers. It defines no -//! optional capabilities yet; capability names provide an extension point for -//! later versions without weakening the v1 decoder. +//! Protocol version 1 multiplexes session operations and delegate callbacks by +//! request ID over one ordered connection. It defines no optional capabilities +//! yet; capability names provide an extension point for later versions without +//! weakening the v1 decoder. +mod codec; mod error; mod message; +mod payload; mod types; +pub use codec::FramedReader; +pub use codec::FramedWriter; +pub use codec::MAX_FRAME_BYTES; pub use error::HandshakeRejectReason; pub use message::ClientHello; pub use message::ClientHelloError; pub use message::ClientToHost; +pub use message::DelegateRequest; +pub use message::DelegateResponse; pub use message::HostHello; +pub use message::HostRequest; +pub use message::HostResponse; pub use message::HostToClient; +pub use message::WireResult; +pub use payload::WireCellId; +pub use payload::WireContentItem; +pub use payload::WireExecuteRequest; +pub use payload::WireImageDetail; +pub use payload::WireNestedToolCall; +pub use payload::WireRuntimeResponse; +pub use payload::WireToolDefinition; +pub use payload::WireToolKind; +pub use payload::WireToolName; +pub use payload::WireWaitOutcome; +pub use payload::WireWaitRequest; pub use types::Capability; pub use types::CapabilitySet; +pub use types::DelegateRequestId; pub use types::DuplicateCapability; pub use types::InvalidIdentifier; pub use types::InvalidSupportedProtocolVersions; pub use types::ProtocolVersion; +pub use types::RequestId; pub use types::SessionId; pub use types::SupportedProtocolVersions; #[cfg(test)] #[path = "host_tests.rs"] mod tests; + +#[cfg(test)] +#[path = "codec_tests.rs"] +mod codec_tests; diff --git a/codex-rs/code-mode-protocol/src/host/payload.rs b/codex-rs/code-mode-protocol/src/host/payload.rs new file mode 100644 index 0000000000..a5ba410711 --- /dev/null +++ b/codex-rs/code-mode-protocol/src/host/payload.rs @@ -0,0 +1,412 @@ +use std::num::TryFromIntError; + +use codex_protocol::ToolName; +use serde::Deserialize; +use serde::Serialize; +use serde_json::Value as JsonValue; + +use crate::CellId; +use crate::CodeModeNestedToolCall; +use crate::CodeModeToolKind; +use crate::ExecuteRequest; +use crate::FunctionCallOutputContentItem; +use crate::ImageDetail; +use crate::RuntimeResponse; +use crate::ToolDefinition; +use crate::WaitOutcome; +use crate::WaitRequest; + +/// A cell identifier with a wire representation owned by protocol V1. +#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)] +#[serde(transparent)] +pub struct WireCellId(String); + +impl WireCellId { + pub fn new(value: impl Into) -> Self { + Self(value.into()) + } + + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl From for WireCellId { + fn from(value: CellId) -> Self { + Self(value.as_str().to_string()) + } +} + +impl From<&CellId> for WireCellId { + fn from(value: &CellId) -> Self { + Self(value.as_str().to_string()) + } +} + +impl From for CellId { + fn from(value: WireCellId) -> Self { + Self::new(value.0) + } +} + +/// The V1 wire representation of a tool's stable name. +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct WireToolName { + pub name: String, + pub namespace: Option, +} + +impl From for WireToolName { + fn from(value: ToolName) -> Self { + Self { + name: value.name, + namespace: value.namespace, + } + } +} + +impl From for ToolName { + fn from(value: WireToolName) -> Self { + Self::new(value.namespace, value.name) + } +} + +/// The tool invocation shape supported by protocol V1. +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum WireToolKind { + Function, + Freeform, +} + +impl From for WireToolKind { + fn from(value: CodeModeToolKind) -> Self { + match value { + CodeModeToolKind::Function => Self::Function, + CodeModeToolKind::Freeform => Self::Freeform, + } + } +} + +impl From for CodeModeToolKind { + fn from(value: WireToolKind) -> Self { + match value { + WireToolKind::Function => Self::Function, + WireToolKind::Freeform => Self::Freeform, + } + } +} + +/// A V1 tool definition embedded in an execute request. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct WireToolDefinition { + pub name: String, + pub tool_name: WireToolName, + pub description: String, + pub kind: WireToolKind, + pub input_schema: Option, + pub output_schema: Option, +} + +impl From for WireToolDefinition { + fn from(value: ToolDefinition) -> Self { + Self { + name: value.name, + tool_name: value.tool_name.into(), + description: value.description, + kind: value.kind.into(), + input_schema: value.input_schema, + output_schema: value.output_schema, + } + } +} + +impl From for ToolDefinition { + fn from(value: WireToolDefinition) -> Self { + Self { + name: value.name, + tool_name: value.tool_name.into(), + description: value.description, + kind: value.kind.into(), + input_schema: value.input_schema, + output_schema: value.output_schema, + } + } +} + +/// The complete execute request shape supported by protocol V1. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct WireExecuteRequest { + pub tool_call_id: String, + pub enabled_tools: Vec, + pub source: String, + pub yield_time_ms: Option, + pub max_output_tokens: Option, +} + +impl TryFrom for WireExecuteRequest { + type Error = TryFromIntError; + + fn try_from(value: ExecuteRequest) -> Result { + Ok(Self { + tool_call_id: value.tool_call_id, + enabled_tools: value.enabled_tools.into_iter().map(Into::into).collect(), + source: value.source, + yield_time_ms: value.yield_time_ms, + max_output_tokens: value.max_output_tokens.map(i32::try_from).transpose()?, + }) + } +} + +impl TryFrom for ExecuteRequest { + type Error = TryFromIntError; + + fn try_from(value: WireExecuteRequest) -> Result { + Ok(Self { + tool_call_id: value.tool_call_id, + enabled_tools: value.enabled_tools.into_iter().map(Into::into).collect(), + source: value.source, + yield_time_ms: value.yield_time_ms, + max_output_tokens: value.max_output_tokens.map(usize::try_from).transpose()?, + }) + } +} + +/// The complete wait request shape supported by protocol V1. +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct WireWaitRequest { + pub cell_id: WireCellId, + pub yield_time_ms: u64, +} + +impl From for WireWaitRequest { + fn from(value: WaitRequest) -> Self { + Self { + cell_id: value.cell_id.into(), + yield_time_ms: value.yield_time_ms, + } + } +} + +impl From for WaitRequest { + fn from(value: WireWaitRequest) -> Self { + Self { + cell_id: value.cell_id.into(), + yield_time_ms: value.yield_time_ms, + } + } +} + +/// Image detail values accepted in a V1 runtime response. +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum WireImageDetail { + Auto, + Low, + High, + Original, +} + +impl From for WireImageDetail { + fn from(value: ImageDetail) -> Self { + match value { + ImageDetail::Auto => Self::Auto, + ImageDetail::Low => Self::Low, + ImageDetail::High => Self::High, + ImageDetail::Original => Self::Original, + } + } +} + +impl From for ImageDetail { + fn from(value: WireImageDetail) -> Self { + match value { + WireImageDetail::Auto => Self::Auto, + WireImageDetail::Low => Self::Low, + WireImageDetail::High => Self::High, + WireImageDetail::Original => Self::Original, + } + } +} + +/// One output item emitted by a V1 runtime response. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields, tag = "type", rename_all = "snake_case")] +pub enum WireContentItem { + InputText { + text: String, + }, + InputImage { + image_url: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + detail: Option, + }, +} + +impl From for WireContentItem { + fn from(value: FunctionCallOutputContentItem) -> Self { + match value { + FunctionCallOutputContentItem::InputText { text } => Self::InputText { text }, + FunctionCallOutputContentItem::InputImage { image_url, detail } => Self::InputImage { + image_url, + detail: detail.map(Into::into), + }, + } + } +} + +impl From for FunctionCallOutputContentItem { + fn from(value: WireContentItem) -> Self { + match value { + WireContentItem::InputText { text } => Self::InputText { text }, + WireContentItem::InputImage { image_url, detail } => Self::InputImage { + image_url, + detail: detail.map(Into::into), + }, + } + } +} + +/// Runtime output returned over the V1 host connection. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub enum WireRuntimeResponse { + Yielded { + cell_id: WireCellId, + content_items: Vec, + }, + Terminated { + cell_id: WireCellId, + content_items: Vec, + }, + Result { + cell_id: WireCellId, + content_items: Vec, + error_text: Option, + }, +} + +impl From for WireRuntimeResponse { + fn from(value: RuntimeResponse) -> Self { + match value { + RuntimeResponse::Yielded { + cell_id, + content_items, + } => Self::Yielded { + cell_id: cell_id.into(), + content_items: content_items.into_iter().map(Into::into).collect(), + }, + RuntimeResponse::Terminated { + cell_id, + content_items, + } => Self::Terminated { + cell_id: cell_id.into(), + content_items: content_items.into_iter().map(Into::into).collect(), + }, + RuntimeResponse::Result { + cell_id, + content_items, + error_text, + } => Self::Result { + cell_id: cell_id.into(), + content_items: content_items.into_iter().map(Into::into).collect(), + error_text, + }, + } + } +} + +impl From for RuntimeResponse { + fn from(value: WireRuntimeResponse) -> Self { + match value { + WireRuntimeResponse::Yielded { + cell_id, + content_items, + } => Self::Yielded { + cell_id: cell_id.into(), + content_items: content_items.into_iter().map(Into::into).collect(), + }, + WireRuntimeResponse::Terminated { + cell_id, + content_items, + } => Self::Terminated { + cell_id: cell_id.into(), + content_items: content_items.into_iter().map(Into::into).collect(), + }, + WireRuntimeResponse::Result { + cell_id, + content_items, + error_text, + } => Self::Result { + cell_id: cell_id.into(), + content_items: content_items.into_iter().map(Into::into).collect(), + error_text, + }, + } + } +} + +/// Whether a waited-for cell remained live in protocol V1. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub enum WireWaitOutcome { + LiveCell(WireRuntimeResponse), + MissingCell(WireRuntimeResponse), +} + +impl From for WireWaitOutcome { + fn from(value: WaitOutcome) -> Self { + match value { + WaitOutcome::LiveCell(response) => Self::LiveCell(response.into()), + WaitOutcome::MissingCell(response) => Self::MissingCell(response.into()), + } + } +} + +impl From for WaitOutcome { + fn from(value: WireWaitOutcome) -> Self { + match value { + WireWaitOutcome::LiveCell(response) => Self::LiveCell(response.into()), + WireWaitOutcome::MissingCell(response) => Self::MissingCell(response.into()), + } + } +} + +/// A nested tool invocation sent over the V1 host connection. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct WireNestedToolCall { + pub cell_id: WireCellId, + pub runtime_tool_call_id: String, + pub tool_name: WireToolName, + pub tool_kind: WireToolKind, + pub input: Option, +} + +impl From for WireNestedToolCall { + fn from(value: CodeModeNestedToolCall) -> Self { + Self { + cell_id: value.cell_id.into(), + runtime_tool_call_id: value.runtime_tool_call_id, + tool_name: value.tool_name.into(), + tool_kind: value.tool_kind.into(), + input: value.input, + } + } +} + +impl From for CodeModeNestedToolCall { + fn from(value: WireNestedToolCall) -> Self { + Self { + cell_id: value.cell_id.into(), + runtime_tool_call_id: value.runtime_tool_call_id, + tool_name: value.tool_name.into(), + tool_kind: value.tool_kind.into(), + input: value.input, + } + } +} diff --git a/codex-rs/code-mode-protocol/src/host/types.rs b/codex-rs/code-mode-protocol/src/host/types.rs index 69c8fd0881..40c69df90b 100644 --- a/codex-rs/code-mode-protocol/src/host/types.rs +++ b/codex-rs/code-mode-protocol/src/host/types.rs @@ -8,6 +8,28 @@ use serde::Serialize; use serde::Serializer; use serde::de::Error as _; +/// Correlates one client operation request with the host's response. +#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] +#[serde(transparent)] +pub struct RequestId(i64); + +impl RequestId { + pub const fn new(value: i64) -> Self { + Self(value) + } +} + +/// Correlates one host delegate request with the client's response. +#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] +#[serde(transparent)] +pub struct DelegateRequestId(i64); + +impl DelegateRequestId { + pub const fn new(value: i64) -> Self { + Self(value) + } +} + #[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] #[serde(transparent)] pub struct ProtocolVersion(NonZeroU32); From e7ae783bffb534b1ade78ac2bf5382ea9b607f6a Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Wed, 24 Jun 2026 20:15:15 -0700 Subject: [PATCH 7/7] core: raise token budget message limits --- codex-rs/core/config.schema.json | 4 ++-- codex-rs/core/src/config/mod.rs | 4 ++-- codex-rs/features/src/feature_configs.rs | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 842f4d5cc5..f29ba59f64 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -2889,12 +2889,12 @@ }, "guidance_message": { "description": "Guidance appended to the context-window metadata in a developer message.", - "maxLength": 1000, + "maxLength": 4096, "type": "string" }, "reminder_message_template": { "description": "Reminder template. `{n_remaining}` is replaced with the tokens remaining before auto-compaction.", - "maxLength": 1000, + "maxLength": 4096, "minLength": 1, "type": "string" }, diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index ed565c6ad5..69a9bac41b 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -1088,8 +1088,8 @@ pub(crate) const DEFAULT_TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE: &str = concat!( "Your context window is nearly exhausted (only {n_remaining} tokens remaining) and will be automatically reset for you soon. ", "Once reset, message items in current context window will be cleared in the new window, but notes and history items will be persistent across windows." ); -const TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE_MAX_BYTES: usize = 1000; -const TOKEN_BUDGET_GUIDANCE_MESSAGE_MAX_BYTES: usize = 1000; +const TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE_MAX_BYTES: usize = 4096; +const TOKEN_BUDGET_GUIDANCE_MESSAGE_MAX_BYTES: usize = 4096; #[derive(Debug, Clone, PartialEq, Eq, Serialize)] pub struct TokenBudgetConfig { diff --git a/codex-rs/features/src/feature_configs.rs b/codex-rs/features/src/feature_configs.rs index a5e9104bce..6f158f3e73 100644 --- a/codex-rs/features/src/feature_configs.rs +++ b/codex-rs/features/src/feature_configs.rs @@ -86,11 +86,11 @@ pub struct TokenBudgetConfigToml { /// Reminder template. `{n_remaining}` is replaced with the tokens remaining before /// auto-compaction. #[serde(skip_serializing_if = "Option::is_none")] - #[schemars(length(min = 1, max = 1000))] + #[schemars(length(min = 1, max = 4096))] pub reminder_message_template: Option, /// Guidance appended to the context-window metadata in a developer message. #[serde(skip_serializing_if = "Option::is_none")] - #[schemars(length(max = 1000))] + #[schemars(length(max = 4096))] pub guidance_message: Option, }