Coverage for src/kimi_code_usage/main.py: 100%
657 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-20 20:05 +0800
1import argparse
2import asyncio
3import json
4import os
5import select as _select_module
6import sys
8from dotenv import load_dotenv
9from rich.console import Console, Group
10from rich.live import Live
11from rich.panel import Panel
12from rich.text import Text
14from kimi_code_usage.config import AppConfig, ConfigResolver, save_theme
15from kimi_code_usage.providers import DailyUsage, ProviderUsage, dispatch_all
17# --- i18n ---
18LANG = os.getenv("LANG", "en")
19IS_ZH = "zh" in LANG.lower()
21L_EN = {
22 "title": "AI Quota Monitor",
23 "weekly_limit": "Weekly Usage",
24 "limit_fallback": "Limit",
25 "remaining": "remaining",
26 "countdown": "Countdown",
27 "reset": "Reset",
28 "no_data": "No usage data found or no providers configured.",
29 "error_api": "API Error",
30}
32L_ZH = {
33 "title": "AI 用量配额监控",
34 "weekly_limit": "周用量限额",
35 "limit_fallback": "限额",
36 "remaining": "剩余",
37 "countdown": "重置倒计时",
38 "reset": "重置时间",
39 "no_data": "未找到用量数据，或未配置任何服务商。",
40 "error_api": "API 错误",
41}
43L = L_ZH if IS_ZH else L_EN
45def _get_visual_width(s: str) -> int:
46 import unicodedata
47 width = 0
48 for char in s:
49 if unicodedata.east_asian_width(char) in ("W", "F", "A"):
50 width += 2
51 else:
52 width += 1
53 return width
55def _get_localized_label(label: str, lang_zh: bool = IS_ZH) -> str:
56 _L = L_ZH if lang_zh else L_EN
58 # 1. Dynamic translations first
59 if lang_zh:
60 translations = {
61 "Credits": "额度",
62 "Key Name": "密钥名称",
63 "Rate Limit": "速率限制",
64 "Usage": "周期已用",
65 "周期已用": "周期已用",
66 "Free Tier": "免费额度",
67 "Limit Reset": "限额重置周期",
68 "Expires At": "过期时间",
69 "Is Provisioning": "配给密钥",
70 "Management Key": "管理密钥",
71 }
72 if label in translations:
73 return translations[label]
74 else:
75 translations = {
76 "额度": "Credits",
77 "密钥名称": "Key Name",
78 "速率限制": "Rate Limit",
79 "周期已用": "Usage",
80 "Usage": "Usage",
81 "免费额度": "Free Tier",
82 "限额重置周期": "Limit Reset",
83 "过期时间": "Expires At",
84 "配给密钥": "Is Provisioning",
85 "管理密钥": "Management Key",
86 }
87 if label in translations:
88 return translations[label]
90 # 2. Standard rule replacements
91 if label == "Weekly Usage":
92 return _L["weekly_limit"]
93 if "h Limit" in label:
94 h = label.split("h")[0]
95 return f"{h}小时限额" if lang_zh else label
96 if "d Limit" in label:
97 d = label.split("d")[0]
98 return f"{d}天限额" if lang_zh else label
99 if "mo Limit" in label:
100 mo = label.split("mo")[0]
101 return f"{mo}个月限额" if lang_zh else label
102 if "m Limit" in label:
103 m = label.split("m")[0]
104 return f"{m}分钟限额" if lang_zh else label
105 if "Limit" in label:
106 return label.replace("Limit", _L["limit_fallback"])
107 return label
110def _get_localized_text_value(text_val: str, lang_zh: bool) -> str:
111 if not text_val:
112 return text_val
113 # OpenRouter period usage pattern
114 if "Daily: $" in text_val or "今日: $" in text_val:
115 import re
116 floats = re.findall(r"\d+\.\d+", text_val)
117 if len(floats) == 3:
118 u_daily, u_weekly, u_monthly = float(floats[0]), float(floats[1]), float(floats[2])
119 if lang_zh:
120 return f"今日: ${u_daily:.4f} | 本周: ${u_weekly:.4f} | 本月: ${u_monthly:.4f}"
121 else:
122 return f"Daily: ${u_daily:.4f} | Weekly: ${u_weekly:.4f} | Monthly: ${u_monthly:.4f}"
124 # Rate limit localization (e.g., Unlimited/10s <-> 无限制/10秒, 20 req/1s <-> 20次/1秒)
125 if lang_zh:
126 if text_val.startswith("Unlimited/"):
127 interval = text_val.split("/", 1)[1]
128 interval_zh = interval.replace("s", "秒").replace("m", "分钟").replace("h", "小时")
129 return f"无限制/{interval_zh}"
130 elif " req/" in text_val:
131 reqs, interval = text_val.split(" req/", 1)
132 interval_zh = interval.replace("s", "秒").replace("m", "分钟").replace("h", "小时")
133 return f"{reqs}次/{interval_zh}"
134 else:
135 if text_val.startswith("无限制/"):
136 interval = text_val.split("/", 1)[1]
137 interval_en = interval.replace("秒", "s").replace("分钟", "m").replace("小时", "h")
138 return f"Unlimited/{interval_en}"
139 elif "次/" in text_val:
140 reqs, interval = text_val.split("次/", 1)
141 interval_en = interval.replace("秒", "s").replace("分钟", "m").replace("小时", "h")
142 return f"{reqs} req/{interval_en}"
144 # Yes/No localization
145 if lang_zh:
146 if text_val == "Yes":
147 return "是"
148 if text_val == "No":
149 return "否"
150 else:
151 if text_val == "是":
152 return "Yes"
153 if text_val == "否":
154 return "No"
156 return text_val
158# (typing already imported at top)
160THEME_MAP = {
161 # ── Classic Blue on Dark ──
162 "blue-dark": {
163 "title": "bold dodger_blue2",
164 "label": "cornflower_blue",
165 "meta": "grey62",
166 "ok": "medium_spring_green",
167 "warning": "gold1",
168 "danger": "indian_red1",
169 },
170 # ── Classic Blue on Light ──
171 "blue-light": {
172 "title": "bold blue",
173 "label": "dark_blue",
174 "meta": "grey35",
175 "ok": "dark_green",
176 "warning": "dark_orange",
177 "danger": "red3",
178 },
179 # ── Royal Sky Blue on Dark ──
180 "sky-dark": {
181 "title": "bold royal_blue1",
182 "label": "sky_blue1",
183 "meta": "grey58",
184 "ok": "spring_green2",
185 "warning": "orange1",
186 "danger": "light_coral",
187 },
188 # ── Salmon Saffron on Dark ──
189 "salmon-dark": {
190 "title": "bold light_salmon3",
191 "label": "sandy_brown",
192 "meta": "grey46",
193 "ok": "dark_sea_green2",
194 "warning": "gold3",
195 "danger": "indian_red",
196 },
197 # ── Turquoise Aquamarine on Dark ──
198 "turquoise-dark": {
199 "title": "bold turquoise2",
200 "label": "medium_aquamarine",
201 "meta": "grey53",
202 "ok": "aquamarine1",
203 "warning": "khaki1",
204 "danger": "hot_pink",
205 },
206 # ── Rose Pink on Light ──
207 "pink-light": {
208 "title": "bold deep_pink3",
209 "label": "hot_pink3",
210 "meta": "grey37",
211 "ok": "dark_cyan",
212 "warning": "dark_goldenrod",
213 "danger": "red3",
214 },
215 # ── Deep Purple on Dark ──
216 "violet-dark": {
217 "title": "bold medium_purple2",
218 "label": "medium_purple1",
219 "meta": "grey54",
220 "ok": "spring_green3",
221 "warning": "orange1",
222 "danger": "plum1",
223 },
224 # ── Warm Saffron Orange on Dark ──
225 "amber-dark": {
226 "title": "bold dark_orange",
227 "label": "gold1",
228 "meta": "grey54",
229 "ok": "green_yellow",
230 "warning": "gold3",
231 "danger": "red1",
232 },
233 # ── Mint Teal on Dark ──
234 "mint-dark": {
235 "title": "bold dark_cyan",
236 "label": "light_sea_green",
237 "meta": "grey54",
238 "ok": "medium_spring_green",
239 "warning": "khaki1",
240 "danger": "indian_red1",
241 },
242 # ── Monochrome ──
243 "monochrome": {
244 "title": "bold white",
245 "label": "white",
246 "meta": "grey46",
247 "ok": "white",
248 "warning": "grey74",
249 "danger": "bold reverse",
250 },
251 # ── Red-Green Blind Friendly ──
252 "blind-deuteranopia": {
253 "title": "bold dodger_blue1",
254 "label": "cornflower_blue",
255 "meta": "grey62",
256 "ok": "dodger_blue1",
257 "warning": "gold1",
258 "danger": "dark_orange",
259 },
260 # ── Blue-Yellow Blind Friendly ──
261 "blind-tritanopia": {
262 "title": "bold red1",
263 "label": "light_coral",
264 "meta": "grey62",
265 "ok": "chartreuse3",
266 "warning": "pale_green3",
267 "danger": "bold red1",
268 },
269}
272def _handle_key(ch: str, idx: int, n: int) -> tuple[int, bool, bool, int | None, bool, bool, bool]:
273 """Map a keypress to a TUI action.
275 Returns:
276 (new_idx, should_quit, should_refresh, toggle_provider_num, lang_toggle, metric_toggle, days_toggle)
277 """
278 if ch in ('q', 'Q', '\x03', '\x04'): # q Ctrl-C Ctrl-D
279 return idx, True, False, None, False, False, False
280 if ch in (']', 'n', '\t', '\x1b[C'): # ] n Tab →
281 return (idx + 1) % n, False, False, None, False, False, False
282 if ch in ('[', 'p', '\x1b[D'): # [ p ←
283 return (idx - 1) % n, False, False, None, False, False, False
284 if ch in ('r', 'R'): # r → refresh data
285 return idx, False, True, None, False, False, False
286 if ch in ('1', '2', '3', '4'): # 1-4 → toggle provider panel
287 return idx, False, False, int(ch) - 1, False, False, False
288 if ch in ('l', 'L'): # l → toggle language zh/en
289 return idx, False, False, None, True, False, False
290 if ch in ('m', 'M'): # m → toggle OpenRouter metric
291 return idx, False, False, None, False, True, False
292 if ch in ('d', 'D'): # d → toggle days window
293 return idx, False, False, None, False, False, True
294 return idx, False, False, None, False, False, False
297def _format_tokens(n: float) -> str:
298 if n >= 1_000_000:
299 return f"{n/1_000_000:.1f}M"
300 if n >= 1_000:
301 return f"{n/1_000:.1f}K"
302 return f"{n:.0f}"
305def _model_short_name(model: str) -> str:
306 if "/" in model:
307 return model.rsplit("/", 1)[-1]
308 return model
311# --- OpenRouter metric helpers ---
312OR_METRIC_SPEND = "spend"
313OR_METRIC_REQUESTS = "requests"
314OR_METRIC_TOKENS = "tokens"
315VALID_OR_METRICS = (OR_METRIC_SPEND, OR_METRIC_REQUESTS, OR_METRIC_TOKENS)
317# --- OpenRouter days window helpers ---
318DAYS_WINDOWS = (7, 14, 30, 60, 90)
321def _parse_days_window(value) -> int:
322 try:
323 value = int(value)
324 except (TypeError, ValueError):
325 return 30
326 return value if value in DAYS_WINDOWS else 30
329def _next_days_window(days_window: int) -> int:
330 try:
331 idx = DAYS_WINDOWS.index(days_window)
332 except ValueError:
333 return 30
334 return DAYS_WINDOWS[(idx + 1) % len(DAYS_WINDOWS)]
337def _parse_or_metric(value: str | None) -> str:
338 if value in VALID_OR_METRICS:
339 return value
340 return OR_METRIC_REQUESTS
343def _next_or_metric(metric: str) -> str:
344 order = (OR_METRIC_SPEND, OR_METRIC_REQUESTS, OR_METRIC_TOKENS)
345 try:
346 idx = order.index(metric)
347 except ValueError:
348 return OR_METRIC_REQUESTS
349 return order[(idx + 1) % len(order)]
352def _metric_label(metric: str, lang_zh: bool) -> str:
353 labels = {
354 OR_METRIC_SPEND: ("支出", "Spend"),
355 OR_METRIC_REQUESTS: ("请求", "Requests"),
356 OR_METRIC_TOKENS: ("Tokens", "Tokens"),
357 }
358 return labels.get(metric, labels[OR_METRIC_SPEND])[0 if lang_zh else 1]
361def _metric_value_model(mu, metric: str) -> float:
362 if metric == OR_METRIC_REQUESTS:
363 return mu.requests
364 if metric == OR_METRIC_TOKENS:
365 return mu.prompt_tokens + mu.completion_tokens + mu.reasoning_tokens
366 return mu.spend
369def _metric_value_day(day, metric: str) -> float:
370 if metric == OR_METRIC_SPEND:
371 return day.total
372 return sum(_metric_value_model(m, metric) for m in day.models)
375def _format_metric_value(value: float, metric: str) -> str:
376 if metric == OR_METRIC_SPEND:
377 return f"${value:.2f}"
378 if metric == OR_METRIC_TOKENS:
379 return _format_tokens(value)
380 return f"{value:,.0f}"
383def _render_activity_totals(totals, lang_zh: bool, theme: dict) -> Text:
384 text = Text()
385 title = "活动" if lang_zh else "Activity"
386 text.append(f" {title}\n", style=theme.get("label", "cornflower_blue"))
387 parts = []
388 if totals.requests:
389 req_label = "请求" if lang_zh else "Req"
390 parts.append(f"{req_label}: {totals.requests:,.0f}")
391 if totals.prompt_tokens or totals.completion_tokens:
392 in_label = "输入" if lang_zh else "In"
393 out_label = "输出" if lang_zh else "Out"
394 tok = f"{in_label}: {_format_tokens(totals.prompt_tokens)} / {out_label}: {_format_tokens(totals.completion_tokens)}"
395 if totals.reasoning_tokens:
396 reason_label = "推理" if lang_zh else "reason"
397 tok += f" (+ {_format_tokens(totals.reasoning_tokens)} {reason_label})"
398 parts.append(tok)
399 if totals.spend:
400 spend_label = "支出" if lang_zh else "Spend"
401 parts.append(f"{spend_label}: ${totals.spend:.2f}")
402 text.append(" " + " | ".join(parts), style=theme.get("meta", "grey62"))
403 return text
406def _render_daily_chart(daily_activity, lang_zh: bool, theme: dict, metric: str = OR_METRIC_REQUESTS, days_window: int = 30) -> Text:
407 if not daily_activity:
408 return Text()
409 metric = _parse_or_metric(metric)
410 days_window = _parse_days_window(days_window)
412 from datetime import datetime, timedelta
414 # Build a contiguous date range of `days_window` days ending at the latest activity date.
415 sorted_days = sorted(daily_activity, key=lambda d: d.date)
416 latest_date = datetime.strptime(sorted_days[-1].date[:10], "%Y-%m-%d")
417 start_date = latest_date - timedelta(days=days_window - 1)
419 day_lookup: dict[str, DailyUsage] = {}
420 for d in sorted_days:
421 key = d.date[:10] if len(d.date) >= 10 else d.date
422 if key in day_lookup:
423 existing = day_lookup[key]
424 merged_models = existing.models + d.models
425 day_lookup[key] = DailyUsage(date=key, models=merged_models, total=existing.total + d.total)
426 else:
427 day_lookup[key] = DailyUsage(date=key, models=list(d.models), total=d.total)
429 days: list[DailyUsage] = []
430 for i in range(days_window):
431 date = (start_date + timedelta(days=i)).strftime("%Y-%m-%d")
432 days.append(day_lookup.get(date, DailyUsage(date=date, models=[], total=0.0)))
434 # Chart dimensions inspired by tokentop: clamp column width to [1, 4] and
435 # cap the total rendered width so the chart fits a typical terminal panel.
436 height = 9
437 top_models_limit = 6
438 max_chart_width = 70
439 col_width = 1
440 if len(days) > 0:
441 col_width = max_chart_width // len(days)
442 if col_width > 4:
443 col_width = 4
444 if col_width < 1:
445 col_width = 1
446 # If the full window would exceed the max width, trim to the most recent days.
447 if len(days) * col_width > max_chart_width:
448 keep = max_chart_width // col_width
449 days = days[-keep:]
451 max_total = max(_metric_value_day(d, metric) for d in days) or 1.0
453 # Find top models across all days for consistent coloring
454 model_totals_all: dict[str, float] = {}
455 for day in days:
456 for model in day.models:
457 model_totals_all[model.model] = model_totals_all.get(model.model, 0.0) + _metric_value_model(model, metric)
458 top_models = _top_n_models(model_totals_all, top_models_limit)
459 top_model_set = set(top_models)
460 others_color_idx = len(top_models) % len(_model_colors(theme))
462 # Pre-compute each column as color indices (bottom -> top)
463 columns: list[list[int | None]] = []
464 for day in days:
465 value_by_model: dict[str, float] = {}
466 others_value = 0.0
467 for model in day.models:
468 v = _metric_value_model(model, metric)
469 if model.model in top_model_set:
470 value_by_model[model.model] = value_by_model.get(model.model, 0.0) + v
471 else:
472 others_value += v
474 segments: list[tuple[int, float]] = []
475 for i, name in enumerate(top_models):
476 if v := value_by_model.get(name, 0.0):
477 segments.append((i, v))
478 if others_value > 0:
479 segments.append((others_color_idx, others_value))
481 day_total = sum(v for _, v in segments)
482 column: list[int | None] = [None] * height
483 if day_total > 0 and max_total > 0:
484 total_cells = max(1, min(height, int(round(day_total / max_total * height))))
485 exacts = [seg[1] / day_total * total_cells for seg in segments]
486 counts = [int(x) for x in exacts]
487 remainders = [exacts[i] - counts[i] for i in range(len(segments))]
488 allocated = sum(counts)
489 while allocated < total_cells:
490 best = max(range(len(remainders)), key=lambda i: remainders[i])
491 counts[best] += 1
492 remainders[best] = -1
493 allocated += 1
495 cell_idx = 0
496 for seg_idx, (cidx, _) in enumerate(segments):
497 for _ in range(counts[seg_idx]):
498 if cell_idx < height:
499 column[cell_idx] = cidx
500 cell_idx += 1
501 columns.append(column)
503 text = Text()
504 title_label = _metric_label(metric, lang_zh)
505 title = f"每日{title_label}" if lang_zh else f"Daily {title_label.title()}"
506 text.append(f" {title}\n\n", style=theme.get("label", "cornflower_blue"))
508 # Y-axis gutter
509 y_labels = [max_total, max_total / 2, max_total / float(height)]
510 y_label_strs = [_format_metric_value(v, metric) for v in y_labels]
511 gutter_width = max(len(s) for s in y_label_strs) + 1
512 gutter_fmt = f"{{:>{gutter_width}s}}"
513 empty_gutter = " " * gutter_width
514 # Indent the whole chart body to align with the title above
515 chart_indent = " "
517 filled_glyph = "▐" + "█" * (col_width - 1)
518 empty_glyph = " " * col_width
520 colors = _model_colors(theme)
522 # Render rows top -> bottom
523 for row in range(height - 1, -1, -1):
524 text.append(chart_indent, style=theme.get("meta", "grey62"))
525 if row == height - 1:
526 text.append(gutter_fmt.format(y_label_strs[0]), style=theme.get("meta", "grey62"))
527 elif row == height // 2:
528 text.append(gutter_fmt.format(y_label_strs[1]), style=theme.get("meta", "grey62"))
529 elif row == 0:
530 text.append(gutter_fmt.format(y_label_strs[2]), style=theme.get("meta", "grey62"))
531 else:
532 text.append(empty_gutter, style=theme.get("meta", "grey62"))
534 text.append("│", style=theme.get("meta", "grey62"))
535 for col in columns:
536 cidx = col[row]
537 if cidx is not None:
538 color = colors[cidx % len(colors)]
539 text.append(filled_glyph, style=color)
540 else:
541 text.append(empty_glyph)
542 text.append("\n")
544 # X-axis line
545 text.append(chart_indent + empty_gutter + "└" + "─" * (col_width * len(days)) + "\n", style=theme.get("meta", "grey62"))
547 # Date labels: distribute evenly based on how many column slots a label
548 # occupies. A date label is 5 chars; we want at least 2 chars of visual gap.
549 if len(days) > 0:
550 date_line = Text(chart_indent + empty_gutter + " ")
551 label_len = len(short_date(days[0].date))
552 # Number of column slots needed per label (including desired gap)
553 slot_gap = max(1, (label_len + 2 + col_width - 1) // col_width)
555 label_positions = set(range(0, len(days), slot_gap))
556 # Always include the last day if not already present and it fits.
557 if len(days) > 1 and (len(days) - 1) % slot_gap != 0:
558 label_positions.add(len(days) - 1)
559 label_positions = sorted(label_positions)
561 max_line_chars = len(days) * col_width
562 skip = 0
563 for i, day in enumerate(days):
564 if skip > 0:
565 skip -= 1
566 continue
567 if i in label_positions:
568 lbl = short_date(day.date)
569 date_line.append(lbl, style=theme.get("meta", "grey62"))
570 consumed_cols = max(1, (len(lbl) + col_width - 1) // col_width)
571 # Pad to the column boundary, but never beyond the chart width
572 current_pos = i * col_width + len(lbl)
573 target_pos = min(i * col_width + consumed_cols * col_width, max_line_chars)
574 padding = target_pos - current_pos
575 if padding > 0:
576 date_line.append(" " * padding)
577 skip = consumed_cols - 1
578 else:
579 date_line.append(" " * col_width)
580 text.append(date_line)
581 text.append("\n")
583 # Legend
584 if top_models:
585 text.append("\n")
586 text.append(" ", style=theme.get("meta", "grey62"))
587 for i, model in enumerate(top_models[:8]):
588 name = truncate(_model_short_name(model), 14)
589 color = colors[i % len(colors)]
590 text.append("█", style=color)
591 text.append(f" {name} ", style=theme.get("meta", "grey62"))
592 text.append("\n")
594 return text
597def _model_colors(theme: dict) -> list[str]:
598 return [
599 theme.get("ok", "medium_spring_green"),
600 theme.get("warning", "gold1"),
601 theme.get("danger", "indian_red1"),
602 "dodger_blue1",
603 "medium_purple1",
604 "dark_orange",
605 "turquoise2",
606 "hot_pink",
607 "spring_green2",
608 "gold3",
609 ]
612def _top_n_models(model_totals: dict[str, float], n: int) -> list[str]:
613 items = sorted(model_totals.items(), key=lambda x: (-x[1], x[0]))
614 return [m for m, _ in items[:n]]
617def short_date(date: str) -> str:
618 if len(date) >= 10:
619 return date[5:10]
620 return date
623def truncate(s: str, max_len: int) -> str:
624 if len(s) > max_len:
625 return s[:max_len - 3] + "..."
626 return s
630def _render_top_models(top_models, lang_zh: bool, theme: dict, metric: str = OR_METRIC_REQUESTS, chart_width: int = 20) -> Text:
631 if not top_models:
632 return Text()
633 metric = _parse_or_metric(metric)
634 max_val = max(_metric_value_model(m, metric) for m in top_models) or 1.0
635 text = Text()
636 title = "Top 模型" if lang_zh else "Top Models"
637 text.append(f" {title}\n", style=theme.get("label", "cornflower_blue"))
638 colors = [
639 theme.get("ok", "medium_spring_green"),
640 theme.get("warning", "gold1"),
641 theme.get("danger", "indian_red1"),
642 "dodger_blue1",
643 "medium_purple1",
644 "dark_orange",
645 "turquoise2",
646 ]
647 for i, m in enumerate(top_models[:7]):
648 val = _metric_value_model(m, metric)
649 bar_len = int(val / max_val * chart_width) if max_val > 0 else 0
650 name = _model_short_name(m.model)
651 if len(name) > 22:
652 name = name[:19] + "..."
653 color = colors[i % len(colors)]
654 text.append(f" {name:22} ", style=theme.get("meta", "grey62"))
655 text.append(f"{_format_metric_value(val, metric):>8} ", style=theme.get("meta", "grey62"))
656 text.append("█" * bar_len, style=color)
657 text.append("░" * (chart_width - bar_len) + "\n", style="grey50")
658 return text
662def _format_aggregated_results(
663 results: dict[str, list[ProviderUsage]],
664 errors: dict[str, str],
665 order: list[str] | None = None,
666 theme_name: str = "blue-dark",
667 lang_zh: bool = IS_ZH,
668 enabled_providers: set | None = None,
669 or_metric: str = OR_METRIC_REQUESTS,
670 days_window: int = 30,
671) -> Text:
672 _L = L_ZH if lang_zh else L_EN
673 result = Text()
674 if order is None:
675 order = ["anthropic", "openai", "openrouter", "kimi"]
677 theme = THEME_MAP.get(theme_name)
678 if not theme:
679 theme = THEME_MAP["blue-dark"]
681 # First pass: pre-render all bodies and find global_max_width
682 provider_bodies = {}
683 error_msgs = {}
684 global_max_width = 0
686 for p in order:
687 p_items = results.get(p)
688 if not p_items:
689 continue
691 visual_widths = [_get_visual_width(_get_localized_label(r.label, lang_zh)) for r in p_items]
692 # Include metadata labels if applicable to ensure correct alignment padding
693 has_countdown = any(r.countdown for r in p_items)
694 has_reset = any(r.reset_at for r in p_items)
695 if has_countdown:
696 visual_widths.append(_get_visual_width(_L['countdown']))
697 if has_reset:
698 visual_widths.append(_get_visual_width(_L['reset']))
700 max_visual_width = max(visual_widths) if visual_widths else 0
701 max_visual_width = max(max_visual_width, 6)
702 bar_width = 20
704 body_text = Text()
705 for i, row in enumerate(p_items):
706 if i > 0:
707 body_text.append("\n")
709 has_structured = row.activity_totals or row.daily_activity or row.top_models
711 if has_structured:
712 # Structured rows render only their visualizations, no generic label line
713 if row.activity_totals:
714 if i > 0:
715 body_text.append("\n")
716 body_text.append(_render_activity_totals(row.activity_totals, lang_zh, theme))
717 if row.daily_activity:
718 if i > 0 or row.activity_totals:
719 body_text.append("\n")
720 body_text.append(_render_daily_chart(row.daily_activity, lang_zh, theme, metric=or_metric, days_window=days_window))
721 if row.top_models:
722 if i > 0 or row.activity_totals or row.daily_activity:
723 body_text.append("\n")
724 body_text.append(_render_top_models(row.top_models, lang_zh, theme, metric=or_metric, chart_width=20))
725 continue
727 loc_label = _get_localized_label(row.label, lang_zh)
728 label_v_width = _get_visual_width(loc_label)
729 padding = " " * (max_visual_width - label_v_width)
731 body_text.append(f" {loc_label}{padding} ", style=theme["label"])
733 if row.limit is not None and row.limit > 0:
734 used_ratio = row.used / row.limit
735 remaining_percent = max(0.0, 100.0 - (used_ratio * 100.0))
736 color = theme["danger"] if used_ratio > 0.9 else theme["warning"] if used_ratio > 0.7 else theme["ok"]
737 filled = min(bar_width, int(used_ratio * bar_width))
739 body_text.append("█" * filled, style=color)
740 body_text.append("·" * (bar_width - filled), style="grey50")
742 if row.unit == "%":
743 body_text.append(f" {used_ratio * 100:.0f}% {remaining_percent:.0f}% {_L['remaining']}", style="bold")
744 elif row.unit == "$":
745 body_text.append(f" ${row.used:.2f} / ${row.limit:.2f} ({remaining_percent:.0f}% {_L['remaining']})", style="bold")
746 else:
747 body_text.append(f" {row.used:,.0f} / {row.limit:,.0f} {row.unit} ({remaining_percent:.0f}% {_L['remaining']})", style="bold")
748 else:
749 if row.unit == "$":
750 body_text.append(f" ${row.used:.2f}", style="bold")
751 elif row.unit == "text":
752 if row.text_value:
753 loc_val = _get_localized_text_value(row.text_value, lang_zh)
754 body_text.append(f" {loc_val}", style="bold")
755 else:
756 body_text.append(f" {row.used:,.0f} {row.unit}", style="bold")
758 if row.countdown:
759 body_text.append("\n")
760 meta_label = _L['countdown']
761 meta_padding = " " * (max_visual_width - _get_visual_width(meta_label))
762 body_text.append(f" {meta_label}{meta_padding} ", style=theme["label"])
763 body_text.append(row.countdown, style="bold")
764 if row.reset_at:
765 body_text.append("\n")
766 meta_label = _L['reset']
767 meta_padding = " " * (max_visual_width - _get_visual_width(meta_label))
768 body_text.append(f" {meta_label}{meta_padding} ", style=theme["label"])
769 body_text.append(row.reset_at, style="bold")
771 provider_bodies[p] = body_text
772 lines = body_text.plain.split("\n")
773 max_line_width = max((_get_visual_width(line) for line in lines), default=0)
774 if max_line_width > global_max_width:
775 global_max_width = max_line_width
777 # 2. Process errors and unconfigured visible providers
778 for p in order:
779 if p in errors:
780 err = errors[p]
781 if err == "Not configured" and lang_zh:
782 translated_err = "未配置"
783 else:
784 translated_err = err
785 err_msg = f" ⚠ {translated_err}"
786 error_msgs[p] = err_msg
787 err_width = _get_visual_width(err_msg)
788 if err_width > global_max_width:
789 global_max_width = err_width
790 elif p not in results:
791 if enabled_providers is not None and p not in enabled_providers:
792 continue
793 err_msg = " ⚠ 未配置" if lang_zh else " ⚠ Not configured"
794 error_msgs[p] = err_msg
795 err_width = _get_visual_width(err_msg)
796 if err_width > global_max_width:
797 global_max_width = err_width
799 max(50, global_max_width + 4)
801 # Second pass: construct final result with simple divider titles
802 for p in order:
803 body_text = provider_bodies.get(p)
804 err_msg = error_msgs.get(p)
806 title_str = "Kimi" if p == "kimi" else p.capitalize()
807 divider_line = f"───────── {title_str} ─────────"
809 if body_text:
810 result.append(f"{divider_line}\n", style=theme["title"])
811 result.append("\n") # Empty line between divider and text
812 result.append(body_text)
814 if err_msg:
815 if not body_text:
816 result.append(f"{divider_line}\n", style=theme["danger"])
817 result.append("\n") # Empty line between divider and error
818 else:
819 result.append("\n\n") # Empty line between body and error
820 result.append(f"{err_msg}", style=theme["danger"])
822 result.append("\n\n") # Empty line below provider section
824 return result
827async def _interactive_mode(config: "AppConfig", initial_theme: str) -> None:
828 """Live Rich TUI with keyboard theme cycling and provider panel toggling.
830 Keys:
831 ←/→ or [ / ] — cycle theme
832 1 / 2 / 3 / 4 — toggle provider panels
833 r — refresh data
834 q / Ctrl-C — quit
835 """
836 import termios
837 import tty
839 themes = list(THEME_MAP.keys())
840 try:
841 idx = themes.index(initial_theme)
842 except ValueError:
843 idx = 0
845 results: dict[str, list[ProviderUsage]] = {}
846 errors: dict[str, str] = {}
848 # Determine initial visible providers and language
849 if config.visible_providers is not None:
850 visible_providers: set = {p for p in config.visible_providers if p in config.provider_order}
851 else:
852 visible_providers: set = set(config.provider_order)
854 if config.language == "zh":
855 lang_zh: bool = True
856 elif config.language == "en":
857 lang_zh: bool = False
858 else:
859 lang_zh: bool = IS_ZH
861 or_metric: str = _parse_or_metric(config.or_metric)
862 days_window: int = _parse_days_window(config.days_window)
864 saved_notice: list = [None] # holds the saved theme name briefly, then None
866 _SHORT = {
867 "anthropic": "Anthropic", "openai": "Openai",
868 "openrouter": "Openrouter", "kimi": "Kimi",
869 }
871 def _build_panel() -> Panel:
872 visible_order = [p for p in config.provider_order if p in visible_providers]
873 _L = L_ZH if lang_zh else L_EN
874 body = _format_aggregated_results(results, errors, visible_order, themes[idx], lang_zh, enabled_providers=set(config.enabled_providers), or_metric=or_metric, days_window=days_window)
876 top_bar = Text()
877 # Provider toggles
878 for i, p in enumerate(config.provider_order, 1):
879 short = _SHORT.get(p, p[:4].title())
880 if p in visible_providers:
881 top_bar.append("● ", style="bold")
882 top_bar.append(f"[{i}]{short} ", style="bold")
883 else:
884 top_bar.append("○ ", style="dim")
885 top_bar.append(f"[{i}]{short} ", style="dim italic")
886 # Theme name
887 top_bar.append("│ ", style="dim")
888 top_bar.append("主题: " if lang_zh else "theme: ", style="dim")
889 top_bar.append(themes[idx], style="bold")
890 top_bar.justify = "center"
892 hint = Text()
893 # Keybindings grouped logically: navigation / view / data / persistence
894 bindings = [
895 ("[q]", "退出" if lang_zh else "quit"),
896 ("[r]", "刷新" if lang_zh else "refresh"),
897 ("[←/→]", "主题" if lang_zh else "theme"),
898 ("[1-4]", "面板" if lang_zh else "panels"),
899 ("[l]", "英" if lang_zh else "ZH"),
900 ("[m]", _metric_label(or_metric, lang_zh)),
901 ("[d]", f"{days_window}d"),
902 ("[⏎ ]", "保存" if lang_zh else "save"),
903 ]
904 for i, (key, label) in enumerate(bindings):
905 if i > 0:
906 hint.append(" ", style="dim")
907 hint.append(key, style="bold")
908 hint.append(f" {label}", style="dim")
909 if saved_notice[0]:
910 hint.append(f" ✓ {'已保存' if lang_zh else 'Saved'}: {saved_notice[0]}", style="bold green")
911 hint.justify = "center"
913 panel_content = Group(
914 top_bar,
915 Text(""),
916 body,
917 hint
918 )
920 return Panel(panel_content, title=f"[bold]{_L['title']}[/bold]", subtitle=None,
921 expand=True, padding=(1, 2, 1, 2))
923 fd = sys.stdin.fileno()
924 old_settings = termios.tcgetattr(fd)
926 def _read_char() -> str:
927 # Prevent input buffering split by using os.read in raw terminal; fall back to stdin.read under tests
928 is_mock = hasattr(sys.stdin.read, "mock") or hasattr(sys.stdin.read, "_mock_self") or "mock" in type(sys.stdin.read).__name__.lower()
929 if not is_mock and hasattr(sys.stdin, "isatty") and sys.stdin.isatty(): # pragma: no cover
930 import os
931 try:
932 b = os.read(fd, 1)
933 if b:
934 return b.decode('utf-8', errors='ignore')
935 except Exception:
936 pass
937 return sys.stdin.read(1)
939 try:
940 tty.setcbreak(fd)
941 results, errors = await dispatch_all(config)
942 console = Console()
943 with Live(_build_panel(), console=console, refresh_per_second=4) as live:
944 running = True
945 while running:
946 rlist, _, _ = _select_module.select([sys.stdin], [], [], 0.15)
947 if rlist:
948 ch = _read_char()
949 if not ch: # pragma: no cover
950 continue
952 if ch == '\x1b':
953 r2, _, _ = _select_module.select([sys.stdin], [], [], 0.05)
954 if r2:
955 ch2 = _read_char()
956 if ch2 == '[':
957 r3, _, _ = _select_module.select([sys.stdin], [], [], 0.05)
958 if r3:
959 ch3 = _read_char()
960 ch = f'\x1b[{ch3}'
962 if saved_notice[0]: # clear notice on any new keypress
963 saved_notice[0] = None
964 if ch in ('\r', '\n'): # Enter → save settings to config file
965 save_theme(
966 themes[idx],
967 language="zh" if lang_zh else "en",
968 visible_providers=list(visible_providers),
969 or_metric=or_metric,
970 days_window=days_window
971 )
972 saved_notice[0] = themes[idx]
973 live.update(_build_panel())
974 # Flush any consecutive \r or \n (like \r\n from terminal)
975 while True:
976 r_extra, _, _ = _select_module.select([sys.stdin], [], [], 0.0)
977 if r_extra: # pragma: no cover
978 extra_ch = _read_char()
979 if extra_ch in ('\r', '\n'): # pragma: no cover
980 continue
981 break
982 continue
983 idx, quit_flag, refresh_flag, toggle_num, lang_toggle, metric_toggle, days_toggle = _handle_key(ch, idx, len(themes))
984 if quit_flag:
985 running = False
986 elif refresh_flag:
987 results, errors = await dispatch_all(config)
988 elif toggle_num is not None:
989 providers = list(config.provider_order)
990 if toggle_num < len(providers):
991 p = providers[toggle_num]
992 if p in visible_providers:
993 visible_providers.discard(p)
994 else:
995 visible_providers.add(p)
996 elif lang_toggle:
997 lang_zh = not lang_zh
998 elif metric_toggle:
999 or_metric = _next_or_metric(or_metric)
1000 elif days_toggle:
1001 days_window = _next_days_window(days_window)
1002 live.update(_build_panel())
1003 await asyncio.sleep(0)
1004 finally:
1005 termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
1008async def main():
1009 load_dotenv()
1010 parser = argparse.ArgumentParser(description="Multi-Provider AI Quota Tracker CLI")
1011 parser.add_argument("--json", action="store_true")
1012 parser.add_argument("--plain", action="store_true")
1013 parser.add_argument("-i", "--interactive", action="store_true",
1014 help="Interactive mode: ←/→ or [/] to cycle themes, r refresh, q quit")
1015 parser.add_argument("--provider", help="Comma-separated providers to query (kimi,openai,anthropic,openrouter)")
1016 parser.add_argument("--config", help="Custom configuration file path")
1017 parser.add_argument("--theme", help="Specify color theme: blue-dark, blue-light, sky-dark, salmon-dark, turquoise-dark, pink-light, violet-dark, amber-dark, mint-dark, monochrome, blind-deuteranopia, blind-tritanopia")
1018 args = parser.parse_args()
1020 resolver = ConfigResolver(config_path=args.config)
1021 config = resolver.resolve()
1023 # Determine theme
1024 theme_name = args.theme if args.theme else config.theme
1026 # If --provider filter is provided, override enabled providers
1027 allowed_providers = None
1028 if args.provider:
1029 allowed_providers = [p.strip().lower() for p in args.provider.split(",") if p.strip()]
1030 # Filter config providers
1031 for p in list(config.providers.keys()):
1032 if p not in allowed_providers:
1033 config.providers[p].api_key = None
1035 # Interactive mode: hand off to TUI loop (handles its own dispatch)
1036 if args.interactive:
1037 await _interactive_mode(config, theme_name)
1038 return
1040 results, errors = await dispatch_all(config)
1042 # Filter outputs
1043 if allowed_providers:
1044 results = {k: v for k, v in results.items() if k in allowed_providers}
1045 errors = {k: v for k, v in errors.items() if k in allowed_providers}
1047 if args.json:
1048 output_data = {}
1049 for p, items in results.items():
1050 output_data[p] = [
1051 {
1052 "label": item.label,
1053 "used": item.used,
1054 "limit": item.limit,
1055 "remaining": item.remaining,
1056 "percent": item.percent,
1057 "reset_at": item.reset_at,
1058 "unit": item.unit
1059 }
1060 for item in items
1061 ]
1062 # Include errors in JSON if any
1063 if errors:
1064 output_data["errors"] = errors
1065 print(json.dumps(output_data, ensure_ascii=False))
1066 elif args.plain:
1067 if not results and not errors:
1068 print(L["no_data"])
1069 else:
1070 for p in config.provider_order:
1071 items = results.get(p)
1072 if items:
1073 for item in items:
1074 if item.limit is not None and item.limit > 0:
1075 pct_used = item.used / item.limit * 100
1076 print(f"{p.capitalize()} - {item.label}: {item.used}/{item.limit} ({pct_used:.0f}% used)")
1077 else:
1078 if item.unit == "text":
1079 if item.text_value:
1080 print(f"{p.capitalize()} - {item.label}: {item.text_value}")
1081 else:
1082 print(f"{p.capitalize()} - {item.label}")
1083 else:
1084 print(f"{p.capitalize()} - {item.label}: {item.used} ({item.unit})")
1085 for p in config.provider_order:
1086 err = errors.get(p)
1087 if err:
1088 print(f"{p.capitalize()} - Error: {err}", file=sys.stderr)
1089 else:
1090 console = Console()
1091 if not results and not errors:
1092 console.print(Panel(Text(L["no_data"], style="dim"), title=f"[bold]{L['title']}[/bold]"))
1093 else:
1094 display_order = config.provider_order
1095 if config.visible_providers is not None:
1096 display_order = [p for p in config.provider_order if p in config.visible_providers]
1097 or_metric = _parse_or_metric(config.or_metric)
1098 days_window = _parse_days_window(config.days_window)
1099 console.print(Panel(_format_aggregated_results(results, errors, display_order, theme_name, enabled_providers=set(config.enabled_providers), or_metric=or_metric, days_window=days_window), title=f"[bold]{L['title']}[/bold]", expand=True, padding=(1, 2, 1, 2)))
1101def run_cli():
1102 asyncio.run(main())
1104if __name__ == "__main__":
1105 run_cli()