feat: add stateful runtime listings
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user