From ad1d8fc387457dbcebcad90e0e6e6dbff62b6bba Mon Sep 17 00:00:00 2001
From: Zubin Duggal <zubin@cmi.ac.in>
Date: Mon, 25 Jan 2021 02:51:35 +0530
Subject: [PATCH] track progress sessions

---
 src/Language/LSP/Test.hs         |  6 ++++++
 src/Language/LSP/Test/Session.hs | 10 +++++++++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/src/Language/LSP/Test.hs b/src/Language/LSP/Test.hs
index 89abfcf..90000fd 100644
--- a/src/Language/LSP/Test.hs
+++ b/src/Language/LSP/Test.hs
@@ -62,6 +62,7 @@ module Language.LSP.Test
   , waitForDiagnosticsSource
   , noDiagnostics
   , getCurrentDiagnostics
+  , getIncompleteProgressSessions
   -- ** Commands
   , executeCommand
   -- ** Code Actions
@@ -101,6 +102,7 @@ import Control.Monad.IO.Class
 import Control.Exception
 import Control.Lens hiding ((.=), List, Empty)
 import qualified Data.Map.Strict as Map
+import qualified Data.Set as Set
 import qualified Data.Text as T
 import qualified Data.Text.IO as T
 import Data.Aeson
@@ -534,6 +536,10 @@ getCodeActionContext doc = do
 getCurrentDiagnostics :: TextDocumentIdentifier -> Session [Diagnostic]
 getCurrentDiagnostics doc = fromMaybe [] . Map.lookup (toNormalizedUri $ doc ^. uri) . curDiagnostics <$> get
 
+-- | Returns the tokens of all progress sessions that have started but not yet ended.
+getIncompleteProgressSessions :: Session (Set.Set ProgressToken)
+getIncompleteProgressSessions = curProgressSessions <$> get
+
 -- | Executes a command.
 executeCommand :: Command -> Session ()
 executeCommand cmd = do
diff --git a/src/Language/LSP/Test/Session.hs b/src/Language/LSP/Test/Session.hs
index ac668e6..4e3b177 100644
--- a/src/Language/LSP/Test/Session.hs
+++ b/src/Language/LSP/Test/Session.hs
@@ -57,6 +57,7 @@ import Data.Default
 import Data.Foldable
 import Data.List
 import qualified Data.Map as Map
+import qualified Data.Set as Set
 import qualified Data.Text as T
 import qualified Data.Text.IO as T
 import qualified Data.HashMap.Strict as HashMap
@@ -175,6 +176,7 @@ data SessionState = SessionState
   , curDynCaps :: Map.Map T.Text SomeRegistration
   -- ^ The capabilities that the server has dynamically registered with us so
   -- far
+  , curProgressSessions :: Set.Set ProgressToken
   }
 
 class Monad m => HasState s m where
@@ -265,7 +267,7 @@ runSession' serverIn serverOut mServerProc serverHandler config caps rootDir exi
   mainThreadId <- myThreadId
 
   let context = SessionContext serverIn absRootDir messageChan timeoutIdVar reqMap initRsp config caps
-      initState vfs = SessionState 0 vfs mempty False Nothing mempty
+      initState vfs = SessionState 0 vfs mempty False Nothing mempty mempty
       runSession' ses = initVFS $ \vfs -> runSessionMonad context (initState vfs) ses
 
       errorHandler = throwTo mainThreadId :: SessionException -> IO ()
@@ -307,6 +309,12 @@ documentChangeUri (InR (InR (InR x))) = x ^. uri
 
 updateState :: (MonadIO m, HasReader SessionContext m, HasState SessionState m)
             => FromServerMessage -> m ()
+updateState (FromServerMess SProgress req) = case req ^. params . value of
+  Begin _ ->
+    modify $ \s -> s { curProgressSessions = Set.insert (req ^. params . token) $ curProgressSessions s }
+  End _ ->
+    modify $ \s -> s { curProgressSessions = Set.delete (req ^. params . token) $ curProgressSessions s }
+  _ -> pure ()
 
 -- Keep track of dynamic capability registration
 updateState (FromServerMess SClientRegisterCapability req) = do
-- 
2.30.2