From 0585f95bb7c3e1ba4183fca1a004900702052d56 Mon Sep 17 00:00:00 2001 From: Rohit Arunachalam Date: Wed, 17 Jun 2026 11:14:25 -0700 Subject: [PATCH] [codex] add sampling input extension hook --- codex-rs/core/src/session/turn.rs | 23 +++- .../tests/suite/extension_sampling_input.rs | 116 ++++++++++++++++++ codex-rs/core/tests/suite/mod.rs | 1 + .../ext/extension-api/src/contributors.rs | 14 +++ .../src/contributors/sampling_input.rs | 21 ++++ codex-rs/ext/extension-api/src/lib.rs | 2 + codex-rs/ext/extension-api/src/registry.rs | 15 +++ codex-rs/ext/extension-api/tests/registry.rs | 13 ++ 8 files changed, 204 insertions(+), 1 deletion(-) create mode 100644 codex-rs/core/tests/suite/extension_sampling_input.rs create mode 100644 codex-rs/ext/extension-api/src/contributors/sampling_input.rs diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 4beef792bb..d2240b4ca0 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/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(), diff --git a/codex-rs/core/tests/suite/extension_sampling_input.rs b/codex-rs/core/tests/suite/extension_sampling_input.rs new file mode 100644 index 0000000000..a73835ab97 --- /dev/null +++ b/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::::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::>(); + 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(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 87dc0ffce3..14b8e336f3 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/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; diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 3a0fee49a9..18d18d9cd4 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/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>>; } +/// 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 diff --git a/codex-rs/ext/extension-api/src/contributors/sampling_input.rs b/codex-rs/ext/extension-api/src/contributors/sampling_input.rs new file mode 100644 index 0000000000..06212da759 --- /dev/null +++ b/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, +} diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index e646a53a52..79d28c36bf 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/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; diff --git a/codex-rs/ext/extension-api/src/registry.rs b/codex-rs/ext/extension-api/src/registry.rs index e8cc989337..fc5f47d7a0 100644 --- a/codex-rs/ext/extension-api/src/registry.rs +++ b/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 { context_contributors: Vec>, mcp_server_contributors: Vec>>, turn_input_contributors: Vec>, + sampling_input_contributors: Vec>, tool_contributors: Vec>, tool_lifecycle_contributors: Vec>, turn_item_contributors: Vec>, @@ -45,6 +47,7 @@ impl Default for ExtensionRegistryBuilder { 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 ExtensionRegistryBuilder { self.turn_input_contributors.push(contributor); } + /// Registers one request-local sampling-input contributor. + pub fn sampling_input_contributor(&mut self, contributor: Arc) { + self.sampling_input_contributors.push(contributor); + } + /// Registers one native tool contributor. pub fn tool_contributor(&mut self, contributor: Arc) { self.tool_contributors.push(contributor); @@ -141,6 +149,7 @@ impl ExtensionRegistryBuilder { 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 { context_contributors: Vec>, mcp_server_contributors: Vec>>, turn_input_contributors: Vec>, + sampling_input_contributors: Vec>, tool_contributors: Vec>, tool_lifecycle_contributors: Vec>, turn_item_contributors: Vec>, @@ -225,6 +235,11 @@ impl ExtensionRegistry { &self.turn_input_contributors } + /// Returns the registered request-local sampling-input contributors. + pub fn sampling_input_contributors(&self) -> &[Arc] { + &self.sampling_input_contributors + } + /// Returns the registered native tool contributors. pub fn tool_contributors(&self) -> &[Arc] { &self.tool_contributors diff --git a/codex-rs/ext/extension-api/tests/registry.rs b/codex-rs/ext/extension-api/tests/registry.rs index 79773a98cf..7c04baa8f1 100644 --- a/codex-rs/ext/extension-api/tests/registry.rs +++ b/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);