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
```
This commit is contained in:
jif-oai
2026-03-23 18:53:54 +00:00
committed by GitHub
Unverified
parent e838645fa2
commit 37ac0c093c
17 changed files with 994 additions and 66 deletions
+62 -14
View File
@@ -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<UserInput>,
) -> CodexResult<String> {
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<String> {
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<String> {
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<ThreadManagerState>,
result: CodexResult<String>,
) -> CodexResult<String> {
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<String> {
+55
View File
@@ -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;
@@ -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<AgentPath>,
content: String,
}
impl InterAgentInstruction {
pub(crate) fn new(
author: AgentPath,
recipient: AgentPath,
other_recipients: Vec<AgentPath>,
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::<Vec<_>>()
.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: ")
}
}
+1
View File
@@ -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;
+56 -25
View File
@@ -798,6 +798,7 @@ pub(crate) struct Session {
pending_mcp_server_refresh_config: Mutex<Option<McpServerRefreshConfig>>,
pub(crate) conversation: Arc<RealtimeConversationManager>,
pub(crate) active_turn: Mutex<Option<ActiveTurn>>,
idle_pending_input: Mutex<Vec<ResponseInputItem>>,
pub(crate) guardian_review_session: GuardianReviewSessionManager,
pub(crate) services: SessionServices,
js_repl: Arc<JsReplHandle>,
@@ -809,6 +810,7 @@ pub(crate) struct TurnSkillsContext {
pub(crate) outcome: Arc<SkillLoadOutcome>,
pub(crate) implicit_invocation_seen_skills: Arc<Mutex<HashSet<String>>>,
}
impl TurnSkillsContext {
pub(crate) fn new(outcome: Arc<SkillLoadOutcome>) -> 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<ResponseInputItem>,
@@ -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<ResponseInputItem>) {
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<ResponseInputItem> {
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<ModelClientSession>,
cancellation_token: CancellationToken,
) -> Option<String> {
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::<Vec<_>>();
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<String> = 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<String> = 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.
@@ -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()
@@ -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;
+28
View File
@@ -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;
+88 -13
View File
@@ -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<ResponseItem> = 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<String> {
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<String> {
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<String> {
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<PathBuf> {
@@ -198,3 +261,15 @@ impl CodexThread {
Ok(*guard)
}
}
fn pending_message_input_item(message: &ResponseItem) -> CodexResult<ResponseInputItem> {
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(),
)),
}
}
+11 -4
View File
@@ -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<usize> {
@@ -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<ResponseItem>) -> 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")]);
+1 -1
View File
@@ -72,7 +72,7 @@ fn parse_agent_message(
let mut content: Vec<AgentMessageContent> = 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() });
}
_ => {
+37
View File
@@ -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::<Vec<_>>();
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();
+30
View File
@@ -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<T: SessionTask>(
self: &Arc<Self>,
turn_context: Arc<TurnContext>,
input: Vec<UserInput>,
task: T,
) {
let task: Arc<dyn SessionTask> = 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<Self>) {
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<Self>, 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);
+27
View File
@@ -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<String> {
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<String> {
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<Arc<CodexThread>> {
self.threads.write().await.remove(thread_id)
@@ -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
@@ -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<crate::codex::Session>,
@@ -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<Self>,
_session: Arc<SessionTaskContext>,
_ctx: Arc<TurnContext>,
_input: Vec<UserInput>,
cancellation_token: CancellationToken,
) -> Option<String> {
cancellation_token.cancelled().await;
None
}
}
fn expect_text_output<T>(output: T) -> (String, Option<bool>)
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]