cowork-local/infrastructure/telemetry/usage_sink.py

"""UsageEventSink - where a turn's token usage goes (R03-T06).

Today each provider records its own usage inline, in the middle of the streaming
loop::

    # providers/openai_compat.py
    def _record_usage(self, messages, text_parts, tool_acc, usage_seen):
        from ..core import usage_tracker as ut
        ...
        ut.record(self.name, self.model, ...)

Three problems with that shape:

1. **Hidden side effect.** ``chat()`` looks like a pure request/response call but
   also writes to the Dashboard's store, so a test of a provider silently
   appends rows to the developer's real usage history.
2. **Duplicated estimation.** The "no usage block from the server, so estimate
   at ~4 chars/token" fallback is copy-pasted per provider and can drift.
3. **One hard-wired destination.** Usage can only ever go to
   ``core.usage_tracker``; a run that wants to bill a workflow, or a test that
   wants to assert on token counts, has nowhere to plug in.

This module introduces the seam: providers build a :class:`UsageEvent` and hand
it to a :class:`UsageEventSink`. Production wires :class:`UsageTrackerSink`
(same destination, same numbers as before); tests wire
:class:`RecordingUsageSink` or :class:`NullUsageSink`.
"""
from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Protocol, Sequence

logger = logging.getLogger("cowork_local.telemetry")

# Rough characters-per-token ratio used when the gateway sends no usage block.
# Matches the constant behaviour of ``core.usage_tracker.estimate_tokens`` so
# moving the estimation here does not change a single recorded number.
_CHARS_PER_TOKEN = 4


@dataclass(frozen=True)
class UsageEvent:
    """Token usage for exactly one provider round trip.

    ``estimated`` marks a record derived from text length rather than reported by
    the server. The Dashboard shows the two differently, and conflating them
    would make cost figures look more precise than they are.
    """

    provider: str
    model: str
    input_tokens: int = 0
    output_tokens: int = 0
    cached_tokens: int = 0
    estimated: bool = False

    @property
    def total_tokens(self) -> int:
        """Input + output. Cached tokens are a subset of input, not an addition,
        so adding them here would double-count a cache hit."""
        return self.input_tokens + self.output_tokens

    def to_dict(self) -> Dict[str, Any]:
        """JSON-safe projection for logs and for sinks that persist raw events."""
        return {
            "provider": self.provider,
            "model": self.model,
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "cached_tokens": self.cached_tokens,
            "estimated": self.estimated,
        }


class UsageEventSink(Protocol):
    """Anything that can absorb a :class:`UsageEvent`.

    Implementations MUST NOT raise: telemetry is observability, and a failure to
    record usage must never abort the turn that produced it.
    """

    def record(self, event: UsageEvent) -> None:
        """Absorb one usage event."""


class NullUsageSink:
    """Discards everything. The default for tests and headless tooling, so a
    unit test never writes into the developer's real usage history."""

    def record(self, event: UsageEvent) -> None:  # noqa: D102 - see protocol
        return None


class RecordingUsageSink:
    """Keeps events in memory so a test can assert on what was recorded."""

    def __init__(self) -> None:
        self.events: List[UsageEvent] = []

    def record(self, event: UsageEvent) -> None:  # noqa: D102 - see protocol
        self.events.append(event)

    @property
    def total_tokens(self) -> int:
        """Sum across every recorded event."""
        return sum(e.total_tokens for e in self.events)


class UsageTrackerSink:
    """Forwards to ``core.usage_tracker`` - the Dashboard's store.

    This is the production sink and the only place that still knows about the
    legacy tracker module, which is what lets EPIC R10 replace the storage
    without touching a single provider.
    """

    def __init__(self, tracker: Optional[Any] = None) -> None:
        # Injectable for tests; imported lazily otherwise because the tracker
        # touches the config directory at import time.
        self._tracker = tracker

    def _resolve(self) -> Any:
        if self._tracker is None:
            from cowork_local.core import usage_tracker

            self._tracker = usage_tracker
        return self._tracker

    def record(self, event: UsageEvent) -> None:
        """Write the event to the usage tracker, swallowing any failure.

        The bare except mirrors the behaviour this replaces (each provider
        already wrapped its ``ut.record`` call in ``try/except: pass``) but logs
        at debug level instead of discarding the reason entirely, so a broken
        Dashboard store can at least be diagnosed.
        """
        try:
            self._resolve().record(
                event.provider, event.model,
                event.input_tokens, event.output_tokens, event.cached_tokens,
                estimated=event.estimated,
            )
        except Exception:  # noqa: BLE001 - telemetry must never break a turn
            logger.debug("usage sink: failed to record %s", event.to_dict(), exc_info=True)


def estimate_tokens(text: str) -> int:
    """Approximate token count for ``text`` (~4 characters per token).

    Deliberately identical to ``core.usage_tracker.estimate_tokens`` so that
    moving estimation into this layer changes no recorded number. Duplicated
    rather than imported to keep this module free of the legacy dependency;
    :class:`UsageTrackerSink` is the only bridge back to it.
    """
    return max(0, len(text or "") // _CHARS_PER_TOKEN)


def estimated_event(provider: str, model: str, sent: str, received: str) -> UsageEvent:
    """Build an estimated :class:`UsageEvent` from the raw text of a round trip.

    Used when the gateway sends no usage block - most self-hosted OpenAI-compatible
    servers and Ollama do not.
    """
    return UsageEvent(
        provider=provider, model=model,
        input_tokens=estimate_tokens(sent),
        output_tokens=estimate_tokens(received),
        cached_tokens=0,
        estimated=True,
    )


def openai_usage_event(provider: str, model: str, usage: Dict[str, Any]) -> UsageEvent:
    """Build a reported :class:`UsageEvent` from an OpenAI-style usage block."""
    details = usage.get("prompt_tokens_details") or {}
    return UsageEvent(
        provider=provider, model=model,
        input_tokens=int(usage.get("prompt_tokens", 0) or 0),
        output_tokens=int(usage.get("completion_tokens", 0) or 0),
        cached_tokens=int(details.get("cached_tokens", 0) or 0),
        estimated=False,
    )


def anthropic_usage_event(provider: str, model: str, usage: Dict[str, Any]) -> UsageEvent:
    """Build a reported :class:`UsageEvent` from Anthropic's usage accumulator.

    Anthropic reports input tokens on ``message_start`` and output tokens on
    ``message_delta``, so ``providers/anthropic.py`` accumulates them into a dict
    keyed ``in``/``out``/``cache`` - this reads that shape.
    """
    return UsageEvent(
        provider=provider, model=model,
        input_tokens=int(usage.get("in", 0) or 0),
        output_tokens=int(usage.get("out", 0) or 0),
        cached_tokens=int(usage.get("cache", 0) or 0),
        estimated=False,
    )


# The sink providers use unless one is injected. A module-level default keeps
# the change to the provider classes to a single attribute, and lets a test swap
# the destination process-wide with one monkeypatch.
default_sink: UsageEventSink = UsageTrackerSink()


def set_default_sink(sink: UsageEventSink) -> UsageEventSink:
    """Replace the process-wide default sink; returns the previous one so a
    caller (or fixture) can restore it."""
    global default_sink
    previous = default_sink
    default_sink = sink
    return previous


__all__ = [
    "UsageEvent",
    "UsageEventSink",
    "UsageTrackerSink",
    "NullUsageSink",
    "RecordingUsageSink",
    "estimate_tokens",
    "estimated_event",
    "openai_usage_event",
    "anthropic_usage_event",
    "default_sink",
    "set_default_sink",
]