openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-05 20:10:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add host-owned turn admission to the agent runtime (#58963)

Browse Source 
## What changed

Add `AgentRuntimeOwnership::Host` and an `AgentTurnAdmission` provider so hosts can acquire execution capacity before runtime tasks run, including root turns. Require the provider at startup and retain its execution guard for the task lifetime. Controller ownership remains the default; host ownership skips controller rollout-budget accounting, budget reminders, and subagent roster context.

Keep pending admission interruptible without blocking input submission. Report admission failures through the turn lifecycle, record accepted input before exposing the error, and unwind task setup such as review mode. Preserve recoverable input during suspension. On interruption, skip prompts that still require a `UserPromptSubmit` hook rather than recording them without inspection.

## Testing

Add integration tests for missing-provider startup cleanup, denial and immediate retry, steering during admission, cancellation, execution guard lifetime, submit-hook handling, suspension recovery, and review-mode cleanup. Add a context snapshot verifying retained input reaches the next admitted turn.

GitOrigin-RevId: 9ee85fad2ecfbfb38735bac8ef88888fb179a967
This commit is contained in:
rhan-oai
2026-11-02 20:10:37 +00:00
committed by copyberry
parent 9ede2e87ce
commit d893b14475
25 changed files with 822 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/agent/host/admission.rs
+30
View File
@@ -0,0 +1,30 @@
//! Connects host-owned turn capacity to Core task lifetimes.
use crate::agent::types::AgentExecutionGuard;
use codex_protocol::ThreadId;
use codex_protocol::error::Result;
use futures::future::BoxFuture;
/// Identifies the execution for which Core requests host capacity.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AgentTurnAdmissionRequest {
pub thread_id: ThreadId,
pub turn_id: String,
}
/// Acquires host-owned capacity before each configured runtime task executes, including roots.
///
/// Select host ownership and install this provider before thread startup finishes. Return
/// `AgentLimitReached` when capacity is exhausted rather than waiting for a slot.
/// Cancellation drops the acquisition future. Errors fail the attempted turn
/// without falling back to Local admission. The returned guard releases or arranges
/// release of capacity when Core ends the task.
/// A submitted turn can be denied after its input was accepted. Core records that
/// input and reports the error through the ordinary turn lifecycle.
/// Hosts own remote retries, abandoned acquisitions and worker-loss recovery.
pub trait AgentTurnAdmission: Send + Sync {
fn admit_turn(
&self,
request: AgentTurnAdmissionRequest,
) -> BoxFuture<'_, Result<AgentExecutionGuard>>;
}
codex-rs/core/src/agent/host/mod.rs
+6
View File
@@ -1,8 +1,12 @@
mod admission;
mod invocation;
mod mcp_metadata;
mod ownership;
mod spawn;
mod types;
pub use admission::AgentTurnAdmission;
pub use admission::AgentTurnAdmissionRequest;
pub use invocation::HostToolInvocationLookup;
pub(crate) use invocation::InvocationReservation;
pub use invocation::ToolInvocationContext;
@@ -10,6 +14,8 @@ pub(crate) use invocation::create;
pub(crate) use invocation::prepare;
pub(crate) use invocation::tool_invocation;
pub use mcp_metadata::AgentMcpMetadataRecipient;
pub use ownership::AgentRuntimeOwnership;
pub(crate) use ownership::AgentRuntimeServices;
pub use types::HostThreadSpawnReservation;
pub use types::PreparedChildContext;
pub(crate) use types::PreparedChildInheritance;
codex-rs/core/src/agent/host/ownership.rs
+43
View File
@@ -0,0 +1,43 @@
//! Captures coordination providers once, after the thread's startup contributors finish.
use super::AgentTurnAdmission;
use codex_extension_api::ExtensionData;
use codex_protocol::error::CodexErr;
use codex_protocol::error::Result;
use std::sync::Arc;
/// Selects the runtime's coordination owner at startup. Installing providers alone
/// does not select host ownership.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum AgentRuntimeOwnership {
#[default]
Controller,
/// Requires an [`AgentTurnAdmission`] provider. Core skips controller rollout-budget
/// accounting, budget reminders, and subagent roster context for this runtime.
Host,
}
pub(crate) enum AgentRuntimeServices {
Controller,
Host(HostAgentServices),
}
pub(crate) struct HostAgentServices {
pub(crate) turn_admission: Arc<dyn AgentTurnAdmission>,
}
impl AgentRuntimeServices {
pub(crate) fn from_extensions(data: &ExtensionData) -> Result<Self> {
if data.get::<AgentRuntimeOwnership>().as_deref() != Some(&AgentRuntimeOwnership::Host) {
return Ok(Self::Controller);
}
let turn_admission = data.get::<Arc<dyn AgentTurnAdmission>>().ok_or_else(|| {
CodexErr::InvalidRequest(
"host coordination requires AgentTurnAdmission before startup".into(),
)
})?;
Ok(Self::Host(HostAgentServices {
turn_admission: Arc::clone(turn_admission.as_ref()),
}))
}
}
codex-rs/core/src/agent/host/spawn.rs
+3
View File
@@ -132,6 +132,9 @@ pub(super) async fn create_spawn(
.into());
}
options.reserved_thread_id = Some(reservation.thread_id);
if let Some(admission) = reservation.turn_admission {
options.thread_extension_init.insert(admission);
}
let child = state
.spawn_host_child_thread(
options,
codex-rs/core/src/agent/host/spawn_tests.rs
+1
View File
@@ -198,6 +198,7 @@ async fn same_manager_creation_retains_step_executor_and_shares_exec_policy() ->
HostThreadSpawnReservation {
thread_id: child_id,
control: Arc::clone(&session.services.agent_control),
turn_admission: None,
},
)
.await?;
codex-rs/core/src/agent/host/types.rs
+5 -1
View File
@@ -1,5 +1,6 @@
//! Prepared inputs and host reservations for native child startup.
use super::admission::AgentTurnAdmission;
use crate::agent::api::AgentControl;
use crate::agents_md_manager::SessionInstructions;
use crate::exec_policy::ExecPolicyManager;
@@ -18,10 +19,13 @@ use std::sync::Arc;
/// A fresh child ID reserved in the host's tree.
/// The host commits membership, serializes startup attempts and deduplicates retries.
/// Core derives the parent, session and path from the invocation and child request.
/// The controller supplies runtime coordination.
/// The controller supplies runtime coordination by default. For host-owned admission, select
/// `AgentRuntimeOwnership::Host` before startup finishes.
pub struct HostThreadSpawnReservation {
pub thread_id: ThreadId,
pub control: Arc<dyn AgentControl>,
/// Admission provider for a host-owned child. A lifecycle contributor may also install it.
pub turn_admission: Option<Arc<dyn AgentTurnAdmission>>,
}
/// Parent inputs captured from the tool invocation for a host starting or forking a child.
codex-rs/core/src/codex_thread.rs
+4
View File
@@ -1408,6 +1408,10 @@ impl CodexThread {
&self,
control: &dyn AgentControl,
) -> CodexResult<()> {
// The host checks and reserves capacity when the task starts.
if self.session.host_turn_admission().is_some() {
return Ok(());
}
if self.session.active_turn.lock().await.is_some() {
return Ok(());
}
codex-rs/core/src/guardian/input_budget.rs
+4 -1
View File
@@ -242,7 +242,10 @@ pub(crate) async fn finalize(
usize::try_from(estimate_item_token_count(&reminder)).unwrap_or(usize::MAX),
);
}
if let Some(reminder) = session
if matches!(
&session.agent_runtime,
crate::agent::host::AgentRuntimeServices::Controller
) && let Some(reminder) = session
.services
.agent_control
.pending_budget_reminder(session.thread_id(), &session.current_window_id().await)
codex-rs/core/src/hook_runtime.rs
+2
View File
@@ -761,6 +761,8 @@ pub(crate) async fn record_pending_input(
client_id,
metadata,
} => {
// Prompt telemetry can contain full text. Emit it only for prompts selected for recording.
turn_context.session_telemetry.user_prompt(&content);
sess.record_user_prompt_and_emit_turn_item(
turn_context.as_ref(),
model_info,
codex-rs/core/src/lib.rs
+3
View File
@@ -157,6 +157,9 @@ mod thread_manager;
mod thread_spawn;
mod thread_store_builder;
pub use agent::host::AgentMcpMetadataRecipient;
pub use agent::host::AgentRuntimeOwnership;
pub use agent::host::AgentTurnAdmission;
pub use agent::host::AgentTurnAdmissionRequest;
pub use agent::host::HostThreadSpawnReservation;
pub use agent::host::HostToolInvocationLookup;
pub use agent::host::PreparedChildContext;
codex-rs/core/src/session/mod.rs
+7
View File
@@ -2529,6 +2529,13 @@ impl Session {
}
pub(crate) async fn emit_turn_started(&self, turn_context: &TurnContext, task_kind: TaskKind) {
if turn_context
.extension_data
.get::<crate::tasks::HostTurnStarted>()
.is_some()
{
return;
}
let event = TurnStartedEvent {
turn_attribution: (task_kind == TaskKind::Regular).then(|| turn_context.attribution()),
turn_id: turn_context.sub_id.clone(),
codex-rs/core/src/session/rollout_budget.rs
+7
View File
@@ -1,5 +1,6 @@
use super::session::Session;
use super::turn_context::TurnContext;
use crate::agent::host::AgentRuntimeServices;
use crate::context::ContextualUserFragment;
use codex_protocol::error::Result as CodexResult;
use codex_protocol::protocol::TokenUsage;
@@ -9,6 +10,9 @@ pub(super) async fn maybe_record_reminder(
turn_context: &TurnContext,
window_id: &str,
) {
if matches!(&sess.agent_runtime, AgentRuntimeServices::Host(_)) {
return;
}
let agent_control = &sess.services.agent_control;
let Some(reminder) = agent_control
.pending_budget_reminder(sess.thread_id(), window_id)
@@ -32,6 +36,9 @@ pub(super) async fn maybe_record_reminder(
impl Session {
pub(crate) async fn record_rollout_budget_usage(&self, usage: &TokenUsage) -> CodexResult<()> {
if matches!(&self.agent_runtime, AgentRuntimeServices::Host(_)) {
return Ok(());
}
self.services
.agent_control
.record_usage(usage.clone())
codex-rs/core/src/session/session.rs
+17 -10
View File
@@ -8,6 +8,7 @@ use super::step_settings::StepSettingsConstraints;
use super::step_settings::StepSettingsUpdate;
use super::*;
use crate::agent::api::AgentControl;
use crate::agent::host::AgentRuntimeServices;
use crate::agents_md_manager::AgentsMdManager;
use crate::agents_md_manager::SessionInstructions;
use crate::config::ConstraintError;
@@ -60,6 +61,7 @@ type McpToolApprovalMetadataMap =
/// A session has at most 1 running task at a time, and can be interrupted by user input.
pub(crate) struct Session {
pub(crate) thread_id: ThreadId,
pub(crate) agent_runtime: AgentRuntimeServices,
pub(crate) installation_id: String,
pub(super) tx_event: Sender<Event>,
/// Installed by Guardian before its first review; forwards items and cancellation without
@@ -1755,6 +1757,12 @@ impl Session {
.await;
}
let (agent_runtime, mut startup_error) =
match AgentRuntimeServices::from_extensions(&thread_extension_data) {
Ok(runtime) => (runtime, None),
// This placeholder is used only for teardown, before the session is published.
Err(err) => (AgentRuntimeServices::Controller, Some(err)),
};
let executed_tool_calls =
crate::state::ExecutedToolCalls::new(&config.features, &initial_history);
let codex_responses_headers =
@@ -1784,12 +1792,12 @@ impl Session {
} else {
ThreadTraceContext::start_root(trace_metadata)
};
// Extensions have already started. On error, finish constructing the
// session so normal teardown can release everything they registered.
let (rollout_thread_trace, trace_error) = match rollout_thread_trace {
Ok(trace) => (trace, None),
Err(err) => (ThreadTraceContext::disabled(), Some(err)),
};
let rollout_thread_trace = rollout_thread_trace.unwrap_or_else(|err| {
startup_error.get_or_insert_with(|| {
CodexErr::Fatal(format!("failed to start rollout recording: {err:#}"))
});
ThreadTraceContext::disabled()
});
let services = SessionServices {
// Start with an empty connection set. The initialized set is
// published after SessionConfigured so MCP events follow it.
@@ -1896,6 +1904,7 @@ impl Session {
let (mcp_prewarm_tx, mcp_prewarm_rx) = async_channel::bounded(1);
let sess = Arc::new(Session {
thread_id,
agent_runtime,
installation_id,
tx_event: tx_event.clone(),
parent_observer_event_tx: OnceLock::new(),
@@ -1937,13 +1946,11 @@ impl Session {
if let Some(startup) = &startup {
startup.set_session(Arc::clone(&sess));
}
if let Some(err) = trace_error {
if let Some(err) = startup_error {
if startup.is_none() {
super::handlers::shutdown_session_runtime(&sess).await;
}
return Err(
CodexErr::Fatal(format!("failed to start rollout recording: {err:#}")).into(),
);
return Err(err.into());
}
if let Some(network_policy_decider_session) = network_policy_decider_session {
let mut guard = network_policy_decider_session.write().await;
codex-rs/core/src/session/tests.rs
+2
View File
@@ -7143,6 +7143,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
let session = Session {
thread_id,
agent_runtime: crate::agent::host::AgentRuntimeServices::Controller,
installation_id: "11111111-1111-4111-8111-111111111111".to_string(),
tx_event,
parent_observer_event_tx: OnceLock::new(),
@@ -9466,6 +9467,7 @@ where
let session = Arc::new(Session {
thread_id,
agent_runtime: crate::agent::host::AgentRuntimeServices::Controller,
installation_id: "11111111-1111-4111-8111-111111111111".to_string(),
tx_event,
parent_observer_event_tx: OnceLock::new(),
codex-rs/core/src/session/turn.rs
+42 -43
View File
@@ -1,3 +1,4 @@
use std::borrow::Cow;
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;
@@ -970,55 +971,22 @@ pub(crate) async fn run_hooks_and_record_inputs(
input: &[TurnInput],
persist_context: PersistContext,
) -> CodexResult<bool> {
let rejected_uploaded_audio = input.iter().any(turn_input_has_uploaded_audio_file);
// Cancellation can reach this path before Guardian's tools and context are
// resolved. Only finalized evidence may enter reusable reviewer history.
if sess
.services
.thread_extension_data
.get::<crate::guardian::PendingReviewContext>()
.is_some()
{
return if rejected_uploaded_audio {
Err(uploaded_audio_disabled_error())
} else {
Ok(false)
};
}
let mut blocked_input = false;
let mut accepted_user_input = false;
for input_item in input {
let filtered_input;
let input_item = match input_item {
TurnInput::UserInput { content, .. } if content.iter().any(is_uploaded_audio_file) => {
let mut filtered = input_item.clone();
let TurnInput::UserInput { content, .. } = &mut filtered else {
unreachable!("cloned user input must remain user input");
};
content.retain(|item| !is_uploaded_audio_file(item));
if content.is_empty() {
continue;
}
filtered_input = filtered;
&filtered_input
}
_ => input_item,
};
let hook_outcome = inspect_pending_input(sess, turn_context, input_item).await;
for input_item in recordable_inputs(sess, input) {
let input_item = input_item?;
let hook_outcome = inspect_pending_input(sess, turn_context, &input_item).await;
if hook_outcome.should_stop {
blocked_input = true;
record_additional_contexts(sess, turn_context, hook_outcome.additional_contexts).await;
} else {
if let TurnInput::UserInput { content, .. } = input_item {
// Prompt telemetry can contain the full text, so emit it only after hooks accept it.
turn_context.session_telemetry.user_prompt(content);
}
if matches!(input_item, TurnInput::UserInput { content, .. } if !content.is_empty()) {
if matches!(input_item.as_ref(), TurnInput::UserInput { content, .. } if !content.is_empty())
{
accepted_user_input = true;
}
// Tool outputs retain their durability barrier, including in mixed input batches.
let input_persist_context = if persist_context == PersistContext::SteeredUserInput
&& matches!(input_item, TurnInput::FunctionCallOutput(_))
&& matches!(input_item.as_ref(), TurnInput::FunctionCallOutput(_))
{
PersistContext::Standard
} else {
@@ -1028,19 +996,50 @@ pub(crate) async fn run_hooks_and_record_inputs(
sess,
turn_context,
model_info,
input_item.clone(),
input_item.into_owned(),
hook_outcome.additional_contexts,
input_persist_context,
)
.await;
}
}
if rejected_uploaded_audio {
return Err(uploaded_audio_disabled_error());
}
Ok(blocked_input && !accepted_user_input)
}
/// Yields usable input in order, then reports unsupported audio after valid siblings can be recorded.
pub(crate) fn recordable_inputs<'a>(
sess: &Session,
input: &'a [TurnInput],
) -> impl Iterator<Item = CodexResult<Cow<'a, TurnInput>>> {
let rejected_audio = input.iter().any(turn_input_has_uploaded_audio_file);
// Only finalized evidence may enter reusable Guardian history.
let items: &[TurnInput] = if sess
.services
.thread_extension_data
.get::<crate::guardian::PendingReviewContext>()
.is_some()
{
&[]
} else {
input
};
items
.iter()
.filter_map(|item| {
let mut item = Cow::Borrowed(item);
if turn_input_has_uploaded_audio_file(&item)
&& let TurnInput::UserInput { content, .. } = item.to_mut()
{
content.retain(|part| !is_uploaded_audio_file(part));
if content.is_empty() {
return None;
}
}
Some(Ok(item))
})
.chain(rejected_audio.then(|| Err(uploaded_audio_disabled_error())))
}
fn turn_input_has_uploaded_audio_file(input: &TurnInput) -> bool {
matches!(input, TurnInput::UserInput { content, .. } if content.iter().any(is_uploaded_audio_file))
}
codex-rs/core/src/session/turn_suspension.rs
+12 -2
View File
@@ -80,8 +80,17 @@ pub(super) async fn suspend_turn_and_shutdown(
let _ = task_handle.await;
}
}
// Pending accepted input and interactive waiters live only in this process. Handoff
// intentionally drops that state; persisting or replaying it needs a separate protocol.
// The task did not run while admission was pending. Preserve its initial input
// before closing the writer so a new worker can recover the same turn.
let recorded = session
.record_unadmitted_input(
&task.turn_context,
&turn.turn_state,
task.pending_admission_input,
/*abort_reason*/ None,
)
.await;
// Other pending input and interactive waiters remain best effort across handoff.
session.input_queue.clear_pending(&turn).await;
// The root task has stopped without a terminal event. Fence new starts and stop every
@@ -99,6 +108,7 @@ pub(super) async fn suspend_turn_and_shutdown(
live_thread.shutdown().await.map_err(|error| {
CodexErr::Fatal(format!("close suspended root turn writer failed: {error}"))
})?;
recorded?;
// Announce completion only after extension cleanup and writer closure so a
// replacement worker cannot write the same thread concurrently.
session
codex-rs/core/src/session/world_state.rs
+5 -1
View File
@@ -134,7 +134,11 @@ impl Session {
model_instructions,
)
};
let environment_subagents = if turn_context.config.include_environment_context {
let environment_subagents = if turn_context.config.include_environment_context
&& matches!(
self.agent_runtime,
crate::agent::host::AgentRuntimeServices::Controller
) {
match turn_context.multi_agent_version {
MultiAgentVersion::V2 => {
let agent_paths = self
codex-rs/core/src/state/turn.rs
+3
View File
@@ -21,6 +21,7 @@ use tokio::sync::oneshot;
use super::TurnTokenUsage;
use crate::agent::types::AgentExecutionGuard;
use crate::session::TurnInput;
use crate::session::TurnInputQueue;
use crate::session::step_context::StepContext;
use crate::session::turn_context::TurnContext;
@@ -85,6 +86,8 @@ pub(crate) struct RunningTask {
pub(crate) handle: AbortOnDropHandle<()>,
pub(crate) turn_context: Arc<TurnContext>,
pub(crate) _agent_execution_guard: Option<AgentExecutionGuard>,
/// Initial input still owned by Core while the host decides whether the task may run.
pub(crate) pending_admission_input: Option<Vec<TurnInput>>,
pub(crate) _diagnostics_guard: GaugeGuard,
// Timer recorded when the task drops to capture the full turn duration.
pub(crate) _timer: Option<codex_otel::Timer>,
codex-rs/core/src/tasks/admission.rs
+134
View File
@@ -0,0 +1,134 @@
//! Acquires host execution capacity without blocking the submission dispatcher.
//! Input stays with the registered task until admission succeeds, so abort owns cleanup.
use std::sync::Arc;
use codex_async_utils::OrCancelExt;
use codex_protocol::error::CodexErr;
use codex_protocol::error::CodexErrorDetails;
use codex_protocol::error::Result;
use tokio_util::sync::CancellationToken;
use crate::agent::host::AgentTurnAdmission;
use crate::agent::host::AgentTurnAdmissionRequest;
use crate::hook_environment_executor::with_environment_hooks;
use crate::hook_runtime::record_pending_input;
use crate::session::TurnInput;
use crate::session::session::Session;
use crate::session::turn::recordable_inputs;
use crate::session::turn::run_hooks_and_record_inputs;
use crate::session::turn_context::TurnContext;
use crate::state::TaskKind;
use crate::state::TurnState;
use codex_protocol::protocol::HookEventName;
use codex_protocol::protocol::TurnAbortReason;
use codex_thread_store::PersistContext;
use tokio::sync::Mutex;
/// The host admission path already emitted the logical turn's start event.
pub(crate) struct HostTurnStarted;
pub(super) async fn admit(
session: &Arc<Session>,
context: &Arc<TurnContext>,
task_kind: TaskKind,
provider: &dyn AgentTurnAdmission,
cancellation: &CancellationToken,
) -> Result<Vec<TurnInput>> {
// Review forwards the delegate's start event. Other tasks announce the attempt
// here so an unresolved admission remains visible and interruptible.
if task_kind != TaskKind::Review {
session.emit_turn_started(context, task_kind).await;
context.extension_data.insert(HostTurnStarted);
}
let guard = provider
.admit_turn(AgentTurnAdmissionRequest {
thread_id: session.thread_id,
turn_id: context.sub_id.clone(),
})
.or_cancel(cancellation)
.await
.map_err(|_| CodexErr::new(CodexErrorDetails::TurnAborted))??;
// Acquiring this lock waits for turn registration to finish before attaching the permit.
// A cancelled/replaced task must never install its permit into a newer task,
// even if recovery reused the turn ID.
let mut active = session.active_turn.lock().await;
let task = active
.as_mut()
.and_then(|turn| turn.task.as_mut())
.filter(|task| Arc::ptr_eq(&task.turn_context, context))
.filter(|_| !cancellation.is_cancelled())
.ok_or_else(|| CodexErr::new(CodexErrorDetails::TurnAborted))?;
task._agent_execution_guard = Some(guard);
task.pending_admission_input.take().ok_or_else(|| {
CodexErr::Fatal("initial input missing from the task awaiting host admission".into())
})
}
impl Session {
pub(crate) fn host_turn_admission(&self) -> Option<Arc<dyn AgentTurnAdmission>> {
match &self.agent_runtime {
crate::agent::host::AgentRuntimeServices::Controller => None,
crate::agent::host::AgentRuntimeServices::Host(services) => {
Some(Arc::clone(&services.turn_admission))
}
}
}
/// Interruption and suspension stop the task before it can persist its initial input.
pub(crate) async fn record_unadmitted_input(
self: &Arc<Self>,
context: &Arc<TurnContext>,
turn_state: &Mutex<TurnState>,
input: Option<Vec<TurnInput>>,
abort_reason: Option<&TurnAbortReason>,
) -> Result<()> {
let Some(mut input) = input else {
return Ok(());
};
input.extend(
self.input_queue
.take_pending_input_for_turn_state(turn_state)
.await,
);
let model_info = context.capture_current_model_info();
if abort_reason != Some(&TurnAbortReason::Interrupted) {
run_hooks_and_record_inputs(
self,
context,
&model_info,
&input,
PersistContext::Standard,
)
.await?;
return Ok(());
}
for item in recordable_inputs(self, &input) {
let item = item?;
// The user must resubmit a prompt that still needs a submit hook.
if matches!(item.as_ref(), TurnInput::UserInput { .. })
&& with_environment_hooks(
self,
context,
&context.initial_environments.refresh_readiness(),
)
.execution_mode_for_event(HookEventName::UserPromptSubmit)
.is_some()
{
continue;
}
record_pending_input(
self,
context,
&model_info,
item.into_owned(),
Vec::new(),
PersistContext::Standard,
)
.await;
}
Ok(())
}
}
codex-rs/core/src/tasks/mod.rs
+77 -40
View File
@@ -1,3 +1,4 @@
mod admission;
mod compact;
mod lifecycle;
mod regular;
@@ -60,6 +61,7 @@ use codex_protocol::protocol::TurnCompleteEvent;
use codex_protocol::protocol::WarningEvent;
use codex_thread_store::PersistContext;
pub(crate) use admission::HostTurnStarted;
use codex_features::Feature;
use codex_protocol::error::CodexErrorDetails;
use codex_protocol::error::Result as CodexResult;
@@ -354,15 +356,22 @@ impl Session {
(task_kind == TaskKind::Regular).then(|| turn_context.attribution()),
)
.await;
let agent_execution_guard = self.services.agent_control.admit_turn(
turn_context.multi_agent_version,
&turn_context.session_source,
);
let host_admission = self.host_turn_admission();
let (agent_execution_guard, task_input, pending_admission_input) = match &host_admission {
Some(_) => (None, Vec::new(), Some(input)),
None => (
self.services.agent_control.admit_turn(
turn_context.multi_agent_version,
&turn_context.session_source,
),
input,
None,
),
};
let done_clone = Arc::clone(&done);
let session = Arc::clone(self);
let ctx = Arc::clone(&turn_context);
let task_for_run = Arc::clone(&task);
let task_input = input;
let task_cancellation_token = cancellation_token.child_token();
// Task-owned turn spans keep a core-owned span open for the
// full task lifecycle after the submission dispatch span ends.
@@ -385,15 +394,27 @@ impl Session {
let storage_originator = AuthStorageOriginator::from_client_name(&turn_context.originator);
let task_future = async move {
let ctx_for_finish = Arc::clone(&ctx);
let mut task_result = task_for_run
.run(
Arc::clone(&session),
ctx,
task_input,
task_cancellation_token.child_token(),
)
.instrument(trace_span!("session_task.run"))
.await;
let mut task_result = async {
let input = match host_admission {
Some(provider) => admission::admit(
&session,
&ctx,
task_kind,
provider.as_ref(),
&task_cancellation_token,
).await?,
None => task_input,
};
task_for_run
.run(
Arc::clone(&session),
ctx,
input,
task_cancellation_token.child_token(),
)
.instrument(trace_span!("session_task.run"))
.await
}.await;
loop {
// Private reviewers save their transcript together with the terminal event.
// Errors and cancellation retain their existing save path.
@@ -448,6 +469,7 @@ impl Session {
cancellation_token,
turn_context: Arc::clone(&turn_context),
_agent_execution_guard: agent_execution_guard,
pending_admission_input,
_diagnostics_guard: ACTIVE_TURNS.track(),
_timer: timer,
};
@@ -677,33 +699,21 @@ impl Session {
turn_context: Arc<TurnContext>,
task_result: SessionTaskResult,
) -> Option<McpStartupRequirements> {
let (regular_output, abort_reason) = match task_result {
Ok(output) => (output, None),
let (regular_output, abort_reason, task_error) = match task_result {
Ok(output) => (output, None, None),
Err(err) if matches!(err.details(), CodexErrorDetails::TurnAborted) => {
(None, Some(TurnAbortReason::Interrupted))
(None, Some(TurnAbortReason::Interrupted), None)
}
Err(err) => {
warn!(%err, "session task returned an unexpected error");
self.emit_turn_error_lifecycle(
turn_context.as_ref(),
err.to_codex_protocol_error(),
err.details(),
)
.await;
self.track_turn_codex_error(turn_context.as_ref(), &err);
self.send_event(
turn_context.as_ref(),
EventMsg::Error(err.to_error_event(/*message_prefix*/ None)),
)
.await;
(None, None)
(None, None, Some(err))
}
};
let can_continue =
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 = {
let (turn_state, unadmitted_task, mut pending_input) = {
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) {
@@ -720,16 +730,21 @@ impl Session {
let task = active_turn.task.take()?;
task.handle.detach();
active_turn.finishing = Some(finishing);
Arc::clone(&active_turn.turn_state)
(
Arc::clone(&active_turn.turn_state),
task.pending_admission_input.is_some().then_some(task.task),
task.pending_admission_input.unwrap_or_default(),
)
};
turn_context
.turn_metadata_state
.cancel_git_enrichment_task();
let last_agent_message = regular_output.and_then(|output| output.last_agent_message);
let pending_input = self
.input_queue
.take_pending_input_for_turn_state(turn_state.as_ref())
.await;
pending_input.extend(
self.input_queue
.take_pending_input_for_turn_state(turn_state.as_ref())
.await,
);
let (
turn_had_memory_citation,
turn_tool_calls,
@@ -744,15 +759,22 @@ impl Session {
std::mem::take(&mut ts.token_usage_by_model),
)
};
if let Err(err) = run_hooks_and_record_inputs(
let record_result = run_hooks_and_record_inputs(
self,
&turn_context,
&turn_context.capture_current_model_info(),
&pending_input,
PersistContext::Standard,
)
.await
{
.await;
if let Some(task) = unadmitted_task {
// The task never ran. Unwind setup such as entering review mode.
task.abort(Arc::clone(self), Arc::clone(&turn_context))
.await;
}
drop(finish_recording);
// A retry on the first error must start a new turn with the denied turn's input in history.
for err in task_error.into_iter().chain(record_result.err()) {
self.emit_turn_error_lifecycle(
turn_context.as_ref(),
err.to_codex_protocol_error(),
@@ -766,7 +788,6 @@ impl Session {
)
.await;
}
drop(finish_recording);
let turn_telemetry = &turn_context.session_telemetry;
// Emit token usage metrics.
{
@@ -1052,6 +1073,22 @@ impl Session {
// Aborting schedules destruction; join before freezing the turn's timing.
let _ = task.handle.await;
if let Err(error) = self
.record_unadmitted_input(
&task.turn_context,
turn_state,
task.pending_admission_input,
Some(&reason),
)
.await
{
self.send_event(
&task.turn_context,
EventMsg::Error(error.to_error_event(/*message_prefix*/ None)),
)
.await;
}
session_task
.abort(Arc::clone(self), Arc::clone(&task.turn_context))
.await;
codex-rs/core/tests/suite/agent_control_tests.rs
+48
View File
@@ -21,6 +21,12 @@ use codex_core::SpawnRequest;
use codex_core::StartThreadOptions;
use codex_core::ThreadConfigSnapshot;
use codex_core::TurnInputRequest;
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_features::Feature;
use codex_history::InitialHistory;
use codex_history::ResumedHistory;
@@ -699,3 +705,45 @@ async fn agent_mail_wakes_completed_root_when_enabled(
assert!(controller.mail.lock().expect("mail lock").is_empty());
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn host_start_requires_turn_admission() -> anyhow::Result<()> {
struct HostWithoutAdmission(Mutex<Vec<&'static str>>);
impl ThreadLifecycleContributor<Config> for HostWithoutAdmission {
fn on_thread_start<'a>(
&'a self,
input: ThreadStartInput<'a, Config>,
) -> ExtensionFuture<'a, ()> {
Box::pin(async move {
self.0.lock().unwrap().push("start");
input
.thread_store
.insert(codex_core::AgentRuntimeOwnership::Host);
})
}
fn on_thread_stop<'a>(&'a self, _input: ThreadStopInput<'a>) -> ExtensionFuture<'a, ()> {
Box::pin(async move { self.0.lock().unwrap().push("stop") })
}
}
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let host = Arc::new(HostWithoutAdmission(Mutex::default()));
let mut extensions = ExtensionRegistryBuilder::new();
extensions.thread_lifecycle_contributor(host.clone());
let error = test_codex()
.with_extensions(Arc::new(extensions.build()))
.build_with_auto_env(&server)
.await
.err()
.expect("missing admission")
.downcast::<CodexErr>()?;
assert!(
matches!(error.details(), CodexErrorDetails::InvalidRequest(message) if message.contains("AgentTurnAdmission")),
"{error:?}"
);
assert_eq!(*host.0.lock().unwrap(), ["start", "stop"]);
Ok(())
}
codex-rs/core/tests/suite/host_turn_admission.rs
+329
View File
@@ -0,0 +1,329 @@
//! Host capacity is acquired by the ordinary task startup path, including on a root.
use std::sync::Arc;
use std::time::Duration;
use anyhow::Context;
use anyhow::Result;
use codex_core::AgentExecutionGuard;
use codex_core::AgentTurnAdmission;
use codex_core::AgentTurnAdmissionRequest;
use codex_core::CodexThread;
use codex_core::SuspendTurnOutcome;
use codex_core::TurnInputRequest;
use codex_core::TurnInputSubmission;
use codex_history::RolloutItem;
use codex_protocol::dynamic_tools::DynamicToolFunctionSpec;
use codex_protocol::dynamic_tools::DynamicToolResponse;
use codex_protocol::dynamic_tools::DynamicToolSpec;
use codex_protocol::error::CodexErr;
use codex_protocol::error::CodexErrorDetails;
use codex_protocol::error::Result as CodexResult;
use codex_protocol::items::TurnItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::models::input_text;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::Op;
use codex_protocol::protocol::ReviewRequest;
use codex_protocol::protocol::ReviewTarget;
use codex_protocol::user_input::UserInput;
use codex_thread_store::LoadThreadHistoryParams;
use core_test_support::context_snapshot;
use core_test_support::context_snapshot::ContextSnapshotOptions;
use core_test_support::hooks::trust_discovered_hooks;
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::TestCodexBuilder;
use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event;
use core_test_support::wait_for_successful_turn;
use futures::future::BoxFuture;
use pretty_assertions::assert_eq;
use serde_json::json;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;
use wiremock::MockServer;
type Decision = oneshot::Sender<CodexResult<AgentExecutionGuard>>;
type Attempt = (AgentTurnAdmissionRequest, Decision);
struct HostAdmission(mpsc::UnboundedSender<Attempt>);
impl AgentTurnAdmission for HostAdmission {
fn admit_turn(
&self,
request: AgentTurnAdmissionRequest,
) -> BoxFuture<'_, CodexResult<AgentExecutionGuard>> {
Box::pin(async move {
let (reply, decision) = oneshot::channel();
self.0.send((request, reply)).expect("admission receiver");
decision.await.expect("host decision")
})
}
}
async fn root_with_admission(
server: &MockServer,
mut builder: TestCodexBuilder,
) -> Result<(TestCodex, mpsc::UnboundedReceiver<Attempt>)> {
let mut test = builder.build_with_auto_env(server).await?;
let mut options = test.start_thread_options().await;
let (attempts, requests) = mpsc::unbounded_channel();
let admission: Arc<dyn AgentTurnAdmission> = Arc::new(HostAdmission(attempts));
options
.thread_extension_init
.insert(codex_core::AgentRuntimeOwnership::Host);
options.thread_extension_init.insert(admission);
options.dynamic_tools = vec![DynamicToolSpec::Function(DynamicToolFunctionSpec {
name: "hold_turn".into(),
description: "Wait for the test to release this turn.".into(),
input_schema: json!({"type": "object", "properties": {}}),
defer_loading: false,
})];
test.codex.shutdown_and_wait().await?;
let root = test.thread_manager.start_thread(options).await?;
test.codex = root.thread;
test.session_configured = root.session_configured;
Ok((test, requests))
}
async fn submit(thread: &CodexThread, text: &str) -> CodexResult<TurnInputSubmission> {
thread
.start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text {
text: text.into(),
text_elements: vec![],
}]))
.await
}
async fn start(
test: &TestCodex,
attempts: &mut mpsc::UnboundedReceiver<Attempt>,
text: &str,
) -> Result<(String, Decision)> {
let submitted =
tokio::time::timeout(Duration::from_secs(10), submit(&test.codex, text)).await??;
let TurnInputSubmission::Started { turn_id, .. } = submitted else {
anyhow::bail!("expected a new turn, got {submitted:?}");
};
let (request, decision) = tokio::time::timeout(Duration::from_secs(10), attempts.recv())
.await?
.context("admission provider dropped")?;
assert_eq!(request.thread_id, test.session_configured.thread_id);
assert_eq!(request.turn_id, turn_id);
Ok((turn_id, decision))
}
fn deny(decision: Decision) {
let error = CodexErr::new(CodexErrorDetails::AgentLimitReached { max_threads: 1 });
assert!(decision.send(Err(error)).is_ok());
}
fn is_terminal(event: &EventMsg) -> bool {
matches!(event, EventMsg::TurnComplete(_) | EventMsg::TurnAborted(_))
}
fn assert_denied(event: &EventMsg, turn_id: &str) {
assert!(
matches!(event, EventMsg::TurnComplete(done) if done.turn_id == turn_id && done.error.is_some()),
"expected failed turn {turn_id}, got {event:?}"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn root_denial_and_interruption_retain_input_and_grants_span_execution() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let (test, mut attempts) = root_with_admission(&server, test_codex()).await?;
let (turn_id, decision) = start(&test, &mut attempts, "Denied initial prompt.").await?;
assert!(
matches!(submit(&test.codex, "Steered while pending.").await?,
TurnInputSubmission::Steered { turn_id: steered, .. } if steered == turn_id)
);
deny(decision);
let mut recorded = Vec::new();
wait_for_event(&test.codex, |event| {
if let EventMsg::ItemCompleted(done) = event
&& let TurnItem::UserMessage(item) = &done.item
{
recorded.push(item.message());
}
matches!(event, EventMsg::Error(_))
})
.await;
assert_eq!(
recorded,
["Denied initial prompt.", "Steered while pending."]
);
// Retry on the first visible error, while the failed turn may still be finishing.
let (retry_id, pending) = start(&test, &mut attempts, "Interrupted before admission.").await?;
assert_denied(&wait_for_event(&test.codex, is_terminal).await, &turn_id);
test.codex.submit(Op::Interrupt).await?;
let event = wait_for_event(&test.codex, is_terminal).await;
assert!(
matches!(&event, EventMsg::TurnAborted(done) if done.turn_id.as_ref() == Some(&retry_id)),
"expected interrupted turn {retry_id}, got {event:?}"
);
assert!(pending.is_closed());
responses::mount_function_call_agent_response(&server, "hold", "{}", "hold_turn").await;
let (_, grant) = start(&test, &mut attempts, "Continue the work.").await?;
let released = CancellationToken::new();
assert!(
grant
.send(Ok(AgentExecutionGuard::new(released.clone().drop_guard())))
.is_ok()
);
wait_for_event(
&test.codex,
|event| matches!(event, EventMsg::DynamicToolCallRequest(call) if call.call_id == "hold"),
)
.await;
assert!(!released.is_cancelled());
test.codex
.submit(Op::DynamicToolResponse {
id: "hold".into(),
response: DynamicToolResponse {
content_items: vec![],
success: true,
},
})
.await?;
wait_for_successful_turn(&test.codex).await;
assert!(released.is_cancelled());
insta::assert_snapshot!(
"retained_input_after_admission",
context_snapshot::format_request_history_snapshot(
"The admitted turn receives input retained after denial and interruption, then completes its tool call.",
&responses::received_responses_requests(&server).await,
&ContextSnapshotOptions::default().rewrite_known_segments(),
)
);
test.codex.shutdown_and_wait().await?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn interrupt_before_admission_skips_submit_hook_and_uninspected_prompt() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let builder = test_codex()
.with_pre_build_hook(|home| {
let hooks = json!({"hooks": {"UserPromptSubmit": [{"hooks": [
{"type": "command", "command": "echo {}"}
]}]}});
std::fs::write(home.join("hooks.json"), hooks.to_string()).unwrap();
})
.with_config(trust_discovered_hooks);
let (test, mut attempts) = root_with_admission(&server, builder).await?;
let interrupted = "Interrupted before inspection.";
let (turn_id, _pending) = start(&test, &mut attempts, interrupted).await?;
test.codex.submit(Op::Interrupt).await?;
wait_for_event(&test.codex, |event| {
assert!(!matches!(event, EventMsg::HookStarted(_)), "{event:?}");
matches!(event, EventMsg::TurnAborted(done) if done.turn_id.as_ref() == Some(&turn_id))
})
.await;
let model = responses::mount_sse_once(
&server,
responses::sse(vec![responses::ev_completed("done")]),
)
.await;
let (turn_id, grant) = start(&test, &mut attempts, "Continue.").await?;
assert!(grant.send(Ok(AgentExecutionGuard::new(()))).is_ok());
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::HookStarted(start) if start.turn_id.as_ref() == Some(&turn_id))
})
.await;
wait_for_successful_turn(&test.codex).await;
let input = model.single_request().message_input_texts("user");
assert!(input.iter().any(|text| text == "Continue."));
assert!(input.iter().all(|text| text != interrupted));
test.codex.shutdown_and_wait().await?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn suspending_root_during_admission_preserves_the_recoverable_prompt() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let (test, mut attempts) = root_with_admission(&server, test_codex()).await?;
let prompt = "Save this prompt across handoff.";
let (turn_id, pending) = start(&test, &mut attempts, prompt).await?;
assert_eq!(
test.codex.suspend_turn_and_shutdown().await?,
SuspendTurnOutcome::Suspended { turn_id }
);
assert!(pending.is_closed());
let history = test
.thread_store
.load_latest_model_context(LoadThreadHistoryParams {
thread_id: test.session_configured.thread_id,
include_archived: false,
})
.await?;
let expected_content = vec![input_text(prompt)];
let saved_prompts = history
.items
.iter()
.filter(|entry| {
let RolloutItem::ResponseItem(item) = entry else {
return false;
};
matches!(&item.item, ResponseItem::Message { role, content, .. }
if role == "user" && content == &expected_content)
})
.count();
assert_eq!(saved_prompts, 1);
assert!(
responses::received_responses_requests(&server)
.await
.is_empty()
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn denied_review_exits_review_mode() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let (test, mut attempts) = root_with_admission(&server, test_codex()).await?;
test.codex
.submit(Op::Review {
review_request: ReviewRequest {
target: ReviewTarget::Custom {
instructions: "Review this change.".into(),
},
user_facing_hint: None,
},
})
.await?;
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::ItemCompleted(done)
if matches!(done.item, TurnItem::EnteredReviewMode(_)))
})
.await;
let (review, decision) = tokio::time::timeout(Duration::from_secs(10), attempts.recv())
.await?
.context("review admission provider dropped")?;
deny(decision);
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::ItemCompleted(done)
if matches!(done.item, TurnItem::ExitedReviewMode(_)))
})
.await;
assert_denied(
&wait_for_event(&test.codex, is_terminal).await,
&review.turn_id,
);
assert!(
responses::received_responses_requests(&server)
.await
.is_empty()
);
test.codex.shutdown_and_wait().await?;
Ok(())
}
codex-rs/core/tests/suite/scenarios.rs
+3
View File
@@ -155,6 +155,9 @@ mod agent_eviction;
#[path = "scenarios_reserved_spawn.rs"]
mod reserved_spawn;
#[path = "host_turn_admission.rs"]
mod host_turn_admission;
#[path = "scenarios_partial_answers.rs"]
mod partial_answers;
codex-rs/core/tests/suite/scenarios_reserved_spawn.rs
+2
View File
@@ -229,6 +229,7 @@ async fn same_manager_creation_preserves_live_bindings_and_attachments(
let reservation = || HostThreadSpawnReservation {
thread_id: child_id,
control: controller.clone(),
turn_admission: None,
};
let mut mismatched = reservation();
mismatched.control = Arc::new(TestAgentControl::new(ThreadId::new()));
@@ -351,6 +352,7 @@ async fn cancelled_creation_cleans_registered_child() -> Result<()> {
HostThreadSpawnReservation {
thread_id: child_id,
control: controller,
turn_admission: None,
},
)
.await
codex-rs/core/tests/suite/snapshots/all__suite__scenarios__host_turn_admission__retained_input_after_admission.snap
+33
View File
@@ -0,0 +1,33 @@
---
source: core/tests/suite/host_turn_admission.rs
expression: "context_snapshot::format_request_history_snapshot(\"The admitted turn receives input retained after denial and interruption, then completes its tool call.\",\n&responses::received_responses_requests(&server).await,\n&ContextSnapshotOptions::default().rewrite_known_segments(),)"
---
Scenario: The admitted turn receives input retained after denial and interruption, then completes its tool call.
## Window 1
-- request 1 (turn) --
00:message/developer:
<MODEL_INSTRUCTIONS>
01:message/user:
Denied initial prompt.
02:message/user:
Steered while pending.
03:message/user:
Interrupted before admission.
04:message/user:
<turn_aborted>
The user interrupted the previous turn on purpose. Any running unified exec processes may still be runni...mands were aborted, they may have partially executed. [hash=8a46c957943b77f4]
</turn_aborted>
05:message/developer[2]:
[01] <PERMISSIONS_INSTRUCTIONS>
[02] <tools>
Deferred tool namespaces:
- multi_agent_v1: Tools for spawning and managing sub-agents.
</tools>
06:message/user:
<ENVIRONMENT_CONTEXT>
07:message/user:
Continue the work.
-- request 2 (turn) --
08:function_call/hold_turn:{}
09:function_call_output: