codex/codex-rs/app-server/src/request_processors/initialize_processor.rs

use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;

use axum::http::HeaderValue;
use codex_analytics::AppServerRpcTransport;
use codex_login::default_client::SetOriginatorError;
use codex_login::default_client::USER_AGENT_SUFFIX;
use codex_login::default_client::get_codex_user_agent;
use codex_login::default_client::set_default_client_residency_requirement;
use codex_login::default_client::set_default_originator;
use codex_protocol::mcp::ClientMcpExtensions;
use codex_protocol::mcp::OPENAI_ELICITATION_EXTENSION_ID;

use super::*;
use crate::message_processor::ConnectionSessionState;
use crate::message_processor::InitializedConnectionSessionState;
use crate::transport::ConnectionOrigin;

const NON_ORIGINATING_CLIENT_NAMES: &[&str] = &["codex_app_server_daemon", "codex-backend"];

#[derive(Clone)]
pub(crate) struct InitializeRequestProcessor {
    gateway_login_control: Arc<codex_login::GatewayLoginControl>,
    gateway_login_initialized: Arc<std::sync::OnceLock<()>>,
    outgoing: Arc<OutgoingMessageSender>,
    analytics_events_client: AnalyticsEventsClient,
    config: Arc<Config>,
    config_warnings: Arc<Vec<ConfigWarningNotification>>,
    rpc_transport: AppServerRpcTransport,
    user_verification: Arc<crate::user_verification::Service>,
}

impl InitializeRequestProcessor {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn new(
        gateway_login_control: Arc<codex_login::GatewayLoginControl>,
        outgoing: Arc<OutgoingMessageSender>,
        analytics_events_client: AnalyticsEventsClient,
        config: Arc<Config>,
        config_warnings: Vec<ConfigWarningNotification>,
        rpc_transport: AppServerRpcTransport,
        user_verification: Arc<crate::user_verification::Service>,
    ) -> Self {
        Self {
            gateway_login_control,
            gateway_login_initialized: Arc::new(std::sync::OnceLock::new()),
            outgoing,
            analytics_events_client,
            config,
            config_warnings: Arc::new(config_warnings),
            rpc_transport,
            user_verification,
        }
    }

    pub(crate) async fn initialize(
        &self,
        connection_id: ConnectionId,
        request_id: RequestId,
        params: InitializeParams,
        session: &ConnectionSessionState,
        // `Some(...)` means the caller wants initialize to immediately mark the
        // connection outbound-ready. Websocket JSON-RPC calls pass `None` so
        // lib.rs can deliver connection-scoped initialize notifications first.
        outbound_initialized: Option<&AtomicBool>,
    ) -> Result<bool, JSONRPCErrorError> {
        let connection_request_id = ConnectionRequestId {
            connection_id,
            request_id,
        };
        if session.initialized() {
            return Err(invalid_request("Already initialized"));
        }

        // TODO(maxj): Revisit capability scoping for `experimental_api_enabled`.
        // Current behavior is per-connection. Reviewer feedback notes this can
        // create odd cross-client behavior (for example dynamic tool calls on a
        // shared thread when another connected client did not opt into
        // experimental API). Proposed direction is instance-global first-write-wins
        // with initialize-time mismatch rejection.
        let analytics_initialize_params = params.clone();
        let capabilities = params.capabilities.unwrap_or_default();
        let experimental_api_enabled = capabilities.experimental_api;
        let request_attestation = capabilities.request_attestation;
        let extensions = capabilities.extensions.as_ref();
        let mut client_mcp_extensions = codex_mcp::client_mcp_extensions(
            extensions,
            capabilities.mcp_server_openai_form_elicitation,
        );
        let opt_out_notification_methods = capabilities
            .opt_out_notification_methods
            .unwrap_or_default();
        let ClientInfo {
            name,
            title: _title,
            version,
        } = params.client_info;
        // Validate before committing; set_default_originator validates while
        // mutating process-global metadata.
        if HeaderValue::from_str(&name).is_err() {
            return Err(invalid_request(format!(
                "Invalid clientInfo.name: '{name}'. Must be a valid HTTP header value."
            )));
        }
        // Activate only the embedded TUI and local desktop host. Client-supplied
        // extensions cannot opt other hosts into verification.
        let user_verification_enabled = experimental_api_enabled
            && matches!(
                (session.origin, name.as_str()),
                (ConnectionOrigin::InProcess, "codex-tui")
                    | (ConnectionOrigin::Stdio, "Codex Desktop")
            )
            && tokio::task::spawn_blocking(self.user_verification.device_supported)
                .await
                .unwrap_or(false);
        if user_verification_enabled {
            let mut extensions = client_mcp_extensions
                .iter()
                .map(|(id, value)| (id.to_string(), value.clone()))
                .collect::<std::collections::HashMap<_, _>>();
            let settings = extensions
                .entry(OPENAI_ELICITATION_EXTENSION_ID.to_string())
                .or_insert_with(|| serde_json::json!({}));
            if !settings.is_object() {
                *settings = serde_json::json!({});
            }
            settings["userVerification"] = serde_json::json!({});
            client_mcp_extensions = ClientMcpExtensions::new(extensions);
        }
        let originator = name.clone();
        let user_agent_suffix = format!("{name}; {version}");
        let mutates_global_identity = !NON_ORIGINATING_CLIENT_NAMES.contains(&name.as_str());
        let codex_home = self.config.codex_home.clone();
        if session
            .initialize(InitializedConnectionSessionState {
                experimental_api_enabled,
                opted_out_notification_methods: opt_out_notification_methods.into_iter().collect(),
                app_server_client_name: name.clone(),
                client_version: version,
                request_attestation,
                client_mcp_extensions,
            })
            .is_err()
        {
            return Err(invalid_request("Already initialized"));
        }
        if user_verification_enabled {
            self.outgoing
                .enable_user_verification_connection(connection_id)
                .await;
        }

        if capabilities.explicit_gateway_oauth || mutates_global_identity {
            // Only the first originating client may restore legacy automatic login.
            // Any explicit opt-in is sticky across subsequent connections.
            self.gateway_login_initialized.get_or_init(|| {
                if !capabilities.explicit_gateway_oauth {
                    self.gateway_login_control.allow_automatic_login();
                }
            });
            if capabilities.explicit_gateway_oauth {
                self.gateway_login_control.require_explicit_login();
            }
        }

        if mutates_global_identity {
            // Only real client initialization may mutate process-global client metadata.
            if let Err(error) = set_default_originator(originator.clone()) {
                match error {
                    SetOriginatorError::InvalidHeaderValue => {
                        tracing::warn!(
                            client_info_name = %name,
                            "validated clientInfo.name was rejected while setting originator"
                        );
                    }
                    SetOriginatorError::AlreadyInitialized => {
                        // No-op. This is expected to happen if the originator is already set via env var.
                        // TODO(owen): Once we remove support for CODEX_INTERNAL_ORIGINATOR_OVERRIDE,
                        // this will be an unexpected state and we can return a JSON-RPC error indicating
                        // internal server error.
                    }
                }
            }
        }
        self.analytics_events_client.track_initialize(
            connection_id.0,
            analytics_initialize_params,
            originator,
            self.rpc_transport,
        );
        set_default_client_residency_requirement(self.config.enforce_residency.value());
        if mutates_global_identity && let Ok(mut suffix) = USER_AGENT_SUFFIX.lock() {
            *suffix = Some(user_agent_suffix);
        }

        #[cfg(windows)]
        if matches!(session.origin, ConnectionOrigin::Stdio) && name == "Codex Desktop" {
            // Uninstall ownership must not depend on account sign-in or sandbox setup.
            // Keep this bounded attempt ahead of the response; background registration can race uninstall.
            let home = codex_home.clone();
            if !matches!(
                tokio::task::spawn_blocking(move || {
                    codex_windows_sandbox::register_desktop_installation(&home)
                })
                .await,
                Ok(Ok(()))
            ) {
                tracing::warn!("could not register desktop uninstall ownership");
            }
        }

        let user_agent = get_codex_user_agent();
        let response = InitializeResponse {
            user_agent,
            codex_home,
            platform_family: std::env::consts::FAMILY.to_string(),
            platform_os: std::env::consts::OS.to_string(),
        };

        self.outgoing
            .send_response(connection_request_id, response)
            .await;

        if let Some(outbound_initialized) = outbound_initialized {
            outbound_initialized.store(true, Ordering::Release);
            return Ok(true);
        }

        Ok(false)
    }

    pub(crate) async fn send_initialize_notifications_to_connection(
        &self,
        connection_id: ConnectionId,
    ) {
        for notification in self.config_warnings.iter().cloned() {
            self.outgoing
                .send_server_notification_to_connections(
                    &[connection_id],
                    ServerNotification::ConfigWarning(notification),
                )
                .await;
        }
    }

    pub(crate) async fn send_initialize_notifications(&self) {
        for notification in self.config_warnings.iter().cloned() {
            self.outgoing
                .send_server_notification(ServerNotification::ConfigWarning(notification))
                .await;
        }
    }

    pub(crate) fn track_initialized_request(
        &self,
        connection_id: ConnectionId,
        request_id: RequestId,
        request: &ClientRequest,
    ) {
        self.analytics_events_client
            .track_request(connection_id.0, request_id, request);
    }
}