src/Language/Haskell/LSP/Test/Session.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2 {-# LANGUAGE FlexibleInstances #-}
   3 {-# LANGUAGE MultiParamTypeClasses #-}
   4 {-# LANGUAGE FlexibleContexts #-}
   5
   6 module Language.Haskell.LSP.Test.Session
   7   ( Session
   8   , SessionConfig(..)
   9   , SessionMessage(..)
  10   , SessionContext(..)
  11   , SessionState(..)
  12   , runSessionWithHandles
  13   , get
  14   , put
  15   , modify
  16   , modifyM
  17   , ask
  18   , asks
  19   , sendMessage
  20   , updateState
  21   , withTimeout
  22   )
  23
  24 where
  25
  26 import Control.Concurrent hiding (yield)
  27 import Control.Exception
  28 import Control.Lens hiding (List)
  29 import Control.Monad
  30 import Control.Monad.IO.Class
  31 import Control.Monad.Except
  32 import Control.Monad.Trans.Reader (ReaderT, runReaderT)
  33 import qualified Control.Monad.Trans.Reader as Reader (ask)
  34 import Control.Monad.Trans.State (StateT, runStateT)
  35 import qualified Control.Monad.Trans.State as State (get, put)
  36 import qualified Data.ByteString.Lazy.Char8 as B
  37 import Data.Aeson
  38 import Data.Conduit as Conduit
  39 import Data.Conduit.Parser as Parser
  40 import Data.Default
  41 import Data.Foldable
  42 import Data.List
  43 import qualified Data.Map as Map
  44 import qualified Data.Text as T
  45 import qualified Data.Text.IO as T
  46 import qualified Data.HashMap.Strict as HashMap
  47 import Data.Maybe
  48 import Data.Function
  49 import Language.Haskell.LSP.Messages
  50 import Language.Haskell.LSP.Types.Capabilities
  51 import Language.Haskell.LSP.Types hiding (error)
  52 import Language.Haskell.LSP.VFS
  53 import Language.Haskell.LSP.Test.Decoding
  54 import Language.Haskell.LSP.Test.Exceptions
  55 import System.Console.ANSI
  56 import System.Directory
  57 import System.IO
  58
  59 -- | A session representing one instance of launching and connecting to a server.
  60 --
  61 -- You can send and receive messages to the server within 'Session' via 'getMessage',
  62 -- 'sendRequest' and 'sendNotification'.
  63 --
  64 -- @
  65 -- runSession \"path\/to\/root\/dir\" $ do
  66 --   docItem <- getDocItem "Desktop/simple.hs" "haskell"
  67 --   sendNotification TextDocumentDidOpen (DidOpenTextDocumentParams docItem)
  68 --   diagnostics <- getMessage :: Session PublishDiagnosticsNotification
  69 -- @
  70 type Session = ParserStateReader FromServerMessage SessionState SessionContext IO
  71
  72 -- | Stuff you can configure for a 'Session'.
  73 data SessionConfig = SessionConfig
  74   {
  75     messageTimeout :: Int -- ^ Maximum time to wait for a message in seconds. Defaults to 60.
  76   , logStdErr :: Bool -- ^ When True redirects the servers stderr output to haskell-lsp-test's stdout. Defaults to False.
  77   , logMessages :: Bool -- ^ When True traces the communication between client and server to stdout. Defaults to True.
  78   }
  79
  80 instance Default SessionConfig where
  81   def = SessionConfig 60 False True
  82
  83 data SessionMessage = ServerMessage FromServerMessage
  84                     | TimeoutMessage Int
  85   deriving Show
  86
  87 data SessionContext = SessionContext
  88   {
  89     serverIn :: Handle
  90   , rootDir :: FilePath
  91   , messageChan :: Chan SessionMessage
  92   , requestMap :: MVar RequestMap
  93   , initRsp :: MVar InitializeResponse
  94   , config :: SessionConfig
  95   , sessionCapabilities :: ClientCapabilities
  96   }
  97
  98 class Monad m => HasReader r m where
  99   ask :: m r
 100   asks :: (r -> b) -> m b
 101   asks f = f <$> ask
 102
 103 instance Monad m => HasReader r (ParserStateReader a s r m) where
 104   ask = lift $ lift Reader.ask
 105
 106 instance Monad m => HasReader SessionContext (ConduitM a b (StateT s (ReaderT SessionContext m))) where
 107   ask = lift $ lift Reader.ask
 108
 109 data SessionState = SessionState
 110   {
 111     curReqId :: LspId
 112   , vfs :: VFS
 113   , curDiagnostics :: Map.Map Uri [Diagnostic]
 114   , curTimeoutId :: Int
 115   , overridingTimeout :: Bool
 116   -- ^ The last received message from the server.
 117   -- Used for providing exception information
 118   , lastReceivedMessage :: Maybe FromServerMessage
 119   }
 120
 121 class Monad m => HasState s m where
 122   get :: m s
 123
 124   put :: s -> m ()
 125
 126   modify :: (s -> s) -> m ()
 127   modify f = get >>= put . f
 128
 129   modifyM :: (HasState s m, Monad m) => (s -> m s) -> m ()
 130   modifyM f = get >>= f >>= put
 131
 132 instance Monad m => HasState s (ParserStateReader a s r m) where
 133   get = lift State.get
 134   put = lift . State.put
 135
 136 instance Monad m => HasState SessionState (ConduitM a b (StateT SessionState m))
 137  where
 138   get = lift State.get
 139   put = lift . State.put
 140
 141 type ParserStateReader a s r m = ConduitParser a (StateT s (ReaderT r m))
 142
 143 runSession :: SessionContext -> SessionState -> Session a -> IO (a, SessionState)
 144 runSession context state session = runReaderT (runStateT conduit state) context
 145   where
 146     conduit = runConduit $ chanSource .| watchdog .| updateStateC .| runConduitParser (catchError session handler)
 147
 148     handler (Unexpected "ConduitParser.empty") = do
 149       lastMsg <- fromJust . lastReceivedMessage <$> get
 150       name <- getParserName
 151       liftIO $ throw (UnexpectedMessage (T.unpack name) lastMsg)
 152
 153     handler e = throw e
 154
 155     chanSource = do
 156       msg <- liftIO $ readChan (messageChan context)
 157       yield msg
 158       chanSource
 159
 160
 161     watchdog :: ConduitM SessionMessage FromServerMessage (StateT SessionState (ReaderT SessionContext IO)) ()
 162     watchdog = Conduit.awaitForever $ \msg -> do
 163       curId <- curTimeoutId <$> get
 164       case msg of
 165         ServerMessage sMsg -> yield sMsg
 166         TimeoutMessage tId -> when (curId == tId) $ throw Timeout
 167
 168 -- | An internal version of 'runSession' that allows for a custom handler to listen to the server.
 169 -- It also does not automatically send initialize and exit messages.
 170 runSessionWithHandles :: Handle -- ^ Server in
 171                       -> Handle -- ^ Server out
 172                       -> (Handle -> SessionContext -> IO ()) -- ^ Server listener
 173                       -> SessionConfig
 174                       -> ClientCapabilities
 175                       -> FilePath -- ^ Root directory
 176                       -> Session a
 177                       -> IO a
 178 runSessionWithHandles serverIn serverOut serverHandler config caps rootDir session = do
 179   absRootDir <- canonicalizePath rootDir
 180
 181   hSetBuffering serverIn  NoBuffering
 182   hSetBuffering serverOut NoBuffering
 183
 184   reqMap <- newMVar newRequestMap
 185   messageChan <- newChan
 186   initRsp <- newEmptyMVar
 187
 188   let context = SessionContext serverIn absRootDir messageChan reqMap initRsp config caps
 189       initState = SessionState (IdInt 0) mempty mempty 0 False Nothing
 190
 191   threadId <- forkIO $ void $ serverHandler serverOut context
 192   (result, _) <- runSession context initState session
 193
 194   killThread threadId
 195
 196   return result
 197
 198 updateStateC :: ConduitM FromServerMessage FromServerMessage (StateT SessionState (ReaderT SessionContext IO)) ()
 199 updateStateC = awaitForever $ \msg -> do
 200   updateState msg
 201   yield msg
 202
 203 updateState :: (MonadIO m, HasReader SessionContext m, HasState SessionState m) => FromServerMessage -> m ()
 204 updateState (NotPublishDiagnostics n) = do
 205   let List diags = n ^. params . diagnostics
 206       doc = n ^. params . uri
 207   modify (\s ->
 208     let newDiags = Map.insert doc diags (curDiagnostics s)
 209       in s { curDiagnostics = newDiags })
 210
 211 updateState (ReqApplyWorkspaceEdit r) = do
 212
 213   allChangeParams <- case r ^. params . edit . documentChanges of
 214     Just (List cs) -> do
 215       mapM_ (checkIfNeedsOpened . (^. textDocument . uri)) cs
 216       return $ map getParams cs
 217     Nothing -> case r ^. params . edit . changes of
 218       Just cs -> do
 219         mapM_ checkIfNeedsOpened (HashMap.keys cs)
 220         return $ concatMap (uncurry getChangeParams) (HashMap.toList cs)
 221       Nothing -> error "No changes!"
 222
 223   modifyM $ \s -> do
 224     newVFS <- liftIO $ changeFromServerVFS (vfs s) r
 225     return $ s { vfs = newVFS }
 226
 227   let groupedParams = groupBy (\a b -> (a ^. textDocument == b ^. textDocument)) allChangeParams
 228       mergedParams = map mergeParams groupedParams
 229
 230   -- TODO: Don't do this when replaying a session
 231   forM_ mergedParams (sendMessage . NotificationMessage "2.0" TextDocumentDidChange)
 232
 233   -- Update VFS to new document versions
 234   let sortedVersions = map (sortBy (compare `on` (^. textDocument . version))) groupedParams
 235       latestVersions = map ((^. textDocument) . last) sortedVersions
 236       bumpedVersions = map (version . _Just +~ 1) latestVersions
 237
 238   forM_ bumpedVersions $ \(VersionedTextDocumentIdentifier uri v) ->
 239     modify $ \s ->
 240       let oldVFS = vfs s
 241           update (VirtualFile oldV t) = VirtualFile (fromMaybe oldV v) t
 242           newVFS = Map.adjust update uri oldVFS
 243       in s { vfs = newVFS }
 244
 245   where checkIfNeedsOpened uri = do
 246           oldVFS <- vfs <$> get
 247           ctx <- ask
 248
 249           -- if its not open, open it
 250           unless (uri `Map.member` oldVFS) $ do
 251             let fp = fromJust $ uriToFilePath uri
 252             contents <- liftIO $ T.readFile fp
 253             let item = TextDocumentItem (filePathToUri fp) "" 0 contents
 254                 msg = NotificationMessage "2.0" TextDocumentDidOpen (DidOpenTextDocumentParams item)
 255             liftIO $ B.hPut (serverIn ctx) $ addHeader (encode msg)
 256
 257             modifyM $ \s -> do
 258               newVFS <- liftIO $ openVFS (vfs s) msg
 259               return $ s { vfs = newVFS }
 260
 261         getParams (TextDocumentEdit docId (List edits)) =
 262           let changeEvents = map (\e -> TextDocumentContentChangeEvent (Just (e ^. range)) Nothing (e ^. newText)) edits
 263             in DidChangeTextDocumentParams docId (List changeEvents)
 264
 265         textDocumentVersions uri = map (VersionedTextDocumentIdentifier uri . Just) [0..]
 266
 267         textDocumentEdits uri edits = map (\(v, e) -> TextDocumentEdit v (List [e])) $ zip (textDocumentVersions uri) edits
 268
 269         getChangeParams uri (List edits) = map getParams (textDocumentEdits uri (reverse edits))
 270
 271         mergeParams :: [DidChangeTextDocumentParams] -> DidChangeTextDocumentParams
 272         mergeParams params = let events = concat (toList (map (toList . (^. contentChanges)) params))
 273                               in DidChangeTextDocumentParams (head params ^. textDocument) (List events)
 274 updateState _ = return ()
 275
 276 sendMessage :: (MonadIO m, HasReader SessionContext m, ToJSON a) => a -> m ()
 277 sendMessage msg = do
 278   h <- serverIn <$> ask
 279   let encoded = encode msg
 280   liftIO $ do
 281
 282     setSGR [SetColor Foreground Vivid Cyan]
 283     putStrLn $ "--> " ++ B.unpack encoded
 284     setSGR [Reset]
 285
 286     B.hPut h (addHeader encoded)
 287
 288 -- | Execute a block f that will throw a 'TimeoutException'
 289 -- after duration seconds. This will override the global timeout
 290 -- for waiting for messages to arrive defined in 'SessionConfig'.
 291 withTimeout :: Int -> Session a -> Session a
 292 withTimeout duration f = do
 293   chan <- asks messageChan
 294   timeoutId <- curTimeoutId <$> get
 295   modify $ \s -> s { overridingTimeout = True }
 296   liftIO $ forkIO $ do
 297     threadDelay (duration * 1000000)
 298     writeChan chan (TimeoutMessage timeoutId)
 299   res <- f
 300   modify $ \s -> s { curTimeoutId = timeoutId + 1,
 301                      overridingTimeout = False
 302                    }
 303   return res