{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TypeFamilies #-}

-- | Wrapper to provide safer interface in constructing an MCP server.
module Unison.MCP.Wrapper
  ( Tool (..),
    Prompt (..),
    HasInputSchema (..),
    mkServer,
    CallToolResult (..),
    PromptArgument (..),
    StaticResources,
    Server,
    MCP.ServerCapabilities (..),
    MCP.ToolAnnotations (..),
    MCP.Implementation (..),
    MCP.ResourcesCapability (..),
    MCP.ToolsCapability (..),
    MCP.PromptsCapability (..),
    MCP.PromptContentType (..),
    errorToolResult,
    textToolResult,
    jsonToolResult,
  )
where

import Data.Aeson (FromJSON)
import Data.Aeson qualified as Aeson
import Data.ByteString.Lazy.Char8 qualified as BL
import Data.Data (Proxy)
import Data.Map qualified as Map
import Data.Text qualified as Text
import Network.MCP.Server
import Network.MCP.Types (CallToolResult (CallToolResult))
import Network.MCP.Types qualified as MCP
import Unison.Prelude
import UnliftIO qualified

type StaticResources = Map Text (MCP.Resource, MCP.ResourceContent)

class HasInputSchema arg where
  toInputSchema :: Proxy arg -> Aeson.Value

instance HasInputSchema () where
  toInputSchema :: Proxy () -> Value
toInputSchema Proxy ()
_ =
    [Pair] -> Value
Aeson.object
      [ (Key
"type", Text -> Value
Aeson.String Text
"object"),
        (Key
"properties", [Pair] -> Value
Aeson.object []),
        (Key
"required", Array -> Value
Aeson.Array Array
forall a. Monoid a => a
mempty)
      ]

data Tool m = forall arg.
  (FromJSON arg, HasInputSchema arg) =>
  Tool
  { forall (m :: * -> *). Tool m -> Text
toolName :: Text,
    forall (m :: * -> *). Tool m -> Text
toolDescription :: Text,
    forall (m :: * -> *). Tool m -> ToolAnnotations
toolAnnotations :: MCP.ToolAnnotations,
    ()
toolArgType :: Proxy arg,
    ()
toolHandler :: arg -> m MCP.CallToolResult
  }

data Prompt m = Prompt
  { forall (m :: * -> *). Prompt m -> Text
promptName :: Text,
    forall (m :: * -> *). Prompt m -> Text
promptDescription :: Text,
    forall (m :: * -> *). Prompt m -> Map Text PromptArgument
promptArgs :: Map Text PromptArgument,
    forall (m :: * -> *).
Prompt m -> Map Text Text -> m GetPromptResult
promptHandler :: Map Text Text -> m MCP.GetPromptResult
  }

data PromptArgument = PromptArgument
  { PromptArgument -> Text
promptArgumentDescription :: Text,
    -- | Whether the argument is required
    PromptArgument -> Bool
promptArgumentRequired :: Bool
  }

mkServer :: (MonadUnliftIO m) => MCP.ServerInfo -> Text -> StaticResources -> [Tool m] -> [Prompt m] -> m Server
mkServer :: forall (m :: * -> *).
MonadUnliftIO m =>
ServerInfo
-> Text -> StaticResources -> [Tool m] -> [Prompt m] -> m Server
mkServer ServerInfo
serverInfo Text
serverDescription StaticResources
staticResources [Tool m]
tools [Prompt m]
prompts = do
  let serverCapabilities :: ServerCapabilities
serverCapabilities =
        MCP.ServerCapabilities
          { $sel:resourcesCapability:ServerCapabilities :: Maybe ResourcesCapability
resourcesCapability = ResourcesCapability -> Maybe ResourcesCapability
forall a. a -> Maybe a
Just (ResourcesCapability -> Maybe ResourcesCapability)
-> ResourcesCapability -> Maybe ResourcesCapability
forall a b. (a -> b) -> a -> b
$ Bool -> ResourcesCapability
MCP.ResourcesCapability (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ StaticResources -> Bool
forall k a. Map k a -> Bool
Map.null StaticResources
staticResources),
            $sel:toolsCapability:ServerCapabilities :: Maybe ToolsCapability
toolsCapability = ToolsCapability -> Maybe ToolsCapability
forall a. a -> Maybe a
Just (ToolsCapability -> Maybe ToolsCapability)
-> ToolsCapability -> Maybe ToolsCapability
forall a b. (a -> b) -> a -> b
$ Bool -> ToolsCapability
MCP.ToolsCapability (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Tool m] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Tool m]
tools),
            $sel:promptsCapability:ServerCapabilities :: Maybe PromptsCapability
promptsCapability = PromptsCapability -> Maybe PromptsCapability
forall a. a -> Maybe a
Just (PromptsCapability -> Maybe PromptsCapability)
-> PromptsCapability -> Maybe PromptsCapability
forall a b. (a -> b) -> a -> b
$ Bool -> PromptsCapability
MCP.PromptsCapability (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [Prompt m] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Prompt m]
prompts)
          }
  Server
server <- IO Server -> m Server
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Server -> m Server) -> IO Server -> m Server
forall a b. (a -> b) -> a -> b
$ ServerInfo -> ServerCapabilities -> Text -> IO Server
createServer ServerInfo
serverInfo ServerCapabilities
serverCapabilities Text
serverDescription

  Server -> StaticResources -> m ()
forall (m :: * -> *).
MonadUnliftIO m =>
Server -> StaticResources -> m ()
doResources Server
server StaticResources
staticResources
  Server -> [Tool m] -> m ()
forall (m :: * -> *). MonadUnliftIO m => Server -> [Tool m] -> m ()
doTools Server
server [Tool m]
tools
  Server -> [Prompt m] -> m ()
forall (m :: * -> *).
MonadUnliftIO m =>
Server -> [Prompt m] -> m ()
doPrompts Server
server [Prompt m]
prompts

  Server -> m Server
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Server
server

doResources :: (MonadUnliftIO m) => Server -> StaticResources -> m ()
doResources :: forall (m :: * -> *).
MonadUnliftIO m =>
Server -> StaticResources -> m ()
doResources Server
server StaticResources
staticResources = do
  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Server -> [Resource] -> IO ()
registerResources Server
server ((Resource, ResourceContent) -> Resource
forall a b. (a, b) -> a
fst ((Resource, ResourceContent) -> Resource)
-> [(Resource, ResourceContent)] -> [Resource]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StaticResources -> [(Resource, ResourceContent)]
forall k a. Map k a -> [a]
Map.elems StaticResources
staticResources)

  -- Register resource read handler
  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Server -> ResourceReadHandler -> IO ()
registerResourceReadHandler Server
server (ResourceReadHandler -> IO ()) -> ResourceReadHandler -> IO ()
forall a b. (a -> b) -> a -> b
$ \(MCP.ReadResourceRequest {Text
resourceReadUri :: Text
$sel:resourceReadUri:ReadResourceRequest :: ReadResourceRequest -> Text
resourceReadUri}) -> do
    case Text -> StaticResources -> Maybe (Resource, ResourceContent)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
resourceReadUri StaticResources
staticResources of
      Just (Resource
_, ResourceContent
content) ->
        ReadResourceResult -> IO ReadResourceResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ReadResourceResult -> IO ReadResourceResult)
-> ([ResourceContent] -> ReadResourceResult)
-> [ResourceContent]
-> IO ReadResourceResult
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ResourceContent] -> ReadResourceResult
MCP.ReadResourceResult ([ResourceContent] -> IO ReadResourceResult)
-> [ResourceContent] -> IO ReadResourceResult
forall a b. (a -> b) -> a -> b
$ [ResourceContent
content]
      Maybe (Resource, ResourceContent)
_ -> ReadResourceResult -> IO ReadResourceResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ReadResourceResult -> IO ReadResourceResult)
-> ReadResourceResult -> IO ReadResourceResult
forall a b. (a -> b) -> a -> b
$ [ResourceContent] -> ReadResourceResult
MCP.ReadResourceResult []

doTools :: (MonadUnliftIO m) => Server -> [Tool m] -> m ()
doTools :: forall (m :: * -> *). MonadUnliftIO m => Server -> [Tool m] -> m ()
doTools Server
server [Tool m]
tools = do
  m CallToolResult -> IO CallToolResult
runInIO <- m (m CallToolResult -> IO CallToolResult)
forall (m :: * -> *) a. MonadUnliftIO m => m (m a -> IO a)
askRunInIO
  let toolMap :: Map Text (Tool m)
toolMap = [(Text, Tool m)] -> Map Text (Tool m)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([Tool m]
tools [Tool m] -> (Tool m -> (Text, Tool m)) -> [(Text, Tool m)]
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> (\Tool m
tool -> (Tool m -> Text
forall (m :: * -> *). Tool m -> Text
toolName Tool m
tool, Tool m
tool)))
  let mcpTools :: [Tool]
mcpTools =
        [Tool m]
tools [Tool m] -> (Tool m -> Tool) -> [Tool]
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \(Tool {Text
$sel:toolName:Tool :: forall (m :: * -> *). Tool m -> Text
toolName :: Text
toolName, Text
$sel:toolDescription:Tool :: forall (m :: * -> *). Tool m -> Text
toolDescription :: Text
toolDescription, ToolAnnotations
$sel:toolAnnotations:Tool :: forall (m :: * -> *). Tool m -> ToolAnnotations
toolAnnotations :: ToolAnnotations
toolAnnotations, Proxy arg
$sel:toolArgType:Tool :: ()
toolArgType :: Proxy arg
toolArgType}) ->
          MCP.Tool
            { Text
toolName :: Text
$sel:toolName:Tool :: Text
MCP.toolName,
              $sel:toolDescription:Tool :: Maybe Text
MCP.toolDescription = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
toolDescription,
              $sel:toolInputSchema:Tool :: Value
MCP.toolInputSchema = Proxy arg -> Value
forall {k} (arg :: k). HasInputSchema arg => Proxy arg -> Value
toInputSchema Proxy arg
toolArgType,
              $sel:toolAnnotations:Tool :: Maybe ToolAnnotations
MCP.toolAnnotations = ToolAnnotations -> Maybe ToolAnnotations
forall a. a -> Maybe a
Just ToolAnnotations
toolAnnotations
            }
  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Server -> [Tool] -> IO ()
registerTools Server
server [Tool]
mcpTools
  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Server -> ToolCallHandler -> IO ()
registerToolCallHandler Server
server \(MCP.CallToolRequest {Text
callToolName :: Text
$sel:callToolName:CallToolRequest :: CallToolRequest -> Text
callToolName, Value
callToolArguments :: Value
$sel:callToolArguments:CallToolRequest :: CallToolRequest -> Value
callToolArguments}) -> m CallToolResult -> IO CallToolResult
runInIO (m CallToolResult -> IO CallToolResult)
-> m CallToolResult -> IO CallToolResult
forall a b. (a -> b) -> a -> b
$ do
    case Text -> Map Text (Tool m) -> Maybe (Tool m)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
callToolName Map Text (Tool m)
toolMap of
      Just Tool {arg -> m CallToolResult
$sel:toolHandler:Tool :: ()
toolHandler :: arg -> m CallToolResult
toolHandler} -> do
        case Value -> Result arg
forall a. FromJSON a => Value -> Result a
Aeson.fromJSON Value
callToolArguments of
          Aeson.Success arg
arg ->
            Int -> m CallToolResult -> m (Maybe CallToolResult)
forall (m :: * -> *) a.
MonadUnliftIO m =>
Int -> m a -> m (Maybe a)
UnliftIO.timeout (Int
60 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1_000_000) (arg -> m CallToolResult
toolHandler arg
arg) m (Maybe CallToolResult)
-> (Maybe CallToolResult -> m CallToolResult) -> m CallToolResult
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
              Maybe CallToolResult
Nothing -> CallToolResult -> m CallToolResult
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (CallToolResult -> m CallToolResult)
-> CallToolResult -> m CallToolResult
forall a b. (a -> b) -> a -> b
$ Text -> CallToolResult
errorToolResult (Text -> CallToolResult) -> Text -> CallToolResult
forall a b. (a -> b) -> a -> b
$ Text
"Tool '" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
callToolName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"' timed out after 1 minute."
              Just CallToolResult
result -> CallToolResult -> m CallToolResult
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure CallToolResult
result
          Aeson.Error String
err -> CallToolResult -> m CallToolResult
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (CallToolResult -> m CallToolResult)
-> CallToolResult -> m CallToolResult
forall a b. (a -> b) -> a -> b
$ Text -> CallToolResult
errorToolResult (Text -> CallToolResult) -> Text -> CallToolResult
forall a b. (a -> b) -> a -> b
$ Text
"Failed to parse arguments for tool '" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
callToolName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"': " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
Text.pack String
err
      Maybe (Tool m)
Nothing -> CallToolResult -> m CallToolResult
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (CallToolResult -> m CallToolResult)
-> CallToolResult -> m CallToolResult
forall a b. (a -> b) -> a -> b
$ Text -> CallToolResult
errorToolResult (Text -> CallToolResult) -> Text -> CallToolResult
forall a b. (a -> b) -> a -> b
$ Text
"Tool '" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
callToolName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"' not found."

errorToolResult :: Text -> MCP.CallToolResult
errorToolResult :: Text -> CallToolResult
errorToolResult Text
errMsg =
  MCP.CallToolResult
    { $sel:callToolContent:CallToolResult :: [ToolContent]
MCP.callToolContent = [ToolContentType -> Maybe Text -> ToolContent
MCP.ToolContent ToolContentType
MCP.TextualContent (Maybe Text -> ToolContent) -> Maybe Text -> ToolContent
forall a b. (a -> b) -> a -> b
$ Text -> Maybe Text
forall a. a -> Maybe a
Just Text
errMsg],
      $sel:callToolIsError:CallToolResult :: Bool
MCP.callToolIsError = Bool
True
    }

doPrompts :: (MonadUnliftIO m) => Server -> [Prompt m] -> m ()
doPrompts :: forall (m :: * -> *).
MonadUnliftIO m =>
Server -> [Prompt m] -> m ()
doPrompts Server
server [Prompt m]
prompts = do
  let mcpPrompts :: [Prompt]
mcpPrompts =
        [Prompt m]
prompts [Prompt m] -> (Prompt m -> Prompt) -> [Prompt]
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \(Prompt {Text
$sel:promptName:Prompt :: forall (m :: * -> *). Prompt m -> Text
promptName :: Text
promptName, Text
$sel:promptDescription:Prompt :: forall (m :: * -> *). Prompt m -> Text
promptDescription :: Text
promptDescription, Map Text PromptArgument
$sel:promptArgs:Prompt :: forall (m :: * -> *). Prompt m -> Map Text PromptArgument
promptArgs :: Map Text PromptArgument
promptArgs}) ->
          MCP.Prompt
            { Text
promptName :: Text
$sel:promptName:Prompt :: Text
MCP.promptName,
              $sel:promptDescription:Prompt :: Maybe Text
MCP.promptDescription = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
promptDescription,
              $sel:promptArguments:Prompt :: [PromptArgument]
MCP.promptArguments =
                Map Text PromptArgument
promptArgs
                  Map Text PromptArgument
-> (Map Text PromptArgument -> [(Text, PromptArgument)])
-> [(Text, PromptArgument)]
forall a b. a -> (a -> b) -> b
& Map Text PromptArgument -> [(Text, PromptArgument)]
forall k a. Map k a -> [(k, a)]
Map.toList
                  [(Text, PromptArgument)]
-> ((Text, PromptArgument) -> PromptArgument) -> [PromptArgument]
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \(Text
argName, PromptArgument {Text
$sel:promptArgumentDescription:PromptArgument :: PromptArgument -> Text
promptArgumentDescription :: Text
promptArgumentDescription, Bool
$sel:promptArgumentRequired:PromptArgument :: PromptArgument -> Bool
promptArgumentRequired :: Bool
promptArgumentRequired}) ->
                    MCP.PromptArgument
                      { $sel:promptArgumentName:PromptArgument :: Text
MCP.promptArgumentName = Text
argName,
                        $sel:promptArgumentDescription:PromptArgument :: Maybe Text
MCP.promptArgumentDescription = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
promptArgumentDescription,
                        $sel:promptArgumentRequired:PromptArgument :: Bool
MCP.promptArgumentRequired = Bool
promptArgumentRequired
                      }
            }
  let promptsMap :: Map Text (Prompt m)
promptsMap = [(Text, Prompt m)] -> Map Text (Prompt m)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(Text, Prompt m)] -> Map Text (Prompt m))
-> [(Text, Prompt m)] -> Map Text (Prompt m)
forall a b. (a -> b) -> a -> b
$ [Prompt m]
prompts [Prompt m] -> (Prompt m -> (Text, Prompt m)) -> [(Text, Prompt m)]
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> (\Prompt m
p -> (Prompt m -> Text
forall (m :: * -> *). Prompt m -> Text
promptName Prompt m
p, Prompt m
p))
  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Server -> [Prompt] -> IO ()
registerPrompts Server
server [Prompt]
mcpPrompts
  m GetPromptResult -> IO GetPromptResult
runInIO <- m (m GetPromptResult -> IO GetPromptResult)
forall (m :: * -> *) a. MonadUnliftIO m => m (m a -> IO a)
askRunInIO
  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Server -> PromptHandler -> IO ()
registerPromptHandler Server
server (PromptHandler -> IO ()) -> PromptHandler -> IO ()
forall a b. (a -> b) -> a -> b
$ \(MCP.GetPromptRequest {Text
getPromptName :: Text
$sel:getPromptName:GetPromptRequest :: GetPromptRequest -> Text
getPromptName, Map Text Text
getPromptArguments :: Map Text Text
$sel:getPromptArguments:GetPromptRequest :: GetPromptRequest -> Map Text Text
getPromptArguments}) -> m GetPromptResult -> IO GetPromptResult
runInIO do
    case Text -> Map Text (Prompt m) -> Maybe (Prompt m)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
getPromptName Map Text (Prompt m)
promptsMap of
      Maybe (Prompt m)
Nothing -> String -> m GetPromptResult
forall a. HasCallStack => String -> a
error (String -> m GetPromptResult) -> String -> m GetPromptResult
forall a b. (a -> b) -> a -> b
$ String
"Prompt '" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Text -> String
Text.unpack Text
getPromptName String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"' not found."
      Just (Prompt {Map Text Text -> m GetPromptResult
$sel:promptHandler:Prompt :: forall (m :: * -> *).
Prompt m -> Map Text Text -> m GetPromptResult
promptHandler :: Map Text Text -> m GetPromptResult
promptHandler}) -> do
        Map Text Text -> m GetPromptResult
promptHandler Map Text Text
getPromptArguments

textToolResult :: Text -> MCP.CallToolResult
textToolResult :: Text -> CallToolResult
textToolResult Text
msg =
  MCP.CallToolResult
    { $sel:callToolContent:CallToolResult :: [ToolContent]
MCP.callToolContent = [ToolContentType -> Maybe Text -> ToolContent
MCP.ToolContent ToolContentType
MCP.TextualContent (Maybe Text -> ToolContent) -> Maybe Text -> ToolContent
forall a b. (a -> b) -> a -> b
$ Text -> Maybe Text
forall a. a -> Maybe a
Just Text
msg],
      $sel:callToolIsError:CallToolResult :: Bool
MCP.callToolIsError = Bool
False
    }

jsonToolResult :: (Aeson.ToJSON a) => a -> MCP.CallToolResult
jsonToolResult :: forall a. ToJSON a => a -> CallToolResult
jsonToolResult a
msg = Text -> CallToolResult
textToolResult (Text -> CallToolResult) -> Text -> CallToolResult
forall a b. (a -> b) -> a -> b
$ String -> Text
Text.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ ByteString -> String
BL.unpack (ByteString -> String) -> ByteString -> String
forall a b. (a -> b) -> a -> b
$ a -> ByteString
forall a. ToJSON a => a -> ByteString
Aeson.encode a
msg