Server.hs

{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}

module MCP.Server
  ( -- * Server Runtime
    runMcpServer
  , runMcpServerStdIn
  , handleMcpMessage

    -- * Re-exports
  , module MCP.Server.Types
  ) where

import           Control.Monad          (when)
import           Control.Monad.IO.Class (MonadIO, liftIO)
import           Data.Aeson
import qualified Data.ByteString.Lazy   as BSL
import qualified Data.Map               as Map
import           Data.Text              (Text)
import qualified Data.Text              as T
import qualified Data.Text.Encoding     as TE
import qualified Data.Text.IO           as TIO
import           System.IO              (hFlush, hPutStrLn, stderr, stdout)

import           MCP.Server.JsonRpc
import           MCP.Server.Protocol
import           MCP.Server.Types

-- | Extract a brief summary of a JSON-RPC message for logging
getMessageSummary :: JsonRpcMessage -> String
getMessageSummary (JsonRpcMessageRequest req) = 
  "Request[" ++ show (requestId req) ++ "] " ++ T.unpack (requestMethod req)
getMessageSummary (JsonRpcMessageNotification notif) = 
  "Notification " ++ T.unpack (notificationMethod notif)
getMessageSummary (JsonRpcMessageResponse resp) = 
  "Response[" ++ show (responseId resp) ++ "]"

-- | Validate protocol version and return negotiated version
validateProtocolVersion :: Text -> Either Text Text
validateProtocolVersion clientVersion
  | clientVersion == protocolVersion = Right protocolVersion  -- Exact match
  | clientVersion == "2024-11-05" = Right "2024-11-05"        -- Accept this version
  | otherwise = Left $ "Unsupported protocol version: " <> clientVersion <> ". Server supports: " <> protocolVersion

-- | Run an MCP server with custom input/output handlers
runMcpServer :: (MonadIO m)
             => McpServerInfo
             -> McpServerHandlers m
             -> (Text -> m ())      -- ^ Output handler
             -> m Text              -- ^ Input handler
             -> m ()
runMcpServer serverInfo handlers outputHandler inputHandler = do
  loop
  where
    loop = do
      input <- inputHandler
      when (not $ T.null $ T.strip input) $ do
        liftIO $ hPutStrLn stderr $ "Received request: " ++ T.unpack input
        case eitherDecode (BSL.fromStrict $ TE.encodeUtf8 input) of
          Left err -> liftIO $ hPutStrLn stderr $ "Parse error: " ++ err
          Right jsonValue -> do
            case parseJsonRpcMessage jsonValue of
              Left err -> liftIO $ hPutStrLn stderr $ "JSON-RPC parse error: " ++ err
              Right message -> do
                liftIO $ hPutStrLn stderr $ "Processing message: " ++ show (getMessageSummary message)
                response <- handleMcpMessage serverInfo handlers message
                case response of
                  Just responseMsg -> do
                    liftIO $ hPutStrLn stderr $ "Sending response for: " ++ show (getMessageSummary message)
                    outputHandler $ TE.decodeUtf8 $ BSL.toStrict $ encode $ encodeJsonRpcMessage responseMsg
                  Nothing -> liftIO $ hPutStrLn stderr $ "No response needed for: " ++ show (getMessageSummary message)
        loop

-- | Run an MCP server using stdin/stdout
runMcpServerStdIn :: McpServerInfo -> McpServerHandlers IO -> IO ()
runMcpServerStdIn serverInfo handlers =
  runMcpServer serverInfo handlers outputHandler inputHandler
  where
    outputHandler text = do
      TIO.putStrLn text
      hFlush stdout
    inputHandler = TIO.getLine

-- | Handle an MCP message and return a response if needed
handleMcpMessage :: (MonadIO m)
                 => McpServerInfo
                 -> McpServerHandlers m
                 -> JsonRpcMessage
                 -> m (Maybe JsonRpcMessage)
handleMcpMessage serverInfo handlers (JsonRpcMessageRequest req) = do
  response <- case requestMethod req of
    "initialize" -> handleInitialize serverInfo req
    "ping" -> handlePing req
    "prompts/list" -> handlePromptsList handlers req
    "prompts/get" -> handlePromptsGet handlers req
    "resources/list" -> handleResourcesList handlers req
    "resources/read" -> handleResourcesRead handlers req
    "tools/list" -> handleToolsList handlers req
    "tools/call" -> handleToolsCall handlers req
    method -> return $ makeErrorResponse (requestId req) $ JsonRpcError
      { errorCode = -32601
      , errorMessage = "Method not found: " <> method
      , errorData = Nothing
      }
  return $ Just $ JsonRpcMessageResponse response

handleMcpMessage _ _ (JsonRpcMessageNotification notif) = do
  case notificationMethod notif of
    "notifications/initialized" -> do
      liftIO $ hPutStrLn stderr "Received initialized notification - server is ready for operation"
      return ()
    _ -> do
      liftIO $ hPutStrLn stderr $ "Received unknown notification: " ++ T.unpack (notificationMethod notif)
      return ()
  return Nothing

handleMcpMessage _ _ (JsonRpcMessageResponse _) =
  return Nothing

-- | Handle initialize request
handleInitialize :: (MonadIO m) => McpServerInfo -> JsonRpcRequest -> m JsonRpcResponse
handleInitialize serverInfo req = do
  case requestParams req of
    Nothing -> return $ makeErrorResponse (requestId req) $ JsonRpcError
      { errorCode = -32602
      , errorMessage = "Missing required parameters for initialize"
      , errorData = Nothing
      }
    Just params ->
      case fromJSON params of
        Error err -> return $ makeErrorResponse (requestId req) $ JsonRpcError
          { errorCode = -32602
          , errorMessage = "Invalid initialize parameters: " <> T.pack err
          , errorData = Nothing
          }
        Success initReq -> do
          -- Check protocol version compatibility
          let clientVersion = initProtocolVersion initReq
          case validateProtocolVersion clientVersion of
            Left errorMsg -> return $ makeErrorResponse (requestId req) $ JsonRpcError
              { errorCode = -32602
              , errorMessage = errorMsg
              , errorData = Nothing
              }
            Right negotiatedVersion -> do
              liftIO $ hPutStrLn stderr $ "Client version: " ++ T.unpack clientVersion ++ ", using: " ++ T.unpack negotiatedVersion
              let capabilities = ServerCapabilities
                    { capabilityPrompts = Just $ PromptCapabilities { promptListChanged = Nothing }
                    , capabilityResources = Just $ ResourceCapabilities { resourceSubscribe = Nothing, resourceListChanged = Nothing }
                    , capabilityTools = Just $ ToolCapabilities { toolListChanged = Nothing }
                    , capabilityLogging = Nothing  -- Not supported yet
                    }
              let response = InitializeResponse
                    { initRespProtocolVersion = negotiatedVersion
                    , initRespCapabilities = capabilities
                    , initRespServerInfo = serverInfo
                    }
              return $ makeSuccessResponse (requestId req) (toJSON response)

-- | Handle ping request
handlePing :: (MonadIO m) => JsonRpcRequest -> m JsonRpcResponse
handlePing req = return $ makeSuccessResponse (requestId req) (toJSON PongResponse)

-- | Handle prompts/list request
handlePromptsList :: (MonadIO m) => McpServerHandlers m -> JsonRpcRequest -> m JsonRpcResponse
handlePromptsList handlers req =
  case prompts handlers of
    Nothing -> return $ makeErrorResponse (requestId req) $ JsonRpcError
      { errorCode = -32601
      , errorMessage = "Prompts not supported"
      , errorData = Nothing
      }
    Just (listHandler, _) -> do
      let cursor = case requestParams req of
            Just params -> case fromJSON params of
              Success listReq -> promptsListCursor listReq
              _               -> Nothing
            Nothing -> Nothing
      paginatedResult <- listHandler cursor
      let response = PromptsListResponse
            { promptsListPrompts = paginatedItems paginatedResult
            , promptsListNextCursor = paginatedNextCursor paginatedResult
            }
      return $ makeSuccessResponse (requestId req) (toJSON response)

-- | Handle prompts/get request
handlePromptsGet :: (MonadIO m) => McpServerHandlers m -> JsonRpcRequest -> m JsonRpcResponse
handlePromptsGet handlers req =
  case prompts handlers of
    Nothing -> return $ makeErrorResponse (requestId req) $ JsonRpcError
      { errorCode = -32601
      , errorMessage = "Prompts not supported"
      , errorData = Nothing
      }
    Just (_, getHandler) -> do
      case requestParams req of
        Nothing -> return $ makeErrorResponse (requestId req) $ JsonRpcError
          { errorCode = -32602
          , errorMessage = "Missing parameters"
          , errorData = Nothing
          }
        Just params ->
          case fromJSON params of
            Error err -> return $ makeErrorResponse (requestId req) $ JsonRpcError
              { errorCode = -32602
              , errorMessage = "Invalid parameters: " <> T.pack err
              , errorData = Nothing
              }
            Success getReq -> do
              let args = maybe [] (map (\(k, v) -> (k, T.pack $ show v)) . Map.toList) (promptsGetArguments getReq)
              result <- getHandler (promptsGetName getReq) args
              case result of
                Left err -> return $ makeErrorResponse (requestId req) $ JsonRpcError
                  { errorCode = errorCodeFromMcpError err
                  , errorMessage = errorMessageFromMcpError err
                  , errorData = Nothing
                  }
                Right content -> do
                  let response = PromptsGetResponse
                        { promptsGetDescription = Nothing
                        , promptsGetMessages = [PromptMessage RoleUser content]
                        }
                  return $ makeSuccessResponse (requestId req) (toJSON response)

-- | Handle resources/list request
handleResourcesList :: (MonadIO m) => McpServerHandlers m -> JsonRpcRequest -> m JsonRpcResponse
handleResourcesList handlers req =
  case resources handlers of
    Nothing -> return $ makeErrorResponse (requestId req) $ JsonRpcError
      { errorCode = -32601
      , errorMessage = "Resources not supported"
      , errorData = Nothing
      }
    Just (listHandler, _) -> do
      let cursor = case requestParams req of
            Just params -> case fromJSON params of
              Success listReq -> resourcesListCursor listReq
              _               -> Nothing
            Nothing -> Nothing
      paginatedResult <- listHandler cursor
      let response = ResourcesListResponse
            { resourcesListResources = paginatedItems paginatedResult
            , resourcesListNextCursor = paginatedNextCursor paginatedResult
            }
      return $ makeSuccessResponse (requestId req) (toJSON response)

-- | Handle resources/read request
handleResourcesRead :: (MonadIO m) => McpServerHandlers m -> JsonRpcRequest -> m JsonRpcResponse
handleResourcesRead handlers req =
  case resources handlers of
    Nothing -> return $ makeErrorResponse (requestId req) $ JsonRpcError
      { errorCode = -32601
      , errorMessage = "Resources not supported"
      , errorData = Nothing
      }
    Just (_, readHandler) -> do
      case requestParams req of
        Nothing -> return $ makeErrorResponse (requestId req) $ JsonRpcError
          { errorCode = -32602
          , errorMessage = "Missing parameters"
          , errorData = Nothing
          }
        Just params ->
          case fromJSON params of
            Error err -> return $ makeErrorResponse (requestId req) $ JsonRpcError
              { errorCode = -32602
              , errorMessage = "Invalid parameters: " <> T.pack err
              , errorData = Nothing
              }
            Success readReq -> do
              result <- readHandler (resourcesReadUri readReq)
              case result of
                Left err -> return $ makeErrorResponse (requestId req) $ JsonRpcError
                  { errorCode = errorCodeFromMcpError err
                  , errorMessage = errorMessageFromMcpError err
                  , errorData = Nothing
                  }
                Right content -> do
                  let response = ResourcesReadResponse
                        { resourcesReadContents = [content]
                        }
                  return $ makeSuccessResponse (requestId req) (toJSON response)

-- | Handle tools/list request
handleToolsList :: (MonadIO m) => McpServerHandlers m -> JsonRpcRequest -> m JsonRpcResponse
handleToolsList handlers req =
  case tools handlers of
    Nothing -> return $ makeErrorResponse (requestId req) $ JsonRpcError
      { errorCode = -32601
      , errorMessage = "Tools not supported"
      , errorData = Nothing
      }
    Just (listHandler, _) -> do
      let cursor = case requestParams req of
            Just params -> case fromJSON params of
              Success listReq -> toolsListCursor listReq
              _               -> Nothing
            Nothing -> Nothing
      paginatedResult <- listHandler cursor
      let response = ToolsListResponse
            { toolsListTools = paginatedItems paginatedResult
            , toolsListNextCursor = paginatedNextCursor paginatedResult
            }
      return $ makeSuccessResponse (requestId req) (toJSON response)

-- | Handle tools/call request
handleToolsCall :: (MonadIO m) => McpServerHandlers m -> JsonRpcRequest -> m JsonRpcResponse
handleToolsCall handlers req =
  case tools handlers of
    Nothing -> return $ makeErrorResponse (requestId req) $ JsonRpcError
      { errorCode = -32601
      , errorMessage = "Tools not supported"
      , errorData = Nothing
      }
    Just (_, callHandler) -> do
      case requestParams req of
        Nothing -> return $ makeErrorResponse (requestId req) $ JsonRpcError
          { errorCode = -32602
          , errorMessage = "Missing parameters"
          , errorData = Nothing
          }
        Just params ->
          case fromJSON params of
            Error err -> return $ makeErrorResponse (requestId req) $ JsonRpcError
              { errorCode = -32602
              , errorMessage = "Invalid parameters: " <> T.pack err
              , errorData = Nothing
              }
            Success callReq -> do
              let args = maybe [] (map (\(k, v) -> (k, T.pack $ show v)) . Map.toList) (toolsCallArguments callReq)
              result <- callHandler (toolsCallName callReq) args
              case result of
                Left err -> return $ makeErrorResponse (requestId req) $ JsonRpcError
                  { errorCode = errorCodeFromMcpError err
                  , errorMessage = errorMessageFromMcpError err
                  , errorData = Nothing
                  }
                Right content -> do
                  let response = ToolsCallResponse
                        { toolsCallContent = [content]
                        , toolsCallIsError = Nothing
                        }
                  return $ makeSuccessResponse (requestId req) (toJSON response)

-- | Convert MCP error to JSON-RPC error code
errorCodeFromMcpError :: Error -> Int
errorCodeFromMcpError (InvalidPromptName _)     = -32602
errorCodeFromMcpError (MissingRequiredParams _) = -32602
errorCodeFromMcpError (ResourceNotFound _)      = -32602
errorCodeFromMcpError (InternalError _)         = -32603
errorCodeFromMcpError (UnknownTool _)           = -32602
errorCodeFromMcpError (InvalidRequest _)        = -32600
errorCodeFromMcpError (MethodNotFound _)        = -32601
errorCodeFromMcpError (InvalidParams _)         = -32602

-- | Convert MCP error to JSON-RPC error message
errorMessageFromMcpError :: Error -> Text
errorMessageFromMcpError (InvalidPromptName msg) = "Invalid prompt name: " <> msg
errorMessageFromMcpError (MissingRequiredParams msg) = "Missing required parameters: " <> msg
errorMessageFromMcpError (ResourceNotFound msg) = "Resource not found: " <> msg
errorMessageFromMcpError (InternalError msg) = "Internal error: " <> msg
errorMessageFromMcpError (UnknownTool msg) = "Unknown tool: " <> msg
errorMessageFromMcpError (InvalidRequest msg) = "Invalid request: " <> msg
errorMessageFromMcpError (MethodNotFound msg) = "Method not found: " <> msg
errorMessageFromMcpError (InvalidParams msg) = "Invalid parameters: " <> msg