mcp_protocol.lua

--[[
    MCP Protocol Handler
    Implements the Model Context Protocol (MCP) using JSON-RPC 2.0
    Handles initialization, capability negotiation, and request routing
--]]

local rapidjson = require("rapidjson")
local logger = require("logger")

local MCPProtocol = {
    version = "2025-03-26", -- MCP protocol version
    serverInfo = {
        name = "koreader-mcp",
        version = "1.0.0",
    },
    capabilities = {},
    resources = nil,
    tools = nil,
    prompts = nil,
    clientFeatures = nil, -- For sampling, elicitation, logging
    initialized = false,

    -- Callback for sending notifications/requests to client
    sendToClient = nil,
}

function MCPProtocol:new(o)
    o = o or {}
    setmetatable(o, self)
    self.__index = self
    return o
end

function MCPProtocol:setResources(resources)
    self.resources = resources
    self.capabilities.resources = {
        subscribe = true,
        listChanged = true,
    }
end

function MCPProtocol:setTools(tools)
    self.tools = tools
    self.capabilities.tools = {
        listChanged = false,
    }
end

function MCPProtocol:setPrompts(prompts)
    self.prompts = prompts
    self.capabilities.prompts = {
        listChanged = false,
    }
end

function MCPProtocol:setClientFeatures(clientFeatures)
    self.clientFeatures = clientFeatures
    -- Server declares logging capability
    self.capabilities.logging = {}
end

-- Set callback for sending messages to client (notifications, requests)
function MCPProtocol:setSendCallback(callback)
    self.sendToClient = callback
    if self.clientFeatures then
        self.clientFeatures:setSendCallback(callback)
    end
end

function MCPProtocol:handleRequest(request)
    -- Parse JSON-RPC request
    local ok, jsonRequest = pcall(rapidjson.decode, request.body)
    if not ok or not jsonRequest then
        -- Log detailed error information for debugging
        local body_preview = request.body or "(empty)"
        if #body_preview > 200 then
            body_preview = body_preview:sub(1, 200) .. "... (truncated, total " .. #request.body .. " bytes)"
        end
        local error_msg = not ok and tostring(jsonRequest) or "JSON decoded to nil/false"
        logger.warn("Invalid JSON in MCP request:")
        logger.warn("  Error:", error_msg)
        logger.warn("  Body:", body_preview)
        logger.warn("  Content-Length:", request.headers and request.headers["content-length"] or "unknown")
        return self:createErrorResponse(nil, -32700, "Parse error")
    end

    -- Validate JSON-RPC structure
    if jsonRequest.jsonrpc ~= "2.0" then
        return self:createErrorResponse(jsonRequest.id, -32600, "Invalid Request")
    end

    -- Route request to appropriate handler
    local method = jsonRequest.method
    local id = jsonRequest.id
    local params = jsonRequest.params or {}

    logger.dbg("MCP Request:", method, "ID:", id)

    -- Handle different MCP methods
    if method == "initialize" then
        return self:handleInitialize(id, params)
    elseif method == "initialized" then
        -- Notification, no response needed
        self.initialized = true
        return self:createNotificationResponse()
    elseif method == "ping" then
        return self:handlePing(id)
    elseif method == "logging/setLevel" then
        return self:handleLoggingSetLevel(id, params)
    elseif method == "resources/list" then
        return self:handleResourcesList(id, params)
    elseif method == "resources/read" then
        return self:handleResourcesRead(id, params)
    elseif method == "resources/templates/list" then
        return self:handleResourcesTemplatesList(id, params)
    elseif method == "resources/subscribe" then
        return self:handleResourcesSubscribe(id, params)
    elseif method == "resources/unsubscribe" then
        return self:handleResourcesUnsubscribe(id, params)
    elseif method == "tools/list" then
        return self:handleToolsList(id, params)
    elseif method == "tools/call" then
        return self:handleToolsCall(id, params)
    elseif method == "prompts/list" then
        return self:handlePromptsList(id, params)
    elseif method == "prompts/get" then
        return self:handlePromptsGet(id, params)
    else
        return self:createErrorResponse(id, -32601, "Method not found")
    end
end

function MCPProtocol:handleInitialize(id, params)
    logger.info("MCP Initialize request from:", params.clientInfo and params.clientInfo.name or "unknown")

    -- Store client capabilities for client features
    if self.clientFeatures and params.capabilities then
        self.clientFeatures:setClientCapabilities(params.capabilities)
        logger.dbg("MCP: Client capabilities:", rapidjson.encode(params.capabilities))
    end

    local result = {
        protocolVersion = self.version,
        capabilities = self.capabilities,
        serverInfo = self.serverInfo,
    }

    return self:createSuccessResponse(id, result)
end

function MCPProtocol:handlePing(id)
    return self:createSuccessResponse(id, {})
end

function MCPProtocol:handleLoggingSetLevel(id, params)
    local level = params.level
    if not level then
        return self:createErrorResponse(id, -32602, "Invalid params: missing level")
    end

    if self.clientFeatures then
        local ok, err = self.clientFeatures:setLogLevel(level)
        if not ok then
            return self:createErrorResponse(id, -32602, err)
        end
    end

    return self:createSuccessResponse(id, {})
end

function MCPProtocol:handleResourcesList(id, params)
    if not self.resources then
        return self:createErrorResponse(id, -32603, "Resources not available")
    end

    local ok, result = pcall(function()
        return self.resources:list()
    end)

    if not ok then
        logger.err("Error listing resources:", result)
        return self:createErrorResponse(id, -32603, "Internal error")
    end

    return self:createSuccessResponse(id, { resources = result })
end

function MCPProtocol:handleResourcesRead(id, params)
    if not self.resources then
        return self:createErrorResponse(id, -32603, "Resources not available")
    end

    local uri = params.uri
    if not uri then
        return self:createErrorResponse(id, -32602, "Invalid params: missing uri")
    end

    local ok, result = pcall(function()
        return self.resources:read(uri)
    end)

    if not ok then
        logger.err("Error reading resource:", result)
        return self:createErrorResponse(id, -32603, "Internal error")
    end

    if not result then
        return self:createErrorResponse(id, -32602, "Resource not found")
    end

    return self:createSuccessResponse(id, { contents = result })
end

function MCPProtocol:handleResourcesTemplatesList(id, params)
    if not self.resources then
        return self:createErrorResponse(id, -32603, "Resources not available")
    end

    local ok, result = pcall(function()
        return self.resources:listTemplates()
    end)

    if not ok then
        logger.err("Error listing resource templates:", result)
        return self:createErrorResponse(id, -32603, "Internal error")
    end

    return self:createSuccessResponse(id, { resourceTemplates = result or {} })
end

function MCPProtocol:handleResourcesSubscribe(id, params)
    if not self.resources then
        return self:createErrorResponse(id, -32603, "Resources not available")
    end

    local uri = params.uri
    if not uri then
        return self:createErrorResponse(id, -32602, "Invalid params: missing uri")
    end

    local ok, result = pcall(function()
        return self.resources:subscribe(uri)
    end)

    if not ok then
        logger.err("Error subscribing to resource:", result)
        return self:createErrorResponse(id, -32603, "Internal error")
    end

    return self:createSuccessResponse(id, {})
end

function MCPProtocol:handleResourcesUnsubscribe(id, params)
    if not self.resources then
        return self:createErrorResponse(id, -32603, "Resources not available")
    end

    local uri = params.uri
    if not uri then
        return self:createErrorResponse(id, -32602, "Invalid params: missing uri")
    end

    local ok, result = pcall(function()
        return self.resources:unsubscribe(uri)
    end)

    if not ok then
        logger.err("Error unsubscribing from resource:", result)
        return self:createErrorResponse(id, -32603, "Internal error")
    end

    return self:createSuccessResponse(id, {})
end

function MCPProtocol:handleToolsList(id, params)
    if not self.tools then
        return self:createErrorResponse(id, -32603, "Tools not available")
    end

    local ok, result = pcall(function()
        return self.tools:list()
    end)

    if not ok then
        logger.err("Error listing tools:", result)
        return self:createErrorResponse(id, -32603, "Internal error")
    end

    return self:createSuccessResponse(id, { tools = result })
end

function MCPProtocol:handleToolsCall(id, params)
    if not self.tools then
        return self:createErrorResponse(id, -32603, "Tools not available")
    end

    local name = params.name
    local arguments = params.arguments or {}

    if not name then
        return self:createErrorResponse(id, -32602, "Invalid params: missing name")
    end

    local ok, result = pcall(function()
        return self.tools:call(name, arguments)
    end)

    if not ok then
        logger.err("Error calling tool:", result)
        return self:createErrorResponse(id, -32603, tostring(result))
    end

    if not result then
        return self:createErrorResponse(id, -32602, "Tool not found")
    end

    return self:createSuccessResponse(id, result)
end

function MCPProtocol:handlePromptsList(id, params)
    if not self.prompts then
        return self:createErrorResponse(id, -32603, "Prompts not available")
    end

    local ok, result = pcall(function()
        return self.prompts:list()
    end)

    if not ok then
        logger.err("Error listing prompts:", result)
        return self:createErrorResponse(id, -32603, "Internal error")
    end

    return self:createSuccessResponse(id, { prompts = result })
end

function MCPProtocol:handlePromptsGet(id, params)
    if not self.prompts then
        return self:createErrorResponse(id, -32603, "Prompts not available")
    end

    local name = params.name
    local arguments = params.arguments or {}

    if not name then
        return self:createErrorResponse(id, -32602, "Invalid params: missing name")
    end

    local ok, result, err = pcall(function()
        return self.prompts:get(name, arguments)
    end)

    if not ok then
        logger.err("Error getting prompt:", result)
        return self:createErrorResponse(id, -32603, "Internal error")
    end

    if not result then
        -- result is nil, err contains the error message from prompts:get
        return self:createErrorResponse(id, -32602, err or "Prompt not found")
    end

    return self:createSuccessResponse(id, result)
end

function MCPProtocol:createSuccessResponse(id, result)
    local response = {
        jsonrpc = "2.0",
        id = id,
        result = result,
    }

    return {
        status = 200,
        statusText = "OK",
        headers = {},
        body = rapidjson.encode(response),
    }
end

function MCPProtocol:createErrorResponse(id, code, message)
    local response = {
        jsonrpc = "2.0",
        id = id or rapidjson.null,
        error = {
            code = code,
            message = message,
        },
    }

    return {
        status = 200, -- JSON-RPC errors still use HTTP 200
        statusText = "OK",
        headers = {},
        body = rapidjson.encode(response),
    }
end

function MCPProtocol:createNotificationResponse()
    -- Notifications don't get responses, return empty 204
    return {
        status = 204,
        statusText = "No Content",
        headers = {},
        body = "",
    }
end

--------------------------------------------------------------------------------
-- Server → Client Notifications
--------------------------------------------------------------------------------

-- Send a resource update notification
function MCPProtocol:notifyResourceUpdated(uri)
    if self.clientFeatures then
        self.clientFeatures:notifyResourceUpdated(uri)
    end
end

-- Send a resource list changed notification
function MCPProtocol:notifyResourceListChanged()
    if self.clientFeatures then
        self.clientFeatures:notifyResourceListChanged()
    end
end

-- Send a progress notification
function MCPProtocol:notifyProgress(progressToken, progress, total, message)
    if self.clientFeatures then
        self.clientFeatures:notifyProgress(progressToken, progress, total, message)
    end
end

-- Log a message to the client
function MCPProtocol:log(level, loggerName, data)
    if self.clientFeatures then
        self.clientFeatures:log(level, loggerName, data)
    end
end

--------------------------------------------------------------------------------
-- Server → Client Requests (via clientFeatures)
--------------------------------------------------------------------------------

-- Request sampling from the client
function MCPProtocol:requestSampling(options, callback)
    if self.clientFeatures then
        return self.clientFeatures:requestSampling(options, callback)
    else
        if callback then callback(nil, "Client features not available") end
    end
end

-- Simple ask helper
function MCPProtocol:ask(prompt, callback)
    if self.clientFeatures then
        return self.clientFeatures:ask(prompt, callback)
    else
        if callback then callback(nil, "Client features not available") end
    end
end

-- Request form elicitation
function MCPProtocol:requestFormElicitation(options, callback)
    if self.clientFeatures then
        return self.clientFeatures:requestFormElicitation(options, callback)
    else
        if callback then callback(nil, "Client features not available") end
    end
end

-- Request user confirmation
function MCPProtocol:confirm(message, callback)
    if self.clientFeatures then
        return self.clientFeatures:confirm(message, callback)
    else
        if callback then callback(false, "Client features not available") end
    end
end

-- Request text input
function MCPProtocol:requestText(message, options, callback)
    if self.clientFeatures then
        return self.clientFeatures:requestText(message, options, callback)
    else
        if callback then callback(nil, "Client features not available") end
    end
end

-- Request choice from options
function MCPProtocol:requestChoice(message, options, callback)
    if self.clientFeatures then
        return self.clientFeatures:requestChoice(message, options, callback)
    else
        if callback then callback(nil, "Client features not available") end
    end
end

-- Handle response from client (for pending sampling/elicitation requests)
function MCPProtocol:handleClientResponse(response)
    if self.clientFeatures then
        return self.clientFeatures:handleResponse(response)
    end
    return false
end

-- Handle notification from client
function MCPProtocol:handleClientNotification(notification)
    if self.clientFeatures then
        return self.clientFeatures:handleNotification(notification)
    end
    return false
end

return MCPProtocol