mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat: add Anthropic Claude Code subscription window tracking
- New headroom/subscription/ package:
- models.py: RateLimitWindow, ExtraUsage (cents->USD), SubscriptionSnapshot
with per-model 7d windows (opus/sonnet), HeadroomContribution,
WindowDiscrepancy, SubscriptionState
- client.py: async httpx client for GET /api/oauth/usage with OAuth token
resolution (env var -> credentials file -> live proxy header)
- tracker.py: background polling singleton (asyncio Task), notify_active(),
update_contribution(), anomaly detection (surge pricing, cache miss),
atomic persistence, OTEL callback
- session_tracking.py: JSONL transcript reader, Sonnet-normalised model
weights, compute_window_tokens() for 5h/7d window boundaries
- __init__.py: clean public re-exports
- Integration:
- proxy/models.py: subscription_tracking_enabled, poll_interval_s,
active_window_s config fields
- proxy/server.py: tracker startup/shutdown, /subscription-window endpoint,
subscription_window key in /stats
- proxy/handlers/anthropic.py: OAuth Bearer detection, notify_active() +
update_contribution() per request
- observability/metrics.py: 5 observable gauges (5h/7d utilisation %,
5h/7d seconds-to-reset, overage USD) using Observation type
- dashboard/templates/dashboard.html: subscription window panel with 5h/7d
progress bars, countdown, overage, Headroom contribution bar + anomaly
alerts
- Tests: 33 new unit tests in tests/test_subscription_tracker.py (all pass)
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
JerrettDavis
parent
1fee4ac8fd
commit
4329213e57
11 changed files with 1896 additions and 1 deletions
|
|
@ -167,7 +167,136 @@
|
|||
</div>
|
||||
</template>
|
||||
|
||||
<!-- Prefix Cache Impact (Provider-Aware) -->
|
||||
<!-- Anthropic Subscription Window -->
|
||||
<template x-if="stats.subscription_window && stats.subscription_window.latest">
|
||||
<div class="bg-surface rounded-lg p-4 border border-border mb-6">
|
||||
<div class="flex justify-between items-center mb-4">
|
||||
<div class="text-sm font-medium text-gray-300">Anthropic Subscription Window</div>
|
||||
<div class="flex items-center gap-2">
|
||||
<span class="w-2 h-2 rounded-full pulse-live"
|
||||
:class="stats.subscription_window.last_active_at ? 'bg-emerald-400' : 'bg-gray-600'"></span>
|
||||
<span class="text-xs text-gray-500" x-text="stats.subscription_window.last_active_at ? 'Active' : 'Idle'"></span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- 5h / 7d progress bars -->
|
||||
<div class="grid grid-cols-1 md:grid-cols-2 gap-4 mb-4">
|
||||
<!-- 5-hour window -->
|
||||
<template x-if="stats.subscription_window.latest.five_hour">
|
||||
<div>
|
||||
<div class="flex justify-between text-xs mb-1">
|
||||
<span class="text-gray-400 uppercase tracking-wide">5-Hour Window</span>
|
||||
<span class="tabular-nums font-mono"
|
||||
:class="stats.subscription_window.latest.five_hour.utilization_pct > 80 ? 'text-red-400' : stats.subscription_window.latest.five_hour.utilization_pct > 60 ? 'text-amber-400' : 'text-emerald-400'"
|
||||
x-text="(stats.subscription_window.latest.five_hour.utilization_pct || 0).toFixed(1) + '%'"></span>
|
||||
</div>
|
||||
<div class="w-full h-3 bg-border rounded-full overflow-hidden">
|
||||
<div class="h-full rounded-full transition-all duration-500"
|
||||
:class="stats.subscription_window.latest.five_hour.utilization_pct > 80 ? 'bg-red-500' : stats.subscription_window.latest.five_hour.utilization_pct > 60 ? 'bg-amber-500' : 'bg-emerald-500'"
|
||||
:style="'width: ' + Math.min(stats.subscription_window.latest.five_hour.utilization_pct || 0, 100) + '%'"></div>
|
||||
</div>
|
||||
<div class="text-xs text-gray-500 mt-1"
|
||||
x-text="'Resets in ' + formatResetTime(stats.subscription_window.latest.five_hour.seconds_to_reset)"></div>
|
||||
</div>
|
||||
</template>
|
||||
|
||||
<!-- 7-day window -->
|
||||
<template x-if="stats.subscription_window.latest.seven_day">
|
||||
<div>
|
||||
<div class="flex justify-between text-xs mb-1">
|
||||
<span class="text-gray-400 uppercase tracking-wide">7-Day Window</span>
|
||||
<span class="tabular-nums font-mono"
|
||||
:class="stats.subscription_window.latest.seven_day.utilization_pct > 80 ? 'text-red-400' : stats.subscription_window.latest.seven_day.utilization_pct > 60 ? 'text-amber-400' : 'text-emerald-400'"
|
||||
x-text="(stats.subscription_window.latest.seven_day.utilization_pct || 0).toFixed(1) + '%'"></span>
|
||||
</div>
|
||||
<div class="w-full h-3 bg-border rounded-full overflow-hidden">
|
||||
<div class="h-full rounded-full transition-all duration-500"
|
||||
:class="stats.subscription_window.latest.seven_day.utilization_pct > 80 ? 'bg-red-500' : stats.subscription_window.latest.seven_day.utilization_pct > 60 ? 'bg-amber-500' : 'bg-emerald-500'"
|
||||
:style="'width: ' + Math.min(stats.subscription_window.latest.seven_day.utilization_pct || 0, 100) + '%'"></div>
|
||||
</div>
|
||||
<div class="text-xs text-gray-500 mt-1"
|
||||
x-text="'Resets in ' + formatResetTime(stats.subscription_window.latest.seven_day.seconds_to_reset)"></div>
|
||||
</div>
|
||||
</template>
|
||||
</div>
|
||||
|
||||
<!-- Overage / extra usage -->
|
||||
<template x-if="stats.subscription_window.latest.extra_usage && stats.subscription_window.latest.extra_usage.is_enabled">
|
||||
<div class="border-t border-border pt-3 mb-3">
|
||||
<div class="flex justify-between text-xs mb-1">
|
||||
<span class="text-gray-400 uppercase tracking-wide">Extra Usage (Overage)</span>
|
||||
<span class="text-amber-400 tabular-nums font-mono"
|
||||
x-text="'$' + formatCurrency(stats.subscription_window.latest.extra_usage.used_credits_usd || 0) + (stats.subscription_window.latest.extra_usage.monthly_limit_usd ? ' / $' + formatCurrency(stats.subscription_window.latest.extra_usage.monthly_limit_usd) : '')"></span>
|
||||
</div>
|
||||
<template x-if="stats.subscription_window.latest.extra_usage.utilization_pct != null">
|
||||
<div class="w-full h-2 bg-border rounded-full overflow-hidden">
|
||||
<div class="h-full bg-amber-500 rounded-full transition-all duration-500"
|
||||
:style="'width: ' + Math.min(stats.subscription_window.latest.extra_usage.utilization_pct || 0, 100) + '%'"></div>
|
||||
</div>
|
||||
</template>
|
||||
</div>
|
||||
</template>
|
||||
|
||||
<!-- Headroom contribution row -->
|
||||
<template x-if="stats.subscription_window.contribution && stats.subscription_window.contribution.tokens_submitted > 0">
|
||||
<div class="border-t border-border pt-3">
|
||||
<div class="text-xs text-gray-400 uppercase tracking-wide mb-2">Headroom Contribution This Window</div>
|
||||
<div class="grid grid-cols-2 md:grid-cols-4 gap-3">
|
||||
<div>
|
||||
<div class="text-xs text-gray-500 mb-0.5">Efficiency</div>
|
||||
<div class="text-lg font-light tabular-nums text-accent"
|
||||
x-text="(stats.subscription_window.contribution.efficiency_pct || 0).toFixed(1) + '%'"></div>
|
||||
</div>
|
||||
<div>
|
||||
<div class="text-xs text-gray-500 mb-0.5">Tokens Saved</div>
|
||||
<div class="text-lg font-light tabular-nums text-emerald-400"
|
||||
x-text="formatNumber(stats.subscription_window.contribution.tokens_saved?.total || 0)"></div>
|
||||
</div>
|
||||
<div>
|
||||
<div class="text-xs text-gray-500 mb-0.5">Compression</div>
|
||||
<div class="text-sm tabular-nums text-gray-300"
|
||||
x-text="formatNumber(stats.subscription_window.contribution.tokens_saved?.compression || 0) + ' tok'"></div>
|
||||
</div>
|
||||
<div>
|
||||
<div class="text-xs text-gray-500 mb-0.5">Cache Reads</div>
|
||||
<div class="text-sm tabular-nums text-cyan-400"
|
||||
x-text="formatNumber(stats.subscription_window.contribution.tokens_saved?.cache_reads || 0) + ' tok'"></div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Efficiency bar: tokens submitted vs what would have been sent without headroom -->
|
||||
<template x-if="stats.subscription_window.contribution.raw_without_headroom > 0">
|
||||
<div class="mt-3">
|
||||
<div class="flex justify-between text-xs text-gray-500 mb-1">
|
||||
<span>Raw vs Submitted</span>
|
||||
<span x-text="formatNumber(stats.subscription_window.contribution.tokens_submitted) + ' / ' + formatNumber(stats.subscription_window.contribution.raw_without_headroom) + ' tokens forwarded'"></span>
|
||||
</div>
|
||||
<div class="w-full h-2 bg-border rounded-full overflow-hidden flex">
|
||||
<div class="h-full bg-accent rounded-l-full"
|
||||
:style="'width: ' + ((stats.subscription_window.contribution.tokens_submitted / stats.subscription_window.contribution.raw_without_headroom) * 100).toFixed(1) + '%'"></div>
|
||||
</div>
|
||||
<div class="flex gap-4 mt-1 text-xs text-gray-500">
|
||||
<span class="flex items-center gap-1"><span class="w-2 h-2 rounded-full bg-accent inline-block"></span>Forwarded to Anthropic</span>
|
||||
<span class="flex items-center gap-1"><span class="w-2 h-2 rounded-full bg-border inline-block"></span>Saved by Headroom</span>
|
||||
</div>
|
||||
</div>
|
||||
</template>
|
||||
</div>
|
||||
</template>
|
||||
|
||||
<!-- Discrepancies -->
|
||||
<template x-if="stats.subscription_window.discrepancies && stats.subscription_window.discrepancies.length > 0">
|
||||
<div class="border-t border-border pt-3 mt-3">
|
||||
<div class="text-xs text-amber-400 uppercase tracking-wide mb-2">⚠ Anomalies Detected</div>
|
||||
<template x-for="d in stats.subscription_window.discrepancies" :key="d.kind + d.description">
|
||||
<div class="text-xs mb-1"
|
||||
:class="d.severity === 'alert' ? 'text-red-400' : 'text-amber-400'"
|
||||
x-text="d.description"></div>
|
||||
</template>
|
||||
</div>
|
||||
</template>
|
||||
</div>
|
||||
</template>
|
||||
<template x-if="(stats.prefix_cache?.totals?.requests || 0) > 0">
|
||||
<div class="bg-surface rounded-lg p-4 border border-border mb-6">
|
||||
<div class="flex justify-between items-center mb-3">
|
||||
|
|
@ -948,6 +1077,15 @@
|
|||
return '0.00';
|
||||
},
|
||||
|
||||
formatResetTime(seconds) {
|
||||
if (seconds == null || seconds <= 0) return 'now';
|
||||
const h = Math.floor(seconds / 3600);
|
||||
const m = Math.floor((seconds % 3600) / 60);
|
||||
if (h > 0) return h + 'h ' + m + 'm';
|
||||
if (m > 0) return m + 'm ' + Math.floor(seconds % 60) + 's';
|
||||
return Math.floor(seconds) + 's';
|
||||
},
|
||||
|
||||
formatTime(ts) {
|
||||
if (!ts) return '-';
|
||||
const d = new Date(ts);
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from threading import Lock
|
|||
from typing import Any, Literal
|
||||
|
||||
from opentelemetry import metrics
|
||||
from opentelemetry.metrics import CallbackOptions, Observation
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
@ -260,6 +261,60 @@ class HeadroomOtelMetrics:
|
|||
unit="1",
|
||||
)
|
||||
|
||||
# Backing values updated by record_subscription_window()
|
||||
self._sub_5h_util_val: float = 0.0
|
||||
self._sub_7d_util_val: float = 0.0
|
||||
self._sub_5h_reset_val: float = 0.0
|
||||
self._sub_7d_reset_val: float = 0.0
|
||||
self._sub_overage_val: float = 0.0
|
||||
|
||||
# Subscription window gauges (Anthropic OAuth accounts)
|
||||
def _cb_5h_util(opts: CallbackOptions) -> list[Observation]:
|
||||
return [Observation(self._sub_5h_util_val)]
|
||||
|
||||
def _cb_7d_util(opts: CallbackOptions) -> list[Observation]:
|
||||
return [Observation(self._sub_7d_util_val)]
|
||||
|
||||
def _cb_5h_reset(opts: CallbackOptions) -> list[Observation]:
|
||||
return [Observation(self._sub_5h_reset_val)]
|
||||
|
||||
def _cb_7d_reset(opts: CallbackOptions) -> list[Observation]:
|
||||
return [Observation(self._sub_7d_reset_val)]
|
||||
|
||||
def _cb_overage(opts: CallbackOptions) -> list[Observation]:
|
||||
return [Observation(self._sub_overage_val)]
|
||||
|
||||
self._meter.create_observable_gauge(
|
||||
"headroom.subscription.5h_utilization_pct",
|
||||
description="Anthropic 5-hour rate-limit window utilisation (0–100%).",
|
||||
unit="1",
|
||||
callbacks=[_cb_5h_util],
|
||||
)
|
||||
self._meter.create_observable_gauge(
|
||||
"headroom.subscription.7d_utilization_pct",
|
||||
description="Anthropic 7-day rate-limit window utilisation (0–100%).",
|
||||
unit="1",
|
||||
callbacks=[_cb_7d_util],
|
||||
)
|
||||
self._meter.create_observable_gauge(
|
||||
"headroom.subscription.5h_seconds_to_reset",
|
||||
description="Seconds until the Anthropic 5-hour window resets.",
|
||||
unit="s",
|
||||
callbacks=[_cb_5h_reset],
|
||||
)
|
||||
self._meter.create_observable_gauge(
|
||||
"headroom.subscription.7d_seconds_to_reset",
|
||||
description="Seconds until the Anthropic 7-day window resets.",
|
||||
unit="s",
|
||||
callbacks=[_cb_7d_reset],
|
||||
)
|
||||
self._meter.create_observable_gauge(
|
||||
"headroom.subscription.overage_usd",
|
||||
description="Anthropic extra-usage (overage) credits consumed in USD.",
|
||||
unit="USD",
|
||||
callbacks=[_cb_overage],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _attrs(**attrs: Any) -> dict[str, Any]:
|
||||
filtered: dict[str, Any] = {}
|
||||
|
|
@ -390,6 +445,24 @@ class HeadroomOtelMetrics:
|
|||
self._attrs(model=model, operation=operation, error_type=error_type),
|
||||
)
|
||||
|
||||
def record_subscription_window(self, state: dict[str, Any]) -> None:
|
||||
"""Update OTEL subscription gauge backing values from the tracker state dict."""
|
||||
latest = state.get("latest") or {}
|
||||
|
||||
five_hour = latest.get("five_hour") or {}
|
||||
if five_hour:
|
||||
self._sub_5h_util_val = float(five_hour.get("utilization_pct", 0.0))
|
||||
self._sub_5h_reset_val = float(five_hour.get("seconds_to_reset") or 0.0)
|
||||
|
||||
seven_day = latest.get("seven_day") or {}
|
||||
if seven_day:
|
||||
self._sub_7d_util_val = float(seven_day.get("utilization_pct", 0.0))
|
||||
self._sub_7d_reset_val = float(seven_day.get("seconds_to_reset") or 0.0)
|
||||
|
||||
extra = latest.get("extra_usage") or {}
|
||||
if extra.get("is_enabled"):
|
||||
self._sub_overage_val = float(extra.get("used_credits_usd") or 0.0)
|
||||
|
||||
|
||||
def get_otel_metrics() -> HeadroomOtelMetrics:
|
||||
global _global_metrics
|
||||
|
|
|
|||
|
|
@ -383,6 +383,15 @@ class AnthropicHandlerMixin:
|
|||
headers.pop("content-length", None)
|
||||
tags = self._extract_tags(headers)
|
||||
|
||||
# Subscription tracker: notify on OAuth requests (not API-key requests)
|
||||
_auth_header = headers.get("authorization", "")
|
||||
if _auth_header.startswith("Bearer ") and not _auth_header.startswith("Bearer sk-ant-api"):
|
||||
from headroom.subscription.tracker import get_subscription_tracker as _get_sub_tracker
|
||||
|
||||
_sub_tracker = _get_sub_tracker()
|
||||
if _sub_tracker is not None:
|
||||
_sub_tracker.notify_active(_auth_header)
|
||||
|
||||
# Rate limiting
|
||||
if self.rate_limiter:
|
||||
api_key = headers.get("x-api-key", "")
|
||||
|
|
@ -1359,6 +1368,22 @@ class AnthropicHandlerMixin:
|
|||
uncached_input_tokens=uncached_input_tokens,
|
||||
)
|
||||
|
||||
# Subscription tracker: update headroom contribution counters
|
||||
if _auth_header.startswith("Bearer ") and not _auth_header.startswith(
|
||||
"Bearer sk-ant-api"
|
||||
):
|
||||
from headroom.subscription.tracker import (
|
||||
get_subscription_tracker as _get_sub_tracker,
|
||||
)
|
||||
|
||||
_sub_tracker = _get_sub_tracker()
|
||||
if _sub_tracker is not None:
|
||||
_sub_tracker.update_contribution(
|
||||
tokens_submitted=optimized_tokens,
|
||||
tokens_saved_compression=tokens_saved,
|
||||
tokens_saved_cache_reads=cr_tokens,
|
||||
)
|
||||
|
||||
# Log request
|
||||
if self.logger:
|
||||
self.logger.log(
|
||||
|
|
|
|||
|
|
@ -199,3 +199,8 @@ class ProxyConfig:
|
|||
|
||||
# Compression Hooks
|
||||
hooks: Any = None
|
||||
|
||||
# Subscription Window Tracking (Anthropic OAuth accounts)
|
||||
subscription_tracking_enabled: bool = True
|
||||
subscription_poll_interval_s: int = 10
|
||||
subscription_active_window_s: int = 60
|
||||
|
|
|
|||
|
|
@ -127,6 +127,11 @@ from headroom.proxy.prometheus_metrics import PrometheusMetrics # noqa: F401
|
|||
from headroom.proxy.rate_limiter import TokenBucketRateLimiter # noqa: F401
|
||||
from headroom.proxy.request_logger import RequestLogger # noqa: F401
|
||||
from headroom.proxy.semantic_cache import SemanticCache # noqa: F401
|
||||
from headroom.subscription.tracker import (
|
||||
configure_subscription_tracker,
|
||||
get_subscription_tracker,
|
||||
shutdown_subscription_tracker,
|
||||
)
|
||||
from headroom.telemetry import get_telemetry_collector
|
||||
from headroom.telemetry.beacon import is_telemetry_enabled
|
||||
from headroom.telemetry.toin import get_toin
|
||||
|
|
@ -662,6 +667,21 @@ class HeadroomProxy(
|
|||
logger.info("CCR: DISABLED")
|
||||
logger.info(f"Savings history: {self.metrics.savings_tracker.storage_path}")
|
||||
|
||||
# Subscription window tracker (Anthropic OAuth accounts)
|
||||
if self.config.subscription_tracking_enabled:
|
||||
tracker = configure_subscription_tracker(
|
||||
poll_interval_s=self.config.subscription_poll_interval_s,
|
||||
active_window_s=self.config.subscription_active_window_s,
|
||||
)
|
||||
await tracker.start()
|
||||
logger.info(
|
||||
"Subscription tracking: ENABLED "
|
||||
f"(poll_interval={self.config.subscription_poll_interval_s}s, "
|
||||
f"active_window={self.config.subscription_active_window_s}s)"
|
||||
)
|
||||
else:
|
||||
logger.info("Subscription tracking: DISABLED")
|
||||
|
||||
# Log anonymous telemetry status so operators can see it in the log stream
|
||||
if is_telemetry_enabled():
|
||||
logger.info(
|
||||
|
|
@ -680,6 +700,9 @@ class HeadroomProxy(
|
|||
if self.memory_handler and hasattr(self.memory_handler, "close"):
|
||||
await self.memory_handler.close()
|
||||
|
||||
# Stop subscription tracker
|
||||
await shutdown_subscription_tracker()
|
||||
|
||||
# Print final stats
|
||||
self._print_summary()
|
||||
|
||||
|
|
@ -1367,6 +1390,9 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
"cache": await proxy.cache.stats() if proxy.cache else None,
|
||||
"rate_limiter": await proxy.rate_limiter.stats() if proxy.rate_limiter else None,
|
||||
"recent_requests": proxy.logger.get_recent(10) if proxy.logger else [],
|
||||
"subscription_window": get_subscription_tracker().state
|
||||
if get_subscription_tracker()
|
||||
else None,
|
||||
}
|
||||
|
||||
@app.get("/stats-history")
|
||||
|
|
@ -1385,6 +1411,17 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
|
||||
return proxy.metrics.savings_tracker.history_response()
|
||||
|
||||
@app.get("/subscription-window")
|
||||
async def subscription_window():
|
||||
"""Current Anthropic subscription window utilisation and Headroom contribution."""
|
||||
tracker = get_subscription_tracker()
|
||||
if tracker is None:
|
||||
return JSONResponse(
|
||||
status_code=503,
|
||||
content={"error": "Subscription tracking is not enabled"},
|
||||
)
|
||||
return JSONResponse(content=tracker.state)
|
||||
|
||||
@app.get("/metrics")
|
||||
async def metrics():
|
||||
"""Prometheus metrics endpoint."""
|
||||
|
|
|
|||
34
headroom/subscription/__init__.py
Normal file
34
headroom/subscription/__init__.py
Normal file
|
|
@ -0,0 +1,34 @@
|
|||
"""Subscription window tracking for Anthropic Claude Code accounts."""
|
||||
|
||||
from headroom.subscription.client import SubscriptionClient, read_cached_oauth_token
|
||||
from headroom.subscription.models import (
|
||||
ExtraUsage,
|
||||
HeadroomContribution,
|
||||
RateLimitWindow,
|
||||
SubscriptionSnapshot,
|
||||
SubscriptionState,
|
||||
WindowDiscrepancy,
|
||||
WindowTokens,
|
||||
)
|
||||
from headroom.subscription.tracker import (
|
||||
SubscriptionTracker,
|
||||
configure_subscription_tracker,
|
||||
get_subscription_tracker,
|
||||
shutdown_subscription_tracker,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"ExtraUsage",
|
||||
"HeadroomContribution",
|
||||
"RateLimitWindow",
|
||||
"SubscriptionClient",
|
||||
"SubscriptionSnapshot",
|
||||
"SubscriptionState",
|
||||
"SubscriptionTracker",
|
||||
"WindowDiscrepancy",
|
||||
"WindowTokens",
|
||||
"configure_subscription_tracker",
|
||||
"get_subscription_tracker",
|
||||
"read_cached_oauth_token",
|
||||
"shutdown_subscription_tracker",
|
||||
]
|
||||
131
headroom/subscription/client.py
Normal file
131
headroom/subscription/client.py
Normal file
|
|
@ -0,0 +1,131 @@
|
|||
"""Async HTTP client for Anthropic's OAuth usage API.
|
||||
|
||||
Endpoint: GET https://api.anthropic.com/api/oauth/usage
|
||||
Required header: anthropic-beta: oauth-2025-04-20
|
||||
Auth: Authorization: Bearer <oauth_access_token>
|
||||
|
||||
Token resolution order (highest → lowest priority):
|
||||
1. Explicit token passed to :meth:`fetch`
|
||||
2. ``CLAUDE_CODE_OAUTH_TOKEN`` env-var
|
||||
3. ``~/.claude/.credentials.json`` → ``claudeAiOauth.accessToken``
|
||||
(respects ``CLAUDE_CONFIG_DIR`` env-var override)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from headroom.subscription.models import SubscriptionSnapshot
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_USAGE_URL = "https://api.anthropic.com/api/oauth/usage"
|
||||
_BETA_HEADER = "oauth-2025-04-20"
|
||||
_TOKEN_EXPIRY_BUFFER_S = 60
|
||||
|
||||
|
||||
def _credentials_path() -> Path:
|
||||
base = os.environ.get("CLAUDE_CONFIG_DIR") or os.path.expanduser("~/.claude")
|
||||
return Path(base) / ".credentials.json"
|
||||
|
||||
|
||||
def _load_credentials_file() -> dict[str, Any] | None:
|
||||
"""Load raw credentials dict from the Claude Code credentials file."""
|
||||
path = _credentials_path()
|
||||
try:
|
||||
with path.open() as fh:
|
||||
return json.load(fh) # type: ignore[no-any-return]
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
except Exception as exc:
|
||||
logger.debug("Cannot read credentials file %s: %s", path, exc)
|
||||
return None
|
||||
|
||||
|
||||
def read_cached_oauth_token() -> str | None:
|
||||
"""Resolve a stored OAuth token for background polling (no request needed).
|
||||
|
||||
Returns the raw access token string if found and not expired, else None.
|
||||
"""
|
||||
# 1. Env var
|
||||
env_token = os.environ.get("CLAUDE_CODE_OAUTH_TOKEN", "").strip()
|
||||
if env_token:
|
||||
return env_token
|
||||
|
||||
# 2. Credentials file
|
||||
creds = _load_credentials_file()
|
||||
if not creds:
|
||||
return None
|
||||
oauth = creds.get("claudeAiOauth") or {}
|
||||
token = oauth.get("accessToken") or ""
|
||||
if not token:
|
||||
return None
|
||||
|
||||
# Check expiry (Anthropic stores timestamp in milliseconds)
|
||||
expires_at_ms = oauth.get("expiresAt")
|
||||
if expires_at_ms is not None:
|
||||
import time
|
||||
|
||||
now_ms = time.time() * 1000
|
||||
if now_ms >= (expires_at_ms - _TOKEN_EXPIRY_BUFFER_S * 1000):
|
||||
logger.debug("Cached OAuth token expired; skipping background poll")
|
||||
return None
|
||||
|
||||
return token
|
||||
|
||||
|
||||
class SubscriptionClient:
|
||||
"""Thin async wrapper around the Anthropic OAuth usage endpoint."""
|
||||
|
||||
def __init__(self, timeout: float = 10.0) -> None:
|
||||
self._timeout = timeout
|
||||
|
||||
async def fetch(self, token: str | None = None) -> SubscriptionSnapshot | None:
|
||||
"""Fetch current subscription window data.
|
||||
|
||||
:param token: OAuth access token. When *None*, falls back to
|
||||
:func:`read_cached_oauth_token`.
|
||||
:returns: :class:`SubscriptionSnapshot` or *None* on auth failure /
|
||||
unsupported account.
|
||||
"""
|
||||
resolved = (token or "").strip() or read_cached_oauth_token()
|
||||
if not resolved:
|
||||
logger.debug("No OAuth token available for subscription polling")
|
||||
return None
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {resolved}",
|
||||
"anthropic-beta": _BETA_HEADER,
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=self._timeout) as client:
|
||||
resp = await client.get(_USAGE_URL, headers=headers)
|
||||
|
||||
if resp.status_code == 401:
|
||||
logger.debug("OAuth token rejected (401) by Anthropic usage API")
|
||||
return None
|
||||
if resp.status_code == 404:
|
||||
# API key accounts (non-subscription) return 404
|
||||
logger.debug("Subscription usage API returned 404; likely API-key account")
|
||||
return None
|
||||
if resp.status_code != 200:
|
||||
logger.warning("Anthropic usage API returned %s", resp.status_code)
|
||||
return None
|
||||
|
||||
data: dict[str, Any] = resp.json()
|
||||
return SubscriptionSnapshot.from_api_response(data, token=resolved)
|
||||
|
||||
except httpx.TimeoutException:
|
||||
logger.debug("Timeout fetching Anthropic subscription window")
|
||||
return None
|
||||
except Exception as exc:
|
||||
logger.warning("Error fetching subscription window: %s", exc)
|
||||
return None
|
||||
395
headroom/subscription/models.py
Normal file
395
headroom/subscription/models.py
Normal file
|
|
@ -0,0 +1,395 @@
|
|||
"""Data models for Anthropic subscription window tracking.
|
||||
|
||||
Mirrors the Anthropic OAuth usage API response exactly, including:
|
||||
- five_hour / seven_day rolling windows (utilization + reset times)
|
||||
- seven_day_opus / seven_day_sonnet per-model 7-day windows
|
||||
- extra_usage overage block (credits stored in cents by Anthropic)
|
||||
- Headroom contribution: tokens conserved by compression, rtk, cache
|
||||
- Window discrepancy detection (surge pricing, cache-miss anomalies)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _utc_now() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
def _to_utc_iso(dt: datetime) -> str:
|
||||
return dt.astimezone(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
|
||||
|
||||
|
||||
def _parse_timestamp(value: Any) -> datetime | None:
|
||||
if not isinstance(value, str) or not value:
|
||||
return None
|
||||
normalized = value.replace("Z", "+00:00")
|
||||
try:
|
||||
dt = datetime.fromisoformat(normalized)
|
||||
except ValueError:
|
||||
return None
|
||||
if dt.tzinfo is None:
|
||||
dt = dt.replace(tzinfo=timezone.utc)
|
||||
return dt.astimezone(timezone.utc)
|
||||
|
||||
|
||||
def _safe_float(value: Any) -> float | None:
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
return float(value)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def _safe_int(value: Any) -> int | None:
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Rate-limit window (five_hour / seven_day / seven_day_opus / seven_day_sonnet)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class RateLimitWindow:
|
||||
"""A single rolling rate-limit window returned by the Anthropic usage API.
|
||||
|
||||
``used`` and ``limit`` are in Anthropic's internal token-equivalent units
|
||||
(not raw tokens; Anthropic weights tokens differently per model family).
|
||||
``utilization_pct`` is the authoritative 0–100 % figure from the API.
|
||||
"""
|
||||
|
||||
used: int = 0
|
||||
limit: int = 0
|
||||
utilization_pct: float = 0.0
|
||||
resets_at: datetime | None = None
|
||||
|
||||
@classmethod
|
||||
def from_api_dict(cls, data: dict[str, Any]) -> RateLimitWindow:
|
||||
return cls(
|
||||
used=int(data.get("used") or 0),
|
||||
limit=int(data.get("limit") or 0),
|
||||
utilization_pct=float(data.get("utilization") or 0.0),
|
||||
resets_at=_parse_timestamp(data.get("resets_at")),
|
||||
)
|
||||
|
||||
def seconds_to_reset(self, *, now: datetime | None = None) -> float | None:
|
||||
if self.resets_at is None:
|
||||
return None
|
||||
return max((self.resets_at - (now or _utc_now())).total_seconds(), 0.0)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"used": self.used,
|
||||
"limit": self.limit,
|
||||
"utilization_pct": round(self.utilization_pct, 2),
|
||||
"resets_at": _to_utc_iso(self.resets_at) if self.resets_at else None,
|
||||
"seconds_to_reset": self.seconds_to_reset(),
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Extra-usage / overage block
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExtraUsage:
|
||||
"""Overage / extra-usage block from the Anthropic usage API.
|
||||
|
||||
``monthly_limit_cents`` and ``used_credits_cents`` are in US cents as
|
||||
returned by the API (divide by 100 for USD).
|
||||
"""
|
||||
|
||||
is_enabled: bool = False
|
||||
monthly_limit_cents: int | None = None
|
||||
used_credits_cents: int | None = None
|
||||
utilization_pct: float | None = None
|
||||
|
||||
@classmethod
|
||||
def from_api_dict(cls, data: dict[str, Any]) -> ExtraUsage:
|
||||
return cls(
|
||||
is_enabled=bool(data.get("is_enabled", False)),
|
||||
monthly_limit_cents=_safe_int(data.get("monthly_limit")),
|
||||
used_credits_cents=_safe_int(data.get("used_credits")),
|
||||
utilization_pct=_safe_float(data.get("utilization")),
|
||||
)
|
||||
|
||||
@property
|
||||
def monthly_limit_usd(self) -> float | None:
|
||||
if self.monthly_limit_cents is None:
|
||||
return None
|
||||
return self.monthly_limit_cents / 100.0
|
||||
|
||||
@property
|
||||
def used_credits_usd(self) -> float | None:
|
||||
if self.used_credits_cents is None:
|
||||
return None
|
||||
return self.used_credits_cents / 100.0
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"is_enabled": self.is_enabled,
|
||||
"monthly_limit_usd": round(self.monthly_limit_usd, 2)
|
||||
if self.monthly_limit_usd is not None
|
||||
else None,
|
||||
"used_credits_usd": round(self.used_credits_usd, 4)
|
||||
if self.used_credits_usd is not None
|
||||
else None,
|
||||
"utilization_pct": round(self.utilization_pct, 2)
|
||||
if self.utilization_pct is not None
|
||||
else None,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Full snapshot from one API poll
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class SubscriptionSnapshot:
|
||||
"""One complete poll of GET /api/oauth/usage."""
|
||||
|
||||
five_hour: RateLimitWindow = field(default_factory=RateLimitWindow)
|
||||
seven_day: RateLimitWindow = field(default_factory=RateLimitWindow)
|
||||
seven_day_opus: RateLimitWindow | None = None
|
||||
seven_day_sonnet: RateLimitWindow | None = None
|
||||
extra_usage: ExtraUsage = field(default_factory=ExtraUsage)
|
||||
polled_at: datetime = field(default_factory=_utc_now)
|
||||
token_prefix: str = ""
|
||||
"""First 8 chars of the OAuth token (for multi-account detection)."""
|
||||
|
||||
@classmethod
|
||||
def from_api_response(cls, data: dict[str, Any], *, token: str = "") -> SubscriptionSnapshot:
|
||||
snap = cls(token_prefix=token[:8] if token else "")
|
||||
if "five_hour" in data and data["five_hour"]:
|
||||
snap.five_hour = RateLimitWindow.from_api_dict(data["five_hour"])
|
||||
if "seven_day" in data and data["seven_day"]:
|
||||
snap.seven_day = RateLimitWindow.from_api_dict(data["seven_day"])
|
||||
if "seven_day_opus" in data and data["seven_day_opus"]:
|
||||
snap.seven_day_opus = RateLimitWindow.from_api_dict(data["seven_day_opus"])
|
||||
if "seven_day_sonnet" in data and data["seven_day_sonnet"]:
|
||||
snap.seven_day_sonnet = RateLimitWindow.from_api_dict(data["seven_day_sonnet"])
|
||||
if "extra_usage" in data and data["extra_usage"]:
|
||||
snap.extra_usage = ExtraUsage.from_api_dict(data["extra_usage"])
|
||||
return snap
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
d: dict[str, Any] = {
|
||||
"five_hour": self.five_hour.to_dict(),
|
||||
"seven_day": self.seven_day.to_dict(),
|
||||
"extra_usage": self.extra_usage.to_dict(),
|
||||
"polled_at": _to_utc_iso(self.polled_at),
|
||||
"token_prefix": self.token_prefix,
|
||||
}
|
||||
if self.seven_day_opus:
|
||||
d["seven_day_opus"] = self.seven_day_opus.to_dict()
|
||||
if self.seven_day_sonnet:
|
||||
d["seven_day_sonnet"] = self.seven_day_sonnet.to_dict()
|
||||
return d
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Transcript-based window token breakdown
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class WindowTokens:
|
||||
"""Token breakdown from Claude transcript JSONL files for one time window."""
|
||||
|
||||
input: int = 0
|
||||
output: int = 0
|
||||
cache_reads: int = 0
|
||||
cache_writes_5m: int = 0
|
||||
cache_writes_1h: int = 0
|
||||
cache_writes_total: int = 0
|
||||
by_model: dict[str, dict[str, int]] = field(default_factory=dict)
|
||||
weighted_token_equivalent: float = 0.0
|
||||
"""Sonnet-normalised weighted total (opus×2, sonnet×1, haiku×0.5)."""
|
||||
|
||||
def total_raw(self) -> int:
|
||||
return self.input + self.output + self.cache_reads + self.cache_writes_total
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"input": self.input,
|
||||
"output": self.output,
|
||||
"cache_reads": self.cache_reads,
|
||||
"cache_writes_5m": self.cache_writes_5m,
|
||||
"cache_writes_1h": self.cache_writes_1h,
|
||||
"cache_writes_total": self.cache_writes_total,
|
||||
"total_raw": self.total_raw(),
|
||||
"weighted_token_equivalent": round(self.weighted_token_equivalent, 1),
|
||||
"by_model": self.by_model,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Headroom contribution estimate
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class HeadroomContribution:
|
||||
"""Tokens conserved within the current 5h window by Headroom's layers.
|
||||
|
||||
These are cumulative counters reset when the 5h window rolls over.
|
||||
"""
|
||||
|
||||
tokens_submitted: int = 0
|
||||
"""Raw input tokens actually forwarded to Anthropic by the proxy."""
|
||||
|
||||
tokens_saved_compression: int = 0
|
||||
"""Input tokens removed by proxy compression."""
|
||||
|
||||
tokens_saved_rtk: int = 0
|
||||
"""Tokens avoided by CLI filtering (rtk) before reaching context."""
|
||||
|
||||
tokens_saved_cache_reads: int = 0
|
||||
"""Input tokens served from Anthropic prefix-cache (discounted reads)."""
|
||||
|
||||
compression_savings_usd: float = 0.0
|
||||
cache_savings_usd: float = 0.0
|
||||
|
||||
def total_saved(self) -> int:
|
||||
return self.tokens_saved_compression + self.tokens_saved_rtk + self.tokens_saved_cache_reads
|
||||
|
||||
def total_savings_usd(self) -> float:
|
||||
return self.compression_savings_usd + self.cache_savings_usd
|
||||
|
||||
def raw_without_headroom(self) -> int:
|
||||
return self.tokens_submitted + self.tokens_saved_compression + self.tokens_saved_rtk
|
||||
|
||||
def efficiency_pct(self) -> float:
|
||||
raw = self.raw_without_headroom()
|
||||
if raw == 0:
|
||||
return 0.0
|
||||
return round(self.total_saved() / raw * 100, 1)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"tokens_submitted": self.tokens_submitted,
|
||||
"tokens_saved": {
|
||||
"compression": self.tokens_saved_compression,
|
||||
"rtk": self.tokens_saved_rtk,
|
||||
"cache_reads": self.tokens_saved_cache_reads,
|
||||
"total": self.total_saved(),
|
||||
},
|
||||
"raw_without_headroom": self.raw_without_headroom(),
|
||||
"efficiency_pct": self.efficiency_pct(),
|
||||
"savings_usd": {
|
||||
"compression": round(self.compression_savings_usd, 4),
|
||||
"cache": round(self.cache_savings_usd, 4),
|
||||
"total": round(self.total_savings_usd(), 4),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Anomaly / discrepancy record
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class WindowDiscrepancy:
|
||||
"""Detected anomaly between expected and API-reported utilization."""
|
||||
|
||||
kind: str
|
||||
"""'surge_pricing' | 'cache_miss' | 'none'"""
|
||||
|
||||
description: str = ""
|
||||
severity: str = "info"
|
||||
"""'info' | 'warning' | 'alert'"""
|
||||
|
||||
expected_utilization_pct: float | None = None
|
||||
actual_utilization_pct: float | None = None
|
||||
delta_pct: float | None = None
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"kind": self.kind,
|
||||
"description": self.description,
|
||||
"severity": self.severity,
|
||||
"expected_utilization_pct": self.expected_utilization_pct,
|
||||
"actual_utilization_pct": self.actual_utilization_pct,
|
||||
"delta_pct": self.delta_pct,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Full tracker state
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class SubscriptionState:
|
||||
"""Persistent state for the subscription tracker."""
|
||||
|
||||
latest: SubscriptionSnapshot | None = None
|
||||
window_tokens: WindowTokens | None = None
|
||||
"""Transcript-derived token breakdown for the current 5h window."""
|
||||
|
||||
contribution: HeadroomContribution = field(default_factory=HeadroomContribution)
|
||||
discrepancies: list[WindowDiscrepancy] = field(default_factory=list)
|
||||
history: list[SubscriptionSnapshot] = field(default_factory=list)
|
||||
|
||||
poll_count: int = 0
|
||||
poll_errors: int = 0
|
||||
last_error: str | None = None
|
||||
last_active_at: datetime | None = None
|
||||
|
||||
_MAX_HISTORY: int = field(default=100, init=False, repr=False)
|
||||
_MAX_DISCREPANCIES: int = field(default=20, init=False, repr=False)
|
||||
|
||||
def add_snapshot(self, snapshot: SubscriptionSnapshot) -> None:
|
||||
self.latest = snapshot
|
||||
self.history.append(snapshot)
|
||||
if len(self.history) > self._MAX_HISTORY:
|
||||
self.history = self.history[-self._MAX_HISTORY :]
|
||||
self.poll_count += 1
|
||||
|
||||
def mark_error(self, msg: str) -> None:
|
||||
self.poll_errors += 1
|
||||
self.last_error = msg
|
||||
|
||||
def add_discrepancy(self, d: WindowDiscrepancy) -> None:
|
||||
self.discrepancies.append(d)
|
||||
if len(self.discrepancies) > self._MAX_DISCREPANCIES:
|
||||
self.discrepancies = self.discrepancies[-self._MAX_DISCREPANCIES :]
|
||||
|
||||
def is_active(self, *, active_window_s: float = 60.0) -> bool:
|
||||
if self.last_active_at is None:
|
||||
return False
|
||||
return (_utc_now() - self.last_active_at).total_seconds() <= active_window_s
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"latest": self.latest.to_dict() if self.latest else None,
|
||||
"window_tokens": self.window_tokens.to_dict() if self.window_tokens else None,
|
||||
"contribution": self.contribution.to_dict(),
|
||||
"discrepancies": [d.to_dict() for d in self.discrepancies[-5:]],
|
||||
"poll_count": self.poll_count,
|
||||
"poll_errors": self.poll_errors,
|
||||
"last_error": self.last_error,
|
||||
"last_active_at": _to_utc_iso(self.last_active_at) if self.last_active_at else None,
|
||||
}
|
||||
|
||||
def to_persist_dict(self) -> dict[str, Any]:
|
||||
d = self.to_dict()
|
||||
d["history"] = [s.to_dict() for s in self.history[-20:]]
|
||||
return d
|
||||
189
headroom/subscription/session_tracking.py
Normal file
189
headroom/subscription/session_tracking.py
Normal file
|
|
@ -0,0 +1,189 @@
|
|||
"""Parse Claude Code transcript JSONL files for per-window token breakdowns.
|
||||
|
||||
Mirrors the approach in the ClaudeCacheTTLStatusLine TypeScript reference
|
||||
implementation (session-tracking.ts). Reads ~/.claude/projects/**/*.jsonl
|
||||
and aggregates token usage for entries whose timestamp falls within a window.
|
||||
|
||||
Model weights (Sonnet-normalised, empirical estimates):
|
||||
opus: 2.0× (higher rate-limit cost)
|
||||
sonnet: 1.0× (baseline)
|
||||
haiku: 0.5× (cheaper, lower rate-limit cost)
|
||||
|
||||
The weighted_token_equivalent lets callers detect surge pricing by comparing
|
||||
it against the API-reported utilisation × window_limit.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from headroom.subscription.models import WindowTokens
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Maximum bytes to read per transcript file (10 MB cap — generous, typical files <1 MB)
|
||||
_MAX_FILE_BYTES = 10 * 1024 * 1024
|
||||
|
||||
# Sonnet-normalised model family weights
|
||||
MODEL_FAMILY_WEIGHTS: dict[str, float] = {
|
||||
"opus": 2.0,
|
||||
"sonnet": 1.0,
|
||||
"haiku": 0.5,
|
||||
}
|
||||
DEFAULT_MODEL_WEIGHT: float = 1.0
|
||||
|
||||
|
||||
def _claude_config_dir() -> Path:
|
||||
base = os.environ.get("CLAUDE_CONFIG_DIR") or os.path.expanduser("~/.claude")
|
||||
return Path(base)
|
||||
|
||||
|
||||
def get_model_weight(model_id: str) -> float:
|
||||
"""Return the Sonnet-normalised weight for a model ID.
|
||||
|
||||
Matches against known family names using a word-boundary check.
|
||||
Falls back to DEFAULT_MODEL_WEIGHT for unrecognised models.
|
||||
"""
|
||||
lower = model_id.lower()
|
||||
import re
|
||||
|
||||
for family, weight in MODEL_FAMILY_WEIGHTS.items():
|
||||
if re.search(rf"(?<![a-z]){family}(?![a-z])", lower):
|
||||
return weight
|
||||
return DEFAULT_MODEL_WEIGHT
|
||||
|
||||
|
||||
def find_transcript_files() -> list[Path]:
|
||||
"""Return all .jsonl files under ~/.claude/projects."""
|
||||
projects = _claude_config_dir() / "projects"
|
||||
results: list[Path] = []
|
||||
_walk_jsonl(projects, results)
|
||||
return results
|
||||
|
||||
|
||||
def _walk_jsonl(directory: Path, results: list[Path]) -> None:
|
||||
try:
|
||||
entries = list(directory.iterdir())
|
||||
except (OSError, PermissionError):
|
||||
return
|
||||
for entry in entries:
|
||||
try:
|
||||
if entry.is_dir():
|
||||
_walk_jsonl(entry, results)
|
||||
elif entry.suffix == ".jsonl":
|
||||
results.append(entry)
|
||||
except OSError:
|
||||
continue
|
||||
|
||||
|
||||
def _read_transcript_lines(path: Path) -> list[str]:
|
||||
try:
|
||||
size = path.stat().st_size
|
||||
read_size = min(size, _MAX_FILE_BYTES)
|
||||
with path.open("rb") as fh:
|
||||
raw = fh.read(read_size)
|
||||
return [line for line in raw.decode("utf-8", errors="replace").splitlines() if line.strip()]
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
def _add_usage_to_tokens(dest: WindowTokens, usage: dict[str, Any]) -> None:
|
||||
dest.input += int(usage.get("input_tokens") or 0)
|
||||
dest.output += int(usage.get("output_tokens") or 0)
|
||||
dest.cache_reads += int(usage.get("cache_read_input_tokens") or 0)
|
||||
|
||||
cache_creation = usage.get("cache_creation") or {}
|
||||
w5m = int(cache_creation.get("ephemeral_5m_input_tokens") or 0)
|
||||
w1h = int(cache_creation.get("ephemeral_1h_input_tokens") or 0)
|
||||
total_writes = int(usage.get("cache_creation_input_tokens") or (w5m + w1h))
|
||||
|
||||
dest.cache_writes_5m += w5m
|
||||
dest.cache_writes_1h += w1h
|
||||
dest.cache_writes_total += total_writes
|
||||
|
||||
|
||||
def compute_window_tokens(start_ts: float, end_ts: float) -> WindowTokens:
|
||||
"""Sum transcript token usage for entries in [start_ts, end_ts).
|
||||
|
||||
Args:
|
||||
start_ts: Window start as a Unix timestamp (seconds).
|
||||
end_ts: Window end as a Unix timestamp (seconds).
|
||||
|
||||
Returns:
|
||||
:class:`WindowTokens` with aggregate + per-model breakdown and
|
||||
``weighted_token_equivalent`` (Sonnet-normalised).
|
||||
"""
|
||||
totals = WindowTokens()
|
||||
by_model: dict[str, WindowTokens] = {}
|
||||
unattributed = WindowTokens()
|
||||
|
||||
for path in find_transcript_files():
|
||||
for line in _read_transcript_lines(path):
|
||||
try:
|
||||
entry: dict[str, Any] = json.loads(line)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
continue
|
||||
|
||||
ts_str = entry.get("timestamp")
|
||||
if not ts_str:
|
||||
continue
|
||||
try:
|
||||
from datetime import datetime
|
||||
|
||||
dt = datetime.fromisoformat(str(ts_str).replace("Z", "+00:00"))
|
||||
ts = dt.timestamp()
|
||||
except (ValueError, TypeError):
|
||||
continue
|
||||
|
||||
if ts < start_ts or ts >= end_ts:
|
||||
continue
|
||||
|
||||
msg = entry.get("message") or {}
|
||||
usage = msg.get("usage")
|
||||
if not usage:
|
||||
continue
|
||||
|
||||
_add_usage_to_tokens(totals, usage)
|
||||
|
||||
model_id: str | None = msg.get("model")
|
||||
if model_id:
|
||||
if model_id not in by_model:
|
||||
by_model[model_id] = WindowTokens()
|
||||
_add_usage_to_tokens(by_model[model_id], usage)
|
||||
else:
|
||||
_add_usage_to_tokens(unattributed, usage)
|
||||
|
||||
# Compute Sonnet-normalised weighted equivalent
|
||||
model_weights: dict[str, float] = {}
|
||||
weighted = 0.0
|
||||
|
||||
for model_id, model_tokens in by_model.items():
|
||||
w = get_model_weight(model_id)
|
||||
model_weights[model_id] = w
|
||||
weighted += _total_token_count(model_tokens) * w
|
||||
|
||||
weighted += _total_token_count(unattributed) * DEFAULT_MODEL_WEIGHT
|
||||
|
||||
totals.weighted_token_equivalent = weighted
|
||||
totals.by_model = {mid: _window_tokens_to_dict(mt) for mid, mt in by_model.items()}
|
||||
|
||||
return totals
|
||||
|
||||
|
||||
def _total_token_count(t: WindowTokens) -> int:
|
||||
return t.input + t.output + t.cache_reads + t.cache_writes_total
|
||||
|
||||
|
||||
def _window_tokens_to_dict(t: WindowTokens) -> dict[str, int]:
|
||||
return {
|
||||
"input": t.input,
|
||||
"output": t.output,
|
||||
"cache_reads": t.cache_reads,
|
||||
"cache_writes_5m": t.cache_writes_5m,
|
||||
"cache_writes_1h": t.cache_writes_1h,
|
||||
"cache_writes_total": t.cache_writes_total,
|
||||
}
|
||||
434
headroom/subscription/tracker.py
Normal file
434
headroom/subscription/tracker.py
Normal file
|
|
@ -0,0 +1,434 @@
|
|||
"""Background subscription window tracker for Anthropic OAuth accounts.
|
||||
|
||||
Polls GET https://api.anthropic.com/api/oauth/usage on a configurable interval
|
||||
while there has been at least one active OAuth session within the last minute.
|
||||
Falls back to a stored token from ~/.claude/.credentials.json when no live
|
||||
request has come through the proxy recently.
|
||||
|
||||
Architecture:
|
||||
- Single asyncio.Task polling loop (started in start(), stopped via asyncio.Event)
|
||||
- Thread-safe state updates via threading.Lock (consistent with headroom patterns)
|
||||
- Atomic JSON persistence via tempfile + os.replace()
|
||||
- Module-level singleton via get_subscription_tracker() / configure_subscription_tracker()
|
||||
|
||||
Also reads Claude transcript JSONL files (via session_tracking module) to provide
|
||||
token breakdowns per window that enable:
|
||||
- Headroom efficiency metrics (tokens saved = raw - what proxy sent)
|
||||
- Surge pricing detection (API utilization vs expected from weighted tokens)
|
||||
- Cache miss detection (low cache_reads despite high input tokens)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from headroom.subscription.client import SubscriptionClient
|
||||
from headroom.subscription.models import (
|
||||
HeadroomContribution,
|
||||
SubscriptionSnapshot,
|
||||
SubscriptionState,
|
||||
WindowDiscrepancy,
|
||||
WindowTokens,
|
||||
_utc_now,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_DEFAULT_POLL_INTERVAL_S = 10
|
||||
_DEFAULT_ACTIVE_WINDOW_S = 60
|
||||
_PERSIST_FILE_ENV = "HEADROOM_SUBSCRIPTION_STATE_PATH"
|
||||
_DEFAULT_PERSIST_DIR = ".headroom"
|
||||
_DEFAULT_PERSIST_FILE = "subscription_state.json"
|
||||
|
||||
# Surge pricing threshold: if actual utilization is >N% higher than expected,
|
||||
# flag it as a potential surge pricing event.
|
||||
_SURGE_THRESHOLD_PCT = 15.0
|
||||
|
||||
# Cache miss threshold: if cache_reads < N% of total input when we expect
|
||||
# heavy caching (>50k input tokens in window), flag it.
|
||||
_CACHE_MISS_RATIO_THRESHOLD = 0.10
|
||||
|
||||
|
||||
def _get_persist_path() -> Path:
|
||||
env = os.environ.get(_PERSIST_FILE_ENV, "").strip()
|
||||
if env:
|
||||
return Path(env)
|
||||
return Path.home() / _DEFAULT_PERSIST_DIR / _DEFAULT_PERSIST_FILE
|
||||
|
||||
|
||||
class SubscriptionTracker:
|
||||
"""Background tracker for Anthropic Claude Code subscription windows.
|
||||
|
||||
Args:
|
||||
poll_interval_s: Seconds between polls while active (1–300).
|
||||
active_window_s: Seconds since last notify_active call that keeps
|
||||
polling alive (default 60 = 1 minute).
|
||||
persist_path: Where to persist state across restarts.
|
||||
client: Injected client (for testing); defaults to SubscriptionClient().
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
poll_interval_s: int = _DEFAULT_POLL_INTERVAL_S,
|
||||
active_window_s: float = _DEFAULT_ACTIVE_WINDOW_S,
|
||||
persist_path: Path | None = None,
|
||||
client: SubscriptionClient | None = None,
|
||||
) -> None:
|
||||
self._poll_interval_s = max(1, min(poll_interval_s, 300))
|
||||
self._active_window_s = max(5.0, active_window_s)
|
||||
self._persist_path = persist_path or _get_persist_path()
|
||||
self._client = client or SubscriptionClient()
|
||||
|
||||
self._lock = threading.Lock()
|
||||
self._state = SubscriptionState()
|
||||
self._current_token: str | None = None
|
||||
self._full_tokens: dict[str, int] = {} # token_prefix -> count of requests
|
||||
|
||||
self._stop_event: asyncio.Event | None = None
|
||||
self._poll_task: asyncio.Task[None] | None = None
|
||||
|
||||
self._load_persisted_state()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Lifecycle
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def start(self) -> None:
|
||||
"""Start the background polling loop."""
|
||||
if self._poll_task and not self._poll_task.done():
|
||||
return
|
||||
self._stop_event = asyncio.Event()
|
||||
self._poll_task = asyncio.create_task(self._poll_loop(), name="subscription-tracker")
|
||||
logger.info("Subscription tracker started (poll_interval=%ds)", self._poll_interval_s)
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the background polling loop and persist current state."""
|
||||
if self._stop_event:
|
||||
self._stop_event.set()
|
||||
if self._poll_task:
|
||||
try:
|
||||
await asyncio.wait_for(self._poll_task, timeout=5.0)
|
||||
except (asyncio.TimeoutError, asyncio.CancelledError):
|
||||
self._poll_task.cancel()
|
||||
self._persist_state()
|
||||
logger.info("Subscription tracker stopped")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Proxy integration hooks
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def notify_active(self, token: str) -> None:
|
||||
"""Called by the proxy handler when an OAuth request comes through.
|
||||
|
||||
Stores the token for polling and marks the tracker as recently active.
|
||||
Only processes Bearer tokens that look like OAuth (not API keys).
|
||||
"""
|
||||
if not token or not token.startswith("Bearer "):
|
||||
return
|
||||
raw = token[len("Bearer ") :]
|
||||
# Skip raw API keys (not OAuth tokens)
|
||||
if raw.startswith("sk-ant-api"):
|
||||
return
|
||||
with self._lock:
|
||||
self._current_token = raw
|
||||
self._state.last_active_at = _utc_now()
|
||||
prefix = raw[:8]
|
||||
self._full_tokens[prefix] = self._full_tokens.get(prefix, 0) + 1
|
||||
|
||||
def update_contribution(
|
||||
self,
|
||||
*,
|
||||
tokens_submitted: int = 0,
|
||||
tokens_saved_compression: int = 0,
|
||||
tokens_saved_rtk: int = 0,
|
||||
tokens_saved_cache_reads: int = 0,
|
||||
compression_savings_usd: float = 0.0,
|
||||
cache_savings_usd: float = 0.0,
|
||||
) -> None:
|
||||
"""Update headroom contribution counters for the current session window.
|
||||
|
||||
Called after each proxy request completes with the actual token deltas.
|
||||
"""
|
||||
with self._lock:
|
||||
c = self._state.contribution
|
||||
c.tokens_submitted += max(tokens_submitted, 0)
|
||||
c.tokens_saved_compression += max(tokens_saved_compression, 0)
|
||||
c.tokens_saved_rtk += max(tokens_saved_rtk, 0)
|
||||
c.tokens_saved_cache_reads += max(tokens_saved_cache_reads, 0)
|
||||
c.compression_savings_usd += max(compression_savings_usd, 0.0)
|
||||
c.cache_savings_usd += max(cache_savings_usd, 0.0)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# State access
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@property
|
||||
def state(self) -> dict[str, Any]:
|
||||
"""Return current tracker state as a serialisable dict."""
|
||||
with self._lock:
|
||||
return self._state.to_dict()
|
||||
|
||||
@property
|
||||
def latest_snapshot(self) -> SubscriptionSnapshot | None:
|
||||
with self._lock:
|
||||
return self._state.latest
|
||||
|
||||
def is_active(self) -> bool:
|
||||
with self._lock:
|
||||
return self._state.is_active(active_window_s=self._active_window_s)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Poll loop
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def _poll_loop(self) -> None:
|
||||
assert self._stop_event is not None
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
await self._maybe_poll()
|
||||
except Exception as exc:
|
||||
logger.warning("Subscription tracker poll error: %s", exc)
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
asyncio.shield(self._stop_event.wait()),
|
||||
timeout=self._poll_interval_s,
|
||||
)
|
||||
break # stop event was set
|
||||
except asyncio.TimeoutError:
|
||||
pass # normal: poll interval elapsed
|
||||
|
||||
async def _maybe_poll(self) -> None:
|
||||
with self._lock:
|
||||
is_active = self._state.is_active(active_window_s=self._active_window_s)
|
||||
token = self._current_token
|
||||
|
||||
if not is_active:
|
||||
# Try background poll using credentials file token
|
||||
from headroom.subscription.client import read_cached_oauth_token
|
||||
|
||||
bg_token = read_cached_oauth_token()
|
||||
if not bg_token:
|
||||
return
|
||||
token = token or bg_token
|
||||
|
||||
snapshot = await self._client.fetch(token)
|
||||
if snapshot is None:
|
||||
with self._lock:
|
||||
self._state.mark_error("fetch returned None")
|
||||
return
|
||||
|
||||
# Read transcript-based window tokens
|
||||
window_tokens = _compute_window_tokens_for_snapshot(snapshot)
|
||||
|
||||
# Detect anomalies
|
||||
discrepancies = _detect_discrepancies(snapshot, window_tokens)
|
||||
|
||||
with self._lock:
|
||||
self._state.add_snapshot(snapshot)
|
||||
self._state.window_tokens = window_tokens
|
||||
for d in discrepancies:
|
||||
self._state.add_discrepancy(d)
|
||||
self._state.last_error = None
|
||||
# Reset contribution when 5h window rolls over
|
||||
self._maybe_reset_contribution(snapshot)
|
||||
|
||||
self._persist_state()
|
||||
logger.debug(
|
||||
"Subscription poll: 5h=%.1f%% 7d=%.1f%%",
|
||||
snapshot.five_hour.utilization_pct,
|
||||
snapshot.seven_day.utilization_pct,
|
||||
)
|
||||
|
||||
# Update OTEL metrics if configured
|
||||
try:
|
||||
from headroom.observability.metrics import get_otel_metrics
|
||||
|
||||
get_otel_metrics().record_subscription_window(self._state.to_dict())
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _maybe_reset_contribution(self, snapshot: SubscriptionSnapshot) -> None:
|
||||
"""Reset contribution counters when the 5h window rolls over."""
|
||||
prev = self._state.history[-2] if len(self._state.history) >= 2 else None
|
||||
if prev is None:
|
||||
return
|
||||
prev_resets_at = prev.five_hour.resets_at
|
||||
curr_resets_at = snapshot.five_hour.resets_at
|
||||
if (
|
||||
prev_resets_at is not None
|
||||
and curr_resets_at is not None
|
||||
and curr_resets_at != prev_resets_at
|
||||
):
|
||||
logger.info("5h window rolled over; resetting headroom contribution counters")
|
||||
self._state.contribution = HeadroomContribution()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Persistence
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _persist_state(self) -> None:
|
||||
try:
|
||||
self._persist_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with self._lock:
|
||||
data = self._state.to_persist_dict()
|
||||
with tempfile.NamedTemporaryFile(
|
||||
mode="w",
|
||||
dir=self._persist_path.parent,
|
||||
delete=False,
|
||||
suffix=".tmp",
|
||||
encoding="utf-8",
|
||||
) as fh:
|
||||
json.dump(data, fh, indent=2)
|
||||
tmp_path = fh.name
|
||||
os.replace(tmp_path, self._persist_path)
|
||||
except Exception as exc:
|
||||
logger.debug("Failed to persist subscription state: %s", exc)
|
||||
|
||||
def _load_persisted_state(self) -> None:
|
||||
try:
|
||||
with open(self._persist_path, encoding="utf-8") as fh:
|
||||
raw = json.load(fh)
|
||||
# Restore only the contribution counters and poll counts for now;
|
||||
# snapshot data is re-fetched on first active poll.
|
||||
contrib = raw.get("contribution", {})
|
||||
c = self._state.contribution
|
||||
c.tokens_submitted = int(contrib.get("tokens_submitted", 0))
|
||||
saved = contrib.get("tokens_saved", {})
|
||||
c.tokens_saved_compression = int(saved.get("compression", 0))
|
||||
c.tokens_saved_rtk = int(saved.get("rtk", 0))
|
||||
c.tokens_saved_cache_reads = int(saved.get("cache_reads", 0))
|
||||
savings_usd = contrib.get("savings_usd", {})
|
||||
c.compression_savings_usd = float(savings_usd.get("compression", 0.0))
|
||||
c.cache_savings_usd = float(savings_usd.get("cache", 0.0))
|
||||
self._state.poll_count = int(raw.get("poll_count", 0))
|
||||
logger.debug("Loaded persisted subscription state from %s", self._persist_path)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
except Exception as exc:
|
||||
logger.debug("Could not load persisted subscription state: %s", exc)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Transcript-based window token computation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _compute_window_tokens_for_snapshot(snapshot: SubscriptionSnapshot) -> WindowTokens:
|
||||
"""Read Claude transcript files and sum tokens for the current 5h window."""
|
||||
try:
|
||||
from headroom.subscription import session_tracking
|
||||
|
||||
resets_at = snapshot.five_hour.resets_at
|
||||
if resets_at is None:
|
||||
return WindowTokens()
|
||||
window_duration_s = 5 * 3600 # 5-hour window
|
||||
start_ts = resets_at.timestamp() - window_duration_s
|
||||
end_ts = resets_at.timestamp()
|
||||
return session_tracking.compute_window_tokens(start_ts, end_ts)
|
||||
except Exception as exc:
|
||||
logger.debug("Could not compute window tokens from transcripts: %s", exc)
|
||||
return WindowTokens()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Anomaly detection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _detect_discrepancies(
|
||||
snapshot: SubscriptionSnapshot,
|
||||
window_tokens: WindowTokens,
|
||||
) -> list[WindowDiscrepancy]:
|
||||
"""Detect surge pricing or cache miss anomalies in the snapshot."""
|
||||
discrepancies: list[WindowDiscrepancy] = []
|
||||
|
||||
if snapshot.five_hour.limit > 0 and window_tokens.weighted_token_equivalent > 0:
|
||||
expected_pct = window_tokens.weighted_token_equivalent / snapshot.five_hour.limit * 100.0
|
||||
actual_pct = snapshot.five_hour.utilization_pct
|
||||
delta = actual_pct - expected_pct
|
||||
|
||||
if delta > _SURGE_THRESHOLD_PCT:
|
||||
discrepancies.append(
|
||||
WindowDiscrepancy(
|
||||
kind="surge_pricing",
|
||||
description=(
|
||||
f"API 5h utilization ({actual_pct:.1f}%) is "
|
||||
f"{delta:.1f}% higher than transcript-implied "
|
||||
f"({expected_pct:.1f}%); possible surge weighting."
|
||||
),
|
||||
severity="warning" if delta < 30 else "alert",
|
||||
expected_utilization_pct=round(expected_pct, 2),
|
||||
actual_utilization_pct=round(actual_pct, 2),
|
||||
delta_pct=round(delta, 2),
|
||||
)
|
||||
)
|
||||
|
||||
total_input = window_tokens.input
|
||||
total_cache_reads = window_tokens.cache_reads
|
||||
if total_input > 50_000 and total_cache_reads < total_input * _CACHE_MISS_RATIO_THRESHOLD:
|
||||
cache_ratio = total_cache_reads / total_input if total_input else 0
|
||||
discrepancies.append(
|
||||
WindowDiscrepancy(
|
||||
kind="cache_miss",
|
||||
description=(
|
||||
f"Cache-read ratio is {cache_ratio:.1%} (threshold "
|
||||
f"{_CACHE_MISS_RATIO_THRESHOLD:.0%}); system may not be "
|
||||
"using prefix cache effectively."
|
||||
),
|
||||
severity="warning",
|
||||
expected_utilization_pct=None,
|
||||
actual_utilization_pct=None,
|
||||
delta_pct=None,
|
||||
)
|
||||
)
|
||||
|
||||
return discrepancies
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Module-level singleton
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_tracker_lock = threading.Lock()
|
||||
_tracker_instance: SubscriptionTracker | None = None
|
||||
|
||||
|
||||
def get_subscription_tracker() -> SubscriptionTracker | None:
|
||||
"""Return the global singleton tracker, or None if not configured."""
|
||||
return _tracker_instance
|
||||
|
||||
|
||||
def configure_subscription_tracker(
|
||||
poll_interval_s: int = _DEFAULT_POLL_INTERVAL_S,
|
||||
active_window_s: float = _DEFAULT_ACTIVE_WINDOW_S,
|
||||
persist_path: Path | None = None,
|
||||
client: SubscriptionClient | None = None,
|
||||
) -> SubscriptionTracker:
|
||||
"""Create (or return existing) global tracker singleton."""
|
||||
global _tracker_instance
|
||||
with _tracker_lock:
|
||||
if _tracker_instance is None:
|
||||
_tracker_instance = SubscriptionTracker(
|
||||
poll_interval_s=poll_interval_s,
|
||||
active_window_s=active_window_s,
|
||||
persist_path=persist_path,
|
||||
client=client,
|
||||
)
|
||||
return _tracker_instance
|
||||
|
||||
|
||||
async def shutdown_subscription_tracker() -> None:
|
||||
"""Stop and clean up the global tracker."""
|
||||
global _tracker_instance
|
||||
with _tracker_lock:
|
||||
tracker = _tracker_instance
|
||||
_tracker_instance = None
|
||||
if tracker:
|
||||
await tracker.stop()
|
||||
434
tests/test_subscription_tracker.py
Normal file
434
tests/test_subscription_tracker.py
Normal file
|
|
@ -0,0 +1,434 @@
|
|||
"""Tests for the Anthropic subscription window tracking feature."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.subscription.models import (
|
||||
ExtraUsage,
|
||||
HeadroomContribution,
|
||||
RateLimitWindow,
|
||||
SubscriptionSnapshot,
|
||||
SubscriptionState,
|
||||
WindowDiscrepancy,
|
||||
WindowTokens,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RateLimitWindow
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestRateLimitWindow:
|
||||
def test_from_api_dict_full(self):
|
||||
now = datetime.now(timezone.utc)
|
||||
resets = (now + timedelta(hours=3)).isoformat()
|
||||
w = RateLimitWindow.from_api_dict(
|
||||
{"utilization": 42.5, "resets_at": resets, "used": 100, "limit": 235}
|
||||
)
|
||||
assert w.utilization_pct == pytest.approx(42.5)
|
||||
assert w.used == 100
|
||||
assert w.limit == 235
|
||||
assert w.resets_at is not None
|
||||
|
||||
def test_from_api_dict_minimal(self):
|
||||
# Real API responses only include utilization + resets_at
|
||||
w = RateLimitWindow.from_api_dict({"utilization": 0.0, "resets_at": None})
|
||||
assert w.utilization_pct == 0.0
|
||||
assert w.resets_at is None
|
||||
assert w.used == 0
|
||||
assert w.limit == 0
|
||||
|
||||
def test_seconds_to_reset(self):
|
||||
now = datetime.now(timezone.utc)
|
||||
future = now + timedelta(hours=2, minutes=30)
|
||||
w = RateLimitWindow(resets_at=future)
|
||||
secs = w.seconds_to_reset(now=now)
|
||||
assert secs is not None
|
||||
assert 9000 <= secs <= 9001
|
||||
|
||||
def test_seconds_to_reset_none_when_no_reset(self):
|
||||
w = RateLimitWindow()
|
||||
assert w.seconds_to_reset() is None
|
||||
|
||||
def test_to_dict_contains_expected_keys(self):
|
||||
w = RateLimitWindow(utilization_pct=55.0)
|
||||
d = w.to_dict()
|
||||
assert "utilization_pct" in d
|
||||
assert "resets_at" in d
|
||||
assert "seconds_to_reset" in d
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ExtraUsage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestExtraUsage:
|
||||
def test_from_api_dict_with_cents(self):
|
||||
eu = ExtraUsage.from_api_dict(
|
||||
{
|
||||
"is_enabled": True,
|
||||
"monthly_limit": 5000, # $50.00
|
||||
"used_credits": 123, # $1.23
|
||||
"utilization": 2.46,
|
||||
}
|
||||
)
|
||||
assert eu.is_enabled is True
|
||||
assert eu.monthly_limit_cents == 5000
|
||||
assert eu.used_credits_cents == 123
|
||||
assert eu.monthly_limit_usd == pytest.approx(50.0)
|
||||
assert eu.used_credits_usd == pytest.approx(1.23)
|
||||
assert eu.utilization_pct == pytest.approx(2.46)
|
||||
|
||||
def test_from_api_dict_disabled(self):
|
||||
eu = ExtraUsage.from_api_dict({"is_enabled": False})
|
||||
assert eu.is_enabled is False
|
||||
assert eu.monthly_limit_cents is None
|
||||
assert eu.used_credits_cents is None
|
||||
assert eu.monthly_limit_usd is None
|
||||
|
||||
def test_to_dict_converts_to_usd(self):
|
||||
eu = ExtraUsage(is_enabled=True, monthly_limit_cents=10000, used_credits_cents=499)
|
||||
d = eu.to_dict()
|
||||
assert d["monthly_limit_usd"] == pytest.approx(100.0, abs=0.01)
|
||||
assert d["used_credits_usd"] == pytest.approx(4.99, abs=0.001)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SubscriptionSnapshot
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestSubscriptionSnapshot:
|
||||
def test_from_api_response_all_fields(self):
|
||||
now = datetime.now(timezone.utc)
|
||||
data = {
|
||||
"five_hour": {"utilization": 30.0, "resets_at": (now + timedelta(hours=2)).isoformat()},
|
||||
"seven_day": {"utilization": 10.0, "resets_at": (now + timedelta(days=3)).isoformat()},
|
||||
"seven_day_opus": {"utilization": 5.0, "resets_at": None},
|
||||
"seven_day_sonnet": {"utilization": 8.0, "resets_at": None},
|
||||
"extra_usage": {
|
||||
"is_enabled": True,
|
||||
"monthly_limit": 5000,
|
||||
"used_credits": 250,
|
||||
"utilization": 5.0,
|
||||
},
|
||||
}
|
||||
snap = SubscriptionSnapshot.from_api_response(data, token="tok_test_abc")
|
||||
assert snap.five_hour.utilization_pct == pytest.approx(30.0)
|
||||
assert snap.seven_day.utilization_pct == pytest.approx(10.0)
|
||||
assert snap.seven_day_opus is not None
|
||||
assert snap.seven_day_sonnet is not None
|
||||
assert snap.extra_usage.is_enabled is True
|
||||
assert snap.extra_usage.used_credits_usd == pytest.approx(2.50)
|
||||
assert snap.token_prefix == "tok_test"
|
||||
|
||||
def test_from_api_response_missing_optional_fields(self):
|
||||
snap = SubscriptionSnapshot.from_api_response(
|
||||
{"five_hour": {"utilization": 0.0, "resets_at": None}}
|
||||
)
|
||||
assert snap.seven_day_opus is None
|
||||
assert snap.seven_day_sonnet is None
|
||||
assert snap.extra_usage.is_enabled is False
|
||||
|
||||
def test_to_dict_round_trip(self):
|
||||
snap = SubscriptionSnapshot.from_api_response(
|
||||
{
|
||||
"five_hour": {"utilization": 42.0, "resets_at": None},
|
||||
"seven_day": {"utilization": 15.0, "resets_at": None},
|
||||
"extra_usage": {"is_enabled": False},
|
||||
}
|
||||
)
|
||||
d = snap.to_dict()
|
||||
assert d["five_hour"]["utilization_pct"] == pytest.approx(42.0)
|
||||
assert d["seven_day"]["utilization_pct"] == pytest.approx(15.0)
|
||||
assert "seven_day_opus" not in d # absent when None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HeadroomContribution
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestHeadroomContribution:
|
||||
def test_efficiency_pct_no_savings(self):
|
||||
c = HeadroomContribution(tokens_submitted=100)
|
||||
assert c.efficiency_pct() == 0.0
|
||||
|
||||
def test_efficiency_pct_with_savings(self):
|
||||
c = HeadroomContribution(
|
||||
tokens_submitted=70,
|
||||
tokens_saved_compression=20,
|
||||
tokens_saved_rtk=10,
|
||||
)
|
||||
# raw_without_headroom = 70 + 20 + 10 = 100
|
||||
# total_saved = 30
|
||||
assert c.efficiency_pct() == pytest.approx(30.0)
|
||||
|
||||
def test_total_savings_usd(self):
|
||||
c = HeadroomContribution(compression_savings_usd=1.5, cache_savings_usd=0.75)
|
||||
assert c.total_savings_usd() == pytest.approx(2.25)
|
||||
|
||||
def test_to_dict_structure(self):
|
||||
c = HeadroomContribution(
|
||||
tokens_submitted=200,
|
||||
tokens_saved_compression=50,
|
||||
compression_savings_usd=0.10,
|
||||
)
|
||||
d = c.to_dict()
|
||||
assert d["tokens_submitted"] == 200
|
||||
assert d["tokens_saved"]["compression"] == 50
|
||||
assert d["savings_usd"]["compression"] == pytest.approx(0.10)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SubscriptionState
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestSubscriptionState:
|
||||
def test_is_active_recent(self):
|
||||
state = SubscriptionState()
|
||||
state.last_active_at = datetime.now(timezone.utc)
|
||||
assert state.is_active(active_window_s=60.0) is True
|
||||
|
||||
def test_is_active_stale(self):
|
||||
state = SubscriptionState()
|
||||
state.last_active_at = datetime.now(timezone.utc) - timedelta(minutes=5)
|
||||
assert state.is_active(active_window_s=60.0) is False
|
||||
|
||||
def test_is_active_none(self):
|
||||
state = SubscriptionState()
|
||||
assert state.is_active() is False
|
||||
|
||||
def test_add_snapshot_caps_history(self):
|
||||
state = SubscriptionState()
|
||||
state._MAX_HISTORY = 3
|
||||
for _ in range(5):
|
||||
state.add_snapshot(SubscriptionSnapshot.from_api_response({}))
|
||||
assert len(state.history) == 3
|
||||
assert state.poll_count == 5
|
||||
|
||||
def test_add_discrepancy_caps_list(self):
|
||||
state = SubscriptionState()
|
||||
state._MAX_DISCREPANCIES = 2
|
||||
for i in range(4):
|
||||
state.add_discrepancy(WindowDiscrepancy(kind=f"kind_{i}"))
|
||||
assert len(state.discrepancies) == 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SubscriptionTracker — unit tests with mocked client
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestSubscriptionTrackerNotifyActive:
|
||||
def _make_tracker(self, tmp_path: Path):
|
||||
from headroom.subscription.tracker import SubscriptionTracker
|
||||
return SubscriptionTracker(
|
||||
poll_interval_s=30,
|
||||
active_window_s=60,
|
||||
persist_path=tmp_path / "state.json",
|
||||
client=MagicMock(),
|
||||
)
|
||||
|
||||
def test_notify_active_oauth_token(self, tmp_path):
|
||||
t = self._make_tracker(tmp_path)
|
||||
t.notify_active("Bearer sk-ant-oat01-sometoken")
|
||||
assert t.is_active() is True
|
||||
assert t._current_token == "sk-ant-oat01-sometoken"
|
||||
|
||||
def test_notify_active_ignores_api_key(self, tmp_path):
|
||||
t = self._make_tracker(tmp_path)
|
||||
t.notify_active("Bearer sk-ant-api03-key")
|
||||
assert t.is_active() is False
|
||||
assert t._current_token is None
|
||||
|
||||
def test_notify_active_ignores_non_bearer(self, tmp_path):
|
||||
t = self._make_tracker(tmp_path)
|
||||
t.notify_active("x-api-key somevalue")
|
||||
assert t.is_active() is False
|
||||
|
||||
def test_update_contribution(self, tmp_path):
|
||||
t = self._make_tracker(tmp_path)
|
||||
t.update_contribution(
|
||||
tokens_submitted=500,
|
||||
tokens_saved_compression=100,
|
||||
tokens_saved_cache_reads=50,
|
||||
)
|
||||
c = t._state.contribution
|
||||
assert c.tokens_submitted == 500
|
||||
assert c.tokens_saved_compression == 100
|
||||
assert c.tokens_saved_cache_reads == 50
|
||||
|
||||
def test_update_contribution_accumulates(self, tmp_path):
|
||||
t = self._make_tracker(tmp_path)
|
||||
t.update_contribution(tokens_submitted=100)
|
||||
t.update_contribution(tokens_submitted=200)
|
||||
assert t._state.contribution.tokens_submitted == 300
|
||||
|
||||
def test_state_dict_structure(self, tmp_path):
|
||||
t = self._make_tracker(tmp_path)
|
||||
d = t.state
|
||||
assert "latest" in d
|
||||
assert "contribution" in d
|
||||
assert "poll_count" in d
|
||||
assert "last_active_at" in d
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SubscriptionTracker — poll loop integration test
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestSubscriptionTrackerPollLoop:
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_called_when_active(self, tmp_path):
|
||||
"""When notify_active is called, the poll loop should fetch a snapshot."""
|
||||
from headroom.subscription.tracker import SubscriptionTracker
|
||||
|
||||
mock_snapshot = SubscriptionSnapshot.from_api_response(
|
||||
{
|
||||
"five_hour": {"utilization": 25.0, "resets_at": None},
|
||||
"seven_day": {"utilization": 5.0, "resets_at": None},
|
||||
}
|
||||
)
|
||||
mock_client = MagicMock()
|
||||
mock_client.fetch = AsyncMock(return_value=mock_snapshot)
|
||||
|
||||
tracker = SubscriptionTracker(
|
||||
poll_interval_s=1,
|
||||
active_window_s=60,
|
||||
persist_path=tmp_path / "state.json",
|
||||
client=mock_client,
|
||||
)
|
||||
|
||||
# Mark active so the poll proceeds
|
||||
tracker.notify_active("Bearer sk-ant-oat01-testtoken")
|
||||
|
||||
# Run poll once directly (bypasses loop timing)
|
||||
await tracker._maybe_poll()
|
||||
|
||||
assert tracker.latest_snapshot is not None
|
||||
assert tracker.latest_snapshot.five_hour.utilization_pct == pytest.approx(25.0)
|
||||
mock_client.fetch.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_poll_skipped_when_no_token_and_inactive(self, tmp_path):
|
||||
"""When not active and no credentials file token, poll should skip."""
|
||||
from headroom.subscription.tracker import SubscriptionTracker
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.fetch = AsyncMock(return_value=None)
|
||||
|
||||
tracker = SubscriptionTracker(
|
||||
poll_interval_s=1,
|
||||
active_window_s=60,
|
||||
persist_path=tmp_path / "state.json",
|
||||
client=mock_client,
|
||||
)
|
||||
|
||||
with patch("headroom.subscription.client.read_cached_oauth_token", return_value=None):
|
||||
await tracker._maybe_poll()
|
||||
|
||||
mock_client.fetch.assert_not_called()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Anomaly detection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestAnomalyDetection:
|
||||
def test_surge_pricing_detection(self):
|
||||
from headroom.subscription.tracker import _detect_discrepancies
|
||||
|
||||
snap = SubscriptionSnapshot.from_api_response(
|
||||
{
|
||||
"five_hour": {"utilization": 80.0, "resets_at": None},
|
||||
"seven_day": {"utilization": 20.0, "resets_at": None},
|
||||
}
|
||||
)
|
||||
# Simulate API limit known; weighted tokens imply only 50% should be used
|
||||
snap.five_hour.limit = 1000
|
||||
window_tokens = WindowTokens(
|
||||
input=300, output=100, weighted_token_equivalent=500.0
|
||||
)
|
||||
|
||||
discrepancies = _detect_discrepancies(snap, window_tokens)
|
||||
kinds = [d.kind for d in discrepancies]
|
||||
assert "surge_pricing" in kinds
|
||||
|
||||
def test_no_surge_when_limit_unknown(self):
|
||||
from headroom.subscription.tracker import _detect_discrepancies
|
||||
|
||||
snap = SubscriptionSnapshot.from_api_response(
|
||||
{"five_hour": {"utilization": 90.0, "resets_at": None}}
|
||||
)
|
||||
snap.five_hour.limit = 0 # Unknown
|
||||
window_tokens = WindowTokens(weighted_token_equivalent=500.0)
|
||||
|
||||
discrepancies = _detect_discrepancies(snap, window_tokens)
|
||||
assert not any(d.kind == "surge_pricing" for d in discrepancies)
|
||||
|
||||
def test_cache_miss_detection(self):
|
||||
from headroom.subscription.tracker import _detect_discrepancies
|
||||
|
||||
snap = SubscriptionSnapshot.from_api_response(
|
||||
{"five_hour": {"utilization": 40.0, "resets_at": None}}
|
||||
)
|
||||
# High input, very few cache reads
|
||||
window_tokens = WindowTokens(input=100_000, cache_reads=500)
|
||||
|
||||
discrepancies = _detect_discrepancies(snap, window_tokens)
|
||||
kinds = [d.kind for d in discrepancies]
|
||||
assert "cache_miss" in kinds
|
||||
|
||||
def test_no_cache_miss_below_threshold(self):
|
||||
from headroom.subscription.tracker import _detect_discrepancies
|
||||
|
||||
snap = SubscriptionSnapshot.from_api_response(
|
||||
{"five_hour": {"utilization": 30.0, "resets_at": None}}
|
||||
)
|
||||
# Low input total — doesn't trigger cache miss check
|
||||
window_tokens = WindowTokens(input=30_000, cache_reads=0)
|
||||
|
||||
discrepancies = _detect_discrepancies(snap, window_tokens)
|
||||
assert not any(d.kind == "cache_miss" for d in discrepancies)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Persistence round-trip
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestPersistence:
|
||||
@pytest.mark.asyncio
|
||||
async def test_persist_and_reload(self, tmp_path):
|
||||
from headroom.subscription.tracker import SubscriptionTracker
|
||||
|
||||
persist_file = tmp_path / "sub_state.json"
|
||||
|
||||
mock_client = MagicMock()
|
||||
snap = SubscriptionSnapshot.from_api_response(
|
||||
{"five_hour": {"utilization": 55.0, "resets_at": None}}
|
||||
)
|
||||
mock_client.fetch = AsyncMock(return_value=snap)
|
||||
|
||||
t1 = SubscriptionTracker(
|
||||
persist_path=persist_file,
|
||||
client=mock_client,
|
||||
)
|
||||
t1.notify_active("Bearer sk-ant-oat01-tok")
|
||||
t1.update_contribution(tokens_submitted=1000, tokens_saved_compression=200)
|
||||
await t1._maybe_poll()
|
||||
t1._persist_state()
|
||||
|
||||
assert persist_file.exists()
|
||||
|
||||
# Load a new tracker from the same file
|
||||
t2 = SubscriptionTracker(
|
||||
persist_path=persist_file,
|
||||
client=MagicMock(),
|
||||
)
|
||||
assert t2._state.contribution.tokens_submitted == 1000
|
||||
assert t2._state.contribution.tokens_saved_compression == 200
|
||||
Loading…
Add table
Add a link
Reference in a new issue