refactor: add typed mcp source connection seam
This commit is contained in:
@@ -21,23 +21,31 @@ from .auth import (
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .connections import (
|
||||
McpSourceConnection,
|
||||
mcp_source_connection_from_connection_config,
|
||||
mcp_source_connection_from_registry_entry,
|
||||
)
|
||||
from .source_registry import (
|
||||
FileSourceRegistryStore,
|
||||
HttpSourceTransport,
|
||||
McpSourceRegistryEntry,
|
||||
SourceRegistryFile,
|
||||
SourceRegistryStore,
|
||||
SourceTransport,
|
||||
StdioSourceTransport,
|
||||
connection_config_to_registry_entry,
|
||||
registry_entry_to_connection_config,
|
||||
workflow_mcp_source_to_connection_config,
|
||||
)
|
||||
from .transports import (
|
||||
HttpSourceTransport,
|
||||
SourceTransport,
|
||||
StdioSourceTransport,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AuthRecord",
|
||||
"FileSourceRegistryStore",
|
||||
"HttpSourceTransport",
|
||||
"McpSourceConnection",
|
||||
"McpSourceRegistryEntry",
|
||||
"SourceRegistryFile",
|
||||
"SourceRegistryStore",
|
||||
@@ -50,6 +58,8 @@ __all__ = [
|
||||
"mcp_auth_env",
|
||||
"mcp_auth_from_neutral",
|
||||
"mcp_auth_headers",
|
||||
"mcp_source_connection_from_connection_config",
|
||||
"mcp_source_connection_from_registry_entry",
|
||||
"neutral_auth_from_mcp",
|
||||
"registry_entry_to_connection_config",
|
||||
"workflow_mcp_source_to_connection_config",
|
||||
@@ -57,14 +67,19 @@ __all__ = [
|
||||
|
||||
|
||||
def __getattr__(name: str) -> object:
|
||||
if name in {
|
||||
"McpSourceConnection",
|
||||
"mcp_source_connection_from_connection_config",
|
||||
"mcp_source_connection_from_registry_entry",
|
||||
}:
|
||||
from . import connections
|
||||
|
||||
return getattr(connections, name)
|
||||
if name in {
|
||||
"FileSourceRegistryStore",
|
||||
"HttpSourceTransport",
|
||||
"McpSourceRegistryEntry",
|
||||
"SourceRegistryFile",
|
||||
"SourceRegistryStore",
|
||||
"SourceTransport",
|
||||
"StdioSourceTransport",
|
||||
"connection_config_to_registry_entry",
|
||||
"registry_entry_to_connection_config",
|
||||
"workflow_mcp_source_to_connection_config",
|
||||
@@ -72,4 +87,12 @@ def __getattr__(name: str) -> object:
|
||||
from . import source_registry
|
||||
|
||||
return getattr(source_registry, name)
|
||||
if name in {
|
||||
"HttpSourceTransport",
|
||||
"SourceTransport",
|
||||
"StdioSourceTransport",
|
||||
}:
|
||||
from . import transports
|
||||
|
||||
return getattr(transports, name)
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
|
||||
+14
-11
@@ -1,22 +1,18 @@
|
||||
"""MCP upstream-source auth helpers.
|
||||
|
||||
This module is canonical for MCP-as-source auth interpretation. The temporary
|
||||
TYPE_CHECKING dependency on `wf_mcp.broker.models.ConnectionConfig` exists until
|
||||
connection runtime DTOs move out of the compatibility MCP facade.
|
||||
This module is canonical for MCP-as-source auth interpretation. Runtime-facing
|
||||
helpers consume source-connection-like objects instead of broker config DTOs.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from typing import Any, Protocol
|
||||
|
||||
from wf_api.auth import AuthRecord as NeutralAuthRecord
|
||||
from wf_artifacts import DependencyDiagnostic, DiagnosticSeverity
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from wf_mcp.broker.models import ConnectionConfig
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AuthRecord:
|
||||
@@ -85,11 +81,18 @@ def mcp_auth_env(auth: AuthRecord | None) -> dict[str, str]:
|
||||
}
|
||||
|
||||
|
||||
def auth_ref_for_connection(connection: ConnectionConfig) -> str | None:
|
||||
class SourceConnectionLike(Protocol):
|
||||
@property
|
||||
def id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def auth_ref(self) -> str | None: ...
|
||||
|
||||
|
||||
def auth_ref_for_connection(connection: SourceConnectionLike) -> str | None:
|
||||
"""Return the explicit auth ref for one source connection, if present."""
|
||||
|
||||
auth_ref = connection.metadata.get("auth_ref")
|
||||
return auth_ref if isinstance(auth_ref, str) else None
|
||||
return connection.auth_ref
|
||||
|
||||
|
||||
def auth_missing_diagnostic(
|
||||
@@ -117,7 +120,7 @@ def auth_missing_diagnostic(
|
||||
|
||||
|
||||
def connection_auth_diagnostic(
|
||||
connection: ConnectionConfig,
|
||||
connection: SourceConnectionLike,
|
||||
*,
|
||||
load_auth_ref: Callable[[str], AuthRecord | None],
|
||||
logical_ref: str | None = None,
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from wf_sources_mcp.ids import parse_connection_id
|
||||
from wf_sources_mcp.source_registry import McpSourceRegistryEntry
|
||||
from wf_sources_mcp.transports import (
|
||||
HttpSourceTransport,
|
||||
SourceTransport,
|
||||
StdioSourceTransport,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from wf_mcp.broker.models import ConnectionConfig
|
||||
|
||||
_FLAT_HTTP_TRANSPORTS = {"http", "streamable-http", "streamable_http", "sse"}
|
||||
_CONNECTION_METADATA_KEYS = {
|
||||
"transport",
|
||||
"command",
|
||||
"args",
|
||||
"env",
|
||||
"cwd",
|
||||
"url",
|
||||
"headers",
|
||||
"profile",
|
||||
"auth_ref",
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class McpSourceConnection:
|
||||
"""Typed runtime-facing MCP source connection.
|
||||
|
||||
This is the object runtime/session code should consume. Legacy broker
|
||||
`ConnectionConfig.metadata` remains at the compatibility edge only.
|
||||
"""
|
||||
|
||||
id: str
|
||||
provider: str
|
||||
account: str
|
||||
transport: SourceTransport
|
||||
enabled: bool = True
|
||||
profile: str | None = None
|
||||
auth_ref: str | None = None
|
||||
metadata: dict[str, object] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
provider, account = parse_connection_id(self.id)
|
||||
if not self.provider:
|
||||
raise ValueError("provider must not be empty")
|
||||
if not self.account:
|
||||
raise ValueError("account must not be empty")
|
||||
if provider != self.provider or account != self.account:
|
||||
raise ValueError(
|
||||
"MCP source connection id must match provider/account fields"
|
||||
)
|
||||
|
||||
|
||||
def mcp_source_connection_from_registry_entry(
|
||||
entry: McpSourceRegistryEntry,
|
||||
) -> McpSourceConnection:
|
||||
"""Adapt persisted desired-source registry state to runtime source shape."""
|
||||
|
||||
return McpSourceConnection(
|
||||
id=entry.id,
|
||||
provider=entry.provider,
|
||||
account=entry.account,
|
||||
enabled=entry.enabled,
|
||||
profile=entry.profile,
|
||||
transport=entry.transport,
|
||||
auth_ref=entry.auth_ref,
|
||||
metadata=dict(entry.metadata),
|
||||
)
|
||||
|
||||
|
||||
def mcp_source_connection_from_connection_config(
|
||||
connection: ConnectionConfig,
|
||||
) -> McpSourceConnection:
|
||||
"""Adapt legacy broker connection config into typed source shape.
|
||||
|
||||
Keep all metadata-bag reads in this compatibility converter. Runtime/session
|
||||
code should use `McpSourceConnection.transport` directly.
|
||||
"""
|
||||
|
||||
transport = _transport_from_connection_metadata(connection)
|
||||
profile = connection.metadata.get("profile")
|
||||
auth_ref = connection.metadata.get("auth_ref")
|
||||
metadata = {
|
||||
str(key): value
|
||||
for key, value in connection.metadata.items()
|
||||
if key not in _CONNECTION_METADATA_KEYS
|
||||
}
|
||||
return McpSourceConnection(
|
||||
id=connection.id,
|
||||
provider=connection.server,
|
||||
account=connection.account,
|
||||
enabled=connection.enabled,
|
||||
profile=profile if isinstance(profile, str) else None,
|
||||
transport=transport,
|
||||
auth_ref=auth_ref if isinstance(auth_ref, str) else None,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
|
||||
def _transport_from_connection_metadata(connection: ConnectionConfig) -> SourceTransport:
|
||||
transport = connection.metadata.get("transport")
|
||||
if isinstance(transport, dict):
|
||||
kind = transport.get("kind")
|
||||
if kind == "stdio":
|
||||
return StdioSourceTransport.model_validate(transport)
|
||||
if kind == "http":
|
||||
return HttpSourceTransport.model_validate(transport)
|
||||
raise ValueError(
|
||||
f"connection {connection.id!r} has unsupported metadata.transport.kind {kind!r}"
|
||||
)
|
||||
if isinstance(transport, str):
|
||||
if transport == "stdio":
|
||||
return StdioSourceTransport(
|
||||
command=str(connection.metadata.get("command", "")),
|
||||
args=tuple(str(arg) for arg in connection.metadata.get("args", ())),
|
||||
env={
|
||||
str(key): str(value)
|
||||
for key, value in dict(connection.metadata.get("env", {})).items()
|
||||
},
|
||||
)
|
||||
if transport in _FLAT_HTTP_TRANSPORTS:
|
||||
url = connection.metadata.get("url", "")
|
||||
return HttpSourceTransport(
|
||||
url=url if isinstance(url, str) else str(url), # type: ignore[arg-type]
|
||||
headers={
|
||||
str(key): str(value)
|
||||
for key, value in dict(
|
||||
connection.metadata.get("headers", {})
|
||||
).items()
|
||||
},
|
||||
)
|
||||
raise ValueError(
|
||||
f"connection {connection.id!r} has unrecognized metadata.transport {transport!r}"
|
||||
)
|
||||
raise ValueError(f"connection {connection.id!r} requires metadata.transport")
|
||||
|
||||
|
||||
__all__ = [
|
||||
"McpSourceConnection",
|
||||
"mcp_source_connection_from_connection_config",
|
||||
"mcp_source_connection_from_registry_entry",
|
||||
]
|
||||
@@ -0,0 +1,35 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
CONNECTION_ID_PATTERN = r"^[A-Za-z0-9_][A-Za-z0-9_.-]*$"
|
||||
|
||||
RESERVED_CONNECTION_IDS = frozenset({"wf.admin", "wf.mcp"})
|
||||
"""Source ids reserved by built-in workflow/MCP control surfaces."""
|
||||
|
||||
|
||||
def parse_connection_id(connection_id: str) -> tuple[str, str]:
|
||||
"""Validate and split one MCP source id into provider/account parts.
|
||||
|
||||
Source ids also key persisted auth, registry, and catalog files. Keep this
|
||||
conservative so unsafe ids are rejected before reaching store boundaries.
|
||||
"""
|
||||
|
||||
if not re.fullmatch(CONNECTION_ID_PATTERN, connection_id):
|
||||
raise ValueError(
|
||||
"connection id must start with alphanumeric or underscore and contain "
|
||||
"only [A-Za-z0-9_.-]"
|
||||
)
|
||||
if "." not in connection_id:
|
||||
raise ValueError("connection id must look like '<server>.<account>'")
|
||||
server, account = connection_id.split(".", 1)
|
||||
if not server or not account:
|
||||
raise ValueError("connection id must look like '<server>.<account>'")
|
||||
return server, account
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CONNECTION_ID_PATTERN",
|
||||
"RESERVED_CONNECTION_IDS",
|
||||
"parse_connection_id",
|
||||
]
|
||||
@@ -1,19 +1,13 @@
|
||||
"""Protocol/result contracts for MCP upstream source providers.
|
||||
|
||||
The temporary `wf_mcp.broker.models.ConnectionConfig` dependency remains until
|
||||
broker runtime connection DTOs move to a neutral/source-provider package.
|
||||
"""
|
||||
"""Protocol/result contracts for MCP upstream source providers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any, Protocol
|
||||
from typing import Any, Protocol
|
||||
|
||||
from wf_sources_mcp.auth import AuthRecord
|
||||
from wf_sources_mcp.catalog import DiscoveredPrompt, DiscoveredResource, DiscoveredTool
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from wf_mcp.broker.models import ConnectionConfig
|
||||
from wf_sources_mcp.connections import McpSourceConnection
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -26,38 +20,38 @@ class ToolCallResult:
|
||||
class BackendAdapter(Protocol):
|
||||
async def list_tools(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None,
|
||||
) -> list[DiscoveredTool]: ...
|
||||
|
||||
async def list_resources(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None,
|
||||
) -> list[DiscoveredResource]: ...
|
||||
|
||||
async def list_prompts(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None,
|
||||
) -> list[DiscoveredPrompt]: ...
|
||||
|
||||
async def get_connection_metadata(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None,
|
||||
) -> dict[str, Any]: ...
|
||||
|
||||
async def read_resource(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None,
|
||||
uri: str,
|
||||
) -> dict[str, Any]: ...
|
||||
|
||||
async def get_prompt(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None,
|
||||
prompt_name: str,
|
||||
arguments: dict[str, str] | None = None,
|
||||
@@ -65,7 +59,7 @@ class BackendAdapter(Protocol):
|
||||
|
||||
async def invoke_method(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None,
|
||||
method: str,
|
||||
params: dict[str, Any] | None = None,
|
||||
@@ -73,7 +67,7 @@ class BackendAdapter(Protocol):
|
||||
|
||||
async def send_notification(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None,
|
||||
method: str,
|
||||
params: dict[str, Any] | None = None,
|
||||
@@ -81,7 +75,7 @@ class BackendAdapter(Protocol):
|
||||
|
||||
async def call_tool(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None,
|
||||
tool_name: str,
|
||||
payload: dict[str, Any],
|
||||
@@ -98,7 +92,7 @@ class ToolExecutor(Protocol):
|
||||
|
||||
async def call_tool(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
connection: McpSourceConnection,
|
||||
auth: AuthRecord | None,
|
||||
tool_name: str,
|
||||
payload: dict[str, Any],
|
||||
|
||||
@@ -8,14 +8,9 @@ runtime DTOs move out of the compatibility MCP facade.
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Annotated, Literal, Protocol
|
||||
from typing import TYPE_CHECKING, Literal, Protocol
|
||||
|
||||
from pydantic import (
|
||||
AnyHttpUrl,
|
||||
Field,
|
||||
field_validator,
|
||||
model_validator,
|
||||
)
|
||||
from pydantic import Field, field_validator, model_validator
|
||||
|
||||
from wf_api.source_registry import (
|
||||
AtomicJsonRegistryStore,
|
||||
@@ -25,12 +20,12 @@ from wf_api.source_registry import (
|
||||
from wf_api.source_registry import (
|
||||
SourceRegistryStore as GenericSourceRegistryStore,
|
||||
)
|
||||
|
||||
# Temporary low-level compatibility imports. `wf_mcp.shared.names` currently
|
||||
# pulls in FastMCP transitively; keep this visible until reserved-name parsing
|
||||
# moves to a neutral/source package.
|
||||
from wf_mcp.connections import parse_connection_id
|
||||
from wf_mcp.shared.names import RESERVED_CONNECTION_IDS
|
||||
from wf_sources_mcp.ids import RESERVED_CONNECTION_IDS, parse_connection_id
|
||||
from wf_sources_mcp.transports import (
|
||||
HttpSourceTransport,
|
||||
SourceTransport,
|
||||
StdioSourceTransport,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from wf_mcp.models import ConnectionConfig
|
||||
@@ -50,25 +45,6 @@ _TRANSPORT_METADATA_KEYS = {
|
||||
}
|
||||
|
||||
|
||||
class StdioSourceTransport(SourceRegistryBaseModel):
|
||||
kind: Literal["stdio"] = "stdio"
|
||||
command: str = Field(min_length=1)
|
||||
args: tuple[str, ...] = ()
|
||||
env: dict[str, str] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class HttpSourceTransport(SourceRegistryBaseModel):
|
||||
kind: Literal["http"] = "http"
|
||||
url: AnyHttpUrl
|
||||
headers: dict[str, str] = Field(default_factory=dict)
|
||||
|
||||
|
||||
SourceTransport = Annotated[
|
||||
StdioSourceTransport | HttpSourceTransport,
|
||||
Field(discriminator="kind"),
|
||||
]
|
||||
|
||||
|
||||
class McpSourceRegistryEntry(SourceRegistryBaseModel):
|
||||
"""Desired MCP source configuration persisted by server-owned mutation."""
|
||||
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated, Literal
|
||||
|
||||
from pydantic import AnyHttpUrl, Field
|
||||
|
||||
from wf_api.source_registry import SourceRegistryBaseModel
|
||||
|
||||
|
||||
class StdioSourceTransport(SourceRegistryBaseModel):
|
||||
kind: Literal["stdio"] = "stdio"
|
||||
command: str = Field(min_length=1)
|
||||
args: tuple[str, ...] = ()
|
||||
env: dict[str, str] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class HttpSourceTransport(SourceRegistryBaseModel):
|
||||
kind: Literal["http"] = "http"
|
||||
url: AnyHttpUrl
|
||||
headers: dict[str, str] = Field(default_factory=dict)
|
||||
|
||||
|
||||
SourceTransport = Annotated[
|
||||
StdioSourceTransport | HttpSourceTransport,
|
||||
Field(discriminator="kind"),
|
||||
]
|
||||
|
||||
|
||||
__all__ = [
|
||||
"HttpSourceTransport",
|
||||
"SourceTransport",
|
||||
"StdioSourceTransport",
|
||||
]
|
||||
Reference in New Issue
Block a user