openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-07 20:30:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Render host-owned V2 subagent rosters in environment context (#59531)

Browse Source 
## 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
This commit is contained in:
rhan-oai
2026-11-05 00:39:39 +00:00
committed by copyberry
parent 9a7213ca37
commit f3074e35c0
33 changed files with 642 additions and 210 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/agent/host/mod.rs
+3
View File
@@ -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;
codex-rs/core/src/agent/host/ownership.rs
+8 -2
View File
@@ -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<Arc<dyn AgentSessionReader>>,
pub(crate) turn_admission: Arc<dyn AgentTurnAdmission>,
pub(crate) usage_accounting: Option<Arc<dyn AgentUsageAccounting>>,
}
@@ -40,6 +42,10 @@ impl AgentRuntimeServices {
)
})?;
Ok(Self::Host(HostAgentServices {
session_reader: data
.get::<Arc<dyn AgentSessionReader>>()
.as_deref()
.cloned(),
turn_admission: Arc::clone(turn_admission.as_ref()),
usage_accounting: data
.get::<Arc<dyn AgentUsageAccounting>>()
codex-rs/core/src/agent/host/session_reader.rs
+32
View File
@@ -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<AgentRef>>;
/// Returns the structural parent; only the root has no parent.
fn parent(&self) -> BoxFuture<'_, Result<Option<AgentRef>>>;
/// 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<Vec<AgentRef>>>;
}
codex-rs/core/src/compact.rs
+7 -8
View File
@@ -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(),
codex-rs/core/src/compact_remote_v2_installation_tests.rs
+86 -43
View File
@@ -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<AgentRef>> {
unreachable!("roster rendering reads children")
}
fn parent(&self) -> BoxFuture<'_, CodexResult<Option<AgentRef>>> {
unreachable!("roster rendering reads children")
}
fn children(&self) -> BoxFuture<'_, CodexResult<Vec<AgentRef>>> {
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::<Vec<_>>();
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::<Vec<_>>();
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, ..
codex-rs/core/src/context/token_budget_context.rs
+2 -9
View File
@@ -97,15 +97,8 @@ impl WorldStateSection for TokenBudgetContext {
) -> SectionTransition<Self::Snapshot> {
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<dyn ContextualUserFragment>);
(
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())
}
}
codex-rs/core/src/context/world_state/compact_permissions.rs
+8 -10
View File
@@ -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<Box<dyn ContextualUserFragment>> =
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<super::WorldStateFragment> = 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),
codex-rs/core/src/context/world_state/environment.rs
+31 -8
View File
@@ -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<dyn ContextualUserFragment>
});
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<String, EnvironmentUpdate>,
legacy_single: bool,
include_primary: bool,
@@ -235,7 +239,7 @@ struct RenderedEnvironments {
timezone: Option<String>,
network: Option<NetworkContext>,
filesystem: Option<FileSystemContext>,
subagents: Option<String>,
pub(crate) subagents: Option<String>,
}
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<Item = AgentPath>) -> String {
const MAX_AGENTS: usize = 8;
const MAX_BYTES: usize = 1_024;
let mut lines = Vec::new();
let mut rendered_bytes = " <subagents>\n </subagents>\n".len();
for path in paths {
if lines.len() == MAX_AGENTS {
break;
}
let line = format!(r#"<agent name="{path}" />"#);
let line_bytes = " \n".len() + line.len();
if rendered_bytes + line_bytes <= MAX_BYTES {
rendered_bytes += line_bytes;
lines.push(line);
}
}
lines.join("\n")
}
codex-rs/core/src/context/world_state/mod.rs
+58 -10
View File
@@ -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<dyn ContextualUserFragment + Send>;
pub(crate) enum WorldStateUpdateContent {
Fragment(Box<dyn ContextualUserFragment>),
Fragment(WorldStateFragment),
Environment(Box<RenderedEnvironments>),
Item(Box<ResponseItem>),
}
impl WorldStateUpdateContent {
pub(crate) fn into_fragment(self) -> Result<WorldStateFragment, Box<ResponseItem>> {
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<dyn ContextualUserFragment>) -> Self {
fn boxed_fragment(fragment: WorldStateFragment) -> Self {
Self {
placement: Placement::Mergeable,
content: WorldStateUpdateContent::Fragment(fragment),
}
}
pub(crate) fn optional_boxed_fragment(
fragment: Option<Box<dyn ContextualUserFragment>>,
) -> Vec<Self> {
pub(crate) fn optional_boxed_fragment(fragment: Option<WorldStateFragment>) -> Vec<Self> {
fragment.into_iter().map(Self::boxed_fragment).collect()
}
pub(crate) fn optional_standalone_boxed_fragment(
fragment: Option<Box<dyn ContextualUserFragment>>,
fragment: Option<WorldStateFragment>,
) -> Vec<Self> {
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<Arc<dyn AgentSessionReader>>,
sections: IndexMap<&'static str, Box<dyn ErasedWorldStateSection>>,
}
@@ -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<Item = &'static str> + '_ {
@@ -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<WorldStateUpdate>) {
let (snapshot, mut updates) = self.render_full();
self.resolve_host_roster(&mut updates).await;
(snapshot, updates)
}
pub(crate) fn render_full(&self) -> (WorldStateSnapshot, Vec<WorldStateUpdate>) {
self.render_with(RenderScope::FullContext, /*stored*/ None, |_, _| {
PreviousSectionState::Absent
codex-rs/core/src/context/world_state/model.rs
+2 -2
View File
@@ -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<Box<dyn ContextualUserFragment>> = if instructions_hash.is_some() {
let fragment: Option<super::WorldStateFragment> = 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<dyn ContextualUserFragment>
as super::WorldStateFragment
})
};
(
codex-rs/core/src/context/world_state/multi_agent_mode.rs
+1 -1
View File
@@ -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<dyn ContextualUserFragment>
as super::WorldStateFragment
});
(
Some(self.clone()),
codex-rs/core/src/context/world_state/multi_agent_usage_hint.rs
+1 -1
View File
@@ -46,7 +46,7 @@ impl WorldStateSection for MultiAgentUsageHintState {
&self,
previous: PreviousSectionState<'_, Self::Snapshot>,
) -> SectionTransition<Self::Snapshot> {
let fragment: Option<Box<dyn ContextualUserFragment>> = match previous {
let fragment: Option<super::WorldStateFragment> = match previous {
PreviousSectionState::Known(previous) if previous == &self.fingerprint => {
return (None, Vec::new());
}
codex-rs/core/src/context/world_state/realtime.rs
+2 -2
View File
@@ -35,14 +35,14 @@ impl RealtimeState {
}
}
fn render_start(&self) -> Box<dyn ContextualUserFragment> {
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<Box<dyn ContextualUserFragment>> {
fn render_transition(&self, previous_active: bool) -> Option<super::WorldStateFragment> {
match (previous_active, self.snapshot.active) {
(false, true) => Some(self.render_start()),
(true, false) => Some(match self.end_instructions.as_deref() {
codex-rs/core/src/context/world_state/test_support.rs
+6 -6
View File
@@ -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<dyn ContextualUserFragment> {
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:?}"),
}
}
codex-rs/core/src/context/world_state/world_state_tests.rs
+4 -4
View File
@@ -28,7 +28,7 @@ impl WorldStateSection for TestSection {
previous: PreviousSectionState<'_, Self::Snapshot>,
) -> SectionTransition<Self::Snapshot> {
let current = self.clone();
let fragment: Option<Box<dyn ContextualUserFragment>> = match previous {
let fragment: Option<super::WorldStateFragment> = 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<_>>(),
vec![
codex-rs/core/src/context_manager/history.rs
-11
View File
@@ -556,17 +556,6 @@ impl ContextManager {
self.reference_context_item.clone()
}
pub(crate) fn update_world_state(
&mut self,
world_state: &WorldState,
) -> (Vec<WorldStateUpdate>, Option<WorldStateItem>) {
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,
codex-rs/core/src/context_manager/history_tests.rs
+12
View File
@@ -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<WorldStateUpdate>, Option<WorldStateItem>) {
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:?}"),
})
codex-rs/core/src/context_manager/updates.rs
+3 -4
View File
@@ -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<WorldStateUpdate>) -> Vec<R
let mut items = Vec::new();
let mut pending = Vec::<RenderedFragment>::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;
codex-rs/core/src/context_manager/updates_tests.rs
+1
View File
@@ -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;
codex-rs/core/src/lib.rs
+2
View File
@@ -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;
codex-rs/core/src/session/mod.rs
+33 -22
View File
@@ -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<WorldStateUpdate>, 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<WorldStateUpdate>),
) -> (Vec<WorldStateUpdate>, WorldStateSnapshot) {
let turn_context = step_context.turn.as_ref();
let mut developer_sections = Vec::<RenderedFragment>::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(
codex-rs/core/src/session/turn.rs
+50 -6
View File
@@ -165,6 +165,40 @@ pub(crate) struct McpStartupRequirements {
required_plugins: HashSet<String>,
}
/// Keep the initial input recoverable while startup can block on a host roster read.
async fn preserve_initial_input_during<T>(
sess: &Session,
turn_context: &Arc<TurnContext>,
input: &mut Vec<TurnInput>,
prepare: impl Future<Output = T>,
) -> CodexResult<T> {
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,
codex-rs/core/src/session/turn_suspension.rs
+3 -4
View File
@@ -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;
codex-rs/core/src/session/world_state.rs
+26 -27
View File
@@ -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<dyn AgentSessionReader>> {
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 = " <subagents>\n </subagents>\n".len();
for agent_path in agent_paths {
if lines.len() == MAX_ENVIRONMENT_SUBAGENTS {
break;
}
let line = format!(r#"<agent name="{agent_path}" />"#);
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?
codex-rs/core/src/state/turn.rs
+2 -2
View File
@@ -86,8 +86,8 @@ pub(crate) struct RunningTask {
pub(crate) handle: AbortOnDropHandle<()>,
pub(crate) turn_context: Arc<TurnContext>,
pub(crate) _agent_execution_guard: Option<AgentExecutionGuard>,
/// Initial input still owned by Core while the host decides whether the task may run.
pub(crate) pending_admission_input: Option<Vec<TurnInput>>,
/// Unrecorded input retained during host admission or preparation that can wait on the roster.
pub(crate) pending_initial_input: Option<Vec<TurnInput>>,
pub(crate) _diagnostics_guard: GaugeGuard,
// Timer recorded when the task drops to capture the full turn duration.
pub(crate) _timer: Option<codex_otel::Timer>,
codex-rs/core/src/tasks/admission.rs
+2 -2
View File
@@ -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<Self>,
context: &Arc<TurnContext>,
turn_state: &Mutex<TurnState>,
codex-rs/core/src/tasks/mod.rs
+7 -6
View File
@@ -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