Files
lda-wf/tests/wf_mcp/test_broker_server.py
T

708 lines
23 KiB
Python

from __future__ import annotations
import asyncio
import json
from pathlib import Path
from typing import Any, cast
from wf_artifacts import (
FileDraftWorkspaceStore,
FileRunStore,
FileWorkflowArtifactStore,
RequiredCapability,
WorkflowArtifact,
WorkflowDeployment,
)
from wf_mcp.broker import (
WfMcpService,
build_service_from_config,
create_broker_server,
load_broker_config,
)
from wf_mcp.models import BrokerConfig, ConnectionConfig
from wf_mcp.source_registry import (
FileSourceRegistryStore,
McpSourceRegistryEntry,
SourceRegistryFile,
StdioSourceTransport,
)
from wf_mcp.storage import FileStore
from .test_support import (
FailingDiscoveryAdapter,
FakeAdapter,
echo_tool,
input_binding,
output_binding,
)
def test_load_broker_config_resolves_relative_store_root(tmp_path: Path) -> None:
tmp_path = tmp_path / "broker_config_test"
tmp_path.mkdir(parents=True, exist_ok=True)
config_path = tmp_path / "wf_mcp.config.json"
config_path.write_text(
json.dumps(
{
"store_root": ".broker-store",
"connections": [
{
"id": "demo.personal",
"server": "demo",
"account": "personal",
}
],
}
),
encoding="utf-8",
)
config = load_broker_config(config_path)
assert config.store_root == (tmp_path / ".broker-store").resolve()
assert [connection.id for connection in config.connections] == ["demo.personal"]
def test_create_broker_server_exposes_tools_resources_and_prompts(
tmp_path: Path,
) -> None:
service = WfMcpService(store=FileStore(tmp_path / "broker_server_store"))
service.register_connection(
ConnectionConfig(id="demo.personal", server="demo", account="personal")
)
service.register_adapter("demo", FakeAdapter())
asyncio.run(service.refresh_connection_catalog("demo.personal"))
server = create_broker_server(service)
tools = asyncio.run(server.list_tools())
resources = asyncio.run(server.list_resources())
prompts = asyncio.run(server.list_prompts())
tool_names = {tool.name for tool in tools}
resource_names = {resource.name for resource in resources}
prompt_names = {prompt.name for prompt in prompts}
assert "get_connection_statuses" in tool_names
assert "refresh_connection_catalog" in tool_names
assert "get_planner_catalog" in tool_names
assert "list_sources" in tool_names
assert "invoke_broker_method" in tool_names
assert "catalog.all" in resource_names
assert "events.all" in resource_names
assert "status.all" in resource_names
assert "workflow_authoring_guide" in prompt_names
assert "plan_with_catalog" not in prompt_names
_content, planner_catalog_raw = asyncio.run(
server.call_tool("get_planner_catalog", {})
)
planner_catalog = cast(dict[str, Any], cast(object, planner_catalog_raw))
planner_names = [node["qualified_name"] for node in planner_catalog["nodes"]]
assert "demo.personal.echo_tool" in planner_names
assert "wf.std.runtime_error" in planner_names
_content, all_sources_payload_raw = asyncio.run(
server.call_tool("list_sources", {})
)
all_sources_payload = cast(dict[str, Any], cast(object, all_sources_payload_raw))
all_sources = all_sources_payload["sources"]
all_source_ids = {source["id"] for source in all_sources}
assert "wf.admin" in all_source_ids
assert "demo.personal" in all_source_ids
def test_broker_admin_tools_are_backed_by_wf_admin_source(tmp_path: Path) -> None:
service = WfMcpService(store=FileStore(tmp_path / "broker_admin_source"))
server = create_broker_server(service)
tools = asyncio.run(server.list_tools())
tool_names = {tool.name for tool in tools}
assert "get_planner_catalog" in tool_names
assert (
"wf.admin.list_sources"
in service.capability_sources["wf.admin"].capabilities.tools
)
def test_build_service_from_config_registers_connections(tmp_path: Path) -> None:
config = BrokerConfig(
store_root=tmp_path / "broker_config_store",
connections=[
ConnectionConfig(id="demo.personal", server="demo", account="personal"),
ConnectionConfig(id="demo.work", server="demo", account="work"),
],
)
service = build_service_from_config(config)
ids = [connection.id for connection in service.connections.list_all()]
assert ids == ["demo.personal", "demo.work"]
def test_broker_refresh_tool_returns_structured_error(tmp_path: Path) -> None:
service = WfMcpService(store=FileStore(tmp_path / "broker_fail_store"))
service.register_connection(
ConnectionConfig(id="demo.personal", server="demo", account="personal")
)
service.register_adapter("demo", FailingDiscoveryAdapter())
server = create_broker_server(service)
_content, structured = asyncio.run(
server.call_tool(
"refresh_connection_catalog", {"connection_id": "demo.personal"}
)
)
assert structured == {
"connection_id": "demo.personal",
"refreshed": False,
"error_type": "PermissionError",
"error": "Access is denied",
}
def test_broker_lists_workflow_artifacts_from_artifact_store(tmp_path: Path) -> None:
artifact_store = FileWorkflowArtifactStore(tmp_path / "broker_artifacts")
artifact_store.save_artifact(_artifact())
service = WfMcpService(
store=FileStore(tmp_path / "broker_artifacts_mcp_store"),
artifact_store=artifact_store,
)
server = create_broker_server(service)
_content, structured = asyncio.run(server.call_tool("list_workflow_artifacts", {}))
payload = cast(dict[str, Any], cast(object, structured))
nodes = payload["nodes"]
assert len(nodes) == 1
assert nodes[0]["name"] == "workflow.summarize_docs.v1"
assert nodes[0]["required_sources"] == ["context7"]
assert "plan" not in nodes[0]
def test_broker_inspects_workflow_artifact_from_artifact_store(tmp_path: Path) -> None:
artifact_store = FileWorkflowArtifactStore(tmp_path / "broker_inspect_artifacts")
artifact_store.save_artifact(_artifact())
service = WfMcpService(
store=FileStore(tmp_path / "broker_inspect_mcp_store"),
artifact_store=artifact_store,
)
server = create_broker_server(service)
_content, structured = asyncio.run(
server.call_tool(
"inspect_workflow_artifact",
{"artifact_id": "summarize_docs", "version": 1},
)
)
artifact = cast(dict[str, Any], cast(object, structured))
assert artifact["id"] == "summarize_docs"
assert artifact["version"] == 1
assert artifact["plan"]["name"] == "summarize_docs"
def test_broker_validates_workflow_deployment_from_artifact_store(
tmp_path: Path,
) -> None:
artifact_store = FileWorkflowArtifactStore(tmp_path / "broker_validate_artifacts")
artifact_store.save_artifact(_artifact())
artifact_store.save_deployment(
WorkflowDeployment(
id="summarize_docs.personal",
artifact_id="summarize_docs",
artifact_version=1,
bindings=[
{"logical_source": "context7", "concrete_source": "context7.personal"}
],
)
)
service = WfMcpService(
store=FileStore(tmp_path / "broker_validate_mcp_store"),
artifact_store=artifact_store,
)
server = create_broker_server(service)
_content, structured = asyncio.run(
server.call_tool(
"validate_workflow_deployment",
{"deployment_id": "summarize_docs.personal"},
)
)
payload = cast(dict[str, Any], cast(object, structured))
assert payload["deployment_id"] == "summarize_docs.personal"
assert payload["artifact_id"] == "summarize_docs"
assert payload["status"] == "unrunnable"
assert payload["diagnostics"][0]["code"] == "source_missing"
def test_broker_saves_workflow_artifact(tmp_path: Path) -> None:
artifact_store = FileWorkflowArtifactStore(tmp_path / "broker_save_artifacts")
service = WfMcpService(
store=FileStore(tmp_path / "broker_save_mcp_store"),
artifact_store=artifact_store,
)
server = create_broker_server(service)
_content, structured = asyncio.run(
server.call_tool(
"save_workflow_artifact",
{"artifact": _artifact().model_dump(mode="json")},
)
)
payload = cast(dict[str, Any], cast(object, structured))
loaded = artifact_store.get_artifact("summarize_docs", 1)
assert payload["artifact_id"] == "summarize_docs"
assert payload["version"] == 1
assert loaded.title == "Summarize Docs"
def test_broker_creates_workflow_artifact_from_plan(tmp_path: Path) -> None:
artifact_store = FileWorkflowArtifactStore(
tmp_path / "broker_create_artifact_from_plan"
)
service = WfMcpService(
store=FileStore(tmp_path / "broker_create_artifact_mcp_store"),
artifact_store=artifact_store,
)
server = create_broker_server(service)
_content, structured = asyncio.run(
server.call_tool(
"create_workflow_artifact_from_plan",
{
"artifact_id": "echo",
"version": 1,
"title": "Echo",
"description": "Echo through a saved plan.",
"plan": _echo_artifact().plan,
"outcomes": ["done"],
"required_capabilities": {
"demo.echo_tool": {
"logical_source": "demo",
"capability_name": "echo_tool",
"kind": "node_spec",
}
},
"created_from_catalog_version": "catalog-1",
},
)
)
payload = cast(dict[str, Any], cast(object, structured))
loaded = artifact_store.get_artifact("echo", 1)
assert payload["artifact_id"] == "echo"
assert payload["version"] == 1
assert payload["saved"] is True
assert loaded.input_schema["properties"]["text"]["type"] == "string"
assert loaded.output_schema["properties"]["echoed"]["type"] == "string"
assert loaded.required_capability_map()["demo.echo_tool"].logical_source == "demo"
def test_broker_saves_and_lists_workflow_deployments(tmp_path: Path) -> None:
artifact_store = FileWorkflowArtifactStore(tmp_path / "broker_save_deployments")
service = WfMcpService(
store=FileStore(tmp_path / "broker_save_deployments_mcp_store"),
artifact_store=artifact_store,
)
server = create_broker_server(service)
_content, save_structured = asyncio.run(
server.call_tool(
"save_workflow_deployment",
{
"deployment": WorkflowDeployment(
id="summarize_docs.personal",
artifact_id="summarize_docs",
artifact_version=1,
bindings=[
{
"logical_source": "context7",
"concrete_source": "context7.personal",
}
],
).model_dump(mode="json")
},
)
)
save_payload = cast(dict[str, Any], cast(object, save_structured))
_content, list_structured = asyncio.run(
server.call_tool("list_workflow_deployments", {})
)
list_payload = cast(dict[str, Any], cast(object, list_structured))
assert save_payload["deployment_id"] == "summarize_docs.personal"
assert list_payload["deployments"][0]["id"] == "summarize_docs.personal"
assert "bindings" not in list_payload["deployments"][0]
def test_broker_runs_non_interrupting_workflow_deployment(tmp_path: Path) -> None:
artifact_store = FileWorkflowArtifactStore(tmp_path / "broker_run_artifacts")
artifact_store.save_artifact(_echo_artifact())
artifact_store.save_deployment(
WorkflowDeployment(
id="echo.personal",
artifact_id="echo",
artifact_version=1,
bindings=[{"logical_source": "demo", "concrete_source": "demo.personal"}],
)
)
service = WfMcpService(
store=FileStore(tmp_path / "broker_run_mcp_store"),
artifact_store=artifact_store,
run_store=FileRunStore(tmp_path / "broker_run_mcp_store"),
)
service.register_connection(
ConnectionConfig(id="demo.personal", server="demo", account="personal")
)
service.register_specs("demo.personal", echo_tool)
server = create_broker_server(service)
_content, structured = asyncio.run(
server.call_tool(
"run_workflow_deployment",
{
"deployment_id": "echo.personal",
"workflow_input": {"text": "hello"},
},
)
)
payload = cast(dict[str, Any], cast(object, structured))
assert payload["deployment_id"] == "echo.personal"
assert payload["artifact_id"] == "echo"
assert payload["status"] == "completed"
assert payload["output"]["echoed"] == "hello"
assert payload["diagnostics"] == []
assert payload["trace_count"] > 0
def test_broker_run_deployment_returns_unrunnable_for_dependency_errors(
tmp_path: Path,
) -> None:
artifact_store = FileWorkflowArtifactStore(
tmp_path / "broker_run_unrunnable_artifacts"
)
artifact_store.save_artifact(_artifact())
artifact_store.save_deployment(
WorkflowDeployment(
id="summarize_docs.personal",
artifact_id="summarize_docs",
artifact_version=1,
bindings=[
{"logical_source": "context7", "concrete_source": "context7.personal"}
],
)
)
service = WfMcpService(
store=FileStore(tmp_path / "broker_run_unrunnable_mcp_store"),
artifact_store=artifact_store,
)
server = create_broker_server(service)
_content, structured = asyncio.run(
server.call_tool(
"run_workflow_deployment",
{
"deployment_id": "summarize_docs.personal",
"workflow_input": {},
},
)
)
payload = cast(dict[str, Any], cast(object, structured))
assert payload["status"] == "unrunnable"
assert payload["output"] is None
assert payload["diagnostics"][0]["code"] == "source_missing"
def test_broker_run_deployment_pauses_and_resumes_interrupting_artifacts(
tmp_path: Path,
) -> None:
artifact_store = FileWorkflowArtifactStore(
tmp_path / "broker_run_interrupt_artifacts"
)
artifact_store.save_artifact(_interrupt_artifact())
artifact_store.save_deployment(
WorkflowDeployment(
id="approval.personal",
artifact_id="approval",
artifact_version=1,
bindings=[],
)
)
service = WfMcpService(
store=FileStore(tmp_path / "broker_run_interrupt_mcp_store"),
artifact_store=artifact_store,
run_store=FileRunStore(tmp_path / "broker_run_interrupt_mcp_store"),
)
server = create_broker_server(service)
_content, structured = asyncio.run(
server.call_tool(
"run_workflow_deployment",
{
"deployment_id": "approval.personal",
"workflow_input": {"message": "send?"},
},
)
)
payload = cast(dict[str, Any], cast(object, structured))
assert payload["status"] == "interrupted"
assert payload["output"] == {}
assert isinstance(payload["run_id"], str)
assert payload["interrupt"]["payload"]["message"] == "send?"
_content, structured = asyncio.run(
server.call_tool(
"resume_workflow_run",
{
"run_id": payload["run_id"],
"resume_payload": {},
},
)
)
resumed = cast(dict[str, Any], cast(object, structured))
assert resumed["status"] == "completed"
assert resumed["outcome"] == "submitted"
assert resumed["run_id"] == payload["run_id"]
assert resumed["resume_readiness"] == "not_applicable"
def test_build_service_from_config_uses_store_root_for_workflow_stores(
tmp_path: Path,
) -> None:
store_root = tmp_path / "broker_config_workflow_stores"
config = BrokerConfig(store_root=store_root, connections=[])
service = build_service_from_config(config)
assert isinstance(service.artifact_store, FileWorkflowArtifactStore)
assert isinstance(service.draft_workspace_store, FileDraftWorkspaceStore)
assert isinstance(service.run_store, FileRunStore)
assert service.artifact_store.root == store_root
assert service.draft_workspace_store.root == store_root
assert service.run_store.root == store_root
# ---------------------------------------------------------------------------
# Source registry integration tests
# ---------------------------------------------------------------------------
def _registry_entry(
source_id: str = "demo.registry",
*,
enabled: bool = True,
) -> McpSourceRegistryEntry:
return McpSourceRegistryEntry(
id=source_id,
kind="mcp",
enabled=enabled,
provider="demo",
account=source_id.rsplit(".", 1)[-1],
transport=StdioSourceTransport(command="demo-server"),
)
def test_build_service_from_config_loads_source_registry_entries(
tmp_path: Path,
) -> None:
tmp_path = tmp_path / "broker_config_registry_load"
tmp_path.mkdir(parents=True, exist_ok=True)
config = BrokerConfig(store_root=tmp_path, connections=[])
FileSourceRegistryStore(tmp_path).save_registry(
SourceRegistryFile(sources=[_registry_entry("fixture.registry")])
)
service = build_service_from_config(config)
assert service.connections.get("fixture.registry").server == "demo"
assert "demo" in service.adapters
assert "fixture.registry" in service.capability_sources
def test_build_service_from_config_config_shadows_registry(tmp_path: Path) -> None:
tmp_path = tmp_path / "broker_config_registry_shadow"
tmp_path.mkdir(parents=True, exist_ok=True)
config = BrokerConfig(
store_root=tmp_path,
connections=[
ConnectionConfig(id="fixture.same", server="fixture", account="config"),
],
)
FileSourceRegistryStore(tmp_path).save_registry(
SourceRegistryFile(sources=[_registry_entry("fixture.same")])
)
service = build_service_from_config(config)
assert service.connections.get("fixture.same").account == "config"
assert any(
event.kind == "source_registry_ignored_config_shadow"
and event.connection_id == "fixture.same"
for event in service.list_events()
)
def test_broker_config_connection_defaults_to_locked(tmp_path: Path) -> None:
config_path = tmp_path / "wf_mcp.config.json"
config_path.write_text(
json.dumps(
{
"store_root": str(tmp_path / "store"),
"connections": [
{"id": "demo.default", "server": "demo", "account": "default"}
],
}
),
encoding="utf-8",
)
config = load_broker_config(config_path)
assert config.connections[0].source_config_ownership == "locked"
def test_broker_config_connection_accepts_seed_policy(tmp_path: Path) -> None:
config_path = tmp_path / "wf_mcp.config.json"
config_path.write_text(
json.dumps(
{
"store_root": str(tmp_path / "store"),
"connections": [
{
"id": "demo.default",
"server": "demo",
"account": "default",
"source_config_ownership": "seed",
}
],
}
),
encoding="utf-8",
)
config = load_broker_config(config_path)
assert config.connections[0].source_config_ownership == "seed"
def _artifact() -> WorkflowArtifact:
return WorkflowArtifact(
id="summarize_docs",
version=1,
title="Summarize Docs",
description="Summarize retrieved documentation.",
input_schema={"type": "object", "properties": {}},
output_schema={"type": "object", "properties": {}},
outcomes=("done",),
plan={"name": "summarize_docs", "nodes": [], "edges": []},
required_capabilities={
"context7.query-docs": RequiredCapability(
ref="context7.query-docs",
kind="tool",
input_schema_hash="sha256:input",
output_schema_hash="sha256:output",
)
},
)
def _echo_artifact() -> WorkflowArtifact:
return WorkflowArtifact(
id="echo",
version=1,
title="Echo",
description="Echo text through a demo capability.",
input_schema={
"type": "object",
"properties": {"text": {"type": "string"}},
"required": ["text"],
},
output_schema={
"type": "object",
"properties": {"echoed": {"type": "string"}},
"required": ["echoed"],
},
outcomes=("completed",),
plan={
"name": "echo",
"input_schema": {
"type": "object",
"properties": {"text": {"type": "string"}},
"required": ["text"],
},
"state_schema": {"fields": {"echoed": {"type": "string"}}},
"output_schema": {
"type": "object",
"properties": {"echoed": {"type": "string"}},
"required": ["echoed"],
},
"start": "echo",
"nodes": [
{
"id": "echo",
"type": "node",
"node": "demo.personal.echo_tool",
"input": [input_binding("input.text", "text")],
"output": [output_binding("echoed", "state.echoed")],
}
],
"edges": [{"from": "echo", "outcome": "ok", "to": "__end__"}],
},
required_capabilities={
"demo.echo_tool": RequiredCapability(
ref="demo.echo_tool",
kind="node_spec",
)
},
)
def _interrupt_artifact() -> WorkflowArtifact:
return WorkflowArtifact(
id="approval",
version=1,
title="Approval",
input_schema={
"type": "object",
"properties": {"message": {"type": "string"}},
"required": ["message"],
},
output_schema={"type": "object", "properties": {}},
outcomes=("submitted",),
plan={
"name": "approval",
"input_schema": {
"type": "object",
"properties": {"message": {"type": "string"}},
"required": ["message"],
},
"state_schema": {"fields": {}},
"output_schema": {"type": "object", "properties": {}},
"outcomes": ["submitted"],
"start": "approval",
"nodes": [
{
"id": "approval",
"type": "interrupt",
"kind": "approval",
"request": [input_binding("input.message", "message")],
"resume": [],
"outcomes": ["submitted"],
},
{"id": "end_submitted", "type": "end", "outcome": "submitted"},
],
"edges": [
{"from": "approval", "outcome": "submitted", "to": "end_submitted"}
],
},
)