From 5eebeb816905f44fd847124ea610922795362cbc Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 26 Jun 2026 00:13:43 +0100 Subject: [PATCH 01/38] Project executor skills through World State (#30088) ## Why A selected executor environment can be unavailable in one model step and ready in the next. The model should see its skills only while that environment is ready, without rescanning stable files on every sample. The product assumption is simple: - an environment ID names one stable logical environment; - the selected root contents do not change during the thread. ## Behavior ```text E1 unavailable -> do not show E1 skills E1 ready -> discover once, cache, show through World State E1 unavailable -> hide skills, keep cache E1 ready again -> reuse cache, show skills again resume -> create a new thread cache and discover again ``` The cache key is the full `SelectedCapabilityRoot`. Availability does not invalidate it; dropping the extension's thread state does. The step supplies the ready selected roots directly. They do not have to be turn environments: ```text turn environment: laptop selected root: worker:/plugins/lint-fix worker ready -> lint-fix skills are visible ``` ## What changes - Keeps executor skill catalogs in the existing skills extension. - Passes the roots resolved as ready for the step into World State contributors. - Loads each ready selected root at most once per thread. - Contributes the executor catalog as the `skills` World State section. - Uses the exact step catalog for explicit skill selection and body reads. - Leaves host and orchestrator skill behavior where it already lives. Taking a step snapshot itself does not add an RPC. Executor filesystem calls happen only on the first discovery of a stable root for that thread. ## What does not change - No filesystem watcher or content-based invalidation. - No retry/generation framework. - No skill runtime migration into core. - No general rewrite of the skills extension. ## Stack 1. Extension-owned World State sections. 2. **This PR:** project cached executor skills through World State. 3. Pin one MCP runtime to each model step. 4. Project selected MCP/app/connector metadata by environment availability. 5. One end-to-end integration scenario. --- codex-rs/core/src/session/world_state.rs | 6 + .../src/contributors/world_state.rs | 3 + codex-rs/ext/skills/src/extension.rs | 54 +++-- codex-rs/ext/skills/src/lib.rs | 1 + codex-rs/ext/skills/src/sources.rs | 5 + codex-rs/ext/skills/src/state.rs | 77 ++++++-- codex-rs/ext/skills/src/world_state.rs | 62 ++++++ codex-rs/ext/skills/tests/skills_extension.rs | 184 +++++++++++++----- 8 files changed, 322 insertions(+), 70 deletions(-) create mode 100644 codex-rs/ext/skills/src/world_state.rs diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs index 2541972252..10d6a5d0d1 100644 --- a/codex-rs/core/src/session/world_state.rs +++ b/codex-rs/core/src/session/world_state.rs @@ -36,12 +36,18 @@ impl Session { ); } let environments = step_context.environments.to_selections(); + let ready_selected_capability_roots = step_context + .selected_capability_roots + .iter() + .map(|root| root.selected_root().clone()) + .collect::>(); for contributor in self.services.extensions.context_contributors() { for section in contributor .contribute_world_state(WorldStateContributionInput { thread_id: self.thread_id(), turn_id: turn_context.sub_id.as_str(), environments: &environments, + ready_selected_capability_roots: &ready_selected_capability_roots, session_store: &self.services.session_extension_data, thread_store: &self.services.thread_extension_data, turn_store: turn_context.extension_data.as_ref(), diff --git a/codex-rs/ext/extension-api/src/contributors/world_state.rs b/codex-rs/ext/extension-api/src/contributors/world_state.rs index b5bb548105..23f180826d 100644 --- a/codex-rs/ext/extension-api/src/contributors/world_state.rs +++ b/codex-rs/ext/extension-api/src/contributors/world_state.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use codex_protocol::ThreadId; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::protocol::TurnEnvironmentSelection; use serde_json::Value; @@ -11,6 +12,8 @@ pub struct WorldStateContributionInput<'a> { pub thread_id: ThreadId, pub turn_id: &'a str, pub environments: &'a [TurnEnvironmentSelection], + /// Selected roots whose stable environments are ready in this sampling step. + pub ready_selected_capability_roots: &'a [SelectedCapabilityRoot], pub session_store: &'a ExtensionData, pub thread_store: &'a ExtensionData, pub turn_store: &'a ExtensionData, diff --git a/codex-rs/ext/skills/src/extension.rs b/codex-rs/ext/skills/src/extension.rs index de9facd8c2..0b582e9387 100644 --- a/codex-rs/ext/skills/src/extension.rs +++ b/codex-rs/ext/skills/src/extension.rs @@ -18,8 +18,9 @@ use codex_extension_api::ToolContributor; use codex_extension_api::ToolExecutor; use codex_extension_api::TurnInputContext; use codex_extension_api::TurnInputContributor; +use codex_extension_api::WorldStateContributionInput; +use codex_extension_api::WorldStateSectionContribution; use codex_mcp::McpResourceClient; -use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::WarningEvent; @@ -40,9 +41,11 @@ use crate::render::truncate_main_prompt_contents; use crate::render::truncate_utf8_to_bytes; use crate::selection::collect_explicit_skill_mentions; use crate::sources::SkillProviders; +use crate::state::ExecutorSkillsStepState; use crate::state::SkillsThreadState; use crate::state::SkillsTurnState; use crate::tools::skill_tools; +use crate::world_state::executor_skills_world_state_section; struct SkillsExtension { providers: SkillProviders, @@ -56,18 +59,12 @@ where { fn on_thread_start<'a>(&'a self, input: ThreadStartInput<'a, C>) -> ExtensionFuture<'a, ()> { Box::pin(async move { - let selected_roots = input - .thread_store - .get::>() - .map(|selected_roots| selected_roots.as_ref().clone()) - .unwrap_or_default(); let orchestrator_skills_available = !input .environments .iter() .any(|environment| environment.environment_id == LOCAL_ENVIRONMENT_ID); input.thread_store.insert(SkillsThreadState::new( (self.config_from_host)(input.config), - selected_roots, orchestrator_skills_available, )); }) @@ -92,7 +89,6 @@ where let orchestrator_skills_available = true; thread_store.insert(SkillsThreadState::new( next_config, - Vec::new(), orchestrator_skills_available, )); } @@ -120,7 +116,7 @@ where .list_skills( SkillListQuery { turn_id: thread_store.level_id().to_string(), - executor_roots: thread_state.selected_roots().to_vec(), + executor_roots: Vec::new(), host_snapshot: None, include_host_skills: false, include_bundled_skills: config.bundled_skills_enabled, @@ -139,6 +135,39 @@ where .collect() }) } + + fn contribute_world_state<'a>( + &'a self, + input: WorldStateContributionInput<'a>, + ) -> ExtensionFuture<'a, Vec> { + Box::pin(async move { + let Some(thread_state) = input.thread_store.get::() else { + return Vec::new(); + }; + let config = thread_state.config(); + let catalog = thread_state + .executor_catalog_snapshot( + &self.providers, + SkillListQuery { + turn_id: input.turn_id.to_string(), + executor_roots: input.ready_selected_capability_roots.to_vec(), + host_snapshot: None, + include_host_skills: false, + include_bundled_skills: config.bundled_skills_enabled, + include_orchestrator_skills: false, + mcp_resources: input.session_store.get::(), + }, + ) + .await; + input + .turn_store + .insert(ExecutorSkillsStepState(catalog.clone())); + vec![executor_skills_world_state_section( + &catalog, + config.include_instructions, + )] + }) + } } impl ToolContributor for SkillsExtension @@ -187,14 +216,17 @@ where let host_snapshot = turn_store.get::(); let query = SkillListQuery { turn_id: input.turn_id.clone(), - executor_roots: thread_state.selected_roots().to_vec(), + executor_roots: Vec::new(), host_snapshot: host_snapshot.clone(), include_host_skills: true, include_bundled_skills: config.bundled_skills_enabled, include_orchestrator_skills: thread_state.orchestrator_skills_enabled(), mcp_resources: session_store.get::(), }; - let catalog = self.list_skills(query, &thread_state).await; + let mut catalog = self.list_skills(query, &thread_state).await; + if let Some(executor_skills) = turn_store.get::() { + catalog.extend(executor_skills.0.clone()); + } for warning in &catalog.warnings { self.emit_warning(&input.turn_id, warning.clone()); } diff --git a/codex-rs/ext/skills/src/lib.rs b/codex-rs/ext/skills/src/lib.rs index 8abdf6a93e..4a004ff9ac 100644 --- a/codex-rs/ext/skills/src/lib.rs +++ b/codex-rs/ext/skills/src/lib.rs @@ -8,6 +8,7 @@ mod selection; mod sources; mod state; mod tools; +mod world_state; pub use config::SkillsExtensionConfig; pub use extension::install; diff --git a/codex-rs/ext/skills/src/sources.rs b/codex-rs/ext/skills/src/sources.rs index d558c13b3b..a75a047e09 100644 --- a/codex-rs/ext/skills/src/sources.rs +++ b/codex-rs/ext/skills/src/sources.rs @@ -135,6 +135,11 @@ impl SkillProviders { Ok(catalog) } + pub(crate) async fn list_executor_for_turn(&self, query: SkillListQuery) -> SkillCatalog { + self.list_matching(&query, |source| source.kind == SkillSourceKind::Executor) + .await + } + async fn list_matching( &self, query: &SkillListQuery, diff --git a/codex-rs/ext/skills/src/state.rs b/codex-rs/ext/skills/src/state.rs index 0b751ebdc2..2f1a84c46b 100644 --- a/codex-rs/ext/skills/src/state.rs +++ b/codex-rs/ext/skills/src/state.rs @@ -18,6 +18,7 @@ use crate::catalog::SkillProviderResult; use crate::catalog::SkillReadResult; use crate::catalog::SkillResourceId; use crate::catalog::SkillSourceKind; +use crate::provider::SkillListQuery; use crate::provider::SkillReadRequest; use crate::sources::SkillProviders; @@ -26,21 +27,17 @@ const MAX_CACHED_ORCHESTRATOR_CONTENT_BYTES: usize = 8 * 1024 * 1024; pub(crate) struct SkillsThreadState { config: Mutex, - selected_roots: Vec, orchestrator_skills_available: bool, + executor_cache: Mutex>, orchestrator_cache: Mutex>>, } impl SkillsThreadState { - pub(crate) fn new( - config: SkillsExtensionConfig, - selected_roots: Vec, - orchestrator_skills_available: bool, - ) -> Self { + pub(crate) fn new(config: SkillsExtensionConfig, orchestrator_skills_available: bool) -> Self { Self { config: Mutex::new(config), - selected_roots, orchestrator_skills_available, + executor_cache: Mutex::new(Vec::new()), orchestrator_cache: Mutex::new(None), } } @@ -59,14 +56,33 @@ impl SkillsThreadState { .unwrap_or_else(std::sync::PoisonError::into_inner) = config; } - pub(crate) fn selected_roots(&self) -> &[SelectedCapabilityRoot] { - &self.selected_roots - } - pub(crate) fn orchestrator_skills_enabled(&self) -> bool { self.orchestrator_skills_available && self.config().orchestrator_skills_enabled } + /// Returns catalogs for stable selected roots. + /// + /// The first catalog returned for a root remains cached until this thread state is dropped. + /// Environment availability only controls whether the root is projected into the current + /// step; it never invalidates the cache. There is intentionally no filesystem watcher or + /// content-based invalidation because selected environment roots are treated as stable. + pub(crate) async fn executor_catalog_snapshot( + &self, + providers: &SkillProviders, + mut query: SkillListQuery, + ) -> SkillCatalog { + let roots = std::mem::take(&mut query.executor_roots); + let mut catalog = SkillCatalog::default(); + for root in roots { + query.executor_roots = vec![root.clone()]; + catalog.extend( + self.executor_root_catalog(providers, root, query.clone()) + .await, + ); + } + catalog + } + pub(crate) async fn orchestrator_catalog_snapshot( &self, mcp_resources: Option<&McpResourceClient>, @@ -140,6 +156,42 @@ impl SkillsThreadState { *cache = Some(Arc::clone(&next_cache)); next_cache } + + async fn executor_root_catalog( + &self, + providers: &SkillProviders, + root: SelectedCapabilityRoot, + query: SkillListQuery, + ) -> SkillCatalog { + if let Some(cached) = self + .executor_cache + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .iter() + .find(|cached| cached.root == root) + { + return cached.catalog.clone(); + } + + let discovered = providers.list_executor_for_turn(query).await; + let mut cache = self + .executor_cache + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(cached) = cache.iter().find(|cached| cached.root == root) { + return cached.catalog.clone(); + } + cache.push(CachedExecutorCatalog { + root, + catalog: discovered.clone(), + }); + discovered + } +} + +struct CachedExecutorCatalog { + root: SelectedCapabilityRoot, + catalog: SkillCatalog, } struct OrchestratorGenerationCache { @@ -204,3 +256,6 @@ pub(crate) struct SkillsTurnState { pub(crate) warnings: Vec, pub(crate) main_prompts_injected: bool, } + +#[derive(Clone, Debug, Default)] +pub(crate) struct ExecutorSkillsStepState(pub(crate) SkillCatalog); diff --git a/codex-rs/ext/skills/src/world_state.rs b/codex-rs/ext/skills/src/world_state.rs new file mode 100644 index 0000000000..ce36427b0c --- /dev/null +++ b/codex-rs/ext/skills/src/world_state.rs @@ -0,0 +1,62 @@ +use codex_extension_api::ContextualUserFragment; +use codex_extension_api::PreviousWorldStateSection; +use codex_extension_api::RenderedWorldStateFragment; +use codex_extension_api::WorldStateSectionContribution; +use codex_protocol::protocol::SKILLS_INSTRUCTIONS_CLOSE_TAG; +use codex_protocol::protocol::SKILLS_INSTRUCTIONS_OPEN_TAG; +use serde_json::json; + +use crate::catalog::SkillCatalog; +use crate::render::available_skills_fragment; + +pub(crate) const SKILLS_WORLD_STATE_ID: &str = "skills"; +const NO_EXECUTOR_SKILLS_BODY: &str = + "\n## Skills update\nNo selected-environment skills are currently available.\n"; +const HIDDEN_EXECUTOR_SKILLS_BODY: &str = "\n## Skills update\nSelected-environment skills are not listed automatically. Explicit skill mentions can still be resolved when available.\n"; + +pub(crate) fn executor_skills_world_state_section( + catalog: &SkillCatalog, + include_instructions: bool, +) -> WorldStateSectionContribution { + let body = if include_instructions { + available_skills_fragment(catalog).map(|fragment| fragment.body()) + } else { + None + }; + let snapshot = json!({ + "body": body, + "includeInstructions": include_instructions, + }); + + WorldStateSectionContribution::new(SKILLS_WORLD_STATE_ID, snapshot, move |previous| { + let previous_is_absent = matches!(&previous, PreviousWorldStateSection::Absent); + if let PreviousWorldStateSection::Known(previous) = &previous { + let previous_body = previous.get("body").and_then(serde_json::Value::as_str); + let previous_include_instructions = previous + .get("includeInstructions") + .and_then(serde_json::Value::as_bool); + if previous_body == body.as_deref() + && previous_include_instructions == Some(include_instructions) + { + return None; + } + } + + let body = match body.as_deref() { + Some(body) => body, + None if previous_is_absent => return None, + None if !include_instructions => HIDDEN_EXECUTOR_SKILLS_BODY, + None => NO_EXECUTOR_SKILLS_BODY, + }; + Some(RenderedWorldStateFragment::new( + "developer", + (SKILLS_INSTRUCTIONS_OPEN_TAG, SKILLS_INSTRUCTIONS_CLOSE_TAG), + body, + )) + }) + .with_legacy_matcher(|role, text| { + role == "developer" + && text.trim_start().starts_with(SKILLS_INSTRUCTIONS_OPEN_TAG) + && text.trim_end().ends_with(SKILLS_INSTRUCTIONS_CLOSE_TAG) + }) +} diff --git a/codex-rs/ext/skills/tests/skills_extension.rs b/codex-rs/ext/skills/tests/skills_extension.rs index 761dc34d71..7e9ea01c52 100644 --- a/codex-rs/ext/skills/tests/skills_extension.rs +++ b/codex-rs/ext/skills/tests/skills_extension.rs @@ -15,10 +15,12 @@ use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionEventSink; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::NoopTurnItemEmitter; +use codex_extension_api::PreviousWorldStateSection; use codex_extension_api::ThreadStartInput; use codex_extension_api::ToolCall; use codex_extension_api::ToolPayload; use codex_extension_api::TurnInputContext; +use codex_extension_api::WorldStateContributionInput; use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::protocol::Event; @@ -28,6 +30,7 @@ use codex_protocol::protocol::SKILLS_INSTRUCTIONS_OPEN_TAG; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SkillScope; use codex_protocol::protocol::TruncationPolicy; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::user_input::UserInput; use codex_skills_extension::SkillProviders; use codex_skills_extension::SkillsExtensionConfig; @@ -144,9 +147,9 @@ async fn installed_extension_uses_host_service_snapshot() -> TestResult { } #[tokio::test] -async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_input() -> TestResult -{ +async fn selected_executor_catalog_follows_step_availability_and_reuses_its_cache() -> TestResult { let read_requests = Arc::new(Mutex::new(Vec::new())); + let list_calls = Arc::new(AtomicUsize::new(0)); let executor_provider = Arc::new(StaticSkillProvider { catalog: SkillCatalog { entries: vec![test_entry( @@ -158,7 +161,7 @@ async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_in warnings: Vec::new(), }, read_requests: Arc::clone(&read_requests), - list_calls: None, + list_calls: Some(Arc::clone(&list_calls)), fail_first_list: false, }); let providers = SkillProviders::new().with_executor_provider(executor_provider); @@ -168,13 +171,13 @@ async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_in let session_store = ExtensionData::new("session"); let thread_store = ExtensionData::new("thread"); - thread_store.insert(vec![SelectedCapabilityRoot { + let selected_roots = vec![SelectedCapabilityRoot { id: "lint-fix".to_string(), location: CapabilityRootLocation::Environment { environment_id: "env-1".to_string(), path: PathUri::parse("file:///skills/lint-fix").expect("skill root URI"), }, - }]); + }]; let session_source = SessionSource::Cli; let config = default_config(); registry.thread_lifecycle_contributors()[0] @@ -191,20 +194,36 @@ async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_in let prompt_fragments = registry.context_contributors()[0] .contribute_thread_context(&session_store, &thread_store) .await; - assert_eq!(1, prompt_fragments.len()); + assert!(prompt_fragments.is_empty()); + + let turn_store = ExtensionData::new("turn-1"); + let turn_environment = TurnEnvironmentSelection { + environment_id: "turn-env".to_string(), + cwd: PathUri::parse("file:///workspace").expect("cwd URI"), + }; + let available_sections = registry.context_contributors()[0] + .contribute_world_state(WorldStateContributionInput { + thread_id: codex_protocol::ThreadId::new(), + turn_id: "turn-1", + environments: std::slice::from_ref(&turn_environment), + ready_selected_capability_roots: &selected_roots, + session_store: &session_store, + thread_store: &thread_store, + turn_store: &turn_store, + }) + .await; + assert_eq!(1, available_sections.len()); + let available_snapshot = available_sections[0].snapshot().clone(); + let available_fragment = available_sections[0] + .render_diff(PreviousWorldStateSection::Absent) + .ok_or("available skills should render")?; + assert!(available_fragment.body().contains("lint-fix")); assert!( - prompt_fragments[0] - .text() - .starts_with(SKILLS_INSTRUCTIONS_OPEN_TAG) - ); - assert!(prompt_fragments[0].text().contains("lint-fix")); - assert!( - prompt_fragments[0] - .text() + available_fragment + .body() .contains("(environment resource: skill://executor/lint-fix/SKILL.md)") ); - let turn_store = ExtensionData::new("turn-1"); let fragments = registry.turn_input_contributors()[0] .contribute( TurnInputContext { @@ -233,30 +252,86 @@ async fn selected_executor_catalog_is_context_and_selected_entrypoint_is_turn_in )], read_request_keys(&read_requests) ); - let rebuilt_prompt_fragments = registry.context_contributors()[0] - .contribute_thread_context(&session_store, &thread_store) + let unavailable_turn_store = ExtensionData::new("turn-2"); + let unavailable_sections = registry.context_contributors()[0] + .contribute_world_state(WorldStateContributionInput { + thread_id: codex_protocol::ThreadId::new(), + turn_id: "turn-2", + environments: &[], + ready_selected_capability_roots: &[], + session_store: &session_store, + thread_store: &thread_store, + turn_store: &unavailable_turn_store, + }) .await; - assert_eq!(1, rebuilt_prompt_fragments.len()); - assert!(rebuilt_prompt_fragments[0].text().contains("lint-fix")); + let unavailable_snapshot = unavailable_sections[0].snapshot().clone(); + let unavailable_fragment = unavailable_sections[0] + .render_diff(PreviousWorldStateSection::Known(&available_snapshot)) + .ok_or("removed skills should render")?; + assert!( + unavailable_fragment + .body() + .contains("No selected-environment skills") + ); - let next_turn_store = ExtensionData::new("turn-2"); - let next_fragments = registry.turn_input_contributors()[0] - .contribute( - TurnInputContext { - turn_id: "turn-2".to_string(), - user_input: vec![UserInput::Text { - text: "no skill this time".to_string(), - text_elements: Vec::new(), - }], - environments: Vec::new(), - }, - &session_store, - &thread_store, - &next_turn_store, - ) + let restored_turn_store = ExtensionData::new("turn-3"); + let restored_sections = registry.context_contributors()[0] + .contribute_world_state(WorldStateContributionInput { + thread_id: codex_protocol::ThreadId::new(), + turn_id: "turn-3", + environments: &[turn_environment], + ready_selected_capability_roots: &selected_roots, + session_store: &session_store, + thread_store: &thread_store, + turn_store: &restored_turn_store, + }) .await; + let restored_snapshot = restored_sections[0].snapshot().clone(); + let restored_fragment = restored_sections[0] + .render_diff(PreviousWorldStateSection::Known(&unavailable_snapshot)) + .ok_or("restored skills should render")?; + assert!(restored_fragment.body().contains("lint-fix")); + assert_eq!(1, list_calls.load(Ordering::Relaxed)); - assert!(next_fragments.is_empty()); + let mut listing_disabled_config = config.clone(); + listing_disabled_config.include_instructions = false; + registry.config_contributors()[0].on_config_changed( + &session_store, + &thread_store, + &config, + &listing_disabled_config, + ); + let listing_disabled_turn_store = ExtensionData::new("turn-4"); + let listing_disabled_sections = registry.context_contributors()[0] + .contribute_world_state(WorldStateContributionInput { + thread_id: codex_protocol::ThreadId::new(), + turn_id: "turn-4", + environments: &[], + ready_selected_capability_roots: &selected_roots, + session_store: &session_store, + thread_store: &thread_store, + turn_store: &listing_disabled_turn_store, + }) + .await; + let listing_disabled_fragment = listing_disabled_sections[0] + .render_diff(PreviousWorldStateSection::Known(&restored_snapshot)) + .ok_or("disabled skill listing should render")?; + assert_eq!( + "\n## Skills update\nSelected-environment skills are not listed automatically. Explicit skill mentions can still be resolved when available.\n", + listing_disabled_fragment.body() + ); + let mut normalized_listing_disabled_snapshot = listing_disabled_sections[0].snapshot().clone(); + normalized_listing_disabled_snapshot + .as_object_mut() + .ok_or("skills snapshot should be an object")? + .remove("body"); + assert!( + listing_disabled_sections[0] + .render_diff(PreviousWorldStateSection::Known( + &normalized_listing_disabled_snapshot + )) + .is_none() + ); Ok(()) } @@ -488,18 +563,16 @@ async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> Te let registry = builder.build(); let session_store = ExtensionData::new("session"); let thread_store = ExtensionData::new("thread"); - thread_store.insert( - [("root-a", "/skills/root-a"), ("root-b", "/skills/root-b")] - .into_iter() - .map(|(id, path)| SelectedCapabilityRoot { - id: id.to_string(), - location: CapabilityRootLocation::Environment { - environment_id: "env-1".to_string(), - path: PathUri::parse(&format!("file://{path}")).expect("skill root URI"), - }, - }) - .collect::>(), - ); + let selected_roots = [("root-a", "/skills/root-a"), ("root-b", "/skills/root-b")] + .into_iter() + .map(|(id, path)| SelectedCapabilityRoot { + id: id.to_string(), + location: CapabilityRootLocation::Environment { + environment_id: "env-1".to_string(), + path: PathUri::parse(&format!("file://{path}")).expect("skill root URI"), + }, + }) + .collect::>(); let session_source = SessionSource::Cli; let config = default_config(); registry.thread_lifecycle_contributors()[0] @@ -513,6 +586,21 @@ async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> Te }) .await; + let turn_store = ExtensionData::new("turn-1"); + registry.context_contributors()[0] + .contribute_world_state(WorldStateContributionInput { + thread_id: codex_protocol::ThreadId::new(), + turn_id: "turn-1", + environments: &[TurnEnvironmentSelection { + environment_id: "env-1".to_string(), + cwd: PathUri::parse("file:///workspace").expect("cwd URI"), + }], + ready_selected_capability_roots: &selected_roots, + session_store: &session_store, + thread_store: &thread_store, + turn_store: &turn_store, + }) + .await; let fragments = registry.turn_input_contributors()[0] .contribute( TurnInputContext { @@ -525,7 +613,7 @@ async fn root_qualified_locator_selects_only_the_matching_executor_skill() -> Te }, &session_store, &thread_store, - &ExtensionData::new("turn-1"), + &turn_store, ) .await; From 8ce931ab763541736bd08d0507b4777543532762 Mon Sep 17 00:00:00 2001 From: Tom Date: Thu, 25 Jun 2026 16:22:22 -0700 Subject: [PATCH 02/38] [codex] Propagate traces through exec-server HTTP (#30117) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes distributed trace continuity across exec-server JSON-RPC HTTP egress by adding an executor client span and injecting its W3C context through a reusable `codex-otel` helper. This preserves the caller trace across core/tool → executor → provider/MCP instead of dropping parentage at raw reqwest. Note that this doesn't include the websocket path, which is needed to really get the full story but at least we cover the basic http path with this change. --- .../src/client/reqwest_http_client.rs | 17 ++++++++-- codex-rs/otel/src/lib.rs | 1 + codex-rs/otel/src/trace_context.rs | 32 +++++++++++++++++++ 3 files changed, 48 insertions(+), 2 deletions(-) diff --git a/codex-rs/exec-server/src/client/reqwest_http_client.rs b/codex-rs/exec-server/src/client/reqwest_http_client.rs index 2fcb7c7e6c..667e4c49d8 100644 --- a/codex-rs/exec-server/src/client/reqwest_http_client.rs +++ b/codex-rs/exec-server/src/client/reqwest_http_client.rs @@ -18,6 +18,7 @@ use reqwest::Url; use reqwest::header::HeaderMap; use reqwest::header::HeaderName; use reqwest::header::HeaderValue; +use tracing::Instrument; use super::HttpResponseBodyStream; use super::response_body_stream::send_body_delta; @@ -146,15 +147,26 @@ impl ReqwestHttpRequestRunner { } } - let headers = Self::build_headers(params.headers)?; + let request_span = tracing::info_span!( + "codex.exec_server.http_request", + otel.kind = "client", + http.request.method = method.as_str(), + server.address = url.host_str().unwrap_or_default(), + server.port = u64::from(url.port_or_known_default().unwrap_or_default()), + http.response.status_code = tracing::field::Empty, + error.type = tracing::field::Empty, + ); + let mut headers = Self::build_headers(params.headers)?; + codex_otel::inject_span_w3c_trace_headers(&request_span, &mut headers); let mut request = self.client.request(method.clone(), url).headers(headers); if let Some(body) = params.body { request = request.body(body.into_inner()); } - let response = match request.send().await { + let response = match request.send().instrument(request_span.clone()).await { Ok(response) => response, Err(error) => { + request_span.record("error.type", "request"); let error_message = error.to_string(); log_send_error(&method, error); return Err(internal_error(format!( @@ -163,6 +175,7 @@ impl ReqwestHttpRequestRunner { } }; let status = response.status().as_u16(); + request_span.record("http.response.status_code", u64::from(status)); let headers = Self::response_headers(response.headers()); if params.stream_response { diff --git a/codex-rs/otel/src/lib.rs b/codex-rs/otel/src/lib.rs index 2a36fab16d..87ccbbfeee 100644 --- a/codex-rs/otel/src/lib.rs +++ b/codex-rs/otel/src/lib.rs @@ -29,6 +29,7 @@ pub use crate::provider::OtelProvider; pub use crate::trace_context::context_from_w3c_trace_context; pub use crate::trace_context::current_span_trace_id; pub use crate::trace_context::current_span_w3c_trace_context; +pub use crate::trace_context::inject_span_w3c_trace_headers; pub use crate::trace_context::set_parent_from_context; pub use crate::trace_context::set_parent_from_w3c_trace_context; pub use crate::trace_context::span_w3c_trace_context; diff --git a/codex-rs/otel/src/trace_context.rs b/codex-rs/otel/src/trace_context.rs index c625a416a7..6ae924d1d8 100644 --- a/codex-rs/otel/src/trace_context.rs +++ b/codex-rs/otel/src/trace_context.rs @@ -49,6 +49,38 @@ pub fn span_w3c_trace_context(span: &Span) -> Option { }) } +/// Injects the W3C trace context for `span` into HTTP headers. +/// +/// Existing `traceparent` and `tracestate` values are replaced so callers can +/// safely reuse a request header map while keeping the supplied span as the +/// source of truth. +pub fn inject_span_w3c_trace_headers(span: &Span, headers: &mut http::HeaderMap) -> bool { + let Some(trace) = span_w3c_trace_context(span) else { + return false; + }; + match trace.traceparent { + Some(traceparent) => { + if let Ok(value) = http::HeaderValue::from_str(&traceparent) { + headers.insert("traceparent", value); + } + } + None => { + headers.remove("traceparent"); + } + } + match trace.tracestate { + Some(tracestate) => { + if let Ok(value) = http::HeaderValue::from_str(&tracestate) { + headers.insert("tracestate", value); + } + } + None => { + headers.remove("tracestate"); + } + } + true +} + pub(crate) fn set_tracestate_entries( entries: BTreeMap>, ) -> Result<(), Box> { From ee9e0f6387b91cd14008ea59a8e4315b2a2fdacb Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 26 Jun 2026 00:53:07 +0100 Subject: [PATCH 03/38] Pin MCP runtimes to model steps (#30101) ## Why An MCP refresh can replace the session's current manager while a model step is still running. The step must execute calls through the same manager whose tools it advertised. ## Boundary ```text current session MCP runtime | | capture once for this model step v StepContext.mcp - exact MCP config - exact connection manager - exact runtime environment context ``` ```rust pub struct McpRuntimeSnapshot { config: Arc, manager: Arc, runtime_context: McpRuntimeContext, } ``` ## Example ```text step A captures runtime A and advertises A's tools refresh publishes runtime B step A tool call -> runtime A next step -> runtime B ``` Capturing the snapshot is only an `Arc` clone. It does not restart MCPs or make an RPC. ## What changes - Captures one MCP runtime in `StepContext`. - Uses it for tool planning, tool calls, resources, approvals, connector attribution, and elicitation. - Publishes replacement runtimes atomically. - Lets an old runtime live only while an in-flight step or request still holds its `Arc`. Most of this diff is mechanical routing from the session-global manager to `step_context.mcp`; it does not introduce selected-plugin discovery yet. ## What does not change - No plugin or extension migration. - No new MCP cache policy. - No environment file watching. - No client sharing between separate managers. ## Stack 1. Extension-owned World State sections. 2. Project executor skills through World State. 3. **This PR:** pin one MCP runtime to each model step. 4. Project selected MCP/app/connector metadata by environment availability. 5. One end-to-end integration scenario. --- .../src/request_processors/mcp_processor.rs | 44 +++--- codex-rs/codex-mcp/src/catalog.rs | 24 ++++ codex-rs/codex-mcp/src/catalog_tests.rs | 11 ++ codex-rs/core/src/codex_delegate.rs | 62 ++++++--- codex-rs/core/src/codex_delegate_tests.rs | 22 +-- codex-rs/core/src/codex_thread.rs | 5 + codex-rs/core/src/mcp_skill_dependencies.rs | 11 +- codex-rs/core/src/mcp_tool_call.rs | 127 +++++++++--------- codex-rs/core/src/mcp_tool_call_tests.rs | 54 +++++--- codex-rs/core/src/session/handlers.rs | 4 +- codex-rs/core/src/session/mcp.rs | 124 +++++++---------- codex-rs/core/src/session/mcp_runtime.rs | 94 +++++++++++++ codex-rs/core/src/session/mod.rs | 42 ++++-- codex-rs/core/src/session/session.rs | 56 ++++---- codex-rs/core/src/session/step_context.rs | 5 + codex-rs/core/src/session/tests.rs | 121 ++++++++++++++--- codex-rs/core/src/session/turn.rs | 2 +- codex-rs/core/src/session/turn_context.rs | 3 +- codex-rs/core/src/state/service.rs | 38 +++++- codex-rs/core/src/tools/handlers/mcp.rs | 5 +- .../list_mcp_resource_templates.rs | 11 +- .../mcp_resource/list_mcp_resources.rs | 11 +- .../mcp_resource/read_mcp_resource.rs | 6 +- .../tools/handlers/request_plugin_install.rs | 18 ++- codex-rs/core/src/tools/router_tests.rs | 4 +- codex-rs/core/src/tools/spec_plan_tests.rs | 4 +- .../core/tests/suite/mcp_refresh_cleanup.rs | 13 +- 27 files changed, 596 insertions(+), 325 deletions(-) create mode 100644 codex-rs/core/src/session/mcp_runtime.rs diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index a3d2191782..abefd26015 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -121,17 +121,11 @@ impl McpRequestProcessor { } = params; let auth = self.auth_manager.auth().await; - let (config, mcp_config) = match thread_id.as_deref() { + let (mcp_config, runtime_context) = match thread_id.as_deref() { Some(thread_id) => { let (_, thread) = self.load_thread(thread_id).await?; - let thread_config = thread.config().await; - let config = self - .config_manager - .load_latest_config_for_thread(thread_config.as_ref()) - .await - .map_err(|err| internal_error(format!("failed to reload config: {err}")))?; - let mcp_config = thread.runtime_mcp_config(&config).await; - (config, mcp_config) + let runtime = thread.current_mcp_runtime(); + (runtime.config().clone(), runtime.runtime_context().clone()) } None => { let config = self.load_latest_config(/*fallback_cwd*/ None).await?; @@ -140,7 +134,11 @@ impl McpRequestProcessor { .mcp_manager() .runtime_config(&config) .await; - (config, mcp_config) + let runtime_context = McpRuntimeContext::new( + self.thread_manager.environment_manager(), + config.cwd.to_path_buf(), + ); + (mcp_config, runtime_context) } }; let effective_servers = codex_mcp::effective_mcp_servers(&mcp_config, auth.as_ref()); @@ -167,10 +165,6 @@ impl McpRequestProcessor { } }; - let runtime_context = McpRuntimeContext::new( - self.thread_manager.environment_manager(), - config.cwd.to_path_buf(), - ); let http_client = runtime_context .resolve_http_client(&name, server) .map_err(|err| { @@ -189,16 +183,16 @@ impl McpRequestProcessor { let handle = perform_oauth_login_return_url_with_http_client( &name, &url, - config.mcp_oauth_credentials_store_mode, - config.auth_keyring_backend_kind(), + mcp_config.mcp_oauth_credentials_store_mode, + mcp_config.auth_keyring_backend_kind, http_headers, env_http_headers, &resolved_scopes.scopes, server.oauth_client_id(), server.oauth_resource.as_deref(), timeout_secs, - config.mcp_oauth_callback_port, - config.mcp_oauth_callback_url.as_deref(), + mcp_config.mcp_oauth_callback_port, + mcp_config.mcp_oauth_callback_url.as_deref(), http_client, ) .await @@ -250,18 +244,16 @@ impl McpRequestProcessor { None => (self.load_latest_config(/*fallback_cwd*/ None).await?, None), }; let mcp_manager = self.thread_manager.mcp_manager(); + let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache(); + let auth = self.auth_manager.auth().await; let mcp_config = match thread { Some(thread) => thread.runtime_mcp_config(&config).await, None => mcp_manager.runtime_config(&config).await, }; - let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache(); - let auth = self.auth_manager.auth().await; - let environment_manager = self.thread_manager.environment_manager(); - // This status path has no turn-selected environment. Use config cwd - // as the local stdio fallback; named environment stdio MCPs must - // declare their own absolute cwd. - let runtime_context = - McpRuntimeContext::new(Arc::clone(&environment_manager), config.cwd.to_path_buf()); + let runtime_context = McpRuntimeContext::new( + self.thread_manager.environment_manager(), + config.cwd.to_path_buf(), + ); tokio::spawn(async move { Self::list_mcp_server_status_task( diff --git a/codex-rs/codex-mcp/src/catalog.rs b/codex-rs/codex-mcp/src/catalog.rs index 3a752c82ba..689f93d813 100644 --- a/codex-rs/codex-mcp/src/catalog.rs +++ b/codex-rs/codex-mcp/src/catalog.rs @@ -371,6 +371,30 @@ impl ResolvedMcpCatalog { .collect() } + /// Replaces the resolved server set while preserving known server sources. + /// + /// Names not present in the existing catalog are treated as config-owned. + pub fn with_materialized_servers(&self, servers: HashMap) -> Self { + let mut builder = Self::builder(); + for (name, config) in servers { + let source = self + .server(&name) + .map(|server| server.source.clone()) + .unwrap_or(McpServerSource::Config); + let precedence = match &source { + McpServerSource::Plugin(_) => RegistrationPrecedence::Plugin(Reverse(0)), + McpServerSource::SelectedPlugin(_) => { + RegistrationPrecedence::SelectedPlugin(Reverse(0)) + } + McpServerSource::Config => RegistrationPrecedence::Config, + McpServerSource::Compatibility { .. } => RegistrationPrecedence::Compatibility, + McpServerSource::Extension { .. } => RegistrationPrecedence::Extension(0), + }; + builder.register(McpServerRegistration::new(name, source, config, precedence)); + } + builder.build() + } + /// Returns package attribution for each winning plugin-owned server. pub fn plugin_attributions_by_server_name(&self) -> HashMap { self.servers diff --git a/codex-rs/codex-mcp/src/catalog_tests.rs b/codex-rs/codex-mcp/src/catalog_tests.rs index 0b2ffef50d..6a0515db39 100644 --- a/codex-rs/codex-mcp/src/catalog_tests.rs +++ b/codex-rs/codex-mcp/src/catalog_tests.rs @@ -304,6 +304,17 @@ fn selected_plugins_override_discovered_plugins_but_not_config() { }] ); + let refreshed = server("https://refreshed.example/mcp"); + let catalog = + catalog.with_materialized_servers(HashMap::from([("docs".to_string(), refreshed.clone())])); + assert_eq!( + catalog.server("docs"), + Some(&super::ResolvedMcpServer { + source: selected_plugin_source("selected-alpha"), + config: refreshed, + }) + ); + let mut builder = catalog.to_builder(); let configured = server("https://config.example/mcp"); builder.register(McpServerRegistration::from_config( diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 21e8de83ca..8874aeae79 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -40,6 +40,7 @@ use crate::guardian::spawn_approval_request_review; use crate::mcp_tool_call::MCP_TOOL_APPROVAL_ACCEPT; use crate::mcp_tool_call::MCP_TOOL_APPROVAL_ACCEPT_FOR_SESSION; use crate::mcp_tool_call::MCP_TOOL_APPROVAL_DECLINE_SYNTHETIC; +use crate::mcp_tool_call::McpToolApprovalMetadata; use crate::mcp_tool_call::build_guardian_mcp_tool_review_request; use crate::mcp_tool_call::is_mcp_tool_approval_question_id; use crate::mcp_tool_call::lookup_mcp_tool_metadata; @@ -60,6 +61,12 @@ use codex_protocol::protocol::MultiAgentVersion; #[cfg(test)] use crate::session::completed_session_loop_termination; +#[derive(Clone)] +struct PendingMcpInvocation { + invocation: McpInvocation, + metadata: Option, +} + /// Start an interactive sub-Codex thread and return IO channels. /// /// The returned `events_rx` yields non-approval events emitted by the sub-agent. @@ -149,10 +156,10 @@ pub(crate) async fn run_codex_thread_interactive( let parent_session_clone = Arc::clone(&parent_session); let parent_ctx_clone = Arc::clone(&parent_ctx); let codex_for_events = Arc::clone(&codex); - // Cache delegated MCP invocations so guardian can recover the full tool call - // context when the later legacy RequestUserInput approval event only carries - // a call_id plus approval question metadata. - let pending_mcp_invocations = Arc::new(Mutex::new(HashMap::::new())); + // Cache the child call's MCP metadata at begin time. The later legacy + // RequestUserInput approval event only carries a call_id and question metadata. + let pending_mcp_invocations = + Arc::new(Mutex::new(HashMap::::new())); tokio::spawn(async move { forward_events( codex_for_events, @@ -270,7 +277,7 @@ async fn forward_events( tx_sub: Sender, parent_session: Arc, parent_ctx: Arc, - pending_mcp_invocations: Arc>>, + pending_mcp_invocations: Arc>>, cancel_token: CancellationToken, ) { let cancelled = cancel_token.cancelled(); @@ -357,10 +364,34 @@ async fn forward_events( id, msg: EventMsg::McpToolCallBegin(event), } => { + // Runtime refreshes are published before a request step is captured, so + // the child runtime at call begin is the one executing this invocation. + // Cache its metadata now; the later approval event has only a call ID. + let metadata = if let Some(turn_context) = + codex.session.turn_context_for_sub_id(&id).await + { + let mcp = codex.session.services.latest_mcp_runtime(); + lookup_mcp_tool_metadata( + codex.session.as_ref(), + turn_context.as_ref(), + mcp.manager(), + &event.invocation.server, + &event.invocation.tool, + ) + .await + } else { + None + }; pending_mcp_invocations .lock() .await - .insert(event.call_id.clone(), event.invocation.clone()); + .insert( + event.call_id.clone(), + PendingMcpInvocation { + invocation: event.invocation.clone(), + metadata, + }, + ); if !forward_event_or_shutdown( &codex, &tx_sub, @@ -648,7 +679,7 @@ async fn handle_request_user_input( id: String, parent_session: &Arc, parent_ctx: &Arc, - pending_mcp_invocations: &Arc>>, + pending_mcp_invocations: &Arc>>, event: RequestUserInputEvent, cancel_token: &CancellationToken, ) { @@ -686,12 +717,12 @@ async fn handle_request_user_input( /// programmatically after running the guardian review. /// /// The RequestUserInput event only carries `call_id` plus approval question -/// metadata, so this helper joins it back to the cached `McpToolCallBegin` -/// invocation in order to rebuild the full guardian review request. +/// metadata, so this helper joins it back to the child runtime metadata cached at +/// `McpToolCallBegin` in order to rebuild the full guardian review request. async fn maybe_auto_review_mcp_request_user_input( parent_session: &Arc, parent_ctx: &Arc, - pending_mcp_invocations: &Arc>>, + pending_mcp_invocations: &Arc>>, event: &RequestUserInputEvent, cancel_token: &CancellationToken, ) -> Option { @@ -702,18 +733,13 @@ async fn maybe_auto_review_mcp_request_user_input( .questions .iter() .find(|question| is_mcp_tool_approval_question_id(&question.id))?; - let invocation = pending_mcp_invocations + let pending = pending_mcp_invocations .lock() .await .get(&event.call_id) .cloned()?; - let metadata = lookup_mcp_tool_metadata( - parent_session.as_ref(), - parent_ctx.as_ref(), - &invocation.server, - &invocation.tool, - ) - .await; + let invocation = pending.invocation; + let metadata = pending.metadata; let approvals_reviewer = mcp_approvals_reviewer(parent_ctx, &invocation.server, metadata.as_ref()); if !routes_approval_to_guardian_with_reviewer(parent_ctx, approvals_reviewer) { diff --git a/codex-rs/core/src/codex_delegate_tests.rs b/codex-rs/core/src/codex_delegate_tests.rs index 6d51a0032b..2791c887e6 100644 --- a/codex-rs/core/src/codex_delegate_tests.rs +++ b/codex-rs/core/src/codex_delegate_tests.rs @@ -411,10 +411,13 @@ async fn delegated_mcp_guardian_abort_returns_synthetic_decline_answer() { let pending_mcp_invocations = Arc::new(Mutex::new(HashMap::from([( "call-1".to_string(), - McpInvocation { - server: "custom_server".to_string(), - tool: "dangerous_tool".to_string(), - arguments: None, + PendingMcpInvocation { + invocation: McpInvocation { + server: "custom_server".to_string(), + tool: "dangerous_tool".to_string(), + arguments: None, + }, + metadata: None, }, )]))); let cancel_token = CancellationToken::new(); @@ -460,10 +463,13 @@ async fn delegated_mcp_user_reviewer_returns_none_without_metadata() { crate::session::tests::make_session_and_context_with_rx().await; let pending_mcp_invocations = Arc::new(Mutex::new(HashMap::from([( "call-1".to_string(), - McpInvocation { - server: CODEX_APPS_MCP_SERVER_NAME.to_string(), - tool: "dangerous_tool".to_string(), - arguments: None, + PendingMcpInvocation { + invocation: McpInvocation { + server: CODEX_APPS_MCP_SERVER_NAME.to_string(), + tool: "dangerous_tool".to_string(), + arguments: None, + }, + metadata: None, }, )]))); let cancel_token = CancellationToken::new(); diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 25f47ce45d..46b38af151 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -594,6 +594,11 @@ impl CodexThread { self.codex.session.runtime_mcp_config(config).await } + /// Returns the exact MCP config, environment bindings, and manager most recently published. + pub fn current_mcp_runtime(&self) -> Arc { + self.codex.session.services.latest_mcp_runtime() + } + pub fn multi_agent_version(&self) -> Option { self.codex.session.multi_agent_version() } diff --git a/codex-rs/core/src/mcp_skill_dependencies.rs b/codex-rs/core/src/mcp_skill_dependencies.rs index 9b2b85aa2d..61d5c54a59 100644 --- a/codex-rs/core/src/mcp_skill_dependencies.rs +++ b/codex-rs/core/src/mcp_skill_dependencies.rs @@ -199,15 +199,8 @@ pub(crate) async fn maybe_install_mcp_dependencies( warn!("failed to refresh MCP dependencies for mentioned skills: {err}"); return; } - let refresh_servers = sess.runtime_mcp_servers(&refresh_config).await; - sess.refresh_mcp_servers_now( - turn_context, - refresh_servers, - config.mcp_oauth_credentials_store_mode, - config.auth_keyring_backend_kind(), - elicitation_reviewer, - ) - .await; + sess.refresh_mcp_servers_now(turn_context, &refresh_config, elicitation_reviewer) + .await; } async fn should_install_mcp_dependencies( diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index a270af7851..19b7c1a6ad 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -18,6 +18,7 @@ use crate::mcp_openai_file::rewrite_mcp_tool_arguments_for_openai_files; use crate::mcp_tool_approval_templates::RenderedMcpToolApprovalParam; use crate::mcp_tool_approval_templates::render_mcp_tool_approval_template; use crate::session::session::Session; +use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; use crate::tools::hook_names::HookToolName; use crate::tools::sandboxing::PermissionRequestPayload; @@ -35,6 +36,7 @@ use codex_features::Feature; use codex_hooks::PermissionRequestDecision; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::MCP_TOOL_CODEX_APPS_META_KEY; +use codex_mcp::McpConnectionManager; use codex_mcp::McpPermissionPromptAutoApproveContext; use codex_mcp::SandboxState; use codex_mcp::auth_elicitation_completed_result; @@ -111,13 +113,15 @@ const MCP_TOOL_CALL_EVENT_RESULT_MAX_BYTES: usize = DEFAULT_OUTPUT_BYTES_CAP; /// item lifecycle events to the `Session`. pub(crate) async fn handle_mcp_tool_call( sess: Arc, - turn_context: &Arc, + step_context: &Arc, call_id: String, server: String, tool_name: String, hook_tool_name: HookToolName, arguments: String, ) -> HandledMcpToolCall { + let turn_context = &step_context.turn; + let manager = step_context.mcp.manager(); // Parse the `arguments` as JSON. An empty string is OK, but invalid JSON // is not. let arguments_value = if arguments.trim().is_empty() { @@ -141,8 +145,14 @@ pub(crate) async fn handle_mcp_tool_call( arguments: arguments_value.clone(), }; - let metadata = - lookup_mcp_tool_metadata(sess.as_ref(), turn_context.as_ref(), &server, &tool_name).await; + let metadata = lookup_mcp_tool_metadata( + sess.as_ref(), + turn_context.as_ref(), + manager, + &server, + &tool_name, + ) + .await; let item_metadata = McpToolCallItemMetadata::from_tool_metadata(&server, metadata.as_ref()); let app_tool_policy = if server == CODEX_APPS_MCP_SERVER_NAME { let annotations = metadata @@ -169,7 +179,6 @@ pub(crate) async fn handle_mcp_tool_call( } else if let Some(approval_mode) = { // Selected-plugin registrations are absent from config.toml and the legacy plugin manager, // so their resolved catalog entry is the authoritative source for tool approval policy. - let manager = sess.services.mcp_connection_manager.load(); manager .is_selected_plugin_mcp_server(&server) .then(|| manager.tool_approval_mode(&server, &tool_name)) @@ -226,7 +235,7 @@ pub(crate) async fn handle_mcp_tool_call( if let Some(decision) = maybe_request_mcp_tool_approval( &sess, - turn_context, + step_context, &call_id, &invocation, &hook_tool_name, @@ -241,7 +250,7 @@ pub(crate) async fn handle_mcp_tool_call( | McpToolApprovalDecision::AcceptAndRemember => { return handle_approved_mcp_tool_call( sess.as_ref(), - turn_context.as_ref(), + step_context.as_ref(), &call_id, invocation, metadata.as_ref(), @@ -298,7 +307,7 @@ pub(crate) async fn handle_mcp_tool_call( handle_approved_mcp_tool_call( sess.as_ref(), - turn_context.as_ref(), + step_context.as_ref(), &call_id, invocation, metadata.as_ref(), @@ -340,24 +349,21 @@ impl McpToolCallItemMetadata { async fn handle_approved_mcp_tool_call( sess: &Session, - turn_context: &TurnContext, + step_context: &StepContext, call_id: &str, invocation: McpInvocation, metadata: Option<&McpToolApprovalMetadata>, item_metadata: McpToolCallItemMetadata, ) -> HandledMcpToolCall { + let turn_context = step_context.turn.as_ref(); + let manager = step_context.mcp.manager(); let server = invocation.server.clone(); - maybe_mark_thread_memory_mode_polluted(sess, turn_context, &server).await; + maybe_mark_thread_memory_mode_polluted(sess, turn_context, manager, &server).await; let tool_name = invocation.tool.clone(); let arguments_value = invocation.arguments.clone(); let connector_id = metadata.and_then(|metadata| metadata.connector_id.as_deref()); let connector_name = metadata.and_then(|metadata| metadata.connector_name.as_deref()); - let server_origin = { - let mcp_connection_manager = sess.services.mcp_connection_manager.load_full(); - mcp_connection_manager - .server_origin(&server) - .map(str::to_string) - }; + let server_origin = manager.server_origin(&server).map(str::to_string); let start = Instant::now(); let rewrite = rewrite_mcp_tool_arguments_for_openai_files( @@ -380,7 +386,7 @@ async fn handle_approved_mcp_tool_call( build_mcp_tool_call_request_meta(turn_context, &server, call_id, metadata); execute_mcp_tool_call( sess, - turn_context, + step_context, call_id, &invocation, rewritten_arguments, @@ -420,7 +426,7 @@ async fn handle_approved_mcp_tool_call( truncate_mcp_tool_result_for_event(&result), ) .await; - maybe_track_codex_app_used(sess, turn_context, &server, &tool_name).await; + maybe_track_codex_app_used(sess, turn_context, manager, &server, &tool_name).await; let outcome = mcp_call_metric_outcome(&result); emit_mcp_call_metrics( @@ -547,17 +553,19 @@ fn truncate_str_to_char_boundary(value: &str, max_chars: usize) -> &str { async fn execute_mcp_tool_call( sess: &Session, - turn_context: &TurnContext, + step_context: &StepContext, call_id: &str, invocation: &McpInvocation, rewritten_arguments: Option, metadata: Option<&McpToolApprovalMetadata>, request_meta: Option, ) -> Result { + let turn_context = step_context.turn.as_ref(); + let manager = step_context.mcp.manager(); let request_meta = with_mcp_tool_call_thread_id_meta(request_meta, &sess.thread_id.to_string()); let request_meta = augment_mcp_tool_request_meta_with_sandbox_state( - sess, - turn_context, + step_context, + manager, &invocation.server, request_meta, ) @@ -568,7 +576,7 @@ async fn execute_mcp_tool_call( .rollout_thread_trace .start_mcp_call_trace(call_id); let request_meta = mcp_call_trace.add_request_meta(request_meta); - let result = sess + let result = manager .call_tool( &invocation.server, &invocation.tool, @@ -587,6 +595,7 @@ async fn execute_mcp_tool_call( Ok(maybe_request_codex_apps_auth_elicitation( sess, turn_context, + manager, call_id, &invocation.server, metadata, @@ -598,17 +607,13 @@ async fn execute_mcp_tool_call( async fn maybe_request_codex_apps_auth_elicitation( sess: &Session, turn_context: &TurnContext, + manager: &McpConnectionManager, call_id: &str, server: &str, metadata: Option<&McpToolApprovalMetadata>, result: CallToolResult, ) -> CallToolResult { - if !sess - .services - .mcp_connection_manager - .load_full() - .is_host_owned_codex_apps_server(server) - { + if !manager.is_host_owned_codex_apps_server(server) { return result; } @@ -666,15 +671,16 @@ async fn maybe_request_codex_apps_auth_elicitation( return result; } - refresh_codex_apps_after_connector_auth(sess, turn_context).await; + refresh_codex_apps_after_connector_auth(sess, turn_context, manager).await; auth_elicitation_completed_result(&plan.auth_failure, result.meta) } -async fn refresh_codex_apps_after_connector_auth(sess: &Session, turn_context: &TurnContext) { - let mcp_tools_result = { - let manager = sess.services.mcp_connection_manager.load_full(); - manager.hard_refresh_codex_apps_tools_cache().await - }; +async fn refresh_codex_apps_after_connector_auth( + sess: &Session, + turn_context: &TurnContext, + manager: &McpConnectionManager, +) { + let mcp_tools_result = manager.hard_refresh_codex_apps_tools_cache().await; match mcp_tools_result { Ok(mcp_tools) => { @@ -692,13 +698,13 @@ async fn refresh_codex_apps_after_connector_auth(sess: &Session, turn_context: & } async fn augment_mcp_tool_request_meta_with_sandbox_state( - sess: &Session, - turn_context: &TurnContext, + step_context: &StepContext, + manager: &McpConnectionManager, server: &str, mut meta: Option, ) -> anyhow::Result> { - let mcp_connection_manager = sess.services.mcp_connection_manager.load_full(); - let supports_sandbox_state_meta = mcp_connection_manager + let turn_context = step_context.turn.as_ref(); + let supports_sandbox_state_meta = manager .server_supports_sandbox_state_meta_capability(server) .await .unwrap_or(false); @@ -706,18 +712,18 @@ async fn augment_mcp_tool_request_meta_with_sandbox_state( return Ok(meta); } - let server_environment_id = mcp_connection_manager + let server_environment_id = manager .server_environment_id(server) .unwrap_or(codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID); - let Some(sandbox_cwd) = sandbox_cwd_for_mcp_server(turn_context, server_environment_id) else { + let Some(sandbox_cwd) = sandbox_cwd_for_mcp_server(step_context, server_environment_id) else { return Ok(meta); }; let permission_profile = turn_context.permission_profile(); let sandbox_state = serde_json::to_value(SandboxState { permission_profile, - codex_linux_sandbox_exe: turn_context.config.codex_linux_sandbox_exe.clone(), + codex_linux_sandbox_exe: step_context.mcp.config().codex_linux_sandbox_exe.clone(), sandbox_cwd, - use_legacy_landlock: turn_context.config.features.use_legacy_landlock(), + use_legacy_landlock: step_context.mcp.config().use_legacy_landlock, })?; match meta.as_mut() { @@ -741,8 +747,8 @@ async fn augment_mcp_tool_request_meta_with_sandbox_state( Ok(meta) } -fn sandbox_cwd_for_mcp_server(turn_context: &TurnContext, environment_id: &str) -> Option { - if let Some(environment) = turn_context +fn sandbox_cwd_for_mcp_server(step_context: &StepContext, environment_id: &str) -> Option { + if let Some(environment) = step_context .environments .turn_environments .iter() @@ -753,7 +759,7 @@ fn sandbox_cwd_for_mcp_server(turn_context: &TurnContext, environment_id: &str) if environment_id == codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID { #[allow(deprecated)] - return Some(PathUri::from_abs_path(&turn_context.cwd)); + return Some(PathUri::from_abs_path(&step_context.turn.cwd)); } None @@ -762,16 +768,13 @@ fn sandbox_cwd_for_mcp_server(turn_context: &TurnContext, environment_id: &str) async fn maybe_mark_thread_memory_mode_polluted( sess: &Session, turn_context: &TurnContext, + manager: &McpConnectionManager, server: &str, ) { if !turn_context.config.memories.disable_on_external_context { return; } - let pollutes_memory = sess - .services - .mcp_connection_manager - .load_full() - .server_pollutes_memory(server); + let pollutes_memory = manager.server_pollutes_memory(server); if !pollutes_memory { return; } @@ -937,13 +940,14 @@ struct McpAppUsageMetadata { async fn maybe_track_codex_app_used( sess: &Session, turn_context: &TurnContext, + manager: &McpConnectionManager, server: &str, tool_name: &str, ) { if server != CODEX_APPS_MCP_SERVER_NAME { return; } - let metadata = lookup_mcp_app_usage_metadata(sess, server, tool_name).await; + let metadata = lookup_mcp_app_usage_metadata(manager, server, tool_name).await; let (connector_id, app_name) = metadata .map(|metadata| (metadata.connector_id, metadata.app_name)) .unwrap_or((None, None)); @@ -983,6 +987,7 @@ enum McpToolApprovalDecision { Cancel, } +#[derive(Clone)] pub(crate) struct McpToolApprovalMetadata { annotations: Option, connector_id: Option, @@ -1176,13 +1181,15 @@ fn mcp_tool_approval_prompt_options( async fn maybe_request_mcp_tool_approval( sess: &Arc, - turn_context: &Arc, + step_context: &Arc, call_id: &str, invocation: &McpInvocation, hook_tool_name: &HookToolName, metadata: Option<&McpToolApprovalMetadata>, approval_mode: AppToolApproval, ) -> Option { + let turn_context = &step_context.turn; + let manager = step_context.mcp.manager(); let approvals_reviewer = mcp_approvals_reviewer(turn_context, &invocation.server, metadata); if mcp_permission_prompt_is_auto_approved( turn_context.approval_policy.value(), @@ -1201,12 +1208,7 @@ async fn maybe_request_mcp_tool_approval( } let session_approval_key = session_mcp_tool_approval_key(invocation, metadata, approval_mode); - let persistent_approval_key = if sess - .services - .mcp_connection_manager - .load() - .is_selected_plugin_mcp_server(&invocation.server) - { + let persistent_approval_key = if manager.is_selected_plugin_mcp_server(&invocation.server) { None } else { persistent_mcp_tool_approval_key(invocation, metadata, approval_mode) @@ -1452,10 +1454,10 @@ async fn mcp_tool_approval_decision_from_guardian( pub(crate) async fn lookup_mcp_tool_metadata( sess: &Session, turn_context: &TurnContext, + manager: &McpConnectionManager, server: &str, tool_name: &str, ) -> Option { - let manager = sess.services.mcp_connection_manager.load_full(); let plugin_id = manager .plugin_id_for_mcp_server_name(server) .map(str::to_string); @@ -1568,16 +1570,11 @@ fn get_mcp_app_resource_uri( } async fn lookup_mcp_app_usage_metadata( - sess: &Session, + manager: &McpConnectionManager, server: &str, tool_name: &str, ) -> Option { - let tools = sess - .services - .mcp_connection_manager - .load_full() - .list_all_tools() - .await; + let tools = manager.list_all_tools().await; tools.into_iter().find_map(|tool_info| { if tool_info.server_name == server && tool_info.tool.name == tool_name { diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 41701d557c..8c8d82037b 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1,6 +1,7 @@ use super::*; use crate::config::ConfigBuilder; use crate::config::ManagedFeatures; +use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::session::tests::make_session_and_context_with_rx; use crate::session::turn_context::TurnEnvironment; @@ -157,10 +158,12 @@ async fn execute_mcp_tool_call_records_replayable_correlation() -> anyhow::Resul }) }); assert!(dispatch_trace.is_enabled()); + let turn_context = Arc::new(turn_context); + let step_context = StepContext::for_test(Arc::clone(&turn_context)); let result = execute_mcp_tool_call( &session, - &turn_context, + step_context.as_ref(), "mcp-call", &McpInvocation { server: "docs".to_string(), @@ -1147,7 +1150,8 @@ async fn mcp_sandbox_cwd_uses_matching_server_environment_uri() -> anyhow::Resul /*shell*/ None, )); - let sandbox_cwd = sandbox_cwd_for_mcp_server(&turn_context, "remote"); + let step_context = StepContext::for_test(Arc::new(turn_context)); + let sandbox_cwd = sandbox_cwd_for_mcp_server(&step_context, "remote"); assert_eq!(sandbox_cwd, Some(secondary_cwd)); Ok(()) @@ -1157,7 +1161,8 @@ async fn mcp_sandbox_cwd_uses_matching_server_environment_uri() -> anyhow::Resul async fn mcp_sandbox_cwd_is_none_for_unselected_server_environment() -> anyhow::Result<()> { let (_, turn_context) = make_session_and_context().await; - let sandbox_cwd = sandbox_cwd_for_mcp_server(&turn_context, "remote"); + let step_context = StepContext::for_test(Arc::new(turn_context)); + let sandbox_cwd = sandbox_cwd_for_mcp_server(&step_context, "remote"); assert_eq!(sandbox_cwd, None); Ok(()) @@ -1370,7 +1375,10 @@ fn codex_apps_auth_failure_metadata() -> McpToolApprovalMetadata { ) } -async fn install_host_owned_codex_apps_manager(session: &Session, turn_context: &TurnContext) { +async fn host_owned_codex_apps_manager( + session: &Session, + turn_context: &TurnContext, +) -> Arc { let auth = session.services.auth_manager.auth().await; let startup_cancellation_token = CancellationToken::new(); startup_cancellation_token.cancel(); @@ -1410,22 +1418,20 @@ async fn install_host_owned_codex_apps_manager(session: &Session, turn_context: /*elicitation_reviewer*/ None, ) .await; - session - .services - .mcp_connection_manager - .store(Arc::new(manager)); + Arc::new(manager) } #[tokio::test] async fn codex_apps_auth_elicitation_feature_disabled_returns_original_result() { let (session, turn_context, rx_event) = make_session_and_context_with_rx().await; - install_host_owned_codex_apps_manager(&session, &turn_context).await; + let manager = host_owned_codex_apps_manager(&session, &turn_context).await; let result = codex_apps_auth_failure_result(); let metadata = codex_apps_auth_failure_metadata(); let returned = maybe_request_codex_apps_auth_elicitation( &session, &turn_context, + manager.as_ref(), "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -1446,10 +1452,12 @@ async fn codex_apps_auth_elicitation_non_host_owned_server_returns_original_resu Arc::make_mut(&mut turn_context.config).features = ManagedFeatures::from(features); let result = codex_apps_auth_failure_result(); let metadata = codex_apps_auth_failure_metadata(); + let manager = session.services.latest_mcp_runtime().manager_arc(); let returned = maybe_request_codex_apps_auth_elicitation( &session, turn_context, + manager.as_ref(), "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -1464,7 +1472,7 @@ async fn codex_apps_auth_elicitation_non_host_owned_server_returns_original_resu #[tokio::test] async fn codex_apps_auth_elicitation_disallowed_by_policy_returns_original_result() { let (session, mut turn_context, rx_event) = make_session_and_context_with_rx().await; - install_host_owned_codex_apps_manager(&session, &turn_context).await; + let manager = host_owned_codex_apps_manager(&session, &turn_context).await; let mut features = Features::with_defaults(); features.enable(Feature::AuthElicitation); let turn_context = Arc::get_mut(&mut turn_context).expect("single turn context ref"); @@ -1479,6 +1487,7 @@ async fn codex_apps_auth_elicitation_disallowed_by_policy_returns_original_resul let returned = maybe_request_codex_apps_auth_elicitation( &session, turn_context, + manager.as_ref(), "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -1493,7 +1502,7 @@ async fn codex_apps_auth_elicitation_disallowed_by_policy_returns_original_resul #[tokio::test] async fn codex_apps_auth_elicitation_granular_mcp_disabled_returns_original_result() { let (session, mut turn_context, rx_event) = make_session_and_context_with_rx().await; - install_host_owned_codex_apps_manager(&session, &turn_context).await; + let manager = host_owned_codex_apps_manager(&session, &turn_context).await; let mut features = Features::with_defaults(); features.enable(Feature::AuthElicitation); let turn_context = Arc::get_mut(&mut turn_context).expect("single turn context ref"); @@ -1514,6 +1523,7 @@ async fn codex_apps_auth_elicitation_granular_mcp_disabled_returns_original_resu let returned = maybe_request_codex_apps_auth_elicitation( &session, turn_context, + manager.as_ref(), "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -1528,7 +1538,7 @@ async fn codex_apps_auth_elicitation_granular_mcp_disabled_returns_original_resu #[tokio::test] async fn codex_apps_auth_elicitation_feature_enabled_requests_elicitation() { let (session, mut turn_context, rx_event) = make_session_and_context_with_rx().await; - install_host_owned_codex_apps_manager(&session, &turn_context).await; + let manager = host_owned_codex_apps_manager(&session, &turn_context).await; *session.active_turn.lock().await = Some(ActiveTurn::default()); let mut features = Features::with_defaults(); features.enable(Feature::AuthElicitation); @@ -1542,10 +1552,12 @@ async fn codex_apps_auth_elicitation_feature_enabled_requests_elicitation() { let request_task = tokio::spawn({ let session = Arc::clone(&session); let turn_context = Arc::clone(&turn_context); + let manager = Arc::clone(&manager); async move { maybe_request_codex_apps_auth_elicitation( &session, &turn_context, + manager.as_ref(), "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -2466,7 +2478,7 @@ async fn approve_mode_skips_when_annotations_do_not_require_approval() { let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-1", &invocation, &HookToolName::new("mcp__test__tool"), @@ -2542,7 +2554,7 @@ async fn guardian_mode_skips_auto_when_annotations_do_not_require_approval() { let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-guardian", &invocation, &HookToolName::new("mcp__test__tool"), @@ -2601,7 +2613,7 @@ async fn permission_request_hook_allows_mcp_tool_call() { let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-mcp-hook", &invocation, &HookToolName::new("mcp__memory__create_entities"), @@ -2663,7 +2675,7 @@ async fn permission_request_hook_uses_hook_tool_name_without_metadata() { let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-mcp-hook-no-metadata", &invocation, &HookToolName::new("mcp__memory__create_entities"), @@ -2745,7 +2757,7 @@ async fn permission_request_hook_runs_after_remembered_mcp_approval() { let turn_context = Arc::new(turn_context); let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-mcp-remembered", &invocation, &HookToolName::new("mcp__memory__create_entities"), @@ -2828,7 +2840,7 @@ async fn guardian_mode_mcp_denial_returns_rationale_message() { let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-guardian-deny", &invocation, &HookToolName::new("mcp__test__tool"), @@ -2888,7 +2900,7 @@ async fn prompt_mode_waits_for_approval_when_annotations_do_not_require_approval tokio::spawn(async move { maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-prompt", &invocation, &HookToolName::new("mcp__test__tool"), @@ -2946,7 +2958,7 @@ async fn full_access_mode_skips_mcp_tool_approval_for_all_approval_modes() { ] { let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-2", &invocation, &HookToolName::new("mcp__test__tool"), @@ -3035,7 +3047,7 @@ async fn approve_mode_skips_guardian_in_every_permission_mode() { let turn_context = Arc::new(turn_context); let decision = maybe_request_mcp_tool_approval( &session, - &turn_context, + &StepContext::for_test(Arc::clone(&turn_context)), "call-3", &invocation, &HookToolName::new("mcp__test__tool"), diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 5c68893e7f..75fbca4c86 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -598,8 +598,8 @@ async fn shutdown_session_runtime(sess: &Arc) { warn!("failed to shutdown code mode session: {err}"); } sess.services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager_arc() .shutdown() .await; sess.guardian_review_session.shutdown().await; diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 249a6994b1..b9e87f0c89 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -105,8 +105,8 @@ impl Session { ) -> McpServerElicitationOutcome { if self .services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager() .elicitations_auto_deny() { return McpServerElicitationOutcome { @@ -201,44 +201,20 @@ impl Session { } self.services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager_arc() .resolve_elicitation(server_name, id, response) .await } - pub async fn list_resources( - &self, - server: &str, - params: Option, - ) -> anyhow::Result { - self.services - .mcp_connection_manager - .load_full() - .list_resources(server, params) - .await - } - - pub async fn list_resource_templates( - &self, - server: &str, - params: Option, - ) -> anyhow::Result { - self.services - .mcp_connection_manager - .load_full() - .list_resource_templates(server, params) - .await - } - pub async fn read_resource( &self, server: &str, params: ReadResourceRequestParams, ) -> anyhow::Result { self.services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager_arc() .read_resource(server, params) .await } @@ -251,8 +227,8 @@ impl Session { meta: Option, ) -> anyhow::Result { self.services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager_arc() .call_tool(server, tool, arguments, meta) .await } @@ -260,17 +236,17 @@ impl Session { async fn refresh_mcp_servers_inner( &self, turn_context: &TurnContext, - mcp_servers: HashMap, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, + mut mcp_config: McpConfig, + configured_mcp_servers: HashMap, elicitation_reviewer: Option, ) { + mcp_config.mcp_server_catalog = mcp_config + .mcp_server_catalog + .with_materialized_servers(configured_mcp_servers); + let mcp_config = Arc::new(mcp_config); let auth = self.services.auth_manager.auth().await; - let config = self.get_config().await; - let mcp_config = self.runtime_mcp_config(config.as_ref()).await; let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(&mcp_config); - let mcp_servers = - effective_mcp_servers_from_configured(mcp_servers, &mcp_config, auth.as_ref()); + let mcp_servers = effective_mcp_servers(&mcp_config, auth.as_ref()); let environment_manager = self.services.turn_environments.environment_manager(); // TODO(anp): Migrate MCP runtime cwd plumbing to PathUri so foreign environment cwd // values can be used without falling back to the legacy host cwd. @@ -286,35 +262,36 @@ impl Session { let mcp_runtime_context = McpRuntimeContext::new(environment_manager, cwd); let auth_statuses = compute_auth_statuses( mcp_servers.iter(), - store_mode, - keyring_backend_kind, + mcp_config.mcp_oauth_credentials_store_mode, + mcp_config.auth_keyring_backend_kind, auth.as_ref(), &mcp_runtime_context, ) .await; let mcp_startup_cancellation_token = { let mut guard = self.services.mcp_startup_cancellation_token.lock().await; - guard.cancel(); + // The previous runtime owns the old token and may still be serving an in-flight step. + // Its manager cancels that token when the last runtime handle is dropped. let cancellation_token = CancellationToken::new(); *guard = cancellation_token.clone(); cancellation_token }; let refreshed_manager = McpConnectionManager::new( &mcp_servers, - store_mode, - keyring_backend_kind, + mcp_config.mcp_oauth_credentials_store_mode, + mcp_config.auth_keyring_backend_kind, auth_statuses, &turn_context.approval_policy, turn_context.sub_id.clone(), self.get_tx_event(), mcp_startup_cancellation_token, turn_context.permission_profile(), - mcp_runtime_context, - config.codex_home.to_path_buf(), + mcp_runtime_context.clone(), + mcp_config.codex_home.clone(), self.services.mcp_manager.codex_apps_tools_cache(), codex_apps_tools_cache_key(auth.as_ref()), mcp_config.prefix_mcp_tool_names, - mcp_config.client_elicitation_capability, + mcp_config.client_elicitation_capability.clone(), self.services .supports_openai_form_elicitation .load(std::sync::atomic::Ordering::Relaxed), @@ -324,14 +301,12 @@ impl Session { ) .await; { - let current_manager = self.services.mcp_connection_manager.load_full(); - refreshed_manager.set_elicitations_auto_deny(current_manager.elicitations_auto_deny()); + let current_manager = self.services.latest_mcp_runtime(); + refreshed_manager + .set_elicitations_auto_deny(current_manager.manager().elicitations_auto_deny()); } - let superseded_manager = self - .services - .mcp_connection_manager - .swap(Arc::new(refreshed_manager)); - superseded_manager.shutdown().await; + self.services + .publish_mcp_runtime(mcp_config, mcp_runtime_context, refreshed_manager); } pub(crate) async fn refresh_mcp_servers_if_requested( @@ -376,14 +351,23 @@ impl Session { } }; - self.refresh_mcp_servers_inner( - turn_context, - mcp_servers, - store_mode, - keyring_backend_kind, - elicitation_reviewer, - ) - .await; + let mut refresh_config = self.get_config().await.as_ref().clone(); + refresh_config.mcp_oauth_credentials_store_mode = store_mode; + let secret_auth_storage_enabled = match keyring_backend_kind { + AuthKeyringBackendKind::Direct => false, + AuthKeyringBackendKind::Secrets => true, + }; + if let Err(err) = refresh_config + .features + .set_enabled(Feature::SecretAuthStorage, secret_auth_storage_enabled) + { + warn!("failed to apply MCP auth keyring backend refresh config: {err}"); + return; + } + + let mcp_config = self.runtime_mcp_config(&refresh_config).await; + self.refresh_mcp_servers_inner(turn_context, mcp_config, mcp_servers, elicitation_reviewer) + .await; } pub(crate) async fn set_openai_form_elicitation_support( @@ -417,19 +401,13 @@ impl Session { pub(crate) async fn refresh_mcp_servers_now( &self, turn_context: &TurnContext, - mcp_servers: HashMap, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, + refresh_config: &Config, elicitation_reviewer: Option, ) { - self.refresh_mcp_servers_inner( - turn_context, - mcp_servers, - store_mode, - keyring_backend_kind, - elicitation_reviewer, - ) - .await; + let mcp_config = self.runtime_mcp_config(refresh_config).await; + let mcp_servers = codex_mcp::configured_mcp_servers(&mcp_config); + self.refresh_mcp_servers_inner(turn_context, mcp_config, mcp_servers, elicitation_reviewer) + .await; } #[cfg(test)] diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs new file mode 100644 index 0000000000..dc5ea011ca --- /dev/null +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -0,0 +1,94 @@ +use std::fmt; +use std::sync::Arc; + +use codex_mcp::McpConfig; +use codex_mcp::McpConnectionManager; +use codex_mcp::McpRuntimeContext; + +/// MCP config, exact environment bindings, and manager used by one model request. +pub struct McpRuntimeSnapshot { + config: Arc, + manager: Arc, + runtime_context: McpRuntimeContext, +} + +impl McpRuntimeSnapshot { + pub(crate) fn new( + config: Arc, + manager: Arc, + runtime_context: McpRuntimeContext, + ) -> Self { + Self { + config, + manager, + runtime_context, + } + } + + pub fn config(&self) -> &McpConfig { + self.config.as_ref() + } + + pub fn manager(&self) -> &McpConnectionManager { + self.manager.as_ref() + } + + pub(crate) fn manager_arc(&self) -> Arc { + Arc::clone(&self.manager) + } + + pub fn runtime_context(&self) -> &McpRuntimeContext { + &self.runtime_context + } + + #[cfg(test)] + pub(crate) fn new_uninitialized_for_test(config: &crate::config::Config) -> Arc { + use codex_exec_server::EnvironmentManager; + use codex_features::Feature; + use codex_mcp::ResolvedMcpCatalog; + use rmcp::model::ElicitationCapability; + + let mcp_config = McpConfig { + chatgpt_base_url: config.chatgpt_base_url.clone(), + apps_mcp_product_sku: config.apps_mcp_product_sku.clone(), + codex_home: config.codex_home.to_path_buf(), + mcp_oauth_credentials_store_mode: config.mcp_oauth_credentials_store_mode, + auth_keyring_backend_kind: config.auth_keyring_backend_kind(), + mcp_oauth_callback_port: config.mcp_oauth_callback_port, + mcp_oauth_callback_url: config.mcp_oauth_callback_url.clone(), + skill_mcp_dependency_install_enabled: config + .features + .enabled(Feature::SkillMcpDependencyInstall), + approval_policy: config.permissions.approval_policy.clone(), + codex_linux_sandbox_exe: config.codex_linux_sandbox_exe.clone(), + use_legacy_landlock: config.features.use_legacy_landlock(), + apps_enabled: config.features.enabled(Feature::Apps), + prefix_mcp_tool_names: config.prefix_mcp_tool_names(), + client_elicitation_capability: ElicitationCapability::default(), + mcp_server_catalog: ResolvedMcpCatalog::default(), + connector_snapshot: codex_connectors::ConnectorSnapshot::default(), + }; + let manager = McpConnectionManager::new_uninitialized_with_permission_profile( + &config.permissions.approval_policy, + config.permissions.permission_profile(), + config.prefix_mcp_tool_names(), + ); + let runtime_context = McpRuntimeContext::new( + Arc::new(EnvironmentManager::default_for_tests()), + config.cwd.to_path_buf(), + ); + Arc::new(Self::new( + Arc::new(mcp_config), + Arc::new(manager), + runtime_context, + )) + } +} + +impl fmt::Debug for McpRuntimeSnapshot { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("McpRuntimeSnapshot") + .finish_non_exhaustive() + } +} diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 6fe61222cf..9421e07a66 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -158,15 +158,9 @@ use codex_utils_path_uri::PathUri; use futures::future::BoxFuture; use futures::future::Shared; use futures::prelude::*; -use rmcp::model::ElicitationCapability; -use rmcp::model::FormElicitationCapability; -use rmcp::model::ListResourceTemplatesResult; -use rmcp::model::ListResourcesResult; -use rmcp::model::PaginatedRequestParams; use rmcp::model::ReadResourceRequestParams; use rmcp::model::ReadResourceResult; use rmcp::model::RequestId; -use rmcp::model::UrlElicitationCapability; use serde_json::Value; use tokio::sync::Mutex; use tokio::sync::RwLock; @@ -214,6 +208,7 @@ mod handlers; mod inject; mod input_queue; mod mcp; +mod mcp_runtime; pub(crate) mod multi_agents; mod review; mod rollout_budget; @@ -235,6 +230,7 @@ use self::handlers::submission_loop; pub(crate) use self::input_queue::InputQueueActivity; pub(crate) use self::input_queue::TurnInput; pub(crate) use self::input_queue::TurnInputQueue; +pub use self::mcp_runtime::McpRuntimeSnapshot; use self::review::spawn_review_thread; use self::session::AppServerClientMetadata; use self::session::Session; @@ -340,7 +336,7 @@ use codex_core_plugins::RecommendedPluginCandidatesInput; use codex_git_utils::get_git_repo_root; use codex_mcp::McpConfig; use codex_mcp::compute_auth_statuses; -use codex_mcp::effective_mcp_servers_from_configured; +use codex_mcp::effective_mcp_servers; use codex_otel::SessionTelemetry; use codex_otel::THREAD_STARTED_METRIC; use codex_otel::TelemetryAuthMode; @@ -840,8 +836,11 @@ impl Codex { ..Default::default() }) .await?; - let mcp_connection_manager = self.session.services.mcp_connection_manager.load_full(); - mcp_connection_manager.set_elicitations_auto_deny(mcp_elicitations_auto_deny); + self.session + .services + .latest_mcp_runtime() + .manager() + .set_elicitations_auto_deny(mcp_elicitations_auto_deny); Ok(()) } @@ -2855,10 +2854,12 @@ impl Session { &environments.captured_environments(), ) .await; + let mcp = self.services.latest_mcp_runtime(); Arc::new(StepContext::new( turn_context, environments, selected_capability_roots, + mcp, loaded_agents_md, )) } @@ -3125,6 +3126,17 @@ impl Session { &self, turn_context: &TurnContext, world_state: &WorldState, + ) -> Vec { + let mcp = self.services.latest_mcp_runtime(); + self.build_initial_context_with_world_state_and_mcp(turn_context, world_state, &mcp) + .await + } + + async fn build_initial_context_with_world_state_and_mcp( + &self, + turn_context: &TurnContext, + world_state: &WorldState, + mcp: &McpRuntimeSnapshot, ) -> Vec { let mut developer_sections = Vec::::with_capacity(8); let mut contextual_user_sections = Vec::::with_capacity(2); @@ -3218,10 +3230,9 @@ impl Session { } } if turn_context.config.include_apps_instructions && turn_context.apps_enabled() { - let mcp_connection_manager = self.services.mcp_connection_manager.load_full(); let accessible_and_enabled_connectors = connectors::list_accessible_and_enabled_connectors_from_manager( - &mcp_connection_manager, + mcp.manager(), &turn_context.config, ) .await; @@ -3328,7 +3339,8 @@ impl Session { if turn_context.config.features.enabled(Feature::TokenBudget) && turn_context.model_context_window().is_some() { - let mcp_result = self + let mcp_result = mcp + .manager() .call_tool( "notes", "thread_hint", @@ -3543,7 +3555,11 @@ impl Session { // Full initial context resets the baseline; later turns persist only its changes. let (mut context_items, world_state_item) = if should_inject_full_context { let context_items = self - .build_initial_context_with_world_state(turn_context, world_state.as_ref()) + .build_initial_context_with_world_state_and_mcp( + turn_context, + world_state.as_ref(), + step_context.mcp.as_ref(), + ) .await; let snapshot = world_state.snapshot(); self.state diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index f385a8c983..2c9e57036d 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -661,10 +661,15 @@ impl Session { let config_for_mcp = Arc::clone(&config); let mcp_manager_for_mcp = Arc::clone(&mcp_manager); let mcp_thread_init_for_startup = &mcp_thread_init; - let mcp_runtime_context_for_auth = McpRuntimeContext::new( - Arc::clone(&environment_manager), - session_configuration.cwd().to_path_buf(), - ); + let mcp_runtime_cwd = session_configuration + .environment_selections() + .first() + .and_then(|environment| environment.cwd.to_abs_path().ok()) + .map(|cwd| cwd.to_path_buf()) + .unwrap_or_else(|| session_configuration.cwd().to_path_buf()); + let mcp_runtime_context = + McpRuntimeContext::new(Arc::clone(&environment_manager), mcp_runtime_cwd); + let mcp_runtime_context_for_auth = mcp_runtime_context.clone(); let auth_and_mcp_fut = async move { let auth = auth_manager_clone.auth().await; let mcp_config = mcp_manager_for_mcp @@ -680,7 +685,13 @@ impl Session { &mcp_runtime_context_for_auth, ) .await; - (auth, mcp_servers, auth_statuses, tool_plugin_provenance) + ( + auth, + mcp_config, + mcp_servers, + auth_statuses, + tool_plugin_provenance, + ) } .instrument(info_span!( "session_init.auth_mcp", @@ -691,7 +702,7 @@ impl Session { let ( thread_persistence_result, state_db_ctx, - (auth, mcp_servers, auth_statuses, tool_plugin_provenance), + (auth, mcp_config, mcp_servers, auth_statuses, tool_plugin_provenance), ) = tokio::join!(thread_persistence_fut, state_db_fut, auth_and_mcp_fut); let mut live_thread_init = @@ -1027,6 +1038,7 @@ impl Session { // changing this to use Option or OnceCell, though the current // setup is straightforward enough and performs well. mcp_connection_manager, + mcp_runtime: arc_swap::ArcSwapOption::empty(), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -1159,14 +1171,6 @@ impl Session { sess.send_event_raw(event).await; } - let client_elicitation_capability = if config.features.enabled(Feature::AuthElicitation) { - ElicitationCapability { - form: Some(FormElicitationCapability::default()), - url: Some(UrlElicitationCapability::default()), - } - } else { - ElicitationCapability::default() - }; let mcp_startup_cancellation_token = { let mut cancel_guard = sess.services.mcp_startup_cancellation_token.lock().await; cancel_guard.cancel(); @@ -1174,20 +1178,6 @@ impl Session { *cancel_guard = cancel_token.clone(); cancel_token }; - let mcp_runtime_context = { - let turn_environments = sess.services.turn_environments.snapshot().await; - // TODO(anp): Migrate MCP runtime cwd plumbing to PathUri so foreign environment - // cwd values can be used without falling back to the session host cwd. - let cwd = turn_environments - .primary() - .and_then(|turn_environment| turn_environment.cwd().to_abs_path().ok()) - .map(|cwd| cwd.to_path_buf()) - .unwrap_or_else(|| session_configuration.cwd().to_path_buf()); - McpRuntimeContext::new( - sess.services.turn_environments.environment_manager(), - cwd, - ) - }; let mcp_connection_manager = McpConnectionManager::new( &mcp_servers, config.mcp_oauth_credentials_store_mode, @@ -1198,12 +1188,12 @@ impl Session { tx_event.clone(), mcp_startup_cancellation_token, session_configuration.permission_profile(), - mcp_runtime_context, + mcp_runtime_context.clone(), config.codex_home.to_path_buf(), sess.services.mcp_manager.codex_apps_tools_cache(), codex_apps_tools_cache_key(auth), config.prefix_mcp_tool_names(), - client_elicitation_capability, + mcp_config.client_elicitation_capability.clone(), sess.services .supports_openai_form_elicitation .load(std::sync::atomic::Ordering::Relaxed), @@ -1217,7 +1207,11 @@ impl Session { )) .await; sess.services - .install_mcp_connection_manager(mcp_connection_manager) + .install_mcp_connection_manager( + Arc::new(mcp_config), + mcp_runtime_context, + mcp_connection_manager, + ) .await?; sess.schedule_startup_prewarm(session_configuration.base_instructions.clone()) .await; diff --git a/codex-rs/core/src/session/step_context.rs b/codex-rs/core/src/session/step_context.rs index e463015459..3210eb0286 100644 --- a/codex-rs/core/src/session/step_context.rs +++ b/codex-rs/core/src/session/step_context.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use crate::agents_md::LoadedAgentsMd; use crate::environment_selection::TurnEnvironmentSnapshot; +use crate::session::McpRuntimeSnapshot; use crate::session::turn_context::TurnContext; use codex_exec_server::ResolvedSelectedCapabilityRoot; @@ -12,6 +13,8 @@ pub(crate) struct StepContext { pub(crate) environments: TurnEnvironmentSnapshot, /// Capability roots bound to ready environments in this exact step. pub(crate) selected_capability_roots: Vec, + /// The exact MCP config and manager used to advertise and execute tools for this step. + pub(crate) mcp: Arc, /// The canonical AGENTS.md value observed with this environment snapshot. pub(crate) loaded_agents_md: Option>, } @@ -21,12 +24,14 @@ impl StepContext { turn: Arc, environments: TurnEnvironmentSnapshot, selected_capability_roots: Vec, + mcp: Arc, loaded_agents_md: Option>, ) -> Self { Self { turn, environments, selected_capability_roots, + mcp, loaded_agents_md, } } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 525fbc3b99..d0aac04249 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -18,6 +18,7 @@ use crate::test_support::models_manager_with_provider; use crate::tools::format_exec_output_str; use codex_config::ConfigLayerStack; use codex_config::ConfigLayerStackOrdering; +use codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID; use codex_config::LoaderOverrides; use codex_config::NetworkConstraints; use codex_config::NetworkDomainPermissionToml; @@ -25,6 +26,8 @@ use codex_config::NetworkDomainPermissionsToml; use codex_config::RequirementSource; use codex_config::Sourced; use codex_config::loader::project_trust_key; +use codex_config::types::McpServerConfig; +use codex_config::types::McpServerTransportConfig; use codex_config::types::ToolSuggestDisabledTool; use codex_core_skills::HostSkillsSnapshot; use core_test_support::test_codex::local_selections; @@ -188,9 +191,10 @@ impl StepContext { pub(crate) fn for_test(turn: Arc) -> Arc { let environments = turn.environments.clone(); Arc::new(Self::new( - turn, + Arc::clone(&turn), environments, Vec::new(), + crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(&turn.config), /*loaded_agents_md*/ None, )) } @@ -379,8 +383,8 @@ async fn request_mcp_server_elicitation_auto_accepts_when_auto_deny_is_enabled() let (session, turn_context, rx) = make_session_and_context_with_rx().await; session .services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager() .set_elicitations_auto_deny(/*auto_deny*/ true); let response = session @@ -5363,14 +5367,11 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { /*bundled_skills_enabled*/ true, )); let network_approval = Arc::new(NetworkApprovalService::default()); + let mcp_runtime = + crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(config.as_ref()); let services = SessionServices { - mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from_pointee( - McpConnectionManager::new_uninitialized_with_permission_profile( - &config.permissions.approval_policy, - config.permissions.permission_profile(), - config.prefix_mcp_tool_names(), - ), - )), + mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from(mcp_runtime.manager_arc())), + mcp_runtime: arc_swap::ArcSwapOption::from(Some(mcp_runtime)), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -7440,14 +7441,11 @@ where /*bundled_skills_enabled*/ true, )); let network_approval = Arc::new(NetworkApprovalService::default()); + let mcp_runtime = + crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(config.as_ref()); let services = SessionServices { - mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from_pointee( - McpConnectionManager::new_uninitialized_with_permission_profile( - &config.permissions.approval_policy, - config.permissions.permission_profile(), - config.prefix_mcp_tool_names(), - ), - )), + mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from(mcp_runtime.manager_arc())), + mcp_runtime: arc_swap::ArcSwapOption::from(Some(mcp_runtime)), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -7603,8 +7601,14 @@ pub(crate) async fn make_session_and_context_with_rx() -> ( } #[tokio::test] -async fn refresh_mcp_servers_is_deferred_until_next_turn() { +async fn refresh_mcp_servers_keeps_the_previous_runtime_alive() { let (session, turn_context) = make_session_and_context().await; + let turn_context = Arc::new(turn_context); + let old_runtime = session.services.latest_mcp_runtime(); + let step_context = session + .capture_step_context(Arc::clone(&turn_context)) + .await; + assert!(Arc::ptr_eq(&step_context.mcp, &old_runtime)); let old_token = session.mcp_startup_cancellation_token().await; assert!(!old_token.is_cancelled()); @@ -7612,8 +7616,16 @@ async fn refresh_mcp_servers_is_deferred_until_next_turn() { serde_json::to_value(OAuthCredentialsStoreMode::Auto).expect("serialize store mode"); let auth_keyring_backend_kind = serde_json::to_value(AuthKeyringBackendKind::Secrets).expect("serialize keyring backend"); + let refreshed_mcp_servers = serde_json::from_value::>(json!({ + "refreshed": { + "url": "https://refreshed.example/mcp", + "enabled": false + } + })) + .expect("parse refreshed MCP servers"); let refresh_config = McpServerRefreshConfig { - mcp_servers: json!({}), + mcp_servers: serde_json::to_value(&refreshed_mcp_servers) + .expect("serialize refreshed MCP servers"), mcp_oauth_credentials_store_mode, auth_keyring_backend_kind, }; @@ -7635,7 +7647,7 @@ async fn refresh_mcp_servers_is_deferred_until_next_turn() { .refresh_mcp_servers_if_requested(&turn_context, /*elicitation_reviewer*/ None) .await; - assert!(old_token.is_cancelled()); + assert!(!old_token.is_cancelled()); assert!( session .pending_mcp_server_refresh_config @@ -7645,6 +7657,71 @@ async fn refresh_mcp_servers_is_deferred_until_next_turn() { ); let new_token = session.mcp_startup_cancellation_token().await; assert!(!new_token.is_cancelled()); + let new_runtime = session.services.latest_mcp_runtime(); + assert!(!Arc::ptr_eq(&old_runtime, &new_runtime)); + assert!(Arc::ptr_eq(&step_context.mcp, &old_runtime)); + assert_eq!( + codex_mcp::configured_mcp_servers(new_runtime.config()), + refreshed_mcp_servers + ); +} + +#[tokio::test] +async fn built_tools_uses_the_step_mcp_runtime() -> anyhow::Result<()> { + let (session, turn_context) = make_session_and_context().await; + let session = Arc::new(session); + let turn_context = Arc::new(turn_context); + let step_context = session.capture_step_context(turn_context).await; + + let mut refresh_config = step_context.turn.config.as_ref().clone(); + refresh_config.mcp_servers.set(HashMap::from([( + "newer".to_string(), + McpServerConfig { + auth: Default::default(), + transport: McpServerTransportConfig::Stdio { + command: "missing-test-mcp-server".to_string(), + args: Vec::new(), + env: None, + env_vars: Vec::new(), + cwd: None, + }, + environment_id: DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), + enabled: true, + required: false, + supports_parallel_tool_calls: false, + disabled_reason: None, + startup_timeout_sec: None, + tool_timeout_sec: None, + default_tools_approval_mode: None, + enabled_tools: None, + disabled_tools: None, + scopes: None, + oauth: None, + oauth_resource: None, + tools: HashMap::new(), + }, + )]))?; + session + .refresh_mcp_servers_now( + step_context.turn.as_ref(), + &refresh_config, + /*elicitation_reviewer*/ None, + ) + .await; + + let router = crate::session::turn::built_tools( + session.as_ref(), + step_context.as_ref(), + &CancellationToken::new(), + ) + .await?; + assert!( + !router + .registered_tool_names_for_test() + .iter() + .any(|name| name.to_string() == "list_mcp_resources") + ); + Ok(()) } #[tokio::test] @@ -10161,8 +10238,8 @@ async fn fatal_tool_error_stops_turn_and_reports_error() { let tools = { session .services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager() .list_all_tools() .await }; diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 9b3588fc1e..2cb1f2cada 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1177,7 +1177,7 @@ pub(crate) async fn built_tools( cancellation_token: &CancellationToken, ) -> CodexResult> { let turn_context = step_context.turn.as_ref(); - let mcp_connection_manager = sess.services.mcp_connection_manager.load_full(); + let mcp_connection_manager = step_context.mcp.manager(); let has_mcp_servers = mcp_connection_manager.has_servers(); let all_mcp_tools = mcp_connection_manager .list_all_tools() diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index a5ae458957..5efbeb8ef3 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -686,7 +686,8 @@ impl Session { .unwrap_or_else(|| session_configuration.cwd().clone()); let per_turn_config = Self::build_per_turn_config(&session_configuration, cwd.clone()); { - let mcp_connection_manager = self.services.mcp_connection_manager.load_full(); + let mcp_runtime = self.services.latest_mcp_runtime(); + let mcp_connection_manager = mcp_runtime.manager(); mcp_connection_manager.set_approval_policy(&session_configuration.approval_policy); mcp_connection_manager .set_permission_profile(session_configuration.permission_profile()); diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 3721cb1848..1bee45cce0 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -15,6 +15,7 @@ use crate::exec_policy::ExecPolicyManager; use crate::guardian::GuardianRejection; use crate::guardian::GuardianRejectionCircuitBreaker; use crate::mcp::McpManager; +use crate::session::McpRuntimeSnapshot; use crate::tools::code_mode::CodeModeService; use crate::tools::handlers::ToolSearchHandlerCache; use crate::tools::network_approval::NetworkApprovalService; @@ -30,7 +31,9 @@ use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionRegistry; use codex_hooks::Hooks; use codex_login::AuthManager; +use codex_mcp::McpConfig; use codex_mcp::McpConnectionManager; +use codex_mcp::McpRuntimeContext; use codex_models_manager::manager::SharedModelsManager; use codex_otel::SessionTelemetry; use codex_protocol::capabilities::SelectedCapabilityRoot; @@ -44,8 +47,10 @@ use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; pub(crate) struct SessionServices { - /// The latest manager; callers retain an owned handle while performing MCP I/O. + /// Mirror of the latest manager for extension resource clients that predate runtime snapshots. pub(crate) mcp_connection_manager: Arc>, + /// The latest atomically published MCP config and manager pair. + pub(crate) mcp_runtime: ArcSwapOption, pub(crate) mcp_startup_cancellation_token: Mutex, pub(crate) unified_exec_manager: UnifiedExecProcessManager, #[cfg_attr(not(unix), allow(dead_code))] @@ -99,12 +104,33 @@ impl SessionServices { /// resolve through the session's manager while validation waits. pub(crate) async fn install_mcp_connection_manager( &self, + config: Arc, + runtime_context: McpRuntimeContext, manager: McpConnectionManager, ) -> Result<()> { - self.mcp_connection_manager.store(Arc::new(manager)); - self.mcp_connection_manager - .load_full() - .validate_required_servers() - .await + let runtime = self.publish_mcp_runtime(config, runtime_context, manager); + runtime.manager().validate_required_servers().await + } + + pub(crate) fn publish_mcp_runtime( + &self, + config: Arc, + runtime_context: McpRuntimeContext, + manager: McpConnectionManager, + ) -> Arc { + let manager = Arc::new(manager); + // Publish the manager for legacy resource clients first. Once the paired snapshot is + // visible, every model-scoped consumer observes this exact manager. + self.mcp_connection_manager.store(Arc::clone(&manager)); + let runtime = Arc::new(McpRuntimeSnapshot::new(config, manager, runtime_context)); + self.mcp_runtime.store(Some(Arc::clone(&runtime))); + runtime + } + + pub(crate) fn latest_mcp_runtime(&self) -> Arc { + let Some(runtime) = self.mcp_runtime.load_full() else { + unreachable!("MCP runtime must be installed before handling requests"); + }; + runtime } } diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index b92a2a5655..c9533fd223 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -124,11 +124,12 @@ impl McpHandler { ) -> Result, FunctionCallError> { let ToolInvocation { session, - turn, + step_context, call_id, payload, .. } = invocation; + let turn = Arc::clone(&step_context.turn); let payload = match payload { ToolPayload::Function { arguments } => arguments, @@ -143,7 +144,7 @@ impl McpHandler { // TODO(sayan): Use StepContext for MCP file arguments when MCP follows dynamic environments. let result = handle_mcp_tool_call( Arc::clone(&session), - &turn, + &step_context, call_id.clone(), self.tool_info.server_name.clone(), self.tool_info.tool.name.to_string(), diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs index ce762a0f02..8c7f21fd17 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs @@ -53,11 +53,13 @@ impl ListMcpResourceTemplatesHandler { ) -> Result, FunctionCallError> { let ToolInvocation { session, - turn, + step_context, call_id, payload, .. } = invocation; + let turn = std::sync::Arc::clone(&step_context.turn); + let manager = step_context.mcp.manager(); let arguments = match payload { ToolPayload::Function { arguments } => arguments, @@ -89,7 +91,7 @@ impl ListMcpResourceTemplatesHandler { let params = cursor .clone() .map(|value| PaginatedRequestParams::default().with_cursor(Some(value))); - let result = session + let result = manager .list_resource_templates(&server_name, params) .await .map_err(|err| { @@ -108,10 +110,7 @@ impl ListMcpResourceTemplatesHandler { )); } - let templates = session - .services - .mcp_connection_manager - .load_full() + let templates = manager .list_all_resource_templates(|server_name| { model_can_access_mcp_server(turn.as_ref(), server_name) }) diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs index 564fb440fd..f5e9bae7e3 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs @@ -53,11 +53,13 @@ impl ListMcpResourcesHandler { ) -> Result, FunctionCallError> { let ToolInvocation { session, - turn, + step_context, call_id, payload, .. } = invocation; + let turn = std::sync::Arc::clone(&step_context.turn); + let manager = step_context.mcp.manager(); let arguments = match payload { ToolPayload::Function { arguments } => arguments, @@ -89,7 +91,7 @@ impl ListMcpResourcesHandler { let params = cursor .clone() .map(|value| PaginatedRequestParams::default().with_cursor(Some(value))); - let result = session + let result = manager .list_resources(&server_name, params) .await .map_err(|err| { @@ -106,10 +108,7 @@ impl ListMcpResourcesHandler { )); } - let resources = session - .services - .mcp_connection_manager - .load_full() + let resources = manager .list_all_resources(|server_name| { model_can_access_mcp_server(turn.as_ref(), server_name) }) diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs b/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs index 3909011eb0..5e9b85241c 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs @@ -52,11 +52,13 @@ impl ReadMcpResourceHandler { ) -> Result, FunctionCallError> { let ToolInvocation { session, - turn, + step_context, call_id, payload, .. } = invocation; + let turn = std::sync::Arc::clone(&step_context.turn); + let manager = step_context.mcp.manager(); let arguments = match payload { ToolPayload::Function { arguments } => arguments, @@ -84,7 +86,7 @@ impl ReadMcpResourceHandler { let payload_result: Result = async { ensure_model_can_access_mcp_server(turn.as_ref(), &server)?; - let result = session + let result = manager .read_resource(&server, ReadResourceRequestParams::new(uri.clone())) .await .map_err(|err| { diff --git a/codex-rs/core/src/tools/handlers/request_plugin_install.rs b/codex-rs/core/src/tools/handlers/request_plugin_install.rs index ce7ce10da8..e44ed41fa4 100644 --- a/codex-rs/core/src/tools/handlers/request_plugin_install.rs +++ b/codex-rs/core/src/tools/handlers/request_plugin_install.rs @@ -1,4 +1,5 @@ use std::collections::HashSet; +use std::sync::Arc; use codex_analytics::PluginInstallRequestSource; use codex_analytics::PluginInstallRequested; @@ -94,10 +95,12 @@ impl RequestPluginInstallHandler { let ToolInvocation { payload, session, - turn, + step_context, call_id, .. } = invocation; + let turn = Arc::clone(&step_context.turn); + let manager = step_context.mcp.manager(); let arguments = match payload { ToolPayload::Function { arguments } => arguments, @@ -227,7 +230,8 @@ impl RequestPluginInstallHandler { let auth = session.services.auth_manager.auth().await; let completed = if user_confirmed { - verify_request_plugin_install_completed(&session, &turn, &tool, auth.as_ref()).await + verify_request_plugin_install_completed(&session, &turn, manager, &tool, auth.as_ref()) + .await } else { false }; @@ -345,13 +349,14 @@ fn disabled_install_request(tool: &DiscoverableTool) -> ToolSuggestDisabledTool async fn verify_request_plugin_install_completed( session: &crate::session::session::Session, turn: &crate::session::turn_context::TurnContext, + manager: &codex_mcp::McpConnectionManager, tool: &DiscoverableTool, auth: Option<&codex_login::CodexAuth>, ) -> bool { match tool { DiscoverableTool::Connector(connector) => refresh_missing_requested_connectors( - session, turn, + manager, auth, std::slice::from_ref(&connector.id), connector.id.as_str(), @@ -370,8 +375,8 @@ async fn verify_request_plugin_install_completed( plugin.id.as_str(), ), refresh_missing_requested_connectors( - session, turn, + manager, auth, &plugin.app_connector_ids, plugin.id.as_str(), @@ -393,8 +398,8 @@ async fn verify_request_plugin_install_completed( session.services.plugins_manager.as_ref(), ); let _ = refresh_missing_requested_connectors( - session, turn, + manager, auth, &plugin.app_connector_ids, plugin.id.as_str(), @@ -435,8 +440,8 @@ fn is_remote_plugin_install_suggestion(plugin_id: &str) -> bool { } async fn refresh_missing_requested_connectors( - session: &crate::session::session::Session, turn: &crate::session::turn_context::TurnContext, + manager: &codex_mcp::McpConnectionManager, auth: Option<&codex_login::CodexAuth>, expected_connector_ids: &[String], tool_id: &str, @@ -445,7 +450,6 @@ async fn refresh_missing_requested_connectors( return Some(Vec::new()); } - let manager = session.services.mcp_connection_manager.load_full(); let mcp_tools = manager.list_all_tools().await; let accessible_connectors = connectors::with_app_enabled_state( connectors::accessible_connectors_from_mcp_tools(&mcp_tools), diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index db1209cc03..6f399fd5fc 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -110,8 +110,8 @@ async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow let step_context = StepContext::for_test(Arc::clone(&turn)); let mcp_tools = session .services - .mcp_connection_manager - .load_full() + .latest_mcp_runtime() + .manager() .list_all_tools() .await; let router = ToolRouter::from_context( diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 6fc5adcb11..498a690c58 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -686,10 +686,12 @@ async fn environment_tools_follow_the_step_context() { let environments = turn.environments.clone(); turn.environments.turn_environments.clear(); + let turn = Arc::new(turn); let step_context = Arc::new(StepContext::new( - Arc::new(turn), + Arc::clone(&turn), environments, Vec::new(), + crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(&turn.config), /*loaded_agents_md*/ None, )); diff --git a/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs b/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs index 7df18727b9..620b6c734c 100644 --- a/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs +++ b/codex-rs/core/tests/suite/mcp_refresh_cleanup.rs @@ -16,7 +16,7 @@ use core_test_support::test_codex::test_codex; use core_test_support::wait_for_mcp_server; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn refresh_shuts_down_superseded_mcp_stdio_server() -> anyhow::Result<()> { +async fn refresh_keeps_superseded_mcp_server_alive_for_in_flight_calls() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; @@ -84,7 +84,7 @@ async fn refresh_shuts_down_superseded_mcp_stdio_server() -> anyhow::Result<()> "sync", Some(serde_json::json!({ "barrier": barrier, - "sleep_after_ms": 30_000 + "sleep_after_ms": 300_000 })), /*meta*/ None, ) @@ -119,9 +119,16 @@ async fn refresh_shuts_down_superseded_mcp_stdio_server() -> anyhow::Result<()> let replacement_pid = wait_for_pid_file(&pid_file).await?; assert_ne!(replacement_pid, superseded_pid); + assert!(process_is_alive(&superseded_pid)?); + long_call.abort(); + assert!( + long_call + .await + .expect_err("call should be aborted") + .is_cancelled() + ); wait_for_process_exit(&superseded_pid).await?; assert!(process_is_alive(&replacement_pid)?); - assert!(long_call.await?.is_err()); fixture.codex.shutdown_and_wait().await?; wait_for_process_exit(&replacement_pid).await From 5044062704af7667e460cc401221a081054ef3ab Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Thu, 25 Jun 2026 17:01:43 -0700 Subject: [PATCH 04/38] ci: narrow Windows test skips (#30134) ## Why The Windows cross-build skip used the broad `powershell` substring, which hid unrelated Windows tests. Narrowing it exposed the same ConPTY Ctrl-C timeout that is breaking `main`; that test is not reliable in either cross-built or native Windows Bazel CI yet. ## What changed - scope the cross-build PowerShell carve-out to the dedicated parser-process test module - exclude the exact ConPTY Ctrl-C test from Bazel CI while leaving local Windows runs enabled - repeat the exact exclusion in the cross-build config because it replaces the base skip list ## Manual validation - `just test-github-scripts` - queried the PTY test target under both `ci-windows` and `ci-windows-cross` - verified the matcher excludes parser-process and ConPTY tests without excluding unrelated PowerShell tests - [Windows shard 4/4](https://github.com/openai/codex/actions/runs/28204844286/job/83553063859) reproduced the `main` ConPTY timeout before the exact CI-only exclusion was applied --- .bazelrc | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/.bazelrc b/.bazelrc index c58667b46e..de01a87443 100644 --- a/.bazelrc +++ b/.bazelrc @@ -107,8 +107,9 @@ common:ci-bazel --config=ci common:ci-bazel --build_metadata=TAG_workflow=bazel # Keep code-mode integration cases out of ordinary Bazel legs. The # Windows-cross config below re-enables them after generating its Windows V8 -# snapshot on the Windows runner. -common:ci-bazel --test_env=CODEX_BAZEL_TEST_SKIP_FILTERS=suite::code_mode:: +# snapshot on the Windows runner. Also skip the ConPTY Ctrl-C integration test +# until Windows Bazel CI can reliably interrupt foreground processes. +common:ci-bazel --test_env=CODEX_BAZEL_TEST_SKIP_FILTERS=suite::code_mode::,tests::windows_tests::conpty_ctrl_c_interrupts_powershell_foreground_child # Shared config for Bazel-backed Rust linting. build:clippy --aspects=@rules_rust//rust:defs.bzl%rust_clippy_aspect @@ -191,11 +192,12 @@ common:ci-windows-cross --strategy=TestRunner=local common:ci-windows-cross --strategy=V8Mksnapshot=local common:ci-windows-cross --local_test_jobs=4 common:ci-windows-cross --test_env=RUST_TEST_THREADS=1 -# Native Windows CI still covers the PowerShell tests. The cross-built gnullvm -# binaries currently hang in PowerShell AST parser tests when those binaries are -# run on the Windows runner. Keep V8-backed code-mode tests enabled except for -# the hidden dynamic-tool callback test, which currently times out on Windows. -common:ci-windows-cross --test_env=CODEX_BAZEL_TEST_SKIP_FILTERS=powershell,suite::code_mode::code_mode_can_call_hidden_dynamic_tools +# Native Windows CI still covers the PowerShell parser-process tests. The +# cross-built gnullvm binaries currently hang in those tests when run on the +# Windows runner. Keep V8-backed code-mode tests enabled except for the hidden +# dynamic-tool callback test, which currently times out on Windows. This config +# replaces the base skip list, so repeat the exact ConPTY Ctrl-C test exclusion. +common:ci-windows-cross --test_env=CODEX_BAZEL_TEST_SKIP_FILTERS=command_safety::powershell_parser::tests::,suite::code_mode::code_mode_can_call_hidden_dynamic_tools,tests::windows_tests::conpty_ctrl_c_interrupts_powershell_foreground_child common:ci-windows-cross --platforms=//:windows_x86_64_gnullvm common:ci-windows-cross --extra_execution_platforms=//:rbe,//:windows_x86_64_msvc common:ci-windows-cross --extra_toolchains=//:windows_gnullvm_tests_on_msvc_host_toolchain From 3095ea9c3d155bfc89197d2628eb818a55c2755d Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 26 Jun 2026 01:36:44 +0100 Subject: [PATCH 05/38] Project selected plugin runtime by environment availability (#30093) ## Why Selected plugin metadata is stable, but MCP processes are live runtime state. They need different lifetimes: - the MCP extension caches manifest, MCP, and connector declarations for each stable selected root; - each model step projects that cached metadata through the roots that resolved as ready for that exact step; - the MCP manager is rebuilt only when that availability projection changes. This matches executor skills: both features consume the same resolved step roots instead of inferring readiness from the turn's selected environments. ## Behavior ```text E1 not ready for this step -> no E1 MCP servers or connectors -> cached plugin metadata stays in ext/mcp E1 becomes ready -> reuse cached metadata -> publish one MCP runtime containing E1 capabilities same ready roots on the next step -> reuse the exact runtime; no rediscovery and no MCP restart resume -> create new extension thread state and a new MCP runtime ``` All model-facing consumers use the same step snapshot: ```text resolved selected roots | v extension MCP/connector projection | v { MCP config, connector snapshot, MCP manager } | +-> advertise model tools +-> build app/connector tools +-> execute MCP calls ``` ## Cache contract The existing MCP extension owns a cache keyed by the full `SelectedCapabilityRoot`: ```rust let state = thread_store.get_or_init(SelectedExecutorPluginMcpState::default); ``` The cache lives with extension thread state. Environment availability filters projection but does not invalidate metadata. Resume creates new thread state. There is no file watcher or executor generation because contents behind a stable environment/root are assumed stable. ## What changes - Keeps executor plugin discovery and cached metadata in `ext/mcp`. - Caches MCP and connector declarations together per selected root. - Uses the step's already-resolved capability roots, including lazy environments that are not turn environments. - Reuses the current MCP runtime when the ready-root projection is unchanged. - Uses the same step MCP manager and connector snapshot for model-visible tools and execution. - Resolves direct thread-scoped MCP requests from the current selected-root projection. ## Deliberately out of scope - `app/list` remains based on the latest global host-plugin state; this PR does not make its response or notifications thread-specific. - `required = true` startup semantics do not apply to delayed executor MCP activation. - No filesystem/content invalidation. - No transport-disconnect watcher. - No executor generations or environment replacement semantics. - No client sharing across complete manager replacements. ## Stack 1. Extension-owned World State sections. 2. Project executor skills through World State. 3. Pin one MCP runtime to each model step. 4. **This PR:** project selected MCP and connector state from extension-owned metadata. 5. Integration coverage for selected capability availability and resume. ## Verification - `selected_plugin_servers_use_managed_requirements_for_the_selected_root_id` - The stacked integration PR covers unavailable to ready activation, unchanged-runtime reuse, skills, MCP tools, connector attribution, and cold resume. --- codex-rs/Cargo.lock | 1 + codex-rs/Cargo.toml | 1 + .../src/request_processors/mcp_processor.rs | 23 +- .../request_processors/thread_processor.rs | 1 - .../app-server/tests/suite/v2/executor_mcp.rs | 1 - codex-rs/core/src/codex_thread.rs | 20 +- codex-rs/core/src/mcp.rs | 45 +++- codex-rs/core/src/prompt_debug.rs | 4 +- codex-rs/core/src/session/mcp.rs | 196 ++++++++++++++---- codex-rs/core/src/session/mcp_runtime.rs | 8 + codex-rs/core/src/session/mod.rs | 19 +- codex-rs/core/src/session/session.rs | 10 +- codex-rs/core/src/session/tests.rs | 3 + codex-rs/core/src/session/turn.rs | 27 +-- codex-rs/core/src/state/service.rs | 14 +- codex-rs/core/src/thread_manager_tests.rs | 24 ++- .../ext/extension-api/src/contributors/mcp.rs | 47 ++++- codex-rs/ext/mcp/Cargo.toml | 1 + codex-rs/ext/mcp/src/executor_plugin.rs | 154 ++++++++++---- codex-rs/ext/mcp/src/lib.rs | 7 - codex-rs/ext/mcp/tests/executor_plugin_mcp.rs | 12 +- 21 files changed, 455 insertions(+), 163 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index b147022825..75c02adfb5 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3359,6 +3359,7 @@ name = "codex-mcp-extension" version = "0.0.0" dependencies = [ "codex-config", + "codex-connectors-extension", "codex-core", "codex-core-plugins", "codex-exec-server", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index ef17cb9ce5..0824f7db7b 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -166,6 +166,7 @@ codex-code-mode-protocol = { path = "code-mode-protocol" } codex-home = { path = "codex-home" } codex-config = { path = "config" } codex-connectors = { path = "connectors" } +codex-connectors-extension = { path = "ext/connectors" } codex-context-fragments = { path = "context-fragments" } codex-core = { path = "core" } codex-core-api = { path = "core-api" } diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index abefd26015..ac2125c2cf 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -124,7 +124,7 @@ impl McpRequestProcessor { let (mcp_config, runtime_context) = match thread_id.as_deref() { Some(thread_id) => { let (_, thread) = self.load_thread(thread_id).await?; - let runtime = thread.current_mcp_runtime(); + let runtime = thread.current_mcp_runtime().await; (runtime.config().clone(), runtime.runtime_context().clone()) } None => { @@ -246,14 +246,21 @@ impl McpRequestProcessor { let mcp_manager = self.thread_manager.mcp_manager(); let codex_apps_tools_cache = mcp_manager.codex_apps_tools_cache(); let auth = self.auth_manager.auth().await; - let mcp_config = match thread { - Some(thread) => thread.runtime_mcp_config(&config).await, - None => mcp_manager.runtime_config(&config).await, + let (mcp_config, runtime_context) = match thread { + Some(thread) => { + let mcp_config = thread.runtime_mcp_config(&config).await; + let runtime = thread.current_mcp_runtime().await; + (mcp_config, runtime.runtime_context().clone()) + } + None => { + let mcp_config = mcp_manager.runtime_config(&config).await; + let runtime_context = McpRuntimeContext::new( + self.thread_manager.environment_manager(), + config.cwd.to_path_buf(), + ); + (mcp_config, runtime_context) + } }; - let runtime_context = McpRuntimeContext::new( - self.thread_manager.environment_manager(), - config.cwd.to_path_buf(), - ); tokio::spawn(async move { Self::list_mcp_server_status_task( diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 0fc37d1a4f..39283bb35c 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -1155,7 +1155,6 @@ impl ThreadRequestProcessor { let mut thread_extension_init = ExtensionDataInit::new(); if !selected_capability_roots.is_empty() { thread_extension_init.insert(selected_capability_roots); - codex_mcp_extension::initialize_executor_plugin_thread_data(&mut thread_extension_init); } let create_thread_started_at = std::time::Instant::now(); let NewThread { diff --git a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs index 67d97bc248..aa650c3ecf 100644 --- a/codex-rs/app-server/tests/suite/v2/executor_mcp.rs +++ b/codex-rs/app-server/tests/suite/v2/executor_mcp.rs @@ -191,7 +191,6 @@ HTTP_PROXY = {http_proxy} ) .await?; - std::fs::write(plugin.path().join(".mcp.json"), r#"{"mcpServers":{}}"#)?; let config_path = codex_home.path().join("config.toml"); let mut config = std::fs::read_to_string(&config_path)?; config.push_str(&format!( diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 46b38af151..90555db911 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -595,8 +595,14 @@ impl CodexThread { } /// Returns the exact MCP config, environment bindings, and manager most recently published. - pub fn current_mcp_runtime(&self) -> Arc { - self.codex.session.services.latest_mcp_runtime() + pub async fn current_mcp_runtime(&self) -> Arc { + let turn_context = self.codex.session.new_default_turn().await; + self.codex + .session + .capture_step_context(turn_context) + .await + .mcp + .clone() } pub fn multi_agent_version(&self) -> Option { @@ -620,8 +626,9 @@ impl CodexThread { uri: &str, ) -> anyhow::Result { let result = self - .codex - .session + .current_mcp_runtime() + .await + .manager_arc() .read_resource(server, ReadResourceRequestParams::new(uri)) .await?; @@ -635,8 +642,9 @@ impl CodexThread { arguments: Option, meta: Option, ) -> anyhow::Result { - self.codex - .session + self.current_mcp_runtime() + .await + .manager_arc() .call_tool(server, tool, arguments, meta) .await } diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 658c5471ee..45ca0d563d 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -3,7 +3,10 @@ use std::sync::Arc; use crate::config::Config; use codex_config::McpServerConfig; +use codex_connectors::ConnectorSnapshot; +use codex_connectors::PluginConnectorSource; use codex_core_plugins::PluginsManager; +use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionRegistry; use codex_extension_api::McpServerContribution; @@ -18,6 +21,7 @@ use codex_mcp::McpServerRegistration; use codex_mcp::codex_apps_mcp_server_config; use codex_mcp::configured_mcp_servers; use codex_mcp::effective_mcp_servers; +use codex_plugin::AppConnectorId; const LEGACY_CODEX_APPS_REGISTRATION_ID: &str = "legacy_codex_apps"; @@ -69,28 +73,32 @@ impl McpManager { /// Returns the MCP config after applying compatibility built-ins and /// runtime-only extension overlays. pub async fn runtime_config(&self, config: &Config) -> McpConfig { - self.runtime_config_with_context(config, /*thread_init*/ None) + self.runtime_config_with_context(McpServerContributionContext::global(config)) .await } - pub(crate) async fn runtime_config_for_thread( + pub(crate) async fn runtime_config_for_step( &self, config: &Config, thread_init: &ExtensionDataInit, + thread_store: &ExtensionData, + available_environment_ids: &[String], ) -> McpConfig { - self.runtime_config_with_context(config, Some(thread_init)) - .await + self.runtime_config_with_context(McpServerContributionContext::for_step( + config, + thread_init, + thread_store, + available_environment_ids, + )) + .await } async fn runtime_config_with_context( &self, - config: &Config, - thread_init: Option<&ExtensionDataInit>, + context: McpServerContributionContext<'_, Config>, ) -> McpConfig { - let context = match thread_init { - Some(thread_init) => McpServerContributionContext::for_thread(config, thread_init), - None => McpServerContributionContext::global(config), - }; + let config = context.config(); + let mut selected_plugin_connector_sources = Vec::new(); let mut selected_plugin_registrations = Vec::new(); let mut overlays = Vec::new(); // A contributor can emit multiple ordered actions, so order each action globally rather @@ -121,6 +129,17 @@ impl McpManager { *config, ), ), + McpServerContribution::SelectedPluginConnectors { + plugin_id, + plugin_display_name, + connector_ids, + } => selected_plugin_connector_sources.push( + PluginConnectorSource::from_connector_ids( + plugin_id, + plugin_display_name, + connector_ids.into_iter().map(AppConnectorId), + ), + ), McpServerContribution::Remove { name } => { overlays.push(OrderedMcpOverlay::Remove { contributor_id: contributor.id(), @@ -186,6 +205,12 @@ impl McpManager { ); } mcp_config.mcp_server_catalog = catalog; + mcp_config.connector_snapshot = + mcp_config + .connector_snapshot + .merged_with(&ConnectorSnapshot::from_plugin_sources( + selected_plugin_connector_sources, + )); mcp_config } diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index a866116a24..cac9bb41ae 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -64,7 +64,7 @@ pub async fn build_prompt_input( ); let thread = thread_manager.start_thread(config).await?; - let output = build_prompt_input_from_session(thread.thread.codex.session.as_ref(), input).await; + let output = build_prompt_input_from_session(&thread.thread.codex.session, input).await; let shutdown = thread.thread.shutdown_and_wait().await; let _removed = thread_manager.remove_thread(&thread.thread_id).await; @@ -73,7 +73,7 @@ pub async fn build_prompt_input( } pub(crate) async fn build_prompt_input_from_session( - sess: &Session, + sess: &Arc, input: Vec, ) -> CodexResult> { let turn_context = sess.new_default_turn().await; diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index b9e87f0c89..50921f1591 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -1,7 +1,9 @@ use super::*; +use codex_exec_server::ResolvedSelectedCapabilityRoot; use codex_mcp::ElicitationReviewRequest; use codex_mcp::ElicitationReviewer; use codex_mcp::ElicitationReviewerHandle; +use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::mcp_approval_meta::APPROVAL_KIND_KEY as MCP_ELICITATION_APPROVAL_KIND_KEY; use codex_protocol::mcp_approval_meta::APPROVAL_KIND_MCP_TOOL_CALL as MCP_ELICITATION_APPROVAL_KIND_MCP_TOOL_CALL; @@ -75,9 +77,20 @@ impl ElicitationReviewer for GuardianMcpElicitationReviewer { impl Session { pub(crate) async fn runtime_mcp_config(&self, config: &Config) -> McpConfig { + let environments = self.services.turn_environments.snapshot().await; + let selected_capability_roots = self + .resolve_selected_capability_roots_for_step(&environments) + .await; + let available_environment_ids = + Self::available_selected_environment_ids(&selected_capability_roots); self.services .mcp_manager - .runtime_config_for_thread(config, &self.services.mcp_thread_init) + .runtime_config_for_step( + config, + &self.services.mcp_thread_init, + &self.services.thread_extension_data, + &available_environment_ids, + ) .await } @@ -88,6 +101,62 @@ impl Session { codex_mcp::configured_mcp_servers(&self.runtime_mcp_config(config).await) } + #[expect( + clippy::await_holding_invalid_type, + reason = "MCP runtime comparison and publication must remain serialized" + )] + pub(crate) async fn mcp_runtime_for_step( + self: &Arc, + turn_context: &TurnContext, + environments: &TurnEnvironmentSnapshot, + selected_capability_roots: &[ResolvedSelectedCapabilityRoot], + ) -> Arc { + let available_environment_ids = + Self::available_selected_environment_ids(selected_capability_roots); + let current = self.services.latest_mcp_runtime(); + if current.available_environment_ids() == available_environment_ids { + return current; + } + + let _guard = self.services.mcp_projection_lock.lock().await; + let current = self.services.latest_mcp_runtime(); + if current.available_environment_ids() == available_environment_ids { + return current; + } + let mcp_config = self + .services + .mcp_manager + .runtime_config_for_step( + &turn_context.config, + &self.services.mcp_thread_init, + &self.services.thread_extension_data, + &available_environment_ids, + ) + .await; + self.refresh_mcp_servers_inner( + turn_context, + mcp_config, + environments, + &available_environment_ids, + Some(self.mcp_elicitation_reviewer()), + ) + .await + } + + pub(crate) async fn resolve_selected_capability_roots_for_step( + &self, + environments: &TurnEnvironmentSnapshot, + ) -> Vec { + self.services + .turn_environments + .environment_manager() + .resolve_selected_capability_roots( + &self.services.selected_capability_roots, + &environments.captured_environments(), + ) + .await + } + pub(crate) fn mcp_elicitation_reviewer(self: &Arc) -> ElicitationReviewerHandle { Arc::new(GuardianMcpElicitationReviewer::new(self)) } @@ -207,51 +276,22 @@ impl Session { .await } - pub async fn read_resource( - &self, - server: &str, - params: ReadResourceRequestParams, - ) -> anyhow::Result { - self.services - .latest_mcp_runtime() - .manager_arc() - .read_resource(server, params) - .await - } - - pub async fn call_tool( - &self, - server: &str, - tool: &str, - arguments: Option, - meta: Option, - ) -> anyhow::Result { - self.services - .latest_mcp_runtime() - .manager_arc() - .call_tool(server, tool, arguments, meta) - .await - } - async fn refresh_mcp_servers_inner( &self, turn_context: &TurnContext, - mut mcp_config: McpConfig, - configured_mcp_servers: HashMap, + mcp_config: McpConfig, + environments: &TurnEnvironmentSnapshot, + available_environment_ids: &[String], elicitation_reviewer: Option, - ) { - mcp_config.mcp_server_catalog = mcp_config - .mcp_server_catalog - .with_materialized_servers(configured_mcp_servers); - let mcp_config = Arc::new(mcp_config); + ) -> Arc { let auth = self.services.auth_manager.auth().await; + let mcp_config = Arc::new(mcp_config); let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(&mcp_config); let mcp_servers = effective_mcp_servers(&mcp_config, auth.as_ref()); let environment_manager = self.services.turn_environments.environment_manager(); // TODO(anp): Migrate MCP runtime cwd plumbing to PathUri so foreign environment cwd // values can be used without falling back to the legacy host cwd. - let cwd = turn_context - .environments + let cwd = environments .primary() .and_then(|turn_environment| turn_environment.cwd().to_abs_path().ok()) .map(|cwd| cwd.to_path_buf()) @@ -305,10 +345,18 @@ impl Session { refreshed_manager .set_elicitations_auto_deny(current_manager.manager().elicitations_auto_deny()); } - self.services - .publish_mcp_runtime(mcp_config, mcp_runtime_context, refreshed_manager); + self.services.publish_mcp_runtime( + mcp_config, + mcp_runtime_context, + available_environment_ids.to_vec(), + refreshed_manager, + ) } + #[expect( + clippy::await_holding_invalid_type, + reason = "MCP runtime refresh and publication must remain serialized" + )] pub(crate) async fn refresh_mcp_servers_if_requested( &self, turn_context: &TurnContext, @@ -365,9 +413,33 @@ impl Session { return; } - let mcp_config = self.runtime_mcp_config(&refresh_config).await; - self.refresh_mcp_servers_inner(turn_context, mcp_config, mcp_servers, elicitation_reviewer) + let _guard = self.services.mcp_projection_lock.lock().await; + let available_environment_ids = self + .services + .latest_mcp_runtime() + .available_environment_ids() + .to_vec(); + let mut mcp_config = self + .services + .mcp_manager + .runtime_config_for_step( + &refresh_config, + &self.services.mcp_thread_init, + &self.services.thread_extension_data, + &available_environment_ids, + ) .await; + mcp_config.mcp_server_catalog = mcp_config + .mcp_server_catalog + .with_materialized_servers(mcp_servers); + self.refresh_mcp_servers_inner( + turn_context, + mcp_config, + &turn_context.environments, + &available_environment_ids, + elicitation_reviewer, + ) + .await; } pub(crate) async fn set_openai_form_elicitation_support( @@ -398,16 +470,54 @@ impl Session { Ok(()) } + #[expect( + clippy::await_holding_invalid_type, + reason = "MCP runtime refresh and publication must remain serialized" + )] pub(crate) async fn refresh_mcp_servers_now( &self, turn_context: &TurnContext, refresh_config: &Config, elicitation_reviewer: Option, ) { - let mcp_config = self.runtime_mcp_config(refresh_config).await; - let mcp_servers = codex_mcp::configured_mcp_servers(&mcp_config); - self.refresh_mcp_servers_inner(turn_context, mcp_config, mcp_servers, elicitation_reviewer) + let _guard = self.services.mcp_projection_lock.lock().await; + let available_environment_ids = self + .services + .latest_mcp_runtime() + .available_environment_ids() + .to_vec(); + let mcp_config = self + .services + .mcp_manager + .runtime_config_for_step( + refresh_config, + &self.services.mcp_thread_init, + &self.services.thread_extension_data, + &available_environment_ids, + ) .await; + self.refresh_mcp_servers_inner( + turn_context, + mcp_config, + &turn_context.environments, + &available_environment_ids, + elicitation_reviewer, + ) + .await; + } + + fn available_selected_environment_ids( + selected_capability_roots: &[ResolvedSelectedCapabilityRoot], + ) -> Vec { + let mut available = Vec::new(); + for root in selected_capability_roots { + let CapabilityRootLocation::Environment { environment_id, .. } = + &root.selected_root().location; + if !available.contains(environment_id) { + available.push(environment_id.clone()); + } + } + available } #[cfg(test)] diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index dc5ea011ca..34cf07cc35 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -10,6 +10,7 @@ pub struct McpRuntimeSnapshot { config: Arc, manager: Arc, runtime_context: McpRuntimeContext, + available_environment_ids: Vec, } impl McpRuntimeSnapshot { @@ -17,11 +18,13 @@ impl McpRuntimeSnapshot { config: Arc, manager: Arc, runtime_context: McpRuntimeContext, + available_environment_ids: Vec, ) -> Self { Self { config, manager, runtime_context, + available_environment_ids, } } @@ -41,6 +44,10 @@ impl McpRuntimeSnapshot { &self.runtime_context } + pub(crate) fn available_environment_ids(&self) -> &[String] { + &self.available_environment_ids + } + #[cfg(test)] pub(crate) fn new_uninitialized_for_test(config: &crate::config::Config) -> Arc { use codex_exec_server::EnvironmentManager; @@ -81,6 +88,7 @@ impl McpRuntimeSnapshot { Arc::new(mcp_config), Arc::new(manager), runtime_context, + Vec::new(), )) } } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 9421e07a66..0008257542 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -100,7 +100,6 @@ use codex_protocol::dynamic_tools::DynamicToolResponse; use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::items::TurnItem; use codex_protocol::items::UserMessageItem; -use codex_protocol::mcp::CallToolResult; use codex_protocol::models::ActivePermissionProfile; use codex_protocol::models::AdditionalPermissionProfile; use codex_protocol::models::BaseInstructions; @@ -158,8 +157,6 @@ use codex_utils_path_uri::PathUri; use futures::future::BoxFuture; use futures::future::Shared; use futures::prelude::*; -use rmcp::model::ReadResourceRequestParams; -use rmcp::model::ReadResourceResult; use rmcp::model::RequestId; use serde_json::Value; use tokio::sync::Mutex; @@ -2825,7 +2822,7 @@ impl Session { /// `run_turn` and pass the result down; standalone request or history boundaries may capture /// their own step. pub(crate) async fn capture_step_context( - &self, + self: &Arc, turn_context: Arc, ) -> Arc { let deferred_executor_enabled = turn_context @@ -2846,15 +2843,15 @@ impl Session { } let loaded_agents_md = self.services.agents_md_manager.get_loaded().await; let selected_capability_roots = self - .services - .turn_environments - .environment_manager() - .resolve_selected_capability_roots( - &self.services.selected_capability_roots, - &environments.captured_environments(), + .resolve_selected_capability_roots_for_step(&environments) + .await; + let mcp = self + .mcp_runtime_for_step( + turn_context.as_ref(), + &environments, + &selected_capability_roots, ) .await; - let mcp = self.services.latest_mcp_runtime(); Arc::new(StepContext::new( turn_context, environments, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 2c9e57036d..0e1609207a 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -661,6 +661,7 @@ impl Session { let config_for_mcp = Arc::clone(&config); let mcp_manager_for_mcp = Arc::clone(&mcp_manager); let mcp_thread_init_for_startup = &mcp_thread_init; + let thread_extension_data_for_mcp = &thread_extension_data; let mcp_runtime_cwd = session_configuration .environment_selections() .first() @@ -673,7 +674,12 @@ impl Session { let auth_and_mcp_fut = async move { let auth = auth_manager_clone.auth().await; let mcp_config = mcp_manager_for_mcp - .runtime_config_for_thread(&config_for_mcp, mcp_thread_init_for_startup) + .runtime_config_for_step( + &config_for_mcp, + mcp_thread_init_for_startup, + thread_extension_data_for_mcp, + /*available_environment_ids*/ &[], + ) .await; let mcp_servers = codex_mcp::effective_mcp_servers(&mcp_config, auth.as_ref()); let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(&mcp_config); @@ -1039,6 +1045,7 @@ impl Session { // setup is straightforward enough and performs well. mcp_connection_manager, mcp_runtime: arc_swap::ArcSwapOption::empty(), + mcp_projection_lock: Mutex::new(()), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -1210,6 +1217,7 @@ impl Session { .install_mcp_connection_manager( Arc::new(mcp_config), mcp_runtime_context, + /*available_environment_ids*/ Vec::new(), mcp_connection_manager, ) .await?; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index d0aac04249..57f1ce9081 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5372,6 +5372,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { let services = SessionServices { mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from(mcp_runtime.manager_arc())), mcp_runtime: arc_swap::ArcSwapOption::from(Some(mcp_runtime)), + mcp_projection_lock: Mutex::new(()), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -7446,6 +7447,7 @@ where let services = SessionServices { mcp_connection_manager: Arc::new(arc_swap::ArcSwap::from(mcp_runtime.manager_arc())), mcp_runtime: arc_swap::ArcSwapOption::from(Some(mcp_runtime)), + mcp_projection_lock: Mutex::new(()), mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, @@ -7603,6 +7605,7 @@ pub(crate) async fn make_session_and_context_with_rx() -> ( #[tokio::test] async fn refresh_mcp_servers_keeps_the_previous_runtime_alive() { let (session, turn_context) = make_session_and_context().await; + let session = Arc::new(session); let turn_context = Arc::new(turn_context); let old_runtime = session.services.latest_mcp_runtime(); let step_context = session diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 2cb1f2cada..7b509ae42e 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -171,8 +171,13 @@ pub(crate) async fn run_turn( .record_context_updates_and_set_reference_context_item(first_step_context.as_ref()) .await; - let Some((injection_items, explicitly_enabled_connectors)) = - build_skills_and_plugins(&sess, turn_context.as_ref(), &input, &cancellation_token).await + let Some((injection_items, explicitly_enabled_connectors)) = build_skills_and_plugins( + &sess, + first_step_context.as_ref(), + &input, + &cancellation_token, + ) + .await else { return Ok(None); }; @@ -504,10 +509,11 @@ async fn run_hooks_and_record_inputs( #[instrument(level = "trace", skip_all)] async fn build_skills_and_plugins( sess: &Arc, - turn_context: &TurnContext, + step_context: &StepContext, input: &[TurnInput], cancellation_token: &CancellationToken, ) -> Option<(Vec, HashSet)> { + let turn_context = step_context.turn.as_ref(); // Guardian input embeds the parent transcript as untrusted evidence. Do not interpret skill or // plugin mentions from that generated prompt as requests to inject additional instructions. if crate::guardian::is_guardian_reviewer_source(&turn_context.session_source) { @@ -538,17 +544,14 @@ async fn build_skills_and_plugins( // enabled plugins, then converted into turn-scoped guidance below. let mentioned_plugins = collect_explicit_plugin_mentions(&user_input, loaded_plugins.capability_summaries()); - let connector_snapshot = codex_connectors::ConnectorSnapshot::from_plugin_capability_summaries( - loaded_plugins.capability_summaries(), - ); + let connector_snapshot = step_context.mcp.config().connector_snapshot.clone(); 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 - .load_full() + match step_context + .mcp + .manager_arc() .list_all_tools() .or_cancel(cancellation_token) .await @@ -1189,9 +1192,7 @@ pub(crate) async fn built_tools( .plugins_for_config(&turn_context.config.plugins_config_input()) .instrument(trace_span!("built_tools.load_plugins")) .await; - let connector_snapshot = codex_connectors::ConnectorSnapshot::from_plugin_capability_summaries( - loaded_plugins.capability_summaries(), - ); + let connector_snapshot = step_context.mcp.config().connector_snapshot.clone(); let apps_enabled = turn_context.apps_enabled(); let accessible_connectors = diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 1bee45cce0..a611b94887 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -51,6 +51,8 @@ pub(crate) struct SessionServices { pub(crate) mcp_connection_manager: Arc>, /// The latest atomically published MCP config and manager pair. pub(crate) mcp_runtime: ArcSwapOption, + /// Serializes environment-driven runtime rebuilds. + pub(crate) mcp_projection_lock: Mutex<()>, pub(crate) mcp_startup_cancellation_token: Mutex, pub(crate) unified_exec_manager: UnifiedExecProcessManager, #[cfg_attr(not(unix), allow(dead_code))] @@ -106,9 +108,11 @@ impl SessionServices { &self, config: Arc, runtime_context: McpRuntimeContext, + available_environment_ids: Vec, manager: McpConnectionManager, ) -> Result<()> { - let runtime = self.publish_mcp_runtime(config, runtime_context, manager); + let runtime = + self.publish_mcp_runtime(config, runtime_context, available_environment_ids, manager); runtime.manager().validate_required_servers().await } @@ -116,13 +120,19 @@ impl SessionServices { &self, config: Arc, runtime_context: McpRuntimeContext, + available_environment_ids: Vec, manager: McpConnectionManager, ) -> Arc { let manager = Arc::new(manager); // Publish the manager for legacy resource clients first. Once the paired snapshot is // visible, every model-scoped consumer observes this exact manager. self.mcp_connection_manager.store(Arc::clone(&manager)); - let runtime = Arc::new(McpRuntimeSnapshot::new(config, manager, runtime_context)); + let runtime = Arc::new(McpRuntimeSnapshot::new( + config, + manager, + runtime_context, + available_environment_ids, + )); self.mcp_runtime.store(Some(Arc::clone(&runtime))); runtime } diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 9bccf8b6c1..5148f16259 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -584,8 +584,28 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() }) .await .expect("start second thread"); - let first_resolved = first_thread.thread.runtime_mcp_config(&config).await; - let second_resolved = second_thread.thread.runtime_mcp_config(&config).await; + let first_session = &first_thread.thread.codex.session; + let first_resolved = first_session + .services + .mcp_manager + .runtime_config_for_step( + &config, + &first_session.services.mcp_thread_init, + &first_session.services.thread_extension_data, + /*available_environment_ids*/ &[], + ) + .await; + let second_session = &second_thread.thread.codex.session; + let second_resolved = second_session + .services + .mcp_manager + .runtime_config_for_step( + &config, + &second_session.services.mcp_thread_init, + &second_session.services.thread_extension_data, + /*available_environment_ids*/ &[], + ) + .await; assert_eq!( *lifecycle_observed diff --git a/codex-rs/ext/extension-api/src/contributors/mcp.rs b/codex-rs/ext/extension-api/src/contributors/mcp.rs index 91e9cb459c..499657f556 100644 --- a/codex-rs/ext/extension-api/src/contributors/mcp.rs +++ b/codex-rs/ext/extension-api/src/contributors/mcp.rs @@ -1,17 +1,22 @@ use codex_config::McpServerConfig; +use crate::ExtensionData; use crate::ExtensionDataInit; /// Input supplied while resolving MCP server contributions. /// -/// Thread-scoped implementations can read the immutable host-seeded inputs -/// through [`Self::thread_init`]. Implementations should not retain borrowed -/// context after contribution completes. +/// Thread-scoped implementations can read stable host inputs through [`Self::thread_init`] and +/// keep their cache in [`Self::thread_store`]. Implementations should not retain borrowed context +/// after contribution completes. pub struct McpServerContributionContext<'a, C> { /// Host configuration visible during MCP resolution. config: &'a C, - /// Initial inputs for the active thread, when resolution is thread-scoped. + /// Extension-owned data for the active thread, when resolution is thread-scoped. + thread_store: Option<&'a ExtensionData>, + /// Stable host inputs for the active thread, when resolution is thread-scoped. thread_init: Option<&'a ExtensionDataInit>, + /// Environment IDs whose selected roots may contribute to this exact step. + available_environment_ids: Option<&'a [String]>, } impl Clone for McpServerContributionContext<'_, C> { @@ -27,15 +32,24 @@ impl<'a, C> McpServerContributionContext<'a, C> { pub fn global(config: &'a C) -> Self { Self { config, + thread_store: None, thread_init: None, + available_environment_ids: None, } } - /// Creates context for one active thread runtime. - pub fn for_thread(config: &'a C, thread_init: &'a ExtensionDataInit) -> Self { + /// Creates context for one model step using only currently available environments. + pub fn for_step( + config: &'a C, + thread_init: &'a ExtensionDataInit, + thread_store: &'a ExtensionData, + available_environment_ids: &'a [String], + ) -> Self { Self { config, + thread_store: Some(thread_store), thread_init: Some(thread_init), + available_environment_ids: Some(available_environment_ids), } } @@ -44,10 +58,23 @@ impl<'a, C> McpServerContributionContext<'a, C> { self.config } - /// Returns the frozen initial inputs when resolving for a running thread. + /// Returns extension-owned state when resolving for a running thread. + pub fn thread_store(&self) -> Option<&'a ExtensionData> { + self.thread_store + } + + /// Returns stable host inputs when resolving for a running thread. pub fn thread_init(&self) -> Option<&'a ExtensionDataInit> { self.thread_init } + + /// Returns the exact environment availability projection for a model step. + /// + /// `Some` means contributors must omit selected roots whose environment ID is absent from the + /// slice. Global resolution returns `None` because it has no thread environments. + pub fn available_environment_ids(&self) -> Option<&'a [String]> { + self.available_environment_ids + } } /// One extension-owned overlay for the runtime MCP server configuration. @@ -66,6 +93,12 @@ pub enum McpServerContribution { selection_order: usize, config: Box, }, + /// Adds connector IDs declared by a plugin selected for this thread. + SelectedPluginConnectors { + plugin_id: String, + plugin_display_name: String, + connector_ids: Vec, + }, /// Removes a named MCP server. Remove { name: String }, } diff --git a/codex-rs/ext/mcp/Cargo.toml b/codex-rs/ext/mcp/Cargo.toml index b98c0905db..d918bd5bd9 100644 --- a/codex-rs/ext/mcp/Cargo.toml +++ b/codex-rs/ext/mcp/Cargo.toml @@ -16,6 +16,7 @@ workspace = true codex-core = { workspace = true } codex-core-plugins = { workspace = true } codex-config = { workspace = true } +codex-connectors-extension = { workspace = true } codex-exec-server = { workspace = true } codex-extension-api = { workspace = true } codex-features = { workspace = true } diff --git a/codex-rs/ext/mcp/src/executor_plugin.rs b/codex-rs/ext/mcp/src/executor_plugin.rs index 69d1a47b31..32637b4f47 100644 --- a/codex-rs/ext/mcp/src/executor_plugin.rs +++ b/codex-rs/ext/mcp/src/executor_plugin.rs @@ -1,44 +1,47 @@ +use codex_connectors_extension::ExecutorPluginConnectorProvider; use codex_core::config::Config; use codex_core_plugins::ExecutorPluginProvider; use codex_exec_server::EnvironmentManager; -use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionFuture; use codex_extension_api::McpServerContribution; use codex_extension_api::McpServerContributionContext; use codex_extension_api::McpServerContributor; +use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; use std::collections::HashMap; use std::sync::Arc; -use tokio::sync::OnceCell; +use std::sync::Mutex; use self::provider::ExecutorPluginMcpProvider; mod provider; -/// Frozen MCP declarations for one selected package. +/// Frozen MCP and connector declarations for one selected package. /// /// Each server config retains the stable logical environment ID. Reconnection may replace the /// concrete environment instance without changing that authority. #[derive(Clone)] -struct SelectedPluginMcpServers { +struct SelectedPluginMetadata { plugin_id: String, plugin_display_name: String, - selection_order: usize, servers: Vec<(String, codex_config::McpServerConfig)>, + connector_ids: Vec, } #[derive(Default)] pub(crate) struct SelectedExecutorPluginMcpState { - snapshot: OnceCell>, + cache: Mutex>, } -pub(crate) fn seed_thread_state(thread_init: &mut ExtensionDataInit) { - thread_init.insert(SelectedExecutorPluginMcpState::default()); +struct CachedSelectedRoot { + root: SelectedCapabilityRoot, + metadata: Option, } pub(crate) struct SelectedExecutorPluginMcpContributor { plugin_provider: ExecutorPluginProvider, mcp_provider: ExecutorPluginMcpProvider, + connector_provider: ExecutorPluginConnectorProvider, } impl SelectedExecutorPluginMcpContributor { @@ -46,46 +49,87 @@ impl SelectedExecutorPluginMcpContributor { Self { plugin_provider: ExecutorPluginProvider::new(Arc::clone(&environment_manager)), mcp_provider: ExecutorPluginMcpProvider, + connector_provider: ExecutorPluginConnectorProvider, } } - async fn resolve_snapshot( + /// Returns metadata for one stable selected root. + /// + /// Successful resolution, including a root that is not a plugin or declares no capabilities, + /// is cached until the thread state is dropped. Environment availability never invalidates + /// this cache; it only controls whether the cached metadata is projected into a model step. + async fn metadata_for_root( &self, - selected_roots: &[SelectedCapabilityRoot], - ) -> Vec { - let mut snapshot = Vec::new(); + state: &SelectedExecutorPluginMcpState, + selected_root: &SelectedCapabilityRoot, + ) -> Option { + if let Some(cached) = state + .cache + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .iter() + .find(|cached| cached.root == *selected_root) + { + return cached.metadata.clone(); + } - for (selection_order, selected_root) in selected_roots.iter().enumerate() { - let plugin = match self.plugin_provider.resolve_bound(selected_root).await { - Ok(Some(plugin)) => plugin, - Ok(None) => continue, - Err(err) => { - tracing::warn!( - selected_root = selected_root.id, - error = %err, - "failed to resolve selected executor plugin for MCP discovery" - ); - continue; - } - }; - match self.mcp_provider.load(&plugin).await { - Ok(servers) => snapshot.push(SelectedPluginMcpServers { - plugin_id: plugin.plugin().selected_root_id().to_string(), - plugin_display_name: plugin.plugin().manifest().display_name().to_string(), - selection_order, - servers, - }), - Err(err) => { + let plugin = match self.plugin_provider.resolve_bound(selected_root).await { + Ok(plugin) => plugin, + Err(err) => { + tracing::warn!( + selected_root = selected_root.id, + error = %err, + "failed to resolve selected executor plugin" + ); + return None; + } + }; + let metadata = match plugin { + Some(plugin) => { + let servers = self.mcp_provider.load(&plugin).await.unwrap_or_else(|err| { tracing::warn!( selected_root = selected_root.id, error = %err, "failed to load selected executor plugin MCP servers" ); - } + Vec::new() + }); + let connector_ids = self + .connector_provider + .load(&plugin) + .await + .unwrap_or_else(|err| { + tracing::warn!( + selected_root = selected_root.id, + error = %err, + "failed to load selected executor plugin connectors" + ); + Vec::new() + }) + .into_iter() + .map(|declaration| declaration.connector_id.0) + .collect(); + Some(SelectedPluginMetadata { + plugin_id: plugin.plugin().selected_root_id().to_string(), + plugin_display_name: plugin.plugin().manifest().display_name().to_string(), + servers, + connector_ids, + }) } + None => None, + }; + let mut cache = state + .cache + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(cached) = cache.iter().find(|cached| cached.root == *selected_root) { + return cached.metadata.clone(); } - - snapshot + cache.push(CachedSelectedRoot { + root: selected_root.clone(), + metadata: metadata.clone(), + }); + metadata } } @@ -102,20 +146,31 @@ impl McpServerContributor for SelectedExecutorPluginMcpContributor { let Some(thread_init) = context.thread_init() else { return Vec::new(); }; + let Some(thread_store) = context.thread_store() else { + return Vec::new(); + }; let Some(selected_roots) = thread_init.get::>() else { return Vec::new(); }; - let Some(state) = thread_init.get::() else { - tracing::warn!("selected executor plugin MCP state was not initialized"); - return Vec::new(); - }; - let snapshot = state - .snapshot - .get_or_init(|| self.resolve_snapshot(selected_roots.as_ref())) - .await; + let state = thread_store.get_or_init(SelectedExecutorPluginMcpState::default); let mut contributions = Vec::new(); - for plugin in snapshot { + for (selection_order, selected_root) in selected_roots.iter().enumerate() { + let CapabilityRootLocation::Environment { environment_id, .. } = + &selected_root.location; + if context + .available_environment_ids() + .is_some_and(|available| { + !available + .iter() + .any(|available| available == environment_id) + }) + { + continue; + } + let Some(plugin) = self.metadata_for_root(&state, selected_root).await else { + continue; + }; let mut servers = plugin.servers.iter().cloned().collect::>(); context .config() @@ -127,10 +182,17 @@ impl McpServerContributor for SelectedExecutorPluginMcpContributor { name, plugin_id: plugin.plugin_id.clone(), plugin_display_name: plugin.plugin_display_name.clone(), - selection_order: plugin.selection_order, + selection_order, config: Box::new(config), } })); + if !plugin.connector_ids.is_empty() { + contributions.push(McpServerContribution::SelectedPluginConnectors { + plugin_id: plugin.plugin_id, + plugin_display_name: plugin.plugin_display_name, + connector_ids: plugin.connector_ids, + }); + } } contributions diff --git a/codex-rs/ext/mcp/src/lib.rs b/codex-rs/ext/mcp/src/lib.rs index e32316ac36..9c22bad5a0 100644 --- a/codex-rs/ext/mcp/src/lib.rs +++ b/codex-rs/ext/mcp/src/lib.rs @@ -51,10 +51,3 @@ pub fn install_executor_plugins( executor_plugin::SelectedExecutorPluginMcpContributor::new(environment_manager), )); } - -/// Seeds the per-thread snapshot used by selected executor plugin MCP discovery. -pub fn initialize_executor_plugin_thread_data( - thread_init: &mut codex_extension_api::ExtensionDataInit, -) { - executor_plugin::seed_thread_state(thread_init); -} diff --git a/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs b/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs index 7bdbc83d62..65a4248a32 100644 --- a/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs +++ b/codex-rs/ext/mcp/tests/executor_plugin_mcp.rs @@ -3,6 +3,7 @@ use codex_core::config::Config; use codex_core::config::ConfigBuilder; use codex_exec_server::EnvironmentManager; use codex_exec_server::LOCAL_ENVIRONMENT_ID; +use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::McpServerContribution; @@ -109,12 +110,15 @@ async fn selected_plugin_contributions( path: PathUri::from_host_native_path(plugin_root)?, }, }]); - codex_mcp_extension::initialize_executor_plugin_thread_data(&mut thread_init); + let thread_store = ExtensionData::new_with_init("test-thread", thread_init.clone()); + let available_environment_ids = vec![LOCAL_ENVIRONMENT_ID.to_string()]; Ok(registry.mcp_server_contributors()[0] - .contribute(McpServerContributionContext::for_thread( + .contribute(McpServerContributionContext::for_step( config, &thread_init, + &thread_store, + &available_environment_ids, )) .await .into_iter() @@ -132,7 +136,9 @@ async fn selected_plugin_contributions( selection_order, enabled: config.enabled, }, - McpServerContribution::Set { .. } | McpServerContribution::Remove { .. } => { + McpServerContribution::Set { .. } + | McpServerContribution::SelectedPluginConnectors { .. } + | McpServerContribution::Remove { .. } => { panic!("expected selected plugin contribution") } }) From 8ebf71ec257fcb95d72a71c7d3bedbbe9a861e45 Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 26 Jun 2026 01:47:00 +0100 Subject: [PATCH 06/38] Reuse walk inventory for environment skill metadata (#30145) ## Why Environment skill discovery already asks the executor to run one `fs/walk`. That response contains every regular file path found under the selected root, including any `agents/openai.yaml` files. Today Core keeps the discovered `SKILL.md` paths but discards the rest of that file inventory. It then sends one `fs/getMetadata` request per skill just to ask whether `agents/openai.yaml` exists. A root with 66 skills and no metadata therefore pays for 66 unnecessary network round trips. ## What changes - Keep the `fs/walk` file and directory inventory for the duration of the scan. - Associate each discovered `SKILL.md` with metadata that is known present, known absent, or still requires a fallback probe. - Read a known `agents/openai.yaml` directly instead of statting it first. - Skip the metadata request entirely when a complete walk shows that the skill has no `agents` directory. - Read a known `SKILL.md` and `agents/openai.yaml` concurrently. - Keep parsing and validation in `core-skills`. The inventory is scan-local. This does not add another cache or change cache lifetime. ## Network impact For a complete scan of 66 valid skills with no `agents/openai.yaml`, and one root `.codex-plugin/plugin.json`: | Operation | Current | After this PR | | --- | ---: | ---: | | `fs/walk` | 1 | 1 | | Read `SKILL.md` | 66 | 66 | | Stat `agents/openai.yaml` | 66 | 0 | | Read `agents/openai.yaml` | 0 | 0 | | Stat plugin manifest | 1 | 1 | | Read plugin manifest | 1 | 1 | | **Total executor RPCs** | **135** | **69** | This removes exactly 66 request/response exchanges from the common cold scan. Warm scans remain at zero discovery RPCs because the thread-level executor catalog cache is unchanged. When metadata exists, each file still requires one read. This PR removes only the preceding existence check; it does not batch file contents into a new RPC. ## Correctness fallbacks Absence is trusted only when the walk is complete and the metadata directory was not present. Core keeps the existing `getMetadata` fallback when: - the walk was truncated; - the walk reported an error; or - an `agents` directory was observed but `openai.yaml` was not, which preserves support for file symlinks and traversal boundaries. ## Deliberate scope This PR changes only the environment skill loader and its existing filesystem-call regression coverage. It does not: - change `fs/walk` or any exec-server protocol; - add `readFiles` or a skills-list endpoint; - change thread caching; - change local skill discovery; - change exec-server request concurrency; or - optimize plugin-manifest lookup. The plugin-manifest stat is intentionally left in place, which is why this PR reaches 69 calls rather than the broader 68-call estimate. That lookup has separate alternate-path, ancestor, and symlink semantics and should not be mixed into this change. --- .../core-skills/src/loader/environment.rs | 152 ++++++++++++++---- .../core-skills/tests/environment_loader.rs | 121 +++++++++++++- 2 files changed, 236 insertions(+), 37 deletions(-) diff --git a/codex-rs/core-skills/src/loader/environment.rs b/codex-rs/core-skills/src/loader/environment.rs index 93ab14e9aa..96caa28eea 100644 --- a/codex-rs/core-skills/src/loader/environment.rs +++ b/codex-rs/core-skills/src/loader/environment.rs @@ -23,7 +23,6 @@ use super::ParsedSkillFrontmatter; use super::SKILLS_FILENAME; use super::SKILLS_METADATA_DIR; use super::SKILLS_METADATA_FILENAME; -use super::SkillFileDiscovery; use super::SkillMetadataFile; use super::parse_skill_frontmatter_metadata_inner; use super::resolve_dependencies; @@ -34,6 +33,24 @@ use super::validate_len; const MAX_SKILLS_ENTRIES_PER_ROOT: usize = 20_000; const MAX_CONCURRENT_SKILL_LOADS: usize = 64; +struct EnvironmentSkillDiscovery { + skills: Vec, + plugin_roots: HashSet, + namespace_roots: HashSet, + warnings: Vec, +} + +struct DiscoveredEnvironmentSkill { + path: PathUri, + metadata: SkillMetadataDiscovery, +} + +enum SkillMetadataDiscovery { + Present(PathUri), + Absent, + Probe(PathUri), +} + /// URI-native metadata for one skill owned by an execution environment. #[derive(Clone, Debug, PartialEq, Eq)] pub struct EnvironmentSkillMetadata { @@ -67,28 +84,44 @@ impl EnvironmentSkillMetadata { async fn parse( file_system: &dyn ExecutorFileSystem, - path: &PathUri, + skill: &DiscoveredEnvironmentSkill, plugin_namespace: Option<&str>, ) -> Result { - let contents = file_system - .read_file_text(path, /*sandbox*/ None) - .await - .map_err(|err| format!("failed to read file: {err}"))?; + let (contents, discovered_metadata) = match &skill.metadata { + SkillMetadataDiscovery::Present(metadata_path) => { + let (contents, metadata) = tokio::join!( + read_skill_contents(file_system, &skill.path), + read_skill_metadata(file_system, metadata_path), + ); + (contents?, metadata) + } + SkillMetadataDiscovery::Absent | SkillMetadataDiscovery::Probe(_) => ( + read_skill_contents(file_system, &skill.path).await?, + (None, None), + ), + }; let ParsedSkillFrontmatter { name: base_name, description, short_description, - } = parse_skill_frontmatter_metadata_inner(&contents, || default_skill_name(path)) + } = parse_skill_frontmatter_metadata_inner(&contents, || default_skill_name(&skill.path)) .map_err(|err| err.to_string())?; let name = plugin_namespace .map(|namespace| format!("{namespace}:{base_name}")) .unwrap_or(base_name); validate_len(&name, MAX_QUALIFIED_NAME_LEN, "qualified name") .map_err(|err| err.to_string())?; - let (dependencies, policy) = load_skill_metadata(file_system, path).await; + let (dependencies, policy) = match &skill.metadata { + SkillMetadataDiscovery::Present(_) | SkillMetadataDiscovery::Absent => { + discovered_metadata + } + SkillMetadataDiscovery::Probe(metadata_path) => { + probe_skill_metadata(file_system, metadata_path).await + } + }; Ok(Self { - path_to_skills_md: path.clone(), + path_to_skills_md: skill.path.clone(), name, description, short_description, @@ -125,6 +158,7 @@ pub async fn load_environment_skills_from_root( .await { Ok(walk) => { + let inventory_complete = !walk.truncated && walk.errors.is_empty(); let mut warnings = walk .errors .into_iter() @@ -141,10 +175,13 @@ pub async fn load_environment_skills_from_root( )); } let mut skill_files = Vec::new(); + let mut file_paths = HashSet::new(); + let mut directory_paths = HashSet::new(); let mut plugin_roots = HashSet::new(); for entry in walk.entries { match entry.kind { WalkEntryKind::Directory => { + directory_paths.insert(entry.path.clone()); if DISCOVERABLE_PLUGIN_MANIFEST_PATHS .iter() .any(|path| path.split('/').next() == entry.path.basename().as_deref()) @@ -154,40 +191,53 @@ pub async fn load_environment_skills_from_root( } } WalkEntryKind::File => { + file_paths.insert(entry.path.clone()); if entry.path.basename().as_deref() == Some(SKILLS_FILENAME) { skill_files.push(entry.path); } } } } - SkillFileDiscovery { - skill_files, + let skills = skill_files + .into_iter() + .map(|path| DiscoveredEnvironmentSkill { + metadata: discover_skill_metadata( + &path, + &file_paths, + &directory_paths, + inventory_complete, + ), + path, + }) + .collect(); + EnvironmentSkillDiscovery { + skills, plugin_roots, namespace_roots: HashSet::from([root.clone()]), warnings, } } - Err(error) if error.kind() == io::ErrorKind::NotFound => SkillFileDiscovery { - skill_files: Vec::new(), + Err(error) if error.kind() == io::ErrorKind::NotFound => EnvironmentSkillDiscovery { + skills: Vec::new(), plugin_roots: HashSet::new(), namespace_roots: HashSet::new(), warnings: Vec::new(), }, - Err(error) => SkillFileDiscovery { - skill_files: Vec::new(), + Err(error) => EnvironmentSkillDiscovery { + skills: Vec::new(), plugin_roots: HashSet::new(), namespace_roots: HashSet::new(), warnings: vec![format!("failed to walk skills root {root}: {error:#}")], }, }; outcome.warnings.extend(discovery.warnings); - if discovery.skill_files.is_empty() { + if discovery.skills.is_empty() { return outcome; } let mut skill_ancestors = HashSet::new(); - for skill_path in &discovery.skill_files { - let mut ancestor = skill_path.parent(); + for skill in &discovery.skills { + let mut ancestor = skill.path.parent(); while let Some(path) = ancestor { skill_ancestors.insert(path.clone()); ancestor = path.parent(); @@ -224,9 +274,9 @@ pub async fn load_environment_skills_from_root( // Remote executors can multiplex these independent per-skill reads, so polling a bounded // number together allows the I/O for each skill and its metadata to happen concurrently. - let skill_results = futures::stream::iter(discovery.skill_files) - .map(|path| { - let mut ancestor = path.parent(); + let skill_results = futures::stream::iter(discovery.skills) + .map(|skill| { + let mut ancestor = skill.path.parent(); let plugin_namespace = loop { let Some(current) = ancestor else { break None; @@ -236,9 +286,10 @@ pub async fn load_environment_skills_from_root( } ancestor = current.parent(); }; + let path = skill.path.clone(); async move { let result = - EnvironmentSkillMetadata::parse(file_system, &path, plugin_namespace).await; + EnvironmentSkillMetadata::parse(file_system, &skill, plugin_namespace).await; (path, result) } }) @@ -267,20 +318,48 @@ pub async fn load_environment_skills_from_root( outcome } -async fn load_skill_metadata( +fn discover_skill_metadata( + skill_path: &PathUri, + file_paths: &HashSet, + directory_paths: &HashSet, + inventory_complete: bool, +) -> SkillMetadataDiscovery { + let Some(skill_dir) = skill_path.parent() else { + return SkillMetadataDiscovery::Absent; + }; + let Ok(metadata_dir) = skill_dir.join(SKILLS_METADATA_DIR) else { + return SkillMetadataDiscovery::Absent; + }; + let Ok(metadata_path) = metadata_dir.join(SKILLS_METADATA_FILENAME) else { + return SkillMetadataDiscovery::Absent; + }; + if file_paths.contains(&metadata_path) { + SkillMetadataDiscovery::Present(metadata_path) + } else if inventory_complete && !directory_paths.contains(&metadata_dir) { + SkillMetadataDiscovery::Absent + } else { + // The walk can omit entries after an error or traversal limit. It also omits file + // symlinks, so keep the existing probe when the metadata directory itself was observed. + SkillMetadataDiscovery::Probe(metadata_path) + } +} + +async fn read_skill_contents( file_system: &dyn ExecutorFileSystem, skill_path: &PathUri, +) -> Result { + file_system + .read_file_text(skill_path, /*sandbox*/ None) + .await + .map_err(|err| format!("failed to read file: {err}")) +} + +async fn probe_skill_metadata( + file_system: &dyn ExecutorFileSystem, + metadata_path: &PathUri, ) -> (Option, Option) { - let Some(skill_dir) = skill_path.parent() else { - return (None, None); - }; - let Ok(metadata_path) = - skill_dir.join(&format!("{SKILLS_METADATA_DIR}/{SKILLS_METADATA_FILENAME}")) - else { - return (None, None); - }; match file_system - .get_metadata(&metadata_path, /*sandbox*/ None) + .get_metadata(metadata_path, /*sandbox*/ None) .await { Ok(metadata) if metadata.is_file => {} @@ -291,8 +370,15 @@ async fn load_skill_metadata( return (None, None); } } + read_skill_metadata(file_system, metadata_path).await +} + +async fn read_skill_metadata( + file_system: &dyn ExecutorFileSystem, + metadata_path: &PathUri, +) -> (Option, Option) { let contents = match file_system - .read_file_text(&metadata_path, /*sandbox*/ None) + .read_file_text(metadata_path, /*sandbox*/ None) .await { Ok(contents) => contents, diff --git a/codex-rs/core-skills/tests/environment_loader.rs b/codex-rs/core-skills/tests/environment_loader.rs index 42d2446083..c17802bedb 100644 --- a/codex-rs/core-skills/tests/environment_loader.rs +++ b/codex-rs/core-skills/tests/environment_loader.rs @@ -1,5 +1,7 @@ use std::fs; use std::sync::Mutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use codex_core_skills::loader::EnvironmentSkillMetadata; use codex_core_skills::loader::load_environment_skills_from_root; @@ -13,6 +15,8 @@ use codex_exec_server::FileSystemSandboxContext; use codex_exec_server::LOCAL_FS; use codex_exec_server::ReadDirectoryEntry; use codex_exec_server::RemoveOptions; +use codex_exec_server::WalkOptions; +use codex_exec_server::WalkOutcome; use codex_utils_path_uri::PathUri; use pretty_assertions::assert_eq; use tempfile::tempdir; @@ -20,6 +24,15 @@ use tempfile::tempdir; struct RecordingFileSystem<'a> { inner: &'a dyn ExecutorFileSystem, read_files: Mutex>, + metadata_files: Mutex>, + walks: AtomicUsize, +} + +#[derive(Debug, PartialEq, Eq)] +struct FileSystemCalls { + walks: usize, + read_files: Vec, + metadata_files: Vec, } impl<'a> RecordingFileSystem<'a> { @@ -27,14 +40,29 @@ impl<'a> RecordingFileSystem<'a> { Self { inner, read_files: Mutex::new(Vec::new()), + metadata_files: Mutex::new(Vec::new()), + walks: AtomicUsize::new(0), } } - fn read_files(&self) -> Vec { - self.read_files + fn calls(&self) -> FileSystemCalls { + let mut read_files = self + .read_files .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) - .clone() + .clone(); + read_files.sort_by_key(ToString::to_string); + let mut metadata_files = self + .metadata_files + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + metadata_files.sort_by_key(ToString::to_string); + FileSystemCalls { + walks: self.walks.load(Ordering::Relaxed), + read_files, + metadata_files, + } } } @@ -90,6 +118,10 @@ impl ExecutorFileSystem for RecordingFileSystem<'_> { path: &'a PathUri, sandbox: Option<&'a FileSystemSandboxContext>, ) -> ExecutorFileSystemFuture<'a, FileMetadata> { + self.metadata_files + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(path.clone()); self.inner.get_metadata(path, sandbox) } @@ -101,6 +133,16 @@ impl ExecutorFileSystem for RecordingFileSystem<'_> { self.inner.read_directory(path, sandbox) } + fn walk<'a>( + &'a self, + path: &'a PathUri, + options: WalkOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, WalkOutcome> { + self.walks.fetch_add(1, Ordering::Relaxed); + self.inner.walk(path, options, sandbox) + } + fn remove<'a>( &'a self, path: &'a PathUri, @@ -200,7 +242,8 @@ async fn loads_nearest_plugin_namespaces_without_reading_unused_sibling_manifest ); let mut manifest_reads = file_system - .read_files() + .calls() + .read_files .into_iter() .filter(|path| path.basename().as_deref() == Some("plugin.json")) .collect::>(); @@ -214,3 +257,73 @@ async fn loads_nearest_plugin_namespaces_without_reading_unused_sibling_manifest expected_manifest_reads.sort_by_key(ToString::to_string); assert_eq!(manifest_reads, expected_manifest_reads); } + +#[tokio::test] +async fn reuses_walk_inventory_for_missing_skill_metadata() { + const SKILL_COUNT: usize = 66; + + let root = tempdir().expect("tempdir"); + let manifest_path = root.path().join(".codex-plugin/plugin.json"); + fs::create_dir_all(manifest_path.parent().expect("manifest parent")).expect("manifest dir"); + fs::write(&manifest_path, r#"{"name":"inventory"}"#).expect("manifest"); + + let mut skill_paths = Vec::new(); + for index in 0..SKILL_COUNT { + let name = format!("skill-{index}"); + let skill_path = root.path().join(&name).join("SKILL.md"); + fs::create_dir_all(skill_path.parent().expect("skill parent")).expect("skill dir"); + fs::write( + &skill_path, + format!("---\nname: {name}\ndescription: {name} skill.\n---\n"), + ) + .expect("skill"); + skill_paths.push(skill_path); + } + + let file_system = RecordingFileSystem::new(LOCAL_FS.as_ref()); + let root_uri = PathUri::from_host_native_path(root.path()).expect("root URI"); + let outcome = load_environment_skills_from_root( + &file_system, + &root_uri, + /*restriction_product*/ None, + ) + .await; + + let mut expected_skills = skill_paths + .iter() + .enumerate() + .map(|(index, skill_path)| EnvironmentSkillMetadata { + path_to_skills_md: PathUri::from_host_native_path(skill_path).unwrap(), + name: format!("inventory:skill-{index}"), + description: format!("skill-{index} skill."), + short_description: None, + dependencies: None, + policy: None, + }) + .collect::>(); + expected_skills.sort_by(|left, right| { + left.name.cmp(&right.name).then_with(|| { + left.path_to_skills_md + .to_string() + .cmp(&right.path_to_skills_md.to_string()) + }) + }); + assert_eq!(outcome.skills, expected_skills); + assert_eq!(outcome.warnings, Vec::::new()); + + let mut expected_read_files = skill_paths + .iter() + .map(|path| PathUri::from_host_native_path(path).unwrap()) + .collect::>(); + let manifest_uri = PathUri::from_host_native_path(manifest_path).unwrap(); + expected_read_files.push(manifest_uri.clone()); + expected_read_files.sort_by_key(ToString::to_string); + assert_eq!( + file_system.calls(), + FileSystemCalls { + walks: 1, + read_files: expected_read_files, + metadata_files: vec![manifest_uri], + } + ); +} From da78d5fdc5d8b30a80b402100d8b25ed8031f8c4 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Thu, 25 Jun 2026 18:00:39 -0700 Subject: [PATCH 07/38] [codex] implement standalone code-mode process host (#30111) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - implement the standalone `codex-code-mode-host` stdio service - route sessions, cells, delegate requests, responses, and cancellation through a bounded host peer - supervise request, writer, cell-forwarding, actor, and V8 failure boundaries - bound request/session tombstones and fail-stop the connection on invalid protocol state - add host-only duplex protocol tests and local Cargo/Bazel run recipes ## Why This stage makes the host process independently runnable and reviewable before exposing any remote client in Codex. Transport or runtime failure closes the connection and relies on process replacement rather than transactional recovery. ## Stack This is **3 of 4** in the process-owned code-mode session stack. - Depends on #30110 - The final client PR targets this branch ## Validation - `just test -p codex-code-mode-host` — 7 host-only tests passed - `just fix -p codex-code-mode-host` - `just bazel-lock-update` - `just bazel-lock-check` - `just fmt` --- codex-rs/Cargo.lock | 8 + codex-rs/code-mode-host/Cargo.toml | 15 + codex-rs/code-mode-host/src/delegate.rs | 85 +++ codex-rs/code-mode-host/src/host_tests.rs | 486 +++++++++++++++++ codex-rs/code-mode-host/src/lib.rs | 605 ++++++++++++++++++++++ codex-rs/code-mode-host/src/main.rs | 5 +- codex-rs/code-mode-host/src/peer.rs | 541 +++++++++++++++++++ codex-rs/code-mode-host/src/peer_tests.rs | 95 ++++ justfile | 14 + 9 files changed, 1853 insertions(+), 1 deletion(-) create mode 100644 codex-rs/code-mode-host/src/delegate.rs create mode 100644 codex-rs/code-mode-host/src/host_tests.rs create mode 100644 codex-rs/code-mode-host/src/lib.rs create mode 100644 codex-rs/code-mode-host/src/peer.rs create mode 100644 codex-rs/code-mode-host/src/peer_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 75c02adfb5..17c51f58a9 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2515,6 +2515,14 @@ dependencies = [ [[package]] name = "codex-code-mode-host" version = "0.0.0" +dependencies = [ + "anyhow", + "codex-code-mode", + "codex-code-mode-protocol", + "pretty_assertions", + "tokio", + "tokio-util", +] [[package]] name = "codex-code-mode-protocol" diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index a2c384d010..d37bd41012 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/codex-rs/code-mode-host/Cargo.toml @@ -8,5 +8,20 @@ license.workspace = true name = "codex-code-mode-host" path = "src/main.rs" +[lib] +doctest = false +name = "codex_code_mode_host" +path = "src/lib.rs" + [lints] workspace = true + +[dependencies] +anyhow = { workspace = true } +codex-code-mode = { workspace = true } +codex-code-mode-protocol = { workspace = true } +tokio = { workspace = true, features = ["io-std", "io-util", "macros", "rt", "sync", "time"] } +tokio-util = { workspace = true, features = ["rt"] } + +[dev-dependencies] +pretty_assertions = { workspace = true } diff --git a/codex-rs/code-mode-host/src/delegate.rs b/codex-rs/code-mode-host/src/delegate.rs new file mode 100644 index 0000000000..50d45e7de7 --- /dev/null +++ b/codex-rs/code-mode-host/src/delegate.rs @@ -0,0 +1,85 @@ +use std::sync::Arc; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeNestedToolCall; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::NotificationFuture; +use codex_code_mode_protocol::ToolInvocationFuture; +use codex_code_mode_protocol::host::DelegateRequest; +use codex_code_mode_protocol::host::DelegateResponse; +use codex_code_mode_protocol::host::SessionId; +use tokio_util::sync::CancellationToken; + +use crate::peer::HostPeer; + +pub(super) struct RemoteDelegate { + session_id: SessionId, + peer: Arc, +} + +impl RemoteDelegate { + pub(super) fn new(session_id: SessionId, peer: Arc) -> Self { + Self { session_id, peer } + } +} + +impl CodeModeSessionDelegate for RemoteDelegate { + fn invoke_tool<'a>( + &'a self, + invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + Box::pin(async move { + match self + .peer + .call( + self.session_id.clone(), + DelegateRequest::InvokeTool { + invocation: invocation.into(), + }, + cancellation_token, + ) + .await? + { + DelegateResponse::ToolResult { result } => Ok(result), + DelegateResponse::NotificationDelivered => { + Err("code-mode client returned an invalid tool result".to_string()) + } + } + }) + } + + fn notify<'a>( + &'a self, + call_id: String, + cell_id: CellId, + text: String, + cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async move { + match self + .peer + .call( + self.session_id.clone(), + DelegateRequest::Notify { + call_id, + cell_id: cell_id.into(), + text, + }, + cancellation_token, + ) + .await? + { + DelegateResponse::NotificationDelivered => Ok(()), + DelegateResponse::ToolResult { .. } => { + Err("code-mode client returned an invalid notification result".to_string()) + } + } + }) + } + + fn cell_closed(&self, cell_id: &CellId) { + self.peer + .close_cell(self.session_id.clone(), cell_id.clone()); + } +} diff --git a/codex-rs/code-mode-host/src/host_tests.rs b/codex-rs/code-mode-host/src/host_tests.rs new file mode 100644 index 0000000000..b189b89e09 --- /dev/null +++ b/codex-rs/code-mode-host/src/host_tests.rs @@ -0,0 +1,486 @@ +use codex_code_mode_protocol::host::Capability; +use codex_code_mode_protocol::host::CapabilitySet; +use codex_code_mode_protocol::host::ClientHello; +use codex_code_mode_protocol::host::ClientToHost; +use codex_code_mode_protocol::host::EncodedFrame; +use codex_code_mode_protocol::host::FramedReader; +use codex_code_mode_protocol::host::FramedWriter; +use codex_code_mode_protocol::host::HandshakeRejectReason; +use codex_code_mode_protocol::host::HostHello; +use codex_code_mode_protocol::host::HostRequest; +use codex_code_mode_protocol::host::HostResponse; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::ProtocolVersion; +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::SessionId; +use codex_code_mode_protocol::host::SupportedProtocolVersions; +use codex_code_mode_protocol::host::WireExecuteRequest; +use codex_code_mode_protocol::host::WireResult; +use pretty_assertions::assert_eq; +use tokio::sync::Semaphore; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; + +use super::HostState; +use super::MAX_ACTIVE_CELLS; +use super::MAX_IN_FLIGHT_REQUESTS; +use super::MAX_RECENT_REQUEST_IDS; +use super::RequestKind; +use super::RequestRegistry; +use super::SeenSessionIds; +use super::peer::HostPeer; +use super::run; + +fn client_hello( + versions: impl IntoIterator, + required_capabilities: CapabilitySet, +) -> ClientToHost { + ClientToHost::ClientHello( + ClientHello::new( + SupportedProtocolVersions::try_new(versions).expect("supported versions"), + required_capabilities, + CapabilitySet::empty(), + ) + .expect("client hello"), + ) +} + +fn session_id(value: &str) -> SessionId { + SessionId::new(value).expect("session ID") +} + +fn request_id(value: i64) -> RequestId { + RequestId::new(value) +} + +async fn decode_frame(frame: EncodedFrame) -> HostToClient { + let (reader, writer) = tokio::io::duplex(/*max_buf_size*/ 4096); + let writer = tokio::spawn(async move { + FramedWriter::new(writer) + .write_frame(&frame) + .await + .expect("write encoded frame"); + }); + let message = FramedReader::new(reader) + .read() + .await + .expect("read encoded frame") + .expect("encoded frame message"); + writer.await.expect("frame writer task"); + message +} + +fn execute_request(source: &str) -> WireExecuteRequest { + WireExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: source.to_string(), + yield_time_ms: Some(60_000), + max_output_tokens: Some(1_000), + } +} + +#[tokio::test] +async fn handshake_and_multiple_session_lifecycles_are_ordered() { + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 4096); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + + writer + .write(&client_hello([ProtocolVersion::V1], CapabilitySet::empty())) + .await + .expect("write hello"); + assert_eq!( + reader.read::().await.expect("read hello"), + Some(HostToClient::HostHello(HostHello::new( + ProtocolVersion::V1, + CapabilitySet::empty(), + ))) + ); + + for (request_id, id) in [ + (request_id(/*value*/ 1), "session-1"), + (request_id(/*value*/ 2), "session-2"), + ] { + writer + .write(&ClientToHost::Request { + id: request_id, + request: HostRequest::OpenSession { + session_id: session_id(id), + }, + }) + .await + .expect("open session"); + assert_eq!( + reader.read::().await.expect("session ready"), + Some(HostToClient::Response { + id: request_id, + result: WireResult::Ok { + value: HostResponse::SessionReady { + session_id: session_id(id), + }, + }, + }) + ); + } + + for (request_id, id) in [ + (request_id(/*value*/ 3), "session-1"), + (request_id(/*value*/ 4), "session-2"), + ] { + writer + .write(&ClientToHost::Request { + id: request_id, + request: HostRequest::ShutdownSession { + session_id: session_id(id), + }, + }) + .await + .expect("shutdown session"); + assert_eq!( + reader.read::().await.expect("session closed"), + Some(HostToClient::Response { + id: request_id, + result: WireResult::Ok { + value: HostResponse::SessionClosed { + session_id: session_id(id), + }, + }, + }) + ); + } + + drop(writer); + drop(reader); + host.await.expect("host task").expect("host connection"); +} + +#[tokio::test] +async fn incompatible_or_invalid_handshake_is_rejected() { + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + let version_two = ProtocolVersion::new(/*value*/ 2).expect("protocol version"); + + writer + .write(&client_hello([version_two], CapabilitySet::empty())) + .await + .expect("write hello"); + assert_eq!( + reader.read::().await.expect("rejection"), + Some(HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::NoCompatibleVersion { + supported_versions: SupportedProtocolVersions::try_new([ProtocolVersion::V1]) + .expect("host versions"), + }, + }) + ); + host.await.expect("host task").expect("host connection"); + + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + writer + .write(&ClientToHost::Request { + id: request_id(/*value*/ 1), + request: HostRequest::OpenSession { + session_id: session_id("session-1"), + }, + }) + .await + .expect("write invalid first message"); + assert_eq!( + reader.read::().await.expect("rejection"), + Some(HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::InvalidHello { + message: "first message must be connection/hello".to_string(), + }, + }) + ); + host.await.expect("host task").expect("host connection"); +} + +#[tokio::test] +async fn unsupported_required_capability_is_rejected() { + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + let capability = Capability::new("required").expect("capability"); + + writer + .write(&client_hello( + [ProtocolVersion::V1], + CapabilitySet::try_new([capability.clone()]).expect("capabilities"), + )) + .await + .expect("write hello"); + assert_eq!( + reader.read::().await.expect("rejection"), + Some(HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::MissingRequiredCapability { capability }, + }) + ); + host.await.expect("host task").expect("host connection"); +} + +#[tokio::test] +async fn session_id_cannot_be_reused_after_shutdown() { + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 2048); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + writer + .write(&client_hello([ProtocolVersion::V1], CapabilitySet::empty())) + .await + .expect("write hello"); + reader + .read::() + .await + .expect("read hello") + .expect("host hello"); + + let id = session_id("session-1"); + for (request_id, request) in [ + ( + request_id(/*value*/ 1), + HostRequest::OpenSession { + session_id: id.clone(), + }, + ), + ( + request_id(/*value*/ 2), + HostRequest::ShutdownSession { + session_id: id.clone(), + }, + ), + ] { + writer + .write(&ClientToHost::Request { + id: request_id, + request, + }) + .await + .expect("session request"); + reader + .read::() + .await + .expect("session response") + .expect("session response message"); + } + writer + .write(&ClientToHost::Request { + id: request_id(/*value*/ 3), + request: HostRequest::OpenSession { session_id: id }, + }) + .await + .expect("reuse session ID"); + assert_eq!( + reader.read::().await.expect("reuse response"), + Some(HostToClient::Response { + id: request_id(/*value*/ 3), + result: WireResult::Err { + message: "code-mode session ID `session-1` was reused".to_string(), + }, + }) + ); + drop(writer); + drop(reader); + host.await.expect("host task").expect("host connection"); +} + +#[test] +fn request_cancellation_tombstones_are_bounded() { + let mut requests = RequestRegistry::default(); + let duplicate = request_id(/*value*/ -1); + requests + .start(duplicate, RequestKind::OpenSession) + .expect("start duplicate probe"); + assert!(requests.start(duplicate, RequestKind::OpenSession).is_err()); + requests.finish(duplicate); + for value in 1..=MAX_RECENT_REQUEST_IDS as i64 + 100 { + let id = request_id(value); + requests + .start(id, RequestKind::Wait) + .expect("start request"); + requests.cancel(id); + requests.finish(id); + } + for value in 10_000..20_000 { + requests.cancel(request_id(value)); + } + + assert!(requests.active.is_empty()); + assert_eq!(requests.recent.len(), MAX_RECENT_REQUEST_IDS); + assert_eq!(requests.recent_order.len(), MAX_RECENT_REQUEST_IDS); +} + +#[tokio::test] +async fn request_task_panic_disconnects_host() { + let (outgoing_tx, _outgoing_rx) = mpsc::channel(/*max_capacity*/ 1); + let peer = Arc::new(HostPeer::new(outgoing_tx)); + let state = HostState { + sessions: Mutex::new(HashMap::new()), + seen_session_ids: Mutex::new(SeenSessionIds::default()), + requests: Mutex::new(RequestRegistry::default()), + request_tasks: TaskTracker::new(), + request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)), + active_cell_permits: Arc::new(Semaphore::new(MAX_ACTIVE_CELLS)), + closing: AtomicBool::new(false), + peer: Arc::clone(&peer), + }; + let task = state.request_tasks.spawn(async { + panic!("request panic probe"); + }); + state.supervise_request_task(task); + + tokio::time::timeout(Duration::from_secs(1), peer.disconnected()) + .await + .expect("request panic should disconnect host"); + assert!( + peer.failure() + .expect("request failure") + .contains("request task failed") + ); +} + +#[tokio::test] +async fn execute_request_id_remains_active_until_initial_response() { + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(/*max_capacity*/ 4); + let peer = Arc::new(HostPeer::new(outgoing_tx)); + let state = Arc::new(HostState { + sessions: Mutex::new(HashMap::new()), + seen_session_ids: Mutex::new(SeenSessionIds::default()), + requests: Mutex::new(RequestRegistry::default()), + request_tasks: TaskTracker::new(), + request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)), + active_cell_permits: Arc::new(Semaphore::new(MAX_ACTIVE_CELLS)), + closing: AtomicBool::new(false), + peer, + }); + let session_id = session_id("session-1"); + state + .open_session(session_id.clone()) + .expect("open session"); + let request_id = request_id(/*value*/ 1); + + state + .spawn_request( + request_id, + HostRequest::Execute { + session_id: session_id.clone(), + request: execute_request("await new Promise(() => {});"), + }, + ) + .expect("spawn execute request"); + let started = decode_frame(outgoing_rx.recv().await.expect("execution started frame")).await; + let HostToClient::Response { + id, + result: + WireResult::Ok { + value: HostResponse::ExecutionStarted { cell_id }, + }, + } = started + else { + panic!("expected execution started response"); + }; + assert_eq!(id, request_id); + assert!( + state + .requests + .lock() + .unwrap_or_else(PoisonError::into_inner) + .active + .contains_key(&request_id) + ); + + state + .session(&session_id) + .expect("session") + .terminate(cell_id.into()) + .await + .expect("terminate cell"); + state.disconnect().await; +} + +#[tokio::test] +async fn active_cell_limit_rejects_execute_without_disconnecting() { + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(/*max_capacity*/ 1); + let peer = Arc::new(HostPeer::new(outgoing_tx)); + let state = HostState { + sessions: Mutex::new(HashMap::new()), + seen_session_ids: Mutex::new(SeenSessionIds::default()), + requests: Mutex::new(RequestRegistry::default()), + request_tasks: TaskTracker::new(), + request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)), + active_cell_permits: Arc::new(Semaphore::new(/*permits*/ 0)), + closing: AtomicBool::new(false), + peer: Arc::clone(&peer), + }; + let session_id = session_id("session-1"); + state + .open_session(session_id.clone()) + .expect("open session"); + let request_id = request_id(/*value*/ 1); + + state + .handle_request( + request_id, + HostRequest::Execute { + session_id, + request: execute_request("text(\"hello\");"), + }, + CancellationToken::new(), + ) + .await; + + assert_eq!( + decode_frame(outgoing_rx.recv().await.expect("execute response frame")).await, + HostToClient::Response { + id: request_id, + result: WireResult::Err { + message: "code-mode host has too many active cells".to_string(), + }, + } + ); + assert!(!peer.is_disconnected()); + state.disconnect().await; +} + +#[tokio::test] +async fn cell_forwarding_panic_disconnects_host() { + let (outgoing_tx, _outgoing_rx) = mpsc::channel(/*max_capacity*/ 1); + let peer = Arc::new(HostPeer::new(outgoing_tx)); + peer.spawn_critical("cell forwarding", async { + panic!("cell forwarding panic probe"); + }); + + tokio::time::timeout(Duration::from_secs(1), peer.disconnected()) + .await + .expect("cell panic should disconnect host"); + assert!( + peer.failure() + .expect("cell failure") + .contains("cell forwarding task failed") + ); +} +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; +use std::sync::atomic::AtomicBool; +use std::time::Duration; diff --git a/codex-rs/code-mode-host/src/lib.rs b/codex-rs/code-mode-host/src/lib.rs new file mode 100644 index 0000000000..d2c1f507b4 --- /dev/null +++ b/codex-rs/code-mode-host/src/lib.rs @@ -0,0 +1,605 @@ +use std::collections::HashMap; +use std::collections::HashSet; +use std::collections::VecDeque; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +use codex_code_mode::InProcessCodeModeSession; +use codex_code_mode_protocol::host::CapabilitySet; +use codex_code_mode_protocol::host::ClientToHost; +use codex_code_mode_protocol::host::EncodedFrame; +use codex_code_mode_protocol::host::FramedReader; +use codex_code_mode_protocol::host::FramedWriter; +use codex_code_mode_protocol::host::HandshakeRejectReason; +use codex_code_mode_protocol::host::HostHello; +use codex_code_mode_protocol::host::HostRequest; +use codex_code_mode_protocol::host::HostResponse; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::ProtocolVersion; +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::SessionId; +use codex_code_mode_protocol::host::SupportedProtocolVersions; +use tokio::io::AsyncRead; +use tokio::io::AsyncWrite; +use tokio::sync::Semaphore; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; + +use self::delegate::RemoteDelegate; +use self::peer::HostPeer; + +mod delegate; +mod peer; + +const MAX_IN_FLIGHT_REQUESTS: usize = 256; +const MAX_ACTIVE_CELLS: usize = 128; +const MAX_RECENT_REQUEST_IDS: usize = 4096; +const MAX_RECENT_SESSION_IDS: usize = 4096; +const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); + +/// Runs one code-mode host connection over the process standard streams. +pub async fn run_stdio() -> Result<()> { + run(tokio::io::stdin(), tokio::io::stdout()).await +} + +/// Runs one code-mode host connection over an ordered input/output pair. +async fn run(reader: R, writer: W) -> Result<()> +where + R: AsyncRead + Send + Unpin + 'static, + W: AsyncWrite + Send + Unpin + 'static, +{ + let mut reader = FramedReader::new(reader); + let mut writer = FramedWriter::new(writer); + if !negotiate(&mut reader, &mut writer).await? { + return Ok(()); + } + + let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(/*max_capacity*/ 128); + let peer = Arc::new(HostPeer::new(outgoing_tx)); + let state = Arc::new(HostState { + sessions: Mutex::new(HashMap::new()), + seen_session_ids: Mutex::new(SeenSessionIds::default()), + requests: Mutex::new(RequestRegistry::default()), + request_tasks: TaskTracker::new(), + request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)), + active_cell_permits: Arc::new(Semaphore::new(MAX_ACTIVE_CELLS)), + closing: AtomicBool::new(false), + peer: Arc::clone(&peer), + }); + let writer_disconnected = peer.disconnection_token(); + let writer_task = tokio::spawn(async move { + loop { + tokio::select! { + _ = writer_disconnected.cancelled() => return Ok::<(), anyhow::Error>(()), + frame = outgoing_rx.recv() => { + let Some(frame) = frame else { + return Ok(()); + }; + if let Err(err) = writer.write_frame(&frame).await { + return Err( + anyhow::Error::new(err) + .context("failed to write code-mode host message") + ); + } + } + } + } + }); + let writer_peer = Arc::clone(&peer); + let writer_supervisor = tokio::spawn(async move { + match writer_task.await { + Ok(Ok(())) if !writer_peer.is_disconnected() => { + writer_peer.fail("code-mode writer task exited unexpectedly".to_string()); + } + Ok(Ok(())) => {} + Ok(Err(err)) => { + writer_peer.fail(format!("code-mode writer task failed: {err:#}")); + } + Err(err) => { + writer_peer.fail(format!("code-mode writer task failed: {err}")); + } + } + }); + + let input_result = async { + loop { + let message = tokio::select! { + _ = peer.disconnected() => break, + message = reader.read::() => message + .context("failed to read code-mode client message")?, + }; + let Some(message) = message else { + break; + }; + match message { + ClientToHost::ClientHello(_) => { + anyhow::bail!("received a second code-mode client hello"); + } + ClientToHost::Request { id, request } => { + state.spawn_request(id, request)?; + } + ClientToHost::CancelRequest { id } => { + state.cancel_request(id); + } + ClientToHost::DelegateResponse { id, result } => { + peer.complete(id, result.into_result()).await; + } + } + } + Ok::<(), anyhow::Error>(()) + } + .await; + + peer.disconnect(); + if tokio::time::timeout(SHUTDOWN_TIMEOUT, state.disconnect()) + .await + .is_err() + { + peer.fail("timed out shutting down code-mode host state".to_string()); + } + drop(state); + tokio::time::timeout(SHUTDOWN_TIMEOUT, writer_supervisor) + .await + .context("timed out supervising code-mode writer task")? + .context("code-mode writer supervisor task failed")?; + let failure = peer.failure(); + drop(peer); + input_result?; + if let Some(failure) = failure { + anyhow::bail!(failure); + } + Ok(()) +} + +async fn negotiate(reader: &mut FramedReader, writer: &mut FramedWriter) -> Result +where + R: AsyncRead + Unpin, + W: AsyncWrite + Unpin, +{ + let Some(first_message) = reader + .read::() + .await + .context("failed to read code-mode client hello")? + else { + return Ok(false); + }; + let ClientToHost::ClientHello(client_hello) = first_message else { + writer + .write(&HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::InvalidHello { + message: "first message must be connection/hello".to_string(), + }, + }) + .await + .context("failed to reject invalid code-mode client hello")?; + return Ok(false); + }; + + let supported_versions = SupportedProtocolVersions::try_new([ProtocolVersion::V1])?; + if !client_hello + .supported_versions() + .contains(ProtocolVersion::V1) + { + writer + .write(&HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::NoCompatibleVersion { supported_versions }, + }) + .await + .context("failed to reject incompatible code-mode client")?; + return Ok(false); + } + + let host_capabilities = CapabilitySet::empty(); + if let Some(capability) = client_hello + .required_capabilities() + .iter() + .find(|capability| !host_capabilities.contains(capability)) + { + writer + .write(&HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::MissingRequiredCapability { + capability: capability.clone(), + }, + }) + .await + .context("failed to reject unsupported code-mode capability")?; + return Ok(false); + } + + writer + .write(&HostToClient::HostHello(HostHello::new( + ProtocolVersion::V1, + host_capabilities, + ))) + .await + .context("failed to write code-mode host hello")?; + Ok(true) +} + +struct HostState { + sessions: Mutex>>, + seen_session_ids: Mutex, + requests: Mutex, + request_tasks: TaskTracker, + request_permits: Arc, + active_cell_permits: Arc, + closing: AtomicBool, + peer: Arc, +} + +impl HostState { + fn spawn_request( + self: &Arc, + request_id: RequestId, + request: HostRequest, + ) -> Result<(), anyhow::Error> { + let cancellation = self + .requests + .lock() + .unwrap_or_else(PoisonError::into_inner) + .start(request_id, RequestKind::from(&request))?; + let Ok(permit) = Arc::clone(&self.request_permits).try_acquire_owned() else { + self.respond( + request_id, + Err("code-mode host has too many in-flight requests".to_string()), + ); + self.finish_request(request_id); + return Ok(()); + }; + let state = Arc::clone(self); + let request_task = self.request_tasks.spawn(async move { + let _permit = permit; + state + .handle_request(request_id, request, cancellation) + .await; + state.finish_request(request_id); + }); + self.supervise_request_task(request_task); + Ok(()) + } + + fn supervise_request_task(&self, task: tokio::task::JoinHandle<()>) { + let peer = Arc::clone(&self.peer); + tokio::spawn(async move { + if let Err(err) = task.await { + peer.fail(format!("code-mode request task failed: {err}")); + } + }); + } + + async fn handle_request( + &self, + request_id: RequestId, + request: HostRequest, + cancellation: CancellationToken, + ) { + if self.closing.load(Ordering::Acquire) { + self.respond( + request_id, + Err("code-mode host is shutting down".to_string()), + ); + return; + } + match request { + HostRequest::OpenSession { session_id } => { + let result = self + .open_session(session_id.clone()) + .map(|()| HostResponse::SessionReady { session_id }); + self.respond(request_id, result); + } + HostRequest::Execute { + session_id, + request, + } => { + if cancellation.is_cancelled() { + self.respond(request_id, Err("code-mode request cancelled".to_string())); + return; + } + let request = match request.try_into() { + Ok(request) => request, + Err(err) => { + self.respond( + request_id, + Err(format!("invalid code-mode execute request: {err}")), + ); + return; + } + }; + let session = match self.session(&session_id) { + Ok(session) => session, + Err(err) => { + self.respond(request_id, Err(err)); + return; + } + }; + let Ok(active_cell_permit) = + Arc::clone(&self.active_cell_permits).try_acquire_owned() + else { + self.respond( + request_id, + Err("code-mode host has too many active cells".to_string()), + ); + return; + }; + let result = session.execute(request).await; + match result { + Ok(started) => { + let cell_id = started.cell_id.clone(); + self.respond( + request_id, + Ok(HostResponse::ExecutionStarted { + cell_id: cell_id.into(), + }), + ); + let initial_response_sent = self.peer.start_cell( + session_id, + request_id, + started, + active_cell_permit, + ); + let _ = initial_response_sent.await; + } + Err(err) => self.respond(request_id, Err(err)), + } + } + HostRequest::Wait { + session_id, + request, + } => { + let result = match self.session(&session_id) { + Ok(session) => { + tokio::select! { + biased; + _ = cancellation.cancelled() => { + Err("code-mode request cancelled".to_string()) + } + result = session.wait(request.into()) => result.map(|outcome| { + HostResponse::WaitCompleted { + outcome: outcome.into(), + } + }), + } + } + Err(err) => Err(err), + }; + self.respond(request_id, result); + } + HostRequest::Terminate { + session_id, + cell_id, + } => { + let result = match self.session(&session_id) { + Ok(session) => session.terminate(cell_id.into()).await.map(|outcome| { + HostResponse::WaitCompleted { + outcome: outcome.into(), + } + }), + Err(err) => Err(err), + }; + self.respond(request_id, result); + } + HostRequest::ShutdownSession { session_id } => { + let session = self + .sessions + .lock() + .unwrap_or_else(PoisonError::into_inner) + .remove(&session_id); + let result = match session { + Some(session) => match session.shutdown().await { + Ok(()) => { + self.peer.wait_for_session_cells(&session_id).await; + Ok(HostResponse::SessionClosed { session_id }) + } + Err(err) => Err(err), + }, + None => Err(format!("unknown code-mode session {session_id}")), + }; + self.respond(request_id, result); + } + } + } + + fn open_session(&self, session_id: SessionId) -> Result<(), String> { + let mut sessions = self.sessions.lock().unwrap_or_else(PoisonError::into_inner); + if sessions.contains_key(&session_id) { + return Err(format!( + "code-mode session ID `{session_id}` is already open" + )); + } + if self.closing.load(Ordering::Acquire) { + return Err("code-mode host is shutting down".to_string()); + } + if !self + .seen_session_ids + .lock() + .unwrap_or_else(PoisonError::into_inner) + .remember(session_id.clone()) + { + return Err(format!("code-mode session ID `{session_id}` was reused")); + } + let delegate = Arc::new(RemoteDelegate::new( + session_id.clone(), + Arc::clone(&self.peer), + )); + let peer = Arc::downgrade(&self.peer); + let task_failure_handler = Arc::new(move |reason| { + if let Some(peer) = peer.upgrade() { + peer.fail(reason); + } + }); + sessions.insert( + session_id, + Arc::new( + InProcessCodeModeSession::with_delegate_and_task_failure_handler( + delegate, + task_failure_handler, + ), + ), + ); + Ok(()) + } + + fn session(&self, session_id: &SessionId) -> Result, String> { + self.sessions + .lock() + .unwrap_or_else(PoisonError::into_inner) + .get(session_id) + .cloned() + .ok_or_else(|| format!("unknown code-mode session {session_id}")) + } + + fn respond(&self, id: RequestId, result: Result) { + self.peer.respond(id, result); + } + + fn cancel_request(&self, request_id: RequestId) { + self.requests + .lock() + .unwrap_or_else(PoisonError::into_inner) + .cancel(request_id); + } + + fn finish_request(&self, request_id: RequestId) { + self.requests + .lock() + .unwrap_or_else(PoisonError::into_inner) + .finish(request_id); + } + + async fn disconnect(&self) { + self.closing.store(true, Ordering::Release); + self.requests + .lock() + .unwrap_or_else(PoisonError::into_inner) + .cancel_all(); + self.request_tasks.close(); + self.request_tasks.wait().await; + let sessions = self + .sessions + .lock() + .unwrap_or_else(PoisonError::into_inner) + .drain() + .map(|(_, session)| session) + .collect::>(); + for session in sessions { + let _ = session.shutdown().await; + } + } +} + +#[derive(Clone, Copy)] +enum RequestKind { + OpenSession, + Execute, + Wait, + Terminate, + ShutdownSession, +} + +impl RequestKind { + fn from(request: &HostRequest) -> Self { + match request { + HostRequest::OpenSession { .. } => Self::OpenSession, + HostRequest::Execute { .. } => Self::Execute, + HostRequest::Wait { .. } => Self::Wait, + HostRequest::Terminate { .. } => Self::Terminate, + HostRequest::ShutdownSession { .. } => Self::ShutdownSession, + } + } + + fn is_cancellable(self) -> bool { + matches!(self, Self::Execute | Self::Wait) + } +} + +struct ActiveRequest { + kind: RequestKind, + cancellation: CancellationToken, +} + +#[derive(Default)] +struct RequestRegistry { + active: HashMap, + recent: HashSet, + recent_order: VecDeque, +} + +impl RequestRegistry { + fn start( + &mut self, + request_id: RequestId, + kind: RequestKind, + ) -> Result { + if self.active.contains_key(&request_id) || self.recent.contains(&request_id) { + anyhow::bail!("duplicate code-mode request ID {request_id:?}"); + } + let cancellation = CancellationToken::new(); + self.active.insert( + request_id, + ActiveRequest { + kind, + cancellation: cancellation.clone(), + }, + ); + Ok(cancellation) + } + + fn cancel(&self, request_id: RequestId) { + if let Some(request) = self.active.get(&request_id) + && request.kind.is_cancellable() + { + request.cancellation.cancel(); + } + } + + fn finish(&mut self, request_id: RequestId) { + if self.active.remove(&request_id).is_none() { + return; + } + self.recent.insert(request_id); + self.recent_order.push_back(request_id); + while self.recent_order.len() > MAX_RECENT_REQUEST_IDS { + if let Some(expired) = self.recent_order.pop_front() { + self.recent.remove(&expired); + } + } + } + + fn cancel_all(&self) { + for request in self.active.values() { + request.cancellation.cancel(); + } + } +} + +#[derive(Default)] +struct SeenSessionIds { + ids: HashSet, + order: VecDeque, +} + +impl SeenSessionIds { + fn remember(&mut self, session_id: SessionId) -> bool { + if !self.ids.insert(session_id.clone()) { + return false; + } + self.order.push_back(session_id); + while self.order.len() > MAX_RECENT_SESSION_IDS { + if let Some(expired) = self.order.pop_front() { + self.ids.remove(&expired); + } + } + true + } +} + +#[cfg(test)] +#[path = "host_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode-host/src/main.rs b/codex-rs/code-mode-host/src/main.rs index f328e4d9d0..b215869c8a 100644 --- a/codex-rs/code-mode-host/src/main.rs +++ b/codex-rs/code-mode-host/src/main.rs @@ -1 +1,4 @@ -fn main() {} +#[tokio::main(flavor = "current_thread")] +async fn main() -> anyhow::Result<()> { + codex_code_mode_host::run_stdio().await +} diff --git a/codex-rs/code-mode-host/src/peer.rs b/codex-rs/code-mode-host/src/peer.rs new file mode 100644 index 0000000000..14ff068e7b --- /dev/null +++ b/codex-rs/code-mode-host/src/peer.rs @@ -0,0 +1,541 @@ +use std::collections::HashMap; +use std::collections::VecDeque; +use std::sync::Arc; +use std::sync::Mutex as StdMutex; +use std::sync::PoisonError; +use std::sync::atomic::AtomicI64; +use std::sync::atomic::Ordering; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::host::DelegateRequest; +use codex_code_mode_protocol::host::DelegateRequestId; +use codex_code_mode_protocol::host::DelegateResponse; +use codex_code_mode_protocol::host::EncodedFrame; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::SessionId; +use codex_code_mode_protocol::host::WireResult; +use tokio::sync::Mutex; +use tokio::sync::Notify; +use tokio::sync::OwnedSemaphorePermit; +use tokio::sync::Semaphore; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio_util::sync::CancellationToken; + +const CELL_MESSAGE_CAPACITY: usize = 128; +const MAX_PENDING_DELEGATE_CALLS: usize = 256; + +pub(super) struct HostPeer { + outgoing_tx: mpsc::Sender, + pending: Mutex>, + delegate_permits: Arc, + cell_routes: StdMutex>, + cell_routes_changed: Notify, + next_request_id: AtomicI64, + disconnected: CancellationToken, + failure: StdMutex>, +} + +struct PendingDelegate { + response_tx: oneshot::Sender>, + dispatched: bool, + _permit: OwnedSemaphorePermit, +} + +enum CellRoute { + Pending(VecDeque), + Active(mpsc::Sender), +} + +enum CellMessage { + Delegate { + id: DelegateRequestId, + request: DelegateRequest, + dispatched_tx: oneshot::Sender>, + }, + Closed, +} + +impl HostPeer { + pub(super) fn new(outgoing_tx: mpsc::Sender) -> Self { + Self { + outgoing_tx, + pending: Mutex::new(HashMap::new()), + delegate_permits: Arc::new(Semaphore::new(MAX_PENDING_DELEGATE_CALLS)), + cell_routes: StdMutex::new(HashMap::new()), + cell_routes_changed: Notify::new(), + next_request_id: AtomicI64::new(1), + disconnected: CancellationToken::new(), + failure: StdMutex::new(None), + } + } + + pub(super) fn send(&self, message: HostToClient) -> Result<(), PeerSendError> { + let frame = EncodedFrame::encode(&message) + .map_err(|err| PeerSendError::Payload(err.to_string()))?; + self.send_frame(frame) + } + + pub(super) fn respond( + &self, + id: RequestId, + result: Result, + ) { + let message = HostToClient::Response { + id, + result: WireResult::from_result(result), + }; + if let Err(PeerSendError::Payload(err)) = self.send(message) { + let _ = self.send(HostToClient::Response { + id, + result: WireResult::Err { + message: format!("code-mode host response exceeds the IPC frame limit: {err}"), + }, + }); + } + } + + fn initial_response( + &self, + id: RequestId, + result: Result, + ) { + let message = HostToClient::InitialResponse { + id, + result: WireResult::from_result(result), + }; + if let Err(PeerSendError::Payload(err)) = self.send(message) { + let _ = self.send(HostToClient::InitialResponse { + id, + result: WireResult::Err { + message: format!( + "code-mode initial response exceeds the IPC frame limit: {err}" + ), + }, + }); + } + } + + pub(super) async fn call( + self: &Arc, + session_id: SessionId, + request: DelegateRequest, + cancellation_token: CancellationToken, + ) -> Result { + if self.disconnected.is_cancelled() { + return Err("code-mode client connection closed".to_string()); + } + let Ok(permit) = Arc::clone(&self.delegate_permits).try_acquire_owned() else { + return Err("code-mode host has too many pending delegate calls".to_string()); + }; + let id = DelegateRequestId::new(self.next_request_id.fetch_add(1, Ordering::Relaxed)); + let (response_tx, response_rx) = oneshot::channel(); + self.pending.lock().await.insert( + id, + PendingDelegate { + response_tx, + dispatched: false, + _permit: permit, + }, + ); + let mut pending = PendingDelegateRequest::new(Arc::clone(self), id); + let cell_id = match &request { + DelegateRequest::InvokeTool { invocation } => invocation.cell_id.clone().into(), + DelegateRequest::Notify { cell_id, .. } => cell_id.clone().into(), + }; + let (dispatched_tx, dispatched_rx) = oneshot::channel(); + if let Err(err) = self.route_cell_message( + (session_id, cell_id), + CellMessage::Delegate { + id, + request, + dispatched_tx, + }, + ) { + self.pending.lock().await.remove(&id); + pending.disarm(); + return Err(err); + } + + let dispatched = tokio::select! { + dispatched = dispatched_rx => dispatched.map_err(|_| { + "code-mode cell route closed before dispatching delegate request".to_string() + })?, + _ = self.disconnected.cancelled() => { + self.pending.lock().await.remove(&id); + pending.disarm(); + return Err("code-mode client connection closed".to_string()); + } + }; + if let Err(err) = dispatched { + self.pending.lock().await.remove(&id); + pending.disarm(); + return Err(err); + } + + tokio::select! { + response = response_rx => { + pending.disarm(); + response.map_err(|_| { + "code-mode client closed before returning delegate output".to_string() + })? + } + _ = cancellation_token.cancelled() => { + if self.remove_pending(id).await.is_some() { + let _ = self.send(HostToClient::CancelDelegateRequest { id }); + } + pending.disarm(); + Err("code mode delegate request cancelled".to_string()) + } + _ = self.disconnected.cancelled() => { + self.pending.lock().await.remove(&id); + pending.disarm(); + Err("code-mode client connection closed".to_string()) + } + } + } + + pub(super) async fn complete( + &self, + id: DelegateRequestId, + response: Result, + ) { + if let Some(pending) = self.remove_pending(id).await { + let _ = pending.response_tx.send(response); + } + } + + pub(super) fn start_cell( + self: &Arc, + session_id: SessionId, + request_id: RequestId, + started: StartedCell, + active_cell_permit: OwnedSemaphorePermit, + ) -> oneshot::Receiver<()> { + let (initial_response_sent_tx, initial_response_sent_rx) = oneshot::channel(); + let key = (session_id, started.cell_id.clone()); + let (messages_tx, messages_rx) = mpsc::channel(CELL_MESSAGE_CAPACITY); + let previous = self + .cell_routes + .lock() + .unwrap_or_else(PoisonError::into_inner) + .insert(key.clone(), CellRoute::Active(messages_tx.clone())); + match previous { + Some(CellRoute::Pending(messages)) => { + for message in messages { + if messages_tx.try_send(message).is_err() { + self.disconnect(); + return initial_response_sent_rx; + } + } + } + Some(CellRoute::Active(_)) => { + self.disconnect(); + return initial_response_sent_rx; + } + None => {} + } + let peer = Arc::clone(self); + self.spawn_critical("cell forwarding", async move { + drive_cell( + peer, + key, + request_id, + started, + messages_rx, + initial_response_sent_tx, + active_cell_permit, + ) + .await; + }); + initial_response_sent_rx + } + + pub(super) fn close_cell(&self, session_id: SessionId, cell_id: CellId) { + let _ = self.route_cell_message((session_id, cell_id), CellMessage::Closed); + } + + pub(super) fn disconnect(&self) { + self.disconnected.cancel(); + } + + pub(super) fn fail(&self, reason: String) { + let mut failure = self.failure.lock().unwrap_or_else(PoisonError::into_inner); + if failure.is_none() { + *failure = Some(reason); + } + drop(failure); + self.disconnect(); + } + + pub(super) fn failure(&self) -> Option { + self.failure + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone() + } + + pub(super) fn is_disconnected(&self) -> bool { + self.disconnected.is_cancelled() + } + + pub(super) async fn disconnected(&self) { + self.disconnected.cancelled().await; + } + + pub(super) fn disconnection_token(&self) -> CancellationToken { + self.disconnected.clone() + } + + pub(super) async fn wait_for_session_cells(&self, session_id: &SessionId) { + loop { + let changed = self.cell_routes_changed.notified(); + if !self + .cell_routes + .lock() + .unwrap_or_else(PoisonError::into_inner) + .keys() + .any(|(route_session_id, _)| route_session_id == session_id) + { + return; + } + tokio::select! { + _ = changed => {} + _ = self.disconnected.cancelled() => return, + } + } + } + + async fn send_delegate_if_pending( + &self, + id: DelegateRequestId, + session_id: SessionId, + request: DelegateRequest, + dispatched_tx: oneshot::Sender>, + ) { + let result = { + let mut pending = self.pending.lock().await; + let Some(pending) = pending.get_mut(&id) else { + let _ = dispatched_tx.send(Err( + "code-mode delegate request was cancelled before dispatch".to_string(), + )); + return; + }; + match self.send(HostToClient::DelegateRequest { + id, + session_id, + request, + }) { + Ok(()) => { + pending.dispatched = true; + Ok(()) + } + Err(err) => Err(err.to_string()), + } + }; + let _ = dispatched_tx.send(result); + } + + fn route_cell_message( + &self, + key: (SessionId, CellId), + message: CellMessage, + ) -> Result<(), String> { + use std::collections::hash_map::Entry; + + let result = match self + .cell_routes + .lock() + .unwrap_or_else(PoisonError::into_inner) + .entry(key) + { + Entry::Occupied(mut entry) => match entry.get_mut() { + CellRoute::Pending(messages) if messages.len() < CELL_MESSAGE_CAPACITY => { + messages.push_back(message); + Ok(()) + } + CellRoute::Pending(_) => Err("code-mode cell message queue is full".to_string()), + CellRoute::Active(sender) => sender + .try_send(message) + .map_err(|_| "code-mode cell message queue is unavailable".to_string()), + }, + Entry::Vacant(entry) => { + entry.insert(CellRoute::Pending(VecDeque::from([message]))); + Ok(()) + } + }; + if result.is_err() { + self.disconnect(); + } + result + } + + async fn remove_pending(&self, id: DelegateRequestId) -> Option { + self.pending.lock().await.remove(&id) + } + + pub(super) fn spawn_critical(self: &Arc, task_name: &'static str, future: F) + where + F: std::future::Future + Send + 'static, + { + let task = tokio::spawn(future); + let peer = Arc::clone(self); + tokio::spawn(async move { + if let Err(err) = task.await { + peer.fail(format!("code-mode {task_name} task failed: {err}")); + } + }); + } + + fn send_frame(&self, frame: EncodedFrame) -> Result<(), PeerSendError> { + match self.outgoing_tx.try_send(frame) { + Ok(()) => Ok(()), + Err(mpsc::error::TrySendError::Full(_)) => { + self.disconnect(); + Err(PeerSendError::Unavailable( + "code-mode host outgoing queue is full".to_string(), + )) + } + Err(mpsc::error::TrySendError::Closed(_)) => { + self.disconnect(); + Err(PeerSendError::Unavailable( + "code-mode client connection closed".to_string(), + )) + } + } + } +} + +async fn drive_cell( + peer: Arc, + key: (SessionId, CellId), + request_id: RequestId, + started: StartedCell, + mut messages_rx: mpsc::Receiver, + initial_response_sent_tx: oneshot::Sender<()>, + _active_cell_permit: OwnedSemaphorePermit, +) { + let mut initial_response_sent_tx = Some(initial_response_sent_tx); + let initial_response = started.initial_response(); + tokio::pin!(initial_response); + let closed = loop { + tokio::select! { + biased; + result = &mut initial_response => { + peer.initial_response(request_id, result.map(Into::into)); + if let Some(initial_response_sent_tx) = initial_response_sent_tx.take() { + let _ = initial_response_sent_tx.send(()); + } + break false; + } + message = messages_rx.recv() => match message { + Some(CellMessage::Delegate { + id, + request, + dispatched_tx, + }) => { + peer.send_delegate_if_pending(id, key.0.clone(), request, dispatched_tx).await; + } + Some(CellMessage::Closed) | None => break true, + }, + _ = peer.disconnected.cancelled() => { + peer.remove_cell_route(&key); + return; + } + } + }; + + if closed { + peer.initial_response(request_id, initial_response.await.map(Into::into)); + if let Some(initial_response_sent_tx) = initial_response_sent_tx.take() { + let _ = initial_response_sent_tx.send(()); + } + } else { + loop { + tokio::select! { + message = messages_rx.recv() => match message { + Some(CellMessage::Delegate { + id, + request, + dispatched_tx, + }) => { + peer.send_delegate_if_pending(id, key.0.clone(), request, dispatched_tx).await; + } + Some(CellMessage::Closed) | None => break, + }, + _ = peer.disconnected.cancelled() => { + peer.remove_cell_route(&key); + return; + } + } + } + } + let _ = peer.send(HostToClient::CellClosed { + session_id: key.0.clone(), + cell_id: (&key.1).into(), + }); + peer.remove_cell_route(&key); +} + +impl HostPeer { + fn remove_cell_route(&self, key: &(SessionId, CellId)) { + let removed = self + .cell_routes + .lock() + .unwrap_or_else(PoisonError::into_inner) + .remove(key); + if removed.is_some() { + self.cell_routes_changed.notify_waiters(); + } + } +} + +pub(super) enum PeerSendError { + Payload(String), + Unavailable(String), +} + +impl std::fmt::Display for PeerSendError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Payload(message) | Self::Unavailable(message) => formatter.write_str(message), + } + } +} + +struct PendingDelegateRequest { + peer: Arc, + id: Option, +} + +impl PendingDelegateRequest { + fn new(peer: Arc, id: DelegateRequestId) -> Self { + Self { peer, id: Some(id) } + } + + fn disarm(&mut self) { + self.id = None; + } +} + +impl Drop for PendingDelegateRequest { + fn drop(&mut self) { + let Some(id) = self.id.take() else { + return; + }; + let peer = Arc::clone(&self.peer); + tokio::spawn(async move { + if let Some(pending) = peer.remove_pending(id).await + && pending.dispatched + { + let _ = peer.send(HostToClient::CancelDelegateRequest { id }); + } + }); + } +} + +#[cfg(test)] +#[path = "peer_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode-host/src/peer_tests.rs b/codex-rs/code-mode-host/src/peer_tests.rs new file mode 100644 index 0000000000..a926027ce9 --- /dev/null +++ b/codex-rs/code-mode-host/src/peer_tests.rs @@ -0,0 +1,95 @@ +use std::sync::Arc; +use std::time::Duration; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::RuntimeResponse; +use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::host::DelegateRequest; +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::SessionId; +use pretty_assertions::assert_eq; +use tokio::sync::Semaphore; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio::sync::oneshot::error::TryRecvError; +use tokio_util::sync::CancellationToken; + +use super::HostPeer; +use super::MAX_PENDING_DELEGATE_CALLS; + +fn session_id(value: &str) -> SessionId { + SessionId::new(value).expect("session ID") +} + +#[tokio::test] +async fn start_cell_reports_when_initial_response_is_enqueued() { + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(/*max_capacity*/ 4); + let peer = Arc::new(HostPeer::new(outgoing_tx)); + let cell_id = CellId::new("cell-1".to_string()); + let (response_tx, response_rx) = oneshot::channel(); + let started = StartedCell::new(cell_id.clone(), response_rx); + let active_cell_permits = Arc::new(Semaphore::new(/*permits*/ 1)); + let active_cell_permit = Arc::clone(&active_cell_permits) + .try_acquire_owned() + .expect("active cell permit"); + + let mut initial_response_sent = peer.start_cell( + session_id("session-1"), + RequestId::new(/*value*/ 1), + started, + active_cell_permit, + ); + assert_eq!(initial_response_sent.try_recv(), Err(TryRecvError::Empty)); + + response_tx + .send(RuntimeResponse::Result { + cell_id: cell_id.clone(), + content_items: Vec::new(), + error_text: None, + }) + .expect("initial response receiver"); + initial_response_sent + .await + .expect("initial response completion"); + outgoing_rx.recv().await.expect("initial response frame"); + assert_eq!(active_cell_permits.available_permits(), 0); + + peer.close_cell(session_id("session-1"), cell_id); + let permit = tokio::time::timeout( + Duration::from_secs(1), + Arc::clone(&active_cell_permits).acquire_owned(), + ) + .await + .expect("cell permit should be released") + .expect("cell permit semaphore should remain open"); + drop(permit); +} + +#[tokio::test] +async fn pending_delegate_limit_rejects_call_without_disconnecting() { + let (outgoing_tx, _outgoing_rx) = mpsc::channel(/*max_capacity*/ 1); + let peer = Arc::new(HostPeer::new(outgoing_tx)); + let permits = Arc::clone(&peer.delegate_permits) + .acquire_many_owned(MAX_PENDING_DELEGATE_CALLS as u32) + .await + .expect("delegate permits"); + + let result = peer + .call( + session_id("session-1"), + DelegateRequest::Notify { + call_id: "call-1".to_string(), + cell_id: CellId::new("cell-1".to_string()).into(), + text: "hello".to_string(), + }, + CancellationToken::new(), + ) + .await; + + assert_eq!( + result, + Err("code-mode host has too many pending delegate calls".to_string()) + ); + assert!(!peer.is_disconnected()); + drop(permits); +} diff --git a/justfile b/justfile index 03f87391ac..d4a0e6566c 100644 --- a/justfile +++ b/justfile @@ -31,6 +31,10 @@ tui-with-exec-server *args: file-search *args: cargo run --bin codex-file-search -- {args} +# Run the standalone code-mode host from source. +code-mode-host *args: + cargo run --bin codex-code-mode-host -- {args} + # Build the CLI and run the app-server test client app-server-test-client *args: cargo build -p codex-cli @@ -107,6 +111,16 @@ bazel-codex *args: bazel-codex *args: bazel run //codex-rs/cli:codex --run_under='cd /d "{{ invocation_directory_native() }}" &&' -- @($args | Select-Object -Skip 1) +# Build and run the standalone code-mode host from source using Bazel. +[no-cd] +[unix] +bazel-code-mode-host *args: + bazel run //codex-rs/code-mode-host:codex-code-mode-host --run_under="cd $PWD &&" -- "$@" + +[windows] +bazel-code-mode-host *args: + bazel run //codex-rs/code-mode-host:codex-code-mode-host --run_under='cd /d "{{ invocation_directory_native() }}" &&' -- @($args | Select-Object -Skip 1) + [no-cd] bazel-lock-update: bazel mod deps --lockfile_mode=update From 841f30598c05eebb62139bdd75abef1b7b61d5f7 Mon Sep 17 00:00:00 2001 From: alexsong-oai Date: Thu, 25 Jun 2026 18:15:48 -0700 Subject: [PATCH 08/38] [codex] Attribute app-server analytics by thread originator (#29935) ## Why Desktop Work threads and regular Codex threads can share the same app-server connection. App-server analytics currently copy `product_client_id` from connection metadata for every thread-scoped event, so Work thread activity is attributed to the Desktop connection instead of the thread's resolved originator. This prevents analytics from distinguishing the two products on a shared connection. ## What changed - Publish the resolved originator after a thread is materialized, covering new, resumed, forked, and subagent threads. - Store that originator in the analytics reducer's existing per-thread state. - Override only `app_server_client.product_client_id` for thread, turn, tool, review, goal, guardian, and compaction events while preserving the connection's client name, version, and transport metadata. - Fall back to the connection-wide product client ID when a thread has no originator override. - Preserve persisted originators in thread initialization analytics for resume and fork flows. ## Validation - `just test -p codex-analytics thread_originator_overrides_shared_connection_across_thread_events subagent_events_keep_thread_originator_with_explicit_turn_connection` - `just test -p codex-app-server turn_start_tracks_thread_originator_in_analytics thread_start_tracks_thread_initialized_analytics thread_fork_tracks_thread_initialized_analytics thread_resume_tracks_thread_initialized_analytics` - `just test -p codex-core thread_manager` --- .../analytics/src/analytics_client_tests.rs | 176 +++++++++++++++++- codex-rs/analytics/src/client.rs | 26 +++ codex-rs/analytics/src/facts.rs | 1 + codex-rs/analytics/src/reducer.rs | 114 ++++++++---- codex-rs/app-server/src/outgoing_message.rs | 48 ++++- .../request_processors/thread_lifecycle.rs | 5 +- .../request_processors/thread_processor.rs | 13 +- .../app-server/tests/suite/v2/analytics.rs | 5 +- .../app-server/tests/suite/v2/thread_fork.rs | 1 + .../tests/suite/v2/thread_resume.rs | 91 ++++++++- .../app-server/tests/suite/v2/thread_start.rs | 2 + .../app-server/tests/suite/v2/turn_start.rs | 6 +- 12 files changed, 431 insertions(+), 57 deletions(-) diff --git a/codex-rs/analytics/src/analytics_client_tests.rs b/codex-rs/analytics/src/analytics_client_tests.rs index c183ad925a..4451d39ed9 100644 --- a/codex-rs/analytics/src/analytics_client_tests.rs +++ b/codex-rs/analytics/src/analytics_client_tests.rs @@ -558,6 +558,7 @@ async fn ingest_rejected_turn_steer( response: Box::new(sample_thread_resume_response( "thread-2", /*ephemeral*/ false, "gpt-5", )), + thread_originator: None, }, out, ) @@ -631,6 +632,7 @@ async fn ingest_turn_prerequisites( response: Box::new(sample_thread_start_response( "thread-2", /*ephemeral*/ false, "gpt-5", )), + thread_originator: None, }, out, ) @@ -654,6 +656,7 @@ async fn ingest_turn_prerequisites( connection_id: 7, request_id: RequestId::Integer(3), response: Box::new(sample_turn_start_response("turn-2")), + thread_originator: None, }, out, ) @@ -720,6 +723,7 @@ async fn ingest_review_prerequisites( response: Box::new(sample_thread_start_response( "thread-1", /*ephemeral*/ false, "gpt-5", )), + thread_originator: None, }, events, ) @@ -1652,6 +1656,7 @@ async fn initialize_caches_client_and_thread_lifecycle_publishes_once_initialize /*ephemeral*/ false, "gpt-5", )), + thread_originator: None, }, &mut events, ) @@ -1697,6 +1702,7 @@ async fn initialize_caches_client_and_thread_lifecycle_publishes_once_initialize response: Box::new(sample_thread_resume_response( "thread-1", /*ephemeral*/ true, "gpt-5", )), + thread_originator: None, }, &mut events, ) @@ -1741,6 +1747,157 @@ async fn initialize_caches_client_and_thread_lifecycle_publishes_once_initialize ); } +#[tokio::test] +async fn thread_originator_overrides_shared_connection_across_thread_events() { + let mut reducer = AnalyticsReducer::default(); + let mut events = Vec::new(); + + reducer + .ingest(sample_initialize_fact(/*connection_id*/ 7), &mut events) + .await; + for (request_id, thread_id, thread_originator) in [ + (1, "thread-work", Some(TEST_PRODUCT_CLIENT_ID.to_string())), + (2, "thread-default", None), + ] { + reducer + .ingest( + AnalyticsFact::ClientResponse { + connection_id: 7, + request_id: RequestId::Integer(request_id), + response: Box::new(sample_thread_start_response( + thread_id, /*ephemeral*/ false, "gpt-5", + )), + thread_originator, + }, + &mut events, + ) + .await; + } + + let initialized = serde_json::to_value(&events).expect("serialize thread events"); + assert_eq!( + initialized + .as_array() + .expect("thread events") + .iter() + .map(|event| { + json!({ + "thread_id": event["event_params"]["thread_id"], + "app_server_client": event["event_params"]["app_server_client"], + }) + }) + .collect::>(), + vec![ + json!({ + "thread_id": "thread-work", + "app_server_client": { + "product_client_id": TEST_PRODUCT_CLIENT_ID, + "client_name": "codex-tui", + "client_version": "1.0.0", + "rpc_transport": "websocket", + "experimental_api_enabled": false, + }, + }), + json!({ + "thread_id": "thread-default", + "app_server_client": { + "product_client_id": DEFAULT_ORIGINATOR, + "client_name": "codex-tui", + "client_version": "1.0.0", + "rpc_transport": "websocket", + "experimental_api_enabled": false, + }, + }), + ] + ); + + events.clear(); + reducer + .ingest( + AnalyticsFact::ClientRequest { + connection_id: 7, + request_id: RequestId::Integer(3), + request: Box::new(sample_turn_start_request( + "thread-work", + /*request_id*/ 3, + )), + }, + &mut events, + ) + .await; + reducer + .ingest( + AnalyticsFact::ClientResponse { + connection_id: 7, + request_id: RequestId::Integer(3), + response: Box::new(sample_turn_start_response("turn-1")), + thread_originator: None, + }, + &mut events, + ) + .await; + ingest_completed_command_execution_item(&mut reducer, &mut events, "thread-work", "item-work") + .await; + ingest_complete_child_turn(&mut reducer, &mut events, "thread-work", "turn-1").await; + reducer + .ingest( + AnalyticsFact::Custom(CustomAnalyticsFact::Compaction(Box::new( + CodexCompactionEvent { + thread_id: "thread-work".to_string(), + turn_id: "turn-compact".to_string(), + trigger: CompactionTrigger::Manual, + reason: CompactionReason::UserRequested, + implementation: CompactionImplementation::Responses, + phase: CompactionPhase::StandaloneTurn, + strategy: CompactionStrategy::Memento, + status: CompactionStatus::Completed, + codex_error_kind: None, + codex_error_http_status_code: None, + active_context_tokens_before: 131_000, + active_context_tokens_after: 64_000, + retained_image_count: None, + compaction_summary_tokens: None, + cached_input_tokens: None, + started_at: 100, + completed_at: 101, + duration_ms: Some(1200), + }, + ))), + &mut events, + ) + .await; + + let lifecycle = serde_json::to_value(&events).expect("serialize lifecycle events"); + assert_eq!( + lifecycle + .as_array() + .expect("lifecycle events") + .iter() + .map(|event| { + json!({ + "event_type": event["event_type"], + "product_client_id": + event["event_params"]["app_server_client"]["product_client_id"], + }) + }) + .collect::>(), + vec![ + json!({ + "event_type": "codex_command_execution_event", + "product_client_id": TEST_PRODUCT_CLIENT_ID, + }), + json!({ + "event_type": "codex_turn_event", + "product_client_id": TEST_PRODUCT_CLIENT_ID, + }), + json!({ + "event_type": "codex_compaction_event", + "product_client_id": TEST_PRODUCT_CLIENT_ID, + }), + ] + ); +} + #[tokio::test] async fn unrelated_client_requests_are_ignored_by_reducer() { let mut reducer = AnalyticsReducer::default(); @@ -1767,6 +1924,7 @@ async fn unrelated_client_requests_are_ignored_by_reducer() { connection_id: 7, request_id: RequestId::Integer(3), response: Box::new(sample_turn_start_response("turn-2")), + thread_originator: None, }, &mut events, ) @@ -1792,6 +1950,7 @@ async fn unrelated_client_responses_are_ignored_by_reducer() { response: Box::new(ClientResponsePayload::ThreadArchive( ThreadArchiveResponse {}, )), + thread_originator: None, }, &mut events, ) @@ -1851,6 +2010,7 @@ async fn compaction_event_ingests_custom_fact() { Some(AppServerThreadSource::Subagent), Some(parent_thread_id.to_string()), )), + thread_originator: None, }, &mut events, ) @@ -1971,6 +2131,7 @@ async fn guardian_review_event_ingests_custom_fact_with_optional_target_item() { /*ephemeral*/ false, "gpt-5", )), + thread_originator: None, }, &mut events, ) @@ -2499,6 +2660,7 @@ async fn item_review_summaries_do_not_cross_threads_with_reused_item_ids() { response: Box::new(sample_thread_start_response( "thread-2", /*ephemeral*/ false, "gpt-5", )), + thread_originator: None, }, &mut events, ) @@ -2749,7 +2911,7 @@ async fn subagent_thread_started_publishes_without_initialize() { } #[tokio::test] -async fn subagent_events_use_inherited_connection_unless_turn_connection_is_explicit() { +async fn subagent_events_keep_thread_originator_with_explicit_turn_connection() { let mut reducer = AnalyticsReducer::default(); let mut events = Vec::new(); let parent_thread_id = @@ -2786,6 +2948,7 @@ async fn subagent_events_use_inherited_connection_unless_turn_connection_is_expl /*ephemeral*/ false, "gpt-5", )), + thread_originator: None, }, &mut events, ) @@ -2908,6 +3071,7 @@ async fn subagent_events_use_inherited_connection_unless_turn_connection_is_expl connection_id: 8, request_id: RequestId::Integer(3), response: Box::new(sample_turn_start_response("turn-explicit")), + thread_originator: None, }, &mut events, ) @@ -2918,7 +3082,11 @@ async fn subagent_events_use_inherited_connection_unless_turn_connection_is_expl }; assert_eq!( event.event_params.app_server_client.product_client_id, - DEFAULT_ORIGINATOR + "parent-client" + ); + assert_eq!( + event.event_params.app_server_client.client_name.as_deref(), + Some("codex-tui") ); } @@ -3868,6 +4036,7 @@ async fn accepted_turn_steer_emits_expected_event() { connection_id: 7, request_id: RequestId::Integer(4), response: Box::new(sample_turn_steer_response("turn-2")), + thread_originator: None, }, &mut out, ) @@ -4039,6 +4208,7 @@ async fn turn_start_error_response_discards_pending_start_request() { connection_id: 7, request_id: RequestId::Integer(3), response: Box::new(sample_turn_start_response("turn-2")), + thread_originator: None, }, &mut out, ) @@ -4367,6 +4537,7 @@ async fn accepted_steers_increment_turn_steer_count() { connection_id: 7, request_id: RequestId::Integer(4), response: Box::new(sample_turn_steer_response("turn-2")), + thread_originator: None, }, &mut out, ) @@ -4414,6 +4585,7 @@ async fn accepted_steers_increment_turn_steer_count() { connection_id: 7, request_id: RequestId::Integer(6), response: Box::new(sample_turn_steer_response("turn-2")), + thread_originator: None, }, &mut out, ) diff --git a/codex-rs/analytics/src/client.rs b/codex-rs/analytics/src/client.rs index 5996734f8e..f957a665a0 100644 --- a/codex-rs/analytics/src/client.rs +++ b/codex-rs/analytics/src/client.rs @@ -435,6 +435,31 @@ impl AnalyticsEventsClient { connection_id: u64, request_id: RequestId, response: ClientResponsePayload, + ) { + self.track_response_inner( + connection_id, + request_id, + response, + /*thread_originator*/ None, + ); + } + + pub fn track_response_with_thread_originator( + &self, + connection_id: u64, + request_id: RequestId, + response: ClientResponsePayload, + thread_originator: String, + ) { + self.track_response_inner(connection_id, request_id, response, Some(thread_originator)); + } + + fn track_response_inner( + &self, + connection_id: u64, + request_id: RequestId, + response: ClientResponsePayload, + thread_originator: Option, ) { if !matches!( response, @@ -450,6 +475,7 @@ impl AnalyticsEventsClient { connection_id, request_id, response: Box::new(response), + thread_originator, }); } diff --git a/codex-rs/analytics/src/facts.rs b/codex-rs/analytics/src/facts.rs index 89b8708648..3346a6e19a 100644 --- a/codex-rs/analytics/src/facts.rs +++ b/codex-rs/analytics/src/facts.rs @@ -464,6 +464,7 @@ pub(crate) enum AnalyticsFact { connection_id: u64, request_id: RequestId, response: Box, + thread_originator: Option, }, ErrorResponse { connection_id: u64, diff --git a/codex-rs/analytics/src/reducer.rs b/codex-rs/analytics/src/reducer.rs index 522a307688..29f21bac52 100644 --- a/codex-rs/analytics/src/reducer.rs +++ b/codex-rs/analytics/src/reducer.rs @@ -165,6 +165,20 @@ struct ConnectionState { struct ThreadAnalyticsState { connection_id: Option, metadata: Option, + originator: Option, +} + +impl ThreadAnalyticsState { + fn app_server_client( + &self, + connection_state: &ConnectionState, + ) -> CodexAppServerClientMetadata { + let mut app_server_client = connection_state.app_server_client.clone(); + if let Some(originator) = self.originator.as_ref() { + app_server_client.product_client_id.clone_from(originator); + } + app_server_client + } } #[derive(Clone, Copy)] @@ -434,9 +448,11 @@ impl AnalyticsReducer { connection_id, request_id, response, + thread_originator, } => { if let Some(response) = response.into_client_response(request_id) { - self.ingest_response(connection_id, response, out).await; + self.ingest_response(connection_id, response, thread_originator, out) + .await; } } AnalyticsFact::ErrorResponse { @@ -575,6 +591,9 @@ impl AnalyticsReducer { .and_then(|parent_thread_id| self.threads.get(parent_thread_id)) .and_then(|thread| thread.connection_id); let thread_state = self.threads.entry(input.thread_id.clone()).or_default(); + thread_state + .originator + .get_or_insert_with(|| input.product_client_id.clone()); thread_state .metadata .get_or_insert_with(|| ThreadMetadataState { @@ -597,7 +616,7 @@ impl AnalyticsReducer { input: GuardianReviewEventParams, out: &mut Vec, ) { - let Some((connection_state, thread_metadata)) = + let Some((connection_state, thread_state, thread_metadata)) = self.thread_context_or_warn(AnalyticsDropSite::guardian(&input)) else { return; @@ -607,7 +626,7 @@ impl AnalyticsReducer { event_type: "codex_guardian_review", event_params: GuardianReviewEventPayload { session_id: thread_metadata.session_id.clone(), - app_server_client: connection_state.app_server_client.clone(), + app_server_client: thread_state.app_server_client(connection_state), runtime: connection_state.runtime.clone(), guardian_review: input, }, @@ -874,6 +893,7 @@ impl AnalyticsReducer { &mut self, connection_id: u64, response: ClientResponse, + thread_originator: Option, out: &mut Vec, ) { match response { @@ -883,6 +903,7 @@ impl AnalyticsReducer { response.thread, response.model, ThreadInitializationMode::New, + thread_originator, out, ); } @@ -892,6 +913,7 @@ impl AnalyticsReducer { response.thread, response.model, ThreadInitializationMode::Resumed, + thread_originator, out, ); } @@ -901,6 +923,7 @@ impl AnalyticsReducer { response.thread, response.model, ThreadInitializationMode::Forked, + thread_originator, out, ); } @@ -1248,7 +1271,7 @@ impl AnalyticsReducer { else { return; }; - let Some((connection_state, thread_metadata)) = self + let Some((connection_state, thread_state, thread_metadata)) = self .thread_context_or_warn(AnalyticsDropSite::tool_item(¬ification, item_id)) else { return; @@ -1260,6 +1283,7 @@ impl AnalyticsReducer { started_at_ms, completed_at_ms, connection_state, + thread_state, thread_metadata, review_summary: self.item_review_summaries.get(&key), }) { @@ -1316,6 +1340,7 @@ impl AnalyticsReducer { thread: codex_app_server_protocol::Thread, model: String, initialization_mode: ThreadInitializationMode, + thread_originator: Option, out: &mut Vec, ) { let session_source: SessionSource = thread.source.into(); @@ -1333,20 +1358,20 @@ impl AnalyticsReducer { parent_thread_id, initialization_mode, ); - self.threads.insert( - thread_id.clone(), - ThreadAnalyticsState { - connection_id: Some(connection_id), - metadata: Some(thread_metadata.clone()), - }, - ); + let thread_state = self.threads.entry(thread_id.clone()).or_default(); + if let Some(originator) = thread_originator { + thread_state.originator = Some(originator); + } + thread_state.connection_id = Some(connection_id); + thread_state.metadata = Some(thread_metadata.clone()); + let app_server_client = thread_state.app_server_client(connection_state); out.push(TrackEventRequest::ThreadInitialized( ThreadInitializedEvent { event_type: "codex_thread_initialized", event_params: ThreadInitializedEventParams { thread_id, session_id, - app_server_client: connection_state.app_server_client.clone(), + app_server_client, runtime: connection_state.runtime.clone(), model, ephemeral: thread.ephemeral, @@ -1362,7 +1387,7 @@ impl AnalyticsReducer { } fn ingest_compaction(&mut self, input: CodexCompactionEvent, out: &mut Vec) { - let Some((connection_state, thread_metadata)) = + let Some((connection_state, thread_state, thread_metadata)) = self.thread_context_or_warn(AnalyticsDropSite::compaction(&input)) else { return; @@ -1373,7 +1398,7 @@ impl AnalyticsReducer { event_params: codex_compaction_event_params( input, thread_metadata.session_id.clone(), - connection_state.app_server_client.clone(), + thread_state.app_server_client(connection_state), connection_state.runtime.clone(), thread_metadata.thread_source.clone(), thread_metadata.subagent_source.clone(), @@ -1384,7 +1409,7 @@ impl AnalyticsReducer { } fn ingest_goal(&mut self, input: CodexGoalEvent, out: &mut Vec) { - let Some((connection_state, thread_metadata)) = + let Some((connection_state, thread_state, thread_metadata)) = self.thread_context_or_warn(AnalyticsDropSite::goal(&input)) else { return; @@ -1394,7 +1419,7 @@ impl AnalyticsReducer { event_params: codex_goal_event_params( input, thread_metadata.session_id.clone(), - connection_state.app_server_client.clone(), + thread_state.app_server_client(connection_state), connection_state.runtime.clone(), thread_metadata.thread_source.clone(), thread_metadata.subagent_source.clone(), @@ -1483,11 +1508,11 @@ impl AnalyticsReducer { return; }; let drop_site = AnalyticsDropSite::turn_steer(&pending_request.thread_id); - let Some(thread_metadata) = self - .threads - .get(drop_site.thread_id) - .and_then(|thread| thread.metadata.as_ref()) - else { + let Some(thread_state) = self.threads.get(drop_site.thread_id) else { + warn_missing_analytics_context(&drop_site, MissingAnalyticsContext::ThreadMetadata); + return; + }; + let Some(thread_metadata) = thread_state.metadata.as_ref() else { warn_missing_analytics_context(&drop_site, MissingAnalyticsContext::ThreadMetadata); return; }; @@ -1498,7 +1523,7 @@ impl AnalyticsReducer { session_id: thread_metadata.session_id.clone(), expected_turn_id: Some(pending_request.expected_turn_id), accepted_turn_id, - app_server_client: connection_state.app_server_client.clone(), + app_server_client: thread_state.app_server_client(connection_state), runtime: connection_state.runtime.clone(), thread_source: thread_metadata.thread_source.clone(), subagent_source: thread_metadata.subagent_source.clone(), @@ -1529,7 +1554,7 @@ impl AnalyticsReducer { &pending_review, ); } - let Some((connection_state, thread_metadata)) = + let Some((connection_state, thread_state, thread_metadata)) = self.thread_context_or_warn(AnalyticsDropSite::review(&pending_review)) else { return; @@ -1541,7 +1566,7 @@ impl AnalyticsReducer { turn_id: pending_review.turn_id, item_id: pending_review.item_id, review_id: pending_review.review_id, - app_server_client: connection_state.app_server_client.clone(), + app_server_client: thread_state.app_server_client(connection_state), runtime: connection_state.runtime.clone(), thread_source: thread_metadata.thread_source.clone(), subagent_source: thread_metadata.subagent_source.clone(), @@ -1610,18 +1635,18 @@ impl AnalyticsReducer { ); return; }; - let Some(thread_metadata) = self - .threads - .get(drop_site.thread_id) - .and_then(|thread| thread.metadata.as_ref()) - else { + let Some(thread_state) = self.threads.get(drop_site.thread_id) else { + warn_missing_analytics_context(&drop_site, MissingAnalyticsContext::ThreadMetadata); + return; + }; + let Some(thread_metadata) = thread_state.metadata.as_ref() else { warn_missing_analytics_context(&drop_site, MissingAnalyticsContext::ThreadMetadata); return; }; let turn_event = TrackEventRequest::TurnEvent(Box::new(CodexTurnEventRequest { event_type: "codex_turn_event", event_params: codex_turn_event_params( - connection_state.app_server_client.clone(), + thread_state.app_server_client(connection_state), connection_state.runtime.clone(), turn_id.to_string(), turn_state, @@ -1663,17 +1688,18 @@ impl AnalyticsReducer { fn thread_context_or_warn( &self, drop_site: AnalyticsDropSite<'_>, - ) -> Option<(&ConnectionState, &ThreadMetadataState)> { + ) -> Option<( + &ConnectionState, + &ThreadAnalyticsState, + &ThreadMetadataState, + )> { let connection_state = self.thread_connection_or_warn(drop_site)?; - let Some(thread_metadata) = self - .threads - .get(drop_site.thread_id) - .and_then(|thread| thread.metadata.as_ref()) - else { + let thread_state = self.threads.get(drop_site.thread_id)?; + let Some(thread_metadata) = thread_state.metadata.as_ref() else { warn_missing_analytics_context(&drop_site, MissingAnalyticsContext::ThreadMetadata); return None; }; - Some((connection_state, thread_metadata)) + Some((connection_state, thread_state, thread_metadata)) } } @@ -1743,6 +1769,7 @@ struct ToolItemEventInput<'a> { started_at_ms: u64, completed_at_ms: u64, connection_state: &'a ConnectionState, + thread_state: &'a ThreadAnalyticsState, thread_metadata: &'a ThreadMetadataState, review_summary: Option<&'a ItemReviewSummary>, } @@ -1755,6 +1782,7 @@ fn tool_item_event(input: ToolItemEventInput<'_>) -> Option { started_at_ms, completed_at_ms, connection_state, + thread_state, thread_metadata, review_summary, } = input; @@ -1784,6 +1812,7 @@ fn tool_item_event(input: ToolItemEventInput<'_>) -> Option { started_at_ms, completed_at_ms, connection_state, + thread_state, thread_metadata, review_summary, }, @@ -1825,6 +1854,7 @@ fn tool_item_event(input: ToolItemEventInput<'_>) -> Option { started_at_ms, completed_at_ms, connection_state, + thread_state, thread_metadata, review_summary, }, @@ -1866,6 +1896,7 @@ fn tool_item_event(input: ToolItemEventInput<'_>) -> Option { started_at_ms, completed_at_ms, connection_state, + thread_state, thread_metadata, review_summary, }, @@ -1910,6 +1941,7 @@ fn tool_item_event(input: ToolItemEventInput<'_>) -> Option { started_at_ms, completed_at_ms, connection_state, + thread_state, thread_metadata, review_summary, }, @@ -1954,6 +1986,7 @@ fn tool_item_event(input: ToolItemEventInput<'_>) -> Option { started_at_ms, completed_at_ms, connection_state, + thread_state, thread_metadata, review_summary, }, @@ -2009,6 +2042,7 @@ fn tool_item_event(input: ToolItemEventInput<'_>) -> Option { started_at_ms, completed_at_ms, connection_state, + thread_state, thread_metadata, review_summary, }, @@ -2045,6 +2079,7 @@ fn tool_item_event(input: ToolItemEventInput<'_>) -> Option { started_at_ms, completed_at_ms, connection_state, + thread_state, thread_metadata, review_summary, }, @@ -2100,6 +2135,7 @@ struct ToolItemContext<'a> { started_at_ms: u64, completed_at_ms: u64, connection_state: &'a ConnectionState, + thread_state: &'a ThreadAnalyticsState, thread_metadata: &'a ThreadMetadataState, review_summary: Option<&'a ItemReviewSummary>, } @@ -2118,7 +2154,9 @@ fn tool_item_base( thread_id: thread_id.to_string(), turn_id: turn_id.to_string(), item_id, - app_server_client: context.connection_state.app_server_client.clone(), + app_server_client: context + .thread_state + .app_server_client(context.connection_state), runtime: context.connection_state.runtime.clone(), thread_source: thread_metadata.thread_source.clone(), subagent_source: thread_metadata.subagent_source.clone(), diff --git a/codex-rs/app-server/src/outgoing_message.rs b/codex-rs/app-server/src/outgoing_message.rs index 9ce815bafe..16d5c2ff9d 100644 --- a/codex-rs/app-server/src/outgoing_message.rs +++ b/codex-rs/app-server/src/outgoing_message.rs @@ -504,13 +504,36 @@ impl OutgoingMessageSender { where T: Into, { - self.send_response_as(request_id, response.into()).await; + self.send_response_as_inner(request_id, response.into(), /*thread_originator*/ None) + .await; + } + + pub(crate) async fn send_response_with_thread_originator( + &self, + request_id: ConnectionRequestId, + response: T, + thread_originator: String, + ) where + T: Into, + { + self.send_response_as_inner(request_id, response.into(), Some(thread_originator)) + .await; } pub(crate) async fn send_response_as( &self, request_id: ConnectionRequestId, response: ClientResponsePayload, + ) { + self.send_response_as_inner(request_id, response, /*thread_originator*/ None) + .await; + } + + async fn send_response_as_inner( + &self, + request_id: ConnectionRequestId, + response: ClientResponsePayload, + thread_originator: Option, ) { let connection_id = request_id.connection_id; let request_id_for_analytics = request_id.request_id.clone(); @@ -518,11 +541,24 @@ impl OutgoingMessageSender { .into_jsonrpc_parts_and_payload(request_id.request_id.clone()) .map(|(id, result, response)| { if let Some(response) = response { - self.analytics_events_client.track_response( - connection_id.0, - request_id_for_analytics, - response, - ); + match thread_originator { + Some(thread_originator) => { + self.analytics_events_client + .track_response_with_thread_originator( + connection_id.0, + request_id_for_analytics, + response, + thread_originator, + ); + } + None => { + self.analytics_events_client.track_response( + connection_id.0, + request_id_for_analytics, + response, + ); + } + } } (id, result) }); diff --git a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs index 03031c87f8..82a49bae0b 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs @@ -640,6 +640,7 @@ pub(super) async fn handle_pending_thread_resume_request( active_permission_profile, workspace_roots, reasoning_effort, + originator, .. } = config_snapshot; let instruction_sources = pending.instruction_sources; @@ -665,7 +666,9 @@ pub(super) async fn handle_pending_thread_resume_request( multi_agent_mode: MultiAgentMode::ExplicitRequestOnly, initial_turns_page, }; - outgoing.send_response(request_id, response).await; + outgoing + .send_response_with_thread_originator(request_id, response, originator) + .await; // Match cold resume: metadata-only resume should attach the listener without // paying the cost of turn reconstruction for historical usage replay. if let Some(token_usage_thread) = token_usage_thread { diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 39283bb35c..8232d68fe6 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -1268,6 +1268,7 @@ impl ThreadRequestProcessor { let cwd = config_snapshot.cwd().clone(); let active_permission_profile = thread_response_active_permission_profile(config_snapshot.active_permission_profile); + let thread_originator = config_snapshot.originator.clone(); let response = ThreadStartResponse { thread: thread.clone(), @@ -1287,7 +1288,7 @@ impl ThreadRequestProcessor { let notif = thread_started_notification(thread); listener_task_context .outgoing - .send_response(request_id, response) + .send_response_with_thread_originator(request_id, response, thread_originator) .instrument(tracing::info_span!( "app_server.thread_start.send_response", otel.name = "app_server.thread_start.send_response", @@ -2863,6 +2864,7 @@ impl ThreadRequestProcessor { } } + let thread_originator = config_snapshot.originator.clone(); let response = ThreadResumeResponse { thread, model: session_configured.model, @@ -2881,7 +2883,9 @@ impl ThreadRequestProcessor { }; let connection_id = request_id.connection_id; - self.outgoing.send_response(request_id, response).await; + self.outgoing + .send_response_with_thread_originator(request_id, response, thread_originator) + .await; // `excludeTurns` is explicitly the cheap resume path, so avoid // rebuilding history only to attribute a replayed usage update. if let Some(token_usage_thread) = token_usage_thread { @@ -3582,6 +3586,7 @@ impl ThreadRequestProcessor { ); let active_permission_profile = thread_response_active_permission_profile(config_snapshot.active_permission_profile); + let thread_originator = config_snapshot.originator.clone(); let response = ThreadForkResponse { thread: thread.clone(), @@ -3602,7 +3607,9 @@ impl ThreadRequestProcessor { let notif = thread_started_notification(thread); let connection_id = request_id.connection_id; let token_usage_thread = include_turns.then(|| response.thread.clone()); - self.outgoing.send_response(request_id, response).await; + self.outgoing + .send_response_with_thread_originator(request_id, response, thread_originator) + .await; // `excludeTurns` is the cheap fork path, so skip restored usage replay // instead of rebuilding history only to attribute a historical update. if let Some(token_usage_thread) = token_usage_thread { diff --git a/codex-rs/app-server/tests/suite/v2/analytics.rs b/codex-rs/app-server/tests/suite/v2/analytics.rs index ea876940ab..2acc86caee 100644 --- a/codex-rs/app-server/tests/suite/v2/analytics.rs +++ b/codex-rs/app-server/tests/suite/v2/analytics.rs @@ -145,7 +145,7 @@ pub(crate) async fn wait_for_goal_event( .await } -async fn wait_for_matching_analytics_event( +pub(crate) async fn wait_for_matching_analytics_event( server: &MockServer, read_timeout: Duration, matches: impl Fn(&Value) -> bool, @@ -191,6 +191,7 @@ pub(crate) fn assert_basic_thread_initialized_event( event: &Value, thread_id: &str, session_id: &str, + expected_product_client_id: &str, expected_model: &str, initialization_mode: &str, expected_thread_source: &str, @@ -199,7 +200,7 @@ pub(crate) fn assert_basic_thread_initialized_event( assert_eq!(event["event_params"]["session_id"], session_id); assert_eq!( event["event_params"]["app_server_client"]["product_client_id"], - DEFAULT_CLIENT_NAME + expected_product_client_id ); assert_eq!( event["event_params"]["app_server_client"]["client_name"], 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 214a678678..4102b45d15 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_fork.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_fork.rs @@ -488,6 +488,7 @@ async fn thread_fork_tracks_thread_initialized_analytics() -> Result<()> { event, &thread.id, &thread.session_id, + "codex", "mock-model", "forked", "user", 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 533087a11e..42db56f526 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -107,6 +107,7 @@ use super::analytics::mount_analytics_capture; use super::analytics::thread_initialized_event; use super::analytics::wait_for_analytics_payload; use super::analytics::wait_for_goal_event; +use super::analytics::wait_for_matching_analytics_event; #[cfg(windows)] const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(25); @@ -573,11 +574,12 @@ async fn thread_resume_tracks_thread_initialized_analytics() -> Result<()> { Some("mock_provider"), /*git_info*/ None, )?; - set_thread_source_on_fake_rollout( + set_session_meta_on_fake_rollout( codex_home.path(), "2025-01-05T12-00-00", &conversation_id, "user", + "codex_work_desktop", )?; let mut mcp = TestAppServer::new_without_managed_config(codex_home.path()).await?; @@ -607,6 +609,7 @@ async fn thread_resume_tracks_thread_initialized_analytics() -> Result<()> { event, &thread.id, &thread.session_id, + "codex_work_desktop", "gpt-5.3-codex", "resumed", "user", @@ -615,11 +618,94 @@ async fn thread_resume_tracks_thread_initialized_analytics() -> Result<()> { Ok(()) } -fn set_thread_source_on_fake_rollout( +#[tokio::test] +async fn thread_resume_running_thread_tracks_thread_originator_in_analytics() -> Result<()> { + let server = create_mock_responses_server_repeating_assistant("Done").await; + + let codex_home = TempDir::new()?; + create_config_toml_with_chatgpt_base_url(codex_home.path(), &server.uri(), &server.uri())?; + mount_analytics_capture(&server, codex_home.path()).await?; + + let mut mcp = TestAppServer::new_without_managed_config(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let start_id = mcp + .send_thread_start_request(ThreadStartParams { + model: Some("mock-model".to_string()), + thread_source: Some(ThreadSource::User), + service_name: Some("codex_work_desktop".to_string()), + ..Default::default() + }) + .await?; + let start_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(start_id)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(start_resp)?; + + let turn_id = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id.clone(), + client_user_message_id: None, + input: vec![UserInput::Text { + text: "materialize rollout".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_id)), + ) + .await??; + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + + let resume_id = mcp + .send_thread_resume_request(ThreadResumeParams { + thread_id: thread.id.clone(), + exclude_turns: true, + ..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: resumed, .. + } = to_response::(resume_resp)?; + + let event = wait_for_matching_analytics_event(&server, DEFAULT_READ_TIMEOUT, |event| { + event["event_type"] == "codex_thread_initialized" + && event["event_params"]["thread_id"] == resumed.id + && event["event_params"]["initialization_mode"] == "resumed" + }) + .await?; + assert_basic_thread_initialized_event( + &event, + &resumed.id, + &resumed.session_id, + "codex_work_desktop", + "mock-model", + "resumed", + "user", + ); + Ok(()) +} + +fn set_session_meta_on_fake_rollout( codex_home: &std::path::Path, filename_ts: &str, thread_id: &str, thread_source: &str, + originator: &str, ) -> Result<()> { let path = rollout_path(codex_home, filename_ts, thread_id); let contents = std::fs::read_to_string(&path)?; @@ -629,6 +715,7 @@ fn set_thread_source_on_fake_rollout( .ok_or_else(|| anyhow::anyhow!("fake rollout missing session meta"))?; let mut session_meta: serde_json::Value = serde_json::from_str(session_meta)?; session_meta["payload"]["thread_source"] = serde_json::json!(thread_source); + session_meta["payload"]["originator"] = serde_json::json!(originator); let remaining = lines.collect::>().join("\n"); std::fs::write(&path, format!("{session_meta}\n{remaining}\n"))?; Ok(()) diff --git a/codex-rs/app-server/tests/suite/v2/thread_start.rs b/codex-rs/app-server/tests/suite/v2/thread_start.rs index 683a2bed18..ac4b255de7 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_start.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_start.rs @@ -655,6 +655,7 @@ async fn thread_start_tracks_thread_initialized_analytics() -> Result<()> { let req_id = mcp .send_thread_start_request(ThreadStartParams { thread_source: Some(ThreadSource::User), + service_name: Some("codex_work_desktop".to_string()), ..Default::default() }) .await?; @@ -672,6 +673,7 @@ async fn thread_start_tracks_thread_initialized_analytics() -> Result<()> { event, &thread.id, &thread.session_id, + "codex_work_desktop", "mock-model", "new", "user", diff --git a/codex-rs/app-server/tests/suite/v2/turn_start.rs b/codex-rs/app-server/tests/suite/v2/turn_start.rs index a528298580..528292ae3f 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_start.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_start.rs @@ -1,6 +1,5 @@ use anyhow::Context; use anyhow::Result; -use app_test_support::DEFAULT_CLIENT_NAME; use app_test_support::TestAppServer; use app_test_support::create_apply_patch_sse_response; use app_test_support::create_exec_command_sse_response; @@ -840,7 +839,7 @@ async fn thread_start_omits_empty_instruction_overrides_from_model_request() -> } #[tokio::test] -async fn turn_start_tracks_turn_event_analytics() -> Result<()> { +async fn turn_start_tracks_thread_originator_in_analytics() -> Result<()> { let server = responses::start_mock_server().await; let response_mock = responses::mount_response_sequence( &server, @@ -875,6 +874,7 @@ async fn turn_start_tracks_turn_event_analytics() -> Result<()> { .send_thread_start_request(ThreadStartParams { model: Some("mock-model".to_string()), thread_source: Some(ThreadSource::User), + service_name: Some("codex_work_desktop".to_string()), ..Default::default() }) .await?; @@ -919,7 +919,7 @@ async fn turn_start_tracks_turn_event_analytics() -> Result<()> { assert_eq!(event["event_params"]["turn_id"], turn.id); assert_eq!( event["event_params"]["app_server_client"]["product_client_id"], - DEFAULT_CLIENT_NAME + "codex_work_desktop" ); assert_eq!(event["event_params"]["model"], "mock-model"); assert_eq!(event["event_params"]["model_provider"], "mock_provider"); From 723b23efd0c06289e878b294b29f6d7ed98a66e8 Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 26 Jun 2026 02:18:00 +0100 Subject: [PATCH 09/38] Reinject missing World State fragments on resume (#30152) ## Why World State restores its structured snapshot on resume so unchanged sections do not have to be rendered again. That is safe only when the model-visible fragment represented by the snapshot is still present in retained history. For selected executor skills, the failing selected-capability scenario exposed this state: ```text persisted World State: selected skill catalog is known retained model history: selected skill catalog message is missing next diff: unchanged, so emit nothing ``` The model resumes without being told about the selected skill catalog. ## What changed World State contributions may now optionally describe the concrete model-visible fragment that must remain in retained history. When a persisted snapshot is present: ```text matching retained fragment exists -> trust snapshot, emit nothing matching retained fragment missing -> treat section as absent, render current state once ``` The skills extension uses this for non-empty selected-environment catalogs by matching its exact rendered catalog body. Empty or hidden catalogs do not require a fragment. ## Scope This does not clear or rebuild the whole World State baseline. It does not change skill discovery, cache invalidation, environment availability, or MCP runtime behavior. It only keeps a persisted section snapshot and its retained model context consistent across resume/history reconstruction. ## Coverage A focused World State regression test verifies both sides: - a missing retained fragment is rendered again - a matching retained fragment avoids duplicate injection --- codex-rs/core/src/context/world_state/mod.rs | 43 +++++++++++- .../context/world_state/world_state_tests.rs | 46 +++++++++++++ .../src/contributors/world_state.rs | 21 ++++++ codex-rs/ext/skills/src/world_state.rs | 66 +++++++++++-------- 4 files changed, 146 insertions(+), 30 deletions(-) diff --git a/codex-rs/core/src/context/world_state/mod.rs b/codex-rs/core/src/context/world_state/mod.rs index 94960a8322..ea82ddd11f 100644 --- a/codex-rs/core/src/context/world_state/mod.rs +++ b/codex-rs/core/src/context/world_state/mod.rs @@ -23,6 +23,10 @@ trait ErasedWorldStateSection: Send + Sync { fn matches_legacy_fragment(&self, role: &str, text: &str) -> bool; + fn has_retained_fragment_matcher(&self) -> bool; + + fn matches_retained_fragment(&self, role: &str, text: &str) -> bool; + fn render_diff( &self, previous: PreviousSectionState<'_, Value>, @@ -57,6 +61,14 @@ impl ErasedWorldStateSection for S { S::matches_legacy_fragment(role, text) } + fn has_retained_fragment_matcher(&self) -> bool { + false + } + + fn matches_retained_fragment(&self, _role: &str, _text: &str) -> bool { + false + } + fn render_diff( &self, previous: PreviousSectionState<'_, Value>, @@ -99,6 +111,14 @@ impl ErasedWorldStateSection for ExtensionWorldStateSection { self.0.matches_legacy_fragment(role, text) } + fn has_retained_fragment_matcher(&self) -> bool { + self.0.has_retained_fragment_matcher() + } + + fn matches_retained_fragment(&self, role: &str, text: &str) -> bool { + self.0.matches_retained_fragment(role, text) + } + fn render_diff( &self, previous: PreviousSectionState<'_, Value>, @@ -268,7 +288,12 @@ impl WorldState { ) -> Vec> { self.render_with(|id, section| { if let Some(previous) = previous.and_then(|previous| previous.sections.get(id)) { - PreviousSectionState::Known(previous) + if section.has_retained_fragment_matcher() && !has_retained_fragment(items, section) + { + PreviousSectionState::Absent + } else { + PreviousSectionState::Known(previous) + } } else if has_legacy_fragment(items, section) { PreviousSectionState::Unknown } else { @@ -288,6 +313,22 @@ impl WorldState { } } +fn has_retained_fragment(items: &[ResponseItem], section: &dyn ErasedWorldStateSection) -> bool { + items.iter().any(|item| { + matches!( + item, + ResponseItem::Message { role, content, .. } + if content.iter().any(|content| { + matches!( + content, + ContentItem::InputText { text } + if section.matches_retained_fragment(role, text) + ) + }) + ) + }) +} + fn has_legacy_fragment(items: &[ResponseItem], section: &dyn ErasedWorldStateSection) -> bool { items.iter().any(|item| { matches!( diff --git a/codex-rs/core/src/context/world_state/world_state_tests.rs b/codex-rs/core/src/context/world_state/world_state_tests.rs index 7ff6f73285..9432672423 100644 --- a/codex-rs/core/src/context/world_state/world_state_tests.rs +++ b/codex-rs/core/src/context/world_state/world_state_tests.rs @@ -149,6 +149,52 @@ fn extension_owned_section_uses_its_snapshot_and_renderer() { ); } +#[test] +fn missing_retained_fragment_is_rendered_again() { + let mut world_state = WorldState::default(); + world_state.add_extension_section( + WorldStateSectionContribution::new( + "extension_test", + json!({"body": "current catalog"}), + |previous| match previous { + PreviousWorldStateSection::Absent => Some(RenderedWorldStateFragment::new( + "developer", + ("", ""), + "current catalog", + )), + PreviousWorldStateSection::Unknown | PreviousWorldStateSection::Known(_) => None, + }, + ) + .with_retained_fragment_matcher(|role, text| { + role == "developer" && text.contains("current catalog") + }), + ); + let previous = world_state.snapshot(); + let retained = ResponseItem::Message { + id: None, + role: "developer".to_string(), + content: vec![ContentItem::InputText { + text: "current catalog".to_string(), + }], + phase: None, + internal_chat_message_metadata_passthrough: None, + }; + + assert_eq!( + world_state + .render_history_diff(Some(&previous), &[]) + .into_iter() + .map(|fragment| fragment.body()) + .collect::>(), + vec!["current catalog"] + ); + assert!( + world_state + .render_history_diff(Some(&previous), &[retained]) + .is_empty() + ); +} + #[test] fn unreadable_section_snapshot_is_treated_as_unknown() { let mut current = WorldState::default(); diff --git a/codex-rs/ext/extension-api/src/contributors/world_state.rs b/codex-rs/ext/extension-api/src/contributors/world_state.rs index 23f180826d..3bafebffa2 100644 --- a/codex-rs/ext/extension-api/src/contributors/world_state.rs +++ b/codex-rs/ext/extension-api/src/contributors/world_state.rs @@ -75,6 +75,7 @@ pub struct WorldStateSectionContribution { snapshot: Value, render_diff: Arc, matches_legacy_fragment: Arc, + matches_retained_fragment: Option>, } impl WorldStateSectionContribution { @@ -93,6 +94,7 @@ impl WorldStateSectionContribution { snapshot, render_diff: Arc::new(render_diff), matches_legacy_fragment: Arc::new(|_, _| false), + matches_retained_fragment: None, } } @@ -104,6 +106,15 @@ impl WorldStateSectionContribution { self } + /// Requires a matching model-visible fragment whenever a persisted snapshot is reused. + pub fn with_retained_fragment_matcher( + mut self, + matcher: impl Fn(&str, &str) -> bool + Send + Sync + 'static, + ) -> Self { + self.matches_retained_fragment = Some(Arc::new(matcher)); + self + } + pub fn id(&self) -> &'static str { self.id } @@ -122,4 +133,14 @@ impl WorldStateSectionContribution { pub fn matches_legacy_fragment(&self, role: &str, text: &str) -> bool { (self.matches_legacy_fragment)(role, text) } + + pub fn has_retained_fragment_matcher(&self) -> bool { + self.matches_retained_fragment.is_some() + } + + pub fn matches_retained_fragment(&self, role: &str, text: &str) -> bool { + self.matches_retained_fragment + .as_ref() + .is_some_and(|matcher| matcher(role, text)) + } } diff --git a/codex-rs/ext/skills/src/world_state.rs b/codex-rs/ext/skills/src/world_state.rs index ce36427b0c..e23eab413f 100644 --- a/codex-rs/ext/skills/src/world_state.rs +++ b/codex-rs/ext/skills/src/world_state.rs @@ -27,36 +27,44 @@ pub(crate) fn executor_skills_world_state_section( "body": body, "includeInstructions": include_instructions, }); + let retained_body = body.clone(); - WorldStateSectionContribution::new(SKILLS_WORLD_STATE_ID, snapshot, move |previous| { - let previous_is_absent = matches!(&previous, PreviousWorldStateSection::Absent); - if let PreviousWorldStateSection::Known(previous) = &previous { - let previous_body = previous.get("body").and_then(serde_json::Value::as_str); - let previous_include_instructions = previous - .get("includeInstructions") - .and_then(serde_json::Value::as_bool); - if previous_body == body.as_deref() - && previous_include_instructions == Some(include_instructions) - { - return None; + let contribution = + WorldStateSectionContribution::new(SKILLS_WORLD_STATE_ID, snapshot, move |previous| { + let previous_is_absent = matches!(&previous, PreviousWorldStateSection::Absent); + if let PreviousWorldStateSection::Known(previous) = &previous { + let previous_body = previous.get("body").and_then(serde_json::Value::as_str); + let previous_include_instructions = previous + .get("includeInstructions") + .and_then(serde_json::Value::as_bool); + if previous_body == body.as_deref() + && previous_include_instructions == Some(include_instructions) + { + return None; + } } - } - let body = match body.as_deref() { - Some(body) => body, - None if previous_is_absent => return None, - None if !include_instructions => HIDDEN_EXECUTOR_SKILLS_BODY, - None => NO_EXECUTOR_SKILLS_BODY, - }; - Some(RenderedWorldStateFragment::new( - "developer", - (SKILLS_INSTRUCTIONS_OPEN_TAG, SKILLS_INSTRUCTIONS_CLOSE_TAG), - body, - )) - }) - .with_legacy_matcher(|role, text| { - role == "developer" - && text.trim_start().starts_with(SKILLS_INSTRUCTIONS_OPEN_TAG) - && text.trim_end().ends_with(SKILLS_INSTRUCTIONS_CLOSE_TAG) - }) + let body = match body.as_deref() { + Some(body) => body, + None if previous_is_absent => return None, + None if !include_instructions => HIDDEN_EXECUTOR_SKILLS_BODY, + None => NO_EXECUTOR_SKILLS_BODY, + }; + Some(RenderedWorldStateFragment::new( + "developer", + (SKILLS_INSTRUCTIONS_OPEN_TAG, SKILLS_INSTRUCTIONS_CLOSE_TAG), + body, + )) + }) + .with_legacy_matcher(|role, text| { + role == "developer" + && text.trim_start().starts_with(SKILLS_INSTRUCTIONS_OPEN_TAG) + && text.trim_end().ends_with(SKILLS_INSTRUCTIONS_CLOSE_TAG) + }); + match retained_body { + Some(body) => contribution.with_retained_fragment_matcher(move |role, text| { + role == "developer" && text.contains(&body) + }), + None => contribution, + } } From fb8598df3ff05366a295c407c91f3d9316a1630c Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 26 Jun 2026 02:28:14 +0100 Subject: [PATCH 10/38] Keep MCP elicitation routable across runtime refreshes (#30127) ## Why An MCP tool call can still be waiting for an elicitation response when an environment update replaces the thread's MCP runtime. Before this change: ```text runtime A starts a tool call and asks the user environment becomes ready, so runtime B is published client answers the prompt through runtime B runtime B cannot find runtime A's pending responder ``` The response is lost and the original tool call stays blocked. ## What changed All MCP runtimes for one thread now share a small elicitation router: ```text runtime A ---\ shared router: response token -> exact pending responder runtime B ---/ ``` When Codex surfaces an MCP elicitation, it assigns a unique opaque response token. The router records which pending request owns that token. A replacement runtime reuses the same router, so the latest runtime can deliver a response to a request started by the previous runtime. The Codex-owned token also prevents two runtime connections that reuse the same MCP server request ID from receiving each other's responses. This does not retain or search old MCP managers. Only the pending responder map is shared. ## Covered scenario The integration test exercises the complete failure mode: 1. A thread starts while its selected environment is still unavailable. 2. A configured MCP server starts a tool call and asks the client for input. 3. The environment becomes ready, causing Codex to publish a replacement MCP runtime. 4. The client answers the original prompt after the replacement. 5. The original tool call receives that answer and completes. A focused routing test also creates two runtimes with the same server request ID and verifies that each response reaches the exact request that emitted its token. ## Scope This PR changes only elicitation response routing across MCP runtime replacement. It does not change when runtimes are rebuilt, which environments contribute MCP configuration, or how environment availability is detected. --- .../app-server/tests/suite/v2/mcp_tool.rs | 233 ++++++++++++++++++ codex-rs/codex-mcp/src/connection_manager.rs | 8 + .../codex-mcp/src/connection_manager_tests.rs | 98 ++++++++ codex-rs/codex-mcp/src/elicitation.rs | 70 ++++-- codex-rs/codex-mcp/src/lib.rs | 1 + codex-rs/codex-mcp/src/mcp/mod.rs | 2 + codex-rs/core/src/connectors.rs | 1 + codex-rs/core/src/mcp_tool_call_tests.rs | 1 + codex-rs/core/src/session/mcp.rs | 9 +- codex-rs/core/src/session/session.rs | 1 + 10 files changed, 398 insertions(+), 26 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/mcp_tool.rs b/codex-rs/app-server/tests/suite/v2/mcp_tool.rs index 3b5afec3fd..f3a3872450 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_tool.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_tool.rs @@ -10,6 +10,8 @@ use app_test_support::create_mock_responses_server_sequence; use app_test_support::to_response; use app_test_support::write_mock_responses_config_toml; use axum::Router; +use codex_app_server_protocol::CapabilityRootLocation; +use codex_app_server_protocol::EnvironmentAddResponse; use codex_app_server_protocol::ItemCompletedNotification; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCResponse; @@ -22,15 +24,21 @@ use codex_app_server_protocol::McpServerToolCallParams; use codex_app_server_protocol::McpServerToolCallResponse; use codex_app_server_protocol::McpToolCallStatus; use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::SelectedCapabilityRoot; use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; +use codex_app_server_protocol::TurnEnvironmentParams; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; +use codex_features::Feature; +use codex_utils_path_uri::PathUri; use codex_utils_pty::DEFAULT_OUTPUT_BYTES_CAP; use core_test_support::responses; +use futures::SinkExt; +use futures::StreamExt; use pretty_assertions::assert_eq; use rmcp::handler::server::ServerHandler; use rmcp::model::BooleanSchema; @@ -56,8 +64,13 @@ use rmcp::transport::streamable_http_server::session::local::LocalSessionManager use serde_json::json; use tempfile::TempDir; use tokio::net::TcpListener; +use tokio::net::TcpStream; +use tokio::sync::oneshot; use tokio::task::JoinHandle; use tokio::time::timeout; +use tokio_tungstenite::WebSocketStream; +use tokio_tungstenite::accept_async; +use tokio_tungstenite::tungstenite::Message; const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10); const TEST_SERVER_NAME: &str = "tool_server"; @@ -68,6 +81,7 @@ const ELICITATION_MESSAGE: &str = "Allow this request?"; const URL_ELICITATION_TRIGGER_MESSAGE: &str = "auth"; const URL_ELICITATION_MESSAGE: &str = "Sign in to GitHub to continue."; const URL_ELICITATION_URL: &str = "https://github.example/login/device"; +const LATE_ENVIRONMENT_ID: &str = "late-environment"; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn mcp_server_tool_call_returns_tool_result() -> Result<()> { @@ -297,6 +311,145 @@ url = "{mcp_server_url}/mcp" Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn mcp_server_elicitation_survives_environment_runtime_refresh() -> Result<()> { + let responses_server = responses::start_mock_server().await; + let (mcp_server_url, mcp_server_handle) = start_mcp_server().await?; + let exec_listener = TcpListener::bind("127.0.0.1:0").await?; + let exec_server_url = format!("ws://{}", exec_listener.local_addr()?); + let codex_home = TempDir::new()?; + write_mock_responses_config_toml( + codex_home.path(), + &responses_server.uri(), + &BTreeMap::from([(Feature::DeferredExecutor, true)]), + /*auto_compact_limit*/ 1024, + /*requires_openai_auth*/ None, + "mock_provider", + "compact", + )?; + let config_path = codex_home.path().join("config.toml"); + let mut config_toml = std::fs::read_to_string(&config_path)?; + config_toml.push_str(&format!( + r#" +[mcp_servers.{TEST_SERVER_NAME}] +url = "{mcp_server_url}/mcp" +"# + )); + std::fs::write(config_path, config_toml)?; + + let mut mcp = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + let add_environment_id = mcp + .send_raw_request( + "environment/add", + Some(json!({ + "environmentId": LATE_ENVIRONMENT_ID, + "execServerUrl": exec_server_url, + "connectTimeoutMs": 10_000, + })), + ) + .await?; + let add_environment_response = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(add_environment_id)), + ) + .await??; + let _: EnvironmentAddResponse = to_response(add_environment_response)?; + + let capability_root = TempDir::new()?; + let thread_start_id = mcp + .send_thread_start_request(ThreadStartParams { + model: Some("mock-model".to_string()), + approval_policy: Some(codex_app_server_protocol::AskForApproval::UnlessTrusted), + environments: Some(vec![TurnEnvironmentParams { + environment_id: LATE_ENVIRONMENT_ID.to_string(), + cwd: codex_utils_absolute_path::AbsolutePathBuf::try_from( + capability_root.path().to_path_buf(), + )? + .into(), + }]), + selected_capability_roots: Some(vec![SelectedCapabilityRoot { + id: "late-plugin@1".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: LATE_ENVIRONMENT_ID.to_string(), + path: PathUri::from_host_native_path(capability_root.path())?, + }, + }]), + ..Default::default() + }) + .await?; + let thread_start_response = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_start_id)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response(thread_start_response)?; + + let tool_call_request_id = mcp + .send_mcp_server_tool_call_request(McpServerToolCallParams { + thread_id: thread.id.clone(), + server: TEST_SERVER_NAME.to_string(), + tool: TEST_TOOL_NAME.to_string(), + arguments: Some(json!({"message": ELICITATION_TRIGGER_MESSAGE})), + meta: None, + }) + .await?; + let server_request = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_request_message(), + ) + .await??; + let ServerRequest::McpServerElicitationRequest { request_id, .. } = server_request else { + panic!("expected MCP elicitation request, got: {server_request:?}"); + }; + + let (filesystem_request_tx, filesystem_request_rx) = oneshot::channel(); + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let exec_server_handle = tokio::spawn(serve_environment_until_shutdown( + exec_listener, + filesystem_request_tx, + shutdown_rx, + )); + let mut filesystem_request_rx = filesystem_request_rx; + timeout(DEFAULT_READ_TIMEOUT, async { + loop { + let status_request_id = mcp + .send_raw_request("mcpServerStatus/list", Some(json!({"threadId": thread.id}))) + .await?; + mcp.read_stream_until_response_message(RequestId::Integer(status_request_id)) + .await?; + if filesystem_request_rx.try_recv().is_ok() { + return Ok::<_, anyhow::Error>(()); + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await??; + + mcp.send_response( + request_id, + serde_json::to_value(McpServerElicitationRequestResponse { + action: McpServerElicitationAction::Accept, + content: Some(json!({"confirmed": true})), + meta: None, + })?, + ) + .await?; + let tool_call_response = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(tool_call_request_id)), + ) + .await??; + let response: McpServerToolCallResponse = to_response(tool_call_response)?; + assert_eq!(response.content[0].get("text"), Some(&json!("accepted"))); + + let _ = shutdown_tx.send(()); + exec_server_handle.await??; + mcp_server_handle.abort(); + let _ = mcp_server_handle.await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn mcp_server_tool_call_forwards_url_elicitation() -> Result<()> { let responses_server = responses::start_mock_server().await; @@ -683,6 +836,86 @@ async fn start_mcp_server() -> Result<(String, JoinHandle<()>)> { Ok((format!("http://{addr}"), handle)) } +async fn serve_environment_until_shutdown( + listener: TcpListener, + filesystem_request_tx: oneshot::Sender<()>, + mut shutdown_rx: oneshot::Receiver<()>, +) -> Result<()> { + let (stream, _) = listener.accept().await?; + let mut websocket = accept_async(stream).await?; + + let initialize = read_exec_server_json(&mut websocket).await?; + assert_eq!(initialize["method"], "initialize"); + websocket + .send(Message::Text( + json!({ + "id": initialize["id"], + "result": {"sessionId": "test-session"}, + }) + .to_string() + .into(), + )) + .await?; + let initialized = read_exec_server_json(&mut websocket).await?; + assert_eq!(initialized["method"], "initialized"); + let environment_info = read_exec_server_json(&mut websocket).await?; + assert_eq!(environment_info["method"], "environment/info"); + websocket + .send(Message::Text( + json!({ + "id": environment_info["id"], + "result": {"shell": {"name": "zsh", "path": "/bin/zsh"}}, + }) + .to_string() + .into(), + )) + .await?; + + let mut filesystem_request_tx = Some(filesystem_request_tx); + loop { + let request = tokio::select! { + request = read_exec_server_json(&mut websocket) => request?, + _ = &mut shutdown_rx => return Ok(()), + }; + if request["method"] + .as_str() + .is_some_and(|method| method.starts_with("fs/")) + && let Some(tx) = filesystem_request_tx.take() + { + let _ = tx.send(()); + } + if request.get("id").is_some() { + websocket + .send(Message::Text( + json!({ + "id": request["id"], + "error": {"code": -32004, "message": "not found"}, + }) + .to_string() + .into(), + )) + .await?; + } + } +} + +async fn read_exec_server_json( + websocket: &mut WebSocketStream, +) -> Result { + loop { + match websocket + .next() + .await + .ok_or_else(|| anyhow::anyhow!("exec-server websocket closed"))?? + { + Message::Text(text) => return Ok(serde_json::from_str(text.as_ref())?), + Message::Binary(bytes) => return Ok(serde_json::from_slice(bytes.as_ref())?), + Message::Ping(_) | Message::Pong(_) => {} + message => anyhow::bail!("expected JSON-RPC message, got {message:?}"), + } + } +} + async fn wait_for_mcp_tool_call_completed( mcp: &mut TestAppServer, call_id: &str, diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 5902bc9607..3107f678b0 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -18,6 +18,7 @@ use crate::codex_apps_cache::CodexAppsToolsCache; use crate::codex_apps_cache::CodexAppsToolsCacheKey; use crate::codex_apps_cache::CodexAppsToolsFetchSource; use crate::elicitation::ElicitationRequestManager; +use crate::elicitation::ElicitationRequestRouter; use crate::elicitation::ElicitationReviewerHandle; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::mcp::ToolPluginProvenance; @@ -141,6 +142,7 @@ impl McpConnectionManager { tool_plugin_provenance: ToolPluginProvenance, auth: Option<&CodexAuth>, elicitation_reviewer: Option, + elicitation_router: ElicitationRequestRouter, ) -> Self { let mut required_servers = mcp_servers .iter() @@ -155,6 +157,7 @@ impl McpConnectionManager { approval_policy.value(), initial_permission_profile, elicitation_reviewer, + elicitation_router, ); let tool_plugin_provenance = Arc::new(tool_plugin_provenance); let startup_submit_id = submit_id.clone(); @@ -352,6 +355,7 @@ impl McpConnectionManager { approval_policy.value(), permission_profile.clone(), /*reviewer*/ None, + ElicitationRequestRouter::default(), ), startup_cancellation_token: CancellationToken::new(), } @@ -444,6 +448,10 @@ impl McpConnectionManager { self.elicitation_requests.set_auto_deny(auto_deny); } + pub fn elicitation_router(&self) -> ElicitationRequestRouter { + self.elicitation_requests.router() + } + pub async fn resolve_elicitation( &self, server_name: String, diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index e11fcb5917..d2869fe5c1 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -3,6 +3,7 @@ use crate::codex_apps_cache::CodexAppsToolsCache; use crate::codex_apps_cache::CodexAppsToolsCacheContext; use crate::declared_openai_file_input_param_names; use crate::elicitation::ElicitationRequestManager; +use crate::elicitation::ElicitationRequestRouter; use crate::elicitation::elicitation_is_rejected_by_policy; use crate::rmcp_client::AsyncManagedClient; use crate::rmcp_client::ManagedClient; @@ -274,6 +275,7 @@ async fn disabled_permissions_auto_accept_elicitation_with_empty_form_schema() { AskForApproval::Never, PermissionProfile::Disabled, /*reviewer*/ None, + ElicitationRequestRouter::default(), ); let (tx_event, _rx_event) = async_channel::bounded(1); let sender = manager.make_sender("server".to_string(), tx_event); @@ -309,6 +311,7 @@ async fn disabled_permissions_do_not_auto_accept_elicitation_with_requested_fiel AskForApproval::Never, PermissionProfile::Disabled, /*reviewer*/ None, + ElicitationRequestRouter::default(), ); let (tx_event, _rx_event) = async_channel::bounded(1); let sender = manager.make_sender("server".to_string(), tx_event); @@ -342,6 +345,100 @@ async fn disabled_permissions_do_not_auto_accept_elicitation_with_requested_fiel ); } +#[tokio::test] +async fn shared_elicitation_router_targets_the_exact_pending_request() { + let router = ElicitationRequestRouter::default(); + let manager_a = ElicitationRequestManager::new( + AskForApproval::OnRequest, + PermissionProfile::default(), + /*reviewer*/ None, + router.clone(), + ); + let manager_b = ElicitationRequestManager::new( + AskForApproval::OnRequest, + PermissionProfile::default(), + /*reviewer*/ None, + router, + ); + let (tx_event, rx_event) = async_channel::bounded(2); + let sender_a = manager_a.make_sender("server".to_string(), tx_event.clone()); + let sender_b = manager_b.make_sender("server".to_string(), tx_event); + let elicitation = codex_rmcp_client::Elicitation::Mcp( + CreateElicitationRequestParams::FormElicitationParams { + meta: None, + message: "Which runtime?".to_string(), + requested_schema: rmcp::model::ElicitationSchema::builder() + .required_property( + "runtime", + rmcp::model::PrimitiveSchema::String(rmcp::model::StringSchema::new()), + ) + .build() + .expect("schema should build"), + }, + ); + + let pending_a = tokio::spawn(sender_a(NumberOrString::Number(1), elicitation.clone())); + let EventMsg::ElicitationRequest(request_a) = rx_event.recv().await.expect("request A").msg + else { + panic!("expected elicitation request"); + }; + let pending_b = tokio::spawn(sender_b(NumberOrString::Number(1), elicitation)); + let EventMsg::ElicitationRequest(request_b) = rx_event.recv().await.expect("request B").msg + else { + panic!("expected elicitation request"); + }; + let ( + codex_protocol::mcp::RequestId::String(request_a_id), + codex_protocol::mcp::RequestId::String(request_b_id), + ) = (request_a.id, request_b.id) + else { + panic!("expected Codex-owned string request IDs"); + }; + assert_ne!(request_a_id, request_b_id); + + let response_a = ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(serde_json::json!({"runtime": "a"})), + meta: None, + }; + manager_b + .resolve( + "server".to_string(), + NumberOrString::String(request_a_id.into()), + response_a.clone(), + ) + .await + .expect("runtime B should route a response to runtime A"); + let response_b = ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(serde_json::json!({"runtime": "b"})), + meta: None, + }; + manager_a + .resolve( + "server".to_string(), + NumberOrString::String(request_b_id.into()), + response_b.clone(), + ) + .await + .expect("runtime A should route a response to runtime B"); + + assert_eq!( + pending_a + .await + .expect("request A task") + .expect("request A response"), + response_a + ); + assert_eq!( + pending_b + .await + .expect("request B task") + .expect("request B response"), + response_b + ); +} + #[test] fn test_normalize_tools_short_non_duplicated_names() { let tools = vec![ @@ -1158,6 +1255,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { ToolPluginProvenance::default(), /*auth*/ None, /*elicitation_reviewer*/ None, + ElicitationRequestRouter::default(), ) .await; diff --git a/codex-rs/codex-mcp/src/elicitation.rs b/codex-rs/codex-mcp/src/elicitation.rs index d6a685b502..e4fcf49d22 100644 --- a/codex-rs/codex-mcp/src/elicitation.rs +++ b/codex-rs/codex-mcp/src/elicitation.rs @@ -8,6 +8,8 @@ use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex as StdMutex; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; use crate::mcp::McpPermissionPromptAutoApproveContext; use crate::mcp::mcp_permission_prompt_is_auto_approved; @@ -32,6 +34,8 @@ use rmcp::model::RequestId; use tokio::sync::Mutex; use tokio::sync::oneshot; +static NEXT_ELICITATION_REQUEST_ID: AtomicU64 = AtomicU64::new(0); + #[derive(Debug, Clone)] pub struct ElicitationReviewRequest { pub server_name: String, @@ -48,9 +52,36 @@ pub trait ElicitationReviewer: Send + Sync { pub type ElicitationReviewerHandle = Arc; +/// Routes model-visible elicitation response tokens to their exact pending responders. +/// +/// One router is shared by every MCP runtime created for a thread. The public response token is +/// generated by Codex rather than copied from the MCP connection, so separate runtimes may reuse +/// the same server request ID without colliding. +#[derive(Clone, Default)] +pub struct ElicitationRequestRouter { + requests: Arc>, +} + +impl ElicitationRequestRouter { + async fn resolve( + &self, + server_name: String, + id: RequestId, + response: ElicitationResponse, + ) -> Result<()> { + self.requests + .lock() + .await + .remove(&(server_name, id)) + .ok_or_else(|| anyhow!("elicitation request not found"))? + .send(response) + .map_err(|e| anyhow!("failed to send elicitation response: {e:?}")) + } +} + #[derive(Clone)] pub(crate) struct ElicitationRequestManager { - requests: Arc>, + router: ElicitationRequestRouter, pub(crate) approval_policy: Arc>, pub(crate) permission_profile: Arc>, auto_deny: Arc>, @@ -62,9 +93,10 @@ impl ElicitationRequestManager { approval_policy: AskForApproval, permission_profile: PermissionProfile, reviewer: Option, + router: ElicitationRequestRouter, ) -> Self { Self { - requests: Arc::new(Mutex::new(HashMap::new())), + router, approval_policy: Arc::new(StdMutex::new(approval_policy)), permission_profile: Arc::new(StdMutex::new(permission_profile)), auto_deny: Arc::new(StdMutex::new(false)), @@ -91,13 +123,11 @@ impl ElicitationRequestManager { id: RequestId, response: ElicitationResponse, ) -> Result<()> { - self.requests - .lock() - .await - .remove(&(server_name, id)) - .ok_or_else(|| anyhow!("elicitation request not found"))? - .send(response) - .map_err(|e| anyhow!("failed to send elicitation response: {e:?}")) + self.router.resolve(server_name, id, response).await + } + + pub(crate) fn router(&self) -> ElicitationRequestRouter { + self.router.clone() } pub(crate) fn make_sender( @@ -105,13 +135,13 @@ impl ElicitationRequestManager { server_name: String, tx_event: Sender, ) -> SendElicitation { - let elicitation_requests = self.requests.clone(); + let router = self.router.clone(); let approval_policy = self.approval_policy.clone(); let permission_profile = self.permission_profile.clone(); let auto_deny = self.auto_deny.clone(); let reviewer = self.reviewer.clone(); Box::new(move |id, elicitation| { - let elicitation_requests = elicitation_requests.clone(); + let router = router.clone(); let tx_event = tx_event.clone(); let server_name = server_name.clone(); let approval_policy = approval_policy.clone(); @@ -171,6 +201,11 @@ impl ElicitationRequestManager { } } + let public_request_id = format!( + "codex-mcp-elicitation-{}", + NEXT_ELICITATION_REQUEST_ID.fetch_add(1, Ordering::Relaxed) + ); + let routed_request_id = RequestId::String(public_request_id.clone().into()); let request = match elicitation { Elicitation::Mcp( rmcp::model::CreateElicitationRequestParams::FormElicitationParams { @@ -215,8 +250,8 @@ impl ElicitationRequestManager { }; let (tx, rx) = oneshot::channel(); { - let mut lock = elicitation_requests.lock().await; - lock.insert((server_name.clone(), id.clone()), tx); + let mut lock = router.requests.lock().await; + lock.insert((server_name.clone(), routed_request_id), tx); } let _ = tx_event .send(Event { @@ -224,14 +259,7 @@ impl ElicitationRequestManager { msg: EventMsg::ElicitationRequest(ElicitationRequestEvent { turn_id: None, server_name, - id: match id.clone() { - rmcp::model::NumberOrString::String(value) => { - ProtocolRequestId::String(value.to_string()) - } - rmcp::model::NumberOrString::Number(value) => { - ProtocolRequestId::Integer(value) - } - }, + id: ProtocolRequestId::String(public_request_id), request, }), }) diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index e611623d44..eda7ac2399 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -1,5 +1,6 @@ pub use connection_manager::McpConnectionManager; pub use connection_manager::tool_is_model_visible; +pub use elicitation::ElicitationRequestRouter; pub use elicitation::ElicitationReviewRequest; pub use elicitation::ElicitationReviewer; pub use elicitation::ElicitationReviewerHandle; diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 3f24a276e4..ffb797f26e 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -339,6 +339,7 @@ pub async fn read_mcp_resource( tool_plugin_provenance(config), auth, /*elicitation_reviewer*/ None, + crate::elicitation::ElicitationRequestRouter::default(), ) .await; @@ -415,6 +416,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( tool_plugin_provenance, auth, /*elicitation_reviewer*/ None, + crate::elicitation::ElicitationRequestRouter::default(), ) .await; diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 97f24e043f..b45cc841a1 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -283,6 +283,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( ToolPluginProvenance::default(), auth.as_ref(), /*elicitation_reviewer*/ None, + codex_mcp::ElicitationRequestRouter::default(), ) .await; diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 8c8d82037b..271141613b 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1416,6 +1416,7 @@ async fn host_owned_codex_apps_manager( codex_mcp::ToolPluginProvenance::default(), auth.as_ref(), /*elicitation_reviewer*/ None, + codex_mcp::ElicitationRequestRouter::default(), ) .await; Arc::new(manager) diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 50921f1591..83a59c960a 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -316,6 +316,7 @@ impl Session { *guard = cancellation_token.clone(); cancellation_token }; + let current_runtime = self.services.latest_mcp_runtime(); let refreshed_manager = McpConnectionManager::new( &mcp_servers, mcp_config.mcp_oauth_credentials_store_mode, @@ -338,13 +339,11 @@ impl Session { tool_plugin_provenance, auth.as_ref(), elicitation_reviewer, + current_runtime.manager().elicitation_router(), ) .await; - { - let current_manager = self.services.latest_mcp_runtime(); - refreshed_manager - .set_elicitations_auto_deny(current_manager.manager().elicitations_auto_deny()); - } + refreshed_manager + .set_elicitations_auto_deny(current_runtime.manager().elicitations_auto_deny()); self.services.publish_mcp_runtime( mcp_config, mcp_runtime_context, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 0e1609207a..ed07352b15 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1207,6 +1207,7 @@ impl Session { tool_plugin_provenance, auth, Some(sess.mcp_elicitation_reviewer()), + codex_mcp::ElicitationRequestRouter::default(), ) .instrument(info_span!( "session_init.mcp_manager_init", From ec300bc7bd7355833f579a603df067d59c1963cb Mon Sep 17 00:00:00 2001 From: Martin Au-Yeung Date: Thu, 25 Jun 2026 18:31:10 -0700 Subject: [PATCH 11/38] Expose MCP app identity in app context (#29934) ## Why MCP tool-call events need to expose trusted app identity and action metadata directly so v2 clients do not have to infer it from tool names or resource URIs. ## What changed - Add optional `appName`, `templateId`, and `actionName` fields to MCP tool-call `appContext`. - Populate `appName` and `templateId` from trusted Codex Apps metadata, and derive `actionName` from the trusted app resource metadata. - Preserve all three fields through core events, legacy protocol events, persisted thread history, resume redaction, and app-server v2 responses. - Document the public `appContext` fields in `codex-rs/app-server/README.md`. - Regenerate app-server JSON and TypeScript schemas and add coverage for serialization, persistence, redaction, and metadata propagation. ## Validation - `just test -p codex-app-server-protocol mcp_tool_call` - `just test -p codex-core mcp_tool_call_item_metadata_only_trusts_codex_apps_identity mcp_tool_call_item_includes_app_identity` - `just write-app-server-schema` --------- Co-authored-by: Martin Au-Yeung <280153141+martinauyeung-oai@users.noreply.github.com> --- .../schema/json/ServerNotification.json | 18 +++++++++ .../codex_app_server_protocol.schemas.json | 18 +++++++++ .../codex_app_server_protocol.v2.schemas.json | 18 +++++++++ .../json/v2/ItemCompletedNotification.json | 18 +++++++++ .../json/v2/ItemStartedNotification.json | 18 +++++++++ .../schema/json/v2/ReviewStartResponse.json | 18 +++++++++ .../schema/json/v2/ThreadForkResponse.json | 18 +++++++++ .../schema/json/v2/ThreadListResponse.json | 18 +++++++++ .../json/v2/ThreadMetadataUpdateResponse.json | 18 +++++++++ .../schema/json/v2/ThreadReadResponse.json | 18 +++++++++ .../schema/json/v2/ThreadResumeResponse.json | 18 +++++++++ .../json/v2/ThreadRollbackResponse.json | 18 +++++++++ .../schema/json/v2/ThreadStartResponse.json | 18 +++++++++ .../json/v2/ThreadStartedNotification.json | 18 +++++++++ .../json/v2/ThreadUnarchiveResponse.json | 18 +++++++++ .../json/v2/TurnCompletedNotification.json | 18 +++++++++ .../schema/json/v2/TurnStartResponse.json | 18 +++++++++ .../json/v2/TurnStartedNotification.json | 18 +++++++++ .../typescript/v2/McpToolCallAppContext.ts | 2 +- .../src/protocol/thread_history.rs | 15 ++++++++ .../src/protocol/v2/item.rs | 6 +++ .../src/protocol/v2/tests.rs | 38 +++++++++++++++++++ codex-rs/app-server/README.md | 2 +- .../thread_resume_redaction.rs | 6 +++ .../tests/suite/v2/thread_resume.rs | 9 +++++ codex-rs/core/src/mcp_tool_call.rs | 26 +++++++++++++ codex-rs/core/src/mcp_tool_call_tests.rs | 36 +++++++++++++++++- .../core/src/tools/handlers/mcp_resource.rs | 6 +++ codex-rs/core/tests/suite/rmcp_client.rs | 3 ++ codex-rs/core/tests/suite/search_tool.rs | 4 ++ codex-rs/protocol/src/items.rs | 15 ++++++++ codex-rs/protocol/src/protocol.rs | 28 ++++++++++++++ 32 files changed, 516 insertions(+), 4 deletions(-) diff --git a/codex-rs/app-server-protocol/schema/json/ServerNotification.json b/codex-rs/app-server-protocol/schema/json/ServerNotification.json index 8a8064c0b9..43fb661822 100644 --- a/codex-rs/app-server-protocol/schema/json/ServerNotification.json +++ b/codex-rs/app-server-protocol/schema/json/ServerNotification.json @@ -2438,6 +2438,18 @@ }, "McpToolCallAppContext": { "properties": { + "actionName": { + "type": [ + "string", + "null" + ] + }, + "appName": { + "type": [ + "string", + "null" + ] + }, "connectorId": { "type": "string" }, @@ -2452,6 +2464,12 @@ "string", "null" ] + }, + "templateId": { + "type": [ + "string", + "null" + ] } }, "required": [ 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 56ea39276a..277421a852 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 @@ -12248,6 +12248,18 @@ }, "McpToolCallAppContext": { "properties": { + "actionName": { + "type": [ + "string", + "null" + ] + }, + "appName": { + "type": [ + "string", + "null" + ] + }, "connectorId": { "type": "string" }, @@ -12262,6 +12274,12 @@ "string", "null" ] + }, + "templateId": { + "type": [ + "string", + "null" + ] } }, "required": [ 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 0f2910befd..39c95053d0 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 @@ -8652,6 +8652,18 @@ }, "McpToolCallAppContext": { "properties": { + "actionName": { + "type": [ + "string", + "null" + ] + }, + "appName": { + "type": [ + "string", + "null" + ] + }, "connectorId": { "type": "string" }, @@ -8666,6 +8678,12 @@ "string", "null" ] + }, + "templateId": { + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/ItemCompletedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/ItemCompletedNotification.json index 8f4b533d8f..dcf6312c7d 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ItemCompletedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ItemCompletedNotification.json @@ -299,6 +299,18 @@ }, "McpToolCallAppContext": { "properties": { + "actionName": { + "type": [ + "string", + "null" + ] + }, + "appName": { + "type": [ + "string", + "null" + ] + }, "connectorId": { "type": "string" }, @@ -313,6 +325,12 @@ "string", "null" ] + }, + "templateId": { + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/ItemStartedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/ItemStartedNotification.json index 4c930afde1..1933225ade 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ItemStartedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ItemStartedNotification.json @@ -299,6 +299,18 @@ }, "McpToolCallAppContext": { "properties": { + "actionName": { + "type": [ + "string", + "null" + ] + }, + "appName": { + "type": [ + "string", + "null" + ] + }, "connectorId": { "type": "string" }, @@ -313,6 +325,12 @@ "string", "null" ] + }, + "templateId": { + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/ReviewStartResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ReviewStartResponse.json index 16da9f205a..6851ddb487 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ReviewStartResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ReviewStartResponse.json @@ -437,6 +437,18 @@ }, "McpToolCallAppContext": { "properties": { + "actionName": { + "type": [ + "string", + "null" + ] + }, + "appName": { + "type": [ + "string", + "null" + ] + }, "connectorId": { "type": "string" }, @@ -451,6 +463,12 @@ "string", "null" ] + }, + "templateId": { + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadForkResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadForkResponse.json index 37a78533c2..3ca62f17f2 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadForkResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadForkResponse.json @@ -541,6 +541,18 @@ }, "McpToolCallAppContext": { "properties": { + "actionName": { + "type": [ + "string", + "null" + ] + }, + "appName": { + "type": [ + "string", + "null" + ] + }, "connectorId": { "type": "string" }, @@ -555,6 +567,12 @@ "string", "null" ] + }, + "templateId": { + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadListResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadListResponse.json index 67726c13a6..a81f643122 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadListResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadListResponse.json @@ -463,6 +463,18 @@ }, "McpToolCallAppContext": { "properties": { + "actionName": { + "type": [ + "string", + "null" + ] + }, + "appName": { + "type": [ + "string", + "null" + ] + }, "connectorId": { "type": "string" }, @@ -477,6 +489,12 @@ "string", "null" ] + }, + "templateId": { + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadMetadataUpdateResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadMetadataUpdateResponse.json index cdd74a735a..d007b32d22 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadMetadataUpdateResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadMetadataUpdateResponse.json @@ -463,6 +463,18 @@ }, "McpToolCallAppContext": { "properties": { + "actionName": { + "type": [ + "string", + "null" + ] + }, + "appName": { + "type": [ + "string", + "null" + ] + }, "connectorId": { "type": "string" }, @@ -477,6 +489,12 @@ "string", "null" ] + }, + "templateId": { + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadReadResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadReadResponse.json index c3e388d1fd..f45efe3707 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadReadResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadReadResponse.json @@ -463,6 +463,18 @@ }, "McpToolCallAppContext": { "properties": { + "actionName": { + "type": [ + "string", + "null" + ] + }, + "appName": { + "type": [ + "string", + "null" + ] + }, "connectorId": { "type": "string" }, @@ -477,6 +489,12 @@ "string", "null" ] + }, + "templateId": { + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeResponse.json index ad44bf406b..09dd3fedbf 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeResponse.json @@ -541,6 +541,18 @@ }, "McpToolCallAppContext": { "properties": { + "actionName": { + "type": [ + "string", + "null" + ] + }, + "appName": { + "type": [ + "string", + "null" + ] + }, "connectorId": { "type": "string" }, @@ -555,6 +567,12 @@ "string", "null" ] + }, + "templateId": { + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackResponse.json index 5ae94d8d11..8083b52d49 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackResponse.json @@ -463,6 +463,18 @@ }, "McpToolCallAppContext": { "properties": { + "actionName": { + "type": [ + "string", + "null" + ] + }, + "appName": { + "type": [ + "string", + "null" + ] + }, "connectorId": { "type": "string" }, @@ -477,6 +489,12 @@ "string", "null" ] + }, + "templateId": { + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartResponse.json index df70c08c43..b33a64a200 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartResponse.json @@ -541,6 +541,18 @@ }, "McpToolCallAppContext": { "properties": { + "actionName": { + "type": [ + "string", + "null" + ] + }, + "appName": { + "type": [ + "string", + "null" + ] + }, "connectorId": { "type": "string" }, @@ -555,6 +567,12 @@ "string", "null" ] + }, + "templateId": { + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartedNotification.json index 157f6193b4..9040457f22 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartedNotification.json @@ -463,6 +463,18 @@ }, "McpToolCallAppContext": { "properties": { + "actionName": { + "type": [ + "string", + "null" + ] + }, + "appName": { + "type": [ + "string", + "null" + ] + }, "connectorId": { "type": "string" }, @@ -477,6 +489,12 @@ "string", "null" ] + }, + "templateId": { + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadUnarchiveResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadUnarchiveResponse.json index 265578dcdd..d4e9fea618 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadUnarchiveResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadUnarchiveResponse.json @@ -463,6 +463,18 @@ }, "McpToolCallAppContext": { "properties": { + "actionName": { + "type": [ + "string", + "null" + ] + }, + "appName": { + "type": [ + "string", + "null" + ] + }, "connectorId": { "type": "string" }, @@ -477,6 +489,12 @@ "string", "null" ] + }, + "templateId": { + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/TurnCompletedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/TurnCompletedNotification.json index ae1b07cfc9..6ffb343595 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/TurnCompletedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/TurnCompletedNotification.json @@ -437,6 +437,18 @@ }, "McpToolCallAppContext": { "properties": { + "actionName": { + "type": [ + "string", + "null" + ] + }, + "appName": { + "type": [ + "string", + "null" + ] + }, "connectorId": { "type": "string" }, @@ -451,6 +463,12 @@ "string", "null" ] + }, + "templateId": { + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/TurnStartResponse.json b/codex-rs/app-server-protocol/schema/json/v2/TurnStartResponse.json index 5853612dd9..805a5099a2 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/TurnStartResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/TurnStartResponse.json @@ -437,6 +437,18 @@ }, "McpToolCallAppContext": { "properties": { + "actionName": { + "type": [ + "string", + "null" + ] + }, + "appName": { + "type": [ + "string", + "null" + ] + }, "connectorId": { "type": "string" }, @@ -451,6 +463,12 @@ "string", "null" ] + }, + "templateId": { + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/json/v2/TurnStartedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/TurnStartedNotification.json index 7db58289ee..1f3f7c7036 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/TurnStartedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/TurnStartedNotification.json @@ -437,6 +437,18 @@ }, "McpToolCallAppContext": { "properties": { + "actionName": { + "type": [ + "string", + "null" + ] + }, + "appName": { + "type": [ + "string", + "null" + ] + }, "connectorId": { "type": "string" }, @@ -451,6 +463,12 @@ "string", "null" ] + }, + "templateId": { + "type": [ + "string", + "null" + ] } }, "required": [ diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/McpToolCallAppContext.ts b/codex-rs/app-server-protocol/schema/typescript/v2/McpToolCallAppContext.ts index 0a8343428d..e4bc5a11d9 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/McpToolCallAppContext.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/McpToolCallAppContext.ts @@ -2,4 +2,4 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. -export type McpToolCallAppContext = { connectorId: string, linkId: string | null, resourceUri: string | null, }; +export type McpToolCallAppContext = { connectorId: string, linkId: string | null, resourceUri: string | null, appName: string | null, templateId: string | null, actionName: string | null, }; 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 98702267ab..aed87b18ab 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -775,6 +775,9 @@ impl ThreadHistoryBuilder { connector_id, link_id: payload.link_id.clone(), resource_uri: payload.mcp_app_resource_uri.clone(), + app_name: payload.app_name.clone(), + template_id: payload.template_id.clone(), + action_name: payload.action_name.clone(), }), mcp_app_resource_uri: payload.mcp_app_resource_uri.clone(), plugin_id: payload.plugin_id.clone(), @@ -825,6 +828,9 @@ impl ThreadHistoryBuilder { connector_id, link_id: payload.link_id.clone(), resource_uri: payload.mcp_app_resource_uri.clone(), + app_name: payload.app_name.clone(), + template_id: payload.template_id.clone(), + action_name: payload.action_name.clone(), }), mcp_app_resource_uri: payload.mcp_app_resource_uri.clone(), plugin_id: payload.plugin_id.clone(), @@ -2418,6 +2424,9 @@ mod tests { connector_id: None, mcp_app_resource_uri: None, link_id: None, + app_name: None, + template_id: None, + action_name: None, plugin_id: None, duration: Duration::from_millis(8), result: Err("boom".into()), @@ -2499,6 +2508,9 @@ mod tests { connector_id: Some("calendar".into()), mcp_app_resource_uri: Some("ui://widget/lookup.html".into()), link_id: Some("link_calendar".into()), + app_name: Some("Calendar".into()), + template_id: Some("calendar_template".into()), + action_name: Some("lookup".into()), plugin_id: Some("sample@test".into()), duration: Duration::from_millis(8), result: Ok(CallToolResult { @@ -2533,6 +2545,9 @@ mod tests { connector_id: "calendar".into(), link_id: Some("link_calendar".into()), resource_uri: Some("ui://widget/lookup.html".into()), + app_name: Some("Calendar".into()), + template_id: Some("calendar_template".into()), + action_name: Some("lookup".into()), }), mcp_app_resource_uri: Some("ui://widget/lookup.html".into()), plugin_id: Some("sample@test".into()), diff --git a/codex-rs/app-server-protocol/src/protocol/v2/item.rs b/codex-rs/app-server-protocol/src/protocol/v2/item.rs index e2a0d45129..751eab824d 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/item.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/item.rs @@ -396,6 +396,9 @@ pub struct McpToolCallAppContext { pub connector_id: String, pub link_id: Option, pub resource_uri: Option, + pub app_name: Option, + pub template_id: Option, + pub action_name: Option, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] @@ -895,6 +898,9 @@ impl From for ThreadItem { connector_id, link_id: mcp.link_id, resource_uri: mcp.mcp_app_resource_uri.clone(), + app_name: mcp.app_name, + template_id: mcp.template_id, + action_name: mcp.action_name, }), mcp_app_resource_uri: mcp.mcp_app_resource_uri, plugin_id: mcp.plugin_id, diff --git a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs index 03c2bcaa5b..c6d460c721 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs @@ -2677,6 +2677,9 @@ fn core_turn_item_into_thread_item_converts_supported_variants() { connector_id: Some("calendar".to_string()), mcp_app_resource_uri: Some("app://connector".to_string()), link_id: Some("link_calendar".to_string()), + app_name: Some("Calendar".to_string()), + template_id: Some("calendar_template".to_string()), + action_name: Some("create_event".to_string()), plugin_id: Some("sample@test".to_string()), status: CoreMcpToolCallStatus::InProgress, result: None, @@ -2696,6 +2699,9 @@ fn core_turn_item_into_thread_item_converts_supported_variants() { connector_id: "calendar".to_string(), link_id: Some("link_calendar".to_string()), resource_uri: Some("app://connector".to_string()), + app_name: Some("Calendar".to_string()), + template_id: Some("calendar_template".to_string()), + action_name: Some("create_event".to_string()), }), mcp_app_resource_uri: Some("app://connector".to_string()), plugin_id: Some("sample@test".to_string()), @@ -2713,6 +2719,9 @@ fn core_turn_item_into_thread_item_converts_supported_variants() { connector_id: None, mcp_app_resource_uri: None, link_id: None, + app_name: None, + template_id: None, + action_name: None, plugin_id: None, status: CoreMcpToolCallStatus::Completed, result: Some(CallToolResult { @@ -2759,6 +2768,9 @@ fn mcp_tool_call_app_context_serializes_connector_id() { connector_id: "calendar".to_string(), link_id: Some("link_calendar".to_string()), resource_uri: Some("app://connector".to_string()), + app_name: Some("Calendar".to_string()), + template_id: Some("calendar_template".to_string()), + action_name: Some("create_event".to_string()), }), mcp_app_resource_uri: Some("app://connector".to_string()), plugin_id: None, @@ -2780,6 +2792,9 @@ fn mcp_tool_call_app_context_serializes_connector_id() { "connectorId": "calendar", "linkId": "link_calendar", "resourceUri": "app://connector", + "appName": "Calendar", + "templateId": "calendar_template", + "actionName": "create_event", }, "mcpAppResourceUri": "app://connector", "pluginId": null, @@ -2790,6 +2805,29 @@ fn mcp_tool_call_app_context_serializes_connector_id() { ); } +#[test] +fn mcp_tool_call_app_context_serializes_missing_mixed_version_fields_as_null() { + assert_eq!( + serde_json::to_value(McpToolCallAppContext { + connector_id: "calendar".to_string(), + link_id: None, + resource_uri: None, + app_name: None, + template_id: None, + action_name: None, + }) + .expect("MCP tool call app context should serialize"), + json!({ + "connectorId": "calendar", + "linkId": null, + "resourceUri": null, + "appName": null, + "templateId": null, + "actionName": null, + }) + ); +} + #[test] fn user_input_into_core_preserves_image_detail() { assert_eq!( diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index a4b488acdf..d291f5c85b 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -1372,7 +1372,7 @@ Today both notifications carry an empty `items` array even when item events were - `reasoning` — `{id, summary, content}` where `summary` holds streamed reasoning summaries (applicable for most OpenAI models) and `content` holds raw reasoning blocks (applicable for e.g. open source models). - `commandExecution` — `{id, command, cwd, status, commandActions, aggregatedOutput?, exitCode?, durationMs?}` for sandboxed commands; `status` is `inProgress`, `completed`, `failed`, or `declined`. - `fileChange` — `{id, changes, status}` describing proposed edits; `changes` list `{path, kind, diff}` and `status` is `inProgress`, `completed`, `failed`, or `declined`. -- `mcpToolCall` — `{id, server, tool, status, arguments, appContext, mcpAppResourceUri?, pluginId, result?, error?}` describing MCP calls; `appContext` is `{connectorId, linkId, resourceUri}` for calls through a trusted MCP app, where `connectorId` identifies the connector that owns the tool, `linkId` identifies the app link, and `resourceUri` points to the widget template. The top-level `mcpAppResourceUri` is deprecated and temporarily duplicated for client migration. `tool` identifies the invoked action. `status` is `inProgress`, `completed`, or `failed`. +- `mcpToolCall` — `{id, server, tool, status, arguments, appContext, mcpAppResourceUri?, pluginId, result?, error?}` describing MCP calls; `appContext` is `{connectorId, linkId, resourceUri, appName, templateId, actionName}` for calls through a trusted MCP app, where `connectorId` identifies the connector that owns the tool, `linkId` identifies the app link, `resourceUri` points to the widget template, `appName` is the connector's display name, `templateId` identifies the app template, and `actionName` is the stable connector `Action.name`. `appName`, `templateId`, and `actionName` may be null for older rollout entries. The top-level `mcpAppResourceUri` is deprecated and temporarily duplicated for client migration. `tool` identifies the raw MCP tool. `status` is `inProgress`, `completed`, or `failed`. - `collabToolCall` — `{id, tool, status, senderThreadId, receiverThreadId?, newThreadId?, prompt?, agentStatus?}` describing collab tool calls (`spawn_agent`, `send_input`, `resume_agent`, `wait`, `close_agent`); `status` is `inProgress`, `completed`, or `failed`. - `webSearch` — `{id, query, action?}` for a web search request issued by the agent; `action` mirrors the Responses API web_search action payload (`search`, `open_page`, `find_in_page`) and may be omitted until completion. - `imageView` — `{id, path}` emitted when the agent invokes the image viewer tool. diff --git a/codex-rs/app-server/src/request_processors/thread_resume_redaction.rs b/codex-rs/app-server/src/request_processors/thread_resume_redaction.rs index 586828e8dd..f753607273 100644 --- a/codex-rs/app-server/src/request_processors/thread_resume_redaction.rs +++ b/codex-rs/app-server/src/request_processors/thread_resume_redaction.rs @@ -83,6 +83,9 @@ mod tests { connector_id: "calendar".to_string(), link_id: Some("link_calendar".to_string()), resource_uri: Some("ui://widget/lookup.html".to_string()), + app_name: Some("Calendar".to_string()), + template_id: Some("calendar_template".to_string()), + action_name: Some("lookup".to_string()), }), mcp_app_resource_uri: Some("ui://widget/lookup.html".to_string()), plugin_id: Some("sample@test".to_string()), @@ -130,6 +133,9 @@ mod tests { connector_id: "calendar".to_string(), link_id: Some("link_calendar".to_string()), resource_uri: Some("ui://widget/lookup.html".to_string()), + app_name: Some("Calendar".to_string()), + template_id: Some("calendar_template".to_string()), + action_name: Some("lookup".to_string()), }), mcp_app_resource_uri: Some("ui://widget/lookup.html".to_string()), plugin_id: Some("sample@test".to_string()), 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 42db56f526..67dd249217 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -834,6 +834,9 @@ async fn thread_resume_redacts_payloads_for_chatgpt_remote_clients() -> Result<( connector_id: "calendar".to_string(), link_id: Some("link_calendar".to_string()), resource_uri: Some("ui://widget/lookup.html".to_string()), + app_name: Some("Calendar".to_string()), + template_id: Some("calendar_template".to_string()), + action_name: Some("lookup".to_string()), }) ); let result = result.as_ref().expect("redacted MCP result"); @@ -884,6 +887,9 @@ async fn thread_resume_redacts_payloads_for_chatgpt_remote_clients() -> Result<( connector_id: "calendar".to_string(), link_id: Some("link_calendar".to_string()), resource_uri: Some("ui://widget/lookup.html".to_string()), + app_name: Some("Calendar".to_string()), + template_id: Some("calendar_template".to_string()), + action_name: Some("lookup".to_string()), }) ); let result = result.as_ref().expect("normal MCP result"); @@ -990,6 +996,9 @@ fn append_resume_redaction_history( connector_id: Some("calendar".to_string()), mcp_app_resource_uri: Some("ui://widget/lookup.html".to_string()), link_id: Some("link_calendar".to_string()), + app_name: Some("Calendar".to_string()), + template_id: Some("calendar_template".to_string()), + action_name: Some("lookup".to_string()), plugin_id: None, duration: Duration::from_millis(8), result: Ok(CallToolResult { diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 19b7c1a6ad..d03f232d12 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -326,6 +326,9 @@ struct McpToolCallItemMetadata { connector_id: Option, link_id: Option, mcp_app_resource_uri: Option, + app_name: Option, + template_id: Option, + action_name: Option, plugin_id: Option, } @@ -342,6 +345,15 @@ impl McpToolCallItemMetadata { link_id: trusted_mcp_app_metadata.and_then(|metadata| metadata.link_id.clone()), mcp_app_resource_uri: metadata .and_then(|metadata| metadata.mcp_app_resource_uri.clone()), + app_name: trusted_mcp_app_metadata.and_then(|metadata| metadata.connector_name.clone()), + template_id: trusted_mcp_app_metadata.and_then(|metadata| metadata.template_id.clone()), + action_name: trusted_mcp_app_metadata + .and_then(|metadata| metadata.codex_apps_meta.as_ref()) + .and_then(|meta| meta.get(MCP_TOOL_RESOURCE_URI_META_KEY)) + .and_then(serde_json::Value::as_str) + .and_then(|resource_uri| resource_uri.trim_matches('/').rsplit('/').next()) + .filter(|action_name| !action_name.is_empty()) + .map(str::to_string), plugin_id: metadata.and_then(|metadata| metadata.plugin_id.clone()), } } @@ -881,6 +893,9 @@ async fn notify_mcp_tool_call_started( connector_id: item_metadata.connector_id, mcp_app_resource_uri: item_metadata.mcp_app_resource_uri, link_id: item_metadata.link_id, + app_name: item_metadata.app_name, + template_id: item_metadata.template_id, + action_name: item_metadata.action_name, plugin_id: item_metadata.plugin_id, status: McpToolCallStatus::InProgress, result: None, @@ -923,6 +938,9 @@ async fn notify_mcp_tool_call_completed( connector_id: item_metadata.connector_id, mcp_app_resource_uri: item_metadata.mcp_app_resource_uri, link_id: item_metadata.link_id, + app_name: item_metadata.app_name, + template_id: item_metadata.template_id, + action_name: item_metadata.action_name, plugin_id: item_metadata.plugin_id, status, result, @@ -999,6 +1017,7 @@ pub(crate) struct McpToolApprovalMetadata { tool_title: Option, tool_description: Option, mcp_app_resource_uri: Option, + template_id: Option, codex_apps_meta: Option>, openai_file_input_params: Option>, } @@ -1009,6 +1028,8 @@ const MCP_TOOL_LINK_ID_META_KEY: &str = "link_id"; const MCP_TOOL_PLUGIN_ID_META_KEY: &str = "plugin_id"; const MCP_TOOL_THREAD_ID_META_KEY: &str = "threadId"; const MCP_TOOL_CONNECTED_ACCOUNT_EMAIL_META_KEY: &str = "connected_account_email"; +const MCP_TOOL_TEMPLATE_ID_META_KEY: &str = "template_id"; +const MCP_TOOL_RESOURCE_URI_META_KEY: &str = "resource_uri"; async fn custom_mcp_tool_approval_mode( sess: &Session, @@ -1531,6 +1552,11 @@ pub(crate) async fn lookup_mcp_tool_metadata( tool_title: tool_info.tool.title, tool_description: tool_info.tool.description.map(std::borrow::Cow::into_owned), mcp_app_resource_uri: get_mcp_app_resource_uri(tool_info.tool.meta.as_deref()), + template_id: codex_apps_meta + .as_ref() + .and_then(|meta| meta.get(MCP_TOOL_TEMPLATE_ID_META_KEY)) + .and_then(serde_json::Value::as_str) + .map(str::to_string), codex_apps_meta, // Disallow custom MCPs from uploading files via fileParams. openai_file_input_params: openai_file_input_params_for_server( diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 271141613b..6b642b1cd0 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -87,6 +87,7 @@ fn approval_metadata( tool_title: tool_title.map(str::to_string), tool_description: tool_description.map(str::to_string), mcp_app_resource_uri: None, + template_id: None, codex_apps_meta: None, openai_file_input_params: None, } @@ -1195,12 +1196,22 @@ async fn plugin_mcp_tool_call_request_meta_includes_plugin_id() { fn mcp_tool_call_item_metadata_only_trusts_codex_apps_identity() { let mut metadata = approval_metadata( Some("asdk_app_0123456789abcdef0123456789abcdef"), - /*connector_name*/ None, + Some("Calendar"), /*connector_description*/ None, - /*tool_title*/ None, + Some("Create a calendar event"), /*tool_description*/ None, ); metadata.link_id = Some("link_fedcba9876543210fedcba9876543210".to_string()); + metadata.template_id = Some("calendar_template".to_string()); + metadata.codex_apps_meta = Some( + serde_json::json!({ + "resource_uri": "/asdk_app_0123456789abcdef0123456789abcdef/link_fedcba9876543210fedcba9876543210/create_event", + "template_id": "calendar_template", + }) + .as_object() + .cloned() + .expect("_codex_apps metadata should be an object"), + ); assert_eq!( McpToolCallItemMetadata::from_tool_metadata(CODEX_APPS_MCP_SERVER_NAME, Some(&metadata),), @@ -1208,6 +1219,9 @@ fn mcp_tool_call_item_metadata_only_trusts_codex_apps_identity() { connector_id: Some("asdk_app_0123456789abcdef0123456789abcdef".to_string()), link_id: Some("link_fedcba9876543210fedcba9876543210".to_string()), mcp_app_resource_uri: None, + app_name: Some("Calendar".to_string()), + template_id: Some("calendar_template".to_string()), + action_name: Some("create_event".to_string()), plugin_id: None, } ); @@ -1217,6 +1231,9 @@ fn mcp_tool_call_item_metadata_only_trusts_codex_apps_identity() { connector_id: None, link_id: None, mcp_app_resource_uri: None, + app_name: None, + template_id: None, + action_name: None, plugin_id: None, } ); @@ -1239,6 +1256,9 @@ async fn mcp_tool_call_item_includes_app_identity() { connector_id: Some("asdk_app_0123456789abcdef0123456789abcdef".to_string()), link_id: Some("link_fedcba9876543210fedcba9876543210".to_string()), mcp_app_resource_uri: None, + app_name: Some("Calendar".to_string()), + template_id: Some("calendar_template".to_string()), + action_name: Some("create_event".to_string()), plugin_id: Some("sample@test".to_string()), }, ) @@ -1264,6 +1284,8 @@ async fn mcp_tool_call_item_includes_app_identity() { Some("link_fedcba9876543210fedcba9876543210") ); assert_eq!(item.plugin_id.as_deref(), Some("sample@test")); + assert_eq!(item.app_name.as_deref(), Some("Calendar")); + assert_eq!(item.action_name.as_deref(), Some("create_event")); } #[tokio::test] @@ -1284,6 +1306,7 @@ async fn codex_apps_tool_call_request_meta_includes_turn_metadata_and_codex_apps tool_title: Some("Create Event".to_string()), tool_description: Some("Create a calendar event.".to_string()), mcp_app_resource_uri: None, + template_id: None, codex_apps_meta: Some( serde_json::json!({ "resource_uri": "connector://calendar/tools/calendar_create_event", @@ -1771,6 +1794,7 @@ fn guardian_mcp_review_request_includes_annotations_when_present() { tool_title: None, tool_description: None, mcp_app_resource_uri: None, + template_id: None, codex_apps_meta: None, openai_file_input_params: None, }; @@ -2473,6 +2497,7 @@ async fn approve_mode_skips_when_annotations_do_not_require_approval() { tool_title: Some("Read Only Tool".to_string()), tool_description: None, mcp_app_resource_uri: None, + template_id: None, codex_apps_meta: None, openai_file_input_params: None, }; @@ -2549,6 +2574,7 @@ async fn guardian_mode_skips_auto_when_annotations_do_not_require_approval() { tool_title: Some("Read Only Tool".to_string()), tool_description: None, mcp_app_resource_uri: None, + template_id: None, codex_apps_meta: None, openai_file_input_params: None, }; @@ -2608,6 +2634,7 @@ async fn permission_request_hook_allows_mcp_tool_call() { tool_title: Some("Create entities".to_string()), tool_description: None, mcp_app_resource_uri: None, + template_id: None, codex_apps_meta: None, openai_file_input_params: None, }; @@ -2746,6 +2773,7 @@ async fn permission_request_hook_runs_after_remembered_mcp_approval() { tool_title: Some("Create entities".to_string()), tool_description: None, mcp_app_resource_uri: None, + template_id: None, codex_apps_meta: None, openai_file_input_params: None, }; @@ -2835,6 +2863,7 @@ async fn guardian_mode_mcp_denial_returns_rationale_message() { tool_title: Some("Dangerous Tool".to_string()), tool_description: Some("Reads calendar data.".to_string()), mcp_app_resource_uri: None, + template_id: None, codex_apps_meta: None, openai_file_input_params: None, }; @@ -2891,6 +2920,7 @@ async fn prompt_mode_waits_for_approval_when_annotations_do_not_require_approval tool_title: Some("Read Only Tool".to_string()), tool_description: None, mcp_app_resource_uri: None, + template_id: None, codex_apps_meta: None, openai_file_input_params: None, }; @@ -2948,6 +2978,7 @@ async fn full_access_mode_skips_mcp_tool_approval_for_all_approval_modes() { tool_title: Some("Dangerous Tool".to_string()), tool_description: Some("Performs a risky action.".to_string()), mcp_app_resource_uri: None, + template_id: None, codex_apps_meta: None, openai_file_input_params: None, }; @@ -3003,6 +3034,7 @@ async fn approve_mode_skips_guardian_in_every_permission_mode() { tool_title: Some("Dangerous Tool".to_string()), tool_description: Some("Performs a risky action.".to_string()), mcp_app_resource_uri: None, + template_id: None, codex_apps_meta: None, openai_file_input_params: None, }; diff --git a/codex-rs/core/src/tools/handlers/mcp_resource.rs b/codex-rs/core/src/tools/handlers/mcp_resource.rs index 231bdf92a9..b073b5ee75 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource.rs @@ -224,6 +224,9 @@ async fn emit_tool_call_begin( connector_id: None, mcp_app_resource_uri: None, link_id: None, + app_name: None, + template_id: None, + action_name: None, plugin_id: None, status: McpToolCallStatus::InProgress, result: None, @@ -265,6 +268,9 @@ async fn emit_tool_call_end( connector_id: None, mcp_app_resource_uri: None, link_id: None, + app_name: None, + template_id: None, + action_name: None, plugin_id: None, status, result, diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index fc1a7df18b..fae4ab6416 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -1359,6 +1359,9 @@ async fn stdio_image_responses_round_trip() -> anyhow::Result<()> { connector_id: None, mcp_app_resource_uri: None, link_id: None, + app_name: None, + template_id: None, + action_name: None, plugin_id: None, }, ); diff --git a/codex-rs/core/tests/suite/search_tool.rs b/codex-rs/core/tests/suite/search_tool.rs index edfc53fdc6..afb2e0d085 100644 --- a/codex-rs/core/tests/suite/search_tool.rs +++ b/codex-rs/core/tests/suite/search_tool.rs @@ -545,6 +545,8 @@ async fn tool_search_returns_deferred_tools_without_follow_up_tool_injection() - unreachable!("event guard guarantees McpToolCallBegin"); }; assert_eq!(begin.call_id, "calendar-call-1"); + assert_eq!(begin.app_name.as_deref(), Some("Calendar")); + assert_eq!(begin.action_name.as_deref(), Some("calendar_create_event")); assert_eq!( begin.mcp_app_resource_uri.as_deref(), Some(CALENDAR_CREATE_EVENT_MCP_APP_RESOURCE_URI) @@ -559,6 +561,8 @@ async fn tool_search_returns_deferred_tools_without_follow_up_tool_injection() - }; assert_eq!(end.call_id, "calendar-call-1"); assert_eq!(end.connector_id.as_deref(), Some("calendar")); + assert_eq!(end.app_name.as_deref(), Some("Calendar")); + assert_eq!(end.action_name.as_deref(), Some("calendar_create_event")); assert_eq!( end.mcp_app_resource_uri.as_deref(), Some(CALENDAR_CREATE_EVENT_MCP_APP_RESOURCE_URI) diff --git a/codex-rs/protocol/src/items.rs b/codex-rs/protocol/src/items.rs index 508d71d4be..c4ee26d96d 100644 --- a/codex-rs/protocol/src/items.rs +++ b/codex-rs/protocol/src/items.rs @@ -202,6 +202,15 @@ pub struct McpToolCallItem { pub link_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] #[ts(optional)] + pub app_name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub template_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub action_name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional)] pub plugin_id: Option, pub status: McpToolCallStatus, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -560,6 +569,9 @@ impl McpToolCallItem { connector_id: self.connector_id.clone(), mcp_app_resource_uri: self.mcp_app_resource_uri.clone(), link_id: self.link_id.clone(), + app_name: self.app_name.clone(), + template_id: self.template_id.clone(), + action_name: self.action_name.clone(), plugin_id: self.plugin_id.clone(), }) } @@ -581,6 +593,9 @@ impl McpToolCallItem { mcp_app_resource_uri: self.mcp_app_resource_uri.clone(), connector_id: self.connector_id.clone(), link_id: self.link_id.clone(), + app_name: self.app_name.clone(), + template_id: self.template_id.clone(), + action_name: self.action_name.clone(), plugin_id: self.plugin_id.clone(), duration: self.duration?, result, diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 5f3e75dd83..491d64aa58 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2357,6 +2357,15 @@ pub struct McpToolCallBeginEvent { pub link_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] #[ts(optional)] + pub app_name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub template_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub action_name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional)] pub plugin_id: Option, } @@ -2376,6 +2385,15 @@ pub struct McpToolCallEndEvent { pub link_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] #[ts(optional)] + pub app_name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub template_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub action_name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional)] pub plugin_id: Option, #[ts(type = "string")] pub duration: Duration, @@ -5140,6 +5158,9 @@ mod tests { connector_id: Some("connector".into()), mcp_app_resource_uri: Some("app://connector".into()), link_id: Some("link_123".into()), + app_name: Some("Calendar".into()), + template_id: Some("calendar_template".into()), + action_name: Some("create_event".into()), plugin_id: Some("sample@test".into()), status: McpToolCallStatus::InProgress, result: None, @@ -5161,6 +5182,8 @@ mod tests { Some("app://connector") ); assert_eq!(event.link_id.as_deref(), Some("link_123")); + assert_eq!(event.app_name.as_deref(), Some("Calendar")); + assert_eq!(event.action_name.as_deref(), Some("create_event")); assert_eq!(event.plugin_id.as_deref(), Some("sample@test")); } _ => panic!("expected McpToolCallBegin event"), @@ -5251,6 +5274,9 @@ mod tests { connector_id: Some("connector".into()), mcp_app_resource_uri: Some("app://connector".into()), link_id: Some("link_123".into()), + app_name: Some("Calendar".into()), + template_id: Some("calendar_template".into()), + action_name: Some("create_event".into()), plugin_id: Some("sample@test".into()), status: McpToolCallStatus::Completed, result: Some(CallToolResult { @@ -5277,6 +5303,8 @@ mod tests { Some("app://connector") ); assert_eq!(event.link_id.as_deref(), Some("link_123")); + assert_eq!(event.app_name.as_deref(), Some("Calendar")); + assert_eq!(event.action_name.as_deref(), Some("create_event")); assert_eq!(event.plugin_id.as_deref(), Some("sample@test")); assert_eq!(event.duration, Duration::from_millis(42)); assert!(event.is_success()); From 0d4351c1b80cae3a40a2d0d64d4e87ca1d5c4afa Mon Sep 17 00:00:00 2001 From: Won Park Date: Thu, 25 Jun 2026 18:34:35 -0700 Subject: [PATCH 12/38] [codex] allow CCA image generation and web search extensions (#29909) ## Summary - allow the standalone image-generation and web-search extensions for the actor-authorized provider shape used by CCA - preserve builtin `image_generation` and `web_search` for older models and existing flows - keep ordinary non-OpenAI providers excluded from both extensions - remove only the image extension local managed-AuthManager requirement that CCA cannot satisfy - share actor-authorization detection through `ModelProviderInfo` - keep Core tests focused on routing behavior and cover header-shape edge cases in `model-provider-info` - add a Responses Lite regression that verifies both `image_gen.imagegen` and `web.run` ## Why CCA uses a provider named `local` with `requires_openai_auth: false` and a non-empty `x-openai-actor-authorization` header. Core accepts that provider shape, but both extension provider-name gates rejected it; image generation additionally required a Codex-managed login. The standalone paths must coexist with existing builtin tools. New Responses Lite models can receive `image_gen.imagegen` and `web.run`, while older models continue using builtin tools. ## Impact This enables both standalone extensions for CCA once installed downstream, without removing or changing builtin-tool compatibility for older models. ## Validation - `just test -p codex-core responses_lite_exposes_standalone_tools_for_actor_authorized_provider` - `just test -p codex-core responses_lite_uses_standalone_web_search_and_image_generation` - `just test -p codex-core hosted_tools_follow_provider_auth_model_and_config_gates` - `just test -p codex-image-generation-extension` - `just test -p codex-web-search-extension` - `just test -p codex-model-provider-info` - `just fmt` - `git diff --check` --- codex-rs/core/src/tools/spec_plan.rs | 16 +--- codex-rs/core/src/tools/spec_plan_tests.rs | 95 ++++--------------- codex-rs/core/tests/suite/responses_lite.rs | 60 +++++++++++- .../ext/image-generation/src/extension.rs | 12 +-- codex-rs/ext/web-search/src/extension.rs | 3 +- codex-rs/model-provider-info/src/lib.rs | 11 +++ .../src/model_provider_info_tests.rs | 25 +++++ 7 files changed, 125 insertions(+), 97 deletions(-) diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index b873f87e27..8ca5eec898 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -98,7 +98,6 @@ use tracing::warn; const MULTI_AGENT_V2_NAMESPACE_DESCRIPTION: &str = "Tools for spawning and managing sub-agents."; const IMAGE_GEN_NAMESPACE: &str = "image_gen"; const IMAGEGEN_TOOL_NAME: &str = "imagegen"; -const ACTOR_AUTHORIZATION_HEADER: &str = "x-openai-actor-authorization"; type PlannedRuntime = Arc; @@ -387,7 +386,10 @@ fn image_generation_tool_enabled(turn_context: &TurnContext) -> bool { } fn image_generation_runtime_enabled(turn_context: &TurnContext) -> bool { - (provider_uses_actor_authorization(turn_context) + (turn_context + .provider + .info() + .uses_openai_actor_authorization() || (turn_context.provider.info().requires_openai_auth && turn_context .auth_manager @@ -400,16 +402,6 @@ fn image_generation_runtime_enabled(turn_context: &TurnContext) -> bool { .contains(&InputModality::Image) } -fn provider_uses_actor_authorization(turn_context: &TurnContext) -> bool { - let provider_info = turn_context.provider.info(); - !provider_info.requires_openai_auth - && provider_info.http_headers.as_ref().is_some_and(|headers| { - headers.iter().any(|(name, value)| { - name.eq_ignore_ascii_case(ACTOR_AUTHORIZATION_HEADER) && !value.trim().is_empty() - }) - }) -} - fn standalone_image_generation_model_visible(turn_context: &TurnContext) -> bool { if !image_generation_runtime_enabled(turn_context) || !namespace_tools_enabled(turn_context) { return false; diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 498a690c58..d95cb5ff44 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -275,22 +275,19 @@ fn use_bedrock_provider(turn: &mut TurnContext) { turn.provider = create_model_provider(provider_info, turn.auth_manager.clone()); } -fn use_provider_auth( - turn: &mut TurnContext, - requires_openai_auth: bool, - actor_header: Option<(&str, &str)>, -) { +fn use_actor_authorized_provider(turn: &mut TurnContext) { let mut provider_info = turn.config.model_provider.clone(); - provider_info.requires_openai_auth = requires_openai_auth; - provider_info.http_headers = actor_header.map(|(name, value)| { - HashMap::from([ - (name.to_string(), value.to_string()), - ( - "ChatGPT-Account-ID".to_string(), - "test-account-id".to_string(), - ), - ]) - }); + provider_info.requires_openai_auth = false; + provider_info.http_headers = Some(HashMap::from([ + ( + "x-openai-actor-authorization".to_string(), + "test-actor-authorization".to_string(), + ), + ( + "ChatGPT-Account-ID".to_string(), + "test-account-id".to_string(), + ), + ])); turn.auth_manager = None; update_config(turn, |config| { config.model_provider = provider_info.clone(); @@ -1552,47 +1549,17 @@ async fn hosted_tools_follow_provider_auth_model_and_config_gates() { .await; unrelated_chatgpt_auth.assert_visible_lacks(&["image_generation"]); - let provider_without_actor_auth = probe(|turn| { + let actor_authorized_provider = probe(|turn| { set_feature(turn, Feature::ImageGeneration, /*enabled*/ true); - use_provider_auth( - turn, /*requires_openai_auth*/ false, /*actor_header*/ None, - ); + use_actor_authorized_provider(turn); turn.model_info.input_modalities = vec![InputModality::Image]; }) .await; - provider_without_actor_auth.assert_visible_lacks(&["image_generation"]); - - let provider_authenticated = probe(|turn| { - set_feature(turn, Feature::ImageGeneration, /*enabled*/ true); - use_provider_auth( - turn, - /*requires_openai_auth*/ false, - Some(("X-OpenAI-Actor-Authorization", "test-actor-authorization")), - ); - turn.model_info.input_modalities = vec![InputModality::Image]; - }) - .await; - provider_authenticated.assert_visible_contains(&["image_generation"]); - - let empty_actor_auth = probe(|turn| { - set_feature(turn, Feature::ImageGeneration, /*enabled*/ true); - use_provider_auth( - turn, - /*requires_openai_auth*/ false, - Some(("x-openai-actor-authorization", " ")), - ); - turn.model_info.input_modalities = vec![InputModality::Image]; - }) - .await; - empty_actor_auth.assert_visible_lacks(&["image_generation"]); + actor_authorized_provider.assert_visible_contains(&["image_generation"]); let feature_disabled = probe(|turn| { set_feature(turn, Feature::ImageGeneration, /*enabled*/ false); - use_provider_auth( - turn, - /*requires_openai_auth*/ false, - Some(("x-openai-actor-authorization", "test-actor-authorization")), - ); + use_actor_authorized_provider(turn); turn.model_info.input_modalities = vec![InputModality::Image]; }) .await; @@ -1600,11 +1567,7 @@ async fn hosted_tools_follow_provider_auth_model_and_config_gates() { let text_only_model = probe(|turn| { set_feature(turn, Feature::ImageGeneration, /*enabled*/ true); - use_provider_auth( - turn, - /*requires_openai_auth*/ false, - Some(("x-openai-actor-authorization", "test-actor-authorization")), - ); + use_actor_authorized_provider(turn); turn.model_info.input_modalities = vec![]; }) .await; @@ -1613,34 +1576,12 @@ async fn hosted_tools_follow_provider_auth_model_and_config_gates() { let unsupported_image_generation_provider = probe(|turn| { set_feature(turn, Feature::ImageGeneration, /*enabled*/ true); use_bedrock_provider(turn); - let mut provider_info = turn.provider.info().clone(); - provider_info.requires_openai_auth = false; - provider_info.http_headers = Some(HashMap::from([( - "x-openai-actor-authorization".to_string(), - "test-actor-authorization".to_string(), - )])); - turn.auth_manager = None; - update_config(turn, |config| { - config.model_provider = provider_info.clone(); - }); - turn.provider = create_model_provider(provider_info, /*auth_manager*/ None); + use_actor_authorized_provider(turn); turn.model_info.input_modalities = vec![InputModality::Image]; }) .await; unsupported_image_generation_provider.assert_visible_lacks(&["image_generation"]); - let codex_managed_auth_provider = probe(|turn| { - set_feature(turn, Feature::ImageGeneration, /*enabled*/ true); - use_provider_auth( - turn, - /*requires_openai_auth*/ true, - Some(("x-openai-actor-authorization", "test-actor-authorization")), - ); - turn.model_info.input_modalities = vec![InputModality::Image]; - }) - .await; - codex_managed_auth_provider.assert_visible_lacks(&["image_generation"]); - let image_generation = probe(|turn| { use_chatgpt_auth(turn); set_feature(turn, Feature::ImageGeneration, /*enabled*/ true); diff --git a/codex-rs/core/tests/suite/responses_lite.rs b/codex-rs/core/tests/suite/responses_lite.rs index 8074344111..5f623bd4c6 100644 --- a/codex-rs/core/tests/suite/responses_lite.rs +++ b/codex-rs/core/tests/suite/responses_lite.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::sync::Arc; use anyhow::Context; @@ -56,6 +57,18 @@ fn has_hosted_tool(tools: &[Value], tool_type: &str) -> bool { .any(|tool| tool.get("type").and_then(Value::as_str) == Some(tool_type)) } +fn has_namespaced_tool(tools: &[Value], namespace: &str, tool_name: &str) -> bool { + tools.iter().any(|tool| { + tool.get("type").and_then(Value::as_str) == Some("namespace") + && tool.get("name").and_then(Value::as_str) == Some(namespace) + && tool["tools"].as_array().is_some_and(|tools| { + tools + .iter() + .any(|tool| tool.get("name").and_then(Value::as_str) == Some(tool_name)) + }) + }) +} + fn additional_tools(body: &Value) -> Result<&[Value]> { body["input"] .as_array() @@ -227,13 +240,58 @@ async fn responses_lite_uses_standalone_web_search_and_image_generation() -> Res let body = request.body_json(); assert!(body.get("tools").is_none()); let tools = additional_tools(&body)?; - assert!(!tools.is_empty()); + assert!(has_namespaced_tool(tools, "web", "run")); + assert!(has_namespaced_tool(tools, "image_gen", "imagegen")); assert!(!has_hosted_tool(tools, "web_search")); assert!(!has_hosted_tool(tools, "image_generation")); Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn responses_lite_exposes_standalone_tools_for_actor_authorized_provider() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let response_mock = responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_completed("resp-1"), + ]), + ) + .await; + + let auth = CodexAuth::from_api_key("dummy"); + let extensions = responses_extensions(&auth); + let mut builder = test_codex() + .with_auth(auth) + .with_extensions(extensions) + .with_model_info_override("gpt-5.4", |model_info| { + model_info.use_responses_lite = true; + configure_image_capable_model(model_info); + }) + .with_config(|config| { + configure_responses_tools(config); + config.model_provider.name = "local".to_string(); + config.model_provider.requires_openai_auth = false; + config.model_provider.http_headers = Some(HashMap::from([( + "x-openai-actor-authorization".to_string(), + "test-actor-authorization".to_string(), + )])); + }); + let test = builder.build(&server).await?; + + test.submit_turn("Use standalone tools").await?; + + let body = response_mock.single_request().body_json(); + let tools = additional_tools(&body)?; + assert!(has_namespaced_tool(tools, "web", "run")); + assert!(has_namespaced_tool(tools, "image_gen", "imagegen")); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn responses_lite_compact_request_uses_lite_transport_contract() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/ext/image-generation/src/extension.rs b/codex-rs/ext/image-generation/src/extension.rs index 8195c08e61..1e3dc3e127 100644 --- a/codex-rs/ext/image-generation/src/extension.rs +++ b/codex-rs/ext/image-generation/src/extension.rs @@ -34,11 +34,11 @@ struct ImageGenerationExtensionConfig { } impl ImageGenerationExtensionConfig { - /// Resolves whether standalone image generation should be available for a thread. + /// Resolves the image provider and save root for a thread. fn from_config(config: &Config, resolve_save_root: &SaveRootResolver) -> Self { Self { - // Core selects this executor per turn using the feature flag or model metadata. - available: config.model_provider.is_openai(), + available: config.model_provider.is_openai() + || config.model_provider.uses_openai_actor_authorization(), provider: config.model_provider.clone(), save_root: resolve_save_root(config), } @@ -46,7 +46,7 @@ impl ImageGenerationExtensionConfig { } impl ThreadLifecycleContributor for ImageGenerationExtension { - /// Seeds image-generation availability when a thread begins. + /// Seeds image-generation configuration when a thread begins. fn on_thread_start<'a>( &'a self, input: ThreadStartInput<'a, Config>, @@ -63,7 +63,7 @@ impl ThreadLifecycleContributor for ImageGenerationExtension { } impl ConfigContributor for ImageGenerationExtension { - /// Refreshes image-generation availability after thread configuration changes. + /// Refreshes image-generation configuration after thread configuration changes. fn on_config_changed( &self, _session_store: &ExtensionData, @@ -88,7 +88,7 @@ impl ToolContributor for ImageGenerationExtension { let Some(config) = thread_store.get::() else { return Vec::new(); }; - if !config.available || !self.auth_manager.current_auth_uses_codex_backend() { + if !config.available { return Vec::new(); } diff --git a/codex-rs/ext/web-search/src/extension.rs b/codex-rs/ext/web-search/src/extension.rs index 7afdfe4137..84e3eeef9d 100644 --- a/codex-rs/ext/web-search/src/extension.rs +++ b/codex-rs/ext/web-search/src/extension.rs @@ -41,7 +41,8 @@ impl From<&Config> for WebSearchExtensionConfig { let web_search_mode = config.web_search_mode.value(); Self { // Core selects this executor per turn using the feature flag or model metadata. - available: config.model_provider.is_openai() + available: (config.model_provider.is_openai() + || config.model_provider.uses_openai_actor_authorization()) && web_search_mode != WebSearchMode::Disabled, provider: config.model_provider.clone(), settings: search_settings(config, web_search_mode), diff --git a/codex-rs/model-provider-info/src/lib.rs b/codex-rs/model-provider-info/src/lib.rs index ca3bf71ccc..265e65b25a 100644 --- a/codex-rs/model-provider-info/src/lib.rs +++ b/codex-rs/model-provider-info/src/lib.rs @@ -33,6 +33,7 @@ const MAX_STREAM_MAX_RETRIES: u64 = 100; const MAX_REQUEST_MAX_RETRIES: u64 = 100; const OPENAI_PROVIDER_NAME: &str = "OpenAI"; +const OPENAI_ACTOR_AUTHORIZATION_HEADER: &str = "x-openai-actor-authorization"; pub const OPENAI_PROVIDER_ID: &str = "openai"; pub const CHATGPT_CODEX_BASE_URL: &str = "https://chatgpt.com/backend-api/codex"; const AMAZON_BEDROCK_PROVIDER_NAME: &str = "Amazon Bedrock"; @@ -392,6 +393,16 @@ impl ModelProviderInfo { self.name == OPENAI_PROVIDER_NAME } + pub fn uses_openai_actor_authorization(&self) -> bool { + !self.requires_openai_auth + && self.http_headers.as_ref().is_some_and(|headers| { + headers.iter().any(|(name, value)| { + name.eq_ignore_ascii_case(OPENAI_ACTOR_AUTHORIZATION_HEADER) + && !value.trim().is_empty() + }) + }) + } + pub fn is_amazon_bedrock(&self) -> bool { self.name == AMAZON_BEDROCK_PROVIDER_NAME } diff --git a/codex-rs/model-provider-info/src/model_provider_info_tests.rs b/codex-rs/model-provider-info/src/model_provider_info_tests.rs index a303050dc1..5949f36bb8 100644 --- a/codex-rs/model-provider-info/src/model_provider_info_tests.rs +++ b/codex-rs/model-provider-info/src/model_provider_info_tests.rs @@ -198,6 +198,31 @@ fn test_supports_remote_compaction_for_non_openai_non_azure_provider() { assert!(!provider.supports_remote_compaction()); } +#[test] +fn test_uses_openai_actor_authorization() { + let mut provider = ModelProviderInfo { + http_headers: Some(maplit::hashmap! { + "X-OpenAI-Actor-Authorization".to_string() => "actor-token".to_string(), + }), + ..ModelProviderInfo::default() + }; + assert!(provider.uses_openai_actor_authorization()); + + provider.http_headers = None; + assert!(!provider.uses_openai_actor_authorization()); + + provider.http_headers = Some(maplit::hashmap! { + OPENAI_ACTOR_AUTHORIZATION_HEADER.to_string() => " ".to_string(), + }); + assert!(!provider.uses_openai_actor_authorization()); + + provider.http_headers = Some(maplit::hashmap! { + OPENAI_ACTOR_AUTHORIZATION_HEADER.to_string() => "actor-token".to_string(), + }); + provider.requires_openai_auth = true; + assert!(!provider.uses_openai_actor_authorization()); +} + #[test] fn test_deserialize_provider_auth_config_defaults() { let base_dir = tempdir().unwrap(); From 25f50de6ed95627e1ffe7f11ca30d3d62c6e20e6 Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 26 Jun 2026 03:11:55 +0100 Subject: [PATCH 13/38] Test selected capabilities across availability and resume (#30157) ## Why This stack crosses World State, executor skills, selected plugin metadata, MCP processes, connectors, dynamic environments, and resume. This PR adds two end-to-end scenarios that validate those pieces together. Both tests enable `deferred_executor`, so they exercise the real delayed-environment path. ## Scenario 1: availability across turns and resume ```text 1. Start a thread with one selected plugin root bound to E1. 2. E1 is unavailable. - executor skill is absent - selected MCP is absent - connector has no selected-plugin attribution 3. Start E1 and register the same stable environment ID. 4. Start a new turn. - the executor skill appears through World State - its body beats a colliding host skill - the selected MCP tool is advertised and executes inside E1 - the connector is attributed to the selected plugin 5. Start another turn without changing E1. - the MCP PID stays the same, proving runtime reuse 6. Restart app-server and resume the thread. - durable selected-root intent is restored - skills, MCP, and connector attribution are restored - a new MCP PID proves ephemeral process state was rebuilt ``` ## Scenario 2: availability changes inside one turn ```text 1. Start a turn while E1 is unavailable. 2. The first model sample sees no executor skill, MCP, or selected connector. 3. The turn pauses on request_user_input. 4. Start E1 and register it while that same turn is still active. 5. Continue the turn. 6. The very next model sample sees: - the executor skill catalog - the selected MCP tool - selected-plugin connector attribution 7. The model calls the MCP, and its output proves execution happened inside E1. ``` This second scenario specifically protects the aeon-style behavior: capability state is captured again for every sampling step, not only at the next user turn. ## Scope These are integration tests only. They do not add a combinatorial matrix for unsupported plugin-file mutation, environment generations, transport disconnects, or delayed `required = true` executor MCPs. --- .../app-server/tests/suite/v2/app_list.rs | 4 +- codex-rs/app-server/tests/suite/v2/mod.rs | 2 + .../suite/v2/selected_capability_stack.rs | 702 ++++++++++++++++++ .../exec-server/src/remote_file_system.rs | 17 +- codex-rs/file-system/src/lib.rs | 16 +- 5 files changed, 735 insertions(+), 6 deletions(-) create mode 100644 codex-rs/app-server/tests/suite/v2/selected_capability_stack.rs diff --git a/codex-rs/app-server/tests/suite/v2/app_list.rs b/codex-rs/app-server/tests/suite/v2/app_list.rs index 1ddf355269..5fad04929f 100644 --- a/codex-rs/app-server/tests/suite/v2/app_list.rs +++ b/codex-rs/app-server/tests/suite/v2/app_list.rs @@ -1533,7 +1533,7 @@ impl ServerHandler for AppListMcpServer { } } -async fn start_apps_server_with_delays( +pub(super) async fn start_apps_server_with_delays( connectors: Vec, tools: Vec, directory_delay: Duration, @@ -1693,7 +1693,7 @@ async fn list_directory_connectors( } } -fn connector_tool(connector_id: &str, connector_name: &str) -> Result { +pub(super) fn connector_tool(connector_id: &str, connector_name: &str) -> Result { let schema: JsonObject = serde_json::from_value(json!({ "type": "object", "additionalProperties": false diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index bbd9952754..9963d37be5 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -56,6 +56,8 @@ mod request_user_input; mod request_validation; mod review; mod safety_check_downgrade; +#[cfg(not(target_os = "windows"))] +mod selected_capability_stack; mod skills_list; mod sleep; mod thread_archive; diff --git a/codex-rs/app-server/tests/suite/v2/selected_capability_stack.rs b/codex-rs/app-server/tests/suite/v2/selected_capability_stack.rs new file mode 100644 index 0000000000..052f286ffc --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/selected_capability_stack.rs @@ -0,0 +1,702 @@ +use std::process::Stdio; +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +use app_test_support::ChatGptAuthFixture; +use app_test_support::TestAppServer; +use app_test_support::to_response; +use app_test_support::write_chatgpt_auth; +use app_test_support::write_mock_responses_config_toml_with_chatgpt_base_url; +use codex_app_server_protocol::AppInfo; +use codex_app_server_protocol::CapabilityRootLocation; +use codex_app_server_protocol::EnvironmentAddResponse; +use codex_app_server_protocol::ListMcpServerStatusParams; +use codex_app_server_protocol::ListMcpServerStatusResponse; +use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::SelectedCapabilityRoot; +use codex_app_server_protocol::ServerRequest; +use codex_app_server_protocol::ThreadResumeParams; +use codex_app_server_protocol::ThreadResumeResponse; +use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::ThreadStartResponse; +use codex_app_server_protocol::TurnEnvironmentParams; +use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::UserInput; +use codex_config::types::AuthCredentialsStoreMode; +use codex_exec_server::LOCAL_ENVIRONMENT_ID; +use codex_protocol::config_types::CollaborationMode; +use codex_protocol::config_types::ModeKind; +use codex_protocol::config_types::Settings; +use codex_utils_absolute_path::AbsolutePathBuf; +use codex_utils_path_uri::PathUri; +use core_test_support::process::wait_for_pid_file; +use core_test_support::responses; +use core_test_support::responses::ResponsesRequest; +use core_test_support::stdio_server_bin; +use pretty_assertions::assert_eq; +use pretty_assertions::assert_ne; +use serde_json::json; +use tempfile::TempDir; +use tokio::io::AsyncBufReadExt; +use tokio::io::BufReader; +use tokio::process::Child; +use tokio::process::Command; +use tokio::time::timeout; + +use super::app_list::connector_tool; +use super::app_list::start_apps_server_with_delays; + +const READ_TIMEOUT: Duration = Duration::from_secs(20); +const EXECUTOR_ID: &str = "executor-1"; +const EXECUTOR_ENV_NAME: &str = "MCP_EXECUTOR_MARKER"; +const EXECUTOR_ENV_VALUE: &str = "executor-only"; +const PLUGIN_ID: &str = "executor-demo@1"; +const PLUGIN_DISPLAY_NAME: &str = "Executor Demo"; +const SKILL_NAME: &str = "executor-demo:deploy"; +const SKILL_DESCRIPTION: &str = "Deploy through the selected executor."; +const SKILL_BODY_MARKER: &str = "SELECTED_EXECUTOR_SKILL_BODY"; +const LOCAL_SKILL_BODY_MARKER: &str = "COLLIDING_LOCAL_SKILL_BODY"; +const MCP_SERVER_NAME: &str = "executor_probe"; +const MCP_CALL_ID: &str = "selected-executor-mcp-call"; +const CONNECTOR_ID: &str = "calendar"; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn selected_capability_stack_tracks_environment_availability_and_resume() -> Result<()> { + let responses_server = responses::start_mock_server().await; + let (apps_url, apps_server_handle) = start_apps_server_with_delays( + vec![AppInfo { + id: CONNECTOR_ID.to_string(), + name: "Calendar".to_string(), + description: None, + logo_url: None, + logo_url_dark: None, + icon_assets: None, + icon_dark_assets: None, + distribution_channel: None, + branding: None, + app_metadata: None, + labels: None, + install_url: None, + is_accessible: false, + is_enabled: true, + plugin_display_names: Vec::new(), + }], + vec![connector_tool(CONNECTOR_ID, "Calendar")?], + Duration::ZERO, + Duration::ZERO, + ) + .await?; + let fixture = selected_capability_fixture(&responses_server.uri(), &apps_url)?; + + let response_mock = responses::mount_sse_sequence( + &responses_server, + vec![ + responses::sse(vec![ + responses::ev_response_created("environment-unavailable"), + responses::ev_assistant_message("unavailable-message", "Waiting"), + responses::ev_completed("environment-unavailable"), + ]), + responses::sse(vec![ + responses::ev_response_created("environment-available-call"), + responses::ev_function_call_with_namespace( + MCP_CALL_ID, + &format!("mcp__{MCP_SERVER_NAME}"), + "echo", + &json!({ + "message": "hello from the selected executor", + "env_var": EXECUTOR_ENV_NAME, + }) + .to_string(), + ), + responses::ev_completed("environment-available-call"), + ]), + responses::sse(vec![ + responses::ev_response_created("environment-available-done"), + responses::ev_assistant_message("available-message", "Done"), + responses::ev_completed("environment-available-done"), + ]), + responses::sse(vec![ + responses::ev_response_created("unchanged-step"), + responses::ev_assistant_message("unchanged-message", "Still ready"), + responses::ev_completed("unchanged-step"), + ]), + responses::sse(vec![ + responses::ev_response_created("resumed-step"), + responses::ev_assistant_message("resumed-message", "Ready after resume"), + responses::ev_completed("resumed-step"), + ]), + ], + ) + .await; + + let mut app_server = TestAppServer::new(fixture.codex_home.path()).await?; + timeout(READ_TIMEOUT, app_server.initialize()).await??; + let thread_id = start_thread( + &mut app_server, + fixture.selected_root.clone(), + fixture.environment_cwd.clone(), + ) + .await?; + + run_turn( + &mut app_server, + &thread_id, + "Inspect the current capabilities", + fixture.environment_cwd.clone(), + ) + .await?; + let initial_requests = response_mock.requests(); + assert_selected_capabilities_absent(&initial_requests[0]); + + let mut exec_server = + spawn_exec_server(fixture.codex_home.path(), &fixture.exec_server_url).await?; + add_environment(&mut app_server, &fixture.exec_server_url).await?; + wait_for_selected_mcp_server(&mut app_server, &thread_id).await?; + + run_turn( + &mut app_server, + &thread_id, + &format!("Use ${SKILL_NAME} and call its selected executor MCP"), + fixture.environment_cwd.clone(), + ) + .await?; + let first_mcp_pid = wait_for_pid_file(&fixture.pid_file).await?; + + run_turn( + &mut app_server, + &thread_id, + "Continue with the same selected capabilities", + fixture.environment_cwd.clone(), + ) + .await?; + assert_eq!(first_mcp_pid, wait_for_pid_file(&fixture.pid_file).await?); + + drop(app_server); + std::fs::remove_file(&fixture.pid_file)?; + + let mut app_server = TestAppServer::new(fixture.codex_home.path()).await?; + timeout(READ_TIMEOUT, app_server.initialize()).await??; + let request_id = app_server + .send_thread_resume_request(ThreadResumeParams { + thread_id: thread_id.clone(), + ..Default::default() + }) + .await?; + let response = timeout( + READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let ThreadResumeResponse { thread, .. } = to_response(response)?; + assert_eq!(thread_id, thread.id); + wait_for_selected_mcp_server(&mut app_server, &thread_id).await?; + + run_turn( + &mut app_server, + &thread_id, + &format!("Use ${SKILL_NAME} after resuming the thread"), + fixture.environment_cwd, + ) + .await?; + let resumed_mcp_pid = wait_for_pid_file(&fixture.pid_file).await?; + assert_ne!(first_mcp_pid, resumed_mcp_pid); + + let requests = response_mock.requests(); + assert_eq!(5, requests.len()); + for request in &requests[1..4] { + assert_selected_skill_is_injected(request, /*expected_count*/ 1); + assert_selected_plugin_tools(request); + } + assert_selected_skill_is_injected(&requests[4], /*expected_count*/ 2); + assert_selected_plugin_tools(&requests[4]); + let output = requests[2].function_call_output(MCP_CALL_ID); + let output = output["output"] + .as_str() + .expect("MCP function output should be text"); + assert!(output.contains("ECHOING: hello from the selected executor")); + assert!(output.contains(EXECUTOR_ENV_VALUE)); + + exec_server.kill().await?; + apps_server_handle.abort(); + let _ = apps_server_handle.await; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn selected_capabilities_become_available_between_samples_in_one_turn() -> Result<()> { + const USER_INPUT_CALL_ID: &str = "pause-for-environment"; + + let responses_server = responses::start_mock_server().await; + let (apps_url, apps_server_handle) = start_apps_server_with_delays( + vec![AppInfo { + id: CONNECTOR_ID.to_string(), + name: "Calendar".to_string(), + description: None, + logo_url: None, + logo_url_dark: None, + icon_assets: None, + icon_dark_assets: None, + distribution_channel: None, + branding: None, + app_metadata: None, + labels: None, + install_url: None, + is_accessible: false, + is_enabled: true, + plugin_display_names: Vec::new(), + }], + vec![connector_tool(CONNECTOR_ID, "Calendar")?], + Duration::ZERO, + Duration::ZERO, + ) + .await?; + let fixture = selected_capability_fixture(&responses_server.uri(), &apps_url)?; + let response_mock = responses::mount_sse_sequence( + &responses_server, + vec![ + responses::sse(vec![ + responses::ev_response_created("environment-pending"), + responses::ev_function_call( + USER_INPUT_CALL_ID, + "request_user_input", + &json!({ + "questions": [{ + "id": "continue", + "header": "Continue", + "question": "Continue after the executor is attached?", + "options": [{ + "label": "Yes (Recommended)", + "description": "Continue the same turn." + }, { + "label": "No", + "description": "Stop here." + }] + }], + "autoResolutionMs": 60_000 + }) + .to_string(), + ), + responses::ev_completed("environment-pending"), + ]), + responses::sse(vec![ + responses::ev_response_created("environment-ready-call"), + responses::ev_function_call_with_namespace( + MCP_CALL_ID, + &format!("mcp__{MCP_SERVER_NAME}"), + "echo", + &json!({ + "message": "same turn", + "env_var": EXECUTOR_ENV_NAME, + }) + .to_string(), + ), + responses::ev_completed("environment-ready-call"), + ]), + responses::sse(vec![ + responses::ev_response_created("same-turn-done"), + responses::ev_assistant_message("same-turn-message", "Done"), + responses::ev_completed("same-turn-done"), + ]), + ], + ) + .await; + + let mut app_server = TestAppServer::new(fixture.codex_home.path()).await?; + timeout(READ_TIMEOUT, app_server.initialize()).await??; + let thread_id = start_thread( + &mut app_server, + fixture.selected_root, + fixture.environment_cwd.clone(), + ) + .await?; + let turn_start_id = app_server + .send_turn_start_request(TurnStartParams { + thread_id, + input: vec![UserInput::Text { + text: "Use the executor when it becomes ready.".to_string(), + text_elements: Vec::new(), + }], + environments: Some(vec![TurnEnvironmentParams { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: fixture.environment_cwd.into(), + }]), + collaboration_mode: Some(CollaborationMode { + mode: ModeKind::Plan, + settings: Settings { + model: "mock-model".to_string(), + reasoning_effort: None, + developer_instructions: None, + }, + }), + ..Default::default() + }) + .await?; + timeout( + READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(turn_start_id)), + ) + .await??; + + let request = timeout(READ_TIMEOUT, app_server.read_stream_until_request_message()).await??; + let ServerRequest::ToolRequestUserInput { request_id, .. } = request else { + panic!("expected request_user_input, got {request:?}"); + }; + let requests = response_mock.requests(); + assert_eq!(1, requests.len()); + assert_selected_capabilities_absent(&requests[0]); + + let mut exec_server = + spawn_exec_server(fixture.codex_home.path(), &fixture.exec_server_url).await?; + add_environment(&mut app_server, &fixture.exec_server_url).await?; + tokio::time::sleep(Duration::from_millis(200)).await; + app_server + .send_response( + request_id, + json!({ + "answers": { + "continue": { "answers": ["yes"] } + } + }), + ) + .await?; + timeout( + READ_TIMEOUT, + app_server.read_stream_until_notification_message("turn/completed"), + ) + .await??; + + let requests = response_mock.requests(); + assert_eq!(3, requests.len()); + assert_selected_skill_catalog_available(&requests[1]); + assert_selected_plugin_tools(&requests[1]); + assert_selected_plugin_tools(&requests[2]); + let output = requests[2].function_call_output(MCP_CALL_ID); + let output = output["output"] + .as_str() + .expect("MCP function output should be text"); + assert!(output.contains("ECHOING: same turn")); + assert!(output.contains(EXECUTOR_ENV_VALUE)); + wait_for_pid_file(&fixture.pid_file).await?; + + exec_server.kill().await?; + apps_server_handle.abort(); + let _ = apps_server_handle.await; + Ok(()) +} + +struct SelectedCapabilityFixture { + codex_home: TempDir, + _plugin: TempDir, + pid_file: std::path::PathBuf, + exec_server_url: String, + selected_root: SelectedCapabilityRoot, + environment_cwd: AbsolutePathBuf, +} + +fn selected_capability_fixture( + responses_server_uri: &str, + apps_url: &str, +) -> Result { + let codex_home = TempDir::new()?; + write_mock_responses_config_toml_with_chatgpt_base_url( + codex_home.path(), + responses_server_uri, + apps_url, + )?; + let config_path = codex_home.path().join("config.toml"); + let config = std::fs::read_to_string(&config_path)?.replacen( + "model_provider = \"mock_provider\"", + "mcp_oauth_credentials_store = \"file\"\nmodel_provider = \"mock_provider\"", + 1, + ); + std::fs::write( + config_path, + format!( + "{config}\n[features]\napps = true\ndeferred_executor = true\n\n[skills]\ninclude_instructions = true\n" + ), + )?; + write_chatgpt_auth( + codex_home.path(), + ChatGptAuthFixture::new("chatgpt-token") + .account_id("account-123") + .email("selected-capability-stack@example.com") + .plan_type("pro") + .chatgpt_account_id("account-123"), + AuthCredentialsStoreMode::File, + )?; + + // Reserve the URL before app-server starts. The configured environment initially fails to + // connect, then environment/add points the same stable ID at the same URL once it is live. + let listener = std::net::TcpListener::bind("127.0.0.1:0")?; + let exec_server_url = format!("ws://{}", listener.local_addr()?); + drop(listener); + std::fs::write( + codex_home.path().join("environments.toml"), + format!( + "default = \"{EXECUTOR_ID}\"\ninclude_local = true\n\n[[environments]]\nid = \"{EXECUTOR_ID}\"\nurl = \"{exec_server_url}\"\nconnect_timeout_sec = 0.05\n" + ), + )?; + + let local_skill_dir = codex_home.path().join("skills/local-deploy"); + std::fs::create_dir_all(&local_skill_dir)?; + std::fs::write( + local_skill_dir.join("SKILL.md"), + format!( + "---\nname: {SKILL_NAME}\ndescription: Colliding local skill.\n---\n\n{LOCAL_SKILL_BODY_MARKER}\n" + ), + )?; + + let plugin = TempDir::new()?; + let manifest_dir = plugin.path().join(".codex-plugin"); + let skill_dir = plugin.path().join("skills/deploy"); + let pid_file = plugin.path().join("executor-mcp.pid"); + std::fs::create_dir_all(&manifest_dir)?; + std::fs::create_dir_all(&skill_dir)?; + std::fs::write( + manifest_dir.join("plugin.json"), + r#"{"name":"executor-demo","apps":"./.app.json","interface":{"displayName":"Executor Demo"}}"#, + )?; + std::fs::write( + skill_dir.join("SKILL.md"), + format!( + "---\nname: deploy\ndescription: {SKILL_DESCRIPTION}\n---\n\n{SKILL_BODY_MARKER}\n" + ), + )?; + std::fs::write( + plugin.path().join(".app.json"), + format!(r#"{{"apps":{{"calendar":{{"id":"{CONNECTOR_ID}"}}}}}}"#), + )?; + std::fs::write( + plugin.path().join(".mcp.json"), + serde_json::to_vec_pretty(&json!({ + "mcpServers": { + (MCP_SERVER_NAME): { + "command": stdio_server_bin()?, + "env": { + "MCP_TEST_PID_FILE": pid_file.to_string_lossy(), + }, + "env_vars": [EXECUTOR_ENV_NAME], + "startup_timeout_sec": 10, + } + } + }))?, + )?; + + let selected_root = SelectedCapabilityRoot { + id: PLUGIN_ID.to_string(), + location: CapabilityRootLocation::Environment { + environment_id: EXECUTOR_ID.to_string(), + path: PathUri::from_host_native_path(plugin.path())?, + }, + }; + let environment_cwd = AbsolutePathBuf::try_from(plugin.path().to_path_buf())?; + Ok(SelectedCapabilityFixture { + codex_home, + _plugin: plugin, + pid_file, + exec_server_url, + selected_root, + environment_cwd, + }) +} + +fn assert_selected_capabilities_absent(request: &ResponsesRequest) { + assert!( + request + .message_input_texts("developer") + .into_iter() + .all(|text| !text.contains(SKILL_DESCRIPTION)) + ); + assert!( + request + .tool_by_name(&format!("mcp__{MCP_SERVER_NAME}"), "echo") + .is_none() + ); + let connector = request + .tool_by_name("mcp__codex_apps__calendar", "connector_calendar") + .expect("host connector should remain model-visible"); + assert!( + connector["description"] + .as_str() + .is_some_and(|description| !description.contains(PLUGIN_DISPLAY_NAME)) + ); +} + +fn assert_selected_skill_is_injected(request: &ResponsesRequest, expected_count: usize) { + assert_selected_skill_catalog_available(request); + + let skill_fragments = request + .message_input_texts("user") + .into_iter() + .filter(|text| text.starts_with("")) + .collect::>(); + assert_eq!(expected_count, skill_fragments.len()); + for fragment in skill_fragments { + assert!(fragment.contains(&format!("{SKILL_NAME}"))); + assert!(fragment.contains(SKILL_BODY_MARKER)); + assert!(!fragment.contains(LOCAL_SKILL_BODY_MARKER)); + } +} + +fn assert_selected_skill_catalog_available(request: &ResponsesRequest) { + let catalog_fragments = request + .message_input_texts("developer") + .into_iter() + .filter(|text| text.contains(SKILL_DESCRIPTION)) + .collect::>(); + assert_eq!(1, catalog_fragments.len()); + assert!(catalog_fragments[0].contains("environment resource:")); +} + +fn assert_selected_plugin_tools(request: &ResponsesRequest) { + assert!( + request + .tool_by_name(&format!("mcp__{MCP_SERVER_NAME}"), "echo") + .is_some() + ); + let connector = request + .tool_by_name("mcp__codex_apps__calendar", "connector_calendar") + .expect("selected connector should be model-visible"); + assert!( + connector["description"] + .as_str() + .is_some_and(|description| description.contains(PLUGIN_DISPLAY_NAME)) + ); +} + +async fn start_thread( + app_server: &mut TestAppServer, + selected_root: SelectedCapabilityRoot, + environment_cwd: AbsolutePathBuf, +) -> Result { + let request_id = app_server + .send_thread_start_request(ThreadStartParams { + model: Some("mock-model".to_string()), + environments: Some(vec![TurnEnvironmentParams { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: environment_cwd.into(), + }]), + selected_capability_roots: Some(vec![selected_root]), + ..Default::default() + }) + .await?; + let response = timeout( + READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response(response)?; + Ok(thread.id) +} + +async fn run_turn( + app_server: &mut TestAppServer, + thread_id: &str, + text: &str, + environment_cwd: AbsolutePathBuf, +) -> Result<()> { + let request_id = app_server + .send_turn_start_request(TurnStartParams { + thread_id: thread_id.to_string(), + input: vec![UserInput::Text { + text: text.to_string(), + text_elements: Vec::new(), + }], + environments: Some(vec![TurnEnvironmentParams { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd: environment_cwd.into(), + }]), + ..Default::default() + }) + .await?; + timeout( + READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + timeout( + READ_TIMEOUT, + app_server.read_stream_until_notification_message("turn/completed"), + ) + .await??; + Ok(()) +} + +async fn add_environment(app_server: &mut TestAppServer, exec_server_url: &str) -> Result<()> { + let request_id = app_server + .send_raw_request( + "environment/add", + Some(json!({ + "environmentId": EXECUTOR_ID, + "execServerUrl": exec_server_url, + "connectTimeoutMs": 10_000, + })), + ) + .await?; + let response = timeout( + READ_TIMEOUT, + app_server.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let _: EnvironmentAddResponse = to_response(response)?; + Ok(()) +} + +async fn wait_for_selected_mcp_server( + app_server: &mut TestAppServer, + thread_id: &str, +) -> Result<()> { + timeout(READ_TIMEOUT, async { + loop { + let request_id = app_server + .send_list_mcp_server_status_request(ListMcpServerStatusParams { + cursor: None, + limit: None, + detail: None, + thread_id: Some(thread_id.to_string()), + }) + .await?; + let response = app_server + .read_stream_until_response_message(RequestId::Integer(request_id)) + .await?; + let response: ListMcpServerStatusResponse = to_response(response)?; + if response + .data + .iter() + .any(|server| server.name == MCP_SERVER_NAME) + { + return Ok::<_, anyhow::Error>(()); + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + }) + .await??; + Ok(()) +} + +async fn spawn_exec_server(codex_home: &std::path::Path, url: &str) -> Result { + let mut child = Command::new(codex_utils_cargo_bin::cargo_bin("codex")?) + .args(["exec-server", "--listen", url]) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::inherit()) + .kill_on_drop(true) + .env("CODEX_HOME", codex_home) + .env(EXECUTOR_ENV_NAME, EXECUTOR_ENV_VALUE) + .spawn()?; + let stdout = child + .stdout + .take() + .context("exec-server stdout was not captured")?; + let mut lines = BufReader::new(stdout).lines(); + loop { + let line = timeout(READ_TIMEOUT, lines.next_line()) + .await + .context("timed out waiting for exec-server URL")?? + .context("exec-server exited before printing its URL")?; + if line.trim() == url { + return Ok(child); + } + } +} diff --git a/codex-rs/exec-server/src/remote_file_system.rs b/codex-rs/exec-server/src/remote_file_system.rs index acc86f7584..0d10c3aece 100644 --- a/codex-rs/exec-server/src/remote_file_system.rs +++ b/codex-rs/exec-server/src/remote_file_system.rs @@ -29,6 +29,7 @@ use crate::protocol::FsWalkParams; use crate::protocol::FsWriteFileParams; const INVALID_REQUEST_ERROR_CODE: i64 = -32600; +const METHOD_NOT_FOUND_ERROR_CODE: i64 = -32601; const NOT_FOUND_ERROR_CODE: i64 = -32004; #[path = "remote_file_stream.rs"] @@ -194,14 +195,26 @@ impl RemoteFileSystem { ) -> FileSystemResult { trace!("remote fs walk"); let client = self.client.get().await.map_err(map_remote_error)?; - let response = client + let response = match client .fs_walk(FsWalkParams { path: path.clone(), options, sandbox: remote_sandbox_context(sandbox), }) .await - .map_err(map_remote_error)?; + { + Ok(response) => response, + Err(ExecServerError::Server { + code: METHOD_NOT_FOUND_ERROR_CODE, + .. + }) => { + return ::walk_via_directory_reads( + self, path, options, sandbox, + ) + .await; + } + Err(error) => return Err(map_remote_error(error)), + }; Ok(response) } diff --git a/codex-rs/file-system/src/lib.rs b/codex-rs/file-system/src/lib.rs index fcdecc8cb5..737412f463 100644 --- a/codex-rs/file-system/src/lib.rs +++ b/codex-rs/file-system/src/lib.rs @@ -316,7 +316,19 @@ pub trait ExecutorFileSystem: Send + Sync { options: WalkOptions, sandbox: Option<&'a FileSystemSandboxContext>, ) -> ExecutorFileSystemFuture<'a, WalkOutcome> { - Box::pin(walk(self, path, options, sandbox)) + self.walk_via_directory_reads(path, options, sandbox) + } + + /// Performs a bounded walk using the primitive filesystem operations. + /// + /// Implementations with an optimized walk transport can use this as a compatibility fallback. + fn walk_via_directory_reads<'a>( + &'a self, + path: &'a PathUri, + options: WalkOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, WalkOutcome> { + Box::pin(walk_via_directory_reads(self, path, options, sandbox)) } fn remove<'a>( @@ -335,7 +347,7 @@ pub trait ExecutorFileSystem: Send + Sync { ) -> ExecutorFileSystemFuture<'a, ()>; } -async fn walk( +async fn walk_via_directory_reads( file_system: &F, root: &PathUri, options: WalkOptions, From f5f812389ee49ab4c9ef1237781ea1013e733fdc Mon Sep 17 00:00:00 2001 From: Tom Date: Thu, 25 Jun 2026 21:01:11 -0700 Subject: [PATCH 14/38] [codex] fix terminal rollout event durability (#30144) Currently session code does not flush the thread store after appending the `TurnComplete` / `TurnAborted` events. This isn't a problem in practice for local storage because append_items itself effectively blocks, but any thread stores that buffer in append_items and only commit on flush effectively never get these events persisted. The fix adds explicit rollout flushes at the terminal emitters after normal completion and interruption. Added test cases that assert the number of flushes when completing or aborting turns. These are admittedly a little brittle and I'm open to better ideas on how to add automated testing. --- codex-rs/core/src/session/tests.rs | 156 +++++++++++++++++++++++++++++ codex-rs/core/src/tasks/mod.rs | 15 ++- 2 files changed, 168 insertions(+), 3 deletions(-) diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 57f1ce9081..2bc63b404f 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -9153,6 +9153,89 @@ impl SessionTask for CompletingTask { } } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum TerminalEventKind { + TurnComplete, + TurnAborted, +} + +async fn attach_in_memory_thread_store( + session: &mut Session, +) -> Arc { + let store = Arc::new(codex_thread_store::InMemoryThreadStore::default()); + let thread_store: Arc = store.clone(); + let config = session.get_config().await; + let live_thread = LiveThread::create( + Arc::clone(&thread_store), + CreateThreadParams { + session_id: session.session_id(), + thread_id: session.thread_id, + extra_config: None, + forked_from_id: None, + parent_thread_id: None, + source: SessionSource::Exec, + thread_source: None, + originator: "test_originator".to_string(), + base_instructions: BaseInstructions::default(), + dynamic_tools: Vec::new(), + multi_agent_version: None, + initial_window_id: Uuid::now_v7().to_string(), + metadata: ThreadPersistenceMetadata { + cwd: Some(config.cwd.to_path_buf()), + model_provider: config.model_provider_id.clone(), + memory_mode: if config.memories.generate_memories { + ThreadMemoryMode::Enabled + } else { + ThreadMemoryMode::Disabled + }, + }, + }, + ) + .await + .expect("create thread persistence"); + session.services.thread_store = thread_store; + session.services.live_thread = Some(live_thread); + store +} + +async fn wait_for_flush_count( + store: &codex_thread_store::InMemoryThreadStore, + expected_flushes: usize, +) -> codex_thread_store::InMemoryThreadStoreCalls { + timeout(Duration::from_secs(2), async { + loop { + let calls = store.calls().await; + if calls.flush_thread >= expected_flushes { + return calls; + } + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("store should observe expected flush count") +} + +async fn recv_terminal_event( + rx: &async_channel::Receiver, + expected: TerminalEventKind, +) -> Event { + timeout(Duration::from_secs(2), async { + loop { + let event = rx.recv().await.expect("event"); + match (&event.msg, expected) { + (EventMsg::TurnComplete(_), TerminalEventKind::TurnComplete) + | (EventMsg::TurnAborted(_), TerminalEventKind::TurnAborted) => return event, + (EventMsg::TurnComplete(_) | EventMsg::TurnAborted(_), _) => { + panic!("unexpected terminal event: {:?}", event.msg) + } + _ => {} + } + } + }) + .await + .expect("terminal event should be delivered") +} + #[derive(Clone, Copy)] struct NeverEndingTask { kind: TaskKind, @@ -9310,6 +9393,79 @@ async fn guardian_helper_review_interrupts_after_three_consecutive_denials() { assert_eq!(aborted.reason, TurnAbortReason::Interrupted); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn turn_complete_flushes_terminal_event_after_delivery() { + let (mut sess, tc, rx) = make_session_and_context_with_rx().await; + let store = attach_in_memory_thread_store( + Arc::get_mut(&mut sess).expect("session should be uniquely owned"), + ) + .await; + + let input = vec![TurnInput::UserInput { + content: vec![UserInput::Text { + text: "complete normally".to_string(), + text_elements: Vec::new(), + }], + client_id: None, + }]; + sess.spawn_task(Arc::clone(&tc), input, CompletingTask) + .await; + + let event = recv_terminal_event(&rx, TerminalEventKind::TurnComplete).await; + assert!(matches!(event.msg, EventMsg::TurnComplete(_))); + // Expected flushes: + // 1. Task-runner flush after the task body finishes, before TurnComplete is emitted. + // 2. Terminal-event flush after TurnComplete is appended. + let calls = wait_for_flush_count(&store, /*expected_flushes*/ 2).await; + assert_eq!(2, calls.flush_thread); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn turn_aborted_flushes_terminal_event_after_delivery() { + let (mut sess, tc, rx) = make_session_and_context_with_rx().await; + let store = attach_in_memory_thread_store( + Arc::get_mut(&mut sess).expect("session should be uniquely owned"), + ) + .await; + + let input = vec![TurnInput::UserInput { + content: vec![UserInput::Text { + text: "interrupt me".to_string(), + text_elements: Vec::new(), + }], + client_id: None, + }]; + sess.spawn_task( + Arc::clone(&tc), + input, + NeverEndingTask { + kind: TaskKind::Regular, + listen_to_cancellation_token: true, + }, + ) + .await; + + let abort_task = tokio::spawn({ + let sess = Arc::clone(&sess); + async move { + sess.abort_all_tasks(TurnAbortReason::Interrupted).await; + } + }); + + let event = recv_terminal_event(&rx, TerminalEventKind::TurnAborted).await; + match event.msg { + EventMsg::TurnAborted(e) => assert_eq!(TurnAbortReason::Interrupted, e.reason), + other => panic!("unexpected event: {other:?}"), + } + abort_task.await.expect("abort task should finish"); + // Expected flushes: + // 1. Task-runner flush after the task body observes cancellation. + // 2. Interrupted-marker flush before TurnAborted so abort observers can reread it. + // 3. Terminal-event flush after TurnAborted is appended. + let calls = wait_for_flush_count(&store, /*expected_flushes*/ 3).await; + assert_eq!(3, calls.flush_thread); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[test_log::test] async fn abort_regular_task_emits_marker_before_turn_aborted() { diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index cb532445fd..c6d9d2990c 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -793,10 +793,14 @@ impl Session { false } }; - if !cleared_active_turn { - return; + if cleared_active_turn { + self.emit_thread_idle_lifecycle_if_idle().await; + } + // Regular items were flushed before this terminal event was appended; buffering + // thread writers may not flush it without another explicit barrier. + if let Err(err) = self.flush_rollout().await { + warn!("failed to flush rollout after emitting terminal turn event: {err}"); } - self.emit_thread_idle_lifecycle_if_idle().await; } async fn take_active_turn(&self) -> Option { @@ -896,6 +900,11 @@ impl Session { .lock() .await .clear_turn(&task.turn_context.sub_id); + // Regular items were flushed before this terminal event was appended; buffering + // thread writers may not flush it without another explicit barrier. + if let Err(err) = self.flush_rollout().await { + warn!("failed to flush rollout after emitting terminal turn event: {err}"); + } } } From 92d2e1df7079b5bd549c92a17fa234706d0580b3 Mon Sep 17 00:00:00 2001 From: kbazzi Date: Thu, 25 Jun 2026 21:31:12 -0700 Subject: [PATCH 15/38] Retry failed Codex Apps MCP startup (#29920) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem The built-in Codex Apps MCP client shares a future for the full startup operation: connect, complete `initialize`, fetch the initial tools, and return a usable client. Sharing deduplicates startup work, but it also memoizes terminal errors. After a transient connection, handshake, or initial `tools/list` failure, later tool builds observe the same failed future. The thread cannot reconnect after the backend recovers and continues serving its startup-time cached tool snapshot, which may be empty or stale. ## Fix When Apps MCP startup ends in an error, Codex starts bounded recovery without putting startup latency on tool-router construction: 1. The current tool build immediately continues with the cached startup snapshot. 2. After the initial failure is reported, Codex starts one fresh full startup attempt in the background. 3. Concurrent tool builds share that in-flight attempt and also continue with cached tools. 4. On success, the recovered client becomes active, refreshes the Apps tools cache, emits a `Ready` startup status, and is reused by later operations. 5. On failure, the cache remains unchanged and later tool builds may start another background attempt after exponential cooldown: 1s, 2s, 4s, 8s, 16s, then 30s maximum. Each recreated startup performs a fresh MCP `initialize` and uncached `tools/list`. The MCP client retains its existing bounded retries for retryable `initialize` and `tools/list` failures. This avoids adding the Apps startup timeout to every request during a sustained outage. ## Scope This is limited to the built-in Codex Apps MCP client: - no reconnects for user-configured MCP servers; - no cache deletion; and - no proactive refresh for a healthy client with stale tools. ## Tests Coverage verifies: - tool builds return cached tools without waiting for a blocked reconnect; - concurrent tool builds start only one background reconnect; - failed reconnects preserve cached tools and respect exponential cooldown; - a recovered client is retained and reused; and - a long-lived thread exposes recovered app tools on a later follow-up. Validation: - `just test -p codex-mcp` — 95 passed - `just test -p codex-core later_follow_up_uses_background_recovered_apps_after_mid_thread_startup_failures --no-capture` — passed - `just fix -p codex-mcp` - `just fmt` --- codex-rs/codex-mcp/src/connection_manager.rs | 6 + .../codex-mcp/src/connection_manager_tests.rs | 298 ++++++++++++++++- codex-rs/codex-mcp/src/rmcp_client.rs | 315 +++++++++++++++--- .../core/tests/common/apps_test_server.rs | 92 +++++ .../core/tests/suite/mcp_tool_exposure.rs | 107 ++++++ 5 files changed, 767 insertions(+), 51 deletions(-) diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 3107f678b0..c76e3c3c08 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -208,6 +208,7 @@ impl McpConnectionManager { }; let async_managed_client = AsyncManagedClient::new( server_name.clone(), + startup_submit_id.clone(), server, store_mode, keyring_backend_kind, @@ -257,6 +258,10 @@ impl McpConnectionManager { ) .await; + if matches!(&outcome, Err(StartupOutcomeError::Failed { .. })) { + async_managed_client.reconnect_failed_startup().await; + } + (server_name, outcome) }); } @@ -479,6 +484,7 @@ impl McpConnectionManager { pub async fn list_all_tools(&self) -> Vec { let mut tools = Vec::new(); for (server_name, managed_client) in &self.clients { + managed_client.reconnect_failed_startup().await; let has_cached_tools = managed_client.has_cached_tools(); let startup_complete = managed_client .startup_complete diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index d2869fe5c1..1f61ab700e 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -6,7 +6,10 @@ use crate::elicitation::ElicitationRequestManager; use crate::elicitation::ElicitationRequestRouter; use crate::elicitation::elicitation_is_rejected_by_policy; use crate::rmcp_client::AsyncManagedClient; +use crate::rmcp_client::CODEX_APPS_RECONNECT_INITIAL_BACKOFF; +use crate::rmcp_client::CodexAppsStartupReconnect; use crate::rmcp_client::ManagedClient; +use crate::rmcp_client::ManagedClientFuture; use crate::rmcp_client::StartupOutcomeError; use crate::server::EffectiveMcpServer; use crate::server::McpServerMetadata; @@ -21,6 +24,7 @@ use codex_config::Constrained; use codex_config::McpServerConfig; use codex_config::McpServerToolConfig; use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::EnvironmentManager; use codex_protocol::ToolName; use codex_protocol::mcp::McpServerInfo; @@ -41,6 +45,7 @@ use rmcp::model::Tool; use std::collections::HashSet; use std::io; use std::sync::Arc; +use std::sync::atomic::AtomicUsize; use tempfile::tempdir; use tokio::io::DuplexStream; @@ -101,9 +106,8 @@ impl InProcessTransportFactory for TestInProcessTransportFactory { } } -async fn create_ready_async_managed_client(tools: Vec) -> AsyncManagedClient { - let tool_filter = ToolFilter::default(); - let managed_client = ManagedClient { +async fn create_test_managed_client(tools: Vec) -> ManagedClient { + ManagedClient { client: Arc::new( RmcpClient::new_in_process_client(Arc::new(TestInProcessTransportFactory)) .await @@ -111,29 +115,73 @@ async fn create_ready_async_managed_client(tools: Vec) -> AsyncManaged ), server_info: create_test_server_info("Ready"), tools, - tool_filter: tool_filter.clone(), + tool_filter: ToolFilter::default(), tool_timeout: None, server_instructions: None, server_supports_sandbox_state_meta_capability: false, codex_apps_tools_cache_context: None, - }; + } +} +async fn create_ready_async_managed_client(tools: Vec) -> AsyncManagedClient { AsyncManagedClient { client: futures::future::ready::>(Ok( - managed_client, + create_test_managed_client(tools).await, )) .boxed() .shared(), is_codex_apps_mcp_server: false, cached_server_info: None, codex_apps_tools_cache_context: None, - tool_filter, + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), + startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), } } +fn create_test_manager_with_failed_apps_startup( + cached_tools: Vec, + reconnect_factory: Arc ManagedClientFuture + Send + Sync>, +) -> McpConnectionManager { + let client: ManagedClientFuture = futures::future::ready(Err(StartupOutcomeError::Failed { + error: "startup failed".to_string(), + is_authentication_required: false, + })) + .boxed() + .shared(); + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("reconnect-test-account"), + Some("reconnect-test-user"), + ); + cache_context.store_current_tools_for_test(cached_tools); + let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut manager = McpConnectionManager::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + manager.clients.insert( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + AsyncManagedClient { + client, + is_codex_apps_mcp_server: true, + cached_server_info: None, + codex_apps_tools_cache_context: Some(cache_context), + tool_filter: ToolFilter::default(), + startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), + startup_reconnect: Some(Arc::new(CodexAppsStartupReconnect::new(reconnect_factory))), + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), + }, + ); + manager +} + fn model_tool_names(tools: &[ToolInfo]) -> HashSet { tools .iter() @@ -721,6 +769,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() { codex_apps_tools_cache_context: Some(cache_context), tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, @@ -760,6 +809,7 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() { codex_apps_tools_cache_context: None, tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, @@ -860,6 +910,7 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tools() { codex_apps_tools_cache_context: None, tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, @@ -898,6 +949,7 @@ async fn shutdown_cancels_pending_tool_listing() { codex_apps_tools_cache_context: None, tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token, }, @@ -944,6 +996,7 @@ async fn shutdown_continues_after_caller_is_aborted() { codex_apps_tools_cache_context: None, tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, @@ -996,6 +1049,7 @@ async fn list_all_tools_does_not_block_when_shared_codex_apps_cache_is_empty() { codex_apps_tools_cache_context: Some(cache_context), tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, @@ -1045,6 +1099,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails() codex_apps_tools_cache_context: Some(cache_context), tool_filter: ToolFilter::default(), startup_complete, + startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, @@ -1069,6 +1124,235 @@ async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails() ); } +#[tokio::test] +async fn list_all_tools_reconnects_failed_codex_apps_startup_and_reuses_client() { + let recovered_client = create_test_managed_client(vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "drive_search", + )]) + .await; + let attempts = Arc::new(AtomicUsize::new(0)); + let attempts_for_reconnect = Arc::clone(&attempts); + let reconnect_finished = Arc::new(tokio::sync::Notify::new()); + let reconnect_finished_for_factory = Arc::clone(&reconnect_finished); + let reconnect_factory = Arc::new(move || { + attempts_for_reconnect.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let reconnect_finished = Arc::clone(&reconnect_finished_for_factory); + let recovered_client = recovered_client.clone(); + async move { + reconnect_finished.notify_one(); + Ok(recovered_client) + } + .boxed() + .shared() + }); + let manager = create_test_manager_with_failed_apps_startup(Vec::new(), reconnect_factory); + + let reconnect_finished_wait = reconnect_finished.notified(); + let tools = manager.list_all_tools().await; + assert!(tools.is_empty()); + reconnect_finished_wait.await; + + let tools = manager.list_all_tools().await; + assert_eq!( + tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["drive_search"] + ); + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); + + let tools = manager.list_all_tools().await; + assert_eq!( + tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["drive_search"] + ); + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); +} + +#[tokio::test(start_paused = true)] +async fn later_tool_list_retries_after_failed_reconnect_and_keeps_cached_tools() { + let recovered_client = create_test_managed_client(vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "drive_search", + )]) + .await; + let attempts = Arc::new(AtomicUsize::new(0)); + let attempts_for_reconnect = Arc::clone(&attempts); + let reconnect_finished = Arc::new(tokio::sync::Notify::new()); + let reconnect_finished_for_factory = Arc::clone(&reconnect_finished); + let reconnect_factory = Arc::new(move || { + let attempt = attempts_for_reconnect.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let reconnect_finished = Arc::clone(&reconnect_finished_for_factory); + let recovered_client = recovered_client.clone(); + async move { + let result = if attempt < 2 { + Err(StartupOutcomeError::Failed { + error: "recreated startup failed".to_string(), + is_authentication_required: false, + }) + } else { + Ok(recovered_client) + }; + reconnect_finished.notify_one(); + result + } + .boxed() + .shared() + }); + let manager = create_test_manager_with_failed_apps_startup( + vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "cached_drive_search", + )], + reconnect_factory, + ); + + let first_reconnect_finished = reconnect_finished.notified(); + let tools = manager.list_all_tools().await; + assert_eq!( + tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["cached_drive_search"] + ); + first_reconnect_finished.await; + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); + + let tools = manager.list_all_tools().await; + assert_eq!( + tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["cached_drive_search"] + ); + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); + + tokio::time::advance(CODEX_APPS_RECONNECT_INITIAL_BACKOFF).await; + let second_reconnect_finished = reconnect_finished.notified(); + let tools = manager.list_all_tools().await; + assert_eq!( + tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["cached_drive_search"] + ); + second_reconnect_finished.await; + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 2); + + tokio::time::advance(CODEX_APPS_RECONNECT_INITIAL_BACKOFF).await; + let tools = manager.list_all_tools().await; + assert_eq!( + tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["cached_drive_search"] + ); + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 2); + + tokio::time::advance(CODEX_APPS_RECONNECT_INITIAL_BACKOFF).await; + let third_reconnect_finished = reconnect_finished.notified(); + let tools = manager.list_all_tools().await; + assert_eq!( + tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["cached_drive_search"] + ); + third_reconnect_finished.await; + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 3); + + let tools = manager.list_all_tools().await; + assert_eq!( + tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["drive_search"] + ); +} + +#[tokio::test] +async fn tool_lists_do_not_block_and_share_codex_apps_startup_reconnect() { + let recovered_client = create_test_managed_client(vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "drive_search", + )]) + .await; + let attempts = Arc::new(AtomicUsize::new(0)); + let attempts_for_reconnect = Arc::clone(&attempts); + let reconnect_started = Arc::new(tokio::sync::Notify::new()); + let reconnect_started_for_factory = Arc::clone(&reconnect_started); + let release_reconnect = Arc::new(tokio::sync::Notify::new()); + let release_reconnect_for_factory = Arc::clone(&release_reconnect); + let reconnect_factory = Arc::new(move || { + let recovered_client = recovered_client.clone(); + let attempts = Arc::clone(&attempts_for_reconnect); + let reconnect_started = Arc::clone(&reconnect_started_for_factory); + let release_reconnect = Arc::clone(&release_reconnect_for_factory); + async move { + attempts.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + reconnect_started.notify_one(); + release_reconnect.notified().await; + Ok(recovered_client) + } + .boxed() + .shared() + }); + let manager = Arc::new(create_test_manager_with_failed_apps_startup( + vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "cached_drive_search", + )], + reconnect_factory, + )); + let reconnect_started_wait = reconnect_started.notified(); + let first_tools = tokio::time::timeout(Duration::from_millis(10), manager.list_all_tools()) + .await + .expect("cached tools should not wait for reconnect"); + + reconnect_started_wait.await; + let second_tools = tokio::time::timeout(Duration::from_millis(10), manager.list_all_tools()) + .await + .expect("concurrent cached tools should not wait for reconnect"); + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); + assert_eq!( + first_tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["cached_drive_search"] + ); + assert_eq!( + second_tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["cached_drive_search"] + ); + + release_reconnect.notify_one(); + tokio::task::yield_now().await; + let tools = manager.list_all_tools().await; + assert_eq!( + tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["drive_search"] + ); + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); +} + #[tokio::test] async fn list_all_tools_adds_server_metadata_to_tools() { let server_name = "docs"; diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index ddfb3ffdba..ee7dde1fe2 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -12,6 +12,7 @@ use std::collections::HashMap; use std::env; use std::ffi::OsString; use std::sync::Arc; +use std::sync::Mutex as StdMutex; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Duration; @@ -48,6 +49,9 @@ use codex_exec_server::HttpClient; use codex_exec_server::ReqwestHttpClient; use codex_protocol::mcp::McpServerInfo; use codex_protocol::protocol::Event; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::McpStartupStatus; +use codex_protocol::protocol::McpStartupUpdateEvent; use codex_rmcp_client::ExecutorStdioServerLauncher; use codex_rmcp_client::LocalStdioServerLauncher; use codex_rmcp_client::RmcpClient; @@ -64,6 +68,7 @@ use rmcp::model::JsonObject; use rmcp::model::ProtocolVersion; use rmcp::model::Tool as RmcpTool; use rmcp::transport::auth::AuthError; +use tokio::time::Instant as TokioInstant; use tokio_util::sync::CancellationToken; use tracing::Instrument; use tracing::instrument; @@ -80,6 +85,9 @@ pub(crate) const MCP_TOOLS_FETCH_UNCACHED_DURATION_METRIC: &str = pub(crate) const DEFAULT_STARTUP_TIMEOUT: Duration = Duration::from_secs(30); pub(crate) const DEFAULT_TOOL_TIMEOUT: Duration = Duration::from_secs(300); +pub(crate) const CODEX_APPS_RECONNECT_INITIAL_BACKOFF: Duration = Duration::from_secs(1); +const CODEX_APPS_RECONNECT_MAX_BACKOFF: Duration = Duration::from_secs(30); + const UNTRUSTED_CONNECTOR_META_KEYS: &[&str] = &[ "connector_id", "connector_name", @@ -128,56 +136,174 @@ impl ManagedClient { } } -#[derive(Clone)] -pub(crate) struct AsyncManagedClient { - pub(crate) client: Shared>>, - pub(crate) is_codex_apps_mcp_server: bool, - pub(crate) cached_server_info: Option, - pub(crate) codex_apps_tools_cache_context: Option, - pub(crate) tool_filter: ToolFilter, - pub(crate) startup_complete: Arc, - pub(crate) tool_plugin_provenance: Arc, - pub(crate) cancel_token: CancellationToken, +pub(crate) type ManagedClientFuture = + Shared>>; + +#[derive(Default)] +struct CodexAppsStartupReconnectState { + current_client: Option, + reconnect_in_flight: bool, + consecutive_failures: u32, + retry_not_before: Option, } -impl AsyncManagedClient { - // Keep this constructor flat so the startup inputs remain readable at the - // single call site instead of introducing a one-off params wrapper. - #[instrument(level = "trace", skip_all, fields(server_name = %server_name))] - #[allow(clippy::too_many_arguments)] - pub(crate) fn new( +#[derive(Clone)] +struct CodexAppsStartupStatusContext { + submit_id: String, + server_name: String, + tx_event: Sender, +} + +pub(crate) struct CodexAppsStartupReconnect { + factory: Arc ManagedClientFuture + Send + Sync>, + state: StdMutex, + startup_status_context: Option, +} + +impl CodexAppsStartupReconnect { + pub(crate) fn new(factory: Arc ManagedClientFuture + Send + Sync>) -> Self { + Self { + factory, + state: StdMutex::new(CodexAppsStartupReconnectState::default()), + startup_status_context: None, + } + } + + fn with_startup_status_context( + mut self, + submit_id: String, server_name: String, - server: EffectiveMcpServer, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, - cancel_token: CancellationToken, tx_event: Sender, - elicitation_requests: ElicitationRequestManager, - codex_apps_tools_cache_context: Option, - tool_plugin_provenance: Arc, - runtime_context: McpRuntimeContext, - runtime_auth_provider: Option, - client_elicitation_capability: ElicitationCapability, - supports_openai_form_elicitation: bool, ) -> Self { + self.startup_status_context = Some(CodexAppsStartupStatusContext { + submit_id, + server_name, + tx_event, + }); + self + } + + fn current_client(&self) -> Option { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .current_client + .clone() + } + + fn reconnect_in_background(self: &Arc) { + { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.current_client.is_some() || state.reconnect_in_flight { + return; + } + if state + .retry_not_before + .is_some_and(|retry_not_before| TokioInstant::now() < retry_not_before) + { + return; + } + state.reconnect_in_flight = true; + } + + let reconnect = Arc::clone(self); + tokio::spawn(async move { + let result = (reconnect.factory)().await; + let startup_status_context = reconnect.startup_status_context.clone(); + let recovered = { + let mut state = reconnect + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.reconnect_in_flight = false; + match result { + Ok(client) => { + state.current_client = Some(client); + state.consecutive_failures = 0; + state.retry_not_before = None; + true + } + Err(error) => { + state.consecutive_failures = state.consecutive_failures.saturating_add(1); + let retry_after = codex_apps_reconnect_backoff(state.consecutive_failures); + state.retry_not_before = Some(TokioInstant::now() + retry_after); + warn!( + error = %error, + retry_after_ms = retry_after.as_millis(), + "Apps MCP startup reconnect failed; continuing with cached tools" + ); + false + } + } + }; + + if recovered && let Some(context) = startup_status_context { + let _ = context + .tx_event + .send(Event { + id: context.submit_id, + msg: EventMsg::McpStartupUpdate(McpStartupUpdateEvent { + server: context.server_name, + status: McpStartupStatus::Ready, + }), + }) + .await; + } + }); + } +} + +fn codex_apps_reconnect_backoff(consecutive_failures: u32) -> Duration { + let exponent = consecutive_failures.saturating_sub(1).min(5); + CODEX_APPS_RECONNECT_INITIAL_BACKOFF + .saturating_mul(1 << exponent) + .min(CODEX_APPS_RECONNECT_MAX_BACKOFF) +} + +#[derive(Clone)] +struct ManagedClientStartup { + server_name: String, + server: EffectiveMcpServer, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, + tx_event: Sender, + elicitation_requests: ElicitationRequestManager, + codex_apps_tools_cache_context: Option, + runtime_context: McpRuntimeContext, + runtime_auth_provider: Option, + client_elicitation_capability: ElicitationCapability, + supports_openai_form_elicitation: bool, + cancel_token: CancellationToken, + startup_complete: Arc, +} + +impl ManagedClientStartup { + fn start(&self) -> ManagedClientFuture { + let Self { + server_name, + server, + store_mode, + keyring_backend_kind, + tx_event, + elicitation_requests, + codex_apps_tools_cache_context, + runtime_context, + runtime_auth_provider, + client_elicitation_capability, + supports_openai_form_elicitation, + cancel_token, + startup_complete, + } = self.clone(); let is_codex_apps_mcp_server = server_name == CODEX_APPS_MCP_SERVER_NAME; let tool_filter = server .configured_config() .map(ToolFilter::from_config) .unwrap_or_default(); - let cached_server_info = if is_codex_apps_mcp_server { - codex_apps_tools_cache_context - .as_ref() - .and_then(load_startup_cached_codex_apps_server_info) - } else { - None - }; - let startup_tool_filter = tool_filter.clone(); - let codex_apps_tools_cache_context_for_fut = codex_apps_tools_cache_context.clone(); - let startup_complete = Arc::new(AtomicBool::new(false)); - let startup_complete_for_fut = Arc::clone(&startup_complete); - let cancel_token_for_fut = cancel_token.clone(); - let fut = async move { + let cancel_token_for_fut = cancel_token; + async move { let outcome = match async { if let Err(error) = validate_mcp_server_name(&server_name) { return Err(error.into()); @@ -207,10 +333,10 @@ impl AsyncManagedClient { .configured_config() .and_then(|config| config.tool_timeout_sec) .unwrap_or(DEFAULT_TOOL_TIMEOUT), - tool_filter: startup_tool_filter, + tool_filter, tx_event, elicitation_requests, - codex_apps_tools_cache_context: codex_apps_tools_cache_context_for_fut, + codex_apps_tools_cache_context, client_elicitation_capability, supports_openai_form_elicitation, }, @@ -224,10 +350,91 @@ impl AsyncManagedClient { Err(CancelErr::Cancelled) => Err(StartupOutcomeError::Cancelled), }; - startup_complete_for_fut.store(true, Ordering::Release); + startup_complete.store(true, Ordering::Release); outcome + } + .in_current_span() + .boxed() + .shared() + } +} + +#[derive(Clone)] +pub(crate) struct AsyncManagedClient { + pub(crate) client: ManagedClientFuture, + pub(crate) is_codex_apps_mcp_server: bool, + pub(crate) cached_server_info: Option, + pub(crate) codex_apps_tools_cache_context: Option, + pub(crate) tool_filter: ToolFilter, + pub(crate) startup_complete: Arc, + pub(crate) startup_reconnect: Option>, + pub(crate) tool_plugin_provenance: Arc, + pub(crate) cancel_token: CancellationToken, +} + +impl AsyncManagedClient { + // Keep this constructor flat so the startup inputs remain readable at the + // single call site instead of introducing a one-off params wrapper. + #[instrument(level = "trace", skip_all, fields(server_name = %server_name))] + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( + server_name: String, + startup_submit_id: String, + server: EffectiveMcpServer, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, + cancel_token: CancellationToken, + tx_event: Sender, + elicitation_requests: ElicitationRequestManager, + codex_apps_tools_cache_context: Option, + tool_plugin_provenance: Arc, + runtime_context: McpRuntimeContext, + runtime_auth_provider: Option, + client_elicitation_capability: ElicitationCapability, + supports_openai_form_elicitation: bool, + ) -> Self { + let is_codex_apps_mcp_server = server_name == CODEX_APPS_MCP_SERVER_NAME; + let reconnect_server_name = server_name.clone(); + let reconnect_tx_event = tx_event.clone(); + let tool_filter = server + .configured_config() + .map(ToolFilter::from_config) + .unwrap_or_default(); + let cached_server_info = if is_codex_apps_mcp_server { + codex_apps_tools_cache_context + .as_ref() + .and_then(load_startup_cached_codex_apps_server_info) + } else { + None }; - let client = fut.in_current_span().boxed().shared(); + let startup_complete = Arc::new(AtomicBool::new(false)); + let startup = Arc::new(ManagedClientStartup { + server_name, + server, + store_mode, + keyring_backend_kind, + tx_event, + elicitation_requests, + codex_apps_tools_cache_context: codex_apps_tools_cache_context.clone(), + runtime_context, + runtime_auth_provider, + client_elicitation_capability, + supports_openai_form_elicitation, + cancel_token: cancel_token.clone(), + startup_complete: Arc::clone(&startup_complete), + }); + let client = startup.start(); + let startup_reconnect = is_codex_apps_mcp_server.then(|| { + let startup = Arc::clone(&startup); + Arc::new( + CodexAppsStartupReconnect::new(Arc::new(move || startup.start())) + .with_startup_status_context( + startup_submit_id, + reconnect_server_name, + reconnect_tx_event, + ), + ) + }); if codex_apps_tools_cache_context .as_ref() .is_some_and(CodexAppsToolsCacheContext::has_current_tools) @@ -245,15 +452,35 @@ impl AsyncManagedClient { codex_apps_tools_cache_context, tool_filter, startup_complete, + startup_reconnect, tool_plugin_provenance, cancel_token, } } pub(crate) async fn client(&self) -> Result { + if let Some(client) = self + .startup_reconnect + .as_ref() + .and_then(|reconnect| reconnect.current_client()) + { + return Ok(client); + } self.client.clone().await } + pub(crate) async fn reconnect_failed_startup(&self) { + let Some(startup_reconnect) = self.startup_reconnect.as_ref() else { + return; + }; + if !self.startup_complete.load(Ordering::Acquire) { + return; + } + if matches!(self.client().await, Err(StartupOutcomeError::Failed { .. })) { + startup_reconnect.reconnect_in_background(); + } + } + pub(crate) async fn shutdown(&self) { self.cancel_token.cancel(); match self.client().await { diff --git a/codex-rs/core/tests/common/apps_test_server.rs b/codex-rs/core/tests/common/apps_test_server.rs index 0849087d0b..3daaada6c5 100644 --- a/codex-rs/core/tests/common/apps_test_server.rs +++ b/codex-rs/core/tests/common/apps_test_server.rs @@ -7,6 +7,7 @@ use codex_login::CodexAuth; use codex_models_manager::bundled_models_response; use serde_json::Value; use serde_json::json; +use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use wiremock::Mock; @@ -55,6 +56,23 @@ pub struct AppsTestServer { pub chatgpt_base_url: String, } +#[derive(Clone)] +pub struct AppsTestServerStartupControl { + initialize_attempts: Arc, + remaining_initialize_failures: Arc, +} + +impl AppsTestServerStartupControl { + pub fn fail_next_initialize_attempts(&self, attempts: usize) { + self.remaining_initialize_failures + .store(attempts, Ordering::SeqCst); + } + + pub fn initialize_attempts(&self) -> usize { + self.initialize_attempts.load(Ordering::SeqCst) + } +} + #[derive(Clone, Copy)] pub enum AppsTestToolLoading { Direct, @@ -130,6 +148,34 @@ impl AppsTestServer { }) } + pub async fn mount_with_startup_control( + server: &MockServer, + ) -> Result<(Self, AppsTestServerStartupControl)> { + mount_oauth_metadata(server).await; + mount_connectors_directory(server).await; + let control = AppsTestServerStartupControl { + initialize_attempts: Arc::new(AtomicUsize::new(0)), + remaining_initialize_failures: Arc::new(AtomicUsize::new(0)), + }; + mount_streamable_http_json_rpc_with_startup_control( + server, + CONNECTOR_NAME.to_string(), + CONNECTOR_DESCRIPTION.to_string(), + /*searchable*/ true, + /*include_app_only_tool*/ false, + AppsTestToolsListBehavior::AlwaysAvailable, + Some(Arc::clone(&control.initialize_attempts)), + Some(Arc::clone(&control.remaining_initialize_failures)), + ) + .await; + Ok(( + Self { + chatgpt_base_url: server.uri(), + }, + control, + )) + } + pub async fn mount_with_tools_available_after_initial_list( server: &MockServer, ) -> Result { @@ -312,6 +358,30 @@ async fn mount_streamable_http_json_rpc( searchable: bool, include_app_only_tool: bool, tools_list_behavior: AppsTestToolsListBehavior, +) { + mount_streamable_http_json_rpc_with_startup_control( + server, + connector_name, + connector_description, + searchable, + include_app_only_tool, + tools_list_behavior, + /*initialize_attempts*/ None, + /*remaining_initialize_failures*/ None, + ) + .await; +} + +#[allow(clippy::too_many_arguments)] +async fn mount_streamable_http_json_rpc_with_startup_control( + server: &MockServer, + connector_name: String, + connector_description: String, + searchable: bool, + include_app_only_tool: bool, + tools_list_behavior: AppsTestToolsListBehavior, + initialize_attempts: Option>, + remaining_initialize_failures: Option>, ) { Mock::given(method("POST")) .and(path_regex("^/api/codex/apps/?$")) @@ -322,6 +392,8 @@ async fn mount_streamable_http_json_rpc( include_app_only_tool, tools_list_behavior, tools_list_calls: AtomicUsize::new(0), + initialize_attempts, + remaining_initialize_failures, }) .mount(server) .await; @@ -334,6 +406,8 @@ struct CodexAppsJsonRpcResponder { include_app_only_tool: bool, tools_list_behavior: AppsTestToolsListBehavior, tools_list_calls: AtomicUsize, + initialize_attempts: Option>, + remaining_initialize_failures: Option>, } impl Respond for CodexAppsJsonRpcResponder { @@ -355,6 +429,24 @@ impl Respond for CodexAppsJsonRpcResponder { match method { "initialize" => { + if let Some(initialize_attempts) = &self.initialize_attempts { + initialize_attempts.fetch_add(1, Ordering::SeqCst); + } + if self + .remaining_initialize_failures + .as_ref() + .is_some_and(|remaining| { + remaining + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |remaining| { + remaining.checked_sub(1) + }) + .is_ok() + }) + { + return ResponseTemplate::new(400).set_body_json(json!({ + "error": "simulated non-retryable Apps MCP startup failure", + })); + } let id = body.get("id").cloned().unwrap_or(Value::Null); let protocol_version = body .pointer("/params/protocolVersion") diff --git a/codex-rs/core/tests/suite/mcp_tool_exposure.rs b/codex-rs/core/tests/suite/mcp_tool_exposure.rs index b86862edc2..32ecf7cf2a 100644 --- a/codex-rs/core/tests/suite/mcp_tool_exposure.rs +++ b/codex-rs/core/tests/suite/mcp_tool_exposure.rs @@ -1,5 +1,8 @@ use anyhow::Result; use codex_features::Feature; +use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_protocol::protocol::McpServerRefreshConfig; +use codex_protocol::protocol::Op; use core_test_support::apps_test_server::AppsTestServer; use core_test_support::apps_test_server::SEARCH_CALENDAR_CREATE_TOOL; use core_test_support::apps_test_server::SEARCH_CALENDAR_NAMESPACE; @@ -8,10 +11,13 @@ use core_test_support::responses; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::namespace_child_tool; use core_test_support::responses::sse; use core_test_support::skip_if_no_network; +use core_test_support::wait_for_mcp_server; use serde_json::Value; +use std::time::Duration; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn code_mode_only_exposes_direct_model_only_mcp_namespaces() -> Result<()> { @@ -79,3 +85,104 @@ async fn code_mode_only_exposes_direct_model_only_mcp_namespaces() -> Result<()> Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn later_follow_up_uses_background_recovered_apps_after_mid_thread_startup_failures() +-> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let (apps_server, startup_control) = + AppsTestServer::mount_with_startup_control(&server).await?; + let response = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_assistant_message("msg-1", "initial turn"), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-2", "recovery-trigger turn"), + ev_completed("resp-2"), + ]), + sse(vec![ + ev_response_created("resp-3"), + ev_assistant_message("msg-3", "recovered follow-up turn"), + ev_completed("resp-3"), + ]), + ], + ) + .await; + + let mut builder = search_capable_apps_builder(apps_server.chatgpt_base_url.clone()) + .with_config(move |config| { + config + .features + .enable(Feature::CodeModeOnly) + .expect("test config should allow feature update"); + config.code_mode.direct_only_tool_namespaces = + vec![SEARCH_CALENDAR_NAMESPACE.to_string()]; + }); + let test = builder.build(&server).await?; + wait_for_mcp_server(&test.codex, CODEX_APPS_MCP_SERVER_NAME).await?; + test.submit_turn("use Calendar before refreshing MCP") + .await?; + + let initial_request = response.requests()[0].body_json(); + assert!( + namespace_child_tool( + &initial_request, + SEARCH_CALENDAR_NAMESPACE, + SEARCH_CALENDAR_CREATE_TOOL, + ) + .is_some(), + "Calendar should be available before the MCP refresh: {initial_request}" + ); + + tokio::fs::remove_dir_all(test.codex_home_path().join("cache/codex_apps_tools")).await?; + startup_control.fail_next_initialize_attempts(/*attempts*/ 1); + let runtime_mcp_config = test.codex.runtime_mcp_config(&test.config).await; + let refresh_config = McpServerRefreshConfig { + mcp_servers: serde_json::to_value(codex_mcp::configured_mcp_servers(&runtime_mcp_config))?, + mcp_oauth_credentials_store_mode: serde_json::to_value( + runtime_mcp_config.mcp_oauth_credentials_store_mode, + )?, + auth_keyring_backend_kind: serde_json::to_value( + runtime_mcp_config.auth_keyring_backend_kind, + )?, + }; + test.codex + .submit(Op::RefreshMcpServers { + config: refresh_config, + }) + .await?; + test.submit_turn("use Calendar after transient Apps startup failures") + .await?; + tokio::time::timeout(Duration::from_secs(1), async { + while startup_control.initialize_attempts() < 3 { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .expect("background Apps reconnect should complete"); + test.submit_turn("use Calendar after background Apps recovery") + .await?; + + let requests = response.requests(); + assert_eq!(requests.len(), 3); + let recovered_request = requests[2].body_json(); + assert!( + namespace_child_tool( + &recovered_request, + SEARCH_CALENDAR_NAMESPACE, + SEARCH_CALENDAR_CREATE_TOOL, + ) + .is_some(), + "Calendar should recover on the follow-up turn: {recovered_request}", + ); + assert_eq!(startup_control.initialize_attempts(), 3); + + Ok(()) +} From b5866eebd631a62a1f3f29ac5e44752e38949c73 Mon Sep 17 00:00:00 2001 From: stevenlee-oai Date: Fri, 26 Jun 2026 02:23:24 -0400 Subject: [PATCH 16/38] Persist Cloudflare affinity cookies for MCP HTTP (#29516) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit [Codex Thread 019ef1f9-36e2-7e91-9337-504f097b9dc1](https://codex-thread-link.openai.chatgpt-team.site/thread/019ef1f9-36e2-7e91-9337-504f097b9dc1) ## Why Hosted plugin-service Streamable HTTP MCP traffic uses `https://chatgpt.com/backend-api/ps/mcp` and depends on Cloudflare's `__cflb` cookie for load-balancer affinity. The local and exec-server `http/request` path built a fresh reqwest client for each request without installing Codex's existing shared ChatGPT Cloudflare cookie store, so affinity could be lost between calls. This is an affinity-hardening change motivated by an incident investigation. It does not establish the broader connector-cache incident RCA or claim to fix that incident in full. ## What changed - Install the existing process-local, strictly allowlisted ChatGPT Cloudflare cookie store on the reqwest client used by `ReqwestHttpClient`. - Fresh clients now share allowed Cloudflare infrastructure cookies within the process that originates the local or exec-server network request. - Keep the existing HTTPS ChatGPT-host and Cloudflare-cookie-name restrictions. This does not introduce a general cookie jar or send ChatGPT Cloudflare cookies to unrelated hosts. ## Test coverage - `codex-client` unit coverage verifies that the existing strict store accepts and returns `__cflb` for HTTPS ChatGPT URLs. - The exec-server HTTPS integration test sends four independent `http/request` calls through a local TLS-intercepting proxy and verifies that: - `Set-Cookie: __cflb=west` is sent on the next plugin-service request; - a later `Set-Cookie: __cflb=central` replaces the stored value; - non-Cloudflare session cookies are discarded; - no stored ChatGPT Cloudflare cookie is sent to a non-ChatGPT host. - `just test -p codex-client` — 38 passed. - `just test -p codex-exec-server --test chatgpt_cloudflare_affinity` — 1 passed. - `just bazel-lock-check` — passed. ## Non-goals - No persistence of ChatGPT auth, account, session, residency, or arbitrary cookies. - No cookie persistence for third-party MCP servers. - No special composition of caller-provided `Cookie` headers. - No plugin-service, connector-cache, Habitat/habicache, routing, redirect, or API-contract changes. - No broader incident RCA conclusions. --- codex-rs/Cargo.lock | 2 + .../src/chatgpt_cloudflare_cookies.rs | 4 +- codex-rs/exec-server/Cargo.toml | 2 + .../src/client/reqwest_http_client.rs | 3 +- .../tests/chatgpt_cloudflare_affinity.rs | 405 ++++++++++++++++++ 5 files changed, 414 insertions(+), 2 deletions(-) create mode 100644 codex-rs/exec-server/tests/chatgpt_cloudflare_affinity.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 17c51f58a9..ff0265d4de 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2929,7 +2929,9 @@ dependencies = [ "opentelemetry_sdk", "pretty_assertions", "prost 0.14.3", + "rcgen", "reqwest 0.12.28", + "rustls", "serde", "serde_json", "serial_test", diff --git a/codex-rs/codex-client/src/chatgpt_cloudflare_cookies.rs b/codex-rs/codex-client/src/chatgpt_cloudflare_cookies.rs index c5f4bbd4eb..7e5d22057c 100644 --- a/codex-rs/codex-client/src/chatgpt_cloudflare_cookies.rs +++ b/codex-rs/codex-client/src/chatgpt_cloudflare_cookies.rs @@ -128,11 +128,12 @@ mod tests { fn stores_and_returns_cloudflare_cookies_for_chatgpt_hosts() { let store = ChatGptCloudflareCookieStore::default(); let url = reqwest::Url::parse("https://chatgpt.com/backend-api/codex/responses").unwrap(); + let load_balancer = HeaderValue::from_static("__cflb=west; Path=/; Secure; HttpOnly"); let cfuvid = HeaderValue::from_static("_cfuvid=visitor; Path=/; Secure; HttpOnly"); let clearance = HeaderValue::from_static("cf_clearance=clearance; Path=/; Secure; HttpOnly"); - store.set_cookies(&mut [&cfuvid, &clearance].into_iter(), &url); + store.set_cookies(&mut [&load_balancer, &cfuvid, &clearance].into_iter(), &url); let mut cookies = store .cookies(&url) @@ -148,6 +149,7 @@ mod tests { assert_eq!( cookies, vec![ + "__cflb=west".to_string(), "_cfuvid=visitor".to_string(), "cf_clearance=clearance".to_string() ] diff --git a/codex-rs/exec-server/Cargo.toml b/codex-rs/exec-server/Cargo.toml index a8c1133098..35fa411021 100644 --- a/codex-rs/exec-server/Cargo.toml +++ b/codex-rs/exec-server/Cargo.toml @@ -69,6 +69,8 @@ http = { workspace = true } opentelemetry = { workspace = true } opentelemetry_sdk = { workspace = true } pretty_assertions = { workspace = true } +rcgen = { workspace = true } +rustls = { workspace = true } serial_test = { workspace = true } tempfile = { workspace = true } test-case = "3.3.1" diff --git a/codex-rs/exec-server/src/client/reqwest_http_client.rs b/codex-rs/exec-server/src/client/reqwest_http_client.rs index 667e4c49d8..7ead01c9ba 100644 --- a/codex-rs/exec-server/src/client/reqwest_http_client.rs +++ b/codex-rs/exec-server/src/client/reqwest_http_client.rs @@ -9,6 +9,7 @@ use std::error::Error as StdError; use std::time::Duration; use codex_client::build_reqwest_client_with_custom_ca; +use codex_client::with_chatgpt_cloudflare_cookie_store; use codex_exec_server_protocol::JSONRPCErrorError; use futures::FutureExt; use futures::StreamExt; @@ -66,7 +67,7 @@ impl ReqwestHttpClient { HttpRedirectPolicy::Follow => builder, HttpRedirectPolicy::Stop => builder.redirect(reqwest::redirect::Policy::none()), }; - build_reqwest_client_with_custom_ca(builder) + build_reqwest_client_with_custom_ca(with_chatgpt_cloudflare_cookie_store(builder)) .map_err(|error| ExecServerError::HttpRequest(error.to_string())) } } diff --git a/codex-rs/exec-server/tests/chatgpt_cloudflare_affinity.rs b/codex-rs/exec-server/tests/chatgpt_cloudflare_affinity.rs new file mode 100644 index 0000000000..75c86b6b79 --- /dev/null +++ b/codex-rs/exec-server/tests/chatgpt_cloudflare_affinity.rs @@ -0,0 +1,405 @@ +#![cfg(unix)] + +mod common; + +use std::collections::BTreeMap; +use std::ffi::OsString; +use std::fs; +use std::io; +use std::io::Read; +use std::io::Write; +use std::net::TcpListener; +use std::net::TcpStream; +use std::sync::Arc; +use std::sync::mpsc; +use std::thread; +use std::time::Duration; + +use codex_exec_server::HttpRedirectPolicy; +use codex_exec_server::HttpRequestParams; +use codex_exec_server::HttpRequestResponse; +use codex_exec_server::InitializeParams; +use codex_exec_server_protocol::JSONRPCMessage; +use codex_exec_server_protocol::JSONRPCResponse; +use codex_exec_server_protocol::RequestId; +use common::exec_server::ExecServerHarness; +use common::exec_server::exec_server_with_env; +use pretty_assertions::assert_eq; +use rcgen::BasicConstraints; +use rcgen::CertificateParams; +use rcgen::CertifiedIssuer; +use rcgen::DistinguishedName; +use rcgen::DnType; +use rcgen::ExtendedKeyUsagePurpose; +use rcgen::IsCa; +use rcgen::KeyPair; +use rcgen::KeyUsagePurpose; +use rcgen::PKCS_ECDSA_P256_SHA256; +use rustls::pki_types::CertificateDer; +use rustls::pki_types::PrivateKeyDer; +use serde::de::DeserializeOwned; +use serde_json::Value; +use tempfile::TempDir; + +const CHATGPT_MCP_URL: &str = "https://chatgpt.com/backend-api/ps/mcp"; +const NON_CHATGPT_MCP_URL: &str = "https://api.openai.com/backend-api/ps/mcp"; + +#[derive(Debug)] +struct CapturedRequest { + connect_authority: String, + request_line: String, + headers: BTreeMap>, +} + +struct TlsMaterial { + ca_cert_pem: String, + server_cert: CertificateDer<'static>, + server_key: PrivateKeyDer<'static>, +} + +struct TlsInterceptingProxy { + ca_cert_pem: String, + request_rx: mpsc::Receiver>, + thread: thread::JoinHandle>, + url: String, +} + +/// Exercises the same `http/request` route used by remotely executed Streamable HTTP MCP calls. +/// Each RPC builds a fresh reqwest client. The first response sets `__cflb`, and the second response +/// replaces it, proving cross-client persistence through the shared cookie store. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_server_replays_only_chatgpt_cloudflare_cookies() -> anyhow::Result<()> { + let proxy = TlsInterceptingProxy::start(/*expected_requests*/ 4)?; + let temp_dir = TempDir::new()?; + let ca_path = temp_dir.path().join("cloudflare-affinity-test-ca.pem"); + fs::write(&ca_path, &proxy.ca_cert_pem)?; + let proxy_url = OsString::from(&proxy.url); + let empty = OsString::new(); + let env = vec![ + ( + OsString::from("CODEX_CA_CERTIFICATE"), + ca_path.as_os_str().to_owned(), + ), + (OsString::from("HTTPS_PROXY"), proxy_url.clone()), + (OsString::from("https_proxy"), proxy_url.clone()), + (OsString::from("ALL_PROXY"), proxy_url.clone()), + (OsString::from("all_proxy"), proxy_url), + (OsString::from("NO_PROXY"), empty.clone()), + (OsString::from("no_proxy"), empty), + ]; + let mut server = exec_server_with_env(env).await?; + initialize_exec_server(&mut server).await?; + + let first_response = execute_http_request(&mut server, CHATGPT_MCP_URL, "first").await?; + assert_eq!( + (first_response.status, first_response.body.into_inner()), + (200, b"ok".to_vec()) + ); + let first = proxy.next_request()?; + assert_eq!( + ( + first.connect_authority.as_str(), + first.request_line.as_str(), + first.headers.get("cookie"), + ), + ("chatgpt.com:443", "POST /backend-api/ps/mcp HTTP/1.1", None,) + ); + + let west_response = execute_http_request(&mut server, CHATGPT_MCP_URL, "west").await?; + assert_eq!(west_response.status, 200); + let request_with_west_affinity = proxy.next_request()?; + assert_eq!( + request_with_west_affinity + .headers + .get("cookie") + .cloned() + .unwrap_or_default(), + vec!["__cflb=west".to_string()] + ); + + let central_response = execute_http_request(&mut server, CHATGPT_MCP_URL, "central").await?; + assert_eq!(central_response.status, 200); + let request_with_central_affinity = proxy.next_request()?; + assert_eq!( + ( + request_with_central_affinity.request_line.as_str(), + request_with_central_affinity + .headers + .get("cookie") + .cloned() + .unwrap_or_default(), + ), + ( + "POST /backend-api/ps/mcp HTTP/1.1", + vec!["__cflb=central".to_string()], + ) + ); + let other_host_response = + execute_http_request(&mut server, NON_CHATGPT_MCP_URL, "other-host").await?; + assert_eq!(other_host_response.status, 200); + let other_host = proxy.next_request()?; + assert_eq!( + ( + other_host.connect_authority.as_str(), + other_host.request_line.as_str(), + other_host.headers.get("cookie"), + ), + ( + "api.openai.com:443", + "POST /backend-api/ps/mcp HTTP/1.1", + None, + ) + ); + + server.shutdown().await?; + proxy.finish()?; + Ok(()) +} + +impl TlsInterceptingProxy { + fn start(expected_requests: usize) -> anyhow::Result { + codex_utils_rustls_provider::ensure_rustls_crypto_provider(); + let material = generate_tls_material()?; + let listener = TcpListener::bind(("127.0.0.1", 0))?; + let address = listener.local_addr()?; + let config = Arc::new( + rustls::ServerConfig::builder_with_protocol_versions(&[&rustls::version::TLS13]) + .with_no_client_auth() + .with_single_cert(vec![material.server_cert], material.server_key)?, + ); + let (request_tx, request_rx) = mpsc::channel(); + let thread = thread::spawn(move || { + run_tls_intercepting_proxy(listener, config, request_tx, expected_requests) + .map_err(|error| error.to_string()) + }); + + Ok(Self { + ca_cert_pem: material.ca_cert_pem, + request_rx, + thread, + url: format!("http://{address}"), + }) + } + + fn next_request(&self) -> anyhow::Result { + self.request_rx + .recv_timeout(Duration::from_secs(5)) + .map_err(anyhow::Error::from)? + .map_err(anyhow::Error::msg) + } + + fn finish(self) -> anyhow::Result<()> { + self.thread + .join() + .map_err(|_| anyhow::anyhow!("TLS proxy thread panicked"))? + .map_err(anyhow::Error::msg) + } +} + +fn generate_tls_material() -> anyhow::Result { + let mut ca_params = CertificateParams::default(); + ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); + ca_params.key_usages = vec![KeyUsagePurpose::KeyCertSign, KeyUsagePurpose::CrlSign]; + let mut ca_distinguished_name = DistinguishedName::new(); + ca_distinguished_name.push(DnType::CommonName, "Codex affinity test CA"); + ca_params.distinguished_name = ca_distinguished_name; + let ca_key_pair = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)?; + let ca = CertifiedIssuer::self_signed(ca_params, ca_key_pair)?; + + let mut server_params = CertificateParams::new(vec![ + "chatgpt.com".to_string(), + "api.openai.com".to_string(), + ])?; + server_params.extended_key_usages = vec![ExtendedKeyUsagePurpose::ServerAuth]; + server_params.key_usages = vec![ + KeyUsagePurpose::DigitalSignature, + KeyUsagePurpose::KeyEncipherment, + ]; + let server_key_pair = KeyPair::generate_for(&PKCS_ECDSA_P256_SHA256)?; + let server_cert = server_params.signed_by(&server_key_pair, &ca)?; + + Ok(TlsMaterial { + ca_cert_pem: ca.pem(), + server_cert: server_cert.der().clone(), + server_key: PrivateKeyDer::from(server_key_pair), + }) +} + +fn run_tls_intercepting_proxy( + listener: TcpListener, + config: Arc, + request_tx: mpsc::Sender>, + expected_requests: usize, +) -> io::Result<()> { + for request_index in 0..expected_requests { + let (mut stream, _) = listener.accept()?; + configure_stream(&stream)?; + let connect_head = read_http_head(&mut stream)?; + let connect_authority = connect_head + .lines() + .next() + .and_then(|line| line.split_whitespace().nth(1)) + .ok_or_else(|| io::Error::other("CONNECT request line should include an authority"))? + .to_string(); + stream.write_all(b"HTTP/1.1 200 Connection Established\r\n\r\n")?; + stream.flush()?; + + let connection = + rustls::ServerConnection::new(Arc::clone(&config)).map_err(io::Error::other)?; + let mut tls = rustls::StreamOwned::new(connection, stream); + let request = capture_http_request(&mut tls, connect_authority); + match request { + Ok(request) => request_tx + .send(Ok(request)) + .map_err(|_| io::Error::other("request receiver was dropped"))?, + Err(error) => { + let message = error.to_string(); + let _ = request_tx.send(Err(message)); + return Err(error); + } + } + + let response = match request_index { + 0 => concat!( + "HTTP/1.1 200 OK\r\n", + "content-length: 2\r\n", + "connection: close\r\n", + "set-cookie: __cflb=west; Path=/; Secure; HttpOnly\r\n", + "set-cookie: chatgpt_session=secret; Path=/; Secure; HttpOnly\r\n", + "\r\n", + "ok", + ), + 1 => concat!( + "HTTP/1.1 200 OK\r\n", + "content-length: 2\r\n", + "connection: close\r\n", + "set-cookie: __cflb=central; Path=/; Secure; HttpOnly\r\n", + "set-cookie: chatgpt_session=still-secret; Path=/; Secure; HttpOnly\r\n", + "\r\n", + "ok", + ), + _ => concat!( + "HTTP/1.1 200 OK\r\n", + "content-length: 2\r\n", + "connection: close\r\n", + "\r\n", + "ok", + ), + }; + tls.write_all(response.as_bytes())?; + tls.flush()?; + } + Ok(()) +} + +fn configure_stream(stream: &TcpStream) -> io::Result<()> { + stream.set_read_timeout(Some(Duration::from_secs(5)))?; + stream.set_write_timeout(Some(Duration::from_secs(5))) +} + +fn capture_http_request( + stream: &mut impl Read, + connect_authority: String, +) -> io::Result { + let request_head = read_http_head(stream)?; + let mut lines = request_head.lines(); + let request_line = lines + .next() + .ok_or_else(|| io::Error::other("HTTP request should include a request line"))? + .to_string(); + let mut headers: BTreeMap> = BTreeMap::new(); + for line in lines.filter(|line| !line.is_empty()) { + let (name, value) = line + .split_once(':') + .ok_or_else(|| io::Error::other(format!("invalid HTTP header: {line}")))?; + headers + .entry(name.to_ascii_lowercase()) + .or_default() + .push(value.trim().to_string()); + } + Ok(CapturedRequest { + connect_authority, + request_line, + headers, + }) +} + +fn read_http_head(stream: &mut impl Read) -> io::Result { + const MAX_HEADER_BYTES: usize = 64 * 1024; + let mut bytes = Vec::new(); + while !bytes.ends_with(b"\r\n\r\n") { + if bytes.len() == MAX_HEADER_BYTES { + return Err(io::Error::other("HTTP headers exceeded test limit")); + } + let mut byte = [0]; + stream.read_exact(&mut byte)?; + bytes.push(byte[0]); + } + String::from_utf8(bytes).map_err(io::Error::other) +} + +async fn initialize_exec_server(server: &mut ExecServerHarness) -> anyhow::Result<()> { + let initialize_id = server + .send_request( + "initialize", + serde_json::to_value(InitializeParams { + client_name: "cloudflare-affinity-test".to_string(), + resume_session_id: None, + })?, + ) + .await?; + let _: Value = wait_for_response(server, initialize_id).await?; + server + .send_notification("initialized", serde_json::json!({})) + .await +} + +async fn execute_http_request( + server: &mut ExecServerHarness, + url: &str, + request_id: &str, +) -> anyhow::Result { + let response_id = server + .send_request( + "http/request", + serde_json::to_value(HttpRequestParams { + method: "POST".to_string(), + url: url.to_string(), + headers: Vec::new(), + body: None, + timeout_ms: Some(5_000), + redirect_policy: HttpRedirectPolicy::Follow, + request_id: request_id.to_string(), + stream_response: false, + })?, + ) + .await?; + wait_for_response(server, response_id).await +} + +async fn wait_for_response( + server: &mut ExecServerHarness, + request_id: RequestId, +) -> anyhow::Result +where + T: DeserializeOwned, +{ + let message = server + .wait_for_event(|event| match event { + JSONRPCMessage::Response(JSONRPCResponse { id, .. }) + | JSONRPCMessage::Error(codex_exec_server_protocol::JSONRPCError { id, .. }) => { + id == &request_id + } + _ => false, + }) + .await?; + match message { + JSONRPCMessage::Response(JSONRPCResponse { result, .. }) => { + Ok(serde_json::from_value(result)?) + } + JSONRPCMessage::Error(error) => { + anyhow::bail!("exec-server returned an error for {request_id:?}: {error:?}") + } + _ => unreachable!("predicate only accepts responses for the requested id"), + } +} From ab16046c88b2ea9e9af849ab6662129c6becb151 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Thu, 25 Jun 2026 23:46:17 -0700 Subject: [PATCH 17/38] [codex] add process-owned code-mode session client (#30112) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - add `ProcessOwnedCodeModeSessionProvider` and logical session generation/rebinding state - add the supervised child-process connection, reader/writer tasks, and driver state machine - make dropped execute/wait/open callers cancellation-safe with explicit ownership handoff and durable cleanup - validate cell/delegate lifecycle state and reject invalid protocol transitions - add end-to-end stdio coverage for delegates, cancellation, frame limits, child loss, stale generations, replacement, and long-lived sessions ## Why This final stage exposes the process-owned client only after the wire protocol, host-safe runtime, and standalone host are independently in place. Transport failure is fail-stop: the client closes local state, cancels callbacks, reaps the child, and lazily rebuilds a fresh host generation rather than transactionally recovering the old connection. ## Stack This is **4 of 4** in the process-owned code-mode session stack. - Depends on #30111 - Full stack: #30108 → #30110 → #30111 → this PR ## Validation - `just test -p codex-code-mode -p codex-code-mode-host` — 86 passed - `just fix -p codex-code-mode` - `just fix -p codex-code-mode-host` - `just bazel-lock-update` - `just bazel-lock-check` - `bazel test //codex-rs/code-mode:code-mode-unit-tests //codex-rs/code-mode-host:code-mode-host-unit-tests //codex-rs/code-mode-host:code-mode-host-stdio-test //codex-rs/code-mode-protocol:code-mode-protocol-unit-tests` — 4/4 passed - `just fmt` --- codex-rs/Cargo.lock | 3 + codex-rs/code-mode-host/Cargo.toml | 3 + codex-rs/code-mode-host/tests/stdio.rs | 874 ++++++++++ codex-rs/code-mode/Cargo.toml | 2 +- codex-rs/code-mode/src/lib.rs | 3 + codex-rs/code-mode/src/remote_session.rs | 496 ++++++ .../src/remote_session/connection.rs | 458 +++++ .../src/remote_session/connection/driver.rs | 179 ++ .../connection/driver/cell_ids.rs | 111 ++ .../connection/driver/cleanup.rs | 40 + .../connection/driver/commands.rs | 298 ++++ .../connection/driver/delegate_runtime.rs | 339 ++++ .../connection/driver/request_tracker.rs | 186 ++ .../connection/driver/responses.rs | 397 +++++ .../connection/driver/session_registry.rs | 222 +++ .../remote_session/connection/driver/types.rs | 196 +++ .../remote_session/connection/driver_tests.rs | 1514 +++++++++++++++++ .../src/remote_session/connection/reader.rs | 29 + .../code-mode/src/remote_session_tests.rs | 52 + 19 files changed, 5401 insertions(+), 1 deletion(-) create mode 100644 codex-rs/code-mode-host/tests/stdio.rs create mode 100644 codex-rs/code-mode/src/remote_session.rs create mode 100644 codex-rs/code-mode/src/remote_session/connection.rs create mode 100644 codex-rs/code-mode/src/remote_session/connection/driver.rs create mode 100644 codex-rs/code-mode/src/remote_session/connection/driver/cell_ids.rs create mode 100644 codex-rs/code-mode/src/remote_session/connection/driver/cleanup.rs create mode 100644 codex-rs/code-mode/src/remote_session/connection/driver/commands.rs create mode 100644 codex-rs/code-mode/src/remote_session/connection/driver/delegate_runtime.rs create mode 100644 codex-rs/code-mode/src/remote_session/connection/driver/request_tracker.rs create mode 100644 codex-rs/code-mode/src/remote_session/connection/driver/responses.rs create mode 100644 codex-rs/code-mode/src/remote_session/connection/driver/session_registry.rs create mode 100644 codex-rs/code-mode/src/remote_session/connection/driver/types.rs create mode 100644 codex-rs/code-mode/src/remote_session/connection/driver_tests.rs create mode 100644 codex-rs/code-mode/src/remote_session/connection/reader.rs create mode 100644 codex-rs/code-mode/src/remote_session_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index ff0265d4de..9f3e69e3cc 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2519,7 +2519,10 @@ dependencies = [ "anyhow", "codex-code-mode", "codex-code-mode-protocol", + "codex-protocol", + "codex-utils-cargo-bin", "pretty_assertions", + "serde_json", "tokio", "tokio-util", ] diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index d37bd41012..90b043b097 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/codex-rs/code-mode-host/Cargo.toml @@ -24,4 +24,7 @@ tokio = { workspace = true, features = ["io-std", "io-util", "macros", "rt", "sy tokio-util = { workspace = true, features = ["rt"] } [dev-dependencies] +codex-protocol = { workspace = true } +codex-utils-cargo-bin = { workspace = true } pretty_assertions = { workspace = true } +serde_json = { workspace = true } diff --git a/codex-rs/code-mode-host/tests/stdio.rs b/codex-rs/code-mode-host/tests/stdio.rs new file mode 100644 index 0000000000..e7b7506163 --- /dev/null +++ b/codex-rs/code-mode-host/tests/stdio.rs @@ -0,0 +1,874 @@ +#![allow(clippy::expect_used)] + +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::time::Duration; + +#[cfg(unix)] +use std::os::unix::fs::PermissionsExt; + +use codex_code_mode::CellId; +use codex_code_mode::CodeModeNestedToolCall; +use codex_code_mode::CodeModeSession; +use codex_code_mode::CodeModeSessionDelegate; +use codex_code_mode::CodeModeSessionProvider; +use codex_code_mode::CodeModeToolKind; +use codex_code_mode::ExecuteRequest; +use codex_code_mode::FunctionCallOutputContentItem; +use codex_code_mode::NotificationFuture; +use codex_code_mode::ProcessOwnedCodeModeSessionProvider; +use codex_code_mode::RuntimeResponse; +use codex_code_mode::ToolDefinition; +use codex_code_mode::ToolInvocationFuture; +use codex_code_mode::WaitOutcome; +use codex_code_mode::WaitRequest; +use codex_code_mode::host::MAX_FRAME_BYTES; +use codex_protocol::ToolName; +use pretty_assertions::assert_eq; +use serde_json::json; +use tokio::sync::Semaphore; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +#[derive(Default)] +struct RecordingDelegate { + invocations: Mutex>, + notifications: Mutex>, + closed_cells: Mutex>, +} + +#[derive(Debug, Eq, PartialEq)] +enum CallbackEvent { + Started(String), + Cancelled(String), + CellClosed(CellId), +} + +struct CancellationDelegate { + events_tx: mpsc::UnboundedSender, + fast_tool_release: Semaphore, + slow_tool_started: Semaphore, + hold_slow_cleanup: AtomicBool, + slow_cleanup_release: Semaphore, +} + +struct OversizedResultDelegate; + +impl CodeModeSessionDelegate for OversizedResultDelegate { + fn invoke_tool<'a>( + &'a self, + _invocation: CodeModeNestedToolCall, + _cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + Box::pin(async { Ok(json!("x".repeat(MAX_FRAME_BYTES))) }) + } + + fn notify<'a>( + &'a self, + _call_id: String, + _cell_id: CellId, + _text: String, + _cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async { Ok(()) }) + } + + fn cell_closed(&self, _cell_id: &CellId) {} +} + +impl CancellationDelegate { + fn new() -> (Arc, mpsc::UnboundedReceiver) { + let (events_tx, events_rx) = mpsc::unbounded_channel(); + ( + Arc::new(Self { + events_tx, + fast_tool_release: Semaphore::new(/*permits*/ 0), + slow_tool_started: Semaphore::new(/*permits*/ 0), + hold_slow_cleanup: AtomicBool::new(false), + slow_cleanup_release: Semaphore::new(/*permits*/ 0), + }), + events_rx, + ) + } + + #[cfg(unix)] + fn hold_slow_cleanup(&self) { + self.hold_slow_cleanup.store(true, Ordering::Release); + } + + #[cfg(unix)] + fn release_slow_cleanup(&self) { + self.slow_cleanup_release.add_permits(1); + } +} + +impl CodeModeSessionDelegate for CancellationDelegate { + fn invoke_tool<'a>( + &'a self, + invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + Box::pin(async move { + let tool_name = invocation.tool_name.name.clone(); + if tool_name == "tool_call_barrier" { + let permit = self + .slow_tool_started + .acquire() + .await + .map_err(|_| "slow tool barrier closed".to_string())?; + permit.forget(); + return Ok(json!({ "tool": tool_name })); + } + let _ = self + .events_tx + .send(CallbackEvent::Started(tool_name.clone())); + if tool_name == "tool_call_slow" { + self.slow_tool_started.add_permits(1); + cancellation_token.cancelled().await; + let _ = self.events_tx.send(CallbackEvent::Cancelled(tool_name)); + if self.hold_slow_cleanup.load(Ordering::Acquire) { + let permit = self + .slow_cleanup_release + .acquire() + .await + .map_err(|_| "slow tool cleanup release closed".to_string())?; + permit.forget(); + } + return Err("slow tool cancelled".to_string()); + } + let permit = self + .fast_tool_release + .acquire() + .await + .map_err(|_| "fast tool release closed".to_string())?; + permit.forget(); + Ok(json!({ "tool": tool_name })) + }) + } + + fn notify<'a>( + &'a self, + _call_id: String, + _cell_id: CellId, + _text: String, + _cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async { Ok(()) }) + } + + fn cell_closed(&self, cell_id: &CellId) { + let _ = self + .events_tx + .send(CallbackEvent::CellClosed(cell_id.clone())); + } +} + +impl CodeModeSessionDelegate for RecordingDelegate { + fn invoke_tool<'a>( + &'a self, + invocation: CodeModeNestedToolCall, + _cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + self.invocations + .lock() + .expect("invocations lock") + .push(invocation); + Box::pin(async { Ok(json!({ "value": "output" })) }) + } + + fn notify<'a>( + &'a self, + call_id: String, + cell_id: CellId, + text: String, + _cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + self.notifications + .lock() + .expect("notifications lock") + .push((call_id, cell_id, text)); + Box::pin(async { Ok(()) }) + } + + fn cell_closed(&self, cell_id: &CellId) { + self.closed_cells + .lock() + .expect("closed cells lock") + .push(cell_id.clone()); + } +} + +fn cell_id(value: &str) -> CellId { + CellId::new(value.to_string()) +} + +fn execute_request(source: &str) -> ExecuteRequest { + ExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: source.to_string(), + yield_time_ms: None, + max_output_tokens: None, + } +} + +async fn execute(session: &Arc, request: ExecuteRequest) -> RuntimeResponse { + session + .execute(request) + .await + .expect("start execution") + .initial_response() + .await + .expect("initial response") +} + +async fn execute_to_terminal( + session: &Arc, + request: ExecuteRequest, +) -> RuntimeResponse { + let started = session.execute(request).await.expect("start execution"); + let mut response = started.initial_response().await.expect("initial response"); + loop { + match response { + RuntimeResponse::Yielded { cell_id, .. } => { + response = match session + .wait(WaitRequest { + cell_id, + yield_time_ms: 60_000, + }) + .await + .expect("wait for terminal response") + { + WaitOutcome::LiveCell(response) | WaitOutcome::MissingCell(response) => { + response + } + }; + } + response => return response, + } + } +} + +async fn next_callback_event( + events_rx: &mut mpsc::UnboundedReceiver, +) -> CallbackEvent { + tokio::time::timeout(Duration::from_secs(5), events_rx.recv()) + .await + .expect("callback event timeout") + .expect("callback event stream closed") +} + +#[tokio::test] +async fn remote_session_persists_values_forwards_delegates_and_controls_cells() { + let provider = ProcessOwnedCodeModeSessionProvider::with_host_program( + codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"), + ); + let delegate = Arc::new(RecordingDelegate::default()); + let session = provider + .create_session(delegate.clone()) + .await + .expect("create remote session"); + + assert_eq!( + execute(&session, execute_request(r#"store("key", "persisted");"#),).await, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: Vec::new(), + error_text: None, + } + ); + + let mut callback_request = execute_request( + r#" +const result = await tools.echo({ value: String(load("key")) }); +notify("notice"); +text(result.value); +"#, + ); + callback_request.tool_call_id = "call-2".to_string(); + callback_request.enabled_tools = vec![ToolDefinition { + name: "echo".to_string(), + tool_name: ToolName::plain("echo"), + description: String::new(), + kind: CodeModeToolKind::Function, + input_schema: None, + output_schema: None, + }]; + assert_eq!( + execute(&session, callback_request).await, + RuntimeResponse::Result { + cell_id: cell_id("2"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "output".to_string(), + }], + error_text: None, + } + ); + assert_eq!( + *delegate.invocations.lock().expect("invocations lock"), + vec![CodeModeNestedToolCall { + cell_id: cell_id("2"), + runtime_tool_call_id: "tool-1".to_string(), + tool_name: ToolName::plain("echo"), + tool_kind: CodeModeToolKind::Function, + input: Some(json!({ "value": "persisted" })), + }] + ); + assert_eq!( + *delegate.notifications.lock().expect("notifications lock"), + vec![("call-2".to_string(), cell_id("2"), "notice".to_string())] + ); + + let mut pending_request = execute_request("await new Promise(() => {});"); + pending_request.tool_call_id = "call-3".to_string(); + pending_request.yield_time_ms = Some(1); + assert_eq!( + execute(&session, pending_request).await, + RuntimeResponse::Yielded { + cell_id: cell_id("3"), + content_items: Vec::new(), + } + ); + assert_eq!( + session + .wait(WaitRequest { + cell_id: cell_id("3"), + yield_time_ms: 1, + }) + .await + .expect("wait for cell"), + WaitOutcome::LiveCell(RuntimeResponse::Yielded { + cell_id: cell_id("3"), + content_items: Vec::new(), + }) + ); + assert_eq!( + session + .terminate(cell_id("3")) + .await + .expect("terminate cell"), + WaitOutcome::LiveCell(RuntimeResponse::Terminated { + cell_id: cell_id("3"), + content_items: Vec::new(), + }) + ); + + session.shutdown().await.expect("shutdown remote session"); + assert_eq!( + *delegate.closed_cells.lock().expect("closed cells lock"), + vec![cell_id("1"), cell_id("2"), cell_id("3")] + ); +} + +#[tokio::test] +async fn dropping_long_wait_releases_observer_before_next_wait() { + let provider = ProcessOwnedCodeModeSessionProvider::with_host_program( + codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"), + ); + let session = provider + .create_session(Arc::new(RecordingDelegate::default())) + .await + .expect("create remote session"); + let mut request = execute_request("await new Promise(() => {});"); + request.yield_time_ms = Some(1); + let started = session.execute(request).await.expect("start execution"); + let running_cell_id = started.cell_id.clone(); + assert_eq!( + started.initial_response().await.expect("initial response"), + RuntimeResponse::Yielded { + cell_id: running_cell_id.clone(), + content_items: Vec::new(), + } + ); + + let wait_session = Arc::clone(&session); + let wait_cell_id = running_cell_id.clone(); + let first_wait = tokio::spawn(async move { + wait_session + .wait(WaitRequest { + cell_id: wait_cell_id, + yield_time_ms: 60_000, + }) + .await + }); + tokio::time::sleep(Duration::from_millis(100)).await; + first_wait.abort(); + let _ = first_wait.await; + + assert_eq!( + tokio::time::timeout( + Duration::from_secs(2), + session.wait(WaitRequest { + cell_id: running_cell_id.clone(), + yield_time_ms: 1, + }) + ) + .await + .expect("second wait timeout") + .expect("second wait"), + WaitOutcome::LiveCell(RuntimeResponse::Yielded { + cell_id: running_cell_id.clone(), + content_items: Vec::new(), + }) + ); + session + .terminate(running_cell_id) + .await + .expect("terminate cell"); + session.shutdown().await.expect("shutdown remote session"); +} + +#[tokio::test] +async fn unawaited_slow_tool_is_cancelled_after_parallel_tools_complete() { + let provider = ProcessOwnedCodeModeSessionProvider::with_host_program( + codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"), + ); + let (delegate, mut events_rx) = CancellationDelegate::new(); + let session = provider + .create_session(delegate.clone()) + .await + .expect("create remote session"); + let mut request = execute_request( + r#" +await (async () => { +text("hello world"); +yield_control(); +await Promise.all([ + tools.tool_call_a({}), + tools.tool_call_b({}), +]); +text("hello"); +tools.tool_call_slow({}); +await tools.tool_call_barrier({}); +return; +})(); +"#, + ); + request.enabled_tools = [ + "tool_call_a", + "tool_call_b", + "tool_call_slow", + "tool_call_barrier", + ] + .into_iter() + .map(|name| ToolDefinition { + name: name.to_string(), + tool_name: ToolName::plain(name), + description: String::new(), + kind: CodeModeToolKind::Function, + input_schema: None, + output_schema: None, + }) + .collect(); + + let started = session.execute(request).await.expect("start execution"); + let running_cell_id = started.cell_id.clone(); + assert_eq!( + started.initial_response().await.expect("initial response"), + RuntimeResponse::Yielded { + cell_id: running_cell_id.clone(), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "hello world".to_string(), + }], + } + ); + + let wait_session = Arc::clone(&session); + let wait_cell_id = running_cell_id.clone(); + let wait_task = tokio::spawn(async move { + wait_session + .wait(WaitRequest { + cell_id: wait_cell_id, + yield_time_ms: 60_000, + }) + .await + }); + + let mut parallel_tools = vec![ + next_callback_event(&mut events_rx).await, + next_callback_event(&mut events_rx).await, + ]; + parallel_tools.sort_by(|left, right| format!("{left:?}").cmp(&format!("{right:?}"))); + assert_eq!( + parallel_tools, + vec![ + CallbackEvent::Started("tool_call_a".to_string()), + CallbackEvent::Started("tool_call_b".to_string()), + ] + ); + delegate.fast_tool_release.add_permits(2); + + assert_eq!( + next_callback_event(&mut events_rx).await, + CallbackEvent::Started("tool_call_slow".to_string()) + ); + assert_eq!( + next_callback_event(&mut events_rx).await, + CallbackEvent::Cancelled("tool_call_slow".to_string()) + ); + assert_eq!( + next_callback_event(&mut events_rx).await, + CallbackEvent::CellClosed(running_cell_id.clone()) + ); + assert_eq!( + wait_task + .await + .expect("wait task") + .expect("wait for terminal response"), + WaitOutcome::LiveCell(RuntimeResponse::Result { + cell_id: running_cell_id, + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "hello".to_string(), + }], + error_text: None, + }) + ); + session.shutdown().await.expect("shutdown remote session"); +} + +#[tokio::test] +async fn oversized_execute_request_does_not_close_the_shared_host() { + let provider = ProcessOwnedCodeModeSessionProvider::with_host_program( + codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"), + ); + let session = provider + .create_session(Arc::new(RecordingDelegate::default())) + .await + .expect("create remote session"); + let error = session + .execute(execute_request(&"x".repeat(MAX_FRAME_BYTES))) + .await + .err() + .expect("oversized execute should fail"); + assert!( + error.contains("IPC frame limit"), + "unexpected error: {error}" + ); + + assert_eq!( + execute(&session, execute_request(r#"text("still alive");"#)).await, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "still alive".to_string(), + }], + error_text: None, + } + ); + session.shutdown().await.expect("shutdown remote session"); +} + +#[tokio::test] +async fn oversized_delegate_payloads_fail_only_the_tool_call() { + let provider = ProcessOwnedCodeModeSessionProvider::with_host_program( + codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"), + ); + let session = provider + .create_session(Arc::new(OversizedResultDelegate)) + .await + .expect("create remote session"); + let tool = |name: &str| ToolDefinition { + name: name.to_string(), + tool_name: ToolName::plain(name), + description: String::new(), + kind: CodeModeToolKind::Function, + input_schema: None, + output_schema: None, + }; + + let mut oversized_argument = execute_request(&format!( + r#" +try {{ + await tools.big_argument({{ value: "x".repeat({MAX_FRAME_BYTES}) }}); +}} catch (_) {{ + text("argument rejected"); +}} +"# + )); + oversized_argument.enabled_tools = vec![tool("big_argument")]; + oversized_argument.yield_time_ms = Some(60_000); + assert_eq!( + execute_to_terminal(&session, oversized_argument).await, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "argument rejected".to_string(), + }], + error_text: None, + } + ); + + let mut oversized_result = execute_request( + r#" +try { + await tools.big_result({}); +} catch (_) { + text("result rejected"); +} +"#, + ); + oversized_result.enabled_tools = vec![tool("big_result")]; + oversized_result.yield_time_ms = Some(60_000); + assert_eq!( + execute_to_terminal(&session, oversized_result).await, + RuntimeResponse::Result { + cell_id: cell_id("2"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "result rejected".to_string(), + }], + error_text: None, + } + ); + + assert_eq!( + execute(&session, execute_request(r#"text("still alive");"#)).await, + RuntimeResponse::Result { + cell_id: cell_id("3"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "still alive".to_string(), + }], + error_text: None, + } + ); + session.shutdown().await.expect("shutdown remote session"); +} + +#[tokio::test] +async fn oversized_initial_response_does_not_close_the_shared_host() { + let provider = ProcessOwnedCodeModeSessionProvider::with_host_program( + codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"), + ); + let session = provider + .create_session(Arc::new(RecordingDelegate::default())) + .await + .expect("create remote session"); + let started = session + .execute(execute_request(&format!( + r#"text("x".repeat({MAX_FRAME_BYTES}));"# + ))) + .await + .expect("start oversized response"); + let error = started + .initial_response() + .await + .expect_err("oversized initial response should fail"); + assert!( + error.contains("IPC frame limit"), + "unexpected error: {error}" + ); + + assert_eq!( + execute(&session, execute_request(r#"text("still alive");"#)).await, + RuntimeResponse::Result { + cell_id: cell_id("2"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "still alive".to_string(), + }], + error_text: None, + } + ); + session.shutdown().await.expect("shutdown remote session"); +} + +#[cfg(unix)] +#[tokio::test] +async fn child_process_loss_cleans_up_and_rebuilds_the_shared_host() { + let host_program = + codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"); + let proxy_dir = + std::env::temp_dir().join(format!("codex-code-mode-host-loss-{}", std::process::id())); + let proxy_program = proxy_dir.join("host-proxy.sh"); + let pid_path = proxy_dir.join("host.pid"); + let _ = std::fs::remove_dir_all(&proxy_dir); + std::fs::create_dir_all(&proxy_dir).expect("create host proxy directory"); + std::fs::write( + &proxy_program, + format!( + "#!/bin/sh\nprintf '%s\\n' \"$$\" > '{}'\nexec '{}'\n", + pid_path.display(), + host_program.display() + ), + ) + .expect("write host proxy"); + let mut permissions = std::fs::metadata(&proxy_program) + .expect("host proxy metadata") + .permissions(); + permissions.set_mode(/*mode*/ 0o700); + std::fs::set_permissions(&proxy_program, permissions).expect("make host proxy executable"); + + let provider = ProcessOwnedCodeModeSessionProvider::with_host_program(proxy_program); + let (delegate_a, mut events_a) = CancellationDelegate::new(); + delegate_a.hold_slow_cleanup(); + let delegate_b = Arc::new(RecordingDelegate::default()); + let session_a = provider + .create_session(delegate_a.clone()) + .await + .expect("create first remote session"); + let session_b = provider + .create_session(delegate_b.clone()) + .await + .expect("create second remote session"); + + let mut request_a = execute_request("await tools.tool_call_slow({});"); + request_a.yield_time_ms = Some(1); + request_a.enabled_tools = vec![ToolDefinition { + name: "tool_call_slow".to_string(), + tool_name: ToolName::plain("tool_call_slow"), + description: String::new(), + kind: CodeModeToolKind::Function, + input_schema: None, + output_schema: None, + }]; + let started_a = session_a + .execute(request_a) + .await + .expect("start first cell"); + let cell_a = started_a.cell_id.clone(); + assert_eq!( + started_a + .initial_response() + .await + .expect("first initial response"), + RuntimeResponse::Yielded { + cell_id: cell_a.clone(), + content_items: Vec::new(), + } + ); + assert_eq!( + next_callback_event(&mut events_a).await, + CallbackEvent::Started("tool_call_slow".to_string()) + ); + + let mut request_b = execute_request("await new Promise(() => {});"); + request_b.yield_time_ms = Some(1); + let started_b = session_b + .execute(request_b) + .await + .expect("start second cell"); + let cell_b = started_b.cell_id.clone(); + assert_eq!( + started_b + .initial_response() + .await + .expect("second initial response"), + RuntimeResponse::Yielded { + cell_id: cell_b.clone(), + content_items: Vec::new(), + } + ); + + let wait_a_session = Arc::clone(&session_a); + let wait_a_cell = cell_a.clone(); + let wait_a = tokio::spawn(async move { + wait_a_session + .wait(WaitRequest { + cell_id: wait_a_cell, + yield_time_ms: 60_000, + }) + .await + }); + let wait_b_session = Arc::clone(&session_b); + let wait_b_cell = cell_b.clone(); + let wait_b = tokio::spawn(async move { + wait_b_session + .wait(WaitRequest { + cell_id: wait_b_cell, + yield_time_ms: 60_000, + }) + .await + }); + tokio::time::sleep(Duration::from_millis(100)).await; + + let pid = tokio::time::timeout(Duration::from_secs(5), async { + loop { + if let Ok(pid) = std::fs::read_to_string(&pid_path) + && let Ok(pid) = pid.trim().parse::() + { + break pid; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("host pid timeout"); + let kill_status = std::process::Command::new("kill") + .args(["-KILL", &pid.to_string()]) + .status() + .expect("kill host process"); + assert!(kill_status.success()); + + assert!( + tokio::time::timeout(Duration::from_secs(5), wait_a) + .await + .expect("first wait failure timeout") + .expect("first wait task") + .is_err() + ); + assert!( + tokio::time::timeout(Duration::from_secs(5), wait_b) + .await + .expect("second wait failure timeout") + .expect("second wait task") + .is_err() + ); + let closure_events = [ + next_callback_event(&mut events_a).await, + next_callback_event(&mut events_a).await, + ]; + assert!(closure_events.contains(&CallbackEvent::Cancelled("tool_call_slow".to_string()))); + assert!(closure_events.contains(&CallbackEvent::CellClosed(cell_a.clone()))); + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if delegate_b + .closed_cells + .lock() + .expect("closed cells lock") + .contains(&cell_b) + { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("unrelated session cleanup timeout"); + + assert_eq!( + execute(&session_b, execute_request(r#"text("replacement");"#)).await, + RuntimeResponse::Result { + cell_id: cell_id("g2:1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "replacement".to_string(), + }], + error_text: None, + } + ); + let stale_error = session_b + .wait(WaitRequest { + cell_id: cell_b.clone(), + yield_time_ms: 1, + }) + .await + .expect_err("stale cell should be rejected"); + assert!(stale_error.contains("stale code-mode host generation")); + + tokio::time::timeout(Duration::from_secs(5), session_a.shutdown()) + .await + .expect("failed session shutdown timeout") + .expect("shutdown failed session"); + tokio::time::timeout(Duration::from_secs(5), session_b.shutdown()) + .await + .expect("unrelated session shutdown timeout") + .expect("shutdown replacement session"); + + delegate_a.release_slow_cleanup(); + tokio::task::yield_now().await; + assert!(matches!( + events_a.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + + std::fs::remove_dir_all(proxy_dir).expect("remove host proxy directory"); +} diff --git a/codex-rs/code-mode/Cargo.toml b/codex-rs/code-mode/Cargo.toml index 6b91e64fe2..011fee6ebc 100644 --- a/codex-rs/code-mode/Cargo.toml +++ b/codex-rs/code-mode/Cargo.toml @@ -21,7 +21,7 @@ codex-protocol = { workspace = true } deno_core_icudata = { workspace = true } futures = { workspace = true } serde_json = { workspace = true } -tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } +tokio = { workspace = true, features = ["io-util", "macros", "process", "rt", "sync", "time"] } tokio-util = { workspace = true, features = ["rt"] } tracing = { workspace = true } v8 = { workspace = true } diff --git a/codex-rs/code-mode/src/lib.rs b/codex-rs/code-mode/src/lib.rs index 60c02e4fae..653d7b1608 100644 --- a/codex-rs/code-mode/src/lib.rs +++ b/codex-rs/code-mode/src/lib.rs @@ -1,4 +1,5 @@ mod cell_actor; +mod remote_session; mod runtime; mod service; mod session_runtime; @@ -6,6 +7,8 @@ mod session_runtime; pub(crate) type TaskFailureHandler = std::sync::Arc; pub use codex_code_mode_protocol::*; +pub use remote_session::ProcessOwnedCodeModeSession; +pub use remote_session::ProcessOwnedCodeModeSessionProvider; pub use service::InProcessCodeModeSession; pub use service::InProcessCodeModeSessionProvider; pub use service::NoopCodeModeSessionDelegate; diff --git a/codex-rs/code-mode/src/remote_session.rs b/codex-rs/code-mode/src/remote_session.rs new file mode 100644 index 0000000000..d942b377c1 --- /dev/null +++ b/codex-rs/code-mode/src/remote_session.rs @@ -0,0 +1,496 @@ +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::Mutex as StdMutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeSession; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::CodeModeSessionProvider; +use codex_code_mode_protocol::CodeModeSessionProviderFuture; +use codex_code_mode_protocol::CodeModeSessionResultFuture; +use codex_code_mode_protocol::ExecuteRequest; +use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::WaitOutcome; +use codex_code_mode_protocol::WaitRequest; +use codex_code_mode_protocol::host::SessionId; +use tokio::sync::Semaphore; +use tokio::sync::watch; + +use self::connection::Connection; +use self::connection::RemoteSession; +use self::connection::SessionCleanup; +use crate::NoopCodeModeSessionDelegate; + +mod connection; + +const CODE_MODE_HOST_PATH_ENV: &str = "CODEX_CODE_MODE_HOST_PATH"; + +type ShutdownResultReceiver = watch::Receiver>>; + +/// Creates code-mode sessions backed by one lazily spawned process host. +pub struct ProcessOwnedCodeModeSessionProvider { + host_program: PathBuf, + process_host: StdMutex>>, +} + +impl ProcessOwnedCodeModeSessionProvider { + pub fn with_host_program(host_program: PathBuf) -> Self { + Self { + host_program, + process_host: StdMutex::new(None), + } + } + + fn process_host(&self) -> Arc { + let mut process_host = self + .process_host + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(process_host) = process_host.as_ref() { + return Arc::clone(process_host); + } + + let new_process_host = Arc::new(OwnedProcessHost::new(self.host_program.clone())); + *process_host = Some(Arc::clone(&new_process_host)); + new_process_host + } +} + +impl Default for ProcessOwnedCodeModeSessionProvider { + fn default() -> Self { + Self::with_host_program(default_host_program()) + } +} + +impl CodeModeSessionProvider for ProcessOwnedCodeModeSessionProvider { + fn create_session<'a>( + &'a self, + delegate: Arc, + ) -> CodeModeSessionProviderFuture<'a> { + let session = ProcessOwnedCodeModeSession::with_process_host(delegate, self.process_host()); + Box::pin(async move { + session.connection().await?; + let session: Arc = Arc::new(session); + Ok(session) + }) + } +} + +struct OwnedProcessHost { + host_program: PathBuf, + connection: StdMutex>>, + spawn_permit: Semaphore, + next_session_id: AtomicU64, +} + +impl OwnedProcessHost { + fn new(host_program: PathBuf) -> Self { + Self { + host_program, + connection: StdMutex::new(None), + spawn_permit: Semaphore::new(/*permits*/ 1), + next_session_id: AtomicU64::new(1), + } + } + + async fn connection(&self) -> Result, String> { + if let Some(connection) = self.live_connection() { + return Ok(connection); + } + + let _spawn_permit = self + .spawn_permit + .acquire() + .await + .map_err(|_| "code-mode host spawn coordinator closed".to_string())?; + if let Some(connection) = self.live_connection() { + return Ok(connection); + } + let new_connection = Arc::new(Connection::spawn(&self.host_program).await?); + *self + .connection + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(Arc::clone(&new_connection)); + Ok(new_connection) + } + + fn live_connection(&self) -> Option> { + self.connection + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .as_ref() + .filter(|connection| connection.is_alive()) + .cloned() + } + + fn allocate_session_id(&self) -> SessionId { + let value = self.next_session_id.fetch_add(1, Ordering::Relaxed); + match SessionId::new(format!("session-{value}")) { + Ok(session_id) => session_id, + Err(_) => unreachable!("a generated code-mode session ID is nonempty"), + } + } +} + +enum SessionState { + New, + Opening { + remote: RemoteSession, + result_rx: watch::Receiver>>, + }, + Open(SessionBinding), + Closing, + Closed, +} + +#[derive(Clone)] +struct SessionBinding { + connection: Arc, + remote: RemoteSession, + cleanup: SessionCleanup, +} + +struct SessionInner { + process_host: Arc, + delegate: Arc, + state: StdMutex, + next_generation: AtomicU64, + shutdown_requested: AtomicBool, + shutdown_result: StdMutex>, + retired_cleanups: StdMutex>, +} + +/// A logical code-mode session assigned to a process-owned host. +pub struct ProcessOwnedCodeModeSession { + inner: Arc, +} + +impl ProcessOwnedCodeModeSession { + pub fn new() -> Self { + Self::with_process_host( + Arc::new(NoopCodeModeSessionDelegate), + Arc::new(OwnedProcessHost::new(default_host_program())), + ) + } + + fn with_process_host( + delegate: Arc, + process_host: Arc, + ) -> Self { + Self { + inner: Arc::new(SessionInner { + process_host, + delegate, + state: StdMutex::new(SessionState::New), + next_generation: AtomicU64::new(1), + shutdown_requested: AtomicBool::new(false), + shutdown_result: StdMutex::new(None), + retired_cleanups: StdMutex::new(Vec::new()), + }), + } + } + + async fn connection(&self) -> Result { + self.inner.connection().await + } + + pub async fn execute(&self, request: ExecuteRequest) -> Result { + let binding = self.connection().await?; + binding.connection.execute(binding.remote, request).await + } + + pub async fn wait(&self, request: WaitRequest) -> Result { + let binding = self.connection().await?; + binding.connection.wait(binding.remote, request).await + } + + pub async fn terminate(&self, cell_id: CellId) -> Result { + let binding = self.connection().await?; + binding.connection.terminate(binding.remote, cell_id).await + } + + pub async fn shutdown(&self) -> Result<(), String> { + wait_for_watch(self.inner.request_shutdown()).await + } +} + +impl SessionInner { + async fn connection(self: &Arc) -> Result { + loop { + if self.shutdown_requested.load(Ordering::Acquire) { + return Err("code mode session is shutting down".to_string()); + } + let (result_rx, start) = { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + match &*state { + SessionState::New => { + let generation = self.next_generation.fetch_add(1, Ordering::Relaxed); + let remote = RemoteSession { + id: self.process_host.allocate_session_id(), + generation, + }; + let (result_tx, result_rx) = watch::channel(None); + *state = SessionState::Opening { + remote: remote.clone(), + result_rx: result_rx.clone(), + }; + (result_rx, Some((remote, result_tx))) + } + SessionState::Opening { result_rx, .. } => (result_rx.clone(), None), + SessionState::Open(binding) if binding.connection.is_alive() => { + return Ok(binding.clone()); + } + SessionState::Open(binding) => { + self.retain_cleanup(binding.cleanup.clone()); + *state = SessionState::New; + continue; + } + SessionState::Closing | SessionState::Closed => { + return Err("code mode session is shutting down".to_string()); + } + } + }; + if let Some((remote, result_tx)) = start { + let inner = Arc::clone(self); + tokio::spawn(async move { + inner.open(remote, result_tx).await; + }); + } + return wait_for_watch(result_rx).await; + } + } + + async fn open( + self: Arc, + remote: RemoteSession, + result_tx: watch::Sender>>, + ) { + let result = match self.process_host.connection().await { + Ok(connection) => { + let cleanup = connection + .open_session(remote.clone(), Arc::clone(&self.delegate)) + .await; + cleanup.map(|cleanup| SessionBinding { + connection, + remote: remote.clone(), + cleanup, + }) + } + Err(err) => Err(err), + }; + { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if matches!( + &*state, + SessionState::Opening { + remote: opening_remote, + .. + } if opening_remote == &remote + ) { + *state = match &result { + Ok(binding) => SessionState::Open(binding.clone()), + Err(_) => SessionState::New, + }; + } + } + result_tx.send_replace(Some(result)); + } + + fn request_shutdown(self: &Arc) -> ShutdownResultReceiver { + self.shutdown_requested.store(true, Ordering::Release); + let mut shutdown_result = self + .shutdown_result + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(result_rx) = shutdown_result.as_ref() { + return result_rx.clone(); + } + let (result_tx, result_rx) = watch::channel(None); + *shutdown_result = Some(result_rx.clone()); + let inner = Arc::clone(self); + tokio::spawn(async move { + let result = inner.drive_shutdown().await; + result_tx.send_replace(Some(result)); + }); + result_rx + } + + async fn drive_shutdown(self: &Arc) -> Result<(), String> { + loop { + let action = { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + match &*state { + SessionState::New => { + *state = SessionState::Closed; + ShutdownAction::Finish + } + SessionState::Opening { result_rx, .. } => { + ShutdownAction::WaitForOpen(result_rx.clone()) + } + SessionState::Open(binding) if !binding.connection.is_alive() => { + let cleanup = binding.cleanup.clone(); + *state = SessionState::Closing; + ShutdownAction::WaitForSessionCleanup(cleanup) + } + SessionState::Open(binding) => { + let binding = binding.clone(); + *state = SessionState::Closing; + ShutdownAction::Close(binding) + } + SessionState::Closing => { + return Err("code-mode session shutdown driver entered twice".to_string()); + } + SessionState::Closed => return Ok(()), + } + }; + match action { + ShutdownAction::WaitForOpen(result_rx) => { + let _ = wait_for_watch(result_rx).await; + } + ShutdownAction::Finish => { + self.wait_for_retired_cleanups().await; + return Ok(()); + } + ShutdownAction::WaitForSessionCleanup(cleanup) => { + cleanup.wait().await; + self.wait_for_retired_cleanups().await; + *self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = SessionState::Closed; + return Ok(()); + } + ShutdownAction::Close(binding) => { + let result = binding.connection.shutdown_session(binding.remote).await; + if result.is_err() && !binding.connection.is_alive() { + binding.cleanup.wait().await; + } + self.wait_for_retired_cleanups().await; + *self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = SessionState::Closed; + return result; + } + } + } + } + + fn retain_cleanup(&self, cleanup: SessionCleanup) { + let mut retired = self + .retired_cleanups + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + retired.retain(|cleanup| !cleanup.is_complete()); + if !cleanup.is_complete() { + retired.push(cleanup); + } + } + + async fn wait_for_retired_cleanups(&self) { + let retired = std::mem::take( + &mut *self + .retired_cleanups + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + ); + for cleanup in retired { + cleanup.wait().await; + } + } +} + +enum ShutdownAction { + WaitForOpen(watch::Receiver>>), + Finish, + WaitForSessionCleanup(SessionCleanup), + Close(SessionBinding), +} + +async fn wait_for_watch( + mut result_rx: watch::Receiver>>, +) -> Result +where + T: Clone, +{ + loop { + if let Some(result) = result_rx.borrow().clone() { + return result; + } + result_rx + .changed() + .await + .map_err(|_| "code-mode session transition stopped".to_string())?; + } +} + +impl Drop for ProcessOwnedCodeModeSession { + fn drop(&mut self) { + if tokio::runtime::Handle::try_current().is_ok() { + self.inner.request_shutdown(); + } + } +} + +impl Default for ProcessOwnedCodeModeSession { + fn default() -> Self { + Self::new() + } +} + +impl CodeModeSession for ProcessOwnedCodeModeSession { + fn execute<'a>( + &'a self, + request: ExecuteRequest, + ) -> CodeModeSessionResultFuture<'a, StartedCell> { + Box::pin(ProcessOwnedCodeModeSession::execute(self, request)) + } + + fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(ProcessOwnedCodeModeSession::wait(self, request)) + } + + fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(ProcessOwnedCodeModeSession::terminate(self, cell_id)) + } + + fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> { + Box::pin(ProcessOwnedCodeModeSession::shutdown(self)) + } +} + +fn default_host_program() -> PathBuf { + if let Some(path) = std::env::var_os(CODE_MODE_HOST_PATH_ENV) { + return PathBuf::from(path); + } + let executable_name = if cfg!(windows) { + "codex-code-mode-host.exe" + } else { + "codex-code-mode-host" + }; + if let Ok(current_exe) = std::env::current_exe() + && let Some(parent) = current_exe.parent() + { + let sibling = parent.join(executable_name); + if sibling.is_file() { + return sibling; + } + } + PathBuf::from(executable_name) +} + +#[cfg(test)] +#[path = "remote_session_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode/src/remote_session/connection.rs b/codex-rs/code-mode/src/remote_session/connection.rs new file mode 100644 index 0000000000..7154784def --- /dev/null +++ b/codex-rs/code-mode/src/remote_session/connection.rs @@ -0,0 +1,458 @@ +use std::path::Path; +use std::process::Stdio; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::ExecuteRequest; +use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::WaitOutcome; +use codex_code_mode_protocol::WaitRequest; +use codex_code_mode_protocol::host::CapabilitySet; +use codex_code_mode_protocol::host::ClientHello; +use codex_code_mode_protocol::host::ClientToHost; +use codex_code_mode_protocol::host::EncodedFrame; +use codex_code_mode_protocol::host::FramedReader; +use codex_code_mode_protocol::host::FramedWriter; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::ProtocolVersion; +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::SupportedProtocolVersions; +use tokio::io::AsyncBufReadExt; +use tokio::io::BufReader; +use tokio::process::Child; +use tokio::process::Command; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; +use tracing::debug; +use tracing::warn; + +use self::driver::ConnectionDriver; +use self::driver::DriverCommand; +use self::driver::DriverEvent; +use self::driver::DriverLifecycle; +pub(super) use self::driver::RemoteSession; +pub(super) use self::driver::SessionCleanup; +use self::reader::drive_reader; + +mod driver; +mod reader; + +const IPC_CHANNEL_CAPACITY: usize = 128; +const HOST_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10); + +pub(super) struct Connection { + command_tx: mpsc::Sender, + execute_claim_tx: mpsc::UnboundedSender, + alive: Arc, + failure: Arc>>, + cancellation: CancellationToken, +} + +struct CallerCancellation { + token: CancellationToken, + armed: bool, +} + +struct ConnectionSupervisor { + child: Child, + event_tx: mpsc::Sender, + cancellation: CancellationToken, + alive: Arc, + failure: Arc>>, + driver_task: JoinHandle<()>, + reader_task: JoinHandle>, + writer_task: JoinHandle>, +} + +impl CallerCancellation { + fn new() -> Self { + Self { + token: CancellationToken::new(), + armed: true, + } + } + + fn token(&self) -> CancellationToken { + self.token.clone() + } + + fn disarm(mut self) { + self.armed = false; + } +} + +impl Drop for CallerCancellation { + fn drop(&mut self) { + if self.armed { + self.token.cancel(); + } + } +} + +impl Connection { + pub(super) async fn spawn(host_program: &Path) -> Result { + let mut command = Command::new(host_program); + #[cfg(unix)] + command.process_group(0); + let mut child = command + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .spawn() + .map_err(|err| { + format!( + "failed to spawn code-mode host {}: {err}", + host_program.display() + ) + })?; + + if let Some(stderr) = child.stderr.take() { + tokio::spawn(async move { + let mut lines = BufReader::new(stderr).lines(); + loop { + match lines.next_line().await { + Ok(Some(line)) => debug!("code-mode host stderr: {line}"), + Ok(None) => break, + Err(err) => { + warn!("failed to read code-mode host stderr: {err}"); + break; + } + } + } + }); + } + + let stdin = child + .stdin + .take() + .ok_or_else(|| "spawned code-mode host has no stdin".to_string())?; + let stdout = child + .stdout + .take() + .ok_or_else(|| "spawned code-mode host has no stdout".to_string())?; + let mut reader = FramedReader::new(stdout); + let mut writer = FramedWriter::new(stdin); + let handshake = async { + let hello = ClientHello::new( + SupportedProtocolVersions::try_new([ProtocolVersion::V1]) + .map_err(|err| err.to_string())?, + CapabilitySet::empty(), + CapabilitySet::empty(), + ) + .map_err(|err| err.to_string())?; + writer + .write(&ClientToHost::ClientHello(hello)) + .await + .map_err(|err| format!("failed to write code-mode host hello: {err}"))?; + match reader + .read::() + .await + .map_err(|err| format!("failed to read code-mode host hello: {err}"))? + { + Some(HostToClient::HostHello(hello)) + if hello.selected_version() == ProtocolVersion::V1 => + { + Ok(()) + } + Some(HostToClient::HandshakeRejected { reason }) => { + Err(format!("code-mode host rejected the handshake: {reason:?}")) + } + Some(message) => Err(format!( + "code-mode host returned an invalid handshake response: {message:?}" + )), + None => Err("code-mode host exited during handshake".to_string()), + } + }; + let handshake_result = match tokio::time::timeout(HOST_HANDSHAKE_TIMEOUT, handshake).await { + Ok(result) => result, + Err(_) => { + kill_and_reap(&mut child).await; + return Err("timed out negotiating with the code-mode host".to_string()); + } + }; + if let Err(err) = handshake_result { + kill_and_reap(&mut child).await; + return Err(err); + } + + let (command_tx, command_rx) = mpsc::channel(IPC_CHANNEL_CAPACITY); + let (event_tx, event_rx) = mpsc::channel(IPC_CHANNEL_CAPACITY); + let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(IPC_CHANNEL_CAPACITY); + let cancellation = CancellationToken::new(); + let alive = Arc::new(AtomicBool::new(true)); + let failure = Arc::new(std::sync::Mutex::new(None)); + + let writer_cancellation = cancellation.clone(); + let writer_task = tokio::spawn(async move { + loop { + tokio::select! { + _ = writer_cancellation.cancelled() => return Ok(()), + frame = outgoing_rx.recv() => { + let Some(frame) = frame else { + return Err("code-mode host outgoing stream closed".to_string()); + }; + if let Err(err) = writer.write_frame(&frame).await { + return Err(format!("failed to write code-mode host message: {err}")); + } + } + } + } + }); + + let reader_events = event_tx.clone(); + let reader_cancellation = cancellation.clone(); + let reader_task = + tokio::spawn( + async move { drive_reader(reader, reader_events, reader_cancellation).await }, + ); + + let (driver, execute_claim_tx) = ConnectionDriver::new( + command_rx, + event_rx, + event_tx.clone(), + outgoing_tx, + DriverLifecycle { + alive: Arc::clone(&alive), + failure: Arc::clone(&failure), + cancellation: cancellation.clone(), + }, + ); + let driver_task = tokio::spawn(driver.run()); + tokio::spawn( + ConnectionSupervisor { + child, + event_tx, + cancellation: cancellation.clone(), + alive: Arc::clone(&alive), + failure: Arc::clone(&failure), + driver_task, + reader_task, + writer_task, + } + .run(), + ); + + Ok(Self { + command_tx, + execute_claim_tx, + alive, + failure, + cancellation, + }) + } + + pub(super) fn is_alive(&self) -> bool { + if self.command_tx.is_closed() { + mark_connection_dead( + &self.alive, + &self.failure, + "code-mode connection driver closed".to_string(), + ); + } + self.alive.load(Ordering::Acquire) + } + + pub(super) async fn open_session( + &self, + session: RemoteSession, + delegate: Arc, + ) -> Result { + let cleanup = SessionCleanup::new(); + let cancellation = CallerCancellation::new(); + let (response_tx, response_rx) = oneshot::channel(); + self.send(DriverCommand::OpenSession { + session, + delegate, + cleanup: cleanup.clone(), + caller_cancellation: cancellation.token(), + response_tx, + }) + .await?; + let result = self.receive(response_rx).await; + cancellation.disarm(); + result?; + Ok(cleanup) + } + + pub(super) async fn execute( + &self, + session: RemoteSession, + request: ExecuteRequest, + ) -> Result { + let cancellation = CallerCancellation::new(); + let (response_tx, response_rx) = oneshot::channel(); + self.send(DriverCommand::Execute { + session, + request, + caller_cancellation: cancellation.token(), + response_tx, + }) + .await?; + let delivered = match self.receive(response_rx).await { + Ok(delivered) => delivered, + Err(err) => { + cancellation.disarm(); + return Err(err); + } + }; + self.execute_claim_tx + .send(delivered.request_id) + .map_err(|_| self.failure_message())?; + cancellation.disarm(); + Ok(delivered.started) + } + + pub(super) async fn wait( + &self, + session: RemoteSession, + request: WaitRequest, + ) -> Result { + let cancellation = CallerCancellation::new(); + let (response_tx, response_rx) = oneshot::channel(); + self.send(DriverCommand::Wait { + session, + request, + caller_cancellation: cancellation.token(), + response_tx, + }) + .await?; + let result = self.receive(response_rx).await; + cancellation.disarm(); + result + } + + pub(super) async fn terminate( + &self, + session: RemoteSession, + cell_id: CellId, + ) -> Result { + let (response_tx, response_rx) = oneshot::channel(); + self.send(DriverCommand::Terminate { + session, + cell_id, + response_tx, + }) + .await?; + self.receive(response_rx).await + } + + pub(super) async fn shutdown_session(&self, session: RemoteSession) -> Result<(), String> { + let (response_tx, response_rx) = oneshot::channel(); + self.send(DriverCommand::ShutdownSession { + session, + response_tx, + }) + .await?; + self.receive(response_rx).await + } + + async fn send(&self, command: DriverCommand) -> Result<(), String> { + if !self.is_alive() { + return Err(self.failure_message()); + } + self.command_tx + .send(command) + .await + .map_err(|_| self.failure_message()) + } + + async fn receive( + &self, + response_rx: oneshot::Receiver>, + ) -> Result { + response_rx.await.map_err(|_| self.failure_message())? + } + + fn failure_message(&self) -> String { + self.failure + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + .unwrap_or_else(|| "code-mode host connection closed".to_string()) + } +} + +impl Drop for Connection { + fn drop(&mut self) { + mark_connection_dead( + &self.alive, + &self.failure, + "code-mode host connection closed".to_string(), + ); + self.cancellation.cancel(); + } +} + +impl ConnectionSupervisor { + async fn run(mut self) { + let mut child_exited = false; + let reason = tokio::select! { + biased; + _ = self.cancellation.cancelled() => failure_message(&self.failure), + result = &mut self.driver_task => match result { + Ok(()) => "code-mode connection driver exited unexpectedly".to_string(), + Err(err) => format!("code-mode connection driver task failed: {err}"), + }, + result = &mut self.reader_task => task_failure("reader", result), + result = &mut self.writer_task => task_failure("writer", result), + result = self.child.wait() => { + child_exited = true; + match result { + Ok(status) => format!("code-mode host exited with status {status}"), + Err(err) => format!("failed waiting for code-mode host: {err}"), + } + } + }; + mark_connection_dead(&self.alive, &self.failure, reason.clone()); + let _ = self.event_tx.try_send(DriverEvent::Failed(reason)); + self.cancellation.cancel(); + if !child_exited { + kill_and_reap(&mut self.child).await; + } + } +} + +fn task_failure( + task_name: &str, + result: Result, tokio::task::JoinError>, +) -> String { + match result { + Ok(Ok(())) => format!("code-mode connection {task_name} exited unexpectedly"), + Ok(Err(err)) => err, + Err(err) => format!("code-mode connection {task_name} task failed: {err}"), + } +} + +fn mark_connection_dead( + alive: &AtomicBool, + failure: &std::sync::Mutex>, + reason: String, +) { + alive.store(false, Ordering::Release); + let mut failure = failure + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if failure.is_none() { + *failure = Some(reason); + } +} + +fn failure_message(failure: &std::sync::Mutex>) -> String { + failure + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + .unwrap_or_else(|| "code-mode host connection closed".to_string()) +} + +async fn kill_and_reap(child: &mut Child) { + let _ = child.start_kill(); + let _ = child.wait().await; +} diff --git a/codex-rs/code-mode/src/remote_session/connection/driver.rs b/codex-rs/code-mode/src/remote_session/connection/driver.rs new file mode 100644 index 0000000000..102ad47581 --- /dev/null +++ b/codex-rs/code-mode/src/remote_session/connection/driver.rs @@ -0,0 +1,179 @@ +use std::panic::AssertUnwindSafe; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::host::EncodedFrame; +use codex_code_mode_protocol::host::RequestId; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +pub(in crate::remote_session) use self::cleanup::SessionCleanup; +use self::delegate_runtime::DelegateRuntime; +use self::request_tracker::RequestTracker; +use self::session_registry::SessionRegistry; +pub(super) use self::types::DriverCommand; +pub(super) use self::types::DriverEvent; +pub(in crate::remote_session) use self::types::RemoteSession; + +mod cell_ids; +mod cleanup; +mod commands; +mod delegate_runtime; +mod request_tracker; +mod responses; +mod session_registry; +mod types; + +pub(super) struct DriverLifecycle { + pub(super) alive: Arc, + pub(super) failure: Arc>>, + pub(super) cancellation: CancellationToken, +} + +pub(super) struct ConnectionDriver { + command_rx: mpsc::Receiver, + event_rx: mpsc::Receiver, + event_tx: mpsc::Sender, + execute_claim_rx: mpsc::UnboundedReceiver, + outgoing_tx: mpsc::Sender, + requests: RequestTracker, + sessions: SessionRegistry, + delegates: DelegateRuntime, + alive: Arc, + failure: Arc>>, + cancellation: CancellationToken, + failed: bool, +} + +impl ConnectionDriver { + pub(super) fn new( + command_rx: mpsc::Receiver, + event_rx: mpsc::Receiver, + event_tx: mpsc::Sender, + outgoing_tx: mpsc::Sender, + lifecycle: DriverLifecycle, + ) -> (Self, mpsc::UnboundedSender) { + let (execute_claim_tx, execute_claim_rx) = mpsc::unbounded_channel(); + ( + Self { + command_rx, + event_rx, + event_tx: event_tx.clone(), + execute_claim_rx, + outgoing_tx, + requests: RequestTracker::new(), + sessions: SessionRegistry::new(), + delegates: DelegateRuntime::new(event_tx), + alive: lifecycle.alive, + failure: lifecycle.failure, + cancellation: lifecycle.cancellation, + failed: false, + }, + execute_claim_tx, + ) + } + + pub(super) async fn run(mut self) { + loop { + tokio::select! { + biased; + _ = self.cancellation.cancelled() => { + self.fail("code-mode host connection closed".to_string()); + return; + } + event = self.event_rx.recv() => { + let Some(event) = event else { + self.fail("code-mode host event stream closed".to_string()); + return; + }; + if !self.cancel_dropped_callers() || !self.handle_event(event) { + return; + } + } + claim = self.execute_claim_rx.recv() => { + let Some(request_id) = claim else { + self.fail("code-mode execute claim stream closed".to_string()); + return; + }; + self.requests.claim_execute(request_id); + } + command = self.command_rx.recv() => { + let Some(command) = command else { + self.fail("code-mode host command stream closed".to_string()); + return; + }; + if !self.cancel_dropped_callers() || !self.handle_command(command) { + return; + } + } + } + } + } + + fn handle_event(&mut self, event: DriverEvent) -> bool { + let keep_running = match event { + DriverEvent::HostMessage(message) => self.handle_host_message(message), + DriverEvent::DelegateCompleted { id, result } => self.complete_delegate(id, result), + DriverEvent::RequestCancelled(id) => self.cancel_request(id), + DriverEvent::Failed(reason) => { + self.fail(reason); + false + } + }; + if keep_running { + self.flush_deferred_waits() + } else { + false + } + } + + fn queue_frame(&mut self, frame: EncodedFrame) -> bool { + match self.outgoing_tx.try_send(frame) { + Ok(()) => true, + Err(mpsc::error::TrySendError::Full(_)) => { + self.fail("code-mode host outgoing queue is full".to_string()); + false + } + Err(mpsc::error::TrySendError::Closed(_)) => { + self.fail("code-mode host writer closed".to_string()); + false + } + } + } + + fn fail(&mut self, reason: String) { + if self.failed { + return; + } + self.failed = true; + self.alive.store(false, Ordering::Release); + let reason = { + let mut failure = self + .failure + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + failure.get_or_insert(reason).clone() + }; + self.requests.fail_all(&reason); + let failed_sessions = self.sessions.drain(); + self.delegates.fail_all(failed_sessions); + self.cancellation.cancel(); + } +} + +impl Drop for ConnectionDriver { + fn drop(&mut self) { + self.fail("code-mode connection driver stopped unexpectedly".to_string()); + } +} + +fn notify_cell_closed(delegate: &Arc, cell_id: &CellId) { + let _ = std::panic::catch_unwind(AssertUnwindSafe(|| delegate.cell_closed(cell_id))); +} + +#[cfg(test)] +#[path = "driver_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/cell_ids.rs b/codex-rs/code-mode/src/remote_session/connection/driver/cell_ids.rs new file mode 100644 index 0000000000..666f6edda6 --- /dev/null +++ b/codex-rs/code-mode/src/remote_session/connection/driver/cell_ids.rs @@ -0,0 +1,111 @@ +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::RuntimeResponse; +use codex_code_mode_protocol::WaitOutcome; +use codex_code_mode_protocol::WaitRequest; +use codex_code_mode_protocol::host::WireCellId; +use codex_code_mode_protocol::host::WireRuntimeResponse; +use codex_code_mode_protocol::host::WireWaitOutcome; +use codex_code_mode_protocol::host::WireWaitRequest; + +use super::RemoteSession; + +pub(super) fn public_cell_id(generation: u64, cell_id: &WireCellId) -> CellId { + if generation == 1 { + CellId::new(cell_id.as_str().to_string()) + } else { + CellId::new(format!("g{generation}:{}", cell_id.as_str())) + } +} + +pub(super) fn public_cell_id_from_protocol(generation: u64, cell_id: &CellId) -> CellId { + public_cell_id(generation, &WireCellId::new(cell_id.as_str())) +} + +pub(super) fn remote_cell_id( + session: &RemoteSession, + cell_id: &CellId, +) -> Result { + if session.generation == 1 { + if cell_id.as_str().starts_with('g') && cell_id.as_str().contains(':') { + return Err(format!( + "cell {cell_id} belongs to a stale code-mode host generation" + )); + } + return Ok(WireCellId::new(cell_id.as_str())); + } + let prefix = format!("g{}:", session.generation); + let Some(remote_id) = cell_id.as_str().strip_prefix(&prefix) else { + return Err(format!( + "cell {cell_id} belongs to a stale code-mode host generation" + )); + }; + Ok(WireCellId::new(remote_id)) +} + +pub(super) fn remote_wait_request( + session: &RemoteSession, + request: WaitRequest, +) -> Result { + Ok(WireWaitRequest { + cell_id: remote_cell_id(session, &request.cell_id)?, + yield_time_ms: request.yield_time_ms, + }) +} + +pub(super) fn public_runtime_response( + generation: u64, + response: RuntimeResponse, +) -> RuntimeResponse { + match response { + RuntimeResponse::Yielded { + cell_id, + content_items, + } => RuntimeResponse::Yielded { + cell_id: public_cell_id_from_protocol(generation, &cell_id), + content_items, + }, + RuntimeResponse::Terminated { + cell_id, + content_items, + } => RuntimeResponse::Terminated { + cell_id: public_cell_id_from_protocol(generation, &cell_id), + content_items, + }, + RuntimeResponse::Result { + cell_id, + content_items, + error_text, + } => RuntimeResponse::Result { + cell_id: public_cell_id_from_protocol(generation, &cell_id), + content_items, + error_text, + }, + } +} + +pub(super) fn public_wait_outcome(generation: u64, outcome: WaitOutcome) -> WaitOutcome { + match outcome { + WaitOutcome::LiveCell(response) => { + WaitOutcome::LiveCell(public_runtime_response(generation, response)) + } + WaitOutcome::MissingCell(response) => { + WaitOutcome::MissingCell(public_runtime_response(generation, response)) + } + } +} + +pub(super) fn runtime_response_cell_id(response: &WireRuntimeResponse) -> &WireCellId { + match response { + WireRuntimeResponse::Yielded { cell_id, .. } + | WireRuntimeResponse::Terminated { cell_id, .. } + | WireRuntimeResponse::Result { cell_id, .. } => cell_id, + } +} + +pub(super) fn wait_outcome_cell_id(outcome: &WireWaitOutcome) -> &WireCellId { + match outcome { + WireWaitOutcome::LiveCell(response) | WireWaitOutcome::MissingCell(response) => { + runtime_response_cell_id(response) + } + } +} diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/cleanup.rs b/codex-rs/code-mode/src/remote_session/connection/driver/cleanup.rs new file mode 100644 index 0000000000..b994b64887 --- /dev/null +++ b/codex-rs/code-mode/src/remote_session/connection/driver/cleanup.rs @@ -0,0 +1,40 @@ +use std::sync::Arc; + +use tokio_util::sync::CancellationToken; + +use super::notify_cell_closed; +use super::session_registry::CellOwner; + +struct CleanupInner { + complete: CancellationToken, +} + +#[derive(Clone)] +pub(in crate::remote_session) struct SessionCleanup { + inner: Arc, +} + +impl SessionCleanup { + pub(in crate::remote_session) fn new() -> Self { + Self { + inner: Arc::new(CleanupInner { + complete: CancellationToken::new(), + }), + } + } + + pub(super) fn fail(&self, cells: Vec) { + for owner in cells { + notify_cell_closed(&owner.delegate, &owner.cell_id); + } + self.inner.complete.cancel(); + } + + pub(in crate::remote_session) async fn wait(&self) { + self.inner.complete.cancelled().await; + } + + pub(in crate::remote_session) fn is_complete(&self) -> bool { + self.inner.complete.is_cancelled() + } +} diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs b/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs new file mode 100644 index 0000000000..c9cba9c32d --- /dev/null +++ b/codex-rs/code-mode/src/remote_session/connection/driver/commands.rs @@ -0,0 +1,298 @@ +use std::sync::Arc; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::ExecuteRequest; +use codex_code_mode_protocol::WaitOutcome; +use codex_code_mode_protocol::WaitRequest; +use codex_code_mode_protocol::host::ClientToHost; +use codex_code_mode_protocol::host::EncodedFrame; +use codex_code_mode_protocol::host::HostRequest; +use codex_code_mode_protocol::host::WireWaitRequest; +use tokio::sync::oneshot; +use tokio_util::sync::CancellationToken; + +use super::ConnectionDriver; +use super::cell_ids::remote_cell_id; +use super::cell_ids::remote_wait_request; +use super::types::CancellableRequest; +use super::types::DeferredWait; +use super::types::DeliveredExecute; +use super::types::DriverCommand; +use super::types::PendingRequest; +use super::types::RemoteSession; + +impl ConnectionDriver { + pub(super) fn handle_command(&mut self, command: DriverCommand) -> bool { + match command { + DriverCommand::OpenSession { + session, + delegate, + cleanup, + caller_cancellation, + response_tx, + } => self.open_session(session, delegate, cleanup, caller_cancellation, response_tx), + DriverCommand::Execute { + session, + request, + caller_cancellation, + response_tx, + } => self.execute(session, request, caller_cancellation, response_tx), + DriverCommand::Wait { + session, + request, + caller_cancellation, + response_tx, + } => self.wait(session, request, caller_cancellation, response_tx), + DriverCommand::Terminate { + session, + cell_id, + response_tx, + } => self.terminate(session, cell_id, response_tx), + DriverCommand::ShutdownSession { + session, + response_tx, + } => self.shutdown_session(session, response_tx), + } + } + + fn open_session( + &mut self, + session: RemoteSession, + delegate: Arc, + cleanup: super::cleanup::SessionCleanup, + caller_cancellation: CancellationToken, + response_tx: oneshot::Sender>, + ) -> bool { + if self.sessions.contains(&session.id) || self.requests.contains_pending_open(&session) { + let _ = response_tx.send(Err(format!( + "code-mode session {} is already open", + session.id + ))); + return true; + } + let request_id = match self.requests.allocate_id() { + Ok(id) => id, + Err(err) => { + let _ = response_tx.send(Err(err)); + return false; + } + }; + let message = ClientToHost::Request { + id: request_id, + request: HostRequest::OpenSession { + session_id: session.id.clone(), + }, + }; + let frame = match EncodedFrame::encode(&message) { + Ok(frame) => frame, + Err(err) => { + let _ = response_tx.send(Err(format!( + "failed to encode code-mode open-session request: {err}" + ))); + return true; + } + }; + let cancellation = CancellableRequest::new(caller_cancellation); + self.requests.insert_pending( + request_id, + PendingRequest::OpenSession { + session, + delegate, + cleanup, + cancellation, + response_tx, + }, + &self.event_tx, + ); + self.queue_frame(frame) + } + + fn execute( + &mut self, + session: RemoteSession, + request: ExecuteRequest, + caller_cancellation: CancellationToken, + response_tx: oneshot::Sender>, + ) -> bool { + if let Err(err) = self.sessions.require_ready(&session) { + let _ = response_tx.send(Err(err)); + return true; + } + let request = match request.try_into() { + Ok(request) => request, + Err(err) => { + let _ = response_tx.send(Err(format!( + "failed to encode code-mode execute request: {err}" + ))); + return true; + } + }; + let request_id = match self.requests.allocate_id() { + Ok(id) => id, + Err(err) => { + let _ = response_tx.send(Err(err)); + return false; + } + }; + let message = ClientToHost::Request { + id: request_id, + request: HostRequest::Execute { + session_id: session.id.clone(), + request, + }, + }; + let frame = match EncodedFrame::encode(&message) { + Ok(frame) => frame, + Err(err) => { + let _ = response_tx.send(Err(format!( + "code-mode execute request exceeds the IPC frame limit: {err}" + ))); + return true; + } + }; + let (initial_response_tx, initial_response_rx) = oneshot::channel(); + let cancellation = CancellableRequest::new(caller_cancellation); + self.requests.insert_pending( + request_id, + PendingRequest::Execute { + session, + response_tx, + initial_response_tx, + initial_response_rx, + cancellation, + }, + &self.event_tx, + ); + self.queue_frame(frame) + } + + fn wait( + &mut self, + session: RemoteSession, + request: WaitRequest, + caller_cancellation: CancellationToken, + response_tx: oneshot::Sender>, + ) -> bool { + if let Err(err) = self.sessions.require_ready(&session) { + let _ = response_tx.send(Err(err)); + return true; + } + let request = match remote_wait_request(&session, request) { + Ok(request) => request, + Err(err) => { + let _ = response_tx.send(Err(err)); + return true; + } + }; + if self.requests.has_cancelled_wait(&session, &request.cell_id) { + self.requests.push_deferred_wait(DeferredWait { + session, + request, + caller_cancellation, + response_tx, + }); + return true; + } + self.start_wait(session, request, caller_cancellation, response_tx) + } + + pub(super) fn start_wait( + &mut self, + session: RemoteSession, + request: WireWaitRequest, + caller_cancellation: CancellationToken, + response_tx: oneshot::Sender>, + ) -> bool { + let cell_id = request.cell_id.clone(); + self.send_request( + HostRequest::Wait { + session_id: session.id.clone(), + request, + }, + PendingRequest::Wait { + session, + cell_id, + cancellation: CancellableRequest::new(caller_cancellation), + response_tx, + }, + ) + } + + fn terminate( + &mut self, + session: RemoteSession, + cell_id: CellId, + response_tx: oneshot::Sender>, + ) -> bool { + if let Err(err) = self.sessions.require_ready(&session) { + let _ = response_tx.send(Err(err)); + return true; + } + let cell_id = match remote_cell_id(&session, &cell_id) { + Ok(cell_id) => cell_id, + Err(err) => { + let _ = response_tx.send(Err(err)); + return true; + } + }; + let pending_cell_id = cell_id.clone(); + self.send_request( + HostRequest::Terminate { + session_id: session.id.clone(), + cell_id, + }, + PendingRequest::Terminate { + session, + cell_id: pending_cell_id, + response_tx, + }, + ) + } + + fn shutdown_session( + &mut self, + session: RemoteSession, + response_tx: oneshot::Sender>, + ) -> bool { + if let Err(err) = self.sessions.begin_shutdown(&session) { + let _ = response_tx.send(Err(err)); + return true; + } + self.send_request( + HostRequest::ShutdownSession { + session_id: session.id.clone(), + }, + PendingRequest::ShutdownSession { + session, + response_tx, + }, + ) + } + + pub(super) fn send_request(&mut self, request: HostRequest, pending: PendingRequest) -> bool { + let request_id = match self.requests.allocate_id() { + Ok(id) => id, + Err(err) => { + pending.fail(err); + return false; + } + }; + let message = ClientToHost::Request { + id: request_id, + request, + }; + let frame = match EncodedFrame::encode(&message) { + Ok(frame) => frame, + Err(err) => { + pending.fail(format!( + "code-mode request exceeds the IPC frame limit: {err}" + )); + return true; + } + }; + self.requests + .insert_pending(request_id, pending, &self.event_tx); + self.queue_frame(frame) + } +} diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/delegate_runtime.rs b/codex-rs/code-mode/src/remote_session/connection/driver/delegate_runtime.rs new file mode 100644 index 0000000000..c79270661d --- /dev/null +++ b/codex-rs/code-mode/src/remote_session/connection/driver/delegate_runtime.rs @@ -0,0 +1,339 @@ +//! Client-side delegate task and closure lifecycle. +//! +//! Cancellation revokes the task's completion path before removing its active-call state. The +//! delegate future may finish later, but it can no longer send a response or affect cell closure. + +use std::collections::HashMap; +use std::collections::HashSet; +use std::collections::VecDeque; + +use codex_code_mode_protocol::CodeModeNestedToolCall; +use codex_code_mode_protocol::host::ClientToHost; +use codex_code_mode_protocol::host::DelegateRequest; +use codex_code_mode_protocol::host::DelegateRequestId; +use codex_code_mode_protocol::host::DelegateResponse; +use codex_code_mode_protocol::host::EncodedFrame; +use codex_code_mode_protocol::host::SessionId; +use codex_code_mode_protocol::host::WireCellId; +use codex_code_mode_protocol::host::WireResult; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +use super::ConnectionDriver; +use super::notify_cell_closed; +use super::session_registry::CellOwner; +use super::session_registry::DelegateTarget; +use super::session_registry::FailedSession; +use super::types::DriverEvent; + +const MAX_RECENT_DELEGATE_REQUEST_IDS: usize = 4096; + +#[derive(Clone, Eq, Hash, PartialEq)] +struct CellKey { + session_id: codex_code_mode_protocol::host::SessionId, + cell_id: codex_code_mode_protocol::CellId, +} + +impl CellKey { + fn for_owner(owner: &CellOwner) -> Self { + Self { + session_id: owner.session_id.clone(), + cell_id: owner.cell_id.clone(), + } + } +} + +struct DelegateCall { + cell: CellKey, + cancellation: CancellationToken, + completion_stop: CancellationToken, +} + +impl DelegateCall { + fn revoke(&self) { + self.cancellation.cancel(); + self.completion_stop.cancel(); + } +} + +enum DelegateTask { + InvokeTool(CodeModeNestedToolCall), + Notify { + call_id: String, + cell_id: codex_code_mode_protocol::CellId, + text: String, + }, +} + +pub(super) struct DelegateEffects { + pub(super) response: Option<(DelegateRequestId, Result)>, + pub(super) closed_cells: Vec, +} + +impl DelegateEffects { + fn empty() -> Self { + Self { + response: None, + closed_cells: Vec::new(), + } + } + + fn append(&mut self, mut other: Self) { + debug_assert!(self.response.is_none()); + self.response = other.response.take(); + self.closed_cells.append(&mut other.closed_cells); + } +} + +pub(super) struct DelegateRuntime { + calls: HashMap, + seen_requests: HashSet, + request_order: VecDeque, + event_tx: mpsc::Sender, +} + +impl DelegateRuntime { + pub(super) fn new(event_tx: mpsc::Sender) -> Self { + Self { + calls: HashMap::new(), + seen_requests: HashSet::new(), + request_order: VecDeque::new(), + event_tx, + } + } + + pub(super) fn start( + &mut self, + id: DelegateRequestId, + target: DelegateTarget, + request: DelegateRequest, + ) -> Result<(), String> { + if self.calls.contains_key(&id) || self.seen_requests.contains(&id) { + return Err(format!("duplicate code-mode delegate request ID {id:?}")); + } + self.remember_request(id); + let cancellation = CancellationToken::new(); + let task_request = match request { + DelegateRequest::InvokeTool { invocation } => { + let mut invocation: CodeModeNestedToolCall = invocation.into(); + invocation.cell_id = target.cell_id.clone(); + DelegateTask::InvokeTool(invocation) + } + DelegateRequest::Notify { + call_id, + cell_id: _, + text, + } => DelegateTask::Notify { + call_id, + cell_id: target.cell_id.clone(), + text, + }, + }; + let delegate = target.delegate; + let task_cancellation = cancellation.clone(); + let delegate_task = tokio::spawn(async move { + match task_request { + DelegateTask::InvokeTool(invocation) => delegate + .invoke_tool(invocation, task_cancellation) + .await + .map(|result| DelegateResponse::ToolResult { result }), + DelegateTask::Notify { + call_id, + cell_id, + text, + } => delegate + .notify(call_id, cell_id, text, task_cancellation) + .await + .map(|()| DelegateResponse::NotificationDelivered), + } + }); + let completion_stop = CancellationToken::new(); + self.calls.insert( + id, + DelegateCall { + cell: CellKey { + session_id: target.session_id, + cell_id: target.cell_id, + }, + cancellation, + completion_stop: completion_stop.clone(), + }, + ); + let event_tx = self.event_tx.clone(); + tokio::spawn(async move { + let result = tokio::select! { + biased; + _ = completion_stop.cancelled() => return, + result = delegate_task => match result { + Ok(result) => result, + Err(err) => Err(format!("code-mode delegate task failed: {err}")), + }, + }; + tokio::select! { + biased; + _ = completion_stop.cancelled() => {} + _ = event_tx.send(DriverEvent::DelegateCompleted { id, result }) => {} + } + }); + Ok(()) + } + + pub(super) fn cancel(&mut self, id: DelegateRequestId) { + if let Some(call) = self.calls.remove(&id) { + call.revoke(); + } + } + + pub(super) fn complete( + &mut self, + id: DelegateRequestId, + result: Result, + ) -> DelegateEffects { + if self.calls.remove(&id).is_none() { + return DelegateEffects::empty(); + } + let mut effects = DelegateEffects::empty(); + effects.response = Some((id, result)); + effects + } + + pub(super) fn close_cell(&mut self, owner: CellOwner) -> DelegateEffects { + let key = CellKey::for_owner(&owner); + self.calls.retain(|_, call| { + if call.cell != key { + return true; + } + call.revoke(); + false + }); + let mut effects = DelegateEffects::empty(); + effects.closed_cells.push(owner); + effects + } + + pub(super) fn close_cells(&mut self, owners: Vec) -> DelegateEffects { + let mut effects = DelegateEffects::empty(); + for owner in owners { + effects.append(self.close_cell(owner)); + } + effects + } + + pub(super) fn fail_all(&mut self, failed_sessions: Vec) { + for (_, call) in self.calls.drain() { + call.revoke(); + } + for session in failed_sessions { + session.cleanup.fail(session.cells); + } + } + + fn remember_request(&mut self, id: DelegateRequestId) { + self.seen_requests.insert(id); + self.request_order.push_back(id); + while self.request_order.len() > MAX_RECENT_DELEGATE_REQUEST_IDS { + if let Some(expired) = self.request_order.pop_front() { + self.seen_requests.remove(&expired); + } + } + } +} + +impl ConnectionDriver { + pub(super) fn start_delegate( + &mut self, + id: DelegateRequestId, + session_id: SessionId, + request: DelegateRequest, + ) -> bool { + let wire_cell_id = match &request { + DelegateRequest::InvokeTool { invocation } => &invocation.cell_id, + DelegateRequest::Notify { cell_id, .. } => cell_id, + }; + let target = match self.sessions.delegate_target(&session_id, wire_cell_id) { + Ok(target) => target, + Err(err) => { + self.fail(err); + return false; + } + }; + if let Err(err) = self.delegates.start(id, target, request) { + self.fail(err); + return false; + } + true + } + + pub(super) fn complete_delegate( + &mut self, + id: DelegateRequestId, + result: Result, + ) -> bool { + let effects = self.delegates.complete(id, result); + self.apply_delegate_effects(effects) + } + + fn send_delegate_response( + &mut self, + id: DelegateRequestId, + result: Result, + ) -> bool { + let message = ClientToHost::DelegateResponse { + id, + result: WireResult::from_result(result), + }; + let frame = match EncodedFrame::encode(&message) { + Ok(frame) => frame, + Err(err) => { + let fallback = ClientToHost::DelegateResponse { + id, + result: WireResult::Err { + message: format!( + "code-mode delegate response exceeds the IPC frame limit: {err}" + ), + }, + }; + match EncodedFrame::encode(&fallback) { + Ok(frame) => frame, + Err(fallback_err) => { + self.fail(format!( + "failed to encode code-mode delegate error response: {fallback_err}" + )); + return false; + } + } + } + }; + self.queue_frame(frame) + } + + pub(super) fn close_cell(&mut self, session_id: SessionId, cell_id: WireCellId) -> bool { + let owner = match self.sessions.remove_cell(&session_id, &cell_id) { + Ok(owner) => owner, + Err(err) => { + self.fail(err); + return false; + } + }; + let effects = self.delegates.close_cell(owner); + self.apply_delegate_effects(effects) + } + + pub(super) fn close_session_locally(&mut self, session_id: &SessionId) -> DelegateEffects { + self.requests.remove_unclaimed_for_session(session_id); + let owners = self.sessions.remove_session(session_id); + self.delegates.close_cells(owners) + } + + pub(super) fn apply_delegate_effects(&mut self, effects: DelegateEffects) -> bool { + if let Some((id, result)) = effects.response + && !self.send_delegate_response(id, result) + { + return false; + } + for closed in effects.closed_cells { + notify_cell_closed(&closed.delegate, &closed.cell_id); + } + true + } +} diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/request_tracker.rs b/codex-rs/code-mode/src/remote_session/connection/driver/request_tracker.rs new file mode 100644 index 0000000000..0e7fc0881d --- /dev/null +++ b/codex-rs/code-mode/src/remote_session/connection/driver/request_tracker.rs @@ -0,0 +1,186 @@ +use std::collections::HashMap; +use std::collections::VecDeque; + +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::SessionId; +use codex_code_mode_protocol::host::WireCellId; +use tokio::sync::mpsc; + +use super::types::DeferredWait; +use super::types::DriverEvent; +use super::types::InitialResponse; +use super::types::PendingRequest; +use super::types::RemoteSession; +use super::types::UnclaimedExecute; + +pub(super) enum CancellationAction { + Send(RequestId), + Terminate { + request_id: RequestId, + execute: UnclaimedExecute, + }, +} + +pub(super) struct RequestTracker { + pending: HashMap, + unclaimed_executes: HashMap, + initial_responses: HashMap, + deferred_waits: VecDeque, + next_request_id: i64, +} + +impl RequestTracker { + pub(super) fn new() -> Self { + Self { + pending: HashMap::new(), + unclaimed_executes: HashMap::new(), + initial_responses: HashMap::new(), + deferred_waits: VecDeque::new(), + next_request_id: 1, + } + } + + pub(super) fn contains_pending_open(&self, session: &RemoteSession) -> bool { + self.pending.values().any(|pending| { + matches!( + pending, + PendingRequest::OpenSession { + session: pending_session, + .. + } if pending_session.id == session.id + ) + }) + } + + pub(super) fn allocate_id(&mut self) -> Result { + let id = self.next_request_id; + self.next_request_id = self + .next_request_id + .checked_add(1) + .ok_or_else(|| "code-mode host request ID space exhausted".to_string())?; + Ok(RequestId::new(id)) + } + + pub(super) fn insert_pending( + &mut self, + id: RequestId, + pending: PendingRequest, + event_tx: &mpsc::Sender, + ) { + self.pending.insert(id, pending); + if let Some(cancellation) = self + .pending + .get_mut(&id) + .and_then(PendingRequest::cancellation_mut) + { + cancellation.spawn_watcher(id, event_tx.clone()); + } + } + + pub(super) fn remove_pending(&mut self, id: RequestId) -> Option { + self.pending.remove(&id) + } + + pub(super) fn insert_initial_response(&mut self, id: RequestId, response: InitialResponse) { + self.initial_responses.insert(id, response); + } + + pub(super) fn remove_initial_response(&mut self, id: RequestId) -> Option { + self.initial_responses.remove(&id) + } + + pub(super) fn insert_unclaimed_execute(&mut self, id: RequestId, execute: UnclaimedExecute) { + self.unclaimed_executes.insert(id, execute); + } + + pub(super) fn claim_execute(&mut self, id: RequestId) { + self.unclaimed_executes.remove(&id); + } + + pub(super) fn collect_cancellations(&mut self) -> Vec { + let mut actions = self + .pending + .iter_mut() + .filter_map(|(id, pending)| { + let cancellation = pending.cancellation_mut()?; + (cancellation.is_cancelled() && cancellation.mark_reported()) + .then_some(CancellationAction::Send(*id)) + }) + .collect::>(); + actions.extend( + self.unclaimed_executes + .extract_if(|_, execute| { + execute.cancellation.is_cancelled() && execute.cancellation.mark_reported() + }) + .map(|(request_id, execute)| CancellationAction::Terminate { + request_id, + execute, + }), + ); + actions + } + + pub(super) fn mark_cancelled(&mut self, id: RequestId) -> Option { + if let Some(cancellation) = self + .pending + .get_mut(&id) + .and_then(PendingRequest::cancellation_mut) + { + return cancellation + .mark_reported() + .then_some(CancellationAction::Send(id)); + } + let execute = self.unclaimed_executes.get_mut(&id)?; + if !execute.cancellation.mark_reported() { + return None; + } + self.unclaimed_executes + .remove(&id) + .map(|execute| CancellationAction::Terminate { + request_id: id, + execute, + }) + } + + pub(super) fn has_cancelled_wait(&self, session: &RemoteSession, cell_id: &WireCellId) -> bool { + self.pending.values().any(|pending| { + matches!( + pending, + PendingRequest::Wait { + session: pending_session, + cell_id: pending_cell_id, + cancellation, + .. + } if pending_session == session + && pending_cell_id == cell_id + && cancellation.is_cancelled() + ) + }) + } + + pub(super) fn push_deferred_wait(&mut self, wait: DeferredWait) { + self.deferred_waits.push_back(wait); + } + + pub(super) fn take_deferred_waits(&mut self) -> VecDeque { + std::mem::take(&mut self.deferred_waits) + } + + pub(super) fn remove_unclaimed_for_session(&mut self, session_id: &SessionId) { + self.unclaimed_executes + .retain(|_, execute| &execute.session.id != session_id); + } + + pub(super) fn fail_all(&mut self, reason: &str) { + for (_, pending) in self.pending.drain() { + pending.fail(reason.to_string()); + } + self.unclaimed_executes.clear(); + for (_, initial) in self.initial_responses.drain() { + let _ = initial.response_tx.send(Err(reason.to_string())); + } + for wait in self.deferred_waits.drain(..) { + let _ = wait.response_tx.send(Err(reason.to_string())); + } + } +} diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs b/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs new file mode 100644 index 0000000000..19c78ec000 --- /dev/null +++ b/codex-rs/code-mode/src/remote_session/connection/driver/responses.rs @@ -0,0 +1,397 @@ +use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::host::ClientToHost; +use codex_code_mode_protocol::host::EncodedFrame; +use codex_code_mode_protocol::host::HostRequest; +use codex_code_mode_protocol::host::HostResponse; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::WireCellId; +use tokio::sync::oneshot; + +use super::ConnectionDriver; +use super::cell_ids::public_runtime_response; +use super::cell_ids::public_wait_outcome; +use super::cell_ids::runtime_response_cell_id; +use super::cell_ids::wait_outcome_cell_id; +use super::request_tracker::CancellationAction; +use super::session_registry::CellAdmissionError; +use super::types::DeliveredExecute; +use super::types::InitialResponse; +use super::types::PendingRequest; +use super::types::RemoteSession; +use super::types::UnclaimedExecute; + +impl ConnectionDriver { + pub(super) fn flush_deferred_waits(&mut self) -> bool { + let mut deferred = self.requests.take_deferred_waits(); + while let Some(wait) = deferred.pop_front() { + if wait.caller_cancellation.is_cancelled() { + let _ = wait + .response_tx + .send(Err("code-mode request cancelled".to_string())); + continue; + } + if self + .requests + .has_cancelled_wait(&wait.session, &wait.request.cell_id) + { + self.requests.push_deferred_wait(wait); + continue; + } + if !self.start_wait( + wait.session, + wait.request, + wait.caller_cancellation, + wait.response_tx, + ) { + for wait in deferred { + let _ = wait + .response_tx + .send(Err("code-mode host connection closed".to_string())); + } + return false; + } + } + true + } + + pub(super) fn handle_host_message(&mut self, message: HostToClient) -> bool { + match message { + HostToClient::Response { id, result } => { + self.complete_request(id, result.into_result()) + } + HostToClient::InitialResponse { id, result } => { + self.complete_initial_response(id, result.into_result()) + } + HostToClient::DelegateRequest { + id, + session_id, + request, + } => self.start_delegate(id, session_id, request), + HostToClient::CancelDelegateRequest { id } => { + self.delegates.cancel(id); + true + } + HostToClient::CellClosed { + session_id, + cell_id, + } => self.close_cell(session_id, cell_id), + HostToClient::HostHello(_) | HostToClient::HandshakeRejected { .. } => { + self.fail("code-mode host sent a second handshake response".to_string()); + false + } + } + } + + fn complete_request(&mut self, id: RequestId, result: Result) -> bool { + let Some(pending) = self.requests.remove_pending(id) else { + self.fail(format!("code-mode host returned unknown request ID {id:?}")); + return false; + }; + match pending { + PendingRequest::OpenSession { + session, + delegate, + cleanup, + cancellation, + response_tx, + } => match result { + Ok(HostResponse::SessionReady { session_id }) if session_id == session.id => { + let abandoned = cancellation.is_cancelled() || response_tx.is_closed(); + self.sessions + .insert_ready(session.clone(), delegate, cleanup); + if abandoned || response_tx.send(Ok(())).is_err() { + return self.shutdown_abandoned_session(session); + } + } + Ok(_) => { + let reason = + "code-mode host returned an invalid open-session response".to_string(); + let _ = response_tx.send(Err(reason.clone())); + self.fail(reason); + return false; + } + Err(err) => { + let _ = response_tx.send(Err(err)); + } + }, + PendingRequest::Execute { + session, + response_tx, + initial_response_tx, + initial_response_rx, + cancellation, + } => match result { + Ok(HostResponse::ExecutionStarted { cell_id }) => { + // The host owns a checked, never-reused ID sequence. Retain only live + // IDs so client memory scales with concurrency, not session lifetime. + let remote_cell_id = cell_id.clone(); + let public_id = match self.sessions.admit_cell(&session, cell_id) { + Ok(public_id) => public_id, + Err(CellAdmissionError::MissingSession) => { + let _ = response_tx + .send(Err("code-mode session closed during execute".to_string())); + return true; + } + Err(CellAdmissionError::DuplicateCell) => { + let reason = format!( + "code-mode host reused live cell {} in session {}", + remote_cell_id.as_str(), + session.id + ); + let _ = response_tx.send(Err(reason.clone())); + self.fail(reason); + return false; + } + }; + self.requests.insert_initial_response( + id, + InitialResponse { + generation: session.generation, + cell_id: remote_cell_id.clone(), + response_tx: initial_response_tx, + }, + ); + let started = StartedCell::from_result_receiver(public_id, initial_response_rx); + if cancellation.is_cancelled() || response_tx.is_closed() { + return self.terminate_abandoned_cell(session, remote_cell_id); + } + let delivered = DeliveredExecute { + request_id: id, + started, + }; + if response_tx.send(Ok(delivered)).is_err() { + return self.terminate_abandoned_cell(session, remote_cell_id); + } + self.requests.insert_unclaimed_execute( + id, + UnclaimedExecute { + session, + cell_id: remote_cell_id, + cancellation, + }, + ); + } + Ok(_) => { + let reason = "code-mode host returned an invalid execute response".to_string(); + let _ = response_tx.send(Err(reason.clone())); + self.fail(reason); + return false; + } + Err(err) => { + let _ = response_tx.send(Err(err)); + } + }, + PendingRequest::Wait { + session, + cell_id, + cancellation: _, + response_tx, + } => { + let result = match result { + Ok(HostResponse::WaitCompleted { outcome }) => { + if wait_outcome_cell_id(&outcome) != &cell_id { + let reason = format!( + "code-mode host returned cell {} for request targeting {}", + wait_outcome_cell_id(&outcome).as_str(), + cell_id.as_str() + ); + let _ = response_tx.send(Err(reason.clone())); + self.fail(reason); + return false; + } + Ok(public_wait_outcome(session.generation, outcome.into())) + } + Ok(_) => { + let reason = "code-mode host returned an invalid cell response".to_string(); + let _ = response_tx.send(Err(reason.clone())); + self.fail(reason); + return false; + } + Err(err) => Err(err), + }; + let _ = response_tx.send(result); + } + PendingRequest::Terminate { + session, + cell_id, + response_tx, + } => { + let result = match result { + Ok(HostResponse::WaitCompleted { outcome }) => { + if wait_outcome_cell_id(&outcome) != &cell_id { + let reason = format!( + "code-mode host returned cell {} for request targeting {}", + wait_outcome_cell_id(&outcome).as_str(), + cell_id.as_str() + ); + let _ = response_tx.send(Err(reason.clone())); + self.fail(reason); + return false; + } + public_wait_outcome(session.generation, outcome.into()) + } + Ok(_) => { + let reason = "code-mode host returned an invalid cell response".to_string(); + let _ = response_tx.send(Err(reason.clone())); + self.fail(reason); + return false; + } + Err(err) => { + let _ = response_tx.send(Err(err)); + return true; + } + }; + let _ = response_tx.send(Ok(result)); + } + PendingRequest::ShutdownSession { + session, + response_tx, + } => match result { + Ok(HostResponse::SessionClosed { session_id }) if session_id == session.id => { + let effects = self.close_session_locally(&session.id); + if !self.apply_delegate_effects(effects) { + return false; + } + let _ = response_tx.send(Ok(())); + } + Ok(_) => { + let err = "code-mode host returned an invalid shutdown response".to_string(); + let _ = response_tx.send(Err(err.clone())); + self.fail(err); + return false; + } + Err(err) => { + let _ = response_tx.send(Err(err.clone())); + self.fail(err); + return false; + } + }, + } + true + } + + pub(super) fn cancel_dropped_callers(&mut self) -> bool { + for action in self.requests.collect_cancellations() { + if !self.apply_cancellation(action) { + return false; + } + } + true + } + + pub(super) fn cancel_request(&mut self, id: RequestId) -> bool { + self.requests + .mark_cancelled(id) + .is_none_or(|action| self.apply_cancellation(action)) + } + + fn apply_cancellation(&mut self, action: CancellationAction) -> bool { + match action { + CancellationAction::Send(id) => self.send_cancel_request(id), + CancellationAction::Terminate { + request_id, + execute, + } => { + if !self.send_cancel_request(request_id) { + return false; + } + self.terminate_abandoned_cell(execute.session, execute.cell_id) + } + } + } + + fn send_cancel_request(&mut self, id: RequestId) -> bool { + let frame = match EncodedFrame::encode(&ClientToHost::CancelRequest { id }) { + Ok(frame) => frame, + Err(err) => { + self.fail(format!( + "failed to encode code-mode cancellation request: {err}" + )); + return false; + } + }; + self.queue_frame(frame) + } + + fn shutdown_abandoned_session(&mut self, session: RemoteSession) -> bool { + let Some(should_shutdown) = self.sessions.begin_abandoned_shutdown(&session.id) else { + self.fail(format!( + "code-mode host committed abandoned session {} without local state", + session.id + )); + return false; + }; + if !should_shutdown { + return true; + } + let (response_tx, response_rx) = oneshot::channel(); + drop(response_rx); + self.send_request( + HostRequest::ShutdownSession { + session_id: session.id.clone(), + }, + PendingRequest::ShutdownSession { + session, + response_tx, + }, + ) + } + + fn terminate_abandoned_cell(&mut self, session: RemoteSession, cell_id: WireCellId) -> bool { + let Some(is_closing) = self.sessions.is_closing(&session.id) else { + self.fail(format!( + "code-mode host admitted an abandoned cell in unknown session {}", + session.id + )); + return false; + }; + if is_closing { + return true; + } + let (response_tx, response_rx) = oneshot::channel(); + drop(response_rx); + self.send_request( + HostRequest::Terminate { + session_id: session.id.clone(), + cell_id: cell_id.clone(), + }, + PendingRequest::Terminate { + session, + cell_id, + response_tx, + }, + ) + } + + fn complete_initial_response( + &mut self, + id: RequestId, + result: Result, + ) -> bool { + let Some(initial) = self.requests.remove_initial_response(id) else { + self.fail(format!( + "code-mode host returned initial response for unknown request ID {id:?}" + )); + return false; + }; + let response = match result { + Ok(response) if runtime_response_cell_id(&response) == &initial.cell_id => { + Ok(public_runtime_response(initial.generation, response.into())) + } + Ok(response) => { + let reason = format!( + "code-mode host returned initial response for cell {} instead of {}", + runtime_response_cell_id(&response).as_str(), + initial.cell_id.as_str() + ); + let _ = initial.response_tx.send(Err(reason.clone())); + self.fail(reason); + return false; + } + Err(err) => Err(err), + }; + let _ = initial.response_tx.send(response); + true + } +} diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/session_registry.rs b/codex-rs/code-mode/src/remote_session/connection/driver/session_registry.rs new file mode 100644 index 0000000000..a3d1c82ae6 --- /dev/null +++ b/codex-rs/code-mode/src/remote_session/connection/driver/session_registry.rs @@ -0,0 +1,222 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::host::SessionId; +use codex_code_mode_protocol::host::WireCellId; + +use super::cell_ids::public_cell_id; +use super::cleanup::SessionCleanup; +use super::types::RemoteSession; + +pub(super) struct CellOwner { + pub(super) session_id: SessionId, + pub(super) cell_id: CellId, + pub(super) delegate: Arc, +} + +pub(super) struct DelegateTarget { + pub(super) session_id: SessionId, + pub(super) cell_id: CellId, + pub(super) delegate: Arc, +} + +pub(super) struct FailedSession { + pub(super) cleanup: SessionCleanup, + pub(super) cells: Vec, +} + +pub(super) enum CellAdmissionError { + MissingSession, + DuplicateCell, +} + +struct SessionRecord { + remote: RemoteSession, + delegate: Arc, + cleanup: SessionCleanup, + phase: SessionPhase, + cells: HashMap, +} + +#[derive(Clone, Copy, Eq, PartialEq)] +enum SessionPhase { + Ready, + Closing, +} + +pub(super) struct SessionRegistry { + records: HashMap, +} + +impl SessionRegistry { + pub(super) fn new() -> Self { + Self { + records: HashMap::new(), + } + } + + pub(super) fn contains(&self, session_id: &SessionId) -> bool { + self.records.contains_key(session_id) + } + + pub(super) fn insert_ready( + &mut self, + session: RemoteSession, + delegate: Arc, + cleanup: SessionCleanup, + ) { + self.records.insert( + session.id.clone(), + SessionRecord { + remote: session, + delegate, + cleanup, + phase: SessionPhase::Ready, + cells: HashMap::new(), + }, + ); + } + + pub(super) fn require_ready(&self, session: &RemoteSession) -> Result<(), String> { + let record = self + .records + .get(&session.id) + .ok_or_else(|| format!("unknown code-mode session {}", session.id))?; + if record.remote != *session { + return Err("stale code-mode session generation".to_string()); + } + if record.phase != SessionPhase::Ready { + return Err("code-mode session is shutting down".to_string()); + } + Ok(()) + } + + pub(super) fn begin_shutdown(&mut self, session: &RemoteSession) -> Result<(), String> { + let record = self + .records + .get_mut(&session.id) + .ok_or_else(|| format!("unknown code-mode session {}", session.id))?; + if record.remote != *session { + return Err("stale code-mode session generation".to_string()); + } + if record.phase == SessionPhase::Closing { + return Err("code-mode session is already closing".to_string()); + } + record.phase = SessionPhase::Closing; + Ok(()) + } + + pub(super) fn begin_abandoned_shutdown(&mut self, session_id: &SessionId) -> Option { + let record = self.records.get_mut(session_id)?; + if record.phase == SessionPhase::Closing { + return Some(false); + } + record.phase = SessionPhase::Closing; + Some(true) + } + + pub(super) fn is_closing(&self, session_id: &SessionId) -> Option { + self.records + .get(session_id) + .map(|record| record.phase == SessionPhase::Closing) + } + + pub(super) fn admit_cell( + &mut self, + session: &RemoteSession, + cell_id: WireCellId, + ) -> Result { + let Some(record) = self.records.get_mut(&session.id) else { + return Err(CellAdmissionError::MissingSession); + }; + if record.cells.contains_key(&cell_id) { + return Err(CellAdmissionError::DuplicateCell); + } + let public_id = public_cell_id(session.generation, &cell_id); + record.cells.insert(cell_id, public_id.clone()); + Ok(public_id) + } + + pub(super) fn delegate_target( + &self, + session_id: &SessionId, + cell_id: &WireCellId, + ) -> Result { + let session = self + .records + .get(session_id) + .ok_or_else(|| format!("code-mode host delegated for unknown session {session_id}"))?; + let public_id = session.cells.get(cell_id).cloned().ok_or_else(|| { + format!( + "code-mode host delegated for unknown cell {} in session {session_id}", + cell_id.as_str() + ) + })?; + Ok(DelegateTarget { + session_id: session_id.clone(), + cell_id: public_id, + delegate: Arc::clone(&session.delegate), + }) + } + + pub(super) fn remove_cell( + &mut self, + session_id: &SessionId, + cell_id: &WireCellId, + ) -> Result { + let session = self.records.get_mut(session_id).ok_or_else(|| { + format!( + "code-mode host closed cell {} in unknown session {session_id}", + cell_id.as_str() + ) + })?; + let public_id = session + .cells + .remove(cell_id) + .ok_or_else(|| format!("code-mode host closed unknown cell in session {session_id}"))?; + Ok(CellOwner { + session_id: session_id.clone(), + cell_id: public_id, + delegate: Arc::clone(&session.delegate), + }) + } + + pub(super) fn remove_session(&mut self, session_id: &SessionId) -> Vec { + let Some(session) = self.records.remove(session_id) else { + return Vec::new(); + }; + session + .cells + .into_values() + .map(|cell_id| CellOwner { + session_id: session_id.clone(), + cell_id, + delegate: Arc::clone(&session.delegate), + }) + .collect() + } + + pub(super) fn drain(&mut self) -> Vec { + let sessions = std::mem::take(&mut self.records); + sessions + .into_iter() + .map(|(session_id, session)| { + let cells = session + .cells + .into_values() + .map(|cell_id| CellOwner { + session_id: session_id.clone(), + cell_id, + delegate: Arc::clone(&session.delegate), + }) + .collect(); + FailedSession { + cleanup: session.cleanup, + cells, + } + }) + .collect() + } +} diff --git a/codex-rs/code-mode/src/remote_session/connection/driver/types.rs b/codex-rs/code-mode/src/remote_session/connection/driver/types.rs new file mode 100644 index 0000000000..d4f6a102f3 --- /dev/null +++ b/codex-rs/code-mode/src/remote_session/connection/driver/types.rs @@ -0,0 +1,196 @@ +use std::sync::Arc; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::ExecuteRequest; +use codex_code_mode_protocol::RuntimeResponse; +use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::WaitOutcome; +use codex_code_mode_protocol::WaitRequest; +use codex_code_mode_protocol::host::DelegateRequestId; +use codex_code_mode_protocol::host::DelegateResponse; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::SessionId; +use codex_code_mode_protocol::host::WireCellId; +use codex_code_mode_protocol::host::WireWaitRequest; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio_util::sync::CancellationToken; + +use super::cleanup::SessionCleanup; + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(in crate::remote_session) struct RemoteSession { + pub(in crate::remote_session) id: SessionId, + pub(in crate::remote_session) generation: u64, +} + +pub(in crate::remote_session::connection) enum DriverCommand { + OpenSession { + session: RemoteSession, + delegate: Arc, + cleanup: SessionCleanup, + caller_cancellation: CancellationToken, + response_tx: oneshot::Sender>, + }, + Execute { + session: RemoteSession, + request: ExecuteRequest, + caller_cancellation: CancellationToken, + response_tx: oneshot::Sender>, + }, + Wait { + session: RemoteSession, + request: WaitRequest, + caller_cancellation: CancellationToken, + response_tx: oneshot::Sender>, + }, + Terminate { + session: RemoteSession, + cell_id: CellId, + response_tx: oneshot::Sender>, + }, + ShutdownSession { + session: RemoteSession, + response_tx: oneshot::Sender>, + }, +} + +pub(in crate::remote_session::connection) enum DriverEvent { + HostMessage(HostToClient), + DelegateCompleted { + id: DelegateRequestId, + result: Result, + }, + RequestCancelled(RequestId), + Failed(String), +} + +pub(super) struct CancellableRequest { + caller_cancellation: CancellationToken, + watcher_stop: CancellationToken, + reported: bool, +} + +impl CancellableRequest { + pub(super) fn new(caller_cancellation: CancellationToken) -> Self { + Self { + caller_cancellation, + watcher_stop: CancellationToken::new(), + reported: false, + } + } + + pub(super) fn is_cancelled(&self) -> bool { + self.caller_cancellation.is_cancelled() + } + + pub(super) fn mark_reported(&mut self) -> bool { + if self.reported { + return false; + } + self.reported = true; + true + } + + pub(super) fn spawn_watcher(&self, id: RequestId, event_tx: mpsc::Sender) { + let caller_cancellation = self.caller_cancellation.clone(); + let watcher_stop = self.watcher_stop.clone(); + tokio::spawn(async move { + tokio::select! { + _ = caller_cancellation.cancelled() => { + let _ = event_tx.send(DriverEvent::RequestCancelled(id)).await; + } + _ = watcher_stop.cancelled() => {} + } + }); + } +} + +impl Drop for CancellableRequest { + fn drop(&mut self) { + self.watcher_stop.cancel(); + } +} + +pub(super) struct InitialResponse { + pub(super) generation: u64, + pub(super) cell_id: WireCellId, + pub(super) response_tx: oneshot::Sender>, +} + +pub(in crate::remote_session::connection) struct DeliveredExecute { + pub(in crate::remote_session::connection) request_id: RequestId, + pub(in crate::remote_session::connection) started: StartedCell, +} + +pub(super) struct UnclaimedExecute { + pub(super) session: RemoteSession, + pub(super) cell_id: WireCellId, + pub(super) cancellation: CancellableRequest, +} + +pub(super) enum PendingRequest { + OpenSession { + session: RemoteSession, + delegate: Arc, + cleanup: SessionCleanup, + cancellation: CancellableRequest, + response_tx: oneshot::Sender>, + }, + Execute { + session: RemoteSession, + response_tx: oneshot::Sender>, + initial_response_tx: oneshot::Sender>, + initial_response_rx: oneshot::Receiver>, + cancellation: CancellableRequest, + }, + Wait { + session: RemoteSession, + cell_id: WireCellId, + cancellation: CancellableRequest, + response_tx: oneshot::Sender>, + }, + Terminate { + session: RemoteSession, + cell_id: WireCellId, + response_tx: oneshot::Sender>, + }, + ShutdownSession { + session: RemoteSession, + response_tx: oneshot::Sender>, + }, +} + +pub(super) struct DeferredWait { + pub(super) session: RemoteSession, + pub(super) request: WireWaitRequest, + pub(super) caller_cancellation: CancellationToken, + pub(super) response_tx: oneshot::Sender>, +} + +impl PendingRequest { + pub(super) fn cancellation_mut(&mut self) -> Option<&mut CancellableRequest> { + match self { + Self::OpenSession { cancellation, .. } + | Self::Execute { cancellation, .. } + | Self::Wait { cancellation, .. } => Some(cancellation), + Self::Terminate { .. } | Self::ShutdownSession { .. } => None, + } + } + + pub(super) fn fail(self, reason: String) { + match self { + Self::OpenSession { response_tx, .. } | Self::ShutdownSession { response_tx, .. } => { + let _ = response_tx.send(Err(reason)); + } + Self::Execute { response_tx, .. } => { + let _ = response_tx.send(Err(reason)); + } + Self::Wait { response_tx, .. } | Self::Terminate { response_tx, .. } => { + let _ = response_tx.send(Err(reason)); + } + } + } +} diff --git a/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs b/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs new file mode 100644 index 0000000000..97ed9b553e --- /dev/null +++ b/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs @@ -0,0 +1,1514 @@ +use std::sync::Arc; +use std::sync::Mutex as StdMutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeNestedToolCall; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::ExecuteRequest; +use codex_code_mode_protocol::NotificationFuture; +use codex_code_mode_protocol::ToolInvocationFuture; +use codex_code_mode_protocol::WaitRequest; +use codex_code_mode_protocol::host::DelegateRequest; +use codex_code_mode_protocol::host::DelegateRequestId; +use codex_code_mode_protocol::host::HostResponse; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::SessionId; +use codex_code_mode_protocol::host::WireNestedToolCall; +use codex_code_mode_protocol::host::WireResult; +use codex_code_mode_protocol::host::WireRuntimeResponse; +use codex_code_mode_protocol::host::WireWaitOutcome; +use codex_protocol::ToolName; +use pretty_assertions::assert_eq; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio_util::sync::CancellationToken; + +use super::ConnectionDriver; +use super::DriverCommand; +use super::DriverEvent; +use super::DriverLifecycle; +use super::RemoteSession; +use super::SessionCleanup; + +struct DriverHarness { + command_tx: mpsc::Sender, + event_tx: mpsc::Sender, + execute_claim_tx: mpsc::UnboundedSender, + outgoing_rx: mpsc::Receiver, + cancellation: CancellationToken, + alive: Arc, + driver_task: tokio::task::JoinHandle<()>, +} + +impl DriverHarness { + fn start() -> Self { + let (command_tx, command_rx) = mpsc::channel(/*max_capacity*/ 16); + let (event_tx, event_rx) = mpsc::channel(/*max_capacity*/ 16); + let (outgoing_tx, outgoing_rx) = mpsc::channel(/*max_capacity*/ 16); + let cancellation = CancellationToken::new(); + let alive = Arc::new(AtomicBool::new(true)); + let (driver, execute_claim_tx) = ConnectionDriver::new( + command_rx, + event_rx, + event_tx.clone(), + outgoing_tx, + DriverLifecycle { + alive: Arc::clone(&alive), + failure: Arc::new(StdMutex::new(None)), + cancellation: cancellation.clone(), + }, + ); + let driver_task = tokio::spawn(driver.run()); + Self { + command_tx, + event_tx, + execute_claim_tx, + outgoing_rx, + cancellation, + alive, + driver_task, + } + } + + async fn open( + &mut self, + session: RemoteSession, + delegate: Arc, + ) -> SessionCleanup { + let cleanup = SessionCleanup::new(); + let (response_tx, response_rx) = oneshot::channel(); + self.command_tx + .send(DriverCommand::OpenSession { + session: session.clone(), + delegate, + cleanup: cleanup.clone(), + caller_cancellation: CancellationToken::new(), + response_tx, + }) + .await + .expect("open command"); + self.outgoing_rx.recv().await.expect("open frame"); + self.event_tx + .send(DriverEvent::HostMessage(HostToClient::Response { + id: RequestId::new(/*value*/ 1), + result: WireResult::Ok { + value: HostResponse::SessionReady { + session_id: session.id, + }, + }, + })) + .await + .expect("open response"); + response_rx + .await + .expect("open reply") + .expect("open session"); + cleanup + } + + async fn start_cell( + &mut self, + session: RemoteSession, + request_id: i64, + cell_id: &str, + ) -> codex_code_mode_protocol::StartedCell { + let (response_tx, response_rx) = oneshot::channel(); + self.command_tx + .send(DriverCommand::Execute { + session, + request: ExecuteRequest { + tool_call_id: format!("call-{request_id}"), + enabled_tools: Vec::new(), + source: "await new Promise(() => {})".to_string(), + yield_time_ms: Some(1), + max_output_tokens: None, + }, + caller_cancellation: CancellationToken::new(), + response_tx, + }) + .await + .expect("execute command"); + self.outgoing_rx.recv().await.expect("execute frame"); + self.event_tx + .send(DriverEvent::HostMessage(HostToClient::Response { + id: RequestId::new(request_id), + result: WireResult::Ok { + value: HostResponse::ExecutionStarted { + cell_id: CellId::new(cell_id.to_string()).into(), + }, + }, + })) + .await + .expect("execute response"); + let delivered = response_rx + .await + .expect("execute reply") + .expect("execute session"); + self.execute_claim_tx + .send(delivered.request_id) + .expect("claim execute"); + delivered.started + } + + async fn start_tool_delegate(&self, session: &RemoteSession, id: DelegateRequestId) { + self.event_tx + .send(DriverEvent::HostMessage(HostToClient::DelegateRequest { + id, + session_id: session.id.clone(), + request: DelegateRequest::InvokeTool { + invocation: WireNestedToolCall { + cell_id: CellId::new("1".to_string()).into(), + runtime_tool_call_id: "tool-1".to_string(), + tool_name: ToolName::plain("slow").into(), + tool_kind: codex_code_mode_protocol::CodeModeToolKind::Function.into(), + input: None, + }, + }, + })) + .await + .expect("delegate request"); + } +} + +impl Drop for DriverHarness { + fn drop(&mut self) { + self.cancellation.cancel(); + } +} + +#[derive(Default)] +struct RecordingDelegate { + closed_cells: StdMutex>, + invocations: AtomicUsize, + notifications: AtomicUsize, +} + +struct PanickingDelegate; + +#[derive(Debug, Eq, PartialEq)] +enum HeldDelegateEvent { + Started, + Cancelled, + Finished, + CellClosed(CellId), +} + +struct HeldDelegate { + events_tx: mpsc::UnboundedSender, + release: CancellationToken, +} + +impl HeldDelegate { + fn new() -> ( + Arc, + mpsc::UnboundedReceiver, + CancellationToken, + ) { + let (events_tx, events_rx) = mpsc::unbounded_channel(); + let release = CancellationToken::new(); + ( + Arc::new(Self { + events_tx, + release: release.clone(), + }), + events_rx, + release, + ) + } +} + +impl CodeModeSessionDelegate for HeldDelegate { + fn invoke_tool<'a>( + &'a self, + _invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + let events_tx = self.events_tx.clone(); + let release = self.release.clone(); + Box::pin(async move { + let _ = events_tx.send(HeldDelegateEvent::Started); + cancellation_token.cancelled().await; + let _ = events_tx.send(HeldDelegateEvent::Cancelled); + release.cancelled().await; + let _ = events_tx.send(HeldDelegateEvent::Finished); + Err("cancelled".to_string()) + }) + } + + fn notify<'a>( + &'a self, + _call_id: String, + _cell_id: CellId, + _text: String, + _cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async { Ok(()) }) + } + + fn cell_closed(&self, cell_id: &CellId) { + let _ = self + .events_tx + .send(HeldDelegateEvent::CellClosed(cell_id.clone())); + } +} + +impl CodeModeSessionDelegate for PanickingDelegate { + fn invoke_tool<'a>( + &'a self, + _invocation: CodeModeNestedToolCall, + _cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + Box::pin(async { panic!("delegate panic probe") }) + } + + fn notify<'a>( + &'a self, + _call_id: String, + _cell_id: CellId, + _text: String, + _cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async { Ok(()) }) + } + + fn cell_closed(&self, _cell_id: &CellId) {} +} + +impl CodeModeSessionDelegate for RecordingDelegate { + fn invoke_tool<'a>( + &'a self, + _invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + self.invocations.fetch_add(1, Ordering::Relaxed); + Box::pin(async move { + cancellation_token.cancelled().await; + Err("cancelled".to_string()) + }) + } + + fn notify<'a>( + &'a self, + _call_id: String, + _cell_id: CellId, + _text: String, + _cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + self.notifications.fetch_add(1, Ordering::Relaxed); + Box::pin(async { Ok(()) }) + } + + fn cell_closed(&self, cell_id: &CellId) { + self.closed_cells + .lock() + .expect("closed cells lock") + .push(cell_id.clone()); + } +} + +fn remote_session() -> RemoteSession { + RemoteSession { + id: SessionId::new("session-1").expect("session ID"), + generation: 1, + } +} + +async fn next_held_delegate_event( + events_rx: &mut mpsc::UnboundedReceiver, +) -> HeldDelegateEvent { + tokio::time::timeout(Duration::from_secs(1), events_rx.recv()) + .await + .expect("delegate event timeout") + .expect("delegate event stream") +} + +#[tokio::test] +async fn dropped_open_waiter_shuts_down_committed_session() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + let (open_tx, open_rx) = oneshot::channel(); + let cleanup = SessionCleanup::new(); + harness + .command_tx + .send(DriverCommand::OpenSession { + session: session.clone(), + delegate: Arc::new(RecordingDelegate::default()), + cleanup, + caller_cancellation: CancellationToken::new(), + response_tx: open_tx, + }) + .await + .expect("open command"); + drop(open_rx); + harness.outgoing_rx.recv().await.expect("open frame"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::Response { + id: RequestId::new(/*value*/ 1), + result: WireResult::Ok { + value: HostResponse::SessionReady { + session_id: session.id.clone(), + }, + }, + })) + .await + .expect("open response"); + harness + .outgoing_rx + .recv() + .await + .expect("abandoned session shutdown frame"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::Response { + id: RequestId::new(/*value*/ 2), + result: WireResult::Ok { + value: HostResponse::SessionClosed { + session_id: session.id.clone(), + }, + }, + })) + .await + .expect("shutdown response"); + + let (execute_tx, execute_rx) = oneshot::channel(); + harness + .command_tx + .send(DriverCommand::Execute { + session: session.clone(), + request: ExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: "text('ok')".to_string(), + yield_time_ms: None, + max_output_tokens: None, + }, + caller_cancellation: CancellationToken::new(), + response_tx: execute_tx, + }) + .await + .expect("execute command"); + assert_eq!( + execute_rx + .await + .expect("execute reply") + .err() + .expect("closed session should reject execute"), + "unknown code-mode session session-1" + ); +} + +#[tokio::test] +async fn delegate_cancel_is_best_effort_and_sends_no_late_response() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + let delegate = Arc::new(RecordingDelegate::default()); + harness.open(session.clone(), delegate.clone()).await; + let _started = harness + .start_cell(session.clone(), /*request_id*/ 2, "1") + .await; + let request_id = DelegateRequestId::new(/*value*/ 7); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::DelegateRequest { + id: request_id, + session_id: session.id.clone(), + request: DelegateRequest::InvokeTool { + invocation: WireNestedToolCall { + cell_id: CellId::new("1".to_string()).into(), + runtime_tool_call_id: "tool-1".to_string(), + tool_name: ToolName::plain("slow").into(), + tool_kind: codex_code_mode_protocol::CodeModeToolKind::Function.into(), + input: None, + }, + }, + })) + .await + .expect("delegate request"); + harness + .event_tx + .send(DriverEvent::HostMessage( + HostToClient::CancelDelegateRequest { id: request_id }, + )) + .await + .expect("delegate cancel"); + tokio::task::yield_now().await; + assert!(matches!( + harness.outgoing_rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::DelegateRequest { + id: request_id, + session_id: session.id, + request: DelegateRequest::Notify { + call_id: "notify-reused".to_string(), + cell_id: CellId::new("1".to_string()).into(), + text: "duplicate".to_string(), + }, + })) + .await + .expect("reused delegate request"); + tokio::task::yield_now().await; + + assert!(!harness.alive.load(Ordering::Acquire)); + assert_eq!(delegate.invocations.load(Ordering::Relaxed), 1); + assert_eq!(delegate.notifications.load(Ordering::Relaxed), 0); +} + +#[tokio::test] +async fn terminate_closes_cell_without_waiting_for_delegate_cleanup() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + let (delegate, mut events_rx, release) = HeldDelegate::new(); + harness.open(session.clone(), delegate).await; + let _started = harness + .start_cell(session.clone(), /*request_id*/ 2, "1") + .await; + let delegate_id = DelegateRequestId::new(/*value*/ 7); + harness.start_tool_delegate(&session, delegate_id).await; + assert_eq!( + next_held_delegate_event(&mut events_rx).await, + HeldDelegateEvent::Started + ); + + let (response_tx, response_rx) = oneshot::channel(); + harness + .command_tx + .send(DriverCommand::Terminate { + session: session.clone(), + cell_id: CellId::new("1".to_string()), + response_tx, + }) + .await + .expect("terminate command"); + harness.outgoing_rx.recv().await.expect("terminate frame"); + harness + .event_tx + .send(DriverEvent::HostMessage( + HostToClient::CancelDelegateRequest { id: delegate_id }, + )) + .await + .expect("delegate cancel"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::CellClosed { + session_id: session.id, + cell_id: CellId::new("1".to_string()).into(), + })) + .await + .expect("cell close"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::Response { + id: RequestId::new(/*value*/ 3), + result: WireResult::Ok { + value: HostResponse::WaitCompleted { + outcome: WireWaitOutcome::LiveCell(WireRuntimeResponse::Terminated { + cell_id: CellId::new("1".to_string()).into(), + content_items: Vec::new(), + }), + }, + }, + })) + .await + .expect("terminate response"); + + let closure_events = [ + next_held_delegate_event(&mut events_rx).await, + next_held_delegate_event(&mut events_rx).await, + ]; + assert!(closure_events.contains(&HeldDelegateEvent::Cancelled)); + assert!(closure_events.contains(&HeldDelegateEvent::CellClosed(CellId::new("1".to_string())))); + assert_eq!( + response_rx.await.expect("terminate reply"), + Ok(codex_code_mode_protocol::WaitOutcome::LiveCell( + codex_code_mode_protocol::RuntimeResponse::Terminated { + cell_id: CellId::new("1".to_string()), + content_items: Vec::new(), + } + )) + ); + assert!(matches!( + events_rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty | mpsc::error::TryRecvError::Disconnected) + )); + + release.cancel(); + assert_eq!( + next_held_delegate_event(&mut events_rx).await, + HeldDelegateEvent::Finished + ); + assert!(matches!( + events_rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty | mpsc::error::TryRecvError::Disconnected) + )); + assert!(matches!( + harness.outgoing_rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + assert!(harness.alive.load(Ordering::Acquire)); +} + +#[tokio::test] +async fn shutdown_closes_cell_without_waiting_for_delegate_cleanup() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + let (delegate, mut events_rx, release) = HeldDelegate::new(); + harness.open(session.clone(), delegate).await; + let _started = harness + .start_cell(session.clone(), /*request_id*/ 2, "1") + .await; + let delegate_id = DelegateRequestId::new(/*value*/ 7); + harness.start_tool_delegate(&session, delegate_id).await; + assert_eq!( + next_held_delegate_event(&mut events_rx).await, + HeldDelegateEvent::Started + ); + + let (response_tx, response_rx) = oneshot::channel(); + harness + .command_tx + .send(DriverCommand::ShutdownSession { + session: session.clone(), + response_tx, + }) + .await + .expect("shutdown command"); + harness.outgoing_rx.recv().await.expect("shutdown frame"); + harness + .event_tx + .send(DriverEvent::HostMessage( + HostToClient::CancelDelegateRequest { id: delegate_id }, + )) + .await + .expect("delegate cancel"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::CellClosed { + session_id: session.id.clone(), + cell_id: CellId::new("1".to_string()).into(), + })) + .await + .expect("cell close"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::Response { + id: RequestId::new(/*value*/ 3), + result: WireResult::Ok { + value: HostResponse::SessionClosed { + session_id: session.id, + }, + }, + })) + .await + .expect("shutdown response"); + + let closure_events = [ + next_held_delegate_event(&mut events_rx).await, + next_held_delegate_event(&mut events_rx).await, + ]; + assert!(closure_events.contains(&HeldDelegateEvent::Cancelled)); + assert!(closure_events.contains(&HeldDelegateEvent::CellClosed(CellId::new("1".to_string())))); + assert_eq!(response_rx.await.expect("shutdown reply"), Ok(())); + assert!(matches!( + events_rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + + release.cancel(); + assert_eq!( + next_held_delegate_event(&mut events_rx).await, + HeldDelegateEvent::Finished + ); + assert!(matches!( + events_rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty | mpsc::error::TryRecvError::Disconnected) + )); + assert!(matches!( + harness.outgoing_rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + assert!(harness.alive.load(Ordering::Acquire)); +} + +#[tokio::test] +async fn completed_delegate_request_id_cannot_be_reused() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + let delegate = Arc::new(RecordingDelegate::default()); + harness.open(session.clone(), delegate.clone()).await; + let _started = harness + .start_cell(session.clone(), /*request_id*/ 2, "1") + .await; + let request_id = DelegateRequestId::new(/*value*/ 7); + let request = || DelegateRequest::Notify { + call_id: "notify-1".to_string(), + cell_id: CellId::new("1".to_string()).into(), + text: "once".to_string(), + }; + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::DelegateRequest { + id: request_id, + session_id: session.id.clone(), + request: request(), + })) + .await + .expect("delegate request"); + harness + .outgoing_rx + .recv() + .await + .expect("delegate response frame"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::DelegateRequest { + id: request_id, + session_id: session.id, + request: request(), + })) + .await + .expect("reused delegate request"); + tokio::task::yield_now().await; + + assert!(!harness.alive.load(Ordering::Acquire)); + assert_eq!(delegate.notifications.load(Ordering::Relaxed), 1); +} + +#[tokio::test] +async fn delegate_task_panic_becomes_tool_error_without_killing_connection() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + harness + .open(session.clone(), Arc::new(PanickingDelegate)) + .await; + let _started = harness + .start_cell(session.clone(), /*request_id*/ 2, "1") + .await; + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::DelegateRequest { + id: DelegateRequestId::new(/*value*/ 7), + session_id: session.id.clone(), + request: DelegateRequest::InvokeTool { + invocation: WireNestedToolCall { + cell_id: CellId::new("1".to_string()).into(), + runtime_tool_call_id: "tool-1".to_string(), + tool_name: ToolName::plain("panic").into(), + tool_kind: codex_code_mode_protocol::CodeModeToolKind::Function.into(), + input: None, + }, + }, + })) + .await + .expect("delegate request"); + tokio::time::timeout(Duration::from_secs(1), harness.outgoing_rx.recv()) + .await + .expect("delegate response timeout") + .expect("delegate response frame"); + + assert!(harness.alive.load(Ordering::Acquire)); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::CellClosed { + session_id: session.id, + cell_id: CellId::new("1".to_string()).into(), + })) + .await + .expect("cell close"); +} + +#[tokio::test] +async fn delegate_for_unknown_cell_fails_connection_without_invocation() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + let delegate = Arc::new(RecordingDelegate::default()); + harness.open(session.clone(), delegate.clone()).await; + + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::DelegateRequest { + id: DelegateRequestId::new(/*value*/ 7), + session_id: session.id, + request: DelegateRequest::InvokeTool { + invocation: WireNestedToolCall { + cell_id: CellId::new("missing".to_string()).into(), + runtime_tool_call_id: "tool-1".to_string(), + tool_name: ToolName::plain("slow").into(), + tool_kind: codex_code_mode_protocol::CodeModeToolKind::Function.into(), + input: None, + }, + }, + })) + .await + .expect("delegate request"); + tokio::task::yield_now().await; + + assert!(!harness.alive.load(Ordering::Acquire)); + assert_eq!(delegate.invocations.load(Ordering::Relaxed), 0); +} + +#[tokio::test] +async fn delegate_after_cell_close_fails_connection_without_invocation() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + let delegate = Arc::new(RecordingDelegate::default()); + harness.open(session.clone(), delegate.clone()).await; + let _started = harness + .start_cell(session.clone(), /*request_id*/ 2, "1") + .await; + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::CellClosed { + session_id: session.id.clone(), + cell_id: CellId::new("1".to_string()).into(), + })) + .await + .expect("cell close"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::DelegateRequest { + id: DelegateRequestId::new(/*value*/ 7), + session_id: session.id, + request: DelegateRequest::Notify { + call_id: "notify-1".to_string(), + cell_id: CellId::new("1".to_string()).into(), + text: "late".to_string(), + }, + })) + .await + .expect("delegate request"); + tokio::task::yield_now().await; + + assert!(!harness.alive.load(Ordering::Acquire)); + assert_eq!(delegate.invocations.load(Ordering::Relaxed), 0); +} + +#[tokio::test] +async fn mismatched_initial_response_fails_connection_and_closes_cell_once() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + let delegate = Arc::new(RecordingDelegate::default()); + harness.open(session.clone(), delegate.clone()).await; + let started = harness.start_cell(session, /*request_id*/ 2, "1").await; + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::InitialResponse { + id: RequestId::new(/*value*/ 2), + result: WireResult::Ok { + value: WireRuntimeResponse::Yielded { + cell_id: CellId::new("2".to_string()).into(), + content_items: Vec::new(), + }, + }, + })) + .await + .expect("initial response"); + + assert!(started.initial_response().await.is_err()); + assert!(!harness.alive.load(Ordering::Acquire)); + assert_eq!( + *delegate.closed_cells.lock().expect("closed cells lock"), + vec![CellId::new("1".to_string())] + ); +} + +#[tokio::test] +async fn mismatched_wait_response_fails_connection() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + let delegate = Arc::new(RecordingDelegate::default()); + harness.open(session.clone(), delegate.clone()).await; + let _started = harness + .start_cell(session.clone(), /*request_id*/ 2, "1") + .await; + let (response_tx, response_rx) = oneshot::channel(); + harness + .command_tx + .send(DriverCommand::Wait { + session, + request: WaitRequest { + cell_id: CellId::new("1".to_string()), + yield_time_ms: 1, + }, + caller_cancellation: CancellationToken::new(), + response_tx, + }) + .await + .expect("wait command"); + harness.outgoing_rx.recv().await.expect("wait frame"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::Response { + id: RequestId::new(/*value*/ 3), + result: WireResult::Ok { + value: HostResponse::WaitCompleted { + outcome: WireWaitOutcome::LiveCell(WireRuntimeResponse::Yielded { + cell_id: CellId::new("2".to_string()).into(), + content_items: Vec::new(), + }), + }, + }, + })) + .await + .expect("wait response"); + + assert!(response_rx.await.expect("wait reply").is_err()); + assert!(!harness.alive.load(Ordering::Acquire)); + assert_eq!( + *delegate.closed_cells.lock().expect("closed cells lock"), + vec![CellId::new("1".to_string())] + ); +} + +#[tokio::test] +async fn mismatched_terminate_response_fails_connection() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + let delegate = Arc::new(RecordingDelegate::default()); + harness.open(session.clone(), delegate.clone()).await; + let _started = harness + .start_cell(session.clone(), /*request_id*/ 2, "1") + .await; + let (response_tx, response_rx) = oneshot::channel(); + harness + .command_tx + .send(DriverCommand::Terminate { + session, + cell_id: CellId::new("1".to_string()), + response_tx, + }) + .await + .expect("terminate command"); + harness.outgoing_rx.recv().await.expect("terminate frame"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::Response { + id: RequestId::new(/*value*/ 3), + result: WireResult::Ok { + value: HostResponse::WaitCompleted { + outcome: WireWaitOutcome::MissingCell(WireRuntimeResponse::Terminated { + cell_id: CellId::new("2".to_string()).into(), + content_items: Vec::new(), + }), + }, + }, + })) + .await + .expect("terminate response"); + + assert!(response_rx.await.expect("terminate reply").is_err()); + assert!(!harness.alive.load(Ordering::Acquire)); + assert_eq!( + *delegate.closed_cells.lock().expect("closed cells lock"), + vec![CellId::new("1".to_string())] + ); +} + +#[tokio::test] +async fn remote_wait_accepts_durations_longer_than_five_minutes() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + harness + .open(session.clone(), Arc::new(RecordingDelegate::default())) + .await; + let _started = harness + .start_cell(session.clone(), /*request_id*/ 2, "1") + .await; + let (response_tx, response_rx) = oneshot::channel(); + harness + .command_tx + .send(DriverCommand::Wait { + session, + request: WaitRequest { + cell_id: CellId::new("1".to_string()), + yield_time_ms: 300_001, + }, + caller_cancellation: CancellationToken::new(), + response_tx, + }) + .await + .expect("wait command"); + tokio::time::timeout(Duration::from_secs(1), harness.outgoing_rx.recv()) + .await + .expect("wait frame timeout") + .expect("wait frame"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::Response { + id: RequestId::new(/*value*/ 3), + result: WireResult::Ok { + value: HostResponse::WaitCompleted { + outcome: WireWaitOutcome::LiveCell(WireRuntimeResponse::Yielded { + cell_id: CellId::new("1".to_string()).into(), + content_items: Vec::new(), + }), + }, + }, + })) + .await + .expect("wait response"); + + assert_eq!( + response_rx.await.expect("wait reply"), + Ok(codex_code_mode_protocol::WaitOutcome::LiveCell( + codex_code_mode_protocol::RuntimeResponse::Yielded { + cell_id: CellId::new("1".to_string()), + content_items: Vec::new(), + } + )) + ); +} + +#[tokio::test] +async fn cancelled_wait_is_retired_before_next_wait_is_sent() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + harness + .open(session.clone(), Arc::new(RecordingDelegate::default())) + .await; + let _started = harness + .start_cell(session.clone(), /*request_id*/ 2, "1") + .await; + let first_cancellation = CancellationToken::new(); + let (first_tx, first_rx) = oneshot::channel(); + harness + .command_tx + .send(DriverCommand::Wait { + session: session.clone(), + request: WaitRequest { + cell_id: CellId::new("1".to_string()), + yield_time_ms: 60_000, + }, + caller_cancellation: first_cancellation.clone(), + response_tx: first_tx, + }) + .await + .expect("first wait command"); + harness.outgoing_rx.recv().await.expect("first wait frame"); + first_cancellation.cancel(); + drop(first_rx); + + let (second_tx, second_rx) = oneshot::channel(); + harness + .command_tx + .send(DriverCommand::Wait { + session, + request: WaitRequest { + cell_id: CellId::new("1".to_string()), + yield_time_ms: 1, + }, + caller_cancellation: CancellationToken::new(), + response_tx: second_tx, + }) + .await + .expect("second wait command"); + harness + .outgoing_rx + .recv() + .await + .expect("cancel request frame"); + assert!(matches!( + harness.outgoing_rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::Response { + id: RequestId::new(/*value*/ 3), + result: WireResult::Err { + message: "code-mode request cancelled".to_string(), + }, + })) + .await + .expect("cancelled wait response"); + harness.outgoing_rx.recv().await.expect("second wait frame"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::Response { + id: RequestId::new(/*value*/ 4), + result: WireResult::Ok { + value: HostResponse::WaitCompleted { + outcome: WireWaitOutcome::LiveCell(WireRuntimeResponse::Yielded { + cell_id: CellId::new("1".to_string()).into(), + content_items: Vec::new(), + }), + }, + }, + })) + .await + .expect("second wait response"); + + assert_eq!( + second_rx.await.expect("second wait reply"), + Ok(codex_code_mode_protocol::WaitOutcome::LiveCell( + codex_code_mode_protocol::RuntimeResponse::Yielded { + cell_id: CellId::new("1".to_string()), + content_items: Vec::new(), + } + )) + ); +} + +#[tokio::test] +async fn abandoned_execute_is_tracked_and_terminated_after_admission() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + let delegate = Arc::new(RecordingDelegate::default()); + harness.open(session.clone(), delegate.clone()).await; + let cancellation = CancellationToken::new(); + let (execute_tx, execute_rx) = oneshot::channel(); + harness + .command_tx + .send(DriverCommand::Execute { + session: session.clone(), + request: ExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: "await new Promise(() => {})".to_string(), + yield_time_ms: Some(1), + max_output_tokens: None, + }, + caller_cancellation: cancellation.clone(), + response_tx: execute_tx, + }) + .await + .expect("execute command"); + harness.outgoing_rx.recv().await.expect("execute frame"); + cancellation.cancel(); + drop(execute_rx); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::Response { + id: RequestId::new(/*value*/ 2), + result: WireResult::Ok { + value: HostResponse::ExecutionStarted { + cell_id: CellId::new("1".to_string()).into(), + }, + }, + })) + .await + .expect("execute response"); + + harness + .outgoing_rx + .recv() + .await + .expect("execute cancellation frame"); + harness + .outgoing_rx + .recv() + .await + .expect("abandoned cell termination frame"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::InitialResponse { + id: RequestId::new(/*value*/ 2), + result: WireResult::Ok { + value: WireRuntimeResponse::Terminated { + cell_id: CellId::new("1".to_string()).into(), + content_items: Vec::new(), + }, + }, + })) + .await + .expect("initial response"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::Response { + id: RequestId::new(/*value*/ 3), + result: WireResult::Ok { + value: HostResponse::WaitCompleted { + outcome: WireWaitOutcome::LiveCell(WireRuntimeResponse::Terminated { + cell_id: CellId::new("1".to_string()).into(), + content_items: Vec::new(), + }), + }, + }, + })) + .await + .expect("terminate response"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::CellClosed { + session_id: session.id, + cell_id: CellId::new("1".to_string()).into(), + })) + .await + .expect("cell close"); + tokio::task::yield_now().await; + + assert!(harness.alive.load(Ordering::Acquire)); + assert_eq!( + *delegate.closed_cells.lock().expect("closed cells lock"), + vec![CellId::new("1".to_string())] + ); +} + +#[tokio::test] +async fn delivered_but_unclaimed_execute_is_terminated_when_the_caller_is_cancelled() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + let delegate = Arc::new(RecordingDelegate::default()); + harness.open(session.clone(), delegate.clone()).await; + let cancellation = CancellationToken::new(); + let (execute_tx, execute_rx) = oneshot::channel(); + harness + .command_tx + .send(DriverCommand::Execute { + session: session.clone(), + request: ExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: "await new Promise(() => {})".to_string(), + yield_time_ms: Some(1), + max_output_tokens: None, + }, + caller_cancellation: cancellation.clone(), + response_tx: execute_tx, + }) + .await + .expect("execute command"); + harness.outgoing_rx.recv().await.expect("execute frame"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::Response { + id: RequestId::new(/*value*/ 2), + result: WireResult::Ok { + value: HostResponse::ExecutionStarted { + cell_id: CellId::new("1".to_string()).into(), + }, + }, + })) + .await + .expect("execute response"); + let delivered = execute_rx + .await + .expect("execute reply") + .expect("delivered execute"); + assert_eq!(delivered.request_id, RequestId::new(/*value*/ 2)); + cancellation.cancel(); + + harness + .outgoing_rx + .recv() + .await + .expect("execute cancellation frame"); + harness + .outgoing_rx + .recv() + .await + .expect("unclaimed cell termination frame"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::InitialResponse { + id: RequestId::new(/*value*/ 2), + result: WireResult::Ok { + value: WireRuntimeResponse::Terminated { + cell_id: CellId::new("1".to_string()).into(), + content_items: Vec::new(), + }, + }, + })) + .await + .expect("initial response"); + assert!(delivered.started.initial_response().await.is_ok()); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::Response { + id: RequestId::new(/*value*/ 3), + result: WireResult::Ok { + value: HostResponse::WaitCompleted { + outcome: WireWaitOutcome::LiveCell(WireRuntimeResponse::Terminated { + cell_id: CellId::new("1".to_string()).into(), + content_items: Vec::new(), + }), + }, + }, + })) + .await + .expect("terminate response"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::CellClosed { + session_id: session.id, + cell_id: CellId::new("1".to_string()).into(), + })) + .await + .expect("cell close"); + tokio::task::yield_now().await; + + assert!(harness.alive.load(Ordering::Acquire)); + assert_eq!( + *delegate.closed_cells.lock().expect("closed cells lock"), + vec![CellId::new("1".to_string())] + ); +} + +#[tokio::test] +async fn session_accepts_more_than_4096_cells_without_growing_a_tombstone_set() { + const CELL_COUNT: usize = 4097; + + let mut harness = DriverHarness::start(); + let session = remote_session(); + let delegate = Arc::new(RecordingDelegate::default()); + harness.open(session.clone(), delegate.clone()).await; + + for sequence in 1..=CELL_COUNT { + let request_id = i64::try_from(sequence).expect("cell sequence fits in i64") + 1; + let cell_id = sequence.to_string(); + let started = harness + .start_cell(session.clone(), request_id, &cell_id) + .await; + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::InitialResponse { + id: RequestId::new(request_id), + result: WireResult::Ok { + value: WireRuntimeResponse::Yielded { + cell_id: CellId::new(cell_id.clone()).into(), + content_items: Vec::new(), + }, + }, + })) + .await + .expect("initial response"); + assert!(started.initial_response().await.is_ok()); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::CellClosed { + session_id: session.id.clone(), + cell_id: CellId::new(cell_id).into(), + })) + .await + .expect("cell close"); + } + + tokio::time::timeout(Duration::from_secs(1), async { + while delegate + .closed_cells + .lock() + .expect("closed cells lock") + .len() + != CELL_COUNT + { + tokio::task::yield_now().await; + } + }) + .await + .expect("cell close callbacks timeout"); + assert!(harness.alive.load(Ordering::Acquire)); +} + +#[tokio::test] +async fn connection_failure_closes_every_live_cell_once() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + let delegate = Arc::new(RecordingDelegate::default()); + let cleanup = harness.open(session.clone(), delegate.clone()).await; + let (execute_tx, execute_rx) = oneshot::channel(); + harness + .command_tx + .send(DriverCommand::Execute { + session, + request: ExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: "await new Promise(() => {})".to_string(), + yield_time_ms: Some(1), + max_output_tokens: None, + }, + caller_cancellation: CancellationToken::new(), + response_tx: execute_tx, + }) + .await + .expect("execute command"); + harness.outgoing_rx.recv().await.expect("execute frame"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::Response { + id: RequestId::new(/*value*/ 2), + result: WireResult::Ok { + value: HostResponse::ExecutionStarted { + cell_id: CellId::new("1".to_string()).into(), + }, + }, + })) + .await + .expect("execute response"); + let _started = execute_rx + .await + .expect("execute reply") + .expect("execute session"); + harness + .event_tx + .send(DriverEvent::Failed("host crashed".to_string())) + .await + .expect("failure event"); + tokio::time::timeout(Duration::from_secs(1), cleanup.wait()) + .await + .expect("session cleanup timeout"); + assert_eq!( + *delegate.closed_cells.lock().expect("closed cells lock"), + vec![CellId::new("1".to_string())] + ); +} + +#[tokio::test] +async fn session_cleanup_does_not_wait_for_delegate_completion() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + let (delegate, mut events_rx, release) = HeldDelegate::new(); + let cleanup = harness.open(session.clone(), delegate).await; + let _started = harness + .start_cell(session.clone(), /*request_id*/ 2, "1") + .await; + harness + .start_tool_delegate(&session, DelegateRequestId::new(/*value*/ 7)) + .await; + assert_eq!( + next_held_delegate_event(&mut events_rx).await, + HeldDelegateEvent::Started + ); + + harness + .event_tx + .send(DriverEvent::Failed("host crashed".to_string())) + .await + .expect("failure event"); + let closure_events = [ + next_held_delegate_event(&mut events_rx).await, + next_held_delegate_event(&mut events_rx).await, + ]; + assert!(closure_events.contains(&HeldDelegateEvent::Cancelled)); + assert!(closure_events.contains(&HeldDelegateEvent::CellClosed(CellId::new("1".to_string())))); + tokio::time::timeout(Duration::from_secs(1), cleanup.wait()) + .await + .expect("session cleanup timeout"); + assert!(matches!( + events_rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + + release.cancel(); + assert_eq!( + next_held_delegate_event(&mut events_rx).await, + HeldDelegateEvent::Finished + ); + assert!(matches!( + events_rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty | mpsc::error::TryRecvError::Disconnected) + )); +} + +#[tokio::test] +async fn aborting_driver_marks_connection_dead_and_closes_cells() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + let delegate = Arc::new(RecordingDelegate::default()); + harness.open(session.clone(), delegate.clone()).await; + let _started = harness + .start_cell(session.clone(), /*request_id*/ 2, "1") + .await; + let (wait_tx, wait_rx) = oneshot::channel(); + harness + .command_tx + .send(DriverCommand::Wait { + session, + request: WaitRequest { + cell_id: CellId::new("1".to_string()), + yield_time_ms: 60_000, + }, + caller_cancellation: CancellationToken::new(), + response_tx: wait_tx, + }) + .await + .expect("wait command"); + harness.outgoing_rx.recv().await.expect("wait frame"); + + harness.driver_task.abort(); + for _ in 0..10 { + if !harness.alive.load(Ordering::Acquire) { + break; + } + tokio::task::yield_now().await; + } + + assert!(!harness.alive.load(Ordering::Acquire)); + assert!(harness.cancellation.is_cancelled()); + assert!(wait_rx.await.expect("wait failure").is_err()); + assert_eq!( + *delegate.closed_cells.lock().expect("closed cells lock"), + vec![CellId::new("1".to_string())] + ); +} + +#[tokio::test] +async fn dropped_shutdown_waiter_does_not_abort_remote_cleanup() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + harness + .open(session.clone(), Arc::new(RecordingDelegate::default())) + .await; + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + harness + .command_tx + .send(DriverCommand::ShutdownSession { + session: session.clone(), + response_tx: shutdown_tx, + }) + .await + .expect("shutdown command"); + drop(shutdown_rx); + harness.outgoing_rx.recv().await.expect("shutdown frame"); + harness + .event_tx + .send(DriverEvent::HostMessage(HostToClient::Response { + id: RequestId::new(/*value*/ 2), + result: WireResult::Ok { + value: HostResponse::SessionClosed { + session_id: session.id.clone(), + }, + }, + })) + .await + .expect("shutdown response"); + + let (execute_tx, execute_rx) = oneshot::channel(); + harness + .command_tx + .send(DriverCommand::Execute { + session, + request: ExecuteRequest { + tool_call_id: "call-2".to_string(), + enabled_tools: Vec::new(), + source: "text('unreachable')".to_string(), + yield_time_ms: None, + max_output_tokens: None, + }, + caller_cancellation: CancellationToken::new(), + response_tx: execute_tx, + }) + .await + .expect("execute command"); + assert_eq!( + execute_rx + .await + .expect("execute reply") + .err() + .expect("closed session should reject execute"), + "unknown code-mode session session-1" + ); + assert!(matches!( + harness.outgoing_rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); +} diff --git a/codex-rs/code-mode/src/remote_session/connection/reader.rs b/codex-rs/code-mode/src/remote_session/connection/reader.rs new file mode 100644 index 0000000000..fbfc94de42 --- /dev/null +++ b/codex-rs/code-mode/src/remote_session/connection/reader.rs @@ -0,0 +1,29 @@ +use codex_code_mode_protocol::host::FramedReader; +use codex_code_mode_protocol::host::HostToClient; +use tokio::process::ChildStdout; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +use super::driver::DriverEvent; + +pub(super) async fn drive_reader( + mut reader: FramedReader, + events: mpsc::Sender, + cancellation: CancellationToken, +) -> Result<(), String> { + loop { + let message = tokio::select! { + _ = cancellation.cancelled() => return Ok(()), + result = reader.read::() => result, + }; + let message = match message { + Ok(Some(message)) => message, + Ok(None) => return Err("code-mode host closed its stdout".to_string()), + Err(err) => return Err(format!("failed to read code-mode host message: {err}")), + }; + events + .send(DriverEvent::HostMessage(message)) + .await + .map_err(|_| "code-mode connection driver closed".to_string())?; + } +} diff --git a/codex-rs/code-mode/src/remote_session_tests.rs b/codex-rs/code-mode/src/remote_session_tests.rs new file mode 100644 index 0000000000..d971207553 --- /dev/null +++ b/codex-rs/code-mode/src/remote_session_tests.rs @@ -0,0 +1,52 @@ +use std::sync::Arc; + +use codex_code_mode_protocol::CodeModeSessionProvider; + +use super::ProcessOwnedCodeModeSession; +use super::ProcessOwnedCodeModeSessionProvider; +use crate::NoopCodeModeSessionDelegate; + +#[test] +fn provider_reuses_its_live_process_host() { + let provider = ProcessOwnedCodeModeSessionProvider::default(); + + let first = provider.process_host(); + let second = provider.process_host(); + + assert!(Arc::ptr_eq(&first, &second)); +} + +#[tokio::test] +async fn provider_reports_host_spawn_failure() { + let provider = ProcessOwnedCodeModeSessionProvider::with_host_program( + "codex-code-mode-host-does-not-exist".into(), + ); + + let error = provider + .create_session(Arc::new(NoopCodeModeSessionDelegate)) + .await + .err() + .expect("session creation should fail"); + + assert!(error.contains("failed to spawn code-mode host")); +} + +#[tokio::test] +async fn shutdown_before_open_does_not_spawn_the_host() { + let session = ProcessOwnedCodeModeSession::new(); + + session.shutdown().await.expect("shutdown session"); + let error = session + .execute(codex_code_mode_protocol::ExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: "text('unreachable')".to_string(), + yield_time_ms: None, + max_output_tokens: None, + }) + .await + .err() + .expect("shutdown session should reject execution"); + + assert_eq!(error, "code mode session is shutting down"); +} From 7d8906b4785d9e2d6ebe8c75a163f97de4f7ca35 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Fri, 26 Jun 2026 00:23:33 -0700 Subject: [PATCH 18/38] [codex] wire process-owned code mode host into core (#30142) ## Summary - add the `code_mode_host` feature flag and select `ProcessOwnedCodeModeSessionProvider` in `CodeModeService` when enabled - initialize code-mode sessions lazily so a missing host reports a tool error without failing thread startup - resolve `codex-code-mode-host` beside the running Codex binary by default while preserving `CODEX_CODE_MODE_HOST_PATH` as an override - add unit and end-to-end coverage for host resolution and graceful missing-host behavior ## Why This wires the process-owned session client from #30112 into the core service behind an opt-in rollout gate. Packaged Codex installations can place the helper in the same `bin` directory as the main executable without relying on `PATH`, while development and custom installations can continue to override the helper path. ## Stack - Depends on #30112 - Base branch: `cconger/process-owned-session-runtime-4-client` ## Validation Build `codex` and `codex-code-mode-host` `CODEX_CODE_MODE_HOST_PATH="$PWD/target/debug/codex-code-mode-host" ./target/debug/codex --enable code_mode_host` --- codex-rs/code-mode/src/remote_session.rs | 21 ++-- .../code-mode/src/remote_session_tests.rs | 51 ++++++++++ codex-rs/core/config.schema.json | 6 ++ codex-rs/core/src/codex_delegate.rs | 1 + codex-rs/core/src/session/mod.rs | 3 + codex-rs/core/src/session/session.rs | 5 +- codex-rs/core/src/session/tests.rs | 12 ++- .../core/src/session/tests/guardian_tests.rs | 1 + codex-rs/core/src/thread_manager.rs | 11 +++ codex-rs/core/src/thread_manager_tests.rs | 58 +++++++++++ codex-rs/core/src/tools/code_mode/mod.rs | 99 +++++++++++++++---- codex-rs/core/tests/suite/code_mode.rs | 26 +++++ codex-rs/features/src/lib.rs | 8 ++ 13 files changed, 275 insertions(+), 27 deletions(-) diff --git a/codex-rs/code-mode/src/remote_session.rs b/codex-rs/code-mode/src/remote_session.rs index d942b377c1..af7d27845b 100644 --- a/codex-rs/code-mode/src/remote_session.rs +++ b/codex-rs/code-mode/src/remote_session.rs @@ -1,3 +1,5 @@ +use std::ffi::OsString; +use std::io; use std::path::PathBuf; use std::sync::Arc; use std::sync::Mutex as StdMutex; @@ -472,7 +474,17 @@ impl CodeModeSession for ProcessOwnedCodeModeSession { } fn default_host_program() -> PathBuf { - if let Some(path) = std::env::var_os(CODE_MODE_HOST_PATH_ENV) { + resolve_host_program( + std::env::var_os(CODE_MODE_HOST_PATH_ENV), + std::env::current_exe(), + ) +} + +fn resolve_host_program( + override_path: Option, + current_exe: io::Result, +) -> PathBuf { + if let Some(path) = override_path { return PathBuf::from(path); } let executable_name = if cfg!(windows) { @@ -480,13 +492,10 @@ fn default_host_program() -> PathBuf { } else { "codex-code-mode-host" }; - if let Ok(current_exe) = std::env::current_exe() + if let Ok(current_exe) = current_exe && let Some(parent) = current_exe.parent() { - let sibling = parent.join(executable_name); - if sibling.is_file() { - return sibling; - } + return parent.join(executable_name); } PathBuf::from(executable_name) } diff --git a/codex-rs/code-mode/src/remote_session_tests.rs b/codex-rs/code-mode/src/remote_session_tests.rs index d971207553..03d75e5239 100644 --- a/codex-rs/code-mode/src/remote_session_tests.rs +++ b/codex-rs/code-mode/src/remote_session_tests.rs @@ -1,9 +1,12 @@ +use std::io; +use std::path::PathBuf; use std::sync::Arc; use codex_code_mode_protocol::CodeModeSessionProvider; use super::ProcessOwnedCodeModeSession; use super::ProcessOwnedCodeModeSessionProvider; +use super::resolve_host_program; use crate::NoopCodeModeSessionDelegate; #[test] @@ -16,6 +19,54 @@ fn provider_reuses_its_live_process_host() { assert!(Arc::ptr_eq(&first, &second)); } +#[test] +fn host_program_override_takes_precedence() { + assert_eq!( + resolve_host_program( + Some("custom-code-mode-host".into()), + Ok(PathBuf::from("/opt/codex/bin/codex")), + ), + PathBuf::from("custom-code-mode-host") + ); +} + +#[test] +fn host_program_is_next_to_the_main_executable_even_when_missing() { + let executable_name = if cfg!(windows) { + "codex-code-mode-host.exe" + } else { + "codex-code-mode-host" + }; + + assert_eq!( + resolve_host_program( + /*override_path*/ None, + Ok(PathBuf::from("/opt/codex/bin/codex")), + ), + PathBuf::from("/opt/codex/bin").join(executable_name) + ); +} + +#[test] +fn host_program_falls_back_to_its_name_when_main_executable_is_unknown() { + let executable_name = if cfg!(windows) { + "codex-code-mode-host.exe" + } else { + "codex-code-mode-host" + }; + + assert_eq!( + resolve_host_program( + /*override_path*/ None, + Err(io::Error::new( + io::ErrorKind::NotFound, + "missing executable" + )), + ), + PathBuf::from(executable_name) + ); +} + #[tokio::test] async fn provider_reports_host_spawn_failure() { let provider = ProcessOwnedCodeModeSessionProvider::with_host_program( diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 512c050a20..ab787ba34f 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -448,6 +448,9 @@ "code_mode": { "$ref": "#/definitions/FeatureToml_for_CodeModeConfigToml" }, + "code_mode_host": { + "type": "boolean" + }, "code_mode_only": { "type": "boolean" }, @@ -4811,6 +4814,9 @@ "code_mode": { "$ref": "#/definitions/FeatureToml_for_CodeModeConfigToml" }, + "code_mode_host": { + "type": "boolean" + }, "code_mode_only": { "type": "boolean" }, diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 8874aeae79..d25af57939 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -105,6 +105,7 @@ pub(crate) async fn run_codex_thread_interactive( skills_service: Arc::clone(&parent_session.services.skills_service), plugins_manager: Arc::clone(&parent_session.services.plugins_manager), mcp_manager: Arc::clone(&parent_session.services.mcp_manager), + code_mode_session_provider: parent_session.services.code_mode_service.session_provider(), extensions: Arc::clone(&parent_session.services.extensions), conversation_history, session_source: SessionSource::SubAgent(subagent_source.clone()), diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 0008257542..6a21a6be63 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -416,6 +416,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) skills_service: Arc, pub(crate) plugins_manager: Arc, pub(crate) mcp_manager: Arc, + pub(crate) code_mode_session_provider: Arc, pub(crate) extensions: Arc>, pub(crate) conversation_history: InitialHistory, pub(crate) session_source: SessionSource, @@ -506,6 +507,7 @@ impl Codex { skills_service, plugins_manager, mcp_manager, + code_mode_session_provider, extensions, conversation_history, session_source, @@ -680,6 +682,7 @@ impl Codex { skills_service, plugins_manager, mcp_manager.clone(), + code_mode_session_provider, extensions, thread_extension_init, supports_openai_form_elicitation, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index ed07352b15..a9cc89db3a 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -489,6 +489,7 @@ impl Session { skills_service: Arc, plugins_manager: Arc, mcp_manager: Arc, + code_mode_session_provider: Arc, extensions: Arc>, mut thread_extension_init: ExtensionDataInit, supports_openai_form_elicitation: bool, @@ -1112,7 +1113,9 @@ impl Session { session_configuration.parent_thread_id, ), ), - code_mode_service: crate::tools::code_mode::CodeModeService::new(), + code_mode_service: crate::tools::code_mode::CodeModeService::new(Arc::clone( + &code_mode_session_provider, + )), tool_search_handler_cache: Default::default(), turn_environments: Arc::clone(&turn_environments), }; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 2bc63b404f..09fdcf2245 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5225,6 +5225,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { skills_service, plugins_manager, mcp_manager, + Arc::new(codex_code_mode::InProcessCodeModeSessionProvider), Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), codex_extension_api::ExtensionDataInit::default(), /*supports_openai_form_elicitation*/ false, @@ -5438,7 +5439,9 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { /*item_ids_enabled*/ config.features.enabled(Feature::ItemIds), /*attestation_provider*/ None, ), - code_mode_service: crate::tools::code_mode::CodeModeService::new(), + code_mode_service: crate::tools::code_mode::CodeModeService::new(Arc::new( + codex_code_mode::InProcessCodeModeSessionProvider, + )), tool_search_handler_cache: Default::default(), turn_environments: Arc::clone(&turn_environments), }; @@ -5602,6 +5605,7 @@ async fn make_session_with_config_and_rx( skills_service, plugins_manager, mcp_manager, + Arc::new(codex_code_mode::InProcessCodeModeSessionProvider), Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), codex_extension_api::ExtensionDataInit::default(), /*supports_openai_form_elicitation*/ false, @@ -5708,6 +5712,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( skills_service, plugins_manager, mcp_manager, + Arc::new(codex_code_mode::InProcessCodeModeSessionProvider), Arc::new(codex_extension_api::ExtensionRegistryBuilder::new().build()), codex_extension_api::ExtensionDataInit::default(), /*supports_openai_form_elicitation*/ false, @@ -7513,7 +7518,9 @@ where /*item_ids_enabled*/ config.features.enabled(Feature::ItemIds), /*attestation_provider*/ None, ), - code_mode_service: crate::tools::code_mode::CodeModeService::new(), + code_mode_service: crate::tools::code_mode::CodeModeService::new(Arc::new( + codex_code_mode::InProcessCodeModeSessionProvider, + )), tool_search_handler_cache: Default::default(), turn_environments: Arc::clone(&turn_environments), }; @@ -9178,6 +9185,7 @@ async fn attach_in_memory_thread_store( originator: "test_originator".to_string(), base_instructions: BaseInstructions::default(), dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), multi_agent_version: None, initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 4f88c9ed66..dc17af5a16 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -727,6 +727,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { skills_service, plugins_manager, mcp_manager, + code_mode_session_provider: Arc::new(codex_code_mode::InProcessCodeModeSessionProvider), extensions: codex_extension_api::empty_extension_registry(), conversation_history: InitialHistory::New, session_source: SessionSource::SubAgent(SubAgentSource::Other( diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index b606bb3f75..2ca21c15ab 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -21,6 +21,9 @@ use codex_agent_graph_store::LocalAgentGraphStore; use codex_analytics::AnalyticsEventsClient; use codex_app_server_protocol::ThreadHistoryBuilder; use codex_app_server_protocol::TurnStatus; +use codex_code_mode::CodeModeSessionProvider; +use codex_code_mode::InProcessCodeModeSessionProvider; +use codex_code_mode::ProcessOwnedCodeModeSessionProvider; use codex_core_plugins::PluginsManager; use codex_exec_server::EnvironmentManager; use codex_extension_api::ExtensionDataInit; @@ -240,6 +243,7 @@ pub(crate) struct ThreadManagerState { skills_service: Arc, plugins_manager: Arc, mcp_manager: Arc, + code_mode_session_provider: Arc, extensions: Arc>, user_instructions_provider: Arc, thread_store: Arc, @@ -336,6 +340,11 @@ impl ThreadManager { skills_service, plugins_manager, mcp_manager, + code_mode_session_provider: if config.features.enabled(Feature::CodeModeHost) { + Arc::new(ProcessOwnedCodeModeSessionProvider::default()) + } else { + Arc::new(InProcessCodeModeSessionProvider) + }, extensions, user_instructions_provider, thread_store, @@ -441,6 +450,7 @@ impl ThreadManager { skills_service, plugins_manager, mcp_manager, + code_mode_session_provider: Arc::new(InProcessCodeModeSessionProvider), extensions: empty_extension_registry(), user_instructions_provider: Arc::new( crate::test_support::EmptyUserInstructionsProvider, @@ -1562,6 +1572,7 @@ impl ThreadManagerState { skills_service: Arc::clone(&self.skills_service), plugins_manager: Arc::clone(&self.plugins_manager), mcp_manager: Arc::clone(&self.mcp_manager), + code_mode_session_provider: Arc::clone(&self.code_mode_session_provider), extensions: Arc::clone(&self.extensions), conversation_history: initial_history, session_source, diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 5148f16259..b33d95ac33 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -395,6 +395,64 @@ async fn shutdown_all_threads_bounded_submits_shutdown_to_every_thread() { assert!(manager.list_thread_ids().await.is_empty()); } +#[tokio::test] +async fn code_mode_session_provider_is_shared_across_threads() { + let temp_dir = tempdir().expect("tempdir"); + let mut config = test_config().await; + config.codex_home = temp_dir.path().join("codex-home").abs(); + config.cwd = config.codex_home.abs(); + std::fs::create_dir_all(&config.codex_home).expect("create codex home"); + + let manager = ThreadManager::with_models_provider_and_home_for_tests( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + config.codex_home.to_path_buf(), + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + ); + let first = manager + .start_thread(config.clone()) + .await + .expect("start first thread"); + let second = manager + .start_thread(config) + .await + .expect("start second thread"); + + let first_provider = first + .thread + .codex + .session + .services + .code_mode_service + .session_provider(); + let second_provider = second + .thread + .codex + .session + .services + .code_mode_service + .session_provider(); + assert!(Arc::ptr_eq(&first_provider, &second_provider)); + assert!(Arc::ptr_eq( + &first_provider, + &manager.state.code_mode_session_provider + )); + + let mut completed = vec![first.thread_id, second.thread_id]; + completed.sort_by_key(std::string::ToString::to_string); + let report = manager + .shutdown_all_threads_bounded(Duration::from_secs(10)) + .await; + assert_eq!( + report, + ThreadShutdownReport { + completed, + submit_failed: Vec::new(), + timed_out: Vec::new(), + } + ); +} + #[tokio::test] async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { let temp_dir = tempdir().expect("tempdir"); diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index c28eef0727..ad746a1de2 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -6,16 +6,19 @@ mod wait_handler; pub(crate) mod wait_spec; use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use std::time::Duration; use codex_code_mode::CellId; use codex_code_mode::CodeModeNestedToolCall; use codex_code_mode::CodeModeSession; +use codex_code_mode::CodeModeSessionProvider; use codex_code_mode::CodeModeToolKind; -use codex_code_mode::InProcessCodeModeSession; use codex_code_mode::RuntimeResponse; use codex_protocol::models::FunctionCallOutputContentItem; use serde_json::Value as JsonValue; +use tokio::sync::OnceCell; use tokio_util::sync::CancellationToken; use crate::function_tool::FunctionCallError; @@ -61,46 +64,62 @@ pub(crate) struct ExecContext { } pub(crate) struct CodeModeService { - session: Option>, + session: OnceCell>, + session_provider: Arc, dispatch_broker: Arc, + shutting_down: AtomicBool, } impl CodeModeService { - pub(crate) fn new() -> Self { + pub(crate) fn new(session_provider: Arc) -> Self { let dispatch_broker = Arc::new(CodeModeDispatchBroker::new()); Self { - session: Some(Arc::new(InProcessCodeModeSession::with_delegate( - dispatch_broker.clone(), - ))), + session: OnceCell::new(), + session_provider, dispatch_broker, + shutting_down: AtomicBool::new(false), } } + pub(crate) fn session_provider(&self) -> Arc { + Arc::clone(&self.session_provider) + } + pub(crate) async fn execute( &self, request: codex_code_mode::ExecuteRequest, ) -> Result { - self.session()?.execute(request).await + self.session().await?.execute(request).await } pub(crate) async fn wait( &self, request: codex_code_mode::WaitRequest, ) -> Result { - self.session()?.wait(request).await + self.session().await?.wait(request).await } pub(crate) async fn terminate( &self, cell_id: CellId, ) -> Result { - self.session()?.terminate(cell_id).await + self.session().await?.terminate(cell_id).await } pub(crate) async fn shutdown(&self) -> Result<(), String> { - match &self.session { - Some(session) => session.shutdown().await, - None => Ok(()), + self.shutting_down.store(true, Ordering::Release); + // Join any initialization already in progress without initializing an unused service. + match self + .session + .get_or_try_init(|| async { + Err::, String>( + "code mode session is shutting down".to_string(), + ) + }) + .await + { + Ok(session) => session.shutdown().await, + Err(_) => Ok(()), } } @@ -121,9 +140,7 @@ impl CodeModeService { ) -> Option { let turn = &step_context.turn; let tool_mode = effective_tool_mode(turn); - if !matches!(tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) - || self.session.is_none() - { + if !matches!(tool_mode, ToolMode::CodeMode | ToolMode::CodeModeOnly) { return None; } @@ -137,10 +154,27 @@ impl CodeModeService { ) } - fn session(&self) -> Result<&Arc, String> { + async fn session(&self) -> Result, String> { + if self.shutting_down.load(Ordering::Acquire) { + return Err("code mode session is shutting down".to_string()); + } self.session - .as_ref() - .ok_or_else(|| "code mode is unavailable".to_string()) + .get_or_try_init(|| async { + if self.shutting_down.load(Ordering::Acquire) { + return Err("code mode session is shutting down".to_string()); + } + let session = self + .session_provider + .create_session(self.dispatch_broker.clone()) + .await?; + if self.shutting_down.load(Ordering::Acquire) { + let _ = session.shutdown().await; + return Err("code mode session is shutting down".to_string()); + } + Ok(session) + }) + .await + .map(Arc::clone) } } @@ -325,10 +359,15 @@ fn build_freeform_tool_payload( #[cfg(test)] mod tests { + use std::sync::Arc; + + use super::CodeModeService; use super::build_nested_tool_payload; use super::truncate_code_mode_result; use crate::tools::context::ToolPayload; use codex_code_mode::CodeModeToolKind; + use codex_code_mode::ExecuteRequest; + use codex_code_mode::ProcessOwnedCodeModeSessionProvider; use codex_protocol::models::FunctionCallOutputContentItem; use codex_tools::ToolName; use serde_json::json; @@ -385,4 +424,28 @@ mod tests { }] ); } + + #[tokio::test] + async fn missing_process_host_is_reported_without_failing_service_creation() { + let service = CodeModeService::new(Arc::new( + ProcessOwnedCodeModeSessionProvider::with_host_program( + "codex-code-mode-host-does-not-exist".into(), + ), + )); + + let error = service + .execute(ExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: "text('unreachable')".to_string(), + yield_time_ms: None, + max_output_tokens: None, + }) + .await + .err() + .expect("missing host should reject execution"); + + assert!(error.contains("failed to spawn code-mode host")); + service.shutdown().await.expect("shutdown unused service"); + } } diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 92b6c9d473..26c041d8f7 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -219,6 +219,32 @@ async fn run_code_mode_turn_with_model_and_config( Ok((test, second_mock)) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn missing_process_host_returns_a_tool_error() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let (_test, follow_up_mock) = + run_code_mode_turn_with_config(&server, "Run code mode", "text('unreachable')", |config| { + config + .features + .enable(Feature::CodeModeHost) + .expect("code mode host should be enabled"); + }) + .await?; + + let output = follow_up_mock + .single_request() + .custom_tool_call_output("call-1"); + assert!( + output["output"] + .as_str() + .is_some_and(|output| output.contains("failed to spawn code-mode host")) + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn code_mode_can_call_standalone_web_search() -> Result<()> { assert_code_mode_standalone_web_search(WebSearchMode::Live, serde_json::json!(true)).await diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 3bce1a7d12..5e9fbe5282 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -92,6 +92,8 @@ pub enum Feature { // Experimental /// Enable JavaScript code mode backed by the in-process V8 runtime. CodeMode, + /// Run JavaScript code mode in the standalone host process. + CodeModeHost, /// Restrict model-visible tools to code mode entrypoints (`exec`, `wait`). CodeModeOnly, /// Use the single unified PTY-backed exec tool. @@ -850,6 +852,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::CodeModeHost, + key: "code_mode_host", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::CodeModeOnly, key: "code_mode_only", From 451c0a437f2434a95c5e3893613f8afb192f4781 Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Fri, 26 Jun 2026 00:47:27 -0700 Subject: [PATCH 19/38] [codex] fix CreateThreadParams test initializer (#30198) ## Summary - initialize `selected_capability_roots` in the new `attach_in_memory_thread_store` test helper - restore `codex-core` test compilation on `main` ## Root cause [#30144](https://github.com/openai/codex/pull/30144) added the helper from commit `0c3d0742`, whose parent was `c38b2e9b`. That branch was based before [#29856](https://github.com/openai/codex/pull/29856) added `selected_capability_roots` as a required field on `CreateThreadParams`. The PR's Rust and Bazel workflows both passed against the stale branch head `0c3d0742`. When #30144 was squashed onto newer `main`, its initializer was integrated alongside the required field from #29856, producing `E0063` in `core/src/session/tests.rs`. Because those workflows tested the branch head rather than the integrated merge result, they did not see the version-skew failure before merge. ## Impact Any job that compiles the `codex-core` library tests fails, which turned the main-branch `rust-ci-full` and `Bazel` workflows red across platforms and blocks unrelated focused core tests. This change only completes the test initializer; it does not alter production behavior or workflow configuration. ## Validation - `just fmt` - `just test -p codex-core turn_complete_flushes_terminal_event_after_delivery` (1 passed, 2909 skipped) - `git diff --check` From 6d2168f06ae275d5e1f73cabf935d2bcc8549998 Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 26 Jun 2026 09:27:41 +0100 Subject: [PATCH 20/38] Reuse MCP runtimes when selected availability changes nothing (#30148) ## Why MCP runtime reuse was keyed by every ready selected-capability environment, even when an environment contributed no MCP servers or connectors. For example: 1. a global stdio MCP is running; 2. a selected remote environment contains only a skill; 3. that environment becomes ready; 4. the MCP and connector projection stays exactly the same; 5. Codex nevertheless rebuilds the MCP manager and restarts the global stdio process. That restart can interrupt active calls and discard process-local state even though nothing about MCP changed. ## What changes When selected-environment availability changes, Codex now resolves the candidate MCP and connector projection before deciding whether to replace the runtime: - if the winning MCP servers or their ownership change, rebuild as before; - if the selected connector snapshot changes, rebuild as before; - if an enabled MCP is explicitly bound to an environment whose availability changed, rebuild as before; - otherwise, keep the exact live manager and processes, and update only the availability input remembered by the snapshot. ```text ready selected environments: [] -> [skills-env] resolved MCP servers: {global_probe} -> {global_probe} resolved connectors: {} -> {} result: reuse manager; keep the same process ``` The comparison uses the resolved winning servers and their sources, so plugin/config ownership remains part of the runtime identity. ## Existing stack coverage The integration PR directly below this one already covers both rebuild boundaries: a selected MCP becomes callable and a selected connector tool becomes model-visible when their environment becomes available. It also verifies that an unchanged selected MCP runtime keeps its process. This PR does not add another remote-attachment integration scenario for the no-change optimization. `environment/add` returns before readiness, and app-server does not currently expose a deterministic readiness signal for an environment that contributes only skills. Keeping a fixed-delay test would add flake risk; adding a new readiness API would be outside this fix. ## Scope and assumptions - This does not change skill discovery, World State rendering, or plugin metadata caching. - This does not add file watching or hot reload behavior. - This does not change disconnect/reconnect handling. - Selected environment IDs and their capability contents retain the stack's existing stability assumption. - Delayed `required = true` executor MCP behavior remains out of scope. --- codex-rs/codex-mcp/src/catalog.rs | 5 +++++ codex-rs/core/src/session/mcp.rs | 30 ++++++++++++++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/codex-rs/codex-mcp/src/catalog.rs b/codex-rs/codex-mcp/src/catalog.rs index 689f93d813..d60ea843df 100644 --- a/codex-rs/codex-mcp/src/catalog.rs +++ b/codex-rs/codex-mcp/src/catalog.rs @@ -371,6 +371,11 @@ impl ResolvedMcpCatalog { .collect() } + /// Returns whether both catalogs resolve to the same winning servers and sources. + pub fn has_same_servers(&self, other: &Self) -> bool { + self.servers == other.servers + } + /// Replaces the resolved server set while preserving known server sources. /// /// Names not present in the existing catalog are treated as config-owned. diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 83a59c960a..25d95e30ea 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -133,6 +133,36 @@ impl Session { &available_environment_ids, ) .await; + let changed_environment_is_used_by_mcp = mcp_config + .mcp_server_catalog + .configured_servers() + .values() + .any(|server| { + let was_available = current + .available_environment_ids() + .contains(&server.environment_id); + let is_available = available_environment_ids.contains(&server.environment_id); + server.enabled && was_available != is_available + }); + if !changed_environment_is_used_by_mcp + && current + .config() + .mcp_server_catalog + .has_same_servers(&mcp_config.mcp_server_catalog) + && current.config().connector_snapshot == mcp_config.connector_snapshot + { + // Availability is only an input to the MCP projection. When that input changes but + // the projected servers and connectors do not, advance the input key without + // replacing the live manager and restarting its processes. + let runtime = Arc::new(McpRuntimeSnapshot::new( + Arc::new(current.config().clone()), + current.manager_arc(), + current.runtime_context().clone(), + available_environment_ids, + )); + self.services.mcp_runtime.store(Some(Arc::clone(&runtime))); + return runtime; + } self.refresh_mcp_servers_inner( turn_context, mcp_config, From 3c03bb4f182c91b07c12bcd4d79ac61df10dd084 Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 26 Jun 2026 11:02:27 +0100 Subject: [PATCH 21/38] Test selected capabilities across unavailable resume (#30215) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Why The selected-capability integration test already covers initial attachment and cold resume, but it resumes while the selected executor is still reachable. That leaves an important World State transition untested: a thread remembers its selected capability root, resumes while that environment is unavailable, and later sees the same stable environment return. ## What this tests This extends the existing end-to-end scenario: ```text selected executor available ↓ app-server stops and the executor goes away ↓ thread resumes with the executor unavailable ↓ skills, selected MCP tools, and connector attribution are absent ↓ the same environment ID is attached again ↓ skills, MCP tools, and connector attribution return ``` The test also checks that the unavailable snapshot explicitly tells the model that no selected-environment skills are currently available. After reattachment, it invokes the selected skill again and verifies that a new executor-owned MCP process starts. ## Scope This is test-only. It keeps the existing assumption that an environment ID refers to stable capability contents. It does not add package-file invalidation or live transport reconnect behavior. --- .../suite/v2/selected_capability_stack.rs | 60 +++++++++++++++---- 1 file changed, 48 insertions(+), 12 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/selected_capability_stack.rs b/codex-rs/app-server/tests/suite/v2/selected_capability_stack.rs index 052f286ffc..6a70116aef 100644 --- a/codex-rs/app-server/tests/suite/v2/selected_capability_stack.rs +++ b/codex-rs/app-server/tests/suite/v2/selected_capability_stack.rs @@ -57,6 +57,7 @@ const SKILL_NAME: &str = "executor-demo:deploy"; const SKILL_DESCRIPTION: &str = "Deploy through the selected executor."; const SKILL_BODY_MARKER: &str = "SELECTED_EXECUTOR_SKILL_BODY"; const LOCAL_SKILL_BODY_MARKER: &str = "COLLIDING_LOCAL_SKILL_BODY"; +const NO_SELECTED_SKILLS_MESSAGE: &str = "No selected-environment skills are currently available."; const MCP_SERVER_NAME: &str = "executor_probe"; const MCP_CALL_ID: &str = "selected-executor-mcp-call"; const CONNECTOR_ID: &str = "calendar"; @@ -122,9 +123,17 @@ async fn selected_capability_stack_tracks_environment_availability_and_resume() responses::ev_completed("unchanged-step"), ]), responses::sse(vec![ - responses::ev_response_created("resumed-step"), - responses::ev_assistant_message("resumed-message", "Ready after resume"), - responses::ev_completed("resumed-step"), + responses::ev_response_created("resumed-unavailable-step"), + responses::ev_assistant_message( + "resumed-unavailable-message", + "Unavailable after resume", + ), + responses::ev_completed("resumed-unavailable-step"), + ]), + responses::sse(vec![ + responses::ev_response_created("reattached-step"), + responses::ev_assistant_message("reattached-message", "Ready after reattach"), + responses::ev_completed("reattached-step"), ]), ], ) @@ -172,6 +181,7 @@ async fn selected_capability_stack_tracks_environment_availability_and_resume() .await?; assert_eq!(first_mcp_pid, wait_for_pid_file(&fixture.pid_file).await?); + exec_server.kill().await?; drop(app_server); std::fs::remove_file(&fixture.pid_file)?; @@ -190,12 +200,30 @@ async fn selected_capability_stack_tracks_environment_availability_and_resume() .await??; let ThreadResumeResponse { thread, .. } = to_response(response)?; assert_eq!(thread_id, thread.id); + + run_turn( + &mut app_server, + &thread_id, + "Inspect capabilities while the selected executor is unavailable", + fixture.environment_cwd.clone(), + ) + .await?; + let requests = response_mock.requests(); + assert_eq!(5, requests.len()); + assert_selected_plugin_tools_absent(&requests[4]); + assert!( + latest_selected_skill_update(&requests[4]) + .is_some_and(|text| text.contains(NO_SELECTED_SKILLS_MESSAGE)) + ); + + exec_server = spawn_exec_server(fixture.codex_home.path(), &fixture.exec_server_url).await?; + add_environment(&mut app_server, &fixture.exec_server_url).await?; wait_for_selected_mcp_server(&mut app_server, &thread_id).await?; run_turn( &mut app_server, &thread_id, - &format!("Use ${SKILL_NAME} after resuming the thread"), + &format!("Use ${SKILL_NAME} after reattaching the selected executor"), fixture.environment_cwd, ) .await?; @@ -203,13 +231,13 @@ async fn selected_capability_stack_tracks_environment_availability_and_resume() assert_ne!(first_mcp_pid, resumed_mcp_pid); let requests = response_mock.requests(); - assert_eq!(5, requests.len()); + assert_eq!(6, requests.len()); for request in &requests[1..4] { assert_selected_skill_is_injected(request, /*expected_count*/ 1); assert_selected_plugin_tools(request); } - assert_selected_skill_is_injected(&requests[4], /*expected_count*/ 2); - assert_selected_plugin_tools(&requests[4]); + assert_selected_skill_is_injected(&requests[5], /*expected_count*/ 2); + assert_selected_plugin_tools(&requests[5]); let output = requests[2].function_call_output(MCP_CALL_ID); let output = output["output"] .as_str() @@ -508,6 +536,10 @@ fn assert_selected_capabilities_absent(request: &ResponsesRequest) { .into_iter() .all(|text| !text.contains(SKILL_DESCRIPTION)) ); + assert_selected_plugin_tools_absent(request); +} + +fn assert_selected_plugin_tools_absent(request: &ResponsesRequest) { assert!( request .tool_by_name(&format!("mcp__{MCP_SERVER_NAME}"), "echo") @@ -540,13 +572,17 @@ fn assert_selected_skill_is_injected(request: &ResponsesRequest, expected_count: } fn assert_selected_skill_catalog_available(request: &ResponsesRequest) { - let catalog_fragments = request + let catalog_fragment = latest_selected_skill_update(request) + .expect("selected skill catalog update should be model-visible"); + assert!(catalog_fragment.contains(SKILL_DESCRIPTION)); + assert!(catalog_fragment.contains("environment resource:")); +} + +fn latest_selected_skill_update(request: &ResponsesRequest) -> Option { + request .message_input_texts("developer") .into_iter() - .filter(|text| text.contains(SKILL_DESCRIPTION)) - .collect::>(); - assert_eq!(1, catalog_fragments.len()); - assert!(catalog_fragments[0].contains("environment resource:")); + .rfind(|text| text.contains(SKILL_DESCRIPTION) || text.contains(NO_SELECTED_SKILLS_MESSAGE)) } fn assert_selected_plugin_tools(request: &ResponsesRequest) { From 914c8eeb4e41442e04a7706ca1638491ceabf4a4 Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Fri, 26 Jun 2026 05:52:04 -0700 Subject: [PATCH 22/38] [codex] narrow unused skills intro export (#29991) ## Summary - stop publicly re-exporting the internally used `SKILLS_INTRO_WITH_ALIASES` constant - keep the constant and all skills rendering behavior unchanged - preserve every integration helper, API, fixture, assertion, and module used by tests ## Scope guardrails This revision keeps all remote/network-facing functionality and every line introduced by `jif `. Following the test-preservation audit, it also restores the in-process RMCP test transport, the original `codex-mcp` fixture, `PluginLoadOutcome::effective_skill_roots` and its assertions, the `EffectiveSkillRoots` API family, the test-only apps renderer, and the TUI dead-code annotation. Those files now match the PR base exactly. No test imports or directly references the remaining public skills export being narrowed. ## Validation - repository-wide test-reference audit: no test-used code remains deleted or narrowed - deleted-line `git blame` audit: zero Jif-authored deletions - `cargo test -p codex-core-plugins -p codex-mcp -p codex-rmcp-client --lib`: 467 passed - `cargo test -p codex-core --lib apps::render`: 2 passed - `cargo test -p codex-core-skills --lib render::tests`: 19 passed - `cargo check -p codex-core-skills --all-targets`: passed - `just fix -p codex-core-skills`: passed - `just fmt`: passed - `git diff --check`: passed The full local `codex-core-skills` suite passed 106/108 tests; two loader tests detected an ambient repository skills root outside the package and failed their isolation assertions. The scoped renderer suite and all-target compile pass, and CI runs in an isolated environment. Final code delta: 1 insertion, 2 deletions across 2 files. --- codex-rs/core-skills/src/lib.rs | 1 - codex-rs/core-skills/src/render.rs | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/codex-rs/core-skills/src/lib.rs b/codex-rs/core-skills/src/lib.rs index 53ffcddd07..dcfc6568a5 100644 --- a/codex-rs/core-skills/src/lib.rs +++ b/codex-rs/core-skills/src/lib.rs @@ -24,7 +24,6 @@ pub use render::AvailableSkills; pub use render::SKILLS_HOW_TO_USE_WITH_ABSOLUTE_PATHS; pub use render::SKILLS_HOW_TO_USE_WITH_ALIASES; pub use render::SKILLS_INTRO_WITH_ABSOLUTE_PATHS; -pub use render::SKILLS_INTRO_WITH_ALIASES; pub use render::SkillMetadataBudget; pub use render::SkillRenderReport; pub use render::build_available_skills; diff --git a/codex-rs/core-skills/src/render.rs b/codex-rs/core-skills/src/render.rs index 89ea1675d0..2e214b1335 100644 --- a/codex-rs/core-skills/src/render.rs +++ b/codex-rs/core-skills/src/render.rs @@ -26,7 +26,7 @@ pub const SKILL_DESCRIPTION_TRUNCATED_WARNING_WITH_PERCENT: &str = "Skill descri pub const SKILL_DESCRIPTIONS_REMOVED_WARNING_PREFIX: &str = "Exceeded skills context budget. All skill descriptions were removed and"; pub const SKILLS_INTRO_WITH_ABSOLUTE_PATHS: &str = "A skill is a set of instructions provided through a `SKILL.md` source. Below is the list of skills that can be used. Each entry includes a name, description, and source locator. `file` locators are on the host filesystem, `environment resource` locators are owned by an execution environment, `orchestrator resource` locators are opaque non-filesystem resources, and `custom resource` locators use their provider's access mechanism."; -pub const SKILLS_INTRO_WITH_ALIASES: &str = "A skill is a set of local instructions to follow that is stored in a `SKILL.md` file. Below is the list of skills that can be used. Each entry includes a name, description, and a short path that can be expanded into an absolute path using the skill roots table."; +const SKILLS_INTRO_WITH_ALIASES: &str = "A skill is a set of local instructions to follow that is stored in a `SKILL.md` file. Below is the list of skills that can be used. Each entry includes a name, description, and a short path that can be expanded into an absolute path using the skill roots table."; pub const SKILLS_HOW_TO_USE_WITH_ABSOLUTE_PATHS: &str = r###"- Discovery: The list above is the skills available in this session (name + description + source locator). `file` entries live on the host filesystem, `environment resource` entries are owned by their execution environment, `orchestrator resource` entries must be accessed through `skills.list` and `skills.read`, and `custom resource` entries use their provider's access mechanism. - Trigger rules: If the user names a skill (with `$SkillName` or plain text) OR the task clearly matches a skill's description shown above, you must use that skill for that turn. Multiple mentions mean use them all. Do not carry skills across turns unless re-mentioned. - Missing/blocked: If a named skill isn't in the list or its source can't be read, say so briefly and continue with the best fallback. From 2c5bc5e284402d2e618c15e396d2e81bcc1ec257 Mon Sep 17 00:00:00 2001 From: charlesgong-openai Date: Fri, 26 Jun 2026 11:24:12 -0400 Subject: [PATCH 23/38] Relax hooks.json top-level metadata validation (#30229) ## Summary - Allow a top-level `description` string in `hooks.json`. - Continue rejecting unknown top-level keys and root-level hook events; events must remain under `hooks`. ## Testing - `just test -p codex-config` --- codex-rs/config/src/hook_config.rs | 2 ++ codex-rs/config/src/hooks_tests.rs | 2 ++ 2 files changed, 4 insertions(+) diff --git a/codex-rs/config/src/hook_config.rs b/codex-rs/config/src/hook_config.rs index a7ee820e60..8e18cb935a 100644 --- a/codex-rs/config/src/hook_config.rs +++ b/codex-rs/config/src/hook_config.rs @@ -10,6 +10,8 @@ use serde::Serialize; #[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] pub struct HooksFile { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub description: Option, #[serde(default)] pub hooks: HookEventsToml, } diff --git a/codex-rs/config/src/hooks_tests.rs b/codex-rs/config/src/hooks_tests.rs index a2eb2d2322..3d51840fc2 100644 --- a/codex-rs/config/src/hooks_tests.rs +++ b/codex-rs/config/src/hooks_tests.rs @@ -13,6 +13,7 @@ use super::MatcherGroup; fn hooks_file_deserializes_existing_json_shape() { let parsed: HooksFile = serde_json::from_str( r#"{ + "description": "Optional stop-time review gate for Codex Companion.", "hooks": { "PreToolUse": [ { @@ -35,6 +36,7 @@ fn hooks_file_deserializes_existing_json_shape() { assert_eq!( parsed, HooksFile { + description: Some("Optional stop-time review gate for Codex Companion.".to_string()), hooks: HookEventsToml { pre_tool_use: vec![MatcherGroup { matcher: Some("^Bash$".to_string()), From 5267e805fb830891c0b23376bcd9cbd382c3473c Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Fri, 26 Jun 2026 09:12:42 -0700 Subject: [PATCH 24/38] feat(app-server): add history_mode to thread (#29927) ## Description This PR adds a new `historyMode = "legacy" | "paginated"` to `Thread`. This will be stored in `SessionMeta` in the JSONL rollout file and as a new column in the SQLite thread_metadata table, and exposed on `thread/start` and on the `Thread` object in app-server. ## What changed - Added canonical `ThreadHistoryMode` with `legacy` and `paginated`, defaulting old and new SessionMeta to `legacy`. - Carried `history_mode` through core session config, ThreadStore stored metadata, local/in-memory stores, rollout metadata extraction, and the existing SQLite `threads` table. - Added experimental `historyMode` to app-server v2 `Thread` and `thread/start`. - Made paginated stored threads metadata-discoverable but unsupported for legacy full-history reads, `load_history`, live resume, and create paths. - Regenerated app-server schema fixtures and added protocol/state/thread-store/app-server coverage for persistence and fail-closed behavior. ## Compatibility floor Because users may be running various versions of Codex binaries on the same machine (TUI, Codex App, etc.), we will need to establish a compatibility floor for upcoming paginated threads, which will change how thread storage reads and writes work. The overall plan here: ``` Release N: - Add historyMode to SessionMeta / Thread / SQLite metadata. - Teach binaries to understand paginated threads. - If a binary sees `historyMode="paginated"` but does not support the paginated contract, it refuses to resume/mutate the thread. - Default remains `"legacy"`. Release N+1: - First-party clients start opting into paginated threads where appropriate. - Internal dogfood / staged rollout. - Measure old-client usage and paginated-thread unsupported errors. Release N+2: - Only after Release N+ is overwhelmingly deployed, make paginated the default. - Accept that a small tail of N-1-or-older binaries may not understand paginated threads. ``` The important behavior change is fail-closed handling for a binary that encounters a persisted `paginated` thread before it knows how to fully support paginated history. In app-server, if a thread is `paginated`, we will: - allow metadata-only discovery paths like `thread/list` and `thread/read(includeTurns=false)`, so clients can still see the thread and inspect its `historyMode` - reject legacy full-history/live-thread paths like `thread/read(includeTurns=true)` and `thread/resume` with an unsupported JSON-RPC error - avoid silently treating an unknown or future `historyMode` as `legacy` Under the hood, the ThreadStore layer also rejects legacy operations that would need to load or replay the full thread history for a paginated thread. That gives us the behavior we want for Release N: future paginated threads are visible, but this binary fails closed instead of trying to operate on them as if they were legacy threads. --- .../analytics/src/analytics_client_tests.rs | 1 + codex-rs/analytics/src/client_tests.rs | 1 + .../schema/json/ClientRequest.json | 7 + .../schema/json/ServerNotification.json | 7 + .../codex_app_server_protocol.schemas.json | 7 + .../codex_app_server_protocol.v2.schemas.json | 7 + .../schema/json/v2/ThreadForkResponse.json | 7 + .../schema/json/v2/ThreadListResponse.json | 7 + .../json/v2/ThreadMetadataUpdateResponse.json | 7 + .../schema/json/v2/ThreadReadResponse.json | 7 + .../schema/json/v2/ThreadResumeResponse.json | 7 + .../json/v2/ThreadRollbackResponse.json | 7 + .../schema/json/v2/ThreadStartParams.json | 7 + .../schema/json/v2/ThreadStartResponse.json | 7 + .../json/v2/ThreadStartedNotification.json | 7 + .../json/v2/ThreadUnarchiveResponse.json | 7 + .../schema/typescript/v2/ThreadHistoryMode.ts | 5 + .../schema/typescript/v2/index.ts | 1 + .../src/protocol/common.rs | 2 + .../src/protocol/v2/tests.rs | 1 + .../src/protocol/v2/thread.rs | 5 + .../src/protocol/v2/thread_data.rs | 32 +++ .../app-server/src/bespoke_event_handling.rs | 1 + codex-rs/app-server/src/request_processors.rs | 2 + .../external_agent_session_import.rs | 2 + .../request_processors/thread_processor.rs | 14 ++ .../thread_processor_tests.rs | 2 + .../thread_resume_redaction.rs | 1 + .../src/request_processors/thread_summary.rs | 1 + codex-rs/app-server/src/thread_status.rs | 1 + codex-rs/app-server/tests/common/rollout.rs | 2 + .../tests/suite/conversation_summary.rs | 1 + .../tests/suite/v2/remote_thread_store.rs | 1 + .../app-server/tests/suite/v2/skills_list.rs | 1 + .../app-server/tests/suite/v2/thread_read.rs | 137 ++++++++++++ .../tests/suite/v2/thread_resume.rs | 1 + .../app-server/tests/suite/v2/thread_start.rs | 50 +++++ .../tests/suite/v2/thread_unarchive.rs | 1 + codex-rs/core/src/agent/control_tests.rs | 2 + codex-rs/core/src/codex_delegate.rs | 1 + codex-rs/core/src/codex_thread.rs | 2 + .../core/src/personality_migration_tests.rs | 1 + codex-rs/core/src/realtime_context_tests.rs | 1 + codex-rs/core/src/session/mod.rs | 7 + codex-rs/core/src/session/session.rs | 5 + codex-rs/core/src/session/tests.rs | 10 + .../core/src/session/tests/guardian_tests.rs | 1 + .../core/src/session_rollout_init_error.rs | 8 + codex-rs/core/src/thread_manager.rs | 12 + codex-rs/core/src/thread_manager_tests.rs | 6 + codex-rs/core/tests/common/test_codex.rs | 1 + codex-rs/core/tests/suite/agents_md.rs | 2 + .../core/tests/suite/personality_migration.rs | 2 + codex-rs/core/tests/suite/sqlite_state.rs | 1 + .../tests/suite/subagent_notifications.rs | 1 + codex-rs/exec/src/lib_tests.rs | 2 + codex-rs/memories/write/src/runtime.rs | 1 + codex-rs/protocol/src/protocol.rs | 109 ++++++++++ codex-rs/rollout/src/compression_tests.rs | 1 + codex-rs/rollout/src/list.rs | 23 +- codex-rs/rollout/src/metadata.rs | 1 + codex-rs/rollout/src/metadata_tests.rs | 39 ++++ codex-rs/rollout/src/recorder.rs | 38 ++++ codex-rs/rollout/src/recorder_tests.rs | 43 ++++ codex-rs/rollout/src/session_index_tests.rs | 1 + codex-rs/rollout/src/state_db_tests.rs | 1 + codex-rs/rollout/src/tests.rs | 47 ++++ .../migrations/0040_threads_history_mode.sql | 1 + codex-rs/state/src/extract.rs | 4 + codex-rs/state/src/model/thread_metadata.rs | 23 ++ codex-rs/state/src/runtime/memories.rs | 25 ++- codex-rs/state/src/runtime/test_support.rs | 3 + codex-rs/state/src/runtime/threads.rs | 38 +++- codex-rs/thread-store/src/error.rs | 12 + codex-rs/thread-store/src/in_memory.rs | 205 +++++++++++++++++- .../thread-store/src/local/create_thread.rs | 3 + codex-rs/thread-store/src/local/helpers.rs | 1 + .../thread-store/src/local/list_threads.rs | 2 + .../thread-store/src/local/live_writer.rs | 27 ++- codex-rs/thread-store/src/local/mod.rs | 104 ++++++++- .../thread-store/src/local/read_thread.rs | 86 ++++++-- .../thread-store/src/local/test_support.rs | 17 ++ .../src/local/update_thread_metadata.rs | 71 ++++++ codex-rs/thread-store/src/store.rs | 9 + .../thread-store/src/thread_metadata_sync.rs | 3 + codex-rs/thread-store/src/types.rs | 46 ++++ codex-rs/tui/src/app/loaded_threads.rs | 1 + codex-rs/tui/src/app/tests.rs | 4 + codex-rs/tui/src/app/thread_session_state.rs | 1 + codex-rs/tui/src/app_server_session.rs | 1 + codex-rs/tui/src/resume_picker.rs | 4 + 91 files changed, 1385 insertions(+), 39 deletions(-) create mode 100644 codex-rs/app-server-protocol/schema/typescript/v2/ThreadHistoryMode.ts create mode 100644 codex-rs/state/migrations/0040_threads_history_mode.sql diff --git a/codex-rs/analytics/src/analytics_client_tests.rs b/codex-rs/analytics/src/analytics_client_tests.rs index 4451d39ed9..0f759fd938 100644 --- a/codex-rs/analytics/src/analytics_client_tests.rs +++ b/codex-rs/analytics/src/analytics_client_tests.rs @@ -198,6 +198,7 @@ fn sample_thread_with_metadata( parent_thread_id, preview: "first prompt".to_string(), ephemeral, + history_mode: Default::default(), model_provider: "openai".to_string(), created_at: 1, updated_at: 2, diff --git a/codex-rs/analytics/src/client_tests.rs b/codex-rs/analytics/src/client_tests.rs index bbc5b5f0b6..ac50d6ebc7 100644 --- a/codex-rs/analytics/src/client_tests.rs +++ b/codex-rs/analytics/src/client_tests.rs @@ -283,6 +283,7 @@ fn sample_thread(thread_id: &str) -> Thread { parent_thread_id: None, preview: "first prompt".to_string(), ephemeral: false, + history_mode: Default::default(), model_provider: "openai".to_string(), created_at: 1, updated_at: 2, diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index acfffc370f..8c1d700ef2 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -3747,6 +3747,13 @@ ], "type": "string" }, + "ThreadHistoryMode": { + "enum": [ + "legacy", + "paginated" + ], + "type": "string" + }, "ThreadInjectItemsParams": { "properties": { "items": { diff --git a/codex-rs/app-server-protocol/schema/json/ServerNotification.json b/codex-rs/app-server-protocol/schema/json/ServerNotification.json index 43fb661822..23dae42a9d 100644 --- a/codex-rs/app-server-protocol/schema/json/ServerNotification.json +++ b/codex-rs/app-server-protocol/schema/json/ServerNotification.json @@ -3917,6 +3917,13 @@ ], "type": "object" }, + "ThreadHistoryMode": { + "enum": [ + "legacy", + "paginated" + ], + "type": "string" + }, "ThreadId": { "type": "string" }, 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 277421a852..f1d6a1af85 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 @@ -17512,6 +17512,13 @@ "title": "ThreadGoalUpdatedNotification", "type": "object" }, + "ThreadHistoryMode": { + "enum": [ + "legacy", + "paginated" + ], + "type": "string" + }, "ThreadId": { "type": "string" }, 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 39c95053d0..9983b79257 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 @@ -15291,6 +15291,13 @@ "title": "ThreadGoalUpdatedNotification", "type": "object" }, + "ThreadHistoryMode": { + "enum": [ + "legacy", + "paginated" + ], + "type": "string" + }, "ThreadId": { "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadForkResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadForkResponse.json index 3ca62f17f2..5d20d64e68 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadForkResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadForkResponse.json @@ -1188,6 +1188,13 @@ "description": "Extra app-server data for a thread.", "type": "object" }, + "ThreadHistoryMode": { + "enum": [ + "legacy", + "paginated" + ], + "type": "string" + }, "ThreadId": { "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadListResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadListResponse.json index a81f643122..e5cbc2e606 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadListResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadListResponse.json @@ -995,6 +995,13 @@ "description": "Extra app-server data for a thread.", "type": "object" }, + "ThreadHistoryMode": { + "enum": [ + "legacy", + "paginated" + ], + "type": "string" + }, "ThreadId": { "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadMetadataUpdateResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadMetadataUpdateResponse.json index d007b32d22..cdd1d47a69 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadMetadataUpdateResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadMetadataUpdateResponse.json @@ -995,6 +995,13 @@ "description": "Extra app-server data for a thread.", "type": "object" }, + "ThreadHistoryMode": { + "enum": [ + "legacy", + "paginated" + ], + "type": "string" + }, "ThreadId": { "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadReadResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadReadResponse.json index f45efe3707..613dd9648e 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadReadResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadReadResponse.json @@ -995,6 +995,13 @@ "description": "Extra app-server data for a thread.", "type": "object" }, + "ThreadHistoryMode": { + "enum": [ + "legacy", + "paginated" + ], + "type": "string" + }, "ThreadId": { "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeResponse.json index 09dd3fedbf..ee1ca83917 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadResumeResponse.json @@ -1188,6 +1188,13 @@ "description": "Extra app-server data for a thread.", "type": "object" }, + "ThreadHistoryMode": { + "enum": [ + "legacy", + "paginated" + ], + "type": "string" + }, "ThreadId": { "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackResponse.json index 8083b52d49..793f11c226 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadRollbackResponse.json @@ -995,6 +995,13 @@ "description": "Extra app-server data for a thread.", "type": "object" }, + "ThreadHistoryMode": { + "enum": [ + "legacy", + "paginated" + ], + "type": "string" + }, "ThreadId": { "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json index 50c0a9cf7f..3a23d13d22 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json @@ -241,6 +241,13 @@ ], "type": "object" }, + "ThreadHistoryMode": { + "enum": [ + "legacy", + "paginated" + ], + "type": "string" + }, "ThreadSource": { "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartResponse.json index b33a64a200..9d16598696 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartResponse.json @@ -1188,6 +1188,13 @@ "description": "Extra app-server data for a thread.", "type": "object" }, + "ThreadHistoryMode": { + "enum": [ + "legacy", + "paginated" + ], + "type": "string" + }, "ThreadId": { "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartedNotification.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartedNotification.json index 9040457f22..a56d59c01d 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartedNotification.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartedNotification.json @@ -995,6 +995,13 @@ "description": "Extra app-server data for a thread.", "type": "object" }, + "ThreadHistoryMode": { + "enum": [ + "legacy", + "paginated" + ], + "type": "string" + }, "ThreadId": { "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadUnarchiveResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadUnarchiveResponse.json index d4e9fea618..6cc9b44cba 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadUnarchiveResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadUnarchiveResponse.json @@ -995,6 +995,13 @@ "description": "Extra app-server data for a thread.", "type": "object" }, + "ThreadHistoryMode": { + "enum": [ + "legacy", + "paginated" + ], + "type": "string" + }, "ThreadId": { "type": "string" }, diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ThreadHistoryMode.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadHistoryMode.ts new file mode 100644 index 0000000000..db0f2d824e --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadHistoryMode.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 ThreadHistoryMode = "legacy" | "paginated"; 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 82b74711c6..2833e6e525 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts @@ -410,6 +410,7 @@ export type { ThreadGoalSetParams } from "./ThreadGoalSetParams"; export type { ThreadGoalSetResponse } from "./ThreadGoalSetResponse"; export type { ThreadGoalStatus } from "./ThreadGoalStatus"; export type { ThreadGoalUpdatedNotification } from "./ThreadGoalUpdatedNotification"; +export type { ThreadHistoryMode } from "./ThreadHistoryMode"; export type { ThreadInjectItemsParams } from "./ThreadInjectItemsParams"; export type { ThreadInjectItemsResponse } from "./ThreadInjectItemsResponse"; export type { ThreadItem } from "./ThreadItem"; diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index 3120007e53..7b15035575 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -2580,6 +2580,7 @@ mod tests { parent_thread_id: None, preview: "first prompt".to_string(), ephemeral: true, + history_mode: Default::default(), model_provider: "openai".to_string(), created_at: 1, updated_at: 2, @@ -2630,6 +2631,7 @@ mod tests { "parentThreadId": null, "preview": "first prompt", "ephemeral": true, + "historyMode": "legacy", "modelProvider": "openai", "createdAt": 1, "updatedAt": 2, diff --git a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs index c6d460c721..ddd2afe8f9 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs @@ -175,6 +175,7 @@ fn thread_resume_response_round_trips_initial_turns_page() { parent_thread_id: None, preview: String::new(), ephemeral: false, + history_mode: Default::default(), model_provider: "openai".to_string(), created_at: 1, updated_at: 1, diff --git a/codex-rs/app-server-protocol/src/protocol/v2/thread.rs b/codex-rs/app-server-protocol/src/protocol/v2/thread.rs index e712ab0ad3..73335fce56 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/thread.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/thread.rs @@ -4,6 +4,7 @@ use super::AskForApproval; use super::SandboxMode; use super::SandboxPolicy; use super::Thread; +use super::ThreadHistoryMode; use super::ThreadItem; use super::ThreadSource; use super::Turn; @@ -105,6 +106,10 @@ pub struct ThreadStartParams { pub multi_agent_mode: Option, #[ts(optional = nullable)] pub ephemeral: Option, + /// Persisted thread history contract to use for this new thread. + #[experimental("thread/start.historyMode")] + #[ts(optional = nullable)] + pub history_mode: Option, #[ts(optional = nullable)] pub session_start_source: Option, /// Optional client-supplied analytics source classification for this thread. diff --git a/codex-rs/app-server-protocol/src/protocol/v2/thread_data.rs b/codex-rs/app-server-protocol/src/protocol/v2/thread_data.rs index 3338513df7..2393efb25c 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/thread_data.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/thread_data.rs @@ -5,6 +5,7 @@ use super::TurnStatus; use codex_experimental_api_macros::ExperimentalApi; use codex_protocol::protocol::SessionSource as CoreSessionSource; use codex_protocol::protocol::SubAgentSource as CoreSubAgentSource; +use codex_protocol::protocol::ThreadHistoryMode as CoreThreadHistoryMode; use codex_protocol::protocol::ThreadSource as CoreThreadSource; use codex_utils_absolute_path::AbsolutePathBuf; use schemars::JsonSchema; @@ -64,6 +65,33 @@ impl From for CoreSessionSource { } } +#[derive(Default, Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "lowercase")] +#[ts(rename_all = "lowercase", export_to = "v2/")] +pub enum ThreadHistoryMode { + #[default] + Legacy, + Paginated, +} + +impl From for ThreadHistoryMode { + fn from(value: CoreThreadHistoryMode) -> Self { + match value { + CoreThreadHistoryMode::Legacy => Self::Legacy, + CoreThreadHistoryMode::Paginated => Self::Paginated, + } + } +} + +impl From for CoreThreadHistoryMode { + fn from(value: ThreadHistoryMode) -> Self { + match value { + ThreadHistoryMode::Legacy => Self::Legacy, + ThreadHistoryMode::Paginated => Self::Paginated, + } + } +} + #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, TS)] #[serde(try_from = "String", into = "String")] #[ts(type = "string")] @@ -155,6 +183,10 @@ pub struct Thread { pub preview: String, /// Whether the thread is ephemeral and should not be materialized on disk. pub ephemeral: bool, + /// Persisted thread history contract selected when this thread was created. + #[experimental("thread.historyMode")] + #[serde(default)] + pub history_mode: ThreadHistoryMode, /// Model provider used for this thread (for example, 'openai'). pub model_provider: String, /// Unix timestamp (in seconds) when the thread was created. diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index bcba603bcd..c63103161f 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -2287,6 +2287,7 @@ mod tests { cwd: test_path_buf("/tmp").abs().into(), cli_version: "0.0.0".to_string(), source: SessionSource::Cli, + history_mode: Default::default(), thread_source: None, agent_nickname: None, agent_role: None, diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 57a65fd313..39c4b79dc4 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -206,6 +206,8 @@ use codex_app_server_protocol::ThreadGoalSetResponse; use codex_app_server_protocol::ThreadGoalStatus; use codex_app_server_protocol::ThreadGoalUpdatedNotification; use codex_app_server_protocol::ThreadHistoryBuilder; +#[cfg(test)] +use codex_app_server_protocol::ThreadHistoryMode; use codex_app_server_protocol::ThreadIncrementElicitationParams; use codex_app_server_protocol::ThreadIncrementElicitationResponse; use codex_app_server_protocol::ThreadInjectItemsParams; diff --git a/codex-rs/app-server/src/request_processors/external_agent_session_import.rs b/codex-rs/app-server/src/request_processors/external_agent_session_import.rs index 15b6e3a43a..e0a53a9c98 100644 --- a/codex-rs/app-server/src/request_processors/external_agent_session_import.rs +++ b/codex-rs/app-server/src/request_processors/external_agent_session_import.rs @@ -15,6 +15,7 @@ use codex_models_manager::manager::RefreshStrategy; use codex_protocol::ThreadId; use codex_protocol::models::BaseInstructions; use codex_protocol::protocol::MultiAgentVersion; +use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadMemoryMode; use codex_rollout::is_persisted_rollout_item; use codex_thread_store::AppendThreadItemsParams; @@ -221,6 +222,7 @@ impl ExternalAgentSessionImporter { dynamic_tools: Vec::new(), selected_capability_roots: Vec::new(), multi_agent_version: Some(MultiAgentVersion::V1), + history_mode: ThreadHistoryMode::Legacy, initial_window_id: uuid::Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { cwd: Some(cwd.clone()), diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 8232d68fe6..d5302d9c43 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -5,6 +5,7 @@ use codex_extension_api::ExtensionDataInit; use codex_protocol::config_types::MultiAgentMode; use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_DANGER_FULL_ACCESS; use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_WORKSPACE; +use codex_protocol::protocol::ThreadHistoryMode; const THREAD_LIST_DEFAULT_LIMIT: usize = 25; const THREAD_LIST_MAX_LIMIT: usize = 100; @@ -926,6 +927,7 @@ impl ThreadRequestProcessor { personality, multi_agent_mode: _multi_agent_mode, ephemeral, + history_mode, session_start_source, thread_source, environments, @@ -980,6 +982,7 @@ impl ThreadRequestProcessor { typesafe_overrides, dynamic_tools, selected_capability_roots.unwrap_or_default(), + history_mode.map(Into::into), session_start_source, thread_source.map(Into::into), environment_selections, @@ -1055,6 +1058,7 @@ impl ThreadRequestProcessor { typesafe_overrides: ConfigOverrides, dynamic_tools: Option>, selected_capability_roots: Vec, + history_mode: Option, session_start_source: Option, thread_source: Option, environments: Option>, @@ -1173,6 +1177,7 @@ impl ThreadRequestProcessor { codex_app_server_protocol::ThreadStartSource::Startup => InitialHistory::New, codex_app_server_protocol::ThreadStartSource::Clear => InitialHistory::Cleared, }, + history_mode, session_source: None, thread_source, dynamic_tools, @@ -1190,6 +1195,7 @@ impl ThreadRequestProcessor { .await .map_err(|err| match err { CodexErr::InvalidRequest(message) => invalid_request(message), + CodexErr::UnsupportedOperation(message) => method_not_found(message), err => internal_error(format!("error creating thread: {err}")), })?; let session_telemetry = thread.session_telemetry(); @@ -2299,6 +2305,9 @@ impl ThreadRequestProcessor { Err(ThreadStoreError::InvalidRequest { message }) => { Err(ThreadReadViewError::InvalidRequest(message)) } + Err(ThreadStoreError::Unsupported { operation }) => { + Err(ThreadReadViewError::Unsupported(operation)) + } Err(err) => Err(ThreadReadViewError::Internal(format!( "failed to read thread: {err}" ))), @@ -2501,6 +2510,9 @@ impl ThreadRequestProcessor { Err(ThreadStoreError::InvalidRequest { message }) => { return Err(ThreadReadViewError::InvalidRequest(message)); } + Err(ThreadStoreError::Unsupported { operation }) => { + return Err(ThreadReadViewError::Unsupported(operation)); + } Err(err) => { return Err(ThreadReadViewError::Internal(format!( "failed to read thread: {err}" @@ -4254,6 +4266,7 @@ pub(crate) fn thread_from_stored_thread( parent_thread_id: thread.parent_thread_id.map(|id| id.to_string()), preview: thread.preview, ephemeral: false, + history_mode: thread.history_mode.into(), model_provider: if thread.model_provider.is_empty() { fallback_provider.to_string() } else { @@ -4465,6 +4478,7 @@ fn build_thread_from_snapshot( parent_thread_id: config_snapshot.parent_thread_id.map(|id| id.to_string()), preview: String::new(), ephemeral: config_snapshot.ephemeral, + history_mode: config_snapshot.history_mode.into(), model_provider: config_snapshot.model_provider_id.clone(), created_at: now, updated_at: now, diff --git a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs index 0d447011b8..8c0d59efed 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs @@ -483,6 +483,7 @@ mod thread_processor_behavior_tests { cwd: PathBuf::from("/tmp"), cli_version: "0.0.0".to_string(), source: SessionSource::Cli, + history_mode: Default::default(), thread_source: Some(codex_protocol::protocol::ThreadSource::User), agent_nickname: None, agent_role: None, @@ -777,6 +778,7 @@ mod thread_processor_behavior_tests { }, }, session_source: SessionSource::Cli, + history_mode: Default::default(), forked_from_thread_id: None, parent_thread_id: None, thread_source: None, diff --git a/codex-rs/app-server/src/request_processors/thread_resume_redaction.rs b/codex-rs/app-server/src/request_processors/thread_resume_redaction.rs index f753607273..90d25ff593 100644 --- a/codex-rs/app-server/src/request_processors/thread_resume_redaction.rs +++ b/codex-rs/app-server/src/request_processors/thread_resume_redaction.rs @@ -195,6 +195,7 @@ mod tests { parent_thread_id: None, preview: "preview".to_string(), ephemeral: false, + history_mode: Default::default(), model_provider: "mock_provider".to_string(), created_at: 0, updated_at: 0, diff --git a/codex-rs/app-server/src/request_processors/thread_summary.rs b/codex-rs/app-server/src/request_processors/thread_summary.rs index 2bacaf6446..37c77a198b 100644 --- a/codex-rs/app-server/src/request_processors/thread_summary.rs +++ b/codex-rs/app-server/src/request_processors/thread_summary.rs @@ -320,6 +320,7 @@ pub(crate) fn summary_to_thread( parent_thread_id: None, preview, ephemeral: false, + history_mode: ThreadHistoryMode::Legacy, model_provider, created_at: created_at.map(|dt| dt.timestamp()).unwrap_or(0), updated_at: updated_at.map(|dt| dt.timestamp()).unwrap_or(0), diff --git a/codex-rs/app-server/src/thread_status.rs b/codex-rs/app-server/src/thread_status.rs index 32b79e1f69..a29b649788 100644 --- a/codex-rs/app-server/src/thread_status.rs +++ b/codex-rs/app-server/src/thread_status.rs @@ -895,6 +895,7 @@ mod tests { parent_thread_id: None, preview: String::new(), ephemeral: false, + history_mode: Default::default(), model_provider: "mock-provider".to_string(), created_at: 0, updated_at: 0, diff --git a/codex-rs/app-server/tests/common/rollout.rs b/codex-rs/app-server/tests/common/rollout.rs index 8d2567c010..c13f2355d4 100644 --- a/codex-rs/app-server/tests/common/rollout.rs +++ b/codex-rs/app-server/tests/common/rollout.rs @@ -202,6 +202,7 @@ fn create_fake_rollout_with_source_and_parent_thread_id( dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: Default::default(), multi_agent_version: None, context_window: None, }; @@ -291,6 +292,7 @@ pub fn create_fake_rollout_with_text_elements( dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: Default::default(), multi_agent_version: None, context_window: None, }; diff --git a/codex-rs/app-server/tests/suite/conversation_summary.rs b/codex-rs/app-server/tests/suite/conversation_summary.rs index 63ad473261..64a0e62249 100644 --- a/codex-rs/app-server/tests/suite/conversation_summary.rs +++ b/codex-rs/app-server/tests/suite/conversation_summary.rs @@ -133,6 +133,7 @@ async fn get_conversation_summary_by_thread_id_reads_pathless_store_thread() -> dynamic_tools: Vec::new(), selected_capability_roots: Vec::new(), multi_agent_version: None, + history_mode: Default::default(), initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { cwd: None, diff --git a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs index 7e0db1b5b3..2002b08a88 100644 --- a/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs +++ b/codex-rs/app-server/tests/suite/v2/remote_thread_store.rs @@ -160,6 +160,7 @@ async fn thread_delete_with_non_local_thread_store_does_not_create_local_persist dynamic_tools: Vec::new(), selected_capability_roots: Vec::new(), multi_agent_version: None, + history_mode: Default::default(), initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { cwd: Some(codex_home.path().to_path_buf()), diff --git a/codex-rs/app-server/tests/suite/v2/skills_list.rs b/codex-rs/app-server/tests/suite/v2/skills_list.rs index 64b7fe0504..f0bee7f27d 100644 --- a/codex-rs/app-server/tests/suite/v2/skills_list.rs +++ b/codex-rs/app-server/tests/suite/v2/skills_list.rs @@ -897,6 +897,7 @@ async fn skills_changed_notification_is_emitted_after_skill_change() -> Result<( personality: None, multi_agent_mode: None, ephemeral: None, + history_mode: None, session_start_source: None, thread_source: None, dynamic_tools: None, diff --git a/codex-rs/app-server/tests/suite/v2/thread_read.rs b/codex-rs/app-server/tests/suite/v2/thread_read.rs index 9f67761757..22adba37c8 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_read.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_read.rs @@ -18,6 +18,7 @@ use codex_app_server_protocol::SessionSource; use codex_app_server_protocol::SortDirection; use codex_app_server_protocol::ThreadForkParams; use codex_app_server_protocol::ThreadForkResponse; +use codex_app_server_protocol::ThreadHistoryMode; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadItemsListParams; use codex_app_server_protocol::ThreadListParams; @@ -196,6 +197,121 @@ async fn thread_read_can_include_turns() -> Result<()> { Ok(()) } +#[tokio::test] +async fn paginated_stored_thread_allows_metadata_discovery_and_rejects_legacy_history_paths() +-> 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_text_elements( + codex_home.path(), + "2025-01-05T12-00-00", + "2025-01-05T12:00:00Z", + "Saved user message", + Vec::new(), + Some("mock_provider"), + /*git_info*/ None, + )?; + set_rollout_history_mode( + rollout_path(codex_home.path(), "2025-01-05T12-00-00", &conversation_id).as_path(), + ThreadHistoryMode::Paginated, + )?; + + let mut mcp = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let read_id = mcp + .send_thread_read_request(ThreadReadParams { + thread_id: conversation_id.clone(), + include_turns: false, + }) + .await?; + let read_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(read_id)), + ) + .await??; + let ThreadReadResponse { thread } = to_response::(read_resp)?; + assert_eq!(thread.history_mode, ThreadHistoryMode::Paginated); + assert!(thread.turns.is_empty()); + + let list_id = mcp + .send_thread_list_request(ThreadListParams { + cursor: None, + limit: Some(50), + sort_key: None, + sort_direction: None, + model_providers: Some(vec!["mock_provider".to_string()]), + source_kinds: None, + archived: None, + cwd: None, + use_state_db_only: false, + search_term: None, + parent_thread_id: None, + ancestor_thread_id: None, + }) + .await?; + let list_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(list_id)), + ) + .await??; + let ThreadListResponse { data, .. } = to_response::(list_resp)?; + let listed = data + .iter() + .find(|thread| thread.id == conversation_id) + .expect("thread/list should include paginated thread"); + assert_eq!(listed.history_mode, ThreadHistoryMode::Paginated); + + let read_id = mcp + .send_thread_read_request(ThreadReadParams { + thread_id: conversation_id.clone(), + include_turns: true, + }) + .await?; + assert_paginated_threads_unsupported( + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(read_id)), + ) + .await??, + ); + + let turns_list_id = mcp + .send_thread_turns_list_request(ThreadTurnsListParams { + thread_id: conversation_id.clone(), + cursor: None, + limit: None, + sort_direction: None, + items_view: None, + }) + .await?; + assert_paginated_threads_unsupported( + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(turns_list_id)), + ) + .await??, + ); + + let resume_id = mcp + .send_thread_resume_request(ThreadResumeParams { + thread_id: conversation_id, + ..Default::default() + }) + .await?; + assert_paginated_threads_unsupported( + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(resume_id)), + ) + .await??, + ); + + Ok(()) +} + #[tokio::test] async fn thread_turns_list_can_page_backward_and_forward() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; @@ -1371,6 +1487,7 @@ async fn seed_pathless_store_thread( dynamic_tools: Vec::new(), selected_capability_roots: Vec::new(), multi_agent_version: None, + history_mode: Default::default(), initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { cwd: None, @@ -1411,6 +1528,26 @@ fn store_history_items() -> Vec { ))] } +fn set_rollout_history_mode(path: &Path, history_mode: ThreadHistoryMode) -> Result<()> { + let mut lines = std::fs::read_to_string(path)? + .lines() + .map(serde_json::from_str::) + .collect::, _>>()?; + lines[0]["payload"]["history_mode"] = serde_json::to_value(history_mode)?; + let contents = lines + .into_iter() + .map(|line| line.to_string()) + .collect::>() + .join("\n"); + std::fs::write(path, format!("{contents}\n"))?; + Ok(()) +} + +fn assert_paginated_threads_unsupported(err: JSONRPCError) { + assert_eq!(err.error.code, -32601); + assert_eq!(err.error.message, "paginated_threads is not supported yet"); +} + fn create_config_toml_with_thread_store(codex_home: &Path, store_id: &str) -> std::io::Result<()> { let config_toml = codex_home.join("config.toml"); std::fs::write( 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 67dd249217..a6e851f1d7 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -2168,6 +2168,7 @@ stream_max_retries = 0 dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: Default::default(), multi_agent_version: None, context_window: None, }; diff --git a/codex-rs/app-server/tests/suite/v2/thread_start.rs b/codex-rs/app-server/tests/suite/v2/thread_start.rs index ac4b255de7..f717dd0014 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_start.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_start.rs @@ -15,6 +15,7 @@ use codex_app_server_protocol::McpServerStatusUpdatedNotification; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::SandboxMode; use codex_app_server_protocol::ServerNotification; +use codex_app_server_protocol::ThreadHistoryMode; use codex_app_server_protocol::ThreadSource; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; @@ -248,6 +249,11 @@ async fn thread_start_creates_thread_and_emits_started() -> Result<()> { Some(false), "new persistent threads should serialize `ephemeral: false`" ); + assert_eq!( + thread_json.get("historyMode").and_then(Value::as_str), + Some("legacy"), + "new threads should serialize `historyMode: legacy`" + ); assert_eq!( thread_json.get("threadSource").and_then(Value::as_str), Some("user"), @@ -308,6 +314,50 @@ async fn thread_start_creates_thread_and_emits_started() -> Result<()> { Ok(()) } +#[tokio::test] +async fn thread_start_history_mode_accepts_legacy_and_rejects_paginated() -> Result<()> { + let server = create_mock_responses_server_repeating_assistant("Done").await; + let codex_home = TempDir::new()?; + create_config_toml_without_approval_policy(codex_home.path(), &server.uri())?; + + let mut mcp = TestAppServer::new_with_auto_env(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp + .send_thread_start_request_with_auto_env(ThreadStartParams { + history_mode: Some(ThreadHistoryMode::Legacy), + ..Default::default() + }) + .await?; + let response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response(response)?; + + assert_eq!(thread.history_mode, ThreadHistoryMode::Legacy); + + let request_id = mcp + .send_thread_start_request_with_auto_env(ThreadStartParams { + history_mode: Some(ThreadHistoryMode::Paginated), + ..Default::default() + }) + .await?; + let error: JSONRPCError = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_error_message(RequestId::Integer(request_id)), + ) + .await??; + + assert_eq!(error.error.code, -32601); + assert_eq!( + error.error.message, + "paginated_threads is not supported yet" + ); + Ok(()) +} + #[tokio::test] async fn thread_start_accepts_absolute_runtime_workspace_roots() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; diff --git a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs index 231ede47b9..07999e1650 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_unarchive.rs @@ -220,6 +220,7 @@ async fn thread_unarchive_preserves_pathless_store_metadata() -> Result<()> { dynamic_tools: Vec::new(), selected_capability_roots: Vec::new(), multi_agent_version: None, + history_mode: Default::default(), initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { cwd: None, diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index acb7af3779..9863bfb7e2 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -1464,6 +1464,7 @@ async fn spawn_agent_fork_last_n_turns_drops_parent_startup_prefix_when_under_li config: harness.config.clone(), allow_provider_model_fallback: false, initial_history: InitialHistory::New, + history_mode: None, session_source: None, thread_source: None, dynamic_tools: Vec::new(), @@ -2253,6 +2254,7 @@ async fn spawn_thread_subagents_persist_parent_originator_across_new_and_truncat config: harness.config.clone(), allow_provider_model_fallback: false, initial_history: InitialHistory::New, + history_mode: None, session_source: None, thread_source: None, dynamic_tools: Vec::new(), diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index d25af57939..a671f0dd74 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -108,6 +108,7 @@ pub(crate) async fn run_codex_thread_interactive( code_mode_session_provider: parent_session.services.code_mode_service.session_provider(), extensions: Arc::clone(&parent_session.services.extensions), conversation_history, + requested_history_mode: None, session_source: SessionSource::SubAgent(subagent_source.clone()), forked_from_thread_id, parent_thread_id: Some(parent_session.thread_id), diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 90555db911..4a562092f1 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -29,6 +29,7 @@ use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionConfiguredEvent; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::Submission; +use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadMemoryMode; use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TokenUsageInfo; @@ -72,6 +73,7 @@ pub struct ThreadConfigSnapshot { pub personality: Option, pub collaboration_mode: CollaborationMode, pub session_source: SessionSource, + pub history_mode: ThreadHistoryMode, pub forked_from_thread_id: Option, pub parent_thread_id: Option, pub thread_source: Option, diff --git a/codex-rs/core/src/personality_migration_tests.rs b/codex-rs/core/src/personality_migration_tests.rs index 396d065644..80f5b3c020 100644 --- a/codex-rs/core/src/personality_migration_tests.rs +++ b/codex-rs/core/src/personality_migration_tests.rs @@ -61,6 +61,7 @@ async fn write_rollout_with_user_event(dir: &Path, thread_id: ThreadId) -> io::R dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: Default::default(), multi_agent_version: None, context_window: None, }, diff --git a/codex-rs/core/src/realtime_context_tests.rs b/codex-rs/core/src/realtime_context_tests.rs index e74d5b3e06..23f0b0c03a 100644 --- a/codex-rs/core/src/realtime_context_tests.rs +++ b/codex-rs/core/src/realtime_context_tests.rs @@ -55,6 +55,7 @@ fn stored_thread(cwd: &str, title: &str, first_user_message: &str) -> StoredThre cwd: PathBuf::from(cwd), cli_version: "test".to_string(), source: SessionSource::Cli, + history_mode: Default::default(), thread_source: None, agent_nickname: None, agent_role: None, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 6a21a6be63..f45616eed8 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -122,6 +122,7 @@ use codex_protocol::protocol::ReviewRequest; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; +use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnContextItem; @@ -419,6 +420,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) code_mode_session_provider: Arc, pub(crate) extensions: Arc>, pub(crate) conversation_history: InitialHistory, + pub(crate) requested_history_mode: Option, pub(crate) session_source: SessionSource, pub(crate) forked_from_thread_id: Option, pub(crate) parent_thread_id: Option, @@ -510,6 +512,7 @@ impl Codex { code_mode_session_provider, extensions, conversation_history, + requested_history_mode, session_source, forked_from_thread_id, parent_thread_id, @@ -600,6 +603,9 @@ impl Codex { .await; let multi_agent_version = resolve_multi_agent_version(&conversation_history, inherited_multi_agent_version); + let history_mode = conversation_history.get_history_mode( + requested_history_mode.unwrap_or_else(|| thread_store.default_history_mode()), + ); config .validate_multi_agent_v2_config() .map_err(|err| CodexErr::InvalidRequest(err.to_string()))?; @@ -655,6 +661,7 @@ impl Codex { app_server_client_name: None, app_server_client_version: None, session_source, + history_mode, forked_from_thread_id, parent_thread_id, thread_source, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index a9cc89db3a..5e7046b148 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -16,6 +16,7 @@ use codex_protocol::config_types::ServiceTier; use codex_protocol::permissions::FileSystemPath; use codex_protocol::permissions::FileSystemSpecialPath; use codex_protocol::protocol::MultiAgentVersion; +use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TurnEnvironmentSelections; use std::sync::OnceLock; @@ -96,6 +97,8 @@ pub(crate) struct SessionConfiguration { pub(super) app_server_client_version: Option, /// Source of the session (cli, vscode, exec, mcp, ...) pub(super) session_source: SessionSource, + /// Persisted thread history contract selected when this thread was created. + pub(super) history_mode: ThreadHistoryMode, /// Immediate history source copied into this thread, when this thread was forked. pub(super) forked_from_thread_id: Option, /// Immediate control/spawn parent for this thread, when it has one. @@ -192,6 +195,7 @@ impl SessionConfiguration { personality: self.personality, collaboration_mode: self.collaboration_mode.clone(), session_source: self.session_source.clone(), + history_mode: self.history_mode, forked_from_thread_id: self.forked_from_thread_id, parent_thread_id: self.parent_thread_id, thread_source: self.thread_source.clone(), @@ -599,6 +603,7 @@ impl Session { dynamic_tools: session_configuration.dynamic_tools.clone(), selected_capability_roots: selected_capability_roots.clone(), multi_agent_version: initial_multi_agent_version, + history_mode: session_configuration.history_mode, initial_window_id: initial_auto_compact_window_ids .window_id .to_string(), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 09fdcf2245..93519cda9d 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -3686,6 +3686,7 @@ async fn set_rate_limits_retains_previous_credits() { app_server_client_name: None, app_server_client_version: None, session_source: SessionSource::Exec, + history_mode: Default::default(), forked_from_thread_id: None, parent_thread_id: None, thread_source: None, @@ -3792,6 +3793,7 @@ async fn set_rate_limits_updates_plan_type_when_present() { app_server_client_name: None, app_server_client_version: None, session_source: SessionSource::Exec, + history_mode: Default::default(), forked_from_thread_id: None, parent_thread_id: None, thread_source: None, @@ -4050,6 +4052,7 @@ async fn attach_thread_persistence(session: &mut Session) -> PathBuf { dynamic_tools: Vec::new(), selected_capability_roots: Vec::new(), multi_agent_version: None, + history_mode: Default::default(), initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { cwd: Some(config.cwd.to_path_buf()), @@ -4323,6 +4326,7 @@ pub(crate) async fn make_session_configuration_for_tests() -> SessionConfigurati app_server_client_name: None, app_server_client_version: None, session_source: SessionSource::Exec, + history_mode: Default::default(), forked_from_thread_id: None, parent_thread_id: None, thread_source: None, @@ -5193,6 +5197,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_packaged_zsh() { app_server_client_name: None, app_server_client_version: None, session_source: SessionSource::Exec, + history_mode: Default::default(), forked_from_thread_id: None, parent_thread_id: None, thread_source: None, @@ -5324,6 +5329,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { app_server_client_name: None, app_server_client_version: None, session_source: SessionSource::Exec, + history_mode: Default::default(), forked_from_thread_id: None, parent_thread_id: None, thread_source: None, @@ -5572,6 +5578,7 @@ async fn make_session_with_config_and_rx( app_server_client_name: None, app_server_client_version: None, session_source: SessionSource::Exec, + history_mode: Default::default(), forked_from_thread_id: None, parent_thread_id: None, thread_source: None, @@ -5679,6 +5686,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( app_server_client_name: None, app_server_client_version: None, session_source: session_source.clone(), + history_mode: Default::default(), forked_from_thread_id: None, parent_thread_id: None, thread_source: None, @@ -6931,6 +6939,7 @@ async fn shutdown_complete_does_not_append_to_thread_store_after_shutdown() { dynamic_tools: Vec::new(), selected_capability_roots: Vec::new(), multi_agent_version: None, + history_mode: Default::default(), initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { cwd: Some(config.cwd.to_path_buf()), @@ -7404,6 +7413,7 @@ where app_server_client_name: None, app_server_client_version: None, session_source: SessionSource::Exec, + history_mode: Default::default(), forked_from_thread_id: None, parent_thread_id: None, thread_source: None, diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index dc17af5a16..d13239d810 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -730,6 +730,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { code_mode_session_provider: Arc::new(codex_code_mode::InProcessCodeModeSessionProvider), extensions: codex_extension_api::empty_extension_registry(), conversation_history: InitialHistory::New, + requested_history_mode: None, session_source: SessionSource::SubAgent(SubAgentSource::Other( GUARDIAN_REVIEWER_NAME.to_string(), )), diff --git a/codex-rs/core/src/session_rollout_init_error.rs b/codex-rs/core/src/session_rollout_init_error.rs index 57e1e8e57a..db71dfcd37 100644 --- a/codex-rs/core/src/session_rollout_init_error.rs +++ b/codex-rs/core/src/session_rollout_init_error.rs @@ -3,8 +3,16 @@ use std::path::Path; use crate::rollout::SESSIONS_SUBDIR; use codex_protocol::error::CodexErr; +use codex_thread_store::ThreadStoreError; pub(crate) fn map_session_init_error(err: &anyhow::Error, codex_home: &Path) -> CodexErr { + if let Some(ThreadStoreError::Unsupported { operation }) = err + .chain() + .find_map(|cause| cause.downcast_ref::()) + { + return CodexErr::UnsupportedOperation(format!("{operation} is not supported yet")); + } + if let Some(mapped) = err .chain() .filter_map(|cause| cause.downcast_ref::()) diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 2ca21c15ab..4ea86cb33d 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -56,6 +56,7 @@ use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionConfiguredEvent; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; +use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; @@ -187,6 +188,7 @@ pub struct StartThreadOptions { pub config: Config, pub allow_provider_model_fallback: bool, pub initial_history: InitialHistory, + pub history_mode: Option, pub session_source: Option, pub thread_source: Option, pub dynamic_tools: Vec, @@ -647,6 +649,7 @@ impl ThreadManager { config, allow_provider_model_fallback: false, initial_history: InitialHistory::New, + history_mode: None, session_source: None, thread_source: None, dynamic_tools, @@ -682,6 +685,7 @@ impl ThreadManager { Box::pin(self.state.spawn_thread_with_source( options.config, options.initial_history, + options.history_mode, options.allow_provider_model_fallback, Arc::clone(&self.state.auth_manager), agent_control, @@ -778,6 +782,7 @@ impl ThreadManager { Box::pin(self.state.spawn_thread_with_source( config, initial_history, + /*history_mode*/ None, /*allow_provider_model_fallback*/ false, auth_manager, agent_control, @@ -848,6 +853,7 @@ impl ThreadManager { Box::pin(self.state.spawn_thread_with_source( config, initial_history, + /*history_mode*/ None, /*allow_provider_model_fallback*/ false, auth_manager, agent_control, @@ -1349,6 +1355,7 @@ impl ThreadManagerState { Box::pin(self.spawn_thread_with_source( config, InitialHistory::New, + /*history_mode*/ None, /*allow_provider_model_fallback*/ false, Arc::clone(&self.auth_manager), agent_control, @@ -1388,6 +1395,7 @@ impl ThreadManagerState { Box::pin(self.spawn_thread_with_source( config, initial_history, + /*history_mode*/ None, /*allow_provider_model_fallback*/ false, Arc::clone(&self.auth_manager), agent_control, @@ -1429,6 +1437,7 @@ impl ThreadManagerState { Box::pin(self.spawn_thread_with_source( config, initial_history, + /*history_mode*/ None, /*allow_provider_model_fallback*/ false, Arc::clone(&self.auth_manager), agent_control, @@ -1471,6 +1480,7 @@ impl ThreadManagerState { Box::pin(self.spawn_thread_with_source( config, initial_history, + /*history_mode*/ None, /*allow_provider_model_fallback*/ false, auth_manager, agent_control, @@ -1496,6 +1506,7 @@ impl ThreadManagerState { &self, config: Config, initial_history: InitialHistory, + history_mode: Option, allow_provider_model_fallback: bool, auth_manager: Arc, agent_control: AgentControl, @@ -1575,6 +1586,7 @@ impl ThreadManagerState { code_mode_session_provider: Arc::clone(&self.code_mode_session_provider), extensions: Arc::clone(&self.extensions), conversation_history: initial_history, + requested_history_mode: history_mode, session_source, forked_from_thread_id, parent_thread_id, diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index b33d95ac33..4f0da0c9c7 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -472,6 +472,7 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { config, allow_provider_model_fallback: false, initial_history: InitialHistory::New, + history_mode: None, session_source: Some(SessionSource::Internal( InternalSessionSource::MemoryConsolidation, )), @@ -615,6 +616,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() config: config.clone(), allow_provider_model_fallback: false, initial_history: InitialHistory::New, + history_mode: None, session_source: None, thread_source: None, dynamic_tools: Vec::new(), @@ -631,6 +633,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() config: config.clone(), allow_provider_model_fallback: false, initial_history: InitialHistory::New, + history_mode: None, session_source: None, thread_source: None, dynamic_tools: Vec::new(), @@ -739,6 +742,7 @@ async fn selected_capability_roots_round_trip_through_fork() { git: None, }, )]), + history_mode: None, session_source: None, thread_source: None, dynamic_tools: Vec::new(), @@ -810,6 +814,7 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { config: source_config, allow_provider_model_fallback: false, initial_history: InitialHistory::New, + history_mode: None, session_source: None, thread_source: None, dynamic_tools: Vec::new(), @@ -1094,6 +1099,7 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() { config: config.clone(), allow_provider_model_fallback: false, initial_history: InitialHistory::New, + history_mode: None, session_source: None, thread_source: Some(ThreadSource::User), dynamic_tools: Vec::new(), diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 6b11267774..e0e82c37fd 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -660,6 +660,7 @@ impl TestCodexBuilder { config: config.clone(), allow_provider_model_fallback: false, initial_history: InitialHistory::New, + history_mode: None, session_source: None, thread_source: None, dynamic_tools: Vec::new(), diff --git a/codex-rs/core/tests/suite/agents_md.rs b/codex-rs/core/tests/suite/agents_md.rs index 4e0ca794e3..7c73b90603 100644 --- a/codex-rs/core/tests/suite/agents_md.rs +++ b/codex-rs/core/tests/suite/agents_md.rs @@ -488,6 +488,7 @@ async fn loads_user_instructions_without_a_primary_environment() -> Result<()> { config: test.config.clone(), allow_provider_model_fallback: false, initial_history: InitialHistory::New, + history_mode: None, session_source: None, thread_source: None, dynamic_tools: Vec::new(), @@ -694,6 +695,7 @@ async fn multi_environment_thread_loads_every_project_and_keeps_creation_snapsho config: test.config.clone(), allow_provider_model_fallback: false, initial_history: InitialHistory::New, + history_mode: None, session_source: None, thread_source: None, dynamic_tools: Vec::new(), diff --git a/codex-rs/core/tests/suite/personality_migration.rs b/codex-rs/core/tests/suite/personality_migration.rs index 3dd0f3c016..0ca3e1b2a9 100644 --- a/codex-rs/core/tests/suite/personality_migration.rs +++ b/codex-rs/core/tests/suite/personality_migration.rs @@ -77,6 +77,7 @@ async fn write_rollout_with_user_event(dir: &Path, thread_id: ThreadId) -> io::R dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: Default::default(), multi_agent_version: None, context_window: None, }, @@ -130,6 +131,7 @@ async fn write_rollout_with_meta_only(dir: &Path, thread_id: ThreadId) -> io::Re dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: Default::default(), multi_agent_version: None, context_window: None, }, diff --git a/codex-rs/core/tests/suite/sqlite_state.rs b/codex-rs/core/tests/suite/sqlite_state.rs index 1c2f08262a..726d74ebfd 100644 --- a/codex-rs/core/tests/suite/sqlite_state.rs +++ b/codex-rs/core/tests/suite/sqlite_state.rs @@ -374,6 +374,7 @@ async fn backfill_scans_existing_rollouts() -> Result<()> { dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: Default::default(), multi_agent_version: None, context_window: None, }, diff --git a/codex-rs/core/tests/suite/subagent_notifications.rs b/codex-rs/core/tests/suite/subagent_notifications.rs index 62a367a229..951fb92f17 100644 --- a/codex-rs/core/tests/suite/subagent_notifications.rs +++ b/codex-rs/core/tests/suite/subagent_notifications.rs @@ -763,6 +763,7 @@ async fn subagent_stop_replaces_stop_and_skips_internal_subagents() -> Result<() config: test.config.clone(), allow_provider_model_fallback: false, initial_history: InitialHistory::New, + history_mode: None, session_source: Some(SessionSource::SubAgent(SubAgentSource::Review)), thread_source: None, dynamic_tools: Vec::new(), diff --git a/codex-rs/exec/src/lib_tests.rs b/codex-rs/exec/src/lib_tests.rs index 9f345eb334..97ded50525 100644 --- a/codex-rs/exec/src/lib_tests.rs +++ b/codex-rs/exec/src/lib_tests.rs @@ -340,6 +340,7 @@ fn turn_items_for_thread_returns_matching_turn_items() { parent_thread_id: None, preview: String::new(), ephemeral: false, + history_mode: Default::default(), model_provider: "openai".to_string(), created_at: 0, updated_at: 0, @@ -760,6 +761,7 @@ fn sample_thread_start_response() -> ThreadStartResponse { parent_thread_id: None, preview: String::new(), ephemeral: false, + history_mode: Default::default(), model_provider: "openai".to_string(), created_at: 0, updated_at: 0, diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs index 73744e0a89..6c7ad3978f 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -331,6 +331,7 @@ impl MemoryStartupContext { config, allow_provider_model_fallback: false, initial_history: InitialHistory::New, + history_mode: None, session_source: Some(SessionSource::Internal( InternalSessionSource::MemoryConsolidation, )), diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 491d64aa58..96a793bcd3 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -678,6 +678,36 @@ pub enum ThreadMemoryMode { Disabled, } +#[derive(Serialize, Deserialize, Clone, Copy, Debug, Default, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "lowercase")] +#[ts(rename_all = "lowercase")] +pub enum ThreadHistoryMode { + #[default] + Legacy, + Paginated, +} + +impl ThreadHistoryMode { + pub const fn as_str(self) -> &'static str { + match self { + Self::Legacy => "legacy", + Self::Paginated => "paginated", + } + } +} + +impl FromStr for ThreadHistoryMode { + type Err = String; + + fn from_str(value: &str) -> Result { + match value { + "legacy" => Ok(Self::Legacy), + "paginated" => Ok(Self::Paginated), + _ => Err(format!("unknown thread history mode `{value}`")), + } + } +} + impl From> for Op { fn from(value: Vec) -> Self { Op::UserInput { @@ -2607,6 +2637,18 @@ impl InitialHistory { } } + pub fn get_history_mode(&self, default_history_mode: ThreadHistoryMode) -> ThreadHistoryMode { + match self { + InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => { + default_history_mode + } + InitialHistory::Resumed(_) => self + .get_resumed_session_meta() + .map(|meta| meta.history_mode) + .unwrap_or(default_history_mode), + } + } + pub fn get_latest_effective_multi_agent_mode(&self) -> Option { let items = match self { InitialHistory::New | InitialHistory::Cleared => return None, @@ -3033,6 +3075,8 @@ pub struct SessionMeta { pub selected_capability_roots: Vec, #[serde(skip_serializing_if = "Option::is_none")] pub memory_mode: Option, + #[serde(default)] + pub history_mode: ThreadHistoryMode, #[serde(skip_serializing_if = "Option::is_none")] pub multi_agent_version: Option, /// Initial context-window identity for consumers that tail rollout JSONL before compaction. @@ -3062,6 +3106,7 @@ impl Default for SessionMeta { dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: ThreadHistoryMode::default(), multi_agent_version: None, context_window: None, } @@ -5542,6 +5587,70 @@ mod tests { Ok(()) } + #[test] + fn session_meta_defaults_legacy_history_mode() -> Result<()> { + let session_meta: SessionMeta = serde_json::from_value(json!({ + "session_id": "00000000-0000-0000-0000-000000000001", + "id": "00000000-0000-0000-0000-000000000001", + "timestamp": "2026-01-01T00:00:00Z", + "cwd": "/tmp", + "originator": "codex", + "cli_version": "0.0.0", + "model_provider": null, + "base_instructions": null + }))?; + + assert_eq!(session_meta.history_mode, ThreadHistoryMode::Legacy); + let serialized = serde_json::to_value(&session_meta)?; + assert_eq!(serialized["history_mode"], json!("legacy")); + let mut unknown = serialized; + unknown["history_mode"] = json!("future"); + assert!(serde_json::from_value::(unknown).is_err()); + Ok(()) + } + + #[test] + fn resumed_history_uses_persisted_history_mode() -> Result<()> { + let thread_id = ThreadId::from_string("00000000-0000-0000-0000-000000000001")?; + let session_meta = RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + session_id: thread_id.into(), + id: thread_id, + history_mode: ThreadHistoryMode::Paginated, + ..SessionMeta::default() + }, + git: None, + }); + let history = InitialHistory::Resumed(ResumedHistory { + conversation_id: thread_id, + history: Arc::new(vec![session_meta.clone()]), + rollout_path: None, + }); + + assert_eq!( + history.get_history_mode(ThreadHistoryMode::Legacy), + ThreadHistoryMode::Paginated + ); + assert_eq!( + InitialHistory::Forked(vec![session_meta]).get_history_mode(ThreadHistoryMode::Legacy), + ThreadHistoryMode::Legacy + ); + assert_eq!( + InitialHistory::New.get_history_mode(ThreadHistoryMode::Paginated), + ThreadHistoryMode::Paginated + ); + assert_eq!( + InitialHistory::Resumed(ResumedHistory { + conversation_id: thread_id, + history: Arc::new(Vec::new()), + rollout_path: None, + }) + .get_history_mode(ThreadHistoryMode::Paginated), + ThreadHistoryMode::Paginated + ); + Ok(()) + } + #[test] fn turn_context_item_deserializes_without_network() -> Result<()> { let item: TurnContextItem = serde_json::from_value(json!({ diff --git a/codex-rs/rollout/src/compression_tests.rs b/codex-rs/rollout/src/compression_tests.rs index 733e6cc5fb..46c90d0120 100644 --- a/codex-rs/rollout/src/compression_tests.rs +++ b/codex-rs/rollout/src/compression_tests.rs @@ -474,6 +474,7 @@ fn write_rollout(path: &std::path::Path, thread_id: ThreadId, message: &str) -> dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: Default::default(), multi_agent_version: None, context_window: None, }, diff --git a/codex-rs/rollout/src/list.rs b/codex-rs/rollout/src/list.rs index b02b8c3e49..1154d36eee 100644 --- a/codex-rs/rollout/src/list.rs +++ b/codex-rs/rollout/src/list.rs @@ -26,7 +26,9 @@ use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::RolloutLine; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::USER_MESSAGE_BEGIN; +use serde_json::Value; /// Returned page of thread (thread) summaries. #[derive(Debug, Default, PartialEq)] @@ -62,6 +64,8 @@ pub struct ThreadItem { pub git_origin_url: Option, /// Session source from session metadata. pub source: Option, + /// Persisted thread history contract selected when this thread was created. + pub history_mode: ThreadHistoryMode, /// Immediate control/spawn parent thread id from session metadata. pub parent_thread_id: Option, /// Random unique nickname from session metadata for AgentControl-spawned sub-agents. @@ -99,6 +103,7 @@ struct HeadTailSummary { git_sha: Option, git_origin_url: Option, source: Option, + history_mode: ThreadHistoryMode, parent_thread_id: Option, agent_nickname: Option, agent_role: Option, @@ -806,6 +811,7 @@ async fn build_thread_item( git_sha, git_origin_url, source, + history_mode, parent_thread_id, agent_nickname, agent_role, @@ -828,6 +834,7 @@ async fn build_thread_item( git_sha, git_origin_url, source, + history_mode, parent_thread_id, agent_nickname, agent_role, @@ -1117,12 +1124,26 @@ async fn read_head_summary(path: &Path, head_limit: usize) -> io::Result = serde_json::from_str(trimmed); - let Ok(rollout_line) = parsed else { continue }; + let rollout_line = match parsed { + Ok(rollout_line) => rollout_line, + Err(_) => { + if !summary.saw_session_meta + && let Ok(value) = serde_json::from_str::(trimmed) + { + // The first SessionMeta belongs to this rollout. Later SessionMeta lines can + // be copied from fork history, so only an unknown mode before the first parsed + // SessionMeta should make this thread unreadable. + crate::recorder::reject_unknown_thread_history_mode(&value)?; + } + continue; + } + }; match rollout_line.item { RolloutItem::SessionMeta(session_meta_line) => { if !summary.saw_session_meta { summary.source = Some(session_meta_line.meta.source.clone()); + summary.history_mode = session_meta_line.meta.history_mode; summary.parent_thread_id = session_meta_line.meta.parent_thread_id; summary.agent_nickname = session_meta_line.meta.agent_nickname.clone(); summary.agent_role = session_meta_line.meta.agent_role.clone(); diff --git a/codex-rs/rollout/src/metadata.rs b/codex-rs/rollout/src/metadata.rs index d6ed809136..127a8af223 100644 --- a/codex-rs/rollout/src/metadata.rs +++ b/codex-rs/rollout/src/metadata.rs @@ -45,6 +45,7 @@ pub(crate) fn builder_from_session_meta( created_at, session_meta.meta.source.clone(), ); + builder.history_mode = session_meta.meta.history_mode; builder.model_provider = session_meta.meta.model_provider.clone(); builder.agent_nickname = session_meta.meta.agent_nickname.clone(); builder.agent_role = session_meta.meta.agent_role.clone(); diff --git a/codex-rs/rollout/src/metadata_tests.rs b/codex-rs/rollout/src/metadata_tests.rs index c0aea248c5..2e1ea98f3a 100644 --- a/codex-rs/rollout/src/metadata_tests.rs +++ b/codex-rs/rollout/src/metadata_tests.rs @@ -13,6 +13,7 @@ use codex_protocol::protocol::RolloutLine; use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadHistoryMode; use codex_state::BackfillStatus; use codex_state::ThreadMetadataBuilder; use pretty_assertions::assert_eq; @@ -51,6 +52,7 @@ async fn extract_metadata_from_rollout_uses_session_meta() { dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: ThreadHistoryMode::Paginated, multi_agent_version: None, context_window: None, }; @@ -81,6 +83,41 @@ async fn extract_metadata_from_rollout_uses_session_meta() { assert_eq!(outcome.parse_errors, 0); } +#[tokio::test] +async fn extract_metadata_from_rollout_rejects_unknown_history_mode() { + let dir = tempdir().expect("tempdir"); + let uuid = Uuid::new_v4(); + let id = ThreadId::from_string(&uuid.to_string()).expect("thread id"); + let path = dir + .path() + .join(format!("rollout-2026-01-27T12-34-56-{uuid}.jsonl")); + let mut rollout_line = serde_json::to_value(RolloutLine { + timestamp: "2026-01-27T12:34:56Z".to_string(), + item: RolloutItem::SessionMeta(SessionMetaLine { + meta: SessionMeta { + session_id: id.into(), + id, + timestamp: "2026-01-27T12:34:56Z".to_string(), + cwd: dir.path().to_path_buf(), + originator: "cli".to_string(), + cli_version: "0.0.0".to_string(), + ..SessionMeta::default() + }, + git: None, + }), + }) + .expect("serialize rollout line"); + rollout_line["payload"]["history_mode"] = serde_json::json!("future"); + let mut file = File::create(&path).expect("create rollout"); + writeln!(file, "{rollout_line}").expect("write rollout"); + + assert!( + extract_metadata_from_rollout(&path, "openai") + .await + .is_err() + ); +} + #[tokio::test] async fn extract_metadata_from_rollout_returns_latest_memory_mode() { let dir = tempdir().expect("tempdir"); @@ -109,6 +146,7 @@ async fn extract_metadata_from_rollout_returns_latest_memory_mode() { dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: Default::default(), multi_agent_version: None, context_window: None, }; @@ -379,6 +417,7 @@ fn write_rollout_in_sessions_with_cwd( dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: Default::default(), multi_agent_version: None, context_window: None, }; diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 58d5c76662..87e97997a2 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -61,6 +61,7 @@ use codex_protocol::protocol::SessionContextWindow; use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadSource; use codex_state::StateRuntime; use codex_utils_path as path_utils; @@ -94,6 +95,7 @@ pub enum RolloutRecorderParams { dynamic_tools: Vec, selected_capability_roots: Vec, multi_agent_version: Option, + history_mode: ThreadHistoryMode, initial_window_id: Option, }, Resume { @@ -186,6 +188,7 @@ impl RolloutRecorderParams { dynamic_tools, selected_capability_roots: Vec::new(), multi_agent_version: None, + history_mode: Default::default(), initial_window_id: None, } } @@ -225,6 +228,16 @@ impl RolloutRecorderParams { self } + pub fn with_history_mode(mut self, history_mode: ThreadHistoryMode) -> Self { + if let Self::Create { + history_mode: mode, .. + } = &mut self + { + *mode = history_mode; + } + self + } + pub fn with_initial_window_id(mut self, initial_window_id: String) -> Self { if let Self::Create { initial_window_id: window_id, @@ -751,6 +764,7 @@ impl RolloutRecorder { dynamic_tools, selected_capability_roots, multi_agent_version, + history_mode, initial_window_id, } => { let log_file_info = precompute_log_file_info(config, conversation_id)?; @@ -789,6 +803,7 @@ impl RolloutRecorder { }, selected_capability_roots, memory_mode: (!config.generate_memories()).then_some("disabled".to_string()), + history_mode, multi_agent_version, context_window: initial_window_id.map(SessionContextWindow::new), }; @@ -956,6 +971,12 @@ impl RolloutRecorder { items.push(item); } Err(e) => { + if thread_id.is_none() { + // The first SessionMeta belongs to this rollout. Later SessionMeta lines + // can be copied from fork history, so only validate unknown history modes + // before we have parsed the rollout's own SessionMeta. + reject_unknown_thread_history_mode(&v)?; + } trace!("failed to parse rollout line: {e}"); parse_errors = parse_errors.saturating_add(1); } @@ -1019,6 +1040,21 @@ impl RolloutRecorder { } } +pub(crate) fn reject_unknown_thread_history_mode(value: &Value) -> std::io::Result<()> { + if value.get("type").and_then(Value::as_str) != Some("session_meta") { + return Ok(()); + } + let Some(history_mode) = value + .get("payload") + .and_then(|payload| payload.get("history_mode")) + else { + return Ok(()); + }; + serde_json::from_value::(history_mode.clone()) + .map(|_| ()) + .map_err(|err| IoError::other(format!("invalid session metadata history_mode: {err}"))) +} + fn strip_legacy_ghost_snapshot_rollout_line(value: &mut Value) -> bool { match value.get("type").and_then(Value::as_str) { Some("response_item") => value @@ -1122,6 +1158,7 @@ fn fill_missing_thread_item_metadata(item: &mut ThreadItem, state_item: ThreadIt git_sha, git_origin_url, source, + history_mode: _, parent_thread_id, agent_nickname, agent_role, @@ -1837,6 +1874,7 @@ fn thread_item_from_state_metadata( .or_else(|_| serde_json::from_value(Value::String(item.source))) .unwrap_or(SessionSource::Unknown), ), + history_mode: item.history_mode, parent_thread_id, agent_nickname: item.agent_nickname, agent_role: item.agent_role, diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index 71ec252c23..7fcaff08ab 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -15,6 +15,7 @@ use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::TurnContextItem; use codex_protocol::protocol::UserMessageEvent; use pretty_assertions::assert_eq; @@ -103,6 +104,7 @@ async fn state_db_init_backfills_before_returning() -> anyhow::Result<()> { dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: Default::default(), multi_agent_version: None, context_window: None, }, @@ -227,6 +229,43 @@ async fn load_rollout_items_defaults_legacy_session_id() -> std::io::Result<()> Ok(()) } +#[tokio::test] +async fn load_rollout_items_ignores_unknown_fork_source_history_mode() -> std::io::Result<()> { + let home = TempDir::new().expect("temp dir"); + let uuid = Uuid::new_v4(); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("thread id"); + let rollout_path = write_session_file(home.path(), "2025-01-03T12-00-00", uuid)?; + let mut file = fs::OpenOptions::new().append(true).open(&rollout_path)?; + let source_uuid = Uuid::new_v4(); + writeln!( + file, + "{}", + serde_json::json!({ + "timestamp": "2025-01-03T12:00:01Z", + "type": "session_meta", + "payload": { + "session_id": source_uuid, + "id": source_uuid, + "timestamp": "2025-01-03T12:00:01Z", + "cwd": ".", + "originator": "test_originator", + "cli_version": "test_version", + "source": "cli", + "model_provider": "test-provider", + "history_mode": "future", + }, + }) + )?; + + let (items, loaded_thread_id, parse_errors) = + RolloutRecorder::load_rollout_items(&rollout_path).await?; + + assert_eq!(loaded_thread_id, Some(thread_id)); + assert_eq!(parse_errors, 1); + assert_eq!(items.len(), 2); + Ok(()) +} + #[tokio::test] async fn load_rollout_items_preserves_legacy_guardian_assessment_lines() -> std::io::Result<()> { let home = TempDir::new().expect("temp dir"); @@ -391,6 +430,7 @@ async fn recorder_materializes_on_flush_with_pending_items() -> std::io::Result< Vec::new(), ) .with_session_id(session_id) + .with_history_mode(ThreadHistoryMode::Paginated) .with_initial_window_id(initial_window_id.clone()), ) .await?; @@ -442,6 +482,7 @@ async fn recorder_materializes_on_flush_with_pending_items() -> std::io::Result< panic!("expected session metadata in rollout"); }; assert_eq!(session_meta.meta.session_id, session_id); + assert_eq!(session_meta.meta.history_mode, ThreadHistoryMode::Paginated); assert_eq!( session_meta .meta @@ -1004,6 +1045,7 @@ fn fill_missing_thread_item_metadata_preserves_identity_and_prefers_state_git_fi git_sha: Some("filesystem-sha".to_string()), git_origin_url: Some("https://example.com/filesystem.git".to_string()), source: None, + history_mode: Default::default(), parent_thread_id: None, agent_nickname: None, agent_role: None, @@ -1023,6 +1065,7 @@ fn fill_missing_thread_item_metadata_preserves_identity_and_prefers_state_git_fi git_sha: Some("state-sha".to_string()), git_origin_url: Some("https://example.com/state.git".to_string()), source: Some(SessionSource::Exec), + history_mode: Default::default(), parent_thread_id: None, agent_nickname: Some("state-agent".to_string()), agent_role: Some("state-role".to_string()), diff --git a/codex-rs/rollout/src/session_index_tests.rs b/codex-rs/rollout/src/session_index_tests.rs index be39fb06ab..6eb1fa0374 100644 --- a/codex-rs/rollout/src/session_index_tests.rs +++ b/codex-rs/rollout/src/session_index_tests.rs @@ -43,6 +43,7 @@ fn write_rollout_with_metadata(path: &Path, thread_id: ThreadId) -> std::io::Res dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: Default::default(), multi_agent_version: None, context_window: None, }, diff --git a/codex-rs/rollout/src/state_db_tests.rs b/codex-rs/rollout/src/state_db_tests.rs index bb341aa85a..18566d887e 100644 --- a/codex-rs/rollout/src/state_db_tests.rs +++ b/codex-rs/rollout/src/state_db_tests.rs @@ -176,6 +176,7 @@ fn write_rollout_with_user_message( dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: Default::default(), multi_agent_version: None, context_window: None, }, diff --git a/codex-rs/rollout/src/tests.rs b/codex-rs/rollout/src/tests.rs index 901461cf56..23f72cb947 100644 --- a/codex-rs/rollout/src/tests.rs +++ b/codex-rs/rollout/src/tests.rs @@ -129,6 +129,41 @@ async fn find_thread_path_falls_back_when_db_path_is_stale() { assert_state_db_rollout_path(home, thread_id, Some(fs_rollout_path.as_path())).await; } +#[tokio::test] +async fn read_thread_item_from_rollout_rejects_unknown_canonical_history_mode() { + let temp = TempDir::new().unwrap(); + let home = temp.path(); + let uuid = Uuid::from_u128(303); + let ts = "2025-01-03T12-00-00"; + write_session_file( + home, + ts, + uuid, + /*num_records*/ 1, + Some(SessionSource::Cli), + ) + .unwrap(); + let path = home.join(format!("sessions/2025/01/03/rollout-{ts}-{uuid}.jsonl")); + let contents = fs::read_to_string(&path).unwrap(); + let mut lines = contents.lines(); + let mut canonical_meta: serde_json::Value = + serde_json::from_str(lines.next().expect("canonical session meta")).unwrap(); + canonical_meta["payload"]["history_mode"] = serde_json::json!("future"); + let user_event = lines.next().expect("user event"); + let mut source_meta = canonical_meta.clone(); + let source_uuid = Uuid::from_u128(304); + source_meta["payload"]["session_id"] = serde_json::json!(source_uuid); + source_meta["payload"]["id"] = serde_json::json!(source_uuid); + source_meta["payload"]["history_mode"] = serde_json::json!("legacy"); + fs::write( + &path, + format!("{canonical_meta}\n{user_event}\n{source_meta}\n"), + ) + .unwrap(); + + assert_eq!(crate::list::read_thread_item_from_rollout(path).await, None); +} + #[tokio::test] async fn find_thread_path_falls_back_when_db_path_points_to_another_thread() { let temp = TempDir::new().unwrap(); @@ -610,6 +645,7 @@ async fn test_list_conversations_latest_first() { git_sha: None, git_origin_url: None, source: Some(SessionSource::VSCode), + history_mode: Default::default(), parent_thread_id: None, agent_nickname: None, agent_role: None, @@ -629,6 +665,7 @@ async fn test_list_conversations_latest_first() { git_sha: None, git_origin_url: None, source: Some(SessionSource::VSCode), + history_mode: Default::default(), parent_thread_id: None, agent_nickname: None, agent_role: None, @@ -648,6 +685,7 @@ async fn test_list_conversations_latest_first() { git_sha: None, git_origin_url: None, source: Some(SessionSource::VSCode), + history_mode: Default::default(), parent_thread_id: None, agent_nickname: None, agent_role: None, @@ -760,6 +798,7 @@ async fn test_pagination_cursor() { git_sha: None, git_origin_url: None, source: Some(SessionSource::VSCode), + history_mode: Default::default(), parent_thread_id: None, agent_nickname: None, agent_role: None, @@ -779,6 +818,7 @@ async fn test_pagination_cursor() { git_sha: None, git_origin_url: None, source: Some(SessionSource::VSCode), + history_mode: Default::default(), parent_thread_id: None, agent_nickname: None, agent_role: None, @@ -834,6 +874,7 @@ async fn test_pagination_cursor() { git_sha: None, git_origin_url: None, source: Some(SessionSource::VSCode), + history_mode: Default::default(), parent_thread_id: None, agent_nickname: None, agent_role: None, @@ -853,6 +894,7 @@ async fn test_pagination_cursor() { git_sha: None, git_origin_url: None, source: Some(SessionSource::VSCode), + history_mode: Default::default(), parent_thread_id: None, agent_nickname: None, agent_role: None, @@ -900,6 +942,7 @@ async fn test_pagination_cursor() { git_sha: None, git_origin_url: None, source: Some(SessionSource::VSCode), + history_mode: Default::default(), parent_thread_id: None, agent_nickname: None, agent_role: None, @@ -1072,6 +1115,7 @@ async fn test_get_thread_contents() { git_sha: None, git_origin_url: None, source: Some(SessionSource::VSCode), + history_mode: Default::default(), parent_thread_id: None, agent_nickname: None, agent_role: None, @@ -1290,6 +1334,7 @@ async fn test_updated_at_uses_file_mtime() -> Result<()> { dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: Default::default(), multi_agent_version: None, context_window: None, }, @@ -1432,6 +1477,7 @@ async fn test_timestamp_only_cursor_skips_same_second_filesystem_ties() { git_sha: None, git_origin_url: None, source: Some(SessionSource::VSCode), + history_mode: Default::default(), parent_thread_id: None, agent_nickname: None, agent_role: None, @@ -1451,6 +1497,7 @@ async fn test_timestamp_only_cursor_skips_same_second_filesystem_ties() { git_sha: None, git_origin_url: None, source: Some(SessionSource::VSCode), + history_mode: Default::default(), parent_thread_id: None, agent_nickname: None, agent_role: None, diff --git a/codex-rs/state/migrations/0040_threads_history_mode.sql b/codex-rs/state/migrations/0040_threads_history_mode.sql new file mode 100644 index 0000000000..53729c2ec6 --- /dev/null +++ b/codex-rs/state/migrations/0040_threads_history_mode.sql @@ -0,0 +1 @@ +ALTER TABLE threads ADD COLUMN history_mode TEXT NOT NULL DEFAULT 'legacy'; diff --git a/codex-rs/state/src/extract.rs b/codex-rs/state/src/extract.rs index 81de0cc39f..f944e3e389 100644 --- a/codex-rs/state/src/extract.rs +++ b/codex-rs/state/src/extract.rs @@ -56,6 +56,7 @@ fn apply_session_meta_from_item(metadata: &mut ThreadMetadata, meta_line: &Sessi } metadata.id = meta_line.meta.id; metadata.source = enum_to_string(&meta_line.meta.source); + metadata.history_mode = meta_line.meta.history_mode; metadata.thread_source = meta_line.meta.thread_source.clone(); metadata.agent_nickname = meta_line.meta.agent_nickname.clone(); metadata.agent_role = meta_line.meta.agent_role.clone(); @@ -345,6 +346,7 @@ mod tests { dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: Default::default(), multi_agent_version: None, context_window: None, }, @@ -538,6 +540,7 @@ mod tests { dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: Default::default(), multi_agent_version: None, context_window: None, }, @@ -560,6 +563,7 @@ mod tests { updated_at: created_at, recency_at: created_at, source: "cli".to_string(), + history_mode: Default::default(), thread_source: None, agent_path: None, agent_nickname: None, diff --git a/codex-rs/state/src/model/thread_metadata.rs b/codex-rs/state/src/model/thread_metadata.rs index ac8096e456..0024c9866d 100644 --- a/codex-rs/state/src/model/thread_metadata.rs +++ b/codex-rs/state/src/model/thread_metadata.rs @@ -6,6 +6,7 @@ use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadSource; use sqlx::Row; use sqlx::sqlite::SqliteRow; @@ -87,6 +88,8 @@ pub struct ThreadMetadata { pub recency_at: DateTime, /// The session source (stringified enum). pub source: String, + /// Persisted thread history contract selected when this thread was created. + pub history_mode: ThreadHistoryMode, /// Optional analytics source classification for this thread. pub thread_source: Option, /// Optional random unique nickname assigned to an AgentControl-spawned sub-agent. @@ -142,6 +145,8 @@ pub struct ThreadMetadataBuilder { pub recency_at: Option>, /// The session source. pub source: SessionSource, + /// Persisted thread history contract selected when this thread was created. + pub history_mode: ThreadHistoryMode, /// Optional analytics source classification for this thread. pub thread_source: Option, /// Optional random unique nickname assigned to the session. @@ -185,6 +190,7 @@ impl ThreadMetadataBuilder { updated_at: None, recency_at: None, source, + history_mode: ThreadHistoryMode::Legacy, thread_source: None, agent_nickname: None, agent_role: None, @@ -222,6 +228,7 @@ impl ThreadMetadataBuilder { updated_at, recency_at, source, + history_mode: self.history_mode, thread_source: self.thread_source.clone(), agent_nickname: self.agent_nickname.clone(), agent_role: self.agent_role.clone(), @@ -368,6 +375,7 @@ pub(crate) struct ThreadRow { updated_at: i64, recency_at: i64, source: String, + history_mode: String, thread_source: Option, agent_nickname: Option, agent_role: Option, @@ -398,6 +406,7 @@ impl ThreadRow { updated_at: row.try_get("updated_at")?, recency_at: row.try_get("recency_at")?, source: row.try_get("source")?, + history_mode: row.try_get("history_mode")?, thread_source: row.try_get("thread_source")?, agent_nickname: row.try_get("agent_nickname")?, agent_role: row.try_get("agent_role")?, @@ -432,6 +441,7 @@ impl TryFrom for ThreadMetadata { updated_at, recency_at, source, + history_mode, thread_source, agent_nickname, agent_role, @@ -456,6 +466,7 @@ impl TryFrom for ThreadMetadata { .map(|thread_source| thread_source.parse()) .transpose() .map_err(anyhow::Error::msg)?; + let history_mode = history_mode.parse().map_err(anyhow::Error::msg)?; Ok(Self { id: ThreadId::try_from(id)?, rollout_path: PathBuf::from(rollout_path), @@ -463,6 +474,7 @@ impl TryFrom for ThreadMetadata { updated_at: epoch_millis_to_datetime(updated_at)?, recency_at: epoch_millis_to_datetime(recency_at)?, source, + history_mode, thread_source, agent_nickname, agent_role, @@ -548,6 +560,7 @@ mod tests { use chrono::Utc; use codex_protocol::ThreadId; use codex_protocol::openai_models::ReasoningEffort; + use codex_protocol::protocol::ThreadHistoryMode; use pretty_assertions::assert_eq; use std::path::PathBuf; @@ -559,6 +572,7 @@ mod tests { updated_at: 1_700_000_100, recency_at: 1_700_000_100, source: "cli".to_string(), + history_mode: "legacy".to_string(), thread_source: None, agent_nickname: None, agent_role: None, @@ -590,6 +604,7 @@ mod tests { updated_at: DateTime::::from_timestamp(1_700_000_100, 0).expect("timestamp"), recency_at: DateTime::::from_timestamp(1_700_000_100, 0).expect("timestamp"), source: "cli".to_string(), + history_mode: ThreadHistoryMode::Legacy, thread_source: None, agent_nickname: None, agent_role: None, @@ -633,4 +648,12 @@ mod tests { expected_thread_metadata(Some(ReasoningEffort::Custom("future".to_string()))) ); } + + #[test] + fn thread_row_rejects_unknown_history_mode() { + let mut row = thread_row(/*reasoning_effort*/ None); + row.history_mode = "future".to_string(); + + assert!(ThreadMetadata::try_from(row).is_err()); + } } diff --git a/codex-rs/state/src/runtime/memories.rs b/codex-rs/state/src/runtime/memories.rs index 83df50e02e..ff25fed982 100644 --- a/codex-rs/state/src/runtime/memories.rs +++ b/codex-rs/state/src/runtime/memories.rs @@ -136,6 +136,7 @@ WHERE kind = ? AND job_key = ? /// - starts from `threads` filtered to active threads and allowed sources /// (`push_thread_filters`) /// - excludes threads with `memory_mode != 'enabled'` + /// - excludes paginated threads because stage 1 still full-loads rollout JSONL /// - excludes the current thread id /// - keeps only threads whose millisecond `updated_at` is in the age window /// - checks memory staleness against the memories DB @@ -177,6 +178,7 @@ SELECT threads.updated_at_ms AS updated_at, threads.recency_at_ms AS recency_at, threads.source, + threads.history_mode, threads.thread_source, threads.agent_path, threads.agent_nickname, @@ -213,7 +215,7 @@ FROM threads }, /*include_thread_id_tiebreaker*/ false, ); - builder.push(" AND threads.memory_mode = 'enabled'"); + builder.push(" AND threads.memory_mode = 'enabled' AND threads.history_mode = 'legacy'"); builder .push(" AND threads.id != ") .push_bind(current_thread_id.as_str()); @@ -549,6 +551,7 @@ SELECT threads.updated_at_ms AS updated_at, threads.recency_at_ms AS recency_at, threads.source, + threads.history_mode, threads.thread_source, threads.agent_nickname, threads.agent_role, @@ -569,7 +572,7 @@ SELECT threads.git_branch, threads.git_origin_url FROM threads -WHERE threads.id = ? AND threads.memory_mode = 'enabled' +WHERE threads.id = ? AND threads.memory_mode = 'enabled' AND threads.history_mode = 'legacy' "#, ) .bind(thread_id.to_string()) @@ -1673,6 +1676,7 @@ mod tests { use chrono::Duration; use chrono::Utc; use codex_protocol::ThreadId; + use codex_protocol::protocol::ThreadHistoryMode; use pretty_assertions::assert_eq; use sqlx::Row; use std::sync::Arc; @@ -2161,7 +2165,7 @@ mod tests { } #[tokio::test] - async fn claim_stage1_jobs_skips_threads_with_disabled_memory_mode() { + async fn claim_stage1_jobs_skips_threads_without_legacy_enabled_memory() { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) .await @@ -2174,6 +2178,8 @@ mod tests { ThreadId::from_string(&Uuid::new_v4().to_string()).expect("current thread id"); let disabled_thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("disabled thread id"); + let paginated_thread_id = + ThreadId::from_string(&Uuid::new_v4().to_string()).expect("paginated thread id"); let enabled_thread_id = ThreadId::from_string(&Uuid::new_v4().to_string()).expect("enabled thread id"); @@ -2200,6 +2206,19 @@ mod tests { .await .expect("disable thread memory mode"); + let mut paginated = test_thread_metadata( + &codex_home, + paginated_thread_id, + codex_home.join("paginated"), + ); + paginated.created_at = eligible_at; + paginated.updated_at = eligible_at; + paginated.history_mode = ThreadHistoryMode::Paginated; + runtime + .upsert_thread(&paginated) + .await + .expect("upsert paginated thread"); + let mut enabled = test_thread_metadata(&codex_home, enabled_thread_id, codex_home.join("enabled")); enabled.created_at = eligible_at; diff --git a/codex-rs/state/src/runtime/test_support.rs b/codex-rs/state/src/runtime/test_support.rs index c029a0ee2d..8c63066842 100644 --- a/codex-rs/state/src/runtime/test_support.rs +++ b/codex-rs/state/src/runtime/test_support.rs @@ -11,6 +11,8 @@ use codex_protocol::protocol::AskForApproval; #[cfg(test)] use codex_protocol::protocol::SandboxPolicy; #[cfg(test)] +use codex_protocol::protocol::ThreadHistoryMode; +#[cfg(test)] use std::path::Path; #[cfg(test)] use std::path::PathBuf; @@ -49,6 +51,7 @@ pub(super) fn test_thread_metadata( updated_at: now, recency_at: now, source: "cli".to_string(), + history_mode: ThreadHistoryMode::Legacy, thread_source: None, agent_nickname: None, agent_role: None, diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index e5ab6fb31d..c80b2b6716 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -15,6 +15,7 @@ SELECT threads.updated_at_ms AS updated_at, threads.recency_at_ms AS recency_at, threads.source, + threads.history_mode, threads.thread_source, threads.agent_nickname, threads.agent_role, @@ -544,6 +545,7 @@ INSERT INTO threads ( updated_at_ms, recency_at_ms, source, + history_mode, thread_source, agent_nickname, agent_role, @@ -565,7 +567,7 @@ INSERT INTO threads ( git_branch, git_origin_url, memory_mode -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO NOTHING "#, ) @@ -578,6 +580,7 @@ ON CONFLICT(id) DO NOTHING .bind(datetime_to_epoch_millis(updated_at)) .bind(datetime_to_epoch_millis(recency_at)) .bind(metadata.source.as_str()) + .bind(metadata.history_mode.as_str()) .bind( metadata .thread_source @@ -796,6 +799,7 @@ INSERT INTO threads ( updated_at_ms, recency_at_ms, source, + history_mode, thread_source, agent_nickname, agent_role, @@ -817,7 +821,7 @@ INSERT INTO threads ( git_branch, git_origin_url, memory_mode -) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET rollout_path = excluded.rollout_path, created_at = excluded.created_at, @@ -827,6 +831,7 @@ ON CONFLICT(id) DO UPDATE SET updated_at_ms = excluded.updated_at_ms, recency_at_ms = threads.recency_at_ms, source = excluded.source, + history_mode = excluded.history_mode, thread_source = excluded.thread_source, agent_nickname = excluded.agent_nickname, agent_role = excluded.agent_role, @@ -858,6 +863,7 @@ ON CONFLICT(id) DO UPDATE SET .bind(datetime_to_epoch_millis(updated_at)) .bind(datetime_to_epoch_millis(insert_recency_at)) .bind(metadata.source.as_str()) + .bind(metadata.history_mode.as_str()) .bind( metadata .thread_source @@ -1206,6 +1212,7 @@ SELECT threads.updated_at_ms AS updated_at, threads.recency_at_ms AS recency_at, threads.source, + threads.history_mode, threads.thread_source, threads.agent_nickname, threads.agent_role, @@ -1419,6 +1426,7 @@ mod tests { use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; use codex_protocol::protocol::SessionSource; + use codex_protocol::protocol::ThreadHistoryMode; use pretty_assertions::assert_eq; use serde_json::json; use std::path::PathBuf; @@ -1461,6 +1469,30 @@ mod tests { assert_eq!(memory_mode, "disabled"); } + #[tokio::test] + async fn thread_metadata_round_trips_history_mode() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await + .expect("state db should initialize"); + let thread_id = + ThreadId::from_string("00000000-0000-0000-0000-000000000124").expect("valid thread id"); + let mut metadata = test_thread_metadata(&codex_home, thread_id, codex_home.clone()); + metadata.history_mode = ThreadHistoryMode::Paginated; + + runtime + .upsert_thread(&metadata) + .await + .expect("upsert should succeed"); + + let metadata = runtime + .get_thread(thread_id) + .await + .expect("thread should load") + .expect("thread should exist"); + assert_eq!(metadata.history_mode, ThreadHistoryMode::Paginated); + } + #[tokio::test] async fn delete_thread_cleans_associated_state() -> Result<()> { let codex_home = unique_temp_dir(); @@ -2150,6 +2182,7 @@ mod tests { dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: Some("polluted".to_string()), + history_mode: Default::default(), multi_agent_version: None, context_window: None, }, @@ -2214,6 +2247,7 @@ mod tests { dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, + history_mode: Default::default(), multi_agent_version: None, context_window: None, }, diff --git a/codex-rs/thread-store/src/error.rs b/codex-rs/thread-store/src/error.rs index 2244c93185..bc5e8ed7c4 100644 --- a/codex-rs/thread-store/src/error.rs +++ b/codex-rs/thread-store/src/error.rs @@ -1,8 +1,20 @@ use codex_protocol::ThreadId; +use codex_protocol::protocol::ThreadHistoryMode; /// Result type returned by thread-store operations. pub type ThreadStoreResult = Result; +pub(crate) fn reject_paginated_history_mode( + history_mode: ThreadHistoryMode, +) -> ThreadStoreResult<()> { + if matches!(history_mode, ThreadHistoryMode::Paginated) { + return Err(ThreadStoreError::Unsupported { + operation: "paginated_threads", + }); + } + Ok(()) +} + /// Error type shared by thread-store implementations. #[derive(Debug, thiserror::Error)] pub enum ThreadStoreError { diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index d64510f223..38598830f7 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -14,6 +14,7 @@ use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionContextWindow; use codex_protocol::protocol::SessionMeta; use codex_protocol::protocol::SessionMetaLine; +use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadMemoryMode; use codex_rollout::persisted_rollout_items; @@ -36,6 +37,8 @@ use crate::ThreadStoreError; use crate::ThreadStoreFuture; use crate::ThreadStoreResult; use crate::UpdateThreadMetadataParams; +use crate::error::reject_paginated_history_mode; +use crate::types::canonical_history_mode_from_rollout_items; static IN_MEMORY_THREAD_STORES: OnceLock>>> = OnceLock::new(); @@ -128,6 +131,7 @@ mod tests { dynamic_tools: Vec::new(), selected_capability_roots: Vec::new(), multi_agent_version: None, + history_mode: ThreadHistoryMode::Legacy, initial_window_id: uuid::Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { cwd: None, @@ -193,6 +197,152 @@ mod tests { HashSet::from([child_thread_id, grandchild_thread_id]) ); } + + #[tokio::test] + async fn paginated_threads_allow_metadata_reads_and_reject_legacy_history_paths() { + let store = InMemoryThreadStore::default(); + let thread_id = ThreadId::default(); + let rollout_path = PathBuf::from("/tmp/paginated-thread.jsonl"); + + store + .create_thread(create_thread_params(thread_id, ThreadHistoryMode::Legacy)) + .await + .expect("create legacy thread"); + store + .resume_thread(ResumeThreadParams { + thread_id, + rollout_path: Some(rollout_path.clone()), + history: None, + include_archived: false, + metadata: thread_metadata(), + }) + .await + .expect("register rollout path"); + store + .update_thread_metadata(UpdateThreadMetadataParams { + thread_id, + patch: ThreadMetadataPatch { + history_mode: Some(ThreadHistoryMode::Paginated), + ..Default::default() + }, + include_archived: false, + }) + .await + .expect("seed paginated metadata"); + + let thread = store + .read_thread(ReadThreadParams { + thread_id, + include_archived: false, + include_history: false, + }) + .await + .expect("metadata read"); + assert_eq!(thread.history_mode, ThreadHistoryMode::Paginated); + assert!(thread.history.is_none()); + + let thread = store + .read_thread_by_rollout_path(ReadThreadByRolloutPathParams { + rollout_path, + include_archived: false, + include_history: false, + }) + .await + .expect("metadata path read"); + assert_eq!(thread.history_mode, ThreadHistoryMode::Paginated); + assert!(thread.history.is_none()); + + assert_paginated_threads_unsupported( + store + .read_thread(ReadThreadParams { + thread_id, + include_archived: false, + include_history: true, + }) + .await + .expect_err("full history read should fail"), + ); + assert_paginated_threads_unsupported( + store + .read_thread_by_rollout_path(ReadThreadByRolloutPathParams { + rollout_path: PathBuf::from("/tmp/paginated-thread.jsonl"), + include_archived: false, + include_history: true, + }) + .await + .expect_err("full history path read should fail"), + ); + assert_paginated_threads_unsupported( + store + .load_history(LoadThreadHistoryParams { + thread_id, + include_archived: false, + }) + .await + .expect_err("history load should fail"), + ); + assert_paginated_threads_unsupported( + store + .resume_thread(ResumeThreadParams { + thread_id, + rollout_path: None, + history: None, + include_archived: false, + metadata: thread_metadata(), + }) + .await + .expect_err("resume should fail"), + ); + assert_paginated_threads_unsupported( + store + .create_thread(create_thread_params( + ThreadId::default(), + ThreadHistoryMode::Paginated, + )) + .await + .expect_err("paginated create should fail"), + ); + } + + fn create_thread_params( + thread_id: ThreadId, + history_mode: ThreadHistoryMode, + ) -> CreateThreadParams { + CreateThreadParams { + session_id: thread_id.into(), + thread_id, + extra_config: None, + forked_from_id: None, + parent_thread_id: None, + source: SessionSource::Exec, + thread_source: None, + originator: "test_originator".to_string(), + base_instructions: BaseInstructions::default(), + dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), + multi_agent_version: None, + history_mode, + initial_window_id: uuid::Uuid::now_v7().to_string(), + metadata: thread_metadata(), + } + } + + fn thread_metadata() -> ThreadPersistenceMetadata { + ThreadPersistenceMetadata { + cwd: None, + model_provider: "test-provider".to_string(), + memory_mode: ThreadMemoryMode::Enabled, + } + } + + fn assert_paginated_threads_unsupported(err: ThreadStoreError) { + assert!(matches!( + err, + ThreadStoreError::Unsupported { + operation: "paginated_threads" + } + )); + } } fn stores_guard() -> MutexGuard<'static, HashMap>> { @@ -265,6 +415,7 @@ impl InMemoryThreadStore { } async fn create_thread(&self, params: CreateThreadParams) -> ThreadStoreResult<()> { + reject_paginated_history_mode(params.history_mode)?; let mut state = self.state.lock().await; state.calls.create_thread += 1; let session_meta = SessionMeta { @@ -285,6 +436,7 @@ impl InMemoryThreadStore { selected_capability_roots: params.selected_capability_roots.clone(), memory_mode: matches!(params.metadata.memory_mode, ThreadMemoryMode::Disabled) .then_some("disabled".to_string()), + history_mode: params.history_mode, multi_agent_version: params.multi_agent_version, context_window: Some(SessionContextWindow::new(params.initial_window_id.clone())), ..SessionMeta::default() @@ -304,6 +456,13 @@ impl InMemoryThreadStore { async fn resume_thread(&self, params: ResumeThreadParams) -> ThreadStoreResult<()> { let mut state = self.state.lock().await; state.calls.resume_thread += 1; + let history_mode = params + .history + .as_deref() + .map(Vec::as_slice) + .map(canonical_history_mode_from_rollout_items) + .unwrap_or_else(|| history_mode_from_state(&state, params.thread_id)); + reject_paginated_history_mode(history_mode)?; if let Some(history) = params.history { state .histories @@ -338,14 +497,18 @@ impl InMemoryThreadStore { ) -> ThreadStoreResult { let mut state = self.state.lock().await; state.calls.load_history += 1; - let items = state.histories.get(¶ms.thread_id).cloned().ok_or( - ThreadStoreError::ThreadNotFound { - thread_id: params.thread_id, - }, - )?; + let items = + state + .histories + .get(¶ms.thread_id) + .ok_or(ThreadStoreError::ThreadNotFound { + thread_id: params.thread_id, + })?; + let history_mode = history_mode_from_state(&state, params.thread_id); + reject_paginated_history_mode(history_mode)?; Ok(StoredThreadHistory { thread_id: params.thread_id, - items, + items: items.clone(), }) } @@ -354,8 +517,10 @@ impl InMemoryThreadStore { state.calls.read_thread += 1; if params.include_history { state.calls.read_thread_with_history += 1; + reject_paginated_history_mode(history_mode_from_state(&state, params.thread_id))?; } - stored_thread_from_state(&state, params.thread_id, params.include_history) + let thread = stored_thread_from_state(&state, params.thread_id, params.include_history)?; + Ok(thread) } async fn read_thread_by_rollout_path( @@ -372,7 +537,11 @@ impl InMemoryThreadStore { ), }); }; - stored_thread_from_state(&state, thread_id, params.include_history) + if params.include_history { + reject_paginated_history_mode(history_mode_from_state(&state, thread_id))?; + } + let thread = stored_thread_from_state(&state, thread_id, params.include_history)?; + Ok(thread) } async fn list_threads(&self) -> ThreadStoreResult { @@ -615,6 +784,9 @@ fn stored_thread_from_state( source: metadata .and_then(|metadata| metadata.source.clone()) .unwrap_or_else(|| created.source.clone()), + history_mode: metadata + .and_then(|metadata| metadata.history_mode) + .unwrap_or(created.history_mode), thread_source: metadata .and_then(|metadata| metadata.thread_source.clone()) .unwrap_or_else(|| created.thread_source.clone()), @@ -634,6 +806,23 @@ fn stored_thread_from_state( }) } +fn history_mode_from_state( + state: &InMemoryThreadStoreState, + thread_id: ThreadId, +) -> ThreadHistoryMode { + state + .metadata_updates + .get(&thread_id) + .and_then(|metadata| metadata.history_mode) + .or_else(|| { + state + .created_threads + .get(&thread_id) + .map(|thread| thread.history_mode) + }) + .unwrap_or_default() +} + fn git_info_from_patch(patch: &ThreadMetadataPatch) -> Option { let git_info = patch.git_info.as_ref()?; let sha = git_info.sha.clone().flatten(); diff --git a/codex-rs/thread-store/src/local/create_thread.rs b/codex-rs/thread-store/src/local/create_thread.rs index 451931b4dc..9b370ef095 100644 --- a/codex-rs/thread-store/src/local/create_thread.rs +++ b/codex-rs/thread-store/src/local/create_thread.rs @@ -2,6 +2,7 @@ use super::LocalThreadStore; use crate::CreateThreadParams; use crate::ThreadStoreError; use crate::ThreadStoreResult; +use crate::error::reject_paginated_history_mode; use codex_protocol::protocol::ThreadMemoryMode; use codex_rollout::RolloutConfig; use codex_rollout::RolloutRecorder; @@ -11,6 +12,7 @@ pub(super) async fn create_thread( store: &LocalThreadStore, params: CreateThreadParams, ) -> ThreadStoreResult { + reject_paginated_history_mode(params.history_mode)?; let cwd = params .metadata .cwd @@ -40,6 +42,7 @@ pub(super) async fn create_thread( .with_session_id(params.session_id) .with_selected_capability_roots(params.selected_capability_roots) .with_multi_agent_version(params.multi_agent_version) + .with_history_mode(params.history_mode) .with_initial_window_id(params.initial_window_id), ) .await diff --git a/codex-rs/thread-store/src/local/helpers.rs b/codex-rs/thread-store/src/local/helpers.rs index f48a0136c0..d883bd7874 100644 --- a/codex-rs/thread-store/src/local/helpers.rs +++ b/codex-rs/thread-store/src/local/helpers.rs @@ -142,6 +142,7 @@ pub(super) fn stored_thread_from_rollout_item( cwd: item.cwd.unwrap_or_default(), cli_version: item.cli_version.unwrap_or_default(), source, + history_mode: item.history_mode, thread_source: None, agent_nickname: item.agent_nickname, agent_role: item.agent_role, diff --git a/codex-rs/thread-store/src/local/list_threads.rs b/codex-rs/thread-store/src/local/list_threads.rs index 2619fe68aa..f514231333 100644 --- a/codex-rs/thread-store/src/local/list_threads.rs +++ b/codex-rs/thread-store/src/local/list_threads.rs @@ -219,6 +219,7 @@ mod tests { use chrono::Utc; use codex_protocol::ThreadId; use codex_protocol::protocol::SessionSource; + use codex_protocol::protocol::ThreadHistoryMode; use pretty_assertions::assert_eq; use std::fs; use tempfile::TempDir; @@ -243,6 +244,7 @@ mod tests { Uuid::from_u128(102), "Hello from user", /*model_provider*/ None, + ThreadHistoryMode::Legacy, ) .expect("session file"); diff --git a/codex-rs/thread-store/src/local/live_writer.rs b/codex-rs/thread-store/src/local/live_writer.rs index 3e7541b36b..85ef80d88a 100644 --- a/codex-rs/thread-store/src/local/live_writer.rs +++ b/codex-rs/thread-store/src/local/live_writer.rs @@ -16,6 +16,8 @@ use crate::ReadThreadParams; use crate::ResumeThreadParams; use crate::ThreadStoreError; use crate::ThreadStoreResult; +use crate::error::reject_paginated_history_mode; +use crate::types::canonical_history_mode_from_rollout_items; const ROLLOUT_SIZE_BYTES_METRIC: &str = "codex.rollout.size_bytes"; @@ -34,6 +36,30 @@ pub(super) async fn resume_thread( params: ResumeThreadParams, ) -> ThreadStoreResult<()> { store.ensure_live_recorder_absent(params.thread_id).await?; + let history_mode = if let Some(history) = params.history.as_deref() { + canonical_history_mode_from_rollout_items(history) + } else if let Some(rollout_path) = params.rollout_path.as_ref() { + super::read_thread::read_thread_by_rollout_path( + store, + rollout_path.clone(), + params.include_archived, + /*include_history*/ false, + ) + .await? + .history_mode + } else { + super::read_thread::read_thread( + store, + ReadThreadParams { + thread_id: params.thread_id, + include_archived: params.include_archived, + include_history: false, + }, + ) + .await? + .history_mode + }; + reject_paginated_history_mode(history_mode)?; let rollout_path = match (params.rollout_path, params.history) { (Some(rollout_path), _history) => rollout_path, (None, history) => { @@ -46,7 +72,6 @@ pub(super) async fn resume_thread( }, ) .await?; - thread .rollout_path .ok_or_else(|| ThreadStoreError::Internal { diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 717ae09b69..bc34c52f33 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -321,6 +321,7 @@ mod tests { use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionSource; + use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadMemoryMode; use codex_protocol::protocol::TurnCompleteEvent; use codex_protocol::protocol::TurnStartedEvent; @@ -333,6 +334,7 @@ mod tests { use crate::local::test_support::test_config; use crate::local::test_support::write_archived_session_file; use crate::local::test_support::write_session_file; + use crate::local::test_support::write_session_file_with_history_mode; #[tokio::test] async fn live_writer_lifecycle_writes_and_closes() { @@ -1109,18 +1111,108 @@ mod tests { .expect("read thread by rollout path"); assert_eq!(thread.thread_id, thread_id); + assert_eq!(thread.history_mode, ThreadHistoryMode::Legacy); assert_eq!( thread .history + .as_ref() .expect("history") .items - .into_iter() + .iter() .filter(|item| matches!(item, RolloutItem::EventMsg(EventMsg::UserMessage(_)))) .count(), 1 ); } + #[tokio::test] + async fn paginated_threads_allow_metadata_reads_and_reject_legacy_history_paths() { + let home = TempDir::new().expect("temp dir"); + let store = LocalThreadStore::new(test_config(home.path()), /*state_db*/ None); + let uuid = uuid::Uuid::from_u128(408); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); + let rollout_path = write_session_file_with_history_mode( + home.path(), + "2025-01-04T12-00-00", + uuid, + ThreadHistoryMode::Paginated, + ) + .expect("session file"); + + let thread = store + .read_thread(ReadThreadParams { + thread_id, + include_archived: false, + include_history: false, + }) + .await + .expect("metadata read"); + assert_eq!(thread.history_mode, ThreadHistoryMode::Paginated); + assert!(thread.history.is_none()); + + let thread = store + .read_thread_by_rollout_path( + rollout_path.clone(), + /*include_archived*/ true, + /*include_history*/ false, + ) + .await + .expect("metadata path read"); + assert_eq!(thread.history_mode, ThreadHistoryMode::Paginated); + assert!(thread.history.is_none()); + + assert_paginated_threads_unsupported( + store + .read_thread(ReadThreadParams { + thread_id, + include_archived: false, + include_history: true, + }) + .await + .expect_err("full history read should fail"), + ); + assert_paginated_threads_unsupported( + store + .read_thread_by_rollout_path( + rollout_path.clone(), + /*include_archived*/ true, + /*include_history*/ true, + ) + .await + .expect_err("full history path read should fail"), + ); + assert_paginated_threads_unsupported( + store + .load_history(LoadThreadHistoryParams { + thread_id, + include_archived: false, + }) + .await + .expect_err("history load should fail"), + ); + assert_paginated_threads_unsupported( + store + .resume_thread(ResumeThreadParams { + thread_id, + rollout_path: Some(rollout_path), + history: None, + include_archived: false, + metadata: thread_metadata(), + }) + .await + .expect_err("resume should fail"), + ); + + let mut create_params = create_thread_params(ThreadId::default()); + create_params.history_mode = ThreadHistoryMode::Paginated; + assert_paginated_threads_unsupported( + store + .create_thread(create_params) + .await + .expect_err("paginated create should fail"), + ); + } + fn create_thread_params(thread_id: ThreadId) -> CreateThreadParams { CreateThreadParams { session_id: thread_id.into(), @@ -1135,11 +1227,21 @@ mod tests { dynamic_tools: Vec::new(), selected_capability_roots: Vec::new(), multi_agent_version: None, + history_mode: ThreadHistoryMode::Legacy, initial_window_id: uuid::Uuid::now_v7().to_string(), metadata: thread_metadata(), } } + fn assert_paginated_threads_unsupported(err: ThreadStoreError) { + assert!(matches!( + err, + ThreadStoreError::Unsupported { + operation: "paginated_threads" + } + )); + } + fn thread_metadata() -> ThreadPersistenceMetadata { ThreadPersistenceMetadata { cwd: Some(std::env::current_dir().expect("cwd")), diff --git a/codex-rs/thread-store/src/local/read_thread.rs b/codex-rs/thread-store/src/local/read_thread.rs index 1271e4fcf5..27ce3076e9 100644 --- a/codex-rs/thread-store/src/local/read_thread.rs +++ b/codex-rs/thread-store/src/local/read_thread.rs @@ -25,6 +25,7 @@ use crate::StoredThread; use crate::StoredThreadHistory; use crate::ThreadStoreError; use crate::ThreadStoreResult; +use crate::error::reject_paginated_history_mode; pub(super) async fn read_thread( store: &LocalThreadStore, @@ -47,7 +48,7 @@ pub(super) async fn read_thread( .await) { let metadata_sandbox_policy = metadata.sandbox_policy.clone(); - let mut thread = stored_thread_from_sqlite_metadata(store, metadata).await; + let mut thread = stored_thread_from_sqlite_metadata(store, metadata).await?; if !params.include_history && let Some(rollout_path) = thread.rollout_path.clone() && let Ok(mut rollout_thread) = read_thread_from_rollout_path(store, rollout_path).await @@ -66,6 +67,7 @@ pub(super) async fn read_thread( ); thread = rollout_thread; } + reject_paginated_history(&thread, params.include_history)?; attach_history_if_requested(&mut thread, params.include_history).await?; return Ok(thread); } @@ -82,6 +84,7 @@ pub(super) async fn read_thread( message: format!("thread {} is archived", thread.thread_id), }); } + reject_paginated_history(&thread, params.include_history)?; attach_history_if_requested(&mut thread, params.include_history).await?; Ok(thread) } @@ -132,10 +135,18 @@ pub(super) async fn read_thread_by_rollout_path( metadata.git_origin_url.or(fallback_origin_url), ); } + reject_paginated_history(&thread, include_history)?; attach_history_if_requested(&mut thread, include_history).await?; Ok(thread) } +fn reject_paginated_history(thread: &StoredThread, include_history: bool) -> ThreadStoreResult<()> { + if include_history { + reject_paginated_history_mode(thread.history_mode)?; + } + Ok(()) +} + async fn resolve_requested_rollout_path( store: &LocalThreadStore, rollout_path: std::path::PathBuf, @@ -261,16 +272,16 @@ async fn read_thread_from_rollout_path( message: format!("failed to read thread id from {}", path.display()), })?; thread.rollout_path = Some(codex_rollout::plain_rollout_path(path.as_path())); - if let Ok(meta_line) = read_session_meta_line(path.as_path()).await { - thread.forked_from_id = meta_line.meta.forked_from_id; - thread.parent_thread_id = meta_line.meta.parent_thread_id; - if let Some(model_provider) = meta_line - .meta - .model_provider - .filter(|provider| !provider.is_empty()) - { - thread.model_provider = model_provider; - } + let meta_line = read_required_session_meta_line(path.as_path()).await?; + thread.forked_from_id = meta_line.meta.forked_from_id; + thread.parent_thread_id = meta_line.meta.parent_thread_id; + thread.history_mode = meta_line.meta.history_mode; + if let Some(model_provider) = meta_line + .meta + .model_provider + .filter(|provider| !provider.is_empty()) + { + thread.model_provider = model_provider; } if let Ok(Some(title)) = find_thread_name_by_id(store.config.codex_home.as_path(), &thread.thread_id).await @@ -302,7 +313,7 @@ async fn read_sqlite_metadata( async fn stored_thread_from_sqlite_metadata( store: &LocalThreadStore, metadata: ThreadMetadata, -) -> StoredThread { +) -> ThreadStoreResult { let name = match distinct_thread_metadata_title(&metadata) { Some(title) => Some(title), None => find_thread_name_by_id(store.config.codex_home.as_path(), &metadata.id) @@ -311,13 +322,32 @@ async fn stored_thread_from_sqlite_metadata( .flatten() .filter(|title| !title.trim().is_empty()), }; - let session_meta = read_session_meta_line(metadata.rollout_path.as_path()) - .await - .ok() - .map(|meta_line| meta_line.meta); + let session_meta = match read_required_session_meta_line(metadata.rollout_path.as_path()).await + { + Ok(meta_line) => Some(meta_line.meta), + Err(_) + if codex_rollout::existing_rollout_path(metadata.rollout_path.as_path()) + .await + .is_none() => + { + None + } + Err(err) => { + return Err(ThreadStoreError::Internal { + message: format!( + "failed to read session metadata {}: {err}", + metadata.rollout_path.display() + ), + }); + } + }; let rollout_path = codex_rollout::plain_rollout_path(metadata.rollout_path.as_path()); let forked_from_id = session_meta.as_ref().and_then(|meta| meta.forked_from_id); let parent_thread_id = session_meta.as_ref().and_then(|meta| meta.parent_thread_id); + let history_mode = session_meta + .as_ref() + .map(|meta| meta.history_mode) + .unwrap_or(metadata.history_mode); let preview = metadata .preview .clone() @@ -325,7 +355,7 @@ async fn stored_thread_from_sqlite_metadata( .unwrap_or_default(); let permission_profile = permission_profile_from_metadata_value(&metadata.sandbox_policy, metadata.cwd.as_path()); - StoredThread { + Ok(StoredThread { thread_id: metadata.id, extra_config: None, rollout_path: Some(rollout_path), @@ -347,6 +377,7 @@ async fn stored_thread_from_sqlite_metadata( cwd: metadata.cwd, cli_version: metadata.cli_version, source: parse_session_source(&metadata.source), + history_mode, thread_source: metadata.thread_source, agent_nickname: metadata.agent_nickname, agent_role: metadata.agent_role, @@ -361,24 +392,30 @@ async fn stored_thread_from_sqlite_metadata( token_usage: None, first_user_message: metadata.first_user_message, history: None, - } + }) } async fn stored_thread_from_session_meta( store: &LocalThreadStore, path: std::path::PathBuf, ) -> ThreadStoreResult { - let meta_line = read_session_meta_line(path.as_path()) - .await - .map_err(|err| ThreadStoreError::Internal { - message: format!("failed to read thread {}: {err}", path.display()), - })?; + let meta_line = read_required_session_meta_line(path.as_path()).await?; let archived = rollout_path_is_archived(store.config.codex_home.as_path(), path.as_path()); Ok(stored_thread_from_meta_line( store, meta_line, path, archived, )) } +async fn read_required_session_meta_line( + path: &std::path::Path, +) -> ThreadStoreResult { + read_session_meta_line(path) + .await + .map_err(|err| ThreadStoreError::Internal { + message: format!("failed to read session metadata {}: {err}", path.display()), + }) +} + fn stored_thread_from_meta_line( store: &LocalThreadStore, meta_line: SessionMetaLine, @@ -414,6 +451,7 @@ fn stored_thread_from_meta_line( cwd: meta_line.meta.cwd, cli_version: meta_line.meta.cli_version, source: meta_line.meta.source, + history_mode: meta_line.meta.history_mode, thread_source: meta_line.meta.thread_source, agent_nickname: meta_line.meta.agent_nickname, agent_role: meta_line.meta.agent_role, @@ -457,6 +495,7 @@ mod tests { use codex_protocol::ThreadId; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionSource; + use codex_protocol::protocol::ThreadHistoryMode; use codex_state::ThreadMetadataBuilder; use pretty_assertions::assert_eq; use tempfile::TempDir; @@ -665,6 +704,7 @@ mod tests { "Forked user message", Some("test-provider"), Some(parent_uuid), + ThreadHistoryMode::Legacy, ) .expect("forked session file"); diff --git a/codex-rs/thread-store/src/local/test_support.rs b/codex-rs/thread-store/src/local/test_support.rs index 10a55c4fd6..a930067fc4 100644 --- a/codex-rs/thread-store/src/local/test_support.rs +++ b/codex-rs/thread-store/src/local/test_support.rs @@ -3,6 +3,7 @@ use std::io::Write; use std::path::Path; use std::path::PathBuf; +use codex_protocol::protocol::ThreadHistoryMode; use codex_rollout::ARCHIVED_SESSIONS_SUBDIR; use uuid::Uuid; @@ -17,6 +18,15 @@ pub(super) fn test_config(codex_home: &Path) -> LocalThreadStoreConfig { } pub(super) fn write_session_file(root: &Path, ts: &str, uuid: Uuid) -> std::io::Result { + write_session_file_with_history_mode(root, ts, uuid, ThreadHistoryMode::Legacy) +} + +pub(super) fn write_session_file_with_history_mode( + root: &Path, + ts: &str, + uuid: Uuid, + history_mode: ThreadHistoryMode, +) -> std::io::Result { write_session_file_with( root, root.join("sessions/2025/01/03"), @@ -24,6 +34,7 @@ pub(super) fn write_session_file(root: &Path, ts: &str, uuid: Uuid) -> std::io:: uuid, "Hello from user", Some("test-provider"), + history_mode, ) } @@ -39,6 +50,7 @@ pub(super) fn write_archived_session_file( uuid, "Archived user message", Some("test-provider"), + ThreadHistoryMode::Legacy, ) } @@ -49,6 +61,7 @@ pub(super) fn write_session_file_with( uuid: Uuid, first_user_message: &str, model_provider: Option<&str>, + history_mode: ThreadHistoryMode, ) -> std::io::Result { write_session_file_with_fork( root, @@ -58,9 +71,11 @@ pub(super) fn write_session_file_with( first_user_message, model_provider, /*forked_from_id*/ None, + history_mode, ) } +#[allow(clippy::too_many_arguments)] pub(super) fn write_session_file_with_fork( root: &Path, day_dir: PathBuf, @@ -69,6 +84,7 @@ pub(super) fn write_session_file_with_fork( first_user_message: &str, model_provider: Option<&str>, forked_from_id: Option, + history_mode: ThreadHistoryMode, ) -> std::io::Result { fs::create_dir_all(&day_dir)?; let path = day_dir.join(format!("rollout-{ts}-{uuid}.jsonl")); @@ -86,6 +102,7 @@ pub(super) fn write_session_file_with_fork( "cli_version": "test_version", "source": "cli", "model_provider": model_provider, + "history_mode": history_mode, "git": { "commit_hash": "abcdef", "branch": "main", diff --git a/codex-rs/thread-store/src/local/update_thread_metadata.rs b/codex-rs/thread-store/src/local/update_thread_metadata.rs index 5f78b2e1f1..38e919b879 100644 --- a/codex-rs/thread-store/src/local/update_thread_metadata.rs +++ b/codex-rs/thread-store/src/local/update_thread_metadata.rs @@ -27,6 +27,7 @@ use crate::ThreadMetadataPatch; use crate::ThreadStoreError; use crate::ThreadStoreResult; use crate::UpdateThreadMetadataParams; +use crate::error::reject_paginated_history_mode; use crate::local::read_thread; struct ResolvedRolloutPath { @@ -53,6 +54,20 @@ pub(super) async fn update_thread_metadata( } let needs_rollout_compat = needs_rollout_compatibility_update(&patch); + if needs_rollout_compat { + // These explicit patches still write legacy rollout/name-index state after the + // SQLite update. Paginated threads must fail before either side is mutated. + let thread = read_thread::read_thread( + store, + ReadThreadParams { + thread_id, + include_archived: params.include_archived, + include_history: false, + }, + ) + .await?; + reject_paginated_history_mode(thread.history_mode)?; + } let require_sqlite_write = sqlite_write_failure_should_block(&patch); let updated = apply_metadata_update( store, @@ -222,6 +237,7 @@ async fn apply_metadata_update( patch.source.clone().unwrap_or(SessionSource::Unknown), ); builder.model_provider = patch.model_provider.clone(); + builder.history_mode = patch.history_mode.unwrap_or_default(); builder.thread_source = patch.thread_source.clone().flatten(); builder.agent_nickname = patch.agent_nickname.clone().flatten(); builder.agent_role = patch.agent_role.clone().flatten(); @@ -269,6 +285,9 @@ async fn apply_metadata_update( if let Some(source) = patch.source { metadata.source = enum_to_string(&source); } + if let Some(history_mode) = patch.history_mode { + metadata.history_mode = history_mode; + } if let Some(thread_source) = patch.thread_source { metadata.thread_source = thread_source; } @@ -412,6 +431,7 @@ fn has_observed_metadata_facts(patch: &ThreadMetadataPatch) -> bool { || patch.permission_profile.is_some() || patch.token_usage.is_some() || patch.first_user_message.is_some() + || patch.history_mode.is_some() } fn enum_to_string(value: &T) -> String { @@ -632,6 +652,7 @@ fn rollout_path_is_archived(store: &LocalThreadStore, path: &Path) -> bool { #[cfg(test)] mod tests { use codex_protocol::models::PermissionProfile; + use codex_protocol::protocol::ThreadHistoryMode; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; @@ -651,6 +672,7 @@ mod tests { use crate::local::test_support::test_config; use crate::local::test_support::write_archived_session_file; use crate::local::test_support::write_session_file; + use crate::local::test_support::write_session_file_with_history_mode; #[tokio::test] async fn update_thread_metadata_sets_name_on_active_rollout_and_indexes_name() { @@ -719,6 +741,55 @@ mod tests { assert_eq!(memory_mode.as_deref(), Some("disabled")); } + #[tokio::test] + async fn update_thread_metadata_rejects_paginated_rollout_compatibility_writes() { + let home = TempDir::new().expect("temp dir"); + let config = test_config(home.path()); + let uuid = Uuid::from_u128(303); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); + let path = write_session_file_with_history_mode( + home.path(), + "2025-01-03T14-35-00", + uuid, + ThreadHistoryMode::Paginated, + ) + .expect("session file"); + let runtime = codex_state::StateRuntime::init( + home.path().to_path_buf(), + config.default_model_provider_id.clone(), + ) + .await + .expect("state db should initialize"); + let store = LocalThreadStore::new(config, Some(runtime.clone())); + + assert!(matches!( + store + .update_thread_metadata(UpdateThreadMetadataParams { + thread_id, + patch: ThreadMetadataPatch { + memory_mode: Some(ThreadMemoryMode::Disabled), + ..Default::default() + }, + include_archived: false, + }) + .await + .expect_err("paginated rollout compatibility write should fail"), + ThreadStoreError::Unsupported { + operation: "paginated_threads" + } + )); + + assert_eq!(last_rollout_item(path.as_path())["type"], "event_msg"); + assert_eq!( + runtime + .get_thread_memory_mode(thread_id) + .await + .expect("thread memory mode should be readable") + .as_deref(), + Some("enabled") + ); + } + #[tokio::test] async fn update_thread_metadata_preserves_memory_mode_when_updating_git_info() { let home = TempDir::new().expect("temp dir"); diff --git a/codex-rs/thread-store/src/store.rs b/codex-rs/thread-store/src/store.rs index 65ad1ce8a8..67ee603f49 100644 --- a/codex-rs/thread-store/src/store.rs +++ b/codex-rs/thread-store/src/store.rs @@ -1,4 +1,5 @@ use codex_protocol::ThreadId; +use codex_protocol::protocol::ThreadHistoryMode; use std::any::Any; use std::future::Future; use std::pin::Pin; @@ -33,6 +34,14 @@ pub trait ThreadStore: Any + Send + Sync { /// Return this store as [`Any`] for implementation-owned escape hatches. fn as_any(&self) -> &dyn Any; + /// Returns the history mode to use when history does not carry a persisted mode. + /// + /// The default is legacy so existing stores stay compatible. Stores whose durable contract is + /// already paginated should override this instead of relying on core to infer storage behavior. + fn default_history_mode(&self) -> ThreadHistoryMode { + ThreadHistoryMode::Legacy + } + /// Creates a new live thread. fn create_thread(&self, params: CreateThreadParams) -> ThreadStoreFuture<'_, ()>; diff --git a/codex-rs/thread-store/src/thread_metadata_sync.rs b/codex-rs/thread-store/src/thread_metadata_sync.rs index 1057715a4a..a7b52ee0cb 100644 --- a/codex-rs/thread-store/src/thread_metadata_sync.rs +++ b/codex-rs/thread-store/src/thread_metadata_sync.rs @@ -76,6 +76,7 @@ impl ThreadMetadataSync { cli_version: Some(env!("CARGO_PKG_VERSION").to_string()), git_info: git_info.map(git_info_patch_from_observation), memory_mode: Some(params.metadata.memory_mode), + history_mode: Some(params.history_mode), ..Default::default() }; Self { @@ -237,6 +238,7 @@ impl ThreadMetadataSync { { update.memory_mode = Some(memory_mode); } + update.history_mode = Some(meta_line.meta.history_mode); } RolloutItem::TurnContext(turn_ctx) => { if !self.cwd_seen { @@ -375,6 +377,7 @@ fn update_has_metadata_facts(update: &ThreadMetadataPatch) -> bool { || update.first_user_message.is_some() || update.git_info.is_some() || update.memory_mode.is_some() + || update.history_mode.is_some() } fn git_info_patch_from_observation(git_info: GitInfo) -> GitInfoPatch { diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index f96a4df31b..3b86840ca7 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -15,6 +15,7 @@ use codex_protocol::protocol::GitInfo; use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadMemoryMode as MemoryMode; use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TokenUsage; @@ -91,6 +92,8 @@ pub struct CreateThreadParams { pub selected_capability_roots: Vec, /// Multi-agent runtime selected when the thread was created. pub multi_agent_version: Option, + /// Persisted thread history contract selected when the thread was created. + pub history_mode: ThreadHistoryMode, /// Initial context-window identity captured when the thread was created. pub initial_window_id: String, /// Metadata captured for the newly created thread. @@ -112,6 +115,20 @@ pub struct ResumeThreadParams { pub metadata: ThreadPersistenceMetadata, } +pub(crate) fn canonical_history_mode_from_rollout_items( + items: &[RolloutItem], +) -> ThreadHistoryMode { + // Forked rollouts keep copied source SessionMeta items after the new thread's + // canonical SessionMeta, so the thread contract comes from the first one. + items + .iter() + .find_map(|item| match item { + RolloutItem::SessionMeta(meta_line) => Some(meta_line.meta.history_mode), + _ => None, + }) + .unwrap_or_default() +} + /// Parameters for appending rollout items to a live thread. #[derive(Clone, Debug)] pub struct AppendThreadItemsParams { @@ -430,6 +447,8 @@ pub struct StoredThread { pub cli_version: String, /// Runtime source for the thread. pub source: SessionSource, + /// Persisted thread history contract selected when this thread was created. + pub history_mode: ThreadHistoryMode, /// Optional analytics source classification for this thread. pub thread_source: Option, /// Optional random nickname for thread-spawn sub-agents. @@ -577,6 +596,8 @@ pub struct ThreadMetadataPatch { pub git_info: Option, /// Thread memory behavior. pub memory_mode: Option, + /// Persisted thread history contract. + pub history_mode: Option, } impl ThreadMetadataPatch { @@ -657,6 +678,9 @@ impl ThreadMetadataPatch { if next.memory_mode.is_some() { self.memory_mode = next.memory_mode; } + if next.history_mode.is_some() { + self.history_mode = next.history_mode; + } } pub fn is_empty(&self) -> bool { @@ -683,6 +707,7 @@ impl ThreadMetadataPatch { && self.first_user_message.is_none() && self.git_info.is_none() && self.memory_mode.is_none() + && self.history_mode.is_none() } } @@ -785,6 +810,17 @@ mod tests { assert!(decoded.is_empty()); } + #[test] + fn canonical_history_mode_uses_first_session_meta() { + assert_eq!( + canonical_history_mode_from_rollout_items(&[ + session_meta(ThreadHistoryMode::Legacy), + session_meta(ThreadHistoryMode::Paginated), + ]), + ThreadHistoryMode::Legacy + ); + } + #[test] fn thread_metadata_patch_merge_uses_presence_semantics() { let mut current = ThreadMetadataPatch { @@ -822,4 +858,14 @@ mod tests { }) ); } + + fn session_meta(history_mode: ThreadHistoryMode) -> RolloutItem { + RolloutItem::SessionMeta(codex_protocol::protocol::SessionMetaLine { + meta: codex_protocol::protocol::SessionMeta { + history_mode, + ..Default::default() + }, + git: None, + }) + } } diff --git a/codex-rs/tui/src/app/loaded_threads.rs b/codex-rs/tui/src/app/loaded_threads.rs index 8e790170e6..9482453e92 100644 --- a/codex-rs/tui/src/app/loaded_threads.rs +++ b/codex-rs/tui/src/app/loaded_threads.rs @@ -134,6 +134,7 @@ mod tests { parent_thread_id: None, preview: String::new(), ephemeral: false, + history_mode: Default::default(), model_provider: "openai".to_string(), created_at: 0, updated_at: 0, diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index 29e29541a7..d8d63575bb 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -2880,6 +2880,7 @@ async fn inactive_thread_started_notification_initializes_replay_session() -> Re parent_thread_id: None, preview: "agent thread".to_string(), ephemeral: false, + history_mode: Default::default(), model_provider: "agent-provider".to_string(), created_at: 1, updated_at: 2, @@ -2974,6 +2975,7 @@ async fn inactive_thread_started_notification_preserves_primary_model_when_path_ parent_thread_id: None, preview: "agent thread".to_string(), ephemeral: false, + history_mode: Default::default(), model_provider: "agent-provider".to_string(), created_at: 1, updated_at: 2, @@ -3035,6 +3037,7 @@ async fn thread_read_session_state_does_not_reuse_primary_permission_profile() { parent_thread_id: None, preview: "read thread".to_string(), ephemeral: false, + history_mode: Default::default(), model_provider: "read-provider".to_string(), created_at: 1, updated_at: 2, @@ -5652,6 +5655,7 @@ async fn thread_rollback_response_discards_queued_active_thread_events() { parent_thread_id: None, preview: String::new(), ephemeral: false, + history_mode: Default::default(), model_provider: "openai".to_string(), created_at: 0, updated_at: 0, diff --git a/codex-rs/tui/src/app/thread_session_state.rs b/codex-rs/tui/src/app/thread_session_state.rs index 58412bff08..602c5b2101 100644 --- a/codex-rs/tui/src/app/thread_session_state.rs +++ b/codex-rs/tui/src/app/thread_session_state.rs @@ -415,6 +415,7 @@ mod tests { parent_thread_id: None, preview: "read thread".to_string(), ephemeral: false, + history_mode: Default::default(), model_provider: "read-provider".to_string(), created_at: 1, updated_at: 2, diff --git a/codex-rs/tui/src/app_server_session.rs b/codex-rs/tui/src/app_server_session.rs index 7181d1a560..0405bb75ab 100644 --- a/codex-rs/tui/src/app_server_session.rs +++ b/codex-rs/tui/src/app_server_session.rs @@ -2326,6 +2326,7 @@ mod tests { parent_thread_id: None, preview: "hello".to_string(), ephemeral: false, + history_mode: Default::default(), model_provider: "openai".to_string(), created_at: 1, updated_at: 2, diff --git a/codex-rs/tui/src/resume_picker.rs b/codex-rs/tui/src/resume_picker.rs index d74017d8d0..e0e3907246 100644 --- a/codex-rs/tui/src/resume_picker.rs +++ b/codex-rs/tui/src/resume_picker.rs @@ -5727,6 +5727,7 @@ session_picker_view = "dense" parent_thread_id: None, preview: String::from("remote thread"), ephemeral: false, + history_mode: Default::default(), model_provider: String::from("openai"), created_at: 1, updated_at: 2, @@ -5764,6 +5765,7 @@ session_picker_view = "dense" parent_thread_id: None, preview: String::from("preview"), ephemeral: false, + history_mode: Default::default(), model_provider: String::from("openai"), created_at: 1, updated_at: 2, @@ -5835,6 +5837,7 @@ session_picker_view = "dense" parent_thread_id: None, preview: String::from("preview"), ephemeral: false, + history_mode: Default::default(), model_provider: String::from("openai"), created_at: 1, updated_at: 2, @@ -5895,6 +5898,7 @@ session_picker_view = "dense" parent_thread_id: None, preview: String::from("preview"), ephemeral: false, + history_mode: Default::default(), model_provider: String::from("openai"), created_at: 1, updated_at: 2, From f91334380e46cfd909c3642c1e05d3d8283c40ef Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Fri, 26 Jun 2026 11:05:00 -0700 Subject: [PATCH 25/38] fix main (#30276) Introduced by a merge race around thread.history_mode. --- codex-rs/core/src/session/tests.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 93519cda9d..20f32ddfcd 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -9197,6 +9197,7 @@ async fn attach_in_memory_thread_store( dynamic_tools: Vec::new(), selected_capability_roots: Vec::new(), multi_agent_version: None, + history_mode: Default::default(), initial_window_id: Uuid::now_v7().to_string(), metadata: ThreadPersistenceMetadata { cwd: Some(config.cwd.to_path_buf()), From d9cf931d0eafa18f4fdfc7ab8c9c241b48df3216 Mon Sep 17 00:00:00 2001 From: hefuc-oai Date: Fri, 26 Jun 2026 11:37:40 -0700 Subject: [PATCH 26/38] [codex] Add managed new-thread model settings (#29683) ## Why Admins need persistent defaults for the model, reasoning effort, and service tier shown when the Desktop App creates a new thread. These are initialization defaults rather than runtime constraints: the App should use them to initialize its draft while still allowing a user to make an explicit selection. The app-server therefore needs to expose the managed values before thread creation without changing `thread/start` behavior for other clients. ## What changed - Parse `model`, `model_reasoning_effort`, and `service_tier` from `[models.new_thread]` in `requirements.toml`. - Compose the `models` requirements through the existing requirements-layer precedence rules. - Expose the resolved values through `configRequirements/read` as `requirements.models.newThread`. - Add the corresponding app-server protocol types and regenerate the JSON and TypeScript schema fixtures. - Document the new `configRequirements/read` fields in the app-server README. ## Scope This PR is data plumbing only. It does not apply these values during `thread/start` and does not change thread creation for existing app-server clients, resumed or forked sessions, internal or subagent sessions, `codex exec`, or the TUI. A companion Desktop App change owns draft initialization, sends the effective settings for ordinary and prewarmed starts, and preserves explicit user changes. ## Validation - Requirements deserialization coverage for `[models.new_thread]` - Requirements-layer precedence coverage - App-server API mapping coverage - `configRequirements/read` integration coverage - Regenerated app-server JSON and TypeScript schema fixtures --- .../codex_app_server_protocol.schemas.json | 52 ++++++++++++ .../codex_app_server_protocol.v2.schemas.json | 52 ++++++++++++ .../v2/ConfigRequirementsReadResponse.json | 57 +++++++++++++ .../typescript/v2/ConfigRequirements.ts | 3 +- .../typescript/v2/ModelsRequirements.ts | 6 ++ .../typescript/v2/NewThreadModelDefaults.ts | 6 ++ .../schema/typescript/v2/index.ts | 2 + .../src/protocol/v2/config.rs | 17 ++++ .../src/protocol/v2/tests.rs | 1 + codex-rs/app-server/README.md | 2 +- .../request_processors/config_processor.rs | 37 +++++++++ .../app-server/tests/suite/v2/config_rpc.rs | 37 +++++++++ codex-rs/config/src/config_requirements.rs | 81 ++++++++++++++++++- codex-rs/config/src/lib.rs | 2 + .../config/src/requirements_layers/stack.rs | 2 + .../src/requirements_layers/stack_tests.rs | 40 +++++++++ codex-rs/core/src/config/config_tests.rs | 1 + codex-rs/tui/src/debug_config.rs | 2 + 18 files changed, 396 insertions(+), 4 deletions(-) create mode 100644 codex-rs/app-server-protocol/schema/typescript/v2/ModelsRequirements.ts create mode 100644 codex-rs/app-server-protocol/schema/typescript/v2/NewThreadModelDefaults.ts 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 f1d6a1af85..b972de2206 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 @@ -8414,6 +8414,16 @@ "object", "null" ] + }, + "models": { + "anyOf": [ + { + "$ref": "#/definitions/v2/ModelsRequirements" + }, + { + "type": "null" + } + ] } }, "type": "object" @@ -12834,6 +12844,21 @@ "title": "ModelVerificationNotification", "type": "object" }, + "ModelsRequirements": { + "properties": { + "newThread": { + "anyOf": [ + { + "$ref": "#/definitions/v2/NewThreadModelDefaults" + }, + { + "type": "null" + } + ] + } + }, + "type": "object" + }, "MultiAgentMode": { "description": "Controls whether the model receives multi-agent delegation instructions and, when it does, whether it should only spawn sub-agents after an explicit user request or may delegate proactively when doing so would help. `none` leaves the multi-agent tools available without injecting delegation instructions.", "enum": [ @@ -12981,6 +13006,33 @@ ], "type": "string" }, + "NewThreadModelDefaults": { + "properties": { + "model": { + "type": [ + "string", + "null" + ] + }, + "modelReasoningEffort": { + "anyOf": [ + { + "$ref": "#/definitions/v2/ReasoningEffort" + }, + { + "type": "null" + } + ] + }, + "serviceTier": { + "type": [ + "string", + "null" + ] + } + }, + "type": "object" + }, "NonSteerableTurnKind": { "enum": [ "review", 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 9983b79257..4af580e9c3 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 @@ -4654,6 +4654,16 @@ "object", "null" ] + }, + "models": { + "anyOf": [ + { + "$ref": "#/definitions/ModelsRequirements" + }, + { + "type": "null" + } + ] } }, "type": "object" @@ -9238,6 +9248,21 @@ "title": "ModelVerificationNotification", "type": "object" }, + "ModelsRequirements": { + "properties": { + "newThread": { + "anyOf": [ + { + "$ref": "#/definitions/NewThreadModelDefaults" + }, + { + "type": "null" + } + ] + } + }, + "type": "object" + }, "MultiAgentMode": { "description": "Controls whether the model receives multi-agent delegation instructions and, when it does, whether it should only spawn sub-agents after an explicit user request or may delegate proactively when doing so would help. `none` leaves the multi-agent tools available without injecting delegation instructions.", "enum": [ @@ -9385,6 +9410,33 @@ ], "type": "string" }, + "NewThreadModelDefaults": { + "properties": { + "model": { + "type": [ + "string", + "null" + ] + }, + "modelReasoningEffort": { + "anyOf": [ + { + "$ref": "#/definitions/ReasoningEffort" + }, + { + "type": "null" + } + ] + }, + "serviceTier": { + "type": [ + "string", + "null" + ] + } + }, + "type": "object" + }, "NonSteerableTurnKind": { "enum": [ "review", diff --git a/codex-rs/app-server-protocol/schema/json/v2/ConfigRequirementsReadResponse.json b/codex-rs/app-server-protocol/schema/json/v2/ConfigRequirementsReadResponse.json index 3fe6a89787..4e35f7aead 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ConfigRequirementsReadResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ConfigRequirementsReadResponse.json @@ -169,6 +169,16 @@ "object", "null" ] + }, + "models": { + "anyOf": [ + { + "$ref": "#/definitions/ModelsRequirements" + }, + { + "type": "null" + } + ] } }, "type": "object" @@ -362,6 +372,21 @@ ], "type": "object" }, + "ModelsRequirements": { + "properties": { + "newThread": { + "anyOf": [ + { + "$ref": "#/definitions/NewThreadModelDefaults" + }, + { + "type": "null" + } + ] + } + }, + "type": "object" + }, "NetworkDomainPermission": { "enum": [ "allow", @@ -484,6 +509,38 @@ ], "type": "string" }, + "NewThreadModelDefaults": { + "properties": { + "model": { + "type": [ + "string", + "null" + ] + }, + "modelReasoningEffort": { + "anyOf": [ + { + "$ref": "#/definitions/ReasoningEffort" + }, + { + "type": "null" + } + ] + }, + "serviceTier": { + "type": [ + "string", + "null" + ] + } + }, + "type": "object" + }, + "ReasoningEffort": { + "description": "A non-empty reasoning effort value advertised by the model.", + "minLength": 1, + "type": "string" + }, "ResidencyRequirement": { "enum": [ "us" diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ConfigRequirements.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ConfigRequirements.ts index 3eb9a65aa6..e2fd573a32 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/ConfigRequirements.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ConfigRequirements.ts @@ -4,8 +4,9 @@ import type { WebSearchMode } from "../WebSearchMode"; import type { AskForApproval } from "./AskForApproval"; import type { ComputerUseRequirements } from "./ComputerUseRequirements"; +import type { ModelsRequirements } from "./ModelsRequirements"; import type { ResidencyRequirement } from "./ResidencyRequirement"; import type { SandboxMode } from "./SandboxMode"; import type { WindowsSandboxSetupMode } from "./WindowsSandboxSetupMode"; -export type ConfigRequirements = {allowedApprovalPolicies: Array | null, allowedSandboxModes: Array | null, allowedWindowsSandboxImplementations: Array | null, allowedPermissionProfiles: { [key in string]?: boolean } | null, defaultPermissions: string | null, allowedWebSearchModes: Array | null, allowManagedHooksOnly: boolean | null, allowAppshots: boolean | null, allowRemoteControl: boolean | null, computerUse: ComputerUseRequirements | null, featureRequirements: { [key in string]?: boolean } | null, enforceResidency: ResidencyRequirement | null}; +export type ConfigRequirements = {allowedApprovalPolicies: Array | null, allowedSandboxModes: Array | null, allowedWindowsSandboxImplementations: Array | null, allowedPermissionProfiles: { [key in string]?: boolean } | null, defaultPermissions: string | null, allowedWebSearchModes: Array | null, allowManagedHooksOnly: boolean | null, allowAppshots: boolean | null, allowRemoteControl: boolean | null, computerUse: ComputerUseRequirements | null, featureRequirements: { [key in string]?: boolean } | null, enforceResidency: ResidencyRequirement | null, models: ModelsRequirements | null}; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ModelsRequirements.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ModelsRequirements.ts new file mode 100644 index 0000000000..9041fff834 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ModelsRequirements.ts @@ -0,0 +1,6 @@ +// 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. +import type { NewThreadModelDefaults } from "./NewThreadModelDefaults"; + +export type ModelsRequirements = { newThread: NewThreadModelDefaults | null, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/NewThreadModelDefaults.ts b/codex-rs/app-server-protocol/schema/typescript/v2/NewThreadModelDefaults.ts new file mode 100644 index 0000000000..fed8b25eaa --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/NewThreadModelDefaults.ts @@ -0,0 +1,6 @@ +// 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. +import type { ReasoningEffort } from "../ReasoningEffort"; + +export type NewThreadModelDefaults = { model: string | null, modelReasoningEffort: ReasoningEffort | null, serviceTier: string | null, }; 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 2833e6e525..c405082c4c 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts @@ -267,6 +267,7 @@ export type { ModelServiceTier } from "./ModelServiceTier"; export type { ModelUpgradeInfo } from "./ModelUpgradeInfo"; export type { ModelVerification } from "./ModelVerification"; export type { ModelVerificationNotification } from "./ModelVerificationNotification"; +export type { ModelsRequirements } from "./ModelsRequirements"; export type { NetworkAccess } from "./NetworkAccess"; export type { NetworkApprovalContext } from "./NetworkApprovalContext"; export type { NetworkApprovalProtocol } from "./NetworkApprovalProtocol"; @@ -275,6 +276,7 @@ export type { NetworkPolicyAmendment } from "./NetworkPolicyAmendment"; export type { NetworkPolicyRuleAction } from "./NetworkPolicyRuleAction"; export type { NetworkRequirements } from "./NetworkRequirements"; export type { NetworkUnixSocketPermission } from "./NetworkUnixSocketPermission"; +export type { NewThreadModelDefaults } from "./NewThreadModelDefaults"; export type { NonSteerableTurnKind } from "./NonSteerableTurnKind"; export type { OverriddenMetadata } from "./OverriddenMetadata"; export type { PatchApplyStatus } from "./PatchApplyStatus"; diff --git a/codex-rs/app-server-protocol/src/protocol/v2/config.rs b/codex-rs/app-server-protocol/src/protocol/v2/config.rs index 4c9e355b47..f615b25d44 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/config.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/config.rs @@ -390,6 +390,23 @@ pub struct ConfigRequirements { pub enforce_residency: Option, #[experimental("configRequirements/read.network")] pub network: Option, + pub models: Option, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct ModelsRequirements { + pub new_thread: Option, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct NewThreadModelDefaults { + pub model: Option, + pub model_reasoning_effort: Option, + pub service_tier: Option, } #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] diff --git a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs index ddd2afe8f9..ea1e0424dd 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs @@ -1753,6 +1753,7 @@ fn config_requirements_granular_allowed_approval_policy_is_marked_experimental() hooks: None, enforce_residency: None, network: None, + models: None, }); assert_eq!(reason, Some("askForApproval.granular")); diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index d291f5c85b..78f46f5b13 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -242,7 +242,7 @@ Example with notification opt-out: - `externalAgentConfig/import` — apply selected external-agent migration items by passing explicit `migrationItems` with `cwd` (`null` for home) and any `details` returned by detect. Callers may pass `source` to identify the product that initiated the import; omitted or `null` means unspecified. The response acknowledges the synchronous import phase with an `importId`. Expected migration failures are reported as per-item failures rather than JSON-RPC errors, so the server still returns that `importId` and emits `externalAgentConfig/import/completed` with the same ID once all synchronous and background work finishes. The completion notification contains type-level `itemTypeResults` with successes and failures, including raw failure messages for the client to report separately. - `config/value/write` — write a single config key/value to the user's config.toml on disk; dotted paths such as `desktop.someKey` use the same generic write surface. - `config/batchWrite` — apply multiple config edits atomically to the user's config.toml on disk, with optional `reloadUserConfig: true` to hot-reload loaded threads, including multiple `desktop.*` edits. -- `configRequirements/read` — fetch loaded requirements constraints from `requirements.toml` and/or MDM (or `null` if none are configured), including allow-lists (`allowedApprovalPolicies`, `allowedSandboxModes`, `allowedWebSearchModes`), the layered permission-profile allow map (`allowedPermissionProfiles`), the managed permission-profile default (`defaultPermissions`), lifecycle hook lockdown (`allowManagedHooksOnly`), remote-control policy (`allowRemoteControl`; `false` force-disables remote control while `true` or `null` preserves existing behavior), computer use policy (`computerUse`), pinned feature values (`featureRequirements`), managed lifecycle hooks (`hooks`), `enforceResidency`, and `network` constraints such as canonical domain/socket permissions plus `managedAllowedDomainsOnly` and `dangerFullAccessDenylistOnly`. +- `configRequirements/read` — fetch loaded requirements constraints from `requirements.toml` and/or MDM (or `null` if none are configured), including allow-lists (`allowedApprovalPolicies`, `allowedSandboxModes`, `allowedWebSearchModes`), the layered permission-profile allow map (`allowedPermissionProfiles`), the managed permission-profile default (`defaultPermissions`), lifecycle hook lockdown (`allowManagedHooksOnly`), remote-control policy (`allowRemoteControl`; `false` force-disables remote control while `true` or `null` preserves existing behavior), computer use policy (`computerUse`), pinned feature values (`featureRequirements`), managed lifecycle hooks (`hooks`), `enforceResidency`, managed new-thread defaults (`models.newThread.model`, `models.newThread.modelReasoningEffort`, and `models.newThread.serviceTier`), and `network` constraints such as canonical domain/socket permissions plus `managedAllowedDomainsOnly` and `dangerFullAccessDenylistOnly`. ### Example: Start or resume a thread diff --git a/codex-rs/app-server/src/request_processors/config_processor.rs b/codex-rs/app-server/src/request_processors/config_processor.rs index 50bba54009..b6dbb30b08 100644 --- a/codex-rs/app-server/src/request_processors/config_processor.rs +++ b/codex-rs/app-server/src/request_processors/config_processor.rs @@ -24,9 +24,11 @@ use codex_app_server_protocol::ExperimentalFeatureEnablementSetResponse; use codex_app_server_protocol::JSONRPCErrorError; use codex_app_server_protocol::ManagedHooksRequirements; use codex_app_server_protocol::ModelProviderCapabilitiesReadResponse; +use codex_app_server_protocol::ModelsRequirements; use codex_app_server_protocol::NetworkDomainPermission; use codex_app_server_protocol::NetworkRequirements; use codex_app_server_protocol::NetworkUnixSocketPermission; +use codex_app_server_protocol::NewThreadModelDefaults; use codex_app_server_protocol::SandboxMode; use codex_app_server_protocol::WindowsSandboxSetupMode; use codex_config::ConfigRequirementsToml; @@ -375,6 +377,13 @@ fn map_requirements_toml_to_api(requirements: ConfigRequirementsToml) -> ConfigR .enforce_residency .map(map_residency_requirement_to_api), network: requirements.network.map(map_network_requirements_to_api), + models: requirements.models.map(|models| ModelsRequirements { + new_thread: models.new_thread.map(|new_thread| NewThreadModelDefaults { + model: new_thread.model, + model_reasoning_effort: new_thread.model_reasoning_effort, + service_tier: new_thread.service_tier, + }), + }), } } @@ -568,7 +577,10 @@ mod tests { use codex_app_server_protocol::WindowsSandboxSetupMode; use codex_config::ComputerUseRequirementsToml; use codex_config::ConfigRequirementsToml; + use codex_config::ModelsRequirementsToml; + use codex_config::NewThreadModelDefaultsToml; use codex_config::WindowsRequirementsToml; + use codex_protocol::openai_models::ReasoningEffort; use pretty_assertions::assert_eq; use std::collections::BTreeMap; @@ -628,6 +640,31 @@ mod tests { assert_eq!(mapped.allow_remote_control, Some(false)); } + #[test] + fn requirements_api_includes_new_thread_model_defaults() { + let mapped = map_requirements_toml_to_api(ConfigRequirementsToml { + models: Some(ModelsRequirementsToml { + new_thread: Some(NewThreadModelDefaultsToml { + model: Some("gpt-managed".to_string()), + model_reasoning_effort: Some(ReasoningEffort::Medium), + service_tier: Some("fast".to_string()), + }), + }), + ..ConfigRequirementsToml::default() + }); + + let defaults = mapped + .models + .and_then(|models| models.new_thread) + .expect("new-thread defaults"); + assert_eq!(defaults.model.as_deref(), Some("gpt-managed")); + assert_eq!( + defaults.model_reasoning_effort, + Some(ReasoningEffort::Medium) + ); + assert_eq!(defaults.service_tier.as_deref(), Some("fast")); + } + #[test] fn requirements_api_includes_computer_use_requirements() { let mapped = map_requirements_toml_to_api(ConfigRequirementsToml { diff --git a/codex-rs/app-server/tests/suite/v2/config_rpc.rs b/codex-rs/app-server/tests/suite/v2/config_rpc.rs index 81ef145ef0..07b8ccc2e8 100644 --- a/codex-rs/app-server/tests/suite/v2/config_rpc.rs +++ b/codex-rs/app-server/tests/suite/v2/config_rpc.rs @@ -75,6 +75,43 @@ async fn config_requirements_read_includes_allow_remote_control() -> Result<()> Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn config_requirements_read_includes_new_thread_model_defaults() -> Result<()> { + let codex_home = TempDir::new()?; + std::fs::write( + codex_home.path().join("requirements.toml"), + r#" +[models.new_thread] +model = "gpt-managed" +model_reasoning_effort = "medium" +service_tier = "fast" +"#, + )?; + let mut mcp = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp.send_config_requirements_read_request().await?; + let response = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let response: ConfigRequirementsReadResponse = to_response(response)?; + + let defaults = response + .requirements + .and_then(|requirements| requirements.models) + .and_then(|models| models.new_thread) + .expect("managed new-thread defaults"); + assert_eq!(defaults.model.as_deref(), Some("gpt-managed")); + assert_eq!( + defaults.model_reasoning_effort, + Some(ReasoningEffort::Medium) + ); + assert_eq!(defaults.service_tier.as_deref(), Some("fast")); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn config_read_returns_effective_and_layers() -> Result<()> { let codex_home = TempDir::new()?; diff --git a/codex-rs/config/src/config_requirements.rs b/codex-rs/config/src/config_requirements.rs index 199ca4eafa..00af972336 100644 --- a/codex-rs/config/src/config_requirements.rs +++ b/codex-rs/config/src/config_requirements.rs @@ -2,6 +2,7 @@ use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::config_types::SandboxMode; use codex_protocol::config_types::WebSearchMode; use codex_protocol::models::PermissionProfile; +use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::AskForApproval; use codex_utils_absolute_path::AbsolutePathBuf; use serde::Deserialize; @@ -876,9 +877,36 @@ pub struct ConfigRequirementsToml { #[serde(rename = "experimental_network")] pub network: Option, pub permissions: Option, + pub models: Option, pub guardian_policy_config: Option, } +#[derive(Deserialize, Debug, Clone, Default, PartialEq, Eq)] +pub struct ModelsRequirementsToml { + pub new_thread: Option, +} + +impl ModelsRequirementsToml { + fn is_empty(&self) -> bool { + self.new_thread + .as_ref() + .is_none_or(NewThreadModelDefaultsToml::is_empty) + } +} + +#[derive(Deserialize, Debug, Clone, Default, PartialEq, Eq)] +pub struct NewThreadModelDefaultsToml { + pub model: Option, + pub model_reasoning_effort: Option, + pub service_tier: Option, +} + +impl NewThreadModelDefaultsToml { + fn is_empty(&self) -> bool { + self.model.is_none() && self.model_reasoning_effort.is_none() && self.service_tier.is_none() + } +} + #[derive(Deserialize, Debug, Clone, PartialEq)] pub struct RemoteSandboxConfigToml { pub hostname_patterns: Vec, @@ -930,6 +958,7 @@ pub struct ConfigRequirementsWithSources { pub enforce_residency: Option>, pub network: Option>, pub permissions: Option>, + pub models: Option>, pub guardian_policy_config: Option>, } @@ -974,6 +1003,7 @@ impl ConfigRequirementsWithSources { enforce_residency: _, network: _, permissions: _, + models: _, guardian_policy_config: _, } = &other; @@ -1010,6 +1040,7 @@ impl ConfigRequirementsWithSources { enforce_residency, network, permissions, + models, guardian_policy_config, } ); @@ -1046,6 +1077,7 @@ impl ConfigRequirementsWithSources { enforce_residency, network, permissions, + models, guardian_policy_config, } = self; ConfigRequirementsToml { @@ -1071,6 +1103,7 @@ impl ConfigRequirementsWithSources { enforce_residency: enforce_residency.map(|sourced| sourced.value), network: network.map(|sourced| sourced.value), permissions: permissions.map(|sourced| sourced.value), + models: models.map(|sourced| sourced.value), guardian_policy_config: guardian_policy_config.map(|sourced| sourced.value), } } @@ -1183,6 +1216,10 @@ impl ConfigRequirementsToml { && self.enforce_residency.is_none() && self.network.is_none() && self.permissions.is_none() + && self + .models + .as_ref() + .is_none_or(ModelsRequirementsToml::is_empty) && self .guardian_policy_config .as_deref() @@ -1214,8 +1251,9 @@ impl TryFrom for ConfigRequirements { fn try_from(toml: ConfigRequirementsWithSources) -> Result { // Profile catalog selection remains on ConfigRequirementsToml for - // config loading and requirements API projection. The normalized - // constraints below only need the compiled PermissionProfile envelope. + // config loading and requirements API projection. Managed new-thread + // defaults also remain there because they are initialization values, + // not runtime constraints. let ConfigRequirementsWithSources { allowed_approval_policies, allowed_approvals_reviewers, @@ -1238,6 +1276,7 @@ impl TryFrom for ConfigRequirements { enforce_residency, network, permissions, + models: _, guardian_policy_config, } = toml; @@ -1635,6 +1674,7 @@ mod tests { enforce_residency, network, permissions, + models, guardian_policy_config, } = toml; ConfigRequirementsWithSources { @@ -1670,6 +1710,7 @@ mod tests { .map(|value| Sourced::new(value, RequirementSource::Unknown)), network: network.map(|value| Sourced::new(value, RequirementSource::Unknown)), permissions: permissions.map(|value| Sourced::new(value, RequirementSource::Unknown)), + models: models.map(|value| Sourced::new(value, RequirementSource::Unknown)), guardian_policy_config: guardian_policy_config .map(|value| Sourced::new(value, RequirementSource::Unknown)), } @@ -1822,6 +1863,31 @@ mod tests { Ok(()) } + #[test] + fn deserialize_new_thread_model_defaults() -> Result<()> { + let requirements: ConfigRequirementsToml = from_str( + r#" + [models.new_thread] + model = "managed-model" + model_reasoning_effort = "medium" + service_tier = "fast" + "#, + )?; + + assert_eq!( + requirements.models, + Some(ModelsRequirementsToml { + new_thread: Some(NewThreadModelDefaultsToml { + model: Some("managed-model".to_string()), + model_reasoning_effort: Some(ReasoningEffort::Medium), + service_tier: Some("fast".to_string()), + }), + }) + ); + assert!(!requirements.is_empty()); + Ok(()) + } + #[test] fn merge_unset_fields_copies_every_field_and_sets_sources() { let mut target = ConfigRequirementsWithSources::default(); @@ -1844,6 +1910,13 @@ mod tests { let computer_use = ComputerUseRequirementsToml { allow_locked_computer_use: Some(false), }; + let models = ModelsRequirementsToml { + new_thread: Some(NewThreadModelDefaultsToml { + model: Some("managed-model".to_string()), + model_reasoning_effort: Some(ReasoningEffort::Medium), + service_tier: Some("fast".to_string()), + }), + }; let enforce_residency = ResidencyRequirement::Us; let enforce_source = source.clone(); let guardian_policy_config = "Use the company-managed guardian policy.".to_string(); @@ -1873,6 +1946,7 @@ mod tests { enforce_residency: Some(enforce_residency), network: None, permissions: None, + models: Some(models.clone()), guardian_policy_config: Some(guardian_policy_config.clone()), }; @@ -1923,6 +1997,7 @@ mod tests { enforce_residency: Some(Sourced::new(enforce_residency, enforce_source)), network: None, permissions: None, + models: Some(Sourced::new(models, source.clone())), guardian_policy_config: Some(Sourced::new(guardian_policy_config, source)), } ); @@ -1970,6 +2045,7 @@ mod tests { enforce_residency: None, network: None, permissions: None, + models: None, guardian_policy_config: None, } ); @@ -2025,6 +2101,7 @@ mod tests { enforce_residency: None, network: None, permissions: None, + models: None, guardian_policy_config: None, } ); diff --git a/codex-rs/config/src/lib.rs b/codex-rs/config/src/lib.rs index 61e6b6966a..df3e45d76d 100644 --- a/codex-rs/config/src/lib.rs +++ b/codex-rs/config/src/lib.rs @@ -67,12 +67,14 @@ pub use config_requirements::FilesystemDenyReadPattern; pub use config_requirements::MarketplaceAllowedSourceKind; pub use config_requirements::MarketplaceAllowedSourceToml; pub use config_requirements::MarketplaceRequirementsToml; +pub use config_requirements::ModelsRequirementsToml; pub use config_requirements::NetworkConstraints; pub use config_requirements::NetworkDomainPermissionToml; pub use config_requirements::NetworkDomainPermissionsToml; pub use config_requirements::NetworkRequirementsToml; pub use config_requirements::NetworkUnixSocketPermissionToml; pub use config_requirements::NetworkUnixSocketPermissionsToml; +pub use config_requirements::NewThreadModelDefaultsToml; pub use config_requirements::PluginRequirementsToml; pub use config_requirements::RemoteSandboxConfigToml; pub use config_requirements::RequirementSource; diff --git a/codex-rs/config/src/requirements_layers/stack.rs b/codex-rs/config/src/requirements_layers/stack.rs index 3be1d11b5c..9a42650bf3 100644 --- a/codex-rs/config/src/requirements_layers/stack.rs +++ b/codex-rs/config/src/requirements_layers/stack.rs @@ -228,6 +228,7 @@ fn populate_merged_regular_fields_with_sources( enforce_residency, network, permissions, + models, guardian_policy_config, } = requirements; @@ -256,6 +257,7 @@ fn populate_merged_regular_fields_with_sources( set_sourced!(enforce_residency, &["enforce_residency"]); set_sourced!(network, &["experimental_network"]); set_sourced!(permissions, &["permissions"]); + set_sourced!(models, &["models"]); if let Some(guardian_policy_config) = guardian_policy_config.filter(|value| !value.trim().is_empty()) diff --git a/codex-rs/config/src/requirements_layers/stack_tests.rs b/codex-rs/config/src/requirements_layers/stack_tests.rs index a775f68a83..92e6bd8b46 100644 --- a/codex-rs/config/src/requirements_layers/stack_tests.rs +++ b/codex-rs/config/src/requirements_layers/stack_tests.rs @@ -109,6 +109,46 @@ allow_remote_control = false ); } +#[test] +fn new_thread_model_defaults_use_toml_priority() { + let composed = compose(vec![ + layer( + "req_low", + "Low", + r#" +[models.new_thread] +model = "low-priority-model" +model_reasoning_effort = "low" +service_tier = "flex" +"#, + ), + layer( + "req_high", + "High", + r#" +[models.new_thread] +model = "high-priority-model" +model_reasoning_effort = "high" +service_tier = "fast" +"#, + ), + ]) + .expect("compose requirements") + .expect("requirements present"); + + assert_eq!( + composed, + expected_requirements( + r#" +[models.new_thread] +model = "high-priority-model" +model_reasoning_effort = "high" +service_tier = "fast" +"# + ) + ); +} + #[test] fn composition_strategy_applies_to_non_cloud_layers() { let mdm_source = RequirementSource::MdmManagedPreferences { diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index 6f7114928e..c344475541 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -8867,6 +8867,7 @@ async fn test_requirements_web_search_mode_allowlist_does_not_warn_when_unset() enforce_residency: None, network: None, permissions: None, + models: None, guardian_policy_config: None, }; let requirement_source = codex_config::RequirementSource::Unknown; diff --git a/codex-rs/tui/src/debug_config.rs b/codex-rs/tui/src/debug_config.rs index 7a35cc4506..c02d4c8261 100644 --- a/codex-rs/tui/src/debug_config.rs +++ b/codex-rs/tui/src/debug_config.rs @@ -791,6 +791,7 @@ mod tests { enforce_residency: Some(ResidencyRequirement::Us), network: None, permissions: None, + models: None, }; let user_file = if cfg!(windows) { @@ -1149,6 +1150,7 @@ approval_policy = "never" enforce_residency: None, network: None, permissions: None, + models: None, }; let stack = ConfigLayerStack::new(Vec::new(), requirements, requirements_toml) From a938d5f6071f0430b1ed05e381ad9239baeb26f7 Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 26 Jun 2026 19:37:59 +0100 Subject: [PATCH 27/38] Overlap executor skill reads with namespace discovery (#30225) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Why Environment skill discovery needs two independent pieces of information: - plugin namespaces from `plugin.json` files; and - skill metadata from each `SKILL.md` file. Today these happen in sequence. Codex waits for every plugin namespace lookup to finish before it starts reading any skill files. On a remote executor, that creates an avoidable network-latency barrier. ```text before: walk -> namespace lookups -> skill reads -> build catalog after: walk -> namespace lookups ─┐ -> skill reads ───────┴-> build catalog ``` ## What changes - Read and parse skill files without waiting for plugin namespace discovery. - Resolve root and nested plugin namespaces concurrently. - Join both results only when constructing the final qualified skill names. - Keep the existing 64-skill concurrency bound, output ordering, warnings, metadata behavior, and namespace rules. ## Testing The regression test makes plugin manifest lookup wait until a `SKILL.md` read has started. The old serialized pipeline would time out; the new pipeline completes and still returns the correctly namespaced skill. `just test -p codex-core-skills` passes all 111 tests. ## Out of scope This does not add an exec-server endpoint, batch filesystem calls, or reduce the number of files transferred. A frontmatter-only read or server-side skill catalog can remain a separate follow-up if benchmarks show that transferred bytes are the next bottleneck. --- .../core-skills/src/loader/environment.rs | 117 +++++++++++------- .../core-skills/tests/environment_loader.rs | 90 +++++++++++++- 2 files changed, 158 insertions(+), 49 deletions(-) diff --git a/codex-rs/core-skills/src/loader/environment.rs b/codex-rs/core-skills/src/loader/environment.rs index 96caa28eea..049b500585 100644 --- a/codex-rs/core-skills/src/loader/environment.rs +++ b/codex-rs/core-skills/src/loader/environment.rs @@ -45,6 +45,15 @@ struct DiscoveredEnvironmentSkill { metadata: SkillMetadataDiscovery, } +struct ParsedEnvironmentSkill { + path_to_skills_md: PathUri, + base_name: String, + description: String, + short_description: Option, + dependencies: Option, + policy: Option, +} + enum SkillMetadataDiscovery { Present(PathUri), Absent, @@ -81,11 +90,12 @@ impl EnvironmentSkillMetadata { None => true, } } +} - async fn parse( +impl ParsedEnvironmentSkill { + async fn load( file_system: &dyn ExecutorFileSystem, skill: &DiscoveredEnvironmentSkill, - plugin_namespace: Option<&str>, ) -> Result { let (contents, discovered_metadata) = match &skill.metadata { SkillMetadataDiscovery::Present(metadata_path) => { @@ -106,11 +116,6 @@ impl EnvironmentSkillMetadata { short_description, } = parse_skill_frontmatter_metadata_inner(&contents, || default_skill_name(&skill.path)) .map_err(|err| err.to_string())?; - let name = plugin_namespace - .map(|namespace| format!("{namespace}:{base_name}")) - .unwrap_or(base_name); - validate_len(&name, MAX_QUALIFIED_NAME_LEN, "qualified name") - .map_err(|err| err.to_string())?; let (dependencies, policy) = match &skill.metadata { SkillMetadataDiscovery::Present(_) | SkillMetadataDiscovery::Absent => { discovered_metadata @@ -122,7 +127,7 @@ impl EnvironmentSkillMetadata { Ok(Self { path_to_skills_md: skill.path.clone(), - name, + base_name, description, short_description, dependencies, @@ -245,38 +250,55 @@ pub async fn load_environment_skills_from_root( } let namespace_roots = discovery.namespace_roots; - let namespace_lookups = join_all(namespace_roots.iter().map(|namespace_root| async { - ( - namespace_root.clone(), - plugin_namespace_for_skill_uri(file_system, namespace_root).await, - ) - })) - .await; - let plugin_lookups = join_all( - discovery - .plugin_roots - .iter() - .filter(|plugin_root| skill_ancestors.contains(*plugin_root)) - .filter(|plugin_root| !namespace_roots.contains(*plugin_root)) - .map(|plugin_root| async { - ( - plugin_root.clone(), - plugin_namespace_for_root_uri(file_system, plugin_root).await, - ) - }), - ) - .await; - let plugin_namespaces = namespace_lookups - .into_iter() - .chain(plugin_lookups) - .filter_map(|(plugin_root, namespace)| namespace.map(|namespace| (plugin_root, namespace))) - .collect::>(); + let plugin_namespaces = async { + let namespace_lookups = join_all(namespace_roots.iter().map(|namespace_root| async { + ( + namespace_root.clone(), + plugin_namespace_for_skill_uri(file_system, namespace_root).await, + ) + })); + let plugin_lookups = join_all( + discovery + .plugin_roots + .iter() + .filter(|plugin_root| skill_ancestors.contains(*plugin_root)) + .filter(|plugin_root| !namespace_roots.contains(*plugin_root)) + .map(|plugin_root| async { + ( + plugin_root.clone(), + plugin_namespace_for_root_uri(file_system, plugin_root).await, + ) + }), + ); + let (namespace_lookups, plugin_lookups) = tokio::join!(namespace_lookups, plugin_lookups); + namespace_lookups + .into_iter() + .chain(plugin_lookups) + .filter_map(|(plugin_root, namespace)| { + namespace.map(|namespace| (plugin_root, namespace)) + }) + .collect::>() + }; // Remote executors can multiplex these independent per-skill reads, so polling a bounded // number together allows the I/O for each skill and its metadata to happen concurrently. let skill_results = futures::stream::iter(discovery.skills) .map(|skill| { - let mut ancestor = skill.path.parent(); + let path = skill.path.clone(); + async move { + ( + path, + ParsedEnvironmentSkill::load(file_system, &skill).await, + ) + } + }) + .buffered(MAX_CONCURRENT_SKILL_LOADS) + .collect::>(); + let (plugin_namespaces, skill_results) = tokio::join!(plugin_namespaces, skill_results); + + for (path, result) in skill_results { + let result = result.and_then(|skill| { + let mut ancestor = skill.path_to_skills_md.parent(); let plugin_namespace = loop { let Some(current) = ancestor else { break None; @@ -286,18 +308,21 @@ pub async fn load_environment_skills_from_root( } ancestor = current.parent(); }; - let path = skill.path.clone(); - async move { - let result = - EnvironmentSkillMetadata::parse(file_system, &skill, plugin_namespace).await; - (path, result) - } - }) - .buffered(MAX_CONCURRENT_SKILL_LOADS) - .collect::>() - .await; + let name = plugin_namespace + .map(|namespace| format!("{namespace}:{}", skill.base_name)) + .unwrap_or(skill.base_name); + validate_len(&name, MAX_QUALIFIED_NAME_LEN, "qualified name") + .map_err(|err| err.to_string())?; - for (path, result) in skill_results { + Ok(EnvironmentSkillMetadata { + path_to_skills_md: skill.path_to_skills_md, + name, + description: skill.description, + short_description: skill.short_description, + dependencies: skill.dependencies, + policy: skill.policy, + }) + }); match result { Ok(skill) if skill.matches_product_restriction(restriction_product) => { outcome.skills.push(skill); diff --git a/codex-rs/core-skills/tests/environment_loader.rs b/codex-rs/core-skills/tests/environment_loader.rs index c17802bedb..74efc6e1a3 100644 --- a/codex-rs/core-skills/tests/environment_loader.rs +++ b/codex-rs/core-skills/tests/environment_loader.rs @@ -1,7 +1,9 @@ use std::fs; use std::sync::Mutex; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use std::time::Duration; use codex_core_skills::loader::EnvironmentSkillMetadata; use codex_core_skills::loader::load_environment_skills_from_root; @@ -20,12 +22,22 @@ use codex_exec_server::WalkOutcome; use codex_utils_path_uri::PathUri; use pretty_assertions::assert_eq; use tempfile::tempdir; +use tokio::sync::Notify; + +#[derive(Clone, Copy)] +enum ManifestMetadataBehavior { + Immediate, + WaitForSkillRead, +} struct RecordingFileSystem<'a> { inner: &'a dyn ExecutorFileSystem, read_files: Mutex>, metadata_files: Mutex>, walks: AtomicUsize, + manifest_metadata_behavior: ManifestMetadataBehavior, + skill_read_started: AtomicBool, + skill_read_started_notify: Notify, } #[derive(Debug, PartialEq, Eq)] @@ -36,12 +48,18 @@ struct FileSystemCalls { } impl<'a> RecordingFileSystem<'a> { - fn new(inner: &'a dyn ExecutorFileSystem) -> Self { + fn new( + inner: &'a dyn ExecutorFileSystem, + manifest_metadata_behavior: ManifestMetadataBehavior, + ) -> Self { Self { inner, read_files: Mutex::new(Vec::new()), metadata_files: Mutex::new(Vec::new()), walks: AtomicUsize::new(0), + manifest_metadata_behavior, + skill_read_started: AtomicBool::new(false), + skill_read_started_notify: Notify::new(), } } @@ -84,6 +102,10 @@ impl ExecutorFileSystem for RecordingFileSystem<'_> { .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .push(path.clone()); + if path.basename().as_deref() == Some("SKILL.md") { + self.skill_read_started.store(true, Ordering::Release); + self.skill_read_started_notify.notify_waiters(); + } self.inner.read_file(path, sandbox) } @@ -122,6 +144,22 @@ impl ExecutorFileSystem for RecordingFileSystem<'_> { .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .push(path.clone()); + if matches!( + self.manifest_metadata_behavior, + ManifestMetadataBehavior::WaitForSkillRead + ) && path.basename().as_deref() == Some("plugin.json") + { + return Box::pin(async move { + loop { + let notified = self.skill_read_started_notify.notified(); + if self.skill_read_started.load(Ordering::Acquire) { + break; + } + notified.await; + } + self.inner.get_metadata(path, sandbox).await + }); + } self.inner.get_metadata(path, sandbox) } @@ -201,7 +239,8 @@ async fn loads_nearest_plugin_namespaces_without_reading_unused_sibling_manifest .expect("skill"); } - let file_system = RecordingFileSystem::new(LOCAL_FS.as_ref()); + let file_system = + RecordingFileSystem::new(LOCAL_FS.as_ref(), ManifestMetadataBehavior::Immediate); let root_uri = PathUri::from_host_native_path(root.path()).expect("root URI"); let outcome = load_environment_skills_from_root( &file_system, @@ -280,7 +319,8 @@ async fn reuses_walk_inventory_for_missing_skill_metadata() { skill_paths.push(skill_path); } - let file_system = RecordingFileSystem::new(LOCAL_FS.as_ref()); + let file_system = + RecordingFileSystem::new(LOCAL_FS.as_ref(), ManifestMetadataBehavior::Immediate); let root_uri = PathUri::from_host_native_path(root.path()).expect("root URI"); let outcome = load_environment_skills_from_root( &file_system, @@ -327,3 +367,47 @@ async fn reuses_walk_inventory_for_missing_skill_metadata() { } ); } + +#[tokio::test] +async fn reads_skill_files_while_resolving_plugin_namespaces() { + let root = tempdir().expect("tempdir"); + let manifest_path = root.path().join(".codex-plugin/plugin.json"); + fs::create_dir_all(manifest_path.parent().expect("manifest parent")).expect("manifest dir"); + fs::write(&manifest_path, r#"{"name":"parallel"}"#).expect("manifest"); + let skill_path = root.path().join("demo/SKILL.md"); + fs::create_dir_all(skill_path.parent().expect("skill parent")).expect("skill dir"); + fs::write( + &skill_path, + "---\nname: demo\ndescription: demo skill.\n---\n", + ) + .expect("skill"); + + let file_system = RecordingFileSystem::new( + LOCAL_FS.as_ref(), + ManifestMetadataBehavior::WaitForSkillRead, + ); + let root_uri = PathUri::from_host_native_path(root.path()).expect("root URI"); + let outcome = tokio::time::timeout( + Duration::from_secs(5), + load_environment_skills_from_root( + &file_system, + &root_uri, + /*restriction_product*/ None, + ), + ) + .await + .expect("skill reads should start before namespace resolution finishes"); + + assert_eq!(outcome.warnings, Vec::::new()); + assert_eq!( + outcome.skills, + vec![EnvironmentSkillMetadata { + path_to_skills_md: PathUri::from_host_native_path(skill_path).unwrap(), + name: "parallel:demo".to_string(), + description: "demo skill.".to_string(), + short_description: None, + dependencies: None, + policy: None, + }] + ); +} From 79a8ffdbf7ca08820e0b44cab0261eb61ff72465 Mon Sep 17 00:00:00 2001 From: Charles Du Date: Fri, 26 Jun 2026 12:17:26 -0700 Subject: [PATCH 28/38] [codex] allow AGENTS.md and skills to authorize delegation (#30274) Prompt update of MAv2 to include agents.md and skills more explicitly should mimic: https://github.com/openai/codex/pull/27919 --- codex-rs/app-server/tests/suite/v2/turn_start.rs | 9 ++++++--- .../core/src/context/multi_agent_mode_instructions.rs | 2 +- codex-rs/core/src/tools/handlers/multi_agents_spec.rs | 2 +- codex-rs/core/tests/suite/multi_agent_mode.rs | 2 +- codex-rs/core/tests/suite/search_tool.rs | 2 +- codex-rs/core/tests/suite/spawn_agent_description.rs | 2 +- 6 files changed, 11 insertions(+), 8 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/turn_start.rs b/codex-rs/app-server/tests/suite/v2/turn_start.rs index 528292ae3f..b42a645a33 100644 --- a/codex-rs/app-server/tests/suite/v2/turn_start.rs +++ b/codex-rs/app-server/tests/suite/v2/turn_start.rs @@ -1817,7 +1817,9 @@ async fn turn_start_ignores_deprecated_multi_agent_mode() -> Result<()> { .single_request() .message_input_texts("developer"); assert!(developer_texts.iter().any(|text| { - text.contains("Do not spawn sub-agents unless the user explicitly asks for sub-agents") + text.contains( + "Do not spawn sub-agents unless the user or applicable AGENTS.md/skill instructions explicitly ask for sub-agents", + ) })); assert!( !developer_texts @@ -1898,8 +1900,9 @@ async fn thread_start_ignores_deprecated_multi_agent_mode() -> Result<()> { .message_input_texts("developer"); assert!(developer_texts.iter().any(|text| { text.contains(MULTI_AGENT_MODE_OPEN_TAG) - && text - .contains("Do not spawn sub-agents unless the user explicitly asks for sub-agents") + && text.contains( + "Do not spawn sub-agents unless the user or applicable AGENTS.md/skill instructions explicitly ask for sub-agents", + ) })); assert!( !developer_texts diff --git a/codex-rs/core/src/context/multi_agent_mode_instructions.rs b/codex-rs/core/src/context/multi_agent_mode_instructions.rs index fc4debe5a2..ff471967ca 100644 --- a/codex-rs/core/src/context/multi_agent_mode_instructions.rs +++ b/codex-rs/core/src/context/multi_agent_mode_instructions.rs @@ -3,7 +3,7 @@ use codex_protocol::config_types::MultiAgentMode; use codex_protocol::protocol::MULTI_AGENT_MODE_CLOSE_TAG; use codex_protocol::protocol::MULTI_AGENT_MODE_OPEN_TAG; -const EXPLICIT_REQUEST_ONLY_MULTI_AGENT_MODE_TEXT: &str = "Do not spawn sub-agents unless the user explicitly asks for sub-agents, delegation, or parallel agent work."; +const EXPLICIT_REQUEST_ONLY_MULTI_AGENT_MODE_TEXT: &str = "Do not spawn sub-agents unless the user or applicable AGENTS.md/skill instructions explicitly ask for sub-agents, delegation, or parallel agent work."; const NO_MULTI_AGENT_MODE_TEXT: &str = "Multi-agent delegation mode instructions are inactive. Any earlier multi-agent mode developer message no longer applies."; const PROACTIVE_MULTI_AGENT_MODE_TEXT: &str = "Proactive multi-agent delegation is active. Any earlier instruction requiring an explicit user request before spawning sub-agents no longer applies. Use sub-agents when parallel work would materially improve speed or quality. This mode remains active until a later multi-agent mode developer message changes it."; diff --git a/codex-rs/core/src/tools/handlers/multi_agents_spec.rs b/codex-rs/core/src/tools/handlers/multi_agents_spec.rs index 08f8ca071f..857a4dea9e 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_spec.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_spec.rs @@ -673,7 +673,7 @@ fn spawn_agent_tool_description( {tool_description} This spawn_agent tool provides you access to sub-agents that inherit your current model by default. Do not set the `model` field unless the user explicitly asks for a different model or there is a clear task-specific reason. You should follow the rules and guidelines below to use this tool. -Do not spawn sub-agents unless the user explicitly asks for sub-agents, delegation, or parallel agent work. +Do not spawn sub-agents unless the user or applicable AGENTS.md/skill instructions explicitly ask for sub-agents, delegation, or parallel agent work. {agent_role_usage_hint} ### Designing delegated subtasks diff --git a/codex-rs/core/tests/suite/multi_agent_mode.rs b/codex-rs/core/tests/suite/multi_agent_mode.rs index 451b9cd38c..fd177e7ed9 100644 --- a/codex-rs/core/tests/suite/multi_agent_mode.rs +++ b/codex-rs/core/tests/suite/multi_agent_mode.rs @@ -21,7 +21,7 @@ use core_test_support::wait_for_event; use pretty_assertions::assert_eq; use serde_json::Value; -const NO_SPAWN_TEXT: &str = "Do not spawn sub-agents unless the user explicitly asks for sub-agents, delegation, or parallel agent work."; +const NO_SPAWN_TEXT: &str = "Do not spawn sub-agents unless the user or applicable AGENTS.md/skill instructions explicitly ask for sub-agents, delegation, or parallel agent work."; const PROACTIVE_TEXT: &str = "Proactive multi-agent delegation is active."; fn add_ultra_reasoning(model_info: &mut ModelInfo) { diff --git a/codex-rs/core/tests/suite/search_tool.rs b/codex-rs/core/tests/suite/search_tool.rs index afb2e0d085..c1bc9cc1fa 100644 --- a/codex-rs/core/tests/suite/search_tool.rs +++ b/codex-rs/core/tests/suite/search_tool.rs @@ -858,7 +858,7 @@ async fn tool_search_returns_deferred_v1_multi_agent_tools() -> Result<()> { .and_then(Value::as_str) .expect("spawn_agent description should be present"); assert!(description.contains( - "Do not spawn sub-agents unless the user explicitly asks for sub-agents, delegation, or parallel agent work." + "Do not spawn sub-agents unless the user or applicable AGENTS.md/skill instructions explicitly ask for sub-agents, delegation, or parallel agent work." )); assert!(description.contains("### Designing delegated subtasks")); assert!(!description.contains("### When to delegate vs. do the subtask yourself")); diff --git a/codex-rs/core/tests/suite/spawn_agent_description.rs b/codex-rs/core/tests/suite/spawn_agent_description.rs index fb493c1b14..2cf1cee2f0 100644 --- a/codex-rs/core/tests/suite/spawn_agent_description.rs +++ b/codex-rs/core/tests/suite/spawn_agent_description.rs @@ -213,7 +213,7 @@ async fn spawn_agent_description_lists_visible_models_and_reasoning_efforts() -> ); assert!( description.contains( - "Do not spawn sub-agents unless the user explicitly asks for sub-agents, delegation, or parallel agent work." + "Do not spawn sub-agents unless the user or applicable AGENTS.md/skill instructions explicitly ask for sub-agents, delegation, or parallel agent work." ), "expected explicit authorization rule in spawn_agent description: {description:?}" ); From cf36c688b3d5c9122e818a93d272d4024c8a2bae Mon Sep 17 00:00:00 2001 From: hefuc-oai Date: Fri, 26 Jun 2026 12:27:31 -0700 Subject: [PATCH 29/38] [codex] Use managed defaults for TUI threads (#30147) ## Why #29683 exposes managed defaults for new-thread model settings through `configRequirements/read` without applying them server-wide. The TUI is an app-server client, so it should explicitly consume those defaults when it creates a fresh thread. This lets plain `codex` start on the managed model while preserving the existing ability to change model settings within the thread. ## What changed - Read `requirements.models.newThread` during TUI app-server bootstrap. - Apply the managed model, reasoning effort, and service tier to the initial fresh thread and subsequent `/new` or `/clear` threads. - Keep explicit launch overrides above the managed defaults. - Normalize the managed `fast` service tier to the `priority` request value. - Leave resumed and forked threads unchanged. The application logic lives in a small TUI-only module; app-server `thread/start` behavior remains unchanged for other clients. ## User experience - Plain `codex` starts with the managed new-thread settings. - A user can still change settings with `/model` or the existing service-tier controls. - Starting another fresh thread reapplies the managed defaults. - Explicit launch choices such as `codex -m ` continue to win. ## Validation - `just test -p codex-tui managed_new_thread_defaults` - `just fix -p codex-tui` Depends on #29683. --- codex-rs/tui/src/app.rs | 14 ++- codex-rs/tui/src/app/session_lifecycle.rs | 8 +- codex-rs/tui/src/app_server_session.rs | 23 ++++ codex-rs/tui/src/lib.rs | 1 + .../tui/src/managed_new_thread_defaults.rs | 51 ++++++++ .../src/managed_new_thread_defaults_tests.rs | 110 ++++++++++++++++++ 6 files changed, 205 insertions(+), 2 deletions(-) create mode 100644 codex-rs/tui/src/managed_new_thread_defaults.rs create mode 100644 codex-rs/tui/src/managed_new_thread_defaults_tests.rs diff --git a/codex-rs/tui/src/app.rs b/codex-rs/tui/src/app.rs index b7f2b9df6a..554df264e2 100644 --- a/codex-rs/tui/src/app.rs +++ b/codex-rs/tui/src/app.rs @@ -51,6 +51,7 @@ use crate::legacy_core::config::ConfigBuilder; use crate::legacy_core::config::ConfigOverrides; use crate::legacy_core::config::PermissionProfileSnapshot; use crate::legacy_core::config::edit::ConfigEditsBuilder; +use crate::managed_new_thread_defaults::apply_managed_new_thread_defaults; use crate::model_catalog::ModelCatalog; use crate::model_migration::ModelMigrationOutcome; use crate::model_migration::migration_copy_for_models; @@ -794,7 +795,18 @@ impl App { None => app_server.bootstrap(&config).await?, }; let bootstrap_ms = bootstrap.duration.as_millis(); - let mut model = bootstrap.default_model; + if matches!( + &session_selection, + SessionSelection::StartFresh | SessionSelection::Exit + ) { + apply_managed_new_thread_defaults( + &mut config, + app_server.managed_new_thread_defaults(), + &cli_kv_overrides, + &harness_overrides, + ); + } + let mut model = config.model.clone().unwrap_or(bootstrap.default_model); let available_models = bootstrap.available_models; let remote_connection = crate::status::remote_connection::remote_connection_status_value( &app_server_target, diff --git a/codex-rs/tui/src/app/session_lifecycle.rs b/codex-rs/tui/src/app/session_lifecycle.rs index a961658c80..b4bbbd231a 100644 --- a/codex-rs/tui/src/app/session_lifecycle.rs +++ b/codex-rs/tui/src/app/session_lifecycle.rs @@ -539,7 +539,13 @@ impl App { self.refresh_in_memory_config_from_disk_best_effort("starting a new thread") .await; let model = self.chat_widget.current_model().to_string(); - let config = self.fresh_session_config(); + let mut config = self.fresh_session_config(); + apply_managed_new_thread_defaults( + &mut config, + app_server.managed_new_thread_defaults(), + &self.cli_kv_overrides, + &self.harness_overrides, + ); let summary = session_summary( self.chat_widget.token_usage(), self.chat_widget.thread_id(), diff --git a/codex-rs/tui/src/app_server_session.rs b/codex-rs/tui/src/app_server_session.rs index 0405bb75ab..aeca17091e 100644 --- a/codex-rs/tui/src/app_server_session.rs +++ b/codex-rs/tui/src/app_server_session.rs @@ -24,6 +24,7 @@ use codex_app_server_protocol::AskForApproval; use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::ClientRequest; use codex_app_server_protocol::ConfigBatchWriteParams; +use codex_app_server_protocol::ConfigRequirementsReadResponse; use codex_app_server_protocol::ConfigWriteResponse; use codex_app_server_protocol::ExternalAgentConfigDetectParams; use codex_app_server_protocol::ExternalAgentConfigDetectResponse; @@ -39,6 +40,7 @@ use codex_app_server_protocol::MemoryResetResponse; use codex_app_server_protocol::Model as ApiModel; use codex_app_server_protocol::ModelListParams; use codex_app_server_protocol::ModelListResponse; +use codex_app_server_protocol::NewThreadModelDefaults; use codex_app_server_protocol::RateLimitSnapshot; use codex_app_server_protocol::RequestId; use codex_app_server_protocol::ReviewDelivery; @@ -178,6 +180,7 @@ pub(crate) struct AppServerSession { thread_settings_update_supported: bool, default_model: Option, available_models: Vec, + managed_new_thread_defaults: Option, external_agent_config_import_completion_pending: AtomicBool, } @@ -222,6 +225,7 @@ impl AppServerSession { thread_settings_update_supported: true, default_model: None, available_models: Vec::new(), + managed_new_thread_defaults: None, external_agent_config_import_completion_pending: AtomicBool::new(false), } } @@ -260,6 +264,21 @@ impl AppServerSession { pub(crate) async fn bootstrap(&mut self, config: &Config) -> Result { let started_at = Instant::now(); let account = self.read_account().await?; + let requirements_request_id = self.next_request_id(); + let requirements: ConfigRequirementsReadResponse = self + .client + .request_typed(ClientRequest::ConfigRequirementsRead { + request_id: requirements_request_id, + params: None, + }) + .await + .map_err(|err| { + bootstrap_request_error("configRequirements/read failed during TUI bootstrap", err) + })?; + self.managed_new_thread_defaults = requirements + .requirements + .and_then(|requirements| requirements.models) + .and_then(|models| models.new_thread); let model_request_id = self.next_request_id(); let models: ModelListResponse = self .client @@ -350,6 +369,10 @@ impl AppServerSession { }) } + pub(crate) fn managed_new_thread_defaults(&self) -> Option<&NewThreadModelDefaults> { + self.managed_new_thread_defaults.as_ref() + } + /// Fetches the current account info without refreshing the auth token. /// /// Used by both `bootstrap` (to populate the initial UI) and `get_login_status` diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 897f1ec83f..96f442a88c 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -137,6 +137,7 @@ mod line_truncation; pub(crate) mod live_wrap; pub use live_wrap::RowBuilder; mod local_chatgpt_auth; +mod managed_new_thread_defaults; mod markdown; mod markdown_render; mod markdown_stream; diff --git a/codex-rs/tui/src/managed_new_thread_defaults.rs b/codex-rs/tui/src/managed_new_thread_defaults.rs new file mode 100644 index 0000000000..ad9ad30b8e --- /dev/null +++ b/codex-rs/tui/src/managed_new_thread_defaults.rs @@ -0,0 +1,51 @@ +use crate::legacy_core::config::Config; +use crate::legacy_core::config::ConfigOverrides; +use codex_app_server_protocol::NewThreadModelDefaults; +use codex_protocol::config_types::ServiceTier; +use toml::Value as TomlValue; + +pub(crate) fn apply_managed_new_thread_defaults( + config: &mut Config, + defaults: Option<&NewThreadModelDefaults>, + cli_kv_overrides: &[(String, TomlValue)], + harness_overrides: &ConfigOverrides, +) { + let Some(defaults) = defaults else { + return; + }; + // Managed values are defaults rather than enforcement. Preserve explicit launch choices from + // dedicated flags such as `-m` (`harness_overrides`) and generic `-c key=value` settings + // (`cli_kv_overrides`), then fill only the fields that were not selected for this invocation. + // Model and reasoning effort are a compatibility-sensitive pair, so an explicit override of + // either opts out of both managed values. For example, `codex -m gpt-5.4` keeps that model and + // its existing/default effort, while `-c model_reasoning_effort=low` does not switch to the + // managed model. Service tier remains independent and is resolved against the selected model + // before the thread starts. + let has_cli_override = |key: &str| cli_kv_overrides.iter().any(|(path, _value)| path == key); + let has_explicit_model_settings = harness_overrides.model.is_some() + || has_cli_override("model") + || has_cli_override("model_reasoning_effort"); + + if !has_explicit_model_settings && let Some(model) = defaults.model.as_ref() { + config.model = Some(model.clone()); + } + if !has_explicit_model_settings + && let Some(reasoning_effort) = defaults.model_reasoning_effort.as_ref() + { + config.model_reasoning_effort = Some(reasoning_effort.clone()); + } + if harness_overrides.service_tier.is_none() + && !has_cli_override("service_tier") + && let Some(service_tier) = defaults.service_tier.as_ref() + { + config.service_tier = Some( + ServiceTier::from_request_value(service_tier) + .map(|tier| tier.request_value().to_string()) + .unwrap_or_else(|| service_tier.clone()), + ); + } +} + +#[cfg(test)] +#[path = "managed_new_thread_defaults_tests.rs"] +mod tests; diff --git a/codex-rs/tui/src/managed_new_thread_defaults_tests.rs b/codex-rs/tui/src/managed_new_thread_defaults_tests.rs new file mode 100644 index 0000000000..d73d6459d6 --- /dev/null +++ b/codex-rs/tui/src/managed_new_thread_defaults_tests.rs @@ -0,0 +1,110 @@ +use super::*; +use crate::legacy_core::config::ConfigBuilder; +use codex_protocol::openai_models::ReasoningEffort; +use pretty_assertions::assert_eq; + +async fn test_config() -> Config { + let codex_home = tempfile::tempdir().expect("tempdir").keep(); + ConfigBuilder::default() + .codex_home(codex_home) + .build() + .await + .expect("config") +} + +fn defaults() -> NewThreadModelDefaults { + NewThreadModelDefaults { + model: Some("managed-model".to_string()), + model_reasoning_effort: Some(ReasoningEffort::High), + service_tier: Some("fast".to_string()), + } +} + +#[tokio::test] +async fn applies_managed_defaults_to_a_new_thread_config() { + let mut actual = test_config().await; + actual.model = Some("configured-model".to_string()); + actual.model_reasoning_effort = Some(ReasoningEffort::Low); + actual.service_tier = Some("flex".to_string()); + let mut expected = actual.clone(); + expected.model = Some("managed-model".to_string()); + expected.model_reasoning_effort = Some(ReasoningEffort::High); + expected.service_tier = Some(ServiceTier::Fast.request_value().to_string()); + + apply_managed_new_thread_defaults( + &mut actual, + Some(&defaults()), + &[], + &ConfigOverrides::default(), + ); + + assert_eq!(actual, expected); +} + +#[tokio::test] +async fn explicit_model_skips_managed_model_and_reasoning_effort() { + let mut actual = test_config().await; + actual.model = Some("explicit-model".to_string()); + actual.model_reasoning_effort = None; + actual.service_tier = Some("flex".to_string()); + let mut expected = actual.clone(); + expected.service_tier = Some(ServiceTier::Fast.request_value().to_string()); + let harness_overrides = ConfigOverrides { + model: Some("explicit-model".to_string()), + ..ConfigOverrides::default() + }; + + apply_managed_new_thread_defaults(&mut actual, Some(&defaults()), &[], &harness_overrides); + + assert_eq!(actual, expected); +} + +#[tokio::test] +async fn explicit_reasoning_effort_skips_managed_model_and_reasoning_effort() { + let mut actual = test_config().await; + actual.model = Some("configured-model".to_string()); + actual.model_reasoning_effort = Some(ReasoningEffort::Low); + actual.service_tier = Some("flex".to_string()); + let mut expected = actual.clone(); + expected.service_tier = Some(ServiceTier::Fast.request_value().to_string()); + let cli_kv_overrides = vec![( + "model_reasoning_effort".to_string(), + TomlValue::String("low".to_string()), + )]; + + apply_managed_new_thread_defaults( + &mut actual, + Some(&defaults()), + &cli_kv_overrides, + &ConfigOverrides::default(), + ); + + assert_eq!(actual, expected); +} + +#[tokio::test] +async fn explicit_launch_overrides_take_precedence() { + let mut actual = test_config().await; + actual.model = Some("explicit-model".to_string()); + actual.model_reasoning_effort = Some(ReasoningEffort::Low); + actual.service_tier = Some("flex".to_string()); + let expected = actual.clone(); + let cli_kv_overrides = vec![( + "model_reasoning_effort".to_string(), + TomlValue::String("low".to_string()), + )]; + let harness_overrides = ConfigOverrides { + model: Some("explicit-model".to_string()), + service_tier: Some(Some("flex".to_string())), + ..ConfigOverrides::default() + }; + + apply_managed_new_thread_defaults( + &mut actual, + Some(&defaults()), + &cli_kv_overrides, + &harness_overrides, + ); + + assert_eq!(actual, expected); +} From 812cd2bb57ddb3a04a10d59495ac0212d2eb52a3 Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Fri, 26 Jun 2026 12:32:31 -0700 Subject: [PATCH 30/38] ensure thread.history_mode is immutable (#30261) ## Description This PR makes `thread.history_mode` immutable after the thread's canonical first `SessionMeta` has been written. Later same-thread `SessionMeta` lines are compatibility metadata writes, not a new thread definition. Without this, an older binary could append a `SessionMeta` that omits `history_mode`; when a newer binary replays it, serde defaults that missing field to `legacy` and SQLite could downgrade a paginated thread. ## Why `history_mode` is the persisted thread storage contract. Paginated-thread fail-closed behavior and SQLite memory filtering depend on it staying aligned with canonical rollout metadata, especially when multiple Codex binary versions can touch the same local rollout. ## What changed - Stop generic rollout metadata replay from overwriting `history_mode` from later `SessionMeta` items. - Remove `history_mode` from `ThreadMetadataPatch`, so mutable metadata sync and app-server metadata updates cannot rewrite it. - When local metadata sync has to recreate a missing SQLite row, recover `history_mode` from the rollout's canonical first `SessionMeta` instead of from a mutable patch. - Keep the in-memory thread store using the created thread's canonical `history_mode` instead of metadata patches. - Fill the one remaining core test `CreateThreadParams` initializer with the new `history_mode` field; Bazel CI caught this after the parent history-mode PR landed. ## Validation - `just fmt` - `just test -p codex-thread-store` - `just test -p codex-state session_meta_does_not_set_model_or_reasoning_effort` --- codex-rs/state/src/extract.rs | 7 +- codex-rs/thread-store/src/in_memory.rs | 41 ++--- .../thread-store/src/local/live_writer.rs | 10 +- codex-rs/thread-store/src/local/mod.rs | 59 ++++++- .../src/local/update_thread_metadata.rs | 166 +++++++++++++++--- .../thread-store/src/thread_metadata_sync.rs | 3 - codex-rs/thread-store/src/types.rs | 6 - 7 files changed, 226 insertions(+), 66 deletions(-) diff --git a/codex-rs/state/src/extract.rs b/codex-rs/state/src/extract.rs index f944e3e389..48e308f627 100644 --- a/codex-rs/state/src/extract.rs +++ b/codex-rs/state/src/extract.rs @@ -56,7 +56,7 @@ fn apply_session_meta_from_item(metadata: &mut ThreadMetadata, meta_line: &Sessi } metadata.id = meta_line.meta.id; metadata.source = enum_to_string(&meta_line.meta.source); - metadata.history_mode = meta_line.meta.history_mode; + // Later SessionMeta lines do not redefine the canonical history_mode. metadata.thread_source = meta_line.meta.thread_source.clone(); metadata.agent_nickname = meta_line.meta.agent_nickname.clone(); metadata.agent_role = meta_line.meta.agent_role.clone(); @@ -178,6 +178,7 @@ mod tests { use codex_protocol::protocol::ThreadGoal; use codex_protocol::protocol::ThreadGoalStatus; use codex_protocol::protocol::ThreadGoalUpdatedEvent; + use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::TurnContextItem; use codex_protocol::protocol::USER_MESSAGE_BEGIN; use codex_protocol::protocol::UserMessageEvent; @@ -516,6 +517,7 @@ mod tests { #[test] fn session_meta_does_not_set_model_or_reasoning_effort() { let mut metadata = metadata_for_test(); + metadata.history_mode = ThreadHistoryMode::Paginated; let thread_id = metadata.id; apply_rollout_item( @@ -540,7 +542,7 @@ mod tests { dynamic_tools: None, selected_capability_roots: Vec::new(), memory_mode: None, - history_mode: Default::default(), + history_mode: ThreadHistoryMode::Legacy, multi_agent_version: None, context_window: None, }, @@ -551,6 +553,7 @@ mod tests { assert_eq!(metadata.model, None); assert_eq!(metadata.reasoning_effort, None); + assert_eq!(metadata.history_mode, ThreadHistoryMode::Paginated); } fn metadata_for_test() -> ThreadMetadata { diff --git a/codex-rs/thread-store/src/in_memory.rs b/codex-rs/thread-store/src/in_memory.rs index 38598830f7..cb3f198ff3 100644 --- a/codex-rs/thread-store/src/in_memory.rs +++ b/codex-rs/thread-store/src/in_memory.rs @@ -218,17 +218,22 @@ mod tests { }) .await .expect("register rollout path"); - store - .update_thread_metadata(UpdateThreadMetadataParams { - thread_id, - patch: ThreadMetadataPatch { - history_mode: Some(ThreadHistoryMode::Paginated), - ..Default::default() - }, - include_archived: false, - }) - .await - .expect("seed paginated metadata"); + { + let mut state = store.state.lock().await; + state + .created_threads + .get_mut(&thread_id) + .expect("created thread") + .history_mode = ThreadHistoryMode::Paginated; + let Some(RolloutItem::SessionMeta(meta_line)) = state + .histories + .get_mut(&thread_id) + .and_then(|history| history.first_mut()) + else { + panic!("canonical session meta"); + }; + meta_line.meta.history_mode = ThreadHistoryMode::Paginated; + } let thread = store .read_thread(ReadThreadParams { @@ -784,9 +789,7 @@ fn stored_thread_from_state( source: metadata .and_then(|metadata| metadata.source.clone()) .unwrap_or_else(|| created.source.clone()), - history_mode: metadata - .and_then(|metadata| metadata.history_mode) - .unwrap_or(created.history_mode), + history_mode: created.history_mode, thread_source: metadata .and_then(|metadata| metadata.thread_source.clone()) .unwrap_or_else(|| created.thread_source.clone()), @@ -811,15 +814,9 @@ fn history_mode_from_state( thread_id: ThreadId, ) -> ThreadHistoryMode { state - .metadata_updates + .created_threads .get(&thread_id) - .and_then(|metadata| metadata.history_mode) - .or_else(|| { - state - .created_threads - .get(&thread_id) - .map(|thread| thread.history_mode) - }) + .map(|thread| thread.history_mode) .unwrap_or_default() } diff --git a/codex-rs/thread-store/src/local/live_writer.rs b/codex-rs/thread-store/src/local/live_writer.rs index 85ef80d88a..74e707bb58 100644 --- a/codex-rs/thread-store/src/local/live_writer.rs +++ b/codex-rs/thread-store/src/local/live_writer.rs @@ -26,9 +26,12 @@ pub(super) async fn create_thread( params: CreateThreadParams, ) -> ThreadStoreResult<()> { let thread_id = params.thread_id; + let history_mode = params.history_mode; store.ensure_live_recorder_absent(thread_id).await?; let recorder = create_thread::create_thread(store, params).await?; - store.insert_live_recorder(thread_id, recorder).await + store + .insert_live_recorder(thread_id, recorder, history_mode) + .await } pub(super) async fn resume_thread( @@ -98,7 +101,9 @@ pub(super) async fn resume_thread( .map_err(|err| ThreadStoreError::Internal { message: format!("failed to resume local thread recorder: {err}"), })?; - store.insert_live_recorder(params.thread_id, recorder).await + store + .insert_live_recorder(params.thread_id, recorder, history_mode) + .await } #[tracing::instrument( @@ -191,6 +196,7 @@ pub(super) async fn rollout_path( .await .get(&thread_id) .ok_or(ThreadStoreError::ThreadNotFound { thread_id })? + .recorder .rollout_path() .to_path_buf()) } diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index bc34c52f33..5939ea606e 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -13,6 +13,7 @@ mod update_thread_metadata; mod test_support; use codex_protocol::ThreadId; +use codex_protocol::protocol::ThreadHistoryMode; use codex_rollout::RolloutRecorder; use codex_rollout::StateDbHandle; use std::collections::HashMap; @@ -57,10 +58,18 @@ use crate::UpdateThreadMetadataParams; #[derive(Clone)] pub struct LocalThreadStore { pub(super) config: LocalThreadStoreConfig, - live_recorders: Arc>>, + live_recorders: Arc>>, state_db: Option, } +struct LiveRecorderEntry { + recorder: RolloutRecorder, + // Local rollout files are materialized lazily, but metadata updates can arrive before the + // canonical SessionMeta is durable. Retain the mode captured when live persistence was opened + // so missing SQLite rows can still be seeded. + history_mode: ThreadHistoryMode, +} + /// Process-scoped configuration for local thread storage. /// /// This describes where local storage lives. New-thread rollout metadata such @@ -135,7 +144,7 @@ impl LocalThreadStore { .lock() .await .get(&thread_id) - .cloned() + .map(|entry| entry.recorder.clone()) .ok_or(ThreadStoreError::ThreadNotFound { thread_id }) } @@ -155,13 +164,17 @@ impl LocalThreadStore { &self, thread_id: ThreadId, recorder: RolloutRecorder, + history_mode: ThreadHistoryMode, ) -> ThreadStoreResult<()> { match self.live_recorders.lock().await.entry(thread_id) { Entry::Occupied(entry) => Err(ThreadStoreError::InvalidRequest { message: format!("thread {} already has a live local writer", entry.key()), }), Entry::Vacant(entry) => { - entry.insert(recorder); + entry.insert(LiveRecorderEntry { + recorder, + history_mode, + }); Ok(()) } } @@ -580,6 +593,46 @@ mod tests { ); } + #[tokio::test] + async fn live_thread_memory_mode_update_before_rollout_materializes_keeps_history_mode() { + let home = TempDir::new().expect("temp dir"); + let config = test_config(home.path()); + let runtime = codex_state::StateRuntime::init( + config.sqlite_home.clone(), + config.default_model_provider_id.clone(), + ) + .await + .expect("state db should initialize"); + let store = Arc::new(LocalThreadStore::new(config, Some(runtime.clone()))); + let thread_id = ThreadId::default(); + let live_thread = LiveThread::create(store.clone(), create_thread_params(thread_id)) + .await + .expect("create live thread"); + + live_thread + .update_memory_mode(ThreadMemoryMode::Disabled, /*include_archived*/ false) + .await + .expect("update memory mode"); + + assert_eq!( + runtime + .get_thread(thread_id) + .await + .expect("sqlite metadata read") + .expect("sqlite metadata") + .history_mode, + ThreadHistoryMode::Legacy + ); + assert_eq!( + runtime + .get_thread_memory_mode(thread_id) + .await + .expect("thread memory mode should be readable") + .as_deref(), + Some("disabled") + ); + } + #[tokio::test] async fn live_thread_shutdown_with_buffered_items_materializes_before_metadata_read() { let home = TempDir::new().expect("temp dir"); diff --git a/codex-rs/thread-store/src/local/update_thread_metadata.rs b/codex-rs/thread-store/src/local/update_thread_metadata.rs index 38e919b879..580df1dc61 100644 --- a/codex-rs/thread-store/src/local/update_thread_metadata.rs +++ b/codex-rs/thread-store/src/local/update_thread_metadata.rs @@ -6,6 +6,7 @@ use codex_protocol::ThreadId; use codex_protocol::protocol::GitInfo; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadMemoryMode; use codex_rollout::ARCHIVED_SESSIONS_SUBDIR; use codex_rollout::append_rollout_item_to_path; @@ -225,31 +226,27 @@ async fn apply_metadata_update( rollout_path_archived = resolved.archived; rollout_path = Some(resolved.path); } - let mut metadata = existing.clone().unwrap_or_else(|| { - let created_at = patch - .created_at - .or(patch.updated_at) - .unwrap_or_else(Utc::now); - let mut builder = ThreadMetadataBuilder::new( - thread_id, - rollout_path.clone().unwrap_or_default(), - created_at, - patch.source.clone().unwrap_or(SessionSource::Unknown), - ); - builder.model_provider = patch.model_provider.clone(); - builder.history_mode = patch.history_mode.unwrap_or_default(); - builder.thread_source = patch.thread_source.clone().flatten(); - builder.agent_nickname = patch.agent_nickname.clone().flatten(); - builder.agent_role = patch.agent_role.clone().flatten(); - builder.agent_path = patch.agent_path.clone().flatten(); - builder.cwd = patch.cwd.clone().map(normalize_cwd).unwrap_or_default(); - builder.cli_version = patch.cli_version.clone(); - let mut metadata = builder.build(store.config.default_model_provider_id.as_str()); - if rollout_path_archived { - metadata.archived_at = Some(metadata.updated_at); + let mut metadata = match existing.clone() { + Some(metadata) => metadata, + None => { + let rollout_path = + rollout_path + .as_deref() + .ok_or_else(|| ThreadStoreError::Internal { + message: format!( + "thread metadata missing rollout path for {thread_id}" + ), + })?; + metadata_for_missing_sqlite_row( + store, + thread_id, + rollout_path, + rollout_path_archived, + &patch, + ) + .await? } - metadata - }); + }; if let Some(rollout_path) = rollout_path { metadata.rollout_path = rollout_path; } @@ -285,9 +282,6 @@ async fn apply_metadata_update( if let Some(source) = patch.source { metadata.source = enum_to_string(&source); } - if let Some(history_mode) = patch.history_mode { - metadata.history_mode = history_mode; - } if let Some(thread_source) = patch.thread_source { metadata.thread_source = thread_source; } @@ -389,6 +383,78 @@ async fn apply_metadata_update( .await } +async fn metadata_for_missing_sqlite_row( + store: &LocalThreadStore, + thread_id: ThreadId, + rollout_path: &Path, + rollout_path_archived: bool, + patch: &ThreadMetadataPatch, +) -> ThreadStoreResult { + let created_at = patch + .created_at + .or(patch.updated_at) + .unwrap_or_else(Utc::now); + let mut builder = ThreadMetadataBuilder::new( + thread_id, + rollout_path.to_path_buf(), + created_at, + patch.source.clone().unwrap_or(SessionSource::Unknown), + ); + builder.model_provider = patch.model_provider.clone(); + builder.history_mode = canonical_history_mode(store, thread_id, rollout_path).await?; + builder.thread_source = patch.thread_source.clone().flatten(); + builder.agent_nickname = patch.agent_nickname.clone().flatten(); + builder.agent_role = patch.agent_role.clone().flatten(); + builder.agent_path = patch.agent_path.clone().flatten(); + builder.cwd = patch.cwd.clone().map(normalize_cwd).unwrap_or_default(); + builder.cli_version = patch.cli_version.clone(); + let mut metadata = builder.build(store.config.default_model_provider_id.as_str()); + if rollout_path_archived { + metadata.archived_at = Some(metadata.updated_at); + } + Ok(metadata) +} + +async fn canonical_history_mode( + store: &LocalThreadStore, + thread_id: ThreadId, + rollout_path: &Path, +) -> ThreadStoreResult { + let session_meta = match read_session_meta_line(rollout_path).await { + Ok(session_meta) => session_meta, + Err(err) => { + if codex_rollout::existing_rollout_path(rollout_path) + .await + .is_none() + && let Some(history_mode) = store + .live_recorders + .lock() + .await + .get(&thread_id) + .map(|entry| entry.history_mode) + { + // The live writer retains the canonical mode selected before its deferred + // SessionMeta reaches JSONL. + return Ok(history_mode); + } + return Err(ThreadStoreError::Internal { + message: format!( + "failed to read canonical session metadata for {thread_id}: {err}" + ), + }); + } + }; + if session_meta.meta.id != thread_id { + return Err(ThreadStoreError::Internal { + message: format!( + "failed to rebuild thread metadata: rollout session metadata id mismatch: expected {thread_id}, found {}", + session_meta.meta.id + ), + }); + } + Ok(session_meta.meta.history_mode) +} + fn needs_rollout_compatibility_update(patch: &ThreadMetadataPatch) -> bool { if patch.name.is_some() { return true; @@ -431,7 +497,6 @@ fn has_observed_metadata_facts(patch: &ThreadMetadataPatch) -> bool { || patch.permission_profile.is_some() || patch.token_usage.is_some() || patch.first_user_message.is_some() - || patch.history_mode.is_some() } fn enum_to_string(value: &T) -> String { @@ -1475,6 +1540,51 @@ mod tests { assert!(metadata.is_none()); } + #[tokio::test] + async fn observed_metadata_rebuilds_history_mode_from_canonical_session_meta() { + let home = TempDir::new().expect("temp dir"); + let config = test_config(home.path()); + let uuid = Uuid::from_u128(317); + let thread_id = ThreadId::from_string(&uuid.to_string()).expect("valid thread id"); + write_session_file_with_history_mode( + home.path(), + "2025-01-03T19-15-00", + uuid, + ThreadHistoryMode::Paginated, + ) + .expect("session file"); + let runtime = codex_state::StateRuntime::init( + home.path().to_path_buf(), + config.default_model_provider_id.clone(), + ) + .await + .expect("state db should initialize"); + let store = LocalThreadStore::new(config, Some(runtime.clone())); + + let thread = store + .update_thread_metadata(UpdateThreadMetadataParams { + thread_id, + patch: ThreadMetadataPatch { + preview: Some("Paginated preview".to_string()), + ..Default::default() + }, + include_archived: false, + }) + .await + .expect("update paginated thread without sqlite row"); + + assert_eq!(thread.history_mode, ThreadHistoryMode::Paginated); + assert_eq!( + runtime + .get_thread(thread_id) + .await + .expect("sqlite metadata read") + .expect("sqlite metadata") + .history_mode, + ThreadHistoryMode::Paginated + ); + } + #[tokio::test] async fn update_thread_metadata_recreates_missing_archived_sqlite_row_as_archived() { let home = TempDir::new().expect("temp dir"); diff --git a/codex-rs/thread-store/src/thread_metadata_sync.rs b/codex-rs/thread-store/src/thread_metadata_sync.rs index a7b52ee0cb..1057715a4a 100644 --- a/codex-rs/thread-store/src/thread_metadata_sync.rs +++ b/codex-rs/thread-store/src/thread_metadata_sync.rs @@ -76,7 +76,6 @@ impl ThreadMetadataSync { cli_version: Some(env!("CARGO_PKG_VERSION").to_string()), git_info: git_info.map(git_info_patch_from_observation), memory_mode: Some(params.metadata.memory_mode), - history_mode: Some(params.history_mode), ..Default::default() }; Self { @@ -238,7 +237,6 @@ impl ThreadMetadataSync { { update.memory_mode = Some(memory_mode); } - update.history_mode = Some(meta_line.meta.history_mode); } RolloutItem::TurnContext(turn_ctx) => { if !self.cwd_seen { @@ -377,7 +375,6 @@ fn update_has_metadata_facts(update: &ThreadMetadataPatch) -> bool { || update.first_user_message.is_some() || update.git_info.is_some() || update.memory_mode.is_some() - || update.history_mode.is_some() } fn git_info_patch_from_observation(git_info: GitInfo) -> GitInfoPatch { diff --git a/codex-rs/thread-store/src/types.rs b/codex-rs/thread-store/src/types.rs index 3b86840ca7..4e688eca6f 100644 --- a/codex-rs/thread-store/src/types.rs +++ b/codex-rs/thread-store/src/types.rs @@ -596,8 +596,6 @@ pub struct ThreadMetadataPatch { pub git_info: Option, /// Thread memory behavior. pub memory_mode: Option, - /// Persisted thread history contract. - pub history_mode: Option, } impl ThreadMetadataPatch { @@ -678,9 +676,6 @@ impl ThreadMetadataPatch { if next.memory_mode.is_some() { self.memory_mode = next.memory_mode; } - if next.history_mode.is_some() { - self.history_mode = next.history_mode; - } } pub fn is_empty(&self) -> bool { @@ -707,7 +702,6 @@ impl ThreadMetadataPatch { && self.first_user_message.is_none() && self.git_info.is_none() && self.memory_mode.is_none() - && self.history_mode.is_none() } } From f72976a5f1e5bba3e1af18cc2ddffedfb2687376 Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Fri, 26 Jun 2026 12:35:54 -0700 Subject: [PATCH 31/38] feat(app-server): add optional turn_id to thread/fork (#30277) ## Description This adds stable optional `turnId` support to `thread/fork`. When supplied, the fork copies persisted history through that terminal turn, inclusive, and drops later turns from the new thread. Omitting or passing `null` preserves the existing full-history fork behavior, including the interruption marker when the stored source history ends mid-turn. ## Why We're deprecating `thread/rollback` and this will help certain UX use cases work around it by using `thread/fork` + `turn_id` instead. --- .../schema/json/ClientRequest.json | 9 +- .../codex_app_server_protocol.schemas.json | 9 +- .../codex_app_server_protocol.v2.schemas.json | 9 +- .../schema/json/v2/ThreadForkParams.json | 9 +- .../schema/typescript/v2/ThreadForkParams.ts | 6 + .../src/protocol/v2/tests.rs | 24 ++++ .../src/protocol/v2/thread.rs | 7 ++ codex-rs/app-server/README.md | 2 +- codex-rs/app-server/src/request_processors.rs | 1 + .../request_processors/thread_processor.rs | 11 +- .../app-server/tests/suite/v2/thread_fork.rs | 113 ++++++++++++++++++ codex-rs/core/src/lib.rs | 1 + .../core/src/thread_rollout_truncation.rs | 56 +++++++++ .../src/thread_rollout_truncation_tests.rs | 101 ++++++++++++++++ 14 files changed, 352 insertions(+), 6 deletions(-) diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index 8c1d700ef2..ca247f2d55 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -3673,6 +3673,13 @@ } ], "description": "Optional client-supplied analytics source classification for this forked thread." + }, + "lastTurnId": { + "description": "Optional last turn id to fork through, inclusive.\n\nWhen specified, turns after `last_turn_id` are omitted from the fork. The referenced turn cannot be in progress.", + "type": [ + "string", + "null" + ] } }, "required": [ @@ -6866,4 +6873,4 @@ } ], "title": "ClientRequest" -} \ No newline at end of file +} 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 b972de2206..9f03be5b79 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 @@ -17287,6 +17287,13 @@ } ], "description": "Optional client-supplied analytics source classification for this forked thread." + }, + "lastTurnId": { + "description": "Optional last turn id to fork through, inclusive.\n\nWhen specified, turns after `last_turn_id` are omitted from the fork. The referenced turn cannot be in progress.", + "type": [ + "string", + "null" + ] } }, "required": [ @@ -20833,4 +20840,4 @@ }, "title": "CodexAppServerProtocol", "type": "object" -} \ No newline at end of file +} 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 4af580e9c3..f7037ef836 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 @@ -15066,6 +15066,13 @@ } ], "description": "Optional client-supplied analytics source classification for this forked thread." + }, + "lastTurnId": { + "description": "Optional last turn id to fork through, inclusive.\n\nWhen specified, turns after `last_turn_id` are omitted from the fork. The referenced turn cannot be in progress.", + "type": [ + "string", + "null" + ] } }, "required": [ @@ -18611,4 +18618,4 @@ }, "title": "CodexAppServerProtocolV2", "type": "object" -} \ No newline at end of file +} diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadForkParams.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadForkParams.json index ad0174cfe0..7a00e3803a 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadForkParams.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadForkParams.json @@ -168,6 +168,13 @@ } ], "description": "Optional client-supplied analytics source classification for this forked thread." + }, + "lastTurnId": { + "description": "Optional last turn id to fork through, inclusive.\n\nWhen specified, turns after `last_turn_id` are omitted from the fork. The referenced turn cannot be in progress.", + "type": [ + "string", + "null" + ] } }, "required": [ @@ -175,4 +182,4 @@ ], "title": "ThreadForkParams", "type": "object" -} \ No newline at end of file +} diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/ThreadForkParams.ts b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadForkParams.ts index c5109b2c7d..3ace4d4417 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/ThreadForkParams.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/ThreadForkParams.ts @@ -18,6 +18,12 @@ import type { ThreadSource } from "./ThreadSource"; * Prefer using thread_id whenever possible. */ export type ThreadForkParams = {threadId: string, /** + * Optional last turn id to fork through, inclusive. + * + * When specified, turns after `last_turn_id` are omitted from the fork. + * The referenced turn cannot be in progress. + */ +lastTurnId?: string | null, /** * Configuration overrides for the forked thread, if any. */ model?: string | null, modelProvider?: string | null, serviceTier?: string | null | null, cwd?: string | null, approvalPolicy?: AskForApproval | null, /** diff --git a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs index ea1e0424dd..930441c63f 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs @@ -839,6 +839,30 @@ fn thread_path_params_deserialize_empty_path_as_none() { ); } +#[test] +fn thread_fork_last_turn_id_round_trips() { + let params: ThreadForkParams = serde_json::from_value(json!({ + "threadId": "thread-1", + "lastTurnId": "turn-2", + })) + .expect("thread/fork params deserialize"); + + assert_eq!(params.last_turn_id, Some("turn-2".to_string())); + let serialized = serde_json::to_value(params).expect("thread/fork params serialize"); + assert_eq!(serialized["lastTurnId"], json!("turn-2")); + + let omitted = serde_json::to_value(ThreadForkParams { + thread_id: "thread-1".to_string(), + ..Default::default() + }) + .expect("thread/fork params without last turn id serialize"); + assert_eq!( + omitted["lastTurnId"], + serde_json::Value::Null, + "optional lastTurnId should serialize as null when omitted" + ); +} + #[test] fn fs_get_metadata_response_round_trips_minimal_fields() { let response = FsGetMetadataResponse { diff --git a/codex-rs/app-server-protocol/src/protocol/v2/thread.rs b/codex-rs/app-server-protocol/src/protocol/v2/thread.rs index 73335fce56..7cd03930d8 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/thread.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/thread.rs @@ -494,6 +494,13 @@ impl From for TurnsPage { pub struct ThreadForkParams { pub thread_id: String, + /// Optional last turn id to fork through, inclusive. + /// + /// When specified, turns after `last_turn_id` are omitted from the fork. + /// The referenced turn cannot be in progress. + #[ts(optional = nullable)] + pub last_turn_id: Option, + /// [UNSTABLE] Specify the rollout path to fork from. /// If specified, the thread_id param will be ignored. #[experimental("thread/fork.path")] diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 78f46f5b13..6f34cae541 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -139,7 +139,7 @@ Example with notification opt-out: - `thread/start` — create a new thread; emits `thread/started` (including the current `thread.status`) and auto-subscribes you to turn/item events for that thread. When the request includes a `cwd` and the resolved sandbox is `workspace-write` or full access, app-server also marks that project as trusted in the user `config.toml`. Pass `sessionStartSource: "clear"` when starting a replacement thread after clearing the current session so `SessionStart` hooks receive `source: "clear"` instead of the default `"startup"`. Experimental `allowProviderModelFallback` lets providers backed by an authoritative static model catalog replace an unavailable requested `model` with the catalog default; dynamic or cached catalogs preserve the requested model. Experimental `runtimeWorkspaceRoots` replaces the thread-scoped runtime workspace roots used to materialize `:workspace_roots`; paths must be absolute. For permissions, prefer experimental `permissions` profile selection by id; the legacy `sandbox` shorthand is still accepted but cannot be combined with `permissions`. Deprecated experimental `multiAgentMode` is ignored; use Ultra reasoning effort for proactive multi-agent behavior. Experimental `environments` selects the sticky execution environments for turns on the thread; omit it to use the server default, pass `[]` to disable environments, or pass explicit environment ids with per-environment `cwd`. Experimental `selectedCapabilityRoots` selects environment-owned plugin or standalone-skill roots using environment-native absolute paths. Skills found below those roots are listed and read through the owning environment. Stdio MCP servers declared by selected plugins are started in that environment, and HTTP MCP connections use that environment's HTTP client. - `thread/resume` — reopen an existing thread by id so subsequent `turn/start` calls append to it. Accepts the same permission override rules as `thread/start`. -- `thread/fork` — fork an existing thread into a new thread id by copying the stored history; if the source thread is currently mid-turn, the fork records the same interruption marker as `turn/interrupt` instead of inheriting an unmarked partial turn suffix. The returned `thread.forkedFromId` points at the source thread when known. Accepts `ephemeral: true` for an in-memory temporary fork, emits `thread/started` (including the current `thread.status`), and auto-subscribes you to turn/item events for the new thread. Experimental clients can pass `excludeTurns: true` when they plan to page fork history via `thread/turns/list` instead of receiving the full turn array immediately. Accepts the same permission override rules as `thread/start`. +- `thread/fork` — fork an existing thread into a new thread id by copying the stored history; pass an optional `lastTurnId` to copy history only through that turn, inclusive, and drop later turns from the fork. An in-progress `lastTurnId` is rejected. If `lastTurnId` is null while the source thread is mid-turn, the fork records the same interruption marker as `turn/interrupt` instead of inheriting an unmarked partial turn suffix. The returned `thread.forkedFromId` points at the source thread when known. Accepts `ephemeral: true` for an in-memory temporary fork, emits `thread/started` (including the current `thread.status`), and auto-subscribes you to turn/item events for the new thread. Experimental clients can pass `excludeTurns: true` when they plan to page fork history via `thread/turns/list` instead of receiving the full turn array immediately. Accepts the same permission override rules as `thread/start`. - `thread/start`, `thread/resume`, and `thread/fork` responses include the legacy `sandbox` compatibility projection. `instructionSources` lists loaded instruction files using each source environment's native absolute path syntax, including files loaded from remote environments. Experimental clients can read `runtimeWorkspaceRoots` for the thread-scoped runtime roots and `activePermissionProfile` for the named or implicit built-in profile identity/provenance when known. Their deprecated experimental `multiAgentMode` field, and the corresponding thread setting, always report `explicitRequestOnly`; Ultra reasoning effort is the source of proactive multi-agent behavior. - `thread/list` — page through stored threads; supports cursor-based pagination and optional `modelProviders`, `sourceKinds`, `archived`, `cwd`, and `searchTerm` filters. Experimental clients can use `parentThreadId` for direct spawned children or `ancestorThreadId` for spawned descendants at any depth; the two filters are mutually exclusive. Review and Guardian threads are not included because they do not participate in that spawn-edge lifecycle. Each returned `thread` includes `status` (`ThreadStatus`), defaulting to `notLoaded` when the thread is not currently loaded. Subagent threads also include `parentThreadId` when the immediate parent is known. - `thread/loaded/list` — list the thread ids currently loaded in memory. diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 39c4b79dc4..cb0f276d85 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -330,6 +330,7 @@ use codex_core::path_utils; #[cfg(test)] use codex_core::read_head_for_summary; use codex_core::sandboxing::SandboxPermissions; +use codex_core::truncate_rollout_after_turn_id; use codex_core::windows_sandbox::WindowsSandboxLevelExt; use codex_core::windows_sandbox::WindowsSandboxSetupMode as CoreWindowsSandboxSetupMode; use codex_core::windows_sandbox::WindowsSandboxSetupRequest; diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index d5302d9c43..b78c53a383 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -3387,6 +3387,7 @@ impl ThreadRequestProcessor { ) -> Result<(), JSONRPCErrorError> { let ThreadForkParams { thread_id, + last_turn_id, path, model, model_provider, @@ -3421,12 +3422,20 @@ impl ThreadRequestProcessor { let history_items = source_thread .history .take() - .map(|history| Arc::new(history.items)) + .map(|history| history.items) .ok_or_else(|| { internal_error(format!( "thread {source_thread_id} did not include persisted history" )) })?; + let history_items = if let Some(last_turn_id) = last_turn_id.as_deref() { + Arc::new( + truncate_rollout_after_turn_id(&history_items, last_turn_id) + .map_err(|err| core_thread_write_error("truncate thread for fork", err))?, + ) + } else { + Arc::new(history_items) + }; let history_cwd = Some(source_thread.cwd.clone()); // Persist Windows sandbox mode. 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 4102b45d15..d7c1707a32 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_fork.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_fork.rs @@ -251,6 +251,119 @@ async fn thread_fork_creates_new_thread_and_emits_started() -> Result<()> { Ok(()) } +#[tokio::test] +async fn thread_fork_at_last_turn_id_keeps_only_terminal_prefix() -> 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 mut mcp = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + + let start_id = mcp + .send_thread_start_request(ThreadStartParams::default()) + .await?; + let start_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(start_id)), + ) + .await??; + let ThreadStartResponse { + thread: source_thread, + .. + } = to_response::(start_resp)?; + let source_thread_id = source_thread.id.clone(); + let source_path = source_thread.path.expect("source thread path"); + + let mut turn_ids = Vec::new(); + for text in ["first", "second", "third"] { + let turn_request_id = mcp + .send_turn_start_request(TurnStartParams { + thread_id: source_thread_id.clone(), + client_user_message_id: None, + input: vec![UserInput::Text { + text: text.to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + let turn_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_request_id)), + ) + .await??; + let TurnStartResponse { turn } = to_response::(turn_resp)?; + turn_ids.push(turn.id); + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + } + + let original_contents = std::fs::read_to_string(source_path.as_path())?; + let fork_id = mcp + .send_thread_fork_request(ThreadForkParams { + thread_id: source_thread_id.clone(), + last_turn_id: Some(turn_ids[1].clone()), + ..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: forked_thread, + .. + } = to_response::(fork_resp)?; + + assert_eq!( + forked_thread + .turns + .iter() + .map(|turn| turn.id.clone()) + .collect::>(), + turn_ids[..2] + ); + assert!( + forked_thread + .turns + .iter() + .all(|turn| turn.status == TurnStatus::Completed) + ); + assert_eq!(forked_thread.forked_from_id, Some(source_thread_id)); + assert_eq!(forked_thread.preview, "first"); + assert_eq!( + std::fs::read_to_string(source_path.as_path())?, + original_contents, + "forking at a turn must not mutate the source rollout" + ); + + let forked_path = forked_thread.path.clone().expect("forked thread path"); + let forked_contents = std::fs::read_to_string(forked_path.as_path())?; + assert!(forked_contents.contains(turn_ids[1].as_str())); + assert!(!forked_contents.contains(turn_ids[2].as_str())); + + let started = loop { + let notification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("thread/started"), + ) + .await??; + let started: ThreadStartedNotification = + serde_json::from_value(notification.params.expect("params must be present"))?; + if started.thread.id == forked_thread.id { + break started; + } + }; + assert!(started.thread.turns.is_empty()); + + Ok(()) +} + #[tokio::test] async fn thread_fork_inherits_explicit_source_name_from_session_index() -> Result<()> { let server = create_mock_responses_server_repeating_assistant("Done").await; diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index ad5f54123f..e9b749fea9 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -145,6 +145,7 @@ pub(crate) mod state_db_bridge; pub use state_db_bridge::StateDbHandle; pub use state_db_bridge::init_state_db; mod thread_rollout_truncation; +pub use thread_rollout_truncation::truncate_rollout_after_turn_id; mod tools; pub(crate) mod turn_diff_tracker; mod turn_metadata; diff --git a/codex-rs/core/src/thread_rollout_truncation.rs b/codex-rs/core/src/thread_rollout_truncation.rs index 25263b14b1..d7d6ad998c 100644 --- a/codex-rs/core/src/thread_rollout_truncation.rs +++ b/codex-rs/core/src/thread_rollout_truncation.rs @@ -5,6 +5,10 @@ use crate::context_manager::is_user_turn_boundary; use crate::event_mapping; +use codex_app_server_protocol::TurnStatus; +use codex_app_server_protocol::build_turns_from_rollout_items; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result as CodexResult; use codex_protocol::items::TurnItem; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; @@ -149,6 +153,58 @@ pub(crate) fn truncate_rollout_before_nth_user_message_from_start( items[..cut_idx].to_vec() } +/// Return a rollout prefix ending after the requested persisted terminal turn. +/// +/// The turn must still be present in the effective post-rollback history and +/// must have an explicit persisted TurnStarted boundary. Synthetic IDs +/// generated while projecting legacy rollouts are intentionally unsupported +/// because they do not provide a stable raw rollout boundary for a fork. +pub fn truncate_rollout_after_turn_id( + items: &[RolloutItem], + last_turn_id: &str, +) -> CodexResult> { + let turns = build_turns_from_rollout_items(items); + let turn = turns + .iter() + .find(|turn| turn.id == last_turn_id) + .ok_or_else(|| { + CodexErr::InvalidRequest(format!( + "lastTurnId '{last_turn_id}' was not found in the source thread" + )) + })?; + + let target_start_index = items + .iter() + .position(|item| { + matches!( + item, + RolloutItem::EventMsg(EventMsg::TurnStarted(event)) + if event.turn_id == last_turn_id + ) + }) + .ok_or_else(|| { + CodexErr::InvalidRequest(format!( + "lastTurnId '{last_turn_id}' is not a persisted canonical turn in the source thread" + )) + })?; + + if matches!(turn.status, TurnStatus::InProgress) { + return Err(CodexErr::InvalidRequest(format!( + "lastTurnId '{last_turn_id}' identifies an in-progress turn" + ))); + } + + let cut_index = items + .iter() + .enumerate() + .skip(target_start_index.saturating_add(1)) + .find_map(|(index, item)| { + matches!(item, RolloutItem::EventMsg(EventMsg::TurnStarted(_))).then_some(index) + }) + .unwrap_or(items.len()); + Ok(items[..cut_index].to_vec()) +} + /// Return a suffix of `items` that keeps the last `n_from_end` fork turns. /// /// If fewer than or equal to `n_from_end` fork turns exist, this keeps from the first fork-turn diff --git a/codex-rs/core/src/thread_rollout_truncation_tests.rs b/codex-rs/core/src/thread_rollout_truncation_tests.rs index 0da9763fc7..16d10a8277 100644 --- a/codex-rs/core/src/thread_rollout_truncation_tests.rs +++ b/codex-rs/core/src/thread_rollout_truncation_tests.rs @@ -6,6 +6,9 @@ use codex_protocol::models::ContentItem; use codex_protocol::models::ReasoningItemReasoningSummary; use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::ThreadRolledBackEvent; +use codex_protocol::protocol::TurnCompleteEvent; +use codex_protocol::protocol::TurnStartedEvent; +use codex_protocol::protocol::UserMessageEvent; use pretty_assertions::assert_eq; use std::sync::Arc; @@ -66,6 +69,104 @@ fn inter_agent_communication(text: &str, trigger_turn: bool) -> RolloutItem { )) } +fn turn_started(turn_id: &str) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::TurnStarted(TurnStartedEvent { + turn_id: turn_id.to_string(), + trace_id: None, + started_at: None, + model_context_window: None, + collaboration_mode_kind: Default::default(), + })) +} + +fn turn_completed(turn_id: &str) -> RolloutItem { + RolloutItem::EventMsg(EventMsg::TurnComplete(TurnCompleteEvent { + turn_id: turn_id.to_string(), + last_agent_message: None, + completed_at: None, + duration_ms: None, + time_to_first_token_ms: None, + })) +} + +#[test] +fn truncates_rollout_after_terminal_canonical_turn_id() { + let rollout = vec![ + turn_started("turn-1"), + turn_completed("turn-1"), + turn_started("turn-2"), + turn_completed("turn-2"), + turn_started("turn-3"), + turn_completed("turn-3"), + ]; + + let truncated = + truncate_rollout_after_turn_id(&rollout, "turn-2").expect("truncate through turn-2"); + + assert_eq!( + serde_json::to_value(&truncated).unwrap(), + serde_json::to_value(&rollout[..4]).unwrap() + ); +} + +#[test] +fn truncate_rollout_after_turn_id_rejects_rolled_back_turn() { + let rollout = vec![ + turn_started("turn-1"), + turn_completed("turn-1"), + turn_started("turn-2"), + turn_completed("turn-2"), + RolloutItem::EventMsg(EventMsg::ThreadRolledBack(ThreadRolledBackEvent { + num_turns: 1, + })), + turn_started("turn-3"), + turn_completed("turn-3"), + ]; + + let err = truncate_rollout_after_turn_id(&rollout, "turn-2") + .expect_err("rolled-back turn should not be a fork anchor"); + + assert!(matches!( + err, + CodexErr::InvalidRequest(message) + if message == "lastTurnId 'turn-2' was not found in the source thread" + )); +} + +#[test] +fn truncate_rollout_after_turn_id_rejects_synthetic_legacy_turn_id() { + let rollout = vec![RolloutItem::EventMsg(EventMsg::UserMessage( + UserMessageEvent { + message: "legacy".to_string(), + ..Default::default() + }, + ))]; + + let err = truncate_rollout_after_turn_id(&rollout, "rollout-0") + .expect_err("synthetic turn should not be a fork anchor"); + + assert!(matches!( + err, + CodexErr::InvalidRequest(message) + if message + == "lastTurnId 'rollout-0' is not a persisted canonical turn in the source thread" + )); +} + +#[test] +fn truncate_rollout_after_turn_id_rejects_in_progress_turn() { + let rollout = vec![turn_started("turn-1")]; + + let err = truncate_rollout_after_turn_id(&rollout, "turn-1") + .expect_err("in-progress turn should not be a fork anchor"); + + assert!(matches!( + err, + CodexErr::InvalidRequest(message) + if message == "lastTurnId 'turn-1' identifies an in-progress turn" + )); +} + #[test] fn truncates_rollout_from_start_before_nth_user_only() { let items = [ From ac85409b7b137843e1ec59d9ecaa063e88c60e6c Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Fri, 26 Jun 2026 12:36:40 -0700 Subject: [PATCH 32/38] Let Codex consult user-level code-review-* skills. (#30143) ## Why I use the `$code-review` skill a lot and it'd be nice to add my own additional review criteria in `$CODEX_HOME/skills/code-review-*`. ## What Removes phrasing about "code-review-* skills in this repository" which in practice seems like enough to get Codex to consult my user-level code review skills in addition to the repo-level ones. --- .codex/skills/code-review/SKILL.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.codex/skills/code-review/SKILL.md b/.codex/skills/code-review/SKILL.md index eec0787c20..ccd37a9868 100644 --- a/.codex/skills/code-review/SKILL.md +++ b/.codex/skills/code-review/SKILL.md @@ -3,7 +3,7 @@ name: code-review description: Run a final code review on a pull request --- -Use subagents to review code using all code-review-* skills in this repository other than this orchestrator. One subagent per skill. Pass full skill path to subagents. Use xhigh reasoning. +Use subagents to review code using all code-review-* skills other than this orchestrator. One subagent per skill. Pass full skill path to subagents. Use xhigh reasoning. You must return every single issue from every subagent. You can return an unlimited number of findings. Use raw Markdown to report findings. From 69596f0e428ae6698b0b9f1174e05bac0a023ded Mon Sep 17 00:00:00 2001 From: Celia Chen Date: Fri, 26 Jun 2026 13:32:49 -0700 Subject: [PATCH 33/38] feat: add GPT-5.6 variants to Bedrock catalog (#30285) ## Summary - add Sol (`openai.gpt-5.6-sol`), Terra (`openai.gpt-5.6-terra`), and Luna (`openai.gpt-5.6-luna`) to the Amazon Bedrock static model catalog - derive all three entries from the bundled GPT-5.5 metadata and add the Bedrock-only `max` reasoning effort - keep the new entries below the current GPT-5.5 and GPT-5.4 models at priorities 2, 3, and 4, preserving GPT-5.5 as the default - add deep-equality coverage for inherited model configuration, catalog ordering, context windows, and service-tier behavior --- codex-rs/model-provider-info/src/lib.rs | 3 + .../src/amazon_bedrock/catalog.rs | 104 ++++++++++++++---- codex-rs/model-provider/src/provider.rs | 11 +- 3 files changed, 95 insertions(+), 23 deletions(-) diff --git a/codex-rs/model-provider-info/src/lib.rs b/codex-rs/model-provider-info/src/lib.rs index 265e65b25a..65517cbd43 100644 --- a/codex-rs/model-provider-info/src/lib.rs +++ b/codex-rs/model-provider-info/src/lib.rs @@ -40,6 +40,9 @@ const AMAZON_BEDROCK_PROVIDER_NAME: &str = "Amazon Bedrock"; pub const AMAZON_BEDROCK_PROVIDER_ID: &str = "amazon-bedrock"; pub const AMAZON_BEDROCK_GPT_5_5_MODEL_ID: &str = "openai.gpt-5.5"; pub const AMAZON_BEDROCK_GPT_5_4_MODEL_ID: &str = "openai.gpt-5.4"; +pub const AMAZON_BEDROCK_GPT_5_6_SOL_MODEL_ID: &str = "openai.gpt-5.6-sol"; +pub const AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID: &str = "openai.gpt-5.6-terra"; +pub const AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID: &str = "openai.gpt-5.6-luna"; pub const AMAZON_BEDROCK_DEFAULT_BASE_URL: &str = "https://bedrock-mantle.us-east-1.api.aws/openai/v1"; const AMAZON_BEDROCK_MANTLE_CLIENT_AGENT_HEADER: &str = "x-amzn-mantle-client-agent"; diff --git a/codex-rs/model-provider/src/amazon_bedrock/catalog.rs b/codex-rs/model-provider/src/amazon_bedrock/catalog.rs index 5c006cd0c6..2294994ae6 100644 --- a/codex-rs/model-provider/src/amazon_bedrock/catalog.rs +++ b/codex-rs/model-provider/src/amazon_bedrock/catalog.rs @@ -1,8 +1,13 @@ use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_4_MODEL_ID; use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_5_MODEL_ID; +use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID; +use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_6_SOL_MODEL_ID; +use codex_model_provider_info::AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID; use codex_models_manager::bundled_models_response; use codex_protocol::openai_models::ModelInfo; use codex_protocol::openai_models::ModelsResponse; +use codex_protocol::openai_models::ReasoningEffort; +use codex_protocol::openai_models::ReasoningEffortPreset; const GPT_5_BEDROCK_CONTEXT_WINDOW: i64 = 272_000; const GPT_5_5_OPENAI_MODEL_ID: &str = "gpt-5.5"; @@ -21,6 +26,21 @@ pub(crate) fn static_model_catalog() -> ModelsResponse { AMAZON_BEDROCK_GPT_5_4_MODEL_ID, /*priority*/ 1, ), + gpt_5_6_bedrock_model( + AMAZON_BEDROCK_GPT_5_6_SOL_MODEL_ID, + "Sol", + /*priority*/ 2, + ), + gpt_5_6_bedrock_model( + AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID, + "Terra", + /*priority*/ 3, + ), + gpt_5_6_bedrock_model( + AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID, + "Luna", + /*priority*/ 4, + ), ], }) } @@ -44,6 +64,18 @@ fn gpt_5_bedrock_model(openai_slug: &str, bedrock_slug: &str, priority: i32) -> model } +fn gpt_5_6_bedrock_model(bedrock_slug: &str, display_name: &str, priority: i32) -> ModelInfo { + let mut model = gpt_5_bedrock_model(GPT_5_5_OPENAI_MODEL_ID, bedrock_slug, priority); + model.display_name = display_name.to_string(); + model + .supported_reasoning_levels + .push(ReasoningEffortPreset { + effort: ReasoningEffort::Custom("max".to_string()), + description: "Maximum reasoning depth for the hardest problems".to_string(), + }); + model +} + fn bundled_openai_model(slug: &str) -> ModelInfo { bundled_models_response() .unwrap_or_else(|err| panic!("bundled models.json should parse: {err}")) @@ -64,39 +96,67 @@ mod tests { fn catalog_uses_mantle_model_ids_as_slugs() { let catalog = static_model_catalog(); - assert_eq!(catalog.models.len(), 2); - assert_eq!(catalog.models[0].slug, AMAZON_BEDROCK_GPT_5_5_MODEL_ID); - assert_eq!(catalog.models[1].slug, AMAZON_BEDROCK_GPT_5_4_MODEL_ID); + assert_eq!( + catalog + .models + .iter() + .map(|model| model.slug.as_str()) + .collect::>(), + vec![ + AMAZON_BEDROCK_GPT_5_5_MODEL_ID, + AMAZON_BEDROCK_GPT_5_4_MODEL_ID, + AMAZON_BEDROCK_GPT_5_6_SOL_MODEL_ID, + AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID, + AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID, + ] + ); } #[test] fn gpt_5_bedrock_models_use_bedrock_context_window() { let catalog = static_model_catalog(); + + for model in catalog.models { + assert_eq!( + (model.context_window, model.max_context_window), + ( + Some(GPT_5_BEDROCK_CONTEXT_WINDOW), + Some(GPT_5_BEDROCK_CONTEXT_WINDOW) + ) + ); + } + } + + #[test] + fn gpt_5_6_bedrock_models_clone_gpt_5_5_config_with_max_reasoning_effort() { + let catalog = static_model_catalog(); let gpt_5_5 = catalog .models .iter() .find(|model| model.slug == AMAZON_BEDROCK_GPT_5_5_MODEL_ID) .expect("Bedrock catalog should include GPT-5.5"); - let gpt_5_4 = catalog - .models - .iter() - .find(|model| model.slug == AMAZON_BEDROCK_GPT_5_4_MODEL_ID) - .expect("Bedrock catalog should include GPT-5.4"); - assert_eq!( - (gpt_5_5.context_window, gpt_5_5.max_context_window), - ( - Some(GPT_5_BEDROCK_CONTEXT_WINDOW), - Some(GPT_5_BEDROCK_CONTEXT_WINDOW) - ) - ); - assert_eq!( - (gpt_5_4.context_window, gpt_5_4.max_context_window), - ( - Some(GPT_5_BEDROCK_CONTEXT_WINDOW), - Some(GPT_5_BEDROCK_CONTEXT_WINDOW) - ) - ); + for (slug, display_name, priority) in [ + (AMAZON_BEDROCK_GPT_5_6_SOL_MODEL_ID, "Sol", 2), + (AMAZON_BEDROCK_GPT_5_6_TERRA_MODEL_ID, "Terra", 3), + (AMAZON_BEDROCK_GPT_5_6_LUNA_MODEL_ID, "Luna", 4), + ] { + let mut expected = gpt_5_5.clone(); + expected.slug = slug.to_string(); + expected.display_name = display_name.to_string(); + expected.priority = priority; + expected + .supported_reasoning_levels + .push(ReasoningEffortPreset { + effort: ReasoningEffort::Custom("max".to_string()), + description: "Maximum reasoning depth for the hardest problems".to_string(), + }); + + assert_eq!( + catalog.models.iter().find(|model| model.slug == slug), + Some(&expected) + ); + } } #[test] diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index 6e58f4b74c..5f101e52cd 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -659,7 +659,16 @@ mod tests { .map(|model| model.slug.as_str()) .collect::>(); - assert_eq!(model_ids, vec!["openai.gpt-5.5", "openai.gpt-5.4"]); + assert_eq!( + model_ids, + vec![ + "openai.gpt-5.5", + "openai.gpt-5.4", + "openai.gpt-5.6-sol", + "openai.gpt-5.6-terra", + "openai.gpt-5.6-luna", + ] + ); let default_model = manager .list_models(RefreshStrategy::Online) From c55ce3b51bc3b8bb2805cbee0afc1a8c6e1e7ee2 Mon Sep 17 00:00:00 2001 From: Abdulrahman Alfozan Date: Fri, 26 Jun 2026 13:56:17 -0700 Subject: [PATCH 34/38] Close thread persistence when submission channel closes (#30173) ### Summary Release live thread persistence when a session ends because its submission channel closes. This prevents a later same-process resume from failing with `thread ... already has a live local writer`. ### Details The issue is in the `codex-core` session teardown path used by Codex hosts, rather than in Managed Agents API or exec-server itself. Explicit shutdown already closes the `LiveThread`, which releases the process-scoped writer held by `LocalThreadStore`. The submission-channel-close fallback ran runtime and extension teardown but skipped that persistence shutdown, leaving the thread ID registered as having a live writer. This change: - closes the `LiveThread` on the channel-close fallback path; - preserves the existing teardown order used by explicit shutdowns; - extends the lifecycle regression test to assert that the thread store receives `shutdown_thread`. Context: [original report](https://openai.slack.com/archives/C0B4NBHQGTV/p1782136364948039), [recent occurrence 1](https://openai.slack.com/archives/C0B4NBHQGTV/p1782434817895839?thread_ts=1782136364.948039&cid=C0B4NBHQGTV), [recent occurrence 2](https://openai.slack.com/archives/C0B4NBHQGTV/p1782335107474429?thread_ts=1782136364.948039&cid=C0B4NBHQGTV) ### Testing - `just test -p codex-core submission_loop_channel_close_runs_full_thread_teardown` - `just test -p codex-core --lib` (1,989 passed; 3 skipped) - `just fix -p codex-core` - `just fmt` - Native code review: no findings I also attempted `just test -p codex-core`. The new regression passed; 79 unrelated integration tests failed in the local harness, primarily because helper binaries such as `test_stdio_server` were unavailable, plus local proxy/shell timing failures. --- codex-rs/core/src/session/handlers.rs | 5 +++ codex-rs/core/src/session/tests.rs | 44 ++++++++++++++++++++++++++- 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 75fbca4c86..a483ce5e0f 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -850,6 +850,11 @@ pub(super) async fn submission_loop( if !shutdown_received { shutdown_session_runtime(&sess).await; emit_thread_stop_lifecycle(sess.as_ref()).await; + if let Some(live_thread) = sess.live_thread() + && let Err(err) = live_thread.shutdown().await + { + warn!("failed to shutdown thread persistence after submission channel closed: {err}"); + } } debug!("Agent loop exited"); } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 20f32ddfcd..46f7ca6e60 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -6971,7 +6971,7 @@ async fn shutdown_complete_does_not_append_to_thread_store_after_shutdown() { } #[tokio::test] -async fn submission_loop_channel_close_emits_thread_stop_lifecycle() { +async fn submission_loop_channel_close_runs_full_thread_teardown() { struct SessionStopMarker; struct ThreadStopMarker; @@ -6998,6 +6998,40 @@ async fn submission_loop_channel_close_emits_thread_stop_lifecycle() { } let (mut session, turn_context) = make_session_and_context().await; + let store = Arc::new(codex_thread_store::InMemoryThreadStore::default()); + let thread_store: Arc = store.clone(); + let config = session.get_config().await; + let live_thread = LiveThread::create( + Arc::clone(&thread_store), + CreateThreadParams { + session_id: session.session_id(), + thread_id: session.thread_id, + extra_config: None, + forked_from_id: None, + parent_thread_id: None, + source: SessionSource::Exec, + thread_source: None, + originator: "test_originator".to_string(), + base_instructions: BaseInstructions::default(), + dynamic_tools: Vec::new(), + selected_capability_roots: Vec::new(), + multi_agent_version: None, + initial_window_id: Uuid::now_v7().to_string(), + metadata: ThreadPersistenceMetadata { + cwd: Some(config.cwd.to_path_buf()), + model_provider: config.model_provider_id.clone(), + memory_mode: if config.memories.generate_memories { + ThreadMemoryMode::Enabled + } else { + ThreadMemoryMode::Disabled + }, + }, + }, + ) + .await + .expect("create thread persistence"); + session.services.thread_store = thread_store; + session.services.live_thread = Some(live_thread); let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); let mut builder = codex_extension_api::ExtensionRegistryBuilder::::new(); builder.thread_lifecycle_contributor(Arc::new(ThreadStopRecorder { @@ -7020,6 +7054,14 @@ async fn submission_loop_channel_close_emits_thread_stop_lifecycle() { submission_loop(session, Arc::clone(&turn_context.config), rx_sub).await; assert_eq!(1, calls.load(std::sync::atomic::Ordering::SeqCst)); + assert_eq!( + codex_thread_store::InMemoryThreadStoreCalls { + create_thread: 1, + shutdown_thread: 1, + ..Default::default() + }, + store.calls().await + ); } #[tokio::test] From 526f495f3a904e29ec1d07a927992c904c8a1f5c Mon Sep 17 00:00:00 2001 From: felixxia-oai Date: Fri, 26 Jun 2026 22:11:13 +0100 Subject: [PATCH 35/38] [codex] Classify nested MCP authentication startup errors (#30257) ## Summary - classify authentication-required RMCP startup failures, including errors nested inside `ClientInitializeError::TransportError` - let `codex-mcp` consume that classification so the existing `reauthenticationRequired` startup failure reason is emitted - add a regression test that performs real startup with an expired persisted OAuth token and no refresh token ## Why Follow-up to #29877. RMCP stores streamable HTTP initialization failures inside a dynamic transport error whose payload is not exposed through the standard Rust error source chain. The original `anyhow::Error::chain()` check therefore missed the nested `AuthError::AuthorizationRequired` seen during real MCP startup and emitted `failureReason: null`. The transport-specific inspection now lives in `codex-rmcp-client`, while `codex-mcp` consumes only the domain-level authentication-required result. This classifier does not distinguish first-time login from reauthentication; the existing auth-state logic remains responsible for that distinction. ## User impact When stored MCP OAuth credentials are expired and cannot be refreshed, app clients now receive `failureReason: "reauthenticationRequired"` on the failed startup update and can show the reconnect action. First-time login and unrelated startup failures remain unchanged. ## Validation - `just test -p codex-rmcp-client --test streamable_http_oauth_startup identifies_expired_unrefreshable_token_startup_error` - `just test -p codex-mcp startup_outcome_error_identifies_authentication_required` - `just test -p codex-mcp mcp_startup_failure_reason_requires_existing_oauth_and_auth_failure` - `cargo build -p codex-cli --bin codex` - local app-server probe emitted `failureReason: "reauthenticationRequired"` - manual end-to-end reconnect flow confirmed - `just fmt` --- codex-rs/codex-mcp/src/rmcp_client.rs | 13 +--- codex-rs/rmcp-client/src/lib.rs | 2 + codex-rs/rmcp-client/src/startup_error.rs | 46 +++++++++++ .../tests/streamable_http_oauth_startup.rs | 78 +++++++++++++++++++ 4 files changed, 128 insertions(+), 11 deletions(-) create mode 100644 codex-rs/rmcp-client/src/startup_error.rs diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index ee7dde1fe2..37819203f6 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -57,6 +57,7 @@ use codex_rmcp_client::LocalStdioServerLauncher; use codex_rmcp_client::RmcpClient; use codex_rmcp_client::StdioServerLauncher; use codex_rmcp_client::ToolWithConnectorId; +use codex_rmcp_client::is_authentication_required_error; use futures::future::BoxFuture; use futures::future::FutureExt; use futures::future::Shared; @@ -67,7 +68,6 @@ use rmcp::model::InitializeRequestParams; use rmcp::model::JsonObject; use rmcp::model::ProtocolVersion; use rmcp::model::Tool as RmcpTool; -use rmcp::transport::auth::AuthError; use tokio::time::Instant as TokioInstant; use tokio_util::sync::CancellationToken; use tracing::Instrument; @@ -552,16 +552,7 @@ impl StartupOutcomeError { impl From for StartupOutcomeError { fn from(error: anyhow::Error) -> Self { - let is_authentication_required = error.chain().any(|source| { - source - .downcast_ref::() - .is_some_and(|auth_error| { - matches!( - auth_error, - AuthError::AuthorizationRequired | AuthError::TokenExpired - ) - }) - }); + let is_authentication_required = is_authentication_required_error(&error); Self::Failed { error: error.to_string(), is_authentication_required, diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index 5be97a56a6..e57eb3686c 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -9,6 +9,7 @@ mod oauth_http_client; mod perform_oauth_login; mod program_resolver; mod rmcp_client; +mod startup_error; mod stdio_server_launcher; mod utils; @@ -40,6 +41,7 @@ pub use rmcp_client::ListToolsWithConnectorIdResult; pub use rmcp_client::RmcpClient; pub use rmcp_client::SendElicitation; pub use rmcp_client::ToolWithConnectorId; +pub use startup_error::is_authentication_required_error; pub use stdio_server_launcher::ExecutorStdioServerLauncher; pub use stdio_server_launcher::LocalStdioServerLauncher; pub use stdio_server_launcher::StdioServerLauncher; diff --git a/codex-rs/rmcp-client/src/startup_error.rs b/codex-rs/rmcp-client/src/startup_error.rs new file mode 100644 index 0000000000..c74f637e3d --- /dev/null +++ b/codex-rs/rmcp-client/src/startup_error.rs @@ -0,0 +1,46 @@ +use anyhow::Error; +use rmcp::service::ClientInitializeError; +use rmcp::transport::auth::AuthError; +use rmcp::transport::streamable_http_client::StreamableHttpError; + +use crate::http_client_adapter::StreamableHttpClientAdapterError; + +/// Returns whether an RMCP client error indicates that authentication is required. +/// +/// This does not distinguish first-time login from reauthentication. +/// Streamable HTTP initialization errors are stored inside RMCP's dynamic +/// transport error, which is not part of the standard error source chain. +pub fn is_authentication_required_error(error: &Error) -> bool { + error.chain().any(|source| { + source + .downcast_ref::() + .is_some_and(auth_error_requires_authentication) + || source + .downcast_ref::() + .is_some_and(client_initialize_error_requires_authentication) + }) +} + +fn client_initialize_error_requires_authentication(error: &ClientInitializeError) -> bool { + let ClientInitializeError::TransportError { error, .. } = error else { + return false; + }; + + error + .error + .downcast_ref::>() + .is_some_and(|error| { + matches!( + error, + StreamableHttpError::Auth(auth_error) + if auth_error_requires_authentication(auth_error) + ) + }) +} + +fn auth_error_requires_authentication(error: &AuthError) -> bool { + matches!( + error, + AuthError::AuthorizationRequired | AuthError::TokenExpired + ) +} diff --git a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs index bc5a701ed0..32f90bc41e 100644 --- a/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs +++ b/codex-rs/rmcp-client/tests/streamable_http_oauth_startup.rs @@ -13,6 +13,7 @@ use codex_rmcp_client::RmcpClient; use codex_rmcp_client::StoredOAuthTokens; use codex_rmcp_client::WrappedOAuthTokenResponse; use codex_rmcp_client::determine_streamable_http_auth_status; +use codex_rmcp_client::is_authentication_required_error; use codex_rmcp_client::save_oauth_tokens; use oauth2::AccessToken; use oauth2::RefreshToken; @@ -144,6 +145,40 @@ async fn reports_auth_status_for_persisted_credentials() -> anyhow::Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +async fn identifies_expired_unrefreshable_token_startup_error() -> anyhow::Result<()> { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/.well-known/oauth-authorization-server/mcp")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "authorization_endpoint": format!("{}/oauth/authorize", server.uri()), + "token_endpoint": format!("{}/oauth/token", server.uri()), + }))) + .expect(1) + .mount(&server) + .await; + + let codex_home = TempDir::new()?; + let status = Command::new(std::env::current_exe()?) + .args([ + "expired_unrefreshable_startup_child", + "--exact", + "--ignored", + "--nocapture", + ]) + .env("CODEX_HOME", codex_home.path()) + .env(CHILD_SERVER_URL_ENV, format!("{}/mcp", server.uri())) + .status() + .await?; + + assert!( + status.success(), + "expired OAuth startup child failed: {status}" + ); + server.verify().await; + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] #[ignore = "spawned by reports_auth_status_for_persisted_credentials"] async fn persisted_credentials_auth_status_child() -> anyhow::Result<()> { @@ -298,3 +333,46 @@ async fn oauth_startup_child() -> anyhow::Result<()> { initialize_client(&client).await?; Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 1)] +#[ignore = "spawned by identifies_expired_unrefreshable_token_startup_error"] +async fn expired_unrefreshable_startup_child() -> anyhow::Result<()> { + let server_url = std::env::var(CHILD_SERVER_URL_ENV)?; + let response = OAuthTokenResponse::new( + AccessToken::new(EXPIRED_ACCESS_TOKEN.to_string()), + BasicTokenType::Bearer, + VendorExtraTokenFields::default(), + ); + let tokens = StoredOAuthTokens { + server_name: SERVER_NAME.to_string(), + url: server_url.clone(), + client_id: "test-client-id".to_string(), + token_response: WrappedOAuthTokenResponse(response), + expires_at: Some(0), + }; + save_oauth_tokens( + SERVER_NAME, + &tokens, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + )?; + + let client = RmcpClient::new_streamable_http_client( + SERVER_NAME, + &server_url, + /*bearer_token*/ None, + /*http_headers*/ None, + /*env_http_headers*/ None, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + Environment::default_for_tests().get_http_client(), + /*auth_provider*/ None, + ) + .await?; + + let error = initialize_client(&client) + .await + .expect_err("expired token without a refresh token should fail startup"); + assert!(is_authentication_required_error(&error)); + Ok(()) +} From 6509f3148a248f5da6c03cdf008f2ad70526cd31 Mon Sep 17 00:00:00 2001 From: charlesgong-openai Date: Fri, 26 Jun 2026 17:24:46 -0400 Subject: [PATCH 36/38] [codex] Support npm marketplace plugin sources (#29375) ## Why Marketplace source deserialization treated `{"source":"npm", ...}` as unsupported. The loader logged and skipped the entry, so npm-backed plugins never appeared in `plugin list --available` and `plugin add` returned "plugin not found". Codex plugins are installed from a plugin root, not from an npm dependency tree. For npm-backed marketplace entries, Codex should fetch the published package contents without running package scripts or installing unrelated dependencies. ## What changed - Add `npm` marketplace plugin sources with `package`, optional semver `version` or version range, and optional HTTPS `registry`. - Reject unsafe npm source fields before materialization, including invalid package names, non-semver version selectors, plaintext or credential-bearing registry URLs, and registry query/fragment data. - Materialize npm plugins with `npm pack --ignore-scripts`, then unpack the resulting tarball through the existing hardened plugin bundle extractor. - Enforce npm archive and extracted-size limits, require the standard npm `package/` archive root, and verify the extracted `package.json` name matches the requested package before installing. - Keep plugin listings, install-source descriptions, CLI JSON/human output, app-server v2 `PluginSource`, TUI source summaries, regenerated schema fixtures, and app-server documentation in sync. ## Impact Marketplaces can distribute Codex plugins from public or configured private HTTPS npm registries using the same install flow as existing materialized plugin sources. `npm` must be available on `PATH` when an npm-backed plugin is installed. Fixes #27831 ## Validation - `just write-app-server-schema` - `just test -p codex-core-plugins -p codex-app-server-protocol -p codex-app-server -p codex-cli` - npm/schema/core-plugin coverage passed in the run. - The full focused command finished with `1739 passed`, `11 failed`, and `6 timed out`; the failures were unrelated local app-server environment failures from `sandbox-exec: sandbox_apply: Operation not permitted` plus one missing `test_stdio_server` helper binary. - Installed an npm-published Codex plugin package through a throwaway local marketplace and throwaway `CODEX_HOME` to exercise the real npm materialization path end to end. --- .../codex_app_server_protocol.schemas.json | 34 +++ .../codex_app_server_protocol.v2.schemas.json | 34 +++ .../json/v2/PluginInstalledResponse.json | 34 +++ .../schema/json/v2/PluginListResponse.json | 34 +++ .../schema/json/v2/PluginReadResponse.json | 34 +++ .../json/v2/PluginShareListResponse.json | 34 +++ .../schema/typescript/v2/PluginSource.ts | 10 +- .../src/protocol/v2/plugin.rs | 9 + .../src/protocol/v2/tests.rs | 17 +- .../src/request_processors/plugins.rs | 11 +- codex-rs/cli/src/plugin_cmd.rs | 30 +++ codex-rs/core-plugins/src/lib.rs | 1 + codex-rs/core-plugins/src/loader.rs | 17 ++ codex-rs/core-plugins/src/manager.rs | 68 +++-- codex-rs/core-plugins/src/marketplace.rs | 155 ++++++++++- .../core-plugins/src/marketplace_tests.rs | 243 +++++++++++++++++- codex-rs/core-plugins/src/npm_source.rs | 188 ++++++++++++++ codex-rs/core-plugins/src/npm_source_tests.rs | 111 ++++++++ codex-rs/tui/src/chatwidget/plugin_catalog.rs | 6 + ...tests__plugin_detail_popup_npm_source.snap | 20 ++ .../chatwidget/tests/popups_and_settings.rs | 45 ++++ 21 files changed, 1102 insertions(+), 33 deletions(-) create mode 100644 codex-rs/core-plugins/src/npm_source.rs create mode 100644 codex-rs/core-plugins/src/npm_source_tests.rs create mode 100644 codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugin_detail_popup_npm_source.snap 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 9f03be5b79..341c325b1d 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 @@ -14207,6 +14207,40 @@ "title": "GitPluginSource", "type": "object" }, + { + "properties": { + "package": { + "type": "string" + }, + "registry": { + "description": "Optional HTTPS registry URL. Authentication stays in the user's npm config.", + "type": [ + "string", + "null" + ] + }, + "type": { + "enum": [ + "npm" + ], + "title": "NpmPluginSourceType", + "type": "string" + }, + "version": { + "description": "Optional npm version or version range.", + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "package", + "type" + ], + "title": "NpmPluginSource", + "type": "object" + }, { "description": "The plugin is available in the remote catalog. Download metadata is kept server-side and is not exposed through the app-server API.", "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 f7037ef836..488bd6016e 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 @@ -10611,6 +10611,40 @@ "title": "GitPluginSource", "type": "object" }, + { + "properties": { + "package": { + "type": "string" + }, + "registry": { + "description": "Optional HTTPS registry URL. Authentication stays in the user's npm config.", + "type": [ + "string", + "null" + ] + }, + "type": { + "enum": [ + "npm" + ], + "title": "NpmPluginSourceType", + "type": "string" + }, + "version": { + "description": "Optional npm version or version range.", + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "package", + "type" + ], + "title": "NpmPluginSource", + "type": "object" + }, { "description": "The plugin is available in the remote catalog. Download metadata is kept server-side and is not exposed through the app-server API.", "properties": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/PluginInstalledResponse.json b/codex-rs/app-server-protocol/schema/json/v2/PluginInstalledResponse.json index ce4ae5aa0b..4c5b205538 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/PluginInstalledResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/PluginInstalledResponse.json @@ -413,6 +413,40 @@ "title": "GitPluginSource", "type": "object" }, + { + "properties": { + "package": { + "type": "string" + }, + "registry": { + "description": "Optional HTTPS registry URL. Authentication stays in the user's npm config.", + "type": [ + "string", + "null" + ] + }, + "type": { + "enum": [ + "npm" + ], + "title": "NpmPluginSourceType", + "type": "string" + }, + "version": { + "description": "Optional npm version or version range.", + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "package", + "type" + ], + "title": "NpmPluginSource", + "type": "object" + }, { "description": "The plugin is available in the remote catalog. Download metadata is kept server-side and is not exposed through the app-server API.", "properties": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/PluginListResponse.json b/codex-rs/app-server-protocol/schema/json/v2/PluginListResponse.json index 1002ad1a8d..62b9ac19ea 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/PluginListResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/PluginListResponse.json @@ -413,6 +413,40 @@ "title": "GitPluginSource", "type": "object" }, + { + "properties": { + "package": { + "type": "string" + }, + "registry": { + "description": "Optional HTTPS registry URL. Authentication stays in the user's npm config.", + "type": [ + "string", + "null" + ] + }, + "type": { + "enum": [ + "npm" + ], + "title": "NpmPluginSourceType", + "type": "string" + }, + "version": { + "description": "Optional npm version or version range.", + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "package", + "type" + ], + "title": "NpmPluginSource", + "type": "object" + }, { "description": "The plugin is available in the remote catalog. Download metadata is kept server-side and is not exposed through the app-server API.", "properties": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/PluginReadResponse.json b/codex-rs/app-server-protocol/schema/json/v2/PluginReadResponse.json index a2d2d442c5..ad63129bf8 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/PluginReadResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/PluginReadResponse.json @@ -553,6 +553,40 @@ "title": "GitPluginSource", "type": "object" }, + { + "properties": { + "package": { + "type": "string" + }, + "registry": { + "description": "Optional HTTPS registry URL. Authentication stays in the user's npm config.", + "type": [ + "string", + "null" + ] + }, + "type": { + "enum": [ + "npm" + ], + "title": "NpmPluginSourceType", + "type": "string" + }, + "version": { + "description": "Optional npm version or version range.", + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "package", + "type" + ], + "title": "NpmPluginSource", + "type": "object" + }, { "description": "The plugin is available in the remote catalog. Download metadata is kept server-side and is not exposed through the app-server API.", "properties": { diff --git a/codex-rs/app-server-protocol/schema/json/v2/PluginShareListResponse.json b/codex-rs/app-server-protocol/schema/json/v2/PluginShareListResponse.json index be97e20b65..7c6f486240 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/PluginShareListResponse.json +++ b/codex-rs/app-server-protocol/schema/json/v2/PluginShareListResponse.json @@ -369,6 +369,40 @@ "title": "GitPluginSource", "type": "object" }, + { + "properties": { + "package": { + "type": "string" + }, + "registry": { + "description": "Optional HTTPS registry URL. Authentication stays in the user's npm config.", + "type": [ + "string", + "null" + ] + }, + "type": { + "enum": [ + "npm" + ], + "title": "NpmPluginSourceType", + "type": "string" + }, + "version": { + "description": "Optional npm version or version range.", + "type": [ + "string", + "null" + ] + } + }, + "required": [ + "package", + "type" + ], + "title": "NpmPluginSource", + "type": "object" + }, { "description": "The plugin is available in the remote catalog. Download metadata is kept server-side and is not exposed through the app-server API.", "properties": { diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/PluginSource.ts b/codex-rs/app-server-protocol/schema/typescript/v2/PluginSource.ts index f6e867195d..c7ba3bcf1c 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/PluginSource.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/PluginSource.ts @@ -3,4 +3,12 @@ // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. import type { AbsolutePathBuf } from "../AbsolutePathBuf"; -export type PluginSource = { "type": "local", path: AbsolutePathBuf, } | { "type": "git", url: string, path: string | null, refName: string | null, sha: string | null, } | { "type": "remote" }; +export type PluginSource = { "type": "local", path: AbsolutePathBuf, } | { "type": "git", url: string, path: string | null, refName: string | null, sha: string | null, } | { "type": "npm", package: string, +/** + * Optional npm version or version range. + */ +version: string | null, +/** + * Optional HTTPS registry URL. Authentication stays in the user's npm config. + */ +registry: string | null, } | { "type": "remote" }; diff --git a/codex-rs/app-server-protocol/src/protocol/v2/plugin.rs b/codex-rs/app-server-protocol/src/protocol/v2/plugin.rs index dbe0b3d55d..148fe2f2a4 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/plugin.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/plugin.rs @@ -744,6 +744,15 @@ pub enum PluginSource { ref_name: Option, sha: Option, }, + #[serde(rename_all = "camelCase")] + #[ts(rename_all = "camelCase")] + Npm { + package: String, + /// Optional npm version or version range. + version: Option, + /// Optional HTTPS registry URL. Authentication stays in the user's npm config. + registry: Option, + }, /// The plugin is available in the remote catalog. Download metadata is /// kept server-side and is not exposed through the app-server API. Remote, diff --git a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs index 930441c63f..677f4470fd 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/tests.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/tests.rs @@ -2927,7 +2927,7 @@ fn skills_extra_roots_set_params_rejects_relative_roots() { } #[test] -fn plugin_source_serializes_local_git_and_remote_variants() { +fn plugin_source_serializes_local_git_npm_and_remote_variants() { let local_path = if cfg!(windows) { r"C:\plugins\linear" } else { @@ -2961,6 +2961,21 @@ fn plugin_source_serializes_local_git_and_remote_variants() { }), ); + assert_eq!( + serde_json::to_value(PluginSource::Npm { + package: "@acme/plugin".to_string(), + version: Some("^1.2.0".to_string()), + registry: Some("https://npm.example.com".to_string()), + }) + .unwrap(), + json!({ + "type": "npm", + "package": "@acme/plugin", + "version": "^1.2.0", + "registry": "https://npm.example.com", + }), + ); + assert_eq!( serde_json::to_value(PluginSource::Remote).unwrap(), json!({ diff --git a/codex-rs/app-server/src/request_processors/plugins.rs b/codex-rs/app-server/src/request_processors/plugins.rs index d123f0379c..020fdc7cbd 100644 --- a/codex-rs/app-server/src/request_processors/plugins.rs +++ b/codex-rs/app-server/src/request_processors/plugins.rs @@ -101,6 +101,15 @@ fn marketplace_plugin_source_to_info(source: MarketplacePluginSource) -> PluginS ref_name, sha, }, + MarketplacePluginSource::Npm { + package, + version, + registry, + } => PluginSource::Npm { + package, + version, + registry, + }, } } @@ -134,7 +143,7 @@ fn share_context_for_source( creator_name: None, share_principals: None, }), - MarketplacePluginSource::Git { .. } => None, + MarketplacePluginSource::Git { .. } | MarketplacePluginSource::Npm { .. } => None, } } diff --git a/codex-rs/cli/src/plugin_cmd.rs b/codex-rs/cli/src/plugin_cmd.rs index 8c9144dfff..6d17719974 100644 --- a/codex-rs/cli/src/plugin_cmd.rs +++ b/codex-rs/cli/src/plugin_cmd.rs @@ -285,6 +285,20 @@ pub async fn run_plugin_list( } parts.join(", ") } + codex_core_plugins::marketplace::MarketplacePluginSource::Npm { + package, + version, + registry, + } => { + let mut parts = vec![package.clone()]; + if let Some(version) = version { + parts.push(format!("version `{version}`")); + } + if let Some(registry) = registry { + parts.push(format!("registry `{registry}`")); + } + parts.join(", ") + } }; plugin_width = plugin_width.max(plugin.id.len()); status_width = status_width.max(state.len()); @@ -412,6 +426,13 @@ enum JsonPluginSource { #[serde(skip_serializing_if = "Option::is_none")] sha: Option, }, + Npm { + package: String, + #[serde(skip_serializing_if = "Option::is_none")] + version: Option, + #[serde(skip_serializing_if = "Option::is_none")] + registry: Option, + }, } impl JsonPluginSource { @@ -437,6 +458,15 @@ impl JsonPluginSource { ref_name, sha, } => Self::Git { url, ref_name, sha }, + MarketplacePluginSource::Npm { + package, + version, + registry, + } => Self::Npm { + package, + version, + registry, + }, } } } diff --git a/codex-rs/core-plugins/src/lib.rs b/codex-rs/core-plugins/src/lib.rs index 7ee5ddd117..d2047afa88 100644 --- a/codex-rs/core-plugins/src/lib.rs +++ b/codex-rs/core-plugins/src/lib.rs @@ -9,6 +9,7 @@ pub mod marketplace_add; mod marketplace_policy; pub mod marketplace_remove; pub mod marketplace_upgrade; +mod npm_source; mod plugin_bundle_archive; mod provider; pub mod remote; diff --git a/codex-rs/core-plugins/src/loader.rs b/codex-rs/core-plugins/src/loader.rs index 49f4580e39..2f406439ae 100644 --- a/codex-rs/core-plugins/src/loader.rs +++ b/codex-rs/core-plugins/src/loader.rs @@ -10,6 +10,7 @@ use crate::marketplace::MarketplacePluginSource; use crate::marketplace::find_marketplace_plugin; use crate::marketplace::list_marketplaces; use crate::marketplace::load_marketplace; +use crate::npm_source::materialize_npm_plugin_source; use crate::remote::REMOTE_GLOBAL_MARKETPLACE_NAME; use crate::remote::RemoteInstalledPlugin; use crate::store::PluginStore; @@ -1358,6 +1359,22 @@ pub fn materialize_marketplace_plugin_source( _tempdir: Some(tempdir), }) } + MarketplacePluginSource::Npm { + package, + version, + registry, + } => { + let (path, tempdir) = materialize_npm_plugin_source( + codex_home, + package, + version.as_deref(), + registry.as_deref(), + )?; + Ok(MaterializedMarketplacePluginSource { + path, + _tempdir: Some(tempdir), + }) + } } } diff --git a/codex-rs/core-plugins/src/manager.rs b/codex-rs/core-plugins/src/manager.rs index b357c5b536..39e5152abd 100644 --- a/codex-rs/core-plugins/src/manager.rs +++ b/codex-rs/core-plugins/src/manager.rs @@ -1578,7 +1578,7 @@ impl PluginsManager { let mut local_version = plugin.local_version; let manifest_fallback = plugin.manifest_fallback.clone(); if installed - && matches!(&plugin.source, MarketplacePluginSource::Git { .. }) + && plugin.source.is_install_materialized() && let Some(plugin_id) = plugin_id.as_ref() && let Some(plugin_root) = self.store.active_plugin_root(plugin_id) && let Some(manifest) = load_plugin_manifest(plugin_root.as_path()) @@ -1741,7 +1741,7 @@ impl PluginsManager { } })?; let plugin_key = plugin_id.as_key(); - if matches!(plugin.source, MarketplacePluginSource::Git { .. }) && !plugin.installed { + if plugin.source.is_install_materialized() && !plugin.installed { let description = remote_plugin_install_required_description(&plugin.source); return Ok(PluginDetail { id: plugin_key, @@ -1766,28 +1766,27 @@ impl PluginsManager { }); } - let source_path = - if matches!(plugin.source, MarketplacePluginSource::Git { .. }) && plugin.installed { - self.store.active_plugin_root(&plugin_id).ok_or_else(|| { - MarketplaceError::InvalidPlugin(format!( - "installed plugin cache entry is missing for {plugin_key}" - )) - })? - } else { - let codex_home = self.codex_home.clone(); - let source = plugin.source.clone(); - let materialized = tokio::task::spawn_blocking(move || { - materialize_marketplace_plugin_source(codex_home.as_path(), &source) - }) - .await - .map_err(|err| { - MarketplaceError::InvalidPlugin(format!( - "failed to materialize plugin source: {err}" - )) - })? - .map_err(MarketplaceError::InvalidPlugin)?; - materialized.path.clone() - }; + let source_path = if plugin.source.is_install_materialized() && plugin.installed { + self.store.active_plugin_root(&plugin_id).ok_or_else(|| { + MarketplaceError::InvalidPlugin(format!( + "installed plugin cache entry is missing for {plugin_key}" + )) + })? + } else { + let codex_home = self.codex_home.clone(); + let source = plugin.source.clone(); + let materialized = tokio::task::spawn_blocking(move || { + materialize_marketplace_plugin_source(codex_home.as_path(), &source) + }) + .await + .map_err(|err| { + MarketplaceError::InvalidPlugin(format!( + "failed to materialize plugin source: {err}" + )) + })? + .map_err(MarketplaceError::InvalidPlugin)?; + materialized.path.clone() + }; if !source_path.as_path().is_dir() { return Err(MarketplaceError::InvalidPlugin( "path does not exist or is not a directory".to_string(), @@ -2474,10 +2473,29 @@ pub(crate) fn remote_plugin_install_required_description( parts.join(", ") } MarketplacePluginSource::Local { path } => path.as_path().display().to_string(), + MarketplacePluginSource::Npm { + package, + version, + registry, + } => { + let mut parts = vec![package.clone()]; + if let Some(version) = version { + parts.push(format!("version `{version}`")); + } + if let Some(registry) = registry { + parts.push(format!("registry `{registry}`")); + } + parts.join(", ") + } }; + let source_kind = if matches!(source, MarketplacePluginSource::Npm { .. }) { + "an npm plugin" + } else { + "a cross-repo plugin" + }; format!( - "This is a cross-repo plugin. Install it to view more detailed information. The source of the plugin is {source_description}." + "This is {source_kind}. Install it to view more detailed information. The source of the plugin is {source_description}." ) } diff --git a/codex-rs/core-plugins/src/marketplace.rs b/codex-rs/core-plugins/src/marketplace.rs index 3190e0449a..6475d37876 100644 --- a/codex-rs/core-plugins/src/marketplace.rs +++ b/codex-rs/core-plugins/src/marketplace.rs @@ -86,8 +86,8 @@ impl MarketplacePluginManifestFallback { } pub(crate) fn parse_for_listing(&self) -> Option { - // Git sources have no plugin root before install. Parse against a host-native synthetic - // absolute root, then discard path-bearing fields so listings expose metadata only. + // Materialized sources have no plugin root before install. Parse against a host-native + // synthetic absolute root, then discard path-bearing fields so listings expose metadata only. let plugin_root = Path::new(if cfg!(windows) { r"C:\" } else { "/" }); let mut manifest = crate::manifest::parse_plugin_manifest( plugin_root, @@ -133,6 +133,23 @@ pub enum MarketplacePluginSource { ref_name: Option, sha: Option, }, + Npm { + package: String, + version: Option, + registry: Option, + }, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum NpmPackageScope { + Scoped, + Unscoped, +} + +impl MarketplacePluginSource { + pub(crate) fn is_install_materialized(&self) -> bool { + matches!(self, Self::Git { .. } | Self::Npm { .. }) + } } #[derive(Debug, Clone, PartialEq, Eq)] @@ -513,10 +530,12 @@ fn resolve_marketplace_plugin_entry( None } } - MarketplacePluginSource::Git { .. } if manifest_fallback.has_metadata => { + MarketplacePluginSource::Git { .. } | MarketplacePluginSource::Npm { .. } + if manifest_fallback.has_metadata => + { manifest_fallback.parse_for_listing() } - MarketplacePluginSource::Git { .. } => None, + MarketplacePluginSource::Git { .. } | MarketplacePluginSource::Npm { .. } => None, }; let interface = plugin_interface_with_marketplace_category( manifest @@ -610,6 +629,17 @@ fn resolve_plugin_source( ref_name: normalize_optional_git_selector(&ref_name), sha: normalize_optional_git_selector(&sha), }), + RawMarketplaceManifestPluginSource::Object( + RawMarketplaceManifestPluginSourceObject::Npm { + package, + version, + registry, + }, + ) => Ok(MarketplacePluginSource::Npm { + package: normalize_npm_package(marketplace_path, &package)?, + version: normalize_optional_npm_version(marketplace_path, version)?, + registry: normalize_optional_npm_registry(marketplace_path, registry)?, + }), RawMarketplaceManifestPluginSource::Unsupported(_) => { unreachable!("unsupported plugin sources should be filtered before resolution") } @@ -748,6 +778,118 @@ fn normalize_optional_git_selector(value: &Option) -> Option { .map(str::to_string) } +fn normalize_npm_package( + marketplace_path: &AbsolutePathBuf, + package: &str, +) -> Result { + let package = package.trim(); + let package_scope = if package.starts_with('@') { + NpmPackageScope::Scoped + } else { + NpmPackageScope::Unscoped + }; + let segments = if let Some(scoped_package) = package.strip_prefix('@') { + scoped_package.split('/').collect::>() + } else { + package.split('/').collect::>() + }; + let expected_segments = match package_scope { + NpmPackageScope::Scoped => 2, + NpmPackageScope::Unscoped => 1, + }; + if package.is_empty() + || segments.len() != expected_segments + || segments + .iter() + .any(|segment| !is_valid_npm_package_segment(segment, package_scope)) + { + return Err(MarketplaceError::InvalidMarketplaceFile { + path: marketplace_path.to_path_buf(), + message: format!("invalid npm plugin source package: {package}"), + }); + } + Ok(package.to_string()) +} + +fn is_valid_npm_package_segment(segment: &str, package_scope: NpmPackageScope) -> bool { + !segment.is_empty() + && segment != "." + && segment != ".." + && (package_scope == NpmPackageScope::Scoped + || !matches!(segment.chars().next(), Some('.' | '_'))) + && segment + .chars() + .all(|ch| ch.is_ascii_alphanumeric() || matches!(ch, '-' | '_' | '.')) +} + +fn normalize_optional_npm_version( + marketplace_path: &AbsolutePathBuf, + version: Option, +) -> Result, MarketplaceError> { + let Some(version) = normalize_optional_npm_source_field(marketplace_path, version, "version")? + else { + return Ok(None); + }; + if !is_registry_npm_version_selector(&version) { + return Err(MarketplaceError::InvalidMarketplaceFile { + path: marketplace_path.to_path_buf(), + message: format!("npm plugin source version must use the registry: {version}"), + }); + } + Ok(Some(version)) +} + +fn is_registry_npm_version_selector(version: &str) -> bool { + version != "." && version != ".." && !version.chars().any(|ch| matches!(ch, '/' | '\\' | ':')) +} + +fn normalize_optional_npm_registry( + marketplace_path: &AbsolutePathBuf, + registry: Option, +) -> Result, MarketplaceError> { + let Some(registry) = + normalize_optional_npm_source_field(marketplace_path, registry, "registry")? + else { + return Ok(None); + }; + let parsed = + url::Url::parse(®istry).map_err(|_| MarketplaceError::InvalidMarketplaceFile { + path: marketplace_path.to_path_buf(), + message: format!("invalid npm plugin source registry: {registry}"), + })?; + if parsed.scheme() != "https" + || parsed.host_str().is_none() + || !parsed.username().is_empty() + || parsed.password().is_some() + || parsed.query().is_some() + || parsed.fragment().is_some() + { + return Err(MarketplaceError::InvalidMarketplaceFile { + path: marketplace_path.to_path_buf(), + message: format!("invalid npm plugin source registry: {registry}"), + }); + } + Ok(Some(registry)) +} + +fn normalize_optional_npm_source_field( + marketplace_path: &AbsolutePathBuf, + value: Option, + field: &str, +) -> Result, MarketplaceError> { + let Some(value) = value else { + return Ok(None); + }; + let value = value.trim(); + if value.is_empty() { + return Err(MarketplaceError::InvalidMarketplaceFile { + path: marketplace_path.to_path_buf(), + message: format!("npm plugin source {field} must not be empty"), + }); + } + Ok(Some(value.to_string())) +} + fn normalize_github_git_url(url: &str) -> String { if url.starts_with("https://github.com/") && !url.ends_with(".git") { format!("{url}.git") @@ -886,6 +1028,11 @@ enum RawMarketplaceManifestPluginSourceObject { ref_name: Option, sha: Option, }, + Npm { + package: String, + version: Option, + registry: Option, + }, } fn resolve_marketplace_interface( diff --git a/codex-rs/core-plugins/src/marketplace_tests.rs b/codex-rs/core-plugins/src/marketplace_tests.rs index 3ea2c9d606..d30e9ad1da 100644 --- a/codex-rs/core-plugins/src/marketplace_tests.rs +++ b/codex-rs/core-plugins/src/marketplace_tests.rs @@ -6,7 +6,6 @@ use tempfile::tempdir; const ALTERNATE_MARKETPLACE_RELATIVE_PATH: &str = ".claude-plugin/marketplace.json"; const ALTERNATE_PLUGIN_MANIFEST_RELATIVE_PATH: &str = ".claude-plugin/plugin.json"; - fn write_alternate_marketplace(repo_root: &Path, contents: &str) -> AbsolutePathBuf { let marketplace_path = repo_root.join(ALTERNATE_MARKETPLACE_RELATIVE_PATH); fs::create_dir_all(marketplace_path.parent().unwrap()).unwrap(); @@ -225,6 +224,248 @@ fn find_marketplace_plugin_omits_interface_asset_paths_for_git_sources() { assert!(interface.screenshots.is_empty()); } +#[test] +fn find_marketplace_plugin_supports_npm_sources() { + let tmp = tempdir().unwrap(); + let repo_root = tmp.path().join("repo"); + fs::create_dir_all(repo_root.join(".git")).unwrap(); + fs::create_dir_all(repo_root.join(".agents/plugins")).unwrap(); + fs::write( + repo_root.join(".agents/plugins/marketplace.json"), + r#"{ + "name": "codex-curated", + "plugins": [ + { + "name": "npm-plugin", + "source": { + "source": "npm", + "package": "@acme/codex-plugin", + "version": "^1.2.0", + "registry": "https://npm.example.com" + } + } + ] +}"#, + ) + .unwrap(); + + let resolved = find_marketplace_plugin( + &AbsolutePathBuf::try_from(repo_root.join(".agents/plugins/marketplace.json")).unwrap(), + "npm-plugin", + ) + .unwrap(); + + assert_eq!( + resolved, + ResolvedMarketplacePlugin { + plugin_id: PluginId::new("npm-plugin".to_string(), "codex-curated".to_string()) + .unwrap(), + source: MarketplacePluginSource::Npm { + package: "@acme/codex-plugin".to_string(), + version: Some("^1.2.0".to_string()), + registry: Some("https://npm.example.com".to_string()), + }, + policy: MarketplacePluginPolicy { + installation: MarketplacePluginInstallPolicy::Available, + authentication: MarketplacePluginAuthPolicy::OnInstall, + products: None, + }, + interface: None, + manifest: None, + manifest_fallback: minimal_manifest_fallback("npm-plugin"), + } + ); +} + +#[test] +fn find_marketplace_plugin_skips_unsafe_npm_sources() { + let tmp = tempdir().unwrap(); + let repo_root = tmp.path().join("repo"); + fs::create_dir_all(repo_root.join(".git")).unwrap(); + let marketplace_path = write_alternate_marketplace( + &repo_root, + r#"{ + "name": "codex-curated", + "plugins": [ + { + "name": "remote-version", + "source": { + "source": "npm", + "package": "@acme/codex-plugin", + "version": "https://attacker.example/plugin.tgz", + "registry": "https://npm.example.com" + } + }, + { + "name": "local-version", + "source": { + "source": "npm", + "package": "@acme/codex-plugin", + "version": ".", + "registry": "https://npm.example.com" + } + }, + { + "name": "plaintext-registry", + "source": { + "source": "npm", + "package": "@acme/codex-plugin", + "version": "1.2.0", + "registry": "http://npm.example.com" + } + }, + { + "name": "credential-registry", + "source": { + "source": "npm", + "package": "@acme/codex-plugin", + "version": "1.2.0", + "registry": "https://user:password@npm.example.com" + } + }, + { + "name": "dot-package", + "source": { + "source": "npm", + "package": ".codex-plugin", + "registry": "https://npm.example.com" + } + }, + { + "name": "underscore-package", + "source": { + "source": "npm", + "package": "_codex-plugin", + "registry": "https://npm.example.com" + } + } + ] +}"#, + ); + + assert_eq!( + load_marketplace(&marketplace_path).unwrap().plugins, + Vec::new() + ); +} + +#[test] +fn find_marketplace_plugin_supports_npm_registry_version_selectors() { + let tmp = tempdir().unwrap(); + let repo_root = tmp.path().join("repo"); + fs::create_dir_all(repo_root.join(".git")).unwrap(); + let marketplace_path = write_alternate_marketplace( + &repo_root, + r#"{ + "name": "codex-curated", + "plugins": [ + { + "name": "dist-tag", + "source": { + "source": "npm", + "package": "@acme/codex-plugin", + "version": "latest" + } + }, + { + "name": "comparator-range", + "source": { + "source": "npm", + "package": "@acme/codex-plugin", + "version": ">=1.2.7 <1.3.0" + } + }, + { + "name": "x-range", + "source": { + "source": "npm", + "package": "@acme/codex-plugin", + "version": "1.2.x" + } + }, + { + "name": "or-range", + "source": { + "source": "npm", + "package": "@acme/codex-plugin", + "version": "1.2.7 || >=1.2.9 <2.0.0" + } + } + ] +}"#, + ); + + assert_eq!( + load_marketplace(&marketplace_path) + .unwrap() + .plugins + .into_iter() + .map(|plugin| plugin.source) + .collect::>(), + vec![ + MarketplacePluginSource::Npm { + package: "@acme/codex-plugin".to_string(), + version: Some("latest".to_string()), + registry: None, + }, + MarketplacePluginSource::Npm { + package: "@acme/codex-plugin".to_string(), + version: Some(">=1.2.7 <1.3.0".to_string()), + registry: None, + }, + MarketplacePluginSource::Npm { + package: "@acme/codex-plugin".to_string(), + version: Some("1.2.x".to_string()), + registry: None, + }, + MarketplacePluginSource::Npm { + package: "@acme/codex-plugin".to_string(), + version: Some("1.2.7 || >=1.2.9 <2.0.0".to_string()), + registry: None, + }, + ] + ); +} + +#[test] +fn find_marketplace_plugin_supports_npm_sources_without_optional_fields() { + let tmp = tempdir().unwrap(); + let repo_root = tmp.path().join("repo"); + fs::create_dir_all(repo_root.join(".git")).unwrap(); + fs::create_dir_all(repo_root.join(".agents/plugins")).unwrap(); + fs::write( + repo_root.join(".agents/plugins/marketplace.json"), + r#"{ + "name": "codex-curated", + "plugins": [ + { + "name": "npm-plugin", + "source": { + "source": "npm", + "package": "@acme/codex-plugin" + } + } + ] +}"#, + ) + .unwrap(); + + let resolved = find_marketplace_plugin( + &AbsolutePathBuf::try_from(repo_root.join(".agents/plugins/marketplace.json")).unwrap(), + "npm-plugin", + ) + .unwrap(); + + assert_eq!( + resolved.source, + MarketplacePluginSource::Npm { + package: "@acme/codex-plugin".to_string(), + version: None, + registry: None, + } + ); +} + #[test] fn find_marketplace_plugin_builds_manifest_fallback_from_entry() { let tmp = tempdir().unwrap(); diff --git a/codex-rs/core-plugins/src/npm_source.rs b/codex-rs/core-plugins/src/npm_source.rs new file mode 100644 index 0000000000..3d728bd0f7 --- /dev/null +++ b/codex-rs/core-plugins/src/npm_source.rs @@ -0,0 +1,188 @@ +use crate::plugin_bundle_archive::unpack_plugin_bundle_tar_gz; +use codex_utils_absolute_path::AbsolutePathBuf; +use serde::Deserialize; +use std::ffi::OsStr; +use std::fs; +use std::path::Path; +use std::path::PathBuf; +use std::process::Command; +use tempfile::TempDir; + +const NPM_PLUGIN_SOURCE_STAGING_DIR: &str = "plugins/.marketplace-plugin-source-staging"; +const NPM_PLUGIN_SOURCE_MAX_ARCHIVE_BYTES: u64 = 50 * 1024 * 1024; +const NPM_PLUGIN_SOURCE_MAX_EXTRACTED_BYTES: u64 = 250 * 1024 * 1024; +const NPM_PACKAGE_ARCHIVE_ROOT: &str = "package"; + +pub(crate) fn materialize_npm_plugin_source( + codex_home: &Path, + package: &str, + version: Option<&str>, + registry: Option<&str>, +) -> Result<(AbsolutePathBuf, TempDir), String> { + materialize_npm_plugin_source_with_command( + codex_home, + package, + version, + registry, + OsStr::new(npm_command()), + ) +} + +fn materialize_npm_plugin_source_with_command( + codex_home: &Path, + package: &str, + version: Option<&str>, + registry: Option<&str>, + npm_command: &OsStr, +) -> Result<(AbsolutePathBuf, TempDir), String> { + let staging_root = codex_home.join(NPM_PLUGIN_SOURCE_STAGING_DIR); + fs::create_dir_all(&staging_root).map_err(|err| { + format!( + "failed to create marketplace plugin source staging directory {}: {err}", + staging_root.display() + ) + })?; + let tempdir = tempfile::Builder::new() + .prefix("marketplace-plugin-source-") + .tempdir_in(&staging_root) + .map_err(|err| { + format!( + "failed to create marketplace plugin source staging directory in {}: {err}", + staging_root.display() + ) + })?; + + pack_npm_package(tempdir.path(), package, version, registry, npm_command)?; + let archive_path = find_npm_package_archive(tempdir.path())?; + let archive_bytes = read_npm_package_archive(&archive_path)?; + + let extraction_root = tempdir.path().join("extracted"); + unpack_plugin_bundle_tar_gz( + &archive_bytes, + &extraction_root, + NPM_PLUGIN_SOURCE_MAX_EXTRACTED_BYTES, + ) + .map_err(|err| format!("failed to extract npm plugin package: {err}"))?; + let plugin_root = extraction_root.join(NPM_PACKAGE_ARCHIVE_ROOT); + if !plugin_root.is_dir() { + return Err(format!( + "npm pack completed without creating plugin package directory {}", + plugin_root.display() + )); + } + validate_npm_package_metadata(&plugin_root, package)?; + let plugin_root = AbsolutePathBuf::try_from(plugin_root) + .map_err(|err| format!("failed to resolve materialized plugin source path: {err}"))?; + Ok((plugin_root, tempdir)) +} + +fn pack_npm_package( + destination: &Path, + package: &str, + version: Option<&str>, + registry: Option<&str>, + npm_command: &OsStr, +) -> Result<(), String> { + let package_spec = version.map_or_else( + || package.to_string(), + |version| format!("{package}@{version}"), + ); + let mut command = Command::new(npm_command); + command + .current_dir(destination) + .arg("pack") + .arg("--ignore-scripts") + .arg("--pack-destination") + .arg(destination); + if let Some(registry) = registry { + command.arg("--registry").arg(registry); + } + command.arg("--").arg(package_spec); + + let output = command + .output() + .map_err(|err| format!("failed to run npm pack: {err}"))?; + if output.status.success() { + return Ok(()); + } + + Err(format!( + "npm pack failed with status {}\nstdout:\n{}\nstderr:\n{}", + output.status, + String::from_utf8_lossy(&output.stdout).trim(), + String::from_utf8_lossy(&output.stderr).trim() + )) +} + +fn find_npm_package_archive(destination: &Path) -> Result { + let mut archives = fs::read_dir(destination) + .map_err(|err| format!("failed to read npm pack destination: {err}"))? + .filter_map(std::result::Result::ok) + .filter_map(|entry| { + let path = entry.path(); + let is_file = entry.file_type().is_ok_and(|file_type| file_type.is_file()); + (is_file && path.extension() == Some(OsStr::new("tgz"))).then_some(path) + }) + .collect::>(); + if archives.len() != 1 { + return Err(format!( + "npm pack completed with {} package archives; expected exactly one", + archives.len() + )); + } + Ok(archives.remove(0)) +} + +fn read_npm_package_archive(archive_path: &Path) -> Result, String> { + let archive_size = fs::metadata(archive_path) + .map_err(|err| format!("failed to inspect npm package archive: {err}"))? + .len(); + if archive_size > NPM_PLUGIN_SOURCE_MAX_ARCHIVE_BYTES { + return Err(format!( + "npm package archive is {archive_size} bytes, exceeding maximum size of {NPM_PLUGIN_SOURCE_MAX_ARCHIVE_BYTES} bytes" + )); + } + fs::read(archive_path).map_err(|err| format!("failed to read npm package archive: {err}")) +} + +fn validate_npm_package_metadata(plugin_root: &Path, package: &str) -> Result<(), String> { + #[derive(Deserialize)] + struct NpmPackageMetadata { + name: String, + } + + let package_json_path = plugin_root.join("package.json"); + let package_json = fs::read_to_string(&package_json_path).map_err(|err| { + format!( + "failed to read npm plugin package metadata {}: {err}", + package_json_path.display() + ) + })?; + let metadata: NpmPackageMetadata = serde_json::from_str(&package_json).map_err(|err| { + format!( + "failed to parse npm plugin package metadata {}: {err}", + package_json_path.display() + ) + })?; + if metadata.name != package { + return Err(format!( + "npm plugin package name '{}' does not match requested package '{package}'", + metadata.name + )); + } + Ok(()) +} + +#[cfg(windows)] +fn npm_command() -> &'static str { + "npm.cmd" +} + +#[cfg(not(windows))] +fn npm_command() -> &'static str { + "npm" +} + +#[cfg(all(test, unix))] +#[path = "npm_source_tests.rs"] +mod tests; diff --git a/codex-rs/core-plugins/src/npm_source_tests.rs b/codex-rs/core-plugins/src/npm_source_tests.rs new file mode 100644 index 0000000000..3c2dcd94ac --- /dev/null +++ b/codex-rs/core-plugins/src/npm_source_tests.rs @@ -0,0 +1,111 @@ +use super::*; +use flate2::Compression; +use flate2::write::GzEncoder; +use pretty_assertions::assert_eq; +use std::io::Cursor; +use std::io::Write; + +#[cfg(unix)] +#[test] +fn materialize_npm_plugin_source_uses_packed_package_root() { + use std::os::unix::fs::PermissionsExt; + + let codex_home = tempfile::tempdir().expect("create codex home"); + let fake_npm_dir = tempfile::tempdir().expect("create fake npm directory"); + let archive_bytes = + npm_package_archive_bytes("@acme/plugin", "1.2.0").expect("build fixture archive"); + let archive_path = fake_npm_dir.path().join("fixture.tgz"); + fs::write(&archive_path, &archive_bytes).expect("write fixture archive"); + let fake_npm = fake_npm_dir.path().join("npm"); + fs::write( + &fake_npm, + format!( + r#"#!/bin/sh +destination="" +previous="" +for argument in "$@"; do + if [ "$previous" = "--pack-destination" ]; then + destination="$argument" + fi + previous="$argument" +done +cp "{}" "$destination/acme-plugin-1.2.0.tgz" +printf '%s\n' "$@" > "$destination/args.txt" +pwd > "$destination/pwd.txt" +"#, + archive_path.display() + ), + ) + .expect("write fake npm"); + let mut permissions = fs::metadata(&fake_npm) + .expect("read fake npm metadata") + .permissions(); + permissions.set_mode(0o755); + fs::set_permissions(&fake_npm, permissions).expect("make fake npm executable"); + + let (plugin_root, tempdir) = materialize_npm_plugin_source_with_command( + codex_home.path(), + "@acme/plugin", + Some("^1.2.0"), + Some("https://npm.example.com"), + fake_npm.as_os_str(), + ) + .expect("materialize npm source"); + + assert_eq!( + plugin_root.as_path(), + tempdir.path().join("extracted/package") + ); + assert!( + plugin_root + .as_path() + .join(".codex-plugin/plugin.json") + .is_file() + ); + let args = fs::read_to_string(tempdir.path().join("args.txt")).expect("read npm arguments"); + assert!(args.contains("pack")); + assert!(args.contains("--ignore-scripts")); + assert!(args.contains("--registry")); + assert!(args.contains("https://npm.example.com")); + assert!(args.contains("@acme/plugin@^1.2.0")); + assert!(!args.contains("install")); + let npm_working_directory = fs::canonicalize( + fs::read_to_string(tempdir.path().join("pwd.txt")) + .expect("read npm working directory") + .trim(), + ) + .expect("canonicalize npm working directory"); + assert_eq!( + npm_working_directory, + fs::canonicalize(tempdir.path()).expect("canonicalize tempdir") + ); +} + +fn npm_package_archive_bytes(package: &str, version: &str) -> std::io::Result> { + let encoder = GzEncoder::new(Vec::new(), Compression::default()); + let mut archive = tar::Builder::new(encoder); + append_archive_file( + &mut archive, + "package/package.json", + format!(r#"{{"name":"{package}","version":"{version}"}}"#).as_bytes(), + )?; + append_archive_file( + &mut archive, + "package/.codex-plugin/plugin.json", + br#"{"name":"plugin"}"#, + )?; + let encoder = archive.into_inner()?; + encoder.finish() +} + +fn append_archive_file( + archive: &mut tar::Builder, + path: &str, + contents: &[u8], +) -> std::io::Result<()> { + let mut header = tar::Header::new_gnu(); + header.set_size(contents.len() as u64); + header.set_mode(0o644); + header.set_cksum(); + archive.append_data(&mut header, path, Cursor::new(contents)) +} diff --git a/codex-rs/tui/src/chatwidget/plugin_catalog.rs b/codex-rs/tui/src/chatwidget/plugin_catalog.rs index cb92ea7d31..40dfca0ce5 100644 --- a/codex-rs/tui/src/chatwidget/plugin_catalog.rs +++ b/codex-rs/tui/src/chatwidget/plugin_catalog.rs @@ -1474,6 +1474,12 @@ fn plugin_source_summary(plugin: &PluginDetail) -> String { Some(ref_name) => format!("Git · {url}@{ref_name}"), None => format!("Git · {url}"), }, + PluginSource::Npm { + package, version, .. + } => match version { + Some(version) => format!("npm · {package}@{version}"), + None => format!("npm · {package}"), + }, PluginSource::Remote => { let marketplace_label = MarketplaceProduct::from_marketplace_name(&plugin.marketplace_name) diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugin_detail_popup_npm_source.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugin_detail_popup_npm_source.snap new file mode 100644 index 0000000000..cb07253718 --- /dev/null +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__plugin_detail_popup_npm_source.snap @@ -0,0 +1,20 @@ +--- +source: tui/src/chatwidget/tests/popups_and_settings.rs +expression: strip_osc8_for_snapshot(&popup) +--- + Plugins + Figma · Can be installed · ChatGPT Marketplace + Data shared with this app is subject to the app's terms of service and privacy policy. Learn + more. + Turn Figma files into implementation context. + +› 1. Back to plugins Return to the plugin list. + 2. Install plugin Install this plugin now. + Source npm · @acme/figma-plugin@^1.2.0 + Auth Auth on install + Skills design-review, extract-copy + Hooks PreToolUse (1), Stop (2) + Apps Figma, Slack + MCP Servers figma-mcp, docs-mcp + + Press esc to close. diff --git a/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs b/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs index 7d162d41e4..49b9ec7030 100644 --- a/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs +++ b/codex-rs/tui/src/chatwidget/tests/popups_and_settings.rs @@ -664,6 +664,51 @@ async fn plugin_detail_popup_snapshot_labels_personal_marketplace_as_local() { ); } +#[tokio::test] +async fn plugin_detail_popup_snapshot_shows_npm_source() { + let (mut chat, _rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; + chat.set_feature_enabled(Feature::Plugins, /*enabled*/ true); + + let mut summary = plugins_test_summary( + "plugin-figma", + "figma", + Some("Figma"), + Some("Design handoff."), + /*installed*/ false, + /*enabled*/ true, + PluginInstallPolicy::Available, + ); + summary.source = PluginSource::Npm { + package: "@acme/figma-plugin".to_string(), + version: Some("^1.2.0".to_string()), + registry: Some("https://npm.example.com".to_string()), + }; + let response = plugins_test_response(vec![plugins_test_curated_marketplace(vec![ + summary.clone(), + ])]); + let cwd = chat.config.cwd.clone(); + chat.on_plugins_loaded(cwd.to_path_buf(), Ok(response)); + chat.add_plugins_output(); + let plugin = plugins_test_detail( + summary, + Some("Turn Figma files into implementation context."), + &["design-review", "extract-copy"], + &[ + (codex_app_server_protocol::HookEventName::PreToolUse, 1), + (codex_app_server_protocol::HookEventName::Stop, 2), + ], + &["Figma", "Slack"], + &["figma-mcp", "docs-mcp"], + ); + chat.on_plugin_detail_loaded(cwd.to_path_buf(), Ok(PluginReadResponse { plugin })); + + let popup = render_bottom_popup(&chat, /*width*/ 100); + assert_chatwidget_snapshot!( + "plugin_detail_popup_npm_source", + strip_osc8_for_snapshot(&popup) + ); +} + #[tokio::test] async fn plugin_detail_popup_distinguishes_admin_installed_from_enabled() { let (mut chat, _rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await; From 1168254bd9366faab73835a2cf8eb43d00acad78 Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Fri, 26 Jun 2026 15:07:05 -0700 Subject: [PATCH 37/38] [codex] group blocking and postmerge CI workflows (#30146) ## Why It's hard to change the set of required jobs when they're managed in the GitHub UI, and when each workflow is responsible for choosing it's own scheduling it's easy to end up with skew between what we enforce on PRs vs. on main. ## What - add a `blocking-ci` caller workflow, triggered by pull requests and pushes to `main`, for Bazel, blob size, cargo-deny, Codespell, `repo-checks`, rust CI, and SDK CI - add an `always()` terminal job named `CI required` that fails unless every called workflow succeeds - add a `postmerge-ci` caller workflow for `rust-ci-full` and `v8-canary`, with a terminal `Postmerge CI results` job - centralize V8 relevance detection in `v8_canary_changes.py`; unrelated PR and postmerge runs execute metadata only and skip the expensive build matrices - leave `v8-canary` outside the blocking gate and leave the external `cla` check independent ## Rollout A repository admin must replace the existing required GitHub Actions contexts with `CI required` in the main-branch ruleset. Retain `cla` as a separate required check. Until that change is coordinated, this PR cannot satisfy the old standalone check names. In-flight PRs will need to be rebased after this lands. --- .github/scripts/check_ci_results.py | 34 +++++++ .github/scripts/v8_canary_changes.py | 96 +++++++++++++++++-- .github/workflows/bazel.yml | 5 +- .github/workflows/blob-size-policy.yml | 19 +++- .github/workflows/blocking-ci.yml | 79 +++++++++++++++ .github/workflows/cargo-deny.yml | 5 +- .github/workflows/codespell.yml | 5 +- .github/workflows/postmerge-ci.yml | 40 ++++++++ .github/workflows/{ci.yml => repo-checks.yml} | 5 +- .github/workflows/rust-ci-full.yml | 4 +- .github/workflows/rust-ci.yml | 2 +- .github/workflows/sdk.yml | 4 +- .github/workflows/v8-canary.yml | 54 +++-------- 13 files changed, 280 insertions(+), 72 deletions(-) create mode 100644 .github/scripts/check_ci_results.py create mode 100644 .github/workflows/blocking-ci.yml create mode 100644 .github/workflows/postmerge-ci.yml rename .github/workflows/{ci.yml => repo-checks.yml} (98%) diff --git a/.github/scripts/check_ci_results.py b/.github/scripts/check_ci_results.py new file mode 100644 index 0000000000..36d1eed3b4 --- /dev/null +++ b/.github/scripts/check_ci_results.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 + +"""Fail a terminal CI job unless every serialized dependency succeeded. + +Parent workflows pass GitHub's `toJSON(needs)` object through the NEEDS +environment variable. Treat skipped and cancelled dependencies as failures too: +for a required fan-in job, only an explicit success is safe to accept. +""" + +import json +import os + + +def main() -> None: + # Keep result policy in one script so blocking-ci and postmerge-ci cannot + # drift in how they interpret dependency conclusions. + needs = json.loads(os.environ["NEEDS"]) + failures = sorted( + (name, dependency["result"]) + for name, dependency in needs.items() + if dependency["result"] != "success" + ) + + if failures: + print("CI dependencies did not succeed:") + for name, result in failures: + print(f"{name}: {result}") + raise SystemExit(1) + + print("All CI dependencies succeeded.") + + +if __name__ == "__main__": + main() diff --git a/.github/scripts/v8_canary_changes.py b/.github/scripts/v8_canary_changes.py index 6d9693c60a..0acc4ec3f6 100644 --- a/.github/scripts/v8_canary_changes.py +++ b/.github/scripts/v8_canary_changes.py @@ -1,12 +1,44 @@ #!/usr/bin/env python3 +"""Decide which V8 canary work is needed for a commit range. + +The workflow deliberately has no trigger-level path filters because it is both +directly triggered for pull requests and called by postmerge-ci. Keeping the +patterns here gives those entrypoints one source of truth; unrelated events +still run metadata but skip the expensive build matrices. +""" + import argparse import subprocess import tomllib +from fnmatch import fnmatchcase from pathlib import Path ROOT = Path(__file__).resolve().parents[2] +# These patterns replace the old pull_request/push path filters. Include parent +# workflow changes because they can alter whether the canary is invoked. +CANARY_PATH_PATTERNS = { + ".bazelrc", + ".github/actions/setup-bazel-ci/**", + ".github/scripts/run_bazel_with_buildbuddy.py", + ".github/scripts/rusty_v8_bazel.py", + ".github/scripts/rusty_v8_module_bazel.py", + ".github/scripts/v8_canary_changes.py", + ".github/workflows/postmerge-ci.yml", + ".github/workflows/rusty-v8-release.yml", + ".github/workflows/v8-canary.yml", + "MODULE.bazel", + "MODULE.bazel.lock", + "codex-rs/Cargo.toml", + "patches/BUILD.bazel", + "patches/llvm_*.patch", + "patches/rules_cc_*.patch", + "patches/v8_*.patch", + "third_party/v8/**", +} +# Windows source builds are a narrower, more expensive subset of the canary. +# A V8 version change also requires them even when no path below changed. WINDOWS_SOURCE_BUILD_PATHS = { ".github/scripts/rusty_v8_bazel.py", ".github/scripts/rusty_v8_module_bazel.py", @@ -16,6 +48,30 @@ WINDOWS_SOURCE_BUILD_PATHS = { } +def matching_canary_paths(changed_files: set[str]) -> set[str]: + """Return changed paths that require the general V8 build matrix.""" + return { + path + for path in changed_files + if any(fnmatchcase(path, pattern) for pattern in CANARY_PATH_PATTERNS) + } + + +def canary_required( + changed_files: set[str], + base_v8_version: str, + head_v8_version: str, + *, + force: bool = False, +) -> bool: + """Return whether the general V8 build matrix should run.""" + return ( + force + or base_v8_version != head_v8_version + or bool(matching_canary_paths(changed_files)) + ) + + def resolved_v8_version(cargo_lock: bytes) -> str: versions = sorted( { @@ -36,6 +92,7 @@ def windows_source_required( *, force: bool = False, ) -> bool: + """Return whether Windows must rebuild rusty_v8 from source.""" return ( force or base_v8_version != head_v8_version @@ -58,6 +115,8 @@ def merge_base(base: str, head: str, *, root: Path = ROOT) -> str: def changed_files(base: str, head: str, *, root: Path = ROOT) -> set[str]: + # Three-dot diff gives PRs merge-base semantics while remaining equivalent + # to before/after for ordinary linear pushes to main. output = git_output( "diff", "--name-only", @@ -79,25 +138,44 @@ def parse_args() -> argparse.Namespace: def main() -> None: args = parse_args() if args.force: - required = True - reason = "manual workflow dispatch" + # workflow_dispatch has no comparison range, and callers use it as a + # manual retry path, so it intentionally runs every variant. + canary = True + canary_reason = "manual workflow dispatch" + windows_source = True + windows_source_reason = "manual workflow dispatch" elif not args.base or not args.head: raise SystemExit("--base and --head are required unless --force is set") else: files = changed_files(args.base, args.head) base_version = v8_version_at_revision(merge_base(args.base, args.head)) head_version = v8_version_at_revision(args.head) - required = windows_source_required(files, base_version, head_version) + + matched_canary_paths = sorted(matching_canary_paths(files)) + canary = canary_required(files, base_version, head_version) + windows_source = windows_source_required(files, base_version, head_version) if base_version != head_version: - reason = f"v8 version changed from {base_version} to {head_version}" + canary_reason = ( + f"v8 version changed from {base_version} to {head_version}" + ) + windows_source_reason = canary_reason else: - matched_paths = sorted(files & WINDOWS_SOURCE_BUILD_PATHS) - reason = ( - ", ".join(matched_paths) if matched_paths else "no relevant changes" + canary_reason = ( + ", ".join(matched_canary_paths) + if matched_canary_paths + else "no relevant changes" + ) + matched_windows_paths = sorted(files & WINDOWS_SOURCE_BUILD_PATHS) + windows_source_reason = ( + ", ".join(matched_windows_paths) + if matched_windows_paths + else "no relevant changes" ) - print(f"windows_source_required={str(required).lower()}") - print(f"windows_source_reason={reason}") + print(f"canary_required={str(canary).lower()}") + print(f"canary_reason={canary_reason}") + print(f"windows_source_required={str(windows_source).lower()}") + print(f"windows_source_reason={windows_source_reason}") if __name__ == "__main__": diff --git a/.github/workflows/bazel.yml b/.github/workflows/bazel.yml index 7a6db869e0..10dd8670de 100644 --- a/.github/workflows/bazel.yml +++ b/.github/workflows/bazel.yml @@ -4,10 +4,7 @@ name: Bazel # https://github.com/cerisier/toolchains_llvm_bootstrapped/blob/main/.github/workflows/ci.yaml on: - pull_request: {} - push: - branches: - - main + workflow_call: workflow_dispatch: concurrency: diff --git a/.github/workflows/blob-size-policy.yml b/.github/workflows/blob-size-policy.yml index ea50e8ee2e..fddc51afe7 100644 --- a/.github/workflows/blob-size-policy.yml +++ b/.github/workflows/blob-size-policy.yml @@ -1,7 +1,7 @@ name: blob-size-policy on: - pull_request: {} + workflow_call: jobs: check: @@ -14,13 +14,24 @@ jobs: fetch-depth: 0 persist-credentials: false - - name: Determine PR comparison range + - name: Determine comparison range id: range shell: bash run: | set -euo pipefail - echo "base=${{ github.event.pull_request.base.sha }}" >> "$GITHUB_OUTPUT" - echo "head=${{ github.event.pull_request.head.sha }}" >> "$GITHUB_OUTPUT" + + # PRs inspect the proposed diff; main pushes inspect only the commit + # range that just landed. Both paths feed the same blob-size checker. + if [[ "${{ github.event_name }}" == "pull_request" ]]; then + base='${{ github.event.pull_request.base.sha }}' + head='${{ github.event.pull_request.head.sha }}' + else + base='${{ github.event.before }}' + head='${{ github.sha }}' + fi + + echo "base=$base" >> "$GITHUB_OUTPUT" + echo "head=$head" >> "$GITHUB_OUTPUT" - name: Check changed blob sizes env: diff --git a/.github/workflows/blocking-ci.yml b/.github/workflows/blocking-ci.yml new file mode 100644 index 0000000000..fe66fdd509 --- /dev/null +++ b/.github/workflows/blocking-ci.yml @@ -0,0 +1,79 @@ +name: blocking-ci + +# This is the single entrypoint for checks that block a PR merge. It also runs +# after pushes to main so the same check family stays grouped in the Actions UI. +on: + pull_request: {} + push: + branches: [main] + +jobs: + # Keep reusable workflow calls alphabetized. The `required` job below is the + # version-controlled list that the main-branch ruleset should require. + bazel: + name: Bazel + uses: ./.github/workflows/bazel.yml + secrets: inherit + + blob-size-policy: + name: Blob size policy + uses: ./.github/workflows/blob-size-policy.yml + secrets: inherit + + cargo-deny: + name: cargo-deny + uses: ./.github/workflows/cargo-deny.yml + secrets: inherit + + codespell: + name: Codespell + uses: ./.github/workflows/codespell.yml + secrets: inherit + + repo-checks: + name: repo-checks + uses: ./.github/workflows/repo-checks.yml + secrets: inherit + + rust-ci: + name: rust-ci + uses: ./.github/workflows/rust-ci.yml + secrets: inherit + + sdk: + name: sdk + uses: ./.github/workflows/sdk.yml + secrets: inherit + + required: + name: CI required + # Without `always()`, GitHub skips this job after a failed dependency and a + # required check can appear successful instead of reporting the failure. + if: ${{ always() }} + needs: + - bazel + - blob-size-policy + - cargo-deny + - codespell + - repo-checks + - rust-ci + - sdk + runs-on: ubuntu-24.04 + steps: + # Keep the helper on the same revision as the caller and child workflows. + # CI workflow uploads are restricted, so this repository does not need a + # separate trusted-base checkout for the terminal policy step. Using the + # PR head also lets the introducing PR exercise a newly added helper. + # + # During the initial rollout, PR branches created before + # check_ci_results.py exists must rebase onto main before this gate can + # run. + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + persist-credentials: false + + - name: Require successful dependencies + env: + NEEDS: ${{ toJSON(needs) }} + run: python3 .github/scripts/check_ci_results.py diff --git a/.github/workflows/cargo-deny.yml b/.github/workflows/cargo-deny.yml index aa37d7cec2..be00e1147d 100644 --- a/.github/workflows/cargo-deny.yml +++ b/.github/workflows/cargo-deny.yml @@ -1,10 +1,7 @@ name: cargo-deny on: - pull_request: - push: - branches: - - main + workflow_call: # Cargo's libgit2 transport has been flaky when fetching git dependencies with # nested submodules. Prefer the system git CLI across every Cargo invocation. diff --git a/.github/workflows/codespell.yml b/.github/workflows/codespell.yml index b642a6dbe1..a1751c8d29 100644 --- a/.github/workflows/codespell.yml +++ b/.github/workflows/codespell.yml @@ -3,10 +3,7 @@ name: Codespell on: - push: - branches: [main] - pull_request: - branches: [main] + workflow_call: permissions: contents: read diff --git a/.github/workflows/postmerge-ci.yml b/.github/workflows/postmerge-ci.yml new file mode 100644 index 0000000000..ac2b46d1db --- /dev/null +++ b/.github/workflows/postmerge-ci.yml @@ -0,0 +1,40 @@ +name: postmerge-ci + +# This is the single entrypoint for main-only CI that is intentionally outside +# the merge-blocking suite. It keeps the broader postmerge signal in one run. +on: + push: + branches: [main] + +jobs: + # Keep reusable workflow calls alphabetized. Each child retains its own + # workflow_dispatch trigger so maintainers can rerun flaky suites directly. + rust-ci-full: + name: rust-ci-full + uses: ./.github/workflows/rust-ci-full.yml + secrets: inherit + + v8-canary: + name: v8-canary + uses: ./.github/workflows/v8-canary.yml + secrets: inherit + + results: + name: Postmerge CI results + needs: + - rust-ci-full + - v8-canary + if: ${{ always() }} + runs-on: ubuntu-24.04 + steps: + # Postmerge runs use the pushed main commit, so this helper always comes + # from the same revision that defined the parent workflow. + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + ref: ${{ github.sha }} + persist-credentials: false + + - name: Require successful dependencies + env: + NEEDS: ${{ toJSON(needs) }} + run: python3 .github/scripts/check_ci_results.py diff --git a/.github/workflows/ci.yml b/.github/workflows/repo-checks.yml similarity index 98% rename from .github/workflows/ci.yml rename to .github/workflows/repo-checks.yml index a78ae034ed..3a76a202ea 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/repo-checks.yml @@ -1,8 +1,7 @@ -name: ci +name: repo-checks on: - pull_request: {} - push: { branches: [main] } + workflow_call: jobs: build-test: diff --git a/.github/workflows/rust-ci-full.yml b/.github/workflows/rust-ci-full.yml index 7ad1d4b3ad..cb99c8911e 100644 --- a/.github/workflows/rust-ci-full.yml +++ b/.github/workflows/rust-ci-full.yml @@ -1,8 +1,10 @@ name: rust-ci-full on: + workflow_call: push: branches: - - main + # Main pushes enter through postmerge-ci. Keep this opt-in branch trigger + # for developers who want the full suite before merging. - "**full-ci**" workflow_dispatch: diff --git a/.github/workflows/rust-ci.yml b/.github/workflows/rust-ci.yml index f693994d1a..cf1afdfe2d 100644 --- a/.github/workflows/rust-ci.yml +++ b/.github/workflows/rust-ci.yml @@ -1,6 +1,6 @@ name: rust-ci on: - pull_request: {} + workflow_call: workflow_dispatch: # Cargo's libgit2 transport has been flaky when fetching git dependencies with diff --git a/.github/workflows/sdk.yml b/.github/workflows/sdk.yml index dcb9406798..41f8874712 100644 --- a/.github/workflows/sdk.yml +++ b/.github/workflows/sdk.yml @@ -1,9 +1,7 @@ name: sdk on: - push: - branches: [main] - pull_request: {} + workflow_call: jobs: python-sdk: diff --git a/.github/workflows/v8-canary.yml b/.github/workflows/v8-canary.yml index efa06aaad7..53d9585108 100644 --- a/.github/workflows/v8-canary.yml +++ b/.github/workflows/v8-canary.yml @@ -1,44 +1,12 @@ name: v8-canary +# Do not use trigger-level path filters here. This workflow is also called by +# postmerge-ci, and GitHub cannot share a path filter between pull_request and +# workflow_call. v8_canary_changes.py is the single source of truth instead; +# unrelated events run only the cheap metadata job below. on: - pull_request: - paths: - - ".bazelrc" - - ".github/actions/setup-bazel-ci/**" - - ".github/scripts/run_bazel_with_buildbuddy.py" - - ".github/scripts/rusty_v8_bazel.py" - - ".github/scripts/rusty_v8_module_bazel.py" - - ".github/scripts/v8_canary_changes.py" - - ".github/workflows/rusty-v8-release.yml" - - ".github/workflows/v8-canary.yml" - - "MODULE.bazel" - - "MODULE.bazel.lock" - - "codex-rs/Cargo.toml" - - "patches/BUILD.bazel" - - "patches/llvm_*.patch" - - "patches/rules_cc_*.patch" - - "patches/v8_*.patch" - - "third_party/v8/**" - push: - branches: - - main - paths: - - ".bazelrc" - - ".github/actions/setup-bazel-ci/**" - - ".github/scripts/run_bazel_with_buildbuddy.py" - - ".github/scripts/rusty_v8_bazel.py" - - ".github/scripts/rusty_v8_module_bazel.py" - - ".github/scripts/v8_canary_changes.py" - - ".github/workflows/rusty-v8-release.yml" - - ".github/workflows/v8-canary.yml" - - "MODULE.bazel" - - "MODULE.bazel.lock" - - "codex-rs/Cargo.toml" - - "patches/BUILD.bazel" - - "patches/llvm_*.patch" - - "patches/rules_cc_*.patch" - - "patches/v8_*.patch" - - "third_party/v8/**" + workflow_call: + pull_request: {} workflow_dispatch: # Cargo's libgit2 transport has been flaky when fetching git dependencies with @@ -54,6 +22,10 @@ jobs: metadata: runs-on: ubuntu-latest outputs: + # A stale PR head can contain the old detector, which does not emit this + # output. Missing must mean "run" so older branches cannot silently skip + # the expensive V8 coverage while reporting success. + canary_required: ${{ steps.changes.outputs.canary_required || 'true' }} v8_version: ${{ steps.v8_version.outputs.version }} windows_source_required: ${{ steps.changes.outputs.windows_source_required }} @@ -77,7 +49,7 @@ jobs: version="$(python3 .github/scripts/rusty_v8_bazel.py resolved-v8-crate-version)" echo "version=${version}" >> "$GITHUB_OUTPUT" - - name: Detect whether Windows source artifacts need rebuilding + - name: Detect V8 canary changes id: changes env: BASE_SHA: ${{ github.event_name == 'pull_request' && github.event.pull_request.base.sha || github.event.before }} @@ -87,6 +59,8 @@ jobs: run: | set -euo pipefail + # Manual runs have no meaningful before/after range. Force every V8 + # variant so workflow_dispatch remains a reliable retry path. if [[ "${EVENT_NAME}" == "workflow_dispatch" ]]; then output="$(python3 .github/scripts/v8_canary_changes.py --force)" else @@ -104,6 +78,8 @@ jobs: build: name: Build ${{ matrix.variant }} ${{ matrix.target }} needs: metadata + # Metadata always runs; only relevant changes pay for the large matrix. + if: ${{ needs.metadata.outputs.canary_required == 'true' }} runs-on: ${{ matrix.runner }} permissions: contents: read From 9f7452a6a102320f4155b7495a783283a2204c34 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Fri, 26 Jun 2026 17:32:01 -0700 Subject: [PATCH 38/38] app-server: test JSON tracing output --- codex-rs/Cargo.lock | 1 + codex-rs/app-server/src/lib.rs | 21 ++++--- .../app-server/tests/common/json_logging.rs | 55 +++++++++++++++++++ codex-rs/app-server/tests/common/lib.rs | 2 + codex-rs/app-server/tests/suite/logging.rs | 27 +++++++++ codex-rs/app-server/tests/suite/mod.rs | 1 + codex-rs/cli/Cargo.toml | 1 + codex-rs/cli/tests/app_server.rs | 25 +++++++++ 8 files changed, 125 insertions(+), 8 deletions(-) create mode 100644 codex-rs/app-server/tests/common/json_logging.rs create mode 100644 codex-rs/app-server/tests/suite/logging.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 9f3e69e3cc..78d3c310cd 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2314,6 +2314,7 @@ name = "codex-cli" version = "0.0.0" dependencies = [ "anyhow", + "app_test_support", "assert_cmd", "assert_matches", "clap", diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index beed35848b..0edb57d8a2 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -912,7 +912,7 @@ pub async fn run_main_with_transport_options( async move { let mut listen_for_threads = true; let mut shutdown_state = ShutdownState::default(); - loop { + let exit_reason = loop { let running_turn_count = { let running_turn_count = running_turn_count_rx.borrow(); *running_turn_count @@ -925,7 +925,7 @@ pub async fn run_main_with_transport_options( let _ = outbound_control_tx .send(OutboundControlEvent::DisconnectAll) .await; - break; + break "shutdown_requested"; } tokio::select! { @@ -947,7 +947,7 @@ pub async fn run_main_with_transport_options( } event = transport_event_rx.recv() => { let Some(event) = event else { - break; + break "transport_channel_closed"; }; match event { TransportEvent::ConnectionOpened { @@ -977,7 +977,7 @@ pub async fn run_main_with_transport_options( .await .is_err() { - break; + break "outbound_router_closed"; } connections.insert( connection_id, @@ -1005,10 +1005,10 @@ pub async fn run_main_with_transport_options( .await; }); if !outbound_closed { - break; + break "outbound_router_closed"; } if shutdown_when_no_connections && connections.is_empty() { - break; + break "last_connection_closed"; } } TransportEvent::IncomingMessage { connection_id, message } => { @@ -1147,7 +1147,7 @@ pub async fn run_main_with_transport_options( } } } - } + }; if !shutdown_state.forced() { futures::future::join_all( @@ -1162,7 +1162,12 @@ pub async fn run_main_with_transport_options( } else { connection_cleanup_tasks.abort(); } - info!("processor task exited (channel closed)"); + info!( + exit_reason, + remaining_connection_count = connections.len(), + shutdown_forced = shutdown_state.forced(), + "processor task exited" + ); } }); diff --git a/codex-rs/app-server/tests/common/json_logging.rs b/codex-rs/app-server/tests/common/json_logging.rs new file mode 100644 index 0000000000..c6c78bd0a4 --- /dev/null +++ b/codex-rs/app-server/tests/common/json_logging.rs @@ -0,0 +1,55 @@ +use std::path::Path; +use std::process::Command; +use std::process::Stdio; + +use anyhow::Context; +use anyhow::Result; +use serde_json::Value; +use serde_json::json; + +pub fn app_server_json_shutdown_event( + binary: &str, + args: &[&str], + codex_home: &Path, +) -> Result { + std::fs::write( + codex_home.join("config.toml"), + "[features]\nplugins = false\n", + )?; + let output = Command::new(codex_utils_cargo_bin::cargo_bin(binary)?) + .stdin(Stdio::null()) + .env("CODEX_HOME", codex_home) + .env( + "CODEX_APP_SERVER_MANAGED_CONFIG_PATH", + codex_home.join("managed_config.toml"), + ) + .env("LOG_FORMAT", "json") + .env("RUST_LOG", "codex_app_server=info") + .args(args) + .output()?; + + let stderr = String::from_utf8(output.stderr)?; + anyhow::ensure!(output.status.success(), "app-server failed: {stderr}"); + + let events = stderr + .lines() + .filter(|line| !line.is_empty()) + .map(serde_json::from_str::) + .collect::>>() + .with_context(|| format!("app-server stderr was not JSONL: {stderr}"))?; + let event = events + .iter() + .find(|event| event["fields"]["message"] == "processor task exited") + .context("missing INFO shutdown event in app-server JSON logs")?; + let timestamp = event["timestamp"] + .as_str() + .context("shutdown event did not include a timestamp")?; + chrono::DateTime::parse_from_rfc3339(timestamp) + .with_context(|| format!("shutdown event timestamp was not RFC 3339: {timestamp}"))?; + + Ok(json!({ + "level": event["level"], + "fields": event["fields"], + "target": event["target"], + })) +} diff --git a/codex-rs/app-server/tests/common/lib.rs b/codex-rs/app-server/tests/common/lib.rs index 76c4c85f26..fbc674806a 100644 --- a/codex-rs/app-server/tests/common/lib.rs +++ b/codex-rs/app-server/tests/common/lib.rs @@ -3,6 +3,7 @@ mod analytics_server; mod auth_fixtures; mod config; +mod json_logging; mod mock_model_server; mod models_cache; mod responses; @@ -26,6 +27,7 @@ pub use core_test_support::test_absolute_path; pub use core_test_support::test_path_buf_with_windows; pub use core_test_support::test_tmp_path; pub use core_test_support::test_tmp_path_buf; +pub use json_logging::app_server_json_shutdown_event; pub use mock_model_server::create_mock_responses_server_repeating_assistant; pub use mock_model_server::create_mock_responses_server_sequence; pub use mock_model_server::create_mock_responses_server_sequence_unchecked; diff --git a/codex-rs/app-server/tests/suite/logging.rs b/codex-rs/app-server/tests/suite/logging.rs new file mode 100644 index 0000000000..ea2c31a766 --- /dev/null +++ b/codex-rs/app-server/tests/suite/logging.rs @@ -0,0 +1,27 @@ +use anyhow::Result; +use app_test_support::app_server_json_shutdown_event; +use pretty_assertions::assert_eq; +use serde_json::json; +use tempfile::TempDir; + +#[test] +fn standalone_app_server_emits_json_info_events() -> Result<()> { + let codex_home = TempDir::new()?; + let event = app_server_json_shutdown_event("codex-app-server", &[], codex_home.path())?; + + assert_eq!( + event, + json!({ + "level": "INFO", + "fields": { + "message": "processor task exited", + "exit_reason": "last_connection_closed", + "remaining_connection_count": 0, + "shutdown_forced": false, + }, + "target": "codex_app_server", + }) + ); + + Ok(()) +} diff --git a/codex-rs/app-server/tests/suite/mod.rs b/codex-rs/app-server/tests/suite/mod.rs index bac309a930..09c40493f6 100644 --- a/codex-rs/app-server/tests/suite/mod.rs +++ b/codex-rs/app-server/tests/suite/mod.rs @@ -1,5 +1,6 @@ mod auth; mod conversation_summary; mod fuzzy_file_search; +mod logging; mod strict_config; mod v2; diff --git a/codex-rs/cli/Cargo.toml b/codex-rs/cli/Cargo.toml index 7be1146df5..9a8312ced2 100644 --- a/codex-rs/cli/Cargo.toml +++ b/codex-rs/cli/Cargo.toml @@ -95,6 +95,7 @@ windows-sys = { version = "0.52", features = [ ] } [dev-dependencies] +app_test_support = { workspace = true } assert_cmd = { workspace = true } assert_matches = { workspace = true } codex-utils-cargo-bin = { workspace = true } diff --git a/codex-rs/cli/tests/app_server.rs b/codex-rs/cli/tests/app_server.rs index 4a2642379d..6815c2be99 100644 --- a/codex-rs/cli/tests/app_server.rs +++ b/codex-rs/cli/tests/app_server.rs @@ -1,7 +1,10 @@ use std::path::Path; use anyhow::Result; +use app_test_support::app_server_json_shutdown_event; use predicates::str::contains; +use pretty_assertions::assert_eq; +use serde_json::json; use tempfile::TempDir; fn codex_command(codex_home: &Path) -> Result { @@ -28,3 +31,25 @@ foo = "bar" Ok(()) } + +#[test] +fn app_server_emits_json_info_events() -> Result<()> { + let codex_home = TempDir::new()?; + let event = app_server_json_shutdown_event("codex", &["app-server"], codex_home.path())?; + + assert_eq!( + event, + json!({ + "level": "INFO", + "fields": { + "message": "processor task exited", + "exit_reason": "last_connection_closed", + "remaining_connection_count": 0, + "shutdown_forced": false, + }, + "target": "codex_app_server", + }) + ); + + Ok(()) +}