310 lines
11 KiB
Python
310 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
from wf_authoring import NodeSpec
|
|
from wf_mcp.broker import WfMcpService
|
|
from wf_mcp.models import ConnectionConfig
|
|
from wf_mcp.storage import FileStore
|
|
from wf_platform import (
|
|
CapabilityBuckets,
|
|
CapabilitySource,
|
|
SourceVisibility,
|
|
)
|
|
|
|
from ..test_support import (
|
|
FakeAdapter,
|
|
echo_tool,
|
|
finalize_tool,
|
|
local_temp_root,
|
|
)
|
|
from .conftest import raw_plan, single_echo_plan
|
|
|
|
|
|
async def test_service_compiles_and_runs_raw_plan() -> None:
|
|
service = WfMcpService(store=FileStore(local_temp_root() / "run_store"))
|
|
service.register_connection(
|
|
ConnectionConfig(id="demo.personal", server="demo", account="personal")
|
|
)
|
|
service.register_specs("demo.personal", echo_tool)
|
|
|
|
plan = raw_plan(
|
|
name="demo_plan",
|
|
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": [
|
|
{
|
|
"target": {"root": "local", "parts": ["text"]},
|
|
"path": {"root": "input", "parts": ["text"]},
|
|
}
|
|
],
|
|
"output": [
|
|
{
|
|
"source": {"root": "local", "parts": ["echoed"]},
|
|
"target": {"root": "state", "parts": ["echoed"]},
|
|
}
|
|
],
|
|
}
|
|
],
|
|
edges=[{"from": "echo", "outcome": "ok", "to": "__end__"}],
|
|
)
|
|
|
|
run = await service.run_workflow_from_plan(plan, {"text": "hello"})
|
|
|
|
assert run.status == "completed"
|
|
assert run.output["echoed"] == "hello"
|
|
|
|
|
|
async def test_service_preserves_raw_plan_root_output_bindings() -> None:
|
|
service = WfMcpService(store=FileStore(local_temp_root() / "root_output_store"))
|
|
service.register_connection(
|
|
ConnectionConfig(id="demo.personal", server="demo", account="personal")
|
|
)
|
|
service.register_specs("demo.personal", echo_tool)
|
|
|
|
plan = raw_plan(
|
|
name="root_output_plan",
|
|
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": [
|
|
{
|
|
"target": {"root": "local", "parts": ["text"]},
|
|
"path": {"root": "input", "parts": ["text"]},
|
|
}
|
|
],
|
|
"output": [
|
|
{
|
|
"source": {"root": "local", "parts": ["echoed"]},
|
|
"target": {"root": "state", "parts": ["echoed"]},
|
|
}
|
|
],
|
|
}
|
|
],
|
|
edges=[{"from": "echo", "outcome": "ok", "to": "__end__"}],
|
|
)
|
|
|
|
run = await service.run_workflow_from_plan(plan, {"text": "hello"})
|
|
|
|
assert run.output["echoed"] == "hello"
|
|
|
|
|
|
async def test_service_resolves_registered_spec_with_dotted_local_name() -> None:
|
|
service = WfMcpService(store=FileStore(local_temp_root() / "dotted_store"))
|
|
service.register_connection(
|
|
ConnectionConfig(id="demo.personal", server="demo", account="personal")
|
|
)
|
|
service.register_specs("demo.personal", echo_tool)
|
|
|
|
plan = single_echo_plan("dotted_plan", "demo.personal.echo_tool")
|
|
|
|
run = await service.run_workflow_from_plan(plan, {"text": "hello"})
|
|
|
|
assert run.status == "completed"
|
|
assert run.output["echoed"] == "hello"
|
|
|
|
|
|
async def test_service_runs_logical_source_plan_with_dotted_local_name() -> None:
|
|
import shutil
|
|
|
|
from wf_artifacts import WorkflowDeployment
|
|
from wf_authoring import NodeSpec
|
|
|
|
from ..test_support import echo_tool
|
|
|
|
store = local_temp_root() / "logical_source_store"
|
|
shutil.rmtree(store, ignore_errors=True)
|
|
service = WfMcpService(store=FileStore(store))
|
|
service.register_connection(
|
|
ConnectionConfig(id="demo.personal", server="demo", account="personal")
|
|
)
|
|
dotted_echo_tool = NodeSpec(
|
|
name="foo.bar",
|
|
input_model=echo_tool.input_model,
|
|
output_model=echo_tool.output_model,
|
|
outcomes=echo_tool.outcomes,
|
|
fn=echo_tool.fn,
|
|
description=echo_tool.description,
|
|
is_async=echo_tool.is_async,
|
|
accepts_context=echo_tool.accepts_context,
|
|
input_schema_contract=echo_tool.input_schema_contract,
|
|
output_schema_contract=echo_tool.output_schema_contract,
|
|
)
|
|
service.register_specs("demo.personal", dotted_echo_tool)
|
|
|
|
plan = single_echo_plan("logical_plan", "demo.foo.bar")
|
|
|
|
run = await service.run_workflow_from_plan(
|
|
plan,
|
|
{"text": "hello"},
|
|
deployment=WorkflowDeployment(
|
|
id="logical_dotted.personal",
|
|
artifact_id="logical_dotted",
|
|
artifact_version=1,
|
|
bindings=[{"logical_source": "demo", "concrete_source": "demo.personal"}],
|
|
),
|
|
)
|
|
|
|
assert run.status == "completed"
|
|
assert run.output["echoed"] == "hello"
|
|
|
|
|
|
async def test_service_binds_longest_logical_source_prefix_first() -> None:
|
|
service = WfMcpService(store=FileStore(local_temp_root() / "prefix_store"))
|
|
service.register_connection(
|
|
ConnectionConfig(id="demo.personal", server="demo", account="personal")
|
|
)
|
|
service.register_connection(
|
|
ConnectionConfig(id="demo.work", server="demo", account="work")
|
|
)
|
|
service.register_specs("demo.personal", echo_tool)
|
|
|
|
plan = single_echo_plan("prefix_plan", "demo.personal.echo_tool")
|
|
|
|
run = await service.run_workflow_from_plan(plan, {"text": "hello"})
|
|
|
|
assert run.status == "completed"
|
|
|
|
|
|
def test_service_does_not_resolve_specs_hidden_from_planner() -> None:
|
|
service = WfMcpService(store=FileStore(local_temp_root() / "hidden_resolve_store"))
|
|
hidden_echo_tool = NodeSpec(
|
|
name="hidden.source.echo_tool",
|
|
input_model=echo_tool.input_model,
|
|
output_model=echo_tool.output_model,
|
|
outcomes=echo_tool.outcomes,
|
|
fn=echo_tool.fn,
|
|
description=echo_tool.description,
|
|
is_async=echo_tool.is_async,
|
|
accepts_context=echo_tool.accepts_context,
|
|
input_schema_contract=echo_tool.input_schema_contract,
|
|
output_schema_contract=echo_tool.output_schema_contract,
|
|
)
|
|
service.register_capability_source(
|
|
CapabilitySource(
|
|
id="hidden.source",
|
|
kind="system",
|
|
capabilities=CapabilityBuckets(
|
|
node_specs={"hidden.source.echo_tool": hidden_echo_tool}
|
|
),
|
|
visibility=SourceVisibility(planner=False, admin_dashboard=False),
|
|
)
|
|
)
|
|
|
|
try:
|
|
service._get_qualified_spec("hidden.source.echo_tool")
|
|
except KeyError:
|
|
pass
|
|
else:
|
|
raise AssertionError("expected hidden spec to be unresolvable")
|
|
|
|
|
|
def test_service_excludes_disabled_connection_specs_from_planner_catalog() -> None:
|
|
service = WfMcpService(store=FileStore(local_temp_root() / "disabled_store"))
|
|
service.register_connection(
|
|
ConnectionConfig(id="demo.personal", server="demo", account="personal")
|
|
)
|
|
service.register_specs("demo.personal", echo_tool)
|
|
service.capability_sources["demo.personal"].enabled = False
|
|
|
|
planner_names = {
|
|
entry.qualified_name for entry in service.get_planner_catalog().entries()
|
|
}
|
|
assert "demo.personal.echo_tool" not in planner_names
|
|
|
|
|
|
def test_service_preserves_disabled_connection_source_on_reregistration() -> None:
|
|
service = WfMcpService(store=FileStore(local_temp_root() / "disabled_reregister"))
|
|
service.register_connection(
|
|
ConnectionConfig(id="demo.personal", server="demo", account="personal")
|
|
)
|
|
service.register_specs("demo.personal", echo_tool)
|
|
service.capability_sources["demo.personal"].enabled = False
|
|
service.register_specs("demo.personal", finalize_tool)
|
|
|
|
source = service.capability_sources["demo.personal"]
|
|
assert source.enabled is False
|
|
assert "demo.personal.finalize_tool" in source.capabilities.node_specs
|
|
assert "demo.personal.echo_tool" not in source.capabilities.node_specs
|
|
|
|
|
|
def test_service_excludes_planner_hidden_connection_specs_from_planner_catalog() -> (
|
|
None
|
|
):
|
|
service = WfMcpService(store=FileStore(local_temp_root() / "planner_hidden_store"))
|
|
service.register_connection(
|
|
ConnectionConfig(id="demo.personal", server="demo", account="personal")
|
|
)
|
|
service.register_specs("demo.personal", echo_tool)
|
|
service.capability_sources["demo.personal"].visibility = SourceVisibility(
|
|
planner=False,
|
|
client=True,
|
|
admin_dashboard=True,
|
|
)
|
|
|
|
planner_names = {
|
|
entry.qualified_name for entry in service.get_planner_catalog().entries()
|
|
}
|
|
assert "demo.personal.echo_tool" not in planner_names
|
|
|
|
|
|
def test_service_preserves_planner_hidden_connection_source_on_reregistration() -> None:
|
|
service = WfMcpService(store=FileStore(local_temp_root() / "planner_hidden_rereg"))
|
|
service.register_connection(
|
|
ConnectionConfig(id="demo.personal", server="demo", account="personal")
|
|
)
|
|
service.register_specs("demo.personal", echo_tool)
|
|
service.capability_sources["demo.personal"].visibility = SourceVisibility(
|
|
planner=False,
|
|
client=True,
|
|
admin_dashboard=True,
|
|
)
|
|
service.register_specs("demo.personal", finalize_tool)
|
|
|
|
source = service.capability_sources["demo.personal"]
|
|
assert source.visibility.planner is False
|
|
assert source.visibility.client is True
|
|
assert source.visibility.admin_dashboard is True
|
|
assert "demo.personal.finalize_tool" in source.capabilities.node_specs
|
|
assert "demo.personal.echo_tool" not in source.capabilities.node_specs
|
|
|
|
|
|
async def test_service_refreshes_catalog_from_adapter() -> None:
|
|
service = WfMcpService(store=FileStore(local_temp_root() / "refresh_store"))
|
|
service.register_connection(
|
|
ConnectionConfig(id="demo.personal", server="demo", account="personal")
|
|
)
|
|
service.register_adapter("demo", FakeAdapter())
|
|
|
|
await service.refresh_connection_catalog("demo.personal")
|
|
|
|
source = service.capability_sources["demo.personal"]
|
|
assert "demo.personal.echo_tool" in source.capabilities.node_specs
|