Skip to content
Draft
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
2 changes: 2 additions & 0 deletions codex-rs/core/src/agent/control.rs
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,8 @@ pub(crate) struct SpawnAgentOptions {
pub(crate) fork_mode: Option<SpawnAgentForkMode>,
pub(crate) parent_thread_id: Option<ThreadId>,
pub(crate) environments: Option<Vec<TurnEnvironmentSelection>>,
/// Explicit child environment subset used to filter fork-inherited capability roots.
pub(crate) restricted_environment_ids: Option<Vec<String>>,
}

#[derive(Clone, Debug)]
Expand Down
9 changes: 8 additions & 1 deletion codex-rs/core/src/agent/control/spawn.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use super::residency::is_v2_resident_session_source;
use super::*;
use codex_extension_api::ExtensionDataInit;
use codex_protocol::capabilities::CapabilityRootLocation;

const AGENT_NAMES: &str = include_str!("../agent_names.txt");

Expand Down Expand Up @@ -480,7 +481,7 @@ impl AgentControl {
))
})?;

let selected_capability_roots = parent_history
let mut selected_capability_roots = parent_history
.items
.iter()
.find_map(|item| {
Expand All @@ -490,6 +491,12 @@ impl AgentControl {
Some(meta_line.meta.selected_capability_roots.clone())
})
.unwrap_or_default();
if let Some(restricted_environment_ids) = options.restricted_environment_ids.as_ref() {
selected_capability_roots.retain(|root| {
let CapabilityRootLocation::Environment { environment_id, .. } = &root.location;
restricted_environment_ids.contains(environment_id)
});
}
let mut forked_rollout_items = parent_history.items;
if let SpawnAgentForkMode::LastNTurns(last_n_turns) = fork_mode {
forked_rollout_items =
Expand Down
90 changes: 90 additions & 0 deletions codex-rs/core/src/agent/control_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1599,6 +1599,96 @@ async fn spawn_agent_fork_last_n_turns_drops_parent_startup_prefix_when_under_li
.expect("parent shutdown should submit");
}

#[tokio::test]
async fn spawn_agent_restricted_fork_filters_selected_capability_roots() {
let harness = AgentControlHarness::new().await;
let allowed_root = SelectedCapabilityRoot {
id: "local@1".to_string(),
location: CapabilityRootLocation::Environment {
environment_id: "local".to_string(),
path: PathUri::parse("file:///local").expect("local root URI"),
},
};
let omitted_root = SelectedCapabilityRoot {
id: "remote@1".to_string(),
location: CapabilityRootLocation::Environment {
environment_id: "remote".to_string(),
path: PathUri::parse("file:///remote").expect("remote root URI"),
},
};
let mut thread_extension_init = ExtensionDataInit::new();
thread_extension_init.insert(vec![allowed_root.clone(), omitted_root]);
let parent = harness
.manager
.start_thread_with_options(StartThreadOptions {
config: harness.config.clone(),
allow_provider_model_fallback: false,
initial_history: InitialHistory::New,
history_mode: None,
session_source: None,
thread_source: None,
dynamic_tools: Vec::new(),
metrics_service_name: None,
parent_trace: None,
environments: harness
.manager
.default_environment_selections(&harness.config.cwd),
thread_extension_init,
supports_openai_form_elicitation: false,
})
.await
.expect("start parent thread");
let parent_thread_id = parent.thread_id;
let parent_thread = parent.thread;

let child_thread_id = harness
.control
.spawn_agent_with_metadata(
harness.config.clone(),
text_input("child task"),
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id,
depth: 1,
agent_path: None,
agent_nickname: None,
agent_role: None,
})),
SpawnAgentOptions {
fork_parent_spawn_call_id: Some("spawn-call-restricted-roots".to_string()),
fork_mode: Some(SpawnAgentForkMode::FullHistory),
restricted_environment_ids: Some(vec!["local".to_string()]),
..Default::default()
},
)
.await
.expect("restricted forked spawn should succeed")
.thread_id;

let child_thread = harness
.manager
.get_thread(child_thread_id)
.await
.expect("child thread should be registered");
assert_eq!(
child_thread
.codex
.session
.services
.selected_capability_roots,
vec![allowed_root]
);

let _ = harness
.control
.shutdown_live_agent(child_thread_id)
.await
.expect("child shutdown should submit");
let _ = parent_thread
.submit(Op::Shutdown {})
.await
.expect("parent shutdown should submit");
}

#[tokio::test]
async fn spawn_agent_fork_last_n_turns_strips_parent_usage_hints() {
let harness = AgentControlHarness::new().await;
Expand Down
6 changes: 5 additions & 1 deletion codex-rs/core/src/tools/handlers/multi_agents/spawn.rs
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,8 @@ async fn handle_spawn_agent(
)
.await?;
apply_spawn_agent_runtime_overrides(&mut config, turn.as_ref())?;
let environments =
spawn_agent_environment_selections(turn.as_ref(), args.environment_ids.as_deref())?;

let result = Box::pin(session.services.agent_control.spawn_agent_with_metadata(
config,
Expand All @@ -133,7 +135,8 @@ async fn handle_spawn_agent(
fork_parent_spawn_call_id: args.fork_context.then(|| call_id.clone()),
fork_mode: args.fork_context.then_some(SpawnAgentForkMode::FullHistory),
parent_thread_id: Some(session.thread_id),
environments: Some(turn.environments.to_selections()),
environments: Some(environments),
restricted_environment_ids: args.environment_ids.clone(),
},
))
.await
Expand Down Expand Up @@ -236,6 +239,7 @@ struct SpawnAgentArgs {
model: Option<String>,
reasoning_effort: Option<ReasoningEffort>,
service_tier: Option<String>,
environment_ids: Option<Vec<String>>,
#[serde(default)]
fork_context: bool,
}
Expand Down
38 changes: 38 additions & 0 deletions codex-rs/core/src/tools/handlers/multi_agents_common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,11 @@ use codex_protocol::openai_models::ReasoningEffort;
use codex_protocol::openai_models::ReasoningEffortPreset;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use codex_protocol::protocol::TurnEnvironmentSelection;
use codex_protocol::user_input::UserInput;
use serde::Serialize;
use serde_json::Value as JsonValue;
use std::collections::HashSet;

/// Minimum wait timeout to prevent tight polling loops from burning CPU.
pub(crate) const MIN_WAIT_TIMEOUT_MS: i64 = DEFAULT_MULTI_AGENT_V2_MIN_WAIT_TIMEOUT_MS;
Expand Down Expand Up @@ -151,6 +153,42 @@ pub(crate) fn parse_collab_input(
}
}

/// Returns the initial environment selections a spawned agent should inherit.
///
/// An explicit subset may only narrow the parent's ready turn environments. Keep the parent's
/// selection order so narrowing a child does not implicitly reorder its primary environment.
pub(crate) fn spawn_agent_environment_selections(
turn: &TurnContext,
environment_ids: Option<&[String]>,
) -> Result<Vec<TurnEnvironmentSelection>, FunctionCallError> {
let parent_selections = turn.environments.to_selections();
let Some(environment_ids) = environment_ids else {
return Ok(parent_selections);
};

let mut requested_ids = HashSet::with_capacity(environment_ids.len());
for environment_id in environment_ids {
if !requested_ids.insert(environment_id.as_str()) {
return Err(FunctionCallError::RespondToModel(format!(
"duplicate spawn environment id `{environment_id}`"
)));
}
if !parent_selections
.iter()
.any(|selection| selection.environment_id.as_str() == environment_id.as_str())
{
return Err(FunctionCallError::RespondToModel(format!(
"spawn environment id `{environment_id}` is not a ready parent-turn environment"
)));
}
}

Ok(parent_selections
.into_iter()
.filter(|selection| requested_ids.contains(selection.environment_id.as_str()))
.collect())
}

/// Builds the base config snapshot for a newly spawned sub-agent.
///
/// The returned config starts from the parent's effective config and then refreshes the
Expand Down
17 changes: 16 additions & 1 deletion codex-rs/core/src/tools/handlers/multi_agents_spec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ const SPAWN_AGENT_MODEL_OVERRIDE_DESCRIPTION: &str =
"Model override for the new agent. Omit unless an explicit override is needed.";
const SPAWN_AGENT_SERVICE_TIER_OVERRIDE_DESCRIPTION: &str =
"Service tier override for the new agent. Omit unless explicitly requested.";
const SPAWN_AGENT_ENVIRONMENT_IDS_DESCRIPTION: &str = "Optional subset of ready environment ids attached to the parent turn to attach when starting the child. Omit to inherit all ready parent environments; pass an empty list to attach none. The child keeps matching environments in the parent's order.";
const MAX_MODEL_OVERRIDES_IN_SPAWN_AGENT_DESCRIPTION: usize = 5;
const MAX_REASONING_EFFORT_CHARS_IN_SPAWN_AGENT_DESCRIPTION: usize = 64;

Expand Down Expand Up @@ -589,6 +590,13 @@ fn spawn_agent_common_properties_v1(agent_type_description: &str) -> BTreeMap<St
SPAWN_AGENT_SERVICE_TIER_OVERRIDE_DESCRIPTION.to_string(),
)),
),
(
"environment_ids".to_string(),
JsonSchema::array(
JsonSchema::string(None),
Some(SPAWN_AGENT_ENVIRONMENT_IDS_DESCRIPTION.to_string()),
),
),
])
}

Expand Down Expand Up @@ -631,6 +639,13 @@ fn spawn_agent_common_properties_v2(agent_type_description: &str) -> BTreeMap<St
SPAWN_AGENT_SERVICE_TIER_OVERRIDE_DESCRIPTION.to_string(),
)),
),
(
"environment_ids".to_string(),
JsonSchema::array(
JsonSchema::string(None),
Some(SPAWN_AGENT_ENVIRONMENT_IDS_DESCRIPTION.to_string()),
),
),
])
}

Expand Down Expand Up @@ -721,7 +736,7 @@ fn spawn_agent_tool_description_v2(
{agent_role_guidance}
Spawns an agent to work on the specified task. If your current task is `/root/task1` and you spawn_agent with task_name "task_3" the agent will have canonical task name `/root/task1/task_3`.
You are then able to refer to this agent as `task_3` or `/root/task1/task_3` interchangeably. However an agent `/root/task2/task_3` would only be able to communicate with this agent via its canonical name `/root/task1/task_3`.
The spawned agent will have the same tools as you and the ability to spawn its own subagents.
The spawned agent will have the same tools as you and the ability to spawn its own subagents, unless you choose a subset of its initially attached execution environments with `environment_ids`.
{inherited_model_guidance}
Only call this tool for a concrete, bounded subtask that can run independently alongside useful local work; otherwise continue locally.
It will be able to send you and other running agents messages, and its final answer will be provided to you when it finishes.
Expand Down
2 changes: 2 additions & 0 deletions codex-rs/core/src/tools/handlers/multi_agents_spec_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,7 @@ fn spawn_agent_tool_v2_requires_task_name_and_lists_visible_models() {
Some(true)
);
assert!(properties.contains_key("fork_turns"));
assert!(properties.contains_key("environment_ids"));
assert!(!properties.contains_key("items"));
assert!(!properties.contains_key("fork_context"));
assert_eq!(
Expand Down Expand Up @@ -142,6 +143,7 @@ fn spawn_agent_tool_v1_keeps_legacy_fork_context_field() {
.expect("spawn_agent should use object params");

assert!(properties.contains_key("fork_context"));
assert!(properties.contains_key("environment_ids"));
assert!(!properties.contains_key("fork_turns"));
assert_eq!(
properties
Expand Down
87 changes: 87 additions & 0 deletions codex-rs/core/src/tools/handlers/multi_agents_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,18 @@ fn thread_manager() -> ThreadManager {
)
}

fn with_windows_environment(mut turn: TurnContext) -> TurnContext {
let mut windows = turn
.environments
.turn_environments
.first()
.expect("test turn should have a local environment")
.clone();
windows.environment_id = "windows".to_string();
turn.environments.turn_environments.push(windows);
turn
}

async fn install_role_with_model_override(turn: &mut TurnContext) -> String {
let role_name = "fork-context-role".to_string();
tokio::fs::create_dir_all(&turn.config.codex_home)
Expand Down Expand Up @@ -171,6 +183,81 @@ where
}
}

#[tokio::test]
async fn spawn_agent_environment_ids_omitted_inherits_parent_environments() {
let (_session, turn) = make_session_and_context().await;
let turn = with_windows_environment(turn);

let selections =
spawn_agent_environment_selections(&turn, None).expect("omitted ids should inherit");

assert_eq!(selections, turn.environments.to_selections());
}

#[tokio::test]
async fn spawn_agent_environment_ids_restrict_child_to_requested_environment() {
let (_session, turn) = make_session_and_context().await;
let turn = with_windows_environment(turn);
let environment_ids = vec!["windows".to_string()];
let expected = vec![
turn.environments
.turn_environments
.iter()
.find(|environment| environment.environment_id == "windows")
.expect("windows environment should exist")
.selection(),
];

let selections = spawn_agent_environment_selections(&turn, Some(&environment_ids))
.expect("attached environment id should be accepted");

assert_eq!(selections, expected);
}

#[tokio::test]
async fn spawn_agent_environment_ids_empty_attaches_no_environments() {
let (_session, turn) = make_session_and_context().await;
let turn = with_windows_environment(turn);
let environment_ids = Vec::new();

let selections = spawn_agent_environment_selections(&turn, Some(&environment_ids))
.expect("empty environment ids should be accepted");

assert_eq!(selections, Vec::new());
}

#[tokio::test]
async fn spawn_agent_environment_ids_rejects_unknown_environment() {
let (_session, turn) = make_session_and_context().await;
let turn = with_windows_environment(turn);
let environment_ids = vec!["missing".to_string()];

let err = spawn_agent_environment_selections(&turn, Some(&environment_ids))
.expect_err("unknown environment id should be rejected");

assert_eq!(
err,
FunctionCallError::RespondToModel(
"spawn environment id `missing` is not a ready parent-turn environment".to_string()
)
);
}

#[tokio::test]
async fn spawn_agent_environment_ids_rejects_duplicate_environment() {
let (_session, turn) = make_session_and_context().await;
let turn = with_windows_environment(turn);
let environment_ids = vec!["windows".to_string(), "windows".to_string()];

let err = spawn_agent_environment_selections(&turn, Some(&environment_ids))
.expect_err("duplicate environment id should be rejected");

assert_eq!(
err,
FunctionCallError::RespondToModel("duplicate spawn environment id `windows`".to_string())
);
}

#[derive(Debug, Deserialize)]
struct ListAgentsResult {
agents: Vec<ListedAgentResult>,
Expand Down
Loading
Loading