fix: close run budget and MCP migration gaps

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,