This commit is contained in:
lda
2026-04-29 17:29:48 +07:00 Verified
parent 7e020c1291
commit a0819e8988
5 changed files with 240 additions and 1 deletions
+90 -1
View File
@@ -2,12 +2,21 @@ from __future__ import annotations
import asyncio
from pathlib import Path
from typing import Any
from pydantic import BaseModel
from wf_authoring import NodeReturn, node
from wf_core import END, RuntimeContext, RunStatus
from wf_mcp import AuthRecord, ConnectionConfig, FileStore, RawWorkflowPlan, WfMcpService
from wf_mcp import (
AuthRecord,
ConnectionConfig,
DiscoveredTool,
FileStore,
RawWorkflowPlan,
ToolCallResult,
WfMcpService,
)
class EchoInput(BaseModel):
@@ -47,6 +56,44 @@ def _local_temp_root() -> Path:
return root
class FakeAdapter:
async def list_tools(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
) -> list[DiscoveredTool]:
return [
DiscoveredTool(
name="echo_tool",
description="Echo text back",
input_schema={
"type": "object",
"properties": {"text": {"type": "string"}},
"required": ["text"],
},
output_schema={
"type": "object",
"properties": {"echoed": {"type": "string"}},
"required": ["echoed"],
},
)
]
async def call_tool(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
tool_name: str,
payload: dict[str, Any],
) -> ToolCallResult:
if tool_name != "echo_tool":
raise KeyError(tool_name)
return ToolCallResult(
outcome="ok",
output={"echoed": str(payload["text"])},
)
def test_file_store_round_trips_auth() -> None:
store = FileStore(_local_temp_root() / "auth_store")
record = AuthRecord(
@@ -129,3 +176,45 @@ def test_service_compiles_and_runs_raw_plan() -> None:
assert run.status == RunStatus.COMPLETED
assert run.output == {"result": "final:hello"}
def test_service_refreshes_catalog_from_adapter() -> None:
service = WfMcpService(store=FileStore(_local_temp_root() / "adapter_store"))
service.register_connection(
ConnectionConfig(id="demo.personal", server="demo", account="personal")
)
service.save_auth(
AuthRecord(
connection_id="demo.personal",
scheme="token",
payload={"token": "abc"},
)
)
service.register_adapter("demo", FakeAdapter())
asyncio.run(service.refresh_connection_catalog("demo.personal"))
payload = service.get_catalog().as_payload()
assert payload["nodes"] == [
{
"qualified_name": "demo.personal.echo_tool",
"connection_id": "demo.personal",
"local_name": "echo_tool",
"description": "Echo text back",
"outcomes": ["ok"],
"input_schema": {
"additionalProperties": True,
"properties": {"text": {"title": "Text"}},
"required": ["text"],
"title": "demo.personal_echo_tool_Input",
"type": "object",
},
"output_schema": {
"additionalProperties": True,
"properties": {"echoed": {"title": "Echoed"}},
"required": ["echoed"],
"title": "demo.personal_echo_tool_Output",
"type": "object",
},
}
]
+6
View File
@@ -1,3 +1,4 @@
from .adapters import BackendAdapter, DiscoveredTool, ToolCallResult
from .catalog import CombinedCatalog
from .connections import ConnectionRegistry, parse_connection_id, qualify_node_name
from .models import (
@@ -9,18 +10,23 @@ from .models import (
)
from .service import WfMcpService
from .store import FileStore, Store
from .wrappers import wrap_discovered_tool
__all__ = [
"AuthRecord",
"BackendAdapter",
"CatalogNodeEntry",
"CatalogSnapshot",
"CombinedCatalog",
"ConnectionConfig",
"ConnectionRegistry",
"DiscoveredTool",
"FileStore",
"RawWorkflowPlan",
"Store",
"ToolCallResult",
"WfMcpService",
"parse_connection_id",
"qualify_node_name",
"wrap_discovered_tool",
]
+39
View File
@@ -0,0 +1,39 @@
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Protocol
from .models import AuthRecord, ConnectionConfig
@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 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 call_tool(
self,
connection: ConnectionConfig,
auth: AuthRecord | None,
tool_name: str,
payload: dict[str, Any],
) -> ToolCallResult: ...
+34
View File
@@ -7,10 +7,12 @@ 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 .models import AuthRecord, CatalogSnapshot, ConnectionConfig, RawWorkflowPlan
from .store import Store
from .wrappers import wrap_discovered_tool
def _qualify_spec(connection_id: str, spec: NodeSpec[Any, Any]) -> NodeSpec[Any, Any]:
@@ -30,6 +32,7 @@ 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
)
@@ -38,6 +41,9 @@ class WfMcpService:
parse_connection_id(connection.id)
self.connections.register(connection)
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)
@@ -72,6 +78,34 @@ class WfMcpService:
snapshots[connection.id] = snapshot
return CombinedCatalog(snapshots=snapshots)
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)
tools = await adapter.list_tools(connection, auth)
specs = [
wrap_discovered_tool(
connection=connection,
auth=auth,
adapter=adapter,
tool=tool,
)
for tool in tools
]
self.register_specs(
connection_id,
*specs,
max_age_seconds=max_age_seconds,
)
def compile_plan(self, plan: RawWorkflowPlan) -> Workflow:
node_defs: dict[str, Any] = {}
for step in plan.nodes:
+71
View File
@@ -0,0 +1,71 @@
from __future__ import annotations
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 .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,
) -> 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]:
result = await adapter.call_tool(
connection=connection,
auth=auth,
tool_name=tool.name,
payload=payload.model_dump(),
)
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,
)