cowork-local/tests/unit/test_usage_sink.py

"""Unit tests for :mod:`infrastructure.telemetry.usage_sink` (R03-T06).

Two things are being protected here:

1. The **numbers do not change**. Extracting usage recording out of the two
   providers is only safe if the events built from each wire format carry
   exactly what ``core.usage_tracker.record`` used to receive - a silent change
   would corrupt the Dashboard's cost history.
2. The **sink can never break a turn**. Telemetry is observability; a broken
   store must be swallowed (and logged), never raised into a chat turn.
"""
from __future__ import annotations

import json
from typing import Any, Dict, List

import pytest

from cowork_local.infrastructure.telemetry import usage_sink as telemetry
from cowork_local.providers.anthropic import AnthropicProvider
from cowork_local.providers.base import Provider
from cowork_local.providers.openai_compat import OpenAICompatProvider


# --------------------------------------------------------------------------- #
# Event construction - one per wire format
# --------------------------------------------------------------------------- #
def test_openai_usage_block_maps_onto_the_canonical_event():
    event = telemetry.openai_usage_event("openai_compat", "gpt-4o-mini", {
        "prompt_tokens": 120,
        "completion_tokens": 45,
        "prompt_tokens_details": {"cached_tokens": 100},
    })

    assert (event.input_tokens, event.output_tokens, event.cached_tokens) == (120, 45, 100)
    assert event.estimated is False
    # Cached tokens are a SUBSET of input, so adding them would double-count.
    assert event.total_tokens == 165


def test_anthropic_usage_accumulator_maps_onto_the_canonical_event():
    """Anthropic reports input on message_start and output on message_delta, so
    providers/anthropic.py accumulates them into in/out/cache keys."""
    event = telemetry.anthropic_usage_event("anthropic", "claude", {
        "in": 200, "out": 80, "cache": 150,
    })

    assert (event.input_tokens, event.output_tokens, event.cached_tokens) == (200, 80, 150)
    assert event.estimated is False


def test_a_missing_usage_block_produces_an_estimated_event():
    event = telemetry.estimated_event("ollama", "llama3.1", "x" * 400, "y" * 40)

    assert event.estimated is True
    assert event.input_tokens == 100      # ~4 characters per token
    assert event.output_tokens == 10
    assert event.cached_tokens == 0


def test_estimation_matches_the_legacy_tracker_formula():
    """The extraction must not shift a single recorded number, so the estimator
    is pinned against the one it replaced."""
    from cowork_local.core import usage_tracker

    for text in ("", "short", "x" * 4001, "unicode - tiếng Việt"):
        assert telemetry.estimate_tokens(text) == usage_tracker.estimate_tokens(text)


# --------------------------------------------------------------------------- #
# Sinks
# --------------------------------------------------------------------------- #
def test_recording_sink_collects_events_for_assertions():
    sink = telemetry.RecordingUsageSink()

    sink.record(telemetry.UsageEvent("p", "m", input_tokens=10, output_tokens=5))
    sink.record(telemetry.UsageEvent("p", "m", input_tokens=1, output_tokens=1))

    assert len(sink.events) == 2
    assert sink.total_tokens == 17


def test_null_sink_discards_without_error():
    telemetry.NullUsageSink().record(telemetry.UsageEvent("p", "m"))


def test_tracker_sink_forwards_every_field_positionally():
    """``core.usage_tracker.record`` takes positional counts plus an ``estimated``
    keyword; the adapter has to preserve that exact call shape."""
    seen: Dict[str, Any] = {}

    class _Tracker:
        @staticmethod
        def record(provider, model, input_tokens, output_tokens, cached_tokens,
                   estimated=False):
            # Fields captured explicitly rather than via locals(), which would
            # also drag in the closed-over `seen` binding itself.
            seen.update({"provider": provider, "model": model,
                         "input_tokens": input_tokens, "output_tokens": output_tokens,
                         "cached_tokens": cached_tokens, "estimated": estimated})

    telemetry.UsageTrackerSink(tracker=_Tracker()).record(
        telemetry.UsageEvent("anthropic", "claude", 7, 3, 2, estimated=True))

    assert seen == {"provider": "anthropic", "model": "claude", "input_tokens": 7,
                    "output_tokens": 3, "cached_tokens": 2, "estimated": True}


def test_a_failing_tracker_never_raises_into_the_turn():
    class _Broken:
        @staticmethod
        def record(*_args, **_kwargs):
            raise OSError("usage store is read-only")

    # Must not raise - the turn that produced this event has already succeeded.
    telemetry.UsageTrackerSink(tracker=_Broken()).record(telemetry.UsageEvent("p", "m"))


def test_set_default_sink_returns_the_previous_one_for_restoration():
    replacement = telemetry.RecordingUsageSink()

    previous = telemetry.set_default_sink(replacement)
    try:
        assert telemetry.default_sink is replacement
    finally:
        telemetry.set_default_sink(previous)
    assert telemetry.default_sink is previous


# --------------------------------------------------------------------------- #
# Provider integration - the seam actually being used
# --------------------------------------------------------------------------- #
class _StubResponse:
    """The few members the provider streaming loop touches."""

    def __init__(self, lines: List[str]) -> None:
        self._lines = lines
        self.status_code = 200
        self.headers: Dict[str, str] = {}
        self.encoding = "utf-8"
        self.text = ""

    def iter_lines(self, decode_unicode: bool = False):
        yield from self._lines

    def close(self) -> None:
        return None


@pytest.fixture
def sink(monkeypatch):
    """A per-instance recording sink, so nothing touches the real usage store."""
    return telemetry.RecordingUsageSink()


@pytest.fixture
def canned(monkeypatch):
    def _install(lines: List[str]):
        monkeypatch.setattr(Provider, "_request",
                            lambda self, method, url, **kw: _StubResponse(lines))
    return _install


def test_openai_provider_reports_server_counts_to_its_sink(canned, sink):
    canned([
        'data: ' + json.dumps({"choices": [{"delta": {"content": "hi"}}],
                               "usage": {"prompt_tokens": 11, "completion_tokens": 2,
                                         "prompt_tokens_details": {"cached_tokens": 4}}}),
        "data: [DONE]",
    ])
    provider = OpenAICompatProvider({"base_url": "https://x.invalid/v1",
                                     "api_key": "k", "model": "gpt-4o-mini"})
    provider.usage_sink = sink

    provider.chat([{"role": "user", "content": "hi"}])

    assert len(sink.events) == 1
    event = sink.events[0]
    assert (event.input_tokens, event.output_tokens, event.cached_tokens) == (11, 2, 4)
    assert event.estimated is False
    assert event.model == "gpt-4o-mini"


def test_openai_provider_estimates_when_the_gateway_sends_no_usage(canned, sink):
    canned(['data: ' + json.dumps({"choices": [{"delta": {"content": "hello"}}]}),
            "data: [DONE]"])
    provider = OpenAICompatProvider({"base_url": "https://x.invalid/v1",
                                     "api_key": "k", "model": "m"})
    provider.usage_sink = sink

    provider.chat([{"role": "user", "content": "hi"}])

    assert sink.events[0].estimated is True
    assert sink.events[0].output_tokens >= 1


def test_anthropic_provider_reports_stream_counts_to_its_sink(canned, sink):
    canned(['data: ' + json.dumps(p) for p in (
        {"type": "message_start", "message": {"usage": {"input_tokens": 30,
                                                        "cache_read_input_tokens": 10}}},
        {"type": "content_block_delta", "index": 0,
         "delta": {"type": "text_delta", "text": "ok"}},
        {"type": "message_delta", "usage": {"output_tokens": 5}},
        {"type": "message_stop"},
    )])
    provider = AnthropicProvider({"base_url": "https://x.invalid", "api_key": "k",
                                  "model": "claude"})
    provider.usage_sink = sink

    provider.chat([{"role": "user", "content": "hi"}])

    event = sink.events[0]
    assert (event.input_tokens, event.output_tokens, event.cached_tokens) == (30, 5, 10)
    assert event.estimated is False


def test_a_provider_without_an_explicit_sink_uses_the_process_default(canned):
    """Existing call sites set no sink, so the default has to keep working -
    that is what makes this extraction a no-op for production behaviour."""
    canned(['data: ' + json.dumps({"choices": [{"delta": {"content": "x"}}]}),
            "data: [DONE]"])
    recorder = telemetry.RecordingUsageSink()
    previous = telemetry.set_default_sink(recorder)
    try:
        OpenAICompatProvider({"base_url": "https://x.invalid/v1", "api_key": "k",
                              "model": "m"}).chat([{"role": "user", "content": "hi"}])
    finally:
        telemetry.set_default_sink(previous)

    assert len(recorder.events) == 1


def test_a_sink_that_raises_does_not_fail_the_turn(canned):
    """The answer has already been produced by the time usage is recorded;
    losing the telemetry is strictly better than losing the answer."""
    canned(['data: ' + json.dumps({"choices": [{"delta": {"content": "x"}}]}),
            "data: [DONE]"])

    class _Exploding:
        def record(self, _event):
            raise RuntimeError("sink is down")

    provider = OpenAICompatProvider({"base_url": "https://x.invalid/v1",
                                     "api_key": "k", "model": "m"})
    provider.usage_sink = _Exploding()

    result = provider.chat([{"role": "user", "content": "hi"}])

    assert result["content"] == "x"