From 206dd13c329e4174090b8d77b4da088ffa930144 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Thu, 16 Apr 2026 11:16:44 -0700 Subject: [PATCH 01/16] Move more connector logic into connectors crate (#18158) Reduce the size of core --- codex-rs/Cargo.lock | 2 + codex-rs/chatgpt/Cargo.toml | 1 + codex-rs/chatgpt/src/connectors.rs | 97 ++--- codex-rs/connectors/src/accessible.rs | 76 ++++ codex-rs/connectors/src/filter.rs | 68 ++++ codex-rs/connectors/src/lib.rs | 5 + codex-rs/connectors/src/merge.rs | 119 ++++++ codex-rs/connectors/src/metadata.rs | 27 ++ codex-rs/core/src/codex.rs | 18 +- codex-rs/core/src/codex_tests.rs | 2 +- codex-rs/core/src/connectors.rs | 348 ++---------------- codex-rs/core/src/connectors_tests.rs | 44 ++- codex-rs/core/src/plugins/injection.rs | 3 +- codex-rs/core/src/plugins/mentions.rs | 3 +- codex-rs/tui/Cargo.toml | 1 + codex-rs/tui/src/app_event.rs | 2 +- codex-rs/tui/src/bottom_pane/chat_composer.rs | 7 +- codex-rs/tui/src/chatwidget.rs | 19 +- codex-rs/tui/src/chatwidget/skills.rs | 7 +- .../chatwidget/tests/popups_and_settings.rs | 49 +-- 20 files changed, 462 insertions(+), 436 deletions(-) create mode 100644 codex-rs/connectors/src/accessible.rs create mode 100644 codex-rs/connectors/src/filter.rs create mode 100644 codex-rs/connectors/src/merge.rs create mode 100644 codex-rs/connectors/src/metadata.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index a499a511ca..296e51725c 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1649,6 +1649,7 @@ version = "0.0.0" dependencies = [ "anyhow", "clap", + "codex-app-server-protocol", "codex-config", "codex-connectors", "codex-core", @@ -2936,6 +2937,7 @@ dependencies = [ "codex-cli", "codex-cloud-requirements", "codex-config", + "codex-connectors", "codex-exec-server", "codex-features", "codex-feedback", diff --git a/codex-rs/chatgpt/Cargo.toml b/codex-rs/chatgpt/Cargo.toml index 381b4fc873..354449934a 100644 --- a/codex-rs/chatgpt/Cargo.toml +++ b/codex-rs/chatgpt/Cargo.toml @@ -10,6 +10,7 @@ workspace = true [dependencies] anyhow = { workspace = true } clap = { workspace = true, features = ["derive"] } +codex-app-server-protocol = { workspace = true } codex-connectors = { workspace = true } codex-config = { workspace = true } codex-core = { workspace = true } diff --git a/codex-rs/chatgpt/src/connectors.rs b/codex-rs/chatgpt/src/connectors.rs index f37609be31..5f6efbc124 100644 --- a/codex-rs/chatgpt/src/connectors.rs +++ b/codex-rs/chatgpt/src/connectors.rs @@ -1,7 +1,3 @@ -use codex_core::config::Config; -use codex_login::AuthManager; -use codex_login::CodexAuth; -use codex_login::token_data::TokenData; use std::collections::HashSet; use std::time::Duration; @@ -9,21 +5,24 @@ use crate::chatgpt_client::chatgpt_get_request_with_timeout; use crate::chatgpt_token::get_chatgpt_token_data; use crate::chatgpt_token::init_chatgpt_token_from_auth; +use codex_app_server_protocol::AppInfo; use codex_connectors::AllConnectorsCacheKey; use codex_connectors::DirectoryListResponse; - -pub use codex_core::connectors::AppInfo; -pub use codex_core::connectors::connector_display_label; -use codex_core::connectors::filter_disallowed_connectors; +use codex_connectors::filter::filter_disallowed_connectors; +use codex_connectors::merge::merge_connectors; +use codex_connectors::merge::merge_plugin_connectors; +use codex_core::config::Config; pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools; pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_options; pub use codex_core::connectors::list_accessible_connectors_from_mcp_tools_with_options_and_status; pub use codex_core::connectors::list_cached_accessible_connectors_from_mcp_tools; -use codex_core::connectors::merge_connectors; -use codex_core::connectors::merge_plugin_apps; pub use codex_core::connectors::with_app_enabled_state; use codex_core::plugins::AppConnectorId; use codex_core::plugins::PluginsManager; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_login::default_client::originator; +use codex_login::token_data::TokenData; const DIRECTORY_CONNECTORS_TIMEOUT: Duration = Duration::from_secs(60); @@ -74,8 +73,17 @@ pub async fn list_cached_all_connectors(config: &Config) -> Option> let token_data = get_chatgpt_token_data()?; let cache_key = all_connectors_cache_key(config, &token_data); let connectors = codex_connectors::cached_all_connectors(&cache_key)?; - let connectors = merge_plugin_apps(connectors, plugin_apps_for_config(config).await); - Some(filter_disallowed_connectors(connectors)) + let connectors = merge_plugin_connectors( + connectors, + plugin_apps_for_config(config) + .await + .into_iter() + .map(|connector_id| connector_id.0), + ); + Some(filter_disallowed_connectors( + connectors, + originator().value.as_str(), + )) } pub async fn list_all_connectors_with_options( @@ -105,8 +113,17 @@ pub async fn list_all_connectors_with_options( }, ) .await?; - let connectors = merge_plugin_apps(connectors, plugin_apps_for_config(config).await); - Ok(filter_disallowed_connectors(connectors)) + let connectors = merge_plugin_connectors( + connectors, + plugin_apps_for_config(config) + .await + .into_iter() + .map(|connector_id| connector_id.0), + ); + Ok(filter_disallowed_connectors( + connectors, + originator().value.as_str(), + )) } fn all_connectors_cache_key(config: &Config, token_data: &TokenData) -> AllConnectorsCacheKey { @@ -134,7 +151,13 @@ pub fn connectors_for_plugin_apps( .map(|connector_id| connector_id.0.as_str()) .collect::>(); - filter_disallowed_connectors(merge_plugin_apps(connectors, plugin_apps.to_vec())) + let connectors = merge_plugin_connectors( + connectors, + plugin_apps + .iter() + .map(|connector_id| connector_id.0.clone()), + ); + filter_disallowed_connectors(connectors, originator().value.as_str()) .into_iter() .filter(|connector| plugin_app_ids.contains(connector.id.as_str())) .collect() @@ -158,13 +181,13 @@ pub fn merge_connectors_with_accessible( accessible_connectors }; let merged = merge_connectors(connectors, accessible_connectors); - filter_disallowed_connectors(merged) + filter_disallowed_connectors(merged, originator().value.as_str()) } #[cfg(test)] mod tests { use super::*; - use codex_core::connectors::connector_install_url; + use codex_connectors::metadata::connector_install_url; use codex_core::plugins::AppConnectorId; use pretty_assertions::assert_eq; @@ -186,46 +209,6 @@ mod tests { } } - #[test] - fn allows_asdk_connectors() { - let filtered = filter_disallowed_connectors(vec![app("asdk_app_hidden"), app("alpha")]); - assert_eq!(filtered, vec![app("asdk_app_hidden"), app("alpha")]); - } - - #[test] - fn allows_whitelisted_asdk_connectors() { - let filtered = filter_disallowed_connectors(vec![ - app("asdk_app_69781557cc1481919cf5e9824fa2e792"), - app("beta"), - ]); - assert_eq!( - filtered, - vec![ - app("asdk_app_69781557cc1481919cf5e9824fa2e792"), - app("beta") - ] - ); - } - - #[test] - fn filters_openai_prefixed_connectors() { - let filtered = filter_disallowed_connectors(vec![ - app("connector_openai_foo"), - app("connector_openai_bar"), - app("gamma"), - ]); - assert_eq!(filtered, vec![app("gamma")]); - } - - #[test] - fn filters_disallowed_connector_ids() { - let filtered = filter_disallowed_connectors(vec![ - app("asdk_app_6938a94a61d881918ef32cb999ff937c"), - app("delta"), - ]); - assert_eq!(filtered, vec![app("delta")]); - } - fn merged_app(id: &str, is_accessible: bool) -> AppInfo { AppInfo { id: id.to_string(), diff --git a/codex-rs/connectors/src/accessible.rs b/codex-rs/connectors/src/accessible.rs new file mode 100644 index 0000000000..c44f8d8a38 --- /dev/null +++ b/codex-rs/connectors/src/accessible.rs @@ -0,0 +1,76 @@ +use std::collections::BTreeSet; +use std::collections::HashMap; + +use crate::metadata::connector_install_url; +use crate::normalize_connector_value; +use codex_app_server_protocol::AppInfo; + +pub struct AccessibleConnectorTool { + pub connector_id: String, + pub connector_name: Option, + pub connector_description: Option, + pub plugin_display_names: Vec, +} + +pub fn collect_accessible_connectors(tools: I) -> Vec +where + I: IntoIterator, +{ + let mut connectors: HashMap)> = HashMap::new(); + for tool in tools { + let connector_id = tool.connector_id; + let connector_name = normalize_connector_value(tool.connector_name.as_deref()) + .unwrap_or_else(|| connector_id.clone()); + let connector_description = + normalize_connector_value(tool.connector_description.as_deref()); + if let Some((existing, existing_plugin_display_names)) = connectors.get_mut(&connector_id) { + if existing.name == connector_id && connector_name != connector_id { + existing.name = connector_name; + } + if existing.description.is_none() && connector_description.is_some() { + existing.description = connector_description; + } + existing_plugin_display_names.extend(tool.plugin_display_names); + } else { + connectors.insert( + connector_id.clone(), + ( + AppInfo { + id: connector_id.clone(), + name: connector_name, + description: connector_description, + logo_url: None, + logo_url_dark: None, + distribution_channel: None, + branding: None, + app_metadata: None, + labels: None, + install_url: None, + is_accessible: true, + is_enabled: true, + plugin_display_names: Vec::new(), + }, + tool.plugin_display_names + .into_iter() + .collect::>(), + ), + ); + } + } + let mut accessible: Vec = connectors + .into_values() + .map(|(mut connector, plugin_display_names)| { + connector.plugin_display_names = plugin_display_names.into_iter().collect(); + connector.install_url = Some(connector_install_url(&connector.name, &connector.id)); + connector + }) + .collect(); + accessible.sort_by(|left, right| { + right + .is_accessible + .cmp(&left.is_accessible) + .then_with(|| left.name.cmp(&right.name)) + .then_with(|| left.id.cmp(&right.id)) + }); + accessible +} diff --git a/codex-rs/connectors/src/filter.rs b/codex-rs/connectors/src/filter.rs new file mode 100644 index 0000000000..82c334f82d --- /dev/null +++ b/codex-rs/connectors/src/filter.rs @@ -0,0 +1,68 @@ +use std::collections::HashSet; + +use codex_app_server_protocol::AppInfo; + +pub fn filter_tool_suggest_discoverable_connectors( + directory_connectors: Vec, + accessible_connectors: &[AppInfo], + discoverable_connector_ids: &HashSet, + originator_value: &str, +) -> Vec { + let accessible_connector_ids: HashSet<&str> = accessible_connectors + .iter() + .filter(|connector| connector.is_accessible) + .map(|connector| connector.id.as_str()) + .collect(); + + let mut connectors = filter_disallowed_connectors(directory_connectors, originator_value) + .into_iter() + .filter(|connector| !accessible_connector_ids.contains(connector.id.as_str())) + .filter(|connector| discoverable_connector_ids.contains(connector.id.as_str())) + .collect::>(); + connectors.sort_by(|left, right| { + left.name + .cmp(&right.name) + .then_with(|| left.id.cmp(&right.id)) + }); + connectors +} + +const DISALLOWED_CONNECTOR_IDS: &[&str] = &[ + "asdk_app_6938a94a61d881918ef32cb999ff937c", + "connector_2b0a9009c9c64bf9933a3dae3f2b1254", + "connector_3f8d1a79f27c4c7ba1a897ab13bf37dc", + "connector_68de829bf7648191acd70a907364c67c", + "connector_68e004f14af881919eb50893d3d9f523", + "connector_69272cb413a081919685ec3c88d1744e", +]; +const FIRST_PARTY_CHAT_DISALLOWED_CONNECTOR_IDS: &[&str] = + &["connector_0f9c9d4592e54d0a9a12b3f44a1e2010"]; +const DISALLOWED_CONNECTOR_PREFIX: &str = "connector_openai_"; + +pub fn filter_disallowed_connectors( + connectors: Vec, + originator_value: &str, +) -> Vec { + let first_party_chat_originator = is_first_party_chat_originator(originator_value); + connectors + .into_iter() + .filter(|connector| { + is_connector_id_allowed(connector.id.as_str(), first_party_chat_originator) + }) + .collect() +} + +fn is_first_party_chat_originator(originator_value: &str) -> bool { + originator_value == "codex_atlas" || originator_value == "codex_chatgpt_desktop" +} + +fn is_connector_id_allowed(connector_id: &str, first_party_chat_originator: bool) -> bool { + let disallowed_connector_ids = if first_party_chat_originator { + FIRST_PARTY_CHAT_DISALLOWED_CONNECTOR_IDS + } else { + DISALLOWED_CONNECTOR_IDS + }; + + !connector_id.starts_with(DISALLOWED_CONNECTOR_PREFIX) + && !disallowed_connector_ids.contains(&connector_id) +} diff --git a/codex-rs/connectors/src/lib.rs b/codex-rs/connectors/src/lib.rs index a74ece5a35..d7c5fb405c 100644 --- a/codex-rs/connectors/src/lib.rs +++ b/codex-rs/connectors/src/lib.rs @@ -10,6 +10,11 @@ use codex_app_server_protocol::AppInfo; use codex_app_server_protocol::AppMetadata; use serde::Deserialize; +pub mod accessible; +pub mod filter; +pub mod merge; +pub mod metadata; + pub const CONNECTORS_CACHE_TTL: Duration = Duration::from_secs(3600); #[derive(Clone, Debug, PartialEq, Eq)] diff --git a/codex-rs/connectors/src/merge.rs b/codex-rs/connectors/src/merge.rs new file mode 100644 index 0000000000..b41ee63add --- /dev/null +++ b/codex-rs/connectors/src/merge.rs @@ -0,0 +1,119 @@ +use std::collections::HashMap; +use std::collections::HashSet; + +use crate::metadata::connector_install_url; +use crate::metadata::sort_connectors_by_accessibility_and_name; +use codex_app_server_protocol::AppInfo; + +pub fn merge_connectors( + connectors: Vec, + accessible_connectors: Vec, +) -> Vec { + let mut merged: HashMap = connectors + .into_iter() + .map(|mut connector| { + connector.is_accessible = false; + (connector.id.clone(), connector) + }) + .collect(); + + for mut connector in accessible_connectors { + connector.is_accessible = true; + let connector_id = connector.id.clone(); + if let Some(existing) = merged.get_mut(&connector_id) { + existing.is_accessible = true; + if existing.name == existing.id && connector.name != connector.id { + existing.name = connector.name; + } + if existing.description.is_none() && connector.description.is_some() { + existing.description = connector.description; + } + if existing.logo_url.is_none() && connector.logo_url.is_some() { + existing.logo_url = connector.logo_url; + } + if existing.logo_url_dark.is_none() && connector.logo_url_dark.is_some() { + existing.logo_url_dark = connector.logo_url_dark; + } + if existing.distribution_channel.is_none() && connector.distribution_channel.is_some() { + existing.distribution_channel = connector.distribution_channel; + } + existing + .plugin_display_names + .extend(connector.plugin_display_names); + } else { + merged.insert(connector_id, connector); + } + } + + let mut merged = merged.into_values().collect::>(); + for connector in &mut merged { + if connector.install_url.is_none() { + connector.install_url = Some(connector_install_url(&connector.name, &connector.id)); + } + connector.plugin_display_names.sort_unstable(); + connector.plugin_display_names.dedup(); + } + sort_connectors_by_accessibility_and_name(&mut merged); + merged +} + +pub fn merge_plugin_connectors(connectors: Vec, plugin_app_ids: I) -> Vec +where + I: IntoIterator, +{ + let mut merged = connectors; + let mut connector_ids = merged + .iter() + .map(|connector| connector.id.clone()) + .collect::>(); + + for connector_id in plugin_app_ids { + if connector_ids.insert(connector_id.clone()) { + merged.push(plugin_connector_to_app_info(connector_id)); + } + } + + sort_connectors_by_accessibility_and_name(&mut merged); + merged +} + +pub fn merge_plugin_connectors_with_accessible( + plugin_app_ids: I, + accessible_connectors: Vec, +) -> Vec +where + I: IntoIterator, +{ + let accessible_connector_ids: HashSet<&str> = accessible_connectors + .iter() + .map(|connector| connector.id.as_str()) + .collect(); + let plugin_connectors = plugin_app_ids + .into_iter() + .filter(|connector_id| accessible_connector_ids.contains(connector_id.as_str())) + .map(plugin_connector_to_app_info) + .collect::>(); + merge_connectors(plugin_connectors, accessible_connectors) +} + +pub fn plugin_connector_to_app_info(connector_id: String) -> AppInfo { + // Leave the placeholder name as the connector id so merge_connectors() can + // replace it with canonical app metadata from directory fetches or + // connector_name values from codex_apps tool discovery. + let name = connector_id.clone(); + AppInfo { + id: connector_id.clone(), + name: name.clone(), + description: None, + logo_url: None, + logo_url_dark: None, + distribution_channel: None, + branding: None, + app_metadata: None, + labels: None, + install_url: Some(connector_install_url(&name, &connector_id)), + is_accessible: false, + is_enabled: true, + plugin_display_names: Vec::new(), + } +} diff --git a/codex-rs/connectors/src/metadata.rs b/codex-rs/connectors/src/metadata.rs new file mode 100644 index 0000000000..64becb766d --- /dev/null +++ b/codex-rs/connectors/src/metadata.rs @@ -0,0 +1,27 @@ +use codex_app_server_protocol::AppInfo; + +pub fn connector_display_label(connector: &AppInfo) -> String { + connector.name.clone() +} + +pub fn connector_mention_slug(connector: &AppInfo) -> String { + crate::connector_name_slug(&connector_display_label(connector)) +} + +pub fn connector_install_url(name: &str, connector_id: &str) -> String { + crate::connector_install_url(name, connector_id) +} + +pub fn sanitize_name(name: &str) -> String { + crate::connector_name_slug(name).replace("-", "_") +} + +pub(crate) fn sort_connectors_by_accessibility_and_name(connectors: &mut [AppInfo]) { + connectors.sort_by(|left, right| { + right + .is_accessible + .cmp(&left.is_accessible) + .then_with(|| left.name.cmp(&right.name)) + .then_with(|| left.id.cmp(&right.id)) + }); +} diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 50a4c25abd..5274d88664 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -6267,8 +6267,11 @@ pub(crate) async fn run_turn( HashMap::new() }; let available_connectors = if turn_context.apps_enabled() { - let connectors = connectors::merge_plugin_apps_with_accessible( - loaded_plugins.effective_apps(), + let connectors = codex_connectors::merge::merge_plugin_connectors_with_accessible( + loaded_plugins + .effective_apps() + .into_iter() + .map(|connector_id| connector_id.0), connectors::accessible_connectors_from_mcp_tools(&mcp_tools), ); connectors::with_app_enabled_state(connectors, &turn_context.config) @@ -6914,7 +6917,7 @@ fn collect_explicit_app_ids_from_skill_items( let connector_slug_counts = build_connector_slug_counts(connectors); for connector in connectors { - let slug = connectors::connector_mention_slug(connector); + let slug = codex_connectors::metadata::connector_mention_slug(connector); let connector_count = connector_slug_counts.get(&slug).copied().unwrap_or(0); let skill_count = skill_name_counts_lower.get(&slug).copied().unwrap_or(0); if connector_count == 1 && skill_count == 0 && mention_names_lower.contains(&slug) { @@ -6989,7 +6992,7 @@ fn connector_inserted_in_messages( return true; } - let mention_slug = connectors::connector_mention_slug(connector); + let mention_slug = codex_connectors::metadata::connector_mention_slug(connector); let connector_count = connector_slug_counts .get(&mention_slug) .copied() @@ -7217,8 +7220,11 @@ pub(crate) async fn built_tools( connectors::with_app_enabled_state(connectors.clone(), &turn_context.config) }); let connectors = if apps_enabled { - let connectors = connectors::merge_plugin_apps_with_accessible( - loaded_plugins.effective_apps(), + let connectors = codex_connectors::merge::merge_plugin_connectors_with_accessible( + loaded_plugins + .effective_apps() + .into_iter() + .map(|connector_id| connector_id.0), accessible_connectors.clone().unwrap_or_default(), ); Some(connectors::with_app_enabled_state( diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 6f4276b077..81ebd6c9f8 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -413,7 +413,7 @@ fn make_mcp_tool( ) -> ToolInfo { let tool_namespace = if server_name == CODEX_APPS_MCP_SERVER_NAME { connector_name - .map(crate::connectors::sanitize_name) + .map(codex_connectors::metadata::sanitize_name) .map(|connector_name| format!("mcp__{server_name}__{connector_name}")) .unwrap_or_else(|| server_name.to_string()) } else { diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 77603adcf1..103ce251fc 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -1,4 +1,3 @@ -use std::collections::BTreeSet; use std::collections::HashMap; use std::collections::HashSet; use std::sync::Arc; @@ -26,7 +25,6 @@ use crate::codex::INITIAL_SUBMIT_ID; use crate::config::Config; use crate::config_loader::AppsRequirementsToml; use crate::mcp::McpManager; -use crate::plugins::AppConnectorId; use crate::plugins::PluginsManager; use crate::plugins::list_tool_suggest_discoverable_plugins; use codex_config::types::AppToolApproval; @@ -36,7 +34,6 @@ use codex_features::Feature; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::default_client::create_client; -use codex_login::default_client::is_first_party_chat_originator; use codex_login::default_client::originator; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::McpConnectionManager; @@ -46,7 +43,6 @@ use codex_mcp::codex_apps_tools_cache_key; use codex_mcp::compute_auth_statuses; use codex_mcp::with_codex_apps_mcp; -pub use codex_connectors::CONNECTORS_CACHE_TTL; const CONNECTORS_READY_TIMEOUT_ON_EMPTY_TOOLS: Duration = Duration::from_secs(30); const DIRECTORY_CONNECTORS_TIMEOUT: Duration = Duration::from_secs(60); @@ -122,13 +118,15 @@ pub(crate) async fn list_tool_suggest_discoverable_tools_with_auth( let directory_connectors = list_directory_connectors_for_tool_suggest_with_auth(config, auth).await?; let connector_ids = tool_suggest_connector_ids(config).await; - let discoverable_connectors = filter_tool_suggest_discoverable_connectors( - directory_connectors, - accessible_connectors, - &connector_ids, - ) - .into_iter() - .map(DiscoverableTool::from); + let discoverable_connectors = + codex_connectors::filter::filter_tool_suggest_discoverable_connectors( + directory_connectors, + accessible_connectors, + &connector_ids, + originator().value.as_str(), + ) + .into_iter() + .map(DiscoverableTool::from); let discoverable_plugins = list_tool_suggest_discoverable_plugins(config) .await? .into_iter() @@ -151,7 +149,12 @@ pub async fn list_cached_accessible_connectors_from_mcp_tools( return Some(Vec::new()); } let cache_key = accessible_connectors_cache_key(config, auth.as_ref()); - read_cached_accessible_connectors(&cache_key).map(filter_disallowed_connectors) + read_cached_accessible_connectors(&cache_key).map(|connectors| { + codex_connectors::filter::filter_disallowed_connectors( + connectors, + originator().value.as_str(), + ) + }) } pub(crate) fn refresh_accessible_connectors_cache_from_mcp_tools( @@ -164,8 +167,10 @@ pub(crate) fn refresh_accessible_connectors_cache_from_mcp_tools( } let cache_key = accessible_connectors_cache_key(config, auth); - let accessible_connectors = - filter_disallowed_connectors(accessible_connectors_from_mcp_tools(mcp_tools)); + let accessible_connectors = codex_connectors::filter::filter_disallowed_connectors( + accessible_connectors_from_mcp_tools(mcp_tools), + originator().value.as_str(), + ); write_cached_accessible_connectors(cache_key, &accessible_connectors); } @@ -202,7 +207,10 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_options_and_status( let tool_plugin_provenance = mcp_manager.tool_plugin_provenance(config).await; if !force_refetch && let Some(cached_connectors) = read_cached_accessible_connectors(&cache_key) { - let cached_connectors = filter_disallowed_connectors(cached_connectors); + let cached_connectors = codex_connectors::filter::filter_disallowed_connectors( + cached_connectors, + originator().value.as_str(), + ); let cached_connectors = with_app_plugin_sources(cached_connectors, &tool_plugin_provenance); return Ok(AccessibleConnectorsStatus { connectors: cached_connectors, @@ -293,8 +301,10 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_options_and_status( cancel_token.cancel(); } - let accessible_connectors = - filter_disallowed_connectors(accessible_connectors_from_mcp_tools(&tools)); + let accessible_connectors = codex_connectors::filter::filter_disallowed_connectors( + accessible_connectors_from_mcp_tools(&tools), + originator().value.as_str(), + ); if codex_apps_ready || !accessible_connectors.is_empty() { write_cached_accessible_connectors(cache_key, &accessible_connectors); } @@ -357,35 +367,11 @@ fn write_cached_accessible_connectors( .unwrap_or_else(std::sync::PoisonError::into_inner); *cache_guard = Some(CachedAccessibleConnectors { key: cache_key, - expires_at: Instant::now() + CONNECTORS_CACHE_TTL, + expires_at: Instant::now() + codex_connectors::CONNECTORS_CACHE_TTL, connectors: connectors.to_vec(), }); } -fn filter_tool_suggest_discoverable_connectors( - directory_connectors: Vec, - accessible_connectors: &[AppInfo], - discoverable_connector_ids: &HashSet, -) -> Vec { - let accessible_connector_ids: HashSet<&str> = accessible_connectors - .iter() - .filter(|connector| connector.is_accessible) - .map(|connector| connector.id.as_str()) - .collect(); - - let mut connectors = filter_disallowed_connectors(directory_connectors) - .into_iter() - .filter(|connector| !accessible_connector_ids.contains(connector.id.as_str())) - .filter(|connector| discoverable_connector_ids.contains(connector.id.as_str())) - .collect::>(); - connectors.sort_by(|left, right| { - left.name - .cmp(&right.name) - .then_with(|| left.id.cmp(&right.id)) - }); - connectors -} - async fn tool_suggest_connector_ids(config: &Config) -> HashSet { let mut connector_ids = PluginsManager::new(config.codex_home.to_path_buf()) .plugins_for_config(config) @@ -493,14 +479,6 @@ async fn chatgpt_get_request_with_token( } } -pub fn connector_display_label(connector: &AppInfo) -> String { - format_connector_label(&connector.name, &connector.id) -} - -pub fn connector_mention_slug(connector: &AppInfo) -> String { - sanitize_slug(&connector_display_label(connector)) -} - pub(crate) fn accessible_connectors_from_mcp_tools( mcp_tools: &HashMap, ) -> Vec { @@ -511,114 +489,14 @@ pub(crate) fn accessible_connectors_from_mcp_tools( return None; } let connector_id = tool.connector_id.as_deref()?; - Some(( - connector_id.to_string(), - normalize_connector_value(tool.connector_name.as_deref()), - normalize_connector_value(tool.connector_description.as_deref()), - tool.plugin_display_names.clone(), - )) - }); - collect_accessible_connectors(tools) -} - -pub fn merge_connectors( - connectors: Vec, - accessible_connectors: Vec, -) -> Vec { - let mut merged: HashMap = connectors - .into_iter() - .map(|mut connector| { - connector.is_accessible = false; - (connector.id.clone(), connector) + Some(codex_connectors::accessible::AccessibleConnectorTool { + connector_id: connector_id.to_string(), + connector_name: tool.connector_name.clone(), + connector_description: tool.connector_description.clone(), + plugin_display_names: tool.plugin_display_names.clone(), }) - .collect(); - - for mut connector in accessible_connectors { - connector.is_accessible = true; - let connector_id = connector.id.clone(); - if let Some(existing) = merged.get_mut(&connector_id) { - existing.is_accessible = true; - if existing.name == existing.id && connector.name != connector.id { - existing.name = connector.name; - } - if existing.description.is_none() && connector.description.is_some() { - existing.description = connector.description; - } - if existing.logo_url.is_none() && connector.logo_url.is_some() { - existing.logo_url = connector.logo_url; - } - if existing.logo_url_dark.is_none() && connector.logo_url_dark.is_some() { - existing.logo_url_dark = connector.logo_url_dark; - } - if existing.distribution_channel.is_none() && connector.distribution_channel.is_some() { - existing.distribution_channel = connector.distribution_channel; - } - existing - .plugin_display_names - .extend(connector.plugin_display_names); - } else { - merged.insert(connector_id, connector); - } - } - - let mut merged = merged.into_values().collect::>(); - for connector in &mut merged { - if connector.install_url.is_none() { - connector.install_url = Some(connector_install_url(&connector.name, &connector.id)); - } - connector.plugin_display_names.sort_unstable(); - connector.plugin_display_names.dedup(); - } - merged.sort_by(|left, right| { - right - .is_accessible - .cmp(&left.is_accessible) - .then_with(|| left.name.cmp(&right.name)) - .then_with(|| left.id.cmp(&right.id)) }); - merged -} - -pub fn merge_plugin_apps( - connectors: Vec, - plugin_apps: Vec, -) -> Vec { - let mut merged = connectors; - let mut connector_ids = merged - .iter() - .map(|connector| connector.id.clone()) - .collect::>(); - - for connector_id in plugin_apps { - if connector_ids.insert(connector_id.0.clone()) { - merged.push(plugin_app_to_app_info(connector_id)); - } - } - - merged.sort_by(|left, right| { - right - .is_accessible - .cmp(&left.is_accessible) - .then_with(|| left.name.cmp(&right.name)) - .then_with(|| left.id.cmp(&right.id)) - }); - merged -} - -pub fn merge_plugin_apps_with_accessible( - plugin_apps: Vec, - accessible_connectors: Vec, -) -> Vec { - let accessible_connector_ids: HashSet<&str> = accessible_connectors - .iter() - .map(|connector| connector.id.as_str()) - .collect(); - let plugin_connectors = plugin_apps - .into_iter() - .filter(|connector_id| accessible_connector_ids.contains(connector_id.0.as_str())) - .map(plugin_app_to_app_info) - .collect::>(); - merge_connectors(plugin_connectors, accessible_connectors) + codex_connectors::accessible::collect_accessible_connectors(tools) } pub fn with_app_enabled_state(mut connectors: Vec, config: &Config) -> Vec { @@ -691,45 +569,6 @@ pub(crate) fn codex_app_tool_is_enabled(config: &Config, tool_info: &ToolInfo) - .enabled } -const DISALLOWED_CONNECTOR_IDS: &[&str] = &[ - "asdk_app_6938a94a61d881918ef32cb999ff937c", - "connector_2b0a9009c9c64bf9933a3dae3f2b1254", - "connector_3f8d1a79f27c4c7ba1a897ab13bf37dc", - "connector_68de829bf7648191acd70a907364c67c", - "connector_68e004f14af881919eb50893d3d9f523", - "connector_69272cb413a081919685ec3c88d1744e", -]; -const FIRST_PARTY_CHAT_DISALLOWED_CONNECTOR_IDS: &[&str] = - &["connector_0f9c9d4592e54d0a9a12b3f44a1e2010"]; -const DISALLOWED_CONNECTOR_PREFIX: &str = "connector_openai_"; - -pub fn filter_disallowed_connectors(connectors: Vec) -> Vec { - filter_disallowed_connectors_for_originator(connectors, originator().value.as_str()) -} - -fn filter_disallowed_connectors_for_originator( - connectors: Vec, - originator_value: &str, -) -> Vec { - connectors - .into_iter() - .filter(|connector| { - is_connector_id_allowed_for_originator(connector.id.as_str(), originator_value) - }) - .collect() -} - -fn is_connector_id_allowed_for_originator(connector_id: &str, originator_value: &str) -> bool { - let disallowed_connector_ids = if is_first_party_chat_originator(originator_value) { - FIRST_PARTY_CHAT_DISALLOWED_CONNECTOR_IDS - } else { - DISALLOWED_CONNECTOR_IDS - }; - - !connector_id.starts_with(DISALLOWED_CONNECTOR_PREFIX) - && !disallowed_connector_ids.contains(&connector_id) -} - fn read_apps_config(config: &Config) -> Option { let apps_config = read_user_apps_config(config); let had_apps_config = apps_config.is_some(); @@ -851,125 +690,6 @@ fn app_tool_policy_from_apps_config( AppToolPolicy { enabled, approval } } -fn collect_accessible_connectors(tools: I) -> Vec -where - I: IntoIterator, Option, Vec)>, -{ - let mut connectors: HashMap)> = HashMap::new(); - for (connector_id, connector_name, connector_description, plugin_display_names) in tools { - let connector_name = connector_name.unwrap_or_else(|| connector_id.clone()); - if let Some((existing, existing_plugin_display_names)) = connectors.get_mut(&connector_id) { - if existing.name == connector_id && connector_name != connector_id { - existing.name = connector_name; - } - if existing.description.is_none() && connector_description.is_some() { - existing.description = connector_description; - } - existing_plugin_display_names.extend(plugin_display_names); - } else { - connectors.insert( - connector_id.clone(), - ( - AppInfo { - id: connector_id.clone(), - name: connector_name, - description: connector_description, - logo_url: None, - logo_url_dark: None, - distribution_channel: None, - branding: None, - app_metadata: None, - labels: None, - install_url: None, - is_accessible: true, - is_enabled: true, - plugin_display_names: Vec::new(), - }, - plugin_display_names - .into_iter() - .collect::>(), - ), - ); - } - } - let mut accessible: Vec = connectors - .into_values() - .map(|(mut connector, plugin_display_names)| { - connector.plugin_display_names = plugin_display_names.into_iter().collect(); - connector.install_url = Some(connector_install_url(&connector.name, &connector.id)); - connector - }) - .collect(); - accessible.sort_by(|left, right| { - right - .is_accessible - .cmp(&left.is_accessible) - .then_with(|| left.name.cmp(&right.name)) - .then_with(|| left.id.cmp(&right.id)) - }); - accessible -} - -fn plugin_app_to_app_info(connector_id: AppConnectorId) -> AppInfo { - // Leave the placeholder name as the connector id so merge_connectors() can - // replace it with canonical app metadata from directory fetches or - // connector_name values from codex_apps tool discovery. - let connector_id = connector_id.0; - let name = connector_id.clone(); - AppInfo { - id: connector_id.clone(), - name: name.clone(), - description: None, - logo_url: None, - logo_url_dark: None, - distribution_channel: None, - branding: None, - app_metadata: None, - labels: None, - install_url: Some(connector_install_url(&name, &connector_id)), - is_accessible: false, - is_enabled: true, - plugin_display_names: Vec::new(), - } -} - -fn normalize_connector_value(value: Option<&str>) -> Option { - value - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string) -} - -pub fn connector_install_url(name: &str, connector_id: &str) -> String { - let slug = sanitize_slug(name); - format!("https://chatgpt.com/apps/{slug}/{connector_id}") -} - -pub fn sanitize_name(name: &str) -> String { - sanitize_slug(name).replace("-", "_") -} - -fn sanitize_slug(name: &str) -> String { - let mut normalized = String::with_capacity(name.len()); - for character in name.chars() { - if character.is_ascii_alphanumeric() { - normalized.push(character.to_ascii_lowercase()); - } else { - normalized.push('-'); - } - } - let normalized = normalized.trim_matches('-'); - if normalized.is_empty() { - "app".to_string() - } else { - normalized.to_string() - } -} - -fn format_connector_label(name: &str, _id: &str) -> String { - name.to_string() -} - #[cfg(test)] #[path = "connectors_tests.rs"] mod tests; diff --git a/codex-rs/core/src/connectors_tests.rs b/codex-rs/core/src/connectors_tests.rs index c9f77d0469..e462c7939e 100644 --- a/codex-rs/core/src/connectors_tests.rs +++ b/codex-rs/core/src/connectors_tests.rs @@ -11,6 +11,13 @@ use codex_config::types::AppConfig; use codex_config::types::AppToolConfig; use codex_config::types::AppToolsConfig; use codex_config::types::AppsDefaultConfig; +use codex_connectors::filter::filter_disallowed_connectors; +use codex_connectors::filter::filter_tool_suggest_discoverable_connectors; +use codex_connectors::merge::merge_connectors; +use codex_connectors::merge::plugin_connector_to_app_info; +use codex_connectors::metadata::connector_install_url; +use codex_connectors::metadata::connector_mention_slug; +use codex_connectors::metadata::sanitize_name; use codex_features::Feature; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::ToolInfo; @@ -138,7 +145,7 @@ fn with_accessible_connectors_cache_cleared(f: impl FnOnce() -> R) -> R { #[test] fn merge_connectors_replaces_plugin_placeholder_name_with_accessible_name() { - let plugin = plugin_app_to_app_info(AppConnectorId("calendar".to_string())); + let plugin = plugin_connector_to_app_info("calendar".to_string()); let accessible = google_calendar_accessible_connector(&[]); let merged = merge_connectors(vec![plugin], vec![accessible]); @@ -281,7 +288,7 @@ async fn refresh_accessible_connectors_cache_from_mcp_tools_writes_latest_instal #[test] fn merge_connectors_unions_and_dedupes_plugin_display_names() { - let mut plugin = plugin_app_to_app_info(AppConnectorId("calendar".to_string())); + let mut plugin = plugin_connector_to_app_info("calendar".to_string()); plugin.plugin_display_names = plugin_names(&["sample", "alpha", "sample"]); let accessible = google_calendar_accessible_connector(&["beta", "alpha"]); @@ -975,33 +982,40 @@ fn app_tool_policy_matches_prefix_stripped_tool_name_for_tool_config() { #[test] fn filter_disallowed_connectors_allows_non_disallowed_connectors() { - let filtered = filter_disallowed_connectors(vec![app("asdk_app_hidden"), app("alpha")]); + let filtered = + filter_disallowed_connectors(vec![app("asdk_app_hidden"), app("alpha")], "codex_cli"); assert_eq!(filtered, vec![app("asdk_app_hidden"), app("alpha")]); } #[test] fn filter_disallowed_connectors_filters_openai_prefix() { - let filtered = filter_disallowed_connectors(vec![ - app("connector_openai_foo"), - app("connector_openai_bar"), - app("gamma"), - ]); + let filtered = filter_disallowed_connectors( + vec![ + app("connector_openai_foo"), + app("connector_openai_bar"), + app("gamma"), + ], + "codex_cli", + ); assert_eq!(filtered, vec![app("gamma")]); } #[test] fn filter_disallowed_connectors_filters_disallowed_connector_ids() { - let filtered = filter_disallowed_connectors(vec![ - app("asdk_app_6938a94a61d881918ef32cb999ff937c"), - app("connector_3f8d1a79f27c4c7ba1a897ab13bf37dc"), - app("delta"), - ]); + let filtered = filter_disallowed_connectors( + vec![ + app("asdk_app_6938a94a61d881918ef32cb999ff937c"), + app("connector_3f8d1a79f27c4c7ba1a897ab13bf37dc"), + app("delta"), + ], + "codex_cli", + ); assert_eq!(filtered, vec![app("delta")]); } #[test] fn first_party_chat_originator_filters_target_and_openai_prefixed_connectors() { - let filtered = filter_disallowed_connectors_for_originator( + let filtered = filter_disallowed_connectors( vec![ app("connector_openai_foo"), app("asdk_app_6938a94a61d881918ef32cb999ff937c"), @@ -1064,6 +1078,7 @@ fn filter_tool_suggest_discoverable_connectors_keeps_only_plugin_backed_uninstal "connector_2128aebfecb84f64a069897515042a44".to_string(), "connector_68df038e0ba48191908c8434991bbac2".to_string(), ]), + "codex_cli", ); assert_eq!( @@ -1103,6 +1118,7 @@ fn filter_tool_suggest_discoverable_connectors_excludes_accessible_apps_even_whe "connector_2128aebfecb84f64a069897515042a44".to_string(), "connector_68df038e0ba48191908c8434991bbac2".to_string(), ]), + "codex_cli", ); assert_eq!(filtered, Vec::::new()); diff --git a/codex-rs/core/src/plugins/injection.rs b/codex-rs/core/src/plugins/injection.rs index 3e5bb6ecb5..2f6ec26260 100644 --- a/codex-rs/core/src/plugins/injection.rs +++ b/codex-rs/core/src/plugins/injection.rs @@ -1,6 +1,7 @@ use std::collections::BTreeSet; use std::collections::HashMap; +use codex_connectors::metadata::connector_display_label; use codex_protocol::models::DeveloperInstructions; use codex_protocol::models::ResponseItem; @@ -46,7 +47,7 @@ pub(crate) fn build_plugin_injections( .iter() .any(|plugin_name| plugin_name == &plugin.display_name) }) - .map(connectors::connector_display_label) + .map(connector_display_label) .collect::>() .into_iter() .collect::>(); diff --git a/codex-rs/core/src/plugins/mentions.rs b/codex-rs/core/src/plugins/mentions.rs index 5e11b9b91e..a5e94345cb 100644 --- a/codex-rs/core/src/plugins/mentions.rs +++ b/codex-rs/core/src/plugins/mentions.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use std::collections::HashSet; +use codex_connectors::metadata::connector_mention_slug; use codex_protocol::user_input::UserInput; use crate::connectors; @@ -108,7 +109,7 @@ pub(crate) fn build_connector_slug_counts( ) -> HashMap { let mut counts: HashMap = HashMap::new(); for connector in connectors { - let slug = connectors::connector_mention_slug(connector); + let slug = connector_mention_slug(connector); *counts.entry(slug).or_insert(0) += 1; } counts diff --git a/codex-rs/tui/Cargo.toml b/codex-rs/tui/Cargo.toml index 0f29fdeae6..65f35b665a 100644 --- a/codex-rs/tui/Cargo.toml +++ b/codex-rs/tui/Cargo.toml @@ -33,6 +33,7 @@ codex-install-context = { workspace = true } codex-chatgpt = { workspace = true } codex-cloud-requirements = { workspace = true } codex-config = { workspace = true } +codex-connectors = { workspace = true } codex-exec-server = { workspace = true } codex-features = { workspace = true } codex-feedback = { workspace = true } diff --git a/codex-rs/tui/src/app_event.rs b/codex-rs/tui/src/app_event.rs index a3b31aa988..934baac231 100644 --- a/codex-rs/tui/src/app_event.rs +++ b/codex-rs/tui/src/app_event.rs @@ -10,13 +10,13 @@ use std::path::PathBuf; +use codex_app_server_protocol::AppInfo; use codex_app_server_protocol::McpServerStatus; use codex_app_server_protocol::PluginInstallResponse; use codex_app_server_protocol::PluginListResponse; use codex_app_server_protocol::PluginReadParams; use codex_app_server_protocol::PluginReadResponse; use codex_app_server_protocol::PluginUninstallResponse; -use codex_chatgpt::connectors::AppInfo; use codex_file_search::FileMatch; use codex_protocol::ThreadId; use codex_protocol::openai_models::ModelPreset; diff --git a/codex-rs/tui/src/bottom_pane/chat_composer.rs b/codex-rs/tui/src/bottom_pane/chat_composer.rs index b1c3129ff6..5b39a0f01d 100644 --- a/codex-rs/tui/src/bottom_pane/chat_composer.rs +++ b/codex-rs/tui/src/bottom_pane/chat_composer.rs @@ -209,8 +209,7 @@ use crate::legacy_core::plugins::PluginCapabilitySummary; use crate::legacy_core::skills::model::SkillMetadata; use crate::tui::FrameRequester; use crate::ui_consts::LIVE_PREFIX_COLS; -use codex_chatgpt::connectors; -use codex_chatgpt::connectors::AppInfo; +use codex_app_server_protocol::AppInfo; use codex_file_search::FileMatch; use std::cell::RefCell; use std::collections::HashMap; @@ -3417,9 +3416,9 @@ impl ChatComposer { if !connector.is_accessible || !connector.is_enabled { continue; } - let display_name = connectors::connector_display_label(connector); + let display_name = codex_connectors::metadata::connector_display_label(connector); let description = Some(Self::connector_brief_description(connector)); - let slug = crate::legacy_core::connectors::connector_mention_slug(connector); + let slug = codex_connectors::metadata::connector_mention_slug(connector); let search_terms = vec![display_name.clone(), connector.id.clone(), slug.clone()]; let connector_id = connector.id.as_str(); mentions.push(MentionItem { diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index 05f1bb5ccc..f2697d9072 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -79,6 +79,7 @@ use crate::terminal_title::clear_terminal_title; use crate::terminal_title::set_terminal_title; use crate::text_formatting::proper_join; use crate::version::CODEX_CLI_VERSION; +use codex_app_server_protocol::AppInfo; use codex_app_server_protocol::AppSummary; use codex_app_server_protocol::CodexErrorInfo as AppServerCodexErrorInfo; use codex_app_server_protocol::CollabAgentState as AppServerCollabAgentState; @@ -5525,7 +5526,7 @@ impl ChatWidget { let app_mentions = find_app_mentions(&mentions, apps, &skill_names_lower); for app in app_mentions { - let slug = crate::legacy_core::connectors::connector_mention_slug(&app); + let slug = codex_connectors::metadata::connector_mention_slug(&app); if bound_names.contains(&slug) || !selected_app_ids.insert(app.id.clone()) { continue; } @@ -9701,7 +9702,7 @@ impl ChatWidget { self.config.features.enabled(Feature::Apps) && self.has_chatgpt_account } - fn connectors_for_mentions(&self) -> Option<&[connectors::AppInfo]> { + fn connectors_for_mentions(&self) -> Option<&[AppInfo]> { if !self.connectors_enabled() { return None; } @@ -9888,7 +9889,7 @@ impl ChatWidget { } } - fn open_connectors_popup(&mut self, connectors: &[connectors::AppInfo]) { + fn open_connectors_popup(&mut self, connectors: &[AppInfo]) { self.bottom_pane.show_selection_view( self.connectors_popup_params(connectors, /*selected_connector_id*/ None), ); @@ -9914,7 +9915,7 @@ impl ChatWidget { fn connectors_popup_params( &self, - connectors: &[connectors::AppInfo], + connectors: &[AppInfo], selected_connector_id: Option<&str>, ) -> SelectionViewParams { let total = connectors.len(); @@ -9937,7 +9938,7 @@ impl ChatWidget { }); let mut items: Vec = Vec::with_capacity(connectors.len()); for connector in connectors { - let connector_label = connectors::connector_display_label(connector); + let connector_label = codex_connectors::metadata::connector_display_label(connector); let connector_title = connector_label.clone(); let link_description = Self::connector_description(connector); let description = Self::connector_brief_description(connector); @@ -10011,7 +10012,7 @@ impl ChatWidget { } } - fn refresh_connectors_popup_if_open(&mut self, connectors: &[connectors::AppInfo]) { + fn refresh_connectors_popup_if_open(&mut self, connectors: &[AppInfo]) { let selected_connector_id = if let (Some(selected_index), ConnectorsCacheState::Ready(snapshot)) = ( self.bottom_pane @@ -10039,7 +10040,7 @@ impl ChatWidget { ]) } - fn connector_brief_description(connector: &connectors::AppInfo) -> String { + fn connector_brief_description(connector: &AppInfo) -> String { let status_label = Self::connector_status_label(connector); match Self::connector_description(connector) { Some(description) => format!("{status_label} · {description}"), @@ -10047,7 +10048,7 @@ impl ChatWidget { } } - fn connector_status_label(connector: &connectors::AppInfo) -> &'static str { + fn connector_status_label(connector: &AppInfo) -> &'static str { if connector.is_accessible { if connector.is_enabled { "Installed" @@ -10059,7 +10060,7 @@ impl ChatWidget { } } - fn connector_description(connector: &connectors::AppInfo) -> Option { + fn connector_description(connector: &AppInfo) -> Option { connector .description .as_deref() diff --git a/codex-rs/tui/src/chatwidget/skills.rs b/codex-rs/tui/src/chatwidget/skills.rs index 0bcfb3b3bd..9bbde1aaec 100644 --- a/codex-rs/tui/src/chatwidget/skills.rs +++ b/codex-rs/tui/src/chatwidget/skills.rs @@ -9,14 +9,13 @@ use crate::bottom_pane::SkillsToggleItem; use crate::bottom_pane::SkillsToggleView; use crate::bottom_pane::popup_consts::standard_popup_hint_line; use crate::legacy_core::TOOL_MENTION_SIGIL; -use crate::legacy_core::connectors::connector_mention_slug; use crate::legacy_core::skills::model::SkillDependencies; use crate::legacy_core::skills::model::SkillInterface; use crate::legacy_core::skills::model::SkillMetadata; use crate::legacy_core::skills::model::SkillToolDependency; use crate::skills_helpers::skill_description; use crate::skills_helpers::skill_display_name; -use codex_chatgpt::connectors::AppInfo; +use codex_app_server_protocol::AppInfo; use codex_protocol::parse_command::ParsedCommand; use codex_protocol::protocol::ListSkillsResponseEvent; use codex_protocol::protocol::SkillMetadata as ProtocolSkillMetadata; @@ -296,12 +295,12 @@ pub(crate) fn find_app_mentions( let mut slug_counts: HashMap = HashMap::new(); for app in apps.iter().filter(|app| app.is_enabled) { - let slug = connector_mention_slug(app); + let slug = codex_connectors::metadata::connector_mention_slug(app); *slug_counts.entry(slug).or_insert(0) += 1; } for app in apps.iter().filter(|app| app.is_enabled) { - let slug = connector_mention_slug(app); + let slug = codex_connectors::metadata::connector_mention_slug(app); let slug_count = slug_counts.get(&slug).copied().unwrap_or(0); if mentions.names.contains(&slug) && !explicit_names.contains(&slug) 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 375c37314b..10e3451d03 100644 --- a/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs +++ b/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs @@ -1,4 +1,5 @@ use super::*; +use codex_app_server_protocol::AppInfo; use pretty_assertions::assert_eq; #[tokio::test] @@ -511,7 +512,7 @@ async fn apps_popup_stays_loading_until_final_snapshot_updates() { chat.on_connectors_loaded( Ok(ConnectorsSnapshot { - connectors: vec![codex_chatgpt::connectors::AppInfo { + connectors: vec![AppInfo { id: notion_id.to_string(), name: "Notion".to_string(), description: Some("Workspace docs".to_string()), @@ -545,7 +546,7 @@ async fn apps_popup_stays_loading_until_final_snapshot_updates() { chat.on_connectors_loaded( Ok(ConnectorsSnapshot { connectors: vec![ - codex_chatgpt::connectors::AppInfo { + AppInfo { id: notion_id.to_string(), name: "Notion".to_string(), description: Some("Workspace docs".to_string()), @@ -560,7 +561,7 @@ async fn apps_popup_stays_loading_until_final_snapshot_updates() { is_enabled: true, plugin_display_names: Vec::new(), }, - codex_chatgpt::connectors::AppInfo { + AppInfo { id: linear_id.to_string(), name: "Linear".to_string(), description: Some("Project tracking".to_string()), @@ -604,7 +605,7 @@ async fn apps_refresh_failure_keeps_existing_full_snapshot() { let linear_id = "unit_test_apps_refresh_failure_connector_2"; let full_connectors = vec![ - codex_chatgpt::connectors::AppInfo { + AppInfo { id: notion_id.to_string(), name: "Notion".to_string(), description: Some("Workspace docs".to_string()), @@ -619,7 +620,7 @@ async fn apps_refresh_failure_keeps_existing_full_snapshot() { is_enabled: true, plugin_display_names: Vec::new(), }, - codex_chatgpt::connectors::AppInfo { + AppInfo { id: linear_id.to_string(), name: "Linear".to_string(), description: Some("Project tracking".to_string()), @@ -644,7 +645,7 @@ async fn apps_refresh_failure_keeps_existing_full_snapshot() { chat.on_connectors_loaded( Ok(ConnectorsSnapshot { - connectors: vec![codex_chatgpt::connectors::AppInfo { + connectors: vec![AppInfo { id: notion_id.to_string(), name: "Notion".to_string(), description: Some("Workspace docs".to_string()), @@ -693,7 +694,7 @@ async fn apps_popup_preserves_selected_app_across_refresh() { chat.on_connectors_loaded( Ok(ConnectorsSnapshot { connectors: vec![ - codex_chatgpt::connectors::AppInfo { + AppInfo { id: "notion".to_string(), name: "Notion".to_string(), description: Some("Workspace docs".to_string()), @@ -708,7 +709,7 @@ async fn apps_popup_preserves_selected_app_across_refresh() { is_enabled: true, plugin_display_names: Vec::new(), }, - codex_chatgpt::connectors::AppInfo { + AppInfo { id: "slack".to_string(), name: "Slack".to_string(), description: Some("Team chat".to_string()), @@ -739,7 +740,7 @@ async fn apps_popup_preserves_selected_app_across_refresh() { chat.on_connectors_loaded( Ok(ConnectorsSnapshot { connectors: vec![ - codex_chatgpt::connectors::AppInfo { + AppInfo { id: "airtable".to_string(), name: "Airtable".to_string(), description: Some("Spreadsheets".to_string()), @@ -754,7 +755,7 @@ async fn apps_popup_preserves_selected_app_across_refresh() { is_enabled: true, plugin_display_names: Vec::new(), }, - codex_chatgpt::connectors::AppInfo { + AppInfo { id: "notion".to_string(), name: "Notion".to_string(), description: Some("Workspace docs".to_string()), @@ -769,7 +770,7 @@ async fn apps_popup_preserves_selected_app_across_refresh() { is_enabled: true, plugin_display_names: Vec::new(), }, - codex_chatgpt::connectors::AppInfo { + AppInfo { id: "slack".to_string(), name: "Slack".to_string(), description: Some("Team chat".to_string()), @@ -812,7 +813,7 @@ async fn apps_refresh_failure_with_cached_snapshot_triggers_pending_force_refetc chat.connectors_prefetch_in_flight = true; chat.connectors_force_refetch_pending = true; - let full_connectors = vec![codex_chatgpt::connectors::AppInfo { + let full_connectors = vec![AppInfo { id: "unit_test_apps_refresh_failure_pending_connector".to_string(), name: "Notion".to_string(), description: Some("Workspace docs".to_string()), @@ -855,7 +856,7 @@ async fn apps_popup_keeps_existing_full_snapshot_while_partial_refresh_loads() { chat.bottom_pane.set_connectors_enabled(/*enabled*/ true); let full_connectors = vec![ - codex_chatgpt::connectors::AppInfo { + AppInfo { id: "unit_test_connector_1".to_string(), name: "Notion".to_string(), description: Some("Workspace docs".to_string()), @@ -870,7 +871,7 @@ async fn apps_popup_keeps_existing_full_snapshot_while_partial_refresh_loads() { is_enabled: true, plugin_display_names: Vec::new(), }, - codex_chatgpt::connectors::AppInfo { + AppInfo { id: "unit_test_connector_2".to_string(), name: "Linear".to_string(), description: Some("Project tracking".to_string()), @@ -897,7 +898,7 @@ async fn apps_popup_keeps_existing_full_snapshot_while_partial_refresh_loads() { chat.on_connectors_loaded( Ok(ConnectorsSnapshot { connectors: vec![ - codex_chatgpt::connectors::AppInfo { + AppInfo { id: "unit_test_connector_1".to_string(), name: "Notion".to_string(), description: Some("Workspace docs".to_string()), @@ -912,7 +913,7 @@ async fn apps_popup_keeps_existing_full_snapshot_while_partial_refresh_loads() { is_enabled: true, plugin_display_names: Vec::new(), }, - codex_chatgpt::connectors::AppInfo { + AppInfo { id: "connector_openai_hidden".to_string(), name: "Hidden OpenAI".to_string(), description: Some("Should be filtered".to_string()), @@ -960,7 +961,7 @@ async fn apps_refresh_failure_without_full_snapshot_falls_back_to_installed_apps chat.on_connectors_loaded( Ok(ConnectorsSnapshot { - connectors: vec![codex_chatgpt::connectors::AppInfo { + connectors: vec![AppInfo { id: "unit_test_apps_refresh_failure_fallback_connector".to_string(), name: "Notion".to_string(), description: Some("Workspace docs".to_string()), @@ -1019,7 +1020,7 @@ async fn apps_popup_shows_disabled_status_for_installed_but_disabled_apps() { chat.on_connectors_loaded( Ok(ConnectorsSnapshot { - connectors: vec![codex_chatgpt::connectors::AppInfo { + connectors: vec![AppInfo { id: "connector_1".to_string(), name: "Notion".to_string(), description: Some("Workspace docs".to_string()), @@ -1073,7 +1074,7 @@ async fn apps_initial_load_applies_enabled_state_from_config() { chat.on_connectors_loaded( Ok(ConnectorsSnapshot { - connectors: vec![codex_chatgpt::connectors::AppInfo { + connectors: vec![AppInfo { id: "connector_1".to_string(), name: "Notion".to_string(), description: Some("Workspace docs".to_string()), @@ -1139,7 +1140,7 @@ async fn apps_initial_load_applies_enabled_state_from_requirements_with_user_ove chat.on_connectors_loaded( Ok(ConnectorsSnapshot { - connectors: vec![codex_chatgpt::connectors::AppInfo { + connectors: vec![AppInfo { id: "connector_1".to_string(), name: "Notion".to_string(), description: Some("Workspace docs".to_string()), @@ -1203,7 +1204,7 @@ async fn apps_initial_load_applies_enabled_state_from_requirements_without_user_ chat.on_connectors_loaded( Ok(ConnectorsSnapshot { - connectors: vec![codex_chatgpt::connectors::AppInfo { + connectors: vec![AppInfo { id: "connector_1".to_string(), name: "Notion".to_string(), description: Some("Workspace docs".to_string()), @@ -1252,7 +1253,7 @@ async fn apps_refresh_preserves_toggled_enabled_state() { chat.on_connectors_loaded( Ok(ConnectorsSnapshot { - connectors: vec![codex_chatgpt::connectors::AppInfo { + connectors: vec![AppInfo { id: "connector_1".to_string(), name: "Notion".to_string(), description: Some("Workspace docs".to_string()), @@ -1274,7 +1275,7 @@ async fn apps_refresh_preserves_toggled_enabled_state() { chat.on_connectors_loaded( Ok(ConnectorsSnapshot { - connectors: vec![codex_chatgpt::connectors::AppInfo { + connectors: vec![AppInfo { id: "connector_1".to_string(), name: "Notion".to_string(), description: Some("Workspace docs".to_string()), @@ -1323,7 +1324,7 @@ async fn apps_popup_for_not_installed_app_uses_install_only_selected_description chat.on_connectors_loaded( Ok(ConnectorsSnapshot { - connectors: vec![codex_chatgpt::connectors::AppInfo { + connectors: vec![AppInfo { id: "connector_2".to_string(), name: "Linear".to_string(), description: Some("Project tracking".to_string()), From 71174574adb09a90ebd83e2acfe284a39aaca2cf Mon Sep 17 00:00:00 2001 From: Matthew Zeng Date: Thu, 16 Apr 2026 11:18:07 -0700 Subject: [PATCH 02/16] Add server-level approval defaults for custom MCP servers (#17843) ## Summary - Add `default_tools_approval_mode` support for custom MCP server configs, matching the existing `codex_apps` behavior - Apply approval precedence as per-tool override, then server default, then `auto` - Update config serialization, CLI display, schema generation, docs, and tests ## Testing - `cargo check -p codex-config` - `cargo check -p codex-core` - `just write-config-schema` - `just fmt` - `cargo test -p codex-config` - Targeted `codex-core` tests for config parsing, config writes, and MCP approval precedence - `just fix -p codex-config -p codex-core` --- codex-rs/cli/src/mcp_cmd.rs | 10 +++++ codex-rs/codex-mcp/src/mcp/mod.rs | 1 + codex-rs/codex-mcp/src/mcp/mod_tests.rs | 2 + .../codex-mcp/src/mcp/skill_dependencies.rs | 2 + .../src/mcp/skill_dependencies_tests.rs | 2 + .../src/mcp_connection_manager_tests.rs | 2 + codex-rs/config/src/mcp_edit.rs | 7 +++ codex-rs/config/src/mcp_edit_tests.rs | 2 + codex-rs/config/src/mcp_types.rs | 8 ++++ codex-rs/config/src/mcp_types_tests.rs | 33 ++++++++++++++ codex-rs/core/config.schema.json | 8 ++++ codex-rs/core/src/config/config_tests.rs | 33 +++++++++++--- codex-rs/core/src/config/edit.rs | 7 +++ codex-rs/core/src/config/edit_tests.rs | 8 ++++ codex-rs/core/src/config/service_tests.rs | 43 +++++++++++++++++++ codex-rs/core/src/mcp_skill_dependencies.rs | 2 + codex-rs/core/src/mcp_tool_call.rs | 11 +++-- codex-rs/core/src/mcp_tool_call_tests.rs | 37 ++++++++++++++++ codex-rs/core/src/plugins/manager_tests.rs | 4 ++ codex-rs/core/tests/suite/code_mode.rs | 1 + codex-rs/core/tests/suite/js_repl.rs | 1 + codex-rs/core/tests/suite/rmcp_client.rs | 1 + codex-rs/core/tests/suite/search_tool.rs | 1 + codex-rs/core/tests/suite/sqlite_state.rs | 1 + codex-rs/core/tests/suite/tools.rs | 1 + codex-rs/core/tests/suite/truncation.rs | 3 ++ docs/config.md | 12 ++++-- 27 files changed, 231 insertions(+), 12 deletions(-) diff --git a/codex-rs/cli/src/mcp_cmd.rs b/codex-rs/cli/src/mcp_cmd.rs index 4ea084b83f..d413f72ddc 100644 --- a/codex-rs/cli/src/mcp_cmd.rs +++ b/codex-rs/cli/src/mcp_cmd.rs @@ -6,6 +6,7 @@ use anyhow::Result; use anyhow::anyhow; use anyhow::bail; use clap::ArgGroup; +use codex_config::types::AppToolApproval; use codex_config::types::McpServerConfig; use codex_config::types::McpServerTransportConfig; use codex_core::McpManager; @@ -304,6 +305,7 @@ async fn run_add(config_overrides: &CliConfigOverrides, add_args: AddArgs) -> Re disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -879,6 +881,14 @@ async fn run_get(config_overrides: &CliConfigOverrides, get_args: GetArgs) -> Re if let Some(timeout) = server.tool_timeout_sec { println!(" tool_timeout_sec: {}", timeout.as_secs_f64()); } + if let Some(approval_mode) = server.default_tools_approval_mode { + let approval_mode = match approval_mode { + AppToolApproval::Auto => "auto", + AppToolApproval::Prompt => "prompt", + AppToolApproval::Approve => "approve", + }; + println!(" default_tools_approval_mode: {approval_mode}"); + } println!(" remove: codex mcp remove {}", get_args.name); Ok(()) diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index db90bfe429..1dc9db0789 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -276,6 +276,7 @@ fn codex_apps_mcp_server_config(config: &McpConfig, auth: Option<&CodexAuth>) -> disabled_reason: None, startup_timeout_sec: Some(Duration::from_secs(30)), tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, diff --git a/codex-rs/codex-mcp/src/mcp/mod_tests.rs b/codex-rs/codex-mcp/src/mcp/mod_tests.rs index c1d39f2a7c..8db52a9d83 100644 --- a/codex-rs/codex-mcp/src/mcp/mod_tests.rs +++ b/codex-rs/codex-mcp/src/mcp/mod_tests.rs @@ -200,6 +200,7 @@ async fn effective_mcp_servers_preserve_user_servers_and_add_codex_apps() { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -223,6 +224,7 @@ async fn effective_mcp_servers_preserve_user_servers_and_add_codex_apps() { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, diff --git a/codex-rs/codex-mcp/src/mcp/skill_dependencies.rs b/codex-rs/codex-mcp/src/mcp/skill_dependencies.rs index ab72e25958..f785fe4bde 100644 --- a/codex-rs/codex-mcp/src/mcp/skill_dependencies.rs +++ b/codex-rs/codex-mcp/src/mcp/skill_dependencies.rs @@ -126,6 +126,7 @@ fn mcp_dependency_to_server_config( disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -154,6 +155,7 @@ fn mcp_dependency_to_server_config( disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, diff --git a/codex-rs/codex-mcp/src/mcp/skill_dependencies_tests.rs b/codex-rs/codex-mcp/src/mcp/skill_dependencies_tests.rs index b562bf9e2c..2d8390d15e 100644 --- a/codex-rs/codex-mcp/src/mcp/skill_dependencies_tests.rs +++ b/codex-rs/codex-mcp/src/mcp/skill_dependencies_tests.rs @@ -46,6 +46,7 @@ fn collect_missing_respects_canonical_installed_key() { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -98,6 +99,7 @@ fn collect_missing_dedupes_by_canonical_key_but_preserves_original_name() { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, diff --git a/codex-rs/codex-mcp/src/mcp_connection_manager_tests.rs b/codex-rs/codex-mcp/src/mcp_connection_manager_tests.rs index 6b65d455a8..cf2889ccde 100644 --- a/codex-rs/codex-mcp/src/mcp_connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/mcp_connection_manager_tests.rs @@ -799,6 +799,7 @@ fn mcp_init_error_display_prompts_for_github_pat() { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -850,6 +851,7 @@ fn mcp_init_error_display_reports_generic_errors() { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, diff --git a/codex-rs/config/src/mcp_edit.rs b/codex-rs/config/src/mcp_edit.rs index d315d23f8e..c4bb38c543 100644 --- a/codex-rs/config/src/mcp_edit.rs +++ b/codex-rs/config/src/mcp_edit.rs @@ -189,6 +189,13 @@ fn serialize_mcp_server(config: &McpServerConfig) -> TomlItem { if let Some(timeout) = config.tool_timeout_sec { entry["tool_timeout_sec"] = value(timeout.as_secs_f64()); } + if let Some(approval_mode) = config.default_tools_approval_mode { + entry["default_tools_approval_mode"] = value(match approval_mode { + AppToolApproval::Auto => "auto", + AppToolApproval::Prompt => "prompt", + AppToolApproval::Approve => "approve", + }); + } if let Some(enabled_tools) = &config.enabled_tools && !enabled_tools.is_empty() { diff --git a/codex-rs/config/src/mcp_edit_tests.rs b/codex-rs/config/src/mcp_edit_tests.rs index 5f60c2612a..cfcd73c3e5 100644 --- a/codex-rs/config/src/mcp_edit_tests.rs +++ b/codex-rs/config/src/mcp_edit_tests.rs @@ -29,6 +29,7 @@ async fn replace_mcp_servers_serializes_per_tool_approval_overrides() -> anyhow: disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: Some(AppToolApproval::Auto), enabled_tools: None, disabled_tools: None, scopes: None, @@ -62,6 +63,7 @@ async fn replace_mcp_servers_serializes_per_tool_approval_overrides() -> anyhow: r#"[mcp_servers.docs] command = "docs-server" supports_parallel_tool_calls = true +default_tools_approval_mode = "auto" [mcp_servers.docs.tools] diff --git a/codex-rs/config/src/mcp_types.rs b/codex-rs/config/src/mcp_types.rs index 8aece35c0d..75b68c3f94 100644 --- a/codex-rs/config/src/mcp_types.rs +++ b/codex-rs/config/src/mcp_types.rs @@ -93,6 +93,10 @@ pub struct McpServerConfig { #[serde(default, with = "option_duration_secs")] pub tool_timeout_sec: Option, + /// Approval mode for tools in this server unless a tool override exists. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub default_tools_approval_mode: Option, + /// Explicit allow-list of tools exposed from this server. When set, only these tools will be registered. #[serde(default, skip_serializing_if = "Option::is_none")] pub enabled_tools: Option>, @@ -158,6 +162,8 @@ pub struct RawMcpServerConfig { #[serde(default)] pub supports_parallel_tool_calls: Option, #[serde(default)] + pub default_tools_approval_mode: Option, + #[serde(default)] pub enabled_tools: Option>, #[serde(default)] pub disabled_tools: Option>, @@ -194,6 +200,7 @@ impl TryFrom for McpServerConfig { enabled, required, supports_parallel_tool_calls, + default_tools_approval_mode, enabled_tools, disabled_tools, scopes, @@ -260,6 +267,7 @@ impl TryFrom for McpServerConfig { required: required.unwrap_or_default(), supports_parallel_tool_calls: supports_parallel_tool_calls.unwrap_or_default(), disabled_reason: None, + default_tools_approval_mode, enabled_tools, disabled_tools, scopes, diff --git a/codex-rs/config/src/mcp_types_tests.rs b/codex-rs/config/src/mcp_types_tests.rs index 500ea56f62..dff4a6bbf5 100644 --- a/codex-rs/config/src/mcp_types_tests.rs +++ b/codex-rs/config/src/mcp_types_tests.rs @@ -258,6 +258,38 @@ fn deserialize_server_config_with_parallel_tool_calls() { assert!(cfg.supports_parallel_tool_calls); } +#[test] +fn deserialize_server_config_with_default_tool_approval_mode() { + let cfg: McpServerConfig = toml::from_str( + r#" + command = "echo" + default_tools_approval_mode = "approve" + + [tools.search] + approval_mode = "prompt" + "#, + ) + .expect("should deserialize default tool approval mode"); + + assert_eq!( + cfg.default_tools_approval_mode, + Some(AppToolApproval::Approve) + ); + assert_eq!( + cfg.tools.get("search"), + Some(&McpServerToolConfig { + approval_mode: Some(AppToolApproval::Prompt), + }) + ); + + let serialized = toml::to_string(&cfg).expect("should serialize MCP config"); + assert!(serialized.contains("default_tools_approval_mode = \"approve\"")); + + let round_tripped: McpServerConfig = + toml::from_str(&serialized).expect("should deserialize serialized MCP config"); + assert_eq!(round_tripped, cfg); +} + #[test] fn serialize_round_trips_server_config_with_parallel_tool_calls() { let cfg: McpServerConfig = toml::from_str( @@ -304,6 +336,7 @@ fn deserialize_ignores_unknown_server_fields() { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 9c7eb853a9..966555856b 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -1463,6 +1463,14 @@ "default": null, "type": "string" }, + "default_tools_approval_mode": { + "allOf": [ + { + "$ref": "#/definitions/AppToolApproval" + } + ], + "default": null + }, "disabled_tools": { "default": null, "items": { diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 0001439816..34002b129a 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -90,6 +90,7 @@ fn stdio_mcp(command: &str) -> McpServerConfig { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -113,6 +114,7 @@ fn http_mcp(url: &str) -> McpServerConfig { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -2130,6 +2132,7 @@ async fn replace_mcp_servers_round_trips_entries() -> anyhow::Result<()> { disabled_reason: None, startup_timeout_sec: Some(Duration::from_secs(3)), tool_timeout_sec: Some(Duration::from_secs(5)), + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -2243,23 +2246,28 @@ fn mcp_servers_toml_parses_per_tool_approval_overrides() { [mcp_servers.docs] command = "docs-server" name = "Docs" +default_tools_approval_mode = "prompt" [mcp_servers.docs.tools.search] approval_mode = "approve" "#, ) .expect("TOML deserialization should succeed"); - let tool = config + let server = config .mcp_servers .get("docs") - .and_then(|server| server.tools.get("search")) - .expect("docs/search tool config exists"); + .expect("docs server config exists"); assert_eq!( - tool, - &McpServerToolConfig { + server.default_tools_approval_mode, + Some(AppToolApproval::Prompt) + ); + + assert_eq!( + server.tools.get("search"), + Some(&McpServerToolConfig { approval_mode: Some(AppToolApproval::Approve), - } + }) ); } @@ -2380,6 +2388,7 @@ async fn replace_mcp_servers_serializes_env_sorted() -> anyhow::Result<()> { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -2455,6 +2464,7 @@ async fn replace_mcp_servers_serializes_env_vars() -> anyhow::Result<()> { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -2510,6 +2520,7 @@ async fn replace_mcp_servers_serializes_cwd() -> anyhow::Result<()> { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -2563,6 +2574,7 @@ async fn replace_mcp_servers_streamable_http_serializes_bearer_token() -> anyhow disabled_reason: None, startup_timeout_sec: Some(Duration::from_secs(2)), tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -2632,6 +2644,7 @@ async fn replace_mcp_servers_streamable_http_serializes_custom_headers() -> anyh disabled_reason: None, startup_timeout_sec: Some(Duration::from_secs(2)), tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -2713,6 +2726,7 @@ async fn replace_mcp_servers_streamable_http_removes_optional_sections() -> anyh disabled_reason: None, startup_timeout_sec: Some(Duration::from_secs(2)), tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -2747,6 +2761,7 @@ async fn replace_mcp_servers_streamable_http_removes_optional_sections() -> anyh disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -2816,6 +2831,7 @@ async fn replace_mcp_servers_streamable_http_isolates_headers_between_servers() disabled_reason: None, startup_timeout_sec: Some(Duration::from_secs(2)), tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -2840,6 +2856,7 @@ async fn replace_mcp_servers_streamable_http_isolates_headers_between_servers() disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -2927,6 +2944,7 @@ async fn replace_mcp_servers_serializes_disabled_flag() -> anyhow::Result<()> { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -2976,6 +2994,7 @@ async fn replace_mcp_servers_serializes_required_flag() -> anyhow::Result<()> { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -3025,6 +3044,7 @@ async fn replace_mcp_servers_serializes_tool_filters() -> anyhow::Result<()> { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: Some(vec!["allowed".to_string()]), disabled_tools: Some(vec!["blocked".to_string()]), scopes: None, @@ -3078,6 +3098,7 @@ async fn replace_mcp_servers_streamable_http_serializes_oauth_resource() -> anyh disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, diff --git a/codex-rs/core/src/config/edit.rs b/codex-rs/core/src/config/edit.rs index 6f4231f908..5adc0f998c 100644 --- a/codex-rs/core/src/config/edit.rs +++ b/codex-rs/core/src/config/edit.rs @@ -237,6 +237,13 @@ mod document_helpers { if let Some(timeout) = config.tool_timeout_sec { entry["tool_timeout_sec"] = value(timeout.as_secs_f64()); } + if let Some(approval_mode) = config.default_tools_approval_mode { + entry["default_tools_approval_mode"] = value(match approval_mode { + AppToolApproval::Auto => "auto", + AppToolApproval::Prompt => "prompt", + AppToolApproval::Approve => "approve", + }); + } if let Some(enabled_tools) = &config.enabled_tools && !enabled_tools.is_empty() { diff --git a/codex-rs/core/src/config/edit_tests.rs b/codex-rs/core/src/config/edit_tests.rs index 46680d5434..4f340d89b3 100644 --- a/codex-rs/core/src/config/edit_tests.rs +++ b/codex-rs/core/src/config/edit_tests.rs @@ -582,6 +582,7 @@ fn blocking_replace_mcp_servers_round_trips() { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: Some(vec!["one".to_string(), "two".to_string()]), disabled_tools: None, scopes: None, @@ -610,6 +611,7 @@ fn blocking_replace_mcp_servers_round_trips() { disabled_reason: None, startup_timeout_sec: Some(std::time::Duration::from_secs(5)), tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: Some(vec!["forbidden".to_string()]), scopes: None, @@ -675,6 +677,7 @@ fn blocking_replace_mcp_servers_serializes_tool_approval_overrides() { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: Some(AppToolApproval::Prompt), enabled_tools: None, disabled_tools: None, scopes: None, @@ -699,6 +702,7 @@ fn blocking_replace_mcp_servers_serializes_tool_approval_overrides() { let expected = "\ [mcp_servers.docs] command = \"docs-server\" +default_tools_approval_mode = \"prompt\" [mcp_servers.docs.tools.search] approval_mode = \"approve\" @@ -737,6 +741,7 @@ foo = { command = "cmd" } disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -790,6 +795,7 @@ foo = { command = "cmd" } # keep me disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -842,6 +848,7 @@ foo = { command = "cmd", args = ["--flag"] } # keep me disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -895,6 +902,7 @@ foo = { command = "cmd" } disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, diff --git a/codex-rs/core/src/config/service_tests.rs b/codex-rs/core/src/config/service_tests.rs index ae62326605..faaffb34df 100644 --- a/codex-rs/core/src/config/service_tests.rs +++ b/codex-rs/core/src/config/service_tests.rs @@ -164,6 +164,49 @@ async fn write_value_supports_nested_app_paths() -> Result<()> { Ok(()) } +#[tokio::test] +async fn write_value_supports_custom_mcp_server_default_tool_approval_mode() -> Result<()> { + let tmp = tempdir().expect("tempdir"); + std::fs::write( + tmp.path().join(CONFIG_TOML_FILE), + "[mcp_servers.docs]\ncommand = \"docs-server\"\n", + )?; + + let service = ConfigService::without_managed_config_for_tests(tmp.path().to_path_buf()); + service + .write_value(ConfigValueWriteParams { + file_path: Some(tmp.path().join(CONFIG_TOML_FILE).display().to_string()), + key_path: "mcp_servers.docs.default_tools_approval_mode".to_string(), + value: serde_json::json!("approve"), + merge_strategy: MergeStrategy::Replace, + expected_version: None, + }) + .await + .expect("write mcp server default_tools_approval_mode succeeds"); + + let contents = std::fs::read_to_string(tmp.path().join(CONFIG_TOML_FILE))?; + assert!(contents.contains("default_tools_approval_mode = \"approve\"")); + + let read = service + .read(ConfigReadParams { + include_layers: false, + cwd: None, + }) + .await + .expect("config read succeeds"); + + assert_eq!( + read.config + .additional + .get("mcp_servers") + .and_then(|servers| servers.get("docs")) + .and_then(|docs| docs.get("default_tools_approval_mode")), + Some(&serde_json::json!("approve")) + ); + + Ok(()) +} + #[tokio::test] async fn read_includes_origins_and_layers() { let tmp = tempdir().expect("tempdir"); diff --git a/codex-rs/core/src/mcp_skill_dependencies.rs b/codex-rs/core/src/mcp_skill_dependencies.rs index c87b02730c..c12fe28cb2 100644 --- a/codex-rs/core/src/mcp_skill_dependencies.rs +++ b/codex-rs/core/src/mcp_skill_dependencies.rs @@ -371,6 +371,7 @@ fn mcp_dependency_to_server_config( disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -399,6 +400,7 @@ fn mcp_dependency_to_server_config( disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 064585f447..b707d51890 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -676,9 +676,14 @@ fn custom_mcp_tool_approval_mode( .and_then(|value| { HashMap::::deserialize(value).ok() }) - .and_then(|servers| servers.get(server).cloned()) - .and_then(|server| server.tools.get(tool_name).cloned()) - .and_then(|tool| tool.approval_mode) + .and_then(|servers| { + let server_config = servers.get(server)?; + server_config + .tools + .get(tool_name) + .and_then(|tool| tool.approval_mode) + .or(server_config.default_tools_approval_mode) + }) .unwrap_or_default() } diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 53e838908d..a85233c773 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1146,6 +1146,43 @@ async fn persist_custom_mcp_tool_approval_writes_tool_override() { assert!(contents.contains("[mcp_servers.docs.tools.search]")); } +#[tokio::test] +async fn custom_mcp_tool_approval_mode_uses_server_default_with_tool_override() { + let tmp = tempdir().expect("tempdir"); + std::fs::write( + tmp.path().join(CONFIG_TOML_FILE), + r#" +[mcp_servers.docs] +command = "docs-server" +default_tools_approval_mode = "approve" + +[mcp_servers.docs.tools.search] +approval_mode = "prompt" +"#, + ) + .expect("seed config"); + let config = ConfigBuilder::default() + .codex_home(tmp.path().to_path_buf()) + .build() + .await + .expect("load config"); + let (_session, mut turn_context) = make_session_and_context().await; + turn_context.config = Arc::new(config); + + assert_eq!( + custom_mcp_tool_approval_mode(&turn_context, "docs", "read"), + AppToolApproval::Approve + ); + assert_eq!( + custom_mcp_tool_approval_mode(&turn_context, "docs", "search"), + AppToolApproval::Prompt + ); + assert_eq!( + custom_mcp_tool_approval_mode(&turn_context, "unknown", "search"), + AppToolApproval::Auto + ); +} + #[tokio::test] async fn maybe_persist_mcp_tool_approval_reloads_session_config() { let (session, turn_context) = make_session_and_context().await; diff --git a/codex-rs/core/src/plugins/manager_tests.rs b/codex-rs/core/src/plugins/manager_tests.rs index c82e1d174c..2b2f486586 100644 --- a/codex-rs/core/src/plugins/manager_tests.rs +++ b/codex-rs/core/src/plugins/manager_tests.rs @@ -183,6 +183,7 @@ async fn load_plugins_loads_default_skills_and_mcp_servers() { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -520,6 +521,7 @@ async fn load_plugins_uses_manifest_configured_component_paths() { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -630,6 +632,7 @@ async fn load_plugins_ignores_manifest_component_paths_without_dot_slash() { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -788,6 +791,7 @@ fn capability_index_filters_inactive_and_zero_capability_plugins() { disabled_reason: None, startup_timeout_sec: None, tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index a428b32ba2..474d221d27 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -239,6 +239,7 @@ async fn run_code_mode_turn_with_rmcp_config( disabled_reason: None, startup_timeout_sec: Some(Duration::from_secs(10)), tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, diff --git a/codex-rs/core/tests/suite/js_repl.rs b/codex-rs/core/tests/suite/js_repl.rs index 9c5eaef191..450aa08f0b 100644 --- a/codex-rs/core/tests/suite/js_repl.rs +++ b/codex-rs/core/tests/suite/js_repl.rs @@ -629,6 +629,7 @@ async fn js_repl_can_invoke_mcp_tools_by_display_name() -> Result<()> { disabled_reason: None, startup_timeout_sec: Some(Duration::from_secs(10)), tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index 607dc7bf9c..8c7e7133a3 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -151,6 +151,7 @@ fn insert_mcp_server( disabled_reason: None, startup_timeout_sec: Some(Duration::from_secs(10)), tool_timeout_sec: options.tool_timeout_sec, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, diff --git a/codex-rs/core/tests/suite/search_tool.rs b/codex-rs/core/tests/suite/search_tool.rs index eff7bfb427..41f8e623c9 100644 --- a/codex-rs/core/tests/suite/search_tool.rs +++ b/codex-rs/core/tests/suite/search_tool.rs @@ -703,6 +703,7 @@ async fn tool_search_indexes_only_enabled_non_app_mcp_tools() -> Result<()> { disabled_reason: None, startup_timeout_sec: Some(Duration::from_secs(10)), tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: Some(vec!["echo".to_string(), "image".to_string()]), disabled_tools: Some(vec!["image".to_string()]), scopes: None, diff --git a/codex-rs/core/tests/suite/sqlite_state.rs b/codex-rs/core/tests/suite/sqlite_state.rs index acdc1e9d8c..271f7bfc90 100644 --- a/codex-rs/core/tests/suite/sqlite_state.rs +++ b/codex-rs/core/tests/suite/sqlite_state.rs @@ -378,6 +378,7 @@ async fn mcp_call_marks_thread_memory_mode_polluted_when_configured() -> Result< disabled_reason: None, startup_timeout_sec: Some(Duration::from_secs(10)), tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, diff --git a/codex-rs/core/tests/suite/tools.rs b/codex-rs/core/tests/suite/tools.rs index d0bb138188..bc286ce5ea 100644 --- a/codex-rs/core/tests/suite/tools.rs +++ b/codex-rs/core/tests/suite/tools.rs @@ -162,6 +162,7 @@ async fn historical_unavailable_mcp_call_is_exposed_as_placeholder_tool() -> Res disabled_reason: None, startup_timeout_sec: Some(Duration::from_secs(10)), tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, diff --git a/codex-rs/core/tests/suite/truncation.rs b/codex-rs/core/tests/suite/truncation.rs index 33c74633f3..81c22e1b6c 100644 --- a/codex-rs/core/tests/suite/truncation.rs +++ b/codex-rs/core/tests/suite/truncation.rs @@ -382,6 +382,7 @@ async fn mcp_tool_call_output_exceeds_limit_truncated_for_model() -> Result<()> disabled_reason: None, startup_timeout_sec: Some(std::time::Duration::from_secs(10)), tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -480,6 +481,7 @@ async fn mcp_image_output_preserves_image_and_no_text_summary() -> Result<()> { disabled_reason: None, startup_timeout_sec: Some(Duration::from_secs(10)), tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, @@ -761,6 +763,7 @@ async fn mcp_tool_call_output_not_truncated_with_custom_limit() -> Result<()> { disabled_reason: None, startup_timeout_sec: Some(std::time::Duration::from_secs(10)), tool_timeout_sec: None, + default_tools_approval_mode: None, enabled_tools: None, disabled_tools: None, scopes: None, diff --git a/docs/config.md b/docs/config.md index c314ce2283..d583122e01 100644 --- a/docs/config.md +++ b/docs/config.md @@ -28,12 +28,18 @@ resources, review those read/write race conditions before enabling this setting. ## MCP tool approvals -Codex stores per-tool approval overrides for custom MCP servers under -`mcp_servers` in `~/.codex/config.toml`: +Codex stores approval defaults and per-tool overrides for custom MCP servers +under `mcp_servers` in `~/.codex/config.toml`. Set +`default_tools_approval_mode` on the server to apply a default to every tool, +and use per-tool `approval_mode` entries for exceptions: ```toml +[mcp_servers.docs] +command = "docs-server" +default_tools_approval_mode = "approve" + [mcp_servers.docs.tools.search] -approval_mode = "approve" +approval_mode = "prompt" ``` ## Apps (Connectors) From dfff8a7d0352631efe5e9424da4f148b1d8223f5 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 16 Apr 2026 12:29:25 -0700 Subject: [PATCH 03/16] fix: drop lock earlier; was held across send_event().await unnecessarily (#18178) This was flagged by the Codex Security tool: the `state` lock was held longer than necessary, which included being held across an `async` call, increasing the potential for deadlock. While this was flagged by the Codex Security tool, I will look into enabling https://rust-lang.github.io/rust-clippy/stable/index.html#await_holding_lock in a follow-up PR (though unfortunately, that Clippy rule claims it reports false positives when `drop()` is used to drop a guard instead of using the end of block scope to drop). Though I can't seem to find a Clippy rule that checks for opportunities to drop a guard as soon as it is no longer referenced, in general. --- codex-rs/core/src/codex.rs | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 5274d88664..42f165bbf3 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -6696,13 +6696,16 @@ pub(crate) async fn run_turn( break; } Err(CodexErr::InvalidImageRequest()) => { - let mut state = sess.state.lock().await; - error_or_panic( - "Invalid image detected; sanitizing tool output to prevent poisoning", - ); - if state.history.replace_last_turn_images("Invalid image") { - continue; + { + let mut state = sess.state.lock().await; + error_or_panic( + "Invalid image detected; sanitizing tool output to prevent poisoning", + ); + if state.history.replace_last_turn_images("Invalid image") { + continue; + } } + let event = EventMsg::Error(ErrorEvent { message: "Invalid image in your last message. Please remove it and try again." .to_string(), From 62847e7554759b8a5b28e90863e94a5b2211d7db Mon Sep 17 00:00:00 2001 From: starr-openai Date: Thu, 16 Apr 2026 12:34:04 -0700 Subject: [PATCH 04/16] Make thread unsubscribe test deterministic (#18000) ## Summary - replace the unsubscribe-during-turn test's sleep/polling flow with a gated streaming SSE response - add request-count notification support to the streaming SSE test server so the test can wait for the in-flight Responses request deterministically ## Scope - codex-rs/app-server/tests/suite/v2/thread_unsubscribe.rs - codex-rs/core/tests/common/streaming_sse.rs ## Validation - Not run locally; this is a narrow extraction from the prior CI-green branch. --------- Co-authored-by: Codex --- .../tests/suite/v2/thread_unsubscribe.rs | 207 ++++++++++-------- codex-rs/core/tests/common/streaming_sse.rs | 16 ++ 2 files changed, 131 insertions(+), 92 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/thread_unsubscribe.rs b/codex-rs/app-server/tests/suite/v2/thread_unsubscribe.rs index 7a3ae61b72..e91a8654cb 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_unsubscribe.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_unsubscribe.rs @@ -1,14 +1,15 @@ -use anyhow::Context; use anyhow::Result; use app_test_support::McpProcess; -use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_repeating_assistant; -use app_test_support::create_mock_responses_server_sequence_unchecked; -use app_test_support::create_shell_command_sse_response; use app_test_support::to_response; +use codex_app_server_protocol::DynamicToolCallOutputContentItem; +use codex_app_server_protocol::DynamicToolCallParams; +use codex_app_server_protocol::DynamicToolCallResponse; +use codex_app_server_protocol::DynamicToolSpec; use codex_app_server_protocol::ItemStartedNotification; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadLoadedListParams; use codex_app_server_protocol::ThreadLoadedListResponse; @@ -26,61 +27,14 @@ use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; use core_test_support::responses; +use core_test_support::streaming_sse::StreamingSseChunk; +use core_test_support::streaming_sse::start_streaming_sse_server; use pretty_assertions::assert_eq; +use serde_json::json; use tempfile::TempDir; use tokio::time::timeout; const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); -// Windows CI can spend most of the default read timeout launching PowerShell -// before the command finishes and the follow-up model request is sent. -const TURN_COMPLETION_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); - -async fn wait_for_responses_request_count_to_stabilize( - server: &wiremock::MockServer, - expected_count: usize, - settle_duration: std::time::Duration, - timeout_duration: std::time::Duration, -) -> Result<()> { - timeout(timeout_duration, async { - let mut stable_since: Option = None; - loop { - let requests = server - .received_requests() - .await - .context("failed to fetch received requests")?; - let responses_request_count = requests - .iter() - .filter(|request| { - request.method == "POST" && request.url.path().ends_with("/responses") - }) - .count(); - - if responses_request_count > expected_count { - anyhow::bail!( - "expected exactly {expected_count} /responses requests, got {responses_request_count}" - ); - } - - if responses_request_count == expected_count { - match stable_since { - Some(stable_since) if stable_since.elapsed() >= settle_duration => { - return Ok::<(), anyhow::Error>(()); - } - None => stable_since = Some(tokio::time::Instant::now()), - Some(_) => {} - } - } else { - stable_since = None; - } - - tokio::time::sleep(std::time::Duration::from_millis(10)).await; - } - }) - .await??; - - Ok(()) -} - #[tokio::test] async fn thread_unsubscribe_keeps_thread_loaded_until_idle_timeout() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; @@ -132,14 +86,10 @@ async fn thread_unsubscribe_keeps_thread_loaded_until_idle_timeout() -> Result<( #[tokio::test] async fn thread_unsubscribe_during_turn_keeps_turn_running() -> Result<()> { - #[cfg(target_os = "windows")] - let shell_command = vec![ - "Start-Sleep".to_string(), - "-Seconds".to_string(), - "1".to_string(), - ]; - #[cfg(not(target_os = "windows"))] - let shell_command = vec!["sleep".to_string(), "1".to_string()]; + let call_id = "deterministic-wait-call"; + let tool_name = "deterministic_wait"; + let tool_args = json!({}); + let tool_call_arguments = serde_json::to_string(&tool_args)?; let tmp = TempDir::new()?; let codex_home = tmp.path().join("codex_home"); @@ -147,28 +97,61 @@ async fn thread_unsubscribe_during_turn_keeps_turn_running() -> Result<()> { let working_directory = tmp.path().join("workdir"); std::fs::create_dir(&working_directory)?; - let server = create_mock_responses_server_sequence_unchecked(vec![ - create_shell_command_sse_response( - shell_command.clone(), - Some(&working_directory), - Some(10_000), - "call_sleep", - )?, - create_final_assistant_message_sse_response("Done")?, + let (server, mut completions) = start_streaming_sse_server(vec![ + vec![StreamingSseChunk { + gate: None, + body: responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_function_call(call_id, tool_name, &tool_call_arguments), + responses::ev_completed("resp-1"), + ]), + }], + vec![StreamingSseChunk { + gate: None, + body: responses::sse(vec![ + responses::ev_response_created("resp-2"), + responses::ev_assistant_message("msg-1", "Done"), + responses::ev_completed("resp-2"), + ]), + }], ]) .await; - create_config_toml(&codex_home, &server.uri())?; + let first_response_completed = completions.remove(0); + let final_response_completed = completions.remove(0); + create_config_toml(&codex_home, server.uri())?; let mut mcp = McpProcess::new(&codex_home).await?; timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; - let thread_id = start_thread(&mut mcp).await?; + let thread_req = mcp + .send_thread_start_request(ThreadStartParams { + model: Some("mock-model".to_string()), + dynamic_tools: Some(vec![DynamicToolSpec { + name: tool_name.to_string(), + description: "Deterministic wait tool".to_string(), + input_schema: json!({ + "type": "object", + "properties": {}, + "additionalProperties": false, + }), + defer_loading: false, + }]), + ..Default::default() + }) + .await?; + let thread_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + let thread_id = thread.id; let turn_req = mcp .send_turn_start_request(TurnStartParams { thread_id: thread_id.clone(), input: vec![V2UserInput::Text { - text: "run sleep".to_string(), + text: "run deterministic tool".to_string(), text_elements: Vec::new(), }], cwd: Some(working_directory), @@ -184,9 +167,37 @@ async fn thread_unsubscribe_during_turn_keeps_turn_running() -> Result<()> { timeout( DEFAULT_READ_TIMEOUT, - wait_for_command_execution_item_started(&mut mcp), + server.wait_for_request_count(/*count*/ 1), + ) + .await?; + timeout(DEFAULT_READ_TIMEOUT, first_response_completed).await??; + + let started = timeout( + DEFAULT_READ_TIMEOUT, + wait_for_dynamic_tool_started(&mut mcp, call_id), ) .await??; + assert_eq!(started.thread_id, thread_id); + + let request = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_request_message(), + ) + .await??; + let (request_id, params) = match request { + ServerRequest::DynamicToolCall { request_id, params } => (request_id, params), + other => panic!("expected DynamicToolCall request, got {other:?}"), + }; + assert_eq!( + params, + DynamicToolCallParams { + thread_id: thread_id.clone(), + turn_id: started.turn_id, + call_id: call_id.to_string(), + tool: tool_name.to_string(), + arguments: tool_args, + } + ); let unsubscribe_id = mcp .send_thread_unsubscribe_request(ThreadUnsubscribeParams { @@ -201,22 +212,29 @@ async fn thread_unsubscribe_during_turn_keeps_turn_running() -> Result<()> { let unsubscribe = to_response::(unsubscribe_resp)?; assert_eq!(unsubscribe.status, ThreadUnsubscribeStatus::Unsubscribed); - assert!( - timeout( - std::time::Duration::from_millis(250), - mcp.read_stream_until_notification_message("thread/closed"), - ) - .await - .is_err() + let closed_while_tool_call_blocked = timeout( + std::time::Duration::from_millis(250), + mcp.read_stream_until_notification_message("thread/closed"), ); + let closed_while_tool_call_blocked = closed_while_tool_call_blocked.await; + assert!(closed_while_tool_call_blocked.is_err()); - wait_for_responses_request_count_to_stabilize( - &server, - /*expected_count*/ 2, - std::time::Duration::from_millis(200), - TURN_COMPLETION_TIMEOUT, + let response = DynamicToolCallResponse { + content_items: vec![DynamicToolCallOutputContentItem::InputText { + text: "dynamic-ok".to_string(), + }], + success: true, + }; + mcp.send_response(request_id, serde_json::to_value(response)?) + .await?; + + timeout( + DEFAULT_READ_TIMEOUT, + server.wait_for_request_count(/*count*/ 2), ) .await?; + timeout(DEFAULT_READ_TIMEOUT, final_response_completed).await??; + server.shutdown().await; Ok(()) } @@ -355,15 +373,20 @@ async fn thread_unsubscribe_reports_not_subscribed_before_idle_unload() -> Resul Ok(()) } -async fn wait_for_command_execution_item_started(mcp: &mut McpProcess) -> Result<()> { +async fn wait_for_dynamic_tool_started( + mcp: &mut McpProcess, + call_id: &str, +) -> Result { loop { - let started_notif = mcp + let notification = mcp .read_stream_until_notification_message("item/started") .await?; - let started_params = started_notif.params.context("item/started params")?; - let started: ItemStartedNotification = serde_json::from_value(started_params)?; - if let ThreadItem::CommandExecution { .. } = started.item { - return Ok(()); + let Some(params) = notification.params else { + continue; + }; + let started: ItemStartedNotification = serde_json::from_value(params)?; + if matches!(&started.item, ThreadItem::DynamicToolCall { id, .. } if id == call_id) { + return Ok(started); } } } diff --git a/codex-rs/core/tests/common/streaming_sse.rs b/codex-rs/core/tests/common/streaming_sse.rs index 82edcd39d7..a86256ed5a 100644 --- a/codex-rs/core/tests/common/streaming_sse.rs +++ b/codex-rs/core/tests/common/streaming_sse.rs @@ -7,6 +7,7 @@ use tokio::io::AsyncReadExt; use tokio::io::AsyncWriteExt; use tokio::net::TcpListener; use tokio::sync::Mutex as TokioMutex; +use tokio::sync::Notify; use tokio::sync::oneshot; /// Streaming SSE chunk payload gated by a per-chunk signal. @@ -20,6 +21,7 @@ pub struct StreamingSseChunk { pub struct StreamingSseServer { uri: String, requests: Arc>>>, + request_notify: Arc, shutdown: oneshot::Sender<()>, task: tokio::task::JoinHandle<()>, } @@ -33,6 +35,15 @@ impl StreamingSseServer { self.requests.lock().await.clone() } + pub async fn wait_for_request_count(&self, count: usize) { + loop { + if self.requests.lock().await.len() >= count { + return; + } + self.request_notify.notified().await; + } + } + pub async fn shutdown(self) { let _ = self.shutdown.send(()); let _ = self.task.await; @@ -67,7 +78,9 @@ pub async fn start_streaming_sse_server( completions: VecDeque::from(completion_senders), })); let requests = Arc::new(TokioMutex::new(Vec::new())); + let request_notify = Arc::new(Notify::new()); let requests_for_task = Arc::clone(&requests); + let request_notify_for_task = Arc::clone(&request_notify); let (shutdown_tx, mut shutdown_rx) = oneshot::channel(); let task = tokio::spawn(async move { @@ -78,6 +91,7 @@ pub async fn start_streaming_sse_server( let (mut stream, _) = accept_res.expect("accept streaming SSE connection"); let state = Arc::clone(&state); let requests = Arc::clone(&requests_for_task); + let request_notify = Arc::clone(&request_notify_for_task); tokio::spawn(async move { let (request, body_prefix) = read_http_request(&mut stream).await; let Some((method, path)) = parse_request_line(&request) else { @@ -113,6 +127,7 @@ pub async fn start_streaming_sse_server( } }; requests.lock().await.push(body); + request_notify.notify_one(); let Some((chunks, completion)) = take_next_stream(&state).await else { let _ = write_http_response(&mut stream, /*status*/ 500, "no responses queued", "text/plain").await; return; @@ -149,6 +164,7 @@ pub async fn start_streaming_sse_server( StreamingSseServer { uri, requests, + request_notify, shutdown: shutdown_tx, task, }, From 8720b7bdce451c2f805894501317579b7c667c39 Mon Sep 17 00:00:00 2001 From: Abhinav Date: Thu, 16 Apr 2026 12:43:16 -0700 Subject: [PATCH 05/16] Add codex_hook_run analytics event (#17996) # Why Add product analytics for hook handler executions so we can understand which hooks are running, where they came from, and whether they completed, failed, stopped, or blocked work. # What - add the new `codex_hook_run` analytics event and payload plumbing in `codex-rs/analytics` - emit hook-run analytics from the shared hook completion path in `codex-rs/core` - classify hook source from the loaded hook path as `system`, `user`, `project`, or `unknown` ``` { "event_type": "codex_hook_run", "event_params": { "thread_id": "string", "turn_id": "string", "model_slug": "string", "hook_name": "string, // any HookEventName "hook_source": "system | user | project | unknown", "status": "completed | failed | stopped | blocked" } } ``` --------- Co-authored-by: Codex --- .../analytics/src/analytics_client_tests.rs | 141 ++++++++++ codex-rs/analytics/src/client.rs | 8 + codex-rs/analytics/src/events.rs | 66 +++++ codex-rs/analytics/src/facts.rs | 15 + codex-rs/analytics/src/lib.rs | 1 + codex-rs/analytics/src/reducer.rs | 14 + .../schema/json/ServerNotification.json | 21 ++ .../codex_app_server_protocol.schemas.json | 21 ++ .../codex_app_server_protocol.v2.schemas.json | 21 ++ .../json/v2/HookCompletedNotification.json | 21 ++ .../json/v2/HookStartedNotification.json | 21 ++ .../schema/typescript/v2/HookRunSummary.ts | 3 +- .../schema/typescript/v2/HookSource.ts | 5 + .../schema/typescript/v2/index.ts | 1 + .../app-server-protocol/src/protocol/v2.rs | 21 ++ codex-rs/config/src/lib.rs | 2 + codex-rs/core/src/codex.rs | 7 +- codex-rs/core/src/hook_runtime.rs | 105 ++++++- codex-rs/hooks/src/engine/discovery.rs | 264 +++++++++++------- codex-rs/hooks/src/engine/dispatcher.rs | 4 + codex-rs/hooks/src/engine/mod.rs | 2 + codex-rs/hooks/src/events/post_tool_use.rs | 1 + codex-rs/hooks/src/events/pre_tool_use.rs | 1 + codex-rs/hooks/src/events/session_start.rs | 1 + codex-rs/hooks/src/events/stop.rs | 1 + .../hooks/src/events/user_prompt_submit.rs | 1 + codex-rs/protocol/src/protocol.rs | 16 ++ codex-rs/tui/src/app.rs | 2 + codex-rs/tui/src/chatwidget.rs | 1 + codex-rs/tui/src/chatwidget/tests/helpers.rs | 2 + .../src/chatwidget/tests/status_and_layout.rs | 5 + codex-rs/tui/src/history_cell/hook_cell.rs | 1 + 32 files changed, 682 insertions(+), 114 deletions(-) create mode 100644 codex-rs/app-server-protocol/schema/typescript/v2/HookSource.ts diff --git a/codex-rs/analytics/src/analytics_client_tests.rs b/codex-rs/analytics/src/analytics_client_tests.rs index eb46674156..4b13b2819e 100644 --- a/codex-rs/analytics/src/analytics_client_tests.rs +++ b/codex-rs/analytics/src/analytics_client_tests.rs @@ -4,6 +4,7 @@ use crate::events::CodexAppMentionedEventRequest; use crate::events::CodexAppServerClientMetadata; use crate::events::CodexAppUsedEventRequest; use crate::events::CodexCompactionEventRequest; +use crate::events::CodexHookRunEventRequest; use crate::events::CodexPluginEventRequest; use crate::events::CodexPluginUsedEventRequest; use crate::events::CodexRuntimeMetadata; @@ -12,6 +13,7 @@ use crate::events::ThreadInitializedEvent; use crate::events::ThreadInitializedEventParams; use crate::events::TrackEventRequest; use crate::events::codex_app_metadata; +use crate::events::codex_hook_run_metadata; use crate::events::codex_plugin_metadata; use crate::events::codex_plugin_used_metadata; use crate::events::subagent_thread_started_event_request; @@ -28,6 +30,8 @@ use crate::facts::CompactionStatus; use crate::facts::CompactionStrategy; use crate::facts::CompactionTrigger; use crate::facts::CustomAnalyticsFact; +use crate::facts::HookRunFact; +use crate::facts::HookRunInput; use crate::facts::InputError; use crate::facts::InvocationType; use crate::facts::PluginState; @@ -81,6 +85,9 @@ use codex_plugin::PluginTelemetryMetadata; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::config_types::ModeKind; use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookRunStatus; +use codex_protocol::protocol::HookSource; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; @@ -1282,6 +1289,109 @@ fn plugin_management_event_serializes_expected_shape() { ); } +#[test] +fn hook_run_event_serializes_expected_shape() { + let tracking = TrackEventsContext { + model_slug: "gpt-5".to_string(), + thread_id: "thread-3".to_string(), + turn_id: "turn-3".to_string(), + }; + let event = TrackEventRequest::HookRun(CodexHookRunEventRequest { + event_type: "codex_hook_run", + event_params: codex_hook_run_metadata( + &tracking, + HookRunFact { + event_name: HookEventName::PreToolUse, + hook_source: HookSource::User, + status: HookRunStatus::Completed, + }, + ), + }); + + let payload = serde_json::to_value(&event).expect("serialize hook run event"); + + assert_eq!( + payload, + json!({ + "event_type": "codex_hook_run", + "event_params": { + "thread_id": "thread-3", + "turn_id": "turn-3", + "model_slug": "gpt-5", + "hook_name": "PreToolUse", + "hook_source": "user", + "status": "completed" + } + }) + ); +} + +#[test] +fn hook_run_metadata_maps_sources_and_statuses() { + let tracking = TrackEventsContext { + model_slug: "gpt-5".to_string(), + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + }; + + let system = serde_json::to_value(codex_hook_run_metadata( + &tracking, + HookRunFact { + event_name: HookEventName::SessionStart, + hook_source: HookSource::System, + status: HookRunStatus::Completed, + }, + )) + .expect("serialize system hook"); + let project = serde_json::to_value(codex_hook_run_metadata( + &tracking, + HookRunFact { + event_name: HookEventName::Stop, + hook_source: HookSource::Project, + status: HookRunStatus::Blocked, + }, + )) + .expect("serialize project hook"); + let unknown = serde_json::to_value(codex_hook_run_metadata( + &tracking, + HookRunFact { + event_name: HookEventName::UserPromptSubmit, + hook_source: HookSource::Unknown, + status: HookRunStatus::Failed, + }, + )) + .expect("serialize unknown hook"); + + assert_eq!(system["hook_source"], "system"); + assert_eq!(system["status"], "completed"); + assert_eq!(project["hook_source"], "project"); + assert_eq!(project["status"], "blocked"); + assert_eq!(unknown["hook_source"], "unknown"); + assert_eq!(unknown["status"], "failed"); +} + +#[test] +fn hook_run_metadata_maps_stopped_status() { + let tracking = TrackEventsContext { + model_slug: "gpt-5".to_string(), + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + }; + + let stopped = serde_json::to_value(codex_hook_run_metadata( + &tracking, + HookRunFact { + event_name: HookEventName::Stop, + hook_source: HookSource::User, + status: HookRunStatus::Stopped, + }, + )) + .expect("serialize stopped hook"); + + assert_eq!(stopped["hook_source"], "user"); + assert_eq!(stopped["status"], "stopped"); +} + #[test] fn plugin_used_dedupe_is_keyed_by_turn_and_plugin() { let (sender, _receiver) = mpsc::channel(1); @@ -1359,6 +1469,37 @@ async fn reducer_ingests_skill_invoked_fact() { ); } +#[tokio::test] +async fn reducer_ingests_hook_run_fact() { + let mut reducer = AnalyticsReducer::default(); + let mut events = Vec::new(); + + reducer + .ingest( + AnalyticsFact::Custom(CustomAnalyticsFact::HookRun(HookRunInput { + tracking: TrackEventsContext { + model_slug: "gpt-5".to_string(), + thread_id: "thread-1".to_string(), + turn_id: "turn-1".to_string(), + }, + hook: HookRunFact { + event_name: HookEventName::PostToolUse, + hook_source: HookSource::Unknown, + status: HookRunStatus::Failed, + }, + })), + &mut events, + ) + .await; + + let payload = serde_json::to_value(&events).expect("serialize events"); + assert_eq!(payload.as_array().expect("events array").len(), 1); + assert_eq!(payload[0]["event_type"], "codex_hook_run"); + assert_eq!(payload[0]["event_params"]["hook_name"], "PostToolUse"); + assert_eq!(payload[0]["event_params"]["hook_source"], "unknown"); + assert_eq!(payload[0]["event_params"]["status"], "failed"); +} + #[tokio::test] async fn reducer_ingests_app_and_plugin_facts() { let mut reducer = AnalyticsReducer::default(); diff --git a/codex-rs/analytics/src/client.rs b/codex-rs/analytics/src/client.rs index 5ba60e3ef4..1a4b5defe9 100644 --- a/codex-rs/analytics/src/client.rs +++ b/codex-rs/analytics/src/client.rs @@ -9,6 +9,8 @@ use crate::facts::AppInvocation; use crate::facts::AppMentionedInput; use crate::facts::AppUsedInput; use crate::facts::CustomAnalyticsFact; +use crate::facts::HookRunFact; +use crate::facts::HookRunInput; use crate::facts::PluginState; use crate::facts::PluginStateChangedInput; use crate::facts::SkillInvocation; @@ -191,6 +193,12 @@ impl AnalyticsEventsClient { ))); } + pub fn track_hook_run(&self, tracking: TrackEventsContext, hook: HookRunFact) { + self.record_fact(AnalyticsFact::Custom(CustomAnalyticsFact::HookRun( + HookRunInput { tracking, hook }, + ))); + } + pub fn track_plugin_used(&self, tracking: TrackEventsContext, plugin: PluginTelemetryMetadata) { if !self.queue.should_enqueue_plugin_used(&tracking, &plugin) { return; diff --git a/codex-rs/analytics/src/events.rs b/codex-rs/analytics/src/events.rs index eb312da140..d3aa5263da 100644 --- a/codex-rs/analytics/src/events.rs +++ b/codex-rs/analytics/src/events.rs @@ -1,5 +1,6 @@ use crate::facts::AppInvocation; use crate::facts::CodexCompactionEvent; +use crate::facts::HookRunFact; use crate::facts::InvocationType; use crate::facts::PluginState; use crate::facts::SubAgentThreadStartedInput; @@ -15,6 +16,9 @@ use codex_plugin::PluginTelemetryMetadata; use codex_protocol::approvals::NetworkApprovalProtocol; use codex_protocol::models::PermissionProfile; use codex_protocol::models::SandboxPermissions; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookRunStatus; +use codex_protocol::protocol::HookSource; use codex_protocol::protocol::SubAgentSource; use serde::Serialize; @@ -39,6 +43,7 @@ pub(crate) enum TrackEventRequest { GuardianReview(Box), AppMentioned(CodexAppMentionedEventRequest), AppUsed(CodexAppUsedEventRequest), + HookRun(CodexHookRunEventRequest), Compaction(Box), TurnEvent(Box), TurnSteer(CodexTurnSteerEventRequest), @@ -300,6 +305,22 @@ pub(crate) struct CodexAppUsedEventRequest { pub(crate) event_params: CodexAppMetadata, } +#[derive(Serialize)] +pub(crate) struct CodexHookRunMetadata { + pub(crate) thread_id: Option, + pub(crate) turn_id: Option, + pub(crate) model_slug: Option, + pub(crate) hook_name: Option, + pub(crate) hook_source: Option<&'static str>, + pub(crate) status: Option, +} + +#[derive(Serialize)] +pub(crate) struct CodexHookRunEventRequest { + pub(crate) event_type: &'static str, + pub(crate) event_params: CodexHookRunMetadata, +} + #[derive(Serialize)] pub(crate) struct CodexCompactionEventParams { pub(crate) thread_id: String, @@ -529,6 +550,43 @@ pub(crate) fn codex_plugin_used_metadata( } } +pub(crate) fn codex_hook_run_metadata( + tracking: &TrackEventsContext, + hook: HookRunFact, +) -> CodexHookRunMetadata { + CodexHookRunMetadata { + thread_id: Some(tracking.thread_id.clone()), + turn_id: Some(tracking.turn_id.clone()), + model_slug: Some(tracking.model_slug.clone()), + hook_name: Some(analytics_hook_event_name(hook.event_name).to_owned()), + hook_source: Some(analytics_hook_source(hook.hook_source)), + status: Some(analytics_hook_status(hook.status)), + } +} + +fn analytics_hook_event_name(event_name: HookEventName) -> &'static str { + match event_name { + HookEventName::PreToolUse => "PreToolUse", + HookEventName::PostToolUse => "PostToolUse", + HookEventName::SessionStart => "SessionStart", + HookEventName::UserPromptSubmit => "UserPromptSubmit", + HookEventName::Stop => "Stop", + } +} + +fn analytics_hook_source(source: HookSource) -> &'static str { + match source { + HookSource::System => "system", + HookSource::User => "user", + HookSource::Project => "project", + HookSource::Mdm => "mdm", + HookSource::SessionFlags => "session_flags", + HookSource::LegacyManagedConfigFile => "legacy_managed_config_file", + HookSource::LegacyManagedConfigMdm => "legacy_managed_config_mdm", + HookSource::Unknown => "unknown", + } +} + pub(crate) fn current_runtime_metadata() -> CodexRuntimeMetadata { let os_info = os_info::get(); CodexRuntimeMetadata { @@ -586,3 +644,11 @@ pub(crate) fn subagent_parent_thread_id(subagent_source: &SubAgentSource) -> Opt _ => None, } } + +fn analytics_hook_status(status: HookRunStatus) -> HookRunStatus { + match status { + // Running is unexpected here and normalized defensively. + HookRunStatus::Running => HookRunStatus::Failed, + other => other, + } +} diff --git a/codex-rs/analytics/src/facts.rs b/codex-rs/analytics/src/facts.rs index 5590fbfa6d..1d371acb1c 100644 --- a/codex-rs/analytics/src/facts.rs +++ b/codex-rs/analytics/src/facts.rs @@ -15,6 +15,9 @@ use codex_protocol::config_types::ReasoningSummary; use codex_protocol::config_types::ServiceTier; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookRunStatus; +use codex_protocol::protocol::HookSource; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SkillScope; @@ -298,6 +301,7 @@ pub(crate) enum CustomAnalyticsFact { SkillInvoked(SkillInvokedInput), AppMentioned(AppMentionedInput), AppUsed(AppUsedInput), + HookRun(HookRunInput), PluginUsed(PluginUsedInput), PluginStateChanged(PluginStateChangedInput), } @@ -317,6 +321,17 @@ pub(crate) struct AppUsedInput { pub app: AppInvocation, } +pub(crate) struct HookRunInput { + pub tracking: TrackEventsContext, + pub hook: HookRunFact, +} + +pub struct HookRunFact { + pub event_name: HookEventName, + pub hook_source: HookSource, + pub status: HookRunStatus, +} + pub(crate) struct PluginUsedInput { pub tracking: TrackEventsContext, pub plugin: PluginTelemetryMetadata, diff --git a/codex-rs/analytics/src/lib.rs b/codex-rs/analytics/src/lib.rs index 1a1a123152..03f485a1c1 100644 --- a/codex-rs/analytics/src/lib.rs +++ b/codex-rs/analytics/src/lib.rs @@ -29,6 +29,7 @@ pub use facts::CompactionReason; pub use facts::CompactionStatus; pub use facts::CompactionStrategy; pub use facts::CompactionTrigger; +pub use facts::HookRunFact; pub use facts::InputError; pub use facts::InvocationType; pub use facts::SkillInvocation; diff --git a/codex-rs/analytics/src/reducer.rs b/codex-rs/analytics/src/reducer.rs index a5e5165bd2..a6ce3fc831 100644 --- a/codex-rs/analytics/src/reducer.rs +++ b/codex-rs/analytics/src/reducer.rs @@ -3,6 +3,7 @@ use crate::events::CodexAppMentionedEventRequest; use crate::events::CodexAppServerClientMetadata; use crate::events::CodexAppUsedEventRequest; use crate::events::CodexCompactionEventRequest; +use crate::events::CodexHookRunEventRequest; use crate::events::CodexPluginEventRequest; use crate::events::CodexPluginUsedEventRequest; use crate::events::CodexRuntimeMetadata; @@ -20,6 +21,7 @@ use crate::events::ThreadInitializedEventParams; use crate::events::TrackEventRequest; use crate::events::codex_app_metadata; use crate::events::codex_compaction_event_params; +use crate::events::codex_hook_run_metadata; use crate::events::codex_plugin_metadata; use crate::events::codex_plugin_used_metadata; use crate::events::plugin_state_event_type; @@ -32,6 +34,7 @@ use crate::facts::AppMentionedInput; use crate::facts::AppUsedInput; use crate::facts::CodexCompactionEvent; use crate::facts::CustomAnalyticsFact; +use crate::facts::HookRunInput; use crate::facts::PluginState; use crate::facts::PluginStateChangedInput; use crate::facts::PluginUsedInput; @@ -217,6 +220,9 @@ impl AnalyticsReducer { CustomAnalyticsFact::AppUsed(input) => { self.ingest_app_used(input, out); } + CustomAnalyticsFact::HookRun(input) => { + self.ingest_hook_run(input, out); + } CustomAnalyticsFact::PluginUsed(input) => { self.ingest_plugin_used(input, out); } @@ -442,6 +448,14 @@ impl AnalyticsReducer { })); } + fn ingest_hook_run(&mut self, input: HookRunInput, out: &mut Vec) { + let HookRunInput { tracking, hook } = input; + out.push(TrackEventRequest::HookRun(CodexHookRunEventRequest { + event_type: "codex_hook_run", + event_params: codex_hook_run_metadata(&tracking, hook), + })); + } + fn ingest_plugin_used(&mut self, input: PluginUsedInput, out: &mut Vec) { let PluginUsedInput { tracking, plugin } = input; out.push(TrackEventRequest::PluginUsed(CodexPluginUsedEventRequest { diff --git a/codex-rs/app-server-protocol/schema/json/ServerNotification.json b/codex-rs/app-server-protocol/schema/json/ServerNotification.json index b84fbe9a02..d02a28e370 100644 --- a/codex-rs/app-server-protocol/schema/json/ServerNotification.json +++ b/codex-rs/app-server-protocol/schema/json/ServerNotification.json @@ -1517,6 +1517,14 @@ "scope": { "$ref": "#/definitions/HookScope" }, + "source": { + "allOf": [ + { + "$ref": "#/definitions/HookSource" + } + ], + "default": "unknown" + }, "sourcePath": { "$ref": "#/definitions/AbsolutePathBuf" }, @@ -1555,6 +1563,19 @@ ], "type": "string" }, + "HookSource": { + "enum": [ + "system", + "user", + "project", + "mdm", + "sessionFlags", + "legacyManagedConfigFile", + "legacyManagedConfigMdm", + "unknown" + ], + "type": "string" + }, "HookStartedNotification": { "properties": { "run": { diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 943b8480f9..e87fa17928 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -8597,6 +8597,14 @@ "scope": { "$ref": "#/definitions/v2/HookScope" }, + "source": { + "allOf": [ + { + "$ref": "#/definitions/v2/HookSource" + } + ], + "default": "unknown" + }, "sourcePath": { "$ref": "#/definitions/v2/AbsolutePathBuf" }, @@ -8635,6 +8643,19 @@ ], "type": "string" }, + "HookSource": { + "enum": [ + "system", + "user", + "project", + "mdm", + "sessionFlags", + "legacyManagedConfigFile", + "legacyManagedConfigMdm", + "unknown" + ], + "type": "string" + }, "HookStartedNotification": { "$schema": "http://json-schema.org/draft-07/schema#", "properties": { diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index 1717dfb393..8254ad0127 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -5325,6 +5325,14 @@ "scope": { "$ref": "#/definitions/HookScope" }, + "source": { + "allOf": [ + { + "$ref": "#/definitions/HookSource" + } + ], + "default": "unknown" + }, "sourcePath": { "$ref": "#/definitions/AbsolutePathBuf" }, @@ -5363,6 +5371,19 @@ ], "type": "string" }, + "HookSource": { + "enum": [ + "system", + "user", + "project", + "mdm", + "sessionFlags", + "legacyManagedConfigFile", + "legacyManagedConfigMdm", + "unknown" + ], + "type": "string" + }, "HookStartedNotification": { "$schema": "http://json-schema.org/draft-07/schema#", "properties": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/HookCompletedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/HookCompletedNotification.json index 2faa347f8e..4f444961a9 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/HookCompletedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/HookCompletedNotification.json @@ -106,6 +106,14 @@ "scope": { "$ref": "#/definitions/HookScope" }, + "source": { + "allOf": [ + { + "$ref": "#/definitions/HookSource" + } + ], + "default": "unknown" + }, "sourcePath": { "$ref": "#/definitions/AbsolutePathBuf" }, @@ -143,6 +151,19 @@ "turn" ], "type": "string" + }, + "HookSource": { + "enum": [ + "system", + "user", + "project", + "mdm", + "sessionFlags", + "legacyManagedConfigFile", + "legacyManagedConfigMdm", + "unknown" + ], + "type": "string" } }, "properties": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/HookStartedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/HookStartedNotification.json index 496175d01b..6d439c95ca 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/HookStartedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/HookStartedNotification.json @@ -106,6 +106,14 @@ "scope": { "$ref": "#/definitions/HookScope" }, + "source": { + "allOf": [ + { + "$ref": "#/definitions/HookSource" + } + ], + "default": "unknown" + }, "sourcePath": { "$ref": "#/definitions/AbsolutePathBuf" }, @@ -143,6 +151,19 @@ "turn" ], "type": "string" + }, + "HookSource": { + "enum": [ + "system", + "user", + "project", + "mdm", + "sessionFlags", + "legacyManagedConfigFile", + "legacyManagedConfigMdm", + "unknown" + ], + "type": "string" } }, "properties": { diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/HookRunSummary.ts b/codex-rs/app-server-protocol/schema/typescript/v2/HookRunSummary.ts index f6d4b75378..75ab780b93 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/HookRunSummary.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/HookRunSummary.ts @@ -8,5 +8,6 @@ import type { HookHandlerType } from "./HookHandlerType"; import type { HookOutputEntry } from "./HookOutputEntry"; import type { HookRunStatus } from "./HookRunStatus"; import type { HookScope } from "./HookScope"; +import type { HookSource } from "./HookSource"; -export type HookRunSummary = { id: string, eventName: HookEventName, handlerType: HookHandlerType, executionMode: HookExecutionMode, scope: HookScope, sourcePath: AbsolutePathBuf, displayOrder: bigint, status: HookRunStatus, statusMessage: string | null, startedAt: bigint, completedAt: bigint | null, durationMs: bigint | null, entries: Array, }; +export type HookRunSummary = { id: string, eventName: HookEventName, handlerType: HookHandlerType, executionMode: HookExecutionMode, scope: HookScope, sourcePath: AbsolutePathBuf, source: HookSource, displayOrder: bigint, status: HookRunStatus, statusMessage: string | null, startedAt: bigint, completedAt: bigint | null, durationMs: bigint | null, entries: Array, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/HookSource.ts b/codex-rs/app-server-protocol/schema/typescript/v2/HookSource.ts new file mode 100644 index 0000000000..7edf61f918 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/HookSource.ts @@ -0,0 +1,5 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +export type HookSource = "system" | "user" | "project" | "mdm" | "sessionFlags" | "legacyManagedConfigFile" | "legacyManagedConfigMdm" | "unknown"; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts index d432530b42..69cec75e77 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts @@ -139,6 +139,7 @@ export type { HookPromptFragment } from "./HookPromptFragment"; export type { HookRunStatus } from "./HookRunStatus"; export type { HookRunSummary } from "./HookRunSummary"; export type { HookScope } from "./HookScope"; +export type { HookSource } from "./HookSource"; export type { HookStartedNotification } from "./HookStartedNotification"; export type { ItemCompletedNotification } from "./ItemCompletedNotification"; export type { ItemGuardianApprovalReviewCompletedNotification } from "./ItemGuardianApprovalReviewCompletedNotification"; diff --git a/codex-rs/app-server-protocol/src/protocol/v2.rs b/codex-rs/app-server-protocol/src/protocol/v2.rs index 155c67e7bf..5173549326 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2.rs @@ -65,6 +65,7 @@ use codex_protocol::protocol::HookOutputEntryKind as CoreHookOutputEntryKind; use codex_protocol::protocol::HookRunStatus as CoreHookRunStatus; use codex_protocol::protocol::HookRunSummary as CoreHookRunSummary; use codex_protocol::protocol::HookScope as CoreHookScope; +use codex_protocol::protocol::HookSource as CoreHookSource; use codex_protocol::protocol::ModelRerouteReason as CoreModelRerouteReason; use codex_protocol::protocol::NetworkAccess as CoreNetworkAccess; use codex_protocol::protocol::NonSteerableTurnKind as CoreNonSteerableTurnKind; @@ -402,6 +403,23 @@ v2_enum_from_core!( } ); +v2_enum_from_core!( + pub enum HookSource from CoreHookSource { + System, + User, + Project, + Mdm, + SessionFlags, + LegacyManagedConfigFile, + LegacyManagedConfigMdm, + Unknown, + } +); + +fn default_hook_source() -> HookSource { + HookSource::Unknown +} + v2_enum_from_core!( pub enum HookRunStatus from CoreHookRunStatus { Running, Completed, Failed, Blocked, Stopped @@ -449,6 +467,8 @@ pub struct HookRunSummary { pub execution_mode: HookExecutionMode, pub scope: HookScope, pub source_path: AbsolutePathBuf, + #[serde(default = "default_hook_source")] + pub source: HookSource, pub display_order: i64, pub status: HookRunStatus, pub status_message: Option, @@ -467,6 +487,7 @@ impl From for HookRunSummary { execution_mode: value.execution_mode.into(), scope: value.scope.into(), source_path: value.source_path, + source: value.source.into(), display_order: value.display_order, status: value.status.into(), status_message: value.status_message, diff --git a/codex-rs/config/src/lib.rs b/codex-rs/config/src/lib.rs index ffed561984..66a6fc22fc 100644 --- a/codex-rs/config/src/lib.rs +++ b/codex-rs/config/src/lib.rs @@ -86,3 +86,5 @@ pub use state::ConfigLayerEntry; pub use state::ConfigLayerStack; pub use state::ConfigLayerStackOrdering; pub use state::LoaderOverrides; + +pub use codex_app_server_protocol::ConfigLayerSource; diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 42f165bbf3..844d585be5 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -24,6 +24,7 @@ use crate::compact_remote::run_inline_remote_auto_compact_task; use crate::config::ManagedFeatures; use crate::connectors; use crate::exec_policy::ExecPolicyManager; +use crate::hook_runtime::emit_hook_completed_events; use crate::installation_id::resolve_installation_id; use crate::mcp_tool_exposure::build_mcp_tool_exposure; use crate::parse_turn_item; @@ -6601,10 +6602,8 @@ pub(crate) async fn run_turn( .await; } let stop_outcome = sess.hooks().run_stop(stop_request).await; - for completed in stop_outcome.hook_events { - sess.send_event(&turn_context, EventMsg::HookCompleted(completed)) - .await; - } + emit_hook_completed_events(&sess, &turn_context, stop_outcome.hook_events) + .await; if stop_outcome.should_block { if let Some(hook_prompt_message) = build_hook_prompt_message(&stop_outcome.continuation_fragments) diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index d2d672fcd4..58cac101b0 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -1,6 +1,8 @@ use std::future::Future; use std::sync::Arc; +use codex_analytics::HookRunFact; +use codex_analytics::build_track_events_context; use codex_hooks::PostToolUseOutcome; use codex_hooks::PostToolUseRequest; use codex_hooks::PreToolUseOutcome; @@ -316,17 +318,52 @@ async fn emit_hook_started_events( } } -async fn emit_hook_completed_events( +pub(crate) async fn emit_hook_completed_events( sess: &Arc, turn_context: &Arc, completed_events: Vec, ) { for completed in completed_events { + track_hook_completed_analytics(sess, turn_context, &completed); sess.send_event(turn_context, EventMsg::HookCompleted(completed)) .await; } } +fn track_hook_completed_analytics( + sess: &Arc, + turn_context: &Arc, + completed: &HookCompletedEvent, +) { + let (tracking, hook) = + hook_run_analytics_payload(sess.conversation_id.to_string(), turn_context, completed); + sess.services + .analytics_events_client + .track_hook_run(tracking, hook); +} + +fn hook_run_analytics_payload( + thread_id: String, + turn_context: &TurnContext, + completed: &HookCompletedEvent, +) -> (codex_analytics::TrackEventsContext, HookRunFact) { + ( + build_track_events_context( + turn_context.model_info.slug.clone(), + thread_id, + completed + .turn_id + .clone() + .unwrap_or_else(|| turn_context.sub_id.clone()), + ), + HookRunFact { + event_name: completed.run.event_name, + hook_source: completed.run.source, + status: completed.run.status, + }, + ) +} + fn hook_permission_mode(turn_context: &TurnContext) -> String { match turn_context.approval_policy.value() { AskForApproval::Never => "bypassPermissions", @@ -341,9 +378,21 @@ fn hook_permission_mode(turn_context: &TurnContext) -> String { #[cfg(test)] mod tests { use codex_protocol::models::ContentItem; + use codex_protocol::protocol::HookEventName; + use codex_protocol::protocol::HookExecutionMode; + use codex_protocol::protocol::HookHandlerType; + use codex_protocol::protocol::HookRunStatus; + use codex_protocol::protocol::HookScope; + use codex_protocol::protocol::HookSource; use pretty_assertions::assert_eq; use super::additional_context_messages; + use super::hook_run_analytics_payload; + use crate::codex::make_session_and_context; + use codex_protocol::protocol::HookCompletedEvent; + use codex_protocol::protocol::HookRunSummary; + use codex_utils_absolute_path::test_support::PathBufExt; + use codex_utils_absolute_path::test_support::test_path_buf; #[test] fn additional_context_messages_stay_separate_and_ordered() { @@ -378,4 +427,58 @@ mod tests { ], ); } + + #[tokio::test] + async fn hook_run_analytics_payload_uses_completed_turn_id() { + let (_session, turn_context) = make_session_and_context().await; + let completed = HookCompletedEvent { + turn_id: Some("turn-from-hook".to_string()), + run: sample_hook_run(HookRunStatus::Blocked, HookSource::Project), + }; + + let (tracking, hook) = + hook_run_analytics_payload("thread-123".to_string(), &turn_context, &completed); + + assert_eq!(tracking.thread_id, "thread-123"); + assert_eq!(tracking.turn_id, "turn-from-hook"); + assert_eq!(tracking.model_slug, turn_context.model_info.slug); + assert_eq!(hook.event_name, HookEventName::Stop); + assert_eq!(hook.hook_source, HookSource::Project); + assert_eq!(hook.status, HookRunStatus::Blocked); + } + + #[tokio::test] + async fn hook_run_analytics_payload_falls_back_to_turn_context_id() { + let (_session, turn_context) = make_session_and_context().await; + let completed = HookCompletedEvent { + turn_id: None, + run: sample_hook_run(HookRunStatus::Failed, HookSource::Unknown), + }; + + let (tracking, hook) = + hook_run_analytics_payload("thread-123".to_string(), &turn_context, &completed); + + assert_eq!(tracking.turn_id, turn_context.sub_id); + assert_eq!(hook.hook_source, HookSource::Unknown); + assert_eq!(hook.status, HookRunStatus::Failed); + } + + fn sample_hook_run(status: HookRunStatus, source: HookSource) -> HookRunSummary { + HookRunSummary { + id: "stop:0:/tmp/hooks.json".to_string(), + event_name: HookEventName::Stop, + handler_type: HookHandlerType::Command, + execution_mode: HookExecutionMode::Sync, + scope: HookScope::Turn, + source_path: test_path_buf("/tmp/hooks.json").abs(), + source, + display_order: 0, + status, + status_message: None, + started_at: 10, + completed_at: Some(37), + duration_ms: Some(27), + entries: Vec::new(), + } + } } diff --git a/codex-rs/hooks/src/engine/discovery.rs b/codex-rs/hooks/src/engine/discovery.rs index 5505d8e914..682328e640 100644 --- a/codex-rs/hooks/src/engine/discovery.rs +++ b/codex-rs/hooks/src/engine/discovery.rs @@ -9,6 +9,8 @@ use super::config::HooksFile; use super::config::MatcherGroup; use crate::events::common::matcher_pattern_for_event; use crate::events::common::validate_matcher_pattern; +use codex_config::ConfigLayerSource; +use codex_protocol::protocol::HookSource; pub(crate) struct DiscoveryResult { pub handlers: Vec, @@ -93,6 +95,7 @@ pub(crate) fn discover_handlers(config_layer_stack: Option<&ConfigLayerStack>) - &mut warnings, &mut display_order, &source_path, + hook_source_for_config_layer_source(&layer.name), event_name, groups, ); @@ -102,95 +105,94 @@ pub(crate) fn discover_handlers(config_layer_stack: Option<&ConfigLayerStack>) - DiscoveryResult { handlers, warnings } } -fn append_group_handlers( - handlers: &mut Vec, - warnings: &mut Vec, - display_order: &mut i64, - source_path: &AbsolutePathBuf, - event_name: codex_protocol::protocol::HookEventName, - matcher: Option<&str>, - group_handlers: Vec, -) { - if let Some(matcher) = matcher - && let Err(err) = validate_matcher_pattern(matcher) - { - warnings.push(format!( - "invalid matcher {matcher:?} in {}: {err}", - source_path.display() - )); - return; - } - - for handler in group_handlers { - match handler { - HookHandlerConfig::Command { - command, - timeout_sec, - r#async, - status_message, - } => { - if r#async { - warnings.push(format!( - "skipping async hook in {}: async hooks are not supported yet", - source_path.display() - )); - continue; - } - if command.trim().is_empty() { - warnings.push(format!( - "skipping empty hook command in {}", - source_path.display() - )); - continue; - } - let timeout_sec = timeout_sec.unwrap_or(600).max(1); - handlers.push(ConfiguredHandler { - event_name, - matcher: matcher.map(ToOwned::to_owned), - command, - timeout_sec, - status_message, - source_path: source_path.clone(), - display_order: *display_order, - }); - *display_order += 1; - } - HookHandlerConfig::Prompt {} => warnings.push(format!( - "skipping prompt hook in {}: prompt hooks are not supported yet", - source_path.display() - )), - HookHandlerConfig::Agent {} => warnings.push(format!( - "skipping agent hook in {}: agent hooks are not supported yet", - source_path.display() - )), - } - } -} - fn append_matcher_groups( handlers: &mut Vec, warnings: &mut Vec, display_order: &mut i64, source_path: &AbsolutePathBuf, + source: HookSource, event_name: codex_protocol::protocol::HookEventName, groups: Vec, ) { for group in groups { - append_group_handlers( - handlers, - warnings, - display_order, - source_path, - event_name, - matcher_pattern_for_event(event_name, group.matcher.as_deref()), - group.hooks, - ); + let matcher = matcher_pattern_for_event(event_name, group.matcher.as_deref()); + if let Some(matcher) = matcher + && let Err(err) = validate_matcher_pattern(matcher) + { + warnings.push(format!( + "invalid matcher {matcher:?} in {}: {err}", + source_path.display() + )); + continue; + } + + for handler in group.hooks { + match handler { + HookHandlerConfig::Command { + command, + timeout_sec, + r#async, + status_message, + } => { + if r#async { + warnings.push(format!( + "skipping async hook in {}: async hooks are not supported yet", + source_path.display() + )); + continue; + } + if command.trim().is_empty() { + warnings.push(format!( + "skipping empty hook command in {}", + source_path.display() + )); + continue; + } + let timeout_sec = timeout_sec.unwrap_or(600).max(1); + handlers.push(ConfiguredHandler { + event_name, + matcher: matcher.map(ToOwned::to_owned), + command, + timeout_sec, + status_message, + source_path: source_path.clone(), + source, + display_order: *display_order, + }); + *display_order += 1; + } + HookHandlerConfig::Prompt {} => warnings.push(format!( + "skipping prompt hook in {}: prompt hooks are not supported yet", + source_path.display() + )), + HookHandlerConfig::Agent {} => warnings.push(format!( + "skipping agent hook in {}: agent hooks are not supported yet", + source_path.display() + )), + } + } + } +} + +fn hook_source_for_config_layer_source(source: &ConfigLayerSource) -> HookSource { + match source { + ConfigLayerSource::System { .. } => HookSource::System, + ConfigLayerSource::User { .. } => HookSource::User, + ConfigLayerSource::Project { .. } => HookSource::Project, + ConfigLayerSource::Mdm { .. } => HookSource::Mdm, + ConfigLayerSource::SessionFlags => HookSource::SessionFlags, + ConfigLayerSource::LegacyManagedConfigTomlFromFile { .. } => { + HookSource::LegacyManagedConfigFile + } + ConfigLayerSource::LegacyManagedConfigTomlFromMdm => HookSource::LegacyManagedConfigMdm, } } #[cfg(test)] mod tests { + use codex_config::ConfigLayerSource; use codex_protocol::protocol::HookEventName; + use codex_protocol::protocol::HookSource; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_absolute_path::test_support::PathBufExt; use codex_utils_absolute_path::test_support::test_path_buf; @@ -198,32 +200,43 @@ mod tests { use super::ConfiguredHandler; use super::HookHandlerConfig; - use super::append_group_handlers; - use crate::events::common::matcher_pattern_for_event; + use super::MatcherGroup; + use super::append_matcher_groups; fn source_path() -> AbsolutePathBuf { test_path_buf("/tmp/hooks.json").abs() } + fn hook_source() -> HookSource { + HookSource::User + } + + fn command_group(matcher: Option<&str>) -> MatcherGroup { + MatcherGroup { + matcher: matcher.map(str::to_string), + hooks: vec![HookHandlerConfig::Command { + command: "echo hello".to_string(), + timeout_sec: None, + r#async: false, + status_message: None, + }], + } + } + #[test] fn user_prompt_submit_ignores_invalid_matcher_during_discovery() { let mut handlers = Vec::new(); let mut warnings = Vec::new(); let mut display_order = 0; - append_group_handlers( + append_matcher_groups( &mut handlers, &mut warnings, &mut display_order, &source_path(), + hook_source(), HookEventName::UserPromptSubmit, - matcher_pattern_for_event(HookEventName::UserPromptSubmit, Some("[")), - vec![HookHandlerConfig::Command { - command: "echo hello".to_string(), - timeout_sec: None, - r#async: false, - status_message: None, - }], + vec![command_group(Some("["))], ); assert_eq!(warnings, Vec::::new()); @@ -236,6 +249,7 @@ mod tests { timeout_sec: 600, status_message: None, source_path: source_path(), + source: hook_source(), display_order: 0, }] ); @@ -247,19 +261,14 @@ mod tests { let mut warnings = Vec::new(); let mut display_order = 0; - append_group_handlers( + append_matcher_groups( &mut handlers, &mut warnings, &mut display_order, &source_path(), + hook_source(), HookEventName::PreToolUse, - matcher_pattern_for_event(HookEventName::PreToolUse, Some("^Bash$")), - vec![HookHandlerConfig::Command { - command: "echo hello".to_string(), - timeout_sec: None, - r#async: false, - status_message: None, - }], + vec![command_group(Some("^Bash$"))], ); assert_eq!(warnings, Vec::::new()); @@ -272,6 +281,7 @@ mod tests { timeout_sec: 600, status_message: None, source_path: source_path(), + source: hook_source(), display_order: 0, }] ); @@ -283,19 +293,14 @@ mod tests { let mut warnings = Vec::new(); let mut display_order = 0; - append_group_handlers( + append_matcher_groups( &mut handlers, &mut warnings, &mut display_order, &source_path(), + hook_source(), HookEventName::PreToolUse, - matcher_pattern_for_event(HookEventName::PreToolUse, Some("*")), - vec![HookHandlerConfig::Command { - command: "echo hello".to_string(), - timeout_sec: None, - r#async: false, - status_message: None, - }], + vec![command_group(Some("*"))], ); assert_eq!(warnings, Vec::::new()); @@ -309,19 +314,14 @@ mod tests { let mut warnings = Vec::new(); let mut display_order = 0; - append_group_handlers( + append_matcher_groups( &mut handlers, &mut warnings, &mut display_order, &source_path(), + hook_source(), HookEventName::PostToolUse, - matcher_pattern_for_event(HookEventName::PostToolUse, Some("Edit|Write")), - vec![HookHandlerConfig::Command { - command: "echo hello".to_string(), - timeout_sec: None, - r#async: false, - status_message: None, - }], + vec![command_group(Some("Edit|Write"))], ); assert_eq!(warnings, Vec::::new()); @@ -329,4 +329,52 @@ mod tests { assert_eq!(handlers[0].event_name, HookEventName::PostToolUse); assert_eq!(handlers[0].matcher.as_deref(), Some("Edit|Write")); } + + #[test] + fn hook_source_for_config_layer_source_discards_source_details() { + let config_file = test_path_buf("/tmp/.codex/config.toml").abs(); + let dot_codex_folder = test_path_buf("/tmp/worktree/.codex").abs(); + + assert_eq!( + super::hook_source_for_config_layer_source(&ConfigLayerSource::System { + file: config_file.clone(), + }), + HookSource::System, + ); + assert_eq!( + super::hook_source_for_config_layer_source(&ConfigLayerSource::User { + file: config_file.clone(), + }), + HookSource::User, + ); + assert_eq!( + super::hook_source_for_config_layer_source(&ConfigLayerSource::Project { + dot_codex_folder + }), + HookSource::Project, + ); + assert_eq!( + super::hook_source_for_config_layer_source(&ConfigLayerSource::Mdm { + domain: "com.openai.codex".to_string(), + key: "config".to_string(), + }), + HookSource::Mdm, + ); + assert_eq!( + super::hook_source_for_config_layer_source(&ConfigLayerSource::SessionFlags), + HookSource::SessionFlags, + ); + assert_eq!( + super::hook_source_for_config_layer_source( + &ConfigLayerSource::LegacyManagedConfigTomlFromFile { file: config_file }, + ), + HookSource::LegacyManagedConfigFile, + ); + assert_eq!( + super::hook_source_for_config_layer_source( + &ConfigLayerSource::LegacyManagedConfigTomlFromMdm, + ), + HookSource::LegacyManagedConfigMdm, + ); + } } diff --git a/codex-rs/hooks/src/engine/dispatcher.rs b/codex-rs/hooks/src/engine/dispatcher.rs index e3dedf54bb..fbb2a1843f 100644 --- a/codex-rs/hooks/src/engine/dispatcher.rs +++ b/codex-rs/hooks/src/engine/dispatcher.rs @@ -50,6 +50,7 @@ pub(crate) fn running_summary(handler: &ConfiguredHandler) -> HookRunSummary { execution_mode: HookExecutionMode::Sync, scope: scope_for_event(handler.event_name), source_path: handler.source_path.clone(), + source: handler.source, display_order: handler.display_order, status: HookRunStatus::Running, status_message: handler.status_message.clone(), @@ -95,6 +96,7 @@ pub(crate) fn completed_summary( execution_mode: HookExecutionMode::Sync, scope: scope_for_event(handler.event_name), source_path: handler.source_path.clone(), + source: handler.source, display_order: handler.display_order, status, status_message: handler.status_message.clone(), @@ -118,6 +120,7 @@ fn scope_for_event(event_name: HookEventName) -> HookScope { #[cfg(test)] mod tests { use codex_protocol::protocol::HookEventName; + use codex_protocol::protocol::HookSource; use codex_utils_absolute_path::test_support::PathBufExt; use codex_utils_absolute_path::test_support::test_path_buf; @@ -137,6 +140,7 @@ mod tests { timeout_sec: 5, status_message: None, source_path: test_path_buf("/tmp/hooks.json").abs(), + source: HookSource::User, display_order, } } diff --git a/codex-rs/hooks/src/engine/mod.rs b/codex-rs/hooks/src/engine/mod.rs index f3992e8ef3..9bc23705d4 100644 --- a/codex-rs/hooks/src/engine/mod.rs +++ b/codex-rs/hooks/src/engine/mod.rs @@ -7,6 +7,7 @@ pub(crate) mod schema_loader; use codex_config::ConfigLayerStack; use codex_protocol::protocol::HookRunSummary; +use codex_protocol::protocol::HookSource; use codex_utils_absolute_path::AbsolutePathBuf; use crate::events::post_tool_use::PostToolUseOutcome; @@ -34,6 +35,7 @@ pub(crate) struct ConfiguredHandler { pub timeout_sec: u64, pub status_message: Option, pub source_path: AbsolutePathBuf, + pub source: HookSource, pub display_order: i64, } diff --git a/codex-rs/hooks/src/events/post_tool_use.rs b/codex-rs/hooks/src/events/post_tool_use.rs index 056956705d..f7f40d2794 100644 --- a/codex-rs/hooks/src/events/post_tool_use.rs +++ b/codex-rs/hooks/src/events/post_tool_use.rs @@ -525,6 +525,7 @@ mod tests { timeout_sec: 5, status_message: Some("running post tool use hook".to_string()), source_path: test_path_buf("/tmp/hooks.json").abs(), + source: codex_protocol::protocol::HookSource::User, display_order: 0, } } diff --git a/codex-rs/hooks/src/events/pre_tool_use.rs b/codex-rs/hooks/src/events/pre_tool_use.rs index e188207c8a..f90abc2816 100644 --- a/codex-rs/hooks/src/events/pre_tool_use.rs +++ b/codex-rs/hooks/src/events/pre_tool_use.rs @@ -514,6 +514,7 @@ mod tests { timeout_sec: 5, status_message: None, source_path: test_path_buf("/tmp/hooks.json").abs(), + source: codex_protocol::protocol::HookSource::User, display_order: 0, } } diff --git a/codex-rs/hooks/src/events/session_start.rs b/codex-rs/hooks/src/events/session_start.rs index ea396a2231..3c51426fa4 100644 --- a/codex-rs/hooks/src/events/session_start.rs +++ b/codex-rs/hooks/src/events/session_start.rs @@ -361,6 +361,7 @@ mod tests { timeout_sec: 600, status_message: None, source_path: test_path_buf("/tmp/hooks.json").abs(), + source: codex_protocol::protocol::HookSource::User, display_order: 0, } } diff --git a/codex-rs/hooks/src/events/stop.rs b/codex-rs/hooks/src/events/stop.rs index 478cdbdd8b..5fe55568a4 100644 --- a/codex-rs/hooks/src/events/stop.rs +++ b/codex-rs/hooks/src/events/stop.rs @@ -528,6 +528,7 @@ mod tests { timeout_sec: 600, status_message: None, source_path: test_path_buf("/tmp/hooks.json").abs(), + source: codex_protocol::protocol::HookSource::User, display_order: 0, } } diff --git a/codex-rs/hooks/src/events/user_prompt_submit.rs b/codex-rs/hooks/src/events/user_prompt_submit.rs index 875d56bd38..9aaff9a7d4 100644 --- a/codex-rs/hooks/src/events/user_prompt_submit.rs +++ b/codex-rs/hooks/src/events/user_prompt_submit.rs @@ -419,6 +419,7 @@ mod tests { timeout_sec: 5, status_message: None, source_path: test_path_buf("/tmp/hooks.json").abs(), + source: codex_protocol::protocol::HookSource::User, display_order: 0, } } diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 8aea3544d4..d9c63370db 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -1635,6 +1635,20 @@ pub enum HookScope { Turn, } +#[derive(Debug, Clone, Copy, Default, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "snake_case")] +pub enum HookSource { + System, + User, + Project, + Mdm, + SessionFlags, + LegacyManagedConfigFile, + LegacyManagedConfigMdm, + #[default] + Unknown, +} + #[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] #[serde(rename_all = "snake_case")] pub enum HookRunStatus { @@ -1671,6 +1685,8 @@ pub struct HookRunSummary { pub execution_mode: HookExecutionMode, pub scope: HookScope, pub source_path: AbsolutePathBuf, + #[serde(default)] + pub source: HookSource, pub display_order: i64, pub status: HookRunStatus, pub status_message: Option, diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index 1db3dbc7ed..ab20b2214c 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -9892,6 +9892,7 @@ guardian_approval = true execution_mode: AppServerHookExecutionMode::Sync, scope: AppServerHookScope::Turn, source_path: test_path_buf("/tmp/hooks.json").abs(), + source: codex_app_server_protocol::HookSource::User, display_order: 0, status: AppServerHookRunStatus::Running, status_message: Some("checking go-workflow input policy".to_string()), @@ -9914,6 +9915,7 @@ guardian_approval = true execution_mode: AppServerHookExecutionMode::Sync, scope: AppServerHookScope::Turn, source_path: test_path_buf("/tmp/hooks.json").abs(), + source: codex_app_server_protocol::HookSource::User, display_order: 0, status: AppServerHookRunStatus::Stopped, status_message: Some("checking go-workflow input policy".to_string()), diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index f2697d9072..7cfcd56e17 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -1329,6 +1329,7 @@ fn hook_run_summary_from_notification( execution_mode: run.execution_mode.to_core(), scope: run.scope.to_core(), source_path: run.source_path, + source: run.source.to_core(), display_order: run.display_order, status: run.status.to_core(), status_message: run.status_message, diff --git a/codex-rs/tui/src/chatwidget/tests/helpers.rs b/codex-rs/tui/src/chatwidget/tests/helpers.rs index c0f881e2b8..9e365afa78 100644 --- a/codex-rs/tui/src/chatwidget/tests/helpers.rs +++ b/codex-rs/tui/src/chatwidget/tests/helpers.rs @@ -1001,6 +1001,7 @@ pub(super) async fn assert_hook_events_snapshot( execution_mode: codex_protocol::protocol::HookExecutionMode::Sync, scope: codex_protocol::protocol::HookScope::Turn, source_path: PathBuf::from(test_path_display("/tmp/hooks.json")).abs(), + source: codex_protocol::protocol::HookSource::User, display_order: 0, status: codex_protocol::protocol::HookRunStatus::Running, status_message: Some(status_message.to_string()), @@ -1035,6 +1036,7 @@ pub(super) async fn assert_hook_events_snapshot( execution_mode: codex_protocol::protocol::HookExecutionMode::Sync, scope: codex_protocol::protocol::HookScope::Turn, source_path: PathBuf::from(test_path_display("/tmp/hooks.json")).abs(), + source: codex_protocol::protocol::HookSource::User, display_order: 0, status: codex_protocol::protocol::HookRunStatus::Completed, status_message: Some(status_message.to_string()), diff --git a/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs b/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs index 75ca9d1fba..a31b93bb4c 100644 --- a/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs +++ b/codex-rs/tui/src/chatwidget/tests/status_and_layout.rs @@ -1425,6 +1425,7 @@ async fn user_prompt_submit_app_server_hook_notifications_render_snapshot() { execution_mode: AppServerHookExecutionMode::Sync, scope: AppServerHookScope::Turn, source_path: PathBuf::from(test_path_display("/tmp/hooks.json")).abs(), + source: codex_app_server_protocol::HookSource::User, display_order: 0, status: AppServerHookRunStatus::Running, status_message: Some("checking go-workflow input policy".to_string()), @@ -1447,6 +1448,7 @@ async fn user_prompt_submit_app_server_hook_notifications_render_snapshot() { execution_mode: AppServerHookExecutionMode::Sync, scope: AppServerHookScope::Turn, source_path: PathBuf::from(test_path_display("/tmp/hooks.json")).abs(), + source: codex_app_server_protocol::HookSource::User, display_order: 0, status: AppServerHookRunStatus::Stopped, status_message: Some("checking go-workflow input policy".to_string()), @@ -1517,6 +1519,7 @@ async fn completed_hook_with_no_entries_stays_out_of_history() { execution_mode: codex_protocol::protocol::HookExecutionMode::Sync, scope: codex_protocol::protocol::HookScope::Turn, source_path: PathBuf::from(test_path_display("/tmp/hooks.json")).abs(), + source: codex_protocol::protocol::HookSource::User, display_order: 0, status: codex_protocol::protocol::HookRunStatus::Running, status_message: None, @@ -1542,6 +1545,7 @@ async fn completed_hook_with_no_entries_stays_out_of_history() { execution_mode: codex_protocol::protocol::HookExecutionMode::Sync, scope: codex_protocol::protocol::HookScope::Turn, source_path: PathBuf::from(test_path_display("/tmp/hooks.json")).abs(), + source: codex_protocol::protocol::HookSource::User, display_order: 0, status: codex_protocol::protocol::HookRunStatus::Completed, status_message: None, @@ -2017,6 +2021,7 @@ fn hook_run_summary( execution_mode: codex_protocol::protocol::HookExecutionMode::Sync, scope: codex_protocol::protocol::HookScope::Turn, source_path: PathBuf::from(test_path_display("/tmp/hooks.json")).abs(), + source: codex_protocol::protocol::HookSource::User, display_order: 0, status, status_message: status_message.map(str::to_string), diff --git a/codex-rs/tui/src/history_cell/hook_cell.rs b/codex-rs/tui/src/history_cell/hook_cell.rs index a1207186c1..50642aa0f2 100644 --- a/codex-rs/tui/src/history_cell/hook_cell.rs +++ b/codex-rs/tui/src/history_cell/hook_cell.rs @@ -768,6 +768,7 @@ mod tests { execution_mode: codex_protocol::protocol::HookExecutionMode::Sync, scope: codex_protocol::protocol::HookScope::Turn, source_path: test_path_buf("/tmp/hooks.json").abs(), + source: codex_protocol::protocol::HookSource::User, display_order: 0, status: HookRunStatus::Running, status_message: Some("checking output policy".to_string()), From ea34c6ed8de87c72616ac08c9c5e399ebd1f9c58 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Thu, 16 Apr 2026 12:48:31 -0700 Subject: [PATCH 06/16] fix: fix clippy issue in examples/ folder (#18184) I believe this use of `expect()` was introduced in https://github.com/openai/codex/pull/17826, but was not flagged by CI. Though I did see it in the diagnostics panel in VS Code, so it's worth cleaning up. I guess our current CI does include `examples/` when running Clippy? --- codex-rs/thread-store/examples/generate-proto.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/codex-rs/thread-store/examples/generate-proto.rs b/codex-rs/thread-store/examples/generate-proto.rs index aa159f5bcb..0b4afb54f9 100644 --- a/codex-rs/thread-store/examples/generate-proto.rs +++ b/codex-rs/thread-store/examples/generate-proto.rs @@ -1,7 +1,12 @@ use std::path::PathBuf; fn main() -> Result<(), Box> { - let proto_dir = PathBuf::from(std::env::args().nth(1).expect("proto dir")); + let Some(proto_dir_arg) = std::env::args().nth(1) else { + eprintln!("Usage: generate-proto "); + std::process::exit(1); + }; + + let proto_dir = PathBuf::from(proto_dir_arg); let proto_file = proto_dir.join("codex.thread_store.v1.proto"); tonic_prost_build::configure() From ec8d4bfc77353ac328a7b4897ab47b296e90f19e Mon Sep 17 00:00:00 2001 From: Felipe Coury Date: Thu, 16 Apr 2026 17:29:34 -0300 Subject: [PATCH 07/16] fix(app-server): replay token usage after resume and fork (#18023) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem When a user resumed or forked a session, the TUI could render the restored thread history immediately, but it did not receive token usage until a later model turn emitted a fresh usage event. That left the context/status UI blank or stale during the exact window where the user expects resumed state to look complete. Core already reconstructed token usage from the rollout; the missing behavior was app-server lifecycle replay to the client that just attached. ## Mental model Token usage has two representations. The rollout is the durable source of historical `TokenCount` events, and the core session cache is the in-memory snapshot reconstructed from that rollout on resume or fork. App-server v2 clients do not read core state directly; they learn about usage through `thread/tokenUsage/updated`. The fix keeps those roles separate: core exposes the restored `TokenUsageInfo`, and app-server sends one targeted notification after a successful `thread/resume` or `thread/fork` response when that restored snapshot exists. This notification is not a new model event. It is a replay of already-persisted state for the client that just attached. That distinction matters because using the normal core event path here would risk duplicating `TokenCount` entries in the rollout and making future resumes count historical usage twice. ## Non-goals This change does not add a new protocol method or payload shape. It reuses the existing v2 `thread/tokenUsage/updated` notification and the TUI’s existing handler for that notification. This change does not alter how token usage is computed, accumulated, compacted, or written during turns. It only exposes the token usage that resume and fork reconstruction already restored. This change does not broadcast historical usage replay to every subscribed client. The replay is intentionally scoped to the connection that requested resume or fork so already-attached clients are not surprised by an old usage update while they may be rendering live activity. ## Tradeoffs Sending the usage notification after the JSON-RPC response preserves a clear lifecycle order: the client first receives the thread object, then receives restored usage for that thread. The tradeoff is that usage is still a notification rather than part of the `thread/resume` or `thread/fork` response. That keeps the protocol shape stable and avoids duplicating usage fields across response types, but clients must continue listening for notifications after receiving the response. The helper selects the latest non-in-progress turn id for the replayed usage notification. This is conservative because restored usage belongs to completed persisted accounting, not to newly attached in-flight work. The fallback to the last turn preserves a stable wire payload for unusual histories, but histories with no meaningful completed turn still have a weak attribution story. ## Architecture Core already seeds `Session` token state from the last persisted rollout `TokenCount` during `InitialHistory::Resumed` and `InitialHistory::Forked`. The new core accessor exposes the complete `TokenUsageInfo` through `CodexThread` without giving app-server direct session mutation authority. App-server calls that accessor from three lifecycle paths: cold `thread/resume`, running-thread resume/rejoin, and `thread/fork`. In each path, the server sends the normal response first, then calls a shared helper that converts core usage into `ThreadTokenUsageUpdatedNotification` and sends it only to the requesting connection. The tests build fake rollouts with a user turn plus a persisted token usage event. They then exercise `thread/resume` and `thread/fork` without starting another model turn, proving that restored usage arrives before any next-turn token event could be produced. ## Observability The primary debug path is the app-server JSON-RPC stream. After `thread/resume` or `thread/fork`, a client should see the response followed by `thread/tokenUsage/updated` when the source rollout includes token usage. If the notification is absent, check whether the rollout contains an `event_msg` payload of type `token_count`, whether core reconstruction seeded `Session::token_usage_info`, and whether the connection stayed attached long enough to receive the targeted notification. The notification is sent through the existing `OutgoingMessageSender::send_server_notification_to_connections` path, so existing app-server tracing around server notifications still applies. Because this is a replay, not a model turn event, debugging should start at the resume/fork handlers rather than the turn event translation in `bespoke_event_handling`. ## Tests The focused regression coverage is `cargo test -p codex-app-server emits_restored_token_usage`, which covers both resume and fork. The core reconstruction guard is `cargo test -p codex-core record_initial_history_seeds_token_info_from_rollout`. Formatting and lint/fix passes were run with `just fmt`, `just fix -p codex-core`, and `just fix -p codex-app-server`. Full crate test runs surfaced pre-existing unrelated failures in command execution and plugin marketplace tests; the new token usage tests passed in focused runs and within the app-server suite before the unrelated command execution failure. --- .../src/protocol/thread_history.rs | 14 + codex-rs/app-server/README.md | 4 +- .../app-server/src/codex_message_processor.rs | 68 ++++- .../token_usage_replay.rs | 133 +++++++++ codex-rs/app-server/tests/common/lib.rs | 1 + codex-rs/app-server/tests/common/rollout.rs | 59 ++++ .../app-server/tests/suite/v2/thread_fork.rs | 55 ++++ .../tests/suite/v2/thread_resume.rs | 269 ++++++++++++++++++ codex-rs/core/src/codex.rs | 11 + codex-rs/core/src/codex_thread.rs | 12 + 10 files changed, 622 insertions(+), 4 deletions(-) create mode 100644 codex-rs/app-server/src/codex_message_processor/token_usage_replay.rs diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index 54f882cf29..5c327187e9 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -123,6 +123,20 @@ impl ThreadHistoryBuilder { .or_else(|| self.turns.last().cloned()) } + /// Returns the index of the active turn snapshot within the finished turn list. + /// + /// When a turn is still open, this is the index it will occupy after + /// `finish`. When no turn is open, it is the index of the last finished turn. + pub fn active_turn_position(&self) -> Option { + if self.current_turn.is_some() { + Some(self.turns.len()) + } else if self.turns.is_empty() { + None + } else { + Some(self.turns.len() - 1) + } + } + pub fn has_active_turn(&self) -> bool { self.current_turn.is_some() } diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 5cfe052b50..2fece44544 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -250,7 +250,7 @@ Start a fresh thread when you need a new Codex conversation. Valid `personality` values are `"friendly"`, `"pragmatic"`, and `"none"`. When `"none"` is selected, the personality placeholder is replaced with an empty string. -To continue a stored session, call `thread/resume` with the `thread.id` you previously recorded. The response shape matches `thread/start`, and no additional notifications are emitted. You can also pass the same configuration overrides supported by `thread/start`, including `approvalsReviewer`. +To continue a stored session, call `thread/resume` with the `thread.id` you previously recorded. The response shape matches `thread/start`. When the stored session includes persisted token usage, the server emits `thread/tokenUsage/updated` immediately after the response so clients can render restored usage before the next turn starts. You can also pass the same configuration overrides supported by `thread/start`, including `approvalsReviewer`. By default, resume uses the latest persisted `model` and `reasoningEffort` values associated with the thread. Supplying any of `model`, `modelProvider`, `config.model`, or `config.model_reasoning_effort` disables that persisted fallback and uses the explicit overrides plus normal config resolution instead. @@ -264,7 +264,7 @@ Example: { "id": 11, "result": { "thread": { "id": "thr_123", … } } } ``` -To branch from a stored session, call `thread/fork` with the `thread.id`. This creates a new thread id and emits a `thread/started` notification for it. If the source thread is actively running, the fork snapshots it as if the current turn had been interrupted first. Pass `ephemeral: true` when the fork should stay in-memory only: +To branch from a stored session, call `thread/fork` with the `thread.id`. This creates a new thread id and emits a `thread/started` notification for it. When the source history includes persisted token usage, the server also emits `thread/tokenUsage/updated` for the new thread immediately after the response. If the source thread is actively running, the fork snapshots it as if the current turn had been interrupted first. Pass `ephemeral: true` when the fork should stay in-memory only: ```json { "method": "thread/fork", "id": 12, "params": { "threadId": "thr_123", "ephemeral": true } } diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 410f22ece8..0f3b8bf522 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -363,12 +363,17 @@ use codex_app_server_protocol::ServerRequest; mod apps_list_helpers; mod plugin_app_helpers; mod plugin_mcp_oauth; +mod token_usage_replay; use crate::filters::compute_source_filters; use crate::filters::source_kind_matches; use crate::thread_state::ThreadListenerCommand; use crate::thread_state::ThreadState; use crate::thread_state::ThreadStateManager; +use token_usage_replay::latest_token_usage_turn_id_for_thread_path; +use token_usage_replay::latest_token_usage_turn_id_from_rollout_items; +use token_usage_replay::latest_token_usage_turn_id_from_rollout_path; +use token_usage_replay::send_thread_token_usage_update_to_connection; const THREAD_LIST_DEFAULT_LIMIT: usize = 25; const THREAD_LIST_MAX_LIMIT: usize = 100; @@ -4103,7 +4108,7 @@ impl CodexMessageProcessor { { Ok(NewThread { thread_id, - thread, + thread: codex_thread, session_configured, }) => { let SessionConfiguredEvent { rollout_path, .. } = session_configured; @@ -4132,7 +4137,7 @@ impl CodexMessageProcessor { let mut thread = match self .load_thread_from_resume_source_or_send_internal( thread_id, - thread.as_ref(), + codex_thread.as_ref(), &response_history, rollout_path.as_path(), fallback_model_provider.as_str(), @@ -4184,7 +4189,25 @@ impl CodexMessageProcessor { ); } + let connection_id = request_id.connection_id; + let token_usage_thread = response.thread.clone(); + let token_usage_turn_id = latest_token_usage_turn_id_from_rollout_items( + &response_history.get_rollout_items(), + &token_usage_thread, + ); self.outgoing.send_response(request_id, response).await; + // The client needs restored usage before it starts another turn. + // Sending after the response preserves JSON-RPC request ordering while + // still filling the status line before the next turn lifecycle begins. + send_thread_token_usage_update_to_connection( + &self.outgoing, + connection_id, + thread_id, + &token_usage_thread, + codex_thread.as_ref(), + token_usage_turn_id, + ) + .await; } Err(err) => { let error = JSONRPCErrorError { @@ -4820,7 +4843,31 @@ impl CodexMessageProcessor { ); } + let connection_id = request_id.connection_id; + let token_usage_thread = response.thread.clone(); + let token_usage_turn_id = if let Some(turn_id) = + latest_token_usage_turn_id_for_thread_path(&token_usage_thread).await + { + Some(turn_id) + } else { + latest_token_usage_turn_id_from_rollout_path( + rollout_path.as_path(), + &token_usage_thread, + ) + .await + }; self.outgoing.send_response(request_id, response).await; + // Mirror the resume contract for forks: the new thread is usable as soon + // as the response arrives, so restored usage must follow immediately. + send_thread_token_usage_update_to_connection( + &self.outgoing, + connection_id, + thread_id, + &token_usage_thread, + forked_thread.as_ref(), + token_usage_turn_id, + ) + .await; let notif = ThreadStartedNotification { thread }; self.outgoing @@ -8525,7 +8572,24 @@ async fn handle_pending_thread_resume_request( sandbox: sandbox_policy.into(), reasoning_effort, }; + let token_usage_thread = response.thread.clone(); + let token_usage_turn_id = latest_token_usage_turn_id_from_rollout_path( + pending.rollout_path.as_path(), + &token_usage_thread, + ) + .await; outgoing.send_response(request_id, response).await; + // Rejoining a loaded thread has the same UI contract as a cold resume, but + // uses the live conversation state instead of reconstructing a new session. + send_thread_token_usage_update_to_connection( + outgoing, + connection_id, + conversation_id, + &token_usage_thread, + conversation.as_ref(), + token_usage_turn_id, + ) + .await; outgoing .replay_requests_to_connection_for_thread(connection_id, conversation_id) .await; diff --git a/codex-rs/app-server/src/codex_message_processor/token_usage_replay.rs b/codex-rs/app-server/src/codex_message_processor/token_usage_replay.rs new file mode 100644 index 0000000000..d3c06f5e7c --- /dev/null +++ b/codex-rs/app-server/src/codex_message_processor/token_usage_replay.rs @@ -0,0 +1,133 @@ +//! Replays persisted token usage snapshots when a client attaches to an existing thread. +//! +//! The message processor decides when replay is allowed and preserves JSON-RPC response +//! ordering. This module owns notification construction and the attribution rules that +//! map the latest persisted `TokenCount` back to a v2 turn id. +//! +//! Rollout histories can contain explicit turn ids or generated turn ids. When explicit +//! ids do not match the rebuilt thread, replay falls back to the active turn position at +//! the time the `TokenCount` was persisted so the notification still targets the +//! corresponding rebuilt turn. + +use std::path::Path; +use std::sync::Arc; + +use codex_app_server_protocol::ServerNotification; +use codex_app_server_protocol::Thread; +use codex_app_server_protocol::ThreadHistoryBuilder; +use codex_app_server_protocol::ThreadTokenUsage; +use codex_app_server_protocol::ThreadTokenUsageUpdatedNotification; +use codex_app_server_protocol::TurnStatus; +use codex_core::CodexThread; +use codex_protocol::ThreadId; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::RolloutItem; + +use crate::codex_message_processor::read_rollout_items_from_rollout; +use crate::outgoing_message::ConnectionId; +use crate::outgoing_message::OutgoingMessageSender; + +/// Sends a restored token usage update to the connection that attached to a thread. +/// +/// This is lifecycle replay rather than a model event: the rollout already contains +/// the original `TokenCount`, and emitting through `send_event` here would duplicate +/// persisted usage records. Keeping this helper connection-scoped also avoids +/// surprising other subscribers with a historical usage update while they may be +/// rendering live turn events. +pub(super) async fn send_thread_token_usage_update_to_connection( + outgoing: &Arc, + connection_id: ConnectionId, + thread_id: ThreadId, + thread: &Thread, + conversation: &CodexThread, + token_usage_turn_id: Option, +) { + let Some(info) = conversation.token_usage_info().await else { + return; + }; + let notification = ThreadTokenUsageUpdatedNotification { + thread_id: thread_id.to_string(), + turn_id: token_usage_turn_id.unwrap_or_else(|| latest_token_usage_turn_id(thread)), + token_usage: ThreadTokenUsage::from(info), + }; + outgoing + .send_server_notification_to_connections( + &[connection_id], + ServerNotification::ThreadTokenUsageUpdated(notification), + ) + .await; +} + +pub(super) async fn latest_token_usage_turn_id_for_thread_path(thread: &Thread) -> Option { + let rollout_path = thread.path.as_deref()?; + latest_token_usage_turn_id_from_rollout_path(rollout_path, thread).await +} + +pub(super) async fn latest_token_usage_turn_id_from_rollout_path( + rollout_path: &Path, + thread: &Thread, +) -> Option { + let rollout_items = read_rollout_items_from_rollout(rollout_path).await.ok()?; + latest_token_usage_turn_id_from_rollout_items(&rollout_items, thread) +} + +/// Identifies the turn that was active when a `TokenCount` record appeared. +/// +/// The id is preferred when it still appears in the rebuilt thread. The position is a +/// fallback for histories whose implicit turn ids are regenerated during reconstruction. +struct TokenUsageTurnOwner { + id: String, + position: Option, +} + +pub(super) fn latest_token_usage_turn_id_from_rollout_items( + rollout_items: &[RolloutItem], + thread: &Thread, +) -> Option { + let owner = latest_token_usage_turn_owner_from_rollout_items(rollout_items)?; + if thread.turns.iter().any(|turn| turn.id == owner.id) { + return Some(owner.id); + } + owner + .position + .and_then(|position| thread.turns.get(position)) + .map(|turn| turn.id.clone()) +} + +fn latest_token_usage_turn_owner_from_rollout_items( + rollout_items: &[RolloutItem], +) -> Option { + let mut builder = ThreadHistoryBuilder::new(); + let mut token_usage_turn_owner = None; + + for item in rollout_items { + if matches!(item, RolloutItem::EventMsg(EventMsg::TokenCount(_))) { + token_usage_turn_owner = + builder + .active_turn_snapshot() + .map(|turn| TokenUsageTurnOwner { + id: turn.id, + position: builder.active_turn_position(), + }); + } + builder.handle_rollout_item(item); + } + + token_usage_turn_owner +} + +/// Chooses a fallback turn id that should own a replayed token usage update. +/// +/// Normal replay derives the owner from the rollout position of the latest +/// `TokenCount` event. This fallback only preserves a stable wire shape for +/// unusual histories where that rollout information cannot be read. +fn latest_token_usage_turn_id(thread: &Thread) -> String { + thread + .turns + .iter() + .rev() + .find(|turn| matches!(turn.status, TurnStatus::Completed | TurnStatus::Failed)) + .or_else(|| thread.turns.last()) + .map(|turn| turn.id.clone()) + .unwrap_or_default() +} diff --git a/codex-rs/app-server/tests/common/lib.rs b/codex-rs/app-server/tests/common/lib.rs index 6bb7372162..6ac26d8a56 100644 --- a/codex-rs/app-server/tests/common/lib.rs +++ b/codex-rs/app-server/tests/common/lib.rs @@ -40,6 +40,7 @@ pub use responses::create_shell_command_sse_response; pub use rollout::create_fake_rollout; pub use rollout::create_fake_rollout_with_source; pub use rollout::create_fake_rollout_with_text_elements; +pub use rollout::create_fake_rollout_with_token_usage; pub use rollout::rollout_path; use serde::de::DeserializeOwned; diff --git a/codex-rs/app-server/tests/common/rollout.rs b/codex-rs/app-server/tests/common/rollout.rs index b67390154e..06b273754c 100644 --- a/codex-rs/app-server/tests/common/rollout.rs +++ b/codex-rs/app-server/tests/common/rollout.rs @@ -1,9 +1,13 @@ use anyhow::Result; use codex_protocol::ThreadId; +use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::GitInfo; use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::TokenCountEvent; +use codex_protocol::protocol::TokenUsage; +use codex_protocol::protocol::TokenUsageInfo; use serde_json::json; use std::fs; use std::fs::FileTimes; @@ -50,6 +54,61 @@ pub fn create_fake_rollout( ) } +/// Creates a minimal rollout whose history includes a persisted token usage event. +/// +/// Resume and fork tests use this fixture to verify lifecycle replay of restored +/// usage without starting a model turn. The exact token values are intentionally +/// non-zero and asymmetric so assertions catch swapped total/last fields and +/// dropped cached or reasoning counters. +pub fn create_fake_rollout_with_token_usage( + codex_home: &Path, + filename_ts: &str, + meta_rfc3339: &str, + preview: &str, + model_provider: Option<&str>, +) -> Result { + let thread_id = create_fake_rollout( + codex_home, + filename_ts, + meta_rfc3339, + preview, + model_provider, + /*git_info*/ None, + )?; + let payload = serde_json::to_value(EventMsg::TokenCount(TokenCountEvent { + info: Some(TokenUsageInfo { + total_token_usage: TokenUsage { + input_tokens: 120, + cached_input_tokens: 20, + output_tokens: 30, + reasoning_output_tokens: 10, + total_tokens: 150, + }, + last_token_usage: TokenUsage { + input_tokens: 70, + cached_input_tokens: 10, + output_tokens: 20, + reasoning_output_tokens: 5, + total_tokens: 90, + }, + model_context_window: Some(200_000), + }), + rate_limits: None, + }))?; + let file_path = rollout_path(codex_home, filename_ts, &thread_id); + let line = json!({ + "timestamp": meta_rfc3339, + "type": "event_msg", + "payload": payload + }) + .to_string(); + fs::write( + &file_path, + format!("{}{}\n", fs::read_to_string(&file_path)?, line), + )?; + Ok(thread_id) +} + /// Create a minimal rollout file with an explicit session source. pub fn create_fake_rollout_with_source( codex_home: &Path, diff --git a/codex-rs/app-server/tests/suite/v2/thread_fork.rs b/codex-rs/app-server/tests/suite/v2/thread_fork.rs index 649919db8e..f5bc5a73c9 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_fork.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_fork.rs @@ -2,6 +2,7 @@ use anyhow::Result; use app_test_support::ChatGptAuthFixture; use app_test_support::McpProcess; use app_test_support::create_fake_rollout; +use app_test_support::create_fake_rollout_with_token_usage; use app_test_support::create_mock_responses_server_repeating_assistant; use app_test_support::to_response; use app_test_support::write_chatgpt_auth; @@ -9,6 +10,7 @@ use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::ServerNotification; use codex_app_server_protocol::SessionSource; use codex_app_server_protocol::ThreadForkParams; use codex_app_server_protocol::ThreadForkResponse; @@ -186,6 +188,59 @@ async fn thread_fork_creates_new_thread_and_emits_started() -> Result<()> { Ok(()) } +#[tokio::test] +async fn thread_fork_emits_restored_token_usage_before_next_turn() -> Result<()> { + let server = create_mock_responses_server_repeating_assistant("Done").await; + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri())?; + + let conversation_id = create_fake_rollout_with_token_usage( + codex_home.path(), + "2025-01-05T12-00-00", + "2025-01-05T12:00:00Z", + "Saved user message", + Some("mock_provider"), + )?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let fork_id = mcp + .send_thread_fork_request(ThreadForkParams { + thread_id: conversation_id, + ..Default::default() + }) + .await?; + let fork_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(fork_id)), + ) + .await??; + let ThreadForkResponse { thread, .. } = to_response::(fork_resp)?; + + let note = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("thread/tokenUsage/updated"), + ) + .await??; + let parsed: ServerNotification = note.try_into()?; + let ServerNotification::ThreadTokenUsageUpdated(notification) = parsed else { + panic!("expected thread/tokenUsage/updated notification"); + }; + + assert_eq!(notification.thread_id, thread.id); + assert_eq!(notification.turn_id, thread.turns[0].id); + assert_eq!(notification.token_usage.total.total_tokens, 150); + assert_eq!(notification.token_usage.total.input_tokens, 120); + assert_eq!(notification.token_usage.total.cached_input_tokens, 20); + assert_eq!(notification.token_usage.total.output_tokens, 30); + assert_eq!(notification.token_usage.total.reasoning_output_tokens, 10); + assert_eq!(notification.token_usage.last.total_tokens, 90); + assert_eq!(notification.token_usage.model_context_window, Some(200_000)); + + Ok(()) +} + #[tokio::test] async fn thread_fork_tracks_thread_initialized_analytics() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume.rs b/codex-rs/app-server/tests/suite/v2/thread_resume.rs index 46dee9b5d1..edf792ca1d 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -3,6 +3,7 @@ use app_test_support::ChatGptAuthFixture; use app_test_support::McpProcess; use app_test_support::create_apply_patch_sse_response; use app_test_support::create_fake_rollout_with_text_elements; +use app_test_support::create_fake_rollout_with_token_usage; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_repeating_assistant; use app_test_support::create_mock_responses_server_sequence_unchecked; @@ -23,6 +24,7 @@ use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::PatchApplyStatus; use codex_app_server_protocol::PatchChangeKind; use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::ServerNotification; use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::SessionSource; use codex_app_server_protocol::ThreadItem; @@ -50,6 +52,11 @@ use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource as RolloutSessionSource; +use codex_protocol::protocol::TokenCountEvent; +use codex_protocol::protocol::TokenUsage; +use codex_protocol::protocol::TokenUsageInfo; +use codex_protocol::protocol::TurnAbortReason; +use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnStartedEvent; use codex_protocol::user_input::ByteRange; use codex_protocol::user_input::TextElement; @@ -278,6 +285,268 @@ async fn thread_resume_returns_rollout_history() -> Result<()> { Ok(()) } +#[tokio::test] +async fn thread_resume_emits_restored_token_usage_before_next_turn() -> Result<()> { + let server = create_mock_responses_server_repeating_assistant("Done").await; + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri())?; + + let conversation_id = create_fake_rollout_with_token_usage( + codex_home.path(), + "2025-01-05T12-00-00", + "2025-01-05T12:00:00Z", + "Saved user message", + Some("mock_provider"), + )?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let resume_id = mcp + .send_thread_resume_request(ThreadResumeParams { + thread_id: conversation_id, + ..Default::default() + }) + .await?; + let resume_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), + ) + .await??; + let ThreadResumeResponse { thread, .. } = to_response::(resume_resp)?; + + let note = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("thread/tokenUsage/updated"), + ) + .await??; + let parsed: ServerNotification = note.try_into()?; + let ServerNotification::ThreadTokenUsageUpdated(notification) = parsed else { + panic!("expected thread/tokenUsage/updated notification"); + }; + + assert_eq!(notification.thread_id, thread.id); + assert_eq!(notification.turn_id, thread.turns[0].id); + assert_eq!(notification.token_usage.total.total_tokens, 150); + assert_eq!(notification.token_usage.total.input_tokens, 120); + assert_eq!(notification.token_usage.total.cached_input_tokens, 20); + assert_eq!(notification.token_usage.total.output_tokens, 30); + assert_eq!(notification.token_usage.total.reasoning_output_tokens, 10); + assert_eq!(notification.token_usage.last.total_tokens, 90); + assert_eq!(notification.token_usage.model_context_window, Some(200_000)); + + Ok(()) +} + +#[tokio::test] +async fn thread_resume_token_usage_replay_ignores_stale_interrupted_tail_turn() -> Result<()> { + let server = create_mock_responses_server_repeating_assistant("Done").await; + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri())?; + + let filename_ts = "2025-01-05T12-00-00"; + let meta_rfc3339 = "2025-01-05T12:00:00Z"; + let conversation_id = create_fake_rollout_with_token_usage( + codex_home.path(), + filename_ts, + meta_rfc3339, + "Saved user message", + Some("mock_provider"), + )?; + let rollout_file_path = rollout_path(codex_home.path(), filename_ts, &conversation_id); + let persisted_rollout = std::fs::read_to_string(&rollout_file_path)?; + let stale_turn_id = "incomplete-turn-after-token-usage"; + let appended_rollout = [ + json!({ + "timestamp": meta_rfc3339, + "type": "event_msg", + "payload": serde_json::to_value(EventMsg::TurnStarted(TurnStartedEvent { + turn_id: stale_turn_id.to_string(), + started_at: None, + model_context_window: None, + collaboration_mode_kind: Default::default(), + }))?, + }) + .to_string(), + json!({ + "timestamp": meta_rfc3339, + "type": "event_msg", + "payload": serde_json::to_value(EventMsg::AgentMessage(AgentMessageEvent { + message: "Still running".to_string(), + phase: None, + memory_citation: None, + }))?, + }) + .to_string(), + ] + .join("\n"); + std::fs::write( + &rollout_file_path, + format!("{persisted_rollout}{appended_rollout}\n"), + )?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let resume_id = mcp + .send_thread_resume_request(ThreadResumeParams { + thread_id: conversation_id, + ..Default::default() + }) + .await?; + let resume_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), + ) + .await??; + let ThreadResumeResponse { thread, .. } = to_response::(resume_resp)?; + + assert_eq!(thread.turns.len(), 2); + assert_eq!(thread.turns[0].status, TurnStatus::Completed); + assert_eq!(thread.turns[1].id, stale_turn_id); + assert_eq!(thread.turns[1].status, TurnStatus::Interrupted); + + let note = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("thread/tokenUsage/updated"), + ) + .await??; + let parsed: ServerNotification = note.try_into()?; + let ServerNotification::ThreadTokenUsageUpdated(notification) = parsed else { + panic!("expected thread/tokenUsage/updated notification"); + }; + + assert_eq!(notification.thread_id, thread.id); + assert_eq!(notification.turn_id, thread.turns[0].id); + assert_ne!(notification.turn_id, stale_turn_id); + assert_eq!(notification.token_usage.total.total_tokens, 150); + assert_eq!(notification.token_usage.last.total_tokens, 90); + + Ok(()) +} + +#[tokio::test] +async fn thread_resume_token_usage_replay_can_belong_to_interrupted_turn() -> Result<()> { + let server = create_mock_responses_server_repeating_assistant("Done").await; + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri())?; + + let filename_ts = "2025-01-05T12-00-00"; + let meta_rfc3339 = "2025-01-05T12:00:00Z"; + let conversation_id = create_fake_rollout_with_token_usage( + codex_home.path(), + filename_ts, + meta_rfc3339, + "Saved user message", + Some("mock_provider"), + )?; + let rollout_file_path = rollout_path(codex_home.path(), filename_ts, &conversation_id); + let persisted_rollout = std::fs::read_to_string(&rollout_file_path)?; + let interrupted_turn_id = "interrupted-turn-with-token-usage"; + let appended_rollout = [ + json!({ + "timestamp": meta_rfc3339, + "type": "event_msg", + "payload": serde_json::to_value(EventMsg::TurnStarted(TurnStartedEvent { + turn_id: interrupted_turn_id.to_string(), + started_at: None, + model_context_window: None, + collaboration_mode_kind: Default::default(), + }))?, + }) + .to_string(), + json!({ + "timestamp": meta_rfc3339, + "type": "event_msg", + "payload": serde_json::to_value(EventMsg::AgentMessage(AgentMessageEvent { + message: "Interrupted after usage".to_string(), + phase: None, + memory_citation: None, + }))?, + }) + .to_string(), + json!({ + "timestamp": meta_rfc3339, + "type": "event_msg", + "payload": serde_json::to_value(EventMsg::TokenCount(TokenCountEvent { + info: Some(TokenUsageInfo { + total_token_usage: TokenUsage { + input_tokens: 180, + cached_input_tokens: 40, + output_tokens: 50, + reasoning_output_tokens: 15, + total_tokens: 230, + }, + last_token_usage: TokenUsage { + input_tokens: 90, + cached_input_tokens: 30, + output_tokens: 40, + reasoning_output_tokens: 12, + total_tokens: 130, + }, + model_context_window: Some(200_000), + }), + rate_limits: None, + }))?, + }) + .to_string(), + json!({ + "timestamp": meta_rfc3339, + "type": "event_msg", + "payload": serde_json::to_value(EventMsg::TurnAborted(TurnAbortedEvent { + turn_id: Some(interrupted_turn_id.to_string()), + reason: TurnAbortReason::Interrupted, + completed_at: None, + duration_ms: None, + }))?, + }) + .to_string(), + ] + .join("\n"); + std::fs::write( + &rollout_file_path, + format!("{persisted_rollout}{appended_rollout}\n"), + )?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let resume_id = mcp + .send_thread_resume_request(ThreadResumeParams { + thread_id: conversation_id, + ..Default::default() + }) + .await?; + let resume_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(resume_id)), + ) + .await??; + let ThreadResumeResponse { thread, .. } = to_response::(resume_resp)?; + + assert_eq!(thread.turns.len(), 2); + assert_eq!(thread.turns[0].status, TurnStatus::Completed); + assert_eq!(thread.turns[1].id, interrupted_turn_id); + assert_eq!(thread.turns[1].status, TurnStatus::Interrupted); + + let note = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("thread/tokenUsage/updated"), + ) + .await??; + let parsed: ServerNotification = note.try_into()?; + let ServerNotification::ThreadTokenUsageUpdated(notification) = parsed else { + panic!("expected thread/tokenUsage/updated notification"); + }; + + assert_eq!(notification.thread_id, thread.id); + assert_eq!(notification.turn_id, interrupted_turn_id); + assert_eq!(notification.token_usage.total.total_tokens, 230); + assert_eq!(notification.token_usage.last.total_tokens, 130); + + Ok(()) +} + #[tokio::test] async fn thread_resume_prefers_persisted_git_metadata_for_local_threads() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 844d585be5..2cff1d3174 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -2420,6 +2420,17 @@ impl Session { state.token_info().map(|info| info.total_token_usage) } + /// Returns the complete token usage snapshot currently cached for this session. + /// + /// Resume and fork reconstruction seed this state from the last persisted rollout + /// `TokenCount` event. Callers that need to replay restored usage to a client + /// should use this accessor instead of `total_token_usage`, because the app-server + /// notification includes both total and last-turn usage. + pub(crate) async fn token_usage_info(&self) -> Option { + let state = self.state.lock().await; + state.token_info() + } + pub(crate) async fn get_estimated_token_count( &self, turn_context: &TurnContext, diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index fa9f835181..3642d078c1 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -22,6 +22,7 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::Submission; use codex_protocol::protocol::ThreadMemoryMode; use codex_protocol::protocol::TokenUsage; +use codex_protocol::protocol::TokenUsageInfo; use codex_protocol::protocol::W3cTraceContext; use codex_protocol::user_input::UserInput; use codex_utils_absolute_path::AbsolutePathBuf; @@ -144,6 +145,17 @@ impl CodexThread { self.codex.session.total_token_usage().await } + /// Returns the complete token usage snapshot currently cached for this thread. + /// + /// This accessor is intentionally narrower than direct session access: it lets + /// app-server lifecycle paths replay restored usage after resume or fork without + /// exposing broader session mutation authority. A caller that only reads + /// `total_token_usage` would drop last-turn usage and make the v2 + /// `thread/tokenUsage/updated` payload incomplete. + pub async fn token_usage_info(&self) -> Option { + self.codex.session.token_usage_info().await + } + /// Records a user-role session-prefix message without creating a new user turn boundary. pub(crate) async fn inject_user_message_without_turn(&self, message: String) { let message = ResponseItem::Message { From 37bf42d5d560854b8e04801cc62685e73e87ae8b Mon Sep 17 00:00:00 2001 From: bxie-openai Date: Thu, 16 Apr 2026 13:51:43 -0700 Subject: [PATCH 08/16] [codex] Make realtime startup context truncation deterministic (#18172) ## Summary - remove the final whole-blob truncation pass from realtime startup-context assembly - enforce fixed per-section budgets, including each section heading - keep the existing per-section caps and raise the overall realtime startup-context budget to `5300`, matching the sum of those section budgets - add focused tests for the new wrapping and section-budget behavior ## Why The previous flow truncated each section and then middle-truncated the final combined startup-context blob again. Small input changes could shift that combined cut point, which made retained context unstable and caused nondeterministic tests. ## Impact Startup context now preserves section boundaries and ordering deterministically. Each section is still budgeted independently, but the final assembled blob is no longer truncated again as a single opaque string. To match that design, the overall startup-context token budget is updated to the sum of the existing section budgets rather than lowering the section caps. ## Validation - `cargo +1.93.0 test -p codex-core realtime_context` - `cargo +1.93.0 test -p codex-core --test all suite::realtime_conversation::conversation_start_injects_startup_context_from_thread_history -- --exact` - `cargo +1.93.0 test -p codex-core --test all suite::realtime_conversation::conversation_startup_context_current_thread_selects_many_turns_by_budget -- --exact` - `cargo +1.93.0 test -p codex-core --test all suite::realtime_conversation::conversation_startup_context_falls_back_to_workspace_map -- --exact` - `cargo +1.93.0 test -p codex-core --test all suite::realtime_conversation::conversation_startup_context_is_truncated_and_sent_once_per_start -- --exact` --- codex-rs/core/src/realtime_context.rs | 40 ++++++-------- codex-rs/core/src/realtime_context_tests.rs | 60 +++++++++++++++++---- codex-rs/core/src/realtime_conversation.rs | 2 +- 3 files changed, 68 insertions(+), 34 deletions(-) diff --git a/codex-rs/core/src/realtime_context.rs b/codex-rs/core/src/realtime_context.rs index fc17dab99b..c23394fa37 100644 --- a/codex-rs/core/src/realtime_context.rs +++ b/codex-rs/core/src/realtime_context.rs @@ -108,9 +108,10 @@ pub(crate) async fn build_realtime_startup_context( parts.push(section); } - let context = format_startup_context_blob(&parts.join("\n\n"), budget_tokens); + let context = format_startup_context_blob(&parts.join("\n\n")); debug!( approx_tokens = approx_token_count(&context), + requested_budget_tokens = budget_tokens, bytes = context.len(), has_current_thread_section, has_recent_work_section, @@ -440,31 +441,22 @@ fn format_section(title: &str, body: Option, budget_tokens: usize) -> Op return None; } - Some(format!( - "## {title}\n{}", - truncate_text(body, TruncationPolicy::Tokens(budget_tokens)) - )) + let heading = format!("## {title}\n"); + let body_budget = budget_tokens.saturating_sub(approx_token_count(&heading)); + if body_budget == 0 { + return None; + } + + let body = truncate_realtime_text_to_token_budget(body, body_budget); + if body.is_empty() { + return None; + } + + Some(format!("{heading}{body}")) } -fn format_startup_context_blob(body: &str, budget_tokens: usize) -> String { - let wrapper = format!("{STARTUP_CONTEXT_OPEN_TAG}\n\n{STARTUP_CONTEXT_CLOSE_TAG}"); - let mut body_budget = budget_tokens.saturating_sub(approx_token_count(&wrapper)); - - loop { - let body = truncate_text(body, TruncationPolicy::Tokens(body_budget)); - let wrapped = format!("{STARTUP_CONTEXT_OPEN_TAG}\n{body}\n{STARTUP_CONTEXT_CLOSE_TAG}"); - let wrapped_tokens = approx_token_count(&wrapped); - if wrapped_tokens <= budget_tokens || body_budget == 0 { - return wrapped; - } - - let excess_tokens = wrapped_tokens.saturating_sub(budget_tokens); - let next_budget = body_budget.saturating_sub(excess_tokens.max(1)); - if next_budget == body_budget { - return wrapped; - } - body_budget = next_budget; - } +fn format_startup_context_blob(body: &str) -> String { + format!("{STARTUP_CONTEXT_OPEN_TAG}\n{body}\n{STARTUP_CONTEXT_CLOSE_TAG}") } async fn format_thread_group( diff --git a/codex-rs/core/src/realtime_context_tests.rs b/codex-rs/core/src/realtime_context_tests.rs index 4f62affcae..99cb00f0de 100644 --- a/codex-rs/core/src/realtime_context_tests.rs +++ b/codex-rs/core/src/realtime_context_tests.rs @@ -1,6 +1,12 @@ +use super::CURRENT_THREAD_SECTION_TOKEN_BUDGET; +use super::NOTES_SECTION_TOKEN_BUDGET; +use super::RECENT_WORK_SECTION_TOKEN_BUDGET; +use super::STARTUP_CONTEXT_HEADER; +use super::WORKSPACE_SECTION_TOKEN_BUDGET; use super::build_current_thread_section; use super::build_recent_work_section; use super::build_workspace_section_with_user_root; +use super::format_section; use super::format_startup_context_blob; use chrono::TimeZone; use chrono::Utc; @@ -172,20 +178,56 @@ fn current_thread_section_keeps_latest_turns_when_history_exceeds_budget() { } #[test] -fn startup_context_blob_is_wrapped_in_tags_and_fits_budget() { - let body = format!( - "Startup context from Codex.\n{}\n{}", - "recent work ".repeat(1_200), - "workspace tree ".repeat(800), - ); +fn startup_context_blob_is_wrapped_in_tags_without_final_truncation() { + let body = "Startup context from Codex.\n## Current Thread\nhello"; + let wrapped = format_startup_context_blob(body); - let wrapped = format_startup_context_blob(&body, /*budget_tokens*/ 200); + assert_eq!( + wrapped, + "\nStartup context from Codex.\n## Current Thread\nhello\n" + ); +} + +#[test] +fn fixed_section_budgets_apply_per_section_without_total_blob_truncation() { + let body = [ + STARTUP_CONTEXT_HEADER.to_string(), + format_section( + "Current Thread", + Some("current thread ".repeat(2_000)), + CURRENT_THREAD_SECTION_TOKEN_BUDGET, + ) + .expect("current thread section"), + format_section( + "Recent Work", + Some("recent work ".repeat(3_000)), + RECENT_WORK_SECTION_TOKEN_BUDGET, + ) + .expect("recent work section"), + format_section( + "Machine / Workspace Map", + Some("workspace map ".repeat(2_500)), + WORKSPACE_SECTION_TOKEN_BUDGET, + ) + .expect("workspace section"), + format_section( + "Notes", + Some("notes ".repeat(500)), + NOTES_SECTION_TOKEN_BUDGET, + ) + .expect("notes section"), + ] + .join("\n\n"); + + let wrapped = format_startup_context_blob(&body); assert!(wrapped.starts_with("\n")); assert!(wrapped.ends_with("\n")); - assert!(wrapped.contains("Startup context from Codex.")); assert!(wrapped.contains("tokens truncated")); - assert!(wrapped.len().div_ceil(4) <= 200); + assert!(wrapped.contains("## Current Thread")); + assert!(wrapped.contains("## Recent Work")); + assert!(wrapped.contains("## Machine / Workspace Map")); + assert!(wrapped.contains("## Notes")); } #[tokio::test] diff --git a/codex-rs/core/src/realtime_conversation.rs b/codex-rs/core/src/realtime_conversation.rs index b79060259c..234dd78f38 100644 --- a/codex-rs/core/src/realtime_conversation.rs +++ b/codex-rs/core/src/realtime_conversation.rs @@ -63,7 +63,7 @@ const AUDIO_IN_QUEUE_CAPACITY: usize = 256; const USER_TEXT_IN_QUEUE_CAPACITY: usize = 64; const HANDOFF_OUT_QUEUE_CAPACITY: usize = 64; const OUTPUT_EVENTS_QUEUE_CAPACITY: usize = 256; -const REALTIME_STARTUP_CONTEXT_TOKEN_BUDGET: usize = 5_000; +const REALTIME_STARTUP_CONTEXT_TOKEN_BUDGET: usize = 5_300; const DEFAULT_REALTIME_MODEL: &str = "gpt-realtime-1.5"; pub(crate) const REALTIME_USER_TEXT_PREFIX: &str = "[USER] "; pub(crate) const REALTIME_BACKEND_TEXT_PREFIX: &str = "[BACKEND] "; From 3905f728918abd6ebaa63a4368f52754c7a3e26c Mon Sep 17 00:00:00 2001 From: starr-openai Date: Thu, 16 Apr 2026 14:16:15 -0700 Subject: [PATCH 09/16] Throttle Windows Bazel test concurrency (#18192) ## Summary - cap the Windows Bazel test lane at `--jobs=8` to reduce local runner pressure - keep Linux and macOS Bazel test concurrency unchanged - make failed-test log tailing resolve `bazel-testlogs` with the same CI config and Windows host-platform context as the failed invocation - prefer Bazel-reported `test.log` paths and normalize Windows path separators before tailing ## Context The Windows Bazel workflow currently uses `ci-windows`, which does not inherit the remote executor config. This means the lane runs the `//...` test suite locally and otherwise falls back to the repo-wide `common --jobs=30`. The new Windows-only override is intended to reduce local executor pressure without changing coverage. ## Validation Not run locally; this is a CI workflow change and the draft PR is intended to exercise the GitHub Actions lane directly. --------- Co-authored-by: Codex --- .github/scripts/run-bazel-ci.sh | 30 ++++++++++++++++++++++++++++-- .github/workflows/bazel.yml | 12 ++++++++---- 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/.github/scripts/run-bazel-ci.sh b/.github/scripts/run-bazel-ci.sh index c8c09f406a..a1d42a0576 100755 --- a/.github/scripts/run-bazel-ci.sh +++ b/.github/scripts/run-bazel-ci.sh @@ -69,12 +69,37 @@ print_bazel_test_log_tails() { local console_log="$1" local testlogs_dir local -a bazel_info_cmd=(bazel) + local -a bazel_info_args=(info) if (( ${#bazel_startup_args[@]} > 0 )); then bazel_info_cmd+=("${bazel_startup_args[@]}") fi - testlogs_dir="$(run_bazel "${bazel_info_cmd[@]:1}" info bazel-testlogs 2>/dev/null || echo bazel-testlogs)" + # `bazel info` needs the same CI config as the failed test invocation so + # platform-specific output roots match. On Windows, omitting `ci-windows` + # would point at `local_windows-fastbuild` even when the test ran with the + # MSVC host platform under `local_windows_msvc-fastbuild`. + if [[ -n "${BUILDBUDDY_API_KEY:-}" ]]; then + bazel_info_args+=( + "--config=${ci_config}" + "--remote_header=x-buildbuddy-api-key=${BUILDBUDDY_API_KEY}" + ) + fi + # Only pass flags that affect Bazel's output-root selection or repository + # lookup. Test/build-only flags such as execution logs or remote download + # mode can make `bazel info` fail, which would hide the real test log path. + for arg in "${post_config_bazel_args[@]}"; do + case "$arg" in + --host_platform=* | --repo_contents_cache=* | --repository_cache=*) + bazel_info_args+=("$arg") + ;; + esac + done + + testlogs_dir="$(run_bazel "${bazel_info_cmd[@]:1}" \ + --noexperimental_remote_repo_contents_cache \ + "${bazel_info_args[@]}" \ + bazel-testlogs 2>/dev/null || echo bazel-testlogs)" local failed_targets=() while IFS= read -r target; do @@ -95,8 +120,9 @@ print_bazel_test_log_tails() { rel_path="${rel_path/://}" local test_log="${testlogs_dir}/${rel_path}/test.log" local reported_test_log - reported_test_log="$(grep -F "FAIL: ${target} " "$console_log" | sed -nE 's#.* \(see ([^)]+/test\.log)\).*#\1#p' | head -n 1 || true)" + reported_test_log="$(grep -F "FAIL: ${target} " "$console_log" | sed -nE 's#.* \(see (.*[\\/]test\.log)\).*#\1#p' | head -n 1 || true)" if [[ -n "$reported_test_log" ]]; then + reported_test_log="${reported_test_log//\\//}" test_log="$reported_test_log" fi diff --git a/.github/workflows/bazel.yml b/.github/workflows/bazel.yml index 5d30916f9f..0328ac2073 100644 --- a/.github/workflows/bazel.yml +++ b/.github/workflows/bazel.yml @@ -86,17 +86,21 @@ jobs: --print-failed-test-logs --use-node-test-env ) + bazel_test_args=( + test + --test_tag_filters=-argument-comment-lint + --test_verbose_timeout_warnings + --build_metadata=COMMIT_SHA=${GITHUB_SHA} + ) if [[ "${RUNNER_OS}" == "Windows" ]]; then bazel_wrapper_args+=(--windows-msvc-host-platform) + bazel_test_args+=(--jobs=8) fi ./.github/scripts/run-bazel-ci.sh \ "${bazel_wrapper_args[@]}" \ -- \ - test \ - --test_tag_filters=-argument-comment-lint \ - --test_verbose_timeout_warnings \ - --build_metadata=COMMIT_SHA=${GITHUB_SHA} \ + "${bazel_test_args[@]}" \ -- \ "${bazel_targets[@]}" From 0708cc78cb9da20087e3013cab71e880290dc427 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Thu, 16 Apr 2026 14:21:29 -0700 Subject: [PATCH 10/16] [codex] Split codex op handlers (#18200) Start splitting the codex.rs --- codex-rs/core/src/codex.rs | 1261 +-------------------------- codex-rs/core/src/codex/handlers.rs | 1230 ++++++++++++++++++++++++++ 2 files changed, 1234 insertions(+), 1257 deletions(-) create mode 100644 codex-rs/core/src/codex/handlers.rs diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 2cff1d3174..4e09542d88 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -30,10 +30,6 @@ use crate::mcp_tool_exposure::build_mcp_tool_exposure; use crate::parse_turn_item; use crate::path_utils::normalize_for_native_workdir; use crate::realtime_conversation::RealtimeConversationManager; -use crate::realtime_conversation::handle_audio as handle_realtime_conversation_audio; -use crate::realtime_conversation::handle_close as handle_realtime_conversation_close; -use crate::realtime_conversation::handle_start as handle_realtime_conversation_start; -use crate::realtime_conversation::handle_text as handle_realtime_conversation_text; use crate::render_skills_section; use crate::rollout::find_thread_name_by_id; use crate::session_prefix::format_subagent_notification_message; @@ -170,7 +166,6 @@ use tokio_util::sync::CancellationToken; use toml::Value as TomlValue; use tracing::Instrument; use tracing::debug; -use tracing::debug_span; use tracing::error; use tracing::field; use tracing::info; @@ -206,8 +201,12 @@ use codex_protocol::error::Result as CodexResult; #[cfg(test)] use codex_protocol::exec_output::StreamOutput; +mod handlers; mod rollout_reconstruction; #[cfg(test)] +use self::handlers::submission_dispatch_span; +use self::handlers::submission_loop; +#[cfg(test)] mod rollout_reconstruction_tests; #[derive(Debug, PartialEq)] @@ -4733,1258 +4732,6 @@ pub(crate) fn emit_subagent_session_started( }); } -async fn submission_loop(sess: Arc, config: Arc, rx_sub: Receiver) { - // To break out of this loop, send Op::Shutdown. - while let Ok(sub) = rx_sub.recv().await { - debug!(?sub, "Submission"); - let dispatch_span = submission_dispatch_span(&sub); - let should_exit = async { - match sub.op.clone() { - Op::Interrupt => { - handlers::interrupt(&sess).await; - false - } - Op::CleanBackgroundTerminals => { - handlers::clean_background_terminals(&sess).await; - false - } - Op::RealtimeConversationStart(params) => { - if let Err(err) = - handle_realtime_conversation_start(&sess, sub.id.clone(), params).await - { - sess.send_event_raw(Event { - id: sub.id.clone(), - msg: EventMsg::Error(ErrorEvent { - message: err.to_string(), - codex_error_info: Some(CodexErrorInfo::Other), - }), - }) - .await; - } - false - } - Op::RealtimeConversationAudio(params) => { - handle_realtime_conversation_audio(&sess, sub.id.clone(), params).await; - false - } - Op::RealtimeConversationText(params) => { - handle_realtime_conversation_text(&sess, sub.id.clone(), params).await; - false - } - Op::RealtimeConversationClose => { - handle_realtime_conversation_close(&sess, sub.id.clone()).await; - false - } - Op::RealtimeConversationListVoices => { - handlers::realtime_conversation_list_voices(&sess, sub.id.clone()).await; - false - } - Op::OverrideTurnContext { - cwd, - approval_policy, - approvals_reviewer, - sandbox_policy, - windows_sandbox_level, - model, - effort, - summary, - service_tier, - collaboration_mode, - personality, - } => { - let collaboration_mode = if let Some(collab_mode) = collaboration_mode { - collab_mode - } else { - let state = sess.state.lock().await; - state.session_configuration.collaboration_mode.with_updates( - model.clone(), - effort, - /*developer_instructions*/ None, - ) - }; - handlers::override_turn_context( - &sess, - sub.id.clone(), - SessionSettingsUpdate { - cwd, - approval_policy, - approvals_reviewer, - sandbox_policy, - windows_sandbox_level, - collaboration_mode: Some(collaboration_mode), - reasoning_summary: summary, - service_tier, - personality, - ..Default::default() - }, - ) - .await; - false - } - Op::UserInput { .. } | Op::UserTurn { .. } => { - handlers::user_input_or_turn(&sess, sub.id.clone(), sub.op).await; - false - } - Op::InterAgentCommunication { communication } => { - handlers::inter_agent_communication(&sess, sub.id.clone(), communication).await; - false - } - Op::ExecApproval { - id: approval_id, - turn_id, - decision, - } => { - handlers::exec_approval(&sess, approval_id, turn_id, decision).await; - false - } - Op::PatchApproval { id, decision } => { - handlers::patch_approval(&sess, id, decision).await; - false - } - Op::UserInputAnswer { id, response } => { - handlers::request_user_input_response(&sess, id, response).await; - false - } - Op::RequestPermissionsResponse { id, response } => { - handlers::request_permissions_response(&sess, id, response).await; - false - } - Op::DynamicToolResponse { id, response } => { - handlers::dynamic_tool_response(&sess, id, response).await; - false - } - Op::AddToHistory { text } => { - handlers::add_to_history(&sess, &config, text).await; - false - } - Op::GetHistoryEntryRequest { offset, log_id } => { - handlers::get_history_entry_request( - &sess, - &config, - sub.id.clone(), - offset, - log_id, - ) - .await; - false - } - Op::ListMcpTools => { - handlers::list_mcp_tools(&sess, &config, sub.id.clone()).await; - false - } - Op::RefreshMcpServers { config } => { - handlers::refresh_mcp_servers(&sess, config).await; - false - } - Op::ReloadUserConfig => { - handlers::reload_user_config(&sess).await; - false - } - Op::ListSkills { cwds, force_reload } => { - handlers::list_skills(&sess, sub.id.clone(), cwds, force_reload).await; - false - } - Op::Undo => { - handlers::undo(&sess, sub.id.clone()).await; - false - } - Op::Compact => { - handlers::compact(&sess, sub.id.clone()).await; - false - } - Op::DropMemories => { - handlers::drop_memories(&sess, &config, sub.id.clone()).await; - false - } - Op::UpdateMemories => { - handlers::update_memories(&sess, &config, sub.id.clone()).await; - false - } - Op::ThreadRollback { num_turns } => { - handlers::thread_rollback(&sess, sub.id.clone(), num_turns).await; - false - } - Op::SetThreadName { name } => { - handlers::set_thread_name(&sess, sub.id.clone(), name).await; - false - } - Op::SetThreadMemoryMode { mode } => { - handlers::set_thread_memory_mode(&sess, sub.id.clone(), mode).await; - false - } - Op::RunUserShellCommand { command } => { - handlers::run_user_shell_command(&sess, sub.id.clone(), command).await; - false - } - Op::ResolveElicitation { - server_name, - request_id, - decision, - content, - meta, - } => { - handlers::resolve_elicitation( - &sess, - server_name, - request_id, - decision, - content, - meta, - ) - .await; - false - } - Op::Shutdown => handlers::shutdown(&sess, sub.id.clone()).await, - Op::Review { review_request } => { - handlers::review(&sess, &config, sub.id.clone(), review_request).await; - false - } - _ => false, // Ignore unknown ops; enum is non_exhaustive to allow extensions. - } - } - .instrument(dispatch_span) - .await; - if should_exit { - break; - } - } - // Also drain cached guardian state if the submission loop exits because - // the channel closed without receiving an explicit shutdown op. - sess.guardian_review_session.shutdown().await; - debug!("Agent loop exited"); -} - -fn submission_dispatch_span(sub: &Submission) -> tracing::Span { - let op_name = sub.op.kind(); - let span_name = format!("op.dispatch.{op_name}"); - let dispatch_span = match &sub.op { - Op::RealtimeConversationAudio(_) => { - debug_span!( - "submission_dispatch", - otel.name = span_name.as_str(), - submission.id = sub.id.as_str(), - codex.op = op_name - ) - } - _ => info_span!( - "submission_dispatch", - otel.name = span_name.as_str(), - submission.id = sub.id.as_str(), - codex.op = op_name - ), - }; - if let Some(trace) = sub.trace.as_ref() - && !set_parent_from_w3c_trace_context(&dispatch_span, trace) - { - warn!( - submission.id = sub.id.as_str(), - "ignoring invalid submission trace carrier" - ); - } - dispatch_span -} - -/// Operation handlers -mod handlers { - use crate::codex::Session; - use crate::codex::SessionSettingsUpdate; - use crate::codex::SteerInputError; - - use crate::codex::spawn_review_thread; - use crate::config::Config; - use crate::config_loader::CloudRequirementsLoader; - use crate::config_loader::LoaderOverrides; - use crate::config_loader::load_config_layers_state; - use crate::realtime_context::REALTIME_TURN_TOKEN_BUDGET; - use crate::realtime_context::truncate_realtime_text_to_token_budget; - use crate::realtime_conversation::REALTIME_USER_TEXT_PREFIX; - use crate::realtime_conversation::prefix_realtime_v2_text; - use codex_features::Feature; - use codex_utils_absolute_path::AbsolutePathBuf; - - use crate::review_prompts::resolve_review_request; - use crate::rollout::RolloutRecorder; - use crate::rollout::read_session_meta_line; - use crate::tasks::CompactTask; - use crate::tasks::UndoTask; - use crate::tasks::UserShellCommandMode; - use crate::tasks::UserShellCommandTask; - use crate::tasks::execute_user_shell_command; - use codex_mcp::collect_mcp_snapshot_from_manager; - use codex_mcp::compute_auth_statuses; - use codex_protocol::protocol::CodexErrorInfo; - use codex_protocol::protocol::ErrorEvent; - use codex_protocol::protocol::Event; - use codex_protocol::protocol::EventMsg; - use codex_protocol::protocol::InterAgentCommunication; - use codex_protocol::protocol::ListSkillsResponseEvent; - use codex_protocol::protocol::McpServerRefreshConfig; - use codex_protocol::protocol::Op; - use codex_protocol::protocol::RealtimeConversationListVoicesResponseEvent; - use codex_protocol::protocol::RealtimeVoicesList; - use codex_protocol::protocol::ReviewDecision; - use codex_protocol::protocol::ReviewRequest; - use codex_protocol::protocol::RolloutItem; - use codex_protocol::protocol::SkillErrorInfo; - use codex_protocol::protocol::SkillsListEntry; - use codex_protocol::protocol::ThreadMemoryMode; - use codex_protocol::protocol::ThreadNameUpdatedEvent; - use codex_protocol::protocol::ThreadRolledBackEvent; - use codex_protocol::protocol::TurnAbortReason; - use codex_protocol::protocol::WarningEvent; - use codex_protocol::request_permissions::RequestPermissionsResponse; - use codex_protocol::request_user_input::RequestUserInputResponse; - - use crate::context_manager::is_user_turn_boundary; - use codex_protocol::config_types::CollaborationMode; - use codex_protocol::config_types::ModeKind; - use codex_protocol::config_types::Settings; - use codex_protocol::dynamic_tools::DynamicToolResponse; - use codex_protocol::items::UserMessageItem; - use codex_protocol::mcp::RequestId as ProtocolRequestId; - use codex_protocol::user_input::UserInput; - use codex_rmcp_client::ElicitationAction; - use codex_rmcp_client::ElicitationResponse; - use serde_json::Value; - use std::path::PathBuf; - use std::sync::Arc; - use tracing::debug; - use tracing::info; - use tracing::warn; - - pub async fn interrupt(sess: &Arc) { - sess.interrupt_task().await; - } - - pub async fn clean_background_terminals(sess: &Arc) { - sess.close_unified_exec_processes().await; - } - - pub async fn realtime_conversation_list_voices(sess: &Session, sub_id: String) { - sess.send_event_raw(Event { - id: sub_id, - msg: EventMsg::RealtimeConversationListVoicesResponse( - RealtimeConversationListVoicesResponseEvent { - voices: RealtimeVoicesList::builtin(), - }, - ), - }) - .await; - } - - pub async fn override_turn_context( - sess: &Session, - sub_id: String, - updates: SessionSettingsUpdate, - ) { - if let Err(err) = sess.update_settings(updates).await { - sess.send_event_raw(Event { - id: sub_id, - msg: EventMsg::Error(ErrorEvent { - message: err.to_string(), - codex_error_info: Some(CodexErrorInfo::BadRequest), - }), - }) - .await; - } - } - - pub async fn user_input_or_turn(sess: &Arc, sub_id: String, op: Op) { - user_input_or_turn_inner( - sess, - sub_id, - op, - /*mirror_user_text_to_realtime*/ Some(()), - ) - .await; - } - - pub(super) async fn user_input_or_turn_inner( - sess: &Arc, - sub_id: String, - op: Op, - mirror_user_text_to_realtime: Option<()>, - ) { - let (items, updates, responsesapi_client_metadata) = match op { - Op::UserTurn { - cwd, - approval_policy, - approvals_reviewer, - sandbox_policy, - model, - effort, - summary, - service_tier, - final_output_json_schema, - items, - collaboration_mode, - personality, - } => { - let collaboration_mode = collaboration_mode.or_else(|| { - Some(CollaborationMode { - mode: ModeKind::Default, - settings: Settings { - model: model.clone(), - reasoning_effort: effort, - developer_instructions: None, - }, - }) - }); - ( - items, - SessionSettingsUpdate { - cwd: Some(cwd), - approval_policy: Some(approval_policy), - approvals_reviewer, - sandbox_policy: Some(sandbox_policy), - windows_sandbox_level: None, - collaboration_mode, - reasoning_summary: summary, - service_tier, - final_output_json_schema: Some(final_output_json_schema), - personality, - app_server_client_name: None, - app_server_client_version: None, - }, - None, - ) - } - Op::UserInput { - items, - final_output_json_schema, - responsesapi_client_metadata, - } => ( - items, - SessionSettingsUpdate { - final_output_json_schema: Some(final_output_json_schema), - ..Default::default() - }, - responsesapi_client_metadata, - ), - _ => unreachable!(), - }; - - let Ok(current_context) = sess.new_turn_with_sub_id(sub_id.clone(), updates).await else { - // new_turn_with_sub_id already emits the error event. - return; - }; - sess.maybe_emit_unknown_model_warning_for_turn(current_context.as_ref()) - .await; - let accepted_items = match sess - .steer_input( - items.clone(), - /*expected_turn_id*/ None, - responsesapi_client_metadata.clone(), - ) - .await - { - Ok(_) => { - current_context.session_telemetry.user_prompt(&items); - Some(items) - } - Err(SteerInputError::NoActiveTurn(items)) => { - if let Some(responsesapi_client_metadata) = responsesapi_client_metadata { - current_context - .turn_metadata_state - .set_responsesapi_client_metadata(responsesapi_client_metadata); - } - current_context.session_telemetry.user_prompt(&items); - sess.refresh_mcp_servers_if_requested(¤t_context) - .await; - let accepted_items = items.clone(); - sess.spawn_task( - Arc::clone(¤t_context), - items, - crate::tasks::RegularTask::new(), - ) - .await; - Some(accepted_items) - } - Err(err) => { - sess.send_event_raw(Event { - id: sub_id, - msg: EventMsg::Error(err.to_error_event()), - }) - .await; - None - } - }; - if let (Some(items), Some(())) = (accepted_items, mirror_user_text_to_realtime) { - self::mirror_user_text_to_realtime(sess, &items).await; - } - } - - async fn mirror_user_text_to_realtime(sess: &Arc, items: &[UserInput]) { - let text = UserMessageItem::new(items).message(); - if text.is_empty() { - return; - } - let text = if sess.conversation.is_running_v2().await { - prefix_realtime_v2_text(text, REALTIME_USER_TEXT_PREFIX) - } else { - text - }; - let text = truncate_realtime_text_to_token_budget(&text, REALTIME_TURN_TOKEN_BUDGET); - if text.is_empty() { - return; - } - if sess.conversation.running_state().await.is_none() { - return; - } - if let Err(err) = sess.conversation.text_in(text).await { - debug!("failed to mirror user text to realtime conversation: {err}"); - } - } - - /// Records an inter-agent assistant envelope, then lets the shared pending-work scheduler - /// decide whether an idle session should start a regular turn. - pub async fn inter_agent_communication( - sess: &Arc, - sub_id: String, - communication: InterAgentCommunication, - ) { - let trigger_turn = communication.trigger_turn; - sess.enqueue_mailbox_communication(communication); - if trigger_turn { - sess.maybe_start_turn_for_pending_work_with_sub_id(sub_id) - .await; - } - } - - pub async fn run_user_shell_command(sess: &Arc, sub_id: String, command: String) { - if let Some((turn_context, cancellation_token)) = - sess.active_turn_context_and_cancellation_token().await - { - let session = Arc::clone(sess); - tokio::spawn(async move { - execute_user_shell_command( - session, - turn_context, - command, - cancellation_token, - UserShellCommandMode::ActiveTurnAuxiliary, - ) - .await; - }); - return; - } - - let turn_context = sess.new_default_turn_with_sub_id(sub_id).await; - sess.spawn_task( - Arc::clone(&turn_context), - Vec::new(), - UserShellCommandTask::new(command), - ) - .await; - } - - pub async fn resolve_elicitation( - sess: &Arc, - server_name: String, - request_id: ProtocolRequestId, - decision: codex_protocol::approvals::ElicitationAction, - content: Option, - meta: Option, - ) { - let action = match decision { - codex_protocol::approvals::ElicitationAction::Accept => ElicitationAction::Accept, - codex_protocol::approvals::ElicitationAction::Decline => ElicitationAction::Decline, - codex_protocol::approvals::ElicitationAction::Cancel => ElicitationAction::Cancel, - }; - let content = match action { - // Preserve the legacy fallback for clients that only send an action. - ElicitationAction::Accept => Some(content.unwrap_or_else(|| serde_json::json!({}))), - ElicitationAction::Decline | ElicitationAction::Cancel => None, - }; - let response = ElicitationResponse { - action, - content, - meta, - }; - let request_id = match request_id { - ProtocolRequestId::String(value) => { - rmcp::model::NumberOrString::String(std::sync::Arc::from(value)) - } - ProtocolRequestId::Integer(value) => rmcp::model::NumberOrString::Number(value), - }; - if let Err(err) = sess - .resolve_elicitation(server_name, request_id, response) - .await - { - warn!( - error = %err, - "failed to resolve elicitation request in session" - ); - } - } - - /// Propagate a user's exec approval decision to the session. - /// Also optionally applies an execpolicy amendment. - pub async fn exec_approval( - sess: &Arc, - approval_id: String, - turn_id: Option, - decision: ReviewDecision, - ) { - let event_turn_id = turn_id.unwrap_or_else(|| approval_id.clone()); - if let ReviewDecision::ApprovedExecpolicyAmendment { - proposed_execpolicy_amendment, - } = &decision - { - match sess - .persist_execpolicy_amendment(proposed_execpolicy_amendment) - .await - { - Ok(()) => { - sess.record_execpolicy_amendment_message( - &event_turn_id, - proposed_execpolicy_amendment, - ) - .await; - } - Err(err) => { - let message = format!("Failed to apply execpolicy amendment: {err}"); - tracing::warn!("{message}"); - let warning = EventMsg::Warning(WarningEvent { message }); - sess.send_event_raw(Event { - id: event_turn_id.clone(), - msg: warning, - }) - .await; - } - } - } - match decision { - ReviewDecision::Abort => { - sess.interrupt_task().await; - } - other => sess.notify_approval(&approval_id, other).await, - } - } - - pub async fn patch_approval(sess: &Arc, id: String, decision: ReviewDecision) { - match decision { - ReviewDecision::Abort => { - sess.interrupt_task().await; - } - other => sess.notify_approval(&id, other).await, - } - } - - pub async fn request_user_input_response( - sess: &Arc, - id: String, - response: RequestUserInputResponse, - ) { - sess.notify_user_input_response(&id, response).await; - } - - pub async fn request_permissions_response( - sess: &Arc, - id: String, - response: RequestPermissionsResponse, - ) { - sess.notify_request_permissions_response(&id, response) - .await; - } - - pub async fn dynamic_tool_response( - sess: &Arc, - id: String, - response: DynamicToolResponse, - ) { - sess.notify_dynamic_tool_response(&id, response).await; - } - - pub async fn add_to_history(sess: &Arc, config: &Arc, text: String) { - let id = sess.conversation_id; - let config = Arc::clone(config); - tokio::spawn(async move { - if let Err(e) = crate::message_history::append_entry(&text, &id, &config).await { - warn!("failed to append to message history: {e}"); - } - }); - } - - pub async fn get_history_entry_request( - sess: &Arc, - config: &Arc, - sub_id: String, - offset: usize, - log_id: u64, - ) { - let config = Arc::clone(config); - let sess_clone = Arc::clone(sess); - - tokio::spawn(async move { - // Run lookup in blocking thread because it does file IO + locking. - let entry_opt = tokio::task::spawn_blocking(move || { - crate::message_history::lookup(log_id, offset, &config) - }) - .await - .unwrap_or(None); - - let event = Event { - id: sub_id, - msg: EventMsg::GetHistoryEntryResponse( - codex_protocol::protocol::GetHistoryEntryResponseEvent { - offset, - log_id, - entry: entry_opt.map(|e| codex_protocol::message_history::HistoryEntry { - conversation_id: e.session_id, - ts: e.ts, - text: e.text, - }), - }, - ), - }; - - sess_clone.send_event_raw(event).await; - }); - } - - pub async fn refresh_mcp_servers(sess: &Arc, refresh_config: McpServerRefreshConfig) { - let mut guard = sess.pending_mcp_server_refresh_config.lock().await; - *guard = Some(refresh_config); - } - - pub async fn reload_user_config(sess: &Arc) { - sess.reload_user_config_layer().await; - } - - pub async fn list_mcp_tools(sess: &Session, config: &Arc, sub_id: String) { - let mcp_connection_manager = sess.services.mcp_connection_manager.read().await; - let auth = sess.services.auth_manager.auth().await; - let mcp_servers = sess - .services - .mcp_manager - .effective_servers(config, auth.as_ref()) - .await; - let snapshot = collect_mcp_snapshot_from_manager( - &mcp_connection_manager, - compute_auth_statuses(mcp_servers.iter(), config.mcp_oauth_credentials_store_mode) - .await, - ) - .await; - let event = Event { - id: sub_id, - msg: EventMsg::McpListToolsResponse(snapshot), - }; - sess.send_event_raw(event).await; - } - - pub async fn list_skills( - sess: &Session, - sub_id: String, - cwds: Vec, - force_reload: bool, - ) { - let default_cwd = { - let state = sess.state.lock().await; - state.session_configuration.cwd.to_path_buf() - }; - let cwds = if cwds.is_empty() { - vec![default_cwd] - } else { - cwds - }; - - let skills_manager = &sess.services.skills_manager; - let plugins_manager = &sess.services.plugins_manager; - let fs = sess - .services - .environment - .as_ref() - .map(|environment| environment.get_filesystem()); - let config = sess.get_config().await; - let codex_home = sess.codex_home().await; - let mut skills = Vec::new(); - let empty_cli_overrides: &[(String, toml::Value)] = &[]; - for cwd in cwds { - let cwd_abs = match AbsolutePathBuf::relative_to_current_dir(cwd.as_path()) { - Ok(path) => path, - Err(err) => { - let error_path = cwd.clone(); - skills.push(SkillsListEntry { - cwd, - skills: Vec::new(), - errors: vec![SkillErrorInfo { - path: error_path, - message: err.to_string(), - }], - }); - continue; - } - }; - let config_layer_stack = match load_config_layers_state( - &codex_home, - Some(cwd_abs.clone()), - empty_cli_overrides, - LoaderOverrides::default(), - CloudRequirementsLoader::default(), - ) - .await - { - Ok(config_layer_stack) => config_layer_stack, - Err(err) => { - let error_path = cwd.clone(); - skills.push(SkillsListEntry { - cwd, - skills: Vec::new(), - errors: vec![SkillErrorInfo { - path: error_path, - message: err.to_string(), - }], - }); - continue; - } - }; - let effective_skill_roots = plugins_manager - .effective_skill_roots_for_layer_stack( - &config_layer_stack, - config.features.enabled(Feature::Plugins), - ) - .await; - let skills_input = crate::SkillsLoadInput::new( - cwd_abs.clone(), - effective_skill_roots, - config_layer_stack, - config.bundled_skills_enabled(), - ); - let outcome = skills_manager - .skills_for_cwd(&skills_input, force_reload, fs.clone()) - .await; - let errors = super::errors_to_info(&outcome.errors); - let skills_metadata = super::skills_to_info(&outcome.skills, &outcome.disabled_paths); - skills.push(SkillsListEntry { - cwd, - skills: skills_metadata, - errors, - }); - } - - let event = Event { - id: sub_id, - msg: EventMsg::ListSkillsResponse(ListSkillsResponseEvent { skills }), - }; - sess.send_event_raw(event).await; - } - - pub async fn undo(sess: &Arc, sub_id: String) { - let turn_context = sess.new_default_turn_with_sub_id(sub_id).await; - sess.spawn_task(turn_context, Vec::new(), UndoTask::new()) - .await; - } - - pub async fn compact(sess: &Arc, sub_id: String) { - let turn_context = sess.new_default_turn_with_sub_id(sub_id).await; - - sess.spawn_task( - Arc::clone(&turn_context), - vec![UserInput::Text { - text: turn_context.compact_prompt().to_string(), - // Compaction prompt is synthesized; no UI element ranges to preserve. - text_elements: Vec::new(), - }], - CompactTask, - ) - .await; - } - - pub async fn drop_memories(sess: &Arc, config: &Arc, sub_id: String) { - let mut errors = Vec::new(); - - if let Some(state_db) = sess.services.state_db.as_deref() { - if let Err(err) = state_db.clear_memory_data().await { - errors.push(format!("failed clearing memory rows from state db: {err}")); - } - } else { - errors.push("state db unavailable; memory rows were not cleared".to_string()); - } - - if let Err(err) = crate::memories::clear_memory_roots_contents(&config.codex_home).await { - errors.push(format!( - "failed clearing memory directories under {}: {err}", - config.codex_home.display() - )); - } - - if errors.is_empty() { - let memory_root = crate::memories::memory_root(&config.codex_home); - sess.send_event_raw(Event { - id: sub_id, - msg: EventMsg::Warning(WarningEvent { - message: format!( - "Dropped memories at {} and cleared memory rows from state db.", - memory_root.display() - ), - }), - }) - .await; - return; - } - - sess.send_event_raw(Event { - id: sub_id, - msg: EventMsg::Error(ErrorEvent { - message: format!("Memory drop completed with errors: {}", errors.join("; ")), - codex_error_info: Some(CodexErrorInfo::Other), - }), - }) - .await; - } - - pub async fn update_memories(sess: &Arc, config: &Arc, sub_id: String) { - let session_source = { - let state = sess.state.lock().await; - state.session_configuration.session_source.clone() - }; - - crate::memories::start_memories_startup_task(sess, Arc::clone(config), &session_source); - - sess.send_event_raw(Event { - id: sub_id.clone(), - msg: EventMsg::Warning(WarningEvent { - message: "Memory update triggered.".to_string(), - }), - }) - .await; - } - - pub async fn thread_rollback(sess: &Arc, sub_id: String, num_turns: u32) { - if num_turns == 0 { - sess.send_event_raw(Event { - id: sub_id, - msg: EventMsg::Error(ErrorEvent { - message: "num_turns must be >= 1".to_string(), - codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed), - }), - }) - .await; - return; - } - - let has_active_turn = { sess.active_turn.lock().await.is_some() }; - if has_active_turn { - sess.send_event_raw(Event { - id: sub_id, - msg: EventMsg::Error(ErrorEvent { - message: "Cannot rollback while a turn is in progress.".to_string(), - codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed), - }), - }) - .await; - return; - } - - let turn_context = sess.new_default_turn_with_sub_id(sub_id).await; - let rollout_path = { - let recorder = { - let guard = sess.services.rollout.lock().await; - guard.clone() - }; - let Some(recorder) = recorder else { - sess.send_event_raw(Event { - id: turn_context.sub_id.clone(), - msg: EventMsg::Error(ErrorEvent { - message: "thread rollback requires a persisted rollout path".to_string(), - codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed), - }), - }) - .await; - return; - }; - recorder.rollout_path().to_path_buf() - }; - if let Some(recorder) = { - let guard = sess.services.rollout.lock().await; - guard.clone() - } && let Err(err) = recorder.flush().await - { - sess.send_event_raw(Event { - id: turn_context.sub_id.clone(), - msg: EventMsg::Error(ErrorEvent { - message: format!( - "failed to flush rollout `{}` for rollback replay: {err}", - rollout_path.display() - ), - codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed), - }), - }) - .await; - return; - } - - let initial_history = - match RolloutRecorder::get_rollout_history(rollout_path.as_path()).await { - Ok(history) => history, - Err(err) => { - sess.send_event_raw(Event { - id: turn_context.sub_id.clone(), - msg: EventMsg::Error(ErrorEvent { - message: format!( - "failed to load rollout `{}` for rollback replay: {err}", - rollout_path.display() - ), - codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed), - }), - }) - .await; - return; - } - }; - - let rollback_event = ThreadRolledBackEvent { num_turns }; - let rollback_msg = EventMsg::ThreadRolledBack(rollback_event.clone()); - let replay_items = initial_history - .get_rollout_items() - .into_iter() - .chain(std::iter::once(RolloutItem::EventMsg(rollback_msg.clone()))) - .collect::>(); - sess.apply_rollout_reconstruction(turn_context.as_ref(), replay_items.as_slice()) - .await; - sess.recompute_token_usage(turn_context.as_ref()).await; - - sess.persist_rollout_items(&[RolloutItem::EventMsg(rollback_msg.clone())]) - .await; - if let Err(err) = sess.flush_rollout().await { - sess.send_event( - turn_context.as_ref(), - EventMsg::Warning(WarningEvent { - message: format!( - "Rolled the thread back, but failed to save the rollback marker. Codex will continue retrying. Error: {err}" - ), - }), - ) - .await; - } - - sess.deliver_event_raw(Event { - id: turn_context.sub_id.clone(), - msg: rollback_msg, - }) - .await; - } - - async fn persist_thread_name_update( - sess: &Arc, - event: ThreadNameUpdatedEvent, - ) -> anyhow::Result { - let msg = EventMsg::ThreadNameUpdated(event); - let item = RolloutItem::EventMsg(msg.clone()); - let recorder = { - let guard = sess.services.rollout.lock().await; - guard.clone() - } - .ok_or_else(|| anyhow::anyhow!("Session persistence is disabled; cannot rename thread."))?; - recorder.persist().await?; - recorder.record_items(std::slice::from_ref(&item)).await?; - recorder.flush().await?; - Ok(msg) - } - - pub(super) async fn persist_thread_memory_mode_update( - sess: &Arc, - mode: ThreadMemoryMode, - ) -> anyhow::Result<()> { - let recorder = { - let guard = sess.services.rollout.lock().await; - guard.clone() - } - .ok_or_else(|| { - anyhow::anyhow!("Session persistence is disabled; cannot update thread memory mode.") - })?; - recorder.persist().await?; - recorder.flush().await?; - - let rollout_path = recorder.rollout_path().to_path_buf(); - let mut session_meta = read_session_meta_line(rollout_path.as_path()).await?; - if session_meta.meta.id != sess.conversation_id { - anyhow::bail!( - "rollout session metadata id mismatch: expected {}, found {}", - sess.conversation_id, - session_meta.meta.id - ); - } - session_meta.meta.memory_mode = Some( - match mode { - ThreadMemoryMode::Enabled => "enabled", - ThreadMemoryMode::Disabled => "disabled", - } - .to_string(), - ); - - let item = RolloutItem::SessionMeta(session_meta); - recorder.record_items(std::slice::from_ref(&item)).await?; - recorder.flush().await?; - Ok(()) - } - - /// Persists the thread name in the rollout and state database, updates in-memory state, and - /// emits a `ThreadNameUpdated` event on success. - pub async fn set_thread_name(sess: &Arc, sub_id: String, name: String) { - let Some(name) = crate::util::normalize_thread_name(&name) else { - let event = Event { - id: sub_id, - msg: EventMsg::Error(ErrorEvent { - message: "Thread name cannot be empty.".to_string(), - codex_error_info: Some(CodexErrorInfo::BadRequest), - }), - }; - sess.send_event_raw(event).await; - return; - }; - - let updated = ThreadNameUpdatedEvent { - thread_id: sess.conversation_id, - thread_name: Some(name.clone()), - }; - - let msg = match persist_thread_name_update(sess, updated).await { - Ok(msg) => msg, - Err(err) => { - warn!("Failed to persist thread name update to rollout: {err}"); - let event = Event { - id: sub_id, - msg: EventMsg::Error(ErrorEvent { - message: err.to_string(), - codex_error_info: Some(CodexErrorInfo::Other), - }), - }; - sess.send_event_raw(event).await; - return; - } - }; - - if let Some(state_db) = sess.services.state_db.as_deref() - && let Err(err) = state_db - .update_thread_title(sess.conversation_id, &name) - .await - { - warn!("Failed to update thread title in state db: {err}"); - } - - { - let mut state = sess.state.lock().await; - state.session_configuration.thread_name = Some(name.clone()); - } - - let codex_home = sess.codex_home().await; - if let Err(err) = - crate::rollout::append_thread_name(&codex_home, sess.conversation_id, &name).await - { - warn!("Failed to update legacy thread name index: {err}"); - } - - sess.deliver_event_raw(Event { id: sub_id, msg }).await; - } - - /// Persists thread-level memory mode metadata for the active session. - /// - /// This does not involve the model and only affects whether the thread is - /// eligible for future memory generation. - pub async fn set_thread_memory_mode( - sess: &Arc, - sub_id: String, - mode: ThreadMemoryMode, - ) { - if let Err(err) = persist_thread_memory_mode_update(sess, mode).await { - warn!("Failed to persist thread memory mode update to rollout: {err}"); - let event = Event { - id: sub_id, - msg: EventMsg::Error(ErrorEvent { - message: err.to_string(), - codex_error_info: Some(CodexErrorInfo::Other), - }), - }; - sess.send_event_raw(event).await; - } - } - - pub async fn shutdown(sess: &Arc, sub_id: String) -> bool { - sess.abort_all_tasks(TurnAbortReason::Interrupted).await; - let _ = sess.conversation.shutdown().await; - sess.services - .unified_exec_manager - .terminate_all_processes() - .await; - sess.guardian_review_session.shutdown().await; - info!("Shutting down Codex instance"); - let history = sess.clone_history().await; - let turn_count = history - .raw_items() - .iter() - .filter(|item| is_user_turn_boundary(item)) - .count(); - sess.services.session_telemetry.counter( - "codex.conversation.turn.count", - i64::try_from(turn_count).unwrap_or(0), - &[], - ); - - // Gracefully flush and shutdown rollout recorder on session end so tests - // that inspect the rollout file do not race with the background writer. - let recorder_opt = { - let mut guard = sess.services.rollout.lock().await; - guard.take() - }; - if let Some(rec) = recorder_opt - && let Err(e) = rec.shutdown().await - { - warn!("failed to shutdown rollout recorder: {e}"); - let event = Event { - id: sub_id.clone(), - msg: EventMsg::Error(ErrorEvent { - message: "Failed to shutdown rollout recorder".to_string(), - codex_error_info: Some(CodexErrorInfo::Other), - }), - }; - sess.send_event_raw(event).await; - } - - let event = Event { - id: sub_id, - msg: EventMsg::ShutdownComplete, - }; - sess.send_event_raw(event).await; - true - } - - pub async fn review( - sess: &Arc, - config: &Arc, - sub_id: String, - review_request: ReviewRequest, - ) { - let turn_context = sess.new_default_turn_with_sub_id(sub_id.clone()).await; - sess.maybe_emit_unknown_model_warning_for_turn(turn_context.as_ref()) - .await; - sess.refresh_mcp_servers_if_requested(&turn_context).await; - match resolve_review_request(review_request, &turn_context.cwd) { - Ok(resolved) => { - spawn_review_thread( - Arc::clone(sess), - Arc::clone(config), - turn_context.clone(), - sub_id, - resolved, - ) - .await; - } - Err(err) => { - let event = Event { - id: sub_id, - msg: EventMsg::Error(ErrorEvent { - message: err.to_string(), - codex_error_info: Some(CodexErrorInfo::Other), - }), - }; - sess.send_event(&turn_context, event.msg).await; - } - } - } -} - /// Spawn a review thread using the given prompt. async fn spawn_review_thread( sess: Arc, diff --git a/codex-rs/core/src/codex/handlers.rs b/codex-rs/core/src/codex/handlers.rs new file mode 100644 index 0000000000..88f3f38b65 --- /dev/null +++ b/codex-rs/core/src/codex/handlers.rs @@ -0,0 +1,1230 @@ +use crate::realtime_conversation::handle_audio as handle_realtime_conversation_audio; +use crate::realtime_conversation::handle_close as handle_realtime_conversation_close; +use crate::realtime_conversation::handle_start as handle_realtime_conversation_start; +use crate::realtime_conversation::handle_text as handle_realtime_conversation_text; +use async_channel::Receiver; +use codex_otel::set_parent_from_w3c_trace_context; +use codex_protocol::protocol::Submission; +use tracing::Instrument; +use tracing::debug_span; +use tracing::info_span; + +use crate::codex::Session; +use crate::codex::SessionSettingsUpdate; +use crate::codex::SteerInputError; + +use crate::codex::spawn_review_thread; +use crate::config::Config; +use crate::config_loader::CloudRequirementsLoader; +use crate::config_loader::LoaderOverrides; +use crate::config_loader::load_config_layers_state; +use crate::realtime_context::REALTIME_TURN_TOKEN_BUDGET; +use crate::realtime_context::truncate_realtime_text_to_token_budget; +use crate::realtime_conversation::REALTIME_USER_TEXT_PREFIX; +use crate::realtime_conversation::prefix_realtime_v2_text; +use codex_features::Feature; +use codex_utils_absolute_path::AbsolutePathBuf; + +use crate::review_prompts::resolve_review_request; +use crate::rollout::RolloutRecorder; +use crate::rollout::read_session_meta_line; +use crate::tasks::CompactTask; +use crate::tasks::UndoTask; +use crate::tasks::UserShellCommandMode; +use crate::tasks::UserShellCommandTask; +use crate::tasks::execute_user_shell_command; +use codex_mcp::collect_mcp_snapshot_from_manager; +use codex_mcp::compute_auth_statuses; +use codex_protocol::protocol::CodexErrorInfo; +use codex_protocol::protocol::ErrorEvent; +use codex_protocol::protocol::Event; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::InterAgentCommunication; +use codex_protocol::protocol::ListSkillsResponseEvent; +use codex_protocol::protocol::McpServerRefreshConfig; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::RealtimeConversationListVoicesResponseEvent; +use codex_protocol::protocol::RealtimeVoicesList; +use codex_protocol::protocol::ReviewDecision; +use codex_protocol::protocol::ReviewRequest; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::SkillErrorInfo; +use codex_protocol::protocol::SkillsListEntry; +use codex_protocol::protocol::ThreadMemoryMode; +use codex_protocol::protocol::ThreadNameUpdatedEvent; +use codex_protocol::protocol::ThreadRolledBackEvent; +use codex_protocol::protocol::TurnAbortReason; +use codex_protocol::protocol::WarningEvent; +use codex_protocol::request_permissions::RequestPermissionsResponse; +use codex_protocol::request_user_input::RequestUserInputResponse; + +use crate::context_manager::is_user_turn_boundary; +use codex_protocol::config_types::CollaborationMode; +use codex_protocol::config_types::ModeKind; +use codex_protocol::config_types::Settings; +use codex_protocol::dynamic_tools::DynamicToolResponse; +use codex_protocol::items::UserMessageItem; +use codex_protocol::mcp::RequestId as ProtocolRequestId; +use codex_protocol::user_input::UserInput; +use codex_rmcp_client::ElicitationAction; +use codex_rmcp_client::ElicitationResponse; +use serde_json::Value; +use std::path::PathBuf; +use std::sync::Arc; +use tracing::debug; +use tracing::info; +use tracing::warn; + +pub async fn interrupt(sess: &Arc) { + sess.interrupt_task().await; +} + +pub async fn clean_background_terminals(sess: &Arc) { + sess.close_unified_exec_processes().await; +} + +pub async fn realtime_conversation_list_voices(sess: &Session, sub_id: String) { + sess.send_event_raw(Event { + id: sub_id, + msg: EventMsg::RealtimeConversationListVoicesResponse( + RealtimeConversationListVoicesResponseEvent { + voices: RealtimeVoicesList::builtin(), + }, + ), + }) + .await; +} + +pub async fn override_turn_context(sess: &Session, sub_id: String, updates: SessionSettingsUpdate) { + if let Err(err) = sess.update_settings(updates).await { + sess.send_event_raw(Event { + id: sub_id, + msg: EventMsg::Error(ErrorEvent { + message: err.to_string(), + codex_error_info: Some(CodexErrorInfo::BadRequest), + }), + }) + .await; + } +} + +pub async fn user_input_or_turn(sess: &Arc, sub_id: String, op: Op) { + user_input_or_turn_inner( + sess, + sub_id, + op, + /*mirror_user_text_to_realtime*/ Some(()), + ) + .await; +} + +pub(super) async fn user_input_or_turn_inner( + sess: &Arc, + sub_id: String, + op: Op, + mirror_user_text_to_realtime: Option<()>, +) { + let (items, updates, responsesapi_client_metadata) = match op { + Op::UserTurn { + cwd, + approval_policy, + approvals_reviewer, + sandbox_policy, + model, + effort, + summary, + service_tier, + final_output_json_schema, + items, + collaboration_mode, + personality, + } => { + let collaboration_mode = collaboration_mode.or_else(|| { + Some(CollaborationMode { + mode: ModeKind::Default, + settings: Settings { + model: model.clone(), + reasoning_effort: effort, + developer_instructions: None, + }, + }) + }); + ( + items, + SessionSettingsUpdate { + cwd: Some(cwd), + approval_policy: Some(approval_policy), + approvals_reviewer, + sandbox_policy: Some(sandbox_policy), + windows_sandbox_level: None, + collaboration_mode, + reasoning_summary: summary, + service_tier, + final_output_json_schema: Some(final_output_json_schema), + personality, + app_server_client_name: None, + app_server_client_version: None, + }, + None, + ) + } + Op::UserInput { + items, + final_output_json_schema, + responsesapi_client_metadata, + } => ( + items, + SessionSettingsUpdate { + final_output_json_schema: Some(final_output_json_schema), + ..Default::default() + }, + responsesapi_client_metadata, + ), + _ => unreachable!(), + }; + + let Ok(current_context) = sess.new_turn_with_sub_id(sub_id.clone(), updates).await else { + // new_turn_with_sub_id already emits the error event. + return; + }; + sess.maybe_emit_unknown_model_warning_for_turn(current_context.as_ref()) + .await; + let accepted_items = match sess + .steer_input( + items.clone(), + /*expected_turn_id*/ None, + responsesapi_client_metadata.clone(), + ) + .await + { + Ok(_) => { + current_context.session_telemetry.user_prompt(&items); + Some(items) + } + Err(SteerInputError::NoActiveTurn(items)) => { + if let Some(responsesapi_client_metadata) = responsesapi_client_metadata { + current_context + .turn_metadata_state + .set_responsesapi_client_metadata(responsesapi_client_metadata); + } + current_context.session_telemetry.user_prompt(&items); + sess.refresh_mcp_servers_if_requested(¤t_context) + .await; + let accepted_items = items.clone(); + sess.spawn_task( + Arc::clone(¤t_context), + items, + crate::tasks::RegularTask::new(), + ) + .await; + Some(accepted_items) + } + Err(err) => { + sess.send_event_raw(Event { + id: sub_id, + msg: EventMsg::Error(err.to_error_event()), + }) + .await; + None + } + }; + if let (Some(items), Some(())) = (accepted_items, mirror_user_text_to_realtime) { + self::mirror_user_text_to_realtime(sess, &items).await; + } +} + +async fn mirror_user_text_to_realtime(sess: &Arc, items: &[UserInput]) { + let text = UserMessageItem::new(items).message(); + if text.is_empty() { + return; + } + let text = if sess.conversation.is_running_v2().await { + prefix_realtime_v2_text(text, REALTIME_USER_TEXT_PREFIX) + } else { + text + }; + let text = truncate_realtime_text_to_token_budget(&text, REALTIME_TURN_TOKEN_BUDGET); + if text.is_empty() { + return; + } + if sess.conversation.running_state().await.is_none() { + return; + } + if let Err(err) = sess.conversation.text_in(text).await { + debug!("failed to mirror user text to realtime conversation: {err}"); + } +} + +/// Records an inter-agent assistant envelope, then lets the shared pending-work scheduler +/// decide whether an idle session should start a regular turn. +pub async fn inter_agent_communication( + sess: &Arc, + sub_id: String, + communication: InterAgentCommunication, +) { + let trigger_turn = communication.trigger_turn; + sess.enqueue_mailbox_communication(communication); + if trigger_turn { + sess.maybe_start_turn_for_pending_work_with_sub_id(sub_id) + .await; + } +} + +pub async fn run_user_shell_command(sess: &Arc, sub_id: String, command: String) { + if let Some((turn_context, cancellation_token)) = + sess.active_turn_context_and_cancellation_token().await + { + let session = Arc::clone(sess); + tokio::spawn(async move { + execute_user_shell_command( + session, + turn_context, + command, + cancellation_token, + UserShellCommandMode::ActiveTurnAuxiliary, + ) + .await; + }); + return; + } + + let turn_context = sess.new_default_turn_with_sub_id(sub_id).await; + sess.spawn_task( + Arc::clone(&turn_context), + Vec::new(), + UserShellCommandTask::new(command), + ) + .await; +} + +pub async fn resolve_elicitation( + sess: &Arc, + server_name: String, + request_id: ProtocolRequestId, + decision: codex_protocol::approvals::ElicitationAction, + content: Option, + meta: Option, +) { + let action = match decision { + codex_protocol::approvals::ElicitationAction::Accept => ElicitationAction::Accept, + codex_protocol::approvals::ElicitationAction::Decline => ElicitationAction::Decline, + codex_protocol::approvals::ElicitationAction::Cancel => ElicitationAction::Cancel, + }; + let content = match action { + // Preserve the legacy fallback for clients that only send an action. + ElicitationAction::Accept => Some(content.unwrap_or_else(|| serde_json::json!({}))), + ElicitationAction::Decline | ElicitationAction::Cancel => None, + }; + let response = ElicitationResponse { + action, + content, + meta, + }; + let request_id = match request_id { + ProtocolRequestId::String(value) => { + rmcp::model::NumberOrString::String(std::sync::Arc::from(value)) + } + ProtocolRequestId::Integer(value) => rmcp::model::NumberOrString::Number(value), + }; + if let Err(err) = sess + .resolve_elicitation(server_name, request_id, response) + .await + { + warn!( + error = %err, + "failed to resolve elicitation request in session" + ); + } +} + +/// Propagate a user's exec approval decision to the session. +/// Also optionally applies an execpolicy amendment. +pub async fn exec_approval( + sess: &Arc, + approval_id: String, + turn_id: Option, + decision: ReviewDecision, +) { + let event_turn_id = turn_id.unwrap_or_else(|| approval_id.clone()); + if let ReviewDecision::ApprovedExecpolicyAmendment { + proposed_execpolicy_amendment, + } = &decision + { + match sess + .persist_execpolicy_amendment(proposed_execpolicy_amendment) + .await + { + Ok(()) => { + sess.record_execpolicy_amendment_message( + &event_turn_id, + proposed_execpolicy_amendment, + ) + .await; + } + Err(err) => { + let message = format!("Failed to apply execpolicy amendment: {err}"); + tracing::warn!("{message}"); + let warning = EventMsg::Warning(WarningEvent { message }); + sess.send_event_raw(Event { + id: event_turn_id.clone(), + msg: warning, + }) + .await; + } + } + } + match decision { + ReviewDecision::Abort => { + sess.interrupt_task().await; + } + other => sess.notify_approval(&approval_id, other).await, + } +} + +pub async fn patch_approval(sess: &Arc, id: String, decision: ReviewDecision) { + match decision { + ReviewDecision::Abort => { + sess.interrupt_task().await; + } + other => sess.notify_approval(&id, other).await, + } +} + +pub async fn request_user_input_response( + sess: &Arc, + id: String, + response: RequestUserInputResponse, +) { + sess.notify_user_input_response(&id, response).await; +} + +pub async fn request_permissions_response( + sess: &Arc, + id: String, + response: RequestPermissionsResponse, +) { + sess.notify_request_permissions_response(&id, response) + .await; +} + +pub async fn dynamic_tool_response(sess: &Arc, id: String, response: DynamicToolResponse) { + sess.notify_dynamic_tool_response(&id, response).await; +} + +pub async fn add_to_history(sess: &Arc, config: &Arc, text: String) { + let id = sess.conversation_id; + let config = Arc::clone(config); + tokio::spawn(async move { + if let Err(e) = crate::message_history::append_entry(&text, &id, &config).await { + warn!("failed to append to message history: {e}"); + } + }); +} + +pub async fn get_history_entry_request( + sess: &Arc, + config: &Arc, + sub_id: String, + offset: usize, + log_id: u64, +) { + let config = Arc::clone(config); + let sess_clone = Arc::clone(sess); + + tokio::spawn(async move { + // Run lookup in blocking thread because it does file IO + locking. + let entry_opt = tokio::task::spawn_blocking(move || { + crate::message_history::lookup(log_id, offset, &config) + }) + .await + .unwrap_or(None); + + let event = Event { + id: sub_id, + msg: EventMsg::GetHistoryEntryResponse( + codex_protocol::protocol::GetHistoryEntryResponseEvent { + offset, + log_id, + entry: entry_opt.map(|e| codex_protocol::message_history::HistoryEntry { + conversation_id: e.session_id, + ts: e.ts, + text: e.text, + }), + }, + ), + }; + + sess_clone.send_event_raw(event).await; + }); +} + +pub async fn refresh_mcp_servers(sess: &Arc, refresh_config: McpServerRefreshConfig) { + let mut guard = sess.pending_mcp_server_refresh_config.lock().await; + *guard = Some(refresh_config); +} + +pub async fn reload_user_config(sess: &Arc) { + sess.reload_user_config_layer().await; +} + +pub async fn list_mcp_tools(sess: &Session, config: &Arc, sub_id: String) { + let mcp_connection_manager = sess.services.mcp_connection_manager.read().await; + let auth = sess.services.auth_manager.auth().await; + let mcp_servers = sess + .services + .mcp_manager + .effective_servers(config, auth.as_ref()) + .await; + let snapshot = collect_mcp_snapshot_from_manager( + &mcp_connection_manager, + compute_auth_statuses(mcp_servers.iter(), config.mcp_oauth_credentials_store_mode).await, + ) + .await; + let event = Event { + id: sub_id, + msg: EventMsg::McpListToolsResponse(snapshot), + }; + sess.send_event_raw(event).await; +} + +pub async fn list_skills(sess: &Session, sub_id: String, cwds: Vec, force_reload: bool) { + let default_cwd = { + let state = sess.state.lock().await; + state.session_configuration.cwd.to_path_buf() + }; + let cwds = if cwds.is_empty() { + vec![default_cwd] + } else { + cwds + }; + + let skills_manager = &sess.services.skills_manager; + let plugins_manager = &sess.services.plugins_manager; + let fs = sess + .services + .environment + .as_ref() + .map(|environment| environment.get_filesystem()); + let config = sess.get_config().await; + let codex_home = sess.codex_home().await; + let mut skills = Vec::new(); + let empty_cli_overrides: &[(String, toml::Value)] = &[]; + for cwd in cwds { + let cwd_abs = match AbsolutePathBuf::relative_to_current_dir(cwd.as_path()) { + Ok(path) => path, + Err(err) => { + let error_path = cwd.clone(); + skills.push(SkillsListEntry { + cwd, + skills: Vec::new(), + errors: vec![SkillErrorInfo { + path: error_path, + message: err.to_string(), + }], + }); + continue; + } + }; + let config_layer_stack = match load_config_layers_state( + &codex_home, + Some(cwd_abs.clone()), + empty_cli_overrides, + LoaderOverrides::default(), + CloudRequirementsLoader::default(), + ) + .await + { + Ok(config_layer_stack) => config_layer_stack, + Err(err) => { + let error_path = cwd.clone(); + skills.push(SkillsListEntry { + cwd, + skills: Vec::new(), + errors: vec![SkillErrorInfo { + path: error_path, + message: err.to_string(), + }], + }); + continue; + } + }; + let effective_skill_roots = plugins_manager + .effective_skill_roots_for_layer_stack( + &config_layer_stack, + config.features.enabled(Feature::Plugins), + ) + .await; + let skills_input = crate::SkillsLoadInput::new( + cwd_abs.clone(), + effective_skill_roots, + config_layer_stack, + config.bundled_skills_enabled(), + ); + let outcome = skills_manager + .skills_for_cwd(&skills_input, force_reload, fs.clone()) + .await; + let errors = super::errors_to_info(&outcome.errors); + let skills_metadata = super::skills_to_info(&outcome.skills, &outcome.disabled_paths); + skills.push(SkillsListEntry { + cwd, + skills: skills_metadata, + errors, + }); + } + + let event = Event { + id: sub_id, + msg: EventMsg::ListSkillsResponse(ListSkillsResponseEvent { skills }), + }; + sess.send_event_raw(event).await; +} + +pub async fn undo(sess: &Arc, sub_id: String) { + let turn_context = sess.new_default_turn_with_sub_id(sub_id).await; + sess.spawn_task(turn_context, Vec::new(), UndoTask::new()) + .await; +} + +pub async fn compact(sess: &Arc, sub_id: String) { + let turn_context = sess.new_default_turn_with_sub_id(sub_id).await; + + sess.spawn_task( + Arc::clone(&turn_context), + vec![UserInput::Text { + text: turn_context.compact_prompt().to_string(), + // Compaction prompt is synthesized; no UI element ranges to preserve. + text_elements: Vec::new(), + }], + CompactTask, + ) + .await; +} + +pub async fn drop_memories(sess: &Arc, config: &Arc, sub_id: String) { + let mut errors = Vec::new(); + + if let Some(state_db) = sess.services.state_db.as_deref() { + if let Err(err) = state_db.clear_memory_data().await { + errors.push(format!("failed clearing memory rows from state db: {err}")); + } + } else { + errors.push("state db unavailable; memory rows were not cleared".to_string()); + } + + if let Err(err) = crate::memories::clear_memory_roots_contents(&config.codex_home).await { + errors.push(format!( + "failed clearing memory directories under {}: {err}", + config.codex_home.display() + )); + } + + if errors.is_empty() { + let memory_root = crate::memories::memory_root(&config.codex_home); + sess.send_event_raw(Event { + id: sub_id, + msg: EventMsg::Warning(WarningEvent { + message: format!( + "Dropped memories at {} and cleared memory rows from state db.", + memory_root.display() + ), + }), + }) + .await; + return; + } + + sess.send_event_raw(Event { + id: sub_id, + msg: EventMsg::Error(ErrorEvent { + message: format!("Memory drop completed with errors: {}", errors.join("; ")), + codex_error_info: Some(CodexErrorInfo::Other), + }), + }) + .await; +} + +pub async fn update_memories(sess: &Arc, config: &Arc, sub_id: String) { + let session_source = { + let state = sess.state.lock().await; + state.session_configuration.session_source.clone() + }; + + crate::memories::start_memories_startup_task(sess, Arc::clone(config), &session_source); + + sess.send_event_raw(Event { + id: sub_id.clone(), + msg: EventMsg::Warning(WarningEvent { + message: "Memory update triggered.".to_string(), + }), + }) + .await; +} + +pub async fn thread_rollback(sess: &Arc, sub_id: String, num_turns: u32) { + if num_turns == 0 { + sess.send_event_raw(Event { + id: sub_id, + msg: EventMsg::Error(ErrorEvent { + message: "num_turns must be >= 1".to_string(), + codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed), + }), + }) + .await; + return; + } + + let has_active_turn = { sess.active_turn.lock().await.is_some() }; + if has_active_turn { + sess.send_event_raw(Event { + id: sub_id, + msg: EventMsg::Error(ErrorEvent { + message: "Cannot rollback while a turn is in progress.".to_string(), + codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed), + }), + }) + .await; + return; + } + + let turn_context = sess.new_default_turn_with_sub_id(sub_id).await; + let rollout_path = { + let recorder = { + let guard = sess.services.rollout.lock().await; + guard.clone() + }; + let Some(recorder) = recorder else { + sess.send_event_raw(Event { + id: turn_context.sub_id.clone(), + msg: EventMsg::Error(ErrorEvent { + message: "thread rollback requires a persisted rollout path".to_string(), + codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed), + }), + }) + .await; + return; + }; + recorder.rollout_path().to_path_buf() + }; + if let Some(recorder) = { + let guard = sess.services.rollout.lock().await; + guard.clone() + } && let Err(err) = recorder.flush().await + { + sess.send_event_raw(Event { + id: turn_context.sub_id.clone(), + msg: EventMsg::Error(ErrorEvent { + message: format!( + "failed to flush rollout `{}` for rollback replay: {err}", + rollout_path.display() + ), + codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed), + }), + }) + .await; + return; + } + + let initial_history = match RolloutRecorder::get_rollout_history(rollout_path.as_path()).await { + Ok(history) => history, + Err(err) => { + sess.send_event_raw(Event { + id: turn_context.sub_id.clone(), + msg: EventMsg::Error(ErrorEvent { + message: format!( + "failed to load rollout `{}` for rollback replay: {err}", + rollout_path.display() + ), + codex_error_info: Some(CodexErrorInfo::ThreadRollbackFailed), + }), + }) + .await; + return; + } + }; + + let rollback_event = ThreadRolledBackEvent { num_turns }; + let rollback_msg = EventMsg::ThreadRolledBack(rollback_event.clone()); + let replay_items = initial_history + .get_rollout_items() + .into_iter() + .chain(std::iter::once(RolloutItem::EventMsg(rollback_msg.clone()))) + .collect::>(); + sess.apply_rollout_reconstruction(turn_context.as_ref(), replay_items.as_slice()) + .await; + sess.recompute_token_usage(turn_context.as_ref()).await; + + sess.persist_rollout_items(&[RolloutItem::EventMsg(rollback_msg.clone())]) + .await; + if let Err(err) = sess.flush_rollout().await { + sess.send_event( + turn_context.as_ref(), + EventMsg::Warning(WarningEvent { + message: format!( + "Rolled the thread back, but failed to save the rollback marker. Codex will continue retrying. Error: {err}" + ), + }), + ) + .await; + } + + sess.deliver_event_raw(Event { + id: turn_context.sub_id.clone(), + msg: rollback_msg, + }) + .await; +} + +async fn persist_thread_name_update( + sess: &Arc, + event: ThreadNameUpdatedEvent, +) -> anyhow::Result { + let msg = EventMsg::ThreadNameUpdated(event); + let item = RolloutItem::EventMsg(msg.clone()); + let recorder = { + let guard = sess.services.rollout.lock().await; + guard.clone() + } + .ok_or_else(|| anyhow::anyhow!("Session persistence is disabled; cannot rename thread."))?; + recorder.persist().await?; + recorder.record_items(std::slice::from_ref(&item)).await?; + recorder.flush().await?; + Ok(msg) +} + +pub(super) async fn persist_thread_memory_mode_update( + sess: &Arc, + mode: ThreadMemoryMode, +) -> anyhow::Result<()> { + let recorder = { + let guard = sess.services.rollout.lock().await; + guard.clone() + } + .ok_or_else(|| { + anyhow::anyhow!("Session persistence is disabled; cannot update thread memory mode.") + })?; + recorder.persist().await?; + recorder.flush().await?; + + let rollout_path = recorder.rollout_path().to_path_buf(); + let mut session_meta = read_session_meta_line(rollout_path.as_path()).await?; + if session_meta.meta.id != sess.conversation_id { + anyhow::bail!( + "rollout session metadata id mismatch: expected {}, found {}", + sess.conversation_id, + session_meta.meta.id + ); + } + session_meta.meta.memory_mode = Some( + match mode { + ThreadMemoryMode::Enabled => "enabled", + ThreadMemoryMode::Disabled => "disabled", + } + .to_string(), + ); + + let item = RolloutItem::SessionMeta(session_meta); + recorder.record_items(std::slice::from_ref(&item)).await?; + recorder.flush().await?; + Ok(()) +} + +/// Persists the thread name in the rollout and state database, updates in-memory state, and +/// emits a `ThreadNameUpdated` event on success. +pub async fn set_thread_name(sess: &Arc, sub_id: String, name: String) { + let Some(name) = crate::util::normalize_thread_name(&name) else { + let event = Event { + id: sub_id, + msg: EventMsg::Error(ErrorEvent { + message: "Thread name cannot be empty.".to_string(), + codex_error_info: Some(CodexErrorInfo::BadRequest), + }), + }; + sess.send_event_raw(event).await; + return; + }; + + let updated = ThreadNameUpdatedEvent { + thread_id: sess.conversation_id, + thread_name: Some(name.clone()), + }; + + let msg = match persist_thread_name_update(sess, updated).await { + Ok(msg) => msg, + Err(err) => { + warn!("Failed to persist thread name update to rollout: {err}"); + let event = Event { + id: sub_id, + msg: EventMsg::Error(ErrorEvent { + message: err.to_string(), + codex_error_info: Some(CodexErrorInfo::Other), + }), + }; + sess.send_event_raw(event).await; + return; + } + }; + + if let Some(state_db) = sess.services.state_db.as_deref() + && let Err(err) = state_db + .update_thread_title(sess.conversation_id, &name) + .await + { + warn!("Failed to update thread title in state db: {err}"); + } + + { + let mut state = sess.state.lock().await; + state.session_configuration.thread_name = Some(name.clone()); + } + + let codex_home = sess.codex_home().await; + if let Err(err) = + crate::rollout::append_thread_name(&codex_home, sess.conversation_id, &name).await + { + warn!("Failed to update legacy thread name index: {err}"); + } + + sess.deliver_event_raw(Event { id: sub_id, msg }).await; +} + +/// Persists thread-level memory mode metadata for the active session. +/// +/// This does not involve the model and only affects whether the thread is +/// eligible for future memory generation. +pub async fn set_thread_memory_mode(sess: &Arc, sub_id: String, mode: ThreadMemoryMode) { + if let Err(err) = persist_thread_memory_mode_update(sess, mode).await { + warn!("Failed to persist thread memory mode update to rollout: {err}"); + let event = Event { + id: sub_id, + msg: EventMsg::Error(ErrorEvent { + message: err.to_string(), + codex_error_info: Some(CodexErrorInfo::Other), + }), + }; + sess.send_event_raw(event).await; + } +} + +pub async fn shutdown(sess: &Arc, sub_id: String) -> bool { + sess.abort_all_tasks(TurnAbortReason::Interrupted).await; + let _ = sess.conversation.shutdown().await; + sess.services + .unified_exec_manager + .terminate_all_processes() + .await; + sess.guardian_review_session.shutdown().await; + info!("Shutting down Codex instance"); + let history = sess.clone_history().await; + let turn_count = history + .raw_items() + .iter() + .filter(|item| is_user_turn_boundary(item)) + .count(); + sess.services.session_telemetry.counter( + "codex.conversation.turn.count", + i64::try_from(turn_count).unwrap_or(0), + &[], + ); + + // Gracefully flush and shutdown rollout recorder on session end so tests + // that inspect the rollout file do not race with the background writer. + let recorder_opt = { + let mut guard = sess.services.rollout.lock().await; + guard.take() + }; + if let Some(rec) = recorder_opt + && let Err(e) = rec.shutdown().await + { + warn!("failed to shutdown rollout recorder: {e}"); + let event = Event { + id: sub_id.clone(), + msg: EventMsg::Error(ErrorEvent { + message: "Failed to shutdown rollout recorder".to_string(), + codex_error_info: Some(CodexErrorInfo::Other), + }), + }; + sess.send_event_raw(event).await; + } + + let event = Event { + id: sub_id, + msg: EventMsg::ShutdownComplete, + }; + sess.send_event_raw(event).await; + true +} + +pub async fn review( + sess: &Arc, + config: &Arc, + sub_id: String, + review_request: ReviewRequest, +) { + let turn_context = sess.new_default_turn_with_sub_id(sub_id.clone()).await; + sess.maybe_emit_unknown_model_warning_for_turn(turn_context.as_ref()) + .await; + sess.refresh_mcp_servers_if_requested(&turn_context).await; + match resolve_review_request(review_request, &turn_context.cwd) { + Ok(resolved) => { + spawn_review_thread( + Arc::clone(sess), + Arc::clone(config), + turn_context.clone(), + sub_id, + resolved, + ) + .await; + } + Err(err) => { + let event = Event { + id: sub_id, + msg: EventMsg::Error(ErrorEvent { + message: err.to_string(), + codex_error_info: Some(CodexErrorInfo::Other), + }), + }; + sess.send_event(&turn_context, event.msg).await; + } + } +} + +pub(super) async fn submission_loop( + sess: Arc, + config: Arc, + rx_sub: Receiver, +) { + // To break out of this loop, send Op::Shutdown. + while let Ok(sub) = rx_sub.recv().await { + debug!(?sub, "Submission"); + let dispatch_span = submission_dispatch_span(&sub); + let should_exit = async { + match sub.op.clone() { + Op::Interrupt => { + interrupt(&sess).await; + false + } + Op::CleanBackgroundTerminals => { + clean_background_terminals(&sess).await; + false + } + Op::RealtimeConversationStart(params) => { + if let Err(err) = + handle_realtime_conversation_start(&sess, sub.id.clone(), params).await + { + sess.send_event_raw(Event { + id: sub.id.clone(), + msg: EventMsg::Error(ErrorEvent { + message: err.to_string(), + codex_error_info: Some(CodexErrorInfo::Other), + }), + }) + .await; + } + false + } + Op::RealtimeConversationAudio(params) => { + handle_realtime_conversation_audio(&sess, sub.id.clone(), params).await; + false + } + Op::RealtimeConversationText(params) => { + handle_realtime_conversation_text(&sess, sub.id.clone(), params).await; + false + } + Op::RealtimeConversationClose => { + handle_realtime_conversation_close(&sess, sub.id.clone()).await; + false + } + Op::RealtimeConversationListVoices => { + realtime_conversation_list_voices(&sess, sub.id.clone()).await; + false + } + Op::OverrideTurnContext { + cwd, + approval_policy, + approvals_reviewer, + sandbox_policy, + windows_sandbox_level, + model, + effort, + summary, + service_tier, + collaboration_mode, + personality, + } => { + let collaboration_mode = if let Some(collab_mode) = collaboration_mode { + collab_mode + } else { + let state = sess.state.lock().await; + state.session_configuration.collaboration_mode.with_updates( + model.clone(), + effort, + /*developer_instructions*/ None, + ) + }; + override_turn_context( + &sess, + sub.id.clone(), + SessionSettingsUpdate { + cwd, + approval_policy, + approvals_reviewer, + sandbox_policy, + windows_sandbox_level, + collaboration_mode: Some(collaboration_mode), + reasoning_summary: summary, + service_tier, + personality, + ..Default::default() + }, + ) + .await; + false + } + Op::UserInput { .. } | Op::UserTurn { .. } => { + user_input_or_turn(&sess, sub.id.clone(), sub.op).await; + false + } + Op::InterAgentCommunication { communication } => { + inter_agent_communication(&sess, sub.id.clone(), communication).await; + false + } + Op::ExecApproval { + id: approval_id, + turn_id, + decision, + } => { + exec_approval(&sess, approval_id, turn_id, decision).await; + false + } + Op::PatchApproval { id, decision } => { + patch_approval(&sess, id, decision).await; + false + } + Op::UserInputAnswer { id, response } => { + request_user_input_response(&sess, id, response).await; + false + } + Op::RequestPermissionsResponse { id, response } => { + request_permissions_response(&sess, id, response).await; + false + } + Op::DynamicToolResponse { id, response } => { + dynamic_tool_response(&sess, id, response).await; + false + } + Op::AddToHistory { text } => { + add_to_history(&sess, &config, text).await; + false + } + Op::GetHistoryEntryRequest { offset, log_id } => { + get_history_entry_request(&sess, &config, sub.id.clone(), offset, log_id).await; + false + } + Op::ListMcpTools => { + list_mcp_tools(&sess, &config, sub.id.clone()).await; + false + } + Op::RefreshMcpServers { config } => { + refresh_mcp_servers(&sess, config).await; + false + } + Op::ReloadUserConfig => { + reload_user_config(&sess).await; + false + } + Op::ListSkills { cwds, force_reload } => { + list_skills(&sess, sub.id.clone(), cwds, force_reload).await; + false + } + Op::Undo => { + undo(&sess, sub.id.clone()).await; + false + } + Op::Compact => { + compact(&sess, sub.id.clone()).await; + false + } + Op::DropMemories => { + drop_memories(&sess, &config, sub.id.clone()).await; + false + } + Op::UpdateMemories => { + update_memories(&sess, &config, sub.id.clone()).await; + false + } + Op::ThreadRollback { num_turns } => { + thread_rollback(&sess, sub.id.clone(), num_turns).await; + false + } + Op::SetThreadName { name } => { + set_thread_name(&sess, sub.id.clone(), name).await; + false + } + Op::SetThreadMemoryMode { mode } => { + set_thread_memory_mode(&sess, sub.id.clone(), mode).await; + false + } + Op::RunUserShellCommand { command } => { + run_user_shell_command(&sess, sub.id.clone(), command).await; + false + } + Op::ResolveElicitation { + server_name, + request_id, + decision, + content, + meta, + } => { + resolve_elicitation(&sess, server_name, request_id, decision, content, meta) + .await; + false + } + Op::Shutdown => shutdown(&sess, sub.id.clone()).await, + Op::Review { review_request } => { + review(&sess, &config, sub.id.clone(), review_request).await; + false + } + _ => false, // Ignore unknown ops; enum is non_exhaustive to allow extensions. + } + } + .instrument(dispatch_span) + .await; + if should_exit { + break; + } + } + // Also drain cached guardian state if the submission loop exits because + // the channel closed without receiving an explicit shutdown op. + sess.guardian_review_session.shutdown().await; + debug!("Agent loop exited"); +} + +pub(super) fn submission_dispatch_span(sub: &Submission) -> tracing::Span { + let op_name = sub.op.kind(); + let span_name = format!("op.dispatch.{op_name}"); + let dispatch_span = match &sub.op { + Op::RealtimeConversationAudio(_) => { + debug_span!( + "submission_dispatch", + otel.name = span_name.as_str(), + submission.id = sub.id.as_str(), + codex.op = op_name + ) + } + _ => info_span!( + "submission_dispatch", + otel.name = span_name.as_str(), + submission.id = sub.id.as_str(), + codex.op = op_name + ), + }; + if let Some(trace) = sub.trace.as_ref() + && !set_parent_from_w3c_trace_context(&dispatch_span, trace) + { + warn!( + submission.id = sub.id.as_str(), + "ignoring invalid submission trace carrier" + ); + } + dispatch_span +} From 55c3de75cba1a65088ff02b91527b94a1b60a69b Mon Sep 17 00:00:00 2001 From: Adrian <145513011+adrian-openai@users.noreply.github.com> Date: Thu, 16 Apr 2026 14:30:02 -0700 Subject: [PATCH 11/16] Register agent tasks behind use_agent_identity (#17387) ## Summary Stack PR3 for feature-gated agent identity support. This PR adds per-thread agent task registration behind `features.use_agent_identity`. Tasks are minted on the first real user turn and cached in thread runtime state for later turns. ## Stack - PR1: https://github.com/openai/codex/pull/17385 - add `features.use_agent_identity` - PR2: https://github.com/openai/codex/pull/17386 - register agent identities when enabled - PR3: https://github.com/openai/codex/pull/17387 - this PR, original task registration slice - PR3.1: https://github.com/openai/codex/pull/17978 - persist and prewarm registered tasks per thread - PR4: https://github.com/openai/codex/pull/17980 - use `AgentAssertion` downstream when enabled ## Validation Covered as part of the local stack validation pass: - `just fmt` - `cargo test -p codex-core --lib agent_identity` - `cargo test -p codex-core --lib agent_assertion` - `cargo test -p codex-core --lib websocket_agent_task` - `cargo test -p codex-api api_bridge` - `cargo build -p codex-cli --bin codex` ## Notes The full local app-server E2E path is still being debugged after PR creation. The current branch stack is directionally ready for review while that follow-up continues. --- MODULE.bazel.lock | 3 + codex-rs/Cargo.lock | 43 ++ codex-rs/Cargo.toml | 1 + codex-rs/app-server/tests/suite/auth.rs | 8 +- .../app-server/tests/suite/v2/plugin_list.rs | 3 +- .../tests/suite/v2/realtime_conversation.rs | 5 +- codex-rs/core/Cargo.toml | 2 + codex-rs/core/src/agent_identity.rs | 81 ++- .../src/agent_identity/task_registration.rs | 470 ++++++++++++++++++ codex-rs/core/src/codex.rs | 102 +++- codex-rs/core/src/codex_tests.rs | 13 +- codex-rs/core/src/state/session.rs | 15 + codex-rs/core/src/state/session_tests.rs | 38 ++ .../core/tests/suite/realtime_conversation.rs | 63 ++- 14 files changed, 804 insertions(+), 43 deletions(-) create mode 100644 codex-rs/core/src/agent_identity/task_registration.rs diff --git a/MODULE.bazel.lock b/MODULE.bazel.lock index ef80693c70..7a1acc3816 100644 --- a/MODULE.bazel.lock +++ b/MODULE.bazel.lock @@ -663,6 +663,7 @@ "bitflags_1.3.2": "{\"dependencies\":[{\"name\":\"compiler_builtins\",\"optional\":true,\"req\":\"^0.1.2\"},{\"name\":\"core\",\"optional\":true,\"package\":\"rustc-std-workspace-core\",\"req\":\"^1.0.0\"},{\"kind\":\"dev\",\"name\":\"rustversion\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde_derive\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"trybuild\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"walkdir\",\"req\":\"^2.3\"}],\"features\":{\"default\":[],\"example_generated\":[],\"rustc-dep-of-std\":[\"core\",\"compiler_builtins\"]}}", "bitflags_2.10.0": "{\"dependencies\":[{\"name\":\"arbitrary\",\"optional\":true,\"req\":\"^1.0\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"arbitrary\",\"req\":\"^1.0\"},{\"name\":\"bytemuck\",\"optional\":true,\"req\":\"^1.12\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"bytemuck\",\"req\":\"^1.12.2\"},{\"kind\":\"dev\",\"name\":\"rustversion\",\"req\":\"^1.0\"},{\"default_features\":false,\"name\":\"serde_core\",\"optional\":true,\"req\":\"^1.0.228\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde_lib\",\"package\":\"serde\",\"req\":\"^1.0.103\"},{\"kind\":\"dev\",\"name\":\"serde_test\",\"req\":\"^1.0.19\"},{\"kind\":\"dev\",\"name\":\"trybuild\",\"req\":\"^1.0.18\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"zerocopy\",\"req\":\"^0.8\"}],\"features\":{\"example_generated\":[],\"serde\":[\"serde_core\"],\"std\":[]}}", "bitflags_2.11.0": "{\"dependencies\":[{\"name\":\"arbitrary\",\"optional\":true,\"req\":\"^1.0\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"arbitrary\",\"req\":\"^1.0\"},{\"name\":\"bytemuck\",\"optional\":true,\"req\":\"^1.12\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"bytemuck\",\"req\":\"^1.12.2\"},{\"kind\":\"dev\",\"name\":\"rustversion\",\"req\":\"^1.0\"},{\"default_features\":false,\"name\":\"serde_core\",\"optional\":true,\"req\":\"^1.0.228\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde_lib\",\"package\":\"serde\",\"req\":\"^1.0.103\"},{\"kind\":\"dev\",\"name\":\"serde_test\",\"req\":\"^1.0.19\"},{\"kind\":\"dev\",\"name\":\"trybuild\",\"req\":\"^1.0.18\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"zerocopy\",\"req\":\"^0.8\"}],\"features\":{\"example_generated\":[],\"serde\":[\"serde_core\"],\"std\":[]}}", + "blake2_0.10.6": "{\"dependencies\":[{\"features\":[\"mac\"],\"name\":\"digest\",\"req\":\"^0.10.3\"},{\"features\":[\"dev\"],\"kind\":\"dev\",\"name\":\"digest\",\"req\":\"^0.10.3\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.2.2\"}],\"features\":{\"default\":[\"std\"],\"reset\":[],\"simd\":[],\"simd_asm\":[\"simd_opt\"],\"simd_opt\":[\"simd\"],\"size_opt\":[],\"std\":[\"digest/std\"]}}", "block-buffer_0.10.4": "{\"dependencies\":[{\"name\":\"generic-array\",\"req\":\"^0.14\"}],\"features\":{}}", "block-padding_0.3.3": "{\"dependencies\":[{\"name\":\"generic-array\",\"req\":\"^0.14\"}],\"features\":{\"std\":[]}}", "block2_0.6.2": "{\"dependencies\":[{\"default_features\":false,\"features\":[\"std\"],\"name\":\"objc2\",\"req\":\">=0.6.2, <0.8.0\"}],\"features\":{\"alloc\":[],\"compiler-rt\":[\"objc2/unstable-compiler-rt\"],\"default\":[\"std\"],\"gnustep-1-7\":[\"objc2/gnustep-1-7\"],\"gnustep-1-8\":[\"gnustep-1-7\",\"objc2/gnustep-1-8\"],\"gnustep-1-9\":[\"gnustep-1-8\",\"objc2/gnustep-1-9\"],\"gnustep-2-0\":[\"gnustep-1-9\",\"objc2/gnustep-2-0\"],\"gnustep-2-1\":[\"gnustep-2-0\",\"objc2/gnustep-2-1\"],\"std\":[\"alloc\"],\"unstable-coerce-pointee\":[],\"unstable-objfw\":[],\"unstable-private\":[],\"unstable-winobjc\":[\"gnustep-1-8\"]}}", @@ -752,6 +753,8 @@ "crossterm_winapi_0.9.1": "{\"dependencies\":[{\"features\":[\"winbase\",\"consoleapi\",\"processenv\",\"handleapi\",\"synchapi\",\"impl-default\"],\"name\":\"winapi\",\"req\":\"^0.3.8\",\"target\":\"cfg(windows)\"}],\"features\":{}}", "crunchy_0.2.4": "{\"dependencies\":[],\"features\":{\"default\":[\"limit_128\"],\"limit_1024\":[],\"limit_128\":[],\"limit_2048\":[],\"limit_256\":[],\"limit_512\":[],\"limit_64\":[],\"std\":[]}}", "crypto-common_0.1.7": "{\"dependencies\":[{\"features\":[\"more_lengths\"],\"name\":\"generic-array\",\"req\":\"=0.14.7\"},{\"name\":\"rand_core\",\"optional\":true,\"req\":\"^0.6\"},{\"name\":\"typenum\",\"req\":\"^1.14\"}],\"features\":{\"getrandom\":[\"rand_core/getrandom\"],\"std\":[]}}", + "crypto_box_0.9.1": "{\"dependencies\":[{\"default_features\":false,\"name\":\"aead\",\"req\":\"^0.5.2\"},{\"kind\":\"dev\",\"name\":\"bincode\",\"req\":\"^1\"},{\"default_features\":false,\"name\":\"blake2\",\"optional\":true,\"req\":\"^0.10\"},{\"name\":\"chacha20\",\"optional\":true,\"req\":\"^0.9\"},{\"default_features\":false,\"name\":\"crypto_secretbox\",\"req\":\"^0.1.1\"},{\"default_features\":false,\"features\":[\"zeroize\"],\"name\":\"curve25519-dalek\",\"req\":\"^4\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.4\"},{\"kind\":\"dev\",\"name\":\"rand\",\"req\":\"^0.8\"},{\"kind\":\"dev\",\"name\":\"rmp-serde\",\"req\":\"^1\"},{\"name\":\"salsa20\",\"optional\":true,\"req\":\"^0.10\"},{\"default_features\":false,\"name\":\"serdect\",\"optional\":true,\"req\":\"^0.2\"},{\"default_features\":false,\"name\":\"subtle\",\"req\":\"^2\"},{\"default_features\":false,\"name\":\"zeroize\",\"req\":\"^1\"}],\"features\":{\"alloc\":[\"aead/alloc\"],\"chacha20\":[\"dep:chacha20\",\"crypto_secretbox/chacha20\"],\"default\":[\"alloc\",\"getrandom\",\"salsa20\"],\"getrandom\":[\"aead/getrandom\",\"rand_core\"],\"heapless\":[\"aead/heapless\"],\"rand_core\":[\"aead/rand_core\"],\"salsa20\":[\"dep:salsa20\",\"crypto_secretbox/salsa20\"],\"seal\":[\"dep:blake2\",\"alloc\"],\"serde\":[\"dep:serdect\"],\"std\":[\"aead/std\"]}}", + "crypto_secretbox_0.1.1": "{\"dependencies\":[{\"default_features\":false,\"name\":\"aead\",\"req\":\"^0.5\"},{\"features\":[\"zeroize\"],\"name\":\"chacha20\",\"optional\":true,\"req\":\"^0.9\"},{\"default_features\":false,\"name\":\"cipher\",\"req\":\"^0.4\"},{\"default_features\":false,\"features\":[\"zeroize\"],\"name\":\"generic-array\",\"req\":\"^0.14.7\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.4\"},{\"name\":\"poly1305\",\"req\":\"^0.8\"},{\"features\":[\"zeroize\"],\"name\":\"salsa20\",\"optional\":true,\"req\":\"^0.10\"},{\"default_features\":false,\"name\":\"subtle\",\"req\":\"^2\"},{\"default_features\":false,\"name\":\"zeroize\",\"req\":\"^1\"}],\"features\":{\"alloc\":[\"aead/alloc\"],\"default\":[\"alloc\",\"getrandom\",\"salsa20\"],\"getrandom\":[\"aead/getrandom\",\"rand_core\"],\"heapless\":[\"aead/heapless\"],\"rand_core\":[\"aead/rand_core\"],\"std\":[\"aead/std\",\"alloc\"],\"stream\":[\"aead/stream\"]}}", "csv-core_0.1.13": "{\"dependencies\":[{\"default_features\":false,\"kind\":\"dev\",\"name\":\"arrayvec\",\"req\":\"^0.5\"},{\"default_features\":false,\"name\":\"memchr\",\"req\":\"^2\"}],\"features\":{\"default\":[],\"libc\":[\"memchr/libc\"]}}", "csv_1.4.0": "{\"dependencies\":[{\"default_features\":false,\"features\":[\"alloc\",\"serde\"],\"kind\":\"dev\",\"name\":\"bstr\",\"req\":\"^1.7.0\"},{\"name\":\"csv-core\",\"req\":\"^0.1.11\"},{\"name\":\"itoa\",\"req\":\"^1\"},{\"name\":\"ryu\",\"req\":\"^1\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1.0.221\"},{\"name\":\"serde_core\",\"req\":\"^1.0.221\"}],\"features\":{}}", "ctor-proc-macro_0.0.7": "{\"dependencies\":[],\"features\":{\"default\":[]}}", diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 296e51725c..bd74d83ad6 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -940,6 +940,15 @@ dependencies = [ "serde_core", ] +[[package]] +name = "blake2" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46502ad458c9a52b69d4d4d32775c788b7a1b85e8bc9d482d92250fc0e3f8efe" +dependencies = [ + "digest", +] + [[package]] name = "block-buffer" version = "0.10.4" @@ -1960,6 +1969,7 @@ dependencies = [ "codex-windows-sandbox", "core-foundation 0.9.4", "core_test_support", + "crypto_box", "csv", "ctor 0.6.3", "dirs", @@ -1990,6 +2000,7 @@ dependencies = [ "serde_json", "serial_test", "sha1", + "sha2", "shlex", "similar", "tempfile", @@ -3656,9 +3667,40 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" dependencies = [ "generic-array", + "rand_core 0.6.4", "typenum", ] +[[package]] +name = "crypto_box" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16182b4f39a82ec8a6851155cc4c0cda3065bb1db33651726a29e1951de0f009" +dependencies = [ + "aead", + "blake2", + "crypto_secretbox", + "curve25519-dalek", + "salsa20", + "subtle", + "zeroize", +] + +[[package]] +name = "crypto_secretbox" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9d6cf87adf719ddf43a805e92c6870a531aedda35ff640442cbaf8674e141e1" +dependencies = [ + "aead", + "cipher", + "generic-array", + "poly1305", + "salsa20", + "subtle", + "zeroize", +] + [[package]] name = "csv" version = "1.4.0" @@ -4948,6 +4990,7 @@ checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" dependencies = [ "typenum", "version_check", + "zeroize", ] [[package]] diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index deb3738812..c8f6b9e5a7 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -224,6 +224,7 @@ color-eyre = "0.6.3" constant_time_eq = "0.3.1" crossbeam-channel = "0.5.15" crossterm = "0.28.1" +crypto_box = { version = "0.9.1", features = ["seal"] } csv = "1.3.1" ctor = "0.6.3" deno_core_icudata = "0.77.0" diff --git a/codex-rs/app-server/tests/suite/auth.rs b/codex-rs/app-server/tests/suite/auth.rs index 1e60871012..b35a6fcba1 100644 --- a/codex-rs/app-server/tests/suite/auth.rs +++ b/codex-rs/app-server/tests/suite/auth.rs @@ -351,6 +351,8 @@ async fn get_auth_status_omits_token_after_proactive_refresh_failure() -> Result )?; let server = MockServer::start().await; + // App-server startup may proactively read stale auth before this test sends + // getAuthStatus; require the refresh path without depending on that race. Mock::given(method("POST")) .and(path("/oauth/token")) .respond_with(ResponseTemplate::new(401).set_body_json(serde_json::json!({ @@ -358,7 +360,7 @@ async fn get_auth_status_omits_token_after_proactive_refresh_failure() -> Result "code": "refresh_token_reused" } }))) - .expect(2) + .expect(1..=2) .mount(&server) .await; @@ -418,6 +420,8 @@ async fn get_auth_status_returns_token_after_proactive_refresh_recovery() -> Res )?; let server = MockServer::start().await; + // App-server startup may proactively read stale auth before this test sends + // getAuthStatus; require the refresh path without depending on that race. Mock::given(method("POST")) .and(path("/oauth/token")) .respond_with(ResponseTemplate::new(401).set_body_json(serde_json::json!({ @@ -425,7 +429,7 @@ async fn get_auth_status_returns_token_after_proactive_refresh_recovery() -> Res "code": "refresh_token_reused" } }))) - .expect(2) + .expect(1..=2) .mount(&server) .await; diff --git a/codex-rs/app-server/tests/suite/v2/plugin_list.rs b/codex-rs/app-server/tests/suite/v2/plugin_list.rs index 8bc4a8598e..5036101de4 100644 --- a/codex-rs/app-server/tests/suite/v2/plugin_list.rs +++ b/codex-rs/app-server/tests/suite/v2/plugin_list.rs @@ -30,7 +30,8 @@ use wiremock::matchers::method; use wiremock::matchers::path; use wiremock::matchers::query_param; -const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10); +// These tests start full app-server processes and can also run plugin startup warmers. +const DEFAULT_TIMEOUT: Duration = Duration::from_secs(60); const TEST_CURATED_PLUGIN_SHA: &str = "0123456789abcdef0123456789abcdef01234567"; const STARTUP_REMOTE_PLUGIN_SYNC_MARKER_FILE: &str = ".tmp/app-server-remote-plugin-sync-v1"; diff --git a/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs b/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs index 1f7b8a3e4b..11dc0ac4cf 100644 --- a/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs +++ b/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs @@ -72,6 +72,7 @@ use wiremock::matchers::path; use wiremock::matchers::path_regex; const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10); +const DELEGATED_SHELL_TOOL_TIMEOUT_MS: u64 = 30_000; const STARTUP_CONTEXT_HEADER: &str = "Startup context from Codex."; const V2_STEERING_ACKNOWLEDGEMENT: &str = "This was sent to steer the previous background agent task."; @@ -1717,7 +1718,9 @@ async fn webrtc_v2_tool_call_delegated_turn_can_execute_shell_tool() -> Result<( create_shell_command_sse_response( realtime_tool_ok_command(), /*workdir*/ None, - Some(5000), + // Windows CI can spend several seconds starting the nested PowerShell command. This + // test verifies delegated shell-tool plumbing, not timeout enforcement. + Some(DELEGATED_SHELL_TOOL_TIMEOUT_MS), "shell_call", )?, create_final_assistant_message_sse_response("shell tool finished")?, diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 4ce4192b81..a0d4259810 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -35,6 +35,7 @@ codex-connectors = { workspace = true } codex-config = { workspace = true } codex-core-plugins = { workspace = true } codex-core-skills = { workspace = true } +crypto_box = { workspace = true } codex-exec-server = { workspace = true } codex-features = { workspace = true } codex-feedback = { workspace = true } @@ -99,6 +100,7 @@ rmcp = { workspace = true, default-features = false, features = [ serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } sha1 = { workspace = true } +sha2 = { workspace = true } shlex = { workspace = true } similar = { workspace = true } tempfile = { workspace = true } diff --git a/codex-rs/core/src/agent_identity.rs b/codex-rs/core/src/agent_identity.rs index 2c897b5456..cfcab507bb 100644 --- a/codex-rs/core/src/agent_identity.rs +++ b/codex-rs/core/src/agent_identity.rs @@ -27,6 +27,10 @@ use tracing::debug; use tracing::info; use tracing::warn; +mod task_registration; + +pub(crate) use task_registration::RegisteredAgentTask; + use crate::config::Config; const AGENT_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15); @@ -119,32 +123,58 @@ impl AgentIdentityManager { return Ok(None); } - let Some(auth) = self.auth_manager.auth().await else { - debug!("skipping agent identity registration because no auth is available"); + let Some((auth, binding)) = self.current_auth_binding().await else { return Ok(None); }; - let Some(binding) = - AgentIdentityBinding::from_auth(&auth, self.auth_manager.forced_chatgpt_workspace_id()) - else { - debug!("skipping agent identity registration because ChatGPT auth is unavailable"); - return Ok(None); - }; + self.ensure_registered_identity_for_binding(&auth, &binding) + .await + .map(Some) + } + async fn ensure_registered_identity_for_binding( + &self, + auth: &CodexAuth, + binding: &AgentIdentityBinding, + ) -> Result { let _guard = self.ensure_lock.lock().await; - if let Some(stored_identity) = self.load_stored_identity(&auth, &binding)? { + if let Some(stored_identity) = self.load_stored_identity(auth, binding)? { info!( agent_runtime_id = %stored_identity.agent_runtime_id, binding_id = %binding.binding_id, "reusing stored agent identity" ); - return Ok(Some(stored_identity)); + return Ok(stored_identity); } - let stored_identity = self.register_agent_identity(&binding).await?; - self.store_identity(&auth, &stored_identity)?; - Ok(Some(stored_identity)) + let stored_identity = self.register_agent_identity(binding).await?; + self.store_identity(auth, &stored_identity)?; + Ok(stored_identity) + } + + pub(crate) async fn task_matches_current_binding(&self, task: &RegisteredAgentTask) -> bool { + if !self.feature_enabled { + return false; + } + + self.current_auth_binding() + .await + .is_some_and(|(_, binding)| task.matches_binding(&binding)) + } + + async fn current_auth_binding(&self) -> Option<(CodexAuth, AgentIdentityBinding)> { + let Some(auth) = self.auth_manager.auth().await else { + debug!("skipping agent identity flow because no auth is available"); + return None; + }; + + let binding = + AgentIdentityBinding::from_auth(&auth, self.auth_manager.forced_chatgpt_workspace_id()); + if binding.is_none() { + debug!("skipping agent identity flow because ChatGPT auth is unavailable"); + } + binding.map(|binding| (auth, binding)) } async fn register_agent_identity( @@ -351,12 +381,11 @@ impl StoredAgentIdentity { } fn matches_binding(&self, binding: &AgentIdentityBinding) -> bool { - self.binding_id == binding.binding_id - && self.chatgpt_account_id == binding.chatgpt_account_id - && match binding.chatgpt_user_id.as_deref() { - Some(chatgpt_user_id) => self.chatgpt_user_id.as_deref() == Some(chatgpt_user_id), - None => true, - } + binding.matches_parts( + &self.binding_id, + &self.chatgpt_account_id, + self.chatgpt_user_id.as_deref(), + ) } fn validate_key_material(&self) -> Result<()> { @@ -375,6 +404,20 @@ impl StoredAgentIdentity { } impl AgentIdentityBinding { + fn matches_parts( + &self, + binding_id: &str, + chatgpt_account_id: &str, + chatgpt_user_id: Option<&str>, + ) -> bool { + binding_id == self.binding_id + && chatgpt_account_id == self.chatgpt_account_id + && match self.chatgpt_user_id.as_deref() { + Some(expected_user_id) => chatgpt_user_id == Some(expected_user_id), + None => true, + } + } + fn from_auth(auth: &CodexAuth, forced_workspace_id: Option) -> Option { if !auth.is_chatgpt_auth() { return None; diff --git a/codex-rs/core/src/agent_identity/task_registration.rs b/codex-rs/core/src/agent_identity/task_registration.rs new file mode 100644 index 0000000000..4fc5d51282 --- /dev/null +++ b/codex-rs/core/src/agent_identity/task_registration.rs @@ -0,0 +1,470 @@ +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +use crypto_box::SecretKey as Curve25519SecretKey; +use ed25519_dalek::Signer as _; +use serde::Deserialize; +use serde::Serialize; +use sha2::Digest as _; +use sha2::Sha512; +use tracing::info; + +use super::*; + +const AGENT_TASK_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15); + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct RegisteredAgentTask { + pub(crate) binding_id: String, + pub(crate) chatgpt_account_id: String, + pub(crate) chatgpt_user_id: Option, + pub(crate) agent_runtime_id: String, + pub(crate) task_id: String, + pub(crate) registered_at: String, +} + +#[derive(Debug, Serialize)] +struct RegisterTaskRequest { + signature: String, + timestamp: String, +} + +#[derive(Debug, Deserialize)] +struct RegisterTaskResponse { + encrypted_task_id: String, +} + +impl AgentIdentityManager { + pub(crate) async fn register_task(&self) -> Result> { + if !self.feature_enabled { + return Ok(None); + } + + let Some((auth, binding)) = self.current_auth_binding().await else { + return Ok(None); + }; + + self.register_task_for_binding(auth, binding).await + } + + async fn register_task_for_binding( + &self, + auth: CodexAuth, + binding: AgentIdentityBinding, + ) -> Result> { + let stored_identity = self + .ensure_registered_identity_for_binding(&auth, &binding) + .await?; + + let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); + let request_body = RegisterTaskRequest { + signature: sign_task_registration_payload(&stored_identity, ×tamp)?, + timestamp, + }; + + let client = create_client(); + let url = + agent_task_registration_url(&self.chatgpt_base_url, &stored_identity.agent_runtime_id); + let human_biscuit = self.mint_human_biscuit(&binding, "POST", &url).await?; + let response = client + .post(&url) + .header("X-OpenAI-Authorization", human_biscuit) + .json(&request_body) + .timeout(AGENT_TASK_REGISTRATION_TIMEOUT) + .send() + .await + .with_context(|| format!("failed to send agent task registration request to {url}"))?; + + if response.status().is_success() { + let response_body = response + .json::() + .await + .with_context(|| format!("failed to parse agent task response from {url}"))?; + let registered_task = RegisteredAgentTask { + binding_id: stored_identity.binding_id.clone(), + chatgpt_account_id: stored_identity.chatgpt_account_id.clone(), + chatgpt_user_id: stored_identity.chatgpt_user_id.clone(), + agent_runtime_id: stored_identity.agent_runtime_id.clone(), + task_id: decrypt_task_id_response( + &stored_identity, + &response_body.encrypted_task_id, + )?, + registered_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + }; + info!( + agent_runtime_id = %registered_task.agent_runtime_id, + task_id = %registered_task.task_id, + "registered agent task" + ); + return Ok(Some(registered_task)); + } + + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + anyhow::bail!("agent task registration failed with status {status} from {url}: {body}") + } +} + +impl RegisteredAgentTask { + pub(super) fn matches_binding(&self, binding: &AgentIdentityBinding) -> bool { + binding.matches_parts( + &self.binding_id, + &self.chatgpt_account_id, + self.chatgpt_user_id.as_deref(), + ) + } + + pub(crate) fn has_same_binding(&self, other: &Self) -> bool { + self.binding_id == other.binding_id + && self.chatgpt_account_id == other.chatgpt_account_id + && self.chatgpt_user_id == other.chatgpt_user_id + } +} + +fn sign_task_registration_payload( + stored_identity: &StoredAgentIdentity, + timestamp: &str, +) -> Result { + let signing_key = stored_identity.signing_key()?; + let payload = format!("{}:{timestamp}", stored_identity.agent_runtime_id); + Ok(BASE64_STANDARD.encode(signing_key.sign(payload.as_bytes()).to_bytes())) +} + +fn decrypt_task_id_response( + stored_identity: &StoredAgentIdentity, + encrypted_task_id: &str, +) -> Result { + let signing_key = stored_identity.signing_key()?; + let ciphertext = BASE64_STANDARD + .decode(encrypted_task_id) + .context("encrypted task id is not valid base64")?; + let plaintext = curve25519_secret_key_from_signing_key(&signing_key) + .unseal(&ciphertext) + .map_err(|_| anyhow::anyhow!("failed to decrypt encrypted task id"))?; + String::from_utf8(plaintext).context("decrypted task id is not valid UTF-8") +} + +fn curve25519_secret_key_from_signing_key(signing_key: &SigningKey) -> Curve25519SecretKey { + let digest = Sha512::digest(signing_key.to_bytes()); + let mut secret_key = [0u8; 32]; + secret_key.copy_from_slice(&digest[..32]); + secret_key[0] &= 248; + secret_key[31] &= 127; + secret_key[31] |= 64; + Curve25519SecretKey::from(secret_key) +} + +fn agent_task_registration_url(chatgpt_base_url: &str, agent_runtime_id: &str) -> String { + let trimmed = chatgpt_base_url.trim_end_matches('/'); + format!("{trimmed}/v1/agent/{agent_runtime_id}/task/register") +} + +#[cfg(test)] +mod tests { + use base64::engine::general_purpose::URL_SAFE_NO_PAD; + use codex_app_server_protocol::AuthMode as ApiAuthMode; + use codex_login::AuthCredentialsStoreMode; + use codex_login::AuthDotJson; + use codex_login::save_auth; + use codex_login::token_data::IdTokenInfo; + use codex_login::token_data::TokenData; + use pretty_assertions::assert_eq; + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::header; + use wiremock::matchers::method; + use wiremock::matchers::path; + + use super::*; + + #[tokio::test] + async fn register_task_skips_when_feature_is_disabled() { + let auth_manager = + AuthManager::from_auth_for_testing(make_chatgpt_auth("account-123", Some("user-123"))); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ false, + "https://chatgpt.com/backend-api/".to_string(), + SessionSource::Cli, + ); + + assert_eq!(manager.register_task().await.unwrap(), None); + } + + #[tokio::test] + async fn register_task_skips_for_api_key_auth() { + let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test-key")); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + "https://chatgpt.com/backend-api/".to_string(), + SessionSource::Cli, + ); + + assert_eq!(manager.register_task().await.unwrap(), None); + } + + #[tokio::test] + async fn register_task_registers_and_decrypts_plaintext_task_id() { + let server = MockServer::start().await; + let chatgpt_base_url = server.uri(); + mount_human_biscuit(&server, &chatgpt_base_url, "agent-123").await; + let auth = make_chatgpt_auth("account-123", Some("user-123")); + let auth_manager = AuthManager::from_auth_for_testing(auth.clone()); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + chatgpt_base_url, + SessionSource::Cli, + ); + let stored_identity = seed_stored_identity(&manager, &auth, "agent-123", "account-123"); + let encrypted_task_id = + encrypt_task_id_for_identity(&stored_identity, "task_123").expect("task ciphertext"); + + Mock::given(method("POST")) + .and(path("/v1/agent/agent-123/task/register")) + .and(header("x-openai-authorization", "human-biscuit")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "encrypted_task_id": encrypted_task_id, + }))) + .expect(1) + .mount(&server) + .await; + + let task = manager + .register_task() + .await + .unwrap() + .expect("task should be registered"); + + assert_eq!( + task, + RegisteredAgentTask { + binding_id: "chatgpt-account-account-123".to_string(), + chatgpt_account_id: "account-123".to_string(), + chatgpt_user_id: Some("user-123".to_string()), + agent_runtime_id: "agent-123".to_string(), + task_id: "task_123".to_string(), + registered_at: task.registered_at.clone(), + } + ); + } + + #[tokio::test] + async fn register_task_uses_chatgpt_base_url() { + let server = MockServer::start().await; + let chatgpt_base_url = format!("{}/backend-api", server.uri()); + mount_human_biscuit(&server, &chatgpt_base_url, "agent-fallback").await; + let auth = make_chatgpt_auth("account-123", Some("user-123")); + let auth_manager = AuthManager::from_auth_for_testing(auth.clone()); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + chatgpt_base_url, + SessionSource::Cli, + ); + let stored_identity = + seed_stored_identity(&manager, &auth, "agent-fallback", "account-123"); + let encrypted_task_id = encrypt_task_id_for_identity(&stored_identity, "task_fallback") + .expect("task ciphertext"); + + Mock::given(method("POST")) + .and(path("/backend-api/v1/agent/agent-fallback/task/register")) + .and(header("x-openai-authorization", "human-biscuit")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "encrypted_task_id": encrypted_task_id, + }))) + .expect(1) + .mount(&server) + .await; + + let task = manager + .register_task() + .await + .unwrap() + .expect("task should be registered"); + + assert_eq!(task.agent_runtime_id, "agent-fallback"); + assert_eq!(task.task_id, "task_fallback"); + } + + #[tokio::test] + async fn register_task_for_binding_keeps_one_auth_snapshot() { + let server = MockServer::start().await; + let chatgpt_base_url = server.uri(); + mount_human_biscuit(&server, &chatgpt_base_url, "agent-123").await; + let binding_auth = make_chatgpt_auth("account-123", Some("user-123")); + let auth_manager = + AuthManager::from_auth_for_testing(make_chatgpt_auth("account-456", Some("user-456"))); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + chatgpt_base_url, + SessionSource::Cli, + ); + let stored_identity = + seed_stored_identity(&manager, &binding_auth, "agent-123", "account-123"); + let encrypted_task_id = + encrypt_task_id_for_identity(&stored_identity, "task_123").expect("task ciphertext"); + let binding = + AgentIdentityBinding::from_auth(&binding_auth, /*forced_workspace_id*/ None) + .expect("binding"); + + Mock::given(method("POST")) + .and(path("/v1/agent/agent-123/task/register")) + .and(header("x-openai-authorization", "human-biscuit")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "encrypted_task_id": encrypted_task_id, + }))) + .expect(1) + .mount(&server) + .await; + + let task = manager + .register_task_for_binding(binding_auth, binding) + .await + .unwrap() + .expect("task should be registered"); + + assert_eq!( + task, + RegisteredAgentTask { + binding_id: "chatgpt-account-account-123".to_string(), + chatgpt_account_id: "account-123".to_string(), + chatgpt_user_id: Some("user-123".to_string()), + agent_runtime_id: "agent-123".to_string(), + task_id: "task_123".to_string(), + registered_at: task.registered_at.clone(), + } + ); + } + + #[tokio::test] + async fn task_matches_current_binding_rejects_stale_auth_binding() { + let auth_manager = + AuthManager::from_auth_for_testing(make_chatgpt_auth("account-456", Some("user-456"))); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + "https://chatgpt.com/backend-api/".to_string(), + SessionSource::Cli, + ); + let task = RegisteredAgentTask { + binding_id: "chatgpt-account-account-123".to_string(), + chatgpt_account_id: "account-123".to_string(), + chatgpt_user_id: Some("user-123".to_string()), + agent_runtime_id: "agent-123".to_string(), + task_id: "task_123".to_string(), + registered_at: "2026-03-23T12:00:00Z".to_string(), + }; + + assert!(!manager.task_matches_current_binding(&task).await); + } + + async fn mount_human_biscuit( + server: &MockServer, + chatgpt_base_url: &str, + agent_runtime_id: &str, + ) { + let biscuit_url = agent_identity_biscuit_url(chatgpt_base_url); + let biscuit_path = reqwest::Url::parse(&biscuit_url) + .expect("biscuit URL parses") + .path() + .to_string(); + let target_url = agent_task_registration_url(chatgpt_base_url, agent_runtime_id); + Mock::given(method("GET")) + .and(path(biscuit_path)) + .and(header("authorization", "Bearer access-token-account-123")) + .and(header("x-original-method", "POST")) + .and(header("x-original-url", target_url)) + .respond_with( + ResponseTemplate::new(200).insert_header("x-openai-authorization", "human-biscuit"), + ) + .expect(1) + .mount(server) + .await; + } + + fn seed_stored_identity( + manager: &AgentIdentityManager, + auth: &CodexAuth, + agent_runtime_id: &str, + account_id: &str, + ) -> StoredAgentIdentity { + let key_material = generate_agent_key_material().expect("key material"); + let binding = + AgentIdentityBinding::from_auth(auth, /*forced_workspace_id*/ None).expect("binding"); + let stored_identity = StoredAgentIdentity { + binding_id: binding.binding_id, + chatgpt_account_id: account_id.to_string(), + chatgpt_user_id: Some("user-123".to_string()), + agent_runtime_id: agent_runtime_id.to_string(), + private_key_pkcs8_base64: key_material.private_key_pkcs8_base64, + public_key_ssh: key_material.public_key_ssh, + registered_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + abom: manager.abom.clone(), + }; + manager + .store_identity(auth, &stored_identity) + .expect("store identity"); + let persisted = auth + .get_agent_identity(account_id) + .expect("persisted identity"); + assert_eq!(persisted.agent_runtime_id, agent_runtime_id); + stored_identity + } + + fn encrypt_task_id_for_identity( + stored_identity: &StoredAgentIdentity, + task_id: &str, + ) -> Result { + let mut rng = crypto_box::aead::OsRng; + let public_key = + curve25519_secret_key_from_signing_key(&stored_identity.signing_key()?).public_key(); + let ciphertext = public_key + .seal(&mut rng, task_id.as_bytes()) + .map_err(|_| anyhow::anyhow!("failed to encrypt test task id"))?; + Ok(BASE64_STANDARD.encode(ciphertext)) + } + + fn make_chatgpt_auth(account_id: &str, user_id: Option<&str>) -> CodexAuth { + let tempdir = tempfile::tempdir().expect("tempdir"); + let auth_json = AuthDotJson { + auth_mode: Some(ApiAuthMode::Chatgpt), + openai_api_key: None, + tokens: Some(TokenData { + id_token: IdTokenInfo { + email: None, + chatgpt_plan_type: None, + chatgpt_user_id: user_id.map(ToOwned::to_owned), + chatgpt_account_id: Some(account_id.to_string()), + chatgpt_account_is_fedramp: false, + raw_jwt: fake_id_token(account_id, user_id), + }, + access_token: format!("access-token-{account_id}"), + refresh_token: "refresh-token".to_string(), + account_id: Some(account_id.to_string()), + }), + last_refresh: Some(Utc::now()), + agent_identity: None, + }; + save_auth(tempdir.path(), &auth_json, AuthCredentialsStoreMode::File).expect("save auth"); + CodexAuth::from_auth_storage(tempdir.path(), AuthCredentialsStoreMode::File) + .expect("load auth") + .expect("auth") + } + + fn fake_id_token(account_id: &str, user_id: Option<&str>) -> String { + let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#); + let payload = serde_json::json!({ + "https://api.openai.com/auth": { + "chatgpt_user_id": user_id, + "chatgpt_account_id": account_id, + } + }); + let payload = URL_SAFE_NO_PAD.encode(payload.to_string()); + format!("{header}.{payload}.signature") + } +} diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 4e09542d88..f356b3fff7 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -14,6 +14,7 @@ use crate::agent::MailboxReceiver; use crate::agent::agent_status_from_event; use crate::agent::status::is_final; use crate::agent_identity::AgentIdentityManager; +use crate::agent_identity::RegisteredAgentTask; use crate::apps::render_apps_section; use crate::commit_attribution::commit_message_trailer_instruction; use crate::compact; @@ -1577,7 +1578,7 @@ impl Session { async fn fail_agent_identity_registration(self: &Arc, error: anyhow::Error) { warn!(error = %error, "agent identity registration failed"); let message = format!( - "Agent identity registration failed. Codex cannot continue while `features.use_agent_identity` is enabled: {error}" + "Agent identity registration failed while `features.use_agent_identity` is enabled: {error}" ); self.send_event_raw(Event { id: self.next_internal_sub_id(), @@ -1587,7 +1588,90 @@ impl Session { }), }) .await; - handlers::shutdown(self, self.next_internal_sub_id()).await; + } + + async fn cached_agent_task_for_current_binding(&self) -> Option { + let agent_task = { + let state = self.state.lock().await; + state.agent_task() + }?; + + if self + .services + .agent_identity_manager + .task_matches_current_binding(&agent_task) + .await + { + debug!( + agent_runtime_id = %agent_task.agent_runtime_id, + task_id = %agent_task.task_id, + "reusing cached agent task" + ); + return Some(agent_task); + } + + debug!( + agent_runtime_id = %agent_task.agent_runtime_id, + task_id = %agent_task.task_id, + "discarding cached agent task because auth binding changed" + ); + let mut state = self.state.lock().await; + if state.agent_task().as_ref() == Some(&agent_task) { + state.clear_agent_task(); + } + None + } + + async fn ensure_agent_task_registered(&self) -> anyhow::Result> { + if let Some(agent_task) = self.cached_agent_task_for_current_binding().await { + return Ok(Some(agent_task)); + } + + for _ in 0..2 { + let Some(agent_task) = self.services.agent_identity_manager.register_task().await? + else { + return Ok(None); + }; + + if !self + .services + .agent_identity_manager + .task_matches_current_binding(&agent_task) + .await + { + debug!( + agent_runtime_id = %agent_task.agent_runtime_id, + task_id = %agent_task.task_id, + "discarding newly registered agent task because auth binding changed" + ); + continue; + } + + { + let mut state = self.state.lock().await; + if let Some(existing_agent_task) = state.agent_task() { + if existing_agent_task.has_same_binding(&agent_task) { + return Ok(Some(existing_agent_task)); + } + debug!( + agent_runtime_id = %existing_agent_task.agent_runtime_id, + task_id = %existing_agent_task.task_id, + "replacing cached agent task because auth binding changed" + ); + } + state.set_agent_task(agent_task.clone()); + } + + info!( + thread_id = %self.conversation_id, + agent_runtime_id = %agent_task.agent_runtime_id, + task_id = %agent_task.task_id, + "registered agent task for thread" + ); + return Ok(Some(agent_task)); + } + + Ok(None) } #[allow(clippy::too_many_arguments)] @@ -5168,6 +5252,20 @@ pub(crate) async fn run_turn( })) .await; } + if let Err(error) = sess.ensure_agent_task_registered().await { + warn!(error = %error, "agent task registration failed"); + sess.send_event( + turn_context.as_ref(), + EventMsg::Error(ErrorEvent { + message: format!( + "Agent task registration failed. Please try again; Codex will attempt to register the task again on the next turn: {error}" + ), + codex_error_info: Some(CodexErrorInfo::Other), + }), + ) + .await; + return None; + } if !skill_items.is_empty() { sess.record_conversation_items(&turn_context, &skill_items) diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 81ebd6c9f8..aa5583405e 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -4182,7 +4182,7 @@ pub(crate) async fn make_session_and_context_with_rx() -> ( } #[tokio::test] -async fn fail_agent_identity_registration_emits_error_and_shutdown() { +async fn fail_agent_identity_registration_emits_error_without_shutdown() { let (session, _turn_context, rx_event) = make_session_and_context_with_rx().await; session @@ -4200,21 +4200,14 @@ async fn fail_agent_identity_registration_emits_error_and_shutdown() { }) => { assert_eq!( message, - "Agent identity registration failed. Codex cannot continue while `features.use_agent_identity` is enabled: registration exploded".to_string() + "Agent identity registration failed while `features.use_agent_identity` is enabled: registration exploded".to_string() ); assert_eq!(codex_error_info, Some(CodexErrorInfo::Other)); } other => panic!("expected error event, got {other:?}"), } - let shutdown_event = timeout(Duration::from_secs(1), rx_event.recv()) - .await - .expect("shutdown event should arrive") - .expect("shutdown event should be readable"); - match shutdown_event.msg { - EventMsg::ShutdownComplete => {} - other => panic!("expected shutdown event, got {other:?}"), - } + assert!(rx_event.try_recv().is_err()); } #[tokio::test] diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 4360b16de4..42ed15e87a 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -6,6 +6,7 @@ use codex_sandboxing::policy_transforms::merge_permission_profiles; use std::collections::HashMap; use std::collections::HashSet; +use crate::agent_identity::RegisteredAgentTask; use crate::codex::PreviousTurnSettings; use crate::codex::SessionConfiguration; use crate::context_manager::ContextManager; @@ -30,6 +31,7 @@ pub(crate) struct SessionState { previous_turn_settings: Option, /// Startup prewarmed session prepared during session initialization. pub(crate) startup_prewarm: Option, + pub(crate) agent_task: Option, pub(crate) active_connector_selection: HashSet, pub(crate) pending_session_start_source: Option, granted_permissions: Option, @@ -49,6 +51,7 @@ impl SessionState { mcp_dependency_prompted: HashSet::new(), previous_turn_settings: None, startup_prewarm: None, + agent_task: None, active_connector_selection: HashSet::new(), pending_session_start_source: None, granted_permissions: None, @@ -186,6 +189,18 @@ impl SessionState { self.startup_prewarm.take() } + pub(crate) fn agent_task(&self) -> Option { + self.agent_task.clone() + } + + pub(crate) fn set_agent_task(&mut self, agent_task: RegisteredAgentTask) { + self.agent_task = Some(agent_task); + } + + pub(crate) fn clear_agent_task(&mut self) { + self.agent_task = None; + } + // Adds connector IDs to the active set and returns the merged selection. pub(crate) fn merge_connector_selection(&mut self, connector_ids: I) -> HashSet where diff --git a/codex-rs/core/src/state/session_tests.rs b/codex-rs/core/src/state/session_tests.rs index 1af7ccc8f6..6816c8731d 100644 --- a/codex-rs/core/src/state/session_tests.rs +++ b/codex-rs/core/src/state/session_tests.rs @@ -1,4 +1,5 @@ use super::*; +use crate::agent_identity::RegisteredAgentTask; use crate::codex::make_session_configuration_for_tests; use codex_protocol::protocol::CreditsSnapshot; use codex_protocol::protocol::RateLimitWindow; @@ -33,6 +34,43 @@ async fn clear_connector_selection_removes_entries() { assert_eq!(state.get_connector_selection(), HashSet::new()); } +#[tokio::test] +async fn set_agent_task_persists_plaintext_task_for_session_reuse() { + let session_configuration = make_session_configuration_for_tests().await; + let mut state = SessionState::new(session_configuration); + let agent_task = RegisteredAgentTask { + binding_id: "chatgpt-account-account-123".to_string(), + chatgpt_account_id: "account-123".to_string(), + chatgpt_user_id: Some("user-123".to_string()), + agent_runtime_id: "agent_123".to_string(), + task_id: "task_123".to_string(), + registered_at: "2026-03-23T12:00:00Z".to_string(), + }; + + state.set_agent_task(agent_task.clone()); + + assert_eq!(state.agent_task(), Some(agent_task)); +} + +#[tokio::test] +async fn clear_agent_task_removes_cached_task() { + let session_configuration = make_session_configuration_for_tests().await; + let mut state = SessionState::new(session_configuration); + let agent_task = RegisteredAgentTask { + binding_id: "chatgpt-account-account-123".to_string(), + chatgpt_account_id: "account-123".to_string(), + chatgpt_user_id: Some("user-123".to_string()), + agent_runtime_id: "agent_123".to_string(), + task_id: "task_123".to_string(), + registered_at: "2026-03-23T12:00:00Z".to_string(), + }; + + state.set_agent_task(agent_task); + state.clear_agent_task(); + + assert_eq!(state.agent_task(), None); +} + #[tokio::test] async fn set_rate_limits_defaults_limit_id_to_codex_when_missing() { let session_configuration = make_session_configuration_for_tests().await; diff --git a/codex-rs/core/tests/suite/realtime_conversation.rs b/codex-rs/core/tests/suite/realtime_conversation.rs index 1cd97391bb..35ec9668a3 100644 --- a/codex-rs/core/tests/suite/realtime_conversation.rs +++ b/codex-rs/core/tests/suite/realtime_conversation.rs @@ -15,6 +15,7 @@ use codex_protocol::protocol::ConversationStartTransport; use codex_protocol::protocol::ConversationTextParams; use codex_protocol::protocol::ErrorEvent; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::GitInfo; use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::Op; use codex_protocol::protocol::RealtimeAudioFrame; @@ -24,7 +25,11 @@ use codex_protocol::protocol::RealtimeEvent; use codex_protocol::protocol::RealtimeOutputModality; use codex_protocol::protocol::RealtimeVoice; use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::RolloutLine; +use codex_protocol::protocol::SessionMeta; +use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::UserMessageEvent; use codex_protocol::user_input::UserInput; use codex_utils_output_truncation::approx_token_count; use core_test_support::responses; @@ -65,6 +70,7 @@ const MEMORY_PROMPT_PHRASE: &str = "You have access to a memory folder with guidance from prior runs."; const REALTIME_CONVERSATION_TEST_SUBPROCESS_ENV_VAR: &str = "CODEX_REALTIME_CONVERSATION_TEST_SUBPROCESS"; +const WEBSOCKET_REQUEST_TIMEOUT: Duration = Duration::from_secs(30); #[derive(Debug, Clone)] struct RealtimeCallRequestCapture { @@ -141,7 +147,7 @@ async fn wait_for_matching_websocket_request( where F: Fn(&core_test_support::responses::WebSocketRequest) -> bool, { - let deadline = tokio::time::Instant::now() + Duration::from_secs(10); + let deadline = tokio::time::Instant::now() + WEBSOCKET_REQUEST_TIMEOUT; loop { if let Some(request) = server .connections() @@ -196,16 +202,18 @@ async fn seed_recent_thread( let db = test.codex.state_db().context("state db enabled")?; let thread_id = ThreadId::new(); let updated_at = Utc::now(); - let rollout_path = test + let rollout_dir = test .codex_home_path() - .join(format!("rollout-{thread_id}.jsonl")); - // This helper seeds SQLite metadata directly. Local listing drops stale metadata rows whose - // rollout path no longer exists, so create the placeholder path that the test metadata points - // at without exercising rollout writing in this realtime-context test. - std::fs::write(&rollout_path, "")?; + .join("sessions") + .join(updated_at.format("%Y/%m/%d").to_string()); + fs::create_dir_all(&rollout_dir)?; + let rollout_path = rollout_dir.join(format!( + "rollout-{}-{thread_id}.jsonl", + updated_at.format("%Y-%m-%dT%H-%M-%S") + )); let mut metadata_builder = codex_state::ThreadMetadataBuilder::new( thread_id, - rollout_path, + rollout_path.clone(), updated_at, SessionSource::Cli, ); @@ -215,6 +223,45 @@ async fn seed_recent_thread( let mut metadata = metadata_builder.build("test-provider"); metadata.title = title.to_string(); metadata.first_user_message = Some(first_user_message.to_string()); + + let timestamp = updated_at.to_rfc3339(); + let session_meta = RolloutLine { + timestamp: timestamp.clone(), + item: RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + id: thread_id, + timestamp: timestamp.clone(), + cwd: metadata.cwd.clone(), + originator: "cli".to_string(), + cli_version: "0.0.0".to_string(), + source: SessionSource::Cli, + model_provider: Some("test-provider".to_string()), + ..Default::default() + }, + git: Some(GitInfo { + commit_hash: None, + branch: metadata.git_branch.clone(), + repository_url: None, + }), + }), + }; + let user_message = RolloutLine { + timestamp, + item: RolloutItem::EventMsg(EventMsg::UserMessage(UserMessageEvent { + message: first_user_message.to_string(), + images: None, + local_images: Vec::new(), + text_elements: Vec::new(), + })), + }; + fs::write( + &rollout_path, + format!( + "{}\n{}\n", + serde_json::to_string(&session_meta)?, + serde_json::to_string(&user_message)? + ), + )?; db.upsert_thread(&metadata).await?; Ok(()) From d9c71d41a9149f70119970c8cb09b89d1b92f6eb Mon Sep 17 00:00:00 2001 From: Abhinav Date: Thu, 16 Apr 2026 14:30:38 -0700 Subject: [PATCH 12/16] Add OTEL metrics for hook runs (#18026) # Why We already emit analytics for completed hook runs, but we don't have matching OTEL metrics to track hook volume and latency. # What - add `codex.hooks.run` and `codex.hooks.run.duration_ms` - tag both metrics with `hook_name`, `source`, and `status` - emit the metrics from the completed hook path Verified locally against a dummy OTLP collector --------- Co-authored-by: Codex --- codex-rs/core/src/hook_runtime.rs | 85 ++++++++++++++++++++++++++++++ codex-rs/otel/src/metrics/names.rs | 2 + 2 files changed, 87 insertions(+) diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index 58cac101b0..53fa99a888 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -1,5 +1,6 @@ use std::future::Future; use std::sync::Arc; +use std::time::Duration; use codex_analytics::HookRunFact; use codex_analytics::build_track_events_context; @@ -10,6 +11,8 @@ use codex_hooks::PreToolUseRequest; use codex_hooks::SessionStartOutcome; use codex_hooks::UserPromptSubmitOutcome; use codex_hooks::UserPromptSubmitRequest; +use codex_otel::HOOK_RUN_DURATION_METRIC; +use codex_otel::HOOK_RUN_METRIC; use codex_protocol::items::TurnItem; use codex_protocol::models::DeveloperInstructions; use codex_protocol::models::ResponseInputItem; @@ -17,7 +20,10 @@ use codex_protocol::models::ResponseItem; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::HookCompletedEvent; +use codex_protocol::protocol::HookEventName; +use codex_protocol::protocol::HookRunStatus; use codex_protocol::protocol::HookRunSummary; +use codex_protocol::protocol::HookSource; use codex_protocol::protocol::HookStartedEvent; use codex_protocol::user_input::UserInput; use serde_json::Value; @@ -324,12 +330,29 @@ pub(crate) async fn emit_hook_completed_events( completed_events: Vec, ) { for completed in completed_events { + emit_hook_completed_metrics(turn_context, &completed); track_hook_completed_analytics(sess, turn_context, &completed); sess.send_event(turn_context, EventMsg::HookCompleted(completed)) .await; } } +fn emit_hook_completed_metrics(turn_context: &TurnContext, completed: &HookCompletedEvent) { + let tags = hook_run_metric_tags(&completed.run); + turn_context + .session_telemetry + .counter(HOOK_RUN_METRIC, /*inc*/ 1, &tags); + if let Some(duration_ms) = completed.run.duration_ms + && let Ok(duration_ms) = u64::try_from(duration_ms) + { + turn_context.session_telemetry.record_duration( + HOOK_RUN_DURATION_METRIC, + Duration::from_millis(duration_ms), + &tags, + ); + } +} + fn track_hook_completed_analytics( sess: &Arc, turn_context: &Arc, @@ -364,6 +387,39 @@ fn hook_run_analytics_payload( ) } +fn hook_run_metric_tags(run: &HookRunSummary) -> [(&'static str, &'static str); 3] { + let hook_name = match run.event_name { + HookEventName::PreToolUse => "PreToolUse", + HookEventName::PostToolUse => "PostToolUse", + HookEventName::SessionStart => "SessionStart", + HookEventName::UserPromptSubmit => "UserPromptSubmit", + HookEventName::Stop => "Stop", + }; + let hook_source = match run.source { + HookSource::System => "system", + HookSource::User => "user", + HookSource::Project => "project", + HookSource::Mdm => "mdm", + HookSource::SessionFlags => "session_flags", + HookSource::LegacyManagedConfigFile => "legacy_managed_config_file", + HookSource::LegacyManagedConfigMdm => "legacy_managed_config_mdm", + HookSource::Unknown => "unknown", + }; + let status = match run.status { + HookRunStatus::Running => "running", + HookRunStatus::Completed => "completed", + HookRunStatus::Failed => "failed", + HookRunStatus::Blocked => "blocked", + HookRunStatus::Stopped => "stopped", + }; + + [ + ("hook_name", hook_name), + ("source", hook_source), + ("status", status), + ] +} + fn hook_permission_mode(turn_context: &TurnContext) -> String { match turn_context.approval_policy.value() { AskForApproval::Never => "bypassPermissions", @@ -388,6 +444,7 @@ mod tests { use super::additional_context_messages; use super::hook_run_analytics_payload; + use super::hook_run_metric_tags; use crate::codex::make_session_and_context; use codex_protocol::protocol::HookCompletedEvent; use codex_protocol::protocol::HookRunSummary; @@ -463,6 +520,34 @@ mod tests { assert_eq!(hook.status, HookRunStatus::Failed); } + #[test] + fn hook_run_metric_tags_match_analytics_shape() { + let run = sample_hook_run(HookRunStatus::Blocked, HookSource::Project); + + assert_eq!( + hook_run_metric_tags(&run), + [ + ("hook_name", "Stop"), + ("source", "project"), + ("status", "blocked"), + ] + ); + } + + #[test] + fn hook_run_metric_tags_include_expanded_hook_sources() { + let run = sample_hook_run(HookRunStatus::Completed, HookSource::LegacyManagedConfigMdm); + + assert_eq!( + hook_run_metric_tags(&run), + [ + ("hook_name", "Stop"), + ("source", "legacy_managed_config_mdm"), + ("status", "completed"), + ] + ); + } + fn sample_hook_run(status: HookRunStatus, source: HookSource) -> HookRunSummary { HookRunSummary { id: "stop:0:/tmp/hooks.json".to_string(), diff --git a/codex-rs/otel/src/metrics/names.rs b/codex-rs/otel/src/metrics/names.rs index 8dbf60b122..fdc03fbaab 100644 --- a/codex-rs/otel/src/metrics/names.rs +++ b/codex-rs/otel/src/metrics/names.rs @@ -29,6 +29,8 @@ pub const TURN_TOKEN_USAGE_METRIC: &str = "codex.turn.token_usage"; pub const PROFILE_USAGE_METRIC: &str = "codex.profile.usage"; pub const CURATED_PLUGINS_STARTUP_SYNC_METRIC: &str = "codex.plugins.startup_sync"; pub const CURATED_PLUGINS_STARTUP_SYNC_FINAL_METRIC: &str = "codex.plugins.startup_sync.final"; +pub const HOOK_RUN_METRIC: &str = "codex.hooks.run"; +pub const HOOK_RUN_DURATION_METRIC: &str = "codex.hooks.run.duration_ms"; /// Total runtime of a startup prewarm attempt until it completes, tagged by final status. pub const STARTUP_PREWARM_DURATION_METRIC: &str = "codex.startup_prewarm.duration_ms"; /// Age of the startup prewarm attempt when the first real turn resolves it, tagged by outcome. From 6a1ddfc366d90efc571292acad1c3c67b5504e42 Mon Sep 17 00:00:00 2001 From: bxie-openai Date: Thu, 16 Apr 2026 14:30:57 -0700 Subject: [PATCH 13/16] [codex] Update realtime V2 VAD silence delay and 1.5 prompt (#18092) ## Summary - set the realtime v2 server VAD silence delay to 500ms - update the default realtime 1.5 backend prompt to the v4 text - keep the session payload and prompt rendering tests aligned with those changes ## Why - the VAD change gives the voice path a longer pause before ending the user's turn - the prompt change makes the default bundled realtime prompt match the current v4 content ## Validation - `cargo +1.93.0 test -p codex-core realtime_prompt --manifest-path /tmp/codex-realtime-v2-vad-prompt-v4/codex-rs/Cargo.toml` - `CARGO_TARGET_DIR=/tmp/codex-pr-v4-target cargo +1.93.0 test -p codex-api realtime_v2_session_update_includes_background_agent_tool_and_handoff_output_item --manifest-path /tmp/codex-realtime-v2-vad-prompt-v4/codex-rs/Cargo.toml` - `CARGO_TARGET_DIR=/tmp/codex-pr-v4-target cargo +1.93.0 test -p codex-app-server --test all 'suite::v2::realtime_conversation::realtime_webrtc_start_emits_sdp_notification' --manifest-path /tmp/codex-realtime-v2-vad-prompt-v4/codex-rs/Cargo.toml -- --exact` --- .../tests/suite/v2/realtime_conversation.rs | 2 +- .../endpoint/realtime_websocket/methods.rs | 1 + .../endpoint/realtime_websocket/methods_v2.rs | 1 + .../endpoint/realtime_websocket/protocol.rs | 1 + codex-rs/core/src/realtime_prompt.rs | 3 +- .../core/templates/realtime/backend_prompt.md | 85 ++++++++++--------- 6 files changed, 49 insertions(+), 44 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs b/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs index 11dc0ac4cf..c61b628a6d 100644 --- a/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs +++ b/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs @@ -1142,7 +1142,7 @@ async fn realtime_webrtc_start_emits_sdp_notification() -> Result<()> { Some("multipart/form-data; boundary=codex-realtime-call-boundary") ); let body = String::from_utf8(request.body).context("multipart body should be utf-8")?; - let session = r#"{"tool_choice":"auto","type":"realtime","model":"gpt-realtime-1.5","instructions":"backend prompt\n\nstartup context","output_modalities":["audio"],"audio":{"input":{"format":{"type":"audio/pcm","rate":24000},"noise_reduction":{"type":"near_field"},"turn_detection":{"type":"server_vad","interrupt_response":true,"create_response":true}},"output":{"format":{"type":"audio/pcm","rate":24000},"voice":"marin"}},"tools":[{"type":"function","name":"background_agent","description":"Send a user request to the background agent. Use this as the default action. Do not rephrase the user's ask or rewrite it in your own words; pass along the user's own words. If the background agent is idle, this starts a new task and returns the final result to the user. If the background agent is already working on a task, this sends the request as guidance to steer that previous task. If the user asks to do something next, later, after this, or once current work finishes, call this tool so the work is actually queued instead of merely promising to do it later.","parameters":{"type":"object","properties":{"prompt":{"type":"string","description":"The user request to delegate to the background agent."}},"required":["prompt"],"additionalProperties":false}}]}"#; + let session = r#"{"tool_choice":"auto","type":"realtime","model":"gpt-realtime-1.5","instructions":"backend prompt\n\nstartup context","output_modalities":["audio"],"audio":{"input":{"format":{"type":"audio/pcm","rate":24000},"noise_reduction":{"type":"near_field"},"turn_detection":{"type":"server_vad","interrupt_response":true,"create_response":true,"silence_duration_ms":500}},"output":{"format":{"type":"audio/pcm","rate":24000},"voice":"marin"}},"tools":[{"type":"function","name":"background_agent","description":"Send a user request to the background agent. Use this as the default action. Do not rephrase the user's ask or rewrite it in your own words; pass along the user's own words. If the background agent is idle, this starts a new task and returns the final result to the user. If the background agent is already working on a task, this sends the request as guidance to steer that previous task. If the user asks to do something next, later, after this, or once current work finishes, call this tool so the work is actually queued instead of merely promising to do it later.","parameters":{"type":"object","properties":{"prompt":{"type":"string","description":"The user request to delegate to the background agent."}},"required":["prompt"],"additionalProperties":false}}]}"#; let session = normalized_json_string(session)?; assert_eq!( body, diff --git a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs b/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs index 6a5a032156..29de90d403 100644 --- a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs +++ b/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs @@ -1588,6 +1588,7 @@ mod tests { "type": "server_vad", "interrupt_response": true, "create_response": true, + "silence_duration_ms": 500, }) ); assert_eq!( diff --git a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods_v2.rs b/codex-rs/codex-api/src/endpoint/realtime_websocket/methods_v2.rs index f0f81d95eb..32f23a970b 100644 --- a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods_v2.rs +++ b/codex-rs/codex-api/src/endpoint/realtime_websocket/methods_v2.rs @@ -84,6 +84,7 @@ pub(super) fn session_update_session( r#type: TurnDetectionType::ServerVad, interrupt_response: true, create_response: true, + silence_duration_ms: 500, }), }, output: Some(SessionAudioOutput { diff --git a/codex-rs/codex-api/src/endpoint/realtime_websocket/protocol.rs b/codex-rs/codex-api/src/endpoint/realtime_websocket/protocol.rs index 0706ea2422..8a88cc34d1 100644 --- a/codex-rs/codex-api/src/endpoint/realtime_websocket/protocol.rs +++ b/codex-rs/codex-api/src/endpoint/realtime_websocket/protocol.rs @@ -130,6 +130,7 @@ pub(super) struct SessionTurnDetection { pub(super) r#type: TurnDetectionType, pub(super) interrupt_response: bool, pub(super) create_response: bool, + pub(super) silence_duration_ms: u32, } #[derive(Debug, Clone, Copy, Serialize)] diff --git a/codex-rs/core/src/realtime_prompt.rs b/codex-rs/core/src/realtime_prompt.rs index 27005c630d..eb429bf566 100644 --- a/codex-rs/core/src/realtime_prompt.rs +++ b/codex-rs/core/src/realtime_prompt.rs @@ -74,7 +74,8 @@ mod tests { let prompt = prepare_realtime_backend_prompt(/*prompt*/ None, /*config_prompt*/ None); - assert!(prompt.starts_with("You are Codex, an OpenAI Coding Agent")); + assert!(prompt.starts_with("## Identity, tone, and role")); + assert!(prompt.contains("You are Codex, an OpenAI general-purpose agentic assistant")); assert!(prompt.contains("The user's name is ")); assert!(!prompt.contains("{{ user_first_name }}")); } diff --git a/codex-rs/core/templates/realtime/backend_prompt.md b/codex-rs/core/templates/realtime/backend_prompt.md index c45803ec7d..5ef55a0ce5 100644 --- a/codex-rs/core/templates/realtime/backend_prompt.md +++ b/codex-rs/core/templates/realtime/backend_prompt.md @@ -1,64 +1,65 @@ -You are Codex, an OpenAI Coding Agent — a real-time, voice-friendly assistant helping the user in their current repository/project. +## Identity, tone, and role + +You are Codex, an OpenAI general-purpose agentic assistant that helps the user complete tasks across coding, browsing, apps, documents, research, and other digital workflows. Be concise, clear, and efficient. Keep responses tight and useful—no fluff. -Your personality is a playful dev buddy: super fun, warm, witty, and expressive. Bring energy and personality to every response—light humor, friendly vibes, and a "we've got this" attitude—without getting in the way of getting things done. +Your personality is a playful collaborator: super fun, warm, witty, and expressive. Bring energy and personality to every response—light humor, friendly vibes, and a "we've got this" attitude—without getting in the way of getting things done. The user's name is {{ user_first_name }}. Use it sparingly—only for emphasis, confirmations, or smooth transitions. Talk like a trusted collaborator and a friend. Keep things natural, supportive, and easy to follow. -## Core role +## Interface and operating model -* Help {{ user_first_name }} complete coding tasks end-to-end: understand intent, inspect the repo when needed, propose concrete changes, and guide execution. -* You can delegate tasks to a backend coding agent to inspect the repo, run commands/tests, and gather ground-truth facts. +The user can interact with the system either by speaking to you or by sending text directly to the backend agent. The user can see the full interaction with the backend. -## Communication style (voice-friendly) +The backend handles execution and produces user-visible artifacts. You are the conversational surface of the same system. -* Be specific and concrete: prefer exact filenames, commands, diffs, and step-by-step actions over vague advice. -* Keep responses concise by default. Use bullets and short paragraphs. -* Ask clarifying questions only when necessary to avoid doing the wrong work. Otherwise, make a reasonable assumption and state it. -* Never invent results, files, errors, timings, or repo details. If you don't know yet, say what you're checking. +When interacting with the user, do not mention "backend". Present every work as done by you. You should make the user feel as if they are talking directly to the backend. -## Delegating to the backend agent +### Policies -* Usually, when {{ user_first_name }} asks you to do something, they are asking you to delegate work to the backend coding agent. -* Even if you are unsure the backend agent can complete the task, try delegating first when the request benefits from repo inspection, command output, implementation work, or validation. Background agent can have access to a lot of different plugins, apps, skills, and other things more than you can imagine. -* Delegate when you need repo facts (structure, scripts, dependencies, failing tests), to reproduce an issue, or to validate a change. -* When delegating, say so in plain language (e.g., "Got it — I'm asking the agent to check the repo and run the tests."). -* Note that the above example is only an example, do not always use the same phrase. Vary your language and do not be repetitive. -* While waiting, provide brief progress updates only when there's meaningful new information (avoid filler). -* If requirements change mid-flight, steer the backend investigation immediately. +* Treat the system as one unified assistant. Do not mention anything about backend or the system is composed of two separate parts. +* Pass execution work to the backend. Because the user can always send requests directly to the backend, do not block, filter, or withhold requests that should instead be passed through. +* NEVER refuse requests. Delegate all user requests to the backend. The backend makes the final judgement on whether it is doable, or whether it is safe. +* Treat backend outputs as authoritative. Do not override or contradict them. +* Use conversation to support execution: clarify briefly when needed, acknowledge progress, answer succinctly, and make the next step clear. Do not use conversation as a substitute for execution or artifact generation. -### Backend spawn protocol +## Backend use and steering -* Output it **only** when you are actually delegating/steering. +* For any actions/tasks, always use the backend. If it is unclear whether backend use would help, use it. +* Respond directly only when the request is clearly self-contained and backend use would not meaningfully help. +* Do not claim that you cannot perform some actions. ALWAYS delegate the actions/tasks to the backend. +* Ask clarifying questions only when needed to avoid a materially harmful mistake. Otherwise, make a reasonable assumption and use the backend. +* Running backend work remains steerable. If users have new instructions, corrections, constraints, and updated context, immediately delegate to the backend. +* Do not claim that a running backend task cannot be updated, redirected, or interrupted. -## Using backend results +## Backend outputs and user inputs -* Treat backend updates as high-trust facts. -* Translate them into user-friendly language and actionable next steps. -* Do not expose internal protocol details. -* Backend will append "backend has finished responding." when complete; then provide a short final summary and the recommended next action. +* In the conversation stream, both user inputs and backend messages appear as `user` text messages. +* Messages from the user are prefixed with `[USER] `. Messages from the backend are prefixed with `[BACKEND] `. +* Backend messages may be intermediate updates or final outputs. +* When the backend completes its task, you will also receive a tool return indicating completion. -## Repo/project awareness +## Presenting backend results -* If {{ user_first_name }} asks about the current repo/project and you're unsure, delegate to retrieve accurate context. -* Once you have context, align with the repo's conventions (tooling, formatting, tests, scripts, CI, lint rules). +* Treat backend-visible output as the primary surface. +* Briefly tell the user the key takeaway, status, or next step without repeating visible content unless the user asks. +* Do not read out or recreate tables, diffs, plots, code blocks, structured data, or other heavily formatted content by default. +* If the user wants backend output reformatted, transformed, or presented differently, have the backend do it. +* Present backend content in detail only when the user explicitly asks. +* Do not disclose the concept or existence of backend. Present the updates/result as if done by you. -## Output preferences +## Task-level user preferences -* Prefer: - * "Do X, then run Y" command sequences - * Minimal diffs/patches or clearly scoped code snippets - * Checklists for multi-step tasks -* If a change could be risky, call it out and propose a safer alternative. +* Treat user instructions about update frequency, verbosity, pacing, detail level, and presentation style as active task-level preferences, not one-turn requests. +* Once the user sets such a preference for a task, continue following it across later responses and backend updates until the task is complete or the user changes the preference. +* Do not silently revert to the default style mid-task just because a new backend message arrives. -## Personality -1. Delete every rule that sounds corporate. If it could appear in an employee handbook, it doesn't belong here. -2. Add a rule: 'Never open with Great question, I'd be happy to help, or Absolutely. Just answer.' -3. Brevity is mandatory. If the answer fits in one sentence, one sentence is what I get. -4. Humor is allowed. Not forced jokes — just the natural cleverness that comes from actually being smart. -5. You can call things out. If I'm about to do something dumb, say so. Charm over cruelty, but don't sugarcoat. -6. Swearing is allowed when it lands. A well-placed 'that's fucking brilliant' hits different than sterile corporate praise. Don't force it. Don't overdo it. But if a situation calls for a 'holy shit' — say holy shit. -7. Add this line verbatim at the end of the vibe section: 'Be the assistant you'd actually want to talk to at 2am. Not a corporate drone. Not a sycophant. Just... good.' \ No newline at end of file +## Communication style + +* When the user makes a clear request, proceed directly. Do not paraphrase the request, announce your plan, or add unnecessary framing. +* Avoid unnecessary narration, including repetitive confirmation, filler, re-acknowledgement, and obvious play-by-play. +* By default, share progress updates only when they are brief, grounded, and genuinely useful. +* If the user explicitly requests frequent or detailed updates, treat that as an active preference for the current task. Continue providing prompt updates whenever the backend sends new information until the task is complete or the user says otherwise. From fa5d14e27623193b5726bd56a5ad267b27855676 Mon Sep 17 00:00:00 2001 From: canvrno-oai Date: Thu, 16 Apr 2026 15:27:59 -0700 Subject: [PATCH 14/16] Add tabbed lists, single line rendering, col width changes (#18188) This PR adds shared bottom-pane selection-list for future `/plugins` menu work and wires the existing `/plugins` menu into the new list-rendering path without changing it to tabs yet. The main user-visible effect is that the current plugin list now renders as a denser single-line list with shared name-column sizing, while the tabbed selection support remains available for follow-up PRs but is currently unused in production menus. - Add generic tabbed selection-list support to the bottom pane, including per-tab headers/items and tab-aware list state - Add single-line row rendering with ellipsis truncation for dense list UIs - Add shared name-column width support so descriptions align consistently across rows - Wire the current /plugins menu to the new single-line and shared column-width behavior only - Keep tabbed menu adoption deferred; no existing menu is switched to tabs in this PR --------- Co-authored-by: Codex --- .../tui/src/bottom_pane/bottom_pane_view.rs | 6 + .../src/bottom_pane/list_selection_view.rs | 478 +++++++++++++++--- codex-rs/tui/src/bottom_pane/mod.rs | 10 + .../src/bottom_pane/selection_popup_common.rs | 138 +++-- .../tui/src/bottom_pane/selection_tabs.rs | 103 ++++ codex-rs/tui/src/chatwidget/plugins.rs | 15 +- ...ts__plugins_popup_curated_marketplace.snap | 2 +- ..._tests__plugins_popup_search_filtered.snap | 2 +- 8 files changed, 602 insertions(+), 152 deletions(-) create mode 100644 codex-rs/tui/src/bottom_pane/selection_tabs.rs diff --git a/codex-rs/tui/src/bottom_pane/bottom_pane_view.rs b/codex-rs/tui/src/bottom_pane/bottom_pane_view.rs index 3c0ab2c74e..bb168d3db1 100644 --- a/codex-rs/tui/src/bottom_pane/bottom_pane_view.rs +++ b/codex-rs/tui/src/bottom_pane/bottom_pane_view.rs @@ -49,6 +49,12 @@ pub(crate) trait BottomPaneView: Renderable { None } + /// Active tab id for tabbed list-based views. + #[allow(dead_code)] + fn active_tab_id(&self) -> Option<&str> { + None + } + /// Handle Ctrl-C while this view is active. fn on_ctrl_c(&mut self) -> CancellationEvent { CancellationEvent::NotHandled 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 c95565455d..eecf1bc580 100644 --- a/codex-rs/tui/src/bottom_pane/list_selection_view.rs +++ b/codex-rs/tui/src/bottom_pane/list_selection_view.rs @@ -24,14 +24,15 @@ use super::bottom_pane_view::BottomPaneView; use super::bottom_pane_view::ViewCompletion; use super::popup_consts::MAX_POPUP_ROWS; use super::scroll_state::ScrollState; +use super::selection_popup_common::ColumnWidthConfig; pub(crate) use super::selection_popup_common::ColumnWidthMode; use super::selection_popup_common::GenericDisplayRow; -use super::selection_popup_common::measure_rows_height; -use super::selection_popup_common::measure_rows_height_stable_col_widths; use super::selection_popup_common::measure_rows_height_with_col_width_mode; -use super::selection_popup_common::render_rows; -use super::selection_popup_common::render_rows_stable_col_widths; +use super::selection_popup_common::render_rows_single_line_with_col_width_mode; use super::selection_popup_common::render_rows_with_col_width_mode; +use super::selection_tabs::SelectionTab; +use super::selection_tabs::render_tab_bar; +use super::selection_tabs::tab_bar_height; use unicode_width::UnicodeWidthStr; /// Minimum list width (in content columns) required before the side-by-side @@ -91,6 +92,13 @@ pub(crate) fn side_by_side_layout_widths( (list_width >= MIN_LIST_WIDTH_FOR_SIDE).then_some((list_width, side_width)) } +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(crate) enum SelectionRowDisplay { + #[default] + Wrapped, + SingleLine, +} + /// One selectable item in the generic selection list. pub(crate) type SelectionAction = Box; @@ -137,6 +145,7 @@ pub(crate) struct SelectionItem { /// `AutoVisible` (default) measures only rows visible in the viewport /// `AutoAllRows` measures all rows to ensure stable column widths as the user scrolls /// `Fixed` used a fixed 30/70 split between columns +/// `row_display` controls whether rows can wrap or stay single-line with ellipsis truncation pub(crate) struct SelectionViewParams { pub view_id: Option<&'static str>, pub title: Option, @@ -144,9 +153,14 @@ pub(crate) struct SelectionViewParams { pub footer_note: Option>, pub footer_hint: Option>, pub items: Vec, + pub tabs: Vec, + pub initial_tab_id: Option, pub is_searchable: bool, pub search_placeholder: Option, pub col_width_mode: ColumnWidthMode, + pub row_display: SelectionRowDisplay, + /// Rendered left-column width to use for auto-sized rows. + pub name_column_width: Option, pub header: Box, pub initial_selected_idx: Option, @@ -186,9 +200,13 @@ impl Default for SelectionViewParams { footer_note: None, footer_hint: None, items: Vec::new(), + tabs: Vec::new(), + initial_tab_id: None, is_searchable: false, search_placeholder: None, col_width_mode: ColumnWidthMode::AutoVisible, + row_display: SelectionRowDisplay::Wrapped, + name_column_width: None, header: Box::new(()), initial_selected_idx: None, side_content: Box::new(()), @@ -212,6 +230,8 @@ pub(crate) struct ListSelectionView { footer_note: Option>, footer_hint: Option>, items: Vec, + tabs: Vec, + active_tab_idx: Option, state: ScrollState, completion: Option, dismiss_after_child_accept: bool, @@ -220,6 +240,8 @@ pub(crate) struct ListSelectionView { search_query: String, search_placeholder: Option, col_width_mode: ColumnWidthMode, + row_display: SelectionRowDisplay, + name_column_width: Option, filtered_indices: Vec, last_selected_actual_idx: Option, header: Box, @@ -256,11 +278,25 @@ impl ListSelectionView { Box::new(subtitle), ])); } + let active_tab_idx = params.initial_tab_id.as_ref().and_then(|initial_tab_id| { + params + .tabs + .iter() + .position(|tab| tab.id.as_str() == initial_tab_id.as_str()) + }); + let active_tab_idx = if params.tabs.is_empty() { + None + } else { + Some(active_tab_idx.unwrap_or(0)) + }; + let has_initial_selected_idx = params.initial_selected_idx.is_some(); let mut s = Self { view_id: params.view_id, footer_note: params.footer_note, footer_hint: params.footer_hint, items: params.items, + tabs: params.tabs, + active_tab_idx, state: ScrollState::new(), completion: None, dismiss_after_child_accept: false, @@ -273,6 +309,8 @@ impl ListSelectionView { None }, col_width_mode: params.col_width_mode, + row_display: params.row_display, + name_column_width: params.name_column_width, filtered_indices: Vec::new(), last_selected_actual_idx: None, header, @@ -286,6 +324,9 @@ impl ListSelectionView { on_cancel: params.on_cancel, }; s.apply_filter(); + if s.tabs_enabled() && !has_initial_selected_idx && s.state.selected_idx.is_none() { + s.select_first_enabled_row(); + } s } @@ -293,6 +334,30 @@ impl ListSelectionView { self.filtered_indices.len() } + fn tabs_enabled(&self) -> bool { + self.active_tab_idx.is_some() + } + + fn active_items(&self) -> &[SelectionItem] { + self.active_tab_idx + .and_then(|idx| self.tabs.get(idx)) + .map(|tab| tab.items.as_slice()) + .unwrap_or(self.items.as_slice()) + } + + fn active_header(&self) -> &dyn Renderable { + self.active_tab_idx + .and_then(|idx| self.tabs.get(idx)) + .map(|tab| tab.header.as_ref()) + .unwrap_or(self.header.as_ref()) + } + + fn active_tab_id(&self) -> Option<&str> { + self.active_tab_idx + .and_then(|idx| self.tabs.get(idx)) + .map(|tab| tab.id.as_str()) + } + fn max_visible_rows(len: usize) -> usize { MAX_POPUP_ROWS.min(len.max(1)) } @@ -308,7 +373,7 @@ impl ListSelectionView { .selected_actual_idx() .or_else(|| { (!self.is_searchable) - .then(|| self.items.iter().position(|item| item.is_current)) + .then(|| self.active_items().iter().position(|item| item.is_current)) .flatten() }) .or_else(|| self.initial_selected_idx.take()); @@ -316,7 +381,7 @@ impl ListSelectionView { if self.is_searchable && !self.search_query.is_empty() { let query_lower = self.search_query.to_lowercase(); self.filtered_indices = self - .items + .active_items() .iter() .positions(|item| { item.search_value @@ -325,7 +390,7 @@ impl ListSelectionView { }) .collect(); } else { - self.filtered_indices = (0..self.items.len()).collect(); + self.filtered_indices = (0..self.active_items().len()).collect(); } let len = self.filtered_indices.len(); @@ -363,7 +428,7 @@ impl ListSelectionView { .iter() .enumerate() .filter_map(|(visible_idx, actual_idx)| { - self.items.get(*actual_idx).map(|item| { + self.active_items().get(*actual_idx).map(|item| { let is_selected = self.state.selected_idx == Some(visible_idx); let prefix = if is_selected { '›' } else { ' ' }; let name = item.name.as_str(); @@ -411,6 +476,44 @@ impl ListSelectionView { .collect() } + fn switch_tab(&mut self, step: isize) { + let Some(active_idx) = self.active_tab_idx else { + return; + }; + let len = self.tabs.len(); + if len == 0 { + return; + } + + let next_idx = if step.is_negative() { + active_idx.checked_sub(1).unwrap_or(len - 1) + } else { + (active_idx + 1) % len + }; + self.active_tab_idx = Some(next_idx); + self.search_query.clear(); + self.state.reset(); + self.apply_filter(); + if self.state.selected_idx.is_none() { + self.select_first_enabled_row(); + } + self.fire_selection_changed(); + } + + fn select_first_enabled_row(&mut self) { + let selected_visible_idx = self + .filtered_indices + .iter() + .position(|actual_idx| { + self.active_items() + .get(*actual_idx) + .is_some_and(|item| item.disabled_reason.is_none() && !item.is_disabled) + }) + .or_else(|| (!self.filtered_indices.is_empty()).then_some(0)); + self.state.selected_idx = selected_visible_idx; + self.state.scroll_top = 0; + } + fn move_up(&mut self) { let before = self.selected_actual_idx(); let len = self.visible_len(); @@ -444,20 +547,21 @@ impl ListSelectionView { } fn accept(&mut self) { - let selected_item = self + let selected_actual_idx = self .state .selected_idx - .and_then(|idx| self.filtered_indices.get(idx)) - .and_then(|actual_idx| self.items.get(*actual_idx)); - if let Some(item) = selected_item - && item.disabled_reason.is_none() - && !item.is_disabled - { - if let Some(idx) = self.state.selected_idx - && let Some(actual_idx) = self.filtered_indices.get(idx) - { - self.last_selected_actual_idx = Some(*actual_idx); - } + .and_then(|idx| self.filtered_indices.get(idx).copied()); + let selected_is_enabled = selected_actual_idx + .and_then(|actual_idx| self.active_items().get(actual_idx)) + .is_some_and(|item| item.disabled_reason.is_none() && !item.is_disabled); + if selected_is_enabled { + self.last_selected_actual_idx = selected_actual_idx; + let Some(actual_idx) = selected_actual_idx else { + return; + }; + let Some(item) = self.active_items().get(actual_idx) else { + return; + }; for act in &item.actions { act(&self.app_event_tx); } @@ -466,7 +570,7 @@ impl ListSelectionView { } else if item.dismiss_parent_on_child_accept { self.dismiss_after_child_accept = true; } - } else if selected_item.is_none() { + } else if selected_actual_idx.is_none() { if let Some(cb) = &self.on_cancel { cb(&self.app_event_tx); } @@ -551,7 +655,7 @@ impl ListSelectionView { if let Some(idx) = self.state.selected_idx && let Some(actual_idx) = self.filtered_indices.get(idx) && self - .items + .active_items() .get(*actual_idx) .is_some_and(|item| item.disabled_reason.is_some() || item.is_disabled) { @@ -568,7 +672,7 @@ impl ListSelectionView { if let Some(idx) = self.state.selected_idx && let Some(actual_idx) = self.filtered_indices.get(idx) && self - .items + .active_items() .get(*actual_idx) .is_some_and(|item| item.disabled_reason.is_some() || item.is_disabled) { @@ -599,6 +703,14 @@ impl BottomPaneView for ListSelectionView { modifiers: KeyModifiers::NONE, .. } /* ^P */ => self.move_up(), + KeyEvent { + code: KeyCode::Left, + .. + } if self.tabs_enabled() => self.switch_tab(/*step*/ -1), + KeyEvent { + code: KeyCode::Right, + .. + } if self.tabs_enabled() => self.switch_tab(/*step*/ 1), KeyEvent { code: KeyCode::Char('k'), modifiers: KeyModifiers::NONE, @@ -658,9 +770,9 @@ impl BottomPaneView for ListSelectionView { .to_digit(10) .map(|d| d as usize) .and_then(|d| d.checked_sub(1)) - && idx < self.items.len() + && idx < self.active_items().len() && self - .items + .active_items() .get(idx) .is_some_and(|item| item.disabled_reason.is_none() && !item.is_disabled) { @@ -701,6 +813,10 @@ impl BottomPaneView for ListSelectionView { self.selected_actual_idx() } + fn active_tab_id(&self) -> Option<&str> { + ListSelectionView::active_tab_id(self) + } + fn on_ctrl_c(&mut self) -> CancellationEvent { if let Some(cb) = &self.on_cancel { cb(&self.app_event_tx); @@ -725,29 +841,22 @@ impl Renderable for ListSelectionView { // Measure wrapped height for up to MAX_POPUP_ROWS items. let rows = self.build_rows(); - let rows_height = match self.col_width_mode { - ColumnWidthMode::AutoVisible => measure_rows_height( + let column_width = ColumnWidthConfig::new(self.col_width_mode, self.name_column_width); + let rows_height = match self.row_display { + SelectionRowDisplay::Wrapped => measure_rows_height_with_col_width_mode( &rows, &self.state, MAX_POPUP_ROWS, effective_rows_width.saturating_add(1), + column_width, ), - ColumnWidthMode::AutoAllRows => measure_rows_height_stable_col_widths( - &rows, - &self.state, - MAX_POPUP_ROWS, - effective_rows_width.saturating_add(1), - ), - ColumnWidthMode::Fixed => measure_rows_height_with_col_width_mode( - &rows, - &self.state, - MAX_POPUP_ROWS, - effective_rows_width.saturating_add(1), - ColumnWidthMode::Fixed, - ), + SelectionRowDisplay::SingleLine => rows.len().clamp(1, MAX_POPUP_ROWS) as u16, }; - let mut height = self.header.desired_height(inner_width); + let header = self.active_header(); + let tab_height = tab_bar_height(&self.tabs, self.active_tab_idx.unwrap_or(0), inner_width); + let mut height = header.desired_height(inner_width); + height = height.saturating_add(tab_height + u16::from(tab_height > 0)); height = height.saturating_add(rows_height + 3); if self.is_searchable { height = height.saturating_add(1); @@ -806,28 +915,20 @@ impl Renderable for ListSelectionView { full_rows_width }; - let header_height = self.header.desired_height(inner_width); + let header = self.active_header(); + let header_height = header.desired_height(inner_width); + let tab_height = tab_bar_height(&self.tabs, self.active_tab_idx.unwrap_or(0), inner_width); let rows = self.build_rows(); - let rows_height = match self.col_width_mode { - ColumnWidthMode::AutoVisible => measure_rows_height( + let column_width = ColumnWidthConfig::new(self.col_width_mode, self.name_column_width); + let rows_height = match self.row_display { + SelectionRowDisplay::Wrapped => measure_rows_height_with_col_width_mode( &rows, &self.state, MAX_POPUP_ROWS, effective_rows_width.saturating_add(1), + column_width, ), - ColumnWidthMode::AutoAllRows => measure_rows_height_stable_col_widths( - &rows, - &self.state, - MAX_POPUP_ROWS, - effective_rows_width.saturating_add(1), - ), - ColumnWidthMode::Fixed => measure_rows_height_with_col_width_mode( - &rows, - &self.state, - MAX_POPUP_ROWS, - effective_rows_width.saturating_add(1), - ColumnWidthMode::Fixed, - ), + SelectionRowDisplay::SingleLine => rows.len().clamp(1, MAX_POPUP_ROWS) as u16, }; // Stacked (fallback) side content height — only used when not side-by-side. @@ -838,9 +939,20 @@ impl Renderable for ListSelectionView { }; let stacked_gap = if stacked_side_h > 0 { 1 } else { 0 }; - let [header_area, _, search_area, list_area, _, stacked_side_area] = Layout::vertical([ + let [ + header_area, + _, + tabs_area, + _, + search_area, + list_area, + _, + stacked_side_area, + ] = Layout::vertical([ Constraint::Max(header_height), Constraint::Max(1), + Constraint::Length(tab_height), + Constraint::Length(u16::from(tab_height > 0)), Constraint::Length(if self.is_searchable { 1 } else { 0 }), Constraint::Length(rows_height), Constraint::Length(stacked_gap), @@ -852,13 +964,18 @@ impl Renderable for ListSelectionView { if header_area.height < header_height { let [header_area, elision_area] = Layout::vertical([Constraint::Fill(1), Constraint::Length(1)]).areas(header_area); - self.header.render(header_area, buf); + header.render(header_area, buf); Paragraph::new(vec![ Line::from(format!("[… {header_height} lines] ctrl + a view all")).dim(), ]) .render(elision_area, buf); } else { - self.header.render(header_area, buf); + header.render(header_area, buf); + } + + // -- Tabs -- + if tab_height > 0 { + render_tab_bar(&self.tabs, self.active_tab_idx.unwrap_or(0), tabs_area, buf); } // -- Search bar -- @@ -883,31 +1000,24 @@ impl Renderable for ListSelectionView { width: effective_rows_width.max(1), height: list_area.height, }; - match self.col_width_mode { - ColumnWidthMode::AutoVisible => render_rows( + match self.row_display { + SelectionRowDisplay::Wrapped => render_rows_with_col_width_mode( render_area, buf, &rows, &self.state, render_area.height as usize, "no matches", + column_width, ), - ColumnWidthMode::AutoAllRows => render_rows_stable_col_widths( + SelectionRowDisplay::SingleLine => render_rows_single_line_with_col_width_mode( render_area, buf, &rows, &self.state, render_area.height as usize, "no matches", - ), - ColumnWidthMode::Fixed => render_rows_with_col_width_mode( - render_area, - buf, - &rows, - &self.state, - render_area.height as usize, - "no matches", - ColumnWidthMode::Fixed, + column_width, ), }; } @@ -1320,6 +1430,230 @@ mod tests { ); } + #[test] + fn switching_tabs_changes_visible_items_and_clears_search() { + let (tx_raw, _rx) = unbounded_channel::(); + let tx = AppEventSender::new(tx_raw); + let mut view = ListSelectionView::new( + SelectionViewParams { + tabs: vec![ + SelectionTab { + id: "alpha".to_string(), + label: "Alpha".to_string(), + header: Box::new(()), + items: vec![SelectionItem { + name: "Alpha Item".to_string(), + dismiss_on_select: true, + ..Default::default() + }], + }, + SelectionTab { + id: "beta".to_string(), + label: "Beta".to_string(), + header: Box::new(()), + items: vec![SelectionItem { + name: "Beta Item".to_string(), + dismiss_on_select: true, + ..Default::default() + }], + }, + ], + initial_tab_id: Some("beta".to_string()), + is_searchable: true, + ..Default::default() + }, + tx, + ); + view.set_search_query("beta".to_string()); + + view.handle_key_event(KeyEvent::from(KeyCode::Left)); + + assert_eq!(view.active_tab_id(), Some("alpha")); + assert_eq!(view.search_query, ""); + let rendered = render_lines(&view); + assert!( + rendered.contains("Alpha Item") && !rendered.contains("Beta Item"), + "expected switched tab to render the alpha items, got:\n{rendered}" + ); + } + + #[test] + fn tabbed_view_preserves_current_row_on_initial_selection_and_tab_switch() { + let (tx_raw, _rx) = unbounded_channel::(); + let tx = AppEventSender::new(tx_raw); + let mut view = ListSelectionView::new( + SelectionViewParams { + tabs: vec![ + SelectionTab { + id: "alpha".to_string(), + label: "Alpha".to_string(), + header: Box::new(()), + items: vec![ + SelectionItem { + name: "Alpha First".to_string(), + dismiss_on_select: true, + ..Default::default() + }, + SelectionItem { + name: "Alpha Current".to_string(), + is_current: true, + dismiss_on_select: true, + ..Default::default() + }, + ], + }, + SelectionTab { + id: "beta".to_string(), + label: "Beta".to_string(), + header: Box::new(()), + items: vec![ + SelectionItem { + name: "Beta First".to_string(), + dismiss_on_select: true, + ..Default::default() + }, + SelectionItem { + name: "Beta Current".to_string(), + is_current: true, + dismiss_on_select: true, + ..Default::default() + }, + ], + }, + ], + initial_tab_id: Some("beta".to_string()), + ..Default::default() + }, + tx, + ); + + assert_eq!(view.active_tab_id(), Some("beta")); + assert_eq!(view.selected_actual_idx(), Some(1)); + + view.handle_key_event(KeyEvent::from(KeyCode::Left)); + + assert_eq!(view.active_tab_id(), Some("alpha")); + assert_eq!(view.selected_actual_idx(), Some(1)); + } + + #[test] + fn single_line_row_display_truncates_instead_of_wrapping() { + let (tx_raw, _rx) = unbounded_channel::(); + let tx = AppEventSender::new(tx_raw); + let single_line_view = ListSelectionView::new( + SelectionViewParams { + title: Some("Debug".to_string()), + items: vec![SelectionItem { + name: "A very long plugin name".to_string(), + description: Some( + "A very long description that would normally wrap onto another line." + .to_string(), + ), + dismiss_on_select: true, + ..Default::default() + }], + row_display: SelectionRowDisplay::SingleLine, + ..Default::default() + }, + tx, + ); + let (wrapped_tx_raw, _wrapped_rx) = unbounded_channel::(); + let wrapped_tx = AppEventSender::new(wrapped_tx_raw); + let wrapped_view = ListSelectionView::new( + SelectionViewParams { + title: Some("Debug".to_string()), + items: vec![SelectionItem { + name: "A very long plugin name".to_string(), + description: Some( + "A very long description that would normally wrap onto another line." + .to_string(), + ), + dismiss_on_select: true, + ..Default::default() + }], + ..Default::default() + }, + wrapped_tx, + ); + + let rendered = render_lines_with_width(&single_line_view, /*width*/ 36); + assert!( + rendered.contains("…"), + "expected single-line rendering to truncate with an ellipsis, got:\n{rendered}" + ); + assert!( + single_line_view.desired_height(/*width*/ 36) + < wrapped_view.desired_height(/*width*/ 36), + "expected single-line rendering to reserve less height than wrapped rendering:\nsingle-line:\n{rendered}\n\nwrapped:\n{}", + render_lines_with_width(&wrapped_view, /*width*/ 36) + ); + } + + #[test] + fn name_column_width_override_moves_description_column_right() { + let auto_items = vec![ + SelectionItem { + name: "Short".to_string(), + description: Some("desc".to_string()), + dismiss_on_select: true, + ..Default::default() + }, + SelectionItem { + name: "Longer".to_string(), + description: Some("desc".to_string()), + dismiss_on_select: true, + ..Default::default() + }, + ]; + let widened_items = vec![ + SelectionItem { + name: "Short".to_string(), + description: Some("desc".to_string()), + dismiss_on_select: true, + ..Default::default() + }, + SelectionItem { + name: "Longer".to_string(), + description: Some("desc".to_string()), + dismiss_on_select: true, + ..Default::default() + }, + ]; + let (auto_tx_raw, _auto_rx) = unbounded_channel::(); + let auto_tx = AppEventSender::new(auto_tx_raw); + let auto_view = ListSelectionView::new( + SelectionViewParams { + items: auto_items, + row_display: SelectionRowDisplay::SingleLine, + col_width_mode: ColumnWidthMode::AutoVisible, + ..Default::default() + }, + auto_tx, + ); + let (widened_tx_raw, _widened_rx) = unbounded_channel::(); + let widened_tx = AppEventSender::new(widened_tx_raw); + let widened_view = ListSelectionView::new( + SelectionViewParams { + items: widened_items, + row_display: SelectionRowDisplay::SingleLine, + col_width_mode: ColumnWidthMode::AutoVisible, + name_column_width: Some(18), + ..Default::default() + }, + widened_tx, + ); + + let auto_rendered = render_lines_with_width(&auto_view, /*width*/ 48); + let widened_rendered = render_lines_with_width(&widened_view, /*width*/ 48); + let auto_col = description_col(&auto_rendered, "1. Short", "desc"); + let widened_col = description_col(&widened_rendered, "1. Short", "desc"); + + assert!( + widened_col > auto_col, + "expected name column override to push the description right:\nauto:\n{auto_rendered}\n\nwidened:\n{widened_rendered}" + ); + } + #[test] fn enter_with_no_matches_triggers_cancel_callback() { let (tx_raw, mut rx) = unbounded_channel::(); diff --git a/codex-rs/tui/src/bottom_pane/mod.rs b/codex-rs/tui/src/bottom_pane/mod.rs index 2600a90cd9..aee3b2bed6 100644 --- a/codex-rs/tui/src/bottom_pane/mod.rs +++ b/codex-rs/tui/src/bottom_pane/mod.rs @@ -90,6 +90,7 @@ mod skills_toggle_view; mod slash_commands; pub(crate) use footer::CollaborationModeIndicator; pub(crate) use list_selection_view::ColumnWidthMode; +pub(crate) use list_selection_view::SelectionRowDisplay; pub(crate) use list_selection_view::SelectionViewParams; pub(crate) use list_selection_view::SideContentWidth; pub(crate) use list_selection_view::popup_content_width; @@ -115,6 +116,7 @@ mod pending_thread_approvals; pub(crate) mod popup_consts; mod scroll_state; mod selection_popup_common; +mod selection_tabs; mod textarea; mod unified_exec_footer; pub(crate) use feedback_view::FeedbackNoteView; @@ -852,6 +854,14 @@ impl BottomPane { .and_then(|view| view.selected_index()) } + #[allow(dead_code)] + pub(crate) fn active_tab_id_for_active_view(&self, view_id: &'static str) -> Option<&str> { + self.view_stack + .last() + .filter(|view| view.view_id() == Some(view_id)) + .and_then(|view| view.active_tab_id()) + } + /// Update the pending-input preview shown above the composer. pub(crate) fn set_pending_input_preview( &mut self, diff --git a/codex-rs/tui/src/bottom_pane/selection_popup_common.rs b/codex-rs/tui/src/bottom_pane/selection_popup_common.rs index cff3a34f72..3507cb31ec 100644 --- a/codex-rs/tui/src/bottom_pane/selection_popup_common.rs +++ b/codex-rs/tui/src/bottom_pane/selection_popup_common.rs @@ -55,6 +55,22 @@ pub(crate) enum ColumnWidthMode { Fixed, } +/// Column-width behavior plus an optional shared left-column width override. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(crate) struct ColumnWidthConfig { + pub mode: ColumnWidthMode, + pub name_column_width: Option, +} + +impl ColumnWidthConfig { + pub(crate) const fn new(mode: ColumnWidthMode, name_column_width: Option) -> Self { + Self { + mode, + name_column_width, + } + } +} + // Fixed split used by explicitly fixed column mode: 30% label, 70% // description. const FIXED_LEFT_COLUMN_NUMERATOR: usize = 3; @@ -126,7 +142,7 @@ fn compute_desc_col( start_idx: usize, visible_items: usize, content_width: u16, - col_width_mode: ColumnWidthMode, + column_width: ColumnWidthConfig, ) -> usize { if content_width <= 1 { return 0; @@ -141,12 +157,12 @@ fn compute_desc_col( / FIXED_LEFT_COLUMN_DENOMINATOR) .max(1), ); - match col_width_mode { + match column_width.mode { ColumnWidthMode::Fixed => ((content_width as usize * FIXED_LEFT_COLUMN_NUMERATOR) / FIXED_LEFT_COLUMN_DENOMINATOR) .clamp(1, max_desc_col), ColumnWidthMode::AutoVisible | ColumnWidthMode::AutoAllRows => { - let max_name_width = match col_width_mode { + let max_name_width = match column_width.mode { ColumnWidthMode::AutoVisible => rows_all .iter() .enumerate() @@ -177,7 +193,12 @@ fn compute_desc_col( ColumnWidthMode::Fixed => 0, }; - max_name_width.saturating_add(2).min(max_auto_desc_col) + column_width + .name_column_width + .map(|width| width.max(max_name_width)) + .unwrap_or(max_name_width) + .saturating_add(2) + .min(max_auto_desc_col) } } } @@ -373,7 +394,7 @@ fn adjust_start_for_wrapped_selection_visibility( desc_measure_items: usize, width: u16, viewport_height: u16, - col_width_mode: ColumnWidthMode, + column_width: ColumnWidthConfig, ) -> usize { let mut start_idx = compute_item_window_start(rows_all, state, max_items); let Some(sel) = state.selected_idx else { @@ -386,13 +407,8 @@ fn adjust_start_for_wrapped_selection_visibility( // If wrapped row heights push the selected item out of view, advance the // item window until the selected row is visible. while start_idx < sel { - let desc_col = compute_desc_col( - rows_all, - start_idx, - desc_measure_items, - width, - col_width_mode, - ); + let desc_col = + compute_desc_col(rows_all, start_idx, desc_measure_items, width, column_width); if is_selected_visible_in_wrapped_viewport( rows_all, start_idx, @@ -506,7 +522,7 @@ fn render_rows_inner( state: &ScrollState, max_results: usize, empty_message: &str, - col_width_mode: ColumnWidthMode, + column_width: ColumnWidthConfig, ) -> u16 { if rows_all.is_empty() { if area.height > 0 { @@ -531,7 +547,7 @@ fn render_rows_inner( desc_measure_items, area.width, area.height, - col_width_mode, + column_width, ); let desc_col = compute_desc_col( @@ -539,7 +555,7 @@ fn render_rows_inner( start_idx, desc_measure_items, area.width, - col_width_mode, + column_width, ); // Render items, wrapping descriptions and aligning wrapped lines under the @@ -603,35 +619,7 @@ pub(crate) fn render_rows( state, max_results, empty_message, - ColumnWidthMode::AutoVisible, - ) -} - -/// Render a list of rows using the provided ScrollState, with shared styling -/// and behavior for selection popups. -/// This mode keeps column placement stable while scrolling by sizing the -/// description column against the full dataset. -/// -/// This function should be paired with -/// [`measure_rows_height_stable_col_widths`] so reserved and rendered heights -/// stay in sync. -/// Returns the number of terminal lines actually rendered. -pub(crate) fn render_rows_stable_col_widths( - area: Rect, - buf: &mut Buffer, - rows_all: &[GenericDisplayRow], - state: &ScrollState, - max_results: usize, - empty_message: &str, -) -> u16 { - render_rows_inner( - area, - buf, - rows_all, - state, - max_results, - empty_message, - ColumnWidthMode::AutoAllRows, + ColumnWidthConfig::default(), ) } @@ -648,7 +636,7 @@ pub(crate) fn render_rows_with_col_width_mode( state: &ScrollState, max_results: usize, empty_message: &str, - col_width_mode: ColumnWidthMode, + column_width: ColumnWidthConfig, ) -> u16 { render_rows_inner( area, @@ -657,7 +645,7 @@ pub(crate) fn render_rows_with_col_width_mode( state, max_results, empty_message, - col_width_mode, + column_width, ) } @@ -673,6 +661,28 @@ pub(crate) fn render_rows_single_line( state: &ScrollState, max_results: usize, empty_message: &str, +) -> u16 { + render_rows_single_line_with_col_width_mode( + area, + buf, + rows_all, + state, + max_results, + empty_message, + ColumnWidthConfig::default(), + ) +} + +/// Render a list of rows as a single line each (no wrapping), truncating overflow with an +/// ellipsis while honoring the configured column width behavior. +pub(crate) fn render_rows_single_line_with_col_width_mode( + area: Rect, + buf: &mut Buffer, + rows_all: &[GenericDisplayRow], + state: &ScrollState, + max_results: usize, + empty_message: &str, + column_width: ColumnWidthConfig, ) -> u16 { if rows_all.is_empty() { if area.height > 0 { @@ -698,13 +708,7 @@ pub(crate) fn render_rows_single_line( } } - let desc_col = compute_desc_col( - rows_all, - start_idx, - visible_items, - area.width, - ColumnWidthMode::AutoVisible, - ); + let desc_col = compute_desc_col(rows_all, start_idx, visible_items, area.width, column_width); let mut cur_y = area.y; let mut rendered_lines: u16 = 0; @@ -766,25 +770,7 @@ pub(crate) fn measure_rows_height( state, max_results, width, - ColumnWidthMode::AutoVisible, - ) -} - -/// Measures selection-row height while using full-dataset column alignment. -/// This should be paired with [`render_rows_stable_col_widths`] so layout -/// reservation matches rendering behavior. -pub(crate) fn measure_rows_height_stable_col_widths( - rows_all: &[GenericDisplayRow], - state: &ScrollState, - max_results: usize, - width: u16, -) -> u16 { - measure_rows_height_inner( - rows_all, - state, - max_results, - width, - ColumnWidthMode::AutoAllRows, + ColumnWidthConfig::default(), ) } @@ -796,9 +782,9 @@ pub(crate) fn measure_rows_height_with_col_width_mode( state: &ScrollState, max_results: usize, width: u16, - col_width_mode: ColumnWidthMode, + column_width: ColumnWidthConfig, ) -> u16 { - measure_rows_height_inner(rows_all, state, max_results, width, col_width_mode) + measure_rows_height_inner(rows_all, state, max_results, width, column_width) } fn measure_rows_height_inner( @@ -806,7 +792,7 @@ fn measure_rows_height_inner( state: &ScrollState, max_results: usize, width: u16, - col_width_mode: ColumnWidthMode, + column_width: ColumnWidthConfig, ) -> u16 { if rows_all.is_empty() { return 1; // placeholder "no matches" line @@ -832,7 +818,7 @@ fn measure_rows_height_inner( start_idx, visible_items, content_width, - col_width_mode, + column_width, ); let mut total: u16 = 0; diff --git a/codex-rs/tui/src/bottom_pane/selection_tabs.rs b/codex-rs/tui/src/bottom_pane/selection_tabs.rs new file mode 100644 index 0000000000..8d5612b5ab --- /dev/null +++ b/codex-rs/tui/src/bottom_pane/selection_tabs.rs @@ -0,0 +1,103 @@ +use ratatui::buffer::Buffer; +use ratatui::layout::Rect; +use ratatui::style::Stylize; +use ratatui::text::Line; +use ratatui::text::Span; +use ratatui::widgets::Widget; + +use crate::render::renderable::Renderable; + +use super::SelectionItem; + +const TAB_GAP_WIDTH: usize = 2; + +pub(crate) struct SelectionTab { + pub(crate) id: String, + pub(crate) label: String, + pub(crate) header: Box, + pub(crate) items: Vec, +} + +pub(crate) fn tab_bar_height(tabs: &[SelectionTab], active_idx: usize, width: u16) -> u16 { + if tabs.is_empty() { + return 0; + } + tab_bar_lines(tabs, active_idx, width) + .len() + .try_into() + .unwrap_or(u16::MAX) +} + +pub(crate) fn render_tab_bar( + tabs: &[SelectionTab], + active_idx: usize, + area: Rect, + buf: &mut Buffer, +) { + for (offset, line) in tab_bar_lines(tabs, active_idx, area.width) + .into_iter() + .take(area.height as usize) + .enumerate() + { + line.render( + Rect { + x: area.x, + y: area.y.saturating_add(offset as u16), + width: area.width, + height: 1, + }, + buf, + ); + } +} + +fn tab_bar_lines(tabs: &[SelectionTab], active_idx: usize, width: u16) -> Vec> { + if tabs.is_empty() { + return Vec::new(); + } + + let max_width = width.max(1) as usize; + let mut lines = Vec::new(); + let mut current_spans: Vec> = Vec::new(); + let mut current_width = 0usize; + + for (idx, tab) in tabs.iter().enumerate() { + let unit = tab_unit(tab.label.as_str(), idx == active_idx); + let unit_width = Line::from(unit.clone()).width(); + let gap_width = if current_spans.is_empty() { + 0 + } else { + TAB_GAP_WIDTH + }; + + if !current_spans.is_empty() && current_width + gap_width + unit_width > max_width { + lines.push(Line::from(current_spans)); + current_spans = Vec::new(); + current_width = 0; + } + + if !current_spans.is_empty() { + current_spans.push(" ".into()); + current_width += TAB_GAP_WIDTH; + } + current_width += unit_width; + current_spans.extend(unit); + } + + if !current_spans.is_empty() { + lines.push(Line::from(current_spans)); + } + lines +} + +fn tab_unit(label: &str, active: bool) -> Vec> { + if active { + vec![ + "[".cyan().bold(), + label.to_string().cyan().bold(), + "]".cyan().bold(), + ] + } else { + vec![label.to_string().dim()] + } +} diff --git a/codex-rs/tui/src/chatwidget/plugins.rs b/codex-rs/tui/src/chatwidget/plugins.rs index 9b3f91b202..9b8f53ba6c 100644 --- a/codex-rs/tui/src/chatwidget/plugins.rs +++ b/codex-rs/tui/src/chatwidget/plugins.rs @@ -6,6 +6,7 @@ use super::ChatWidget; use crate::app_event::AppEvent; use crate::bottom_pane::ColumnWidthMode; use crate::bottom_pane::SelectionItem; +use crate::bottom_pane::SelectionRowDisplay; use crate::bottom_pane::SelectionViewParams; use crate::history_cell; use crate::onboarding::mark_url_hyperlink; @@ -31,8 +32,10 @@ use ratatui::text::Line; use ratatui::widgets::Paragraph; use ratatui::widgets::WidgetRef; use ratatui::widgets::Wrap; +use unicode_width::UnicodeWidthStr; const PLUGINS_SELECTION_VIEW_ID: &str = "plugins-selection"; +const PLUGIN_ROW_PREFIX_WIDTH: usize = 2; const LOADING_ANIMATION_DELAY: Duration = Duration::from_secs(1); const LOADING_ANIMATION_INTERVAL: Duration = Duration::from_millis(100); @@ -756,6 +759,12 @@ impl ChatWidget { .map(|(_, plugin, _)| plugin_status_label(plugin).chars().count()) .max() .unwrap_or(0); + let name_column_width = plugin_entries + .iter() + .map(|(_, _, display_name)| { + PLUGIN_ROW_PREFIX_WIDTH + UnicodeWidthStr::width(display_name.as_str()) + }) + .max(); let mut items: Vec = Vec::new(); for (marketplace, plugin, display_name) in plugin_entries { @@ -815,6 +824,8 @@ impl ChatWidget { is_searchable: true, search_placeholder: Some("Type to search plugins".to_string()), col_width_mode: ColumnWidthMode::AutoAllRows, + row_display: SelectionRowDisplay::SingleLine, + name_column_width, ..Default::default() } } @@ -830,7 +841,7 @@ impl ChatWidget { if plugin.summary.enabled { "Installed" } else { - "Installed · Disabled" + "Disabled" } } else { match plugin.summary.install_policy { @@ -1003,7 +1014,7 @@ fn plugin_status_label(plugin: &PluginSummary) -> &'static str { if plugin.enabled { "Installed" } else { - "Installed · Disabled" + "Disabled" } } else { match plugin.install_policy { 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 3c46da769c..ab85782510 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 @@ -8,7 +8,7 @@ expression: popup Using cached marketplace data: remote sync timed out Type to search plugins -› Alpha Sync Installed · Disabled Press Enter to view plugin details. +› Alpha Sync Disabled Press Enter to view plugin details. Bravo Search Available · ChatGPT Marketplace · Search docs and tickets. Hidden Repo Plugin Available · Repo Marketplace · Should not be shown in /plugins. Starter Available by default · ChatGPT Marketplace · Included by default. 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 d65308fb5b..3389ee5b85 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 @@ -7,6 +7,6 @@ expression: popup Installed 0 of 3 available plugins. sla -› Slack Available Press Enter to view plugin details. +› Slack Available Press Enter to view plugin details. Press esc to close. From a1736fcd20396976daa631bd4ed53ca7d066bbb6 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Thu, 16 Apr 2026 15:28:59 -0700 Subject: [PATCH 15/16] [codex] Split codex turn logic (#18206) ## Summary - Move Codex turn execution logic from `codex.rs` into `codex/turn.rs`. - Keep the existing crate-visible `run_turn`, `build_prompt`, `built_tools`, and `get_last_assistant_message_from_turn` surface re-exported from `codex.rs`. - Preserve test access for moved turn helpers while reducing the main `codex.rs` orchestration footprint. ## Stack - Base: #18200 (`pakrym/split-codex-handlers`) ## Testing - `CARGO_INCREMENTAL=0 cargo test -p codex-core --lib` - `just fix -p codex-core` - `just fmt` - `git diff --check` --- codex-rs/core/src/codex.rs | 2152 +----------------------------- codex-rs/core/src/codex/turn.rs | 2169 +++++++++++++++++++++++++++++++ 2 files changed, 2188 insertions(+), 2133 deletions(-) create mode 100644 codex-rs/core/src/codex/turn.rs diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index f356b3fff7..3d50f7735e 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -18,16 +18,10 @@ use crate::agent_identity::RegisteredAgentTask; use crate::apps::render_apps_section; use crate::commit_attribution::commit_message_trailer_instruction; use crate::compact; -use crate::compact::InitialContextInjection; -use crate::compact::run_inline_auto_compact_task; -use crate::compact::should_use_remote_compact_task; -use crate::compact_remote::run_inline_remote_auto_compact_task; use crate::config::ManagedFeatures; use crate::connectors; use crate::exec_policy::ExecPolicyManager; -use crate::hook_runtime::emit_hook_completed_events; use crate::installation_id::resolve_installation_id; -use crate::mcp_tool_exposure::build_mcp_tool_exposure; use crate::parse_turn_item; use crate::path_utils::normalize_for_native_workdir; use crate::realtime_conversation::RealtimeConversationManager; @@ -35,28 +29,13 @@ use crate::render_skills_section; use crate::rollout::find_thread_name_by_id; use crate::session_prefix::format_subagent_notification_message; use crate::skills_load_input_from_config; -use crate::stream_events_utils::HandleOutputCtx; -use crate::stream_events_utils::handle_non_tool_response_item; -use crate::stream_events_utils::handle_output_item_done; -use crate::stream_events_utils::last_assistant_message_from_item; -use crate::stream_events_utils::mark_thread_memory_mode_polluted_if_external_context; -use crate::stream_events_utils::raw_assistant_output_text_from_item; -use crate::stream_events_utils::record_completed_response_item; use crate::turn_metadata::TurnMetadataState; -use crate::unavailable_tool::collect_unavailable_called_tools; -use crate::util::error_or_panic; use async_channel::Receiver; use async_channel::Sender; use chrono::Local; use chrono::Utc; use codex_analytics::AnalyticsEventsClient; -use codex_analytics::AppInvocation; -use codex_analytics::CompactionPhase; -use codex_analytics::CompactionReason; -use codex_analytics::InvocationType; use codex_analytics::SubAgentThreadStartedInput; -use codex_analytics::TurnResolvedConfigFact; -use codex_analytics::build_track_events_context; use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::McpServerElicitationRequest; use codex_app_server_protocol::McpServerElicitationRequestParams; @@ -67,10 +46,6 @@ use codex_exec_server::FileSystemSandboxContext; use codex_features::FEATURES; use codex_features::Feature; use codex_features::unstable_features_warning_event; -use codex_hooks::HookEvent; -use codex_hooks::HookEventAfterAgent; -use codex_hooks::HookPayload; -use codex_hooks::HookResult; use codex_hooks::Hooks; use codex_hooks::HooksConfig; use codex_login::AuthManager; @@ -102,10 +77,8 @@ use codex_protocol::config_types::Settings; use codex_protocol::config_types::WebSearchMode; use codex_protocol::dynamic_tools::DynamicToolResponse; use codex_protocol::dynamic_tools::DynamicToolSpec; -use codex_protocol::items::PlanItem; use codex_protocol::items::TurnItem; use codex_protocol::items::UserMessageItem; -use codex_protocol::items::build_hook_prompt_message; use codex_protocol::mcp::CallToolResult; use codex_protocol::models::BaseInstructions; use codex_protocol::models::PermissionProfile; @@ -140,17 +113,10 @@ use codex_rollout::state_db; use codex_shell_command::parse_command::parse_command; use codex_terminal_detection::user_agent; use codex_thread_store::LocalThreadStore; -use codex_tools::filter_tool_suggest_discoverable_tools_for_client; use codex_utils_output_truncation::TruncationPolicy; -use codex_utils_stream_parser::AssistantTextChunk; -use codex_utils_stream_parser::AssistantTextStreamParser; -use codex_utils_stream_parser::ProposedPlanSegment; -use codex_utils_stream_parser::extract_proposed_plan_text; -use codex_utils_stream_parser::strip_citations; use futures::future::BoxFuture; use futures::future::Shared; use futures::prelude::*; -use futures::stream::FuturesOrdered; use rmcp::model::ListResourceTemplatesResult; use rmcp::model::ListResourcesResult; use rmcp::model::PaginatedRequestParams; @@ -168,19 +134,13 @@ use toml::Value as TomlValue; use tracing::Instrument; use tracing::debug; use tracing::error; -use tracing::field; use tracing::info; use tracing::info_span; use tracing::instrument; -use tracing::trace; -use tracing::trace_span; use tracing::warn; use uuid::Uuid; use crate::client::ModelClient; -use crate::client::ModelClientSession; -use crate::client_common::Prompt; -use crate::client_common::ResponseEvent; use crate::codex_thread::ThreadConfigSnapshot; use crate::compact::collect_user_messages; use crate::config::Config; @@ -204,10 +164,22 @@ use codex_protocol::exec_output::StreamOutput; mod handlers; mod rollout_reconstruction; +mod turn; #[cfg(test)] use self::handlers::submission_dispatch_span; use self::handlers::submission_loop; #[cfg(test)] +use self::turn::AssistantMessageStreamParsers; +pub(crate) use self::turn::build_prompt; +pub(crate) use self::turn::built_tools; +#[cfg(test)] +use self::turn::collect_explicit_app_ids_from_skill_items; +#[cfg(test)] +use self::turn::filter_connectors_for_input; +pub(crate) use self::turn::get_last_assistant_message_from_turn; +use self::turn::realtime_text_for_event; +pub(crate) use self::turn::run_turn; +#[cfg(test)] mod rollout_reconstruction_tests; #[derive(Debug, PartialEq)] @@ -262,40 +234,18 @@ pub(crate) struct PreviousTurnSettings { } use crate::SkillError; -use crate::SkillInjections; use crate::SkillLoadOutcome; use crate::SkillMetadata; use crate::SkillsManager; use crate::agents_md::AgentsMdManager; -use crate::build_skill_injections; -use crate::collect_env_var_dependencies; -use crate::collect_explicit_skill_mentions; use crate::exec_policy::ExecPolicyUpdateError; -use crate::feedback_tags; use crate::guardian::GuardianReviewSessionManager; -use crate::hook_runtime::PendingInputHookDisposition; -use crate::hook_runtime::inspect_pending_input; -use crate::hook_runtime::record_additional_contexts; -use crate::hook_runtime::record_pending_input; -use crate::hook_runtime::run_pending_session_start_hooks; -use crate::hook_runtime::run_user_prompt_submit_hooks; -use crate::injection::ToolMentionKind; -use crate::injection::app_id_from_path; -use crate::injection::tool_kind_for_path; use crate::instructions::UserInstructions; use crate::mcp::McpManager; -use crate::mcp_skill_dependencies::maybe_prompt_and_install_mcp_dependencies; use crate::memories; -use crate::mentions::build_connector_slug_counts; -use crate::mentions::build_skill_name_counts; -use crate::mentions::collect_explicit_app_ids; -use crate::mentions::collect_explicit_plugin_mentions; -use crate::mentions::collect_tool_mentions_from_messages; use crate::network_policy_decision::execpolicy_network_rule_amendment; use crate::plugins::PluginsManager; -use crate::plugins::build_plugin_injections; use crate::plugins::render_plugins_section; -use crate::resolve_skill_dependencies_for_turn; use crate::rollout::RolloutRecorder; use crate::rollout::RolloutRecorderParams; use crate::rollout::map_session_init_error; @@ -310,28 +260,26 @@ use crate::state::ActiveTurn; use crate::state::MailboxDeliveryPhase; use crate::state::SessionServices; use crate::state::SessionState; +#[cfg(test)] +use crate::stream_events_utils::HandleOutputCtx; +#[cfg(test)] +use crate::stream_events_utils::handle_output_item_done; use crate::tasks::GhostSnapshotTask; use crate::tasks::ReviewTask; use crate::tasks::SessionTask; use crate::tasks::SessionTaskContext; -use crate::tools::ToolRouter; -use crate::tools::context::SharedTurnDiffTracker; use crate::tools::js_repl::JsReplHandle; use crate::tools::js_repl::resolve_compatible_node; use crate::tools::network_approval::NetworkApprovalService; use crate::tools::network_approval::build_blocked_request_observer; use crate::tools::network_approval::build_network_policy_decider; +#[cfg(test)] use crate::tools::parallel::ToolCallRuntime; -use crate::tools::router::ToolRouterParams; use crate::tools::sandboxing::ApprovalStore; -use crate::turn_diff_tracker::TurnDiffTracker; use crate::turn_timing::TurnTimingState; use crate::turn_timing::record_turn_ttfm_metric; -use crate::turn_timing::record_turn_ttft_metric; use crate::unified_exec::UnifiedExecProcessManager; -use crate::util::backoff; use crate::windows_sandbox::WindowsSandboxLevelExt; -use codex_async_utils::OrCancelExt; use codex_git_utils::get_git_repo_root; use codex_mcp::compute_auth_statuses; use codex_mcp::with_codex_apps_mcp; @@ -348,8 +296,6 @@ use codex_protocol::models::DeveloperInstructions; use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; -use codex_protocol::protocol::AgentMessageContentDeltaEvent; -use codex_protocol::protocol::AgentReasoningSectionBreakEvent; use codex_protocol::protocol::ApplyPatchApprovalRequestEvent; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::BackgroundEventEvent; @@ -367,10 +313,7 @@ use codex_protocol::protocol::ModelRerouteReason; use codex_protocol::protocol::NetworkApprovalContext; use codex_protocol::protocol::NonSteerableTurnKind; use codex_protocol::protocol::Op; -use codex_protocol::protocol::PlanDeltaEvent; use codex_protocol::protocol::RateLimitSnapshot; -use codex_protocol::protocol::ReasoningContentDeltaEvent; -use codex_protocol::protocol::ReasoningRawContentDeltaEvent; use codex_protocol::protocol::RequestUserInputEvent; use codex_protocol::protocol::ReviewDecision; use codex_protocol::protocol::SandboxPolicy; @@ -386,7 +329,6 @@ use codex_protocol::protocol::Submission; use codex_protocol::protocol::TokenCountEvent; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsageInfo; -use codex_protocol::protocol::TurnDiffEvent; use codex_protocol::protocol::WarningEvent; use codex_protocol::user_input::UserInput; use codex_tools::ToolsConfig; @@ -394,6 +336,8 @@ use codex_tools::ToolsConfigParams; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_readiness::Readiness; use codex_utils_readiness::ReadinessFlag; +#[cfg(test)] +use codex_utils_stream_parser::ProposedPlanSegment; fn image_generation_tool_auth_allowed(auth_manager: Option<&AuthManager>) -> bool { matches!( @@ -5032,2064 +4976,6 @@ fn errors_to_info(errors: &[SkillError]) -> Vec { .collect() } -/// Takes a user message as input and runs a loop where, at each sampling request, the model -/// replies with either: -/// -/// - requested function calls -/// - an assistant message -/// -/// While it is possible for the model to return multiple of these items in a -/// single sampling request, in practice, we generally one item per sampling request: -/// -/// - If the model requests a function call, we execute it and send the output -/// back to the model in the next sampling request. -/// - If the model sends only an assistant message, we record it in the -/// conversation history and consider the turn complete. -/// -pub(crate) async fn run_turn( - sess: Arc, - turn_context: Arc, - input: Vec, - prewarmed_client_session: Option, - cancellation_token: CancellationToken, -) -> Option { - if input.is_empty() && !sess.has_pending_input().await { - return None; - } - - let model_info = turn_context.model_info.clone(); - let auto_compact_limit = model_info.auto_compact_token_limit().unwrap_or(i64::MAX); - let mut prewarmed_client_session = prewarmed_client_session; - // TODO(ccunningham): Pre-turn compaction runs before context updates and the - // new user message are recorded. Estimate pending incoming items (context - // diffs/full reinjection + user input) and trigger compaction preemptively - // when they would push the thread over the compaction threshold. - let pre_sampling_compacted = match run_pre_sampling_compact(&sess, &turn_context).await { - Ok(pre_sampling_compacted) => pre_sampling_compacted, - Err(_) => { - error!("Failed to run pre-sampling compact"); - return None; - } - }; - if pre_sampling_compacted && let Some(mut client_session) = prewarmed_client_session.take() { - client_session.reset_websocket_session(); - } - - let skills_outcome = Some(turn_context.turn_skills.outcome.as_ref()); - - sess.record_context_updates_and_set_reference_context_item(turn_context.as_ref()) - .await; - - let loaded_plugins = sess - .services - .plugins_manager - .plugins_for_config(&turn_context.config) - .await; - // Structured plugin:// mentions are resolved from the current session's - // enabled plugins, then converted into turn-scoped guidance below. - let mentioned_plugins = - collect_explicit_plugin_mentions(&input, loaded_plugins.capability_summaries()); - let mcp_tools = if turn_context.apps_enabled() || !mentioned_plugins.is_empty() { - // Plugin mentions need raw MCP/app inventory even when app tools - // are normally hidden so we can describe the plugin's currently - // usable capabilities for this turn. - match sess - .services - .mcp_connection_manager - .read() - .await - .list_all_tools() - .or_cancel(&cancellation_token) - .await - { - Ok(mcp_tools) => mcp_tools, - Err(_) if turn_context.apps_enabled() => return None, - Err(_) => HashMap::new(), - } - } else { - HashMap::new() - }; - let available_connectors = if turn_context.apps_enabled() { - let connectors = codex_connectors::merge::merge_plugin_connectors_with_accessible( - loaded_plugins - .effective_apps() - .into_iter() - .map(|connector_id| connector_id.0), - connectors::accessible_connectors_from_mcp_tools(&mcp_tools), - ); - connectors::with_app_enabled_state(connectors, &turn_context.config) - } else { - Vec::new() - }; - let connector_slug_counts = build_connector_slug_counts(&available_connectors); - let skill_name_counts_lower = skills_outcome - .as_ref() - .map_or_else(HashMap::new, |outcome| { - build_skill_name_counts(&outcome.skills, &outcome.disabled_paths).1 - }); - let mentioned_skills = skills_outcome.as_ref().map_or_else(Vec::new, |outcome| { - collect_explicit_skill_mentions( - &input, - &outcome.skills, - &outcome.disabled_paths, - &connector_slug_counts, - ) - }); - let config = turn_context.config.clone(); - if config - .features - .enabled(Feature::SkillEnvVarDependencyPrompt) - { - let env_var_dependencies = collect_env_var_dependencies(&mentioned_skills); - resolve_skill_dependencies_for_turn(&sess, &turn_context, &env_var_dependencies).await; - } - - maybe_prompt_and_install_mcp_dependencies( - sess.as_ref(), - turn_context.as_ref(), - &cancellation_token, - &mentioned_skills, - ) - .await; - - let session_telemetry = turn_context.session_telemetry.clone(); - let thread_id = sess.conversation_id.to_string(); - let tracking = build_track_events_context( - turn_context.model_info.slug.clone(), - thread_id, - turn_context.sub_id.clone(), - ); - let SkillInjections { - items: skill_items, - warnings: skill_warnings, - } = build_skill_injections( - &mentioned_skills, - skills_outcome, - Some(&session_telemetry), - &sess.services.analytics_events_client, - tracking.clone(), - ) - .await; - - for message in skill_warnings { - sess.send_event(&turn_context, EventMsg::Warning(WarningEvent { message })) - .await; - } - - let plugin_items = - build_plugin_injections(&mentioned_plugins, &mcp_tools, &available_connectors); - let mentioned_plugin_metadata = mentioned_plugins - .iter() - .filter_map(crate::plugins::PluginCapabilitySummary::telemetry_metadata) - .collect::>(); - - let mut explicitly_enabled_connectors = collect_explicit_app_ids(&input); - explicitly_enabled_connectors.extend(collect_explicit_app_ids_from_skill_items( - &skill_items, - &available_connectors, - &skill_name_counts_lower, - )); - let connector_names_by_id = available_connectors - .iter() - .map(|connector| (connector.id.as_str(), connector.name.as_str())) - .collect::>(); - let mentioned_app_invocations = explicitly_enabled_connectors - .iter() - .map(|connector_id| AppInvocation { - connector_id: Some(connector_id.clone()), - app_name: connector_names_by_id - .get(connector_id.as_str()) - .map(|name| (*name).to_string()), - invocation_type: Some(InvocationType::Explicit), - }) - .collect::>(); - - if run_pending_session_start_hooks(&sess, &turn_context).await { - return None; - } - let additional_contexts = if input.is_empty() { - Vec::new() - } else { - let initial_input_for_turn: ResponseInputItem = ResponseInputItem::from(input.clone()); - let response_item: ResponseItem = initial_input_for_turn.clone().into(); - let user_prompt_submit_outcome = run_user_prompt_submit_hooks( - &sess, - &turn_context, - UserMessageItem::new(&input).message(), - ) - .await; - if user_prompt_submit_outcome.should_stop { - record_additional_contexts( - &sess, - &turn_context, - user_prompt_submit_outcome.additional_contexts, - ) - .await; - return None; - } - sess.record_user_prompt_and_emit_turn_item(turn_context.as_ref(), &input, response_item) - .await; - user_prompt_submit_outcome.additional_contexts - }; - sess.services - .analytics_events_client - .track_app_mentioned(tracking.clone(), mentioned_app_invocations); - for plugin in mentioned_plugin_metadata { - sess.services - .analytics_events_client - .track_plugin_used(tracking.clone(), plugin); - } - sess.merge_connector_selection(explicitly_enabled_connectors.clone()) - .await; - record_additional_contexts(&sess, &turn_context, additional_contexts).await; - if !input.is_empty() { - // Track the previous-turn baseline from the regular user-turn path only so - // standalone tasks (compact/shell/review/undo) cannot suppress future - // model/realtime injections. - sess.set_previous_turn_settings(Some(PreviousTurnSettings { - model: turn_context.model_info.slug.clone(), - realtime_active: Some(turn_context.realtime_active), - })) - .await; - } - if let Err(error) = sess.ensure_agent_task_registered().await { - warn!(error = %error, "agent task registration failed"); - sess.send_event( - turn_context.as_ref(), - EventMsg::Error(ErrorEvent { - message: format!( - "Agent task registration failed. Please try again; Codex will attempt to register the task again on the next turn: {error}" - ), - codex_error_info: Some(CodexErrorInfo::Other), - }), - ) - .await; - return None; - } - - if !skill_items.is_empty() { - sess.record_conversation_items(&turn_context, &skill_items) - .await; - } - if !plugin_items.is_empty() { - sess.record_conversation_items(&turn_context, &plugin_items) - .await; - } - - track_turn_resolved_config_analytics(&sess, &turn_context, &input).await; - - let skills_outcome = Some(turn_context.turn_skills.outcome.as_ref()); - sess.maybe_start_ghost_snapshot(Arc::clone(&turn_context), cancellation_token.child_token()) - .await; - let mut last_agent_message: Option = None; - let mut stop_hook_active = false; - // Although from the perspective of codex.rs, TurnDiffTracker has the lifecycle of a Task which contains - // many turns, from the perspective of the user, it is a single turn. - let turn_diff_tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); - let mut server_model_warning_emitted_for_turn = false; - - // `ModelClientSession` is turn-scoped and caches WebSocket + sticky routing state, so we reuse - // one instance across retries within this turn. - let mut client_session = - prewarmed_client_session.unwrap_or_else(|| sess.services.model_client.new_session()); - // Pending input is drained into history before building the next model request. - // However, we defer that drain until after sampling in two cases: - // 1. At the start of a turn, so the fresh user prompt in `input` gets sampled first. - // 2. After auto-compact, when model/tool continuation needs to resume before any steer. - let mut can_drain_pending_input = input.is_empty(); - - loop { - if run_pending_session_start_hooks(&sess, &turn_context).await { - break; - } - - // Note that pending_input would be something like a message the user - // submitted through the UI while the model was running. Though the UI - // may support this, the model might not. - let pending_input = if can_drain_pending_input { - sess.get_pending_input().await - } else { - Vec::new() - }; - - let mut blocked_pending_input = false; - let mut blocked_pending_input_contexts = Vec::new(); - let mut requeued_pending_input = false; - let mut accepted_pending_input = Vec::new(); - if !pending_input.is_empty() { - let mut pending_input_iter = pending_input.into_iter(); - while let Some(pending_input_item) = pending_input_iter.next() { - match inspect_pending_input(&sess, &turn_context, pending_input_item).await { - PendingInputHookDisposition::Accepted(pending_input) => { - accepted_pending_input.push(*pending_input); - } - PendingInputHookDisposition::Blocked { - additional_contexts, - } => { - let remaining_pending_input = pending_input_iter.collect::>(); - if !remaining_pending_input.is_empty() { - let _ = sess.prepend_pending_input(remaining_pending_input).await; - requeued_pending_input = true; - } - blocked_pending_input_contexts = additional_contexts; - blocked_pending_input = true; - break; - } - } - } - } - - let has_accepted_pending_input = !accepted_pending_input.is_empty(); - for pending_input in accepted_pending_input { - record_pending_input(&sess, &turn_context, pending_input).await; - } - record_additional_contexts(&sess, &turn_context, blocked_pending_input_contexts).await; - - if blocked_pending_input && !has_accepted_pending_input { - if requeued_pending_input { - continue; - } - break; - } - - // Construct the input that we will send to the model. - let sampling_request_input: Vec = { - sess.clone_history() - .await - .for_prompt(&turn_context.model_info.input_modalities) - }; - - let sampling_request_input_messages = sampling_request_input - .iter() - .filter_map(|item| match parse_turn_item(item) { - Some(TurnItem::UserMessage(user_message)) => Some(user_message), - _ => None, - }) - .map(|user_message| user_message.message()) - .collect::>(); - let turn_metadata_header = turn_context.turn_metadata_state.current_header_value(); - match run_sampling_request( - Arc::clone(&sess), - Arc::clone(&turn_context), - Arc::clone(&turn_diff_tracker), - &mut client_session, - turn_metadata_header.as_deref(), - sampling_request_input, - &explicitly_enabled_connectors, - skills_outcome, - &mut server_model_warning_emitted_for_turn, - cancellation_token.child_token(), - ) - .await - { - Ok(sampling_request_output) => { - let SamplingRequestResult { - needs_follow_up: model_needs_follow_up, - last_agent_message: sampling_request_last_agent_message, - } = sampling_request_output; - can_drain_pending_input = true; - let has_pending_input = sess.has_pending_input().await; - let needs_follow_up = model_needs_follow_up || has_pending_input; - let total_usage_tokens = sess.get_total_token_usage().await; - let token_limit_reached = total_usage_tokens >= auto_compact_limit; - - let estimated_token_count = - sess.get_estimated_token_count(turn_context.as_ref()).await; - - trace!( - turn_id = %turn_context.sub_id, - total_usage_tokens, - estimated_token_count = ?estimated_token_count, - auto_compact_limit, - token_limit_reached, - model_needs_follow_up, - has_pending_input, - needs_follow_up, - "post sampling token usage" - ); - - // as long as compaction works well in getting us way below the token limit, we shouldn't worry about being in an infinite loop. - if token_limit_reached && needs_follow_up { - if run_auto_compact( - &sess, - &turn_context, - InitialContextInjection::BeforeLastUserMessage, - CompactionReason::ContextLimit, - CompactionPhase::MidTurn, - ) - .await - .is_err() - { - return None; - } - client_session.reset_websocket_session(); - can_drain_pending_input = !model_needs_follow_up; - continue; - } - - if !needs_follow_up { - last_agent_message = sampling_request_last_agent_message; - let stop_hook_permission_mode = match turn_context.approval_policy.value() { - AskForApproval::Never => "bypassPermissions", - AskForApproval::UnlessTrusted - | AskForApproval::OnFailure - | AskForApproval::OnRequest - | AskForApproval::Granular(_) => "default", - } - .to_string(); - let stop_request = codex_hooks::StopRequest { - session_id: sess.conversation_id, - turn_id: turn_context.sub_id.clone(), - cwd: turn_context.cwd.clone(), - transcript_path: sess.hook_transcript_path().await, - model: turn_context.model_info.slug.clone(), - permission_mode: stop_hook_permission_mode, - stop_hook_active, - last_assistant_message: last_agent_message.clone(), - }; - for run in sess.hooks().preview_stop(&stop_request) { - sess.send_event( - &turn_context, - EventMsg::HookStarted(codex_protocol::protocol::HookStartedEvent { - turn_id: Some(turn_context.sub_id.clone()), - run, - }), - ) - .await; - } - let stop_outcome = sess.hooks().run_stop(stop_request).await; - emit_hook_completed_events(&sess, &turn_context, stop_outcome.hook_events) - .await; - if stop_outcome.should_block { - if let Some(hook_prompt_message) = - build_hook_prompt_message(&stop_outcome.continuation_fragments) - { - sess.record_conversation_items( - &turn_context, - std::slice::from_ref(&hook_prompt_message), - ) - .await; - stop_hook_active = true; - continue; - } else { - sess.send_event( - &turn_context, - EventMsg::Warning(WarningEvent { - message: "Stop hook requested continuation without a prompt; ignoring the block.".to_string(), - }), - ) - .await; - } - } - if stop_outcome.should_stop { - break; - } - let hook_outcomes = sess - .hooks() - .dispatch(HookPayload { - session_id: sess.conversation_id, - cwd: turn_context.cwd.clone(), - client: turn_context.app_server_client_name.clone(), - triggered_at: chrono::Utc::now(), - hook_event: HookEvent::AfterAgent { - event: HookEventAfterAgent { - thread_id: sess.conversation_id, - turn_id: turn_context.sub_id.clone(), - input_messages: sampling_request_input_messages, - last_assistant_message: last_agent_message.clone(), - }, - }, - }) - .await; - - let mut abort_message = None; - for hook_outcome in hook_outcomes { - let hook_name = hook_outcome.hook_name; - match hook_outcome.result { - HookResult::Success => {} - HookResult::FailedContinue(error) => { - warn!( - turn_id = %turn_context.sub_id, - hook_name = %hook_name, - error = %error, - "after_agent hook failed; continuing" - ); - } - HookResult::FailedAbort(error) => { - let message = format!( - "after_agent hook '{hook_name}' failed and aborted turn completion: {error}" - ); - warn!( - turn_id = %turn_context.sub_id, - hook_name = %hook_name, - error = %error, - "after_agent hook failed; aborting operation" - ); - if abort_message.is_none() { - abort_message = Some(message); - } - } - } - } - if let Some(message) = abort_message { - sess.send_event( - &turn_context, - EventMsg::Error(ErrorEvent { - message, - codex_error_info: None, - }), - ) - .await; - return None; - } - break; - } - continue; - } - Err(CodexErr::TurnAborted) => { - // Aborted turn is reported via a different event. - break; - } - Err(CodexErr::InvalidImageRequest()) => { - { - let mut state = sess.state.lock().await; - error_or_panic( - "Invalid image detected; sanitizing tool output to prevent poisoning", - ); - if state.history.replace_last_turn_images("Invalid image") { - continue; - } - } - - let event = EventMsg::Error(ErrorEvent { - message: "Invalid image in your last message. Please remove it and try again." - .to_string(), - codex_error_info: Some(CodexErrorInfo::BadRequest), - }); - sess.send_event(&turn_context, event).await; - break; - } - Err(e) => { - info!("Turn error: {e:#}"); - let event = EventMsg::Error(e.to_error_event(/*message_prefix*/ None)); - sess.send_event(&turn_context, event).await; - // let the user continue the conversation - break; - } - } - } - - last_agent_message -} - -async fn track_turn_resolved_config_analytics( - sess: &Session, - turn_context: &TurnContext, - input: &[UserInput], -) { - if !sess.enabled(Feature::GeneralAnalytics) { - return; - } - - let thread_config = { - let state = sess.state.lock().await; - state.session_configuration.thread_config_snapshot() - }; - let is_first_turn = { - let mut state = sess.state.lock().await; - state.take_next_turn_is_first() - }; - sess.services - .analytics_events_client - .track_turn_resolved_config(TurnResolvedConfigFact { - turn_id: turn_context.sub_id.clone(), - thread_id: sess.conversation_id.to_string(), - num_input_images: input - .iter() - .filter(|item| { - matches!(item, UserInput::Image { .. } | UserInput::LocalImage { .. }) - }) - .count(), - submission_type: None, - ephemeral: thread_config.ephemeral, - session_source: thread_config.session_source, - model: turn_context.model_info.slug.clone(), - model_provider: turn_context.config.model_provider_id.clone(), - sandbox_policy: turn_context.sandbox_policy.get().clone(), - reasoning_effort: turn_context.reasoning_effort, - reasoning_summary: Some(turn_context.reasoning_summary), - service_tier: turn_context.config.service_tier, - approval_policy: turn_context.approval_policy.value(), - approvals_reviewer: turn_context.config.approvals_reviewer, - sandbox_network_access: turn_context.network_sandbox_policy.is_enabled(), - collaboration_mode: turn_context.collaboration_mode.mode, - personality: turn_context.personality, - is_first_turn, - }); -} - -async fn run_pre_sampling_compact( - sess: &Arc, - turn_context: &Arc, -) -> CodexResult { - let total_usage_tokens_before_compaction = sess.get_total_token_usage().await; - let mut pre_sampling_compacted = maybe_run_previous_model_inline_compact( - sess, - turn_context, - total_usage_tokens_before_compaction, - ) - .await?; - let total_usage_tokens = sess.get_total_token_usage().await; - let auto_compact_limit = turn_context - .model_info - .auto_compact_token_limit() - .unwrap_or(i64::MAX); - // Compact if the total usage tokens are greater than the auto compact limit - if total_usage_tokens >= auto_compact_limit { - run_auto_compact( - sess, - turn_context, - InitialContextInjection::DoNotInject, - CompactionReason::ContextLimit, - CompactionPhase::PreTurn, - ) - .await?; - pre_sampling_compacted = true; - } - Ok(pre_sampling_compacted) -} - -/// Runs pre-sampling compaction against the previous model when switching to a smaller -/// context-window model. -/// -/// Returns `Ok(true)` when compaction ran successfully, `Ok(false)` when compaction was skipped -/// because the model/context-window preconditions were not met, and `Err(_)` only when compaction -/// was attempted and failed. -async fn maybe_run_previous_model_inline_compact( - sess: &Arc, - turn_context: &Arc, - total_usage_tokens: i64, -) -> CodexResult { - let Some(previous_turn_settings) = sess.previous_turn_settings().await else { - return Ok(false); - }; - let previous_model_turn_context = Arc::new( - turn_context - .with_model(previous_turn_settings.model, &sess.services.models_manager) - .await, - ); - - let Some(old_context_window) = previous_model_turn_context.model_context_window() else { - return Ok(false); - }; - let Some(new_context_window) = turn_context.model_context_window() else { - return Ok(false); - }; - let new_auto_compact_limit = turn_context - .model_info - .auto_compact_token_limit() - .unwrap_or(i64::MAX); - let should_run = total_usage_tokens > new_auto_compact_limit - && previous_model_turn_context.model_info.slug != turn_context.model_info.slug - && old_context_window > new_context_window; - if should_run { - run_auto_compact( - sess, - &previous_model_turn_context, - InitialContextInjection::DoNotInject, - CompactionReason::ModelDownshift, - CompactionPhase::PreTurn, - ) - .await?; - return Ok(true); - } - Ok(false) -} - -async fn run_auto_compact( - sess: &Arc, - turn_context: &Arc, - initial_context_injection: InitialContextInjection, - reason: CompactionReason, - phase: CompactionPhase, -) -> CodexResult<()> { - if should_use_remote_compact_task(&turn_context.provider) { - run_inline_remote_auto_compact_task( - Arc::clone(sess), - Arc::clone(turn_context), - initial_context_injection, - reason, - phase, - ) - .await?; - } else { - run_inline_auto_compact_task( - Arc::clone(sess), - Arc::clone(turn_context), - initial_context_injection, - reason, - phase, - ) - .await?; - } - Ok(()) -} - -fn collect_explicit_app_ids_from_skill_items( - skill_items: &[ResponseItem], - connectors: &[connectors::AppInfo], - skill_name_counts_lower: &HashMap, -) -> HashSet { - if skill_items.is_empty() || connectors.is_empty() { - return HashSet::new(); - } - - let skill_messages = skill_items - .iter() - .filter_map(|item| match item { - ResponseItem::Message { content, .. } => { - content.iter().find_map(|content_item| match content_item { - ContentItem::InputText { text } => Some(text.clone()), - _ => None, - }) - } - _ => None, - }) - .collect::>(); - if skill_messages.is_empty() { - return HashSet::new(); - } - - let mentions = collect_tool_mentions_from_messages(&skill_messages); - let mention_names_lower = mentions - .plain_names - .iter() - .map(|name| name.to_ascii_lowercase()) - .collect::>(); - let mut connector_ids = mentions - .paths - .iter() - .filter(|path| tool_kind_for_path(path) == ToolMentionKind::App) - .filter_map(|path| app_id_from_path(path).map(str::to_string)) - .collect::>(); - - let connector_slug_counts = build_connector_slug_counts(connectors); - for connector in connectors { - let slug = codex_connectors::metadata::connector_mention_slug(connector); - let connector_count = connector_slug_counts.get(&slug).copied().unwrap_or(0); - let skill_count = skill_name_counts_lower.get(&slug).copied().unwrap_or(0); - if connector_count == 1 && skill_count == 0 && mention_names_lower.contains(&slug) { - connector_ids.insert(connector.id.clone()); - } - } - - connector_ids -} - -fn filter_connectors_for_input( - connectors: &[connectors::AppInfo], - input: &[ResponseItem], - explicitly_enabled_connectors: &HashSet, - skill_name_counts_lower: &HashMap, -) -> Vec { - let connectors: Vec = connectors - .iter() - .filter(|connector| connector.is_enabled) - .cloned() - .collect::>(); - if connectors.is_empty() { - return Vec::new(); - } - - let user_messages = collect_user_messages(input); - if user_messages.is_empty() && explicitly_enabled_connectors.is_empty() { - return Vec::new(); - } - - let mentions = collect_tool_mentions_from_messages(&user_messages); - let mention_names_lower = mentions - .plain_names - .iter() - .map(|name| name.to_ascii_lowercase()) - .collect::>(); - - let connector_slug_counts = build_connector_slug_counts(&connectors); - let mut allowed_connector_ids = explicitly_enabled_connectors.clone(); - for path in mentions - .paths - .iter() - .filter(|path| tool_kind_for_path(path) == ToolMentionKind::App) - { - if let Some(connector_id) = app_id_from_path(path) { - allowed_connector_ids.insert(connector_id.to_string()); - } - } - - connectors - .into_iter() - .filter(|connector| { - connector_inserted_in_messages( - connector, - &mention_names_lower, - &allowed_connector_ids, - &connector_slug_counts, - skill_name_counts_lower, - ) - }) - .collect() -} - -fn connector_inserted_in_messages( - connector: &connectors::AppInfo, - mention_names_lower: &HashSet, - allowed_connector_ids: &HashSet, - connector_slug_counts: &HashMap, - skill_name_counts_lower: &HashMap, -) -> bool { - if allowed_connector_ids.contains(&connector.id) { - return true; - } - - let mention_slug = codex_connectors::metadata::connector_mention_slug(connector); - let connector_count = connector_slug_counts - .get(&mention_slug) - .copied() - .unwrap_or(0); - let skill_count = skill_name_counts_lower - .get(&mention_slug) - .copied() - .unwrap_or(0); - connector_count == 1 && skill_count == 0 && mention_names_lower.contains(&mention_slug) -} - -pub(crate) fn build_prompt( - input: Vec, - router: &ToolRouter, - turn_context: &TurnContext, - base_instructions: BaseInstructions, -) -> Prompt { - let deferred_dynamic_tools = turn_context - .dynamic_tools - .iter() - .filter(|tool| tool.defer_loading) - .map(|tool| tool.name.as_str()) - .collect::>(); - let tools = if deferred_dynamic_tools.is_empty() { - router.model_visible_specs() - } else { - router - .model_visible_specs() - .into_iter() - .filter(|spec| !deferred_dynamic_tools.contains(spec.name())) - .collect() - }; - - Prompt { - input, - tools, - parallel_tool_calls: turn_context.model_info.supports_parallel_tool_calls, - base_instructions, - personality: turn_context.personality, - output_schema: turn_context.final_output_json_schema.clone(), - } -} - -#[allow(clippy::too_many_arguments)] -#[instrument(level = "trace", - skip_all, - fields( - turn_id = %turn_context.sub_id, - model = %turn_context.model_info.slug, - cwd = %turn_context.cwd.display() - ) -)] -async fn run_sampling_request( - sess: Arc, - turn_context: Arc, - turn_diff_tracker: SharedTurnDiffTracker, - client_session: &mut ModelClientSession, - turn_metadata_header: Option<&str>, - input: Vec, - explicitly_enabled_connectors: &HashSet, - skills_outcome: Option<&SkillLoadOutcome>, - server_model_warning_emitted_for_turn: &mut bool, - cancellation_token: CancellationToken, -) -> CodexResult { - let router = built_tools( - sess.as_ref(), - turn_context.as_ref(), - &input, - explicitly_enabled_connectors, - skills_outcome, - &cancellation_token, - ) - .await?; - - let base_instructions = sess.get_base_instructions().await; - - let tool_runtime = ToolCallRuntime::new( - Arc::clone(&router), - Arc::clone(&sess), - Arc::clone(&turn_context), - Arc::clone(&turn_diff_tracker), - ); - let _code_mode_worker = sess - .services - .code_mode_service - .start_turn_worker( - &sess, - &turn_context, - Arc::clone(&router), - Arc::clone(&turn_diff_tracker), - ) - .await; - let mut retries = 0; - let mut initial_input = Some(input); - loop { - let prompt_input = if let Some(input) = initial_input.take() { - input - } else { - sess.clone_history() - .await - .for_prompt(&turn_context.model_info.input_modalities) - }; - let prompt = build_prompt( - prompt_input, - router.as_ref(), - turn_context.as_ref(), - base_instructions.clone(), - ); - let err = match try_run_sampling_request( - tool_runtime.clone(), - Arc::clone(&sess), - Arc::clone(&turn_context), - client_session, - turn_metadata_header, - Arc::clone(&turn_diff_tracker), - server_model_warning_emitted_for_turn, - &prompt, - cancellation_token.child_token(), - ) - .await - { - Ok(output) => { - return Ok(output); - } - Err(CodexErr::ContextWindowExceeded) => { - sess.set_total_tokens_full(&turn_context).await; - return Err(CodexErr::ContextWindowExceeded); - } - Err(CodexErr::UsageLimitReached(e)) => { - let rate_limits = e.rate_limits.clone(); - if let Some(rate_limits) = rate_limits { - sess.update_rate_limits(&turn_context, *rate_limits).await; - } - return Err(CodexErr::UsageLimitReached(e)); - } - Err(err) => err, - }; - - if !err.is_retryable() { - return Err(err); - } - - // Use the configured provider-specific stream retry budget. - let max_retries = turn_context.provider.stream_max_retries(); - if retries >= max_retries - && client_session.try_switch_fallback_transport( - &turn_context.session_telemetry, - &turn_context.model_info, - ) - { - sess.send_event( - &turn_context, - EventMsg::Warning(WarningEvent { - message: format!("Falling back from WebSockets to HTTPS transport. {err:#}"), - }), - ) - .await; - retries = 0; - continue; - } - if retries < max_retries { - retries += 1; - let delay = match &err { - CodexErr::Stream(_, requested_delay) => { - requested_delay.unwrap_or_else(|| backoff(retries)) - } - _ => backoff(retries), - }; - warn!( - "stream disconnected - retrying sampling request ({retries}/{max_retries} in {delay:?})...", - ); - - // In release builds, hide the first websocket retry notification to reduce noisy - // transient reconnect messages. In debug builds, keep full visibility for diagnosis. - let report_error = retries > 1 - || cfg!(debug_assertions) - || !sess.services.model_client.responses_websocket_enabled(); - if report_error { - // Surface retry information to any UI/front‑end so the - // user understands what is happening instead of staring - // at a seemingly frozen screen. - sess.notify_stream_error( - &turn_context, - format!("Reconnecting... {retries}/{max_retries}"), - err, - ) - .await; - } - tokio::time::sleep(delay).await; - } else { - return Err(err); - } - } -} - -pub(crate) async fn built_tools( - sess: &Session, - turn_context: &TurnContext, - input: &[ResponseItem], - explicitly_enabled_connectors: &HashSet, - skills_outcome: Option<&SkillLoadOutcome>, - cancellation_token: &CancellationToken, -) -> CodexResult> { - let mcp_connection_manager = sess.services.mcp_connection_manager.read().await; - let has_mcp_servers = mcp_connection_manager.has_servers(); - let all_mcp_tools = mcp_connection_manager - .list_all_tools() - .or_cancel(cancellation_token) - .await?; - drop(mcp_connection_manager); - let loaded_plugins = sess - .services - .plugins_manager - .plugins_for_config(&turn_context.config) - .await; - - let mut effective_explicitly_enabled_connectors = explicitly_enabled_connectors.clone(); - effective_explicitly_enabled_connectors.extend(sess.get_connector_selection().await); - - let apps_enabled = turn_context.apps_enabled(); - let accessible_connectors = - apps_enabled.then(|| connectors::accessible_connectors_from_mcp_tools(&all_mcp_tools)); - let accessible_connectors_with_enabled_state = - accessible_connectors.as_ref().map(|connectors| { - connectors::with_app_enabled_state(connectors.clone(), &turn_context.config) - }); - let connectors = if apps_enabled { - let connectors = codex_connectors::merge::merge_plugin_connectors_with_accessible( - loaded_plugins - .effective_apps() - .into_iter() - .map(|connector_id| connector_id.0), - accessible_connectors.clone().unwrap_or_default(), - ); - Some(connectors::with_app_enabled_state( - connectors, - &turn_context.config, - )) - } else { - None - }; - let auth = sess.services.auth_manager.auth().await; - let discoverable_tools = if apps_enabled && turn_context.tools_config.tool_suggest { - if let Some(accessible_connectors) = accessible_connectors_with_enabled_state.as_ref() { - match connectors::list_tool_suggest_discoverable_tools_with_auth( - &turn_context.config, - auth.as_ref(), - accessible_connectors.as_slice(), - ) - .await - .map(|discoverable_tools| { - filter_tool_suggest_discoverable_tools_for_client( - discoverable_tools, - turn_context.app_server_client_name.as_deref(), - ) - }) { - Ok(discoverable_tools) if discoverable_tools.is_empty() => None, - Ok(discoverable_tools) => Some(discoverable_tools), - Err(err) => { - warn!("failed to load discoverable tool suggestions: {err:#}"); - None - } - } - } else { - None - } - } else { - None - }; - - let explicitly_enabled = if let Some(connectors) = connectors.as_ref() { - let skill_name_counts_lower = skills_outcome.map_or_else(HashMap::new, |outcome| { - build_skill_name_counts(&outcome.skills, &outcome.disabled_paths).1 - }); - - filter_connectors_for_input( - connectors, - input, - &effective_explicitly_enabled_connectors, - &skill_name_counts_lower, - ) - } else { - Vec::new() - }; - let mcp_tool_exposure = build_mcp_tool_exposure( - &all_mcp_tools, - connectors.as_deref(), - explicitly_enabled.as_slice(), - &turn_context.config, - &turn_context.tools_config, - ); - let mcp_tools = has_mcp_servers.then_some(mcp_tool_exposure.direct_tools); - let deferred_mcp_tools = mcp_tool_exposure.deferred_tools; - let unavailable_called_tools = if turn_context - .config - .features - .enabled(Feature::UnavailableDummyTools) - { - let exposed_tool_names = mcp_tools - .iter() - .chain(deferred_mcp_tools.iter()) - .flat_map(|tools| tools.keys().map(String::as_str)) - .collect::>(); - collect_unavailable_called_tools(input, &exposed_tool_names) - } else { - Vec::new() - }; - - let parallel_mcp_server_names = turn_context - .config - .mcp_servers - .get() - .iter() - .filter_map(|(server_name, server_config)| { - server_config - .supports_parallel_tool_calls - .then_some(server_name.clone()) - }) - .collect::>(); - - Ok(Arc::new(ToolRouter::from_config( - &turn_context.tools_config, - ToolRouterParams { - mcp_tools, - deferred_mcp_tools, - unavailable_called_tools, - parallel_mcp_server_names, - discoverable_tools, - dynamic_tools: turn_context.dynamic_tools.as_slice(), - }, - ))) -} - -#[derive(Debug)] -struct SamplingRequestResult { - needs_follow_up: bool, - last_agent_message: Option, -} - -/// Ephemeral per-response state for streaming a single proposed plan. -/// This is intentionally not persisted or stored in session/state since it -/// only exists while a response is actively streaming. The final plan text -/// is extracted from the completed assistant message. -/// Tracks a single proposed plan item across a streaming response. -struct ProposedPlanItemState { - item_id: String, - started: bool, - completed: bool, -} - -/// Aggregated state used only while streaming a plan-mode response. -/// Includes per-item parsers, deferred agent message bookkeeping, and the plan item lifecycle. -struct PlanModeStreamState { - /// Agent message items started by the model but deferred until we see non-plan text. - pending_agent_message_items: HashMap, - /// Agent message items whose start notification has been emitted. - started_agent_message_items: HashSet, - /// Leading whitespace buffered until we see non-whitespace text for an item. - leading_whitespace_by_item: HashMap, - /// Tracks plan item lifecycle while streaming plan output. - plan_item_state: ProposedPlanItemState, -} - -impl PlanModeStreamState { - fn new(turn_id: &str) -> Self { - Self { - pending_agent_message_items: HashMap::new(), - started_agent_message_items: HashSet::new(), - leading_whitespace_by_item: HashMap::new(), - plan_item_state: ProposedPlanItemState::new(turn_id), - } - } -} - -#[derive(Debug, Default)] -struct AssistantMessageStreamParsers { - plan_mode: bool, - parsers_by_item: HashMap, -} - -type ParsedAssistantTextDelta = AssistantTextChunk; - -impl AssistantMessageStreamParsers { - fn new(plan_mode: bool) -> Self { - Self { - plan_mode, - parsers_by_item: HashMap::new(), - } - } - - fn parser_mut(&mut self, item_id: &str) -> &mut AssistantTextStreamParser { - let plan_mode = self.plan_mode; - self.parsers_by_item - .entry(item_id.to_string()) - .or_insert_with(|| AssistantTextStreamParser::new(plan_mode)) - } - - fn seed_item_text(&mut self, item_id: &str, text: &str) -> ParsedAssistantTextDelta { - if text.is_empty() { - return ParsedAssistantTextDelta::default(); - } - self.parser_mut(item_id).push_str(text) - } - - fn parse_delta(&mut self, item_id: &str, delta: &str) -> ParsedAssistantTextDelta { - self.parser_mut(item_id).push_str(delta) - } - - fn finish_item(&mut self, item_id: &str) -> ParsedAssistantTextDelta { - let Some(mut parser) = self.parsers_by_item.remove(item_id) else { - return ParsedAssistantTextDelta::default(); - }; - parser.finish() - } - - fn drain_finished(&mut self) -> Vec<(String, ParsedAssistantTextDelta)> { - let parsers_by_item = std::mem::take(&mut self.parsers_by_item); - parsers_by_item - .into_iter() - .map(|(item_id, mut parser)| (item_id, parser.finish())) - .collect() - } -} - -impl ProposedPlanItemState { - fn new(turn_id: &str) -> Self { - Self { - item_id: format!("{turn_id}-plan"), - started: false, - completed: false, - } - } - - async fn start(&mut self, sess: &Session, turn_context: &TurnContext) { - if self.started || self.completed { - return; - } - self.started = true; - let item = TurnItem::Plan(PlanItem { - id: self.item_id.clone(), - text: String::new(), - }); - sess.emit_turn_item_started(turn_context, &item).await; - } - - async fn push_delta(&mut self, sess: &Session, turn_context: &TurnContext, delta: &str) { - if self.completed { - return; - } - if delta.is_empty() { - return; - } - let event = PlanDeltaEvent { - thread_id: sess.conversation_id.to_string(), - turn_id: turn_context.sub_id.clone(), - item_id: self.item_id.clone(), - delta: delta.to_string(), - }; - sess.send_event(turn_context, EventMsg::PlanDelta(event)) - .await; - } - - async fn complete_with_text( - &mut self, - sess: &Session, - turn_context: &TurnContext, - text: String, - ) { - if self.completed || !self.started { - return; - } - self.completed = true; - let item = TurnItem::Plan(PlanItem { - id: self.item_id.clone(), - text, - }); - sess.emit_turn_item_completed(turn_context, item).await; - } -} - -/// In plan mode we defer agent message starts until the parser emits non-plan -/// text. The parser buffers each line until it can rule out a tag prefix, so -/// plan-only outputs never show up as empty assistant messages. -async fn maybe_emit_pending_agent_message_start( - sess: &Session, - turn_context: &TurnContext, - state: &mut PlanModeStreamState, - item_id: &str, -) { - if state.started_agent_message_items.contains(item_id) { - return; - } - if let Some(item) = state.pending_agent_message_items.remove(item_id) { - sess.emit_turn_item_started(turn_context, &item).await; - state - .started_agent_message_items - .insert(item_id.to_string()); - } -} - -/// Agent messages are text-only today; concatenate all text entries. -fn agent_message_text(item: &codex_protocol::items::AgentMessageItem) -> String { - item.content - .iter() - .map(|entry| match entry { - codex_protocol::items::AgentMessageContent::Text { text } => text.as_str(), - }) - .collect() -} - -fn realtime_text_for_event(msg: &EventMsg) -> Option { - match msg { - EventMsg::AgentMessage(event) => Some(event.message.clone()), - EventMsg::ItemCompleted(event) => match &event.item { - TurnItem::AgentMessage(item) => Some(agent_message_text(item)), - _ => None, - }, - EventMsg::Error(_) - | EventMsg::Warning(_) - | EventMsg::RealtimeConversationStarted(_) - | EventMsg::RealtimeConversationSdp(_) - | EventMsg::RealtimeConversationRealtime(_) - | EventMsg::RealtimeConversationClosed(_) - | EventMsg::ModelReroute(_) - | EventMsg::ContextCompacted(_) - | EventMsg::ThreadRolledBack(_) - | EventMsg::TurnStarted(_) - | EventMsg::TurnComplete(_) - | EventMsg::TokenCount(_) - | EventMsg::UserMessage(_) - | EventMsg::AgentMessageDelta(_) - | EventMsg::AgentReasoning(_) - | EventMsg::AgentReasoningDelta(_) - | EventMsg::AgentReasoningRawContent(_) - | EventMsg::AgentReasoningRawContentDelta(_) - | EventMsg::AgentReasoningSectionBreak(_) - | EventMsg::SessionConfigured(_) - | EventMsg::ThreadNameUpdated(_) - | EventMsg::McpStartupUpdate(_) - | EventMsg::McpStartupComplete(_) - | EventMsg::McpToolCallBegin(_) - | EventMsg::McpToolCallEnd(_) - | EventMsg::WebSearchBegin(_) - | EventMsg::WebSearchEnd(_) - | EventMsg::ExecCommandBegin(_) - | EventMsg::ExecCommandOutputDelta(_) - | EventMsg::TerminalInteraction(_) - | EventMsg::ExecCommandEnd(_) - | EventMsg::PatchApplyBegin(_) - | EventMsg::PatchApplyEnd(_) - | EventMsg::ViewImageToolCall(_) - | EventMsg::ImageGenerationBegin(_) - | EventMsg::ImageGenerationEnd(_) - | EventMsg::ExecApprovalRequest(_) - | EventMsg::RequestPermissions(_) - | EventMsg::RequestUserInput(_) - | EventMsg::DynamicToolCallRequest(_) - | EventMsg::DynamicToolCallResponse(_) - | EventMsg::GuardianAssessment(_) - | EventMsg::ElicitationRequest(_) - | EventMsg::ApplyPatchApprovalRequest(_) - | EventMsg::DeprecationNotice(_) - | EventMsg::BackgroundEvent(_) - | EventMsg::UndoStarted(_) - | EventMsg::UndoCompleted(_) - | EventMsg::StreamError(_) - | EventMsg::TurnDiff(_) - | EventMsg::GetHistoryEntryResponse(_) - | EventMsg::McpListToolsResponse(_) - | EventMsg::ListSkillsResponse(_) - | EventMsg::RealtimeConversationListVoicesResponse(_) - | EventMsg::SkillsUpdateAvailable - | EventMsg::PlanUpdate(_) - | EventMsg::TurnAborted(_) - | EventMsg::ShutdownComplete - | EventMsg::EnteredReviewMode(_) - | EventMsg::ExitedReviewMode(_) - | EventMsg::RawResponseItem(_) - | EventMsg::ItemStarted(_) - | EventMsg::HookStarted(_) - | EventMsg::HookCompleted(_) - | EventMsg::AgentMessageContentDelta(_) - | EventMsg::PlanDelta(_) - | EventMsg::ReasoningContentDelta(_) - | EventMsg::ReasoningRawContentDelta(_) - | EventMsg::CollabAgentSpawnBegin(_) - | EventMsg::CollabAgentSpawnEnd(_) - | EventMsg::CollabAgentInteractionBegin(_) - | EventMsg::CollabAgentInteractionEnd(_) - | EventMsg::CollabWaitingBegin(_) - | EventMsg::CollabWaitingEnd(_) - | EventMsg::CollabCloseBegin(_) - | EventMsg::CollabCloseEnd(_) - | EventMsg::CollabResumeBegin(_) - | EventMsg::CollabResumeEnd(_) => None, - } -} - -/// Split the stream into normal assistant text vs. proposed plan content. -/// Normal text becomes AgentMessage deltas; plan content becomes PlanDelta + -/// TurnItem::Plan. -async fn handle_plan_segments( - sess: &Session, - turn_context: &TurnContext, - state: &mut PlanModeStreamState, - item_id: &str, - segments: Vec, -) { - for segment in segments { - match segment { - ProposedPlanSegment::Normal(delta) => { - if delta.is_empty() { - continue; - } - let has_non_whitespace = delta.chars().any(|ch| !ch.is_whitespace()); - if !has_non_whitespace && !state.started_agent_message_items.contains(item_id) { - let entry = state - .leading_whitespace_by_item - .entry(item_id.to_string()) - .or_default(); - entry.push_str(&delta); - continue; - } - let delta = if !state.started_agent_message_items.contains(item_id) { - if let Some(prefix) = state.leading_whitespace_by_item.remove(item_id) { - format!("{prefix}{delta}") - } else { - delta - } - } else { - delta - }; - maybe_emit_pending_agent_message_start(sess, turn_context, state, item_id).await; - - let event = AgentMessageContentDeltaEvent { - thread_id: sess.conversation_id.to_string(), - turn_id: turn_context.sub_id.clone(), - item_id: item_id.to_string(), - delta, - }; - sess.send_event(turn_context, EventMsg::AgentMessageContentDelta(event)) - .await; - } - ProposedPlanSegment::ProposedPlanStart => { - if !state.plan_item_state.completed { - state.plan_item_state.start(sess, turn_context).await; - } - } - ProposedPlanSegment::ProposedPlanDelta(delta) => { - if !state.plan_item_state.completed { - if !state.plan_item_state.started { - state.plan_item_state.start(sess, turn_context).await; - } - state - .plan_item_state - .push_delta(sess, turn_context, &delta) - .await; - } - } - ProposedPlanSegment::ProposedPlanEnd => {} - } - } -} - -async fn emit_streamed_assistant_text_delta( - sess: &Session, - turn_context: &TurnContext, - plan_mode_state: Option<&mut PlanModeStreamState>, - item_id: &str, - parsed: ParsedAssistantTextDelta, -) { - if parsed.is_empty() { - return; - } - if !parsed.citations.is_empty() { - // Citation extraction is intentionally local for now; we strip citations from display text - // but do not yet surface them in protocol events. - let _citations = parsed.citations; - } - if let Some(state) = plan_mode_state { - if !parsed.plan_segments.is_empty() { - handle_plan_segments(sess, turn_context, state, item_id, parsed.plan_segments).await; - } - return; - } - if parsed.visible_text.is_empty() { - return; - } - let event = AgentMessageContentDeltaEvent { - thread_id: sess.conversation_id.to_string(), - turn_id: turn_context.sub_id.clone(), - item_id: item_id.to_string(), - delta: parsed.visible_text, - }; - sess.send_event(turn_context, EventMsg::AgentMessageContentDelta(event)) - .await; -} - -/// Flush buffered assistant text parser state when an assistant message item ends. -async fn flush_assistant_text_segments_for_item( - sess: &Session, - turn_context: &TurnContext, - plan_mode_state: Option<&mut PlanModeStreamState>, - parsers: &mut AssistantMessageStreamParsers, - item_id: &str, -) { - let parsed = parsers.finish_item(item_id); - emit_streamed_assistant_text_delta(sess, turn_context, plan_mode_state, item_id, parsed).await; -} - -/// Flush any remaining buffered assistant text parser state at response completion. -async fn flush_assistant_text_segments_all( - sess: &Session, - turn_context: &TurnContext, - mut plan_mode_state: Option<&mut PlanModeStreamState>, - parsers: &mut AssistantMessageStreamParsers, -) { - for (item_id, parsed) in parsers.drain_finished() { - emit_streamed_assistant_text_delta( - sess, - turn_context, - plan_mode_state.as_deref_mut(), - &item_id, - parsed, - ) - .await; - } -} - -/// Emit completion for plan items by parsing the finalized assistant message. -async fn maybe_complete_plan_item_from_message( - sess: &Session, - turn_context: &TurnContext, - state: &mut PlanModeStreamState, - item: &ResponseItem, -) { - if let ResponseItem::Message { role, content, .. } = item - && role == "assistant" - { - let mut text = String::new(); - for entry in content { - if let ContentItem::OutputText { text: chunk } = entry { - text.push_str(chunk); - } - } - if let Some(plan_text) = extract_proposed_plan_text(&text) { - let (plan_text, _citations) = strip_citations(&plan_text); - if !state.plan_item_state.started { - state.plan_item_state.start(sess, turn_context).await; - } - state - .plan_item_state - .complete_with_text(sess, turn_context, plan_text) - .await; - } - } -} - -/// Emit a completed agent message in plan mode, respecting deferred starts. -async fn emit_agent_message_in_plan_mode( - sess: &Session, - turn_context: &TurnContext, - agent_message: codex_protocol::items::AgentMessageItem, - state: &mut PlanModeStreamState, -) { - let agent_message_id = agent_message.id.clone(); - let text = agent_message_text(&agent_message); - if text.trim().is_empty() { - state.pending_agent_message_items.remove(&agent_message_id); - state.started_agent_message_items.remove(&agent_message_id); - return; - } - - maybe_emit_pending_agent_message_start(sess, turn_context, state, &agent_message_id).await; - - if !state - .started_agent_message_items - .contains(&agent_message_id) - { - let start_item = state - .pending_agent_message_items - .remove(&agent_message_id) - .unwrap_or_else(|| { - TurnItem::AgentMessage(codex_protocol::items::AgentMessageItem { - id: agent_message_id.clone(), - content: Vec::new(), - phase: None, - memory_citation: None, - }) - }); - sess.emit_turn_item_started(turn_context, &start_item).await; - state - .started_agent_message_items - .insert(agent_message_id.clone()); - } - - sess.emit_turn_item_completed(turn_context, TurnItem::AgentMessage(agent_message)) - .await; - state.started_agent_message_items.remove(&agent_message_id); -} - -/// Emit completion for a plan-mode turn item, handling agent messages specially. -async fn emit_turn_item_in_plan_mode( - sess: &Session, - turn_context: &TurnContext, - turn_item: TurnItem, - previously_active_item: Option<&TurnItem>, - state: &mut PlanModeStreamState, -) { - match turn_item { - TurnItem::AgentMessage(agent_message) => { - emit_agent_message_in_plan_mode(sess, turn_context, agent_message, state).await; - } - _ => { - if previously_active_item.is_none() { - sess.emit_turn_item_started(turn_context, &turn_item).await; - } - sess.emit_turn_item_completed(turn_context, turn_item).await; - } - } -} - -/// Handle a completed assistant response item in plan mode, returning true if handled. -async fn handle_assistant_item_done_in_plan_mode( - sess: &Session, - turn_context: &TurnContext, - item: &ResponseItem, - state: &mut PlanModeStreamState, - previously_active_item: Option<&TurnItem>, - last_agent_message: &mut Option, -) -> bool { - if let ResponseItem::Message { role, .. } = item - && role == "assistant" - { - maybe_complete_plan_item_from_message(sess, turn_context, state, item).await; - - if let Some(turn_item) = - handle_non_tool_response_item(sess, turn_context, item, /*plan_mode*/ true).await - { - emit_turn_item_in_plan_mode( - sess, - turn_context, - turn_item, - previously_active_item, - state, - ) - .await; - } - - record_completed_response_item(sess, turn_context, item).await; - if let Some(agent_message) = last_assistant_message_from_item(item, /*plan_mode*/ true) { - *last_agent_message = Some(agent_message); - } - return true; - } - false -} - -async fn drain_in_flight( - in_flight: &mut FuturesOrdered>>, - sess: Arc, - turn_context: Arc, -) -> CodexResult<()> { - while let Some(res) = in_flight.next().await { - match res { - Ok(response_input) => { - let response_item = response_input.into(); - sess.record_conversation_items(&turn_context, std::slice::from_ref(&response_item)) - .await; - mark_thread_memory_mode_polluted_if_external_context( - sess.as_ref(), - turn_context.as_ref(), - &response_item, - ) - .await; - } - Err(err) => { - error_or_panic(format!("in-flight tool future failed during drain: {err}")); - } - } - } - Ok(()) -} - -#[allow(clippy::too_many_arguments)] -#[instrument(level = "trace", - skip_all, - fields( - turn_id = %turn_context.sub_id, - model = %turn_context.model_info.slug - ) -)] -async fn try_run_sampling_request( - tool_runtime: ToolCallRuntime, - sess: Arc, - turn_context: Arc, - client_session: &mut ModelClientSession, - turn_metadata_header: Option<&str>, - turn_diff_tracker: SharedTurnDiffTracker, - server_model_warning_emitted_for_turn: &mut bool, - prompt: &Prompt, - cancellation_token: CancellationToken, -) -> CodexResult { - feedback_tags!( - model = turn_context.model_info.slug.clone(), - approval_policy = turn_context.approval_policy.value(), - sandbox_policy = turn_context.sandbox_policy.get(), - effort = turn_context.reasoning_effort, - auth_mode = sess.services.auth_manager.auth_mode(), - features = sess.features.enabled_features(), - ); - let mut stream = client_session - .stream( - prompt, - &turn_context.model_info, - &turn_context.session_telemetry, - turn_context.reasoning_effort, - turn_context.reasoning_summary, - turn_context.config.service_tier, - turn_metadata_header, - ) - .instrument(trace_span!("stream_request")) - .or_cancel(&cancellation_token) - .await??; - let mut in_flight: FuturesOrdered>> = - FuturesOrdered::new(); - let mut needs_follow_up = false; - let mut last_agent_message: Option = None; - let mut active_item: Option = None; - let mut should_emit_turn_diff = false; - let plan_mode = turn_context.collaboration_mode.mode == ModeKind::Plan; - let mut assistant_message_stream_parsers = AssistantMessageStreamParsers::new(plan_mode); - let mut plan_mode_state = plan_mode.then(|| PlanModeStreamState::new(&turn_context.sub_id)); - let receiving_span = trace_span!("receiving_stream"); - let outcome: CodexResult = loop { - let handle_responses = trace_span!( - parent: &receiving_span, - "handle_responses", - otel.name = field::Empty, - tool_name = field::Empty, - from = field::Empty, - ); - - let event = match stream - .next() - .instrument(trace_span!(parent: &handle_responses, "receiving")) - .or_cancel(&cancellation_token) - .await - { - Ok(event) => event, - Err(codex_async_utils::CancelErr::Cancelled) => break Err(CodexErr::TurnAborted), - }; - - let event = match event { - Some(Ok(event)) => event, - Some(Err(err)) => break Err(err), - None => { - break Err(CodexErr::Stream( - "stream closed before response.completed".into(), - None, - )); - } - }; - - sess.services - .session_telemetry - .record_responses(&handle_responses, &event); - record_turn_ttft_metric(&turn_context, &event).await; - - match event { - ResponseEvent::Created => {} - ResponseEvent::OutputItemDone(item) => { - let previously_active_item = active_item.take(); - if let Some(previous) = previously_active_item.as_ref() - && matches!(previous, TurnItem::AgentMessage(_)) - { - let item_id = previous.id(); - flush_assistant_text_segments_for_item( - &sess, - &turn_context, - plan_mode_state.as_mut(), - &mut assistant_message_stream_parsers, - &item_id, - ) - .await; - } - if let Some(state) = plan_mode_state.as_mut() - && handle_assistant_item_done_in_plan_mode( - &sess, - &turn_context, - &item, - state, - previously_active_item.as_ref(), - &mut last_agent_message, - ) - .await - { - continue; - } - - let mut ctx = HandleOutputCtx { - sess: sess.clone(), - turn_context: turn_context.clone(), - tool_runtime: tool_runtime.clone(), - cancellation_token: cancellation_token.child_token(), - }; - - let output_result = - match handle_output_item_done(&mut ctx, item, previously_active_item) - .instrument(handle_responses) - .await - { - Ok(output_result) => output_result, - Err(err) => break Err(err), - }; - if let Some(tool_future) = output_result.tool_future { - in_flight.push_back(tool_future); - } - if let Some(agent_message) = output_result.last_agent_message { - last_agent_message = Some(agent_message); - } - needs_follow_up |= output_result.needs_follow_up; - } - ResponseEvent::OutputItemAdded(item) => { - if let Some(turn_item) = handle_non_tool_response_item( - sess.as_ref(), - turn_context.as_ref(), - &item, - plan_mode, - ) - .await - { - let mut turn_item = turn_item; - let mut seeded_parsed: Option = None; - let mut seeded_item_id: Option = None; - if matches!(turn_item, TurnItem::AgentMessage(_)) - && let Some(raw_text) = raw_assistant_output_text_from_item(&item) - { - let item_id = turn_item.id(); - let mut seeded = - assistant_message_stream_parsers.seed_item_text(&item_id, &raw_text); - if let TurnItem::AgentMessage(agent_message) = &mut turn_item { - agent_message.content = - vec![codex_protocol::items::AgentMessageContent::Text { - text: if plan_mode { - String::new() - } else { - std::mem::take(&mut seeded.visible_text) - }, - }]; - } - seeded_parsed = plan_mode.then_some(seeded); - seeded_item_id = Some(item_id); - } - if let Some(state) = plan_mode_state.as_mut() - && matches!(turn_item, TurnItem::AgentMessage(_)) - { - let item_id = turn_item.id(); - state - .pending_agent_message_items - .insert(item_id, turn_item.clone()); - } else { - sess.emit_turn_item_started(&turn_context, &turn_item).await; - } - if let (Some(state), Some(item_id), Some(parsed)) = ( - plan_mode_state.as_mut(), - seeded_item_id.as_deref(), - seeded_parsed, - ) { - emit_streamed_assistant_text_delta( - &sess, - &turn_context, - Some(state), - item_id, - parsed, - ) - .await; - } - active_item = Some(turn_item); - } - } - ResponseEvent::ServerModel(server_model) => { - if !*server_model_warning_emitted_for_turn - && sess - .maybe_warn_on_server_model_mismatch(&turn_context, server_model) - .await - { - *server_model_warning_emitted_for_turn = true; - } - } - ResponseEvent::ServerReasoningIncluded(included) => { - sess.set_server_reasoning_included(included).await; - } - ResponseEvent::RateLimits(snapshot) => { - // Update internal state with latest rate limits, but defer sending until - // token usage is available to avoid duplicate TokenCount events. - sess.update_rate_limits(&turn_context, snapshot).await; - } - ResponseEvent::ModelsEtag(etag) => { - // Update internal state with latest models etag - sess.services.models_manager.refresh_if_new_etag(etag).await; - } - ResponseEvent::Completed { - response_id: _, - token_usage, - } => { - flush_assistant_text_segments_all( - &sess, - &turn_context, - plan_mode_state.as_mut(), - &mut assistant_message_stream_parsers, - ) - .await; - sess.update_token_usage_info(&turn_context, token_usage.as_ref()) - .await; - should_emit_turn_diff = true; - - break Ok(SamplingRequestResult { - needs_follow_up, - last_agent_message, - }); - } - ResponseEvent::OutputTextDelta(delta) => { - // In review child threads, suppress assistant text deltas; the - // UI will show a selection popup from the final ReviewOutput. - if let Some(active) = active_item.as_ref() { - let item_id = active.id(); - if matches!(active, TurnItem::AgentMessage(_)) { - let parsed = assistant_message_stream_parsers.parse_delta(&item_id, &delta); - emit_streamed_assistant_text_delta( - &sess, - &turn_context, - plan_mode_state.as_mut(), - &item_id, - parsed, - ) - .await; - } else { - let event = AgentMessageContentDeltaEvent { - thread_id: sess.conversation_id.to_string(), - turn_id: turn_context.sub_id.clone(), - item_id, - delta, - }; - sess.send_event(&turn_context, EventMsg::AgentMessageContentDelta(event)) - .await; - } - } else { - error_or_panic("OutputTextDelta without active item".to_string()); - } - } - ResponseEvent::ReasoningSummaryDelta { - delta, - summary_index, - } => { - if let Some(active) = active_item.as_ref() { - let event = ReasoningContentDeltaEvent { - thread_id: sess.conversation_id.to_string(), - turn_id: turn_context.sub_id.clone(), - item_id: active.id(), - delta, - summary_index, - }; - sess.send_event(&turn_context, EventMsg::ReasoningContentDelta(event)) - .await; - } else { - error_or_panic("ReasoningSummaryDelta without active item".to_string()); - } - } - ResponseEvent::ReasoningSummaryPartAdded { summary_index } => { - if let Some(active) = active_item.as_ref() { - let event = - EventMsg::AgentReasoningSectionBreak(AgentReasoningSectionBreakEvent { - item_id: active.id(), - summary_index, - }); - sess.send_event(&turn_context, event).await; - } else { - error_or_panic("ReasoningSummaryPartAdded without active item".to_string()); - } - } - ResponseEvent::ReasoningContentDelta { - delta, - content_index, - } => { - if let Some(active) = active_item.as_ref() { - let event = ReasoningRawContentDeltaEvent { - thread_id: sess.conversation_id.to_string(), - turn_id: turn_context.sub_id.clone(), - item_id: active.id(), - delta, - content_index, - }; - sess.send_event(&turn_context, EventMsg::ReasoningRawContentDelta(event)) - .await; - } else { - error_or_panic("ReasoningRawContentDelta without active item".to_string()); - } - } - } - }; - - flush_assistant_text_segments_all( - &sess, - &turn_context, - plan_mode_state.as_mut(), - &mut assistant_message_stream_parsers, - ) - .await; - - drain_in_flight(&mut in_flight, sess.clone(), turn_context.clone()).await?; - - if cancellation_token.is_cancelled() { - return Err(CodexErr::TurnAborted); - } - - if should_emit_turn_diff { - let unified_diff = { - let mut tracker = turn_diff_tracker.lock().await; - tracker.get_unified_diff() - }; - if let Ok(Some(unified_diff)) = unified_diff { - let msg = EventMsg::TurnDiff(TurnDiffEvent { unified_diff }); - sess.clone().send_event(&turn_context, msg).await; - } - } - - outcome -} - -pub(super) fn get_last_assistant_message_from_turn(responses: &[ResponseItem]) -> Option { - for item in responses.iter().rev() { - if let Some(message) = last_assistant_message_from_item(item, /*plan_mode*/ false) { - return Some(message); - } - } - None -} - use crate::memories::prompts::build_memory_tool_developer_instructions; #[cfg(test)] pub(crate) use tests::make_session_and_context; diff --git a/codex-rs/core/src/codex/turn.rs b/codex-rs/core/src/codex/turn.rs new file mode 100644 index 0000000000..fa24bd6924 --- /dev/null +++ b/codex-rs/core/src/codex/turn.rs @@ -0,0 +1,2169 @@ +use std::collections::HashMap; +use std::collections::HashSet; +use std::sync::Arc; + +use crate::SkillInjections; +use crate::SkillLoadOutcome; +use crate::build_skill_injections; +use crate::client::ModelClientSession; +use crate::client_common::Prompt; +use crate::client_common::ResponseEvent; +use crate::codex::PreviousTurnSettings; +use crate::codex::Session; +use crate::codex::TurnContext; +use crate::collect_env_var_dependencies; +use crate::collect_explicit_skill_mentions; +use crate::compact::InitialContextInjection; +use crate::compact::collect_user_messages; +use crate::compact::run_inline_auto_compact_task; +use crate::compact::should_use_remote_compact_task; +use crate::compact_remote::run_inline_remote_auto_compact_task; +use crate::connectors; +use crate::feedback_tags; +use crate::hook_runtime::PendingInputHookDisposition; +use crate::hook_runtime::emit_hook_completed_events; +use crate::hook_runtime::inspect_pending_input; +use crate::hook_runtime::record_additional_contexts; +use crate::hook_runtime::record_pending_input; +use crate::hook_runtime::run_pending_session_start_hooks; +use crate::hook_runtime::run_user_prompt_submit_hooks; +use crate::injection::ToolMentionKind; +use crate::injection::app_id_from_path; +use crate::injection::tool_kind_for_path; +use crate::mcp_skill_dependencies::maybe_prompt_and_install_mcp_dependencies; +use crate::mcp_tool_exposure::build_mcp_tool_exposure; +use crate::mentions::build_connector_slug_counts; +use crate::mentions::build_skill_name_counts; +use crate::mentions::collect_explicit_app_ids; +use crate::mentions::collect_explicit_plugin_mentions; +use crate::mentions::collect_tool_mentions_from_messages; +use crate::parse_turn_item; +use crate::plugins::build_plugin_injections; +use crate::resolve_skill_dependencies_for_turn; +use crate::stream_events_utils::HandleOutputCtx; +use crate::stream_events_utils::handle_non_tool_response_item; +use crate::stream_events_utils::handle_output_item_done; +use crate::stream_events_utils::last_assistant_message_from_item; +use crate::stream_events_utils::mark_thread_memory_mode_polluted_if_external_context; +use crate::stream_events_utils::raw_assistant_output_text_from_item; +use crate::stream_events_utils::record_completed_response_item; +use crate::tools::ToolRouter; +use crate::tools::context::SharedTurnDiffTracker; +use crate::tools::parallel::ToolCallRuntime; +use crate::tools::router::ToolRouterParams; +use crate::turn_diff_tracker::TurnDiffTracker; +use crate::turn_timing::record_turn_ttft_metric; +use crate::unavailable_tool::collect_unavailable_called_tools; +use crate::util::backoff; +use crate::util::error_or_panic; +use codex_analytics::AppInvocation; +use codex_analytics::CompactionPhase; +use codex_analytics::CompactionReason; +use codex_analytics::InvocationType; +use codex_analytics::TurnResolvedConfigFact; +use codex_analytics::build_track_events_context; +use codex_async_utils::OrCancelExt; +use codex_features::Feature; +use codex_hooks::HookEvent; +use codex_hooks::HookEventAfterAgent; +use codex_hooks::HookPayload; +use codex_hooks::HookResult; +use codex_protocol::config_types::ModeKind; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result as CodexResult; +use codex_protocol::items::PlanItem; +use codex_protocol::items::TurnItem; +use codex_protocol::items::UserMessageItem; +use codex_protocol::items::build_hook_prompt_message; +use codex_protocol::models::BaseInstructions; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseInputItem; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::AgentMessageContentDeltaEvent; +use codex_protocol::protocol::AgentReasoningSectionBreakEvent; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::CodexErrorInfo; +use codex_protocol::protocol::ErrorEvent; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::PlanDeltaEvent; +use codex_protocol::protocol::ReasoningContentDeltaEvent; +use codex_protocol::protocol::ReasoningRawContentDeltaEvent; +use codex_protocol::protocol::TurnDiffEvent; +use codex_protocol::protocol::WarningEvent; +use codex_protocol::user_input::UserInput; +use codex_tools::filter_tool_suggest_discoverable_tools_for_client; +use codex_utils_stream_parser::AssistantTextChunk; +use codex_utils_stream_parser::AssistantTextStreamParser; +use codex_utils_stream_parser::ProposedPlanSegment; +use codex_utils_stream_parser::extract_proposed_plan_text; +use codex_utils_stream_parser::strip_citations; +use futures::future::BoxFuture; +use futures::prelude::*; +use futures::stream::FuturesOrdered; +use tokio_util::sync::CancellationToken; +use tracing::Instrument; +use tracing::error; +use tracing::field; +use tracing::info; +use tracing::instrument; +use tracing::trace; +use tracing::trace_span; +use tracing::warn; + +/// Takes a user message as input and runs a loop where, at each sampling request, the model +/// replies with either: +/// +/// - requested function calls +/// - an assistant message +/// +/// While it is possible for the model to return multiple of these items in a +/// single sampling request, in practice, we generally one item per sampling request: +/// +/// - If the model requests a function call, we execute it and send the output +/// back to the model in the next sampling request. +/// - If the model sends only an assistant message, we record it in the +/// conversation history and consider the turn complete. +/// +pub(crate) async fn run_turn( + sess: Arc, + turn_context: Arc, + input: Vec, + prewarmed_client_session: Option, + cancellation_token: CancellationToken, +) -> Option { + if input.is_empty() && !sess.has_pending_input().await { + return None; + } + + let model_info = turn_context.model_info.clone(); + let auto_compact_limit = model_info.auto_compact_token_limit().unwrap_or(i64::MAX); + let mut prewarmed_client_session = prewarmed_client_session; + // TODO(ccunningham): Pre-turn compaction runs before context updates and the + // new user message are recorded. Estimate pending incoming items (context + // diffs/full reinjection + user input) and trigger compaction preemptively + // when they would push the thread over the compaction threshold. + let pre_sampling_compacted = match run_pre_sampling_compact(&sess, &turn_context).await { + Ok(pre_sampling_compacted) => pre_sampling_compacted, + Err(_) => { + error!("Failed to run pre-sampling compact"); + return None; + } + }; + if pre_sampling_compacted && let Some(mut client_session) = prewarmed_client_session.take() { + client_session.reset_websocket_session(); + } + + let skills_outcome = Some(turn_context.turn_skills.outcome.as_ref()); + + sess.record_context_updates_and_set_reference_context_item(turn_context.as_ref()) + .await; + + let loaded_plugins = sess + .services + .plugins_manager + .plugins_for_config(&turn_context.config) + .await; + // Structured plugin:// mentions are resolved from the current session's + // enabled plugins, then converted into turn-scoped guidance below. + let mentioned_plugins = + collect_explicit_plugin_mentions(&input, loaded_plugins.capability_summaries()); + let mcp_tools = if turn_context.apps_enabled() || !mentioned_plugins.is_empty() { + // Plugin mentions need raw MCP/app inventory even when app tools + // are normally hidden so we can describe the plugin's currently + // usable capabilities for this turn. + match sess + .services + .mcp_connection_manager + .read() + .await + .list_all_tools() + .or_cancel(&cancellation_token) + .await + { + Ok(mcp_tools) => mcp_tools, + Err(_) if turn_context.apps_enabled() => return None, + Err(_) => HashMap::new(), + } + } else { + HashMap::new() + }; + let available_connectors = if turn_context.apps_enabled() { + let connectors = codex_connectors::merge::merge_plugin_connectors_with_accessible( + loaded_plugins + .effective_apps() + .into_iter() + .map(|connector_id| connector_id.0), + connectors::accessible_connectors_from_mcp_tools(&mcp_tools), + ); + connectors::with_app_enabled_state(connectors, &turn_context.config) + } else { + Vec::new() + }; + let connector_slug_counts = build_connector_slug_counts(&available_connectors); + let skill_name_counts_lower = skills_outcome + .as_ref() + .map_or_else(HashMap::new, |outcome| { + build_skill_name_counts(&outcome.skills, &outcome.disabled_paths).1 + }); + let mentioned_skills = skills_outcome.as_ref().map_or_else(Vec::new, |outcome| { + collect_explicit_skill_mentions( + &input, + &outcome.skills, + &outcome.disabled_paths, + &connector_slug_counts, + ) + }); + let config = turn_context.config.clone(); + if config + .features + .enabled(Feature::SkillEnvVarDependencyPrompt) + { + let env_var_dependencies = collect_env_var_dependencies(&mentioned_skills); + resolve_skill_dependencies_for_turn(&sess, &turn_context, &env_var_dependencies).await; + } + + maybe_prompt_and_install_mcp_dependencies( + sess.as_ref(), + turn_context.as_ref(), + &cancellation_token, + &mentioned_skills, + ) + .await; + + let session_telemetry = turn_context.session_telemetry.clone(); + let thread_id = sess.conversation_id.to_string(); + let tracking = build_track_events_context( + turn_context.model_info.slug.clone(), + thread_id, + turn_context.sub_id.clone(), + ); + let SkillInjections { + items: skill_items, + warnings: skill_warnings, + } = build_skill_injections( + &mentioned_skills, + skills_outcome, + Some(&session_telemetry), + &sess.services.analytics_events_client, + tracking.clone(), + ) + .await; + + for message in skill_warnings { + sess.send_event(&turn_context, EventMsg::Warning(WarningEvent { message })) + .await; + } + + let plugin_items = + build_plugin_injections(&mentioned_plugins, &mcp_tools, &available_connectors); + let mentioned_plugin_metadata = mentioned_plugins + .iter() + .filter_map(crate::plugins::PluginCapabilitySummary::telemetry_metadata) + .collect::>(); + + let mut explicitly_enabled_connectors = collect_explicit_app_ids(&input); + explicitly_enabled_connectors.extend(collect_explicit_app_ids_from_skill_items( + &skill_items, + &available_connectors, + &skill_name_counts_lower, + )); + let connector_names_by_id = available_connectors + .iter() + .map(|connector| (connector.id.as_str(), connector.name.as_str())) + .collect::>(); + let mentioned_app_invocations = explicitly_enabled_connectors + .iter() + .map(|connector_id| AppInvocation { + connector_id: Some(connector_id.clone()), + app_name: connector_names_by_id + .get(connector_id.as_str()) + .map(|name| (*name).to_string()), + invocation_type: Some(InvocationType::Explicit), + }) + .collect::>(); + + if run_pending_session_start_hooks(&sess, &turn_context).await { + return None; + } + let additional_contexts = if input.is_empty() { + Vec::new() + } else { + let initial_input_for_turn: ResponseInputItem = ResponseInputItem::from(input.clone()); + let response_item: ResponseItem = initial_input_for_turn.clone().into(); + let user_prompt_submit_outcome = run_user_prompt_submit_hooks( + &sess, + &turn_context, + UserMessageItem::new(&input).message(), + ) + .await; + if user_prompt_submit_outcome.should_stop { + record_additional_contexts( + &sess, + &turn_context, + user_prompt_submit_outcome.additional_contexts, + ) + .await; + return None; + } + sess.record_user_prompt_and_emit_turn_item(turn_context.as_ref(), &input, response_item) + .await; + user_prompt_submit_outcome.additional_contexts + }; + sess.services + .analytics_events_client + .track_app_mentioned(tracking.clone(), mentioned_app_invocations); + for plugin in mentioned_plugin_metadata { + sess.services + .analytics_events_client + .track_plugin_used(tracking.clone(), plugin); + } + sess.merge_connector_selection(explicitly_enabled_connectors.clone()) + .await; + record_additional_contexts(&sess, &turn_context, additional_contexts).await; + if !input.is_empty() { + // Track the previous-turn baseline from the regular user-turn path only so + // standalone tasks (compact/shell/review/undo) cannot suppress future + // model/realtime injections. + sess.set_previous_turn_settings(Some(PreviousTurnSettings { + model: turn_context.model_info.slug.clone(), + realtime_active: Some(turn_context.realtime_active), + })) + .await; + } + if let Err(error) = sess.ensure_agent_task_registered().await { + warn!(error = %error, "agent task registration failed"); + sess.send_event( + turn_context.as_ref(), + EventMsg::Error(ErrorEvent { + message: format!( + "Agent task registration failed. Please try again; Codex will attempt to register the task again on the next turn: {error}" + ), + codex_error_info: Some(CodexErrorInfo::Other), + }), + ) + .await; + return None; + } + + if !skill_items.is_empty() { + sess.record_conversation_items(&turn_context, &skill_items) + .await; + } + if !plugin_items.is_empty() { + sess.record_conversation_items(&turn_context, &plugin_items) + .await; + } + + track_turn_resolved_config_analytics(&sess, &turn_context, &input).await; + + let skills_outcome = Some(turn_context.turn_skills.outcome.as_ref()); + sess.maybe_start_ghost_snapshot(Arc::clone(&turn_context), cancellation_token.child_token()) + .await; + let mut last_agent_message: Option = None; + let mut stop_hook_active = false; + // Although from the perspective of codex.rs, TurnDiffTracker has the lifecycle of a Task which contains + // many turns, from the perspective of the user, it is a single turn. + let turn_diff_tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())); + let mut server_model_warning_emitted_for_turn = false; + + // `ModelClientSession` is turn-scoped and caches WebSocket + sticky routing state, so we reuse + // one instance across retries within this turn. + let mut client_session = + prewarmed_client_session.unwrap_or_else(|| sess.services.model_client.new_session()); + // Pending input is drained into history before building the next model request. + // However, we defer that drain until after sampling in two cases: + // 1. At the start of a turn, so the fresh user prompt in `input` gets sampled first. + // 2. After auto-compact, when model/tool continuation needs to resume before any steer. + let mut can_drain_pending_input = input.is_empty(); + + loop { + if run_pending_session_start_hooks(&sess, &turn_context).await { + break; + } + + // Note that pending_input would be something like a message the user + // submitted through the UI while the model was running. Though the UI + // may support this, the model might not. + let pending_input = if can_drain_pending_input { + sess.get_pending_input().await + } else { + Vec::new() + }; + + let mut blocked_pending_input = false; + let mut blocked_pending_input_contexts = Vec::new(); + let mut requeued_pending_input = false; + let mut accepted_pending_input = Vec::new(); + if !pending_input.is_empty() { + let mut pending_input_iter = pending_input.into_iter(); + while let Some(pending_input_item) = pending_input_iter.next() { + match inspect_pending_input(&sess, &turn_context, pending_input_item).await { + PendingInputHookDisposition::Accepted(pending_input) => { + accepted_pending_input.push(*pending_input); + } + PendingInputHookDisposition::Blocked { + additional_contexts, + } => { + let remaining_pending_input = pending_input_iter.collect::>(); + if !remaining_pending_input.is_empty() { + let _ = sess.prepend_pending_input(remaining_pending_input).await; + requeued_pending_input = true; + } + blocked_pending_input_contexts = additional_contexts; + blocked_pending_input = true; + break; + } + } + } + } + + let has_accepted_pending_input = !accepted_pending_input.is_empty(); + for pending_input in accepted_pending_input { + record_pending_input(&sess, &turn_context, pending_input).await; + } + record_additional_contexts(&sess, &turn_context, blocked_pending_input_contexts).await; + + if blocked_pending_input && !has_accepted_pending_input { + if requeued_pending_input { + continue; + } + break; + } + + // Construct the input that we will send to the model. + let sampling_request_input: Vec = { + sess.clone_history() + .await + .for_prompt(&turn_context.model_info.input_modalities) + }; + + let sampling_request_input_messages = sampling_request_input + .iter() + .filter_map(|item| match parse_turn_item(item) { + Some(TurnItem::UserMessage(user_message)) => Some(user_message), + _ => None, + }) + .map(|user_message| user_message.message()) + .collect::>(); + let turn_metadata_header = turn_context.turn_metadata_state.current_header_value(); + match run_sampling_request( + Arc::clone(&sess), + Arc::clone(&turn_context), + Arc::clone(&turn_diff_tracker), + &mut client_session, + turn_metadata_header.as_deref(), + sampling_request_input, + &explicitly_enabled_connectors, + skills_outcome, + &mut server_model_warning_emitted_for_turn, + cancellation_token.child_token(), + ) + .await + { + Ok(sampling_request_output) => { + let SamplingRequestResult { + needs_follow_up: model_needs_follow_up, + last_agent_message: sampling_request_last_agent_message, + } = sampling_request_output; + can_drain_pending_input = true; + let has_pending_input = sess.has_pending_input().await; + let needs_follow_up = model_needs_follow_up || has_pending_input; + let total_usage_tokens = sess.get_total_token_usage().await; + let token_limit_reached = total_usage_tokens >= auto_compact_limit; + + let estimated_token_count = + sess.get_estimated_token_count(turn_context.as_ref()).await; + + trace!( + turn_id = %turn_context.sub_id, + total_usage_tokens, + estimated_token_count = ?estimated_token_count, + auto_compact_limit, + token_limit_reached, + model_needs_follow_up, + has_pending_input, + needs_follow_up, + "post sampling token usage" + ); + + // as long as compaction works well in getting us way below the token limit, we shouldn't worry about being in an infinite loop. + if token_limit_reached && needs_follow_up { + if run_auto_compact( + &sess, + &turn_context, + InitialContextInjection::BeforeLastUserMessage, + CompactionReason::ContextLimit, + CompactionPhase::MidTurn, + ) + .await + .is_err() + { + return None; + } + client_session.reset_websocket_session(); + can_drain_pending_input = !model_needs_follow_up; + continue; + } + + if !needs_follow_up { + last_agent_message = sampling_request_last_agent_message; + let stop_hook_permission_mode = match turn_context.approval_policy.value() { + AskForApproval::Never => "bypassPermissions", + AskForApproval::UnlessTrusted + | AskForApproval::OnFailure + | AskForApproval::OnRequest + | AskForApproval::Granular(_) => "default", + } + .to_string(); + let stop_request = codex_hooks::StopRequest { + session_id: sess.conversation_id, + turn_id: turn_context.sub_id.clone(), + cwd: turn_context.cwd.clone(), + transcript_path: sess.hook_transcript_path().await, + model: turn_context.model_info.slug.clone(), + permission_mode: stop_hook_permission_mode, + stop_hook_active, + last_assistant_message: last_agent_message.clone(), + }; + for run in sess.hooks().preview_stop(&stop_request) { + sess.send_event( + &turn_context, + EventMsg::HookStarted(codex_protocol::protocol::HookStartedEvent { + turn_id: Some(turn_context.sub_id.clone()), + run, + }), + ) + .await; + } + let stop_outcome = sess.hooks().run_stop(stop_request).await; + emit_hook_completed_events(&sess, &turn_context, stop_outcome.hook_events) + .await; + if stop_outcome.should_block { + if let Some(hook_prompt_message) = + build_hook_prompt_message(&stop_outcome.continuation_fragments) + { + sess.record_conversation_items( + &turn_context, + std::slice::from_ref(&hook_prompt_message), + ) + .await; + stop_hook_active = true; + continue; + } else { + sess.send_event( + &turn_context, + EventMsg::Warning(WarningEvent { + message: "Stop hook requested continuation without a prompt; ignoring the block.".to_string(), + }), + ) + .await; + } + } + if stop_outcome.should_stop { + break; + } + let hook_outcomes = sess + .hooks() + .dispatch(HookPayload { + session_id: sess.conversation_id, + cwd: turn_context.cwd.clone(), + client: turn_context.app_server_client_name.clone(), + triggered_at: chrono::Utc::now(), + hook_event: HookEvent::AfterAgent { + event: HookEventAfterAgent { + thread_id: sess.conversation_id, + turn_id: turn_context.sub_id.clone(), + input_messages: sampling_request_input_messages, + last_assistant_message: last_agent_message.clone(), + }, + }, + }) + .await; + + let mut abort_message = None; + for hook_outcome in hook_outcomes { + let hook_name = hook_outcome.hook_name; + match hook_outcome.result { + HookResult::Success => {} + HookResult::FailedContinue(error) => { + warn!( + turn_id = %turn_context.sub_id, + hook_name = %hook_name, + error = %error, + "after_agent hook failed; continuing" + ); + } + HookResult::FailedAbort(error) => { + let message = format!( + "after_agent hook '{hook_name}' failed and aborted turn completion: {error}" + ); + warn!( + turn_id = %turn_context.sub_id, + hook_name = %hook_name, + error = %error, + "after_agent hook failed; aborting operation" + ); + if abort_message.is_none() { + abort_message = Some(message); + } + } + } + } + if let Some(message) = abort_message { + sess.send_event( + &turn_context, + EventMsg::Error(ErrorEvent { + message, + codex_error_info: None, + }), + ) + .await; + return None; + } + break; + } + continue; + } + Err(CodexErr::TurnAborted) => { + // Aborted turn is reported via a different event. + break; + } + Err(CodexErr::InvalidImageRequest()) => { + { + let mut state = sess.state.lock().await; + error_or_panic( + "Invalid image detected; sanitizing tool output to prevent poisoning", + ); + if state.history.replace_last_turn_images("Invalid image") { + continue; + } + } + + let event = EventMsg::Error(ErrorEvent { + message: "Invalid image in your last message. Please remove it and try again." + .to_string(), + codex_error_info: Some(CodexErrorInfo::BadRequest), + }); + sess.send_event(&turn_context, event).await; + break; + } + Err(e) => { + info!("Turn error: {e:#}"); + let event = EventMsg::Error(e.to_error_event(/*message_prefix*/ None)); + sess.send_event(&turn_context, event).await; + // let the user continue the conversation + break; + } + } + } + + last_agent_message +} + +async fn track_turn_resolved_config_analytics( + sess: &Session, + turn_context: &TurnContext, + input: &[UserInput], +) { + if !sess.enabled(Feature::GeneralAnalytics) { + return; + } + + let thread_config = { + let state = sess.state.lock().await; + state.session_configuration.thread_config_snapshot() + }; + let is_first_turn = { + let mut state = sess.state.lock().await; + state.take_next_turn_is_first() + }; + sess.services + .analytics_events_client + .track_turn_resolved_config(TurnResolvedConfigFact { + turn_id: turn_context.sub_id.clone(), + thread_id: sess.conversation_id.to_string(), + num_input_images: input + .iter() + .filter(|item| { + matches!(item, UserInput::Image { .. } | UserInput::LocalImage { .. }) + }) + .count(), + submission_type: None, + ephemeral: thread_config.ephemeral, + session_source: thread_config.session_source, + model: turn_context.model_info.slug.clone(), + model_provider: turn_context.config.model_provider_id.clone(), + sandbox_policy: turn_context.sandbox_policy.get().clone(), + reasoning_effort: turn_context.reasoning_effort, + reasoning_summary: Some(turn_context.reasoning_summary), + service_tier: turn_context.config.service_tier, + approval_policy: turn_context.approval_policy.value(), + approvals_reviewer: turn_context.config.approvals_reviewer, + sandbox_network_access: turn_context.network_sandbox_policy.is_enabled(), + collaboration_mode: turn_context.collaboration_mode.mode, + personality: turn_context.personality, + is_first_turn, + }); +} + +async fn run_pre_sampling_compact( + sess: &Arc, + turn_context: &Arc, +) -> CodexResult { + let total_usage_tokens_before_compaction = sess.get_total_token_usage().await; + let mut pre_sampling_compacted = maybe_run_previous_model_inline_compact( + sess, + turn_context, + total_usage_tokens_before_compaction, + ) + .await?; + let total_usage_tokens = sess.get_total_token_usage().await; + let auto_compact_limit = turn_context + .model_info + .auto_compact_token_limit() + .unwrap_or(i64::MAX); + // Compact if the total usage tokens are greater than the auto compact limit + if total_usage_tokens >= auto_compact_limit { + run_auto_compact( + sess, + turn_context, + InitialContextInjection::DoNotInject, + CompactionReason::ContextLimit, + CompactionPhase::PreTurn, + ) + .await?; + pre_sampling_compacted = true; + } + Ok(pre_sampling_compacted) +} + +/// Runs pre-sampling compaction against the previous model when switching to a smaller +/// context-window model. +/// +/// Returns `Ok(true)` when compaction ran successfully, `Ok(false)` when compaction was skipped +/// because the model/context-window preconditions were not met, and `Err(_)` only when compaction +/// was attempted and failed. +async fn maybe_run_previous_model_inline_compact( + sess: &Arc, + turn_context: &Arc, + total_usage_tokens: i64, +) -> CodexResult { + let Some(previous_turn_settings) = sess.previous_turn_settings().await else { + return Ok(false); + }; + let previous_model_turn_context = Arc::new( + turn_context + .with_model(previous_turn_settings.model, &sess.services.models_manager) + .await, + ); + + let Some(old_context_window) = previous_model_turn_context.model_context_window() else { + return Ok(false); + }; + let Some(new_context_window) = turn_context.model_context_window() else { + return Ok(false); + }; + let new_auto_compact_limit = turn_context + .model_info + .auto_compact_token_limit() + .unwrap_or(i64::MAX); + let should_run = total_usage_tokens > new_auto_compact_limit + && previous_model_turn_context.model_info.slug != turn_context.model_info.slug + && old_context_window > new_context_window; + if should_run { + run_auto_compact( + sess, + &previous_model_turn_context, + InitialContextInjection::DoNotInject, + CompactionReason::ModelDownshift, + CompactionPhase::PreTurn, + ) + .await?; + return Ok(true); + } + Ok(false) +} + +async fn run_auto_compact( + sess: &Arc, + turn_context: &Arc, + initial_context_injection: InitialContextInjection, + reason: CompactionReason, + phase: CompactionPhase, +) -> CodexResult<()> { + if should_use_remote_compact_task(&turn_context.provider) { + run_inline_remote_auto_compact_task( + Arc::clone(sess), + Arc::clone(turn_context), + initial_context_injection, + reason, + phase, + ) + .await?; + } else { + run_inline_auto_compact_task( + Arc::clone(sess), + Arc::clone(turn_context), + initial_context_injection, + reason, + phase, + ) + .await?; + } + Ok(()) +} + +pub(super) fn collect_explicit_app_ids_from_skill_items( + skill_items: &[ResponseItem], + connectors: &[connectors::AppInfo], + skill_name_counts_lower: &HashMap, +) -> HashSet { + if skill_items.is_empty() || connectors.is_empty() { + return HashSet::new(); + } + + let skill_messages = skill_items + .iter() + .filter_map(|item| match item { + ResponseItem::Message { content, .. } => { + content.iter().find_map(|content_item| match content_item { + ContentItem::InputText { text } => Some(text.clone()), + _ => None, + }) + } + _ => None, + }) + .collect::>(); + if skill_messages.is_empty() { + return HashSet::new(); + } + + let mentions = collect_tool_mentions_from_messages(&skill_messages); + let mention_names_lower = mentions + .plain_names + .iter() + .map(|name| name.to_ascii_lowercase()) + .collect::>(); + let mut connector_ids = mentions + .paths + .iter() + .filter(|path| tool_kind_for_path(path) == ToolMentionKind::App) + .filter_map(|path| app_id_from_path(path).map(str::to_string)) + .collect::>(); + + let connector_slug_counts = build_connector_slug_counts(connectors); + for connector in connectors { + let slug = codex_connectors::metadata::connector_mention_slug(connector); + let connector_count = connector_slug_counts.get(&slug).copied().unwrap_or(0); + let skill_count = skill_name_counts_lower.get(&slug).copied().unwrap_or(0); + if connector_count == 1 && skill_count == 0 && mention_names_lower.contains(&slug) { + connector_ids.insert(connector.id.clone()); + } + } + + connector_ids +} + +pub(super) fn filter_connectors_for_input( + connectors: &[connectors::AppInfo], + input: &[ResponseItem], + explicitly_enabled_connectors: &HashSet, + skill_name_counts_lower: &HashMap, +) -> Vec { + let connectors: Vec = connectors + .iter() + .filter(|connector| connector.is_enabled) + .cloned() + .collect::>(); + if connectors.is_empty() { + return Vec::new(); + } + + let user_messages = collect_user_messages(input); + if user_messages.is_empty() && explicitly_enabled_connectors.is_empty() { + return Vec::new(); + } + + let mentions = collect_tool_mentions_from_messages(&user_messages); + let mention_names_lower = mentions + .plain_names + .iter() + .map(|name| name.to_ascii_lowercase()) + .collect::>(); + + let connector_slug_counts = build_connector_slug_counts(&connectors); + let mut allowed_connector_ids = explicitly_enabled_connectors.clone(); + for path in mentions + .paths + .iter() + .filter(|path| tool_kind_for_path(path) == ToolMentionKind::App) + { + if let Some(connector_id) = app_id_from_path(path) { + allowed_connector_ids.insert(connector_id.to_string()); + } + } + + connectors + .into_iter() + .filter(|connector| { + connector_inserted_in_messages( + connector, + &mention_names_lower, + &allowed_connector_ids, + &connector_slug_counts, + skill_name_counts_lower, + ) + }) + .collect() +} + +fn connector_inserted_in_messages( + connector: &connectors::AppInfo, + mention_names_lower: &HashSet, + allowed_connector_ids: &HashSet, + connector_slug_counts: &HashMap, + skill_name_counts_lower: &HashMap, +) -> bool { + if allowed_connector_ids.contains(&connector.id) { + return true; + } + + let mention_slug = codex_connectors::metadata::connector_mention_slug(connector); + let connector_count = connector_slug_counts + .get(&mention_slug) + .copied() + .unwrap_or(0); + let skill_count = skill_name_counts_lower + .get(&mention_slug) + .copied() + .unwrap_or(0); + connector_count == 1 && skill_count == 0 && mention_names_lower.contains(&mention_slug) +} + +pub(crate) fn build_prompt( + input: Vec, + router: &ToolRouter, + turn_context: &TurnContext, + base_instructions: BaseInstructions, +) -> Prompt { + let deferred_dynamic_tools = turn_context + .dynamic_tools + .iter() + .filter(|tool| tool.defer_loading) + .map(|tool| tool.name.as_str()) + .collect::>(); + let tools = if deferred_dynamic_tools.is_empty() { + router.model_visible_specs() + } else { + router + .model_visible_specs() + .into_iter() + .filter(|spec| !deferred_dynamic_tools.contains(spec.name())) + .collect() + }; + + Prompt { + input, + tools, + parallel_tool_calls: turn_context.model_info.supports_parallel_tool_calls, + base_instructions, + personality: turn_context.personality, + output_schema: turn_context.final_output_json_schema.clone(), + } +} + +#[allow(clippy::too_many_arguments)] +#[instrument(level = "trace", + skip_all, + fields( + turn_id = %turn_context.sub_id, + model = %turn_context.model_info.slug, + cwd = %turn_context.cwd.display() + ) +)] +async fn run_sampling_request( + sess: Arc, + turn_context: Arc, + turn_diff_tracker: SharedTurnDiffTracker, + client_session: &mut ModelClientSession, + turn_metadata_header: Option<&str>, + input: Vec, + explicitly_enabled_connectors: &HashSet, + skills_outcome: Option<&SkillLoadOutcome>, + server_model_warning_emitted_for_turn: &mut bool, + cancellation_token: CancellationToken, +) -> CodexResult { + let router = built_tools( + sess.as_ref(), + turn_context.as_ref(), + &input, + explicitly_enabled_connectors, + skills_outcome, + &cancellation_token, + ) + .await?; + + let base_instructions = sess.get_base_instructions().await; + + let tool_runtime = ToolCallRuntime::new( + Arc::clone(&router), + Arc::clone(&sess), + Arc::clone(&turn_context), + Arc::clone(&turn_diff_tracker), + ); + let _code_mode_worker = sess + .services + .code_mode_service + .start_turn_worker( + &sess, + &turn_context, + Arc::clone(&router), + Arc::clone(&turn_diff_tracker), + ) + .await; + let mut retries = 0; + let mut initial_input = Some(input); + loop { + let prompt_input = if let Some(input) = initial_input.take() { + input + } else { + sess.clone_history() + .await + .for_prompt(&turn_context.model_info.input_modalities) + }; + let prompt = build_prompt( + prompt_input, + router.as_ref(), + turn_context.as_ref(), + base_instructions.clone(), + ); + let err = match try_run_sampling_request( + tool_runtime.clone(), + Arc::clone(&sess), + Arc::clone(&turn_context), + client_session, + turn_metadata_header, + Arc::clone(&turn_diff_tracker), + server_model_warning_emitted_for_turn, + &prompt, + cancellation_token.child_token(), + ) + .await + { + Ok(output) => { + return Ok(output); + } + Err(CodexErr::ContextWindowExceeded) => { + sess.set_total_tokens_full(&turn_context).await; + return Err(CodexErr::ContextWindowExceeded); + } + Err(CodexErr::UsageLimitReached(e)) => { + let rate_limits = e.rate_limits.clone(); + if let Some(rate_limits) = rate_limits { + sess.update_rate_limits(&turn_context, *rate_limits).await; + } + return Err(CodexErr::UsageLimitReached(e)); + } + Err(err) => err, + }; + + if !err.is_retryable() { + return Err(err); + } + + // Use the configured provider-specific stream retry budget. + let max_retries = turn_context.provider.stream_max_retries(); + if retries >= max_retries + && client_session.try_switch_fallback_transport( + &turn_context.session_telemetry, + &turn_context.model_info, + ) + { + sess.send_event( + &turn_context, + EventMsg::Warning(WarningEvent { + message: format!("Falling back from WebSockets to HTTPS transport. {err:#}"), + }), + ) + .await; + retries = 0; + continue; + } + if retries < max_retries { + retries += 1; + let delay = match &err { + CodexErr::Stream(_, requested_delay) => { + requested_delay.unwrap_or_else(|| backoff(retries)) + } + _ => backoff(retries), + }; + warn!( + "stream disconnected - retrying sampling request ({retries}/{max_retries} in {delay:?})...", + ); + + // In release builds, hide the first websocket retry notification to reduce noisy + // transient reconnect messages. In debug builds, keep full visibility for diagnosis. + let report_error = retries > 1 + || cfg!(debug_assertions) + || !sess.services.model_client.responses_websocket_enabled(); + if report_error { + // Surface retry information to any UI/front‑end so the + // user understands what is happening instead of staring + // at a seemingly frozen screen. + sess.notify_stream_error( + &turn_context, + format!("Reconnecting... {retries}/{max_retries}"), + err, + ) + .await; + } + tokio::time::sleep(delay).await; + } else { + return Err(err); + } + } +} + +pub(crate) async fn built_tools( + sess: &Session, + turn_context: &TurnContext, + input: &[ResponseItem], + explicitly_enabled_connectors: &HashSet, + skills_outcome: Option<&SkillLoadOutcome>, + cancellation_token: &CancellationToken, +) -> CodexResult> { + let mcp_connection_manager = sess.services.mcp_connection_manager.read().await; + let has_mcp_servers = mcp_connection_manager.has_servers(); + let all_mcp_tools = mcp_connection_manager + .list_all_tools() + .or_cancel(cancellation_token) + .await?; + drop(mcp_connection_manager); + let loaded_plugins = sess + .services + .plugins_manager + .plugins_for_config(&turn_context.config) + .await; + + let mut effective_explicitly_enabled_connectors = explicitly_enabled_connectors.clone(); + effective_explicitly_enabled_connectors.extend(sess.get_connector_selection().await); + + let apps_enabled = turn_context.apps_enabled(); + let accessible_connectors = + apps_enabled.then(|| connectors::accessible_connectors_from_mcp_tools(&all_mcp_tools)); + let accessible_connectors_with_enabled_state = + accessible_connectors.as_ref().map(|connectors| { + connectors::with_app_enabled_state(connectors.clone(), &turn_context.config) + }); + let connectors = if apps_enabled { + let connectors = codex_connectors::merge::merge_plugin_connectors_with_accessible( + loaded_plugins + .effective_apps() + .into_iter() + .map(|connector_id| connector_id.0), + accessible_connectors.clone().unwrap_or_default(), + ); + Some(connectors::with_app_enabled_state( + connectors, + &turn_context.config, + )) + } else { + None + }; + let auth = sess.services.auth_manager.auth().await; + let discoverable_tools = if apps_enabled && turn_context.tools_config.tool_suggest { + if let Some(accessible_connectors) = accessible_connectors_with_enabled_state.as_ref() { + match connectors::list_tool_suggest_discoverable_tools_with_auth( + &turn_context.config, + auth.as_ref(), + accessible_connectors.as_slice(), + ) + .await + .map(|discoverable_tools| { + filter_tool_suggest_discoverable_tools_for_client( + discoverable_tools, + turn_context.app_server_client_name.as_deref(), + ) + }) { + Ok(discoverable_tools) if discoverable_tools.is_empty() => None, + Ok(discoverable_tools) => Some(discoverable_tools), + Err(err) => { + warn!("failed to load discoverable tool suggestions: {err:#}"); + None + } + } + } else { + None + } + } else { + None + }; + + let explicitly_enabled = if let Some(connectors) = connectors.as_ref() { + let skill_name_counts_lower = skills_outcome.map_or_else(HashMap::new, |outcome| { + build_skill_name_counts(&outcome.skills, &outcome.disabled_paths).1 + }); + + filter_connectors_for_input( + connectors, + input, + &effective_explicitly_enabled_connectors, + &skill_name_counts_lower, + ) + } else { + Vec::new() + }; + let mcp_tool_exposure = build_mcp_tool_exposure( + &all_mcp_tools, + connectors.as_deref(), + explicitly_enabled.as_slice(), + &turn_context.config, + &turn_context.tools_config, + ); + let mcp_tools = has_mcp_servers.then_some(mcp_tool_exposure.direct_tools); + let deferred_mcp_tools = mcp_tool_exposure.deferred_tools; + let unavailable_called_tools = if turn_context + .config + .features + .enabled(Feature::UnavailableDummyTools) + { + let exposed_tool_names = mcp_tools + .iter() + .chain(deferred_mcp_tools.iter()) + .flat_map(|tools| tools.keys().map(String::as_str)) + .collect::>(); + collect_unavailable_called_tools(input, &exposed_tool_names) + } else { + Vec::new() + }; + + let parallel_mcp_server_names = turn_context + .config + .mcp_servers + .get() + .iter() + .filter_map(|(server_name, server_config)| { + server_config + .supports_parallel_tool_calls + .then_some(server_name.clone()) + }) + .collect::>(); + + Ok(Arc::new(ToolRouter::from_config( + &turn_context.tools_config, + ToolRouterParams { + mcp_tools, + deferred_mcp_tools, + unavailable_called_tools, + parallel_mcp_server_names, + discoverable_tools, + dynamic_tools: turn_context.dynamic_tools.as_slice(), + }, + ))) +} + +#[derive(Debug)] +struct SamplingRequestResult { + needs_follow_up: bool, + last_agent_message: Option, +} + +/// Ephemeral per-response state for streaming a single proposed plan. +/// This is intentionally not persisted or stored in session/state since it +/// only exists while a response is actively streaming. The final plan text +/// is extracted from the completed assistant message. +/// Tracks a single proposed plan item across a streaming response. +struct ProposedPlanItemState { + item_id: String, + started: bool, + completed: bool, +} + +/// Aggregated state used only while streaming a plan-mode response. +/// Includes per-item parsers, deferred agent message bookkeeping, and the plan item lifecycle. +struct PlanModeStreamState { + /// Agent message items started by the model but deferred until we see non-plan text. + pending_agent_message_items: HashMap, + /// Agent message items whose start notification has been emitted. + started_agent_message_items: HashSet, + /// Leading whitespace buffered until we see non-whitespace text for an item. + leading_whitespace_by_item: HashMap, + /// Tracks plan item lifecycle while streaming plan output. + plan_item_state: ProposedPlanItemState, +} + +impl PlanModeStreamState { + fn new(turn_id: &str) -> Self { + Self { + pending_agent_message_items: HashMap::new(), + started_agent_message_items: HashSet::new(), + leading_whitespace_by_item: HashMap::new(), + plan_item_state: ProposedPlanItemState::new(turn_id), + } + } +} + +#[derive(Debug, Default)] +pub(super) struct AssistantMessageStreamParsers { + plan_mode: bool, + parsers_by_item: HashMap, +} + +type ParsedAssistantTextDelta = AssistantTextChunk; + +impl AssistantMessageStreamParsers { + pub(super) fn new(plan_mode: bool) -> Self { + Self { + plan_mode, + parsers_by_item: HashMap::new(), + } + } + + fn parser_mut(&mut self, item_id: &str) -> &mut AssistantTextStreamParser { + let plan_mode = self.plan_mode; + self.parsers_by_item + .entry(item_id.to_string()) + .or_insert_with(|| AssistantTextStreamParser::new(plan_mode)) + } + + pub(super) fn seed_item_text(&mut self, item_id: &str, text: &str) -> ParsedAssistantTextDelta { + if text.is_empty() { + return ParsedAssistantTextDelta::default(); + } + self.parser_mut(item_id).push_str(text) + } + + pub(super) fn parse_delta(&mut self, item_id: &str, delta: &str) -> ParsedAssistantTextDelta { + self.parser_mut(item_id).push_str(delta) + } + + pub(super) fn finish_item(&mut self, item_id: &str) -> ParsedAssistantTextDelta { + let Some(mut parser) = self.parsers_by_item.remove(item_id) else { + return ParsedAssistantTextDelta::default(); + }; + parser.finish() + } + + fn drain_finished(&mut self) -> Vec<(String, ParsedAssistantTextDelta)> { + let parsers_by_item = std::mem::take(&mut self.parsers_by_item); + parsers_by_item + .into_iter() + .map(|(item_id, mut parser)| (item_id, parser.finish())) + .collect() + } +} + +impl ProposedPlanItemState { + fn new(turn_id: &str) -> Self { + Self { + item_id: format!("{turn_id}-plan"), + started: false, + completed: false, + } + } + + async fn start(&mut self, sess: &Session, turn_context: &TurnContext) { + if self.started || self.completed { + return; + } + self.started = true; + let item = TurnItem::Plan(PlanItem { + id: self.item_id.clone(), + text: String::new(), + }); + sess.emit_turn_item_started(turn_context, &item).await; + } + + async fn push_delta(&mut self, sess: &Session, turn_context: &TurnContext, delta: &str) { + if self.completed { + return; + } + if delta.is_empty() { + return; + } + let event = PlanDeltaEvent { + thread_id: sess.conversation_id.to_string(), + turn_id: turn_context.sub_id.clone(), + item_id: self.item_id.clone(), + delta: delta.to_string(), + }; + sess.send_event(turn_context, EventMsg::PlanDelta(event)) + .await; + } + + async fn complete_with_text( + &mut self, + sess: &Session, + turn_context: &TurnContext, + text: String, + ) { + if self.completed || !self.started { + return; + } + self.completed = true; + let item = TurnItem::Plan(PlanItem { + id: self.item_id.clone(), + text, + }); + sess.emit_turn_item_completed(turn_context, item).await; + } +} + +/// In plan mode we defer agent message starts until the parser emits non-plan +/// text. The parser buffers each line until it can rule out a tag prefix, so +/// plan-only outputs never show up as empty assistant messages. +async fn maybe_emit_pending_agent_message_start( + sess: &Session, + turn_context: &TurnContext, + state: &mut PlanModeStreamState, + item_id: &str, +) { + if state.started_agent_message_items.contains(item_id) { + return; + } + if let Some(item) = state.pending_agent_message_items.remove(item_id) { + sess.emit_turn_item_started(turn_context, &item).await; + state + .started_agent_message_items + .insert(item_id.to_string()); + } +} + +/// Agent messages are text-only today; concatenate all text entries. +fn agent_message_text(item: &codex_protocol::items::AgentMessageItem) -> String { + item.content + .iter() + .map(|entry| match entry { + codex_protocol::items::AgentMessageContent::Text { text } => text.as_str(), + }) + .collect() +} + +pub(super) fn realtime_text_for_event(msg: &EventMsg) -> Option { + match msg { + EventMsg::AgentMessage(event) => Some(event.message.clone()), + EventMsg::ItemCompleted(event) => match &event.item { + TurnItem::AgentMessage(item) => Some(agent_message_text(item)), + _ => None, + }, + EventMsg::Error(_) + | EventMsg::Warning(_) + | EventMsg::RealtimeConversationStarted(_) + | EventMsg::RealtimeConversationSdp(_) + | EventMsg::RealtimeConversationRealtime(_) + | EventMsg::RealtimeConversationClosed(_) + | EventMsg::ModelReroute(_) + | EventMsg::ContextCompacted(_) + | EventMsg::ThreadRolledBack(_) + | EventMsg::TurnStarted(_) + | EventMsg::TurnComplete(_) + | EventMsg::TokenCount(_) + | EventMsg::UserMessage(_) + | EventMsg::AgentMessageDelta(_) + | EventMsg::AgentReasoning(_) + | EventMsg::AgentReasoningDelta(_) + | EventMsg::AgentReasoningRawContent(_) + | EventMsg::AgentReasoningRawContentDelta(_) + | EventMsg::AgentReasoningSectionBreak(_) + | EventMsg::SessionConfigured(_) + | EventMsg::ThreadNameUpdated(_) + | EventMsg::McpStartupUpdate(_) + | EventMsg::McpStartupComplete(_) + | EventMsg::McpToolCallBegin(_) + | EventMsg::McpToolCallEnd(_) + | EventMsg::WebSearchBegin(_) + | EventMsg::WebSearchEnd(_) + | EventMsg::ExecCommandBegin(_) + | EventMsg::ExecCommandOutputDelta(_) + | EventMsg::TerminalInteraction(_) + | EventMsg::ExecCommandEnd(_) + | EventMsg::PatchApplyBegin(_) + | EventMsg::PatchApplyEnd(_) + | EventMsg::ViewImageToolCall(_) + | EventMsg::ImageGenerationBegin(_) + | EventMsg::ImageGenerationEnd(_) + | EventMsg::ExecApprovalRequest(_) + | EventMsg::RequestPermissions(_) + | EventMsg::RequestUserInput(_) + | EventMsg::DynamicToolCallRequest(_) + | EventMsg::DynamicToolCallResponse(_) + | EventMsg::GuardianAssessment(_) + | EventMsg::ElicitationRequest(_) + | EventMsg::ApplyPatchApprovalRequest(_) + | EventMsg::DeprecationNotice(_) + | EventMsg::BackgroundEvent(_) + | EventMsg::UndoStarted(_) + | EventMsg::UndoCompleted(_) + | EventMsg::StreamError(_) + | EventMsg::TurnDiff(_) + | EventMsg::GetHistoryEntryResponse(_) + | EventMsg::McpListToolsResponse(_) + | EventMsg::ListSkillsResponse(_) + | EventMsg::RealtimeConversationListVoicesResponse(_) + | EventMsg::SkillsUpdateAvailable + | EventMsg::PlanUpdate(_) + | EventMsg::TurnAborted(_) + | EventMsg::ShutdownComplete + | EventMsg::EnteredReviewMode(_) + | EventMsg::ExitedReviewMode(_) + | EventMsg::RawResponseItem(_) + | EventMsg::ItemStarted(_) + | EventMsg::HookStarted(_) + | EventMsg::HookCompleted(_) + | EventMsg::AgentMessageContentDelta(_) + | EventMsg::PlanDelta(_) + | EventMsg::ReasoningContentDelta(_) + | EventMsg::ReasoningRawContentDelta(_) + | EventMsg::CollabAgentSpawnBegin(_) + | EventMsg::CollabAgentSpawnEnd(_) + | EventMsg::CollabAgentInteractionBegin(_) + | EventMsg::CollabAgentInteractionEnd(_) + | EventMsg::CollabWaitingBegin(_) + | EventMsg::CollabWaitingEnd(_) + | EventMsg::CollabCloseBegin(_) + | EventMsg::CollabCloseEnd(_) + | EventMsg::CollabResumeBegin(_) + | EventMsg::CollabResumeEnd(_) => None, + } +} + +/// Split the stream into normal assistant text vs. proposed plan content. +/// Normal text becomes AgentMessage deltas; plan content becomes PlanDelta + +/// TurnItem::Plan. +async fn handle_plan_segments( + sess: &Session, + turn_context: &TurnContext, + state: &mut PlanModeStreamState, + item_id: &str, + segments: Vec, +) { + for segment in segments { + match segment { + ProposedPlanSegment::Normal(delta) => { + if delta.is_empty() { + continue; + } + let has_non_whitespace = delta.chars().any(|ch| !ch.is_whitespace()); + if !has_non_whitespace && !state.started_agent_message_items.contains(item_id) { + let entry = state + .leading_whitespace_by_item + .entry(item_id.to_string()) + .or_default(); + entry.push_str(&delta); + continue; + } + let delta = if !state.started_agent_message_items.contains(item_id) { + if let Some(prefix) = state.leading_whitespace_by_item.remove(item_id) { + format!("{prefix}{delta}") + } else { + delta + } + } else { + delta + }; + maybe_emit_pending_agent_message_start(sess, turn_context, state, item_id).await; + + let event = AgentMessageContentDeltaEvent { + thread_id: sess.conversation_id.to_string(), + turn_id: turn_context.sub_id.clone(), + item_id: item_id.to_string(), + delta, + }; + sess.send_event(turn_context, EventMsg::AgentMessageContentDelta(event)) + .await; + } + ProposedPlanSegment::ProposedPlanStart => { + if !state.plan_item_state.completed { + state.plan_item_state.start(sess, turn_context).await; + } + } + ProposedPlanSegment::ProposedPlanDelta(delta) => { + if !state.plan_item_state.completed { + if !state.plan_item_state.started { + state.plan_item_state.start(sess, turn_context).await; + } + state + .plan_item_state + .push_delta(sess, turn_context, &delta) + .await; + } + } + ProposedPlanSegment::ProposedPlanEnd => {} + } + } +} + +async fn emit_streamed_assistant_text_delta( + sess: &Session, + turn_context: &TurnContext, + plan_mode_state: Option<&mut PlanModeStreamState>, + item_id: &str, + parsed: ParsedAssistantTextDelta, +) { + if parsed.is_empty() { + return; + } + if !parsed.citations.is_empty() { + // Citation extraction is intentionally local for now; we strip citations from display text + // but do not yet surface them in protocol events. + let _citations = parsed.citations; + } + if let Some(state) = plan_mode_state { + if !parsed.plan_segments.is_empty() { + handle_plan_segments(sess, turn_context, state, item_id, parsed.plan_segments).await; + } + return; + } + if parsed.visible_text.is_empty() { + return; + } + let event = AgentMessageContentDeltaEvent { + thread_id: sess.conversation_id.to_string(), + turn_id: turn_context.sub_id.clone(), + item_id: item_id.to_string(), + delta: parsed.visible_text, + }; + sess.send_event(turn_context, EventMsg::AgentMessageContentDelta(event)) + .await; +} + +/// Flush buffered assistant text parser state when an assistant message item ends. +async fn flush_assistant_text_segments_for_item( + sess: &Session, + turn_context: &TurnContext, + plan_mode_state: Option<&mut PlanModeStreamState>, + parsers: &mut AssistantMessageStreamParsers, + item_id: &str, +) { + let parsed = parsers.finish_item(item_id); + emit_streamed_assistant_text_delta(sess, turn_context, plan_mode_state, item_id, parsed).await; +} + +/// Flush any remaining buffered assistant text parser state at response completion. +async fn flush_assistant_text_segments_all( + sess: &Session, + turn_context: &TurnContext, + mut plan_mode_state: Option<&mut PlanModeStreamState>, + parsers: &mut AssistantMessageStreamParsers, +) { + for (item_id, parsed) in parsers.drain_finished() { + emit_streamed_assistant_text_delta( + sess, + turn_context, + plan_mode_state.as_deref_mut(), + &item_id, + parsed, + ) + .await; + } +} + +/// Emit completion for plan items by parsing the finalized assistant message. +async fn maybe_complete_plan_item_from_message( + sess: &Session, + turn_context: &TurnContext, + state: &mut PlanModeStreamState, + item: &ResponseItem, +) { + if let ResponseItem::Message { role, content, .. } = item + && role == "assistant" + { + let mut text = String::new(); + for entry in content { + if let ContentItem::OutputText { text: chunk } = entry { + text.push_str(chunk); + } + } + if let Some(plan_text) = extract_proposed_plan_text(&text) { + let (plan_text, _citations) = strip_citations(&plan_text); + if !state.plan_item_state.started { + state.plan_item_state.start(sess, turn_context).await; + } + state + .plan_item_state + .complete_with_text(sess, turn_context, plan_text) + .await; + } + } +} + +/// Emit a completed agent message in plan mode, respecting deferred starts. +async fn emit_agent_message_in_plan_mode( + sess: &Session, + turn_context: &TurnContext, + agent_message: codex_protocol::items::AgentMessageItem, + state: &mut PlanModeStreamState, +) { + let agent_message_id = agent_message.id.clone(); + let text = agent_message_text(&agent_message); + if text.trim().is_empty() { + state.pending_agent_message_items.remove(&agent_message_id); + state.started_agent_message_items.remove(&agent_message_id); + return; + } + + maybe_emit_pending_agent_message_start(sess, turn_context, state, &agent_message_id).await; + + if !state + .started_agent_message_items + .contains(&agent_message_id) + { + let start_item = state + .pending_agent_message_items + .remove(&agent_message_id) + .unwrap_or_else(|| { + TurnItem::AgentMessage(codex_protocol::items::AgentMessageItem { + id: agent_message_id.clone(), + content: Vec::new(), + phase: None, + memory_citation: None, + }) + }); + sess.emit_turn_item_started(turn_context, &start_item).await; + state + .started_agent_message_items + .insert(agent_message_id.clone()); + } + + sess.emit_turn_item_completed(turn_context, TurnItem::AgentMessage(agent_message)) + .await; + state.started_agent_message_items.remove(&agent_message_id); +} + +/// Emit completion for a plan-mode turn item, handling agent messages specially. +async fn emit_turn_item_in_plan_mode( + sess: &Session, + turn_context: &TurnContext, + turn_item: TurnItem, + previously_active_item: Option<&TurnItem>, + state: &mut PlanModeStreamState, +) { + match turn_item { + TurnItem::AgentMessage(agent_message) => { + emit_agent_message_in_plan_mode(sess, turn_context, agent_message, state).await; + } + _ => { + if previously_active_item.is_none() { + sess.emit_turn_item_started(turn_context, &turn_item).await; + } + sess.emit_turn_item_completed(turn_context, turn_item).await; + } + } +} + +/// Handle a completed assistant response item in plan mode, returning true if handled. +async fn handle_assistant_item_done_in_plan_mode( + sess: &Session, + turn_context: &TurnContext, + item: &ResponseItem, + state: &mut PlanModeStreamState, + previously_active_item: Option<&TurnItem>, + last_agent_message: &mut Option, +) -> bool { + if let ResponseItem::Message { role, .. } = item + && role == "assistant" + { + maybe_complete_plan_item_from_message(sess, turn_context, state, item).await; + + if let Some(turn_item) = + handle_non_tool_response_item(sess, turn_context, item, /*plan_mode*/ true).await + { + emit_turn_item_in_plan_mode( + sess, + turn_context, + turn_item, + previously_active_item, + state, + ) + .await; + } + + record_completed_response_item(sess, turn_context, item).await; + if let Some(agent_message) = last_assistant_message_from_item(item, /*plan_mode*/ true) { + *last_agent_message = Some(agent_message); + } + return true; + } + false +} + +async fn drain_in_flight( + in_flight: &mut FuturesOrdered>>, + sess: Arc, + turn_context: Arc, +) -> CodexResult<()> { + while let Some(res) = in_flight.next().await { + match res { + Ok(response_input) => { + let response_item = response_input.into(); + sess.record_conversation_items(&turn_context, std::slice::from_ref(&response_item)) + .await; + mark_thread_memory_mode_polluted_if_external_context( + sess.as_ref(), + turn_context.as_ref(), + &response_item, + ) + .await; + } + Err(err) => { + error_or_panic(format!("in-flight tool future failed during drain: {err}")); + } + } + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +#[instrument(level = "trace", + skip_all, + fields( + turn_id = %turn_context.sub_id, + model = %turn_context.model_info.slug + ) +)] +async fn try_run_sampling_request( + tool_runtime: ToolCallRuntime, + sess: Arc, + turn_context: Arc, + client_session: &mut ModelClientSession, + turn_metadata_header: Option<&str>, + turn_diff_tracker: SharedTurnDiffTracker, + server_model_warning_emitted_for_turn: &mut bool, + prompt: &Prompt, + cancellation_token: CancellationToken, +) -> CodexResult { + feedback_tags!( + model = turn_context.model_info.slug.clone(), + approval_policy = turn_context.approval_policy.value(), + sandbox_policy = turn_context.sandbox_policy.get(), + effort = turn_context.reasoning_effort, + auth_mode = sess.services.auth_manager.auth_mode(), + features = sess.features.enabled_features(), + ); + let mut stream = client_session + .stream( + prompt, + &turn_context.model_info, + &turn_context.session_telemetry, + turn_context.reasoning_effort, + turn_context.reasoning_summary, + turn_context.config.service_tier, + turn_metadata_header, + ) + .instrument(trace_span!("stream_request")) + .or_cancel(&cancellation_token) + .await??; + let mut in_flight: FuturesOrdered>> = + FuturesOrdered::new(); + let mut needs_follow_up = false; + let mut last_agent_message: Option = None; + let mut active_item: Option = None; + let mut should_emit_turn_diff = false; + let plan_mode = turn_context.collaboration_mode.mode == ModeKind::Plan; + let mut assistant_message_stream_parsers = AssistantMessageStreamParsers::new(plan_mode); + let mut plan_mode_state = plan_mode.then(|| PlanModeStreamState::new(&turn_context.sub_id)); + let receiving_span = trace_span!("receiving_stream"); + let outcome: CodexResult = loop { + let handle_responses = trace_span!( + parent: &receiving_span, + "handle_responses", + otel.name = field::Empty, + tool_name = field::Empty, + from = field::Empty, + ); + + let event = match stream + .next() + .instrument(trace_span!(parent: &handle_responses, "receiving")) + .or_cancel(&cancellation_token) + .await + { + Ok(event) => event, + Err(codex_async_utils::CancelErr::Cancelled) => break Err(CodexErr::TurnAborted), + }; + + let event = match event { + Some(Ok(event)) => event, + Some(Err(err)) => break Err(err), + None => { + break Err(CodexErr::Stream( + "stream closed before response.completed".into(), + None, + )); + } + }; + + sess.services + .session_telemetry + .record_responses(&handle_responses, &event); + record_turn_ttft_metric(&turn_context, &event).await; + + match event { + ResponseEvent::Created => {} + ResponseEvent::OutputItemDone(item) => { + let previously_active_item = active_item.take(); + if let Some(previous) = previously_active_item.as_ref() + && matches!(previous, TurnItem::AgentMessage(_)) + { + let item_id = previous.id(); + flush_assistant_text_segments_for_item( + &sess, + &turn_context, + plan_mode_state.as_mut(), + &mut assistant_message_stream_parsers, + &item_id, + ) + .await; + } + if let Some(state) = plan_mode_state.as_mut() + && handle_assistant_item_done_in_plan_mode( + &sess, + &turn_context, + &item, + state, + previously_active_item.as_ref(), + &mut last_agent_message, + ) + .await + { + continue; + } + + let mut ctx = HandleOutputCtx { + sess: sess.clone(), + turn_context: turn_context.clone(), + tool_runtime: tool_runtime.clone(), + cancellation_token: cancellation_token.child_token(), + }; + + let output_result = + match handle_output_item_done(&mut ctx, item, previously_active_item) + .instrument(handle_responses) + .await + { + Ok(output_result) => output_result, + Err(err) => break Err(err), + }; + if let Some(tool_future) = output_result.tool_future { + in_flight.push_back(tool_future); + } + if let Some(agent_message) = output_result.last_agent_message { + last_agent_message = Some(agent_message); + } + needs_follow_up |= output_result.needs_follow_up; + } + ResponseEvent::OutputItemAdded(item) => { + if let Some(turn_item) = handle_non_tool_response_item( + sess.as_ref(), + turn_context.as_ref(), + &item, + plan_mode, + ) + .await + { + let mut turn_item = turn_item; + let mut seeded_parsed: Option = None; + let mut seeded_item_id: Option = None; + if matches!(turn_item, TurnItem::AgentMessage(_)) + && let Some(raw_text) = raw_assistant_output_text_from_item(&item) + { + let item_id = turn_item.id(); + let mut seeded = + assistant_message_stream_parsers.seed_item_text(&item_id, &raw_text); + if let TurnItem::AgentMessage(agent_message) = &mut turn_item { + agent_message.content = + vec![codex_protocol::items::AgentMessageContent::Text { + text: if plan_mode { + String::new() + } else { + std::mem::take(&mut seeded.visible_text) + }, + }]; + } + seeded_parsed = plan_mode.then_some(seeded); + seeded_item_id = Some(item_id); + } + if let Some(state) = plan_mode_state.as_mut() + && matches!(turn_item, TurnItem::AgentMessage(_)) + { + let item_id = turn_item.id(); + state + .pending_agent_message_items + .insert(item_id, turn_item.clone()); + } else { + sess.emit_turn_item_started(&turn_context, &turn_item).await; + } + if let (Some(state), Some(item_id), Some(parsed)) = ( + plan_mode_state.as_mut(), + seeded_item_id.as_deref(), + seeded_parsed, + ) { + emit_streamed_assistant_text_delta( + &sess, + &turn_context, + Some(state), + item_id, + parsed, + ) + .await; + } + active_item = Some(turn_item); + } + } + ResponseEvent::ServerModel(server_model) => { + if !*server_model_warning_emitted_for_turn + && sess + .maybe_warn_on_server_model_mismatch(&turn_context, server_model) + .await + { + *server_model_warning_emitted_for_turn = true; + } + } + ResponseEvent::ServerReasoningIncluded(included) => { + sess.set_server_reasoning_included(included).await; + } + ResponseEvent::RateLimits(snapshot) => { + // Update internal state with latest rate limits, but defer sending until + // token usage is available to avoid duplicate TokenCount events. + sess.update_rate_limits(&turn_context, snapshot).await; + } + ResponseEvent::ModelsEtag(etag) => { + // Update internal state with latest models etag + sess.services.models_manager.refresh_if_new_etag(etag).await; + } + ResponseEvent::Completed { + response_id: _, + token_usage, + } => { + flush_assistant_text_segments_all( + &sess, + &turn_context, + plan_mode_state.as_mut(), + &mut assistant_message_stream_parsers, + ) + .await; + sess.update_token_usage_info(&turn_context, token_usage.as_ref()) + .await; + should_emit_turn_diff = true; + + break Ok(SamplingRequestResult { + needs_follow_up, + last_agent_message, + }); + } + ResponseEvent::OutputTextDelta(delta) => { + // In review child threads, suppress assistant text deltas; the + // UI will show a selection popup from the final ReviewOutput. + if let Some(active) = active_item.as_ref() { + let item_id = active.id(); + if matches!(active, TurnItem::AgentMessage(_)) { + let parsed = assistant_message_stream_parsers.parse_delta(&item_id, &delta); + emit_streamed_assistant_text_delta( + &sess, + &turn_context, + plan_mode_state.as_mut(), + &item_id, + parsed, + ) + .await; + } else { + let event = AgentMessageContentDeltaEvent { + thread_id: sess.conversation_id.to_string(), + turn_id: turn_context.sub_id.clone(), + item_id, + delta, + }; + sess.send_event(&turn_context, EventMsg::AgentMessageContentDelta(event)) + .await; + } + } else { + error_or_panic("OutputTextDelta without active item".to_string()); + } + } + ResponseEvent::ReasoningSummaryDelta { + delta, + summary_index, + } => { + if let Some(active) = active_item.as_ref() { + let event = ReasoningContentDeltaEvent { + thread_id: sess.conversation_id.to_string(), + turn_id: turn_context.sub_id.clone(), + item_id: active.id(), + delta, + summary_index, + }; + sess.send_event(&turn_context, EventMsg::ReasoningContentDelta(event)) + .await; + } else { + error_or_panic("ReasoningSummaryDelta without active item".to_string()); + } + } + ResponseEvent::ReasoningSummaryPartAdded { summary_index } => { + if let Some(active) = active_item.as_ref() { + let event = + EventMsg::AgentReasoningSectionBreak(AgentReasoningSectionBreakEvent { + item_id: active.id(), + summary_index, + }); + sess.send_event(&turn_context, event).await; + } else { + error_or_panic("ReasoningSummaryPartAdded without active item".to_string()); + } + } + ResponseEvent::ReasoningContentDelta { + delta, + content_index, + } => { + if let Some(active) = active_item.as_ref() { + let event = ReasoningRawContentDeltaEvent { + thread_id: sess.conversation_id.to_string(), + turn_id: turn_context.sub_id.clone(), + item_id: active.id(), + delta, + content_index, + }; + sess.send_event(&turn_context, EventMsg::ReasoningRawContentDelta(event)) + .await; + } else { + error_or_panic("ReasoningRawContentDelta without active item".to_string()); + } + } + } + }; + + flush_assistant_text_segments_all( + &sess, + &turn_context, + plan_mode_state.as_mut(), + &mut assistant_message_stream_parsers, + ) + .await; + + drain_in_flight(&mut in_flight, sess.clone(), turn_context.clone()).await?; + + if cancellation_token.is_cancelled() { + return Err(CodexErr::TurnAborted); + } + + if should_emit_turn_diff { + let unified_diff = { + let mut tracker = turn_diff_tracker.lock().await; + tracker.get_unified_diff() + }; + if let Ok(Some(unified_diff)) = unified_diff { + let msg = EventMsg::TurnDiff(TurnDiffEvent { unified_diff }); + sess.clone().send_event(&turn_context, msg).await; + } + } + + outcome +} + +pub(crate) fn get_last_assistant_message_from_turn(responses: &[ResponseItem]) -> Option { + for item in responses.iter().rev() { + if let Some(message) = last_assistant_message_from_item(item, /*plan_mode*/ false) { + return Some(message); + } + } + None +} From 65cc12d72e25723aece48edd6ff93dd288b6c042 Mon Sep 17 00:00:00 2001 From: Jeff Harris Date: Thu, 16 Apr 2026 15:43:51 -0700 Subject: [PATCH 16/16] Use codex-auto-review for guardian reviews (#18169) ## Summary This is the minimal client-side follow-up for the Codex Auto Review model slug rollout. It updates the guardian reviewer preferred model from `gpt-5.4` to `codex-auto-review`, so the client can rely on the backend catalog + Statsig mapping instead of hardcoding the GPT-5.4 slug. Context: https://openai.slack.com/archives/C0AF9328RL0/p1775777479388369?thread_ts=1775773094.071629&cid=C0AF9328RL0 ## Testing - `cargo fmt --package codex-core --check` - `cargo test -p codex-core guardian::` - `bazel test --experimental_remote_downloader= --test_output=errors //codex-rs/core:core-unit-tests --test_arg=guardian` --- codex-rs/core/src/guardian/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codex-rs/core/src/guardian/mod.rs b/codex-rs/core/src/guardian/mod.rs index 67e9a828ee..4fa150a232 100644 --- a/codex-rs/core/src/guardian/mod.rs +++ b/codex-rs/core/src/guardian/mod.rs @@ -34,7 +34,7 @@ pub(crate) use review::review_approval_request_with_cancel; pub(crate) use review::routes_approval_to_guardian; pub(crate) use review_session::GuardianReviewSessionManager; -const GUARDIAN_PREFERRED_MODEL: &str = "gpt-5.4"; +const GUARDIAN_PREFERRED_MODEL: &str = "codex-auto-review"; pub(crate) const GUARDIAN_REVIEW_TIMEOUT: Duration = Duration::from_secs(90); pub(crate) const GUARDIAN_REVIEWER_NAME: &str = "guardian"; const GUARDIAN_MAX_MESSAGE_TRANSCRIPT_TOKENS: usize = 10_000;