Skip to content

feat: mcp server framework implementation #12168

New issue

Have a question about this project? # for a free GitHub account to open an issue and contact its maintainers and the community.

#

By clicking “#”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? # to your account

Draft
wants to merge 8 commits into
base: master
Choose a base branch
from
+465 −276
Draft

feat: mcp server framework implementation #12168

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
5102e81
feat: add mcp server implementation
bzp2010 Apr 21, 2025
08e4a72
update
bzp2010 Apr 22, 2025
4f02ba7
chore: add license header
bzp2010 Apr 25, 2025
c7c6f89
chore: add lua install
bzp2010 Apr 25, 2025
c744183
fix: lint
bzp2010 Apr 25, 2025
0fa0564
fix: lint
bzp2010 Apr 25, 2025
040d6ac
fix: install
bzp2010 Apr 25, 2025
e08b714
fix: broker utils
bzp2010 Apr 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion Makefile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -382,8 +382,11 @@ install: runtime
$(ENV_INSTALL) -d $(ENV_INST_LUADIR)/apisix/plugins/ai-rag/vector-search
$(ENV_INSTALL) apisix/plugins/ai-rag/vector-search/*.lua $(ENV_INST_LUADIR)/apisix/plugins/ai-rag/vector-search

$(ENV_INSTALL) -d $(ENV_INST_LUADIR)/apisix/plugins/mcp
$(ENV_INSTALL) -d $(ENV_INST_LUADIR)/apisix/plugins/mcp/broker
$(ENV_INSTALL) -d $(ENV_INST_LUADIR)/apisix/plugins/mcp/transport
$(ENV_INSTALL) apisix/plugins/mcp/*.lua $(ENV_INST_LUADIR)/apisix/plugins/mcp
$(ENV_INSTALL) apisix/plugins/mcp/broker/*.lua $(ENV_INST_LUADIR)/apisix/plugins/mcp/broker
$(ENV_INSTALL) apisix/plugins/mcp/transport/*.lua $(ENV_INST_LUADIR)/apisix/plugins/mcp/transport

$(ENV_INSTALL) bin/apisix $(ENV_INST_BINDIR)/apisix

Expand Down
239 changes: 86 additions & 153 deletions apisix/plugins/mcp-bridge.lua
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,18 +15,14 @@
-- limitations under the License.
--
local unpack = unpack
local type = type
local tostring = tostring
local ngx = ngx
local re_match = ngx.re.match
local thread_spawn = ngx.thread.spawn
local thread_kill = ngx.thread.kill
local resty_signal = require("resty.signal")
local core = require("apisix.core")
local pipe = require("ngx.pipe")

local mcp_session_manager = require("apisix.plugins.mcp.session")

local V241105_ENDPOINT_SSE = "sse"
local V241105_ENDPOINT_MESSAGE = "message"
local mcp_server_wrapper = require("apisix.plugins.mcp.server_wrapper")

local schema = {
type = "object",
Expand Down Expand Up @@ -68,171 +64,108 @@ function _M.check_schema(conf, schema_type)
end


local function sse_send(id, event, data)
local ok, err = ngx.print((id and "id: " .. id .. "\n" or "") ..
"event: " .. event .. "\ndata: " .. data .. "\n\n")
if not ok then
return ok, "failed to write buffer: " .. err
end
return ngx.flush(true)
end


local function sse_handler(conf, ctx)
-- TODO: recover by Last-Event-ID
local session = mcp_session_manager.new()

-- spawn subprocess
local proc, err = pipe.spawn({conf.command, unpack(conf.args or {})})
if not proc then
core.log.error("failed to spawn mcp process: ", err)
return 500
end
proc:set_timeouts(nil, 100, 100)

core.response.set_header("Content-Type", "text/event-stream")
core.response.set_header("Cache-Control", "no-cache")

-- send endpoint event to advertise the message endpoint
sse_send(nil, "endpoint", conf.base_uri .. "/message?sessionId=" .. session.id)

local stdout_partial, stderr_partial

-- enter loop
while true do
if session:session_need_ping() then
local next_ping_id, err = session:session_next_ping_id()
if not next_ping_id then
core.log.error("session ", session.id, " exit, failed to get next ping id: ", err)
break
end
local ok, err = sse_send(nil, "message",
'{"jsonrpc": "2.0","method": "ping","id":"ping:'..next_ping_id..'"}')
if not ok then
core.log.info("session ", session.id, " exit, failed to send ping message: ", err)
break
end
local function on_connect(conf, ctx)
return function(additional)
local proc, err = pipe.spawn({conf.command, unpack(conf.args or {})})
if not proc then
core.log.error("failed to spawn mcp process: ", err)
return 500
end
if session:session_timed_out() then
core.log.info("session ", session.id, " exit, timed out")
break
end

-- pop the message from client in the queue and send it to the mcp server
repeat
local queue_item, err = session:pop_message_queue()
if err then
core.log.info("session ", session.id,
" exit, failed to pop message from queue: ", err)
break
end
-- write task message to stdio
if queue_item and type(queue_item) == "string" then
core.log.info("session ", session.id, " send message to mcp server: ", queue_item)
proc:write(queue_item .. "\n")
end
until not queue_item

-- read all the messages in stdout's pipe, line by line
-- if there is an incomplete message it is buffered and spliced before the next message
repeat
local line, _
line, _, stdout_partial = proc:stdout_read_line()
if line then
local ok, err = sse_send(nil, "message",
stdout_partial and stdout_partial .. line or line)
if not ok then
core.log.info("session ", session.id,
" exit, failed to send response message: ", err)
proc:set_timeouts(nil, 100, 100)
ctx.mcp_bridge_proc = proc

local server = additional.server

-- ngx_pipe is a yield operation, so we no longer need
-- to explicitly yield to other threads by ngx_sleep
ctx.mcp_bridge_proc_event_loop = thread_spawn(function ()
local stdout_partial, stderr_partial, need_exit
while true do
-- read all the messages in stdout's pipe, line by line
-- if there is an incomplete message it is buffered and
-- spliced before the next message
repeat
local line, _
line, _, stdout_partial = proc:stdout_read_line()
if line then
local ok, err = server.transport:send(
stdout_partial and stdout_partial .. line or line
)
if not ok then
core.log.info("session ", server.session_id,
" exit, failed to send response message: ", err)
need_exit = true
break
end
stdout_partial = nil -- luacheck: ignore
end
until not line
if need_exit then
break
end
stdout_partial = nil -- luacheck: ignore
end
until not line

repeat
local line, _
line, _, stderr_partial = proc:stderr_read_line()
if line then
local ok, err = sse_send(nil, "message",
'{"jsonrpc":"2.0","method":"notifications/stderr","params":{"content":"'
.. (stderr_partial and stderr_partial .. line or line) .. '"}}'
)
if not ok then
core.log.info("session ", session.id,
" exit, failed to send response message: ", err)

repeat
local line, _
line, _, stderr_partial = proc:stderr_read_line()
if line then
local ok, err = server.transport:send(
'{"jsonrpc":"2.0","method":"notifications/stderr","params":{"content":"'
.. (stderr_partial and stderr_partial .. line or line) .. '"}}')
if not ok then
core.log.info("session ", server.session_id,
" exit, failed to send response message: ", err)
need_exit = true
break
end
stderr_partial = "" -- luacheck: ignore
end
until not line
if need_exit then
break
end
stderr_partial = "" -- luacheck: ignore
end
until not line
end

session:destroy()

-- shutdown the subprocess
proc:shutdown("stdin")
proc:wait()
local _, err = proc:wait() -- check if process not exited then kill it
if err ~= "exited" then
proc:kill(resty_signal.signum("KILL") or 9)
end)
end
end


local function message_handler(conf, ctx)
local session_id = ctx.var.arg_sessionId
local session, err = mcp_session_manager.recover(session_id)

if not session then
core.log.error("failed to recover session: ", err)
return 404
local function on_client_message(conf, ctx)
return function(message, additional)
core.log.info("session ", additional.server.session_id,
" send message to mcp server: ", additional.raw)
ctx.mcp_bridge_proc:write(additional.raw .. "\n")
end
end

local body = core.request.get_body(nil, ctx)
if not body then
return 400
end

local body_json = core.json.decode(body)
if not body_json then
return 400
end
if core.string.has_prefix(tostring(body_json.id), "ping") then --TODO check client pong
session:on_session_pong()
return 202
end
local function on_disconnect(conf, ctx)
return function()
if ctx.mcp_bridge_proc_event_loop then
thread_kill(ctx.mcp_bridge_proc_event_loop)
ctx.mcp_bridge_proc_event_loop = nil
end

local ok, err = session:push_message_queue(body)
if not ok then
core.log.error("failed to add task to queue: ", err)
return 500
local proc = ctx.mcp_bridge_proc
if proc then
proc:shutdown("stdin")
proc:wait()
local _, err = proc:wait() -- check if process not exited then kill it
if err ~= "exited" then
proc:kill(resty_signal.signum("KILL") or 9)
end
end
end

return 202
end


function _M.access(conf, ctx)
local m, err = re_match(ctx.var.uri, "^" .. conf.base_uri .. "/(.*)", "jo")
if err then
core.log.info("failed to mcp base uri: ", err)
return 404
end
local action = m and m[1] or false
if not action then
return 404
end

if action == V241105_ENDPOINT_SSE and core.request.get_method() == "GET" then
return sse_handler(conf, ctx)
end

if action == V241105_ENDPOINT_MESSAGE and core.request.get_method() == "POST" then
return message_handler(conf, ctx)
end

return 200
return mcp_server_wrapper.access(conf, ctx, {
event_handler = {
on_connect = on_connect(conf, ctx),
on_client_message = on_client_message(conf, ctx),
on_disconnect = on_disconnect(conf, ctx),
},
})
end


Expand Down
89 changes: 89 additions & 0 deletions apisix/plugins/mcp/broker/shared_dict.lua
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
--
-- Licensed to the Apache Software Foundation (ASF) under one or more
-- contributor license agreements. See the NOTICE file distributed with
-- this work for additional information regarding copyright ownership.
-- The ASF licenses this file to You under the Apache License, Version 2.0
-- (the "License"); you may not use this file except in compliance with
-- the License. You may obtain a copy of the License at
--
-- http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.
--
local type = type
local setmetatable = setmetatable
local ngx = ngx
local ngx_sleep = ngx.sleep
local thread_spawn = ngx.thread.spawn
local thread_kill = ngx.thread.kill
local shared_dict = ngx.shared["mcp-session"] -- TODO: rename to something like mcp-broker
local core = require("apisix.core")
local broker_utils = require("apisix.plugins.mcp.broker.utils")

local _M = {}
local mt = { __index = _M }


local STORAGE_SUFFIX_QUEUE = ":queue"


function _M.new(opts)
return setmetatable({
session_id = opts.session_id,
event_handler = {}
}, mt)
end


function _M.on(self, event, cb)
self.event_handler[event] = cb
end


function _M.push(self, message)
if not message then
return nil, "message is nil"
end
local ok, err = shared_dict:rpush(self.session_id .. STORAGE_SUFFIX_QUEUE, message)
if not ok then
return nil, "failed to push message to queue: " .. err
end
return true
end


function _M.start(self)
self.thread = thread_spawn(function()
while true do
local item, err = shared_dict:lpop(self.session_id .. STORAGE_SUFFIX_QUEUE)
if err then
core.log.info("session ", self.session_id,
" exit, failed to pop message from queue: ", err)
break
end
if item and type(item) == "string"
and type(self.event_handler[broker_utils.EVENT_MESSAGE]) == "function" then
self.event_handler[broker_utils.EVENT_MESSAGE](
core.json.decode(item), { raw = item }
)
end

ngx_sleep(0.1) -- yield to other light threads
end
end)
end


function _M.close(self)
if self.thread then
thread_kill(self.thread)
self.thread = nil
end
end


return _M
Loading
Loading