discovery
This commit is contained in:
@@ -1,419 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
from wf_authoring import NodeReturn, node
|
||||
from wf_core import END, RuntimeContext, RunStatus
|
||||
from wf_mcp import (
|
||||
AuthRecord,
|
||||
ConnectionConfig,
|
||||
DiscoveredPrompt,
|
||||
DiscoveredResource,
|
||||
DiscoveredTool,
|
||||
FileStore,
|
||||
McpSdkAdapter,
|
||||
RawWorkflowPlan,
|
||||
ToolCallResult,
|
||||
WfMcpService,
|
||||
)
|
||||
|
||||
|
||||
class EchoInput(BaseModel):
|
||||
text: str
|
||||
|
||||
|
||||
class EchoOutput(BaseModel):
|
||||
echoed: str
|
||||
|
||||
|
||||
class FinalizeInput(BaseModel):
|
||||
echoed: str
|
||||
|
||||
|
||||
class FinalizeOutput(BaseModel):
|
||||
result: str
|
||||
|
||||
|
||||
@node()
|
||||
async def echo_tool(payload: EchoInput, ctx: RuntimeContext) -> EchoOutput:
|
||||
return EchoOutput(echoed=payload.text)
|
||||
|
||||
|
||||
@node(outcomes=("done",))
|
||||
def finalize_tool(
|
||||
payload: FinalizeInput, ctx: RuntimeContext
|
||||
) -> NodeReturn[FinalizeOutput]:
|
||||
return NodeReturn(
|
||||
outcome="done",
|
||||
output=FinalizeOutput(result=f"final:{payload.echoed}"),
|
||||
)
|
||||
|
||||
|
||||
def _local_temp_root() -> Path:
|
||||
root = Path("test-artifacts") / "wf_mcp_store"
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
return root
|
||||
|
||||
|
||||
def _fixture_server_path() -> str:
|
||||
return str((Path(__file__).parent / "fixtures" / "mcp_echo_server.py").resolve())
|
||||
|
||||
|
||||
def _everything_server_connection() -> ConnectionConfig | None:
|
||||
transport = os.environ.get("MCP_EVERYTHING_TRANSPORT", "stdio")
|
||||
if transport == "stdio":
|
||||
command = os.environ.get("MCP_EVERYTHING_COMMAND")
|
||||
if not command:
|
||||
return None
|
||||
|
||||
raw_args = os.environ.get("MCP_EVERYTHING_ARGS", "")
|
||||
args = [arg for arg in raw_args.split(" ") if arg]
|
||||
|
||||
metadata: dict[str, Any] = {
|
||||
"transport": transport,
|
||||
"command": command,
|
||||
"args": args,
|
||||
}
|
||||
|
||||
elif transport == "streamable_http":
|
||||
url = os.environ.get("MCP_EVERYTHING_URL")
|
||||
if not url:
|
||||
raise AssertionError(
|
||||
"MCP_EVERYTHING_URL must be set when MCP_EVERYTHING_TRANSPORT=streamable_http"
|
||||
)
|
||||
metadata = {
|
||||
"transport": "streamable_http",
|
||||
"url": url,
|
||||
}
|
||||
else:
|
||||
return None
|
||||
return ConnectionConfig(
|
||||
id="everything.default",
|
||||
server="everything",
|
||||
account="default",
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
|
||||
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 list_resources(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
) -> list[DiscoveredResource]:
|
||||
return [
|
||||
DiscoveredResource(
|
||||
uri="demo://docs/welcome",
|
||||
name="resource.welcome",
|
||||
description="Welcome resource",
|
||||
mime_type="text/plain",
|
||||
metadata={"kind": "static"},
|
||||
)
|
||||
]
|
||||
|
||||
async def list_prompts(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
) -> list[DiscoveredPrompt]:
|
||||
return [
|
||||
DiscoveredPrompt(
|
||||
name="prompt.summarize",
|
||||
description="Summarize text",
|
||||
arguments=[
|
||||
{
|
||||
"name": "text",
|
||||
"required": True,
|
||||
"description": "Text to summarize",
|
||||
}
|
||||
],
|
||||
metadata={"kind": "template"},
|
||||
)
|
||||
]
|
||||
|
||||
async def get_connection_metadata(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"server": connection.server,
|
||||
"account": connection.account,
|
||||
"auth_scheme": auth.scheme if auth is not None else None,
|
||||
}
|
||||
|
||||
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(
|
||||
connection_id="demo.personal",
|
||||
scheme="oauth",
|
||||
payload={"token": "secret"},
|
||||
)
|
||||
|
||||
store.save_auth(record)
|
||||
loaded = store.load_auth("demo.personal")
|
||||
|
||||
assert loaded == record
|
||||
|
||||
|
||||
def test_service_builds_namespaced_catalog() -> None:
|
||||
service = WfMcpService(store=FileStore(_local_temp_root() / "catalog_store"))
|
||||
service.register_connection(
|
||||
ConnectionConfig(id="demo.personal", server="demo", account="personal")
|
||||
)
|
||||
service.register_specs("demo.personal", echo_tool, finalize_tool)
|
||||
|
||||
payload = service.get_catalog().as_payload()
|
||||
names = [node["qualified_name"] for node in payload["nodes"]]
|
||||
|
||||
assert names == [
|
||||
"demo.personal.echo_tool",
|
||||
"demo.personal.finalize_tool",
|
||||
]
|
||||
|
||||
|
||||
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, finalize_tool)
|
||||
|
||||
plan = RawWorkflowPlan(
|
||||
name="demo_plan",
|
||||
input_schema={
|
||||
"type": "object",
|
||||
"properties": {"text": {"type": "string"}},
|
||||
"required": ["text"],
|
||||
},
|
||||
state_schema={
|
||||
"fields": {
|
||||
"echoed": {"type": "string"},
|
||||
"result": {"type": "string"},
|
||||
}
|
||||
},
|
||||
output_schema={
|
||||
"type": "object",
|
||||
"properties": {"result": {"type": "string"}},
|
||||
"required": ["result"],
|
||||
},
|
||||
start="echo",
|
||||
nodes=[
|
||||
{
|
||||
"id": "echo",
|
||||
"type": "node",
|
||||
"node": "demo.personal.echo_tool",
|
||||
"in_map": {"input.text": "text"},
|
||||
"out_map": {"echoed": "state.echoed"},
|
||||
},
|
||||
{
|
||||
"id": "finalize",
|
||||
"type": "node",
|
||||
"node": "demo.personal.finalize_tool",
|
||||
"in_map": {"state.echoed": "echoed"},
|
||||
"out_map": {"result": "state.result"},
|
||||
},
|
||||
],
|
||||
edges=[
|
||||
{"from": "echo", "outcome": "ok", "to": "finalize"},
|
||||
{"from": "finalize", "outcome": "done", "to": END},
|
||||
],
|
||||
)
|
||||
|
||||
run = asyncio.run(service.run_workflow_from_plan(plan, {"text": "hello"}))
|
||||
|
||||
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",
|
||||
},
|
||||
}
|
||||
]
|
||||
assert payload["resources"] == [
|
||||
{
|
||||
"qualified_name": "demo.personal.resource.welcome",
|
||||
"connection_id": "demo.personal",
|
||||
"local_name": "resource.welcome",
|
||||
"uri": "demo://docs/welcome",
|
||||
"description": "Welcome resource",
|
||||
"mime_type": "text/plain",
|
||||
"metadata": {"kind": "static"},
|
||||
}
|
||||
]
|
||||
assert payload["prompts"] == [
|
||||
{
|
||||
"qualified_name": "demo.personal.prompt.summarize",
|
||||
"connection_id": "demo.personal",
|
||||
"local_name": "prompt.summarize",
|
||||
"description": "Summarize text",
|
||||
"arguments": [
|
||||
{
|
||||
"name": "text",
|
||||
"required": True,
|
||||
"description": "Text to summarize",
|
||||
}
|
||||
],
|
||||
"metadata": {"kind": "template"},
|
||||
}
|
||||
]
|
||||
assert payload["connections"] == [
|
||||
{
|
||||
"connection_id": "demo.personal",
|
||||
"fetched_at_epoch_ms": payload["connections"][0]["fetched_at_epoch_ms"],
|
||||
"max_age_seconds": 300,
|
||||
"metadata": {
|
||||
"server": "demo",
|
||||
"account": "personal",
|
||||
"auth_scheme": "token",
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_mcp_sdk_adapter_lists_and_calls_stdio_tool() -> None:
|
||||
service = WfMcpService(store=FileStore(_local_temp_root() / "sdk_adapter_store"))
|
||||
service.register_connection(
|
||||
ConnectionConfig(
|
||||
id="fixture.personal",
|
||||
server="fixture",
|
||||
account="personal",
|
||||
metadata={
|
||||
"transport": "stdio",
|
||||
"command": sys.executable,
|
||||
"args": [_fixture_server_path()],
|
||||
},
|
||||
)
|
||||
)
|
||||
service.register_adapter("fixture", McpSdkAdapter())
|
||||
|
||||
try:
|
||||
asyncio.run(service.refresh_connection_catalog("fixture.personal"))
|
||||
except PermissionError as exc:
|
||||
pytest.skip(f"stdio MCP transport is not permitted in this environment: {exc}")
|
||||
|
||||
payload = service.get_catalog().as_payload()
|
||||
assert payload["nodes"][0]["qualified_name"] == "fixture.personal.echo_tool"
|
||||
|
||||
adapter = McpSdkAdapter()
|
||||
try:
|
||||
result = asyncio.run(
|
||||
adapter.call_tool(
|
||||
connection=service.connections.get("fixture.personal"),
|
||||
auth=None,
|
||||
tool_name="echo_tool",
|
||||
payload={"text": "hello"},
|
||||
)
|
||||
)
|
||||
except PermissionError as exc:
|
||||
pytest.skip(f"stdio MCP transport is not permitted in this environment: {exc}")
|
||||
assert result.outcome == "ok"
|
||||
assert result.output == {"echoed": "hello"}
|
||||
|
||||
|
||||
def test_mcp_sdk_adapter_can_probe_everything_server() -> None:
|
||||
connection = _everything_server_connection()
|
||||
if connection is None:
|
||||
pytest.skip(
|
||||
"set MCP_EVERYTHING_COMMAND to enable the live everything-server integration test"
|
||||
)
|
||||
|
||||
service = WfMcpService(
|
||||
store=FileStore(_local_temp_root() / "everything_server_store")
|
||||
)
|
||||
service.register_connection(connection)
|
||||
service.register_adapter("everything", McpSdkAdapter())
|
||||
|
||||
try:
|
||||
asyncio.run(service.refresh_connection_catalog("everything.default"))
|
||||
except PermissionError as exc:
|
||||
pytest.skip(f"live MCP transport is not permitted in this environment: {exc}")
|
||||
|
||||
payload = service.get_catalog().as_payload()
|
||||
assert payload["nodes"], "everything-server should expose at least one tool"
|
||||
assert all(
|
||||
node["qualified_name"].startswith("everything.default.")
|
||||
for node in payload["nodes"]
|
||||
)
|
||||
assert "resources" in payload
|
||||
assert "prompts" in payload
|
||||
@@ -0,0 +1,82 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from wf_mcp import ConnectionConfig, FileStore, McpSdkAdapter, WfMcpService
|
||||
|
||||
from test_wf_mcp_support import (
|
||||
everything_server_connection,
|
||||
fixture_server_path,
|
||||
local_temp_root,
|
||||
sys,
|
||||
)
|
||||
|
||||
|
||||
def test_mcp_sdk_adapter_lists_and_calls_stdio_tool() -> None:
|
||||
service = WfMcpService(store=FileStore(local_temp_root() / "sdk_adapter_store"))
|
||||
service.register_connection(
|
||||
ConnectionConfig(
|
||||
id="fixture.personal",
|
||||
server="fixture",
|
||||
account="personal",
|
||||
metadata={
|
||||
"transport": "stdio",
|
||||
"command": sys.executable,
|
||||
"args": [fixture_server_path()],
|
||||
},
|
||||
)
|
||||
)
|
||||
service.register_adapter("fixture", McpSdkAdapter())
|
||||
|
||||
try:
|
||||
asyncio.run(service.refresh_connection_catalog("fixture.personal"))
|
||||
except PermissionError as exc:
|
||||
pytest.skip(f"stdio MCP transport is not permitted in this environment: {exc}")
|
||||
|
||||
payload = service.get_catalog().as_payload()
|
||||
assert payload["nodes"][0]["qualified_name"] == "fixture.personal.echo_tool"
|
||||
|
||||
adapter = McpSdkAdapter()
|
||||
try:
|
||||
result = asyncio.run(
|
||||
adapter.call_tool(
|
||||
connection=service.connections.get("fixture.personal"),
|
||||
auth=None,
|
||||
tool_name="echo_tool",
|
||||
payload={"text": "hello"},
|
||||
)
|
||||
)
|
||||
except PermissionError as exc:
|
||||
pytest.skip(f"stdio MCP transport is not permitted in this environment: {exc}")
|
||||
assert result.outcome == "ok"
|
||||
assert result.output == {"echoed": "hello"}
|
||||
|
||||
|
||||
def test_mcp_sdk_adapter_can_probe_everything_server() -> None:
|
||||
connection = everything_server_connection()
|
||||
if connection is None:
|
||||
pytest.skip(
|
||||
"set MCP_EVERYTHING_COMMAND to enable the live everything-server integration test"
|
||||
)
|
||||
|
||||
service = WfMcpService(
|
||||
store=FileStore(local_temp_root() / "everything_server_store")
|
||||
)
|
||||
service.register_connection(connection)
|
||||
service.register_adapter("everything", McpSdkAdapter())
|
||||
|
||||
try:
|
||||
asyncio.run(service.refresh_connection_catalog("everything.default"))
|
||||
except PermissionError as exc:
|
||||
pytest.skip(f"live MCP transport is not permitted in this environment: {exc}")
|
||||
|
||||
payload = service.get_catalog().as_payload()
|
||||
assert payload["nodes"], "everything-server should expose at least one tool"
|
||||
assert all(
|
||||
node["qualified_name"].startswith("everything.default.")
|
||||
for node in payload["nodes"]
|
||||
)
|
||||
assert "resources" in payload
|
||||
assert "prompts" in payload
|
||||
@@ -0,0 +1,165 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
from wf_core import END, RunStatus
|
||||
from wf_mcp import (
|
||||
AuthRecord,
|
||||
ConnectionConfig,
|
||||
FileStore,
|
||||
RawWorkflowPlan,
|
||||
WfMcpService,
|
||||
)
|
||||
|
||||
from test_wf_mcp_support import FakeAdapter, echo_tool, finalize_tool, local_temp_root
|
||||
|
||||
|
||||
def test_service_builds_namespaced_catalog() -> None:
|
||||
service = WfMcpService(store=FileStore(local_temp_root() / "catalog_store"))
|
||||
service.register_connection(
|
||||
ConnectionConfig(id="demo.personal", server="demo", account="personal")
|
||||
)
|
||||
service.register_specs("demo.personal", echo_tool, finalize_tool)
|
||||
|
||||
payload = service.get_catalog().as_payload()
|
||||
names = [node["qualified_name"] for node in payload["nodes"]]
|
||||
|
||||
assert names == [
|
||||
"demo.personal.echo_tool",
|
||||
"demo.personal.finalize_tool",
|
||||
]
|
||||
|
||||
|
||||
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, finalize_tool)
|
||||
|
||||
plan = RawWorkflowPlan(
|
||||
name="demo_plan",
|
||||
input_schema={
|
||||
"type": "object",
|
||||
"properties": {"text": {"type": "string"}},
|
||||
"required": ["text"],
|
||||
},
|
||||
state_schema={
|
||||
"fields": {
|
||||
"echoed": {"type": "string"},
|
||||
"result": {"type": "string"},
|
||||
}
|
||||
},
|
||||
output_schema={
|
||||
"type": "object",
|
||||
"properties": {"result": {"type": "string"}},
|
||||
"required": ["result"],
|
||||
},
|
||||
start="echo",
|
||||
nodes=[
|
||||
{
|
||||
"id": "echo",
|
||||
"type": "node",
|
||||
"node": "demo.personal.echo_tool",
|
||||
"in_map": {"input.text": "text"},
|
||||
"out_map": {"echoed": "state.echoed"},
|
||||
},
|
||||
{
|
||||
"id": "finalize",
|
||||
"type": "node",
|
||||
"node": "demo.personal.finalize_tool",
|
||||
"in_map": {"state.echoed": "echoed"},
|
||||
"out_map": {"result": "state.result"},
|
||||
},
|
||||
],
|
||||
edges=[
|
||||
{"from": "echo", "outcome": "ok", "to": "finalize"},
|
||||
{"from": "finalize", "outcome": "done", "to": END},
|
||||
],
|
||||
)
|
||||
|
||||
run = asyncio.run(service.run_workflow_from_plan(plan, {"text": "hello"}))
|
||||
|
||||
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",
|
||||
},
|
||||
}
|
||||
]
|
||||
assert payload["resources"] == [
|
||||
{
|
||||
"qualified_name": "demo.personal.resource.welcome",
|
||||
"connection_id": "demo.personal",
|
||||
"local_name": "resource.welcome",
|
||||
"uri": "demo://docs/welcome",
|
||||
"description": "Welcome resource",
|
||||
"mime_type": "text/plain",
|
||||
"metadata": {"kind": "static"},
|
||||
}
|
||||
]
|
||||
assert payload["prompts"] == [
|
||||
{
|
||||
"qualified_name": "demo.personal.prompt.summarize",
|
||||
"connection_id": "demo.personal",
|
||||
"local_name": "prompt.summarize",
|
||||
"description": "Summarize text",
|
||||
"arguments": [
|
||||
{
|
||||
"name": "text",
|
||||
"required": True,
|
||||
"description": "Text to summarize",
|
||||
}
|
||||
],
|
||||
"metadata": {"kind": "template"},
|
||||
}
|
||||
]
|
||||
assert payload["connections"] == [
|
||||
{
|
||||
"connection_id": "demo.personal",
|
||||
"fetched_at_epoch_ms": payload["connections"][0]["fetched_at_epoch_ms"],
|
||||
"max_age_seconds": 300,
|
||||
"metadata": {
|
||||
"server": "demo",
|
||||
"account": "personal",
|
||||
"auth_scheme": "token",
|
||||
},
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,19 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from wf_mcp import AuthRecord, FileStore
|
||||
|
||||
from test_wf_mcp_support import local_temp_root
|
||||
|
||||
|
||||
def test_file_store_round_trips_auth() -> None:
|
||||
store = FileStore(local_temp_root() / "auth_store")
|
||||
record = AuthRecord(
|
||||
connection_id="demo.personal",
|
||||
scheme="oauth",
|
||||
payload={"token": "secret"},
|
||||
)
|
||||
|
||||
store.save_auth(record)
|
||||
loaded = store.load_auth("demo.personal")
|
||||
|
||||
assert loaded == record
|
||||
@@ -0,0 +1,191 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from wf_authoring import NodeReturn, node
|
||||
from wf_core import RuntimeContext
|
||||
from wf_mcp import (
|
||||
AuthRecord,
|
||||
ConnectionConfig,
|
||||
DiscoveredPrompt,
|
||||
DiscoveredResource,
|
||||
DiscoveredTool,
|
||||
ToolCallResult,
|
||||
)
|
||||
|
||||
|
||||
class EchoInput(BaseModel):
|
||||
text: str
|
||||
|
||||
|
||||
class EchoOutput(BaseModel):
|
||||
echoed: str
|
||||
|
||||
|
||||
class FinalizeInput(BaseModel):
|
||||
echoed: str
|
||||
|
||||
|
||||
class FinalizeOutput(BaseModel):
|
||||
result: str
|
||||
|
||||
|
||||
@node()
|
||||
async def echo_tool(payload: EchoInput, ctx: RuntimeContext) -> EchoOutput:
|
||||
return EchoOutput(echoed=payload.text)
|
||||
|
||||
|
||||
@node(outcomes=("done",))
|
||||
def finalize_tool(
|
||||
payload: FinalizeInput, ctx: RuntimeContext
|
||||
) -> NodeReturn[FinalizeOutput]:
|
||||
return NodeReturn(
|
||||
outcome="done",
|
||||
output=FinalizeOutput(result=f"final:{payload.echoed}"),
|
||||
)
|
||||
|
||||
|
||||
def local_temp_root() -> Path:
|
||||
root = Path("test-artifacts") / "wf_mcp_store"
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
return root
|
||||
|
||||
|
||||
def fixture_server_path() -> str:
|
||||
return str((Path(__file__).resolve().parent / "fixtures" / "mcp_echo_server.py"))
|
||||
|
||||
|
||||
def everything_server_connection() -> ConnectionConfig | None:
|
||||
transport = os.environ.get("MCP_EVERYTHING_TRANSPORT", "stdio")
|
||||
if transport == "stdio":
|
||||
command = os.environ.get("MCP_EVERYTHING_COMMAND")
|
||||
if not command:
|
||||
return None
|
||||
|
||||
raw_args = os.environ.get("MCP_EVERYTHING_ARGS", "")
|
||||
args = [arg for arg in raw_args.split(" ") if arg]
|
||||
|
||||
metadata: dict[str, Any] = {
|
||||
"transport": transport,
|
||||
"command": command,
|
||||
"args": args,
|
||||
}
|
||||
elif transport == "streamable_http":
|
||||
url = os.environ.get("MCP_EVERYTHING_URL")
|
||||
if not url:
|
||||
raise AssertionError(
|
||||
"MCP_EVERYTHING_URL must be set when MCP_EVERYTHING_TRANSPORT=streamable_http"
|
||||
)
|
||||
metadata = {
|
||||
"transport": "streamable_http",
|
||||
"url": url,
|
||||
}
|
||||
else:
|
||||
return None
|
||||
|
||||
return ConnectionConfig(
|
||||
id="everything.default",
|
||||
server="everything",
|
||||
account="default",
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
|
||||
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 list_resources(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
) -> list[DiscoveredResource]:
|
||||
return [
|
||||
DiscoveredResource(
|
||||
uri="demo://docs/welcome",
|
||||
name="resource.welcome",
|
||||
description="Welcome resource",
|
||||
mime_type="text/plain",
|
||||
metadata={"kind": "static"},
|
||||
)
|
||||
]
|
||||
|
||||
async def list_prompts(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
) -> list[DiscoveredPrompt]:
|
||||
return [
|
||||
DiscoveredPrompt(
|
||||
name="prompt.summarize",
|
||||
description="Summarize text",
|
||||
arguments=[
|
||||
{
|
||||
"name": "text",
|
||||
"required": True,
|
||||
"description": "Text to summarize",
|
||||
}
|
||||
],
|
||||
metadata={"kind": "template"},
|
||||
)
|
||||
]
|
||||
|
||||
async def get_connection_metadata(
|
||||
self,
|
||||
connection: ConnectionConfig,
|
||||
auth: AuthRecord | None,
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"server": connection.server,
|
||||
"account": connection.account,
|
||||
"auth_scheme": auth.scheme if auth is not None else None,
|
||||
}
|
||||
|
||||
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"])},
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"FakeAdapter",
|
||||
"echo_tool",
|
||||
"everything_server_connection",
|
||||
"finalize_tool",
|
||||
"fixture_server_path",
|
||||
"local_temp_root",
|
||||
"sys",
|
||||
]
|
||||
Reference in New Issue
Block a user