tracer-cloud--opensre/tools/system/fleet_monitoring/meters/__init__.py

"""Token meters: stateless parsers extracting token counts from CLI output.

Cost calculation is deliberately separate — per-token rates change
per model and per counter (cache reads at 0.1×, cache writes at
1.25×), so binding cost to the parser would couple ``tokens/min``
to ``$/hr`` in a way that grows brittle.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Protocol


@dataclass(frozen=True)
class TokenUsage:
    """Structured token counters extracted from one meter sample.

    ``cached_input_tokens`` is a discounted subset of ``input_tokens``
    for Codex/OpenAI-style usage. Claude's cache counters are emitted
    as separate read/write fields and are additive for display.
    """

    input_tokens: float = 0
    output_tokens: float = 0
    cached_input_tokens: float = 0
    cache_read_input_tokens: float = 0
    cache_creation_input_tokens: float = 0

    @classmethod
    def from_total(cls, tokens: float) -> TokenUsage:
        return cls(input_tokens=max(0.0, tokens))

    @property
    def tokens(self) -> float:
        """Visible total for the dashboard's ``tokens/min`` cell."""
        return (
            max(0.0, self.input_tokens)
            + max(0.0, self.output_tokens)
            + max(0.0, self.cache_read_input_tokens)
            + max(0.0, self.cache_creation_input_tokens)
        )

    def clamped(self) -> TokenUsage:
        return TokenUsage(
            input_tokens=max(0.0, self.input_tokens),
            output_tokens=max(0.0, self.output_tokens),
            cached_input_tokens=max(0.0, self.cached_input_tokens),
            cache_read_input_tokens=max(0.0, self.cache_read_input_tokens),
            cache_creation_input_tokens=max(0.0, self.cache_creation_input_tokens),
        )

    def scaled(self, factor: float) -> TokenUsage:
        return TokenUsage(
            input_tokens=self.input_tokens * factor,
            output_tokens=self.output_tokens * factor,
            cached_input_tokens=self.cached_input_tokens * factor,
            cache_read_input_tokens=self.cache_read_input_tokens * factor,
            cache_creation_input_tokens=self.cache_creation_input_tokens * factor,
        )

    def __add__(self, other: TokenUsage) -> TokenUsage:
        return TokenUsage(
            input_tokens=self.input_tokens + other.input_tokens,
            output_tokens=self.output_tokens + other.output_tokens,
            cached_input_tokens=self.cached_input_tokens + other.cached_input_tokens,
            cache_read_input_tokens=self.cache_read_input_tokens + other.cache_read_input_tokens,
            cache_creation_input_tokens=(
                self.cache_creation_input_tokens + other.cache_creation_input_tokens
            ),
        )


@dataclass(frozen=True)
class TokenSample:
    """A meter's read of a stdout chunk: usage buckets + optional model hint."""

    usage: TokenUsage = field(default_factory=TokenUsage)
    model: str | None = None

    @classmethod
    def from_tokens(cls, tokens: float, model: str | None = None) -> TokenSample:
        return cls(usage=TokenUsage.from_total(tokens).clamped(), model=model)

    @property
    def tokens(self) -> int:
        return int(self.usage.tokens)


class TokenMeter(Protocol):
    """A token-count parser over a CLI stdout chunk.

    Implementations must be safe to call with partial chunks — chunks
    coming from a streaming subprocess split at arbitrary byte offsets
    and may not align with line or JSON-document boundaries.
    """

    def parse_chunk(self, chunk: str, /) -> int:
        """Return newly observed token count from ``chunk``; must not mutate any per-PID parser state."""
        raise NotImplementedError

    def sample_chunk(self, chunk: str, /, *, pid: int | None = None) -> TokenSample:
        """Return parsed usage/model data for ``chunk`` and optional ``pid`` without negative counts."""
        raise NotImplementedError

    def forget(self, pid: int, /) -> None:
        """Drop any parser state for ``pid`` so future samples start from a clean stream boundary."""
        raise NotImplementedError

    def known_pids(self) -> list[int]:
        """Return PIDs with retained parser state; callers may use this for cleanup bookkeeping."""
        raise NotImplementedError


class NullMeter:
    """Always returns 0 / ``None``."""

    def parse_chunk(self, _chunk: str, /) -> int:
        return 0

    def sample_chunk(self, _chunk: str, /, *, pid: int | None = None) -> TokenSample:  # noqa: ARG002
        return TokenSample()

    def forget(self, _pid: int, /) -> None:
        return None

    def known_pids(self) -> list[int]:
        return []


null_meter: TokenMeter = NullMeter()


def safe_int(value: object) -> int:
    """Coerce ``value`` to a non-negative int.

    ``bool`` is rejected explicitly because ``isinstance(True, int)``
    is ``True`` — a stray ``"input_tokens": true`` must not add 1.
    """
    if isinstance(value, bool):
        return 0
    if isinstance(value, int):
        return value if value >= 0 else 0
    return 0


__all__ = [
    "NullMeter",
    "TokenMeter",
    "TokenSample",
    "TokenUsage",
    "null_meter",
    "safe_int",
]