Simplify queued user message admission (#38092)

## What changed

- Resolve user message admission when Core accepts the input as a new turn or a steer, without waiting for rollout persistence.
- Remove persistence- and hook-specific admission errors and task bookkeeping.
- Delete queued messages after Core accepts them for turn processing, including messages subsequently stopped by a prompt hook.

## Testing

- Update concurrent admission and queue service tests for acceptance-based admission.

GitOrigin-RevId: 7dcc311cfecd3d710733b420f2add546edb86153
This commit is contained in:
Owen Lin
2026-08-12 00:41:53 +00:00
committed by copyberry
parent 4c89139da9
commit da2803c73c
14 changed files with 41 additions and 600 deletions
+7 -68
View File
@@ -7,9 +7,7 @@ use crate::session::SessionSettingsUpdate;
use crate::session::SteerInputError;
use crate::session::TurnInput;
use crate::session::session::Session;
use crate::user_message_admission::PendingUserMessageAdmissionState;
use crate::user_message_admission::UserMessageAdmission;
use crate::user_message_admission::UserMessageAdmissionError;
use codex_diagnostics::Gauge;
use codex_diagnostics::GaugeGuard;
use codex_exec_server::SelectedCapabilityRootsStatus;
@@ -109,12 +107,6 @@ pub enum TryStartTurnIfIdleRejectionReason {
/// Another turn or task is active, or the idle reservation was lost before
/// the automatic turn could start.
Busy,
/// A user-prompt hook consumed and rejected the submitted input.
RejectedByHook,
/// The automatic turn ended before its initial input was persisted.
TaskEndedBeforePersistence,
/// The initial input could not be durably written to the rollout.
PersistenceFailed,
}
/// Rejection returned when an extension asks to start automatic idle work but
@@ -329,72 +321,21 @@ impl CodexThread {
trace: Option<W3cTraceContext>,
client_user_message_id: Option<String>,
) -> CodexResult<UserMessageAdmission> {
self.submit_user_input_and_wait_for_admission_inner(
op,
trace,
client_user_message_id,
PendingUserMessageAdmissionState::Immediate,
)
.await
.map_err(Into::into)
}
/// Waits for durable admission and preserves its typed failure outcome.
///
/// A client user-message id is required to identify the persisted message.
pub async fn submit_user_input_and_wait_for_persisted_admission(
&self,
op: Op,
trace: Option<W3cTraceContext>,
client_user_message_id: Option<String>,
) -> Result<UserMessageAdmission, UserMessageAdmissionError> {
if client_user_message_id.is_none() {
return Err(UserMessageAdmissionError::Admission(
CodexErr::InvalidRequest(
"persisted user message admission requires a client user message id"
.to_string(),
),
));
}
if matches!(&op, Op::UserInput { items, .. } if items.is_empty()) {
return Err(UserMessageAdmissionError::Admission(
CodexErr::InvalidRequest(
"user message admission requires nonempty user input".to_string(),
),
));
}
self.submit_user_input_and_wait_for_admission_inner(
op,
trace,
client_user_message_id,
PendingUserMessageAdmissionState::WaitingForAdmission,
)
.await
}
async fn submit_user_input_and_wait_for_admission_inner(
&self,
op: Op,
trace: Option<W3cTraceContext>,
client_user_message_id: Option<String>,
state: PendingUserMessageAdmissionState,
) -> Result<UserMessageAdmission, UserMessageAdmissionError> {
let Op::UserInput { .. } = &op else {
return Err(UserMessageAdmissionError::Admission(
CodexErr::InvalidRequest("user message admission requires user input".to_string()),
return Err(CodexErr::InvalidRequest(
"user message admission requires user input".to_string(),
));
};
self.session
.services
.agent_control
.ensure_execution_capacity_for_op(self.session.thread_id(), &op)
.await
.map_err(UserMessageAdmissionError::Admission)?;
.await?;
let submission_id = crate::session::new_submission_id();
let (_pending_admission, admission) = self
.session
.pending_user_message_admissions
.register(submission_id.clone(), client_user_message_id.clone(), state);
.register(submission_id.clone());
self.io
.submit_with_id(Submission {
id: submission_id.clone(),
@@ -403,14 +344,12 @@ impl CodexThread {
trace,
parent_turn_id: None,
})
.await
.map_err(UserMessageAdmissionError::Admission)?;
.await?;
tokio::select! {
biased;
result = admission => result
.unwrap_or(Err(UserMessageAdmissionError::TaskEndedBeforePersistence)),
result = admission => result.unwrap_or(Err(CodexErr::InternalAgentDied)),
() = self.io.session_loop_termination.clone() => {
Err(UserMessageAdmissionError::TaskEndedBeforePersistence)
Err(CodexErr::InternalAgentDied)
},
}
}
+1 -86
View File
@@ -43,7 +43,6 @@ use codex_thread_store::ReadThreadParams;
use serde_json::Value;
use tracing::instrument;
use crate::codex_thread::TryStartTurnIfIdleRejectionReason;
use crate::context::ContextualUserFragment;
use crate::context::HookAdditionalContext;
use crate::event_mapping::parse_turn_item;
@@ -52,7 +51,6 @@ use crate::session::session::Session;
use crate::session::turn_context::TurnContext;
use crate::tools::hook_names::HookToolName;
use crate::tools::sandboxing::PermissionRequestPayload;
use crate::user_message_admission::UserMessageAdmissionError;
pub(crate) struct HookRuntimeOutcome {
pub should_stop: bool,
@@ -574,72 +572,15 @@ pub(crate) async fn inspect_pending_input(
}
}
#[expect(
clippy::await_holding_invalid_type,
reason = "durable queue admission must settle before its active turn can be aborted"
)]
pub(crate) async fn record_pending_input(
sess: &Arc<Session>,
turn_context: &Arc<TurnContext>,
pending_input: TurnInput,
additional_contexts: Vec<String>,
persist_context: PersistContext,
) -> Result<(), TryStartTurnIfIdleRejectionReason> {
) {
match pending_input {
TurnInput::UserInput { content, client_id } => {
let awaiting_admission = client_id.as_deref().is_some_and(|client_id| {
sess.pending_user_message_admissions
.contains_client_id(client_id)
});
let mut active = sess.active_turn.lock().await;
let input_persisted = active
.as_mut()
.and_then(|turn| turn.task.as_mut())
.filter(|task| Arc::ptr_eq(&task.turn_context, turn_context))
.and_then(|task| task.input_persisted.take());
if awaiting_admission || input_persisted.is_some() {
// An admission acknowledgment always requires synchronous persistence.
sess.record_user_prompt_and_emit_turn_item(
turn_context.as_ref(),
content.as_slice(),
client_id.clone(),
PersistContext::Standard,
)
.await;
record_additional_contexts(sess, turn_context, additional_contexts).await;
match sess.flush_rollout().await {
Ok(()) => {
if awaiting_admission && let Some(client_id) = client_id.as_deref() {
sess.pending_user_message_admissions
.complete_persistence(client_id, Ok(()));
}
if let Some(sender) = input_persisted {
let _ = sender.send(Ok(()));
}
}
Err(error) => {
let error = codex_protocol::error::CodexErr::from(error);
sess.send_event(
turn_context.as_ref(),
EventMsg::Error(error.to_error_event(/*message_prefix*/ None)),
)
.await;
if awaiting_admission && let Some(client_id) = client_id.as_deref() {
sess.pending_user_message_admissions.complete_persistence(
client_id,
Err(UserMessageAdmissionError::PersistenceFailed(error)),
);
}
if let Some(sender) = input_persisted {
let _ = sender
.send(Err(TryStartTurnIfIdleRejectionReason::PersistenceFailed));
}
return Err(TryStartTurnIfIdleRejectionReason::PersistenceFailed);
}
}
return Ok(());
}
drop(active);
sess.record_user_prompt_and_emit_turn_item(
turn_context.as_ref(),
content.as_slice(),
@@ -658,32 +599,6 @@ pub(crate) async fn record_pending_input(
}
}
record_additional_contexts(sess, turn_context, additional_contexts).await;
Ok(())
}
pub(crate) async fn reject_pending_input(
sess: &Arc<Session>,
turn_context: &Arc<TurnContext>,
pending_input: &TurnInput,
) {
let TurnInput::UserInput { client_id, .. } = pending_input else {
return;
};
if let Some(client_id) = client_id.as_deref() {
sess.pending_user_message_admissions
.complete_persistence(client_id, Err(UserMessageAdmissionError::RejectedByHook));
}
let input_persisted = {
let mut active = sess.active_turn.lock().await;
active
.as_mut()
.and_then(|turn| turn.task.as_mut())
.filter(|task| Arc::ptr_eq(&task.turn_context, turn_context))
.and_then(|task| task.input_persisted.take())
};
if let Some(sender) = input_persisted {
let _ = sender.send(Err(TryStartTurnIfIdleRejectionReason::RejectedByHook));
}
}
/// Processes finished async hook results at a safe turn boundary.
-1
View File
@@ -35,7 +35,6 @@ pub use codex_thread::TryStartTurnIfIdleRejectionReason;
pub use environment_config::EnvironmentConfig;
pub use session::turn_context::TurnContext;
pub use user_message_admission::UserMessageAdmission;
pub use user_message_admission::UserMessageAdmissionError;
mod agent;
mod agent_communication;
mod attestation;
-12
View File
@@ -123,9 +123,6 @@ impl Session {
));
}
let (input_persisted_sender, input_persisted_receiver) =
has_user_input.then(tokio::sync::oneshot::channel).unzip();
let original_input = input.clone();
let task_input = if has_user_input {
self.clear_connector_selection().await;
for item in &input {
@@ -144,18 +141,9 @@ impl Session {
turn_context,
task_input,
RegularTask::new(),
input_persisted_sender,
MailboxParentProvenance::Ignore,
)
.await;
if let Some(receiver) = input_persisted_receiver {
return receiver
.await
.unwrap_or(Err(
TryStartTurnIfIdleRejectionReason::TaskEndedBeforePersistence,
))
.map_err(|reason| TryStartTurnIfIdleError::new(reason, original_input));
}
Ok(())
}
-4
View File
@@ -4124,10 +4124,6 @@ impl Session {
pending_input,
)
.await;
if let Some(client_id) = client_user_message_id.as_deref() {
self.pending_user_message_admissions
.associate_steered_by_client_id(client_id, active_turn_id);
}
Ok(active_turn_id.clone())
}
-85
View File
@@ -93,8 +93,6 @@ use crate::tools::handlers::ShellCommandHandler;
use crate::tools::registry::ToolExecutor;
use crate::tools::router::ToolCallSource;
use crate::turn_diff_tracker::TurnDiffTracker;
use crate::user_message_admission::PendingUserMessageAdmissionState;
use crate::user_message_admission::UserMessageAdmissionError;
use codex_config::config_toml::ConfigToml;
use codex_config::config_toml::ProjectConfig;
use codex_config::permissions_toml::FilesystemPermissionToml;
@@ -10015,89 +10013,6 @@ async fn attach_in_memory_thread_store(
store
}
#[tokio::test]
async fn failed_user_message_persistence_stops_turn_processing() {
let (mut session, turn_context, _events) = make_session_and_context_with_rx().await;
open_thread_persistence(
Arc::get_mut(&mut session).unwrap_or_else(|| panic!("session should be uniquely owned")),
)
.await;
session
.services
.live_thread
.as_ref()
.unwrap_or_else(|| panic!("session should have a live thread"))
.shutdown()
.await
.unwrap_or_else(|error| panic!("live thread should shut down: {error}"));
let client_id = "persistence-failure".to_string();
let (_pending_admission, admission) = session.pending_user_message_admissions.register(
turn_context.sub_id.clone(),
Some(client_id.clone()),
PendingUserMessageAdmissionState::WaitingForAdmission,
);
let failed_queued_prompt = vec![UserInput::Text {
text: "message with a failed rollout flush".to_string(),
text_elements: Vec::new(),
}];
let later_steer = vec![UserInput::Text {
text: "later ordinary steer".to_string(),
text_elements: Vec::new(),
}];
let later_response = user_message("later response item");
let later_mailbox = InterAgentCommunication::new(
AgentPath::root().join("worker").expect("worker path"),
AgentPath::root(),
Vec::new(),
"later mailbox message".to_string(),
/*trigger_turn*/ false,
);
let input = [
TurnInput::UserInput {
content: failed_queued_prompt,
client_id: Some(client_id),
},
TurnInput::UserInput {
content: later_steer.clone(),
client_id: None,
},
TurnInput::ResponseItem(later_response.clone()),
TurnInput::InterAgentCommunication(later_mailbox.clone()),
];
assert!(
super::turn::run_hooks_and_record_inputs(
&session,
&turn_context,
&input,
PersistContext::Standard,
)
.await,
"a rollout flush failure should stop turn processing"
);
assert!(matches!(
admission.await,
Ok(Err(UserMessageAdmissionError::PersistenceFailed(_)))
));
assert!(turn_context.terminal_error.lock().await.is_some());
let mut expected_suffix = vec![
session.response_item_from_user_input(later_steer),
later_response,
later_mailbox.to_model_input_item(),
];
for item in &mut expected_suffix {
item.set_turn_id_if_missing(&turn_context.sub_id);
}
let history = session.clone_history().await;
let history_items = strip_response_item_ids(&raw_history_items(&history));
assert!(
history_items.ends_with(&strip_response_item_ids(&expected_suffix)),
"later steer, response, and mailbox inputs should survive the failed queued message"
);
}
#[tokio::test]
async fn hook_transcript_path_does_not_persist_non_local_thread_store() {
let (mut session, _) = make_session_and_context().await;
@@ -463,7 +463,6 @@ async fn strict_auto_review_turn_grant_forces_guardian_for_shell_command_policy_
kind: crate::state::TaskKind::Regular,
listen_to_cancellation_token: true,
},
/*input_persisted*/ None,
crate::tasks::MailboxParentProvenance::Ignore,
)
.await;
+3 -12
View File
@@ -18,7 +18,6 @@ use crate::hook_runtime::drain_async_hook_results;
use crate::hook_runtime::inspect_pending_input;
use crate::hook_runtime::record_additional_contexts;
use crate::hook_runtime::record_pending_input;
use crate::hook_runtime::reject_pending_input;
use crate::hook_runtime::run_legacy_after_agent_hook;
use crate::hook_runtime::run_pending_session_start_hooks;
use crate::hook_runtime::run_turn_stop_hooks;
@@ -603,34 +602,26 @@ pub(crate) async fn run_hooks_and_record_inputs(
) -> bool {
let mut blocked_input = false;
let mut accepted_user_input = false;
let mut persistence_failed = false;
for input_item in input {
let hook_outcome = inspect_pending_input(sess, turn_context, input_item).await;
if hook_outcome.should_stop {
blocked_input = true;
reject_pending_input(sess, turn_context, input_item).await;
record_additional_contexts(sess, turn_context, hook_outcome.additional_contexts).await;
} else {
if matches!(input_item, TurnInput::UserInput { content, .. } if !content.is_empty()) {
accepted_user_input = true;
}
if record_pending_input(
record_pending_input(
sess,
turn_context,
input_item.clone(),
hook_outcome.additional_contexts,
persist_context,
)
.await
.is_err()
{
// Preserve later drained inputs, but stop before sampling so a
// later queue retry cannot execute this message twice.
persistence_failed = true;
}
.await;
}
}
persistence_failed || (blocked_input && !accepted_user_input)
blocked_input && !accepted_user_input
}
fn turn_user_input(input: &[TurnInput]) -> Vec<UserInput> {
-3
View File
@@ -19,7 +19,6 @@ use rmcp::model::RequestId;
use tokio::sync::oneshot;
use crate::agent::control::AgentExecutionGuard;
use crate::codex_thread::TryStartTurnIfIdleRejectionReason;
use crate::mcp_tool_call::McpToolApprovalMetadata;
use crate::session::TurnInputQueue;
use crate::session::turn_context::TurnContext;
@@ -76,8 +75,6 @@ pub(crate) struct RunningTask {
pub(crate) done: Arc<Notify>,
pub(crate) kind: TaskKind,
pub(crate) task: Arc<dyn AnySessionTask>,
pub(crate) input_persisted:
Option<oneshot::Sender<Result<(), TryStartTurnIfIdleRejectionReason>>>,
pub(crate) cancellation_token: CancellationToken,
pub(crate) handle: AbortOnDropHandle<()>,
pub(crate) turn_context: Arc<TurnContext>,
+7 -43
View File
@@ -24,7 +24,6 @@ use tracing::trace_span;
use tracing::warn;
use crate::codex_thread::BackgroundTerminalInfo;
use crate::codex_thread::TryStartTurnIfIdleRejectionReason;
use crate::config::Config;
use crate::context::ContextualUserFragment;
use crate::session::TurnInput;
@@ -285,14 +284,8 @@ impl Session {
) {
self.abort_all_tasks(TurnAbortReason::Replaced).await;
self.clear_connector_selection().await;
self.start_task(
turn_context,
input,
task,
/*input_persisted*/ None,
MailboxParentProvenance::Ignore,
)
.await;
self.start_task(turn_context, input, task, MailboxParentProvenance::Ignore)
.await;
}
pub(crate) async fn start_task<T: SessionTask>(
@@ -300,9 +293,6 @@ impl Session {
turn_context: Arc<TurnContext>,
input: Vec<TurnInput>,
task: T,
input_persisted: Option<
tokio::sync::oneshot::Sender<Result<(), TryStartTurnIfIdleRejectionReason>>,
>,
mailbox_parent_provenance: MailboxParentProvenance,
) {
let task: Arc<dyn AnySessionTask> = Arc::new(task);
@@ -421,7 +411,6 @@ impl Session {
handle: AbortOnDropHandle::new(handle),
kind: task_kind,
task,
input_persisted,
cancellation_token,
turn_context: Arc::clone(&turn_context),
_agent_execution_guard: agent_execution_guard,
@@ -486,7 +475,6 @@ impl Session {
turn_context,
Vec::new(),
RegularTask::new(),
/*input_persisted*/ None,
MailboxParentProvenance::Attribute,
)
.await;
@@ -596,24 +584,12 @@ impl Session {
let turn_state = {
let mut active = self.active_turn.lock().await;
active.as_mut().and_then(|active_turn| {
let mut task = active_turn.task.take()?;
let task_ended_before_persistence =
if let Some(sender) = task.input_persisted.take() {
let _ = sender.send(Err(
TryStartTurnIfIdleRejectionReason::TaskEndedBeforePersistence,
));
true
} else {
false
};
let task = active_turn.task.take()?;
task.handle.detach();
Some((
Arc::clone(&active_turn.turn_state),
task_ended_before_persistence,
))
Some(Arc::clone(&active_turn.turn_state))
})
};
let Some((turn_state, mut task_ended_before_persistence)) = turn_state else {
let Some(turn_state) = turn_state else {
return;
};
let pending_input = self
@@ -635,9 +611,6 @@ impl Session {
PersistContext::Standard,
)
.await;
task_ended_before_persistence |= self
.pending_user_message_admissions
.complete_task_end(&turn_context.sub_id);
// Emit token usage metrics.
{
// TODO(jif): drop this
@@ -787,9 +760,7 @@ impl Session {
});
let idle_cause = if matches!(abort_reason.as_ref(), Some(TurnAbortReason::Interrupted)) {
ThreadIdleCause::Interrupted
} else if task_ended_before_persistence
|| (abort_reason.is_none() && turn_context.terminal_error.lock().await.is_some())
{
} else if abort_reason.is_none() && turn_context.terminal_error.lock().await.is_some() {
ThreadIdleCause::Failed
} else {
ThreadIdleCause::Completed
@@ -877,19 +848,12 @@ impl Session {
.await
}
async fn handle_task_abort(self: &Arc<Self>, mut task: RunningTask, reason: TurnAbortReason) {
async fn handle_task_abort(self: &Arc<Self>, task: RunningTask, reason: TurnAbortReason) {
let sub_id = task.turn_context.sub_id.clone();
if task.cancellation_token.is_cancelled() {
return;
}
if let Some(sender) = task.input_persisted.take() {
let _ = sender.send(Err(
TryStartTurnIfIdleRejectionReason::TaskEndedBeforePersistence,
));
}
self.pending_user_message_admissions
.complete_task_end(&sub_id);
trace!(task_kind = ?task.kind, sub_id, "aborting running task");
task.cancellation_token.cancel();
if reason == TurnAbortReason::Interrupted
+7 -192
View File
@@ -1,4 +1,3 @@
use codex_protocol::error::CodexErr;
use codex_protocol::error::Result as CodexResult;
use std::collections::HashMap;
use std::sync::Mutex;
@@ -13,72 +12,24 @@ pub enum UserMessageAdmission {
Steered { turn_id: String },
}
/// Explains why a submitted user message did not become durable.
#[derive(Debug)]
pub enum UserMessageAdmissionError {
/// Core rejected the submission before admitting it to a turn.
Admission(CodexErr),
/// A user-prompt hook rejected the submitted message.
RejectedByHook,
/// The owning task finished before the submitted message was persisted.
TaskEndedBeforePersistence,
/// The admitted message could not be durably written to the rollout.
PersistenceFailed(CodexErr),
}
impl From<UserMessageAdmissionError> for CodexErr {
fn from(error: UserMessageAdmissionError) -> Self {
match error {
UserMessageAdmissionError::Admission(error) => error,
UserMessageAdmissionError::RejectedByHook => {
Self::InvalidRequest("user message was rejected by a hook".to_string())
}
UserMessageAdmissionError::TaskEndedBeforePersistence => Self::InternalAgentDied,
UserMessageAdmissionError::PersistenceFailed(error) => error,
}
}
}
#[derive(Default)]
pub(crate) struct PendingUserMessageAdmissions {
pending: Mutex<HashMap<String, PendingUserMessageAdmission>>,
}
struct PendingUserMessageAdmission {
client_id: Option<String>,
sender: oneshot::Sender<Result<UserMessageAdmission, UserMessageAdmissionError>>,
state: PendingUserMessageAdmissionState,
}
pub(crate) enum PendingUserMessageAdmissionState {
Immediate,
WaitingForAdmission,
Admitted(UserMessageAdmission),
Persisted,
pending: Mutex<HashMap<String, oneshot::Sender<CodexResult<UserMessageAdmission>>>>,
}
impl PendingUserMessageAdmissions {
pub(crate) fn register(
&self,
submission_id: String,
client_id: Option<String>,
state: PendingUserMessageAdmissionState,
) -> (
PendingUserMessageAdmissionGuard<'_>,
oneshot::Receiver<Result<UserMessageAdmission, UserMessageAdmissionError>>,
oneshot::Receiver<CodexResult<UserMessageAdmission>>,
) {
let (sender, receiver) = oneshot::channel();
self.pending
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(
submission_id.clone(),
PendingUserMessageAdmission {
client_id,
sender,
state,
},
);
.insert(submission_id.clone(), sender);
(
PendingUserMessageAdmissionGuard {
admissions: self,
@@ -93,152 +44,16 @@ impl PendingUserMessageAdmissions {
submission_id: &str,
admission: CodexResult<UserMessageAdmission>,
) {
let mut pending = self
.pending
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let Some(entry) = pending.get_mut(submission_id) else {
return;
};
match admission {
Ok(admission) => match &entry.state {
PendingUserMessageAdmissionState::WaitingForAdmission => {
entry.state = PendingUserMessageAdmissionState::Admitted(admission);
}
PendingUserMessageAdmissionState::Immediate
| PendingUserMessageAdmissionState::Persisted => {
if let Some(entry) = pending.remove(submission_id) {
let _ = entry.sender.send(Ok(admission));
}
}
PendingUserMessageAdmissionState::Admitted(_) => {}
},
Err(error) => {
if let Some(entry) = pending.remove(submission_id) {
let _ = entry
.sender
.send(Err(UserMessageAdmissionError::Admission(error)));
}
}
}
}
pub(crate) fn contains_client_id(&self, client_id: &str) -> bool {
self.pending
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.values()
.any(|entry| {
!matches!(&entry.state, PendingUserMessageAdmissionState::Immediate)
&& entry.client_id.as_deref() == Some(client_id)
})
}
pub(crate) fn associate_steered_by_client_id(&self, client_id: &str, turn_id: &str) {
let submission_id = self
let sender = self
.pending
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.iter()
.find_map(|(submission_id, entry)| {
(matches!(
&entry.state,
PendingUserMessageAdmissionState::WaitingForAdmission
| PendingUserMessageAdmissionState::Persisted
) && entry.client_id.as_deref() == Some(client_id))
.then(|| submission_id.clone())
});
if let Some(submission_id) = submission_id {
self.complete(
&submission_id,
Ok(UserMessageAdmission::Steered {
turn_id: turn_id.to_string(),
}),
);
.remove(submission_id);
if let Some(sender) = sender {
let _ = sender.send(admission);
}
}
pub(crate) fn complete_persistence(
&self,
client_id: &str,
result: Result<(), UserMessageAdmissionError>,
) {
let mut pending = self
.pending
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let Some(submission_id) = pending.iter().find_map(|(submission_id, entry)| {
(!matches!(&entry.state, PendingUserMessageAdmissionState::Immediate)
&& entry.client_id.as_deref() == Some(client_id))
.then(|| submission_id.clone())
}) else {
return;
};
if let Err(error) = result {
if let Some(entry) = pending.remove(&submission_id) {
let _ = entry.sender.send(Err(error));
}
return;
}
let Some(entry) = pending.get_mut(&submission_id) else {
return;
};
match std::mem::replace(
&mut entry.state,
PendingUserMessageAdmissionState::Persisted,
) {
PendingUserMessageAdmissionState::Admitted(admission) => {
if let Some(entry) = pending.remove(&submission_id) {
let _ = entry.sender.send(Ok(admission));
}
}
PendingUserMessageAdmissionState::Immediate
| PendingUserMessageAdmissionState::WaitingForAdmission
| PendingUserMessageAdmissionState::Persisted => {}
}
}
pub(crate) fn complete_task_end(&self, turn_id: &str) -> bool {
let mut pending = self
.pending
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let submission_ids = pending
.iter()
.filter_map(|(submission_id, entry)| {
let belongs_to_turn = match &entry.state {
PendingUserMessageAdmissionState::WaitingForAdmission => {
submission_id == turn_id
}
PendingUserMessageAdmissionState::Admitted(
UserMessageAdmission::Started {
turn_id: admission_turn_id,
}
| UserMessageAdmission::Steered {
turn_id: admission_turn_id,
},
) => admission_turn_id == turn_id,
PendingUserMessageAdmissionState::Immediate
| PendingUserMessageAdmissionState::Persisted => false,
};
belongs_to_turn.then(|| submission_id.clone())
})
.collect::<Vec<_>>();
let had_pending_admissions = !submission_ids.is_empty();
for submission_id in submission_ids {
if let Some(entry) = pending.remove(&submission_id) {
let _ = entry
.sender
.send(Err(UserMessageAdmissionError::TaskEndedBeforePersistence));
}
}
had_pending_admissions
}
fn remove(&self, submission_id: &str) {
self.pending
.lock()
@@ -1,8 +1,5 @@
use codex_core::UserMessageAdmission;
use codex_core::UserMessageAdmissionError;
use codex_core::config::Constrained;
use codex_history::RolloutItem;
use codex_history::RolloutLine;
use codex_protocol::error::CodexErrorDetails;
use codex_protocol::items::TurnItem;
use codex_protocol::models::ContentItem;
@@ -104,34 +101,6 @@ async fn user_message_admission_starts_turn_for_empty_input() {
assert!(user_input_groups[0][0].starts_with("<environment_context>"));
}
#[tokio::test]
async fn persisted_user_message_admission_rejects_empty_input() {
let server = responses::start_mock_server().await;
let test = test_codex()
.with_model("gpt-5.4")
.build_with_auto_env(&server)
.await
.expect("build persisted user-message admission session");
let error = timeout(
Duration::from_secs(5),
test.codex
.submit_user_input_and_wait_for_persisted_admission(
empty_user_input(),
/*trace*/ None,
Some("empty-client-message".to_string()),
),
)
.await
.expect("empty persisted input should be rejected promptly")
.expect_err("empty input cannot produce a persisted user message");
assert!(matches!(
error,
UserMessageAdmissionError::Admission(ref error)
if matches!(error.details(), CodexErrorDetails::InvalidRequest(_))
));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn default_turn_start_persists_developer_and_user_input_before_model_request() {
let (release_response, response_gate) = oneshot::channel();
@@ -189,7 +158,7 @@ async fn default_turn_start_persists_developer_and_user_input_before_model_reque
server.shutdown().await;
}
/// Concurrent submissions must start exactly one turn, steer the other message, and persist both client IDs.
/// Concurrent submissions must start exactly one turn and steer the other message.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn user_message_admission_reports_started_and_steered_for_concurrent_submissions() {
let (release_response, response_gate) = oneshot::channel();
@@ -225,7 +194,7 @@ async fn user_message_admission_reports_started_and_steered_for_concurrent_submi
async move {
barrier.wait().await;
codex
.submit_user_input_and_wait_for_persisted_admission(
.submit_user_input_and_wait_for_admission(
user_input("first message"),
/*trace*/ None,
Some("client-message-1".to_string()),
@@ -239,7 +208,7 @@ async fn user_message_admission_reports_started_and_steered_for_concurrent_submi
async move {
barrier.wait().await;
codex
.submit_user_input_and_wait_for_persisted_admission(
.submit_user_input_and_wait_for_admission(
user_input("second message"),
/*trace*/ None,
Some("client-message-2".to_string()),
@@ -263,7 +232,7 @@ async fn user_message_admission_reports_started_and_steered_for_concurrent_submi
tokio::join!(first_submission, second_submission)
})
.await
.expect("both concurrent admissions should resolve once their messages are persisted");
.expect("both concurrent admissions should resolve");
let first_admission = first_admission
.expect("first submission task should finish")
.expect("first user message should be admitted");
@@ -287,33 +256,6 @@ async fn user_message_admission_reports_started_and_steered_for_concurrent_submi
};
assert_eq!(started_turn_id, steered_turn_id);
let rollout_path = codex.rollout_path().expect("user-message rollout path");
let rollout = tokio::fs::read_to_string(rollout_path)
.await
.expect("read persisted user-message rollout");
let mut persisted_client_message_ids: Vec<_> = rollout
.lines()
.map(serde_json::from_str::<RolloutLine>)
.collect::<Result<Vec<_>, _>>()
.expect("parse persisted user-message rollout")
.into_iter()
.filter_map(|line| match line.item {
RolloutItem::EventMsg(EventMsg::ItemCompleted(event)) => match event.item {
TurnItem::UserMessage(message) => message.client_id,
_ => None,
},
_ => None,
})
.collect();
persisted_client_message_ids.sort();
assert_eq!(
persisted_client_message_ids,
vec![
"client-message-1".to_string(),
"client-message-2".to_string()
]
);
let mut observed_client_message_ids = Vec::new();
loop {
let event = timeout(Duration::from_secs(10), codex.next_event())
+4 -18
View File
@@ -6,10 +6,8 @@ use std::sync::Weak;
use codex_core::CodexThread;
use codex_core::ThreadManager;
use codex_core::TryStartTurnIfIdleRejectionReason;
use codex_core::TurnInput;
use codex_core::UserMessageAdmission;
use codex_core::UserMessageAdmissionError;
use codex_extension_api::ExtensionEventSink;
use codex_extension_api::ExtensionFuture;
use codex_extension_api::ThreadIdleCause;
@@ -191,7 +189,8 @@ impl QueuedItemService {
Ok(())
}
/// Starts or steers the exact queued message and removes it after admission.
/// Starts or steers the exact queued message and removes it after Core
/// accepts the input for turn processing.
pub async fn start(
&self,
thread: &CodexThread,
@@ -212,7 +211,7 @@ impl QueuedItemService {
return Err(QueueServiceError::InvalidInput);
};
let admission = thread
.submit_user_input_and_wait_for_persisted_admission(
.submit_user_input_and_wait_for_admission(
Op::UserInput {
items: content,
final_output_json_schema: None,
@@ -223,16 +222,7 @@ impl QueuedItemService {
trace,
client_id,
)
.await;
let admission = match admission {
Ok(admission) => admission,
Err(error) => {
if matches!(&error, UserMessageAdmissionError::RejectedByHook) {
self.delete_locked(thread_id, queued_item_id).await?;
}
return Err(QueueServiceError::Admission(error.into()));
}
};
.await?;
self.delete_locked(thread_id, queued_item_id).await?;
Ok(admission)
}
@@ -272,10 +262,6 @@ impl QueuedItemService {
}
if let Err(error) = thread.try_start_turn_if_idle(vec![input]).await {
if error.reason() == TryStartTurnIfIdleRejectionReason::RejectedByHook {
self.delete_locked(thread_id, queued_item_id).await?;
continue;
}
tracing::warn!(
%thread_id,
%queued_item_id,
+8 -13
View File
@@ -453,21 +453,16 @@ async fn rejected_queue_messages_are_consumed_without_retrying_or_blocking_follo
assert!(queue.list(thread_id).await?.is_empty());
let rejected = staging.enqueue(thread_id, user_input("blocked")).await?;
let error = tokio::time::timeout(
let admission = tokio::time::timeout(
Duration::from_secs(10),
queue.start(test.codex.as_ref(), rejected.id, /*trace*/ None),
)
.await?
.expect_err("explicitly started blocked input should be rejected");
assert!(matches!(
error,
QueueServiceError::Admission(error)
if matches!(
error.details(),
codex_protocol::error::CodexErrorDetails::InvalidRequest(message)
if message == "user message was rejected by a hook"
)
));
.await??;
assert!(matches!(admission, UserMessageAdmission::Started { .. }));
wait_for_event_match(test.codex.as_ref(), |event| {
matches!(event, EventMsg::TurnComplete(_)).then_some(())
})
.await;
assert!(queue.list(thread_id).await?.is_empty());
let hook_log = std::fs::read_to_string(test.codex_home_path().join("queue_prompt_hook.log"))?;
assert_eq!(
@@ -732,7 +727,7 @@ async fn invalid_head_is_skipped_and_a_live_user_turn_is_accepted() -> anyhow::R
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn resumed_idle_dispatches_input_persisted_without_a_loaded_manager() -> anyhow::Result<()> {
async fn resumed_idle_dispatches_input_without_a_loaded_manager() -> anyhow::Result<()> {
let server = start_mock_server().await;
responses::mount_sse_once(&server, responses::sse_completed("resumed-turn")).await;
let test = test_codex().build_with_auto_env(&server).await?;