From e07c7e33ce829344c6173a7a6b0b7accbdb86d93 Mon Sep 17 00:00:00 2001 From: mangoiv Date: Fri, 8 Aug 2025 17:45:43 +0200 Subject: [PATCH] [feat] websockets Allow the user to set up an lsp server that talks to the user using the websocket protocol instead of standard IO and do enough refactoring to make the setup with websockets possible --- lsp/lsp.cabal | 1 + lsp/src/Language/LSP/Server/Control.hs | 256 ++++++++++++++++++++----- 2 files changed, 207 insertions(+), 50 deletions(-) diff --git a/lsp/lsp.cabal b/lsp/lsp.cabal index d1f6fa2c..825e5f9a 100644 --- a/lsp/lsp.cabal +++ b/lsp/lsp.cabal @@ -79,6 +79,7 @@ library , unliftio ^>=0.2 , unliftio-core ^>=0.2 , unordered-containers ^>=0.2 + , websockets ^>=0.13 executable lsp-demo-reactor-server import: warnings diff --git a/lsp/src/Language/LSP/Server/Control.hs b/lsp/src/Language/LSP/Server/Control.hs index 283f6c05..a1ec3d34 100644 --- a/lsp/src/Language/LSP/Server/Control.hs +++ b/lsp/src/Language/LSP/Server/Control.hs @@ -1,18 +1,35 @@ +{-# LANGUAGE DerivingStrategies #-} {-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} module Language.LSP.Server.Control ( -- * Running - runServer, runServerWith, - runServerWithHandles, + runServerWithConfig, + ServerConfig (..), LspServerLog (..), + + -- ** Using standard 'IO' 'Handle's + runServer, + + -- ** Using 'Handle's + runServerWithHandles, + prependHeader, + parseHeaders, + + -- ** Using websockets + WebsocketConfig (..), + withWebsocket, + withWebsocketRunServer, ) where import Colog.Core (LogAction (..), Severity (..), WithSeverity (..), (<&)) import Colog.Core qualified as L import Control.Applicative ((<|>)) import Control.Concurrent +import Control.Concurrent.Async import Control.Concurrent.STM.TChan +import Control.Exception (finally) import Control.Monad import Control.Monad.IO.Class import Control.Monad.STM @@ -31,6 +48,7 @@ import Language.LSP.Protocol.Message import Language.LSP.Server.Core import Language.LSP.Server.Processing qualified as Processing import Language.LSP.VFS +import Network.WebSockets qualified as WS import Prettyprinter import System.IO @@ -42,6 +60,7 @@ data LspServerLog | Starting | ParsedMsg T.Text | SendMsg TL.Text + | WebsocketLog WebsocketLog deriving (Show) instance Pretty LspServerLog where @@ -60,6 +79,7 @@ instance Pretty LspServerLog where pretty Starting = "Starting server" pretty (ParsedMsg msg) = "---> " <> pretty msg pretty (SendMsg msg) = "<--2-- " <> pretty msg + pretty (WebsocketLog msg) = "Websocket:" <+> pretty msg -- --------------------------------------------------------------------- @@ -71,18 +91,20 @@ instance Pretty LspServerLog where runServer :: forall config. ServerDefinition config -> IO Int runServer = runServerWithHandles - ioLogger - lspLogger + defaultIOLogger + defaultLspLogger stdin stdout + +defaultIOLogger :: LogAction IO (WithSeverity LspServerLog) +defaultIOLogger = L.cmap (show . prettyMsg) L.logStringStderr where prettyMsg l = "[" <> viaShow (L.getSeverity l) <> "] " <> pretty (L.getMsg l) - ioLogger :: LogAction IO (WithSeverity LspServerLog) - ioLogger = L.cmap (show . prettyMsg) L.logStringStderr - lspLogger :: LogAction (LspM config) (WithSeverity LspServerLog) - lspLogger = - let clientLogger = L.cmap (fmap (T.pack . show . pretty)) defaultClientLogger - in clientLogger <> L.hoistLogAction liftIO ioLogger + +defaultLspLogger :: LogAction (LspM config) (WithSeverity LspServerLog) +defaultLspLogger = + let clientLogger = L.cmap (fmap (T.pack . show . pretty)) defaultClientLogger + in clientLogger <> L.hoistLogAction liftIO defaultIOLogger {- | Starts a language server over the specified handles. This function will return once the @exit@ notification is received. @@ -116,6 +138,9 @@ runServerWithHandles ioLogger logger hin hout serverDefinition = do {- | Starts listening and sending requests and responses using the specified I/O. + + Assumes that the client sends (and wants to receive) the Content-Length + header. If you do not want this to be the case, use 'runServerWithConfig' -} runServerWith :: -- | The logger to use outside the main body of the server where we can't assume the ability to send messages. @@ -123,20 +148,44 @@ runServerWith :: -- | The logger to use once the server has started and can successfully send messages. LogAction (LspM config) (WithSeverity LspServerLog) -> -- | Client input. - IO BS.ByteString -> + IO BS.StrictByteString -> -- | Function to provide output to. - (BSL.ByteString -> IO ()) -> + (BSL.LazyByteString -> IO ()) -> ServerDefinition config -> IO Int -- exit code -runServerWith ioLogger logger clientIn clientOut serverDefinition = do +runServerWith ioLogger lspLogger inwards outwards = + runServerWithConfig ServerConfig{prepareOutwards = prependHeader, parseInwards = parseHeaders, ..} + +-- --------------------------------------------------------------------- + +data ServerConfig config = ServerConfig + { ioLogger :: LogAction IO (WithSeverity LspServerLog) + -- ^ The logger to use outside the main body of the server where we can't assume the ability to send messages. + , lspLogger :: LogAction (LspM config) (WithSeverity LspServerLog) + -- ^ The logger to use once the server has started and can successfully send messages. + , inwards :: IO BS.StrictByteString + -- ^ Client input. + , outwards :: BSL.LazyByteString -> IO () + -- ^ Function to provide output to. + , prepareOutwards :: BSL.LazyByteString -> BSL.LazyByteString + -- ^ how to prepare an outgoing response for sending. This can be used, to e.g. prepend the Content-Length header, c.f. 'prependHeader' + , parseInwards :: Attoparsec.Parser BS.StrictByteString + -- ^ how to parse the input. This can be used to consume the Content-Length and Content-Type headers, c.f. 'parseHeaders' + } + +runServerWithConfig :: + ServerConfig config -> + ServerDefinition config -> + IO Int +runServerWithConfig ServerConfig{..} serverDefinition = do ioLogger <& Starting `WithSeverity` Info cout <- atomically newTChan :: IO (TChan J.Value) - _rhpid <- forkIO $ sendServer ioLogger cout clientOut + _rhpid <- forkIO $ sendServer ioLogger cout outwards prepareOutwards let sendMsg msg = atomically $ writeTChan cout $ J.toJSON msg - ioLoop ioLogger logger clientIn serverDefinition emptyVFS sendMsg + ioLoop ioLogger lspLogger inwards parseInwards serverDefinition emptyVFS sendMsg return 1 @@ -146,12 +195,13 @@ ioLoop :: forall config. LogAction IO (WithSeverity LspServerLog) -> LogAction (LspM config) (WithSeverity LspServerLog) -> - IO BS.ByteString -> + IO BS.StrictByteString -> + Attoparsec.Parser BS.StrictByteString -> ServerDefinition config -> VFS -> (FromServerMessage -> IO ()) -> IO () -ioLoop ioLogger logger clientIn serverDefinition vfs sendMsg = do +ioLoop ioLogger logger clientIn parser serverDefinition vfs sendMsg = do minitialize <- parseOne ioLogger clientIn (parse parser "") case minitialize of Nothing -> pure () @@ -167,7 +217,7 @@ ioLoop ioLogger logger clientIn serverDefinition vfs sendMsg = do pioLogger = L.cmap (fmap LspProcessingLog) ioLogger pLogger = L.cmap (fmap LspProcessingLog) logger - loop :: Result BS.ByteString -> LspM config () + loop :: Result BS.StrictByteString -> LspM config () loop = go where go r = do @@ -178,30 +228,11 @@ ioLoop ioLogger logger clientIn serverDefinition vfs sendMsg = do Processing.processMessage pLogger $ BSL.fromStrict msg go (parse parser remainder) - parser = do - try contentType <|> return () - len <- contentLength - try contentType <|> return () - _ <- string _ONE_CRLF - Attoparsec.take len - - contentLength = do - _ <- string "Content-Length: " - len <- decimal - _ <- string _ONE_CRLF - return len - - contentType = do - _ <- string "Content-Type: " - skipWhile (/= '\r') - _ <- string _ONE_CRLF - return () - parseOne :: MonadIO m => LogAction m (WithSeverity LspServerLog) -> - IO BS.ByteString -> - Result BS.ByteString -> + IO BS.StrictByteString -> + Result BS.StrictByteString -> m (Maybe (BS.ByteString, BS.ByteString)) parseOne logger clientIn = go where @@ -223,22 +254,115 @@ parseOne logger clientIn = go -- --------------------------------------------------------------------- +data WebsocketLog + = WebsocketShutDown + | WebsocketNewConnection + | WebsocketConnectionClosed + | WebsocketPing + | WebsocketStarted + | WebsocketIncomingRequest + | WebsocketOutgoingResponse + deriving stock (Show) + +instance Pretty WebsocketLog where + pretty l = case l of + WebsocketPing -> "Ping" + WebsocketStarted -> "Started Server, waiting for connections" + WebsocketShutDown -> "Shut down server" + WebsocketNewConnection -> "New connection established" + WebsocketIncomingRequest -> "Received request" + WebsocketConnectionClosed -> "Closed connection to client" + WebsocketOutgoingResponse -> "Sent response" + +-- | 'host' and 'port' of the websocket server to set up +data WebsocketConfig = WebsocketConfig + { host :: !String + -- ^ the host of the websocket server, e.g. @"localhost"@ + , port :: !Int + -- ^ the port of the websocket server, e.g. @8080@ + } + +-- | Set up a websocket server, then call call the continuation (in our case this corresponds to the language server) after accepting a connection +withWebsocket :: + -- | The logger + LogAction IO (WithSeverity LspServerLog) -> + -- | The configuration of the websocket server + WebsocketConfig -> + -- | invoke the lsp server, passing communication functions + (IO BS.StrictByteString -> (BSL.LazyByteString -> IO ()) -> IO r) -> + IO () +withWebsocket logger conf startLspServer = do + let wsLogger = L.cmap (fmap WebsocketLog) logger + + WS.runServer (host conf) (port conf) $ \pending -> do + conn <- WS.acceptRequest pending + wsLogger <& WebsocketNewConnection `WithSeverity` Debug + + outChan <- newChan + inChan <- newChan + + let inwards = readChan inChan + outwards = writeChan outChan + + WS.withPingThread conn 30 (wsLogger <& WebsocketPing `WithSeverity` Debug) $ do + withAsync (startLspServer inwards outwards) $ \_lspAsync -> + race_ + ( forever $ do + msg <- readChan outChan + wsLogger <& WebsocketOutgoingResponse `WithSeverity` Debug + WS.sendTextData conn msg + ) + ( forever $ do + msg <- WS.receiveData conn + wsLogger <& WebsocketIncomingRequest `WithSeverity` Debug + writeChan inChan msg + ) + `finally` do + wsLogger <& WebsocketConnectionClosed `WithSeverity` Debug + +withWebsocketRunServer :: + -- | Configuration for the websocket + WebsocketConfig -> + -- | How to set up a new 'ServerDefinition' for a specific configuration. z + -- This is passed as CPS'd 'IO' to allow for setting (- and cleaning) up + -- a server per websocket connection + ((ServerDefinition config -> IO Int) -> IO Int) -> + -- | The 'IO' logger + LogAction IO (WithSeverity LspServerLog) -> + -- | The logger that logs in 'LspM' to the client + LogAction (LspM config) (WithSeverity LspServerLog) -> + IO () +withWebsocketRunServer wsConf withLspDefinition ioLogger lspLogger = + withWebsocket ioLogger wsConf $ \inwards outwards -> do + withLspDefinition $ \lspDefinition -> + runServerWithConfig + ServerConfig + { ioLogger + , lspLogger + , inwards + , outwards + , -- NOTE: if you run the language server on websockets, you do not + -- need to prepend headers to requests and responses, because + -- the chunking is already handled by the websocket, i.e. there's + -- no situation where the client or the server has to rely on input/ + -- output chunking + prepareOutwards = id + , parseInwards = Attoparsec.takeByteString + } + lspDefinition + +-- --------------------------------------------------------------------- + -- | Simple server to make sure all output is serialised -sendServer :: LogAction IO (WithSeverity LspServerLog) -> TChan J.Value -> (BSL.ByteString -> IO ()) -> IO () -sendServer _logger msgChan clientOut = do +sendServer :: LogAction IO (WithSeverity LspServerLog) -> TChan J.Value -> (BSL.LazyByteString -> IO ()) -> (BSL.LazyByteString -> BSL.LazyByteString) -> IO () +sendServer _logger msgChan clientOut prepareMessage = do forever $ do msg <- atomically $ readTChan msgChan -- We need to make sure we only send over the content of the message, -- and no other tags/wrapper stuff let str = J.encode msg - - let out = - BSL.concat - [ TL.encodeUtf8 $ TL.pack $ "Content-Length: " ++ show (BSL.length str) - , BSL.fromStrict _TWO_CRLF - , str - ] + let out = prepareMessage str clientOut out @@ -246,7 +370,39 @@ sendServer _logger msgChan clientOut = do -- This can lead to infinite recursion in logging, see https://github.com/haskell/lsp/issues/447 -- logger <& SendMsg (TL.decodeUtf8 str) `WithSeverity` Debug -_ONE_CRLF :: BS.ByteString +-- | prepend a Content-Length header to the given message +prependHeader :: BSL.LazyByteString -> BSL.LazyByteString +prependHeader str = + BSL.concat + [ TL.encodeUtf8 $ TL.pack $ "Content-Length: " ++ show (BSL.length str) + , BSL.fromStrict _TWO_CRLF + , str + ] + +{- | parse Content-Length and Content-Type headers and then consume + input with length of the Content-Length +-} +parseHeaders :: Attoparsec.Parser BS.StrictByteString +parseHeaders = do + try contentType <|> return () + len <- contentLength + try contentType <|> return () + _ <- string _ONE_CRLF + Attoparsec.take len + where + contentLength = do + _ <- string "Content-Length: " + len <- decimal + _ <- string _ONE_CRLF + return len + + contentType = do + _ <- string "Content-Type: " + skipWhile (/= '\r') + _ <- string _ONE_CRLF + return () + +_ONE_CRLF :: BS.StrictByteString _ONE_CRLF = "\r\n" -_TWO_CRLF :: BS.ByteString +_TWO_CRLF :: BS.StrictByteString _TWO_CRLF = "\r\n\r\n"