This commit is contained in:
lda
2026-05-25 23:27:11 +07:00 Verified
parent dc0799b2ff
commit 334d69a98a
9 changed files with 170 additions and 15 deletions
@@ -0,0 +1,49 @@
# Stateful Transparent Proxy Implementation Plan
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
**Goal:** Preserve one upstream MCP session for visible proxy operations made by one connected downstream client, while documenting that generic upstream list/resource notifications are still not relayed.
**Architecture:** The workflow execution pool already owns background/offline sessions. The transparent proxy should use FastMCP's `StatefulProxyClient`, which is designed for Playwright-like upstreams and scopes reuse to one downstream MCP session. This slice does not unify interactive proxy sessions with deployment runtimes and does not implement arbitrary notification rebroadcast.
**Tech Stack:** Python 3.14, FastMCP `StatefulProxyClient` and `FastMCPProxy`, pytest fixture MCP server, ruff, basedpyright.
---
### Task 1: Prove Stateful Proxy Behavior
**Files:**
- Modify: `tests/fixtures/mcp_echo_server.py`
- Modify: `tests/wf_mcp/test_proxy.py`
- Modify: `tests/wf_mcp/test_protocol_relay.py`
- [ ] Add fixture tools that store and read a value in the upstream server process.
- [ ] Add a proxy test that writes the value through one proxied request and reads it through another request in the same downstream client session.
- [ ] Change protocol-relay coverage to assert that `tools/list_changed`, `resources/list_changed`, `prompts/list_changed`, and `resources/updated` are not yet relayed. Keep string-valued logging forwarding as a strict expected-failure tripwire because the installed FastMCP `StatefulProxyClient` handler currently assumes mapping-valued MCP log data.
- [ ] Run the focused tests and confirm they fail before proxy construction changes.
### Task 2: Use FastMCP Stateful Proxy Sessions
**Files:**
- Modify: `src/wf_mcp/proxy/mounts.py`
- [ ] Replace `create_proxy(Client(...))` with `StatefulProxyClient(...)` and `FastMCPProxy(client_factory=client.new_stateful, ...)`.
- [ ] Preserve the existing `ProxyNamespace` and `ResourceLinkNamespace` transforms exactly as mounted-provider output transforms.
- [ ] Add a docstring/comment stating that FastMCP owns the interactive session lifecycle and this is intentionally separate from offline workflow execution sessions.
### Task 3: Verification
**Files:**
- Test: `tests/wf_mcp/test_proxy.py`
- Test: `tests/wf_mcp/test_protocol_relay.py`
- [ ] Run focused proxy/protocol tests.
- [ ] Run `uv run pytest -q`.
- [ ] Run `uvx ruff check`.
- [ ] Run `uv run basedpyright --level error`.
## Scope Boundary
- Interactive visible proxy calls share state within one downstream MCP client session.
- Deployment execution continues to use the owned runtime pool because scheduled/background runs may exist without a downstream client session.
- Generic upstream notification relay remains separate work. FastMCP's stateful path is intended to forward logs/progress/elicitation, but string-valued MCP log data currently exposes an upstream FastMCP handler bug and remains documented with an expected-failure test.
+6 -1
View File
@@ -19,7 +19,12 @@ from .models import (
def build_workflow_from_draft(draft: WorkflowDraft) -> Workflow:
"""Adapt one typed draft through `WorkflowBuilder` into a core workflow."""
"""Adapt one typed draft through `WorkflowBuilder` into a core workflow.
Draft step `output` bindings become node-output-to-state writes. Final
workflow output projection stays in core runtime and uses output schema
property names as state keys.
"""
builder = WorkflowBuilder(
name=draft.name,
input_schema=draft.input_schema,
+13 -2
View File
@@ -28,7 +28,12 @@ STEP_KIND_KEYS = frozenset(
class DraftUseStep(BaseModel):
"""Draft step that calls one externally resolvable workflow capability."""
"""Draft step that calls one externally resolvable workflow capability.
`output` writes node-local output fields into workflow state. It does not
define final workflow output; core currently projects final output from
state keys whose names match `WorkflowDraft.output_schema.properties`.
"""
model_config = ConfigDict(extra="forbid", populate_by_name=True)
@@ -270,7 +275,13 @@ DraftStep = (
class WorkflowDraft(BaseModel):
"""Patch-friendly JSON authoring document for one workflow graph."""
"""Patch-friendly JSON authoring document for one workflow graph.
There is intentionally no top-level output map in this draft shape. Final
workflow output is projected from state by matching output-schema property
names, so terminal steps should write required output fields to same-named
state paths.
"""
name: str
input_schema: JsonObject
+5
View File
@@ -438,6 +438,11 @@ def reducer_for_state_path(
def project_output(workflow: Workflow, state: dict[str, Any]) -> dict[str, Any]:
"""Project final workflow output from same-named state fields.
Node output bindings write into state during execution. At END, the runtime
exposes only state keys declared by `workflow.output_schema.properties`.
"""
return {
key: state[key] for key in workflow.output_schema.properties if key in state
}
+14 -5
View File
@@ -7,9 +7,8 @@ from pathlib import Path
from typing import Any, Generic, TypeVar
from fastmcp import FastMCP
from fastmcp.client import Client
from fastmcp.client.transports.config import MCPConfigTransport
from fastmcp.server import create_proxy
from fastmcp.server.providers.proxy import FastMCPProxy, StatefulProxyClient
from ..models import BrokerConfig, ConnectionConfig
from ..proxy_results import ResourceLinkNamespace
from ..proxy_config import broker_config_to_fastmcp_config
@@ -90,13 +89,23 @@ def create_proxy_mount(
connection: ConnectionConfig,
store_root: Path,
) -> ProxyMount[FastMCP[Any]]:
"""Create one FastMCP proxy mount for an enabled upstream connection."""
"""Create one interactive stateful FastMCP proxy mount for a connection.
FastMCP's `StatefulProxyClient` owns the upstream session per downstream
MCP client and restores request context for relayed progress, logging,
elicitation, and sampling interactions. Offline workflow executions use
the separate runtime pool because they do not have an interactive client
session to scope this lifetime to.
"""
server_config = broker_config_to_fastmcp_config(
BrokerConfig(store_root=store_root, connections=[connection])
)
transport = MCPConfigTransport(server_config, name_as_prefix=False)
client = Client(transport=transport, name=f"wf-mcp:{connection.id}")
proxy: FastMCP[Any] = create_proxy(client, name=f"Proxy-{connection.id}")
client = StatefulProxyClient(transport=transport, name=f"wf-mcp:{connection.id}")
proxy: FastMCP[Any] = FastMCPProxy(
client_factory=client.new_stateful,
name=f"Proxy-{connection.id}",
)
proxy.add_transform(ProxyNamespace(connection.id))
proxy.add_transform(ResourceLinkNamespace(connection.id))
return ProxyMount(
+17
View File
@@ -9,6 +9,7 @@ from pydantic import Field
server = FastMCP("echo-fixture")
_remembered_value: str | None = None
class EchoToolResult(TypedDict):
@@ -22,6 +23,22 @@ async def echo_tool(
return {"echoed": text}
@server.tool(title="Remember value tool")
async def remember_value_tool(
value: Annotated[str, Field(description="Value held in this server process.")],
) -> dict[str, str]:
"""Store state so proxy tests can distinguish reused and fresh sessions."""
global _remembered_value
_remembered_value = value
return {"remembered": value}
@server.tool(title="Recall value tool")
async def recall_value_tool() -> dict[str, str | None]:
"""Return process-local state written by `remember_value_tool`."""
return {"remembered": _remembered_value}
@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."""
+27 -3
View File
@@ -70,7 +70,7 @@ def test_fixture_server_emits_observable_protocol_notifications_directly() -> No
assert "notifications/message" in methods
def test_proxy_does_not_relay_upstream_protocol_notifications_yet() -> None:
def _fixture_proxy_notification_methods() -> list[str]:
config = BrokerConfig(
store_root=local_temp_root() / "protocol_relay_store",
connections=[
@@ -102,5 +102,29 @@ def test_proxy_does_not_relay_upstream_protocol_notifications_yet() -> None:
except PermissionError as exc:
pytest.skip(f"stdio MCP transport is not permitted in this environment: {exc}")
methods = _notification_methods(notifications)
assert methods == []
return _notification_methods(notifications)
def test_proxy_does_not_relay_generic_upstream_notifications_yet() -> None:
methods = _fixture_proxy_notification_methods()
# Stateful proxy sessions preserve FastMCP's supported request callbacks,
# but generic upstream change/update notification rebroadcast is separate
# protocol relay work.
assert "notifications/tools/list_changed" not in methods
assert "notifications/resources/list_changed" not in methods
assert "notifications/prompts/list_changed" not in methods
assert "notifications/resources/updated" not in methods
@pytest.mark.xfail(
strict=True,
reason=(
"FastMCP StatefulProxyClient log forwarding assumes mapping-valued log "
"data; valid string-valued MCP logging data is rejected upstream."
),
)
def test_proxy_relays_string_valued_upstream_log_when_fastmcp_supports_it() -> None:
methods = _fixture_proxy_notification_methods()
assert "notifications/message" in methods
+38 -3
View File
@@ -83,14 +83,16 @@ def test_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"] == 3
assert len(proxy_tools) == 3
assert proxy_tools_payload["total"] == 5
assert len(proxy_tools) == 5
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
proxy_names = [tool["proxy_name"] for tool in proxy_tools]
assert "fixture.personal.emit_notifications_tool" in proxy_names
assert "fixture.personal.remember_value_tool" in proxy_names
assert "fixture.personal.recall_value_tool" in proxy_names
assert "fixture.personal.resource_link_tool" in proxy_names
proxy_tool_result = await client.call_tool(
@@ -153,6 +155,39 @@ def test_proxy_rewrites_resource_links_returned_by_tools() -> None:
asyncio.run(run_proxy())
def test_proxy_reuses_one_upstream_session_for_stateful_tools() -> None:
"""Visible proxy tools must share server-local state for one MCP client."""
config = BrokerConfig(
store_root=local_temp_root() / "proxy_stateful_session_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_proxy_client(config)
async with client:
written = await client.call_tool(
"fixture.personal.remember_value_tool",
{"value": "held"},
)
recalled = await client.call_tool("fixture.personal.recall_value_tool")
assert _structured(written)["remembered"] == "held"
assert _structured(recalled)["remembered"] == "held"
asyncio.run(run_proxy())
def test_proxy_rejects_invalid_connection_config() -> None:
config = BrokerConfig(
store_root=local_temp_root() / "proxy_invalid_store",
@@ -355,7 +390,7 @@ def test_proxy_proxy_tool_listing_supports_filters_and_cursor() -> None:
first_page = _structured(first_page_result)
assert len(first_page["tools"]) == 1
assert first_page["nextCursor"] is not None
assert first_page["total"] == 6
assert first_page["total"] == 10
second_page_result = await client.call_tool(
"wf.admin.list_proxy_tools",