discovery

wf_mcp/__init__.py

@@ -7,6 +7,11 @@ from .adapters import (
 )
 from .catalog import CombinedCatalog
 from .connections import ConnectionRegistry, parse_connection_id, qualify_node_name
+from .discovery import (
+    DiscoveredConnectionCapabilities,
+    discover_connection_capabilities,
+    specs_from_discovered_tools,
+)
 from .models import (
     AuthRecord,
     CatalogNodeEntry,
@@ -31,6 +36,7 @@ __all__ = [
     "CombinedCatalog",
     "ConnectionConfig",
     "ConnectionRegistry",
+    "DiscoveredConnectionCapabilities",
     "DiscoveredPrompt",
     "DiscoveredResource",
     "DiscoveredTool",
@@ -40,7 +46,9 @@ __all__ = [
     "Store",
     "ToolCallResult",
     "WfMcpService",
+    "discover_connection_capabilities",
     "parse_connection_id",
     "qualify_node_name",
+    "specs_from_discovered_tools",
     "wrap_discovered_tool",
 ]

wf_mcp/discovery.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from wf_authoring import NodeSpec
+
+from .adapters import (
+    BackendAdapter,
+    DiscoveredPrompt,
+    DiscoveredResource,
+    DiscoveredTool,
+)
+from .models import AuthRecord, ConnectionConfig
+from .wrappers import wrap_discovered_tool
+
+
+@dataclass(slots=True)
+class DiscoveredConnectionCapabilities:
+    tools: list[DiscoveredTool] = field(default_factory=list)
+    resources: list[DiscoveredResource] = field(default_factory=list)
+    prompts: list[DiscoveredPrompt] = field(default_factory=list)
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+async def discover_connection_capabilities(
+    *,
+    connection: ConnectionConfig,
+    auth: AuthRecord | None,
+    adapter: BackendAdapter,
+) -> DiscoveredConnectionCapabilities:
+    tools = await adapter.list_tools(connection, auth)
+    resources = await adapter.list_resources(connection, auth)
+    prompts = await adapter.list_prompts(connection, auth)
+    metadata = await adapter.get_connection_metadata(connection, auth)
+    return DiscoveredConnectionCapabilities(
+        tools=tools,
+        resources=resources,
+        prompts=prompts,
+        metadata=metadata,
+    )
+
+
+def specs_from_discovered_tools(
+    *,
+    connection: ConnectionConfig,
+    auth: AuthRecord | None,
+    adapter: BackendAdapter,
+    tools: list[DiscoveredTool],
+) -> list[NodeSpec[Any, Any]]:
+    return [
+        wrap_discovered_tool(
+            connection=connection,
+            auth=auth,
+            adapter=adapter,
+            tool=tool,
+        )
+        for tool in tools
+    ]

wf_mcp/service.py

@@ -10,9 +10,9 @@ from wf_core import NodeUse, Workflow, execute_workflow_async
 from .adapters import BackendAdapter
 from .catalog import CombinedCatalog, snapshot_from_specs
 from .connections import ConnectionRegistry, parse_connection_id, qualify_node_name
+from .discovery import discover_connection_capabilities, specs_from_discovered_tools
 from .models import AuthRecord, CatalogSnapshot, ConnectionConfig, RawWorkflowPlan
 from .store import Store
-from .wrappers import wrap_discovered_tool
 
 
 def _qualify_spec(connection_id: str, spec: NodeSpec[Any, Any]) -> NodeSpec[Any, Any]:
@@ -92,19 +92,17 @@ class WfMcpService:
             raise KeyError(f"no adapter registered for server {connection.server!r}")
 
         auth = self.load_auth(connection_id)
-        tools = await adapter.list_tools(connection, auth)
-        resources = await adapter.list_resources(connection, auth)
-        prompts = await adapter.list_prompts(connection, auth)
-        metadata = await adapter.get_connection_metadata(connection, auth)
-        specs = [
-            wrap_discovered_tool(
-                connection=connection,
-                auth=auth,
-                adapter=adapter,
-                tool=tool,
-            )
-            for tool in tools
-        ]
+        capabilities = await discover_connection_capabilities(
+            connection=connection,
+            auth=auth,
+            adapter=adapter,
+        )
+        specs = specs_from_discovered_tools(
+            connection=connection,
+            auth=auth,
+            adapter=adapter,
+            tools=capabilities.tools,
+        )
         self.register_specs(
             connection_id,
             *specs,
@@ -113,9 +111,9 @@ class WfMcpService:
         snapshot = snapshot_from_specs(
             connection_id,
             specs=self.specs_by_connection.get(connection_id, {}),
-            resources=resources,
-            prompts=prompts,
-            metadata=metadata,
+            resources=capabilities.resources,
+            prompts=capabilities.prompts,
+            metadata=capabilities.metadata,
             fetched_at_epoch_ms=int(time.time() * 1000),
             max_age_seconds=max_age_seconds or self.default_catalog_max_age_seconds,
         )