diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index e76a4a2f1..a5cac3907 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -363,7 +363,7 @@ impl Codex { /// /// A session has at most 1 running task at a time, and can be interrupted by user input. pub(crate) struct Session { - conversation_id: ThreadId, + pub(crate) conversation_id: ThreadId, tx_event: Sender, agent_status: watch::Sender, state: Mutex, diff --git a/codex-rs/core/src/rollout/policy.rs b/codex-rs/core/src/rollout/policy.rs index 224e45dc5..fafdc8310 100644 --- a/codex-rs/core/src/rollout/policy.rs +++ b/codex-rs/core/src/rollout/policy.rs @@ -90,6 +90,14 @@ pub(crate) fn should_persist_event_msg(ev: &EventMsg) -> bool { | EventMsg::AgentMessageContentDelta(_) | EventMsg::ReasoningContentDelta(_) | EventMsg::ReasoningRawContentDelta(_) - | EventMsg::SkillsUpdateAvailable => false, + | EventMsg::SkillsUpdateAvailable + | EventMsg::CollabAgentSpawnBegin(_) + | EventMsg::CollabAgentSpawnEnd(_) + | EventMsg::CollabAgentInteractionBegin(_) + | EventMsg::CollabAgentInteractionEnd(_) + | EventMsg::CollabWaitingBegin(_) + | EventMsg::CollabWaitingEnd(_) + | EventMsg::CollabCloseBegin(_) + | EventMsg::CollabCloseEnd(_) => false, } } diff --git a/codex-rs/core/src/tools/handlers/collab.rs b/codex-rs/core/src/tools/handlers/collab.rs index 7d254ea9f..805af76b3 100644 --- a/codex-rs/core/src/tools/handlers/collab.rs +++ b/codex-rs/core/src/tools/handlers/collab.rs @@ -1,4 +1,5 @@ use crate::agent::AgentStatus; +use crate::codex::Session; use crate::codex::TurnContext; use crate::config::Config; use crate::error::CodexErr; @@ -11,6 +12,14 @@ use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; use async_trait::async_trait; use codex_protocol::ThreadId; +use codex_protocol::protocol::CollabAgentInteractionBeginEvent; +use codex_protocol::protocol::CollabAgentInteractionEndEvent; +use codex_protocol::protocol::CollabAgentSpawnBeginEvent; +use codex_protocol::protocol::CollabAgentSpawnEndEvent; +use codex_protocol::protocol::CollabCloseBeginEvent; +use codex_protocol::protocol::CollabCloseEndEvent; +use codex_protocol::protocol::CollabWaitingBeginEvent; +use codex_protocol::protocol::CollabWaitingEndEvent; use serde::Deserialize; use serde::Serialize; @@ -40,6 +49,7 @@ impl ToolHandler for CollabHandler { turn, tool_name, payload, + call_id, .. } = invocation; @@ -53,10 +63,10 @@ impl ToolHandler for CollabHandler { }; match tool_name.as_str() { - "spawn_agent" => spawn::handle(session, turn, arguments).await, - "send_input" => send_input::handle(session, arguments).await, - "wait" => wait::handle(session, arguments).await, - "close_agent" => close_agent::handle(session, arguments).await, + "spawn_agent" => spawn::handle(session, turn, call_id, arguments).await, + "send_input" => send_input::handle(session, turn, call_id, arguments).await, + "wait" => wait::handle(session, turn, call_id, arguments).await, + "close_agent" => close_agent::handle(session, turn, call_id, arguments).await, other => Err(FunctionCallError::RespondToModel(format!( "unsupported collab tool {other}" ))), @@ -66,7 +76,6 @@ impl ToolHandler for CollabHandler { mod spawn { use super::*; - use crate::codex::Session; use std::sync::Arc; #[derive(Debug, Deserialize)] @@ -82,24 +91,58 @@ mod spawn { pub async fn handle( session: Arc, turn: Arc, + call_id: String, arguments: String, ) -> Result { let args: SpawnAgentArgs = parse_arguments(&arguments)?; - if args.message.trim().is_empty() { + let prompt = args.message; + if prompt.trim().is_empty() { return Err(FunctionCallError::RespondToModel( "Empty message can't be sent to an agent".to_string(), )); } + session + .send_event( + &turn, + CollabAgentSpawnBeginEvent { + call_id: call_id.clone(), + sender_thread_id: session.conversation_id, + prompt: prompt.clone(), + } + .into(), + ) + .await; let config = build_agent_spawn_config(turn.as_ref())?; let result = session .services .agent_control - .spawn_agent(config, args.message) + .spawn_agent(config, prompt.clone()) .await - .map_err(collab_spawn_error)?; + .map_err(collab_spawn_error); + let (new_thread_id, status) = match &result { + Ok(thread_id) => ( + Some(*thread_id), + session.services.agent_control.get_status(*thread_id).await, + ), + Err(_) => (None, AgentStatus::NotFound), + }; + session + .send_event( + &turn, + CollabAgentSpawnEndEvent { + call_id, + sender_thread_id: session.conversation_id, + new_thread_id, + prompt, + status, + } + .into(), + ) + .await; + let new_thread_id = result?; let content = serde_json::to_string(&SpawnAgentResult { - agent_id: result.to_string(), + agent_id: new_thread_id.to_string(), }) .map_err(|err| { FunctionCallError::Fatal(format!("failed to serialize spawn_agent result: {err}")) @@ -115,7 +158,6 @@ mod spawn { mod send_input { use super::*; - use crate::codex::Session; use std::sync::Arc; #[derive(Debug, Deserialize)] @@ -131,22 +173,55 @@ mod send_input { pub async fn handle( session: Arc, + turn: Arc, + call_id: String, arguments: String, ) -> Result { let args: SendInputArgs = parse_arguments(&arguments)?; - let agent_id = agent_id(&args.id)?; - if args.message.trim().is_empty() { + let receiver_thread_id = agent_id(&args.id)?; + let prompt = args.message; + if prompt.trim().is_empty() { return Err(FunctionCallError::RespondToModel( "Empty message can't be sent to an agent".to_string(), )); } - let agent_id_for_err = agent_id; - let submission_id = session + session + .send_event( + &turn, + CollabAgentInteractionBeginEvent { + call_id: call_id.clone(), + sender_thread_id: session.conversation_id, + receiver_thread_id, + prompt: prompt.clone(), + } + .into(), + ) + .await; + let result = session .services .agent_control - .send_prompt(agent_id, args.message) + .send_prompt(receiver_thread_id, prompt.clone()) .await - .map_err(|err| collab_agent_error(agent_id_for_err, err))?; + .map_err(|err| collab_agent_error(receiver_thread_id, err)); + let status = session + .services + .agent_control + .get_status(receiver_thread_id) + .await; + session + .send_event( + &turn, + CollabAgentInteractionEndEvent { + call_id, + sender_thread_id: session.conversation_id, + receiver_thread_id, + prompt, + status, + } + .into(), + ) + .await; + let submission_id = result?; let content = serde_json::to_string(&SendInputResult { submission_id }).map_err(|err| { FunctionCallError::Fatal(format!("failed to serialize send_input result: {err}")) @@ -163,7 +238,6 @@ mod send_input { mod wait { use super::*; use crate::agent::status::is_final; - use crate::codex::Session; use std::sync::Arc; use std::time::Duration; use tokio::time::Instant; @@ -183,10 +257,12 @@ mod wait { pub async fn handle( session: Arc, + turn: Arc, + call_id: String, arguments: String, ) -> Result { let args: WaitArgs = parse_arguments(&arguments)?; - let agent_id = agent_id(&args.id)?; + let receiver_thread_id = agent_id(&args.id)?; // Validate timeout. let timeout_ms = args.timeout_ms.unwrap_or(DEFAULT_WAIT_TIMEOUT_MS); @@ -199,14 +275,85 @@ mod wait { ms => ms.min(MAX_WAIT_TIMEOUT_MS), }; - let agent_id_for_err = agent_id; - let mut status_rx = session + session + .send_event( + &turn, + CollabWaitingBeginEvent { + sender_thread_id: session.conversation_id, + receiver_thread_id, + call_id: call_id.clone(), + } + .into(), + ) + .await; + let status_rx = match session .services .agent_control - .subscribe_status(agent_id) + .subscribe_status(receiver_thread_id) .await - .map_err(|err| collab_agent_error(agent_id_for_err, err))?; + { + Ok(status_rx) => status_rx, + Err(err) => { + let status = session + .services + .agent_control + .get_status(receiver_thread_id) + .await; + session + .send_event( + &turn, + CollabWaitingEndEvent { + sender_thread_id: session.conversation_id, + receiver_thread_id, + call_id: call_id.clone(), + status, + } + .into(), + ) + .await; + return Err(collab_agent_error(receiver_thread_id, err)); + } + }; + let result = + wait_for_status(session.as_ref(), receiver_thread_id, timeout_ms, status_rx).await; + session + .send_event( + &turn, + CollabWaitingEndEvent { + sender_thread_id: session.conversation_id, + receiver_thread_id, + call_id, + status: result.status.clone(), + } + .into(), + ) + .await; + if matches!(result.status, AgentStatus::NotFound) { + return Err(FunctionCallError::RespondToModel(format!( + "agent with id {receiver_thread_id} not found" + ))); + } + + let content = serde_json::to_string(&result).map_err(|err| { + FunctionCallError::Fatal(format!("failed to serialize wait result: {err}")) + })?; + + let success = !result.timed_out && !matches!(result.status, AgentStatus::Errored(_)); + + Ok(ToolOutput::Function { + content, + success: Some(success), + content_items: None, + }) + } + + async fn wait_for_status( + session: &Session, + agent_id: ThreadId, + timeout_ms: i64, + mut status_rx: tokio::sync::watch::Receiver, + ) -> WaitResult { // Get last known status. let mut status = status_rx.borrow_and_update().clone(); let deadline = Instant::now() + Duration::from_millis(timeout_ms as u64); @@ -231,31 +378,12 @@ mod wait { } }; - if matches!(status, AgentStatus::NotFound) { - return Err(FunctionCallError::RespondToModel(format!( - "agent with id {agent_id} not found" - ))); - } - - let result = WaitResult { status, timed_out }; - - let content = serde_json::to_string(&result).map_err(|err| { - FunctionCallError::Fatal(format!("failed to serialize wait result: {err}")) - })?; - - let success = !result.timed_out && !matches!(result.status, AgentStatus::Errored(_)); - - Ok(ToolOutput::Function { - content, - success: Some(success), - content_items: None, - }) + WaitResult { status, timed_out } } } pub mod close_agent { use super::*; - use crate::codex::Session; use std::sync::Arc; #[derive(Debug, Deserialize, Serialize)] @@ -265,28 +393,71 @@ pub mod close_agent { pub async fn handle( session: Arc, + turn: Arc, + call_id: String, arguments: String, ) -> Result { let args: CloseAgentArgs = parse_arguments(&arguments)?; let agent_id = agent_id(&args.id)?; - let agent_id_for_err = agent_id; - let mut status_rx = session + session + .send_event( + &turn, + CollabCloseBeginEvent { + call_id: call_id.clone(), + sender_thread_id: session.conversation_id, + receiver_thread_id: agent_id, + } + .into(), + ) + .await; + let status = match session .services .agent_control .subscribe_status(agent_id) .await - .map_err(|err| collab_agent_error(agent_id_for_err, err))?; - let status = status_rx.borrow_and_update().clone(); - - if !matches!(status, AgentStatus::Shutdown) { - let agent_id_for_err = agent_id; - let _ = session + { + Ok(mut status_rx) => status_rx.borrow_and_update().clone(), + Err(err) => { + let status = session.services.agent_control.get_status(agent_id).await; + session + .send_event( + &turn, + CollabCloseEndEvent { + call_id: call_id.clone(), + sender_thread_id: session.conversation_id, + receiver_thread_id: agent_id, + status, + } + .into(), + ) + .await; + return Err(collab_agent_error(agent_id, err)); + } + }; + let result = if !matches!(status, AgentStatus::Shutdown) { + session .services .agent_control .shutdown_agent(agent_id) .await - .map_err(|err| collab_agent_error(agent_id_for_err, err))?; - } + .map_err(|err| collab_agent_error(agent_id, err)) + .map(|_| ()) + } else { + Ok(()) + }; + session + .send_event( + &turn, + CollabCloseEndEvent { + call_id, + sender_thread_id: session.conversation_id, + receiver_thread_id: agent_id, + status: status.clone(), + } + .into(), + ) + .await; + result?; let content = serde_json::to_string(&CloseAgentResult { status }).map_err(|err| { FunctionCallError::Fatal(format!("failed to serialize close_agent result: {err}")) diff --git a/codex-rs/exec/src/event_processor_with_human_output.rs b/codex-rs/exec/src/event_processor_with_human_output.rs index f1cba0b9f..42a83e44d 100644 --- a/codex-rs/exec/src/event_processor_with_human_output.rs +++ b/codex-rs/exec/src/event_processor_with_human_output.rs @@ -571,6 +571,16 @@ impl EventProcessor for EventProcessorWithHumanOutput { EventMsg::ContextCompacted(_) => { ts_msg!(self, "context compacted"); } + EventMsg::CollabAgentSpawnBegin(_) + | EventMsg::CollabAgentSpawnEnd(_) + | EventMsg::CollabAgentInteractionBegin(_) + | EventMsg::CollabAgentInteractionEnd(_) + | EventMsg::CollabWaitingBegin(_) + | EventMsg::CollabWaitingEnd(_) + | EventMsg::CollabCloseBegin(_) + | EventMsg::CollabCloseEnd(_) => { + // TODO(jif) handle collab tools. + } EventMsg::ShutdownComplete => return CodexStatus::Shutdown, EventMsg::WebSearchBegin(_) | EventMsg::ExecApprovalRequest(_) diff --git a/codex-rs/mcp-server/src/codex_tool_runner.rs b/codex-rs/mcp-server/src/codex_tool_runner.rs index c7aaa6190..4d826f920 100644 --- a/codex-rs/mcp-server/src/codex_tool_runner.rs +++ b/codex-rs/mcp-server/src/codex_tool_runner.rs @@ -349,6 +349,14 @@ async fn run_codex_tool_session_inner( | EventMsg::ExitedReviewMode(_) | EventMsg::ContextCompacted(_) | EventMsg::ThreadRolledBack(_) + | EventMsg::CollabAgentSpawnBegin(_) + | EventMsg::CollabAgentSpawnEnd(_) + | EventMsg::CollabAgentInteractionBegin(_) + | EventMsg::CollabAgentInteractionEnd(_) + | EventMsg::CollabWaitingBegin(_) + | EventMsg::CollabWaitingEnd(_) + | EventMsg::CollabCloseBegin(_) + | EventMsg::CollabCloseEnd(_) | EventMsg::DeprecationNotice(_) => { // For now, we do not do anything extra for these // events. Note that diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index ca66fc7dc..2116813e1 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -773,6 +773,71 @@ pub enum EventMsg { AgentMessageContentDelta(AgentMessageContentDeltaEvent), ReasoningContentDelta(ReasoningContentDeltaEvent), ReasoningRawContentDelta(ReasoningRawContentDeltaEvent), + + /// Collab interaction: agent spawn begin. + CollabAgentSpawnBegin(CollabAgentSpawnBeginEvent), + /// Collab interaction: agent spawn end. + CollabAgentSpawnEnd(CollabAgentSpawnEndEvent), + /// Collab interaction: agent interaction begin. + CollabAgentInteractionBegin(CollabAgentInteractionBeginEvent), + /// Collab interaction: agent interaction end. + CollabAgentInteractionEnd(CollabAgentInteractionEndEvent), + /// Collab interaction: waiting begin. + CollabWaitingBegin(CollabWaitingBeginEvent), + /// Collab interaction: waiting end. + CollabWaitingEnd(CollabWaitingEndEvent), + /// Collab interaction: close begin. + CollabCloseBegin(CollabCloseBeginEvent), + /// Collab interaction: close end. + CollabCloseEnd(CollabCloseEndEvent), +} + +impl From for EventMsg { + fn from(event: CollabAgentSpawnBeginEvent) -> Self { + EventMsg::CollabAgentSpawnBegin(event) + } +} + +impl From for EventMsg { + fn from(event: CollabAgentSpawnEndEvent) -> Self { + EventMsg::CollabAgentSpawnEnd(event) + } +} + +impl From for EventMsg { + fn from(event: CollabAgentInteractionBeginEvent) -> Self { + EventMsg::CollabAgentInteractionBegin(event) + } +} + +impl From for EventMsg { + fn from(event: CollabAgentInteractionEndEvent) -> Self { + EventMsg::CollabAgentInteractionEnd(event) + } +} + +impl From for EventMsg { + fn from(event: CollabWaitingBeginEvent) -> Self { + EventMsg::CollabWaitingBegin(event) + } +} + +impl From for EventMsg { + fn from(event: CollabWaitingEndEvent) -> Self { + EventMsg::CollabWaitingEnd(event) + } +} + +impl From for EventMsg { + fn from(event: CollabCloseBeginEvent) -> Self { + EventMsg::CollabCloseBegin(event) + } +} + +impl From for EventMsg { + fn from(event: CollabCloseEndEvent) -> Self { + EventMsg::CollabCloseEnd(event) + } } /// Agent lifecycle status, derived from emitted events. @@ -2023,6 +2088,105 @@ pub enum TurnAbortReason { ReviewEnded, } +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema, TS)] +pub struct CollabAgentSpawnBeginEvent { + /// Identifier for the collab tool call. + pub call_id: String, + /// Thread ID of the sender. + pub sender_thread_id: ThreadId, + /// Initial prompt sent to the agent. Can be empty to prevent CoT leaking at the + /// beginning. + pub prompt: String, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema, TS)] +pub struct CollabAgentSpawnEndEvent { + /// Identifier for the collab tool call. + pub call_id: String, + /// Thread ID of the sender. + pub sender_thread_id: ThreadId, + /// Thread ID of the newly spawned agent, if it was created. + pub new_thread_id: Option, + /// Initial prompt sent to the agent. Can be empty to prevent CoT leaking at the + /// beginning. + pub prompt: String, + /// Last known status of the new agent reported to the sender agent. + pub status: AgentStatus, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema, TS)] +pub struct CollabAgentInteractionBeginEvent { + /// Identifier for the collab tool call. + pub call_id: String, + /// Thread ID of the sender. + pub sender_thread_id: ThreadId, + /// Thread ID of the receiver. + pub receiver_thread_id: ThreadId, + /// Prompt sent from the sender to the receiver. Can be empty to prevent CoT + /// leaking at the beginning. + pub prompt: String, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema, TS)] +pub struct CollabAgentInteractionEndEvent { + /// Identifier for the collab tool call. + pub call_id: String, + /// Thread ID of the sender. + pub sender_thread_id: ThreadId, + /// Thread ID of the receiver. + pub receiver_thread_id: ThreadId, + /// Prompt sent from the sender to the receiver. Can be empty to prevent CoT + /// leaking at the beginning. + pub prompt: String, + /// Last known status of the receiver agent reported to the sender agent. + pub status: AgentStatus, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema, TS)] +pub struct CollabWaitingBeginEvent { + /// Thread ID of the sender. + pub sender_thread_id: ThreadId, + /// Thread ID of the receiver. + pub receiver_thread_id: ThreadId, + /// ID of the waiting call. + pub call_id: String, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema, TS)] +pub struct CollabWaitingEndEvent { + /// Thread ID of the sender. + pub sender_thread_id: ThreadId, + /// Thread ID of the receiver. + pub receiver_thread_id: ThreadId, + /// ID of the waiting call. + pub call_id: String, + /// Last known status of the receiver agent reported to the sender agent. + pub status: AgentStatus, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema, TS)] +pub struct CollabCloseBeginEvent { + /// Identifier for the collab tool call. + pub call_id: String, + /// Thread ID of the sender. + pub sender_thread_id: ThreadId, + /// Thread ID of the receiver. + pub receiver_thread_id: ThreadId, +} + +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema, TS)] +pub struct CollabCloseEndEvent { + /// Identifier for the collab tool call. + pub call_id: String, + /// Thread ID of the sender. + pub sender_thread_id: ThreadId, + /// Thread ID of the receiver. + pub receiver_thread_id: ThreadId, + /// Last known status of the receiver agent reported to the sender agent before + /// the close. + pub status: AgentStatus, +} + #[cfg(test)] mod tests { use super::*; diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index f1f532a62..f3c9c2f4c 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -2402,6 +2402,16 @@ impl ChatWidget { } EventMsg::ExitedReviewMode(review) => self.on_exited_review_mode(review), EventMsg::ContextCompacted(_) => self.on_agent_message("Context compacted".to_owned()), + EventMsg::CollabAgentSpawnBegin(_) + | EventMsg::CollabAgentSpawnEnd(_) + | EventMsg::CollabAgentInteractionBegin(_) + | EventMsg::CollabAgentInteractionEnd(_) + | EventMsg::CollabWaitingBegin(_) + | EventMsg::CollabWaitingEnd(_) + | EventMsg::CollabCloseBegin(_) + | EventMsg::CollabCloseEnd(_) => { + // TODO(jif) handle collab tools. + } EventMsg::ThreadRolledBack(_) => {} EventMsg::RawResponseItem(_) | EventMsg::ItemStarted(_) diff --git a/codex-rs/tui2/src/chatwidget.rs b/codex-rs/tui2/src/chatwidget.rs index 9ebf17882..fb36217ee 100644 --- a/codex-rs/tui2/src/chatwidget.rs +++ b/codex-rs/tui2/src/chatwidget.rs @@ -2155,6 +2155,16 @@ impl ChatWidget { } EventMsg::ExitedReviewMode(review) => self.on_exited_review_mode(review), EventMsg::ContextCompacted(_) => self.on_agent_message("Context compacted".to_owned()), + EventMsg::CollabAgentSpawnBegin(_) + | EventMsg::CollabAgentSpawnEnd(_) + | EventMsg::CollabAgentInteractionBegin(_) + | EventMsg::CollabAgentInteractionEnd(_) + | EventMsg::CollabWaitingBegin(_) + | EventMsg::CollabWaitingEnd(_) + | EventMsg::CollabCloseBegin(_) + | EventMsg::CollabCloseEnd(_) => { + // TODO(jif) handle collab tools. + } EventMsg::RawResponseItem(_) | EventMsg::ThreadRolledBack(_) | EventMsg::ItemStarted(_)