refactor: move mcp auth record ownership

src/wf_mcp/broker/__init__.py

@@ -1,20 +1,20 @@
 from .catalog import CombinedCatalog, snapshot_from_specs
+from .config import build_service_from_config, load_broker_config
 from .discovery import (
     DiscoveredConnectionCapabilities,
     discover_connection_capabilities,
     specs_from_discovered_tools,
 )
 from .events import McpEvent, make_event
+from .models import BrokerConfig, ConnectionConfig, SourceConfigOwnership
 from .server import (
     build_workflow_server_from_config,
     build_workflow_server_from_workflow_config,
     create_broker_server,
     workflow_server_from_service,
 )
-from .config import build_service_from_config, load_broker_config
-from .models import BrokerConfig, ConnectionConfig, SourceConfigOwnership
-from .transport import normalize_transport
 from .service import WfMcpService
+from .transport import normalize_transport
 
 __all__ = [
     "BrokerConfig",

src/wf_mcp/broker/config.py

@@ -11,15 +11,19 @@ from ..control import BrokerConfigFile, ConnectionConfigFile
 from ..models import BrokerConfig
 from ..runtime import McpRuntimePool, PersistentSessionFactory
 from ..sdk import McpSdkAdapter
-from ..source_registry import FileSourceRegistryStore, workflow_mcp_source_to_connection_config
+from ..source_registry import (
+    FileSourceRegistryStore,
+    workflow_mcp_source_to_connection_config,
+)
 from ..storage import FileStore
 from .service import WfMcpService
 
-
 _HTTP_TRANSPORTS = {"http", "streamable-http", "streamable_http", "sse"}
 
 
-def _source_metadata_without_transport(metadata: dict[str, object]) -> dict[str, object]:
+def _source_metadata_without_transport(
+    metadata: dict[str, object],
+) -> dict[str, object]:
     return {
         key: value
         for key, value in metadata.items()

src/wf_mcp/broker/discovery.py

@@ -9,8 +9,9 @@ from mcp.types import METHOD_NOT_FOUND
 
 from wf_authoring import NodeSpec
 
+from ..auth import AuthRecord
 from ..capabilities import DiscoveredPrompt, DiscoveredResource, DiscoveredTool
-from ..models import AuthRecord, ConnectionConfig
+from ..models import ConnectionConfig
 from ..runtime import ToolExecutor
 from ..sdk import BackendAdapter
 from ..shared import root_exception

src/wf_mcp/broker/models.py

@@ -4,7 +4,6 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Literal
 
-
 SourceConfigOwnership = Literal["locked", "seed"]
 
 

src/wf_mcp/broker/server.py

@@ -10,9 +10,12 @@ from wf_api import (
     durable_workflow_api,
 )
 from wf_api.stores import WorkflowStores
+from wf_config import WorkflowConfigFile
 from wf_server import WorkflowServer, WorkflowServerConfig
 
-from wf_config import WorkflowConfigFile
+from ..models import BrokerConfig
+from ..sdk.adapter import McpSdkAdapter
+from ..source_registry import FileSourceRegistryStore, SourceRegistryStore
 from .artifact_tools import register_artifact_tools
 from .config import broker_config_from_workflow_config, build_service_from_config
 from .prompts import register_broker_prompts
@@ -22,9 +25,6 @@ from .service.auth_admin import McpAuthAdminProvider
 from .service.source_registry_admin import SourceRegistryAdminProvider
 from .service.workflow_operation_context import context_from_service
 from .tools import register_broker_tools
-from ..models import BrokerConfig
-from ..sdk.adapter import McpSdkAdapter
-from ..source_registry import FileSourceRegistryStore, SourceRegistryStore
 
 
 def create_broker_server(service: WfMcpService) -> FastMCP:
@@ -75,11 +75,11 @@ def workflow_server_from_service(
         config_connections=config.connections,
         connection_service=service.connection_service,
         config=config,
-        ensure_adapter=lambda connection: service.register_adapter(
-            connection.server, McpSdkAdapter()
-        )
-        if connection.server not in service.adapters
-        else None,
+        ensure_adapter=lambda connection: (
+            service.register_adapter(connection.server, McpSdkAdapter())
+            if connection.server not in service.adapters
+            else None
+        ),
         load_auth=service.upstream.load_auth,
     )
     source_registry_admin = WorkflowSourceRegistryApi(
@@ -118,9 +118,7 @@ def build_workflow_server_from_workflow_config(
     config: WorkflowConfigFile,
 ) -> WorkflowServer:
     """Build an MCP-backed WorkflowServer from neutral workflow config sources."""
-    return build_workflow_server_from_config(
-        broker_config_from_workflow_config(config)
-    )
+    return build_workflow_server_from_config(broker_config_from_workflow_config(config))
 
 
 __all__ = [

src/wf_mcp/broker/service/core.py

@@ -3,6 +3,8 @@ from __future__ import annotations
 from dataclasses import dataclass, field
 from typing import Any
 
+from wf_api.models import RawWorkflowPlan
+from wf_api.saved_subgraphs import SavedSubgraphTree
 from wf_artifacts import (
     DraftWorkspaceStore,
     RunStore,
@@ -15,29 +17,32 @@ from wf_core import (
     RunState,
     Workflow,
 )
-from wf_api.models import RawWorkflowPlan
-from wf_mcp.capabilities import CatalogNodeEntry, CatalogPromptEntry, CatalogResourceEntry
+from wf_mcp.capabilities import (
+    CatalogNodeEntry,
+    CatalogPromptEntry,
+    CatalogResourceEntry,
+)
 from wf_platform import (
     CapabilitySource,
 )
+
+from ...auth import AuthRecord
 from ...connections import ConnectionRegistry
 from ...events import EventBus, McpEvent
 from ...models import (
-    AuthRecord,
-    CatalogSnapshot,
     BrokerConfig,
+    CatalogSnapshot,
     ConnectionConfig,
 )
-from ...sdk import BackendAdapter
 from ...runtime import ToolExecutor
+from ...sdk import BackendAdapter
 from ...source_registry import SourceRegistryStore
-from .connection_service import ConnectionService
-from .content_access import ContentAccessService
 from ...storage import Store
-from wf_api.saved_subgraphs import SavedSubgraphTree
 from ..admin_capabilities import admin_source
 from ..catalog import CombinedCatalog
 from .builtins import builtin_sources
+from .connection_service import ConnectionService
+from .content_access import ContentAccessService
 from .events import BrokerEventRecorder
 from .source_catalog import SourceCatalogService
 from .upstream_transport import UpstreamTransportService

src/wf_mcp/broker/service/source_catalog.py

@@ -6,8 +6,13 @@ from dataclasses import dataclass, field
 from typing import Any
 
 from pydantic import BaseModel
+
 from wf_authoring import NodeReturn, NodeSpec
-from wf_mcp.capabilities import CatalogNodeEntry, CatalogPromptEntry, CatalogResourceEntry
+from wf_mcp.capabilities import (
+    CatalogNodeEntry,
+    CatalogPromptEntry,
+    CatalogResourceEntry,
+)
 from wf_platform import (
     CapabilityBuckets,
     CapabilitySource,
@@ -18,10 +23,10 @@ from wf_platform import (
     page_items,
 )
 
+from ...auth import AuthRecord
 from ...connections import ConnectionConfig, qualify_node_name
 from ...events import McpEvent, make_event
 from ...models import (
-    AuthRecord,
     CatalogSnapshot,
 )
 from ...runtime import ToolExecutor
@@ -30,7 +35,6 @@ from ...workflow.wrappers import _model_from_schema
 from ..catalog import CombinedCatalog, snapshot_from_specs
 from .specs import get_qualified_spec, qualify_spec
 
-
 ConnectionLookup = Callable[[str], ConnectionConfig]
 ConnectionList = Callable[[], list[ConnectionConfig]]
 ToolExecutorLookup = Callable[[ConnectionConfig], ToolExecutor]

src/wf_mcp/broker/service/source_registry_admin.py

@@ -6,8 +6,8 @@ from typing import Any
 
 from wf_api.source_registry_admin import WorkflowSourceRegistryMutationProvider
 
-from ...auth import connection_auth_diagnostic
-from ...models import AuthRecord, BrokerConfig, ConnectionConfig
+from ...auth import AuthRecord, connection_auth_diagnostic
+from ...models import BrokerConfig, ConnectionConfig
 from ...source_registry import (
     McpSourceRegistryEntry,
     SourceRegistryFile,
@@ -135,12 +135,18 @@ class SourceRegistryAdminProvider(WorkflowSourceRegistryMutationProvider):
         if self.connection_service is None or self.config is None:
             raise RuntimeError("source registry apply requires runtime service context")
 
-        before = {connection.id: connection for connection in self.connection_service.list_all()}
+        before = {
+            connection.id: connection
+            for connection in self.connection_service.list_all()
+        }
         self.connection_service.sync_connections_from_config(
             self.config,
             source_registry_store=self.source_registry_store,
         )
-        after = {connection.id: connection for connection in self.connection_service.list_all()}
+        after = {
+            connection.id: connection
+            for connection in self.connection_service.list_all()
+        }
 
         if self.ensure_adapter is not None:
             for connection in after.values():

src/wf_mcp/broker/service/specs.py

@@ -4,9 +4,8 @@ from collections.abc import Mapping
 from typing import Any
 
 from wf_authoring import NodeSpec
-
-from wf_platform import CapabilitySource
 from wf_mcp.connections import qualify_node_name
+from wf_platform import CapabilitySource
 
 
 def qualify_spec(connection_id: str, spec: NodeSpec[Any, Any]) -> NodeSpec[Any, Any]:

src/wf_mcp/broker/service/upstream_transport.py

@@ -10,19 +10,21 @@ import anyio
 import httpx
 from mcp.client.streamable_http import StreamableHTTPError
 from mcp.shared.exceptions import McpError
+
 from wf_artifacts import (
     DependencyDiagnostic,
     DiagnosticSeverity,
     WorkflowArtifact,
     WorkflowDeployment,
 )
+from wf_mcp.auth import AuthRecord
 from wf_mcp.broker.catalog import snapshot_from_specs
 from wf_mcp.broker.discovery import (
     discover_connection_capabilities,
     specs_from_discovered_tools,
 )
 from wf_mcp.events import McpEvent, make_event
-from wf_mcp.models import AuthRecord, CatalogSnapshot, ConnectionConfig
+from wf_mcp.models import CatalogSnapshot, ConnectionConfig
 from wf_mcp.runtime import ToolExecutor
 from wf_mcp.sdk import BackendAdapter
 from wf_mcp.shared.errors import error_payload

src/wf_mcp/broker/service/workflow_operation_context.py

@@ -4,8 +4,6 @@ from collections.abc import Sequence
 from dataclasses import dataclass
 from typing import Any
 
-from wf_artifacts import DependencyDiagnostic, WorkflowArtifact, WorkflowDeployment
-from wf_authoring import NodeSpec
 from wf_api.operation_context import (
     WorkflowEventRecorder,
     WorkflowLiveSourceChecker,
@@ -13,6 +11,9 @@ from wf_api.operation_context import (
     WorkflowRuntimeRunner,
     WorkflowSpecProvider,
 )
+from wf_artifacts import DependencyDiagnostic, WorkflowArtifact, WorkflowDeployment
+from wf_authoring import NodeSpec
+
 from .core import WfMcpService
 from .events import BrokerEventRecorder
 from .source_catalog import SourceCatalogService

src/wf_mcp/broker/service/workflow_runtime.py

@@ -4,6 +4,13 @@ from collections.abc import Callable
 from dataclasses import dataclass
 from typing import Any
 
+from wf_api.models import RawWorkflowPlan
+from wf_api.runtime_dependencies import resolve_runtime_dependencies
+from wf_api.saved_subgraphs import (
+    SavedSubgraphTree,
+    prepare_saved_subgraphs,
+    resolve_saved_subgraph_tree,
+)
 from wf_artifacts import WorkflowArtifact, WorkflowArtifactStore, WorkflowDeployment
 from wf_authoring import NodeSpec
 from wf_core import (
@@ -14,13 +21,6 @@ from wf_core import (
     execute_workflow_result_async,
     resume_workflow_result_async,
 )
-from wf_api.models import RawWorkflowPlan
-from wf_api.runtime_dependencies import resolve_runtime_dependencies
-from wf_api.saved_subgraphs import (
-    SavedSubgraphTree,
-    prepare_saved_subgraphs,
-    resolve_saved_subgraph_tree,
-)
 
 from ...events import McpEvent, make_event
 from .source_catalog import SourceCatalogService