From 41c41dece1d97705ddc0d5897ce24f408b2efe2c Mon Sep 17 00:00:00 2001
From: Matthew Bauer <mjbauer95@gmail.com>
Date: Sun, 8 Feb 2026 17:02:08 -0500
Subject: [PATCH] Handle exit message from lsp
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

noticed that some profiling wasn’t getting cleanly output. i think
it’s because the client is sigkilling before static-lsp has a chance
to output profiling.
---
 app/Main.hs              |  2 +-
 src/StaticLS/Handlers.hs |  6 ++++++
 src/StaticLS/Server.hs   | 27 ++++++++++++++++++++++-----
 3 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/app/Main.hs b/app/Main.hs
index 95349cd5..2187c1f6 100644
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -18,5 +18,5 @@ main = do
     Success (App.GHCIDOptions {args}) -> ghcid args
     argsRes -> do
       staticEnvOpts <- App.handleParseResultWithSuppression jsonOrDefaultOpts argsRes
-      _ <- StaticLS.runServer staticEnvOpts logger
+      StaticLS.runServer staticEnvOpts logger
       pure ()
diff --git a/src/StaticLS/Handlers.hs b/src/StaticLS/Handlers.hs
index e99a88a9..6ebb051e 100644
--- a/src/StaticLS/Handlers.hs
+++ b/src/StaticLS/Handlers.hs
@@ -54,6 +54,7 @@ import System.Directory (doesFileExist)
 import System.FSNotify qualified as FSNotify
 import Text.Parsec.Text qualified as Parsec
 import UnliftIO.Exception qualified as Exception
+import Control.Concurrent.MVar
 
 -----------------------------------------------------------------
 --------------------- LSP event handlers ------------------------
@@ -352,3 +353,8 @@ clearDiagnostics = LSP.flushDiagnosticsBySource 100 (Just "haskell")
 
 testing :: (Show a) => [a] -> String
 testing = show
+
+handleExit :: MVar () -> Handlers (LspT c StaticLsM)
+handleExit clientMsgVar = LSP.notificationHandler LSP.SMethod_Exit $ \_ -> do
+  _ <- liftIO $ tryPutMVar clientMsgVar ()
+  pure ()
diff --git a/src/StaticLS/Server.hs b/src/StaticLS/Server.hs
index bd49f434..35a1086b 100644
--- a/src/StaticLS/Server.hs
+++ b/src/StaticLS/Server.hs
@@ -42,6 +42,8 @@ import System.FSNotify qualified as FSNotify
 import System.FilePath qualified as FilePath
 import UnliftIO.Concurrent qualified as Conc
 import UnliftIO.Exception qualified as Exception
+import Control.Concurrent.MVar
+import Control.Concurrent.Async
 
 type LspConfig = ()
 
@@ -168,8 +170,8 @@ initServer reactorChan staticEnvOptions logger serverConfig _ = do
       Just p -> Right p
 #endif
 
-serverDef :: StaticEnvOptions -> LoggerM IO -> IO (ServerDefinition ())
-serverDef options logger = do
+serverDef :: StaticEnvOptions -> LoggerM IO -> MVar () -> IO (ServerDefinition ())
+serverDef options logger clientMsgVar = do
   reactorChan <- liftIO Conc.newChan
   let
     -- TODO: actually respond to the client with an error
@@ -227,6 +229,7 @@ serverDef options logger = do
                 , handleDocumentSymbols
                 , handleCompletion
                 , handleCompletionItemResolve
+                , handleExit clientMsgVar
                 ]
                   <> (if options.provideInlays then [handleInlayHintRequest options, handleResolveInlayHint] else [])
                   <> ( case options.fourmoluCommand of
@@ -276,7 +279,21 @@ lspOptions =
     , LSP.optCompletionTriggerCharacters = Just ['.']
     }
 
-runServer :: StaticEnvOptions -> LoggerM IO -> IO Int
+runServer :: StaticEnvOptions -> LoggerM IO -> IO ()
 runServer argOptions logger = do
-  server <- serverDef argOptions logger
-  LSP.runServer server
+  -- This MVar becomes full when the server thread exits or we receive exit message from client.
+  -- LSP server will be canceled when it's full.
+  clientMsgVar <- newEmptyMVar
+
+  server <- serverDef argOptions logger clientMsgVar
+  untilMVar clientMsgVar $ do
+    _ <- LSP.runServer server
+    pure ()
+
+-- | Runs the action until it ends or until the given MVar is put.
+--   It is important, that the thread that puts the 'MVar' is not dropped before it puts the 'MVar' i.e. it should
+--   occur as the final action in a 'finally' or 'bracket', because otherwise this thread will finish early (as soon
+--   as the thread receives the BlockedIndefinitelyOnMVar exception)
+--   Rethrows any exceptions.
+untilMVar :: MVar () -> IO a -> IO ()
+untilMVar mvar io = race_ (readMVar mvar) io