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