X-Git-Url: http://git.lukelau.me/?a=blobdiff_plain;f=src%2FLanguage%2FHaskell%2FLSP%2FTest%2FSession.hs;h=1777f15f2750a913e812df81c7a42d7c6156fee1;hb=bd2ac7da01645dc111a6912be661bc13b7f9340f;hp=46155f0607cfa1b752b541b755a761b1551b4b30;hpb=a546e59018ab8dbf80d83442e9773f94f293593e;p=lsp-test.git

diff --git a/src/Language/Haskell/LSP/Test/Session.hs b/src/Language/Haskell/LSP/Test/Session.hs
index 46155f0..1777f15 100644
--- a/src/Language/Haskell/LSP/Test/Session.hs
+++ b/src/Language/Haskell/LSP/Test/Session.hs
@@ -190,9 +190,11 @@ runSessionWithHandles :: Handle -- ^ Server in
                       -> SessionConfig
                       -> ClientCapabilities
                       -> FilePath -- ^ Root directory
+                      -> Session () -- ^ To exit Server
                       -> Session a
                       -> IO a
-runSessionWithHandles serverIn serverOut serverHandler config caps rootDir session = do
+runSessionWithHandles serverIn serverOut serverHandler config caps rootDir exitServer session = do
+  
   absRootDir <- canonicalizePath rootDir
 
   hSetBuffering serverIn  NoBuffering
@@ -210,11 +212,13 @@ runSessionWithHandles serverIn serverOut serverHandler config caps rootDir sessi
 
   let context = SessionContext serverIn absRootDir messageChan reqMap initRsp config caps
       initState = SessionState (IdInt 0) mempty mempty 0 False Nothing
-      launchServerHandler = forkIO $ catch (serverHandler serverOut context)
-                                           (throwTo mainThreadId :: SessionException -> IO ())
-  (result, _) <- bracket launchServerHandler killThread $
-    const $ runSession context initState session
+      runSession' = runSession context initState
+      
+      errorHandler = throwTo mainThreadId :: SessionException -> IO()
+      serverLauncher = forkIO $ catch (serverHandler serverOut context) errorHandler
+      serverFinalizer tid = runSession' exitServer >> killThread tid
       
+  (result, _) <- bracket serverLauncher serverFinalizer (const $ runSession' session)
   return result
 
 updateStateC :: ConduitM FromServerMessage FromServerMessage (StateT SessionState (ReaderT SessionContext IO)) ()