mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 07:17:54 +08:00
cli: start MCP servers around chat and add /mcp slash command
This commit is contained in:
+23
-1
@@ -26,7 +26,7 @@ from plyngent.cli.state import ReplState
|
||||
from plyngent.config.models import DatabaseConfig
|
||||
from plyngent.memory import MemoryStore
|
||||
from plyngent.prompting import NonInteractiveBackend, configure_prompting
|
||||
from plyngent.runtime import ProviderNotSupportedError, create_client
|
||||
from plyngent.runtime import McpManager, ProviderNotSupportedError, create_client
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Mapping
|
||||
@@ -134,6 +134,24 @@ def _setup_hooks(*, interactive: bool) -> None:
|
||||
PtyManager.set_limit_continue_hook(None)
|
||||
|
||||
|
||||
async def _start_mcp_manager(store: ConfigStore) -> McpManager | None:
|
||||
"""Start configured MCP servers before the REPL state is built.
|
||||
|
||||
Returns ``None`` when nothing is enabled. Failures degrade to a warning
|
||||
per server (the rest of the session, including LLM traffic, continues).
|
||||
"""
|
||||
manager = McpManager(store.mcp_config)
|
||||
if not manager.enabled_servers():
|
||||
return None
|
||||
await manager.ensure_started()
|
||||
for name, status, count in manager.statuses():
|
||||
if status.startswith("error"):
|
||||
click.secho(f"warning: mcp server {name!r} failed: {status}", fg="yellow", err=True)
|
||||
else:
|
||||
click.secho(f"mcp: {name} ({status}, tools={count})", fg="bright_black", err=True)
|
||||
return manager
|
||||
|
||||
|
||||
async def _bind_session(
|
||||
state: ReplState,
|
||||
*,
|
||||
@@ -285,6 +303,7 @@ async def _run_chat( # noqa: C901, PLR0912, PLR0915 — chat orchestration
|
||||
except ProviderNotSupportedError as exc:
|
||||
raise click.ClickException(str(exc)) from exc
|
||||
|
||||
mcp_manager = await _start_mcp_manager(store)
|
||||
state = ReplState(
|
||||
config=store,
|
||||
memory=memory,
|
||||
@@ -297,6 +316,7 @@ async def _run_chat( # noqa: C901, PLR0912, PLR0915 — chat orchestration
|
||||
stream_enabled=stream,
|
||||
interactive_limits=interactive,
|
||||
yolo=yolo,
|
||||
mcp_manager=mcp_manager,
|
||||
)
|
||||
# Path denylist and policy confirm live on instance.workspace (no process bag).
|
||||
state.instance_state.workspace.path_denylist = tuple(store.agent_config.path_denylist or ())
|
||||
@@ -353,6 +373,8 @@ async def _run_chat( # noqa: C901, PLR0912, PLR0915 — chat orchestration
|
||||
set_fetch_policy_confirm_hook(None, instance=state_obj.instance_state)
|
||||
clear_private_grants(instance=state_obj.instance_state)
|
||||
clear_ssrf_assume_public_cidrs()
|
||||
if state_obj.mcp_manager is not None:
|
||||
await state_obj.mcp_manager.aclose()
|
||||
await state_obj.instance_state.shutdown()
|
||||
else:
|
||||
# No ReplState: only PTY class cleanup (temps require instance allowlist).
|
||||
|
||||
@@ -1147,6 +1147,53 @@ def plugins_slash_cmd(state: ReplState, action: str | None, name: str | None) ->
|
||||
raise click.UsageError(str(exc)) from exc
|
||||
|
||||
|
||||
def _print_mcp_status(state: ReplState) -> None:
|
||||
manager = state.mcp_manager
|
||||
if manager is None:
|
||||
click.echo("mcp: no servers configured (or none enabled)")
|
||||
return
|
||||
rows = manager.statuses()
|
||||
if not rows:
|
||||
click.echo("mcp: no servers configured (or none enabled)")
|
||||
return
|
||||
for name, status, count in rows:
|
||||
click.echo(f" {name}\t{status}\ttools={count}")
|
||||
connection = next(
|
||||
(conn for conn in manager.connections() if conn.name == name),
|
||||
None,
|
||||
)
|
||||
if connection is not None and connection.error is not None:
|
||||
for line in connection.stderr_tail()[-3:]:
|
||||
click.secho(f" {line}", fg="bright_black")
|
||||
|
||||
|
||||
@slash.command("mcp")
|
||||
@click.argument("action", required=False, type=click.Choice(["list", "reconnect"]))
|
||||
@click.pass_obj
|
||||
def mcp_slash_cmd(state: ReplState, action: str | None) -> None:
|
||||
"""List MCP server statuses, or reconnect after config edits.
|
||||
|
||||
``/mcp`` or ``/mcp list`` — per-server connection status and tool count.
|
||||
``/mcp reconnect`` — re-read the config file, restart every enabled
|
||||
server, and rebuild the tool registry (adopts newly added servers too).
|
||||
"""
|
||||
act = (action or "list").lower()
|
||||
if act == "list":
|
||||
_print_mcp_status(state)
|
||||
return
|
||||
from plyngent.runtime import McpManager
|
||||
|
||||
state.config.reload()
|
||||
if state.mcp_manager is None:
|
||||
state.mcp_manager = McpManager(state.config.mcp_config)
|
||||
else:
|
||||
_await(state.mcp_manager.restart(state.config.mcp_config))
|
||||
_await(state.mcp_manager.ensure_started())
|
||||
_slash_rebuild_tools_if_on(state)
|
||||
click.echo(f"mcp reconnected; tools={'on' if state.tools_enabled else 'off'}")
|
||||
_print_mcp_status(state)
|
||||
|
||||
|
||||
@slash.command("yolo")
|
||||
@click.argument("mode", required=False, type=YOLO_MODE, metavar="[on|off|once]")
|
||||
@click.pass_obj
|
||||
|
||||
@@ -32,6 +32,7 @@ if TYPE_CHECKING:
|
||||
from plyngent.config.store import ConfigStore
|
||||
from plyngent.memory import MemoryStore
|
||||
from plyngent.memory.database.schema import Session as SessionRow
|
||||
from plyngent.runtime.mcp_client import McpManager
|
||||
|
||||
type YoloMode = Literal["off", "on", "once"]
|
||||
|
||||
@@ -65,6 +66,8 @@ class ReplState:
|
||||
todo_stack: TodoStack = field(default_factory=TodoStack)
|
||||
instance_state: InstanceState = field(default_factory=InstanceState)
|
||||
session_state: SessionState = field(default_factory=SessionState)
|
||||
# MCP connections started by the host app before this state was built.
|
||||
mcp_manager: McpManager | None = None
|
||||
_todo_persist_tasks: set[object] = field(default_factory=set, init=False, repr=False)
|
||||
# Session ids for Tab complete (updated when listing/creating/resuming).
|
||||
_session_id_cache: list[int] = field(default_factory=list, init=False, repr=False)
|
||||
@@ -215,6 +218,11 @@ class ReplState:
|
||||
plugins_cfg.enable,
|
||||
disable=plugins_cfg.disable,
|
||||
)
|
||||
if self.mcp_manager is not None:
|
||||
from plyngent.tools.mcp import register_mcp_tools
|
||||
|
||||
# Connections were started by the host; registration is sync-only.
|
||||
_ = register_mcp_tools(self.mcp_manager, self.config.mcp_config)
|
||||
# Local surface: builtins + allowlisted plugins (import registers into catalog).
|
||||
tools = catalog.select(surface="local")
|
||||
yolo = self.effective_yolo() != "off"
|
||||
|
||||
@@ -0,0 +1,122 @@
|
||||
"""CLI slash ``/mcp`` command and MCP wiring in ``ReplState``."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pytest
|
||||
import tomlkit
|
||||
|
||||
from plyngent.cli.slash import handle_slash
|
||||
from plyngent.cli.state import ReplState
|
||||
from plyngent.config.models import DatabaseConfig, McpConfig, McpServerConfig, OpenAIProvider
|
||||
from plyngent.config.store import ConfigStore
|
||||
from plyngent.memory import MemoryStore
|
||||
from plyngent.runtime.mcp_client import McpManager
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
_FAKE_SERVER = Path(__file__).parent.parent / "test_runtime" / "mcp_fake_server.py"
|
||||
|
||||
|
||||
def _server_config(**overrides: object) -> McpServerConfig:
|
||||
base: dict[str, object] = {
|
||||
"command": sys.executable,
|
||||
"args": [str(_FAKE_SERVER)],
|
||||
"timeout": 5.0,
|
||||
}
|
||||
return McpServerConfig(**{**base, **overrides}) # type: ignore[arg-type]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def state(tmp_path: Path) -> AsyncIterator[ReplState]:
|
||||
memory = await MemoryStore.open(DatabaseConfig())
|
||||
provider = OpenAIProvider(access_key_or_token="sk-test")
|
||||
config = ConfigStore(path=tmp_path / "plyngent.toml", document=tomlkit.document())
|
||||
config.providers = {"local": provider}
|
||||
st = ReplState(
|
||||
config=config,
|
||||
memory=memory,
|
||||
workspace=tmp_path,
|
||||
provider_name="local",
|
||||
provider=provider,
|
||||
model="gpt-test",
|
||||
tools_enabled=False,
|
||||
)
|
||||
await st.new_session("t")
|
||||
yield st
|
||||
await memory.close()
|
||||
|
||||
|
||||
async def test_slash_mcp_list_without_manager(state: ReplState, capsys: pytest.CaptureFixture[str]) -> None:
|
||||
assert await handle_slash(state, "/mcp") is True
|
||||
assert "no servers configured" in capsys.readouterr().out
|
||||
assert await handle_slash(state, "/mcp list") is True
|
||||
assert "no servers configured" in capsys.readouterr().out
|
||||
|
||||
|
||||
async def test_slash_mcp_list_statuses(state: ReplState, capsys: pytest.CaptureFixture[str]) -> None:
|
||||
manager = McpManager(McpConfig(servers={"fake": _server_config()}, disable=[]))
|
||||
await manager.ensure_started()
|
||||
state.mcp_manager = manager
|
||||
try:
|
||||
assert await handle_slash(state, "/mcp list") is True
|
||||
out = capsys.readouterr().out
|
||||
assert "fake" in out
|
||||
assert "connected" in out
|
||||
assert "tools=3" in out
|
||||
finally:
|
||||
await manager.aclose()
|
||||
|
||||
|
||||
async def test_slash_mcp_reconnect_connects_new_config(
|
||||
state: ReplState,
|
||||
capsys: pytest.CaptureFixture[str],
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
# The fixture's config file does not exist yet; write one with an MCP server.
|
||||
payload = (
|
||||
"[mcp.servers.fake]\n"
|
||||
f"command = {json.dumps(sys.executable)}\n"
|
||||
f"args = [{json.dumps(str(_FAKE_SERVER))}]\n"
|
||||
"timeout = 5.0\n"
|
||||
)
|
||||
_ = (tmp_path / "plyngent.toml").write_text(payload, encoding="utf-8")
|
||||
assert await handle_slash(state, "/mcp reconnect") is True
|
||||
try:
|
||||
assert state.mcp_manager is not None
|
||||
statuses = state.mcp_manager.statuses()
|
||||
assert [name for name, _, _ in statuses] == ["fake"]
|
||||
assert statuses[0][1] == "connected"
|
||||
assert statuses[0][2] == 3
|
||||
out = capsys.readouterr().out
|
||||
assert "tools=off" in out
|
||||
assert "fake" in out
|
||||
finally:
|
||||
if state.mcp_manager is not None:
|
||||
await state.mcp_manager.aclose()
|
||||
|
||||
|
||||
async def test_mcp_tools_registered_on_tools_rebuild(state: ReplState) -> None:
|
||||
manager = McpManager(McpConfig(servers={"fake": _server_config()}, disable=[]))
|
||||
await manager.ensure_started()
|
||||
state.mcp_manager = manager
|
||||
try:
|
||||
from plyngent.tools.catalog import catalog_scope
|
||||
|
||||
state.tools_enabled = True
|
||||
with catalog_scope(empty=True):
|
||||
state.rebuild_client()
|
||||
registry = state.agent.tools
|
||||
assert registry is not None
|
||||
assert registry.get("mcp__fake__echo") is not None
|
||||
assert registry.get("mcp__fake__fail") is not None
|
||||
# Handler proxies to the live server subprocess.
|
||||
out = await registry.execute("mcp__fake__echo", '{"text": "hi"}')
|
||||
assert out == "echo: hi"
|
||||
finally:
|
||||
await manager.aclose()
|
||||
Reference in New Issue
Block a user