"""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", ]