From 3404ecff153241c688fd34307a4049acf92ad561 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Thu, 26 Feb 2026 18:55:34 +0000 Subject: [PATCH] feat: add post-compaction sub-agent infos (#12774) Co-authored-by: Codex --- codex-rs/core/src/agent/control.rs | 35 ++++++++ codex-rs/core/src/codex.rs | 9 +- codex-rs/core/src/compact_remote.rs | 1 - codex-rs/core/src/environment_context.rs | 90 +++++++++++++++---- codex-rs/core/src/session_prefix.rs | 7 ++ codex-rs/core/src/thread_manager.rs | 6 +- .../core/tests/common/context_snapshot.rs | 47 +++++++++- .../core/tests/suite/model_visible_layout.rs | 45 ++++++++++ ...ronment_context_includes_one_subagent.snap | 6 ++ ...onment_context_includes_two_subagents.snap | 6 ++ 10 files changed, 230 insertions(+), 22 deletions(-) create mode 100644 codex-rs/core/tests/suite/snapshots/all__suite__model_visible_layout__model_visible_layout_environment_context_includes_one_subagent.snap create mode 100644 codex-rs/core/tests/suite/snapshots/all__suite__model_visible_layout__model_visible_layout_environment_context_includes_two_subagents.snap diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index bb1488b76..8ea236228 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -5,6 +5,7 @@ use crate::error::CodexErr; use crate::error::Result as CodexResult; use crate::find_thread_path_by_id_str; use crate::rollout::RolloutRecorder; +use crate::session_prefix::format_subagent_context_line; use crate::session_prefix::format_subagent_notification_message; use crate::state_db; use crate::thread_manager::ThreadManagerState; @@ -343,6 +344,40 @@ impl AgentControl { thread.total_token_usage().await } + pub(crate) async fn format_environment_context_subagents( + &self, + parent_thread_id: ThreadId, + ) -> String { + let Ok(state) = self.upgrade() else { + return String::new(); + }; + + let mut agents = Vec::new(); + for thread_id in state.list_thread_ids().await { + let Ok(thread) = state.get_thread(thread_id).await else { + continue; + }; + let snapshot = thread.config_snapshot().await; + let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: agent_parent_thread_id, + agent_nickname, + .. + }) = snapshot.session_source + else { + continue; + }; + if agent_parent_thread_id != parent_thread_id { + continue; + } + agents.push(format_subagent_context_line( + &thread_id.to_string(), + agent_nickname.as_deref(), + )); + } + agents.sort(); + agents.join("\n") + } + /// Starts a detached watcher for sub-agents spawned from another thread. /// /// This is only enabled for `SubAgentSource::ThreadSpawn`, where a parent thread exists and diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 98a39fc7b..a9d5d60f7 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -3090,8 +3090,15 @@ impl Session { .serialize_to_text(), ); } + let subagents = self + .services + .agent_control + .format_environment_context_subagents(self.conversation_id) + .await; contextual_user_sections.push( - EnvironmentContext::from_turn_context(turn_context, shell.as_ref()).serialize_to_xml(), + EnvironmentContext::from_turn_context(turn_context, shell.as_ref()) + .with_subagents(subagents) + .serialize_to_xml(), ); let mut items = Vec::with_capacity(2); diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index c019a58ce..cc5f5164c 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -105,7 +105,6 @@ async fn run_remote_compact_task_inner_impl( "trimmed history items before remote compaction" ); } - // Required to keep `/undo` available after compaction let ghost_snapshots: Vec = history .raw_items() diff --git a/codex-rs/core/src/environment_context.rs b/codex-rs/core/src/environment_context.rs index 8d8d3c6de..c450818d6 100644 --- a/codex-rs/core/src/environment_context.rs +++ b/codex-rs/core/src/environment_context.rs @@ -14,6 +14,7 @@ pub(crate) struct EnvironmentContext { pub cwd: Option, pub shell: Shell, pub network: Option, + pub subagents: Option, } #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)] @@ -23,11 +24,17 @@ pub(crate) struct NetworkContext { } impl EnvironmentContext { - pub fn new(cwd: Option, shell: Shell, network: Option) -> Self { + pub fn new( + cwd: Option, + shell: Shell, + network: Option, + subagents: Option, + ) -> Self { Self { cwd, shell, network, + subagents, } } @@ -38,9 +45,10 @@ impl EnvironmentContext { let EnvironmentContext { cwd, network, + subagents, shell: _, } = other; - self.cwd == *cwd && self.network == *network + self.cwd == *cwd && self.network == *network && self.subagents == *subagents } pub fn diff_from_turn_context_item( @@ -60,7 +68,7 @@ impl EnvironmentContext { } else { before_network }; - EnvironmentContext::new(cwd, shell.clone(), network) + EnvironmentContext::new(cwd, shell.clone(), network, None) } pub fn from_turn_context(turn_context: &TurnContext, shell: &Shell) -> Self { @@ -68,6 +76,7 @@ impl EnvironmentContext { Some(turn_context.cwd.clone()), shell.clone(), Self::network_from_turn_context(turn_context), + None, ) } @@ -76,9 +85,17 @@ impl EnvironmentContext { Some(turn_context_item.cwd.clone()), shell.clone(), Self::network_from_turn_context_item(turn_context_item), + None, ) } + pub fn with_subagents(mut self, subagents: String) -> Self { + if !subagents.is_empty() { + self.subagents = Some(subagents); + } + self + } + fn network_from_turn_context(turn_context: &TurnContext) -> Option { let network = turn_context .config @@ -142,6 +159,11 @@ impl EnvironmentContext { // lines.push(" ".to_string()); } } + if let Some(subagents) = self.subagents { + lines.push(" ".to_string()); + lines.extend(subagents.lines().map(|line| format!(" {line}"))); + lines.push(" ".to_string()); + } ENVIRONMENT_CONTEXT_FRAGMENT.wrap(lines.join("\n")) } } @@ -171,7 +193,7 @@ mod tests { #[test] fn serialize_workspace_write_environment_context() { let cwd = test_path_buf("/repo"); - let context = EnvironmentContext::new(Some(cwd.clone()), fake_shell(), None); + let context = EnvironmentContext::new(Some(cwd.clone()), fake_shell(), None, None); let expected = format!( r#" @@ -190,8 +212,12 @@ mod tests { allowed_domains: vec!["api.example.com".to_string(), "*.openai.com".to_string()], denied_domains: vec!["blocked.example.com".to_string()], }; - let context = - EnvironmentContext::new(Some(test_path_buf("/repo")), fake_shell(), Some(network)); + let context = EnvironmentContext::new( + Some(test_path_buf("/repo")), + fake_shell(), + Some(network), + None, + ); let expected = format!( r#" @@ -211,7 +237,7 @@ mod tests { #[test] fn serialize_read_only_environment_context() { - let context = EnvironmentContext::new(None, fake_shell(), None); + let context = EnvironmentContext::new(None, fake_shell(), None, None); let expected = r#" bash @@ -222,7 +248,7 @@ mod tests { #[test] fn serialize_external_sandbox_environment_context() { - let context = EnvironmentContext::new(None, fake_shell(), None); + let context = EnvironmentContext::new(None, fake_shell(), None, None); let expected = r#" bash @@ -233,7 +259,7 @@ mod tests { #[test] fn serialize_external_sandbox_with_restricted_network_environment_context() { - let context = EnvironmentContext::new(None, fake_shell(), None); + let context = EnvironmentContext::new(None, fake_shell(), None, None); let expected = r#" bash @@ -244,7 +270,7 @@ mod tests { #[test] fn serialize_full_access_environment_context() { - let context = EnvironmentContext::new(None, fake_shell(), None); + let context = EnvironmentContext::new(None, fake_shell(), None, None); let expected = r#" bash @@ -255,23 +281,29 @@ mod tests { #[test] fn equals_except_shell_compares_cwd() { - let context1 = EnvironmentContext::new(Some(PathBuf::from("/repo")), fake_shell(), None); - let context2 = EnvironmentContext::new(Some(PathBuf::from("/repo")), fake_shell(), None); + let context1 = + EnvironmentContext::new(Some(PathBuf::from("/repo")), fake_shell(), None, None); + let context2 = + EnvironmentContext::new(Some(PathBuf::from("/repo")), fake_shell(), None, None); assert!(context1.equals_except_shell(&context2)); } #[test] fn equals_except_shell_ignores_sandbox_policy() { - let context1 = EnvironmentContext::new(Some(PathBuf::from("/repo")), fake_shell(), None); - let context2 = EnvironmentContext::new(Some(PathBuf::from("/repo")), fake_shell(), None); + let context1 = + EnvironmentContext::new(Some(PathBuf::from("/repo")), fake_shell(), None, None); + let context2 = + EnvironmentContext::new(Some(PathBuf::from("/repo")), fake_shell(), None, None); assert!(context1.equals_except_shell(&context2)); } #[test] fn equals_except_shell_compares_cwd_differences() { - let context1 = EnvironmentContext::new(Some(PathBuf::from("/repo1")), fake_shell(), None); - let context2 = EnvironmentContext::new(Some(PathBuf::from("/repo2")), fake_shell(), None); + let context1 = + EnvironmentContext::new(Some(PathBuf::from("/repo1")), fake_shell(), None, None); + let context2 = + EnvironmentContext::new(Some(PathBuf::from("/repo2")), fake_shell(), None, None); assert!(!context1.equals_except_shell(&context2)); } @@ -286,6 +318,7 @@ mod tests { shell_snapshot: crate::shell::empty_shell_snapshot_receiver(), }, None, + None, ); let context2 = EnvironmentContext::new( Some(PathBuf::from("/repo")), @@ -295,8 +328,33 @@ mod tests { shell_snapshot: crate::shell::empty_shell_snapshot_receiver(), }, None, + None, ); assert!(context1.equals_except_shell(&context2)); } + + #[test] + fn serialize_environment_context_with_subagents() { + let context = EnvironmentContext::new( + Some(test_path_buf("/repo")), + fake_shell(), + None, + Some("- agent-1: atlas\n- agent-2".to_string()), + ); + + let expected = format!( + r#" + {} + bash + + - agent-1: atlas + - agent-2 + +"#, + test_path_buf("/repo").display() + ); + + assert_eq!(context.serialize_to_xml(), expected); + } } diff --git a/codex-rs/core/src/session_prefix.rs b/codex-rs/core/src/session_prefix.rs index ebf068894..db3ac00a6 100644 --- a/codex-rs/core/src/session_prefix.rs +++ b/codex-rs/core/src/session_prefix.rs @@ -12,3 +12,10 @@ pub(crate) fn format_subagent_notification_message(agent_id: &str, status: &Agen .to_string(); SUBAGENT_NOTIFICATION_FRAGMENT.wrap(payload_json) } + +pub(crate) fn format_subagent_context_line(agent_id: &str, agent_nickname: Option<&str>) -> String { + match agent_nickname.filter(|nickname| !nickname.is_empty()) { + Some(agent_nickname) => format!("- {agent_id}: {agent_nickname}"), + None => format!("- {agent_id}"), + } +} diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index ab2e4c2d2..0a56cab1d 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -252,7 +252,7 @@ impl ThreadManager { } pub async fn list_thread_ids(&self) -> Vec { - self.state.threads.read().await.keys().copied().collect() + self.state.list_thread_ids().await } pub async fn refresh_mcp_servers(&self, refresh_config: McpServerRefreshConfig) { @@ -412,6 +412,10 @@ impl ThreadManager { } impl ThreadManagerState { + pub(crate) async fn list_thread_ids(&self) -> Vec { + self.threads.read().await.keys().copied().collect() + } + /// Fetch a thread by ID or return ThreadNotFound. pub(crate) async fn get_thread(&self, thread_id: ThreadId) -> CodexResult> { let threads = self.threads.read().await; diff --git a/codex-rs/core/tests/common/context_snapshot.rs b/codex-rs/core/tests/common/context_snapshot.rs index 24442dd4b..addb1bc31 100644 --- a/codex-rs/core/tests/common/context_snapshot.rs +++ b/codex-rs/core/tests/common/context_snapshot.rs @@ -238,15 +238,34 @@ fn canonicalize_snapshot_text(text: &str) -> String { return "".to_string(); } if text.starts_with("") { + let subagent_count = text + .split_once("") + .and_then(|(_, rest)| rest.split_once("")) + .map(|(subagents, _)| { + subagents + .lines() + .filter(|line| line.trim_start().starts_with("- ")) + .count() + }) + .unwrap_or(0); + let subagents_suffix = if subagent_count > 0 { + format!(":subagents={subagent_count}") + } else { + String::new() + }; if let (Some(cwd_start), Some(cwd_end)) = (text.find(""), text.find("")) { let cwd = &text[cwd_start + "".len()..cwd_end]; return if cwd.ends_with("PRETURN_CONTEXT_DIFF_CWD") { - "".to_string() + format!("") } else { - ">".to_string() + format!("{subagents_suffix}>") }; } - return "".to_string(); + return if subagent_count > 0 { + format!("") + } else { + "".to_string() + }; } if text.starts_with("You are performing a CONTEXT CHECKPOINT COMPACTION.") { return "".to_string(); @@ -308,6 +327,28 @@ mod tests { assert_eq!(rendered, "00:message/user:"); } + #[test] + fn redacted_text_mode_normalizes_environment_context_with_subagents() { + let items = vec![json!({ + "type": "message", + "role": "user", + "content": [{ + "type": "input_text", + "text": "\n /tmp/example\n bash\n \n - agent-1: atlas\n - agent-2\n \n" + }] + })]; + + let rendered = format_response_items_snapshot( + &items, + &ContextSnapshotOptions::default().render_mode(ContextSnapshotRenderMode::RedactedText), + ); + + assert_eq!( + rendered, + "00:message/user::subagents=2>" + ); + } + #[test] fn image_only_message_is_rendered_as_non_text_span() { let items = vec![json!({ diff --git a/codex-rs/core/tests/suite/model_visible_layout.rs b/codex-rs/core/tests/suite/model_visible_layout.rs index 3288eb5c0..c8a5a978f 100644 --- a/codex-rs/core/tests/suite/model_visible_layout.rs +++ b/codex-rs/core/tests/suite/model_visible_layout.rs @@ -26,6 +26,7 @@ use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; +use serde_json::json; const PRETURN_CONTEXT_DIFF_CWD: &str = "PRETURN_CONTEXT_DIFF_CWD"; @@ -53,6 +54,30 @@ fn agents_message_count(request: &ResponsesRequest) -> usize { .count() } +fn format_environment_context_subagents_snapshot(subagents: &[&str]) -> String { + let subagents_block = if subagents.is_empty() { + String::new() + } else { + let lines = subagents + .iter() + .map(|line| format!(" {line}")) + .collect::>() + .join("\n"); + format!("\n \n{lines}\n ") + }; + let items = vec![json!({ + "type": "message", + "role": "user", + "content": [{ + "type": "input_text", + "text": format!( + "\n /tmp/example\n bash{subagents_block}\n" + ), + }], + })]; + context_snapshot::format_response_items_snapshot(items.as_slice(), &context_snapshot_options()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn snapshot_model_visible_layout_turn_overrides() -> Result<()> { skip_if_no_network!(Ok(())); @@ -445,3 +470,23 @@ async fn snapshot_model_visible_layout_resume_override_matches_rollout_model() - Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn snapshot_model_visible_layout_environment_context_includes_one_subagent() -> Result<()> { + insta::assert_snapshot!( + "model_visible_layout_environment_context_includes_one_subagent", + format_environment_context_subagents_snapshot(&["- agent-1: Atlas"]) + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn snapshot_model_visible_layout_environment_context_includes_two_subagents() -> Result<()> { + insta::assert_snapshot!( + "model_visible_layout_environment_context_includes_two_subagents", + format_environment_context_subagents_snapshot(&["- agent-1: Atlas", "- agent-2: Juniper"]) + ); + + Ok(()) +} diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__model_visible_layout__model_visible_layout_environment_context_includes_one_subagent.snap b/codex-rs/core/tests/suite/snapshots/all__suite__model_visible_layout__model_visible_layout_environment_context_includes_one_subagent.snap new file mode 100644 index 000000000..3436943cd --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__model_visible_layout__model_visible_layout_environment_context_includes_one_subagent.snap @@ -0,0 +1,6 @@ +--- +source: core/tests/suite/model_visible_layout.rs +assertion_line: 476 +expression: "format_environment_context_subagents_snapshot(&[\"- agent-1: Atlas\"])" +--- +00:message/user::subagents=1> diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__model_visible_layout__model_visible_layout_environment_context_includes_two_subagents.snap b/codex-rs/core/tests/suite/snapshots/all__suite__model_visible_layout__model_visible_layout_environment_context_includes_two_subagents.snap new file mode 100644 index 000000000..105c28515 --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__model_visible_layout__model_visible_layout_environment_context_includes_two_subagents.snap @@ -0,0 +1,6 @@ +--- +source: core/tests/suite/model_visible_layout.rs +assertion_line: 486 +expression: "format_environment_context_subagents_snapshot(&[\"- agent-1: Atlas\",\n\"- agent-2: Juniper\",])" +--- +00:message/user::subagents=2>