diff --git a/agent/cost_ladder_wire.py b/agent/cost_ladder_wire.py index ec76a0d74e06..bc135a5dfabf 100644 --- a/agent/cost_ladder_wire.py +++ b/agent/cost_ladder_wire.py @@ -82,6 +82,8 @@ def record_inference_from_response( provider: Optional[str] = None, base_url: Optional[str] = None, api_mode: Optional[str] = None, + route: str = "unknown", + escalated_to_opus: bool = False, ) -> None: """Feed the cost-ladder estimator from an inference response. @@ -129,6 +131,8 @@ def record_inference_from_response( model_name=resolved_model, provider=provider, base_url=base_url, + route=route, + escalated_to_opus=escalated_to_opus, ) except Exception as exc: # Fail-soft per the contract — estimator failures must not diff --git a/agent/cost_state_holder.py b/agent/cost_state_holder.py index 52d0320fffb3..902d839427bd 100644 --- a/agent/cost_state_holder.py +++ b/agent/cost_state_holder.py @@ -276,6 +276,8 @@ def record_inference( model_name: str, provider: Optional[str] = None, base_url: Optional[str] = None, + route: str = "unknown", + escalated_to_opus: bool = False, ) -> None: """Per-call estimator update. @@ -305,6 +307,21 @@ def record_inference( base_url: Optional base URL override; used when the SDK is configured against a custom endpoint (provider-fronted Anthropic, etc.). + route: KR-CHEAP-COST-TELEMETRY route label per + ``kora_cli.telemetry.cost_telemetry`` taxonomy + (``slack_dm``, ``email_inbound``, ``mcp_tool``, + ``alert_investigation``, ``probe_investigation``, + ``tool_loop_iteration``, ``scheduled_task``, + ``email_outbound_compose``, or ``unknown``). + Defaults to ``"unknown"`` so existing callers keep + working unchanged; they bucket into the unknown + route until tagged explicitly. Telemetry is a + READ-side observer of pricing; it does NOT affect + whether/how much is billed. + escalated_to_opus: Per-call escalation signal — Lock R3-3 + tunable. When True, the telemetry counters increment + ``escalation_count`` for this route in addition to + the normal call+token counters. """ cost_result = estimate_usage_cost( model_name, @@ -312,6 +329,36 @@ def record_inference( provider=provider, base_url=base_url, ) + + # KR-CHEAP-COST-TELEMETRY: tag the per-route counters. This + # is read-only telemetry — does NOT affect billing accumulation + # below. Fail-soft import + record so an unwired test path or + # an import-cycle scenario can't crash the inference handler. + cost_estimate_for_telemetry: Optional[float] + if cost_result.amount_usd is None: + cost_estimate_for_telemetry = None + else: + try: + cost_estimate_for_telemetry = float(cost_result.amount_usd) + except (TypeError, ValueError): + cost_estimate_for_telemetry = None + try: + from kora_cli.telemetry import get_telemetry + + get_telemetry().record_call( + route=route, + model=model_name, + canonical_usage=canonical_usage, + cost_estimate_usd=cost_estimate_for_telemetry, + escalated_to_opus=escalated_to_opus, + ) + except Exception as exc: + logger.debug( + "[kora.cost_ladder] telemetry record_call failed: %r — " + "billing accumulation continues", + exc, + ) + if cost_result.amount_usd is None: logger.debug( "[kora.cost_ladder] no pricing for model=%s provider=%s " diff --git a/kora_cli/handlers/slack_dm_handler.py b/kora_cli/handlers/slack_dm_handler.py index 7c1ebf860160..34c879caf0ba 100644 --- a/kora_cli/handlers/slack_dm_handler.py +++ b/kora_cli/handlers/slack_dm_handler.py @@ -680,6 +680,13 @@ def _record_inference_to_cost_ladder( ), model_name=str(model_name), provider="anthropic", + # KR-CHEAP-COST-TELEMETRY — tag this Kora reply-bill + # under the slack_dm route. First iteration of a + # tool-use loop is attributed to the originating + # route (here slack_dm); iteration 2+ would attribute + # to ``tool_loop_iteration`` once the reasoning + # engine surfaces that signal (deferred follow-on). + route="slack_dm", ) except Exception as exc: logger.warning( diff --git a/kora_cli/listeners/__init__.py b/kora_cli/listeners/__init__.py index 137eeb383542..75f19ebdaaaf 100644 --- a/kora_cli/listeners/__init__.py +++ b/kora_cli/listeners/__init__.py @@ -65,3 +65,9 @@ # upstream holders + the periodic-task scheduler are guaranteed # registered before the snapshot task gets enqueued. from kora_cli.listeners import snapshot_listener # noqa: F401 +# KR-CHEAP-COST-TELEMETRY — per-route cost-counter persistence + +# window-reset tasks. Imported AFTER snapshot_listener so the +# heartbeat scheduler picks up the snapshot task first (snapshot +# reads telemetry; persist task writes telemetry — order doesn't +# matter for correctness, only for boot-log readability). +from kora_cli.listeners import cost_telemetry_listener # noqa: F401 diff --git a/kora_cli/listeners/cost_telemetry_listener.py b/kora_cli/listeners/cost_telemetry_listener.py new file mode 100644 index 000000000000..d7393d623670 --- /dev/null +++ b/kora_cli/listeners/cost_telemetry_listener.py @@ -0,0 +1,350 @@ +"""Cost-telemetry persistence + window-reset listener — KR-CHEAP-COST-TELEMETRY. + +Registers three periodic tasks with the heartbeat scheduler: + + * ``cost_telemetry.persist`` (default 300s / 5 min) — atomic- + write the current counters to disk so the cockpit + future + routing-layer can read them outside of process memory. + * ``cost_telemetry.rolling_24h_reset`` (default 3600s / 1h tick; + actual reset gated on UTC midnight) — clears the rolling-24h + window once per UTC day. + * ``cost_telemetry.monthly_reset`` (default 3600s tick; gated on + month rollover) — clears the monthly window once per UTC + month boundary. + +The two reset tasks use a "watch + act" pattern: they fire on a +sub-window cadence and check whether the boundary has crossed +since the last reset. This avoids the trickiness of asyncio +scheduling at exact midnight while still guaranteeing the +boundary is honored within a small window of when it occurs. + +# Fail-soft + +All three tasks are wrapped — exceptions log + swallow so the +heartbeat scheduler keeps ticking. Persistence failures are +operator-recoverable (manual disk inspection); reset failures +worst-case leave the counter rolling slightly past its boundary +(operator visible via the snapshot's window timestamps). +""" + +from __future__ import annotations + +import json +import logging +import os +import tempfile +from datetime import datetime, timezone +from pathlib import Path +from typing import Optional + +from kora_cli.daemon import DEFAULT_SHUTDOWN_TIMEOUT, register_daemon_listener +from kora_cli.listeners.heartbeat import register_periodic_task +from kora_cli.telemetry import ( + WINDOW_MONTHLY, + WINDOW_ROLLING_24H, + get_telemetry, +) +from utils import atomic_replace + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Path resolution +# --------------------------------------------------------------------------- + + +COST_TELEMETRY_PATH_ENV = "KORA_COST_TELEMETRY_PATH" +_COST_TELEMETRY_RELATIVE_PATH = Path("cache") / "cost_telemetry.json" + + +def cost_telemetry_path() -> Path: + """Resolve the on-disk telemetry path. Env override first, then + ``${KORA_HOME}/cache/cost_telemetry.json``. + + Mirrors the snapshot module's path-resolution pattern from + PR #157 so monkeypatch in tests works without ContextVar + plumbing. + """ + override = os.environ.get(COST_TELEMETRY_PATH_ENV, "").strip() + if override: + return Path(override) + from kora_constants import get_kora_home + + return get_kora_home() / _COST_TELEMETRY_RELATIVE_PATH + + +# --------------------------------------------------------------------------- +# Cadence config +# --------------------------------------------------------------------------- + + +DEFAULT_PERSIST_INTERVAL_SEC: float = 300.0 # 5 min per spec §2(c) +PERSIST_INTERVAL_ENV: str = "KORA_COST_TELEMETRY_PERSIST_INTERVAL_SEC" + +# Reset tasks tick every hour and check whether the boundary has +# crossed. Smaller cadence (≤ boundary granularity) means +# precision; choosing 1h trades sub-hour precision for cheap +# accounting (the boundary lateness within an hour is operator- +# negligible for the 24h / monthly windows). +DEFAULT_RESET_TICK_INTERVAL_SEC: float = 3600.0 +RESET_TICK_INTERVAL_ENV: str = "KORA_COST_TELEMETRY_RESET_TICK_SEC" + + +def _read_positive_interval(env_name: str, default_value: float) -> float: + raw = os.environ.get(env_name, "").strip() + if not raw: + return default_value + try: + value = float(raw) + except ValueError: + logger.warning( + "[kora.cost_telemetry_listener] %s=%r is not numeric; " + "using default %ss", + env_name, + raw, + default_value, + ) + return default_value + if value <= 0: + logger.warning( + "[kora.cost_telemetry_listener] %s=%s must be > 0; using " + "default %ss", + env_name, + value, + default_value, + ) + return default_value + return value + + +def _read_persist_interval() -> float: + return _read_positive_interval( + PERSIST_INTERVAL_ENV, DEFAULT_PERSIST_INTERVAL_SEC + ) + + +def _read_reset_tick_interval() -> float: + return _read_positive_interval( + RESET_TICK_INTERVAL_ENV, DEFAULT_RESET_TICK_INTERVAL_SEC + ) + + +# --------------------------------------------------------------------------- +# Persistence +# --------------------------------------------------------------------------- + + +def write_telemetry_snapshot() -> None: + """Atomic-write the current telemetry snapshot to disk. + + Same pattern as the daemon snapshot from PR #157: write to a + sibling tmp file then ``atomic_replace`` to the target. Parent + dir created if missing. + """ + target = cost_telemetry_path() + target.parent.mkdir(parents=True, exist_ok=True) + snapshot = get_telemetry().snapshot() + payload = { + "schema_version": 1, + "written_at": datetime.now(timezone.utc).strftime( + "%Y-%m-%dT%H:%M:%SZ" + ), + "windows": snapshot, + } + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + delete=False, + dir=target.parent, + prefix="cost_telemetry.", + suffix=".tmp", + ) as fp: + json.dump(payload, fp, indent=2, sort_keys=True) + fp.write("\n") + tmp_path = fp.name + atomic_replace(tmp_path, target) + + +def read_telemetry_snapshot() -> Optional[dict]: + """Read the on-disk telemetry snapshot. Returns the parsed dict + or ``None`` when missing / unreadable / malformed. + + Same fail-soft posture as :func:`kora_cli.snapshot.read_snapshot`.""" + target = cost_telemetry_path() + if not target.is_file(): + return None + try: + raw = target.read_text(encoding="utf-8") + except OSError as exc: + logger.warning( + "[kora.cost_telemetry] read failed for %s: %r", target, exc + ) + return None + try: + snapshot = json.loads(raw) + except json.JSONDecodeError as exc: + logger.warning( + "[kora.cost_telemetry] malformed snapshot at %s: %r", + target, + exc, + ) + return None + if not isinstance(snapshot, dict): + return None + return snapshot + + +# --------------------------------------------------------------------------- +# Window-reset state +# --------------------------------------------------------------------------- + + +# Track the last UTC date/month on which a reset fired so the +# watch-and-act task fires exactly once per boundary crossing. +# Module-level state is fine — the listener is single-process; if a +# future bucket adds multi-process daemons each one resets its own +# in-memory counters independently and that's the correct shape +# (telemetry is per-process today). +_last_24h_reset_date: Optional[datetime] = None +_last_monthly_reset_month: Optional[tuple] = None # (year, month) + + +def _utc_date_now() -> datetime: + return datetime.now(timezone.utc).replace( + hour=0, minute=0, second=0, microsecond=0 + ) + + +def _utc_month_now() -> tuple: + now = datetime.now(timezone.utc) + return (now.year, now.month) + + +# --------------------------------------------------------------------------- +# Periodic-task entry points +# --------------------------------------------------------------------------- + + +async def run_persist_cycle() -> None: + """One scheduler-tick of the persistence task. Atomic-writes the + current counter set to disk. Fail-soft.""" + try: + write_telemetry_snapshot() + except Exception as exc: + logger.warning( + "[kora.cost_telemetry] persist cycle raised %r — counters " + "still in memory; will retry next tick", + exc, + ) + + +async def run_rolling_24h_reset_check() -> None: + """Reset the rolling-24h window if a UTC midnight has crossed + since the last reset. + + First call after process boot stamps the "last reset date" as + today's UTC date WITHOUT firing a reset (counters at zero + anyway). Subsequent calls fire a reset only when the UTC date + changes. + """ + global _last_24h_reset_date + today = _utc_date_now() + if _last_24h_reset_date is None: + _last_24h_reset_date = today + return + if today > _last_24h_reset_date: + try: + get_telemetry().reset_window(WINDOW_ROLLING_24H) + _last_24h_reset_date = today + except Exception as exc: + logger.warning( + "[kora.cost_telemetry] rolling_24h reset raised %r — " + "will retry next tick", + exc, + ) + + +async def run_monthly_reset_check() -> None: + """Reset the monthly window if a UTC month rollover has crossed + since the last reset. + + Same first-tick stamping shape as the 24h reset. + """ + global _last_monthly_reset_month + current = _utc_month_now() + if _last_monthly_reset_month is None: + _last_monthly_reset_month = current + return + if current != _last_monthly_reset_month: + try: + get_telemetry().reset_window(WINDOW_MONTHLY) + _last_monthly_reset_month = current + except Exception as exc: + logger.warning( + "[kora.cost_telemetry] monthly reset raised %r — " + "will retry next tick", + exc, + ) + + +def _reset_window_tracking_for_tests() -> None: + """Test-only: clear the reset-tracking module state. Production + code MUST NOT call this.""" + global _last_24h_reset_date, _last_monthly_reset_month + _last_24h_reset_date = None + _last_monthly_reset_month = None + + +# --------------------------------------------------------------------------- +# Listener +# --------------------------------------------------------------------------- + + +class CostTelemetryListener: + """Stateless lifecycle binding. The three periodic tasks are + registered at module-import time; this listener just emits + boot logs so operators can confirm wiring.""" + + async def startup(self) -> None: + logger.info( + "[kora.cost_telemetry_listener] periodic tasks registered: " + "persist cadence=%ss, reset-tick cadence=%ss", + _read_persist_interval(), + _read_reset_tick_interval(), + ) + + async def shutdown(self) -> None: + logger.info("[kora.cost_telemetry_listener] shutdown") + + +# --------------------------------------------------------------------------- +# Factory + registration (import-time side effect) +# --------------------------------------------------------------------------- + + +def _factory(): + listener = CostTelemetryListener() + return (listener.startup, listener.shutdown, DEFAULT_SHUTDOWN_TIMEOUT) + + +register_daemon_listener("cost_telemetry", _factory) + + +# 5-min persistence cycle. +register_periodic_task( + "cost_telemetry.persist", + interval_seconds=_read_persist_interval(), + callable=run_persist_cycle, +) +# Hourly watch-and-act resets for the two windowed counters. +register_periodic_task( + "cost_telemetry.rolling_24h_reset", + interval_seconds=_read_reset_tick_interval(), + callable=run_rolling_24h_reset_check, +) +register_periodic_task( + "cost_telemetry.monthly_reset", + interval_seconds=_read_reset_tick_interval(), + callable=run_monthly_reset_check, +) diff --git a/kora_cli/snapshot/state_snapshot.py b/kora_cli/snapshot/state_snapshot.py index 7ff12a8d9f87..eb719138d667 100644 --- a/kora_cli/snapshot/state_snapshot.py +++ b/kora_cli/snapshot/state_snapshot.py @@ -55,7 +55,7 @@ logger = logging.getLogger(__name__) -SCHEMA_VERSION = 1 +SCHEMA_VERSION = 2 # v1 → v2: added cost_telemetry section (KR-CHEAP-COST-TELEMETRY) SNAPSHOT_FRESH_THRESHOLD_SECONDS = 600 # 10 min — spec §2(a) is_snapshot_fresh # Probe names the snapshot exposes. Matches the 5 default probes in @@ -318,6 +318,47 @@ def _collect_tasks() -> Dict[str, Any]: return {"open_count": "unknown", "in_progress_count": "unknown"} +def _collect_cost_telemetry() -> Dict[str, Any]: + """Per-route cost counters projection — KR-CHEAP-COST-TELEMETRY. + + Schema v2 addition. Exposes the two operator-facing windows + (``rolling_24h`` + ``monthly``); the ``process_lifetime`` + window is intentionally excluded from the snapshot to keep the + on-disk file size bounded (operator can hit ``/api/cost_telemetry`` + directly for the full window set). + + Fail-soft: missing telemetry singleton (e.g., the cost_telemetry + listener hasn't booted yet) degrades to empty per-window dicts + so the snapshot shape is stable. + """ + try: + from kora_cli.telemetry import ( + WINDOW_MONTHLY, + WINDOW_ROLLING_24H, + get_telemetry, + ) + except Exception as exc: + logger.debug( + "[kora.snapshot] cost_telemetry import failed: %r — " + "degrading section", + exc, + ) + return {"rolling_24h": {}, "monthly": {}} + try: + all_windows = get_telemetry().snapshot() + except Exception as exc: + logger.warning( + "[kora.snapshot] telemetry.snapshot() raised %r — " + "degrading section", + exc, + ) + return {"rolling_24h": {}, "monthly": {}} + return { + "rolling_24h": all_windows.get(WINDOW_ROLLING_24H, {}), + "monthly": all_windows.get(WINDOW_MONTHLY, {}), + } + + # --------------------------------------------------------------------------- # Public surface # --------------------------------------------------------------------------- @@ -329,6 +370,9 @@ def compute_snapshot() -> Dict[str, Any]: Per-source failures degrade in-place (per the collector contracts); this top-level function never raises. Caller can treat the returned dict as a safe-to-serialize wire payload. + + Schema v2 (KR-CHEAP-COST-TELEMETRY) adds ``cost_telemetry`` + alongside the v1 sections. """ from kora_time import now @@ -342,6 +386,7 @@ def compute_snapshot() -> Dict[str, Any]: "cost_ladder": _collect_cost_ladder(), "tasks": _collect_tasks(), "service_health": _collect_service_health(), + "cost_telemetry": _collect_cost_telemetry(), } diff --git a/kora_cli/telemetry/__init__.py b/kora_cli/telemetry/__init__.py new file mode 100644 index 000000000000..2cdaeef0ecee --- /dev/null +++ b/kora_cli/telemetry/__init__.py @@ -0,0 +1,50 @@ +"""Telemetry — KR-CHEAP-COST-TELEMETRY. + +Per-route cost counters tagged on every ``record_inference`` +call. See :mod:`kora_cli.telemetry.cost_telemetry` for the full +contract. + +Public surface: + - :class:`CostRouteTelemetry` — singleton accumulator + - :func:`get_telemetry` — process-global accessor + - ``ROUTE_*`` literal constants — canonical route taxonomy + - ``WINDOW_*`` literal constants — three counter windows +""" + +from kora_cli.telemetry.cost_telemetry import ( + KNOWN_ROUTES, + KNOWN_WINDOWS, + ROUTE_ALERT_INVESTIGATION, + ROUTE_EMAIL_INBOUND, + ROUTE_EMAIL_OUTBOUND_COMPOSE, + ROUTE_MCP_TOOL, + ROUTE_PROBE_INVESTIGATION, + ROUTE_SCHEDULED_TASK, + ROUTE_SLACK_DM, + ROUTE_TOOL_LOOP_ITERATION, + ROUTE_UNKNOWN, + WINDOW_MONTHLY, + WINDOW_PROCESS_LIFETIME, + WINDOW_ROLLING_24H, + CostRouteTelemetry, + get_telemetry, +) + +__all__ = [ + "CostRouteTelemetry", + "KNOWN_ROUTES", + "KNOWN_WINDOWS", + "ROUTE_ALERT_INVESTIGATION", + "ROUTE_EMAIL_INBOUND", + "ROUTE_EMAIL_OUTBOUND_COMPOSE", + "ROUTE_MCP_TOOL", + "ROUTE_PROBE_INVESTIGATION", + "ROUTE_SCHEDULED_TASK", + "ROUTE_SLACK_DM", + "ROUTE_TOOL_LOOP_ITERATION", + "ROUTE_UNKNOWN", + "WINDOW_MONTHLY", + "WINDOW_PROCESS_LIFETIME", + "WINDOW_ROLLING_24H", + "get_telemetry", +] diff --git a/kora_cli/telemetry/cost_telemetry.py b/kora_cli/telemetry/cost_telemetry.py new file mode 100644 index 000000000000..f2c44c958922 --- /dev/null +++ b/kora_cli/telemetry/cost_telemetry.py @@ -0,0 +1,360 @@ +"""Per-route cost telemetry — KR-CHEAP-COST-TELEMETRY (R3-4 #10). + +Tags every ``record_inference`` call with a route label; accumulates +per-route counters; surfaces for cockpit + tuning decisions. + +Zero LLM cost on the accounting itself. The decision layer for any +future tuning (escalation-rate, classifier, route-shape) reads +from this telemetry to evaluate "is cheap-substrate work actually +saving what we expect?" + +# Route taxonomy (v1 — spec §2) + +| Route | When | +|---|---| +| ``slack_dm`` | DM-equivalent traffic the slack handler bills | +| ``email_inbound`` | Inbound email reasoning (reserved — handler doesn't yet write a bill) | +| ``email_outbound_compose`` | Reasoning invoked to draft an outbound email (reserved) | +| ``mcp_tool`` | An MCP-driven invocation reaches reasoning (reserved) | +| ``alert_investigation`` | Alert wakes Kora; investigation reasoning (Lock R3-8 (d); reserved) | +| ``probe_investigation`` | Probe escalates an issue; Kora investigates (Lock R3-8 (b); reserved) | +| ``tool_loop_iteration`` | Tool-use loop iteration 2+ (first iteration attributed to the parent route) | +| ``scheduled_task`` | Cron-fired scheduled task invokes reasoning (reserved) | +| ``unknown`` | route metadata absent / unrecognized (fail-soft; never raises) | + +Routes accepted whether or not a current consumer exists. Reserving +the literal in the taxonomy lets a future bucket wire a new consumer +without touching this module. + +# Concurrency + +A single :class:`threading.RLock` protects all counter mutations + +snapshot reads. ``record_call`` is called from: + + * Daemon asyncio loop (handlers, reasoning, periodic tasks) + * Cron-driven periodic snapshot task (same loop today) + * Possibly future background threads + +Lock-per-update is cheap; the alternative (lock-free dict +mutations under GIL) is technically safe for individual key +mutations but the snapshot/aggregation read is multi-key and a +race could surface an inconsistent picture. Lock is the +conservative choice. + +# Windows + +Three independent counter windows maintained: + + * ``process_lifetime`` — reset on process boot only + * ``rolling_24h`` — reset at midnight UTC (cron-driven) + * ``monthly`` — reset at month rollover (cron-driven) + +The cost-ladder holder already tracks the monthly billing window +for budget-cap purposes; this telemetry's ``monthly`` window +aligns to the same UTC month boundary so operator can correlate +per-route burn to monthly burn cap. +""" + +from __future__ import annotations + +import logging +import threading +from dataclasses import dataclass, field +from typing import Any, Dict, Iterable, Optional + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Canonical route taxonomy (spec §2) +# --------------------------------------------------------------------------- + + +ROUTE_SLACK_DM = "slack_dm" +ROUTE_EMAIL_INBOUND = "email_inbound" +ROUTE_EMAIL_OUTBOUND_COMPOSE = "email_outbound_compose" +ROUTE_MCP_TOOL = "mcp_tool" +ROUTE_ALERT_INVESTIGATION = "alert_investigation" +ROUTE_PROBE_INVESTIGATION = "probe_investigation" +ROUTE_TOOL_LOOP_ITERATION = "tool_loop_iteration" +ROUTE_SCHEDULED_TASK = "scheduled_task" +ROUTE_UNKNOWN = "unknown" + +KNOWN_ROUTES = frozenset( + { + ROUTE_SLACK_DM, + ROUTE_EMAIL_INBOUND, + ROUTE_EMAIL_OUTBOUND_COMPOSE, + ROUTE_MCP_TOOL, + ROUTE_ALERT_INVESTIGATION, + ROUTE_PROBE_INVESTIGATION, + ROUTE_TOOL_LOOP_ITERATION, + ROUTE_SCHEDULED_TASK, + ROUTE_UNKNOWN, + } +) + + +# Window names; document inline for snapshot / endpoint consumers. +WINDOW_PROCESS_LIFETIME = "process_lifetime" +WINDOW_ROLLING_24H = "rolling_24h" +WINDOW_MONTHLY = "monthly" + +KNOWN_WINDOWS = (WINDOW_PROCESS_LIFETIME, WINDOW_ROLLING_24H, WINDOW_MONTHLY) + + +# --------------------------------------------------------------------------- +# Counter shape +# --------------------------------------------------------------------------- + + +@dataclass +class _RouteCounters: + """Mutable counters for one (window, route) pair. + + Not frozen — mutated in-place under the singleton's lock. The + ``snapshot()`` method produces a JSON-serializable dict copy + safe to hand out to readers. + """ + + calls_count: int = 0 + input_tokens_total: int = 0 + output_tokens_total: int = 0 + cache_read_tokens_total: int = 0 + cache_creation_tokens_total: int = 0 + cost_estimate_usd_total: float = 0.0 + escalation_count: int = 0 + model_breakdown: Dict[str, int] = field(default_factory=dict) + + def add( + self, + *, + canonical_usage: Any, + cost_estimate_usd: Optional[float], + model: str, + escalated_to_opus: bool, + ) -> None: + """Increment counters from one call's worth of usage.""" + self.calls_count += 1 + # ``canonical_usage`` is a duck-typed CanonicalUsage — read + # the 4 token-count fields defensively (each may be missing + # on a future shape variant; tolerate via getattr with 0 + # default). The pricing module's canonical struct uses + # ``cache_write_tokens`` for the per-call cost write; we + # surface it under ``cache_creation_tokens_total`` per the + # spec §2 counter naming. + self.input_tokens_total += int( + getattr(canonical_usage, "input_tokens", 0) or 0 + ) + self.output_tokens_total += int( + getattr(canonical_usage, "output_tokens", 0) or 0 + ) + self.cache_read_tokens_total += int( + getattr(canonical_usage, "cache_read_tokens", 0) or 0 + ) + self.cache_creation_tokens_total += int( + getattr(canonical_usage, "cache_write_tokens", 0) or 0 + ) + if cost_estimate_usd is not None: + try: + self.cost_estimate_usd_total += float(cost_estimate_usd) + except (TypeError, ValueError): + pass + if escalated_to_opus: + self.escalation_count += 1 + if model: + self.model_breakdown[model] = ( + self.model_breakdown.get(model, 0) + 1 + ) + + def to_dict(self) -> Dict[str, Any]: + return { + "calls_count": self.calls_count, + "input_tokens_total": self.input_tokens_total, + "output_tokens_total": self.output_tokens_total, + "cache_read_tokens_total": self.cache_read_tokens_total, + "cache_creation_tokens_total": self.cache_creation_tokens_total, + "cost_estimate_usd_total": round( + self.cost_estimate_usd_total, 6 + ), + "escalation_count": self.escalation_count, + "model_breakdown": dict(self.model_breakdown), + } + + +def _empty_window() -> Dict[str, _RouteCounters]: + """Build a fresh per-route counter dict for one window. + + Pre-populated with every known route so the snapshot shape is + stable from process boot — consumers don't have to handle + "route absent" vs "route at zero" distinctly. + """ + return {route: _RouteCounters() for route in KNOWN_ROUTES} + + +# --------------------------------------------------------------------------- +# CostRouteTelemetry — singleton accumulator +# --------------------------------------------------------------------------- + + +class CostRouteTelemetry: + """Per-route counter accumulator. Singleton via :func:`get_telemetry`. + + Counters reset at process boot (``process_lifetime`` window). + Persistent windowed counters (``rolling_24h``, ``monthly``) are + written to disk every 5 min by the snapshot listener; window- + reset periodic tasks clear them at the appropriate boundaries. + """ + + def __init__(self) -> None: + self._lock = threading.RLock() + self._windows: Dict[str, Dict[str, _RouteCounters]] = { + window: _empty_window() for window in KNOWN_WINDOWS + } + + def record_call( + self, + *, + route: str, + model: str, + canonical_usage: Any, + cost_estimate_usd: Optional[float], + escalated_to_opus: bool = False, + ) -> None: + """Increment counters for this route across all live windows. + + Fail-soft: any exception inside is caught + logged so the + caller (a hot-path inference completion handler) never sees + a telemetry failure. Unknown routes silently bucket to + ``"unknown"`` rather than raising. + """ + try: + self._record_call_inner( + route=route, + model=model or "", + canonical_usage=canonical_usage, + cost_estimate_usd=cost_estimate_usd, + escalated_to_opus=bool(escalated_to_opus), + ) + except Exception as exc: + logger.warning( + "[kora.cost_telemetry] record_call raised %r — counters " + "not updated for route=%s", + exc, + route, + ) + + def _record_call_inner( + self, + *, + route: str, + model: str, + canonical_usage: Any, + cost_estimate_usd: Optional[float], + escalated_to_opus: bool, + ) -> None: + normalized_route = ( + route + if isinstance(route, str) and route in KNOWN_ROUTES + else ROUTE_UNKNOWN + ) + with self._lock: + for window in KNOWN_WINDOWS: + self._windows[window][normalized_route].add( + canonical_usage=canonical_usage, + cost_estimate_usd=cost_estimate_usd, + model=model, + escalated_to_opus=escalated_to_opus, + ) + + def snapshot(self) -> Dict[str, Any]: + """Snapshot the current counters for cockpit / on-disk write. + + Returns a JSON-serializable dict shaped: + + .. code-block:: python + + { + "process_lifetime": {: {...}, ...}, + "rolling_24h": {: {...}, ...}, + "monthly": {: {...}, ...}, + } + + Each per-route dict matches :meth:`_RouteCounters.to_dict`. + Snapshot read happens under the lock so a concurrent + ``record_call`` can't surface a half-updated counter. + """ + with self._lock: + return { + window: { + route: counters.to_dict() + for route, counters in routes.items() + } + for window, routes in self._windows.items() + } + + def reset_window(self, window: str) -> None: + """Reset one named window's counters to zero. + + Used by the rolling-24h and monthly reset periodic tasks + when the respective boundary rolls. The + ``process_lifetime`` window MAY be reset via this surface + (test fixtures use it) but production code shouldn't. + """ + if window not in KNOWN_WINDOWS: + logger.warning( + "[kora.cost_telemetry] reset_window: unknown window=%r " + "(known: %s) — ignoring", + window, + KNOWN_WINDOWS, + ) + return + with self._lock: + self._windows[window] = _empty_window() + logger.info( + "[kora.cost_telemetry] window=%s counters reset", window + ) + + def reset_all_for_tests(self) -> None: + """Test-only: reset every window to zero. Production code + uses :meth:`reset_window`.""" + with self._lock: + for window in KNOWN_WINDOWS: + self._windows[window] = _empty_window() + + +# --------------------------------------------------------------------------- +# Singleton accessor +# --------------------------------------------------------------------------- + + +_telemetry_singleton: Optional[CostRouteTelemetry] = None +_singleton_lock = threading.Lock() + + +def get_telemetry() -> CostRouteTelemetry: + """Process-global accessor for :class:`CostRouteTelemetry`. + + Lazily constructs the singleton on first call. Thread-safe + construction via double-checked locking. Subsequent calls + return the same instance. + """ + global _telemetry_singleton + if _telemetry_singleton is not None: + return _telemetry_singleton + with _singleton_lock: + if _telemetry_singleton is None: + _telemetry_singleton = CostRouteTelemetry() + return _telemetry_singleton + + +def _reset_singleton_for_tests() -> None: + """Drop the singleton so the next ``get_telemetry()`` call returns + a fresh instance. Used by test fixtures to isolate state. + + Production code MUST NOT call this — it would zero all live + counters mid-process. + """ + global _telemetry_singleton + with _singleton_lock: + _telemetry_singleton = None diff --git a/kora_cli/web_server.py b/kora_cli/web_server.py index cdec563defd7..7bc3af4c179e 100644 --- a/kora_cli/web_server.py +++ b/kora_cli/web_server.py @@ -5993,6 +5993,38 @@ async def get_daemon_snapshot(): return snap +# --------------------------------------------------------------------------- +# Per-route cost telemetry (KR-CHEAP-COST-TELEMETRY) +# --------------------------------------------------------------------------- + + +@app.get("/api/cost_telemetry") +async def get_cost_telemetry(): + """Per-route cost counters across all 3 windows. + + Source-of-truth for any Kora-cost decisions (escalation rate, + route shape, classifier tuning, etc.). Reads the in-memory + telemetry singleton directly — no disk roundtrip, no LLM + cost. + + Shape: + + .. code-block:: json + + { + "process_lifetime": {"slack_dm": {...}, "unknown": {...}, ...}, + "rolling_24h": {"slack_dm": {...}, ...}, + "monthly": {"slack_dm": {...}, ...} + } + + Per-route counter shape comes from + :class:`kora_cli.telemetry.cost_telemetry._RouteCounters.to_dict`. + """ + from kora_cli.telemetry import get_telemetry + + return get_telemetry().snapshot() + + # --------------------------------------------------------------------------- # Profile management endpoints (minimal — list/create/rename/delete + SOUL.md) # --------------------------------------------------------------------------- diff --git a/tests/kora_cli/snapshot/test_state_snapshot.py b/tests/kora_cli/snapshot/test_state_snapshot.py index 28ac2b1a425b..34869484226b 100644 --- a/tests/kora_cli/snapshot/test_state_snapshot.py +++ b/tests/kora_cli/snapshot/test_state_snapshot.py @@ -100,6 +100,8 @@ def test_compute_snapshot_has_all_required_top_level_keys(env): "cost_ladder", "tasks", "service_health", + # KR-CHEAP-COST-TELEMETRY v2 addition. + "cost_telemetry", } diff --git a/tests/kora_cli/telemetry/__init__.py b/tests/kora_cli/telemetry/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/kora_cli/telemetry/test_cost_telemetry.py b/tests/kora_cli/telemetry/test_cost_telemetry.py new file mode 100644 index 000000000000..f3f3325516e1 --- /dev/null +++ b/tests/kora_cli/telemetry/test_cost_telemetry.py @@ -0,0 +1,422 @@ +"""Tests for KR-CHEAP-COST-TELEMETRY — per-route counters. + +Bucket §2 scenarios: + + Counter shape: + 1. record_call increments calls_count + per-route totals + 2. Token sums accumulate across multiple calls + 3. cost_estimate_usd_total sums; None values skipped + 4. escalation_count increments only when escalated_to_opus=True + 5. model_breakdown tracks per-model call counts + + Route taxonomy: + 6. Each canonical route accepted + bucketed correctly + 7. Unknown route string buckets to "unknown" + 8. Non-string route falls back to "unknown" + + Windows: + 9. record_call updates ALL three windows + 10. reset_window clears one window without affecting others + 11. reset_window with unknown window name is no-op + warns + + Concurrency: + 12. Concurrent record_call from multiple threads doesn't lose counts + + Singleton: + 13. get_telemetry returns same instance across calls + 14. _reset_singleton_for_tests gives a fresh instance + + Snapshot shape: + 15. snapshot returns dict with all 3 windows + 16. Each window has every known route pre-populated (stable shape) + 17. Snapshot is JSON-serializable + + Fail-soft: + 18. record_call with bad canonical_usage doesn't raise +""" + +from __future__ import annotations + +import json +import threading +from dataclasses import dataclass +from typing import Optional + +import pytest + +from kora_cli.telemetry import ( + KNOWN_ROUTES, + KNOWN_WINDOWS, + ROUTE_ALERT_INVESTIGATION, + ROUTE_EMAIL_INBOUND, + ROUTE_MCP_TOOL, + ROUTE_SLACK_DM, + ROUTE_UNKNOWN, + WINDOW_MONTHLY, + WINDOW_PROCESS_LIFETIME, + WINDOW_ROLLING_24H, + CostRouteTelemetry, + get_telemetry, +) +from kora_cli.telemetry.cost_telemetry import _reset_singleton_for_tests + + +@dataclass(frozen=True) +class _FakeUsage: + input_tokens: int = 0 + output_tokens: int = 0 + cache_read_tokens: int = 0 + cache_write_tokens: int = 0 + reasoning_tokens: int = 0 + + +@pytest.fixture +def telemetry(): + """Fresh CostRouteTelemetry per test (bypasses the singleton so + tests don't pollute each other).""" + return CostRouteTelemetry() + + +# =========================================================================== +# Counter shape +# =========================================================================== + + +def test_record_call_increments_calls_count(telemetry): + telemetry.record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=_FakeUsage(input_tokens=100, output_tokens=50), + cost_estimate_usd=0.0042, + ) + snap = telemetry.snapshot() + assert snap[WINDOW_PROCESS_LIFETIME][ROUTE_SLACK_DM]["calls_count"] == 1 + assert snap[WINDOW_ROLLING_24H][ROUTE_SLACK_DM]["calls_count"] == 1 + assert snap[WINDOW_MONTHLY][ROUTE_SLACK_DM]["calls_count"] == 1 + + +def test_token_sums_accumulate(telemetry): + for _ in range(3): + telemetry.record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=_FakeUsage( + input_tokens=100, + output_tokens=50, + cache_read_tokens=200, + cache_write_tokens=300, + ), + cost_estimate_usd=0.01, + ) + row = telemetry.snapshot()[WINDOW_PROCESS_LIFETIME][ROUTE_SLACK_DM] + assert row["calls_count"] == 3 + assert row["input_tokens_total"] == 300 + assert row["output_tokens_total"] == 150 + assert row["cache_read_tokens_total"] == 600 + assert row["cache_creation_tokens_total"] == 900 + assert row["cost_estimate_usd_total"] == pytest.approx(0.03) + + +def test_none_cost_skipped(telemetry): + """Cost-estimate None values don't blow up the sum.""" + telemetry.record_call( + route=ROUTE_SLACK_DM, + model="claude-haiku-4-5", + canonical_usage=_FakeUsage(input_tokens=10), + cost_estimate_usd=None, + ) + telemetry.record_call( + route=ROUTE_SLACK_DM, + model="claude-haiku-4-5", + canonical_usage=_FakeUsage(input_tokens=10), + cost_estimate_usd=0.0001, + ) + row = telemetry.snapshot()[WINDOW_PROCESS_LIFETIME][ROUTE_SLACK_DM] + assert row["calls_count"] == 2 # Both calls counted + assert row["cost_estimate_usd_total"] == pytest.approx(0.0001) + + +def test_escalation_count_only_when_flag_true(telemetry): + telemetry.record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=_FakeUsage(), + cost_estimate_usd=0.01, + escalated_to_opus=False, + ) + telemetry.record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=_FakeUsage(), + cost_estimate_usd=0.01, + escalated_to_opus=True, + ) + telemetry.record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=_FakeUsage(), + cost_estimate_usd=0.01, + escalated_to_opus=True, + ) + row = telemetry.snapshot()[WINDOW_PROCESS_LIFETIME][ROUTE_SLACK_DM] + assert row["calls_count"] == 3 + assert row["escalation_count"] == 2 + + +def test_model_breakdown_tracks_per_model_calls(telemetry): + for model in ["claude-opus-4-7", "claude-sonnet-4-6", "claude-opus-4-7"]: + telemetry.record_call( + route=ROUTE_SLACK_DM, + model=model, + canonical_usage=_FakeUsage(), + cost_estimate_usd=0.001, + ) + row = telemetry.snapshot()[WINDOW_PROCESS_LIFETIME][ROUTE_SLACK_DM] + assert row["model_breakdown"] == { + "claude-opus-4-7": 2, + "claude-sonnet-4-6": 1, + } + + +# =========================================================================== +# Route taxonomy +# =========================================================================== + + +def test_every_canonical_route_accepted(telemetry): + for route in KNOWN_ROUTES: + telemetry.record_call( + route=route, + model="claude-haiku-4-5", + canonical_usage=_FakeUsage(input_tokens=1), + cost_estimate_usd=0.0001, + ) + snap = telemetry.snapshot() + for route in KNOWN_ROUTES: + assert snap[WINDOW_PROCESS_LIFETIME][route]["calls_count"] == 1 + + +def test_unknown_route_string_buckets_to_unknown(telemetry): + telemetry.record_call( + route="this_is_not_a_real_route", + model="claude-haiku-4-5", + canonical_usage=_FakeUsage(), + cost_estimate_usd=0.0001, + ) + snap = telemetry.snapshot() + assert snap[WINDOW_PROCESS_LIFETIME][ROUTE_UNKNOWN]["calls_count"] == 1 + # Other routes untouched. + assert snap[WINDOW_PROCESS_LIFETIME][ROUTE_SLACK_DM]["calls_count"] == 0 + + +def test_non_string_route_buckets_to_unknown(telemetry): + telemetry.record_call( + route=None, # type: ignore[arg-type] + model="claude-haiku-4-5", + canonical_usage=_FakeUsage(), + cost_estimate_usd=0.0001, + ) + assert ( + telemetry.snapshot()[WINDOW_PROCESS_LIFETIME][ROUTE_UNKNOWN]["calls_count"] + == 1 + ) + + +# =========================================================================== +# Windows +# =========================================================================== + + +def test_record_call_updates_all_three_windows(telemetry): + telemetry.record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=_FakeUsage(input_tokens=10), + cost_estimate_usd=0.001, + ) + snap = telemetry.snapshot() + for window in KNOWN_WINDOWS: + assert snap[window][ROUTE_SLACK_DM]["calls_count"] == 1 + assert snap[window][ROUTE_SLACK_DM]["input_tokens_total"] == 10 + + +def test_reset_window_clears_one_window_only(telemetry): + telemetry.record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=_FakeUsage(input_tokens=100), + cost_estimate_usd=0.01, + ) + telemetry.reset_window(WINDOW_ROLLING_24H) + snap = telemetry.snapshot() + # rolling_24h zeroed. + assert snap[WINDOW_ROLLING_24H][ROUTE_SLACK_DM]["calls_count"] == 0 + assert snap[WINDOW_ROLLING_24H][ROUTE_SLACK_DM]["input_tokens_total"] == 0 + # process_lifetime + monthly untouched. + assert snap[WINDOW_PROCESS_LIFETIME][ROUTE_SLACK_DM]["calls_count"] == 1 + assert snap[WINDOW_MONTHLY][ROUTE_SLACK_DM]["calls_count"] == 1 + + +def test_reset_window_unknown_name_is_noop_and_warns(telemetry, caplog): + telemetry.record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=_FakeUsage(input_tokens=10), + cost_estimate_usd=0.001, + ) + with caplog.at_level("WARNING"): + telemetry.reset_window("not_a_real_window") + snap = telemetry.snapshot() + # All windows untouched. + for window in KNOWN_WINDOWS: + assert snap[window][ROUTE_SLACK_DM]["calls_count"] == 1 + assert any("unknown window" in r.message for r in caplog.records) + + +def test_reset_all_for_tests_clears_everything(telemetry): + for route in [ROUTE_SLACK_DM, ROUTE_EMAIL_INBOUND, ROUTE_MCP_TOOL]: + telemetry.record_call( + route=route, + model="claude-haiku-4-5", + canonical_usage=_FakeUsage(input_tokens=10), + cost_estimate_usd=0.001, + ) + telemetry.reset_all_for_tests() + snap = telemetry.snapshot() + for window in KNOWN_WINDOWS: + for route in KNOWN_ROUTES: + assert snap[window][route]["calls_count"] == 0 + + +# =========================================================================== +# Concurrency — basic safety check +# =========================================================================== + + +def test_concurrent_record_call_no_lost_counts(telemetry): + """1000 calls split across 10 threads → final calls_count must + be exactly 1000 (no lost increments under lock).""" + calls_per_thread = 100 + threads = 10 + + def worker(): + for _ in range(calls_per_thread): + telemetry.record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=_FakeUsage(input_tokens=1, output_tokens=1), + cost_estimate_usd=0.0001, + ) + + workers = [threading.Thread(target=worker) for _ in range(threads)] + for w in workers: + w.start() + for w in workers: + w.join() + + row = telemetry.snapshot()[WINDOW_PROCESS_LIFETIME][ROUTE_SLACK_DM] + assert row["calls_count"] == calls_per_thread * threads + assert row["input_tokens_total"] == calls_per_thread * threads + + +# =========================================================================== +# Singleton +# =========================================================================== + + +def test_get_telemetry_returns_same_instance(): + _reset_singleton_for_tests() + t1 = get_telemetry() + t2 = get_telemetry() + assert t1 is t2 + + +def test_reset_singleton_gives_fresh_instance(): + _reset_singleton_for_tests() + t1 = get_telemetry() + t1.record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=_FakeUsage(input_tokens=1), + cost_estimate_usd=0.001, + ) + _reset_singleton_for_tests() + t2 = get_telemetry() + assert t2 is not t1 + snap = t2.snapshot() + # Fresh — counters at zero. + assert snap[WINDOW_PROCESS_LIFETIME][ROUTE_SLACK_DM]["calls_count"] == 0 + + +# =========================================================================== +# Snapshot shape +# =========================================================================== + + +def test_snapshot_has_all_3_windows(telemetry): + snap = telemetry.snapshot() + assert set(snap.keys()) == set(KNOWN_WINDOWS) + + +def test_snapshot_pre_populates_every_route_per_window(telemetry): + """Stable shape: every known route appears as a zero-counter + entry in every window, even with no calls recorded.""" + snap = telemetry.snapshot() + for window in KNOWN_WINDOWS: + assert set(snap[window].keys()) == set(KNOWN_ROUTES) + for route in KNOWN_ROUTES: + row = snap[window][route] + assert row["calls_count"] == 0 + assert row["model_breakdown"] == {} + + +def test_snapshot_is_json_serializable(telemetry): + telemetry.record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=_FakeUsage(input_tokens=10), + cost_estimate_usd=0.001, + escalated_to_opus=True, + ) + snap = telemetry.snapshot() + # Roundtrip — if any value is non-serializable this raises. + raw = json.dumps(snap) + reparsed = json.loads(raw) + assert reparsed[WINDOW_PROCESS_LIFETIME][ROUTE_SLACK_DM]["calls_count"] == 1 + + +# =========================================================================== +# Fail-soft +# =========================================================================== + + +def test_record_call_bad_usage_no_raise(telemetry): + """A canonical_usage that's not the expected shape shouldn't + crash — defensive getattr fallback in _RouteCounters.add.""" + + class BadUsage: + # Missing all expected attrs; getattr fallbacks to 0. + pass + + telemetry.record_call( + route=ROUTE_SLACK_DM, + model="claude-haiku-4-5", + canonical_usage=BadUsage(), + cost_estimate_usd=0.0, + ) + row = telemetry.snapshot()[WINDOW_PROCESS_LIFETIME][ROUTE_SLACK_DM] + assert row["calls_count"] == 1 + assert row["input_tokens_total"] == 0 + + +def test_record_call_bad_cost_no_raise(telemetry): + """Non-numeric cost_estimate_usd doesn't break the sum.""" + telemetry.record_call( + route=ROUTE_SLACK_DM, + model="claude-haiku-4-5", + canonical_usage=_FakeUsage(input_tokens=5), + cost_estimate_usd="not-a-number", # type: ignore[arg-type] + ) + row = telemetry.snapshot()[WINDOW_PROCESS_LIFETIME][ROUTE_SLACK_DM] + assert row["calls_count"] == 1 + assert row["cost_estimate_usd_total"] == 0.0 diff --git a/tests/kora_cli/test_listeners/test_cost_telemetry_listener.py b/tests/kora_cli/test_listeners/test_cost_telemetry_listener.py new file mode 100644 index 000000000000..3cddb49e82a4 --- /dev/null +++ b/tests/kora_cli/test_listeners/test_cost_telemetry_listener.py @@ -0,0 +1,326 @@ +"""Tests for the KR-CHEAP-COST-TELEMETRY daemon listener. + +Scenarios: + 1. Listener registered in LISTENER_REGISTRY at import time + 2. All 3 periodic tasks (persist, rolling_24h_reset, monthly_reset) + registered with the heartbeat scheduler + 3. Default cadences honored; env overrides respected + 4. Listener factory tuple shape correct + 5. Persist task writes atomically to disk + 6. Persist task fail-soft on disk error + 7. Rolling-24h reset fires when UTC date crosses + 8. Rolling-24h reset NO-OP when UTC date stable + 9. Monthly reset fires when UTC month crosses + 10. Monthly reset NO-OP when UTC month stable + 11. read_telemetry_snapshot returns None when file missing + 12. read_telemetry_snapshot returns parsed dict when present +""" + +from __future__ import annotations + +import json +from datetime import datetime, timedelta, timezone +from unittest.mock import patch + +import pytest + +from kora_cli import daemon as daemon_mod +from kora_cli.listeners.heartbeat import PERIODIC_TASK_REGISTRY +from kora_cli.listeners import cost_telemetry_listener +from kora_cli.listeners.cost_telemetry_listener import ( + COST_TELEMETRY_PATH_ENV, + DEFAULT_PERSIST_INTERVAL_SEC, + DEFAULT_RESET_TICK_INTERVAL_SEC, + PERSIST_INTERVAL_ENV, + RESET_TICK_INTERVAL_ENV, + CostTelemetryListener, + _factory, + _read_persist_interval, + _read_reset_tick_interval, + _reset_window_tracking_for_tests, + cost_telemetry_path, + read_telemetry_snapshot, + run_monthly_reset_check, + run_persist_cycle, + run_rolling_24h_reset_check, + write_telemetry_snapshot, +) +from kora_cli.telemetry import ( + ROUTE_SLACK_DM, + WINDOW_MONTHLY, + WINDOW_ROLLING_24H, + get_telemetry, +) +from kora_cli.telemetry.cost_telemetry import _reset_singleton_for_tests + + +@pytest.fixture(autouse=True) +def _isolate(tmp_path, monkeypatch): + """Isolated KORA_HOME + fresh telemetry singleton + cleared + window-reset module state.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + monkeypatch.setenv("KORA_HOME", str(tmp_path)) + monkeypatch.setattr("kora_constants.get_kora_home", lambda: tmp_path) + monkeypatch.delenv(COST_TELEMETRY_PATH_ENV, raising=False) + monkeypatch.delenv(PERSIST_INTERVAL_ENV, raising=False) + monkeypatch.delenv(RESET_TICK_INTERVAL_ENV, raising=False) + _reset_singleton_for_tests() + _reset_window_tracking_for_tests() + yield tmp_path + _reset_singleton_for_tests() + _reset_window_tracking_for_tests() + + +# =========================================================================== +# Registration +# =========================================================================== + + +def test_listener_registered_in_daemon_registry(): + names = {name for name, _f in daemon_mod.LISTENER_REGISTRY} + assert "cost_telemetry" in names + + +def test_all_three_periodic_tasks_registered(): + names = {t.name for t in PERIODIC_TASK_REGISTRY} + assert "cost_telemetry.persist" in names + assert "cost_telemetry.rolling_24h_reset" in names + assert "cost_telemetry.monthly_reset" in names + + +def test_factory_tuple_shape(): + startup, shutdown, timeout = _factory() + assert callable(startup) + assert callable(shutdown) + assert isinstance(timeout, (int, float)) + + +# =========================================================================== +# Cadence +# =========================================================================== + + +def test_persist_interval_default(monkeypatch): + monkeypatch.delenv(PERSIST_INTERVAL_ENV, raising=False) + assert _read_persist_interval() == DEFAULT_PERSIST_INTERVAL_SEC == 300.0 + + +def test_persist_interval_env_override(monkeypatch): + monkeypatch.setenv(PERSIST_INTERVAL_ENV, "60") + assert _read_persist_interval() == 60.0 + + +def test_reset_tick_interval_default(monkeypatch): + monkeypatch.delenv(RESET_TICK_INTERVAL_ENV, raising=False) + assert _read_reset_tick_interval() == DEFAULT_RESET_TICK_INTERVAL_SEC == 3600.0 + + +def test_persist_interval_invalid_falls_back(monkeypatch, caplog): + monkeypatch.setenv(PERSIST_INTERVAL_ENV, "not-numeric") + with caplog.at_level("WARNING"): + assert _read_persist_interval() == DEFAULT_PERSIST_INTERVAL_SEC + assert any("is not numeric" in r.message for r in caplog.records) + + +# =========================================================================== +# Path resolution +# =========================================================================== + + +def test_cost_telemetry_path_default(_isolate, tmp_path): + assert cost_telemetry_path() == tmp_path / "cache" / "cost_telemetry.json" + + +def test_cost_telemetry_path_env_override(_isolate, tmp_path, monkeypatch): + override = tmp_path / "alt_cost.json" + monkeypatch.setenv(COST_TELEMETRY_PATH_ENV, str(override)) + assert cost_telemetry_path() == override + + +# =========================================================================== +# Persistence +# =========================================================================== + + +def test_write_telemetry_snapshot_creates_file_with_payload(_isolate): + # Pre-populate a counter so the write has interesting content. + fake_usage = type("U", (), {"input_tokens": 10, "output_tokens": 5})() + get_telemetry().record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=fake_usage, + cost_estimate_usd=0.001, + ) + write_telemetry_snapshot() + target = cost_telemetry_path() + assert target.exists() + payload = json.loads(target.read_text()) + assert payload["schema_version"] == 1 + assert "written_at" in payload + assert "windows" in payload + # All three windows present. + assert set(payload["windows"].keys()) == { + "process_lifetime", + "rolling_24h", + "monthly", + } + + +@pytest.mark.asyncio +async def test_run_persist_cycle_writes_file_end_to_end(_isolate): + await run_persist_cycle() + assert cost_telemetry_path().exists() + + +@pytest.mark.asyncio +async def test_run_persist_cycle_fail_soft_on_write_error( + _isolate, monkeypatch, caplog +): + def boom(): + raise OSError("disk full") + + monkeypatch.setattr( + "kora_cli.listeners.cost_telemetry_listener.write_telemetry_snapshot", + boom, + ) + with caplog.at_level("WARNING"): + await run_persist_cycle() # must not raise + assert any("persist cycle raised" in r.message for r in caplog.records) + + +# =========================================================================== +# Window-reset checks +# =========================================================================== + + +@pytest.mark.asyncio +async def test_rolling_24h_first_call_stamps_no_reset(_isolate): + """First call after boot stamps the date but does NOT fire a + reset (counters at zero anyway).""" + get_telemetry().record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=type("U", (), {"input_tokens": 10})(), + cost_estimate_usd=0.001, + ) + await run_rolling_24h_reset_check() + # Counter still populated — no reset fired on first call. + snap = get_telemetry().snapshot() + assert snap[WINDOW_ROLLING_24H][ROUTE_SLACK_DM]["calls_count"] == 1 + + +@pytest.mark.asyncio +async def test_rolling_24h_reset_fires_when_utc_date_crosses(_isolate): + """Force the tracked "last reset" date back to yesterday so the + next check fires a reset.""" + get_telemetry().record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=type("U", (), {"input_tokens": 10})(), + cost_estimate_usd=0.001, + ) + # First tick — stamps today. + await run_rolling_24h_reset_check() + # Manually backdate the tracked date so the next tick crosses. + yesterday = datetime.now(timezone.utc).replace( + hour=0, minute=0, second=0, microsecond=0 + ) - timedelta(days=1) + cost_telemetry_listener._last_24h_reset_date = yesterday + await run_rolling_24h_reset_check() + snap = get_telemetry().snapshot() + # rolling_24h zeroed. + assert snap[WINDOW_ROLLING_24H][ROUTE_SLACK_DM]["calls_count"] == 0 + # process_lifetime + monthly untouched. + assert snap["process_lifetime"][ROUTE_SLACK_DM]["calls_count"] == 1 + assert snap[WINDOW_MONTHLY][ROUTE_SLACK_DM]["calls_count"] == 1 + + +@pytest.mark.asyncio +async def test_rolling_24h_reset_noop_when_date_stable(_isolate): + """Two ticks in the same UTC day → second one is a no-op.""" + get_telemetry().record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=type("U", (), {"input_tokens": 10})(), + cost_estimate_usd=0.001, + ) + await run_rolling_24h_reset_check() + await run_rolling_24h_reset_check() + snap = get_telemetry().snapshot() + assert snap[WINDOW_ROLLING_24H][ROUTE_SLACK_DM]["calls_count"] == 1 + + +@pytest.mark.asyncio +async def test_monthly_reset_fires_when_month_crosses(_isolate): + get_telemetry().record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=type("U", (), {"input_tokens": 10})(), + cost_estimate_usd=0.001, + ) + await run_monthly_reset_check() + # Backdate to last month. + now = datetime.now(timezone.utc) + prev_month = (now.year - 1, 12) if now.month == 1 else ( + now.year, now.month - 1 + ) + cost_telemetry_listener._last_monthly_reset_month = prev_month + await run_monthly_reset_check() + snap = get_telemetry().snapshot() + assert snap[WINDOW_MONTHLY][ROUTE_SLACK_DM]["calls_count"] == 0 + # Other windows untouched. + assert snap[WINDOW_ROLLING_24H][ROUTE_SLACK_DM]["calls_count"] == 1 + + +@pytest.mark.asyncio +async def test_monthly_reset_noop_when_month_stable(_isolate): + get_telemetry().record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=type("U", (), {"input_tokens": 10})(), + cost_estimate_usd=0.001, + ) + await run_monthly_reset_check() + await run_monthly_reset_check() + snap = get_telemetry().snapshot() + assert snap[WINDOW_MONTHLY][ROUTE_SLACK_DM]["calls_count"] == 1 + + +# =========================================================================== +# read_telemetry_snapshot +# =========================================================================== + + +def test_read_returns_none_when_missing(_isolate): + assert read_telemetry_snapshot() is None + + +def test_read_returns_parsed_when_present(_isolate): + write_telemetry_snapshot() + snap = read_telemetry_snapshot() + assert snap is not None + assert snap["schema_version"] == 1 + assert "windows" in snap + + +def test_read_returns_none_on_malformed_json(_isolate): + target = cost_telemetry_path() + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text("{ not json") + assert read_telemetry_snapshot() is None + + +# =========================================================================== +# Lifecycle log lines +# =========================================================================== + + +@pytest.mark.asyncio +async def test_listener_lifecycle_emits_info_lines(_isolate, caplog): + listener = CostTelemetryListener() + with caplog.at_level("INFO"): + await listener.startup() + await listener.shutdown() + msgs = " ".join(r.message for r in caplog.records) + assert "periodic tasks registered" in msgs + assert "shutdown" in msgs diff --git a/tests/kora_cli/test_web_server_cost_telemetry.py b/tests/kora_cli/test_web_server_cost_telemetry.py new file mode 100644 index 000000000000..419396cbe222 --- /dev/null +++ b/tests/kora_cli/test_web_server_cost_telemetry.py @@ -0,0 +1,158 @@ +"""Tests for /api/cost_telemetry + snapshot v2 integration. + +Scenarios: + 1. /api/cost_telemetry returns the telemetry snapshot dict + 2. Endpoint surfaces all 3 windows + 3. Snapshot v2 schema_version bumped to 2 + 4. Snapshot includes cost_telemetry section with rolling_24h + monthly + 5. Snapshot cost_telemetry survives telemetry import failure +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest + +from kora_cli.telemetry import ( + ROUTE_SLACK_DM, + WINDOW_MONTHLY, + WINDOW_PROCESS_LIFETIME, + WINDOW_ROLLING_24H, + get_telemetry, +) +from kora_cli.telemetry.cost_telemetry import _reset_singleton_for_tests + + +@dataclass(frozen=True) +class _FakeUsage: + input_tokens: int = 0 + output_tokens: int = 0 + cache_read_tokens: int = 0 + cache_write_tokens: int = 0 + + +@pytest.fixture(autouse=True) +def _isolate(tmp_path, monkeypatch): + """3-namespace get_kora_home isolation + fresh telemetry.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path)) + monkeypatch.setenv("KORA_HOME", str(tmp_path)) + monkeypatch.setattr("kora_constants.get_kora_home", lambda: tmp_path) + monkeypatch.setattr("kora_cli.config.get_kora_home", lambda: tmp_path) + monkeypatch.setattr( + "kora_cli.web_server.get_kora_home", lambda: tmp_path, raising=False + ) + _reset_singleton_for_tests() + yield tmp_path + _reset_singleton_for_tests() + + +# =========================================================================== +# /api/cost_telemetry +# =========================================================================== + + +@pytest.mark.asyncio +async def test_endpoint_returns_telemetry_snapshot(_isolate): + from kora_cli import web_server + + get_telemetry().record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=_FakeUsage(input_tokens=100, output_tokens=50), + cost_estimate_usd=0.005, + ) + result = await web_server.get_cost_telemetry() + assert set(result.keys()) == { + WINDOW_PROCESS_LIFETIME, + WINDOW_ROLLING_24H, + WINDOW_MONTHLY, + } + assert result[WINDOW_PROCESS_LIFETIME][ROUTE_SLACK_DM]["calls_count"] == 1 + assert result[WINDOW_ROLLING_24H][ROUTE_SLACK_DM]["calls_count"] == 1 + + +@pytest.mark.asyncio +async def test_endpoint_returns_zero_counters_when_no_calls(_isolate): + from kora_cli import web_server + + result = await web_server.get_cost_telemetry() + # Stable shape even with no calls — every route at zero. + for window in ( + WINDOW_PROCESS_LIFETIME, + WINDOW_ROLLING_24H, + WINDOW_MONTHLY, + ): + assert result[window][ROUTE_SLACK_DM]["calls_count"] == 0 + + +# =========================================================================== +# Snapshot v2 +# =========================================================================== + + +def test_snapshot_schema_version_bumped_to_v2(_isolate): + from kora_cli.snapshot import compute_snapshot + + snap = compute_snapshot() + assert snap["schema_version"] == 2 + + +def test_snapshot_includes_cost_telemetry_section(_isolate): + from kora_cli.snapshot import compute_snapshot + + get_telemetry().record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=_FakeUsage(input_tokens=10), + cost_estimate_usd=0.001, + ) + snap = compute_snapshot() + assert "cost_telemetry" in snap + # Exposes the two operator-facing windows (NOT process_lifetime + # — see _collect_cost_telemetry docstring). + assert set(snap["cost_telemetry"].keys()) == {"rolling_24h", "monthly"} + assert ( + snap["cost_telemetry"]["rolling_24h"][ROUTE_SLACK_DM]["calls_count"] + == 1 + ) + + +def test_snapshot_cost_telemetry_degrades_when_singleton_unavailable( + _isolate, monkeypatch +): + """If the telemetry singleton's snapshot() raises, the section + degrades to empty dicts rather than failing the whole snapshot.""" + from kora_cli.snapshot import compute_snapshot + + def boom(): + raise RuntimeError("telemetry dead") + + # Get the singleton built, then sabotage its snapshot method. + t = get_telemetry() + monkeypatch.setattr(t, "snapshot", boom) + snap = compute_snapshot() + assert snap["cost_telemetry"] == {"rolling_24h": {}, "monthly": {}} + + +@pytest.mark.asyncio +async def test_api_snapshot_endpoint_includes_cost_telemetry(_isolate): + """End-to-end: /api/snapshot returns the v2 shape with the new + cost_telemetry section.""" + from kora_cli import web_server + from kora_cli.snapshot import compute_snapshot, write_snapshot + + get_telemetry().record_call( + route=ROUTE_SLACK_DM, + model="claude-opus-4-7", + canonical_usage=_FakeUsage(input_tokens=50, output_tokens=25), + cost_estimate_usd=0.002, + ) + write_snapshot(compute_snapshot()) + result = await web_server.get_daemon_snapshot() + assert result["schema_version"] == 2 + assert "cost_telemetry" in result + assert ( + result["cost_telemetry"]["rolling_24h"][ROUTE_SLACK_DM]["calls_count"] + == 1 + ) diff --git a/tests/kora_cli/test_web_server_snapshot.py b/tests/kora_cli/test_web_server_snapshot.py index cb3264f01ed7..f520dd02ad78 100644 --- a/tests/kora_cli/test_web_server_snapshot.py +++ b/tests/kora_cli/test_web_server_snapshot.py @@ -46,12 +46,14 @@ async def test_endpoint_returns_snapshot_when_fresh(_isolate): write_snapshot(compute_snapshot()) result = await web_server.get_daemon_snapshot() assert "error" not in result - assert result["schema_version"] == 1 + # KR-CHEAP-COST-TELEMETRY bumped schema v1 → v2 (added cost_telemetry). + assert result["schema_version"] == 2 assert "computed_at" in result assert "operational_state" in result assert "alerts" in result assert "cost_ladder" in result assert "service_health" in result + assert "cost_telemetry" in result @pytest.mark.asyncio