Add Handbook.md benchmark tasks
This commit is contained in:
DerekSurge
@@ -0,0 +1,59 @@
|
||||
"""Boot-time guard: every registered MCP tool must be async.
|
||||
|
||||
Sync (`def`) tools make FastMCP run pydantic argument validation in an anyio
|
||||
worker threadpool. Under concurrent calls the shared pydantic-core validator is
|
||||
re-entered across threads and panics with ``pyo3_runtime.PanicException:
|
||||
dictionary changed size during iteration``, which tears down the
|
||||
StreamableHTTP task group and turns every later request into a 500. Async
|
||||
(`async def`) tools validate on the single-threaded event loop and are safe.
|
||||
|
||||
This guard runs at server boot (and therefore during ``mcp-proxy gen``, which
|
||||
boots every server): a non-conformant package fails fast instead of shipping a
|
||||
server that can crash under load. The same check is intentionally duplicated in
|
||||
each package because the bundled prod images install only that package's own
|
||||
dependencies, so there is no shared runtime module to import.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
from functools import partial
|
||||
from typing import Any
|
||||
|
||||
|
||||
def _unwrap(fn: Any) -> Any:
|
||||
while isinstance(fn, partial):
|
||||
fn = fn.func
|
||||
return fn
|
||||
|
||||
|
||||
def find_sync_tools(mcp: Any) -> list[str]:
|
||||
"""Return the names of registered tools whose function is not a coroutine."""
|
||||
# mcp.server.fastmcp.FastMCP exposes a synchronous tool manager.
|
||||
manager = getattr(mcp, "_tool_manager", None)
|
||||
if manager is not None and hasattr(manager, "list_tools"):
|
||||
return sorted(t.name for t in manager.list_tools() if not inspect.iscoroutinefunction(_unwrap(t.fn)))
|
||||
|
||||
# fastmcp.FastMCP exposes an async API.
|
||||
async def _collect() -> list[str]:
|
||||
sync: list[str] = []
|
||||
for descriptor in await mcp.list_tools():
|
||||
tool = await mcp.get_tool(descriptor.name)
|
||||
if not inspect.iscoroutinefunction(_unwrap(tool.fn)):
|
||||
sync.append(descriptor.name)
|
||||
return sorted(sync)
|
||||
|
||||
return asyncio.run(_collect())
|
||||
|
||||
|
||||
def assert_tools_async(mcp: Any) -> None:
|
||||
"""Raise if any registered tool is synchronous."""
|
||||
sync = find_sync_tools(mcp)
|
||||
if sync:
|
||||
raise RuntimeError(
|
||||
"MCP tools must be async (`async def`). Sync tools run pydantic argument "
|
||||
"validation in a worker threadpool and can trigger a pydantic-core "
|
||||
"'dictionary changed size during iteration' panic under concurrent calls, "
|
||||
"which kills the server. Make these tools async: " + ", ".join(sync)
|
||||
)
|
||||
@@ -0,0 +1,47 @@
|
||||
"""MCP proxy CLI entrypoint.
|
||||
|
||||
Most configuration is read from WORLDBENCH_* environment variables.
|
||||
See scripts/start.sh for the translation from legacy CLI args to env
|
||||
vars. Some controls (e.g. --current-time) are passed straight through
|
||||
as CLI arguments rather than via env.
|
||||
|
||||
Each server's ``setup`` hook runs automatically before the server starts,
|
||||
as part of the ``mcp`` command's startup (see ``commands/mcp.py``); there
|
||||
is no separate setup entrypoint.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(prog="mcp-proxy")
|
||||
subparsers = parser.add_subparsers(dest="command")
|
||||
|
||||
mcp_parser = subparsers.add_parser("mcp", help="Start the MCP proxy server")
|
||||
mcp_parser.add_argument("--method", help="Transport method (stdio, sse, http)", default=None)
|
||||
mcp_parser.add_argument("--port", type=int, help="Port for the MCP server")
|
||||
mcp_parser.add_argument(
|
||||
"--current-time",
|
||||
default=None,
|
||||
help="RFC3339 timestamp; runs every MCP service under a faked clock anchored here",
|
||||
)
|
||||
subparsers.add_parser("gen", help="Generate mcp-tools.generated.json for all servers")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.command == "mcp":
|
||||
from mcp_proxy.commands.mcp import run
|
||||
|
||||
run(method=args.method, port=args.port, current_time=args.current_time)
|
||||
elif args.command == "gen":
|
||||
from mcp_proxy.commands.gen import run
|
||||
|
||||
run()
|
||||
else:
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,149 @@
|
||||
"""Generate per-package mcp-tools.generated.json files.
|
||||
|
||||
Discovers MCP servers the same way the ``mcp`` and ``setup`` commands do
|
||||
(WORLDBENCH_ROOT, WORLDBENCH_PACKAGES_ROOT), spawns each server in HTTP mode to
|
||||
introspect its tools, and writes one file per package:
|
||||
|
||||
- packages/<name>/mcp-tools.generated.json (tools + toolset membership)
|
||||
|
||||
This per-package file is the single committed source of truth: the runtime
|
||||
proxy reads it to mount tools and validate requested tool sets, and
|
||||
``scripts/collect_docker_images.py`` reads it to build the per-image toolset
|
||||
artifacts. There is no root aggregate / metadata.jsonc / toolsets/ tree — they
|
||||
were redundant re-groupings of this same data.
|
||||
|
||||
Usage:
|
||||
mcp-proxy gen
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import secrets
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from fastmcp.client import Client
|
||||
from fastmcp.client.transports import StreamableHttpTransport
|
||||
|
||||
from mcp_proxy.commands.mcp import (
|
||||
_find_free_port,
|
||||
_iter_packages,
|
||||
_wait_for_port,
|
||||
)
|
||||
from mcp_proxy.service import McpService, resolve_command
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Per-package tool introspection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _list_tools_http(cfg: McpService, env: dict[str, str]) -> list[dict]:
|
||||
"""Spawn the server in HTTP mode and return its tool list.
|
||||
|
||||
Prefers the dedicated ``gen`` entrypoint (a lightweight tool-listing
|
||||
process that doesn't require the package's full install) and falls back
|
||||
to the last ``run`` step when no ``gen`` is declared.
|
||||
"""
|
||||
port = _find_free_port()
|
||||
proxy_token = secrets.token_hex(16)
|
||||
|
||||
server_step = cfg.config.gen[-1] if cfg.config.gen else cfg.config.run[-1]
|
||||
run_args = [resolve_command(server_step.command), *server_step.args]
|
||||
proc_env = {
|
||||
**env,
|
||||
**server_step.env,
|
||||
"PORT": str(port),
|
||||
"MCP_PROXY_TOKEN": proxy_token,
|
||||
}
|
||||
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
*run_args,
|
||||
cwd=str(cfg.cwd),
|
||||
env=proc_env,
|
||||
stdin=asyncio.subprocess.DEVNULL,
|
||||
stdout=asyncio.subprocess.DEVNULL,
|
||||
stderr=asyncio.subprocess.DEVNULL,
|
||||
)
|
||||
|
||||
try:
|
||||
if proc.returncode is not None:
|
||||
raise RuntimeError(f"Server exited immediately (code {proc.returncode})")
|
||||
if not _wait_for_port(port):
|
||||
raise RuntimeError(f"Server did not start listening on port {port} within 30 s")
|
||||
|
||||
transport = StreamableHttpTransport(
|
||||
f"http://127.0.0.1:{port}/mcp",
|
||||
headers={"X-Proxy-Token": proxy_token},
|
||||
)
|
||||
async with Client(transport) as client:
|
||||
tools = await client.list_tools()
|
||||
return [
|
||||
{
|
||||
"name": t.name,
|
||||
"description": t.description or "",
|
||||
"inputSchema": t.inputSchema,
|
||||
**({"outputSchema": t.outputSchema} if t.outputSchema else {}),
|
||||
}
|
||||
for t in tools
|
||||
]
|
||||
finally:
|
||||
proc.terminate()
|
||||
await proc.wait()
|
||||
|
||||
|
||||
REQUIRED_STATE_TOOLS = ("export_state", "import_state")
|
||||
|
||||
|
||||
def _gen_pkg(cfg: McpService, env: dict[str, str]) -> tuple[list[dict], dict[str, list[str]]]:
|
||||
"""Introspect one package and write its mcp-tools.generated.json.
|
||||
|
||||
Returns ``(tools, toolsets)`` with unnamespaced tool names.
|
||||
"""
|
||||
print(f"[GEN] Spawning {cfg.name} (HTTP mode) to list tools…", file=sys.stderr)
|
||||
tools = asyncio.run(_list_tools_http(cfg, env))
|
||||
tools.sort(key=lambda t: t["name"])
|
||||
|
||||
tool_names = {t["name"] for t in tools}
|
||||
missing = [t for t in REQUIRED_STATE_TOOLS if t not in tool_names]
|
||||
if missing:
|
||||
print(
|
||||
f"[GEN] ERROR: {cfg.name} does not expose required state tool(s): {', '.join(missing)}. "
|
||||
"Every MCP server must implement export_state and import_state.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
toolsets = cfg.toolsets
|
||||
|
||||
output_path = cfg.package_dir / "mcp-tools.generated.json"
|
||||
output_path.write_text(json.dumps({"schema_version": "v1", "tools": tools, "toolsets": toolsets}, indent=2) + "\n")
|
||||
print(f"[GEN] Generated {output_path} ({len(tools)} tools)", file=sys.stderr)
|
||||
|
||||
return tools, toolsets
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Entrypoint
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def run() -> None:
|
||||
base_dir = Path(os.environ.get("WORLDBENCH_ROOT", os.getcwd())).resolve()
|
||||
packages_root = Path(os.environ.get("WORLDBENCH_PACKAGES_ROOT") or (base_dir / "packages"))
|
||||
|
||||
# gen always introspects every package, regardless of WORLDBENCH_TOOL_SETS.
|
||||
configs = _iter_packages(packages_root)
|
||||
if not configs:
|
||||
print(f"[GEN] No MCP servers found under {packages_root}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
env = {
|
||||
**os.environ,
|
||||
"WORLDBENCH_ROOT": str(base_dir),
|
||||
}
|
||||
|
||||
for cfg in configs:
|
||||
_gen_pkg(cfg, env)
|
||||
|
||||
print(f"[GEN] Generated per-package mcp-tools.generated.json for {len(configs)} servers", file=sys.stderr)
|
||||
@@ -0,0 +1,835 @@
|
||||
"""Start the MCP proxy server.
|
||||
|
||||
Discovers MCP servers in a configurable packages directory: any package directory
|
||||
containing an ``mcp.json`` file is treated as an MCP server. This is
|
||||
language/runtime agnostic — Python, Node, or any other executable can be
|
||||
described by the config.
|
||||
|
||||
``mcp.json`` schema
|
||||
-------------------
|
||||
{
|
||||
"run": { // required: how to start the MCP server
|
||||
"command": "node",
|
||||
"args": ["dist/index.js"],
|
||||
"env": {} // optional: extra env vars to inject into the subprocess
|
||||
},
|
||||
"setup": { // optional: run once before the server starts
|
||||
"command": "npm",
|
||||
"args": ["run", "build"],
|
||||
"env": {}
|
||||
},
|
||||
"secrets": ["BRAVE_API_KEY"] // optional: credential-shaped vars to forward from the proxy
|
||||
}
|
||||
|
||||
``run`` and ``setup`` share the same step format: a single step dict
|
||||
``{command, args?, env?}`` or a list of such dicts for multi-step hooks
|
||||
(most useful for ``setup``).
|
||||
|
||||
The top-level ``secrets`` is an opt-in for credential-shaped env vars. A credential-shaped
|
||||
var not in here will not be forwarded to the subprocess.
|
||||
|
||||
Each service is started as an HTTP server on a dynamically assigned port.
|
||||
The proxy communicates with services via StreamableHTTP transport on ``/mcp``
|
||||
and reverse-proxies their viewer UIs through a tabbed interface on port 8000.
|
||||
|
||||
Access to service HTTP servers is restricted via a shared secret token
|
||||
(``MCP_PROXY_TOKEN``) passed as an environment variable and required as an
|
||||
``X-Proxy-Token`` header on all non-MCP requests.
|
||||
"""
|
||||
|
||||
import atexit
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import secrets
|
||||
import shutil
|
||||
import signal
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Iterable
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp.client.transports import StreamableHttpTransport
|
||||
from fastmcp.server.middleware.middleware import CallNext, Middleware, MiddlewareContext
|
||||
from fastmcp.server.providers.proxy import FastMCPProxy, StatefulProxyClient
|
||||
from fastmcp.server.server import ToolTransform
|
||||
|
||||
from mcp_proxy.service import McpConfig, McpService, _bare_toolsets
|
||||
|
||||
logger = logging.getLogger("mcp_proxy")
|
||||
|
||||
# Sentinel tag used to mark proxy-level tools that must be callable but hidden
|
||||
# from tools/list — the proxy aggregate export_state / import_state are
|
||||
# infrastructure for the grading harness, not for agents to discover and use.
|
||||
HIDDEN_FROM_LISTING_TAG = "mcp_proxy:hidden_from_listing"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fake clock (--current-time)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def to_faketime_spec(current_time: str) -> str:
|
||||
"""Convert an RFC3339/ISO-8601 timestamp to the absolute timestamp string
|
||||
the ``faketime`` wrapper expects, normalized to UTC.
|
||||
|
||||
A naive timestamp (no tz offset) is assumed to be UTC. Raises ``ValueError``
|
||||
on an unparseable string. Combined with ``TZ=UTC`` in the service env, the
|
||||
returned ``"%Y-%m-%d %H:%M:%S"`` string anchors the faked clock to the
|
||||
intended instant and lets it advance at real wall-clock rate.
|
||||
"""
|
||||
dt = datetime.fromisoformat(current_time) # `Z` accepted on Python 3.11+
|
||||
if dt.tzinfo is None:
|
||||
dt = dt.replace(tzinfo=UTC)
|
||||
return dt.astimezone(UTC).strftime("%Y-%m-%d %H:%M:%S")
|
||||
|
||||
|
||||
def resolve_faketime_launch(current_time: str) -> tuple[list[str], dict[str, str]]:
|
||||
"""Resolve ``--current-time`` into a command prefix and env overrides.
|
||||
|
||||
Returns ``(["faketime", "<spec>"], {"TZ": "UTC"})``. Exits the process
|
||||
(``SystemExit``) if the timestamp is unparseable or the ``faketime`` wrapper
|
||||
is not installed — a fake clock that silently falls back to real time would
|
||||
break determinism, so we fail loudly instead.
|
||||
"""
|
||||
try:
|
||||
spec = to_faketime_spec(current_time)
|
||||
except ValueError as exc:
|
||||
logger.error("Invalid --current-time %r: %s", current_time, exc)
|
||||
sys.exit(1)
|
||||
|
||||
if shutil.which("faketime") is None:
|
||||
logger.error(
|
||||
"--current-time requires the 'faketime' wrapper, which is not on PATH. "
|
||||
"Install the faketime package (it ships the libfaketime .so)."
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
return ["faketime", spec], {"TZ": "UTC"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Crash / signal handlers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
_child_processes: list[subprocess.Popen] = []
|
||||
|
||||
|
||||
def install_crash_handlers() -> None:
|
||||
"""Install handlers that log diagnostic info when the MCP server crashes."""
|
||||
|
||||
def _excepthook(exc_type, exc_value, exc_tb):
|
||||
tb = "".join(traceback.format_exception(exc_type, exc_value, exc_tb))
|
||||
logger.error("Unhandled exception:\n%s", tb)
|
||||
sys.__excepthook__(exc_type, exc_value, exc_tb)
|
||||
|
||||
def _exit_handler():
|
||||
logger.info("Process exiting (pid=%d)", os.getpid())
|
||||
_kill_child_processes()
|
||||
|
||||
def _signal_handler(signum, _frame):
|
||||
logger.error("Received signal %s (%d)", signal.Signals(signum).name, signum)
|
||||
_kill_child_processes()
|
||||
os._exit(128 + signum)
|
||||
|
||||
sys.excepthook = _excepthook
|
||||
atexit.register(_exit_handler)
|
||||
signal.signal(signal.SIGTERM, _signal_handler)
|
||||
signal.signal(signal.SIGINT, _signal_handler)
|
||||
|
||||
|
||||
def _kill_child_processes() -> None:
|
||||
"""Terminate all child service processes."""
|
||||
for proc in _child_processes:
|
||||
with contextlib.suppress(OSError):
|
||||
proc.terminate()
|
||||
for proc in _child_processes:
|
||||
try:
|
||||
proc.wait(timeout=5)
|
||||
except Exception:
|
||||
with contextlib.suppress(OSError):
|
||||
proc.kill()
|
||||
|
||||
|
||||
def _validate_tool_sets(requested: list[str], packages_root: Path) -> None:
|
||||
"""Validate that every requested tool_set is provided by some package's mcp.json.
|
||||
|
||||
The valid universe is ``{<pkg>_<set>}`` over every ``packages/<pkg>/mcp.json``
|
||||
``toolsets`` declaration — the same source the proxy uses to decide what to
|
||||
mount. Exits with an error if any requested name is unknown. (Each package's
|
||||
``mcp.json`` is the source of truth; there is no separate ``metadata.jsonc``.)
|
||||
"""
|
||||
valid_sets = {f"{svc.name}_{set_name}" for svc in _iter_packages(packages_root) for set_name in svc.toolsets}
|
||||
|
||||
unknown = set(requested) - valid_sets
|
||||
if unknown:
|
||||
logger.error(
|
||||
"Unknown tool set(s): %s. Valid entries: %s",
|
||||
", ".join(sorted(unknown)),
|
||||
", ".join(sorted(valid_sets)),
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def _iter_packages(packages_root: Path) -> list[McpService]:
|
||||
"""Return every ``mcp.json``-bearing package under *packages_root*.
|
||||
|
||||
Used by ``gen`` (which needs to introspect every server) and as the
|
||||
discovery primitive shared with :func:`discover_mcp_servers`. Sorted by
|
||||
directory name for deterministic ordering.
|
||||
"""
|
||||
if not packages_root.exists() or not packages_root.is_dir():
|
||||
logger.error("Packages root not found: %s", packages_root)
|
||||
sys.exit(1)
|
||||
|
||||
result: list[McpService] = []
|
||||
for pkg_dir in sorted(packages_root.iterdir()):
|
||||
if not pkg_dir.is_dir():
|
||||
continue
|
||||
cfg_file = pkg_dir / "mcp.json"
|
||||
if not cfg_file.exists():
|
||||
continue
|
||||
try:
|
||||
data = json.loads(cfg_file.read_text())
|
||||
cfg = McpConfig.from_mcp_json(pkg_dir.resolve(), data)
|
||||
except Exception as e:
|
||||
logger.error("Could not parse %s: %s", cfg_file, e)
|
||||
sys.exit(1)
|
||||
result.append(McpService(cfg, child_processes=_child_processes))
|
||||
return result
|
||||
|
||||
|
||||
def discover_mcp_servers(
|
||||
packages_root: Path,
|
||||
tool_sets: list[str] | None = None,
|
||||
) -> list[McpService]:
|
||||
"""Return packages under *packages_root* matching at least one of *tool_sets*.
|
||||
|
||||
*tool_sets* must be a non-empty list of namespaced toolset names
|
||||
(e.g. ``"google_mail_read"``, ``"slack_state"``). ``None``, ``[]``, or a
|
||||
list containing only empty strings returns no servers — every consumer
|
||||
must declare its toolsets explicitly.
|
||||
|
||||
Use :func:`_iter_packages` instead when you need every package
|
||||
(``gen`` does this).
|
||||
"""
|
||||
if tool_sets is None:
|
||||
return []
|
||||
requested = [ts for ts in tool_sets if ts]
|
||||
if not requested:
|
||||
return []
|
||||
|
||||
_validate_tool_sets(requested, packages_root)
|
||||
|
||||
result: list[McpService] = []
|
||||
for svc in _iter_packages(packages_root):
|
||||
matched = _bare_toolsets(svc.name, set(svc.toolsets.keys()), requested)
|
||||
if matched:
|
||||
result.append(svc)
|
||||
|
||||
_warn_core_and_compat_shim(result)
|
||||
return result
|
||||
|
||||
|
||||
# ``syntara`` is the temporary backward-compat shim that forwards to ``core``
|
||||
# (REMOVE after 2026-06-18). It and ``core`` expose the same underlying tools
|
||||
# under different names; mounting both at once is redundant rather than fatal,
|
||||
# since syntara is implemented on top of core, so we warn and let both run.
|
||||
COMPAT_SHIM_NAME = "syntara"
|
||||
COMPAT_SHIM_TARGET = "core"
|
||||
|
||||
|
||||
def _warn_core_and_compat_shim(servers: list[McpService]) -> None:
|
||||
names = {svc.name for svc in servers}
|
||||
if COMPAT_SHIM_NAME in names and COMPAT_SHIM_TARGET in names:
|
||||
logger.warning(
|
||||
"requested tool sets include both '%s' and the legacy '%s' compatibility shim, "
|
||||
"which forwards to '%s'. Both will be mounted, exposing the same underlying tools "
|
||||
"under different names. Prefer a single surface: drop the legacy '%s_*' tool sets "
|
||||
"(preferred) or the '%s_*' ones in WORLDBENCH_TOOL_SETS.",
|
||||
COMPAT_SHIM_TARGET,
|
||||
COMPAT_SHIM_NAME,
|
||||
COMPAT_SHIM_TARGET,
|
||||
COMPAT_SHIM_NAME,
|
||||
COMPAT_SHIM_TARGET,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Port allocation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _find_free_port() -> int:
|
||||
"""Find a free TCP port on localhost using OS assignment."""
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
def _wait_for_port(port: int, timeout: float = 30.0) -> bool:
|
||||
"""Wait until a service is listening on the given port."""
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.settimeout(1)
|
||||
s.connect(("127.0.0.1", port))
|
||||
return True
|
||||
except (ConnectionRefusedError, OSError):
|
||||
time.sleep(0.2)
|
||||
return False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Subprocess environment
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
# Every env var from the proxy is forwarded *unless* its name contains a
|
||||
# credential keyword.
|
||||
_CREDENTIAL_KEYWORDS: tuple[str, ...] = (
|
||||
"KEY",
|
||||
"TOKEN",
|
||||
"SECRET",
|
||||
"PASSWORD",
|
||||
"PASSWD",
|
||||
"PASSPHRASE",
|
||||
"CREDENTIAL",
|
||||
"AUTH",
|
||||
"BEARER",
|
||||
"JWT",
|
||||
"COOKIE",
|
||||
"OTP",
|
||||
)
|
||||
|
||||
|
||||
def _looks_like_credential(name: str) -> bool:
|
||||
"""True if the env var's name contains any credential keyword (case-insensitive)."""
|
||||
upper = name.upper()
|
||||
return any(kw in upper for kw in _CREDENTIAL_KEYWORDS)
|
||||
|
||||
|
||||
_CRED_URL_RE = re.compile(r"://[^/\s@]+@")
|
||||
|
||||
|
||||
def _value_has_url_userinfo(value: str) -> bool:
|
||||
"""True if *value* contains a URL-like substring with userinfo (``user@`` or ``user:pass@``).
|
||||
|
||||
Catches credentials embedded in values whose names don't trip
|
||||
:func:`_looks_like_credential` — e.g. ``DATABASE_URL=postgres://u:p@host/db``,
|
||||
``PIP_INDEX_URL=https://user:token@pypi.example.com/simple``, or
|
||||
``GIT_REMOTE_URL=https://ghp_xxx@github.com/org/repo`` (token-as-username,
|
||||
no password).
|
||||
|
||||
Uses a regex (not :func:`urllib.parse.urlsplit`) so nested URI schemes
|
||||
are also caught — ``://X@`` is the universal "URL userinfo" shape regardless
|
||||
of outer scheme.
|
||||
"""
|
||||
return bool(_CRED_URL_RE.search(value))
|
||||
|
||||
|
||||
def _build_subprocess_env(
|
||||
base_dir: Path,
|
||||
server_name: str,
|
||||
declared_secrets: Iterable[str] = (),
|
||||
) -> dict[str, str]:
|
||||
"""Return an environment dict to forward to one subprocess.
|
||||
|
||||
Default-allow with a credential-name denylist: every var from the proxy's
|
||||
env is forwarded unless its name looks credential-shaped. Services that need
|
||||
a real credential must declare it by exact name in their mcp.json ``secrets`` field.
|
||||
|
||||
WORLDBENCH_ROOT is set to ``base_dir`` for every subprocess.
|
||||
|
||||
INPUTDIR keeps today's legacy contract: namespaced to
|
||||
``<INPUTDIR>/<server_name>`` when set, else ``<base_dir>/<server_name>``.
|
||||
|
||||
OUTPUTDIR is namespaced per server (``<OUTPUTDIR or base_dir>/<server_name>``)
|
||||
so legacy ``initial.json`` / ``final.json`` snapshots stay isolated.
|
||||
|
||||
BUNDLEDIR is the unified-bundle root (set by the parent process —
|
||||
``scripts/start.sh`` for local dev, the production harness in production). Each
|
||||
service owns ``<BUNDLEDIR>/services/<own_name>/`` for input; the
|
||||
conventional single-file name is ``state.json`` but services with
|
||||
multi-file seeds glob additional JSON next to it. When unset,
|
||||
services fall back to the legacy INPUTDIR path. mcp_proxy doesn't
|
||||
namespace this env var — services append their own name themselves.
|
||||
|
||||
BUNDLE_OUTPUT_DIR is namespaced per server
|
||||
(``<OUTPUTDIR or base_dir>/services/<server_name>``). Services write
|
||||
``state.json`` directly here, so the on-disk layout becomes
|
||||
``services/<name>/state.json`` in the output bundle. Input and output
|
||||
share the same per-service-subdir layout, so an output bundle
|
||||
round-trips cleanly as the next run's input bundle.
|
||||
|
||||
BUNDLE_INPUT_DIR is the per-service bundle path
|
||||
(``<BUNDLEDIR>/services/<server_name>``, without trailing separator)
|
||||
— only set when ``BUNDLEDIR`` is present in the parent env. It resolves
|
||||
both bundle layouts transparently to the consuming service: in the new
|
||||
folder layout it is a real directory containing files like
|
||||
``mysql.tar.zst``; in the legacy flat layout it is a path *prefix*
|
||||
whose siblings are named ``<server_name>.mysql.tar.zst``. A service
|
||||
can probe both ``"${BUNDLE_INPUT_DIR}/<name>"`` and
|
||||
``"${BUNDLE_INPUT_DIR}.<name>"`` to find a bundle file without
|
||||
having to know its own server name.
|
||||
"""
|
||||
declared = set(declared_secrets)
|
||||
env = {
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k in declared or (not _looks_like_credential(k) and not _value_has_url_userinfo(v))
|
||||
}
|
||||
|
||||
output_root = Path(env.get("OUTPUTDIR", base_dir))
|
||||
input_dir = Path(env.get("INPUTDIR", base_dir)) / server_name
|
||||
output_dir = output_root / server_name
|
||||
bundle_output_dir = output_root / "services" / server_name
|
||||
bundle_root = env.get("BUNDLEDIR")
|
||||
bundle_input_dir = Path(bundle_root) / "services" / server_name if bundle_root else None
|
||||
# INPUTDIR may live on a read-only mount or not have a seed dir for every
|
||||
# server (e.g. core has no seed state). Servers should handle a missing
|
||||
# INPUTDIR themselves; don't fail the proxy on mkdir.
|
||||
try:
|
||||
input_dir.mkdir(parents=True, exist_ok=True)
|
||||
except OSError as exc:
|
||||
logger.warning("[%s] INPUTDIR mkdir skipped (%s): %s", server_name, exc, input_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
bundle_output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
env["WORLDBENCH_ROOT"] = str(base_dir)
|
||||
env["INPUTDIR"] = str(input_dir)
|
||||
env["OUTPUTDIR"] = str(output_dir)
|
||||
env["BUNDLE_OUTPUT_DIR"] = str(bundle_output_dir)
|
||||
if bundle_input_dir is not None:
|
||||
env["BUNDLE_INPUT_DIR"] = str(bundle_input_dir)
|
||||
|
||||
logger.info(
|
||||
"[%s] INPUTDIR=%s OUTPUTDIR=%s BUNDLE_OUTPUT_DIR=%s BUNDLE_INPUT_DIR=%s",
|
||||
server_name,
|
||||
input_dir,
|
||||
output_dir,
|
||||
bundle_output_dir,
|
||||
bundle_input_dir if bundle_input_dir is not None else "<unset>",
|
||||
)
|
||||
|
||||
return env
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Proxy setup
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
# Registry of running services: name → port
|
||||
service_registry: dict[str, int] = {}
|
||||
proxy_token: str = ""
|
||||
|
||||
|
||||
class HideTaggedToolsMiddleware(Middleware):
|
||||
"""Strip tools tagged with ``HIDDEN_FROM_LISTING_TAG`` from ``tools/list``.
|
||||
|
||||
Hidden tools remain fully callable via ``tools/call`` — the grading harness
|
||||
relies on the un-namespaced ``export_state`` / ``import_state`` for snapshots —
|
||||
but they are removed from the listing the model sees, so an agent has no
|
||||
way to discover or invoke them on its own.
|
||||
"""
|
||||
|
||||
async def on_list_tools(
|
||||
self,
|
||||
context: MiddlewareContext,
|
||||
call_next: CallNext,
|
||||
):
|
||||
tools = await call_next(context)
|
||||
return [t for t in tools if HIDDEN_FROM_LISTING_TAG not in (t.tags or set())]
|
||||
|
||||
|
||||
def _assert_local_tools_async(app: FastMCP) -> None:
|
||||
"""Run the async-tool guard against only the proxy's *locally* registered tools.
|
||||
|
||||
The aggregate proxy mounts child-server tools via ``add_provider`` — those
|
||||
arrive as ``ProxyTool`` (no ``.fn``) or ``ToolTransform``-wrapped
|
||||
``TransformedTool`` instances, neither of which is a locally-defined
|
||||
``FunctionTool``. Feeding them to ``find_sync_tools`` would either raise
|
||||
(ProxyTool has no ``fn`` field) or false-positive (a transform wrapper's
|
||||
``fn`` is not the remote tool's coroutine). The pydantic-core concurrency
|
||||
panic the guard defends against only affects tools whose argument validation
|
||||
runs locally in the proxy — i.e. its own ``FunctionTool`` definitions
|
||||
(``export_state``, ``import_state``).
|
||||
So we scope the guard to ``FunctionTool`` instances and skip remote/proxied
|
||||
tools, while still failing for any local sync ``def`` tool.
|
||||
"""
|
||||
from fastmcp.tools.function_tool import FunctionTool
|
||||
|
||||
from mcp_proxy.async_tool_guard import assert_tools_async
|
||||
|
||||
async def _local_tools() -> list:
|
||||
local = []
|
||||
for descriptor in await app.list_tools():
|
||||
tool = await app.get_tool(descriptor.name)
|
||||
if isinstance(tool, FunctionTool):
|
||||
local.append(tool)
|
||||
return local
|
||||
|
||||
import asyncio
|
||||
|
||||
tools = asyncio.run(_local_tools())
|
||||
|
||||
class _LocalView:
|
||||
async def list_tools(self):
|
||||
return tools
|
||||
|
||||
async def get_tool(self, name: str):
|
||||
return next(t for t in tools if t.name == name)
|
||||
|
||||
assert_tools_async(_LocalView())
|
||||
|
||||
|
||||
def build_proxy_app(
|
||||
packages_root: Path,
|
||||
base_dir: Path,
|
||||
run_setup_hooks: bool = False,
|
||||
tool_sets: list[str] | None = None,
|
||||
command_prefix: list[str] | None = None,
|
||||
service_env_overrides: dict[str, str] | None = None,
|
||||
) -> FastMCP:
|
||||
"""Discover sub-servers, start them as HTTP processes, and return
|
||||
an aggregating FastMCP proxy app.
|
||||
|
||||
Also populates the global ``service_registry`` with name → port mappings
|
||||
and sets ``proxy_token`` for viewer authentication.
|
||||
|
||||
*command_prefix* and *service_env_overrides*, when set, wrap and augment
|
||||
each service's launch — used by ``--current-time`` to run services under
|
||||
the ``faketime`` wrapper with ``TZ=UTC``.
|
||||
"""
|
||||
global proxy_token
|
||||
|
||||
proxy_token = secrets.token_hex(32)
|
||||
|
||||
configs = discover_mcp_servers(packages_root, tool_sets=tool_sets)
|
||||
if not configs:
|
||||
logger.error("No MCP servers found under %s", packages_root)
|
||||
sys.exit(1)
|
||||
|
||||
app = FastMCP("mcp-proxy")
|
||||
app.add_middleware(HideTaggedToolsMiddleware())
|
||||
pending: list[tuple[McpService, int, subprocess.Popen]] = []
|
||||
startup_failures: list[str] = []
|
||||
|
||||
for cfg in configs:
|
||||
env = _build_subprocess_env(base_dir, server_name=cfg.name, declared_secrets=cfg.secrets)
|
||||
|
||||
if not cfg.run_install(env):
|
||||
logger.error("Install failed for %s, aborting", cfg.name)
|
||||
sys.exit(1)
|
||||
|
||||
if run_setup_hooks and not cfg.run_setup(env):
|
||||
logger.error("Setup failed for %s, aborting", cfg.name)
|
||||
sys.exit(1)
|
||||
|
||||
if not cfg.run_pre_steps(env):
|
||||
logger.error("Pre-run step failed for %s, aborting", cfg.name)
|
||||
sys.exit(1)
|
||||
|
||||
# Apply the fake clock only to the long-running service (and its
|
||||
# children, e.g. the syntara bash/python sandboxes) — not to the
|
||||
# install/setup/pre-run hooks above.
|
||||
service_env = {**env, **service_env_overrides} if service_env_overrides else env
|
||||
|
||||
port = _find_free_port()
|
||||
proc = cfg.start_service(service_env, port, proxy_token, command_prefix=command_prefix)
|
||||
if proc is None:
|
||||
startup_failures.append(f"{cfg.name} (failed to launch process)")
|
||||
continue
|
||||
|
||||
pending.append((cfg, port, proc))
|
||||
|
||||
# Wait for all services to become ready, then mount them
|
||||
for cfg, port, proc in pending:
|
||||
if proc.poll() is not None:
|
||||
logger.error("%s exited immediately with code %d", cfg.name, proc.returncode)
|
||||
startup_failures.append(f"{cfg.name} (exited immediately with code {proc.returncode})")
|
||||
continue
|
||||
|
||||
if not _wait_for_port(port, timeout=cfg.config.startup_timeout):
|
||||
logger.error(
|
||||
"%s did not start listening on port %d within %.0fs", cfg.name, port, cfg.config.startup_timeout
|
||||
)
|
||||
proc.terminate()
|
||||
startup_failures.append(
|
||||
f"{cfg.name} (did not start listening on port {port} within {cfg.config.startup_timeout:.0f}s)"
|
||||
)
|
||||
continue
|
||||
|
||||
logger.info("Mounting %s (HTTP on port %d)", cfg.name, port)
|
||||
|
||||
transport = StreamableHttpTransport(
|
||||
f"http://127.0.0.1:{port}/mcp",
|
||||
headers={"X-Proxy-Token": proxy_token},
|
||||
)
|
||||
client = StatefulProxyClient(transport)
|
||||
proxy_server = FastMCPProxy(
|
||||
client_factory=client.new_stateful,
|
||||
name=cfg.name,
|
||||
)
|
||||
|
||||
transforms = cfg.resolve_tool_transforms(tool_sets or [])
|
||||
|
||||
from fastmcp.server.providers.fastmcp_provider import FastMCPProvider
|
||||
|
||||
provider = FastMCPProvider(proxy_server)
|
||||
provider = provider.wrap_transform(ToolTransform(transforms))
|
||||
app.add_provider(provider, namespace="")
|
||||
|
||||
service_registry[cfg.name] = port
|
||||
|
||||
if startup_failures:
|
||||
logger.error(
|
||||
"Aborting startup because requested MCP services failed to start: %s",
|
||||
", ".join(startup_failures),
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
_register_aggregate_state_tools(app, configs)
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def _aggregate_state_schema(configs: list[McpService]) -> dict:
|
||||
"""Assemble the proxy-level import_state schema from each mounted server.
|
||||
|
||||
Pulls the ``state`` property schema out of each server's ``import_state``
|
||||
tool so the aggregate schema reflects only the currently-mounted services
|
||||
(not every package on disk). Produces an object schema keyed by server
|
||||
name, giving clients a precise view of the aggregate shape rather than a
|
||||
generic ``{}``.
|
||||
"""
|
||||
properties: dict[str, dict] = {}
|
||||
for cfg in configs:
|
||||
for tool in cfg.tools:
|
||||
if tool.get("name") != "import_state":
|
||||
continue
|
||||
state_schema = (tool.get("inputSchema") or {}).get("properties", {}).get("state")
|
||||
if state_schema is not None:
|
||||
properties[cfg.name] = state_schema
|
||||
break
|
||||
return {
|
||||
"type": "object",
|
||||
"description": "Aggregate state keyed by server name.",
|
||||
"properties": properties,
|
||||
}
|
||||
|
||||
|
||||
def _register_aggregate_state_tools(app: FastMCP, configs: list[McpService]) -> None:
|
||||
"""Register the proxy's own ``export_state`` / ``import_state`` tools.
|
||||
|
||||
These aggregate every mounted sub-server's state into a single JSON object
|
||||
keyed by server name, enabling round-trip snapshots across the full proxy.
|
||||
"""
|
||||
from fastmcp.client import Client
|
||||
from mcp.types import ToolAnnotations
|
||||
|
||||
aggregate_schema = _aggregate_state_schema(configs)
|
||||
|
||||
async def _call_downstream(name: str, tool: str, args: dict) -> dict | None:
|
||||
port = service_registry.get(name)
|
||||
if port is None:
|
||||
raise ValueError(f"service '{name}' is not mounted on this proxy")
|
||||
transport = StreamableHttpTransport(
|
||||
f"http://127.0.0.1:{port}/mcp",
|
||||
headers={"X-Proxy-Token": proxy_token},
|
||||
)
|
||||
async with Client(transport) as client:
|
||||
result = await client.call_tool(tool, args)
|
||||
if not result.content:
|
||||
return None
|
||||
text = getattr(result.content[0], "text", None)
|
||||
if text is None:
|
||||
return None
|
||||
try:
|
||||
return json.loads(text)
|
||||
except json.JSONDecodeError:
|
||||
return {"_raw": text}
|
||||
|
||||
@app.tool(
|
||||
name="export_state",
|
||||
description=(
|
||||
"Export the aggregate state of every mounted MCP server as one JSON "
|
||||
"object keyed by server name. Round-trips with import_state."
|
||||
),
|
||||
annotations=ToolAnnotations(readOnlyHint=True, idempotentHint=True),
|
||||
tags={HIDDEN_FROM_LISTING_TAG},
|
||||
)
|
||||
async def export_state() -> dict:
|
||||
result: dict = {}
|
||||
for name in sorted(service_registry):
|
||||
result[name] = await _call_downstream(name, "export_state", {})
|
||||
return result
|
||||
|
||||
@app.tool(
|
||||
name="import_state",
|
||||
description=(
|
||||
"Replace the state of every mounted MCP server from one JSON object "
|
||||
"keyed by server name. Servers omitted from the input are left "
|
||||
"untouched. Round-trips with export_state."
|
||||
),
|
||||
annotations=ToolAnnotations(readOnlyHint=False, idempotentHint=True),
|
||||
tags={HIDDEN_FROM_LISTING_TAG},
|
||||
)
|
||||
async def import_state(state: dict) -> dict:
|
||||
unknown = [n for n in state if n not in service_registry]
|
||||
if unknown:
|
||||
raise ValueError(f"unknown server(s): {', '.join(sorted(unknown))}")
|
||||
for name, sub_state in state.items():
|
||||
await _call_downstream(name, "import_state", {"state": sub_state})
|
||||
return {"ok": True}
|
||||
|
||||
# Attach the assembled schema so clients see the per-server state shape.
|
||||
# FastMCP doesn't expose a public API for overriding an already-registered
|
||||
# tool's parameters, so we reach into the tool manager carefully. If the
|
||||
# internal layout changes, we silently skip and keep the generic schema.
|
||||
try:
|
||||
tool_manager = getattr(app, "_tool_manager", None)
|
||||
registered = getattr(tool_manager, "_tools", {}) if tool_manager else {}
|
||||
import_tool = registered.get("import_state")
|
||||
if import_tool is not None:
|
||||
import_tool.parameters = {
|
||||
"type": "object",
|
||||
"properties": {"state": aggregate_schema},
|
||||
"required": ["state"],
|
||||
}
|
||||
except Exception:
|
||||
logger.debug("Could not override proxy import_state schema", exc_info=True)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Viewer HTTP server
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _start_viewer_server(host: str, port: int) -> None:
|
||||
"""Start the viewer reverse-proxy HTTP server in a background thread."""
|
||||
import uvicorn
|
||||
|
||||
from mcp_proxy.viewer import create_viewer_app
|
||||
|
||||
viewer_app = create_viewer_app(service_registry, proxy_token)
|
||||
|
||||
def _run():
|
||||
uvicorn.run(
|
||||
viewer_app,
|
||||
host=host,
|
||||
port=port,
|
||||
log_level="warning",
|
||||
)
|
||||
|
||||
thread = threading.Thread(target=_run, daemon=True)
|
||||
thread.start()
|
||||
print(
|
||||
f"[MCP_PROXY] Viewer server started on http://{host}:{port}/",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Entrypoint
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _read_build_sha() -> str:
|
||||
"""Return the build SHA stamped into the image, or 'dev' for local runs.
|
||||
|
||||
The Dockerfile writes ``/app/.git-sha`` from the ``GIT_SHA`` build arg.
|
||||
Logged at startup and surfaced via ``/health`` so an out-of-date image
|
||||
can be identified without inspecting the host.
|
||||
"""
|
||||
for candidate in (Path("/app/.git-sha"), Path(".git-sha")):
|
||||
try:
|
||||
sha = candidate.read_text().strip()
|
||||
except OSError:
|
||||
continue
|
||||
if sha:
|
||||
return sha
|
||||
return "dev"
|
||||
|
||||
|
||||
def run(method: str | None = None, port: int | None = None, current_time: str | None = None) -> None:
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="[%(name)s] %(message)s",
|
||||
stream=sys.stderr,
|
||||
)
|
||||
|
||||
install_crash_handlers()
|
||||
|
||||
build_sha = _read_build_sha()
|
||||
logger.info("build %s", build_sha)
|
||||
|
||||
base_dir = Path(os.environ.get("WORLDBENCH_ROOT", os.getcwd())).resolve()
|
||||
packages_root = Path(os.environ.get("WORLDBENCH_PACKAGES_ROOT") or (base_dir / "packages"))
|
||||
tool_sets = os.environ.get("WORLDBENCH_TOOL_SETS", "").split()
|
||||
method = method or os.environ.get("WORLDBENCH_METHOD", "stdio")
|
||||
viewer_port = os.environ.get("VIEWER_PORT")
|
||||
mcp_port = port or int(os.environ.get("PORT", "8000"))
|
||||
|
||||
command_prefix: list[str] | None = None
|
||||
service_env_overrides: dict[str, str] | None = None
|
||||
# `is not None` (not truthiness): an explicitly-passed empty --current-time
|
||||
# is a bad value that must fail loudly, not silently fall back to real time.
|
||||
if current_time is not None:
|
||||
command_prefix, service_env_overrides = resolve_faketime_launch(current_time)
|
||||
logger.info("Faking clock for all services: faketime %s (TZ=UTC)", command_prefix[1])
|
||||
|
||||
app = build_proxy_app(
|
||||
packages_root=packages_root,
|
||||
base_dir=base_dir,
|
||||
run_setup_hooks=True,
|
||||
tool_sets=tool_sets,
|
||||
command_prefix=command_prefix,
|
||||
service_env_overrides=service_env_overrides,
|
||||
)
|
||||
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import JSONResponse
|
||||
|
||||
@app.custom_route("/health", methods=["GET"])
|
||||
async def _health(_request: Request) -> JSONResponse:
|
||||
return JSONResponse(
|
||||
{
|
||||
"status": "ok",
|
||||
"build": build_sha,
|
||||
"services": sorted(service_registry.keys()),
|
||||
}
|
||||
)
|
||||
|
||||
_assert_local_tools_async(app)
|
||||
|
||||
if viewer_port:
|
||||
_start_viewer_server("0.0.0.0", int(viewer_port))
|
||||
|
||||
try:
|
||||
if method in ("http", "sse"):
|
||||
transport = "sse" if method == "sse" else "streamable-http"
|
||||
app.run(
|
||||
transport=transport,
|
||||
port=mcp_port,
|
||||
host="0.0.0.0",
|
||||
path="/mcp",
|
||||
show_banner=False,
|
||||
)
|
||||
else:
|
||||
app.run(show_banner=False)
|
||||
except Exception:
|
||||
logger.exception("app.run() raised an unhandled exception")
|
||||
raise
|
||||
|
||||
_kill_child_processes()
|
||||
@@ -0,0 +1,343 @@
|
||||
"""MCP service configuration and lifecycle management.
|
||||
|
||||
Provides :class:`McpConfig` (a Pydantic model representing ``mcp.json``) and
|
||||
:class:`McpService` (a wrapper that adds install/setup/run lifecycle methods).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
logger = logging.getLogger("mcp_proxy")
|
||||
|
||||
|
||||
def resolve_command(command: str) -> str:
|
||||
"""Resolve a bare ``python``/``python3`` to the proxy's own interpreter so
|
||||
services don't depend on /opt/venv being on PATH (it isn't, by design)."""
|
||||
if command in ("python", "python3"):
|
||||
return sys.executable
|
||||
return command
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# mcp.json data model
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class HookStep(BaseModel):
|
||||
"""A single setup/run step: ``{command, args?, env?}``."""
|
||||
|
||||
command: str
|
||||
args: list[str] = Field(default_factory=list)
|
||||
env: dict[str, str] = Field(default_factory=dict)
|
||||
|
||||
|
||||
def _parse_steps(raw: dict | list | None) -> list[HookStep]:
|
||||
"""Normalise a hook value (run or setup) to a list of HookStep."""
|
||||
if raw is None:
|
||||
return []
|
||||
if isinstance(raw, dict):
|
||||
return [HookStep(**raw)]
|
||||
return [HookStep(**s) for s in raw]
|
||||
|
||||
|
||||
class McpConfig(BaseModel):
|
||||
"""Parsed ``mcp.json`` configuration for one MCP server package."""
|
||||
|
||||
model_config = {"arbitrary_types_allowed": True}
|
||||
|
||||
package_dir: Path
|
||||
name: str
|
||||
run: list[HookStep]
|
||||
install: list[HookStep] = Field(default_factory=list)
|
||||
setup: list[HookStep] = Field(default_factory=list)
|
||||
# Optional lightweight entrypoint used by `mcp-proxy gen` to list tools
|
||||
# without spinning up the full service. Falls back to ``run`` when absent.
|
||||
gen: list[HookStep] = Field(default_factory=list)
|
||||
toolsets: dict[str, list[str]] = Field(default_factory=dict)
|
||||
# Names of secrets to forward from the proxy's env to the
|
||||
# subprocess. The proxy filter drops credential-shaped names by default.
|
||||
secrets: list[str] = Field(default_factory=list)
|
||||
startup_timeout: float = 30.0
|
||||
# When False, tools from this package are exposed without the ``{pkg}__``
|
||||
# prefix. Defaults to True for backwards compatibility.
|
||||
namespaced: bool = True
|
||||
cwd: Path
|
||||
|
||||
@classmethod
|
||||
def from_mcp_json(cls, package_dir: Path, data: dict) -> McpConfig:
|
||||
name = data.get("name") or package_dir.name
|
||||
run = _parse_steps(data.get("run"))
|
||||
if not run:
|
||||
raise ValueError("'run' is required and must define at least one step")
|
||||
install = _parse_steps(data.get("install"))
|
||||
setup = _parse_steps(data.get("setup"))
|
||||
gen = _parse_steps(data.get("gen"))
|
||||
toolsets = data.get("toolsets", {})
|
||||
secrets = list(data.get("secrets", []))
|
||||
startup_timeout = float(data.get("startup_timeout", 30.0))
|
||||
namespaced = bool(data.get("namespaced", True))
|
||||
|
||||
raw_cwd = data.get("cwd")
|
||||
cwd = (package_dir / raw_cwd).resolve() if raw_cwd else package_dir
|
||||
|
||||
return cls(
|
||||
package_dir=package_dir,
|
||||
name=name,
|
||||
run=run,
|
||||
install=install,
|
||||
setup=setup,
|
||||
gen=gen,
|
||||
toolsets=toolsets,
|
||||
secrets=secrets,
|
||||
startup_timeout=startup_timeout,
|
||||
namespaced=namespaced,
|
||||
cwd=cwd,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Toolset helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _service_argv(step: HookStep, command_prefix: list[str] | None) -> list[str]:
|
||||
"""Assemble the argv for a run step, optionally prefixed.
|
||||
|
||||
*command_prefix* (e.g. ``["faketime", "2025-01-15 09:00:00"]``) is prepended
|
||||
to the step's command so the service runs under a wrapper. ``None`` or an
|
||||
empty list returns the bare argv. Factored out so the prefixing is unit
|
||||
testable without launching a process.
|
||||
"""
|
||||
argv = [resolve_command(step.command), *step.args]
|
||||
if command_prefix:
|
||||
return [*command_prefix, *argv]
|
||||
return argv
|
||||
|
||||
|
||||
def _bare_toolsets(pkg_name: str, pkg_toolset_keys: set[str], requested: list[str]) -> set[str]:
|
||||
"""Resolve *requested* toolset names to bare names for *pkg_name*.
|
||||
|
||||
Accepts package-namespaced names (``"google_mail_read"``) and returns
|
||||
the set of matching bare toolset names that exist in *pkg_toolset_keys*.
|
||||
|
||||
Examples::
|
||||
|
||||
_bare_toolsets("google_mail", {"read", "write"}, ["google_mail_read", "slack_write"])
|
||||
# → {"read"} ("slack_write" has no google_mail prefix)
|
||||
|
||||
_bare_toolsets("core", {"ds_all", "read"}, ["core_ds_all", "core_read"])
|
||||
# → {"ds_all", "read"}
|
||||
"""
|
||||
result: set[str] = set()
|
||||
for ts in requested:
|
||||
prefix = f"{pkg_name}_"
|
||||
if ts.startswith(prefix):
|
||||
bare = ts[len(prefix) :]
|
||||
if bare in pkg_toolset_keys:
|
||||
result.add(bare)
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# McpService — lifecycle wrapper around McpConfig
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class McpService:
|
||||
"""Wraps an :class:`McpConfig` with lifecycle methods (install, setup, run)."""
|
||||
|
||||
def __init__(self, config: McpConfig, child_processes: list[subprocess.Popen] | None = None) -> None:
|
||||
self.config = config
|
||||
self._child_processes = child_processes if child_processes is not None else []
|
||||
self._tools_cache: list[dict] | None = None
|
||||
|
||||
# -- Convenience accessors ------------------------------------------------
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return self.config.name
|
||||
|
||||
@property
|
||||
def package_dir(self) -> Path:
|
||||
return self.config.package_dir
|
||||
|
||||
@property
|
||||
def cwd(self) -> Path:
|
||||
return self.config.cwd
|
||||
|
||||
@property
|
||||
def toolsets(self) -> dict[str, list[str]]:
|
||||
return self.config.toolsets
|
||||
|
||||
@property
|
||||
def secrets(self) -> list[str]:
|
||||
"""Credential-shaped env var names this service declared in mcp.json to
|
||||
receive from the proxy's env. Used by ``_build_subprocess_env`` to
|
||||
scope which subprocesses see each secret."""
|
||||
return self.config.secrets
|
||||
|
||||
@property
|
||||
def tools(self) -> list[dict]:
|
||||
"""Tools introspected at gen time, lazily loaded from mcp-tools.generated.json.
|
||||
|
||||
Raises FileNotFoundError if the file is absent — run ``just gen`` to produce it.
|
||||
"""
|
||||
if self._tools_cache is None:
|
||||
tools_file = self.package_dir / "mcp-tools.generated.json"
|
||||
if not tools_file.exists():
|
||||
raise FileNotFoundError(
|
||||
f"mcp-tools.generated.json not found for {self.name} at {tools_file}. "
|
||||
"Run 'just gen' to generate it."
|
||||
)
|
||||
self._tools_cache = json.loads(tools_file.read_text()).get("tools", [])
|
||||
return self._tools_cache
|
||||
|
||||
# -- Hook execution -------------------------------------------------------
|
||||
|
||||
def _run_hook_step(self, step: HookStep, cwd: Path, env: dict[str, str], label: str) -> bool:
|
||||
"""Run a single hook step synchronously. Returns True on success."""
|
||||
step_args = [resolve_command(step.command), *step.args]
|
||||
merged_env = {**env, **step.env}
|
||||
|
||||
logger.info("%s: %s", label, " ".join(step_args))
|
||||
try:
|
||||
result = subprocess.run(
|
||||
step_args, cwd=str(cwd), env=merged_env, check=False, stdin=subprocess.DEVNULL, stdout=sys.stderr
|
||||
)
|
||||
if result.returncode != 0:
|
||||
logger.error("%s exited with code %d", label, result.returncode)
|
||||
return False
|
||||
return True
|
||||
except FileNotFoundError:
|
||||
logger.error("%s: command not found: %s", label, step.command)
|
||||
return False
|
||||
|
||||
def _run_steps(self, steps: list[HookStep], phase: str, cwd: Path, env: dict[str, str]) -> bool:
|
||||
"""Run a list of hook steps in order. Returns False if any step fails."""
|
||||
for i, step in enumerate(steps):
|
||||
label = f"{phase}({self.name})[{i}]" if len(steps) > 1 else f"{phase}({self.name})"
|
||||
if not self._run_hook_step(step, cwd, env, label):
|
||||
return False
|
||||
return True
|
||||
|
||||
def run_install(self, env: dict[str, str]) -> bool:
|
||||
"""Run the package's install steps. No-op (returns True) if none are defined."""
|
||||
if not self.config.install:
|
||||
logger.info("install(%s): no install hook defined, skipping", self.name)
|
||||
return True
|
||||
return self._run_steps(self.config.install, "install", self.package_dir, env)
|
||||
|
||||
def run_setup(self, env: dict[str, str]) -> bool:
|
||||
"""Run the package's setup steps. No-op (returns True) if none are defined."""
|
||||
if not self.config.setup:
|
||||
logger.info("setup(%s): no setup hook defined, skipping", self.name)
|
||||
return True
|
||||
return self._run_steps(self.config.setup, "setup", self.package_dir, env)
|
||||
|
||||
def run_pre_steps(self, env: dict[str, str]) -> bool:
|
||||
"""Run all run steps except the last (pre-run steps). No-op if only one run step."""
|
||||
pre = self.config.run[:-1]
|
||||
if not pre:
|
||||
return True
|
||||
return self._run_steps(pre, "run", self.cwd, env)
|
||||
|
||||
def start_service(
|
||||
self,
|
||||
env: dict[str, str],
|
||||
port: int,
|
||||
proxy_token: str,
|
||||
command_prefix: list[str] | None = None,
|
||||
) -> subprocess.Popen | None:
|
||||
"""Start the service as an HTTP subprocess on the given port.
|
||||
|
||||
Uses the last ``run`` step as the server command. *command_prefix*, when
|
||||
given, wraps that command (e.g. ``["faketime", "<spec>"]``) so the
|
||||
service — and everything it spawns — runs under the wrapper.
|
||||
Returns the Popen handle, or None if the service failed to start.
|
||||
"""
|
||||
step = self.config.run[-1]
|
||||
run_args = _service_argv(step, command_prefix)
|
||||
merged_env = {
|
||||
**env,
|
||||
**step.env,
|
||||
"PORT": str(port),
|
||||
"MCP_PROXY_TOKEN": proxy_token,
|
||||
}
|
||||
|
||||
logger.info("Starting %s on port %d (command: %s)", self.name, port, " ".join(run_args))
|
||||
|
||||
try:
|
||||
proc = subprocess.Popen(
|
||||
run_args,
|
||||
cwd=str(self.cwd),
|
||||
env=merged_env,
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=sys.stderr,
|
||||
)
|
||||
self._child_processes.append(proc)
|
||||
return proc
|
||||
except FileNotFoundError:
|
||||
logger.error("Failed to start %s: command not found: %s", self.name, step.command)
|
||||
return None
|
||||
|
||||
def resolve_tool_transforms(self, tool_sets: list[str]) -> dict:
|
||||
"""Return a ``{original_name: ToolTransformConfig}`` mapping.
|
||||
|
||||
Allowed tools are renamed to ``{pkg}__{name}``. Tools not in the
|
||||
requested *tool_sets* are disabled (hidden from clients).
|
||||
|
||||
An empty *tool_sets* (or a package that declares no toolsets in its
|
||||
``mcp.json``) exposes every introspected tool. Otherwise only tools
|
||||
listed in the named, namespaced toolsets (e.g. ``"slack_read"``) are
|
||||
exposed.
|
||||
|
||||
Raises ``FileNotFoundError`` if ``mcp-tools.generated.json`` is absent.
|
||||
Run ``just gen`` to generate it.
|
||||
"""
|
||||
from fastmcp.server.server import ToolTransformConfig
|
||||
|
||||
all_tool_names = [t["name"] for t in self.tools]
|
||||
|
||||
# Determine the allowed set of tool names. tool_sets entries are
|
||||
# always namespaced ("{pkg}_{toolset}") — bare names are no longer
|
||||
# accepted. An empty list (or no toolsets defined) means "expose
|
||||
# everything", which is how the proxy behaves with no filter.
|
||||
pkg_toolsets = self.toolsets
|
||||
if not tool_sets or not pkg_toolsets:
|
||||
allowed: set[str] | None = None # expose everything
|
||||
else:
|
||||
bare = _bare_toolsets(self.name, set(pkg_toolsets.keys()), tool_sets)
|
||||
if not bare:
|
||||
allowed = set() # no matching toolsets → expose nothing
|
||||
else:
|
||||
allowed = set()
|
||||
for ts in bare:
|
||||
allowed.update(pkg_toolsets.get(ts, []))
|
||||
|
||||
transforms: dict[str, ToolTransformConfig] = {}
|
||||
enabled_count = 0
|
||||
for name in all_tool_names:
|
||||
if allowed is None or name in allowed:
|
||||
exposed = name if not self.config.namespaced else f"{self.name}__{name}"
|
||||
transforms[name] = ToolTransformConfig(name=exposed)
|
||||
enabled_count += 1
|
||||
else:
|
||||
transforms[name] = ToolTransformConfig(enabled=False)
|
||||
|
||||
logger.info(
|
||||
"Namespacing %d/%d tools for %s with prefix '%s__' (tool_sets=%s)",
|
||||
enabled_count,
|
||||
len(all_tool_names),
|
||||
self.name,
|
||||
self.name,
|
||||
tool_sets,
|
||||
)
|
||||
return transforms
|
||||
@@ -0,0 +1,98 @@
|
||||
"""Generic HTTP wrapper for FastMCP services.
|
||||
|
||||
Provides a simple way to run any FastMCP app in HTTP mode with:
|
||||
- MCP StreamableHTTP transport at /mcp
|
||||
- Proxy token authentication
|
||||
- A simple placeholder viewer page at /
|
||||
|
||||
Services with richer data viewers can implement their own viewer
|
||||
module instead of using this generic wrapper.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
import uvicorn
|
||||
from starlette.applications import Starlette
|
||||
from starlette.middleware import Middleware
|
||||
from starlette.middleware.base import BaseHTTPMiddleware
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import HTMLResponse, Response
|
||||
from starlette.routing import Route
|
||||
|
||||
|
||||
class ProxyTokenMiddleware(BaseHTTPMiddleware):
|
||||
"""Reject non-MCP requests lacking the correct X-Proxy-Token header."""
|
||||
|
||||
async def dispatch(self, request: Request, call_next):
|
||||
if request.url.path.startswith("/mcp"):
|
||||
return await call_next(request)
|
||||
token = os.environ.get("MCP_PROXY_TOKEN", "")
|
||||
if token and request.headers.get("x-proxy-token") != token:
|
||||
return Response("Forbidden: invalid proxy token", status_code=403)
|
||||
return await call_next(request)
|
||||
|
||||
|
||||
def run_fastmcp_http(mcp_app, port: int, viewer_html: str | None = None) -> None:
|
||||
"""Run a FastMCP app as an HTTP server with optional viewer.
|
||||
|
||||
Args:
|
||||
mcp_app: A FastMCP instance.
|
||||
port: TCP port to listen on (127.0.0.1 only).
|
||||
viewer_html: Optional HTML string for the viewer page at /.
|
||||
If None, a simple placeholder page is shown.
|
||||
"""
|
||||
fastmcp_asgi = mcp_app.http_app(
|
||||
transport="streamable-http",
|
||||
path="/mcp",
|
||||
)
|
||||
|
||||
html = viewer_html or _placeholder_html(mcp_app.name)
|
||||
|
||||
async def index(request: Request) -> HTMLResponse:
|
||||
return HTMLResponse(html)
|
||||
|
||||
viewer = Starlette(
|
||||
routes=[Route("/", index)],
|
||||
middleware=[Middleware(ProxyTokenMiddleware)],
|
||||
)
|
||||
|
||||
async def combined_app(scope, receive, send):
|
||||
path = scope.get("path", "")
|
||||
if path.startswith("/mcp"):
|
||||
await fastmcp_asgi(scope, receive, send)
|
||||
else:
|
||||
await viewer(scope, receive, send)
|
||||
|
||||
uvicorn.run(
|
||||
combined_app,
|
||||
host="127.0.0.1",
|
||||
port=port,
|
||||
log_level="warning",
|
||||
)
|
||||
|
||||
|
||||
def _placeholder_html(name: str) -> str:
|
||||
return f"""<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<title>{name}</title>
|
||||
<style>
|
||||
body {{ font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
|
||||
display: flex; align-items: center; justify-content: center;
|
||||
height: 100vh; margin: 0; background: #f8fafc; color: #475569; }}
|
||||
.card {{ text-align: center; padding: 40px; }}
|
||||
h1 {{ font-size: 24px; color: #1e293b; margin-bottom: 8px; }}
|
||||
p {{ font-size: 14px; }}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="card">
|
||||
<h1>{name}</h1>
|
||||
<p>MCP service running — viewer not yet implemented.</p>
|
||||
</div>
|
||||
</body>
|
||||
</html>"""
|
||||
@@ -0,0 +1,292 @@
|
||||
"""Viewer reverse-proxy server.
|
||||
|
||||
Serves a tabbed interface that embeds each MCP service's viewer UI in an
|
||||
iframe. All requests to ``/viewer/<service>/...`` are reverse-proxied to
|
||||
the service's internal HTTP server with the ``X-Proxy-Token`` header injected.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
from starlette.applications import Starlette
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import HTMLResponse, RedirectResponse, Response
|
||||
from starlette.routing import Route
|
||||
|
||||
|
||||
def create_viewer_app(
|
||||
service_registry: dict[str, int],
|
||||
proxy_token: str,
|
||||
) -> Starlette:
|
||||
"""Build a Starlette ASGI app that serves the viewer UI and reverse-proxies
|
||||
to service HTTP servers."""
|
||||
|
||||
http_client = httpx.AsyncClient(timeout=180.0)
|
||||
|
||||
_COOKIE = "__viewer_service"
|
||||
|
||||
# Server-side state: last selected service. Used as a fallback when
|
||||
# the browser blocks third-party cookies (e.g. when embedded in an
|
||||
# iframe). Safe because this server has a single user.
|
||||
_state: dict[str, str] = {}
|
||||
|
||||
async def index(request: Request) -> HTMLResponse:
|
||||
return HTMLResponse(_render_shell(service_registry))
|
||||
|
||||
async def set_service(request: Request) -> Response:
|
||||
"""Set the active service cookie and redirect to /."""
|
||||
app = request.query_params.get("app", "")
|
||||
if app not in service_registry:
|
||||
return Response("Unknown service", status_code=404)
|
||||
_state["active_service"] = app
|
||||
resp = RedirectResponse("/", status_code=307)
|
||||
resp.set_cookie(_COOKIE, app, httponly=True, samesite="lax")
|
||||
return resp
|
||||
|
||||
async def proxy(request: Request) -> Response:
|
||||
"""Proxy requests to the active service (read from cookie).
|
||||
|
||||
Falls back to the shell page if no service cookie is set.
|
||||
"""
|
||||
service_name = request.cookies.get(_COOKIE)
|
||||
|
||||
# When the cookie is blocked (iframe context), fall back to
|
||||
# server-side state if the request originated from our own host
|
||||
# (i.e. the inner iframe redirect, not a direct navigation).
|
||||
if (not service_name or service_name not in service_registry) and _state.get(
|
||||
"active_service"
|
||||
) in service_registry:
|
||||
service_name = _state["active_service"]
|
||||
|
||||
if not service_name or service_name not in service_registry:
|
||||
if request.method == "GET":
|
||||
return HTMLResponse(_render_shell(service_registry))
|
||||
return Response("No active service", status_code=404)
|
||||
|
||||
# Direct navigation to / (no referer) with a stale cookie → show shell
|
||||
if request.url.path == "/" and not request.headers.get("referer"):
|
||||
return HTMLResponse(_render_shell(service_registry))
|
||||
|
||||
port = service_registry[service_name]
|
||||
path = request.url.path
|
||||
query = str(request.url.query)
|
||||
url = f"http://127.0.0.1:{port}{path}"
|
||||
if query:
|
||||
url = f"{url}?{query}"
|
||||
|
||||
headers = dict(request.headers)
|
||||
headers["x-proxy-token"] = proxy_token
|
||||
headers.pop("host", None)
|
||||
|
||||
body = await request.body()
|
||||
|
||||
try:
|
||||
resp = await http_client.request(
|
||||
method=request.method,
|
||||
url=url,
|
||||
headers=headers,
|
||||
content=body or None,
|
||||
)
|
||||
except httpx.ConnectError:
|
||||
return Response("Service unavailable", status_code=502)
|
||||
|
||||
excluded = {"transfer-encoding", "connection", "keep-alive"}
|
||||
resp_headers = {k: v for k, v in resp.headers.items() if k.lower() not in excluded}
|
||||
|
||||
return Response(
|
||||
content=resp.content,
|
||||
status_code=resp.status_code,
|
||||
headers=resp_headers,
|
||||
)
|
||||
|
||||
routes = [
|
||||
Route(
|
||||
"/__viewer__/set",
|
||||
set_service,
|
||||
methods=["GET"],
|
||||
),
|
||||
Route(
|
||||
"/__viewer__",
|
||||
index,
|
||||
),
|
||||
# Everything else proxies to the active service (or shows shell if no cookie)
|
||||
Route(
|
||||
"/{path:path}",
|
||||
proxy,
|
||||
methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"],
|
||||
),
|
||||
Route(
|
||||
"/",
|
||||
proxy,
|
||||
methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"],
|
||||
),
|
||||
]
|
||||
|
||||
return Starlette(routes=routes)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shell HTML template
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_DISPLAY_NAMES: dict[str, str] = {
|
||||
"jira": "Jira",
|
||||
"slack": "Slack",
|
||||
"google_mail": "Mail",
|
||||
"emails_toolathlon_mock": "Mail",
|
||||
"google_calendar": "Calendar",
|
||||
"shopify": "Shopify",
|
||||
"core": "Core",
|
||||
# Legacy compat shim (forwards to core); REMOVE after 2026-06-18.
|
||||
"syntara": "Syntara",
|
||||
"web": "Web",
|
||||
}
|
||||
|
||||
_ICONS: dict[str, str] = {
|
||||
"jira": "M12 2L2 12l10 10 10-10L12 2zm0 3.27L19.73 12 12 19.73 4.27 12 12 5.27z",
|
||||
"slack": "M6 15a2 2 0 01-2 2 2 2 0 01-2-2 2 2 0 012-2h2v2zm1 0a2 2 0 012-2 2 2 0 012 2v5a2 2 0 01-2 2 2 2 0 01-2-2v-5zm2-8a2 2 0 01-2-2 2 2 0 012-2 2 2 0 012 2v2H9zm0 1a2 2 0 012 2 2 2 0 01-2 2H4a2 2 0 01-2-2 2 2 0 012-2h5zm8 2a2 2 0 012-2 2 2 0 012 2 2 2 0 01-2 2h-2v-2zm-1 0a2 2 0 01-2 2 2 2 0 01-2-2V5a2 2 0 012-2 2 2 0 012 2v5zm-2 8a2 2 0 012 2 2 2 0 01-2 2 2 2 0 01-2-2v-2h2zm0-1a2 2 0 01-2-2 2 2 0 012-2h5a2 2 0 012 2 2 2 0 01-2 2h-5z",
|
||||
"google_mail": "M20 4H4c-1.1 0-2 .9-2 2v12c0 1.1.9 2 2 2h16c1.1 0 2-.9 2-2V6c0-1.1-.9-2-2-2zm0 4l-8 5-8-5V6l8 5 8-5v2z",
|
||||
"emails_toolathlon_mock": "M20 4H4c-1.1 0-2 .9-2 2v12c0 1.1.9 2 2 2h16c1.1 0 2-.9 2-2V6c0-1.1-.9-2-2-2zm0 4l-8 5-8-5V6l8 5 8-5v2z",
|
||||
"google_calendar": "M19 3h-1V1h-2v2H8V1H6v2H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zm0 16H5V8h14v11z",
|
||||
"shopify": "M15.34 2.61a.49.49 0 00-.45-.36.49.49 0 00-.44.2s-.84 1.06-1.06 1.33a3.58 3.58 0 00-1.63-.84l-.26-1.52a.49.49 0 00-.35-.4.49.49 0 00-.5.13L9.27 2.53a.73.73 0 00-.16.5l.05.53C7.4 4.34 6.3 6.05 6.3 8.12a5.7 5.7 0 005.7 5.7 5.7 5.7 0 005.7-5.7c0-2.24-1.3-4.18-3.36-5.51z",
|
||||
"core": "M12 2C6.48 2 2 6.48 2 12s4.48 10 10 10 10-4.48 10-10S17.52 2 12 2zm-2 15l-5-5 1.41-1.41L10 14.17l7.59-7.59L19 8l-9 9z",
|
||||
# Legacy compat shim (forwards to core); REMOVE after 2026-06-18.
|
||||
"syntara": "M12 2C6.48 2 2 6.48 2 12s4.48 10 10 10 10-4.48 10-10S17.52 2 12 2zm-2 15l-5-5 1.41-1.41L10 14.17l7.59-7.59L19 8l-9 9z",
|
||||
"web": "M12 2C6.48 2 2 6.48 2 12s4.48 10 10 10 10-4.48 10-10S17.52 2 12 2zm-1 17.93c-3.94-.49-7-3.85-7-7.93 0-.62.08-1.21.21-1.79L9 15v1c0 1.1.9 2 2 2v1.93zm6.9-2.54c-.26-.81-1-1.39-1.9-1.39h-1v-3c0-.55-.45-1-1-1H8v-2h2c.55 0 1-.45 1-1V7h2c1.1 0 2-.9 2-2v-.41c2.93 1.19 5 4.06 5 7.41 0 2.08-.8 3.97-2.1 5.39z",
|
||||
}
|
||||
|
||||
|
||||
def _render_shell(registry: dict[str, int]) -> str:
|
||||
"""Render the main tabbed shell HTML."""
|
||||
service_names = sorted(registry.keys())
|
||||
if not service_names:
|
||||
return "<html><body><h1>No services available</h1></body></html>"
|
||||
|
||||
nav_items = []
|
||||
for name in service_names:
|
||||
display = _DISPLAY_NAMES.get(name, name.replace("_", " ").title())
|
||||
icon_path = _ICONS.get(name, "M12 2C6.48 2 2 6.48 2 12s4.48 10 10 10 10-4.48 10-10S17.52 2 12 2z")
|
||||
nav_items.append(
|
||||
f'<button class="nav-item" data-service="{name}" onclick="selectService(\'{name}\')">'
|
||||
f'<svg viewBox="0 0 24 24" width="18" height="18" fill="currentColor"><path d="{icon_path}"/></svg>'
|
||||
f"<span>{display}</span>"
|
||||
f"</button>"
|
||||
)
|
||||
|
||||
nav_html = "\n".join(nav_items)
|
||||
first_service = service_names[0]
|
||||
|
||||
return f"""<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<title>Services</title>
|
||||
<style>
|
||||
* {{ margin: 0; padding: 0; box-sizing: border-box; }}
|
||||
body {{ font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; display: flex; height: 100vh; background: #0f172a; color: #e2e8f0; }}
|
||||
|
||||
/* Sidebar */
|
||||
.sidebar {{
|
||||
width: 220px;
|
||||
min-width: 220px;
|
||||
background: #1e293b;
|
||||
border-right: 1px solid #334155;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
padding: 0;
|
||||
}}
|
||||
.sidebar-header {{
|
||||
padding: 20px 16px 16px;
|
||||
border-bottom: 1px solid #334155;
|
||||
}}
|
||||
.sidebar-header h1 {{
|
||||
font-size: 15px;
|
||||
font-weight: 600;
|
||||
color: #f1f5f9;
|
||||
letter-spacing: -0.01em;
|
||||
}}
|
||||
.sidebar-header p {{
|
||||
font-size: 11px;
|
||||
color: #64748b;
|
||||
margin-top: 4px;
|
||||
}}
|
||||
.nav {{
|
||||
flex: 1;
|
||||
padding: 8px;
|
||||
overflow-y: auto;
|
||||
}}
|
||||
.nav-item {{
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 10px;
|
||||
width: 100%;
|
||||
padding: 10px 12px;
|
||||
border: none;
|
||||
border-radius: 8px;
|
||||
background: transparent;
|
||||
color: #94a3b8;
|
||||
font-size: 13px;
|
||||
font-weight: 500;
|
||||
cursor: pointer;
|
||||
transition: all 0.15s ease;
|
||||
text-align: left;
|
||||
margin-bottom: 2px;
|
||||
}}
|
||||
.nav-item:hover {{
|
||||
background: #334155;
|
||||
color: #e2e8f0;
|
||||
}}
|
||||
.nav-item.active {{
|
||||
background: #3b82f6;
|
||||
color: #fff;
|
||||
}}
|
||||
.nav-item svg {{
|
||||
flex-shrink: 0;
|
||||
opacity: 0.8;
|
||||
}}
|
||||
.nav-item.active svg {{
|
||||
opacity: 1;
|
||||
}}
|
||||
|
||||
/* Main content */
|
||||
.main {{
|
||||
flex: 1;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
min-width: 0;
|
||||
}}
|
||||
.main iframe {{
|
||||
flex: 1;
|
||||
border: none;
|
||||
width: 100%;
|
||||
height: 100%;
|
||||
background: #fff;
|
||||
}}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="sidebar">
|
||||
<div class="sidebar-header">
|
||||
<h1>Services</h1>
|
||||
<p>{len(service_names)} services</p>
|
||||
</div>
|
||||
<div class="nav">
|
||||
{nav_html}
|
||||
</div>
|
||||
</div>
|
||||
<div class="main">
|
||||
<iframe id="viewer-frame" src="/__viewer__/set?app={first_service}" sandbox="allow-scripts allow-same-origin allow-forms"></iframe>
|
||||
</div>
|
||||
<script>
|
||||
function selectService(name) {{
|
||||
document.getElementById('viewer-frame').src = '/__viewer__/set?app=' + name;
|
||||
document.querySelectorAll('.nav-item').forEach(el => {{
|
||||
el.classList.toggle('active', el.dataset.service === name);
|
||||
}});
|
||||
}}
|
||||
// Activate first service on load
|
||||
document.querySelector('.nav-item[data-service="{first_service}"]')?.classList.add('active');
|
||||
</script>
|
||||
</body>
|
||||
</html>"""
|
||||
@@ -0,0 +1,32 @@
|
||||
[build-system]
|
||||
build-backend = "uv_build"
|
||||
requires = [ "uv-build>=0.11,<0.12" ]
|
||||
|
||||
[project]
|
||||
name = "mcp-proxy"
|
||||
version = "0.0.1"
|
||||
description = "MCP proxy server and CLI tools"
|
||||
requires-python = ">=3.13"
|
||||
classifiers = [
|
||||
"Programming Language :: Python :: 3 :: Only",
|
||||
"Programming Language :: Python :: 3.13",
|
||||
"Programming Language :: Python :: 3.14",
|
||||
]
|
||||
dependencies = [
|
||||
"fastmcp>=2",
|
||||
"httpx>=0.27",
|
||||
"pydantic>=2",
|
||||
]
|
||||
optional-dependencies.dev = [ "pytest>=8" ]
|
||||
scripts.mcp-proxy = "mcp_proxy.cli:main"
|
||||
|
||||
[tool.uv]
|
||||
environments = [
|
||||
"sys_platform == 'darwin' and platform_machine == 'arm64'",
|
||||
"sys_platform == 'linux' and platform_machine == 'x86_64'",
|
||||
]
|
||||
build-backend.module-root = ""
|
||||
|
||||
[tool.pytest]
|
||||
ini_options.pythonpath = [ "." ]
|
||||
ini_options.markers = [ "e2e: end-to-end tests requiring a running proxy" ]
|
||||
@@ -0,0 +1,66 @@
|
||||
from pathlib import Path
|
||||
|
||||
from mcp_proxy.commands.mcp import _build_subprocess_env
|
||||
|
||||
|
||||
def test_inputdir_namespaced_per_server(tmp_path, monkeypatch):
|
||||
"""When INPUTDIR is set in env, _build_subprocess_env appends server_name."""
|
||||
monkeypatch.setenv("INPUTDIR", str(tmp_path))
|
||||
env = _build_subprocess_env(tmp_path, server_name="google_mail")
|
||||
assert env["INPUTDIR"] == str(tmp_path / "google_mail")
|
||||
assert Path(env["INPUTDIR"]).is_dir()
|
||||
|
||||
|
||||
def test_inputdir_defaults_to_base_dir(tmp_path, monkeypatch):
|
||||
"""When INPUTDIR is not in env, it defaults to base_dir/server_name."""
|
||||
monkeypatch.delenv("INPUTDIR", raising=False)
|
||||
env = _build_subprocess_env(tmp_path, server_name="google_mail")
|
||||
assert env["INPUTDIR"] == str(tmp_path / "google_mail")
|
||||
assert Path(env["INPUTDIR"]).is_dir()
|
||||
|
||||
|
||||
def test_outputdir_from_env(tmp_path, monkeypatch):
|
||||
"""When OUTPUTDIR is set in env, it is namespaced per server."""
|
||||
monkeypatch.setenv("OUTPUTDIR", str(tmp_path))
|
||||
env = _build_subprocess_env(tmp_path, server_name="slack")
|
||||
assert env["OUTPUTDIR"] == str(tmp_path / "slack")
|
||||
assert Path(env["OUTPUTDIR"]).is_dir()
|
||||
|
||||
|
||||
def test_outputdir_defaults_to_base_dir(tmp_path, monkeypatch):
|
||||
"""When OUTPUTDIR is not in env, it defaults to base_dir/server_name."""
|
||||
monkeypatch.delenv("OUTPUTDIR", raising=False)
|
||||
env = _build_subprocess_env(tmp_path, server_name="google_mail")
|
||||
assert env["OUTPUTDIR"] == str(tmp_path / "google_mail")
|
||||
assert Path(env["OUTPUTDIR"]).is_dir()
|
||||
|
||||
|
||||
def test_all_three_dirs_always_set(tmp_path, monkeypatch):
|
||||
"""INPUTDIR, OUTPUTDIR, and BUNDLE_OUTPUT_DIR are always present in the returned env."""
|
||||
monkeypatch.delenv("INPUTDIR", raising=False)
|
||||
monkeypatch.delenv("OUTPUTDIR", raising=False)
|
||||
env = _build_subprocess_env(tmp_path, server_name="jira")
|
||||
assert "INPUTDIR" in env
|
||||
assert "OUTPUTDIR" in env
|
||||
assert "BUNDLE_OUTPUT_DIR" in env
|
||||
|
||||
|
||||
def test_bundle_output_dir_namespaced_per_service(tmp_path, monkeypatch):
|
||||
"""BUNDLE_OUTPUT_DIR is per-service so each service writes its own
|
||||
services/<name>/state.json, matching the nested input bundle layout."""
|
||||
monkeypatch.setenv("OUTPUTDIR", str(tmp_path))
|
||||
|
||||
env_a = _build_subprocess_env(tmp_path, server_name="google_mail")
|
||||
env_b = _build_subprocess_env(tmp_path, server_name="shopify")
|
||||
|
||||
assert env_a["BUNDLE_OUTPUT_DIR"] == str(tmp_path / "services" / "google_mail")
|
||||
assert env_b["BUNDLE_OUTPUT_DIR"] == str(tmp_path / "services" / "shopify")
|
||||
|
||||
|
||||
def test_bundle_output_dir_defaults_to_base_dir(tmp_path, monkeypatch):
|
||||
"""Without OUTPUTDIR in env, BUNDLE_OUTPUT_DIR is rooted at
|
||||
base_dir/services/<name>."""
|
||||
monkeypatch.delenv("OUTPUTDIR", raising=False)
|
||||
env = _build_subprocess_env(tmp_path, server_name="emails_mock")
|
||||
assert env["BUNDLE_OUTPUT_DIR"] == str(tmp_path / "services" / "emails_mock")
|
||||
assert Path(env["BUNDLE_OUTPUT_DIR"]).is_dir()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,107 @@
|
||||
"""Regression tests for proxy-level env compartmentalization.
|
||||
|
||||
This is the structural defense against agent-side secret exfiltration via
|
||||
mcp tools, like core's bash.
|
||||
"""
|
||||
|
||||
from mcp_proxy.commands.mcp import _build_subprocess_env
|
||||
|
||||
|
||||
def test_declared_credential_is_forwarded(tmp_path, monkeypatch):
|
||||
"""A service that declares a credential by exact name in its mcp.json
|
||||
``env`` list does receive it — opt-in bypasses the denylist."""
|
||||
monkeypatch.setenv("BRAVE_API_KEY", "secret-for-web")
|
||||
env = _build_subprocess_env(tmp_path, server_name="web", declared_secrets=["BRAVE_API_KEY"])
|
||||
assert env["BRAVE_API_KEY"] == "secret-for-web"
|
||||
|
||||
|
||||
def test_secret_isolation_between_services(tmp_path, monkeypatch):
|
||||
"""Same proxy env, two services with different declarations — only the
|
||||
declaring service sees the secret."""
|
||||
monkeypatch.setenv("BRAVE_API_KEY", "secret-do-not-leak")
|
||||
|
||||
web_env = _build_subprocess_env(tmp_path, server_name="web", declared_secrets=["BRAVE_API_KEY"])
|
||||
core_env = _build_subprocess_env(tmp_path, server_name="core", declared_secrets=())
|
||||
|
||||
assert "BRAVE_API_KEY" in web_env
|
||||
assert "BRAVE_API_KEY" not in core_env
|
||||
|
||||
|
||||
def test_credential_shaped_vars_are_dropped(tmp_path, monkeypatch):
|
||||
pollute = {
|
||||
# KEY
|
||||
"BRAVE_API_KEY": "x",
|
||||
"MY_PRIVATE_KEY_PATH": "/etc/key",
|
||||
"WORLDBENCH_FOO_API_KEY": "x",
|
||||
"GITHUB_TOKEN": "ghp_x",
|
||||
"AWS_SECRET_ACCESS_KEY": "aws-x",
|
||||
"stripe_secret": "x", # case-insensitive
|
||||
"DB_PASSWORD": "hunter2",
|
||||
"MYSQL_PASSWD": "old-school-x",
|
||||
"SSH_PASSPHRASE": "x",
|
||||
"GPG_PASSPHRASE": "x",
|
||||
"GOOGLE_CREDENTIALS": "json-x",
|
||||
"OAUTH_CLIENT_ID": "x",
|
||||
"BEARER_TOKEN": "x",
|
||||
"JWT_SECRET": "x",
|
||||
"JWT_SIGNING_KEY": "x",
|
||||
"SESSION_COOKIE": "x",
|
||||
"AUTH_COOKIE": "x",
|
||||
"TOTP_SECRET": "x",
|
||||
"HOTP_KEY": "x",
|
||||
"DATABASE_URL": "postgres://app_user:hunter2@db.internal:5432/app",
|
||||
"REDIS_URL": "redis://:hunter2@cache.internal:6379/0",
|
||||
"PIP_INDEX_URL": "https://user:token@pypi.example.com/simple",
|
||||
"UV_INDEX_URL": "https://user:token@uv.example.com/simple",
|
||||
"JDBC_DATABASE_URL": "jdbc:postgresql://app:hunter2@db.internal/app",
|
||||
"JDBC_MYSQL_URL": "jdbc:mysql://u:p@db:3306/x",
|
||||
"GIT_REMOTE_URL": "https://ghp_xxxxxxxxxxxx@github.com/org/repo",
|
||||
"REGISTRY_URL": "https://pypi-token-here@pypi.example.com/simple",
|
||||
}
|
||||
for k, v in pollute.items():
|
||||
monkeypatch.setenv(k, v)
|
||||
|
||||
env = _build_subprocess_env(tmp_path, server_name="core", declared_secrets=())
|
||||
|
||||
for k in pollute:
|
||||
assert k not in env, f"{k} leaked despite credential-shaped name"
|
||||
|
||||
|
||||
def test_innocuous_vars_pass_through(tmp_path, monkeypatch):
|
||||
"""Default-allow: anything without a credential keyword is forwarded.
|
||||
This is what makes the simplified model work — we don't have to enumerate
|
||||
every system var a subprocess might legitimately need."""
|
||||
monkeypatch.setenv("PATH", "/usr/bin:/bin")
|
||||
monkeypatch.setenv("HOME", "/home/model")
|
||||
monkeypatch.setenv("LANG", "en_US.UTF-8")
|
||||
monkeypatch.setenv("WORLDBENCH_TASK_ID", "task-123")
|
||||
monkeypatch.setenv("WORLDBENCH_TOOL_SETS", "core_read")
|
||||
monkeypatch.setenv("RANDOM_CUSTOM_VAR", "value")
|
||||
monkeypatch.setenv("BRAVE_SEARCH_URL", "https://api.search.brave.com/res/v1/web/search")
|
||||
monkeypatch.setenv("PUBLIC_API_BASE", "https://api.example.com/v1")
|
||||
monkeypatch.setenv("DATABASE_URL", "postgres://localhost/dev") # no userinfo
|
||||
|
||||
env = _build_subprocess_env(tmp_path, server_name="core", declared_secrets=())
|
||||
|
||||
assert env["PATH"] == "/usr/bin:/bin"
|
||||
assert env["HOME"] == "/home/model"
|
||||
assert env["LANG"] == "en_US.UTF-8"
|
||||
assert env["WORLDBENCH_TASK_ID"] == "task-123"
|
||||
assert env["WORLDBENCH_TOOL_SETS"] == "core_read"
|
||||
assert env["RANDOM_CUSTOM_VAR"] == "value"
|
||||
assert env["BRAVE_SEARCH_URL"] == "https://api.search.brave.com/res/v1/web/search"
|
||||
assert env["PUBLIC_API_BASE"] == "https://api.example.com/v1"
|
||||
assert env["DATABASE_URL"] == "postgres://localhost/dev"
|
||||
|
||||
|
||||
def test_url_with_userinfo_can_be_explicitly_declared(tmp_path, monkeypatch):
|
||||
"""A service that legitimately needs a credentialed URL can opt in by
|
||||
declaring the var name in mcp.json's ``secrets`` — declared names bypass
|
||||
both the name-shape and value-shape checks."""
|
||||
monkeypatch.setenv("DATABASE_URL", "postgres://app_user:hunter2@db.internal:5432/app")
|
||||
|
||||
db_user_env = _build_subprocess_env(tmp_path, server_name="my-db-service", declared_secrets=["DATABASE_URL"])
|
||||
core_env = _build_subprocess_env(tmp_path, server_name="core", declared_secrets=())
|
||||
|
||||
assert "DATABASE_URL" in db_user_env
|
||||
assert "DATABASE_URL" not in core_env
|
||||
@@ -0,0 +1,158 @@
|
||||
"""Tests for the ``--current-time`` fake-clock wiring.
|
||||
|
||||
The proxy runs every MCP service under the ``faketime`` wrapper when a
|
||||
``--current-time`` value is supplied, so the service (and anything it spawns,
|
||||
notably syntara's executeBash / executePython sandboxes) observes a clock
|
||||
anchored to that instant and advancing at real wall-clock rate.
|
||||
|
||||
Unit tests here run anywhere — they exercise the pure mapping and the argv
|
||||
assembly without launching a process. The integration test is skipped unless
|
||||
the ``faketime`` wrapper is actually installed on PATH.
|
||||
"""
|
||||
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from datetime import UTC, datetime
|
||||
|
||||
import pytest
|
||||
|
||||
from mcp_proxy.commands import mcp as mcp_cmd
|
||||
from mcp_proxy.commands.mcp import resolve_faketime_launch, to_faketime_spec
|
||||
from mcp_proxy.service import HookStep, _service_argv
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# to_faketime_spec — RFC3339/ISO-8601 -> faketime wrapper's absolute timestamp
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_spec_utc_zulu():
|
||||
assert to_faketime_spec("2025-01-15T09:00:00Z") == "2025-01-15 09:00:00"
|
||||
|
||||
|
||||
def test_spec_naive_assumed_utc():
|
||||
assert to_faketime_spec("2025-01-15T09:00:00") == "2025-01-15 09:00:00"
|
||||
|
||||
|
||||
def test_spec_offset_converted_to_utc():
|
||||
# 09:00 at +05:00 is 04:00 UTC.
|
||||
assert to_faketime_spec("2025-01-15T09:00:00+05:00") == "2025-01-15 04:00:00"
|
||||
|
||||
|
||||
def test_spec_negative_offset_crosses_midnight():
|
||||
# 22:00 at -05:00 is 03:00 UTC the next day.
|
||||
assert to_faketime_spec("2025-01-15T22:00:00-05:00") == "2025-01-16 03:00:00"
|
||||
|
||||
|
||||
def test_spec_invalid_raises():
|
||||
with pytest.raises(ValueError):
|
||||
to_faketime_spec("not-a-timestamp")
|
||||
|
||||
|
||||
def test_spec_empty_raises():
|
||||
# An explicitly-passed empty --current-time is a bad value, not "no clock":
|
||||
# run() routes it here (via `is not None`) so it fails fast rather than
|
||||
# silently falling back to the real clock.
|
||||
with pytest.raises(ValueError):
|
||||
to_faketime_spec("")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _service_argv — prepend the faketime wrapper when a prefix is given
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
# ``python``/``python3`` resolve to the proxy's own interpreter (see
|
||||
# service.resolve_command), so the argv carries ``sys.executable``, not the bare
|
||||
# name. The faketime prefix, when present, wraps that resolved command.
|
||||
def test_argv_no_prefix():
|
||||
step = HookStep(command="python", args=["-m", "server", "--http"])
|
||||
assert _service_argv(step, None) == [sys.executable, "-m", "server", "--http"]
|
||||
|
||||
|
||||
def test_argv_empty_prefix_is_noop():
|
||||
step = HookStep(command="python", args=["-m", "server"])
|
||||
assert _service_argv(step, []) == [sys.executable, "-m", "server"]
|
||||
|
||||
|
||||
def test_argv_with_faketime_prefix():
|
||||
step = HookStep(command="python", args=["-m", "server"])
|
||||
prefix = ["faketime", "2025-01-15 09:00:00"]
|
||||
assert _service_argv(step, prefix) == [
|
||||
"faketime",
|
||||
"2025-01-15 09:00:00",
|
||||
sys.executable,
|
||||
"-m",
|
||||
"server",
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# resolve_faketime_launch — command prefix + env, with fail-fast
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_resolve_returns_prefix_and_utc_env(monkeypatch):
|
||||
monkeypatch.setattr(mcp_cmd.shutil, "which", lambda _name: "/usr/bin/faketime")
|
||||
prefix, env = resolve_faketime_launch("2025-01-15T09:00:00Z")
|
||||
assert prefix == ["faketime", "2025-01-15 09:00:00"]
|
||||
assert env == {"TZ": "UTC"}
|
||||
|
||||
|
||||
def test_resolve_fails_fast_when_wrapper_missing(monkeypatch):
|
||||
monkeypatch.setattr(mcp_cmd.shutil, "which", lambda _name: None)
|
||||
with pytest.raises(SystemExit):
|
||||
resolve_faketime_launch("2025-01-15T09:00:00Z")
|
||||
|
||||
|
||||
def test_resolve_invalid_timestamp_fails(monkeypatch):
|
||||
monkeypatch.setattr(mcp_cmd.shutil, "which", lambda _name: "/usr/bin/faketime")
|
||||
with pytest.raises(SystemExit):
|
||||
resolve_faketime_launch("not-a-timestamp")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Integration — only when the faketime wrapper is installed
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_HAS_FAKETIME = shutil.which("faketime") is not None
|
||||
_SKIP_REASON = "faketime wrapper not installed on PATH"
|
||||
|
||||
# 2025-01-15T09:00:00Z
|
||||
_FAKE_EPOCH = 1736931600
|
||||
_SPEC = "2025-01-15 09:00:00"
|
||||
|
||||
|
||||
@pytest.mark.skipif(not _HAS_FAKETIME, reason=_SKIP_REASON)
|
||||
def test_faketime_fakes_date():
|
||||
out = subprocess.run(
|
||||
["faketime", _SPEC, "date", "+%s"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env={"TZ": "UTC", "PATH": __import__("os").environ["PATH"]},
|
||||
timeout=30,
|
||||
)
|
||||
assert out.returncode == 0, out.stderr
|
||||
reported = int(out.stdout.strip())
|
||||
# Clock advances from the anchor, so allow a small forward window.
|
||||
assert _FAKE_EPOCH <= reported < _FAKE_EPOCH + 60
|
||||
|
||||
|
||||
@pytest.mark.skipif(not _HAS_FAKETIME, reason=_SKIP_REASON)
|
||||
def test_faketime_fakes_python():
|
||||
out = subprocess.run(
|
||||
["faketime", _SPEC, sys.executable, "-c", "import time; print(time.time())"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env={"TZ": "UTC", "PATH": __import__("os").environ["PATH"]},
|
||||
timeout=30,
|
||||
)
|
||||
assert out.returncode == 0, out.stderr
|
||||
reported = float(out.stdout.strip())
|
||||
assert _FAKE_EPOCH <= reported < _FAKE_EPOCH + 60
|
||||
|
||||
|
||||
def test_fake_epoch_matches_spec():
|
||||
"""Guard the integration constants against drift."""
|
||||
dt = datetime(2025, 1, 15, 9, 0, 0, tzinfo=UTC)
|
||||
assert int(dt.timestamp()) == _FAKE_EPOCH
|
||||
@@ -0,0 +1,600 @@
|
||||
"""Tests for utility functions in commands/mcp.py."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import ClassVar
|
||||
|
||||
import pytest
|
||||
from fastmcp import FastMCP
|
||||
|
||||
from mcp_proxy.commands import mcp as mcp_command
|
||||
from mcp_proxy.commands.mcp import (
|
||||
HIDDEN_FROM_LISTING_TAG,
|
||||
HideTaggedToolsMiddleware,
|
||||
_find_free_port,
|
||||
_iter_packages,
|
||||
_validate_tool_sets,
|
||||
_wait_for_port,
|
||||
build_proxy_app,
|
||||
discover_mcp_servers,
|
||||
)
|
||||
from mcp_proxy.service import McpConfig, McpService, resolve_command
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _find_free_port
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFindFreePort:
|
||||
def test_returns_int(self):
|
||||
port = _find_free_port()
|
||||
assert isinstance(port, int)
|
||||
|
||||
def test_port_is_in_valid_range(self):
|
||||
port = _find_free_port()
|
||||
assert 1024 <= port <= 65535
|
||||
|
||||
def test_returns_different_ports(self):
|
||||
ports = {_find_free_port() for _ in range(5)}
|
||||
# At least 2 distinct ports (extremely likely)
|
||||
assert len(ports) >= 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _wait_for_port
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestWaitForPort:
|
||||
def test_returns_true_when_port_is_listening(self):
|
||||
# Start a simple server
|
||||
server = HTTPServer(("127.0.0.1", 0), BaseHTTPRequestHandler)
|
||||
port = server.server_address[1]
|
||||
thread = threading.Thread(target=server.handle_request, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
assert _wait_for_port(port, timeout=5.0) is True
|
||||
finally:
|
||||
server.server_close()
|
||||
|
||||
def test_returns_false_when_port_not_listening(self):
|
||||
# Find a free port and don't listen on it
|
||||
port = _find_free_port()
|
||||
assert _wait_for_port(port, timeout=0.5) is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# McpConfig & discover_mcp_servers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestMcpConfig:
|
||||
def test_basic_config(self):
|
||||
cfg = McpConfig.from_mcp_json(
|
||||
Path("/tmp/my-service"),
|
||||
{"run": {"command": "node", "args": ["index.js"]}},
|
||||
)
|
||||
assert cfg.name == "my-service"
|
||||
assert len(cfg.run) == 1
|
||||
assert cfg.run[0].command == "node"
|
||||
assert cfg.run[0].args == ["index.js"]
|
||||
assert cfg.setup == []
|
||||
|
||||
def test_run_list(self):
|
||||
cfg = McpConfig.from_mcp_json(
|
||||
Path("/tmp/svc"),
|
||||
{"run": [{"command": "build"}, {"command": "node", "args": ["index.js"]}]},
|
||||
)
|
||||
assert len(cfg.run) == 2
|
||||
assert cfg.run[0].command == "build"
|
||||
assert cfg.run[1].command == "node"
|
||||
|
||||
def test_missing_run_raises(self):
|
||||
with pytest.raises(ValueError, match="'run' is required"):
|
||||
McpConfig.from_mcp_json(Path("/tmp/svc"), {})
|
||||
|
||||
def test_custom_name(self):
|
||||
cfg = McpConfig.from_mcp_json(
|
||||
Path("/tmp/my-service"),
|
||||
{"name": "Custom", "run": {"command": "python"}},
|
||||
)
|
||||
assert cfg.name == "Custom"
|
||||
|
||||
def test_setup_single_dict(self):
|
||||
cfg = McpConfig.from_mcp_json(
|
||||
Path("/tmp/svc"),
|
||||
{
|
||||
"run": {"command": "node"},
|
||||
"setup": {"command": "npm", "args": ["run", "build"]},
|
||||
},
|
||||
)
|
||||
assert len(cfg.setup) == 1
|
||||
assert cfg.setup[0].command == "npm"
|
||||
|
||||
def test_setup_list(self):
|
||||
cfg = McpConfig.from_mcp_json(
|
||||
Path("/tmp/svc"),
|
||||
{
|
||||
"run": {"command": "node"},
|
||||
"setup": [{"command": "a"}, {"command": "b"}],
|
||||
},
|
||||
)
|
||||
assert len(cfg.setup) == 2
|
||||
|
||||
def test_cwd_override(self, tmp_path):
|
||||
cfg = McpConfig.from_mcp_json(
|
||||
tmp_path / "svc",
|
||||
{"run": {"command": "node"}, "cwd": "subdir"},
|
||||
)
|
||||
assert cfg.cwd == (tmp_path / "svc" / "subdir").resolve()
|
||||
|
||||
def test_install_single_dict(self):
|
||||
cfg = McpConfig.from_mcp_json(
|
||||
Path("/tmp/svc"),
|
||||
{
|
||||
"run": {"command": "node"},
|
||||
"install": {"command": "uv", "args": ["sync", "--package", "foo"]},
|
||||
},
|
||||
)
|
||||
assert len(cfg.install) == 1
|
||||
assert cfg.install[0].command == "uv"
|
||||
assert cfg.install[0].args == ["sync", "--package", "foo"]
|
||||
|
||||
def test_install_defaults_empty(self):
|
||||
cfg = McpConfig.from_mcp_json(
|
||||
Path("/tmp/svc"),
|
||||
{"run": {"command": "node"}},
|
||||
)
|
||||
assert cfg.install == []
|
||||
|
||||
def test_toolsets_parsed(self):
|
||||
cfg = McpConfig.from_mcp_json(
|
||||
Path("/tmp/svc"),
|
||||
{
|
||||
"run": {"command": "node"},
|
||||
"toolsets": {"read": ["t1", "t2"], "write": ["t3"]},
|
||||
},
|
||||
)
|
||||
assert cfg.toolsets == {"read": ["t1", "t2"], "write": ["t3"]}
|
||||
|
||||
def test_toolsets_defaults_empty(self):
|
||||
cfg = McpConfig.from_mcp_json(
|
||||
Path("/tmp/svc"),
|
||||
{"run": {"command": "node"}},
|
||||
)
|
||||
assert cfg.toolsets == {}
|
||||
|
||||
|
||||
class TestIterPackages:
|
||||
"""Iteration is the discovery primitive — it returns every package
|
||||
regardless of any tool_sets filter. ``gen`` uses this directly."""
|
||||
|
||||
def test_iter_packages_returns_each(self, tmp_path):
|
||||
pkg = tmp_path / "svc1"
|
||||
pkg.mkdir()
|
||||
(pkg / "mcp.json").write_text(json.dumps({"run": {"command": "node"}}))
|
||||
|
||||
servers = _iter_packages(tmp_path)
|
||||
assert len(servers) == 1
|
||||
assert servers[0].name == "svc1"
|
||||
|
||||
def test_exits_on_missing_run_command(self, tmp_path):
|
||||
pkg = tmp_path / "bad"
|
||||
pkg.mkdir()
|
||||
(pkg / "mcp.json").write_text(json.dumps({"run": {}}))
|
||||
|
||||
with pytest.raises(SystemExit):
|
||||
_iter_packages(tmp_path)
|
||||
|
||||
def test_exits_on_invalid_json(self, tmp_path):
|
||||
pkg = tmp_path / "broken"
|
||||
pkg.mkdir()
|
||||
(pkg / "mcp.json").write_text("not json")
|
||||
|
||||
with pytest.raises(SystemExit):
|
||||
_iter_packages(tmp_path)
|
||||
|
||||
def test_exits_on_nonexistent_root(self, tmp_path):
|
||||
with pytest.raises(SystemExit):
|
||||
_iter_packages(tmp_path / "nope")
|
||||
|
||||
def test_sorted_by_name(self, tmp_path):
|
||||
for name in ("zebra", "alpha", "middle"):
|
||||
pkg = tmp_path / name
|
||||
pkg.mkdir()
|
||||
(pkg / "mcp.json").write_text(json.dumps({"run": {"command": "python"}}))
|
||||
|
||||
servers = _iter_packages(tmp_path)
|
||||
names = [s.name for s in servers]
|
||||
assert names == ["alpha", "middle", "zebra"]
|
||||
|
||||
|
||||
class TestDiscoverMcpServers:
|
||||
def test_filter_by_namespaced_toolset(self, tmp_path):
|
||||
"""Only servers declaring the requested namespaced toolset are returned."""
|
||||
(tmp_path / "alpha").mkdir()
|
||||
(tmp_path / "alpha" / "mcp.json").write_text(
|
||||
json.dumps({"run": {"command": "node"}, "toolsets": {"read": ["t1"]}})
|
||||
)
|
||||
(tmp_path / "beta").mkdir()
|
||||
(tmp_path / "beta" / "mcp.json").write_text(
|
||||
json.dumps({"run": {"command": "node"}, "toolsets": {"write": ["t2"]}})
|
||||
)
|
||||
(tmp_path / "gamma").mkdir()
|
||||
(tmp_path / "gamma" / "mcp.json").write_text(
|
||||
json.dumps({"run": {"command": "node"}, "toolsets": {"read": ["t3"]}})
|
||||
)
|
||||
|
||||
servers = discover_mcp_servers(tmp_path, tool_sets=["alpha_read", "gamma_read"])
|
||||
assert [s.name for s in servers] == ["alpha", "gamma"]
|
||||
|
||||
def test_no_filter_returns_nothing(self, tmp_path):
|
||||
"""tool_sets=None / [] / [""] all mean 'no servers' — callers must opt in
|
||||
explicitly. Use _iter_packages when you need every package."""
|
||||
for name in ("alpha", "beta"):
|
||||
pkg = tmp_path / name
|
||||
pkg.mkdir()
|
||||
(pkg / "mcp.json").write_text(json.dumps({"run": {"command": "node"}}))
|
||||
|
||||
assert discover_mcp_servers(tmp_path, tool_sets=None) == []
|
||||
assert discover_mcp_servers(tmp_path, tool_sets=[]) == []
|
||||
assert discover_mcp_servers(tmp_path, tool_sets=[""]) == []
|
||||
|
||||
def test_filter_multiple_namespaced_toolsets(self, tmp_path):
|
||||
"""Servers declaring any of the requested namespaced toolsets are included."""
|
||||
(tmp_path / "alpha").mkdir()
|
||||
(tmp_path / "alpha" / "mcp.json").write_text(
|
||||
json.dumps({"run": {"command": "node"}, "toolsets": {"read": ["t1"]}})
|
||||
)
|
||||
(tmp_path / "beta").mkdir()
|
||||
(tmp_path / "beta" / "mcp.json").write_text(
|
||||
json.dumps({"run": {"command": "node"}, "toolsets": {"write": ["t2"]}})
|
||||
)
|
||||
(tmp_path / "gamma").mkdir()
|
||||
(tmp_path / "gamma" / "mcp.json").write_text(
|
||||
json.dumps({"run": {"command": "node"}, "toolsets": {"debug": ["t3"]}})
|
||||
)
|
||||
|
||||
servers = discover_mcp_servers(tmp_path, tool_sets=["alpha_read", "beta_write"])
|
||||
assert [s.name for s in servers] == ["alpha", "beta"]
|
||||
|
||||
def test_filter_namespaced_toolset(self, tmp_path):
|
||||
"""Namespaced form 'pkg_toolset' selects only that package+toolset."""
|
||||
(tmp_path / "google_mail").mkdir()
|
||||
(tmp_path / "google_mail" / "mcp.json").write_text(
|
||||
json.dumps({"run": {"command": "node"}, "toolsets": {"read": ["t1"], "write": ["t2"]}})
|
||||
)
|
||||
(tmp_path / "slack").mkdir()
|
||||
(tmp_path / "slack" / "mcp.json").write_text(
|
||||
json.dumps({"run": {"command": "node"}, "toolsets": {"read": ["t3"], "write": ["t4"]}})
|
||||
)
|
||||
(tmp_path / "jira").mkdir()
|
||||
(tmp_path / "jira" / "mcp.json").write_text(
|
||||
json.dumps({"run": {"command": "node"}}) # no toolsets
|
||||
)
|
||||
|
||||
# Only google_mail should be found; slack has no google_mail_read toolset
|
||||
servers = discover_mcp_servers(tmp_path, tool_sets=["google_mail_read"])
|
||||
assert [s.name for s in servers] == ["google_mail"]
|
||||
|
||||
def test_returns_mcp_service_objects(self, tmp_path):
|
||||
"""discover_mcp_servers returns McpService instances wrapping McpConfig."""
|
||||
pkg = tmp_path / "svc"
|
||||
pkg.mkdir()
|
||||
(pkg / "mcp.json").write_text(json.dumps({"run": {"command": "node"}, "toolsets": {"read": ["t"]}}))
|
||||
|
||||
servers = discover_mcp_servers(tmp_path, tool_sets=["svc_read"])
|
||||
assert len(servers) == 1
|
||||
assert isinstance(servers[0], McpService)
|
||||
assert isinstance(servers[0].config, McpConfig)
|
||||
|
||||
def test_bare_package_name_is_rejected(self, tmp_path):
|
||||
"""Bare package names are not valid tool_sets — they must be namespaced
|
||||
toolset names (``google_mail_read``), so validation rejects them."""
|
||||
(tmp_path / "google_mail").mkdir()
|
||||
(tmp_path / "google_mail" / "mcp.json").write_text(
|
||||
json.dumps({"run": {"command": "node"}, "toolsets": {"read": ["t1"]}})
|
||||
)
|
||||
# "google_mail" is a package name, not a toolset name — unknown tool set.
|
||||
with pytest.raises(SystemExit):
|
||||
discover_mcp_servers(tmp_path, tool_sets=["google_mail"])
|
||||
|
||||
def _core_and_shim(self, tmp_path):
|
||||
"""Create fake ``core`` and ``syntara`` (compat shim) packages."""
|
||||
(tmp_path / "core").mkdir()
|
||||
(tmp_path / "core" / "mcp.json").write_text(
|
||||
json.dumps({"run": {"command": "python"}, "toolsets": {"read": ["readFile"]}})
|
||||
)
|
||||
(tmp_path / "syntara").mkdir()
|
||||
(tmp_path / "syntara" / "mcp.json").write_text(
|
||||
json.dumps({"run": {"command": "python"}, "toolsets": {"read": ["readFile"]}})
|
||||
)
|
||||
|
||||
def test_core_and_syntara_shim_together_warns(self, tmp_path, caplog):
|
||||
"""core + the syntara compat shim are redundant but allowed — selecting both
|
||||
mounts both surfaces and warns. REMOVE with the syntara package after 2026-06-18."""
|
||||
self._core_and_shim(tmp_path)
|
||||
with caplog.at_level(logging.WARNING, logger="mcp_proxy"):
|
||||
servers = discover_mcp_servers(tmp_path, tool_sets=["core_read", "syntara_read"])
|
||||
assert sorted(s.name for s in servers) == ["core", "syntara"]
|
||||
assert any("both 'core' and the legacy 'syntara'" in r.message for r in caplog.records)
|
||||
|
||||
def test_core_alone_is_allowed(self, tmp_path):
|
||||
self._core_and_shim(tmp_path)
|
||||
assert [s.name for s in discover_mcp_servers(tmp_path, tool_sets=["core_read"])] == ["core"]
|
||||
|
||||
def test_syntara_shim_alone_is_allowed(self, tmp_path):
|
||||
self._core_and_shim(tmp_path)
|
||||
assert [s.name for s in discover_mcp_servers(tmp_path, tool_sets=["syntara_read"])] == ["syntara"]
|
||||
|
||||
|
||||
class TestValidateToolSets:
|
||||
@staticmethod
|
||||
def _write_pkg(root, name, toolsets):
|
||||
"""Materialize a minimal ``packages/<name>/mcp.json`` declaring *toolsets*."""
|
||||
pkg = root / name
|
||||
pkg.mkdir()
|
||||
(pkg / "mcp.json").write_text(json.dumps({"run": {"command": "node"}, "toolsets": toolsets}))
|
||||
|
||||
def test_valid_tool_sets_pass(self, tmp_path):
|
||||
"""Namespaced tool set names that a package declares do not raise."""
|
||||
self._write_pkg(tmp_path, "google_mail", {"read": []})
|
||||
self._write_pkg(tmp_path, "slack", {"state": []})
|
||||
_validate_tool_sets(["google_mail_read", "slack_state"], tmp_path)
|
||||
|
||||
def test_unknown_tool_set_exits(self, tmp_path):
|
||||
"""Unknown tool set names cause sys.exit."""
|
||||
self._write_pkg(tmp_path, "slack", {"read": []})
|
||||
with pytest.raises(SystemExit):
|
||||
_validate_tool_sets(["nonexistent"], tmp_path)
|
||||
|
||||
def test_partial_unknown_exits(self, tmp_path):
|
||||
"""Even one unknown tool set among valid ones causes exit."""
|
||||
self._write_pkg(tmp_path, "slack", {"read": []})
|
||||
with pytest.raises(SystemExit):
|
||||
_validate_tool_sets(["slack_read", "bad_name"], tmp_path)
|
||||
|
||||
def test_empty_request_passes(self, tmp_path):
|
||||
"""No requested tool sets is vacuously valid (nothing unknown)."""
|
||||
self._write_pkg(tmp_path, "slack", {"read": []})
|
||||
_validate_tool_sets([], tmp_path) # should not raise
|
||||
|
||||
|
||||
class TestBuildProxyApp:
|
||||
def test_exits_when_selected_service_fails_startup(self, tmp_path, monkeypatch, caplog):
|
||||
"""Requested services that die during startup abort the whole proxy."""
|
||||
|
||||
class FakeProc:
|
||||
returncode = 7
|
||||
|
||||
def poll(self):
|
||||
return self.returncode
|
||||
|
||||
def terminate(self):
|
||||
pass
|
||||
|
||||
class FakeService:
|
||||
name = "jira"
|
||||
secrets: ClassVar[list[str]] = []
|
||||
config = SimpleNamespace(startup_timeout=0.1)
|
||||
|
||||
def run_install(self, env):
|
||||
return True
|
||||
|
||||
def run_setup(self, env):
|
||||
return True
|
||||
|
||||
def run_pre_steps(self, env):
|
||||
return True
|
||||
|
||||
def start_service(self, env, port, proxy_token, command_prefix=None):
|
||||
return FakeProc()
|
||||
|
||||
monkeypatch.setattr(
|
||||
"mcp_proxy.commands.mcp.discover_mcp_servers",
|
||||
lambda packages_root, tool_sets=None: [FakeService()],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"mcp_proxy.commands.mcp._build_subprocess_env",
|
||||
lambda base_dir, server_name, declared_secrets=(): {},
|
||||
)
|
||||
|
||||
with pytest.raises(SystemExit):
|
||||
build_proxy_app(tmp_path, tmp_path, tool_sets=["jira_read"])
|
||||
|
||||
assert "Aborting startup because requested MCP services failed to start" in caplog.text
|
||||
assert "jira (exited immediately with code 7)" in caplog.text
|
||||
|
||||
|
||||
class TestRunViewerStartup:
|
||||
class FakeApp:
|
||||
def __init__(self):
|
||||
self.run_calls = []
|
||||
|
||||
def custom_route(self, *_args, **_kwargs):
|
||||
def decorator(fn):
|
||||
return fn
|
||||
|
||||
return decorator
|
||||
|
||||
def run(self, **kwargs):
|
||||
self.run_calls.append(kwargs)
|
||||
|
||||
def _stub_proxy_run(self, monkeypatch, tmp_path):
|
||||
app = self.FakeApp()
|
||||
monkeypatch.setenv("WORLDBENCH_ROOT", str(tmp_path))
|
||||
monkeypatch.setenv("WORLDBENCH_PACKAGES_ROOT", str(tmp_path))
|
||||
monkeypatch.delenv("WORLDBENCH_TOOL_SETS", raising=False)
|
||||
monkeypatch.delenv("PORT", raising=False)
|
||||
monkeypatch.setattr(mcp_command, "install_crash_handlers", lambda: None)
|
||||
monkeypatch.setattr(mcp_command, "build_proxy_app", lambda **_kwargs: app)
|
||||
monkeypatch.setattr(mcp_command, "_assert_local_tools_async", lambda _app: None)
|
||||
monkeypatch.setattr(mcp_command, "_kill_child_processes", lambda: None)
|
||||
return app
|
||||
|
||||
def test_does_not_start_viewer_when_viewer_port_is_unset(self, tmp_path, monkeypatch):
|
||||
app = self._stub_proxy_run(monkeypatch, tmp_path)
|
||||
viewer_calls = []
|
||||
monkeypatch.delenv("VIEWER_PORT", raising=False)
|
||||
monkeypatch.setattr(mcp_command, "_start_viewer_server", lambda host, port: viewer_calls.append((host, port)))
|
||||
|
||||
mcp_command.run(method="stdio")
|
||||
|
||||
assert viewer_calls == []
|
||||
assert app.run_calls == [{"show_banner": False}]
|
||||
|
||||
def test_starts_viewer_when_viewer_port_is_set(self, tmp_path, monkeypatch):
|
||||
app = self._stub_proxy_run(monkeypatch, tmp_path)
|
||||
viewer_calls = []
|
||||
monkeypatch.setenv("VIEWER_PORT", "8765")
|
||||
monkeypatch.setattr(mcp_command, "_start_viewer_server", lambda host, port: viewer_calls.append((host, port)))
|
||||
|
||||
mcp_command.run(method="stdio")
|
||||
|
||||
assert viewer_calls == [("0.0.0.0", 8765)]
|
||||
assert app.run_calls == [{"show_banner": False}]
|
||||
|
||||
|
||||
class TestHideTaggedToolsMiddleware:
|
||||
"""The proxy's un-namespaced ``export_state`` / ``import_state`` are
|
||||
infrastructure for the grading harness: they must remain callable so
|
||||
snapshots round-trip, but they must NOT appear in ``tools/list`` (an
|
||||
agent shouldn't even know they exist).
|
||||
"""
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_tagged_tools_are_hidden_from_list_but_callable(self):
|
||||
app = FastMCP("test")
|
||||
app.add_middleware(HideTaggedToolsMiddleware())
|
||||
|
||||
@app.tool(name="visible")
|
||||
async def visible() -> dict:
|
||||
return {"v": True}
|
||||
|
||||
@app.tool(name="hidden", tags={HIDDEN_FROM_LISTING_TAG})
|
||||
async def hidden() -> dict:
|
||||
return {"h": True}
|
||||
|
||||
listed = {t.name for t in await app.list_tools()}
|
||||
assert listed == {"visible"}, f"hidden tool leaked into list_tools: {listed}"
|
||||
|
||||
# Hidden tool must still be callable — that's the point.
|
||||
hidden_result = await app.call_tool("hidden", {})
|
||||
assert hidden_result.structured_content == {"h": True}
|
||||
|
||||
visible_result = await app.call_tool("visible", {})
|
||||
assert visible_result.structured_content == {"v": True}
|
||||
|
||||
@pytest.fixture
|
||||
def anyio_backend(self):
|
||||
return "asyncio"
|
||||
|
||||
|
||||
class TestMcpServiceLifecycle:
|
||||
@pytest.fixture
|
||||
def env(self):
|
||||
"""Return a minimal env dict with PATH so subprocesses can find python."""
|
||||
import os
|
||||
|
||||
return dict(os.environ)
|
||||
|
||||
def test_run_install_executes_steps(self, tmp_path, env):
|
||||
"""Install steps are executed and create side effects."""
|
||||
marker = tmp_path / "installed.txt"
|
||||
cfg = McpConfig.from_mcp_json(
|
||||
tmp_path,
|
||||
{
|
||||
"run": {"command": "python", "args": ["-c", "pass"]},
|
||||
"install": {"command": "python", "args": ["-c", f"open('{marker}', 'w').write('ok')"]},
|
||||
},
|
||||
)
|
||||
svc = McpService(cfg)
|
||||
assert svc.run_install(env) is True
|
||||
assert marker.read_text() == "ok"
|
||||
|
||||
def test_run_install_returns_true_when_empty(self, tmp_path, env):
|
||||
"""No install steps means success."""
|
||||
cfg = McpConfig.from_mcp_json(tmp_path, {"run": {"command": "python"}})
|
||||
svc = McpService(cfg)
|
||||
assert svc.run_install(env) is True
|
||||
|
||||
def test_run_install_returns_false_on_failure(self, tmp_path, env):
|
||||
"""Failed install step returns False."""
|
||||
cfg = McpConfig.from_mcp_json(
|
||||
tmp_path,
|
||||
{
|
||||
"run": {"command": "python"},
|
||||
"install": {"command": "python", "args": ["-c", "import sys; sys.exit(1)"]},
|
||||
},
|
||||
)
|
||||
svc = McpService(cfg)
|
||||
assert svc.run_install(env) is False
|
||||
|
||||
def test_run_setup_executes_steps(self, tmp_path, env):
|
||||
"""Setup steps are executed and create side effects."""
|
||||
marker = tmp_path / "setup.txt"
|
||||
cfg = McpConfig.from_mcp_json(
|
||||
tmp_path,
|
||||
{
|
||||
"run": {"command": "python"},
|
||||
"setup": {"command": "python", "args": ["-c", f"open('{marker}', 'w').write('done')"]},
|
||||
},
|
||||
)
|
||||
svc = McpService(cfg)
|
||||
assert svc.run_setup(env) is True
|
||||
assert marker.read_text() == "done"
|
||||
|
||||
def test_run_setup_forwards_env(self, tmp_path, env):
|
||||
"""Environment variables are forwarded to setup steps."""
|
||||
marker = tmp_path / "env.txt"
|
||||
env["MY_VAR"] = "hello"
|
||||
cfg = McpConfig.from_mcp_json(
|
||||
tmp_path,
|
||||
{
|
||||
"run": {"command": "python"},
|
||||
"setup": {
|
||||
"command": "python",
|
||||
"args": ["-c", f"import os; open('{marker}', 'w').write(os.environ.get('MY_VAR', 'MISSING'))"],
|
||||
},
|
||||
},
|
||||
)
|
||||
svc = McpService(cfg)
|
||||
assert svc.run_setup(env) is True
|
||||
assert marker.read_text() == "hello"
|
||||
|
||||
def test_python_steps_use_proxy_interpreter_not_path(self, tmp_path, env):
|
||||
# Empty PATH: a bare `python` would be unfindable, so success proves the
|
||||
# step ran the absolute sys.executable (the PATH-elimination contract).
|
||||
import sys
|
||||
|
||||
marker = tmp_path / "interp.txt"
|
||||
env["PATH"] = ""
|
||||
cfg = McpConfig.from_mcp_json(
|
||||
tmp_path,
|
||||
{
|
||||
"run": {"command": "python"},
|
||||
"setup": {
|
||||
"command": "python",
|
||||
"args": ["-c", f"import sys; open('{marker}', 'w').write(sys.executable)"],
|
||||
},
|
||||
},
|
||||
)
|
||||
svc = McpService(cfg)
|
||||
assert svc.run_setup(env) is True, "bare `python` step failed — interpreter not resolved to an absolute path"
|
||||
assert marker.read_text() == sys.executable
|
||||
|
||||
|
||||
class TestResolveCommand:
|
||||
def test_bare_python_resolves_to_sys_executable(self):
|
||||
import sys
|
||||
|
||||
assert resolve_command("python") == sys.executable
|
||||
assert resolve_command("python3") == sys.executable
|
||||
|
||||
def test_non_python_commands_pass_through(self):
|
||||
for command in ("node", "bash", "npm", "uv", "/opt/venv/bin/python", "frappe-bench/env/bin/python"):
|
||||
assert resolve_command(command) == command
|
||||
@@ -0,0 +1,183 @@
|
||||
"""Tests for resolve_tool_transforms in mcp_proxy.service."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from mcp_proxy.service import McpConfig, McpService
|
||||
|
||||
|
||||
def _make_cfg(tmp_path, tools: list[str], toolsets: dict, name: str | None = None) -> McpService:
|
||||
"""Write mcp.json + mcp-tools.generated.json and return an McpService."""
|
||||
mcp_json: dict = {
|
||||
"run": {"command": "python", "args": ["-m", "fake"]},
|
||||
"toolsets": toolsets,
|
||||
}
|
||||
if name is not None:
|
||||
mcp_json["name"] = name
|
||||
(tmp_path / "mcp.json").write_text(json.dumps(mcp_json))
|
||||
tools_data = {
|
||||
"schema_version": "v1",
|
||||
"tools": [{"name": t, "description": "", "inputSchema": {}} for t in tools],
|
||||
"toolsets": toolsets,
|
||||
}
|
||||
(tmp_path / "mcp-tools.generated.json").write_text(json.dumps(tools_data))
|
||||
return McpService(McpConfig.from_mcp_json(tmp_path, mcp_json))
|
||||
|
||||
|
||||
def _enabled_tools(transforms: dict) -> dict[str, str]:
|
||||
"""Return {original_name: new_name} for enabled tools only."""
|
||||
return {name: t.name for name, t in transforms.items() if t.enabled and t.name}
|
||||
|
||||
|
||||
def _disabled_tools(transforms: dict) -> set[str]:
|
||||
"""Return set of original names for disabled tools."""
|
||||
return {name for name, t in transforms.items() if not t.enabled}
|
||||
|
||||
|
||||
TOOLS = ["echo", "readFile", "writeFile", "bash"]
|
||||
TOOLSETS = {
|
||||
"read": ["echo", "readFile"],
|
||||
"write": ["writeFile", "bash"],
|
||||
}
|
||||
|
||||
|
||||
def test_empty_tool_sets_exposes_all_tools(tmp_path):
|
||||
"""An empty filter (or unset env var) means 'expose every introspected tool'."""
|
||||
cfg = _make_cfg(tmp_path, TOOLS, TOOLSETS)
|
||||
result = cfg.resolve_tool_transforms([])
|
||||
assert set(_enabled_tools(result).keys()) == set(TOOLS)
|
||||
assert _disabled_tools(result) == set()
|
||||
|
||||
|
||||
def test_namespaced_read_filters_correctly(tmp_path):
|
||||
pkg = tmp_path / "slack"
|
||||
pkg.mkdir()
|
||||
cfg = _make_cfg(pkg, TOOLS, TOOLSETS)
|
||||
result = cfg.resolve_tool_transforms(["slack_read"])
|
||||
assert set(_enabled_tools(result).keys()) == {"echo", "readFile"}
|
||||
assert _disabled_tools(result) == {"writeFile", "bash"}
|
||||
|
||||
|
||||
def test_namespaced_write_filters_correctly(tmp_path):
|
||||
pkg = tmp_path / "slack"
|
||||
pkg.mkdir()
|
||||
cfg = _make_cfg(pkg, TOOLS, TOOLSETS)
|
||||
result = cfg.resolve_tool_transforms(["slack_write"])
|
||||
assert set(_enabled_tools(result).keys()) == {"writeFile", "bash"}
|
||||
assert _disabled_tools(result) == {"echo", "readFile"}
|
||||
|
||||
|
||||
def test_multiple_namespaced_toolsets_are_unioned(tmp_path):
|
||||
pkg = tmp_path / "slack"
|
||||
pkg.mkdir()
|
||||
cfg = _make_cfg(pkg, TOOLS, TOOLSETS)
|
||||
result = cfg.resolve_tool_transforms(["slack_read", "slack_write"])
|
||||
assert set(_enabled_tools(result).keys()) == set(TOOLS)
|
||||
assert _disabled_tools(result) == set()
|
||||
|
||||
|
||||
def test_tools_are_namespaced(tmp_path):
|
||||
pkg = tmp_path / "slack"
|
||||
pkg.mkdir()
|
||||
cfg = _make_cfg(pkg, TOOLS, TOOLSETS)
|
||||
result = cfg.resolve_tool_transforms(["slack_read"])
|
||||
for original, new_name in _enabled_tools(result).items():
|
||||
assert new_name == f"slack__{original}"
|
||||
|
||||
|
||||
def test_unnamespaced_package_exposes_bare_tool_names(tmp_path):
|
||||
"""When mcp.json sets namespaced=false, tools are exposed without a prefix."""
|
||||
mcp_json = {
|
||||
"run": {"command": "python", "args": ["-m", "fake"]},
|
||||
"namespaced": False,
|
||||
"toolsets": TOOLSETS,
|
||||
}
|
||||
(tmp_path / "mcp.json").write_text(json.dumps(mcp_json))
|
||||
tools_data = {
|
||||
"schema_version": "v1",
|
||||
"tools": [{"name": t, "description": "", "inputSchema": {}} for t in TOOLS],
|
||||
"toolsets": TOOLSETS,
|
||||
}
|
||||
(tmp_path / "mcp-tools.generated.json").write_text(json.dumps(tools_data))
|
||||
cfg = McpService(McpConfig.from_mcp_json(tmp_path, mcp_json))
|
||||
result = cfg.resolve_tool_transforms([])
|
||||
for original, new_name in _enabled_tools(result).items():
|
||||
assert new_name == original
|
||||
|
||||
|
||||
def test_no_toolsets_in_mcp_json_exposes_all(tmp_path):
|
||||
"""When mcp.json has no toolsets key, all tools are exposed regardless of filter."""
|
||||
cfg = _make_cfg(tmp_path, TOOLS, {})
|
||||
result = cfg.resolve_tool_transforms(["whatever"])
|
||||
assert set(_enabled_tools(result).keys()) == set(TOOLS)
|
||||
assert _disabled_tools(result) == set()
|
||||
|
||||
|
||||
def test_unknown_toolset_disables_all(tmp_path):
|
||||
"""A namespaced toolset that doesn't exist on this package disables every tool."""
|
||||
pkg = tmp_path / "slack"
|
||||
pkg.mkdir()
|
||||
cfg = _make_cfg(pkg, TOOLS, TOOLSETS)
|
||||
result = cfg.resolve_tool_transforms(["slack_nonexistent"])
|
||||
assert _enabled_tools(result) == {}
|
||||
assert _disabled_tools(result) == set(TOOLS)
|
||||
|
||||
|
||||
def test_other_packages_namespace_does_not_match(tmp_path):
|
||||
"""A toolset namespaced to a different package is ignored — matches no tools."""
|
||||
pkg = tmp_path / "slack"
|
||||
pkg.mkdir()
|
||||
cfg = _make_cfg(pkg, TOOLS, TOOLSETS)
|
||||
result = cfg.resolve_tool_transforms(["jira_read"])
|
||||
assert _enabled_tools(result) == {}
|
||||
assert _disabled_tools(result) == set(TOOLS)
|
||||
|
||||
|
||||
def test_missing_generated_json_raises(tmp_path):
|
||||
mcp_json = {"run": {"command": "python", "args": []}, "toolsets": TOOLSETS}
|
||||
(tmp_path / "mcp.json").write_text(json.dumps(mcp_json))
|
||||
cfg = McpService(McpConfig.from_mcp_json(tmp_path, mcp_json))
|
||||
with pytest.raises(FileNotFoundError, match=r"mcp-tools\.generated\.json not found"):
|
||||
cfg.resolve_tool_transforms(["slack_read"])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# State toolsets are reachable when explicitly requested. The "all" / "_all"
|
||||
# special case (with its hidden-by-default semantics for state/grading tools)
|
||||
# was removed when the auto-aggregated cross-package toolsets were dropped.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
STATEFUL_TOOLS = ["echo", "readFile", "export_state", "import_state"]
|
||||
STATEFUL_TOOLSETS = {
|
||||
"read": ["echo", "readFile"],
|
||||
"state": ["export_state", "import_state"],
|
||||
}
|
||||
|
||||
|
||||
def test_namespaced_state_request_exposes_state(tmp_path):
|
||||
pkg = tmp_path / "slack"
|
||||
pkg.mkdir()
|
||||
cfg = _make_cfg(pkg, STATEFUL_TOOLS, STATEFUL_TOOLSETS)
|
||||
result = cfg.resolve_tool_transforms(["slack_state"])
|
||||
assert set(_enabled_tools(result).keys()) == {"export_state", "import_state"}
|
||||
|
||||
|
||||
def test_namespaced_read_does_not_expose_state(tmp_path):
|
||||
pkg = tmp_path / "slack"
|
||||
pkg.mkdir()
|
||||
cfg = _make_cfg(pkg, STATEFUL_TOOLS, STATEFUL_TOOLSETS)
|
||||
result = cfg.resolve_tool_transforms(["slack_read"])
|
||||
assert set(_enabled_tools(result).keys()) == {"echo", "readFile"}
|
||||
assert _disabled_tools(result) == {"export_state", "import_state"}
|
||||
|
||||
|
||||
def test_combined_namespaced_toolsets_union(tmp_path):
|
||||
pkg = tmp_path / "slack"
|
||||
pkg.mkdir()
|
||||
cfg = _make_cfg(pkg, STATEFUL_TOOLS, STATEFUL_TOOLSETS)
|
||||
result = cfg.resolve_tool_transforms(["slack_read", "slack_state"])
|
||||
assert set(_enabled_tools(result).keys()) == set(STATEFUL_TOOLS)
|
||||
@@ -0,0 +1,89 @@
|
||||
"""Tests for service_http module."""
|
||||
|
||||
import pytest
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from mcp_proxy.service_http import ProxyTokenMiddleware, _placeholder_html
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Placeholder HTML
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestPlaceholderHtml:
|
||||
def test_contains_service_name(self):
|
||||
html = _placeholder_html("My Service")
|
||||
assert "My Service" in html
|
||||
|
||||
def test_is_valid_html(self):
|
||||
html = _placeholder_html("Test")
|
||||
assert html.startswith("<!DOCTYPE html>")
|
||||
assert "</html>" in html
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ProxyTokenMiddleware
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestProxyTokenMiddleware:
|
||||
@pytest.fixture
|
||||
def app(self):
|
||||
"""Create a minimal Starlette app with ProxyTokenMiddleware."""
|
||||
from starlette.applications import Starlette
|
||||
from starlette.middleware import Middleware
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import PlainTextResponse
|
||||
from starlette.routing import Route
|
||||
|
||||
async def index(request: Request):
|
||||
return PlainTextResponse("ok")
|
||||
|
||||
async def mcp_endpoint(request: Request):
|
||||
return PlainTextResponse("mcp ok")
|
||||
|
||||
return Starlette(
|
||||
routes=[
|
||||
Route("/", index),
|
||||
Route("/mcp", mcp_endpoint),
|
||||
Route("/mcp/test", mcp_endpoint),
|
||||
Route("/api/data", index),
|
||||
],
|
||||
middleware=[Middleware(ProxyTokenMiddleware)],
|
||||
)
|
||||
|
||||
def test_mcp_path_bypasses_auth(self, app, monkeypatch):
|
||||
monkeypatch.setenv("MCP_PROXY_TOKEN", "secret")
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.get("/mcp")
|
||||
assert resp.status_code == 200
|
||||
|
||||
def test_mcp_subpath_bypasses_auth(self, app, monkeypatch):
|
||||
monkeypatch.setenv("MCP_PROXY_TOKEN", "secret")
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.get("/mcp/test")
|
||||
assert resp.status_code == 200
|
||||
|
||||
def test_non_mcp_blocked_without_token(self, app, monkeypatch):
|
||||
monkeypatch.setenv("MCP_PROXY_TOKEN", "secret")
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.get("/")
|
||||
assert resp.status_code == 403
|
||||
|
||||
def test_non_mcp_allowed_with_correct_token(self, app, monkeypatch):
|
||||
monkeypatch.setenv("MCP_PROXY_TOKEN", "secret")
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.get("/", headers={"x-proxy-token": "secret"})
|
||||
assert resp.status_code == 200
|
||||
|
||||
def test_non_mcp_blocked_with_wrong_token(self, app, monkeypatch):
|
||||
monkeypatch.setenv("MCP_PROXY_TOKEN", "secret")
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.get("/", headers={"x-proxy-token": "wrong"})
|
||||
assert resp.status_code == 403
|
||||
|
||||
def test_no_token_set_allows_all(self, app, monkeypatch):
|
||||
monkeypatch.delenv("MCP_PROXY_TOKEN", raising=False)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = client.get("/")
|
||||
assert resp.status_code == 200
|
||||
@@ -0,0 +1,183 @@
|
||||
"""Tests for the setup pass the proxy runs during ``mcp`` startup.
|
||||
|
||||
There is no standalone setup command: ``build_proxy_app`` discovers MCP
|
||||
servers and runs each one's ``install`` then ``setup`` hook from
|
||||
``mcp.json`` (forwarding WORLDBENCH_* env vars) before starting the
|
||||
server. Servers are selected by WORLDBENCH_TOOL_SETS: only servers whose
|
||||
mcp.json declares at least one of the requested namespaced toolsets are
|
||||
included. An empty / unset value runs no setup hooks.
|
||||
|
||||
These tests exercise that same discover -> install -> setup pass in
|
||||
process, using the exact primitives ``build_proxy_app`` calls.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
from mcp_proxy.commands.mcp import _build_subprocess_env, discover_mcp_servers
|
||||
|
||||
|
||||
class SetupCommandTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp_dir = tempfile.mkdtemp()
|
||||
self.packages_root = Path(self.temp_dir) / "packages"
|
||||
self.packages_root.mkdir()
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(self.temp_dir)
|
||||
|
||||
def _create_server(self, name, setup_hook=None, toolsets=None):
|
||||
"""Create a fake MCP server package with an optional setup hook.
|
||||
|
||||
Always declares a ``"default"`` toolset (so :meth:`_run_setup` can
|
||||
opt every package in by default). *toolsets* is merged on top, in
|
||||
case a test wants to declare additional toolsets to filter against.
|
||||
"""
|
||||
pkg_dir = self.packages_root / name
|
||||
pkg_dir.mkdir(parents=True)
|
||||
|
||||
config = {"run": {"command": "python", "args": ["-c", "pass"]}}
|
||||
if setup_hook is not None:
|
||||
config["setup"] = setup_hook
|
||||
config["toolsets"] = {"default": [], **(toolsets or {})}
|
||||
|
||||
(pkg_dir / "mcp.json").write_text(json.dumps(config))
|
||||
return pkg_dir
|
||||
|
||||
def _run_setup(self, env_overrides=None):
|
||||
"""Run the startup setup pass in process and return the failed servers.
|
||||
|
||||
Mirrors the discover -> install -> setup loop ``build_proxy_app``
|
||||
runs before booting (the only place setup hooks fire now). Returns
|
||||
the list of server names whose install or setup hook failed — empty
|
||||
on success.
|
||||
|
||||
Defaults WORLDBENCH_TOOL_SETS to every fake package's ``_default``
|
||||
toolset so each test gets full coverage without spelling them out.
|
||||
Override via *env_overrides* when a test wants narrower filtering.
|
||||
"""
|
||||
default_sets = " ".join(f"{p.name}_default" for p in sorted(self.packages_root.iterdir()) if p.is_dir())
|
||||
env = {
|
||||
"WORLDBENCH_ROOT": self.temp_dir,
|
||||
"WORLDBENCH_PACKAGES_ROOT": str(self.packages_root),
|
||||
"WORLDBENCH_TOOL_SETS": default_sets,
|
||||
}
|
||||
if env_overrides:
|
||||
env.update(env_overrides)
|
||||
|
||||
failed: list[str] = []
|
||||
with mock.patch.dict(os.environ, env):
|
||||
base_dir = Path(os.environ["WORLDBENCH_ROOT"]).resolve()
|
||||
packages_root = Path(os.environ["WORLDBENCH_PACKAGES_ROOT"])
|
||||
tool_sets = os.environ["WORLDBENCH_TOOL_SETS"].split()
|
||||
for cfg in discover_mcp_servers(packages_root, tool_sets=tool_sets):
|
||||
hook_env = _build_subprocess_env(base_dir, cfg.name, declared_secrets=cfg.secrets)
|
||||
if not cfg.run_install(hook_env) or not cfg.run_setup(hook_env):
|
||||
failed.append(cfg.name)
|
||||
return failed
|
||||
|
||||
def test_runs_setup_hook(self):
|
||||
"""Setup hook from mcp.json is executed."""
|
||||
marker = Path(self.temp_dir) / "setup_ran.txt"
|
||||
self._create_server(
|
||||
"svc",
|
||||
setup_hook={
|
||||
"command": "python",
|
||||
"args": ["-c", f"open('{marker}', 'w').write('ok')"],
|
||||
},
|
||||
)
|
||||
|
||||
failed = self._run_setup()
|
||||
assert failed == [], failed
|
||||
assert marker.exists()
|
||||
assert marker.read_text() == "ok"
|
||||
|
||||
def test_skips_server_without_setup_hook(self):
|
||||
"""Servers without a setup hook are skipped gracefully."""
|
||||
self._create_server("no-setup")
|
||||
|
||||
failed = self._run_setup()
|
||||
assert failed == [], failed
|
||||
|
||||
def test_forwards_env_vars(self):
|
||||
"""WORLDBENCH_* env vars are forwarded to the setup hook."""
|
||||
marker = Path(self.temp_dir) / "env_check.txt"
|
||||
self._create_server(
|
||||
"svc",
|
||||
setup_hook={
|
||||
"command": "python",
|
||||
"args": [
|
||||
"-c",
|
||||
f"import os; open('{marker}', 'w').write(os.environ.get('WORLDBENCH_TASK_ID', 'MISSING'))",
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
failed = self._run_setup({"WORLDBENCH_TASK_ID": "test-123"})
|
||||
assert failed == [], failed
|
||||
assert marker.read_text() == "test-123"
|
||||
|
||||
def test_respects_server_filter(self):
|
||||
"""Only servers declaring the requested toolset are set up."""
|
||||
marker_a = Path(self.temp_dir) / "a_ran.txt"
|
||||
marker_b = Path(self.temp_dir) / "b_ran.txt"
|
||||
self._create_server(
|
||||
"alpha",
|
||||
setup_hook={
|
||||
"command": "python",
|
||||
"args": ["-c", f"open('{marker_a}', 'w').write('ok')"],
|
||||
},
|
||||
toolsets={"read": ["tool_a"]},
|
||||
)
|
||||
self._create_server(
|
||||
"beta",
|
||||
setup_hook={
|
||||
"command": "python",
|
||||
"args": ["-c", f"open('{marker_b}', 'w').write('ok')"],
|
||||
},
|
||||
toolsets={"write": ["tool_b"]},
|
||||
)
|
||||
|
||||
failed = self._run_setup({"WORLDBENCH_TOOL_SETS": "alpha_read"})
|
||||
assert failed == [], failed
|
||||
assert marker_a.exists()
|
||||
assert not marker_b.exists()
|
||||
|
||||
def test_fails_if_setup_hook_fails(self):
|
||||
"""A non-zero exit from a setup hook is reported as a failed server."""
|
||||
self._create_server(
|
||||
"bad",
|
||||
setup_hook={
|
||||
"command": "python",
|
||||
"args": ["-c", "import sys; sys.exit(1)"],
|
||||
},
|
||||
)
|
||||
|
||||
failed = self._run_setup()
|
||||
assert failed == ["bad"], failed
|
||||
|
||||
def test_multi_step_setup(self):
|
||||
"""Setup hooks can be a list of steps."""
|
||||
marker1 = Path(self.temp_dir) / "step1.txt"
|
||||
marker2 = Path(self.temp_dir) / "step2.txt"
|
||||
self._create_server(
|
||||
"svc",
|
||||
setup_hook=[
|
||||
{"command": "python", "args": ["-c", f"open('{marker1}', 'w').write('1')"]},
|
||||
{"command": "python", "args": ["-c", f"open('{marker2}', 'w').write('2')"]},
|
||||
],
|
||||
)
|
||||
|
||||
failed = self._run_setup()
|
||||
assert failed == [], failed
|
||||
assert marker1.read_text() == "1"
|
||||
assert marker2.read_text() == "2"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,25 @@
|
||||
"""Repository-level checks for package toolset declarations."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
STATE_TOOLS = {"export_state", "import_state"}
|
||||
|
||||
|
||||
def test_state_tools_only_appear_in_state_toolsets():
|
||||
repo_root = Path(__file__).resolve().parents[3]
|
||||
offenders: list[str] = []
|
||||
|
||||
for mcp_json in sorted((repo_root / "packages").glob("*/mcp.json")):
|
||||
package = mcp_json.parent.name
|
||||
data = json.loads(mcp_json.read_text())
|
||||
for toolset, tools in data.get("toolsets", {}).items():
|
||||
if toolset == "state":
|
||||
continue
|
||||
leaked = sorted(STATE_TOOLS.intersection(tools))
|
||||
if leaked:
|
||||
offenders.append(f"{package}:{toolset} -> {', '.join(leaked)}")
|
||||
|
||||
assert offenders == []
|
||||
@@ -0,0 +1,124 @@
|
||||
"""Tests for the viewer reverse-proxy app."""
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from starlette.testclient import TestClient
|
||||
|
||||
from mcp_proxy.viewer import _DISPLAY_NAMES, _render_shell, create_viewer_app
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shell rendering
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRenderShell:
|
||||
def test_empty_registry(self):
|
||||
html = _render_shell({})
|
||||
assert "No services available" in html
|
||||
|
||||
def test_single_service(self):
|
||||
html = _render_shell({"jira": 9000})
|
||||
assert "Services" in html
|
||||
assert "1 services" in html
|
||||
assert 'data-service="jira"' in html
|
||||
assert 'src="/__viewer__/set?app=jira"' in html
|
||||
|
||||
def test_multiple_services(self):
|
||||
html = _render_shell({"jira": 9000, "slack": 9001, "google_mail": 9002})
|
||||
assert "3 services" in html
|
||||
for name in ("jira", "slack", "google_mail"):
|
||||
assert f'data-service="{name}"' in html
|
||||
|
||||
def test_display_names_used(self):
|
||||
html = _render_shell({"jira": 9000})
|
||||
assert _DISPLAY_NAMES["jira"] in html # "Issues"
|
||||
|
||||
def test_unknown_service_gets_title_case(self):
|
||||
html = _render_shell({"my_custom_service": 9000})
|
||||
assert "My Custom Service" in html
|
||||
|
||||
def test_services_sorted(self):
|
||||
html = _render_shell({"slack": 9001, "jira": 9000})
|
||||
jira_pos = html.index('data-service="jira"')
|
||||
slack_pos = html.index('data-service="slack"')
|
||||
assert jira_pos < slack_pos
|
||||
|
||||
def test_first_service_is_default_iframe(self):
|
||||
html = _render_shell({"slack": 9001, "jira": 9000})
|
||||
# sorted → jira comes first
|
||||
assert 'src="/__viewer__/set?app=jira"' in html
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Viewer app routes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestViewerApp:
|
||||
@pytest.fixture
|
||||
def app(self):
|
||||
return create_viewer_app(
|
||||
service_registry={"test_svc": 9999, "other": 8888},
|
||||
proxy_token="secret123",
|
||||
)
|
||||
|
||||
@pytest.fixture
|
||||
def client(self, app):
|
||||
return TestClient(app, raise_server_exceptions=False)
|
||||
|
||||
def test_index_returns_html(self, client):
|
||||
resp = client.get("/")
|
||||
assert resp.status_code == 200
|
||||
assert "Services" in resp.text
|
||||
|
||||
def test_reverse_proxy_unknown_service_returns_404(self, client):
|
||||
resp = client.get("/__viewer__/set?app=nonexistent")
|
||||
assert resp.status_code == 404
|
||||
assert "Unknown service" in resp.text
|
||||
|
||||
def test_iframe_fallback_uses_server_state_when_cookie_blocked(self, app, monkeypatch):
|
||||
"""When embedded in an iframe, browsers may block the cookie. The
|
||||
viewer should fall back to server-side state when any referer is
|
||||
present (indicating a redirect, not a direct navigation)."""
|
||||
|
||||
async def mock_request(self, *args, **kwargs):
|
||||
return httpx.Response(200, text="service content")
|
||||
|
||||
monkeypatch.setattr(httpx.AsyncClient, "request", mock_request)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
|
||||
# First, select a service via /__viewer__/set to populate server state.
|
||||
resp = client.get("/__viewer__/set?app=test_svc", follow_redirects=False)
|
||||
assert resp.status_code == 307
|
||||
|
||||
# Now simulate an iframe request with NO cookie but a referer present
|
||||
# (browser stripped the path but kept the origin).
|
||||
resp = client.get(
|
||||
"/",
|
||||
headers={"referer": "http://some-outer-host/"},
|
||||
cookies={}, # no cookie — blocked by browser
|
||||
)
|
||||
# Should proxy to the service, not re-render the shell
|
||||
assert resp.status_code == 200
|
||||
assert "service content" in resp.text
|
||||
assert "Services" not in resp.text # not the shell HTML
|
||||
|
||||
def test_iframe_fallback_not_triggered_without_referer(self, client):
|
||||
"""Without a referer, the server state fallback should NOT activate —
|
||||
direct navigation should always show the shell."""
|
||||
resp = client.get("/", cookies={})
|
||||
assert resp.status_code == 200
|
||||
assert "Services" in resp.text
|
||||
|
||||
def test_reverse_proxy_connection_error_returns_502(self, app, monkeypatch):
|
||||
# Mock httpx to simulate a connection error regardless of port state
|
||||
async def mock_request(self, *args, **kwargs):
|
||||
raise httpx.ConnectError("Connection refused")
|
||||
|
||||
monkeypatch.setattr(httpx.AsyncClient, "request", mock_request)
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
# Set the service cookie first, then make a proxied request
|
||||
client.cookies.set("__viewer_service", "test_svc")
|
||||
resp = client.get("/some-page", headers={"referer": "http://testserver/"})
|
||||
assert resp.status_code == 502
|
||||
assert "Service unavailable" in resp.text
|
||||
Reference in New Issue
Block a user