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