Download codex-rs/core/src/agent/control.rs from SaylorTwift/codex: direct link, hf CLI and curl.
- Browser
- Download file 34.9 kB
-
https://huggingface.co/SaylorTwift/codex/resolve/main/codex-rs/core/src/agent/control.rs
- Command line
-
hf download hf://SaylorTwift/codex/codex-rs/core/src/agent/control.rs
-
curl -L -o control.rs https://huggingface.co/SaylorTwift/codex/resolve/main/codex-rs/core/src/agent/control.rs
34.9 kB
| use crate::TurnInputRequest; | |
| use crate::TurnInputSubmission; | |
| use crate::TurnStartOptions; | |
| use crate::agent::AgentStatus; | |
| use crate::agent::registry::AgentMetadata; | |
| use crate::agent::registry::AgentRegistry; | |
| use crate::agent::role::DEFAULT_ROLE_NAME; | |
| use crate::agent::role::resolve_role_config; | |
| use crate::agent::status::is_final; | |
| use crate::agent_communication::AgentCommunicationContext; | |
| use crate::agent_communication::AgentCommunicationKind; | |
| use crate::codex_thread::ThreadConfigSnapshot; | |
| use crate::config::Config; | |
| use crate::config::RolloutBudgetConfig; | |
| use crate::context::SubagentNotification; | |
| use crate::environment_selection::TurnEnvironmentSnapshot; | |
| use crate::rollout_budget::RolloutBudget; | |
| use crate::session::emit_subagent_session_started; | |
| use crate::session::multi_agents::ResolvedMultiAgentV2UsageHints; | |
| use crate::session_prefix::format_inter_agent_completion_message; | |
| use crate::session_prefix::format_subagent_context_line; | |
| use crate::thread_manager::ResumeThreadWithHistoryOptions; | |
| use crate::thread_manager::ThreadIdGenerator; | |
| use crate::thread_manager::ThreadManagerState; | |
| use crate::thread_manager::default_thread_id_generator; | |
| use crate::thread_rollout_truncation::truncate_rollout_to_last_n_fork_turns; | |
| use crate::turn_timing::now_unix_timestamp_ms; | |
| use arc_swap::ArcSwapOption; | |
| use codex_history::InitialHistory; | |
| use codex_history::ResumedHistory; | |
| use codex_history::RolloutItem; | |
| use codex_protocol::AgentPath; | |
| use codex_protocol::SessionId; | |
| use codex_protocol::ThreadId; | |
| use codex_protocol::error::CodexErr; | |
| use codex_protocol::error::CodexErrorDetails; | |
| use codex_protocol::error::Result as CodexResult; | |
| use codex_protocol::items::SubAgentActivityItem; | |
| use codex_protocol::items::TurnItem; | |
| use codex_protocol::models::ContentItem; | |
| use codex_protocol::models::MessagePhase; | |
| use codex_protocol::models::ResponseItem; | |
| use codex_protocol::protocol::Event; | |
| use codex_protocol::protocol::EventMsg; | |
| use codex_protocol::protocol::HasLegacyEvent; | |
| use codex_protocol::protocol::InterAgentCommunication; | |
| use codex_protocol::protocol::ItemCompletedEvent; | |
| use codex_protocol::protocol::ItemStartedEvent; | |
| use codex_protocol::protocol::MultiAgentVersion; | |
| use codex_protocol::protocol::Op; | |
| use codex_protocol::protocol::SessionSource; | |
| use codex_protocol::protocol::SubAgentSource; | |
| use codex_protocol::protocol::ThreadHistoryMode; | |
| use codex_protocol::protocol::ThreadSource; | |
| use codex_protocol::protocol::TurnEnvironmentSelection; | |
| use codex_protocol::turn_input::CyberAccessProgram; | |
| use codex_protocol::user_input::UserInput; | |
| use codex_thread_store::LoadThreadHistoryParams; | |
| use codex_thread_store::ReadThreadParams; | |
| use serde::Serialize; | |
| use std::collections::HashMap; | |
| use std::collections::HashSet; | |
| use std::collections::VecDeque; | |
| use std::sync::Arc; | |
| use std::sync::Weak; | |
| use tokio::sync::watch; | |
| use tracing::warn; | |
| use uuid::Uuid; | |
| pub(crate) use self::delivery::AgentMessage; | |
| pub(crate) use self::delivery::MessageDeliveryError; | |
| pub(crate) use self::delivery::MessageDeliveryMode; | |
| pub(crate) use self::execution::AgentExecutionGuard; | |
| use self::execution::AgentExecutionLimiter; | |
| pub(crate) use self::interrupt::AgentInterruptError; | |
| pub(crate) use self::interrupt::AgentInterruptOutcome; | |
| use self::residency::V2Residency; | |
| mod delivery; | |
| mod execution; | |
| mod interrupt; | |
| mod legacy; | |
| mod residency; | |
| mod service_tier; | |
| mod spawn; | |
| mod user_authorization; | |
| const MAX_ENVIRONMENT_SUBAGENTS: usize = 8; | |
| const MAX_ENVIRONMENT_SUBAGENT_BYTES: usize = 1_024; | |
| pub(crate) enum SpawnAgentForkMode { | |
| FullHistory, | |
| LastNTurns(usize), | |
| } | |
| pub(crate) struct SpawnAgentOptions { | |
| pub(crate) fork_parent_spawn_call_id: Option<String>, | |
| pub(crate) fork_mode: Option<SpawnAgentForkMode>, | |
| pub(crate) parent_thread_id: Option<ThreadId>, | |
| pub(crate) parent_turn_id: Option<String>, | |
| /// Attribute delegated usage to the turn that initiated it. | |
| pub(crate) turn_trigger: Option<String>, | |
| pub(crate) root_turn_id: Option<String>, | |
| pub(crate) environments: Option<Vec<TurnEnvironmentSelection>>, | |
| pub(crate) multi_agent_v2_usage_hints: Option<ResolvedMultiAgentV2UsageHints>, | |
| pub(crate) cyber_access_program: Option<CyberAccessProgram>, | |
| } | |
| pub(crate) struct LiveAgent { | |
| pub(crate) thread_id: ThreadId, | |
| pub(crate) metadata: AgentMetadata, | |
| pub(crate) status: AgentStatus, | |
| } | |
| pub(crate) struct ListedAgent { | |
| pub(crate) agent_name: String, | |
| pub(crate) agent_status: AgentStatus, | |
| } | |
| /// Control-plane handle for multi-agent operations. | |
| /// `AgentControl` is held by each session (via `SessionServices`). It provides capability to | |
| /// spawn new agents and the inter-agent communication layer. | |
| /// An `AgentControl` instance is intended to be created at most once per root thread/session | |
| /// tree. That same `AgentControl` is then shared with every sub-agent spawned from that root, | |
| /// which keeps the registry scoped to that root thread rather than the entire `ThreadManager`. | |
| pub(crate) struct AgentControl { | |
| /// session_id is equal to the root thread's ID. | |
| session_id: SessionId, | |
| /// Weak handle back to the global thread registry/state. | |
| /// This is `Weak` to avoid reference cycles and shadow persistence of the form | |
| /// `ThreadManagerState -> CodexThread -> Session -> SessionServices -> ThreadManagerState`. | |
| manager: Weak<ThreadManagerState>, | |
| /// Captured at construction so delegates retain their manager's allocation policy. | |
| thread_id_generator: ThreadIdGenerator, | |
| state: Arc<AgentRegistry>, | |
| v2_residency: Arc<V2Residency>, | |
| agent_execution_limiter: Arc<AgentExecutionLimiter>, | |
| /// Session-scoped state shared by the root thread and every cloned sub-agent control handle. | |
| rollout_budget: Arc<RolloutBudget>, | |
| /// The user-selected root routing tier, shared by the entire agent tree. | |
| root_service_tier: Arc<ArcSwapOption<String>>, | |
| } | |
| impl Default for AgentControl { | |
| fn default() -> Self { | |
| Self::new( | |
| Weak::default(), | |
| default_thread_id_generator(), | |
| /*rollout_budget*/ None, | |
| ) | |
| } | |
| } | |
| impl AgentControl { | |
| /// Construct a new `AgentControl` that can spawn/message agents via the given manager state. | |
| pub(crate) fn new( | |
| manager: Weak<ThreadManagerState>, | |
| thread_id_generator: ThreadIdGenerator, | |
| rollout_budget: Option<RolloutBudgetConfig>, | |
| ) -> Self { | |
| let control = Self { | |
| session_id: SessionId::default(), | |
| manager, | |
| thread_id_generator, | |
| state: Arc::default(), | |
| v2_residency: Arc::default(), | |
| agent_execution_limiter: Arc::default(), | |
| rollout_budget: Arc::default(), | |
| root_service_tier: Arc::new(ArcSwapOption::from(None)), | |
| }; | |
| if let Some(rollout_budget) = rollout_budget { | |
| control.rollout_budget.configure(rollout_budget); | |
| } | |
| control | |
| } | |
| pub(crate) fn with_session_id(mut self, session_id: SessionId, max_threads: usize) -> Self { | |
| self.session_id = session_id; | |
| self.agent_execution_limiter.initialize(max_threads); | |
| self | |
| } | |
| pub(crate) fn session_id(&self) -> SessionId { | |
| self.session_id | |
| } | |
| pub(crate) fn generate_thread_id(&self) -> ThreadId { | |
| (self.thread_id_generator)() | |
| } | |
| pub(crate) fn rollout_budget(&self) -> &RolloutBudget { | |
| self.rollout_budget.as_ref() | |
| } | |
| /// Send rich user input items to an existing agent thread. | |
| pub(crate) async fn send_input( | |
| &self, | |
| agent_id: ThreadId, | |
| input: Vec<UserInput>, | |
| start_options: TurnStartOptions, | |
| ) -> CodexResult<String> { | |
| let state = self.upgrade()?; | |
| let thread = state.get_thread(agent_id).await?; | |
| let result = match thread | |
| .start_or_steer_turn(TurnInputRequest::user_input(input).on_start(start_options)) | |
| .await | |
| { | |
| Ok(TurnInputSubmission::Started { turn_id }) => Ok(turn_id), | |
| Ok(TurnInputSubmission::Steered { .. }) => { | |
| // MAv1 exposes an opaque `submission_id` to the model. The legacy | |
| // `Op::UserInput` path returned a fresh ID for every steer, while the | |
| // turn-input API returns the active turn ID. Keep the tool-visible ID | |
| // unique without adding a submission receipt back to Core. | |
| Ok(Uuid::now_v7().to_string()) | |
| } | |
| Ok(TurnInputSubmission::NotSubmitted { reason }) => Err(CodexErr::InvalidRequest( | |
| format!("turn input was not submitted: {reason:?}"), | |
| )), | |
| Err(err) => Err(err), | |
| }; | |
| self.handle_thread_request_result(agent_id, &state, result) | |
| .await | |
| } | |
| pub(crate) async fn send_inter_agent_communication( | |
| &self, | |
| agent_id: ThreadId, | |
| communication: InterAgentCommunication, | |
| agent_communication_context: AgentCommunicationContext, | |
| start_options: TurnStartOptions, | |
| ) -> CodexResult<String> { | |
| let state = self.upgrade()?; | |
| if communication.trigger_turn { | |
| let thread = state.get_thread(agent_id).await?; | |
| self.ensure_execution_capacity_for_turn_start(&thread) | |
| .await?; | |
| } | |
| self.send_inter_agent_communication_after_capacity_check( | |
| agent_id, | |
| &state, | |
| communication, | |
| agent_communication_context, | |
| start_options, | |
| ) | |
| .await | |
| } | |
| pub(crate) async fn emit_sub_agent_activity( | |
| &self, | |
| thread_id: ThreadId, | |
| turn_id: String, | |
| item: SubAgentActivityItem, | |
| ) -> CodexResult<()> { | |
| let state = self.upgrade()?; | |
| let thread = state.get_thread(thread_id).await?; | |
| let started_at_ms = now_unix_timestamp_ms(); | |
| let item = TurnItem::SubAgentActivity(item); | |
| thread | |
| .session | |
| .send_event_raw(Event { | |
| id: turn_id.clone(), | |
| msg: EventMsg::ItemStarted(ItemStartedEvent { | |
| thread_id, | |
| turn_id: turn_id.clone(), | |
| item: item.clone(), | |
| started_at_ms, | |
| }), | |
| }) | |
| .await; | |
| let completed_at_ms = now_unix_timestamp_ms(); | |
| let completed = ItemCompletedEvent { | |
| thread_id, | |
| turn_id: turn_id.clone(), | |
| item, | |
| started_at_ms: Some(started_at_ms), | |
| completed_at_ms, | |
| }; | |
| thread | |
| .session | |
| .send_event_raw(Event { | |
| id: turn_id.clone(), | |
| msg: EventMsg::ItemCompleted(completed.clone()), | |
| }) | |
| .await; | |
| for legacy in completed.as_legacy_events(/*show_raw_agent_reasoning*/ false) { | |
| thread | |
| .session | |
| .send_event_raw(Event { | |
| id: turn_id.clone(), | |
| msg: legacy, | |
| }) | |
| .await; | |
| } | |
| Ok(()) | |
| } | |
| async fn send_inter_agent_communication_after_capacity_check( | |
| &self, | |
| agent_id: ThreadId, | |
| state: &Arc<ThreadManagerState>, | |
| communication: InterAgentCommunication, | |
| context: AgentCommunicationContext, | |
| start_options: TurnStartOptions, | |
| ) -> CodexResult<String> { | |
| self.submit_inter_agent_communication( | |
| agent_id, | |
| state, | |
| communication, | |
| context, | |
| start_options, | |
| ) | |
| .await | |
| } | |
| async fn submit_inter_agent_communication( | |
| &self, | |
| agent_id: ThreadId, | |
| state: &Arc<ThreadManagerState>, | |
| communication: InterAgentCommunication, | |
| context: AgentCommunicationContext, | |
| start_options: TurnStartOptions, | |
| ) -> CodexResult<String> { | |
| let communication_for_log = | |
| crate::agent_communication::logging_enabled().then(|| communication.clone()); | |
| let (parent_turn_id, root_turn_id) = if communication.trigger_turn { | |
| ( | |
| start_options.parent_turn_id.clone(), | |
| start_options.root_turn_id.clone(), | |
| ) | |
| } else { | |
| (None, None) | |
| }; | |
| let result = self | |
| .handle_thread_request_result( | |
| agent_id, | |
| state, | |
| state | |
| .send_op( | |
| agent_id, | |
| Op::InterAgentCommunication { | |
| communication, | |
| start_options, | |
| }, | |
| parent_turn_id, | |
| root_turn_id, | |
| ) | |
| .await, | |
| ) | |
| .await; | |
| if let (Some(communication), Ok(communication_id)) = | |
| (communication_for_log, result.as_ref()) | |
| { | |
| crate::agent_communication::emit_agent_communication_send( | |
| communication_id, | |
| &context, | |
| &communication, | |
| agent_id, | |
| ); | |
| } | |
| result | |
| } | |
| /// Interrupt the current task for an existing agent thread. | |
| pub(crate) async fn interrupt_agent(&self, agent_id: ThreadId) -> CodexResult<String> { | |
| let state = self.upgrade()?; | |
| self.handle_thread_request_result( | |
| agent_id, | |
| &state, | |
| state | |
| .send_op( | |
| agent_id, | |
| Op::Interrupt, | |
| /*parent_turn_id*/ None, | |
| /*root_turn_id*/ None, | |
| ) | |
| .await, | |
| ) | |
| .await | |
| } | |
| async fn handle_thread_request_result( | |
| &self, | |
| agent_id: ThreadId, | |
| state: &Arc<ThreadManagerState>, | |
| result: CodexResult<String>, | |
| ) -> CodexResult<String> { | |
| if result | |
| .as_ref() | |
| .is_err_and(|err| matches!(err.details(), CodexErrorDetails::InternalAgentDied)) | |
| { | |
| let _ = state.remove_thread(&agent_id).await; | |
| self.forget_v2_residency(agent_id); | |
| self.state.release_spawned_thread(agent_id); | |
| } | |
| result | |
| } | |
| /// Fetch the last known status for `agent_id`, returning `NotFound` when unavailable. | |
| pub(crate) async fn get_status(&self, agent_id: ThreadId) -> AgentStatus { | |
| let Ok(state) = self.upgrade() else { | |
| // No agent available if upgrade fails. | |
| return AgentStatus::NotFound; | |
| }; | |
| let Ok(thread) = state.get_thread(agent_id).await else { | |
| return AgentStatus::NotFound; | |
| }; | |
| thread.agent_status().await | |
| } | |
| pub(crate) fn register_session_root( | |
| &self, | |
| current_thread_id: ThreadId, | |
| current_parent_thread_id: Option<ThreadId>, | |
| ) { | |
| if current_parent_thread_id.is_none() { | |
| self.state.register_root_thread(current_thread_id); | |
| } | |
| } | |
| pub(crate) fn get_agent_metadata(&self, agent_id: ThreadId) -> Option<AgentMetadata> { | |
| self.state.agent_metadata_for_thread(agent_id) | |
| } | |
| pub(crate) fn ensure_agent_known(&self, agent_id: ThreadId) -> CodexResult<AgentMetadata> { | |
| self.state | |
| .agent_metadata_for_thread(agent_id) | |
| .ok_or_else(|| CodexErr::ThreadNotFound(agent_id)) | |
| } | |
| pub(crate) async fn list_live_agent_subtree_thread_ids( | |
| &self, | |
| agent_id: ThreadId, | |
| ) -> CodexResult<Vec<ThreadId>> { | |
| let mut thread_ids = vec![agent_id]; | |
| thread_ids.extend(self.live_thread_spawn_descendants(agent_id).await?); | |
| Ok(thread_ids) | |
| } | |
| pub(crate) async fn get_agent_config_snapshot( | |
| &self, | |
| agent_id: ThreadId, | |
| ) -> Option<ThreadConfigSnapshot> { | |
| let Ok(state) = self.upgrade() else { | |
| return None; | |
| }; | |
| let Ok(thread) = state.get_thread(agent_id).await else { | |
| return None; | |
| }; | |
| Some(thread.config_snapshot().await) | |
| } | |
| pub(crate) async fn resolve_agent_reference( | |
| &self, | |
| _current_thread_id: ThreadId, | |
| current_session_source: &SessionSource, | |
| agent_reference: &str, | |
| ) -> CodexResult<ThreadId> { | |
| let current_agent_path = current_session_source | |
| .get_agent_path() | |
| .unwrap_or_else(AgentPath::root); | |
| let agent_path = current_agent_path | |
| .resolve(agent_reference) | |
| .map_err(CodexErr::UnsupportedOperation)?; | |
| if let Some(thread_id) = self.state.agent_id_for_path(&agent_path) { | |
| return Ok(thread_id); | |
| } | |
| Err(CodexErr::UnsupportedOperation(format!( | |
| "live agent path `{}` not found", | |
| agent_path.as_str() | |
| ))) | |
| } | |
| /// Subscribe to status updates for `agent_id`, yielding the latest value and changes. | |
| pub(crate) async fn subscribe_status( | |
| &self, | |
| agent_id: ThreadId, | |
| ) -> CodexResult<watch::Receiver<AgentStatus>> { | |
| let state = self.upgrade()?; | |
| let thread = state.get_thread(agent_id).await?; | |
| Ok(thread.subscribe_status()) | |
| } | |
| pub(crate) async fn format_environment_context_subagents( | |
| &self, | |
| parent_thread_id: ThreadId, | |
| multi_agent_version: MultiAgentVersion, | |
| ) -> String { | |
| if multi_agent_version != MultiAgentVersion::V2 { | |
| let Ok(agents) = self.open_thread_spawn_children(parent_thread_id).await else { | |
| return String::new(); | |
| }; | |
| return agents | |
| .into_iter() | |
| .map(|(thread_id, metadata)| { | |
| let reference = metadata | |
| .agent_path | |
| .as_ref() | |
| .map(|path| path.name().to_string()) | |
| .unwrap_or_else(|| thread_id.to_string()); | |
| format_subagent_context_line(&reference, metadata.agent_nickname.as_deref()) | |
| }) | |
| .collect::<Vec<_>>() | |
| .join("\n"); | |
| } | |
| let Some(parent_path) = self | |
| .state | |
| .agent_metadata_for_thread(parent_thread_id) | |
| .and_then(|metadata| metadata.agent_path) | |
| else { | |
| return String::new(); | |
| }; | |
| let parent_prefix = format!("{parent_path}/"); | |
| let mut agent_paths = self | |
| .state | |
| .live_agents() | |
| .into_iter() | |
| .filter_map(|metadata| metadata.agent_path) | |
| .filter(|path| { | |
| path.as_str() | |
| .strip_prefix(&parent_prefix) | |
| .is_some_and(|name| !name.contains('/')) | |
| }) | |
| .collect::<Vec<_>>(); | |
| let loaded_paths = self | |
| .open_thread_spawn_children(parent_thread_id) | |
| .await | |
| .unwrap_or_default() | |
| .into_iter() | |
| .filter_map(|(_, metadata)| metadata.agent_path) | |
| .collect::<HashSet<_>>(); | |
| agent_paths.sort(); | |
| // Stable sorting preserves alphabetical order within each group. | |
| agent_paths.sort_by_key(|path| !loaded_paths.contains(path)); | |
| let mut lines = Vec::with_capacity(agent_paths.len().min(MAX_ENVIRONMENT_SUBAGENTS)); | |
| let mut rendered_bytes = " <subagents>\n </subagents>\n".len(); | |
| for agent_path in agent_paths { | |
| if lines.len() == MAX_ENVIRONMENT_SUBAGENTS { | |
| break; | |
| } | |
| let line = format!(r#"<agent name="{agent_path}" />"#); | |
| let line_bytes = " \n".len() + line.len(); | |
| if rendered_bytes + line_bytes <= MAX_ENVIRONMENT_SUBAGENT_BYTES { | |
| rendered_bytes += line_bytes; | |
| lines.push(line); | |
| } | |
| } | |
| lines.join("\n") | |
| } | |
| pub(crate) async fn list_agents( | |
| &self, | |
| current_session_source: &SessionSource, | |
| path_prefix: Option<&str>, | |
| ) -> CodexResult<Vec<ListedAgent>> { | |
| let state = self.upgrade()?; | |
| let resolved_prefix = path_prefix | |
| .map(|prefix| { | |
| current_session_source | |
| .get_agent_path() | |
| .unwrap_or_else(AgentPath::root) | |
| .resolve(prefix) | |
| .map_err(CodexErr::UnsupportedOperation) | |
| }) | |
| .transpose()?; | |
| let mut live_agents = self.state.live_agents(); | |
| live_agents.sort_by(|left, right| { | |
| left.agent_path | |
| .as_deref() | |
| .unwrap_or_default() | |
| .cmp(right.agent_path.as_deref().unwrap_or_default()) | |
| .then_with(|| { | |
| left.agent_id | |
| .map(|id| id.to_string()) | |
| .unwrap_or_default() | |
| .cmp(&right.agent_id.map(|id| id.to_string()).unwrap_or_default()) | |
| }) | |
| }); | |
| let root_path = AgentPath::root(); | |
| let mut agents = Vec::with_capacity(live_agents.len().saturating_add(1)); | |
| if resolved_prefix | |
| .as_ref() | |
| .is_none_or(|prefix| agent_matches_prefix(Some(&root_path), prefix)) | |
| && let Some(root_thread_id) = self.state.agent_id_for_path(&root_path) | |
| && let Ok(root_thread) = state.get_thread(root_thread_id).await | |
| { | |
| agents.push(ListedAgent { | |
| agent_name: root_path.to_string(), | |
| agent_status: root_thread.agent_status().await, | |
| }); | |
| } | |
| for metadata in live_agents { | |
| let Some(thread_id) = metadata.agent_id else { | |
| continue; | |
| }; | |
| if resolved_prefix | |
| .as_ref() | |
| .is_some_and(|prefix| !agent_matches_prefix(metadata.agent_path.as_ref(), prefix)) | |
| { | |
| continue; | |
| } | |
| let Ok(thread) = state.get_thread(thread_id).await else { | |
| continue; | |
| }; | |
| let agent_name = metadata | |
| .agent_path | |
| .as_ref() | |
| .map(ToString::to_string) | |
| .unwrap_or_else(|| thread_id.to_string()); | |
| agents.push(ListedAgent { | |
| agent_name, | |
| agent_status: thread.agent_status().await, | |
| }); | |
| } | |
| Ok(agents) | |
| } | |
| /// Starts a detached watcher for sub-agents spawned from another thread. | |
| /// | |
| /// This is only enabled for `SubAgentSource::ThreadSpawn`, where a parent thread exists and | |
| /// can receive completion notifications. | |
| fn maybe_start_completion_watcher( | |
| &self, | |
| child_thread_id: ThreadId, | |
| session_source: Option<SessionSource>, | |
| child_reference: String, | |
| child_agent_path: Option<AgentPath>, | |
| ) { | |
| let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { | |
| parent_thread_id, .. | |
| })) = session_source | |
| else { | |
| return; | |
| }; | |
| let control = self.clone(); | |
| tokio::spawn(async move { | |
| let status = match control.subscribe_status(child_thread_id).await { | |
| Ok(mut status_rx) => { | |
| let mut status = status_rx.borrow().clone(); | |
| while !is_final(&status) { | |
| if status_rx.changed().await.is_err() { | |
| status = control.get_status(child_thread_id).await; | |
| break; | |
| } | |
| status = status_rx.borrow().clone(); | |
| } | |
| status | |
| } | |
| Err(_) => control.get_status(child_thread_id).await, | |
| }; | |
| if !is_final(&status) { | |
| return; | |
| } | |
| let Ok(state) = control.upgrade() else { | |
| return; | |
| }; | |
| let child_thread = state.get_thread(child_thread_id).await.ok(); | |
| let child_uses_multi_agent_v2 = match child_thread.as_ref() { | |
| Some(child_thread) => { | |
| child_thread.multi_agent_version() == Some(MultiAgentVersion::V2) | |
| } | |
| None => true, | |
| }; | |
| if child_agent_path.is_some() && child_uses_multi_agent_v2 { | |
| let Some(child_agent_path) = child_agent_path.clone() else { | |
| return; | |
| }; | |
| let Some(parent_agent_path) = child_agent_path | |
| .as_str() | |
| .rsplit_once('/') | |
| .and_then(|(parent, _)| AgentPath::try_from(parent).ok()) | |
| else { | |
| return; | |
| }; | |
| let Some(message) = format_inter_agent_completion_message( | |
| parent_agent_path.clone(), | |
| child_agent_path.clone(), | |
| &status, | |
| ) else { | |
| return; | |
| }; | |
| let communication = InterAgentCommunication::new( | |
| child_agent_path, | |
| parent_agent_path, | |
| Vec::new(), | |
| message, | |
| /*trigger_turn*/ false, | |
| ); | |
| let context = | |
| AgentCommunicationContext::new(AgentCommunicationKind::Result, child_thread_id); | |
| let _ = control | |
| .send_inter_agent_communication( | |
| parent_thread_id, | |
| communication, | |
| context, | |
| TurnStartOptions::default(), | |
| ) | |
| .await; | |
| return; | |
| } | |
| let Ok(parent_thread) = state.get_thread(parent_thread_id).await else { | |
| return; | |
| }; | |
| parent_thread | |
| .inject_fragment_without_turn(SubagentNotification::new( | |
| child_reference.as_str(), | |
| status, | |
| )) | |
| .await; | |
| }); | |
| } | |
| fn prepare_agent_metadata( | |
| &self, | |
| reservation: &mut crate::agent::registry::SpawnReservation, | |
| config: &Config, | |
| agent_path: Option<AgentPath>, | |
| agent_role: Option<String>, | |
| preferred_agent_nickname: Option<String>, | |
| ) -> CodexResult<AgentMetadata> { | |
| if let Some(agent_path) = agent_path.as_ref() { | |
| reservation.reserve_agent_path(agent_path)?; | |
| } | |
| let candidate_names = spawn::agent_nickname_candidates(config, agent_role.as_deref()); | |
| let candidate_name_refs: Vec<&str> = candidate_names.iter().map(String::as_str).collect(); | |
| let agent_nickname = Some(reservation.reserve_agent_nickname_with_preference( | |
| &candidate_name_refs, | |
| preferred_agent_nickname.as_deref(), | |
| )?); | |
| Ok(AgentMetadata { | |
| agent_id: None, | |
| agent_path, | |
| agent_nickname, | |
| agent_role, | |
| }) | |
| } | |
| fn prepare_thread_spawn( | |
| &self, | |
| reservation: &mut crate::agent::registry::SpawnReservation, | |
| config: &Config, | |
| parent_thread_id: ThreadId, | |
| depth: i32, | |
| agent_path: Option<AgentPath>, | |
| agent_role: Option<String>, | |
| preferred_agent_nickname: Option<String>, | |
| ) -> CodexResult<(SessionSource, AgentMetadata)> { | |
| if depth == 1 { | |
| self.state.register_root_thread(parent_thread_id); | |
| } | |
| let agent_metadata = self.prepare_agent_metadata( | |
| reservation, | |
| config, | |
| agent_path, | |
| agent_role, | |
| preferred_agent_nickname, | |
| )?; | |
| let session_source = SessionSource::SubAgent(SubAgentSource::ThreadSpawn { | |
| parent_thread_id, | |
| depth, | |
| agent_path: agent_metadata.agent_path.clone(), | |
| agent_nickname: agent_metadata.agent_nickname.clone(), | |
| agent_role: agent_metadata.agent_role.clone(), | |
| }); | |
| Ok((session_source, agent_metadata)) | |
| } | |
| fn upgrade(&self) -> CodexResult<Arc<ThreadManagerState>> { | |
| self.manager | |
| .upgrade() | |
| .ok_or_else(|| CodexErr::UnsupportedOperation("thread manager dropped".to_string())) | |
| } | |
| async fn inherited_environments_for_source( | |
| &self, | |
| state: &Arc<ThreadManagerState>, | |
| session_source: Option<&SessionSource>, | |
| ) -> Option<TurnEnvironmentSnapshot> { | |
| let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { | |
| parent_thread_id, .. | |
| })) = session_source | |
| else { | |
| return None; | |
| }; | |
| let parent_thread = state.get_thread(*parent_thread_id).await.ok()?; | |
| Some( | |
| parent_thread | |
| .session | |
| .services | |
| .turn_environments | |
| .snapshot() | |
| .await, | |
| ) | |
| } | |
| async fn inherited_exec_policy_for_source( | |
| &self, | |
| state: &Arc<ThreadManagerState>, | |
| session_source: Option<&SessionSource>, | |
| child_config: &Config, | |
| ) -> Option<Arc<crate::exec_policy::ExecPolicyManager>> { | |
| let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { | |
| parent_thread_id, .. | |
| })) = session_source | |
| else { | |
| return None; | |
| }; | |
| let parent_thread = state.get_thread(*parent_thread_id).await.ok()?; | |
| let parent_config = parent_thread.session.get_config().await; | |
| if !crate::exec_policy::child_uses_parent_exec_policy(&parent_config, child_config) { | |
| return None; | |
| } | |
| Some(Arc::clone(&parent_thread.session.services.exec_policy)) | |
| } | |
| async fn open_thread_spawn_children( | |
| &self, | |
| parent_thread_id: ThreadId, | |
| ) -> CodexResult<Vec<(ThreadId, AgentMetadata)>> { | |
| let mut children_by_parent = self.live_thread_spawn_children().await?; | |
| Ok(children_by_parent | |
| .remove(&parent_thread_id) | |
| .unwrap_or_default()) | |
| } | |
| async fn live_thread_spawn_children( | |
| &self, | |
| ) -> CodexResult<HashMap<ThreadId, Vec<(ThreadId, AgentMetadata)>>> { | |
| let state = self.upgrade()?; | |
| let mut children_by_parent = HashMap::<ThreadId, Vec<(ThreadId, AgentMetadata)>>::new(); | |
| for (parent_thread_id, child_thread_id) in state.list_live_thread_spawn_edges().await { | |
| children_by_parent | |
| .entry(parent_thread_id) | |
| .or_default() | |
| .push(( | |
| child_thread_id, | |
| self.state | |
| .agent_metadata_for_thread(child_thread_id) | |
| .unwrap_or(AgentMetadata { | |
| agent_id: Some(child_thread_id), | |
| ..Default::default() | |
| }), | |
| )); | |
| } | |
| for children in children_by_parent.values_mut() { | |
| children.sort_by(|left, right| { | |
| left.1 | |
| .agent_path | |
| .as_deref() | |
| .unwrap_or_default() | |
| .cmp(right.1.agent_path.as_deref().unwrap_or_default()) | |
| .then_with(|| left.0.to_string().cmp(&right.0.to_string())) | |
| }); | |
| } | |
| Ok(children_by_parent) | |
| } | |
| async fn persist_thread_spawn_edge_for_source( | |
| &self, | |
| child_thread: &crate::CodexThread, | |
| child_thread_id: ThreadId, | |
| session_source: Option<&SessionSource>, | |
| ) { | |
| let Some(parent_thread_id) = session_source.and_then(SessionSource::parent_thread_id) | |
| else { | |
| return; | |
| }; | |
| if child_thread.config_snapshot().await.ephemeral { | |
| return; | |
| } | |
| let Ok(state) = self.upgrade() else { | |
| return; | |
| }; | |
| let Some(agent_graph_store) = state.agent_graph_store() else { | |
| return; | |
| }; | |
| if let Err(err) = agent_graph_store | |
| .upsert_thread_spawn_edge( | |
| parent_thread_id, | |
| child_thread_id, | |
| codex_agent_graph_store::ThreadSpawnEdgeStatus::Open, | |
| ) | |
| .await | |
| { | |
| warn!("failed to persist thread-spawn edge: {err}"); | |
| } | |
| } | |
| async fn live_thread_spawn_descendants( | |
| &self, | |
| root_thread_id: ThreadId, | |
| ) -> CodexResult<Vec<ThreadId>> { | |
| let mut children_by_parent = self.live_thread_spawn_children().await?; | |
| let mut descendants = Vec::new(); | |
| let mut stack = children_by_parent | |
| .remove(&root_thread_id) | |
| .unwrap_or_default() | |
| .into_iter() | |
| .map(|(child_thread_id, _)| child_thread_id) | |
| .rev() | |
| .collect::<Vec<_>>(); | |
| while let Some(thread_id) = stack.pop() { | |
| descendants.push(thread_id); | |
| if let Some(children) = children_by_parent.remove(&thread_id) { | |
| for (child_thread_id, _) in children.into_iter().rev() { | |
| stack.push(child_thread_id); | |
| } | |
| } | |
| } | |
| Ok(descendants) | |
| } | |
| } | |
| fn agent_matches_prefix(agent_path: Option<&AgentPath>, prefix: &AgentPath) -> bool { | |
| if prefix.is_root() { | |
| return true; | |
| } | |
| agent_path.is_some_and(|agent_path| { | |
| agent_path == prefix | |
| || agent_path | |
| .as_str() | |
| .strip_prefix(prefix.as_str()) | |
| .is_some_and(|suffix| suffix.starts_with('/')) | |
| }) | |
| } | |
| pub(crate) fn render_input_preview(input: &[UserInput]) -> String { | |
| input | |
| .iter() | |
| .map(|item| match item { | |
| UserInput::Text { text, .. } => text.clone(), | |
| UserInput::Image { .. } => "[image]".to_string(), | |
| UserInput::LocalImage { path, .. } => { | |
| format!("[local_image:{}]", path.display()) | |
| } | |
| UserInput::Audio { .. } => "[audio]".to_string(), | |
| UserInput::LocalAudio { path } => { | |
| format!("[local_audio:{}]", path.display()) | |
| } | |
| UserInput::Skill { name, path, .. } => { | |
| format!("[skill:${name}]({})", path.display()) | |
| } | |
| UserInput::Mention { name, path, .. } => format!("[mention:${name}]({path})"), | |
| _ => "[input]".to_string(), | |
| }) | |
| .collect::<Vec<_>>() | |
| .join("\n") | |
| } | |
| fn thread_spawn_depth(session_source: &SessionSource) -> Option<i32> { | |
| match session_source { | |
| SessionSource::SubAgent(SubAgentSource::ThreadSpawn { depth, .. }) => Some(*depth), | |
| _ => None, | |
| } | |
| } | |
| mod tests; | |