From fc6b03c8daf203fd323a74ab9eb290a61013ae5b Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Thu, 20 Feb 2025 16:40:42 +0530
Subject: [PATCH] Handle requests received after shutdown message (#16262)

## Summary

This PR should help in
https://github.com/astral-sh/ruff-vscode/issues/676.

There are two issues that this is trying to fix all related to the way
shutdown should happen as per the protocol:
1. After the server handled the [shutdown
request](https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#shutdown)
and while waiting for the exit notification:

> If a server receives requests after a shutdown request those requests
should error with `InvalidRequest`.

But, we raised an error and exited. This PR fixes it by entering a loop
which responds to any request during this period with `InvalidRequest`

2. If the server received an [exit
notification](https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#exit)
but the shutdown request was never received, the server handled that by
logging and exiting with success but as per the spec:

> The server should exit with success code 0 if the shutdown request has
been received before; otherwise with error code 1.

    So, this PR fixes that as well by raising an error in this case.

## Test Plan

I'm not sure how to go about testing this without using a mock server.
---
 crates/ruff_server/src/server/connection.rs | 37 ++++++++++++++++-----
 1 file changed, 29 insertions(+), 8 deletions(-)

diff --git a/crates/ruff_server/src/server/connection.rs b/crates/ruff_server/src/server/connection.rs
index c04567c57a..5993023be8 100644
--- a/crates/ruff_server/src/server/connection.rs
+++ b/crates/ruff_server/src/server/connection.rs
@@ -91,19 +91,40 @@ impl Connection {
                 self.sender
                     .send(lsp::Response::new_ok(id.clone(), ()).into())?;
                 tracing::info!("Shutdown request received. Waiting for an exit notification...");
-                match self.receiver.recv_timeout(std::time::Duration::from_secs(30))? {
-                    lsp::Message::Notification(lsp::Notification { method, .. }) if method == lsp_types::notification::Exit::METHOD => {
-                        tracing::info!("Exit notification received. Server shutting down...");
-                        Ok(true)
-                    },
-                    message => anyhow::bail!("Server received unexpected message {message:?} while waiting for exit notification")
+
+                loop {
+                    match &self
+                        .receiver
+                        .recv_timeout(std::time::Duration::from_secs(30))?
+                    {
+                        lsp::Message::Notification(lsp::Notification { method, .. })
+                            if method == lsp_types::notification::Exit::METHOD =>
+                        {
+                            tracing::info!("Exit notification received. Server shutting down...");
+                            return Ok(true);
+                        }
+                        lsp::Message::Request(lsp::Request { id, method, .. }) => {
+                            tracing::warn!(
+                                "Server received unexpected request {method} ({id}) while waiting for exit notification",
+                            );
+                            self.sender.send(lsp::Message::Response(lsp::Response::new_err(
+                                id.clone(),
+                                lsp::ErrorCode::InvalidRequest as i32,
+                                "Server received unexpected request while waiting for exit notification".to_string(),
+                            )))?;
+                        }
+                        message => {
+                            tracing::warn!(
+                                "Server received unexpected message while waiting for exit notification: {message:?}"
+                            );
+                        }
+                    }
                 }
             }
             lsp::Message::Notification(lsp::Notification { method, .. })
                 if method == lsp_types::notification::Exit::METHOD =>
             {
-                tracing::error!("Server received an exit notification before a shutdown request was sent. Exiting...");
-                Ok(true)
+                anyhow::bail!("Server received an exit notification before a shutdown request was sent. Exiting...");
             }
             _ => Ok(false),
         }