Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 35 additions & 0 deletions codex-rs/core/src/agent/control.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down
9 changes: 8 additions & 1 deletion codex-rs/core/src/codex.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
1 change: 0 additions & 1 deletion codex-rs/core/src/compact_remote.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
90 changes: 74 additions & 16 deletions codex-rs/core/src/environment_context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ pub(crate) struct EnvironmentContext {
pub cwd: Option<PathBuf>,
pub shell: Shell,
pub network: Option<NetworkContext>,
pub subagents: Option<String>,
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should consider adding subagent info as a seprate message. @charley-oai has a PR where all separate messages are getting merged.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Happy to do it once the PR is there

}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
Expand All @@ -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,
}
}

Expand All @@ -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(
Expand All @@ -60,14 +68,15 @@ 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 {
Self::new(
Some(turn_context.cwd.clone()),
shell.clone(),
Self::network_from_turn_context(turn_context),
None,
)
}

Expand All @@ -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
Expand Down Expand Up @@ -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"))
}
}
Expand Down Expand Up @@ -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>
Expand All @@ -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>
Expand All @@ -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>
Expand All @@ -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>
Expand All @@ -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>
Expand All @@ -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>
Expand All @@ -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));
}
Expand All @@ -286,6 +318,7 @@ mod tests {
shell_snapshot: crate::shell::empty_shell_snapshot_receiver(),
},
None,
None,
);
let context2 = EnvironmentContext::new(
Some(PathBuf::from("/repo")),
Expand All @@ -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 changes: 7 additions & 0 deletions codex-rs/core/src/session_prefix.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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}"),
}
}
6 changes: 5 additions & 1 deletion codex-rs/core/src/thread_manager.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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;
Expand Down
47 changes: 44 additions & 3 deletions codex-rs/core/tests/common/context_snapshot.rs
Original file line number Diff line number Diff line change
Expand Up @@ -238,15 +238,34 @@ fn canonicalize_snapshot_text(text: &str) -> String {
return "<AGENTS_MD>".to_string();
}
if text.starts_with("<environment_context>") {
let subagent_count = text
.split_once("<subagents>")
.and_then(|(_, rest)| rest.split_once("</subagents>"))
.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("<cwd>"), text.find("</cwd>")) {
let cwd = &text[cwd_start + "<cwd>".len()..cwd_end];
return if cwd.ends_with("PRETURN_CONTEXT_DIFF_CWD") {
"<ENVIRONMENT_CONTEXT:cwd=PRETURN_CONTEXT_DIFF_CWD>".to_string()
format!("<ENVIRONMENT_CONTEXT:cwd=PRETURN_CONTEXT_DIFF_CWD{subagents_suffix}>")
} else {
"<ENVIRONMENT_CONTEXT:cwd=<CWD>>".to_string()
format!("<ENVIRONMENT_CONTEXT:cwd=<CWD>{subagents_suffix}>")
};
}
return "<ENVIRONMENT_CONTEXT>".to_string();
return if subagent_count > 0 {
format!("<ENVIRONMENT_CONTEXT{subagents_suffix}>")
} else {
"<ENVIRONMENT_CONTEXT>".to_string()
};
}
if text.starts_with("You are performing a CONTEXT CHECKPOINT COMPACTION.") {
return "<SUMMARIZATION_PROMPT>".to_string();
Expand Down Expand Up @@ -308,6 +327,28 @@ mod tests {
assert_eq!(rendered, "00:message/user:<AGENTS_MD>");
}

#[test]
fn redacted_text_mode_normalizes_environment_context_with_subagents() {
let items = vec![json!({
"type": "message",
"role": "user",
"content": [{
"type": "input_text",
"text": "<environment_context>\n <cwd>/tmp/example</cwd>\n <shell>bash</shell>\n <subagents>\n - agent-1: atlas\n - agent-2\n </subagents>\n</environment_context>"
}]
})];

let rendered = format_response_items_snapshot(
&items,
&ContextSnapshotOptions::default().render_mode(ContextSnapshotRenderMode::RedactedText),
);

assert_eq!(
rendered,
"00:message/user:<ENVIRONMENT_CONTEXT:cwd=<CWD>:subagents=2>"
);
}

#[test]
fn image_only_message_is_rendered_as_non_text_span() {
let items = vec![json!({
Expand Down
Loading
Loading