diff --git a/docs/superpowers/specs/2026-06-27-draft-semantic-authoring-boundary.md b/docs/superpowers/specs/2026-06-27-draft-semantic-authoring-boundary.md index df792de4..734c03ba 100644 --- a/docs/superpowers/specs/2026-06-27-draft-semantic-authoring-boundary.md +++ b/docs/superpowers/specs/2026-06-27-draft-semantic-authoring-boundary.md @@ -85,8 +85,8 @@ semantic operation produces one patch and consumes one revision. clients do not need to know about the internal service split. The service boundary is intentionally not capability-only. The current draft -model also represents `end`, `condition`, `interrupt`, `foreach`, `when`, -`choose`, `match`, and subgraph steps, and core may gain more step kinds. This +model also represents `end`, `interrupt`, `foreach`, `when`, `choose`, `match`, +and subgraph steps, and core may gain more step kinds. This slice adds semantic operations only where required, but new step-kind helpers belong in `WorkflowDraftAuthoringApi` rather than a parallel authoring system. diff --git a/src/wf_api/runs.py b/src/wf_api/runs.py index 46978b7c..1a690d55 100644 --- a/src/wf_api/runs.py +++ b/src/wf_api/runs.py @@ -83,6 +83,9 @@ class WorkflowRunApi: max_steps: int | None = None, ) -> RunResult: trace_values = _trace_range_values(trace_range) + limits = ( + RunLimits(max_steps=max_steps) if max_steps is not None else RunLimits() + ) deployment, artifact, diagnostics, tree = ( self.deployments.deployment_validation(deployment_id) ) @@ -92,12 +95,10 @@ class WorkflowRunApi: artifact=artifact, status="unrunnable", diagnostics=diagnostics, + max_steps=limits.max_steps, ) plan = raw_plan_from_artifact(artifact) - limits = ( - RunLimits(max_steps=max_steps) if max_steps is not None else RunLimits() - ) run = await self.context.runtime.run_workflow_from_plan( plan, workflow_input, diff --git a/src/wf_core/runtime/ops/flow.py b/src/wf_core/runtime/ops/flow.py index 2f7f135e..69a17092 100644 --- a/src/wf_core/runtime/ops/flow.py +++ b/src/wf_core/runtime/ops/flow.py @@ -22,11 +22,14 @@ from wf_core.runtime.scheduler import ( ) from wf_core.tokens import END -# Sentinel for ``append_trace()``: copy the named frame's admitted step number -# (failing closed when unassigned). Interrupt resume passes its stored -# activation number explicitly instead, so one activation keeps one number -# across its interrupt and resume-completion entries without a second admission. -_FROM_FRAME: Any = object() + +class _FromFrame: + """Sentinel type for resolving a trace step number from its frame.""" + + +# Interrupt resume passes its stored activation number explicitly, so one +# activation keeps one number across interrupt and completion trace entries. +_FROM_FRAME = _FromFrame() def append_trace( @@ -40,7 +43,7 @@ def append_trace( next_node_id: str, output: dict[str, Any], state_changes: dict[str, Any], - step_number: int | None | Any = _FROM_FRAME, + step_number: int | None | _FromFrame = _FROM_FRAME, ) -> None: """Append one trace entry carrying its admitted step number. @@ -50,7 +53,7 @@ def append_trace( ``WorkflowExecutionError``. Pass ``step_number`` explicitly only to reuse a persisted activation number (interrupt resume-completion). """ - if step_number is _FROM_FRAME: + if isinstance(step_number, _FromFrame): frame = run.frames.get(frame_id) if frame is None: raise WorkflowExecutionError( @@ -85,7 +88,7 @@ def append_step_result_trace( step_type: str, next_node_id: str, result: StepExecutionResult, - step_number: int | None | Any = _FROM_FRAME, + step_number: int | None | _FromFrame = _FROM_FRAME, ) -> None: append_trace( run, diff --git a/src/wf_mcp/proxy/safe_names.py b/src/wf_mcp/proxy/safe_names.py index 99a7687b..604383b0 100644 --- a/src/wf_mcp/proxy/safe_names.py +++ b/src/wf_mcp/proxy/safe_names.py @@ -2,6 +2,7 @@ from __future__ import annotations import hashlib import re +from asyncio import Lock from collections.abc import Sequence from typing import TYPE_CHECKING, Any @@ -39,6 +40,7 @@ class SafeToolNames(Transform): self._original_to_safe: dict[str, str] = {} self._server = server self._primed = False + self._prime_lock = Lock() async def list_tools(self, tools: Sequence[Tool]) -> Sequence[Tool]: return [ @@ -70,14 +72,18 @@ class SafeToolNames(Transform): them; a failed prime falls through to the identity fallback so the call still ends in a proper unknown-tool error. """ - if self._server is None or self._primed: + if self._server is None: return None - self._primed = True - try: - await self._server.list_tools() - except Exception: - return None - return self._safe_to_original.get(name) + # A direct-call burst must share the first live listing. Publishing + # ``_primed`` without this barrier lets siblings observe empty maps. + async with self._prime_lock: + if not self._primed: + self._primed = True + try: + await self._server.list_tools() + except Exception: + return None + return self._safe_to_original.get(name) def _safe_name(self, original_name: str) -> str: cached = self._original_to_safe.get(original_name) diff --git a/src/wf_sources_mcp/client/source_client.py b/src/wf_sources_mcp/client/source_client.py index d3f9d2f1..e8cdeda4 100644 --- a/src/wf_sources_mcp/client/source_client.py +++ b/src/wf_sources_mcp/client/source_client.py @@ -1,25 +1,25 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Protocol +from typing import TYPE_CHECKING, Any, Protocol, cast from mcp.types import ( CallToolResult, ClientNotification, - ClientRequest, GetPromptResult, ListPromptsResult, ListResourcesResult, ListToolsResult, ReadResourceResult, - ServerResult, - client_notification_adapter, - client_request_adapter, - server_result_adapter, ) -from pydantic import TypeAdapter from wf_sources_mcp.connections import McpSourceConnection +from wf_sources_mcp.raw_messages import ( + RawRequest, + RawResult, + raw_notification, + raw_request, +) if TYPE_CHECKING: from wf_sources_mcp.catalog import ( @@ -58,9 +58,9 @@ class McpClientSession(Protocol): # BaseSession stuff. not even complete signature, thats crazy async def send_request( self, - request: ClientRequest, - result_type: type[ServerResult] | TypeAdapter[ServerResult], - ) -> ServerResult: ... + request: RawRequest, + result_type: type[RawResult], + ) -> RawResult: ... async def send_notification(self, notification: ClientNotification) -> None: ... @@ -138,10 +138,8 @@ class McpSourceClient: params: dict[str, Any] | None = None, ) -> dict[str, Any]: result = await self.session.send_request( - client_request_adapter.validate_python( - {"method": method, "params": params} - ), - server_result_adapter, + raw_request(method, params), + RawResult, ) return result.model_dump(by_alias=True, mode="json", exclude_none=True) @@ -150,11 +148,10 @@ class McpSourceClient: method: str, params: dict[str, Any] | None = None, ) -> None: - await self.session.send_notification( - client_notification_adapter.validate_python( - {"method": method, "params": params} - ) - ) + notification = raw_notification(method, params) + # MCP 2's runtime accepts the generic Notification base class, while + # its public annotation still names only the standard-method union. + await self.session.send_notification(cast(ClientNotification, notification)) async def call_tool( self, diff --git a/src/wf_sources_mcp/raw_messages.py b/src/wf_sources_mcp/raw_messages.py new file mode 100644 index 00000000..b245b33c --- /dev/null +++ b/src/wf_sources_mcp/raw_messages.py @@ -0,0 +1,36 @@ +"""Generic MCP messages for the deliberately untyped extension surface.""" + +from typing import Any + +from mcp.types import Notification, Request +from pydantic import BaseModel, ConfigDict + +type RawParams = dict[str, Any] | None +type RawRequest = Request[RawParams, str] +type RawNotification = Notification[RawParams, str] + + +class RawResult(BaseModel): + """Preserve every field returned by an extension method.""" + + model_config = ConfigDict(extra="allow") + + +def raw_request(method: str, params: RawParams) -> RawRequest: + """Build an extension request without narrowing it to standard methods.""" + return Request[RawParams, str](method=method, params=params) + + +def raw_notification(method: str, params: RawParams) -> RawNotification: + """Build an extension notification without narrowing its method name.""" + return Notification[RawParams, str](method=method, params=params) + + +__all__ = [ + "RawNotification", + "RawParams", + "RawRequest", + "RawResult", + "raw_notification", + "raw_request", +] diff --git a/src/wf_sources_mcp/runtime/session.py b/src/wf_sources_mcp/runtime/session.py index 1675d1ec..4be196f9 100644 --- a/src/wf_sources_mcp/runtime/session.py +++ b/src/wf_sources_mcp/runtime/session.py @@ -2,13 +2,19 @@ from __future__ import annotations from collections.abc import Awaitable, Callable from dataclasses import dataclass -from typing import Any +from typing import Any, cast from mcp.client.session import ClientSession +from mcp.types import ClientNotification from wf_sources_mcp.auth import AuthRecord from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource, DiscoveredTool from wf_sources_mcp.connections import McpSourceConnection +from wf_sources_mcp.raw_messages import ( + RawResult, + raw_notification, + raw_request, +) from wf_sources_mcp.sdk import ToolCallResult from wf_sources_mcp.sdk.converters import tool_result_to_call_result @@ -134,13 +140,9 @@ class PersistentMcpSession: if self.invoke_method_callback is not None: return await self.invoke_method_callback(method, params) if self.client is not None: - from mcp.types import client_request_adapter, server_result_adapter - result = await self.client.send_request( - client_request_adapter.validate_python( - {"method": method, "params": params} - ), - server_result_adapter, + raw_request(method, params), + RawResult, ) return result.model_dump(by_alias=True, mode="json", exclude_none=True) raise RuntimeError("persistent MCP session has no method invoke transport") @@ -155,13 +157,9 @@ class PersistentMcpSession: await self.send_notification_callback(method, params) return if self.client is not None: - from mcp.types import client_notification_adapter - - await self.client.send_notification( - client_notification_adapter.validate_python( - {"method": method, "params": params} - ) - ) + notification = raw_notification(method, params) + # MCP 2's annotation has not widened to its generic runtime shape. + await self.client.send_notification(cast(ClientNotification, notification)) return raise RuntimeError("persistent MCP session has no notification send transport") diff --git a/tests/core/test_structured_context_validation.py b/tests/core/test_structured_context_validation.py index 14afe4ee..70463666 100644 --- a/tests/core/test_structured_context_validation.py +++ b/tests/core/test_structured_context_validation.py @@ -1,9 +1,17 @@ from __future__ import annotations +from collections.abc import Callable + import pytest from wf_core import END, Edge, ForeachNode, NodeDef, NodeUse, SchemaRef, Workflow from wf_core.models.schemas import StateField, StateSchema +from wf_core.models.steps import ( + ConditionNode, + InputPathBinding, + InterruptNode, + SubgraphNode, +) from wf_core.validation.issues import ValidationIssueCode @@ -275,13 +283,7 @@ def test_node_input_surfaces_report_exact_model_paths(make_node, expected_path) ) -def test_all_model_surfaces_reject_missing_foreach_id() -> None: - from wf_core.models.steps import ConditionNode, InterruptNode, SubgraphNode - from wf_core.validation import validate_workflow - - bad = "context.foreach.missing.item" - # Subgraph input - workflow = _base_workflow() +def _missing_foreach_subgraph(workflow: Workflow, bad: str) -> str: workflow.nodes[2] = SubgraphNode.model_validate( { "id": "work", @@ -290,16 +292,10 @@ def test_all_model_surfaces_reject_missing_foreach_id() -> None: "input": [{"target": "order", "path": bad}], } ) - report = validate_workflow(workflow) - assert ( - _issue( - report, ValidationIssueCode.INVALID_CONTEXT_PATH, "nodes[2].input[0].path" - ) - is not None - ) + return "nodes[2].input[0].path" - # Condition check - workflow = _base_workflow() + +def _missing_foreach_condition(workflow: Workflow, bad: str) -> str: workflow.nodes[2] = ConditionNode.model_validate( {"id": "work", "type": "condition", "check": {"op": "exists", "path": bad}} ) @@ -314,25 +310,17 @@ def test_all_model_surfaces_reject_missing_foreach_id() -> None: ), Edge.model_validate({"from": "customers", "outcome": "done", "to": END}), ] - report = validate_workflow(workflow) - assert ( - _issue(report, ValidationIssueCode.INVALID_CONTEXT_PATH, "nodes[2].check.path") - is not None - ) + return "nodes[2].check.path" - # Foreach over - workflow = _base_workflow() + +def _missing_foreach_over(workflow: Workflow, bad: str) -> str: workflow.nodes[1] = ForeachNode.model_validate( {"id": "orders", "type": "foreach", "over": bad, "as": "order"} ) - report = validate_workflow(workflow) - assert ( - _issue(report, ValidationIssueCode.INVALID_CONTEXT_PATH, "nodes[1].over") - is not None - ) + return "nodes[1].over" - # Interrupt request - workflow = _base_workflow() + +def _missing_foreach_interrupt(workflow: Workflow, bad: str) -> str: workflow.nodes[2] = InterruptNode.model_validate( { "id": "work", @@ -351,24 +339,37 @@ def test_all_model_surfaces_reject_missing_foreach_id() -> None: ), Edge.model_validate({"from": "customers", "outcome": "done", "to": END}), ] - report = validate_workflow(workflow) - assert ( - _issue( - report, - ValidationIssueCode.INVALID_CONTEXT_PATH, - "nodes[2].request[0].path", - ) - is not None - ) + return "nodes[2].request[0].path" - # Workflow output - from wf_core.models.steps import InputPathBinding as _IPB + +def _missing_foreach_output(workflow: Workflow, bad: str) -> str: + workflow.output = [ + InputPathBinding.model_validate({"target": "result", "path": bad}) + ] + return "output[0].path" + + +@pytest.mark.parametrize( + "configure_surface", + [ + _missing_foreach_subgraph, + _missing_foreach_condition, + _missing_foreach_over, + _missing_foreach_interrupt, + _missing_foreach_output, + ], + ids=["subgraph", "condition", "foreach", "interrupt", "workflow-output"], +) +def test_all_model_surfaces_reject_missing_foreach_id( + configure_surface: Callable[[Workflow, str], str], +) -> None: + from wf_core.validation import validate_workflow workflow = _base_workflow() - workflow.output = [_IPB.model_validate({"target": "result", "path": bad})] + expected_path = configure_surface(workflow, "context.foreach.missing.item") report = validate_workflow(workflow) assert ( - _issue(report, ValidationIssueCode.INVALID_CONTEXT_PATH, "output[0].path") + _issue(report, ValidationIssueCode.INVALID_CONTEXT_PATH, expected_path) is not None ) diff --git a/tests/wf_api/test_runs.py b/tests/wf_api/test_runs.py index 7783c118..08564cce 100644 --- a/tests/wf_api/test_runs.py +++ b/tests/wf_api/test_runs.py @@ -52,6 +52,24 @@ def _echo_service(root: Path) -> WfMcpService: return service +def _unrunnable_service(root: Path) -> WfMcpService: + artifact_store = FileWorkflowArtifactStore(root) + artifact_store.save_artifact(echo_artifact()) + artifact_store.save_deployment( + WorkflowDeployment( + id="echo.unbound", + artifact_id="echo", + artifact_version=1, + bindings=[], + ) + ) + return WfMcpService( + store=FileStore(root / "mcp"), + artifact_store=artifact_store, + run_store=FileRunStore(root / "mcp"), + ) + + def _interrupt_artifact() -> WorkflowArtifact: return WorkflowArtifact( id="approval", @@ -162,6 +180,36 @@ async def test_run_deployment_rejects_non_positive_max_steps(tmp_path: Path) -> ) +async def test_invalid_budget_wins_over_unrunnable_deployment(tmp_path: Path) -> None: + api = WorkflowRunApi( + context_from_service(_unrunnable_service(tmp_path / "invalid_unrunnable")) + ) + + with pytest.raises(ValueError, match="positive"): + await api.run_deployment( + deployment_id="echo.unbound", + workflow_input={"text": "hello"}, + max_steps=0, + ) + + +async def test_unrunnable_deployment_reports_requested_budget(tmp_path: Path) -> None: + api = WorkflowRunApi( + context_from_service(_unrunnable_service(tmp_path / "requested_unrunnable")) + ) + + result = await api.run_deployment( + deployment_id="echo.unbound", + workflow_input={"text": "hello"}, + max_steps=7, + ) + + assert result["status"] == "unrunnable" + assert result["max_steps"] == 7 + assert result["steps_executed"] == 0 + assert result["steps_remaining"] == 7 + + async def test_inspect_run_reports_effective_budget(tmp_path: Path) -> None: api = WorkflowRunApi(context_from_service(_echo_service(tmp_path / "inspect"))) started = await api.run_deployment( diff --git a/tests/wf_cli/test_app.py b/tests/wf_cli/test_app.py index 5c24ef96..d314cebe 100644 --- a/tests/wf_cli/test_app.py +++ b/tests/wf_cli/test_app.py @@ -1766,8 +1766,10 @@ def test_wf_draft_add_control_command_help_is_type_specific() -> None: interrupt = runner.invoke(app, ["draft", "add", "interrupt", "--help"]) foreach = runner.invoke(app, ["draft", "add", "foreach", "--help"]) end = runner.invoke(app, ["draft", "add", "end", "--help"]) + removed_join = runner.invoke(app, ["draft", "add", "join", "--help"]) assert interrupt.exit_code == foreach.exit_code == end.exit_code == 0 + assert removed_join.exit_code != 0 assert "--request-schema-file" in interrupt.output assert "--resume-schema-file" in interrupt.output assert "--request" in interrupt.output diff --git a/tests/wf_client/test_codec.py b/tests/wf_client/test_codec.py index 15929374..27e32b79 100644 --- a/tests/wf_client/test_codec.py +++ b/tests/wf_client/test_codec.py @@ -152,6 +152,9 @@ def test_decode_run_result_returns_typed_domain_boundary() -> None: assert result.run_id == "run-1" assert result.output == {"result": "hello"} + assert result.max_steps == 10_000 + assert result.steps_executed == 1 + assert result.steps_remaining == 9_999 assert result.diagnostics == () diff --git a/tests/wf_mcp/test_protocol_capabilities.py b/tests/wf_mcp/test_protocol_capabilities.py index a66a8d86..7b63ab31 100644 --- a/tests/wf_mcp/test_protocol_capabilities.py +++ b/tests/wf_mcp/test_protocol_capabilities.py @@ -71,7 +71,7 @@ def test_unified_proxy_initialize_capabilities_reflect_local_surface( capabilities = asyncio.run(inspect_capabilities()) except PermissionError as exc: pytest.skip(f"stdio MCP transport is not permitted in this environment: {exc}") - + # TODO disables most of these + find another way (preferably with _meta) to get these back assert capabilities.tools is not None # assert capabilities.tools.list_changed is True diff --git a/tests/wf_mcp/test_safe_tool_names.py b/tests/wf_mcp/test_safe_tool_names.py index 7817b34e..88f40b78 100644 --- a/tests/wf_mcp/test_safe_tool_names.py +++ b/tests/wf_mcp/test_safe_tool_names.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio +import pytest from fastmcp import FastMCP from wf_mcp.proxy.safe_names import ( @@ -41,6 +42,43 @@ def test_safe_tool_names_hashes_overlength_names() -> None: transform.assert_consistent() +async def test_concurrent_direct_calls_share_cold_start_priming( + monkeypatch: pytest.MonkeyPatch, +) -> None: + server: FastMCP[object] = FastMCP("safe-name-test") + for name in ("demo.one", "demo.two"): + + def handler() -> None: + return None + + server.tool(name=name)(handler) + transform = SafeToolNames(server) + server.add_transform(transform) + + listing_started = asyncio.Event() + release_listing = asyncio.Event() + list_tools = server.list_tools + + async def delayed_list_tools(*, run_middleware: bool = True): + listing_started.set() + await release_listing.wait() + return await list_tools(run_middleware=run_middleware) + + monkeypatch.setattr(server, "list_tools", delayed_list_tools) + first = asyncio.create_task(server.get_tool("demo_one")) + await listing_started.wait() + second = asyncio.create_task(server.get_tool("demo_two")) + await asyncio.sleep(0) + release_listing.set() + + first_tool, second_tool = await asyncio.gather(first, second) + + assert first_tool is not None + assert second_tool is not None + assert first_tool.name == "demo_one" + assert second_tool.name == "demo_two" + + def _server_with_tools( *names: str, transform: SafeToolNames | None = None, diff --git a/tests/wf_sources_mcp/test_runtime.py b/tests/wf_sources_mcp/test_runtime.py index 2267c62d..3a88a170 100644 --- a/tests/wf_sources_mcp/test_runtime.py +++ b/tests/wf_sources_mcp/test_runtime.py @@ -20,10 +20,11 @@ from mcp.types import ( Tool, server_result_adapter, ) -from pydantic import AnyUrl, TypeAdapter +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, @@ -574,10 +575,11 @@ async def test_persistent_session_invoke_method_client_fallback() -> None: class _MinimalClient: async def send_request( self, - request: ClientRequest, - result_type: type[ServerResult] | TypeAdapter[ServerResult], - ) -> ServerResult: - return server_result_adapter.validate_python({"tools": []}) + request: RawRequest, + result_type: type[RawResult], + ) -> RawResult: + assert request.method == "test.method" + return result_type.model_validate({"extension": True}) session = PersistentMcpSession( connection=_connection(), @@ -585,8 +587,8 @@ async def test_persistent_session_invoke_method_client_fallback() -> None: client=_MinimalClient(), # type: ignore[arg-type, ty:invalid-argument-type] ) - result = await session.invoke_method("tools/list") - assert result["tools"] == [] + result = await session.invoke_method("test.method") + assert result["extension"] is True @pytest.mark.asyncio @@ -604,9 +606,10 @@ async def test_persistent_session_send_notification_client_fallback() -> None: client=_MinimalClient(), # type: ignore[arg-type, ty:invalid-argument-type] ) - await session.send_notification("notifications/initialized") + await session.send_notification("test.event") assert len(sent) == 1 + assert sent[0].method == "test.event" @pytest.mark.asyncio diff --git a/tests/wf_sources_mcp/test_sdk_adapter.py b/tests/wf_sources_mcp/test_sdk_adapter.py index 0b7df355..c5c7582e 100644 --- a/tests/wf_sources_mcp/test_sdk_adapter.py +++ b/tests/wf_sources_mcp/test_sdk_adapter.py @@ -5,11 +5,10 @@ from contextlib import asynccontextmanager from typing import Any import pytest -from mcp import GetPromptResult, ReadResourceResult, ServerResult +from mcp import GetPromptResult, ReadResourceResult from mcp.types import ( CallToolResult, ClientNotification, - ClientRequest, ListPromptsResult, ListResourcesResult, ListToolsResult, @@ -18,10 +17,10 @@ from mcp.types import ( TextContent, Tool, ) -from pydantic import TypeAdapter from wf_sources_mcp.client import McpSourceClient from wf_sources_mcp.connections import McpSourceConnection +from wf_sources_mcp.raw_messages import RawRequest, RawResult from wf_sources_mcp.sdk import BackendAdapter, McpSdkAdapter from wf_sources_mcp.transports import StdioSourceTransport @@ -38,7 +37,7 @@ def _connection() -> McpSourceConnection: class _FakeSession: def __init__(self) -> None: self.notifications: list[ClientNotification] = [] - self.requests: list[ClientRequest] = [] + self.requests: list[RawRequest] = [] async def list_tools(self) -> ListToolsResult: return ListToolsResult( @@ -103,16 +102,11 @@ class _FakeSession: async def send_request( self, - request: ClientRequest, - result_type: type[ServerResult] | TypeAdapter[ServerResult], - ) -> Any: - assert isinstance(result_type, TypeAdapter) + request: RawRequest, + result_type: type[RawResult], + ) -> RawResult: self.requests.append(request) - return type( - "ServerResultModel", - (), - {"model_dump": lambda _self, **_kwargs: {"ok": True}}, - )() + return result_type.model_validate({"ok": True}) async def send_notification(self, notification: ClientNotification) -> None: self.notifications.append(notification) diff --git a/tests/wf_sources_mcp/test_source_client.py b/tests/wf_sources_mcp/test_source_client.py index ee74907f..24bee1c5 100644 --- a/tests/wf_sources_mcp/test_source_client.py +++ b/tests/wf_sources_mcp/test_source_client.py @@ -3,11 +3,10 @@ from __future__ import annotations from typing import Any import pytest -from mcp import GetPromptResult, ReadResourceResult, ServerResult +from mcp import GetPromptResult, ReadResourceResult from mcp.types import ( CallToolResult, ClientNotification, - ClientRequest, ListPromptsResult, ListResourcesResult, ListToolsResult, @@ -16,10 +15,10 @@ from mcp.types import ( TextContent, Tool, ) -from pydantic import TypeAdapter from wf_sources_mcp.client import McpSourceClient from wf_sources_mcp.connections import McpSourceConnection +from wf_sources_mcp.raw_messages import RawRequest, RawResult from wf_sources_mcp.transports import StdioSourceTransport @@ -34,7 +33,7 @@ def _connection() -> McpSourceConnection: class _FakeSession: def __init__(self) -> None: - self.requests: list[ClientRequest] = [] + self.requests: list[RawRequest] = [] self.notifications: list[ClientNotification] = [] async def list_tools(self) -> ListToolsResult: @@ -101,16 +100,11 @@ class _FakeSession: async def send_request( self, - request: ClientRequest, - result_type: type[ServerResult] | TypeAdapter[ServerResult], - ) -> Any: - assert isinstance(result_type, TypeAdapter) + request: RawRequest, + result_type: type[RawResult], + ) -> RawResult: self.requests.append(request) - return type( - "ClientResultModel", - (), - {"model_dump": lambda _self, **_kwargs: {"ok": True}}, - )() + return result_type.model_validate({"ok": True}) async def send_notification(self, notification: ClientNotification) -> None: self.notifications.append(notification) @@ -158,16 +152,23 @@ async def test_source_client_reads_resources_and_prompts_as_payloads() -> None: @pytest.mark.asyncio -async def test_source_client_invokes_methods_and_notifications() -> None: +async def test_source_client_invokes_extension_method() -> None: session = _FakeSession() source_client = McpSourceClient(session=session, connection=_connection()) - result = await source_client.invoke_method("ping") - await source_client.send_notification("notifications/initialized") + result = await source_client.invoke_method("test.method", {"value": 1}) assert result == {"ok": True} - assert session.requests, "invoke_method should send a request" - assert session.notifications, "send_notification should send a notification" + assert session.requests[0].method == "test.method" + + +async def test_source_client_sends_extension_notification() -> None: + session = _FakeSession() + source_client = McpSourceClient(session=session, connection=_connection()) + + await source_client.send_notification("test.event", {"value": 1}) + + assert session.notifications[0].method == "test.event" @pytest.mark.asyncio