diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 3df52caa38..48a5f23147 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -3341,9 +3341,13 @@ impl CodexMessageProcessor { self.thread_manager.subscribe_thread_created() } - pub(crate) async fn connection_initialized(&self, connection_id: ConnectionId) { + pub(crate) async fn connection_initialized( + &self, + connection_id: ConnectionId, + experimental_api_enabled: bool, + ) { self.thread_state_manager - .connection_initialized(connection_id) + .connection_initialized(connection_id, experimental_api_enabled) .await; } @@ -8936,7 +8940,9 @@ mod tests { let connection = ConnectionId(1); let (cancel_tx, cancel_rx) = oneshot::channel(); - manager.connection_initialized(connection).await; + manager + .connection_initialized(connection, /*experimental_api_enabled*/ false) + .await; manager .try_ensure_connection_subscribed(thread_id, connection, false) .await @@ -8979,8 +8985,12 @@ mod tests { let connection_b = ConnectionId(2); let (cancel_tx, mut cancel_rx) = oneshot::channel(); - manager.connection_initialized(connection_a).await; - manager.connection_initialized(connection_b).await; + manager + .connection_initialized(connection_a, /*experimental_api_enabled*/ false) + .await; + manager + .connection_initialized(connection_b, /*experimental_api_enabled*/ false) + .await; manager .try_ensure_connection_subscribed(thread_id, connection_a, false) .await @@ -9014,7 +9024,9 @@ mod tests { let thread_id = ThreadId::from_string("ad7f0408-99b8-4f6e-a46f-bd0eec433370")?; let connection = ConnectionId(1); - manager.connection_initialized(connection).await; + manager + .connection_initialized(connection, /*experimental_api_enabled*/ false) + .await; manager.remove_connection(connection).await; assert!( diff --git a/codex-rs/app-server/src/delegated_model.rs b/codex-rs/app-server/src/delegated_model.rs index 9a85630bcb..90da87cb6c 100644 --- a/codex-rs/app-server/src/delegated_model.rs +++ b/codex-rs/app-server/src/delegated_model.rs @@ -67,11 +67,17 @@ impl DelegatedRequestKey { } } +#[derive(Clone)] +struct InflightDelegatedRequest { + connection_id: ConnectionId, + sender: mpsc::Sender, +} + #[derive(Clone)] pub(crate) struct AppServerDelegatedModelTransport { outgoing: Arc, thread_state_manager: ThreadStateManager, - inflight: Arc>>>, + inflight: Arc>>, } impl AppServerDelegatedModelTransport { @@ -88,25 +94,32 @@ impl AppServerDelegatedModelTransport { pub(crate) async fn handle_client_notification( &self, + connection_id: ConnectionId, notification: ClientNotification, ) -> Result<()> { match notification { ClientNotification::Initialized => Ok(()), ClientNotification::ModelStreamMetadata(notification) => { - self.deliver_metadata(notification).await + self.deliver_metadata(connection_id, notification).await } ClientNotification::ModelStreamEvent(notification) => { - self.deliver_stream_event(notification).await + self.deliver_stream_event(connection_id, notification).await } ClientNotification::ModelRequestFailed(notification) => { - self.deliver_request_failed(notification).await + self.deliver_request_failed(connection_id, notification) + .await } } } - async fn deliver_metadata(&self, notification: ModelStreamMetadataNotification) -> Result<()> { + async fn deliver_metadata( + &self, + connection_id: ConnectionId, + notification: ModelStreamMetadataNotification, + ) -> Result<()> { let key = DelegatedRequestKey::from_metadata(¬ification); self.deliver_event( + connection_id, key, DelegatedModelEvent::StreamMetadata(notification.metadata), /*final_event*/ false, @@ -114,10 +127,15 @@ impl AppServerDelegatedModelTransport { .await } - async fn deliver_stream_event(&self, notification: ModelStreamEventNotification) -> Result<()> { + async fn deliver_stream_event( + &self, + connection_id: ConnectionId, + notification: ModelStreamEventNotification, + ) -> Result<()> { let final_event = stream_event_is_final(¬ification.event); let key = DelegatedRequestKey::from_stream_event(¬ification); self.deliver_event( + connection_id, key, DelegatedModelEvent::StreamEvent(notification.event), final_event, @@ -127,10 +145,12 @@ impl AppServerDelegatedModelTransport { async fn deliver_request_failed( &self, + connection_id: ConnectionId, notification: ModelRequestFailedNotification, ) -> Result<()> { let key = DelegatedRequestKey::from_failed(¬ification); self.deliver_event( + connection_id, key, DelegatedModelEvent::RequestFailed(notification.error), /*final_event*/ true, @@ -140,11 +160,12 @@ impl AppServerDelegatedModelTransport { async fn deliver_event( &self, + connection_id: ConnectionId, key: DelegatedRequestKey, event: DelegatedModelEvent, final_event: bool, ) -> Result<()> { - let sender = { + let inflight = { let inflight = self.inflight.lock().await; inflight.get(&key).cloned() } @@ -155,7 +176,14 @@ impl AppServerDelegatedModelTransport { )) })?; - if sender.send(event).await.is_err() { + if inflight.connection_id != connection_id { + return Err(CodexErr::InvalidRequest(format!( + "received delegated model notification from unexpected connection: expected={:?} actual={:?} thread_id={} turn_id={} request_id={}", + inflight.connection_id, connection_id, key.thread_id, key.turn_id, key.request_id + ))); + } + + if inflight.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={}", @@ -171,16 +199,16 @@ impl AppServerDelegatedModelTransport { async fn select_handler_connection(&self, thread_id: ThreadId) -> Result { let mut connection_ids = self .thread_state_manager - .subscribed_connection_ids(thread_id) + .subscribed_connection_ids_with_experimental_api(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" + "cannot delegate model request for thread {thread_id}: no subscribed experimental-api 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" + "cannot delegate model request for thread {thread_id}: multiple subscribed experimental-api clients are attached and PR 1 requires exactly one model/request handler" ))), } } @@ -236,24 +264,28 @@ fn stream_event_is_final(event: &serde_json::Value) -> bool { }) } -#[async_trait] -impl DelegatedModelTransport for AppServerDelegatedModelTransport { - async fn start_model_request( +impl AppServerDelegatedModelTransport { + async fn start_model_request_impl( &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 (tx, rx) = mpsc::channel(1600); + let existing = self.inflight.lock().await.insert( + key.clone(), + InflightDelegatedRequest { + connection_id, + sender: 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 request_sender = ThreadScopedOutgoingMessageSender::new( Arc::clone(&self.outgoing), vec![connection_id], @@ -269,19 +301,22 @@ impl DelegatedModelTransport for AppServerDelegatedModelTransport { .await; self.handle_model_request_response(&request, response_rx.await) .await?; - Ok::<(), CodexErr>(()) + Ok::, CodexErr>(rx) } .await; - if let Err(err) = result { - self.remove_inflight_request(&key).await; - return Err(err); - } + let receiver = match result { + Ok(receiver) => receiver, + Err(err) => { + self.remove_inflight_request(&key).await; + return Err(err); + } + }; - Ok(rx) + Ok(receiver) } - async fn run_model_compact_request( + async fn run_model_compact_request_impl( &self, request: DelegatedModelCompactRequest, ) -> Result> { @@ -321,6 +356,23 @@ impl DelegatedModelTransport for AppServerDelegatedModelTransport { } } +#[async_trait] +impl DelegatedModelTransport for AppServerDelegatedModelTransport { + async fn start_model_request( + &self, + request: DelegatedModelRequest, + ) -> Result> { + self.start_model_request_impl(request).await + } + + async fn run_model_compact_request( + &self, + request: DelegatedModelCompactRequest, + ) -> Result> { + self.run_model_compact_request_impl(request).await + } +} + #[cfg(test)] mod tests { use super::*; @@ -397,7 +449,7 @@ mod tests { let thread_id = ThreadId::new(); let connection_id = ConnectionId(7); thread_state_manager - .connection_initialized(connection_id) + .connection_initialized(connection_id, /*experimental_api_enabled*/ true) .await; assert!( thread_state_manager @@ -450,14 +502,15 @@ mod tests { let metadata = HashMap::from([("x-codex-turn-state".to_string(), "sticky-route".to_string())]); transport - .handle_client_notification(ClientNotification::ModelStreamMetadata( - ModelStreamMetadataNotification { + .handle_client_notification( + connection_id, + 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"); @@ -466,14 +519,15 @@ mod tests { "delta": "hello", }); transport - .handle_client_notification(ClientNotification::ModelStreamEvent( - ModelStreamEventNotification { + .handle_client_notification( + connection_id, + 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"); @@ -482,14 +536,15 @@ mod tests { "response": { "id": "resp-1" }, }); transport - .handle_client_notification(ClientNotification::ModelStreamEvent( - ModelStreamEventNotification { + .handle_client_notification( + connection_id, + 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"); @@ -513,8 +568,9 @@ mod tests { ); let err = transport - .handle_client_notification(ClientNotification::ModelRequestFailed( - ModelRequestFailedNotification { + .handle_client_notification( + connection_id, + ClientNotification::ModelRequestFailed(ModelRequestFailedNotification { thread_id: thread_id.to_string(), turn_id: "turn-1".to_string(), request_id: "req-1".to_string(), @@ -524,8 +580,8 @@ mod tests { param: None, message: "too late".to_string(), }, - }, - )) + }), + ) .await .expect_err("completed requests should be removed from inflight routing"); assert!( @@ -547,10 +603,10 @@ mod tests { let connection_a = ConnectionId(7); let connection_b = ConnectionId(8); thread_state_manager - .connection_initialized(connection_a) + .connection_initialized(connection_a, /*experimental_api_enabled*/ true) .await; thread_state_manager - .connection_initialized(connection_b) + .connection_initialized(connection_b, /*experimental_api_enabled*/ true) .await; assert!( thread_state_manager @@ -586,7 +642,7 @@ mod tests { let thread_id = ThreadId::new(); let connection_id = ConnectionId(7); thread_state_manager - .connection_initialized(connection_id) + .connection_initialized(connection_id, /*experimental_api_enabled*/ true) .await; assert!( thread_state_manager @@ -647,4 +703,122 @@ mod tests { expected_output ); } + + #[tokio::test] + async fn delegated_request_requires_experimental_api_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_id = ConnectionId(7); + thread_state_manager + .connection_initialized(connection_id, /*experimental_api_enabled*/ false) + .await; + assert!( + thread_state_manager + .try_add_connection_to_thread(thread_id, connection_id) + .await + ); + + let err = transport + .start_model_request(test_request(thread_id)) + .await + .expect_err("non-experimental subscribers should be rejected"); + assert!( + err.to_string().contains("experimental-api client"), + "unexpected error: {err}" + ); + } + + #[tokio::test] + async fn delegated_request_rejects_notifications_from_unexpected_connection() { + 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 handler_connection = ConnectionId(7); + let other_connection = ConnectionId(8); + thread_state_manager + .connection_initialized(handler_connection, /*experimental_api_enabled*/ true) + .await; + thread_state_manager + .connection_initialized(other_connection, /*experimental_api_enabled*/ false) + .await; + assert!( + thread_state_manager + .try_add_connection_to_thread(thread_id, handler_connection) + .await + ); + assert!( + thread_state_manager + .try_add_connection_to_thread(thread_id, other_connection) + .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, + message: OutgoingMessage::Request(ServerRequest::ModelRequest { request_id, .. }), + } => { + assert_eq!(connection_id, handler_connection); + request_id + } + other => panic!("expected targeted model/request, got {other:?}"), + }; + + outgoing + .notify_client_response( + request_id, + json!({ + "accepted": true, + "rejectionReason": null, + }), + ) + .await; + + let event_rx = start_task + .await + .expect("request task should join") + .expect("delegated model request should be accepted"); + drop(event_rx); + + let err = transport + .handle_client_notification( + other_connection, + ClientNotification::ModelStreamEvent(ModelStreamEventNotification { + thread_id: thread_id.to_string(), + turn_id: "turn-1".to_string(), + request_id: "req-1".to_string(), + event: json!({ + "type": "response.completed", + "response": { "id": "resp-1" }, + }), + }), + ) + .await + .expect_err("mismatched connection should be rejected"); + assert!( + err.to_string().contains("unexpected connection"), + "unexpected error: {err}" + ); + } } diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index 4288d15393..aa4d663ed3 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -453,7 +453,12 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle { } } Some(ProcessorCommand::Notification(notification)) => { - processor.process_client_notification(notification).await; + processor + .process_client_notification( + IN_PROCESS_CONNECTION_ID, + notification, + ) + .await; } None => { break; diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 7abd2dd8c1..76625ba494 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -769,7 +769,14 @@ pub async fn run_main_with_transport( connection_id, ) .await; - processor.connection_initialized(connection_id).await; + processor + .connection_initialized( + connection_id, + connection_state + .session + .experimental_api_enabled, + ) + .await; connection_state .outbound_initialized .store(true, std::sync::atomic::Ordering::Release); @@ -787,7 +794,9 @@ pub async fn run_main_with_transport( warn!("dropping notification from unknown connection: {connection_id:?}"); continue; } - processor.process_notification(notification).await; + processor + .process_notification(connection_id, notification) + .await; } JSONRPCMessage::Error(err) => { if !connections.contains_key(&connection_id) { diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 457b62db31..18298d1a44 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -396,20 +396,31 @@ impl MessageProcessor { .await; } - pub(crate) async fn process_notification(&self, notification: JSONRPCNotification) { + pub(crate) async fn process_notification( + &self, + connection_id: ConnectionId, + notification: JSONRPCNotification, + ) { tracing::info!("<- notification: {:?}", notification); match ClientNotification::try_from(notification) { - Ok(notification) => self.process_client_notification(notification).await, + Ok(notification) => { + self.process_client_notification(connection_id, 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) { + pub(crate) async fn process_client_notification( + &self, + connection_id: ConnectionId, + notification: ClientNotification, + ) { tracing::info!("<- typed notification: {:?}", notification); if let Err(err) = self .delegated_model_transport - .handle_client_notification(notification) + .handle_client_notification(connection_id, notification) .await { tracing::warn!("failed to process client notification: {err}"); @@ -447,9 +458,13 @@ impl MessageProcessor { } } - pub(crate) async fn connection_initialized(&self, connection_id: ConnectionId) { + pub(crate) async fn connection_initialized( + &self, + connection_id: ConnectionId, + experimental_api_enabled: bool, + ) { self.codex_message_processor - .connection_initialized(connection_id) + .connection_initialized(connection_id, experimental_api_enabled) .await; } @@ -612,7 +627,7 @@ impl MessageProcessor { // initialize handling for the specific connection. outbound_initialized.store(true, Ordering::Release); self.codex_message_processor - .connection_initialized(connection_id) + .connection_initialized(connection_id, session.experimental_api_enabled) .await; } return; diff --git a/codex-rs/app-server/src/thread_state.rs b/codex-rs/app-server/src/thread_state.rs index be5478dd51..ca4dc2d05e 100644 --- a/codex-rs/app-server/src/thread_state.rs +++ b/codex-rs/app-server/src/thread_state.rs @@ -133,6 +133,7 @@ impl Default for ThreadEntry { #[derive(Default)] struct ThreadStateManagerInner { live_connections: HashSet, + experimental_api_enabled_by_connection: HashMap, threads: HashMap, thread_ids_by_connection: HashMap>, } @@ -147,12 +148,16 @@ impl ThreadStateManager { Self::default() } - pub(crate) async fn connection_initialized(&self, connection_id: ConnectionId) { - self.state - .lock() - .await - .live_connections - .insert(connection_id); + pub(crate) async fn connection_initialized( + &self, + connection_id: ConnectionId, + experimental_api_enabled: bool, + ) { + let mut state = self.state.lock().await; + state.live_connections.insert(connection_id); + state + .experimental_api_enabled_by_connection + .insert(connection_id, experimental_api_enabled); } pub(crate) async fn subscribed_connection_ids(&self, thread_id: ThreadId) -> Vec { @@ -164,6 +169,31 @@ impl ThreadStateManager { .unwrap_or_default() } + pub(crate) async fn subscribed_connection_ids_with_experimental_api( + &self, + thread_id: ThreadId, + ) -> Vec { + let state = self.state.lock().await; + state + .threads + .get(&thread_id) + .map(|thread_entry| { + thread_entry + .connection_ids + .iter() + .copied() + .filter(|connection_id| { + state + .experimental_api_enabled_by_connection + .get(connection_id) + .copied() + .unwrap_or(false) + }) + .collect() + }) + .unwrap_or_default() + } + pub(crate) async fn thread_state(&self, thread_id: ThreadId) -> Arc> { let mut state = self.state.lock().await; state.threads.entry(thread_id).or_default().state.clone() @@ -317,6 +347,9 @@ impl ThreadStateManager { let thread_states = { let mut state = self.state.lock().await; state.live_connections.remove(&connection_id); + state + .experimental_api_enabled_by_connection + .remove(&connection_id); let thread_ids = state .thread_ids_by_connection .remove(&connection_id) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index b294b869f9..1ba2d3adb8 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -352,10 +352,14 @@ impl ModelClient { /// This constructor does not perform network I/O itself; the session opens a websocket lazily /// when the first stream request is issued. pub fn new_session(&self) -> ModelClientSession { + self.new_session_with_turn_id(Uuid::new_v4().to_string()) + } + + pub fn new_session_with_turn_id(&self, turn_id: String) -> ModelClientSession { ModelClientSession { client: self.clone(), websocket_session: self.take_cached_websocket_session(), - turn_id: Uuid::new_v4().to_string(), + turn_id, next_request_id: 0, turn_state: Arc::new(OnceLock::new()), } @@ -768,6 +772,11 @@ impl Drop for ModelClientSession { } impl ModelClientSession { + pub(crate) fn set_turn_id(&mut self, turn_id: String) { + self.turn_id = turn_id; + self.next_request_id = 0; + } + fn reset_websocket_session(&mut self) { self.websocket_session.connection = None; self.websocket_session.last_request = None; @@ -944,9 +953,17 @@ impl ModelClientSession { &self, turn_metadata_header: Option<&str>, ) -> HashMap { - let mut headers = self - .client - .build_websocket_headers(Some(&self.turn_state), turn_metadata_header); + let turn_metadata_header = parse_turn_metadata_header(turn_metadata_header); + let conversation_id = self.client.state.conversation_id.to_string(); + let mut headers = build_responses_headers( + self.client.state.beta_features_header.as_deref(), + Some(&self.turn_state), + turn_metadata_header.as_ref(), + ); + if let Ok(header_value) = HeaderValue::from_str(&conversation_id) { + headers.insert("x-client-request-id", header_value); + } + headers.extend(build_conversation_headers(Some(conversation_id))); headers.extend(self.client.build_subagent_headers()); header_map_to_string_map(&headers) } diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 03dad87341..b4dbda69bb 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -96,15 +96,21 @@ fn test_session_telemetry() -> SessionTelemetry { #[derive(Default)] struct RecordingDelegatedTransport { compact_requests: Mutex>, + stream_requests: Mutex>, } #[async_trait] impl DelegatedModelTransport for RecordingDelegatedTransport { async fn start_model_request( &self, - _request: DelegatedModelRequest, + request: DelegatedModelRequest, ) -> crate::error::Result> { - panic!("streaming delegation is not used in this test"); + self.stream_requests + .lock() + .expect("stream request recording mutex") + .push(request); + let (_tx, rx) = mpsc::channel(1); + Ok(rx) } async fn run_model_compact_request( @@ -319,3 +325,82 @@ async fn delegated_compact_conversation_history_routes_through_transport() { ])) ); } + +#[tokio::test] +async fn delegated_stream_request_uses_live_turn_id_and_http_headers() { + let client = test_model_client(SessionSource::SubAgent(SubAgentSource::Other( + "delegate".to_string(), + ))); + 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: "hello".to_string(), + }], + end_turn: None, + phase: None, + }], + tools: Vec::new(), + parallel_tool_calls: false, + base_instructions: BaseInstructions { + text: "stream instructions".to_string(), + }, + personality: None, + output_schema: None, + }; + + let mut session = client.new_session_with_turn_id("turn-live".to_string()); + let _stream = session + .stream( + &prompt, + &test_model_info(), + &test_session_telemetry(), + None, + ReasoningSummaryConfig::None, + None, + None, + ) + .await + .expect("delegated stream request should be accepted"); + + let recorded = transport + .stream_requests + .lock() + .expect("stream 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-live"); + assert_eq!(request.request_id, "turn-live-0"); + let headers = request + .request + .request_headers + .as_ref() + .expect("delegated stream request should include headers"); + assert_eq!( + headers.get("x-client-request-id"), + Some(&client.state.conversation_id.to_string()) + ); + assert_eq!( + headers.get("session_id"), + Some(&client.state.conversation_id.to_string()) + ); + assert_eq!( + headers.get("x-openai-subagent"), + Some(&"delegate".to_string()) + ); + assert_eq!(headers.get(super::OPENAI_BETA_HEADER), None); + assert_eq!( + headers.get(super::X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER), + None + ); +} diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 85402f5f66..721f4f24a4 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -5659,8 +5659,12 @@ pub(crate) async fn run_turn( // `ModelClientSession` is turn-scoped and caches WebSocket + sticky routing state, so we reuse // one instance across retries within this turn. - let mut client_session = - prewarmed_client_session.unwrap_or_else(|| sess.services.model_client.new_session()); + let mut client_session = prewarmed_client_session.unwrap_or_else(|| { + sess.services + .model_client + .new_session_with_turn_id(turn_context.sub_id.clone()) + }); + client_session.set_turn_id(turn_context.sub_id.clone()); loop { if run_pending_session_start_hooks(&sess, &turn_context).await {