diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index b6ece83ddf..3df52caa38 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -417,6 +417,10 @@ pub(crate) struct CodexMessageProcessorArgs { } impl CodexMessageProcessor { + pub(crate) fn thread_state_manager(&self) -> ThreadStateManager { + self.thread_state_manager.clone() + } + pub(crate) fn clear_plugin_related_caches(&self) { self.thread_manager.plugins_manager().clear_cache(); self.thread_manager.skills_manager().clear_cache(); @@ -1839,6 +1843,7 @@ impl CodexMessageProcessor { service_name, base_instructions, developer_instructions, + network_delegation, dynamic_tools, mock_experimental_field: _mock_experimental_field, experimental_raw_events, @@ -1879,6 +1884,7 @@ impl CodexMessageProcessor { config, typesafe_overrides, dynamic_tools, + network_delegation, persist_extended_history, service_name, experimental_raw_events, @@ -1944,6 +1950,7 @@ impl CodexMessageProcessor { config_overrides: Option>, typesafe_overrides: ConfigOverrides, dynamic_tools: Option>, + network_delegation: Option, persist_extended_history: bool, service_name: Option, experimental_raw_events: bool, @@ -2021,6 +2028,15 @@ impl CodexMessageProcessor { session_configured, .. } = new_conv; + if let Err(error) = + Self::set_network_delegation(thread.as_ref(), network_delegation).await + { + listener_task_context + .outgoing + .send_error(request_id, error) + .await; + return; + } let config_snapshot = thread .config_snapshot() .instrument(tracing::info_span!( @@ -3413,6 +3429,7 @@ impl CodexMessageProcessor { base_instructions, developer_instructions, personality, + network_delegation, persist_extended_history, } = params; @@ -3494,6 +3511,12 @@ impl CodexMessageProcessor { thread, session_configured, }) => { + if let Err(error) = + Self::set_network_delegation(thread.as_ref(), network_delegation).await + { + self.outgoing.send_error(request_id, error).await; + return; + } let SessionConfiguredEvent { rollout_path, .. } = session_configured; let Some(rollout_path) = rollout_path else { self.send_internal_error( @@ -3913,6 +3936,7 @@ impl CodexMessageProcessor { config: cli_overrides, base_instructions, developer_instructions, + network_delegation, ephemeral, persist_extended_history, } = params; @@ -4060,6 +4084,12 @@ impl CodexMessageProcessor { return; } }; + if let Err(error) = + Self::set_network_delegation(forked_thread.as_ref(), network_delegation).await + { + self.outgoing.send_error(request_id, error).await; + return; + } // Auto-attach a conversation listener when forking a thread. Self::log_listener_attach_result( @@ -6059,6 +6089,20 @@ impl CodexMessageProcessor { }) } + async fn set_network_delegation( + thread: &CodexThread, + network_delegation: Option, + ) -> Result<(), JSONRPCErrorError> { + thread + .set_network_delegation(network_delegation) + .await + .map_err(|err| JSONRPCErrorError { + code: INTERNAL_ERROR_CODE, + message: format!("failed to set network delegation: {err}"), + data: None, + }) + } + async fn turn_steer(&self, request_id: ConnectionRequestId, params: TurnSteerParams) { let (_, thread) = match self.load_thread(¶ms.thread_id).await { Ok(v) => v, @@ -7438,6 +7482,14 @@ fn collect_resume_override_mismatches( config_snapshot.service_tier )); } + if let Some(requested_network_delegation) = request.network_delegation.as_ref() + && config_snapshot.network_delegation.as_ref() != Some(requested_network_delegation) + { + mismatch_details.push(format!( + "network_delegation requested={requested_network_delegation:?} active={:?}", + config_snapshot.network_delegation + )); + } if let Some(requested_cwd) = request.cwd.as_deref() { let requested_cwd_path = std::path::PathBuf::from(requested_cwd); if requested_cwd_path != config_snapshot.cwd { @@ -8423,6 +8475,7 @@ mod tests { base_instructions: None, developer_instructions: None, personality: None, + network_delegation: None, persist_extended_history: false, }; let config_snapshot = ThreadConfigSnapshot { @@ -8436,6 +8489,7 @@ mod tests { ephemeral: false, reasoning_effort: None, personality: None, + network_delegation: None, session_source: SessionSource::Cli, }; diff --git a/codex-rs/app-server/src/delegated_model.rs b/codex-rs/app-server/src/delegated_model.rs new file mode 100644 index 0000000000..9a85630bcb --- /dev/null +++ b/codex-rs/app-server/src/delegated_model.rs @@ -0,0 +1,650 @@ +use crate::outgoing_message::ClientRequestResult; +use crate::outgoing_message::ConnectionId; +use crate::outgoing_message::OutgoingMessageSender; +use crate::outgoing_message::ThreadScopedOutgoingMessageSender; +use crate::server_request_error::is_turn_transition_server_request_error; +use crate::thread_state::ThreadStateManager; +use async_trait::async_trait; +use codex_app_server_protocol::ClientNotification; +use codex_app_server_protocol::ModelCompactParams; +use codex_app_server_protocol::ModelCompactResponse; +use codex_app_server_protocol::ModelRequestFailedNotification; +use codex_app_server_protocol::ModelRequestParams; +use codex_app_server_protocol::ModelRequestResponse; +use codex_app_server_protocol::ModelStreamEventNotification; +use codex_app_server_protocol::ModelStreamMetadataNotification; +use codex_app_server_protocol::ServerRequestPayload; +use codex_core::delegated_model_transport::DelegatedModelCompactRequest; +use codex_core::delegated_model_transport::DelegatedModelEvent; +use codex_core::delegated_model_transport::DelegatedModelRequest; +use codex_core::delegated_model_transport::DelegatedModelTransport; +use codex_core::error::CodexErr; +use codex_core::error::Result; +use codex_protocol::ThreadId; +use std::collections::HashMap; +use std::sync::Arc; +use tokio::sync::Mutex; +use tokio::sync::mpsc; + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +struct DelegatedRequestKey { + thread_id: String, + turn_id: String, + request_id: String, +} + +impl DelegatedRequestKey { + fn from_request(request: &DelegatedModelRequest) -> Self { + Self { + thread_id: request.thread_id.to_string(), + turn_id: request.turn_id.clone(), + request_id: request.request_id.clone(), + } + } + + fn from_metadata(notification: &ModelStreamMetadataNotification) -> Self { + Self { + thread_id: notification.thread_id.clone(), + turn_id: notification.turn_id.clone(), + request_id: notification.request_id.clone(), + } + } + + fn from_stream_event(notification: &ModelStreamEventNotification) -> Self { + Self { + thread_id: notification.thread_id.clone(), + turn_id: notification.turn_id.clone(), + request_id: notification.request_id.clone(), + } + } + + fn from_failed(notification: &ModelRequestFailedNotification) -> Self { + Self { + thread_id: notification.thread_id.clone(), + turn_id: notification.turn_id.clone(), + request_id: notification.request_id.clone(), + } + } +} + +#[derive(Clone)] +pub(crate) struct AppServerDelegatedModelTransport { + outgoing: Arc, + thread_state_manager: ThreadStateManager, + inflight: Arc>>>, +} + +impl AppServerDelegatedModelTransport { + pub(crate) fn new( + outgoing: Arc, + thread_state_manager: ThreadStateManager, + ) -> Self { + Self { + outgoing, + thread_state_manager, + inflight: Arc::new(Mutex::new(HashMap::new())), + } + } + + pub(crate) async fn handle_client_notification( + &self, + notification: ClientNotification, + ) -> Result<()> { + match notification { + ClientNotification::Initialized => Ok(()), + ClientNotification::ModelStreamMetadata(notification) => { + self.deliver_metadata(notification).await + } + ClientNotification::ModelStreamEvent(notification) => { + self.deliver_stream_event(notification).await + } + ClientNotification::ModelRequestFailed(notification) => { + self.deliver_request_failed(notification).await + } + } + } + + async fn deliver_metadata(&self, notification: ModelStreamMetadataNotification) -> Result<()> { + let key = DelegatedRequestKey::from_metadata(¬ification); + self.deliver_event( + key, + DelegatedModelEvent::StreamMetadata(notification.metadata), + /*final_event*/ false, + ) + .await + } + + async fn deliver_stream_event(&self, notification: ModelStreamEventNotification) -> Result<()> { + let final_event = stream_event_is_final(¬ification.event); + let key = DelegatedRequestKey::from_stream_event(¬ification); + self.deliver_event( + key, + DelegatedModelEvent::StreamEvent(notification.event), + final_event, + ) + .await + } + + async fn deliver_request_failed( + &self, + notification: ModelRequestFailedNotification, + ) -> Result<()> { + let key = DelegatedRequestKey::from_failed(¬ification); + self.deliver_event( + key, + DelegatedModelEvent::RequestFailed(notification.error), + /*final_event*/ true, + ) + .await + } + + async fn deliver_event( + &self, + key: DelegatedRequestKey, + event: DelegatedModelEvent, + final_event: bool, + ) -> Result<()> { + let sender = { + let inflight = self.inflight.lock().await; + inflight.get(&key).cloned() + } + .ok_or_else(|| { + CodexErr::InvalidRequest(format!( + "received delegated model notification for unknown request: thread_id={} turn_id={} request_id={}", + key.thread_id, key.turn_id, key.request_id + )) + })?; + + if sender.send(event).await.is_err() { + self.inflight.lock().await.remove(&key); + return Err(CodexErr::InvalidRequest(format!( + "received delegated model notification for completed request: thread_id={} turn_id={} request_id={}", + key.thread_id, key.turn_id, key.request_id + ))); + } + if final_event { + self.inflight.lock().await.remove(&key); + } + Ok(()) + } + + async fn select_handler_connection(&self, thread_id: ThreadId) -> Result { + let mut connection_ids = self + .thread_state_manager + .subscribed_connection_ids(thread_id) + .await; + connection_ids.sort_by_key(|connection_id| connection_id.0); + match connection_ids.as_slice() { + [] => Err(CodexErr::InvalidRequest(format!( + "cannot delegate model request for thread {thread_id}: no subscribed client is available to handle model/request" + ))), + [connection_id] => Ok(*connection_id), + _ => Err(CodexErr::InvalidRequest(format!( + "cannot delegate model request for thread {thread_id}: multiple subscribed clients are attached and PR 1 requires exactly one model/request handler" + ))), + } + } + + async fn handle_model_request_response( + &self, + request: &DelegatedModelRequest, + response: std::result::Result, + ) -> Result<()> { + match response { + Ok(Ok(value)) => { + let response = serde_json::from_value::(value) + .map_err(CodexErr::Json)?; + if response.accepted { + Ok(()) + } else { + let reason = response.rejection_reason.unwrap_or_else(|| { + "delegated model request was rejected by the client".to_string() + }); + Err(CodexErr::InvalidRequest(reason)) + } + } + Ok(Err(err)) if is_turn_transition_server_request_error(&err) => { + Err(CodexErr::InvalidRequest(format!( + "delegated model request was cancelled during a turn transition: turn_id={}", + request.turn_id + ))) + } + Ok(Err(err)) => Err(CodexErr::InvalidRequest(format!( + "client rejected delegated model request: {}", + err.message + ))), + Err(err) => Err(CodexErr::InvalidRequest(format!( + "client closed delegated model request before responding: {err}" + ))), + } + } + + async fn remove_inflight_request(&self, key: &DelegatedRequestKey) { + self.inflight.lock().await.remove(key); + } +} + +fn stream_event_is_final(event: &serde_json::Value) -> bool { + event + .get("type") + .and_then(serde_json::Value::as_str) + .is_some_and(|event_type| { + matches!( + event_type, + "response.completed" | "response.failed" | "response.incomplete" + ) + }) +} + +#[async_trait] +impl DelegatedModelTransport for AppServerDelegatedModelTransport { + async fn start_model_request( + &self, + request: DelegatedModelRequest, + ) -> Result> { + let key = DelegatedRequestKey::from_request(&request); + let (tx, rx) = mpsc::channel(1600); + let existing = self.inflight.lock().await.insert(key.clone(), tx); + if existing.is_some() { + return Err(CodexErr::InvalidRequest(format!( + "delegated model request already exists: thread_id={} turn_id={} request_id={}", + key.thread_id, key.turn_id, key.request_id + ))); + } + + let result = async { + let connection_id = self.select_handler_connection(request.thread_id).await?; + let request_sender = ThreadScopedOutgoingMessageSender::new( + Arc::clone(&self.outgoing), + vec![connection_id], + request.thread_id, + ); + let (_, response_rx) = request_sender + .send_request(ServerRequestPayload::ModelRequest(ModelRequestParams { + thread_id: key.thread_id.clone(), + turn_id: key.turn_id.clone(), + request_id: key.request_id.clone(), + request: request.request.clone(), + })) + .await; + self.handle_model_request_response(&request, response_rx.await) + .await?; + Ok::<(), CodexErr>(()) + } + .await; + + if let Err(err) = result { + self.remove_inflight_request(&key).await; + return Err(err); + } + + Ok(rx) + } + + async fn run_model_compact_request( + &self, + request: DelegatedModelCompactRequest, + ) -> Result> { + let connection_id = self.select_handler_connection(request.thread_id).await?; + let request_sender = ThreadScopedOutgoingMessageSender::new( + Arc::clone(&self.outgoing), + vec![connection_id], + request.thread_id, + ); + let (_, response_rx) = request_sender + .send_request(ServerRequestPayload::ModelCompact(ModelCompactParams { + thread_id: request.thread_id.to_string(), + turn_id: request.turn_id.clone(), + request_id: request.request_id.clone(), + request: request.request.clone(), + })) + .await; + + match response_rx.await { + Ok(Ok(value)) => serde_json::from_value::(value) + .map(|response| response.output) + .map_err(CodexErr::Json), + Ok(Err(err)) if is_turn_transition_server_request_error(&err) => { + Err(CodexErr::InvalidRequest(format!( + "delegated model compact request was cancelled during a turn transition: turn_id={}", + request.turn_id + ))) + } + Ok(Err(err)) => Err(CodexErr::InvalidRequest(format!( + "client rejected delegated model compact request: {}", + err.message + ))), + Err(err) => Err(CodexErr::InvalidRequest(format!( + "client closed delegated model compact request before responding: {err}" + ))), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::outgoing_message::OutgoingEnvelope; + use crate::outgoing_message::OutgoingMessage; + use crate::outgoing_message::OutgoingMessageSender; + use codex_app_server_protocol::ModelCompactEnvelope; + use codex_app_server_protocol::ModelRequestEnvelope; + use codex_app_server_protocol::ServerRequest; + use codex_protocol::models::ContentItem; + use codex_protocol::models::ResponseItem; + use pretty_assertions::assert_eq; + use serde_json::json; + use std::time::Duration; + use tokio::time::timeout; + + fn test_request(thread_id: ThreadId) -> DelegatedModelRequest { + DelegatedModelRequest { + thread_id, + turn_id: "turn-1".to_string(), + request_id: "req-1".to_string(), + request: ModelRequestEnvelope { + model: "gpt-test".to_string(), + instructions: "delegate this".to_string(), + input: Vec::new(), + tools: Vec::new(), + tool_choice: "auto".to_string(), + parallel_tool_calls: false, + reasoning: None, + store: false, + stream: true, + include: Vec::new(), + service_tier: None, + prompt_cache_key: None, + text: None, + request_headers: Some(HashMap::from([( + "x-codex-turn-state".to_string(), + "sticky-route".to_string(), + )])), + }, + } + } + + fn test_compact_request(thread_id: ThreadId) -> DelegatedModelCompactRequest { + DelegatedModelCompactRequest { + thread_id, + turn_id: "turn-1".to_string(), + request_id: "compact-1".to_string(), + request: ModelCompactEnvelope { + model: "gpt-test".to_string(), + input: Vec::new(), + instructions: "compact this".to_string(), + tools: Vec::new(), + parallel_tool_calls: false, + reasoning: None, + text: None, + request_headers: Some(HashMap::from([( + "session_id".to_string(), + "thread-123".to_string(), + )])), + }, + } + } + + #[tokio::test] + async fn delegated_request_round_trips_ack_and_notifications() { + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(8); + let outgoing = Arc::new(OutgoingMessageSender::new(outgoing_tx)); + let thread_state_manager = ThreadStateManager::new(); + let transport = AppServerDelegatedModelTransport::new( + Arc::clone(&outgoing), + thread_state_manager.clone(), + ); + let thread_id = ThreadId::new(); + let connection_id = ConnectionId(7); + thread_state_manager + .connection_initialized(connection_id) + .await; + assert!( + thread_state_manager + .try_add_connection_to_thread(thread_id, connection_id) + .await + ); + + let request = test_request(thread_id); + let start_task = tokio::spawn({ + let transport = transport.clone(); + let request = request.clone(); + async move { transport.start_model_request(request).await } + }); + + let request_id = match timeout(Duration::from_secs(1), outgoing_rx.recv()) + .await + .expect("receive model request envelope") + .expect("model request should be sent") + { + OutgoingEnvelope::ToConnection { + connection_id: actual_connection_id, + message: + OutgoingMessage::Request(ServerRequest::ModelRequest { request_id, params }), + } => { + assert_eq!(actual_connection_id, connection_id); + assert_eq!(params.thread_id, thread_id.to_string()); + assert_eq!(params.turn_id, "turn-1"); + assert_eq!(params.request_id, "req-1"); + assert_eq!(params.request, request.request); + request_id + } + other => panic!("expected targeted model/request, got {other:?}"), + }; + + outgoing + .notify_client_response( + request_id, + json!({ + "accepted": true, + "rejectionReason": null, + }), + ) + .await; + + let mut event_rx = start_task + .await + .expect("request task should join") + .expect("delegated model request should be accepted"); + + let metadata = + HashMap::from([("x-codex-turn-state".to_string(), "sticky-route".to_string())]); + transport + .handle_client_notification(ClientNotification::ModelStreamMetadata( + ModelStreamMetadataNotification { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + request_id: "req-1".to_string(), + metadata: metadata.clone(), + }, + )) + .await + .expect("metadata notification should route"); + + let delta_event = json!({ + "type": "response.output_text.delta", + "delta": "hello", + }); + transport + .handle_client_notification(ClientNotification::ModelStreamEvent( + ModelStreamEventNotification { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + request_id: "req-1".to_string(), + event: delta_event.clone(), + }, + )) + .await + .expect("stream event notification should route"); + + let completed_event = json!({ + "type": "response.completed", + "response": { "id": "resp-1" }, + }); + transport + .handle_client_notification(ClientNotification::ModelStreamEvent( + ModelStreamEventNotification { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + request_id: "req-1".to_string(), + event: completed_event.clone(), + }, + )) + .await + .expect("final stream event notification should route"); + + assert_eq!( + timeout(Duration::from_secs(1), event_rx.recv()) + .await + .expect("receive metadata"), + Some(DelegatedModelEvent::StreamMetadata(metadata)) + ); + assert_eq!( + timeout(Duration::from_secs(1), event_rx.recv()) + .await + .expect("receive delta event"), + Some(DelegatedModelEvent::StreamEvent(delta_event)) + ); + assert_eq!( + timeout(Duration::from_secs(1), event_rx.recv()) + .await + .expect("receive completed event"), + Some(DelegatedModelEvent::StreamEvent(completed_event)) + ); + + let err = transport + .handle_client_notification(ClientNotification::ModelRequestFailed( + ModelRequestFailedNotification { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + request_id: "req-1".to_string(), + error: codex_app_server_protocol::ModelRequestError { + error_type: "invalid_request_error".to_string(), + code: None, + param: None, + message: "too late".to_string(), + }, + }, + )) + .await + .expect_err("completed requests should be removed from inflight routing"); + assert!( + err.to_string().contains("unknown request"), + "unexpected error: {err}" + ); + } + + #[tokio::test] + async fn delegated_request_requires_exactly_one_subscriber() { + let (outgoing_tx, _outgoing_rx) = mpsc::channel(8); + let outgoing = Arc::new(OutgoingMessageSender::new(outgoing_tx)); + let thread_state_manager = ThreadStateManager::new(); + let transport = AppServerDelegatedModelTransport::new( + Arc::clone(&outgoing), + thread_state_manager.clone(), + ); + let thread_id = ThreadId::new(); + let connection_a = ConnectionId(7); + let connection_b = ConnectionId(8); + thread_state_manager + .connection_initialized(connection_a) + .await; + thread_state_manager + .connection_initialized(connection_b) + .await; + assert!( + thread_state_manager + .try_add_connection_to_thread(thread_id, connection_a) + .await + ); + assert!( + thread_state_manager + .try_add_connection_to_thread(thread_id, connection_b) + .await + ); + + let err = transport + .start_model_request(test_request(thread_id)) + .await + .expect_err("multiple subscribers should be rejected in PR 1"); + assert!( + err.to_string() + .contains("exactly one model/request handler"), + "unexpected error: {err}" + ); + } + + #[tokio::test] + async fn delegated_compact_request_round_trips_output() { + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(8); + let outgoing = Arc::new(OutgoingMessageSender::new(outgoing_tx)); + let thread_state_manager = ThreadStateManager::new(); + let transport = AppServerDelegatedModelTransport::new( + Arc::clone(&outgoing), + thread_state_manager.clone(), + ); + let thread_id = ThreadId::new(); + let connection_id = ConnectionId(7); + thread_state_manager + .connection_initialized(connection_id) + .await; + assert!( + thread_state_manager + .try_add_connection_to_thread(thread_id, connection_id) + .await + ); + + let request = test_compact_request(thread_id); + let compact_task = tokio::spawn({ + let transport = transport.clone(); + let request = request.clone(); + async move { transport.run_model_compact_request(request).await } + }); + + let request_id = match timeout(Duration::from_secs(1), outgoing_rx.recv()) + .await + .expect("receive model compact envelope") + .expect("model compact request should be sent") + { + OutgoingEnvelope::ToConnection { + connection_id: actual_connection_id, + message: + OutgoingMessage::Request(ServerRequest::ModelCompact { request_id, params }), + } => { + assert_eq!(actual_connection_id, connection_id); + assert_eq!(params.thread_id, thread_id.to_string()); + assert_eq!(params.turn_id, "turn-1"); + assert_eq!(params.request_id, "compact-1"); + assert_eq!(params.request, request.request); + request_id + } + other => panic!("expected targeted model/compact, got {other:?}"), + }; + + let expected_output = vec![ResponseItem::Message { + id: None, + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: "compacted".to_string(), + }], + end_turn: None, + phase: None, + }]; + outgoing + .notify_client_response( + request_id, + json!({ + "output": expected_output.clone(), + }), + ) + .await; + + assert_eq!( + compact_task + .await + .expect("compact task should join") + .expect("delegated compact should succeed"), + expected_output + ); + } +} diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 8b4afc23d0..7abd2dd8c1 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -61,6 +61,7 @@ mod bespoke_event_handling; mod codex_message_processor; mod command_exec; mod config_api; +mod delegated_model; mod dynamic_tools; mod error_code; mod external_agent_config_api; diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 3804c4f9b4..457b62db31 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -8,6 +8,7 @@ use std::sync::atomic::Ordering; use crate::codex_message_processor::CodexMessageProcessor; use crate::codex_message_processor::CodexMessageProcessorArgs; use crate::config_api::ConfigApi; +use crate::delegated_model::AppServerDelegatedModelTransport; use crate::error_code::INVALID_REQUEST_ERROR_CODE; use crate::external_agent_config_api::ExternalAgentConfigApi; use crate::fs_api::FsApi; @@ -145,6 +146,7 @@ impl ExternalAuthRefresher for ExternalAuthRefreshBridge { pub(crate) struct MessageProcessor { outgoing: Arc, codex_message_processor: CodexMessageProcessor, + delegated_model_transport: AppServerDelegatedModelTransport, config_api: ConfigApi, external_agent_config_api: ExternalAgentConfigApi, fs_api: FsApi, @@ -246,6 +248,12 @@ impl MessageProcessor { thread_manager .plugins_manager() .maybe_start_curated_repo_sync_for_config(&config, auth_manager.clone()); + let delegated_model_transport = AppServerDelegatedModelTransport::new( + Arc::clone(&outgoing), + codex_message_processor.thread_state_manager(), + ); + thread_manager + .set_delegated_model_transport(Some(Arc::new(delegated_model_transport.clone()))); let config_api = ConfigApi::new( config.codex_home.clone(), cli_overrides, @@ -260,6 +268,7 @@ impl MessageProcessor { Self { outgoing, codex_message_processor, + delegated_model_transport, config_api, external_agent_config_api, fs_api, @@ -388,16 +397,23 @@ impl MessageProcessor { } pub(crate) async fn process_notification(&self, notification: JSONRPCNotification) { - // Currently, we do not expect to receive any notifications from the - // client, so we just log them. tracing::info!("<- notification: {:?}", notification); + match ClientNotification::try_from(notification) { + Ok(notification) => self.process_client_notification(notification).await, + Err(err) => tracing::warn!("failed to decode typed client notification: {err}"), + } } /// Handles typed notifications from in-process clients. pub(crate) async fn process_client_notification(&self, notification: ClientNotification) { - // Currently, we do not expect to receive any typed notifications from - // in-process clients, so we just log them. tracing::info!("<- typed notification: {:?}", notification); + if let Err(err) = self + .delegated_model_transport + .handle_client_notification(notification) + .await + { + tracing::warn!("failed to process client notification: {err}"); + } } async fn run_request_with_context( diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index ba71033c3b..9615509840 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -36,6 +36,10 @@ use crate::api_bridge::map_api_error; use crate::auth::UnauthorizedRecovery; use crate::auth_env_telemetry::AuthEnvTelemetry; use crate::auth_env_telemetry::collect_auth_env_telemetry; +use crate::delegated_model_transport::DelegatedModelCompactRequest; +use crate::delegated_model_transport::DelegatedModelEvent; +use crate::delegated_model_transport::DelegatedModelRequest; +use crate::delegated_model_transport::DelegatedModelTransport; use codex_api::CompactClient as ApiCompactClient; use codex_api::CompactionInput as ApiCompactionInput; use codex_api::MemoriesClient as ApiMemoriesClient; @@ -58,8 +62,15 @@ use codex_api::common::Reasoning; use codex_api::common::ResponsesWsRequest; use codex_api::create_text_param_for_request; use codex_api::error::ApiError; +use codex_api::rate_limits::parse_all_rate_limits; use codex_api::requests::responses::Compression; use codex_api::response_create_client_metadata; +use codex_api::sse::responses::ResponsesStreamEvent; +use codex_api::sse::responses::process_responses_event; +use codex_app_server_protocol::ModelCompactEnvelope; +use codex_app_server_protocol::ModelRequestEnvelope; +use codex_app_server_protocol::NetworkDelegationConfig; +use codex_app_server_protocol::NetworkDelegationMode; use codex_otel::SessionTelemetry; use codex_otel::current_span_w3c_trace_context; @@ -89,6 +100,7 @@ use tokio_tungstenite::tungstenite::Message; use tracing::instrument; use tracing::trace; use tracing::warn; +use uuid::Uuid; use crate::AuthManager; use crate::auth::AuthMode; @@ -129,6 +141,16 @@ pub(crate) const WEBSOCKET_CONNECT_TIMEOUT: Duration = /// /// This is intentionally kept minimal so `ModelClient` does not need to hold a full `Config`. Most /// configuration is per turn and is passed explicitly to streaming/unary methods. +struct DelegatedModelTransportSlot(StdMutex>>); + +impl std::fmt::Debug for DelegatedModelTransportSlot { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_tuple("DelegatedModelTransportSlot") + .field(&self.0.lock().map(|slot| slot.is_some()).unwrap_or(false)) + .finish() + } +} + #[derive(Debug)] struct ModelClientState { auth_manager: Option>, @@ -142,6 +164,8 @@ struct ModelClientState { beta_features_header: Option, disable_websockets: AtomicBool, cached_websocket_session: StdMutex, + delegated_model_transport: DelegatedModelTransportSlot, + network_delegation_config: StdMutex>, } /// Resolved API client setup for a single request attempt. @@ -197,6 +221,8 @@ pub struct ModelClient { pub struct ModelClientSession { client: ModelClient, websocket_session: WebsocketSession, + turn_id: String, + next_request_id: u64, /// Turn state for sticky routing. /// /// This is an `OnceLock` that stores the turn state value received from the server @@ -278,10 +304,49 @@ impl ModelClient { beta_features_header, disable_websockets: AtomicBool::new(false), cached_websocket_session: StdMutex::new(WebsocketSession::default()), + delegated_model_transport: DelegatedModelTransportSlot(StdMutex::new(None)), + network_delegation_config: StdMutex::new(None), }), } } + pub fn set_delegated_model_transport( + &self, + transport: Option>, + ) { + *self + .state + .delegated_model_transport + .0 + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = transport; + } + + pub fn set_network_delegation_config(&self, config: Option) { + *self + .state + .network_delegation_config + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = config; + } + + fn delegated_model_transport(&self) -> Option> { + self.state + .delegated_model_transport + .0 + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + } + + fn network_delegation_config(&self) -> Option { + self.state + .network_delegation_config + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + } + /// Creates a fresh turn-scoped streaming session. /// /// This constructor does not perform network I/O itself; the session opens a websocket lazily @@ -290,6 +355,8 @@ impl ModelClient { ModelClientSession { client: self.clone(), websocket_session: self.take_cached_websocket_session(), + turn_id: Uuid::new_v4().to_string(), + next_request_id: 0, turn_state: Arc::new(OnceLock::new()), } } @@ -341,6 +408,7 @@ impl ModelClient { /// session-scoped. pub async fn compact_conversation_history( &self, + turn_id: &str, prompt: &Prompt, model_info: &ModelInfo, effort: Option, @@ -350,21 +418,6 @@ impl ModelClient { if prompt.input.is_empty() { return Ok(Vec::new()); } - let client_setup = self.current_client_setup().await?; - let transport = ReqwestTransport::new(build_reqwest_client()); - let request_telemetry = Self::build_request_telemetry( - session_telemetry, - AuthRequestTelemetryContext::new( - client_setup.auth.as_ref().map(CodexAuth::auth_mode), - &client_setup.api_auth, - PendingUnauthorizedRetry::default(), - ), - RequestRouteTelemetry::for_endpoint(RESPONSES_COMPACT_ENDPOINT), - self.state.auth_env_telemetry.clone(), - ); - let client = - ApiCompactClient::new(transport, client_setup.api_provider, client_setup.api_auth) - .with_telemetry(Some(request_telemetry)); let instructions = prompt.base_instructions.text.clone(); let input = prompt.get_formatted_input(); @@ -382,6 +435,55 @@ impl ModelClient { None }; let text = create_text_param_for_request(verbosity, &prompt.output_schema); + let mut extra_headers = self.build_subagent_headers(); + extra_headers.extend(build_conversation_headers(Some( + self.state.conversation_id.to_string(), + ))); + + if self + .network_delegation_config() + .is_some_and(|config| config.mode == NetworkDelegationMode::Enabled) + { + let Some(transport) = self.delegated_model_transport() else { + return Err(CodexErr::InvalidRequest( + "network delegation is enabled for this thread, but no delegated model transport is configured".to_string(), + )); + }; + + let request = ModelCompactEnvelope { + model: model_info.slug.clone(), + input, + instructions, + tools, + parallel_tool_calls: prompt.parallel_tool_calls, + reasoning: reasoning.map(serde_json::to_value).transpose()?, + text: text.map(serde_json::to_value).transpose()?, + request_headers: Some(header_map_to_string_map(&extra_headers)), + }; + let delegated_request = DelegatedModelCompactRequest { + thread_id: self.state.conversation_id, + turn_id: turn_id.to_string(), + request_id: format!("{turn_id}-compact-{}", Uuid::new_v4()), + request, + }; + return transport.run_model_compact_request(delegated_request).await; + } + + let client_setup = self.current_client_setup().await?; + let transport = ReqwestTransport::new(build_reqwest_client()); + let request_telemetry = Self::build_request_telemetry( + session_telemetry, + AuthRequestTelemetryContext::new( + client_setup.auth.as_ref().map(CodexAuth::auth_mode), + &client_setup.api_auth, + PendingUnauthorizedRetry::default(), + ), + RequestRouteTelemetry::for_endpoint(RESPONSES_COMPACT_ENDPOINT), + self.state.auth_env_telemetry.clone(), + ); + let client = + ApiCompactClient::new(transport, client_setup.api_provider, client_setup.api_auth) + .with_telemetry(Some(request_telemetry)); let payload = ApiCompactionInput { model: &model_info.slug, input: &input, @@ -391,11 +493,6 @@ impl ModelClient { reasoning, text, }; - - let mut extra_headers = self.build_subagent_headers(); - extra_headers.extend(build_conversation_headers(Some( - self.state.conversation_id.to_string(), - ))); client .compact_input(&payload, extra_headers) .await @@ -821,6 +918,103 @@ impl ModelClientSession { }) } + fn network_delegation_config(&self) -> Option { + self.client.network_delegation_config() + } + + fn delegated_transport_enabled(&self) -> bool { + self.network_delegation_config() + .is_some_and(|config| config.mode == NetworkDelegationMode::Enabled) + } + + fn delegated_stream_idle_timeout(&self) -> Duration { + self.network_delegation_config() + .and_then(|config| config.stream_idle_timeout_ms) + .map(Duration::from_millis) + .unwrap_or_else(|| self.client.state.provider.stream_idle_timeout()) + } + + fn next_delegated_request_id(&mut self) -> String { + let request_id = format!("{}-{}", self.turn_id, self.next_request_id); + self.next_request_id = self.next_request_id.saturating_add(1); + request_id + } + + fn build_delegated_request_headers( + &self, + turn_metadata_header: Option<&str>, + ) -> HashMap { + let mut headers = self + .client + .build_websocket_headers(Some(&self.turn_state), turn_metadata_header); + headers.extend(self.client.build_subagent_headers()); + header_map_to_string_map(&headers) + } + + fn build_model_request_envelope( + &self, + request: ResponsesApiRequest, + request_headers: HashMap, + ) -> Result { + Ok(ModelRequestEnvelope { + model: request.model, + instructions: request.instructions, + input: request.input, + tools: request.tools, + tool_choice: request.tool_choice, + parallel_tool_calls: request.parallel_tool_calls, + reasoning: request.reasoning.map(serde_json::to_value).transpose()?, + store: request.store, + stream: request.stream, + include: request.include, + service_tier: request.service_tier, + prompt_cache_key: request.prompt_cache_key, + text: request.text.map(serde_json::to_value).transpose()?, + request_headers: Some(request_headers), + }) + } + + async fn stream_delegated( + &mut self, + prompt: &Prompt, + model_info: &ModelInfo, + effort: Option, + summary: ReasoningSummaryConfig, + service_tier: Option, + turn_metadata_header: Option<&str>, + ) -> Result { + let Some(transport) = self.client.delegated_model_transport() else { + return Err(CodexErr::InvalidRequest( + "network delegation is enabled for this thread, but no delegated model transport is configured".to_string(), + )); + }; + + let request = self.build_responses_request( + &self.client.state.provider.to_api_provider(None)?, + prompt, + model_info, + effort, + summary, + service_tier, + )?; + let request_headers = self.build_delegated_request_headers(turn_metadata_header); + let request_id = self.next_delegated_request_id(); + let request = self.build_model_request_envelope(request, request_headers)?; + let delegated_request = DelegatedModelRequest { + thread_id: self.client.state.conversation_id, + turn_id: self.turn_id.clone(), + request_id, + request, + }; + let idle_timeout = self.delegated_stream_idle_timeout(); + let receiver = transport.start_model_request(delegated_request).await?; + Ok(spawn_delegated_response_stream( + receiver, + idle_timeout, + Arc::clone(&self.turn_state), + )) + } + fn prepare_websocket_request( &mut self, payload: ResponseCreateWsRequest, @@ -1239,6 +1433,9 @@ impl ModelClientSession { service_tier: Option, turn_metadata_header: Option<&str>, ) -> Result<()> { + if self.delegated_transport_enabled() { + return Ok(()); + } if !self.client.responses_websocket_enabled() { return Ok(()); } @@ -1296,6 +1493,18 @@ impl ModelClientSession { service_tier: Option, turn_metadata_header: Option<&str>, ) -> Result { + if self.delegated_transport_enabled() { + return self + .stream_delegated( + prompt, + model_info, + effort, + summary, + service_tier, + turn_metadata_header, + ) + .await; + } let wire_api = self.client.state.provider.wire_api; match wire_api { WireApi::Responses => { @@ -1402,6 +1611,160 @@ fn build_responses_headers( headers } +fn header_map_to_string_map(headers: &ApiHeaderMap) -> HashMap { + headers + .iter() + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|value| (name.as_str().to_string(), value.to_string())) + }) + .collect() +} + +fn string_map_to_header_map(headers: &HashMap) -> ApiHeaderMap { + let mut header_map = ApiHeaderMap::new(); + for (name, value) in headers { + if let (Ok(header_name), Ok(header_value)) = ( + name.parse::(), + HeaderValue::from_str(value), + ) { + header_map.insert(header_name, header_value); + } + } + header_map +} + +fn metadata_events( + headers: &HashMap, + turn_state: &Arc>, +) -> Vec { + let header_map = string_map_to_header_map(headers); + if let Some(state) = header_map + .get(X_CODEX_TURN_STATE_HEADER) + .and_then(|value| value.to_str().ok()) + { + let _ = turn_state.set(state.to_string()); + } + + let mut events = Vec::new(); + let server_model = header_map + .get("openai-model") + .or_else(|| header_map.get("x-openai-model")) + .and_then(|value| value.to_str().ok()) + .map(ToString::to_string); + if let Some(server_model) = server_model { + events.push(ResponseEvent::ServerModel(server_model)); + } + for snapshot in parse_all_rate_limits(&header_map) { + events.push(ResponseEvent::RateLimits(snapshot)); + } + if let Some(etag) = header_map + .get("X-Models-Etag") + .and_then(|value| value.to_str().ok()) + { + events.push(ResponseEvent::ModelsEtag(etag.to_string())); + } + if header_map + .get("x-reasoning-included") + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| value.eq_ignore_ascii_case("true")) + { + events.push(ResponseEvent::ServerReasoningIncluded(true)); + } + events +} + +fn delegated_request_failed_error(error: codex_app_server_protocol::ModelRequestError) -> CodexErr { + let message = if let Some(code) = error.code { + format!( + "delegated model request failed: type={} code={} message={}", + error.error_type, code, error.message + ) + } else { + format!( + "delegated model request failed: type={} message={}", + error.error_type, error.message + ) + }; + CodexErr::InvalidRequest(message) +} + +fn spawn_delegated_response_stream( + mut receiver: mpsc::Receiver, + idle_timeout: Duration, + turn_state: Arc>, +) -> ResponseStream { + let (tx_event, rx_event) = mpsc::channel::>(1600); + tokio::spawn(async move { + loop { + let next_event = tokio::time::timeout(idle_timeout, receiver.recv()).await; + match next_event { + Ok(Some(DelegatedModelEvent::StreamMetadata(metadata))) => { + for event in metadata_events(&metadata, &turn_state) { + if tx_event.send(Ok(event)).await.is_err() { + return; + } + } + } + Ok(Some(DelegatedModelEvent::StreamEvent(event))) => { + let event = match serde_json::from_value::(event) { + Ok(event) => event, + Err(err) => { + let _ = tx_event.send(Err(CodexErr::Json(err))).await; + return; + } + }; + let response_event = match process_responses_event(event) { + Ok(response_event) => response_event, + Err(err) => { + let _ = tx_event + .send(Err(map_api_error(err.into_api_error()))) + .await; + return; + } + }; + if let Some(response_event) = response_event { + let completed = matches!(response_event, ResponseEvent::Completed { .. }); + if tx_event.send(Ok(response_event)).await.is_err() { + return; + } + if completed { + return; + } + } + } + Ok(Some(DelegatedModelEvent::RequestFailed(error))) => { + let _ = tx_event + .send(Err(delegated_request_failed_error(error))) + .await; + return; + } + Ok(None) => { + let _ = tx_event + .send(Err(CodexErr::Stream( + "delegated model stream closed before completion".to_string(), + None, + ))) + .await; + return; + } + Err(_) => { + let _ = tx_event + .send(Err(CodexErr::Stream( + "delegated model stream timed out waiting for events".to_string(), + None, + ))) + .await; + return; + } + } + } + }); + ResponseStream { rx_event } +} + fn map_response_stream( api_stream: S, session_telemetry: SessionTelemetry, diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 2c07b4fd1d..03dad87341 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -2,13 +2,34 @@ use super::AuthRequestTelemetryContext; use super::ModelClient; use super::PendingUnauthorizedRetry; use super::UnauthorizedRecoveryExecution; +use crate::client_common::Prompt; +use crate::delegated_model_transport::DelegatedModelCompactRequest; +use crate::delegated_model_transport::DelegatedModelEvent; +use crate::delegated_model_transport::DelegatedModelRequest; +use crate::delegated_model_transport::DelegatedModelTransport; +use async_trait::async_trait; +use codex_api::ResponseEvent; +use codex_app_server_protocol::NetworkDelegationConfig; +use codex_app_server_protocol::NetworkDelegationMode; use codex_otel::SessionTelemetry; use codex_protocol::ThreadId; +use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig; +use codex_protocol::models::BaseInstructions; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ModelInfo; +use codex_protocol::protocol::CreditsSnapshot; +use codex_protocol::protocol::RateLimitSnapshot; +use codex_protocol::protocol::RateLimitWindow; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use pretty_assertions::assert_eq; use serde_json::json; +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::OnceLock; +use tokio::sync::mpsc; fn test_model_client(session_source: SessionSource) -> ModelClient { let provider = crate::model_provider_info::create_oss_provider_with_base_url( @@ -72,6 +93,40 @@ fn test_session_telemetry() -> SessionTelemetry { ) } +#[derive(Default)] +struct RecordingDelegatedTransport { + compact_requests: Mutex>, +} + +#[async_trait] +impl DelegatedModelTransport for RecordingDelegatedTransport { + async fn start_model_request( + &self, + _request: DelegatedModelRequest, + ) -> crate::error::Result> { + panic!("streaming delegation is not used in this test"); + } + + async fn run_model_compact_request( + &self, + request: DelegatedModelCompactRequest, + ) -> crate::error::Result> { + self.compact_requests + .lock() + .expect("compact request recording mutex") + .push(request); + Ok(vec![ResponseItem::Message { + id: Some("msg_compact".to_string()), + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: "compacted output".to_string(), + }], + end_turn: None, + phase: None, + }]) + } +} + #[test] fn build_subagent_headers_sets_other_subagent_label() { let client = test_model_client(SessionSource::SubAgent(SubAgentSource::Other( @@ -115,3 +170,152 @@ fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() { assert_eq!(auth_context.recovery_mode, Some("managed")); assert_eq!(auth_context.recovery_phase, Some("refresh_token")); } + +#[test] +fn metadata_events_sets_turn_state_and_emits_header_derived_events() { + let turn_state = Arc::new(OnceLock::new()); + let headers = HashMap::from([ + ("x-codex-turn-state".to_string(), "sticky-route".to_string()), + ("openai-model".to_string(), "gpt-delegated".to_string()), + ( + "x-codex-primary-used-percent".to_string(), + "42.0".to_string(), + ), + ( + "x-codex-primary-window-minutes".to_string(), + "15".to_string(), + ), + ("x-codex-primary-reset-at".to_string(), "1700".to_string()), + ("x-codex-limit-name".to_string(), "Codex".to_string()), + ( + "x-codex-credits-has-credits".to_string(), + "true".to_string(), + ), + ("x-codex-credits-unlimited".to_string(), "false".to_string()), + ("x-codex-credits-balance".to_string(), "10.50".to_string()), + ("X-Models-Etag".to_string(), "etag-123".to_string()), + ("x-reasoning-included".to_string(), "true".to_string()), + ]); + + let events = super::metadata_events(&headers, &turn_state); + + assert_eq!(turn_state.get().map(String::as_str), Some("sticky-route")); + assert_eq!(events.len(), 4); + match &events[0] { + ResponseEvent::ServerModel(model) => assert_eq!(model, "gpt-delegated"), + other => panic!("expected server model event, got {other:?}"), + } + match &events[1] { + ResponseEvent::RateLimits(snapshot) => assert_eq!( + snapshot, + &RateLimitSnapshot { + limit_id: Some("codex".to_string()), + limit_name: Some("Codex".to_string()), + primary: Some(RateLimitWindow { + used_percent: 42.0, + window_minutes: Some(15), + resets_at: Some(1700), + }), + secondary: None, + credits: Some(CreditsSnapshot { + has_credits: true, + unlimited: false, + balance: Some("10.50".to_string()), + }), + plan_type: None, + } + ), + other => panic!("expected rate limits event, got {other:?}"), + } + match &events[2] { + ResponseEvent::ModelsEtag(etag) => assert_eq!(etag, "etag-123"), + other => panic!("expected models etag event, got {other:?}"), + } + match &events[3] { + ResponseEvent::ServerReasoningIncluded(included) => assert_eq!(included, &true), + other => panic!("expected reasoning included event, got {other:?}"), + } +} + +#[tokio::test] +async fn delegated_compact_conversation_history_routes_through_transport() { + let client = test_model_client(SessionSource::SubAgent(SubAgentSource::Compact)); + let transport = Arc::new(RecordingDelegatedTransport::default()); + client.set_delegated_model_transport(Some(transport.clone())); + client.set_network_delegation_config(Some(NetworkDelegationConfig { + mode: NetworkDelegationMode::Enabled, + stream_idle_timeout_ms: None, + })); + + let prompt = Prompt { + input: vec![ResponseItem::Message { + id: Some("msg_user".to_string()), + role: "user".to_string(), + content: vec![ContentItem::InputText { + text: "please compact".to_string(), + }], + end_turn: None, + phase: None, + }], + tools: Vec::new(), + parallel_tool_calls: false, + base_instructions: BaseInstructions { + text: "compact instructions".to_string(), + }, + personality: None, + output_schema: None, + }; + + let output = client + .compact_conversation_history( + "turn-compact", + &prompt, + &test_model_info(), + None, + ReasoningSummaryConfig::None, + &test_session_telemetry(), + ) + .await + .expect("delegated compact request should succeed"); + + assert_eq!( + output, + vec![ResponseItem::Message { + id: Some("msg_compact".to_string()), + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: "compacted output".to_string(), + }], + end_turn: None, + phase: None, + }] + ); + + let recorded = transport + .compact_requests + .lock() + .expect("compact request recording mutex") + .clone(); + assert_eq!(recorded.len(), 1); + let request = &recorded[0]; + assert_eq!(request.thread_id, client.state.conversation_id); + assert_eq!(request.turn_id, "turn-compact"); + assert!(request.request_id.starts_with("turn-compact-compact-")); + assert_eq!(request.request.model, "gpt-test"); + assert_eq!(request.request.instructions, "compact instructions"); + assert_eq!(request.request.input, prompt.input); + assert_eq!(request.request.tools, Vec::::new()); + assert!(!request.request.parallel_tool_calls); + assert_eq!(request.request.reasoning, None); + assert_eq!(request.request.text, None); + assert_eq!( + request.request.request_headers, + Some(HashMap::from([ + ( + "session_id".to_string(), + client.state.conversation_id.to_string(), + ), + ("x-openai-subagent".to_string(), "compact".to_string()), + ])) + ); +} diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index b90d702b2e..85402f5f66 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -26,6 +26,7 @@ use crate::compact::should_use_remote_compact_task; use crate::compact_remote::run_inline_remote_auto_compact_task; use crate::config::ManagedFeatures; use crate::connectors; +use crate::delegated_model_transport::DelegatedModelTransport; use crate::exec_policy::ExecPolicyManager; use crate::features::FEATURES; use crate::features::Feature; @@ -59,6 +60,7 @@ use chrono::Local; use chrono::Utc; use codex_app_server_protocol::McpServerElicitationRequest; use codex_app_server_protocol::McpServerElicitationRequestParams; +use codex_app_server_protocol::NetworkDelegationConfig; use codex_exec_server::Environment; use codex_hooks::HookEvent; use codex_hooks::HookEventAfterAgent; @@ -587,6 +589,7 @@ impl Codex { original_config_do_not_use: Arc::clone(&config), metrics_service_name, app_server_client_name: None, + network_delegation: None, session_source, dynamic_tools, persist_extended_history, @@ -717,6 +720,35 @@ impl Codex { .await } + pub(crate) async fn set_network_delegation( + &self, + network_delegation: Option, + ) -> ConstraintResult<()> { + let Some(network_delegation) = network_delegation else { + return Ok(()); + }; + self.session + .services + .model_client + .set_network_delegation_config(Some(network_delegation.clone())); + self.session + .update_settings(SessionSettingsUpdate { + network_delegation: Some(network_delegation), + ..Default::default() + }) + .await + } + + pub(crate) fn set_delegated_model_transport( + &self, + delegated_model_transport: Option>, + ) { + self.session + .services + .model_client + .set_delegated_model_transport(delegated_model_transport); + } + pub(crate) async fn agent_status(&self) -> AgentStatus { self.agent_status.borrow().clone() } @@ -1051,6 +1083,7 @@ pub(crate) struct SessionConfiguration { /// Optional service name tag for session metrics. metrics_service_name: Option, app_server_client_name: Option, + network_delegation: Option, /// Source of the session (cli, vscode, exec, mcp, ...) session_source: SessionSource, dynamic_tools: Vec, @@ -1069,6 +1102,7 @@ impl SessionConfiguration { model: self.collaboration_mode.model().to_string(), model_provider_id: self.original_config_do_not_use.model_provider_id.clone(), service_tier: self.service_tier, + network_delegation: self.network_delegation.clone(), approval_policy: self.approval_policy.value(), approvals_reviewer: self.approvals_reviewer, sandbox_policy: self.sandbox_policy.get().clone(), @@ -1132,6 +1166,9 @@ impl SessionConfiguration { if let Some(app_server_client_name) = updates.app_server_client_name.clone() { next_configuration.app_server_client_name = Some(app_server_client_name); } + if let Some(network_delegation) = updates.network_delegation.clone() { + next_configuration.network_delegation = Some(network_delegation); + } Ok(next_configuration) } } @@ -1149,6 +1186,7 @@ pub(crate) struct SessionSettingsUpdate { pub(crate) final_output_json_schema: Option>, pub(crate) personality: Option, pub(crate) app_server_client_name: Option, + pub(crate) network_delegation: Option, } impl Session { @@ -4531,6 +4569,7 @@ mod handlers { final_output_json_schema: Some(final_output_json_schema), personality, app_server_client_name: None, + network_delegation: None, }, ) } diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 787ad399b1..51a151e150 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -1655,6 +1655,7 @@ async fn set_rate_limits_retains_previous_credits() { original_config_do_not_use: Arc::clone(&config), metrics_service_name: None, app_server_client_name: None, + network_delegation: None, session_source: SessionSource::Exec, dynamic_tools: Vec::new(), persist_extended_history: false, @@ -1753,6 +1754,7 @@ async fn set_rate_limits_updates_plan_type_when_present() { original_config_do_not_use: Arc::clone(&config), metrics_service_name: None, app_server_client_name: None, + network_delegation: None, session_source: SessionSource::Exec, dynamic_tools: Vec::new(), persist_extended_history: false, @@ -2097,6 +2099,7 @@ pub(crate) async fn make_session_configuration_for_tests() -> SessionConfigurati original_config_do_not_use: Arc::clone(&config), metrics_service_name: None, app_server_client_name: None, + network_delegation: None, session_source: SessionSource::Exec, dynamic_tools: Vec::new(), persist_extended_history: false, @@ -2328,6 +2331,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_zsh_path() { original_config_do_not_use: Arc::clone(&config), metrics_service_name: None, app_server_client_name: None, + network_delegation: None, session_source: SessionSource::Exec, dynamic_tools: Vec::new(), persist_extended_history: false, @@ -2423,6 +2427,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { original_config_do_not_use: Arc::clone(&config), metrics_service_name: None, app_server_client_name: None, + network_delegation: None, session_source: SessionSource::Exec, dynamic_tools: Vec::new(), persist_extended_history: false, @@ -3217,6 +3222,7 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( original_config_do_not_use: Arc::clone(&config), metrics_service_name: None, app_server_client_name: None, + network_delegation: None, session_source: SessionSource::Exec, dynamic_tools, persist_extended_history: false, diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 2bd9608b95..43cfad5c98 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -2,6 +2,7 @@ use crate::agent::AgentStatus; use crate::codex::Codex; use crate::codex::SteerInputError; use crate::config::ConstraintResult; +use crate::delegated_model_transport::DelegatedModelTransport; use crate::error::CodexErr; use crate::error::Result as CodexResult; use crate::features::Feature; @@ -23,16 +24,19 @@ use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::W3cTraceContext; use codex_protocol::user_input::UserInput; use std::path::PathBuf; +use std::sync::Arc; use tokio::sync::Mutex; use tokio::sync::watch; use crate::state_db::StateDbHandle; +use codex_app_server_protocol::NetworkDelegationConfig; #[derive(Clone, Debug)] pub struct ThreadConfigSnapshot { pub model: String, pub model_provider_id: String, pub service_tier: Option, + pub network_delegation: Option, pub approval_policy: AskForApproval, pub approvals_reviewer: ApprovalsReviewer, pub sandbox_policy: SandboxPolicy, @@ -99,6 +103,21 @@ impl CodexThread { .await } + pub async fn set_network_delegation( + &self, + network_delegation: Option, + ) -> ConstraintResult<()> { + self.codex.set_network_delegation(network_delegation).await + } + + pub fn set_delegated_model_transport( + &self, + delegated_model_transport: Option>, + ) { + self.codex + .set_delegated_model_transport(delegated_model_transport); + } + /// Use sparingly: this is intended to be removed soon. pub async fn submit_with_id(&self, sub: Submission) -> CodexResult<()> { self.codex.submit_with_id(sub).await diff --git a/codex-rs/core/src/compact_remote.rs b/codex-rs/core/src/compact_remote.rs index 9439d8125e..e728f0a5ad 100644 --- a/codex-rs/core/src/compact_remote.rs +++ b/codex-rs/core/src/compact_remote.rs @@ -118,6 +118,7 @@ async fn run_remote_compact_task_inner_impl( .services .model_client .compact_conversation_history( + &turn_context.sub_id, &prompt, &turn_context.model_info, turn_context.reasoning_effort, diff --git a/codex-rs/core/src/delegated_model_transport.rs b/codex-rs/core/src/delegated_model_transport.rs new file mode 100644 index 0000000000..f8dfa9f31a --- /dev/null +++ b/codex-rs/core/src/delegated_model_transport.rs @@ -0,0 +1,46 @@ +use crate::error::Result; +use async_trait::async_trait; +use codex_app_server_protocol::ModelCompactEnvelope; +use codex_app_server_protocol::ModelRequestEnvelope; +use codex_app_server_protocol::ModelRequestError; +use codex_protocol::ThreadId; +use codex_protocol::models::ResponseItem; +use serde_json::Value; +use std::collections::HashMap; +use tokio::sync::mpsc; + +#[derive(Debug, Clone, PartialEq)] +pub struct DelegatedModelRequest { + pub thread_id: ThreadId, + pub turn_id: String, + pub request_id: String, + pub request: ModelRequestEnvelope, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct DelegatedModelCompactRequest { + pub thread_id: ThreadId, + pub turn_id: String, + pub request_id: String, + pub request: ModelCompactEnvelope, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum DelegatedModelEvent { + StreamMetadata(HashMap), + StreamEvent(Value), + RequestFailed(ModelRequestError), +} + +#[async_trait] +pub trait DelegatedModelTransport: Send + Sync { + async fn start_model_request( + &self, + request: DelegatedModelRequest, + ) -> Result>; + + async fn run_model_compact_request( + &self, + request: DelegatedModelCompactRequest, + ) -> Result>; +} diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 10a51b23ec..8cd0a70d6a 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -32,6 +32,7 @@ pub mod connectors; mod context_manager; mod contextual_user_message; pub mod custom_prompts; +pub mod delegated_model_transport; pub mod env; mod environment_context; pub mod error; diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index a63cf2cb94..5ba7e407f1 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -9,6 +9,7 @@ use crate::codex::CodexSpawnOk; use crate::codex::INITIAL_SUBMIT_ID; use crate::codex_thread::CodexThread; use crate::config::Config; +use crate::delegated_model_transport::DelegatedModelTransport; use crate::error::CodexErr; use crate::error::Result as CodexResult; use crate::file_watcher::FileWatcher; @@ -38,6 +39,7 @@ use futures::stream::FuturesUnordered; use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; +use std::sync::Mutex as StdMutex; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Duration; @@ -157,6 +159,7 @@ pub(crate) struct ThreadManagerState { mcp_manager: Arc, file_watcher: Arc, session_source: SessionSource, + delegated_model_transport: StdMutex>>, // Captures submitted ops for testing purpose when test mode is enabled. ops_log: Option, } @@ -205,6 +208,7 @@ impl ThreadManager { file_watcher, auth_manager, session_source, + delegated_model_transport: StdMutex::new(None), ops_log: should_use_test_thread_manager_behavior() .then(|| Arc::new(std::sync::Mutex::new(Vec::new()))), }), @@ -269,6 +273,7 @@ impl ThreadManager { file_watcher, auth_manager, session_source: SessionSource::Exec, + delegated_model_transport: StdMutex::new(None), ops_log: should_use_test_thread_manager_behavior() .then(|| Arc::new(std::sync::Mutex::new(Vec::new()))), }), @@ -304,6 +309,17 @@ impl ThreadManager { self.state.models_manager.clone() } + pub fn set_delegated_model_transport( + &self, + transport: Option>, + ) { + *self + .state + .delegated_model_transport + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = transport; + } + pub async fn list_models( &self, refresh_strategy: crate::models_manager::manager::RefreshStrategy, @@ -762,6 +778,11 @@ impl ThreadManagerState { let watch_registration = self .file_watcher .register_config(&config, self.skills_manager.as_ref()); + let delegated_model_transport = self + .delegated_model_transport + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); let CodexSpawnOk { codex, thread_id, .. } = Codex::spawn(CodexSpawnArgs { @@ -784,6 +805,7 @@ impl ThreadManagerState { parent_trace, }) .await?; + codex.set_delegated_model_transport(delegated_model_transport); self.finalize_thread_spawn(codex, thread_id, watch_registration) .await }