lda.chat/lda-wf
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: route all mcp operations through runtime pool

Browse Source
This commit is contained in:
lda
2026-06-08 18:29:22 +07:00 Verified
parent 929103219e
commit 0a8cc4190f
14 changed files with 961 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/wf_mcp/service/test_content_access.py
+53
View File
@@ -181,6 +181,59 @@ class _StatefulRuntime:
]
}
async def list_tools(self, connection, auth):
from wf_sources_mcp.catalog import DiscoveredTool
return [
DiscoveredTool(
name="echo_tool",
title="Echo Tool",
description="Echo text back",
input_schema={"type": "object"},
output_schema={"type": "object"},
)
]
async def list_resources(self, connection, auth):
from wf_sources_mcp.catalog import DiscoveredResource
return [
DiscoveredResource(
uri="demo://docs/welcome",
name="resource.welcome",
title="Welcome Resource",
description="Welcome resource",
mime_type="text/plain",
)
]
async def list_prompts(self, connection, auth):
from wf_sources_mcp.catalog import DiscoveredPrompt
return [
DiscoveredPrompt(
name="prompt.summarize",
title="Summarize Prompt",
description="Summarize text",
arguments=[
{
"name": "text",
"required": True,
"description": "Text to summarize",
}
],
)
]
async def get_connection_metadata(self, connection, auth):
return {"server": connection.server, "transport": "stdio"}
async def invoke_method(self, connection, auth, method, params=None):
raise AssertionError("not used by content access tests")
async def send_notification(self, connection, auth, method, params=None):
raise AssertionError("not used by content access tests")
async def test_content_access_uses_stateful_runtime_for_upstream_content() -> None:
runtime = _StatefulRuntime()
tests/wf_mcp/service/test_upstream_transport.py
+210 -3
View File
@@ -11,6 +11,7 @@ from wf_mcp.events import McpEvent
from wf_mcp.models import AuthRecord, CatalogSnapshot, ConnectionConfig
from wf_mcp.storage import FileAuthStore, FileCatalogStore, FileStore
from wf_platform import CapabilityBuckets, CapabilitySource, SourcePermissions
from wf_sources_mcp.catalog import DiscoveredTool
from ..test_support import FakeAdapter, local_temp_root
from ..workflow_surface.conftest import echo_artifact
@@ -357,6 +358,11 @@ class _StatefulRuntime:
def __init__(self) -> None:
self.resources: list[tuple[str, str]] = []
self.prompts: list[tuple[str, str, dict[str, str] | None]] = []
self.tools_called: list[tuple[str, str]] = []
self.methods_invoked: list[tuple[str, str, str, dict[str, object] | None]] = []
self.notifications_sent: list[
tuple[str, str, str, dict[str, object] | None]
] = []
async def call_tool(self, connection, auth, tool_name, payload):
raise AssertionError("not used by these tests")
@@ -382,6 +388,38 @@ class _StatefulRuntime:
]
}
async def list_tools(self, connection, auth):
self.tools_called.append((connection.id, connection.provider))
return [
DiscoveredTool(
name="stateful_tool",
title="Stateful Tool",
description="A stateful tool",
input_schema={"type": "object"},
output_schema={"type": "object"},
)
]
async def list_resources(self, connection, auth):
return []
async def list_prompts(self, connection, auth):
return []
async def get_connection_metadata(self, connection, auth):
return {"server": connection.provider, "transport": "stdio"}
async def invoke_method(self, connection, auth, method, params=None):
self.methods_invoked.append(
(connection.id, connection.provider, method, params)
)
return {"echoed": (params or {}).get("text", "")}
async def send_notification(self, connection, auth, method, params=None):
self.notifications_sent.append(
(connection.id, connection.provider, method, params)
)
class _ExplodingContentAdapter(FakeAdapter):
async def read_resource(self, connection, auth, uri):
@@ -451,10 +489,179 @@ async def test_upstream_transport_prefers_stateful_runtime_for_prompts(
)
assert result["messages"][0]["content"]["text"] == "stateful prompt"
assert runtime.prompts == [
("demo.personal", "prompt.summarize", {"text": "hello"})
]
assert runtime.prompts == [("demo.personal", "prompt.summarize", {"text": "hello"})]
assert [event.kind for event in events] == [
"prompt_get_started",
"prompt_get_completed",
]
async def test_upstream_transport_prefers_stateful_runtime_for_invoke_method(
tmp_path: Path,
) -> None:
events: list[McpEvent] = []
runtime = _StatefulRuntime()
transport = UpstreamTransportService(
auth_store=FileStore(tmp_path),
catalog_store=FileStore(tmp_path),
event_sink=events.append,
stateful_runtime=runtime,
)
transport.register_adapter("demo", _ExplodingContentAdapter())
connection = ConnectionConfig(
id="demo.personal",
server="demo",
account="personal",
metadata=_fake_transport_metadata(),
)
result = await transport.invoke_method(
connection,
"demo.echo",
params={"text": "hello"},
)
assert result["echoed"] == "hello"
assert runtime.methods_invoked == [
("demo.personal", "demo", "demo.echo", {"text": "hello"})
]
assert [event.kind for event in events] == [
"raw_method_started",
"raw_method_completed",
]
async def test_upstream_transport_prefers_stateful_runtime_for_send_notification(
tmp_path: Path,
) -> None:
events: list[McpEvent] = []
runtime = _StatefulRuntime()
transport = UpstreamTransportService(
auth_store=FileStore(tmp_path),
catalog_store=FileStore(tmp_path),
event_sink=events.append,
stateful_runtime=runtime,
)
transport.register_adapter("demo", _ExplodingContentAdapter())
connection = ConnectionConfig(
id="demo.personal",
server="demo",
account="personal",
metadata=_fake_transport_metadata(),
)
await transport.send_notification(
connection,
"notifications/test",
params={"data": "value"},
)
assert runtime.notifications_sent == [
("demo.personal", "demo", "notifications/test", {"data": "value"})
]
assert [event.kind for event in events] == [
"raw_notification_started",
"raw_notification_completed",
]
async def test_upstream_transport_prefers_stateful_runtime_for_catalog_refresh(
tmp_path: Path,
) -> None:
events: list[McpEvent] = []
store = FileStore(tmp_path)
connections = ConnectionRegistry()
connection = ConnectionConfig(
id="demo.personal",
server="demo",
account="personal",
metadata=_fake_transport_metadata(),
)
connections.register(connection)
runtime = _StatefulRuntime()
transport = UpstreamTransportService(
auth_store=store,
catalog_store=store,
event_sink=events.append,
stateful_runtime=runtime,
)
transport.register_adapter("demo", _ExplodingContentAdapter())
source_catalog = SourceCatalogService(
store=store,
connection_lookup=connections.get,
connection_list_enabled=connections.list_enabled,
connection_list_all=connections.list_all,
tool_executor_for=transport.tool_executor_for,
load_auth=transport.load_connection_auth,
emit_event=events.append,
)
source_catalog.hydrate_connection_source_from_snapshot(connection)
await transport.refresh_connection_catalog(
connection,
source_catalog=source_catalog,
record_catalog_change_events=lambda source_id, snapshot, reason: None,
)
assert runtime.tools_called == [("demo.personal", "demo")]
assert "catalog_refresh_started" in [event.kind for event in events]
assert "catalog_refresh_completed" in [event.kind for event in events]
class _ExplodingAdapterForDiagnostics(FakeAdapter):
async def list_tools(self, connection, auth):
raise AssertionError("adapter list_tools should not be used in diagnostics")
async def test_upstream_transport_prefers_stateful_runtime_for_deployment_diagnostics(
tmp_path: Path,
) -> None:
runtime = _StatefulRuntime()
connections = ConnectionRegistry()
connection = ConnectionConfig(
id="demo.personal",
server="demo",
account="personal",
metadata=_fake_transport_metadata(),
)
connections.register(connection)
transport = UpstreamTransportService(
auth_store=FileStore(tmp_path),
catalog_store=FileStore(tmp_path),
event_sink=lambda event: None,
stateful_runtime=runtime,
)
transport.register_adapter("demo", _ExplodingAdapterForDiagnostics())
source_catalog = SourceCatalogService(
store=transport.catalog_store,
connection_lookup=connections.get,
connection_list_enabled=connections.list_enabled,
connection_list_all=connections.list_all,
tool_executor_for=transport.tool_executor_for,
load_auth=transport.load_connection_auth,
emit_event=lambda event: None,
)
source_catalog.register_capability_source(
CapabilitySource(
id="demo.personal",
kind="connection",
permissions=SourcePermissions(calls_upstream=True),
capabilities=CapabilityBuckets(),
)
)
artifact = echo_artifact()
deployment = WorkflowDeployment(
id="echo.personal",
artifact_id="echo",
artifact_version=1,
bindings=[{"logical_source": "demo", "concrete_source": "demo.personal"}],
)
diagnostics = await transport.deployment_diagnostics(
deployment=deployment,
artifacts=[artifact],
source_catalog=source_catalog,
)
assert diagnostics == []
assert runtime.tools_called == [("demo.personal", "demo")]