From c6c09fe29fb0b926ced5146e6ed47b6ff756c1a4 Mon Sep 17 00:00:00 2001 From: jif Date: Tue, 22 Sep 2026 13:39:29 +0000 Subject: [PATCH] Consolidate agent operations in the `AgentControl` implementation (#47292) ## What changed Move `spawn`, `send`, and `inspect` into the `AgentControl` implementation for `LocalAgentControl`, removing the forwarding methods and the `spawn_request` module. Update tool handlers and existing tests to use the trait methods, and remove redundant boxing at spawn call sites. Preserve spawn validation, message attribution, runtime loading, and queue-only delivery semantics. GitOrigin-RevId: 6027d4d97d0dfb6763dc97ef92a2d7ef58de945d --- codex-rs/core/src/agent/control.rs | 1 - codex-rs/core/src/agent/control/api.rs | 116 +++++++++++++++++- codex-rs/core/src/agent/control/delivery.rs | 80 +----------- codex-rs/core/src/agent/control/inspection.rs | 10 -- codex-rs/core/src/agent/control/interrupt.rs | 3 +- .../core/src/agent/control/spawn_request.rs | 57 --------- codex-rs/core/src/agent/control_tests.rs | 1 + .../src/tools/handlers/multi_agents/spawn.rs | 54 ++++---- .../handlers/multi_agents_v2/message_tool.rs | 1 + .../tools/handlers/multi_agents_v2/spawn.rs | 9 +- 10 files changed, 151 insertions(+), 181 deletions(-) delete mode 100644 codex-rs/core/src/agent/control/spawn_request.rs diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 1b9fc35424..ca27e0208d 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -81,7 +81,6 @@ mod resume; mod sender_context; mod service_tier; mod spawn; -mod spawn_request; mod target; mod user_authorization; mod watch; diff --git a/codex-rs/core/src/agent/control/api.rs b/codex-rs/core/src/agent/control/api.rs index 7144877eb9..4f5e9f985f 100644 --- a/codex-rs/core/src/agent/control/api.rs +++ b/codex-rs/core/src/agent/control/api.rs @@ -1,10 +1,12 @@ -//! Adapts the existing local operations to the shared controller interface. +//! Implements shared controller operations using the existing local runtime helpers. //! Runtime loading, message delivery and shared state remain in their existing modules. use super::LocalAgentControl; +use super::spawn::SpawnInitialInput; use crate::agent::api::AgentConfigUpdate; use crate::agent::api::AgentControl; use crate::agent::api::AgentInfo; +use crate::agent::api::AgentInput; use crate::agent::api::AgentTarget; use crate::agent::api::AgentTurnOutcome; use crate::agent::api::DeliveryReceipt; @@ -13,12 +15,17 @@ use crate::agent::api::SpawnRequest; use crate::agent::api::StatusSubscription; use crate::agent::types::AgentExecutionGuard; use crate::agent::types::LiveAgent; +use crate::agent::types::MessageDeliveryMode; +use crate::agent_communication::AgentCommunicationContext; +use crate::agent_communication::AgentCommunicationKind; use crate::codex_thread::GuardianRootSnapshot; use crate::codex_thread::ThreadConfigSnapshot; use crate::config::Config; use crate::rollout_budget::RolloutBudgetReminder; +use codex_protocol::AgentPath; use codex_protocol::SessionId; use codex_protocol::ThreadId; +use codex_protocol::error::CodexErr; use codex_protocol::error::Result; use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::SessionSource; @@ -35,7 +42,42 @@ impl AgentControl for LocalAgentControl { &self, request: SpawnRequest, ) -> BoxFuture<'_, Result<(LiveAgent, ThreadConfigSnapshot)>> { - Box::pin(LocalAgentControl::spawn(self, request)) + Box::pin(async move { + let SpawnRequest { + caller, + config, + input, + source, + options, + } = request; + let input = match input { + AgentInput::UserInput(input) => SpawnInitialInput::UserInput(input), + AgentInput::Message { message, mode } => { + if mode != MessageDeliveryMode::TriggerTurn { + return Err(CodexErr::InvalidRequest( + "spawn input must start the child turn".to_string(), + )); + } + let recipient = source.get_agent_path().ok_or_else(|| { + CodexErr::InvalidRequest( + "spawned agent is missing a canonical task name".to_string(), + ) + })?; + let author = recipient + .as_str() + .rsplit_once('/') + .and_then(|(parent, _)| AgentPath::try_from(parent).ok()) + .ok_or_else(|| { + CodexErr::InvalidRequest("spawn input needs a child path".to_string()) + })?; + SpawnInitialInput::InterAgentCommunication( + message.into_communication(author, recipient, mode), + AgentCommunicationContext::new(AgentCommunicationKind::Spawn, caller), + ) + } + }; + Box::pin(self.spawn_agent_internal(config, input, Some(source), options)).await + }) } fn resume( @@ -52,7 +94,70 @@ impl AgentControl for LocalAgentControl { } fn send(&self, request: SendRequest) -> BoxFuture<'_, Result> { - Box::pin(LocalAgentControl::send(self, request)) + Box::pin(async move { + let SendRequest { + caller, + target, + resume_config, + input, + mut start_options, + } = request; + let target = self.resolve_target(caller, &target)?; + let (metadata, submission_id) = match input { + AgentInput::UserInput(input) => { + let receiver = self.get_agent_metadata(target); + if receiver.is_some() { + self.ensure_v2_agent_loaded(resume_config, target, /*parent*/ None) + .await?; + } + let submission_id = self.send_input(target, input, start_options).await?; + (receiver.unwrap_or_default(), submission_id) + } + AgentInput::Message { message, mode } => { + let receiver = self.ensure_agent_known(target)?; + let author = self + .ensure_agent_known(caller)? + .agent_path + .unwrap_or_else(AgentPath::root); + if mode == MessageDeliveryMode::TriggerTurn + && receiver.agent_path.as_ref().is_some_and(AgentPath::is_root) + { + return Err(CodexErr::UnsupportedOperation( + "Follow-up tasks can't target the root agent".to_string(), + )); + } + let receiver_path = receiver.agent_path.clone().ok_or_else(|| { + CodexErr::UnsupportedOperation( + "target agent is missing an agent_path".to_string(), + ) + })?; + self.ensure_v2_agent_loaded(resume_config, target, /*parent*/ None) + .await?; + let communication = message.into_communication(author, receiver_path, mode); + let kind = match mode { + MessageDeliveryMode::QueueOnly => { + start_options.parent_turn_id = None; + AgentCommunicationKind::Message + } + MessageDeliveryMode::TriggerTurn => AgentCommunicationKind::Followup, + }; + let submission_id = self + .send_inter_agent_communication( + target, + communication, + AgentCommunicationContext::new(kind, caller), + start_options, + ) + .await?; + (receiver, submission_id) + } + }; + Ok(DeliveryReceipt { + thread_id: target, + metadata, + submission_id, + }) + }) } fn interrupt( @@ -82,7 +187,10 @@ impl AgentControl for LocalAgentControl { } fn inspect(&self, caller: ThreadId, target: AgentTarget) -> BoxFuture<'_, Result> { - Box::pin(LocalAgentControl::inspect(self, caller, target)) + Box::pin(async move { + let target = self.resolve_target(caller, &target)?; + self.inspect_agent(target).await + }) } fn list<'a>( diff --git a/codex-rs/core/src/agent/control/delivery.rs b/codex-rs/core/src/agent/control/delivery.rs index 0d84f14cae..fb55ccca93 100644 --- a/codex-rs/core/src/agent/control/delivery.rs +++ b/codex-rs/core/src/agent/control/delivery.rs @@ -1,21 +1,11 @@ -//! Delivers captured agent input without exposing local loading and eviction to callers. -//! -//! Target checks precede reload, and queue-only messages retain their non-waking semantics. +//! Converts agent messages into attributed input while preserving their wake mode. -use super::LocalAgentControl; -use crate::agent::api::AgentInput; -use crate::agent::api::DeliveryReceipt; -use crate::agent::api::SendRequest; use crate::agent::types::AgentMessage; use crate::agent::types::MessageDeliveryMode; -use crate::agent_communication::AgentCommunicationContext; -use crate::agent_communication::AgentCommunicationKind; use crate::context::ContextualUserFragment; use crate::context::InterAgentMessage; use crate::context::InterAgentMessageType; use codex_protocol::AgentPath; -use codex_protocol::error::CodexErr; -use codex_protocol::error::Result as CodexResult; use codex_protocol::protocol::InterAgentCommunication; impl AgentMessage { @@ -51,71 +41,3 @@ impl AgentMessage { } } } - -impl LocalAgentControl { - /// Resolves and delivers captured input, restoring an evicted runtime when necessary. - pub(crate) async fn send(&self, request: SendRequest) -> CodexResult { - let SendRequest { - caller, - target, - resume_config, - input, - mut start_options, - } = request; - let target = self.resolve_target(caller, &target)?; - let (metadata, submission_id) = match input { - AgentInput::UserInput(input) => { - let receiver = self.get_agent_metadata(target); - if receiver.is_some() { - self.ensure_v2_agent_loaded(resume_config, target, /*parent*/ None) - .await?; - } - let submission_id = self.send_input(target, input, start_options).await?; - (receiver.unwrap_or_default(), submission_id) - } - AgentInput::Message { message, mode } => { - let receiver = self.ensure_agent_known(target)?; - let author = self - .ensure_agent_known(caller)? - .agent_path - .unwrap_or_else(AgentPath::root); - if mode == MessageDeliveryMode::TriggerTurn - && receiver.agent_path.as_ref().is_some_and(AgentPath::is_root) - { - return Err(CodexErr::UnsupportedOperation( - "Follow-up tasks can't target the root agent".to_string(), - )); - } - let receiver_path = receiver.agent_path.clone().ok_or_else(|| { - CodexErr::UnsupportedOperation( - "target agent is missing an agent_path".to_string(), - ) - })?; - self.ensure_v2_agent_loaded(resume_config, target, /*parent*/ None) - .await?; - let communication = message.into_communication(author, receiver_path, mode); - let kind = match mode { - MessageDeliveryMode::QueueOnly => { - start_options.parent_turn_id = None; - AgentCommunicationKind::Message - } - MessageDeliveryMode::TriggerTurn => AgentCommunicationKind::Followup, - }; - let submission_id = self - .send_inter_agent_communication( - target, - communication, - AgentCommunicationContext::new(kind, caller), - start_options, - ) - .await?; - (receiver, submission_id) - } - }; - Ok(DeliveryReceipt { - thread_id: target, - metadata, - submission_id, - }) - } -} diff --git a/codex-rs/core/src/agent/control/inspection.rs b/codex-rs/core/src/agent/control/inspection.rs index 4f0a963997..115037166a 100644 --- a/codex-rs/core/src/agent/control/inspection.rs +++ b/codex-rs/core/src/agent/control/inspection.rs @@ -3,22 +3,12 @@ use super::LocalAgentControl; use crate::agent::api::AgentInfo; -use crate::agent::api::AgentTarget; use crate::agent::types::LiveAgent; use codex_protocol::ThreadId; use codex_protocol::error::CodexErrorDetails; use codex_protocol::error::Result as CodexResult; impl LocalAgentControl { - pub(crate) async fn inspect( - &self, - caller: ThreadId, - target: AgentTarget, - ) -> CodexResult { - let target = self.resolve_target(caller, &target)?; - self.inspect_agent(target).await - } - pub(super) async fn inspect_agent(&self, thread_id: ThreadId) -> CodexResult { let manager = self.upgrade()?; let thread = match manager.get_thread(thread_id).await { diff --git a/codex-rs/core/src/agent/control/interrupt.rs b/codex-rs/core/src/agent/control/interrupt.rs index c729e42a65..5eed6a907a 100644 --- a/codex-rs/core/src/agent/control/interrupt.rs +++ b/codex-rs/core/src/agent/control/interrupt.rs @@ -5,7 +5,6 @@ use super::LocalAgentControl; use crate::agent::api::AgentInfo; -use crate::agent::api::AgentTarget; use codex_protocol::AgentPath; use codex_protocol::ThreadId; use codex_protocol::error::CodexErr; @@ -38,7 +37,7 @@ impl LocalAgentControl { receiver_agent.agent_path.as_ref().ok_or_else(|| { CodexErr::UnsupportedOperation("target agent is missing an agent_path".to_string()) })?; - let snapshot = self.inspect(caller, AgentTarget::Id(target)).await?; + let snapshot = self.inspect_agent(target).await?; match self.interrupt_agent(target).await { Ok(_) => {} Err(err) diff --git a/codex-rs/core/src/agent/control/spawn_request.rs b/codex-rs/core/src/agent/control/spawn_request.rs deleted file mode 100644 index 2813f93016..0000000000 --- a/codex-rs/core/src/agent/control/spawn_request.rs +++ /dev/null @@ -1,57 +0,0 @@ -//! Converts captured spawn input into the existing local startup operation. -//! Agent messages keep their canonical paths and spawn attribution. - -use super::LocalAgentControl; -use super::spawn::SpawnInitialInput; -use crate::agent::api::AgentInput; -use crate::agent::api::SpawnRequest; -use crate::agent::types::LiveAgent; -use crate::agent::types::MessageDeliveryMode; -use crate::agent_communication::AgentCommunicationContext; -use crate::agent_communication::AgentCommunicationKind; -use crate::codex_thread::ThreadConfigSnapshot; -use codex_protocol::AgentPath; -use codex_protocol::error::CodexErr; -use codex_protocol::error::Result as CodexResult; - -impl LocalAgentControl { - pub(crate) async fn spawn( - &self, - request: SpawnRequest, - ) -> CodexResult<(LiveAgent, ThreadConfigSnapshot)> { - let SpawnRequest { - caller, - config, - input, - source, - options, - } = request; - let input = match input { - AgentInput::UserInput(input) => SpawnInitialInput::UserInput(input), - AgentInput::Message { message, mode } => { - if mode != MessageDeliveryMode::TriggerTurn { - return Err(CodexErr::InvalidRequest( - "spawn input must start the child turn".to_string(), - )); - } - let recipient = source.get_agent_path().ok_or_else(|| { - CodexErr::InvalidRequest( - "spawned agent is missing a canonical task name".to_string(), - ) - })?; - let author = recipient - .as_str() - .rsplit_once('/') - .and_then(|(parent, _)| AgentPath::try_from(parent).ok()) - .ok_or_else(|| { - CodexErr::InvalidRequest("spawn input needs a child path".to_string()) - })?; - SpawnInitialInput::InterAgentCommunication( - message.into_communication(author, recipient, mode), - AgentCommunicationContext::new(AgentCommunicationKind::Spawn, caller), - ) - } - }; - Box::pin(self.spawn_agent_internal(config, input, Some(source), options)).await - } -} diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index c14218ae14..df4a842ab9 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -3,6 +3,7 @@ use crate::CodexThread; use crate::StateDbHandle; use crate::ThreadManager; use crate::agent::agent_status_from_event; +use crate::agent::api::AgentControl; use crate::agent::api::AgentInfo; use crate::agent::api::AgentInput; use crate::agent::api::AgentTarget; diff --git a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs index 32a024d8ed..daa8720b8c 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/spawn.rs @@ -1,4 +1,5 @@ use super::*; +use crate::agent::api::AgentControl; use crate::agent::api::AgentInput; use crate::agent::api::SpawnRequest; use crate::agent::child_config::SpawnConfigOptions; @@ -108,31 +109,34 @@ async fn handle_spawn_agent( .await .map_err(FunctionCallError::RespondToModel)?; let config = prepared.config; - let result = Box::pin(session.services.agent_control.spawn(SpawnRequest { - caller: session.thread_id, - config, - input: AgentInput::UserInput(input_items), - source: thread_spawn_source( - session.thread_id, - &turn.session_source, - child_depth, - prepared.role_name.as_deref(), - /*task_name*/ None, - )?, - options: SpawnAgentOptions { - 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), - parent_turn_id: Some(turn.sub_id.clone()), - root_turn_id: turn.turn_metadata_state.root_turn_id(), - turn_trigger: turn.turn_metadata_state.current_turn_trigger(), - environments: Some(step_context.environments.to_selections()), - multi_agent_v2_usage_hints: None, - cyber_access_program: turn.cyber_access_program, - }, - })) - .await - .map_err(collab_spawn_error); + let result = session + .services + .agent_control + .spawn(SpawnRequest { + caller: session.thread_id, + config, + input: AgentInput::UserInput(input_items), + source: thread_spawn_source( + session.thread_id, + &turn.session_source, + child_depth, + prepared.role_name.as_deref(), + /*task_name*/ None, + )?, + options: SpawnAgentOptions { + 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), + parent_turn_id: Some(turn.sub_id.clone()), + root_turn_id: turn.turn_metadata_state.root_turn_id(), + turn_trigger: turn.turn_metadata_state.current_turn_trigger(), + environments: Some(step_context.environments.to_selections()), + multi_agent_v2_usage_hints: None, + cyber_access_program: turn.cyber_access_program, + }, + }) + .await + .map_err(collab_spawn_error); let (new_thread_id, status) = match &result { Ok((spawned_agent, _)) => (Some(spawned_agent.thread_id), spawned_agent.status.clone()), Err(_) => (None, AgentStatus::NotFound), diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs index 26d329fa0e..3b8caaef47 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs @@ -6,6 +6,7 @@ use super::analytics::ToolCallAnalytics; use super::*; use crate::TurnStartOptions; +use crate::agent::api::AgentControl; use crate::agent::api::AgentInput; use crate::agent::api::AgentTarget; use crate::agent::api::SendRequest; diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs index bd852d15a6..f48eb89a8b 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs @@ -1,4 +1,5 @@ use super::*; +use crate::agent::api::AgentControl; use crate::agent::api::AgentInput; use crate::agent::api::SpawnRequest; use crate::agent::child_config::SpawnConfigOptions; @@ -181,8 +182,10 @@ async fn handle_spawn_agent( } else { None }; - let (spawned_agent, agent_snapshot) = - Box::pin(session.services.agent_control.spawn(SpawnRequest { + let (spawned_agent, agent_snapshot) = session + .services + .agent_control + .spawn(SpawnRequest { caller: session.thread_id, config, input: AgentInput::Message { @@ -201,7 +204,7 @@ async fn handle_spawn_agent( multi_agent_v2_usage_hints, cyber_access_program: turn.cyber_access_program, }, - })) + }) .await .map_err(collab_spawn_error)?; let new_thread_id = spawned_agent.thread_id;