unified mode to conditionally disable admin tools via cli
This commit is contained in:
@@ -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.
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|||||||
@@ -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,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -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,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
@@ -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())
|
||||||
|
|||||||
Reference in New Issue
Block a user