diff --git a/README.md b/README.md index ab86a04..04f809b 100644 --- a/README.md +++ b/README.md @@ -84,11 +84,13 @@ Environment variables override the config file: ## Unified Dashboard `shusub2` opens a compact dashboard for Accounts, Keys, Logs, and Errors. Its -summary also reports Pricing Monitor upstream health and CNY total. Press `a`, -`k`, `l`, or `e` to focus a dashboard table; press `p` to open the full -Upstreams source-balance page. `/` filters the active table and `r` refreshes -all data immediately. Automatic Accounts and Upstreams, Keys/Logs, and Errors refresh defaults -to every five minutes. The client requests +summary also reports Pricing Monitor upstream health and CNY total. The +Accounts table attaches a source balance only when an account has a canonical +quota-source name and the sanitized Pricing Monitor projection has one +unambiguous matching source. Press `a`, `k`, `l`, or `e` to focus a dashboard +table; press `p` to open the full Upstreams source-balance page. `/` filters +the active table and `r` refreshes all data immediately. Automatic Accounts +and Upstreams, Keys/Logs, and Errors refresh defaults to every five minutes. The client requests gzip-compressed JSON and transparently decodes it when the upstream supports it. The two-line detail area follows the selected row. Accounts, Keys, Logs, and Errors use green, yellow, magenta, and red section styling respectively. @@ -98,7 +100,7 @@ complete keys, accounts, and models. On a narrow terminal, a table scrolls horizontally instead of truncating a field: ```text -ACCOUNT | Group | Today | Daily | 5h | 7d | Avail +ACCOUNT | Group | Source | Src CNY | Src state | Today | Daily | 5h | 7d | Avail UPSTREAMS (p) | Source | Kind | Balance | CNY | Unit | Status | Updated | Age KEY | Today | Tokens | Req LOG KEY | Account | Model | First | Duration | Tok/s | Input | Output | Cache | Tokens | Cost | Time | Age @@ -252,7 +254,7 @@ panel after the accounts table when a token is configured. The dedicated Accounts page keeps the full columns for scanning inside zellij: ```text -Name | Provider | Group | Daily | Today | Tokens | Req | Kind | 5h | 7d | Reset | Status | Availability +Name | Provider | Group | Source | Src CNY | Src state | Daily | Today | Tokens | Req | Kind | 5h | 7d | Reset | Status | Availability ``` `Provider` distinguishes `openai` and `anthropic` accounts from the public @@ -265,6 +267,15 @@ Name | Provider | Group | Daily | Today | Tokens | Req | Kind | 5h | 7d | Reset are not upstream provider balances and therefore never claim to be a cash or source-account balance. +`Source`, `Src CNY`, and `Src state` use Pricing Monitor's cached, sanitized +`view=accounts` projection. A balance is attached only for exact canonical +account names of the form `{family}-quota-{source}` or +`{family}-quotaonly-{source}`, where `{source}` is one or more lowercase +alphanumeric segments joined by single hyphens, and a single normalized source +name. Provider, URL, display-name, and fuzzy matching are deliberately not used. Accounts +without a unique mapping show `-`; an `error` or `stale` source retains its +source/state label but does not show a CNY amount. + The dedicated `--pricing` page shows Pricing Monitor source rows such as `code-plan`, `codexapis`, `testvideo`, `kedaya`, `ycy`, and `mdkj`: raw available balance with its source unit, derived CNY when supplied, health state, and last diff --git a/pyproject.toml b/pyproject.toml index c68c5f2..ae4a10c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "shusub2" -version = "0.2.14" +version = "0.2.15" description = "Terminal UI for Sub2API account quota and daily usage" readme = "README.md" requires-python = ">=3.11" diff --git a/sub2api_quota_tui.py b/sub2api_quota_tui.py index c3f64ea..e9fb393 100644 --- a/sub2api_quota_tui.py +++ b/sub2api_quota_tui.py @@ -21,7 +21,7 @@ from typing import Any APP_NAME = "shusub2" -FALLBACK_VERSION = "0.2.14" +FALLBACK_VERSION = "0.2.15" DEFAULT_API_URL = "http://127.0.0.1:18318/api/tui/accounts" DEFAULT_CONFIG_FILE = "~/.config/shusub2/api-url" DEFAULT_STATUS_CONFIG_FILE = "~/.config/shusub2/status-url" @@ -47,6 +47,10 @@ DEFAULT_VERSION_CHECK_TIMEOUT_SECONDS = 2 MONITOR_OK_STATUSES = {"operational", "ok", "success"} MONITOR_FAILED_STATUSES = {"error", "failed", "failure"} MONITOR_STOPWORDS = {"response", "responses", "monitor"} +PRICING_ACCOUNT_SOURCE_RE = re.compile( + r"^[a-z0-9]+-quota(?:only)?-(?P[a-z0-9]+(?:-[a-z0-9]+)*)$", + re.IGNORECASE, +) INSTALL_COMMAND = "uv tool install --force git+https://gitea.shujk.top/shujakuin/shusub2.git" INSTALL_COMMAND_ARGS = ["uv", "tool", "install", "--force", "git+https://gitea.shujk.top/shujakuin/shusub2.git"] @@ -404,11 +408,32 @@ def fetch_optional_payload(url: str, timeout: int) -> tuple[dict[str, Any], str] def fetch_pricing_payload(pricing_url: str, timeout: int) -> dict[str, Any]: payload = fetch_payload(pricing_url, timeout) - if not isinstance(payload.get("sources"), list): - raise RuntimeError("Pricing Monitor did not return source balances") + sources = payload.get("sources") + if payload.get("view") != "accounts" or not isinstance(sources, list): + raise RuntimeError("Pricing Monitor did not return an accounts source projection") + for source in sources: + if ( + not isinstance(source, dict) + or not isinstance(source.get("name"), str) + or not source["name"].strip() + or not isinstance(source.get("source_kind"), str) + or not source["source_kind"].strip() + or not isinstance(source.get("health_state"), str) + or not source["health_state"].strip() + or not isinstance(source.get("balance_available"), bool) + or not isinstance(source.get("balance"), dict) + ): + raise RuntimeError("Pricing Monitor returned an invalid accounts source projection") return payload +def fetch_optional_pricing_payload(pricing_url: str, timeout: int) -> tuple[dict[str, Any], str]: + try: + return fetch_pricing_payload(pricing_url, timeout), "" + except Exception: + return {}, "unavailable" + + def logs_request_url(logs_url: str, limit: int) -> str: parsed = urllib.parse.urlparse(logs_url) query = urllib.parse.parse_qsl(parsed.query, keep_blank_values=True) @@ -1132,12 +1157,18 @@ def reset_cell(account: dict[str, Any]) -> str: return "-" -def normalize_account_rows(payload: dict[str, Any], filter_text: str = "") -> list[dict[str, Any]]: +def normalize_account_rows( + payload: dict[str, Any], + filter_text: str = "", + pricing_payload: dict[str, Any] | None = None, +) -> list[dict[str, Any]]: needle = filter_text.strip().lower() + pricing_sources = pricing_source_lookup(pricing_payload or {}) rows = [] for account in payload.get("accounts") or []: if not isinstance(account, dict): continue + pricing_source = pricing_source_for_account(account.get("name"), pricing_sources) haystack = " ".join( str(account.get(key) or "") for key in ( @@ -1153,6 +1184,8 @@ def normalize_account_rows(payload: dict[str, Any], filter_text: str = "") -> li "daily_quota_timezone", ) ).lower() + if pricing_source: + haystack = f"{haystack} {pricing_source.get('name') or ''} {pricing_source.get('status') or ''}".lower() if needle and needle not in haystack: continue routing_group = str(account.get("routing_group") or "-") @@ -1162,6 +1195,12 @@ def normalize_account_rows(payload: dict[str, Any], filter_text: str = "") -> li "name": str(account.get("name") or ""), "routing_group": routing_group, "base_url_hash": str(account.get("base_url_hash") or ""), + "pricing_source": str(pricing_source.get("name") or "-") if pricing_source else "-", + "pricing_cny": pricing_source.get("balance_cny") if pricing_source and pricing_source.get("status") == "healthy" else None, + "pricing_state": str(pricing_source.get("status") or "-") if pricing_source else "-", + "pricing_balance": pricing_source.get("balance") if pricing_source else None, + "pricing_unit": str(pricing_source.get("unit") or "-") if pricing_source else "-", + "pricing_updated": str(pricing_source.get("updated") or "-") if pricing_source else "-", "provider": provider_label(account), "kind": str(account.get("kind") or "unknown"), "kind_label": kind_label(account.get("kind")), @@ -1217,11 +1256,12 @@ def normalize_pricing_rows(payload: dict[str, Any], filter_text: str = "") -> li if not isinstance(source, dict): continue balance = source.get("balance") if isinstance(source.get("balance"), dict) else {} + balance_is_available = source.get("balance_available") is not False row = { "name": str(source.get("name") or "-"), "kind": str(source.get("source_kind") or "-"), - "balance": optional_number(balance.get("available")), - "balance_cny": optional_number(balance.get("available_cny")), + "balance": optional_number(balance.get("available")) if balance_is_available else None, + "balance_cny": optional_number(balance.get("available_cny")) if balance_is_available else None, "unit": str(balance.get("unit") or "-"), "status": pricing_status(source), "updated_at": source.get("last_success_at"), @@ -1240,6 +1280,32 @@ def normalize_pricing_rows(payload: dict[str, Any], filter_text: str = "") -> li return rows +def pricing_source_key(value: Any) -> str: + return re.sub(r"[^a-z0-9]+", "", str(value or "").strip().casefold()) + + +def pricing_source_lookup(payload: dict[str, Any]) -> dict[str, dict[str, Any] | None]: + lookup: dict[str, dict[str, Any] | None] = {} + for source in normalize_pricing_rows(payload): + key = pricing_source_key(source.get("name")) + if not key: + continue + if key in lookup: + lookup[key] = None + else: + lookup[key] = source + return lookup + + +def pricing_source_for_account( + account_name: Any, pricing_sources: dict[str, dict[str, Any] | None] +) -> dict[str, Any] | None: + match = PRICING_ACCOUNT_SOURCE_RE.fullmatch(str(account_name or "").strip()) + if not match: + return None + return pricing_sources.get(pricing_source_key(match.group("source"))) + + def pricing_summary(payload: dict[str, Any], rows: list[dict[str, Any]] | None = None) -> str: source_rows = rows if rows is not None else normalize_pricing_rows(payload) healthy = sum(row["status"] == "healthy" for row in source_rows) @@ -1264,6 +1330,18 @@ def pricing_detail_line(row: dict[str, Any]) -> str: return detail +def account_pricing_detail(row: dict[str, Any]) -> str: + source = str(row.get("pricing_source") or "-") + if source == "-": + return "upstream -" + cny = format_cny(row.get("pricing_cny")) + raw_balance = format_amount(row.get("pricing_balance")) + unit = str(row.get("pricing_unit") or "-") + state = str(row.get("pricing_state") or "-") + updated = str(row.get("pricing_updated") or "-") + return f"upstream {source} | CNY {cny} | raw {raw_balance} {unit} | {state} | updated {updated}" + + def print_pricing_once(payload: dict[str, Any], filter_text: str = "") -> None: rows = normalize_pricing_rows(payload, filter_text) print(f"{short_time(payload.get('generated_at'))} | {pricing_summary(payload, rows)}") @@ -1294,16 +1372,30 @@ def summary_line(payload: dict[str, Any]) -> str: ) -def print_once(payload: dict[str, Any], filter_text: str = "", status_payload: dict[str, Any] | None = None, status_error: str = "") -> None: +def print_once( + payload: dict[str, Any], + filter_text: str = "", + status_payload: dict[str, Any] | None = None, + status_error: str = "", + pricing_payload: dict[str, Any] | None = None, + pricing_error: str = "", +) -> None: print(summary_line(payload)) if status_payload or status_error: print(monitor_summary(status_payload or {}, status_error)) - print("name provider group daily today tokens req kind 5h 7d reset status availability") - for row in normalize_account_rows(payload, filter_text): + if pricing_error: + print(f"upstreams: {pricing_error}") + elif pricing_payload: + print(pricing_summary(pricing_payload)) + print("name provider group source src cny src state daily today tokens req kind 5h 7d reset status availability") + for row in normalize_account_rows(payload, filter_text, pricing_payload): print( f"{row['name'][:30]:<30} " f"{row['provider']:<9} " f"{row['routing_group']:<6} " + f"{row['pricing_source']:<16} " + f"{format_cny(row['pricing_cny']):<13} " + f"{row['pricing_state']:<12} " f"{row['daily_quota_cell']:<10} " f"{format_cost(row['today_cost_usd']):<10} " f"{format_count(row['today_tokens']):<8} " @@ -1404,7 +1496,7 @@ def run_textual( def on_mount(self) -> None: self.configure_table( self.query_one("#accounts", DataTable), - ("ACCOUNT", "Group", "Today", "Daily", "5h", "7d", "Avail"), + ("ACCOUNT", "Group", "Source", "Src CNY", "Src state", "Today", "Daily", "5h", "7d", "Avail"), ) self.configure_table( self.query_one("#keys", DataTable), @@ -1493,12 +1585,10 @@ def run_textual( self.render_meta() def refresh_pricing(self) -> None: - self.pricing_error = "" - try: - self.pricing_payload = fetch_pricing_payload(pricing_url, timeout) - except Exception as exc: - self.pricing_payload = {} - self.pricing_error = str(exc) + self.pricing_payload, self.pricing_error = fetch_optional_pricing_payload( + pricing_url, timeout + ) + self.render_accounts() self.render_meta() def refresh_logs(self) -> None: @@ -1557,7 +1647,9 @@ def run_textual( def render_accounts(self) -> None: filter_text = self.query_one("#filter", Input).value - self.account_rows = normalize_account_rows(self.accounts_payload, filter_text) + self.account_rows = normalize_account_rows( + self.accounts_payload, filter_text, self.pricing_payload + ) table = self.query_one("#accounts", DataTable) table.clear() self.account_by_key = {} @@ -1567,6 +1659,9 @@ def run_textual( table.add_row( row["name"], row["routing_group"], + row["pricing_source"], + format_cny(row["pricing_cny"]), + row["pricing_state"], format_cost(row["today_cost_usd"]), row["daily_quota_cell"], row["five_hour"], @@ -1736,6 +1831,7 @@ def run_textual( f"7d {row['weekly']} | " f"today {format_cost(row['today_cost_usd'])}, " f"{format_count(row['today_tokens'])} tokens, {format_count(row['today_requests'])} req | " + f"{account_pricing_detail(row)} | " f"{monitor_detail(row, self.status_payload)}" ) error = str(row.get("error") or "").strip() @@ -1784,6 +1880,8 @@ def run_textual( super().__init__() self.payload: dict[str, Any] = {} self.status_payload: dict[str, Any] = {} + self.pricing_payload: dict[str, Any] = {} + self.pricing_error = "" self.status_error = "" self.keys_payload: dict[str, Any] = {} self.keys_error = "" @@ -1804,7 +1902,7 @@ def run_textual( table = self.query_one("#accounts", DataTable) table.cursor_type = "row" table.zebra_stripes = True - table.add_columns("Name", "Provider", "Group", "Daily", "Today", "Tokens", "Req", "Kind", "5h", "7d", "Reset", "Status", "Availability") + table.add_columns("Name", "Provider", "Group", "Source", "Src CNY", "Src state", "Daily", "Today", "Tokens", "Req", "Kind", "5h", "7d", "Reset", "Status", "Availability") keys_table = self.query_one("#keys", DataTable) keys_table.cursor_type = "row" keys_table.zebra_stripes = True @@ -1841,10 +1939,17 @@ def run_textual( try: self.payload = fetch_payload(api_url, timeout, refresh=refresh) self.status_payload, self.status_error = fetch_optional_payload(status_url, timeout) + self.pricing_payload, self.pricing_error = fetch_optional_pricing_payload( + pricing_url, timeout + ) self.refresh_keys() self.render_payload() self.render_keys() status_bits = [version_message, monitor_summary(self.status_payload, self.status_error), f"source {self.payload.get('source_name') or '-'}"] + if self.pricing_error: + status_bits.append(f"upstreams: {self.pricing_error}") + else: + status_bits.append(pricing_summary(self.pricing_payload)) if self.keys_error: status_bits.append(f"keys: {self.keys_error}") status.update(" | ".join(bit for bit in status_bits if bit) + f" | {api_url}") @@ -1878,7 +1983,7 @@ def run_textual( def render_payload(self) -> None: filter_text = self.query_one("#filter", Input).value - self.rows = normalize_account_rows(self.payload, filter_text) + self.rows = normalize_account_rows(self.payload, filter_text, self.pricing_payload) table = self.query_one("#accounts", DataTable) table.clear() self.row_by_key = {} @@ -1889,6 +1994,9 @@ def run_textual( row["name"], row["provider"], row["routing_group"], + row["pricing_source"], + format_cny(row["pricing_cny"]), + row["pricing_state"], row["daily_quota_cell"], format_cost(row["today_cost_usd"]), format_count(row["today_tokens"]), @@ -1920,6 +2028,7 @@ def run_textual( f"daily {row['daily_quota_cell']} ({format_percent(row['daily_quota_used_percent'])} used, " f"{row['daily_quota_remaining'] if row['daily_quota_remaining'] is not None else '-'} left) | " f"5h {row['five_hour']} | 7d {row['weekly']} | " + f"{account_pricing_detail(row)} | " f"today {format_cost(row['today_cost_usd'])}, {format_count(row['today_tokens'])} tokens, " f"{format_count(row['today_requests'])} req | latest {row['latest_usage_at']} | " f"priority {row['priority'] if row['priority'] is not None else '-'} | " @@ -1994,14 +2103,14 @@ def run_textual( self.render_payload() status_bits = [version_message, pricing_summary(self.payload), pricing_url] status.update(" | ".join(bit for bit in status_bits if bit)) - except Exception as exc: + except Exception: self.payload = {} self.rows = [] self.row_by_key = {} self.query_one("#pricing", DataTable).clear() self.query_one("#summary", Static).update("upstreams unavailable") self.query_one("#detail", Static).update("no upstream balances") - status.update(f"upstreams error: {exc}") + status.update("upstreams unavailable") def render_payload(self) -> None: filter_text = self.query_one("#filter", Input).value @@ -2430,7 +2539,13 @@ def main(argv: list[str] | None = None) -> int: if version_message: print(version_message) if args.pricing: - print_pricing_once(fetch_pricing_payload(args.pricing_url, args.timeout), args.filter) + try: + print_pricing_once( + fetch_pricing_payload(args.pricing_url, args.timeout), args.filter + ) + except Exception: + print("upstreams unavailable", file=sys.stderr) + return 1 return 0 if args.logs: if not str(args.logs_token or "").strip(): @@ -2454,7 +2569,15 @@ def main(argv: list[str] | None = None) -> int: ) return 0 status_payload, status_error = fetch_optional_payload(status_url, args.timeout) - print_once(fetch_payload(args.api_url, args.timeout, refresh=True), args.filter, status_payload, status_error) + pricing_payload, pricing_error = fetch_optional_pricing_payload(args.pricing_url, args.timeout) + print_once( + fetch_payload(args.api_url, args.timeout, refresh=True), + args.filter, + status_payload, + status_error, + pricing_payload, + pricing_error, + ) if str(args.logs_token or "").strip(): print() try: diff --git a/tests/test_payload.py b/tests/test_payload.py index d9d6497..7136a29 100644 --- a/tests/test_payload.py +++ b/tests/test_payload.py @@ -90,6 +90,55 @@ class Sub2APIQuotaTUITests(unittest.TestCase): self.assertEqual([row["name"] for row in mod.normalize_account_rows(payload, "sfast")], ["beta"]) self.assertEqual([row["name"] for row in mod.normalize_account_rows(payload, "anthropic")], ["alpha"]) + def test_account_rows_attach_only_canonical_pricing_source_balances(self) -> None: + mod = load_module() + accounts_payload = { + "accounts": [ + {"id": 1, "name": "oai-quota-code-plan", "routing_group": "fast", "kind": "quota_limited"}, + {"id": 2, "name": "oai-quotaonly-code-plan", "routing_group": "fast", "kind": "quota_limited"}, + {"id": 3, "name": "oai-sub-code-plan", "routing_group": "fast", "kind": "quota_limited"}, + {"id": 4, "name": "oai-quota-unknown", "routing_group": "fast", "kind": "quota_limited"}, + {"id": 5, "name": "oai-quota-ycy", "routing_group": "fast", "kind": "quota_limited"}, + {"id": 6, "name": "oai-quota-code--plan", "routing_group": "fast", "kind": "quota_limited"}, + {"id": 7, "name": "oai-quota-code-plan-", "routing_group": "fast", "kind": "quota_limited"}, + ] + } + pricing_payload = { + "sources": [ + { + "name": "code-plan", + "source_kind": "newapi", + "health_state": "healthy", + "last_success_at": "2026-08-02T12:00:00Z", + "balance": {"available": 5000000, "available_cny": 10, "unit": "quota"}, + }, + { + "name": "ycy", + "source_kind": "newapi", + "health_state": "error", + "last_error": "upstream request failed", + "balance": {"available": 1000000, "available_cny": 2, "unit": "quota"}, + }, + ] + } + + rows = {row["name"]: row for row in mod.normalize_account_rows(accounts_payload, pricing_payload=pricing_payload)} + + self.assertEqual(rows["oai-quota-code-plan"]["pricing_source"], "code-plan") + self.assertEqual(rows["oai-quota-code-plan"]["pricing_cny"], 10) + self.assertEqual(rows["oai-quotaonly-code-plan"]["pricing_source"], "code-plan") + self.assertEqual(rows["oai-sub-code-plan"]["pricing_source"], "-") + self.assertEqual(rows["oai-quota-unknown"]["pricing_source"], "-") + self.assertEqual(rows["oai-quota-ycy"]["pricing_source"], "ycy") + self.assertEqual(rows["oai-quota-ycy"]["pricing_state"], "error") + self.assertIsNone(rows["oai-quota-ycy"]["pricing_cny"]) + self.assertEqual(rows["oai-quota-code--plan"]["pricing_source"], "-") + self.assertEqual(rows["oai-quota-code-plan-"]["pricing_source"], "-") + + ambiguous = dict(pricing_payload, sources=[*pricing_payload["sources"], {"name": "codeplan", "health_state": "healthy", "balance": {"available_cny": 12}}]) + ambiguous_rows = {row["name"]: row for row in mod.normalize_account_rows(accounts_payload, pricing_payload=ambiguous)} + self.assertEqual(ambiguous_rows["oai-quota-code-plan"]["pricing_source"], "-") + def test_normalize_pricing_rows_uses_pricing_monitor_source_balances(self) -> None: mod = load_module() payload = { @@ -122,6 +171,24 @@ class Sub2APIQuotaTUITests(unittest.TestCase): self.assertEqual(rows[0]["unit"], "quota") self.assertEqual(rows[0]["status"], "healthy") self.assertEqual(rows[1]["status"], "error") + self.assertIsNone(rows[1]["balance"]) + self.assertIsNone(rows[1]["balance_cny"]) + unavailable_rows = mod.normalize_pricing_rows( + { + "sources": [ + { + "name": "unavailable", + "source_kind": "newapi", + "health_state": "healthy", + "balance_available": False, + "balance": {"available": 99, "available_cny": 99, "unit": "quota"}, + } + ] + } + ) + self.assertEqual(unavailable_rows[0]["status"], "healthy") + self.assertIsNone(unavailable_rows[0]["balance"]) + self.assertIsNone(unavailable_rows[0]["balance_cny"]) self.assertEqual([row["name"] for row in mod.normalize_pricing_rows(payload, "keday")], ["kedaya"]) self.assertEqual(mod.pricing_summary(payload, rows), "upstreams 1/2 healthy | CNY ¥10") out = io.StringIO() @@ -161,18 +228,60 @@ class Sub2APIQuotaTUITests(unittest.TestCase): else: os.environ["SHUSUB2_PRICING_URL_FILE"] = old_file - def test_fetch_pricing_payload_requires_sources_list(self) -> None: + def test_fetch_pricing_payload_requires_accounts_projection_shape(self) -> None: mod = load_module() original_fetch = mod.fetch_payload - mod.fetch_payload = lambda *args, **kwargs: {"view": "accounts", "sources": []} + valid_payload = { + "view": "accounts", + "sources": [ + { + "name": "code-plan", + "source_kind": "newapi", + "health_state": "healthy", + "balance_available": True, + "balance": {"available": 1, "available_cny": 1, "unit": "quota"}, + } + ], + } + mod.fetch_payload = lambda *args, **kwargs: valid_payload try: - self.assertEqual(mod.fetch_pricing_payload("https://price.example/api/ui-data?view=accounts", 1)["sources"], []) - mod.fetch_payload = lambda *args, **kwargs: {"view": "accounts"} - with self.assertRaisesRegex(RuntimeError, "source balances"): - mod.fetch_pricing_payload("https://price.example/api/ui-data?view=accounts", 1) + self.assertEqual( + mod.fetch_pricing_payload("https://price.example/api/ui-data?view=accounts", 1), + valid_payload, + ) + for invalid_payload in ( + {"view": "overview", "sources": []}, + {"view": "accounts"}, + {"view": "accounts", "sources": [{"name": "code-plan"}]}, + { + "view": "accounts", + "sources": [ + {"name": "", "source_kind": "newapi", "health_state": "healthy", "balance_available": True, "balance": {}} + ], + }, + ): + mod.fetch_payload = lambda *args, _payload=invalid_payload, **kwargs: _payload + with self.assertRaisesRegex(RuntimeError, "accounts source projection"): + mod.fetch_pricing_payload("https://price.example/api/ui-data?view=accounts", 1) finally: mod.fetch_payload = original_fetch + def test_optional_pricing_failure_is_nonblocking_and_redacted(self) -> None: + mod = load_module() + original_fetch = mod.fetch_pricing_payload + mod.fetch_pricing_payload = lambda *args, **kwargs: (_ for _ in ()).throw( + RuntimeError("https://pricing.example.test/?token=private") + ) + try: + payload, error = mod.fetch_optional_pricing_payload( + "https://pricing.example.test/api/ui-data?view=accounts", 1 + ) + finally: + mod.fetch_pricing_payload = original_fetch + + self.assertEqual(payload, {}) + self.assertEqual(error, "unavailable") + def test_infers_public_status_url_from_accounts_url(self) -> None: mod = load_module() @@ -830,7 +939,7 @@ class DashboardLayoutTests(unittest.IsolatedAsyncioTestCase): "accounts": [ { "id": 1, - "name": "account-one-with-a-complete-expanded-dashboard-name", + "name": "oai-quota-code-plan", "routing_group": "fast", "provider": "openai", "kind": "pay_as_you_go", @@ -936,6 +1045,8 @@ class DashboardLayoutTests(unittest.IsolatedAsyncioTestCase): self.assertGreater(logs_table.max_scroll_x, 0) self.assertGreater(errors_table.max_scroll_x, 0) self.assertEqual(screen.log_rows[0]["cache_tokens"], 4150) + self.assertEqual(screen.account_rows[0]["pricing_source"], "code-plan") + self.assertEqual(screen.account_rows[0]["pricing_cny"], 10) self.assertGreater(screen.log_rows[0]["tokens_per_second"], 0) self.assertEqual(len({table.styles.background for table in tables}), 4) self.assertIn("wmy", str(screen.query_one("#summary").render())) @@ -977,6 +1088,75 @@ class PageSelectionTests(unittest.TestCase): self.assertIn("code-plan", out.getvalue()) self.assertIn("¥10", out.getvalue()) + def test_once_pricing_failure_is_redacted(self) -> None: + mod = load_module() + original_fetch = mod.fetch_pricing_payload + mod.fetch_pricing_payload = lambda *args, **kwargs: (_ for _ in ()).throw( + RuntimeError("https://price.example.test/?token=private") + ) + err = io.StringIO() + try: + with contextlib.redirect_stderr(err): + rc = mod.main(["--once", "--pricing", "--no-version-check"]) + finally: + mod.fetch_pricing_payload = original_fetch + + self.assertEqual(rc, 1) + self.assertEqual(err.getvalue().strip(), "upstreams unavailable") + + def test_once_accounts_includes_only_mapped_pricing_source_balance(self) -> None: + mod = load_module() + original_fetch = mod.fetch_payload + original_optional = mod.fetch_optional_payload + original_pricing_optional = mod.fetch_optional_pricing_payload + accounts_payload = { + "source_name": "test", + "totals": {"total_accounts": 2, "usable_accounts": 2, "today_cost_usd": 0}, + "accounts": [ + {"id": 1, "name": "oai-quota-code-plan", "routing_group": "fast", "kind": "quota_limited"}, + {"id": 2, "name": "oai-sub-code-plan", "routing_group": "fast", "kind": "quota_limited"}, + ], + } + pricing_payload = { + "sources": [ + { + "name": "code-plan", + "source_kind": "newapi", + "health_state": "healthy", + "last_success_at": "2026-08-02T12:00:00Z", + "balance": {"available": 5000000, "available_cny": 10, "unit": "quota"}, + } + ] + } + mod.fetch_payload = lambda *args, **kwargs: accounts_payload + mod.fetch_optional_payload = lambda *args, **kwargs: ({}, "") + mod.fetch_optional_pricing_payload = lambda *args, **kwargs: (pricing_payload, "") + out = io.StringIO() + try: + with contextlib.redirect_stdout(out): + rc = mod.main( + [ + "--once", + "--api-url", + "https://accounts.example.test/api/tui/accounts", + "--pricing-url", + "https://pricing.example.test/api/ui-data?view=accounts", + "--no-version-check", + ] + ) + finally: + mod.fetch_payload = original_fetch + mod.fetch_optional_payload = original_optional + mod.fetch_optional_pricing_payload = original_pricing_optional + + text = out.getvalue() + self.assertEqual(rc, 0) + self.assertIn("src cny", text) + self.assertIn("oai-quota-code-plan", text) + self.assertIn("oai-sub-code-plan", text) + self.assertGreaterEqual(text.count("code-plan"), 2) + self.assertIn("¥10", text) + def test_dedicated_page_flags_are_mutually_exclusive(self) -> None: mod = load_module() old_token = os.environ.pop("SHUSUB2_LOGS_TOKEN", None) diff --git a/uv.lock b/uv.lock index 7684dc0..af63f3f 100644 --- a/uv.lock +++ b/uv.lock @@ -85,7 +85,7 @@ wheels = [ [[package]] name = "shusub2" -version = "0.2.14" +version = "0.2.15" source = { editable = "." } dependencies = [ { name = "textual" },