diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 223d4296d5..356b0887b1 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -871,12 +871,19 @@ impl McpConnectionManager { } async fn client_by_name(&self, name: &str) -> Result { - self.clients + let client = self + .clients .get(name) - .ok_or_else(|| anyhow!("unknown MCP server '{name}'"))? - .client() - .await - .context("failed to get client") + .ok_or_else(|| anyhow!("unknown MCP server '{name}'"))?; + if client.is_codex_apps_mcp_server + && client + .codex_apps_tools_cache_context + .as_ref() + .is_some_and(|context| !context.is_active()) + { + return Err(anyhow!("connector runtime context was discarded")); + } + client.client().await.context("failed to get client") } #[cfg(test)] diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 41fc19a594..4b91e6149a 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -1413,6 +1413,59 @@ async fn list_all_tools_adds_server_metadata_to_tools() { assert_eq!(tool.server_origin.as_deref(), Some("https://docs.example")); } +#[tokio::test] +async fn discarded_codex_apps_context_hides_tools_and_rejects_calls() { + let codex_home = tempdir().expect("tempdir"); + let cache = CodexAppsToolsCache::default(); + let context = cache.context( + codex_home.path().to_path_buf(), + CodexAppsToolsCacheKey { + account_id: Some("account-one".to_string()), + chatgpt_user_id: Some("user-one".to_string()), + is_workspace_account: false, + }, + ); + let mut client = create_ready_async_managed_client(vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "old_tool", + )]) + .await; + client.is_codex_apps_mcp_server = true; + client.codex_apps_tools_cache_context = Some(context); + + let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut manager = McpConnectionManager::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + manager + .clients + .insert(CODEX_APPS_MCP_SERVER_NAME.to_string(), client); + + let _new_context = cache.context( + codex_home.path().to_path_buf(), + CodexAppsToolsCacheKey { + account_id: Some("account-two".to_string()), + chatgpt_user_id: Some("user-two".to_string()), + is_workspace_account: false, + }, + ); + + assert!(manager.list_all_tools().await.is_empty()); + let error = manager + .call_tool( + CODEX_APPS_MCP_SERVER_NAME, + "old_tool", + /*arguments*/ None, + /*meta*/ None, + ) + .await + .expect_err("discarded context should reject tool calls"); + assert_eq!(error.to_string(), "connector runtime context was discarded"); +} + #[test] fn server_metadata_preserves_tool_approval_policy() { let mut config = crate::codex_apps_mcp_server_config( diff --git a/codex-rs/codex-mcp/src/connector_runtime.rs b/codex-rs/codex-mcp/src/connector_runtime.rs index 1fca0a2d35..a13c62645b 100644 --- a/codex-rs/codex-mcp/src/connector_runtime.rs +++ b/codex-rs/codex-mcp/src/connector_runtime.rs @@ -205,6 +205,21 @@ impl ConnectorRuntimeContext { ticket: CodexAppsToolsFetchTicket, server_info: &McpServerInfo, tools: Vec, + ) -> Result, ConnectorRuntimeContextDiscarded> { + self.publish_runtime_if_newest_accepted_with( + ticket, + server_info, + tools, + persist_codex_apps_cache, + ) + } + + fn publish_runtime_if_newest_accepted_with( + &self, + ticket: CodexAppsToolsFetchTicket, + server_info: &McpServerInfo, + tools: Vec, + persist: impl FnOnce(&ConnectorRuntimeContext, &McpServerInfo, &ConnectorRuntimeSnapshot), ) -> Result, ConnectorRuntimeContextDiscarded> { let publish_start = Instant::now(); let active = lock_unpoisoned(&self.manager.active); @@ -241,9 +256,11 @@ impl ConnectorRuntimeContext { self.entry .current_snapshot .store(Some(Arc::clone(&snapshot))); + // Keep both guards through persistence so accepted generations cannot reach disk out of + // order, and the same account cannot be reactivated with stale disk state mid-commit. + persist(self, server_info, snapshot.as_ref()); drop(last_accepted_generation); drop(active); - persist_codex_apps_cache(self, server_info, snapshot.as_ref()); emit_duration( MCP_TOOLS_CACHE_PUBLISH_DURATION_METRIC, publish_start.elapsed(), @@ -275,7 +292,7 @@ impl ConnectorRuntimeContext { self.entry.current_snapshot.store(Some(Arc::new(snapshot))); } - fn is_active(&self) -> bool { + pub(crate) fn is_active(&self) -> bool { self.matches_active(lock_unpoisoned(&self.manager.active).as_ref()) } diff --git a/codex-rs/codex-mcp/src/connector_runtime_tests.rs b/codex-rs/codex-mcp/src/connector_runtime_tests.rs index 1bec68fbef..d654f77d04 100644 --- a/codex-rs/codex-mcp/src/connector_runtime_tests.rs +++ b/codex-rs/codex-mcp/src/connector_runtime_tests.rs @@ -581,3 +581,137 @@ fn cold_loaded_snapshot_uses_cache_modification_time() { assert_eq!(model_tool_names(snapshot.tools()), model_tool_names(&tools)); assert_eq!(snapshot.refreshed_at(), modified_at); } +#[test] +fn accepted_generations_finish_persistence_in_order() { + let codex_home = tempdir().expect("tempdir"); + let context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let older_ticket = context.begin_fetch(CodexAppsToolsFetchSource::Startup); + let newer_ticket = context.begin_fetch(CodexAppsToolsFetchSource::HardRefresh); + let (older_persisting_tx, older_persisting_rx) = std::sync::mpsc::channel(); + let (release_older_tx, release_older_rx) = std::sync::mpsc::channel(); + let older_context = context.clone(); + let older_publish = std::thread::spawn(move || { + older_context.publish_runtime_if_newest_accepted_with( + older_ticket, + &create_test_server_info("Codex Apps"), + vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "older")], + move |_, _, _| { + older_persisting_tx + .send(()) + .expect("signal older persistence"); + release_older_rx.recv().expect("release older persistence"); + }, + ) + }); + older_persisting_rx + .recv_timeout(Duration::from_secs(1)) + .expect("older generation should enter persistence"); + + let (newer_persisting_tx, newer_persisting_rx) = std::sync::mpsc::channel(); + let newer_context = context.clone(); + let newer_publish = std::thread::spawn(move || { + newer_context.publish_runtime_if_newest_accepted_with( + newer_ticket, + &create_test_server_info("Codex Apps"), + vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "newer")], + move |_, _, _| { + newer_persisting_tx + .send(()) + .expect("signal newer persistence"); + }, + ) + }); + + assert!( + newer_persisting_rx + .recv_timeout(Duration::from_millis(20)) + .is_err() + ); + release_older_tx + .send(()) + .expect("allow older persistence to finish"); + newer_persisting_rx + .recv_timeout(Duration::from_secs(1)) + .expect("newer generation should persist after older generation"); + + older_publish + .join() + .expect("join older publish") + .expect("publish older generation"); + let newer_snapshot = newer_publish + .join() + .expect("join newer publish") + .expect("publish newer generation"); + assert_eq!(newer_snapshot.tools()[0].callable_name, "newer"); +} + +#[test] +fn workspace_context_change_discards_matching_account_and_user_ids() { + let codex_home = tempdir().expect("tempdir"); + let manager = ConnectorRuntimeManager::default(); + let personal_context = manager.context( + codex_home.path().to_path_buf(), + ConnectorRuntimeContextKey { + account_id: Some("account".to_string()), + chatgpt_user_id: Some("user".to_string()), + is_workspace_account: false, + }, + ); + personal_context + .publish_runtime_if_newest_accepted( + personal_context.begin_fetch(CodexAppsToolsFetchSource::Startup), + &create_test_server_info("Codex Apps"), + vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "personal")], + ) + .expect("publish personal context"); + + let workspace_context = manager.context( + codex_home.path().to_path_buf(), + ConnectorRuntimeContextKey { + account_id: Some("account".to_string()), + chatgpt_user_id: Some("user".to_string()), + is_workspace_account: true, + }, + ); + + assert!(personal_context.current_snapshot().is_none()); + assert!(workspace_context.current_snapshot().is_none()); +} + +#[test] +fn live_publish_sets_timestamp_and_stale_publish_preserves_it() { + let codex_home = tempdir().expect("tempdir"); + let context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("account-one"), + Some("user-one"), + ); + let stale_ticket = context.begin_fetch(CodexAppsToolsFetchSource::Startup); + let current_ticket = context.begin_fetch(CodexAppsToolsFetchSource::HardRefresh); + let before = SystemTime::now(); + let current = context + .publish_runtime_if_newest_accepted( + current_ticket, + &create_test_server_info("Codex Apps"), + vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "current")], + ) + .expect("publish current generation"); + let after = SystemTime::now(); + + assert!(current.refreshed_at() >= before); + assert!(current.refreshed_at() <= after); + + let stale = context + .publish_runtime_if_newest_accepted( + stale_ticket, + &create_test_server_info("Codex Apps"), + vec![create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "stale")], + ) + .expect("stale publish returns current snapshot"); + assert!(Arc::ptr_eq(¤t, &stale)); + assert_eq!(stale.refreshed_at(), current.refreshed_at()); +} diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 6fe820b4fe..19e7c6c326 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -517,6 +517,14 @@ impl AsyncManagedClient { } pub(crate) async fn listed_tools(&self) -> Option> { + if self.is_codex_apps_mcp_server + && self + .codex_apps_tools_cache_context + .as_ref() + .is_some_and(|context| !context.is_active()) + { + return None; + } // Keep cache payloads raw; plugin provenance is resolved per-session at read time. let tools = if !self.startup_complete.load(Ordering::Acquire) && let Some(startup_tools) = self.cached_tools()