fix: close run budget and MCP migration gaps

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
     )
 

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_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

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

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)

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