Stop doing action for handles and ignore responses after shutdown

lsp/src/Language/LSP/Server/Core.hs

@@ -1024,4 +1024,8 @@ be cancelled when we receive `shutdown`.
 Shutdown is a request, and the client won't send `exit` until a server responds, so if you
 want to be sure that some cleanup happens, you need to ensure we don't respond to `shutdown`
 until it's done. The best way to do this is just to install a specific `shutdown` handler.
+
+After the `shutdown` request, we don't handle any more requests and notifications other than
+`exit`. We also don't handle any more responses to requests we have sent but just throw the 
+responses away. 
 -}

lsp/src/Language/LSP/Server/Processing.hs

@@ -92,6 +92,7 @@ instance Pretty LspProcessingLog where
 processMessage :: (m ~ LspM config) => LogAction m (WithSeverity LspProcessingLog) -> BSL.ByteString -> m ()
 processMessage logger jsonStr = do
   pendingResponsesVar <- LspT $ asks $ resPendingResponses . resState
+  shutdown <- isShuttingDown
   join $ liftIO $ atomically $ fmap handleErrors $ runExceptT $ do
     val <- except $ eitherDecode jsonStr
     pending <- lift $ readTVar pendingResponsesVar
@@ -100,8 +101,10 @@ processMessage logger jsonStr = do
       FromClientMess m mess ->
         pure $ handle logger m mess
       FromClientRsp (P.Pair (ServerResponseCallback f) (Const !newMap)) res -> do
+        -- see Note [Shutdown]
         writeTVar pendingResponsesVar newMap
-        pure $ liftIO $ f (res ^. L.result)
+        unless shutdown <$> do
+          pure $ liftIO $ f (res ^. L.result)
  where
   parser :: ResponseMap -> Value -> Parser (FromClientMessage' (P.Product ServerResponseCallback (Const ResponseMap)))
   parser rm = parseClientMessage $ \i ->
@@ -449,31 +452,35 @@ handle' ::
   TClientMessage meth ->
   m ()
 handle' logger mAction m msg = do
-  maybe (return ()) (\f -> f msg) mAction
+  shutdown <- isShuttingDown
+  -- These are the methods that we are allowed to process during shutdown.
+  -- The reason that we include 'shutdown' itself here is because we change 
+  -- state to "shutting down" just _before_ we process the message itself!
+  let allowedMethod m = case (splitClientMethod m, m) of
+        (IsClientNot, SMethod_Exit) -> True
+        -- we do not allow SMethod_Shutdown to be called during shutdown
+        _ -> False
+
+  case mAction of
+    Just f | not shutdown || allowedMethod m -> f msg
+    _ -> pure ()
 
   dynReqHandlers <- getsState resRegistrationsReq
   dynNotHandlers <- getsState resRegistrationsNot
 
   env <- getLspEnv
   let Handlers{reqHandlers, notHandlers} = resHandlers env
-  shutdown <- isShuttingDown
 
   case splitClientMethod m of
     -- See Note [Shutdown]
     IsClientNot | shutdown, not (allowedMethod m) -> notificationDuringShutdown
-     where
-      allowedMethod SMethod_Exit = True
-      allowedMethod _ = False
     IsClientNot -> case pickHandler dynNotHandlers notHandlers of
       Just h -> liftIO $ h msg
       Nothing
         | SMethod_Exit <- m -> exitNotificationHandler logger msg
         | otherwise -> missingNotificationHandler
     -- See Note [Shutdown]
     IsClientReq | shutdown, not (allowedMethod m) -> requestDuringShutdown msg
-     where
-      allowedMethod SMethod_Shutdown = True
-      allowedMethod _ = False
     IsClientReq -> case pickHandler dynReqHandlers reqHandlers of
       Just h -> liftIO $ h msg (runLspT env . sendResponse msg)
       Nothing