upstream transport
This commit is contained in:
@@ -124,6 +124,10 @@ implementation state.
|
||||
- Workflow runtime execution is being separated from broker coordination.
|
||||
`WorkflowRuntimeService` now owns plan compilation, dependency preparation,
|
||||
run, and resume; `WfMcpService` keeps delegate methods for compatibility.
|
||||
- Upstream MCP transport is being separated from broker coordination.
|
||||
`UpstreamTransportService` now owns adapter registration, auth persistence,
|
||||
catalog refresh I/O, resource/prompt reads, raw method/notification calls,
|
||||
generated-tool executor selection, and live source diagnostics.
|
||||
|
||||
Frame stress points remaining for native subgraphs and future fork/gather:
|
||||
|
||||
|
||||
@@ -1,6 +1,5 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
@@ -35,16 +34,14 @@ from ...models import (
|
||||
)
|
||||
from ...sdk import BackendAdapter
|
||||
from ...runtime import ToolExecutor
|
||||
from ...shared.errors import error_payload
|
||||
from ...shared.names import RESERVED_CONNECTION_IDS
|
||||
from ...storage import Store
|
||||
from wf_api.saved_subgraphs import SavedSubgraphTree
|
||||
from ..admin_capabilities import admin_source
|
||||
from ..catalog import CombinedCatalog, snapshot_from_specs
|
||||
from ..discovery import discover_connection_capabilities, specs_from_discovered_tools
|
||||
from .adapters import require_adapter
|
||||
from ..catalog import CombinedCatalog
|
||||
from .builtins import builtin_sources
|
||||
from .source_catalog import SourceCatalogService
|
||||
from .upstream_transport import UpstreamTransportService
|
||||
from .workflow_runtime import WorkflowRuntimeService
|
||||
|
||||
|
||||
@@ -53,13 +50,13 @@ class WfMcpService:
|
||||
store: Store
|
||||
default_catalog_max_age_seconds: int = 300
|
||||
connections: ConnectionRegistry = field(default_factory=ConnectionRegistry)
|
||||
adapters: dict[str, BackendAdapter] = field(default_factory=dict)
|
||||
event_bus: EventBus = field(default_factory=EventBus)
|
||||
include_builtin_specs: bool = True
|
||||
artifact_store: WorkflowArtifactStore | None = None
|
||||
draft_workspace_store: DraftWorkspaceStore | None = None
|
||||
run_store: RunStore | None = None
|
||||
tool_executor: ToolExecutor | None = None
|
||||
upstream: UpstreamTransportService = field(init=False)
|
||||
source_catalog: SourceCatalogService = field(init=False)
|
||||
workflow_runtime: WorkflowRuntimeService = field(init=False)
|
||||
|
||||
@@ -70,13 +67,18 @@ class WfMcpService:
|
||||
must not guess workflow persistence from the MCP catalog/auth store because
|
||||
CLI, MCP, and future HTTP frontends may share or swap those stores.
|
||||
"""
|
||||
self.upstream = UpstreamTransportService(
|
||||
store=self.store,
|
||||
event_sink=self._record_event,
|
||||
tool_executor=self.tool_executor,
|
||||
)
|
||||
self.source_catalog = SourceCatalogService(
|
||||
store=self.store,
|
||||
connection_lookup=self.connections.get,
|
||||
connection_list_enabled=self.connections.list_enabled,
|
||||
connection_list_all=self.connections.list_all,
|
||||
tool_executor_for=self._tool_executor_for,
|
||||
load_auth=self.load_auth,
|
||||
tool_executor_for=self.upstream.tool_executor_for,
|
||||
load_auth=self.upstream.load_auth,
|
||||
emit_event=self._record_event,
|
||||
default_catalog_max_age_seconds=self.default_catalog_max_age_seconds,
|
||||
)
|
||||
@@ -99,6 +101,11 @@ class WfMcpService:
|
||||
"""
|
||||
return self.source_catalog.capability_sources
|
||||
|
||||
@property
|
||||
def adapters(self) -> dict[str, BackendAdapter]:
|
||||
"""Compatibility view of upstream adapter registry."""
|
||||
return self.upstream.adapters
|
||||
|
||||
def register_connection(self, connection: ConnectionConfig) -> None:
|
||||
parse_connection_id(connection.id)
|
||||
if connection.id in RESERVED_CONNECTION_IDS:
|
||||
@@ -141,31 +148,16 @@ class WfMcpService:
|
||||
source.enabled = connection.enabled
|
||||
|
||||
def register_adapter(self, server: str, adapter: BackendAdapter) -> None:
|
||||
self.adapters[server] = adapter
|
||||
self.upstream.register_adapter(server, adapter)
|
||||
|
||||
def _tool_executor_for(self, connection: ConnectionConfig) -> ToolExecutor:
|
||||
"""Return the executor used by generated workflow NodeSpecs.
|
||||
|
||||
Discovery still uses the short-lived adapter path. Generated workflow
|
||||
nodes use this executor hook so config-built services can swap in a
|
||||
persistent runtime pool for stateful MCP servers.
|
||||
"""
|
||||
if self.tool_executor is not None:
|
||||
return self.tool_executor
|
||||
return require_adapter(connection, self.adapters)
|
||||
return self.upstream.tool_executor_for(connection)
|
||||
|
||||
def save_auth(self, record: AuthRecord) -> None:
|
||||
self.store.save_auth(record)
|
||||
self._record_event(
|
||||
make_event(
|
||||
"auth_saved",
|
||||
connection_id=record.connection_id,
|
||||
payload={"scheme": record.scheme},
|
||||
)
|
||||
)
|
||||
self.upstream.save_auth(record)
|
||||
|
||||
def load_auth(self, connection_id: str) -> AuthRecord | None:
|
||||
return self.store.load_auth(connection_id)
|
||||
return self.upstream.load_auth(connection_id)
|
||||
|
||||
def register_specs(
|
||||
self,
|
||||
@@ -268,26 +260,11 @@ class WfMcpService:
|
||||
|
||||
resource = self.get_resource(qualified_name)
|
||||
connection = self.connections.get(resource.connection_id)
|
||||
adapter = require_adapter(connection, self.adapters)
|
||||
auth = self.load_auth(resource.connection_id)
|
||||
self._record_event(
|
||||
make_event(
|
||||
"resource_read_started",
|
||||
connection_id=resource.connection_id,
|
||||
capability_id=qualified_name,
|
||||
payload={"uri": resource.uri},
|
||||
return await self.upstream.read_resource(
|
||||
connection,
|
||||
qualified_name,
|
||||
resource.uri,
|
||||
)
|
||||
)
|
||||
result = await adapter.read_resource(connection, auth, resource.uri)
|
||||
self._record_event(
|
||||
make_event(
|
||||
"resource_read_completed",
|
||||
connection_id=resource.connection_id,
|
||||
capability_id=qualified_name,
|
||||
payload={"uri": resource.uri},
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
async def invoke_method(
|
||||
self,
|
||||
@@ -297,26 +274,11 @@ class WfMcpService:
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
connection = self.connections.get(connection_id)
|
||||
adapter = require_adapter(connection, self.adapters)
|
||||
auth = self.load_auth(connection_id)
|
||||
self._record_event(
|
||||
make_event(
|
||||
"raw_method_started",
|
||||
connection_id=connection_id,
|
||||
capability_id=method,
|
||||
payload={"params": params or {}},
|
||||
return await self.upstream.invoke_method(
|
||||
connection,
|
||||
method,
|
||||
params=params,
|
||||
)
|
||||
)
|
||||
result = await adapter.invoke_method(connection, auth, method, params)
|
||||
self._record_event(
|
||||
make_event(
|
||||
"raw_method_completed",
|
||||
connection_id=connection_id,
|
||||
capability_id=method,
|
||||
payload={"result_keys": sorted(result.keys())},
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
async def send_notification(
|
||||
self,
|
||||
@@ -326,25 +288,7 @@ class WfMcpService:
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
connection = self.connections.get(connection_id)
|
||||
adapter = require_adapter(connection, self.adapters)
|
||||
auth = self.load_auth(connection_id)
|
||||
self._record_event(
|
||||
make_event(
|
||||
"raw_notification_started",
|
||||
connection_id=connection_id,
|
||||
capability_id=method,
|
||||
payload={"params": params or {}},
|
||||
)
|
||||
)
|
||||
await adapter.send_notification(connection, auth, method, params)
|
||||
self._record_event(
|
||||
make_event(
|
||||
"raw_notification_completed",
|
||||
connection_id=connection_id,
|
||||
capability_id=method,
|
||||
payload={},
|
||||
)
|
||||
)
|
||||
await self.upstream.send_notification(connection, method, params=params)
|
||||
|
||||
async def render_prompt(
|
||||
self,
|
||||
@@ -379,31 +323,12 @@ class WfMcpService:
|
||||
|
||||
prompt = self.get_prompt(qualified_name)
|
||||
connection = self.connections.get(prompt.connection_id)
|
||||
adapter = require_adapter(connection, self.adapters)
|
||||
auth = self.load_auth(prompt.connection_id)
|
||||
self._record_event(
|
||||
make_event(
|
||||
"prompt_get_started",
|
||||
connection_id=prompt.connection_id,
|
||||
capability_id=qualified_name,
|
||||
payload={"argument_keys": sorted((arguments or {}).keys())},
|
||||
)
|
||||
)
|
||||
result = await adapter.get_prompt(
|
||||
return await self.upstream.render_prompt(
|
||||
connection,
|
||||
auth,
|
||||
qualified_name,
|
||||
prompt.local_name,
|
||||
arguments,
|
||||
)
|
||||
self._record_event(
|
||||
make_event(
|
||||
"prompt_get_completed",
|
||||
connection_id=prompt.connection_id,
|
||||
capability_id=qualified_name,
|
||||
payload={"argument_keys": sorted((arguments or {}).keys())},
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
async def refresh_connection_catalog(
|
||||
self,
|
||||
@@ -412,73 +337,19 @@ class WfMcpService:
|
||||
max_age_seconds: int | None = None,
|
||||
) -> None:
|
||||
connection = self.connections.get(connection_id)
|
||||
adapter = require_adapter(connection, self.adapters)
|
||||
|
||||
auth = self.load_auth(connection_id)
|
||||
self._record_event(
|
||||
make_event(
|
||||
"catalog_refresh_started",
|
||||
connection_id=connection_id,
|
||||
payload={"server": connection.server},
|
||||
)
|
||||
)
|
||||
try:
|
||||
capabilities = await discover_connection_capabilities(
|
||||
connection=connection,
|
||||
auth=auth,
|
||||
adapter=adapter,
|
||||
)
|
||||
specs = specs_from_discovered_tools(
|
||||
connection=connection,
|
||||
auth=auth,
|
||||
executor=self._tool_executor_for(connection),
|
||||
tools=capabilities.tools,
|
||||
emit_event=self._record_event,
|
||||
)
|
||||
self.register_specs(
|
||||
connection_id,
|
||||
*specs,
|
||||
await self.upstream.refresh_connection_catalog(
|
||||
connection,
|
||||
source_catalog=self.source_catalog,
|
||||
max_age_seconds=max_age_seconds,
|
||||
emit_change_events=False,
|
||||
)
|
||||
snapshot = snapshot_from_specs(
|
||||
connection_id,
|
||||
specs=self.capability_sources[connection_id].capabilities.node_specs,
|
||||
tool_display_names={
|
||||
tool.name: tool.title for tool in capabilities.tools
|
||||
},
|
||||
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,
|
||||
)
|
||||
self.store.save_catalog(snapshot)
|
||||
default_catalog_max_age_seconds=self.default_catalog_max_age_seconds,
|
||||
record_catalog_change_events=lambda source_id, snapshot, reason: (
|
||||
self._record_catalog_change_events(
|
||||
connection_id,
|
||||
source_id,
|
||||
snapshot,
|
||||
reason="catalog_refresh",
|
||||
reason=reason,
|
||||
)
|
||||
self._record_event(
|
||||
make_event(
|
||||
"catalog_refresh_completed",
|
||||
connection_id=connection_id,
|
||||
payload={
|
||||
"node_count": len(snapshot.nodes),
|
||||
"resource_count": len(snapshot.resources),
|
||||
"prompt_count": len(snapshot.prompts),
|
||||
},
|
||||
),
|
||||
)
|
||||
)
|
||||
except Exception as exc:
|
||||
self._record_event(
|
||||
make_event(
|
||||
"catalog_refresh_failed",
|
||||
connection_id=connection_id,
|
||||
payload=error_payload(exc),
|
||||
)
|
||||
)
|
||||
raise
|
||||
|
||||
def compile_plan(
|
||||
self,
|
||||
|
||||
@@ -0,0 +1,362 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from collections.abc import Callable, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
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.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.runtime import ToolExecutor
|
||||
from wf_mcp.sdk import BackendAdapter
|
||||
from wf_mcp.shared.errors import error_payload
|
||||
from wf_mcp.storage import Store
|
||||
|
||||
from .adapters import require_adapter
|
||||
from .source_catalog import SourceCatalogService
|
||||
|
||||
EventSink = Callable[[McpEvent], None]
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class UpstreamTransportService:
|
||||
"""Own upstream MCP adapter/auth operations for the broker service.
|
||||
|
||||
This is not protocol-neutral. It is the MCP transport implementation used by
|
||||
admin calls, discovery, generated workflow NodeSpecs, and live source checks.
|
||||
"""
|
||||
|
||||
store: Store
|
||||
event_sink: EventSink
|
||||
adapters: dict[str, BackendAdapter] = field(default_factory=dict)
|
||||
tool_executor: ToolExecutor | None = None
|
||||
|
||||
def register_adapter(self, server: str, adapter: BackendAdapter) -> None:
|
||||
self.adapters[server] = adapter
|
||||
|
||||
def save_auth(self, record: AuthRecord) -> None:
|
||||
self.store.save_auth(record)
|
||||
self.event_sink(
|
||||
make_event(
|
||||
"auth_saved",
|
||||
connection_id=record.connection_id,
|
||||
payload={"scheme": record.scheme},
|
||||
)
|
||||
)
|
||||
|
||||
def load_auth(self, connection_id: str) -> AuthRecord | None:
|
||||
return self.store.load_auth(connection_id)
|
||||
|
||||
def tool_executor_for(self, connection: ConnectionConfig) -> ToolExecutor:
|
||||
"""Return the executor used by generated workflow NodeSpecs.
|
||||
|
||||
Discovery uses short-lived adapters. Generated workflow nodes use this
|
||||
hook so config-built services can swap in a persistent runtime pool for
|
||||
stateful MCP servers.
|
||||
"""
|
||||
if self.tool_executor is not None:
|
||||
return self.tool_executor
|
||||
return require_adapter(connection, self.adapters)
|
||||
|
||||
async def read_resource(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
qualified_name: str,
|
||||
uri: str,
|
||||
) -> dict[str, Any]:
|
||||
adapter = require_adapter(connection, self.adapters)
|
||||
auth = self.load_auth(connection.id)
|
||||
self.event_sink(
|
||||
make_event(
|
||||
"resource_read_started",
|
||||
connection_id=connection.id,
|
||||
capability_id=qualified_name,
|
||||
payload={"uri": uri},
|
||||
)
|
||||
)
|
||||
result = await adapter.read_resource(connection, auth, uri)
|
||||
self.event_sink(
|
||||
make_event(
|
||||
"resource_read_completed",
|
||||
connection_id=connection.id,
|
||||
capability_id=qualified_name,
|
||||
payload={"uri": uri},
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
async def render_prompt(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
qualified_name: str,
|
||||
local_name: str,
|
||||
arguments: dict[str, str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
adapter = require_adapter(connection, self.adapters)
|
||||
auth = self.load_auth(connection.id)
|
||||
self.event_sink(
|
||||
make_event(
|
||||
"prompt_get_started",
|
||||
connection_id=connection.id,
|
||||
capability_id=qualified_name,
|
||||
payload={"argument_keys": sorted((arguments or {}).keys())},
|
||||
)
|
||||
)
|
||||
result = await adapter.get_prompt(connection, auth, local_name, arguments)
|
||||
self.event_sink(
|
||||
make_event(
|
||||
"prompt_get_completed",
|
||||
connection_id=connection.id,
|
||||
capability_id=qualified_name,
|
||||
payload={"argument_keys": sorted((arguments or {}).keys())},
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
async def invoke_method(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
method: str,
|
||||
*,
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
adapter = require_adapter(connection, self.adapters)
|
||||
auth = self.load_auth(connection.id)
|
||||
self.event_sink(
|
||||
make_event(
|
||||
"raw_method_started",
|
||||
connection_id=connection.id,
|
||||
capability_id=method,
|
||||
payload={"params": params or {}},
|
||||
)
|
||||
)
|
||||
result = await adapter.invoke_method(connection, auth, method, params)
|
||||
self.event_sink(
|
||||
make_event(
|
||||
"raw_method_completed",
|
||||
connection_id=connection.id,
|
||||
capability_id=method,
|
||||
payload={"result_keys": sorted(result.keys())},
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
async def send_notification(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
method: str,
|
||||
*,
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
adapter = require_adapter(connection, self.adapters)
|
||||
auth = self.load_auth(connection.id)
|
||||
self.event_sink(
|
||||
make_event(
|
||||
"raw_notification_started",
|
||||
connection_id=connection.id,
|
||||
capability_id=method,
|
||||
payload={"params": params or {}},
|
||||
)
|
||||
)
|
||||
await adapter.send_notification(connection, auth, method, params)
|
||||
self.event_sink(
|
||||
make_event(
|
||||
"raw_notification_completed",
|
||||
connection_id=connection.id,
|
||||
capability_id=method,
|
||||
payload={},
|
||||
)
|
||||
)
|
||||
|
||||
async def refresh_connection_catalog(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
*,
|
||||
source_catalog: SourceCatalogService,
|
||||
max_age_seconds: int | None = None,
|
||||
default_catalog_max_age_seconds: int = 300,
|
||||
record_catalog_change_events: Callable[[str, CatalogSnapshot, str], None],
|
||||
) -> None:
|
||||
auth = self.load_auth(connection.id)
|
||||
self.event_sink(
|
||||
make_event(
|
||||
"catalog_refresh_started",
|
||||
connection_id=connection.id,
|
||||
payload={"server": connection.server},
|
||||
)
|
||||
)
|
||||
try:
|
||||
adapter = require_adapter(connection, self.adapters)
|
||||
capabilities = await discover_connection_capabilities(
|
||||
connection=connection,
|
||||
auth=auth,
|
||||
adapter=adapter,
|
||||
)
|
||||
specs = specs_from_discovered_tools(
|
||||
connection=connection,
|
||||
auth=auth,
|
||||
executor=self.tool_executor_for(connection),
|
||||
tools=capabilities.tools,
|
||||
emit_event=self.event_sink,
|
||||
)
|
||||
source_catalog.register_specs(
|
||||
connection.id,
|
||||
*specs,
|
||||
max_age_seconds=max_age_seconds,
|
||||
emit_change_events=False,
|
||||
)
|
||||
snapshot = snapshot_from_specs(
|
||||
connection.id,
|
||||
specs=source_catalog.capability_sources[
|
||||
connection.id
|
||||
].capabilities.node_specs,
|
||||
tool_display_names={
|
||||
tool.name: tool.title for tool in capabilities.tools
|
||||
},
|
||||
resources=capabilities.resources,
|
||||
prompts=capabilities.prompts,
|
||||
metadata=capabilities.metadata,
|
||||
fetched_at_epoch_ms=int(time.time() * 1000),
|
||||
max_age_seconds=max_age_seconds or default_catalog_max_age_seconds,
|
||||
)
|
||||
self.store.save_catalog(snapshot)
|
||||
record_catalog_change_events(connection.id, snapshot, "catalog_refresh")
|
||||
self.event_sink(
|
||||
make_event(
|
||||
"catalog_refresh_completed",
|
||||
connection_id=connection.id,
|
||||
payload={
|
||||
"node_count": len(snapshot.nodes),
|
||||
"resource_count": len(snapshot.resources),
|
||||
"prompt_count": len(snapshot.prompts),
|
||||
},
|
||||
)
|
||||
)
|
||||
except Exception as exc:
|
||||
self.event_sink(
|
||||
make_event(
|
||||
"catalog_refresh_failed",
|
||||
connection_id=connection.id,
|
||||
payload=error_payload(exc),
|
||||
)
|
||||
)
|
||||
raise
|
||||
|
||||
async def deployment_diagnostics(
|
||||
self,
|
||||
*,
|
||||
deployment: WorkflowDeployment,
|
||||
artifacts: Sequence[WorkflowArtifact],
|
||||
source_catalog: SourceCatalogService,
|
||||
) -> list[DependencyDiagnostic]:
|
||||
"""Return opt-in diagnostics for bound upstream sources that cannot answer.
|
||||
|
||||
Static deployment validation only checks the last known source catalog.
|
||||
This probe intentionally performs live upstream I/O, so MCP tools keep it
|
||||
disabled by default and only run it when the caller asks for liveness.
|
||||
"""
|
||||
diagnostics: list[DependencyDiagnostic] = []
|
||||
for source_id, logical_ref in _required_live_sources(
|
||||
deployment, artifacts
|
||||
).items():
|
||||
source = source_catalog.capability_sources.get(source_id)
|
||||
if (
|
||||
source is None
|
||||
or not source.enabled
|
||||
or not source.permissions.calls_upstream
|
||||
):
|
||||
continue
|
||||
try:
|
||||
connection = source_catalog.connection_lookup(source_id)
|
||||
except KeyError as exc:
|
||||
diagnostics.append(
|
||||
_source_unreachable_diagnostic(
|
||||
logical_ref=logical_ref,
|
||||
source_id=source_id,
|
||||
exc=exc,
|
||||
)
|
||||
)
|
||||
continue
|
||||
try:
|
||||
adapter = require_adapter(connection, self.adapters)
|
||||
auth = self.load_auth(source_id)
|
||||
await asyncio.wait_for(
|
||||
adapter.list_tools(connection, auth),
|
||||
timeout=LIVE_SOURCE_CHECK_TIMEOUT_SECONDS,
|
||||
)
|
||||
except _LIVE_SOURCE_CHECK_FAILURES as exc:
|
||||
diagnostics.append(
|
||||
_source_unreachable_diagnostic(
|
||||
logical_ref=logical_ref,
|
||||
source_id=source_id,
|
||||
exc=exc,
|
||||
)
|
||||
)
|
||||
return diagnostics
|
||||
|
||||
|
||||
LIVE_SOURCE_CHECK_TIMEOUT_SECONDS = 8.0
|
||||
_LIVE_SOURCE_CHECK_FAILURES = (
|
||||
TimeoutError,
|
||||
OSError,
|
||||
anyio.ClosedResourceError,
|
||||
anyio.EndOfStream,
|
||||
anyio.BrokenResourceError,
|
||||
httpx.HTTPError,
|
||||
McpError,
|
||||
StreamableHTTPError,
|
||||
)
|
||||
|
||||
|
||||
def _required_live_sources(
|
||||
deployment: WorkflowDeployment,
|
||||
artifacts: Sequence[WorkflowArtifact],
|
||||
) -> dict[str, str]:
|
||||
"""Return concrete upstream source ids to live-check, with one logical ref."""
|
||||
bindings = deployment.binding_map()
|
||||
required: dict[str, str] = {}
|
||||
for artifact in artifacts:
|
||||
for logical_ref, capability in artifact.required_capability_map().items():
|
||||
source_id = bindings.get(capability.logical_source)
|
||||
if source_id is not None:
|
||||
required.setdefault(source_id, logical_ref)
|
||||
return required
|
||||
|
||||
|
||||
def _source_unreachable_diagnostic(
|
||||
*,
|
||||
logical_ref: str,
|
||||
source_id: str,
|
||||
exc: BaseException,
|
||||
) -> DependencyDiagnostic:
|
||||
"""Build a liveness diagnostic without catching unrelated probe bugs."""
|
||||
return DependencyDiagnostic(
|
||||
severity=DiagnosticSeverity.ERROR,
|
||||
code="source_unreachable",
|
||||
logical_ref=logical_ref,
|
||||
bound_source=source_id,
|
||||
message=(
|
||||
f"Live check for upstream source {source_id!r} failed: "
|
||||
f"{type(exc).__name__}: {exc}"
|
||||
),
|
||||
repair_hint=(
|
||||
"Start or reconnect the source, fix its transport/auth "
|
||||
"configuration, or bind this deployment to another source."
|
||||
),
|
||||
)
|
||||
@@ -1,135 +0,0 @@
|
||||
"""MCP-adapter-owned live source diagnostics for deployment validation.
|
||||
|
||||
Static deployment validation only checks the last known source catalog.
|
||||
This probe intentionally performs live upstream I/O, so MCP tools keep it
|
||||
disabled by default and only run it when the caller asks for liveness.
|
||||
|
||||
This module owns the MCP-only imports to avoid a circular import:
|
||||
handlers.py -> workflow_operation_context.py -> handlers.py
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Sequence
|
||||
|
||||
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 .adapters import require_adapter
|
||||
from .core import WfMcpService
|
||||
|
||||
LIVE_SOURCE_CHECK_TIMEOUT_SECONDS = 8.0
|
||||
_LIVE_SOURCE_CHECK_FAILURES = (
|
||||
TimeoutError,
|
||||
OSError,
|
||||
anyio.ClosedResourceError,
|
||||
anyio.EndOfStream,
|
||||
anyio.BrokenResourceError,
|
||||
httpx.HTTPError,
|
||||
McpError,
|
||||
StreamableHTTPError,
|
||||
)
|
||||
|
||||
|
||||
def _required_live_sources(
|
||||
deployment: WorkflowDeployment,
|
||||
artifacts: Sequence[WorkflowArtifact],
|
||||
) -> dict[str, str]:
|
||||
"""Return concrete upstream source ids to live-check, with one logical ref."""
|
||||
bindings = deployment.binding_map()
|
||||
required: dict[str, str] = {}
|
||||
for artifact in artifacts:
|
||||
for logical_ref, capability in artifact.required_capability_map().items():
|
||||
source_id = bindings.get(capability.logical_source)
|
||||
if source_id is not None:
|
||||
required.setdefault(source_id, logical_ref)
|
||||
return required
|
||||
|
||||
|
||||
async def live_source_diagnostics(
|
||||
service: WfMcpService,
|
||||
*,
|
||||
deployment: WorkflowDeployment,
|
||||
artifacts: Sequence[WorkflowArtifact],
|
||||
) -> list[DependencyDiagnostic]:
|
||||
"""Return opt-in diagnostics for bound upstream sources that cannot answer.
|
||||
|
||||
Static deployment validation only checks the last known source catalog.
|
||||
This probe intentionally performs live upstream I/O, so MCP tools keep it
|
||||
disabled by default and only run it when the caller asks for liveness.
|
||||
"""
|
||||
diagnostics: list[DependencyDiagnostic] = []
|
||||
for source_id, logical_ref in _required_live_sources(deployment, artifacts).items():
|
||||
source = service.source_catalog.capability_sources.get(source_id)
|
||||
if (
|
||||
source is None
|
||||
or not source.enabled
|
||||
or not source.permissions.calls_upstream
|
||||
):
|
||||
continue
|
||||
try:
|
||||
connection = service.connections.get(source_id)
|
||||
except KeyError as exc:
|
||||
diagnostics.append(
|
||||
_source_unreachable_diagnostic(
|
||||
logical_ref=logical_ref,
|
||||
source_id=source_id,
|
||||
exc=exc,
|
||||
)
|
||||
)
|
||||
continue
|
||||
try:
|
||||
adapter = require_adapter(connection, service.adapters)
|
||||
auth = service.load_auth(source_id)
|
||||
await asyncio.wait_for(
|
||||
adapter.list_tools(connection, auth),
|
||||
timeout=LIVE_SOURCE_CHECK_TIMEOUT_SECONDS,
|
||||
)
|
||||
except _LIVE_SOURCE_CHECK_FAILURES as exc:
|
||||
diagnostics.append(
|
||||
_source_unreachable_diagnostic(
|
||||
logical_ref=logical_ref,
|
||||
source_id=source_id,
|
||||
exc=exc,
|
||||
)
|
||||
)
|
||||
return diagnostics
|
||||
|
||||
|
||||
def _source_unreachable_diagnostic(
|
||||
*,
|
||||
logical_ref: str,
|
||||
source_id: str,
|
||||
exc: BaseException,
|
||||
) -> DependencyDiagnostic:
|
||||
"""Build a liveness diagnostic without catching unrelated probe bugs."""
|
||||
return DependencyDiagnostic(
|
||||
severity=DiagnosticSeverity.ERROR,
|
||||
code="source_unreachable",
|
||||
logical_ref=logical_ref,
|
||||
bound_source=source_id,
|
||||
message=(
|
||||
f"Live check for upstream source {source_id!r} failed: "
|
||||
f"{type(exc).__name__}: {exc}"
|
||||
),
|
||||
repair_hint=(
|
||||
"Start or reconnect the source, fix its transport/auth "
|
||||
"configuration, or bind this deployment to another source."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"LIVE_SOURCE_CHECK_TIMEOUT_SECONDS",
|
||||
"live_source_diagnostics",
|
||||
]
|
||||
@@ -17,7 +17,6 @@ from wf_api.operation_context import (
|
||||
from wf_mcp.events import make_event
|
||||
|
||||
from .core import WfMcpService
|
||||
from .workflow_live_checks import live_source_diagnostics
|
||||
from .workflow_runtime import WorkflowRuntimeService
|
||||
|
||||
|
||||
@@ -122,10 +121,10 @@ class WfMcpWorkflowLiveSourceChecker(WorkflowLiveSourceChecker):
|
||||
deployment: WorkflowDeployment,
|
||||
artifacts: Sequence[WorkflowArtifact],
|
||||
) -> list[DependencyDiagnostic]:
|
||||
return await live_source_diagnostics(
|
||||
self.service,
|
||||
return await self.service.upstream.deployment_diagnostics(
|
||||
deployment=deployment,
|
||||
artifacts=artifacts,
|
||||
source_catalog=self.service.source_catalog,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,166 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
from wf_artifacts import WorkflowDeployment
|
||||
from wf_platform import CapabilityBuckets, CapabilitySource, SourcePermissions
|
||||
|
||||
from wf_mcp.broker.service.source_catalog import SourceCatalogService
|
||||
from wf_mcp.broker.service.upstream_transport import UpstreamTransportService
|
||||
from wf_mcp.connections import ConnectionRegistry
|
||||
from wf_mcp.events import McpEvent
|
||||
from wf_mcp.models import AuthRecord, ConnectionConfig
|
||||
from wf_mcp.storage import FileStore
|
||||
|
||||
from wf_mcp.broker import WfMcpService
|
||||
|
||||
from ..test_support import FakeAdapter, local_temp_root
|
||||
from ..workflow_surface.conftest import echo_artifact
|
||||
|
||||
|
||||
def test_upstream_transport_registers_adapter() -> None:
|
||||
events: list[McpEvent] = []
|
||||
transport = UpstreamTransportService(
|
||||
store=FileStore(local_temp_root() / "upstream_adapter"),
|
||||
event_sink=events.append,
|
||||
)
|
||||
adapter = FakeAdapter()
|
||||
|
||||
transport.register_adapter("demo", adapter)
|
||||
|
||||
assert transport.adapters["demo"] is adapter
|
||||
|
||||
|
||||
def test_upstream_transport_saves_and_loads_auth_with_event() -> None:
|
||||
events: list[McpEvent] = []
|
||||
transport = UpstreamTransportService(
|
||||
store=FileStore(local_temp_root() / "upstream_auth"),
|
||||
event_sink=events.append,
|
||||
)
|
||||
record = AuthRecord(connection_id="demo.personal", scheme="bearer")
|
||||
|
||||
transport.save_auth(record)
|
||||
loaded = transport.load_auth("demo.personal")
|
||||
|
||||
assert loaded is not None
|
||||
assert loaded.connection_id == "demo.personal"
|
||||
assert events[-1].kind == "auth_saved"
|
||||
assert events[-1].connection_id == "demo.personal"
|
||||
|
||||
|
||||
def test_wfmcpservice_uses_upstream_transport_for_adapters_and_auth() -> None:
|
||||
service = WfMcpService(store=FileStore(local_temp_root() / "service_upstream"))
|
||||
adapter = FakeAdapter()
|
||||
|
||||
service.register_adapter("demo", adapter)
|
||||
service.save_auth(AuthRecord(connection_id="demo.personal", scheme="bearer"))
|
||||
|
||||
assert service.upstream.adapters["demo"] is adapter
|
||||
assert service.adapters is service.upstream.adapters
|
||||
assert service.load_auth("demo.personal") is not None
|
||||
assert service.list_events()[-1].kind == "auth_saved"
|
||||
|
||||
|
||||
def test_upstream_transport_invokes_raw_method_and_records_events() -> None:
|
||||
events: list[McpEvent] = []
|
||||
connections = ConnectionRegistry()
|
||||
connections.register(
|
||||
ConnectionConfig(id="demo.personal", server="demo", account="personal")
|
||||
)
|
||||
transport = UpstreamTransportService(
|
||||
store=FileStore(local_temp_root() / "upstream_raw_method"),
|
||||
event_sink=events.append,
|
||||
)
|
||||
transport.register_adapter("demo", FakeAdapter())
|
||||
|
||||
result = asyncio.run(
|
||||
transport.invoke_method(
|
||||
connections.get("demo.personal"),
|
||||
"demo.echo",
|
||||
params={"text": "hello"},
|
||||
)
|
||||
)
|
||||
|
||||
assert result["echoed"] == "hello"
|
||||
assert [event.kind for event in events] == [
|
||||
"raw_method_started",
|
||||
"raw_method_completed",
|
||||
]
|
||||
|
||||
|
||||
def test_upstream_transport_refreshes_catalog_directly() -> None:
|
||||
events: list[McpEvent] = []
|
||||
store = FileStore(local_temp_root() / "upstream_refresh")
|
||||
connections = ConnectionRegistry()
|
||||
connection = ConnectionConfig(id="demo.personal", server="demo", account="personal")
|
||||
connections.register(connection)
|
||||
transport = UpstreamTransportService(store=store, event_sink=events.append)
|
||||
transport.register_adapter("demo", FakeAdapter())
|
||||
source_catalog = SourceCatalogService(
|
||||
store=store,
|
||||
connection_lookup=connections.get,
|
||||
connection_list_enabled=connections.list_enabled,
|
||||
connection_list_all=connections.list_all,
|
||||
tool_executor_for=transport.tool_executor_for,
|
||||
load_auth=transport.load_auth,
|
||||
emit_event=events.append,
|
||||
)
|
||||
source_catalog.hydrate_connection_source_from_snapshot(connection)
|
||||
|
||||
asyncio.run(
|
||||
transport.refresh_connection_catalog(
|
||||
connection,
|
||||
source_catalog=source_catalog,
|
||||
record_catalog_change_events=lambda source_id, snapshot, reason: None,
|
||||
)
|
||||
)
|
||||
|
||||
snapshot = store.load_catalog("demo.personal")
|
||||
assert snapshot is not None
|
||||
assert len(snapshot.nodes) >= 1
|
||||
assert "catalog_refresh_started" in [event.kind for event in events]
|
||||
assert "catalog_refresh_completed" in [event.kind for event in events]
|
||||
|
||||
|
||||
def test_upstream_transport_live_diagnostics_report_missing_connection() -> None:
|
||||
transport = UpstreamTransportService(
|
||||
store=FileStore(local_temp_root() / "upstream_live_missing"),
|
||||
event_sink=lambda event: None,
|
||||
)
|
||||
source_catalog = SourceCatalogService(
|
||||
store=transport.store,
|
||||
connection_lookup=lambda connection_id: (_ for _ in ()).throw(
|
||||
KeyError(connection_id)
|
||||
),
|
||||
connection_list_enabled=lambda: [],
|
||||
connection_list_all=lambda: [],
|
||||
tool_executor_for=transport.tool_executor_for,
|
||||
load_auth=transport.load_auth,
|
||||
emit_event=lambda event: None,
|
||||
)
|
||||
source_catalog.register_capability_source(
|
||||
CapabilitySource(
|
||||
id="demo.personal",
|
||||
kind="connection",
|
||||
permissions=SourcePermissions(calls_upstream=True),
|
||||
capabilities=CapabilityBuckets(),
|
||||
)
|
||||
)
|
||||
artifact = echo_artifact()
|
||||
deployment = WorkflowDeployment(
|
||||
id="echo.personal",
|
||||
artifact_id="echo",
|
||||
artifact_version=1,
|
||||
bindings=[{"logical_source": "demo", "concrete_source": "demo.personal"}],
|
||||
)
|
||||
|
||||
diagnostics = asyncio.run(
|
||||
transport.deployment_diagnostics(
|
||||
deployment=deployment,
|
||||
artifacts=[artifact],
|
||||
source_catalog=source_catalog,
|
||||
)
|
||||
)
|
||||
|
||||
assert diagnostics[0].code == "source_unreachable"
|
||||
assert diagnostics[0].bound_source == "demo.personal"
|
||||
Reference in New Issue
Block a user