diff --git a/codex-rs/core/src/session/input_queue.rs b/codex-rs/core/src/session/input_queue.rs index b1100b12e8..955491e708 100644 --- a/codex-rs/core/src/session/input_queue.rs +++ b/codex-rs/core/src/session/input_queue.rs @@ -11,6 +11,7 @@ use codex_protocol::turn_input::TurnStartOptions; use codex_protocol::user_input::UserInput; use serde::Deserialize; use serde::Serialize; +use std::collections::BTreeMap; use std::collections::VecDeque; use std::sync::Arc; use tokio::sync::Mutex; @@ -19,8 +20,11 @@ use tokio_util::sync::CancellationToken; use tokio_util::task::AbortOnDropHandle; /// Host capture metadata belonging to one input, including steers within another turn. -#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] pub struct UserInputMetadata { + /// Private extension data carried with this queued input through completion. + #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] + pub extension_data: BTreeMap, #[serde(default, skip_serializing_if = "Option::is_none")] pub acceptance_order: Option, /// Optional host-assigned UUID for this particular Responses user item. @@ -527,6 +531,18 @@ mod tests { use codex_protocol::user_input::UserInput; use pretty_assertions::assert_eq; + #[test] + fn user_input_extension_data_round_trips_without_changing_legacy_input() { + let legacy = serde_json::json!({"UserInput": {"content": [], "client_id": "m"}}); + let input: TurnInput = serde_json::from_value(legacy.clone()).unwrap(); + assert_eq!(serde_json::to_value(input).unwrap(), legacy); + + let mut annotated = legacy; + annotated["UserInput"]["extension_data"] = serde_json::json!({"example": ["private"]}); + let input: TurnInput = serde_json::from_value(annotated.clone()).unwrap(); + assert_eq!(serde_json::to_value(input).unwrap(), annotated); + } + #[test_case::test_case("ResponseItem", TurnInput::ResponseItem)] #[test_case::test_case("FunctionCallOutput", TurnInput::FunctionCallOutput)] fn response_item_serde_preserves_legacy_shape_and_rejects_metadata( diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index a88383e738..133360cd7d 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3,6 +3,7 @@ pub(crate) mod startup; use crate::PromptCacheKey; use crate::WithTurnExtensionData; use crate::runtime_auth::RuntimeAuthSlot; +use std::collections::BTreeMap; use std::collections::HashMap; use std::collections::HashSet; use std::fmt::Debug; @@ -2970,27 +2971,30 @@ impl Session { turn_context: &TurnContext, item: TurnItem, ) { - self.emit_turn_item_completed_with_token_count( + self.emit_turn_item_completed_with_metadata( turn_context, item, /*content_token_count*/ None, + /*user_input_extension_data*/ None, ) .await; } - pub(crate) async fn emit_turn_item_completed_with_token_count( + pub(crate) async fn emit_turn_item_completed_with_metadata( &self, turn_context: &TurnContext, item: TurnItem, content_token_count: Option, + user_input_extension_data: Option<&BTreeMap>, ) { record_turn_ttfm_metric(turn_context, &item).await; for contributor in self.services.extensions.turn_lifecycle_contributors() { contributor - .on_item_completed( + .on_item_completed_with_metadata( &self.services.thread_extension_data, turn_context.extension_data.as_ref(), &item, + user_input_extension_data, ) .await; } @@ -5809,7 +5813,13 @@ impl Session { user_message_item.client_id = client_id; let turn_item = TurnItem::UserMessage(user_message_item); self.emit_turn_item_started(turn_context, &turn_item).await; - self.emit_turn_item_completed(turn_context, turn_item).await; + self.emit_turn_item_completed_with_metadata( + turn_context, + turn_item, + /*content_token_count*/ None, + Some(&metadata.extension_data), + ) + .await; self.ensure_rollout_materialized(persist_context).await; } diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 7261e3378b..c7f1b8a6fc 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -2691,10 +2691,11 @@ async fn emit_agent_message_in_plan_mode( .insert(agent_message_id.clone()); } - sess.emit_turn_item_completed_with_token_count( + sess.emit_turn_item_completed_with_metadata( turn_context, TurnItem::AgentMessage(agent_message), content_token_count, + /*user_input_extension_data*/ None, ) .await; state.started_agent_message_items.remove(&agent_message_id); diff --git a/codex-rs/core/src/session/turn_input.rs b/codex-rs/core/src/session/turn_input.rs index 7c2ac89938..831e719ee5 100644 --- a/codex-rs/core/src/session/turn_input.rs +++ b/codex-rs/core/src/session/turn_input.rs @@ -356,6 +356,7 @@ fn user_input_metadata(request: &TurnInputRequest) -> UserInputMetadata { UserInputMetadata { origin: UserInputOrigin::from_turn_trigger(request.start.turn_trigger.as_deref()), response_item_uuid: request.user_item_uuid, + extension_data: request.user_input_extension_data.clone(), ..Default::default() } } @@ -419,7 +420,7 @@ async fn start_or_steer( /*expected_turn_id*/ None, settings.required_active_final_output_json_schema(), responsesapi_client_metadata.clone(), - input_metadata, + input_metadata.clone(), ) .await { diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 4478f53747..66161a3390 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -415,10 +415,11 @@ pub(crate) async fn handle_output_item_done( } ctx.sess - .emit_turn_item_completed_with_token_count( + .emit_turn_item_completed_with_metadata( &ctx.step_context.turn, finalized_turn_item.turn_item, finalized_turn_item.content_token_count, + /*user_input_extension_data*/ None, ) .await; } diff --git a/codex-rs/core/tests/suite/turn_input_submission.rs b/codex-rs/core/tests/suite/turn_input_submission.rs index 23a6b17c38..e23d3b1660 100644 --- a/codex-rs/core/tests/suite/turn_input_submission.rs +++ b/codex-rs/core/tests/suite/turn_input_submission.rs @@ -12,14 +12,18 @@ use codex_core::config::Constrained; use codex_core::context::ContextualUserFragment; use codex_core::context::InternalContextSource; use codex_core::context::InternalModelContextFragment; +use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionDataInit; +use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::TurnLifecycleContributor; use codex_extension_api::TurnStartAdmission; use codex_features::Feature; use codex_protocol::AgentPath; use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::ModeKind; use codex_protocol::config_types::Settings; +use codex_protocol::items::TurnItem; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::InterAgentCommunication; @@ -47,12 +51,14 @@ use core_test_support::wait_for_event; use core_test_support::wait_for_event_match; use pretty_assertions::assert_eq; use serde_json::Value; +use std::collections::BTreeMap; use std::sync::Arc; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Duration; use test_case::test_case; use tokio::sync::Barrier; +use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio::time::timeout; @@ -1142,7 +1148,32 @@ async fn continue_turn_if_idle_starts_new_turn_with_internal_input() { assert_eq!(mock.requests().len(), 1); } -/// Concurrent submissions must start exactly one turn and steer the other message. +struct CompletedInputRecorder(mpsc::UnboundedSender<(String, String, Option)>); + +impl TurnLifecycleContributor for CompletedInputRecorder { + fn on_item_completed_with_metadata<'a>( + &'a self, + _thread_store: &'a ExtensionData, + turn_store: &'a ExtensionData, + item: &'a TurnItem, + data: Option<&'a BTreeMap>, + ) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + if let TurnItem::UserMessage(message) = item { + self.0 + .send(( + turn_store.level_id().to_owned(), + message.message(), + data.and_then(|data| data.get("test")).cloned(), + )) + .expect("completion receiver"); + } + }) + } +} + +/// Concurrent submissions can reuse a client ID. Both completions must carry their own +/// metadata to the accepting turn, whichever message wins the race to start it. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn turn_input_submission_reports_started_and_steered_for_concurrent_submissions() { let (release_response, response_gate) = oneshot::channel(); @@ -1163,7 +1194,11 @@ async fn turn_input_submission_reports_started_and_steered_for_concurrent_submis }], ]) .await; + let (tx, mut completions) = mpsc::unbounded_channel(); + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.turn_lifecycle_contributor(Arc::new(CompletedInputRecorder(tx))); let test = test_codex() + .with_extensions(Arc::new(extensions.build())) .with_model("gpt-5.4") .build_with_streaming_server(&server) .await @@ -1171,12 +1206,26 @@ async fn turn_input_submission_reports_started_and_steered_for_concurrent_submis let codex = Arc::clone(&test.codex); let barrier = Arc::new(Barrier::new(3)); + let annotated = |text: &str| { + let mut request = TurnInputRequest::new(TurnInput::UserInput { + content: vec![UserInput::Text { + text: text.to_owned(), + text_elements: Vec::new(), + }], + client_id: Some("shared-id".to_owned()), + }); + request.user_input_extension_data.insert( + "test".into(), + Value::from(format!("extension-test-marker:{text}")), + ); + request + }; let first_submission = tokio::spawn({ let codex = Arc::clone(&codex); let barrier = Arc::clone(&barrier); async move { barrier.wait().await; - submit_user_message(codex.as_ref(), "first message").await + codex.start_or_steer_turn(annotated("first message")).await } }); let second_submission = tokio::spawn({ @@ -1184,7 +1233,7 @@ async fn turn_input_submission_reports_started_and_steered_for_concurrent_submis let barrier = Arc::clone(&barrier); async move { barrier.wait().await; - submit_user_message(codex.as_ref(), "second message").await + codex.start_or_steer_turn(annotated("second message")).await } }); barrier.wait().await; @@ -1250,6 +1299,28 @@ async fn turn_input_submission_reports_started_and_steered_for_concurrent_submis assert!(request_bodies[1].to_string().contains("first message")); assert!(request_bodies[1].to_string().contains("second message")); + let mut recorded = [ + completions.try_recv().unwrap(), + completions.try_recv().unwrap(), + ]; + recorded.sort_by(|a, b| a.1.cmp(&b.1)); + for ((turn, message, data), expected) in recorded + .into_iter() + .zip(["first message", "second message"]) + { + assert_eq!(&turn, started_turn_id); + assert_eq!(message, expected); + assert_eq!( + data, + Some(Value::from(format!("extension-test-marker:{expected}"))) + ); + } + assert!( + request_bodies + .iter() + .all(|body| !body.to_string().contains("extension-test-marker:")) + ); + server.shutdown().await; } diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 80fb7cb193..581b5aaab7 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -1,3 +1,4 @@ +use std::collections::BTreeMap; use std::future::Future; use std::pin::Pin; use std::sync::Arc; @@ -255,6 +256,18 @@ pub trait TurnLifecycleContributor: Send + Sync { Box::pin(std::future::ready(())) } + /// Observes the same completion with host-supplied data for a recorded user input. + /// The default forwards to [`Self::on_item_completed`]. + fn on_item_completed_with_metadata<'a>( + &'a self, + thread_store: &'a ExtensionData, + turn_store: &'a ExtensionData, + item: &'a TurnItem, + _user_input_extension_data: Option<&'a BTreeMap>, + ) -> ExtensionFuture<'a, ()> { + self.on_item_completed(thread_store, turn_store, item) + } + /// Called before the host drops the completed turn runtime and turn store. /// The callback completes before the host emits the turn-complete event. fn on_turn_stop<'a>(&'a self, input: TurnStopInput<'a>) -> ExtensionFuture<'a, ()> { diff --git a/codex-rs/protocol/src/turn_input.rs b/codex-rs/protocol/src/turn_input.rs index 453f3d1541..dda1331d75 100644 --- a/codex-rs/protocol/src/turn_input.rs +++ b/codex-rs/protocol/src/turn_input.rs @@ -62,6 +62,9 @@ pub struct TurnInputRequest { /// Optional host-assigned UUID for the Responses message created from [`TurnInput::UserInput`]. /// Ignored for all other inputs, including raw user-role Responses items. pub user_item_uuid: Option, + /// Host-supplied, extension-keyed data for this user input's completion. + /// Travels with the individual input and is never included in model input. + pub user_input_extension_data: BTreeMap, pub thread_settings: ThreadSettingsOverrides, pub start: TurnStartOptions, pub additional_context: BTreeMap, @@ -120,6 +123,7 @@ impl TurnInputRequest { Self { input, user_item_uuid: None, + user_input_extension_data: BTreeMap::new(), thread_settings: ThreadSettingsOverrides::default(), start: TurnStartOptions::default(), additional_context: BTreeMap::new(),