wf-mcp reorg/refactor 1

src/wf_mcp/__init__.py

@@ -48,12 +48,14 @@ from .names import (
     parse_namespaced_tool_name,
 )
 from .proxy_validation import ProxyConfigError, validate_transparent_proxy_config
+from .proxy_config import (
+    broker_config_to_fastmcp_config,
+    connection_to_fastmcp_server_config,
+)
 from .service import WfMcpService
 from .store import FileStore, Store
 from .transparent_proxy import (
     TransparentProxyRuntime,
-    broker_config_to_fastmcp_config,
-    connection_to_fastmcp_server_config,
     create_proxy_admin_server,
     create_transparent_proxy_client,
     create_transparent_proxy_server,

src/wf_mcp/proxy_config.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from typing import Any
+
+from fastmcp.mcp_config import MCPConfig
+
+from .models import BrokerConfig, ConnectionConfig
+from .proxy_validation import validate_transparent_proxy_config
+
+
+def connection_to_fastmcp_server_config(
+    connection: ConnectionConfig,
+) -> dict[str, Any]:
+    """Convert one broker connection into FastMCP client config."""
+    metadata = dict(connection.metadata)
+    transport = metadata.get("transport", "stdio")
+    if transport == "streamable_http":
+        metadata["transport"] = "http"
+    if transport == "stdio":
+        return {
+            "command": metadata["command"],
+            "args": list(metadata.get("args", [])),
+            "env": dict(metadata.get("env", {})),
+            "cwd": metadata.get("cwd"),
+            "transport": "stdio",
+            "description": metadata.get("description"),
+        }
+    if transport in {"http", "streamable-http", "sse"}:
+        return {
+            "url": metadata["url"],
+            "transport": transport,
+            "headers": dict(metadata.get("headers", {})),
+            "description": metadata.get("description"),
+        }
+    raise ValueError(f"unsupported MCP transport {transport!r}")
+
+
+def broker_config_to_fastmcp_config(config: BrokerConfig) -> MCPConfig:
+    """Convert broker config into FastMCP's multi-server config object."""
+    validate_transparent_proxy_config(config)
+    return MCPConfig.from_dict(
+        {
+            "mcpServers": {
+                connection.id: connection_to_fastmcp_server_config(connection)
+                for connection in config.connections
+                if connection.enabled
+            }
+        }
+    )

src/wf_mcp/service/__init__.py

@@ -0,0 +1,3 @@
+from .core import WfMcpService
+
+__all__ = ["WfMcpService"]

src/wf_mcp/service/adapters.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+from ..adapters import BackendAdapter
+from ..models import ConnectionConfig
+
+
+def require_adapter(
+    connection: ConnectionConfig,
+    adapters: Mapping[str, BackendAdapter],
+) -> BackendAdapter:
+    """Return the adapter for a connection or raise a useful lookup error."""
+    adapter = adapters.get(connection.server)
+    if adapter is None:
+        raise KeyError(f"no adapter registered for server {connection.server!r}")
+    return adapter

src/wf_mcp/service/core.py

@@ -7,13 +7,13 @@ from typing import Any
 from wf_authoring import NodeSpec, build_async_registry
 from wf_core import NodeUse, Workflow, execute_workflow_async
 
-from .adapters import BackendAdapter
+from ..adapters import BackendAdapter
-from .catalog import CombinedCatalog, snapshot_from_specs
+from ..catalog import CombinedCatalog, snapshot_from_specs
-from .connections import ConnectionRegistry, parse_connection_id, qualify_node_name
+from ..connections import ConnectionRegistry, parse_connection_id, qualify_node_name
-from .discovery import discover_connection_capabilities, specs_from_discovered_tools
+from ..discovery import discover_connection_capabilities, specs_from_discovered_tools
-from .error_info import error_payload
+from ..error_info import error_payload
-from .events import McpEvent, make_event
+from ..events import McpEvent, make_event
-from .models import (
+from ..models import (
     AuthRecord,
     CatalogPromptEntry,
     CatalogResourceEntry,
@@ -21,19 +21,9 @@ from .models import (
     ConnectionConfig,
     RawWorkflowPlan,
 )
-from .store import Store
+from ..store import Store
-
+from .adapters import require_adapter
-
+from .specs import get_qualified_spec, qualify_spec
-def _qualify_spec(connection_id: str, spec: NodeSpec[Any, Any]) -> NodeSpec[Any, Any]:
-    return NodeSpec(
-        name=qualify_node_name(connection_id, spec.name),
-        input_model=spec.input_model,
-        output_model=spec.output_model,
-        outcomes=spec.outcomes,
-        fn=spec.fn,
-        description=spec.description,
-        is_async=spec.is_async,
-    )
 
 
 @dataclass(slots=True)
@@ -82,9 +72,7 @@ class WfMcpService:
     ) -> None:
         self.connections.get(connection_id)
         qualified_specs = {
-            qualify_node_name(connection_id, spec.name): _qualify_spec(
+            qualify_node_name(connection_id, spec.name): qualify_spec(connection_id, spec)
-                connection_id, spec
-            )
             for spec in specs
         }
         self.specs_by_connection[connection_id] = qualified_specs
@@ -180,9 +168,7 @@ class WfMcpService:
     async def read_resource(self, qualified_name: str) -> dict[str, Any]:
         resource = self.get_resource(qualified_name)
         connection = self.connections.get(resource.connection_id)
-        adapter = self.adapters.get(connection.server)
+        adapter = require_adapter(connection, self.adapters)
-        if adapter is None:
-            raise KeyError(f"no adapter registered for server {connection.server!r}")
         auth = self.load_auth(resource.connection_id)
         self._record_event(
             make_event(
@@ -211,9 +197,7 @@ class WfMcpService:
         params: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
         connection = self.connections.get(connection_id)
-        adapter = self.adapters.get(connection.server)
+        adapter = require_adapter(connection, self.adapters)
-        if adapter is None:
-            raise KeyError(f"no adapter registered for server {connection.server!r}")
         auth = self.load_auth(connection_id)
         self._record_event(
             make_event(
@@ -242,9 +226,7 @@ class WfMcpService:
         arguments: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
         connection = self.connections.get(connection_id)
-        adapter = self.adapters.get(connection.server)
+        adapter = require_adapter(connection, self.adapters)
-        if adapter is None:
-            raise KeyError(f"no adapter registered for server {connection.server!r}")
         auth = self.load_auth(connection_id)
         capability_id = qualify_node_name(connection_id, tool_name)
         payload = arguments or {}
@@ -279,9 +261,7 @@ class WfMcpService:
         params: dict[str, Any] | None = None,
     ) -> None:
         connection = self.connections.get(connection_id)
-        adapter = self.adapters.get(connection.server)
+        adapter = require_adapter(connection, self.adapters)
-        if adapter is None:
-            raise KeyError(f"no adapter registered for server {connection.server!r}")
         auth = self.load_auth(connection_id)
         self._record_event(
             make_event(
@@ -309,9 +289,7 @@ class WfMcpService:
     ) -> dict[str, Any]:
         prompt = self.get_prompt(qualified_name)
         connection = self.connections.get(prompt.connection_id)
-        adapter = self.adapters.get(connection.server)
+        adapter = require_adapter(connection, self.adapters)
-        if adapter is None:
-            raise KeyError(f"no adapter registered for server {connection.server!r}")
         auth = self.load_auth(prompt.connection_id)
         self._record_event(
             make_event(
@@ -344,9 +322,7 @@ class WfMcpService:
         max_age_seconds: int | None = None,
     ) -> None:
         connection = self.connections.get(connection_id)
-        adapter = self.adapters.get(connection.server)
+        adapter = require_adapter(connection, self.adapters)
-        if adapter is None:
-            raise KeyError(f"no adapter registered for server {connection.server!r}")
 
         auth = self.load_auth(connection_id)
         self._record_event(
@@ -462,11 +438,7 @@ class WfMcpService:
         return list(self.events)
 
     def _get_qualified_spec(self, qualified_name: str) -> NodeSpec[Any, Any]:
-        connection_id, _ = qualified_name.rsplit(".", 1)
+        return get_qualified_spec(self.specs_by_connection, qualified_name)
-        specs = self.specs_by_connection.get(connection_id)
-        if specs is None or qualified_name not in specs:
-            raise KeyError(f"unknown qualified node {qualified_name!r}")
-        return specs[qualified_name]
 
     def _record_event(self, event: McpEvent) -> None:
         self.events.append(event)

src/wf_mcp/service/specs.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from typing import Any
+
+from wf_authoring import NodeSpec
+
+from ..connections import qualify_node_name
+
+
+def qualify_spec(connection_id: str, spec: NodeSpec[Any, Any]) -> NodeSpec[Any, Any]:
+    """Return a copy of a spec with its node name scoped to a connection."""
+    return NodeSpec(
+        name=qualify_node_name(connection_id, spec.name),
+        input_model=spec.input_model,
+        output_model=spec.output_model,
+        outcomes=spec.outcomes,
+        fn=spec.fn,
+        description=spec.description,
+        is_async=spec.is_async,
+    )
+
+
+def get_qualified_spec(
+    specs_by_connection: dict[str, dict[str, NodeSpec[Any, Any]]],
+    qualified_name: str,
+) -> NodeSpec[Any, Any]:
+    """Resolve a namespaced node spec from the service's connection cache."""
+    connection_id, _ = qualified_name.rsplit(".", 1)
+    specs = specs_by_connection.get(connection_id)
+    if specs is None or qualified_name not in specs:
+        raise KeyError(f"unknown qualified node {qualified_name!r}")
+    return specs[qualified_name]

src/wf_mcp/transparent_proxy.py

@@ -8,15 +8,15 @@ from fastmcp import FastMCP
 from fastmcp.client import Client
 from fastmcp.client.transports.config import MCPConfigTransport
 from fastmcp.client.transports.memory import FastMCPTransport
-from fastmcp.mcp_config import MCPConfig
 from fastmcp.server import create_proxy
 from fastmcp.server.transforms import Namespace, PromptsAsTools, ResourcesAsTools
 from fastmcp.server.transforms.search import BM25SearchTransform
 
 from .config_manager import BrokerConfigManager, ConfigMutationError
-from .models import BrokerConfig, ConnectionConfig
+from .models import BrokerConfig
 from .names import ADMIN_NAMESPACE, is_admin_tool_name, parse_namespaced_tool_name
 from .pagination import paginate_items
+from .proxy_config import broker_config_to_fastmcp_config
 from .proxy_validation import validate_transparent_proxy_config
 
 _ADMIN_TOOL_NAMES = [
@@ -336,45 +336,6 @@ def create_proxy_admin_server(
     return admin
 
 
-def connection_to_fastmcp_server_config(
-    connection: ConnectionConfig,
-) -> dict[str, Any]:
-    metadata = dict(connection.metadata)
-    transport = metadata.get("transport", "stdio")
-    if transport == "streamable_http":
-        metadata["transport"] = "http"
-    if transport == "stdio":
-        return {
-            "command": metadata["command"],
-            "args": list(metadata.get("args", [])),
-            "env": dict(metadata.get("env", {})),
-            "cwd": metadata.get("cwd"),
-            "transport": "stdio",
-            "description": metadata.get("description"),
-        }
-    if transport in {"http", "streamable-http", "sse"}:
-        return {
-            "url": metadata["url"],
-            "transport": transport,
-            "headers": dict(metadata.get("headers", {})),
-            "description": metadata.get("description"),
-        }
-    raise ValueError(f"unsupported MCP transport {transport!r}")
-
-
-def broker_config_to_fastmcp_config(config: BrokerConfig) -> MCPConfig:
-    validate_transparent_proxy_config(config)
-    return MCPConfig.from_dict(
-        {
-            "mcpServers": {
-                connection.id: connection_to_fastmcp_server_config(connection)
-                for connection in config.connections
-                if connection.enabled
-            }
-        }
-    )
-
-
 def create_transparent_proxy_server(
     config: BrokerConfig,
     *,