and thats a server we can use
This commit is contained in:
@@ -0,0 +1,69 @@
|
||||
from .adapters import (
|
||||
BackendAdapter,
|
||||
ToolCallResult,
|
||||
)
|
||||
from .broker_server import (
|
||||
build_service_from_config,
|
||||
create_broker_server,
|
||||
load_broker_config,
|
||||
)
|
||||
from .capabilities import (
|
||||
CatalogNodeEntry,
|
||||
CatalogPromptEntry,
|
||||
CatalogResourceEntry,
|
||||
DiscoveredPrompt,
|
||||
DiscoveredResource,
|
||||
DiscoveredTool,
|
||||
)
|
||||
from .catalog import CombinedCatalog
|
||||
from .connections import ConnectionRegistry, parse_connection_id, qualify_node_name
|
||||
from .discovery import (
|
||||
DiscoveredConnectionCapabilities,
|
||||
discover_connection_capabilities,
|
||||
specs_from_discovered_tools,
|
||||
)
|
||||
from .events import McpEvent, make_event
|
||||
from .models import (
|
||||
AuthRecord,
|
||||
BrokerConfig,
|
||||
CatalogSnapshot,
|
||||
ConnectionConfig,
|
||||
RawWorkflowPlan,
|
||||
)
|
||||
from .mcp_sdk_adapter import McpSdkAdapter
|
||||
from .service import WfMcpService
|
||||
from .store import FileStore, Store
|
||||
from .wrappers import wrap_discovered_tool
|
||||
|
||||
__all__ = [
|
||||
"AuthRecord",
|
||||
"BackendAdapter",
|
||||
"BrokerConfig",
|
||||
"CatalogNodeEntry",
|
||||
"CatalogPromptEntry",
|
||||
"CatalogResourceEntry",
|
||||
"CatalogSnapshot",
|
||||
"CombinedCatalog",
|
||||
"ConnectionConfig",
|
||||
"ConnectionRegistry",
|
||||
"DiscoveredConnectionCapabilities",
|
||||
"DiscoveredPrompt",
|
||||
"DiscoveredResource",
|
||||
"DiscoveredTool",
|
||||
"FileStore",
|
||||
"McpEvent",
|
||||
"McpSdkAdapter",
|
||||
"RawWorkflowPlan",
|
||||
"Store",
|
||||
"ToolCallResult",
|
||||
"WfMcpService",
|
||||
"build_service_from_config",
|
||||
"create_broker_server",
|
||||
"discover_connection_capabilities",
|
||||
"load_broker_config",
|
||||
"make_event",
|
||||
"parse_connection_id",
|
||||
"qualify_node_name",
|
||||
"specs_from_discovered_tools",
|
||||
"wrap_discovered_tool",
|
||||
]
|
||||
@@ -0,0 +1,79 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Protocol
|
||||
|
||||
from .capabilities import DiscoveredPrompt, DiscoveredResource, DiscoveredTool
|
||||
from .models import AuthRecord, ConnectionConfig
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ToolCallResult:
|
||||
outcome: str
|
||||
output: dict[str, Any] = field(default_factory=dict)
|
||||
meta: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
class BackendAdapter(Protocol):
|
||||
async def list_tools(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
) -> list[DiscoveredTool]: ...
|
||||
|
||||
async def list_resources(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
) -> list[DiscoveredResource]: ...
|
||||
|
||||
async def list_prompts(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
) -> list[DiscoveredPrompt]: ...
|
||||
|
||||
async def get_connection_metadata(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
) -> dict[str, Any]: ...
|
||||
|
||||
async def read_resource(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
uri: str,
|
||||
) -> dict[str, Any]: ...
|
||||
|
||||
async def get_prompt(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
prompt_name: str,
|
||||
arguments: dict[str, str] | None = None,
|
||||
) -> dict[str, Any]: ...
|
||||
|
||||
async def invoke_method(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
method: str,
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]: ...
|
||||
|
||||
async def send_notification(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
method: str,
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> None: ...
|
||||
|
||||
async def call_tool(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
tool_name: str,
|
||||
payload: dict[str, Any],
|
||||
) -> ToolCallResult: ...
|
||||
@@ -0,0 +1,176 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import asdict
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal, cast
|
||||
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
|
||||
from .mcp_sdk_adapter import McpSdkAdapter
|
||||
from .models import BrokerConfig, ConnectionConfig
|
||||
from .service import WfMcpService
|
||||
from .store import FileStore
|
||||
|
||||
|
||||
def load_broker_config(path: str | Path) -> BrokerConfig:
|
||||
config_path = Path(path)
|
||||
data = json.loads(config_path.read_text(encoding="utf-8"))
|
||||
store_root_raw = data.get("store_root", ".wf_mcp_store")
|
||||
store_root = Path(store_root_raw)
|
||||
if not store_root.is_absolute():
|
||||
store_root = (config_path.parent / store_root).resolve()
|
||||
|
||||
connections = [ConnectionConfig(**item) for item in data.get("connections", [])]
|
||||
return BrokerConfig(store_root=store_root, connections=connections)
|
||||
|
||||
|
||||
def build_service_from_config(config: BrokerConfig) -> WfMcpService:
|
||||
service = WfMcpService(store=FileStore(config.store_root))
|
||||
for connection in config.connections:
|
||||
service.register_connection(connection)
|
||||
if connection.server not in service.adapters:
|
||||
service.register_adapter(connection.server, McpSdkAdapter())
|
||||
return service
|
||||
|
||||
|
||||
def create_broker_server(service: WfMcpService) -> FastMCP:
|
||||
server = FastMCP(
|
||||
"wf-mcp-broker",
|
||||
instructions=(
|
||||
"A broker MCP server over one or more upstream MCP connections. "
|
||||
"Use tools for refresh and invocation, resources for snapshots, "
|
||||
"and prompts for planning against available capabilities."
|
||||
),
|
||||
)
|
||||
|
||||
@server.tool()
|
||||
async def list_connections() -> list[dict[str, Any]]:
|
||||
return [
|
||||
asdict(connection)
|
||||
for connection in sorted(
|
||||
service.connections.list_all(),
|
||||
key=lambda connection: connection.id,
|
||||
)
|
||||
]
|
||||
|
||||
@server.tool()
|
||||
async def refresh_connection_catalog(connection_id: str) -> dict[str, Any]:
|
||||
await service.refresh_connection_catalog(connection_id)
|
||||
snapshot = service.get_connection_snapshot(connection_id)
|
||||
if snapshot is None:
|
||||
return {"connection_id": connection_id, "refreshed": False}
|
||||
return {
|
||||
"connection_id": connection_id,
|
||||
"refreshed": True,
|
||||
"node_count": len(snapshot.nodes),
|
||||
"resource_count": len(snapshot.resources),
|
||||
"prompt_count": len(snapshot.prompts),
|
||||
}
|
||||
|
||||
@server.tool()
|
||||
async def get_catalog() -> dict[str, Any]:
|
||||
return service.get_catalog().as_payload()
|
||||
|
||||
@server.tool()
|
||||
async def read_broker_resource(qualified_name: str) -> dict[str, Any]:
|
||||
return await service.read_resource(qualified_name)
|
||||
|
||||
@server.tool()
|
||||
async def render_broker_prompt(
|
||||
qualified_name: str,
|
||||
arguments: dict[str, str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
return await service.render_prompt(qualified_name, arguments=arguments)
|
||||
|
||||
@server.tool()
|
||||
async def invoke_broker_method(
|
||||
connection_id: str,
|
||||
method: str,
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
return await service.invoke_method(connection_id, method, params=params)
|
||||
|
||||
@server.tool()
|
||||
async def get_broker_events() -> list[dict[str, Any]]:
|
||||
return [asdict(event) for event in service.list_events()]
|
||||
|
||||
@server.resource("wf-mcp://catalog", name="catalog.all")
|
||||
def catalog_resource() -> str:
|
||||
return json.dumps(service.get_catalog().as_payload(), indent=2)
|
||||
|
||||
@server.resource(
|
||||
"wf-mcp://connection/{connection_id}/catalog",
|
||||
name="catalog.connection",
|
||||
)
|
||||
def connection_catalog_resource(connection_id: str) -> str:
|
||||
snapshot = service.get_connection_snapshot(connection_id)
|
||||
if snapshot is None:
|
||||
raise KeyError(connection_id)
|
||||
return json.dumps(
|
||||
{
|
||||
"connection_id": snapshot.connection_id,
|
||||
"fetched_at_epoch_ms": snapshot.fetched_at_epoch_ms,
|
||||
"max_age_seconds": snapshot.max_age_seconds,
|
||||
"nodes": [asdict(node) for node in snapshot.nodes],
|
||||
"resources": [asdict(resource) for resource in snapshot.resources],
|
||||
"prompts": [asdict(prompt) for prompt in snapshot.prompts],
|
||||
"metadata": snapshot.metadata,
|
||||
},
|
||||
indent=2,
|
||||
)
|
||||
|
||||
@server.resource("wf-mcp://events", name="events.all")
|
||||
def events_resource() -> str:
|
||||
return json.dumps([asdict(event) for event in service.list_events()], indent=2)
|
||||
|
||||
@server.prompt(
|
||||
name="plan_with_catalog",
|
||||
description="Provide the broker catalog as planning context.",
|
||||
)
|
||||
def plan_with_catalog() -> list[dict[str, str]]:
|
||||
payload = json.dumps(service.get_catalog().as_payload(), indent=2)
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": (
|
||||
"Plan a workflow using this broker catalog. "
|
||||
"Prefer existing namespaced capabilities.\n\n"
|
||||
f"{payload}"
|
||||
),
|
||||
}
|
||||
]
|
||||
|
||||
return server
|
||||
|
||||
|
||||
def main() -> None:
|
||||
config_path = os.environ.get("WF_MCP_CONFIG", "wf_mcp.config.json")
|
||||
transport_env = os.environ.get("WF_MCP_TRANSPORT", "stdio")
|
||||
run_broker_server(config_path, transport_env)
|
||||
|
||||
|
||||
def normalize_transport(
|
||||
transport: str,
|
||||
) -> Literal["stdio", "sse", "streamable-http"]:
|
||||
match transport:
|
||||
case "streamable_http" | "streamable-http":
|
||||
return "streamable-http"
|
||||
case "stdio":
|
||||
return "stdio"
|
||||
case "sse":
|
||||
return "sse"
|
||||
case _:
|
||||
raise ValueError(f"we dont support {transport} yet sry")
|
||||
|
||||
|
||||
def run_broker_server(config_path: str | Path, transport: str = "stdio") -> None:
|
||||
config = load_broker_config(config_path)
|
||||
service = build_service_from_config(config)
|
||||
server = create_broker_server(service)
|
||||
server.run(transport=normalize_transport(transport))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,63 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class DiscoveredTool:
|
||||
name: str
|
||||
description: str | None
|
||||
input_schema: dict[str, Any]
|
||||
output_schema: dict[str, Any]
|
||||
outcomes: tuple[str, ...] = ("ok",)
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class DiscoveredResource:
|
||||
uri: str
|
||||
name: str
|
||||
description: str | None
|
||||
mime_type: str | None = None
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class DiscoveredPrompt:
|
||||
name: str
|
||||
description: str | None
|
||||
arguments: list[dict[str, Any]] = field(default_factory=list)
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class CatalogNodeEntry:
|
||||
qualified_name: str
|
||||
connection_id: str
|
||||
local_name: str
|
||||
description: str | None
|
||||
outcomes: tuple[str, ...]
|
||||
input_schema: dict[str, Any]
|
||||
output_schema: dict[str, Any]
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class CatalogResourceEntry:
|
||||
qualified_name: str
|
||||
connection_id: str
|
||||
local_name: str
|
||||
uri: str
|
||||
description: str | None
|
||||
mime_type: str | None = None
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class CatalogPromptEntry:
|
||||
qualified_name: str
|
||||
connection_id: str
|
||||
local_name: str
|
||||
description: str | None
|
||||
arguments: list[dict[str, Any]] = field(default_factory=list)
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
@@ -0,0 +1,161 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from wf_authoring import NodeCatalog, NodeSpec
|
||||
|
||||
from .capabilities import (
|
||||
CatalogNodeEntry,
|
||||
CatalogPromptEntry,
|
||||
CatalogResourceEntry,
|
||||
DiscoveredPrompt,
|
||||
DiscoveredResource,
|
||||
)
|
||||
from .connections import qualify_node_name
|
||||
from .models import CatalogSnapshot
|
||||
|
||||
|
||||
def snapshot_from_specs(
|
||||
connection_id: str,
|
||||
*,
|
||||
specs: dict[str, NodeSpec[Any, Any]],
|
||||
resources: list[DiscoveredResource] | None = None,
|
||||
prompts: list[DiscoveredPrompt] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
fetched_at_epoch_ms: int,
|
||||
max_age_seconds: int,
|
||||
) -> CatalogSnapshot:
|
||||
catalog = NodeCatalog.from_specs(*specs.values())
|
||||
nodes = [
|
||||
CatalogNodeEntry(
|
||||
qualified_name=entry.name
|
||||
if entry.name.startswith(f"{connection_id}.")
|
||||
else qualify_node_name(connection_id, entry.name),
|
||||
connection_id=connection_id,
|
||||
local_name=entry.name.removeprefix(f"{connection_id}."),
|
||||
description=entry.description,
|
||||
outcomes=entry.outcomes,
|
||||
input_schema=entry.input_schema,
|
||||
output_schema=entry.output_schema,
|
||||
)
|
||||
for entry in catalog.entries()
|
||||
]
|
||||
resource_entries = [
|
||||
CatalogResourceEntry(
|
||||
qualified_name=qualify_node_name(connection_id, resource.name),
|
||||
connection_id=connection_id,
|
||||
local_name=resource.name,
|
||||
uri=resource.uri,
|
||||
description=resource.description,
|
||||
mime_type=resource.mime_type,
|
||||
metadata=resource.metadata,
|
||||
)
|
||||
for resource in resources or []
|
||||
]
|
||||
prompt_entries = [
|
||||
CatalogPromptEntry(
|
||||
qualified_name=qualify_node_name(connection_id, prompt.name),
|
||||
connection_id=connection_id,
|
||||
local_name=prompt.name,
|
||||
description=prompt.description,
|
||||
arguments=prompt.arguments,
|
||||
metadata=prompt.metadata,
|
||||
)
|
||||
for prompt in prompts or []
|
||||
]
|
||||
return CatalogSnapshot(
|
||||
connection_id=connection_id,
|
||||
fetched_at_epoch_ms=fetched_at_epoch_ms,
|
||||
max_age_seconds=max_age_seconds,
|
||||
nodes=nodes,
|
||||
resources=resource_entries,
|
||||
prompts=prompt_entries,
|
||||
metadata=metadata or {},
|
||||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class CombinedCatalog:
|
||||
snapshots: dict[str, CatalogSnapshot] = field(default_factory=dict)
|
||||
|
||||
def entries(self) -> list[CatalogNodeEntry]:
|
||||
result: list[CatalogNodeEntry] = []
|
||||
for snapshot in self.snapshots.values():
|
||||
result.extend(snapshot.nodes)
|
||||
return sorted(result, key=lambda entry: entry.qualified_name)
|
||||
|
||||
def resource_entries(self) -> list[CatalogResourceEntry]:
|
||||
result: list[CatalogResourceEntry] = []
|
||||
for snapshot in self.snapshots.values():
|
||||
result.extend(snapshot.resources)
|
||||
return sorted(result, key=lambda entry: entry.qualified_name)
|
||||
|
||||
def prompt_entries(self) -> list[CatalogPromptEntry]:
|
||||
result: list[CatalogPromptEntry] = []
|
||||
for snapshot in self.snapshots.values():
|
||||
result.extend(snapshot.prompts)
|
||||
return sorted(result, key=lambda entry: entry.qualified_name)
|
||||
|
||||
def find_resource(self, qualified_name: str) -> CatalogResourceEntry | None:
|
||||
for entry in self.resource_entries():
|
||||
if entry.qualified_name == qualified_name:
|
||||
return entry
|
||||
return None
|
||||
|
||||
def find_prompt(self, qualified_name: str) -> CatalogPromptEntry | None:
|
||||
for entry in self.prompt_entries():
|
||||
if entry.qualified_name == qualified_name:
|
||||
return entry
|
||||
return None
|
||||
|
||||
def as_payload(self) -> dict[str, Any]:
|
||||
return {
|
||||
"nodes": [
|
||||
{
|
||||
"qualified_name": entry.qualified_name,
|
||||
"connection_id": entry.connection_id,
|
||||
"local_name": entry.local_name,
|
||||
"description": entry.description,
|
||||
"outcomes": list(entry.outcomes),
|
||||
"input_schema": entry.input_schema,
|
||||
"output_schema": entry.output_schema,
|
||||
}
|
||||
for entry in self.entries()
|
||||
],
|
||||
"resources": [
|
||||
{
|
||||
"qualified_name": entry.qualified_name,
|
||||
"connection_id": entry.connection_id,
|
||||
"local_name": entry.local_name,
|
||||
"uri": entry.uri,
|
||||
"description": entry.description,
|
||||
"mime_type": entry.mime_type,
|
||||
"metadata": entry.metadata,
|
||||
}
|
||||
for entry in self.resource_entries()
|
||||
],
|
||||
"prompts": [
|
||||
{
|
||||
"qualified_name": entry.qualified_name,
|
||||
"connection_id": entry.connection_id,
|
||||
"local_name": entry.local_name,
|
||||
"description": entry.description,
|
||||
"arguments": entry.arguments,
|
||||
"metadata": entry.metadata,
|
||||
}
|
||||
for entry in self.prompt_entries()
|
||||
],
|
||||
"connections": [
|
||||
{
|
||||
"connection_id": snapshot.connection_id,
|
||||
"fetched_at_epoch_ms": snapshot.fetched_at_epoch_ms,
|
||||
"max_age_seconds": snapshot.max_age_seconds,
|
||||
"metadata": snapshot.metadata,
|
||||
}
|
||||
for snapshot in sorted(
|
||||
self.snapshots.values(),
|
||||
key=lambda snapshot: snapshot.connection_id,
|
||||
)
|
||||
],
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from .broker_server import (
|
||||
build_service_from_config,
|
||||
load_broker_config,
|
||||
run_broker_server,
|
||||
)
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(prog="wf-mcp")
|
||||
parser.add_argument(
|
||||
"--config",
|
||||
default="wf_mcp.config.json",
|
||||
help="Path to broker config JSON.",
|
||||
)
|
||||
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
serve = subparsers.add_parser("serve", help="Run the broker MCP server.")
|
||||
serve.add_argument(
|
||||
"--transport",
|
||||
default="stdio",
|
||||
choices=["stdio", "sse", "streamable-http", "streamable_http"],
|
||||
help="Transport to run the broker server with.",
|
||||
)
|
||||
|
||||
subparsers.add_parser("connections", help="List configured connections.")
|
||||
subparsers.add_parser("catalog", help="Print the broker catalog as JSON.")
|
||||
|
||||
refresh = subparsers.add_parser(
|
||||
"refresh",
|
||||
help="Refresh one connection catalog or all configured connections.",
|
||||
)
|
||||
refresh.add_argument("connection_id", nargs="?", help="Connection id to refresh.")
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def _service_from_config(config_path: str | Path):
|
||||
config = load_broker_config(config_path)
|
||||
return build_service_from_config(config)
|
||||
|
||||
|
||||
def _json_dump(data: Any) -> None:
|
||||
print(json.dumps(data, indent=2))
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
if args.command == "serve":
|
||||
run_broker_server(args.config, args.transport)
|
||||
return 0
|
||||
|
||||
service = _service_from_config(args.config)
|
||||
|
||||
if args.command == "connections":
|
||||
_json_dump(
|
||||
[
|
||||
{
|
||||
"id": connection.id,
|
||||
"server": connection.server,
|
||||
"account": connection.account,
|
||||
"enabled": connection.enabled,
|
||||
"metadata": connection.metadata,
|
||||
}
|
||||
for connection in service.connections.list_all()
|
||||
]
|
||||
)
|
||||
return 0
|
||||
|
||||
if args.command == "catalog":
|
||||
_json_dump(service.get_catalog().as_payload())
|
||||
return 0
|
||||
|
||||
if args.command == "refresh":
|
||||
if args.connection_id:
|
||||
asyncio.run(service.refresh_connection_catalog(args.connection_id))
|
||||
else:
|
||||
for connection in service.connections.list_enabled():
|
||||
asyncio.run(service.refresh_connection_catalog(connection.id))
|
||||
_json_dump(service.get_catalog().as_payload())
|
||||
return 0
|
||||
|
||||
parser.error(f"unknown command {args.command!r}")
|
||||
return 2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,41 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from .models import ConnectionConfig
|
||||
|
||||
|
||||
def parse_connection_id(connection_id: str) -> tuple[str, str]:
|
||||
if "." not in connection_id:
|
||||
raise ValueError("connection id must look like '<server>.<account>'")
|
||||
server, account = connection_id.split(".", 1)
|
||||
if not server or not account:
|
||||
raise ValueError("connection id must look like '<server>.<account>'")
|
||||
return server, account
|
||||
|
||||
|
||||
def qualify_node_name(connection_id: str, local_name: str) -> str:
|
||||
parse_connection_id(connection_id)
|
||||
if not local_name:
|
||||
raise ValueError("local node name must not be empty")
|
||||
return f"{connection_id}.{local_name}"
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ConnectionRegistry:
|
||||
connections: dict[str, ConnectionConfig] = field(default_factory=dict)
|
||||
|
||||
def register(self, connection: ConnectionConfig) -> None:
|
||||
parse_connection_id(connection.id)
|
||||
self.connections[connection.id] = connection
|
||||
|
||||
def get(self, connection_id: str) -> ConnectionConfig:
|
||||
return self.connections[connection_id]
|
||||
|
||||
def list_all(self) -> list[ConnectionConfig]:
|
||||
return list(self.connections.values())
|
||||
|
||||
def list_enabled(self) -> list[ConnectionConfig]:
|
||||
return [
|
||||
connection for connection in self.connections.values() if connection.enabled
|
||||
]
|
||||
@@ -0,0 +1,63 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from wf_authoring import NodeSpec
|
||||
|
||||
from .adapters import (
|
||||
BackendAdapter,
|
||||
DiscoveredPrompt,
|
||||
DiscoveredResource,
|
||||
DiscoveredTool,
|
||||
)
|
||||
from .events import McpEvent
|
||||
from .models import AuthRecord, ConnectionConfig
|
||||
from .wrappers import wrap_discovered_tool
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class DiscoveredConnectionCapabilities:
|
||||
tools: list[DiscoveredTool] = field(default_factory=list)
|
||||
resources: list[DiscoveredResource] = field(default_factory=list)
|
||||
prompts: list[DiscoveredPrompt] = field(default_factory=list)
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
async def discover_connection_capabilities(
|
||||
*,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
adapter: BackendAdapter,
|
||||
) -> DiscoveredConnectionCapabilities:
|
||||
tools = await adapter.list_tools(connection, auth)
|
||||
resources = await adapter.list_resources(connection, auth)
|
||||
prompts = await adapter.list_prompts(connection, auth)
|
||||
metadata = await adapter.get_connection_metadata(connection, auth)
|
||||
return DiscoveredConnectionCapabilities(
|
||||
tools=tools,
|
||||
resources=resources,
|
||||
prompts=prompts,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
|
||||
def specs_from_discovered_tools(
|
||||
*,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
adapter: BackendAdapter,
|
||||
tools: list[DiscoveredTool],
|
||||
emit_event: Callable[[McpEvent], None] | None = None,
|
||||
) -> list[NodeSpec[Any, Any]]:
|
||||
return [
|
||||
wrap_discovered_tool(
|
||||
connection=connection,
|
||||
auth=auth,
|
||||
adapter=adapter,
|
||||
tool=tool,
|
||||
emit_event=emit_event,
|
||||
)
|
||||
for tool in tools
|
||||
]
|
||||
@@ -0,0 +1,33 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class McpEvent:
|
||||
kind: str
|
||||
timestamp_epoch_ms: int
|
||||
connection_id: str | None = None
|
||||
capability_id: str | None = None
|
||||
workflow_name: str | None = None
|
||||
payload: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
def make_event(
|
||||
kind: str,
|
||||
*,
|
||||
connection_id: str | None = None,
|
||||
capability_id: str | None = None,
|
||||
workflow_name: str | None = None,
|
||||
payload: dict[str, Any] | None = None,
|
||||
) -> McpEvent:
|
||||
return McpEvent(
|
||||
kind=kind,
|
||||
timestamp_epoch_ms=int(time.time() * 1000),
|
||||
connection_id=connection_id,
|
||||
capability_id=capability_id,
|
||||
workflow_name=workflow_name,
|
||||
payload=payload or {},
|
||||
)
|
||||
@@ -0,0 +1,230 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from mcp import ClientResult
|
||||
from mcp.client.session import ClientSession
|
||||
from mcp.client.stdio import StdioServerParameters, stdio_client
|
||||
from mcp.client.streamable_http import streamable_http_client
|
||||
from mcp.types import CallToolResult as McpCallToolResult
|
||||
from mcp.types import ClientNotification, ClientRequest
|
||||
from mcp.types import ListPromptsResult, ListResourcesResult
|
||||
from mcp.types import ListToolsResult, Tool as McpTool
|
||||
from mcp.types import Prompt as McpPrompt
|
||||
from mcp.types import Resource as McpResource
|
||||
from pydantic import AnyUrl
|
||||
|
||||
from .adapters import (
|
||||
BackendAdapter,
|
||||
DiscoveredPrompt,
|
||||
DiscoveredResource,
|
||||
DiscoveredTool,
|
||||
ToolCallResult,
|
||||
)
|
||||
from .models import AuthRecord, ConnectionConfig
|
||||
|
||||
|
||||
def _auth_headers(auth: AuthRecord | None) -> dict[str, str]:
|
||||
if auth is None:
|
||||
return {}
|
||||
headers = dict(auth.payload.get("headers", {}))
|
||||
token = auth.payload.get("token")
|
||||
if isinstance(token, str) and "Authorization" not in headers:
|
||||
headers["Authorization"] = f"Bearer {token}"
|
||||
return headers
|
||||
|
||||
|
||||
def _tool_to_discovered(tool: McpTool) -> DiscoveredTool:
|
||||
output_schema = tool.outputSchema or {
|
||||
"type": "object",
|
||||
"properties": {"content": {"type": "array"}},
|
||||
}
|
||||
return DiscoveredTool(
|
||||
name=tool.name,
|
||||
description=tool.description,
|
||||
input_schema=tool.inputSchema,
|
||||
output_schema=output_schema,
|
||||
outcomes=("ok", "error"),
|
||||
metadata=tool.model_dump(by_alias=True, mode="json"),
|
||||
)
|
||||
|
||||
|
||||
def _resource_to_discovered(resource: McpResource) -> DiscoveredResource:
|
||||
local_name = resource.name or str(resource.uri)
|
||||
return DiscoveredResource(
|
||||
uri=str(resource.uri),
|
||||
name=local_name,
|
||||
description=resource.description,
|
||||
mime_type=resource.mimeType,
|
||||
metadata=resource.model_dump(by_alias=True, mode="json"),
|
||||
)
|
||||
|
||||
|
||||
def _prompt_to_discovered(prompt: McpPrompt) -> DiscoveredPrompt:
|
||||
arguments = [
|
||||
argument.model_dump(by_alias=True, mode="json")
|
||||
for argument in prompt.arguments or []
|
||||
]
|
||||
return DiscoveredPrompt(
|
||||
name=prompt.name,
|
||||
description=prompt.description,
|
||||
arguments=arguments,
|
||||
metadata=prompt.model_dump(by_alias=True, mode="json"),
|
||||
)
|
||||
|
||||
|
||||
def _tool_result_to_call_result(result: McpCallToolResult) -> ToolCallResult:
|
||||
if result.structuredContent is not None:
|
||||
output = result.structuredContent
|
||||
else:
|
||||
output = {
|
||||
"content": [item.model_dump(by_alias=True) for item in result.content]
|
||||
}
|
||||
return ToolCallResult(
|
||||
outcome="error" if result.isError else "ok",
|
||||
output=output,
|
||||
meta=result.meta or {},
|
||||
)
|
||||
|
||||
|
||||
class McpSdkAdapter(BackendAdapter):
|
||||
@asynccontextmanager
|
||||
async def _session(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
):
|
||||
transport = connection.metadata.get("transport", "stdio")
|
||||
if transport == "stdio":
|
||||
command = connection.metadata["command"]
|
||||
args = list(connection.metadata.get("args", []))
|
||||
env = connection.metadata.get("env")
|
||||
cwd = connection.metadata.get("cwd")
|
||||
if auth is not None:
|
||||
auth_env = auth.payload.get("env")
|
||||
if isinstance(auth_env, dict):
|
||||
env = {**(env or {}), **auth_env}
|
||||
params = StdioServerParameters(
|
||||
command=command,
|
||||
args=args,
|
||||
env=env,
|
||||
cwd=cwd,
|
||||
)
|
||||
async with stdio_client(params) as (read_stream, write_stream):
|
||||
async with ClientSession(read_stream, write_stream) as session:
|
||||
await session.initialize()
|
||||
yield session
|
||||
return
|
||||
|
||||
if transport == "streamable_http":
|
||||
url = connection.metadata["url"]
|
||||
headers = _auth_headers(auth)
|
||||
http_client = httpx.AsyncClient(headers=headers or None)
|
||||
async with http_client:
|
||||
async with streamable_http_client(
|
||||
url,
|
||||
http_client=http_client,
|
||||
) as (read_stream, write_stream, _get_session_id):
|
||||
async with ClientSession(read_stream, write_stream) as session:
|
||||
await session.initialize()
|
||||
yield session
|
||||
return
|
||||
|
||||
raise ValueError(f"unsupported MCP transport {transport!r}")
|
||||
|
||||
async def list_tools(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
) -> list[DiscoveredTool]:
|
||||
async with self._session(connection, auth) as session:
|
||||
result: ListToolsResult = await session.list_tools()
|
||||
return [_tool_to_discovered(tool) for tool in result.tools]
|
||||
|
||||
async def list_resources(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
) -> list[DiscoveredResource]:
|
||||
async with self._session(connection, auth) as session:
|
||||
result: ListResourcesResult = await session.list_resources()
|
||||
return [_resource_to_discovered(resource) for resource in result.resources]
|
||||
|
||||
async def list_prompts(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
) -> list[DiscoveredPrompt]:
|
||||
async with self._session(connection, auth) as session:
|
||||
result: ListPromptsResult = await session.list_prompts()
|
||||
return [_prompt_to_discovered(prompt) for prompt in result.prompts]
|
||||
|
||||
async def get_connection_metadata(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"server": connection.server,
|
||||
"transport": connection.metadata.get("transport", "stdio"),
|
||||
}
|
||||
|
||||
async def read_resource(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
uri: str,
|
||||
) -> dict[str, Any]:
|
||||
async with self._session(connection, auth) as session:
|
||||
result = await session.read_resource(AnyUrl(uri))
|
||||
return result.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
|
||||
async def get_prompt(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
prompt_name: str,
|
||||
arguments: dict[str, str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
async with self._session(connection, auth) as session:
|
||||
result = await session.get_prompt(prompt_name, arguments)
|
||||
return result.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
|
||||
async def invoke_method(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
method: str,
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
async with self._session(connection, auth) as session:
|
||||
result = await session.send_request(
|
||||
ClientRequest.model_validate({"method": method, "params": params}),
|
||||
ClientResult,
|
||||
)
|
||||
return result.model_dump(by_alias=True, mode="json", exclude_none=True)
|
||||
|
||||
async def send_notification(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
method: str,
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
async with self._session(connection, auth) as session:
|
||||
await session.send_notification(
|
||||
ClientNotification.model_validate({"method": method, "params": params})
|
||||
)
|
||||
|
||||
async def call_tool(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
tool_name: str,
|
||||
payload: dict[str, Any],
|
||||
) -> ToolCallResult:
|
||||
async with self._session(connection, auth) as session:
|
||||
result = await session.call_tool(tool_name, payload)
|
||||
return _tool_result_to_call_result(result)
|
||||
@@ -0,0 +1,67 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from .capabilities import CatalogNodeEntry, CatalogPromptEntry, CatalogResourceEntry
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ConnectionConfig:
|
||||
id: str
|
||||
server: str
|
||||
account: str
|
||||
enabled: bool = True
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AuthRecord:
|
||||
connection_id: str
|
||||
scheme: str
|
||||
payload: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class CatalogSnapshot:
|
||||
connection_id: str
|
||||
fetched_at_epoch_ms: int
|
||||
max_age_seconds: int
|
||||
nodes: list[CatalogNodeEntry] = field(default_factory=list)
|
||||
resources: list[CatalogResourceEntry] = field(default_factory=list)
|
||||
prompts: list[CatalogPromptEntry] = field(default_factory=list)
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def is_stale(self, now_epoch_ms: int) -> bool:
|
||||
age_ms = now_epoch_ms - self.fetched_at_epoch_ms
|
||||
return age_ms > self.max_age_seconds * 1000
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class RawWorkflowPlan:
|
||||
name: str
|
||||
input_schema: dict[str, Any]
|
||||
state_schema: dict[str, Any]
|
||||
output_schema: dict[str, Any]
|
||||
start: str
|
||||
nodes: list[dict[str, Any]]
|
||||
edges: list[dict[str, Any]]
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class BrokerConfig:
|
||||
store_root: Path
|
||||
connections: list[ConnectionConfig] = field(default_factory=list)
|
||||
|
||||
|
||||
def dump_catalog_snapshot(snapshot: CatalogSnapshot) -> dict[str, Any]:
|
||||
return {
|
||||
"connection_id": snapshot.connection_id,
|
||||
"fetched_at_epoch_ms": snapshot.fetched_at_epoch_ms,
|
||||
"max_age_seconds": snapshot.max_age_seconds,
|
||||
"nodes": [asdict(node) for node in snapshot.nodes],
|
||||
"resources": [asdict(resource) for resource in snapshot.resources],
|
||||
"prompts": [asdict(prompt) for prompt in snapshot.prompts],
|
||||
"metadata": snapshot.metadata,
|
||||
}
|
||||
@@ -0,0 +1,395 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from wf_authoring import NodeSpec, build_async_registry
|
||||
from wf_core import NodeUse, Workflow, execute_workflow_async
|
||||
|
||||
from .adapters import BackendAdapter
|
||||
from .catalog import CombinedCatalog, snapshot_from_specs
|
||||
from .connections import ConnectionRegistry, parse_connection_id, qualify_node_name
|
||||
from .discovery import discover_connection_capabilities, specs_from_discovered_tools
|
||||
from .events import McpEvent, make_event
|
||||
from .models import (
|
||||
AuthRecord,
|
||||
CatalogPromptEntry,
|
||||
CatalogResourceEntry,
|
||||
CatalogSnapshot,
|
||||
ConnectionConfig,
|
||||
RawWorkflowPlan,
|
||||
)
|
||||
from .store import Store
|
||||
|
||||
|
||||
def _qualify_spec(connection_id: str, spec: NodeSpec[Any, Any]) -> NodeSpec[Any, Any]:
|
||||
return NodeSpec(
|
||||
name=qualify_node_name(connection_id, spec.name),
|
||||
input_model=spec.input_model,
|
||||
output_model=spec.output_model,
|
||||
outcomes=spec.outcomes,
|
||||
fn=spec.fn,
|
||||
description=spec.description,
|
||||
is_async=spec.is_async,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
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)
|
||||
specs_by_connection: dict[str, dict[str, NodeSpec[Any, Any]]] = field(
|
||||
default_factory=dict
|
||||
)
|
||||
events: list[McpEvent] = field(default_factory=list)
|
||||
|
||||
def register_connection(self, connection: ConnectionConfig) -> None:
|
||||
parse_connection_id(connection.id)
|
||||
self.connections.register(connection)
|
||||
self._record_event(
|
||||
make_event(
|
||||
"connection_registered",
|
||||
connection_id=connection.id,
|
||||
payload={"server": connection.server, "account": connection.account},
|
||||
)
|
||||
)
|
||||
|
||||
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._record_event(
|
||||
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 register_specs(
|
||||
self,
|
||||
connection_id: str,
|
||||
*specs: NodeSpec[Any, Any],
|
||||
max_age_seconds: int | None = None,
|
||||
) -> None:
|
||||
self.connections.get(connection_id)
|
||||
qualified_specs = {
|
||||
qualify_node_name(connection_id, spec.name): _qualify_spec(
|
||||
connection_id, spec
|
||||
)
|
||||
for spec in specs
|
||||
}
|
||||
self.specs_by_connection[connection_id] = qualified_specs
|
||||
snapshot = snapshot_from_specs(
|
||||
connection_id,
|
||||
specs=qualified_specs,
|
||||
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)
|
||||
self._record_event(
|
||||
make_event(
|
||||
"specs_registered",
|
||||
connection_id=connection_id,
|
||||
payload={"node_count": len(qualified_specs)},
|
||||
)
|
||||
)
|
||||
|
||||
def get_catalog(self) -> CombinedCatalog:
|
||||
snapshots: dict[str, CatalogSnapshot] = {}
|
||||
for connection in self.connections.list_enabled():
|
||||
snapshot = self.store.load_catalog(connection.id)
|
||||
if snapshot is not None:
|
||||
snapshots[connection.id] = snapshot
|
||||
return CombinedCatalog(snapshots=snapshots)
|
||||
|
||||
def get_connection_snapshot(self, connection_id: str) -> CatalogSnapshot | None:
|
||||
self.connections.get(connection_id)
|
||||
return self.store.load_catalog(connection_id)
|
||||
|
||||
def list_resources(
|
||||
self,
|
||||
*,
|
||||
connection_id: str | None = None,
|
||||
) -> list[CatalogResourceEntry]:
|
||||
if connection_id is None:
|
||||
return self.get_catalog().resource_entries()
|
||||
snapshot = self.get_connection_snapshot(connection_id)
|
||||
if snapshot is None:
|
||||
return []
|
||||
return sorted(snapshot.resources, key=lambda entry: entry.qualified_name)
|
||||
|
||||
def list_prompts(
|
||||
self,
|
||||
*,
|
||||
connection_id: str | None = None,
|
||||
) -> list[CatalogPromptEntry]:
|
||||
if connection_id is None:
|
||||
return self.get_catalog().prompt_entries()
|
||||
snapshot = self.get_connection_snapshot(connection_id)
|
||||
if snapshot is None:
|
||||
return []
|
||||
return sorted(snapshot.prompts, key=lambda entry: entry.qualified_name)
|
||||
|
||||
def get_resource(self, qualified_name: str) -> CatalogResourceEntry:
|
||||
entry = self.get_catalog().find_resource(qualified_name)
|
||||
if entry is None:
|
||||
raise KeyError(f"unknown resource {qualified_name!r}")
|
||||
return entry
|
||||
|
||||
def get_prompt(self, qualified_name: str) -> CatalogPromptEntry:
|
||||
entry = self.get_catalog().find_prompt(qualified_name)
|
||||
if entry is None:
|
||||
raise KeyError(f"unknown prompt {qualified_name!r}")
|
||||
return entry
|
||||
|
||||
async def read_resource(self, qualified_name: str) -> dict[str, Any]:
|
||||
resource = self.get_resource(qualified_name)
|
||||
connection = self.connections.get(resource.connection_id)
|
||||
adapter = self.adapters.get(connection.server)
|
||||
if adapter is None:
|
||||
raise KeyError(f"no adapter registered for server {connection.server!r}")
|
||||
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},
|
||||
)
|
||||
)
|
||||
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,
|
||||
connection_id: str,
|
||||
method: str,
|
||||
*,
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
connection = self.connections.get(connection_id)
|
||||
adapter = self.adapters.get(connection.server)
|
||||
if adapter is None:
|
||||
raise KeyError(f"no adapter registered for server {connection.server!r}")
|
||||
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 {}},
|
||||
)
|
||||
)
|
||||
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,
|
||||
connection_id: str,
|
||||
method: str,
|
||||
*,
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
connection = self.connections.get(connection_id)
|
||||
adapter = self.adapters.get(connection.server)
|
||||
if adapter is None:
|
||||
raise KeyError(f"no adapter registered for server {connection.server!r}")
|
||||
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={},
|
||||
)
|
||||
)
|
||||
|
||||
async def render_prompt(
|
||||
self,
|
||||
qualified_name: str,
|
||||
*,
|
||||
arguments: dict[str, str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
prompt = self.get_prompt(qualified_name)
|
||||
connection = self.connections.get(prompt.connection_id)
|
||||
adapter = self.adapters.get(connection.server)
|
||||
if adapter is None:
|
||||
raise KeyError(f"no adapter registered for server {connection.server!r}")
|
||||
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(
|
||||
connection,
|
||||
auth,
|
||||
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,
|
||||
connection_id: str,
|
||||
*,
|
||||
max_age_seconds: int | None = None,
|
||||
) -> None:
|
||||
connection = self.connections.get(connection_id)
|
||||
adapter = self.adapters.get(connection.server)
|
||||
if adapter is None:
|
||||
raise KeyError(f"no adapter registered for server {connection.server!r}")
|
||||
|
||||
auth = self.load_auth(connection_id)
|
||||
self._record_event(
|
||||
make_event(
|
||||
"catalog_refresh_started",
|
||||
connection_id=connection_id,
|
||||
payload={"server": connection.server},
|
||||
)
|
||||
)
|
||||
capabilities = await discover_connection_capabilities(
|
||||
connection=connection,
|
||||
auth=auth,
|
||||
adapter=adapter,
|
||||
)
|
||||
specs = specs_from_discovered_tools(
|
||||
connection=connection,
|
||||
auth=auth,
|
||||
adapter=adapter,
|
||||
tools=capabilities.tools,
|
||||
emit_event=self._record_event,
|
||||
)
|
||||
self.register_specs(
|
||||
connection_id,
|
||||
*specs,
|
||||
max_age_seconds=max_age_seconds,
|
||||
)
|
||||
snapshot = snapshot_from_specs(
|
||||
connection_id,
|
||||
specs=self.specs_by_connection.get(connection_id, {}),
|
||||
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)
|
||||
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),
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
def compile_plan(self, plan: RawWorkflowPlan) -> Workflow:
|
||||
node_defs: dict[str, Any] = {}
|
||||
for step in plan.nodes:
|
||||
if step.get("type") != "node":
|
||||
continue
|
||||
qualified_name = step["node"]
|
||||
spec = self._get_qualified_spec(qualified_name)
|
||||
node_defs[qualified_name] = spec.to_node_def()
|
||||
|
||||
payload = {
|
||||
"name": plan.name,
|
||||
"input_schema": plan.input_schema,
|
||||
"state_schema": plan.state_schema,
|
||||
"output_schema": plan.output_schema,
|
||||
"start": plan.start,
|
||||
"node_defs": [node.model_dump() for node in node_defs.values()],
|
||||
"nodes": plan.nodes,
|
||||
"edges": plan.edges,
|
||||
}
|
||||
return Workflow.model_validate(payload)
|
||||
|
||||
async def run_workflow_from_plan(
|
||||
self,
|
||||
plan: RawWorkflowPlan,
|
||||
workflow_input: dict[str, Any],
|
||||
):
|
||||
self._record_event(
|
||||
make_event(
|
||||
"workflow_run_started",
|
||||
workflow_name=plan.name,
|
||||
payload={"input_keys": sorted(workflow_input.keys())},
|
||||
)
|
||||
)
|
||||
workflow = self.compile_plan(plan)
|
||||
specs = [
|
||||
self._get_qualified_spec(node.node)
|
||||
for node in workflow.nodes
|
||||
if isinstance(node, NodeUse)
|
||||
]
|
||||
registry = build_async_registry(*specs)
|
||||
run = await execute_workflow_async(workflow, workflow_input, registry)
|
||||
self._record_event(
|
||||
make_event(
|
||||
"workflow_run_completed",
|
||||
workflow_name=plan.name,
|
||||
payload={"status": run.status.value},
|
||||
)
|
||||
)
|
||||
return run
|
||||
|
||||
def list_events(self) -> list[McpEvent]:
|
||||
return list(self.events)
|
||||
|
||||
def _get_qualified_spec(self, qualified_name: str) -> NodeSpec[Any, Any]:
|
||||
connection_id, _ = qualified_name.rsplit(".", 1)
|
||||
specs = self.specs_by_connection.get(connection_id)
|
||||
if specs is None or qualified_name not in specs:
|
||||
raise KeyError(f"unknown qualified node {qualified_name!r}")
|
||||
return specs[qualified_name]
|
||||
|
||||
def _record_event(self, event: McpEvent) -> None:
|
||||
self.events.append(event)
|
||||
@@ -0,0 +1,95 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from .models import (
|
||||
AuthRecord,
|
||||
CatalogNodeEntry,
|
||||
CatalogPromptEntry,
|
||||
CatalogResourceEntry,
|
||||
CatalogSnapshot,
|
||||
dump_catalog_snapshot,
|
||||
)
|
||||
|
||||
|
||||
class Store:
|
||||
def save_auth(self, record: AuthRecord) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def load_auth(self, connection_id: str) -> AuthRecord | None:
|
||||
raise NotImplementedError
|
||||
|
||||
def save_catalog(self, snapshot: CatalogSnapshot) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def load_catalog(self, connection_id: str) -> CatalogSnapshot | None:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class FileStore(Store):
|
||||
def __init__(self, root: Path) -> None:
|
||||
self.root = root
|
||||
self.root.mkdir(parents=True, exist_ok=True)
|
||||
self.auth_dir.mkdir(parents=True, exist_ok=True)
|
||||
self.catalog_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
@property
|
||||
def auth_dir(self) -> Path:
|
||||
return self.root / "auth"
|
||||
|
||||
@property
|
||||
def catalog_dir(self) -> Path:
|
||||
return self.root / "catalog"
|
||||
|
||||
def _auth_path(self, connection_id: str) -> Path:
|
||||
return self.auth_dir / f"{connection_id}.json"
|
||||
|
||||
def _catalog_path(self, connection_id: str) -> Path:
|
||||
return self.catalog_dir / f"{connection_id}.json"
|
||||
|
||||
def save_auth(self, record: AuthRecord) -> None:
|
||||
self._auth_path(record.connection_id).write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"connection_id": record.connection_id,
|
||||
"scheme": record.scheme,
|
||||
"payload": record.payload,
|
||||
},
|
||||
indent=2,
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
def load_auth(self, connection_id: str) -> AuthRecord | None:
|
||||
path = self._auth_path(connection_id)
|
||||
if not path.exists():
|
||||
return None
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
return AuthRecord(**data)
|
||||
|
||||
def save_catalog(self, snapshot: CatalogSnapshot) -> None:
|
||||
self._catalog_path(snapshot.connection_id).write_text(
|
||||
json.dumps(dump_catalog_snapshot(snapshot), indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
def load_catalog(self, connection_id: str) -> CatalogSnapshot | None:
|
||||
path = self._catalog_path(connection_id)
|
||||
if not path.exists():
|
||||
return None
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
return CatalogSnapshot(
|
||||
connection_id=data["connection_id"],
|
||||
fetched_at_epoch_ms=data["fetched_at_epoch_ms"],
|
||||
max_age_seconds=data["max_age_seconds"],
|
||||
nodes=[CatalogNodeEntry(**node) for node in data.get("nodes", [])],
|
||||
resources=[
|
||||
CatalogResourceEntry(**resource)
|
||||
for resource in data.get("resources", [])
|
||||
],
|
||||
prompts=[
|
||||
CatalogPromptEntry(**prompt) for prompt in data.get("prompts", [])
|
||||
],
|
||||
metadata=data.get("metadata", {}),
|
||||
)
|
||||
@@ -0,0 +1,95 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Any, cast
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, create_model
|
||||
|
||||
from wf_authoring import NodeReturn, NodeSpec
|
||||
from wf_core import RuntimeContext
|
||||
|
||||
from .adapters import BackendAdapter, DiscoveredTool
|
||||
from .events import McpEvent, make_event
|
||||
from .models import AuthRecord, ConnectionConfig
|
||||
|
||||
|
||||
def _model_from_schema(name: str, schema: dict[str, Any]) -> type[BaseModel]:
|
||||
properties = cast(dict[str, Any], schema.get("properties", {}))
|
||||
required = set(cast(list[str], schema.get("required", [])))
|
||||
field_defs: dict[str, tuple[object, object]] = {}
|
||||
|
||||
for field_name in properties:
|
||||
default = ... if field_name in required else None
|
||||
field_defs[field_name] = (Any, Field(default=default))
|
||||
|
||||
raw_field_defs = cast(dict[str, Any], field_defs)
|
||||
model = create_model(
|
||||
name,
|
||||
__config__=ConfigDict(extra="allow"),
|
||||
**raw_field_defs,
|
||||
)
|
||||
return cast(type[BaseModel], model)
|
||||
|
||||
|
||||
def wrap_discovered_tool(
|
||||
*,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
adapter: BackendAdapter,
|
||||
tool: DiscoveredTool,
|
||||
emit_event: Callable[[McpEvent], None] | None = None,
|
||||
) -> NodeSpec[BaseModel, BaseModel]:
|
||||
input_model = _model_from_schema(
|
||||
f"{connection.id}_{tool.name}_Input",
|
||||
tool.input_schema,
|
||||
)
|
||||
output_model = _model_from_schema(
|
||||
f"{connection.id}_{tool.name}_Output",
|
||||
tool.output_schema,
|
||||
)
|
||||
|
||||
async def invoke_tool(
|
||||
payload: BaseModel,
|
||||
ctx: RuntimeContext,
|
||||
) -> NodeReturn[BaseModel]:
|
||||
if emit_event is not None:
|
||||
emit_event(
|
||||
make_event(
|
||||
"tool_call_started",
|
||||
connection_id=connection.id,
|
||||
capability_id=f"{connection.id}.{tool.name}",
|
||||
payload={"input": payload.model_dump()},
|
||||
)
|
||||
)
|
||||
result = await adapter.call_tool(
|
||||
connection=connection,
|
||||
auth=auth,
|
||||
tool_name=tool.name,
|
||||
payload=payload.model_dump(),
|
||||
)
|
||||
if emit_event is not None:
|
||||
emit_event(
|
||||
make_event(
|
||||
"tool_call_completed",
|
||||
connection_id=connection.id,
|
||||
capability_id=f"{connection.id}.{tool.name}",
|
||||
payload={
|
||||
"outcome": result.outcome,
|
||||
"meta": result.meta,
|
||||
},
|
||||
)
|
||||
)
|
||||
return NodeReturn(
|
||||
outcome=result.outcome,
|
||||
output=output_model.model_validate(result.output),
|
||||
)
|
||||
|
||||
return NodeSpec(
|
||||
name=tool.name,
|
||||
input_model=input_model,
|
||||
output_model=output_model,
|
||||
outcomes=tool.outcomes,
|
||||
fn=invoke_tool,
|
||||
description=tool.description,
|
||||
is_async=True,
|
||||
)
|
||||
Reference in New Issue
Block a user