diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 70880f31a7..58b932b68e 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2035,6 +2035,7 @@ dependencies = [ "codex-otel", "codex-plugin", "codex-protocol", + "codex-queue-extension", "codex-rmcp-client", "codex-rollout", "codex-sandboxing", @@ -3880,6 +3881,27 @@ dependencies = [ "wildmatch", ] +[[package]] +name = "codex-queue-extension" +version = "0.0.0" +dependencies = [ + "anyhow", + "codex-core", + "codex-extension-api", + "codex-protocol", + "codex-state", + "codex-thread-store", + "codex-utils-absolute-path", + "core_test_support", + "pretty_assertions", + "serde_json", + "tempfile", + "thiserror 2.0.18", + "tokio", + "tracing", + "uuid", +] + [[package]] name = "codex-response-debug-context" version = "0.0.0" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 9f971f7e96..21e9e91f11 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -61,6 +61,7 @@ members = [ "ext/items", "ext/memories", "ext/mcp", + "ext/queue", "ext/skills", "ext/web-search", "external-agent-migration", @@ -226,6 +227,7 @@ codex-model-provider = { path = "model-provider" } codex-process-hardening = { path = "process-hardening" } codex-protocol = { path = "protocol" } codex-prompts = { path = "prompts" } +codex-queue-extension = { path = "ext/queue" } codex-responses-api-proxy = { path = "responses-api-proxy" } codex-response-debug-context = { path = "response-debug-context" } codex-rmcp-client = { path = "rmcp-client" } diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst index 3f5b1db01e..9ec0af06a1 100644 Binary files a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst and b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst differ diff --git a/codex-rs/app-server/Cargo.toml b/codex-rs/app-server/Cargo.toml index 3824d15ccc..0bd41f0968 100644 --- a/codex-rs/app-server/Cargo.toml +++ b/codex-rs/app-server/Cargo.toml @@ -74,6 +74,7 @@ codex-mcp-extension = { workspace = true } codex-model-provider = { workspace = true } codex-models-manager = { workspace = true } codex-protocol = { workspace = true } +codex-queue-extension = { workspace = true } codex-app-server-protocol = { workspace = true } codex-app-server-transport = { workspace = true } codex-feedback = { workspace = true } diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index a14522f62e..d4140102bd 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -1186,6 +1186,7 @@ pub(crate) async fn apply_bespoke_event_handling( )) .await; } + EventMsg::ThreadQueueChanged(_) => {} EventMsg::ThreadSettingsApplied(thread_settings_event) => { let thread_settings = thread_settings_from_core_snapshot(thread_settings_event.thread_settings); diff --git a/codex-rs/app-server/src/extensions.rs b/codex-rs/app-server/src/extensions.rs index c394c8ad81..df0daaee38 100644 --- a/codex-rs/app-server/src/extensions.rs +++ b/codex-rs/app-server/src/extensions.rs @@ -26,7 +26,7 @@ use codex_protocol::error::CodexErr; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_rollout::state_db::StateDbHandle; -use codex_thread_store::ThreadStore; +use codex_thread_store::QueueStore; use crate::outgoing_message::OutgoingMessageSender; use crate::outgoing_message::ThreadScopedOutgoingMessageSender; @@ -44,8 +44,8 @@ pub(crate) struct ThreadExtensionDependencies { pub(crate) executor_skill_provider: Arc, pub(crate) git_attribution_base_url: String, pub(crate) http_client_factory: HttpClientFactory, - /// Process-scoped persistence backend for extensions that need stored thread history. - pub(crate) thread_store: Arc, + /// Process-scoped persistence backend for queued user messages. + pub(crate) queue_store: Option>, } pub(crate) fn thread_extensions( @@ -66,9 +66,17 @@ where executor_skill_provider, git_attribution_base_url, http_client_factory, - thread_store: _thread_store, + queue_store, } = dependencies; - let mut builder = ExtensionRegistryBuilder::::with_event_sink(event_sink); + let mut builder = ExtensionRegistryBuilder::::with_event_sink(Arc::clone(&event_sink)); + if let Some(queue) = queue_store { + let queue_service = Arc::new(codex_queue_extension::QueuedItemService::new( + queue, + thread_manager.clone(), + event_sink, + )); + codex_queue_extension::install(&mut builder, queue_service); + } if let Some(state_db) = state_db { codex_goal_extension::install_with_backend( &mut builder, diff --git a/codex-rs/app-server/src/mcp_refresh.rs b/codex-rs/app-server/src/mcp_refresh.rs index 405f575c16..e1bb838f95 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -339,7 +339,7 @@ enabled = false executor_skill_provider: Arc::clone(&executor_skill_provider), git_attribution_base_url: good_config.chatgpt_base_url.clone(), http_client_factory: good_config.http_client_factory(), - thread_store: Arc::clone(&thread_store), + queue_store: None, }, ), Arc::new(CodexHomeUserInstructionsProvider::new( diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index c11671bb5d..6637345472 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -66,6 +66,7 @@ use codex_chatgpt::workspace_settings; use codex_code_mode::CodeModeSessionProvider; use codex_core::ThreadManager; use codex_core::config::Config; +use codex_core::config::ThreadStoreConfig; use codex_exec_server::EnvironmentManager; use codex_feedback::CodexFeedback; use codex_goal_extension::GoalService; @@ -77,6 +78,8 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::W3cTraceContext; use codex_rollout::StateDbHandle; use codex_state::log_db::LogDbLayer; +use codex_thread_store::LocalQueueStore; +use codex_thread_store::QueueStore; use tokio::sync::Mutex; use tokio::sync::Semaphore; use tokio::sync::broadcast; @@ -247,6 +250,14 @@ impl MessageProcessor { // affect per-thread behavior, but they must not move newly started, // resumed, or forked threads to a different persistence backend/root. let thread_store = codex_core::thread_store_from_config(config.as_ref(), state_db.clone()); + // Queue persistence requires SQLite, so in-memory thread stores and + // app servers without a state database do not have a queue backend. + let queue_store: Option> = match &config.experimental_thread_store { + ThreadStoreConfig::Local => state_db.as_ref().map(|state_db| { + Arc::new(LocalQueueStore::new(Arc::clone(state_db))) as Arc + }), + ThreadStoreConfig::InMemory { .. } => None, + }; let environment_manager_for_requests = Arc::clone(&environment_manager); let environment_manager_for_extensions = Arc::clone(&environment_manager); let restriction_product = session_source.restriction_product(); @@ -281,7 +292,7 @@ impl MessageProcessor { executor_skill_provider: Arc::clone(&executor_skill_provider), git_attribution_base_url: config.chatgpt_base_url.clone(), http_client_factory: config.http_client_factory(), - thread_store: Arc::clone(&thread_store), + queue_store, }, ), Arc::new(CodexHomeUserInstructionsProvider::new( diff --git a/codex-rs/app-server/src/request_processors/thread_goal_processor.rs b/codex-rs/app-server/src/request_processors/thread_goal_processor.rs index 81cb9bf0be..835ae75563 100644 --- a/codex-rs/app-server/src/request_processors/thread_goal_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_goal_processor.rs @@ -65,18 +65,11 @@ impl ThreadGoalRequestProcessor { .map(|()| None) } - pub(crate) async fn emit_resume_goal_snapshot_and_continue( - &self, - thread_id: ThreadId, - thread: &CodexThread, - ) { + pub(crate) async fn emit_resume_goal_snapshot(&self, thread_id: ThreadId) { if !self.config.features.enabled(Feature::Goals) { return; } self.emit_thread_goal_snapshot(thread_id).await; - // App-server owns resume response and snapshot ordering, so wait until - // those are sent before letting extensions react to the idle thread. - thread.emit_thread_idle_lifecycle_if_idle().await; } pub(crate) async fn pending_resume_goal_state( diff --git a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs index 1c9f1668da..1119f1c56b 100644 --- a/codex-rs/app-server/src/request_processors/thread_lifecycle.rs +++ b/codex-rs/app-server/src/request_processors/thread_lifecycle.rs @@ -1,5 +1,6 @@ use super::*; use crate::extensions::send_thread_warning; +use codex_extension_api::ThreadIdleCause; use codex_protocol::config_types::MultiAgentMode; pub(super) const THREAD_UNLOADING_DELAY: Duration = Duration::from_secs(30 * 60); @@ -750,9 +751,9 @@ pub(super) async fn handle_pending_thread_resume_request( .await; // App-server owns resume response and snapshot ordering, so wait until // replay completes before letting extensions react to the idle thread. - if pending.emit_thread_goal_update { - conversation.emit_thread_idle_lifecycle_if_idle().await; - } + conversation + .emit_thread_idle_lifecycle_if_idle(ThreadIdleCause::Completed) + .await; } pub(super) async fn send_thread_goal_snapshot_notification( diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 2164737c21..07c9d1755f 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -8,6 +8,7 @@ use codex_app_server_protocol::ThreadSection; use codex_app_server_protocol::ThreadSectionMoveParams; use codex_app_server_protocol::ThreadSectionMoveResponse; use codex_extension_api::ExtensionDataInit; +use codex_extension_api::ThreadIdleCause; use codex_protocol::config_types::MultiAgentMode; use codex_protocol::error::CodexErrorDetails; use codex_protocol::mcp::ClientMcpExtensions; @@ -3438,7 +3439,10 @@ impl ThreadRequestProcessor { .await; } self.thread_goal_processor - .emit_resume_goal_snapshot_and_continue(thread_id, codex_thread.as_ref()) + .emit_resume_goal_snapshot(thread_id) + .await; + codex_thread + .emit_thread_idle_lifecycle_if_idle(ThreadIdleCause::Completed) .await; } Err(err) => { diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 3e5c3ad948..badcdca87d 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -6,8 +6,11 @@ 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_exec_server::SelectedCapabilityRootsStatus; +use codex_extension_api::ThreadIdleCause; use codex_features::Feature; use codex_otel::SessionTelemetry; use codex_protocol::ThreadId; @@ -100,6 +103,12 @@ 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 @@ -233,6 +242,11 @@ impl CodexThread { self.session.services.session_telemetry.clone() } + /// Returns extension-owned data attached to this thread runtime. + pub fn thread_extension_data(&self) -> &codex_extension_api::ExtensionData { + &self.session.services.thread_extension_data + } + pub async fn shutdown_and_wait(&self) -> CodexResult<()> { self.io.shutdown_and_wait().await } @@ -258,8 +272,8 @@ impl CodexThread { } } - pub async fn emit_thread_idle_lifecycle_if_idle(&self) { - self.session.emit_thread_idle_lifecycle_if_idle().await; + pub async fn emit_thread_idle_lifecycle_if_idle(&self, cause: ThreadIdleCause) { + self.session.emit_thread_idle_lifecycle_if_idle(cause).await; } #[doc(hidden)] @@ -298,28 +312,79 @@ impl CodexThread { .await } - /// Waits until Core has actually started a turn or steered the active turn. + /// Waits until Core has 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.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(), + ), + )); + } + 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 { items, .. } = &op else { + return Err(UserMessageAdmissionError::Admission( + CodexErr::InvalidRequest("user message admission requires user input".to_string()), + )); + }; + if items.is_empty() { + return Err(UserMessageAdmissionError::Admission( + CodexErr::InvalidRequest( + "user message admission requires nonempty user input".to_string(), + ), )); } self.session .services .agent_control .ensure_execution_capacity_for_op(self.session.thread_id(), &op) - .await?; + .await + .map_err(UserMessageAdmissionError::Admission)?; let submission_id = crate::session::new_submission_id(); let (_pending_admission, admission) = self .session .pending_user_message_admissions - .register(submission_id.clone()); + .register(submission_id.clone(), client_user_message_id.clone(), state); self.io .submit_with_id(Submission { id: submission_id.clone(), @@ -328,11 +393,15 @@ impl CodexThread { trace, parent_turn_id: None, }) - .await?; + .await + .map_err(UserMessageAdmissionError::Admission)?; tokio::select! { biased; - result = admission => result.map_err(|_| CodexErr::InternalAgentDied)?, - () = self.io.session_loop_termination.clone() => Err(CodexErr::InternalAgentDied), + result = admission => result + .unwrap_or(Err(UserMessageAdmissionError::TaskEndedBeforePersistence)), + () = self.io.session_loop_termination.clone() => { + Err(UserMessageAdmissionError::TaskEndedBeforePersistence) + }, } } diff --git a/codex-rs/core/src/guardian/review.rs b/codex-rs/core/src/guardian/review.rs index 41446bda71..c1f274c70a 100644 --- a/codex-rs/core/src/guardian/review.rs +++ b/codex-rs/core/src/guardian/review.rs @@ -6,6 +6,7 @@ use codex_analytics::GuardianReviewTerminalStatus; use codex_analytics::GuardianReviewTrackContext; use codex_analytics::GuardianReviewedAction; use codex_core_plugins::PluginCommandAttribution; +use codex_extension_api::ThreadIdleCause; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::openai_models::MODEL_SPECIALTY_CYBER; use codex_protocol::protocol::AskForApproval; @@ -281,7 +282,9 @@ async fn record_guardian_denial(session: &Arc, turn: &Arc, if aborted { // Guardian aborts bypass normal task completion, so emit its idle lifecycle here. // User interrupts deliberately do not take this path. - session.emit_thread_idle_lifecycle_if_idle().await; + session + .emit_thread_idle_lifecycle_if_idle(ThreadIdleCause::Interrupted) + .await; } }); } diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index 420ef3e991..ee4684a8c6 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -39,6 +39,7 @@ 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; @@ -47,6 +48,7 @@ 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, @@ -568,14 +570,69 @@ 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, -) { +) -> 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() { + sess.record_user_prompt_and_emit_turn_item( + turn_context.as_ref(), + content.as_slice(), + client_id.clone(), + ) + .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(), @@ -593,6 +650,32 @@ 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)); + } } async fn run_context_injecting_hook( diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 4cecf46aed..357785cf3d 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -35,6 +35,7 @@ pub use codex_thread::TryStartTurnIfIdleError; pub use codex_thread::TryStartTurnIfIdleRejectionReason; 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 9b73b05370..c1855cf318 100644 --- a/codex-rs/core/src/session/inject.rs +++ b/codex-rs/core/src/session/inject.rs @@ -123,6 +123,9 @@ 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 { @@ -141,9 +144,18 @@ 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/input_queue.rs b/codex-rs/core/src/session/input_queue.rs index d70dbe6306..a133ec106f 100644 --- a/codex-rs/core/src/session/input_queue.rs +++ b/codex-rs/core/src/session/input_queue.rs @@ -4,13 +4,15 @@ use crate::state::TurnState; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::user_input::UserInput; +use serde::Deserialize; +use serde::Serialize; use std::collections::VecDeque; use std::sync::Arc; use tokio::sync::Mutex; use tokio::sync::watch; /// Input consumed by a regular turn. -#[derive(Clone, Debug, PartialEq)] +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub enum TurnInput { UserInput { content: Vec, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 0117b596ff..01c742be89 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1319,6 +1319,15 @@ impl Session { InitialHistory::Resumed(resumed_history) => { let turn_context = self.new_default_turn().await; let rollout_items = resumed_history.history; + if matches!( + rollout_items.iter().rev().find_map(|item| match item { + RolloutItem::EventMsg(event) => agent_status_from_event(event), + _ => None, + }), + Some(AgentStatus::Interrupted) + ) { + self.agent_status.send_replace(AgentStatus::Interrupted); + } let previous_turn_settings = self .apply_rollout_reconstruction(&turn_context, &rollout_items) .await; @@ -4025,7 +4034,7 @@ impl Session { .collect::>(); pending_input.push(TurnInput::UserInput { content: input, - client_id: client_user_message_id, + client_id: client_user_message_id.clone(), }); self.input_queue .extend_pending_input_and_accept_mailbox_delivery_for_turn_state( @@ -4033,6 +4042,10 @@ 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 b2d7a14105..be87b859c7 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -92,6 +92,8 @@ 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; @@ -9865,6 +9867,83 @@ 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).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(history.raw_items()); + 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; @@ -10461,7 +10540,9 @@ async fn thread_idle_lifecycle_waits_for_trigger_turn_mailbox_work() { ) .await; - session.emit_thread_idle_lifecycle_if_idle().await; + session + .emit_thread_idle_lifecycle_if_idle(codex_extension_api::ThreadIdleCause::Completed) + .await; assert_eq!(0, calls.load(std::sync::atomic::Ordering::SeqCst)); } @@ -10960,7 +11041,7 @@ async fn steered_input_reopens_mailbox_delivery_for_current_turn() { text: "follow up".to_string(), text_elements: Vec::new(), }], - client_id: None + client_id: None, }, TurnInput::InterAgentCommunication(communication), ], @@ -11018,7 +11099,7 @@ async fn stale_defer_mailbox_delivery_does_not_override_steered_input() { text: "follow up".to_string(), text_elements: Vec::new(), }], - client_id: None + client_id: None, }, TurnInput::InterAgentCommunication(communication), ], diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 16ef4e7d36..eb18e3b24f 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -20,6 +20,7 @@ use crate::feedback_tags; 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; @@ -583,25 +584,33 @@ 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; } - record_pending_input( + if record_pending_input( sess, turn_context, input_item.clone(), hook_outcome.additional_contexts, ) - .await; + .await + .is_err() + { + // Preserve later drained inputs, but stop before sampling so a + // later queue retry cannot execute this message twice. + persistence_failed = true; + } } } - blocked_input && !accepted_user_input + persistence_failed || (blocked_input && !accepted_user_input) } fn turn_user_input(input: &[TurnInput]) -> Vec { @@ -1750,6 +1759,7 @@ pub(super) fn realtime_text_for_event(msg: &EventMsg) -> Option<(String, Option< | EventMsg::EnvironmentConnected(_) | EventMsg::EnvironmentDisconnected(_) | EventMsg::ThreadGoalUpdated(_) + | EventMsg::ThreadQueueChanged(_) | EventMsg::McpStartupUpdate(_) | EventMsg::McpStartupComplete(_) | EventMsg::McpToolCallBegin(_) diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index 79c103745f..b284199d96 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -18,6 +18,7 @@ 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; @@ -74,6 +75,8 @@ 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/lifecycle.rs b/codex-rs/core/src/tasks/lifecycle.rs index 789955df13..0af3127d95 100644 --- a/codex-rs/core/src/tasks/lifecycle.rs +++ b/codex-rs/core/src/tasks/lifecycle.rs @@ -1,4 +1,5 @@ use codex_extension_api::ExtensionData; +use codex_extension_api::ThreadIdleCause; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TurnAbortReason; @@ -39,7 +40,7 @@ impl Session { } } - pub(crate) async fn emit_thread_idle_lifecycle_if_idle(&self) { + pub(crate) async fn emit_thread_idle_lifecycle_if_idle(&self, cause: ThreadIdleCause) { if self.active_turn.lock().await.is_some() || self.input_queue.has_trigger_turn_mailbox_items().await { @@ -49,6 +50,7 @@ impl Session { for contributor in self.services.extensions.thread_lifecycle_contributors() { contributor .on_thread_idle(codex_extension_api::ThreadIdleInput { + cause, session_store: &self.services.session_extension_data, thread_store: &self.services.thread_extension_data, }) diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index bfd18ef09d..75506c4e85 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use std::time::Duration; use std::time::Instant; +use codex_extension_api::ThreadIdleCause; use futures::future::BoxFuture; use tokio::select; use tokio::sync::Notify; @@ -22,13 +23,12 @@ 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::hook_runtime::inspect_pending_input; -use crate::hook_runtime::record_additional_contexts; -use crate::hook_runtime::record_pending_input; use crate::session::TurnInput; use crate::session::session::Session; +use crate::session::turn::run_hooks_and_record_inputs; use crate::session::turn_context::TurnContext; use crate::state::ActiveTurn; use crate::state::RunningTask; @@ -281,8 +281,14 @@ impl Session { ) { self.abort_all_tasks(TurnAbortReason::Replaced).await; self.clear_connector_selection().await; - self.start_task(turn_context, input, task, MailboxParentProvenance::Ignore) - .await; + self.start_task( + turn_context, + input, + task, + /*input_persisted*/ None, + MailboxParentProvenance::Ignore, + ) + .await; } pub(crate) async fn start_task( @@ -290,6 +296,9 @@ 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); @@ -408,6 +417,7 @@ 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, @@ -471,6 +481,7 @@ impl Session { turn_context, Vec::new(), RegularTask::new(), + /*input_persisted*/ None, MailboxParentProvenance::Attribute, ) .await; @@ -580,12 +591,24 @@ impl Session { let turn_state = { let mut active = self.active_turn.lock().await; active.as_mut().and_then(|active_turn| { - let task = active_turn.task.take()?; + 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 + }; task.handle.detach(); - Some(Arc::clone(&active_turn.turn_state)) + Some(( + Arc::clone(&active_turn.turn_state), + task_ended_before_persistence, + )) }) }; - let Some(turn_state) = turn_state else { + let Some((turn_state, mut task_ended_before_persistence)) = turn_state else { return; }; let pending_input = self @@ -600,28 +623,10 @@ impl Session { ts.token_usage_at_turn_start.clone(), ) }; - if !pending_input.is_empty() { - for pending_input_item in pending_input { - let hook_outcome = - inspect_pending_input(self, &turn_context, &pending_input_item).await; - if hook_outcome.should_stop { - record_additional_contexts( - self, - &turn_context, - hook_outcome.additional_contexts, - ) - .await; - } else { - record_pending_input( - self, - &turn_context, - pending_input_item, - hook_outcome.additional_contexts, - ) - .await; - } - } - } + run_hooks_and_record_inputs(self, &turn_context, &pending_input).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 @@ -764,6 +769,15 @@ impl Session { turn_id: turn_context.sub_id.clone(), profile, }); + 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()) + { + ThreadIdleCause::Failed + } else { + ThreadIdleCause::Completed + }; let event = if let Some(reason) = abort_reason { self.emit_turn_abort_lifecycle(reason.clone(), turn_context.extension_data.as_ref()) .await; @@ -812,7 +826,7 @@ impl Session { } }; if cleared_active_turn { - self.emit_thread_idle_lifecycle_if_idle().await; + self.emit_thread_idle_lifecycle_if_idle(idle_cause).await; } // Regular items were flushed before this terminal event was appended; buffering // thread writers may not flush it without another explicit barrier. @@ -847,12 +861,19 @@ impl Session { .await } - async fn handle_task_abort(self: &Arc, task: RunningTask, reason: TurnAbortReason) { + async fn handle_task_abort(self: &Arc, mut 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(); task.turn_context diff --git a/codex-rs/core/src/user_message_admission.rs b/codex-rs/core/src/user_message_admission.rs index 1fded6b5ce..59196b8a6e 100644 --- a/codex-rs/core/src/user_message_admission.rs +++ b/codex-rs/core/src/user_message_admission.rs @@ -1,3 +1,4 @@ +use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CodexResult; use std::collections::HashMap; use std::sync::Mutex; @@ -12,24 +13,72 @@ 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>>>, + pending: Mutex>, +} + +struct PendingUserMessageAdmission { + client_id: Option, + sender: oneshot::Sender>, + state: PendingUserMessageAdmissionState, +} + +pub(crate) enum PendingUserMessageAdmissionState { + Immediate, + WaitingForAdmission, + Admitted(UserMessageAdmission), + Persisted, } 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(), sender); + .insert( + submission_id.clone(), + PendingUserMessageAdmission { + client_id, + sender, + state, + }, + ); ( PendingUserMessageAdmissionGuard { admissions: self, @@ -44,16 +93,152 @@ impl PendingUserMessageAdmissions { submission_id: &str, admission: CodexResult, ) { - let admission_sender = self + 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 .pending .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) - .remove(submission_id); - if let Some(sender) = admission_sender { - let _ = sender.send(admission); + .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(), + }), + ); } } + 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 c88c14a91c..60f84a0d11 100644 --- a/codex-rs/core/tests/suite/user_message_admission.rs +++ b/codex-rs/core/tests/suite/user_message_admission.rs @@ -73,7 +73,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_admission( + .submit_user_input_and_wait_for_persisted_admission( user_input("first message"), /*trace*/ None, Some("client-message-1".to_string()), @@ -87,7 +87,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_admission( + .submit_user_input_and_wait_for_persisted_admission( user_input("second message"), /*trace*/ None, Some("client-message-2".to_string()), @@ -97,11 +97,21 @@ async fn user_message_admission_reports_started_and_steered_for_concurrent_submi }); barrier.wait().await; + timeout( + Duration::from_secs(5), + server.wait_for_request_count(/*count*/ 1), + ) + .await + .expect("the started turn should reach its first model request"); + release_response + .send(()) + .expect("response gate should remain open"); + 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"); + .expect("both concurrent admissions should resolve once their messages are persisted"); let first_admission = first_admission .expect("first submission task should finish") .expect("first user message should be admitted"); @@ -125,9 +135,33 @@ async fn user_message_admission_reports_started_and_steered_for_concurrent_submi }; assert_eq!(started_turn_id, steered_turn_id); - release_response - .send(()) - .expect("response gate should remain open"); + 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()) @@ -165,36 +199,6 @@ async fn user_message_admission_reports_started_and_steered_for_concurrent_submi 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; } diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 14e0aa32dc..048d614a50 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -29,6 +29,7 @@ pub use prompt::PromptFragment; pub use prompt::PromptSlot; pub use skill_invocation::SkillInvocationInput; pub use skill_invocation::SkillInvocationKind; +pub use thread_lifecycle::ThreadIdleCause; pub use thread_lifecycle::ThreadIdleInput; pub use thread_lifecycle::ThreadOriginator; pub use thread_lifecycle::ThreadResumeInput; diff --git a/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs b/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs index bef595ee19..d2b83bfe79 100644 --- a/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs +++ b/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs @@ -42,8 +42,21 @@ pub struct ThreadResumeInput<'a> { pub thread_store: &'a ExtensionData, } +/// Why a thread has no immediately pending work. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ThreadIdleCause { + /// The previous turn completed and automatic follow-up work can run. + Completed, + /// The user interrupted the previous turn. + Interrupted, + /// The previous turn ended with a terminal error. + Failed, +} + /// Input supplied when the host has no immediately pending thread work. pub struct ThreadIdleInput<'a> { + /// Why the thread became idle. + pub cause: ThreadIdleCause, /// Store scoped to the host session runtime. pub session_store: &'a ExtensionData, /// Store scoped to this thread runtime. diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 8cae0e752b..2f21380b0f 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -47,6 +47,7 @@ pub use contributors::RenderedWorldStateFragment; pub use contributors::SkillInvocationContributor; pub use contributors::SkillInvocationInput; pub use contributors::SkillInvocationKind; +pub use contributors::ThreadIdleCause; pub use contributors::ThreadIdleInput; pub use contributors::ThreadLifecycleContributor; pub use contributors::ThreadOriginator; diff --git a/codex-rs/ext/queue/BUILD.bazel b/codex-rs/ext/queue/BUILD.bazel new file mode 100644 index 0000000000..21f5afd2c9 --- /dev/null +++ b/codex-rs/ext/queue/BUILD.bazel @@ -0,0 +1,6 @@ +load("//:defs.bzl", "codex_rust_crate") + +codex_rust_crate( + name = "queue", + crate_name = "codex_queue_extension", +) diff --git a/codex-rs/ext/queue/Cargo.toml b/codex-rs/ext/queue/Cargo.toml new file mode 100644 index 0000000000..4b5e5d04f4 --- /dev/null +++ b/codex-rs/ext/queue/Cargo.toml @@ -0,0 +1,34 @@ +[package] +edition.workspace = true +license.workspace = true +name = "codex-queue-extension" +version.workspace = true + +[lib] +name = "codex_queue_extension" +path = "src/lib.rs" +test = false +doctest = false + +[lints] +workspace = true + +[dependencies] +codex-core = { workspace = true } +codex-extension-api = { workspace = true } +codex-protocol = { workspace = true } +codex-thread-store = { workspace = true } +serde_json = { workspace = true } +thiserror = { workspace = true } +tokio = { workspace = true, features = ["rt", "sync", "time"] } +tracing = { workspace = true } +uuid = { workspace = true, features = ["v7"] } + +[dev-dependencies] +anyhow = { workspace = true } +codex-state = { workspace = true } +codex-utils-absolute-path = { workspace = true } +core_test_support = { workspace = true } +pretty_assertions = { workspace = true } +tempfile = { workspace = true } +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/codex-rs/ext/queue/src/lib.rs b/codex-rs/ext/queue/src/lib.rs new file mode 100644 index 0000000000..1f7be82988 --- /dev/null +++ b/codex-rs/ext/queue/src/lib.rs @@ -0,0 +1,19 @@ +//! Durable, storage-neutral user-message queue and idle dispatch. + +use std::sync::Arc; + +use codex_extension_api::ExtensionRegistryBuilder; + +mod service; + +pub use service::QueueServiceError; +pub use service::QueuedItem; +pub use service::QueuedItemService; + +/// Registers the caller-owned queue before lower-priority idle contributors. +pub fn install(registry: &mut ExtensionRegistryBuilder, service: Arc) +where + C: Send + Sync + 'static, +{ + registry.thread_lifecycle_contributor(service); +} diff --git a/codex-rs/ext/queue/src/service.rs b/codex-rs/ext/queue/src/service.rs new file mode 100644 index 0000000000..67c6b71958 --- /dev/null +++ b/codex-rs/ext/queue/src/service.rs @@ -0,0 +1,394 @@ +use std::collections::BTreeMap; +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::Mutex as StdMutex; +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; +use codex_extension_api::ThreadIdleInput; +use codex_extension_api::ThreadLifecycleContributor; +use codex_protocol::ThreadId; +use codex_protocol::error::CodexErr; +use codex_protocol::models::snapshot_local_user_input; +use codex_protocol::protocol::AgentStatus; +use codex_protocol::protocol::Event; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::ThreadQueueChangedEvent; +use codex_protocol::protocol::ThreadSettingsOverrides; +use codex_protocol::protocol::W3cTraceContext; +use codex_protocol::user_input::MAX_USER_INPUT_TEXT_CHARS; +use codex_protocol::user_input::UserInput; +use codex_thread_store::MAX_QUEUE_ITEMS; +use codex_thread_store::QueueStore; +use codex_thread_store::QueuedUserSubmissionRecord; +use codex_thread_store::ThreadStoreError; +use thiserror::Error; +use tokio::sync::Mutex; +use tokio::sync::OwnedMutexGuard; +use uuid::Uuid; + +/// One user message waiting to start on its thread. +#[derive(Clone, Debug, PartialEq)] +pub struct QueuedItem { + pub id: String, + pub input: TurnInput, +} + +#[derive(Debug, Error)] +pub enum QueueServiceError { + #[error("queue storage failed: {0}")] + Storage(#[from] ThreadStoreError), + #[error("queued item payload is invalid: {0}")] + InvalidPayload(#[from] serde_json::Error), + #[error("local queued attachment is invalid: {0}")] + InvalidAttachment(#[from] std::io::Error), + #[error("queued user message could not be admitted: {0}")] + Admission(#[from] CodexErr), + #[error("only user input can be added to the user-message queue")] + InvalidInput, + #[error( + "queued user input exceeds the maximum length of {MAX_USER_INPUT_TEXT_CHARS} characters ({actual_chars} provided)" + )] + InputTooLarge { actual_chars: usize }, +} + +#[derive(Clone)] +pub struct QueuedItemService { + queue: Arc, + thread_manager: Weak, + event_sink: Arc, + dispatch_locks: Arc>>>>, +} + +impl QueuedItemService { + pub fn new( + queue: Arc, + thread_manager: Weak, + event_sink: Arc, + ) -> Self { + Self { + queue, + thread_manager, + event_sink, + dispatch_locks: Arc::new(StdMutex::new(HashMap::new())), + } + } + + fn dispatch_lock(&self, thread_id: ThreadId) -> Arc> { + let mut locks = self + .dispatch_locks + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + locks.retain(|_, lock| lock.strong_count() != 0); + if let Some(lock) = locks.get(&thread_id).and_then(Weak::upgrade) { + return lock; + } + let lock = Arc::new(Mutex::new(())); + locks.insert(thread_id, Arc::downgrade(&lock)); + lock + } + + async fn dispatch_guard(&self, thread_id: ThreadId) -> OwnedMutexGuard<()> { + self.dispatch_lock(thread_id).lock_owned().await + } + + pub async fn enqueue( + &self, + thread_id: ThreadId, + input: TurnInput, + ) -> Result { + let input = prepare_queued_user_input(input).await?; + let payload = serde_json::to_string(&input)?; + let item = queued_item_from_record(self.queue.enqueue(thread_id, payload).await?)?; + self.emit_changed(thread_id); + self.wake_if_loaded(thread_id).await; + Ok(item) + } + + pub async fn list(&self, thread_id: ThreadId) -> Result, QueueServiceError> { + self.queue + .list_page(thread_id, /*offset*/ 0, MAX_QUEUE_ITEMS) + .await? + .into_iter() + .map(queued_item_from_record) + .collect() + } + + pub async fn update( + &self, + thread_id: ThreadId, + queued_item_id: String, + mut input: TurnInput, + ) -> Result, QueueServiceError> { + let _dispatch_guard = self.dispatch_guard(thread_id).await; + if let TurnInput::UserInput { client_id, .. } = &mut input + && client_id.is_none() + { + *client_id = self + .list(thread_id) + .await? + .into_iter() + .find_map(|item| match item { + QueuedItem { + id, + input: TurnInput::UserInput { client_id, .. }, + } if id == queued_item_id => client_id, + _ => None, + }); + } + let input = prepare_queued_user_input(input).await?; + let payload = serde_json::to_string(&input)?; + let item = self + .queue + .update(thread_id, queued_item_id, payload) + .await? + .map(queued_item_from_record) + .transpose()?; + if item.is_some() { + self.emit_changed(thread_id); + } + Ok(item) + } + + pub async fn delete( + &self, + thread_id: ThreadId, + queued_item_id: String, + ) -> Result { + let _dispatch_guard = self.dispatch_guard(thread_id).await; + self.delete_locked(thread_id, queued_item_id).await + } + + async fn delete_locked( + &self, + thread_id: ThreadId, + queued_item_id: String, + ) -> Result { + let deleted = self.queue.delete(thread_id, queued_item_id).await?; + if deleted { + self.emit_changed(thread_id); + } + Ok(deleted) + } + + pub async fn reorder( + &self, + thread_id: ThreadId, + ordered_ids: Vec, + ) -> Result<(), QueueServiceError> { + let _dispatch_guard = self.dispatch_guard(thread_id).await; + self.queue.reorder(thread_id, ordered_ids).await?; + self.emit_changed(thread_id); + Ok(()) + } + + /// Starts or steers the exact queued message and removes it after admission. + pub async fn start( + &self, + thread: &CodexThread, + queued_item_id: String, + trace: Option, + ) -> Result { + let thread_id = thread.session_configured().thread_id; + let _dispatch_guard = self.dispatch_guard(thread_id).await; + let item = self + .list(thread_id) + .await? + .into_iter() + .find(|item| item.id == queued_item_id) + .ok_or_else(|| ThreadStoreError::InvalidRequest { + message: format!("queued item not found: {queued_item_id}"), + })?; + let TurnInput::UserInput { content, client_id } = item.input else { + return Err(QueueServiceError::InvalidInput); + }; + let admission = thread + .submit_user_input_and_wait_for_persisted_admission( + Op::UserInput { + items: content, + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: BTreeMap::new(), + thread_settings: ThreadSettingsOverrides::default(), + }, + 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())); + } + }; + self.delete_locked(thread_id, queued_item_id).await?; + Ok(admission) + } + + async fn dispatch_if_idle(&self, thread_id: ThreadId) -> Result<(), QueueServiceError> { + let Some(manager) = self.thread_manager.upgrade() else { + return Ok(()); + }; + let Ok(thread) = manager.get_thread(thread_id).await else { + return Ok(()); + }; + + loop { + let Some(record) = self + .queue + .list_page(thread_id, /*offset*/ 0, /*limit*/ 1) + .await? + .into_iter() + .next() + else { + return Ok(()); + }; + let queued_item_id = record.id.clone(); + + let input = match serde_json::from_str::(&record.payload) { + Ok(input) => input, + Err(error) => { + tracing::warn!(%queued_item_id, %error, "discarding invalid queued item"); + self.delete_locked(thread_id, queued_item_id).await?; + continue; + } + }; + if !matches!(input, TurnInput::UserInput { .. }) { + tracing::warn!(%queued_item_id, "discarding non-user queued input"); + self.delete_locked(thread_id, queued_item_id).await?; + continue; + } + + 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, + reason = ?error.reason(), + "core could not start queued user input" + ); + return Ok(()); + } + + self.delete_locked(thread_id, queued_item_id).await?; + return Ok(()); + } + } + + async fn wake_if_loaded(&self, thread_id: ThreadId) { + let Some(manager) = self.thread_manager.upgrade() else { + return; + }; + if let Ok(thread) = manager.get_thread(thread_id).await + && !matches!(thread.agent_status().await, AgentStatus::Interrupted) + { + thread + .emit_thread_idle_lifecycle_if_idle(ThreadIdleCause::Completed) + .await; + } + } + + fn emit_changed(&self, thread_id: ThreadId) { + self.event_sink.emit(Event { + id: Uuid::now_v7().to_string(), + msg: EventMsg::ThreadQueueChanged(ThreadQueueChangedEvent { thread_id }), + }); + } +} + +async fn prepare_queued_user_input(mut input: TurnInput) -> Result { + validate_queued_user_input(&input)?; + let TurnInput::UserInput { content, client_id } = &mut input else { + return Err(QueueServiceError::InvalidInput); + }; + client_id.get_or_insert_with(|| Uuid::now_v7().to_string()); + if !content.iter().any(|item| { + matches!( + item, + UserInput::LocalImage { .. } | UserInput::LocalAudio { .. } + ) + }) { + return Ok(input); + } + + tokio::task::spawn_blocking(move || { + let mut input = input; + if let TurnInput::UserInput { content, .. } = &mut input { + for item in content { + snapshot_local_user_input(item)?; + } + } + Ok::(input) + }) + .await + .map_err(|error| QueueServiceError::InvalidAttachment(std::io::Error::other(error)))? + .map_err(QueueServiceError::InvalidAttachment) +} + +fn validate_queued_user_input(input: &TurnInput) -> Result<(), QueueServiceError> { + let TurnInput::UserInput { content, .. } = input else { + return Err(QueueServiceError::InvalidInput); + }; + if content.is_empty() { + return Err(QueueServiceError::InvalidInput); + } + let actual_chars: usize = content + .iter() + .filter_map(|item| match item { + UserInput::Text { text, .. } => Some(text.chars().count()), + _ => None, + }) + .sum(); + if actual_chars > MAX_USER_INPUT_TEXT_CHARS { + return Err(QueueServiceError::InputTooLarge { actual_chars }); + } + Ok(()) +} + +impl ThreadLifecycleContributor for QueuedItemService +where + C: Send + Sync + 'static, +{ + fn on_thread_idle<'a>(&'a self, input: ThreadIdleInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + if input.cause != ThreadIdleCause::Completed { + return; + } + let Ok(thread_id) = ThreadId::from_string(input.thread_store.level_id()) else { + tracing::warn!( + level_id = input.thread_store.level_id(), + "queue extension received an invalid thread id" + ); + return; + }; + let _guard = self.dispatch_guard(thread_id).await; + if let Err(error) = self.dispatch_if_idle(thread_id).await { + tracing::warn!(%thread_id, %error, "failed to dispatch queued user input"); + } + }) + } +} + +fn queued_item_from_record( + record: QueuedUserSubmissionRecord, +) -> Result { + Ok(QueuedItem { + id: record.id, + input: serde_json::from_str::(&record.payload)?, + }) +} diff --git a/codex-rs/ext/queue/tests/queue_service.rs b/codex-rs/ext/queue/tests/queue_service.rs new file mode 100644 index 0000000000..70167f643e --- /dev/null +++ b/codex-rs/ext/queue/tests/queue_service.rs @@ -0,0 +1,756 @@ +use std::path::Path; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::OnceLock; +use std::sync::Weak; +use std::time::Duration; + +use anyhow::Context; +use codex_core::TurnInput; +use codex_core::UserMessageAdmission; +use codex_extension_api::ExtensionData; +use codex_extension_api::ExtensionEventSink; +use codex_extension_api::ExtensionFuture; +use codex_extension_api::ExtensionRegistry; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::ExtensionWarning; +use codex_extension_api::NoopExtensionEventSink; +use codex_extension_api::ThreadIdleCause; +use codex_extension_api::ThreadIdleInput; +use codex_extension_api::ThreadLifecycleContributor; +use codex_protocol::ThreadId; +use codex_protocol::items::TurnItem; +use codex_protocol::models::ImageDetail; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::Event; +use codex_protocol::protocol::EventMsg; +use codex_protocol::user_input::MAX_USER_INPUT_TEXT_CHARS; +use codex_protocol::user_input::UserInput; +use codex_queue_extension::QueueServiceError; +use codex_queue_extension::QueuedItemService; +use codex_state::SqliteConfig; +use codex_state::StateRuntime; +use codex_thread_store::LocalQueueStore; +use codex_thread_store::QueueStore; +use codex_thread_store::ThreadStoreError; +use codex_utils_absolute_path::test_support::PathExt; +use core_test_support::hooks::trust_discovered_hooks; +use core_test_support::responses; +use core_test_support::responses::start_mock_server; +use core_test_support::test_codex::TestCodex; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event_match; +use pretty_assertions::assert_eq; +use tempfile::TempDir; + +const TINY_PNG_BYTES: &[u8] = &[ + 137, 80, 78, 71, 13, 10, 26, 10, 0, 0, 0, 13, 73, 72, 68, 82, 0, 0, 0, 1, 0, 0, 0, 1, 8, 6, 0, + 0, 0, 31, 21, 196, 137, 0, 0, 0, 11, 73, 68, 65, 84, 120, 156, 99, 96, 0, 2, 0, 0, 5, 0, 1, + 122, 94, 171, 63, 0, 0, 0, 0, 73, 69, 78, 68, 174, 66, 96, 130, +]; +const TINY_PNG_DATA_URL: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAC0lEQVR4nGNgAAIAAAUAAXpeqz8AAAAASUVORK5CYII="; + +#[derive(Default)] +struct RecordingEventSink { + events: Mutex>, +} + +impl ExtensionEventSink for RecordingEventSink { + fn emit(&self, event: Event) { + self.events + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(event); + } + + fn emit_warning(&self, _warning: ExtensionWarning) {} +} + +#[derive(Default)] +struct InstalledQueue(OnceLock>); + +impl ThreadLifecycleContributor for InstalledQueue { + fn on_thread_idle<'a>(&'a self, input: ThreadIdleInput<'a>) -> ExtensionFuture<'a, ()> { + match self.0.get() { + Some(service) => >::on_thread_idle(service.as_ref(), input), + None => Box::pin(async {}), + } + } +} + +fn registered_queue_extensions() -> ( + Arc, + Arc>, +) { + let installed = Arc::new(InstalledQueue::default()); + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.thread_lifecycle_contributor(installed.clone()); + (installed, Arc::new(extensions.build())) +} + +fn install_registered_queue( + test: &TestCodex, + installed: &InstalledQueue, +) -> anyhow::Result> { + let service = Arc::new(QueuedItemService::new( + loaded_thread_queue(test)?, + Arc::downgrade(&test.thread_manager), + Arc::new(NoopExtensionEventSink), + )); + assert!(installed.0.set(Arc::clone(&service)).is_ok()); + Ok(service) +} + +fn write_rejecting_prompt_hook(home: &Path) { + let script_path = home.join("queue_prompt_hook.py"); + let log_path = home.join("queue_prompt_hook.log"); + let script = format!( + r#"import json +from pathlib import Path +import sys + +payload = json.load(sys.stdin) +with Path(r"{log_path}").open("a", encoding="utf-8") as log: + log.write(payload["prompt"] + "\n") +if payload["prompt"] == "blocked": + print(json.dumps({{"decision": "block", "reason": "blocked by queue hook"}})) +"#, + log_path = log_path.display(), + ); + std::fs::write(&script_path, script) + .unwrap_or_else(|error| panic!("write queue hook script: {error}")); + let hooks = serde_json::json!({ + "hooks": { + "UserPromptSubmit": [{ + "hooks": [{ + "type": "command", + "command": format!("python3 {}", script_path.display()), + }] + }] + } + }); + std::fs::write(home.join("hooks.json"), hooks.to_string()) + .unwrap_or_else(|error| panic!("write queue hooks: {error}")); +} + +async fn test_queue() -> anyhow::Result<(Arc, TempDir)> { + let home = tempfile::tempdir()?; + let sqlite = SqliteConfig::new_for_testing(home.path().abs()); + let runtime = StateRuntime::init(sqlite.clone(), "test-provider".to_string()).await?; + let queue: Arc = Arc::new(LocalQueueStore::new(runtime)); + Ok((queue, home)) +} + +fn loaded_thread_queue(test: &TestCodex) -> anyhow::Result> { + let runtime = test.codex.state_db().context("state runtime unavailable")?; + Ok(Arc::new(LocalQueueStore::new(runtime))) +} + +fn user_input(text: &str) -> TurnInput { + TurnInput::UserInput { + content: vec![UserInput::Text { + text: text.to_string(), + text_elements: Vec::new(), + }], + client_id: None, + } +} + +fn structured_user_input(text: &str) -> TurnInput { + TurnInput::UserInput { + content: vec![UserInput::Text { + text: text.to_string(), + text_elements: Vec::new(), + }], + client_id: Some("stable-client-message".to_string()), + } +} + +fn user_input_with_media(text: &str, image: UserInput, audio: UserInput) -> TurnInput { + let mut input = structured_user_input(text); + if let TurnInput::UserInput { content, .. } = &mut input { + content.extend([image, audio]); + } + input +} + +async fn emit_idle(service: &QueuedItemService, thread_id: ThreadId) { + emit_idle_with_cause(service, thread_id, ThreadIdleCause::Completed).await; +} + +async fn emit_idle_with_cause( + service: &QueuedItemService, + thread_id: ThreadId, + cause: ThreadIdleCause, +) { + let session_store = ExtensionData::new("session"); + let thread_store = ExtensionData::new(thread_id.to_string()); + >::on_thread_idle( + service, + ThreadIdleInput { + cause, + session_store: &session_store, + thread_store: &thread_store, + }, + ) + .await; +} + +#[tokio::test] +async fn queued_input_and_unique_event_ids_round_trip() -> anyhow::Result<()> { + let (queue, _home) = test_queue().await?; + let sink = Arc::new(RecordingEventSink::default()); + let event_sink: Arc = sink.clone(); + let service = QueuedItemService::new(queue, Weak::new(), event_sink); + let thread_id = ThreadId::new(); + let input = structured_user_input("structured message"); + + let first = service.enqueue(thread_id, input.clone()).await?; + let second = service.enqueue(thread_id, user_input("next")).await?; + let TurnInput::UserInput { + client_id: Some(generated_client_id), + .. + } = &second.input + else { + anyhow::bail!("queued message did not receive a client message id"); + }; + assert_eq!( + 7, + uuid::Uuid::parse_str(generated_client_id)?.get_version_num() + ); + assert_eq!(vec![first.clone(), second], service.list(thread_id).await?); + assert_eq!(input, first.input); + + let events = sink + .events + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + assert_eq!(2, events.len()); + for event in events.iter() { + let EventMsg::ThreadQueueChanged(change) = &event.msg else { + anyhow::bail!("event is not a queue change"); + }; + assert_eq!( + serde_json::json!({ "threadId": thread_id }), + serde_json::to_value(change)? + ); + } + assert_ne!(events[0].id, events[1].id); + Ok(()) +} + +#[tokio::test] +async fn editing_reordering_and_deleting_preserve_queue_identity() -> anyhow::Result<()> { + let (queue, _home) = test_queue().await?; + let service = QueuedItemService::new(queue, Weak::new(), Arc::new(NoopExtensionEventSink)); + let thread_id = ThreadId::new(); + let first = service + .enqueue(thread_id, structured_user_input("first")) + .await?; + let second = service.enqueue(thread_id, user_input("second")).await?; + let edited = service + .update(thread_id, first.id.clone(), user_input("edited")) + .await? + .context("queued item missing")?; + assert_eq!(first.id, edited.id); + let TurnInput::UserInput { client_id, .. } = &edited.input else { + anyhow::bail!("edited queue item does not contain user input"); + }; + assert_eq!(Some("stable-client-message"), client_id.as_deref()); + assert!(matches!( + service.reorder(thread_id, vec![first.id.clone()]).await, + Err(QueueServiceError::Storage( + ThreadStoreError::InvalidRequest { .. } + )) + )); + service + .reorder(thread_id, vec![second.id.clone(), first.id.clone()]) + .await?; + assert_eq!( + vec![second.clone(), edited.clone()], + service.list(thread_id).await? + ); + assert!(service.delete(thread_id, second.id).await?); + assert_eq!(vec![edited], service.list(thread_id).await?); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn starting_a_selected_item_preserves_the_remaining_queue() -> anyhow::Result<()> { + let server = start_mock_server().await; + let response = + responses::mount_sse_once(&server, responses::sse_completed("selected-turn")).await; + let test = test_codex().build_with_auto_env(&server).await?; + let thread_id = test.session_configured.thread_id; + let queue = loaded_thread_queue(&test)?; + let service = QueuedItemService::new(queue, Weak::new(), Arc::new(NoopExtensionEventSink)); + let first = service.enqueue(thread_id, user_input("first")).await?; + let second = service + .enqueue(thread_id, structured_user_input("second")) + .await?; + + let admission = service + .start(test.codex.as_ref(), second.id.clone(), /*trace*/ None) + .await?; + + assert!(matches!(admission, UserMessageAdmission::Started { .. })); + assert_eq!(vec![first], service.list(thread_id).await?); + wait_for_event_match(test.codex.as_ref(), |event| match event { + EventMsg::TurnComplete(_) => Some(()), + _ => None, + }) + .await; + assert_eq!( + Some("second"), + response + .single_request() + .message_input_texts("user") + .last() + .map(String::as_str) + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn stopped_turns_leave_queued_messages_in_place() -> anyhow::Result<()> { + let server = start_mock_server().await; + let test = test_codex().build_with_auto_env(&server).await?; + let thread_id = test.session_configured.thread_id; + let queue = loaded_thread_queue(&test)?; + let service = QueuedItemService::new( + queue, + Arc::downgrade(&test.thread_manager), + Arc::new(NoopExtensionEventSink), + ); + let queued = service + .enqueue(thread_id, user_input("keep queued after stopping")) + .await?; + + for cause in [ThreadIdleCause::Interrupted, ThreadIdleCause::Failed] { + emit_idle_with_cause(&service, thread_id, cause).await; + + assert_eq!(vec![queued.clone()], service.list(thread_id).await?); + } + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn registered_queue_lifecycle_starts_messages_in_fifo_order() -> anyhow::Result<()> { + let server = start_mock_server().await; + let responses = responses::mount_sse_sequence( + &server, + ["turn-a", "turn-b", "turn-c"] + .into_iter() + .map(responses::sse_completed) + .collect(), + ) + .await; + let (installed, extensions) = registered_queue_extensions(); + let test = test_codex() + .with_extensions(extensions) + .with_config(|config| config.include_environment_context = false) + .build_with_auto_env(&server) + .await?; + let thread_id = test.session_configured.thread_id; + let staging = QueuedItemService::new( + loaded_thread_queue(&test)?, + Weak::new(), + Arc::new(NoopExtensionEventSink), + ); + for prompt in ["B", "C"] { + staging.enqueue(thread_id, user_input(prompt)).await?; + } + let queue = install_registered_queue(&test, installed.as_ref())?; + + tokio::time::timeout(Duration::from_secs(10), async { + test.submit_text_turn("A").await?; + for _ in 0..2 { + wait_for_event_match(test.codex.as_ref(), |event| { + matches!(event, EventMsg::TurnComplete(_)).then_some(()) + }) + .await; + } + anyhow::Ok(()) + }) + .await??; + + let prompts = responses + .requests() + .into_iter() + .map(|request| { + request + .message_input_texts("user") + .pop() + .context("model request has no user input") + }) + .collect::>>()?; + assert_eq!(vec!["A", "B", "C"], prompts); + assert!(queue.list(thread_id).await?.is_empty()); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn rejected_queue_messages_are_consumed_without_retrying_or_blocking_followups() +-> anyhow::Result<()> { + let server = start_mock_server().await; + let responses = responses::mount_sse_sequence( + &server, + ["initial-turn", "following-turn"] + .into_iter() + .map(responses::sse_completed) + .collect(), + ) + .await; + let (installed, extensions) = registered_queue_extensions(); + let test = test_codex() + .with_extensions(extensions) + .with_pre_build_hook(write_rejecting_prompt_hook) + .with_config(trust_discovered_hooks) + .with_config(|config| config.include_environment_context = false) + .build_with_auto_env(&server) + .await?; + let thread_id = test.session_configured.thread_id; + let staging = QueuedItemService::new( + loaded_thread_queue(&test)?, + Weak::new(), + Arc::new(NoopExtensionEventSink), + ); + for prompt in ["blocked", "C"] { + staging.enqueue(thread_id, user_input(prompt)).await?; + } + let queue = install_registered_queue(&test, installed.as_ref())?; + + tokio::time::timeout(Duration::from_secs(10), async { + test.submit_text_turn("A").await?; + for _ in 0..2 { + wait_for_event_match(test.codex.as_ref(), |event| { + matches!(event, EventMsg::TurnComplete(_)).then_some(()) + }) + .await; + } + anyhow::Ok(()) + }) + .await??; + + let prompts = responses + .requests() + .into_iter() + .map(|request| { + request + .message_input_texts("user") + .pop() + .context("model request has no user input") + }) + .collect::>>()?; + assert_eq!(vec!["A", "C"], prompts); + let hook_log = std::fs::read_to_string(test.codex_home_path().join("queue_prompt_hook.log"))?; + assert_eq!( + vec!["A", "blocked", "C"], + hook_log.lines().collect::>() + ); + assert!(queue.list(thread_id).await?.is_empty()); + + let rejected = staging.enqueue(thread_id, user_input("blocked")).await?; + let error = 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" + ) + )); + 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!( + vec!["A", "blocked", "C", "blocked"], + hook_log.lines().collect::>() + ); + assert_eq!(2, responses.requests().len()); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn local_attachments_are_snapshotted_before_enqueue_and_update() -> anyhow::Result<()> { + let (queue, home) = test_queue().await?; + let service = QueuedItemService::new( + Arc::clone(&queue), + Weak::new(), + Arc::new(NoopExtensionEventSink), + ); + let thread_id = ThreadId::new(); + let image_path = home.path().join("queued.png"); + let audio_path = home.path().join("queued.mp3"); + std::fs::write(&image_path, TINY_PNG_BYTES)?; + std::fs::write(&audio_path, b"audio")?; + let queued_input = user_input_with_media( + "queued attachments", + UserInput::LocalImage { + path: image_path.clone(), + detail: Some(ImageDetail::Original), + }, + UserInput::LocalAudio { + path: audio_path.clone(), + }, + ); + let expected_queued_input = user_input_with_media( + "queued attachments", + UserInput::Image { + image_url: TINY_PNG_DATA_URL.to_string(), + detail: Some(ImageDetail::Original), + }, + UserInput::Audio { + audio_url: "data:audio/mpeg;base64,YXVkaW8=".to_string(), + }, + ); + + let queued = service.enqueue(thread_id, queued_input).await?; + std::fs::remove_file(&image_path)?; + std::fs::remove_file(&audio_path)?; + + assert_eq!(expected_queued_input, queued.input); + assert_eq!(vec![queued.clone()], service.list(thread_id).await?); + let queued_record = queue + .list_page(thread_id, /*offset*/ 0, /*limit*/ 1) + .await? + .into_iter() + .next() + .context("snapshotted queued attachments were not persisted")?; + let persisted_queued_input: TurnInput = serde_json::from_str(&queued_record.payload)?; + assert_eq!(expected_queued_input, persisted_queued_input); + + let edited_image_path = home.path().join("edited.png"); + let edited_audio_path = home.path().join("edited.m4a"); + std::fs::write(&edited_image_path, TINY_PNG_BYTES)?; + std::fs::write(&edited_audio_path, b"edited audio")?; + let edited_input = user_input_with_media( + "edited attachments", + UserInput::LocalImage { + path: edited_image_path.clone(), + detail: Some(ImageDetail::High), + }, + UserInput::LocalAudio { + path: edited_audio_path.clone(), + }, + ); + let expected_edited_input = user_input_with_media( + "edited attachments", + UserInput::Image { + image_url: TINY_PNG_DATA_URL.to_string(), + detail: Some(ImageDetail::High), + }, + UserInput::Audio { + audio_url: "data:audio/mp4;base64,ZWRpdGVkIGF1ZGlv".to_string(), + }, + ); + + let edited = service + .update(thread_id, queued.id.clone(), edited_input) + .await? + .context("snapshotted queued attachments were not updated")?; + std::fs::remove_file(&edited_image_path)?; + std::fs::remove_file(&edited_audio_path)?; + + assert_eq!(expected_edited_input, edited.input); + assert_eq!(vec![edited.clone()], service.list(thread_id).await?); + let edited_record = queue + .list_page(thread_id, /*offset*/ 0, /*limit*/ 1) + .await? + .into_iter() + .next() + .context("snapshotted edited attachments were not persisted")?; + let persisted_edited_input: TurnInput = serde_json::from_str(&edited_record.payload)?; + assert_eq!(expected_edited_input, persisted_edited_input); + + Ok(()) +} + +#[tokio::test] +async fn invalid_local_attachments_do_not_mutate_queue() -> anyhow::Result<()> { + let (queue, home) = test_queue().await?; + let service = QueuedItemService::new(queue, Weak::new(), Arc::new(NoopExtensionEventSink)); + let thread_id = ThreadId::new(); + let existing = service.enqueue(thread_id, user_input("existing")).await?; + let missing_path = home.path().join("missing.png"); + let invalid_input = user_input_with_media( + "invalid attachments", + UserInput::LocalImage { + path: missing_path, + detail: Some(ImageDetail::Original), + }, + UserInput::Audio { + audio_url: "data:audio/mpeg;base64,YXVkaW8=".to_string(), + }, + ); + + assert!(matches!( + service + .enqueue(thread_id, invalid_input.clone()) + .await, + Err(QueueServiceError::InvalidAttachment(error)) + if error.kind() == std::io::ErrorKind::NotFound + )); + assert!(matches!( + service + .update( + thread_id, + existing.id.clone(), + invalid_input, + ) + .await, + Err(QueueServiceError::InvalidAttachment(error)) + if error.kind() == std::io::ErrorKind::NotFound + )); + assert_eq!(vec![existing], service.list(thread_id).await?); + + Ok(()) +} + +#[tokio::test] +async fn non_user_input_cannot_enter_the_user_message_queue() -> anyhow::Result<()> { + let (queue, _home) = test_queue().await?; + let service = QueuedItemService::new(queue, Weak::new(), Arc::new(NoopExtensionEventSink)); + let error = service + .enqueue( + ThreadId::new(), + TurnInput::ResponseItem(ResponseItem::Other), + ) + .await + .expect_err("response item should not enter the user queue"); + assert!(matches!(error, QueueServiceError::InvalidInput)); + Ok(()) +} + +#[tokio::test] +async fn queued_text_limit_is_enforced_across_all_input_items() -> anyhow::Result<()> { + let (queue, _home) = test_queue().await?; + let service = QueuedItemService::new(queue, Weak::new(), Arc::new(NoopExtensionEventSink)); + let thread_id = ThreadId::new(); + let existing = service.enqueue(thread_id, user_input("existing")).await?; + let mut oversized = user_input(&"x".repeat(MAX_USER_INPUT_TEXT_CHARS / 2)); + if let TurnInput::UserInput { content, .. } = &mut oversized { + content.push(UserInput::Text { + text: "y".repeat(MAX_USER_INPUT_TEXT_CHARS / 2 + 1), + text_elements: Vec::new(), + }); + } + + assert!(matches!( + service + .enqueue(thread_id, oversized.clone()) + .await, + Err(QueueServiceError::InputTooLarge { actual_chars }) + if actual_chars == MAX_USER_INPUT_TEXT_CHARS + 1 + )); + assert!(matches!( + service + .update( + thread_id, + existing.id.clone(), + oversized, + ) + .await, + Err(QueueServiceError::InputTooLarge { actual_chars }) + if actual_chars == MAX_USER_INPUT_TEXT_CHARS + 1 + )); + assert_eq!(vec![existing], service.list(thread_id).await?); + + Ok(()) +} + +#[tokio::test] +async fn queued_text_limit_counts_characters_and_ignores_non_text_items() -> anyhow::Result<()> { + let (queue, _home) = test_queue().await?; + let service = QueuedItemService::new(queue, Weak::new(), Arc::new(NoopExtensionEventSink)); + let mut input = structured_user_input(&"é".repeat(MAX_USER_INPUT_TEXT_CHARS)); + if let TurnInput::UserInput { content, .. } = &mut input { + content.push(UserInput::Mention { + name: "Demo App".to_string(), + path: "app://demo-app".to_string(), + }); + } + + let queued = service.enqueue(ThreadId::new(), input.clone()).await?; + assert_eq!(input, queued.input); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn invalid_head_is_skipped_and_a_live_user_turn_is_accepted() -> anyhow::Result<()> { + let server = start_mock_server().await; + let response = + responses::mount_sse_once(&server, responses::sse_completed("queued-turn")).await; + let test = test_codex() + .with_config(|config| config.include_environment_context = false) + .build_with_auto_env(&server) + .await?; + let thread_id = test.session_configured.thread_id; + let queue = loaded_thread_queue(&test)?; + queue + .enqueue(thread_id, r#"{"unsupported":true}"#.to_string()) + .await?; + let service = QueuedItemService::new( + queue, + Arc::downgrade(&test.thread_manager), + Arc::new(NoopExtensionEventSink), + ); + + service + .enqueue(thread_id, structured_user_input("durable follow-up")) + .await?; + emit_idle(&service, thread_id).await; + let client_id = wait_for_event_match(test.codex.as_ref(), |event| match event { + EventMsg::ItemCompleted(event) => match &event.item { + TurnItem::UserMessage(item) => Some(item.client_id.clone()), + _ => None, + }, + _ => None, + }) + .await; + assert_eq!(Some("stable-client-message".to_string()), client_id); + wait_for_event_match(test.codex.as_ref(), |event| match event { + EventMsg::TurnComplete(_) => Some(()), + _ => None, + }) + .await; + assert!(service.list(thread_id).await?.is_empty()); + let request = response.single_request(); + assert_eq!( + vec!["durable follow-up"], + request.message_input_texts("user") + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn resumed_idle_dispatches_input_persisted_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?; + let thread_id = test.session_configured.thread_id; + let queue = loaded_thread_queue(&test)?; + QueuedItemService::new( + Arc::clone(&queue), + Weak::new(), + Arc::new(NoopExtensionEventSink), + ) + .enqueue(thread_id, user_input("queued while unloaded")) + .await?; + let service = QueuedItemService::new( + queue, + Arc::downgrade(&test.thread_manager), + Arc::new(NoopExtensionEventSink), + ); + emit_idle(&service, thread_id).await; + assert!(service.list(thread_id).await?.is_empty()); + Ok(()) +} diff --git a/codex-rs/mcp-server/src/codex_tool_runner.rs b/codex-rs/mcp-server/src/codex_tool_runner.rs index dd551a580a..e056e9b5c5 100644 --- a/codex-rs/mcp-server/src/codex_tool_runner.rs +++ b/codex-rs/mcp-server/src/codex_tool_runner.rs @@ -306,8 +306,8 @@ async fn run_codex_tool_session_inner( EventMsg::SessionConfigured(_) => { tracing::error!("unexpected SessionConfigured event"); } - EventMsg::ThreadGoalUpdated(_) => { - // Ignore thread goal metadata updates in MCP tool runner. + EventMsg::ThreadGoalUpdated(_) | EventMsg::ThreadQueueChanged(_) => { + // Ignore thread-scoped metadata updates in MCP tool runner. } EventMsg::McpStartupUpdate(_) | EventMsg::McpStartupComplete(_) => { // Ignored in MCP tool runner. diff --git a/codex-rs/protocol/src/lib.rs b/codex-rs/protocol/src/lib.rs index 5a969d4661..28fff2120a 100644 --- a/codex-rs/protocol/src/lib.rs +++ b/codex-rs/protocol/src/lib.rs @@ -20,6 +20,7 @@ pub mod error; pub mod exec_output; pub mod items; mod legacy_events; +pub mod local_media; pub mod mcp; pub mod mcp_approval_meta; pub mod memory_citation; diff --git a/codex-rs/protocol/src/local_media.rs b/codex-rs/protocol/src/local_media.rs new file mode 100644 index 0000000000..7e4cbd9dfd --- /dev/null +++ b/codex-rs/protocol/src/local_media.rs @@ -0,0 +1,93 @@ +use std::io; +use std::io::Read; +use std::path::Path; + +use codex_utils_image::MAX_PROMPT_IMAGE_INPUT_BYTES; +use codex_utils_image::PromptImageMode; +use codex_utils_image::data_url_from_bytes; +use codex_utils_image::load_for_prompt_bytes; + +use crate::models::ImageDetail; +use crate::user_input::UserInput; + +/// Maximum accepted decoded byte length for prompt audio inputs. +/// +/// This matches the Responses API audio input limit. +pub const MAX_PROMPT_AUDIO_INPUT_BYTES: usize = 50 * 1024 * 1024; + +/// Snapshots local image and audio input into portable, validated data URLs. +pub fn snapshot_local_user_input(input: &mut UserInput) -> io::Result<()> { + match input { + UserInput::LocalImage { path, detail } => { + let image_detail = *detail; + let mode = match image_detail { + Some(ImageDetail::Original) => PromptImageMode::Original, + Some(ImageDetail::Auto | ImageDetail::Low | ImageDetail::High) | None => { + PromptImageMode::ResizeToFit + } + }; + let file_bytes = read_bounded_local_media(path, MAX_PROMPT_IMAGE_INPUT_BYTES, "image")?; + let image = load_for_prompt_bytes(path, file_bytes, mode).map_err(io::Error::other)?; + *input = UserInput::Image { + image_url: image.into_data_url(), + detail: image_detail, + }; + } + UserInput::LocalAudio { path } => { + let mime = audio_mime_for_path(path).ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidData, "unsupported audio format") + })?; + let file_bytes = read_bounded_local_media(path, MAX_PROMPT_AUDIO_INPUT_BYTES, "audio")?; + *input = UserInput::Audio { + audio_url: data_url_from_bytes(mime, &file_bytes), + }; + } + UserInput::Text { .. } + | UserInput::Image { .. } + | UserInput::Audio { .. } + | UserInput::Skill { .. } + | UserInput::Mention { .. } => {} + } + Ok(()) +} + +fn read_bounded_local_media(path: &Path, max_bytes: usize, kind: &str) -> io::Result> { + let file = std::fs::File::open(path)?; + if file.metadata()?.len() > max_bytes as u64 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("{kind} input exceeds {max_bytes} bytes"), + )); + } + + let mut bytes = Vec::new(); + file.take(max_bytes as u64 + 1).read_to_end(&mut bytes)?; + if bytes.len() > max_bytes { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("{kind} input exceeds {max_bytes} bytes"), + )); + } + Ok(bytes) +} + +pub(crate) fn audio_mime_for_path(path: &Path) -> Option<&'static str> { + let extension = path.extension()?.to_str()?; + if extension.eq_ignore_ascii_case("wav") { + Some("audio/wav") + } else if extension.eq_ignore_ascii_case("mp3") { + Some("audio/mpeg") + } else if extension.eq_ignore_ascii_case("m4a") { + Some("audio/mp4") + } else if extension.eq_ignore_ascii_case("webm") { + Some("audio/webm") + } else if extension.eq_ignore_ascii_case("ogg") { + Some("audio/ogg") + } else { + None + } +} + +#[cfg(test)] +#[path = "local_media_tests.rs"] +mod tests; diff --git a/codex-rs/protocol/src/local_media_tests.rs b/codex-rs/protocol/src/local_media_tests.rs new file mode 100644 index 0000000000..3d03c32620 --- /dev/null +++ b/codex-rs/protocol/src/local_media_tests.rs @@ -0,0 +1,128 @@ +use super::*; +use anyhow::Result; +use pretty_assertions::assert_eq; +use tempfile::tempdir; + +const TINY_PNG_BYTES: &[u8] = &[ + 137, 80, 78, 71, 13, 10, 26, 10, 0, 0, 0, 13, 73, 72, 68, 82, 0, 0, 0, 1, 0, 0, 0, 1, 8, 6, 0, + 0, 0, 31, 21, 196, 137, 0, 0, 0, 11, 73, 68, 65, 84, 120, 156, 99, 96, 0, 2, 0, 0, 5, 0, 1, + 122, 94, 171, 63, 0, 0, 0, 0, 73, 69, 78, 68, 174, 66, 96, 130, +]; + +#[test] +fn snapshots_local_image_user_input_with_requested_detail() -> Result<()> { + let temp_dir = tempdir()?; + let image_path = temp_dir.path().join("sample.png"); + std::fs::write(&image_path, TINY_PNG_BYTES)?; + + for detail in [ + None, + Some(ImageDetail::Auto), + Some(ImageDetail::Low), + Some(ImageDetail::High), + Some(ImageDetail::Original), + ] { + let mut input = UserInput::LocalImage { + path: image_path.clone(), + detail, + }; + + snapshot_local_user_input(&mut input)?; + + assert_eq!( + UserInput::Image { + image_url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAC0lEQVR4nGNgAAIAAAUAAXpeqz8AAAAASUVORK5CYII=" + .to_string(), + detail, + }, + input + ); + } + + Ok(()) +} + +#[test] +fn snapshots_local_audio_user_input_with_supported_mime_types() -> Result<()> { + let temp_dir = tempdir()?; + + for (extension, mime) in [ + ("wav", "audio/wav"), + ("MP3", "audio/mpeg"), + ("m4a", "audio/mp4"), + ("webm", "audio/webm"), + ("ogg", "audio/ogg"), + ] { + let audio_path = temp_dir.path().join(format!("sample.{extension}")); + std::fs::write(&audio_path, b"audio")?; + let mut input = UserInput::LocalAudio { path: audio_path }; + + snapshot_local_user_input(&mut input)?; + + assert_eq!( + UserInput::Audio { + audio_url: format!("data:{mime};base64,YXVkaW8="), + }, + input + ); + } + + Ok(()) +} + +#[test] +fn rejects_invalid_local_media_without_changing_user_input() -> Result<()> { + let temp_dir = tempdir()?; + let invalid_image_path = temp_dir.path().join("invalid.png"); + let unsupported_audio_path = temp_dir.path().join("unsupported.flac"); + std::fs::write(&invalid_image_path, b"invalid image")?; + std::fs::write(&unsupported_audio_path, b"audio")?; + + let mut invalid_image = UserInput::LocalImage { + path: invalid_image_path, + detail: Some(ImageDetail::Original), + }; + let original_image = invalid_image.clone(); + let image_error = snapshot_local_user_input(&mut invalid_image) + .expect_err("an invalid local image must be rejected"); + assert_eq!(io::ErrorKind::Other, image_error.kind()); + assert_eq!(original_image, invalid_image); + + let mut unsupported_audio = UserInput::LocalAudio { + path: unsupported_audio_path, + }; + let original_audio = unsupported_audio.clone(); + let audio_error = snapshot_local_user_input(&mut unsupported_audio) + .expect_err("an unsupported local audio format must be rejected"); + assert_eq!(io::ErrorKind::InvalidData, audio_error.kind()); + assert_eq!(original_audio, unsupported_audio); + + Ok(()) +} + +#[test] +fn rejects_oversized_local_media_without_reading_it() -> Result<()> { + let temp_dir = tempdir()?; + + for (name, limit, is_image) in [ + ("oversized.png", MAX_PROMPT_IMAGE_INPUT_BYTES, true), + ("oversized.mp3", MAX_PROMPT_AUDIO_INPUT_BYTES, false), + ] { + let path = temp_dir.path().join(name); + std::fs::File::create(&path)?.set_len(limit as u64 + 1)?; + let mut input = if is_image { + UserInput::LocalImage { path, detail: None } + } else { + UserInput::LocalAudio { path } + }; + let original = input.clone(); + + let error = snapshot_local_user_input(&mut input) + .expect_err("oversized local media must be rejected before it is read"); + + assert_eq!(io::ErrorKind::InvalidData, error.kind()); + assert_eq!(original, input); + } + + Ok(()) +} diff --git a/codex-rs/protocol/src/models.rs b/codex-rs/protocol/src/models.rs index 2c2a6d6b7a..99d6bf7f50 100644 --- a/codex-rs/protocol/src/models.rs +++ b/codex-rs/protocol/src/models.rs @@ -12,6 +12,7 @@ use serde::Serialize; use serde::ser::Serializer; use ts_rs::TS; +use crate::local_media::audio_mime_for_path; use crate::permissions::FileSystemAccessMode; use crate::permissions::FileSystemPath; use crate::permissions::FileSystemSandboxEntry; @@ -30,6 +31,8 @@ use crate::mcp::CallToolResult; mod executed_tool_calls; +pub use crate::local_media::MAX_PROMPT_AUDIO_INPUT_BYTES; +pub use crate::local_media::snapshot_local_user_input; pub use executed_tool_calls::ExecutedToolCall; pub use executed_tool_calls::ExecutedToolCallArguments; pub use executed_tool_calls::ExecutedToolCallTruncation; @@ -1529,23 +1532,6 @@ pub enum LocalImagePreparation { Defer, } -fn audio_mime_for_path(path: &std::path::Path) -> Option<&'static str> { - let extension = path.extension()?.to_str()?; - if extension.eq_ignore_ascii_case("wav") { - Some("audio/wav") - } else if extension.eq_ignore_ascii_case("mp3") { - Some("audio/mpeg") - } else if extension.eq_ignore_ascii_case("m4a") { - Some("audio/mp4") - } else if extension.eq_ignore_ascii_case("webm") { - Some("audio/webm") - } else if extension.eq_ignore_ascii_case("ogg") { - Some("audio/ogg") - } else { - None - } -} - fn unsupported_audio_error_placeholder(path: &std::path::Path) -> ContentItem { ContentItem::InputText { text: format!( diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 242a5657ae..48f0d36567 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -1371,6 +1371,9 @@ pub enum EventMsg { /// Updated long-running goal metadata for the thread. ThreadGoalUpdated(ThreadGoalUpdatedEvent), + /// A durable thread-scoped user-message queue changed. + ThreadQueueChanged(ThreadQueueChangedEvent), + /// Incremental MCP startup progress updates. McpStartupUpdate(McpStartupUpdateEvent), @@ -2332,7 +2335,7 @@ pub struct AgentMessageEvent { pub memory_citation: Option, } -#[derive(Debug, Clone, Default, Deserialize, Serialize, JsonSchema, TS)] +#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, JsonSchema, TS)] pub struct UserMessageEvent { #[serde(default, skip_serializing_if = "Option::is_none")] pub client_id: Option, @@ -4117,6 +4120,13 @@ pub struct ThreadGoalUpdatedEvent { pub goal: ThreadGoal, } +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "protocol/")] +pub struct ThreadQueueChangedEvent { + pub thread_id: ThreadId, +} + /// User's decision in response to an ExecApprovalRequest. #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, Display, JsonSchema, TS)] #[serde(rename_all = "snake_case")] diff --git a/codex-rs/rollout-trace/src/protocol_event.rs b/codex-rs/rollout-trace/src/protocol_event.rs index 22d2974514..e0a03a3737 100644 --- a/codex-rs/rollout-trace/src/protocol_event.rs +++ b/codex-rs/rollout-trace/src/protocol_event.rs @@ -371,6 +371,7 @@ pub(crate) fn tool_runtime_trace_event(event: &EventMsg) -> Option Option<&'static s | EventMsg::AgentReasoningRawContent(_) | EventMsg::AgentReasoningSectionBreak(_) | EventMsg::ThreadGoalUpdated(_) + | EventMsg::ThreadQueueChanged(_) | EventMsg::McpStartupUpdate(_) | EventMsg::McpStartupComplete(_) | EventMsg::McpToolCallBegin(_) diff --git a/codex-rs/rollout/src/policy.rs b/codex-rs/rollout/src/policy.rs index 65a8213f08..22718f6c95 100644 --- a/codex-rs/rollout/src/policy.rs +++ b/codex-rs/rollout/src/policy.rs @@ -120,6 +120,7 @@ pub fn should_persist_event_msg(ev: &EventMsg, history_mode: ThreadHistoryMode) // Transient, non-durable events. EventMsg::Error(_) + | EventMsg::ThreadQueueChanged(_) | EventMsg::GuardianAssessment(_) | EventMsg::ExecCommandEnd(_) | EventMsg::ViewImageToolCall(_) diff --git a/codex-rs/utils/audio/src/lib.rs b/codex-rs/utils/audio/src/lib.rs index ab836b5631..a2adf86f72 100644 --- a/codex-rs/utils/audio/src/lib.rs +++ b/codex-rs/utils/audio/src/lib.rs @@ -4,6 +4,7 @@ use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use codex_protocol::models::ContentItem; use codex_protocol::models::FunctionCallOutputContentItem; +use codex_protocol::models::MAX_PROMPT_AUDIO_INPUT_BYTES; use codex_protocol::models::ResponseItem; use codex_utils_cache::BlockingLruCache; use codex_utils_cache::sha1_digest; @@ -25,10 +26,6 @@ const AUDIO_TOO_LARGE_PLACEHOLDER: &str = const UNSUPPORTED_AUDIO_FORMAT_PLACEHOLDER: &str = "audio content omitted because its format is not supported; use wav, mp3, m4a, webm, or ogg"; -/// Maximum accepted decoded byte length for prompt audio inputs. -/// -/// This matches the Responses API audio input limit. -const MAX_PROMPT_AUDIO_INPUT_BYTES: usize = 50 * 1024 * 1024; const MAX_PROMPT_AUDIO_BASE64_BYTES: usize = MAX_PROMPT_AUDIO_INPUT_BYTES.div_ceil(3) * 4; const AUDIO_TOKEN_ESTIMATE_CACHE_SIZE: usize = 32; const AUDIO_TOKENS_PER_SECOND: f64 = 10.0;