From f3074e35c079bd8a5e89b0301b5e735d0e618235 Mon Sep 17 00:00:00 2001 From: rhan-oai Date: Thu, 5 Nov 2026 00:24:33 +0000 Subject: [PATCH] Render host-owned V2 subagent rosters in environment context (#59531) ## What changed - Expose `AgentSessionReader` and `AgentRef` for read-only access to host-owned agent session membership. - Read direct children when rendering V2 environment context, using the existing roster size limits. Refresh membership when rebuilding context after compaction, and omit the roster if the provider is absent or the read fails. - Preserve unrecorded initial input across suspension or interruption while waiting for a roster read. Resolve the roster before advancing the compaction window. ## Testing Add coverage for roster refreshes after compaction, unchanged environments avoiding extra reads, unavailable rosters, recoverable input during blocked reads, and stale compaction results rejected after a concurrent history replacement. GitOrigin-RevId: 3c7e3822af09e9295d405e53dabb87b020c8f00c --- codex-rs/core/src/agent/host/mod.rs | 3 + codex-rs/core/src/agent/host/ownership.rs | 10 +- .../core/src/agent/host/session_reader.rs | 32 +++++ codex-rs/core/src/compact.rs | 15 +- .../compact_remote_v2_installation_tests.rs | 129 ++++++++++++------ .../core/src/context/token_budget_context.rs | 11 +- .../world_state/compact_permissions.rs | 18 ++- .../src/context/world_state/environment.rs | 39 ++++-- codex-rs/core/src/context/world_state/mod.rs | 68 +++++++-- .../core/src/context/world_state/model.rs | 4 +- .../context/world_state/multi_agent_mode.rs | 2 +- .../world_state/multi_agent_usage_hint.rs | 2 +- .../core/src/context/world_state/realtime.rs | 4 +- .../src/context/world_state/test_support.rs | 12 +- .../context/world_state/world_state_tests.rs | 8 +- codex-rs/core/src/context_manager/history.rs | 11 -- .../core/src/context_manager/history_tests.rs | 12 ++ codex-rs/core/src/context_manager/updates.rs | 7 +- .../core/src/context_manager/updates_tests.rs | 1 + codex-rs/core/src/lib.rs | 2 + codex-rs/core/src/session/mod.rs | 55 +++++--- codex-rs/core/src/session/turn.rs | 56 +++++++- codex-rs/core/src/session/turn_suspension.rs | 7 +- codex-rs/core/src/session/world_state.rs | 53 ++++--- codex-rs/core/src/state/turn.rs | 4 +- codex-rs/core/src/tasks/admission.rs | 4 +- codex-rs/core/src/tasks/mod.rs | 13 +- .../core/tests/common/context_snapshot.rs | 7 + .../common/context_snapshot/normalize.rs | 1 + .../core/tests/suite/agent_control_tests.rs | 41 ++++++ .../core/tests/suite/host_turn_admission.rs | 92 ++++++++++++- .../tests/suite/scenarios_budget_context.rs | 91 ++++++++++-- ..._host_budget_context_after_compaction.snap | 38 +++++- 33 files changed, 642 insertions(+), 210 deletions(-) create mode 100644 codex-rs/core/src/agent/host/session_reader.rs diff --git a/codex-rs/core/src/agent/host/mod.rs b/codex-rs/core/src/agent/host/mod.rs index aba89129e3..5c01e05660 100644 --- a/codex-rs/core/src/agent/host/mod.rs +++ b/codex-rs/core/src/agent/host/mod.rs @@ -2,6 +2,7 @@ mod admission; mod invocation; mod mcp_metadata; mod ownership; +mod session_reader; mod spawn; mod types; mod usage; @@ -17,6 +18,8 @@ pub(crate) use invocation::tool_invocation; pub use mcp_metadata::AgentMcpMetadataRecipient; pub use ownership::AgentRuntimeOwnership; pub(crate) use ownership::AgentRuntimeServices; +pub use session_reader::AgentRef; +pub use session_reader::AgentSessionReader; pub use types::HostThreadSpawnReservation; pub use types::InheritedSessionId; pub use types::PreparedChildContext; diff --git a/codex-rs/core/src/agent/host/ownership.rs b/codex-rs/core/src/agent/host/ownership.rs index 33005741ce..37125d04cb 100644 --- a/codex-rs/core/src/agent/host/ownership.rs +++ b/codex-rs/core/src/agent/host/ownership.rs @@ -1,5 +1,6 @@ //! Captures coordination providers once, after the thread's startup contributors finish. +use super::AgentSessionReader; use super::AgentTurnAdmission; use super::AgentUsageAccounting; use codex_extension_api::ExtensionData; @@ -14,8 +15,8 @@ pub enum AgentRuntimeOwnership { #[default] Controller, /// Requires an [`AgentTurnAdmission`] provider. Core skips controller rollout-budget - /// accounting, budget reminders, and subagent roster context for this runtime. - /// Install an [`AgentUsageAccounting`] provider to receive usage reports. + /// accounting and budget reminders. Install an [`AgentUsageAccounting`] provider + /// to receive usage reports and an [`AgentSessionReader`] to render the V2 subagent roster. Host, } @@ -25,6 +26,7 @@ pub(crate) enum AgentRuntimeServices { } pub(crate) struct HostAgentServices { + pub(crate) session_reader: Option>, pub(crate) turn_admission: Arc, pub(crate) usage_accounting: Option>, } @@ -40,6 +42,10 @@ impl AgentRuntimeServices { ) })?; Ok(Self::Host(HostAgentServices { + session_reader: data + .get::>() + .as_deref() + .cloned(), turn_admission: Arc::clone(turn_admission.as_ref()), usage_accounting: data .get::>() diff --git a/codex-rs/core/src/agent/host/session_reader.rs b/codex-rs/core/src/agent/host/session_reader.rs new file mode 100644 index 0000000000..28d61afd49 --- /dev/null +++ b/codex-rs/core/src/agent/host/session_reader.rs @@ -0,0 +1,32 @@ +//! Read-only host agent session membership for native context construction. + +use codex_protocol::AgentPath; +use codex_protocol::ThreadId; +use codex_protocol::error::Result; +use futures::future::BoxFuture; + +/// An agent's identity in the selected owner's agent session. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct AgentRef { + pub thread_id: ThreadId, + pub path: AgentPath, +} + +/// An agent session view scoped to one runtime, installed before host-owned startup or resume. +/// +/// Reads must not register, load, or restore agents. Return an error when membership +/// is unavailable. Implementations may read a local cache or query their coordinator. +/// Core reads child paths for its bounded V2 roster when rendering an environment fragment. +/// If the read fails, Core renders the fragment without a roster. +pub trait AgentSessionReader: Send + Sync { + /// Returns the root, including self when this runtime is the root. + fn root(&self) -> BoxFuture<'_, Result>; + + /// Returns the structural parent; only the root has no parent. + fn parent(&self) -> BoxFuture<'_, Result>>; + + /// Returns registered direct children, including unloaded members, excluding + /// provisional reservations and historical/removed relationships. Return loaded + /// children first, then alphabetical paths within each group, as native rosters do. + fn children(&self) -> BoxFuture<'_, Result>>; +} diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index 095b9b68e4..c75de9167c 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -444,15 +444,14 @@ async fn run_compact_task_inner_impl( new_history = crate::compact_pending_input::insert_summary(new_history, summary, pending_sequence); } + let rendered = world_state.render_full_with_host_roster().await; let (window_number, window_ids) = sess.advance_auto_compact_window().await; - - let (new_history, world_state_baseline) = build_compaction_replacement_history( - sess.as_ref(), - &replacement_step_context, - &world_state, - new_history, - ) - .await; + let (updates, world_state_baseline) = sess + .assemble_initial_context(&replacement_step_context, rendered) + .await; + let (prefix, context) = split_prefix_updates(updates); + let new_history = + assemble_compaction_history(new_history, prefix, merge_world_state_updates(context)); sess.replace_compacted_history( new_history, replacement_step_context.to_turn_context_item(), diff --git a/codex-rs/core/src/compact_remote_v2_installation_tests.rs b/codex-rs/core/src/compact_remote_v2_installation_tests.rs index f24d9791c8..48878fa955 100644 --- a/codex-rs/core/src/compact_remote_v2_installation_tests.rs +++ b/codex-rs/core/src/compact_remote_v2_installation_tests.rs @@ -1,14 +1,19 @@ //! Stale prefixes are skipped or replaced with a fresh request according to caller intent. use super::*; +use crate::AgentRef; +use crate::AgentSessionReader; use crate::compact_remote_v2::RemoteCompactV2Attempt; +use crate::context::world_state::EnvironmentsState; use crate::context::world_state::WorldState; use crate::session::CompactionPrefix; use codex_login::CodexAuth; use codex_protocol::models::ResponseItem; use core_test_support::responses; +use futures::future::BoxFuture; use pretty_assertions::assert_eq; use test_case::test_case; +use tokio::sync::Notify; use tokio_util::sync::CancellationToken; #[derive(Clone, Copy)] @@ -17,6 +22,30 @@ enum HistoryChange { Account, } +#[derive(Default)] +struct FailingRosterRead { + started: Notify, + release: Notify, +} + +impl AgentSessionReader for FailingRosterRead { + fn root(&self) -> BoxFuture<'_, CodexResult> { + unreachable!("roster rendering reads children") + } + + fn parent(&self) -> BoxFuture<'_, CodexResult>> { + unreachable!("roster rendering reads children") + } + + fn children(&self) -> BoxFuture<'_, CodexResult>> { + Box::pin(async { + self.started.notify_one(); + self.release.notified().await; + Err(CodexErr::InvalidRequest("roster unavailable".into())) + }) + } +} + #[test_case(HistoryChange::Replacement, CompactionFallback::Skip; "replacement optional")] #[test_case(HistoryChange::Account, CompactionFallback::Skip; "account optional")] #[test_case(HistoryChange::Replacement, CompactionFallback::Foreground; "replacement required")] @@ -68,14 +97,31 @@ async fn prepared_prefix_rejection_respects_caller_intent( .await; let history = sess.clone_history().await; let auth = sess.services.model_client.auth_owner_generation(); - let mut candidate = BackgroundCompactionCandidate { - prefix: CompactionPrefix { - generation: history.history_version(), - len: history.annotated_items().len(), - auth_owner_generation: auth, + let prefix = CompactionPrefix { + generation: history.history_version(), + len: history.annotated_items().len(), + auth_owner_generation: match change { + HistoryChange::Replacement => auth, + HistoryChange::Account => Some(auth.unwrap_or_default().wrapping_add(1)), }, + }; + let roster = Arc::new(FailingRosterRead::default()); + let mut cutoff_world_state = WorldState::default(); + if matches!(change, HistoryChange::Replacement) { + cutoff_world_state.add_section( + EnvironmentsState::from_turn_context_with_environments( + &step.turn, + &step.environments, + /*current_date*/ Some("2026-01-01".into()), + ) + .await, + ); + cutoff_world_state.host_roster_reader = Some(roster.clone()); + } + let candidate = BackgroundCompactionCandidate { + prefix, model_info: Arc::clone(&step.settings.model_info), - cutoff_world_state: Arc::new(WorldState::default()), + cutoff_world_state: Arc::new(cutoff_world_state), attempt: RemoteCompactV2Attempt { rejected_audio_file_ids: Default::default(), input_goal_ids: Default::default(), @@ -97,37 +143,14 @@ async fn prepared_prefix_rejection_respects_caller_intent( owned_client_session: None, }, }; - match change { - HistoryChange::Replacement => { - sess.replace_history( - vec![original.clone(), original.clone()], - /*reference_context_item*/ None, - ) - .await; - } - HistoryChange::Account => { - candidate.prefix.auth_owner_generation = Some(auth.unwrap_or_default().wrapping_add(1)); - } - } - // Rewritten histories still fit the captured length; length alone cannot reject them. - let history = sess.clone_history().await; - assert!(history.annotated_items().len() >= candidate.prefix.len); - if !matches!(change, HistoryChange::Account) { - assert_ne!(history.history_version(), candidate.prefix.generation); - } - let before = sess - .clone_history() - .await - .raw_items() - .cloned() - .collect::>(); + let world_state = WorldState::default(); let mut analytics = CompactionAnalyticsDetails::default(); - let installed = run_compaction( + let mut compaction = Box::pin(run_compaction( &sess, &step, &step, /*client_session*/ None, - &WorldState::default(), + &world_state, CompactionTurnMetadata::new( CompactionTrigger::Auto, CompactionReason::ContextLimit, @@ -138,9 +161,29 @@ async fn prepared_prefix_rejection_respects_caller_intent( /*started_item*/ None, Some(candidate), fallback, - ) - .await - .expect("reject stale result without failing compaction"); + )); + if matches!(change, HistoryChange::Replacement) { + tokio::select! { + _ = roster.started.notified() => {} + result = &mut compaction => panic!("expected a pending roster read, got {result:?}"), + } + sess.replace_history( + vec![original.clone(), original], + /*reference_context_item*/ None, + ) + .await; + } + // The replacement still fits the captured length, but changes the generation. + let before = sess.clone_history().await; + let window_before = sess.current_window().await; + assert!(before.annotated_items().len() >= prefix.len); + if matches!(change, HistoryChange::Replacement) { + assert_ne!(before.history_version(), prefix.generation); + roster.release.notify_one(); + } + let installed = compaction + .await + .expect("reject stale result without failing compaction"); // The foreground response has no usage; stale candidate statistics must not leak. assert_eq!( ( @@ -150,24 +193,24 @@ async fn prepared_prefix_rejection_respects_caller_intent( ), (None, None, None), ); - let after = sess - .clone_history() - .await - .raw_items() - .cloned() - .collect::>(); + let after = sess.clone_history().await; match fallback { CompactionFallback::Skip => { assert!(!installed); assert_eq!(analytics.retained_image_count, None); assert_eq!(requests.requests().len(), 0); - assert_eq!(after, before); + assert_eq!(after.annotated_items(), before.annotated_items()); + assert_eq!(sess.current_window().await, window_before); + assert_eq!( + after.world_state_checkpoint(), + before.world_state_checkpoint() + ); } CompactionFallback::Foreground => { assert!(installed); assert_eq!(requests.requests().len(), 1); let summaries = after - .iter() + .raw_items() .filter_map(|item| match item { ResponseItem::Compaction { encrypted_content, .. diff --git a/codex-rs/core/src/context/token_budget_context.rs b/codex-rs/core/src/context/token_budget_context.rs index 7aa415526d..fea07d3385 100644 --- a/codex-rs/core/src/context/token_budget_context.rs +++ b/codex-rs/core/src/context/token_budget_context.rs @@ -97,15 +97,8 @@ impl WorldStateSection for TokenBudgetContext { ) -> SectionTransition { let current = self.agent_path.clone(); let fragment = matches!(previous, PreviousSectionState::Known(agent_path) if agent_path != &self.agent_path) - .then(|| Box::new(self.clone()) as Box); - ( - Some(current), - fragment - .into_iter() - .map(WorldStateUpdate::boxed_fragment) - .map(WorldStateUpdate::standalone) - .collect(), - ) + .then(|| WorldStateUpdate::fragment(self.clone()).standalone()); + (Some(current), fragment.into_iter().collect()) } } diff --git a/codex-rs/core/src/context/world_state/compact_permissions.rs b/codex-rs/core/src/context/world_state/compact_permissions.rs index c1afa2d4d6..ade6795de7 100644 --- a/codex-rs/core/src/context/world_state/compact_permissions.rs +++ b/codex-rs/core/src/context/world_state/compact_permissions.rs @@ -4,7 +4,6 @@ use super::WorldStateSection; use super::WorldStateUpdate; use crate::context::APPROVED_COMMAND_PREFIX_SAVED_MESSAGE_PREFIX; use crate::context::ApprovedCommandPrefixSaved; -use crate::context::ContextualUserFragment; use codex_execpolicy::Policy; use codex_protocol::models::format_allow_prefixes; use std::collections::BTreeSet; @@ -51,15 +50,14 @@ impl WorldStateSection for CompactPermissionsState { return (Some(current), Vec::new()); } }; - let fragment: Option> = - format_allow_prefixes(added_prefixes) - .filter(|prefixes| !prefixes.is_empty()) - .map(|prefixes| { - Box::new(ApprovedCommandPrefixSaved::new( - prefixes, - self.saved_prefix_heading.as_deref(), - )) as _ - }); + let fragment: Option = format_allow_prefixes(added_prefixes) + .filter(|prefixes| !prefixes.is_empty()) + .map(|prefixes| { + Box::new(ApprovedCommandPrefixSaved::new( + prefixes, + self.saved_prefix_heading.as_deref(), + )) as _ + }); ( Some(current), WorldStateUpdate::optional_boxed_fragment(fragment), diff --git a/codex-rs/core/src/context/world_state/environment.rs b/codex-rs/core/src/context/world_state/environment.rs index ba661c1ee4..f9f8c49883 100644 --- a/codex-rs/core/src/context/world_state/environment.rs +++ b/codex-rs/core/src/context/world_state/environment.rs @@ -1,7 +1,9 @@ +use super::Placement; use super::PreviousSectionState; use super::SectionTransition; use super::WorldStateSection; use super::WorldStateUpdate; +use super::WorldStateUpdateContent; use crate::context::ContextualUserFragment; use crate::context::environment_context::FileSystemContext; use crate::context::environment_context::NetworkContext; @@ -11,6 +13,7 @@ use crate::environment_selection::TurnEnvironmentState; use crate::session::turn_context::TurnContext; use crate::shell::ShellType; use codex_features::Feature; +use codex_protocol::AgentPath; use codex_protocol::models::ContentItem; use codex_protocol::models::ContentItemKind; use codex_protocol::models::input_text; @@ -179,7 +182,7 @@ impl WorldStateSection for EnvironmentsState { .values() .all(|update| matches!(update, EnvironmentUpdate::Current(_))); let fragment = (!updates.is_empty() || turn_context_values_changed).then(|| { - Box::new(RenderedEnvironments { + let environment = Box::new(RenderedEnvironments { updates, legacy_single, include_primary: multiple_environments || previous_multiple_environments, @@ -193,12 +196,13 @@ impl WorldStateSection for EnvironmentsState { network: self.network.clone(), filesystem: self.filesystem.clone(), subagents: self.subagents.clone(), - }) as Box + }); + WorldStateUpdate { + placement: Placement::Mergeable, + content: WorldStateUpdateContent::Environment(environment), + } }); - ( - Some(current), - WorldStateUpdate::optional_boxed_fragment(fragment), - ) + (Some(current), fragment.into_iter().collect()) } } @@ -224,7 +228,7 @@ impl ContextualUserFragment for EnvironmentsState { } } -struct RenderedEnvironments { +pub(crate) struct RenderedEnvironments { updates: BTreeMap, legacy_single: bool, include_primary: bool, @@ -235,7 +239,7 @@ struct RenderedEnvironments { timezone: Option, network: Option, filesystem: Option, - subagents: Option, + pub(crate) subagents: Option, } enum EnvironmentUpdate { @@ -565,3 +569,22 @@ mod tests; #[cfg(test)] #[path = "environment_render_tests.rs"] mod render_tests; + +pub(crate) fn format_agent_paths(paths: impl Iterator) -> String { + const MAX_AGENTS: usize = 8; + const MAX_BYTES: usize = 1_024; + let mut lines = Vec::new(); + let mut rendered_bytes = " \n \n".len(); + for path in paths { + if lines.len() == MAX_AGENTS { + break; + } + let line = format!(r#""#); + let line_bytes = " \n".len() + line.len(); + if rendered_bytes + line_bytes <= MAX_BYTES { + rendered_bytes += line_bytes; + lines.push(line); + } + } + lines.join("\n") +} diff --git a/codex-rs/core/src/context/world_state/mod.rs b/codex-rs/core/src/context/world_state/mod.rs index 5037cc32cb..c83cce0979 100644 --- a/codex-rs/core/src/context/world_state/mod.rs +++ b/codex-rs/core/src/context/world_state/mod.rs @@ -19,6 +19,7 @@ pub(crate) mod test_support; mod tools; mod top_level_tools; +use crate::AgentSessionReader; use crate::context::ContextualUserFragment; use codex_extension_api::PreviousWorldStateSection; use codex_extension_api::RenderedWorldStateFragment; @@ -36,6 +37,7 @@ use sha1::Digest; use sha1::Sha1; use std::collections::BTreeMap; use std::fmt; +use std::sync::Arc; pub(crate) use agents_md::AgentsMdState; pub(crate) use apps_instructions::AppsInstructionsState; @@ -43,6 +45,8 @@ pub(crate) use collaboration_mode::CollaborationModeState; pub(crate) use compact_permissions::CompactPermissionsState; pub(crate) use context_window_guidance::ContextWindowGuidanceState; pub(crate) use environment::EnvironmentsState; +use environment::RenderedEnvironments; +pub(crate) use environment::format_agent_paths; pub(crate) use environments_instructions::EnvironmentsInstructionsState; pub(crate) use managed_developer_instructions::ManagedDeveloperInstructions; pub(crate) use managed_developer_instructions::ManagedDeveloperInstructionsState; @@ -73,31 +77,42 @@ pub(crate) enum Placement { Mergeable, } +pub(crate) type WorldStateFragment = Box; + pub(crate) enum WorldStateUpdateContent { - Fragment(Box), + Fragment(WorldStateFragment), + Environment(Box), Item(Box), } +impl WorldStateUpdateContent { + pub(crate) fn into_fragment(self) -> Result> { + match self { + Self::Fragment(fragment) => Ok(fragment), + Self::Environment(environment) => Ok(environment), + Self::Item(item) => Err(item), + } + } +} + impl WorldStateUpdate { - pub(crate) fn fragment(fragment: impl ContextualUserFragment + 'static) -> Self { + pub(crate) fn fragment(fragment: impl ContextualUserFragment + Send + 'static) -> Self { Self::boxed_fragment(Box::new(fragment)) } - pub(crate) fn boxed_fragment(fragment: Box) -> Self { + fn boxed_fragment(fragment: WorldStateFragment) -> Self { Self { placement: Placement::Mergeable, content: WorldStateUpdateContent::Fragment(fragment), } } - pub(crate) fn optional_boxed_fragment( - fragment: Option>, - ) -> Vec { + pub(crate) fn optional_boxed_fragment(fragment: Option) -> Vec { fragment.into_iter().map(Self::boxed_fragment).collect() } pub(crate) fn optional_standalone_boxed_fragment( - fragment: Option>, + fragment: Option, ) -> Vec { fragment .into_iter() @@ -129,9 +144,9 @@ pub(crate) fn split_prefix_updates( let mut context = Vec::new(); for update in updates { match update.placement { - Placement::Prefix => prefix.push(match update.content { - WorldStateUpdateContent::Item(item) => *item, - WorldStateUpdateContent::Fragment(fragment) => fragment.into_boxed_response_item(), + Placement::Prefix => prefix.push(match update.content.into_fragment() { + Ok(fragment) => fragment.into_boxed_response_item(), + Err(item) => *item, }), Placement::Standalone | Placement::Mergeable => context.push(update), } @@ -408,6 +423,7 @@ fn hash_component(hasher: &mut Sha1, value: &str) { /// Live model-visible state, keyed by the same stable section IDs used in rollouts. #[derive(Default)] pub(crate) struct WorldState { + pub(crate) host_roster_reader: Option>, sections: IndexMap<&'static str, Box>, } @@ -517,6 +533,30 @@ enum RenderScope { } impl WorldState { + /// Read the roster only when rendering an environment fragment. If the read fails, + /// keep rendering the rest of the fragment. + pub(crate) async fn resolve_host_roster(&self, updates: &mut [WorldStateUpdate]) { + if let Some(reader) = &self.host_roster_reader + && let Some(environment) = + updates + .iter_mut() + .find_map(|update| match &mut update.content { + WorldStateUpdateContent::Environment(environment) => Some(environment), + WorldStateUpdateContent::Fragment(_) | WorldStateUpdateContent::Item(_) => { + None + } + }) + { + match reader.children().await { + Ok(children) => { + let roster = format_agent_paths(children.into_iter().map(|agent| agent.path)); + environment.subagents = (!roster.is_empty()).then_some(roster); + } + Err(_) => tracing::warn!("failed to read host agent roster for prompt context"), + } + } + } + pub(crate) fn full_context_excluded_section_ids( &self, ) -> impl Iterator + '_ { @@ -551,6 +591,14 @@ impl WorldState { } /// Renders restorable sections as new, preserving update order and placement. + pub(crate) async fn render_full_with_host_roster( + &self, + ) -> (WorldStateSnapshot, Vec) { + let (snapshot, mut updates) = self.render_full(); + self.resolve_host_roster(&mut updates).await; + (snapshot, updates) + } + pub(crate) fn render_full(&self) -> (WorldStateSnapshot, Vec) { self.render_with(RenderScope::FullContext, /*stored*/ None, |_, _| { PreviousSectionState::Absent diff --git a/codex-rs/core/src/context/world_state/model.rs b/codex-rs/core/src/context/world_state/model.rs index a16fcb90ae..a9fb43f156 100644 --- a/codex-rs/core/src/context/world_state/model.rs +++ b/codex-rs/core/src/context/world_state/model.rs @@ -125,7 +125,7 @@ impl WorldStateSection for ModelInstructionsState { if instructions_hash.is_some() && !announce_model && !instructions_changed { return (model_changed.then(|| self.snapshot.clone()), Vec::new()); } - let fragment: Option> = if instructions_hash.is_some() { + let fragment: Option = if instructions_hash.is_some() { Some(Box::new(ModelInstructions { model: announce_model.then(|| model.clone()), instructions: instructions_changed.then(|| self.instructions.clone()), @@ -133,7 +133,7 @@ impl WorldStateSection for ModelInstructionsState { } else { (model_changed && !self.instructions.is_empty()).then(|| { Box::new(ModelSwitchInstructions::new(self.instructions.clone())) - as Box + as super::WorldStateFragment }) }; ( diff --git a/codex-rs/core/src/context/world_state/multi_agent_mode.rs b/codex-rs/core/src/context/world_state/multi_agent_mode.rs index 3ff875db45..8ae3a15fed 100644 --- a/codex-rs/core/src/context/world_state/multi_agent_mode.rs +++ b/codex-rs/core/src/context/world_state/multi_agent_mode.rs @@ -93,7 +93,7 @@ impl WorldStateSection for MultiAgentModeState { let fragment = MultiAgentModeInstructions::from_mode(mode).map(|instructions| { Box::new(instructions.with_metadata(ContentItemMetadata::tool(self.namespace.clone()))) - as Box + as super::WorldStateFragment }); ( Some(self.clone()), diff --git a/codex-rs/core/src/context/world_state/multi_agent_usage_hint.rs b/codex-rs/core/src/context/world_state/multi_agent_usage_hint.rs index edf40e812e..57d7dadf83 100644 --- a/codex-rs/core/src/context/world_state/multi_agent_usage_hint.rs +++ b/codex-rs/core/src/context/world_state/multi_agent_usage_hint.rs @@ -46,7 +46,7 @@ impl WorldStateSection for MultiAgentUsageHintState { &self, previous: PreviousSectionState<'_, Self::Snapshot>, ) -> SectionTransition { - let fragment: Option> = match previous { + let fragment: Option = match previous { PreviousSectionState::Known(previous) if previous == &self.fingerprint => { return (None, Vec::new()); } diff --git a/codex-rs/core/src/context/world_state/realtime.rs b/codex-rs/core/src/context/world_state/realtime.rs index 52a55768c7..730ed3aef5 100644 --- a/codex-rs/core/src/context/world_state/realtime.rs +++ b/codex-rs/core/src/context/world_state/realtime.rs @@ -35,14 +35,14 @@ impl RealtimeState { } } - fn render_start(&self) -> Box { + fn render_start(&self) -> super::WorldStateFragment { match self.start_instructions.as_deref() { Some(instructions) => Box::new(RealtimeStartWithInstructions::new(instructions)), None => Box::new(RealtimeStartInstructions), } } - fn render_transition(&self, previous_active: bool) -> Option> { + fn render_transition(&self, previous_active: bool) -> Option { match (previous_active, self.snapshot.active) { (false, true) => Some(self.render_start()), (true, false) => Some(match self.end_instructions.as_deref() { diff --git a/codex-rs/core/src/context/world_state/test_support.rs b/codex-rs/core/src/context/world_state/test_support.rs index f627c7f83a..46f55267af 100644 --- a/codex-rs/core/src/context/world_state/test_support.rs +++ b/codex-rs/core/src/context/world_state/test_support.rs @@ -22,8 +22,8 @@ pub(super) fn render_section_cases<'a, S: WorldStateSection>( } else { updates .into_iter() - .map(|update| match update.content { - WorldStateUpdateContent::Fragment(fragment) => { + .map(|update| match update.content.into_fragment() { + Ok(fragment) => { let rendered = match fragment.render() { ContentItem::InputText { text } | ContentItem::OutputText { text, .. } => text, @@ -31,7 +31,7 @@ pub(super) fn render_section_cases<'a, S: WorldStateSection>( }; format!(" (role - {})\n{rendered}", fragment.role()) } - WorldStateUpdateContent::Item(item) => { + Err(item) => { let placement = match update.placement { Placement::Prefix => "prefix", Placement::Standalone => "standalone", @@ -125,9 +125,9 @@ pub(crate) fn expect_fragments( } fn expect_fragment(content: WorldStateUpdateContent) -> Box { - match content { - WorldStateUpdateContent::Fragment(fragment) => fragment, - WorldStateUpdateContent::Item(item) => panic!("expected a context fragment, got {item:?}"), + match content.into_fragment() { + Ok(fragment) => fragment, + Err(item) => panic!("expected a context fragment, got {item:?}"), } } 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 9d31b636c5..20c995d9d2 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 @@ -28,7 +28,7 @@ impl WorldStateSection for TestSection { previous: PreviousSectionState<'_, Self::Snapshot>, ) -> SectionTransition { let current = self.clone(); - let fragment: Option> = match previous { + let fragment: Option = match previous { PreviousSectionState::Known(previous) if self.value != previous.value => { Some(Box::new(TestFragment(self.value.clone()))) } @@ -520,9 +520,9 @@ fn mixed_output_preserves_order_and_restores_its_snapshot() { assert_eq!( context .into_iter() - .map(|update| match update.content { - WorldStateUpdateContent::Fragment(fragment) => fragment.into_boxed_response_item(), - WorldStateUpdateContent::Item(item) => *item, + .map(|update| match update.content.into_fragment() { + Ok(fragment) => fragment.into_boxed_response_item(), + Err(item) => *item, }) .collect::>(), vec![ diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 2673e03a3b..8ead85ad21 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -556,17 +556,6 @@ impl ContextManager { self.reference_context_item.clone() } - pub(crate) fn update_world_state( - &mut self, - world_state: &WorldState, - ) -> (Vec, Option) { - let (snapshot, fragments) = - world_state.render_history_diff(self.world_state_baseline.as_ref(), self.raw_items()); - let rollout_item = self.world_state_item(&snapshot); - self.world_state_baseline = Some(snapshot); - (fragments, rollout_item) - } - pub(crate) fn render_step_world_state( &self, world_state: &WorldState, diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index 5a314620a4..d49bb423b4 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -180,6 +180,17 @@ fn estimate_response_item_model_visible_bytes(item: &ResponseItem) -> i64 { super::estimate_response_item_model_visible_bytes(item, &FileAudioDurations::default()) } +impl ContextManager { + fn update_world_state( + &mut self, + world_state: &WorldState, + ) -> (Vec, Option) { + let (snapshot, updates, item) = self.render_step_world_state(world_state); + self.set_world_state_baseline(snapshot); + (updates, item) + } +} + #[test] fn delivered_assistant_context_invalidates_reviews_without_changing_authorization() { let revisions = |history: &ContextManager| { @@ -765,6 +776,7 @@ fn world_state_reconciles_matching_legacy_history_once() { .into_iter() .map(|update| match update.content { WorldStateUpdateContent::Fragment(fragment) => fragment.body().expect_text(), + WorldStateUpdateContent::Environment(_) => panic!("unexpected environment"), WorldStateUpdateContent::Item(item) => panic!("expected a text fragment, got {item:?}"), }) diff --git a/codex-rs/core/src/context_manager/updates.rs b/codex-rs/core/src/context_manager/updates.rs index fa3ded9147..f2520c49e6 100644 --- a/codex-rs/core/src/context_manager/updates.rs +++ b/codex-rs/core/src/context_manager/updates.rs @@ -2,7 +2,6 @@ use crate::context::world_state::Placement; use crate::context::world_state::WorldStateUpdate; -use crate::context::world_state::WorldStateUpdateContent; use codex_context_fragments::RenderedFragment; use codex_context_fragments::message_from_parts; use codex_protocol::models::ResponseItem; @@ -22,9 +21,9 @@ pub(crate) fn merge_world_state_updates(updates: Vec) -> Vec::new(); for update in updates { - let fragment = match update.content { - WorldStateUpdateContent::Fragment(fragment) => fragment, - WorldStateUpdateContent::Item(item) => { + let fragment = match update.content.into_fragment() { + Ok(fragment) => fragment, + Err(item) => { items.extend(build_rendered_message(std::mem::take(&mut pending))); items.push(*item); continue; diff --git a/codex-rs/core/src/context_manager/updates_tests.rs b/codex-rs/core/src/context_manager/updates_tests.rs index 7f937f7c1d..3cae62964a 100644 --- a/codex-rs/core/src/context_manager/updates_tests.rs +++ b/codex-rs/core/src/context_manager/updates_tests.rs @@ -2,6 +2,7 @@ use super::*; use crate::context::ContextualUserFragment; +use crate::context::world_state::WorldStateUpdateContent; use codex_protocol::models::ContentItem; use codex_protocol::models::ContentItemKind; use codex_protocol::models::InternalChatMessageMetadataPassthrough; diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 4393f14db0..553689953f 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -158,7 +158,9 @@ mod thread_manager; mod thread_spawn; mod thread_store_builder; pub use agent::host::AgentMcpMetadataRecipient; +pub use agent::host::AgentRef; pub use agent::host::AgentRuntimeOwnership; +pub use agent::host::AgentSessionReader; pub use agent::host::AgentTurnAdmission; pub use agent::host::AgentTurnAdmissionRequest; pub use agent::host::AgentUsageAccounting; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 40439db6bc..1997508e72 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -4127,12 +4127,13 @@ impl Session { let world_state = self .build_world_state_for_step(step_context, /*new_window*/ false) .await?; - let (world_state_snapshot, fragments, world_state_item) = self + let (world_state_snapshot, mut fragments, world_state_item) = self .state .lock() .await .history .render_step_world_state(&world_state); + world_state.resolve_host_roster(&mut fragments).await; let items = crate::context_manager::updates::merge_world_state_updates(fragments); if !items.is_empty() { self.record_conversation_items(turn_context, &step_context.settings.model_info, &items) @@ -4967,6 +4968,15 @@ impl Session { &self, step_context: &StepContext, world_state: &WorldState, + ) -> (Vec, WorldStateSnapshot) { + let rendered = world_state.render_full_with_host_roster().await; + self.assemble_initial_context(step_context, rendered).await + } + + pub(crate) async fn assemble_initial_context( + &self, + step_context: &StepContext, + (snapshot, updates): (WorldStateSnapshot, Vec), ) -> (Vec, WorldStateSnapshot) { let turn_context = step_context.turn.as_ref(); let mut developer_sections = Vec::::with_capacity(8); @@ -5114,7 +5124,6 @@ impl Session { .render_with_hints(context_window_hints), ); } - let (snapshot, updates) = world_state.render_full(); let (mut prefix, context) = split_prefix_updates(updates); let mut items = Vec::with_capacity(4); let mut context = context.into_iter(); @@ -5125,9 +5134,9 @@ impl Session { let mut initial_multi_agent_mode = None; let mut managed_developer_instructions = None; for update in context.by_ref() { - let fragment = match update.content { - WorldStateUpdateContent::Fragment(fragment) => fragment, - WorldStateUpdateContent::Item(item) => { + let fragment = match update.content.into_fragment() { + Ok(fragment) => fragment, + Err(item) => { next_item = Some(*item); break; } @@ -5339,14 +5348,11 @@ impl Session { crate::compact_remote_v2::RETAINED_MESSAGE_TOKEN_BUDGET, &history.file_audio_durations, ); - let window = { - let mut state = self.state.lock().await; - state.start_new_context_window() - }; - let (window_number, window_ids) = window; - let (context_updates, world_state_snapshot) = self - .build_initial_context_with_world_state(step_context, world_state.as_ref()) - .await; + // Resolve the roster before allocating a new context window. + let rendered = world_state.render_full_with_host_roster().await; + let (window_number, window_ids) = self.state.lock().await.start_new_context_window(); + let (context_updates, world_state_snapshot) = + self.assemble_initial_context(step_context, rendered).await; let context_items = crate::context_manager::updates::merge_world_state_updates(context_updates) .into_iter() @@ -5424,15 +5430,20 @@ impl Session { Some(WorldStateItem::full(snapshot.into_object())), ) } else { - let (world_state_items, world_state_item) = { - let mut state = self.state.lock().await; - let (fragments, rollout_item) = state.history.update_world_state(&world_state); - ( - crate::context_manager::updates::merge_world_state_updates(fragments), - rollout_item, - ) - }; - (world_state_items, world_state_item) + let (snapshot, mut fragments, world_state_item) = self + .state + .lock() + .await + .history + .render_step_world_state(&world_state); + world_state.resolve_host_roster(&mut fragments).await; + let items = crate::context_manager::updates::merge_world_state_updates(fragments); + self.state + .lock() + .await + .history + .set_world_state_baseline(snapshot); + (items, world_state_item) }; if !should_inject_full_context && turn_context_changed { context_items.extend( diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 30a294cc1a..34ea63e75f 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -165,6 +165,40 @@ pub(crate) struct McpStartupRequirements { required_plugins: HashSet, } +/// Keep the initial input recoverable while startup can block on a host roster read. +async fn preserve_initial_input_during( + sess: &Session, + turn_context: &Arc, + input: &mut Vec, + prepare: impl Future, +) -> CodexResult { + if input.is_empty() || sess.host_roster_reader(turn_context).is_none() { + return Ok(prepare.await); + } + + let exchange_input = async |replacement| { + let mut active = sess.active_turn.lock().await; + let task = active + .as_mut() + .and_then(|turn| turn.task.as_mut()) + .filter(|task| Arc::ptr_eq(&task.turn_context, turn_context)) + .ok_or_else(|| CodexErr::new(CodexErrorDetails::TurnAborted))?; + Ok::<_, CodexErr>(std::mem::replace( + &mut task.pending_initial_input, + replacement, + )) + }; + + let previous = exchange_input(Some(std::mem::take(input))).await?; + debug_assert!(previous.is_none()); + let result = prepare.await; + // If suspension took the task, its cleanup owns the input. + *input = exchange_input(/* replacement */ None) + .await? + .ok_or_else(|| CodexErr::Fatal("initial input missing after turn preparation".into()))?; + Ok(result) +} + /// Takes initial turn input and runs a loop where, at each sampling request, /// the model replies with either: /// @@ -224,13 +258,18 @@ pub(crate) async fn run_turn( // diffs/full reinjection + user input) and trigger compaction preemptively // when they would push the thread over the compaction threshold. if auto_compaction_enabled - && let Err(err) = run_pre_sampling_compact( + && let Err(err) = preserve_initial_input_during( &sess, &turn_context, - &mut client_session, - &cancellation_token, + &mut input, + run_pre_sampling_compact( + &sess, + &turn_context, + &mut client_session, + &cancellation_token, + ), ) - .await + .await? { // Compaction runs before the new input is recorded, so preserve it on every failure. let _ = run_hooks_and_record_inputs( @@ -324,7 +363,12 @@ pub(crate) async fn run_turn( }; // Record initial context while preparing diff display roots. let (record_context, display_roots) = tokio::join!( - sess.record_context_updates_and_set_reference_context_item(first_step_context.as_ref()), + preserve_initial_input_during( + &sess, + &turn_context, + &mut input, + sess.record_context_updates_and_set_reference_context_item(first_step_context.as_ref()), + ), async { // Guardian must not wait for remote Git discovery just to display diff paths. if crate::guardian::is_basic_session_source(&turn_context.session_source) @@ -349,7 +393,7 @@ pub(crate) async fn run_turn( } }, ); - let mut world_state = record_context?; + let mut world_state = record_context??; let Some((injection_items, explicitly_enabled_connectors)) = build_skills_and_plugins( &sess, diff --git a/codex-rs/core/src/session/turn_suspension.rs b/codex-rs/core/src/session/turn_suspension.rs index 233f0bfccb..b5d9cc54de 100644 --- a/codex-rs/core/src/session/turn_suspension.rs +++ b/codex-rs/core/src/session/turn_suspension.rs @@ -80,13 +80,12 @@ pub(super) async fn suspend_turn_and_shutdown( let _ = task_handle.await; } } - // The task did not run while admission was pending. Preserve its initial input - // before closing the writer so a new worker can recover the same turn. + // Save any initial input the task had not claimed so a new worker can recover the turn. let recorded = session - .record_unadmitted_input( + .record_pending_initial_input( &task.turn_context, &turn.turn_state, - task.pending_admission_input, + task.pending_initial_input, /*abort_reason*/ None, ) .await; diff --git a/codex-rs/core/src/session/world_state.rs b/codex-rs/core/src/session/world_state.rs index f4c8ee6249..13204bfa41 100644 --- a/codex-rs/core/src/session/world_state.rs +++ b/codex-rs/core/src/session/world_state.rs @@ -2,6 +2,9 @@ use std::sync::Arc; use super::session::Session; use super::step_context::StepContext; +use super::turn_context::TurnContext; +use crate::agent::host::AgentRuntimeServices; +use crate::agent::host::AgentSessionReader; use crate::connectors; use crate::context::TokenBudgetContext; use crate::context::world_state::AgentsMdState; @@ -23,6 +26,7 @@ use crate::context::world_state::RealtimeState; use crate::context::world_state::ToolsState; use crate::context::world_state::TopLevelToolsState; use crate::context::world_state::WorldState; +use crate::context::world_state::format_agent_paths; use crate::tools::handlers::multi_agents_spec::MULTI_AGENT_V1_NAMESPACE; use codex_connectors::AppToolPolicyEvaluator; use codex_extension_api::WorldStateContributionInput; @@ -37,10 +41,21 @@ use codex_protocol::models::ContentItemNamespace; use codex_protocol::protocol::MultiAgentVersion; use codex_tools::ToolName; -const MAX_ENVIRONMENT_SUBAGENTS: usize = 8; -const MAX_ENVIRONMENT_SUBAGENT_BYTES: usize = 1_024; - impl Session { + pub(super) fn host_roster_reader( + &self, + turn_context: &TurnContext, + ) -> Option<&Arc> { + if turn_context.config.include_environment_context + && turn_context.multi_agent_version == MultiAgentVersion::V2 + && let AgentRuntimeServices::Host(services) = &self.agent_runtime + { + services.session_reader.as_ref() + } else { + None + } + } + pub(crate) async fn current_window_uses_incremental_tools( &self, step_context: &StepContext, @@ -135,33 +150,16 @@ impl Session { ) }; let environment_subagents = if turn_context.config.include_environment_context - && matches!( - self.agent_runtime, - crate::agent::host::AgentRuntimeServices::Controller - ) { + && matches!(self.agent_runtime, AgentRuntimeServices::Controller) + { match turn_context.multi_agent_version { - MultiAgentVersion::V2 => { - let agent_paths = self - .services + MultiAgentVersion::V2 => format_agent_paths( + self.services .agent_control .child_agent_paths(self.thread_id) - .await; - let mut lines = - Vec::with_capacity(agent_paths.len().min(MAX_ENVIRONMENT_SUBAGENTS)); - let mut rendered_bytes = " \n \n".len(); - for agent_path in agent_paths { - if lines.len() == MAX_ENVIRONMENT_SUBAGENTS { - break; - } - let line = format!(r#""#); - let line_bytes = " \n".len() + line.len(); - if rendered_bytes + line_bytes <= MAX_ENVIRONMENT_SUBAGENT_BYTES { - rendered_bytes += line_bytes; - lines.push(line); - } - } - lines.join("\n") - } + .await + .into_iter(), + ), MultiAgentVersion::Disabled | MultiAgentVersion::V1 => { self.services .local_agent_runtime @@ -298,6 +296,7 @@ impl Session { )); } if turn_context.config.include_environment_context { + world_state.host_roster_reader = self.host_roster_reader(turn_context).cloned(); let current_date = if turn_context.config.include_environment_context_time { self.read_clock_for_context(turn_context, "environment_date") .await? diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index 79dad90719..4dedfc461b 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -86,8 +86,8 @@ pub(crate) struct RunningTask { pub(crate) handle: AbortOnDropHandle<()>, pub(crate) turn_context: Arc, pub(crate) _agent_execution_guard: Option, - /// Initial input still owned by Core while the host decides whether the task may run. - pub(crate) pending_admission_input: Option>, + /// Unrecorded input retained during host admission or preparation that can wait on the roster. + pub(crate) pending_initial_input: Option>, pub(crate) _diagnostics_guard: GaugeGuard, // Timer recorded when the task drops to capture the full turn duration. pub(crate) _timer: Option, diff --git a/codex-rs/core/src/tasks/admission.rs b/codex-rs/core/src/tasks/admission.rs index e796ed93dc..40baba5f70 100644 --- a/codex-rs/core/src/tasks/admission.rs +++ b/codex-rs/core/src/tasks/admission.rs @@ -61,7 +61,7 @@ pub(super) async fn admit( .filter(|_| !cancellation.is_cancelled()) .ok_or_else(|| CodexErr::new(CodexErrorDetails::TurnAborted))?; task._agent_execution_guard = Some(guard); - task.pending_admission_input.take().ok_or_else(|| { + task.pending_initial_input.take().ok_or_else(|| { CodexErr::Fatal("initial input missing from the task awaiting host admission".into()) }) } @@ -77,7 +77,7 @@ impl Session { } /// Interruption and suspension stop the task before it can persist its initial input. - pub(crate) async fn record_unadmitted_input( + pub(crate) async fn record_pending_initial_input( self: &Arc, context: &Arc, turn_state: &Mutex, diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 5337d7f364..8f02fd0fe5 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -362,7 +362,7 @@ impl Session { ) .await; let host_admission = self.host_turn_admission(); - let (agent_execution_guard, task_input, pending_admission_input) = match &host_admission { + let (agent_execution_guard, task_input, pending_initial_input) = match &host_admission { Some(_) => (None, Vec::new(), Some(input)), None => ( self.services.agent_control.admit_turn( @@ -474,7 +474,7 @@ impl Session { cancellation_token, turn_context: Arc::clone(&turn_context), _agent_execution_guard: agent_execution_guard, - pending_admission_input, + pending_initial_input, _diagnostics_guard: ACTIVE_TURNS.track(), _timer: timer, }; @@ -739,8 +739,9 @@ impl Session { active_turn.finishing = Some(finishing); ( Arc::clone(&active_turn.turn_state), - task.pending_admission_input.is_some().then_some(task.task), - task.pending_admission_input.unwrap_or_default(), + (task._agent_execution_guard.is_none() && task.pending_initial_input.is_some()) + .then_some(task.task), + task.pending_initial_input.unwrap_or_default(), ) }; turn_context @@ -1078,10 +1079,10 @@ impl Session { let _ = task.handle.await; if let Err(error) = self - .record_unadmitted_input( + .record_pending_initial_input( &task.turn_context, turn_state, - task.pending_admission_input, + task.pending_initial_input, Some(&reason), ) .await diff --git a/codex-rs/core/tests/common/context_snapshot.rs b/codex-rs/core/tests/common/context_snapshot.rs index 7261f0f902..61e27406f8 100644 --- a/codex-rs/core/tests/common/context_snapshot.rs +++ b/codex-rs/core/tests/common/context_snapshot.rs @@ -26,6 +26,7 @@ const MAX_SNAPSHOT_LINE_CHARS: usize = 160; #[derive(Debug, Clone, Default)] pub struct ContextSnapshotOptions { rewrite_known_segments: bool, + preserve_environment_context: bool, include_request_settings: bool, } @@ -38,6 +39,12 @@ impl ContextSnapshotOptions { self } + /// Keep normalized environment text when rewriting other known segments. + pub fn preserve_environment_context(mut self) -> Self { + self.preserve_environment_context = true; + self + } + /// Render model, instructions, tools, and other settings at window boundaries. /// Settings still participate in grouping when this is disabled. pub fn include_request_settings(mut self) -> Self { diff --git a/codex-rs/core/tests/common/context_snapshot/normalize.rs b/codex-rs/core/tests/common/context_snapshot/normalize.rs index 46c21ad053..feef8d610c 100644 --- a/codex-rs/core/tests/common/context_snapshot/normalize.rs +++ b/codex-rs/core/tests/common/context_snapshot/normalize.rs @@ -63,6 +63,7 @@ impl Normalizer { let text = self.normalize_values(&text); if options.rewrite_known_segments && let Some(segment) = segment + && !(options.preserve_environment_context && segment == "ENVIRONMENT_CONTEXT") { let marker = format!("<{segment}>"); if segment == "COMPACTION_SUMMARY" diff --git a/codex-rs/core/tests/suite/agent_control_tests.rs b/codex-rs/core/tests/suite/agent_control_tests.rs index e5368c29ee..2253fb494b 100644 --- a/codex-rs/core/tests/suite/agent_control_tests.rs +++ b/codex-rs/core/tests/suite/agent_control_tests.rs @@ -3,11 +3,14 @@ use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use codex_core::AgentConfigUpdate; use codex_core::AgentControl; use codex_core::AgentExecutionGuard; use codex_core::AgentInfo; +use codex_core::AgentRef; use codex_core::AgentTarget; use codex_core::AgentTurnOutcome; use codex_core::DeliveryReceipt; @@ -71,8 +74,46 @@ impl codex_core::AgentTurnAdmission for HostRuntime { } } +pub(super) struct HostSessionReader { + root: ThreadId, + pub(super) children: Mutex>>, + pub(super) children_reads: AtomicUsize, +} + +impl codex_core::AgentSessionReader for HostSessionReader { + fn root(&self) -> BoxFuture<'_, CodexResult> { + Box::pin(async { + Ok(AgentRef { + thread_id: self.root, + path: AgentPath::root(), + }) + }) + } + + fn parent(&self) -> BoxFuture<'_, CodexResult>> { + Box::pin(async { Ok(None) }) + } + + fn children(&self) -> BoxFuture<'_, CodexResult>> { + Box::pin(async { + self.children_reads.fetch_add(1, Ordering::Relaxed); + self.children + .lock() + .expect("host roster") + .clone() + .ok_or_else(|| CodexErr::InvalidRequest("host agent session unavailable".into())) + }) + } +} + pub(super) fn install_host_runtime(data: &codex_extension_api::ExtensionData) { data.insert(codex_core::AgentRuntimeOwnership::Host); + let reader = data.get_or_init(|| HostSessionReader { + root: ThreadId::from_string(data.level_id()).expect("test thread identity"), + children: Mutex::new(Some(Vec::new())), + children_reads: AtomicUsize::new(0), + }); + data.get_or_init(|| reader as Arc); data.get_or_init(|| Arc::new(HostRuntime) as Arc); } diff --git a/codex-rs/core/tests/suite/host_turn_admission.rs b/codex-rs/core/tests/suite/host_turn_admission.rs index 6566ccbcbb..99e190ee77 100644 --- a/codex-rs/core/tests/suite/host_turn_admission.rs +++ b/codex-rs/core/tests/suite/host_turn_admission.rs @@ -6,12 +6,15 @@ use std::time::Duration; use anyhow::Context; use anyhow::Result; use codex_core::AgentExecutionGuard; +use codex_core::AgentRef; +use codex_core::AgentSessionReader; use codex_core::AgentTurnAdmission; use codex_core::AgentTurnAdmissionRequest; use codex_core::CodexThread; use codex_core::SuspendTurnOutcome; use codex_core::TurnInputRequest; use codex_core::TurnInputSubmission; +use codex_features::Feature; use codex_history::RolloutItem; use codex_protocol::dynamic_tools::DynamicToolFunctionSpec; use codex_protocol::dynamic_tools::DynamicToolResponse; @@ -23,6 +26,7 @@ use codex_protocol::items::TurnItem; use codex_protocol::models::ResponseItem; use codex_protocol::models::input_text; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::Op; use codex_protocol::protocol::ReviewRequest; use codex_protocol::protocol::ReviewTarget; @@ -41,6 +45,7 @@ use core_test_support::wait_for_successful_turn; use futures::future::BoxFuture; use pretty_assertions::assert_eq; use serde_json::json; +use tokio::sync::Notify; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio_util::sync::CancellationToken; @@ -63,9 +68,30 @@ impl AgentTurnAdmission for HostAdmission { } } +#[derive(Default)] +struct PendingRosterRead(Notify); + +impl AgentSessionReader for PendingRosterRead { + fn root(&self) -> BoxFuture<'_, CodexResult> { + unreachable!("roster rendering reads children") + } + + fn parent(&self) -> BoxFuture<'_, CodexResult>> { + unreachable!("roster rendering reads children") + } + + fn children(&self) -> BoxFuture<'_, CodexResult>> { + Box::pin(async { + self.0.notify_one(); + std::future::pending().await + }) + } +} + async fn root_with_admission( server: &MockServer, mut builder: TestCodexBuilder, + session_reader: Option>, ) -> Result<(TestCodex, mpsc::UnboundedReceiver)> { let mut test = builder.build_with_auto_env(server).await?; let mut options = test.start_thread_options().await; @@ -75,6 +101,9 @@ async fn root_with_admission( .thread_extension_init .insert(codex_core::AgentRuntimeOwnership::Host); options.thread_extension_init.insert(admission); + if let Some(session_reader) = session_reader { + options.thread_extension_init.insert(session_reader); + } options.dynamic_tools = vec![DynamicToolSpec::Function(DynamicToolFunctionSpec { name: "hold_turn".into(), description: "Wait for the test to release this turn.".into(), @@ -135,7 +164,8 @@ fn assert_denied(event: &EventMsg, turn_id: &str) { async fn root_denial_and_interruption_retain_input_and_grants_span_execution() -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; - let (test, mut attempts) = root_with_admission(&server, test_codex()).await?; + let (test, mut attempts) = + root_with_admission(&server, test_codex(), /*session_reader*/ None).await?; let (turn_id, decision) = start(&test, &mut attempts, "Denied initial prompt.").await?; assert!( matches!(submit(&test.codex, "Steered while pending.").await?, @@ -216,8 +246,15 @@ async fn interrupt_before_admission_skips_submit_hook_and_uninspected_prompt() - ]}]}}); std::fs::write(home.join("hooks.json"), hooks.to_string()).unwrap(); }) - .with_config(trust_discovered_hooks); - let (test, mut attempts) = root_with_admission(&server, builder).await?; + .with_config(|config| { + trust_discovered_hooks(config); + config + .features + .enable(Feature::MultiAgentV2) + .expect("enable multi-agent V2"); + }); + let (test, mut attempts) = + root_with_admission(&server, builder, /*session_reader*/ None).await?; let interrupted = "Interrupted before inspection."; let (turn_id, _pending) = start(&test, &mut attempts, interrupted).await?; test.codex.submit(Op::Interrupt).await?; @@ -239,9 +276,19 @@ async fn interrupt_before_admission_skips_submit_hook_and_uninspected_prompt() - }) .await; wait_for_successful_turn(&test.codex).await; + assert_eq!( + test.codex.multi_agent_version(), + Some(MultiAgentVersion::V2) + ); let input = model.single_request().message_input_texts("user"); assert!(input.iter().any(|text| text == "Continue.")); assert!(input.iter().all(|text| text != interrupted)); + assert!( + input + .iter() + .any(|text| text.contains("")) + ); + assert!(input.iter().all(|text| !text.contains(""))); test.codex.shutdown_and_wait().await?; Ok(()) } @@ -250,14 +297,44 @@ async fn interrupt_before_admission_skips_submit_hook_and_uninspected_prompt() - async fn suspending_root_during_admission_preserves_the_recoverable_prompt() -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; - let (test, mut attempts) = root_with_admission(&server, test_codex()).await?; + let (test, mut attempts) = + root_with_admission(&server, test_codex(), /*session_reader*/ None).await?; let prompt = "Save this prompt across handoff."; let (turn_id, pending) = start(&test, &mut attempts, prompt).await?; + suspend_and_check_recoverable_prompt(&test, &server, turn_id, prompt).await?; + assert!(pending.is_closed()); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn suspending_root_during_roster_read_preserves_the_recoverable_prompt() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = responses::start_mock_server().await; + let builder = test_codex().with_config(|config| { + config + .features + .enable(Feature::MultiAgentV2) + .expect("enable multi-agent V2"); + }); + let roster = Arc::new(PendingRosterRead::default()); + let (test, mut attempts) = root_with_admission(&server, builder, Some(roster.clone())).await?; + let prompt = "Save this admitted prompt across handoff."; + let (turn_id, grant) = start(&test, &mut attempts, prompt).await?; + assert!(grant.send(Ok(AgentExecutionGuard::new(()))).is_ok()); + tokio::time::timeout(Duration::from_secs(10), roster.0.notified()).await?; + suspend_and_check_recoverable_prompt(&test, &server, turn_id, prompt).await +} + +async fn suspend_and_check_recoverable_prompt( + test: &TestCodex, + server: &MockServer, + turn_id: String, + prompt: &str, +) -> Result<()> { assert_eq!( test.codex.suspend_turn_and_shutdown().await?, SuspendTurnOutcome::Suspended { turn_id } ); - assert!(pending.is_closed()); let history = test .thread_store .load_latest_model_context(LoadThreadHistoryParams { @@ -279,7 +356,7 @@ async fn suspending_root_during_admission_preserves_the_recoverable_prompt() -> .count(); assert_eq!(saved_prompts, 1); assert!( - responses::received_responses_requests(&server) + responses::received_responses_requests(server) .await .is_empty() ); @@ -290,7 +367,8 @@ async fn suspending_root_during_admission_preserves_the_recoverable_prompt() -> async fn denied_review_exits_review_mode() -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; - let (test, mut attempts) = root_with_admission(&server, test_codex()).await?; + let (test, mut attempts) = + root_with_admission(&server, test_codex(), /*session_reader*/ None).await?; test.codex .submit(Op::Review { review_request: ReviewRequest { diff --git a/codex-rs/core/tests/suite/scenarios_budget_context.rs b/codex-rs/core/tests/suite/scenarios_budget_context.rs index 59c262485f..fe99fe6821 100644 --- a/codex-rs/core/tests/suite/scenarios_budget_context.rs +++ b/codex-rs/core/tests/suite/scenarios_budget_context.rs @@ -1,8 +1,10 @@ -//! Host budget reminders use generic World State and refresh after compaction spending. +//! Host budgets and agent session reads use World State. +use super::super::agent_control::HostSessionReader; use super::super::rollout_budget::HostAccounting; use super::super::rollout_budget::rollout_budget; use super::*; +use codex_core::AgentRef; use codex_extension_api::ContextContributor; use codex_extension_api::ExtensionFuture; use codex_extension_api::PreviousWorldStateSection; @@ -10,6 +12,10 @@ use codex_extension_api::RenderedWorldStateFragment; use codex_extension_api::WorldStateContributionInput; use codex_extension_api::WorldStateSectionContribution; use codex_extension_api::input_text; +use codex_protocol::AgentPath; +use codex_protocol::ThreadId; +use codex_protocol::protocol::ThreadSettingsOverrides; +use codex_protocol::protocol::TurnEnvironmentRequests; use core_test_support::responses::received_responses_requests; use core_test_support::responses::sse_response; use core_test_support::responses::strip_metadata_from_json; @@ -124,9 +130,23 @@ async fn host_budget_context_refreshes_after_remote_compaction() -> Result<()> { configure_scenario_catalog(config); config.workspace_roots = vec![config.cwd.clone()]; config.rollout_budget = Some(rollout_budget()); + config + .features + .enable(Feature::MultiAgentV2) + .expect("enable multi-agent V2"); }) .build_with_auto_env(&server) .await?; + let reader = test + .codex + .thread_extension_data() + .get::() + .expect("host agent session reader"); + let child = AgentRef { + thread_id: ThreadId::new(), + path: AgentPath::root().join("release").expect("child path"), + }; + *reader.children.lock().expect("host roster") = Some(vec![child.clone()]); mount_sse_sequence( &server, @@ -139,8 +159,18 @@ async fn host_budget_context_refreshes_after_remote_compaction() -> Result<()> { .await; test.submit_turn("Start planning the release.").await?; remaining.store(900, Ordering::SeqCst); + // An unchanged environment does not need to read the roster. + let reads = reader.children_reads.load(Ordering::Relaxed); test.submit_turn("Continue planning.").await?; + assert_eq!(reader.children_reads.load(Ordering::Relaxed), reads); remaining.store(400, Ordering::SeqCst); + *reader.children.lock().expect("host roster") = Some(vec![ + child, + AgentRef { + thread_id: ThreadId::new(), + path: AgentPath::root().join("testing").expect("child path"), + }, + ]); test.submit_turn("Check the release plan.").await?; let compaction_budget = Arc::clone(&remaining); @@ -177,6 +207,7 @@ async fn host_budget_context_refreshes_after_remote_compaction() -> Result<()> { vec![ sse(vec![ev_completed("fresh-window")]), sse(vec![ev_completed("same-fresh-threshold")]), + sse(vec![ev_completed("roster-unavailable")]), ], ) .await; @@ -185,10 +216,50 @@ async fn host_budget_context_refreshes_after_remote_compaction() -> Result<()> { remaining.store(300, Ordering::SeqCst); test.submit_turn("Finish the release plan.").await?; + // Change the working directory to refresh the environment fragment while + // preserving the executor and its tools. + let mut environment = test.executor_environment().request(); + environment.cwd = test.workspace_path_uri("roster-check")?; + test.fs() + .create_directory( + &environment.cwd, + CreateDirectoryOptions { + recursive: false, + follow_symlinks: true, + }, + /*sandbox*/ None, + ) + .await?; + *reader.children.lock().expect("host roster") = None; + let reads = reader.children_reads.load(Ordering::Relaxed); + test.codex + .start_or_steer_turn( + TurnInputRequest::user_input(vec![UserInput::Text { + text: "Check the unavailable roster.".into(), + text_elements: Vec::new(), + }]) + .with_thread_settings(ThreadSettingsOverrides { + environments: Some(TurnEnvironmentRequests::new( + test.config.cwd.join("roster-check"), + vec![environment], + )), + ..Default::default() + }), + ) + .await?; + wait_for_event(&test.codex, |event| { + assert!(!matches!(event, EventMsg::Error(_)), "{event:?}"); + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + assert_eq!(reader.children_reads.load(Ordering::Relaxed), reads + 1); + let requests = received_responses_requests(&server).await; - assert_eq!(requests.len(), 6); + assert_eq!(requests.len(), 7); for request in &requests { assert!(!request.has_content_kinds(&["rollout_budget.remaining_tokens"])); + let input = serde_json::to_string(&request.input())?; + assert!(input.contains("You are `/root`")); } // Compare the actual wire input without normalization: both unchanged and // threshold-crossing reminders must only append within a context window. @@ -208,6 +279,7 @@ async fn host_budget_context_refreshes_after_remote_compaction() -> Result<()> { ); } } + let resumed_input = requests[4].input(); // Compaction rebuilds item IDs and metadata while preserving the tool // declarations and base instructions ahead of the conversation. @@ -217,13 +289,14 @@ async fn host_budget_context_refreshes_after_remote_compaction() -> Result<()> { strip_metadata_from_json(strip_response_item_ids_from_json(initial_prefix)), strip_metadata_from_json(strip_response_item_ids_from_json(resumed_prefix)), ); - insta::assert_snapshot!( - "host_budget_context_after_compaction", - context_snapshot::format_request_history_snapshot( - "A host contributes budget reminders at thresholds. Remote compaction spends within the same threshold, so its captured reminder is omitted and the next sampling step inserts the fresh remainder once.", - &requests, - &ContextSnapshotOptions::default().rewrite_known_segments(), - ) + + let snapshot = context_snapshot::format_request_history_snapshot( + "Core renders the host's roster with native World State cadence. Compaction rebuilds it from current membership, and the next sampling step refreshes the budget after compaction spending.", + &requests, + &ContextSnapshotOptions::default() + .rewrite_known_segments() + .preserve_environment_context(), ); + insta::assert_snapshot!("host_budget_context_after_compaction", snapshot); Ok(()) } diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__budget_context__host_budget_context_after_compaction.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__budget_context__host_budget_context_after_compaction.snap index 01869797f8..822205d740 100644 --- a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__budget_context__host_budget_context_after_compaction.snap +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__budget_context__host_budget_context_after_compaction.snap @@ -1,8 +1,8 @@ --- source: core/tests/suite/scenarios_budget_context.rs -expression: "context_snapshot::format_request_history_snapshot(\"A host contributes budget reminders at thresholds. Remote compaction spends within the same threshold, so its captured reminder is omitted and the next sampling step inserts the fresh remainder once.\",\n&requests, &ContextSnapshotOptions::default().rewrite_known_segments(),)" +expression: snapshot --- -Scenario: A host contributes budget reminders at thresholds. Remote compaction spends within the same threshold, so its captured reminder is omitted and the next sampling step inserts the fresh remainder once. +Scenario: Core renders the host's roster with native World State cadence. Compaction rebuilds it from current membership, and the next sampling step refreshes the budget after compaction spending. ## Window 1 -- request 1 (turn) -- @@ -31,7 +31,16 @@ Scenario: A host contributes budget reminders at thresholds. Remote compaction s 04:message/developer: 05:message/user: - + + + + + + <...e="unrestricted" /> [hash=d40b2f45cb028888] + + + + 06:message/user: Start planning the release. 07:message/developer: @@ -83,7 +92,17 @@ Scenario: A host contributes budget reminders at thresholds. Remote compaction s 06:message/developer: 07:message/user: - + + + + + + <...e="unrestricted" /> [hash=d40b2f45cb028888] + + + + + 08:message/user: Check the release plan. 09:compaction:encrypted=true; chars=25; hash=455278c90f6b4d93 @@ -96,3 +115,14 @@ Scenario: A host contributes budget reminders at thresholds. Remote compaction s -- request 6 (turn) -- 12:message/user: Finish the release plan. +-- request 7 (turn) -- +13:message/user: + + /roster-check + + + + <...e="unrestricted" /> [hash=d40b2f45cb028888] + +14:message/user: + Check the unavailable roster.