projectAIpopular/tests/test_tools.py

"""工具内核测试：路径关押、文件工具往返、ToolLoop 循环编排（假 chat_fn，零外部依赖）。"""
import asyncio
import json

import pytest


def asyncio_run(coro):
    """与 test_pipeline.py 相同的事件循环包装（不依赖 pytest-asyncio 插件模式）。"""
    return asyncio.run(coro)

from router_system.tools import (
    ToolError,
    ToolLoop,
    WorkspaceTools,
    parse_tool_calls,
)


@pytest.fixture()
def ws(tmp_path):
    return WorkspaceTools(tmp_path / "ws")


# ---------------- WorkspaceTools ----------------

def test_jail_blocks_traversal(ws):
    with pytest.raises(ToolError):
        ws.resolve("../outside.txt")
    with pytest.raises(ToolError):
        ws.resolve("../../etc/passwd")
    with pytest.raises(ToolError):
        ws.resolve("a/../../b.txt")
    # 绝对路径逃逸
    with pytest.raises(ToolError):
        ws.resolve(str(ws.root.parent / "secret.txt"))


def test_write_read_roundtrip(ws):
    r = ws.write_file("dir/a.txt", "你好，工作区")
    assert r["ok"] is True
    got = ws.read_file("dir/a.txt")
    assert got["ok"] is True
    assert got["content"] == "你好，工作区"
    assert got["truncated"] is False
    # 越界写入被折叠为 ok=False（不抛出）
    bad = ws.execute("write_file", {"path": "../evil.txt", "content": "x"})
    assert bad["ok"] is False
    assert not (ws.root.parent / "evil.txt").exists()


def test_list_dir(ws):
    ws.write_file("b.txt", "x" * 10)
    ws.write_file("sub/c.txt", "y")
    r = ws.list_dir("")
    names = [e["name"] for e in r["entries"]]
    assert "b.txt" in names and "sub/" in names
    r2 = ws.list_dir("sub")
    assert r2["entries"][0]["name"] == "c.txt"
    # 不存在的目录
    assert ws.list_dir("nope")["ok"] is False


def test_read_truncation(ws):
    ws.write_file("big.txt", "字" * 20000)
    got = ws.read_file("big.txt")
    assert got["truncated"] is True
    assert len(got["content"]) == 8000


def test_unknown_tool_and_missing_file(ws):
    assert ws.execute("rm_rf", {"path": "."})["ok"] is False
    assert ws.execute("read_file", {"path": "ghost.txt"})["ok"] is False


# ---------------- parse_tool_calls ----------------

def test_parse_tool_calls():
    msg = {"tool_calls": [
        {"id": "c1", "function": {"name": "read_file", "arguments": '{"path": "a.txt"}'}},
        {"id": "c2", "function": {"name": "write_file", "arguments": "{broken json"}},
    ]}
    calls = parse_tool_calls(msg)
    assert calls[0]["name"] == "read_file"
    assert calls[0]["arguments"] == {"path": "a.txt"}
    assert calls[1]["arguments"] == {}            # 非法 JSON 容错
    assert parse_tool_calls({}) == []


# ---------------- ToolLoop ----------------

def _mk_chat(script):
    """script: 依次弹出的响应列表。记录每次收到的 messages。"""
    calls = []

    async def chat_fn(messages, tools_spec):
        calls.append([dict(m) for m in messages])
        return script.pop(0)

    chat_fn.calls = calls
    return chat_fn


def test_toolloop_answer_direct(ws):
    chat = _mk_chat([{"content": "最终答案", "tool_calls": [],
                      "usage": {"prompt_tokens": 10, "completion_tokens": 5}}])
    events = []
    loop = ToolLoop(ws, chat, on_event=events.append)
    result = asyncio_run(loop.run("任务"))
    assert result["response"] == "最终答案"
    assert result["reason"] == "answer"
    assert result["prompt_tokens"] == 10 and result["completion_tokens"] == 5
    assert events[0]["type"] == "round"
    assert events[-1]["type"] == "final" and events[-1]["reason"] == "answer"


def test_toolloop_write_then_answer(ws):
    """第一轮调 write_file，第二轮给最终答复；工具结果应回喂。"""
    chat = _mk_chat([
        {"content": None,
         "tool_calls": [{"id": "c1", "name": "write_file",
                         "arguments": {"path": "hello.txt", "content": "hi"}}],
         "usage": {"prompt_tokens": 20, "completion_tokens": 4}},
        {"content": "已写入 hello.txt", "tool_calls": [],
         "usage": {"prompt_tokens": 30, "completion_tokens": 6}},
    ])
    events = []
    loop = ToolLoop(ws, chat, on_event=events.append)
    result = asyncio_run(loop.run("写个文件"))
    assert result["reason"] == "answer"
    assert (ws.root / "hello.txt").read_text(encoding="utf-8") == "hi"
    kinds = [e["type"] for e in events]
    assert kinds.count("tool_call") == 1 and kinds.count("tool_result") == 1
    # 第二轮 messages 应包含 assistant(tool_calls) + tool 结果
    second = chat.calls[1]
    assert second[1]["role"] == "assistant"
    assert json.loads(second[1]["tool_calls"][0]["function"]["arguments"])["path"] == "hello.txt"
    assert second[2]["role"] == "tool"
    assert second[2]["tool_call_id"] == "c1"


def test_toolloop_max_rounds_forces_summary(ws):
    chat = _mk_chat([
        {"content": None,
         "tool_calls": [{"id": "c1", "name": "list_dir", "arguments": {}}],
         "usage": {"prompt_tokens": 5, "completion_tokens": 1}},
    ] * 3 + [
        {"content": "强制总结", "tool_calls": [], "usage": {"prompt_tokens": 5, "completion_tokens": 2}},
    ])
    loop = ToolLoop(ws, chat, max_rounds=3)
    result = asyncio_run(loop.run("任务"))
    assert result["reason"] == "max_rounds"
    assert result["response"] == "强制总结"
    assert result["error"]


def test_toolloop_token_cap(ws):
    chat = _mk_chat([
        {"content": "x", "tool_calls": [],
         "usage": {"prompt_tokens": 100, "completion_tokens": 100}},
    ] * 5)
    loop = ToolLoop(ws, chat, token_cap=150)
    result = asyncio_run(loop.run("任务"))
    assert result["reason"] == "token_cap"
    assert "熔断" in result["error"]
    assert len(chat.calls) == 1                    # 触顶后不再继续调用


def test_toolloop_chat_error(ws):
    async def bad_chat(messages, tools_spec):
        raise RuntimeError("网络断了")

    loop = ToolLoop(ws, bad_chat)
    result = asyncio_run(loop.run("任务"))
    assert result["reason"] == "error"
    assert "RuntimeError" in result["error"]