From da2803c73cd366b5e01ffe8d0e5f7d396247f827 Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Wed, 12 Aug 2026 00:36:24 +0000 Subject: [PATCH] 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 --- codex-rs/core/src/codex_thread.rs | 75 +------ codex-rs/core/src/hook_runtime.rs | 87 +------- codex-rs/core/src/lib.rs | 1 - codex-rs/core/src/session/inject.rs | 12 -- codex-rs/core/src/session/mod.rs | 4 - codex-rs/core/src/session/tests.rs | 85 -------- .../core/src/session/tests/guardian_tests.rs | 1 - codex-rs/core/src/session/turn.rs | 15 +- codex-rs/core/src/state/turn.rs | 3 - codex-rs/core/src/tasks/mod.rs | 50 +---- codex-rs/core/src/user_message_admission.rs | 199 +----------------- .../tests/suite/user_message_admission.rs | 66 +----- codex-rs/ext/queue/src/service.rs | 22 +- codex-rs/ext/queue/tests/queue_service.rs | 21 +- 14 files changed, 41 insertions(+), 600 deletions(-) diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 310587fad2..d1212f90bf 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -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, client_user_message_id: Option, ) -> CodexResult { - 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, - client_user_message_id: Option, - ) -> Result { - 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, - client_user_message_id: Option, - state: PendingUserMessageAdmissionState, - ) -> Result { 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) }, } } diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index da4a8f2b83..e3628e5b34 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -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, turn_context: &Arc, pending_input: TurnInput, additional_contexts: Vec, 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, - turn_context: &Arc, - 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. diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index b4aebe4a3c..816af473d7 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -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; diff --git a/codex-rs/core/src/session/inject.rs b/codex-rs/core/src/session/inject.rs index c1855cf318..9b73b05370 100644 --- a/codex-rs/core/src/session/inject.rs +++ b/codex-rs/core/src/session/inject.rs @@ -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(()) } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 5cfd3e9f49..3b51352b72 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -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()) } diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 08327d38d0..3a9cec895b 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -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; diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index b9937800a9..c4d31e951f 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -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; diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index e5dad1af6f..bef1dd3398 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -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 { diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index 490dc042d1..056de27955 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -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, pub(crate) kind: TaskKind, pub(crate) task: Arc, - pub(crate) input_persisted: - Option>>, pub(crate) cancellation_token: CancellationToken, pub(crate) handle: AbortOnDropHandle<()>, pub(crate) turn_context: Arc, diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index d4aa3a7256..3418ba83c1 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -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( @@ -300,9 +293,6 @@ impl Session { turn_context: Arc, input: Vec, task: T, - input_persisted: Option< - tokio::sync::oneshot::Sender>, - >, mailbox_parent_provenance: MailboxParentProvenance, ) { let task: Arc = 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, mut task: RunningTask, reason: TurnAbortReason) { + async fn handle_task_abort(self: &Arc, 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 diff --git a/codex-rs/core/src/user_message_admission.rs b/codex-rs/core/src/user_message_admission.rs index 59196b8a6e..e4def5c2e1 100644 --- a/codex-rs/core/src/user_message_admission.rs +++ b/codex-rs/core/src/user_message_admission.rs @@ -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 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>, -} - -struct PendingUserMessageAdmission { - client_id: Option, - sender: oneshot::Sender>, - state: PendingUserMessageAdmissionState, -} - -pub(crate) enum PendingUserMessageAdmissionState { - Immediate, - WaitingForAdmission, - Admitted(UserMessageAdmission), - Persisted, + pending: Mutex>>>, } impl PendingUserMessageAdmissions { pub(crate) fn register( &self, submission_id: String, - client_id: Option, - state: PendingUserMessageAdmissionState, ) -> ( PendingUserMessageAdmissionGuard<'_>, - oneshot::Receiver>, + oneshot::Receiver>, ) { 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, ) { - 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::>(); - - 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() diff --git a/codex-rs/core/tests/suite/user_message_admission.rs b/codex-rs/core/tests/suite/user_message_admission.rs index 355a0cef7b..188acafba9 100644 --- a/codex-rs/core/tests/suite/user_message_admission.rs +++ b/codex-rs/core/tests/suite/user_message_admission.rs @@ -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("")); } -#[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::) - .collect::, _>>() - .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()) diff --git a/codex-rs/ext/queue/src/service.rs b/codex-rs/ext/queue/src/service.rs index 67c6b71958..2d428ac87a 100644 --- a/codex-rs/ext/queue/src/service.rs +++ b/codex-rs/ext/queue/src/service.rs @@ -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, diff --git a/codex-rs/ext/queue/tests/queue_service.rs b/codex-rs/ext/queue/tests/queue_service.rs index 70167f643e..4c0f7a7c07 100644 --- a/codex-rs/ext/queue/tests/queue_service.rs +++ b/codex-rs/ext/queue/tests/queue_service.rs @@ -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?;