cad-agent/tests/test_review.py

"""`dwg_review.phases.review` 单元测试。"""

import json

from dwg_review.cad_tools import READ_ONLY_TOOLS
from dwg_review.mcp_client import McpToolError
from dwg_review.phases.review import (
    _run_readonly_tool,
    evidence_payload,
    review_frame,
)
from tests.fakes import (
    FakeLlmClient,
    FakeMcpClient,
    llm_calls,
    llm_text,
    make_ctx,
    make_evidence,
    ok_image,
    ok_mixed,
    ok_text,
)


# ---------------------------------------------------------------------------
#  evidence_payload
# ---------------------------------------------------------------------------
def test_evidence_payload_strips_image():
    ev = make_evidence(image_b64="QUJD", image_mime="image/png")
    payload = evidence_payload(ev)
    assert "images" not in payload
    assert "image_b64" not in payload
    assert payload["frame_index"] == 1


def test_evidence_payload_keeps_only_frame_metadata():
    ev = make_evidence()
    ev["frame"]["secret"] = "should-not-leak"
    payload = evidence_payload(ev)
    assert "secret" not in payload["frame"]
    assert set(payload["frame"]) == {"sheet_size", "kind", "layer", "bbox", "confidence"}


def test_evidence_payload_carries_structured_data():
    ev = make_evidence(entities={"total": 9, "items": []}, texts={"total": 2})
    payload = evidence_payload(ev)
    assert payload["entities"] == {"total": 9, "items": []}
    assert payload["texts"] == {"total": 2}


# ---------------------------------------------------------------------------
#  _run_readonly_tool
# ---------------------------------------------------------------------------
def test_readonly_rejects_non_whitelisted_tool():
    client = FakeMcpClient()
    ctx = make_ctx(client=client)
    content, budget = _run_readonly_tool(ctx, "get_some_unknown_tool", {}, [3])
    assert "不在只读取证白名单内" in content[0]["text"]
    assert client.calls == []
    assert budget == [3]


def test_readonly_rejects_write_tool():
    client = FakeMcpClient()
    ctx = make_ctx(client=client)
    content, _ = _run_readonly_tool(ctx, "add_leader_annotation", {}, [3])
    assert "不在只读取证白名单内" in content[0]["text"]
    assert client.calls == []


def test_readonly_whitelist_covers_expected_tools():
    assert {"query_entities", "get_entity_detail", "zoom_to_window"} <= READ_ONLY_TOOLS


def test_readonly_handles_tool_error():
    client = FakeMcpClient(responses={"query_entities": McpToolError("炸了")})
    ctx = make_ctx(client=client)
    content, _ = _run_readonly_tool(ctx, "query_entities", {"window": {}}, [3])
    assert "执行失败" in content[0]["text"]


def test_readonly_text_and_image_content():
    client = FakeMcpClient(
        responses={"query_entities": ok_mixed(["hi"], [("QUJD", "image/png")])}
    )
    ctx = make_ctx(client=client)
    content, budget = _run_readonly_tool(ctx, "query_entities", {}, [3])
    assert content[0] == {"type": "text", "text": "hi"}
    assert content[1]["type"] == "image_url"
    assert budget == [2]


def test_readonly_image_budget_exhausted():
    client = FakeMcpClient(
        responses={"query_entities": ok_mixed(["t"], [("QUJD", "image/png")])}
    )
    ctx = make_ctx(client=client)
    content, budget = _run_readonly_tool(ctx, "query_entities", {}, [0])
    assert budget == [0]
    assert any("已省略" in c.get("text", "") for c in content)


def test_readonly_empty_response_marker():
    empty = {"jsonrpc": "2.0", "id": 1, "result": {"content": []}}
    client = FakeMcpClient(responses={"query_entities": empty})
    ctx = make_ctx(client=client)
    content, _ = _run_readonly_tool(ctx, "query_entities", {}, [3])
    assert content == [{"type": "text", "text": "（无返回内容）"}]


# ---------------------------------------------------------------------------
#  review_frame
# ---------------------------------------------------------------------------
def test_review_frame_returns_summary_when_no_tools():
    llm = FakeLlmClient(script=[llm_text("结论A")])
    ctx = make_ctx(client=FakeMcpClient(), llm=llm)
    out = review_frame(ctx, make_evidence(), "系统提示", log=lambda *a: None)
    assert out["summary"] == "结论A"  # 非 JSON → 降级为纯文本 summary
    assert out["issues"] == []
    assert len(llm.calls) == 1
    assert llm.calls[0]["tools"] is not None  # 传了只读工具 schema


def test_review_frame_includes_image_block_when_present():
    llm = FakeLlmClient(script=[llm_text("ok")])
    ctx = make_ctx(client=FakeMcpClient(), llm=llm)
    ev = make_evidence(image_b64="QUJD", image_mime="image/png")
    review_frame(ctx, ev, "sys", log=lambda *a: None)
    user_content = llm.calls[0]["messages"][1]["content"]
    assert any(b.get("type") == "image_url" for b in user_content)


def test_review_frame_omits_image_block_when_absent():
    llm = FakeLlmClient(script=[llm_text("ok")])
    ctx = make_ctx(client=FakeMcpClient(), llm=llm)
    review_frame(ctx, make_evidence(), "sys", log=lambda *a: None)
    user_content = llm.calls[0]["messages"][1]["content"]
    assert all(b.get("type") == "text" for b in user_content)


def test_review_frame_runs_tool_then_concludes():
    llm = FakeLlmClient(
        script=[
            llm_calls([("query_entities", {"window": {}})]),
            llm_text("结论B"),
        ]
    )
    client = FakeMcpClient(responses={"query_entities": ok_text({"total": 1})})
    ctx = make_ctx(client=client, llm=llm)
    out = review_frame(ctx, make_evidence(), "sys", log=lambda *a: None)
    assert out["summary"] == "结论B"
    assert len(llm.calls) == 2
    assert client.called_names() == ["query_entities"]


def test_review_frame_feeds_observation_back_as_tool_message():
    llm = FakeLlmClient(script=[llm_calls([("query_entities", {})]), llm_text("done")])
    ctx = make_ctx(client=FakeMcpClient(), llm=llm)
    review_frame(ctx, make_evidence(), "sys", log=lambda *a: None)
    tool_msgs = [m for m in llm.calls[1]["messages"] if isinstance(m, dict) and m.get("role") == "tool"]
    assert tool_msgs and tool_msgs[0]["tool_call_id"] == "1"


def test_review_frame_invalid_tool_arguments_become_empty():
    llm = FakeLlmClient(
        script=[llm_calls([("query_entities", "not-json")]), llm_text("done")]
    )
    client = FakeMcpClient()
    ctx = make_ctx(client=client, llm=llm)
    review_frame(ctx, make_evidence(), "sys", log=lambda *a: None)
    assert client.arguments_of("query_entities") == {}


def test_review_frame_force_closes_at_iteration_limit():
    llm = FakeLlmClient(
        script=[
            llm_calls([("query_entities", {})]),
            llm_text("收口结论"),
        ]
    )
    ctx = make_ctx(client=FakeMcpClient(), llm=llm, max_iters=1)
    out = review_frame(ctx, make_evidence(), "sys", log=lambda *a: None)
    assert out["summary"] == "收口结论"
    last_msg = llm.calls[-1]["messages"][-1]
    assert last_msg["role"] == "user"
    assert "上限" in last_msg["content"]


def test_review_frame_with_zero_iters_forces_immediate_close():
    llm = FakeLlmClient(script=[llm_text("直接结论")])
    ctx = make_ctx(client=FakeMcpClient(), llm=llm, max_iters=0)
    out = review_frame(ctx, make_evidence(), "sys", log=lambda *a: None)
    assert out["summary"] == "直接结论"
    assert len(llm.calls) == 1


def test_review_frame_logs_tool_usage():
    logs = []
    llm = FakeLlmClient(script=[llm_calls([("get_layers", {})]), llm_text("x")])
    ctx = make_ctx(client=FakeMcpClient(), llm=llm)
    review_frame(ctx, make_evidence(), "sys", log=logs.append)
    assert any("取证" in line and "get_layers" in line for line in logs)


def test_review_frame_payload_is_json_serializable():
    llm = FakeLlmClient(script=[llm_text("ok")])
    ctx = make_ctx(client=FakeMcpClient(), llm=llm)
    review_frame(ctx, make_evidence(), "sys", log=lambda *a: None)
    text = llm.calls[0]["messages"][1]["content"][0]["text"]
    json.loads(text.split("：\n", 1)[1])  # 不抛异常即通过


def test_review_frame_parses_structured_json_conclusion():
    text = (
        '{"summary": "整体尚可", "issues": ['
        '{"category": "几何", "severity": "medium", "message": "零长度", "handle": "A1"}]}'
    )
    llm = FakeLlmClient(script=[llm_text(text)])
    ctx = make_ctx(client=FakeMcpClient(), llm=llm)
    out = review_frame(ctx, make_evidence(), "sys", log=lambda *a: None)
    assert out["summary"] == "整体尚可"
    assert out["issues"] == [
        {"category": "几何", "severity": "medium", "message": "零长度", "handle": "A1"}
    ]


def test_review_frame_parses_fenced_json():
    llm = FakeLlmClient(script=[llm_text('```json\n{"summary": "s", "issues": []}\n```')])
    ctx = make_ctx(client=FakeMcpClient(), llm=llm)
    out = review_frame(ctx, make_evidence(), "sys", log=lambda *a: None)
    assert out["summary"] == "s"
    assert out["issues"] == []


def test_review_frame_intermediate_tool_turn_json_not_parsed():
    # 中间工具回合的 content 里即使塞了 JSON，也不能被当成最终结论
    intermediate = '{"summary": "这不是最终结论", "issues": [{"message": "x"}]}'
    llm = FakeLlmClient(script=[
        llm_calls([("query_entities", {})], content=intermediate),
        llm_text('{"summary": "最终结论", "issues": []}'),
    ])
    ctx = make_ctx(client=FakeMcpClient(), llm=llm)
    out = review_frame(ctx, make_evidence(), "sys", log=lambda *a: None)
    assert out["summary"] == "最终结论"
    assert out["issues"] == []


def test_review_frame_includes_multiple_image_blocks():
    llm = FakeLlmClient(script=[llm_text("ok")])
    ctx = make_ctx(client=FakeMcpClient(), llm=llm)
    ev = make_evidence(
        images=[
            {"b64": "A", "mime": "image/png", "label": "overall"},
            {"b64": "B", "mime": "image/png", "label": "tile-1"},
        ]
    )
    review_frame(ctx, ev, "sys", log=lambda *a: None)
    user_content = llm.calls[0]["messages"][1]["content"]
    image_blocks = [b for b in user_content if b.get("type") == "image_url"]
    assert len(image_blocks) == 2
    assert any(b.get("text") == "（截图：overall）" for b in user_content)
    assert any(b.get("text") == "（截图：tile-1）" for b in user_content)


# ---------------------------------------------------------------------------
#  视口复位（放大取证不残留给用户）
# ---------------------------------------------------------------------------
def test_review_frame_resets_viewport_after_zoom():
    llm = FakeLlmClient(
        script=[
            llm_calls([("zoom_to_window", {"min_x": 1, "min_y": 2, "max_x": 3, "max_y": 4})]),
            llm_text("done"),
        ]
    )
    client = FakeMcpClient()
    ctx = make_ctx(client=client, llm=llm)
    review_frame(ctx, make_evidence(), "sys", log=lambda *a: None)
    zooms = client.called("zoom_to_window")
    assert len(zooms) == 2  # LLM 放大一次 + 结束时复位一次
    reset = zooms[-1]["arguments"]
    assert (reset["min_x"], reset["min_y"]) == (0.0, 0.0)  # 复位到图框整幅视图
    assert (reset["max_x"], reset["max_y"]) == (100.0, 50.0)
    assert reset["margin"] == 0.02


def test_review_frame_resets_viewport_after_zoom_to_entity():
    llm = FakeLlmClient(
        script=[llm_calls([("zoom_to_entity", {"handle": "A1"})]), llm_text("done")]
    )
    client = FakeMcpClient()
    ctx = make_ctx(client=client, llm=llm)
    review_frame(ctx, make_evidence(), "sys", log=lambda *a: None)
    assert client.called_names().count("zoom_to_window") == 1


def test_review_frame_resets_viewport_at_iteration_limit():
    llm = FakeLlmClient(
        script=[llm_calls([("zoom_to_window", {})]), llm_text("收口结论")]
    )
    client = FakeMcpClient()
    ctx = make_ctx(client=client, llm=llm, max_iters=1)
    out = review_frame(ctx, make_evidence(), "sys", log=lambda *a: None)
    assert out["summary"] == "收口结论"
    assert client.called_names().count("zoom_to_window") == 2  # 强制收口路径也复位


def test_review_frame_does_not_touch_viewport_without_zoom():
    llm = FakeLlmClient(script=[llm_calls([("get_layers", {})]), llm_text("x")])
    client = FakeMcpClient()
    ctx = make_ctx(client=client, llm=llm)
    review_frame(ctx, make_evidence(), "sys", log=lambda *a: None)
    assert "zoom_to_window" not in client.called_names()


def test_review_frame_viewport_reset_failure_is_swallowed():
    seen = {"n": 0}

    def zoom(_args):
        seen["n"] += 1
        if seen["n"] > 1:
            raise McpToolError("复位失败")
        return ok_text({})

    llm = FakeLlmClient(
        script=[llm_calls([("zoom_to_window", {})]), llm_text("结论")]  # noqa: E501
    )
    client = FakeMcpClient(responses={"zoom_to_window": zoom})
    ctx = make_ctx(client=client, llm=llm)
    out = review_frame(ctx, make_evidence(), "sys", log=lambda *a: None)
    assert out["summary"] == "结论"  # 复位失败不影响结论
    assert seen["n"] == 2


def test_review_frame_no_viewport_reset_without_window():
    llm = FakeLlmClient(script=[llm_calls([("zoom_to_window", {})]), llm_text("x")])
    client = FakeMcpClient()
    ctx = make_ctx(client=client, llm=llm)
    ev = make_evidence()
    ev["window"] = None  # 非 CAD 源（无窗口）→ 无从复位
    review_frame(ctx, ev, "sys", log=lambda *a: None)
    assert client.called_names().count("zoom_to_window") == 1