feat: add post-compaction sub-agent infos (#12774)

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
jif-oai
2026-02-26 18:55:34 +00:00
committed by GitHub
Unverified
parent eb77db2957
commit 3404ecff15
10 changed files with 230 additions and 22 deletions
+35
View File
@@ -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
+8 -1
View File
@@ -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);
-1
View File
@@ -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<ResponseItem> = history
.raw_items()
+74 -16
View File
@@ -14,6 +14,7 @@ pub(crate) struct EnvironmentContext {
pub cwd: Option<PathBuf>,
pub shell: Shell,
pub network: Option<NetworkContext>,
pub subagents: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
@@ -23,11 +24,17 @@ pub(crate) struct NetworkContext {
}
impl EnvironmentContext {
pub fn new(cwd: Option<PathBuf>, shell: Shell, network: Option<NetworkContext>) -> Self {
pub fn new(
cwd: Option<PathBuf>,
shell: Shell,
network: Option<NetworkContext>,
subagents: Option<String>,
) -> 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<NetworkContext> {
let network = turn_context
.config
@@ -142,6 +159,11 @@ impl EnvironmentContext {
// lines.push(" <network enabled=\"false\" />".to_string());
}
}
if let Some(subagents) = self.subagents {
lines.push(" <subagents>".to_string());
lines.extend(subagents.lines().map(|line| format!(" {line}")));
lines.push(" </subagents>".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#"<environment_context>
@@ -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#"<environment_context>
@@ -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#"<environment_context>
<shell>bash</shell>
@@ -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#"<environment_context>
<shell>bash</shell>
@@ -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#"<environment_context>
<shell>bash</shell>
@@ -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#"<environment_context>
<shell>bash</shell>
@@ -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#"<environment_context>
<cwd>{}</cwd>
<shell>bash</shell>
<subagents>
- agent-1: atlas
- agent-2
</subagents>
</environment_context>"#,
test_path_buf("/repo").display()
);
assert_eq!(context.serialize_to_xml(), expected);
}
}
+7
View File
@@ -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}"),
}
}
+5 -1
View File
@@ -252,7 +252,7 @@ impl ThreadManager {
}
pub async fn list_thread_ids(&self) -> Vec<ThreadId> {
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<ThreadId> {
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<Arc<CodexThread>> {
let threads = self.threads.read().await;