Track multi-agent mode in world state (#34845)

## Why

Multi-agent mode instructions are durable model context and need to survive
history changes without re-emitting unrelated setup hints.

## What changed

- Add a `multi_agent_mode` world-state section that diffs, retains, and restores
  the effective mode instructions.
- Persist mode changes in world-state snapshots while keeping the legacy
  `TurnContextItem.multi_agent_mode` field for reading older rollouts.
- Keep initial mode instructions after the root-agent usage hint so the active
  mode takes precedence, and bound custom mode text before snapshotting it.

## Testing

- Add snapshot coverage for mode transitions, retained history, and custom-text
  truncation.
- Verify live mode changes preserve instruction ordering, avoid duplicating the
  usage hint, and record both modes in world state.

GitOrigin-RevId: 96b9964d6e5299b07a011c559b68e3d948a49226
This commit is contained in:
sayan-oai
2026-07-22 23:57:24 +00:00
committed by copyberry
parent 79500d3cc1
commit 0da13c6c99
12 changed files with 297 additions and 123 deletions
-1
View File
@@ -58,7 +58,6 @@ pub(crate) use legacy_apply_patch_exec_command_warning::LegacyApplyPatchExecComm
pub(crate) use legacy_model_mismatch_warning::LegacyModelMismatchWarning;
pub(crate) use legacy_unified_exec_process_limit_warning::LegacyUnifiedExecProcessLimitWarning;
pub(crate) use model_switch_instructions::ModelSwitchInstructions;
pub(crate) use multi_agent_mode_instructions::MultiAgentModeInstructions;
pub(crate) use network_rule_saved::NetworkRuleSaved;
pub use permissions_instructions::ApprovalPromptContext;
pub use permissions_instructions::PermissionsInstructions;
@@ -7,12 +7,12 @@ const EXPLICIT_REQUEST_ONLY_MULTI_AGENT_MODE_TEXT: &str = "Any earlier instructi
const PROACTIVE_MULTI_AGENT_MODE_TEXT: &str = "Proactive multi-agent delegation is active. Any earlier instruction requiring an explicit user request before spawning sub-agents no longer applies. Use sub-agents when parallel work would materially improve speed or quality. This mode remains active until a later multi-agent mode developer message changes it.";
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct MultiAgentModeInstructions {
pub(super) struct MultiAgentModeInstructions {
multi_agent_mode: MultiAgentMode,
}
impl MultiAgentModeInstructions {
pub(crate) fn from_mode(multi_agent_mode: MultiAgentMode) -> Option<Self> {
pub(super) fn from_mode(multi_agent_mode: MultiAgentMode) -> Option<Self> {
if matches!(
&multi_agent_mode,
MultiAgentMode::Custom(hint_text) if hint_text.is_empty()
@@ -3,6 +3,7 @@ mod apps_instructions;
mod collaboration_mode;
mod environment;
mod environments_instructions;
mod multi_agent_mode;
mod permissions;
mod plugins_instructions;
mod realtime;
@@ -30,6 +31,7 @@ pub(crate) use apps_instructions::AppsInstructionsState;
pub(crate) use collaboration_mode::CollaborationModeState;
pub(crate) use environment::EnvironmentsState;
pub(crate) use environments_instructions::EnvironmentsInstructionsState;
pub(crate) use multi_agent_mode::MultiAgentModeState;
pub(crate) use permissions::PermissionsState;
pub(crate) use plugins_instructions::PluginsInstructionsState;
pub(crate) use realtime::RealtimeState;
@@ -0,0 +1,80 @@
use super::PreviousSectionState;
use super::WorldStateSection;
use crate::context::ContextualUserFragment;
use crate::context::multi_agent_mode_instructions::MultiAgentModeInstructions;
use codex_protocol::config_types::MultiAgentMode;
use codex_utils_output_truncation::TruncationPolicy;
use codex_utils_output_truncation::truncate_text;
use serde::Deserialize;
use serde::Serialize;
const MULTI_AGENT_MODE_MAX_TOKENS: usize = 400;
/// Effective multi-agent mode currently visible to the model.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub(crate) struct MultiAgentModeState {
mode: Option<MultiAgentMode>,
}
impl MultiAgentModeState {
pub(crate) fn new(mode: Option<MultiAgentMode>) -> Self {
Self {
mode: mode.map(|mode| match mode {
MultiAgentMode::Custom(hint_text) => MultiAgentMode::Custom(truncate_text(
&hint_text,
TruncationPolicy::Tokens(MULTI_AGENT_MODE_MAX_TOKENS),
)),
mode @ (MultiAgentMode::ExplicitRequestOnly | MultiAgentMode::Proactive) => mode,
}),
}
}
}
impl WorldStateSection for MultiAgentModeState {
const ID: &'static str = "multi_agent_mode";
type Snapshot = Self;
fn snapshot(&self) -> Self::Snapshot {
self.clone()
}
fn matches_legacy_fragment(role: &str, text: &str) -> bool {
role == "developer" && MultiAgentModeInstructions::matches_text(text)
}
fn has_retained_fragment_matcher() -> bool {
true
}
fn matches_retained_fragment(role: &str, text: &str) -> bool {
Self::matches_legacy_fragment(role, text)
}
fn render_diff(
&self,
previous: PreviousSectionState<'_, Self::Snapshot>,
) -> Option<Box<dyn ContextualUserFragment>> {
let mode = match (&self.mode, previous) {
(Some(mode), PreviousSectionState::Known(previous))
if previous.mode.as_ref() == Some(mode) =>
{
return None;
}
(Some(mode), _) => mode.clone(),
(None, PreviousSectionState::Known(previous))
if previous.mode == Some(MultiAgentMode::Proactive) =>
{
MultiAgentMode::ExplicitRequestOnly
}
(None, PreviousSectionState::Unknown) => MultiAgentMode::ExplicitRequestOnly,
(None, PreviousSectionState::Absent | PreviousSectionState::Known(_)) => return None,
};
MultiAgentModeInstructions::from_mode(mode)
.map(|instructions| Box::new(instructions) as Box<dyn ContextualUserFragment>)
}
}
#[cfg(test)]
#[path = "multi_agent_mode_tests.rs"]
mod tests;
@@ -0,0 +1,81 @@
use super::super::test_support::render_section_cases;
use super::*;
use crate::context::world_state::WorldState;
use codex_protocol::models::ResponseItem;
use codex_utils_output_truncation::approx_token_count;
fn state(mode: Option<MultiAgentMode>) -> MultiAgentModeState {
MultiAgentModeState::new(mode)
}
#[test]
fn snapshots() {
use PreviousSectionState::Absent;
use PreviousSectionState::Known;
use PreviousSectionState::Unknown;
let inactive = state(/*mode*/ None);
let explicit = state(Some(MultiAgentMode::ExplicitRequestOnly));
let proactive = state(Some(MultiAgentMode::Proactive));
let custom = state(Some(MultiAgentMode::Custom(
"use a custom policy".to_string(),
)));
let empty = state(Some(MultiAgentMode::Custom(String::new())));
insta::assert_snapshot!(render_section_cases(&[
(Absent, Absent),
(Absent, Known(&inactive)),
(Absent, Known(&explicit)),
(Known(&explicit), Known(&explicit)),
(Known(&explicit), Known(&proactive)),
(Known(&proactive), Known(&inactive)),
(Known(&explicit), Known(&inactive)),
(Known(&explicit), Known(&custom)),
(Known(&custom), Known(&empty)),
(Unknown, Known(&explicit)),
(Unknown, Known(&inactive)),
]));
}
#[test]
fn persisted_mode_is_restored_only_when_missing_from_history() {
let state = state(Some(MultiAgentMode::ExplicitRequestOnly));
let retained: ResponseItem = ContextualUserFragment::into(
MultiAgentModeInstructions::from_mode(MultiAgentMode::ExplicitRequestOnly)
.expect("explicit mode should render"),
);
let mut world_state = WorldState::default();
world_state.add_section(state);
let snapshot = world_state.snapshot();
assert_eq!(
world_state
.render_history_diff(/*previous*/ None, std::slice::from_ref(&retained))
.len(),
1,
);
assert_eq!(
world_state.render_history_diff(Some(&snapshot), &[]).len(),
1
);
assert!(
world_state
.render_history_diff(Some(&snapshot), &[retained])
.is_empty()
);
}
#[test]
fn custom_mode_is_bounded_before_snapshot_and_rendering() {
let state = state(Some(MultiAgentMode::Custom("custom mode ".repeat(1_000))));
let Some(MultiAgentMode::Custom(snapshot_mode)) = state.snapshot().mode else {
panic!("expected custom multi-agent mode")
};
assert!(approx_token_count(&snapshot_mode) < 1_000);
let rendered = state
.render_diff(PreviousSectionState::Absent)
.expect("custom mode should render")
.render();
assert!(approx_token_count(&rendered) < 1_000);
}
@@ -0,0 +1,36 @@
---
source: core/src/context/world_state/multi_agent_mode_tests.rs
expression: "render_section_cases(&[(Absent, Absent), (Absent, Known(&inactive)),\n(Absent, Known(&explicit)), (Known(&explicit), Known(&explicit)),\n(Known(&explicit), Known(&proactive)), (Known(&proactive), Known(&inactive)),\n(Known(&explicit), Known(&inactive)), (Known(&explicit), Known(&custom)),\n(Known(&custom), Known(&empty)), (Unknown, Known(&explicit)),\n(Unknown, Known(&inactive)),])"
---
Absent -> Absent
None
Absent -> {}
None
Absent -> {"mode":"explicitRequestOnly"} (role - developer)
<multi_agent_mode>Any earlier instruction enabling proactive multi-agent delegation no longer applies. Do not spawn sub-agents unless the user or applicable AGENTS.md/skill instructions explicitly ask for sub-agents, delegation, or parallel agent work.</multi_agent_mode>
{"mode":"explicitRequestOnly"} -> {"mode":"explicitRequestOnly"}
None
{"mode":"explicitRequestOnly"} -> {"mode":"proactive"} (role - developer)
<multi_agent_mode>Proactive multi-agent delegation is active. Any earlier instruction requiring an explicit user request before spawning sub-agents no longer applies. Use sub-agents when parallel work would materially improve speed or quality. This mode remains active until a later multi-agent mode developer message changes it.</multi_agent_mode>
{"mode":"proactive"} -> {} (role - developer)
<multi_agent_mode>Any earlier instruction enabling proactive multi-agent delegation no longer applies. Do not spawn sub-agents unless the user or applicable AGENTS.md/skill instructions explicitly ask for sub-agents, delegation, or parallel agent work.</multi_agent_mode>
{"mode":"explicitRequestOnly"} -> {}
None
{"mode":"explicitRequestOnly"} -> {"mode":{"custom":"use a custom policy"}} (role - developer)
<multi_agent_mode>use a custom policy</multi_agent_mode>
{"mode":{"custom":"use a custom policy"}} -> {"mode":{"custom":""}}
None
Unknown -> {"mode":"explicitRequestOnly"} (role - developer)
<multi_agent_mode>Any earlier instruction enabling proactive multi-agent delegation no longer applies. Do not spawn sub-agents unless the user or applicable AGENTS.md/skill instructions explicitly ask for sub-agents, delegation, or parallel agent work.</multi_agent_mode>
Unknown -> {} (role - developer)
<multi_agent_mode>Any earlier instruction enabling proactive multi-agent delegation no longer applies. Do not spawn sub-agents unless the user or applicable AGENTS.md/skill instructions explicitly ask for sub-agents, delegation, or parallel agent work.</multi_agent_mode>
@@ -1,37 +1,14 @@
use crate::context::ContextualUserFragment;
use crate::context::ModelSwitchInstructions;
use crate::context::MultiAgentModeInstructions;
use crate::context::PersonalitySpecInstructions;
use crate::session::PreviousTurnSettings;
use crate::session::turn_context::TurnContext;
use codex_protocol::config_types::MultiAgentMode;
use codex_protocol::config_types::Personality;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::openai_models::ModelInfo;
use codex_protocol::protocol::TurnContextItem;
fn build_multi_agent_mode_update_item(
previous: Option<&TurnContextItem>,
next: &TurnContext,
) -> Option<String> {
let effective_multi_agent_mode = crate::session::multi_agents::effective_multi_agent_mode(next);
let previous = previous?;
if previous.multi_agent_mode == effective_multi_agent_mode {
return None;
}
match effective_multi_agent_mode {
Some(multi_agent_mode) => MultiAgentModeInstructions::from_mode(multi_agent_mode)
.map(|instructions| instructions.render()),
None if previous.multi_agent_mode == Some(MultiAgentMode::Proactive) => {
MultiAgentModeInstructions::from_mode(MultiAgentMode::ExplicitRequestOnly)
.map(|instructions| instructions.render())
}
None => None,
}
}
fn build_personality_update_item(
previous: Option<&TurnContextItem>,
next: &TurnContext,
@@ -145,7 +122,6 @@ pub(crate) fn build_settings_update_items(
// Keep model-switch instructions first so model-specific guidance is read before
// any other context diffs on this turn.
build_model_instructions_update_item(previous_turn_settings, next),
build_multi_agent_mode_update_item(previous, next),
build_personality_update_item(previous, next, personality_feature_enabled),
]
.into_iter()
+8 -5
View File
@@ -26,7 +26,6 @@ use crate::config::resolve_tool_suggest_config_from_layer_stack;
use crate::context::ApprovedCommandPrefixSaved;
use crate::context::AvailableSkillsInstructions;
use crate::context::ContextualUserFragment;
use crate::context::MultiAgentModeInstructions;
use crate::context::NetworkRuleSaved;
use crate::context::PersonalitySpecInstructions;
use crate::context::RecommendedPluginsInstructions;
@@ -120,6 +119,7 @@ use codex_protocol::protocol::HasLegacyEvent;
use codex_protocol::protocol::InterAgentCommunication;
use codex_protocol::protocol::ItemCompletedEvent;
use codex_protocol::protocol::ItemStartedEvent;
use codex_protocol::protocol::MULTI_AGENT_MODE_OPEN_TAG;
use codex_protocol::protocol::MultiAgentVersion;
use codex_protocol::protocol::RawResponseItemEvent;
use codex_protocol::protocol::RolloutItem;
@@ -3437,8 +3437,13 @@ impl Session {
.push(crate::context::ContextWindowGuidance::new(guidance_message).render());
}
}
// Render the active mode after the usage hint so it can override that hint.
let mut initial_multi_agent_mode = None;
for fragment in world_state.render_full() {
match fragment.role() {
"developer" if fragment.markers().0 == MULTI_AGENT_MODE_OPEN_TAG => {
initial_multi_agent_mode = Some(fragment);
}
"developer" => developer_sections.push(fragment.render()),
"user" => contextual_user_sections.push(fragment.render()),
_ => {}
@@ -3469,10 +3474,8 @@ impl Session {
{
items.push(usage_hint_message);
}
if let Some(multi_agent_mode) = multi_agents::effective_multi_agent_mode(turn_context)
&& let Some(instructions) = MultiAgentModeInstructions::from_mode(multi_agent_mode)
{
items.push(ContextualUserFragment::into(instructions));
if let Some(initial_multi_agent_mode) = initial_multi_agent_mode {
items.push(initial_multi_agent_mode.into_boxed_response_item());
}
if let Some(contextual_user_message) =
crate::context_manager::updates::build_contextual_user_message(contextual_user_sections)
+1 -1
View File
@@ -389,7 +389,7 @@ impl TurnContext {
personality: self.personality,
collaboration_mode: Some(self.collaboration_mode()),
multi_agent_version: Some(self.multi_agent_version),
multi_agent_mode: super::multi_agents::effective_multi_agent_mode(self),
multi_agent_mode: None,
realtime_active: Some(self.realtime_active),
effort: self.reasoning_effort.clone(),
summary: ReasoningSummaryConfig::Auto,
+4
View File
@@ -7,6 +7,7 @@ use crate::context::world_state::AppsInstructionsState;
use crate::context::world_state::CollaborationModeState;
use crate::context::world_state::EnvironmentsInstructionsState;
use crate::context::world_state::EnvironmentsState;
use crate::context::world_state::MultiAgentModeState;
use crate::context::world_state::PermissionsState;
use crate::context::world_state::PluginsInstructionsState;
use crate::context::world_state::RealtimeState;
@@ -129,6 +130,9 @@ impl Session {
world_state.add_extension_section(section);
}
}
world_state.add_section(MultiAgentModeState::new(
super::multi_agents::effective_multi_agent_mode(turn_context),
));
world_state
}
}
+82 -23
View File
@@ -25,6 +25,7 @@ use serde_json::json;
const NO_SPAWN_TEXT: &str = "Any earlier instruction enabling proactive multi-agent delegation no longer applies. Do not spawn sub-agents unless the user or applicable AGENTS.md/skill instructions explicitly ask for sub-agents, delegation, or parallel agent work.";
const PROACTIVE_TEXT: &str = "Proactive multi-agent delegation is active.";
const CUSTOM_MODE_HINT_TEXT: &str = "Use the configured delegation policy.";
const ROOT_USAGE_HINT_TEXT: &str = "Root usage hint.";
fn add_ultra_reasoning(model_info: &mut ModelInfo) {
model_info
@@ -149,12 +150,6 @@ async fn configured_mode_hint_uses_custom_mode_across_reasoning_efforts() -> Res
.with_config(configure_custom_mode_hint)
.build(&server)
.await?;
let rollout_path = test
.session_configured
.rollout_path
.clone()
.expect("rollout path");
submit_turn(&test.codex, "explicit", Some(ReasoningEffort::High)).await?;
submit_turn(&test.codex, "proactive", Some(ReasoningEffort::Ultra)).await?;
@@ -172,23 +167,6 @@ async fn configured_mode_hint_uses_custom_mode_across_reasoning_efforts() -> Res
};
assert_eq!(instruction_counts(&first_texts), (1, 0, 0));
assert_eq!(instruction_counts(&second_texts), (1, 0, 0));
let rollout_values = std::fs::read_to_string(rollout_path)?
.lines()
.map(serde_json::from_str::<Value>)
.collect::<serde_json::Result<Vec<_>>>()?;
let recorded_modes = rollout_values
.iter()
.filter(|value| value.get("type").and_then(Value::as_str) == Some("turn_context"))
.filter_map(|value| value.pointer("/payload/multi_agent_mode").cloned())
.collect::<Vec<_>>();
assert_eq!(
recorded_modes,
[
json!({"custom": CUSTOM_MODE_HINT_TEXT}),
json!({"custom": CUSTOM_MODE_HINT_TEXT}),
]
);
Ok(())
}
@@ -281,6 +259,87 @@ async fn changing_configured_mode_hint_to_empty_emits_no_update() -> Result<()>
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn live_mode_change_appends_mode_without_reappending_usage_hint() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let responses = mount_sse_sequence(
&server,
(1..=2)
.map(|index| {
sse(vec![
ev_response_created(&format!("resp-{index}")),
ev_completed(&format!("resp-{index}")),
])
})
.collect(),
)
.await;
let test = test_codex()
.with_model_info_override("gpt-5.4", add_ultra_reasoning)
.with_config(|config| {
configure_ultra(config);
config.multi_agent_v2.root_agent_usage_hint_text =
Some(ROOT_USAGE_HINT_TEXT.to_string());
})
.build_with_auto_env(&server)
.await?;
let rollout_path = test
.session_configured
.rollout_path
.clone()
.expect("rollout path");
submit_turn(&test.codex, "proactive", /*effort*/ None).await?;
submit_turn(&test.codex, "explicit", Some(ReasoningEffort::High)).await?;
let requests = responses.requests();
let first_input = requests[0].input();
let first_texts = developer_texts(&first_input);
let hint_index = first_texts
.iter()
.position(|text| text.contains(ROOT_USAGE_HINT_TEXT))
.expect("initial usage hint");
let mode_index = first_texts
.iter()
.position(|text| text.contains(PROACTIVE_TEXT))
.expect("initial proactive mode");
assert!(hint_index < mode_index);
let second_input = requests[1].input();
let second_texts = developer_texts(&second_input);
assert_eq!(
(
count_containing(&second_texts, ROOT_USAGE_HINT_TEXT),
count_containing(&second_texts, PROACTIVE_TEXT),
count_containing(&second_texts, NO_SPAWN_TEXT),
),
(1, 1, 1),
);
test.codex.ensure_rollout_materialized().await;
test.codex.flush_rollout().await?;
let rollout_values = std::fs::read_to_string(rollout_path)?
.lines()
.map(serde_json::from_str::<Value>)
.collect::<serde_json::Result<Vec<_>>>()?;
let recorded_modes = rollout_values
.iter()
.filter(|value| value.get("type").and_then(Value::as_str) == Some("world_state"))
.filter_map(|value| {
value
.pointer("/payload/state/multi_agent_mode/mode")
.cloned()
})
.collect::<Vec<_>>();
assert_eq!(
recorded_modes,
[json!("proactive"), json!("explicitRequestOnly")]
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn leaving_ultra_after_cold_resume_emits_explicit_mode() -> Result<()> {
skip_if_no_network!(Ok(()));
+1 -67
View File
@@ -2686,28 +2686,6 @@ impl InitialHistory {
}
}
pub fn get_latest_effective_multi_agent_mode(&self) -> Option<MultiAgentMode> {
let items = match self {
InitialHistory::New | InitialHistory::Cleared => return None,
InitialHistory::Resumed(resumed) => &resumed.history,
InitialHistory::Forked(items) => items,
};
items
.iter()
.rev()
.find_map(|item| match item {
RolloutItem::TurnContext(turn_context) => Some(turn_context),
RolloutItem::SessionMeta(_)
| RolloutItem::ResponseItem(_)
| RolloutItem::InterAgentCommunication(_)
| RolloutItem::InterAgentCommunicationMetadata { .. }
| RolloutItem::Compacted(_)
| RolloutItem::WorldState(_)
| RolloutItem::EventMsg(_) => None,
})
.and_then(|turn_context| turn_context.multi_agent_mode.clone())
}
pub fn get_resumed_session_sources(&self) -> Option<(SessionSource, Option<ThreadSource>)> {
let meta = self.get_resumed_session_meta()?;
Some((meta.source.clone(), meta.thread_source.clone()))
@@ -3318,7 +3296,7 @@ pub struct TurnContextItem {
pub collaboration_mode: Option<CollaborationMode>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub multi_agent_version: Option<MultiAgentVersion>,
/// Effective model-visible mode used as the durable context-diff baseline.
/// Legacy effective model-visible mode retained to deserialize older rollouts.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub multi_agent_mode: Option<MultiAgentMode>,
#[serde(default, skip_serializing_if = "Option::is_none")]
@@ -6062,50 +6040,6 @@ mod tests {
Ok(())
}
#[test]
fn latest_effective_multi_agent_mode_uses_latest_turn_context_even_when_unset() -> Result<()> {
let turn_context_item = |multi_agent_mode| -> Result<RolloutItem> {
let mut value = json!({
"cwd": test_path_buf("/tmp"),
"approval_policy": "never",
"sandbox_policy": { "type": "danger-full-access" },
"model": "gpt-5",
"summary": "auto",
});
value["multi_agent_mode"] = serde_json::to_value(multi_agent_mode)?;
Ok(RolloutItem::TurnContext(serde_json::from_value(value)?))
};
assert_eq!(
InitialHistory::Forked(vec![
turn_context_item(Some(MultiAgentMode::Proactive))?,
turn_context_item(/*multi_agent_mode*/ None)?,
])
.get_latest_effective_multi_agent_mode(),
None
);
Ok(())
}
#[test]
fn latest_effective_multi_agent_mode_maps_legacy_none_to_empty_custom() -> Result<()> {
let value = json!({
"cwd": test_path_buf("/tmp"),
"approval_policy": "never",
"sandbox_policy": { "type": "danger-full-access" },
"model": "gpt-5",
"multi_agent_mode": "none",
"summary": "auto",
});
let item = RolloutItem::TurnContext(serde_json::from_value(value)?);
assert_eq!(
InitialHistory::Forked(vec![item]).get_latest_effective_multi_agent_mode(),
Some(MultiAgentMode::Custom(String::new()))
);
Ok(())
}
#[test]
fn turn_context_item_serializes_network_when_present() -> Result<()> {
let item = TurnContextItem {