mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
Charlie Marsh
41b4fabbb4
## Summary We originally addressed startup prewarming holding the read side of `RwLock<McpConnectionManager>` by snapshotting tool-list state. Review feedback identified the broader ownership problem: the outer synchronization should only publish or retrieve the current manager, while MCP operations rely on the manager's internal synchronization. A follow-up preserved operation retirement with a separate gate, but further review questioned whether that synchronization was actually required and whether we could support latest-wins replacement instead. This PR now stores the current MCP manager in `ArcSwap`. Each operation uses `load_full()` to obtain an owned `Arc<McpConnectionManager>`, then performs MCP I/O without retaining the publication mechanism. Refresh cancels obsolete startup work, constructs a replacement, and atomically publishes it. New operations see the latest manager, while operations that already loaded the previous manager retain a valid handle. Refresh happens at a turn boundary, so there should be no active user tool calls to drain. Git history supports dropping the outer `RwLock`. It was introduced in `03ffe4d595` on November 17, 2025 for non-blocking MCP startup: the session published an empty manager, startup initialized that same object while holding the write lock, and readers waited for initialization. `7cd2e84026` on February 19, 2026 removed that two-phase initialization in favor of constructing a fresh manager and swapping it in, explicitly noting that `Option` or `OnceCell` could replace the placeholder design. Hot reload later reused the existing lock to publish a replacement, but I found no indication that the lock was introduced to guarantee in-flight tool calls finish before refresh or shutdown. Terminal shutdown remains separate from refresh: it aborts startup prewarming and active tasks before shutting down the current manager, so tool calls may be interrupted and no model WebSocket work continues after shutdown. Focused regression coverage exercises pending tool-list cancellation, deferred refresh, and startup-prewarm shutdown.
101 lines
4.4 KiB
Rust
101 lines
4.4 KiB
Rust
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
|
|
use crate::SkillsManager;
|
|
use crate::agent::AgentControl;
|
|
use crate::attestation::AttestationProvider;
|
|
use crate::client::ModelClient;
|
|
use crate::config::NetworkProxyAuditMetadata;
|
|
use crate::config::StartedNetworkProxy;
|
|
use crate::exec_policy::ExecPolicyManager;
|
|
use crate::guardian::GuardianRejection;
|
|
use crate::guardian::GuardianRejectionCircuitBreaker;
|
|
use crate::mcp::McpManager;
|
|
use crate::tools::code_mode::CodeModeService;
|
|
use crate::tools::network_approval::NetworkApprovalService;
|
|
use crate::tools::sandboxing::ApprovalStore;
|
|
use crate::unified_exec::UnifiedExecProcessManager;
|
|
use anyhow::Result;
|
|
use arc_swap::ArcSwap;
|
|
use arc_swap::ArcSwapOption;
|
|
use codex_analytics::AnalyticsEventsClient;
|
|
use codex_core_plugins::PluginsManager;
|
|
use codex_exec_server::EnvironmentManager;
|
|
use codex_extension_api::ExtensionData;
|
|
use codex_extension_api::ExtensionRegistry;
|
|
use codex_hooks::Hooks;
|
|
use codex_login::AuthManager;
|
|
use codex_mcp::McpConnectionManager;
|
|
use codex_models_manager::manager::SharedModelsManager;
|
|
use codex_otel::SessionTelemetry;
|
|
use codex_rollout::state_db::StateDbHandle;
|
|
use codex_rollout_trace::ThreadTraceContext;
|
|
use codex_thread_store::LiveThread;
|
|
use codex_thread_store::ThreadStore;
|
|
use std::path::PathBuf;
|
|
use tokio::runtime::Handle;
|
|
use tokio::sync::Mutex;
|
|
use tokio::sync::watch;
|
|
use tokio_util::sync::CancellationToken;
|
|
|
|
pub(crate) struct SessionServices {
|
|
/// The latest manager; callers retain an owned handle while performing MCP I/O.
|
|
pub(crate) mcp_connection_manager: ArcSwap<McpConnectionManager>,
|
|
pub(crate) mcp_startup_cancellation_token: Mutex<CancellationToken>,
|
|
pub(crate) unified_exec_manager: UnifiedExecProcessManager,
|
|
#[cfg_attr(not(unix), allow(dead_code))]
|
|
pub(crate) shell_zsh_path: Option<PathBuf>,
|
|
#[cfg_attr(not(unix), allow(dead_code))]
|
|
pub(crate) main_execve_wrapper_exe: Option<PathBuf>,
|
|
pub(crate) analytics_events_client: AnalyticsEventsClient,
|
|
pub(crate) hooks: ArcSwap<Hooks>,
|
|
pub(crate) rollout_thread_trace: ThreadTraceContext,
|
|
pub(crate) user_shell: Arc<crate::shell::Shell>,
|
|
pub(crate) shell_snapshot_tx: watch::Sender<Option<Arc<crate::shell_snapshot::ShellSnapshot>>>,
|
|
pub(crate) show_raw_agent_reasoning: bool,
|
|
pub(crate) exec_policy: Arc<ExecPolicyManager>,
|
|
pub(crate) auth_manager: Arc<AuthManager>,
|
|
pub(crate) models_manager: SharedModelsManager,
|
|
pub(crate) session_telemetry: SessionTelemetry,
|
|
pub(crate) tool_approvals: Mutex<ApprovalStore>,
|
|
pub(crate) guardian_rejections: Mutex<HashMap<String, GuardianRejection>>,
|
|
pub(crate) guardian_rejection_circuit_breaker: Mutex<GuardianRejectionCircuitBreaker>,
|
|
pub(crate) runtime_handle: Handle,
|
|
pub(crate) skills_manager: Arc<SkillsManager>,
|
|
pub(crate) plugins_manager: Arc<PluginsManager>,
|
|
pub(crate) mcp_manager: Arc<McpManager>,
|
|
pub(crate) extensions: Arc<ExtensionRegistry<crate::config::Config>>,
|
|
pub(crate) session_extension_data: ExtensionData,
|
|
pub(crate) thread_extension_data: ExtensionData,
|
|
pub(crate) agent_control: AgentControl,
|
|
pub(crate) network_proxy: ArcSwapOption<StartedNetworkProxy>,
|
|
pub(crate) network_proxy_audit_metadata: NetworkProxyAuditMetadata,
|
|
pub(crate) managed_network_requirements_configured: bool,
|
|
pub(crate) network_approval: Arc<NetworkApprovalService>,
|
|
pub(crate) state_db: Option<StateDbHandle>,
|
|
pub(crate) live_thread: Option<LiveThread>,
|
|
pub(crate) thread_store: Arc<dyn ThreadStore>,
|
|
pub(crate) attestation_provider: Option<Arc<dyn AttestationProvider>>,
|
|
/// Session-scoped model client shared across turns.
|
|
pub(crate) model_client: ModelClient,
|
|
pub(crate) code_mode_service: CodeModeService,
|
|
/// Shared process-level environment registry. Sessions carry an `Arc` handle so they can pass
|
|
/// the same manager through child-thread spawn paths without reconstructing it.
|
|
pub(crate) environment_manager: Arc<EnvironmentManager>,
|
|
}
|
|
|
|
impl SessionServices {
|
|
/// Installs the manager before validating required servers so startup-time elicitation can
|
|
/// resolve through the session's manager while validation waits.
|
|
pub(crate) async fn install_mcp_connection_manager(
|
|
&self,
|
|
manager: McpConnectionManager,
|
|
) -> Result<()> {
|
|
self.mcp_connection_manager.store(Arc::new(manager));
|
|
self.mcp_connection_manager
|
|
.load_full()
|
|
.validate_required_servers()
|
|
.await
|
|
}
|
|
}
|