diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 580282d5d3..32783d718a 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -691,7 +691,6 @@ pub async fn run_main_with_transport( warn!("dropping request from unknown connection: {connection_id:?}"); continue; }; - let was_initialized = connection_state.session.initialized; processor .process_request( connection_id, @@ -720,9 +719,6 @@ pub async fn run_main_with_transport( connection_state.session.experimental_api_enabled, std::sync::atomic::Ordering::Release, ); - if !was_initialized && connection_state.session.initialized { - processor.send_initialize_notifications().await; - } } JSONRPCMessage::Response(response) => { if !connections.contains_key(&connection_id) { @@ -732,11 +728,13 @@ pub async fn run_main_with_transport( processor.process_response(response).await; } JSONRPCMessage::Notification(notification) => { - if !connections.contains_key(&connection_id) { + let Some(connection_state) = connections.get_mut(&connection_id) else { warn!("dropping notification from unknown connection: {connection_id:?}"); continue; - } - processor.process_notification(notification).await; + }; + processor + .process_notification(notification, &mut connection_state.session) + .await; } JSONRPCMessage::Error(err) => { if !connections.contains_key(&connection_id) { diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index f56e12467f..9b3e8e19b5 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -141,6 +141,7 @@ pub(crate) struct MessageProcessor { #[derive(Clone, Debug, Default)] pub(crate) struct ConnectionSessionState { pub(crate) initialized: bool, + pub(crate) initialize_notifications_sent: bool, pub(crate) experimental_api_enabled: bool, pub(crate) opted_out_notification_methods: HashSet, pub(crate) app_server_client_name: Option, @@ -457,9 +458,21 @@ impl MessageProcessor { .await; } - pub(crate) async fn process_notification(&self, notification: JSONRPCNotification) { - // Currently, we do not expect to receive any notifications from the - // client, so we just log them. + pub(crate) async fn process_notification( + &self, + notification: JSONRPCNotification, + session: &mut ConnectionSessionState, + ) { + if notification.method == "initialized" { + if session.initialized && !session.initialize_notifications_sent { + self.send_initialize_notifications().await; + session.initialize_notifications_sent = true; + } + return; + } + + // Currently, we do not expect to receive any other notifications from + // the client, so we just log them. tracing::info!("<- notification: {:?}", notification); } diff --git a/codex-rs/app-server/tests/common/mcp_process.rs b/codex-rs/app-server/tests/common/mcp_process.rs index 753816b8d2..ab9b65f62b 100644 --- a/codex-rs/app-server/tests/common/mcp_process.rs +++ b/codex-rs/app-server/tests/common/mcp_process.rs @@ -178,21 +178,49 @@ impl McpProcess { .await } + /// Sends `initialize` and returns the server response without sending the + /// follow-up client `initialized` notification. Tests use this to observe + /// notifications that must be deferred until the handshake is fully + /// completed. + pub async fn initialize_without_initialized_notification( + &mut self, + ) -> anyhow::Result { + self.initialize_with_params( + InitializeParams { + client_info: ClientInfo { + name: DEFAULT_CLIENT_NAME.to_string(), + title: None, + version: "0.1.0".to_string(), + }, + capabilities: Some(InitializeCapabilities { + experimental_api: true, + ..Default::default() + }), + }, + false, + ) + .await + } + pub async fn initialize_with_capabilities( &mut self, client_info: ClientInfo, capabilities: Option, ) -> anyhow::Result { - self.initialize_with_params(InitializeParams { - client_info, - capabilities, - }) + self.initialize_with_params( + InitializeParams { + client_info, + capabilities, + }, + true, + ) .await } async fn initialize_with_params( &mut self, params: InitializeParams, + send_initialized_notification: bool, ) -> anyhow::Result { let params = Some(serde_json::to_value(params)?); let request_id = self.send_request("initialize", params).await?; @@ -207,9 +235,10 @@ impl McpProcess { ); } - // Send notifications/initialized to ack the response. - self.send_notification(ClientNotification::Initialized) - .await?; + if send_initialized_notification { + self.send_notification(ClientNotification::Initialized) + .await?; + } Ok(JSONRPCMessage::Response(response)) } diff --git a/codex-rs/app-server/tests/suite/v2/initialize.rs b/codex-rs/app-server/tests/suite/v2/initialize.rs index de437c5bf5..2f0c03d1a6 100644 --- a/codex-rs/app-server/tests/suite/v2/initialize.rs +++ b/codex-rs/app-server/tests/suite/v2/initialize.rs @@ -4,6 +4,8 @@ use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence_unchecked; use app_test_support::to_response; use codex_app_server_protocol::ClientInfo; +use codex_app_server_protocol::ClientNotification; +use codex_app_server_protocol::ConfigWarningNotification; use codex_app_server_protocol::InitializeCapabilities; use codex_app_server_protocol::InitializeResponse; use codex_app_server_protocol::JSONRPCMessage; @@ -186,6 +188,81 @@ async fn initialize_opt_out_notification_methods_filters_notifications() -> Resu Ok(()) } +#[tokio::test] +async fn initialize_defers_config_warning_until_initialized_notification() -> Result<()> { + let responses = Vec::new(); + let server = create_mock_responses_server_sequence_unchecked(responses).await; + let codex_home = TempDir::new()?; + create_config_toml_with_extra( + codex_home.path(), + &server.uri(), + "never", + r#" +[mcp_servers.figma] +url = "https://mcp.figma.com/mcp" + +[mcp_servers.figma] +url = "https://mcp.figma.com/mcp" +"#, + )?; + let mut mcp = McpProcess::new(codex_home.path()).await?; + + let message = timeout( + DEFAULT_READ_TIMEOUT, + mcp.initialize_without_initialized_notification(), + ) + .await??; + let JSONRPCMessage::Response(response) = message else { + anyhow::bail!("expected initialize response, got {message:?}"); + }; + let _: InitializeResponse = to_response(response)?; + + let early_warning = timeout( + std::time::Duration::from_millis(250), + mcp.read_stream_until_notification_message("configWarning"), + ) + .await; + assert!( + early_warning.is_err(), + "configWarning should not be emitted before the client sends initialized" + ); + + mcp.send_notification(ClientNotification::Initialized) + .await?; + + let warning = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("configWarning"), + ) + .await??; + let payload = warning + .params + .expect("configWarning params should be present"); + let warning: ConfigWarningNotification = serde_json::from_value(payload)?; + + assert_eq!(warning.summary, "Invalid configuration; using defaults."); + assert!( + warning + .details + .as_deref() + .is_some_and(|details| details.contains("duplicate")), + "expected duplicate-table parse error, got {:?}", + warning.details + ); + assert!( + warning + .path + .as_deref() + .is_some_and(|path| path.ends_with("config.toml")) + ); + assert!( + warning.range.is_some(), + "config warning should include a text range" + ); + + Ok(()) +} + #[tokio::test] async fn turn_start_notify_payload_includes_initialize_client_name() -> Result<()> { let responses = vec![create_final_assistant_message_sse_response("Done")?];