262 lines
9.6 KiB
Python
262 lines
9.6 KiB
Python
import asyncio
|
|
from contextlib import AsyncExitStack
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
from unittest.mock import patch
|
|
|
|
import httpx
|
|
from mcp import ClientSession, StdioServerParameters, types
|
|
from mcp.client.stdio import stdio_client
|
|
import pytest
|
|
|
|
from artemis import service_registry as registry
|
|
from mcp_server.bridge import backend_session, create_bridge, load_catalog
|
|
from mcp_server.shared import ToolExecutor, create_shared_mcp
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_catalog_and_shared_tools_match_original_schemas():
|
|
from mcp_server import mcp
|
|
|
|
shared, _ = create_shared_mcp()
|
|
original = [tool.model_dump(exclude_none=True) for tool in await mcp.list_tools()]
|
|
|
|
def by_name(tools):
|
|
return {tool["name"]: tool for tool in tools}
|
|
|
|
assert by_name(load_catalog()["tools"]) == by_name(original)
|
|
assert load_catalog()["instructions"] == mcp.instructions
|
|
assert by_name(
|
|
[tool.model_dump(exclude_none=True) for tool in await shared.list_tools()]
|
|
) == by_name(original)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stdio_survives_offline_backend_and_recovers_without_restart(shared_mcp_backend):
|
|
registration = shared_mcp_backend["registration"]
|
|
registration.close()
|
|
async with AsyncExitStack() as stack:
|
|
streams = await stack.enter_async_context(
|
|
stdio_client(
|
|
StdioServerParameters(
|
|
command=sys.executable,
|
|
args=["-m", "mcp_server"],
|
|
env={**os.environ, **shared_mcp_backend["env"], "PYTHONUTF8": "1"},
|
|
)
|
|
)
|
|
)
|
|
session = await stack.enter_async_context(ClientSession(*streams))
|
|
initialized = await asyncio.wait_for(session.initialize(), 6)
|
|
assert initialized.capabilities.tools is not None
|
|
tools = await session.list_tools()
|
|
assert "mobile_run_task" in {tool.name for tool in tools.tools}
|
|
offline = await session.call_tool("mobile_workspace", {"action": "queue"})
|
|
assert offline.isError
|
|
assert "不会自动启动" in offline.content[0].text
|
|
assert shared_mcp_backend["calls"] == []
|
|
registration.acquire()
|
|
connected = await session.call_tool("shared_probe", {})
|
|
assert connected.structuredContent["pid"] == os.getpid()
|
|
assert connected.structuredContent["count"] == 1
|
|
registration.close()
|
|
assert (await session.call_tool("shared_probe", {})).isError
|
|
registration.acquire()
|
|
assert (await session.call_tool("shared_probe", {})).structuredContent["count"] == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_multiple_stdio_processes_share_backend_and_disconnect_independently(
|
|
shared_mcp_backend,
|
|
):
|
|
async def probe():
|
|
async with (
|
|
stdio_client(
|
|
StdioServerParameters(
|
|
command=sys.executable,
|
|
args=["-m", "mcp_server"],
|
|
env={**os.environ, **shared_mcp_backend["env"], "PYTHONUTF8": "1"},
|
|
)
|
|
) as streams,
|
|
ClientSession(*streams) as session,
|
|
):
|
|
await session.initialize()
|
|
return await session.call_tool("shared_probe", {"delay": 0.05})
|
|
|
|
results = await asyncio.gather(*(probe() for _ in range(3)))
|
|
assert {result.structuredContent["pid"] for result in results} == {os.getpid()}
|
|
assert len({result.structuredContent["thread"] for result in results}) == 1
|
|
assert len(shared_mcp_backend["calls"]) == 3
|
|
assert (await probe()).structuredContent["count"] == 4
|
|
|
|
|
|
def test_bridge_import_does_not_load_runtime():
|
|
result = subprocess.run(
|
|
[
|
|
sys.executable,
|
|
"-c",
|
|
"import mcp_server.bridge,sys; print([n for n in ('artemis.sdk.agent','artemis.runtime','mcp_server.tools','cv2','torch') if n in sys.modules])",
|
|
],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=15,
|
|
)
|
|
assert result.returncode == 0, result.stderr
|
|
assert result.stdout.strip() == "[]"
|
|
|
|
|
|
def test_missing_backend_never_spawns_or_imports_execution(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(tempfile, "tempdir", str(tmp_path))
|
|
with patch("subprocess.Popen") as spawn:
|
|
with pytest.raises(registry.ServiceUnavailableError, match="不会自动启动"):
|
|
registry.discover_service()
|
|
spawn.assert_not_called()
|
|
|
|
|
|
def test_registration_excludes_second_server_and_releases(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(tempfile, "tempdir", str(tmp_path))
|
|
args = dict(port=8100, token="test", app_dir=tmp_path, traces_dir=tmp_path / "traces")
|
|
first = registry.ServiceRegistration(**args)
|
|
second = registry.ServiceRegistration(**args)
|
|
first.acquire()
|
|
try:
|
|
with pytest.raises(registry.ServiceUnavailableError, match="已有主程序"):
|
|
second.acquire()
|
|
finally:
|
|
first.close()
|
|
second.acquire()
|
|
second.close()
|
|
|
|
|
|
def test_registration_excludes_another_process(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(tempfile, "tempdir", str(tmp_path))
|
|
registration = registry.ServiceRegistration(
|
|
port=8100, token="test", app_dir=tmp_path, traces_dir=tmp_path
|
|
)
|
|
registration.acquire()
|
|
try:
|
|
code = "from artemis.service_registry import ServiceRegistration; import pathlib,sys; ServiceRegistration(port=8101,token='test',app_dir=pathlib.Path(sys.argv[1]),traces_dir=pathlib.Path(sys.argv[1])).acquire()"
|
|
result = subprocess.run(
|
|
[sys.executable, "-c", code, str(tmp_path)],
|
|
env={
|
|
**os.environ,
|
|
"TEMP": str(tmp_path),
|
|
"TMP": str(tmp_path),
|
|
"TMPDIR": str(tmp_path),
|
|
"PYTHONUTF8": "1",
|
|
},
|
|
capture_output=True,
|
|
text=True,
|
|
encoding="utf-8",
|
|
timeout=10,
|
|
)
|
|
assert result.returncode != 0
|
|
assert "已有主程序" in result.stderr
|
|
finally:
|
|
registration.close()
|
|
|
|
|
|
def test_discovery_nondefault_port_and_scope(shared_mcp_backend, monkeypatch, tmp_path):
|
|
info = registry.discover_service()
|
|
assert info["port"] == shared_mcp_backend["port"]
|
|
assert info["pid"] == os.getpid()
|
|
monkeypatch.setenv("ARTEMIS_TRACES_DIR", str(tmp_path / "other"))
|
|
with pytest.raises(registry.ServiceUnavailableError):
|
|
registry.discover_service()
|
|
|
|
|
|
def test_explicit_port_is_not_silently_ignored(shared_mcp_backend, monkeypatch):
|
|
monkeypatch.setenv("ARTEMIS_DAEMON_PORT", "1")
|
|
with pytest.raises(registry.ServiceUnavailableError):
|
|
registry.discover_service()
|
|
|
|
|
|
def test_mcp_endpoint_rejects_unauthenticated_call(shared_mcp_backend):
|
|
response = httpx.post(
|
|
f"http://127.0.0.1:{shared_mcp_backend['port']}/api/mcp/",
|
|
json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"},
|
|
trust_env=False,
|
|
)
|
|
assert response.status_code == 403
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_connections_share_executor_and_preserve_structured_results(shared_mcp_backend):
|
|
async def call():
|
|
async with backend_session() as (session, _):
|
|
return await session.call_tool("shared_probe", {"delay": 0.01})
|
|
|
|
results = await asyncio.gather(*(call() for _ in range(5)))
|
|
assert all(not result.isError for result in results)
|
|
assert all(result.structuredContent is not None for result in results), results
|
|
assert {result.structuredContent["pid"] for result in results} == {os.getpid()}
|
|
assert len({result.structuredContent["thread"] for result in results}) == 1
|
|
assert len(shared_mcp_backend["calls"]) == 5
|
|
# Closing five bridge sessions must not stop the common backend.
|
|
assert registry.discover_service()["pid"] == os.getpid()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_bridge_preserves_tool_errors_and_never_replays(shared_mcp_backend):
|
|
bridge = create_bridge()
|
|
handler = bridge.request_handlers[types.CallToolRequest]
|
|
result = await handler(
|
|
types.CallToolRequest(
|
|
method="tools/call", params=types.CallToolRequestParams(name="missing", arguments={})
|
|
)
|
|
)
|
|
assert result.root.isError
|
|
assert shared_mcp_backend["calls"] == []
|
|
result = await handler(
|
|
types.CallToolRequest(
|
|
method="tools/call",
|
|
params=types.CallToolRequestParams(name="shared_probe", arguments={}),
|
|
)
|
|
)
|
|
assert result.root.structuredContent["count"] == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_executor_cancellation_reaches_coroutine():
|
|
executor = ToolExecutor()
|
|
executor.start()
|
|
import threading
|
|
|
|
started, stopped = threading.Event(), threading.Event()
|
|
|
|
async def operation():
|
|
started.set()
|
|
try:
|
|
await asyncio.sleep(60)
|
|
finally:
|
|
stopped.set()
|
|
|
|
try:
|
|
task = asyncio.create_task(executor.invoke(operation))
|
|
assert await asyncio.to_thread(started.wait, 2)
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
assert await asyncio.to_thread(stopped.wait, 2)
|
|
finally:
|
|
await executor.stop()
|
|
|
|
|
|
def test_stale_pid_and_wrong_installation_are_not_reused(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(tempfile, "tempdir", str(tmp_path))
|
|
directory = registry.registry_directory()
|
|
directory.mkdir(parents=True)
|
|
(directory / "stale.json").write_text(
|
|
json.dumps(
|
|
{"installation": registry.installation_id(), "pid": -1, "port": 8100, "token": "x"}
|
|
)
|
|
)
|
|
(directory / "other.json").write_text(
|
|
json.dumps({"installation": "another", "pid": os.getpid(), "port": 8100, "token": "x"})
|
|
)
|
|
with pytest.raises(registry.ServiceUnavailableError):
|
|
registry.discover_service()
|