subclass Transform js to get our shit looking nice holy moly
This commit is contained in:
@@ -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.
|
||||
|
||||
@@ -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,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}"
|
||||
@@ -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),
|
||||
|
||||
Vendored
+16
@@ -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",
|
||||
|
||||
@@ -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"],
|
||||
|
||||
Reference in New Issue
Block a user