fix: close run budget and MCP migration gaps

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`,
+model also represents `end`, `interrupt`, `foreach`, `when`, `choose`, `match`,
-`choose`, `match`, and subgraph steps, and core may gain more step kinds. This
+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.
 

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,

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
+class _FromFrame:
-# activation number explicitly instead, so one activation keeps one number
+    """Sentinel type for resolving a trace step number from its frame."""
-# across its interrupt and resume-completion entries without a second admission.
+
-_FROM_FRAME: Any = object()
+
+# 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,

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,8 +72,12 @@ 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
+        # 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()

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,
+        request: RawRequest,
-        result_type: type[ServerResult] | TypeAdapter[ServerResult],
+        result_type: type[RawResult],
-    ) -> ServerResult: ...
+    ) -> 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(
+            raw_request(method, params),
-                {"method": method, "params": params}
+            RawResult,
-            ),
-            server_result_adapter,
         )
         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(
+        notification = raw_notification(method, params)
-            client_notification_adapter.validate_python(
+        # MCP 2's runtime accepts the generic Notification base class, while
-                {"method": method, "params": params}
+        # its public annotation still names only the standard-method union.
-            )
+        await self.session.send_notification(cast(ClientNotification, notification))
-        )
 
     async def call_tool(
         self,

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",
+]

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(
+                raw_request(method, params),
-                    {"method": method, "params": params}
+                RawResult,
-                ),
-                server_result_adapter,
             )
             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
+            notification = raw_notification(method, params)
-
+            # MCP 2's annotation has not widened to its generic runtime shape.
-            await self.client.send_notification(
+            await self.client.send_notification(cast(ClientNotification, notification))
-                client_notification_adapter.validate_python(
-                    {"method": method, "params": params}
-                )
-            )
             return
         raise RuntimeError("persistent MCP session has no notification send transport")
 

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:
+def _missing_foreach_subgraph(workflow: Workflow, bad: str) -> str:
-    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()
     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)
+    return "nodes[2].input[0].path"
-    assert (
-        _issue(
-            report, ValidationIssueCode.INVALID_CONTEXT_PATH, "nodes[2].input[0].path"
-        )
-        is not None
-    )
 
-    # 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)
+    return "nodes[2].check.path"
-    assert (
-        _issue(report, ValidationIssueCode.INVALID_CONTEXT_PATH, "nodes[2].check.path")
-        is not None
-    )
 
-    # 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)
+    return "nodes[1].over"
-    assert (
-        _issue(report, ValidationIssueCode.INVALID_CONTEXT_PATH, "nodes[1].over")
-        is not None
-    )
 
-    # 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)
+    return "nodes[2].request[0].path"
-    assert (
-        _issue(
-            report,
-            ValidationIssueCode.INVALID_CONTEXT_PATH,
-            "nodes[2].request[0].path",
-        )
-        is not None
-    )
 
-    # 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
     )
 

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(

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

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 == ()
 
 

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,

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,
+            request: RawRequest,
-            result_type: type[ServerResult] | TypeAdapter[ServerResult],
+            result_type: type[RawResult],
-        ) -> ServerResult:
+        ) -> RawResult:
-            return server_result_adapter.validate_python({"tools": []})
+            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")
+    result = await session.invoke_method("test.method")
-    assert result["tools"] == []
+    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

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,
+        request: RawRequest,
-        result_type: type[ServerResult] | TypeAdapter[ServerResult],
+        result_type: type[RawResult],
-    ) -> Any:
+    ) -> RawResult:
-        assert isinstance(result_type, TypeAdapter)
         self.requests.append(request)
-        return type(
+        return result_type.model_validate({"ok": True})
-            "ServerResultModel",
-            (),
-            {"model_dump": lambda _self, **_kwargs: {"ok": True}},
-        )()
 
     async def send_notification(self, notification: ClientNotification) -> None:
         self.notifications.append(notification)

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,
+        request: RawRequest,
-        result_type: type[ServerResult] | TypeAdapter[ServerResult],
+        result_type: type[RawResult],
-    ) -> Any:
+    ) -> RawResult:
-        assert isinstance(result_type, TypeAdapter)
         self.requests.append(request)
-        return type(
+        return result_type.model_validate({"ok": True})
-            "ClientResultModel",
-            (),
-            {"model_dump": lambda _self, **_kwargs: {"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")
+    result = await source_client.invoke_method("test.method", {"value": 1})
-    await source_client.send_notification("notifications/initialized")
 
     assert result == {"ok": True}
-    assert session.requests, "invoke_method should send a request"
+    assert session.requests[0].method == "test.method"
-    assert session.notifications, "send_notification should send a notification"
+
+
+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