optional search + config validation

This commit is contained in:
lda
2026-04-30 05:58:59 +07:00 Verified
parent fc76d6f8fd
commit ef2f179110
8 changed files with 292 additions and 9 deletions
+12
View File
@@ -19,6 +19,12 @@ from .capabilities import (
) )
from .catalog import CombinedCatalog from .catalog import CombinedCatalog
from .connections import ConnectionRegistry, parse_connection_id, qualify_node_name from .connections import ConnectionRegistry, parse_connection_id, qualify_node_name
from .config_models import (
BrokerConfigFile,
ConnectionConfigFile,
HttpConnectionMetadata,
StdioConnectionMetadata,
)
from .discovery import ( from .discovery import (
DiscoveredConnectionCapabilities, DiscoveredConnectionCapabilities,
discover_connection_capabilities, discover_connection_capabilities,
@@ -39,6 +45,7 @@ from .store import FileStore, Store
from .transparent_proxy import ( from .transparent_proxy import (
broker_config_to_fastmcp_config, broker_config_to_fastmcp_config,
connection_to_fastmcp_server_config, connection_to_fastmcp_server_config,
create_proxy_admin_server,
create_transparent_proxy_client, create_transparent_proxy_client,
create_transparent_proxy_server, create_transparent_proxy_server,
) )
@@ -54,23 +61,28 @@ __all__ = [
"CatalogSnapshot", "CatalogSnapshot",
"CombinedCatalog", "CombinedCatalog",
"ConnectionConfig", "ConnectionConfig",
"ConnectionConfigFile",
"ConnectionRegistry", "ConnectionRegistry",
"DiscoveredConnectionCapabilities", "DiscoveredConnectionCapabilities",
"DiscoveredPrompt", "DiscoveredPrompt",
"DiscoveredResource", "DiscoveredResource",
"DiscoveredTool", "DiscoveredTool",
"FileStore", "FileStore",
"BrokerConfigFile",
"HttpConnectionMetadata",
"McpEvent", "McpEvent",
"McpSdkAdapter", "McpSdkAdapter",
"ProxyConfigError", "ProxyConfigError",
"RawWorkflowPlan", "RawWorkflowPlan",
"Store", "Store",
"StdioConnectionMetadata",
"ToolCallResult", "ToolCallResult",
"WfMcpService", "WfMcpService",
"build_service_from_config", "build_service_from_config",
"broker_config_to_fastmcp_config", "broker_config_to_fastmcp_config",
"connection_to_fastmcp_server_config", "connection_to_fastmcp_server_config",
"create_broker_server", "create_broker_server",
"create_proxy_admin_server",
"create_transparent_proxy_client", "create_transparent_proxy_client",
"create_transparent_proxy_server", "create_transparent_proxy_server",
"discover_connection_capabilities", "discover_connection_capabilities",
+5 -8
View File
@@ -8,9 +8,10 @@ from typing import Any, Literal
from mcp.server.fastmcp import FastMCP from mcp.server.fastmcp import FastMCP
from .config_models import BrokerConfigFile
from .error_info import error_payload from .error_info import error_payload
from .mcp_sdk_adapter import McpSdkAdapter from .mcp_sdk_adapter import McpSdkAdapter
from .models import BrokerConfig, ConnectionConfig from .models import BrokerConfig
from .service import WfMcpService from .service import WfMcpService
from .store import FileStore from .store import FileStore
from .transparent_proxy import create_transparent_proxy_server from .transparent_proxy import create_transparent_proxy_server
@@ -19,13 +20,7 @@ from .transparent_proxy import create_transparent_proxy_server
def load_broker_config(path: str | Path) -> BrokerConfig: def load_broker_config(path: str | Path) -> BrokerConfig:
config_path = Path(path) config_path = Path(path)
data = json.loads(config_path.read_text(encoding="utf-8")) data = json.loads(config_path.read_text(encoding="utf-8"))
store_root_raw = data.get("store_root", ".wf_mcp_store") return BrokerConfigFile.model_validate(data).to_runtime(config_path=config_path)
store_root = Path(store_root_raw)
if not store_root.is_absolute():
store_root = (config_path.parent / store_root).resolve()
connections = [ConnectionConfig(**item) for item in data.get("connections", [])]
return BrokerConfig(store_root=store_root, connections=connections)
def build_service_from_config(config: BrokerConfig) -> WfMcpService: def build_service_from_config(config: BrokerConfig) -> WfMcpService:
@@ -228,12 +223,14 @@ def run_transparent_proxy_server(
*, *,
resources_as_tools: bool = False, resources_as_tools: bool = False,
prompts_as_tools: bool = False, prompts_as_tools: bool = False,
search_tools: bool = False,
) -> None: ) -> None:
config = load_broker_config(config_path) config = load_broker_config(config_path)
server = create_transparent_proxy_server( server = create_transparent_proxy_server(
config, config,
resources_as_tools=resources_as_tools, resources_as_tools=resources_as_tools,
prompts_as_tools=prompts_as_tools, prompts_as_tools=prompts_as_tools,
search_tools=search_tools,
) )
server.run(transport=normalize_transport(transport), show_banner=False) server.run(transport=normalize_transport(transport), show_banner=False)
+6
View File
@@ -47,6 +47,11 @@ def build_parser() -> argparse.ArgumentParser:
action="store_true", action="store_true",
help="Expose proxied prompts through list_prompts/get_prompt tools.", help="Expose proxied prompts through list_prompts/get_prompt tools.",
) )
serve.add_argument(
"--search-tools",
action="store_true",
help="Collapse a large tool catalog into a search interface, for discovery on demand",
)
subparsers.add_parser("connections", help="List configured connections.") subparsers.add_parser("connections", help="List configured connections.")
subparsers.add_parser("status", help="Show connection status and snapshot counts.") subparsers.add_parser("status", help="Show connection status and snapshot counts.")
@@ -115,6 +120,7 @@ def main(argv: list[str] | None = None) -> int:
args.transport, args.transport,
resources_as_tools=args.resources_as_tools, resources_as_tools=args.resources_as_tools,
prompts_as_tools=args.prompts_as_tools, prompts_as_tools=args.prompts_as_tools,
search_tools=args.search_tools,
) )
else: else:
run_broker_server(args.config, args.transport) run_broker_server(args.config, args.transport)
+84
View File
@@ -0,0 +1,84 @@
from __future__ import annotations
from pathlib import Path
from typing import Annotated, Any, Literal
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, field_validator
from .models import BrokerConfig, ConnectionConfig
class StdioConnectionMetadata(BaseModel):
model_config = ConfigDict(extra="allow")
transport: Literal["stdio"] = "stdio"
command: str | None = None
args: list[str] = Field(default_factory=list)
env: dict[str, str] = Field(default_factory=dict)
cwd: str | None = None
description: str | None = None
class HttpConnectionMetadata(BaseModel):
model_config = ConfigDict(extra="allow")
transport: Literal["http", "streamable-http", "streamable_http", "sse"]
url: str | None = None
headers: dict[str, str] = Field(default_factory=dict)
description: str | None = None
TypedConnectionMetadata = Annotated[
StdioConnectionMetadata | HttpConnectionMetadata,
Field(discriminator="transport"),
]
_METADATA_ADAPTER = TypeAdapter(TypedConnectionMetadata)
class ConnectionConfigFile(BaseModel):
model_config = ConfigDict(extra="forbid")
id: str
server: str
account: str
enabled: bool = True
metadata: dict[str, Any] = Field(default_factory=dict)
@field_validator("metadata", mode="before")
@classmethod
def validate_metadata(cls, value: object) -> dict[str, Any]:
if value is None:
return {}
if not isinstance(value, dict):
raise ValueError("metadata must be an object")
if not value:
return {}
if "transport" not in value:
value = {**value, "transport": "stdio"}
metadata = _METADATA_ADAPTER.validate_python(value)
return metadata.model_dump(exclude_none=True)
def to_runtime(self) -> ConnectionConfig:
return ConnectionConfig(
id=self.id,
server=self.server,
account=self.account,
enabled=self.enabled,
metadata=self.metadata,
)
class BrokerConfigFile(BaseModel):
model_config = ConfigDict(extra="forbid")
store_root: Path = Path(".wf_mcp_store")
connections: list[ConnectionConfigFile] = Field(default_factory=list)
def to_runtime(self, *, config_path: Path) -> BrokerConfig:
store_root = self.store_root
if not store_root.is_absolute():
store_root = (config_path.parent / store_root).resolve()
return BrokerConfig(
store_root=store_root,
connections=[connection.to_runtime() for connection in self.connections],
)
+3
View File
@@ -7,6 +7,7 @@ from .models import BrokerConfig, ConnectionConfig
_NAMESPACE_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9.-]*$") _NAMESPACE_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9.-]*$")
_SUPPORTED_TRANSPORTS = {"stdio", "http", "streamable-http", "streamable_http", "sse"} _SUPPORTED_TRANSPORTS = {"stdio", "http", "streamable-http", "streamable_http", "sse"}
_RESERVED_CONNECTION_IDS = {"wf.mcp"}
class ProxyConfigError(ValueError): class ProxyConfigError(ValueError):
@@ -50,6 +51,8 @@ def _validate_connection_ids(
if not connection_id: if not connection_id:
errors.append("connection id must not be empty") errors.append("connection id must not be empty")
continue continue
if connection_id in _RESERVED_CONNECTION_IDS:
errors.append(f"connection id {connection_id!r} is reserved by wf-mcp")
if "_" in connection_id: if "_" in connection_id:
errors.append( errors.append(
f"connection id {connection_id!r} must not contain '_' because " f"connection id {connection_id!r} must not contain '_' because "
+53 -1
View File
@@ -1,5 +1,6 @@
from __future__ import annotations from __future__ import annotations
from dataclasses import asdict
from typing import Any from typing import Any
from fastmcp import FastMCP from fastmcp import FastMCP
@@ -9,10 +10,52 @@ from fastmcp.client.transports.memory import FastMCPTransport
from fastmcp.mcp_config import MCPConfig from fastmcp.mcp_config import MCPConfig
from fastmcp.server import create_proxy from fastmcp.server import create_proxy
from fastmcp.server.transforms import Namespace, PromptsAsTools, ResourcesAsTools from fastmcp.server.transforms import Namespace, PromptsAsTools, ResourcesAsTools
from fastmcp.server.transforms.search import BM25SearchTransform
from .models import BrokerConfig, ConnectionConfig from .models import BrokerConfig, ConnectionConfig
from .proxy_validation import validate_transparent_proxy_config from .proxy_validation import validate_transparent_proxy_config
_ADMIN_NAMESPACE = "wf.mcp"
_ADMIN_TOOL_NAMES = [
f"{_ADMIN_NAMESPACE}_list_connections",
f"{_ADMIN_NAMESPACE}_get_connection_statuses",
]
def create_proxy_admin_server(config: BrokerConfig) -> FastMCP[Any]:
admin = FastMCP(
"wf-mcp-admin",
instructions="Administrative tools for this wf-mcp proxy instance.",
)
@admin.tool()
async def list_connections() -> list[dict[str, Any]]:
return [
asdict(connection)
for connection in sorted(
config.connections,
key=lambda connection: connection.id,
)
]
@admin.tool()
async def get_connection_statuses() -> list[dict[str, Any]]:
return [
{
"connection_id": connection.id,
"server": connection.server,
"account": connection.account,
"enabled": connection.enabled,
"transport": connection.metadata.get("transport"),
}
for connection in sorted(
config.connections,
key=lambda connection: connection.id,
)
]
return admin
def connection_to_fastmcp_server_config( def connection_to_fastmcp_server_config(
connection: ConnectionConfig, connection: ConnectionConfig,
@@ -58,11 +101,13 @@ def create_transparent_proxy_server(
*, *,
resources_as_tools: bool = False, resources_as_tools: bool = False,
prompts_as_tools: bool = False, prompts_as_tools: bool = False,
search_tools: bool = False,
) -> FastMCP[Any]: ) -> FastMCP[Any]:
validate_transparent_proxy_config( validate_transparent_proxy_config(
config, config,
resources_as_tools=resources_as_tools, resources_as_tools=resources_as_tools,
prompts_as_tools=prompts_as_tools, prompts_as_tools=prompts_as_tools,
# not yet idk codex help
) )
root = FastMCP( root = FastMCP(
"wf-mcp-transparent-proxy", "wf-mcp-transparent-proxy",
@@ -73,6 +118,10 @@ def create_transparent_proxy_server(
), ),
) )
admin = create_proxy_admin_server(config)
admin.add_transform(Namespace(_ADMIN_NAMESPACE))
root.mount(admin)
for connection in config.connections: for connection in config.connections:
if not connection.enabled: if not connection.enabled:
continue continue
@@ -89,7 +138,8 @@ def create_transparent_proxy_server(
root.add_transform(ResourcesAsTools(root)) root.add_transform(ResourcesAsTools(root))
if prompts_as_tools: if prompts_as_tools:
root.add_transform(PromptsAsTools(root)) root.add_transform(PromptsAsTools(root))
if search_tools:
root.add_transform(BM25SearchTransform(always_visible=_ADMIN_TOOL_NAMES))
return root return root
@@ -98,6 +148,7 @@ def create_transparent_proxy_client(
*, *,
resources_as_tools: bool = False, resources_as_tools: bool = False,
prompts_as_tools: bool = False, prompts_as_tools: bool = False,
search_tools: bool = False,
) -> Client[FastMCPTransport]: ) -> Client[FastMCPTransport]:
return Client( return Client(
FastMCPTransport( FastMCPTransport(
@@ -105,6 +156,7 @@ def create_transparent_proxy_client(
config, config,
resources_as_tools=resources_as_tools, resources_as_tools=resources_as_tools,
prompts_as_tools=prompts_as_tools, prompts_as_tools=prompts_as_tools,
search_tools=search_tools,
) )
) )
) )
+67
View File
@@ -3,7 +3,11 @@ from __future__ import annotations
import json import json
from pathlib import Path from pathlib import Path
import pytest
from pydantic import ValidationError
from wf_mcp.cli import build_parser, main from wf_mcp.cli import build_parser, main
from wf_mcp.broker_server import load_broker_config
from tests.test_wf_mcp_support import local_temp_root from tests.test_wf_mcp_support import local_temp_root
@@ -37,6 +41,7 @@ def test_build_parser_accepts_serve_transport() -> None:
assert args.mode == "proxy" assert args.mode == "proxy"
assert args.resources_as_tools is False assert args.resources_as_tools is False
assert args.prompts_as_tools is False assert args.prompts_as_tools is False
assert args.search_tools is False
def test_build_parser_accepts_proxy_compatibility_flags() -> None: def test_build_parser_accepts_proxy_compatibility_flags() -> None:
@@ -48,6 +53,7 @@ def test_build_parser_accepts_proxy_compatibility_flags() -> None:
"serve", "serve",
"--resources-as-tools", "--resources-as-tools",
"--prompts-as-tools", "--prompts-as-tools",
"--search-tools",
] ]
) )
@@ -55,6 +61,7 @@ def test_build_parser_accepts_proxy_compatibility_flags() -> None:
assert args.mode == "proxy" assert args.mode == "proxy"
assert args.resources_as_tools is True assert args.resources_as_tools is True
assert args.prompts_as_tools is True assert args.prompts_as_tools is True
assert args.search_tools is True
def test_cli_connections_prints_configured_connections(capsys) -> None: def test_cli_connections_prints_configured_connections(capsys) -> None:
@@ -114,3 +121,63 @@ def test_cli_status_prints_connection_statuses(capsys) -> None:
"prompt_count": 0, "prompt_count": 0,
} }
] ]
def test_load_broker_config_normalizes_typed_stdio_metadata() -> None:
tmp_path = local_temp_root() / "cli_typed_stdio_config_test"
tmp_path.mkdir(parents=True, exist_ok=True)
config_path = tmp_path / "wf_mcp.config.json"
config_path.write_text(
json.dumps(
{
"store_root": ".wf_mcp_store",
"connections": [
{
"id": "demo.personal",
"server": "demo",
"account": "personal",
"metadata": {
"command": "python",
"args": ["server.py"],
"env": {"TOKEN": "secret"},
},
}
],
}
),
encoding="utf-8",
)
config = load_broker_config(config_path)
assert config.store_root == (tmp_path / ".wf_mcp_store").resolve()
assert config.connections[0].metadata == {
"transport": "stdio",
"command": "python",
"args": ["server.py"],
"env": {"TOKEN": "secret"},
}
def test_load_broker_config_rejects_bad_metadata_shape() -> None:
tmp_path = local_temp_root() / "cli_bad_config_test"
tmp_path.mkdir(parents=True, exist_ok=True)
config_path = tmp_path / "wf_mcp.config.json"
config_path.write_text(
json.dumps(
{
"connections": [
{
"id": "demo.personal",
"server": "demo",
"account": "personal",
"metadata": {"transport": "stdio", "args": "server.py"},
}
],
}
),
encoding="utf-8",
)
with pytest.raises(ValidationError):
load_broker_config(config_path)
+62
View File
@@ -38,8 +38,27 @@ def test_transparent_proxy_lists_and_calls_upstream_tools() -> None:
async with client: async with client:
tools = await client.list_tools() tools = await client.list_tools()
names = [tool.name for tool in tools] names = [tool.name for tool in tools]
assert "wf.mcp_list_connections" in names
assert "wf.mcp_get_connection_statuses" in names
assert "fixture.personal_echo_tool" in names assert "fixture.personal_echo_tool" in names
connections_result = await client.call_tool("wf.mcp_list_connections")
assert connections_result.structured_content == {
"result": [
{
"id": "fixture.personal",
"server": "fixture",
"account": "personal",
"enabled": True,
"metadata": {
"transport": "stdio",
"command": sys.executable,
"args": [fixture_server_path()],
},
}
]
}
result = await client.call_tool( result = await client.call_tool(
"fixture.personal_echo_tool", "fixture.personal_echo_tool",
{"text": "hello"}, {"text": "hello"},
@@ -77,6 +96,12 @@ def test_transparent_proxy_rejects_invalid_connection_config() -> None:
account="http", account="http",
metadata={"transport": "http"}, metadata={"transport": "http"},
), ),
ConnectionConfig(
id="wf.mcp",
server="wf",
account="mcp",
metadata={"transport": "stdio", "command": sys.executable},
),
], ],
) )
@@ -89,6 +114,7 @@ def test_transparent_proxy_rejects_invalid_connection_config() -> None:
assert "fixture.personal: unsupported MCP transport 'websocket'" in message assert "fixture.personal: unsupported MCP transport 'websocket'" in message
assert "connection id 'bad_scope.personal' must not contain '_'" in message assert "connection id 'bad_scope.personal' must not contain '_'" in message
assert "fixture.http: http transport requires metadata.url" in message assert "fixture.http: http transport requires metadata.url" in message
assert "connection id 'wf.mcp' is reserved by wf-mcp" in message
def test_transparent_proxy_can_expose_resources_and_prompts_as_tools() -> None: def test_transparent_proxy_can_expose_resources_and_prompts_as_tools() -> None:
@@ -123,3 +149,39 @@ def test_transparent_proxy_can_expose_resources_and_prompts_as_tools() -> None:
assert "get_prompt" in names assert "get_prompt" in names
asyncio.run(run_proxy()) asyncio.run(run_proxy())
def test_transparent_proxy_can_collapse_upstream_tools_behind_search() -> None:
config = BrokerConfig(
store_root=local_temp_root() / "transparent_proxy_search_store",
connections=[
ConnectionConfig(
id="fixture.personal",
server="fixture",
account="personal",
metadata={
"transport": "stdio",
"command": sys.executable,
"args": [fixture_server_path()],
},
)
],
)
async def run_proxy() -> None:
client = create_transparent_proxy_client(config, search_tools=True)
async with client:
tools = await client.list_tools()
names = [tool.name for tool in tools]
assert "search_tools" in names
assert "wf.mcp_list_connections" in names
assert "wf.mcp_get_connection_statuses" in names
assert "fixture.personal_echo_tool" not in names
search_result = await client.call_tool(
"search_tools",
{"query": "echo text back"},
)
assert "fixture.personal_echo_tool" in str(search_result)
asyncio.run(run_proxy())