update some proxy tool query methods

This commit is contained in:
lda
2026-04-30 07:35:18 +07:00 Verified
parent 2a1b695213
commit c4a68135fd
5 changed files with 247 additions and 18 deletions
+8
View File
@@ -1,6 +1,9 @@
from __future__ import annotations from __future__ import annotations
from dataclasses import dataclass from dataclasses import dataclass
from re import L
from fastmcp.server.transforms import Namespace
ADMIN_NAMESPACE = "wf.mcp" ADMIN_NAMESPACE = "wf.mcp"
@@ -40,3 +43,8 @@ def parse_namespaced_tool_name(
def is_admin_tool_name(proxy_name: str) -> bool: def is_admin_tool_name(proxy_name: str) -> bool:
return proxy_name.startswith(f"{ADMIN_NAMESPACE}_") return proxy_name.startswith(f"{ADMIN_NAMESPACE}_")
class LdaNamespace(Namespace):
def __init__(self, prefix: str) -> None:
super().__init__(prefix)
self._name_prefix = f"{prefix}." # some good stuff
+44
View File
@@ -0,0 +1,44 @@
from __future__ import annotations
import base64
import json
from typing import TypeVar
T = TypeVar("T")
def parse_cursor(cursor: str | None) -> int:
if cursor is None:
return 0
try:
payload = json.loads(base64.urlsafe_b64decode(cursor.encode()).decode())
except Exception as exc:
raise ValueError("invalid cursor") from exc
start = payload.get("start")
if not isinstance(start, int) or start < 0:
raise ValueError("invalid cursor")
return start
def make_cursor(start: int) -> str:
payload = json.dumps({"start": start}, separators=(",", ":")).encode()
return base64.urlsafe_b64encode(payload).decode()
def clamp_limit(limit: int, *, default: int = 50, maximum: int = 200) -> int:
if limit <= 0:
return default
return min(limit, maximum)
def paginate_items(
items: list[T],
*,
cursor: str | None,
limit: int,
) -> tuple[list[T], str | None]:
page_limit = clamp_limit(limit)
start = parse_cursor(cursor)
end = start + page_limit
next_cursor = make_cursor(end) if end < len(items) else None
return items[start:end], next_cursor
+111 -10
View File
@@ -16,6 +16,7 @@ from fastmcp.server.transforms.search import BM25SearchTransform
from .config_manager import BrokerConfigManager, ConfigMutationError from .config_manager import BrokerConfigManager, ConfigMutationError
from .models import BrokerConfig, ConnectionConfig from .models import BrokerConfig, ConnectionConfig
from .names import ADMIN_NAMESPACE, is_admin_tool_name, parse_namespaced_tool_name from .names import ADMIN_NAMESPACE, is_admin_tool_name, parse_namespaced_tool_name
from .pagination import paginate_items
from .proxy_validation import validate_transparent_proxy_config from .proxy_validation import validate_transparent_proxy_config
_ADMIN_TOOL_NAMES = [ _ADMIN_TOOL_NAMES = [
@@ -24,6 +25,7 @@ _ADMIN_TOOL_NAMES = [
f"{ADMIN_NAMESPACE}_get_config", f"{ADMIN_NAMESPACE}_get_config",
f"{ADMIN_NAMESPACE}_reload_config", f"{ADMIN_NAMESPACE}_reload_config",
f"{ADMIN_NAMESPACE}_list_proxy_tools", f"{ADMIN_NAMESPACE}_list_proxy_tools",
f"{ADMIN_NAMESPACE}_get_proxy_tool",
f"{ADMIN_NAMESPACE}_add_connection", f"{ADMIN_NAMESPACE}_add_connection",
f"{ADMIN_NAMESPACE}_update_connection", f"{ADMIN_NAMESPACE}_update_connection",
f"{ADMIN_NAMESPACE}_enable_connection", f"{ADMIN_NAMESPACE}_enable_connection",
@@ -105,6 +107,9 @@ class TransparentProxyRuntime:
} }
async def list_proxy_tools(self) -> list[dict[str, Any]]: async def list_proxy_tools(self) -> list[dict[str, Any]]:
return await self._list_proxy_tools()
async def _list_proxy_tools(self) -> list[dict[str, Any]]:
config = self.current_config() config = self.current_config()
connection_ids = { connection_ids = {
connection.id for connection in config.connections if connection.enabled connection.id for connection in config.connections if connection.enabled
@@ -118,17 +123,99 @@ class TransparentProxyRuntime:
if parsed is None: if parsed is None:
continue continue
result.append( result.append(
{ _proxy_tool_payload(
"proxy_name": parsed.proxy_name, proxy_name=parsed.proxy_name,
"connection_id": parsed.connection_id, connection_id=parsed.connection_id,
"local_name": parsed.local_name, local_name=parsed.local_name,
"title": tool.title, tool=tool,
"description": tool.description, include_schema=False,
"enabled": True, )
}
) )
return sorted(result, key=lambda item: item["proxy_name"]) return sorted(result, key=lambda item: item["proxy_name"])
async def list_proxy_tools_page(
self,
*,
connection_id: str | None = None,
query: str | None = None,
limit: int = 50,
cursor: str | None = None,
) -> dict[str, Any]:
tools = await self._list_proxy_tools()
if connection_id is not None:
tools = [
tool for tool in tools if tool["connection_id"] == connection_id
]
if query:
needle = query.casefold()
tools = [
tool
for tool in tools
if needle
in " ".join(
str(tool.get(key, ""))
for key in (
"proxy_name",
"connection_id",
"local_name",
"title",
"description",
)
).casefold()
]
page, next_cursor = paginate_items(tools, cursor=cursor, limit=limit)
return {
"tools": page,
"nextCursor": next_cursor,
"total": len(tools),
}
async def get_proxy_tool(self, proxy_name: str) -> dict[str, Any]:
config = self.current_config()
connection_ids = {
connection.id for connection in config.connections if connection.enabled
}
parsed = parse_namespaced_tool_name(proxy_name, connection_ids)
if parsed is None:
raise KeyError(proxy_name)
tools = await self.server.list_tools()
for tool in tools:
if tool.name == proxy_name:
return _proxy_tool_payload(
proxy_name=parsed.proxy_name,
connection_id=parsed.connection_id,
local_name=parsed.local_name,
tool=tool,
include_schema=True,
)
raise KeyError(proxy_name)
def _proxy_tool_payload(
*,
proxy_name: str,
connection_id: str,
local_name: str,
tool: Any,
include_schema: bool,
) -> dict[str, Any]:
payload = {
"proxy_name": proxy_name,
"connection_id": connection_id,
"local_name": local_name,
"title": getattr(tool, "title", None),
"description": getattr(tool, "description", None),
"enabled": True,
}
if include_schema:
payload["input_schema"] = getattr(
tool,
"input_schema",
getattr(tool, "parameters", None),
)
payload["output_schema"] = getattr(tool, "output_schema", None)
return payload
def create_proxy_admin_server( def create_proxy_admin_server(
runtime: TransparentProxyRuntime, runtime: TransparentProxyRuntime,
@@ -179,8 +266,22 @@ def create_proxy_admin_server(
return runtime.reload() return runtime.reload()
@admin.tool() @admin.tool()
async def list_proxy_tools() -> list[dict[str, Any]]: async def list_proxy_tools(
return await runtime.list_proxy_tools() connection_id: str | None = None,
query: str | None = None,
limit: int = 50,
cursor: str | None = None,
) -> dict[str, Any]:
return await runtime.list_proxy_tools_page(
connection_id=connection_id,
query=query,
limit=limit,
cursor=cursor,
)
@admin.tool()
async def get_proxy_tool(proxy_name: str) -> dict[str, Any]:
return await runtime.get_proxy_tool(proxy_name)
@admin.tool() @admin.tool()
async def add_connection( async def add_connection(
+80 -1
View File
@@ -50,6 +50,7 @@ def test_transparent_proxy_lists_and_calls_upstream_tools() -> None:
assert "wf.mcp_list_connections" in names assert "wf.mcp_list_connections" in names
assert "wf.mcp_get_connection_statuses" in names assert "wf.mcp_get_connection_statuses" in names
assert "wf.mcp_list_proxy_tools" in names assert "wf.mcp_list_proxy_tools" in names
assert "wf.mcp_get_proxy_tool" 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") connections_result = await client.call_tool("wf.mcp_list_connections")
@@ -76,13 +77,26 @@ def test_transparent_proxy_lists_and_calls_upstream_tools() -> None:
assert _structured(result) == {"echoed": "hello"} assert _structured(result) == {"echoed": "hello"}
proxy_tools_result = await client.call_tool("wf.mcp_list_proxy_tools") proxy_tools_result = await client.call_tool("wf.mcp_list_proxy_tools")
proxy_tools = _structured(proxy_tools_result)["result"] proxy_tools_payload = _structured(proxy_tools_result)
proxy_tools = proxy_tools_payload["tools"]
assert proxy_tools_payload["nextCursor"] is None
assert proxy_tools_payload["total"] == 1
assert len(proxy_tools) == 1 assert len(proxy_tools) == 1
assert proxy_tools[0]["proxy_name"] == "fixture.personal_echo_tool" assert proxy_tools[0]["proxy_name"] == "fixture.personal_echo_tool"
assert proxy_tools[0]["connection_id"] == "fixture.personal" assert proxy_tools[0]["connection_id"] == "fixture.personal"
assert proxy_tools[0]["local_name"] == "echo_tool" assert proxy_tools[0]["local_name"] == "echo_tool"
assert proxy_tools[0]["enabled"] is True assert proxy_tools[0]["enabled"] is True
proxy_tool_result = await client.call_tool(
"wf.mcp_get_proxy_tool",
{"proxy_name": "fixture.personal_echo_tool"},
)
proxy_tool = _structured(proxy_tool_result)
assert proxy_tool["proxy_name"] == "fixture.personal_echo_tool"
assert proxy_tool["connection_id"] == "fixture.personal"
assert proxy_tool["local_name"] == "echo_tool"
assert proxy_tool["input_schema"]["properties"]["text"]["type"] == "string"
asyncio.run(run_proxy()) asyncio.run(run_proxy())
@@ -206,6 +220,71 @@ def test_transparent_proxy_can_collapse_upstream_tools_behind_search() -> None:
asyncio.run(run_proxy()) asyncio.run(run_proxy())
def test_transparent_proxy_proxy_tool_listing_supports_filters_and_cursor() -> None:
config = BrokerConfig(
store_root=local_temp_root() / "transparent_proxy_paged_tools_store",
connections=[
ConnectionConfig(
id="fixture.personal",
server="fixture",
account="personal",
metadata={
"transport": "stdio",
"command": sys.executable,
"args": [fixture_server_path()],
},
),
ConnectionConfig(
id="fixture.work",
server="fixture",
account="work",
metadata={
"transport": "stdio",
"command": sys.executable,
"args": [fixture_server_path()],
},
),
],
)
async def run_proxy() -> None:
client = create_transparent_proxy_client(config)
async with client:
first_page_result = await client.call_tool(
"wf.mcp_list_proxy_tools",
{"limit": 1},
)
first_page = _structured(first_page_result)
assert len(first_page["tools"]) == 1
assert first_page["nextCursor"] is not None
assert first_page["total"] == 2
second_page_result = await client.call_tool(
"wf.mcp_list_proxy_tools",
{"limit": 1, "cursor": first_page["nextCursor"]},
)
second_page = _structured(second_page_result)
assert len(second_page["tools"]) == 1
assert second_page["tools"][0]["proxy_name"] != first_page["tools"][0][
"proxy_name"
]
filtered_result = await client.call_tool(
"wf.mcp_list_proxy_tools",
{
"connection_id": "fixture.personal",
"query": "echo",
"limit": 10,
},
)
filtered = _structured(filtered_result)
assert filtered["nextCursor"] is None
assert filtered["total"] == 1
assert filtered["tools"][0]["proxy_name"] == "fixture.personal_echo_tool"
asyncio.run(run_proxy())
def test_transparent_proxy_admin_tools_mutate_config_file() -> None: def test_transparent_proxy_admin_tools_mutate_config_file() -> None:
tmp_path = local_temp_root() / "transparent_proxy_admin_store" tmp_path = local_temp_root() / "transparent_proxy_admin_store"
tmp_path.mkdir(parents=True, exist_ok=True) tmp_path.mkdir(parents=True, exist_ok=True)
+4 -7
View File
@@ -9,9 +9,7 @@
"metadata": { "metadata": {
"transport": "stdio", "transport": "stdio",
"command": "pnpx", "command": "pnpx",
"args": [ "args": ["@upstash/context7-mcp"],
"@upstash/context7-mcp"
],
"env": {} "env": {}
} }
}, },
@@ -25,7 +23,8 @@
"command": "serena", "command": "serena",
"args": [ "args": [
"start-mcp-server", "start-mcp-server",
"--enable-web-dashboard=true" "--enable-web-dashboard=true",
"--project-from-cwd"
], ],
"env": {} "env": {}
} }
@@ -38,9 +37,7 @@
"metadata": { "metadata": {
"transport": "stdio", "transport": "stdio",
"command": "pnpx", "command": "pnpx",
"args": [ "args": ["@modelcontextprotocol/server-everything"],
"@modelcontextprotocol/server-everything"
],
"env": {} "env": {}
} }
} }