tests/test_agent_features.py

"""Testy verbosity chatu, konfigurace sub-agentů a token accountingu."""

from types import SimpleNamespace

from fastapi.testclient import TestClient


def _tool_call_resp(name="run_command", args='{"command": "ls"}'):
    tc = SimpleNamespace(id="c1", function=SimpleNamespace(name=name, arguments=args))
    msg = SimpleNamespace(content="", tool_calls=[tc])
    return SimpleNamespace(choices=[SimpleNamespace(message=msg)],
                           usage={"prompt_tokens": 100, "completion_tokens": 10})


def _final_resp(text="HOTOVO"):
    msg = SimpleNamespace(content=text, tool_calls=None)
    return SimpleNamespace(choices=[SimpleNamespace(message=msg)],
                           usage={"prompt_tokens": 150, "completion_tokens": 30})


def _scripted_llm(app_module, responses):
    """Nahradí _llm sekvencí odpovědí; plní stats jako skutečné _llm."""
    it = iter(responses)

    def fake(messages, tools, prefer_api=False, use_cache=True, stats=None):
        resp = next(it)
        app_module.TOKENS.add("local", resp.usage)
        if stats is not None:
            stats.add("local", resp.usage)
        return resp
    return fake


def _run_chat(app_module):
    return list(app_module.agent_chat("udelej ukol", []))


def _setup(app_module, monkeypatch, verbosity):
    app_module.SETTINGS.update({"runner_url": "http://fake-runner",
                                "chat_verbosity": verbosity})
    monkeypatch.setattr(app_module, "_llm",
                        _scripted_llm(app_module, [_tool_call_resp(), _final_resp()]))
    monkeypatch.setattr(app_module, "_exec_tool",
                        lambda name, args, allowed, stats=None: {"ok": True})


# ---------------------------------------------------------------- verbosity

def test_verbosity_full(app_module, monkeypatch):
    _setup(app_module, monkeypatch, "full")
    out = _run_chat(app_module)
    assert "run_command" in out[-1] and "```json" in out[-1]
    assert "HOTOVO" in out[-1]
    assert "📊" in out[-1] and "vstup 250 tok" in out[-1] and "výstup 40 tok" in out[-1]
    assert "kontext 150 tok" in out[-1]


def test_verbosity_compact(app_module, monkeypatch):
    _setup(app_module, monkeypatch, "compact")
    out = _run_chat(app_module)
    assert "run_command" in out[-1]
    assert "```json" not in out[-1]          # bez JSON dumpů
    assert "HOTOVO" in out[-1] and "📊" in out[-1]


def test_verbosity_final(app_module, monkeypatch):
    _setup(app_module, monkeypatch, "final")
    out = _run_chat(app_module)
    assert any("⏳" in chunk for chunk in out[:-1])   # heartbeat během práce
    assert "HOTOVO" in out[-1]
    assert "run_command" not in out[-1]
    assert "📊" not in out[-1]                        # bez patičky


def test_global_token_counter_grows(app_module, monkeypatch):
    _setup(app_module, monkeypatch, "full")
    before = app_module.TOKENS.snapshot()["total_tokens"]
    _run_chat(app_module)
    snap = app_module.TOKENS.snapshot()
    assert snap["total_tokens"] == before + 290
    assert snap["by_source"]["local"]["calls"] == 2


def test_health_exposes_tokens(app_module, auth):
    c = TestClient(app_module.app)
    data = c.get("/health").json()
    assert "tokens" in data
    assert set(data["tokens"]) >= {"calls", "prompt_tokens", "completion_tokens",
                                   "last_context_tokens",
                                   "saved_by_compaction_tokens"}


# ---------------------------------------------------------------- kompakce v loopu

def test_agent_loop_emits_compaction_note(app_module, monkeypatch):
    app_module.SETTINGS.update({"runner_url": "http://fake",
                                "context_budget_tokens": 300,
                                "context_keep_last_steps": 1,
                                "chat_verbosity": "full"})
    monkeypatch.setattr(app_module, "_llm",
                        _scripted_llm(app_module, [_final_resp()]))
    # kotva (system + první user) se nekompaktuje — potřeba starší historie
    msgs = [{"role": "user", "content": "zadani ukolu"},
            {"role": "assistant", "content": "a" * 4000},
            {"role": "user", "content": "b" * 4000},
            {"role": "assistant", "content": "c" * 4000},
            {"role": "user", "content": "pokracuj"}]
    events = list(app_module.agent_loop("sys prompt", msgs, [], set(), 5,
                                        yield_progress=True))
    kinds = [k for k, _ in events]
    assert "note" in kinds
    note = next(d for k, d in events if k == "note")
    assert "🧹" in note and "ušetřeno" in note


def test_compaction_off(app_module, monkeypatch):
    app_module.SETTINGS.update({"context_compaction": "off",
                                "context_budget_tokens": 80})
    monkeypatch.setattr(app_module, "_llm",
                        _scripted_llm(app_module, [_final_resp()]))
    msgs = [{"role": "user", "content": "x" * 4000}]
    events = list(app_module.agent_loop("sys", msgs, [], set(), 5,
                                        yield_progress=True))
    assert all(k != "note" for k, _ in events)


# ---------------------------------------------------------------- sub-agenti

def test_subagent_disabled_returns_hint(app_module):
    app_module.SETTINGS.update({"subagent_coder_enabled": False})
    res = app_module.run_subagent("coder", "udelej neco")
    assert "vypnuty" in res["error"]


def test_subagent_unknown_role(app_module):
    assert "error" in app_module.run_subagent("hacker", "x")


def test_delegate_tool_reflects_enabled_roles(app_module):
    tools = app_module.build_main_tools()
    delegate = [t for t in tools if t["function"]["name"] == "delegate_task"]
    assert delegate[0]["function"]["parameters"]["properties"]["role"]["enum"] == \
        ["explorer", "coder", "reviewer"]

    app_module.SETTINGS.update({"subagent_coder_enabled": False})
    tools = app_module.build_main_tools()
    delegate = [t for t in tools if t["function"]["name"] == "delegate_task"]
    assert delegate[0]["function"]["parameters"]["properties"]["role"]["enum"] == \
        ["explorer", "reviewer"]
    # původní definice nezmutována
    assert app_module.DELEGATE_TOOL["function"]["parameters"]["properties"]["role"]["enum"] == \
        ["explorer", "coder", "reviewer"]

    app_module.SETTINGS.update({"subagent_explorer_enabled": False,
                                "subagent_reviewer_enabled": False})
    tools = app_module.build_main_tools()
    assert all(t["function"]["name"] != "delegate_task" for t in tools)


def test_subagent_prompt_override(app_module, monkeypatch):
    app_module.SETTINGS.update({"subagent_explorer_prompt": "CUSTOM EXPLORER PROMPT"})
    captured = {}

    def fake(messages, tools, prefer_api=False, use_cache=True, stats=None):
        captured["system"] = messages[0]["content"]
        return _final_resp("done")
    monkeypatch.setattr(app_module, "_llm", fake)
    res = app_module.run_subagent("explorer", "prozkoumej repo")
    assert res["result"] == "done"
    assert captured["system"] == "CUSTOM EXPLORER PROMPT"


def test_subagent_default_prompt_when_empty(app_module, monkeypatch):
    captured = {}

    def fake(messages, tools, prefer_api=False, use_cache=True, stats=None):
        captured["system"] = messages[0]["content"]
        return _final_resp("done")
    monkeypatch.setattr(app_module, "_llm", fake)
    app_module.run_subagent("reviewer", "zkontroluj")
    assert captured["system"] == app_module.SUB_PROMPTS["reviewer"]


def test_turn_stats_include_subagent_calls(app_module, monkeypatch):
    """Tokeny sub-agenta se počítají do stats rodičovského běhu."""
    from context import TokenStats
    monkeypatch.setattr(app_module, "_llm",
                        _scripted_llm(app_module, [_final_resp("sub done")]))
    stats = TokenStats()
    app_module.run_subagent("explorer", "ukol", stats=stats)
    assert stats.snapshot()["calls"] == 1
    assert stats.snapshot()["total_tokens"] == 180