chore: import upstream snapshot with attribution

tests/test_cursor_native_usage.py

@@ -0,0 +1,426 @@
+"""Unit tests for cursor-native token-usage capture.
+
+Covers the pure pieces a live cursor-agent isn't needed for: normalizing the
+``stop``-hook payload, the append-only usage log + recorder CLI, the cumulative
+accumulator (per-turn sum + generation-id dedup), state round-trip, the
+``external_session_usage`` POST shape, and the poll loop (POST-on-change,
+no-repost-when-unchanged, persist-only-after-success). The live tmux +
+cursor-agent hook path is exercised by the e2e gate, not here.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import json
+import subprocess
+import sys
+from pathlib import Path
+
+import httpx
+import pytest
+
+from omnigent import cursor_native_usage as usage
+
+# A representative cursor ``stop``-hook payload (live-captured field set).
+_TURN1 = {
+    "generation_id": "g1",
+    "model": "claude-4-sonnet",
+    "status": "completed",
+    "input_tokens": 23666,
+    "output_tokens": 5,
+    "cache_read_tokens": 23617,
+    "cache_write_tokens": 47,
+}
+_TURN2 = {
+    "generation_id": "g2",
+    "model": "claude-4-sonnet",
+    "status": "completed",
+    "input_tokens": 24010,
+    "output_tokens": 120,
+    "cache_read_tokens": 23700,
+    "cache_write_tokens": 60,
+}
+
+
+class TestNormalizeHookPayload:
+    def test_extracts_usage_fields(self) -> None:
+        line = usage.normalize_hook_payload(_TURN1)
+        assert line == {
+            "generation_id": "g1",
+            "model": "claude-4-sonnet",
+            "input_tokens": 23666,
+            "output_tokens": 5,
+            "cache_read_tokens": 23617,
+            "cache_write_tokens": 47,
+        }
+
+    def test_falls_back_to_conversation_id(self) -> None:
+        line = usage.normalize_hook_payload({"conversation_id": "c1", "output_tokens": 3})
+        assert line is not None and line["generation_id"] == "c1"
+
+    def test_skips_when_no_generation_id(self) -> None:
+        assert usage.normalize_hook_payload({"output_tokens": 3}) is None
+
+    def test_skips_when_all_tokens_zero(self) -> None:
+        payload = {"generation_id": "g", "input_tokens": 0, "output_tokens": 0}
+        assert usage.normalize_hook_payload(payload) is None
+
+    def test_omits_model_when_absent(self) -> None:
+        line = usage.normalize_hook_payload({"generation_id": "g", "output_tokens": 4})
+        assert line is not None and "model" not in line
+
+    def test_coerces_and_floors_negative_tokens(self) -> None:
+        line = usage.normalize_hook_payload(
+            {"generation_id": "g", "input_tokens": "7", "output_tokens": -3}
+        )
+        assert line is not None
+        assert line["input_tokens"] == 7  # coerced from str
+        assert line["output_tokens"] == 0  # negative floored
+
+    def test_non_dict_is_skipped(self) -> None:
+        assert usage.normalize_hook_payload("nope") is None
+
+
+class TestRecordUsagePayload:
+    def test_appends_one_line_per_turn(self, tmp_path: Path) -> None:
+        assert usage.record_usage_payload(tmp_path, _TURN1) is True
+        assert usage.record_usage_payload(tmp_path, _TURN2) is True
+        lines = (tmp_path / usage.USAGE_FILE).read_text().splitlines()
+        assert len(lines) == 2
+        assert json.loads(lines[0])["generation_id"] == "g1"
+
+    def test_non_billable_payload_is_skipped(self, tmp_path: Path) -> None:
+        assert usage.record_usage_payload(tmp_path, {"text": "hi"}) is False
+        assert not (tmp_path / usage.USAGE_FILE).exists()
+
+
+class TestRecordUsageCli:
+    def test_cli_reads_stdin_appends_and_emits_continue(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture
+    ) -> None:
+        import io
+
+        monkeypatch.setattr("sys.stdin", io.StringIO(json.dumps(_TURN1)))
+        rc = usage._cli_record_usage(tmp_path)
+        assert rc == 0
+        # cursor reads stdout as the hook response; "{}" means "continue".
+        assert capsys.readouterr().out == "{}"
+        assert (tmp_path / usage.USAGE_FILE).read_text().strip()
+
+    def test_cli_never_fails_on_garbage_stdin(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture
+    ) -> None:
+        import io
+
+        monkeypatch.setattr("sys.stdin", io.StringIO("not json{{{"))
+        assert usage._cli_record_usage(tmp_path) == 0
+        assert capsys.readouterr().out == "{}"
+        assert not (tmp_path / usage.USAGE_FILE).exists()
+
+    def test_module_main_entrypoint(self, tmp_path: Path) -> None:
+        # End-to-end through the real CLI the hooks.json command invokes. ``-I``
+        # is dropped here (the worktree isn't pip-installed); PYTHONPATH makes it
+        # importable, mirroring the installed-package resolution in production.
+        env = {"PYTHONPATH": str(Path(__file__).resolve().parents[1])}
+        proc = subprocess.run(
+            [
+                sys.executable,
+                "-m",
+                "omnigent.cursor_native_usage",
+                "record-usage",
+                "--bridge-dir",
+                str(tmp_path),
+            ],
+            input=json.dumps(_TURN1),
+            capture_output=True,
+            text=True,
+            env=env,
+        )
+        assert proc.returncode == 0, proc.stderr
+        assert proc.stdout == "{}"
+        assert (tmp_path / usage.USAGE_FILE).read_text().strip()
+
+
+class TestUsageAccumulator:
+    def test_sums_per_turn_counts(self) -> None:
+        acc = usage._UsageAccumulator()
+        assert acc.add_line(usage.normalize_hook_payload(_TURN1)) is True  # type: ignore[arg-type]
+        assert acc.add_line(usage.normalize_hook_payload(_TURN2)) is True  # type: ignore[arg-type]
+        assert acc.input_tokens == 23666 + 24010
+        assert acc.output_tokens == 5 + 120
+        assert acc.cache_read_tokens == 23617 + 23700
+        assert acc.model == "claude-4-sonnet"
+
+    def test_dedups_by_generation_id(self) -> None:
+        acc = usage._UsageAccumulator()
+        line = usage.normalize_hook_payload(_TURN1)
+        assert acc.add_line(line) is True  # type: ignore[arg-type]
+        assert acc.add_line(line) is False  # type: ignore[arg-type] — same gen id, ignored
+        assert acc.output_tokens == 5  # not doubled
+
+    def test_latest_model_wins(self) -> None:
+        acc = usage._UsageAccumulator()
+        acc.add_line({"generation_id": "a", "model": "claude-4-sonnet", "output_tokens": 1})
+        acc.add_line({"generation_id": "b", "model": "gpt-5", "output_tokens": 1})
+        assert acc.model == "gpt-5"
+
+    def test_line_without_gen_id_is_ignored(self) -> None:
+        acc = usage._UsageAccumulator()
+        assert acc.add_line({"output_tokens": 9}) is False
+        assert acc.output_tokens == 0
+
+
+class TestStateRoundTrip:
+    def test_write_then_read_preserves_totals_and_seen(self, tmp_path: Path) -> None:
+        acc = usage._UsageAccumulator(
+            input_tokens=10, output_tokens=2, cache_read_tokens=4, model="gpt-5", seen={"g1"}
+        )
+        usage._write_usage_state(tmp_path, acc)
+        loaded = usage._read_usage_state(tmp_path)
+        assert loaded.input_tokens == 10
+        assert loaded.output_tokens == 2
+        assert loaded.cache_read_tokens == 4
+        assert loaded.model == "gpt-5"
+        assert loaded.seen == {"g1"}
+
+    def test_missing_state_is_cold_default(self, tmp_path: Path) -> None:
+        loaded = usage._read_usage_state(tmp_path)
+        assert (loaded.input_tokens, loaded.output_tokens, loaded.cache_read_tokens) == (0, 0, 0)
+        assert loaded.model is None and loaded.seen == set()
+
+
+class TestReadUsageLines:
+    def test_reads_valid_lines_and_skips_garbage(self, tmp_path: Path) -> None:
+        (tmp_path / usage.USAGE_FILE).write_text(
+            json.dumps({"generation_id": "g1", "output_tokens": 1})
+            + "\n\n"  # blank line tolerated
+            + "not json\n"  # garbage tolerated
+            + json.dumps({"generation_id": "g2", "output_tokens": 2})
+            + "\n"
+        )
+        lines = usage._read_usage_lines(tmp_path)
+        assert [line["generation_id"] for line in lines] == ["g1", "g2"]
+
+    def test_missing_file_is_empty(self, tmp_path: Path) -> None:
+        assert usage._read_usage_lines(tmp_path) == []
+
+
+class TestUsagePostBody:
+    def test_shape_with_model(self) -> None:
+        acc = usage._UsageAccumulator(
+            input_tokens=100, output_tokens=20, cache_read_tokens=80, model="gpt-5"
+        )
+        assert usage._usage_post_body(acc) == {
+            "cumulative_input_tokens": 100,
+            "cumulative_output_tokens": 20,
+            "cumulative_cache_read_input_tokens": 80,
+            "model": "gpt-5",
+        }
+
+    def test_omits_model_when_none(self) -> None:
+        body = usage._usage_post_body(usage._UsageAccumulator(output_tokens=5))
+        assert "model" not in body
+        assert body["cumulative_output_tokens"] == 5
+
+
+class TestClearUsageState:
+    def test_removes_log_and_state(self, tmp_path: Path) -> None:
+        usage.record_usage_payload(tmp_path, _TURN1)
+        usage._write_usage_state(tmp_path, usage._UsageAccumulator(output_tokens=1))
+        usage.clear_cursor_usage_state(tmp_path)
+        assert not (tmp_path / usage.USAGE_FILE).exists()
+        assert not (tmp_path / usage._USAGE_STATE_FILE).exists()
+
+    def test_noop_when_absent(self, tmp_path: Path) -> None:
+        usage.clear_cursor_usage_state(tmp_path)  # must not raise
+
+
+class _CtxRecordingClient:
+    """Async httpx-client stub (records POSTs, 200) usable as ``async with``."""
+
+    def __init__(self) -> None:
+        self.posts: list[tuple[str, dict]] = []
+
+    async def __aenter__(self) -> _CtxRecordingClient:
+        return self
+
+    async def __aexit__(self, *exc: object) -> None:
+        return None
+
+    async def post(self, url: str, *, json: dict) -> httpx.Response:
+        self.posts.append((url, json))
+        return httpx.Response(200, request=httpx.Request("POST", url))
+
+
+async def _run_loop_until(
+    monkeypatch: pytest.MonkeyPatch,
+    bridge_dir: Path,
+    until,
+    *,
+    client: _CtxRecordingClient | None = None,
+    max_wait_s: float = 3.0,
+) -> _CtxRecordingClient:
+    """Run the real poll loop with a recording client until *until(client)* holds."""
+    client = client or _CtxRecordingClient()
+    monkeypatch.setattr(httpx, "AsyncClient", lambda **kwargs: client)
+    task = asyncio.create_task(
+        usage.forward_cursor_usage_to_session(
+            base_url="http://test",
+            headers={},
+            session_id="conv_1",
+            bridge_dir=bridge_dir,
+            poll_interval_s=0.01,
+        )
+    )
+    loop = asyncio.get_running_loop()
+    deadline = loop.time() + max_wait_s
+    try:
+        while loop.time() < deadline:
+            if until(client):
+                return client
+            await asyncio.sleep(0.01)
+        raise AssertionError("loop condition was not met before timeout")
+    finally:
+        task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await task
+
+
+def _usage_posts(client: _CtxRecordingClient) -> list[tuple[str, dict]]:
+    """Only the ``external_session_usage`` POSTs (excludes the idle wake edges)."""
+    return [(u, b) for (u, b) in client.posts if b.get("type") == "external_session_usage"]
+
+
+def _idle_posts(client: _CtxRecordingClient) -> list[tuple[str, dict]]:
+    """Only the ``external_session_status: idle`` turn-end wake POSTs."""
+    return [
+        (u, b)
+        for (u, b) in client.posts
+        if b.get("type") == "external_session_status" and b.get("data", {}).get("status") == "idle"
+    ]
+
+
+@pytest.mark.asyncio
+class TestForwardLoop:
+    async def test_posts_cumulative_usage_and_persists(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        usage.record_usage_payload(tmp_path, _TURN1)
+        usage.record_usage_payload(tmp_path, _TURN2)
+        client = await _run_loop_until(
+            monkeypatch,
+            tmp_path,
+            lambda c: _usage_posts(c) and usage._read_usage_state(tmp_path).seen == {"g1", "g2"},
+        )
+        url, body = _usage_posts(client)[0]
+        assert url == "/v1/sessions/conv_1/events"
+        assert body["type"] == "external_session_usage"
+        assert body["data"] == {
+            "cumulative_input_tokens": 23666 + 24010,
+            "cumulative_output_tokens": 5 + 120,
+            "cumulative_cache_read_input_tokens": 23617 + 23700,
+            "model": "claude-4-sonnet",
+        }
+        # State persisted after the successful POST so a restart resumes.
+        persisted = usage._read_usage_state(tmp_path)
+        assert persisted.seen == {"g1", "g2"}
+
+    async def test_no_repost_when_unchanged(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        usage.record_usage_payload(tmp_path, _TURN1)
+        client = await _run_loop_until(monkeypatch, tmp_path, _usage_posts)
+        # Let several more polls run; with no new turns there must be no 2nd
+        # usage POST and no further idle edge.
+        await asyncio.sleep(0.1)
+        assert len(_usage_posts(client)) == 1
+        assert len(_idle_posts(client)) == 1
+
+    async def test_new_turn_triggers_followup_post(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        usage.record_usage_payload(tmp_path, _TURN1)
+        client = _CtxRecordingClient()
+        # First usage POST covers turn 1.
+        await _run_loop_until(monkeypatch, tmp_path, _usage_posts, client=client)
+        # Append a second turn and run again: a fresh cumulative POST must land.
+        usage.record_usage_payload(tmp_path, _TURN2)
+        await _run_loop_until(
+            monkeypatch, tmp_path, lambda c: len(_usage_posts(c)) >= 2, client=client
+        )
+        _, body = _usage_posts(client)[-1]
+        assert body["data"]["cumulative_output_tokens"] == 5 + 120
+
+    async def test_completed_turn_posts_idle_wake_edge(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """A completed cursor turn posts external_session_status: idle.
+
+        This is the edge the runner turns into a parent-inbox wake for a cursor
+        sub-agent. Without it, a cursor child finishes with its review only in
+        its own transcript and the parent orchestrator is never woken.
+        """
+        usage.record_usage_payload(tmp_path, _TURN1)
+        client = await _run_loop_until(monkeypatch, tmp_path, _idle_posts)
+        url, body = _idle_posts(client)[0]
+        assert url == "/v1/sessions/conv_1/events"
+        assert body == {"type": "external_session_status", "data": {"status": "idle"}}
+
+    async def test_idle_posted_once_per_turn(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Each completed turn wakes the parent exactly once within a run."""
+        usage.record_usage_payload(tmp_path, _TURN1)
+        client = _CtxRecordingClient()
+        await _run_loop_until(monkeypatch, tmp_path, _idle_posts, client=client)
+        await asyncio.sleep(0.1)  # extra polls: turn 1 must not re-wake
+        assert len(_idle_posts(client)) == 1
+        # A second turn yields exactly one more idle edge.
+        usage.record_usage_payload(tmp_path, _TURN2)
+        await _run_loop_until(
+            monkeypatch, tmp_path, lambda c: len(_idle_posts(c)) >= 2, client=client
+        )
+        await asyncio.sleep(0.1)
+        assert len(_idle_posts(client)) == 2
+
+    async def test_restart_reposts_idle_for_dedup_not_skip(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """A restart re-posts idle (server dedupes) rather than risk skipping a wake.
+
+        The idle counter is NOT seeded from persisted usage: if it were, a wake
+        whose idle POST crashed after the usage flush persisted would be skipped
+        forever. Seeding at 0 makes a restart re-post at most one idle, which the
+        server treats as an already-delivered no-op.
+        """
+        # First run: turn 1 usage persisted.
+        usage.record_usage_payload(tmp_path, _TURN1)
+        first = _CtxRecordingClient()
+        await _run_loop_until(monkeypatch, tmp_path, _idle_posts, client=first)
+        assert usage._read_usage_state(tmp_path).seen == {"g1"}
+
+        # Fresh loop (simulated restart) over the SAME bridge dir, no new turns:
+        # idle is re-posted so a lost pre-crash wake still lands.
+        second = _CtxRecordingClient()
+        await _run_loop_until(monkeypatch, tmp_path, _idle_posts, client=second)
+        await asyncio.sleep(0.1)
+        # Exactly one re-post — the single already-seen turn, not a loop.
+        assert len(_idle_posts(second)) == 1
+
+    async def test_failed_post_is_not_persisted(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        class _FailingClient(_CtxRecordingClient):
+            async def post(self, url: str, *, json: dict) -> httpx.Response:
+                self.posts.append((url, json))
+                raise httpx.ConnectError("boom")
+
+        usage.record_usage_payload(tmp_path, _TURN1)
+        client = await _run_loop_until(
+            monkeypatch, tmp_path, lambda c: len(c.posts) >= 1, client=_FailingClient()
+        )
+        assert len(client.posts) >= 1  # it tried
+        # A failed flush must NOT advance persisted state — the turn stays unseen
+        # so the next poll retries it (no silent loss).
+        assert usage._read_usage_state(tmp_path).seen == set()