openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-28 18:58:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

core: route delegated model transport through app-server

Browse Source
This commit is contained in:
sdcoffey
2026-03-19 09:27:24 -07:00
parent f5affbd860
commit 73a20bcb37
13 changed files with 1446 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/codex_message_processor.rs
+54
View File
@@ -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<HashMap<String, serde_json::Value>>,
typesafe_overrides: ConfigOverrides,
dynamic_tools: Option<Vec<ApiDynamicToolSpec>>,
network_delegation: Option<codex_app_server_protocol::NetworkDelegationConfig>,
persist_extended_history: bool,
service_name: Option<String>,
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<codex_app_server_protocol::NetworkDelegationConfig>,
) -> 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(&params.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,
};
codex-rs/app-server/src/delegated_model.rs
+650
View File
@@ -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<OutgoingMessageSender>,
thread_state_manager: ThreadStateManager,
inflight: Arc<Mutex<HashMap<DelegatedRequestKey, mpsc::Sender<DelegatedModelEvent>>>>,
}
impl AppServerDelegatedModelTransport {
pub(crate) fn new(
outgoing: Arc<OutgoingMessageSender>,
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(&notification);
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(&notification.event);
let key = DelegatedRequestKey::from_stream_event(&notification);
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(&notification);
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<ConnectionId> {
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<ClientRequestResult, tokio::sync::oneshot::error::RecvError>,
) -> Result<()> {
match response {
Ok(Ok(value)) => {
let response = serde_json::from_value::<ModelRequestResponse>(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<mpsc::Receiver<DelegatedModelEvent>> {
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<Vec<codex_protocol::models::ResponseItem>> {
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::<ModelCompactResponse>(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
);
}
}
codex-rs/app-server/src/lib.rs
+1
View File
@@ -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;
codex-rs/app-server/src/message_processor.rs
+20 -4
View File
@@ -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<OutgoingMessageSender>,
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<F>(
codex-rs/core/src/client.rs
+383 -20
View File
@@ -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<Option<Arc<dyn DelegatedModelTransport>>>);
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<Arc<AuthManager>>,
@@ -142,6 +164,8 @@ struct ModelClientState {
beta_features_header: Option<String>,
disable_websockets: AtomicBool,
cached_websocket_session: StdMutex<WebsocketSession>,
delegated_model_transport: DelegatedModelTransportSlot,
network_delegation_config: StdMutex<Option<NetworkDelegationConfig>>,
}
/// 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<Arc<dyn DelegatedModelTransport>>,
) {
*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<NetworkDelegationConfig>) {
*self
.state
.network_delegation_config
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = config;
}
fn delegated_model_transport(&self) -> Option<Arc<dyn DelegatedModelTransport>> {
self.state
.delegated_model_transport
.0
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone()
}
fn network_delegation_config(&self) -> Option<NetworkDelegationConfig> {
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<ReasoningEffortConfig>,
@@ -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<NetworkDelegationConfig> {
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<String, String> {
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<String, String>,
) -> Result<ModelRequestEnvelope> {
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<ReasoningEffortConfig>,
summary: ReasoningSummaryConfig,
service_tier: Option<ServiceTier>,
turn_metadata_header: Option<&str>,
) -> Result<ResponseStream> {
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<ServiceTier>,
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<ServiceTier>,
turn_metadata_header: Option<&str>,
) -> Result<ResponseStream> {
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<String, String> {
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<String, String>) -> ApiHeaderMap {
let mut header_map = ApiHeaderMap::new();
for (name, value) in headers {
if let (Ok(header_name), Ok(header_value)) = (
name.parse::<http::HeaderName>(),
HeaderValue::from_str(value),
) {
header_map.insert(header_name, header_value);
}
}
header_map
}
fn metadata_events(
headers: &HashMap<String, String>,
turn_state: &Arc<OnceLock<String>>,
) -> Vec<ResponseEvent> {
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<DelegatedModelEvent>,
idle_timeout: Duration,
turn_state: Arc<OnceLock<String>>,
) -> ResponseStream {
let (tx_event, rx_event) = mpsc::channel::<Result<ResponseEvent>>(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::<ResponsesStreamEvent>(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<S>(
api_stream: S,
session_telemetry: SessionTelemetry,
codex-rs/core/src/client_tests.rs
+204
View File
@@ -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<Vec<DelegatedModelCompactRequest>>,
}
#[async_trait]
impl DelegatedModelTransport for RecordingDelegatedTransport {
async fn start_model_request(
&self,
_request: DelegatedModelRequest,
) -> crate::error::Result<mpsc::Receiver<DelegatedModelEvent>> {
panic!("streaming delegation is not used in this test");
}
async fn run_model_compact_request(
&self,
request: DelegatedModelCompactRequest,
) -> crate::error::Result<Vec<ResponseItem>> {
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::<serde_json::Value>::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()),
]))
);
}
codex-rs/core/src/codex.rs
+39
View File
@@ -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<NetworkDelegationConfig>,
) -> 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<Arc<dyn DelegatedModelTransport>>,
) {
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<String>,
app_server_client_name: Option<String>,
network_delegation: Option<NetworkDelegationConfig>,
/// Source of the session (cli, vscode, exec, mcp, ...)
session_source: SessionSource,
dynamic_tools: Vec<DynamicToolSpec>,
@@ -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<Option<Value>>,
pub(crate) personality: Option<Personality>,
pub(crate) app_server_client_name: Option<String>,
pub(crate) network_delegation: Option<NetworkDelegationConfig>,
}
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,
},
)
}
codex-rs/core/src/codex_tests.rs
+6
View File
@@ -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,
codex-rs/core/src/codex_thread.rs
+19
View File
@@ -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<ServiceTier>,
pub network_delegation: Option<NetworkDelegationConfig>,
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<NetworkDelegationConfig>,
) -> ConstraintResult<()> {
self.codex.set_network_delegation(network_delegation).await
}
pub fn set_delegated_model_transport(
&self,
delegated_model_transport: Option<Arc<dyn DelegatedModelTransport>>,
) {
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
codex-rs/core/src/compact_remote.rs
+1
View File
@@ -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,
codex-rs/core/src/delegated_model_transport.rs
+46
View File
@@ -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<String, String>),
StreamEvent(Value),
RequestFailed(ModelRequestError),
}
#[async_trait]
pub trait DelegatedModelTransport: Send + Sync {
async fn start_model_request(
&self,
request: DelegatedModelRequest,
) -> Result<mpsc::Receiver<DelegatedModelEvent>>;
async fn run_model_compact_request(
&self,
request: DelegatedModelCompactRequest,
) -> Result<Vec<ResponseItem>>;
}
codex-rs/core/src/lib.rs
+1
View File
@@ -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;
codex-rs/core/src/thread_manager.rs
+22
View File
@@ -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<McpManager>,
file_watcher: Arc<FileWatcher>,
session_source: SessionSource,
delegated_model_transport: StdMutex<Option<Arc<dyn DelegatedModelTransport>>>,
// Captures submitted ops for testing purpose when test mode is enabled.
ops_log: Option<SharedCapturedOps>,
}
@@ -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<Arc<dyn DelegatedModelTransport>>,
) {
*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
}