{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE GADTs #-}

module Network.MCP.Client
    ( createClient
    , connectClient
    , disconnectClient
    , listTools
    , callTool
    , listResources
    , readResource
    , listPrompts
    , getPrompt
    ) where

import Control.Concurrent.MVar
import Control.Exception (throw)
import Control.Monad (void)
import Data.Aeson
import System.IO
import System.Process
import Network.MCP.Client.Request
import Network.MCP.Client.Types
import qualified Data.Map.Strict as Map
import qualified Data.Text as T
import qualified Data.ByteString.Lazy.Char8 as BLC

-- | Create a new client
createClient :: ClientConfig -> IO Client
createClient :: ClientConfig -> IO Client
createClient ClientConfig
config = ClientConfig
-> MVar (Maybe ProcessHandle)
-> MVar (Maybe Handle)
-> MVar (Maybe Handle)
-> Client
Client
    (ClientConfig
 -> MVar (Maybe ProcessHandle)
 -> MVar (Maybe Handle)
 -> MVar (Maybe Handle)
 -> Client)
-> IO ClientConfig
-> IO
     (MVar (Maybe ProcessHandle)
      -> MVar (Maybe Handle) -> MVar (Maybe Handle) -> Client)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ClientConfig -> IO ClientConfig
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ClientConfig
config
    IO
  (MVar (Maybe ProcessHandle)
   -> MVar (Maybe Handle) -> MVar (Maybe Handle) -> Client)
-> IO (MVar (Maybe ProcessHandle))
-> IO (MVar (Maybe Handle) -> MVar (Maybe Handle) -> Client)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe ProcessHandle -> IO (MVar (Maybe ProcessHandle))
forall a. a -> IO (MVar a)
newMVar Maybe ProcessHandle
forall a. Maybe a
Nothing
    IO (MVar (Maybe Handle) -> MVar (Maybe Handle) -> Client)
-> IO (MVar (Maybe Handle)) -> IO (MVar (Maybe Handle) -> Client)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Handle -> IO (MVar (Maybe Handle))
forall a. a -> IO (MVar a)
newMVar Maybe Handle
forall a. Maybe a
Nothing
    IO (MVar (Maybe Handle) -> Client)
-> IO (MVar (Maybe Handle)) -> IO Client
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Handle -> IO (MVar (Maybe Handle))
forall a. a -> IO (MVar a)
newMVar Maybe Handle
forall a. Maybe a
Nothing

-- | Connect to the MCP server via StdIO
connectClient :: Client -> FilePath -> [String] -> IO Client
connectClient :: Client -> FilePath -> [FilePath] -> IO Client
connectClient client :: Client
client@Client{MVar (Maybe Handle)
MVar (Maybe ProcessHandle)
ClientConfig
clientConfig :: ClientConfig
clientProcess :: MVar (Maybe ProcessHandle)
clientStdin :: MVar (Maybe Handle)
clientStdout :: MVar (Maybe Handle)
$sel:clientConfig:Client :: Client -> ClientConfig
$sel:clientProcess:Client :: Client -> MVar (Maybe ProcessHandle)
$sel:clientStdin:Client :: Client -> MVar (Maybe Handle)
$sel:clientStdout:Client :: Client -> MVar (Maybe Handle)
..} FilePath
cmd [FilePath]
args = do

    -- TODO:
    --
    --   Create a pipe for interprocess communication and return a (readEnd,
    --   writeEnd) Handle pair.
    --
    --   WinIO Support hen this function is used with WinIO enabled it's
    --   Wthe caller's responsibility to register the handles with the
    --   WI/O manager. If this is not done the operation will deadlock.
    --   WAssociation can be done as follows:
    --
    -- https://hackage.haskell.org/package/process-1.6.25.0/docs/System-Process.html#g:8

    (Just Handle
hstdin, Just Handle
hstdout, Maybe Handle
Nothing, ProcessHandle
ph) <-
        CreateProcess
-> IO (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
createProcess (FilePath -> [FilePath] -> CreateProcess
proc FilePath
cmd [FilePath]
args)
            { std_in = CreatePipe
            , std_out = CreatePipe
            -- was std_err meant to be used for debugging or comms?
            -- our implementation requires stderr be use
            , std_err = Inherit
            }

    -- Send initialization message
    let initMessage :: Value
initMessage = [Pair] -> Value
object
            [ Key
"jsonrpc" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Text
"2.0" :: T.Text)
            , Key
"method" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Text
"initialize" :: T.Text)
            , Key
"params" Key -> Value -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= [Pair] -> Value
object
                [ Key
"clientInfo" Key -> Value -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= [Pair] -> Value
object
                    [ Key
"name" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= ClientConfig -> Text
clientName ClientConfig
clientConfig
                    , Key
"version" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= ClientConfig -> Text
clientVersion ClientConfig
clientConfig
                    ]
                , Key
"capabilities" Key -> Value -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= ClientConfig -> Value
clientCapabilities ClientConfig
clientConfig
                ]
            , Key
"id" Key -> Int -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Int
1 :: Int)
            ]

    -- Write initialization message
    Handle -> ByteString -> IO ()
BLC.hPutStrLn Handle
hstdin (Value -> ByteString
forall a. ToJSON a => a -> ByteString
encode Value
initMessage)
    Handle -> IO ()
hFlush Handle
hstdin

    IO (Maybe ProcessHandle) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Maybe ProcessHandle) -> IO ())
-> IO (Maybe ProcessHandle) -> IO ()
forall a b. (a -> b) -> a -> b
$ MVar (Maybe ProcessHandle)
-> Maybe ProcessHandle -> IO (Maybe ProcessHandle)
forall a. MVar a -> a -> IO a
swapMVar MVar (Maybe ProcessHandle)
clientProcess (ProcessHandle -> Maybe ProcessHandle
forall a. a -> Maybe a
Just ProcessHandle
ph)
    IO (Maybe Handle) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Maybe Handle) -> IO ()) -> IO (Maybe Handle) -> IO ()
forall a b. (a -> b) -> a -> b
$ MVar (Maybe Handle) -> Maybe Handle -> IO (Maybe Handle)
forall a. MVar a -> a -> IO a
swapMVar MVar (Maybe Handle)
clientStdin (Handle -> Maybe Handle
forall a. a -> Maybe a
Just Handle
hstdin)
    IO (Maybe Handle) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Maybe Handle) -> IO ()) -> IO (Maybe Handle) -> IO ()
forall a b. (a -> b) -> a -> b
$ MVar (Maybe Handle) -> Maybe Handle -> IO (Maybe Handle)
forall a. MVar a -> a -> IO a
swapMVar MVar (Maybe Handle)
clientStdout (Handle -> Maybe Handle
forall a. a -> Maybe a
Just Handle
hstdout)

    return Client
client

-- | Disconnect from the server
disconnectClient :: Client -> IO ()
disconnectClient :: Client -> IO ()
disconnectClient Client{MVar (Maybe Handle)
MVar (Maybe ProcessHandle)
ClientConfig
$sel:clientConfig:Client :: Client -> ClientConfig
$sel:clientProcess:Client :: Client -> MVar (Maybe ProcessHandle)
$sel:clientStdin:Client :: Client -> MVar (Maybe Handle)
$sel:clientStdout:Client :: Client -> MVar (Maybe Handle)
clientConfig :: ClientConfig
clientProcess :: MVar (Maybe ProcessHandle)
clientStdin :: MVar (Maybe Handle)
clientStdout :: MVar (Maybe Handle)
..} = do
    -- Close stdin and stdout
    MVar (Maybe Handle) -> IO (Maybe Handle)
forall a. MVar a -> IO a
readMVar MVar (Maybe Handle)
clientStdin IO (Maybe Handle) -> (Maybe Handle -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Handle -> IO ()) -> Maybe Handle -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Handle -> IO ()
hClose
    MVar (Maybe Handle) -> IO (Maybe Handle)
forall a. MVar a -> IO a
readMVar MVar (Maybe Handle)
clientStdout IO (Maybe Handle) -> (Maybe Handle -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Handle -> IO ()) -> Maybe Handle -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Handle -> IO ()
hClose

    -- Terminate process
    MVar (Maybe ProcessHandle) -> IO (Maybe ProcessHandle)
forall a. MVar a -> IO a
readMVar MVar (Maybe ProcessHandle)
clientProcess IO (Maybe ProcessHandle) -> (Maybe ProcessHandle -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (ProcessHandle -> IO ()) -> Maybe ProcessHandle -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ProcessHandle -> IO ()
terminateProcess

-- | List available tools
listTools :: Client -> IO [Value]
listTools :: Client -> IO [Value]
listTools Client
client = do
    let message :: Value
message = [Pair] -> Value
object
            [ Key
"jsonrpc" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Text
"2.0" :: T.Text)
            , Key
"method" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Text
"tools/list" :: T.Text)
            , Key
"id" Key -> Int -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Int
2 :: Int)
            ]

    -- Send message and handle response
    Value
response <- Client -> Value -> IO Value
sendClientRequest Client
client Value
message

    -- Extract tools from response
    case Value -> Result [Value]
forall a. FromJSON a => Value -> Result a
fromJSON Value
response of
        Success [Value]
tools -> [Value] -> IO [Value]
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return [Value]
tools
        Error FilePath
err -> McpClientError -> IO [Value]
forall a e. Exception e => e -> a
throw (McpClientError -> IO [Value]) -> McpClientError -> IO [Value]
forall a b. (a -> b) -> a -> b
$ Text -> McpClientError
ProtocolError (FilePath -> Text
T.pack FilePath
err)

-- | Call a tool with arguments
callTool :: Client -> T.Text -> Map.Map T.Text Value -> IO Value
callTool :: Client -> Text -> Map Text Value -> IO Value
callTool Client
client Text
toolName Map Text Value
args = do
    let message :: Value
message = [Pair] -> Value
object
            [ Key
"jsonrpc" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Text
"2.0" :: T.Text)
            , Key
"method" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Text
"tools/call" :: T.Text)
            , Key
"params" Key -> Value -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= [Pair] -> Value
object
                [ Key
"name" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= Text
toolName
                , Key
"arguments" Key -> Map Text Value -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= Map Text Value
args
                ]
            , Key
"id" Key -> Int -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Int
3 :: Int)
            ]

    -- Send message and handle response
    Client -> Value -> IO Value
sendClientRequest Client
client Value
message

-- | List available resources
listResources :: Client -> IO [Value]
listResources :: Client -> IO [Value]
listResources Client
client = do
    let message :: Value
message = [Pair] -> Value
object
            [ Key
"jsonrpc" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Text
"2.0" :: T.Text)
            , Key
"method" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Text
"resources/list" :: T.Text)
            , Key
"id" Key -> Int -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Int
4 :: Int)
            ]

    -- Send message and handle response
    Value
response <- Client -> Value -> IO Value
sendClientRequest Client
client Value
message

    -- Extract resources from response
    case Value -> Result [Value]
forall a. FromJSON a => Value -> Result a
fromJSON Value
response of
        Success [Value]
resources -> [Value] -> IO [Value]
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return [Value]
resources
        Error FilePath
err -> McpClientError -> IO [Value]
forall a e. Exception e => e -> a
throw (McpClientError -> IO [Value]) -> McpClientError -> IO [Value]
forall a b. (a -> b) -> a -> b
$ Text -> McpClientError
ProtocolError (FilePath -> Text
T.pack FilePath
err)

-- | Read a specific resource
readResource :: Client -> T.Text -> IO Value
readResource :: Client -> Text -> IO Value
readResource Client
client Text
resourceUri = do
    let message :: Value
message = [Pair] -> Value
object
            [ Key
"jsonrpc" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Text
"2.0" :: T.Text)
            , Key
"method" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Text
"resources/read" :: T.Text)
            , Key
"params" Key -> Value -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= [Pair] -> Value
object
                [ Key
"uri" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= Text
resourceUri ]
            , Key
"id" Key -> Int -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Int
5 :: Int)
            ]

    -- Send message and handle response
    Client -> Value -> IO Value
sendClientRequest Client
client Value
message




listPrompts :: Client -> IO [Value]
listPrompts :: Client -> IO [Value]
listPrompts Client
client = do
    -- Get stdin handle

    -- Create the request
    let message :: Value
message = [Pair] -> Value
object
            [ Key
"jsonrpc" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Text
"2.0" :: T.Text)
            , Key
"method" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Text
"prompts/list" :: T.Text)
            , Key
"id" Key -> Int -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Int
6 :: Int)
            ]

    -- Send message and handle response
    Value
response <- Client -> Value -> IO Value
sendClientRequest Client
client Value
message

    -- Extract prompts from response
    case Value -> Result [Value]
forall a. FromJSON a => Value -> Result a
fromJSON Value
response of
        Success [Value]
prompts -> [Value] -> IO [Value]
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return [Value]
prompts
        Error FilePath
err -> McpClientError -> IO [Value]
forall a e. Exception e => e -> a
throw (McpClientError -> IO [Value]) -> McpClientError -> IO [Value]
forall a b. (a -> b) -> a -> b
$ Text -> McpClientError
ProtocolError (FilePath -> Text
T.pack FilePath
err)

-- | Get a specific prompt
getPrompt :: Client -> T.Text -> Map.Map T.Text Value -> IO Value
getPrompt :: Client -> Text -> Map Text Value -> IO Value
getPrompt Client
client Text
promptName Map Text Value
args = do
    let message :: Value
message = [Pair] -> Value
object
            [ Key
"jsonrpc" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Text
"2.0" :: T.Text)
            , Key
"method" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Text
"prompts/get" :: T.Text)
            , Key
"params" Key -> Value -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= [Pair] -> Value
object
                [ Key
"name" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= Text
promptName
                , Key
"arguments" Key -> Map Text Value -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= Map Text Value
args
                ]
            , Key
"id" Key -> Int -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
forall v. ToJSON v => Key -> v -> Pair
.= (Int
7 :: Int)
            ]

    -- Send message and handle response
    Client -> Value -> IO Value
sendClientRequest Client
client Value
message