From 0d86844eb443e4966f2fa955e4c04b1fbd23911c Mon Sep 17 00:00:00 2001
From: Stephen Macke <stephen.macke@gmail.com>
Date: Fri, 8 Sep 2023 07:50:16 -0700
Subject: [PATCH] Use invalid request handler rather than raising key error for
 requests after shutdown (#432)

---
 pylsp/lsp.py        | 18 ++++++++++++++++++
 pylsp/python_lsp.py | 10 +++++++++-
 pyproject.toml      |  2 +-
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/pylsp/lsp.py b/pylsp/lsp.py
index d5f4737b..7b3f02ee 100644
--- a/pylsp/lsp.py
+++ b/pylsp/lsp.py
@@ -95,3 +95,21 @@ class TextDocumentSyncKind:
 class NotebookCellKind:
     Markup = 1
     Code = 2
+
+
+# https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#errorCodes
+class ErrorCodes:
+    ParseError = -32700
+    InvalidRequest = -32600
+    MethodNotFound = -32601
+    InvalidParams = -32602
+    InternalError = -32603
+    jsonrpcReservedErrorRangeStart = -32099
+    ServerNotInitialized = -32002
+    UnknownErrorCode = -32001
+    jsonrpcReservedErrorRangeEnd = -32000
+    lspReservedErrorRangeStart = -32899
+    ServerCancelled = -32802
+    ContentModified = -32801
+    RequestCancelled = -32800
+    lspReservedErrorRangeEnd = -32800
diff --git a/pylsp/python_lsp.py b/pylsp/python_lsp.py
index 63d5478d..3eeb2f7f 100644
--- a/pylsp/python_lsp.py
+++ b/pylsp/python_lsp.py
@@ -215,7 +215,7 @@ def __getitem__(self, item):
         if self._shutdown and item != "exit":
             # exit is the only allowed method during shutdown
             log.debug("Ignoring non-exit method during shutdown: %s", item)
-            raise KeyError
+            item = "invalid_request_after_shutdown"
 
         try:
             return super().__getitem__(item)
@@ -234,6 +234,14 @@ def m_shutdown(self, **_kwargs):
             workspace.close()
         self._shutdown = True
 
+    def m_invalid_request_after_shutdown(self, **_kwargs):
+        return {
+            "error": {
+                "code": lsp.ErrorCodes.InvalidRequest,
+                "message": "Requests after shutdown are not valid",
+            }
+        }
+
     def m_exit(self, **_kwargs):
         self._endpoint.shutdown()
         if self._jsonrpc_stream_reader is not None:
diff --git a/pyproject.toml b/pyproject.toml
index d0d7df80..a27c1316 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -17,7 +17,7 @@ dependencies = [
     "importlib_metadata>=4.8.3;python_version<\"3.10\"",
     "jedi>=0.17.2,<0.20.0",
     "pluggy>=1.0.0",
-    "python-lsp-jsonrpc>=1.0.0",
+    "python-lsp-jsonrpc>=1.1.0,<2.0.0",
     "ujson>=3.0.0",
 ]
 dynamic = ["version"]