Add extension hooks for model requests and response streams (#47679)

## What changed

Add `ModelRequestContributor` and `ModelResponseInterceptor` to the extension API. Contributors can add filtered `client_metadata` without overwriting existing or reserved fields and create interceptors for individual requests.

Wire contributors into HTTP and WebSocket requests, distinguishing generation from warmup. Apply response interceptors in registration order before tracing and response bookkeeping so transformed output is recorded consistently. Interceptors can inspect, buffer, or transform events, with stream ownership carrying request state and cancellation.

## Testing

Add tests for metadata protection, incremental event transformation, interceptor ordering, upstream cancellation, and consistency between delivered output, traces, and WebSocket bookkeeping. Integration tests cover output gating and interruption over HTTP and WebSocket, including a contributor that bypasses gating for warmup requests.

GitOrigin-RevId: d4bc1745d2cacf8636728dbc0a58ab9f250e9392

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

@@ -110,3 +110,12 @@ pub use user_instructions::ThreadInstructionsProvider;
 pub use user_instructions::UserInstructionsProvider;
 
 pub use contributors::ToolDispatchInput;
+
+mod model_request;
+pub use model_request::ModelRequestContributor;
+pub use model_request::ModelRequestInput;
+pub use model_request::ModelRequestKind;
+pub use model_request::ModelResponseError;
+pub use model_request::ModelResponseInterceptor;
+pub use model_request::ModelResponseStream;
+pub use model_request::ResponseEvent;

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

@@ -0,0 +1,44 @@
+//! Request-scoped interception of model response streams.
+pub use codex_api::ApiError as ModelResponseError;
+pub use codex_api::ResponseEvent;
+use futures::stream::BoxStream;
+use std::collections::HashMap;
+
+/// Model response events owned by an interceptor or its downstream consumer.
+pub type ModelResponseStream = BoxStream<'static, Result<ResponseEvent, ModelResponseError>>;
+
+/// Whether a request generates output or only prepares a session.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum ModelRequestKind {
+    /// Generate model output.
+    Generation,
+    /// Prepare the session without generating output.
+    Warmup,
+}
+
+/// Request identity and additional metadata supplied before sending to the provider.
+pub struct ModelRequestInput<'a> {
+    /// Whether this request generates output or warms up the session.
+    pub kind: ModelRequestKind,
+    /// Thread making the request.
+    pub thread_id: &'a str,
+    /// Additional provider metadata for request tracking; existing and reserved keys cannot be replaced.
+    pub client_metadata: &'a mut Option<HashMap<String, String>>,
+    /// Requested model identifier.
+    pub model: &'a str,
+}
+
+/// Creates request-scoped interceptors without delaying inference.
+/// Return None for requests this host does not manage.
+pub trait ModelRequestContributor: Send + Sync + std::fmt::Debug {
+    fn request(&self, input: ModelRequestInput<'_>) -> Option<Box<dyn ModelResponseInterceptor>>;
+}
+
+/// Owns request state until consumed by the response stream or dropped on failure.
+/// Interceptors compose in registration order; dropping their stream must release unfinished work.
+/// Trusted host implementations must bound any model-visible content they synthesize.
+pub trait ModelResponseInterceptor: Send + Sync {
+    /// Wrap the response stream, optionally buffering, inspecting, or transforming events.
+    /// The returned stream must own request state and propagate cancellation by dropping upstream.
+    fn intercept(self: Box<Self>, stream: ModelResponseStream) -> ModelResponseStream;
+}

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

@@ -5,6 +5,7 @@ use crate::ConfigContributor;
 use crate::ContextContributor;
 use crate::ExtensionEventSink;
 use crate::McpServerContributor;
+use crate::ModelRequestContributor;
 use crate::NoopExtensionEventSink;
 use crate::SkillInvocationContributor;
 use crate::ThreadLifecycleContributor;
@@ -38,6 +39,7 @@ impl<C: Sync> Default for ExtensionRegistryBuilder<C> {
                 turn_input_contributors: Vec::new(),
                 tool_contributors: Vec::new(),
                 tool_lifecycle_contributors: Vec::new(),
+                model_request_contributors: Vec::new(),
                 turn_item_contributors: Vec::new(),
             },
         }
@@ -122,6 +124,11 @@ impl<C: Sync> ExtensionRegistryBuilder<C> {
         self.registry.turn_input_contributors.push(contributor);
     }
 
+    /// Registers a contributor for request-scoped response interception.
+    pub fn model_request_contributor(&mut self, contributor: Arc<dyn ModelRequestContributor>) {
+        self.registry.model_request_contributors.push(contributor);
+    }
+
     /// Registers one native tool contributor.
     pub fn tool_contributor(&mut self, contributor: Arc<dyn ToolContributor>) {
         self.registry.tool_contributors.push(contributor);
@@ -157,6 +164,7 @@ pub struct ExtensionRegistry<C: Sync> {
     turn_input_contributors: Vec<Arc<dyn TurnInputContributor>>,
     tool_contributors: Vec<Arc<dyn ToolContributor>>,
     tool_lifecycle_contributors: Vec<Arc<dyn ToolLifecycleContributor>>,
+    model_request_contributors: Vec<Arc<dyn ModelRequestContributor>>,
     turn_item_contributors: Vec<Arc<dyn TurnItemContributor>>,
     approval_review_contributors: Vec<Arc<dyn ApprovalReviewContributor>>,
 }
@@ -178,6 +186,7 @@ impl<C: Sync> ExtensionRegistry<C> {
                 turn_input_contributors: self.turn_input_contributors.clone(),
                 tool_contributors: self.tool_contributors.clone(),
                 tool_lifecycle_contributors: self.tool_lifecycle_contributors.clone(),
+                model_request_contributors: self.model_request_contributors.clone(),
                 turn_item_contributors: self.turn_item_contributors.clone(),
                 approval_review_contributors: self.approval_review_contributors.clone(),
             },
@@ -267,6 +276,11 @@ impl<C: Sync> ExtensionRegistry<C> {
         &self.tool_contributors
     }
 
+    /// Returns response interceptor contributors in registration order.
+    pub fn model_request_contributors(&self) -> &[Arc<dyn ModelRequestContributor>] {
+        &self.model_request_contributors
+    }
+
     /// Returns the registered tool-lifecycle contributors.
     pub fn tool_lifecycle_contributors(&self) -> &[Arc<dyn ToolLifecycleContributor>] {
         &self.tool_lifecycle_contributors