115 lines
4.6 KiB
Python
115 lines
4.6 KiB
Python
"""§8.11 delegate 子循环的纯件测试。
|
|
|
|
只覆盖不依赖 litellm 的部分:FilteredExecutor(白名单过滤 / 递归防护)+ Session
|
|
persist=False(内存态不落 DB)。full loop 拦截 + run_delegate 路径要 import core.loop
|
|
(顶层 import litellm,本机导入卡死),不在此单测,靠生产 / 手验——同 salvage 拆分策略。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import unittest
|
|
from pathlib import Path
|
|
from uuid import uuid4
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
from core.delegate import DELEGATE_ALLOWED, FilteredExecutor # noqa: E402
|
|
from core.executor import ExecCtx, Executor, ToolResult # noqa: E402
|
|
|
|
|
|
class _FakeExecutor(Executor):
|
|
"""记录被调用的工具名;schemas 由构造时给的名字列表拼。"""
|
|
|
|
def __init__(self, names):
|
|
self._names = list(names)
|
|
self.calls = []
|
|
|
|
def has_tool(self, name):
|
|
return name in self._names
|
|
|
|
def schemas(self):
|
|
return [
|
|
{"type": "function", "function": {"name": n, "parameters": {}}}
|
|
for n in self._names
|
|
]
|
|
|
|
def call_tool(self, name, args, ctx):
|
|
self.calls.append(name)
|
|
return ToolResult(content=f"ran:{name}", exit_code=0)
|
|
|
|
|
|
def _ctx():
|
|
return ExecCtx(user_id=uuid4(), task_id=uuid4(), working_dir=Path("."))
|
|
|
|
|
|
class TestFilteredExecutor(unittest.TestCase):
|
|
def test_only_whitelisted_and_available_tools_exposed(self):
|
|
# 父有 read(白名单)、write(非白名单)、run_python(非白名单)、document_search(白名单)
|
|
inner = _FakeExecutor(["read", "write", "run_python", "document_search"])
|
|
fe = FilteredExecutor(inner, DELEGATE_ALLOWED)
|
|
|
|
exposed = {s["function"]["name"] for s in fe.schemas()}
|
|
self.assertEqual(exposed, {"read", "document_search"})
|
|
self.assertTrue(fe.has_tool("read"))
|
|
self.assertFalse(fe.has_tool("write"))
|
|
self.assertFalse(fe.has_tool("run_python"))
|
|
|
|
def test_whitelisted_call_passes_through(self):
|
|
inner = _FakeExecutor(["read", "document_search"])
|
|
fe = FilteredExecutor(inner, DELEGATE_ALLOWED)
|
|
out = fe.call_tool("read", {"path": "x"}, _ctx())
|
|
self.assertEqual(out.content, "ran:read")
|
|
self.assertEqual(inner.calls, ["read"])
|
|
|
|
def test_non_whitelisted_call_rejected_without_reaching_inner(self):
|
|
inner = _FakeExecutor(["read", "write", "shell"])
|
|
fe = FilteredExecutor(inner, DELEGATE_ALLOWED)
|
|
for banned in ("write", "shell"):
|
|
out = fe.call_tool(banned, {}, _ctx())
|
|
self.assertTrue(out.content.startswith("[Error]"), out.content)
|
|
self.assertEqual(out.exit_code, 2)
|
|
self.assertEqual(inner.calls, []) # 拒在 FilteredExecutor,没穿到父
|
|
|
|
def test_available_narrows_to_intersection(self):
|
|
# 白名单里有 mp_*,但父没注册(缺 key)→ 不暴露、不可调
|
|
inner = _FakeExecutor(["read"]) # 只有 read
|
|
fe = FilteredExecutor(inner, DELEGATE_ALLOWED)
|
|
self.assertFalse(fe.has_tool("mp_search_summary"))
|
|
exposed = {s["function"]["name"] for s in fe.schemas()}
|
|
self.assertEqual(exposed, {"read"})
|
|
|
|
|
|
class TestDelegateWhitelistInvariants(unittest.TestCase):
|
|
def test_no_write_or_exec_tools_in_whitelist(self):
|
|
# clause ②:不给 write/edit/run_python/shell
|
|
for banned in ("write", "edit", "run_python", "shell"):
|
|
self.assertNotIn(banned, DELEGATE_ALLOWED)
|
|
|
|
def test_delegate_not_in_whitelist_prevents_recursion(self):
|
|
# clause ①:子循环无法再 delegate(schema 不暴露)
|
|
self.assertNotIn("delegate", DELEGATE_ALLOWED)
|
|
|
|
def test_fetch_to_disk_retrieval_tools_allowed(self):
|
|
# 公测决策:给"落盘检索"工具写口(落的是抓取数据非模型内容)
|
|
for t in ("document_download", "mp_get_structure", "mp_get_entries"):
|
|
self.assertIn(t, DELEGATE_ALLOWED)
|
|
|
|
|
|
class TestSessionPersistFlag(unittest.TestCase):
|
|
def test_in_memory_session_appends_without_db(self):
|
|
# persist=False 的 Session.append 不碰 DB(不需要配置连接)
|
|
from core.session import Session
|
|
|
|
s = Session(task_id=uuid4(), system_prompt="SYS", persist=False)
|
|
self.assertEqual(len(s.messages), 1) # system
|
|
mid = s.append({"role": "user", "content": "hi"})
|
|
self.assertIsNone(mid) # 内存态无 message_id
|
|
s.append({"role": "assistant", "content": "yo"})
|
|
roles = [m["role"] for m in s.messages]
|
|
self.assertEqual(roles, ["system", "user", "assistant"])
|
|
self.assertEqual(s.n_user_msgs(), 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|