Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions codex-rs/core/src/session/handlers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -422,7 +422,7 @@ pub async fn dynamic_tool_response(sess: &Arc<Session>, id: String, response: Dy
}

pub fn refresh_mcp_servers(sess: &Session) {
sess.mark_mcp_runtime_dirty();
sess.request_mcp_runtime_refresh();
}

pub async fn reload_user_config(sess: &Arc<Session>) {
Expand Down Expand Up @@ -585,9 +585,12 @@ async fn shutdown_session_runtime(sess: &Arc<Session>) {
if let Err(err) = sess.services.code_mode_service.shutdown().await {
warn!("failed to shutdown code mode session: {err}");
}
let _refresh = sess.mcp_refresh_lock.acquire().await;
sess.mcp_refresh_lock.close();
sess.services.mcp_runtime.shutdown().await;
sess.stop_mcp_prewarm_worker().await;
{
let _refresh = sess.mcp_refresh_lock.acquire().await;
sess.mcp_refresh_lock.close();
sess.services.mcp_runtime.shutdown().await;
}
sess.guardian_review_session.shutdown().await;

crate::hook_runtime::run_session_end_hooks(sess).await;
Expand Down
2 changes: 1 addition & 1 deletion codex-rs/core/src/session/mcp.rs
Original file line number Diff line number Diff line change
Expand Up @@ -502,7 +502,7 @@ impl Session {
self.services
.supports_openai_form_elicitation
.store(supported, std::sync::atomic::Ordering::Relaxed);
self.mark_mcp_runtime_dirty();
self.request_mcp_runtime_refresh();
Ok(())
}

Expand Down
75 changes: 75 additions & 0 deletions codex-rs/core/src/session/mcp_prewarm.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
//! Best-effort MCP prewarming.
//!
//! A bounded channel coalesces refresh requests. The worker only prepares the
//! newest thread state; exact model steps remain the correctness path.

use super::*;

impl Session {
pub(super) fn request_mcp_runtime_refresh(&self) {
self.mark_mcp_runtime_dirty();
self.schedule_mcp_prewarm();
}

pub(super) fn start_mcp_prewarm_worker(
self: &Arc<Self>,
requests: async_channel::Receiver<()>,
mut auth_changes: tokio::sync::watch::Receiver<u64>,
) {
let session = Arc::downgrade(self);
let shutdown = self.mcp_prewarm_shutdown.clone();
let worker = self.services.runtime_handle.spawn(async move {
loop {
let auth_changed = tokio::select! {
biased;
_ = shutdown.cancelled() => break,
request = requests.recv() => {
if request.is_err() {
break;
}
false
},
auth_change = auth_changes.changed() => {
if auth_change.is_err() {
break;
}
true
},
};
let Some(session) = session.upgrade() else {
break;
};
if auth_changed {
session.mark_mcp_runtime_dirty();
}
tokio::select! {
biased;
_ = shutdown.cancelled() => break,
_ = session.refresh_mcp_if_dirty() => {},
}
}
});
*self
.mcp_prewarm_task
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = Some(worker);
}

pub(super) fn schedule_mcp_prewarm(&self) {
let _ = self.mcp_prewarm_tx.try_send(());
}

pub(super) async fn stop_mcp_prewarm_worker(&self) {
self.mcp_prewarm_shutdown.cancel();
let worker = self
.mcp_prewarm_task
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.take();
if let Some(worker) = worker
&& let Err(error) = worker.await
{
warn!(%error, "MCP prewarm worker stopped unexpectedly");
}
}
}
18 changes: 15 additions & 3 deletions codex-rs/core/src/session/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,7 @@ mod handlers;
mod inject;
mod input_queue;
mod mcp;
mod mcp_prewarm;
mod mcp_runtime;
pub(crate) mod multi_agents;
mod review;
Expand Down Expand Up @@ -1483,7 +1484,7 @@ impl Session {
updates: SessionSettingsUpdate,
) -> ConstraintResult<()> {
let notify_config_contributors = !self.services.extensions.config_contributors().is_empty();
let (previous_config, new_config, permission_profile_changed) = {
let (previous_config, new_config, permission_profile_changed, mcp_inputs_changed) = {
let mut state = self.state.lock().await;
let updated = match state.session_configuration.apply(&updates) {
Ok(updated) => updated,
Expand All @@ -1507,17 +1508,25 @@ impl Session {
.turn_environments
.update_selections(updated.environment_selections());
}
state.session_configuration = updated;
if mcp_inputs_changed {
self.mark_mcp_runtime_dirty();
}
state.session_configuration = updated;
(previous_config, new_config, permission_profile_changed)
(
previous_config,
new_config,
permission_profile_changed,
mcp_inputs_changed,
)
};
self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref());
if permission_profile_changed {
self.refresh_managed_network_proxy_for_current_permission_profile()
.await;
}
if mcp_inputs_changed {
self.schedule_mcp_prewarm();
}
Ok(())
}

Expand Down Expand Up @@ -1633,6 +1642,7 @@ impl Session {
(previous_config, new_config, config)
};
self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref());
self.schedule_mcp_prewarm();
let environments = self.services.turn_environments.snapshot().await;
let hooks = build_hooks_for_config(
config.as_ref(),
Expand Down Expand Up @@ -1698,6 +1708,8 @@ impl Session {
}
state.session_configuration.original_config_do_not_use = Arc::new(config);
self.mark_mcp_runtime_dirty();
drop(state);
self.schedule_mcp_prewarm();
}

fn emit_config_changed_contributors(
Expand Down
9 changes: 9 additions & 0 deletions codex-rs/core/src/session/session.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,9 @@ pub(crate) struct Session {
pub(super) mcp_refresh_lock: Semaphore,
pub(super) mcp_elicitation_reviewer_handle: OnceLock<codex_mcp::ElicitationReviewerHandle>,
pub(super) mcp_elicitation_lifecycle_handle: OnceLock<codex_mcp::ElicitationLifecycle>,
pub(super) mcp_prewarm_tx: async_channel::Sender<()>,
pub(super) mcp_prewarm_shutdown: CancellationToken,
pub(super) mcp_prewarm_task: std::sync::Mutex<Option<JoinHandle<()>>>,
pub(crate) conversation: Arc<RealtimeConversationManager>,
pub(crate) active_turn: Mutex<Option<ActiveTurn>>,
pub(crate) input_queue: InputQueue,
Expand Down Expand Up @@ -1145,6 +1148,7 @@ impl Session {
tool_search_handler_cache: Default::default(),
turn_environments: Arc::clone(&turn_environments),
};
let (mcp_prewarm_tx, mcp_prewarm_rx) = async_channel::bounded(1);
let sess = Arc::new(Session {
thread_id,
installation_id,
Expand All @@ -1159,6 +1163,9 @@ impl Session {
mcp_refresh_lock: Semaphore::new(/*permits*/ 1),
mcp_elicitation_reviewer_handle: OnceLock::new(),
mcp_elicitation_lifecycle_handle: OnceLock::new(),
mcp_prewarm_tx,
mcp_prewarm_shutdown: CancellationToken::new(),
mcp_prewarm_task: std::sync::Mutex::new(None),
conversation: Arc::new(RealtimeConversationManager::new()),
active_turn: Mutex::new(None),
input_queue: InputQueue::new(),
Expand Down Expand Up @@ -1209,13 +1216,15 @@ impl Session {
}
turn_environments.start_connection_event_forwarding(tx_event.clone());

let mcp_auth_changes = sess.services.auth_manager.auth_change_receiver();
sess.install_initial_mcp_runtime(
&session_configuration,
mcp_projection,
&resolved_environments,
mcp_runtime_cwd,
)
.await?;
sess.start_mcp_prewarm_worker(mcp_prewarm_rx, mcp_auth_changes);
sess.schedule_startup_prewarm(session_configuration.base_instructions.clone())
.await;
let session_start_source = match &initial_history {
Expand Down
29 changes: 29 additions & 0 deletions codex-rs/core/src/session/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5548,6 +5548,9 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
mcp_refresh_lock: Semaphore::new(/*permits*/ 1),
mcp_elicitation_reviewer_handle: OnceLock::new(),
mcp_elicitation_lifecycle_handle: OnceLock::new(),
mcp_prewarm_tx: async_channel::bounded(1).0,
mcp_prewarm_shutdown: CancellationToken::new(),
mcp_prewarm_task: std::sync::Mutex::new(None),
conversation: Arc::new(RealtimeConversationManager::new()),
active_turn: Mutex::new(None),
input_queue: super::input_queue::InputQueue::new(),
Expand Down Expand Up @@ -7715,6 +7718,9 @@ where
mcp_refresh_lock: Semaphore::new(/*permits*/ 1),
mcp_elicitation_reviewer_handle: OnceLock::new(),
mcp_elicitation_lifecycle_handle: OnceLock::new(),
mcp_prewarm_tx: async_channel::bounded(1).0,
mcp_prewarm_shutdown: CancellationToken::new(),
mcp_prewarm_task: std::sync::Mutex::new(None),
conversation: Arc::new(RealtimeConversationManager::new()),
active_turn: Mutex::new(None),
input_queue: super::input_queue::InputQueue::new(),
Expand Down Expand Up @@ -7907,6 +7913,29 @@ async fn mcp_elicitation_reviewer_is_reused_across_runtime_refreshes() {
assert!(Arc::ptr_eq(&previous, &session.mcp_elicitation_reviewer()));
}

#[tokio::test]
async fn mcp_policy_changes_schedule_runtime_refresh() {
let (session, _turn_context) = make_session_and_context().await;
let session = Arc::new(session);

session
.new_turn_with_sub_id(
"policy-change".to_string(),
SessionSettingsUpdate {
approval_policy: Some(AskForApproval::Never),
..Default::default()
},
)
.await
.expect("approval policy update should succeed");

assert!(
session
.mcp_refresh_pending
.load(std::sync::atomic::Ordering::Acquire)
);
}

struct PendingNoiseConnectProvider;

impl codex_exec_server::NoiseRendezvousConnectProvider for PendingNoiseConnectProvider {
Expand Down
42 changes: 25 additions & 17 deletions codex-rs/core/src/session/turn_context.rs
Original file line number Diff line number Diff line change
Expand Up @@ -638,6 +638,7 @@ impl Session {
state.session_configuration = next.clone();
Ok((
next,
mcp_inputs_changed,
permission_profile_changed,
previous_config,
new_config,
Expand All @@ -647,23 +648,31 @@ impl Session {
}
};

let (session_configuration, permission_profile_changed, previous_config, new_config) =
match update_result {
Ok(update) => update,
Err(err) => {
let message = err.to_string();
self.send_event_raw(Event {
id: sub_id.clone(),
msg: EventMsg::Error(ErrorEvent {
message: message.clone(),
codex_error_info: Some(CodexErrorInfo::BadRequest),
}),
})
.await;
return Err(CodexErr::InvalidRequest(message));
}
};
let (
session_configuration,
mcp_inputs_changed,
permission_profile_changed,
previous_config,
new_config,
) = match update_result {
Ok(update) => update,
Err(err) => {
let message = err.to_string();
self.send_event_raw(Event {
id: sub_id.clone(),
msg: EventMsg::Error(ErrorEvent {
message: message.clone(),
codex_error_info: Some(CodexErrorInfo::BadRequest),
}),
})
.await;
return Err(CodexErr::InvalidRequest(message));
}
};
self.emit_config_changed_contributors(previous_config.as_ref(), new_config.as_ref());
if mcp_inputs_changed {
self.schedule_mcp_prewarm();
}

if permission_profile_changed {
self.refresh_managed_network_proxy_for_current_permission_profile()
Expand Down Expand Up @@ -727,7 +736,6 @@ impl Session {
.and_then(|turn_environment| turn_environment.cwd().to_abs_path().ok())
.unwrap_or_else(|| session_configuration.cwd().clone());
let per_turn_config = Self::build_per_turn_config(&session_configuration, cwd.clone());

let model_info = self
.services
.models_manager
Expand Down
1 change: 1 addition & 0 deletions codex-rs/core/src/tools/spec_plan_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1119,6 +1119,7 @@ async fn tool_search_cache_rebuilds_when_deferred_world_state_changes() {
)],
tool_suggest_candidates: None,
extension_tool_executors: Vec::new(),
wait_for_environment_tool_config: None,
dynamic_tools: &[],
},
&cache,
Expand Down
Loading
Loading