Improve Kimi plan quota reporting

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)

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))

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 `<home>/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.
 

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(

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"

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"