From 8de2d336b880bed7e6135fedfc2e65e722a744b3 Mon Sep 17 00:00:00 2001 From: Ruslan Nigmatullin Date: Wed, 23 Sep 2026 21:52:23 +0000 Subject: [PATCH] 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/Cargo.lock | 2 + codex-rs/core/src/client.rs | 26 ++- codex-rs/core/src/client_tests.rs | 68 ++++++ codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/model_request.rs | 51 ++++ codex-rs/core/src/model_request_tests.rs | 104 +++++++++ codex-rs/core/src/session/session.rs | 3 +- codex-rs/core/src/session/tests.rs | 3 + codex-rs/core/tests/responses_headers.rs | 3 + codex-rs/core/tests/suite/client.rs | 2 + .../core/tests/suite/client_websockets.rs | 1 + codex-rs/core/tests/suite/mod.rs | 1 + codex-rs/core/tests/suite/model_request.rs | 217 ++++++++++++++++++ codex-rs/ext/extension-api/Cargo.toml | 2 + codex-rs/ext/extension-api/src/lib.rs | 9 + .../ext/extension-api/src/model_request.rs | 44 ++++ codex-rs/ext/extension-api/src/registry.rs | 14 ++ codex-rs/memories/write/src/runtime.rs | 1 + 18 files changed, 550 insertions(+), 2 deletions(-) create mode 100644 codex-rs/core/src/model_request.rs create mode 100644 codex-rs/core/src/model_request_tests.rs create mode 100644 codex-rs/core/tests/suite/model_request.rs create mode 100644 codex-rs/ext/extension-api/src/model_request.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index f761c1f8d7..d23a71ba17 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3260,6 +3260,7 @@ dependencies = [ name = "codex-extension-api" version = "0.0.0" dependencies = [ + "codex-api", "codex-config", "codex-context-fragments", "codex-exec-server-protocol", @@ -3270,6 +3271,7 @@ dependencies = [ "codex-tools", "codex-utils-absolute-path", "codex-utils-path-uri", + "futures", "pretty_assertions", "serde_json", "tokio", diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index f182909905..975aad120a 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -266,6 +266,7 @@ pub struct ModelClient { event_sender: Option>, http_client_factory: HttpClientFactory, restored_history: bool, + request_contributors: Vec>, } /// A turn-scoped streaming session created from a [`ModelClient`]. @@ -496,6 +497,7 @@ impl ModelClient { attestation_provider: Option>, http_client_factory: HttpClientFactory, workspace_routing: WorkspaceRoutingContext, + request_contributors: Vec>, ) -> Self { let model_provider = create_model_provider(provider_info, auth_manager); let codex_api_key_env_enabled = model_provider @@ -542,6 +544,7 @@ impl ModelClient { event_sender: None, http_client_factory, restored_history: false, + request_contributors, } } @@ -1731,6 +1734,13 @@ impl ModelClientSession { let request_session_telemetry = session_telemetry_for_request(session_telemetry, &request); options.extra_headers.extend(responses_headers); + let interceptors = crate::model_request::prepare( + &self.client.request_contributors, + &self.client.state.thread_id.to_string(), + &model_info.slug, + codex_extension_api::ModelRequestKind::Generation, + &mut request.client_metadata, + ); let inference_trace_attempt = inference_trace.start_attempt(); inference_trace_attempt.add_request_headers(&mut options.extra_headers); inference_trace_attempt.record_started(&request); @@ -1749,6 +1759,7 @@ impl ModelClientSession { request_session_telemetry, inference_trace_attempt, Arc::clone(&self.client.state.provider), + interceptors, ); return Ok(stream); } @@ -1998,6 +2009,17 @@ impl ModelClientSession { client_setup.auth.as_ref(), &responses_headers, ); + let interceptors = crate::model_request::prepare( + &self.client.request_contributors, + &self.client.state.thread_id.to_string(), + &model_info.slug, + if warmup { + codex_extension_api::ModelRequestKind::Warmup + } else { + codex_extension_api::ModelRequestKind::Generation + }, + &mut ws_payload.client_metadata, + ); let mut ws_request = ResponsesWsRequest::ResponseCreate(ws_payload); stamp_ws_stream_request_start_ms(&mut ws_request); if !previous_response_id_from_untraced_warmup { @@ -2048,6 +2070,7 @@ impl ModelClientSession { request_session_telemetry, inference_trace_attempt, Arc::clone(&self.client.state.provider), + interceptors, ); self.websocket_session.last_response_rx = Some(last_request_rx); return Ok(WebsocketStreamOutcome::Stream(stream)); @@ -2287,6 +2310,7 @@ fn map_response_stream( session_telemetry: SessionTelemetry, inference_trace_attempt: InferenceTraceAttempt, provider: SharedModelProvider, + interceptors: Vec>, ) -> (ResponseStream, oneshot::Receiver) { let codex_api::ResponseStream { rx_event, @@ -2298,7 +2322,7 @@ fn map_response_stream( }; map_response_events( upstream_request_id, - api_stream, + crate::model_request::intercept_stream(Box::pin(api_stream), interceptors), session_telemetry, inference_trace_attempt, provider, diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 8380e365d7..4f809b0a1d 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -130,6 +130,7 @@ fn test_model_client_with_thread_id( codex_model_provider::WorkspaceRoutingContext::new( "https://chatgpt.com/backend-api".into(), ), + Vec::new(), ) } @@ -1729,6 +1730,7 @@ fn model_client_with_counting_attestation( codex_model_provider::WorkspaceRoutingContext::new( "https://chatgpt.com/backend-api".into(), ), + Vec::new(), ); (model_client, attestation_calls) } @@ -1858,3 +1860,69 @@ async fn non_chatgpt_codex_endpoints_omit_attestation_generation() { ); assert_eq!(attestation_calls.load(Ordering::Relaxed), 0); } + +#[tokio::test] +async fn intercepted_output_reaches_trace_and_websocket_bookkeeping() -> anyhow::Result<()> { + struct ReplaceOutput; + impl codex_extension_api::ModelResponseInterceptor for ReplaceOutput { + fn intercept( + self: Box, + stream: codex_extension_api::ModelResponseStream, + ) -> codex_extension_api::ModelResponseStream { + Box::pin(stream.map(|event| { + event.map(|event| match event { + ResponseEvent::OutputItemDone(_) => { + ResponseEvent::OutputItemDone(output_message("1", "transformed")) + } + other => other, + }) + })) + } + } + + let temp = TempDir::new()?; + let attempt = started_inference_attempt(&temp)?; + let (tx_event, rx_event) = tokio::sync::mpsc::channel(2); + tx_event + .send(Ok(ResponseEvent::OutputItemDone(output_message( + "1", "original", + )))) + .await?; + tx_event + .send(Ok(ResponseEvent::Completed { + response_id: "response".into(), + token_usage: None, + usage_metadata: None, + end_turn: None, + })) + .await?; + drop(tx_event); + let (mut stream, last_response) = super::map_response_stream( + codex_api::ResponseStream { + rx_event, + upstream_request_id: None, + }, + test_session_telemetry(), + attempt, + test_model_provider(), + vec![Box::new(ReplaceOutput)], + ); + let mut delivered = Vec::new(); + while let Some(event) = stream.next().await { + if let ResponseEvent::OutputItemDone(item) = event? { + delivered.push(item); + } + } + assert_eq!(delivered, vec![output_message("1", "transformed")]); + assert_eq!(last_response.await?.items_added, delivered); + let rollout = replay_bundle(temp.path())?; + let payload = rollout + .raw_payloads + .values() + .find(|payload| payload.kind == codex_rollout_trace::RawPayloadKind::InferenceResponse) + .expect("response trace payload"); + let recorded: serde_json::Value = + serde_json::from_slice(&std::fs::read(temp.path().join(&payload.path))?)?; + assert_eq!(recorded["output_items"], serde_json::to_value(&delivered)?); + Ok(()) +} diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 61fa679c03..1ce7587fd6 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -9,6 +9,7 @@ mod apply_patch; mod apps; mod client; mod client_common; +mod model_request; mod realtime_context; mod realtime_conversation; mod realtime_history; diff --git a/codex-rs/core/src/model_request.rs b/codex-rs/core/src/model_request.rs new file mode 100644 index 0000000000..8b13422bc7 --- /dev/null +++ b/codex-rs/core/src/model_request.rs @@ -0,0 +1,51 @@ +//! Compose request-scoped response interceptors before tracing and response bookkeeping. +use codex_extension_api::ModelRequestContributor; +use codex_extension_api::ModelRequestInput; +use codex_extension_api::ModelRequestKind; +use codex_extension_api::ModelResponseInterceptor; +use codex_extension_api::ModelResponseStream; +use std::collections::HashMap; +use std::sync::Arc; + +pub(crate) fn prepare( + contributors: &[Arc], + thread_id: &str, + model: &str, + kind: ModelRequestKind, + metadata: &mut Option>, +) -> Vec> { + contributors + .iter() + .filter_map(|contributor| { + let mut additions = None; + let interceptor = contributor.request(ModelRequestInput { + kind, + thread_id, + client_metadata: &mut additions, + model, + }); + if let Some(additions) = additions { + for (key, value) in crate::responses_metadata::filter_extra_metadata(additions) { + if key != "parent_response_id" { + metadata.get_or_insert_default().entry(key).or_insert(value); + } + } + } + interceptor + }) + .collect() +} + +pub(crate) fn intercept_stream( + mut stream: ModelResponseStream, + interceptors: Vec>, +) -> ModelResponseStream { + for interceptor in interceptors { + stream = interceptor.intercept(stream); + } + stream +} + +#[cfg(test)] +#[path = "model_request_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/model_request_tests.rs b/codex-rs/core/src/model_request_tests.rs new file mode 100644 index 0000000000..82da1f65fe --- /dev/null +++ b/codex-rs/core/src/model_request_tests.rs @@ -0,0 +1,104 @@ +//! Interceptors can transform incremental events and retain upstream cancellation. +use super::*; +use crate::client_common::ResponseStream; +use codex_extension_api::ResponseEvent; +use futures::StreamExt; +use pretty_assertions::assert_eq; +use std::time::Duration; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +struct Append(&'static str); +impl ModelResponseInterceptor for Append { + fn intercept(self: Box, stream: ModelResponseStream) -> ModelResponseStream { + Box::pin(stream.map(move |event| { + event.map(|event| match event { + ResponseEvent::OutputTextDelta(text) => { + ResponseEvent::OutputTextDelta(format!("{text}{}", self.0)) + } + other => other, + }) + })) + } +} + +#[tokio::test] +async fn interceptors_forward_before_completion_in_order_and_cancel_upstream() { + let (tx_event, rx_event) = mpsc::channel(1); + let consumer_dropped = CancellationToken::new(); + let upstream_cancelled = consumer_dropped.clone(); + let mut stream = intercept_stream( + Box::pin( + ResponseStream { + rx_event, + consumer_dropped, + } + .map(|event| { + event.map_err(|error| { + codex_extension_api::ModelResponseError::Stream(error.to_string()) + }) + }), + ), + vec![Box::new(Append("a")), Box::new(Append("b"))], + ); + tx_event + .send(Ok(ResponseEvent::OutputTextDelta("text".into()))) + .await + .unwrap(); + // Upstream remains open and has not sent Completed. + let event = tokio::time::timeout(Duration::from_secs(1), stream.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + let ResponseEvent::OutputTextDelta(text) = event else { + panic!("expected text delta"); + }; + assert_eq!(text, "textab"); + drop(stream); + tokio::time::timeout(Duration::from_secs(1), upstream_cancelled.cancelled()) + .await + .unwrap(); +} + +#[derive(Debug)] +struct MetadataContributor; +impl ModelRequestContributor for MetadataContributor { + fn request(&self, input: ModelRequestInput<'_>) -> Option> { + *input.client_metadata = Some(HashMap::from([ + ("thread_id".into(), "replaced".into()), + ("guardian_credits_requested".into(), "true".into()), + ("parent_response_id".into(), "forged".into()), + ("existing".into(), "replaced".into()), + ("x-request-id".into(), "request".into()), + ])); + None + } +} + +#[test] +fn contributor_metadata_preserves_core_fields_and_existing_values() { + let mut metadata = Some(HashMap::from([ + ("thread_id".into(), "thread".into()), + ("existing".into(), "original".into()), + ])); + let contributors: Vec> = vec![Arc::new(MetadataContributor)]; + assert!( + prepare( + &contributors, + "thread", + "model", + ModelRequestKind::Generation, + &mut metadata + ) + .is_empty() + ); + assert_eq!( + metadata, + Some(HashMap::from([ + ("thread_id".into(), "thread".into()), + ("existing".into(), "original".into()), + ("x-request-id".into(), "request".into()), + ])) + ); +} diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 051f6b0a11..30e70ead55 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1689,7 +1689,7 @@ impl Session { agents_md_manager, plugins_manager: Arc::clone(&plugins_manager), mcp_manager: Arc::clone(&mcp_manager), - extensions, + extensions: Arc::clone(&extensions), // TODO(jif): extract session to share between sub-agents session_extension_data, thread_extension_data, @@ -1731,6 +1731,7 @@ impl Session { attestation_provider, config.http_client_factory(), workspace_routing.as_ref().clone(), + extensions.model_request_contributors().to_vec(), ) .with_restored_history(matches!( &initial_history, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 9c81bb39fe..a4efea636f 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -847,6 +847,7 @@ fn test_model_client_session() -> crate::client::ModelClientSession { codex_model_provider::WorkspaceRoutingContext::new( "https://chatgpt.com/backend-api".into(), ), + Vec::new(), ) .new_session() } @@ -6528,6 +6529,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { /*attestation_provider*/ None, config.http_client_factory(), config.workspace_routing_context(), + Vec::new(), ), executed_tool_calls: executed_tool_calls.clone(), code_mode_service: crate::tools::code_mode::CodeModeService::new( @@ -8798,6 +8800,7 @@ where /*attestation_provider*/ None, config.http_client_factory(), config.workspace_routing_context(), + Vec::new(), ), executed_tool_calls: executed_tool_calls.clone(), code_mode_service: crate::tools::code_mode::CodeModeService::new( diff --git a/codex-rs/core/tests/responses_headers.rs b/codex-rs/core/tests/responses_headers.rs index 0b1b91dec3..fe456c9a04 100644 --- a/codex-rs/core/tests/responses_headers.rs +++ b/codex-rs/core/tests/responses_headers.rs @@ -140,6 +140,7 @@ async fn responses_stream_includes_subagent_header_on_review() { /*attestation_provider*/ None, config.http_client_factory(), config.workspace_routing_context(), + Vec::new(), ); let responses_metadata = test_turn_responses_metadata(&client, thread_id, &session_source); let mut client_session = client.new_session(); @@ -281,6 +282,7 @@ async fn responses_stream_includes_subagent_header_on_other() { /*attestation_provider*/ None, config.http_client_factory(), config.workspace_routing_context(), + Vec::new(), ); let responses_metadata = test_turn_responses_metadata(&client, thread_id, &session_source); let mut client_session = client.new_session(); @@ -407,6 +409,7 @@ async fn responses_respects_model_info_overrides_from_config() { /*attestation_provider*/ None, config.http_client_factory(), config.workspace_routing_context(), + Vec::new(), ); let responses_metadata = test_turn_responses_metadata(&client, thread_id, &session_source); let mut client_session = client.new_session(); diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 48cc71457e..e8591fddc9 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -1712,6 +1712,7 @@ async fn send_request_with_provider(provider: ModelProviderInfo) { /*attestation_provider*/ None, config.http_client_factory(), config.workspace_routing_context(), + Vec::new(), ); let responses_metadata = test_turn_responses_metadata(&client, thread_id); let mut client_session = client.new_session(); @@ -3221,6 +3222,7 @@ async fn azure_responses_request_does_not_store_and_preserves_prefixed_item_ids( /*attestation_provider*/ None, config.http_client_factory(), config.workspace_routing_context(), + Vec::new(), ); let responses_metadata = test_turn_responses_metadata(&client, thread_id); let mut client_session = client.new_session(); diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index a652b37558..2da2ff7116 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -2694,6 +2694,7 @@ async fn websocket_harness_with_provider_options_and_auth( /*attestation_provider*/ None, http_client_factory, config.workspace_routing_context(), + Vec::new(), ); WebsocketTestHarness { diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index fcae96b487..96a29dce08 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -129,6 +129,7 @@ mod mcp_user_verification; mod model_overrides; #[path = "model_provider_requirements_tests.rs"] mod model_provider_requirements; +mod model_request; mod model_runtime_selectors; mod model_switching; mod model_visible_layout; diff --git a/codex-rs/core/tests/suite/model_request.rs b/codex-rs/core/tests/suite/model_request.rs new file mode 100644 index 0000000000..0dd5e3d7e7 --- /dev/null +++ b/codex-rs/core/tests/suite/model_request.rs @@ -0,0 +1,217 @@ +//! Output gates allow inference immediately, hold complete output, and cancel with the turn. +use anyhow::Result; +use codex_core::TurnInputRequest; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::ModelRequestContributor; +use codex_extension_api::ModelRequestInput; +use codex_extension_api::ModelRequestKind; +use codex_extension_api::ModelResponseInterceptor; +use codex_extension_api::ModelResponseStream; +use codex_extension_api::ResponseEvent; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::user_input::UserInput; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::sse; +use core_test_support::responses::{self}; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use futures::StreamExt; +use pretty_assertions::assert_eq; +use std::sync::Arc; +use tokio::sync::Semaphore; +use tokio::time::Duration; +use tokio::time::timeout; + +#[derive(Debug)] +struct Gates { + delivered: Semaphore, + output_waiting: Semaphore, + cancelled: Semaphore, + output_items: std::sync::atomic::AtomicUsize, +} +#[derive(Debug)] +struct Contributor(Arc); +impl ModelRequestContributor for Contributor { + fn request(&self, input: ModelRequestInput<'_>) -> Option> { + input + .client_metadata + .get_or_insert_default() + .insert("test-interceptor".into(), "present".into()); + if input.kind == ModelRequestKind::Warmup { + return None; + } + Some(Box::new(Permit { + gates: self.0.clone(), + released: false, + })) + } +} +struct Permit { + gates: Arc, + released: bool, +} +impl ModelResponseInterceptor for Permit { + fn intercept(mut self: Box, mut stream: ModelResponseStream) -> ModelResponseStream { + Box::pin( + futures::stream::once(async move { + let mut events = Vec::new(); + let mut output_items = 0; + while let Some(event) = stream.next().await { + if matches!(&event, Ok(ResponseEvent::OutputItemDone(_))) { + output_items += 1; + } + let completed = matches!(&event, Ok(ResponseEvent::Completed { .. })); + events.push(event); + if completed { + self.gates + .output_items + .store(output_items, std::sync::atomic::Ordering::SeqCst); + self.gates.output_waiting.add_permits(1); + self.gates + .delivered + .acquire() + .await + .expect("delivery semaphore remains open") + .forget(); + self.released = true; + break; + } + } + futures::stream::iter(events) + }) + .flatten(), + ) + } +} +impl Drop for Permit { + fn drop(&mut self) { + if !self.released { + self.gates.cancelled.add_permits(1); + } + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn model_request_gates_output_and_cancellation() -> Result<()> { + let server = responses::start_mock_server().await; + let response = sse(vec![ + ev_response_created("r1"), + ev_assistant_message("m1", "first"), + ev_assistant_message("m2", "done"), + ev_completed("r1"), + ]); + let model = responses::mount_sse_sequence(&server, vec![response.clone(), response]).await; + let gates = Arc::new(Gates { + delivered: Semaphore::new(0), + output_waiting: Semaphore::new(0), + cancelled: Semaphore::new(0), + output_items: std::sync::atomic::AtomicUsize::new(0), + }); + let mut registry = ExtensionRegistryBuilder::new(); + registry.model_request_contributor(Arc::new(Contributor(gates.clone()))); + let mut builder = test_codex().with_extensions(Arc::new(registry.build())); + let test = builder.build_with_auto_env(&server).await?; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "hello".into(), + text_elements: vec![], + }])) + .await?; + timeout(Duration::from_secs(10), gates.output_waiting.acquire()) + .await?? + .forget(); + assert_eq!(model.requests().len(), 1); + assert_eq!( + gates.output_items.load(std::sync::atomic::Ordering::SeqCst), + 2 + ); + assert_eq!( + model.single_request().body_json()["client_metadata"]["test-interceptor"], + "present" + ); + // Drain earlier events and prove no assistant output passes a closed delivery gate. + assert!( + timeout( + Duration::from_millis(100), + wait_for_event(&test.codex, |e| matches!( + e, + EventMsg::AgentMessage(_) | EventMsg::TurnComplete(_) + )) + ) + .await + .is_err() + ); + gates.delivered.add_permits(1); + wait_for_event(&test.codex, |e| matches!(e, EventMsg::TurnComplete(_))).await; + assert_eq!(gates.cancelled.available_permits(), 0); + + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "again".into(), + text_elements: vec![], + }])) + .await?; + timeout(Duration::from_secs(10), gates.output_waiting.acquire()) + .await?? + .forget(); + test.codex.submit(Op::Interrupt).await?; + timeout(Duration::from_secs(10), gates.cancelled.acquire()) + .await?? + .forget(); + assert_eq!(model.requests().len(), 2); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn model_request_gates_websocket_output_but_not_warmup() -> Result<()> { + let server = responses::start_websocket_server(vec![vec![ + vec![ev_response_created("warm"), ev_completed("warm")], + vec![ + ev_response_created("response"), + ev_assistant_message("message", "done"), + ev_completed("response"), + ], + ]]) + .await; + let gates = Arc::new(Gates { + delivered: Semaphore::new(0), + output_waiting: Semaphore::new(0), + cancelled: Semaphore::new(0), + output_items: std::sync::atomic::AtomicUsize::new(0), + }); + let mut registry = ExtensionRegistryBuilder::new(); + registry.model_request_contributor(Arc::new(Contributor(gates.clone()))); + let mut builder = test_codex() + .with_model("gpt-5.2") + .with_extensions(Arc::new(registry.build())); + let test = builder.build_with_websocket_server(&server).await?; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "hello".into(), + text_elements: vec![], + }])) + .await?; + timeout(Duration::from_secs(10), gates.output_waiting.acquire()) + .await?? + .forget(); + let requests = server.single_connection(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[0].body_json()["generate"], false); + assert_eq!( + requests[0].body_json()["client_metadata"]["test-interceptor"], + "present" + ); + assert_eq!( + requests[1].body_json()["client_metadata"]["test-interceptor"], + "present" + ); + test.codex.submit(Op::Interrupt).await?; + timeout(Duration::from_secs(10), gates.cancelled.acquire()) + .await?? + .forget(); + server.shutdown().await; + Ok(()) +} diff --git a/codex-rs/ext/extension-api/Cargo.toml b/codex-rs/ext/extension-api/Cargo.toml index 7769b788a5..e98063a8d2 100644 --- a/codex-rs/ext/extension-api/Cargo.toml +++ b/codex-rs/ext/extension-api/Cargo.toml @@ -14,6 +14,8 @@ doctest = false workspace = true [dependencies] +codex-api = { workspace = true } +futures = { workspace = true } codex-history = { workspace = true } codex-config = { workspace = true } codex-context-fragments = { workspace = true } diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 2e6c65f662..ef2f3886a2 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/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; diff --git a/codex-rs/ext/extension-api/src/model_request.rs b/codex-rs/ext/extension-api/src/model_request.rs new file mode 100644 index 0000000000..67803d28d3 --- /dev/null +++ b/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>; + +/// 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>, + /// 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>; +} + +/// 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, stream: ModelResponseStream) -> ModelResponseStream; +} diff --git a/codex-rs/ext/extension-api/src/registry.rs b/codex-rs/ext/extension-api/src/registry.rs index 357cb681d6..ddb02113ed 100644 --- a/codex-rs/ext/extension-api/src/registry.rs +++ b/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 Default for ExtensionRegistryBuilder { 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 ExtensionRegistryBuilder { self.registry.turn_input_contributors.push(contributor); } + /// Registers a contributor for request-scoped response interception. + pub fn model_request_contributor(&mut self, contributor: Arc) { + self.registry.model_request_contributors.push(contributor); + } + /// Registers one native tool contributor. pub fn tool_contributor(&mut self, contributor: Arc) { self.registry.tool_contributors.push(contributor); @@ -157,6 +164,7 @@ pub struct ExtensionRegistry { turn_input_contributors: Vec>, tool_contributors: Vec>, tool_lifecycle_contributors: Vec>, + model_request_contributors: Vec>, turn_item_contributors: Vec>, approval_review_contributors: Vec>, } @@ -178,6 +186,7 @@ impl ExtensionRegistry { 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 ExtensionRegistry { &self.tool_contributors } + /// Returns response interceptor contributors in registration order. + pub fn model_request_contributors(&self) -> &[Arc] { + &self.model_request_contributors + } + /// Returns the registered tool-lifecycle contributors. pub fn tool_lifecycle_contributors(&self) -> &[Arc] { &self.tool_lifecycle_contributors diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs index af02601b2b..8d09f6dda4 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -339,6 +339,7 @@ impl MemoryStartupContext { /*attestation_provider*/ None, config.http_client_factory(), config.workspace_routing_context(), + Vec::new(), ); let mut client_session = model_client.new_session();