openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-14 21:40:39 +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/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?