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::codex_thread::ThreadConfigSnapshot;use crate::config::Config;use crate::environment_selection::TurnEnvironmentSnapshot;use crate::session::emit_subagent_session_started;use crate::session_prefix::format_inter_agent_completion_message;use crate::session_prefix::format_subagent_context_line;use crate::session_prefix::format_subagent_notification_message;use crate::thread_manager::ResumeThreadWithHistoryOptions;use crate::thread_manager::ThreadManagerState;use crate::thread_rollout_truncation::truncate_rollout_to_last_n_fork_turns;use codex_protocol::AgentPath;use codex_protocol::SessionId;use codex_protocol::ThreadId;use codex_protocol::error::CodexErr;use codex_protocol::error::Result as CodexResult;use codex_protocol::models::ContentItem;use codex_protocol::models::MessagePhase;use codex_protocol::models::ResponseItem;use codex_protocol::protocol::InitialHistory;use codex_protocol::protocol::InterAgentCommunication;use codex_protocol::protocol::MultiAgentVersion;use codex_protocol::protocol::Op;use codex_protocol::protocol::ResumedHistory;use codex_protocol::protocol::RolloutItem;use codex_protocol::protocol::SessionSource;use codex_protocol::protocol::SubAgentSource;use codex_protocol::protocol::ThreadSource;use codex_protocol::protocol::TurnEnvironmentSelection;use codex_protocol::user_input::UserInput;use codex_state::DirectionalThreadSpawnEdgeStatus;use codex_thread_store::ReadThreadParams;use serde::Serialize;use std::collections::HashMap;use std::collections::VecDeque;use std::sync::Arc;use std::sync::Weak;use tokio::sync::watch;use tracing::warn;pub(crate) use self::execution::AgentExecutionGuard;use self::execution::AgentExecutionLimiter;use self::residency::V2Residency;const ROOT_LAST_TASK_MESSAGE: &str = "Main thread";mod execution;mod legacy;mod residency;mod spawn;#[derive(Clone, Debug, PartialEq, Eq)]pub(crate) enum SpawnAgentForkMode {    FullHistory,    LastNTurns(usize),}#[derive(Clone, Debug, Default)]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) environments: Option<Vec<TurnEnvironmentSelection>>,}#[derive(Clone, Debug)]pub(crate) struct LiveAgent {    pub(crate) thread_id: ThreadId,    pub(crate) metadata: AgentMetadata,    pub(crate) status: AgentStatus,}#[derive(Clone, Debug, Serialize, PartialEq, Eq)]pub(crate) struct ListedAgent {    pub(crate) agent_name: String,    pub(crate) agent_status: AgentStatus,    pub(crate) last_task_message: Option<String>,}/// 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`.#[derive(Clone, Default)]pub(crate) struct AgentControl {    /// ID shared by the whole agent control session. This means every sub-agents from a common    /// root share the same session 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>,    state: Arc<AgentRegistry>,    v2_residency: Arc<V2Residency>,    agent_execution_limiter: Arc<AgentExecutionLimiter>,}impl AgentControl {    /// Construct a new `AgentControl` that can spawn/message agents via the given manager state.    pub(crate) fn new(manager: Weak<ThreadManagerState>) -> Self {        Self {            manager,            ..Default::default()        }    }    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    }    /// Send rich user input items to an existing agent thread.    pub(crate) async fn send_input(        &self,        agent_id: ThreadId,        initial_operation: Op,    ) -> CodexResult<String> {        let state = self.upgrade()?;        self.ensure_execution_capacity_for_op(agent_id, &initial_operation)            .await?;        self.send_input_after_capacity_check(agent_id, &state, initial_operation)            .await    }    async fn send_input_after_capacity_check(        &self,        agent_id: ThreadId,        state: &Arc<ThreadManagerState>,        initial_operation: Op,    ) -> CodexResult<String> {        let last_task_message = match &initial_operation {            Op::InterAgentCommunication { communication } => {                last_task_message_from_communication(communication)            }            _ => non_empty_task_message(render_input_preview(&initial_operation)),        };        let result = self            .handle_thread_request_result(                agent_id,                state,                state.send_op(agent_id, initial_operation).await,            )            .await;        if result.is_ok() {            match last_task_message {                Some(last_task_message) => self                    .state                    .update_last_task_message(agent_id, last_task_message),                None => self.state.clear_last_task_message(agent_id),            }        }        result    }    pub(crate) async fn send_inter_agent_communication(        &self,        agent_id: ThreadId,        communication: InterAgentCommunication,    ) -> CodexResult<String> {        let last_task_message = last_task_message_from_communication(&communication);        let state = self.upgrade()?;        let op = Op::InterAgentCommunication { communication };        self.ensure_execution_capacity_for_op(agent_id, &op).await?;        let result = self            .handle_thread_request_result(agent_id, &state, state.send_op(agent_id, op).await)            .await;        if result.is_ok() {            match last_task_message {                Some(last_task_message) => self                    .state                    .update_last_task_message(agent_id, last_task_message),                None => self.state.clear_last_task_message(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).await,        )        .await    }    async fn handle_thread_request_result(        &self,        agent_id: ThreadId,        state: &Arc<ThreadManagerState>,        result: CodexResult<String>,    ) -> CodexResult<String> {        if matches!(result, Err(CodexErr::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(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,    ) -> String {        let Ok(agents) = self.open_thread_spawn_children(parent_thread_id).await else {            return String::new();        };        agents            .into_iter()            .map(|(thread_id, metadata)| {                let reference = metadata                    .agent_path                    .as_ref()                    .map(|agent_path| agent_path.name().to_string())                    .unwrap_or_else(|| thread_id.to_string());                format_subagent_context_line(reference.as_str(), metadata.agent_nickname.as_deref())            })            .collect::<Vec<_>>()            .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,                last_task_message: Some(ROOT_LAST_TASK_MESSAGE.to_string()),            });        }        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());            let last_task_message = metadata.last_task_message.clone();            agents.push(ListedAgent {                agent_name,                agent_status: thread.agent_status().await,                last_task_message,            });        }        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 _ = control                    .send_inter_agent_communication(parent_thread_id, communication)                    .await;                return;            }            let message = format_subagent_notification_message(child_reference.as_str(), &status);            let Ok(parent_thread) = state.get_thread(parent_thread_id).await else {                return;            };            parent_thread                .inject_user_message_without_turn(message)                .await;        });    }    #[allow(clippy::too_many_arguments)]    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);        }        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(),        )?);        let session_source = SessionSource::SubAgent(SubAgentSource::ThreadSpawn {            parent_thread_id,            depth,            agent_path: agent_path.clone(),            agent_nickname: agent_nickname.clone(),            agent_role: agent_role.clone(),        });        let agent_metadata = AgentMetadata {            agent_id: None,            agent_path,            agent_nickname,            agent_role,            last_task_message: None,        };        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                .codex                .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.codex.session.get_config().await;        if !crate::exec_policy::child_uses_parent_exec_policy(&parent_config, child_config) {            return None;        }        Some(Arc::clone(            &parent_thread.codex.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,        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;        };        let Some(state_db_ctx) = thread.state_db() else {            return;        };        if let Err(err) = state_db_ctx            .upsert_thread_spawn_edge(                parent_thread_id,                child_thread_id,                DirectionalThreadSpawnEdgeStatus::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(initial_operation: &Op) -> String {    match initial_operation {        Op::UserInput { items, .. } => items            .iter()            .map(|item| match item {                UserInput::Text { text, .. } => text.clone(),                UserInput::Image { .. } => "[image]".to_string(),                UserInput::LocalImage { path, .. } => {                    format!("[local_image:{}]", 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"),        Op::InterAgentCommunication { communication } => communication.content.clone(),        _ => String::new(),    }}fn last_task_message_from_communication(communication: &InterAgentCommunication) -> Option<String> {    if communication.encrypted_content.is_some() {        return None;    }    non_empty_task_message(communication.content.clone())}fn non_empty_task_message(message: String) -> Option<String> {    (!message.is_empty()).then_some(message)}fn thread_spawn_depth(session_source: &SessionSource) -> Option<i32> {    match session_source {        SessionSource::SubAgent(SubAgentSource::ThreadSpawn { depth, .. }) => Some(*depth),        _ => None,    }}#[cfg(test)]#[path = "control_tests.rs"]mod tests;