735 lines
23 KiB
Python
735 lines
23 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from contextlib import AsyncExitStack
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from mcp import GetPromptResult, ReadResourceResult, ServerResult
|
|
from mcp.client.session import ClientSession
|
|
from mcp.types import CallToolResult as RawCallToolResult
|
|
from mcp.types import (
|
|
ClientNotification,
|
|
ClientRequest,
|
|
ListPromptsResult,
|
|
ListResourcesResult,
|
|
ListToolsResult,
|
|
Prompt,
|
|
Resource,
|
|
TextContent,
|
|
Tool,
|
|
server_result_adapter,
|
|
)
|
|
from pydantic import AnyUrl
|
|
|
|
from wf_sources_mcp.auth import AuthRecord
|
|
from wf_sources_mcp.connections import McpSourceConnection
|
|
from wf_sources_mcp.raw_messages import RawRequest, RawResult
|
|
from wf_sources_mcp.runtime import (
|
|
McpRuntimePool,
|
|
PersistentMcpSession,
|
|
connection_runtime_fingerprint,
|
|
)
|
|
from wf_sources_mcp.runtime.factory import PersistentSessionFactory
|
|
from wf_sources_mcp.sdk import ToolCallResult
|
|
from wf_sources_mcp.transports import StdioSourceTransport
|
|
|
|
|
|
def _connection() -> McpSourceConnection:
|
|
return McpSourceConnection(
|
|
id="demo.personal",
|
|
provider="demo",
|
|
account="personal",
|
|
transport=StdioSourceTransport(command="fake"),
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persistent_session_call_callback_returns_canonical_result() -> None:
|
|
async def call_tool(tool_name: str, payload: dict[str, Any]) -> ToolCallResult:
|
|
assert tool_name == "echo"
|
|
assert payload == {"text": "hi"}
|
|
return ToolCallResult(outcome="ok", output={"echoed": "hi"})
|
|
|
|
session = PersistentMcpSession(
|
|
connection=_connection(),
|
|
auth=AuthRecord(connection_id="demo.personal", scheme="none"),
|
|
call_callback=call_tool,
|
|
)
|
|
|
|
result = await session.call_tool("echo", {"text": "hi"})
|
|
|
|
assert result.outcome == "ok"
|
|
assert result.output == {"echoed": "hi"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persistent_session_raises_without_transport() -> None:
|
|
session = PersistentMcpSession(connection=_connection(), auth=None)
|
|
|
|
with pytest.raises(RuntimeError, match="no tool call transport"):
|
|
await session.call_tool("echo", {})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persistent_session_raises_without_resource_transport() -> None:
|
|
session = PersistentMcpSession(connection=_connection(), auth=None)
|
|
|
|
with pytest.raises(RuntimeError, match="no resource read transport"):
|
|
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):
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, dict[str, object]]] = []
|
|
self.created_connections: list[McpSourceConnection] = []
|
|
|
|
async def _call_tool(
|
|
self, tool_name: str, payload: dict[str, object]
|
|
) -> RawCallToolResult:
|
|
self.calls.append((tool_name, payload))
|
|
return RawCallToolResult(
|
|
content=[TextContent(type="text", text="ok")],
|
|
structured_content={"echoed": payload["text"]},
|
|
)
|
|
|
|
async def _create_with_stack(
|
|
self,
|
|
stack: AsyncExitStack,
|
|
connection: McpSourceConnection,
|
|
auth: AuthRecord | None,
|
|
) -> ClientSession:
|
|
self.created_connections.append(connection)
|
|
factory = self
|
|
|
|
class _FakeClient:
|
|
async def call_tool(
|
|
self, tool_name: str, payload: dict[str, object]
|
|
) -> RawCallToolResult:
|
|
return await factory._call_tool(tool_name, payload)
|
|
|
|
async def read_resource(self, uri: AnyUrl):
|
|
return ReadResourceResult.model_validate(
|
|
{"contents": [{"uri": str(uri), "text": "resource text"}]}
|
|
)
|
|
|
|
async def get_prompt(
|
|
self,
|
|
prompt_name: str,
|
|
arguments: dict[str, str] | None = None,
|
|
) -> GetPromptResult:
|
|
return GetPromptResult.model_validate(
|
|
{
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": {
|
|
"type": "text",
|
|
"text": f"{prompt_name}:{arguments or {}}",
|
|
},
|
|
}
|
|
]
|
|
}
|
|
)
|
|
|
|
async def list_resources(self) -> ListResourcesResult:
|
|
return ListResourcesResult(
|
|
resources=[
|
|
Resource(
|
|
uri=("fixture://docs/runtime"),
|
|
name="resource.runtime",
|
|
title="Runtime Resource",
|
|
description="Runtime-scoped resource.",
|
|
mime_type="text/plain",
|
|
)
|
|
]
|
|
)
|
|
|
|
async def list_prompts(self) -> ListPromptsResult:
|
|
return ListPromptsResult(
|
|
prompts=[
|
|
Prompt(
|
|
name="prompt.runtime",
|
|
title="Runtime Prompt",
|
|
description="Runtime-scoped prompt.",
|
|
arguments=[],
|
|
)
|
|
]
|
|
)
|
|
|
|
async def list_tools(self) -> ListToolsResult:
|
|
return ListToolsResult(
|
|
tools=[
|
|
Tool(
|
|
name="tool.runtime",
|
|
title="Runtime Tool",
|
|
description="Runtime-scoped tool.",
|
|
input_schema={"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[ServerResult],
|
|
) -> ServerResult:
|
|
return server_result_adapter.validate_python(
|
|
{"jsonrpc": "2.0", "id": 1, "result": {}}
|
|
)
|
|
|
|
return _FakeClient() # type: ignore[return-value, ty:invalid-return-type]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persistent_session_factory_serializes_tool_calls() -> None:
|
|
factory = _FakeFactory()
|
|
connection = _connection()
|
|
session = await factory.create(connection, None)
|
|
|
|
first = await session.call_tool("echo", {"text": "one"})
|
|
second = await session.call_tool("echo", {"text": "two"})
|
|
await session.close()
|
|
|
|
assert first.output == {"echoed": "one"}
|
|
assert second.output == {"echoed": "two"}
|
|
assert factory.created_connections == [connection]
|
|
assert factory.calls == [
|
|
("echo", {"text": "one"}),
|
|
("echo", {"text": "two"}),
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_runtime_pool_reuses_unchanged_connection() -> None:
|
|
created: list[McpSourceConnection] = []
|
|
|
|
async def create_session(
|
|
connection: McpSourceConnection, auth: AuthRecord | None
|
|
) -> PersistentMcpSession:
|
|
created.append(connection)
|
|
|
|
async def _call(tool_name: str, payload: dict[str, Any]) -> ToolCallResult:
|
|
return ToolCallResult(outcome="ok", output={"echoed": payload["text"]})
|
|
|
|
return PersistentMcpSession(
|
|
connection=connection,
|
|
auth=auth,
|
|
call_callback=_call,
|
|
)
|
|
|
|
pool = McpRuntimePool(session_factory=create_session)
|
|
connection = _connection()
|
|
|
|
await pool.call_tool(connection, None, "echo", {"text": "one"})
|
|
await pool.call_tool(connection, None, "echo", {"text": "two"})
|
|
|
|
assert created == [connection]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_runtime_pool_serializes_concurrent_session_creation() -> None:
|
|
created: list[McpSourceConnection] = []
|
|
release = asyncio.Event()
|
|
|
|
async def create_session(
|
|
connection: McpSourceConnection, auth: AuthRecord | None
|
|
) -> PersistentMcpSession:
|
|
created.append(connection)
|
|
await release.wait()
|
|
|
|
async def _call(tool_name: str, payload: dict[str, Any]) -> ToolCallResult:
|
|
return ToolCallResult(outcome="ok", output={"echoed": payload["text"]})
|
|
|
|
return PersistentMcpSession(
|
|
connection=connection,
|
|
auth=auth,
|
|
call_callback=_call,
|
|
)
|
|
|
|
pool = McpRuntimePool(session_factory=create_session)
|
|
connection = _connection()
|
|
|
|
first = asyncio.create_task(pool.get_session(connection, None))
|
|
second = asyncio.create_task(pool.get_session(connection, None))
|
|
await asyncio.sleep(0)
|
|
release.set()
|
|
|
|
first_session, second_session = await asyncio.gather(first, second)
|
|
await pool.close_all()
|
|
|
|
assert first_session is second_session
|
|
assert created == [connection]
|
|
|
|
|
|
def test_runtime_fingerprint_changes_when_transport_changes() -> None:
|
|
original = _connection()
|
|
changed = McpSourceConnection(
|
|
id="demo.personal",
|
|
provider="demo",
|
|
account="personal",
|
|
transport=StdioSourceTransport(command="changed"),
|
|
)
|
|
|
|
assert connection_runtime_fingerprint(original) != connection_runtime_fingerprint(
|
|
changed
|
|
)
|
|
|
|
|
|
def test_persistent_session_public_runtime_exposes_safe_read_operations() -> None:
|
|
public_operations = {
|
|
name
|
|
for name in dir(PersistentMcpSession)
|
|
if not name.startswith("_") and callable(getattr(PersistentMcpSession, name))
|
|
}
|
|
|
|
assert "call_tool" in public_operations
|
|
assert "read_resource" in public_operations
|
|
assert "get_prompt" in public_operations
|
|
assert "list_resources" in public_operations
|
|
assert "list_prompts" in public_operations
|
|
assert "list_tools" 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
|
|
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
|
|
async def test_persistent_session_factory_routes_resource_reads_through_owner() -> None:
|
|
factory = _FakeFactory()
|
|
connection = _connection()
|
|
session = await factory.create(connection, None)
|
|
|
|
await session.call_tool("echo", {"text": "one"})
|
|
resource_payload = await session.read_resource("fixture://docs/welcome")
|
|
await session.close()
|
|
|
|
assert factory.created_connections == [connection]
|
|
assert factory.calls == [("echo", {"text": "one"})]
|
|
assert resource_payload["contents"] == [
|
|
{"uri": "fixture://docs/welcome", "text": "resource text"}
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_runtime_pool_reuses_session_for_tool_and_resource_read() -> 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",
|
|
)
|
|
await pool.close_all()
|
|
|
|
assert tool_result.output == {"echoed": "one"}
|
|
assert resource_payload["contents"][0]["text"] == "resource text"
|
|
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]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persistent_session_factory_routes_resource_and_prompt_lists() -> None:
|
|
factory = _FakeFactory()
|
|
session = await factory.create(_connection(), None)
|
|
|
|
resources = await session.list_resources()
|
|
prompts = await session.list_prompts()
|
|
await session.close()
|
|
|
|
assert resources[0].name == "resource.runtime"
|
|
assert resources[0].uri == "fixture://docs/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",
|
|
input_schema={"type": "object"},
|
|
)
|
|
]
|
|
)
|
|
|
|
session = PersistentMcpSession(
|
|
connection=_connection(),
|
|
auth=None,
|
|
client=_MinimalClient(), # type: ignore[arg-type, ty:invalid-argument-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:
|
|
class _MinimalClient:
|
|
async def send_request(
|
|
self,
|
|
request: RawRequest,
|
|
result_type: type[RawResult],
|
|
) -> RawResult:
|
|
assert request.method == "test.method"
|
|
return result_type.model_validate({"extension": True})
|
|
|
|
session = PersistentMcpSession(
|
|
connection=_connection(),
|
|
auth=None,
|
|
client=_MinimalClient(), # type: ignore[arg-type, ty:invalid-argument-type]
|
|
)
|
|
|
|
result = await session.invoke_method("test.method")
|
|
assert result["extension"] is True
|
|
|
|
|
|
@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, ty:invalid-argument-type]
|
|
)
|
|
|
|
await session.send_notification("test.event")
|
|
|
|
assert len(sent) == 1
|
|
assert sent[0].method == "test.event"
|
|
|
|
|
|
@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
|
|
async def test_runtime_pool_reuses_session_for_resource_and_prompt_lists() -> None:
|
|
factory = _FakeFactory()
|
|
pool = McpRuntimePool(factory.create)
|
|
connection = _connection()
|
|
|
|
resources = await pool.list_resources(connection, None)
|
|
prompts = await pool.list_prompts(connection, None)
|
|
await pool.close_all()
|
|
|
|
assert resources[0].name == "resource.runtime"
|
|
assert prompts[0].name == "prompt.runtime"
|
|
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:
|
|
from wf_sources_mcp.sdk import (
|
|
PromptRuntime,
|
|
ResourceRuntime,
|
|
StatefulMcpRuntime,
|
|
ToolRuntime,
|
|
)
|
|
|
|
factory = _FakeFactory()
|
|
pool = McpRuntimePool(factory.create)
|
|
|
|
tool_runtime: ToolRuntime = pool
|
|
resource_runtime: ResourceRuntime = pool
|
|
prompt_runtime: PromptRuntime = pool
|
|
stateful_runtime: StatefulMcpRuntime = pool
|
|
|
|
assert tool_runtime is pool
|
|
assert resource_runtime is pool
|
|
assert prompt_runtime is pool
|
|
assert stateful_runtime is pool
|