fix: close run budget and MCP migration gaps

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
-# 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,

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)

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,

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(
-                    {"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")