Mirrowel commited on
Commit
8b4ff52
·
1 Parent(s): 846c165

feat(quota-viewer): ✨ add quota and usage statistics viewer system

Browse files

This introduces a comprehensive quota and usage statistics monitoring system with API endpoints and a rich terminal UI for real-time credential consumption visibility.

Key components:
- Add GET/POST /v1/quota-stats API endpoints to main proxy for retrieving and refreshing statistics
- Implement full-featured TUI viewer (quota_viewer.py) with provider summaries, credential details, progress bars, and remote management
- Add configuration system (quota_viewer_config.py) supporting multiple remote proxies with API key management
- Extend RotatingClient with get_quota_stats(), reload_usage_from_disk(), and force_refresh_quota() methods
- Enhance usage_manager with get_stats_for_endpoint() for aggregated statistics
- Integrate viewer into launcher TUI as new menu option 5

Features include toggle between current and global/lifetime stats, token formatting, cooldown timers, quota group visualization, and estimated cost tracking. Supports Antigravity quota group enrichment and live API refresh.

Also in this commit:
- fix(antigravity): increase empty response retry attempts from 4 to 6

src/proxy_app/launcher_tui.py CHANGED
@@ -429,9 +429,10 @@ class LauncherTUI:
429
  self.console.print(" 3. 🔑 Manage Credentials")
430
 
431
  self.console.print(" 4. 📊 View Provider & Advanced Settings")
432
- self.console.print(" 5. 🔄 Reload Configuration")
433
- self.console.print(" 6. ℹ️ About")
434
- self.console.print(" 7. 🚪 Exit")
 
435
 
436
  self.console.print()
437
  self.console.print("━" * 70)
@@ -439,7 +440,7 @@ class LauncherTUI:
439
 
440
  choice = Prompt.ask(
441
  "Select option",
442
- choices=["1", "2", "3", "4", "5", "6", "7"],
443
  show_choices=False,
444
  )
445
 
@@ -452,12 +453,14 @@ class LauncherTUI:
452
  elif choice == "4":
453
  self.show_provider_settings_menu()
454
  elif choice == "5":
 
 
455
  load_dotenv(dotenv_path=_get_env_file(), override=True)
456
  self.config = LauncherConfig() # Reload config
457
  self.console.print("\n[green]✅ Configuration reloaded![/green]")
458
- elif choice == "6":
459
- self.show_about()
460
  elif choice == "7":
 
 
461
  self.running = False
462
  sys.exit(0)
463
 
@@ -874,6 +877,20 @@ class LauncherTUI:
874
  # Reload environment after settings tool
875
  load_dotenv(dotenv_path=_get_env_file(), override=True)
876
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
877
  def show_about(self):
878
  """Display About page with project information"""
879
  clear_screen()
 
429
  self.console.print(" 3. 🔑 Manage Credentials")
430
 
431
  self.console.print(" 4. 📊 View Provider & Advanced Settings")
432
+ self.console.print(" 5. 📈 View Quota & Usage Stats")
433
+ self.console.print(" 6. 🔄 Reload Configuration")
434
+ self.console.print(" 7. ℹ️ About")
435
+ self.console.print(" 8. 🚪 Exit")
436
 
437
  self.console.print()
438
  self.console.print("━" * 70)
 
440
 
441
  choice = Prompt.ask(
442
  "Select option",
443
+ choices=["1", "2", "3", "4", "5", "6", "7", "8"],
444
  show_choices=False,
445
  )
446
 
 
453
  elif choice == "4":
454
  self.show_provider_settings_menu()
455
  elif choice == "5":
456
+ self.launch_quota_viewer()
457
+ elif choice == "6":
458
  load_dotenv(dotenv_path=_get_env_file(), override=True)
459
  self.config = LauncherConfig() # Reload config
460
  self.console.print("\n[green]✅ Configuration reloaded![/green]")
 
 
461
  elif choice == "7":
462
+ self.show_about()
463
+ elif choice == "8":
464
  self.running = False
465
  sys.exit(0)
466
 
 
877
  # Reload environment after settings tool
878
  load_dotenv(dotenv_path=_get_env_file(), override=True)
879
 
880
+ def launch_quota_viewer(self):
881
+ """Launch the quota stats viewer"""
882
+ clear_screen()
883
+
884
+ self.console.print("━" * 70)
885
+ self.console.print("Quota & Usage Statistics Viewer")
886
+ self.console.print("━" * 70)
887
+ self.console.print()
888
+
889
+ # Import the lightweight viewer (no heavy imports)
890
+ from proxy_app.quota_viewer import run_quota_viewer
891
+
892
+ run_quota_viewer()
893
+
894
  def show_about(self):
895
  """Display About page with project information"""
896
  clear_screen()
src/proxy_app/main.py CHANGED
@@ -1148,6 +1148,145 @@ async def list_providers(_=Depends(verify_api_key)):
1148
  return list(PROVIDER_PLUGINS.keys())
1149
 
1150
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1151
  @app.post("/v1/token-count")
1152
  async def token_count(
1153
  request: Request,
 
1148
  return list(PROVIDER_PLUGINS.keys())
1149
 
1150
 
1151
+ @app.get("/v1/quota-stats")
1152
+ async def get_quota_stats(
1153
+ request: Request,
1154
+ client: RotatingClient = Depends(get_rotating_client),
1155
+ _=Depends(verify_api_key),
1156
+ provider: str = None,
1157
+ ):
1158
+ """
1159
+ Returns quota and usage statistics for all credentials.
1160
+
1161
+ This returns cached data from the proxy without making external API calls.
1162
+ Use POST to reload from disk or force refresh from external APIs.
1163
+
1164
+ Query Parameters:
1165
+ provider: Optional filter to return stats for a specific provider only
1166
+
1167
+ Returns:
1168
+ {
1169
+ "providers": {
1170
+ "provider_name": {
1171
+ "credential_count": int,
1172
+ "active_count": int,
1173
+ "on_cooldown_count": int,
1174
+ "exhausted_count": int,
1175
+ "total_requests": int,
1176
+ "tokens": {...},
1177
+ "approx_cost": float | null,
1178
+ "quota_groups": {...}, // For Antigravity
1179
+ "credentials": [...]
1180
+ }
1181
+ },
1182
+ "summary": {...},
1183
+ "data_source": "cache",
1184
+ "timestamp": float
1185
+ }
1186
+ """
1187
+ try:
1188
+ stats = await client.get_quota_stats(provider_filter=provider)
1189
+ return stats
1190
+ except Exception as e:
1191
+ logging.error(f"Failed to get quota stats: {e}")
1192
+ raise HTTPException(status_code=500, detail=str(e))
1193
+
1194
+
1195
+ @app.post("/v1/quota-stats")
1196
+ async def refresh_quota_stats(
1197
+ request: Request,
1198
+ client: RotatingClient = Depends(get_rotating_client),
1199
+ _=Depends(verify_api_key),
1200
+ ):
1201
+ """
1202
+ Refresh quota and usage statistics.
1203
+
1204
+ Request body:
1205
+ {
1206
+ "action": "reload" | "force_refresh",
1207
+ "scope": "all" | "provider" | "credential",
1208
+ "provider": "antigravity", // required if scope != "all"
1209
+ "credential": "antigravity_oauth_1.json" // required if scope == "credential"
1210
+ }
1211
+
1212
+ Actions:
1213
+ - reload: Re-read data from disk (no external API calls)
1214
+ - force_refresh: For Antigravity, fetch live quota from API.
1215
+ For other providers, same as reload.
1216
+
1217
+ Returns:
1218
+ Same as GET, plus a "refresh_result" field with operation details.
1219
+ """
1220
+ try:
1221
+ data = await request.json()
1222
+ action = data.get("action", "reload")
1223
+ scope = data.get("scope", "all")
1224
+ provider = data.get("provider")
1225
+ credential = data.get("credential")
1226
+
1227
+ # Validate parameters
1228
+ if action not in ("reload", "force_refresh"):
1229
+ raise HTTPException(
1230
+ status_code=400,
1231
+ detail="action must be 'reload' or 'force_refresh'",
1232
+ )
1233
+
1234
+ if scope not in ("all", "provider", "credential"):
1235
+ raise HTTPException(
1236
+ status_code=400,
1237
+ detail="scope must be 'all', 'provider', or 'credential'",
1238
+ )
1239
+
1240
+ if scope in ("provider", "credential") and not provider:
1241
+ raise HTTPException(
1242
+ status_code=400,
1243
+ detail="'provider' is required when scope is 'provider' or 'credential'",
1244
+ )
1245
+
1246
+ if scope == "credential" and not credential:
1247
+ raise HTTPException(
1248
+ status_code=400,
1249
+ detail="'credential' is required when scope is 'credential'",
1250
+ )
1251
+
1252
+ refresh_result = {
1253
+ "action": action,
1254
+ "scope": scope,
1255
+ "provider": provider,
1256
+ "credential": credential,
1257
+ }
1258
+
1259
+ if action == "reload":
1260
+ # Just reload from disk
1261
+ start_time = time.time()
1262
+ await client.reload_usage_from_disk()
1263
+ refresh_result["duration_ms"] = int((time.time() - start_time) * 1000)
1264
+ refresh_result["success"] = True
1265
+ refresh_result["message"] = "Reloaded usage data from disk"
1266
+
1267
+ elif action == "force_refresh":
1268
+ # Force refresh from external API (for supported providers like Antigravity)
1269
+ result = await client.force_refresh_quota(
1270
+ provider=provider if scope in ("provider", "credential") else None,
1271
+ credential=credential if scope == "credential" else None,
1272
+ )
1273
+ refresh_result.update(result)
1274
+ refresh_result["success"] = result["failed_count"] == 0
1275
+
1276
+ # Get updated stats
1277
+ stats = await client.get_quota_stats(provider_filter=provider)
1278
+ stats["refresh_result"] = refresh_result
1279
+ stats["data_source"] = "refreshed"
1280
+
1281
+ return stats
1282
+
1283
+ except HTTPException:
1284
+ raise
1285
+ except Exception as e:
1286
+ logging.error(f"Failed to refresh quota stats: {e}")
1287
+ raise HTTPException(status_code=500, detail=str(e))
1288
+
1289
+
1290
  @app.post("/v1/token-count")
1291
  async def token_count(
1292
  request: Request,
src/proxy_app/quota_viewer.py ADDED
@@ -0,0 +1,1086 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Lightweight Quota Stats Viewer TUI.
3
+
4
+ Connects to a running proxy to display quota and usage statistics.
5
+ Uses only httpx + rich (no heavy rotator_library imports).
6
+ """
7
+
8
+ import os
9
+ import sys
10
+ import time
11
+ from datetime import datetime, timezone
12
+ from typing import Any, Dict, List, Optional, Tuple
13
+
14
+ import httpx
15
+ from rich.console import Console
16
+ from rich.panel import Panel
17
+ from rich.progress import BarColumn, Progress, TextColumn
18
+ from rich.prompt import Prompt
19
+ from rich.table import Table
20
+ from rich.text import Text
21
+
22
+ from .quota_viewer_config import QuotaViewerConfig
23
+
24
+
25
+ def clear_screen():
26
+ """Clear the terminal screen."""
27
+ os.system("cls" if os.name == "nt" else "clear")
28
+
29
+
30
+ def format_tokens(count: int) -> str:
31
+ """Format token count for display (e.g., 125000 -> 125k)."""
32
+ if count >= 1_000_000:
33
+ return f"{count / 1_000_000:.1f}M"
34
+ elif count >= 1_000:
35
+ return f"{count / 1_000:.0f}k"
36
+ return str(count)
37
+
38
+
39
+ def format_cost(cost: Optional[float]) -> str:
40
+ """Format cost for display."""
41
+ if cost is None or cost == 0:
42
+ return "-"
43
+ if cost < 0.01:
44
+ return f"${cost:.4f}"
45
+ return f"${cost:.2f}"
46
+
47
+
48
+ def format_time_ago(timestamp: Optional[float]) -> str:
49
+ """Format timestamp as relative time (e.g., '5 min ago')."""
50
+ if not timestamp:
51
+ return "Never"
52
+ try:
53
+ delta = time.time() - timestamp
54
+ if delta < 60:
55
+ return f"{int(delta)}s ago"
56
+ elif delta < 3600:
57
+ return f"{int(delta / 60)} min ago"
58
+ elif delta < 86400:
59
+ return f"{int(delta / 3600)}h ago"
60
+ else:
61
+ return f"{int(delta / 86400)}d ago"
62
+ except (ValueError, OSError):
63
+ return "Unknown"
64
+
65
+
66
+ def format_reset_time(iso_time: Optional[str]) -> str:
67
+ """Format ISO time string for display."""
68
+ if not iso_time:
69
+ return "-"
70
+ try:
71
+ dt = datetime.fromisoformat(iso_time.replace("Z", "+00:00"))
72
+ # Convert to local time
73
+ local_dt = dt.astimezone()
74
+ return local_dt.strftime("%b %d %H:%M")
75
+ except (ValueError, AttributeError):
76
+ return iso_time[:16] if iso_time else "-"
77
+
78
+
79
+ def create_progress_bar(percent: Optional[int], width: int = 10) -> str:
80
+ """Create a text-based progress bar."""
81
+ if percent is None:
82
+ return "░" * width
83
+ filled = int(percent / 100 * width)
84
+ return "▓" * filled + "░" * (width - filled)
85
+
86
+
87
+ def format_cooldown(seconds: int) -> str:
88
+ """Format cooldown seconds as human-readable string."""
89
+ if seconds < 60:
90
+ return f"{seconds}s"
91
+ elif seconds < 3600:
92
+ mins = seconds // 60
93
+ secs = seconds % 60
94
+ return f"{mins}m {secs}s" if secs > 0 else f"{mins}m"
95
+ else:
96
+ hours = seconds // 3600
97
+ mins = (seconds % 3600) // 60
98
+ return f"{hours}h {mins}m" if mins > 0 else f"{hours}h"
99
+
100
+
101
+ class QuotaViewer:
102
+ """Main Quota Viewer TUI class."""
103
+
104
+ def __init__(self, config: Optional[QuotaViewerConfig] = None):
105
+ """
106
+ Initialize the viewer.
107
+
108
+ Args:
109
+ config: Optional config object. If not provided, one will be created.
110
+ """
111
+ self.console = Console()
112
+ self.config = config or QuotaViewerConfig()
113
+ self.config.sync_with_launcher_config()
114
+
115
+ self.current_remote: Optional[Dict[str, Any]] = None
116
+ self.cached_stats: Optional[Dict[str, Any]] = None
117
+ self.last_error: Optional[str] = None
118
+ self.running = True
119
+ self.view_mode = "current" # "current" or "global"
120
+
121
+ def _get_headers(self) -> Dict[str, str]:
122
+ """Get HTTP headers including auth if configured."""
123
+ headers = {}
124
+ if self.current_remote and self.current_remote.get("api_key"):
125
+ headers["Authorization"] = f"Bearer {self.current_remote['api_key']}"
126
+ return headers
127
+
128
+ def _get_base_url(self) -> str:
129
+ """Get base URL for the current remote."""
130
+ if not self.current_remote:
131
+ return "http://127.0.0.1:8000"
132
+ host = self.current_remote.get("host", "127.0.0.1")
133
+ port = self.current_remote.get("port", 8000)
134
+ # Use https if port is 443 or host looks like a domain
135
+ scheme = "https" if port == 443 or "." in host else "http"
136
+ return f"{scheme}://{host}:{port}"
137
+
138
+ def check_connection(
139
+ self, remote: Dict[str, Any], timeout: float = 3.0
140
+ ) -> Tuple[bool, str]:
141
+ """
142
+ Check if a remote proxy is reachable.
143
+
144
+ Args:
145
+ remote: Remote configuration dict
146
+ timeout: Connection timeout in seconds
147
+
148
+ Returns:
149
+ Tuple of (is_online, status_message)
150
+ """
151
+ host = remote.get("host", "127.0.0.1")
152
+ port = remote.get("port", 8000)
153
+ scheme = "https" if port == 443 or "." in host else "http"
154
+ url = f"{scheme}://{host}:{port}/"
155
+
156
+ headers = {}
157
+ if remote.get("api_key"):
158
+ headers["Authorization"] = f"Bearer {remote['api_key']}"
159
+
160
+ try:
161
+ with httpx.Client(timeout=timeout) as client:
162
+ response = client.get(url, headers=headers)
163
+ if response.status_code == 200:
164
+ return True, "Online"
165
+ elif response.status_code == 401:
166
+ return False, "Auth failed"
167
+ else:
168
+ return False, f"HTTP {response.status_code}"
169
+ except httpx.ConnectError:
170
+ return False, "Offline"
171
+ except httpx.TimeoutException:
172
+ return False, "Timeout"
173
+ except Exception as e:
174
+ return False, str(e)[:20]
175
+
176
+ def fetch_stats(self, provider: Optional[str] = None) -> Optional[Dict[str, Any]]:
177
+ """
178
+ Fetch quota stats from the current remote.
179
+
180
+ Args:
181
+ provider: Optional provider filter
182
+
183
+ Returns:
184
+ Stats dict or None on failure
185
+ """
186
+ url = f"{self._get_base_url()}/v1/quota-stats"
187
+ if provider:
188
+ url += f"?provider={provider}"
189
+
190
+ try:
191
+ with httpx.Client(timeout=30.0) as client:
192
+ response = client.get(url, headers=self._get_headers())
193
+
194
+ if response.status_code == 401:
195
+ self.last_error = "Authentication failed. Check API key."
196
+ return None
197
+ elif response.status_code != 200:
198
+ self.last_error = (
199
+ f"HTTP {response.status_code}: {response.text[:100]}"
200
+ )
201
+ return None
202
+
203
+ self.cached_stats = response.json()
204
+ self.last_error = None
205
+ return self.cached_stats
206
+
207
+ except httpx.ConnectError:
208
+ self.last_error = "Connection failed. Is the proxy running?"
209
+ return None
210
+ except httpx.TimeoutException:
211
+ self.last_error = "Request timed out."
212
+ return None
213
+ except Exception as e:
214
+ self.last_error = str(e)
215
+ return None
216
+
217
+ def post_action(
218
+ self,
219
+ action: str,
220
+ scope: str = "all",
221
+ provider: Optional[str] = None,
222
+ credential: Optional[str] = None,
223
+ ) -> Optional[Dict[str, Any]]:
224
+ """
225
+ Post a refresh action to the proxy.
226
+
227
+ Args:
228
+ action: "reload" or "force_refresh"
229
+ scope: "all", "provider", or "credential"
230
+ provider: Provider name (required for scope != "all")
231
+ credential: Credential identifier (required for scope == "credential")
232
+
233
+ Returns:
234
+ Response dict or None on failure
235
+ """
236
+ url = f"{self._get_base_url()}/v1/quota-stats"
237
+ payload = {
238
+ "action": action,
239
+ "scope": scope,
240
+ }
241
+ if provider:
242
+ payload["provider"] = provider
243
+ if credential:
244
+ payload["credential"] = credential
245
+
246
+ try:
247
+ with httpx.Client(timeout=60.0) as client:
248
+ response = client.post(url, headers=self._get_headers(), json=payload)
249
+
250
+ if response.status_code == 401:
251
+ self.last_error = "Authentication failed. Check API key."
252
+ return None
253
+ elif response.status_code != 200:
254
+ self.last_error = (
255
+ f"HTTP {response.status_code}: {response.text[:100]}"
256
+ )
257
+ return None
258
+
259
+ result = response.json()
260
+ self.cached_stats = result
261
+ self.last_error = None
262
+ return result
263
+
264
+ except httpx.ConnectError:
265
+ self.last_error = "Connection failed. Is the proxy running?"
266
+ return None
267
+ except httpx.TimeoutException:
268
+ self.last_error = "Request timed out."
269
+ return None
270
+ except Exception as e:
271
+ self.last_error = str(e)
272
+ return None
273
+
274
+ # =========================================================================
275
+ # DISPLAY SCREENS
276
+ # =========================================================================
277
+
278
+ def show_connection_error(self):
279
+ """Display connection error screen."""
280
+ clear_screen()
281
+ self.console.print(
282
+ Panel(
283
+ Text.from_markup(
284
+ "[bold red]Connection Error[/bold red]\n\n"
285
+ f"{self.last_error or 'Unknown error'}\n\n"
286
+ "[bold]This tool requires the proxy to be running.[/bold]\n"
287
+ "Start the proxy first, then try again.\n\n"
288
+ "[dim]Tip: Select option 1 from the main menu to run the proxy.[/dim]"
289
+ ),
290
+ border_style="red",
291
+ expand=False,
292
+ )
293
+ )
294
+ Prompt.ask("\nPress Enter to return to main menu", default="")
295
+
296
+ def show_summary_screen(self):
297
+ """Display the main summary screen with all providers."""
298
+ clear_screen()
299
+
300
+ # Header
301
+ remote_name = (
302
+ self.current_remote.get("name", "Unknown")
303
+ if self.current_remote
304
+ else "None"
305
+ )
306
+ remote_host = self.current_remote.get("host", "") if self.current_remote else ""
307
+ remote_port = self.current_remote.get("port", "") if self.current_remote else ""
308
+
309
+ # Calculate data age
310
+ data_age = ""
311
+ if self.cached_stats and self.cached_stats.get("timestamp"):
312
+ age_seconds = int(time.time() - self.cached_stats["timestamp"])
313
+ data_age = f"Data age: {age_seconds}s"
314
+
315
+ # View mode indicator
316
+ if self.view_mode == "global":
317
+ view_label = "[magenta]📊 Global/Lifetime[/magenta]"
318
+ else:
319
+ view_label = "[cyan]📈 Current Period[/cyan]"
320
+
321
+ self.console.print("━" * 78)
322
+ self.console.print(
323
+ f"[bold cyan]📈 Quota & Usage Statistics[/bold cyan] | {view_label}"
324
+ )
325
+ self.console.print("━" * 78)
326
+ self.console.print(
327
+ f"Connected to: [bold]{remote_name}[/bold] ({remote_host}:{remote_port}) "
328
+ f"[green]✅[/green] | {data_age}"
329
+ )
330
+ self.console.print()
331
+
332
+ if not self.cached_stats:
333
+ self.console.print("[yellow]No data available. Press R to reload.[/yellow]")
334
+ else:
335
+ # Build provider table
336
+ table = Table(box=None, show_header=True, header_style="bold")
337
+ table.add_column("Provider", style="cyan", min_width=12)
338
+ table.add_column("Creds", justify="center", min_width=6)
339
+ table.add_column("Quota Status", min_width=28)
340
+ table.add_column("Requests", justify="right", min_width=9)
341
+ table.add_column("Tokens (in/out)", min_width=22)
342
+ table.add_column("Cost", justify="right", min_width=8)
343
+
344
+ providers = self.cached_stats.get("providers", {})
345
+ provider_list = list(providers.keys())
346
+
347
+ for idx, (provider, prov_stats) in enumerate(providers.items(), 1):
348
+ cred_count = prov_stats.get("credential_count", 0)
349
+
350
+ # Use global stats if in global mode
351
+ if self.view_mode == "global":
352
+ stats_source = prov_stats.get("global", prov_stats)
353
+ total_requests = stats_source.get("total_requests", 0)
354
+ tokens = stats_source.get("tokens", {})
355
+ cost_value = stats_source.get("approx_cost")
356
+ else:
357
+ total_requests = prov_stats.get("total_requests", 0)
358
+ tokens = prov_stats.get("tokens", {})
359
+ cost_value = prov_stats.get("approx_cost")
360
+
361
+ # Format tokens
362
+ input_total = tokens.get("input_cached", 0) + tokens.get(
363
+ "input_uncached", 0
364
+ )
365
+ output = tokens.get("output", 0)
366
+ cache_pct = tokens.get("input_cache_pct", 0)
367
+ token_str = f"{format_tokens(input_total)}/{format_tokens(output)} ({cache_pct}% cached)"
368
+
369
+ # Format cost
370
+ cost_str = format_cost(cost_value)
371
+
372
+ # Build quota status string (for providers with quota groups)
373
+ quota_groups = prov_stats.get("quota_groups", {})
374
+ if quota_groups:
375
+ quota_lines = []
376
+ for group_name, group_stats in quota_groups.items():
377
+ avg_pct = group_stats.get("avg_remaining_pct", 0)
378
+ exhausted = group_stats.get("credentials_exhausted", 0)
379
+ total = group_stats.get("credentials_total", 0)
380
+
381
+ # Determine color based on remaining
382
+ if exhausted > 0:
383
+ color = "red"
384
+ status = f"({exhausted}/{total} exhausted)"
385
+ elif avg_pct < 20:
386
+ color = "yellow"
387
+ status = ""
388
+ else:
389
+ color = "green"
390
+ status = ""
391
+
392
+ bar = create_progress_bar(avg_pct)
393
+ display_name = group_name[:10]
394
+ quota_lines.append(
395
+ f"[{color}]{display_name}: {avg_pct}% {bar}[/{color}] {status}"
396
+ )
397
+
398
+ # First line goes in the main row
399
+ first_quota = quota_lines[0] if quota_lines else "-"
400
+ table.add_row(
401
+ provider,
402
+ str(cred_count),
403
+ first_quota,
404
+ str(total_requests),
405
+ token_str,
406
+ cost_str,
407
+ )
408
+ # Additional quota lines as sub-rows
409
+ for quota_line in quota_lines[1:]:
410
+ table.add_row("", "", quota_line, "", "", "")
411
+ else:
412
+ # No quota groups
413
+ table.add_row(
414
+ provider,
415
+ str(cred_count),
416
+ "-",
417
+ str(total_requests),
418
+ token_str,
419
+ cost_str,
420
+ )
421
+
422
+ # Add separator between providers (except last)
423
+ if idx < len(providers):
424
+ table.add_row(
425
+ "─" * 10, "─" * 4, "─" * 26, "─" * 7, "─" * 20, "─" * 6
426
+ )
427
+
428
+ self.console.print(table)
429
+
430
+ # Summary line - use global_summary if in global mode
431
+ if self.view_mode == "global":
432
+ summary = self.cached_stats.get(
433
+ "global_summary", self.cached_stats.get("summary", {})
434
+ )
435
+ else:
436
+ summary = self.cached_stats.get("summary", {})
437
+
438
+ total_creds = summary.get("total_credentials", 0)
439
+ total_requests = summary.get("total_requests", 0)
440
+ total_tokens = summary.get("tokens", {})
441
+ total_input = total_tokens.get("input_cached", 0) + total_tokens.get(
442
+ "input_uncached", 0
443
+ )
444
+ total_output = total_tokens.get("output", 0)
445
+ total_cost = format_cost(summary.get("approx_total_cost"))
446
+
447
+ self.console.print()
448
+ self.console.print(
449
+ f"[bold]Total:[/bold] {total_creds} credentials | "
450
+ f"{total_requests} requests | "
451
+ f"{format_tokens(total_input)}/{format_tokens(total_output)} tokens | "
452
+ f"{total_cost} cost"
453
+ )
454
+
455
+ # Menu
456
+ self.console.print()
457
+ self.console.print("━" * 78)
458
+ self.console.print()
459
+
460
+ # Build provider menu options
461
+ providers = self.cached_stats.get("providers", {}) if self.cached_stats else {}
462
+ provider_list = list(providers.keys())
463
+
464
+ for idx, provider in enumerate(provider_list, 1):
465
+ self.console.print(f" {idx}. View [cyan]{provider}[/cyan] details")
466
+
467
+ self.console.print()
468
+ self.console.print(" G. Toggle view mode (current/global)")
469
+ self.console.print(" R. Reload all stats (re-read from proxy)")
470
+ self.console.print(" S. Switch remote")
471
+ self.console.print(" M. Manage remotes")
472
+ self.console.print(" B. Back to main menu")
473
+ self.console.print()
474
+ self.console.print("━" * 78)
475
+
476
+ # Get input
477
+ valid_choices = [str(i) for i in range(1, len(provider_list) + 1)]
478
+ valid_choices.extend(["r", "R", "s", "S", "m", "M", "b", "B", "g", "G"])
479
+
480
+ choice = Prompt.ask("Select option", default="B").strip()
481
+
482
+ if choice.lower() == "b":
483
+ self.running = False
484
+ elif choice.lower() == "g":
485
+ # Toggle view mode
486
+ self.view_mode = "global" if self.view_mode == "current" else "current"
487
+ elif choice.lower() == "r":
488
+ with self.console.status("[bold]Reloading stats...", spinner="dots"):
489
+ self.post_action("reload", scope="all")
490
+ elif choice.lower() == "s":
491
+ self.show_switch_remote_screen()
492
+ elif choice.lower() == "m":
493
+ self.show_manage_remotes_screen()
494
+ elif choice.isdigit() and 1 <= int(choice) <= len(provider_list):
495
+ provider = provider_list[int(choice) - 1]
496
+ self.show_provider_detail_screen(provider)
497
+
498
+ def show_provider_detail_screen(self, provider: str):
499
+ """Display detailed stats for a specific provider."""
500
+ while True:
501
+ clear_screen()
502
+
503
+ # View mode indicator
504
+ if self.view_mode == "global":
505
+ view_label = "[magenta]Global/Lifetime[/magenta]"
506
+ else:
507
+ view_label = "[cyan]Current Period[/cyan]"
508
+
509
+ self.console.print("━" * 78)
510
+ self.console.print(
511
+ f"[bold cyan]📊 {provider.title()} - Detailed Stats[/bold cyan] | {view_label}"
512
+ )
513
+ self.console.print("━" * 78)
514
+ self.console.print()
515
+
516
+ if not self.cached_stats:
517
+ self.console.print("[yellow]No data available.[/yellow]")
518
+ else:
519
+ prov_stats = self.cached_stats.get("providers", {}).get(provider, {})
520
+ credentials = prov_stats.get("credentials", [])
521
+
522
+ if not credentials:
523
+ self.console.print(
524
+ "[dim]No credentials configured for this provider.[/dim]"
525
+ )
526
+ else:
527
+ for idx, cred in enumerate(credentials, 1):
528
+ self._render_credential_panel(idx, cred, provider)
529
+ self.console.print()
530
+
531
+ # Menu
532
+ self.console.print("━" * 78)
533
+ self.console.print()
534
+ self.console.print(" G. Toggle view mode (current/global)")
535
+ self.console.print(" R. Reload stats (from proxy cache)")
536
+ self.console.print(" RA. Reload all stats")
537
+
538
+ # Force refresh options (only for providers that support it)
539
+ has_quota_groups = bool(
540
+ self.cached_stats
541
+ and self.cached_stats.get("providers", {})
542
+ .get(provider, {})
543
+ .get("quota_groups")
544
+ )
545
+
546
+ if has_quota_groups:
547
+ self.console.print()
548
+ self.console.print(
549
+ f" F. [yellow]Force refresh ALL {provider} quotas from API[/yellow]"
550
+ )
551
+ credentials = (
552
+ self.cached_stats.get("providers", {})
553
+ .get(provider, {})
554
+ .get("credentials", [])
555
+ if self.cached_stats
556
+ else []
557
+ )
558
+ for idx, cred in enumerate(credentials, 1):
559
+ identifier = cred.get("identifier", f"credential {idx}")
560
+ email = cred.get("email", identifier)
561
+ self.console.print(
562
+ f" F{idx}. Force refresh [{idx}] only ({email})"
563
+ )
564
+
565
+ self.console.print()
566
+ self.console.print(" B. Back to summary")
567
+ self.console.print()
568
+ self.console.print("━" * 78)
569
+
570
+ choice = Prompt.ask("Select option", default="B").strip().upper()
571
+
572
+ if choice == "B":
573
+ break
574
+ elif choice == "G":
575
+ # Toggle view mode
576
+ self.view_mode = "global" if self.view_mode == "current" else "current"
577
+ elif choice == "R":
578
+ with self.console.status(
579
+ f"[bold]Reloading {provider} stats...", spinner="dots"
580
+ ):
581
+ self.post_action("reload", scope="provider", provider=provider)
582
+ elif choice == "RA":
583
+ with self.console.status(
584
+ "[bold]Reloading all stats...", spinner="dots"
585
+ ):
586
+ self.post_action("reload", scope="all")
587
+ elif choice == "F" and has_quota_groups:
588
+ with self.console.status(
589
+ f"[bold]Fetching live quota for ALL {provider} credentials...",
590
+ spinner="dots",
591
+ ):
592
+ result = self.post_action(
593
+ "force_refresh", scope="provider", provider=provider
594
+ )
595
+ if result and result.get("refresh_result"):
596
+ rr = result["refresh_result"]
597
+ self.console.print(
598
+ f"\n[green]Refreshed {rr.get('credentials_refreshed', 0)} credentials "
599
+ f"in {rr.get('duration_ms', 0)}ms[/green]"
600
+ )
601
+ if rr.get("errors"):
602
+ for err in rr["errors"]:
603
+ self.console.print(f"[red] Error: {err}[/red]")
604
+ Prompt.ask("Press Enter to continue", default="")
605
+ elif choice.startswith("F") and choice[1:].isdigit() and has_quota_groups:
606
+ idx = int(choice[1:])
607
+ credentials = (
608
+ self.cached_stats.get("providers", {})
609
+ .get(provider, {})
610
+ .get("credentials", [])
611
+ if self.cached_stats
612
+ else []
613
+ )
614
+ if 1 <= idx <= len(credentials):
615
+ cred = credentials[idx - 1]
616
+ cred_id = cred.get("identifier", "")
617
+ email = cred.get("email", cred_id)
618
+ with self.console.status(
619
+ f"[bold]Fetching live quota for {email}...", spinner="dots"
620
+ ):
621
+ result = self.post_action(
622
+ "force_refresh",
623
+ scope="credential",
624
+ provider=provider,
625
+ credential=cred_id,
626
+ )
627
+ if result and result.get("refresh_result"):
628
+ rr = result["refresh_result"]
629
+ self.console.print(
630
+ f"\n[green]Refreshed in {rr.get('duration_ms', 0)}ms[/green]"
631
+ )
632
+ if rr.get("errors"):
633
+ for err in rr["errors"]:
634
+ self.console.print(f"[red] Error: {err}[/red]")
635
+ Prompt.ask("Press Enter to continue", default="")
636
+
637
+ def _render_credential_panel(self, idx: int, cred: Dict[str, Any], provider: str):
638
+ """Render a single credential as a panel."""
639
+ identifier = cred.get("identifier", f"credential {idx}")
640
+ email = cred.get("email")
641
+ tier = cred.get("tier", "")
642
+ status = cred.get("status", "unknown")
643
+
644
+ # Check for active cooldowns
645
+ key_cooldown = cred.get("key_cooldown_remaining")
646
+ model_cooldowns = cred.get("model_cooldowns", {})
647
+ has_cooldown = key_cooldown or model_cooldowns
648
+
649
+ # Status indicator
650
+ if status == "exhausted":
651
+ status_icon = "[red]⛔ Exhausted[/red]"
652
+ elif status == "cooldown" or has_cooldown:
653
+ if key_cooldown:
654
+ status_icon = f"[yellow]⚠️ Cooldown ({format_cooldown(int(key_cooldown))})[/yellow]"
655
+ else:
656
+ status_icon = "[yellow]⚠️ Cooldown[/yellow]"
657
+ else:
658
+ status_icon = "[green]✅ Active[/green]"
659
+
660
+ # Header line
661
+ display_name = email if email else identifier
662
+ tier_str = f" ({tier})" if tier else ""
663
+ header = f"[{idx}] {display_name}{tier_str} {status_icon}"
664
+
665
+ # Use global stats if in global mode
666
+ if self.view_mode == "global":
667
+ stats_source = cred.get("global", cred)
668
+ else:
669
+ stats_source = cred
670
+
671
+ # Stats line
672
+ last_used = format_time_ago(cred.get("last_used_ts")) # Always from current
673
+ requests = stats_source.get("requests", 0)
674
+ tokens = stats_source.get("tokens", {})
675
+ input_total = tokens.get("input_cached", 0) + tokens.get("input_uncached", 0)
676
+ output = tokens.get("output", 0)
677
+ cost = format_cost(stats_source.get("approx_cost"))
678
+
679
+ stats_line = (
680
+ f"Last used: {last_used} | Requests: {requests} | "
681
+ f"Tokens: {format_tokens(input_total)}/{format_tokens(output)}"
682
+ )
683
+ if cost != "-":
684
+ stats_line += f" | Cost: {cost}"
685
+
686
+ # Build panel content
687
+ content_lines = [
688
+ f"[dim]{stats_line}[/dim]",
689
+ ]
690
+
691
+ # Show model cooldowns if any
692
+ if model_cooldowns:
693
+ content_lines.append("")
694
+ content_lines.append("[yellow]Active Cooldowns:[/yellow]")
695
+ for model_name, cooldown_info in model_cooldowns.items():
696
+ remaining = cooldown_info.get("remaining_seconds", 0)
697
+ if remaining > 0:
698
+ # Shorten model name for display
699
+ short_model = model_name.split("/")[-1][:35]
700
+ content_lines.append(
701
+ f" [yellow]⏱️ {short_model}: {format_cooldown(int(remaining))}[/yellow]"
702
+ )
703
+
704
+ # Model groups (for providers with quota tracking)
705
+ model_groups = cred.get("model_groups", {})
706
+ if model_groups:
707
+ content_lines.append("")
708
+ for group_name, group_stats in model_groups.items():
709
+ remaining_pct = group_stats.get("remaining_pct")
710
+ requests_used = group_stats.get("requests_used", 0)
711
+ requests_max = group_stats.get("requests_max")
712
+ is_exhausted = group_stats.get("is_exhausted", False)
713
+ reset_time = format_reset_time(group_stats.get("reset_time_iso"))
714
+ confidence = group_stats.get("confidence", "low")
715
+
716
+ # Format display
717
+ display = group_stats.get("display", f"{requests_used}/?")
718
+ bar = create_progress_bar(remaining_pct)
719
+
720
+ # Color based on status
721
+ if is_exhausted:
722
+ color = "red"
723
+ status_text = "⛔ EXHAUSTED"
724
+ elif remaining_pct is not None and remaining_pct < 20:
725
+ color = "yellow"
726
+ status_text = "⚠️ LOW"
727
+ else:
728
+ color = "green"
729
+ status_text = f"Resets: {reset_time}"
730
+
731
+ # Confidence indicator
732
+ conf_indicator = ""
733
+ if confidence == "low":
734
+ conf_indicator = " [dim](~)[/dim]"
735
+ elif confidence == "medium":
736
+ conf_indicator = " [dim](?)[/dim]"
737
+
738
+ pct_str = f"{remaining_pct}%" if remaining_pct is not None else "?%"
739
+ content_lines.append(
740
+ f" [{color}]{group_name:<18} {display:<10} {pct_str:>4} {bar}[/{color}] {status_text}{conf_indicator}"
741
+ )
742
+ else:
743
+ # For providers without quota groups, show model breakdown if available
744
+ models = cred.get("models", {})
745
+ if models:
746
+ content_lines.append("")
747
+ content_lines.append(" [dim]Models used:[/dim]")
748
+ for model_name, model_stats in models.items():
749
+ req_count = model_stats.get("success_count", 0)
750
+ model_cost = format_cost(model_stats.get("approx_cost"))
751
+ # Shorten model name for display
752
+ short_name = model_name.split("/")[-1][:30]
753
+ content_lines.append(
754
+ f" {short_name}: {req_count} requests, {model_cost}"
755
+ )
756
+
757
+ self.console.print(
758
+ Panel(
759
+ "\n".join(content_lines),
760
+ title=header,
761
+ title_align="left",
762
+ border_style="dim",
763
+ expand=True,
764
+ )
765
+ )
766
+
767
+ def show_switch_remote_screen(self):
768
+ """Display remote selection screen."""
769
+ clear_screen()
770
+
771
+ self.console.print("━" * 78)
772
+ self.console.print("[bold cyan]🔄 Switch Remote[/bold cyan]")
773
+ self.console.print("━" * 78)
774
+ self.console.print()
775
+
776
+ current_name = self.current_remote.get("name") if self.current_remote else None
777
+ self.console.print(f"Current: [bold]{current_name}[/bold]")
778
+ self.console.print()
779
+ self.console.print("Available remotes:")
780
+
781
+ remotes = self.config.get_remotes()
782
+ remote_status: List[Tuple[Dict, bool, str]] = []
783
+
784
+ # Check status of all remotes
785
+ with self.console.status("[dim]Checking remote status...", spinner="dots"):
786
+ for remote in remotes:
787
+ is_online, status_msg = self.check_connection(remote)
788
+ remote_status.append((remote, is_online, status_msg))
789
+
790
+ for idx, (remote, is_online, status_msg) in enumerate(remote_status, 1):
791
+ name = remote.get("name", "Unknown")
792
+ host = remote.get("host", "")
793
+ port = remote.get("port", 8000)
794
+
795
+ is_current = name == current_name
796
+ current_marker = " (current)" if is_current else ""
797
+
798
+ if is_online:
799
+ status_icon = "[green]✅ Online[/green]"
800
+ else:
801
+ status_icon = f"[red]⚠️ {status_msg}[/red]"
802
+
803
+ self.console.print(
804
+ f" {idx}. {name:<20} {host}:{port:<6} {status_icon}{current_marker}"
805
+ )
806
+
807
+ self.console.print()
808
+ self.console.print("━" * 78)
809
+ self.console.print()
810
+
811
+ choice = Prompt.ask(
812
+ f"Select remote (1-{len(remotes)}) or B to go back", default="B"
813
+ ).strip()
814
+
815
+ if choice.lower() == "b":
816
+ return
817
+
818
+ if choice.isdigit() and 1 <= int(choice) <= len(remotes):
819
+ selected = remotes[int(choice) - 1]
820
+ self.current_remote = selected
821
+ self.config.set_last_used(selected["name"])
822
+ self.cached_stats = None # Clear cache
823
+
824
+ # Try to fetch stats from new remote
825
+ with self.console.status("[bold]Connecting...", spinner="dots"):
826
+ stats = self.fetch_stats()
827
+ if stats is None:
828
+ # Try with API key from .env for Local
829
+ if selected["name"] == "Local" and not selected.get("api_key"):
830
+ env_key = self.config.get_api_key_from_env()
831
+ if env_key:
832
+ self.current_remote["api_key"] = env_key
833
+ stats = self.fetch_stats()
834
+
835
+ if stats is None:
836
+ self.show_api_key_prompt()
837
+
838
+ def show_api_key_prompt(self):
839
+ """Prompt for API key when authentication fails."""
840
+ self.console.print()
841
+ self.console.print(
842
+ "[yellow]Authentication required or connection failed.[/yellow]"
843
+ )
844
+ self.console.print(f"Error: {self.last_error}")
845
+ self.console.print()
846
+
847
+ api_key = Prompt.ask(
848
+ "Enter API key (or press Enter to cancel)", default=""
849
+ ).strip()
850
+
851
+ if api_key:
852
+ self.current_remote["api_key"] = api_key
853
+ # Update config with new API key
854
+ self.config.update_remote(self.current_remote["name"], api_key=api_key)
855
+
856
+ # Try again
857
+ with self.console.status("[bold]Reconnecting...", spinner="dots"):
858
+ if self.fetch_stats() is None:
859
+ self.console.print(f"[red]Still failed: {self.last_error}[/red]")
860
+ Prompt.ask("Press Enter to continue", default="")
861
+ else:
862
+ self.console.print("[dim]Cancelled.[/dim]")
863
+ Prompt.ask("Press Enter to continue", default="")
864
+
865
+ def show_manage_remotes_screen(self):
866
+ """Display remote management screen."""
867
+ while True:
868
+ clear_screen()
869
+
870
+ self.console.print("━" * 78)
871
+ self.console.print("[bold cyan]⚙️ Manage Remotes[/bold cyan]")
872
+ self.console.print("━" * 78)
873
+ self.console.print()
874
+
875
+ remotes = self.config.get_remotes()
876
+
877
+ table = Table(box=None, show_header=True, header_style="bold")
878
+ table.add_column("#", style="dim", width=3)
879
+ table.add_column("Name", min_width=16)
880
+ table.add_column("Host", min_width=24)
881
+ table.add_column("Port", justify="right", width=6)
882
+ table.add_column("Default", width=8)
883
+
884
+ for idx, remote in enumerate(remotes, 1):
885
+ is_default = "★" if remote.get("is_default") else ""
886
+ table.add_row(
887
+ str(idx),
888
+ remote.get("name", ""),
889
+ remote.get("host", ""),
890
+ str(remote.get("port", 8000)),
891
+ is_default,
892
+ )
893
+
894
+ self.console.print(table)
895
+
896
+ self.console.print()
897
+ self.console.print("━" * 78)
898
+ self.console.print()
899
+ self.console.print(" A. Add new remote")
900
+ self.console.print(" E. Edit remote (enter number, e.g., E1)")
901
+ self.console.print(" D. Delete remote (enter number, e.g., D1)")
902
+ self.console.print(" S. Set default remote")
903
+ self.console.print(" B. Back")
904
+ self.console.print()
905
+ self.console.print("━" * 78)
906
+
907
+ choice = Prompt.ask("Select option", default="B").strip().upper()
908
+
909
+ if choice == "B":
910
+ break
911
+ elif choice == "A":
912
+ self._add_remote_dialog()
913
+ elif choice == "S":
914
+ self._set_default_dialog(remotes)
915
+ elif choice.startswith("E") and choice[1:].isdigit():
916
+ idx = int(choice[1:])
917
+ if 1 <= idx <= len(remotes):
918
+ self._edit_remote_dialog(remotes[idx - 1])
919
+ elif choice.startswith("D") and choice[1:].isdigit():
920
+ idx = int(choice[1:])
921
+ if 1 <= idx <= len(remotes):
922
+ self._delete_remote_dialog(remotes[idx - 1])
923
+
924
+ def _add_remote_dialog(self):
925
+ """Dialog to add a new remote."""
926
+ self.console.print()
927
+ self.console.print("[bold]Add New Remote[/bold]")
928
+ self.console.print()
929
+
930
+ name = Prompt.ask("Name", default="").strip()
931
+ if not name:
932
+ self.console.print("[dim]Cancelled.[/dim]")
933
+ return
934
+
935
+ host = Prompt.ask("Host", default="").strip()
936
+ if not host:
937
+ self.console.print("[dim]Cancelled.[/dim]")
938
+ return
939
+
940
+ port_str = Prompt.ask("Port", default="8000").strip()
941
+ try:
942
+ port = int(port_str)
943
+ except ValueError:
944
+ port = 8000
945
+
946
+ api_key = Prompt.ask("API Key (optional)", default="").strip() or None
947
+
948
+ if self.config.add_remote(name, host, port, api_key):
949
+ self.console.print(f"[green]Added remote '{name}'.[/green]")
950
+ else:
951
+ self.console.print(f"[red]Remote '{name}' already exists.[/red]")
952
+
953
+ Prompt.ask("Press Enter to continue", default="")
954
+
955
+ def _edit_remote_dialog(self, remote: Dict[str, Any]):
956
+ """Dialog to edit an existing remote."""
957
+ self.console.print()
958
+ self.console.print(f"[bold]Edit Remote: {remote['name']}[/bold]")
959
+ self.console.print("[dim]Press Enter to keep current value[/dim]")
960
+ self.console.print()
961
+
962
+ new_name = Prompt.ask("Name", default=remote["name"]).strip()
963
+ new_host = Prompt.ask("Host", default=remote.get("host", "")).strip()
964
+ new_port_str = Prompt.ask("Port", default=str(remote.get("port", 8000))).strip()
965
+ try:
966
+ new_port = int(new_port_str)
967
+ except ValueError:
968
+ new_port = remote.get("port", 8000)
969
+
970
+ current_key = remote.get("api_key", "") or ""
971
+ display_key = f"{current_key[:8]}..." if len(current_key) > 8 else current_key
972
+ new_key = Prompt.ask(
973
+ f"API Key (current: {display_key or 'none'})", default=""
974
+ ).strip()
975
+
976
+ updates = {}
977
+ if new_name != remote["name"]:
978
+ updates["new_name"] = new_name
979
+ if new_host != remote.get("host"):
980
+ updates["host"] = new_host
981
+ if new_port != remote.get("port"):
982
+ updates["port"] = new_port
983
+ if new_key:
984
+ updates["api_key"] = new_key
985
+
986
+ if updates:
987
+ if self.config.update_remote(remote["name"], **updates):
988
+ self.console.print("[green]Remote updated.[/green]")
989
+ # Update current_remote if it was the one being edited
990
+ if (
991
+ self.current_remote
992
+ and self.current_remote["name"] == remote["name"]
993
+ ):
994
+ self.current_remote.update(updates)
995
+ if "new_name" in updates:
996
+ self.current_remote["name"] = updates["new_name"]
997
+ else:
998
+ self.console.print("[red]Failed to update remote.[/red]")
999
+ else:
1000
+ self.console.print("[dim]No changes made.[/dim]")
1001
+
1002
+ Prompt.ask("Press Enter to continue", default="")
1003
+
1004
+ def _delete_remote_dialog(self, remote: Dict[str, Any]):
1005
+ """Dialog to delete a remote."""
1006
+ self.console.print()
1007
+ self.console.print(f"[yellow]Delete remote '{remote['name']}'?[/yellow]")
1008
+
1009
+ confirm = Prompt.ask("Type 'yes' to confirm", default="no").strip().lower()
1010
+
1011
+ if confirm == "yes":
1012
+ if self.config.delete_remote(remote["name"]):
1013
+ self.console.print(f"[green]Deleted remote '{remote['name']}'.[/green]")
1014
+ # If deleted current remote, switch to another
1015
+ if (
1016
+ self.current_remote
1017
+ and self.current_remote["name"] == remote["name"]
1018
+ ):
1019
+ self.current_remote = self.config.get_default_remote()
1020
+ self.cached_stats = None
1021
+ else:
1022
+ self.console.print(
1023
+ "[red]Cannot delete. At least one remote must exist.[/red]"
1024
+ )
1025
+ else:
1026
+ self.console.print("[dim]Cancelled.[/dim]")
1027
+
1028
+ Prompt.ask("Press Enter to continue", default="")
1029
+
1030
+ def _set_default_dialog(self, remotes: List[Dict[str, Any]]):
1031
+ """Dialog to set the default remote."""
1032
+ self.console.print()
1033
+ choice = Prompt.ask(f"Set default (1-{len(remotes)})", default="").strip()
1034
+
1035
+ if choice.isdigit() and 1 <= int(choice) <= len(remotes):
1036
+ remote = remotes[int(choice) - 1]
1037
+ if self.config.set_default_remote(remote["name"]):
1038
+ self.console.print(
1039
+ f"[green]'{remote['name']}' is now the default.[/green]"
1040
+ )
1041
+ else:
1042
+ self.console.print("[red]Failed to set default.[/red]")
1043
+ Prompt.ask("Press Enter to continue", default="")
1044
+
1045
+ # =========================================================================
1046
+ # MAIN LOOP
1047
+ # =========================================================================
1048
+
1049
+ def run(self):
1050
+ """Main viewer loop."""
1051
+ # Get initial remote
1052
+ self.current_remote = self.config.get_last_used_remote()
1053
+
1054
+ if not self.current_remote:
1055
+ self.console.print("[red]No remotes configured.[/red]")
1056
+ return
1057
+
1058
+ # For Local remote, try to get API key from .env if not set
1059
+ if self.current_remote["name"] == "Local" and not self.current_remote.get(
1060
+ "api_key"
1061
+ ):
1062
+ env_key = self.config.get_api_key_from_env()
1063
+ if env_key:
1064
+ self.current_remote["api_key"] = env_key
1065
+
1066
+ # Initial fetch
1067
+ with self.console.status("[bold]Connecting to proxy...", spinner="dots"):
1068
+ stats = self.fetch_stats()
1069
+
1070
+ if stats is None:
1071
+ self.show_connection_error()
1072
+ return
1073
+
1074
+ # Main loop
1075
+ while self.running:
1076
+ self.show_summary_screen()
1077
+
1078
+
1079
+ def run_quota_viewer():
1080
+ """Entry point for the quota viewer."""
1081
+ viewer = QuotaViewer()
1082
+ viewer.run()
1083
+
1084
+
1085
+ if __name__ == "__main__":
1086
+ run_quota_viewer()
src/proxy_app/quota_viewer_config.py ADDED
@@ -0,0 +1,288 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Configuration management for the Quota Viewer.
3
+
4
+ Handles remote proxy configurations including:
5
+ - Multiple remote proxies (local, VPS, etc.)
6
+ - API key storage per remote
7
+ - Default and last-used remote tracking
8
+ """
9
+
10
+ import json
11
+ from pathlib import Path
12
+ from typing import Any, Dict, List, Optional
13
+
14
+
15
+ class QuotaViewerConfig:
16
+ """Manages quota viewer configuration including remote proxies."""
17
+
18
+ def __init__(self, config_path: Optional[Path] = None):
19
+ """
20
+ Initialize the config manager.
21
+
22
+ Args:
23
+ config_path: Path to config file. Defaults to quota_viewer_config.json
24
+ in the current directory or EXE directory.
25
+ """
26
+ if config_path is None:
27
+ import sys
28
+
29
+ if getattr(sys, "frozen", False):
30
+ base_dir = Path(sys.executable).parent
31
+ else:
32
+ base_dir = Path.cwd()
33
+ config_path = base_dir / "quota_viewer_config.json"
34
+
35
+ self.config_path = config_path
36
+ self.config = self._load()
37
+
38
+ def _load(self) -> Dict[str, Any]:
39
+ """Load config from file or return defaults."""
40
+ if self.config_path.exists():
41
+ try:
42
+ with open(self.config_path, "r", encoding="utf-8") as f:
43
+ config = json.load(f)
44
+ # Ensure required fields exist
45
+ if "remotes" not in config:
46
+ config["remotes"] = []
47
+ return config
48
+ except (json.JSONDecodeError, IOError):
49
+ pass
50
+
51
+ # Return default config with Local remote
52
+ return {
53
+ "remotes": [
54
+ {
55
+ "name": "Local",
56
+ "host": "127.0.0.1",
57
+ "port": 8000,
58
+ "api_key": None,
59
+ "is_default": True,
60
+ }
61
+ ],
62
+ "last_used": "Local",
63
+ }
64
+
65
+ def _save(self) -> bool:
66
+ """Save config to file. Returns True on success."""
67
+ try:
68
+ with open(self.config_path, "w", encoding="utf-8") as f:
69
+ json.dump(self.config, f, indent=2)
70
+ return True
71
+ except IOError:
72
+ return False
73
+
74
+ def get_remotes(self) -> List[Dict[str, Any]]:
75
+ """Get list of all configured remotes."""
76
+ return self.config.get("remotes", [])
77
+
78
+ def get_remote_by_name(self, name: str) -> Optional[Dict[str, Any]]:
79
+ """Get a remote by name."""
80
+ for remote in self.config.get("remotes", []):
81
+ if remote["name"] == name:
82
+ return remote
83
+ return None
84
+
85
+ def get_default_remote(self) -> Optional[Dict[str, Any]]:
86
+ """Get the default remote."""
87
+ for remote in self.config.get("remotes", []):
88
+ if remote.get("is_default"):
89
+ return remote
90
+ # Fallback to first remote
91
+ remotes = self.config.get("remotes", [])
92
+ return remotes[0] if remotes else None
93
+
94
+ def get_last_used_remote(self) -> Optional[Dict[str, Any]]:
95
+ """Get the last used remote, or default if not set."""
96
+ last_used_name = self.config.get("last_used")
97
+ if last_used_name:
98
+ remote = self.get_remote_by_name(last_used_name)
99
+ if remote:
100
+ return remote
101
+ return self.get_default_remote()
102
+
103
+ def set_last_used(self, name: str) -> bool:
104
+ """Set the last used remote name."""
105
+ self.config["last_used"] = name
106
+ return self._save()
107
+
108
+ def add_remote(
109
+ self,
110
+ name: str,
111
+ host: str,
112
+ port: int = 8000,
113
+ api_key: Optional[str] = None,
114
+ is_default: bool = False,
115
+ ) -> bool:
116
+ """
117
+ Add a new remote configuration.
118
+
119
+ Args:
120
+ name: Display name for the remote
121
+ host: Hostname or IP address
122
+ port: Port number (default 8000)
123
+ api_key: Optional API key for authentication
124
+ is_default: Whether this should be the default remote
125
+
126
+ Returns:
127
+ True on success, False if name already exists
128
+ """
129
+ # Check for duplicate name
130
+ if self.get_remote_by_name(name):
131
+ return False
132
+
133
+ # If setting as default, clear default from others
134
+ if is_default:
135
+ for remote in self.config.get("remotes", []):
136
+ remote["is_default"] = False
137
+
138
+ remote = {
139
+ "name": name,
140
+ "host": host,
141
+ "port": port,
142
+ "api_key": api_key,
143
+ "is_default": is_default,
144
+ }
145
+ self.config.setdefault("remotes", []).append(remote)
146
+ return self._save()
147
+
148
+ def update_remote(self, name: str, **kwargs) -> bool:
149
+ """
150
+ Update an existing remote configuration.
151
+
152
+ Args:
153
+ name: Name of the remote to update
154
+ **kwargs: Fields to update (host, port, api_key, is_default, new_name)
155
+
156
+ Returns:
157
+ True on success, False if remote not found
158
+ """
159
+ remote = self.get_remote_by_name(name)
160
+ if not remote:
161
+ return False
162
+
163
+ # Handle rename
164
+ if "new_name" in kwargs:
165
+ new_name = kwargs.pop("new_name")
166
+ if new_name != name and self.get_remote_by_name(new_name):
167
+ return False # New name already exists
168
+ remote["name"] = new_name
169
+ # Update last_used if it was this remote
170
+ if self.config.get("last_used") == name:
171
+ self.config["last_used"] = new_name
172
+
173
+ # If setting as default, clear default from others
174
+ if kwargs.get("is_default"):
175
+ for r in self.config.get("remotes", []):
176
+ r["is_default"] = False
177
+
178
+ # Update other fields
179
+ for key in ("host", "port", "api_key", "is_default"):
180
+ if key in kwargs:
181
+ remote[key] = kwargs[key]
182
+
183
+ return self._save()
184
+
185
+ def delete_remote(self, name: str) -> bool:
186
+ """
187
+ Delete a remote configuration.
188
+
189
+ Args:
190
+ name: Name of the remote to delete
191
+
192
+ Returns:
193
+ True on success, False if remote not found or is the only one
194
+ """
195
+ remotes = self.config.get("remotes", [])
196
+ if len(remotes) <= 1:
197
+ return False # Don't delete the last remote
198
+
199
+ for i, remote in enumerate(remotes):
200
+ if remote["name"] == name:
201
+ remotes.pop(i)
202
+ # Update last_used if it was this remote
203
+ if self.config.get("last_used") == name:
204
+ self.config["last_used"] = remotes[0]["name"] if remotes else None
205
+ return self._save()
206
+ return False
207
+
208
+ def set_default_remote(self, name: str) -> bool:
209
+ """Set a remote as the default."""
210
+ remote = self.get_remote_by_name(name)
211
+ if not remote:
212
+ return False
213
+
214
+ # Clear default from all remotes
215
+ for r in self.config.get("remotes", []):
216
+ r["is_default"] = False
217
+
218
+ # Set new default
219
+ remote["is_default"] = True
220
+ return self._save()
221
+
222
+ def sync_with_launcher_config(self) -> None:
223
+ """
224
+ Sync the Local remote with launcher_config.json if it exists.
225
+
226
+ This ensures the Local remote always matches the launcher settings.
227
+ """
228
+ import sys
229
+
230
+ if getattr(sys, "frozen", False):
231
+ base_dir = Path(sys.executable).parent
232
+ else:
233
+ base_dir = Path.cwd()
234
+
235
+ launcher_config_path = base_dir / "launcher_config.json"
236
+
237
+ if launcher_config_path.exists():
238
+ try:
239
+ with open(launcher_config_path, "r", encoding="utf-8") as f:
240
+ launcher_config = json.load(f)
241
+
242
+ host = launcher_config.get("host", "127.0.0.1")
243
+ port = launcher_config.get("port", 8000)
244
+
245
+ # Update Local remote
246
+ local_remote = self.get_remote_by_name("Local")
247
+ if local_remote:
248
+ local_remote["host"] = host
249
+ local_remote["port"] = port
250
+ self._save()
251
+ else:
252
+ # Create Local remote if it doesn't exist
253
+ self.add_remote("Local", host, port, is_default=True)
254
+
255
+ except (json.JSONDecodeError, IOError):
256
+ pass
257
+
258
+ def get_api_key_from_env(self) -> Optional[str]:
259
+ """
260
+ Get PROXY_API_KEY from .env file for Local remote.
261
+
262
+ Returns:
263
+ API key string or None
264
+ """
265
+ import sys
266
+
267
+ if getattr(sys, "frozen", False):
268
+ base_dir = Path(sys.executable).parent
269
+ else:
270
+ base_dir = Path.cwd()
271
+
272
+ env_path = base_dir / ".env"
273
+ if not env_path.exists():
274
+ return None
275
+
276
+ try:
277
+ with open(env_path, "r", encoding="utf-8") as f:
278
+ for line in f:
279
+ line = line.strip()
280
+ if line.startswith("PROXY_API_KEY="):
281
+ value = line.split("=", 1)[1].strip()
282
+ # Remove quotes if present
283
+ if value and value[0] in ('"', "'") and value[-1] == value[0]:
284
+ value = value[1:-1]
285
+ return value if value else None
286
+ except IOError:
287
+ pass
288
+ return None
src/rotator_library/client.py CHANGED
@@ -2612,3 +2612,280 @@ class RotatingClient:
2612
  for models in all_provider_models.values():
2613
  flat_models.extend(models)
2614
  return flat_models
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2612
  for models in all_provider_models.values():
2613
  flat_models.extend(models)
2614
  return flat_models
2615
+
2616
+ async def get_quota_stats(
2617
+ self,
2618
+ provider_filter: Optional[str] = None,
2619
+ ) -> Dict[str, Any]:
2620
+ """
2621
+ Get quota and usage stats for all credentials.
2622
+
2623
+ This returns cached/disk data aggregated by provider.
2624
+ For provider-specific quota info (e.g., Antigravity quota groups),
2625
+ it enriches the data from provider plugins.
2626
+
2627
+ Args:
2628
+ provider_filter: If provided, only return stats for this provider
2629
+
2630
+ Returns:
2631
+ Complete stats dict ready for the /v1/quota-stats endpoint
2632
+ """
2633
+ # Get base stats from usage manager
2634
+ stats = await self.usage_manager.get_stats_for_endpoint(provider_filter)
2635
+
2636
+ # Enrich with provider-specific quota data
2637
+ for provider, prov_stats in stats.get("providers", {}).items():
2638
+ provider_class = self._provider_plugins.get(provider)
2639
+ if not provider_class:
2640
+ continue
2641
+
2642
+ # Get or create provider instance
2643
+ if provider not in self._provider_instances:
2644
+ self._provider_instances[provider] = provider_class()
2645
+ provider_instance = self._provider_instances[provider]
2646
+
2647
+ # Check if provider has quota tracking (like Antigravity)
2648
+ if hasattr(provider_instance, "_get_effective_quota_groups"):
2649
+ # Add quota group summary
2650
+ quota_groups = provider_instance._get_effective_quota_groups()
2651
+ prov_stats["quota_groups"] = {}
2652
+
2653
+ for group_name, group_models in quota_groups.items():
2654
+ group_stats = {
2655
+ "models": group_models,
2656
+ "credentials_total": 0,
2657
+ "credentials_exhausted": 0,
2658
+ "avg_remaining_pct": 0,
2659
+ "total_remaining_pcts": [],
2660
+ }
2661
+
2662
+ # Calculate per-credential quota for this group
2663
+ for cred in prov_stats.get("credentials", []):
2664
+ models_data = cred.get("models", {})
2665
+ group_stats["credentials_total"] += 1
2666
+
2667
+ # Find any model from this group
2668
+ for model in group_models:
2669
+ # Try with and without provider prefix
2670
+ prefixed_model = f"{provider}/{model}"
2671
+ model_stats = models_data.get(
2672
+ prefixed_model
2673
+ ) or models_data.get(model)
2674
+
2675
+ if model_stats:
2676
+ baseline = model_stats.get(
2677
+ "baseline_remaining_fraction"
2678
+ )
2679
+ if baseline is not None:
2680
+ remaining_pct = int(baseline * 100)
2681
+ group_stats["total_remaining_pcts"].append(
2682
+ remaining_pct
2683
+ )
2684
+ if baseline <= 0:
2685
+ group_stats["credentials_exhausted"] += 1
2686
+ break
2687
+
2688
+ # Calculate average remaining percentage
2689
+ if group_stats["total_remaining_pcts"]:
2690
+ group_stats["avg_remaining_pct"] = int(
2691
+ sum(group_stats["total_remaining_pcts"])
2692
+ / len(group_stats["total_remaining_pcts"])
2693
+ )
2694
+ del group_stats["total_remaining_pcts"]
2695
+
2696
+ prov_stats["quota_groups"][group_name] = group_stats
2697
+
2698
+ # Also enrich each credential with formatted quota group info
2699
+ for cred in prov_stats.get("credentials", []):
2700
+ cred["model_groups"] = {}
2701
+ models_data = cred.get("models", {})
2702
+
2703
+ for group_name, group_models in quota_groups.items():
2704
+ # Find representative model from this group
2705
+ for model in group_models:
2706
+ prefixed_model = f"{provider}/{model}"
2707
+ model_stats = models_data.get(
2708
+ prefixed_model
2709
+ ) or models_data.get(model)
2710
+
2711
+ if model_stats:
2712
+ baseline = model_stats.get(
2713
+ "baseline_remaining_fraction"
2714
+ )
2715
+ max_req = model_stats.get("quota_max_requests")
2716
+ req_count = model_stats.get("request_count", 0)
2717
+ reset_ts = model_stats.get("quota_reset_ts")
2718
+
2719
+ remaining_pct = (
2720
+ int(baseline * 100)
2721
+ if baseline is not None
2722
+ else None
2723
+ )
2724
+ is_exhausted = baseline is not None and baseline <= 0
2725
+
2726
+ # Format reset time
2727
+ reset_iso = None
2728
+ if reset_ts:
2729
+ try:
2730
+ from datetime import datetime, timezone
2731
+
2732
+ reset_iso = datetime.fromtimestamp(
2733
+ reset_ts, tz=timezone.utc
2734
+ ).isoformat()
2735
+ except (ValueError, OSError):
2736
+ pass
2737
+
2738
+ cred["model_groups"][group_name] = {
2739
+ "remaining_pct": remaining_pct,
2740
+ "requests_used": req_count,
2741
+ "requests_max": max_req,
2742
+ "display": f"{req_count}/{max_req}"
2743
+ if max_req
2744
+ else f"{req_count}/?",
2745
+ "is_exhausted": is_exhausted,
2746
+ "reset_time_iso": reset_iso,
2747
+ "models": group_models,
2748
+ "confidence": self._get_baseline_confidence(
2749
+ model_stats
2750
+ ),
2751
+ }
2752
+ break
2753
+
2754
+ # Try to get email from provider's cache
2755
+ cred_path = cred.get("full_path", "")
2756
+ if hasattr(provider_instance, "project_tier_cache"):
2757
+ tier = provider_instance.project_tier_cache.get(cred_path)
2758
+ if tier:
2759
+ cred["tier"] = tier
2760
+
2761
+ return stats
2762
+
2763
+ def _get_baseline_confidence(self, model_stats: Dict) -> str:
2764
+ """
2765
+ Determine confidence level based on baseline age.
2766
+
2767
+ Args:
2768
+ model_stats: Model statistics dict with baseline_fetched_at
2769
+
2770
+ Returns:
2771
+ "high" | "medium" | "low"
2772
+ """
2773
+ baseline_fetched_at = model_stats.get("baseline_fetched_at")
2774
+ if not baseline_fetched_at:
2775
+ return "low"
2776
+
2777
+ age_seconds = time.time() - baseline_fetched_at
2778
+ if age_seconds < 300: # 5 minutes
2779
+ return "high"
2780
+ elif age_seconds < 1800: # 30 minutes
2781
+ return "medium"
2782
+ return "low"
2783
+
2784
+ async def reload_usage_from_disk(self) -> None:
2785
+ """
2786
+ Force reload usage data from disk.
2787
+
2788
+ Useful when wanting fresh stats without making external API calls.
2789
+ """
2790
+ await self.usage_manager.reload_from_disk()
2791
+
2792
+ async def force_refresh_quota(
2793
+ self,
2794
+ provider: Optional[str] = None,
2795
+ credential: Optional[str] = None,
2796
+ ) -> Dict[str, Any]:
2797
+ """
2798
+ Force refresh quota from external API.
2799
+
2800
+ For Antigravity, this fetches live quota data from the API.
2801
+ For other providers, this is a no-op (just reloads from disk).
2802
+
2803
+ Args:
2804
+ provider: If specified, only refresh this provider
2805
+ credential: If specified, only refresh this specific credential
2806
+
2807
+ Returns:
2808
+ Refresh result dict with success/failure info
2809
+ """
2810
+ result = {
2811
+ "action": "force_refresh",
2812
+ "scope": "credential"
2813
+ if credential
2814
+ else ("provider" if provider else "all"),
2815
+ "provider": provider,
2816
+ "credential": credential,
2817
+ "credentials_refreshed": 0,
2818
+ "success_count": 0,
2819
+ "failed_count": 0,
2820
+ "duration_ms": 0,
2821
+ "errors": [],
2822
+ }
2823
+
2824
+ start_time = time.time()
2825
+
2826
+ # Determine which providers to refresh
2827
+ if provider:
2828
+ providers_to_refresh = (
2829
+ [provider] if provider in self.all_credentials else []
2830
+ )
2831
+ else:
2832
+ providers_to_refresh = list(self.all_credentials.keys())
2833
+
2834
+ for prov in providers_to_refresh:
2835
+ provider_class = self._provider_plugins.get(prov)
2836
+ if not provider_class:
2837
+ continue
2838
+
2839
+ # Get or create provider instance
2840
+ if prov not in self._provider_instances:
2841
+ self._provider_instances[prov] = provider_class()
2842
+ provider_instance = self._provider_instances[prov]
2843
+
2844
+ # Check if provider supports quota refresh (like Antigravity)
2845
+ if hasattr(provider_instance, "fetch_initial_baselines"):
2846
+ # Get credentials to refresh
2847
+ if credential:
2848
+ # Find full path for this credential
2849
+ creds_to_refresh = []
2850
+ for cred_path in self.all_credentials.get(prov, []):
2851
+ if cred_path.endswith(credential) or cred_path == credential:
2852
+ creds_to_refresh.append(cred_path)
2853
+ break
2854
+ else:
2855
+ creds_to_refresh = self.all_credentials.get(prov, [])
2856
+
2857
+ if not creds_to_refresh:
2858
+ continue
2859
+
2860
+ try:
2861
+ # Fetch live quota from API for ALL specified credentials
2862
+ quota_results = await provider_instance.fetch_initial_baselines(
2863
+ creds_to_refresh
2864
+ )
2865
+
2866
+ # Store baselines in usage manager
2867
+ if hasattr(provider_instance, "_store_baselines_to_usage_manager"):
2868
+ stored = (
2869
+ await provider_instance._store_baselines_to_usage_manager(
2870
+ quota_results, self.usage_manager
2871
+ )
2872
+ )
2873
+ result["success_count"] += stored
2874
+
2875
+ result["credentials_refreshed"] += len(creds_to_refresh)
2876
+
2877
+ # Count failures
2878
+ for cred_path, data in quota_results.items():
2879
+ if data.get("status") != "success":
2880
+ result["failed_count"] += 1
2881
+ result["errors"].append(
2882
+ f"{Path(cred_path).name}: {data.get('error', 'Unknown error')}"
2883
+ )
2884
+
2885
+ except Exception as e:
2886
+ lib_logger.error(f"Failed to refresh quota for {prov}: {e}")
2887
+ result["errors"].append(f"{prov}: {str(e)}")
2888
+ result["failed_count"] += len(creds_to_refresh)
2889
+
2890
+ result["duration_ms"] = int((time.time() - start_time) * 1000)
2891
+ return result
src/rotator_library/providers/antigravity_provider.py CHANGED
@@ -103,7 +103,7 @@ DEFAULT_MAX_OUTPUT_TOKENS = 64000
103
  # Empty response retry configuration
104
  # When Antigravity returns an empty response (no content, no tool calls),
105
  # automatically retry up to this many attempts before giving up (minimum 1)
106
- EMPTY_RESPONSE_MAX_ATTEMPTS = max(1, _env_int("ANTIGRAVITY_EMPTY_RESPONSE_ATTEMPTS", 4))
107
  EMPTY_RESPONSE_RETRY_DELAY = _env_int("ANTIGRAVITY_EMPTY_RESPONSE_RETRY_DELAY", 2)
108
 
109
  # Model alias mappings (internal ↔ public)
 
103
  # Empty response retry configuration
104
  # When Antigravity returns an empty response (no content, no tool calls),
105
  # automatically retry up to this many attempts before giving up (minimum 1)
106
+ EMPTY_RESPONSE_MAX_ATTEMPTS = max(1, _env_int("ANTIGRAVITY_EMPTY_RESPONSE_ATTEMPTS", 6))
107
  EMPTY_RESPONSE_RETRY_DELAY = _env_int("ANTIGRAVITY_EMPTY_RESPONSE_RETRY_DELAY", 2)
108
 
109
  # Model alias mappings (internal ↔ public)
src/rotator_library/usage_manager.py CHANGED
@@ -1993,3 +1993,479 @@ class UsageManager:
1993
  """
1994
  # Disabled - see docstring above
1995
  pass
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1993
  """
1994
  # Disabled - see docstring above
1995
  pass
1996
+
1997
+ async def get_stats_for_endpoint(
1998
+ self,
1999
+ provider_filter: Optional[str] = None,
2000
+ include_global: bool = True,
2001
+ ) -> Dict[str, Any]:
2002
+ """
2003
+ Get usage stats formatted for the /v1/quota-stats endpoint.
2004
+
2005
+ Aggregates data from key_usage.json grouped by provider.
2006
+ Includes both current period stats and global (lifetime) stats.
2007
+
2008
+ Args:
2009
+ provider_filter: If provided, only return stats for this provider
2010
+ include_global: If True, include global/lifetime stats alongside current
2011
+
2012
+ Returns:
2013
+ {
2014
+ "providers": {
2015
+ "provider_name": {
2016
+ "credential_count": int,
2017
+ "active_count": int,
2018
+ "on_cooldown_count": int,
2019
+ "total_requests": int,
2020
+ "tokens": {
2021
+ "input_cached": int,
2022
+ "input_uncached": int,
2023
+ "input_cache_pct": float,
2024
+ "output": int
2025
+ },
2026
+ "approx_cost": float | None,
2027
+ "credentials": [...],
2028
+ "global": {...} # If include_global is True
2029
+ }
2030
+ },
2031
+ "summary": {...},
2032
+ "global_summary": {...}, # If include_global is True
2033
+ "timestamp": float
2034
+ }
2035
+ """
2036
+ await self._lazy_init()
2037
+
2038
+ now_ts = time.time()
2039
+ providers: Dict[str, Dict[str, Any]] = {}
2040
+ # Track global stats separately
2041
+ global_providers: Dict[str, Dict[str, Any]] = {}
2042
+
2043
+ async with self._data_lock:
2044
+ if not self._usage_data:
2045
+ return {
2046
+ "providers": {},
2047
+ "summary": {
2048
+ "total_providers": 0,
2049
+ "total_credentials": 0,
2050
+ "active_credentials": 0,
2051
+ "exhausted_credentials": 0,
2052
+ "total_requests": 0,
2053
+ "tokens": {
2054
+ "input_cached": 0,
2055
+ "input_uncached": 0,
2056
+ "input_cache_pct": 0,
2057
+ "output": 0,
2058
+ },
2059
+ "approx_total_cost": 0.0,
2060
+ },
2061
+ "global_summary": {
2062
+ "total_providers": 0,
2063
+ "total_credentials": 0,
2064
+ "total_requests": 0,
2065
+ "tokens": {
2066
+ "input_cached": 0,
2067
+ "input_uncached": 0,
2068
+ "input_cache_pct": 0,
2069
+ "output": 0,
2070
+ },
2071
+ "approx_total_cost": 0.0,
2072
+ },
2073
+ "data_source": "cache",
2074
+ "timestamp": now_ts,
2075
+ }
2076
+
2077
+ for credential, cred_data in self._usage_data.items():
2078
+ # Extract provider from credential path
2079
+ provider = self._get_provider_from_credential(credential)
2080
+ if not provider:
2081
+ continue
2082
+
2083
+ # Apply filter if specified
2084
+ if provider_filter and provider != provider_filter:
2085
+ continue
2086
+
2087
+ # Initialize provider entry
2088
+ if provider not in providers:
2089
+ providers[provider] = {
2090
+ "credential_count": 0,
2091
+ "active_count": 0,
2092
+ "on_cooldown_count": 0,
2093
+ "exhausted_count": 0,
2094
+ "total_requests": 0,
2095
+ "tokens": {
2096
+ "input_cached": 0,
2097
+ "input_uncached": 0,
2098
+ "input_cache_pct": 0,
2099
+ "output": 0,
2100
+ },
2101
+ "approx_cost": 0.0,
2102
+ "credentials": [],
2103
+ }
2104
+ global_providers[provider] = {
2105
+ "total_requests": 0,
2106
+ "tokens": {
2107
+ "input_cached": 0,
2108
+ "input_uncached": 0,
2109
+ "input_cache_pct": 0,
2110
+ "output": 0,
2111
+ },
2112
+ "approx_cost": 0.0,
2113
+ }
2114
+
2115
+ prov_stats = providers[provider]
2116
+ prov_stats["credential_count"] += 1
2117
+
2118
+ # Determine credential status and cooldowns
2119
+ key_cooldown = cred_data.get("key_cooldown_until", 0) or 0
2120
+ model_cooldowns = cred_data.get("model_cooldowns", {})
2121
+
2122
+ # Build active cooldowns with remaining time
2123
+ active_cooldowns = {}
2124
+ for model, cooldown_ts in model_cooldowns.items():
2125
+ if cooldown_ts > now_ts:
2126
+ remaining_seconds = int(cooldown_ts - now_ts)
2127
+ active_cooldowns[model] = {
2128
+ "until_ts": cooldown_ts,
2129
+ "remaining_seconds": remaining_seconds,
2130
+ }
2131
+
2132
+ key_cooldown_remaining = None
2133
+ if key_cooldown > now_ts:
2134
+ key_cooldown_remaining = int(key_cooldown - now_ts)
2135
+
2136
+ has_active_cooldown = key_cooldown > now_ts or len(active_cooldowns) > 0
2137
+
2138
+ # Check if exhausted (all quota groups exhausted for Antigravity)
2139
+ is_exhausted = False
2140
+ models_data = cred_data.get("models", {})
2141
+ if models_data:
2142
+ # Check if any model has remaining quota
2143
+ all_exhausted = True
2144
+ for model_stats in models_data.values():
2145
+ if isinstance(model_stats, dict):
2146
+ baseline = model_stats.get("baseline_remaining_fraction")
2147
+ if baseline is None or baseline > 0:
2148
+ all_exhausted = False
2149
+ break
2150
+ if all_exhausted and len(models_data) > 0:
2151
+ is_exhausted = True
2152
+
2153
+ if is_exhausted:
2154
+ prov_stats["exhausted_count"] += 1
2155
+ status = "exhausted"
2156
+ elif has_active_cooldown:
2157
+ prov_stats["on_cooldown_count"] += 1
2158
+ status = "cooldown"
2159
+ else:
2160
+ prov_stats["active_count"] += 1
2161
+ status = "active"
2162
+
2163
+ # Aggregate token stats (current period)
2164
+ cred_tokens = {
2165
+ "input_cached": 0,
2166
+ "input_uncached": 0,
2167
+ "output": 0,
2168
+ }
2169
+ cred_requests = 0
2170
+ cred_cost = 0.0
2171
+
2172
+ # Aggregate global token stats
2173
+ cred_global_tokens = {
2174
+ "input_cached": 0,
2175
+ "input_uncached": 0,
2176
+ "output": 0,
2177
+ }
2178
+ cred_global_requests = 0
2179
+ cred_global_cost = 0.0
2180
+
2181
+ # Handle per-model structure (current period)
2182
+ if models_data:
2183
+ for model_name, model_stats in models_data.items():
2184
+ if not isinstance(model_stats, dict):
2185
+ continue
2186
+ # Prefer request_count if available and non-zero, else fall back to success+failure
2187
+ req_count = model_stats.get("request_count", 0)
2188
+ if req_count > 0:
2189
+ cred_requests += req_count
2190
+ else:
2191
+ cred_requests += model_stats.get("success_count", 0)
2192
+ cred_requests += model_stats.get("failure_count", 0)
2193
+ # Token stats - track cached separately
2194
+ cred_tokens["input_cached"] += model_stats.get(
2195
+ "prompt_tokens_cached", 0
2196
+ )
2197
+ cred_tokens["input_uncached"] += model_stats.get(
2198
+ "prompt_tokens", 0
2199
+ )
2200
+ cred_tokens["output"] += model_stats.get("completion_tokens", 0)
2201
+ cred_cost += model_stats.get("approx_cost", 0.0)
2202
+
2203
+ # Handle legacy daily structure
2204
+ daily_data = cred_data.get("daily", {})
2205
+ daily_models = daily_data.get("models", {})
2206
+ for model_name, model_stats in daily_models.items():
2207
+ if not isinstance(model_stats, dict):
2208
+ continue
2209
+ cred_requests += model_stats.get("success_count", 0)
2210
+ cred_tokens["input_cached"] += model_stats.get(
2211
+ "prompt_tokens_cached", 0
2212
+ )
2213
+ cred_tokens["input_uncached"] += model_stats.get("prompt_tokens", 0)
2214
+ cred_tokens["output"] += model_stats.get("completion_tokens", 0)
2215
+ cred_cost += model_stats.get("approx_cost", 0.0)
2216
+
2217
+ # Handle global stats
2218
+ global_data = cred_data.get("global", {})
2219
+ global_models = global_data.get("models", {})
2220
+ for model_name, model_stats in global_models.items():
2221
+ if not isinstance(model_stats, dict):
2222
+ continue
2223
+ cred_global_requests += model_stats.get("success_count", 0)
2224
+ cred_global_tokens["input_cached"] += model_stats.get(
2225
+ "prompt_tokens_cached", 0
2226
+ )
2227
+ cred_global_tokens["input_uncached"] += model_stats.get(
2228
+ "prompt_tokens", 0
2229
+ )
2230
+ cred_global_tokens["output"] += model_stats.get(
2231
+ "completion_tokens", 0
2232
+ )
2233
+ cred_global_cost += model_stats.get("approx_cost", 0.0)
2234
+
2235
+ # Add current period stats to global totals
2236
+ cred_global_requests += cred_requests
2237
+ cred_global_tokens["input_cached"] += cred_tokens["input_cached"]
2238
+ cred_global_tokens["input_uncached"] += cred_tokens["input_uncached"]
2239
+ cred_global_tokens["output"] += cred_tokens["output"]
2240
+ cred_global_cost += cred_cost
2241
+
2242
+ # Build credential entry
2243
+ # Mask credential identifier for display
2244
+ if credential.startswith("env://"):
2245
+ identifier = credential
2246
+ else:
2247
+ identifier = Path(credential).name
2248
+
2249
+ cred_entry = {
2250
+ "identifier": identifier,
2251
+ "full_path": credential,
2252
+ "status": status,
2253
+ "last_used_ts": cred_data.get("last_used_ts"),
2254
+ "requests": cred_requests,
2255
+ "tokens": cred_tokens,
2256
+ "approx_cost": cred_cost if cred_cost > 0 else None,
2257
+ }
2258
+
2259
+ # Add cooldown info
2260
+ if key_cooldown_remaining is not None:
2261
+ cred_entry["key_cooldown_remaining"] = key_cooldown_remaining
2262
+ if active_cooldowns:
2263
+ cred_entry["model_cooldowns"] = active_cooldowns
2264
+
2265
+ # Add global stats for this credential
2266
+ if include_global:
2267
+ # Calculate global cache percentage
2268
+ global_total_input = (
2269
+ cred_global_tokens["input_cached"]
2270
+ + cred_global_tokens["input_uncached"]
2271
+ )
2272
+ global_cache_pct = (
2273
+ round(
2274
+ cred_global_tokens["input_cached"]
2275
+ / global_total_input
2276
+ * 100,
2277
+ 1,
2278
+ )
2279
+ if global_total_input > 0
2280
+ else 0
2281
+ )
2282
+
2283
+ cred_entry["global"] = {
2284
+ "requests": cred_global_requests,
2285
+ "tokens": {
2286
+ "input_cached": cred_global_tokens["input_cached"],
2287
+ "input_uncached": cred_global_tokens["input_uncached"],
2288
+ "input_cache_pct": global_cache_pct,
2289
+ "output": cred_global_tokens["output"],
2290
+ },
2291
+ "approx_cost": cred_global_cost
2292
+ if cred_global_cost > 0
2293
+ else None,
2294
+ }
2295
+
2296
+ # Add model-specific data for providers with per-model tracking
2297
+ if models_data:
2298
+ cred_entry["models"] = {}
2299
+ for model_name, model_stats in models_data.items():
2300
+ if not isinstance(model_stats, dict):
2301
+ continue
2302
+ cred_entry["models"][model_name] = {
2303
+ "requests": model_stats.get("success_count", 0)
2304
+ + model_stats.get("failure_count", 0),
2305
+ "request_count": model_stats.get("request_count", 0),
2306
+ "success_count": model_stats.get("success_count", 0),
2307
+ "failure_count": model_stats.get("failure_count", 0),
2308
+ "prompt_tokens": model_stats.get("prompt_tokens", 0),
2309
+ "prompt_tokens_cached": model_stats.get(
2310
+ "prompt_tokens_cached", 0
2311
+ ),
2312
+ "completion_tokens": model_stats.get(
2313
+ "completion_tokens", 0
2314
+ ),
2315
+ "approx_cost": model_stats.get("approx_cost", 0.0),
2316
+ "window_start_ts": model_stats.get("window_start_ts"),
2317
+ "quota_reset_ts": model_stats.get("quota_reset_ts"),
2318
+ # Quota baseline fields (Antigravity-specific)
2319
+ "baseline_remaining_fraction": model_stats.get(
2320
+ "baseline_remaining_fraction"
2321
+ ),
2322
+ "baseline_fetched_at": model_stats.get(
2323
+ "baseline_fetched_at"
2324
+ ),
2325
+ "quota_max_requests": model_stats.get("quota_max_requests"),
2326
+ "quota_display": model_stats.get("quota_display"),
2327
+ }
2328
+
2329
+ prov_stats["credentials"].append(cred_entry)
2330
+
2331
+ # Aggregate to provider totals (current period)
2332
+ prov_stats["total_requests"] += cred_requests
2333
+ prov_stats["tokens"]["input_cached"] += cred_tokens["input_cached"]
2334
+ prov_stats["tokens"]["input_uncached"] += cred_tokens["input_uncached"]
2335
+ prov_stats["tokens"]["output"] += cred_tokens["output"]
2336
+ if cred_cost > 0:
2337
+ prov_stats["approx_cost"] += cred_cost
2338
+
2339
+ # Aggregate to global provider totals
2340
+ global_providers[provider]["total_requests"] += cred_global_requests
2341
+ global_providers[provider]["tokens"]["input_cached"] += (
2342
+ cred_global_tokens["input_cached"]
2343
+ )
2344
+ global_providers[provider]["tokens"]["input_uncached"] += (
2345
+ cred_global_tokens["input_uncached"]
2346
+ )
2347
+ global_providers[provider]["tokens"]["output"] += cred_global_tokens[
2348
+ "output"
2349
+ ]
2350
+ global_providers[provider]["approx_cost"] += cred_global_cost
2351
+
2352
+ # Calculate cache percentages for each provider
2353
+ for provider, prov_stats in providers.items():
2354
+ total_input = (
2355
+ prov_stats["tokens"]["input_cached"]
2356
+ + prov_stats["tokens"]["input_uncached"]
2357
+ )
2358
+ if total_input > 0:
2359
+ prov_stats["tokens"]["input_cache_pct"] = round(
2360
+ prov_stats["tokens"]["input_cached"] / total_input * 100, 1
2361
+ )
2362
+ # Set cost to None if 0
2363
+ if prov_stats["approx_cost"] == 0:
2364
+ prov_stats["approx_cost"] = None
2365
+
2366
+ # Calculate global cache percentages
2367
+ if include_global and provider in global_providers:
2368
+ gp = global_providers[provider]
2369
+ global_total = (
2370
+ gp["tokens"]["input_cached"] + gp["tokens"]["input_uncached"]
2371
+ )
2372
+ if global_total > 0:
2373
+ gp["tokens"]["input_cache_pct"] = round(
2374
+ gp["tokens"]["input_cached"] / global_total * 100, 1
2375
+ )
2376
+ if gp["approx_cost"] == 0:
2377
+ gp["approx_cost"] = None
2378
+ prov_stats["global"] = gp
2379
+
2380
+ # Build summary (current period)
2381
+ total_creds = sum(p["credential_count"] for p in providers.values())
2382
+ active_creds = sum(p["active_count"] for p in providers.values())
2383
+ exhausted_creds = sum(p["exhausted_count"] for p in providers.values())
2384
+ total_requests = sum(p["total_requests"] for p in providers.values())
2385
+ total_input_cached = sum(
2386
+ p["tokens"]["input_cached"] for p in providers.values()
2387
+ )
2388
+ total_input_uncached = sum(
2389
+ p["tokens"]["input_uncached"] for p in providers.values()
2390
+ )
2391
+ total_output = sum(p["tokens"]["output"] for p in providers.values())
2392
+ total_cost = sum(p["approx_cost"] or 0 for p in providers.values())
2393
+
2394
+ total_input = total_input_cached + total_input_uncached
2395
+ input_cache_pct = (
2396
+ round(total_input_cached / total_input * 100, 1) if total_input > 0 else 0
2397
+ )
2398
+
2399
+ result = {
2400
+ "providers": providers,
2401
+ "summary": {
2402
+ "total_providers": len(providers),
2403
+ "total_credentials": total_creds,
2404
+ "active_credentials": active_creds,
2405
+ "exhausted_credentials": exhausted_creds,
2406
+ "total_requests": total_requests,
2407
+ "tokens": {
2408
+ "input_cached": total_input_cached,
2409
+ "input_uncached": total_input_uncached,
2410
+ "input_cache_pct": input_cache_pct,
2411
+ "output": total_output,
2412
+ },
2413
+ "approx_total_cost": total_cost if total_cost > 0 else None,
2414
+ },
2415
+ "data_source": "cache",
2416
+ "timestamp": now_ts,
2417
+ }
2418
+
2419
+ # Build global summary
2420
+ if include_global:
2421
+ global_total_requests = sum(
2422
+ gp["total_requests"] for gp in global_providers.values()
2423
+ )
2424
+ global_total_input_cached = sum(
2425
+ gp["tokens"]["input_cached"] for gp in global_providers.values()
2426
+ )
2427
+ global_total_input_uncached = sum(
2428
+ gp["tokens"]["input_uncached"] for gp in global_providers.values()
2429
+ )
2430
+ global_total_output = sum(
2431
+ gp["tokens"]["output"] for gp in global_providers.values()
2432
+ )
2433
+ global_total_cost = sum(
2434
+ gp["approx_cost"] or 0 for gp in global_providers.values()
2435
+ )
2436
+
2437
+ global_total_input = global_total_input_cached + global_total_input_uncached
2438
+ global_input_cache_pct = (
2439
+ round(global_total_input_cached / global_total_input * 100, 1)
2440
+ if global_total_input > 0
2441
+ else 0
2442
+ )
2443
+
2444
+ result["global_summary"] = {
2445
+ "total_providers": len(global_providers),
2446
+ "total_credentials": total_creds,
2447
+ "total_requests": global_total_requests,
2448
+ "tokens": {
2449
+ "input_cached": global_total_input_cached,
2450
+ "input_uncached": global_total_input_uncached,
2451
+ "input_cache_pct": global_input_cache_pct,
2452
+ "output": global_total_output,
2453
+ },
2454
+ "approx_total_cost": global_total_cost
2455
+ if global_total_cost > 0
2456
+ else None,
2457
+ }
2458
+
2459
+ return result
2460
+
2461
+ async def reload_from_disk(self) -> None:
2462
+ """
2463
+ Force reload usage data from disk.
2464
+
2465
+ Useful when another process may have updated the file.
2466
+ """
2467
+ async with self._init_lock:
2468
+ self._initialized.clear()
2469
+ await self._load_usage()
2470
+ await self._reset_daily_stats_if_needed()
2471
+ self._initialized.set()