roboco/tests/unit/llm/providers/test_kimi_cli_usage.py

"""kimi_cli_usage — resolve the session dir from a run's stdout meta line (or
the newest matching session dir), then sum the real 4-bucket
``usage.record``/``usageScope=="turn"`` events in that session's wire.jsonl.
"""

from __future__ import annotations

import json
import time
from typing import TYPE_CHECKING

from roboco.llm.providers import kimi_cli_usage as ku

if TYPE_CHECKING:
    from pathlib import Path

    import pytest


def _write_jsonl(path: Path, lines: list[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def _usage_record(
    *,
    input_other: int,
    output: int,
    cache_read: int = 0,
    cache_creation: int = 0,
    scope: str = "turn",
) -> str:
    return json.dumps(
        {
            "type": "usage.record",
            "model": "kimi-code/k3",
            "usageScope": scope,
            "usage": {
                "inputOther": input_other,
                "output": output,
                "inputCacheRead": cache_read,
                "inputCacheCreation": cache_creation,
            },
        }
    )


def _resume_hint(session_id: str) -> str:
    return json.dumps(
        {"role": "meta", "type": "session.resume_hint", "session_id": session_id}
    )


# ---------------------------------------------------------------------------
# session_id_from_run_log
# ---------------------------------------------------------------------------


def test_session_id_from_run_log_finds_terminal_meta_line(tmp_path: Path) -> None:
    log = tmp_path / "run.jsonl"
    _write_jsonl(
        log,
        [
            json.dumps({"role": "assistant", "content": "working"}),
            _resume_hint("session_abc123"),
        ],
    )
    assert ku.session_id_from_run_log(log) == "session_abc123"


def test_session_id_from_run_log_keeps_the_last_match(tmp_path: Path) -> None:
    log = tmp_path / "run.jsonl"
    _write_jsonl(log, [_resume_hint("session_first"), _resume_hint("session_second")])
    assert ku.session_id_from_run_log(log) == "session_second"


def test_session_id_from_run_log_none_when_absent(tmp_path: Path) -> None:
    log = tmp_path / "run.jsonl"
    _write_jsonl(log, [json.dumps({"role": "assistant", "content": "hi"})])
    assert ku.session_id_from_run_log(log) is None
    assert ku.session_id_from_run_log(tmp_path / "nope.jsonl") is None


# ---------------------------------------------------------------------------
# resolve_session_dir — primary (known id) + fallback (newest under cwd basename)
# ---------------------------------------------------------------------------


def test_resolve_session_dir_finds_known_session_id(tmp_path: Path) -> None:
    home = tmp_path / ".kimi-code"
    session_dir = home / "sessions" / "wd_myrepo_ab12cd34ef56" / "session_abc123"
    session_dir.mkdir(parents=True)
    resolved = ku.resolve_session_dir(
        session_id="session_abc123",
        workdir="/data/workspaces/myrepo",
        kimi_code_home=home,
    )
    assert resolved == session_dir


def test_resolve_session_dir_falls_back_to_newest(tmp_path: Path) -> None:
    home = tmp_path / ".kimi-code"
    wd_dir = home / "sessions" / "wd_myrepo_ab12cd34ef56"
    old_session = wd_dir / "session_old"
    new_session = wd_dir / "session_new"
    old_session.mkdir(parents=True)
    time.sleep(0.01)
    new_session.mkdir(parents=True)
    resolved = ku.resolve_session_dir(
        session_id=None, workdir="/data/workspaces/myrepo", kimi_code_home=home
    )
    assert resolved == new_session


def test_resolve_session_dir_none_when_sessions_root_absent(tmp_path: Path) -> None:
    home = tmp_path / ".kimi-code"
    assert (
        ku.resolve_session_dir(
            session_id=None, workdir="/x/myrepo", kimi_code_home=home
        )
        is None
    )


def test_resolve_session_dir_falls_back_when_id_not_found(tmp_path: Path) -> None:
    home = tmp_path / ".kimi-code"
    wd_dir = home / "sessions" / "wd_myrepo_ab12cd34ef56"
    only_session = wd_dir / "session_other"
    only_session.mkdir(parents=True)
    resolved = ku.resolve_session_dir(
        session_id="session_missing",
        workdir="/data/workspaces/myrepo",
        kimi_code_home=home,
    )
    assert resolved == only_session


# ---------------------------------------------------------------------------
# aggregate_usage_from_wire
# ---------------------------------------------------------------------------


def test_aggregate_sums_turn_scoped_usage_records(tmp_path: Path) -> None:
    wire = tmp_path / "wire.jsonl"
    _write_jsonl(
        wire,
        [
            _usage_record(input_other=100, output=50, cache_read=10),
            json.dumps({"type": "llm.request", "model": "kimi-code/k3"}),
            _usage_record(input_other=200, output=80, cache_read=20, cache_creation=5),
        ],
    )
    agg = ku.aggregate_usage_from_wire(wire)
    assert agg["inputOther"] == 300  # noqa: PLR2004
    assert agg["output"] == 130  # noqa: PLR2004
    assert agg["inputCacheRead"] == 30  # noqa: PLR2004
    assert agg["inputCacheCreation"] == 5  # noqa: PLR2004
    assert agg["turns"] == 2  # noqa: PLR2004


def test_aggregate_ignores_non_turn_scope_and_bad_lines(tmp_path: Path) -> None:
    wire = tmp_path / "wire.jsonl"
    _write_jsonl(
        wire,
        [
            "not json",
            _usage_record(input_other=5, output=1, scope="session"),
            _usage_record(input_other=10, output=5),
        ],
    )
    agg = ku.aggregate_usage_from_wire(wire)
    assert agg["inputOther"] == 10  # noqa: PLR2004
    assert agg["turns"] == 1


def test_aggregate_zero_for_missing_log(tmp_path: Path) -> None:
    agg = ku.aggregate_usage_from_wire(tmp_path / "nope.jsonl")
    assert agg["turns"] == 0
    assert all(v == 0 for k, v in agg.items() if k != "turns")


# ---------------------------------------------------------------------------
# capture_run_usage / main
# ---------------------------------------------------------------------------


def test_capture_run_usage_writes_usage_json(tmp_path: Path) -> None:
    home = tmp_path / ".kimi-code"
    session_dir = home / "sessions" / "wd_myrepo_hash1" / "session_abc"
    (session_dir / "agents" / "main").mkdir(parents=True)
    wire = session_dir / "agents" / "main" / "wire.jsonl"
    _write_jsonl(wire, [_usage_record(input_other=100, output=50, cache_read=10)])

    run_log = tmp_path / "run.jsonl"
    _write_jsonl(run_log, [_resume_hint("session_abc")])

    out = tmp_path / "usage.json"
    tokens = ku.capture_run_usage(
        run_log=run_log,
        workdir="/data/workspaces/myrepo",
        model="kimi-code/k3",
        out_path=out,
        kimi_code_home=home,
    )
    assert tokens == (100, 50, 10, 0)
    data = json.loads(out.read_text())
    assert data["model"] == "kimi-code/k3"
    assert data["tokens_input"] == 100  # noqa: PLR2004
    assert data["tokens_output"] == 50  # noqa: PLR2004
    assert data["tokens_cache_read"] == 10  # noqa: PLR2004
    assert data["turns"] == 1
    assert data["cost_usd"] > 0.0


def test_capture_run_usage_zero_when_no_session_found(tmp_path: Path) -> None:
    home = tmp_path / ".kimi-code"
    run_log = tmp_path / "run.jsonl"
    run_log.write_text("", encoding="utf-8")
    out = tmp_path / "usage.json"
    tokens = ku.capture_run_usage(
        run_log=run_log,
        workdir="/data/workspaces/myrepo",
        model="kimi-code/k3",
        out_path=out,
        kimi_code_home=home,
    )
    assert tokens == (0, 0, 0, 0)
    data = json.loads(out.read_text())
    assert data["tokens_input"] == 0
    assert data["turns"] == 0


def test_main_writes_usage_file(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    home = tmp_path / ".kimi-code"
    session_dir = home / "sessions" / "wd_myrepo_hash1" / "session_abc"
    (session_dir / "agents" / "main").mkdir(parents=True)
    wire = session_dir / "agents" / "main" / "wire.jsonl"
    _write_jsonl(wire, [_usage_record(input_other=200, output=100)])

    run_log = tmp_path / "run.jsonl"
    _write_jsonl(run_log, [_resume_hint("session_abc")])

    out = tmp_path / "usage.json"
    monkeypatch.setattr(ku, "USAGE_OUT_PATH", out)
    monkeypatch.setattr(ku, "KIMI_CODE_HOME", home)
    monkeypatch.setenv("ROBOCO_KIMI_RUN_LOG", str(run_log))
    monkeypatch.setenv("ROBOCO_KIMI_WORKDIR", "/data/workspaces/myrepo")
    monkeypatch.setenv("ROBOCO_AGENT_MODEL", "kimi-code/k3")

    assert ku.main() == 0
    data = json.loads(out.read_text())
    assert data["tokens_input"] == 200  # noqa: PLR2004
    assert data["tokens_output"] == 100  # noqa: PLR2004


def test_main_warns_when_run_log_env_missing(
    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None:
    monkeypatch.delenv("ROBOCO_KIMI_RUN_LOG", raising=False)
    with caplog.at_level("WARNING", logger="roboco.llm.providers.kimi_cli_usage"):
        assert ku.main() == 0
    assert any("ROBOCO_KIMI_RUN_LOG" in r.message for r in caplog.records)