diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index d2240b4ca0..e84b8db84c 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -1083,28 +1083,37 @@ async fn run_sampling_request( .sampling_input_contributors() .to_vec(); loop { - let mut prompt_input = if let Some(input) = initial_input.take() { + let prompt_input = if let Some(mut input) = initial_input.take() { + let mut contributed_items = Vec::new(); + for contributor in &sampling_input_contributors { + contributed_items.extend( + contributor + .contribute(SamplingInputContext { + turn_id: &turn_context.sub_id, + session_store: &sess.services.session_extension_data, + thread_store: &sess.services.thread_extension_data, + turn_store: turn_store.as_ref(), + }) + .or_cancel(&cancellation_token) + .await? + .map_err(|err| { + CodexErr::Fatal(format!("sampling input contributor failed: {err}")) + })? + .into_iter() + .map(ContextualUserFragment::into_boxed_response_item), + ); + } + if !contributed_items.is_empty() { + sess.record_conversation_items(&turn_context, &contributed_items) + .await; + input.extend(contributed_items); + } input } else { sess.clone_history() .await .for_prompt(&turn_context.model_info.input_modalities) }; - for contributor in &sampling_input_contributors { - contributor - .contribute(SamplingInputContext { - turn_id: &turn_context.sub_id, - session_store: &sess.services.session_extension_data, - thread_store: &sess.services.thread_extension_data, - turn_store: turn_store.as_ref(), - request_input: &mut prompt_input, - }) - .or_cancel(&cancellation_token) - .await? - .map_err(|err| { - CodexErr::Fatal(format!("sampling input contributor failed: {err}")) - })?; - } let prompt = build_prompt( prompt_input, router.as_ref(), diff --git a/codex-rs/core/tests/suite/extension_sampling_input.rs b/codex-rs/core/tests/suite/extension_sampling_input.rs index a73835ab97..660edbe358 100644 --- a/codex-rs/core/tests/suite/extension_sampling_input.rs +++ b/codex-rs/core/tests/suite/extension_sampling_input.rs @@ -4,12 +4,11 @@ use std::sync::atomic::Ordering; use anyhow::Result; use codex_core::config::Config; +use codex_extension_api::ContextualUserFragment; use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::SamplingInputContext; use codex_extension_api::SamplingInputContributor; -use codex_protocol::models::ContentItem; -use codex_protocol::models::ResponseItem; use core_test_support::responses; use core_test_support::skip_if_no_network; use core_test_support::test_codex::test_codex; @@ -23,38 +22,42 @@ struct TimestampLikeContributor { calls: AtomicUsize, } +struct ReminderFragment(usize); + +impl ContextualUserFragment for ReminderFragment { + fn role(&self) -> &'static str { + "developer" + } + + fn markers(&self) -> (&'static str, &'static str) { + Self::type_markers() + } + + fn type_markers() -> (&'static str, &'static str) { + ("", "") + } + + fn body(&self) -> String { + format!("{MARKER_PREFIX}{}]", self.0) + } +} + impl SamplingInputContributor for TimestampLikeContributor { fn contribute<'a>( &'a self, - input: SamplingInputContext<'a>, - ) -> ExtensionFuture<'a, Result<(), String>> { + _input: SamplingInputContext<'a>, + ) -> ExtensionFuture<'a, Result>, String>> { Box::pin(async move { let attempt = self.calls.fetch_add(1, Ordering::Relaxed) + 1; - let Some(content) = input.request_input.iter_mut().rev().find_map(|item| { - let ResponseItem::Message { role, content, .. } = item else { - return None; - }; - (role == "user").then_some(content) - }) else { - return Err("sampling request has no user message".to_string()); - }; - let Some(text) = content.iter_mut().find_map(|item| { - let ContentItem::InputText { text } = item else { - return None; - }; - Some(text) - }) else { - return Err("user message has no input text".to_string()); - }; - text.push_str(&format!("\n{MARKER_PREFIX}{attempt}]")); - Ok(()) + let reminder: Box = + Box::new(ReminderFragment(attempt)); + Ok(vec![reminder]) }) } } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn sampling_input_contributor_runs_for_each_request_without_rewriting_history() -> Result<()> -{ +async fn sampling_input_contributor_appends_items_to_history_before_each_request() -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; @@ -93,21 +96,21 @@ async fn sampling_input_contributor_runs_for_each_request_without_rewriting_hist let requests = response_mock.requests(); assert_eq!(requests.len(), 2); - let user_prompts = requests + let reminder_messages = requests .iter() .map(|request| { request - .message_input_texts("user") + .message_input_texts("developer") .into_iter() - .find(|text| text.starts_with(USER_PROMPT)) - .expect("request should contain the submitted user prompt") + .filter(|text| text.starts_with(MARKER_PREFIX)) + .collect::>() }) .collect::>(); assert_eq!( - user_prompts, + reminder_messages, vec![ - format!("{USER_PROMPT}\n{MARKER_PREFIX}1]"), - format!("{USER_PROMPT}\n{MARKER_PREFIX}2]"), + vec![format!("{MARKER_PREFIX}1]")], + vec![format!("{MARKER_PREFIX}1]"), format!("{MARKER_PREFIX}2]"),], ] ); assert_eq!(contributor.calls.load(Ordering::Relaxed), 2); diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 18d18d9cd4..bba6530914 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -172,16 +172,16 @@ pub trait TurnInputContributor: Send + Sync { ) -> ExtensionFuture<'a, Vec>>; } -/// Extension contribution that can update request-local model input immediately -/// before each sampling attempt. +/// Extension contribution that can append model input immediately before a +/// logical sampling request. /// -/// Implementations should preserve the ordering and provenance of existing items. -/// Returning an error prevents the sampling request from being sent. +/// Returned items are appended to canonical conversation history and included in +/// the outbound request. Returning an error prevents the request from being sent. pub trait SamplingInputContributor: Send + Sync { fn contribute<'a>( &'a self, input: SamplingInputContext<'a>, - ) -> ExtensionFuture<'a, Result<(), String>>; + ) -> ExtensionFuture<'a, Result>, String>>; } /// Contributor for host-owned configuration changes. diff --git a/codex-rs/ext/extension-api/src/contributors/sampling_input.rs b/codex-rs/ext/extension-api/src/contributors/sampling_input.rs index 06212da759..6e344873e5 100644 --- a/codex-rs/ext/extension-api/src/contributors/sampling_input.rs +++ b/codex-rs/ext/extension-api/src/contributors/sampling_input.rs @@ -1,12 +1,6 @@ -use codex_protocol::models::ResponseItem; - use crate::ExtensionData; -/// Input supplied immediately before the host builds one model sampling request. -/// -/// `request_input` is a request-local clone of the conversation history. Mutations -/// affect only the current outbound request and are not persisted to canonical -/// history automatically. +/// Input supplied immediately before the host sends one logical sampling request. pub struct SamplingInputContext<'a> { /// Stable host-owned turn identifier. pub turn_id: &'a str, @@ -16,6 +10,4 @@ pub struct SamplingInputContext<'a> { pub thread_store: &'a ExtensionData, /// Store scoped to this turn runtime. pub turn_store: &'a ExtensionData, - /// Model input for the current sampling request. - pub request_input: &'a mut Vec, } diff --git a/codex-rs/ext/extension-api/src/registry.rs b/codex-rs/ext/extension-api/src/registry.rs index fc5f47d7a0..d5b20a28c1 100644 --- a/codex-rs/ext/extension-api/src/registry.rs +++ b/codex-rs/ext/extension-api/src/registry.rs @@ -117,7 +117,7 @@ impl ExtensionRegistryBuilder { self.turn_input_contributors.push(contributor); } - /// Registers one request-local sampling-input contributor. + /// Registers one pre-sampling input contributor. pub fn sampling_input_contributor(&mut self, contributor: Arc) { self.sampling_input_contributors.push(contributor); } @@ -235,7 +235,7 @@ impl ExtensionRegistry { &self.turn_input_contributors } - /// Returns the registered request-local sampling-input contributors. + /// Returns the registered pre-sampling input contributors. pub fn sampling_input_contributors(&self) -> &[Arc] { &self.sampling_input_contributors } diff --git a/codex-rs/ext/extension-api/tests/registry.rs b/codex-rs/ext/extension-api/tests/registry.rs index 7c04baa8f1..77c73fdc86 100644 --- a/codex-rs/ext/extension-api/tests/registry.rs +++ b/codex-rs/ext/extension-api/tests/registry.rs @@ -73,8 +73,8 @@ impl SamplingInputContributor for AllContributors { fn contribute<'a>( &'a self, _input: SamplingInputContext<'a>, - ) -> ExtensionFuture<'a, Result<(), String>> { - Box::pin(std::future::ready(Ok(()))) + ) -> ExtensionFuture<'a, Result>, String>> { + Box::pin(std::future::ready(Ok(Vec::new()))) } }