[codex] add sampling input extension hook

codex-rs/core/src/session/turn.rs

@@ -75,6 +75,7 @@ use codex_analytics::build_track_events_context;
 use codex_async_utils::OrCancelExt;
 use codex_core_plugins::RecommendedPluginCandidatesInput;
 use codex_core_skills::injection::InjectedHostSkillPrompts;
+use codex_extension_api::SamplingInputContext;
 use codex_extension_api::TurnInputContext;
 use codex_extension_api::TurnInputEnvironment;
 use codex_features::Feature;
@@ -1076,14 +1077,34 @@ async fn run_sampling_request(
     let mut retries = 0;
     let mut initial_input = Some(input);
     let mut original_input = None;
+    let sampling_input_contributors = sess
+        .services
+        .extensions
+        .sampling_input_contributors()
+        .to_vec();
     loop {
-        let prompt_input = if let Some(input) = initial_input.take() {
+        let mut prompt_input = if let Some(input) = initial_input.take() {
             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

@@ -0,0 +1,116 @@
+use std::sync::Arc;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+
+use anyhow::Result;
+use codex_core::config::Config;
+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;
+use pretty_assertions::assert_eq;
+use serde_json::json;
+
+const USER_PROMPT: &str = "exercise sampling input contributors";
+const MARKER_PREFIX: &str = "[sampling-attempt=";
+
+struct TimestampLikeContributor {
+    calls: AtomicUsize,
+}
+
+impl SamplingInputContributor for TimestampLikeContributor {
+    fn contribute<'a>(
+        &'a self,
+        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(())
+        })
+    }
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn sampling_input_contributor_runs_for_each_request_without_rewriting_history() -> Result<()>
+{
+    skip_if_no_network!(Ok(()));
+
+    let server = responses::start_mock_server().await;
+    let plan_args = json!({
+        "plan": [{"step": "exercise follow-up sampling", "status": "in_progress"}],
+    })
+    .to_string();
+    let response_mock = responses::mount_sse_sequence(
+        &server,
+        vec![
+            responses::sse(vec![
+                responses::ev_response_created("resp-1"),
+                responses::ev_function_call("call-1", "update_plan", &plan_args),
+                responses::ev_completed("resp-1"),
+            ]),
+            responses::sse(vec![
+                responses::ev_response_created("resp-2"),
+                responses::ev_assistant_message("msg-1", "done"),
+                responses::ev_completed("resp-2"),
+            ]),
+        ],
+    )
+    .await;
+
+    let contributor = Arc::new(TimestampLikeContributor {
+        calls: AtomicUsize::new(0),
+    });
+    let mut extensions = ExtensionRegistryBuilder::<Config>::new();
+    extensions.sampling_input_contributor(contributor.clone());
+    let test = test_codex()
+        .with_extensions(Arc::new(extensions.build()))
+        .build(&server)
+        .await?;
+
+    test.submit_turn(USER_PROMPT).await?;
+
+    let requests = response_mock.requests();
+    assert_eq!(requests.len(), 2);
+    let user_prompts = requests
+        .iter()
+        .map(|request| {
+            request
+                .message_input_texts("user")
+                .into_iter()
+                .find(|text| text.starts_with(USER_PROMPT))
+                .expect("request should contain the submitted user prompt")
+        })
+        .collect::<Vec<_>>();
+    assert_eq!(
+        user_prompts,
+        vec![
+            format!("{USER_PROMPT}\n{MARKER_PREFIX}1]"),
+            format!("{USER_PROMPT}\n{MARKER_PREFIX}2]"),
+        ]
+    );
+    assert_eq!(contributor.calls.load(Ordering::Relaxed), 2);
+
+    Ok(())
+}

codex-rs/core/tests/suite/mod.rs

@@ -51,6 +51,7 @@ mod compact_resume_fork;
 mod deprecation_notice;
 mod exec;
 mod exec_policy;
+mod extension_sampling_input;
 #[cfg(not(target_os = "windows"))]
 mod extension_sandbox;
 mod fork_thread;

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

@@ -13,6 +13,7 @@ use crate::ExtensionData;
 
 mod mcp;
 mod prompt;
+mod sampling_input;
 mod thread_lifecycle;
 mod tool_lifecycle;
 mod turn_input;
@@ -22,6 +23,7 @@ pub use mcp::McpServerContribution;
 pub use mcp::McpServerContributionContext;
 pub use prompt::PromptFragment;
 pub use prompt::PromptSlot;
+pub use sampling_input::SamplingInputContext;
 pub use thread_lifecycle::ThreadIdleInput;
 pub use thread_lifecycle::ThreadResumeInput;
 pub use thread_lifecycle::ThreadStartInput;
@@ -170,6 +172,18 @@ 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.
+///
+/// Implementations should preserve the ordering and provenance of existing items.
+/// Returning an error prevents the sampling request from being sent.
+pub trait SamplingInputContributor: Send + Sync {
+    fn contribute<'a>(
+        &'a self,
+        input: SamplingInputContext<'a>,
+    ) -> ExtensionFuture<'a, Result<(), String>>;
+}
+
 /// Contributor for host-owned configuration changes.
 ///
 /// Implementations should treat the supplied values as immutable before/after

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

@@ -0,0 +1,21 @@
+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.
+pub struct SamplingInputContext<'a> {
+    /// Stable host-owned turn identifier.
+    pub turn_id: &'a str,
+    /// Store scoped to the host session runtime.
+    pub session_store: &'a ExtensionData,
+    /// Store scoped to this thread runtime.
+    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/lib.rs

@@ -40,6 +40,8 @@ pub use contributors::McpServerContributionContext;
 pub use contributors::McpServerContributor;
 pub use contributors::PromptFragment;
 pub use contributors::PromptSlot;
+pub use contributors::SamplingInputContext;
+pub use contributors::SamplingInputContributor;
 pub use contributors::ThreadIdleInput;
 pub use contributors::ThreadLifecycleContributor;
 pub use contributors::ThreadResumeInput;

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

@@ -9,6 +9,7 @@ use crate::ExtensionData;
 use crate::ExtensionEventSink;
 use crate::McpServerContributor;
 use crate::NoopExtensionEventSink;
+use crate::SamplingInputContributor;
 use crate::ThreadLifecycleContributor;
 use crate::TokenUsageContributor;
 use crate::ToolContributor;
@@ -27,6 +28,7 @@ pub struct ExtensionRegistryBuilder<C: Sync> {
     context_contributors: Vec<Arc<dyn ContextContributor>>,
     mcp_server_contributors: Vec<Arc<dyn McpServerContributor<C>>>,
     turn_input_contributors: Vec<Arc<dyn TurnInputContributor>>,
+    sampling_input_contributors: Vec<Arc<dyn SamplingInputContributor>>,
     tool_contributors: Vec<Arc<dyn ToolContributor>>,
     tool_lifecycle_contributors: Vec<Arc<dyn ToolLifecycleContributor>>,
     turn_item_contributors: Vec<Arc<dyn TurnItemContributor>>,
@@ -45,6 +47,7 @@ impl<C: Sync> Default for ExtensionRegistryBuilder<C> {
             context_contributors: Vec::new(),
             mcp_server_contributors: Vec::new(),
             turn_input_contributors: Vec::new(),
+            sampling_input_contributors: Vec::new(),
             tool_contributors: Vec::new(),
             tool_lifecycle_contributors: Vec::new(),
             turn_item_contributors: Vec::new(),
@@ -114,6 +117,11 @@ impl<C: Sync> ExtensionRegistryBuilder<C> {
         self.turn_input_contributors.push(contributor);
     }
 
+    /// Registers one request-local sampling-input contributor.
+    pub fn sampling_input_contributor(&mut self, contributor: Arc<dyn SamplingInputContributor>) {
+        self.sampling_input_contributors.push(contributor);
+    }
+
     /// Registers one native tool contributor.
     pub fn tool_contributor(&mut self, contributor: Arc<dyn ToolContributor>) {
         self.tool_contributors.push(contributor);
@@ -141,6 +149,7 @@ impl<C: Sync> ExtensionRegistryBuilder<C> {
             context_contributors: self.context_contributors,
             mcp_server_contributors: self.mcp_server_contributors,
             turn_input_contributors: self.turn_input_contributors,
+            sampling_input_contributors: self.sampling_input_contributors,
             tool_contributors: self.tool_contributors,
             tool_lifecycle_contributors: self.tool_lifecycle_contributors,
             turn_item_contributors: self.turn_item_contributors,
@@ -158,6 +167,7 @@ pub struct ExtensionRegistry<C: Sync> {
     context_contributors: Vec<Arc<dyn ContextContributor>>,
     mcp_server_contributors: Vec<Arc<dyn McpServerContributor<C>>>,
     turn_input_contributors: Vec<Arc<dyn TurnInputContributor>>,
+    sampling_input_contributors: Vec<Arc<dyn SamplingInputContributor>>,
     tool_contributors: Vec<Arc<dyn ToolContributor>>,
     tool_lifecycle_contributors: Vec<Arc<dyn ToolLifecycleContributor>>,
     turn_item_contributors: Vec<Arc<dyn TurnItemContributor>>,
@@ -225,6 +235,11 @@ impl<C: Sync> ExtensionRegistry<C> {
         &self.turn_input_contributors
     }
 
+    /// Returns the registered request-local sampling-input contributors.
+    pub fn sampling_input_contributors(&self) -> &[Arc<dyn SamplingInputContributor>] {
+        &self.sampling_input_contributors
+    }
+
     /// Returns the registered native tool contributors.
     pub fn tool_contributors(&self) -> &[Arc<dyn ToolContributor>] {
         &self.tool_contributors

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

@@ -12,6 +12,8 @@ use codex_extension_api::ExtensionEventSink;
 use codex_extension_api::ExtensionFuture;
 use codex_extension_api::ExtensionRegistryBuilder;
 use codex_extension_api::PromptFragment;
+use codex_extension_api::SamplingInputContext;
+use codex_extension_api::SamplingInputContributor;
 use codex_extension_api::ThreadLifecycleContributor;
 use codex_extension_api::TokenUsageContributor;
 use codex_extension_api::ToolCall;
@@ -67,6 +69,15 @@ impl TurnInputContributor for AllContributors {
     }
 }
 
+impl SamplingInputContributor for AllContributors {
+    fn contribute<'a>(
+        &'a self,
+        _input: SamplingInputContext<'a>,
+    ) -> ExtensionFuture<'a, Result<(), String>> {
+        Box::pin(std::future::ready(Ok(())))
+    }
+}
+
 impl ToolContributor for AllContributors {
     fn tools(
         &self,
@@ -117,6 +128,7 @@ async fn build_round_trips_every_contributor_category() {
     builder.token_usage_contributor(contributor.clone());
     builder.prompt_contributor(contributor.clone());
     builder.turn_input_contributor(contributor.clone());
+    builder.sampling_input_contributor(contributor.clone());
     builder.tool_contributor(contributor.clone());
     builder.tool_lifecycle_contributor(contributor.clone());
     builder.turn_item_contributor(contributor.clone());
@@ -129,6 +141,7 @@ async fn build_round_trips_every_contributor_category() {
     assert_eq!(registry.token_usage_contributors().len(), 1);
     assert_eq!(registry.context_contributors().len(), 1);
     assert_eq!(registry.turn_input_contributors().len(), 1);
+    assert_eq!(registry.sampling_input_contributors().len(), 1);
     assert_eq!(registry.tool_contributors().len(), 1);
     assert_eq!(registry.tool_lifecycle_contributors().len(), 1);
     assert_eq!(registry.turn_item_contributors().len(), 1);