diff --git a/bin/omarchy-agent-usage-kimi b/bin/omarchy-agent-usage-kimi index f7d251743d1..0f3322a9234 100755 --- a/bin/omarchy-agent-usage-kimi +++ b/bin/omarchy-agent-usage-kimi @@ -12,6 +12,7 @@ import hashlib import json import math import os +import subprocess import sys import tempfile import threading @@ -99,6 +100,15 @@ def empty_stats() -> dict[str, Any]: } +def plan_label() -> str: + config_home = Path(os.environ.get("XDG_CONFIG_HOME") or Path.home() / ".config") + try: + configured = json.loads((config_home / "omarchy" / "agents" / "kimi.json").read_text()).get("planLabel") + except (OSError, json.JSONDecodeError, AttributeError): + return "" + return " ".join(configured.split())[:40] if isinstance(configured, str) else "" + + def base_record(**overrides: Any) -> dict[str, Any]: record: dict[str, Any] = { "schemaVersion": 1, @@ -111,7 +121,7 @@ def base_record(**overrides: Any) -> dict[str, Any]: # reports the same truth, so aggregation must not sum them. "scope": "account", "hasPromptStats": False, - "tierLabel": "", + "tierLabel": plan_label(), "usageStatusText": "", "authHelpText": "", "limits": [], @@ -295,6 +305,25 @@ def refresh_credentials(auth: dict[str, str] | None = None) -> dict[str, Any]: return credentials +def pi_bearer_token(auth: dict[str, str] | None = None) -> str: + # An explicit Kimi home is account-isolating; never borrow Pi's sign-in then. + if os.environ.get("KIMI_CODE_HOME"): + return "" + try: + result = subprocess.run( + ["pi", "auth", "print-bearer-token", "--provider", "kimi-coding", "--min-expiry", "5m"], + capture_output=True, text=True, timeout=30, check=False, + ) + except (OSError, subprocess.TimeoutExpired): + return "" + key = result.stdout.strip() if result.returncode == 0 else "" + if not key or len(key) > 8192 or any(character.isspace() for character in key): + return "" + if auth is not None: + auth["key"] = key + return key + + def api_key(auth: dict[str, str] | None = None) -> str: key = os.environ.get("KIMI_API_KEY", "").strip() if key: @@ -305,6 +334,8 @@ def api_key(auth: dict[str, str] | None = None) -> str: if auth is not None: auth["key"] = str(credentials.get("access_token") or "") if not credentials.get("access_token"): + if key := pi_bearer_token(auth): + return key raise KimiError(AUTH_HELP) if not token_fresh(credentials): credentials = refresh_credentials(auth) @@ -425,7 +456,9 @@ def to_limits(data: dict[str, Any]) -> list[dict[str, Any]]: except (TypeError, ValueError): continue if math.isfinite(ratio): - row = {"label": label, "percent": min(1.0, max(0.0, ratio)), "resetsAt": reset_at(detail)} + row = {"id": key, "label": label, "percent": min(1.0, max(0.0, ratio)), "resetsAt": reset_at(detail)} + if key == "limit_5h": + row["durationSeconds"] = 5 * 60 * 60 if key == "limit_month_code": row["title"] = "Code Monthly" limits.append(row) diff --git a/shell/plugins/agents/Panel.qml b/shell/plugins/agents/Panel.qml index 8d814e211e6..7696381d7a6 100644 --- a/shell/plugins/agents/Panel.qml +++ b/shell/plugins/agents/Panel.qml @@ -616,8 +616,13 @@ Panel { return value } + function percentNumber(fraction) { + var percent = Number(fraction) * 100 + return percent > 0 && percent < 1 ? "<" + Commons.I18n.formatNumber(1) : Commons.I18n.formatNumber(Math.round(percent)) + } + function scopedAllowanceText(scoped, windowTitle) { - var percent = Commons.I18n.formatNumber(Math.round(scoped.percent * 100)) + var percent = percentNumber(scoped.percent) var value = String(windowTitle || "") if (value === "Session") return Commons.I18n.t("agents.scopedSession").arg(scoped.title).arg(percent) if (value === "Weekly") return Commons.I18n.t("agents.scopedWeekly").arg(scoped.title).arg(percent) @@ -1997,7 +2002,7 @@ Panel { var lines = [] if (compact.window) { var windowName = root.windowDisplayTitle(compact.window.title) - var used = Commons.I18n.formatNumber(Math.round(compact.window.percent * 100)) + var used = root.percentNumber(compact.window.percent) lines.push(compact.resetMs > 0 ? Commons.I18n.t("agents.windowUsedResets").arg(windowName).arg(used).arg(root.formatDuration(compact.resetMs)) : Commons.I18n.t("agents.windowUsed").arg(windowName).arg(used)) diff --git a/shell/plugins/agents/README.md b/shell/plugins/agents/README.md index 9578180d47c..c03017d7206 100644 --- a/shell/plugins/agents/README.md +++ b/shell/plugins/agents/README.md @@ -102,7 +102,13 @@ pi stores in `~/.pi/agent/auth.json` when Fireworks is signed in there (honoring `PI_CODING_AGENT_DIR`, and pi's literal and `$ENV_VAR` key forms), and finally the key opencode stores in `~/.local/share/opencode/auth.json`. -Kimi reads `KIMI_API_KEY` first, then `credentials/kimi-code.json` under an explicit `KIMI_CODE_HOME`. Without that override it prefers `~/.kimi-code/credentials/kimi-code.json`, with the legacy `~/.kimi/credentials/kimi-code.json` as a compatibility fallback. An explicit home never borrows another home's sign-in. Its record reports account quotas, without local token or prompt statistics. +Kimi reads `KIMI_API_KEY` first, then `credentials/kimi-code.json` under an explicit `KIMI_CODE_HOME`. Without that override it prefers `~/.kimi-code/credentials/kimi-code.json`, with the legacy `~/.kimi/credentials/kimi-code.json` as a compatibility fallback. If no Kimi CLI sign-in exists, it can use Pi's `kimi-coding` OAuth through `pi auth print-bearer-token`, which owns token refresh; an explicit home never borrows another home's or Pi's sign-in. Its record reports account plan quotas from the `usages` response, not the separate request-rate allowance, and has no local token or prompt statistics. + +The usage endpoint does not report a plan name. Set one explicitly in `~/.config/omarchy/agents/kimi.json` when you want it shown beside Kimi in the panel and other snapshot readers: + +```json +{ "planLabel": "Ultra" } +``` Expired tokens are refreshed under the current CLI's renewable directory lock at `/oauth/kimi-code.lock`; legacy credentials use their sibling file lock. The rotated token pair is flushed and published atomically with private permissions. Credentials are reread after the exchange so a sign-in or logout completed during the request takes precedence. A failed lock prevents refresh, and credentials are sent only to trusted Kimi HTTPS hosts without following redirects. When sign-in expires, run `kimi login`; an invalid overriding `KIMI_API_KEY` must instead be updated or unset. diff --git a/test/shell.d/agent-usage-claude-limits-test.sh b/test/shell.d/agent-usage-claude-limits-test.sh index df39f9d8707..2f592868f5b 100644 --- a/test/shell.d/agent-usage-claude-limits-test.sh +++ b/test/shell.d/agent-usage-claude-limits-test.sh @@ -362,6 +362,13 @@ panelRoot.nowMs += 1000 assertDeepEqual(limitWindows(measured), [{ title: 'Session', percent: 0, resetAt: '' }], 'usage resets at the deadline while the panel stays open') assertEqual(measured.limits[0].percent, 0.78, 'display reset leaves the measured record intact') +const Commons = { I18n: { formatNumber: value => String(value) } } +eval(source.slice(source.indexOf('function percentNumber'), source.indexOf('function scopedAllowanceText'))) +assertEqual(percentNumber(0), '0', 'true zero remains zero') +assertEqual(percentNumber(0.0033), '<1', 'small nonzero usage does not round to zero') +assertEqual(percentNumber(0.01), '1', 'one percent remains one') +assertEqual(percentNumber(0.126), '13', 'ordinary percentages retain whole-number rounding') + // A model's weekly allowance rides under the Weekly row; one whose window has // no row of its own keeps a row. assertDeepEqual( diff --git a/test/shell.d/agent-usage-kimi-regression-test.sh b/test/shell.d/agent-usage-kimi-regression-test.sh index 9b5f9f92f1f..8ba54302271 100644 --- a/test/shell.d/agent-usage-kimi-regression-test.sh +++ b/test/shell.d/agent-usage-kimi-regression-test.sh @@ -22,7 +22,9 @@ from pathlib import Path requests = [] state = {} -payload = {"usages": {"limit_5h": {"used_ratio": 0.25, "reset_time": "2026-10-15T12:00:00Z"}, +payload = {"limits": [{"window": {"duration": 300, "timeUnit": "TIME_UNIT_MINUTE"}, + "detail": {"limit": "100", "used": "4", "remaining": "96"}}], + "usages": {"limit_5h": {"used_ratio": 0.25, "reset_time": "2026-10-15T12:00:00Z"}, "limit_7d": {"used_ratio": "0.8"}, "limit_month_total": {"used_ratio": 1}, "limit_month_code": {"used_ratio": 0}}} @@ -196,6 +198,9 @@ def modern_quota(root): data = record() assert [row["percent"] for row in data["limits"]] == [0.25, 0.8, 1, 0], "current CLI quota windows were dropped" assert data["limits"][0]["resetsAt"] == "2026-10-15T12:00:00Z" + assert [row.get("id") for row in data["limits"]] == ["limit_5h", "limit_7d", "limit_month_total", "limit_month_code"] + assert data["limits"][0]["durationSeconds"] == 18000 + assert all("used" not in row and "limit" not in row for row in data["limits"]), "request-rate counters leaked into plan quota" assert data["limits"][-1].get("title") == "Code Monthly", "monthly code allowance lost its scope" assert not data["hasLocalStats"] and not data["hasPromptStats"], "quota was advertised as local history" scanner.USAGE_URL = "https://api.kimi.ai/coding/v1/usages" diff --git a/test/shell.d/agent-usage-kimi-test.sh b/test/shell.d/agent-usage-kimi-test.sh index d1627f5c909..185ddb02040 100644 --- a/test/shell.d/agent-usage-kimi-test.sh +++ b/test/shell.d/agent-usage-kimi-test.sh @@ -27,6 +27,7 @@ import sys import time from contextlib import redirect_stdout from pathlib import Path +from types import SimpleNamespace collector_path = str(Path(sys.argv[1])) test_home = Path(os.environ["TEST_HOME"]) @@ -58,6 +59,21 @@ os.environ["KIMI_CODE_HOME"] = str(test_home / "kimi-code") (code_dir / "kimi-code.json").write_text("{}") summary["codeStorePreferred"] = scanner.credentials_dir() == code_dir +# With no Kimi CLI sign-in, Pi's external-client command supplies a fresh OAuth +# bearer without the collector reading or writing Pi's credential files. +del os.environ["KIMI_CODE_HOME"] +calls = [] +def pi_bearer(command, **kwargs): + calls.append((command, kwargs)) + return SimpleNamespace(returncode=0, stdout="pi_oauth\n") +scanner.subprocess.run = pi_bearer +auth = {} +summary["piFallback"] = scanner.api_key(auth) == "pi_oauth" and auth["key"] == "pi_oauth" and calls == [( + ["pi", "auth", "print-bearer-token", "--provider", "kimi-coding", "--min-expiry", "5m"], + {"capture_output": True, "text": True, "timeout": 30, "check": False}, +)] +os.environ["KIMI_CODE_HOME"] = str(test_home / "kimi-code") + # A fresh token is used as-is; an expired one is refreshed under the lock and # the rotated single-use refresh token is persisted back. scanner.CREDENTIALS_DIR = code_dir @@ -112,12 +128,17 @@ summary["limits"] = limits os.environ["KIMI_API_KEY"] = "kimi_env" scanner.quota = lambda key: payload +config_dir = test_home / "config" / "omarchy" / "agents" +config_dir.mkdir(parents=True) +(config_dir / "kimi.json").write_text(json.dumps({"planLabel": "Ultra"})) +os.environ["XDG_CONFIG_HOME"] = str(test_home / "config") sys.argv = ["omarchy-agent-usage-kimi", "--limits-only"] out = io.StringIO() with redirect_stdout(out): scanner.main() summary["record"] = json.loads(out.getvalue()) summary["quotaOnly"] = not summary["record"]["hasLocalStats"] and not summary["record"]["hasPromptStats"] +summary["planLabel"] = summary["record"]["tierLabel"] == "Ultra" print(json.dumps(summary, separators=(",", ":"))) PY @@ -131,6 +152,10 @@ pass "Kimi collector prefers KIMI_API_KEY over stored credentials" fail "Kimi collector prefers the kimi-code store and falls back to legacy kimi-cli" "$result" pass "Kimi collector prefers the kimi-code store and falls back to legacy kimi-cli" +[[ $(jq -r '.piFallback' <<<"$result") == "true" ]] || + fail "Kimi collector uses Pi's refreshed OAuth bearer when Kimi CLI has no sign-in" "$result" +pass "Kimi collector uses Pi's refreshed OAuth bearer when Kimi CLI has no sign-in" + [[ $(jq -r '"\(.freshTokenUsed):\(.refreshRotates)"' <<<"$result") == "true:true" ]] || fail "Kimi collector refreshes expired tokens and persists the rotated pair" "$result" pass "Kimi collector refreshes expired tokens and persists the rotated pair" @@ -147,5 +172,9 @@ pass "Kimi collector renders usage windows into display rows" fail "Kimi collector prints the display-ready record contract" "$result" pass "Kimi collector prints the display-ready record contract" +[[ $(jq -r '.planLabel' <<<"$result") == "true" ]] || + fail "Kimi collector exposes the configured plan label" "$result" +pass "Kimi collector exposes the configured plan label" + [[ $(jq -r '.quotaOnly' <<<"$result") == "true" ]] || fail "Kimi quota is not local usage statistics" "$result" pass "Kimi quota is not local usage statistics"