1014 lines
38 KiB
Python
1014 lines
38 KiB
Python
"""Plugin transport/installation tests use local deterministic MCP servers, never devices."""
|
|
|
|
import asyncio
|
|
import base64
|
|
import io
|
|
import json
|
|
from pathlib import Path
|
|
import sys
|
|
import zipfile
|
|
|
|
from mcp.types import CallToolResult, ImageContent, TextContent
|
|
import pytest
|
|
|
|
from artemis.plugins.installer import install_archive
|
|
from artemis.plugins.models import CatalogEntry, PluginManifest, PluginProfile, PluginSelection
|
|
from artemis.plugins.runtime import PluginRuntime
|
|
from artemis.plugins.session import resolve_server
|
|
from artemis.plugins.store import PluginStore
|
|
from artemis.runtime.device_lock import DeviceBusyError, DeviceExecutionLock
|
|
|
|
|
|
def manifest(**updates):
|
|
return PluginManifest.model_validate(
|
|
{
|
|
"id": "test-plugin",
|
|
"name": "测试插件",
|
|
"version": "1.0.0",
|
|
"platforms": [sys.platform],
|
|
"skill": "SKILL.md",
|
|
"resources": ["desktop:local"],
|
|
"config": {"token": {"label": "密钥", "secret": True, "required": True}},
|
|
**updates,
|
|
}
|
|
)
|
|
|
|
|
|
def archive(plugin, extra=None):
|
|
buffer = io.BytesIO()
|
|
with zipfile.ZipFile(buffer, "w") as output:
|
|
output.writestr("plugin.json", plugin.model_dump_json())
|
|
output.writestr("SKILL.md", "先观察，再操作；报告真实证据。")
|
|
for name, value in (extra or {}).items():
|
|
output.writestr(name, value)
|
|
return buffer.getvalue()
|
|
|
|
|
|
@pytest.fixture
|
|
def store(tmp_path):
|
|
return PluginStore(tmp_path / "plugins")
|
|
|
|
|
|
def enable(store, plugin, **profile):
|
|
install_archive(store, archive(plugin))
|
|
store.save_profile(
|
|
plugin.id, PluginProfile(secrets={"token": "private-value"}, authorized=True, **profile)
|
|
)
|
|
store.set_enabled(plugin.id, True)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_android_task_discovers_cross_device_plugins_and_loads_only_on_demand(store):
|
|
plugin = manifest(targets=["computer", "mixed"])
|
|
enable(store, plugin)
|
|
runtime = PluginRuntime(store=store, target="android")
|
|
try:
|
|
assert list(runtime.packages) == [plugin.id]
|
|
assert runtime.loaded == set()
|
|
assert runtime.sessions == {}
|
|
assert plugin.id in runtime.prompt()
|
|
assert (await runtime.call("plugin_load", {"plugin_id": plugin.id})).status == "success"
|
|
assert runtime.loaded == {plugin.id}
|
|
finally:
|
|
await runtime.aclose()
|
|
|
|
|
|
def test_android_auto_selection_respects_enablement_authorization_and_explicit_opt_out(store):
|
|
enable(store, manifest(targets=["mixed"]))
|
|
assert len(store.pin(None, "android")) == 1
|
|
assert store.pin([], "android") == []
|
|
store.set_enabled("test-plugin", False)
|
|
assert store.pin(None, "android") == []
|
|
with pytest.raises(ValueError, match="未启用"):
|
|
store.pin([PluginSelection(id="test-plugin")], "android")
|
|
store.set_enabled("test-plugin", True)
|
|
store.save_profile("test-plugin", PluginProfile(authorized=False))
|
|
assert store.pin(None, "android") == []
|
|
with pytest.raises(ValueError, match="授权"):
|
|
store.pin([PluginSelection(id="test-plugin")], "android")
|
|
|
|
|
|
def test_android_task_uses_authorized_default_and_allows_explicit_profile(store):
|
|
enable(store, manifest(targets=["mixed"]))
|
|
store.save_profile(
|
|
"test-plugin",
|
|
PluginProfile(name="lab", is_default=True, authorized=True, secrets={"token": "lab"}),
|
|
)
|
|
assert store.pin(None, "android")[0].profile == "lab"
|
|
assert store.pin([PluginSelection(id="test-plugin")], "android")[0].profile == "default"
|
|
|
|
|
|
def test_computer_only_plugin_is_not_automatically_offered_to_android(store):
|
|
enable(store, manifest(targets=["computer"]))
|
|
assert store.pin(None, "android") == []
|
|
with pytest.raises(ValueError, match="不支持"):
|
|
store.pin([PluginSelection(id="test-plugin")], "android")
|
|
|
|
|
|
def test_profile_secrets_are_write_only_and_blank_preserves(store):
|
|
enable(store, manifest())
|
|
assert "private-value" not in json.dumps(store.list())
|
|
assert store.list()[0]["profiles"][0]["configured_secrets"] == ["token"]
|
|
store.save_profile("test-plugin", PluginProfile(secrets={"token": ""}, authorized=True))
|
|
_, profile, _ = store.resolve(PluginSelection(id="test-plugin"))
|
|
assert profile.secrets["token"].get_secret_value() == "private-value"
|
|
|
|
|
|
def test_untrusted_config_cannot_hide_secrets_in_ordinary_fields(store):
|
|
install_archive(store, archive(manifest()))
|
|
with pytest.raises(ValueError, match="普通配置"):
|
|
store.save_profile("test-plugin", PluginProfile(values={"token": "secret"}))
|
|
with pytest.raises(ValueError, match="授权"):
|
|
store.set_enabled("test-plugin", True)
|
|
|
|
|
|
def test_package_versions_immutable_and_old_pin_survives_upgrade(store):
|
|
enable(store, manifest())
|
|
with pytest.raises(ValueError, match="覆盖"):
|
|
install_archive(store, archive(manifest()))
|
|
install_archive(store, archive(manifest(version="1.1.0")))
|
|
assert store.get("test-plugin")["enabled"] is False
|
|
store.set_enabled("test-plugin", True)
|
|
old, _, _ = store.resolve(PluginSelection(id="test-plugin", version="1.0.0"))
|
|
assert old.version == "1.0.0"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"path",
|
|
["../../escape.txt", "C:/escape.txt", "/escape.txt", "dir/../../../escape", "x/../escape"],
|
|
)
|
|
def test_archive_rejects_traversal(store, path):
|
|
with pytest.raises(ValueError):
|
|
install_archive(store, archive(manifest(), {path: "bad"}))
|
|
assert store.list() == []
|
|
|
|
|
|
def test_checksum_failure_keeps_previous_version(store):
|
|
enable(store, manifest())
|
|
entry = CatalogEntry(
|
|
manifest=manifest(version="1.1.0"),
|
|
archive_url="https://example.com/package.zip",
|
|
sha256="0" * 64,
|
|
verification="本地确定性验证",
|
|
)
|
|
with pytest.raises(ValueError, match="SHA-256"):
|
|
install_archive(store, archive(entry.manifest), entry=entry)
|
|
assert store.get("test-plugin")["manifest"]["version"] == "1.0.0"
|
|
assert store.get("test-plugin")["enabled"] is True
|
|
|
|
|
|
def test_server_substitutions_do_not_use_shell_or_expand_unknown_values(tmp_path):
|
|
plugin = manifest(
|
|
server={
|
|
"command": "${plugin_dir}/server.exe",
|
|
"args": ["${config.target}"],
|
|
"env": {"TOKEN": "${secret.token}"},
|
|
},
|
|
config={"token": {"label": "密钥", "secret": True}, "target": {"label": "目标"}},
|
|
)
|
|
profile = PluginProfile(
|
|
values={"target": "a b;$(echo forbidden)"}, secrets={"token": "private-value"}
|
|
)
|
|
config = resolve_server(plugin, profile, tmp_path)
|
|
assert config.args == ["a b;$(echo forbidden)"]
|
|
assert config.env == {"TOKEN": "private-value"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_skill_load_and_reference_cannot_read_outside_package(store):
|
|
enable(store, manifest())
|
|
runtime = PluginRuntime(store=store)
|
|
assert {t.name for t in runtime.declarations()} == {"plugin_load", "plugin_read_file"}
|
|
result = await runtime.call("plugin_load", {"plugin_id": "test-plugin"})
|
|
assert result.status == "success" and "先观察" in result.text()
|
|
assert (
|
|
await runtime.call(
|
|
"plugin_read_file", {"plugin_id": "test-plugin", "path": "../../../registry.sqlite3"}
|
|
)
|
|
).status == "error"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_real_stdio_mcp_errors_images_secrets_and_langchain(store, tmp_path):
|
|
server = tmp_path / "server.py"
|
|
server.write_text(
|
|
"""
|
|
import os
|
|
from mcp.server.fastmcp import FastMCP
|
|
from mcp.types import CallToolResult, TextContent, ImageContent
|
|
server = FastMCP("test")
|
|
@server.tool()
|
|
def observe(app: str) -> CallToolResult:
|
|
return CallToolResult(content=[TextContent(type="text", text=os.environ["TOKEN"]),
|
|
ImageContent(type="image", mimeType="image/png", data="aGVsbG8=")],
|
|
structuredContent={"app": app}, isError=app == "failed")
|
|
server.run()
|
|
""",
|
|
encoding="utf-8",
|
|
)
|
|
plugin = manifest(
|
|
server={
|
|
"command": sys.executable,
|
|
"args": [str(server)],
|
|
"env": {"TOKEN": "${secret.token}"},
|
|
}
|
|
)
|
|
enable(store, plugin, argument_constraints={"observe": {"app": ["ok", "failed"]}})
|
|
runtime = PluginRuntime(store=store)
|
|
try:
|
|
loaded = await runtime.load(plugin.id)
|
|
assert loaded.status == "success"
|
|
tool_name = next(iter(runtime.bindings))
|
|
result = await runtime.execute(tool_name, {"app": "ok"}, "tc-1")
|
|
assert result.status == "success"
|
|
assert json.loads(json.dumps(result.metadata))["resource_ids"] == ["desktop:local"]
|
|
assert result.screenshot_bytes == b"hello"
|
|
assert "private-value" not in result.text_summary
|
|
assert "[已隐藏密钥]" in result.text_summary
|
|
assert any(b["type"] == "image_url" for b in result.raw_result)
|
|
failed = await runtime.call(tool_name, {"app": "failed"})
|
|
assert failed.status == "error"
|
|
assert failed.structured_content == {"app": "failed"}
|
|
blocked = await runtime.call(tool_name, {"app": "not-authorized"})
|
|
assert blocked.status == "error"
|
|
wrapper = next(t for t in runtime.langchain_tools() if t.name == tool_name)
|
|
assert (await wrapper.ainvoke({"app": "failed"})).status == "error"
|
|
finally:
|
|
await runtime.aclose()
|
|
assert runtime.sessions == {}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stdio_timeout_is_bounded_and_never_replays(store, tmp_path, monkeypatch):
|
|
from artemis.config.settings import settings
|
|
|
|
server = tmp_path / "slow.py"
|
|
count = tmp_path / "calls.txt"
|
|
server.write_text(
|
|
f"""
|
|
import asyncio
|
|
from pathlib import Path
|
|
from mcp.server.fastmcp import FastMCP
|
|
server = FastMCP("slow")
|
|
@server.tool()
|
|
async def write_once() -> str:
|
|
with Path({str(count)!r}).open("a") as f: f.write("called\\n")
|
|
await asyncio.sleep(60)
|
|
return "done"
|
|
server.run()
|
|
""",
|
|
encoding="utf-8",
|
|
)
|
|
enable(store, manifest(server={"command": sys.executable, "args": [str(server)]}))
|
|
runtime = PluginRuntime(store=store)
|
|
await runtime.load("test-plugin")
|
|
monkeypatch.setattr(settings, "ARTEMIS_PLUGIN_TOOL_TIMEOUT_SECONDS", 0.2)
|
|
try:
|
|
name = next(iter(runtime.bindings))
|
|
result = await asyncio.wait_for(runtime.call(name, {}), timeout=10)
|
|
assert result.status == "uncertain"
|
|
assert count.read_text().splitlines() == ["called"]
|
|
again = await runtime.call(name, {})
|
|
assert again.status == "uncertain"
|
|
assert count.read_text().splitlines() == ["called"]
|
|
finally:
|
|
await runtime.aclose()
|
|
|
|
|
|
@pytest.mark.parametrize("scope", ["", "test-adb-endpoint"])
|
|
def test_resource_lock_conflicts_across_devices_without_holding_partial_leases(
|
|
tmp_path, monkeypatch, scope
|
|
):
|
|
import artemis.runtime.device_lock as locks
|
|
|
|
monkeypatch.setattr(locks, "get_temp_dir", lambda *_: tmp_path)
|
|
monkeypatch.delenv(DeviceExecutionLock.QUEUE_TICKET_ENV, raising=False)
|
|
monkeypatch.setenv(DeviceExecutionLock.LOCK_SCOPE_ENV, scope)
|
|
first = DeviceExecutionLock(
|
|
"phone-a", session_id="first", resource_ids=["desktop:local"], concurrency_mode="per_device"
|
|
)
|
|
second = DeviceExecutionLock(
|
|
"phone-b",
|
|
session_id="second",
|
|
resource_ids=["desktop:local", "router:lab"],
|
|
concurrency_mode="per_device",
|
|
)
|
|
other = DeviceExecutionLock(
|
|
"phone-c", session_id="other", resource_ids=["router:lab"], concurrency_mode="per_device"
|
|
)
|
|
first.acquire(blocking=False)
|
|
try:
|
|
with pytest.raises(DeviceBusyError):
|
|
second.acquire(blocking=False)
|
|
other.acquire(blocking=False)
|
|
assert len(DeviceExecutionLock.get_active_owners()) == 2
|
|
assert DeviceExecutionLock.get_active_owners()["phone-a"].resource_ids == ("desktop:local",)
|
|
assert DeviceExecutionLock.annotate_active_owner(
|
|
session_id="first", device_id="phone-a", ingress="test"
|
|
)
|
|
assert DeviceExecutionLock.get_active_owners()["phone-a"].resource_ids == ("desktop:local",)
|
|
finally:
|
|
other.release()
|
|
first.release()
|
|
second.acquire(blocking=False)
|
|
second.release()
|
|
assert not DeviceExecutionLock.get_active_owners()
|
|
|
|
|
|
def test_task_reservation_blocks_mutations_and_removes_files_after_release(store):
|
|
enable(store, manifest())
|
|
selection = PluginSelection(id="test-plugin", version="1.0.0")
|
|
assert store.reserve("queued-task", [selection], "mixed") == ["desktop:local"]
|
|
for change in (
|
|
lambda: store.save_profile("test-plugin", PluginProfile()),
|
|
lambda: store.set_enabled("test-plugin", False),
|
|
lambda: store.uninstall("test-plugin"),
|
|
lambda: install_archive(store, archive(manifest(version="1.1.0"))),
|
|
):
|
|
with pytest.raises(ValueError, match="任务"):
|
|
change()
|
|
assert not store.package_dir("test-plugin", "1.1.0").exists()
|
|
store.release("queued-task")
|
|
store.uninstall("test-plugin")
|
|
assert store.list() == []
|
|
assert not store.package_dir("test-plugin", "1.0.0").exists()
|
|
|
|
|
|
@pytest.mark.parametrize("stop_all", [False, True])
|
|
def test_cancelled_queue_entries_release_plugin_pin_but_live_workers_keep_it(
|
|
store, monkeypatch, stop_all
|
|
):
|
|
from types import SimpleNamespace
|
|
from apps.admin_console.services import task_queue_service as queue_module
|
|
from apps.admin_console.services import automation_store
|
|
|
|
enable(store, manifest())
|
|
selection = PluginSelection(id="test-plugin", version="1.0.0")
|
|
store.reserve("pending", [selection], "mixed")
|
|
store.reserve("active", [selection], "mixed")
|
|
pending = {"session_id": "pending", "status": "pending"}
|
|
active = {"session_id": "active", "status": "running"}
|
|
state = SimpleNamespace(
|
|
queue_items=[pending, active],
|
|
active_connections={},
|
|
active_session_id=None,
|
|
current_goal=None,
|
|
current_profile=None,
|
|
wake_event=SimpleNamespace(set=lambda: None),
|
|
)
|
|
state.clear_queue = lambda: state.queue_items.clear()
|
|
monkeypatch.setattr(queue_module, "state", state)
|
|
monkeypatch.setattr("artemis.plugins.store.PluginStore", lambda: store)
|
|
monkeypatch.setattr(DeviceExecutionLock, "cancel_reservation", lambda *_: None)
|
|
monkeypatch.setattr(
|
|
automation_store, "store", SimpleNamespace(records=lambda *_: [], job=lambda *_: None)
|
|
)
|
|
service = queue_module.TaskQueueService
|
|
for name in (
|
|
"_terminate_all_device_owners",
|
|
"_kill_all_local_runs",
|
|
"_clear_pause_file",
|
|
"ensure_worker_running",
|
|
):
|
|
monkeypatch.setattr(service, name, lambda: None)
|
|
if stop_all:
|
|
service._stop_all_tasks()
|
|
else:
|
|
service._remove_stopped_queue_item("pending", False)
|
|
service._remove_stopped_queue_item("active", False)
|
|
with store.connect() as db:
|
|
assert [row[0] for row in db.execute("SELECT task_id FROM task_pins")] == ["active"]
|
|
# Even after stop removes the queue row, the worker finalizer must drop its reservation.
|
|
service._remove_task("active")
|
|
store.assert_idle("test-plugin")
|
|
|
|
|
|
def test_registry_failure_rolls_back_package_files(store, monkeypatch):
|
|
def failed(*_):
|
|
raise OSError("disk failure")
|
|
|
|
monkeypatch.setattr(store, "register", failed)
|
|
with pytest.raises(OSError):
|
|
install_archive(store, archive(manifest()))
|
|
assert not store.package_dir("test-plugin", "1.0.0").exists()
|
|
assert store.list() == []
|
|
|
|
|
|
@pytest.mark.parametrize("path", ["NUL", "dir/COM1.txt", "aux.md", "LPT9/data"])
|
|
def test_archive_rejects_windows_device_names(store, path):
|
|
with pytest.raises(ValueError, match="路径"):
|
|
install_archive(store, archive(manifest(), {path: "bad"}))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_runtime_keeps_version_and_profile_reserved_while_paused(store):
|
|
enable(store, manifest())
|
|
runtime = PluginRuntime(store=store)
|
|
await runtime.aclose(release=False)
|
|
with pytest.raises(ValueError, match="任务"):
|
|
store.save_profile("test-plugin", PluginProfile())
|
|
await runtime.aclose()
|
|
store.save_profile("test-plugin", PluginProfile())
|
|
|
|
|
|
def test_plugin_http_api_redacts_invalid_payload_and_requires_resource_confirmation(
|
|
store, monkeypatch
|
|
):
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
from apps.admin_console.routers import plugins
|
|
|
|
enable(store, manifest())
|
|
monkeypatch.setattr(plugins, "get_store", lambda: store)
|
|
monkeypatch.setattr(DeviceExecutionLock, "get_active_owners", lambda: {})
|
|
app = FastAPI()
|
|
app.include_router(plugins.router)
|
|
with TestClient(app) as client:
|
|
response = client.put(
|
|
"/api/plugins/test-plugin/profiles",
|
|
json={
|
|
"name": "default",
|
|
"secrets": {"token": {"invalid": "DO-NOT-LEAK"}},
|
|
},
|
|
)
|
|
assert response.status_code == 422
|
|
assert "DO-NOT-LEAK" not in response.text
|
|
assert "private-value" not in client.get("/api/plugins").text
|
|
store.block_resources(["desktop:local"], "动作未确认停止")
|
|
assert client.get("/api/plugins").json()["blocked_resources"]
|
|
assert (
|
|
client.post(
|
|
"/api/plugins/resources/confirm-stopped", json={"resource": "desktop:local"}
|
|
).status_code
|
|
== 422
|
|
)
|
|
assert store.blocked_resources()
|
|
assert (
|
|
client.post(
|
|
"/api/plugins/resources/confirm-stopped",
|
|
json={"resource": "desktop:local", "confirmed_stopped": True},
|
|
).status_code
|
|
== 200
|
|
)
|
|
assert store.blocked_resources() == {}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pause_cleanup_error_does_not_release_resource(tmp_path, monkeypatch):
|
|
from unittest.mock import Mock
|
|
from artemis.runtime import task_control
|
|
|
|
monkeypatch.setattr(task_control, "TRACES_PATH", tmp_path)
|
|
lock = Mock()
|
|
control = task_control.TaskControl("cleanup-failed", lock)
|
|
entered = asyncio.Event()
|
|
|
|
async def operation():
|
|
entered.set()
|
|
try:
|
|
await asyncio.sleep(30)
|
|
except asyncio.CancelledError:
|
|
raise RuntimeError("still running remotely") from None
|
|
|
|
future = asyncio.create_task(control.run_interruptibly(operation()))
|
|
await entered.wait()
|
|
task_control.request_pause(control.session_id)
|
|
with pytest.raises(RuntimeError, match="still running"):
|
|
await future
|
|
lock.release.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stdio_pause_yields_resources_and_resume_waits_for_other_task(
|
|
store, tmp_path, monkeypatch
|
|
):
|
|
from artemis.runtime import task_control
|
|
import artemis.runtime.device_lock as locks
|
|
|
|
monkeypatch.setattr(task_control, "TRACES_PATH", tmp_path / "traces")
|
|
monkeypatch.setattr(locks, "get_temp_dir", lambda *_: tmp_path)
|
|
monkeypatch.delenv(DeviceExecutionLock.QUEUE_TICKET_ENV, raising=False)
|
|
started = tmp_path / "started.txt"
|
|
server = tmp_path / "pause_server.py"
|
|
server.write_text(
|
|
f"""
|
|
import asyncio
|
|
from pathlib import Path
|
|
from mcp.server.fastmcp import FastMCP
|
|
server = FastMCP("pause")
|
|
@server.tool()
|
|
async def observe(slow: bool = False) -> str:
|
|
if slow:
|
|
Path({str(started)!r}).write_text("started")
|
|
await asyncio.sleep(60)
|
|
return "fresh observation"
|
|
server.run()
|
|
""",
|
|
encoding="utf-8",
|
|
)
|
|
enable(store, manifest(server={"command": sys.executable, "args": [str(server)]}))
|
|
runtime = PluginRuntime(store=store)
|
|
first = DeviceExecutionLock(
|
|
"computer-a",
|
|
session_id="paused-plugin",
|
|
resource_ids=runtime.resources,
|
|
concurrency_mode="per_device",
|
|
)
|
|
other = DeviceExecutionLock(
|
|
"computer-b",
|
|
session_id="other-task",
|
|
resource_ids=runtime.resources,
|
|
concurrency_mode="per_device",
|
|
)
|
|
first.acquire(blocking=False)
|
|
control = task_control.TaskControl("paused-plugin", first)
|
|
future = None
|
|
|
|
async def until(predicate):
|
|
async with asyncio.timeout(12):
|
|
while not predicate():
|
|
await asyncio.sleep(0.05)
|
|
|
|
async def park():
|
|
await runtime.aclose(release=False)
|
|
|
|
async def resume():
|
|
await runtime.load("test-plugin")
|
|
|
|
try:
|
|
await runtime.load("test-plugin")
|
|
name = next(iter(runtime.bindings))
|
|
future = asyncio.create_task(
|
|
control.run_interruptibly(
|
|
runtime.call(name, {"slow": True}), on_pause=park, on_resume=resume
|
|
)
|
|
)
|
|
await until(started.exists)
|
|
task_control.request_pause(control.session_id)
|
|
await until(
|
|
lambda: task_control.read_state(control.session_id).get("status") == "task_paused"
|
|
)
|
|
assert not runtime.sessions
|
|
with pytest.raises(ValueError, match="任务"):
|
|
store.assert_idle("test-plugin")
|
|
other.acquire(blocking=False)
|
|
task_control.request_resume(control.session_id)
|
|
await until(lambda: task_control.read_state(control.session_id).get("status") == "resuming")
|
|
assert not future.done()
|
|
other.release()
|
|
assert await asyncio.wait_for(future, timeout=15) == (True, None)
|
|
assert DeviceExecutionLock.get_active_owners()["computer-a"].resource_ids == (
|
|
"desktop:local",
|
|
)
|
|
result = await runtime.call(next(iter(runtime.bindings)), {})
|
|
assert result.status == "success"
|
|
assert result.structured_content == {"result": "fresh observation"}
|
|
finally:
|
|
if future and not future.done():
|
|
future.cancel()
|
|
await asyncio.gather(future, return_exceptions=True)
|
|
await runtime.aclose()
|
|
control.release()
|
|
other.release()
|
|
store.assert_idle("test-plugin")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_computer_observation_and_both_action_pipelines_without_adb(
|
|
store, tmp_path, monkeypatch
|
|
):
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
from PIL import Image
|
|
from artemis.context import ArtemisContext, DeviceContext, DevicePlatform
|
|
from artemis.drivers.factory import create_driver
|
|
from artemis.mcp.action_executor import McpActionExecutor
|
|
from artemis.agents.validator.execution_loop import _process_action
|
|
from artemis.agents.checker.checker import build_checker_tools
|
|
|
|
picture = io.BytesIO()
|
|
Image.new("RGB", (80, 60), "white").save(picture, format="PNG")
|
|
encoded = base64.b64encode(picture.getvalue()).decode()
|
|
server = tmp_path / "desktop.py"
|
|
server.write_text(
|
|
f"""
|
|
from mcp.server.fastmcp import FastMCP
|
|
from mcp.types import CallToolResult, TextContent, ImageContent
|
|
server = FastMCP("desktop")
|
|
value = "initial"
|
|
@server.tool()
|
|
def list_apps() -> str:
|
|
return "fixture-window"
|
|
@server.tool()
|
|
def get_app_state(app: str) -> CallToolResult:
|
|
return CallToolResult(content=[TextContent(type="text", text=app + ":" + value), ImageContent(type="image", mimeType="image/png", data={encoded!r})])
|
|
@server.tool()
|
|
def set_value(text: str) -> str:
|
|
global value
|
|
value = text
|
|
return value
|
|
server.run()
|
|
""",
|
|
encoding="utf-8",
|
|
)
|
|
enable(
|
|
store,
|
|
manifest(
|
|
targets=["computer", "mixed"],
|
|
observation={"tool": "list_apps", "refresh_tools": ["get_app_state"]},
|
|
server={"command": sys.executable, "args": [str(server)]},
|
|
),
|
|
)
|
|
ctx = ArtemisContext(
|
|
device=DeviceContext(mobile_platform=DevicePlatform.COMPUTER, device_id="computer:local"),
|
|
execution_target="computer",
|
|
)
|
|
ctx.plugins = PluginRuntime(store=store, ctx=ctx)
|
|
adb = Mock(side_effect=AssertionError("computer tasks cannot create an ADB client"))
|
|
monkeypatch.setattr("artemis.drivers.factory.AdbClient", adb)
|
|
try:
|
|
driver = create_driver(ctx)
|
|
first = await driver.get_screen_data()
|
|
assert first.screenshot_bytes == b"" and "fixture-window" in first.observation_text
|
|
names = {tool.name: name for name, (_, tool) in ctx.plugins.bindings.items()}
|
|
await ctx.plugins.call(names["get_app_state"], {"app": "fixture-window"})
|
|
executor = McpActionExecutor(ctx)
|
|
result = await executor.execute(
|
|
names["set_value"], {"text": "flash-value"}, "flash-call", None
|
|
)
|
|
assert result.status == "success" and result.metadata["plugin_id"] == "test-plugin"
|
|
observed = await driver.get_screen_data()
|
|
assert observed.width == 80 and "flash-value" in observed.observation_text
|
|
node = SimpleNamespace(ctx=ctx)
|
|
outcome = await _process_action(
|
|
node,
|
|
None,
|
|
None,
|
|
{"action": names["set_value"], "args": {"text": "pro-value"}},
|
|
names["set_value"],
|
|
"",
|
|
burst=False,
|
|
)
|
|
assert outcome.success and outcome.action_item["plugin_version"] == "1.0.0"
|
|
observed = await driver.get_screen_data()
|
|
assert "pro-value" in observed.observation_text
|
|
assert "probe_device" not in {tool.name for tool in build_checker_tools(ctx, "final")}
|
|
adb.assert_not_called()
|
|
finally:
|
|
await ctx.plugins.aclose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("target", ["computer", "mixed"])
|
|
async def test_flash_plugin_images_never_replace_the_mixed_phone_observation(monkeypatch, target):
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock, AsyncMock
|
|
from artemis.agents.flash import runner as flash
|
|
from artemis.agents.validator.tool_declarations import ToolExecutionResult
|
|
|
|
name = "plug_test_observe"
|
|
plugins = SimpleNamespace(
|
|
bindings={name: ("test-plugin", None)},
|
|
tool_names={name},
|
|
last_images={"test-plugin": b"desktop-before"},
|
|
)
|
|
monkeypatch.setattr("artemis.plugins.runtime.get_plugin_runtime", lambda ctx: plugins)
|
|
capture = AsyncMock(return_value=("phone-path", b"phone-after", "phone-tree-after"))
|
|
monkeypatch.setattr(flash, "capture_screenshot_and_parse_ui", capture)
|
|
runner = flash.FlashRunner.__new__(flash.FlashRunner)
|
|
runner.ctx = SimpleNamespace(execution_target=target, data_engine=Mock(), task_control=None)
|
|
runner.controller, runner.summarizer, runner._llm = Mock(), None, None
|
|
runner.executor = SimpleNamespace(
|
|
action_tool_names={name},
|
|
execute=AsyncMock(
|
|
return_value=ToolExecutionResult(
|
|
tool_name=name,
|
|
tool_call_id="tc",
|
|
status="success",
|
|
screenshot_bytes=b"desktop-after",
|
|
ui_elements_text="desktop-tree",
|
|
text_summary="Desktop observation",
|
|
metadata={"plugin_id": "test-plugin"},
|
|
)
|
|
),
|
|
)
|
|
runner._build_action_record = Mock(return_value={"action": name})
|
|
runner._record_action_step = Mock(return_value="step")
|
|
state = SimpleNamespace(latest_screenshot="old-phone-path")
|
|
result = await runner._execute_and_record_action(
|
|
name,
|
|
{},
|
|
"tc",
|
|
state,
|
|
[],
|
|
"",
|
|
{},
|
|
b"phone-before",
|
|
"phone-tree-before",
|
|
0,
|
|
flash._TurnRecord(),
|
|
)
|
|
assert runner._record_action_step.call_args.args[5:8] == (
|
|
b"desktop-before",
|
|
"",
|
|
b"desktop-after",
|
|
)
|
|
if target == "mixed":
|
|
assert result[:2] == (b"phone-after", "phone-tree-after")
|
|
assert state.latest_screenshot == "phone-path"
|
|
capture.assert_awaited_once()
|
|
else:
|
|
assert result[:2] == (b"desktop-after", "desktop-tree")
|
|
capture.assert_not_awaited()
|
|
|
|
|
|
def test_named_default_profile_is_used_only_when_authorized(store):
|
|
enable(store, manifest())
|
|
store.save_profile(
|
|
"test-plugin",
|
|
PluginProfile(
|
|
name="lab", is_default=True, authorized=True, secrets={"token": "lab-secret"}
|
|
),
|
|
)
|
|
assert store.selections()[0].profile == "lab"
|
|
store.save_profile("test-plugin", PluginProfile(name="lab", is_default=True, authorized=False))
|
|
assert store.selections() == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streamable_http_timeout_protects_resources_and_never_replays(
|
|
store, tmp_path, monkeypatch
|
|
):
|
|
import socket
|
|
import subprocess
|
|
from artemis.config.settings import settings
|
|
|
|
with socket.socket() as available:
|
|
available.bind(("127.0.0.1", 0))
|
|
port = available.getsockname()[1]
|
|
counter = tmp_path / "writes.txt"
|
|
script = tmp_path / "http_mcp.py"
|
|
script.write_text(
|
|
f"""
|
|
import asyncio
|
|
from pathlib import Path
|
|
from mcp.server.fastmcp import FastMCP
|
|
server = FastMCP("http-test", host="127.0.0.1", port={port})
|
|
@server.tool()
|
|
async def write_once() -> str:
|
|
with Path({str(counter)!r}).open("a") as output: output.write("write\\n")
|
|
await asyncio.sleep(30)
|
|
return "written"
|
|
server.run(transport="streamable-http")
|
|
""",
|
|
encoding="utf-8",
|
|
)
|
|
process = subprocess.Popen(
|
|
[sys.executable, str(script)],
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL,
|
|
creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
|
|
)
|
|
runtime = None
|
|
try:
|
|
for _ in range(100):
|
|
with socket.socket() as probe:
|
|
if probe.connect_ex(("127.0.0.1", port)) == 0:
|
|
break
|
|
assert process.poll() is None
|
|
await asyncio.sleep(0.05)
|
|
else:
|
|
pytest.fail("local MCP fixture did not start")
|
|
enable(
|
|
store,
|
|
manifest(
|
|
server={"transport": "streamable_http", "url": f"http://127.0.0.1:{port}/mcp"}
|
|
),
|
|
)
|
|
runtime = PluginRuntime(store=store)
|
|
await runtime.load("test-plugin")
|
|
monkeypatch.setattr(settings, "ARTEMIS_PLUGIN_TOOL_TIMEOUT_SECONDS", 0.2)
|
|
name = next(iter(runtime.bindings))
|
|
result = await asyncio.wait_for(runtime.call(name, {}), timeout=10)
|
|
assert result.status == "uncertain"
|
|
assert "desktop:local" in store.blocked_resources()
|
|
assert counter.read_text().splitlines() == ["write"]
|
|
with pytest.raises(RuntimeError, match="资源"):
|
|
await runtime.load("test-plugin")
|
|
assert counter.read_text().splitlines() == ["write"]
|
|
finally:
|
|
if runtime:
|
|
await runtime.aclose()
|
|
process.terminate()
|
|
process.wait(timeout=10)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_computer_only_plugin_diagnosis_uses_its_declared_target(store, monkeypatch):
|
|
from apps.admin_console.routers import plugins
|
|
|
|
enable(store, manifest(targets=["computer"], observation={"tool": "observe"}))
|
|
monkeypatch.setattr(plugins, "get_store", lambda: store)
|
|
result = await plugins.diagnose_plugin("test-plugin", plugins.DiagnoseRequest())
|
|
assert result["ok"] is True
|
|
store.assert_idle("test-plugin")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_workspace_computer_discovery_never_queries_adb(monkeypatch):
|
|
from mcp_server.tools import automation
|
|
|
|
routes = []
|
|
|
|
async def request(method, route):
|
|
routes.append(route)
|
|
assert route != "/api/devices"
|
|
return {"plugins": [{"id": "computer"}]} if route == "/api/plugins" else {}
|
|
|
|
monkeypatch.setattr(automation, "_request", request)
|
|
result = await automation.mobile_workspace(include_devices=False)
|
|
assert result["devices"] == []
|
|
assert result["plugins"] == [{"id": "computer"}]
|
|
assert set(routes) == {"/api/status", "/api/plugins"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sdk_setup_failure_releases_plugin_version_pin(store, monkeypatch):
|
|
from artemis.sdk.agent import Agent
|
|
import artemis.plugins.runtime as runtime_module
|
|
|
|
enable(store, manifest(targets=["computer"], observation={"tool": "observe"}))
|
|
monkeypatch.setattr(runtime_module, "PluginStore", lambda: store)
|
|
monkeypatch.setattr(DeviceExecutionLock, "get_active_owner", lambda _: None)
|
|
|
|
def fail_before_tracing(*args, **kwargs):
|
|
raise ValueError("cannot initialize lock")
|
|
|
|
monkeypatch.setattr("artemis.sdk.agent.DeviceExecutionLock", fail_before_tracing)
|
|
agent = Agent()
|
|
await agent.init(execution_target="computer")
|
|
request = agent.new_task("fixture").with_plugins([PluginSelection(id="test-plugin")]).build()
|
|
# Fail while constructing the lock, before device/trace initialization.
|
|
fail_before_tracing.get_active_owner = lambda _: None
|
|
with pytest.raises(ValueError, match="initialize lock"):
|
|
await agent.run_task(request=request)
|
|
store.assert_idle("test-plugin")
|
|
|
|
|
|
def test_plugin_screenshots_export_with_source_and_without_phone_substitution(
|
|
tmp_path, monkeypatch
|
|
):
|
|
from apps.admin_console.services import report_export
|
|
|
|
phone, desktop = tmp_path / "phone.png", tmp_path / "desktop.png"
|
|
phone.write_bytes(b"phone")
|
|
desktop.write_bytes(b"desktop")
|
|
monkeypatch.setattr(
|
|
report_export, "image_path", lambda name: {"phone": phone, "desktop": desktop}.get(name)
|
|
)
|
|
evidence = {"plugin_id": "desktop-control", "image_names": ["desktop"]}
|
|
report = {
|
|
"steps": [
|
|
{
|
|
"step_number": 1,
|
|
"pre_image_name": "phone",
|
|
"generic_tools": [{"payload": evidence}, {"payload": json.dumps(evidence)}],
|
|
}
|
|
]
|
|
}
|
|
images = report_export.session_screenshots(report)
|
|
assert [path for path, _ in images] == [phone, desktop]
|
|
assert "插件_desktop-control" in images[1][1]
|
|
|
|
|
|
def test_plugin_images_live_until_the_last_referencing_task_is_deleted(tmp_path):
|
|
from uuid import uuid4
|
|
from artemis.data_engine.storage import StorageManager
|
|
|
|
storage = StorageManager(tmp_path / "history.db", tmp_path)
|
|
name = "f" * 64
|
|
image = tmp_path / "images" / f"{name}.jpg"
|
|
image.parent.mkdir()
|
|
image.write_bytes(b"plugin image")
|
|
first, second = uuid4(), uuid4()
|
|
storage.reference_plugin_image(first, name)
|
|
storage.reference_plugin_image(second, name)
|
|
storage.delete_session(first)
|
|
assert image.read_bytes() == b"plugin image"
|
|
storage.delete_session(second)
|
|
assert not image.exists()
|
|
|
|
|
|
def test_plugin_operation_export_retains_source_and_failure_evidence():
|
|
from apps.admin_console.services.report_export import plugin_operations
|
|
from apps.admin_console.database.repositories.step_repository import StepRepository
|
|
|
|
operations = plugin_operations(
|
|
[
|
|
{
|
|
"step_number": 3,
|
|
"generic_tools": [
|
|
{
|
|
"name": "plug_click",
|
|
"status": "failed",
|
|
"payload": {
|
|
"plugin_id": "desktop",
|
|
"plugin_version": "1.0.0",
|
|
"resource_ids": ["desktop:local"],
|
|
"error": "窗口已关闭",
|
|
"post_image_name": "evidence",
|
|
},
|
|
},
|
|
],
|
|
}
|
|
]
|
|
)
|
|
trace = StepRepository()._normalize_display_trace(
|
|
{
|
|
"name": "plug_click",
|
|
"type": "tool",
|
|
"status": "failed",
|
|
"payload": {
|
|
"plugin_id": "desktop",
|
|
"plugin_version": "1.0.0",
|
|
"resource_ids": ["desktop:local"],
|
|
"post_image_name": "evidence",
|
|
"image_names": ["evidence"],
|
|
"error": "窗口已关闭",
|
|
},
|
|
}
|
|
)
|
|
assert trace["payload"]["post_image_name"] == "evidence"
|
|
assert plugin_operations([{"step_number": 3, "generic_tools": [trace]}]) == operations
|
|
assert operations == [
|
|
{
|
|
"step": 3,
|
|
"plugin": "desktop",
|
|
"version": "1.0.0",
|
|
"tool": "plug_click",
|
|
"resources": "desktop:local",
|
|
"status": "failed",
|
|
"content": "窗口已关闭",
|
|
}
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pro_text_only_observation_never_sends_an_empty_image():
|
|
from artemis.agents.operator.prompts import ObservationPromptComponent, PromptBuilder
|
|
|
|
builder = PromptBuilder()
|
|
await ObservationPromptComponent()(
|
|
builder, None, None, latest_screenshot_b64="", minimal_list="fixture window"
|
|
)
|
|
assert all(
|
|
not isinstance(part, dict) or part.get("type") != "image_url"
|
|
for part in builder.human_parts
|
|
)
|
|
assert any("fixture window" in part for part in builder.human_parts if isinstance(part, str))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_plugin_image_text_stays_with_its_own_observation(store, tmp_path, monkeypatch):
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock
|
|
from PIL import Image
|
|
from artemis.context import ArtemisContext
|
|
from artemis.data_engine.engine import DataEngine
|
|
from artemis.plugins.models import PluginCallResult
|
|
|
|
monkeypatch.setattr(DataEngine, "_ipc_port_candidates", staticmethod(lambda: []))
|
|
monkeypatch.setattr("artemis.data_engine.engine.PAUSE_FILE", tmp_path / "pause")
|
|
ctx = MagicMock(spec=ArtemisContext)
|
|
ctx.execution_setup = SimpleNamespace(traces_path=str(tmp_path / "traces"))
|
|
ctx.device = None
|
|
engine = DataEngine(ctx)
|
|
engine.start_session("plugin image provenance")
|
|
runtime = PluginRuntime([], store=store, ctx=SimpleNamespace(data_engine=engine))
|
|
images = []
|
|
names = []
|
|
try:
|
|
for text, color in (("baseline", "white"), ("changed", "green")):
|
|
buffer = io.BytesIO()
|
|
Image.new("RGB", (40, 40), color).save(buffer, format="PNG")
|
|
images.append(buffer.getvalue())
|
|
outcome = PluginCallResult(
|
|
status="success",
|
|
content=[
|
|
{"type": "text", "text": text},
|
|
{"type": "image", "data": base64.b64encode(images[-1]).decode()},
|
|
],
|
|
)
|
|
runtime._persist_images(outcome)
|
|
names.extend(outcome.image_names)
|
|
engine.record_step(
|
|
pre_screenshot_bytes=images[0],
|
|
post_screenshot_bytes=images[1],
|
|
ui_tree="changed",
|
|
)
|
|
await engine.shutdown()
|
|
assert engine.storage.get_image(names[0]).ui_tree == "baseline"
|
|
assert engine.storage.get_image(names[1]).ui_tree == "changed"
|
|
finally:
|
|
await runtime.aclose()
|