wire up stuff to the MCP server
This commit is contained in:
@@ -0,0 +1,83 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
from wf_artifacts import (
|
||||
AvailableCapability,
|
||||
AvailableSource,
|
||||
validate_deployment_dependencies,
|
||||
)
|
||||
|
||||
from .service import WfMcpService
|
||||
|
||||
|
||||
def register_artifact_tools(server: FastMCP, service: WfMcpService) -> None:
|
||||
"""Register stable MCP tools for saved workflow artifact inspection."""
|
||||
|
||||
@server.tool()
|
||||
async def list_workflow_artifacts() -> dict[str, Any]:
|
||||
if service.artifact_store is None:
|
||||
return {"nodes": []}
|
||||
entries = [
|
||||
service.workflow_artifact_catalog_entry(artifact).model_dump(mode="json")
|
||||
for artifact in service.artifact_store.list_artifacts()
|
||||
]
|
||||
return {"nodes": entries}
|
||||
|
||||
@server.tool()
|
||||
async def inspect_workflow_artifact(
|
||||
artifact_id: str,
|
||||
version: int,
|
||||
) -> dict[str, Any]:
|
||||
if service.artifact_store is None:
|
||||
raise KeyError("workflow artifact store is not configured")
|
||||
artifact = service.artifact_store.get_artifact(artifact_id, version)
|
||||
return artifact.model_dump(mode="json")
|
||||
|
||||
@server.tool()
|
||||
async def validate_workflow_deployment(deployment_id: str) -> dict[str, Any]:
|
||||
if service.artifact_store is None:
|
||||
raise KeyError("workflow artifact store is not configured")
|
||||
deployment = service.artifact_store.get_deployment(deployment_id)
|
||||
artifact = service.artifact_store.get_artifact(
|
||||
deployment.artifact_id,
|
||||
deployment.artifact_version,
|
||||
)
|
||||
diagnostics = validate_deployment_dependencies(
|
||||
artifact=artifact,
|
||||
deployment=deployment,
|
||||
sources=_available_sources(service),
|
||||
)
|
||||
return {
|
||||
"deployment_id": deployment.id,
|
||||
"artifact_id": artifact.id,
|
||||
"artifact_version": artifact.version,
|
||||
"status": "unrunnable" if diagnostics else "runnable",
|
||||
"diagnostics": [
|
||||
diagnostic.model_dump(mode="json") for diagnostic in diagnostics
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _available_sources(service: WfMcpService) -> list[AvailableSource]:
|
||||
"""Convert broker capability sources into artifact validation snapshots."""
|
||||
sources: list[AvailableSource] = []
|
||||
for source in service.capability_sources.values():
|
||||
capabilities = {
|
||||
spec.name.rsplit(".", maxsplit=1)[-1]: AvailableCapability(
|
||||
name=spec.name.rsplit(".", maxsplit=1)[-1],
|
||||
kind="node_spec",
|
||||
input_schema_hash=None,
|
||||
output_schema_hash=None,
|
||||
)
|
||||
for spec in source.capabilities.node_specs.values()
|
||||
}
|
||||
sources.append(
|
||||
AvailableSource(
|
||||
id=source.id,
|
||||
enabled=source.enabled,
|
||||
capabilities=capabilities,
|
||||
)
|
||||
)
|
||||
return sources
|
||||
@@ -7,6 +7,7 @@ from ..control import BrokerConfigFile
|
||||
from ..models import BrokerConfig
|
||||
from ..sdk import McpSdkAdapter
|
||||
from ..storage import FileStore
|
||||
from wf_artifacts import FileWorkflowArtifactStore
|
||||
from .service import WfMcpService
|
||||
|
||||
|
||||
@@ -19,7 +20,10 @@ def load_broker_config(path: str | Path) -> BrokerConfig:
|
||||
|
||||
def build_service_from_config(config: BrokerConfig) -> WfMcpService:
|
||||
"""Create a broker service with SDK adapters for configured connections."""
|
||||
service = WfMcpService(store=FileStore(config.store_root))
|
||||
service = WfMcpService(
|
||||
store=FileStore(config.store_root),
|
||||
artifact_store=FileWorkflowArtifactStore(config.store_root),
|
||||
)
|
||||
for connection in config.connections:
|
||||
service.register_connection(connection)
|
||||
if connection.server not in service.adapters:
|
||||
|
||||
@@ -6,6 +6,7 @@ from pathlib import Path
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
|
||||
from ..transparent_proxy import create_transparent_proxy_server
|
||||
from .artifact_tools import register_artifact_tools
|
||||
from .config import build_service_from_config, load_broker_config
|
||||
from .prompts import register_broker_prompts
|
||||
from .resources import register_broker_resources
|
||||
@@ -25,6 +26,7 @@ def create_broker_server(service: WfMcpService) -> FastMCP:
|
||||
)
|
||||
|
||||
register_broker_tools(server, service)
|
||||
register_artifact_tools(server, service)
|
||||
register_broker_resources(server, service)
|
||||
register_broker_prompts(server, service)
|
||||
return server
|
||||
|
||||
@@ -2,11 +2,19 @@ from __future__ import annotations
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from wf_authoring import NodeReturn, NodeSpec, build_async_registry
|
||||
from wf_artifacts import (
|
||||
FileWorkflowArtifactStore,
|
||||
WorkflowArtifact,
|
||||
WorkflowArtifactCatalogEntry,
|
||||
WorkflowArtifactStore,
|
||||
artifact_catalog_entry,
|
||||
)
|
||||
from wf_core import NodeUse, Workflow, execute_workflow_async
|
||||
|
||||
from ...connections import ConnectionRegistry, parse_connection_id, qualify_node_name
|
||||
@@ -40,6 +48,12 @@ from .sources import SpecSource
|
||||
from .specs import get_qualified_spec, qualify_spec
|
||||
|
||||
|
||||
def _store_root(store: Store) -> Path:
|
||||
"""Return the file root for stores that expose one, else use local default."""
|
||||
root = getattr(store, "root", None)
|
||||
return root if isinstance(root, Path) else Path(".wf_mcp_store")
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class WfMcpService:
|
||||
store: Store
|
||||
@@ -49,9 +63,12 @@ class WfMcpService:
|
||||
capability_sources: dict[str, CapabilitySource] = field(default_factory=dict)
|
||||
events: list[McpEvent] = field(default_factory=list)
|
||||
include_builtin_specs: bool = True
|
||||
artifact_store: WorkflowArtifactStore | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
"""Install broker-local system specs when enabled."""
|
||||
if self.artifact_store is None:
|
||||
self.artifact_store = FileWorkflowArtifactStore(_store_root(self.store))
|
||||
if self.include_builtin_specs:
|
||||
for source in builtin_sources(self).values():
|
||||
self.register_spec_source(source)
|
||||
@@ -230,6 +247,13 @@ class WfMcpService:
|
||||
"""Return planner-visible node catalog entries from every visible source."""
|
||||
return self.get_planner_catalog().entries()
|
||||
|
||||
def workflow_artifact_catalog_entry(
|
||||
self,
|
||||
artifact: WorkflowArtifact,
|
||||
) -> WorkflowArtifactCatalogEntry:
|
||||
"""Project a saved workflow artifact as a planner catalog entry."""
|
||||
return artifact_catalog_entry(artifact)
|
||||
|
||||
def get_connection_snapshot(self, connection_id: str) -> CatalogSnapshot | None:
|
||||
self.connections.get(connection_id)
|
||||
return self.store.load_catalog(connection_id)
|
||||
|
||||
@@ -4,6 +4,12 @@ import asyncio
|
||||
import json
|
||||
from typing import Any, cast
|
||||
|
||||
from wf_artifacts import (
|
||||
FileWorkflowArtifactStore,
|
||||
RequiredCapability,
|
||||
WorkflowArtifact,
|
||||
WorkflowDeployment,
|
||||
)
|
||||
from wf_mcp.broker import (
|
||||
WfMcpService,
|
||||
build_service_from_config,
|
||||
@@ -174,3 +180,114 @@ def test_broker_call_tool_returns_structured_result() -> None:
|
||||
"output": {"echoed": "hello"},
|
||||
"meta": {},
|
||||
}
|
||||
|
||||
|
||||
def test_broker_lists_workflow_artifacts_from_artifact_store() -> None:
|
||||
artifact_store = FileWorkflowArtifactStore(local_temp_root() / "broker_artifacts")
|
||||
artifact_store.save_artifact(_artifact())
|
||||
service = WfMcpService(
|
||||
store=FileStore(local_temp_root() / "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() -> None:
|
||||
artifact_store = FileWorkflowArtifactStore(
|
||||
local_temp_root() / "broker_inspect_artifacts"
|
||||
)
|
||||
artifact_store.save_artifact(_artifact())
|
||||
service = WfMcpService(
|
||||
store=FileStore(local_temp_root() / "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() -> None:
|
||||
artifact_store = FileWorkflowArtifactStore(
|
||||
local_temp_root() / "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={"context7": "context7.personal"},
|
||||
)
|
||||
)
|
||||
service = WfMcpService(
|
||||
store=FileStore(local_temp_root() / "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_build_service_from_config_uses_store_root_for_artifacts() -> None:
|
||||
store_root = local_temp_root() / "broker_config_artifact_store"
|
||||
service = build_service_from_config(
|
||||
BrokerConfig(
|
||||
store_root=store_root,
|
||||
connections=[],
|
||||
)
|
||||
)
|
||||
|
||||
assert isinstance(service.artifact_store, FileWorkflowArtifactStore)
|
||||
assert service.artifact_store.root == store_root
|
||||
|
||||
|
||||
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(
|
||||
logical_source="context7",
|
||||
capability_name="query-docs",
|
||||
kind="tool",
|
||||
input_schema_hash="sha256:input",
|
||||
output_schema_hash="sha256:output",
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user