From 5b901557bd69565328d60037a17a654f48e8cfc2 Mon Sep 17 00:00:00 2001
From: lda <lda@ldlda.com>
Date: Sun, 6 Sep 2026 11:58:06 +0700
Subject: [PATCH] fix: prime safe tool name map on direct calls without prior
 listing

---
 src/wf_mcp/proxy/runtime.py    |  2 +-
 src/wf_mcp/proxy/safe_names.py | 32 ++++++++++++++++++++++++++++++--
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/src/wf_mcp/proxy/runtime.py b/src/wf_mcp/proxy/runtime.py
index ef8b7b58..e6a41bd8 100644
--- a/src/wf_mcp/proxy/runtime.py
+++ b/src/wf_mcp/proxy/runtime.py
@@ -123,7 +123,7 @@ class ProxyRuntime:
             # Keep this outermost so every previous tool projection, including
             # search mode's synthetic tools and always-visible controls, is
             # adapted for clients with stricter name patterns.
-            self.server.add_transform(SafeToolNames())
+            self.server.add_transform(SafeToolNames(server=self.server))
 
     def current_config(self) -> BrokerConfig:
         if self.manager is None:
diff --git a/src/wf_mcp/proxy/safe_names.py b/src/wf_mcp/proxy/safe_names.py
index 1ac4d9db..99a7687b 100644
--- a/src/wf_mcp/proxy/safe_names.py
+++ b/src/wf_mcp/proxy/safe_names.py
@@ -3,11 +3,15 @@ from __future__ import annotations
 import hashlib
 import re
 from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
 
 from fastmcp.server.transforms import GetToolNext, Transform
 from fastmcp.tools.base import Tool
 from fastmcp.utilities.versions import VersionSpec
 
+if TYPE_CHECKING:
+    from fastmcp import FastMCP
+
 _SAFE_TOOL_NAME_PATTERN = re.compile(r"^[a-zA-Z0-9_-]{1,64}$")
 _MAX_TOOL_NAME_LENGTH = 64
 
@@ -30,9 +34,11 @@ class SafeToolNames(Transform):
     readable.
     """
 
-    def __init__(self) -> None:
+    def __init__(self, server: FastMCP[Any] | None = None) -> None:
         self._safe_to_original: dict[str, str] = {}
         self._original_to_safe: dict[str, str] = {}
+        self._server = server
+        self._primed = False
 
     async def list_tools(self, tools: Sequence[Tool]) -> Sequence[Tool]:
         return [
@@ -47,10 +53,32 @@ class SafeToolNames(Transform):
         *,
         version: VersionSpec | None = None,
     ) -> Tool | None:
-        original_name = self._safe_to_original.get(name) or decode_safe_tool_name(name)
+        original_name = self._safe_to_original.get(name)
+        if original_name is None:
+            original_name = await self._prime_lookup(name)
+        if original_name is None:
+            original_name = decode_safe_tool_name(name)
         tool = await call_next(original_name, version=version)
         return None if tool is None else tool.model_copy(update={"name": name})
 
+    async def _prime_lookup(self, name: str) -> str | None:
+        """Populate the reverse map from a live listing on first miss.
+
+        FastMCP v4 resolves `tools/call` through this chain without requiring
+        a prior `tools/list`, so the tables built during listing may be empty
+        when a directly-called safe name arrives. One live listing primes
+        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:
+            return None
+        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)
         if cached is not None: