subclass Transform js to get our shit looking nice holy moly

This commit is contained in:
lda
2026-05-15 23:37:45 +07:00 Verified
parent 6935329d44
commit 9a035e9557
8 changed files with 224 additions and 24 deletions
+7 -10
View File
@@ -98,12 +98,11 @@ Tool simulate-research-query requires task augmentation (taskSupport: 'required'
- Proxy tool projection: healthy for ordinary tools.
- Proxy resource list/read projection: healthy for listed resources and
templates.
- Tool-result resource links: incomplete; embedded resource URIs need explicit
rewrite or a documented limitation.
- `wf_mcp.proxy_results` now contains pure typed helpers for rewriting
`mcp.types.ResourceLink` content inside `mcp.types.CallToolResult`. These
helpers are not wired into the FastMCP proxy runtime yet because FastMCP does
not currently expose a result-transform hook.
- Tool-result resource links: ordinary embedded resource URIs are now rewritten
by a local proxy transform so downstream clients receive namespaced URIs.
- `wf_mcp.proxy_results` contains pure typed helpers for rewriting
`mcp.types.ResourceLink` content and a small FastMCP workaround transform that
wraps proxied tools until upstream FastMCP handles this projection itself.
- Session resources: unresolved; needs a focused test because session affinity
may matter.
- Tasks: unsupported; task-required tools should remain clearly diagnosed until
@@ -115,10 +114,8 @@ Tool simulate-research-query requires task augmentation (taskSupport: 'required'
- Add automated tests for listed resource/template namespacing through unified
mode.
- Add a fixture or Everything integration test showing tool-returned
`resource_link` URIs are not rewritten today.
- Decide whether to rewrite embedded resource links in tool results or expose a
helper that resolves upstream resource links into proxied resource URIs.
- Keep the fixture regression proving ordinary tool-returned `resource_link`
URIs are rewritten into downstream-facing namespaced URIs.
- Investigate session resource lifecycle separately from URI rewriting.
- Inventory current FastMCP/MCP SDK APIs for tasks, logging notifications, and
resource update notifications before implementing protocol forwarding.
@@ -178,10 +178,28 @@ actually implemented.
- wiring `wf_mcp.proxy_results` helpers into real tool-result handling
- replacing FastMCP behavior that should be fixed upstream instead
## Open Question
## Investigation Result
Whether reused mounts should keep FastMCP `ProxyProvider` component caches across
reload is not yet a promise. Reuse likely does preserve those caches. If that
causes stale catalog behavior, the registry should either invalidate the provider
cache explicitly when supported or treat reload as a reason to rebuild mounts
for affected connections.
FastMCP 3.3.0 makes unchanged-mount reuse reasonable, with caveats:
- `create_proxy_mount()` passes a disconnected `Client` into `create_proxy(...)`.
FastMCP turns that into `client.new()` per request, so reusing a
`FastMCPProxy` does **not** keep one permanently connected upstream session
alive.
- `ProxyProvider` intentionally caches tools, resources, templates, and prompts
for lookup efficiency. The default TTL is 300 seconds.
- Every explicit `list_*` call refreshes the corresponding proxy cache. Direct
lookup paths may use a still-fresh cached list until TTL expiry.
- `FastMCP.mount()` returns no provider handle, and no public general-purpose
unmount API was found. Parent-side provider-list rebuild is still the only
safe visible-surface removal we currently own.
Therefore the current registry behavior is acceptable:
- reuse unchanged enabled mounts
- rebuild changed mounts
- stop remounting disabled or removed mounts
- do not claim safe close/unmount of retired mounts yet
If dynamic upstream catalogs become a practical problem, prefer a documented
refresh/invalidation policy over throwing away unchanged mounts on every reload.
+18 -5
View File
@@ -64,6 +64,20 @@ Unified mode currently reuses `ProxyRuntime` as its proxy mounting engine. The
the place where configured upstream MCP connections become mounted FastMCP
providers. `TransparentProxyRuntime` remains a compatibility alias.
`ProxyRuntime` now owns a small `ProxyMountRegistry`. Reload still clears the
visible mounted provider list and rebuilds it from current config, but unchanged
enabled connections reuse their cached proxy mount instead of recreating a new
client/proxy pair every time. Disabled or removed connections are no longer
mounted after reload, but their cached mounts are only *retired* internally for
now; they are not safely closed or unmounted because FastMCP does not yet expose
the lifecycle hook we need.
The reused FastMCP proxies are not holding one forever-open upstream connection.
`create_proxy_mount()` gives `create_proxy(...)` a disconnected client, and
FastMCP creates fresh request clients from it. What persists across reload is the
proxy/provider object and its component-list caches. FastMCP refreshes those
caches on explicit `list_*` calls and otherwise expires them after its TTL.
After a successful reload, the runtime publishes local `tools_changed`,
`resources_changed`, `prompts_changed`, and `catalog_changed` events when an
event bus is supplied. The admin MCP tool projects the same event kinds into
@@ -75,11 +89,10 @@ typed result to a plain payload at the boundary.
Proxy tool listing similarly uses `ProxyToolPayload` / `ProxyToolsPage`
internally and serializes to admin MCP payloads at the boundary.
Do not memoize mounted proxies or clients without an explicit lifecycle design.
The tempting implementation is a dictionary keyed by connection id around
`create_proxy(client, ...)`, but cached clients need clear close/reconnect/error
semantics. Prefer FastMCP's official unmount/provider lifecycle when it becomes
available.
Do not add more lifecycle behavior outside `ProxyMountRegistry`. Cached clients
still need clear close/reconnect/error semantics, and the registry is the single
place where that future behavior should land. Prefer FastMCP's official
unmount/provider lifecycle when it becomes available.
Do not add notification proxying or long-lived subscription handling across
reloads without first introducing an explicit mount lifecycle boundary.
+3
View File
@@ -3,8 +3,11 @@ from .resource_links import (
rewrite_call_tool_result_resource_links,
rewrite_resource_link_content,
)
from .resource_link_transform import ResourceLinkNamespace, ResourceLinkRewritingTool
__all__ = [
"ResourceLinkNamespace",
"ResourceLinkRewritingTool",
"ResourceUriRewriter",
"rewrite_call_tool_result_resource_links",
"rewrite_resource_link_content",
@@ -0,0 +1,87 @@
from __future__ import annotations
import re
from collections.abc import Callable, Sequence
from typing import Any
from fastmcp.server.transforms import GetToolNext, Transform
from fastmcp.tools.base import Tool, ToolResult
from fastmcp.utilities.versions import VersionSpec
from pydantic import ConfigDict
from pydantic.json_schema import SkipJsonSchema
from .resource_links import rewrite_resource_link_content
_URI_PATTERN = re.compile(r"^([^:]+://)(.*?)$")
class ResourceLinkRewritingTool(Tool):
"""Delegate tool execution while rewriting returned resource-link URIs."""
model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
parent_tool: SkipJsonSchema[Tool]
rewrite_uri: SkipJsonSchema[Callable[[str], str]]
async def run(self, arguments: dict[str, Any]) -> ToolResult:
"""Run the parent tool and project any resource links downstream."""
result = await self.parent_tool.run(arguments)
return ToolResult(
content=[
rewrite_resource_link_content(content, self.rewrite_uri)
for content in result.content
],
structured_content=result.structured_content,
meta=result.meta,
)
@classmethod
def wrap(
cls,
tool: Tool,
rewrite_uri: Callable[[str], str],
) -> ResourceLinkRewritingTool:
"""Copy one tool's public schema while replacing only execution."""
return cls.model_validate(
{
**tool.model_dump(),
"parent_tool": tool,
"rewrite_uri": rewrite_uri,
}
)
class ResourceLinkNamespace(Transform):
"""Rewrite resource links returned by tools into one namespace."""
def __init__(self, prefix: str) -> None:
self._prefix = prefix
def __repr__(self) -> str:
return f"ResourceLinkNamespace({self._prefix!r})"
async def list_tools(self, tools: Sequence[Tool]) -> Sequence[Tool]:
"""Wrap listed tools so downstream callers receive projected links."""
return [self._wrap_tool(tool) for tool in tools]
async def get_tool(
self,
name: str,
call_next: GetToolNext,
*,
version: VersionSpec | None = None,
) -> Tool | None:
"""Wrap fetched tools so direct calls also receive projected links."""
tool = await call_next(name, version=version)
return None if tool is None else self._wrap_tool(tool)
def _wrap_tool(self, tool: Tool) -> ResourceLinkRewritingTool:
return ResourceLinkRewritingTool.wrap(tool, self._transform_uri)
def _transform_uri(self, uri: str) -> str:
"""Match FastMCP Namespace URI projection for tool-returned links."""
match = _URI_PATTERN.match(uri)
if match is None:
return uri
protocol, path = match.groups()
return f"{protocol}{self._prefix}/{path}"
+2
View File
@@ -13,6 +13,7 @@ from fastmcp.server import create_proxy
from fastmcp.server.transforms import Namespace
from ..models import BrokerConfig, ConnectionConfig
from ..proxy_results import ResourceLinkNamespace
from ..proxy_config import broker_config_to_fastmcp_config
ProxyT = TypeVar("ProxyT")
@@ -98,6 +99,7 @@ def create_proxy_mount(
client = Client(transport=transport, name=f"wf-mcp:{connection.id}")
proxy: FastMCP[Any] = create_proxy(client, name=f"Proxy-{connection.id}")
proxy.add_transform(Namespace(connection.id))
proxy.add_transform(ResourceLinkNamespace(connection.id))
return ProxyMount(
connection_id=connection.id,
fingerprint=connection_fingerprint(connection),
+16
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
from typing import Annotated, TypedDict
import mcp.types as mcp_types
from mcp.server.fastmcp import FastMCP
from pydantic import Field
@@ -20,6 +21,21 @@ async def echo_tool(
return {"echoed": text}
@server.tool(title="Resource link tool")
async def resource_link_tool() -> list[mcp_types.ResourceLink]:
"""Return a link to a fixture resource so proxy URI rewriting is testable."""
return [
mcp_types.ResourceLink.model_validate(
{
"type": "resource_link",
"name": "resource.welcome",
"uri": "fixture://docs/welcome",
"mimeType": "text/plain",
}
)
]
@server.resource(
"fixture://docs/welcome",
name="resource.welcome",
+67 -3
View File
@@ -83,12 +83,16 @@ def test_transparent_proxy_lists_and_calls_upstream_tools() -> None:
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 proxy_tools_payload["total"] == 2
assert len(proxy_tools) == 2
assert proxy_tools[0]["proxy_name"] == "fixture.personal_echo_tool"
assert proxy_tools[0]["connection_id"] == "fixture.personal"
assert proxy_tools[0]["local_name"] == "echo_tool"
assert proxy_tools[0]["enabled"] is True
assert (
proxy_tools[1]["proxy_name"]
== "fixture.personal_resource_link_tool"
)
proxy_tool_result = await client.call_tool(
"wf.admin.get_proxy_tool",
@@ -103,6 +107,38 @@ def test_transparent_proxy_lists_and_calls_upstream_tools() -> None:
asyncio.run(run_proxy())
def test_transparent_proxy_rewrites_resource_links_returned_by_tools() -> None:
config = BrokerConfig(
store_root=local_temp_root() / "transparent_proxy_resource_link_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)
async with client:
result = await client.call_tool("fixture.personal_resource_link_tool")
link = result.content[0]
assert isinstance(link, mcp_types.ResourceLink)
assert str(link.uri) == "fixture://fixture.personal/docs/welcome"
contents = await client.read_resource(str(link.uri))
assert isinstance(contents[0], mcp_types.TextResourceContents)
assert contents[0].text == "Welcome from the fixture MCP server."
asyncio.run(run_proxy())
def test_transparent_proxy_rejects_invalid_connection_config() -> None:
config = BrokerConfig(
store_root=local_temp_root() / "transparent_proxy_invalid_store",
@@ -267,7 +303,7 @@ def test_transparent_proxy_proxy_tool_listing_supports_filters_and_cursor() -> N
first_page = _structured(first_page_result)
assert len(first_page["tools"]) == 1
assert first_page["nextCursor"] is not None
assert first_page["total"] == 2
assert first_page["total"] == 4
second_page_result = await client.call_tool(
"wf.admin.list_proxy_tools",
@@ -561,6 +597,34 @@ def test_transparent_proxy_runtime_reload_publishes_local_change_events() -> Non
assert catalog_changed[0].payload["reason"] == "transparent_reload"
def test_transparent_proxy_runtime_reload_reuses_unchanged_mounts() -> None:
config = BrokerConfig(
store_root=local_temp_root() / "transparent_proxy_reuse_store",
connections=[
ConnectionConfig(
id="fixture.personal",
server="fixture",
account="personal",
metadata={
"transport": "stdio",
"command": sys.executable,
"args": [fixture_server_path()],
},
)
],
)
runtime = ProxyRuntime(config)
first_mount = runtime.mounts.active_mounts_for(config)[0]
result = runtime.reload()
second_mount = runtime.mounts.active_mounts_for(config)[0]
assert result["mounted_connections"] == ["fixture.personal"]
assert result["connection_count"] == 1
assert result["enabled_connection_count"] == 1
assert first_mount is second_mount
def test_proxy_reload_result_serializes_and_drives_reload_events() -> None:
result = ProxyReloadResult(
mounted_connections=["fixture.personal"],