codex/codex-rs/core/src/state/service.rs

use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;

use crate::SkillsService;
use crate::agent::AgentControl;
use crate::agents_md_manager::AgentsMdManager;
use crate::attestation::AttestationProvider;
use crate::client::ModelClient;
use crate::config::NetworkProxyAuditMetadata;
use crate::config::StartedNetworkProxy;
use crate::current_time::TimeProvider;
use crate::elicitation::ElicitationService;
use crate::environment_selection::ThreadEnvironments;
use crate::exec_policy::ExecPolicyManager;
use crate::guardian::GuardianRejection;
use crate::guardian::GuardianRejectionCircuitBreaker;
use crate::mcp::McpManager;
use crate::model_provider_runtime::ModelProviderRuntimeSnapshot;
use crate::model_provider_runtime::ModelProviderRuntimeSource;
use crate::session::McpRuntimeSnapshot;
use crate::tools::code_mode::CodeModeService;
use crate::tools::handlers::ToolSearchHandlerCache;
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_extension_api::ExtensionData;
use codex_extension_api::ExtensionDataInit;
use codex_extension_api::ExtensionRegistry;
use codex_hooks::Hooks;
use codex_login::AuthManager;
use codex_mcp::McpConfig;
use codex_mcp::McpConnectionManager;
use codex_mcp::McpRuntimeContext;
use codex_models_manager::manager::SharedModelsManager;
use codex_otel::SessionTelemetry;
use codex_protocol::capabilities::SelectedCapabilityRoot;
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_util::sync::CancellationToken;

pub(crate) struct SessionServices {
    /// Mirror of the latest manager for extension resource clients that predate runtime snapshots.
    pub(crate) mcp_connection_manager: Arc<ArcSwap<McpConnectionManager>>,
    /// The latest atomically published MCP config and manager pair.
    pub(crate) mcp_runtime: ArcSwapOption<McpRuntimeSnapshot>,
    /// Serializes environment-driven runtime rebuilds.
    pub(crate) mcp_projection_lock: Mutex<()>,
    pub(crate) mcp_startup_cancellation_token: Mutex<CancellationToken>,
    pub(crate) unified_exec_manager: UnifiedExecProcessManager,
    pub(crate) elicitations: ElicitationService,
    #[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) show_raw_agent_reasoning: bool,
    pub(crate) exec_policy: Arc<ExecPolicyManager>,
    pub(crate) auth_manager: Arc<AuthManager>,
    /// Atomically replaceable provider runtime used to build future turns. Each turn clones one
    /// immutable runtime so an in-flight turn cannot observe a provider refresh.
    pub(crate) model_runtime: ArcSwap<SessionModelRuntime>,
    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_service: Arc<SkillsService>,
    pub(crate) agents_md_manager: Arc<AgentsMdManager>,
    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) supports_openai_form_elicitation: AtomicBool,
    /// Raw capability selections for this thread. Each model step resolves them against its
    /// current executor environments before using them.
    pub(crate) selected_capability_roots: Vec<SelectedCapabilityRoot>,
    pub(crate) mcp_thread_init: ExtensionDataInit,
    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>>,
    pub(crate) time_provider: Arc<dyn TimeProvider>,
    /// Session-scoped model client shared across turns.
    pub(crate) model_client: ModelClient,
    pub(crate) code_mode_service: CodeModeService,
    pub(crate) tool_search_handler_cache: ToolSearchHandlerCache,
    pub(crate) turn_environments: Arc<ThreadEnvironments>,
}

#[derive(Clone)]
pub(crate) struct SessionModelRuntime {
    pub(crate) source: ModelProviderRuntimeSource,
    pub(crate) snapshot: Arc<ModelProviderRuntimeSnapshot>,
    pub(crate) model_client: ModelClient,
}

impl SessionModelRuntime {
    pub(crate) fn models_manager(&self) -> SharedModelsManager {
        self.snapshot.models_manager()
    }
}

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,
        config: Arc<McpConfig>,
        plugins_available: bool,
        runtime_context: McpRuntimeContext,
        available_environment_ids: Vec<String>,
        manager: McpConnectionManager,
    ) -> Result<()> {
        let runtime = self.publish_mcp_runtime(
            config,
            plugins_available,
            runtime_context,
            available_environment_ids,
            manager,
        );
        runtime.manager().validate_required_servers().await
    }

    pub(crate) fn publish_mcp_runtime(
        &self,
        config: Arc<McpConfig>,
        plugins_available: bool,
        runtime_context: McpRuntimeContext,
        available_environment_ids: Vec<String>,
        manager: McpConnectionManager,
    ) -> Arc<McpRuntimeSnapshot> {
        let manager = Arc::new(manager);
        // Publish the manager for legacy resource clients first. Once the paired snapshot is
        // visible, every model-scoped consumer observes this exact manager.
        self.mcp_connection_manager.store(Arc::clone(&manager));
        let runtime = Arc::new(McpRuntimeSnapshot::new(
            config,
            plugins_available,
            manager,
            runtime_context,
            available_environment_ids,
        ));
        self.mcp_runtime.store(Some(Arc::clone(&runtime)));
        runtime
    }

    pub(crate) fn latest_mcp_runtime(&self) -> Arc<McpRuntimeSnapshot> {
        let Some(runtime) = self.mcp_runtime.load_full() else {
            unreachable!("MCP runtime must be installed before handling requests");
        };
        runtime
    }
}