[codex] persist sampling input contributions

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(),

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<Vec<Box<dyn ContextualUserFragment + Send>>, 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<dyn ContextualUserFragment + Send> =
+                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::<Vec<_>>()
         })
         .collect::<Vec<_>>();
     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);

codex-rs/ext/extension-api/src/contributors.rs

@@ -172,16 +172,16 @@ pub trait TurnInputContributor: Send + Sync {
     ) -> ExtensionFuture<'a, Vec<Box<dyn ContextualUserFragment + Send>>>;
 }
 
-/// 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<Vec<Box<dyn ContextualUserFragment + Send>>, String>>;
 }
 
 /// Contributor for host-owned configuration changes.

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<ResponseItem>,
 }

codex-rs/ext/extension-api/src/registry.rs

@@ -117,7 +117,7 @@ impl<C: Sync> ExtensionRegistryBuilder<C> {
         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<dyn SamplingInputContributor>) {
         self.sampling_input_contributors.push(contributor);
     }
@@ -235,7 +235,7 @@ impl<C: Sync> ExtensionRegistry<C> {
         &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<dyn SamplingInputContributor>] {
         &self.sampling_input_contributors
     }

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<Vec<Box<dyn ContextualUserFragment + Send>>, String>> {
+        Box::pin(std::future::ready(Ok(Vec::new())))
     }
 }