{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- So we can keep using the old prettyprinter modules (which have a better
-- compatibility range) for now.
{-# OPTIONS_GHC -Wno-deprecations #-}

module Language.LSP.Server.Control (
  -- * Running
  runServer,
  runServerWith,
  runServerWithHandles,
  LspServerLog (..),
) where

import Colog.Core (LogAction (..), Severity (..), WithSeverity (..), (<&))
import Colog.Core qualified as L
import Control.Applicative ((<|>))
import Control.Concurrent
import Control.Concurrent.STM.TChan
import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.STM
import Data.Aeson qualified as J
import Data.Attoparsec.ByteString qualified as Attoparsec
import Data.Attoparsec.ByteString.Char8
import Data.ByteString qualified as BS
import Data.ByteString.Builder.Extra (defaultChunkSize)
import Data.ByteString.Lazy qualified as BSL
import Data.List
import Data.Text qualified as T
import Data.Text.Lazy qualified as TL
import Data.Text.Lazy.Encoding qualified as TL
import Data.Text.Prettyprint.Doc
import Language.LSP.Logging (defaultClientLogger)
import Language.LSP.Protocol.Message
import Language.LSP.Server.Core
import Language.LSP.Server.Processing qualified as Processing
import Language.LSP.VFS
import System.IO

data LspServerLog
  = LspProcessingLog Processing.LspProcessingLog
  | DecodeInitializeError String
  | HeaderParseFail [String] String
  | EOF
  | Starting
  | ParsedMsg T.Text
  | SendMsg TL.Text
  deriving (Int -> LspServerLog -> ShowS
[LspServerLog] -> ShowS
LspServerLog -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [LspServerLog] -> ShowS
$cshowList :: [LspServerLog] -> ShowS
show :: LspServerLog -> String
$cshow :: LspServerLog -> String
showsPrec :: Int -> LspServerLog -> ShowS
$cshowsPrec :: Int -> LspServerLog -> ShowS
Show)

instance Pretty LspServerLog where
  pretty :: forall ann. LspServerLog -> Doc ann
pretty (LspProcessingLog LspProcessingLog
l) = forall a ann. Pretty a => a -> Doc ann
pretty LspProcessingLog
l
  pretty (DecodeInitializeError String
err) =
    forall ann. [Doc ann] -> Doc ann
vsep
      [ Doc ann
"Got error while decoding initialize:"
      , forall a ann. Pretty a => a -> Doc ann
pretty String
err
      ]
  pretty (HeaderParseFail [String]
ctxs String
err) =
    forall ann. [Doc ann] -> Doc ann
vsep
      [ Doc ann
"Failed to parse message header:"
      , forall a ann. Pretty a => a -> Doc ann
pretty (forall a. [a] -> [[a]] -> [a]
intercalate String
" > " [String]
ctxs) forall a. Semigroup a => a -> a -> a
<> Doc ann
": " forall ann. Doc ann -> Doc ann -> Doc ann
<+> forall a ann. Pretty a => a -> Doc ann
pretty String
err
      ]
  pretty LspServerLog
EOF = Doc ann
"Got EOF"
  pretty LspServerLog
Starting = Doc ann
"Starting server"
  pretty (ParsedMsg Text
msg) = Doc ann
"---> " forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty Text
msg
  pretty (SendMsg Text
msg) = Doc ann
"<--2-- " forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty Text
msg

-- ---------------------------------------------------------------------

{- | Convenience function for 'runServerWithHandles' which:
     (1) reads from stdin;
     (2) writes to stdout; and
     (3) logs to stderr and to the client, with some basic filtering.
-}
runServer :: forall config. ServerDefinition config -> IO Int
runServer :: forall config. ServerDefinition config -> IO Int
runServer =
  forall config.
LogAction IO (WithSeverity LspServerLog)
-> LogAction (LspM config) (WithSeverity LspServerLog)
-> Handle
-> Handle
-> ServerDefinition config
-> IO Int
runServerWithHandles
    LogAction IO (WithSeverity LspServerLog)
ioLogger
    LogAction (LspM config) (WithSeverity LspServerLog)
lspLogger
    Handle
stdin
    Handle
stdout
 where
  prettyMsg :: WithSeverity a -> Doc ann
prettyMsg WithSeverity a
l = Doc ann
"[" forall a. Semigroup a => a -> a -> a
<> forall a ann. Show a => a -> Doc ann
viaShow (forall msg. WithSeverity msg -> Severity
L.getSeverity WithSeverity a
l) forall a. Semigroup a => a -> a -> a
<> Doc ann
"] " forall a. Semigroup a => a -> a -> a
<> forall a ann. Pretty a => a -> Doc ann
pretty (forall msg. WithSeverity msg -> msg
L.getMsg WithSeverity a
l)
  ioLogger :: LogAction IO (WithSeverity LspServerLog)
  ioLogger :: LogAction IO (WithSeverity LspServerLog)
ioLogger = forall a b (m :: * -> *).
(a -> b) -> LogAction m b -> LogAction m a
L.cmap (forall a. Show a => a -> String
show forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {a} {ann}. Pretty a => WithSeverity a -> Doc ann
prettyMsg) forall (m :: * -> *). MonadIO m => LogAction m String
L.logStringStderr
  lspLogger :: LogAction (LspM config) (WithSeverity LspServerLog)
  lspLogger :: LogAction (LspM config) (WithSeverity LspServerLog)
lspLogger =
    let clientLogger :: LogAction (LspM config) (WithSeverity LspServerLog)
clientLogger = forall a b (m :: * -> *).
(a -> b) -> LogAction m b -> LogAction m a
L.cmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (String -> Text
T.pack forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> String
show forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a ann. Pretty a => a -> Doc ann
pretty)) forall c (m :: * -> *).
MonadLsp c m =>
LogAction m (WithSeverity Text)
defaultClientLogger
     in LogAction (LspM config) (WithSeverity LspServerLog)
clientLogger forall a. Semigroup a => a -> a -> a
<> forall (m :: * -> *) (n :: * -> *) a.
(forall x. m x -> n x) -> LogAction m a -> LogAction n a
L.hoistLogAction forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO LogAction IO (WithSeverity LspServerLog)
ioLogger

{- | Starts a language server over the specified handles.
 This function will return once the @exit@ notification is received.
-}
runServerWithHandles ::
  -- | The logger to use outside the main body of the server where we can't assume the ability to send messages.
  LogAction IO (WithSeverity LspServerLog) ->
  -- | The logger to use once the server has started and can successfully send messages.
  LogAction (LspM config) (WithSeverity LspServerLog) ->
  -- | Handle to read client input from.
  Handle ->
  -- | Handle to write output to.
  Handle ->
  ServerDefinition config ->
  IO Int -- exit code
runServerWithHandles :: forall config.
LogAction IO (WithSeverity LspServerLog)
-> LogAction (LspM config) (WithSeverity LspServerLog)
-> Handle
-> Handle
-> ServerDefinition config
-> IO Int
runServerWithHandles LogAction IO (WithSeverity LspServerLog)
ioLogger LogAction (LspM config) (WithSeverity LspServerLog)
logger Handle
hin Handle
hout ServerDefinition config
serverDefinition = do
  Handle -> BufferMode -> IO ()
hSetBuffering Handle
hin BufferMode
NoBuffering
  Handle -> TextEncoding -> IO ()
hSetEncoding Handle
hin TextEncoding
utf8

  Handle -> BufferMode -> IO ()
hSetBuffering Handle
hout BufferMode
NoBuffering
  Handle -> TextEncoding -> IO ()
hSetEncoding Handle
hout TextEncoding
utf8

  let
    clientIn :: IO ByteString
clientIn = Handle -> Int -> IO ByteString
BS.hGetSome Handle
hin Int
defaultChunkSize

    clientOut :: ByteString -> IO ()
clientOut ByteString
out = do
      Handle -> ByteString -> IO ()
BSL.hPut Handle
hout ByteString
out
      Handle -> IO ()
hFlush Handle
hout

  forall config.
LogAction IO (WithSeverity LspServerLog)
-> LogAction (LspM config) (WithSeverity LspServerLog)
-> IO ByteString
-> (ByteString -> IO ())
-> ServerDefinition config
-> IO Int
runServerWith LogAction IO (WithSeverity LspServerLog)
ioLogger LogAction (LspM config) (WithSeverity LspServerLog)
logger IO ByteString
clientIn ByteString -> IO ()
clientOut ServerDefinition config
serverDefinition

{- | Starts listening and sending requests and responses
 using the specified I/O.
-}
runServerWith ::
  -- | The logger to use outside the main body of the server where we can't assume the ability to send messages.
  LogAction IO (WithSeverity LspServerLog) ->
  -- | The logger to use once the server has started and can successfully send messages.
  LogAction (LspM config) (WithSeverity LspServerLog) ->
  -- | Client input.
  IO BS.ByteString ->
  -- | Function to provide output to.
  (BSL.ByteString -> IO ()) ->
  ServerDefinition config ->
  IO Int -- exit code
runServerWith :: forall config.
LogAction IO (WithSeverity LspServerLog)
-> LogAction (LspM config) (WithSeverity LspServerLog)
-> IO ByteString
-> (ByteString -> IO ())
-> ServerDefinition config
-> IO Int
runServerWith LogAction IO (WithSeverity LspServerLog)
ioLogger LogAction (LspM config) (WithSeverity LspServerLog)
logger IO ByteString
clientIn ByteString -> IO ()
clientOut ServerDefinition config
serverDefinition = do
  LogAction IO (WithSeverity LspServerLog)
ioLogger forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& LspServerLog
Starting forall msg. msg -> Severity -> WithSeverity msg
`WithSeverity` Severity
Info

  TChan Value
cout <- forall a. STM a -> IO a
atomically forall a. STM (TChan a)
newTChan :: IO (TChan J.Value)
  ThreadId
_rhpid <- IO () -> IO ThreadId
forkIO forall a b. (a -> b) -> a -> b
$ LogAction IO (WithSeverity LspServerLog)
-> TChan Value -> (ByteString -> IO ()) -> IO ()
sendServer LogAction IO (WithSeverity LspServerLog)
ioLogger TChan Value
cout ByteString -> IO ()
clientOut

  let sendMsg :: a -> IO ()
sendMsg a
msg = forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TChan a -> a -> STM ()
writeTChan TChan Value
cout forall a b. (a -> b) -> a -> b
$ forall a. ToJSON a => a -> Value
J.toJSON a
msg

  forall config.
LogAction IO (WithSeverity LspServerLog)
-> LogAction (LspM config) (WithSeverity LspServerLog)
-> IO ByteString
-> ServerDefinition config
-> VFS
-> (FromServerMessage -> IO ())
-> IO ()
ioLoop LogAction IO (WithSeverity LspServerLog)
ioLogger LogAction (LspM config) (WithSeverity LspServerLog)
logger IO ByteString
clientIn ServerDefinition config
serverDefinition VFS
emptyVFS forall {a}. ToJSON a => a -> IO ()
sendMsg

  forall (m :: * -> *) a. Monad m => a -> m a
return Int
1

-- ---------------------------------------------------------------------

ioLoop ::
  forall config.
  LogAction IO (WithSeverity LspServerLog) ->
  LogAction (LspM config) (WithSeverity LspServerLog) ->
  IO BS.ByteString ->
  ServerDefinition config ->
  VFS ->
  (FromServerMessage -> IO ()) ->
  IO ()
ioLoop :: forall config.
LogAction IO (WithSeverity LspServerLog)
-> LogAction (LspM config) (WithSeverity LspServerLog)
-> IO ByteString
-> ServerDefinition config
-> VFS
-> (FromServerMessage -> IO ())
-> IO ()
ioLoop LogAction IO (WithSeverity LspServerLog)
ioLogger LogAction (LspM config) (WithSeverity LspServerLog)
logger IO ByteString
clientIn ServerDefinition config
serverDefinition VFS
vfs FromServerMessage -> IO ()
sendMsg = do
  Maybe (ByteString, ByteString)
minitialize <- forall (m :: * -> *).
MonadIO m =>
LogAction m (WithSeverity LspServerLog)
-> IO ByteString
-> Result ByteString
-> m (Maybe (ByteString, ByteString))
parseOne LogAction IO (WithSeverity LspServerLog)
ioLogger IO ByteString
clientIn (forall a. Parser a -> ByteString -> Result a
parse Parser ByteString ByteString
parser ByteString
"")
  case Maybe (ByteString, ByteString)
minitialize of
    Maybe (ByteString, ByteString)
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Just (ByteString
msg, ByteString
remainder) -> do
      case forall a. FromJSON a => ByteString -> Either String a
J.eitherDecode forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BSL.fromStrict ByteString
msg of
        Left String
err -> LogAction IO (WithSeverity LspServerLog)
ioLogger forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& String -> LspServerLog
DecodeInitializeError String
err forall msg. msg -> Severity -> WithSeverity msg
`WithSeverity` Severity
Error
        Right TRequestMessage @'ClientToServer 'Method_Initialize
initialize -> do
          Maybe (LanguageContextEnv config)
mInitResp <- forall config.
LogAction IO (WithSeverity LspProcessingLog)
-> ServerDefinition config
-> VFS
-> (FromServerMessage -> IO ())
-> TMessage @'ClientToServer @'Request 'Method_Initialize
-> IO (Maybe (LanguageContextEnv config))
Processing.initializeRequestHandler LogAction IO (WithSeverity LspProcessingLog)
pioLogger ServerDefinition config
serverDefinition VFS
vfs FromServerMessage -> IO ()
sendMsg TRequestMessage @'ClientToServer 'Method_Initialize
initialize
          case Maybe (LanguageContextEnv config)
mInitResp of
            Maybe (LanguageContextEnv config)
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            Just LanguageContextEnv config
env -> forall config (m :: * -> *) a.
LanguageContextEnv config -> LspT config m a -> m a
runLspT LanguageContextEnv config
env forall a b. (a -> b) -> a -> b
$ Result ByteString -> LspM config ()
loop (forall a. Parser a -> ByteString -> Result a
parse Parser ByteString ByteString
parser ByteString
remainder)
 where
  pioLogger :: LogAction IO (WithSeverity LspProcessingLog)
pioLogger = forall a b (m :: * -> *).
(a -> b) -> LogAction m b -> LogAction m a
L.cmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap LspProcessingLog -> LspServerLog
LspProcessingLog) LogAction IO (WithSeverity LspServerLog)
ioLogger
  pLogger :: LogAction (LspM config) (WithSeverity LspProcessingLog)
pLogger = forall a b (m :: * -> *).
(a -> b) -> LogAction m b -> LogAction m a
L.cmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap LspProcessingLog -> LspServerLog
LspProcessingLog) LogAction (LspM config) (WithSeverity LspServerLog)
logger

  loop :: Result BS.ByteString -> LspM config ()
  loop :: Result ByteString -> LspM config ()
loop = Result ByteString -> LspM config ()
go
   where
    go :: Result ByteString -> LspM config ()
go Result ByteString
r = do
      Maybe (ByteString, ByteString)
res <- forall (m :: * -> *).
MonadIO m =>
LogAction m (WithSeverity LspServerLog)
-> IO ByteString
-> Result ByteString
-> m (Maybe (ByteString, ByteString))
parseOne LogAction (LspM config) (WithSeverity LspServerLog)
logger IO ByteString
clientIn Result ByteString
r
      case Maybe (ByteString, ByteString)
res of
        Maybe (ByteString, ByteString)
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        Just (ByteString
msg, ByteString
remainder) -> do
          forall (m :: * -> *) config.
((m :: (* -> *)) ~ (LspM config :: (* -> *))) =>
LogAction m (WithSeverity LspProcessingLog) -> ByteString -> m ()
Processing.processMessage LogAction (LspM config) (WithSeverity LspProcessingLog)
pLogger forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BSL.fromStrict ByteString
msg
          Result ByteString -> LspM config ()
go (forall a. Parser a -> ByteString -> Result a
parse Parser ByteString ByteString
parser ByteString
remainder)

  parser :: Parser ByteString ByteString
parser = do
    forall i a. Parser i a -> Parser i a
try Parser ByteString ()
contentType forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (forall (m :: * -> *) a. Monad m => a -> m a
return ())
    Int
len <- Parser ByteString Int
contentLength
    forall i a. Parser i a -> Parser i a
try Parser ByteString ()
contentType forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (forall (m :: * -> *) a. Monad m => a -> m a
return ())
    ByteString
_ <- ByteString -> Parser ByteString ByteString
string ByteString
_ONE_CRLF
    Int -> Parser ByteString ByteString
Attoparsec.take Int
len

  contentLength :: Parser ByteString Int
contentLength = do
    ByteString
_ <- ByteString -> Parser ByteString ByteString
string ByteString
"Content-Length: "
    Int
len <- forall a. Integral a => Parser a
decimal
    ByteString
_ <- ByteString -> Parser ByteString ByteString
string ByteString
_ONE_CRLF
    forall (m :: * -> *) a. Monad m => a -> m a
return Int
len

  contentType :: Parser ByteString ()
contentType = do
    ByteString
_ <- ByteString -> Parser ByteString ByteString
string ByteString
"Content-Type: "
    (Char -> Bool) -> Parser ByteString ()
skipWhile (forall a. Eq a => a -> a -> Bool
/= Char
'\r')
    ByteString
_ <- ByteString -> Parser ByteString ByteString
string ByteString
_ONE_CRLF
    forall (m :: * -> *) a. Monad m => a -> m a
return ()

parseOne ::
  MonadIO m =>
  LogAction m (WithSeverity LspServerLog) ->
  IO BS.ByteString ->
  Result BS.ByteString ->
  m (Maybe (BS.ByteString, BS.ByteString))
parseOne :: forall (m :: * -> *).
MonadIO m =>
LogAction m (WithSeverity LspServerLog)
-> IO ByteString
-> Result ByteString
-> m (Maybe (ByteString, ByteString))
parseOne LogAction m (WithSeverity LspServerLog)
logger IO ByteString
clientIn = forall {a}. IResult ByteString a -> m (Maybe (a, ByteString))
go
 where
  go :: IResult ByteString a -> m (Maybe (a, ByteString))
go (Fail ByteString
_ [String]
ctxs String
err) = do
    LogAction m (WithSeverity LspServerLog)
logger forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& [String] -> String -> LspServerLog
HeaderParseFail [String]
ctxs String
err forall msg. msg -> Severity -> WithSeverity msg
`WithSeverity` Severity
Error
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
  go (Partial ByteString -> IResult ByteString a
c) = do
    ByteString
bs <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ByteString
clientIn
    if ByteString -> Bool
BS.null ByteString
bs
      then do
        LogAction m (WithSeverity LspServerLog)
logger forall (m :: * -> *) msg. LogAction m msg -> msg -> m ()
<& LspServerLog
EOF forall msg. msg -> Severity -> WithSeverity msg
`WithSeverity` Severity
Error
        forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
      else IResult ByteString a -> m (Maybe (a, ByteString))
go (ByteString -> IResult ByteString a
c ByteString
bs)
  go (Done ByteString
remainder a
msg) = do
    -- TODO: figure out how to re-enable
    -- This can lead to infinite recursion in logging, see https://github.com/haskell/lsp/issues/447
    -- logger <& ParsedMsg (T.decodeUtf8 msg) `WithSeverity` Debug
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (a
msg, ByteString
remainder)

-- ---------------------------------------------------------------------

-- | Simple server to make sure all output is serialised
sendServer :: LogAction IO (WithSeverity LspServerLog) -> TChan J.Value -> (BSL.ByteString -> IO ()) -> IO ()
sendServer :: LogAction IO (WithSeverity LspServerLog)
-> TChan Value -> (ByteString -> IO ()) -> IO ()
sendServer LogAction IO (WithSeverity LspServerLog)
_logger TChan Value
msgChan ByteString -> IO ()
clientOut = do
  forall (f :: * -> *) a b. Applicative f => f a -> f b
forever forall a b. (a -> b) -> a -> b
$ do
    Value
msg <- forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TChan a -> STM a
readTChan TChan Value
msgChan

    -- We need to make sure we only send over the content of the message,
    -- and no other tags/wrapper stuff
    let str :: ByteString
str = forall a. ToJSON a => a -> ByteString
J.encode Value
msg

    let out :: ByteString
out =
          [ByteString] -> ByteString
BSL.concat
            [ Text -> ByteString
TL.encodeUtf8 forall a b. (a -> b) -> a -> b
$ String -> Text
TL.pack forall a b. (a -> b) -> a -> b
$ String
"Content-Length: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (ByteString -> Int64
BSL.length ByteString
str)
            , ByteString -> ByteString
BSL.fromStrict ByteString
_TWO_CRLF
            , ByteString
str
            ]

    ByteString -> IO ()
clientOut ByteString
out

-- TODO: figure out how to re-enable
-- 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
_ONE_CRLF :: ByteString
_ONE_CRLF = ByteString
"\r\n"
_TWO_CRLF :: BS.ByteString
_TWO_CRLF :: ByteString
_TWO_CRLF = ByteString
"\r\n\r\n"