diff --git a/docs/integrations/mcp.mdx b/docs/integrations/mcp.mdx index d2acca18..c80e2f78 100644 --- a/docs/integrations/mcp.mdx +++ b/docs/integrations/mcp.mdx @@ -71,6 +71,27 @@ Strix reads this file at the start of each run. There is no default file, so no server offers, or set it to a list of tool names to allow only those. + + Free-text notes for the agent about what this connection is and how you want + it used, for example "Staging analytics database, read-only, prefer aggregate + queries." When set, the notes are given to the agent at the start of the run + as a description of the connection. + + +## Choosing connections per run + +By default every connection in the file is used on each run. To narrow it for a +single run without editing the file, use either flag (both repeatable): + +```bash +strix --mcp-server github -t ... # use only the named connection(s) +strix --mcp-exclude staging-db -t ... # use everything except the named one(s) +``` + +`--mcp-server` keeps only the connections you name; `--mcp-exclude` drops the +ones you name. Connection names must be unique in the file; if two entries share +a name, the first is kept and the rest are ignored. + ## Pointing at a different file To read the config from another path instead of `~/.strix/mcp-servers.json`, either pass `--mcp-config ` on the command line: diff --git a/strix/core/runner.py b/strix/core/runner.py index 9dee9779..e90360f1 100644 --- a/strix/core/runner.py +++ b/strix/core/runner.py @@ -74,6 +74,23 @@ def _mcp_startup_summary(connections: list[ConnectedMcpServer]) -> str: return f"MCP: connected {server_count} {servers_word} ({tool_count} {tools_word}): {names}" +def _mcp_connection_notes(connections: list[ConnectedMcpServer]) -> str | None: + """A block describing the connections the user left notes on, for the agent. + + Only connections with notes are listed, so the note describes the connection + once rather than being repeated onto every tool. Returns ``None`` when no + connection has notes. + """ + noted = [(c.name, c.notes) for c in connections if c.notes] + if not noted: + return None + lines = "\n".join(f"- `{name}.*` tools: {notes}" for name, notes in noted) + return ( + "The user connected these MCP servers for this run and left notes on how " + f"to use each:\n{lines}" + ) + + def _merge_root_prompt_context( scope_context: dict[str, Any], extra_system_prompt_context: dict[str, Any] | None, @@ -323,6 +340,9 @@ async def run_strix_scan( mcp_servers = [c.server for c in connections] if connections: report(_mcp_startup_summary(connections)) + notes_block = _mcp_connection_notes(connections) + if notes_block: + root_task = f"{root_task}\n\n{notes_block}" except Exception: logger.exception("Failed to connect user MCP servers; continuing without them") diff --git a/strix/interface/cli_args.py b/strix/interface/cli_args.py index 248bcb70..5b789eda 100644 --- a/strix/interface/cli_args.py +++ b/strix/interface/cli_args.py @@ -227,6 +227,23 @@ Examples: help="Path to an MCP servers JSON file to use instead of ~/.strix/mcp-servers.json.", ) + parser.add_argument( + "--mcp-server", + dest="mcp_server", + action="append", + metavar="NAME", + help="Use only this MCP connection for the run, by its config name " + "(repeatable). Every other configured connection is skipped.", + ) + + parser.add_argument( + "--mcp-exclude", + dest="mcp_exclude", + action="append", + metavar="NAME", + help="Skip this MCP connection for the run, by its config name (repeatable).", + ) + parser.add_argument( "--max-budget", "--max-budget-usd", @@ -283,6 +300,12 @@ Examples: # setting it here makes the flag win over the default location. os.environ["STRIX_MCP_CONFIG"] = str(mcp_config_path) + # The MCP loader reads these as its per-run include/exclude selection. + if args.mcp_server: + os.environ["STRIX_MCP_ONLY"] = ",".join(args.mcp_server) + if args.mcp_exclude: + os.environ["STRIX_MCP_EXCLUDE"] = ",".join(args.mcp_exclude) + if args.update: sys.exit(0 if self_update() else 1) diff --git a/strix/tools/mcp/client.py b/strix/tools/mcp/client.py index 83aa93f4..7f79fb9f 100644 --- a/strix/tools/mcp/client.py +++ b/strix/tools/mcp/client.py @@ -13,6 +13,7 @@ and skipped, so one bad connection never fails the run. from __future__ import annotations +import contextlib import json import logging from typing import TYPE_CHECKING, Any, NamedTuple, cast @@ -54,12 +55,15 @@ class ConnectedMcpServer(NamedTuple): """One successfully connected MCP server and how many tools it registered. ``server`` is kept so the caller can clean it up when the run ends; - ``name`` and ``tool_count`` let the caller show the user a startup summary. + ``name`` and ``tool_count`` let the caller show the user a startup summary; + ``notes`` carries the connection's optional free-text description so the + caller can surface it to the agent as context about the connection. """ server: MCPServer name: str tool_count: int + notes: str | None = None def _auth_headers(config: McpConnectionConfig) -> dict[str, str]: @@ -235,12 +239,28 @@ async def connect_mcp_servers( except Exception: logger.exception("Skipping MCP connection %r", config.name) if server is not None: - await server.cleanup() # type: ignore[no-untyped-call] + with contextlib.suppress(Exception): + await server.cleanup() # type: ignore[no-untyped-call] continue + except BaseException: + # A cancellation (or other non-Exception failure) mid-connect must not + # orphan MCP subprocesses or HTTP sessions. Clean up the server being + # connected and every server already connected, then re-raise so the + # caller still stops. The runner only receives the list on a clean + # return, so on an abnormal exit this function owns the cleanup. + if server is not None: + with contextlib.suppress(Exception): + await server.cleanup() # type: ignore[no-untyped-call] + for established in connected: + with contextlib.suppress(Exception): + await established.server.cleanup() # type: ignore[no-untyped-call] + raise logger.info("Connected MCP server %r (%d tools)", config.name, len(tools)) connected.append( - ConnectedMcpServer(server=server, name=config.name, tool_count=len(tools)) + ConnectedMcpServer( + server=server, name=config.name, tool_count=len(tools), notes=config.notes + ) ) return connected diff --git a/strix/tools/mcp/config.py b/strix/tools/mcp/config.py index 15acbeaa..8df59ac2 100644 --- a/strix/tools/mcp/config.py +++ b/strix/tools/mcp/config.py @@ -59,6 +59,11 @@ class McpConnectionConfig(BaseModel): """Tool allowlist, applied after the server lists its tools. ``None`` (the default) exposes every tool the server lists; a list restricts to it.""" + notes: str | None = None + """Free-text notes for the agent describing what this connection is and how + to use it. When set, they are prefixed onto each of the connection's tool + descriptions so the agent sees them.""" + @model_validator(mode="after") def _check_transport_fields(self) -> McpConnectionConfig: if self.transport == "http" and not self.url: diff --git a/strix/tools/mcp/loader.py b/strix/tools/mcp/loader.py index fa7ccb17..c179981d 100644 --- a/strix/tools/mcp/loader.py +++ b/strix/tools/mcp/loader.py @@ -28,6 +28,10 @@ logger = logging.getLogger(__name__) _DEFAULT_PATH: Path = Path.home() / ".strix" / "mcp-servers.json" _PATH_ENV_VAR = "STRIX_MCP_CONFIG" +# Per-run selection, set by the --mcp-server / --mcp-exclude CLI flags. Each is a +# comma-separated list of connection names. +_ONLY_ENV_VAR = "STRIX_MCP_ONLY" +_EXCLUDE_ENV_VAR = "STRIX_MCP_EXCLUDE" def _resolve_path(path: Path | None) -> Path: @@ -39,6 +43,60 @@ def _resolve_path(path: Path | None) -> Path: return _DEFAULT_PATH +def _dedupe_by_name(configs: list[McpConnectionConfig]) -> list[McpConnectionConfig]: + """Keep the first connection of each name, dropping later duplicates. + + Names namespace a server's tools (``.``), so two connections + sharing a name would collide and the second's tools would be silently + rejected at registration. Drop the duplicate here, with a warning, instead. + """ + seen: set[str] = set() + unique: list[McpConnectionConfig] = [] + for config in configs: + if config.name in seen: + logger.warning( + "Ignoring MCP server %r: another connection already uses that name " + "(names must be unique because they namespace the server's tools).", + config.name, + ) + continue + seen.add(config.name) + unique.append(config) + return unique + + +def _parse_names(env_var: str) -> set[str]: + return {name.strip() for name in os.environ.get(env_var, "").split(",") if name.strip()} + + +def _apply_run_selection(configs: list[McpConnectionConfig]) -> list[McpConnectionConfig]: + """Restrict this run's connections to an optional include/exclude selection. + + ``STRIX_MCP_ONLY`` (if set) keeps only the named connections; then + ``STRIX_MCP_EXCLUDE`` drops any named connection. With neither set, every + connection is kept. + """ + only = _parse_names(_ONLY_ENV_VAR) + exclude = _parse_names(_EXCLUDE_ENV_VAR) + if not only and not exclude: + return configs + + available = {config.name for config in configs} + for name in sorted((only | exclude) - available): + logger.warning( + "MCP connection selection named %r, which is not configured; ignoring it", name + ) + + selected: list[McpConnectionConfig] = [] + for config in configs: + if only and config.name not in only: + continue + if config.name in exclude: + continue + selected.append(config) + return selected + + def load_user_mcp_configs(path: Path | None = None) -> list[McpConnectionConfig]: """Load MCP connection configs from the user's JSON file. @@ -46,7 +104,8 @@ def load_user_mcp_configs(path: Path | None = None) -> list[McpConnectionConfig] ``~/.strix/mcp-servers.json``. The file is a JSON list of server entries. A missing file returns ``[]``; an unreadable or non-list file is logged and returns ``[]``; individual entries that fail validation are logged and - skipped. + skipped. Connections sharing a name are de-duplicated (first wins), and an + optional per-run include/exclude selection is applied last. """ source = _resolve_path(path) if not source.exists(): @@ -70,4 +129,4 @@ def load_user_mcp_configs(path: Path | None = None) -> list[McpConnectionConfig] except ValidationError as exc: logger.warning("Skipping invalid MCP server entry #%d in %s: %s", index, source, exc) - return configs + return _apply_run_selection(_dedupe_by_name(configs)) diff --git a/tests/test_cli_mcp_config.py b/tests/test_cli_mcp_config.py index 12de41d4..a63f6efe 100644 --- a/tests/test_cli_mcp_config.py +++ b/tests/test_cli_mcp_config.py @@ -59,3 +59,30 @@ def test_mcp_config_flag_rejects_missing_file( cli_main.parse_arguments() assert "--mcp-config file not found" in capsys.readouterr().err + + +def test_mcp_server_flags_set_selection_env(monkeypatch: pytest.MonkeyPatch) -> None: + _stub_settings(monkeypatch) + monkeypatch.delenv("STRIX_MCP_ONLY", raising=False) + monkeypatch.delenv("STRIX_MCP_EXCLUDE", raising=False) + monkeypatch.setattr( + sys, + "argv", + [ + "strix", + "-t", + "https://test.com/", + "-n", + "--mcp-server", + "a", + "--mcp-server", + "b", + "--mcp-exclude", + "c", + ], + ) + + cli_main.parse_arguments() + + assert os.environ["STRIX_MCP_ONLY"] == "a,b" + assert os.environ["STRIX_MCP_EXCLUDE"] == "c" diff --git a/tests/test_mcp_client.py b/tests/test_mcp_client.py index 19f66426..3a7386e7 100644 --- a/tests/test_mcp_client.py +++ b/tests/test_mcp_client.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import json from typing import TYPE_CHECKING, Any @@ -12,11 +13,14 @@ from mcp.types import Tool as MCPTool from pydantic import ValidationError from strix.agents import factory +from strix.core.runner import _mcp_connection_notes from strix.tools.mcp import ( BearerAuth, + ConnectedMcpServer, McpConnectionConfig, load_user_mcp_configs, ) +from strix.tools.mcp import client as mcp_client from strix.tools.mcp.client import _auth_headers, _build_server, _register_server_tools @@ -438,3 +442,141 @@ def test_loader_reads_env_var_override(tmp_path: Path, monkeypatch: pytest.Monke configs = load_user_mcp_configs() assert [c.name for c in configs] == ["local_fs"] + + +# --- connection notes -------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_connection_notes_are_carried_on_the_connection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + server = FakeMCPServer("db", [_mcp_tool("query")]) + monkeypatch.setattr(mcp_client, "_build_server", lambda _config: server) + config = McpConnectionConfig( + name="db", + url="https://mcp.example.com", + notes="Staging analytics DB; read-only.", + allowed_tools=["query"], + ) + + connections = await mcp_client.connect_mcp_servers([config]) + + # Notes ride on the connection (surfaced once), not stapled onto each tool. + assert connections[0].notes == "Staging analytics DB; read-only." + + +def test_connection_notes_block_lists_only_noted_connections() -> None: + connections = [ + ConnectedMcpServer( + server=FakeMCPServer("db", []), name="db", tool_count=2, notes="staging, read-only" + ), + ConnectedMcpServer(server=FakeMCPServer("fs", []), name="fs", tool_count=1, notes=None), + ] + + block = _mcp_connection_notes(connections) + + assert block is not None + assert "db" in block + assert "staging, read-only" in block + # A connection without notes is not listed. + assert "fs" not in block + + +def test_connection_notes_block_is_none_without_notes() -> None: + connections = [ + ConnectedMcpServer(server=FakeMCPServer("db", []), name="db", tool_count=1, notes=None) + ] + + assert _mcp_connection_notes(connections) is None + + +# --- cancellation cleanup ---------------------------------------------------- + + +@pytest.mark.asyncio +async def test_connect_cleans_up_when_cancelled_mid_connect( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cleaned: list[str] = [] + + class _Tracking(FakeMCPServer): + def __init__(self, name: str, *, fail_connect: bool = False) -> None: + super().__init__(name, [_mcp_tool("t")]) + self._fail_connect = fail_connect + + async def connect(self) -> None: + if self._fail_connect: + raise asyncio.CancelledError + + async def cleanup(self) -> None: + cleaned.append(self._name) + + servers = {"good": _Tracking("good"), "bad": _Tracking("bad", fail_connect=True)} + monkeypatch.setattr(mcp_client, "_build_server", lambda config: servers[config.name]) + + configs = [ + McpConnectionConfig(name="good", url="https://mcp.example.com", allowed_tools=["t"]), + McpConnectionConfig(name="bad", url="https://mcp.example.com", allowed_tools=["t"]), + ] + + with pytest.raises(asyncio.CancelledError): + await mcp_client.connect_mcp_servers(configs) + + # The server being connected when cancelled, and the one already connected, + # are both cleaned up rather than orphaned. + assert cleaned == ["bad", "good"] + + +# --- duplicate names and run selection --------------------------------------- + + +def _names_file(tmp_path: Path, *names: str) -> Path: + config_file = tmp_path / "mcp-servers.json" + config_file.write_text( + json.dumps([{"name": n, "transport": "stdio", "command": "npx"} for n in names]), + encoding="utf-8", + ) + return config_file + + +def test_loader_drops_duplicate_named_connections(tmp_path: Path) -> None: + config_file = tmp_path / "mcp-servers.json" + config_file.write_text( + json.dumps( + [ + {"name": "dup", "transport": "stdio", "command": "first"}, + {"name": "dup", "transport": "stdio", "command": "second"}, + {"name": "other", "transport": "stdio", "command": "npx"}, + ] + ), + encoding="utf-8", + ) + + configs = load_user_mcp_configs(config_file) + + # Duplicate name is dropped; the first entry wins. + assert [c.name for c in configs] == ["dup", "other"] + assert configs[0].command == "first" + + +def test_loader_include_selection_keeps_only_named( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + config_file = _names_file(tmp_path, "a", "b", "c") + monkeypatch.setenv("STRIX_MCP_ONLY", "a,c") + + configs = load_user_mcp_configs(config_file) + + assert [c.name for c in configs] == ["a", "c"] + + +def test_loader_exclude_selection_drops_named( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + config_file = _names_file(tmp_path, "a", "b", "c") + monkeypatch.setenv("STRIX_MCP_EXCLUDE", "b") + + configs = load_user_mcp_configs(config_file) + + assert [c.name for c in configs] == ["a", "c"]