Skip to content

Commit 238497c

Browse files
shijie-oaishijie-oai
committed
Mcp server refresh sends the latest mcp configs
1 parent 8c46235 commit 238497c

File tree

4 files changed

+107
-28
lines changed

4 files changed

+107
-28
lines changed

‎codex-rs/app-server/src/codex_message_processor.rs‎

Lines changed: 47 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -159,6 +159,7 @@ use codex_protocol::items::TurnItem;
159159
use codex_protocol::models::ResponseItem;
160160
use codex_protocol::protocol::GitInfo as CoreGitInfo;
161161
use codex_protocol::protocol::McpAuthStatus as CoreMcpAuthStatus;
162+
use codex_protocol::protocol::McpServerRefreshConfig;
162163
use codex_protocol::protocol::RateLimitSnapshot as CoreRateLimitSnapshot;
163164
use codex_protocol::protocol::RolloutItem;
164165
use codex_protocol::protocol::SessionMetaLine;
@@ -2308,7 +2309,52 @@ impl CodexMessageProcessor {
23082309
}
23092310

23102311
async fn mcp_server_refresh(&self, request_id: RequestId, _params: McpServerRefreshParams) {
2311-
self.thread_manager.refresh_mcp_servers().await;
2312+
let config = match self.load_latest_config().await {
2313+
Ok(config) => config,
2314+
Err(error) => {
2315+
self.outgoing.send_error(request_id, error).await;
2316+
return;
2317+
}
2318+
};
2319+
2320+
let mcp_servers = match serde_json::to_value(&config.mcp_servers) {
2321+
Ok(value) => value,
2322+
Err(err) => {
2323+
let error = JSONRPCErrorError {
2324+
code: INTERNAL_ERROR_CODE,
2325+
message: format!("failed to serialize MCP servers: {err}"),
2326+
data: None,
2327+
};
2328+
self.outgoing.send_error(request_id, error).await;
2329+
return;
2330+
}
2331+
};
2332+
2333+
let mcp_oauth_credentials_store_mode =
2334+
match serde_json::to_value(config.mcp_oauth_credentials_store_mode) {
2335+
Ok(value) => value,
2336+
Err(err) => {
2337+
let error = JSONRPCErrorError {
2338+
code: INTERNAL_ERROR_CODE,
2339+
message: format!(
2340+
"failed to serialize MCP OAuth credentials store mode: {err}"
2341+
),
2342+
data: None,
2343+
};
2344+
self.outgoing.send_error(request_id, error).await;
2345+
return;
2346+
}
2347+
};
2348+
2349+
let refresh_config = McpServerRefreshConfig {
2350+
mcp_servers,
2351+
mcp_oauth_credentials_store_mode,
2352+
};
2353+
2354+
let thread_manager = Arc::clone(&self.thread_manager);
2355+
tokio::spawn(async move {
2356+
thread_manager.refresh_mcp_servers(refresh_config).await;
2357+
});
23122358
let response = McpServerRefreshResponse {};
23132359
self.outgoing.send_response(request_id, response).await;
23142360
}

‎codex-rs/core/src/codex.rs‎

Lines changed: 44 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ use codex_protocol::protocol::TurnAbortReason;
4848
use codex_protocol::protocol::TurnContextItem;
4949
use codex_protocol::protocol::TurnStartedEvent;
5050
use codex_rmcp_client::ElicitationResponse;
51+
use codex_rmcp_client::OAuthCredentialsStoreMode;
5152
use futures::future::BoxFuture;
5253
use futures::prelude::*;
5354
use futures::stream::FuturesOrdered;
@@ -84,6 +85,7 @@ use crate::config::Config;
8485
use crate::config::Constrained;
8586
use crate::config::ConstraintResult;
8687
use crate::config::GhostSnapshotConfig;
88+
use crate::config::types::McpServerConfig;
8789
use crate::config::types::ShellEnvironmentPolicy;
8890
use crate::context_manager::ContextManager;
8991
use crate::environment_context::EnvironmentContext;
@@ -107,6 +109,7 @@ use crate::protocol::ErrorEvent;
107109
use crate::protocol::Event;
108110
use crate::protocol::EventMsg;
109111
use crate::protocol::ExecApprovalRequestEvent;
112+
use crate::protocol::McpServerRefreshConfig;
110113
use crate::protocol::Op;
111114
use crate::protocol::RateLimitSnapshot;
112115
use crate::protocol::ReasoningContentDeltaEvent;
@@ -361,7 +364,7 @@ pub(crate) struct Session {
361364
/// The set of enabled features should be invariant for the lifetime of the
362365
/// session.
363366
features: Features,
364-
should_refresh_mcp_servers: AtomicBool,
367+
pending_mcp_server_refresh_config: Mutex<Option<McpServerRefreshConfig>>,
365368
pub(crate) active_turn: Mutex<Option<ActiveTurn>>,
366369
pub(crate) services: SessionServices,
367370
next_internal_sub_id: AtomicU64,
@@ -707,7 +710,7 @@ impl Session {
707710
agent_status: Arc::clone(&agent_status),
708711
state: Mutex::new(state),
709712
features: config.features.clone(),
710-
should_refresh_mcp_servers: AtomicBool::new(false),
713+
pending_mcp_server_refresh_config: Mutex::new(None),
711714
active_turn: Mutex::new(None),
712715
services,
713716
next_internal_sub_id: AtomicU64::new(0),
@@ -1635,25 +1638,41 @@ impl Session {
16351638
Arc::clone(&self.services.user_shell)
16361639
}
16371640

1638-
fn request_mcp_server_refresh(&self) {
1639-
self.should_refresh_mcp_servers
1640-
.store(true, Ordering::SeqCst);
1641+
async fn request_mcp_server_refresh(&self, refresh_config: McpServerRefreshConfig) {
1642+
let mut guard = self.pending_mcp_server_refresh_config.lock().await;
1643+
*guard = Some(refresh_config);
16411644
}
16421645

16431646
async fn refresh_mcp_servers_if_requested(&self, turn_context: &TurnContext) {
1644-
if !self
1645-
.should_refresh_mcp_servers
1646-
.swap(false, Ordering::SeqCst)
1647-
{
1647+
let refresh_config = { self.pending_mcp_server_refresh_config.lock().await.take() };
1648+
let Some(refresh_config) = refresh_config else {
16481649
return;
1649-
}
1650+
};
16501651

1651-
let config = turn_context.client.config();
1652-
let auth_statuses = compute_auth_statuses(
1653-
config.mcp_servers.iter(),
1654-
config.mcp_oauth_credentials_store_mode,
1655-
)
1656-
.await;
1652+
let McpServerRefreshConfig {
1653+
mcp_servers,
1654+
mcp_oauth_credentials_store_mode,
1655+
} = refresh_config;
1656+
1657+
let mcp_servers =
1658+
match serde_json::from_value::<HashMap<String, McpServerConfig>>(mcp_servers) {
1659+
Ok(servers) => servers,
1660+
Err(err) => {
1661+
warn!("failed to parse MCP server refresh config: {err}");
1662+
return;
1663+
}
1664+
};
1665+
let store_mode = match serde_json::from_value::<OAuthCredentialsStoreMode>(
1666+
mcp_oauth_credentials_store_mode,
1667+
) {
1668+
Ok(mode) => mode,
1669+
Err(err) => {
1670+
warn!("failed to parse MCP OAuth refresh config: {err}");
1671+
return;
1672+
}
1673+
};
1674+
1675+
let auth_statuses = compute_auth_statuses(mcp_servers.iter(), store_mode).await;
16571676
let sandbox_state = SandboxState {
16581677
sandbox_policy: turn_context.sandbox_policy.clone(),
16591678
codex_linux_sandbox_exe: turn_context.codex_linux_sandbox_exe.clone(),
@@ -1664,8 +1683,8 @@ impl Session {
16641683
let mut refreshed_manager = McpConnectionManager::default();
16651684
refreshed_manager
16661685
.initialize(
1667-
config.mcp_servers.clone(),
1668-
config.mcp_oauth_credentials_store_mode,
1686+
mcp_servers,
1687+
store_mode,
16691688
auth_statuses,
16701689
self.get_tx_event(),
16711690
cancel_token,
@@ -1760,8 +1779,8 @@ async fn submission_loop(sess: Arc<Session>, config: Arc<Config>, rx_sub: Receiv
17601779
Op::ListMcpTools => {
17611780
handlers::list_mcp_tools(&sess, &config, sub.id.clone()).await;
17621781
}
1763-
Op::RefreshMcpServers => {
1764-
handlers::refresh_mcp_servers(&sess).await;
1782+
Op::RefreshMcpServers { config } => {
1783+
handlers::refresh_mcp_servers(&sess, config).await;
17651784
}
17661785
Op::ListCustomPrompts => {
17671786
handlers::list_custom_prompts(&sess, sub.id.clone()).await;
@@ -1831,6 +1850,7 @@ mod handlers {
18311850
use codex_protocol::protocol::EventMsg;
18321851
use codex_protocol::protocol::ListCustomPromptsResponseEvent;
18331852
use codex_protocol::protocol::ListSkillsResponseEvent;
1853+
use codex_protocol::protocol::McpServerRefreshConfig;
18341854
use codex_protocol::protocol::Op;
18351855
use codex_protocol::protocol::ReviewDecision;
18361856
use codex_protocol::protocol::ReviewRequest;
@@ -2062,8 +2082,8 @@ mod handlers {
20622082
});
20632083
}
20642084

2065-
pub async fn refresh_mcp_servers(sess: &Arc<Session>) {
2066-
sess.request_mcp_server_refresh();
2085+
pub async fn refresh_mcp_servers(sess: &Arc<Session>, refresh_config: McpServerRefreshConfig) {
2086+
sess.request_mcp_server_refresh(refresh_config).await;
20672087
}
20682088

20692089
pub async fn list_mcp_tools(sess: &Session, config: &Arc<Config>, sub_id: String) {
@@ -3618,7 +3638,7 @@ mod tests {
36183638
agent_status: Arc::clone(&agent_status),
36193639
state: Mutex::new(state),
36203640
features: config.features.clone(),
3621-
should_refresh_mcp_servers: AtomicBool::new(false),
3641+
pending_mcp_server_refresh_config: Mutex::new(None),
36223642
active_turn: Mutex::new(None),
36233643
services,
36243644
next_internal_sub_id: AtomicU64::new(0),
@@ -3713,7 +3733,7 @@ mod tests {
37133733
agent_status: Arc::clone(&agent_status),
37143734
state: Mutex::new(state),
37153735
features: config.features.clone(),
3716-
should_refresh_mcp_servers: AtomicBool::new(false),
3736+
pending_mcp_server_refresh_config: Mutex::new(None),
37173737
active_turn: Mutex::new(None),
37183738
services,
37193739
next_internal_sub_id: AtomicU64::new(0),

‎codex-rs/core/src/thread_manager.rs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ use crate::skills::SkillsManager;
2121
use codex_protocol::ThreadId;
2222
use codex_protocol::openai_models::ModelPreset;
2323
use codex_protocol::protocol::InitialHistory;
24+
use codex_protocol::protocol::McpServerRefreshConfig;
2425
use codex_protocol::protocol::Op;
2526
use codex_protocol::protocol::RolloutItem;
2627
use codex_protocol::protocol::SessionSource;
@@ -137,7 +138,7 @@ impl ThreadManager {
137138
self.state.threads.read().await.keys().copied().collect()
138139
}
139140

140-
pub async fn refresh_mcp_servers(&self) {
141+
pub async fn refresh_mcp_servers(&self, refresh_config: McpServerRefreshConfig) {
141142
let threads = self
142143
.state
143144
.threads
@@ -147,7 +148,12 @@ impl ThreadManager {
147148
.cloned()
148149
.collect::<Vec<_>>();
149150
for thread in threads {
150-
if let Err(err) = thread.submit(Op::RefreshMcpServers).await {
151+
if let Err(err) = thread
152+
.submit(Op::RefreshMcpServers {
153+
config: refresh_config.clone(),
154+
})
155+
.await
156+
{
151157
warn!("failed to request MCP server refresh: {err}");
152158
}
153159
}

‎codex-rs/protocol/src/protocol.rs‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,13 @@ pub struct Submission {
6060
pub op: Op,
6161
}
6262

63+
/// Config payload for refreshing MCP servers.
64+
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
65+
pub struct McpServerRefreshConfig {
66+
pub mcp_servers: Value,
67+
pub mcp_oauth_credentials_store_mode: Value,
68+
}
69+
6370
/// Submission operation
6471
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, JsonSchema)]
6572
#[serde(tag = "type", rename_all = "snake_case")]
@@ -187,7 +194,7 @@ pub enum Op {
187194
ListMcpTools,
188195

189196
/// Request MCP servers to reinitialize and refresh cached tool lists.
190-
RefreshMcpServers,
197+
RefreshMcpServers { config: McpServerRefreshConfig },
191198

192199
/// Request the list of available custom prompts.
193200
ListCustomPrompts,

0 commit comments

Comments
 (0)