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.