From 37ac0c093cb4be42f7812737366cab181b9d0417 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Mon, 23 Mar 2026 18:53:54 +0000 Subject: [PATCH] feat: structured multi-agent output (#15515) Send input now sends messages as assistant message and with this format: ``` author: /root/worker_a recipient: /root/worker_a/tester other_recipients: [] Content: bla bla bla. Actual content. Only text for now ``` --- codex-rs/core/src/agent/control.rs | 76 +++- codex-rs/core/src/agent/control_tests.rs | 55 +++ .../core/src/agent/inter_agent_instruction.rs | 74 ++++ codex-rs/core/src/agent/mod.rs | 1 + codex-rs/core/src/codex.rs | 81 +++-- .../core/src/codex/rollout_reconstruction.rs | 10 +- .../src/codex/rollout_reconstruction_tests.rs | 102 ++++++ codex-rs/core/src/codex_tests.rs | 28 ++ codex-rs/core/src/codex_thread.rs | 101 +++++- codex-rs/core/src/context_manager/history.rs | 15 +- .../core/src/context_manager/history_tests.rs | 41 +++ codex-rs/core/src/event_mapping.rs | 2 +- codex-rs/core/src/event_mapping_tests.rs | 37 ++ codex-rs/core/src/tasks/mod.rs | 30 ++ codex-rs/core/src/thread_manager.rs | 27 ++ .../tools/handlers/multi_agents/send_input.rs | 48 ++- .../src/tools/handlers/multi_agents_tests.rs | 332 ++++++++++++++++++ 17 files changed, 994 insertions(+), 66 deletions(-) create mode 100644 codex-rs/core/src/agent/inter_agent_instruction.rs diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index d2a7cbf9d..5bc5a0711 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -1,4 +1,6 @@ use crate::agent::AgentStatus; +use crate::agent::inter_agent_instruction::InterAgentDelivery; +use crate::agent::inter_agent_instruction::InterAgentInstruction; use crate::agent::registry::AgentMetadata; use crate::agent::registry::AgentRegistry; use crate::agent::role::DEFAULT_ROLE_NAME; @@ -454,20 +456,53 @@ impl AgentControl { items: Vec, ) -> CodexResult { let state = self.upgrade()?; - let result = state - .send_op( - agent_id, - Op::UserInput { - items, - final_output_json_schema: None, - }, - ) - .await; - if matches!(result, Err(CodexErr::InternalAgentDied)) { - let _ = state.remove_thread(&agent_id).await; - self.state.release_spawned_thread(agent_id); - } - result + self.handle_thread_request_result( + agent_id, + &state, + state + .send_op( + agent_id, + Op::UserInput { + items, + final_output_json_schema: None, + }, + ) + .await, + ) + .await + } + + /// Append a prebuilt message to an existing agent thread outside the normal user-input path. + #[cfg(test)] + pub(crate) async fn append_message( + &self, + agent_id: ThreadId, + message: ResponseItem, + ) -> CodexResult { + let state = self.upgrade()?; + self.handle_thread_request_result( + agent_id, + &state, + state.append_message(agent_id, message).await, + ) + .await + } + + pub(crate) async fn deliver_inter_agent_instruction( + &self, + agent_id: ThreadId, + instruction: InterAgentInstruction, + delivery: InterAgentDelivery, + ) -> CodexResult { + let state = self.upgrade()?; + self.handle_thread_request_result( + agent_id, + &state, + state + .deliver_inter_agent_instruction(agent_id, instruction, delivery) + .await, + ) + .await } /// Interrupt the current task for an existing agent thread. @@ -476,6 +511,19 @@ impl AgentControl { state.send_op(agent_id, Op::Interrupt).await } + async fn handle_thread_request_result( + &self, + agent_id: ThreadId, + state: &Arc, + result: CodexResult, + ) -> CodexResult { + if matches!(result, Err(CodexErr::InternalAgentDied)) { + let _ = state.remove_thread(&agent_id).await; + self.state.release_spawned_thread(agent_id); + } + result + } + /// Submit a shutdown request for a live agent without marking it explicitly closed in /// persisted spawn-edge state. pub(crate) async fn shutdown_live_agent(&self, agent_id: ThreadId) -> CodexResult { diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 20c051f85..677bf5a2f 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -382,6 +382,61 @@ async fn send_input_submits_user_message() { assert_eq!(captured, Some(expected)); } +#[tokio::test] +async fn append_message_records_assistant_message() { + let harness = AgentControlHarness::new().await; + let (thread_id, thread) = harness.start_thread().await; + let message = + "author: /root\nrecipient: /root/worker\nother_recipients: []\nContent: hello from tests"; + + let submission_id = harness + .control + .append_message( + thread_id, + ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::InputText { + text: message.to_string(), + }], + end_turn: None, + phase: None, + }, + ) + .await + .expect("append_message should succeed"); + assert!(!submission_id.is_empty()); + + timeout(Duration::from_secs(5), async { + loop { + let history_items = thread + .codex + .session + .clone_history() + .await + .raw_items() + .to_vec(); + let recorded = history_items.iter().any(|item| { + matches!( + item, + ResponseItem::Message { role, content, .. } + if role == "assistant" + && content.iter().any(|content_item| matches!( + content_item, + ContentItem::InputText { text } if text == message + )) + ) + }); + if recorded { + break; + } + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("assistant message should be recorded"); +} + #[tokio::test] async fn spawn_agent_creates_thread_and_sends_prompt() { let harness = AgentControlHarness::new().await; diff --git a/codex-rs/core/src/agent/inter_agent_instruction.rs b/codex-rs/core/src/agent/inter_agent_instruction.rs new file mode 100644 index 000000000..0eff40beb --- /dev/null +++ b/codex-rs/core/src/agent/inter_agent_instruction.rs @@ -0,0 +1,74 @@ +use codex_protocol::AgentPath; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum InterAgentDelivery { + CurrentTurn, + NextTurn, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) struct InterAgentInstruction { + author: AgentPath, + recipient: AgentPath, + other_recipients: Vec, + content: String, +} + +impl InterAgentInstruction { + pub(crate) fn new( + author: AgentPath, + recipient: AgentPath, + other_recipients: Vec, + content: String, + ) -> Self { + Self { + author, + recipient, + other_recipients, + content, + } + } + + pub(crate) fn to_response_item(&self) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: self.as_text(), + }], + end_turn: None, + phase: None, + } + } + + pub(crate) fn is_message_content(content: &[ContentItem]) -> bool { + content.iter().any(|content_item| match content_item { + ContentItem::InputText { text } | ContentItem::OutputText { text } => { + Self::is_instruction_text(text) + } + _ => false, + }) + } + + fn as_text(&self) -> String { + let other_recipients = self + .other_recipients + .iter() + .map(std::string::ToString::to_string) + .collect::>() + .join(", "); + format!( + "author: {}\nrecipient: {}\nother_recipients: [{other_recipients}]\nContent: {}", + self.author, self.recipient, self.content + ) + } + + fn is_instruction_text(text: &str) -> bool { + text.starts_with("author: ") + && text.contains("\nrecipient: ") + && text.contains("\nother_recipients: [") + && text.contains("]\nContent: ") + } +} diff --git a/codex-rs/core/src/agent/mod.rs b/codex-rs/core/src/agent/mod.rs index 350962dc0..3f14e5c59 100644 --- a/codex-rs/core/src/agent/mod.rs +++ b/codex-rs/core/src/agent/mod.rs @@ -1,5 +1,6 @@ pub(crate) mod agent_resolver; pub(crate) mod control; +pub(crate) mod inter_agent_instruction; mod registry; pub(crate) mod role; pub(crate) mod status; diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 099edb7c0..b040883c2 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -798,6 +798,7 @@ pub(crate) struct Session { pending_mcp_server_refresh_config: Mutex>, pub(crate) conversation: Arc, pub(crate) active_turn: Mutex>, + idle_pending_input: Mutex>, pub(crate) guardian_review_session: GuardianReviewSessionManager, pub(crate) services: SessionServices, js_repl: Arc, @@ -809,6 +810,7 @@ pub(crate) struct TurnSkillsContext { pub(crate) outcome: Arc, pub(crate) implicit_invocation_seen_skills: Arc>>, } + impl TurnSkillsContext { pub(crate) fn new(outcome: Arc) -> Self { Self { @@ -1895,6 +1897,7 @@ impl Session { pending_mcp_server_refresh_config: Mutex::new(None), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), + idle_pending_input: Mutex::new(Vec::new()), guardian_review_session: GuardianReviewSessionManager::default(), services, js_repl, @@ -3912,7 +3915,7 @@ impl Session { Ok(active_turn_id.clone()) } - /// Returns the input if there was no task running to inject into + /// Returns the input if there was no task running to inject into. pub async fn inject_response_items( &self, input: Vec, @@ -3953,6 +3956,24 @@ impl Session { } } + /// Queue response items to be injected into the next active turn created for this session. + pub(crate) async fn queue_response_items_for_next_turn(&self, items: Vec) { + if items.is_empty() { + return; + } + + let mut idle_pending_input = self.idle_pending_input.lock().await; + idle_pending_input.extend(items); + } + + pub(crate) async fn take_queued_response_items_for_next_turn(&self) -> Vec { + std::mem::take(&mut *self.idle_pending_input.lock().await) + } + + pub(crate) async fn has_queued_response_items_for_next_turn(&self) -> bool { + !self.idle_pending_input.lock().await.is_empty() + } + pub async fn has_pending_input(&self) -> bool { let active = self.active_turn.lock().await; match active.as_ref() { @@ -5441,7 +5462,7 @@ pub(crate) async fn run_turn( prewarmed_client_session: Option, cancellation_token: CancellationToken, ) -> Option { - if input.is_empty() { + if input.is_empty() && !sess.has_pending_input().await { return None; } @@ -5583,25 +5604,33 @@ pub(crate) async fn run_turn( }) .collect::>(); - let initial_input_for_turn: ResponseInputItem = ResponseInputItem::from(input.clone()); - let response_item: ResponseItem = initial_input_for_turn.clone().into(); - let mut last_agent_message: Option = None; if run_pending_session_start_hooks(&sess, &turn_context).await { - return last_agent_message; + return None; } - let user_prompt_submit_outcome = - run_user_prompt_submit_hooks(&sess, &turn_context, UserMessageItem::new(&input).message()) - .await; - if user_prompt_submit_outcome.should_stop { - record_additional_contexts( + let additional_contexts = if input.is_empty() { + Vec::new() + } else { + let initial_input_for_turn: ResponseInputItem = ResponseInputItem::from(input.clone()); + let response_item: ResponseItem = initial_input_for_turn.clone().into(); + let user_prompt_submit_outcome = run_user_prompt_submit_hooks( &sess, &turn_context, - user_prompt_submit_outcome.additional_contexts, + UserMessageItem::new(&input).message(), ) .await; - return last_agent_message; - } - let additional_contexts = user_prompt_submit_outcome.additional_contexts; + if user_prompt_submit_outcome.should_stop { + record_additional_contexts( + &sess, + &turn_context, + user_prompt_submit_outcome.additional_contexts, + ) + .await; + return None; + } + sess.record_user_prompt_and_emit_turn_item(turn_context.as_ref(), &input, response_item) + .await; + user_prompt_submit_outcome.additional_contexts + }; sess.services .analytics_events_client .track_app_mentioned(tracking.clone(), mentioned_app_invocations); @@ -5612,17 +5641,17 @@ pub(crate) async fn run_turn( } sess.merge_connector_selection(explicitly_enabled_connectors.clone()) .await; - sess.record_user_prompt_and_emit_turn_item(turn_context.as_ref(), &input, response_item) - .await; record_additional_contexts(&sess, &turn_context, additional_contexts).await; - // Track the previous-turn baseline from the regular user-turn path only so - // standalone tasks (compact/shell/review/undo) cannot suppress future - // model/realtime injections. - sess.set_previous_turn_settings(Some(PreviousTurnSettings { - model: turn_context.model_info.slug.clone(), - realtime_active: Some(turn_context.realtime_active), - })) - .await; + if !input.is_empty() { + // Track the previous-turn baseline from the regular user-turn path only so + // standalone tasks (compact/shell/review/undo) cannot suppress future + // model/realtime injections. + sess.set_previous_turn_settings(Some(PreviousTurnSettings { + model: turn_context.model_info.slug.clone(), + realtime_active: Some(turn_context.realtime_active), + })) + .await; + } if !skill_items.is_empty() { sess.record_conversation_items(&turn_context, &skill_items) @@ -5633,8 +5662,10 @@ pub(crate) async fn run_turn( .await; } + let skills_outcome = Some(turn_context.turn_skills.outcome.as_ref()); sess.maybe_start_ghost_snapshot(Arc::clone(&turn_context), cancellation_token.child_token()) .await; + let mut last_agent_message: Option = None; let mut stop_hook_active = false; // Although from the perspective of codex.rs, TurnDiffTracker has the lifecycle of a Task which contains // many turns, from the perspective of the user, it is a single turn. diff --git a/codex-rs/core/src/codex/rollout_reconstruction.rs b/codex-rs/core/src/codex/rollout_reconstruction.rs index 11ddbc192..a4c042af0 100644 --- a/codex-rs/core/src/codex/rollout_reconstruction.rs +++ b/codex-rs/core/src/codex/rollout_reconstruction.rs @@ -1,4 +1,5 @@ use super::*; +use crate::context_manager::is_user_turn_boundary; // Return value of `Session::reconstruct_history_from_rollout`, bundling the rebuilt history with // the resume/fork hydration metadata derived from the same replay. @@ -201,9 +202,12 @@ impl Session { ); } } - RolloutItem::ResponseItem(_) - | RolloutItem::EventMsg(_) - | RolloutItem::SessionMeta(_) => {} + RolloutItem::ResponseItem(response_item) => { + let active_segment = + active_segment.get_or_insert_with(ActiveReplaySegment::default); + active_segment.counts_as_user_turn |= is_user_turn_boundary(response_item); + } + RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) => {} } if base_replacement_history.is_some() diff --git a/codex-rs/core/src/codex/rollout_reconstruction_tests.rs b/codex-rs/core/src/codex/rollout_reconstruction_tests.rs index 6cc99a290..09b4f4561 100644 --- a/codex-rs/core/src/codex/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/codex/rollout_reconstruction_tests.rs @@ -33,6 +33,18 @@ fn assistant_message(text: &str) -> ResponseItem { } } +fn inter_agent_assistant_message(text: &str) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: text.to_string(), + }], + end_turn: None, + phase: None, + } +} + #[tokio::test] async fn record_initial_history_resumed_bare_turn_context_does_not_hydrate_previous_turn_settings() { @@ -430,6 +442,96 @@ async fn reconstruct_history_rollback_skips_non_user_turns_for_history_and_metad ); } +#[tokio::test] +async fn reconstruct_history_rollback_counts_inter_agent_assistant_turns() { + let (session, turn_context) = make_session_and_context().await; + let first_context_item = turn_context.to_turn_context_item(); + let first_turn_id = first_context_item + .turn_id + .clone() + .expect("turn context should have turn_id"); + let assistant_turn_id = "assistant-instruction-turn".to_string(); + let assistant_turn_context = TurnContextItem { + turn_id: Some(assistant_turn_id.clone()), + ..first_context_item.clone() + }; + let assistant_instruction = inter_agent_assistant_message( + "author: /root\nrecipient: /root/worker\nother_recipients: []\nContent: continue", + ); + let assistant_reply = assistant_message("worker reply"); + + let rollout_items = vec![ + RolloutItem::EventMsg(EventMsg::TurnStarted( + codex_protocol::protocol::TurnStartedEvent { + turn_id: first_turn_id.clone(), + model_context_window: Some(128_000), + collaboration_mode_kind: ModeKind::Default, + }, + )), + RolloutItem::EventMsg(EventMsg::UserMessage( + codex_protocol::protocol::UserMessageEvent { + message: "turn 1 user".to_string(), + images: None, + local_images: Vec::new(), + text_elements: Vec::new(), + }, + )), + RolloutItem::TurnContext(first_context_item.clone()), + RolloutItem::ResponseItem(user_message("turn 1 user")), + RolloutItem::ResponseItem(assistant_message("turn 1 assistant")), + RolloutItem::EventMsg(EventMsg::TurnComplete( + codex_protocol::protocol::TurnCompleteEvent { + turn_id: first_turn_id, + last_agent_message: None, + }, + )), + RolloutItem::EventMsg(EventMsg::TurnStarted( + codex_protocol::protocol::TurnStartedEvent { + turn_id: assistant_turn_id.clone(), + model_context_window: Some(128_000), + collaboration_mode_kind: ModeKind::Default, + }, + )), + RolloutItem::TurnContext(assistant_turn_context), + RolloutItem::ResponseItem(assistant_instruction), + RolloutItem::ResponseItem(assistant_reply), + RolloutItem::EventMsg(EventMsg::TurnComplete( + codex_protocol::protocol::TurnCompleteEvent { + turn_id: assistant_turn_id, + last_agent_message: None, + }, + )), + RolloutItem::EventMsg(EventMsg::ThreadRolledBack( + codex_protocol::protocol::ThreadRolledBackEvent { num_turns: 1 }, + )), + ]; + + let reconstructed = session + .reconstruct_history_from_rollout(&turn_context, &rollout_items) + .await; + + assert_eq!( + reconstructed.history, + vec![ + user_message("turn 1 user"), + assistant_message("turn 1 assistant") + ] + ); + assert_eq!( + reconstructed.previous_turn_settings, + Some(PreviousTurnSettings { + model: turn_context.model_info.slug.clone(), + realtime_active: Some(turn_context.realtime_active), + }) + ); + assert_eq!( + serde_json::to_value(reconstructed.reference_context_item) + .expect("serialize reconstructed reference context item"), + serde_json::to_value(Some(first_context_item)) + .expect("serialize expected reference context item") + ); +} + #[tokio::test] async fn reconstruct_history_rollback_clears_history_and_metadata_when_exceeding_user_turns() { let (session, turn_context) = make_session_and_context().await; diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 2305cb1fa..24470051f 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -2682,6 +2682,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { pending_mcp_server_refresh_config: Mutex::new(None), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), + idle_pending_input: Mutex::new(Vec::new()), guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(), services, js_repl, @@ -3481,6 +3482,7 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( pending_mcp_server_refresh_config: Mutex::new(None), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), + idle_pending_input: Mutex::new(Vec::new()), guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(), services, js_repl, @@ -4631,6 +4633,32 @@ async fn prepend_pending_input_keeps_older_tail_ahead_of_newer_input() { assert_eq!(sess.get_pending_input().await, vec![later, newer]); } +#[tokio::test] +async fn queued_response_items_for_next_turn_move_into_next_active_turn() { + let (sess, tc, _rx) = make_session_and_context_with_rx().await; + let queued_item = ResponseInputItem::Message { + role: "assistant".to_string(), + content: vec![ContentItem::InputText { + text: "queued before wake".to_string(), + }], + }; + + sess.queue_response_items_for_next_turn(vec![queued_item.clone()]) + .await; + + sess.spawn_task( + Arc::clone(&tc), + Vec::new(), + NeverEndingTask { + kind: TaskKind::Regular, + listen_to_cancellation_token: false, + }, + ) + .await; + + assert_eq!(sess.get_pending_input().await, vec![queued_item]); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn abort_review_task_emits_exited_then_aborted_and_records_history() { let (sess, tc, rx) = make_session_and_context_with_rx().await; diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index e016fec97..ff36de6c1 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -1,4 +1,6 @@ use crate::agent::AgentStatus; +use crate::agent::inter_agent_instruction::InterAgentDelivery; +use crate::agent::inter_agent_instruction::InterAgentInstruction; use crate::codex::Codex; use crate::codex::SteerInputError; use crate::config::ConstraintResult; @@ -122,29 +124,90 @@ impl CodexThread { /// Records a user-role session-prefix message without creating a new user turn boundary. pub(crate) async fn inject_user_message_without_turn(&self, message: String) { - let pending_item = ResponseInputItem::Message { + let message = ResponseItem::Message { + id: None, role: "user".to_string(), content: vec![ContentItem::InputText { text: message }], + end_turn: None, + phase: None, }; - let pending_items = vec![pending_item]; - let Err(items_without_active_turn) = self + let pending_item = match pending_message_input_item(&message) { + Ok(pending_item) => pending_item, + Err(err) => { + debug_assert!(false, "session-prefix message append should succeed: {err}"); + return; + } + }; + if self .codex .session - .inject_response_items(pending_items) + .inject_response_items(vec![pending_item]) .await - else { - return; - }; + .is_err() + { + let turn_context = self.codex.session.new_default_turn().await; + self.codex + .session + .record_conversation_items(turn_context.as_ref(), &[message]) + .await; + } + } - let turn_context = self.codex.session.new_default_turn().await; - let items: Vec = items_without_active_turn - .into_iter() - .map(ResponseItem::from) - .collect(); + /// Append a prebuilt message to the thread history without treating it as a user turn. + /// + /// If the thread already has an active turn, the message is queued as pending input for that + /// turn. Otherwise it is queued at session scope and a regular turn is started so the agent + /// can consume that pending input through the normal turn pipeline. + pub(crate) async fn append_message(&self, message: ResponseItem) -> CodexResult { + let submission_id = uuid::Uuid::new_v4().to_string(); + let pending_item = pending_message_input_item(&message)?; + if let Err(items) = self + .codex + .session + .inject_response_items(vec![pending_item]) + .await + { + self.codex + .session + .queue_response_items_for_next_turn(items) + .await; + self.codex + .session + .ensure_task_for_queued_response_items() + .await; + } + + Ok(submission_id) + } + + pub(crate) async fn deliver_inter_agent_instruction( + &self, + instruction: InterAgentInstruction, + delivery: InterAgentDelivery, + ) -> CodexResult { + let message = instruction.to_response_item(); + match delivery { + InterAgentDelivery::CurrentTurn => self.append_message(message).await, + InterAgentDelivery::NextTurn => self.queue_message_for_next_turn(message).await, + } + } + + /// Queue a prebuilt message so the next turn records it before any submitted user input. + pub(crate) async fn queue_message_for_next_turn( + &self, + message: ResponseItem, + ) -> CodexResult { + let submission_id = uuid::Uuid::new_v4().to_string(); + let pending_item = pending_message_input_item(&message)?; self.codex .session - .record_conversation_items(turn_context.as_ref(), &items) + .queue_response_items_for_next_turn(vec![pending_item]) .await; + self.codex + .session + .ensure_task_for_queued_response_items() + .await; + Ok(submission_id) } pub fn rollout_path(&self) -> Option { @@ -198,3 +261,15 @@ impl CodexThread { Ok(*guard) } } + +fn pending_message_input_item(message: &ResponseItem) -> CodexResult { + match message { + ResponseItem::Message { role, content, .. } => Ok(ResponseInputItem::Message { + role: role.clone(), + content: content.clone(), + }), + _ => Err(CodexErr::InvalidRequest( + "append_message only supports ResponseItem::Message".to_string(), + )), + } +} diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index f990a80dc..6ed037048 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -1,3 +1,4 @@ +use crate::agent::inter_agent_instruction::InterAgentInstruction; use crate::codex::TurnContext; use crate::context_manager::normalize; use crate::event_mapping::is_contextual_user_message_content; @@ -204,9 +205,10 @@ impl ContextManager { } } - /// Drop the last `num_turns` user turns from this history. + /// Drop the last `num_turns` instruction turns from this history. /// - /// "User turns" are identified as `ResponseItem::Message` entries whose role is `"user"`. + /// Instruction turns are history messages that should behave like a new prompt boundary: + /// ordinary user messages and structured assistant inter-agent instructions. /// /// This mirrors thread-rollback semantics: /// - `num_turns == 0` is a no-op @@ -248,7 +250,7 @@ impl ContextManager { } fn get_non_last_reasoning_items_tokens(&self) -> i64 { - // Get reasoning items excluding all the ones after the last user message. + // Get reasoning items excluding all the ones after the last instruction boundary. let Some(last_user_index) = self.items.iter().rposition(is_user_turn_boundary) else { return 0; }; @@ -631,7 +633,12 @@ pub(crate) fn is_user_turn_boundary(item: &ResponseItem) -> bool { return false; }; - role == "user" && !is_contextual_user_message_content(content) + (role == "user" && !is_contextual_user_message_content(content)) + || (role == "assistant" && is_inter_agent_instruction_content(content)) +} + +fn is_inter_agent_instruction_content(content: &[ContentItem]) -> bool { + InterAgentInstruction::is_message_content(content) } fn user_message_positions(items: &[ResponseItem]) -> Vec { diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index 71b3aded0..4deb76ed6 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -38,6 +38,18 @@ fn assistant_msg(text: &str) -> ResponseItem { } } +fn inter_agent_assistant_msg(text: &str) -> ResponseItem { + ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: text.to_string(), + }], + end_turn: None, + phase: None, + } +} + fn create_history_with_items(items: Vec) -> ContextManager { let mut h = ContextManager::new(); // Use a generous but fixed token budget; tests only rely on truncation @@ -225,6 +237,35 @@ fn items_after_last_model_generated_tokens_are_zero_without_model_generated_item ); } +#[test] +fn inter_agent_assistant_messages_are_turn_boundaries() { + let item = inter_agent_assistant_msg( + "author: /root\nrecipient: /root/worker\nother_recipients: []\nContent: continue", + ); + + assert!(is_user_turn_boundary(&item)); +} + +#[test] +fn drop_last_n_user_turns_treats_inter_agent_assistant_messages_as_instruction_turns() { + let first_turn = user_input_text_msg("first"); + let first_reply = assistant_msg("done"); + let inter_agent_turn = inter_agent_assistant_msg( + "author: /root\nrecipient: /root/worker\nother_recipients: []\nContent: continue", + ); + let inter_agent_reply = assistant_msg("worker reply"); + let mut history = create_history_with_items(vec![ + first_turn.clone(), + first_reply.clone(), + inter_agent_turn, + inter_agent_reply, + ]); + + history.drop_last_n_user_turns(1); + + assert_eq!(history.raw_items(), &vec![first_turn, first_reply]); +} + #[test] fn total_token_usage_includes_all_items_after_last_model_generated_item() { let mut history = create_history_with_items(vec![assistant_msg("already counted by API")]); diff --git a/codex-rs/core/src/event_mapping.rs b/codex-rs/core/src/event_mapping.rs index ad776d142..8f80547a4 100644 --- a/codex-rs/core/src/event_mapping.rs +++ b/codex-rs/core/src/event_mapping.rs @@ -72,7 +72,7 @@ fn parse_agent_message( let mut content: Vec = Vec::new(); for content_item in message.iter() { match content_item { - ContentItem::OutputText { text } => { + ContentItem::InputText { text } | ContentItem::OutputText { text } => { content.push(AgentMessageContent::Text { text: text.clone() }); } _ => { diff --git a/codex-rs/core/src/event_mapping_tests.rs b/codex-rs/core/src/event_mapping_tests.rs index 553550d74..65c2e7d97 100644 --- a/codex-rs/core/src/event_mapping_tests.rs +++ b/codex-rs/core/src/event_mapping_tests.rs @@ -95,6 +95,43 @@ fn skips_local_image_label_text() { } } +#[test] +fn parses_assistant_message_input_text_for_backward_compatibility() { + let item = ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::InputText { + text: "author: /root\nrecipient: /root/worker\nother_recipients: []\nContent: continue" + .to_string(), + }], + end_turn: None, + phase: None, + }; + + let turn_item = parse_turn_item(&item).expect("expected assistant message turn item"); + + match turn_item { + TurnItem::AgentMessage(message) => { + let rendered = message + .content + .into_iter() + .map(|content| { + let AgentMessageContent::Text { text } = content; + text + }) + .collect::>(); + assert_eq!( + rendered, + vec![ + "author: /root\nrecipient: /root/worker\nother_recipients: []\nContent: continue" + .to_string() + ] + ); + } + other => panic!("expected TurnItem::AgentMessage, got {other:?}"), + } +} + #[test] fn skips_unnamed_image_label_text() { let image_url = "data:image/png;base64,abc".to_string(); diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index b8e1d73b7..b2f110486 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -153,7 +153,15 @@ impl Session { ) { self.abort_all_tasks(TurnAbortReason::Replaced).await; self.clear_connector_selection().await; + self.start_task(turn_context, input, task).await; + } + async fn start_task( + self: &Arc, + turn_context: Arc, + input: Vec, + task: T, + ) { let task: Arc = Arc::new(task); let task_kind = task.kind(); let span_name = task.span_name(); @@ -224,6 +232,22 @@ impl Session { .await; } + pub(crate) async fn ensure_task_for_queued_response_items(self: &Arc) { + if !self.has_queued_response_items_for_next_turn().await { + return; + } + + if self.active_turn.lock().await.is_some() { + return; + } + + let turn_context = self.new_default_turn().await; + self.maybe_emit_unknown_model_warning_for_turn(turn_context.as_ref()) + .await; + self.start_task(turn_context, Vec::new(), RegularTask::new()) + .await; + } + pub async fn abort_all_tasks(self: &Arc, reason: TurnAbortReason) { if let Some(mut active_turn) = self.take_active_turn().await { for task in active_turn.drain_tasks() { @@ -233,6 +257,9 @@ impl Session { // in-flight approval wait can surface as a model-visible rejection before TurnAborted. active_turn.clear_pending().await; } + if reason == TurnAbortReason::Interrupted { + self.ensure_task_for_queued_response_items().await; + } } pub async fn on_task_finished( @@ -371,6 +398,9 @@ impl Session { let mut turn = ActiveTurn::default(); let mut turn_state = turn.turn_state.lock().await; turn_state.token_usage_at_turn_start = token_usage_at_turn_start; + for item in self.take_queued_response_items_for_next_turn().await { + turn_state.push_pending_input(item); + } drop(turn_state); turn.add_task(task); *active = Some(turn); diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index a63cf2cb9..d93bf3718 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -3,6 +3,8 @@ use crate::CodexAuth; use crate::ModelProviderInfo; use crate::OPENAI_PROVIDER_ID; use crate::agent::AgentControl; +use crate::agent::inter_agent_instruction::InterAgentDelivery; +use crate::agent::inter_agent_instruction::InterAgentInstruction; use crate::codex::Codex; use crate::codex::CodexSpawnArgs; use crate::codex::CodexSpawnOk; @@ -26,6 +28,8 @@ use crate::shell_snapshot::ShellSnapshot; use crate::skills::SkillsManager; use codex_protocol::ThreadId; use codex_protocol::config_types::CollaborationModeMask; +#[cfg(test)] +use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ModelPreset; use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::McpServerRefreshConfig; @@ -606,6 +610,29 @@ impl ThreadManagerState { thread.submit(op).await } + #[cfg(test)] + /// Append a prebuilt message to a thread by ID outside the normal user-input path. + pub(crate) async fn append_message( + &self, + thread_id: ThreadId, + message: ResponseItem, + ) -> CodexResult { + let thread = self.get_thread(thread_id).await?; + thread.append_message(message).await + } + + pub(crate) async fn deliver_inter_agent_instruction( + &self, + thread_id: ThreadId, + instruction: InterAgentInstruction, + delivery: InterAgentDelivery, + ) -> CodexResult { + let thread = self.get_thread(thread_id).await?; + thread + .deliver_inter_agent_instruction(instruction, delivery) + .await + } + /// Remove a thread from the manager by ID, returning it when present. pub(crate) async fn remove_thread(&self, thread_id: &ThreadId) -> Option> { self.threads.write().await.remove(thread_id) diff --git a/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs b/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs index 8fc4dd515..25f70c730 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/send_input.rs @@ -1,7 +1,15 @@ use super::*; +use crate::agent::inter_agent_instruction::InterAgentDelivery; +use crate::agent::inter_agent_instruction::InterAgentInstruction; pub(crate) struct Handler; +fn can_use_v2_inter_agent_instruction(items: &[UserInput]) -> bool { + items + .iter() + .all(|item| matches!(item, UserInput::Text { .. })) +} + #[async_trait] impl ToolHandler for Handler { type Output = SendInputResult; @@ -52,12 +60,40 @@ impl ToolHandler for Handler { .into(), ) .await; - let result = session - .services - .agent_control - .send_input(receiver_thread_id, input_items) - .await - .map_err(|err| collab_agent_error(receiver_thread_id, err)); + let agent_control = session.services.agent_control.clone(); + let result = if turn.config.features.enabled(Feature::MultiAgentV2) + && can_use_v2_inter_agent_instruction(&input_items) + { + let receiver_agent_path = receiver_agent.agent_path.clone().ok_or_else(|| { + FunctionCallError::RespondToModel( + "target agent is missing an agent_path".to_string(), + ) + })?; + let instruction = InterAgentInstruction::new( + turn.session_source + .get_agent_path() + .unwrap_or_else(AgentPath::root), + receiver_agent_path, + Vec::new(), + prompt.clone(), + ); + agent_control + .deliver_inter_agent_instruction( + receiver_thread_id, + instruction, + if args.interrupt { + InterAgentDelivery::NextTurn + } else { + InterAgentDelivery::CurrentTurn + }, + ) + .await + } else { + agent_control + .send_input(receiver_thread_id, input_items) + .await + } + .map_err(|err| collab_agent_error(receiver_thread_id, err)); let status = session .services .agent_control diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index d8b2a713a..b5ea255d8 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -7,6 +7,7 @@ use crate::codex::make_session_and_context; use crate::config::DEFAULT_AGENT_MAX_DEPTH; use crate::config::types::ShellEnvironmentPolicy; use crate::function_tool::FunctionCallError; +use crate::protocol::AgentStatus; use crate::protocol::AskForApproval; use crate::protocol::FileSystemSandboxPolicy; use crate::protocol::NetworkSandboxPolicy; @@ -14,6 +15,9 @@ use crate::protocol::Op; use crate::protocol::SandboxPolicy; use crate::protocol::SessionSource; use crate::protocol::SubAgentSource; +use crate::state::TaskKind; +use crate::tasks::SessionTask; +use crate::tasks::SessionTaskContext; use crate::tools::context::ToolOutput; use crate::turn_diff_tracker::TurnDiffTracker; use codex_features::Feature; @@ -24,6 +28,7 @@ use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::InitialHistory; use codex_protocol::protocol::RolloutItem; +use codex_protocol::user_input::UserInput; use pretty_assertions::assert_eq; use serde::Deserialize; use serde_json::json; @@ -33,6 +38,7 @@ use std::sync::Arc; use std::time::Duration; use tokio::sync::Mutex; use tokio::time::timeout; +use tokio_util::sync::CancellationToken; fn invocation( session: Arc, @@ -68,6 +74,31 @@ fn thread_manager() -> ThreadManager { ) } +#[derive(Clone, Copy)] +struct NeverEndingTask; + +#[async_trait::async_trait] +impl SessionTask for NeverEndingTask { + fn kind(&self) -> TaskKind { + TaskKind::Regular + } + + fn span_name(&self) -> &'static str { + "session_task.multi_agent_never_ending" + } + + async fn run( + self: Arc, + _session: Arc, + _ctx: Arc, + _input: Vec, + cancellation_token: CancellationToken, + ) -> Option { + cancellation_token.cancelled().await; + None + } +} + fn expect_text_output(output: T) -> (String, Option) where T: ToolOutput, @@ -337,6 +368,307 @@ async fn multi_agent_v2_spawn_returns_path_and_send_input_accepts_relative_path( )) .await .expect("send_input should accept v2 path"); + + let child_thread = manager + .get_thread(child_thread_id) + .await + .expect("child thread should exist"); + timeout(Duration::from_secs(2), async { + loop { + let history_items = child_thread + .codex + .session + .clone_history() + .await + .raw_items() + .to_vec(); + let recorded = history_items.iter().any(|item| { + matches!( + item, + ResponseItem::Message { role, content, .. } + if role == "assistant" + && content.iter().any(|content_item| matches!( + content_item, + ContentItem::OutputText { text } + if text + == "author: /root\nrecipient: /root/test_process\nother_recipients: []\nContent: continue" + )) + ) + }); + let saw_user_message = history_items.iter().any(|item| { + matches!( + item, + ResponseItem::Message { role, content, .. } + if role == "user" + && content.iter().any(|content_item| matches!( + content_item, + ContentItem::InputText { text } if text == "continue" + )) + ) + }); + if recorded && !saw_user_message { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("v2 send_input should record assistant envelope"); +} + +#[tokio::test] +async fn multi_agent_v2_send_input_accepts_structured_items() { + let (mut session, mut turn) = make_session_and_context().await; + let manager = thread_manager(); + let root = manager + .start_thread((*turn.config).clone()) + .await + .expect("root thread should start"); + session.services.agent_control = manager.agent_control(); + session.conversation_id = root.thread_id; + let mut config = turn.config.as_ref().clone(); + let _ = config.features.enable(Feature::MultiAgentV2); + turn.config = Arc::new(config); + let session = Arc::new(session); + let turn = Arc::new(turn); + + SpawnAgentHandler + .handle(invocation( + session.clone(), + turn.clone(), + "spawn_agent", + function_payload(json!({ + "message": "boot worker", + "task_name": "worker" + })), + )) + .await + .expect("spawn worker"); + let agent_id = session + .services + .agent_control + .resolve_agent_reference(session.conversation_id, &turn.session_source, "worker") + .await + .expect("worker should resolve"); + let thread = manager + .get_thread(agent_id) + .await + .expect("worker thread should exist"); + let invocation = invocation( + session, + turn, + "send_input", + function_payload(json!({ + "target": agent_id.to_string(), + "items": [ + {"type": "mention", "name": "drive", "path": "app://google_drive"}, + {"type": "text", "text": "read the folder"} + ] + })), + ); + + SendInputHandler + .handle(invocation) + .await + .expect("structured items should be accepted in v2"); + + let expected = Op::UserInput { + items: vec![ + UserInput::Mention { + name: "drive".to_string(), + path: "app://google_drive".to_string(), + }, + UserInput::Text { + text: "read the folder".to_string(), + text_elements: Vec::new(), + }, + ], + final_output_json_schema: None, + }; + let captured = manager + .captured_ops() + .into_iter() + .find(|(id, op)| *id == agent_id && *op == expected); + assert_eq!(captured, Some((agent_id, expected))); + + timeout(Duration::from_secs(2), async { + loop { + let history_items = thread + .codex + .session + .clone_history() + .await + .raw_items() + .to_vec(); + let recorded_assistant_envelope = history_items.iter().any(|item| { + matches!( + item, + ResponseItem::Message { role, content, .. } + if role == "assistant" + && content.iter().any(|content_item| matches!( + content_item, + ContentItem::OutputText { text } + if text + == "author: /root\nrecipient: /root/worker\nother_recipients: []\nContent: [mention:$drive](app://google_drive)\nread the folder" + )) + ) + }); + let saw_user_message = history_items.iter().any(|item| { + matches!( + item, + ResponseItem::Message { role, content, .. } + if role == "user" + && content.iter().any(|content_item| matches!( + content_item, + ContentItem::InputText { text } + if text == "read the folder" + || text == "[mention:$drive](app://google_drive)\nread the folder" + )) + ) + }); + if !recorded_assistant_envelope && saw_user_message { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("structured items should stay on the legacy user-input path"); + + let _ = thread + .submit(Op::Shutdown {}) + .await + .expect("shutdown should submit"); +} + +#[tokio::test] +async fn multi_agent_v2_send_input_interrupts_busy_child_without_losing_message() { + let (mut session, mut turn) = make_session_and_context().await; + let manager = thread_manager(); + let root = manager + .start_thread((*turn.config).clone()) + .await + .expect("root thread should start"); + session.services.agent_control = manager.agent_control(); + session.conversation_id = root.thread_id; + let mut config = turn.config.as_ref().clone(); + let _ = config.features.enable(Feature::MultiAgentV2); + turn.config = Arc::new(config); + let session = Arc::new(session); + let turn = Arc::new(turn); + + SpawnAgentHandler + .handle(invocation( + session.clone(), + turn.clone(), + "spawn_agent", + function_payload(json!({ + "message": "boot worker", + "task_name": "worker" + })), + )) + .await + .expect("spawn worker"); + let agent_id = session + .services + .agent_control + .resolve_agent_reference(session.conversation_id, &turn.session_source, "worker") + .await + .expect("worker should resolve"); + let thread = manager + .get_thread(agent_id) + .await + .expect("worker thread should exist"); + + let active_turn = thread.codex.session.new_default_turn().await; + thread + .codex + .session + .spawn_task( + Arc::clone(&active_turn), + vec![UserInput::Text { + text: "working".to_string(), + text_elements: Vec::new(), + }], + NeverEndingTask, + ) + .await; + + SendInputHandler + .handle(invocation( + session, + turn, + "send_input", + function_payload(json!({ + "target": agent_id.to_string(), + "message": "continue", + "interrupt": true + })), + )) + .await + .expect("interrupting v2 send_input should succeed"); + + let ops = manager.captured_ops(); + let ops_for_agent: Vec<&Op> = ops + .iter() + .filter_map(|(id, op)| (*id == agent_id).then_some(op)) + .collect(); + assert!(ops_for_agent.iter().any(|op| matches!(op, Op::Interrupt))); + assert!(!ops_for_agent.iter().any(|op| matches!( + op, + Op::UserInput { items, .. } + if items.iter().any(|item| matches!( + item, + UserInput::Text { text, .. } if text == "continue" + )) + ))); + + timeout(Duration::from_secs(5), async { + loop { + let history_items = thread + .codex + .session + .clone_history() + .await + .raw_items() + .to_vec(); + let saw_envelope = history_items.iter().any(|item| { + matches!( + item, + ResponseItem::Message { role, content, .. } + if role == "assistant" + && content.iter().any(|content_item| matches!( + content_item, + ContentItem::OutputText { text } + if text + == "author: /root\nrecipient: /root/worker\nother_recipients: []\nContent: continue" + )) + ) + }); + let saw_user_message = history_items.iter().any(|item| { + matches!( + item, + ResponseItem::Message { role, content, .. } + if role == "user" + && content.iter().any(|content_item| matches!( + content_item, + ContentItem::InputText { text } if text == "continue" + )) + ) + }); + if saw_envelope && !saw_user_message { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("interrupting v2 send_input should preserve the redirected message"); + + let _ = thread + .submit(Op::Shutdown {}) + .await + .expect("shutdown should submit"); } #[tokio::test]