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     capabilities :: ClientCapabilities -- ^ Specific capabilities the client should advertise. Default is yes to everything.
  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   }
  79
  80 instance Default SessionConfig where
  81   def = SessionConfig def 60 False
  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   }
  96
  97 class Monad m => HasReader r m where
  98   ask :: m r
  99   asks :: (r -> b) -> m b
 100   asks f = f <$> ask
 101
 102 instance Monad m => HasReader r (ParserStateReader a s r m) where
 103   ask = lift $ lift Reader.ask
 104
 105 instance Monad m => HasReader SessionContext (ConduitM a b (StateT s (ReaderT SessionContext m))) where
 106   ask = lift $ lift Reader.ask
 107
 108 data SessionState = SessionState
 109   {
 110     curReqId :: LspId
 111   , vfs :: VFS
 112   , curDiagnostics :: Map.Map Uri [Diagnostic]
 113   , curTimeoutId :: Int
 114   , overridingTimeout :: Bool
 115   -- ^ The last received message from the server.
 116   -- Used for providing exception information
 117   , lastReceivedMessage :: Maybe FromServerMessage
 118   }
 119
 120 class Monad m => HasState s m where
 121   get :: m s
 122
 123   put :: s -> m ()
 124
 125   modify :: (s -> s) -> m ()
 126   modify f = get >>= put . f
 127
 128   modifyM :: (HasState s m, Monad m) => (s -> m s) -> m ()
 129   modifyM f = get >>= f >>= put
 130
 131 instance Monad m => HasState s (ParserStateReader a s r m) where
 132   get = lift State.get
 133   put = lift . State.put
 134
 135 instance Monad m => HasState SessionState (ConduitM a b (StateT SessionState m))
 136  where
 137   get = lift State.get
 138   put = lift . State.put
 139
 140 type ParserStateReader a s r m = ConduitParser a (StateT s (ReaderT r m))
 141
 142 runSession :: SessionContext -> SessionState -> Session a -> IO (a, SessionState)
 143 runSession context state session = runReaderT (runStateT conduit state) context
 144   where
 145     conduit = runConduit $ chanSource .| watchdog .| updateStateC .| runConduitParser (catchError session handler)
 146
 147     handler (Unexpected "ConduitParser.empty") = do
 148       lastMsg <- fromJust . lastReceivedMessage <$> get
 149       name <- getParserName
 150       liftIO $ throw (UnexpectedMessage (T.unpack name) lastMsg)
 151
 152     handler e = throw e
 153
 154     chanSource = do
 155       msg <- liftIO $ readChan (messageChan context)
 156       yield msg
 157       chanSource
 158
 159
 160     watchdog :: ConduitM SessionMessage FromServerMessage (StateT SessionState (ReaderT SessionContext IO)) ()
 161     watchdog = Conduit.awaitForever $ \msg -> do
 162       curId <- curTimeoutId <$> get
 163       case msg of
 164         ServerMessage sMsg -> yield sMsg
 165         TimeoutMessage tId -> when (curId == tId) $ throw Timeout
 166
 167 -- | An internal version of 'runSession' that allows for a custom handler to listen to the server.
 168 -- It also does not automatically send initialize and exit messages.
 169 runSessionWithHandles :: Handle -- ^ Server in
 170                       -> Handle -- ^ Server out
 171                       -> (Handle -> SessionContext -> IO ()) -- ^ Server listener
 172                       -> SessionConfig
 173                       -> FilePath
 174                       -> Session a
 175                       -> IO a
 176 runSessionWithHandles serverIn serverOut serverHandler config rootDir session = do
 177   absRootDir <- canonicalizePath rootDir
 178
 179   hSetBuffering serverIn  NoBuffering
 180   hSetBuffering serverOut NoBuffering
 181
 182   reqMap <- newMVar newRequestMap
 183   messageChan <- newChan
 184   initRsp <- newEmptyMVar
 185
 186   let context = SessionContext serverIn absRootDir messageChan reqMap initRsp config
 187       initState = SessionState (IdInt 0) mempty mempty 0 False Nothing
 188
 189   threadId <- forkIO $ void $ serverHandler serverOut context
 190   (result, _) <- runSession context initState session
 191
 192   killThread threadId
 193
 194   return result
 195
 196 updateStateC :: ConduitM FromServerMessage FromServerMessage (StateT SessionState (ReaderT SessionContext IO)) ()
 197 updateStateC = awaitForever $ \msg -> do
 198   updateState msg
 199   yield msg
 200
 201 updateState :: (MonadIO m, HasReader SessionContext m, HasState SessionState m) => FromServerMessage -> m ()
 202 updateState (NotPublishDiagnostics n) = do
 203   let List diags = n ^. params . diagnostics
 204       doc = n ^. params . uri
 205   modify (\s ->
 206     let newDiags = Map.insert doc diags (curDiagnostics s)
 207       in s { curDiagnostics = newDiags })
 208
 209 updateState (ReqApplyWorkspaceEdit r) = do
 210
 211
 212   allChangeParams <- case r ^. params . edit . documentChanges of
 213     Just (List cs) -> do
 214       mapM_ (checkIfNeedsOpened . (^. textDocument . uri)) cs
 215       return $ map getParams cs
 216     Nothing -> case r ^. params . edit . changes of
 217       Just cs -> do
 218         mapM_ checkIfNeedsOpened (HashMap.keys cs)
 219         return $ concatMap (uncurry getChangeParams) (HashMap.toList cs)
 220       Nothing -> error "No changes!"
 221
 222   modifyM $ \s -> do
 223     newVFS <- liftIO $ changeFromServerVFS (vfs s) r
 224     return $ s { vfs = newVFS }
 225
 226   let groupedParams = groupBy (\a b -> (a ^. textDocument == b ^. textDocument)) allChangeParams
 227       mergedParams = map mergeParams groupedParams
 228
 229   -- TODO: Don't do this when replaying a session
 230   forM_ mergedParams (sendMessage . NotificationMessage "2.0" TextDocumentDidChange)
 231
 232   -- Update VFS to new document versions
 233   let sortedVersions = map (sortBy (compare `on` (^. textDocument . version))) groupedParams
 234       latestVersions = map ((^. textDocument) . last) sortedVersions
 235       bumpedVersions = map (version . _Just +~ 1) latestVersions
 236
 237   forM_ bumpedVersions $ \(VersionedTextDocumentIdentifier uri v) ->
 238     modify $ \s ->
 239       let oldVFS = vfs s
 240           update (VirtualFile oldV t) = VirtualFile (fromMaybe oldV v) t
 241           newVFS = Map.adjust update uri oldVFS
 242       in s { vfs = newVFS }
 243
 244   where checkIfNeedsOpened uri = do
 245           oldVFS <- vfs <$> get
 246           ctx <- ask
 247
 248           -- if its not open, open it
 249           unless (uri `Map.member` oldVFS) $ do
 250             let fp = fromJust $ uriToFilePath uri
 251             contents <- liftIO $ T.readFile fp
 252             let item = TextDocumentItem (filePathToUri fp) "" 0 contents
 253                 msg = NotificationMessage "2.0" TextDocumentDidOpen (DidOpenTextDocumentParams item)
 254             liftIO $ B.hPut (serverIn ctx) $ addHeader (encode msg)
 255
 256             modifyM $ \s -> do
 257               newVFS <- liftIO $ openVFS (vfs s) msg
 258               return $ s { vfs = newVFS }
 259
 260         getParams (TextDocumentEdit docId (List edits)) =
 261           let changeEvents = map (\e -> TextDocumentContentChangeEvent (Just (e ^. range)) Nothing (e ^. newText)) edits
 262             in DidChangeTextDocumentParams docId (List changeEvents)
 263
 264         textDocumentVersions uri = map (VersionedTextDocumentIdentifier uri . Just) [0..]
 265
 266         textDocumentEdits uri edits = map (\(v, e) -> TextDocumentEdit v (List [e])) $ zip (textDocumentVersions uri) edits
 267
 268         getChangeParams uri (List edits) = map getParams (textDocumentEdits uri (reverse edits))
 269
 270         mergeParams :: [DidChangeTextDocumentParams] -> DidChangeTextDocumentParams
 271         mergeParams params = let events = concat (toList (map (toList . (^. contentChanges)) params))
 272                               in DidChangeTextDocumentParams (head params ^. textDocument) (List events)
 273 updateState _ = return ()
 274
 275 sendMessage :: (MonadIO m, HasReader SessionContext m, ToJSON a) => a -> m ()
 276 sendMessage msg = do
 277   h <- serverIn <$> ask
 278   let encoded = encode msg
 279   liftIO $ do
 280
 281     setSGR [SetColor Foreground Vivid Cyan]
 282     putStrLn $ "--> " ++ B.unpack encoded
 283     setSGR [Reset]
 284
 285     B.hPut h (addHeader encoded)
 286
 287 -- | Execute a block f that will throw a 'TimeoutException'
 288 -- after duration seconds. This will override the global timeout
 289 -- for waiting for messages to arrive defined in 'SessionConfig'.
 290 withTimeout :: Int -> Session a -> Session a
 291 withTimeout duration f = do
 292   chan <- asks messageChan
 293   timeoutId <- curTimeoutId <$> get
 294   modify $ \s -> s { overridingTimeout = True }
 295   liftIO $ forkIO $ do
 296     threadDelay (duration * 1000000)
 297     writeChan chan (TimeoutMessage timeoutId)
 298   res <- f
 299   modify $ \s -> s { curTimeoutId = timeoutId + 1,
 300                      overridingTimeout = False
 301                    }
 302   return res