fix: run first pre-init request in service loop (#1263)

crates/rmcp/tests/test_stateless_server_requests.rs

@@ -1,14 +1,18 @@
 #![cfg(all(feature = "server", not(feature = "local")))]
 
-use std::sync::{Arc, Mutex};
+use std::{
+    sync::{Arc, Mutex},
+    time::Duration,
+};
 
 use rmcp::{
     ServerHandler, ServiceExt,
     model::{
         ClientCapabilities, ClientJsonRpcMessage, ClientRequest, DiscoverRequest,
         DiscoverRequestParams, ErrorCode, ErrorData, Implementation, ListToolsRequest,
-        ListToolsResult, PaginatedRequestParams, ProtocolVersion, RequestId, RequestMetaObject,
-        ServerJsonRpcMessage,
+        ListToolsResult, NumberOrString, PaginatedRequestParams, ProgressNotificationParam,
+        ProgressToken, ProtocolVersion, RequestId, RequestMetaObject, ServerJsonRpcMessage,
+        ServerNotification,
     },
     service::{MaybeSendFuture, RequestContext, RoleServer, ServerInitializeError},
     transport::{IntoTransport, Transport},
@@ -212,3 +216,73 @@ async fn stateless_server_rejects_malformed_metadata_opener_with_error_response(
         ServerInitializeError::ExpectedInitializeRequest(Some(_))
     ));
 }
+
+#[derive(Clone)]
+struct ProgressServer;
+
+impl ServerHandler for ProgressServer {
+    async fn list_tools(
+        &self,
+        _request: Option<PaginatedRequestParams>,
+        context: RequestContext<RoleServer>,
+    ) -> Result<ListToolsResult, ErrorData> {
+        let progress_token = context
+            .meta
+            .get_progress_token()
+            .expect("progress token in request meta");
+        context
+            .peer
+            .notify_progress(ProgressNotificationParam::new(progress_token, 1.0))
+            .await
+            .expect("send progress notification");
+        Ok(ListToolsResult::default())
+    }
+}
+
+/// Regression test for issue #1261: the first request of an `initialize`-less
+/// session used to be handled outside the service loop, so a handler that sent
+/// a notification before returning waited forever for the loop to flush it.
+#[tokio::test]
+async fn stateless_server_first_request_handler_can_send_notifications() {
+    let (server_transport, client_transport) = tokio::io::duplex(4096);
+    let server_task = tokio::spawn(async move {
+        ProgressServer
+            .serve(server_transport)
+            .await
+            .expect("server should start")
+    });
+    let mut client = IntoTransport::<rmcp::RoleClient, _, _>::into_transport(client_transport);
+
+    let mut meta = complete_meta();
+    meta.set_progress_token(ProgressToken(NumberOrString::Number(7)));
+    client
+        .send(list_tools_request(meta))
+        .await
+        .expect("send first request");
+
+    let exchange = async {
+        let notification = match client.receive().await {
+            Some(ServerJsonRpcMessage::Notification(notification)) => notification.notification,
+            other => panic!("expected progress notification before the response, got {other:?}"),
+        };
+        assert!(
+            matches!(notification, ServerNotification::ProgressNotification(_)),
+            "expected progress notification, got {notification:?}"
+        );
+        let response = match client.receive().await {
+            Some(ServerJsonRpcMessage::Response(response)) => response,
+            other => panic!("expected list tools response, got {other:?}"),
+        };
+        assert_eq!(response.id, RequestId::Number(1));
+    };
+    tokio::time::timeout(Duration::from_secs(5), exchange)
+        .await
+        .expect("first request must not deadlock the server");
+
+    server_task
+        .await
+        .expect("server task")
+        .cancel()
+        .await
+        .expect("cancel server");
+}