"""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: """Ước lượng số token của một đoạn văn bản theo tỉ lệ 4 ký tự ≈ 1 token. Ước lượng thô là đủ: con số này chỉ dùng để quyết định khi nào nén lịch sử, không dùng để tính tiền (tiền lấy từ số token thật provider trả về). """ 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, }