host
fd37505c48
Agent definitions gain tools (websearch/imagegen/playwright, unknown tool is a 400), stored as agent_tools JSON and returned by GET /tasks. This is deliberately required-tools, not an allowlist: create_thread has no tool allowlist and tenant homes are per-project shared, so blocking other tools cannot be claimed. Instead the worker fails fast before any model call when a declared tool has no [mcp_servers.<tool>] section in the tenant codex-home config.toml (checked after _tenant_dirs seeding, so a section seeded earlier by another worker counts as present). The gate runs only for tenant-isolated Docker workers; local/cloud workers skip it, documented, not verified. Tests: test_agent_tools.py (8 tests).
203 lines
8.0 KiB
Python
203 lines
8.0 KiB
Python
"""agent tools v3 单测：解析、落库、租户必需工具门。
|
|
|
|
运行（在仓库根目录下）：
|
|
python3 -m unittest discover -s cloud/tests -p "test_agent_tools.py" -v
|
|
|
|
说明：tools 是必需工具，不是白名单（见 agent_def.py 的 v3
|
|
边界说明：本代码库没有 tool allowlist 通道，worker 也不能
|
|
阻止 agent 用别的工具）。单测覆盖：未知工具 400、必需工具
|
|
缺失判定（纯函数 + 真实租户 config.toml）、建任务落库与
|
|
回读、agent_name 示例的 tools。run_one 的门接线是单条
|
|
if 调用 _missing_required_tools，完整 Docker 派发需真实
|
|
Docker 主机，不在单测里假装跑通。
|
|
"""
|
|
|
|
import http.client
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
try:
|
|
from cloud.agent_def import KNOWN_TOOLS, missing_required_tools
|
|
from cloud.agent_def import parse_agent_definition
|
|
from cloud.auth import AuthStore
|
|
from cloud.sandbox import DockerExecutor
|
|
from cloud.server import create_server, init_db
|
|
from cloud.worker import Worker
|
|
except ImportError:
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
|
from cloud.agent_def import KNOWN_TOOLS, missing_required_tools
|
|
from cloud.agent_def import parse_agent_definition
|
|
from cloud.auth import AuthStore
|
|
from cloud.sandbox import DockerExecutor
|
|
from cloud.server import create_server, init_db
|
|
from cloud.worker import Worker
|
|
|
|
|
|
def _agent(**overrides):
|
|
agent = {
|
|
"name": "research-brief",
|
|
"version": "v1",
|
|
"instructions": "Cite sources.",
|
|
"skills": ["live-info"],
|
|
"tools": ["websearch"],
|
|
"defaults": {"effort": "medium", "approval": "ask",
|
|
"executor": "docker"},
|
|
}
|
|
agent.update(overrides)
|
|
return agent
|
|
|
|
|
|
class ParseToolsTestCase(unittest.TestCase):
|
|
def test_known_tools_and_dedup(self):
|
|
definition = parse_agent_definition(
|
|
_agent(tools=["websearch", "websearch", "imagegen"]))
|
|
self.assertEqual(definition.tools, ["websearch", "imagegen"])
|
|
self.assertEqual(set(KNOWN_TOOLS),
|
|
{"websearch", "imagegen", "playwright"})
|
|
|
|
def test_unknown_tool_rejected(self):
|
|
with self.assertRaises(ValueError):
|
|
parse_agent_definition(_agent(tools=["shell"]))
|
|
|
|
def test_missing_required_tools_pure(self):
|
|
config = "[mcp_servers.websearch]\ncommand = \"x\"\n"
|
|
self.assertEqual(
|
|
missing_required_tools(["websearch", "imagegen"], config),
|
|
["imagegen"])
|
|
self.assertEqual(missing_required_tools(["websearch"], config), [])
|
|
self.assertEqual(
|
|
missing_required_tools(["playwright"], ""), ["playwright"])
|
|
|
|
|
|
class ServerToolsTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.db_path = os.path.join(self.tmp.name, "control.db")
|
|
self.auth = AuthStore(self.db_path)
|
|
self.owner_id, self.owner_key = self.auth.create_user("owner")
|
|
self.server = create_server(self.db_path, "127.0.0.1", 0)
|
|
self.port = self.server.server_address[1]
|
|
self.thread = threading.Thread(
|
|
target=self.server.serve_forever, daemon=True)
|
|
self.thread.start()
|
|
status, self.project = self._req(
|
|
"POST", "/projects", {"name": "demo"}, self.owner_key)
|
|
self.assertEqual(status, 201)
|
|
|
|
def tearDown(self):
|
|
self.server.shutdown()
|
|
self.server.server_close()
|
|
self.thread.join(timeout=5)
|
|
try:
|
|
self.auth.conn.close()
|
|
except Exception:
|
|
pass
|
|
self.tmp.cleanup()
|
|
|
|
def _req(self, method, path, body=None, key=None):
|
|
conn = http.client.HTTPConnection("127.0.0.1", self.port, timeout=10)
|
|
headers = {"Content-Type": "application/json"}
|
|
if key:
|
|
headers["X-API-Key"] = key
|
|
conn.request(method, path,
|
|
json.dumps(body) if body is not None else None, headers)
|
|
resp = conn.getresponse()
|
|
raw = resp.read()
|
|
conn.close()
|
|
try:
|
|
data = json.loads(raw)
|
|
except ValueError:
|
|
data = None
|
|
return resp.status, data
|
|
|
|
def _list_tasks(self):
|
|
status, data = self._req(
|
|
"GET", "/tasks?project_id=" + self.project["id"],
|
|
key=self.owner_key)
|
|
self.assertEqual(status, 200)
|
|
return data["tasks"]
|
|
|
|
def test_inline_agent_tools_stored(self):
|
|
status, task = self._req("POST", "/tasks", {
|
|
"project_id": self.project["id"], "prompt": "do research",
|
|
"agent": _agent()}, self.owner_key)
|
|
self.assertEqual(status, 201, task)
|
|
self.assertEqual(task["agent_tools"], ["websearch"])
|
|
self.assertEqual(task["agent"]["tools"], ["websearch"])
|
|
self.assertEqual(self._list_tasks()[0]["agent_tools"],
|
|
["websearch"])
|
|
|
|
def test_unknown_tool_rejected_no_task(self):
|
|
status, _data = self._req("POST", "/tasks", {
|
|
"project_id": self.project["id"], "prompt": "do research",
|
|
"agent": _agent(tools=["shell"])}, self.owner_key)
|
|
self.assertEqual(status, 400)
|
|
self.assertEqual(self._list_tasks(), [])
|
|
|
|
def test_agent_name_example_tools(self):
|
|
status, task = self._req("POST", "/tasks", {
|
|
"project_id": self.project["id"], "prompt": "do research",
|
|
"agent_name": "research-brief"}, self.owner_key)
|
|
self.assertEqual(status, 201, task)
|
|
self.assertEqual(task["agent_tools"], ["websearch"])
|
|
|
|
|
|
class WorkerToolGateTestCase(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.db_path = os.path.join(self.tmp.name, "control.db")
|
|
init_db(self.db_path)
|
|
self.tenants_root = os.path.join(self.tmp.name, "tenants")
|
|
self.template_home = os.path.join(self.tmp.name, "template")
|
|
os.makedirs(self.template_home, exist_ok=True)
|
|
with open(os.path.join(self.template_home, "config.toml"),
|
|
"w", encoding="utf-8") as fh:
|
|
fh.write("")
|
|
self.template_repo = os.path.join(self.tmp.name, "test-repo")
|
|
os.makedirs(self.template_repo, exist_ok=True)
|
|
with open(os.path.join(self.template_repo, "sample.py"),
|
|
"w", encoding="utf-8") as fh:
|
|
fh.write("")
|
|
executor = DockerExecutor.__new__(DockerExecutor)
|
|
self.worker = Worker.__new__(Worker)
|
|
self.worker.tenant_isolation = True
|
|
self.worker.executor = executor
|
|
self.worker.tenants_root = self.tenants_root
|
|
self.worker.tenant_template_home = self.template_home
|
|
self.worker.tenant_template_repo = self.template_repo
|
|
self.worker.websearch_mcp = False
|
|
self.worker.imagegen_mcp = False
|
|
self.worker.playwright_mcp = False
|
|
self.worker.seed_skills = False
|
|
|
|
def tearDown(self):
|
|
self.tmp.cleanup()
|
|
|
|
def test_gate_missing_and_present(self):
|
|
# 租户 home 初始只有模板 config.toml，无任何 MCP 段。
|
|
task = {"project_id": "p1", "agent_tools": ["websearch"]}
|
|
self.assertEqual(
|
|
self.worker._missing_required_tools(task), ["websearch"])
|
|
# 模拟之前别的 worker 已播过 websearch：段保留在共享
|
|
# 租户 home 里，本 worker 旗关也不误判。
|
|
home = self.worker._tenant_dirs("p1")["codex_home"]
|
|
with open(os.path.join(home, "config.toml"),
|
|
"a", encoding="utf-8") as fh:
|
|
fh.write("\n[mcp_servers.websearch]\ncommand = \"x\"\n")
|
|
self.assertEqual(
|
|
self.worker._missing_required_tools(task), [])
|
|
|
|
def test_gate_skipped_when_not_applicable(self):
|
|
self.worker.tenant_isolation = False
|
|
task = {"project_id": "p1", "agent_tools": ["websearch"]}
|
|
self.assertEqual(self.worker._missing_required_tools(task), [])
|
|
self.worker.tenant_isolation = True
|
|
self.assertEqual(
|
|
self.worker._missing_required_tools(
|
|
{"project_id": "p1", "agent_tools": []}), [])
|