From bf7804c25418a71f3dacf43c3a13d8778f1757e4 Mon Sep 17 00:00:00 2001 From: keith thornhill Date: Fri, 31 Jul 2026 20:06:24 +0000 Subject: [PATCH] Add acknowledged user message submission to core (#36385) ## What changed - Add `CodexThread::submit_user_input_and_wait_for_admission` and export the `UserMessageAdmission` result through `codex-core-api`. - Resolve submissions only after they start a new turn or steer the active turn, returning the accepting turn ID in either case. - Return errors for invalid operations, rejected thread settings, and session termination instead of leaving admission waiters unresolved. ## Testing Add integration coverage for concurrent start-and-steer submissions, settings rejection and recovery, non-user operations, and submission after shutdown. GitOrigin-RevId: 21f1ede38d399baa9836d61a59b9f9af2dbf83b8 --- codex-rs/core-api/src/lib.rs | 1 + codex-rs/core/src/codex_thread.rs | 39 ++ codex-rs/core/src/lib.rs | 2 + codex-rs/core/src/session/handlers.rs | 31 +- codex-rs/core/src/session/mod.rs | 4 +- codex-rs/core/src/session/session.rs | 3 + codex-rs/core/src/session/tests.rs | 2 + codex-rs/core/src/user_message_admission.rs | 74 ++++ codex-rs/core/tests/suite/mod.rs | 1 + .../tests/suite/user_message_admission.rs | 353 ++++++++++++++++++ 10 files changed, 500 insertions(+), 10 deletions(-) create mode 100644 codex-rs/core/src/user_message_admission.rs create mode 100644 codex-rs/core/tests/suite/user_message_admission.rs diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index de8031f871..c2bd8c93cd 100644 --- a/codex-rs/core-api/src/lib.rs +++ b/codex-rs/core-api/src/lib.rs @@ -35,6 +35,7 @@ pub use codex_core::StartThreadOptions; pub use codex_core::StateDbHandle; pub use codex_core::ThreadManager; pub use codex_core::ThreadShutdownReport; +pub use codex_core::UserMessageAdmission; pub use codex_core::WaitForEnvironmentToolConfig; pub use codex_core::build_models_manager; pub use codex_core::config::Config; diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 074719114d..5bb0228bb4 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -5,6 +5,7 @@ use crate::session::SessionIo; use crate::session::SessionSettingsUpdate; use crate::session::SteerInputError; use crate::session::session::Session; +use crate::user_message_admission::UserMessageAdmission; use codex_exec_server::SelectedCapabilityRootsStatus; use codex_features::Feature; use codex_otel::SessionTelemetry; @@ -296,6 +297,44 @@ impl CodexThread { .await } + /// Waits until Core has actually started a turn or steered the active turn. + pub async fn submit_user_input_and_wait_for_admission( + &self, + op: Op, + trace: Option, + client_user_message_id: Option, + ) -> CodexResult { + if !matches!(op, Op::UserInput { .. }) { + 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?; + let submission_id = crate::session::new_submission_id(); + let (_pending_admission, admission) = self + .session + .pending_user_message_admissions + .register(submission_id.clone()); + self.io + .submit_with_id(Submission { + id: submission_id.clone(), + op, + client_user_message_id, + trace, + parent_turn_id: None, + }) + .await?; + tokio::select! { + biased; + result = admission => result.map_err(|_| CodexErr::InternalAgentDied)?, + () = self.io.session_loop_termination.clone() => Err(CodexErr::InternalAgentDied), + } + } + /// Persist whether this thread is eligible for future memory generation. pub async fn set_thread_memory_mode(&self, mode: ThreadMemoryMode) -> anyhow::Result<()> { self.session.set_thread_memory_mode(mode).await diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 93d861169b..c21bdc8c8d 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -25,6 +25,7 @@ mod compact_remote; mod compact_remote_v2; mod compact_token_budget; mod config_lock; +mod user_message_admission; pub use codex_thread::BackgroundTerminalInfo; pub use codex_thread::CodexThread; pub use codex_thread::CodexThreadSettingsOverrides; @@ -32,6 +33,7 @@ pub use codex_thread::ThreadConfigSnapshot; pub use codex_thread::TryStartTurnIfIdleError; pub use codex_thread::TryStartTurnIfIdleRejectionReason; pub use session::turn_context::TurnContext; +pub use user_message_admission::UserMessageAdmission; mod agent; mod agent_communication; mod attestation; diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 28fdadf60d..dc8b06f02e 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -22,6 +22,9 @@ use crate::tasks::CompactTask; use crate::tasks::UserShellCommandMode; use crate::tasks::UserShellCommandTask; use crate::tasks::execute_user_shell_command; +use crate::user_message_admission::UserMessageAdmission; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result as CodexResult; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; @@ -85,7 +88,16 @@ pub async fn user_input_or_turn( client_user_message_id: Option, parent_turn_id: Option, ) { - user_input_or_turn_inner(sess, sub_id, op, client_user_message_id, parent_turn_id).await; + let admission = user_input_or_turn_inner( + sess, + sub_id.clone(), + op, + client_user_message_id, + parent_turn_id, + ) + .await; + sess.pending_user_message_admissions + .complete(&sub_id, admission); } pub async fn update_thread_settings( @@ -180,7 +192,7 @@ pub(super) async fn user_input_or_turn_inner( op: Op, client_user_message_id: Option, parent_turn_id: Option, -) { +) -> CodexResult { let Op::UserInput { items, final_output_json_schema, @@ -199,10 +211,8 @@ pub(super) async fn user_input_or_turn_inner( }; updates.final_output_json_schema = Some(final_output_json_schema); - let Ok(current_context) = sess.new_turn_with_sub_id(sub_id.clone(), updates).await else { - // new_turn_with_sub_id already emits the error event. - return; - }; + // new_turn_with_sub_id already emits an error event when settings are invalid. + let current_context = sess.new_turn_with_sub_id(sub_id.clone(), updates).await?; if emit_thread_settings_applied { sess.send_event_raw_without_materializing_rollout(Event { id: sub_id.clone(), @@ -222,8 +232,9 @@ pub(super) async fn user_input_or_turn_inner( ) .await { - Ok(_) => { + Ok(turn_id) => { current_context.session_telemetry.user_prompt(&items); + Ok(UserMessageAdmission::Steered { turn_id }) } Err(SteerInputError::NoActiveTurn(items)) => { if let Some(id) = parent_turn_id { @@ -256,13 +267,17 @@ pub(super) async fn user_input_or_turn_inner( crate::tasks::RegularTask::new(), ) .await; + Ok(UserMessageAdmission::Started { turn_id: sub_id }) } Err(err) => { sess.send_event_raw(Event { - id: sub_id, + id: sub_id.clone(), msg: EventMsg::Error(err.to_error_event()), }) .await; + Err(CodexErr::InvalidRequest(format!( + "failed to admit user message: {err:?}" + ))) } } } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index cb25b5706b..01f9f598cd 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -865,7 +865,7 @@ impl SessionIo { /// /// Some use cases take advantage of the fact that these are UUID7 which /// encodes a timestamp, so think carefully before changing this. -fn new_submission_id() -> String { +pub(crate) fn new_submission_id() -> String { Uuid::now_v7().to_string() } @@ -1190,7 +1190,7 @@ impl Session { } pub(crate) async fn route_realtime_text_input(self: &Arc, text: String) { - handlers::user_input_or_turn_inner( + handlers::user_input_or_turn( self, Uuid::now_v7().to_string(), Op::UserInput { diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index c47f704fbd..4b40cbc2e7 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -52,6 +52,8 @@ pub(crate) struct Session { pub(super) mcp_prewarm_task: std::sync::Mutex>>, pub(crate) conversation: Arc, pub(crate) active_turn: Mutex>, + pub(crate) pending_user_message_admissions: + crate::user_message_admission::PendingUserMessageAdmissions, pub(crate) input_queue: InputQueue, pub(crate) guardian_review_session: GuardianReviewSessionManager, pub(crate) services: SessionServices, @@ -1186,6 +1188,7 @@ impl Session { mcp_prewarm_task: std::sync::Mutex::new(None), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), + pending_user_message_admissions: Default::default(), input_queue: InputQueue::new(), guardian_review_session: GuardianReviewSessionManager::default(), services, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 22f07c0785..17dbe0f726 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5835,6 +5835,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { mcp_prewarm_task: std::sync::Mutex::new(None), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), + pending_user_message_admissions: Default::default(), input_queue: super::input_queue::InputQueue::new(), guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(), services, @@ -8023,6 +8024,7 @@ where mcp_prewarm_task: std::sync::Mutex::new(None), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), + pending_user_message_admissions: Default::default(), input_queue: super::input_queue::InputQueue::new(), guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(), services, diff --git a/codex-rs/core/src/user_message_admission.rs b/codex-rs/core/src/user_message_admission.rs new file mode 100644 index 0000000000..1fded6b5ce --- /dev/null +++ b/codex-rs/core/src/user_message_admission.rs @@ -0,0 +1,74 @@ +use codex_protocol::error::Result as CodexResult; +use std::collections::HashMap; +use std::sync::Mutex; +use tokio::sync::oneshot; + +/// The turn that accepted a submitted user message. +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum UserMessageAdmission { + /// Core installed a new turn for the submitted message. + Started { turn_id: String }, + /// Core accepted the submitted message into an already-running turn. + Steered { turn_id: String }, +} + +#[derive(Default)] +pub(crate) struct PendingUserMessageAdmissions { + pending: Mutex>>>, +} + +impl PendingUserMessageAdmissions { + pub(crate) fn register( + &self, + submission_id: String, + ) -> ( + PendingUserMessageAdmissionGuard<'_>, + oneshot::Receiver>, + ) { + let (sender, receiver) = oneshot::channel(); + self.pending + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(submission_id.clone(), sender); + ( + PendingUserMessageAdmissionGuard { + admissions: self, + submission_id, + }, + receiver, + ) + } + + pub(crate) fn complete( + &self, + submission_id: &str, + admission: CodexResult, + ) { + let admission_sender = self + .pending + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(submission_id); + if let Some(sender) = admission_sender { + let _ = sender.send(admission); + } + } + + fn remove(&self, submission_id: &str) { + self.pending + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(submission_id); + } +} + +pub(crate) struct PendingUserMessageAdmissionGuard<'a> { + admissions: &'a PendingUserMessageAdmissions, + submission_id: String, +} + +impl Drop for PendingUserMessageAdmissionGuard<'_> { + fn drop(&mut self) { + self.admissions.remove(&self.submission_id); + } +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index ea28172dc9..7f7311e9b6 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -151,6 +151,7 @@ mod unified_exec_process_events; #[cfg(unix)] mod unified_exec_zsh_fork_approvals; mod unstable_features_warning; +mod user_message_admission; mod user_notification; mod user_shell_cmd; mod view_image; diff --git a/codex-rs/core/tests/suite/user_message_admission.rs b/codex-rs/core/tests/suite/user_message_admission.rs new file mode 100644 index 0000000000..c88c14a91c --- /dev/null +++ b/codex-rs/core/tests/suite/user_message_admission.rs @@ -0,0 +1,353 @@ +use codex_core::UserMessageAdmission; +use codex_core::config::Constrained; +use codex_protocol::error::CodexErrorDetails; +use codex_protocol::items::TurnItem; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::CodexErrorInfo; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::RolloutItem; +use codex_protocol::protocol::RolloutLine; +use codex_protocol::protocol::ThreadHistoryMode; +use codex_protocol::user_input::UserInput; +use core_test_support::responses; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_response_created; +use core_test_support::streaming_sse::StreamingSseChunk; +use core_test_support::streaming_sse::start_streaming_sse_server; +use core_test_support::test_codex::test_codex; +use pretty_assertions::assert_eq; +use serde_json::Value; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::Barrier; +use tokio::sync::oneshot; +use tokio::time::timeout; + +fn user_input(text: &str) -> Op { + Op::UserInput { + items: vec![UserInput::Text { + text: text.to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + } +} + +/// Concurrent submissions must start exactly one turn, steer the other message, and persist both client IDs. +#[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(); + let (server, _completions) = start_streaming_sse_server(vec![ + vec![ + StreamingSseChunk { + gate: None, + body: responses::sse(vec![ev_response_created("resp-1")]), + }, + StreamingSseChunk { + gate: Some(response_gate), + body: responses::sse(vec![ev_completed("resp-1")]), + }, + ], + vec![StreamingSseChunk { + gate: None, + body: responses::sse(vec![ev_response_created("resp-2"), ev_completed("resp-2")]), + }], + ]) + .await; + let test = test_codex() + .with_model("gpt-5.4") + .with_history_mode(ThreadHistoryMode::Paginated) + .build_with_streaming_server(&server) + .await + .expect("build user-message admission session"); + let codex = Arc::clone(&test.codex); + let barrier = Arc::new(Barrier::new(3)); + + let first_submission = tokio::spawn({ + let codex = Arc::clone(&codex); + let barrier = Arc::clone(&barrier); + async move { + barrier.wait().await; + codex + .submit_user_input_and_wait_for_admission( + user_input("first message"), + /*trace*/ None, + Some("client-message-1".to_string()), + ) + .await + } + }); + let second_submission = tokio::spawn({ + let codex = Arc::clone(&codex); + let barrier = Arc::clone(&barrier); + async move { + barrier.wait().await; + codex + .submit_user_input_and_wait_for_admission( + user_input("second message"), + /*trace*/ None, + Some("client-message-2".to_string()), + ) + .await + } + }); + barrier.wait().await; + + let (first_admission, second_admission) = timeout(Duration::from_secs(5), async { + tokio::join!(first_submission, second_submission) + }) + .await + .expect("both concurrent admissions should resolve before sampling finishes"); + let first_admission = first_admission + .expect("first submission task should finish") + .expect("first user message should be admitted"); + let second_admission = second_admission + .expect("second submission task should finish") + .expect("second user message should be admitted"); + let (started_turn_id, steered_turn_id, started_message) = + match (&first_admission, &second_admission) { + ( + UserMessageAdmission::Started { turn_id: started }, + UserMessageAdmission::Steered { turn_id: steered }, + ) => (started, steered, "first message"), + ( + UserMessageAdmission::Steered { turn_id: steered }, + UserMessageAdmission::Started { turn_id: started }, + ) => (started, steered, "second message"), + _ => panic!( + "concurrent messages must start exactly one turn and steer the other: \ + {first_admission:?}, {second_admission:?}" + ), + }; + assert_eq!(started_turn_id, steered_turn_id); + + release_response + .send(()) + .expect("response gate should remain open"); + let mut observed_client_message_ids = Vec::new(); + loop { + let event = timeout(Duration::from_secs(10), codex.next_event()) + .await + .expect("admitted user-message events should arrive promptly") + .expect("admitted user-message event stream should remain available"); + match event.msg { + EventMsg::ItemStarted(item) => { + if let TurnItem::UserMessage(message) = item.item { + observed_client_message_ids.push(message.client_id); + } + } + EventMsg::Error(error) => panic!("admitted user message failed: {error:?}"), + EventMsg::StreamError(error) => panic!("admitted response stream failed: {error:?}"), + EventMsg::TurnComplete(_) => break, + _ => {} + } + } + observed_client_message_ids.sort(); + assert_eq!( + observed_client_message_ids, + vec![ + Some("client-message-1".to_string()), + Some("client-message-2".to_string()), + ] + ); + + let requests = server.requests().await; + assert_eq!(requests.len(), 2); + let request_bodies: Vec = requests + .iter() + .map(|request| serde_json::from_slice(request).expect("parse model request")) + .collect(); + assert!(request_bodies[0].to_string().contains(started_message)); + assert!(request_bodies[1].to_string().contains("first message")); + assert!(request_bodies[1].to_string().contains("second message")); + + codex + .flush_rollout() + .await + .expect("flush persisted user-message history"); + 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() + ] + ); + server.shutdown().await; +} + +/// Handler-side settings rejection must acknowledge failure once and leave the session usable. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn user_message_admission_reports_invalid_settings_and_recovers() { + let (server, _completions) = start_streaming_sse_server(vec![vec![StreamingSseChunk { + gate: None, + body: responses::sse(vec![ev_response_created("resp-1"), ev_completed("resp-1")]), + }]]) + .await; + let test = test_codex() + .with_model("gpt-5.4") + .with_config(|config| { + config.permissions.approval_policy = Constrained::allow_only(AskForApproval::OnRequest); + }) + .build_with_streaming_server(&server) + .await + .expect("build approval-constrained user-message admission session"); + let codex = &test.codex; + let mut invalid_input = user_input("invalid approval policy"); + let Op::UserInput { + thread_settings, .. + } = &mut invalid_input + else { + unreachable!("test helper always produces a user-input operation"); + }; + thread_settings.approval_policy = Some(AskForApproval::Never); + + let error = timeout( + Duration::from_secs(5), + codex.submit_user_input_and_wait_for_admission( + invalid_input, + /*trace*/ None, + Some("invalid-client-message".to_string()), + ), + ) + .await + .expect("handler-side rejection should resolve promptly") + .expect_err("disallowed approval settings should not be admitted"); + assert!(matches!( + error.details(), + CodexErrorDetails::InvalidRequest(_) + )); + + let error_event = loop { + let event = timeout(Duration::from_secs(5), codex.next_event()) + .await + .expect("invalid settings should emit an error event promptly") + .expect("event stream should remain available"); + match event.msg { + EventMsg::Error(error) => break error, + EventMsg::TurnComplete(_) => { + panic!("invalid settings completed a turn without emitting an error") + } + EventMsg::StreamError(error) => { + panic!("invalid settings unexpectedly started a model stream: {error:?}") + } + _ => {} + } + }; + assert_eq!( + error_event.codex_error_info, + Some(CodexErrorInfo::BadRequest) + ); + + let recovered = timeout( + Duration::from_secs(5), + codex.submit_user_input_and_wait_for_admission( + user_input("valid message after rejected settings"), + /*trace*/ None, + Some("valid-client-message".to_string()), + ), + ) + .await + .expect("session should accept a valid message after rejecting invalid settings") + .expect("valid message should start a turn"); + assert!(matches!(recovered, UserMessageAdmission::Started { .. })); + + loop { + let event = timeout(Duration::from_secs(10), codex.next_event()) + .await + .expect("recovered turn should finish promptly") + .expect("event stream should remain available after rejected settings"); + match event.msg { + EventMsg::Error(error) => panic!("settings rejection emitted another error: {error:?}"), + EventMsg::StreamError(error) => panic!("recovered turn stream failed: {error:?}"), + EventMsg::TurnComplete(_) => break, + _ => {} + } + } + assert_eq!(server.requests().await.len(), 1); + server.shutdown().await; +} + +/// Non-user operations must fail validation immediately instead of waiting for a user-message acknowledgement. +#[tokio::test] +async fn user_message_admission_rejects_non_user_operations_without_waiting() { + let (server, _completions) = start_streaming_sse_server(Vec::new()).await; + let test = test_codex() + .with_model("gpt-5.4") + .build_with_streaming_server(&server) + .await + .expect("build user-message admission session"); + let codex = &test.codex; + + let error = timeout( + Duration::from_secs(5), + codex.submit_user_input_and_wait_for_admission( + Op::Interrupt, + /*trace*/ None, + Some("client-message-1".to_string()), + ), + ) + .await + .expect("non-user operation should be rejected promptly") + .expect_err("non-user operation should not use message admission"); + assert!(matches!( + error.details(), + CodexErrorDetails::InvalidRequest(_) + )); + server.shutdown().await; +} + +/// A terminated session must reject admission immediately instead of leaving a retained waiter unresolved. +#[tokio::test] +async fn user_message_admission_fails_promptly_after_session_shutdown() { + let (server, _completions) = start_streaming_sse_server(Vec::new()).await; + let test = test_codex() + .with_model("gpt-5.4") + .build_with_streaming_server(&server) + .await + .expect("build user-message admission session"); + let codex = &test.codex; + codex + .shutdown_and_wait() + .await + .expect("test session should shut down"); + + let error = timeout( + Duration::from_secs(5), + codex.submit_user_input_and_wait_for_admission( + user_input("after shutdown"), + /*trace*/ None, + Some("client-message-1".to_string()), + ), + ) + .await + .expect("closed session should reject admission promptly") + .expect_err("closed session should not accept user messages"); + assert!(matches!( + error.details(), + CodexErrorDetails::InternalAgentDied + )); + server.shutdown().await; +}