Add Kimi Code and Grok Build support to model usage - #6477
Open
dalmasluca wants to merge 7 commits into
Open
Conversation
added 3 commits
August 1, 2026 10:10
The model-usage widget tracked Claude Code and Codex. This adds the two
providers the AI Usage plugin ships on top: Kimi Code and Grok Build.
- providers/Kimi.qml and providers/Grok.qml: quota-only adapters that
surface a flat list of {label, percent, resetsAt} limits instead of the
single short/long window pair Claude and Codex report.
- scripts/kimi_usage.py and scripts/grok_usage.py: read the CLIs' own
credential stores and OAuth flows, so the plugin and the CLIs never
invalidate each other's refresh tokens. Stdlib-only; tokens stay in
memory and are never printed or logged.
- Panel.qml renders any number of limits per provider and keeps the
providers' own window labels (5h limit, 30d limit, On-demand cap).
- Kimi and Grok have no local token history, so they self-hide until a
scan finds numbers, exactly like the local-stat providers.
- Scanner tests spin up a local HTTP server against real recorded payload
shapes and cover the not-logged-in path.
The Kimi and Grok scanners now also walk the CLIs' local session files and report the same token history the Claude and Codex providers do: per-day and per-model totals, prompt and session counts, and active-day ranges. Kimi reads the per-turn usage.record events in ~/.kimi-code/sessions wire transcripts, which carry the full input/output/cache split. Grok derives turn totals from the cumulative totalTokens in ~/.grok/sessions updates with the model id from the matching turn_started event; the CLI records no input/output/cache split, so the turn total lands in inputTokens (the same fallback the Codex scanner uses). Both cap the scan to transcripts touched in the last 30 days.
Contributor
There was a problem hiding this comment.
Pull request overview
Adds Kimi Code and Grok Build quota and local token-history support to the model-usage plugin.
Changes:
- Adds quota and transcript scanners for both providers.
- Integrates providers into the panel, synchronization, and defaults.
- Adds scanner tests and documentation.
Tip
If you aren't ready for review, convert to a draft PR.
Click "Convert to draft" or run gh pr ready --undo.
Click "Ready for review" or run gh pr ready to reengage.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
shell/plugins/model-usage/Main.qml |
Registers and exposes the new providers. |
shell/plugins/model-usage/Panel.qml |
Displays multi-window quotas and provider headers. |
shell/plugins/model-usage/README.md |
Documents Kimi and Grok support. |
shell/plugins/model-usage/manifest.json |
Enables both providers by default. |
shell/plugins/model-usage/providers/Kimi.qml |
Implements the Kimi provider bridge. |
shell/plugins/model-usage/providers/Grok.qml |
Implements the Grok provider bridge. |
shell/plugins/model-usage/scripts/kimi_usage.py |
Fetches Kimi quotas and scans local history. |
shell/plugins/model-usage/scripts/grok_usage.py |
Fetches Grok quotas and scans local history. |
test/shell.d/model-usage-kimi-scanner-test.sh |
Tests Kimi quota and history parsing. |
test/shell.d/model-usage-grok-scanner-test.sh |
Tests Grok quota and history parsing. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+86
to
+89
| for key, entry in store.items(): | ||
| if isinstance(entry, dict) and entry.get("key"): | ||
| return key, entry | ||
| return None, None |
| "token_type": payload.get("token_type") or token.get("token_type") or "Bearer", | ||
| "expires_in": int(expires_in), | ||
| }) | ||
| save_token(merged) |
Comment on lines
+109
to
+111
| // Only a successful fetch replaces the limits; a transient network | ||
| // failure keeps the last good quota on screen and just reports it. | ||
| if (data.available === true) { |
Comment on lines
+110
to
+112
| // Only a successful fetch replaces the limits; a transient network | ||
| // failure keeps the last good quota on screen and just reports it. | ||
| if (data.available === true) { |
Comment on lines
+382
to
+385
| // Kimi and Grok have no shipped mark yet; collapse the icon | ||
| // slot so the hero reads as a clean two-line header. | ||
| width: source === "" ? 0 : Style.font.display | ||
| height: source === "" ? 0 : Style.font.display |
Author
|
Addressed the review comments (commits 358be43, 3119100, 85665dc):
|
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 11 out of 11 changed files in this pull request and generated 1 comment.
Suppressed comments (4)
shell/plugins/model-usage/scripts/grok_usage.py:582
- The repeated-401 path raises
RuntimeError("authorization failed..."), so it falls into thiselseand is marked transient.Grok.qmlconsequently keeps displaying the last quota indefinitely even though authentication is terminal. Treat the explicit authorization error (and invalid-grant HTTP 400 responses from refresh) as non-transient.
if isinstance(exc, urllib.error.HTTPError) and exc.code in (401, 403):
message = "authorization failed (try: grok login)"
else:
# Network, timeout or server error: the last good quota is still valid.
transient = True
shell/plugins/model-usage/Panel.qml:385
Image.sourceis a QMLurlvalue, so strict comparison with the string""does not reliably detect an empty URL. For Kimi/Grok this can leave the image at display size, makingPanelHero.hasIcontrue and preserving the blank icon gap this change intends to remove. Convert the URL to a string before comparing.
width: source === "" ? 0 : Style.font.display
height: source === "" ? 0 : Style.font.display
shell/plugins/model-usage/scripts/kimi_usage.py:546
- An expired/revoked refresh token commonly produces OAuth
invalid_grantas HTTP 400, but this branch marks every status except 401 as transient. The provider then preserves a stale quota instead of clearing it and asking the user to log in again. Classify terminal OAuth client errors separately from network, rate-limit, and server failures.
if isinstance(exc, urllib.error.HTTPError) and exc.code == 401:
message = "authorization failed (try: kimi login)"
else:
# Network, timeout or server error: the last good quota is still valid.
transient = True
shell/plugins/model-usage/scripts/grok_usage.py:5
- Correct the source/project name typo to “grok-build.”
gork-build crates/codegen/xai-grok-shell/src/extensions/billing.rs):
Comment on lines
+159
to
+164
| fresh = store.get(entry_key) | ||
| if not (isinstance(fresh, dict) and fresh.get("key")): | ||
| fresh = current | ||
| if not force and not token_expired(fresh): | ||
| return fresh # the CLI refreshed while we waited | ||
| fresh = refresh_entry(fresh) |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 11 out of 15 changed files in this pull request and generated no new comments.
Suppressed comments (6)
shell/plugins/model-usage/scripts/grok_usage.py:5
- Correct the referenced project name from
gork-buildtogrok-build.
gork-build crates/codegen/xai-grok-shell/src/extensions/billing.rs):
shell/plugins/model-usage/scripts/kimi_usage.py:541
- An expired or revoked refresh token normally makes the OAuth endpoint return HTTP 400 (
invalid_grant). This condition treats that as a transient network failure, soKimi.qmlpreserves the last quota indefinitely instead of clearing stale data and asking the user to log in again. Treat OAuth/authentication 4xx responses as terminal while retaining transient handling for transport errors, 429s, and 5xx responses.
if isinstance(exc, urllib.error.HTTPError) and exc.code == 401:
message = "authorization failed (try: kimi login)"
else:
# Network, timeout or server error: the last good quota is still valid.
transient = True
shell/plugins/model-usage/scripts/grok_usage.py:578
- Authentication failures are misclassified as transient here. A failed refresh commonly returns HTTP 400, and the explicit second-401 path above raises
RuntimeError("authorization failed..."); both enter theelsebranch, causingGrok.qmlto retain stale quota after authentication is definitively invalid. Classify these auth outcomes as terminal and reservetransientfor transport errors, 429s, and 5xx responses.
if isinstance(exc, urllib.error.HTTPError) and exc.code in (401, 403):
message = "authorization failed (try: grok login)"
else:
# Network, timeout or server error: the last good quota is still valid.
transient = True
test/shell.d/model-usage-kimi-scanner-test.sh:24
- This fixture deliberately bypasses the newly added OAuth refresh, locking, rotation, and atomic credential-write path. That path modifies the CLI's live credential store and is a core part of this scanner, so add a test with an expired token that serves a refresh response and verifies the refreshed token is persisted with the expected mode and used for
/usages; a concurrent-refresh case should also verify that a token rotated under the lock is reused.
# The CLI's FileTokenStorage shape: a non-expired token so the scanner goes
# straight to the /usages endpoint instead of the OAuth refresh flow.
test/shell.d/model-usage-grok-scanner-test.sh:24
- This fixture deliberately bypasses the newly added OAuth refresh, advisory lock, token rotation, and atomic
auth.jsonupdate. Because this scanner writes the CLI's credential store, add coverage with an expired entry that serves a refresh response and verifies persistence, permissions, and use of the refreshed token; also cover the forced refresh after a billing 401, which would expose the transient-classification bug.
# The CLI's auth.json shape: an entry with a non-expired token so the scanner
# goes straight to the billing endpoint instead of the OAuth refresh flow.
shell/Ui/PanelHero.qml:43
- This changes the shared hero geometry used by multiple panels, while the PR reports only scanner tests and there is no model-usage or
PanelHeroacceptance coverage. Verify the iconless and icon-present states in the running shell, including the new four-provider switch and both light/dark icons, and inspect screenshots for spacing, clipping, and overlap before merging.
anchors.left: root.hasIcon ? iconLoader.right : parent.left
anchors.leftMargin: root.hasIcon ? Style.space(14) : 0
7 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds official quota and per-model local token history for Kimi Code and Grok Build to the
omarchy.model-usageplugin.scripts/kimi_usage.pyreads the Kimi Code/usagesendpoint (weekly + rolling 5h/7d/30d windows), with token refresh through the CLI's own OAuth flow.scripts/grok_usage.pyreads the Grok Buildbilling?format=creditsendpoint (weekly/monthly credit pool + on-demand cap).~/.kimi-code/sessions/*/session_*/agents/*/wire.jsonl,usage.record) split by model.~/.grok/sessions,turn_startedevents + cumulativetotalTokensinupdates.jsonl).~/.pi/agent/sessions) for thekimi-codingandxaiproviders, merged into the same per-model buckets.Kimi.qml/Grok.qml, wiring inMain.qml/Panel.qml,manifest.json(enabled by default), README.test/shell.d/model-usage-kimi-scanner-test.shandmodel-usage-grok-scanner-test.sh(9 assertions each, all green).Notes
Tests