mirror of
https://github.com/openai/codex.git
synced 2026-11-05 20:10:33 +00:00
Add host reporting for terminal agent turn outcomes (#59725)
## What changed Expose `AgentTurnOutcomeReporter` for hosts to install through `StartThreadOptions::thread_extension_init`. Await its callback after emitting the terminal turn event, independently of native parent notifications, so hosts can handle result delivery and retries. Track outstanding outcome handoffs and wait for them during graceful shutdown before running thread-stop hooks. Prevent new turns from starting once handoff registration closes. ## Testing Add integration coverage for completed and errored outcomes, turn lineage, queued follow-up turns, and shutdown waiting for handoffs. Extend Guardian tests to verify that strict safety stops reach the reporter while default interruptions remain silent. GitOrigin-RevId: 0d825f6208f0250969734488881f32fd1dcd7509
This commit is contained in:
@@ -232,6 +232,7 @@ pub struct DeliveryReceipt {
|
||||
pub submission_id: String,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AgentTurnOutcome {
|
||||
pub thread_id: ThreadId,
|
||||
pub turn_id: String,
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
//! Lets an embedding observe completed agent work for host-managed result delivery.
|
||||
|
||||
use futures::future::BoxFuture;
|
||||
|
||||
use crate::agent::api::AgentTurnOutcome;
|
||||
|
||||
/// Observes the turn outcomes eligible for native parent result delivery.
|
||||
///
|
||||
/// Install an `Arc<dyn AgentTurnOutcomeReporter>` in
|
||||
/// `StartThreadOptions::thread_extension_init` for each thread to observe,
|
||||
/// including on cold resume. Native parent notifications are independent.
|
||||
/// Core awaits the callback after emitting the terminal turn event. Graceful
|
||||
/// shutdown waits for outstanding handoffs before running thread-stop hooks.
|
||||
/// Hand off promptly to host-owned delivery. Waiting for this thread to shut
|
||||
/// down inside the callback would deadlock. The host owns transport and retries.
|
||||
/// The embedding can supply the session ID from its thread metadata.
|
||||
pub trait AgentTurnOutcomeReporter: Send + Sync {
|
||||
fn report(&self, outcome: AgentTurnOutcome) -> BoxFuture<'_, ()>;
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
mod admission;
|
||||
mod completion;
|
||||
mod invocation;
|
||||
mod mcp_metadata;
|
||||
mod ownership;
|
||||
@@ -9,6 +10,7 @@ mod usage;
|
||||
|
||||
pub use admission::AgentTurnAdmission;
|
||||
pub use admission::AgentTurnAdmissionRequest;
|
||||
pub use completion::AgentTurnOutcomeReporter;
|
||||
pub use invocation::HostToolInvocationLookup;
|
||||
pub(crate) use invocation::InvocationReservation;
|
||||
pub use invocation::ToolInvocationContext;
|
||||
|
||||
@@ -164,6 +164,7 @@ pub use agent::host::AgentRuntimeOwnership;
|
||||
pub use agent::host::AgentSessionReader;
|
||||
pub use agent::host::AgentTurnAdmission;
|
||||
pub use agent::host::AgentTurnAdmissionRequest;
|
||||
pub use agent::host::AgentTurnOutcomeReporter;
|
||||
pub use agent::host::AgentUsageAccounting;
|
||||
pub use agent::host::AgentUsageReport;
|
||||
pub use agent::host::HostThreadSpawnReservation;
|
||||
|
||||
@@ -308,6 +308,9 @@ pub async fn set_thread_memory_mode(sess: &Arc<Session>, sub_id: String, mode: T
|
||||
}
|
||||
|
||||
pub(super) async fn shutdown_session_runtime(sess: &Arc<Session>) {
|
||||
if let Some(reporting) = &sess.outcome_reporting {
|
||||
reporting.handoffs.close();
|
||||
}
|
||||
let startup_prewarm = {
|
||||
let mut state = sess.state.lock().await;
|
||||
// Stop admission and take the current warmup together so resume cannot replace it.
|
||||
@@ -319,6 +322,9 @@ pub(super) async fn shutdown_session_runtime(sess: &Arc<Session>) {
|
||||
}
|
||||
let _ = sess.conversation.shutdown().await;
|
||||
sess.abort_all_tasks(TurnAbortReason::Interrupted).await;
|
||||
if let Some(reporting) = &sess.outcome_reporting {
|
||||
reporting.handoffs.wait().await;
|
||||
}
|
||||
super::async_compaction::shutdown(sess).await;
|
||||
let shell_snapshot_prewarm = sess.state.lock().await.shell_snapshot_prewarm.take();
|
||||
if let Some(shell_snapshot_prewarm) = shell_snapshot_prewarm {
|
||||
|
||||
@@ -2697,7 +2697,7 @@ impl Session {
|
||||
.track_guardian_session_event(self.thread_id, &event);
|
||||
}
|
||||
self.send_event_raw(event).await;
|
||||
self.maybe_notify_parent_of_terminal_turn(turn_context, &legacy_source)
|
||||
self.maybe_publish_terminal_agent_result(turn_context, &legacy_source)
|
||||
.await;
|
||||
self.maybe_mirror_event_text_to_realtime(&legacy_source)
|
||||
.await;
|
||||
@@ -2717,33 +2717,47 @@ impl Session {
|
||||
}
|
||||
}
|
||||
|
||||
/// Forwards terminal turn events from spawned MultiAgentV2 children to their direct parent.
|
||||
async fn maybe_notify_parent_of_terminal_turn(
|
||||
/// Classifies terminal results for native parent delivery and opted-in hosts.
|
||||
async fn maybe_publish_terminal_agent_result(
|
||||
&self,
|
||||
turn_context: &TurnContext,
|
||||
msg: &EventMsg,
|
||||
) {
|
||||
if turn_context.multi_agent_version != MultiAgentVersion::V2 {
|
||||
return;
|
||||
}
|
||||
|
||||
if !matches!(msg, EventMsg::TurnComplete(_) | EventMsg::TurnAborted(_)) {
|
||||
return;
|
||||
}
|
||||
|
||||
let SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
|
||||
agent_path: Some(_),
|
||||
..
|
||||
}) = &turn_context.session_source
|
||||
else {
|
||||
let notify_native_parent = turn_context.multi_agent_version == MultiAgentVersion::V2
|
||||
&& matches!(
|
||||
&turn_context.session_source,
|
||||
SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
|
||||
agent_path: Some(_),
|
||||
..
|
||||
})
|
||||
);
|
||||
let reporter = self
|
||||
.outcome_reporting
|
||||
.as_ref()
|
||||
.map(|reporting| &reporting.reporter);
|
||||
if !notify_native_parent && reporter.is_none() {
|
||||
return;
|
||||
};
|
||||
}
|
||||
|
||||
let terminal_error_message = {
|
||||
let mut terminal_error = turn_context.terminal_error.lock().await;
|
||||
if notify_native_parent {
|
||||
terminal_error.take().map(|error| error.message)
|
||||
} else {
|
||||
terminal_error.as_ref().map(|error| error.message.clone())
|
||||
}
|
||||
};
|
||||
let mut error_info = None;
|
||||
let status = match turn_context.terminal_error.lock().await.take() {
|
||||
Some(error) => {
|
||||
let status = AgentStatus::Errored(error.message);
|
||||
self.agent_status.send_replace(status.clone());
|
||||
let status = match terminal_error_message {
|
||||
Some(message) => {
|
||||
let status = AgentStatus::Errored(message);
|
||||
if notify_native_parent {
|
||||
self.agent_status.send_replace(status.clone());
|
||||
}
|
||||
status
|
||||
}
|
||||
None => {
|
||||
@@ -2768,24 +2782,31 @@ impl Session {
|
||||
return;
|
||||
}
|
||||
|
||||
self.services
|
||||
.agent_control
|
||||
.turn_finished(
|
||||
AgentTurnOutcome {
|
||||
thread_id: self.thread_id,
|
||||
turn_id: turn_context.sub_id.clone(),
|
||||
source: turn_context.session_source.clone(),
|
||||
parent_turn_id: turn_context.turn_metadata_state.parent_turn_id(),
|
||||
initiating_agent_path: turn_context
|
||||
.turn_metadata_state
|
||||
.initiating_agent_path()
|
||||
.cloned(),
|
||||
status,
|
||||
error_info,
|
||||
},
|
||||
&self.services.rollout_thread_trace,
|
||||
)
|
||||
.await;
|
||||
let outcome = AgentTurnOutcome {
|
||||
thread_id: self.thread_id,
|
||||
turn_id: turn_context.sub_id.clone(),
|
||||
source: turn_context.session_source.clone(),
|
||||
parent_turn_id: turn_context.turn_metadata_state.parent_turn_id(),
|
||||
initiating_agent_path: turn_context
|
||||
.turn_metadata_state
|
||||
.initiating_agent_path()
|
||||
.cloned(),
|
||||
status,
|
||||
error_info,
|
||||
};
|
||||
let host_outcome = if notify_native_parent {
|
||||
let host_outcome = reporter.is_some().then(|| outcome.clone());
|
||||
self.services
|
||||
.agent_control
|
||||
.turn_finished(outcome, &self.services.rollout_thread_trace)
|
||||
.await;
|
||||
host_outcome
|
||||
} else {
|
||||
Some(outcome)
|
||||
};
|
||||
if let Some((reporter, outcome)) = reporter.zip(host_outcome) {
|
||||
reporter.report(outcome).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn maybe_mirror_event_text_to_realtime(&self, msg: &EventMsg) {
|
||||
|
||||
@@ -9,6 +9,7 @@ use super::step_settings::StepSettingsUpdate;
|
||||
use super::*;
|
||||
use crate::agent::api::AgentControl;
|
||||
use crate::agent::host::AgentRuntimeServices;
|
||||
use crate::agent::host::AgentTurnOutcomeReporter;
|
||||
use crate::agent::host::InheritedSessionId;
|
||||
use crate::agents_md_manager::AgentsMdManager;
|
||||
use crate::agents_md_manager::SessionInstructions;
|
||||
@@ -56,10 +57,17 @@ use codex_utils_git_discovery::GitRootDiscovery;
|
||||
use codex_utils_path::replace_path_and_deduplicate;
|
||||
use std::sync::OnceLock;
|
||||
use tokio::sync::Semaphore;
|
||||
use tokio_util::task::TaskTracker;
|
||||
|
||||
type McpToolApprovalMetadataMap =
|
||||
HashMap<(String, String), std::sync::Weak<McpToolApprovalContext>>;
|
||||
|
||||
/// A thread's installed reporter and the tasks that may still hand it an outcome.
|
||||
pub(crate) struct AgentOutcomeReporting {
|
||||
pub(super) reporter: Arc<dyn AgentTurnOutcomeReporter>,
|
||||
pub(crate) handoffs: TaskTracker,
|
||||
}
|
||||
|
||||
/// Context for an initialized model agent
|
||||
///
|
||||
/// A session has at most 1 running task at a time, and can be interrupted by user input.
|
||||
@@ -67,6 +75,7 @@ pub(crate) struct Session {
|
||||
pub(crate) thread_id: ThreadId,
|
||||
pub(crate) item_id_generator: ItemIdGenerator,
|
||||
pub(crate) agent_runtime: AgentRuntimeServices,
|
||||
pub(crate) outcome_reporting: Option<AgentOutcomeReporting>,
|
||||
pub(crate) installation_id: String,
|
||||
pub(super) tx_event: Sender<Event>,
|
||||
/// Installed by Guardian before its first review; forwards items and cancellation without
|
||||
@@ -711,6 +720,12 @@ async fn warm_plugins_and_skills_for_session_init(
|
||||
}
|
||||
|
||||
impl Session {
|
||||
pub(crate) fn outcome_handoffs_closed(&self) -> bool {
|
||||
self.outcome_reporting
|
||||
.as_ref()
|
||||
.is_some_and(|reporting| reporting.handoffs.is_closed())
|
||||
}
|
||||
|
||||
pub(crate) fn new_item_id_generator() -> ItemIdGenerator {
|
||||
Arc::new(|prefix| ResponseItemId::with_suffix(prefix, uuid::Uuid::now_v7()))
|
||||
}
|
||||
@@ -1842,6 +1857,12 @@ impl Session {
|
||||
.await;
|
||||
}
|
||||
|
||||
let outcome_reporting = thread_extension_data
|
||||
.get::<Arc<dyn AgentTurnOutcomeReporter>>()
|
||||
.map(|reporter| AgentOutcomeReporting {
|
||||
reporter: Arc::clone(&reporter),
|
||||
handoffs: TaskTracker::new(),
|
||||
});
|
||||
let (agent_runtime, mut startup_error) =
|
||||
match AgentRuntimeServices::from_extensions(&thread_extension_data) {
|
||||
Ok(runtime) => (runtime, None),
|
||||
@@ -2006,6 +2027,7 @@ impl Session {
|
||||
thread_id,
|
||||
item_id_generator,
|
||||
agent_runtime,
|
||||
outcome_reporting,
|
||||
installation_id,
|
||||
tx_event: tx_event.clone(),
|
||||
parent_observer_event_tx: OnceLock::new(),
|
||||
|
||||
@@ -7076,6 +7076,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
|
||||
thread_id,
|
||||
item_id_generator,
|
||||
agent_runtime: crate::agent::host::AgentRuntimeServices::Controller,
|
||||
outcome_reporting: None,
|
||||
installation_id: "11111111-1111-4111-8111-111111111111".to_string(),
|
||||
tx_event,
|
||||
parent_observer_event_tx: OnceLock::new(),
|
||||
@@ -9419,6 +9420,7 @@ where
|
||||
thread_id,
|
||||
item_id_generator,
|
||||
agent_runtime: crate::agent::host::AgentRuntimeServices::Controller,
|
||||
outcome_reporting: None,
|
||||
installation_id: "11111111-1111-4111-8111-111111111111".to_string(),
|
||||
tx_event,
|
||||
parent_observer_event_tx: OnceLock::new(),
|
||||
|
||||
@@ -747,7 +747,7 @@ impl Session {
|
||||
&self,
|
||||
) -> Option<(Arc<tokio::sync::Mutex<TurnState>>, TurnStartOptions)> {
|
||||
let mut active_turn = self.active_turn.lock().await;
|
||||
if active_turn.is_some() {
|
||||
if active_turn.is_some() || self.outcome_handoffs_closed() {
|
||||
return None;
|
||||
}
|
||||
let has_trigger_turn = self.input_queue.has_trigger_turn_mailbox_items().await;
|
||||
|
||||
@@ -54,6 +54,8 @@ pub(super) async fn suspend_turn_and_shutdown(
|
||||
CodexErr::Fatal("accepted root turn suspension had no running task".to_string())
|
||||
})?;
|
||||
let turn_id = task.turn_context.sub_id.clone();
|
||||
// A suspended task emits no terminal outcome.
|
||||
drop(task.outcome_handoff);
|
||||
// Normal shutdown records a terminal turn event, preventing another worker from
|
||||
// recovering this turn under its original ID. Cancel the task without that event.
|
||||
task.cancellation_token.cancel();
|
||||
|
||||
@@ -9,6 +9,7 @@ use tokio::sync::Mutex;
|
||||
use tokio::sync::Notify;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tokio_util::task::AbortOnDropHandle;
|
||||
use tokio_util::task::task_tracker::TaskTrackerToken;
|
||||
|
||||
use codex_diagnostics::GaugeGuard;
|
||||
use codex_protocol::dynamic_tools::DynamicToolResponse;
|
||||
@@ -86,6 +87,8 @@ pub(crate) struct RunningTask {
|
||||
pub(crate) handle: AbortOnDropHandle<()>,
|
||||
pub(crate) turn_context: Arc<TurnContext>,
|
||||
pub(crate) _agent_execution_guard: Option<AgentExecutionGuard>,
|
||||
/// Follows the task into whichever finalizer claims its terminal event.
|
||||
pub(crate) outcome_handoff: Option<TaskTrackerToken>,
|
||||
/// Unrecorded input retained during host admission or preparation that can wait on the roster.
|
||||
pub(crate) pending_initial_input: Option<Vec<TurnInput>>,
|
||||
pub(crate) _diagnostics_guard: GaugeGuard,
|
||||
|
||||
@@ -336,6 +336,9 @@ impl Session {
|
||||
let (pending_items, _) = self.input_queue.drain_mailbox_input_items().await;
|
||||
let turn_state = {
|
||||
let mut active = self.active_turn.lock().await;
|
||||
if self.outcome_handoffs_closed() {
|
||||
return;
|
||||
}
|
||||
let turn = active.get_or_insert_with(ActiveTurn::default);
|
||||
debug_assert!(turn.task.is_none());
|
||||
turn.finishing = None;
|
||||
@@ -353,6 +356,15 @@ impl Session {
|
||||
.await;
|
||||
|
||||
let mut active = self.active_turn.lock().await;
|
||||
// Shutdown closes registration before taking this same lock to abort active work.
|
||||
let outcome_handoff = if let Some(reporting) = &self.outcome_reporting {
|
||||
if reporting.handoffs.is_closed() {
|
||||
return;
|
||||
}
|
||||
Some(reporting.handoffs.token())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let turn = active.get_or_insert_with(ActiveTurn::default);
|
||||
debug_assert!(turn.task.is_none());
|
||||
// Record lineage only when the task starts: settings changes can invalidate a reservation.
|
||||
@@ -474,6 +486,7 @@ impl Session {
|
||||
cancellation_token,
|
||||
turn_context: Arc::clone(&turn_context),
|
||||
_agent_execution_guard: agent_execution_guard,
|
||||
outcome_handoff,
|
||||
pending_initial_input,
|
||||
_diagnostics_guard: ACTIVE_TURNS.track(),
|
||||
_timer: timer,
|
||||
@@ -720,7 +733,7 @@ impl Session {
|
||||
regular_output.is_some() && turn_context.terminal_error.lock().await.is_none();
|
||||
let finishing = CancellationToken::new();
|
||||
let finish_recording = finishing.clone().drop_guard();
|
||||
let (turn_state, unadmitted_task, mut pending_input) = {
|
||||
let (turn_state, unadmitted_task, mut pending_input, outcome_handoff) = {
|
||||
let mut active = self.active_turn.lock().await;
|
||||
let active_turn = active.as_mut()?;
|
||||
if !Arc::ptr_eq(&active_turn.task.as_ref()?.turn_context, &turn_context) {
|
||||
@@ -742,6 +755,7 @@ impl Session {
|
||||
(task._agent_execution_guard.is_none() && task.pending_initial_input.is_some())
|
||||
.then_some(task.task),
|
||||
task.pending_initial_input.unwrap_or_default(),
|
||||
task.outcome_handoff,
|
||||
)
|
||||
};
|
||||
turn_context
|
||||
@@ -979,6 +993,7 @@ impl Session {
|
||||
// The parent can request another review as soon as it receives this event.
|
||||
self.send_event(turn_context.as_ref(), event).await;
|
||||
}
|
||||
drop(outcome_handoff);
|
||||
if cleared_active_turn {
|
||||
self.emit_thread_idle_lifecycle_if_idle(idle_cause).await;
|
||||
}
|
||||
|
||||
@@ -5,6 +5,7 @@ use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
use std::sync::atomic::AtomicUsize;
|
||||
use std::sync::atomic::Ordering;
|
||||
use std::time::Duration;
|
||||
|
||||
use codex_core::AgentConfigUpdate;
|
||||
use codex_core::AgentControl;
|
||||
@@ -13,6 +14,7 @@ use codex_core::AgentInfo;
|
||||
use codex_core::AgentRef;
|
||||
use codex_core::AgentTarget;
|
||||
use codex_core::AgentTurnOutcome;
|
||||
use codex_core::AgentTurnOutcomeReporter;
|
||||
use codex_core::DeliveryReceipt;
|
||||
use codex_core::ForkSnapshot;
|
||||
use codex_core::GuardianRootSnapshot;
|
||||
@@ -24,13 +26,18 @@ use codex_core::SendRequest;
|
||||
use codex_core::SpawnRequest;
|
||||
use codex_core::StartThreadOptions;
|
||||
use codex_core::ThreadConfigSnapshot;
|
||||
use codex_core::TurnInput;
|
||||
use codex_core::TurnInputRequest;
|
||||
use codex_core::TurnInputSubmission;
|
||||
use codex_core::TurnStartOptions;
|
||||
use codex_core::config::Config;
|
||||
use codex_extension_api::ExtensionFuture;
|
||||
use codex_extension_api::ExtensionRegistryBuilder;
|
||||
use codex_extension_api::ThreadLifecycleContributor;
|
||||
use codex_extension_api::ThreadStartInput;
|
||||
use codex_extension_api::ThreadStopInput;
|
||||
use codex_extension_api::TurnLifecycleContributor;
|
||||
use codex_extension_api::TurnStopInput;
|
||||
use codex_features::Feature;
|
||||
use codex_history::InitialHistory;
|
||||
use codex_history::ResumedHistory;
|
||||
@@ -42,25 +49,32 @@ use codex_protocol::ThreadId;
|
||||
use codex_protocol::error::CodexErr;
|
||||
use codex_protocol::error::CodexErrorDetails;
|
||||
use codex_protocol::error::Result as CodexResult;
|
||||
use codex_protocol::protocol::AgentStatus;
|
||||
use codex_protocol::protocol::InterAgentCommunication;
|
||||
use codex_protocol::protocol::MultiAgentVersion;
|
||||
use codex_protocol::protocol::Op;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use codex_protocol::protocol::SubAgentSource;
|
||||
use codex_protocol::protocol::ThreadHistoryMode;
|
||||
use codex_protocol::protocol::ThreadSource;
|
||||
use codex_protocol::protocol::TokenUsage;
|
||||
use codex_protocol::user_input::UserInput;
|
||||
use codex_rollout_trace::ThreadTraceContext;
|
||||
use codex_thread_store::LoadThreadHistoryParams;
|
||||
use core_test_support::ThreadIdle;
|
||||
use core_test_support::responses;
|
||||
use core_test_support::skip_if_no_network;
|
||||
use core_test_support::test_codex::TestCodex;
|
||||
use core_test_support::test_codex::test_codex;
|
||||
use core_test_support::wait_for_event;
|
||||
use core_test_support::wait_for_event_match;
|
||||
use futures::future::BoxFuture;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::json;
|
||||
use test_case::test_case;
|
||||
use tokio::sync::Notify;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use wiremock::MockServer;
|
||||
|
||||
pub(super) struct HostRuntime;
|
||||
@@ -642,6 +656,258 @@ async fn fresh_hosted_child_persists_inherited_session_id(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn hosted_child_reports_terminal_turn_outcomes() -> anyhow::Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
|
||||
const ANSWER: &str = "The patch looks good.";
|
||||
const FAILURE: &str = "The request was rejected.";
|
||||
const FOLLOWUP: &str = "Try again with the smaller patch.";
|
||||
const ASSIGNING_TURN: &str = "assigner-turn";
|
||||
|
||||
struct Reporter {
|
||||
outcomes: mpsc::UnboundedSender<AgentTurnOutcome>,
|
||||
gate: Mutex<Option<oneshot::Receiver<()>>>,
|
||||
}
|
||||
|
||||
impl AgentTurnOutcomeReporter for Reporter {
|
||||
fn report(&self, outcome: AgentTurnOutcome) -> BoxFuture<'_, ()> {
|
||||
Box::pin(async move {
|
||||
let failed = matches!(outcome.status, AgentStatus::Errored(_));
|
||||
self.outcomes.send(outcome).expect("report receiver");
|
||||
if failed {
|
||||
let gate = self.gate.lock().expect("report gate").take();
|
||||
gate.expect("one failure report")
|
||||
.await
|
||||
.expect("release failure report");
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
let server = responses::start_mock_server().await;
|
||||
let mut extensions = ExtensionRegistryBuilder::new();
|
||||
extensions.thread_lifecycle_contributor(Arc::new(ThreadIdle));
|
||||
let root = test_codex()
|
||||
.with_extensions(Arc::new(extensions.build()))
|
||||
.build_with_auto_env(&server)
|
||||
.await?;
|
||||
let session_id = root.session_configured.session_id;
|
||||
let (release_report, report_gate) = oneshot::channel();
|
||||
let (outcomes, mut reports) = mpsc::unbounded_channel();
|
||||
let reporter = Arc::new(Reporter {
|
||||
outcomes,
|
||||
gate: Mutex::new(Some(report_gate)),
|
||||
});
|
||||
let source = SessionSource::Custom("host".to_string());
|
||||
let mut options = StartThreadOptions {
|
||||
session_source: Some(source.clone()),
|
||||
thread_source: Some(ThreadSource::Subagent),
|
||||
..root.start_thread_options().await
|
||||
};
|
||||
options
|
||||
.thread_extension_init
|
||||
.insert(InheritedSessionId(session_id));
|
||||
options
|
||||
.thread_extension_init
|
||||
.insert(reporter as Arc<dyn AgentTurnOutcomeReporter>);
|
||||
let child = root.thread_manager.start_thread(options).await?;
|
||||
let thread = &child.thread;
|
||||
|
||||
responses::mount_sse_sequence(
|
||||
&server,
|
||||
vec![
|
||||
responses::sse(vec![
|
||||
responses::ev_response_created("review"),
|
||||
responses::ev_assistant_message("answer", ANSWER),
|
||||
responses::ev_completed("review"),
|
||||
]),
|
||||
responses::sse_failed("failed", "invalid_prompt", FAILURE),
|
||||
responses::sse_completed("silent-followup"),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
let mut expect_report = async |expected_status| {
|
||||
let turn_id = wait_for_event_match(thread, |event| match event {
|
||||
EventMsg::TurnComplete(event) => Some(event.turn_id.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.await;
|
||||
let report = tokio::time::timeout(Duration::from_secs(10), reports.recv())
|
||||
.await
|
||||
.expect("timed out waiting for the terminal turn report")
|
||||
.expect("reporter dropped");
|
||||
assert_eq!(
|
||||
(report.thread_id, &report.turn_id, &report.source),
|
||||
(child.thread_id, &turn_id, &source),
|
||||
);
|
||||
assert_eq!(
|
||||
(&report.status, report.error_info.as_ref()),
|
||||
(&expected_status, None)
|
||||
);
|
||||
report
|
||||
};
|
||||
for (prompt, expected_status) in [
|
||||
(
|
||||
"Review the patch.",
|
||||
AgentStatus::Completed(Some(ANSWER.into())),
|
||||
),
|
||||
("Fail this request.", AgentStatus::Errored(FAILURE.into())),
|
||||
] {
|
||||
thread
|
||||
.start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
|
||||
text: prompt.into(),
|
||||
text_elements: Vec::new(),
|
||||
}]))
|
||||
.await?;
|
||||
let report = expect_report(expected_status).await;
|
||||
if !matches!(report.status, AgentStatus::Errored(_)) {
|
||||
ThreadIdle::wait(thread).await;
|
||||
}
|
||||
}
|
||||
|
||||
// Queue the next task while the previous turn is still reporting its outcome.
|
||||
let assigner = AgentPath::root().join("assigner").expect("assigner path");
|
||||
let mail = InterAgentCommunication::new(
|
||||
assigner.clone(),
|
||||
AgentPath::root().join("reviewer").expect("child path"),
|
||||
Vec::new(),
|
||||
FOLLOWUP.into(),
|
||||
/*trigger_turn*/ true,
|
||||
);
|
||||
let followup = TurnInputRequest::new(TurnInput::InterAgentCommunication(mail)).on_start(
|
||||
TurnStartOptions {
|
||||
parent_turn_id: Some(ASSIGNING_TURN.into()),
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
assert_eq!(
|
||||
thread.start_or_steer_turn(followup).await?,
|
||||
TurnInputSubmission::Accepted
|
||||
);
|
||||
release_report.send(()).expect("release failure report");
|
||||
|
||||
let report = expect_report(AgentStatus::Completed(None)).await;
|
||||
ThreadIdle::wait(thread).await;
|
||||
assert_eq!(
|
||||
(
|
||||
report.parent_turn_id.as_deref(),
|
||||
report.initiating_agent_path
|
||||
),
|
||||
(Some(ASSIGNING_TURN), Some(assigner)),
|
||||
);
|
||||
assert_eq!(
|
||||
reports.try_recv().err(),
|
||||
Some(mpsc::error::TryRecvError::Empty)
|
||||
);
|
||||
thread.shutdown_and_wait().await?;
|
||||
let requests = responses::received_responses_requests(&server).await;
|
||||
assert_eq!(requests.len(), 3);
|
||||
assert_eq!(
|
||||
requests[2].body_json()["client_metadata"]["turn_id"],
|
||||
report.turn_id
|
||||
);
|
||||
let messages = requests[2].inputs_of_type("agent_message");
|
||||
assert_eq!(messages.len(), 1);
|
||||
assert!(
|
||||
messages[0]["content"][0]["text"]
|
||||
.as_str()
|
||||
.is_some_and(|text| text.ends_with(FOLLOWUP))
|
||||
);
|
||||
root.codex.shutdown_and_wait().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shutdown_waits_for_host_outcome_handoff() -> anyhow::Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
|
||||
#[derive(Default)]
|
||||
struct Handoff {
|
||||
entered: Notify,
|
||||
release: Notify,
|
||||
events: Mutex<Vec<&'static str>>,
|
||||
}
|
||||
|
||||
impl Handoff {
|
||||
async fn pause(&self) {
|
||||
self.entered.notify_one();
|
||||
self.release.notified().await;
|
||||
}
|
||||
}
|
||||
|
||||
impl TurnLifecycleContributor for Handoff {
|
||||
fn on_turn_stop<'a>(&'a self, _input: TurnStopInput<'a>) -> ExtensionFuture<'a, ()> {
|
||||
Box::pin(self.pause())
|
||||
}
|
||||
}
|
||||
|
||||
impl AgentTurnOutcomeReporter for Handoff {
|
||||
fn report(&self, _outcome: AgentTurnOutcome) -> BoxFuture<'_, ()> {
|
||||
Box::pin(async move {
|
||||
self.pause().await;
|
||||
self.events.lock().expect("handoff events").push("reported");
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl ThreadLifecycleContributor<Config> for Handoff {
|
||||
fn on_thread_stop<'a>(&'a self, _input: ThreadStopInput<'a>) -> ExtensionFuture<'a, ()> {
|
||||
Box::pin(async move { self.events.lock().expect("handoff events").push("stopped") })
|
||||
}
|
||||
}
|
||||
|
||||
let handoff = Arc::new(Handoff::default());
|
||||
let mut extensions = ExtensionRegistryBuilder::new();
|
||||
extensions.turn_lifecycle_contributor(handoff.clone());
|
||||
extensions.thread_lifecycle_contributor(handoff.clone());
|
||||
let server = responses::start_mock_server().await;
|
||||
let root = test_codex()
|
||||
.with_extensions(Arc::new(extensions.build()))
|
||||
.build_with_auto_env(&server)
|
||||
.await?;
|
||||
let mut options = root.start_thread_options().await;
|
||||
options
|
||||
.thread_extension_init
|
||||
.insert(handoff.clone() as Arc<dyn AgentTurnOutcomeReporter>);
|
||||
let child = root.thread_manager.start_thread(options).await?;
|
||||
let thread = &child.thread;
|
||||
|
||||
responses::mount_sse_once(&server, responses::sse_completed("done")).await;
|
||||
thread
|
||||
.start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
|
||||
text: "Finish the task.".into(),
|
||||
text_elements: Vec::new(),
|
||||
}]))
|
||||
.await?;
|
||||
|
||||
// The finalizer has detached its task, but has not entered the reporter.
|
||||
let wait_for_gate =
|
||||
async || tokio::time::timeout(Duration::from_secs(10), handoff.entered.notified()).await;
|
||||
wait_for_gate().await?;
|
||||
thread.submit(Op::Shutdown).await?;
|
||||
let assert_pending_then_release = async || {
|
||||
assert!(
|
||||
tokio::time::timeout(Duration::from_millis(250), thread.wait_until_terminated())
|
||||
.await
|
||||
.is_err(),
|
||||
"shutdown completed before the local handoff"
|
||||
);
|
||||
handoff.release.notify_one();
|
||||
};
|
||||
assert_pending_then_release().await;
|
||||
// The reporter has now been called, but the host has not accepted the outcome yet.
|
||||
wait_for_gate().await?;
|
||||
assert_pending_then_release().await;
|
||||
tokio::time::timeout(Duration::from_secs(10), thread.wait_until_terminated()).await?;
|
||||
assert_eq!(
|
||||
*handoff.events.lock().expect("handoff events"),
|
||||
["reported", "stopped"]
|
||||
);
|
||||
root.codex.shutdown_and_wait().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn collaboration_tools_dispatch_to_the_host_controller() -> anyhow::Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
|
||||
@@ -2,21 +2,66 @@
|
||||
|
||||
use super::*;
|
||||
use codex_config::config_toml::CircuitBreakAction;
|
||||
use codex_core::AgentTurnOutcome;
|
||||
use codex_core::AgentTurnOutcomeReporter;
|
||||
use codex_core::TurnInputSubmission;
|
||||
use codex_core::config::Constrained;
|
||||
use codex_extension_api::ExtensionFuture;
|
||||
use codex_extension_api::ExtensionRegistryBuilder;
|
||||
use codex_extension_api::ThreadLifecycleContributor;
|
||||
use codex_extension_api::ThreadStartInput;
|
||||
use codex_protocol::config_types::ApprovalsReviewer;
|
||||
use codex_protocol::protocol::CodexErrorInfo;
|
||||
use codex_protocol::protocol::ErrorEvent;
|
||||
use codex_protocol::protocol::MultiAgentVersion;
|
||||
use codex_protocol::protocol::TurnAbortReason;
|
||||
use core_test_support::ThreadIdle;
|
||||
use core_test_support::responses::received_responses_requests;
|
||||
use core_test_support::streaming_sse::StreamingSseChunk;
|
||||
use core_test_support::streaming_sse::start_streaming_sse_server;
|
||||
use futures::future::BoxFuture;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::sync::Arc;
|
||||
use std::sync::Mutex;
|
||||
use test_case::test_case;
|
||||
use tokio::sync::oneshot;
|
||||
|
||||
type ReportedOutcome = (ThreadId, String, AgentStatus, Option<CodexErrorInfo>);
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
struct CompletionReports(Arc<Mutex<Vec<ReportedOutcome>>>);
|
||||
|
||||
impl AgentTurnOutcomeReporter for CompletionReports {
|
||||
fn report(&self, outcome: AgentTurnOutcome) -> BoxFuture<'_, ()> {
|
||||
Box::pin(async move {
|
||||
self.0.lock().expect("completion reports").push((
|
||||
outcome.thread_id,
|
||||
outcome.turn_id,
|
||||
outcome.status,
|
||||
outcome.error_info,
|
||||
));
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl ThreadLifecycleContributor<Config> for CompletionReports {
|
||||
fn on_thread_start<'a>(
|
||||
&'a self,
|
||||
input: ThreadStartInput<'a, Config>,
|
||||
) -> ExtensionFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
if matches!(
|
||||
input.session_source,
|
||||
SessionSource::SubAgent(SubAgentSource::ThreadSpawn { .. })
|
||||
) {
|
||||
input
|
||||
.thread_store
|
||||
.insert(Arc::new(self.clone()) as Arc<dyn AgentTurnOutcomeReporter>);
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[test_case(CircuitBreakAction::Strict; "strict_notifies_parent")]
|
||||
#[test_case(CircuitBreakAction::Default; "default_stays_silent")]
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
@@ -24,7 +69,9 @@ async fn guardian_circuit_breaker_notifies_parent(action: CircuitBreakAction) ->
|
||||
skip_if_no_network!(Ok(()));
|
||||
let server = start_mock_server().await;
|
||||
let mut extensions = ExtensionRegistryBuilder::new();
|
||||
let reports = CompletionReports::default();
|
||||
extensions.thread_lifecycle_contributor(Arc::new(ThreadIdle));
|
||||
extensions.thread_lifecycle_contributor(Arc::new(reports.clone()));
|
||||
let test = test_codex()
|
||||
.with_model_info_override("test-gpt-5.1-codex", |model| {
|
||||
model.multi_agent_version = Some(MultiAgentVersion::V2);
|
||||
@@ -140,14 +187,11 @@ async fn guardian_circuit_breaker_notifies_parent(action: CircuitBreakAction) ->
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let mut created = test.thread_manager.subscribe_thread_created();
|
||||
test.submit_text_turn("Spawn a worker to run the checks.")
|
||||
.await?;
|
||||
ThreadIdle::wait(&test.codex).await;
|
||||
let worker = test
|
||||
.thread_manager
|
||||
.get_thread(created.recv().await?)
|
||||
.await?;
|
||||
let worker_id = test.codex.resolve_agent_reference("/root/worker").await?;
|
||||
let worker = test.thread_manager.get_thread(worker_id).await?;
|
||||
let warning = wait_for_event_match(&worker, |event| match event {
|
||||
EventMsg::GuardianWarning(warning) if warning.message.contains("3 consecutive") => {
|
||||
Some(warning.message.clone())
|
||||
@@ -174,43 +218,47 @@ async fn guardian_circuit_breaker_notifies_parent(action: CircuitBreakAction) ->
|
||||
AgentStatus::Interrupted,
|
||||
),
|
||||
);
|
||||
|
||||
// Delivery is queued. If the parent samples before draining its mailbox,
|
||||
// wait_agent must observe the notification before the next model request.
|
||||
let _wait = if action == CircuitBreakAction::Strict {
|
||||
Some(
|
||||
mount_sse_once_match(
|
||||
&server,
|
||||
move |request: &wiremock::Request| {
|
||||
is_root(request) && !body_contains(request, "Agent interrupted by Guardian:")
|
||||
},
|
||||
sse(vec![
|
||||
ev_function_call_with_namespace(
|
||||
"wait-worker",
|
||||
MULTI_AGENT_V2_NAMESPACE,
|
||||
"wait_agent",
|
||||
"{}",
|
||||
),
|
||||
ev_completed("parent-wait"),
|
||||
]),
|
||||
)
|
||||
.await,
|
||||
let expected_report = (action == CircuitBreakAction::Strict).then(|| {
|
||||
(
|
||||
worker_id,
|
||||
aborted.turn_id.expect("aborted worker turn ID"),
|
||||
AgentStatus::Errored(warning.clone()),
|
||||
Some(CodexErrorInfo::TooManyDenials),
|
||||
)
|
||||
} else {
|
||||
None
|
||||
});
|
||||
assert_eq!(
|
||||
reports.0.lock().expect("completion reports").as_slice(),
|
||||
expected_report.as_slice(),
|
||||
);
|
||||
|
||||
// Core can sample the user prompt and drain queued mail on a later request in the same turn.
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/responses"))
|
||||
.and(is_root)
|
||||
.respond_with(sse_response(sse(vec![ev_completed("parent-continued")])))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let TurnInputSubmission::Started { turn_id, .. } = test
|
||||
.codex
|
||||
.start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
|
||||
text: "Report the worker's status.".into(),
|
||||
text_elements: Vec::new(),
|
||||
}]))
|
||||
.await?
|
||||
else {
|
||||
anyhow::bail!("parent prompt should start a turn");
|
||||
};
|
||||
let continued = mount_sse_once_match(
|
||||
&server,
|
||||
move |request: &wiremock::Request| {
|
||||
is_root(request)
|
||||
&& (action == CircuitBreakAction::Default
|
||||
|| body_contains(request, "Agent interrupted by Guardian:"))
|
||||
},
|
||||
sse(vec![ev_completed("parent-continued")]),
|
||||
)
|
||||
let done = wait_for_event_match(&test.codex, |event| match event {
|
||||
EventMsg::TurnComplete(done) if done.turn_id == turn_id => Some(done.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.await;
|
||||
test.submit_text_turn("Report the worker's status.").await?;
|
||||
let continued_request = continued.single_request();
|
||||
assert_eq!(done.error, None);
|
||||
let continued_request = received_responses_requests(&server)
|
||||
.await
|
||||
.into_iter()
|
||||
.rfind(|request| request.body_json()["client_metadata"]["turn_id"] == turn_id)
|
||||
.expect("model request for the parent's follow-up turn");
|
||||
let messages = continued_request.inputs_of_type("agent_message");
|
||||
let expected = if action == CircuitBreakAction::Strict {
|
||||
vec![json!({
|
||||
|
||||
Reference in New Issue
Block a user