Refresh MCP runtimes when session auth changes (#35146)

## What changed

- Track the authentication state used by each published MCP runtime and mark the runtime dirty when the auth mode, account details, or token changes.
- Refresh plugin auth mode before evaluating pending MCP work, and rebuild the initial MCP projection if authentication changes during session startup.
- Carry one auth snapshot through runtime construction so plugin discovery and MCP connections use consistent credentials.

## Testing

- Add a regression test that replaces an API key, reloads authentication, and verifies that MCP refresh publishes a runtime using the new credentials.

GitOrigin-RevId: 219c2f7a157f53b5467116e18bfe70a8c3dca695
This commit is contained in:
jif
2026-07-24 09:50:41 +00:00
committed by copyberry
parent ef2d3edb95
commit 6c729ef1c1
5 changed files with 160 additions and 23 deletions
+24
View File
@@ -76,6 +76,8 @@ pub struct McpRuntime {
struct PublishedMcpRuntime {
connections: Arc<McpConnectionSet>,
config: Option<Arc<McpConfig>>,
auth: Option<CodexAuth>,
auth_token: Option<String>,
plugins_available: bool,
ready_selected_capability_roots: Vec<SelectedCapabilityRoot>,
}
@@ -125,6 +127,8 @@ impl McpRuntime {
current: ArcSwap::from_pointee(PublishedMcpRuntime {
connections: Arc::new(McpConnectionSet::empty(prefix_mcp_tool_names)),
config: None,
auth: None,
auth_token: None,
plugins_available: false,
ready_selected_capability_roots: Vec::new(),
}),
@@ -154,6 +158,8 @@ impl McpRuntime {
async fn publish(&self, input: McpRuntimeInput, previous: Option<&McpConnectionSet>) {
let (publish, publication_gate) = McpPublicationGate::pending();
let config = Arc::clone(&input.config);
let auth = input.auth.clone();
let auth_token = auth.as_ref().and_then(|auth| auth.get_token().ok());
let plugins_available = input.plugins_available;
let ready_selected_capability_roots = input.ready_selected_capability_roots.clone();
let connections = Arc::new(
@@ -168,6 +174,8 @@ impl McpRuntime {
self.current.store(Arc::new(PublishedMcpRuntime {
connections,
config: Some(config),
auth,
auth_token,
plugins_available,
ready_selected_capability_roots,
}));
@@ -186,6 +194,22 @@ impl McpRuntime {
))
}
/// Returns whether the published snapshot still belongs to the current credentials.
pub fn current_auth_matches(&self, auth: Option<&CodexAuth>) -> bool {
let current = self.current.load();
match (current.auth.as_ref(), auth) {
(Some(previous), Some(latest)) => {
previous == latest
&& previous.get_account_id() == latest.get_account_id()
&& previous.get_chatgpt_user_id() == latest.get_chatgpt_user_id()
&& previous.is_fedramp_account() == latest.is_fedramp_account()
&& current.auth_token == latest.get_token().ok()
}
(None, None) => true,
(Some(_), None) | (None, Some(_)) => false,
}
}
pub fn current_ready_selected_capability_roots(&self) -> Vec<SelectedCapabilityRoot> {
self.current.load().ready_selected_capability_roots.clone()
}