fix type issue
This commit is contained in:
@@ -4,8 +4,11 @@ from collections.abc import Awaitable, Callable
|
|||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
|
from mcp.client.session import ClientSession
|
||||||
|
|
||||||
from ..models import AuthRecord, ConnectionConfig
|
from ..models import AuthRecord, ConnectionConfig
|
||||||
from ..sdk import ToolCallResult
|
from ..sdk import ToolCallResult
|
||||||
|
from ..sdk.converters import tool_result_to_call_result
|
||||||
|
|
||||||
|
|
||||||
@dataclass(slots=True)
|
@dataclass(slots=True)
|
||||||
@@ -14,20 +17,14 @@ class PersistentMcpSession:
|
|||||||
|
|
||||||
connection: ConnectionConfig
|
connection: ConnectionConfig
|
||||||
auth: AuthRecord | None
|
auth: AuthRecord | None
|
||||||
client: Any
|
client: ClientSession
|
||||||
close_callback: Callable[[], Awaitable[None]] | None = None
|
close_callback: Callable[[], Awaitable[None]] | None = None
|
||||||
|
|
||||||
async def call_tool(self, tool_name: str, payload: dict[str, Any]) -> ToolCallResult:
|
async def call_tool(self, tool_name: str, payload: dict[str, Any]) -> ToolCallResult:
|
||||||
return await self.client.call_tool(tool_name, payload)
|
result = await self.client.call_tool(tool_name, payload)
|
||||||
|
return tool_result_to_call_result(result)
|
||||||
|
|
||||||
async def close(self) -> None:
|
async def close(self) -> None:
|
||||||
"""Close this runtime without assuming a specific SDK client shape."""
|
"""Close the transport/session stack owned by the runtime factory."""
|
||||||
if self.close_callback is not None:
|
if self.close_callback is not None:
|
||||||
await self.close_callback()
|
await self.close_callback()
|
||||||
return
|
|
||||||
close = getattr(self.client, "close", None)
|
|
||||||
if close is None:
|
|
||||||
return
|
|
||||||
result = close()
|
|
||||||
if hasattr(result, "__await__"):
|
|
||||||
await result
|
|
||||||
|
|||||||
@@ -2,8 +2,10 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from typing import Any
|
from typing import Any, cast
|
||||||
|
|
||||||
|
from mcp.client.session import ClientSession
|
||||||
|
from mcp.types import CallToolResult
|
||||||
from wf_authoring import build_async_registry
|
from wf_authoring import build_async_registry
|
||||||
from wf_core import RuntimeContext
|
from wf_core import RuntimeContext
|
||||||
from wf_mcp.capabilities import DiscoveredTool
|
from wf_mcp.capabilities import DiscoveredTool
|
||||||
@@ -46,14 +48,26 @@ class FakeStatefulClient:
|
|||||||
closed: bool = False
|
closed: bool = False
|
||||||
calls: list[tuple[str, dict[str, Any]]] = field(default_factory=list)
|
calls: list[tuple[str, dict[str, Any]]] = field(default_factory=list)
|
||||||
|
|
||||||
async def call_tool(self, tool_name: str, payload: dict[str, Any]) -> ToolCallResult:
|
async def call_tool(
|
||||||
|
self, tool_name: str, payload: dict[str, Any]
|
||||||
|
) -> CallToolResult:
|
||||||
self.calls.append((tool_name, payload))
|
self.calls.append((tool_name, payload))
|
||||||
if tool_name == "browser_navigate":
|
if tool_name == "browser_navigate":
|
||||||
self.page_open = True
|
self.page_open = True
|
||||||
return ToolCallResult(outcome="ok", output={"content": "opened"})
|
return CallToolResult(
|
||||||
|
content=[],
|
||||||
|
structuredContent={"content": "opened"},
|
||||||
|
)
|
||||||
if tool_name == "browser_snapshot" and self.page_open:
|
if tool_name == "browser_snapshot" and self.page_open:
|
||||||
return ToolCallResult(outcome="ok", output={"content": "snapshot"})
|
return CallToolResult(
|
||||||
return ToolCallResult(outcome="error", output={"message": "No open page"})
|
content=[],
|
||||||
|
structuredContent={"content": "snapshot"},
|
||||||
|
)
|
||||||
|
return CallToolResult(
|
||||||
|
content=[],
|
||||||
|
structuredContent={"message": "No open page"},
|
||||||
|
isError=True,
|
||||||
|
)
|
||||||
|
|
||||||
async def close(self) -> None:
|
async def close(self) -> None:
|
||||||
self.closed = True
|
self.closed = True
|
||||||
@@ -129,7 +143,12 @@ def test_runtime_pool_reuses_stateful_session_for_same_connection() -> None:
|
|||||||
) -> PersistentMcpSession:
|
) -> PersistentMcpSession:
|
||||||
client = FakeStatefulClient()
|
client = FakeStatefulClient()
|
||||||
created_clients.append(client)
|
created_clients.append(client)
|
||||||
return PersistentMcpSession(connection=connection, auth=auth, client=client)
|
return PersistentMcpSession(
|
||||||
|
connection=connection,
|
||||||
|
auth=auth,
|
||||||
|
client=cast(ClientSession, client),
|
||||||
|
close_callback=client.close,
|
||||||
|
)
|
||||||
|
|
||||||
async def run_calls() -> ToolCallResult:
|
async def run_calls() -> ToolCallResult:
|
||||||
pool = McpRuntimePool(factory)
|
pool = McpRuntimePool(factory)
|
||||||
@@ -164,7 +183,12 @@ def test_runtime_pool_replaces_session_when_fingerprint_changes() -> None:
|
|||||||
) -> PersistentMcpSession:
|
) -> PersistentMcpSession:
|
||||||
client = FakeStatefulClient()
|
client = FakeStatefulClient()
|
||||||
created_clients.append(client)
|
created_clients.append(client)
|
||||||
return PersistentMcpSession(connection=connection, auth=auth, client=client)
|
return PersistentMcpSession(
|
||||||
|
connection=connection,
|
||||||
|
auth=auth,
|
||||||
|
client=cast(ClientSession, client),
|
||||||
|
close_callback=client.close,
|
||||||
|
)
|
||||||
|
|
||||||
async def run_calls() -> None:
|
async def run_calls() -> None:
|
||||||
pool = McpRuntimePool(factory)
|
pool = McpRuntimePool(factory)
|
||||||
|
|||||||
Reference in New Issue
Block a user