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