From a0e1e3ca2baaf6092c773777e769e940ccdbf6c3 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Mon, 15 Jun 2026 20:18:52 +0200 Subject: [PATCH] feat: better MAv2 message rendering --- codex-rs/core/src/agent/control.rs | 10 +++- codex-rs/core/src/agent/control_tests.rs | 8 ++- codex-rs/core/src/config/mod.rs | 3 +- .../context/inter_agent_completion_message.rs | 41 ++++++++++++++ codex-rs/core/src/context/mod.rs | 2 + codex-rs/core/src/guardian/prompt.rs | 20 ++----- codex-rs/core/src/realtime_context.rs | 16 ++---- codex-rs/core/src/session/mod.rs | 10 +++- codex-rs/core/src/session_prefix.rs | 18 ++++++ .../src/tools/handlers/multi_agents_tests.rs | 18 +++--- .../tests/suite/subagent_notifications.rs | 23 +++++--- codex-rs/ext/web-search/src/history.rs | 12 +--- codex-rs/protocol/src/models.rs | 28 ++++++++++ codex-rs/protocol/src/protocol.rs | 55 +++++++++++++++++-- .../rollout-trace/src/reducer/tool/agents.rs | 13 +++-- .../src/reducer/tool/agents_tests.rs | 8 ++- 16 files changed, 214 insertions(+), 71 deletions(-) create mode 100644 codex-rs/core/src/context/inter_agent_completion_message.rs diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index bd982b3c6f93..64e4ca77007e 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -7,6 +7,7 @@ use crate::agent::status::is_final; use crate::codex_thread::ThreadConfigSnapshot; use crate::config::Config; 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::shell_snapshot::ShellSnapshot; @@ -434,7 +435,6 @@ impl AgentControl { return; }; let child_thread = state.get_thread(child_thread_id).await.ok(); - let message = format_subagent_notification_message(child_reference.as_str(), &status); let child_uses_multi_agent_v2 = match child_thread.as_ref() { Some(child_thread) => { child_thread.multi_agent_version() == Some(MultiAgentVersion::V2) @@ -452,6 +452,13 @@ impl AgentControl { 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, @@ -464,6 +471,7 @@ impl AgentControl { .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; }; diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index ac42984b5857..ae681c1590db 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -1990,10 +1990,12 @@ async fn multi_agent_v2_completion_queues_message_for_direct_parent() { ) .await; - let expected_message = crate::session_prefix::format_subagent_notification_message( - tester_path.as_str(), + let expected_message = crate::session_prefix::format_inter_agent_completion_message( + worker_path.clone(), + tester_path.clone(), &AgentStatus::Completed(Some("done".to_string())), - ); + ) + .expect("completed status should render"); let expected = ( worker_thread_id, Op::InterAgentCommunication { diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 1dd4117e6545..04d418117e68 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -210,6 +210,7 @@ You can decide how much context you want to propagate to your sub-agents with th You will receive messages in the analysis channel in the form: ``` Message Type: MESSAGE | FINAL_ANSWER +Task name: Sender: Payload: @@ -228,7 +229,7 @@ When you provide a response in the final channel, that content is immediately de You will receive messages in the analysis channel in the form: ``` Message Type: NEW_TASK | MESSAGE | FINAL_ANSWER -Task name: # only for NEW_TASK -- this determines your identity +Task name: Sender: Payload: diff --git a/codex-rs/core/src/context/inter_agent_completion_message.rs b/codex-rs/core/src/context/inter_agent_completion_message.rs new file mode 100644 index 000000000000..b31e27e1ad3c --- /dev/null +++ b/codex-rs/core/src/context/inter_agent_completion_message.rs @@ -0,0 +1,41 @@ +use codex_protocol::AgentPath; + +use super::ContextualUserFragment; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct InterAgentCompletionMessage { + task_name: AgentPath, + sender: AgentPath, + payload: String, +} + +impl InterAgentCompletionMessage { + pub(crate) fn new(task_name: AgentPath, sender: AgentPath, payload: impl Into) -> Self { + Self { + task_name, + sender, + payload: payload.into(), + } + } +} + +impl ContextualUserFragment for InterAgentCompletionMessage { + fn role(&self) -> &'static str { + "assistant" + } + + fn markers(&self) -> (&'static str, &'static str) { + Self::type_markers() + } + + fn type_markers() -> (&'static str, &'static str) { + ("", "") + } + + fn body(&self) -> String { + format!( + "Message Type: FINAL_ANSWER\nTask name: {}\nSender: {}\nPayload:\n{}", + self.task_name, self.sender, self.payload, + ) + } +} diff --git a/codex-rs/core/src/context/mod.rs b/codex-rs/core/src/context/mod.rs index a470c3afb7b9..bdc90ad250a7 100644 --- a/codex-rs/core/src/context/mod.rs +++ b/codex-rs/core/src/context/mod.rs @@ -10,6 +10,7 @@ mod environment_context; mod guardian_followup_review_reminder; mod hook_additional_context; mod image_generation_instructions; +mod inter_agent_completion_message; mod internal_model_context; mod legacy_apply_patch_exec_command_warning; mod legacy_model_mismatch_warning; @@ -46,6 +47,7 @@ pub(crate) use guardian_followup_review_reminder::GuardianFollowupReviewReminder pub(crate) use hook_additional_context::HookAdditionalContext; pub(crate) use image_generation_instructions::ImageGenerationInstructions; pub use image_generation_instructions::extension_image_generation_output_hint; +pub(crate) use inter_agent_completion_message::InterAgentCompletionMessage; pub use internal_model_context::InternalContextSource; pub use internal_model_context::InternalModelContextFragment; pub use internal_model_context::InvalidInternalContextSource; diff --git a/codex-rs/core/src/guardian/prompt.rs b/codex-rs/core/src/guardian/prompt.rs index c40a03bf5ab2..1031464ca26c 100644 --- a/codex-rs/core/src/guardian/prompt.rs +++ b/codex-rs/core/src/guardian/prompt.rs @@ -1,7 +1,7 @@ use std::collections::HashMap; -use codex_protocol::models::AgentMessageInputContent; use codex_protocol::models::ResponseItem; +use codex_protocol::models::plaintext_agent_message_content; use codex_protocol::protocol::GuardianRiskLevel; use codex_protocol::protocol::GuardianUserAuthorization; use codex_protocol::user_input::UserInput; @@ -455,20 +455,10 @@ pub(crate) fn collect_guardian_transcript_entries( } ResponseItem::AgentMessage { author, content, .. - } => { - let text = content - .iter() - .filter_map(|content| match content { - AgentMessageInputContent::InputText { text } => Some(text.as_str()), - AgentMessageInputContent::EncryptedContent { .. } => None, - }) - .collect::>() - .join("\n"); - (!text.trim().is_empty()).then(|| GuardianTranscriptEntry { - kind: GuardianTranscriptEntryKind::Assistant, - text: format!("Agent message from {author}:\n{text}"), - }) - } + } => plaintext_agent_message_content(content).map(|text| GuardianTranscriptEntry { + kind: GuardianTranscriptEntryKind::Assistant, + text: format!("Agent message from {author}:\n{text}"), + }), ResponseItem::LocalShellCall { action, .. } => serialized_entry( GuardianTranscriptEntryKind::Tool("tool shell call".to_string()), serde_json::to_string(action).ok(), diff --git a/codex-rs/core/src/realtime_context.rs b/codex-rs/core/src/realtime_context.rs index 8c2c5b84883d..857389d54256 100644 --- a/codex-rs/core/src/realtime_context.rs +++ b/codex-rs/core/src/realtime_context.rs @@ -4,8 +4,8 @@ use crate::session::session::Session; use chrono::Utc; use codex_exec_server::LOCAL_FS; use codex_git_utils::resolve_root_git_project_for_trust; -use codex_protocol::models::AgentMessageInputContent; use codex_protocol::models::ResponseItem; +use codex_protocol::models::plaintext_agent_message_content; use codex_thread_store::ListThreadsParams; use codex_thread_store::SortDirection; use codex_thread_store::StoredThread; @@ -243,16 +243,10 @@ fn build_current_thread_section(items: &[ResponseItem]) -> Option { ResponseItem::AgentMessage { author, content, .. } => { - let text = content - .iter() - .filter_map(|content| match content { - AgentMessageInputContent::InputText { text } => Some(text.as_str()), - AgentMessageInputContent::EncryptedContent { .. } => None, - }) - .collect::>() - .join("\n"); - if text.trim().is_empty() || current_user.is_empty() && current_assistant.is_empty() - { + let Some(text) = plaintext_agent_message_content(content) else { + continue; + }; + if current_user.is_empty() && current_assistant.is_empty() { continue; } current_assistant.push(format!("Agent message from {author}:\n{text}")); diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 8e67020883f7..aaff69b2fea0 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -36,7 +36,7 @@ use crate::exec_policy::ExecPolicyManager; use crate::image_preparation::prepare_response_items; use crate::parse_turn_item; use crate::realtime_conversation::RealtimeConversationManager; -use crate::session_prefix::format_subagent_notification_message; +use crate::session_prefix::format_inter_agent_completion_message; use crate::skills::SkillRenderSideEffects; use crate::skills_load_input_from_config; use crate::turn_metadata::TurnMetadataState; @@ -1790,7 +1790,13 @@ impl Session { return; }; - let message = format_subagent_notification_message(child_agent_path.as_str(), &status); + let Some(message) = format_inter_agent_completion_message( + parent_agent_path.clone(), + child_agent_path.clone(), + &status, + ) else { + return; + }; // `communication` owns the message. Keep a second copy only when the // recorder will actually need it after parent delivery succeeds. let trace_message = self diff --git a/codex-rs/core/src/session_prefix.rs b/codex-rs/core/src/session_prefix.rs index b50d1f9070a6..4e6b9da94bee 100644 --- a/codex-rs/core/src/session_prefix.rs +++ b/codex-rs/core/src/session_prefix.rs @@ -1,6 +1,8 @@ +use codex_protocol::AgentPath; use codex_protocol::protocol::AgentStatus; use crate::context::ContextualUserFragment; +use crate::context::InterAgentCompletionMessage; use crate::context::SubagentNotification; // Helpers for model-visible session state markers that are stored in user-role @@ -14,6 +16,22 @@ pub(crate) fn format_subagent_notification_message( SubagentNotification::new(agent_reference, status.clone()).render() } +pub(crate) fn format_inter_agent_completion_message( + task_name: AgentPath, + sender: AgentPath, + status: &AgentStatus, +) -> Option { + let payload = match status { + AgentStatus::Completed(Some(message)) => message.clone(), + AgentStatus::Completed(None) => String::new(), + AgentStatus::Errored(error) => format!("Agent errored: {error}"), + AgentStatus::Shutdown => "Agent shut down.".to_string(), + AgentStatus::NotFound => "Agent was not found.".to_string(), + AgentStatus::PendingInit | AgentStatus::Running | AgentStatus::Interrupted => return None, + }; + Some(InterAgentCompletionMessage::new(task_name, sender, payload).render()) +} + pub(crate) fn format_subagent_context_line( agent_reference: &str, agent_nickname: Option<&str>, 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 88a8f07dcb08..c0a7c853b204 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -5,7 +5,7 @@ use crate::config::DEFAULT_AGENT_MAX_DEPTH; use crate::function_tool::FunctionCallError; use crate::init_state_db; use crate::session::tests::make_session_and_context; -use crate::session_prefix::format_subagent_notification_message; +use crate::session_prefix::format_inter_agent_completion_message; use crate::thread_manager::thread_store_from_config; use crate::tools::context::ToolOutput; use crate::tools::handlers::multi_agents_v2::FollowupTaskHandler as FollowupTaskHandlerV2; @@ -2043,14 +2043,18 @@ async fn multi_agent_v2_followup_task_completion_notifies_parent_on_every_turn() ) .await; - let first_notification = format_subagent_notification_message( - worker_path.as_str(), + let first_notification = format_inter_agent_completion_message( + AgentPath::root(), + worker_path.clone(), &AgentStatus::Completed(Some("first done".to_string())), - ); - let second_notification = format_subagent_notification_message( - worker_path.as_str(), + ) + .expect("completed status should render"); + let second_notification = format_inter_agent_completion_message( + AgentPath::root(), + worker_path.clone(), &AgentStatus::Completed(Some("second done".to_string())), - ); + ) + .expect("completed status should render"); let notifications = timeout(Duration::from_secs(5), async { loop { diff --git a/codex-rs/core/tests/suite/subagent_notifications.rs b/codex-rs/core/tests/suite/subagent_notifications.rs index 40ac2c797209..3a1a6978c101 100644 --- a/codex-rs/core/tests/suite/subagent_notifications.rs +++ b/codex-rs/core/tests/suite/subagent_notifications.rs @@ -1092,10 +1092,16 @@ async fn encrypted_multi_agent_v2_spawn_sends_agent_message_to_child() -> Result "type": "agent_message", "author": "/root", "recipient": "/root/worker", - "content": [{ - "type": "encrypted_content", - "encrypted_content": encrypted_message, - }], + "content": [ + { + "type": "input_text", + "text": "Message Type: NEW_TASK\nTask name: /root/worker\nSender: /root\nPayload:\n", + }, + { + "type": "encrypted_content", + "encrypted_content": encrypted_message, + }, + ], })] ); @@ -1133,7 +1139,7 @@ async fn plaintext_multi_agent_v2_completion_sends_agent_message() -> Result<()> mount_sse_once_match( &server, |req: &wiremock::Request| { - body_contains(req, SPAWN_CALL_ID) && !body_contains(req, "") + body_contains(req, SPAWN_CALL_ID) && !body_contains(req, "Message Type: FINAL_ANSWER") }, sse(vec![ ev_response_created("resp-parent-2"), @@ -1142,14 +1148,15 @@ async fn plaintext_multi_agent_v2_completion_sends_agent_message() -> Result<()> ]), ) .await; - let notification = "\n{\"agent_path\":\"/root/worker\",\"status\":{\"completed\":\"child done\"}}\n"; + let notification = + "Message Type: FINAL_ANSWER\nTask name: /root\nSender: /root/worker\nPayload:\nchild done"; // If the child is still running when the parent turn starts, wait_agent blocks // until mailbox delivery. The follow-up request must then contain that delivery. mount_sse_once_match( &server, |req: &wiremock::Request| { body_contains(req, TURN_2_NO_WAIT_PROMPT) - && !body_contains(req, "") + && !body_contains(req, "Message Type: FINAL_ANSWER") }, sse(vec![ ev_response_created("resp-parent-3"), @@ -1162,7 +1169,7 @@ async fn plaintext_multi_agent_v2_completion_sends_agent_message() -> Result<()> &server, |req: &wiremock::Request| { body_contains(req, TURN_2_NO_WAIT_PROMPT) - && body_contains(req, "") + && body_contains(req, "Message Type: FINAL_ANSWER") }, sse(vec![ ev_response_created("resp-parent-4"), diff --git a/codex-rs/ext/web-search/src/history.rs b/codex-rs/ext/web-search/src/history.rs index 94f9f3977276..4940ad6bc42f 100644 --- a/codex-rs/ext/web-search/src/history.rs +++ b/codex-rs/ext/web-search/src/history.rs @@ -1,9 +1,9 @@ use codex_api::SearchInput; use codex_core::parse_turn_item; use codex_protocol::items::TurnItem; -use codex_protocol::models::AgentMessageInputContent; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; +use codex_protocol::models::plaintext_agent_message_content; use codex_tools::retain_tail_from_last_n_user_messages; use codex_tools::truncate_assistant_output_text_to_token_budget; @@ -34,15 +34,7 @@ fn push_visible_message(messages: &mut Vec, item: &ResponseItem) { ResponseItem::AgentMessage { author, content, .. } => { - let text = content - .iter() - .filter_map(|content| match content { - AgentMessageInputContent::InputText { text } => Some(text.as_str()), - AgentMessageInputContent::EncryptedContent { .. } => None, - }) - .collect::>() - .join("\n"); - if !text.trim().is_empty() { + if let Some(text) = plaintext_agent_message_content(content) { messages.push(ResponseItem::Message { id: None, role: ASSISTANT_ROLE.to_string(), diff --git a/codex-rs/protocol/src/models.rs b/codex-rs/protocol/src/models.rs index 9d4647b30b21..8433e1ab5f70 100644 --- a/codex-rs/protocol/src/models.rs +++ b/codex-rs/protocol/src/models.rs @@ -723,6 +723,20 @@ pub enum AgentMessageInputContent { EncryptedContent { encrypted_content: String }, } +/// Returns the locally readable text when an agent message is entirely plaintext. +pub fn plaintext_agent_message_content(content: &[AgentMessageInputContent]) -> Option { + let mut text_parts = Vec::with_capacity(content.len()); + for part in content { + match part { + AgentMessageInputContent::InputText { text } => text_parts.push(text.as_str()), + AgentMessageInputContent::EncryptedContent { .. } => return None, + } + } + + let text = text_parts.join("\n"); + (!text.trim().is_empty()).then_some(text) +} + #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, JsonSchema, TS)] #[serde(rename_all = "lowercase")] pub enum ImageDetail { @@ -1699,6 +1713,20 @@ mod tests { 1, 122, 94, 171, 63, 0, 0, 0, 0, 73, 69, 78, 68, 174, 66, 96, 130, ]; + #[test] + fn plaintext_agent_message_content_rejects_mixed_encrypted_content() { + let content = vec![ + AgentMessageInputContent::InputText { + text: "Message Type: MESSAGE\nPayload:\n".to_string(), + }, + AgentMessageInputContent::EncryptedContent { + encrypted_content: "encrypted-payload".to_string(), + }, + ]; + + assert_eq!(plaintext_agent_message_content(&content), None); + } + #[test] fn response_input_message_conversion_preserves_phase() { let item = ResponseItem::from(ResponseInputItem::Message { diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 92b8194b4762..ef9f9c98810f 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -733,17 +733,32 @@ impl InterAgentCommunication { pub fn to_model_input_item(&self) -> ResponseItem { let content = match &self.encrypted_content { - Some(encrypted_content) => AgentMessageInputContent::EncryptedContent { - encrypted_content: encrypted_content.clone(), - }, - None => AgentMessageInputContent::InputText { + Some(encrypted_content) => { + let message_type = if self.trigger_turn { + "NEW_TASK" + } else { + "MESSAGE" + }; + vec![ + AgentMessageInputContent::InputText { + text: format!( + "Message Type: {message_type}\nTask name: {}\nSender: {}\nPayload:\n", + self.recipient, self.author + ), + }, + AgentMessageInputContent::EncryptedContent { + encrypted_content: encrypted_content.clone(), + }, + ] + } + None => vec![AgentMessageInputContent::InputText { text: self.content.clone(), - }, + }], }; ResponseItem::AgentMessage { author: self.author.to_string(), recipient: self.recipient.to_string(), - content: vec![content], + content, } } @@ -4246,6 +4261,34 @@ mod tests { ); } + #[test] + fn queued_encrypted_inter_agent_communication_renders_message_envelope() { + let communication = InterAgentCommunication::new_encrypted( + AgentPath::root().join("worker").expect("author path"), + AgentPath::root(), + Vec::new(), + "encrypted payload".to_string(), + /*trigger_turn*/ false, + ); + + assert_eq!( + communication.to_model_input_item(), + ResponseItem::AgentMessage { + author: "/root/worker".to_string(), + recipient: "/root".to_string(), + content: vec![ + AgentMessageInputContent::InputText { + text: "Message Type: MESSAGE\nTask name: /root\nSender: /root/worker\nPayload:\n" + .to_string(), + }, + AgentMessageInputContent::EncryptedContent { + encrypted_content: "encrypted payload".to_string(), + }, + ], + } + ); + } + #[test] fn session_source_from_startup_arg_normalizes_custom_values() { assert_eq!( diff --git a/codex-rs/rollout-trace/src/reducer/tool/agents.rs b/codex-rs/rollout-trace/src/reducer/tool/agents.rs index 381c475b6948..10b02f237ff9 100644 --- a/codex-rs/rollout-trace/src/reducer/tool/agents.rs +++ b/codex-rs/rollout-trace/src/reducer/tool/agents.rs @@ -771,12 +771,13 @@ fn inter_agent_message_fields(item: &ConversationItem) -> Option<(String, String return None; } if let Some(agent_message) = &item.agent_message { - let [content] = item.body.parts.as_slice() else { - return None; - }; - let message_content = match content { - ConversationPart::Text { text } => text, - ConversationPart::Encoded { label, value } if label == "encrypted_content" => value, + let message_content = match item.body.parts.as_slice() { + [ConversationPart::Text { text }] => text, + [ConversationPart::Encoded { label, value }] if label == "encrypted_content" => value, + [ + ConversationPart::Text { .. }, + ConversationPart::Encoded { label, value }, + ] if label == "encrypted_content" => value, _ => return None, }; return Some(( diff --git a/codex-rs/rollout-trace/src/reducer/tool/agents_tests.rs b/codex-rs/rollout-trace/src/reducer/tool/agents_tests.rs index db8fcc007c90..f5a4648f8021 100644 --- a/codex-rs/rollout-trace/src/reducer/tool/agents_tests.rs +++ b/codex-rs/rollout-trace/src/reducer/tool/agents_tests.rs @@ -255,7 +255,13 @@ fn sub_agent_started_activity_creates_spawn_edge() -> anyhow::Result<()> { "type": "agent_message", "author": "/root", "recipient": "/root/reviewer", - "content": [{"type": "input_text", "text": "review this"}] + "content": [ + { + "type": "input_text", + "text": "Message Type: NEW_TASK\nTask name: /root/reviewer\nSender: /root\nPayload:\n" + }, + {"type": "encrypted_content", "encrypted_content": "review this"} + ] })], )?;