projectAIpopular/tests/test_agent_api.py

"""智能体端点测试：注入脚本化 chat_fn，不依赖真实模型/API key。"""
import json
import time

import pytest

pytest.importorskip("fastapi")
pytest.importorskip("httpx")

from fastapi.testclient import TestClient

import gateway.agent as ag
import gateway.api as ga
from gateway.model_pool import PoolStore
import gateway.model_pool as mp


@pytest.fixture()
def agent_env(tmp_path, monkeypatch):
    """隔离：池/服务/工作区全部指向临时目录，chat_fn 用脚本替身。"""
    mp.reset_pool()
    mp._store = PoolStore(path=tmp_path / "pool.json")
    ag.reset_agent_service()
    service = ag.AgentService(run_dir=tmp_path / "agent_runs")
    ag._service = service
    # 快照用户真实设置，测试后原样恢复（settings.json 是活文件，不能污染）
    store = ga.settings_store()
    snapshot = json.loads(json.dumps(store._data, ensure_ascii=False))
    # 工作区指向临时目录 + 给经典回退一个假 key（防止 .env 缺失时 400）
    store.update({"agent": {"workspace_dir": str(tmp_path / "ws")},
                  "architect": {"api_key": "sk-fake-test"}})
    ga.rebuild_pipeline()

    script = []

    def set_script(events):
        script.clear()
        script.extend(events)

    def fake_chat_factory(acfg):
        async def chat_fn(messages, tools_spec):
            if not script:
                return {"content": "（脚本用尽）好的。", "tool_calls": [], "usage": {}}
            return script.pop(0)
        return chat_fn

    monkeypatch.setattr(ga, "build_agent_chat", fake_chat_factory)
    yield {"service": service, "set_script": set_script, "ws": tmp_path / "ws"}

    store._data = snapshot
    store.save()
    mp.reset_pool()
    ag.reset_agent_service()
    ga.rebuild_pipeline()


@pytest.fixture()
def client():
    return TestClient(ga.app)


def _wait_done(service, rid, timeout=10.0):
    t0 = time.time()
    while time.time() - t0 < timeout:
        info = service.get(rid)
        if info and info.state in ("done", "failed"):
            return info
        time.sleep(0.05)
    return service.get(rid)


def test_agent_full_flow(agent_env, client):
    """写文件 -> 最终答复：验证事件、工作区落盘、状态终态。"""
    agent_env["set_script"]([
        {"content": None,
         "tool_calls": [{"id": "c1", "name": "write_file",
                         "arguments": {"path": "notes.md", "content": "# 笔记"}}],
         "usage": {"prompt_tokens": 30, "completion_tokens": 6}},
        {"content": "已创建 notes.md，任务完成。", "tool_calls": [],
         "usage": {"prompt_tokens": 40, "completion_tokens": 8}},
    ])
    r = client.post("/agent", json={"task": "帮我建一个 notes.md"})
    assert r.status_code == 200
    rid = r.json()["request_id"]
    assert r.json()["status"] == "running"

    info = _wait_done(agent_env["service"], rid)
    assert info.state == "done", info.error
    assert "notes.md" in info.response

    # 工作区真实落盘
    assert (agent_env["ws"] / "notes.md").read_text(encoding="utf-8") == "# 笔记"

    # 事件序列
    events = client.get(f"/agent/{rid}/events").json()
    kinds = [e["type"] for e in events]
    assert "tool_call" in kinds and "tool_result" in kinds and "final" in kinds
    assert events[-1]["reason"] == "answer"

    # 状态端点
    st = client.get(f"/agent/{rid}/status").json()
    assert st["state"] == "done"
    assert st["prompt_tokens"] == 70 and st["completion_tokens"] == 14

    # 工作区浏览端点
    ls = client.get("/agent/workspace").json()
    assert ls["ok"] is True
    assert any(e["name"] == "notes.md" for e in ls["entries"])
    f = client.get("/agent/file", params={"path": "notes.md"}).json()
    assert f["content"] == "# 笔记"


def test_agent_jail_via_api(agent_env, client):
    """工具结果为 ok=False（越界被拒），循环仍能继续到最终答复。"""
    agent_env["set_script"]([
        {"content": None,
         "tool_calls": [{"id": "c1", "name": "read_file",
                         "arguments": {"path": "../../secret.txt"}}],
         "usage": {"prompt_tokens": 10, "completion_tokens": 2}},
        {"content": "越界访问被拒绝。", "tool_calls": [], "usage": {}},
    ])
    r = client.post("/agent", json={"task": "读一下上级目录"})
    rid = r.json()["request_id"]
    info = _wait_done(agent_env["service"], rid)
    assert info.state == "done"
    events = agent_env["service"].read_events(rid)
    tool_result = next(e for e in events if e["type"] == "tool_result")
    assert tool_result["ok"] is False

    # 文件读取 API 直接越界 -> 404/400
    r2 = client.get("/agent/file", params={"path": "../../x.txt"})
    assert r2.status_code in (400, 404)


def test_agent_model_from_pool(agent_env, client, monkeypatch):
    """池 agent 角色（或显式 pool_id）应被采用；mock 池模型拒绝。"""
    from gateway.agent import OpenAICompatChat
    captured = {}
    real_factory = None

    # 先放一个 openai 池条目并指派 agent 角色
    client.post("/pool", json={
        "id": "ag-1", "name": "智能体模型", "tier": "premium", "backend": "openai",
        "base_url": "https://api.example.com", "model": "big-model-x",
        "api_key": "sk-abc1234567", "enabled": True,
    })
    client.put("/pool/roles", json={"agent": "ag-1"})

    # /agent 不带 pool_id -> 用池 agent 角色
    r = client.post("/agent", json={"task": "hi"})
    assert r.status_code == 200
    assert r.json()["model"] == "big-model-x"

    # mock 条目 -> 400
    client.post("/pool", json={
        "id": "mk-1", "name": "mock", "tier": "local", "backend": "mock",
        "model": "mock", "enabled": True,
    })
    r2 = client.post("/agent", json={"task": "hi", "pool_id": "mk-1"})
    assert r2.status_code == 400


def test_agent_task_validation(agent_env, client):
    assert client.post("/agent", json={"task": ""}).status_code == 400
    assert client.post("/agent", json={}).status_code == 400


def test_agent_404(agent_env, client):
    assert client.get("/agent/ghost/status").status_code == 404
    assert client.get("/agent/ghost/events").json() == []