feat: route mcp runtime prompt gets

This commit is contained in:
lda
2026-06-07 19:13:23 +07:00 Verified
parent f89801e7a8
commit b6bdae2020
7 changed files with 130 additions and 4 deletions
+3
View File
@@ -250,6 +250,9 @@ implementation state.
the owner-task queue and `McpSourceClient`. This is intentionally a thin the owner-task queue and `McpSourceClient`. This is intentionally a thin
wrapper over the existing source-client facade; prompt/raw method runtime wrapper over the existing source-client facade; prompt/raw method runtime
operations remain separate future slices. operations remain separate future slices.
- Completed: persistent MCP runtime can now route `get_prompt` through
the owner-task queue and `McpSourceClient`. This keeps prompt reads
stateful without adding raw method invocation or notification support.
- Auth/source secrets boundary: keep registry desired state separate from - Auth/source secrets boundary: keep registry desired state separate from
upstream credentials, and surface missing auth as validation diagnostics. upstream credentials, and surface missing auth as validation diagnostics.
The contract is now specified in The contract is now specified in
@@ -116,8 +116,11 @@ First slices should move leaf modules only and leave `wf_mcp` re-export shims:
only `call_tool`; non-tool methods require a separate public-surface slice. only `call_tool`; non-tool methods require a separate public-surface slice.
11. Complete: persistent MCP runtime can route `read_resource` through the 11. Complete: persistent MCP runtime can route `read_resource` through the
owner-task queue and `McpSourceClient`. Runtime still does not expose owner-task queue and `McpSourceClient`. Runtime still does not expose
`get_prompt`, raw method invocation, or notifications. raw method invocation, notifications, or discovery list operations.
12. Upstream transport/discovery/session services. 12. Complete: persistent MCP runtime can route `get_prompt` through the
owner-task queue and `McpSourceClient`. Runtime still does not expose raw
method invocation, notifications, or discovery list operations.
13. Upstream transport/discovery/session services.
Each slice should add import-direction tests so the new source-provider package Each slice should add import-direction tests so the new source-provider package
does not depend on `wf_mcp.workflow_surface`, `wf_mcp.admin_surface`, does not depend on `wf_mcp.workflow_surface`, `wf_mcp.admin_surface`,
+12
View File
@@ -43,6 +43,7 @@ class PersistentSessionFactory:
auth=auth, auth=auth,
call_callback=owner.call_tool, call_callback=owner.call_tool,
read_resource_callback=owner.read_resource, read_resource_callback=owner.read_resource,
get_prompt_callback=owner.get_prompt,
close_callback=owner.close, close_callback=owner.close,
) )
@@ -151,6 +152,17 @@ class _SessionOwner:
run=lambda client: client.read_resource(uri), run=lambda client: client.read_resource(uri),
) )
async def get_prompt(
self,
prompt_name: str,
arguments: dict[str, str] | None = None,
) -> dict[str, Any]:
"""Submit a prompt get through the generic owner-task operation queue."""
return await self.submit(
operation="get_prompt",
run=lambda client: client.get_prompt(prompt_name, arguments),
)
async def close(self) -> None: async def close(self) -> None:
"""Ask the owner task to close the MCP transport in its own scope.""" """Ask the owner task to close the MCP transport in its own scope."""
task = self._task task = self._task
+10
View File
@@ -92,6 +92,16 @@ class McpRuntimePool:
session = await self.get_session(connection, auth) session = await self.get_session(connection, auth)
return await session.read_resource(uri) return await session.read_resource(uri)
async def get_prompt(
self,
connection: McpSourceConnection,
auth: AuthRecord | None,
prompt_name: str,
arguments: dict[str, str] | None = None,
) -> dict[str, Any]:
session = await self.get_session(connection, auth)
return await session.get_prompt(prompt_name, arguments)
async def close_connection(self, connection_id: str) -> None: async def close_connection(self, connection_id: str) -> None:
current = self._sessions.pop(connection_id, None) current = self._sessions.pop(connection_id, None)
if current is not None: if current is not None:
+18
View File
@@ -14,6 +14,10 @@ from wf_sources_mcp.sdk.converters import tool_result_to_call_result
RawToolCaller = Callable[[str, dict[str, Any]], Awaitable[ToolCallResult]] RawToolCaller = Callable[[str, dict[str, Any]], Awaitable[ToolCallResult]]
RawResourceReader = Callable[[str], Awaitable[dict[str, Any]]] RawResourceReader = Callable[[str], Awaitable[dict[str, Any]]]
RawPromptGetter = Callable[
[str, dict[str, str] | None],
Awaitable[dict[str, Any]],
]
@dataclass(slots=True) @dataclass(slots=True)
@@ -35,6 +39,7 @@ class PersistentMcpSession:
client: ClientSession | None = None client: ClientSession | None = None
call_callback: RawToolCaller | None = None call_callback: RawToolCaller | None = None
read_resource_callback: RawResourceReader | None = None read_resource_callback: RawResourceReader | None = None
get_prompt_callback: RawPromptGetter | None = None
close_callback: Callable[[], Awaitable[None]] | None = None close_callback: Callable[[], Awaitable[None]] | None = None
async def call_tool( async def call_tool(
@@ -56,6 +61,19 @@ class PersistentMcpSession:
return result.model_dump(by_alias=True, mode="json", exclude_none=True) return result.model_dump(by_alias=True, mode="json", exclude_none=True)
raise RuntimeError("persistent MCP session has no resource read transport") raise RuntimeError("persistent MCP session has no resource read transport")
async def get_prompt(
self,
prompt_name: str,
arguments: dict[str, str] | None = None,
) -> dict[str, Any]:
"""Get an MCP prompt through the owner task or injected session."""
if self.get_prompt_callback is not None:
return await self.get_prompt_callback(prompt_name, arguments)
if self.client is not None:
result = await self.client.get_prompt(prompt_name, arguments)
return result.model_dump(by_alias=True, mode="json", exclude_none=True)
raise RuntimeError("persistent MCP session has no prompt transport")
async def close(self) -> None: async def close(self) -> None:
"""Close the transport/session stack owned by the runtime factory.""" """Close the transport/session stack owned by the runtime factory."""
if self.close_callback is not None: if self.close_callback is not None:
+82 -2
View File
@@ -65,6 +65,14 @@ async def test_persistent_session_raises_without_resource_transport() -> None:
await session.read_resource("test://x") await session.read_resource("test://x")
@pytest.mark.asyncio
async def test_persistent_session_raises_without_prompt_transport() -> None:
session = PersistentMcpSession(connection=_connection(), auth=None)
with pytest.raises(RuntimeError, match="no prompt transport"):
await session.get_prompt("prompt.summarize")
class _FakeFactory(PersistentSessionFactory): class _FakeFactory(PersistentSessionFactory):
def __init__(self) -> None: def __init__(self) -> None:
self.calls: list[tuple[str, dict[str, object]]] = [] self.calls: list[tuple[str, dict[str, object]]] = []
@@ -105,6 +113,29 @@ class _FakeFactory(PersistentSessionFactory):
}, },
)() )()
async def get_prompt(
self,
prompt_name: str,
arguments: dict[str, str] | None = None,
):
return type(
"GetPromptResult",
(),
{
"model_dump": lambda _self, **_kwargs: {
"messages": [
{
"role": "user",
"content": {
"type": "text",
"text": f"{prompt_name}:{arguments or {}}",
},
}
]
}
},
)()
return _FakeClient() # type: ignore[return-value] return _FakeClient() # type: ignore[return-value]
@@ -168,7 +199,7 @@ def test_runtime_fingerprint_changes_when_transport_changes() -> None:
) )
def test_persistent_session_public_runtime_exposes_only_tool_and_resource_read() -> None: def test_persistent_session_public_runtime_exposes_safe_read_operations() -> None:
public_operations = { public_operations = {
name name
for name in dir(PersistentMcpSession) for name in dir(PersistentMcpSession)
@@ -177,11 +208,32 @@ def test_persistent_session_public_runtime_exposes_only_tool_and_resource_read()
assert "call_tool" in public_operations assert "call_tool" in public_operations
assert "read_resource" in public_operations assert "read_resource" in public_operations
assert "get_prompt" not in public_operations assert "get_prompt" in public_operations
assert "invoke_method" not in public_operations assert "invoke_method" not in public_operations
assert "send_notification" not in public_operations assert "send_notification" not in public_operations
@pytest.mark.asyncio
async def test_persistent_session_factory_routes_prompts_through_owner() -> None:
factory = _FakeFactory()
connection = _connection()
session = await factory.create(connection, None)
await session.call_tool("echo", {"text": "one"})
await session.read_resource("fixture://docs/welcome")
prompt_payload = await session.get_prompt(
"prompt.summarize",
{"text": "hello"},
)
await session.close()
assert factory.created_connections == [connection]
assert factory.calls == [("echo", {"text": "one"})]
assert prompt_payload["messages"][0]["content"]["text"] == (
"prompt.summarize:{'text': 'hello'}"
)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_persistent_session_factory_routes_resource_reads_through_owner() -> None: async def test_persistent_session_factory_routes_resource_reads_through_owner() -> None:
factory = _FakeFactory() factory = _FakeFactory()
@@ -218,3 +270,31 @@ async def test_runtime_pool_reuses_session_for_tool_and_resource_read() -> None:
assert tool_result.output == {"echoed": "one"} assert tool_result.output == {"echoed": "one"}
assert resource_payload["contents"][0]["text"] == "resource text" assert resource_payload["contents"][0]["text"] == "resource text"
assert factory.created_connections == [connection] assert factory.created_connections == [connection]
@pytest.mark.asyncio
async def test_runtime_pool_reuses_session_for_tool_resource_and_prompt() -> None:
factory = _FakeFactory()
pool = McpRuntimePool(factory.create)
connection = _connection()
tool_result = await pool.call_tool(connection, None, "echo", {"text": "one"})
resource_payload = await pool.read_resource(
connection,
None,
"fixture://docs/welcome",
)
prompt_payload = await pool.get_prompt(
connection,
None,
"prompt.summarize",
{"text": "hello"},
)
await pool.close_all()
assert tool_result.output == {"echoed": "one"}
assert resource_payload["contents"][0]["text"] == "resource text"
assert prompt_payload["messages"][0]["content"]["text"] == (
"prompt.summarize:{'text': 'hello'}"
)
assert factory.created_connections == [connection]