This commit is contained in:
lda
2026-04-29 22:34:12 +07:00 Verified
parent 8337d7615c
commit 1e41a94fc6
7 changed files with 196 additions and 6 deletions
+61
View File
@@ -82,6 +82,9 @@ def test_service_compiles_and_runs_raw_plan() -> None:
assert run.status == RunStatus.COMPLETED
assert run.output == {"result": "final:hello"}
event_kinds = [event.kind for event in service.list_events()]
assert "workflow_run_started" in event_kinds
assert "workflow_run_completed" in event_kinds
def test_service_refreshes_catalog_from_adapter() -> None:
@@ -163,3 +166,61 @@ def test_service_refreshes_catalog_from_adapter() -> None:
},
}
]
event_kinds = [event.kind for event in service.list_events()]
assert "catalog_refresh_started" in event_kinds
assert "catalog_refresh_completed" in event_kinds
def test_service_records_tool_call_events() -> None:
service = WfMcpService(store=FileStore(local_temp_root() / "event_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"))
plan = RawWorkflowPlan(
name="tool_only_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",
"in_map": {"input.text": "text"},
"out_map": {"echoed": "state.echoed"},
}
],
edges=[
{"from": "echo", "outcome": "ok", "to": END},
],
)
run = asyncio.run(service.run_workflow_from_plan(plan, {"text": "hello"}))
assert run.status == RunStatus.COMPLETED
tool_events = [
event for event in service.list_events() if "tool_call" in event.kind
]
assert [event.kind for event in tool_events] == [
"tool_call_started",
"tool_call_completed",
]
assert tool_events[0].capability_id == "demo.personal.echo_tool"
assert tool_events[1].payload["outcome"] == "ok"
+3
View File
@@ -12,6 +12,7 @@ from .discovery import (
discover_connection_capabilities,
specs_from_discovered_tools,
)
from .events import McpEvent, make_event
from .models import (
AuthRecord,
CatalogNodeEntry,
@@ -41,12 +42,14 @@ __all__ = [
"DiscoveredResource",
"DiscoveredTool",
"FileStore",
"McpEvent",
"McpSdkAdapter",
"RawWorkflowPlan",
"Store",
"ToolCallResult",
"WfMcpService",
"discover_connection_capabilities",
"make_event",
"parse_connection_id",
"qualify_node_name",
"specs_from_discovered_tools",
+4
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any
@@ -11,6 +12,7 @@ from .adapters import (
DiscoveredResource,
DiscoveredTool,
)
from .events import McpEvent
from .models import AuthRecord, ConnectionConfig
from .wrappers import wrap_discovered_tool
@@ -47,6 +49,7 @@ def specs_from_discovered_tools(
auth: AuthRecord | None,
adapter: BackendAdapter,
tools: list[DiscoveredTool],
emit_event: Callable[[McpEvent], None] | None = None,
) -> list[NodeSpec[Any, Any]]:
return [
wrap_discovered_tool(
@@ -54,6 +57,7 @@ def specs_from_discovered_tools(
auth=auth,
adapter=adapter,
tool=tool,
emit_event=emit_event,
)
for tool in tools
]
+33
View File
@@ -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 {},
)
+7 -5
View File
@@ -44,29 +44,31 @@ def _tool_to_discovered(tool: McpTool) -> DiscoveredTool:
input_schema=tool.inputSchema,
output_schema=output_schema,
outcomes=("ok", "error"),
metadata=tool.model_dump(by_alias=True),
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=str(resource.uri),
name=local_name,
description=resource.description,
mime_type=resource.mimeType,
metadata=resource.model_dump(by_alias=True),
metadata=resource.model_dump(by_alias=True, mode="json"),
)
def _prompt_to_discovered(prompt: McpPrompt) -> DiscoveredPrompt:
arguments = [
argument.model_dump(by_alias=True) for argument in prompt.arguments or []
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),
metadata=prompt.model_dump(by_alias=True, mode="json"),
)
+64 -1
View File
@@ -11,6 +11,7 @@ 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, CatalogSnapshot, ConnectionConfig, RawWorkflowPlan
from .store import Store
@@ -36,16 +37,31 @@ class WfMcpService:
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)
@@ -71,6 +87,13 @@ class WfMcpService:
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] = {}
@@ -92,6 +115,13 @@ class WfMcpService:
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,
@@ -102,6 +132,7 @@ class WfMcpService:
auth=auth,
adapter=adapter,
tools=capabilities.tools,
emit_event=self._record_event,
)
self.register_specs(
connection_id,
@@ -118,6 +149,17 @@ class WfMcpService:
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] = {}
@@ -145,6 +187,13 @@ class WfMcpService:
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)
@@ -152,7 +201,18 @@ class WfMcpService:
if isinstance(node, NodeUse)
]
registry = build_async_registry(*specs)
return await execute_workflow_async(workflow, workflow_input, registry)
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)
@@ -160,3 +220,6 @@ class WfMcpService:
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)
+24
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
from collections.abc import Callable
from typing import Any, cast
from pydantic import BaseModel, ConfigDict, Field, create_model
@@ -8,6 +9,7 @@ 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
@@ -35,6 +37,7 @@ def wrap_discovered_tool(
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",
@@ -49,12 +52,33 @@ def wrap_discovered_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),