Add durable user-message queue dispatch (#37204)

## What changed

- Add a storage-neutral queue extension for listing, adding, editing, reordering, deleting, and explicitly starting queued user messages.
- Dispatch queued messages in FIFO order when a thread becomes idle after a completed turn, while leaving them queued after interruptions or failures.
- Wait for user input to be persisted before removing its queue item, and report hook rejection, task termination, and persistence failures distinctly.
- Snapshot local image and audio attachments into validated data URLs before storing them, enforce input limits, and emit `ThreadQueueChanged` events after mutations.

## Testing

- Cover queue mutations, automatic and explicit dispatch, interruption and failure behavior, hook rejection, invalid records, input limits, and local attachment snapshotting.

GitOrigin-RevId: 14f39d7c8cbaccb767a6da84b82cced0edfcb4b9
This commit is contained in:
efrazer-oai authored and copyberry committed 2026-08-06 03:11:42 +00:00
1 parent 70b4653232
commit bc8b25ea02
42 files changed
+2114 -140

No files matched your search

+22
View File
@@ -2035,6 +2035,7 @@ dependencies = [
"codex-otel", "codex-otel",
"codex-plugin", "codex-plugin",
"codex-protocol", "codex-protocol",
"codex-queue-extension",
"codex-rmcp-client", "codex-rmcp-client",
"codex-rollout", "codex-rollout",
"codex-sandboxing", "codex-sandboxing",
@@ -3880,6 +3881,27 @@ dependencies = [
"wildmatch", "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]] [[package]]
name = "codex-response-debug-context" name = "codex-response-debug-context"
version = "0.0.0" version = "0.0.0"
+2
View File
@@ -61,6 +61,7 @@ members = [
"ext/items", "ext/items",
"ext/memories", "ext/memories",
"ext/mcp", "ext/mcp",
"ext/queue",
"ext/skills", "ext/skills",
"ext/web-search", "ext/web-search",
"external-agent-migration", "external-agent-migration",
@@ -226,6 +227,7 @@ codex-model-provider = { path = "model-provider" }
codex-process-hardening = { path = "process-hardening" } codex-process-hardening = { path = "process-hardening" }
codex-protocol = { path = "protocol" } codex-protocol = { path = "protocol" }
codex-prompts = { path = "prompts" } codex-prompts = { path = "prompts" }
codex-queue-extension = { path = "ext/queue" }
codex-responses-api-proxy = { path = "responses-api-proxy" } codex-responses-api-proxy = { path = "responses-api-proxy" }
codex-response-debug-context = { path = "response-debug-context" } codex-response-debug-context = { path = "response-debug-context" }
codex-rmcp-client = { path = "rmcp-client" } codex-rmcp-client = { path = "rmcp-client" }
+1
View File
@@ -74,6 +74,7 @@ codex-mcp-extension = { workspace = true }
codex-model-provider = { workspace = true } codex-model-provider = { workspace = true }
codex-models-manager = { workspace = true } codex-models-manager = { workspace = true }
codex-protocol = { workspace = true } codex-protocol = { workspace = true }
codex-queue-extension = { workspace = true }
codex-app-server-protocol = { workspace = true } codex-app-server-protocol = { workspace = true }
codex-app-server-transport = { workspace = true } codex-app-server-transport = { workspace = true }
codex-feedback = { workspace = true } codex-feedback = { workspace = true }
@@ -1186,6 +1186,7 @@ pub(crate) async fn apply_bespoke_event_handling(
)) ))
.await; .await;
} }
EventMsg::ThreadQueueChanged(_) => {}
EventMsg::ThreadSettingsApplied(thread_settings_event) => { EventMsg::ThreadSettingsApplied(thread_settings_event) => {
let thread_settings = let thread_settings =
thread_settings_from_core_snapshot(thread_settings_event.thread_settings); thread_settings_from_core_snapshot(thread_settings_event.thread_settings);
+13 -5
View File
@@ -26,7 +26,7 @@ use codex_protocol::error::CodexErr;
use codex_protocol::protocol::Event; use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::EventMsg;
use codex_rollout::state_db::StateDbHandle; 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::OutgoingMessageSender;
use crate::outgoing_message::ThreadScopedOutgoingMessageSender; use crate::outgoing_message::ThreadScopedOutgoingMessageSender;
@@ -44,8 +44,8 @@ pub(crate) struct ThreadExtensionDependencies {
pub(crate) executor_skill_provider: Arc<dyn codex_skills_extension::SkillProvider>, pub(crate) executor_skill_provider: Arc<dyn codex_skills_extension::SkillProvider>,
pub(crate) git_attribution_base_url: String, pub(crate) git_attribution_base_url: String,
pub(crate) http_client_factory: HttpClientFactory, pub(crate) http_client_factory: HttpClientFactory,
/// Process-scoped persistence backend for extensions that need stored thread history. /// Process-scoped persistence backend for queued user messages.
pub(crate) thread_store: Arc<dyn ThreadStore>, pub(crate) queue_store: Option<Arc<dyn QueueStore>>,
} }
pub(crate) fn thread_extensions<S>( pub(crate) fn thread_extensions<S>(
@@ -66,9 +66,17 @@ where
executor_skill_provider, executor_skill_provider,
git_attribution_base_url, git_attribution_base_url,
http_client_factory, http_client_factory,
thread_store: _thread_store, queue_store,
} = dependencies; } = dependencies;
let mut builder = ExtensionRegistryBuilder::<Config>::with_event_sink(event_sink); let mut builder = ExtensionRegistryBuilder::<Config>::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 { if let Some(state_db) = state_db {
codex_goal_extension::install_with_backend( codex_goal_extension::install_with_backend(
&mut builder, &mut builder,
+1 -1
View File
@@ -339,7 +339,7 @@ enabled = false
executor_skill_provider: Arc::clone(&executor_skill_provider), executor_skill_provider: Arc::clone(&executor_skill_provider),
git_attribution_base_url: good_config.chatgpt_base_url.clone(), git_attribution_base_url: good_config.chatgpt_base_url.clone(),
http_client_factory: good_config.http_client_factory(), http_client_factory: good_config.http_client_factory(),
thread_store: Arc::clone(&thread_store), queue_store: None,
}, },
), ),
Arc::new(CodexHomeUserInstructionsProvider::new( Arc::new(CodexHomeUserInstructionsProvider::new(
+12 -1
View File
@@ -66,6 +66,7 @@ use codex_chatgpt::workspace_settings;
use codex_code_mode::CodeModeSessionProvider; use codex_code_mode::CodeModeSessionProvider;
use codex_core::ThreadManager; use codex_core::ThreadManager;
use codex_core::config::Config; use codex_core::config::Config;
use codex_core::config::ThreadStoreConfig;
use codex_exec_server::EnvironmentManager; use codex_exec_server::EnvironmentManager;
use codex_feedback::CodexFeedback; use codex_feedback::CodexFeedback;
use codex_goal_extension::GoalService; use codex_goal_extension::GoalService;
@@ -77,6 +78,8 @@ use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::W3cTraceContext; use codex_protocol::protocol::W3cTraceContext;
use codex_rollout::StateDbHandle; use codex_rollout::StateDbHandle;
use codex_state::log_db::LogDbLayer; use codex_state::log_db::LogDbLayer;
use codex_thread_store::LocalQueueStore;
use codex_thread_store::QueueStore;
use tokio::sync::Mutex; use tokio::sync::Mutex;
use tokio::sync::Semaphore; use tokio::sync::Semaphore;
use tokio::sync::broadcast; use tokio::sync::broadcast;
@@ -247,6 +250,14 @@ impl MessageProcessor {
// affect per-thread behavior, but they must not move newly started, // affect per-thread behavior, but they must not move newly started,
// resumed, or forked threads to a different persistence backend/root. // 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()); 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<Arc<dyn QueueStore>> = match &config.experimental_thread_store {
ThreadStoreConfig::Local => state_db.as_ref().map(|state_db| {
Arc::new(LocalQueueStore::new(Arc::clone(state_db))) as Arc<dyn QueueStore>
}),
ThreadStoreConfig::InMemory { .. } => None,
};
let environment_manager_for_requests = Arc::clone(&environment_manager); let environment_manager_for_requests = Arc::clone(&environment_manager);
let environment_manager_for_extensions = Arc::clone(&environment_manager); let environment_manager_for_extensions = Arc::clone(&environment_manager);
let restriction_product = session_source.restriction_product(); let restriction_product = session_source.restriction_product();
@@ -281,7 +292,7 @@ impl MessageProcessor {
executor_skill_provider: Arc::clone(&executor_skill_provider), executor_skill_provider: Arc::clone(&executor_skill_provider),
git_attribution_base_url: config.chatgpt_base_url.clone(), git_attribution_base_url: config.chatgpt_base_url.clone(),
http_client_factory: config.http_client_factory(), http_client_factory: config.http_client_factory(),
thread_store: Arc::clone(&thread_store), queue_store,
}, },
), ),
Arc::new(CodexHomeUserInstructionsProvider::new( Arc::new(CodexHomeUserInstructionsProvider::new(
@@ -65,18 +65,11 @@ impl ThreadGoalRequestProcessor {
.map(|()| None) .map(|()| None)
} }
pub(crate) async fn emit_resume_goal_snapshot_and_continue( pub(crate) async fn emit_resume_goal_snapshot(&self, thread_id: ThreadId) {
&self,
thread_id: ThreadId,
thread: &CodexThread,
) {
if !self.config.features.enabled(Feature::Goals) { if !self.config.features.enabled(Feature::Goals) {
return; return;
} }
self.emit_thread_goal_snapshot(thread_id).await; 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( pub(crate) async fn pending_resume_goal_state(
@@ -1,5 +1,6 @@
use super::*; use super::*;
use crate::extensions::send_thread_warning; use crate::extensions::send_thread_warning;
use codex_extension_api::ThreadIdleCause;
use codex_protocol::config_types::MultiAgentMode; use codex_protocol::config_types::MultiAgentMode;
pub(super) const THREAD_UNLOADING_DELAY: Duration = Duration::from_secs(30 * 60); 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; .await;
// App-server owns resume response and snapshot ordering, so wait until // App-server owns resume response and snapshot ordering, so wait until
// replay completes before letting extensions react to the idle thread. // replay completes before letting extensions react to the idle thread.
if pending.emit_thread_goal_update { conversation
conversation.emit_thread_idle_lifecycle_if_idle().await; .emit_thread_idle_lifecycle_if_idle(ThreadIdleCause::Completed)
} .await;
} }
pub(super) async fn send_thread_goal_snapshot_notification( pub(super) async fn send_thread_goal_snapshot_notification(
@@ -8,6 +8,7 @@ use codex_app_server_protocol::ThreadSection;
use codex_app_server_protocol::ThreadSectionMoveParams; use codex_app_server_protocol::ThreadSectionMoveParams;
use codex_app_server_protocol::ThreadSectionMoveResponse; use codex_app_server_protocol::ThreadSectionMoveResponse;
use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionDataInit;
use codex_extension_api::ThreadIdleCause;
use codex_protocol::config_types::MultiAgentMode; use codex_protocol::config_types::MultiAgentMode;
use codex_protocol::error::CodexErrorDetails; use codex_protocol::error::CodexErrorDetails;
use codex_protocol::mcp::ClientMcpExtensions; use codex_protocol::mcp::ClientMcpExtensions;
@@ -3438,7 +3439,10 @@ impl ThreadRequestProcessor {
.await; .await;
} }
self.thread_goal_processor 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; .await;
} }
Err(err) => { Err(err) => {
+80 -11
View File
@@ -6,8 +6,11 @@ use crate::session::SessionSettingsUpdate;
use crate::session::SteerInputError; use crate::session::SteerInputError;
use crate::session::TurnInput; use crate::session::TurnInput;
use crate::session::session::Session; use crate::session::session::Session;
use crate::user_message_admission::PendingUserMessageAdmissionState;
use crate::user_message_admission::UserMessageAdmission; use crate::user_message_admission::UserMessageAdmission;
use crate::user_message_admission::UserMessageAdmissionError;
use codex_exec_server::SelectedCapabilityRootsStatus; use codex_exec_server::SelectedCapabilityRootsStatus;
use codex_extension_api::ThreadIdleCause;
use codex_features::Feature; use codex_features::Feature;
use codex_otel::SessionTelemetry; use codex_otel::SessionTelemetry;
use codex_protocol::ThreadId; use codex_protocol::ThreadId;
@@ -100,6 +103,12 @@ pub enum TryStartTurnIfIdleRejectionReason {
/// Another turn or task is active, or the idle reservation was lost before /// Another turn or task is active, or the idle reservation was lost before
/// the automatic turn could start. /// the automatic turn could start.
Busy, 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 /// Rejection returned when an extension asks to start automatic idle work but
@@ -233,6 +242,11 @@ impl CodexThread {
self.session.services.session_telemetry.clone() 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<()> { pub async fn shutdown_and_wait(&self) -> CodexResult<()> {
self.io.shutdown_and_wait().await self.io.shutdown_and_wait().await
} }
@@ -258,8 +272,8 @@ impl CodexThread {
} }
} }
pub async fn emit_thread_idle_lifecycle_if_idle(&self) { pub async fn emit_thread_idle_lifecycle_if_idle(&self, cause: ThreadIdleCause) {
self.session.emit_thread_idle_lifecycle_if_idle().await; self.session.emit_thread_idle_lifecycle_if_idle(cause).await;
} }
#[doc(hidden)] #[doc(hidden)]
@@ -298,28 +312,79 @@ impl CodexThread {
.await .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( pub async fn submit_user_input_and_wait_for_admission(
&self, &self,
op: Op, op: Op,
trace: Option<W3cTraceContext>, trace: Option<W3cTraceContext>,
client_user_message_id: Option<String>, client_user_message_id: Option<String>,
) -> CodexResult<UserMessageAdmission> { ) -> CodexResult<UserMessageAdmission> {
if !matches!(op, Op::UserInput { .. }) { self.submit_user_input_and_wait_for_admission_inner(
return Err(CodexErr::InvalidRequest( op,
"user message admission requires user input".to_string(), trace,
client_user_message_id,
PendingUserMessageAdmissionState::Immediate,
)
.await
.map_err(Into::into)
}
/// Waits for durable admission and preserves its typed failure outcome.
///
/// A client user-message id is required to identify the persisted message.
pub async fn submit_user_input_and_wait_for_persisted_admission(
&self,
op: Op,
trace: Option<W3cTraceContext>,
client_user_message_id: Option<String>,
) -> Result<UserMessageAdmission, UserMessageAdmissionError> {
if client_user_message_id.is_none() {
return Err(UserMessageAdmissionError::Admission(
CodexErr::InvalidRequest(
"persisted user message admission requires a client user message id"
.to_string(),
),
));
}
self.submit_user_input_and_wait_for_admission_inner(
op,
trace,
client_user_message_id,
PendingUserMessageAdmissionState::WaitingForAdmission,
)
.await
}
async fn submit_user_input_and_wait_for_admission_inner(
&self,
op: Op,
trace: Option<W3cTraceContext>,
client_user_message_id: Option<String>,
state: PendingUserMessageAdmissionState,
) -> Result<UserMessageAdmission, UserMessageAdmissionError> {
let Op::UserInput { 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 self.session
.services .services
.agent_control .agent_control
.ensure_execution_capacity_for_op(self.session.thread_id(), &op) .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 submission_id = crate::session::new_submission_id();
let (_pending_admission, admission) = self let (_pending_admission, admission) = self
.session .session
.pending_user_message_admissions .pending_user_message_admissions
.register(submission_id.clone()); .register(submission_id.clone(), client_user_message_id.clone(), state);
self.io self.io
.submit_with_id(Submission { .submit_with_id(Submission {
id: submission_id.clone(), id: submission_id.clone(),
@@ -328,11 +393,15 @@ impl CodexThread {
trace, trace,
parent_turn_id: None, parent_turn_id: None,
}) })
.await?; .await
.map_err(UserMessageAdmissionError::Admission)?;
tokio::select! { tokio::select! {
biased; biased;
result = admission => result.map_err(|_| CodexErr::InternalAgentDied)?, result = admission => result
() = self.io.session_loop_termination.clone() => Err(CodexErr::InternalAgentDied), .unwrap_or(Err(UserMessageAdmissionError::TaskEndedBeforePersistence)),
() = self.io.session_loop_termination.clone() => {
Err(UserMessageAdmissionError::TaskEndedBeforePersistence)
},
} }
} }
+4 -1
View File
@@ -6,6 +6,7 @@ use codex_analytics::GuardianReviewTerminalStatus;
use codex_analytics::GuardianReviewTrackContext; use codex_analytics::GuardianReviewTrackContext;
use codex_analytics::GuardianReviewedAction; use codex_analytics::GuardianReviewedAction;
use codex_core_plugins::PluginCommandAttribution; use codex_core_plugins::PluginCommandAttribution;
use codex_extension_api::ThreadIdleCause;
use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::config_types::ApprovalsReviewer;
use codex_protocol::openai_models::MODEL_SPECIALTY_CYBER; use codex_protocol::openai_models::MODEL_SPECIALTY_CYBER;
use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::AskForApproval;
@@ -281,7 +282,9 @@ async fn record_guardian_denial(session: &Arc<Session>, turn: &Arc<TurnContext>,
if aborted { if aborted {
// Guardian aborts bypass normal task completion, so emit its idle lifecycle here. // Guardian aborts bypass normal task completion, so emit its idle lifecycle here.
// User interrupts deliberately do not take this path. // 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;
} }
}); });
} }
+84 -1
View File
@@ -39,6 +39,7 @@ use codex_thread_store::ReadThreadParams;
use serde_json::Value; use serde_json::Value;
use tracing::instrument; use tracing::instrument;
use crate::codex_thread::TryStartTurnIfIdleRejectionReason;
use crate::context::ContextualUserFragment; use crate::context::ContextualUserFragment;
use crate::context::HookAdditionalContext; use crate::context::HookAdditionalContext;
use crate::event_mapping::parse_turn_item; use crate::event_mapping::parse_turn_item;
@@ -47,6 +48,7 @@ use crate::session::session::Session;
use crate::session::turn_context::TurnContext; use crate::session::turn_context::TurnContext;
use crate::tools::hook_names::HookToolName; use crate::tools::hook_names::HookToolName;
use crate::tools::sandboxing::PermissionRequestPayload; use crate::tools::sandboxing::PermissionRequestPayload;
use crate::user_message_admission::UserMessageAdmissionError;
pub(crate) struct HookRuntimeOutcome { pub(crate) struct HookRuntimeOutcome {
pub should_stop: bool, 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( pub(crate) async fn record_pending_input(
sess: &Arc<Session>, sess: &Arc<Session>,
turn_context: &Arc<TurnContext>, turn_context: &Arc<TurnContext>,
pending_input: TurnInput, pending_input: TurnInput,
additional_contexts: Vec<String>, additional_contexts: Vec<String>,
) { ) -> Result<(), TryStartTurnIfIdleRejectionReason> {
match pending_input { match pending_input {
TurnInput::UserInput { content, client_id } => { 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( sess.record_user_prompt_and_emit_turn_item(
turn_context.as_ref(), turn_context.as_ref(),
content.as_slice(), content.as_slice(),
@@ -593,6 +650,32 @@ pub(crate) async fn record_pending_input(
} }
} }
record_additional_contexts(sess, turn_context, additional_contexts).await; record_additional_contexts(sess, turn_context, additional_contexts).await;
Ok(())
}
pub(crate) async fn reject_pending_input(
sess: &Arc<Session>,
turn_context: &Arc<TurnContext>,
pending_input: &TurnInput,
) {
let TurnInput::UserInput { client_id, .. } = pending_input else {
return;
};
if let Some(client_id) = client_id.as_deref() {
sess.pending_user_message_admissions
.complete_persistence(client_id, Err(UserMessageAdmissionError::RejectedByHook));
}
let input_persisted = {
let mut active = sess.active_turn.lock().await;
active
.as_mut()
.and_then(|turn| turn.task.as_mut())
.filter(|task| Arc::ptr_eq(&task.turn_context, turn_context))
.and_then(|task| task.input_persisted.take())
};
if let Some(sender) = input_persisted {
let _ = sender.send(Err(TryStartTurnIfIdleRejectionReason::RejectedByHook));
}
} }
async fn run_context_injecting_hook<Fut, Outcome>( async fn run_context_injecting_hook<Fut, Outcome>(
+1
View File
@@ -35,6 +35,7 @@ pub use codex_thread::TryStartTurnIfIdleError;
pub use codex_thread::TryStartTurnIfIdleRejectionReason; pub use codex_thread::TryStartTurnIfIdleRejectionReason;
pub use session::turn_context::TurnContext; pub use session::turn_context::TurnContext;
pub use user_message_admission::UserMessageAdmission; pub use user_message_admission::UserMessageAdmission;
pub use user_message_admission::UserMessageAdmissionError;
mod agent; mod agent;
mod agent_communication; mod agent_communication;
mod attestation; mod attestation;
+12
View File
@@ -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 { let task_input = if has_user_input {
self.clear_connector_selection().await; self.clear_connector_selection().await;
for item in &input { for item in &input {
@@ -141,9 +144,18 @@ impl Session {
turn_context, turn_context,
task_input, task_input,
RegularTask::new(), RegularTask::new(),
input_persisted_sender,
MailboxParentProvenance::Ignore, MailboxParentProvenance::Ignore,
) )
.await; .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(()) Ok(())
} }
+3 -1
View File
@@ -4,13 +4,15 @@ use crate::state::TurnState;
use codex_protocol::models::ResponseItem; use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::InterAgentCommunication;
use codex_protocol::user_input::UserInput; use codex_protocol::user_input::UserInput;
use serde::Deserialize;
use serde::Serialize;
use std::collections::VecDeque; use std::collections::VecDeque;
use std::sync::Arc; use std::sync::Arc;
use tokio::sync::Mutex; use tokio::sync::Mutex;
use tokio::sync::watch; use tokio::sync::watch;
/// Input consumed by a regular turn. /// Input consumed by a regular turn.
#[derive(Clone, Debug, PartialEq)] #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum TurnInput { pub enum TurnInput {
UserInput { UserInput {
content: Vec<UserInput>, content: Vec<UserInput>,
+14 -1
View File
@@ -1319,6 +1319,15 @@ impl Session {
InitialHistory::Resumed(resumed_history) => { InitialHistory::Resumed(resumed_history) => {
let turn_context = self.new_default_turn().await; let turn_context = self.new_default_turn().await;
let rollout_items = resumed_history.history; 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 let previous_turn_settings = self
.apply_rollout_reconstruction(&turn_context, &rollout_items) .apply_rollout_reconstruction(&turn_context, &rollout_items)
.await; .await;
@@ -4025,7 +4034,7 @@ impl Session {
.collect::<Vec<_>>(); .collect::<Vec<_>>();
pending_input.push(TurnInput::UserInput { pending_input.push(TurnInput::UserInput {
content: input, content: input,
client_id: client_user_message_id, client_id: client_user_message_id.clone(),
}); });
self.input_queue self.input_queue
.extend_pending_input_and_accept_mailbox_delivery_for_turn_state( .extend_pending_input_and_accept_mailbox_delivery_for_turn_state(
@@ -4033,6 +4042,10 @@ impl Session {
pending_input, pending_input,
) )
.await; .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()) Ok(active_turn_id.clone())
} }
+84 -3
View File
@@ -92,6 +92,8 @@ use crate::tools::handlers::ShellCommandHandler;
use crate::tools::registry::ToolExecutor; use crate::tools::registry::ToolExecutor;
use crate::tools::router::ToolCallSource; use crate::tools::router::ToolCallSource;
use crate::turn_diff_tracker::TurnDiffTracker; 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::ConfigToml;
use codex_config::config_toml::ProjectConfig; use codex_config::config_toml::ProjectConfig;
use codex_config::permissions_toml::FilesystemPermissionToml; use codex_config::permissions_toml::FilesystemPermissionToml;
@@ -9865,6 +9867,83 @@ async fn attach_in_memory_thread_store(
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] #[tokio::test]
async fn hook_transcript_path_does_not_persist_non_local_thread_store() { async fn hook_transcript_path_does_not_persist_non_local_thread_store() {
let (mut session, _) = make_session_and_context().await; let (mut session, _) = make_session_and_context().await;
@@ -10461,7 +10540,9 @@ async fn thread_idle_lifecycle_waits_for_trigger_turn_mailbox_work() {
) )
.await; .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)); 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: "follow up".to_string(),
text_elements: Vec::new(), text_elements: Vec::new(),
}], }],
client_id: None client_id: None,
}, },
TurnInput::InterAgentCommunication(communication), TurnInput::InterAgentCommunication(communication),
], ],
@@ -11018,7 +11099,7 @@ async fn stale_defer_mailbox_delivery_does_not_override_steered_input() {
text: "follow up".to_string(), text: "follow up".to_string(),
text_elements: Vec::new(), text_elements: Vec::new(),
}], }],
client_id: None client_id: None,
}, },
TurnInput::InterAgentCommunication(communication), TurnInput::InterAgentCommunication(communication),
], ],
+13 -3
View File
@@ -20,6 +20,7 @@ use crate::feedback_tags;
use crate::hook_runtime::inspect_pending_input; use crate::hook_runtime::inspect_pending_input;
use crate::hook_runtime::record_additional_contexts; use crate::hook_runtime::record_additional_contexts;
use crate::hook_runtime::record_pending_input; 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_legacy_after_agent_hook;
use crate::hook_runtime::run_pending_session_start_hooks; use crate::hook_runtime::run_pending_session_start_hooks;
use crate::hook_runtime::run_turn_stop_hooks; use crate::hook_runtime::run_turn_stop_hooks;
@@ -583,25 +584,33 @@ pub(crate) async fn run_hooks_and_record_inputs(
) -> bool { ) -> bool {
let mut blocked_input = false; let mut blocked_input = false;
let mut accepted_user_input = false; let mut accepted_user_input = false;
let mut persistence_failed = false;
for input_item in input { for input_item in input {
let hook_outcome = inspect_pending_input(sess, turn_context, input_item).await; let hook_outcome = inspect_pending_input(sess, turn_context, input_item).await;
if hook_outcome.should_stop { if hook_outcome.should_stop {
blocked_input = true; blocked_input = true;
reject_pending_input(sess, turn_context, input_item).await;
record_additional_contexts(sess, turn_context, hook_outcome.additional_contexts).await; record_additional_contexts(sess, turn_context, hook_outcome.additional_contexts).await;
} else { } else {
if matches!(input_item, TurnInput::UserInput { content, .. } if !content.is_empty()) { if matches!(input_item, TurnInput::UserInput { content, .. } if !content.is_empty()) {
accepted_user_input = true; accepted_user_input = true;
} }
record_pending_input( if record_pending_input(
sess, sess,
turn_context, turn_context,
input_item.clone(), input_item.clone(),
hook_outcome.additional_contexts, 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<UserInput> { fn turn_user_input(input: &[TurnInput]) -> Vec<UserInput> {
@@ -1750,6 +1759,7 @@ pub(super) fn realtime_text_for_event(msg: &EventMsg) -> Option<(String, Option<
| EventMsg::EnvironmentConnected(_) | EventMsg::EnvironmentConnected(_)
| EventMsg::EnvironmentDisconnected(_) | EventMsg::EnvironmentDisconnected(_)
| EventMsg::ThreadGoalUpdated(_) | EventMsg::ThreadGoalUpdated(_)
| EventMsg::ThreadQueueChanged(_)
| EventMsg::McpStartupUpdate(_) | EventMsg::McpStartupUpdate(_)
| EventMsg::McpStartupComplete(_) | EventMsg::McpStartupComplete(_)
| EventMsg::McpToolCallBegin(_) | EventMsg::McpToolCallBegin(_)
+3
View File
@@ -18,6 +18,7 @@ use rmcp::model::RequestId;
use tokio::sync::oneshot; use tokio::sync::oneshot;
use crate::agent::control::AgentExecutionGuard; use crate::agent::control::AgentExecutionGuard;
use crate::codex_thread::TryStartTurnIfIdleRejectionReason;
use crate::mcp_tool_call::McpToolApprovalMetadata; use crate::mcp_tool_call::McpToolApprovalMetadata;
use crate::session::TurnInputQueue; use crate::session::TurnInputQueue;
use crate::session::turn_context::TurnContext; use crate::session::turn_context::TurnContext;
@@ -74,6 +75,8 @@ pub(crate) struct RunningTask {
pub(crate) done: Arc<Notify>, pub(crate) done: Arc<Notify>,
pub(crate) kind: TaskKind, pub(crate) kind: TaskKind,
pub(crate) task: Arc<dyn AnySessionTask>, pub(crate) task: Arc<dyn AnySessionTask>,
pub(crate) input_persisted:
Option<oneshot::Sender<Result<(), TryStartTurnIfIdleRejectionReason>>>,
pub(crate) cancellation_token: CancellationToken, pub(crate) cancellation_token: CancellationToken,
pub(crate) handle: AbortOnDropHandle<()>, pub(crate) handle: AbortOnDropHandle<()>,
pub(crate) turn_context: Arc<TurnContext>, pub(crate) turn_context: Arc<TurnContext>,
+3 -1
View File
@@ -1,4 +1,5 @@
use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionData;
use codex_extension_api::ThreadIdleCause;
use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::CodexErrorInfo;
use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TokenUsage;
use codex_protocol::protocol::TurnAbortReason; 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() if self.active_turn.lock().await.is_some()
|| self.input_queue.has_trigger_turn_mailbox_items().await || self.input_queue.has_trigger_turn_mailbox_items().await
{ {
@@ -49,6 +50,7 @@ impl Session {
for contributor in self.services.extensions.thread_lifecycle_contributors() { for contributor in self.services.extensions.thread_lifecycle_contributors() {
contributor contributor
.on_thread_idle(codex_extension_api::ThreadIdleInput { .on_thread_idle(codex_extension_api::ThreadIdleInput {
cause,
session_store: &self.services.session_extension_data, session_store: &self.services.session_extension_data,
thread_store: &self.services.thread_extension_data, thread_store: &self.services.thread_extension_data,
}) })
+53 -32
View File
@@ -8,6 +8,7 @@ use std::sync::Arc;
use std::time::Duration; use std::time::Duration;
use std::time::Instant; use std::time::Instant;
use codex_extension_api::ThreadIdleCause;
use futures::future::BoxFuture; use futures::future::BoxFuture;
use tokio::select; use tokio::select;
use tokio::sync::Notify; use tokio::sync::Notify;
@@ -22,13 +23,12 @@ use tracing::trace_span;
use tracing::warn; use tracing::warn;
use crate::codex_thread::BackgroundTerminalInfo; use crate::codex_thread::BackgroundTerminalInfo;
use crate::codex_thread::TryStartTurnIfIdleRejectionReason;
use crate::config::Config; use crate::config::Config;
use crate::context::ContextualUserFragment; 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::TurnInput;
use crate::session::session::Session; use crate::session::session::Session;
use crate::session::turn::run_hooks_and_record_inputs;
use crate::session::turn_context::TurnContext; use crate::session::turn_context::TurnContext;
use crate::state::ActiveTurn; use crate::state::ActiveTurn;
use crate::state::RunningTask; use crate::state::RunningTask;
@@ -281,8 +281,14 @@ impl Session {
) { ) {
self.abort_all_tasks(TurnAbortReason::Replaced).await; self.abort_all_tasks(TurnAbortReason::Replaced).await;
self.clear_connector_selection().await; self.clear_connector_selection().await;
self.start_task(turn_context, input, task, MailboxParentProvenance::Ignore) self.start_task(
.await; turn_context,
input,
task,
/*input_persisted*/ None,
MailboxParentProvenance::Ignore,
)
.await;
} }
pub(crate) async fn start_task<T: SessionTask>( pub(crate) async fn start_task<T: SessionTask>(
@@ -290,6 +296,9 @@ impl Session {
turn_context: Arc<TurnContext>, turn_context: Arc<TurnContext>,
input: Vec<TurnInput>, input: Vec<TurnInput>,
task: T, task: T,
input_persisted: Option<
tokio::sync::oneshot::Sender<Result<(), TryStartTurnIfIdleRejectionReason>>,
>,
mailbox_parent_provenance: MailboxParentProvenance, mailbox_parent_provenance: MailboxParentProvenance,
) { ) {
let task: Arc<dyn AnySessionTask> = Arc::new(task); let task: Arc<dyn AnySessionTask> = Arc::new(task);
@@ -408,6 +417,7 @@ impl Session {
handle: AbortOnDropHandle::new(handle), handle: AbortOnDropHandle::new(handle),
kind: task_kind, kind: task_kind,
task, task,
input_persisted,
cancellation_token, cancellation_token,
turn_context: Arc::clone(&turn_context), turn_context: Arc::clone(&turn_context),
_agent_execution_guard: agent_execution_guard, _agent_execution_guard: agent_execution_guard,
@@ -471,6 +481,7 @@ impl Session {
turn_context, turn_context,
Vec::new(), Vec::new(),
RegularTask::new(), RegularTask::new(),
/*input_persisted*/ None,
MailboxParentProvenance::Attribute, MailboxParentProvenance::Attribute,
) )
.await; .await;
@@ -580,12 +591,24 @@ impl Session {
let turn_state = { let turn_state = {
let mut active = self.active_turn.lock().await; let mut active = self.active_turn.lock().await;
active.as_mut().and_then(|active_turn| { 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(); 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; return;
}; };
let pending_input = self let pending_input = self
@@ -600,28 +623,10 @@ impl Session {
ts.token_usage_at_turn_start.clone(), ts.token_usage_at_turn_start.clone(),
) )
}; };
if !pending_input.is_empty() { run_hooks_and_record_inputs(self, &turn_context, &pending_input).await;
for pending_input_item in pending_input { task_ended_before_persistence |= self
let hook_outcome = .pending_user_message_admissions
inspect_pending_input(self, &turn_context, &pending_input_item).await; .complete_task_end(&turn_context.sub_id);
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;
}
}
}
// Emit token usage metrics. // Emit token usage metrics.
{ {
// TODO(jif): drop this // TODO(jif): drop this
@@ -764,6 +769,15 @@ impl Session {
turn_id: turn_context.sub_id.clone(), turn_id: turn_context.sub_id.clone(),
profile, 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 { let event = if let Some(reason) = abort_reason {
self.emit_turn_abort_lifecycle(reason.clone(), turn_context.extension_data.as_ref()) self.emit_turn_abort_lifecycle(reason.clone(), turn_context.extension_data.as_ref())
.await; .await;
@@ -812,7 +826,7 @@ impl Session {
} }
}; };
if cleared_active_turn { 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 // Regular items were flushed before this terminal event was appended; buffering
// thread writers may not flush it without another explicit barrier. // thread writers may not flush it without another explicit barrier.
@@ -847,12 +861,19 @@ impl Session {
.await .await
} }
async fn handle_task_abort(self: &Arc<Self>, task: RunningTask, reason: TurnAbortReason) { async fn handle_task_abort(self: &Arc<Self>, mut task: RunningTask, reason: TurnAbortReason) {
let sub_id = task.turn_context.sub_id.clone(); let sub_id = task.turn_context.sub_id.clone();
if task.cancellation_token.is_cancelled() { if task.cancellation_token.is_cancelled() {
return; 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"); trace!(task_kind = ?task.kind, sub_id, "aborting running task");
task.cancellation_token.cancel(); task.cancellation_token.cancel();
task.turn_context task.turn_context
+192 -7
View File
@@ -1,3 +1,4 @@
use codex_protocol::error::CodexErr;
use codex_protocol::error::Result as CodexResult; use codex_protocol::error::Result as CodexResult;
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::Mutex; use std::sync::Mutex;
@@ -12,24 +13,72 @@ pub enum UserMessageAdmission {
Steered { turn_id: String }, Steered { turn_id: String },
} }
/// Explains why a submitted user message did not become durable.
#[derive(Debug)]
pub enum UserMessageAdmissionError {
/// Core rejected the submission before admitting it to a turn.
Admission(CodexErr),
/// A user-prompt hook rejected the submitted message.
RejectedByHook,
/// The owning task finished before the submitted message was persisted.
TaskEndedBeforePersistence,
/// The admitted message could not be durably written to the rollout.
PersistenceFailed(CodexErr),
}
impl From<UserMessageAdmissionError> for CodexErr {
fn from(error: UserMessageAdmissionError) -> Self {
match error {
UserMessageAdmissionError::Admission(error) => error,
UserMessageAdmissionError::RejectedByHook => {
Self::InvalidRequest("user message was rejected by a hook".to_string())
}
UserMessageAdmissionError::TaskEndedBeforePersistence => Self::InternalAgentDied,
UserMessageAdmissionError::PersistenceFailed(error) => error,
}
}
}
#[derive(Default)] #[derive(Default)]
pub(crate) struct PendingUserMessageAdmissions { pub(crate) struct PendingUserMessageAdmissions {
pending: Mutex<HashMap<String, oneshot::Sender<CodexResult<UserMessageAdmission>>>>, pending: Mutex<HashMap<String, PendingUserMessageAdmission>>,
}
struct PendingUserMessageAdmission {
client_id: Option<String>,
sender: oneshot::Sender<Result<UserMessageAdmission, UserMessageAdmissionError>>,
state: PendingUserMessageAdmissionState,
}
pub(crate) enum PendingUserMessageAdmissionState {
Immediate,
WaitingForAdmission,
Admitted(UserMessageAdmission),
Persisted,
} }
impl PendingUserMessageAdmissions { impl PendingUserMessageAdmissions {
pub(crate) fn register( pub(crate) fn register(
&self, &self,
submission_id: String, submission_id: String,
client_id: Option<String>,
state: PendingUserMessageAdmissionState,
) -> ( ) -> (
PendingUserMessageAdmissionGuard<'_>, PendingUserMessageAdmissionGuard<'_>,
oneshot::Receiver<CodexResult<UserMessageAdmission>>, oneshot::Receiver<Result<UserMessageAdmission, UserMessageAdmissionError>>,
) { ) {
let (sender, receiver) = oneshot::channel(); let (sender, receiver) = oneshot::channel();
self.pending self.pending
.lock() .lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) .unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(submission_id.clone(), sender); .insert(
submission_id.clone(),
PendingUserMessageAdmission {
client_id,
sender,
state,
},
);
( (
PendingUserMessageAdmissionGuard { PendingUserMessageAdmissionGuard {
admissions: self, admissions: self,
@@ -44,16 +93,152 @@ impl PendingUserMessageAdmissions {
submission_id: &str, submission_id: &str,
admission: CodexResult<UserMessageAdmission>, admission: CodexResult<UserMessageAdmission>,
) { ) {
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 .pending
.lock() .lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) .unwrap_or_else(std::sync::PoisonError::into_inner)
.remove(submission_id); .iter()
if let Some(sender) = admission_sender { .find_map(|(submission_id, entry)| {
let _ = sender.send(admission); (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::<Vec<_>>();
let had_pending_admissions = !submission_ids.is_empty();
for submission_id in submission_ids {
if let Some(entry) = pending.remove(&submission_id) {
let _ = entry
.sender
.send(Err(UserMessageAdmissionError::TaskEndedBeforePersistence));
}
}
had_pending_admissions
}
fn remove(&self, submission_id: &str) { fn remove(&self, submission_id: &str) {
self.pending self.pending
.lock() .lock()
@@ -73,7 +73,7 @@ async fn user_message_admission_reports_started_and_steered_for_concurrent_submi
async move { async move {
barrier.wait().await; barrier.wait().await;
codex codex
.submit_user_input_and_wait_for_admission( .submit_user_input_and_wait_for_persisted_admission(
user_input("first message"), user_input("first message"),
/*trace*/ None, /*trace*/ None,
Some("client-message-1".to_string()), Some("client-message-1".to_string()),
@@ -87,7 +87,7 @@ async fn user_message_admission_reports_started_and_steered_for_concurrent_submi
async move { async move {
barrier.wait().await; barrier.wait().await;
codex codex
.submit_user_input_and_wait_for_admission( .submit_user_input_and_wait_for_persisted_admission(
user_input("second message"), user_input("second message"),
/*trace*/ None, /*trace*/ None,
Some("client-message-2".to_string()), 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; 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 { let (first_admission, second_admission) = timeout(Duration::from_secs(5), async {
tokio::join!(first_submission, second_submission) tokio::join!(first_submission, second_submission)
}) })
.await .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 let first_admission = first_admission
.expect("first submission task should finish") .expect("first submission task should finish")
.expect("first user message should be admitted"); .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); assert_eq!(started_turn_id, steered_turn_id);
release_response let rollout_path = codex.rollout_path().expect("user-message rollout path");
.send(()) let rollout = tokio::fs::read_to_string(rollout_path)
.expect("response gate should remain open"); .await
.expect("read persisted user-message rollout");
let mut persisted_client_message_ids: Vec<_> = rollout
.lines()
.map(serde_json::from_str::<RolloutLine>)
.collect::<Result<Vec<_>, _>>()
.expect("parse persisted user-message rollout")
.into_iter()
.filter_map(|line| match line.item {
RolloutItem::EventMsg(EventMsg::ItemCompleted(event)) => match event.item {
TurnItem::UserMessage(message) => message.client_id,
_ => None,
},
_ => None,
})
.collect();
persisted_client_message_ids.sort();
assert_eq!(
persisted_client_message_ids,
vec![
"client-message-1".to_string(),
"client-message-2".to_string()
]
);
let mut observed_client_message_ids = Vec::new(); let mut observed_client_message_ids = Vec::new();
loop { loop {
let event = timeout(Duration::from_secs(10), codex.next_event()) 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("first message"));
assert!(request_bodies[1].to_string().contains("second 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::<RolloutLine>)
.collect::<Result<Vec<_>, _>>()
.expect("parse persisted user-message rollout")
.into_iter()
.filter_map(|line| match line.item {
RolloutItem::EventMsg(EventMsg::ItemCompleted(event)) => match event.item {
TurnItem::UserMessage(message) => message.client_id,
_ => None,
},
_ => None,
})
.collect();
persisted_client_message_ids.sort();
assert_eq!(
persisted_client_message_ids,
vec![
"client-message-1".to_string(),
"client-message-2".to_string()
]
);
server.shutdown().await; server.shutdown().await;
} }
@@ -29,6 +29,7 @@ pub use prompt::PromptFragment;
pub use prompt::PromptSlot; pub use prompt::PromptSlot;
pub use skill_invocation::SkillInvocationInput; pub use skill_invocation::SkillInvocationInput;
pub use skill_invocation::SkillInvocationKind; pub use skill_invocation::SkillInvocationKind;
pub use thread_lifecycle::ThreadIdleCause;
pub use thread_lifecycle::ThreadIdleInput; pub use thread_lifecycle::ThreadIdleInput;
pub use thread_lifecycle::ThreadOriginator; pub use thread_lifecycle::ThreadOriginator;
pub use thread_lifecycle::ThreadResumeInput; pub use thread_lifecycle::ThreadResumeInput;
@@ -42,8 +42,21 @@ pub struct ThreadResumeInput<'a> {
pub thread_store: &'a ExtensionData, 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. /// Input supplied when the host has no immediately pending thread work.
pub struct ThreadIdleInput<'a> { pub struct ThreadIdleInput<'a> {
/// Why the thread became idle.
pub cause: ThreadIdleCause,
/// Store scoped to the host session runtime. /// Store scoped to the host session runtime.
pub session_store: &'a ExtensionData, pub session_store: &'a ExtensionData,
/// Store scoped to this thread runtime. /// Store scoped to this thread runtime.
+1
View File
@@ -47,6 +47,7 @@ pub use contributors::RenderedWorldStateFragment;
pub use contributors::SkillInvocationContributor; pub use contributors::SkillInvocationContributor;
pub use contributors::SkillInvocationInput; pub use contributors::SkillInvocationInput;
pub use contributors::SkillInvocationKind; pub use contributors::SkillInvocationKind;
pub use contributors::ThreadIdleCause;
pub use contributors::ThreadIdleInput; pub use contributors::ThreadIdleInput;
pub use contributors::ThreadLifecycleContributor; pub use contributors::ThreadLifecycleContributor;
pub use contributors::ThreadOriginator; pub use contributors::ThreadOriginator;
+6
View File
@@ -0,0 +1,6 @@
load("//:defs.bzl", "codex_rust_crate")
codex_rust_crate(
name = "queue",
crate_name = "codex_queue_extension",
)
+34
View File
@@ -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"] }
+19
View File
@@ -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<C>(registry: &mut ExtensionRegistryBuilder<C>, service: Arc<QueuedItemService>)
where
C: Send + Sync + 'static,
{
registry.thread_lifecycle_contributor(service);
}
+394
View File
@@ -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<dyn QueueStore>,
thread_manager: Weak<ThreadManager>,
event_sink: Arc<dyn ExtensionEventSink>,
dispatch_locks: Arc<StdMutex<HashMap<ThreadId, Weak<Mutex<()>>>>>,
}
impl QueuedItemService {
pub fn new(
queue: Arc<dyn QueueStore>,
thread_manager: Weak<ThreadManager>,
event_sink: Arc<dyn ExtensionEventSink>,
) -> Self {
Self {
queue,
thread_manager,
event_sink,
dispatch_locks: Arc::new(StdMutex::new(HashMap::new())),
}
}
fn dispatch_lock(&self, thread_id: ThreadId) -> Arc<Mutex<()>> {
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<QueuedItem, QueueServiceError> {
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<Vec<QueuedItem>, 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<Option<QueuedItem>, 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<bool, QueueServiceError> {
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<bool, QueueServiceError> {
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<String>,
) -> 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<W3cTraceContext>,
) -> Result<UserMessageAdmission, QueueServiceError> {
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::<TurnInput>(&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<TurnInput, QueueServiceError> {
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::<TurnInput, std::io::Error>(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<C> ThreadLifecycleContributor<C> 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<QueuedItem, QueueServiceError> {
Ok(QueuedItem {
id: record.id,
input: serde_json::from_str::<TurnInput>(&record.payload)?,
})
}
+756
View File
@@ -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<Vec<Event>>,
}
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<Arc<QueuedItemService>>);
impl ThreadLifecycleContributor<codex_core::config::Config> for InstalledQueue {
fn on_thread_idle<'a>(&'a self, input: ThreadIdleInput<'a>) -> ExtensionFuture<'a, ()> {
match self.0.get() {
Some(service) => <QueuedItemService as ThreadLifecycleContributor<
codex_core::config::Config,
>>::on_thread_idle(service.as_ref(), input),
None => Box::pin(async {}),
}
}
}
fn registered_queue_extensions() -> (
Arc<InstalledQueue>,
Arc<ExtensionRegistry<codex_core::config::Config>>,
) {
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<Arc<QueuedItemService>> {
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<dyn QueueStore>, 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<dyn QueueStore> = Arc::new(LocalQueueStore::new(runtime));
Ok((queue, home))
}
fn loaded_thread_queue(test: &TestCodex) -> anyhow::Result<Arc<dyn QueueStore>> {
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());
<QueuedItemService as ThreadLifecycleContributor<()>>::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<dyn ExtensionEventSink> = 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::<anyhow::Result<Vec<_>>>()?;
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::<anyhow::Result<Vec<_>>>()?;
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::<Vec<_>>()
);
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::<Vec<_>>()
);
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(())
}
+2 -2
View File
@@ -306,8 +306,8 @@ async fn run_codex_tool_session_inner(
EventMsg::SessionConfigured(_) => { EventMsg::SessionConfigured(_) => {
tracing::error!("unexpected SessionConfigured event"); tracing::error!("unexpected SessionConfigured event");
} }
EventMsg::ThreadGoalUpdated(_) => { EventMsg::ThreadGoalUpdated(_) | EventMsg::ThreadQueueChanged(_) => {
// Ignore thread goal metadata updates in MCP tool runner. // Ignore thread-scoped metadata updates in MCP tool runner.
} }
EventMsg::McpStartupUpdate(_) | EventMsg::McpStartupComplete(_) => { EventMsg::McpStartupUpdate(_) | EventMsg::McpStartupComplete(_) => {
// Ignored in MCP tool runner. // Ignored in MCP tool runner.
+1
View File
@@ -20,6 +20,7 @@ pub mod error;
pub mod exec_output; pub mod exec_output;
pub mod items; pub mod items;
mod legacy_events; mod legacy_events;
pub mod local_media;
pub mod mcp; pub mod mcp;
pub mod mcp_approval_meta; pub mod mcp_approval_meta;
pub mod memory_citation; pub mod memory_citation;
+93
View File
@@ -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<Vec<u8>> {
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;
+128
View File
@@ -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(())
}
+3 -17
View File
@@ -12,6 +12,7 @@ use serde::Serialize;
use serde::ser::Serializer; use serde::ser::Serializer;
use ts_rs::TS; use ts_rs::TS;
use crate::local_media::audio_mime_for_path;
use crate::permissions::FileSystemAccessMode; use crate::permissions::FileSystemAccessMode;
use crate::permissions::FileSystemPath; use crate::permissions::FileSystemPath;
use crate::permissions::FileSystemSandboxEntry; use crate::permissions::FileSystemSandboxEntry;
@@ -30,6 +31,8 @@ use crate::mcp::CallToolResult;
mod executed_tool_calls; 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::ExecutedToolCall;
pub use executed_tool_calls::ExecutedToolCallArguments; pub use executed_tool_calls::ExecutedToolCallArguments;
pub use executed_tool_calls::ExecutedToolCallTruncation; pub use executed_tool_calls::ExecutedToolCallTruncation;
@@ -1529,23 +1532,6 @@ pub enum LocalImagePreparation {
Defer, 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 { fn unsupported_audio_error_placeholder(path: &std::path::Path) -> ContentItem {
ContentItem::InputText { ContentItem::InputText {
text: format!( text: format!(
+11 -1
View File
@@ -1371,6 +1371,9 @@ pub enum EventMsg {
/// Updated long-running goal metadata for the thread. /// Updated long-running goal metadata for the thread.
ThreadGoalUpdated(ThreadGoalUpdatedEvent), ThreadGoalUpdated(ThreadGoalUpdatedEvent),
/// A durable thread-scoped user-message queue changed.
ThreadQueueChanged(ThreadQueueChangedEvent),
/// Incremental MCP startup progress updates. /// Incremental MCP startup progress updates.
McpStartupUpdate(McpStartupUpdateEvent), McpStartupUpdate(McpStartupUpdateEvent),
@@ -2332,7 +2335,7 @@ pub struct AgentMessageEvent {
pub memory_citation: Option<MemoryCitation>, pub memory_citation: Option<MemoryCitation>,
} }
#[derive(Debug, Clone, Default, Deserialize, Serialize, JsonSchema, TS)] #[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, JsonSchema, TS)]
pub struct UserMessageEvent { pub struct UserMessageEvent {
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub client_id: Option<String>, pub client_id: Option<String>,
@@ -4117,6 +4120,13 @@ pub struct ThreadGoalUpdatedEvent {
pub goal: ThreadGoal, 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. /// User's decision in response to an ExecApprovalRequest.
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, Display, JsonSchema, TS)] #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, Display, JsonSchema, TS)]
#[serde(rename_all = "snake_case")] #[serde(rename_all = "snake_case")]
@@ -371,6 +371,7 @@ pub(crate) fn tool_runtime_trace_event(event: &EventMsg) -> Option<ToolRuntimeTr
| EventMsg::ContextCompacted(_) | EventMsg::ContextCompacted(_)
| EventMsg::ThreadRolledBack(_) | EventMsg::ThreadRolledBack(_)
| EventMsg::ThreadGoalUpdated(_) | EventMsg::ThreadGoalUpdated(_)
| EventMsg::ThreadQueueChanged(_)
| EventMsg::TurnStarted(_) | EventMsg::TurnStarted(_)
| EventMsg::ThreadSettingsApplied(_) | EventMsg::ThreadSettingsApplied(_)
| EventMsg::TurnComplete(_) | EventMsg::TurnComplete(_)
@@ -457,6 +458,7 @@ pub(crate) fn wrapped_protocol_event_type(event: &EventMsg) -> Option<&'static s
| EventMsg::AgentReasoningRawContent(_) | EventMsg::AgentReasoningRawContent(_)
| EventMsg::AgentReasoningSectionBreak(_) | EventMsg::AgentReasoningSectionBreak(_)
| EventMsg::ThreadGoalUpdated(_) | EventMsg::ThreadGoalUpdated(_)
| EventMsg::ThreadQueueChanged(_)
| EventMsg::McpStartupUpdate(_) | EventMsg::McpStartupUpdate(_)
| EventMsg::McpStartupComplete(_) | EventMsg::McpStartupComplete(_)
| EventMsg::McpToolCallBegin(_) | EventMsg::McpToolCallBegin(_)
+1
View File
@@ -120,6 +120,7 @@ pub fn should_persist_event_msg(ev: &EventMsg, history_mode: ThreadHistoryMode)
// Transient, non-durable events. // Transient, non-durable events.
EventMsg::Error(_) EventMsg::Error(_)
| EventMsg::ThreadQueueChanged(_)
| EventMsg::GuardianAssessment(_) | EventMsg::GuardianAssessment(_)
| EventMsg::ExecCommandEnd(_) | EventMsg::ExecCommandEnd(_)
| EventMsg::ViewImageToolCall(_) | EventMsg::ViewImageToolCall(_)
+1 -4
View File
@@ -4,6 +4,7 @@ use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use codex_protocol::models::ContentItem; use codex_protocol::models::ContentItem;
use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputContentItem;
use codex_protocol::models::MAX_PROMPT_AUDIO_INPUT_BYTES;
use codex_protocol::models::ResponseItem; use codex_protocol::models::ResponseItem;
use codex_utils_cache::BlockingLruCache; use codex_utils_cache::BlockingLruCache;
use codex_utils_cache::sha1_digest; use codex_utils_cache::sha1_digest;
@@ -25,10 +26,6 @@ const AUDIO_TOO_LARGE_PLACEHOLDER: &str =
const UNSUPPORTED_AUDIO_FORMAT_PLACEHOLDER: &str = const UNSUPPORTED_AUDIO_FORMAT_PLACEHOLDER: &str =
"audio content omitted because its format is not supported; use wav, mp3, m4a, webm, or ogg"; "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 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_TOKEN_ESTIMATE_CACHE_SIZE: usize = 32;
const AUDIO_TOKENS_PER_SECOND: f64 = 10.0; const AUDIO_TOKENS_PER_SECOND: f64 = 10.0;