68 lines
2.2 KiB
Python
68 lines
2.2 KiB
Python
import httpx
|
|
|
|
from ai_orchestrator.config import (
|
|
ExecutionConfig,
|
|
McpServerConfig,
|
|
ModelsConfig,
|
|
ModelSlotConfig,
|
|
PolicyConfig,
|
|
ProjectConfig,
|
|
ProjectMetadata,
|
|
)
|
|
from ai_orchestrator.infrastructure.mcp_client import McpHttpClient, McpToolGateway
|
|
from ai_orchestrator.infrastructure.storage.memory import (
|
|
InMemoryArtifactStore,
|
|
InMemoryInvocationStore,
|
|
)
|
|
|
|
|
|
def test_mcp_tool_gateway_normalizes_result_and_persists_invocation() -> None:
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
assert request.url.path == "/rpc"
|
|
payload = {
|
|
"jsonrpc": "2.0",
|
|
"id": "1",
|
|
"result": {
|
|
"status": "success",
|
|
"content": {"rows": [{"id": 1}]},
|
|
"artifacts": [],
|
|
"logs": ["done"],
|
|
"error": None,
|
|
},
|
|
}
|
|
return httpx.Response(200, json=payload)
|
|
|
|
transport = httpx.MockTransport(handler)
|
|
project = ProjectConfig(
|
|
project=ProjectMetadata(id="default", name="Default"),
|
|
models=ModelsConfig(
|
|
weak=ModelSlotConfig(provider="local", model="weak"),
|
|
strong=ModelSlotConfig(provider="local", model="strong"),
|
|
vision=ModelSlotConfig(provider="disabled"),
|
|
embedding=ModelSlotConfig(provider="local", model="embed"),
|
|
),
|
|
mcp_servers={"one_c": McpServerConfig(base_url="http://mcp.test")},
|
|
execution=ExecutionConfig(),
|
|
policy=PolicyConfig(),
|
|
)
|
|
invocation_store = InMemoryInvocationStore()
|
|
gateway = McpToolGateway(
|
|
project_configs={"default": project},
|
|
client=McpHttpClient(transport=transport),
|
|
invocation_store=invocation_store,
|
|
artifact_store=InMemoryArtifactStore(),
|
|
)
|
|
|
|
result = gateway.call(
|
|
project_id="default",
|
|
server_id="one_c",
|
|
tool_name="one_c.run_sql",
|
|
args={"query": "select 1"},
|
|
task_context={"task_id": "task_1", "node_id": "node_2"},
|
|
)
|
|
|
|
assert result.status == "success"
|
|
assert result.content["rows"][0]["id"] == 1
|
|
assert result.logs == ["done"]
|
|
assert len(invocation_store.list_tool_invocations("task_1")) == 1
|