From 08ff013685fa0e80e4e73d5573e1bd4c16ffcae9 Mon Sep 17 00:00:00 2001 From: Frost Ming Date: Thu, 14 May 2026 17:44:30 +0800 Subject: [PATCH] feat: add bub-mcp-server package with SSE support and run_model tool Signed-off-by: Frost Ming --- .github/workflows/ci.yml | 2 +- packages/bub-mcp-server/README.md | 45 ++++++++ packages/bub-mcp-server/pyproject.toml | 26 +++++ .../src/bub_mcp_server/__init__.py | 8 ++ .../src/bub_mcp_server/config.py | 14 +++ .../src/bub_mcp_server/plugin.py | 102 ++++++++++++++++++ .../src/bub_mcp_server/py.typed | 1 + packages/bub-mcp-server/tests/test_plugin.py | 79 ++++++++++++++ pyproject.toml | 2 + uv.lock | 28 +++++ 10 files changed, 306 insertions(+), 1 deletion(-) create mode 100644 packages/bub-mcp-server/README.md create mode 100644 packages/bub-mcp-server/pyproject.toml create mode 100644 packages/bub-mcp-server/src/bub_mcp_server/__init__.py create mode 100644 packages/bub-mcp-server/src/bub_mcp_server/config.py create mode 100644 packages/bub-mcp-server/src/bub_mcp_server/plugin.py create mode 100644 packages/bub-mcp-server/src/bub_mcp_server/py.typed create mode 100644 packages/bub-mcp-server/tests/test_plugin.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 0da5829..559b039 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -26,4 +26,4 @@ jobs: python-version: "3.12" - name: Run tests - run: uv run --frozen --group test pytest --import-mode=importlib + run: uv run --frozen --group test pytest diff --git a/packages/bub-mcp-server/README.md b/packages/bub-mcp-server/README.md new file mode 100644 index 0000000..a041d43 --- /dev/null +++ b/packages/bub-mcp-server/README.md @@ -0,0 +1,45 @@ +# bub-mcp-server + +Expose Bub as an SSE MCP server. + +## What It Provides + +- Channel implementation: `MCPServerChannel` (`name = "mcp-server"`) +- FastMCP SSE server lifecycle managed by Bub channel startup and shutdown +- One MCP tool: `run_model` + +## Tool + +`run_model` accepts: + +- `prompt` (required): input text to send through Bub +- `session_id` (optional): Bub session id, default `mcp:default` + +It returns Bub's `model_output` for that turn. + +## Configuration + +Settings are read from environment variables with the `BUB_MCP_SERVER_` prefix. + +- `BUB_MCP_SERVER_HOST`: bind host, default `127.0.0.1` +- `BUB_MCP_SERVER_PORT`: bind port, default `28280` (BUBU0 on 9-keyboard) +- `BUB_MCP_SERVER_PATH`: SSE path, default `/sse` +- `BUB_MCP_SERVER_LOG_LEVEL`: Uvicorn log level, default `info` + +## Installation + +```bash +uv pip install "git+https://github.com/bubbuild/bub-contrib.git#subdirectory=packages/bub-mcp-server" +``` + +In this repository, the package is included in the workspace and root dependencies. + +## Usage + +Start Bub with channels enabled. The MCP SSE endpoint is available at: + +```text +http://127.0.0.1:28280/sse +``` + +Configure your MCP client to use SSE transport against that URL. diff --git a/packages/bub-mcp-server/pyproject.toml b/packages/bub-mcp-server/pyproject.toml new file mode 100644 index 0000000..930069c --- /dev/null +++ b/packages/bub-mcp-server/pyproject.toml @@ -0,0 +1,26 @@ +[project] +name = "bub-mcp-server" +version = "0.1.0" +description = "Expose Bub as an SSE MCP server" +readme = "README.md" +authors = [ + { name = "Frost Ming", email = "me@frostming.com" } +] +requires-python = ">=3.12" +dependencies = [ + "bub", + "fastmcp>=3.2.0,<4", + "pydantic-settings>=2.10.1", +] + +[project.entry-points.bub] +mcp-server = "bub_mcp_server.plugin:MCPServerPlugin" + +[build-system] +requires = ["uv_build>=0.10.4,<0.11.0"] +build-backend = "uv_build" + +[dependency-groups] +dev = [ + "pytest>=9.0.3", +] diff --git a/packages/bub-mcp-server/src/bub_mcp_server/__init__.py b/packages/bub-mcp-server/src/bub_mcp_server/__init__.py new file mode 100644 index 0000000..7730b22 --- /dev/null +++ b/packages/bub-mcp-server/src/bub_mcp_server/__init__.py @@ -0,0 +1,8 @@ +"""Expose Bub as an MCP server.""" + +from __future__ import annotations + +__all__ = ["MCPServerChannel", "MCPServerPlugin", "MCPServerSettings"] + +from bub_mcp_server.config import MCPServerSettings +from bub_mcp_server.plugin import MCPServerChannel, MCPServerPlugin diff --git a/packages/bub-mcp-server/src/bub_mcp_server/config.py b/packages/bub-mcp-server/src/bub_mcp_server/config.py new file mode 100644 index 0000000..cfabe25 --- /dev/null +++ b/packages/bub-mcp-server/src/bub_mcp_server/config.py @@ -0,0 +1,14 @@ +from __future__ import annotations + +import bub +from pydantic_settings import SettingsConfigDict + + +@bub.config(name="mcp-server") +class MCPServerSettings(bub.Settings): + model_config = SettingsConfigDict(env_prefix="BUB_MCP_SERVER_", extra="ignore") + + host: str = "127.0.0.1" + port: int = 28280 + path: str = "/sse" + log_level: str = "info" diff --git a/packages/bub-mcp-server/src/bub_mcp_server/plugin.py b/packages/bub-mcp-server/src/bub_mcp_server/plugin.py new file mode 100644 index 0000000..9c6ba23 --- /dev/null +++ b/packages/bub-mcp-server/src/bub_mcp_server/plugin.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +import asyncio +import contextlib +from typing import TYPE_CHECKING, Any + +import bub +from bub import hookimpl +from bub.channels import Channel +from bub.channels.message import ChannelMessage +from bub.types import MessageHandler +from fastmcp import FastMCP +from loguru import logger + +from bub_mcp_server.config import MCPServerSettings + +if TYPE_CHECKING: + from bub.framework import BubFramework + + +class MCPServerChannel(Channel): + name = "mcp-server" + + def __init__(self, framework: BubFramework) -> None: + self.framework = framework + self.settings = bub.ensure_config(MCPServerSettings) + self._server: FastMCP | None = None + self._task: asyncio.Task[None] | None = None + + @property + def server(self) -> FastMCP | None: + return self._server + + async def start(self, stop_event: asyncio.Event) -> None: + del stop_event + if self._task is not None and not self._task.done(): + return + self._server = self._build_server() + self._task = asyncio.create_task( + self._run_server(self._server), name="bub-mcp-server.sse" + ) + + async def stop(self) -> None: + task = self._task + self._task = None + if task is None: + return + if not task.done(): + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + def _build_server(self) -> FastMCP: + server = FastMCP(name="bub") + + @server.tool( + name="run_model", + description="Run one prompt through Bub and return the model output.", + ) + async def run_model(prompt: str, session_id: str = "mcp:default") -> str: + return await self.run_model(prompt=prompt, session_id=session_id) + + return server + + async def _run_server(self, server: FastMCP) -> None: + logger.info( + "starting bub MCP SSE server on http://{}:{}{}", + self.settings.host, + self.settings.port, + self.settings.path, + ) + await server.run_async( + transport="sse", + host=self.settings.host, + port=self.settings.port, + path=self.settings.path, + log_level=self.settings.log_level, + show_banner=False, + ) + + async def run_model(self, *, prompt: str, session_id: str) -> str: + normalized_session_id = session_id.strip() or "mcp:default" + inbound = ChannelMessage( + session_id=normalized_session_id, + channel=self.name, + chat_id=normalized_session_id, + content=prompt, + is_active=True, + kind="normal", + ) + result = await self.framework.process_inbound(inbound) + return result.model_output + + +class MCPServerPlugin: + def __init__(self, framework: Any) -> None: + self._channel = MCPServerChannel(framework) + + @hookimpl + def provide_channels(self, message_handler: MessageHandler) -> list[Channel]: + del message_handler + return [self._channel] diff --git a/packages/bub-mcp-server/src/bub_mcp_server/py.typed b/packages/bub-mcp-server/src/bub_mcp_server/py.typed new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/packages/bub-mcp-server/src/bub_mcp_server/py.typed @@ -0,0 +1 @@ + diff --git a/packages/bub-mcp-server/tests/test_plugin.py b/packages/bub-mcp-server/tests/test_plugin.py new file mode 100644 index 0000000..fac52d2 --- /dev/null +++ b/packages/bub-mcp-server/tests/test_plugin.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +import asyncio +from dataclasses import dataclass + +from bub_mcp_server import plugin + + +@dataclass +class FakeTurnResult: + model_output: str + + +class FakeFramework: + def __init__(self) -> None: + self.messages: list[object] = [] + + async def process_inbound(self, inbound: object) -> FakeTurnResult: + self.messages.append(inbound) + return FakeTurnResult(model_output=f"reply: {inbound.content}") + + +def test_mcp_server_exposes_only_run_model_tool() -> None: + channel = plugin.MCPServerChannel(FakeFramework()) + server = channel._build_server() + + async def run_test() -> None: + tools = await server.list_tools() + assert [tool.name for tool in tools] == ["run_model"] + + asyncio.run(run_test()) + + +def test_run_model_tool_forwards_prompt_to_bub_framework() -> None: + framework = FakeFramework() + channel = plugin.MCPServerChannel(framework) + server = channel._build_server() + + async def run_test() -> None: + result = await server.call_tool( + "run_model", + {"prompt": "hello", "session_id": "mcp:session-1"}, + ) + assert result.content[0].text == "reply: hello" + + asyncio.run(run_test()) + + assert len(framework.messages) == 1 + inbound = framework.messages[0] + assert inbound.session_id == "mcp:session-1" + assert inbound.channel == "mcp-server" + assert inbound.chat_id == "mcp:session-1" + assert inbound.content == "hello" + assert inbound.is_active is True + + +def test_channel_start_runs_sse_server_and_stop_cancels_task(monkeypatch) -> None: + channel = plugin.MCPServerChannel(FakeFramework()) + calls: list[tuple[str, object]] = [] + entered = asyncio.Event() + + async def fake_run_server(server) -> None: + calls.append(("run", server)) + entered.set() + await asyncio.Event().wait() + + monkeypatch.setattr(channel, "_run_server", fake_run_server) + + async def run_test() -> None: + await channel.start(asyncio.Event()) + assert channel.server is not None + await asyncio.wait_for(entered.wait(), timeout=1) + assert channel._task is not None + await channel.stop() + assert channel._task is None + + asyncio.run(run_test()) + + assert calls == [("run", channel.server)] diff --git a/pyproject.toml b/pyproject.toml index c5aa525..4a98f5b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,7 @@ dependencies = [ "bub-github-copilot", "bub-kimi", "bub-mcp", + "bub-mcp-server", "bub-qq", "bub-schedule", "bub-searxng-search", @@ -39,6 +40,7 @@ bub-dingtalk = { workspace = true } bub-github-copilot = { workspace = true } bub-kimi = { workspace = true } bub-mcp = { workspace = true } +bub-mcp-server = { workspace = true } bub-qq = { workspace = true } bub-schedule = { workspace = true } bub-searxng-search = { workspace = true } diff --git a/uv.lock b/uv.lock index da8bf3d..86ee533 100644 --- a/uv.lock +++ b/uv.lock @@ -17,6 +17,7 @@ members = [ "bub-github-copilot", "bub-kimi", "bub-mcp", + "bub-mcp-server", "bub-qq", "bub-schedule", "bub-searxng-search", @@ -392,6 +393,7 @@ dependencies = [ { name = "bub-github-copilot" }, { name = "bub-kimi" }, { name = "bub-mcp" }, + { name = "bub-mcp-server" }, { name = "bub-qq" }, { name = "bub-schedule" }, { name = "bub-searxng-search" }, @@ -423,6 +425,7 @@ requires-dist = [ { name = "bub-github-copilot", editable = "packages/bub-github-copilot" }, { name = "bub-kimi", editable = "packages/bub-kimi" }, { name = "bub-mcp", editable = "packages/bub-mcp" }, + { name = "bub-mcp-server", editable = "packages/bub-mcp-server" }, { name = "bub-qq", editable = "packages/bub-qq" }, { name = "bub-schedule", editable = "packages/bub-schedule" }, { name = "bub-searxng-search", editable = "packages/bub-searxng-search" }, @@ -532,6 +535,31 @@ requires-dist = [ [package.metadata.requires-dev] dev = [{ name = "pytest", specifier = ">=9.0.3" }] +[[package]] +name = "bub-mcp-server" +version = "0.1.0" +source = { editable = "packages/bub-mcp-server" } +dependencies = [ + { name = "bub" }, + { name = "fastmcp" }, + { name = "pydantic-settings" }, +] + +[package.dev-dependencies] +dev = [ + { name = "pytest" }, +] + +[package.metadata] +requires-dist = [ + { name = "bub", git = "https://github.com/bubbuild/bub.git" }, + { name = "fastmcp", specifier = ">=3.2.0,<4" }, + { name = "pydantic-settings", specifier = ">=2.10.1" }, +] + +[package.metadata.requires-dev] +dev = [{ name = "pytest", specifier = ">=9.0.3" }] + [[package]] name = "bub-qq" version = "0.1.0"