feat: add neutral auth store boundary
This commit is contained in:
@@ -0,0 +1,69 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from wf_api.auth import AuthRecord as NeutralAuthRecord
|
||||
|
||||
from .models import AuthRecord as McpAuthRecord
|
||||
|
||||
|
||||
def mcp_auth_from_neutral(record: NeutralAuthRecord) -> McpAuthRecord:
|
||||
"""Adapt neutral auth to the current MCP compatibility record."""
|
||||
|
||||
return McpAuthRecord(
|
||||
connection_id=record.id,
|
||||
scheme=record.scheme,
|
||||
payload=dict(record.payload),
|
||||
)
|
||||
|
||||
|
||||
def neutral_auth_from_mcp(record: McpAuthRecord) -> NeutralAuthRecord:
|
||||
"""Adapt legacy MCP auth into the neutral record shape."""
|
||||
|
||||
return NeutralAuthRecord(
|
||||
id=record.connection_id,
|
||||
scheme=record.scheme,
|
||||
payload=dict(record.payload),
|
||||
)
|
||||
|
||||
|
||||
def mcp_auth_headers(auth: McpAuthRecord | None) -> dict[str, str]:
|
||||
"""Return HTTP headers understood by MCP HTTP transports.
|
||||
|
||||
This is intentionally MCP-specific. Neutral code must not inspect payload
|
||||
keys such as `headers` or `token`.
|
||||
"""
|
||||
|
||||
if auth is None:
|
||||
return {}
|
||||
raw_headers = auth.payload.get("headers", {})
|
||||
headers = {
|
||||
str(key): str(value)
|
||||
for key, value in raw_headers.items()
|
||||
if isinstance(key, str) and isinstance(value, str)
|
||||
} if isinstance(raw_headers, dict) else {}
|
||||
token = auth.payload.get("token")
|
||||
if isinstance(token, str) and "Authorization" not in headers:
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
return headers
|
||||
|
||||
|
||||
def mcp_auth_env(auth: McpAuthRecord | None) -> dict[str, str]:
|
||||
"""Return environment variables understood by MCP stdio transports."""
|
||||
|
||||
if auth is None:
|
||||
return {}
|
||||
raw_env = auth.payload.get("env", {})
|
||||
if not isinstance(raw_env, dict):
|
||||
return {}
|
||||
return {
|
||||
str(key): str(value)
|
||||
for key, value in raw_env.items()
|
||||
if isinstance(key, str) and isinstance(value, str)
|
||||
}
|
||||
|
||||
|
||||
__all__ = [
|
||||
"mcp_auth_env",
|
||||
"mcp_auth_from_neutral",
|
||||
"mcp_auth_headers",
|
||||
"neutral_auth_from_mcp",
|
||||
]
|
||||
@@ -91,7 +91,7 @@ class WfMcpService:
|
||||
connection_list_enabled=self.connection_service.list_enabled,
|
||||
connection_list_all=self.connection_service.list_all,
|
||||
tool_executor_for=self.upstream.tool_executor_for,
|
||||
load_auth=self.upstream.load_auth,
|
||||
load_auth=self.upstream.load_connection_auth,
|
||||
emit_event=self.events.record_event,
|
||||
default_catalog_max_age_seconds=self.default_catalog_max_age_seconds,
|
||||
)
|
||||
|
||||
@@ -36,7 +36,7 @@ from .specs import get_qualified_spec, qualify_spec
|
||||
ConnectionLookup = Callable[[str], ConnectionConfig]
|
||||
ConnectionList = Callable[[], list[ConnectionConfig]]
|
||||
ToolExecutorLookup = Callable[[ConnectionConfig], ToolExecutor]
|
||||
AuthLoader = Callable[[str], AuthRecord | None]
|
||||
AuthLoader = Callable[[ConnectionConfig], AuthRecord | None]
|
||||
EventEmitter = Callable[[McpEvent], None]
|
||||
|
||||
|
||||
@@ -270,7 +270,7 @@ class SourceCatalogService:
|
||||
|
||||
async def invoke_tool(payload: BaseModel) -> NodeReturn[BaseModel]:
|
||||
connection = self.connection_lookup(entry.connection_id)
|
||||
auth = self.load_auth(entry.connection_id)
|
||||
auth = self.load_auth(connection)
|
||||
result = await self.tool_executor_for(connection).call_tool(
|
||||
connection,
|
||||
auth,
|
||||
|
||||
@@ -63,6 +63,19 @@ class UpstreamTransportService:
|
||||
def load_auth(self, connection_id: str) -> AuthRecord | None:
|
||||
return self.store.load_auth(connection_id)
|
||||
|
||||
def load_connection_auth(self, connection: ConnectionConfig) -> AuthRecord | None:
|
||||
"""Resolve auth for a connection, preferring explicit source auth_ref.
|
||||
|
||||
Legacy MCP auth records are keyed by connection id. New source registry
|
||||
and neutral config entries carry `auth_ref`; keep both paths until the
|
||||
old compatibility surface has no callers.
|
||||
"""
|
||||
|
||||
auth_ref = connection.metadata.get("auth_ref")
|
||||
if isinstance(auth_ref, str):
|
||||
return self.load_auth(auth_ref)
|
||||
return self.load_auth(connection.id)
|
||||
|
||||
def tool_executor_for(self, connection: ConnectionConfig) -> ToolExecutor:
|
||||
"""Return the executor used by generated workflow NodeSpecs.
|
||||
|
||||
@@ -81,7 +94,7 @@ class UpstreamTransportService:
|
||||
uri: str,
|
||||
) -> dict[str, Any]:
|
||||
adapter = require_adapter(connection, self.adapters)
|
||||
auth = self.load_auth(connection.id)
|
||||
auth = self.load_connection_auth(connection)
|
||||
self.event_sink(
|
||||
make_event(
|
||||
"resource_read_started",
|
||||
@@ -109,7 +122,7 @@ class UpstreamTransportService:
|
||||
arguments: dict[str, str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
adapter = require_adapter(connection, self.adapters)
|
||||
auth = self.load_auth(connection.id)
|
||||
auth = self.load_connection_auth(connection)
|
||||
self.event_sink(
|
||||
make_event(
|
||||
"prompt_get_started",
|
||||
@@ -137,7 +150,7 @@ class UpstreamTransportService:
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
adapter = require_adapter(connection, self.adapters)
|
||||
auth = self.load_auth(connection.id)
|
||||
auth = self.load_connection_auth(connection)
|
||||
self.event_sink(
|
||||
make_event(
|
||||
"raw_method_started",
|
||||
@@ -165,7 +178,7 @@ class UpstreamTransportService:
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
adapter = require_adapter(connection, self.adapters)
|
||||
auth = self.load_auth(connection.id)
|
||||
auth = self.load_connection_auth(connection)
|
||||
self.event_sink(
|
||||
make_event(
|
||||
"raw_notification_started",
|
||||
@@ -193,7 +206,7 @@ class UpstreamTransportService:
|
||||
default_catalog_max_age_seconds: int = 300,
|
||||
record_catalog_change_events: Callable[[str, CatalogSnapshot, str], None],
|
||||
) -> None:
|
||||
auth = self.load_auth(connection.id)
|
||||
auth = self.load_connection_auth(connection)
|
||||
self.event_sink(
|
||||
make_event(
|
||||
"catalog_refresh_started",
|
||||
@@ -295,7 +308,7 @@ class UpstreamTransportService:
|
||||
continue
|
||||
try:
|
||||
adapter = require_adapter(connection, self.adapters)
|
||||
auth = self.load_auth(source_id)
|
||||
auth = self.load_connection_auth(connection)
|
||||
await asyncio.wait_for(
|
||||
adapter.list_tools(connection, auth),
|
||||
timeout=LIVE_SOURCE_CHECK_TIMEOUT_SECONDS,
|
||||
|
||||
@@ -10,20 +10,11 @@ from mcp.client.stdio import StdioServerParameters, stdio_client
|
||||
from mcp.client.streamable_http import streamable_http_client
|
||||
from mcp.types import CallToolResult
|
||||
|
||||
from ..auth import mcp_auth_env, mcp_auth_headers
|
||||
from ..models import AuthRecord, ConnectionConfig
|
||||
from .session import PersistentMcpSession
|
||||
|
||||
|
||||
def _auth_headers(auth: AuthRecord | None) -> dict[str, str]:
|
||||
if auth is None:
|
||||
return {}
|
||||
headers = dict(auth.payload.get("headers", {}))
|
||||
token = auth.payload.get("token")
|
||||
if isinstance(token, str) and "Authorization" not in headers:
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
return headers
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PersistentSessionFactory:
|
||||
"""Create initialized persistent MCP sessions for configured connections.
|
||||
@@ -57,10 +48,9 @@ class PersistentSessionFactory:
|
||||
transport = connection.metadata.get("transport", "stdio")
|
||||
if transport == "stdio":
|
||||
env = connection.metadata.get("env")
|
||||
if auth is not None:
|
||||
auth_env = auth.payload.get("env")
|
||||
if isinstance(auth_env, dict):
|
||||
env = {**(env or {}), **auth_env}
|
||||
auth_env = mcp_auth_env(auth)
|
||||
if auth_env:
|
||||
env = {**(env or {}), **auth_env}
|
||||
params = StdioServerParameters(
|
||||
command=connection.metadata["command"],
|
||||
args=list(connection.metadata.get("args", [])),
|
||||
@@ -78,7 +68,7 @@ class PersistentSessionFactory:
|
||||
|
||||
if transport == "streamable_http":
|
||||
http_client = await stack.enter_async_context(
|
||||
httpx.AsyncClient(headers=_auth_headers(auth) or None)
|
||||
httpx.AsyncClient(headers=mcp_auth_headers(auth) or None)
|
||||
)
|
||||
(
|
||||
read_stream,
|
||||
|
||||
@@ -17,6 +17,7 @@ from mcp.types import (
|
||||
)
|
||||
from pydantic import AnyUrl
|
||||
|
||||
from ..auth import mcp_auth_env, mcp_auth_headers
|
||||
from ..capabilities import DiscoveredPrompt, DiscoveredResource, DiscoveredTool
|
||||
from ..models import AuthRecord, ConnectionConfig
|
||||
from .base import BackendAdapter, ToolCallResult
|
||||
@@ -28,16 +29,6 @@ from .converters import (
|
||||
)
|
||||
|
||||
|
||||
def _auth_headers(auth: AuthRecord | None) -> dict[str, str]:
|
||||
if auth is None:
|
||||
return {}
|
||||
headers = dict(auth.payload.get("headers", {}))
|
||||
token = auth.payload.get("token")
|
||||
if isinstance(token, str) and "Authorization" not in headers:
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
return headers
|
||||
|
||||
|
||||
class McpSdkAdapter(BackendAdapter):
|
||||
@asynccontextmanager
|
||||
async def _session(
|
||||
@@ -51,10 +42,9 @@ class McpSdkAdapter(BackendAdapter):
|
||||
args = list(connection.metadata.get("args", []))
|
||||
env = connection.metadata.get("env")
|
||||
cwd = connection.metadata.get("cwd")
|
||||
if auth is not None:
|
||||
auth_env = auth.payload.get("env")
|
||||
if isinstance(auth_env, dict):
|
||||
env = {**(env or {}), **auth_env}
|
||||
auth_env = mcp_auth_env(auth)
|
||||
if auth_env:
|
||||
env = {**(env or {}), **auth_env}
|
||||
params = StdioServerParameters(
|
||||
command=command,
|
||||
args=args,
|
||||
@@ -69,7 +59,7 @@ class McpSdkAdapter(BackendAdapter):
|
||||
|
||||
if transport == "streamable_http":
|
||||
url = connection.metadata["url"]
|
||||
headers = _auth_headers(auth)
|
||||
headers = mcp_auth_headers(auth)
|
||||
http_client = httpx.AsyncClient(headers=headers or None)
|
||||
async with http_client:
|
||||
async with streamable_http_client(
|
||||
|
||||
@@ -3,6 +3,9 @@ from __future__ import annotations
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from wf_api.auth import AuthRecord as NeutralAuthRecord
|
||||
|
||||
from ..auth import mcp_auth_from_neutral, neutral_auth_from_mcp
|
||||
from ..connections import parse_connection_id
|
||||
from ..models import (
|
||||
AuthRecord,
|
||||
@@ -21,6 +24,12 @@ class Store:
|
||||
def load_auth(self, connection_id: str) -> AuthRecord | None:
|
||||
raise NotImplementedError
|
||||
|
||||
def save_auth_record(self, record: NeutralAuthRecord) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def load_auth_record(self, auth_ref: str) -> NeutralAuthRecord | None:
|
||||
raise NotImplementedError
|
||||
|
||||
def save_catalog(self, snapshot: CatalogSnapshot) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
@@ -81,6 +90,19 @@ class FileStore(Store):
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
return AuthRecord(**data)
|
||||
|
||||
def save_auth_record(self, record: NeutralAuthRecord) -> None:
|
||||
"""Save neutral auth through the legacy MCP file shape."""
|
||||
|
||||
self.save_auth(mcp_auth_from_neutral(record))
|
||||
|
||||
def load_auth_record(self, auth_ref: str) -> NeutralAuthRecord | None:
|
||||
"""Load neutral auth from the legacy MCP file shape."""
|
||||
|
||||
record = self.load_auth(auth_ref)
|
||||
if record is None:
|
||||
return None
|
||||
return neutral_auth_from_mcp(record)
|
||||
|
||||
def save_catalog(self, snapshot: CatalogSnapshot) -> None:
|
||||
self._catalog_path(snapshot.connection_id).write_text(
|
||||
json.dumps(dump_catalog_snapshot(snapshot), indent=2),
|
||||
|
||||
Reference in New Issue
Block a user