301 lines
9.5 KiB
Python
301 lines
9.5 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
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 _transport(root: Path) -> UpstreamTransportService:
|
|
events: list[McpEvent] = []
|
|
return UpstreamTransportService(
|
|
store=FileStore(root),
|
|
event_sink=events.append,
|
|
)
|
|
|
|
|
|
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"
|
|
|
|
|
|
async 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 = await 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",
|
|
]
|
|
|
|
|
|
async 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_connection_auth,
|
|
emit_event=events.append,
|
|
)
|
|
source_catalog.hydrate_connection_source_from_snapshot(connection)
|
|
|
|
await 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]
|
|
|
|
|
|
async 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,
|
|
)
|
|
|
|
def _raise_missing_connection(connection_id: str) -> ConnectionConfig:
|
|
raise KeyError(connection_id)
|
|
|
|
source_catalog = SourceCatalogService(
|
|
store=transport.store,
|
|
connection_lookup=_raise_missing_connection,
|
|
connection_list_enabled=lambda: [],
|
|
connection_list_all=lambda: [],
|
|
tool_executor_for=transport.tool_executor_for,
|
|
load_auth=transport.load_connection_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 = await transport.deployment_diagnostics(
|
|
deployment=deployment,
|
|
artifacts=[artifact],
|
|
source_catalog=source_catalog,
|
|
)
|
|
|
|
assert diagnostics[0].code == "source_unreachable"
|
|
assert diagnostics[0].bound_source == "demo.personal"
|
|
|
|
|
|
def test_upstream_load_connection_auth_prefers_auth_ref(tmp_path: Path) -> None:
|
|
service = _transport(tmp_path)
|
|
service.save_auth(
|
|
AuthRecord(
|
|
connection_id="github.creds",
|
|
scheme="bearer",
|
|
payload={"token": "secret"},
|
|
)
|
|
)
|
|
service.save_auth(
|
|
AuthRecord(
|
|
connection_id="github.work",
|
|
scheme="bearer",
|
|
payload={"token": "wrong"},
|
|
)
|
|
)
|
|
connection = ConnectionConfig(
|
|
id="github.work",
|
|
server="github",
|
|
account="work",
|
|
metadata={"auth_ref": "github.creds"},
|
|
)
|
|
|
|
assert service.load_connection_auth(connection) == AuthRecord(
|
|
connection_id="github.creds",
|
|
scheme="bearer",
|
|
payload={"token": "secret"},
|
|
)
|
|
|
|
|
|
def test_upstream_load_connection_auth_falls_back_to_connection_id(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
service = _transport(tmp_path)
|
|
service.save_auth(
|
|
AuthRecord(
|
|
connection_id="github.work",
|
|
scheme="bearer",
|
|
payload={"token": "legacy"},
|
|
)
|
|
)
|
|
connection = ConnectionConfig(
|
|
id="github.work",
|
|
server="github",
|
|
account="work",
|
|
)
|
|
|
|
assert service.load_connection_auth(connection) == AuthRecord(
|
|
connection_id="github.work",
|
|
scheme="bearer",
|
|
payload={"token": "legacy"},
|
|
)
|
|
|
|
|
|
def test_upstream_load_connection_auth_ignores_non_string_auth_ref(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
service = _transport(tmp_path)
|
|
service.save_auth(
|
|
AuthRecord(
|
|
connection_id="github.work",
|
|
scheme="bearer",
|
|
payload={"token": "legacy"},
|
|
)
|
|
)
|
|
connection = ConnectionConfig(
|
|
id="github.work",
|
|
server="github",
|
|
account="work",
|
|
metadata={"auth_ref": 123},
|
|
)
|
|
|
|
assert service.load_connection_auth(connection) == AuthRecord(
|
|
connection_id="github.work",
|
|
scheme="bearer",
|
|
payload={"token": "legacy"},
|
|
)
|
|
|
|
|
|
async def test_upstream_transport_live_diagnostics_report_missing_auth_ref(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
events: list[McpEvent] = []
|
|
store = FileStore(tmp_path)
|
|
connections = ConnectionRegistry()
|
|
connection = ConnectionConfig(
|
|
id="github.work",
|
|
server="demo",
|
|
account="work",
|
|
metadata={"auth_ref": "github.creds"},
|
|
)
|
|
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_connection_auth,
|
|
emit_event=events.append,
|
|
)
|
|
source_catalog.register_capability_source(
|
|
CapabilitySource(
|
|
id="github.work",
|
|
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": "github.work"}],
|
|
)
|
|
|
|
diagnostics = await transport.deployment_diagnostics(
|
|
deployment=deployment,
|
|
artifacts=[artifact],
|
|
source_catalog=source_catalog,
|
|
)
|
|
|
|
assert diagnostics[0].code == "auth_not_found"
|
|
assert diagnostics[0].bound_source == "github.work"
|
|
assert "github.creds" in diagnostics[0].message
|