cowork-local/core/usage_tracker.py

"""Token-usage tracking for the Dashboard tab.

Every provider turn records one event (JSON line, one file per day under
``~/.cowork_local/usage/``): when, which tab/task ("source" + "label"),
provider/model, input/output/cached token counts. Real counts come from the
server's ``usage`` block when the stream includes one; otherwise a ~4 chars ≈
1 token estimate keeps the dashboard useful on gateways that never report
usage (events carry ``"estimated": true`` so the UI can say so).

The turn's source/label is set by the caller ON THE WORKER THREAD via
:func:`set_context` (thread-local — concurrent turns don't mix labels).
"""
from __future__ import annotations

# Giữ đường vào cũ: nhiều nơi import mấy tên này thẳng từ usage_tracker.
from .usage_ai_report import build_ai_analysis_prompt  # noqa: F401
from .usage_cost import (  # noqa: F401
    DEFAULT_PRICING, SUPPORTED_CURRENCIES, cost_usd, cost_usd_events,
    format_cost, format_cost_compact,
)
from .usage_periods import (  # noqa: F401
    bucketed_series, period_bounds, period_breakdown, period_range_label,
    period_totals, period_window,
)

import json
import threading
from datetime import date, datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

from ..config import CONFIG_DIR
from . import model_pricing as mp

USAGE_DIR = CONFIG_DIR / "usage"

_local = threading.local()

# Process-global identity (NOT thread-local — who's logged in and which
# machine this is are fixed for the whole process, set once right after
# login in app.py::run(), unlike source/label which vary per worker turn).
_identity_account = ""
_identity_machine = ""
_identity_shared_dir = ""


def set_identity(account: str, machine: str, shared_dir: str = "") -> None:
    """Called once after login succeeds. ``shared_dir``, when reachable,
    makes every subsequent :func:`record` ALSO best-effort-append to the
    shared cross-machine telemetry store (see :mod:`telemetry_shared`)."""
    global _identity_account, _identity_machine, _identity_shared_dir
    _identity_account = account or ""
    _identity_machine = machine or ""
    _identity_shared_dir = shared_dir or ""


def set_context(source: str, label: str = "") -> None:
    """Tag subsequent :func:`record` calls on THIS thread (e.g. ("cowork",
    "chat title") / ("task", "task title"))."""
    _local.source = source
    _local.label = label


def current_context() -> tuple:
    """The ``(source, label)`` currently tagged on THIS thread.

    Public counterpart to :func:`set_context`, added for
    ``infrastructure/telemetry/usage_sink.py``: a subscriber that needs to
    attribute one event to a different surface must be able to save the
    caller's context and put it back afterwards, instead of leaving the worker
    thread permanently retagged.
    """
    return getattr(_local, "source", "") or "", getattr(_local, "label", "") or ""


# ---- per-thread usage accumulator -----------------------------------------
# A step/run that wants to know its OWN token/cost (not the all-time file total)
# calls begin_accumulation(), reads accumulated() before/after a unit of work,
# and diffs the two. Because record() runs on the same worker thread that drives
# the work (providers are called synchronously inside it), the thread-local
# total is exactly that thread's usage — concurrent flows on other threads
# accumulate independently, with no locking or label collisions. Used by the
# Co4E runner to attach per-step token/cost to each node's output event.
def begin_accumulation() -> None:
    """Start (or reset) this thread's usage accumulator."""
    _local.acc = {"in": 0, "out": 0, "cache": 0, "events": []}


def accumulated() -> Dict[str, Any]:
    """Snapshot of this thread's accumulated usage since :func:`begin_accumulation`
    (all zeros / empty if never started). ``events`` is a per-turn list of
    ``{model, in, out}`` so a caller can price a delta with the per-model table."""
    acc = getattr(_local, "acc", None)
    if acc is None:
        return {"in": 0, "out": 0, "cache": 0, "events": []}
    return {"in": acc["in"], "out": acc["out"], "cache": acc["cache"],
            "events": list(acc["events"])}


def end_accumulation() -> None:
    """Stop accumulating on this thread (subsequent records aren't tallied)."""
    _local.acc = None


def estimate_tokens(text: str) -> int:
    return max(0, len(text or "") // 4)


def record(provider: str, model: str, input_tokens: int, output_tokens: int,
           cached_tokens: int = 0, estimated: bool = False) -> None:
    """Append one usage event. Never raises — usage tracking must never break
    a chat turn."""
    try:
        now = datetime.now()
        event = {
            "ts": now.isoformat(timespec="seconds"),
            "source": getattr(_local, "source", "") or "other",
            "label": getattr(_local, "label", "") or "",
            "provider": provider or "",
            "model": model or "",
            "in": int(input_tokens or 0),
            "out": int(output_tokens or 0),
            "cache": int(cached_tokens or 0),
            "estimated": bool(estimated),
            "account": _identity_account,
            "machine": _identity_machine,
        }
        USAGE_DIR.mkdir(parents=True, exist_ok=True)
        path = USAGE_DIR / f"{now.strftime('%Y-%m-%d')}.jsonl"
        with path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(event, ensure_ascii=False) + "\n")
        _write_shared(event, now)
        # Feed this thread's live accumulator, if one is active (see above).
        acc = getattr(_local, "acc", None)
        if acc is not None:
            acc["in"] += event["in"]
            acc["out"] += event["out"]
            acc["cache"] += event["cache"]
            acc["events"].append({"model": event["model"], "in": event["in"], "out": event["out"]})
    except Exception:  # noqa: BLE001
        pass


def _write_shared(event: Dict[str, Any], now: datetime) -> None:
    """Best-effort mirror of ``event`` into the shared cross-machine store —
    one file PER MACHINE per day, so no two machines ever write the same
    file (avoids any read-modify-write race). Never raises."""
    if not _identity_shared_dir or not _identity_machine:
        return
    try:
        shared = Path(_identity_shared_dir).expanduser() / "telemetry" / "usage"
        shared.mkdir(parents=True, exist_ok=True)
        path = shared / f"{_identity_machine}-{now.strftime('%Y-%m-%d')}.jsonl"
        with path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(event, ensure_ascii=False) + "\n")
    except Exception:  # noqa: BLE001
        pass


def load_events(start: Optional[date] = None, end: Optional[date] = None,
                directory: Path = None) -> List[Dict[str, Any]]:
    """Events between ``start`` and ``end`` (inclusive; None = unbounded)."""
    directory = directory or USAGE_DIR
    if not directory.exists():
        return []
    events: List[Dict[str, Any]] = []
    for path in sorted(directory.glob("*.jsonl")):
        try:
            day = datetime.strptime(path.stem, "%Y-%m-%d").date()
        except ValueError:
            continue
        if (start and day < start) or (end and day > end):
            continue
        try:
            for line in path.read_text(encoding="utf-8").splitlines():
                if line.strip():
                    events.append(json.loads(line))
        except (OSError, json.JSONDecodeError):
            continue
    return events


def summarize(events: List[Dict[str, Any]]) -> Dict[str, Any]:
    """Aggregate a list of events into dashboard numbers + habit stats."""
    total_in = sum(e.get("in", 0) for e in events)
    total_out = sum(e.get("out", 0) for e in events)
    total_cache = sum(e.get("cache", 0) for e in events)
    by_label: Dict[str, int] = {}
    by_source: Dict[str, int] = {}
    by_hour: Dict[int, int] = {}
    by_day: Dict[str, int] = {}
    for e in events:
        tok = e.get("in", 0) + e.get("out", 0)
        key = e.get("label") or e.get("source") or "?"
        by_label[key] = by_label.get(key, 0) + tok
        by_source[e.get("source", "?")] = by_source.get(e.get("source", "?"), 0) + tok
        try:
            dt = datetime.fromisoformat(e.get("ts", ""))
            by_hour[dt.hour] = by_hour.get(dt.hour, 0) + tok
            by_day[dt.strftime("%Y-%m-%d")] = by_day.get(dt.strftime("%Y-%m-%d"), 0) + tok
        except ValueError:
            pass
    return {
        "turns": len(events),
        "in": total_in, "out": total_out, "cache": total_cache,
        "total": total_in + total_out,
        "avg_per_turn": (total_in + total_out) // len(events) if events else 0,
        "estimated_share": (sum(1 for e in events if e.get("estimated")) / len(events)
                            if events else 0.0),
        "top_labels": sorted(by_label.items(), key=lambda kv: -kv[1])[:5],
        "by_source": sorted(by_source.items(), key=lambda kv: -kv[1]),
        "busiest_hour": max(by_hour.items(), key=lambda kv: kv[1])[0] if by_hour else None,
        "busiest_day": max(by_day.items(), key=lambda kv: kv[1])[0] if by_day else None,
    }


# ---- cost ------------------------------------------------------------------


# Currencies the display picker offers — exactly the ones format_cost() can
# actually convert to (symbol/precision above + a usd_to_* rate below).






















def set_budget(config, amount: float, currency: Optional[str] = None) -> None:
    """Set (or reset) the spending budget. ``amount`` is read in ``currency``
    (defaults to the current display currency) and converted + stored as USD.

    Remaining balance is always DERIVED fresh from the usage log — never
    incrementally decremented — so re-entering a budget starts a clean window
    instead of double-subtracting spend the old budget had already accounted
    for. The cutoff is an EVENT-COUNT baseline (how many usage events existed
    at the moment of setting), not a timestamp: events append in chronological
    order and ``budget_set_at`` only has 1-second resolution, so a timestamp
    cutoff could mis-include/exclude an event recorded in that same second —
    the count baseline is exact regardless of timing."""
    usage = config.data.setdefault("usage", {})
    ccy = (currency or usage.get("currency") or "USD").upper()
    usage["budget_amount_usd"] = mp.convert(float(amount or 0), ccy, "USD", config)
    usage["budget_set_at"] = datetime.now().isoformat(timespec="seconds")   # display only
    usage["budget_baseline_count"] = len(load_events())                     # the real cutoff


def clear_budget(config) -> None:
    """Remove the budget entirely (Remaining/Budget box goes back to unset)."""
    usage = config.data.setdefault("usage", {})
    usage.pop("budget_amount_usd", None)
    usage.pop("budget_set_at", None)
    usage.pop("budget_baseline_count", None)


def budget_status(config) -> Optional[Dict[str, Any]]:
    """``None`` when no budget is configured. Else a dict with ``amount_usd``,
    ``spent_usd`` (cost of events recorded AFTER the budget was last set — NOT
    the all-time total, so a reset budget never inherits older spend),
    ``remaining_usd``, ``pct_used`` and ``over_85`` (⚠ the Overview/Dashboard
    balance turns red at this point)."""
    usage = (getattr(config, "data", {}) or {}).get("usage") or {}
    amount = usage.get("budget_amount_usd")
    set_at = usage.get("budget_set_at")
    if not amount or not set_at:
        return None
    all_events = load_events()
    baseline = usage.get("budget_baseline_count")
    if baseline is None:
        # backward-compat: a budget set before this field existed — fall back
        # to the timestamp cutoff (best-effort, may double-count a same-second event).
        events = [e for e in all_events if str(e.get("ts", "")) >= str(set_at)]
    else:
        events = all_events[int(baseline):]
    pricing = {**DEFAULT_PRICING, **usage}
    spent = sum(cost_usd_events(events, pricing).values())
    amount = float(amount)
    pct = (spent / amount) if amount else 0.0
    return {
        "amount_usd": amount, "spent_usd": spent, "remaining_usd": amount - spent,
        "pct_used": pct, "over_85": pct >= 0.85, "set_at": set_at,
    }