projectAIpopular/tests/test_agent_features.py

"""dsh 功能对齐测试（T31）：

- 流式 _stream_partial 复位（首次成功后再次流式失败应回退非流式，而非误抛）
- LLM 调用重试退避（5xx/传输错误重试，非可重试错误不重试）
- 会话重命名（PATCH /agent/sessions/{sid}）
- 重复工具调用提醒（同工具同参数第 3 次起回喂警语）
- search_files 目录修剪（node_modules 等不进入）
- 原子写入（write_file 落盘内容完整、无 .tmp 残留）
- 慢工具线程卸载（run_tool_async 快内联/慢走线程，异常回传）
"""
import asyncio
import json

import httpx
import pytest

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

from fastapi.testclient import TestClient

import gateway.agent as ag
from gateway.agent import OpenAICompatChat


@pytest.fixture()
def ws(tmp_path):
    from router_system.tools import WorkspaceTools
    return WorkspaceTools(tmp_path / "ws")


# ---------------- _stream_partial 复位 ----------------

def _ok_sse_body():
    chunks = [{"choices": [{"delta": {"content": "第一段答复"}}]},
              {"choices": [{"delta": {}}], "usage": {"prompt_tokens": 1,
                                                     "completion_tokens": 1}}]
    lines = [f"data: {json.dumps(c, ensure_ascii=False)}" for c in chunks]
    lines.append("data: [DONE]")
    return ("\n\n".join(lines) + "\n\n").encode("utf-8")


def _nonstream_body(text: str) -> bytes:
    return json.dumps({
        "choices": [{"message": {"content": text}}],
        "usage": {"prompt_tokens": 2, "completion_tokens": 2},
    }).encode("utf-8")


def test_stream_partial_flag_resets_between_calls():
    """首次流式成功（置位）后，第二次流式失败应正常回退非流式。"""
    state = {"n": 0}

    def handler(request: httpx.Request) -> httpx.Response:
        state["n"] += 1
        if state["n"] == 1:
            return httpx.Response(200, content=_ok_sse_body())
        # 第二次：流式 500（无部分输出）-> 应回退非流式（第 3 次请求）
        if state["n"] == 2:
            return httpx.Response(500, content=b"boom")
        return httpx.Response(200, content=_nonstream_body("回退答案"))

    chat = OpenAICompatChat(base_url="http://x", api_key=None, model="m",
                            transport=httpx.MockTransport(handler),
                            retry_delay_s=0)
    r1 = asyncio.run(chat([{"role": "user", "content": "a"}], []))
    assert r1["content"] == "第一段答复"
    r2 = asyncio.run(chat([{"role": "user", "content": "b"}], []))
    assert r2["content"] == "回退答案"          # 不因上次置位而误抛
    assert state["n"] == 3


# ---------------- 重试退避 ----------------

def test_retry_on_5xx_then_success():
    calls = {"n": 0}

    def handler(request: httpx.Request) -> httpx.Response:
        calls["n"] += 1
        if calls["n"] == 1:
            return httpx.Response(502, content=b"bad gateway")
        return httpx.Response(200, content=_nonstream_body("恢复"))

    chat = OpenAICompatChat(base_url="http://x", api_key=None, model="m",
                            stream=False, max_retries=2, retry_delay_s=0,
                            transport=httpx.MockTransport(handler))
    r = asyncio.run(chat([{"role": "user", "content": "x"}], []))
    assert r["content"] == "恢复" and calls["n"] == 2


def test_retry_exhausted_raises():
    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(500, content=b"always down")

    chat = OpenAICompatChat(base_url="http://x", api_key=None, model="m",
                            stream=False, max_retries=1, retry_delay_s=0,
                            transport=httpx.MockTransport(handler))
    with pytest.raises(httpx.HTTPStatusError):
        asyncio.run(chat([{"role": "user", "content": "x"}], []))


def test_no_retry_on_4xx_client_error():
    calls = {"n": 0}

    def handler(request: httpx.Request) -> httpx.Response:
        calls["n"] += 1
        return httpx.Response(401, content=b"unauthorized")

    chat = OpenAICompatChat(base_url="http://x", api_key=None, model="m",
                            stream=False, max_retries=2, retry_delay_s=0,
                            transport=httpx.MockTransport(handler))
    with pytest.raises(httpx.HTTPStatusError):
        asyncio.run(chat([{"role": "user", "content": "x"}], []))
    assert calls["n"] == 1                    # 4xx 不重试


# ---------------- 会话重命名 ----------------

def test_session_rename_endpoint(tmp_path):
    import gateway.api as ga
    ag.reset_session_store()
    ag._session_store = ag.SessionStore(root=tmp_path / "sess")
    client = TestClient(ga.app)
    r = client.post("/agent/sessions", json={"title": "旧名", "workspace": ""})
    sid = r.json()["id"]
    r2 = client.patch(f"/agent/sessions/{sid}", json={"title": "新名字"})
    assert r2.status_code == 200
    assert r2.json()["title"] == "新名字"
    assert client.get(f"/agent/sessions/{sid}").json()["title"] == "新名字"
    # 空标题 400；不存在 404
    assert client.patch(f"/agent/sessions/{sid}", json={"title": " "}).status_code == 400
    assert client.patch("/agent/sessions/asdeadbeef99",
                        json={"title": "x"}).status_code == 404
    ag.reset_session_store()


# ---------------- 重复调用提醒 ----------------

def test_repeat_call_warning(ws, tmp_path):
    """同工具同参数第 3 次调用：回喂内容带系统提示 + repeat_warning 事件。"""
    from router_system.tools import ToolLoop
    seen_msgs = []
    calls = {"n": 0}

    async def chat(messages, tools_spec):
        seen_msgs.append(list(messages))
        calls["n"] += 1
        if calls["n"] <= 3:
            return {"content": None,
                    "tool_calls": [{"id": "c" + str(calls["n"]), "name": "read_file",
                                    "arguments": {"path": "a.txt"}}],
                    "usage": {}}
        return {"content": "收手了", "tool_calls": [], "usage": {}}

    events = []
    loop = ToolLoop(ws, chat, on_event=events.append)
    (tmp_path / "ws" / "a.txt").parent.mkdir(parents=True, exist_ok=True)
    (tmp_path / "ws" / "a.txt").write_text("x", encoding="utf-8")
    result = asyncio.run(loop.run("反复读"))
    assert result["response"] == "收手了"
    # 第 3 次工具结果消息应带提醒
    tool_msgs = [m for m in seen_msgs[3] if m.get("role") == "tool"]
    assert any("系统提示" in m["content"] for m in tool_msgs)
    assert any(e["type"] == "repeat_warning" and e["count"] == 3 for e in events)


# ---------------- search_files 修剪 ----------------

def test_search_files_prunes_skip_dirs(ws, tmp_path):
    root = tmp_path / "ws"
    (root / "node_modules" / "pkg").mkdir(parents=True, exist_ok=True)
    (root / "node_modules" / "pkg" / "dep.js").write_text("NEEDLE", encoding="utf-8")
    (root / "src").mkdir(parents=True, exist_ok=True)
    (root / "src" / "app.js").write_text("NEEDLE", encoding="utf-8")
    r = ws.search_files("NEEDLE")
    files = {m["file"] for m in r["matches"]}
    assert files == {"src/app.js"}             # node_modules 被修剪


# ---------------- 原子写入 ----------------

def test_atomic_write_roundtrip(ws, tmp_path):
    ws.write_file("sub/atomic.txt", "第一版")
    ws.edit_file("sub/atomic.txt", "第一版", "第二版")
    assert (tmp_path / "ws" / "sub" / "atomic.txt").read_text(
        encoding="utf-8") == "第二版"
    # 无 .tmp 残留
    leftovers = [p.name for p in (tmp_path / "ws" / "sub").iterdir()
                 if p.name.endswith(".tmp")]
    assert leftovers == []


# ---------------- 慢工具线程卸载 ----------------

def _boom(*_args):
    raise RuntimeError("线程内炸了")


def test_run_tool_async_fast_inline_and_thread_exception(ws):
    from router_system.tools import run_tool_async
    # 快工具：内联
    ws.write_file("fast.txt", "v")
    r = asyncio.run(run_tool_async(ws, "read_file", {"path": "fast.txt"}))
    assert r["ok"] is True and r["content"] == "v"
    # 异常从线程回传（execute 以属性形式提供）
    class Boom:
        execute = staticmethod(_boom)
    with pytest.raises(RuntimeError):
        asyncio.run(run_tool_async(Boom(), "read_file", {"path": "x"}))