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