unified mode to conditionally disable admin tools via cli

This commit is contained in:
lda
2026-05-13 17:01:59 +07:00 Verified
parent c36614d46a
commit bfd51faa6e
6 changed files with 73 additions and 4 deletions
@@ -268,7 +268,7 @@ config_reloaded
- [x] `wf.workflow.save_deployment` - [x] `wf.workflow.save_deployment`
- [x] `wf.workflow.validate_deployment` - [x] `wf.workflow.validate_deployment`
- [x] `wf.workflow.run_deployment` - [x] `wf.workflow.run_deployment`
- [ ] Register admin tools only when admin exposure is enabled. - [x] Register admin tools only when admin exposure is enabled.
- [x] Project upstream tools using the transparent proxy path. - [x] Project upstream tools using the transparent proxy path.
- [x] Keep existing `broker` and `proxy` CLI modes during migration. - [x] Keep existing `broker` and `proxy` CLI modes during migration.
- [x] Add `unified` CLI mode. - [x] Add `unified` CLI mode.
+8
View File
@@ -53,6 +53,13 @@ def build_parser() -> argparse.ArgumentParser:
action="store_true", action="store_true",
help="Collapse a large tool catalog into a search interface, for discovery on demand", help="Collapse a large tool catalog into a search interface, for discovery on demand",
) )
serve.add_argument(
"--no-admin-tools",
dest="admin_tools",
action="store_false",
help="Hide wf.admin.* tools in unified mode.",
)
serve.set_defaults(admin_tools=True)
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.")
@@ -134,6 +141,7 @@ def main(argv: list[str] | None = None) -> int:
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, search_tools=args.search_tools,
admin_tools=args.admin_tools,
) )
return 0 return 0
+6
View File
@@ -21,6 +21,7 @@ def create_unified_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, search_tools: bool = False,
admin_tools: bool = True,
) -> FastMCP[Any]: ) -> FastMCP[Any]:
"""Create one MCP server with upstream proxy, admin, and workflow tools.""" """Create one MCP server with upstream proxy, admin, and workflow tools."""
runtime = TransparentProxyRuntime( runtime = TransparentProxyRuntime(
@@ -29,6 +30,7 @@ def create_unified_proxy_server(
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, search_tools=search_tools,
admin_tools=admin_tools,
) )
service = build_service_from_config(config) service = build_service_from_config(config)
_register_workflow_tools(runtime.server, WorkflowSurfaceHandlers(service)) _register_workflow_tools(runtime.server, WorkflowSurfaceHandlers(service))
@@ -43,6 +45,7 @@ def run_unified_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, search_tools: bool = False,
admin_tools: bool = True,
) -> None: ) -> None:
server = create_unified_proxy_server( server = create_unified_proxy_server(
config, config,
@@ -50,6 +53,7 @@ def run_unified_proxy_server(
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, search_tools=search_tools,
admin_tools=admin_tools,
) )
server.run(transport=normalize_transport(transport), show_banner=False) server.run(transport=normalize_transport(transport), show_banner=False)
@@ -61,6 +65,7 @@ def create_unified_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, search_tools: bool = False,
admin_tools: bool = True,
) -> Client[FastMCPTransport]: ) -> Client[FastMCPTransport]:
return Client( return Client(
FastMCPTransport( FastMCPTransport(
@@ -70,6 +75,7 @@ def create_unified_proxy_client(
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, search_tools=search_tools,
admin_tools=admin_tools,
) )
) )
) )
+10 -3
View File
@@ -47,6 +47,7 @@ class TransparentProxyRuntime:
resources_as_tools: bool = False, resources_as_tools: bool = False,
prompts_as_tools: bool = False, prompts_as_tools: bool = False,
search_tools: bool = False, search_tools: bool = False,
admin_tools: bool = True,
) -> None: ) -> None:
self.config = config self.config = config
self.manager = None if config_path is None else BrokerConfigManager(config_path) self.manager = None if config_path is None else BrokerConfigManager(config_path)
@@ -58,6 +59,7 @@ class TransparentProxyRuntime:
"broker capabilities with connection-qualified names." "broker capabilities with connection-qualified names."
), ),
) )
self.admin_tools = admin_tools
self.reload() self.reload()
if resources_as_tools: if resources_as_tools:
self.server.add_transform(ResourcesAsTools(self.server)) self.server.add_transform(ResourcesAsTools(self.server))
@@ -86,9 +88,10 @@ class TransparentProxyRuntime:
validate_transparent_proxy_config(config) validate_transparent_proxy_config(config)
self.server.providers[:] = [self.server.local_provider] self.server.providers[:] = [self.server.local_provider]
admin = create_proxy_admin_server(self) if self.admin_tools:
admin.add_transform(LdaNamespace(ADMIN_NAMESPACE)) admin = create_proxy_admin_server(self)
self.server.mount(admin) admin.add_transform(LdaNamespace(ADMIN_NAMESPACE))
self.server.mount(admin)
mounted_connections: list[str] = [] mounted_connections: list[str] = []
for connection in config.connections: for connection in config.connections:
@@ -167,6 +170,7 @@ 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, search_tools: bool = False,
admin_tools: bool = True,
) -> FastMCP[Any]: ) -> FastMCP[Any]:
validate_transparent_proxy_config( validate_transparent_proxy_config(
config, config,
@@ -179,6 +183,7 @@ def create_transparent_proxy_server(
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, search_tools=search_tools,
admin_tools=admin_tools,
).server ).server
@@ -189,6 +194,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, search_tools: bool = False,
admin_tools: bool = True,
) -> Client[FastMCPTransport]: ) -> Client[FastMCPTransport]:
return Client( return Client(
FastMCPTransport( FastMCPTransport(
@@ -198,6 +204,7 @@ def create_transparent_proxy_client(
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, search_tools=search_tools,
admin_tools=admin_tools,
) )
) )
) )
+19
View File
@@ -78,6 +78,25 @@ def test_build_parser_accepts_unified_mode() -> None:
assert args.command == "serve" assert args.command == "serve"
assert args.mode == "unified" assert args.mode == "unified"
assert args.admin_tools is True
def test_build_parser_accepts_no_admin_tools_flag() -> None:
parser = build_parser()
args = parser.parse_args(
[
"--config",
"wf_mcp.config.json",
"serve",
"--mode",
"unified",
"--no-admin-tools",
]
)
assert args.command == "serve"
assert args.mode == "unified"
assert args.admin_tools is False
def test_cli_connections_prints_configured_connections(capsys) -> None: def test_cli_connections_prints_configured_connections(capsys) -> None:
+29
View File
@@ -53,3 +53,32 @@ def test_unified_server_exposes_upstream_admin_and_workflow_tools() -> None:
assert _structured(artifacts_result)["nodes"] == [] assert _structured(artifacts_result)["nodes"] == []
asyncio.run(run_proxy()) asyncio.run(run_proxy())
def test_unified_server_can_hide_admin_tools() -> None:
config = BrokerConfig(
store_root=local_temp_root() / "unified_no_admin_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_unified_proxy_client(config, admin_tools=False)
async with client:
tools = await client.list_tools()
names = [tool.name for tool in tools]
assert "fixture.personal_echo_tool" in names
assert "wf.workflow.list_artifacts" in names
assert "wf.admin.list_connections" not in names
asyncio.run(run_proxy())