refactor: route all mcp operations through runtime pool
This commit is contained in:
@@ -260,6 +260,10 @@ implementation state.
|
|||||||
resources, and prompts, and can route session-scoped `list_resources` and
|
resources, and prompts, and can route session-scoped `list_resources` and
|
||||||
`list_prompts` through the owner task. Catalog refresh still uses one-shot
|
`list_prompts` through the owner task. Catalog refresh still uses one-shot
|
||||||
adapter policy.
|
adapter policy.
|
||||||
|
- Completed: persistent MCP runtime now implements the full upstream MCP
|
||||||
|
operation surface used by the one-shot adapter. Broker upstream operations
|
||||||
|
prefer the shared runtime pool when configured, with one-shot adapters
|
||||||
|
retained as fallback.
|
||||||
- Completed: MCP source catalog aggregation helpers (`CombinedCatalog` and
|
- Completed: MCP source catalog aggregation helpers (`CombinedCatalog` and
|
||||||
`snapshot_from_specs`) now live in `wf_sources_mcp.catalog`; the old
|
`snapshot_from_specs`) now live in `wf_sources_mcp.catalog`; the old
|
||||||
`wf_mcp.broker.catalog` path is a compatibility shim.
|
`wf_mcp.broker.catalog` path is a compatibility shim.
|
||||||
|
|||||||
@@ -157,7 +157,11 @@ First slices should move leaf modules only and leave `wf_mcp` re-export shims:
|
|||||||
23. Complete: broker DTO construction removed from `wf_sources_mcp`.
|
23. Complete: broker DTO construction removed from `wf_sources_mcp`.
|
||||||
`wf_sources_mcp` accepts legacy-shaped inputs structurally, while
|
`wf_sources_mcp` accepts legacy-shaped inputs structurally, while
|
||||||
`wf_mcp.source_registry` owns helpers that construct `ConnectionConfig`.
|
`wf_mcp.source_registry` owns helpers that construct `ConnectionConfig`.
|
||||||
24. Upstream transport/discovery/session services.
|
24. Complete: `McpRuntimePool` implements the full MCP operation surface
|
||||||
|
(`list_tools`, resources, prompts, tools, raw methods, notifications, and
|
||||||
|
local metadata). Broker upstream operations prefer the persistent runtime
|
||||||
|
when configured and fall back to one-shot adapters.
|
||||||
|
25. 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`,
|
||||||
|
|||||||
@@ -158,7 +158,9 @@ class UpstreamTransportService:
|
|||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
adapter = require_adapter(connection, self.adapters)
|
adapter = require_adapter(connection, self.adapters)
|
||||||
result = await adapter.get_prompt(source_connection, auth, local_name, arguments)
|
result = await adapter.get_prompt(
|
||||||
|
source_connection, auth, local_name, arguments
|
||||||
|
)
|
||||||
self.event_sink(
|
self.event_sink(
|
||||||
make_event(
|
make_event(
|
||||||
"prompt_get_completed",
|
"prompt_get_completed",
|
||||||
@@ -176,7 +178,6 @@ class UpstreamTransportService:
|
|||||||
*,
|
*,
|
||||||
params: dict[str, Any] | None = None,
|
params: dict[str, Any] | None = None,
|
||||||
) -> dict[str, Any]:
|
) -> dict[str, Any]:
|
||||||
adapter = require_adapter(connection, self.adapters)
|
|
||||||
auth = self.load_connection_auth(connection)
|
auth = self.load_connection_auth(connection)
|
||||||
# Compatibility boundary: broker callers still pass ConnectionConfig.
|
# Compatibility boundary: broker callers still pass ConnectionConfig.
|
||||||
source_connection = mcp_source_connection_from_connection_config(connection)
|
source_connection = mcp_source_connection_from_connection_config(connection)
|
||||||
@@ -188,7 +189,15 @@ class UpstreamTransportService:
|
|||||||
payload={"params": params or {}},
|
payload={"params": params or {}},
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
result = await adapter.invoke_method(source_connection, auth, method, params)
|
if self.stateful_runtime is not None:
|
||||||
|
result = await self.stateful_runtime.invoke_method(
|
||||||
|
source_connection, auth, method, params
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
adapter = require_adapter(connection, self.adapters)
|
||||||
|
result = await adapter.invoke_method(
|
||||||
|
source_connection, auth, method, params
|
||||||
|
)
|
||||||
self.event_sink(
|
self.event_sink(
|
||||||
make_event(
|
make_event(
|
||||||
"raw_method_completed",
|
"raw_method_completed",
|
||||||
@@ -206,7 +215,6 @@ class UpstreamTransportService:
|
|||||||
*,
|
*,
|
||||||
params: dict[str, Any] | None = None,
|
params: dict[str, Any] | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
adapter = require_adapter(connection, self.adapters)
|
|
||||||
auth = self.load_connection_auth(connection)
|
auth = self.load_connection_auth(connection)
|
||||||
# Compatibility boundary: broker callers still pass ConnectionConfig.
|
# Compatibility boundary: broker callers still pass ConnectionConfig.
|
||||||
source_connection = mcp_source_connection_from_connection_config(connection)
|
source_connection = mcp_source_connection_from_connection_config(connection)
|
||||||
@@ -218,7 +226,13 @@ class UpstreamTransportService:
|
|||||||
payload={"params": params or {}},
|
payload={"params": params or {}},
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
await adapter.send_notification(source_connection, auth, method, params)
|
if self.stateful_runtime is not None:
|
||||||
|
await self.stateful_runtime.send_notification(
|
||||||
|
source_connection, auth, method, params
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
adapter = require_adapter(connection, self.adapters)
|
||||||
|
await adapter.send_notification(source_connection, auth, method, params)
|
||||||
self.event_sink(
|
self.event_sink(
|
||||||
make_event(
|
make_event(
|
||||||
"raw_notification_completed",
|
"raw_notification_completed",
|
||||||
@@ -246,12 +260,15 @@ class UpstreamTransportService:
|
|||||||
)
|
)
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
adapter = require_adapter(connection, self.adapters)
|
|
||||||
source_connection = mcp_source_connection_from_connection_config(connection)
|
source_connection = mcp_source_connection_from_connection_config(connection)
|
||||||
|
if self.stateful_runtime is not None:
|
||||||
|
operations: BackendAdapter = self.stateful_runtime
|
||||||
|
else:
|
||||||
|
operations = require_adapter(connection, self.adapters)
|
||||||
capabilities = await discover_connection_capabilities(
|
capabilities = await discover_connection_capabilities(
|
||||||
connection=source_connection,
|
connection=source_connection,
|
||||||
auth=auth,
|
auth=auth,
|
||||||
adapter=adapter,
|
adapter=operations,
|
||||||
)
|
)
|
||||||
specs = specs_from_discovered_tools(
|
specs = specs_from_discovered_tools(
|
||||||
connection=connection,
|
connection=connection,
|
||||||
@@ -351,12 +368,18 @@ class UpstreamTransportService:
|
|||||||
diagnostics.append(auth_diagnostic)
|
diagnostics.append(auth_diagnostic)
|
||||||
continue
|
continue
|
||||||
try:
|
try:
|
||||||
adapter = require_adapter(connection, self.adapters)
|
|
||||||
auth = self.load_connection_auth(connection)
|
auth = self.load_connection_auth(connection)
|
||||||
await asyncio.wait_for(
|
if self.stateful_runtime is not None:
|
||||||
adapter.list_tools(source_connection, auth),
|
await asyncio.wait_for(
|
||||||
timeout=LIVE_SOURCE_CHECK_TIMEOUT_SECONDS,
|
self.stateful_runtime.list_tools(source_connection, auth),
|
||||||
)
|
timeout=LIVE_SOURCE_CHECK_TIMEOUT_SECONDS,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
adapter = require_adapter(connection, self.adapters)
|
||||||
|
await asyncio.wait_for(
|
||||||
|
adapter.list_tools(source_connection, auth),
|
||||||
|
timeout=LIVE_SOURCE_CHECK_TIMEOUT_SECONDS,
|
||||||
|
)
|
||||||
except _LIVE_SOURCE_CHECK_FAILURES as exc:
|
except _LIVE_SOURCE_CHECK_FAILURES as exc:
|
||||||
diagnostics.append(
|
diagnostics.append(
|
||||||
_source_unreachable_diagnostic(
|
_source_unreachable_diagnostic(
|
||||||
|
|||||||
@@ -11,7 +11,7 @@ from typing import Any, Generic, TypeVar
|
|||||||
from mcp.client.session import ClientSession
|
from mcp.client.session import ClientSession
|
||||||
|
|
||||||
from wf_sources_mcp.auth import AuthRecord
|
from wf_sources_mcp.auth import AuthRecord
|
||||||
from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource
|
from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource, DiscoveredTool
|
||||||
from wf_sources_mcp.client import McpSourceClient, open_mcp_session
|
from wf_sources_mcp.client import McpSourceClient, open_mcp_session
|
||||||
from wf_sources_mcp.connections import McpSourceConnection
|
from wf_sources_mcp.connections import McpSourceConnection
|
||||||
from wf_sources_mcp.sdk import ToolCallResult
|
from wf_sources_mcp.sdk import ToolCallResult
|
||||||
@@ -47,6 +47,10 @@ class PersistentSessionFactory:
|
|||||||
get_prompt_callback=owner.get_prompt,
|
get_prompt_callback=owner.get_prompt,
|
||||||
list_resources_callback=owner.list_resources,
|
list_resources_callback=owner.list_resources,
|
||||||
list_prompts_callback=owner.list_prompts,
|
list_prompts_callback=owner.list_prompts,
|
||||||
|
list_tools_callback=owner.list_tools,
|
||||||
|
get_connection_metadata_callback=owner.get_connection_metadata,
|
||||||
|
invoke_method_callback=owner.invoke_method,
|
||||||
|
send_notification_callback=owner.send_notification,
|
||||||
close_callback=owner.close,
|
close_callback=owner.close,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -180,6 +184,43 @@ class _SessionOwner:
|
|||||||
run=lambda client: client.list_prompts(),
|
run=lambda client: client.list_prompts(),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
async def list_tools(self) -> list[DiscoveredTool]:
|
||||||
|
"""Submit tool listing through the generic owner-task operation queue."""
|
||||||
|
return await self.submit(
|
||||||
|
operation="list_tools",
|
||||||
|
run=lambda client: client.list_tools(),
|
||||||
|
)
|
||||||
|
|
||||||
|
async def get_connection_metadata(self) -> dict[str, Any]:
|
||||||
|
"""Return connection metadata computed locally without an upstream call."""
|
||||||
|
transport = self.connection.transport
|
||||||
|
return {
|
||||||
|
"server": self.connection.provider,
|
||||||
|
"transport": transport.kind if transport is not None else None,
|
||||||
|
}
|
||||||
|
|
||||||
|
async def invoke_method(
|
||||||
|
self,
|
||||||
|
method: str,
|
||||||
|
params: dict[str, Any] | None = None,
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
"""Submit a raw method invocation through the generic owner-task operation queue."""
|
||||||
|
return await self.submit(
|
||||||
|
operation="invoke_method",
|
||||||
|
run=lambda client: client.invoke_method(method, params),
|
||||||
|
)
|
||||||
|
|
||||||
|
async def send_notification(
|
||||||
|
self,
|
||||||
|
method: str,
|
||||||
|
params: dict[str, Any] | None = None,
|
||||||
|
) -> None:
|
||||||
|
"""Submit a notification send through the generic owner-task operation queue."""
|
||||||
|
await self.submit(
|
||||||
|
operation="send_notification",
|
||||||
|
run=lambda client: client.send_notification(method, params),
|
||||||
|
)
|
||||||
|
|
||||||
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
|
||||||
|
|||||||
@@ -7,7 +7,7 @@ from inspect import isawaitable
|
|||||||
from typing import Any, cast
|
from typing import Any, cast
|
||||||
|
|
||||||
from wf_sources_mcp.auth import AuthRecord
|
from wf_sources_mcp.auth import AuthRecord
|
||||||
from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource
|
from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource, DiscoveredTool
|
||||||
from wf_sources_mcp.connections import McpSourceConnection
|
from wf_sources_mcp.connections import McpSourceConnection
|
||||||
from wf_sources_mcp.sdk import ToolCallResult
|
from wf_sources_mcp.sdk import ToolCallResult
|
||||||
|
|
||||||
@@ -119,6 +119,42 @@ class McpRuntimePool:
|
|||||||
session = await self.get_session(connection, auth)
|
session = await self.get_session(connection, auth)
|
||||||
return await session.list_prompts()
|
return await session.list_prompts()
|
||||||
|
|
||||||
|
async def list_tools(
|
||||||
|
self,
|
||||||
|
connection: McpSourceConnection,
|
||||||
|
auth: AuthRecord | None,
|
||||||
|
) -> list[DiscoveredTool]:
|
||||||
|
session = await self.get_session(connection, auth)
|
||||||
|
return await session.list_tools()
|
||||||
|
|
||||||
|
async def get_connection_metadata(
|
||||||
|
self,
|
||||||
|
connection: McpSourceConnection,
|
||||||
|
auth: AuthRecord | None,
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
session = await self.get_session(connection, auth)
|
||||||
|
return await session.get_connection_metadata()
|
||||||
|
|
||||||
|
async def invoke_method(
|
||||||
|
self,
|
||||||
|
connection: McpSourceConnection,
|
||||||
|
auth: AuthRecord | None,
|
||||||
|
method: str,
|
||||||
|
params: dict[str, Any] | None = None,
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
session = await self.get_session(connection, auth)
|
||||||
|
return await session.invoke_method(method, params)
|
||||||
|
|
||||||
|
async def send_notification(
|
||||||
|
self,
|
||||||
|
connection: McpSourceConnection,
|
||||||
|
auth: AuthRecord | None,
|
||||||
|
method: str,
|
||||||
|
params: dict[str, Any] | None = None,
|
||||||
|
) -> None:
|
||||||
|
session = await self.get_session(connection, auth)
|
||||||
|
await session.send_notification(method, params)
|
||||||
|
|
||||||
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:
|
||||||
|
|||||||
@@ -8,7 +8,7 @@ from mcp.client.session import ClientSession
|
|||||||
from pydantic import AnyUrl
|
from pydantic import AnyUrl
|
||||||
|
|
||||||
from wf_sources_mcp.auth import AuthRecord
|
from wf_sources_mcp.auth import AuthRecord
|
||||||
from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource
|
from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource, DiscoveredTool
|
||||||
from wf_sources_mcp.connections import McpSourceConnection
|
from wf_sources_mcp.connections import McpSourceConnection
|
||||||
from wf_sources_mcp.sdk import ToolCallResult
|
from wf_sources_mcp.sdk import ToolCallResult
|
||||||
from wf_sources_mcp.sdk.converters import tool_result_to_call_result
|
from wf_sources_mcp.sdk.converters import tool_result_to_call_result
|
||||||
@@ -21,6 +21,10 @@ RawPromptGetter = Callable[
|
|||||||
]
|
]
|
||||||
RawResourceLister = Callable[[], Awaitable[list[DiscoveredResource]]]
|
RawResourceLister = Callable[[], Awaitable[list[DiscoveredResource]]]
|
||||||
RawPromptLister = Callable[[], Awaitable[list[DiscoveredPrompt]]]
|
RawPromptLister = Callable[[], Awaitable[list[DiscoveredPrompt]]]
|
||||||
|
RawToolLister = Callable[[], Awaitable[list[DiscoveredTool]]]
|
||||||
|
RawMetadataGetter = Callable[[], Awaitable[dict[str, Any]]]
|
||||||
|
RawMethodInvoker = Callable[[str, dict[str, Any] | None], Awaitable[dict[str, Any]]]
|
||||||
|
RawNotificationSender = Callable[[str, dict[str, Any] | None], Awaitable[None]]
|
||||||
|
|
||||||
|
|
||||||
@dataclass(slots=True)
|
@dataclass(slots=True)
|
||||||
@@ -41,6 +45,10 @@ class PersistentMcpSession:
|
|||||||
get_prompt_callback: RawPromptGetter | None = None
|
get_prompt_callback: RawPromptGetter | None = None
|
||||||
list_resources_callback: RawResourceLister | None = None
|
list_resources_callback: RawResourceLister | None = None
|
||||||
list_prompts_callback: RawPromptLister | None = None
|
list_prompts_callback: RawPromptLister | None = None
|
||||||
|
list_tools_callback: RawToolLister | None = None
|
||||||
|
get_connection_metadata_callback: RawMetadataGetter | None = None
|
||||||
|
invoke_method_callback: RawMethodInvoker | None = None
|
||||||
|
send_notification_callback: RawNotificationSender | None = None
|
||||||
close_callback: Callable[[], Awaitable[None]] | None = None
|
close_callback: Callable[[], Awaitable[None]] | None = None
|
||||||
|
|
||||||
async def call_tool(
|
async def call_tool(
|
||||||
@@ -97,6 +105,64 @@ class PersistentMcpSession:
|
|||||||
return [prompt_to_discovered(prompt) for prompt in result.prompts]
|
return [prompt_to_discovered(prompt) for prompt in result.prompts]
|
||||||
raise RuntimeError("persistent MCP session has no prompt list transport")
|
raise RuntimeError("persistent MCP session has no prompt list transport")
|
||||||
|
|
||||||
|
async def list_tools(self) -> list[DiscoveredTool]:
|
||||||
|
"""List MCP tools through the owner task or injected session."""
|
||||||
|
if self.list_tools_callback is not None:
|
||||||
|
return await self.list_tools_callback()
|
||||||
|
if self.client is not None:
|
||||||
|
from wf_sources_mcp.sdk.converters import tool_to_discovered
|
||||||
|
|
||||||
|
result = await self.client.list_tools()
|
||||||
|
return [tool_to_discovered(tool) for tool in result.tools]
|
||||||
|
raise RuntimeError("persistent MCP session has no tools list transport")
|
||||||
|
|
||||||
|
async def get_connection_metadata(self) -> dict[str, Any]:
|
||||||
|
"""Return connection metadata from callback or local connection info."""
|
||||||
|
if self.get_connection_metadata_callback is not None:
|
||||||
|
return await self.get_connection_metadata_callback()
|
||||||
|
transport = self.connection.transport
|
||||||
|
return {
|
||||||
|
"server": self.connection.provider,
|
||||||
|
"transport": transport.kind if transport is not None else None,
|
||||||
|
}
|
||||||
|
|
||||||
|
async def invoke_method(
|
||||||
|
self,
|
||||||
|
method: str,
|
||||||
|
params: dict[str, Any] | None = None,
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
"""Invoke a raw MCP method through the owner task or injected session."""
|
||||||
|
if self.invoke_method_callback is not None:
|
||||||
|
return await self.invoke_method_callback(method, params)
|
||||||
|
if self.client is not None:
|
||||||
|
from mcp import ClientResult
|
||||||
|
from mcp.types import ClientRequest
|
||||||
|
|
||||||
|
result = await self.client.send_request(
|
||||||
|
ClientRequest.model_validate({"method": method, "params": params}),
|
||||||
|
ClientResult,
|
||||||
|
)
|
||||||
|
return result.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||||
|
raise RuntimeError("persistent MCP session has no method invoke transport")
|
||||||
|
|
||||||
|
async def send_notification(
|
||||||
|
self,
|
||||||
|
method: str,
|
||||||
|
params: dict[str, Any] | None = None,
|
||||||
|
) -> None:
|
||||||
|
"""Send an MCP notification through the owner task or injected session."""
|
||||||
|
if self.send_notification_callback is not None:
|
||||||
|
await self.send_notification_callback(method, params)
|
||||||
|
return
|
||||||
|
if self.client is not None:
|
||||||
|
from mcp.types import ClientNotification
|
||||||
|
|
||||||
|
await self.client.send_notification(
|
||||||
|
ClientNotification.model_validate({"method": method, "params": params})
|
||||||
|
)
|
||||||
|
return
|
||||||
|
raise RuntimeError("persistent MCP session has no notification send 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:
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ from .converters import (
|
|||||||
)
|
)
|
||||||
from .protocols import (
|
from .protocols import (
|
||||||
BackendAdapter,
|
BackendAdapter,
|
||||||
|
McpSourceOperations,
|
||||||
PromptRuntime,
|
PromptRuntime,
|
||||||
ResourceRuntime,
|
ResourceRuntime,
|
||||||
StatefulMcpRuntime,
|
StatefulMcpRuntime,
|
||||||
@@ -21,6 +22,7 @@ from .protocols import (
|
|||||||
__all__ = [
|
__all__ = [
|
||||||
"BackendAdapter",
|
"BackendAdapter",
|
||||||
"McpSdkAdapter",
|
"McpSdkAdapter",
|
||||||
|
"McpSourceOperations",
|
||||||
"PromptRuntime",
|
"PromptRuntime",
|
||||||
"ResourceRuntime",
|
"ResourceRuntime",
|
||||||
"StatefulMcpRuntime",
|
"StatefulMcpRuntime",
|
||||||
|
|||||||
@@ -17,7 +17,9 @@ class ToolCallResult:
|
|||||||
meta: dict[str, Any] = field(default_factory=dict)
|
meta: dict[str, Any] = field(default_factory=dict)
|
||||||
|
|
||||||
|
|
||||||
class BackendAdapter(Protocol):
|
class McpSourceOperations(Protocol):
|
||||||
|
"""Full MCP operation surface shared by one-shot adapters and persistent runtimes."""
|
||||||
|
|
||||||
async def list_tools(
|
async def list_tools(
|
||||||
self,
|
self,
|
||||||
connection: McpSourceConnection,
|
connection: McpSourceConnection,
|
||||||
@@ -82,6 +84,10 @@ class BackendAdapter(Protocol):
|
|||||||
) -> ToolCallResult: ...
|
) -> ToolCallResult: ...
|
||||||
|
|
||||||
|
|
||||||
|
class BackendAdapter(McpSourceOperations, Protocol):
|
||||||
|
"""One-shot or adapter-style MCP operation executor."""
|
||||||
|
|
||||||
|
|
||||||
class ToolRuntime(Protocol):
|
class ToolRuntime(Protocol):
|
||||||
"""Runtime boundary for executing MCP tools from workflow nodes."""
|
"""Runtime boundary for executing MCP tools from workflow nodes."""
|
||||||
|
|
||||||
@@ -133,8 +139,8 @@ class PromptRuntime(Protocol):
|
|||||||
) -> dict[str, Any]: ...
|
) -> dict[str, Any]: ...
|
||||||
|
|
||||||
|
|
||||||
class StatefulMcpRuntime(ToolRuntime, ResourceRuntime, PromptRuntime, Protocol):
|
class StatefulMcpRuntime(McpSourceOperations, Protocol):
|
||||||
"""Stateful execution/read/list boundary for configured MCP sources.
|
"""Persistent MCP operation executor for configured sources.
|
||||||
|
|
||||||
Implementations keep source session state across calls. Catalog refresh may
|
Implementations keep source session state across calls. Catalog refresh may
|
||||||
still use one-shot adapters by policy.
|
still use one-shot adapters by policy.
|
||||||
@@ -143,6 +149,7 @@ class StatefulMcpRuntime(ToolRuntime, ResourceRuntime, PromptRuntime, Protocol):
|
|||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"BackendAdapter",
|
"BackendAdapter",
|
||||||
|
"McpSourceOperations",
|
||||||
"PromptRuntime",
|
"PromptRuntime",
|
||||||
"ResourceRuntime",
|
"ResourceRuntime",
|
||||||
"StatefulMcpRuntime",
|
"StatefulMcpRuntime",
|
||||||
|
|||||||
@@ -47,11 +47,20 @@ _TRANSPORT_METADATA_KEYS = {
|
|||||||
class LegacyConnectionConfigLike(Protocol):
|
class LegacyConnectionConfigLike(Protocol):
|
||||||
"""Structural shape needed from legacy broker connection configs."""
|
"""Structural shape needed from legacy broker connection configs."""
|
||||||
|
|
||||||
id: str
|
@property
|
||||||
server: str
|
def id(self) -> str: ...
|
||||||
account: str
|
|
||||||
enabled: bool
|
@property
|
||||||
metadata: Mapping[str, object]
|
def server(self) -> str: ...
|
||||||
|
|
||||||
|
@property
|
||||||
|
def account(self) -> str: ...
|
||||||
|
|
||||||
|
@property
|
||||||
|
def enabled(self) -> bool: ...
|
||||||
|
|
||||||
|
@property
|
||||||
|
def metadata(self) -> Mapping[str, object]: ...
|
||||||
|
|
||||||
|
|
||||||
class McpSourceRegistryEntry(SourceRegistryBaseModel):
|
class McpSourceRegistryEntry(SourceRegistryBaseModel):
|
||||||
|
|||||||
@@ -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:
|
async def test_content_access_uses_stateful_runtime_for_upstream_content() -> None:
|
||||||
runtime = _StatefulRuntime()
|
runtime = _StatefulRuntime()
|
||||||
|
|||||||
@@ -11,6 +11,7 @@ from wf_mcp.events import McpEvent
|
|||||||
from wf_mcp.models import AuthRecord, CatalogSnapshot, ConnectionConfig
|
from wf_mcp.models import AuthRecord, CatalogSnapshot, ConnectionConfig
|
||||||
from wf_mcp.storage import FileAuthStore, FileCatalogStore, FileStore
|
from wf_mcp.storage import FileAuthStore, FileCatalogStore, FileStore
|
||||||
from wf_platform import CapabilityBuckets, CapabilitySource, SourcePermissions
|
from wf_platform import CapabilityBuckets, CapabilitySource, SourcePermissions
|
||||||
|
from wf_sources_mcp.catalog import DiscoveredTool
|
||||||
|
|
||||||
from ..test_support import FakeAdapter, local_temp_root
|
from ..test_support import FakeAdapter, local_temp_root
|
||||||
from ..workflow_surface.conftest import echo_artifact
|
from ..workflow_surface.conftest import echo_artifact
|
||||||
@@ -357,6 +358,11 @@ class _StatefulRuntime:
|
|||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
self.resources: list[tuple[str, str]] = []
|
self.resources: list[tuple[str, str]] = []
|
||||||
self.prompts: list[tuple[str, str, dict[str, str] | None]] = []
|
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):
|
async def call_tool(self, connection, auth, tool_name, payload):
|
||||||
raise AssertionError("not used by these tests")
|
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):
|
class _ExplodingContentAdapter(FakeAdapter):
|
||||||
async def read_resource(self, connection, auth, uri):
|
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 result["messages"][0]["content"]["text"] == "stateful prompt"
|
||||||
assert runtime.prompts == [
|
assert runtime.prompts == [("demo.personal", "prompt.summarize", {"text": "hello"})]
|
||||||
("demo.personal", "prompt.summarize", {"text": "hello"})
|
|
||||||
]
|
|
||||||
assert [event.kind for event in events] == [
|
assert [event.kind for event in events] == [
|
||||||
"prompt_get_started",
|
"prompt_get_started",
|
||||||
"prompt_get_completed",
|
"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")]
|
||||||
|
|||||||
@@ -4,14 +4,19 @@ from contextlib import AsyncExitStack
|
|||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
from mcp import ClientResult
|
||||||
from mcp.client.session import ClientSession
|
from mcp.client.session import ClientSession
|
||||||
from mcp.types import CallToolResult as RawCallToolResult
|
from mcp.types import CallToolResult as RawCallToolResult
|
||||||
from mcp.types import (
|
from mcp.types import (
|
||||||
|
ClientNotification,
|
||||||
|
ClientRequest,
|
||||||
ListPromptsResult,
|
ListPromptsResult,
|
||||||
ListResourcesResult,
|
ListResourcesResult,
|
||||||
|
ListToolsResult,
|
||||||
Prompt,
|
Prompt,
|
||||||
Resource,
|
Resource,
|
||||||
TextContent,
|
TextContent,
|
||||||
|
Tool,
|
||||||
)
|
)
|
||||||
from pydantic import AnyUrl
|
from pydantic import AnyUrl
|
||||||
|
|
||||||
@@ -167,6 +172,44 @@ class _FakeFactory(PersistentSessionFactory):
|
|||||||
]
|
]
|
||||||
)
|
)
|
||||||
|
|
||||||
|
async def list_tools(self) -> ListToolsResult:
|
||||||
|
return ListToolsResult(
|
||||||
|
tools=[
|
||||||
|
Tool(
|
||||||
|
name="tool.runtime",
|
||||||
|
title="Runtime Tool",
|
||||||
|
description="Runtime-scoped tool.",
|
||||||
|
inputSchema={"type": "object"},
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
async def get_connection_metadata(self) -> dict[str, object]:
|
||||||
|
return {"server": "demo", "transport": "stdio"}
|
||||||
|
|
||||||
|
async def invoke_method(
|
||||||
|
self,
|
||||||
|
method: str,
|
||||||
|
params: dict[str, object] | None = None,
|
||||||
|
) -> dict[str, object]:
|
||||||
|
return {"echoed": (params or {}).get("text", "")}
|
||||||
|
|
||||||
|
async def send_notification(
|
||||||
|
self,
|
||||||
|
method: str,
|
||||||
|
params: dict[str, object] | None = None,
|
||||||
|
) -> None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
async def send_request(
|
||||||
|
self,
|
||||||
|
request: ClientRequest,
|
||||||
|
result_type: type[ClientResult],
|
||||||
|
) -> ClientResult:
|
||||||
|
return ClientResult.model_validate(
|
||||||
|
{"jsonrpc": "2.0", "id": 1, "result": {}}
|
||||||
|
)
|
||||||
|
|
||||||
return _FakeClient() # type: ignore[return-value]
|
return _FakeClient() # type: ignore[return-value]
|
||||||
|
|
||||||
|
|
||||||
@@ -242,8 +285,10 @@ def test_persistent_session_public_runtime_exposes_safe_read_operations() -> Non
|
|||||||
assert "get_prompt" in public_operations
|
assert "get_prompt" in public_operations
|
||||||
assert "list_resources" in public_operations
|
assert "list_resources" in public_operations
|
||||||
assert "list_prompts" in public_operations
|
assert "list_prompts" in public_operations
|
||||||
assert "invoke_method" not in public_operations
|
assert "list_tools" in public_operations
|
||||||
assert "send_notification" not in public_operations
|
assert "get_connection_metadata" in public_operations
|
||||||
|
assert "invoke_method" in public_operations
|
||||||
|
assert "send_notification" in public_operations
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@@ -347,6 +392,237 @@ async def test_persistent_session_factory_routes_resource_and_prompt_lists() ->
|
|||||||
assert prompts[0].name == "prompt.runtime"
|
assert prompts[0].name == "prompt.runtime"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_persistent_session_factory_routes_list_tools_through_owner() -> None:
|
||||||
|
factory = _FakeFactory()
|
||||||
|
session = await factory.create(_connection(), None)
|
||||||
|
|
||||||
|
tools = await session.list_tools()
|
||||||
|
await session.close()
|
||||||
|
|
||||||
|
assert tools[0].name == "tool.runtime"
|
||||||
|
assert tools[0].description == "Runtime-scoped tool."
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_persistent_session_factory_routes_metadata_through_owner() -> None:
|
||||||
|
factory = _FakeFactory()
|
||||||
|
session = await factory.create(_connection(), None)
|
||||||
|
|
||||||
|
metadata = await session.get_connection_metadata()
|
||||||
|
await session.close()
|
||||||
|
|
||||||
|
assert metadata["server"] == "demo"
|
||||||
|
assert metadata["transport"] == "stdio"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_persistent_session_factory_routes_invoke_method_through_owner() -> None:
|
||||||
|
factory = _FakeFactory()
|
||||||
|
session = await factory.create(_connection(), None)
|
||||||
|
|
||||||
|
result = await session.invoke_method("ping")
|
||||||
|
await session.close()
|
||||||
|
|
||||||
|
assert isinstance(result, dict)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_persistent_session_factory_routes_send_notification_through_owner() -> (
|
||||||
|
None
|
||||||
|
):
|
||||||
|
factory = _FakeFactory()
|
||||||
|
session = await factory.create(_connection(), None)
|
||||||
|
|
||||||
|
await session.send_notification("notifications/initialized")
|
||||||
|
await session.close()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_persistent_session_list_tools_callback() -> None:
|
||||||
|
async def list_tools_cb() -> list:
|
||||||
|
from wf_sources_mcp.catalog import DiscoveredTool
|
||||||
|
|
||||||
|
return [
|
||||||
|
DiscoveredTool(
|
||||||
|
name="cb_tool",
|
||||||
|
title=None,
|
||||||
|
description="Callback tool",
|
||||||
|
input_schema={"type": "object"},
|
||||||
|
output_schema={"type": "object"},
|
||||||
|
)
|
||||||
|
]
|
||||||
|
|
||||||
|
session = PersistentMcpSession(
|
||||||
|
connection=_connection(),
|
||||||
|
auth=None,
|
||||||
|
list_tools_callback=list_tools_cb,
|
||||||
|
)
|
||||||
|
|
||||||
|
tools = await session.list_tools()
|
||||||
|
|
||||||
|
assert tools[0].name == "cb_tool"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_persistent_session_get_connection_metadata_callback() -> None:
|
||||||
|
async def metadata_cb() -> dict[str, object]:
|
||||||
|
return {"server": "cb_server", "transport": "http"}
|
||||||
|
|
||||||
|
session = PersistentMcpSession(
|
||||||
|
connection=_connection(),
|
||||||
|
auth=None,
|
||||||
|
get_connection_metadata_callback=metadata_cb,
|
||||||
|
)
|
||||||
|
|
||||||
|
metadata = await session.get_connection_metadata()
|
||||||
|
|
||||||
|
assert metadata["server"] == "cb_server"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_persistent_session_invoke_method_callback() -> None:
|
||||||
|
async def invoke_cb(
|
||||||
|
method: str, params: dict[str, object] | None
|
||||||
|
) -> dict[str, object]:
|
||||||
|
return {"method": method, "params": params}
|
||||||
|
|
||||||
|
session = PersistentMcpSession(
|
||||||
|
connection=_connection(),
|
||||||
|
auth=None,
|
||||||
|
invoke_method_callback=invoke_cb,
|
||||||
|
)
|
||||||
|
|
||||||
|
result = await session.invoke_method("test.method", {"key": "val"})
|
||||||
|
|
||||||
|
assert result["method"] == "test.method"
|
||||||
|
assert result["params"] == {"key": "val"}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_persistent_session_send_notification_callback() -> None:
|
||||||
|
sent: list[tuple[str, dict[str, object] | None]] = []
|
||||||
|
|
||||||
|
async def notify_cb(method: str, params: dict[str, object] | None) -> None:
|
||||||
|
sent.append((method, params))
|
||||||
|
|
||||||
|
session = PersistentMcpSession(
|
||||||
|
connection=_connection(),
|
||||||
|
auth=None,
|
||||||
|
send_notification_callback=notify_cb,
|
||||||
|
)
|
||||||
|
|
||||||
|
await session.send_notification("test.event", {"data": 1})
|
||||||
|
|
||||||
|
assert sent == [("test.event", {"data": 1})]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_persistent_session_list_tools_client_fallback() -> None:
|
||||||
|
from mcp.types import ListToolsResult, Tool
|
||||||
|
|
||||||
|
class _MinimalClient:
|
||||||
|
async def list_tools(self) -> ListToolsResult:
|
||||||
|
return ListToolsResult(
|
||||||
|
tools=[
|
||||||
|
Tool(
|
||||||
|
name="client_tool",
|
||||||
|
description="Client tool",
|
||||||
|
inputSchema={"type": "object"},
|
||||||
|
)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
session = PersistentMcpSession(
|
||||||
|
connection=_connection(),
|
||||||
|
auth=None,
|
||||||
|
client=_MinimalClient(), # type: ignore[arg-type]
|
||||||
|
)
|
||||||
|
|
||||||
|
tools = await session.list_tools()
|
||||||
|
|
||||||
|
assert tools[0].name == "client_tool"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_persistent_session_invoke_method_client_fallback() -> None:
|
||||||
|
from mcp import ClientResult
|
||||||
|
|
||||||
|
class _MinimalClient:
|
||||||
|
async def send_request(
|
||||||
|
self,
|
||||||
|
request: ClientRequest,
|
||||||
|
result_type: type[ClientResult],
|
||||||
|
) -> ClientResult:
|
||||||
|
return ClientResult.model_validate(
|
||||||
|
{"jsonrpc": "2.0", "id": 1, "result": {"tools": []}}
|
||||||
|
)
|
||||||
|
|
||||||
|
session = PersistentMcpSession(
|
||||||
|
connection=_connection(),
|
||||||
|
auth=None,
|
||||||
|
client=_MinimalClient(), # type: ignore[arg-type]
|
||||||
|
)
|
||||||
|
|
||||||
|
result = await session.invoke_method("tools/list")
|
||||||
|
|
||||||
|
assert result["result"]["tools"] == []
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_persistent_session_send_notification_client_fallback() -> None:
|
||||||
|
|
||||||
|
sent: list[ClientNotification] = []
|
||||||
|
|
||||||
|
class _MinimalClient:
|
||||||
|
async def send_notification(self, notification: ClientNotification) -> None:
|
||||||
|
sent.append(notification)
|
||||||
|
|
||||||
|
session = PersistentMcpSession(
|
||||||
|
connection=_connection(),
|
||||||
|
auth=None,
|
||||||
|
client=_MinimalClient(), # type: ignore[arg-type]
|
||||||
|
)
|
||||||
|
|
||||||
|
await session.send_notification("notifications/initialized")
|
||||||
|
|
||||||
|
assert len(sent) == 1
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_persistent_session_raises_without_tools_transport() -> None:
|
||||||
|
session = PersistentMcpSession(connection=_connection(), auth=None)
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="no tools list transport"):
|
||||||
|
await session.list_tools()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_persistent_session_raises_without_invoke_transport() -> None:
|
||||||
|
session = PersistentMcpSession(connection=_connection(), auth=None)
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="no method invoke transport"):
|
||||||
|
await session.invoke_method("test.ping")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_persistent_session_raises_without_notification_transport() -> None:
|
||||||
|
session = PersistentMcpSession(connection=_connection(), auth=None)
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="no notification send transport"):
|
||||||
|
await session.send_notification("test.event")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_persistent_session_get_connection_metadata_local_fallback() -> None:
|
||||||
|
session = PersistentMcpSession(connection=_connection(), auth=None)
|
||||||
|
|
||||||
|
metadata = await session.get_connection_metadata()
|
||||||
|
|
||||||
|
assert metadata["server"] == "demo"
|
||||||
|
assert metadata["transport"] == "stdio"
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_runtime_pool_reuses_session_for_resource_and_prompt_lists() -> None:
|
async def test_runtime_pool_reuses_session_for_resource_and_prompt_lists() -> None:
|
||||||
factory = _FakeFactory()
|
factory = _FakeFactory()
|
||||||
@@ -362,6 +638,57 @@ async def test_runtime_pool_reuses_session_for_resource_and_prompt_lists() -> No
|
|||||||
assert factory.created_connections == [connection]
|
assert factory.created_connections == [connection]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_runtime_pool_reuses_session_for_list_tools() -> None:
|
||||||
|
factory = _FakeFactory()
|
||||||
|
pool = McpRuntimePool(factory.create)
|
||||||
|
connection = _connection()
|
||||||
|
|
||||||
|
tools = await pool.list_tools(connection, None)
|
||||||
|
await pool.close_all()
|
||||||
|
|
||||||
|
assert tools[0].name == "tool.runtime"
|
||||||
|
assert factory.created_connections == [connection]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_runtime_pool_reuses_session_for_invoke_method() -> None:
|
||||||
|
factory = _FakeFactory()
|
||||||
|
pool = McpRuntimePool(factory.create)
|
||||||
|
connection = _connection()
|
||||||
|
|
||||||
|
result = await pool.invoke_method(connection, None, "ping")
|
||||||
|
await pool.close_all()
|
||||||
|
|
||||||
|
assert isinstance(result, dict)
|
||||||
|
assert factory.created_connections == [connection]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_runtime_pool_reuses_session_for_send_notification() -> None:
|
||||||
|
factory = _FakeFactory()
|
||||||
|
pool = McpRuntimePool(factory.create)
|
||||||
|
connection = _connection()
|
||||||
|
|
||||||
|
await pool.send_notification(connection, None, "notifications/initialized")
|
||||||
|
await pool.close_all()
|
||||||
|
|
||||||
|
assert factory.created_connections == [connection]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_runtime_pool_reuses_session_for_metadata() -> None:
|
||||||
|
factory = _FakeFactory()
|
||||||
|
pool = McpRuntimePool(factory.create)
|
||||||
|
connection = _connection()
|
||||||
|
|
||||||
|
metadata = await pool.get_connection_metadata(connection, None)
|
||||||
|
await pool.close_all()
|
||||||
|
|
||||||
|
assert metadata["server"] == "demo"
|
||||||
|
assert factory.created_connections == [connection]
|
||||||
|
|
||||||
|
|
||||||
def test_runtime_pool_satisfies_stateful_protocol_static_shape() -> None:
|
def test_runtime_pool_satisfies_stateful_protocol_static_shape() -> None:
|
||||||
from wf_sources_mcp.sdk import (
|
from wf_sources_mcp.sdk import (
|
||||||
PromptRuntime,
|
PromptRuntime,
|
||||||
|
|||||||
@@ -4,7 +4,7 @@ from dataclasses import is_dataclass
|
|||||||
from typing import cast
|
from typing import cast
|
||||||
|
|
||||||
from wf_sources_mcp.auth import AuthRecord
|
from wf_sources_mcp.auth import AuthRecord
|
||||||
from wf_sources_mcp.catalog import DiscoveredTool
|
from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource, DiscoveredTool
|
||||||
from wf_sources_mcp.connections import McpSourceConnection
|
from wf_sources_mcp.connections import McpSourceConnection
|
||||||
from wf_sources_mcp.sdk import BackendAdapter, ToolCallResult, ToolExecutor
|
from wf_sources_mcp.sdk import BackendAdapter, ToolCallResult, ToolExecutor
|
||||||
from wf_sources_mcp.transports import StdioSourceTransport
|
from wf_sources_mcp.transports import StdioSourceTransport
|
||||||
@@ -36,6 +36,90 @@ class EchoAdapter:
|
|||||||
return ToolCallResult(outcome="ok", output={"echoed": payload})
|
return ToolCallResult(outcome="ok", output={"echoed": payload})
|
||||||
|
|
||||||
|
|
||||||
|
class _FullSurfaceAdapter:
|
||||||
|
"""Implements every MCP operation for protocol conformance tests."""
|
||||||
|
|
||||||
|
async def list_tools(
|
||||||
|
self,
|
||||||
|
connection: McpSourceConnection,
|
||||||
|
auth: AuthRecord | None,
|
||||||
|
) -> list[DiscoveredTool]:
|
||||||
|
return [
|
||||||
|
DiscoveredTool(
|
||||||
|
name="echo",
|
||||||
|
title=None,
|
||||||
|
description="Echo",
|
||||||
|
input_schema={"type": "object"},
|
||||||
|
output_schema={"type": "object"},
|
||||||
|
)
|
||||||
|
]
|
||||||
|
|
||||||
|
async def list_resources(
|
||||||
|
self,
|
||||||
|
connection: McpSourceConnection,
|
||||||
|
auth: AuthRecord | None,
|
||||||
|
) -> list[DiscoveredResource]:
|
||||||
|
return []
|
||||||
|
|
||||||
|
async def list_prompts(
|
||||||
|
self,
|
||||||
|
connection: McpSourceConnection,
|
||||||
|
auth: AuthRecord | None,
|
||||||
|
) -> list[DiscoveredPrompt]:
|
||||||
|
return []
|
||||||
|
|
||||||
|
async def get_connection_metadata(
|
||||||
|
self,
|
||||||
|
connection: McpSourceConnection,
|
||||||
|
auth: AuthRecord | None,
|
||||||
|
) -> dict[str, object]:
|
||||||
|
return {"server": "demo"}
|
||||||
|
|
||||||
|
async def read_resource(
|
||||||
|
self,
|
||||||
|
connection: McpSourceConnection,
|
||||||
|
auth: AuthRecord | None,
|
||||||
|
uri: str,
|
||||||
|
) -> dict[str, object]:
|
||||||
|
return {"contents": []}
|
||||||
|
|
||||||
|
async def get_prompt(
|
||||||
|
self,
|
||||||
|
connection: McpSourceConnection,
|
||||||
|
auth: AuthRecord | None,
|
||||||
|
prompt_name: str,
|
||||||
|
arguments: dict[str, str] | None = None,
|
||||||
|
) -> dict[str, object]:
|
||||||
|
return {"messages": []}
|
||||||
|
|
||||||
|
async def invoke_method(
|
||||||
|
self,
|
||||||
|
connection: McpSourceConnection,
|
||||||
|
auth: AuthRecord | None,
|
||||||
|
method: str,
|
||||||
|
params: dict[str, object] | None = None,
|
||||||
|
) -> dict[str, object]:
|
||||||
|
return {}
|
||||||
|
|
||||||
|
async def send_notification(
|
||||||
|
self,
|
||||||
|
connection: McpSourceConnection,
|
||||||
|
auth: AuthRecord | None,
|
||||||
|
method: str,
|
||||||
|
params: dict[str, object] | None = None,
|
||||||
|
) -> None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
async def call_tool(
|
||||||
|
self,
|
||||||
|
connection: McpSourceConnection,
|
||||||
|
auth: AuthRecord | None,
|
||||||
|
tool_name: str,
|
||||||
|
payload: dict[str, object],
|
||||||
|
) -> ToolCallResult:
|
||||||
|
return ToolCallResult(outcome="ok", output={"echoed": payload})
|
||||||
|
|
||||||
|
|
||||||
def test_tool_call_result_is_slots_dataclass_with_empty_defaults() -> None:
|
def test_tool_call_result_is_slots_dataclass_with_empty_defaults() -> None:
|
||||||
result = ToolCallResult(outcome="ok")
|
result = ToolCallResult(outcome="ok")
|
||||||
|
|
||||||
@@ -77,6 +161,72 @@ async def test_tool_executor_protocol_can_describe_tool_calls() -> None:
|
|||||||
assert result.output == {"echoed": {"message": "hello"}}
|
assert result.output == {"echoed": {"message": "hello"}}
|
||||||
|
|
||||||
|
|
||||||
|
async def test_backend_adapter_protocol_full_operation_surface() -> None:
|
||||||
|
adapter = cast(BackendAdapter, _FullSurfaceAdapter())
|
||||||
|
conn = McpSourceConnection(
|
||||||
|
id="demo.default",
|
||||||
|
provider="demo",
|
||||||
|
account="default",
|
||||||
|
transport=StdioSourceTransport(command="echo"),
|
||||||
|
)
|
||||||
|
|
||||||
|
tools = await adapter.list_tools(conn, None)
|
||||||
|
resources = await adapter.list_resources(conn, None)
|
||||||
|
prompts = await adapter.list_prompts(conn, None)
|
||||||
|
metadata = await adapter.get_connection_metadata(conn, None)
|
||||||
|
read_result = await adapter.read_resource(conn, None, "test://x")
|
||||||
|
prompt_result = await adapter.get_prompt(conn, None, "prompt.summarize")
|
||||||
|
invoke_result = await adapter.invoke_method(conn, None, "ping")
|
||||||
|
await adapter.send_notification(conn, None, "test.notify")
|
||||||
|
call_result = await adapter.call_tool(conn, None, "echo", {"text": "hi"})
|
||||||
|
|
||||||
|
assert tools[0].name == "echo"
|
||||||
|
assert resources == []
|
||||||
|
assert prompts == []
|
||||||
|
assert metadata["server"] == "demo"
|
||||||
|
assert read_result == {"contents": []}
|
||||||
|
assert prompt_result == {"messages": []}
|
||||||
|
assert invoke_result == {}
|
||||||
|
assert call_result.outcome == "ok"
|
||||||
|
|
||||||
|
|
||||||
|
async def test_stateful_mcp_runtime_protocol_full_operation_surface() -> None:
|
||||||
|
from wf_sources_mcp.sdk import StatefulMcpRuntime
|
||||||
|
|
||||||
|
runtime = cast(StatefulMcpRuntime, _FullSurfaceAdapter())
|
||||||
|
conn = McpSourceConnection(
|
||||||
|
id="demo.default",
|
||||||
|
provider="demo",
|
||||||
|
account="default",
|
||||||
|
transport=StdioSourceTransport(command="echo"),
|
||||||
|
)
|
||||||
|
|
||||||
|
tools = await runtime.list_tools(conn, None)
|
||||||
|
resources = await runtime.list_resources(conn, None)
|
||||||
|
prompts = await runtime.list_prompts(conn, None)
|
||||||
|
metadata = await runtime.get_connection_metadata(conn, None)
|
||||||
|
read_result = await runtime.read_resource(conn, None, "test://x")
|
||||||
|
prompt_result = await runtime.get_prompt(conn, None, "prompt.summarize")
|
||||||
|
invoke_result = await runtime.invoke_method(conn, None, "ping")
|
||||||
|
await runtime.send_notification(conn, None, "test.notify")
|
||||||
|
call_result = await runtime.call_tool(conn, None, "echo", {"text": "hi"})
|
||||||
|
|
||||||
|
assert tools[0].name == "echo"
|
||||||
|
assert resources == []
|
||||||
|
assert prompts == []
|
||||||
|
assert metadata["server"] == "demo"
|
||||||
|
assert read_result == {"contents": []}
|
||||||
|
assert prompt_result == {"messages": []}
|
||||||
|
assert invoke_result == {}
|
||||||
|
assert call_result.outcome == "ok"
|
||||||
|
|
||||||
|
|
||||||
|
def test_mcp_source_operations_protocol_shape() -> None:
|
||||||
|
from wf_sources_mcp.sdk import McpSourceOperations
|
||||||
|
|
||||||
|
assert McpSourceOperations.__name__ == "McpSourceOperations"
|
||||||
|
|
||||||
|
|
||||||
def test_stateful_mcp_runtime_protocol_shape() -> None:
|
def test_stateful_mcp_runtime_protocol_shape() -> None:
|
||||||
from wf_sources_mcp.sdk import StatefulMcpRuntime, ToolExecutor
|
from wf_sources_mcp.sdk import StatefulMcpRuntime, ToolExecutor
|
||||||
|
|
||||||
@@ -86,6 +236,7 @@ def test_stateful_mcp_runtime_protocol_shape() -> None:
|
|||||||
|
|
||||||
def test_stateful_runtime_protocol_slices_export() -> None:
|
def test_stateful_runtime_protocol_slices_export() -> None:
|
||||||
from wf_sources_mcp.sdk import (
|
from wf_sources_mcp.sdk import (
|
||||||
|
McpSourceOperations,
|
||||||
PromptRuntime,
|
PromptRuntime,
|
||||||
ResourceRuntime,
|
ResourceRuntime,
|
||||||
StatefulMcpRuntime,
|
StatefulMcpRuntime,
|
||||||
@@ -97,4 +248,5 @@ def test_stateful_runtime_protocol_slices_export() -> None:
|
|||||||
assert ResourceRuntime.__name__ == "ResourceRuntime"
|
assert ResourceRuntime.__name__ == "ResourceRuntime"
|
||||||
assert PromptRuntime.__name__ == "PromptRuntime"
|
assert PromptRuntime.__name__ == "PromptRuntime"
|
||||||
assert ToolExecutor.__name__ == "ToolExecutor"
|
assert ToolExecutor.__name__ == "ToolExecutor"
|
||||||
|
assert McpSourceOperations.__name__ == "McpSourceOperations"
|
||||||
assert StatefulMcpRuntime.__name__ == "StatefulMcpRuntime"
|
assert StatefulMcpRuntime.__name__ == "StatefulMcpRuntime"
|
||||||
|
|||||||
Reference in New Issue
Block a user