src/services/billing/usage_watcher.py

"""SSE usage watcher — meters streaming responses by capturing the usage chunk.

OpenAI-compatible streams (incl. OpenRouter :free) emit a final chunk carrying a
`usage` object when the client sets `stream_options.include_usage: true`. This
scanner watches the filtered byte stream as it passes through and extracts
(prompt_tokens, completion_tokens) from that final chunk.

If the upstream doesn't emit usage (e.g. raw subnet nodes), tok_in/tok_out stay
0 and the request is effectively unmetered — acceptable for MVP since upstreams
are free and we control the margin.
"""
from __future__ import annotations

import json


class UsageWatcher:
    """Scan SSE bytes for the OpenAI usage object. Buffers across chunk boundaries."""

    def __init__(self):
        self.tok_in = 0   # captured prompt_tokens (filled when the usage chunk arrives)
        self.tok_out = 0   # captured completion_tokens (filled when the usage chunk arrives)
        self._buf = b""   # leftover bytes that didn't end on a newline

    def feed(self, chunk: bytes) -> None:
        """Scan one chunk of SSE bytes for a usage-bearing data line."""
        self._buf += chunk   # append to the carryover buffer
        while b"\n" in self._buf:   # process complete lines only
            line, self._buf = self._buf.split(b"\n", 1)   # peel one line
            self._scan_line(line)   # check it for usage

    def flush(self) -> None:
        """Process any remaining buffered bytes after the stream ends."""
        if self._buf.strip():   # there's a trailing line without a newline
            self._scan_line(self._buf)   # check it too
        self._buf = b""   # clear the buffer

    def _scan_line(self, line: bytes) -> None:
        """Inspect one SSE line; if it's a data: line with usage, capture it."""
        line = line.strip()   # drop trailing \r and whitespace
        if not line.startswith(b"data:"):   # only data: lines carry JSON payloads
            return   # skip comments, event:, id:, etc.
        if b'"usage"' not in line:   # cheap pre-check before JSON parsing
            return   # most chunks have no usage — avoid the parse cost
        payload = line[5:].strip()   # strip the "data:" prefix
        if payload == b"[DONE]":   # stream terminator, not JSON
            return   # ignore
        try:
            obj = json.loads(payload)   # parse the SSE event JSON
        except Exception:   # malformed JSON (shouldn't happen post-filter, but be safe)
            return   # skip
        if not isinstance(obj, dict):   # not an object
            return   # skip
        u = obj.get("usage")   # the usage sub-object (OpenAI shape)
        if isinstance(u, dict):   # this chunk carries usage
            self.tok_in = int(u.get("prompt_tokens", self.tok_in))   # capture input tokens
            self.tok_out = int(u.get("completion_tokens", self.tok_out))   # capture output tokens

    @property
    def has_usage(self) -> bool:
        """True if we captured a usage object during the stream."""
        return self.tok_in > 0 or self.tok_out > 0