openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-27 18:48:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: make task return a more complex type

Browse Source
This commit is contained in:
jif-oai
2026-03-30 18:50:26 +02:00
parent 716f7b0428
commit 1cd1fbd6b5
12 changed files with 495 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/agent/control.rs
+36
View File
@@ -527,6 +527,37 @@ impl AgentControl {
result
}
pub(crate) async fn forward_child_final_status_to_parent(
&self,
session_source: &SessionSource,
status: &AgentStatus,
) {
let Some(parent_thread_id) = thread_spawn_parent_thread_id(session_source) else {
return;
};
let Some(child_agent_path) = session_source.get_agent_path() else {
return;
};
let Some(parent_agent_path) = child_agent_path
.as_str()
.rsplit_once('/')
.and_then(|(parent, _)| AgentPath::try_from(parent).ok())
else {
return;
};
let message = format_subagent_notification_message(child_agent_path.as_str(), status);
let communication = InterAgentCommunication::new(
child_agent_path,
parent_agent_path,
Vec::new(),
message,
/*trigger_turn*/ false,
);
let _ = self
.send_inter_agent_communication(parent_thread_id, communication)
.await;
}
/// Interrupt the current task for an existing agent thread.
pub(crate) async fn interrupt_agent(&self, agent_id: ThreadId) -> CodexResult<String> {
let state = self.upgrade()?;
@@ -782,6 +813,11 @@ impl AgentControl {
child_reference: String,
child_agent_path: Option<AgentPath>,
) {
// Agent-path children now forward final status directly from `Session::send_event`,
// which keeps reused spawned agents notifying their parent on every completed turn.
if child_agent_path.is_some() {
return;
}
let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id, ..
})) = session_source
codex-rs/core/src/agent/control_tests.rs
+32 -22
View File
@@ -1180,35 +1180,46 @@ async fn multi_agent_v2_completion_ignores_dead_direct_parent() {
#[tokio::test]
async fn multi_agent_v2_completion_queues_message_for_direct_parent() {
let harness = AgentControlHarness::new().await;
let (_root_thread_id, root_thread) = harness.start_thread().await;
let (worker_thread_id, _worker_thread) = harness.start_thread().await;
let (root_thread_id, root_thread) = harness.start_thread().await;
let mut tester_config = harness.config.clone();
let _ = tester_config.features.enable(Feature::MultiAgentV2);
let tester_thread_id = harness
.manager
.start_thread(tester_config)
let worker_path = AgentPath::root().join("worker_a").expect("worker path");
let tester_path = worker_path.join("tester").expect("tester path");
let worker_thread_id = harness
.control
.spawn_agent(
tester_config.clone(),
text_input("hello worker"),
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id: root_thread_id,
depth: 1,
agent_path: Some(worker_path.clone()),
agent_nickname: None,
agent_role: Some("explorer".to_string()),
})),
)
.await
.expect("tester thread should start")
.thread_id;
.expect("worker spawn should succeed");
let tester_thread_id = harness
.control
.spawn_agent(
tester_config,
text_input("hello tester"),
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id: worker_thread_id,
depth: 2,
agent_path: Some(tester_path.clone()),
agent_nickname: None,
agent_role: Some("explorer".to_string()),
})),
)
.await
.expect("tester spawn should succeed");
let tester_thread = harness
.manager
.get_thread(tester_thread_id)
.await
.expect("tester thread should exist");
let worker_path = AgentPath::root().join("worker_a").expect("worker path");
let tester_path = worker_path.join("tester").expect("tester path");
harness.control.maybe_start_completion_watcher(
tester_thread_id,
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id: worker_thread_id,
depth: 2,
agent_path: Some(tester_path.clone()),
agent_nickname: None,
agent_role: Some("explorer".to_string()),
})),
tester_path.to_string(),
Some(tester_path.clone()),
);
let tester_turn = tester_thread.codex.session.new_default_turn().await;
tester_thread
.codex
@@ -1221,7 +1232,6 @@ 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(),
&AgentStatus::Completed(Some("done".to_string())),
codex-rs/core/src/codex.rs
+27 -25
View File
@@ -14,6 +14,7 @@ use crate::agent::AgentStatus;
use crate::agent::Mailbox;
use crate::agent::MailboxReceiver;
use crate::agent::agent_status_from_event;
use crate::agent::status::is_final;
use crate::apps::render_apps_section;
use crate::auth_env_telemetry::collect_auth_env_telemetry;
use crate::commit_attribution::commit_message_trailer_instruction;
@@ -334,6 +335,7 @@ use crate::tasks::GhostSnapshotTask;
use crate::tasks::ReviewTask;
use crate::tasks::SessionTask;
use crate::tasks::SessionTaskContext;
use crate::tasks::TaskCompletion;
use crate::tools::ToolRouter;
use crate::tools::context::SharedTurnDiffTracker;
use crate::tools::js_repl::JsReplHandle;
@@ -2641,6 +2643,15 @@ impl Session {
.await;
self.maybe_clear_realtime_handoff_for_event(&legacy_source)
.await;
if let Some(status) = agent_status_from_event(&legacy_source)
&& is_final(&status)
&& turn_context.session_source.get_agent_path().is_some()
{
self.services
.agent_control
.forward_child_final_status_to_parent(&turn_context.session_source, &status)
.await;
}
let show_raw_agent_reasoning = self.show_raw_agent_reasoning();
for legacy in legacy_source.as_legacy_events(show_raw_agent_reasoning) {
@@ -5586,9 +5597,9 @@ pub(crate) async fn run_turn(
input: Vec<UserInput>,
prewarmed_client_session: Option<ModelClientSession>,
cancellation_token: CancellationToken,
) -> Option<String> {
) -> TaskCompletion {
if input.is_empty() && !sess.has_pending_input().await {
return None;
return TaskCompletion::Completed(None);
}
let model_info = turn_context.model_info.clone();
@@ -5602,7 +5613,7 @@ pub(crate) async fn run_turn(
.is_err()
{
error!("Failed to run pre-sampling compact");
return None;
return TaskCompletion::Completed(None);
}
let skills_outcome = Some(turn_context.turn_skills.outcome.as_ref());
@@ -5632,7 +5643,7 @@ pub(crate) async fn run_turn(
.await
{
Ok(mcp_tools) => mcp_tools,
Err(_) if turn_context.apps_enabled() => return None,
Err(_) if turn_context.apps_enabled() => return TaskCompletion::Completed(None),
Err(_) => HashMap::new(),
}
} else {
@@ -5730,7 +5741,7 @@ pub(crate) async fn run_turn(
.collect::<Vec<_>>();
if run_pending_session_start_hooks(&sess, &turn_context).await {
return None;
return TaskCompletion::Completed(None);
}
let additional_contexts = if input.is_empty() {
Vec::new()
@@ -5750,7 +5761,7 @@ pub(crate) async fn run_turn(
user_prompt_submit_outcome.additional_contexts,
)
.await;
return None;
return TaskCompletion::Completed(None);
}
sess.record_user_prompt_and_emit_turn_item(turn_context.as_ref(), &input, response_item)
.await;
@@ -5913,7 +5924,7 @@ pub(crate) async fn run_turn(
.await
.is_err()
{
return None;
return TaskCompletion::Completed(None);
}
continue;
}
@@ -6025,15 +6036,10 @@ pub(crate) async fn run_turn(
}
}
if let Some(message) = abort_message {
sess.send_event(
&turn_context,
EventMsg::Error(ErrorEvent {
message,
codex_error_info: None,
}),
)
.await;
return None;
return TaskCompletion::Failed(ErrorEvent {
message,
codex_error_info: None,
});
}
break;
}
@@ -6051,25 +6057,21 @@ pub(crate) async fn run_turn(
if state.history.replace_last_turn_images("Invalid image") {
continue;
}
let event = EventMsg::Error(ErrorEvent {
let error = ErrorEvent {
message: "Invalid image in your last message. Please remove it and try again."
.to_string(),
codex_error_info: Some(CodexErrorInfo::BadRequest),
});
sess.send_event(&turn_context, event).await;
break;
};
return TaskCompletion::Failed(error);
}
Err(e) => {
info!("Turn error: {e:#}");
let event = EventMsg::Error(e.to_error_event(/*message_prefix*/ None));
sess.send_event(&turn_context, event).await;
// let the user continue the conversation
break;
return TaskCompletion::Failed(e.to_error_event(/*message_prefix*/ None));
}
}
}
last_agent_message
TaskCompletion::Completed(last_agent_message)
}
async fn run_pre_sampling_compact(
codex-rs/core/src/codex_tests.rs
+61 -5
View File
@@ -35,6 +35,7 @@ use tracing::Span;
use crate::protocol::CompactedItem;
use crate::protocol::CreditsSnapshot;
use crate::protocol::ErrorEvent;
use crate::protocol::InitialHistory;
use crate::protocol::NetworkApprovalProtocol;
use crate::protocol::RateLimitSnapshot;
@@ -52,6 +53,7 @@ use crate::rollout::recorder::RolloutRecorder;
use crate::rollout::recorder::RolloutRecorderParams;
use crate::state::TaskKind;
use crate::tasks::SessionTask;
use crate::tasks::TaskCompletion;
use crate::tasks::SessionTaskContext;
use crate::tools::ToolRouter;
use crate::tools::context::FunctionToolOutput;
@@ -3138,13 +3140,13 @@ async fn spawn_task_turn_span_inherits_dispatch_trace_context() {
_ctx: Arc<TurnContext>,
_input: Vec<UserInput>,
_cancellation_token: CancellationToken,
) -> Option<String> {
) -> TaskCompletion {
let mut trace = self
.captured_trace
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
*trace = current_span_w3c_trace_context();
None
TaskCompletion::Completed(None)
}
}
@@ -4380,6 +4382,9 @@ struct NeverEndingTask {
listen_to_cancellation_token: bool,
}
#[derive(Clone, Copy)]
struct FailingTask;
#[async_trait::async_trait]
impl SessionTask for NeverEndingTask {
fn kind(&self) -> TaskKind {
@@ -4396,10 +4401,10 @@ impl SessionTask for NeverEndingTask {
_ctx: Arc<TurnContext>,
_input: Vec<UserInput>,
cancellation_token: CancellationToken,
) -> Option<String> {
) -> TaskCompletion {
if self.listen_to_cancellation_token {
cancellation_token.cancelled().await;
return None;
return TaskCompletion::Completed(None);
}
loop {
sleep(Duration::from_secs(60)).await;
@@ -4407,6 +4412,30 @@ impl SessionTask for NeverEndingTask {
}
}
#[async_trait::async_trait]
impl SessionTask for FailingTask {
fn kind(&self) -> TaskKind {
TaskKind::Regular
}
fn span_name(&self) -> &'static str {
"session_task.failing"
}
async fn run(
self: Arc<Self>,
_session: Arc<SessionTaskContext>,
_ctx: Arc<TurnContext>,
_input: Vec<UserInput>,
_cancellation_token: CancellationToken,
) -> TaskCompletion {
TaskCompletion::Failed(ErrorEvent {
message: "boom".to_string(),
codex_error_info: None,
})
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[test_log::test]
async fn abort_regular_task_emits_turn_aborted_only() {
@@ -4441,6 +4470,30 @@ async fn abort_regular_task_emits_turn_aborted_only() {
assert!(rx.try_recv().is_err());
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[test_log::test]
async fn failing_regular_task_emits_error_without_turn_complete() {
let (sess, tc, rx) = make_session_and_context_with_rx().await;
let input = vec![UserInput::Text {
text: "hello".to_string(),
text_elements: Vec::new(),
}];
sess.spawn_task(Arc::clone(&tc), input, FailingTask).await;
let evt = tokio::time::timeout(std::time::Duration::from_secs(2), rx.recv())
.await
.expect("timeout waiting for event")
.expect("event");
match evt.msg {
EventMsg::Error(ErrorEvent {
message,
codex_error_info: None,
}) => assert_eq!(message, "boom"),
other => panic!("unexpected event: {other:?}"),
}
assert!(rx.try_recv().is_err());
}
#[tokio::test]
async fn abort_gracefully_emits_turn_aborted_only() {
let (sess, tc, rx) = make_session_and_context_with_rx().await;
@@ -4502,7 +4555,10 @@ async fn task_finish_emits_turn_item_lifecycle_for_leftover_pending_user_input()
.await
.expect("inject pending input into active turn");
sess.on_task_finished(Arc::clone(&tc), /*last_agent_message*/ None)
sess.on_task_finished(
Arc::clone(&tc),
TaskCompletion::Completed(/*last_agent_message*/ None),
)
.await;
let history = sess.clone_history().await;
codex-rs/core/src/tasks/compact.rs
+3 -2
View File
@@ -1,6 +1,7 @@
use std::sync::Arc;
use super::SessionTask;
use super::TaskCompletion;
use super::SessionTaskContext;
use crate::codex::TurnContext;
use crate::state::TaskKind;
@@ -27,7 +28,7 @@ impl SessionTask for CompactTask {
ctx: Arc<TurnContext>,
input: Vec<UserInput>,
_cancellation_token: CancellationToken,
) -> Option<String> {
) -> TaskCompletion {
let session = session.clone_session();
let _ = if crate::compact::should_use_remote_compact_task(&ctx.provider) {
let _ = session.services.session_telemetry.counter(
@@ -44,6 +45,6 @@ impl SessionTask for CompactTask {
);
crate::compact::run_compact_task(session.clone(), ctx, input).await
};
None
TaskCompletion::Completed(None)
}
}
codex-rs/core/src/tasks/ghost_snapshot.rs
+3 -2
View File
@@ -3,6 +3,7 @@ use crate::protocol::EventMsg;
use crate::protocol::WarningEvent;
use crate::state::TaskKind;
use crate::tasks::SessionTask;
use crate::tasks::TaskCompletion;
use crate::tasks::SessionTaskContext;
use async_trait::async_trait;
use codex_git_utils::CreateGhostCommitOptions;
@@ -42,7 +43,7 @@ impl SessionTask for GhostSnapshotTask {
ctx: Arc<TurnContext>,
_input: Vec<UserInput>,
cancellation_token: CancellationToken,
) -> Option<String> {
) -> TaskCompletion {
tokio::task::spawn(async move {
let token = self.token;
let warnings_enabled = !ctx.ghost_snapshot.disable_warnings;
@@ -156,7 +157,7 @@ impl SessionTask for GhostSnapshotTask {
Err(err) => warn!("failed to mark ghost snapshot ready: {err}"),
}
});
None
TaskCompletion::Completed(None)
}
}
codex-rs/core/src/tasks/mod.rs
+21 -10
View File
@@ -29,6 +29,7 @@ use crate::hook_runtime::inspect_pending_input;
use crate::hook_runtime::record_additional_contexts;
use crate::hook_runtime::record_pending_input;
use crate::models_manager::manager::ModelsManager;
use crate::protocol::ErrorEvent;
use crate::protocol::EventMsg;
use crate::protocol::TurnAbortReason;
use crate::protocol::TurnAbortedEvent;
@@ -99,6 +100,11 @@ pub(crate) struct SessionTaskContext {
session: Arc<Session>,
}
pub(crate) enum TaskCompletion {
Completed(Option<String>),
Failed(ErrorEvent),
}
impl SessionTaskContext {
pub(crate) fn new(session: Arc<Session>) -> Self {
Self { session }
@@ -140,15 +146,15 @@ pub(crate) trait SessionTask: Send + Sync + 'static {
/// `ctx`, returning an optional final agent message when finished. The
/// provided `cancellation_token` is cancelled when the session requests an
/// abort; implementers should watch for it and terminate quickly once it
/// fires. Returning [`Some`] yields a final message that
/// [`Session::on_task_finished`] will emit to the client.
/// fires. The returned [`TaskCompletion`] determines the single terminal
/// event emitted by [`Session::on_task_finished`].
async fn run(
self: Arc<Self>,
session: Arc<SessionTaskContext>,
ctx: Arc<TurnContext>,
input: Vec<UserInput>,
cancellation_token: CancellationToken,
) -> Option<String>;
) -> TaskCompletion;
/// Gives the task a chance to perform cleanup after an abort.
///
@@ -209,7 +215,7 @@ impl Session {
tokio::spawn(
async move {
let ctx_for_finish = Arc::clone(&ctx);
let last_agent_message = task_for_run
let completion = task_for_run
.run(
Arc::clone(&session_ctx),
ctx,
@@ -221,7 +227,7 @@ impl Session {
sess.flush_rollout().await;
if !task_cancellation_token.is_cancelled() {
// Emit completion uniformly from spawn site so all tasks share the same lifecycle.
sess.on_task_finished(Arc::clone(&ctx_for_finish), last_agent_message)
sess.on_task_finished(Arc::clone(&ctx_for_finish), completion)
.await;
}
done_clone.notify_waiters();
@@ -314,7 +320,7 @@ impl Session {
pub async fn on_task_finished(
self: &Arc<Self>,
turn_context: Arc<TurnContext>,
last_agent_message: Option<String>,
completion: TaskCompletion,
) {
turn_context
.turn_metadata_state
@@ -431,10 +437,15 @@ impl Session {
&[("token_type", "reasoning_output"), tmp_mem],
);
}
let event = EventMsg::TurnComplete(TurnCompleteEvent {
turn_id: turn_context.sub_id.clone(),
last_agent_message,
});
let event = match completion {
TaskCompletion::Completed(last_agent_message) => {
EventMsg::TurnComplete(TurnCompleteEvent {
turn_id: turn_context.sub_id.clone(),
last_agent_message,
})
}
TaskCompletion::Failed(error) => EventMsg::Error(error),
};
self.send_event(turn_context.as_ref(), event).await;
}
codex-rs/core/src/tasks/regular.rs
+5 -2
View File
@@ -14,6 +14,7 @@ use tracing::Instrument;
use tracing::trace_span;
use super::SessionTask;
use super::TaskCompletion;
use super::SessionTaskContext;
#[derive(Default)]
@@ -41,7 +42,7 @@ impl SessionTask for RegularTask {
ctx: Arc<TurnContext>,
input: Vec<UserInput>,
cancellation_token: CancellationToken,
) -> Option<String> {
) -> TaskCompletion {
let sess = session.clone_session();
let run_turn_span = trace_span!("run_turn");
// Regular turns emit `TurnStarted` inline so first-turn lifecycle does
@@ -57,7 +58,9 @@ impl SessionTask for RegularTask {
.consume_startup_prewarm_for_regular_turn(&cancellation_token)
.await
{
SessionStartupPrewarmResolution::Cancelled => return None,
SessionStartupPrewarmResolution::Cancelled => {
return TaskCompletion::Completed(None);
}
SessionStartupPrewarmResolution::Unavailable { .. } => None,
SessionStartupPrewarmResolution::Ready(prewarmed_client_session) => {
Some(*prewarmed_client_session)
codex-rs/core/src/tasks/review.rs
+3 -2
View File
@@ -30,6 +30,7 @@ use codex_protocol::user_input::UserInput;
use std::sync::LazyLock;
use super::SessionTask;
use super::TaskCompletion;
use super::SessionTaskContext;
static REVIEW_EXIT_SUCCESS_TEMPLATE: LazyLock<Template> = LazyLock::new(|| {
@@ -64,7 +65,7 @@ impl SessionTask for ReviewTask {
ctx: Arc<TurnContext>,
input: Vec<UserInput>,
cancellation_token: CancellationToken,
) -> Option<String> {
) -> TaskCompletion {
let _ = session.session.services.session_telemetry.counter(
"codex.task.review",
/*inc*/ 1,
@@ -86,7 +87,7 @@ impl SessionTask for ReviewTask {
if !cancellation_token.is_cancelled() {
exit_review_mode(session.clone_session(), output.clone(), ctx.clone()).await;
}
None
TaskCompletion::Completed(None)
}
async fn abort(&self, session: Arc<SessionTaskContext>, ctx: Arc<TurnContext>) {
codex-rs/core/src/tasks/undo.rs
+5 -4
View File
@@ -6,6 +6,7 @@ use crate::protocol::UndoCompletedEvent;
use crate::protocol::UndoStartedEvent;
use crate::state::TaskKind;
use crate::tasks::SessionTask;
use crate::tasks::TaskCompletion;
use crate::tasks::SessionTaskContext;
use async_trait::async_trait;
use codex_git_utils::RestoreGhostCommitOptions;
@@ -41,7 +42,7 @@ impl SessionTask for UndoTask {
ctx: Arc<TurnContext>,
_input: Vec<UserInput>,
cancellation_token: CancellationToken,
) -> Option<String> {
) -> TaskCompletion {
let _ = session.session.services.session_telemetry.counter(
"codex.task.undo",
/*inc*/ 1,
@@ -65,7 +66,7 @@ impl SessionTask for UndoTask {
}),
)
.await;
return None;
return TaskCompletion::Completed(None);
}
let history = sess.clone_history().await;
@@ -90,7 +91,7 @@ impl SessionTask for UndoTask {
completed.message = Some("No ghost snapshot available to undo.".to_string());
sess.send_event(ctx.as_ref(), EventMsg::UndoCompleted(completed))
.await;
return None;
return TaskCompletion::Completed(None);
};
let commit_id = ghost_commit.id().to_string();
@@ -126,6 +127,6 @@ impl SessionTask for UndoTask {
sess.send_event(ctx.as_ref(), EventMsg::UndoCompleted(completed))
.await;
None
TaskCompletion::Completed(None)
}
}
codex-rs/core/src/tasks/user_shell.rs
+3 -2
View File
@@ -32,6 +32,7 @@ use crate::user_shell_command::user_shell_command_record_item;
use codex_sandboxing::SandboxType;
use super::SessionTask;
use super::TaskCompletion;
use super::SessionTaskContext;
use crate::codex::Session;
use codex_protocol::models::ResponseInputItem;
@@ -78,7 +79,7 @@ impl SessionTask for UserShellCommandTask {
turn_context: Arc<TurnContext>,
_input: Vec<UserInput>,
cancellation_token: CancellationToken,
) -> Option<String> {
) -> TaskCompletion {
execute_user_shell_command(
session.clone_session(),
turn_context,
@@ -87,7 +88,7 @@ impl SessionTask for UserShellCommandTask {
UserShellCommandMode::StandaloneTurn,
)
.await;
None
TaskCompletion::Completed(None)
}
}
codex-rs/core/src/tools/handlers/multi_agents_tests.rs
+296 -2
View File
@@ -9,6 +9,7 @@ use crate::config::types::ShellEnvironmentPolicy;
use crate::function_tool::FunctionCallError;
use crate::protocol::AgentStatus;
use crate::protocol::AskForApproval;
use crate::protocol::ErrorEvent;
use crate::protocol::EventMsg;
use crate::protocol::FileSystemSandboxPolicy;
use crate::protocol::NetworkSandboxPolicy;
@@ -19,6 +20,7 @@ use crate::protocol::SubAgentSource;
use crate::protocol::TurnCompleteEvent;
use crate::state::TaskKind;
use crate::tasks::SessionTask;
use crate::tasks::TaskCompletion;
use crate::tasks::SessionTaskContext;
use crate::tools::context::ToolOutput;
use crate::tools::handlers::multi_agents_v2::AssignTaskHandler as AssignTaskHandlerV2;
@@ -112,6 +114,9 @@ fn history_contains_inter_agent_communication(
#[derive(Clone, Copy)]
struct NeverEndingTask;
#[derive(Clone, Copy)]
struct FailingTask;
#[async_trait::async_trait]
impl SessionTask for NeverEndingTask {
fn kind(&self) -> TaskKind {
@@ -128,9 +133,33 @@ impl SessionTask for NeverEndingTask {
_ctx: Arc<TurnContext>,
_input: Vec<UserInput>,
cancellation_token: CancellationToken,
) -> Option<String> {
) -> TaskCompletion {
cancellation_token.cancelled().await;
None
TaskCompletion::Completed(None)
}
}
#[async_trait::async_trait]
impl SessionTask for FailingTask {
fn kind(&self) -> TaskKind {
TaskKind::Regular
}
fn span_name(&self) -> &'static str {
"session_task.multi_agent_failing"
}
async fn run(
self: Arc<Self>,
_session: Arc<SessionTaskContext>,
_ctx: Arc<TurnContext>,
_input: Vec<UserInput>,
_cancellation_token: CancellationToken,
) -> TaskCompletion {
TaskCompletion::Failed(ErrorEvent {
message: "boom".to_string(),
codex_error_info: None,
})
}
}
@@ -1119,6 +1148,271 @@ async fn multi_agent_v2_assign_task_interrupts_busy_child_without_losing_message
.expect("shutdown should submit");
}
#[tokio::test]
async fn multi_agent_v2_assign_task_completion_notifies_parent_after_reuse() {
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);
SpawnAgentHandlerV2
.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 worker_path = AgentPath::try_from("/root/worker").expect("agent path");
let first_turn = thread.codex.session.new_default_turn().await;
thread
.codex
.session
.send_event(
first_turn.as_ref(),
EventMsg::TurnComplete(TurnCompleteEvent {
turn_id: first_turn.sub_id.clone(),
last_agent_message: Some("done once".to_string()),
}),
)
.await;
AssignTaskHandlerV2
.handle(invocation(
session,
turn,
"assign_task",
function_payload(json!({
"target": agent_id.to_string(),
"items": [{"type": "text", "text": "continue"}]
})),
))
.await
.expect("assign_task should succeed");
let second_turn = thread.codex.session.new_default_turn().await;
thread
.codex
.session
.send_event(
second_turn.as_ref(),
EventMsg::TurnComplete(TurnCompleteEvent {
turn_id: second_turn.sub_id.clone(),
last_agent_message: Some("done twice".to_string()),
}),
)
.await;
let expected_first = (
root.thread_id,
Op::InterAgentCommunication {
communication: InterAgentCommunication::new(
worker_path.clone(),
AgentPath::root(),
Vec::new(),
crate::session_prefix::format_subagent_notification_message(
worker_path.as_str(),
&AgentStatus::Completed(Some("done once".to_string())),
),
/*trigger_turn*/ false,
),
},
);
let expected_second = (
root.thread_id,
Op::InterAgentCommunication {
communication: InterAgentCommunication::new(
worker_path,
AgentPath::root(),
Vec::new(),
crate::session_prefix::format_subagent_notification_message(
"/root/worker",
&AgentStatus::Completed(Some("done twice".to_string())),
),
/*trigger_turn*/ false,
),
},
);
timeout(Duration::from_secs(5), async {
loop {
let captured = manager.captured_ops();
if captured.contains(&expected_first) && captured.contains(&expected_second) {
break;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("reused worker should notify parent on both completed turns");
}
#[tokio::test]
async fn multi_agent_v2_failed_turn_forwards_only_error_status_after_reuse() {
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);
SpawnAgentHandlerV2
.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 worker_path = AgentPath::try_from("/root/worker").expect("agent path");
let first_turn = thread.codex.session.new_default_turn().await;
thread
.codex
.session
.send_event(
first_turn.as_ref(),
EventMsg::TurnComplete(TurnCompleteEvent {
turn_id: first_turn.sub_id.clone(),
last_agent_message: Some("done once".to_string()),
}),
)
.await;
AssignTaskHandlerV2
.handle(invocation(
session,
turn,
"assign_task",
function_payload(json!({
"target": agent_id.to_string(),
"items": [{"type": "text", "text": "continue"}]
})),
))
.await
.expect("assign_task should succeed");
let second_turn = thread.codex.session.new_default_turn().await;
thread
.codex
.session
.spawn_task(
Arc::clone(&second_turn),
vec![UserInput::Text {
text: "continue".to_string(),
text_elements: Vec::new(),
}],
FailingTask,
)
.await;
let expected_first = (
root.thread_id,
Op::InterAgentCommunication {
communication: InterAgentCommunication::new(
worker_path.clone(),
AgentPath::root(),
Vec::new(),
crate::session_prefix::format_subagent_notification_message(
worker_path.as_str(),
&AgentStatus::Completed(Some("done once".to_string())),
),
/*trigger_turn*/ false,
),
},
);
let expected_error = (
root.thread_id,
Op::InterAgentCommunication {
communication: InterAgentCommunication::new(
worker_path.clone(),
AgentPath::root(),
Vec::new(),
crate::session_prefix::format_subagent_notification_message(
worker_path.as_str(),
&AgentStatus::Errored("boom".to_string()),
),
/*trigger_turn*/ false,
),
},
);
let unexpected_completed = (
root.thread_id,
Op::InterAgentCommunication {
communication: InterAgentCommunication::new(
worker_path,
AgentPath::root(),
Vec::new(),
crate::session_prefix::format_subagent_notification_message(
"/root/worker",
&AgentStatus::Completed(None),
),
/*trigger_turn*/ false,
),
},
);
timeout(Duration::from_secs(5), async {
loop {
let captured = manager.captured_ops();
if captured.contains(&expected_first)
&& captured.contains(&expected_error)
&& !captured.contains(&unexpected_completed)
{
break;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("failed reused worker turn should only notify parent with the errored status");
}
#[tokio::test]
async fn multi_agent_v2_spawn_includes_agent_id_key_when_named() {
let (mut session, mut turn) = make_session_and_context().await;