feat: add stateful runtime listings

This commit is contained in:
lda
2026-06-07 21:22:53 +07:00 Verified
parent 9fea19c550
commit 8ac44fdf2d
13 changed files with 258 additions and 25 deletions
+17
View File
@@ -11,6 +11,7 @@ from typing import Any, Generic, TypeVar
from mcp.client.session import ClientSession
from wf_sources_mcp.auth import AuthRecord
from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource
from wf_sources_mcp.client import McpSourceClient, open_mcp_session
from wf_sources_mcp.connections import McpSourceConnection
from wf_sources_mcp.sdk import ToolCallResult
@@ -44,6 +45,8 @@ class PersistentSessionFactory:
call_callback=owner.call_tool,
read_resource_callback=owner.read_resource,
get_prompt_callback=owner.get_prompt,
list_resources_callback=owner.list_resources,
list_prompts_callback=owner.list_prompts,
close_callback=owner.close,
)
@@ -163,6 +166,20 @@ class _SessionOwner:
run=lambda client: client.get_prompt(prompt_name, arguments),
)
async def list_resources(self) -> list[DiscoveredResource]:
"""Submit resource listing through the generic owner-task operation queue."""
return await self.submit(
operation="list_resources",
run=lambda client: client.list_resources(),
)
async def list_prompts(self) -> list[DiscoveredPrompt]:
"""Submit prompt listing through the generic owner-task operation queue."""
return await self.submit(
operation="list_prompts",
run=lambda client: client.list_prompts(),
)
async def close(self) -> None:
"""Ask the owner task to close the MCP transport in its own scope."""
task = self._task
+17
View File
@@ -7,6 +7,7 @@ from inspect import isawaitable
from typing import Any, cast
from wf_sources_mcp.auth import AuthRecord
from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource
from wf_sources_mcp.connections import McpSourceConnection
from wf_sources_mcp.sdk import ToolCallResult
@@ -102,6 +103,22 @@ class McpRuntimePool:
session = await self.get_session(connection, auth)
return await session.get_prompt(prompt_name, arguments)
async def list_resources(
self,
connection: McpSourceConnection,
auth: AuthRecord | None,
) -> list[DiscoveredResource]:
session = await self.get_session(connection, auth)
return await session.list_resources()
async def list_prompts(
self,
connection: McpSourceConnection,
auth: AuthRecord | None,
) -> list[DiscoveredPrompt]:
session = await self.get_session(connection, auth)
return await session.list_prompts()
async def close_connection(self, connection_id: str) -> None:
current = self._sessions.pop(connection_id, None)
if current is not None:
+31 -8
View File
@@ -8,6 +8,7 @@ from mcp.client.session import ClientSession
from pydantic import AnyUrl
from wf_sources_mcp.auth import AuthRecord
from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource
from wf_sources_mcp.connections import McpSourceConnection
from wf_sources_mcp.sdk import ToolCallResult
from wf_sources_mcp.sdk.converters import tool_result_to_call_result
@@ -18,20 +19,18 @@ RawPromptGetter = Callable[
[str, dict[str, str] | None],
Awaitable[dict[str, Any]],
]
RawResourceLister = Callable[[], Awaitable[list[DiscoveredResource]]]
RawPromptLister = Callable[[], Awaitable[list[DiscoveredPrompt]]]
@dataclass(slots=True)
class PersistentMcpSession:
"""Long-lived MCP execution handle for one configured connection.
Production sessions use `call_callback` because MCP transports are entered
inside an AnyIO cancel scope and must be called and closed by that same
owner task. `client` remains available for simple injected/fake sessions in
tests. `call_tool()` always normalizes SDK results for workflow nodes.
This runtime intentionally exposes only tool calls for now. Shared
non-tool operations live on `McpSourceClient`; routing them through the
owner task is a separate runtime-expansion slice.
Production sessions route all MCP operations through the owner-task queue
for session safety. `client` remains available for simple injected/fake
sessions in tests. `call_tool()` always normalizes SDK results for workflow
nodes.
"""
connection: McpSourceConnection
@@ -40,6 +39,8 @@ class PersistentMcpSession:
call_callback: RawToolCaller | None = None
read_resource_callback: RawResourceReader | None = None
get_prompt_callback: RawPromptGetter | None = None
list_resources_callback: RawResourceLister | None = None
list_prompts_callback: RawPromptLister | None = None
close_callback: Callable[[], Awaitable[None]] | None = None
async def call_tool(
@@ -74,6 +75,28 @@ class PersistentMcpSession:
return result.model_dump(by_alias=True, mode="json", exclude_none=True)
raise RuntimeError("persistent MCP session has no prompt transport")
async def list_resources(self) -> list[DiscoveredResource]:
"""List MCP resources through the owner task or injected session."""
if self.list_resources_callback is not None:
return await self.list_resources_callback()
if self.client is not None:
from wf_sources_mcp.sdk.converters import resource_to_discovered
result = await self.client.list_resources()
return [resource_to_discovered(resource) for resource in result.resources]
raise RuntimeError("persistent MCP session has no resource list transport")
async def list_prompts(self) -> list[DiscoveredPrompt]:
"""List MCP prompts through the owner task or injected session."""
if self.list_prompts_callback is not None:
return await self.list_prompts_callback()
if self.client is not None:
from wf_sources_mcp.sdk.converters import prompt_to_discovered
result = await self.client.list_prompts()
return [prompt_to_discovered(prompt) for prompt in result.prompts]
raise RuntimeError("persistent MCP session has no prompt list transport")
async def close(self) -> None:
"""Close the transport/session stack owned by the runtime factory."""
if self.close_callback is not None:
+12 -1
View File
@@ -8,14 +8,25 @@ from .converters import (
tool_to_discovered,
workflow_output_schema_from_mcp_tool_schema,
)
from .protocols import BackendAdapter, StatefulMcpRuntime, ToolCallResult, ToolExecutor
from .protocols import (
BackendAdapter,
PromptRuntime,
ResourceRuntime,
StatefulMcpRuntime,
ToolCallResult,
ToolExecutor,
ToolRuntime,
)
__all__ = [
"BackendAdapter",
"McpSdkAdapter",
"PromptRuntime",
"ResourceRuntime",
"StatefulMcpRuntime",
"ToolCallResult",
"ToolExecutor",
"ToolRuntime",
"prompt_to_discovered",
"resource_to_discovered",
"tool_result_to_call_result",
+34 -12
View File
@@ -82,13 +82,8 @@ class BackendAdapter(Protocol):
) -> ToolCallResult: ...
class ToolExecutor(Protocol):
"""Runtime boundary for executing MCP tools from workflow nodes.
Discovery can stay one-shot, but workflow execution needs this smaller
protocol so persistent runtime pools can replace one-shot adapters without
changing generated NodeSpecs.
"""
class ToolRuntime(Protocol):
"""Runtime boundary for executing MCP tools from workflow nodes."""
async def call_tool(
self,
@@ -99,12 +94,18 @@ class ToolExecutor(Protocol):
) -> ToolCallResult: ...
class StatefulMcpRuntime(ToolExecutor, Protocol):
"""Stateful execution/read boundary for configured MCP sources.
class ToolExecutor(ToolRuntime, Protocol):
"""Compatibility name for workflow-node tool execution."""
Implementations keep source session state across calls. Discovery/catalog
refresh may still use one-shot adapters by policy.
"""
class ResourceRuntime(Protocol):
"""Stateful resource operations for configured MCP sources."""
async def list_resources(
self,
connection: McpSourceConnection,
auth: AuthRecord | None,
) -> list[DiscoveredResource]: ...
async def read_resource(
self,
@@ -113,6 +114,16 @@ class StatefulMcpRuntime(ToolExecutor, Protocol):
uri: str,
) -> dict[str, Any]: ...
class PromptRuntime(Protocol):
"""Stateful prompt operations for configured MCP sources."""
async def list_prompts(
self,
connection: McpSourceConnection,
auth: AuthRecord | None,
) -> list[DiscoveredPrompt]: ...
async def get_prompt(
self,
connection: McpSourceConnection,
@@ -122,9 +133,20 @@ class StatefulMcpRuntime(ToolExecutor, Protocol):
) -> dict[str, Any]: ...
class StatefulMcpRuntime(ToolRuntime, ResourceRuntime, PromptRuntime, Protocol):
"""Stateful execution/read/list boundary for configured MCP sources.
Implementations keep source session state across calls. Catalog refresh may
still use one-shot adapters by policy.
"""
__all__ = [
"BackendAdapter",
"PromptRuntime",
"ResourceRuntime",
"StatefulMcpRuntime",
"ToolCallResult",
"ToolExecutor",
"ToolRuntime",
]