diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml index 613e1fe..ac8a2b5 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.yml +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -15,7 +15,7 @@ body: - OpenRouter - Warp - z.ai - - Kimi K2 + - Kimrel - Ollama - Multiple / All validations: diff --git a/AGENTS.md b/AGENTS.md index f4ced3a..96ab2ca 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -162,13 +162,18 @@ exhaustive provider coverage. Restore a dropped topic via Stream Deck UI. The SVG owns the value, glyph, and ratio fill; the title bar owns the label text. Send labels in UPPERCASE (`SESSION`, `WEEKLY`, …) to match the title font's expected look. -- Provider button glyphs live in `internal/icons/`. Most come from - lobehub/lobe-icons (MIT) — `internal/icons/lobe_generated.go` is - produced by `go run scripts/sync-lobe-icons.go`. Edit the mapping - table in that script and re-run to add a provider, change a variant - (`mono` vs wordmark `text`), or refresh after upstream changes. The - remaining hand-drawn marks (warp, factory, abacus, augment, - jetbrains, kiro, opencodego, synthetic) live in their own +- Provider button glyphs live in `internal/icons/`. **Source of truth + for new provider glyphs is lobehub/lobe-icons (MIT).** Add the entry + to `scripts/sync-lobe-icons.go` and re-run `go run + scripts/sync-lobe-icons.go` to regenerate `internal/icons/lobe_generated.go`. + Do not hand-draw a new mark when a Lobe icon exists for that brand, + and do not borrow another provider's mark when one doesn't (instead + use a neutral letterform — see `internal/icons/kimrel.go`). Action- + picker SVGs in `io.github.anthonybaldwin.UsageButtons.sdPlugin/assets/` + must use the same path data as the in-Go glyph so the SD store icon + matches what the button renders. The remaining hand-drawn marks + (warp, factory, abacus, augment, jetbrains, kiro, opencodego, + synthetic) predate the lobe-icons sync and stay in their own `.go` files alongside the trimmed `icons.go` literal. ## Browser fetch bridge (Usage Buttons Helper extension) diff --git a/README.md b/README.md index 15d1171..5fb0e1b 100644 --- a/README.md +++ b/README.md @@ -17,7 +17,7 @@ Runs on **Windows and macOS**. ## Settings -Each provider is its own action — drag **Claude Code**, **Codex**, **Gemini**, +Each provider is its own action — drag **Claude**, **Codex**, **Gemini**, **Vertex AI**, etc. onto a key and configure the metric, colors, and thresholds from the Property Inspector. @@ -65,7 +65,7 @@ UsageButtons/ │ │ ├── amp/ # Amp (browser) │ │ ├── antigravity/ # Antigravity (local language server) │ │ ├── augment/ # Augment (CLI/browser) -│ │ ├── claude/ # Claude Code (Claude OAuth + browser web API) +│ │ ├── claude/ # Claude (Claude OAuth + browser web API) │ │ ├── codex/ # Codex (OAuth) │ │ ├── cookieaux/ # cookie-gated provider messaging helpers │ │ ├── copilot/ # GitHub Copilot @@ -78,8 +78,8 @@ UsageButtons/ │ │ ├── nousresearch/ # Nous Research portal (Hermes / Nous Chat, browser) │ │ ├── jetbrains/ # JetBrains AI │ │ ├── kilo/ # Kilo -│ │ ├── kimi/ # Kimi (browser) -│ │ ├── kimik2/ # Kimi K2 (API key) +│ │ ├── kimi/ # Kimi (browser, OAuth fallback) +│ │ ├── kimrel/ # Kimrel — third-party Kimi K2 reseller (API key) │ │ ├── kiro/ # Kiro │ │ ├── minimax/ # MiniMax (browser/API key) │ │ ├── mistral/ # Mistral (browser) @@ -119,7 +119,7 @@ Short version: **UsageButtons-Helper-unpacked.zip** from the same release, unzip it, and **Load unpacked** in `chrome://extensions`. The plugin auto-registers — nothing to configure. -3. Drag a provider (**Claude Code**, **Codex**, **Copilot**, etc.) onto a +3. Drag a provider (**Claude**, **Codex**, **Copilot**, etc.) onto a Stream Deck key and pick a metric from the Property Inspector. ## Build from source @@ -176,7 +176,7 @@ through your real browser session; cookies never leave Chrome. `chrome://extensions`, done. - **Providers that don't need it keep working unchanged** — Gemini, Vertex AI, Copilot, OpenRouter, DeepSeek, Moonshot, Warp, z.ai, - Kimi K2, Synthetic, Kilo, Kiro, JetBrains AI, Anthropic, OpenAI, + Kimrel, Synthetic, Kilo, Kiro, JetBrains AI, Anthropic, OpenAI, and Antigravity never require the extension. - **Waits patiently on cold start.** Cookie-gated buttons stay in a quiet "needs browser extension" state until the extension diff --git a/chrome-extension/manifest.json b/chrome-extension/manifest.json index 61c9841..d02d620 100644 --- a/chrome-extension/manifest.json +++ b/chrome-extension/manifest.json @@ -2,7 +2,7 @@ "manifest_version": 3, "name": "Usage Buttons Helper", "short_name": "UsageButtons", - "version": "0.8.1", + "version": "0.8.0", "description": "Companion extension for the Usage Buttons Stream Deck plugin. Proxies a narrow allowlist of AI usage APIs (abacus.ai, alibabacloud.com, aliyun.com, claude.ai, cursor.com, factory.ai, ollama.com, chatgpt.com, augmentcode.com, ampcode.com, perplexity.ai, grok.com, nousresearch.com, opencode.ai, kimi.com, minimax.io, minimaxi.com, mistral.ai, deepseek.com) using your logged-in browser session — credentials stay in Chrome.", "author": "Anthony Baldwin", "minimum_chrome_version": "127", diff --git a/chrome-extension/service-worker.js b/chrome-extension/service-worker.js index 1dd40e0..2bba287 100644 --- a/chrome-extension/service-worker.js +++ b/chrome-extension/service-worker.js @@ -232,6 +232,46 @@ async function readDeepSeekPlatformToken() { return null; } +// readKimiAccessToken reads localStorage["access_token"] from any open +// kimi.com tab. Kimi migrated off cookie-based auth — the kimi-auth +// JWT cookie is no longer set, and apiv2 endpoints reject cookie-only +// requests with `REASON_INVALID_AUTH_TOKEN`. The session bearer token +// now lives in localStorage and the page's own client passes it via +// `Authorization: Bearer ` on every API call. +// +// Returns null when no kimi.com tab is open or the user is signed out; +// the Go side then falls back to OAuth credentials placed by the +// `kimi login` CLI. +async function readKimiAccessToken() { + let tabs; + try { + tabs = await chrome.tabs.query({ url: "*://*.kimi.com/*" }); + } catch (_e) { + return null; + } + if (!tabs || tabs.length === 0) return null; + for (const t of tabs) { + try { + const results = await chrome.scripting.executeScript({ + target: { tabId: t.id }, + // MAIN world reaches the page's own localStorage rather than + // the isolated-world copy a content script would see. + world: "MAIN", + func: () => localStorage.getItem("access_token") || "", + }); + const raw = results && results[0] && results[0].result; + if (raw && typeof raw === "string" && raw.length > 0) { + return raw; + } + } catch (_e) { + // executeScript can fail if the tab is loading or in a special state; + // try the next one. + continue; + } + } + return null; +} + // augmentHeadersForOrigin attaches site-specific auth/version headers // for hosts whose internal APIs require explicit non-cookie auth. // Returns the final headers object to pass to fetch(). For everything @@ -252,6 +292,16 @@ async function augmentHeadersForOrigin(url, callerHeaders) { headers["x-app-version"] = DEEPSEEK_PLATFORM_APP_VERSION; } } + + if (host === "kimi.com" || host === "www.kimi.com" || host.endsWith(".kimi.com")) { + if (!hasHeader(headers, "authorization")) { + const tok = await readKimiAccessToken(); + if (tok) { + headers["Authorization"] = "Bearer " + tok; + } + } + } + return headers; } diff --git a/cmd/genkeys/main.go b/cmd/genkeys/main.go index 2fce197..cc19f3c 100644 --- a/cmd/genkeys/main.go +++ b/cmd/genkeys/main.go @@ -37,7 +37,8 @@ var providerColors = map[string]string{ "vertexai": "#4285f4", "warp": "#938bb4", "zai": "#e85a6a", - "kimi-k2": "#4c00ff", + "kimi-k2": "#64748b", + "moonshot": "#0a84ff", } func main() { diff --git a/cmd/plugin/main.go b/cmd/plugin/main.go index db4cd0e..0cf14da 100644 --- a/cmd/plugin/main.go +++ b/cmd/plugin/main.go @@ -26,6 +26,7 @@ import ( "github.com/anthonybaldwin/UsageButtons/internal/settings" "github.com/anthonybaldwin/UsageButtons/internal/streamdeck" "github.com/anthonybaldwin/UsageButtons/internal/update" + "github.com/anthonybaldwin/UsageButtons/internal/wsl" // Register all providers via init(). _ "github.com/anthonybaldwin/UsageButtons/internal/providers/abacus" @@ -45,7 +46,7 @@ import ( _ "github.com/anthonybaldwin/UsageButtons/internal/providers/jetbrains" _ "github.com/anthonybaldwin/UsageButtons/internal/providers/kilo" _ "github.com/anthonybaldwin/UsageButtons/internal/providers/kimi" - _ "github.com/anthonybaldwin/UsageButtons/internal/providers/kimik2" + _ "github.com/anthonybaldwin/UsageButtons/internal/providers/kimrel" _ "github.com/anthonybaldwin/UsageButtons/internal/providers/kiro" _ "github.com/anthonybaldwin/UsageButtons/internal/providers/minimax" _ "github.com/anthonybaldwin/UsageButtons/internal/providers/mistral" @@ -741,9 +742,32 @@ func handleSendToPlugin(conn *streamdeck.Connection, ev streamdeck.Event) { go replyUnregisterCookieHost(conn, ev.Context, ev.Action) case "getProviderStatus": go replyProviderStatus(conn, ev.Context, ev.Action) + case "getWSLDistros": + go replyWSLDistros(conn, ev.Context, ev.Action) } } +// replyWSLDistros tells the PI which WSL distributions are currently +// running. The PI uses the response to inject extra "WSL: " +// entries into the metric dropdown for cost-tile-bearing providers. +// +// On non-Windows builds wsl.Sources() returns nil and the PI sees an +// empty list, so nothing changes in the UI. +func replyWSLDistros(conn *streamdeck.Connection, ctxStr, action string) { + sources := wsl.Sources() + distros := make([]map[string]string, 0, len(sources)) + for _, s := range sources { + distros = append(distros, map[string]string{ + "key": s.Key, + "label": s.Label, + }) + } + conn.SendToPropertyInspector(ctxStr, action, map[string]any{ + "action": "wslDistros", + "distros": distros, + }) +} + // cookieHostPayload is the PI → plugin shape for registerCookieHost. func replyCookieStatus(conn *streamdeck.Connection, ctxStr, action string) { pctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) diff --git a/docs/PROVIDERS.md b/docs/PROVIDERS.md index b26c5c4..877a7c0 100644 --- a/docs/PROVIDERS.md +++ b/docs/PROVIDERS.md @@ -32,12 +32,12 @@ account or API response includes that quota lane. | Nous Research | Usage Buttons Helper from `portal.nousresearch.com`. | Subscription credits ($, Hermes Agent + Nous Chat pool), API credits balance ($), all-time totals (spend $, requests, tokens, input/output/cache-read/cache-write tokens) — combined or split by allowance (api / sub). | | JetBrains AI | Local JetBrains IDE quota files. Optional overrides: `CODEXBAR_JETBRAINS_IDE_BASE_PATH` or `JETBRAINS_QUOTA_FILE`. | Current credits remaining %. | | Kilo | Kilo API key from the Provider tab, `KILO_API_KEY`, or `~/.local/share/kilo/auth.json`. | Credits remaining %, Kilo Pass remaining %. | -| Kimi | Usage Buttons Helper from `kimi.com`. | Weekly coding quota remaining %, 5-hour rate limit remaining %. | -| Kimi K2 | Kimi K2 API key from the Provider tab or `KIMI_K2_API_KEY`. | Credits remaining. | +| Kimi | Usage Buttons Helper from `kimi.com` (preferred). Falls back to OAuth credentials placed by `kimi login` at `~/.kimi/credentials/kimi-code.json` when the Helper isn't connected; refresh tokens are exchanged against `auth.kimi.com`. | Weekly coding quota remaining %, 5-hour rate limit remaining %. | +| Kimrel | Kimrel API key from the Provider tab or `KIMREL_API_KEY` (older `KIMI_K2_API_KEY` / `KIMI_API_KEY` / `KIMI_KEY` still resolve). Kimrel (kimrel.com, formerly kimi-k2.ai) is an **independent third-party reseller** of Kimi K2 model access — not affiliated with, endorsed by, or sponsored by Moonshot AI. Use the Moonshot provider for the official Moonshot dev platform. | Credits remaining. | | Kiro | `kiro-cli`; run `kiro-cli login` first. | Monthly credits remaining %, bonus credits remaining %. | | MiniMax | MiniMax API key from the Provider tab / `MINIMAX_API_KEY`, or Usage Buttons Helper from `minimax.io`. Optional region override: `MINIMAX_REGION`. | Coding prompts remaining %. | | Mistral | Usage Buttons Helper from `admin.mistral.ai`. | Monthly billing usage. | -| Moonshot (Kimi platform) | Moonshot API key from the Provider tab or `MOONSHOT_API_KEY` / `KIMI_PLATFORM_API_KEY`. Optional China-region host override: `MOONSHOT_API_HOST=https://api.moonshot.cn`. | Available balance ($), voucher balance ($), cash balance ($). Distinct from the Kimi (chat) provider — Moonshot is the paid developer API platform. | +| Moonshot (Kimi platform) | Moonshot API key from the Provider tab or `MOONSHOT_API_KEY` / `KIMI_PLATFORM_API_KEY`. Optional China-region host override: `MOONSHOT_API_HOST=https://api.moonshot.cn`. | Available balance ($), voucher balance ($), cash balance ($). Distinct from the Kimi provider (per-user CLI quotas) and Kimrel (third-party reseller credits) — Moonshot is the org-wide paid developer API balance. | | Ollama | Usage Buttons Helper from the signed-in Ollama web session. | Session usage remaining %, session pace (burn rate), weekly usage remaining %, weekly pace (burn rate). | | OpenAI | OpenAI admin API key (`sk-admin-…`) from the Provider tab or `OPENAI_ADMIN_API_KEY` (kept namespaced so it doesn't shadow the SDK-standard `OPENAI_API_KEY`). Org admins only — personal `sk-` keys are rejected by the admin endpoints. | Org spend today (UTC, $), yesterday ($), last 7 days ($), month-to-date ($), last 30 days ($), 7-day burn rate ($/day), projected month total ($). Distinct from the Codex provider (per-user session/weekly window from ChatGPT OAuth) — this is the org-wide cost view. | | OpenCode | Usage Buttons Helper from `opencode.ai`. Optional workspace override: `CODEXBAR_OPENCODE_WORKSPACE_ID`. | 5-hour usage remaining %, weekly usage remaining %. | diff --git a/docs/index.html b/docs/index.html index 953ac60..4701aa1 100644 --- a/docs/index.html +++ b/docs/index.html @@ -479,9 +479,9 @@

Standalone native binary

Supported providers

-

35 providers are live, spanning hosted AI assistants (Claude Code, Codex, Cursor, Copilot, Gemini, Grok), self-hosted gateways (OpenClaw, Hermes Agent, Ollama), portal accounts (Nous Research, Perplexity, Mistral, Kimi, MiniMax, Abacus, Alibaba, Augment, Amp, Droid, OpenCode, JetBrains AI, Antigravity, Kilo, Kiro), and direct API keys (OpenRouter, DeepSeek, Moonshot, Anthropic, OpenAI, Vertex AI, Synthetic, Warp, z.ai, Kimi K2).

+

35 providers are live, spanning hosted AI assistants (Claude, Codex, Cursor, Copilot, Gemini, Grok), self-hosted gateways (OpenClaw, Hermes Agent, Ollama), portal accounts (Nous Research, Perplexity, Mistral, Kimi, MiniMax, Abacus, Alibaba, Augment, Amp, Droid, OpenCode, JetBrains AI, Antigravity, Kilo, Kiro), and direct API keys (OpenRouter, DeepSeek, Moonshot, Anthropic, OpenAI, Vertex AI, Synthetic, Warp, z.ai, Kimrel).

- Claude Code + Claude Anthropic Codex OpenAI @@ -513,7 +513,7 @@

Supported providers

Synthetic Warp z.ai - Kimi K2 + Kimrel Ollama Kiro
diff --git a/docs/provider-icons/kimi-k2.svg b/docs/provider-icons/kimi-k2.svg index 0948f79..80e90bc 100644 --- a/docs/provider-icons/kimi-k2.svg +++ b/docs/provider-icons/kimi-k2.svg @@ -1,3 +1,3 @@ - - \ No newline at end of file + + diff --git a/docs/provider-icons/moonshot.svg b/docs/provider-icons/moonshot.svg index 7142857..206e769 100644 --- a/docs/provider-icons/moonshot.svg +++ b/docs/provider-icons/moonshot.svg @@ -1,5 +1,5 @@ - - + + diff --git a/internal/icons/kimrel.go b/internal/icons/kimrel.go new file mode 100644 index 0000000..c3c102f --- /dev/null +++ b/internal/icons/kimrel.go @@ -0,0 +1,22 @@ +package icons + +import "github.com/anthonybaldwin/UsageButtons/internal/render" + +// Kimrel (provider ID "kimi-k2") is a third-party reseller of Kimi K2 +// model access — not affiliated with Moonshot AI. We deliberately do +// not reuse Lobe's Kimi mark (which the auto-generated lobe_generated.go +// would otherwise pull in for this provider ID). +// +// The mark below is a plain bold sans-serif "K" — visually distinct +// from Kimi's flower-K and from Moonshot's lunar mark, signalling +// "credits balance" without implying official Moonshot/Kimi branding. +// +// File name "kimrel.go" sorts after "lobe_generated.go" so its init +// runs second and overwrites any stale entry the generator might +// re-emit on a future sync. +func init() { + ProviderIcons["kimi-k2"] = &render.ProviderGlyph{ + ViewBox: "0 0 24 24", + D: `M5 4h2.6v7.7L14 4h3.2l-6.2 7.4L17.7 20h-3.2L9.6 13.2 7.6 15.5V20H5V4z`, + } +} diff --git a/internal/icons/lobe_generated.go b/internal/icons/lobe_generated.go index dd9e6c5..054e6cc 100644 --- a/internal/icons/lobe_generated.go +++ b/internal/icons/lobe_generated.go @@ -88,14 +88,8 @@ func init() { {D: `M11.065 11.199l7.257-7.2c.137-.136.06-.41-.116-.41H14.3a.164.164 0 00-.117.051l-7.82 7.756c-.122.12-.302.013-.302-.179V3.82c0-.127-.083-.23-.185-.23H3.186c-.103 0-.186.103-.186.23V19.77c0 .128.083.23.186.23h2.69c.103 0 .186-.102.186-.23v-3.25c0-.069.025-.135.069-.178l2.424-2.406a.158.158 0 01.205-.023l6.484 4.772a7.677 7.677 0 003.453 1.283c.108.012.2-.095.2-.23v-3.06c0-.117-.07-.212-.164-.227a5.028 5.028 0 01-2.027-.807l-5.613-4.064c-.117-.078-.132-.279-.028-.381z`}, }, } - // kimi-k2 — lobe-icons kimi.svg (mono) - ProviderIcons["kimi-k2"] = &render.ProviderGlyph{ - ViewBox: "0 0 24 24", - Paths: []render.GlyphPath{ - {D: `M21.846 0a1.923 1.923 0 110 3.846H20.15a.226.226 0 01-.227-.226V1.923C19.923.861 20.784 0 21.846 0z`}, - {D: `M11.065 11.199l7.257-7.2c.137-.136.06-.41-.116-.41H14.3a.164.164 0 00-.117.051l-7.82 7.756c-.122.12-.302.013-.302-.179V3.82c0-.127-.083-.23-.185-.23H3.186c-.103 0-.186.103-.186.23V19.77c0 .128.083.23.186.23h2.69c.103 0 .186-.102.186-.23v-3.25c0-.069.025-.135.069-.178l2.424-2.406a.158.158 0 01.205-.023l6.484 4.772a7.677 7.677 0 003.453 1.283c.108.012.2-.095.2-.23v-3.06c0-.117-.07-.212-.164-.227a5.028 5.028 0 01-2.027-.807l-5.613-4.064c-.117-.078-.132-.279-.028-.381z`}, - }, - } + // kimi-k2 (Kimrel) — custom mark lives in internal/icons/kimrel.go. + // Kimrel is a third-party reseller; not allowed to share Kimi's icon. // minimax — lobe-icons minimax.svg (mono) ProviderIcons["minimax"] = &render.ProviderGlyph{ ViewBox: "0 0 24 24", diff --git a/internal/providers/claude/costs.go b/internal/providers/claude/costs.go index 6a6ff68..6b3866a 100644 --- a/internal/providers/claude/costs.go +++ b/internal/providers/claude/costs.go @@ -12,14 +12,16 @@ import ( "time" "github.com/anthonybaldwin/UsageButtons/internal/providers" + "github.com/anthonybaldwin/UsageButtons/internal/wsl" ) // Cost scan cache — rescan at most once per 5 minutes. var ( // costMu guards costCache and costCacheT. costMu sync.Mutex - // costCache holds the most recent result from scanCosts. - costCache *costResult + // costCache holds the most recent result from scanCosts (covering + // both the Windows-native projects dir and any running WSL distros). + costCache *allCostResults // costCacheT is the wall-clock time of the last successful scan. costCacheT time.Time // unpricedModelMu guards unpricedModelCounts. @@ -115,9 +117,28 @@ type costResult struct { Last30d float64 } -// scanCosts walks ~/.claude/projects and returns aggregated today/30-day -// spend estimates, memoized for costCacheTTL. -func scanCosts() (*costResult, error) { +// allCostResults holds the costResult for the Windows-native projects +// dir plus, on Windows builds with WSL installed and distros running, +// one costResult per running distro keyed by Source.Key. +// +// Each distro is treated as its own "machine" — no aggregation with +// the native scope, since the user explicitly wants them visible +// separately rather than silently summed. +type allCostResults struct { + Native costResult + // WSL is keyed by wsl.Source.Key. Empty/nil when WSL is unavailable + // or no distros are running. + WSL map[string]costResult + // WSLLabels maps Source.Key → friendly distro name (e.g. + // "Ubuntu-22.04") so the metric Caption can identify each scope + // without re-running discovery. + WSLLabels map[string]string +} + +// scanCosts walks ~/.claude/projects on the Windows host plus the +// equivalent path inside every running WSL distro, returning the +// per-scope aggregates memoized for costCacheTTL. +func scanCosts() (*allCostResults, error) { costMu.Lock() defer costMu.Unlock() if costCache != nil && time.Since(costCacheT) < costCacheTTL { @@ -125,25 +146,58 @@ func scanCosts() (*costResult, error) { } resetUnpricedModels() + now := time.Now() + todayStart := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()) + thirtyDaysAgo := now.AddDate(0, 0, -30) + + out := &allCostResults{} + home, err := os.UserHomeDir() if err != nil { return nil, err } - projectsDir := filepath.Join(home, ".claude", "projects") - - entries, err := os.ReadDir(projectsDir) + native, err := scanProjectsDir(filepath.Join(home, ".claude", "projects"), todayStart, thirtyDaysAgo) if err != nil { - if os.IsNotExist(err) { - return &costResult{}, nil - } return nil, err } + out.Native = native + + // wsl.Sources() is a no-op on non-Windows builds and returns nil + // when WSL isn't installed or no distros are running, so this loop + // degrades cleanly to zero extra work in the common case. + if sources := wsl.Sources(); len(sources) > 0 { + out.WSL = make(map[string]costResult, len(sources)) + out.WSLLabels = make(map[string]string, len(sources)) + for _, src := range sources { + r, err := scanProjectsDir(filepath.Join(src.Home, ".claude", "projects"), todayStart, thirtyDaysAgo) + if err != nil { + // One unreachable distro shouldn't poison the whole + // scan; just skip it and keep going. + continue + } + out.WSL[src.Key] = r + out.WSLLabels[src.Key] = src.Label + } + } - now := time.Now() - todayStart := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()) - thirtyDaysAgo := now.AddDate(0, 0, -30) + costCache = out + costCacheT = time.Now() + return out, nil +} +// scanProjectsDir walks one Claude projects directory (on any +// filesystem — native or \\wsl.localhost\...) and returns the +// today/30d cost aggregate. Missing directories return a zero result +// without error so callers can scan optional scopes safely. +func scanProjectsDir(projectsDir string, todayStart, thirtyDaysAgo time.Time) (costResult, error) { var result costResult + entries, err := os.ReadDir(projectsDir) + if err != nil { + if os.IsNotExist(err) { + return result, nil + } + return result, err + } for _, project := range entries { if !project.IsDir() { @@ -158,7 +212,7 @@ func scanCosts() (*costResult, error) { if f.IsDir() || !strings.HasSuffix(f.Name(), ".jsonl") { continue } - // Quick filter: skip files older than 30 days by mod time + // Quick filter: skip files older than 30 days by mod time. info, err := f.Info() if err != nil || info.ModTime().Before(thirtyDaysAgo) { continue @@ -166,10 +220,7 @@ func scanCosts() (*costResult, error) { scanFile(filepath.Join(projPath, f.Name()), todayStart, thirtyDaysAgo, &result) } } - - costCache = &result - costCacheT = time.Now() - return &result, nil + return result, nil } // scanFile parses one session .jsonl file and accumulates token costs @@ -339,8 +390,11 @@ func ptrFloat(v float64) *float64 { return &v } -// costMetrics renders the scanned spend into MetricValue tiles for today -// and the trailing 30 days. +// costMetrics renders the scanned spend into MetricValue tiles for the +// trailing day and 30 days. The Windows-native scope produces the +// stable "cost-today" / "cost-30d" IDs; each running WSL distro +// produces a parallel pair suffixed with the distro key +// (e.g. "cost-today-wsl-Debian"), treated as a separate machine. func costMetrics() []providers.MetricValue { result, err := scanCosts() if err != nil || result == nil { @@ -350,32 +404,47 @@ func costMetrics() []providers.MetricValue { now := time.Now().UTC().Format(time.RFC3339) var out []providers.MetricValue - if result.Today > 0 || result.Last30d > 0 { - today := math.Round(result.Today*100) / 100 - out = append(out, providers.MetricValue{ - ID: "cost-today", - Label: "TODAY", - Name: "Estimated spend today", - Value: fmt.Sprintf("$%.2f", today), - NumericValue: &today, - NumericUnit: "dollars", - NumericGoodWhen: "low", - Caption: "Cost (local)", - UpdatedAt: now, - }) - - last30 := math.Round(result.Last30d*100) / 100 - out = append(out, providers.MetricValue{ - ID: "cost-30d", - Label: "30 DAYS", - Name: "Estimated spend last 30 days", - Value: fmt.Sprintf("$%.2f", last30), - NumericValue: &last30, - NumericUnit: "dollars", - NumericGoodWhen: "low", - Caption: "Cost (local)", - UpdatedAt: now, - }) + emit := func(scopeSuffix, captionSuffix string, r costResult) { + // Skip empty scopes so a fresh WSL distro with no sessions + // doesn't render a fake $0.00 tile. + if r.Today == 0 && r.Last30d == 0 { + return + } + today := math.Round(r.Today*100) / 100 + last30 := math.Round(r.Last30d*100) / 100 + out = append(out, + providers.MetricValue{ + ID: "cost-today" + scopeSuffix, + Label: "TODAY", + Name: "Estimated spend today" + captionSuffix, + Value: fmt.Sprintf("$%.2f", today), + NumericValue: &today, + NumericUnit: "dollars", + NumericGoodWhen: "low", + Caption: "Cost (local)" + captionSuffix, + UpdatedAt: now, + }, + providers.MetricValue{ + ID: "cost-30d" + scopeSuffix, + Label: "30 DAYS", + Name: "Estimated spend last 30 days" + captionSuffix, + Value: fmt.Sprintf("$%.2f", last30), + NumericValue: &last30, + NumericUnit: "dollars", + NumericGoodWhen: "low", + Caption: "Cost (local)" + captionSuffix, + UpdatedAt: now, + }, + ) + } + + emit("", "", result.Native) + for key, r := range result.WSL { + label := result.WSLLabels[key] + if label == "" { + label = key + } + emit("-wsl-"+key, " (WSL: "+label+")", r) } return out diff --git a/internal/providers/codex/costs.go b/internal/providers/codex/costs.go index a653015..511a8a3 100644 --- a/internal/providers/codex/costs.go +++ b/internal/providers/codex/costs.go @@ -12,6 +12,7 @@ import ( "time" "github.com/anthonybaldwin/UsageButtons/internal/providers" + "github.com/anthonybaldwin/UsageButtons/internal/wsl" ) // Codex CLI writes one JSONL file per session under @@ -32,11 +33,12 @@ import ( var ( // codexCostMu guards the cost cache against concurrent scanners. codexCostMu sync.Mutex - // codexCostCache is the most recent scan result. - codexCostCache *codexCostResult + // codexCostCache is the most recent scan result, including any + // running WSL distros on Windows. + codexCostCache *allCodexCostResults // codexCostCacheT is the time of the most recent scan. codexCostCacheT time.Time - // codexCostCacheErr is the error from the most recent scan, if any. + // codexCostCacheErr is the error from the most recent native scan. codexCostCacheErr error ) @@ -188,8 +190,12 @@ func codexCostLog(format string, args ...any) { } } -// sessionsRoot returns the filesystem root under which Codex writes -// session JSONL files, honoring CODEX_HOME when set. +// sessionsRoot returns the Windows-native filesystem root under which +// Codex writes session JSONL files, honoring CODEX_HOME when set. +// +// WSL distros are scanned via scanCodexCosts using their own home paths; +// CODEX_HOME is intentionally NOT propagated into WSL scopes since each +// distro is treated as a separate machine with its own environment. func sessionsRoot() string { if ch := os.Getenv("CODEX_HOME"); ch != "" { return filepath.Join(ch, "sessions") @@ -201,33 +207,74 @@ func sessionsRoot() string { return filepath.Join(home, ".codex", "sessions") } -// scanCodexCosts walks the session tree and returns aggregated token cost -// estimates for today and the last 30 days, memoized for codexCostCacheTTL. -func scanCodexCosts() (*codexCostResult, error) { +// allCodexCostResults holds the codexCostResult for the Windows-native +// session tree plus, on Windows builds with WSL distros running, one +// codexCostResult per running distro keyed by wsl.Source.Key. Each WSL +// scope is treated as a separate machine — never aggregated with native. +type allCodexCostResults struct { + Native codexCostResult + // WSL is keyed by wsl.Source.Key. Empty/nil when WSL is unavailable + // or no distros are running. + WSL map[string]codexCostResult + // WSLLabels maps Source.Key → friendly distro name for UI use. + WSLLabels map[string]string +} + +// scanCodexCosts walks the Windows-native session tree plus the +// equivalent path inside every running WSL distro, returning per-scope +// aggregates memoized for codexCostCacheTTL. The returned error reflects +// only the native scan; failed WSL scopes are silently skipped so a +// flaky distro can't poison the Windows tile. +func scanCodexCosts() (*allCodexCostResults, error) { codexCostMu.Lock() defer codexCostMu.Unlock() if codexCostCache != nil && time.Since(codexCostCacheT) < codexCostCacheTTL { return codexCostCache, codexCostCacheErr } - root := sessionsRoot() + now := time.Now() + todayStart := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()) + thirtyDaysAgo := now.AddDate(0, 0, -30) + + out := &allCodexCostResults{} + native, nativeErr := scanCodexSessionsTree(sessionsRoot(), todayStart, thirtyDaysAgo) + out.Native = native + + if sources := wsl.Sources(); len(sources) > 0 { + out.WSL = make(map[string]codexCostResult, len(sources)) + out.WSLLabels = make(map[string]string, len(sources)) + for _, src := range sources { + r, err := scanCodexSessionsTree(filepath.Join(src.Home, ".codex", "sessions"), todayStart, thirtyDaysAgo) + if err != nil { + codexCostLog("WSL %s scan failed: %v", src.Label, err) + continue + } + out.WSL[src.Key] = r + out.WSLLabels[src.Key] = src.Label + } + } + + codexCostCacheErr = nativeErr + codexCostCache = out + codexCostCacheT = time.Now() + return codexCostCache, codexCostCacheErr +} + +// scanCodexSessionsTree walks one Codex sessions root (native or +// \\wsl.localhost\...) and returns the aggregated token-cost estimate. +// A missing root returns a zero result without error so callers can +// scan optional scopes safely. +func scanCodexSessionsTree(root string, todayStart, thirtyDaysAgo time.Time) (codexCostResult, error) { + var result codexCostResult if root == "" { - codexCostCacheErr = nil - codexCostCache = &codexCostResult{} - codexCostCacheT = time.Now() - return codexCostCache, nil + return result, nil } if _, err := os.Stat(root); os.IsNotExist(err) { - codexCostCacheErr = nil - codexCostCache = &codexCostResult{} - codexCostCacheT = time.Now() - return codexCostCache, nil + return result, nil + } else if err != nil { + return result, err } - now := time.Now() - todayStart := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, now.Location()) - thirtyDaysAgo := now.AddDate(0, 0, -30) - state := &codexScanState{ byID: make(map[string]string), metaByID: make(map[string]codexSessionMeta), @@ -240,7 +287,7 @@ func scanCodexCosts() (*codexCostResult, error) { } var toScan []fileEntry - err := filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error { + walkErr := filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error { if err != nil { codexCostLog("skipping %s: %v", path, err) return nil @@ -272,17 +319,13 @@ func scanCodexCosts() (*codexCostResult, error) { return nil }) - var result codexCostResult for _, fe := range toScan { if scanErr := scanCodexSessionFile(state, fe.path, fe.meta, todayStart, thirtyDaysAgo, &result); scanErr != nil { codexCostLog("skipping %s: %v", fe.path, scanErr) } } - codexCostCacheErr = err - codexCostCache = &result - codexCostCacheT = time.Now() - return codexCostCache, codexCostCacheErr + return result, walkErr } // readCodexSessionMeta reads the first handful of lines of a JSONL file @@ -748,41 +791,63 @@ func ptrFloat(v float64) *float64 { } // codexCostMetrics returns cost-today + cost-30d metrics built from the -// local session-log scan. Returns nil if no session data was found for -// the last 30 days so the renderer draws a dash instead of a fake $0.00. +// local session-log scan. The Windows-native scope produces the stable +// "cost-today" / "cost-30d" IDs; each running WSL distro produces a +// parallel pair suffixed with the distro key, treated as a separate +// machine. Scopes with no priced rows are silently dropped so empty +// distros don't render a fake $0.00. func codexCostMetrics() []providers.MetricValue { result, err := scanCodexCosts() if err != nil || result == nil { return nil } - if !result.Seen { - return nil - } + now := time.Now().UTC().Format(time.RFC3339) - today := math.Round(result.Today*100) / 100 - last30 := math.Round(result.Last30d*100) / 100 - return []providers.MetricValue{ - { - ID: "cost-today", - Label: "TODAY", - Name: "Estimated Codex spend today (local logs)", - Value: fmt.Sprintf("$%.2f", today), - NumericValue: &today, - NumericUnit: "dollars", - NumericGoodWhen: "low", - Caption: "Cost (local)", - UpdatedAt: now, - }, - { - ID: "cost-30d", - Label: "30 DAYS", - Name: "Estimated Codex spend last 30 days (local logs)", - Value: fmt.Sprintf("$%.2f", last30), - NumericValue: &last30, - NumericUnit: "dollars", - NumericGoodWhen: "low", - Caption: "Cost (local)", - UpdatedAt: now, - }, + var out []providers.MetricValue + + emit := func(scopeSuffix, captionSuffix string, r codexCostResult) { + if !r.Seen { + return + } + today := math.Round(r.Today*100) / 100 + last30 := math.Round(r.Last30d*100) / 100 + out = append(out, + providers.MetricValue{ + ID: "cost-today" + scopeSuffix, + Label: "TODAY", + Name: "Estimated Codex spend today (local logs)" + captionSuffix, + Value: fmt.Sprintf("$%.2f", today), + NumericValue: &today, + NumericUnit: "dollars", + NumericGoodWhen: "low", + Caption: "Cost (local)" + captionSuffix, + UpdatedAt: now, + }, + providers.MetricValue{ + ID: "cost-30d" + scopeSuffix, + Label: "30 DAYS", + Name: "Estimated Codex spend last 30 days (local logs)" + captionSuffix, + Value: fmt.Sprintf("$%.2f", last30), + NumericValue: &last30, + NumericUnit: "dollars", + NumericGoodWhen: "low", + Caption: "Cost (local)" + captionSuffix, + UpdatedAt: now, + }, + ) + } + + emit("", "", result.Native) + for key, r := range result.WSL { + label := result.WSLLabels[key] + if label == "" { + label = key + } + emit("-wsl-"+key, " (WSL: "+label+")", r) + } + + if len(out) == 0 { + return nil } + return out } diff --git a/internal/providers/codex/costs_test.go b/internal/providers/codex/costs_test.go index a185553..5eb1789 100644 --- a/internal/providers/codex/costs_test.go +++ b/internal/providers/codex/costs_test.go @@ -277,7 +277,7 @@ func TestScanCodexCostsSkipsMalformedSessionFiles(t *testing.T) { if err != nil { t.Fatalf("expected malformed session to be skipped, got %v", err) } - if result == nil || !result.Seen || result.Last30d <= 0 { + if result == nil || !result.Native.Seen || result.Native.Last30d <= 0 { t.Fatalf("expected valid session to remain priced, got %+v", result) } } diff --git a/internal/providers/kimi/kimi.go b/internal/providers/kimi/kimi.go index 26d0858..2c5af75 100644 --- a/internal/providers/kimi/kimi.go +++ b/internal/providers/kimi/kimi.go @@ -1,9 +1,16 @@ // Package kimi implements the Kimi usage provider. // -// Auth: Usage Buttons Helper extension with the user's kimi.com browser -// session. There is no manual cookie/JWT paste — the extension is the -// only path so credentials never leave Chrome. -// Endpoint: POST https://www.kimi.com/apiv2/kimi.gateway.billing.v1.BillingService/GetUsages. +// Auth (extension-first): primary path is the Usage Buttons Helper with +// the user's kimi.com browser session, hitting +// POST https://www.kimi.com/apiv2/kimi.gateway.billing.v1.BillingService/GetUsages +// for FEATURE_CODING quota and rate limits. +// +// Fallback: when the Helper isn't connected (or the cookie session +// returns 401/403), credentials placed by the `kimi login` CLI at +// ~/.kimi/credentials/kimi-code.json are used to call +// GET https://api.kimi.com/coding/v1/usages directly. Tokens are +// refreshed against https://auth.kimi.com/api/oauth/token within +// 5 minutes of expiry. See oauth.go. package kimi import ( @@ -84,21 +91,35 @@ func (Provider) MetricIDs() []string { return []string{"session-percent", "weekly-percent"} } -// Fetch returns the latest Kimi usage snapshot. +// Fetch returns the latest Kimi usage snapshot. The Helper +// extension is the preferred path; if it's unavailable, or returns +// 401/403 because the kimi.com session has lapsed, the OAuth credential +// blob written by `kimi login` is used as a fallback. func (Provider) Fetch(_ providers.FetchContext) (providers.Snapshot, error) { ctx, cancel := context.WithTimeout(context.Background(), 25*time.Second) defer cancel() - if !cookies.HostAvailable(ctx) { - return errorSnapshot(cookieaux.MissingMessage("kimi.com")), nil - } - usage, err := fetchWithBrowser(ctx) - if err != nil { + + if cookies.HostAvailable(ctx) { + usage, err := fetchWithBrowser(ctx) + if err == nil { + return snapshotFromUsage(usage), nil + } var httpErr *httputil.Error if errors.As(err, &httpErr) && (httpErr.Status == 401 || httpErr.Status == 403) { + // Cookie session stale — try OAuth before surfacing the + // extension-prompt message. + if usage, oauthErr := fetchWithOAuth(ctx); oauthErr == nil { + return snapshotFromUsage(usage), nil + } return errorSnapshot(cookieaux.StaleMessage("kimi.com")), nil } return errorSnapshot(err.Error()), nil } + + usage, err := fetchWithOAuth(ctx) + if err != nil { + return errorSnapshot("Install the Usage Buttons Helper to pull kimi.com usage, or run `kimi login` to use the OAuth fallback."), nil + } return snapshotFromUsage(usage), nil } diff --git a/internal/providers/kimi/kimi_test.go b/internal/providers/kimi/kimi_test.go new file mode 100644 index 0000000..20282b8 --- /dev/null +++ b/internal/providers/kimi/kimi_test.go @@ -0,0 +1,290 @@ +package kimi + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +// TestParseUsageMapsFeatureCoding verifies the cookie-path parser keeps +// returning weekly + 5-hour rate from a FEATURE_CODING usage entry. +func TestParseUsageMapsFeatureCoding(t *testing.T) { + now := time.Date(2026, 5, 1, 12, 0, 0, 0, time.UTC) + resp := usageResponse{Usages: []usageEntry{{ + Scope: "FEATURE_CODING", + Detail: usageDetail{ + Limit: "100", + Remaining: "74", + ResetTime: "2026-05-08T12:00:00Z", + }, + Limits: []rateLimitEntry{{ + Window: rateWindow{Duration: 300, TimeUnit: "TIME_UNIT_MINUTE"}, + Detail: usageDetail{Limit: "20", Remaining: "5", ResetTime: "2026-05-01T15:00:00Z"}, + }}, + }}} + snap, err := parseUsage(resp, now) + if err != nil { + t.Fatalf("parseUsage error: %v", err) + } + if snap.Weekly.Limit != "100" || snap.Weekly.Remaining != "74" { + t.Fatalf("weekly = %+v", snap.Weekly) + } + if snap.Rate == nil || snap.Rate.Limit != "20" || snap.Rate.Remaining != "5" { + t.Fatalf("rate = %+v", snap.Rate) + } +} + +// TestParseOAuthUsageMapsEnvelope verifies the OAuth-direct envelope +// (top-level usage + limits[]) flattens into the same usageSnapshot +// shape as the cookie path so snapshotFromUsage can be reused. +func TestParseOAuthUsageMapsEnvelope(t *testing.T) { + now := time.Date(2026, 5, 1, 12, 0, 0, 0, time.UTC) + env := oauthUsageEnvelope{ + Usage: &usageDetail{Limit: "100", Used: "26", ResetTime: "2026-05-08T12:00:00Z"}, + Limits: []rateLimitEntry{{ + Window: rateWindow{Duration: 300, TimeUnit: "TIME_UNIT_MINUTE"}, + Detail: usageDetail{Limit: "20", Used: "15", ResetTime: "2026-05-01T15:00:00Z"}, + }}, + } + snap, err := parseOAuthUsage(env, now) + if err != nil { + t.Fatalf("parseOAuthUsage error: %v", err) + } + if snap.Weekly.Limit != "100" || snap.Weekly.Used != "26" { + t.Fatalf("weekly = %+v", snap.Weekly) + } + if snap.Rate == nil || snap.Rate.Used != "15" { + t.Fatalf("rate = %+v", snap.Rate) + } +} + +// TestParseOAuthUsageEmptyResponseErrors guards against a 200 with +// neither a usage block nor a limits[] array — happens when the +// account hasn't been provisioned for Kimi. +func TestParseOAuthUsageEmptyResponseErrors(t *testing.T) { + _, err := parseOAuthUsage(oauthUsageEnvelope{}, time.Now()) + if err == nil { + t.Fatal("expected error for empty envelope, got nil") + } +} + +// TestNeedsRefresh verifies the 5-minute proactive refresh window and +// the missing-token early-out both behave correctly. +func TestNeedsRefresh(t *testing.T) { + now := time.Date(2026, 5, 1, 12, 0, 0, 0, time.UTC) + exp := func(offset time.Duration) *float64 { + v := float64(now.Add(offset).Unix()) + return &v + } + + tests := []struct { + name string + creds oauthCreds + want bool + }{ + {"missing access token", oauthCreds{RefreshToken: "r", ExpiresAt: exp(time.Hour)}, true}, + {"missing expires_at", oauthCreds{AccessToken: "a", RefreshToken: "r"}, true}, + {"within buffer", oauthCreds{AccessToken: "a", RefreshToken: "r", ExpiresAt: exp(2 * time.Minute)}, true}, + {"already expired", oauthCreds{AccessToken: "a", RefreshToken: "r", ExpiresAt: exp(-time.Hour)}, true}, + {"fresh", oauthCreds{AccessToken: "a", RefreshToken: "r", ExpiresAt: exp(time.Hour)}, false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := tt.creds.needsRefresh(now); got != tt.want { + t.Fatalf("needsRefresh() = %v, want %v", got, tt.want) + } + }) + } +} + +// TestLoadOAuthCredsMissingFile guards the user-facing message path so +// `kimi login` is mentioned when credentials don't exist yet. +func TestLoadOAuthCredsMissingFile(t *testing.T) { + withTempCredsDir(t) + _, err := loadOAuthCreds() + if err == nil { + t.Fatal("expected error for missing creds file") + } + if !strings.Contains(err.Error(), "kimi login") { + t.Fatalf("error %q does not mention `kimi login`", err.Error()) + } +} + +// TestLoadOAuthCredsParsesValidFile verifies fractional expires_at, +// optional scope, and optional token_type round-trip cleanly. +func TestLoadOAuthCredsParsesValidFile(t *testing.T) { + dir := withTempCredsDir(t) + body := `{ + "access_token": "AT", + "refresh_token": "RT", + "expires_at": 1769861835.261056, + "scope": "kimi-code", + "token_type": "Bearer" + }` + writeCreds(t, dir, body) + creds, err := loadOAuthCreds() + if err != nil { + t.Fatalf("loadOAuthCreds error: %v", err) + } + if creds.AccessToken != "AT" || creds.RefreshToken != "RT" || creds.Scope != "kimi-code" { + t.Fatalf("creds = %+v", creds) + } + if creds.ExpiresAt == nil || *creds.ExpiresAt < 1.7e9 { + t.Fatalf("expires_at = %+v, want fractional unix seconds", creds.ExpiresAt) + } +} + +// TestRefreshOAuthTokenPersistsResponse checks that a 200 response +// updates access_token, optional refresh_token, expires_at, and +// preserves the existing token_type while writing back atomically. +func TestRefreshOAuthTokenPersistsResponse(t *testing.T) { + dir := withTempCredsDir(t) + writeCreds(t, dir, `{"access_token":"old","refresh_token":"R","token_type":"Bearer"}`) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("Content-Type"); !strings.HasPrefix(got, "application/x-www-form-urlencoded") { + t.Errorf("Content-Type = %q", got) + } + _ = r.ParseForm() + if r.PostForm.Get("client_id") != oauthClientID { + t.Errorf("client_id = %q", r.PostForm.Get("client_id")) + } + if r.PostForm.Get("grant_type") != "refresh_token" { + t.Errorf("grant_type = %q", r.PostForm.Get("grant_type")) + } + if r.PostForm.Get("refresh_token") != "R" { + t.Errorf("refresh_token = %q", r.PostForm.Get("refresh_token")) + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{ + "access_token":"NEW", + "expires_in": 3600, + "scope": "kimi-code" + }`)) + })) + defer srv.Close() + withRefreshURL(t, srv.URL) + + creds, err := loadOAuthCreds() + if err != nil { + t.Fatalf("loadOAuthCreds: %v", err) + } + creds, err = refreshOAuthToken(context.Background(), creds) + if err != nil { + t.Fatalf("refreshOAuthToken: %v", err) + } + if creds.AccessToken != "NEW" { + t.Fatalf("access_token = %q", creds.AccessToken) + } + if creds.ExpiresAt == nil { + t.Fatal("expires_at not set") + } + if creds.Scope != "kimi-code" { + t.Fatalf("scope = %q", creds.Scope) + } + + reloaded, err := loadOAuthCreds() + if err != nil { + t.Fatalf("reload: %v", err) + } + if reloaded.AccessToken != "NEW" || reloaded.RefreshToken != "R" || reloaded.TokenType != "Bearer" { + t.Fatalf("reloaded creds did not round-trip: %+v", reloaded) + } +} + +// TestRefreshOAuthToken401Surfaces verifies an auth-status response +// surfaces a `kimi login` hint instead of being swallowed. +func TestRefreshOAuthToken401Surfaces(t *testing.T) { + dir := withTempCredsDir(t) + writeCreds(t, dir, `{"access_token":"old","refresh_token":"R"}`) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + })) + defer srv.Close() + withRefreshURL(t, srv.URL) + + creds, err := loadOAuthCreds() + if err != nil { + t.Fatalf("loadOAuthCreds: %v", err) + } + _, err = refreshOAuthToken(context.Background(), creds) + if err == nil || !strings.Contains(err.Error(), "kimi login") { + t.Fatalf("err = %v, want `kimi login` hint", err) + } +} + +// TestRefreshOAuthTokenTransientFailureIsLenient verifies a 502 leaves +// the existing access token in place so the usage call can still try. +func TestRefreshOAuthTokenTransientFailureIsLenient(t *testing.T) { + dir := withTempCredsDir(t) + writeCreds(t, dir, `{"access_token":"keep","refresh_token":"R"}`) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusBadGateway) + })) + defer srv.Close() + withRefreshURL(t, srv.URL) + + creds, err := loadOAuthCreds() + if err != nil { + t.Fatalf("loadOAuthCreds: %v", err) + } + got, err := refreshOAuthToken(context.Background(), creds) + if err != nil { + t.Fatalf("err = %v, want nil for transient failure", err) + } + if got.AccessToken != "keep" { + t.Fatalf("access_token = %q, want existing token preserved", got.AccessToken) + } +} + +// TestSnapshotFromUsageReportsKimiForCoding verifies the rebrand made +// it to the snapshot ProviderName so the SD cache key matches the new +// label. +func TestSnapshotFromUsageReportsKimiForCoding(t *testing.T) { + snap := snapshotFromUsage(usageSnapshot{ + Weekly: usageDetail{Limit: "100", Used: "10"}, + UpdatedAt: time.Now().UTC(), + }) + if snap.ProviderName != "Kimi" { + t.Fatalf("ProviderName = %q, want %q", snap.ProviderName, "Kimi") + } +} + +// withTempCredsDir redirects oauthCredsPathFn at a fresh temp dir for +// the duration of the test. +func withTempCredsDir(t *testing.T) string { + t.Helper() + dir := t.TempDir() + prev := oauthCredsPathFn + oauthCredsPathFn = func() string { return filepath.Join(dir, "kimi-code.json") } + t.Cleanup(func() { oauthCredsPathFn = prev }) + return dir +} + +// writeCreds places a credential JSON blob at the test temp location. +func writeCreds(t *testing.T, _ string, body string) { + t.Helper() + if !json.Valid([]byte(body)) { + t.Fatalf("test fixture invalid JSON: %s", body) + } + if err := os.WriteFile(oauthCredsPath(), []byte(body), 0o600); err != nil { + t.Fatalf("write creds: %v", err) + } +} + +// withRefreshURL redirects oauthRefreshURLFn at the supplied URL for +// the duration of the test. +func withRefreshURL(t *testing.T, url string) { + t.Helper() + prev := oauthRefreshURLFn + oauthRefreshURLFn = func() string { return url } + t.Cleanup(func() { oauthRefreshURLFn = prev }) +} diff --git a/internal/providers/kimi/oauth.go b/internal/providers/kimi/oauth.go new file mode 100644 index 0000000..fe7cacf --- /dev/null +++ b/internal/providers/kimi/oauth.go @@ -0,0 +1,300 @@ +// Auth fallback for Kimi when the Helper extension is +// unavailable. Uses an OAuth token grant that the `kimi login` CLI +// places at ~/.kimi/credentials/kimi-code.json, refreshes it against +// auth.kimi.com when within 5 minutes of expiry (or on a 401/403), +// and reads usage from api.kimi.com/coding/v1/usages. +// +// The cookie path remains primary (extension-first architecture); +// OAuth is only consulted when HostAvailable is false or when the +// cookie path returns an auth-stale error. + +package kimi + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "os" + "path/filepath" + "strings" + "time" + + "github.com/anthonybaldwin/UsageButtons/internal/httputil" + "github.com/anthonybaldwin/UsageButtons/internal/providers/providerutil" +) + +const ( + // oauthClientID is the Kimi-published OAuth client used by the + // `kimi login` CLI. Refresh-only — initial auth-code grant + PKCE + // is handled by the CLI; we never see the user's password. + oauthClientID = "17e5f671-d194-4dfb-9706-5516cb48c098" + // defaultOAuthRefreshURL is the form-encoded refresh-token endpoint. + defaultOAuthRefreshURL = "https://auth.kimi.com/api/oauth/token" + // oauthUsageURL returns Kimi session/weekly windows. + oauthUsageURL = "https://api.kimi.com/coding/v1/usages" + // oauthRefreshBuffer is how far ahead of expiry we proactively refresh. + oauthRefreshBuffer = 5 * time.Minute + // oauthHTTPTimeout caps refresh + usage HTTP calls. + oauthHTTPTimeout = 30 * time.Second +) + +// oauthCredsPathFn is overridden in tests to redirect the credential +// lookup at a temp dir without depending on HOME/USERPROFILE quirks. +var oauthCredsPathFn = defaultOAuthCredsPath + +// oauthRefreshURLFn is overridden in tests to point refresh requests at +// a httptest server. +var oauthRefreshURLFn = func() string { return defaultOAuthRefreshURL } + +// defaultOAuthCredsPath returns the standard kimi-code OAuth path. +func defaultOAuthCredsPath() string { + home, _ := os.UserHomeDir() + return filepath.Join(home, ".kimi", "credentials", "kimi-code.json") +} + +// oauthCredsPath returns the on-disk location of the kimi-code OAuth blob. +func oauthCredsPath() string { return oauthCredsPathFn() } + +// oauthCreds is the on-disk shape of ~/.kimi/credentials/kimi-code.json. +// Fields not listed here (if Kimi's CLI ever adds new ones) are preserved +// across refresh because saveOAuthCreds round-trips through map[string]any. +type oauthCreds struct { + AccessToken string `json:"access_token,omitempty"` + RefreshToken string `json:"refresh_token,omitempty"` + ExpiresAt *float64 `json:"expires_at,omitempty"` // unix seconds (may be fractional) + Scope string `json:"scope,omitempty"` + TokenType string `json:"token_type,omitempty"` +} + +// loadOAuthCreds reads and validates the kimi-code credential blob. +func loadOAuthCreds() (oauthCreds, error) { + path := oauthCredsPath() + data, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + return oauthCreds{}, fmt.Errorf("Kimi credentials not found at %s. Run `kimi login` to authenticate.", path) + } + return oauthCreds{}, err + } + var c oauthCreds + if err := json.Unmarshal(data, &c); err != nil { + return oauthCreds{}, fmt.Errorf("invalid JSON in %s: %w", path, err) + } + if strings.TrimSpace(c.AccessToken) == "" && strings.TrimSpace(c.RefreshToken) == "" { + return oauthCreds{}, fmt.Errorf("Kimi credentials at %s missing access_token / refresh_token. Run `kimi login` to authenticate.", path) + } + return c, nil +} + +// needsRefresh reports whether the access token is missing or within +// oauthRefreshBuffer of its expiry. +func (c oauthCreds) needsRefresh(now time.Time) bool { + if strings.TrimSpace(c.AccessToken) == "" { + return true + } + if c.ExpiresAt == nil { + return true + } + return now.Add(oauthRefreshBuffer).Unix() >= int64(*c.ExpiresAt) +} + +// refreshResponse is the JSON shape of auth.kimi.com's token endpoint. +type refreshResponse struct { + AccessToken string `json:"access_token"` + RefreshToken string `json:"refresh_token,omitempty"` + ExpiresIn *int64 `json:"expires_in,omitempty"` + Scope string `json:"scope,omitempty"` + TokenType string `json:"token_type,omitempty"` +} + +// refreshOAuthToken exchanges the stored refresh_token for a fresh +// access_token and persists the result. On non-2xx responses other than +// 401/403 the existing access token is kept and the caller proceeds — +// matching openusage's leniency for transient refresh failures. +func refreshOAuthToken(ctx context.Context, creds oauthCreds) (oauthCreds, error) { + if strings.TrimSpace(creds.RefreshToken) == "" { + return creds, fmt.Errorf("Kimi credentials missing refresh_token. Run `kimi login` to authenticate.") + } + form := url.Values{} + form.Set("client_id", oauthClientID) + form.Set("grant_type", "refresh_token") + form.Set("refresh_token", creds.RefreshToken) + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, oauthRefreshURLFn(), strings.NewReader(form.Encode())) + if err != nil { + return creds, err + } + req.Header.Set("Content-Type", "application/x-www-form-urlencoded") + req.Header.Set("Accept", "application/json") + req.Header.Set("User-Agent", httputil.DefaultUserAgent) + + client := &http.Client{Timeout: oauthHTTPTimeout} + resp, err := client.Do(req) + if err != nil { + return creds, fmt.Errorf("Kimi OAuth refresh network error: %w", err) + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + return creds, fmt.Errorf("Kimi OAuth refresh read error: %w", err) + } + if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden { + return creds, fmt.Errorf("Kimi session expired. Run `kimi login` to authenticate.") + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + // Non-auth failure: keep existing token and let the usage call try. + return creds, nil + } + var decoded refreshResponse + if err := json.Unmarshal(body, &decoded); err != nil { + return creds, fmt.Errorf("Kimi OAuth refresh parse error: %w", err) + } + if strings.TrimSpace(decoded.AccessToken) == "" { + return creds, fmt.Errorf("Kimi OAuth refresh missing access_token") + } + creds.AccessToken = strings.TrimSpace(decoded.AccessToken) + if strings.TrimSpace(decoded.RefreshToken) != "" { + creds.RefreshToken = strings.TrimSpace(decoded.RefreshToken) + } + if decoded.ExpiresIn != nil { + exp := float64(time.Now().Unix() + *decoded.ExpiresIn) + creds.ExpiresAt = &exp + } + if strings.TrimSpace(decoded.Scope) != "" { + creds.Scope = strings.TrimSpace(decoded.Scope) + } + if strings.TrimSpace(decoded.TokenType) != "" { + creds.TokenType = strings.TrimSpace(decoded.TokenType) + } + if err := saveOAuthCreds(creds); err != nil { + return creds, fmt.Errorf("save Kimi credentials: %w", err) + } + return creds, nil +} + +// saveOAuthCreds writes refreshed credentials back to disk while +// preserving any extra keys the Kimi CLI may have added since. +func saveOAuthCreds(creds oauthCreds) error { + path := oauthCredsPath() + root := map[string]any{} + if data, err := os.ReadFile(path); err == nil { + _ = json.Unmarshal(data, &root) + } + root["access_token"] = creds.AccessToken + if creds.RefreshToken != "" { + root["refresh_token"] = creds.RefreshToken + } + if creds.ExpiresAt != nil { + root["expires_at"] = *creds.ExpiresAt + } + if creds.Scope != "" { + root["scope"] = creds.Scope + } + if creds.TokenType != "" { + root["token_type"] = creds.TokenType + } + return providerutil.WriteJSONAtomic(path, root) +} + +// oauthUsageEnvelope is the api.kimi.com/coding/v1/usages response shape. +// Differs from the gateway shape parsed by parseUsage in kimi.go (which +// returns {usages:[{scope,detail,limits[]}]}) — the OAuth-direct path +// returns a single top-level usage block plus a windowed limits[] array. +type oauthUsageEnvelope struct { + Usage *usageDetail `json:"usage,omitempty"` + Limits []rateLimitEntry `json:"limits,omitempty"` +} + +// fetchWithOAuth refreshes credentials when needed, then reads the +// kimi-code usage endpoint and normalizes it into a usageSnapshot +// compatible with the cookie path's snapshotFromUsage(). +func fetchWithOAuth(ctx context.Context) (usageSnapshot, error) { + creds, err := loadOAuthCreds() + if err != nil { + return usageSnapshot{}, err + } + if creds.needsRefresh(time.Now()) { + creds, err = refreshOAuthToken(ctx, creds) + if err != nil { + return usageSnapshot{}, err + } + } + usage, fetchErr := readOAuthUsage(ctx, creds.AccessToken) + if fetchErr == nil { + return usage, nil + } + // Reactive refresh on auth failure, retry once. + var httpErr *httputil.Error + if errors.As(fetchErr, &httpErr) && (httpErr.Status == 401 || httpErr.Status == 403) { + refreshed, refreshErr := refreshOAuthToken(ctx, creds) + if refreshErr != nil { + return usageSnapshot{}, refreshErr + } + usage, fetchErr = readOAuthUsage(ctx, refreshed.AccessToken) + if fetchErr == nil { + return usage, nil + } + } + return usageSnapshot{}, fetchErr +} + +// readOAuthUsage performs the GET against api.kimi.com/coding/v1/usages. +func readOAuthUsage(ctx context.Context, accessToken string) (usageSnapshot, error) { + if strings.TrimSpace(accessToken) == "" { + return usageSnapshot{}, fmt.Errorf("Kimi OAuth access token empty after refresh") + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, oauthUsageURL, nil) + if err != nil { + return usageSnapshot{}, err + } + req.Header.Set("Authorization", "Bearer "+accessToken) + req.Header.Set("Accept", "application/json") + req.Header.Set("User-Agent", httputil.DefaultUserAgent) + + client := &http.Client{Timeout: oauthHTTPTimeout} + resp, err := client.Do(req) + if err != nil { + return usageSnapshot{}, fmt.Errorf("Kimi usage fetch error: %w", err) + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + return usageSnapshot{}, fmt.Errorf("Kimi usage read error: %w", err) + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return usageSnapshot{}, &httputil.Error{ + Status: resp.StatusCode, + StatusText: resp.Status, + Body: string(body), + URL: oauthUsageURL, + } + } + var env oauthUsageEnvelope + if err := json.Unmarshal(body, &env); err != nil { + return usageSnapshot{}, fmt.Errorf("invalid Kimi usage JSON: %w", err) + } + return parseOAuthUsage(env, time.Now().UTC()) +} + +// parseOAuthUsage maps the kimi-code envelope into the same usageSnapshot +// shape produced by parseUsage so snapshotFromUsage doesn't care which +// transport produced the data. +func parseOAuthUsage(env oauthUsageEnvelope, now time.Time) (usageSnapshot, error) { + if env.Usage == nil && len(env.Limits) == 0 { + return usageSnapshot{}, fmt.Errorf("Kimi response missing usage and limits") + } + snap := usageSnapshot{UpdatedAt: now} + if env.Usage != nil { + snap.Weekly = *env.Usage + } + if len(env.Limits) > 0 { + detail := env.Limits[0].Detail + snap.Rate = &detail + } + return snap, nil +} + diff --git a/internal/providers/kimik2/kimik2.go b/internal/providers/kimrel/kimrel.go similarity index 76% rename from internal/providers/kimik2/kimik2.go rename to internal/providers/kimrel/kimrel.go index 036e2a2..505e256 100644 --- a/internal/providers/kimik2/kimik2.go +++ b/internal/providers/kimrel/kimrel.go @@ -1,9 +1,24 @@ -// Package kimik2 implements the Kimi K2 credits provider. +// Package kimrel implements the Kimrel credits provider. // -// Auth: Property Inspector settings field or KIMI_K2_API_KEY / -// KIMI_API_KEY / KIMI_KEY environment variable. -// Endpoint: GET https://kimi-k2.ai/api/user/credits -package kimik2 +// Kimrel (kimrel.com, formerly kimi-k2.ai) is an INDEPENDENT THIRD-PARTY +// reseller of Kimi K2 model access. It is NOT affiliated with, endorsed +// by, or sponsored by Moonshot AI — Kimrel's own footer states this. +// Users only see data here if they hold a kimrel.com account; Moonshot +// API keys won't authenticate against this endpoint. +// +// For the official Moonshot dev platform (api.moonshot.ai/cn) use the +// `moonshot` provider instead. For the kimi.com membership (Moderato, +// Allegretto, etc.) use the `kimi` provider. +// +// Auth: Property Inspector settings field or KIMREL_API_KEY (preferred) +// / KIMI_K2_API_KEY / KIMI_API_KEY / KIMI_KEY environment variable. +// Endpoint: GET https://kimi-k2.ai/api/user/credits (308-redirects to +// the kimrel.com production host). +// +// Provider ID stays "kimi-k2" so existing user button settings keep +// working across the rename — only the user-visible label and package +// name moved to "Kimrel." +package kimrel import ( "encoding/json" @@ -18,14 +33,19 @@ import ( "github.com/anthonybaldwin/UsageButtons/internal/settings" ) -// creditsURL is the Kimi K2 credits lookup endpoint. +// creditsURL is the Kimrel credits lookup endpoint. The kimi-k2.ai +// host 308-redirects to kimrel.com today; we point at the original +// host so server-side renames remain transparent. const creditsURL = "https://kimi-k2.ai/api/user/credits" -// getAPIKey resolves a Kimi K2 API key from user settings or env vars. +// getAPIKey resolves a Kimrel API key from user settings or env vars. +// KIMREL_API_KEY is the preferred name; the older KIMI_K2_API_KEY / +// KIMI_API_KEY / KIMI_KEY names still resolve so existing setups keep +// working after the rename. func getAPIKey() string { return settings.ResolveAPIKey( settings.ProviderKeysGet().KimiK2Key, - "KIMI_K2_API_KEY", "KIMI_API_KEY", "KIMI_KEY", + "KIMREL_API_KEY", "KIMI_K2_API_KEY", "KIMI_API_KEY", "KIMI_KEY", ) } @@ -241,37 +261,41 @@ func dateFromNumeric(v float64) (time.Time, bool) { return time.Unix(int64(v), 0), true } -// Provider fetches Kimi K2 usage data. +// Provider fetches Kimrel usage data. type Provider struct{} -// ID returns the provider identifier used by the registry. +// ID returns the provider identifier used by the registry. Kept as +// "kimi-k2" so existing user button settings keep working after the +// rename to Kimrel. func (Provider) ID() string { return "kimi-k2" } // Name returns the human-readable provider name. -func (Provider) Name() string { return "Kimi K2" } +func (Provider) Name() string { return "Kimrel" } -// BrandColor returns the accent color used on button faces. -func (Provider) BrandColor() string { return "#4c00ff" } +// BrandColor returns the accent color used on button faces. Slate gray +// is intentional — Kimrel is third-party and shouldn't borrow Kimi's +// orange or Moonshot's blue, which would imply official affiliation. +func (Provider) BrandColor() string { return "#64748b" } // BrandBg returns the background color used on button faces. -func (Provider) BrandBg() string { return "#0c0324" } +func (Provider) BrandBg() string { return "#1e293b" } // MetricIDs enumerates the metrics this provider can emit. func (Provider) MetricIDs() []string { return []string{"credits-balance"} } -// Fetch returns the latest Kimi K2 credits snapshot. +// Fetch returns the latest Kimrel credits snapshot. func (Provider) Fetch(_ providers.FetchContext) (providers.Snapshot, error) { apiKey := getAPIKey() if apiKey == "" { return providers.Snapshot{ ProviderID: "kimi-k2", - ProviderName: "Kimi K2", + ProviderName: "Kimrel", Source: "none", Metrics: []providers.MetricValue{}, Status: "unknown", - Error: "Enter a Kimi K2 API key in the Kimi K2 tab, or set KIMI_K2_API_KEY.", + Error: "Enter a Kimrel API key in the Kimrel tab, or set KIMREL_API_KEY. Kimrel (kimrel.com) is an independent third-party reseller of Kimi K2 — not affiliated with Moonshot AI. For the official Moonshot platform, use the Moonshot provider instead.", }, nil } @@ -285,11 +309,11 @@ func (Provider) Fetch(_ providers.FetchContext) (providers.Snapshot, error) { if errors.As(err, &httpErr) && (httpErr.Status == 401 || httpErr.Status == 403) { return providers.Snapshot{ ProviderID: "kimi-k2", - ProviderName: "Kimi K2", + ProviderName: "Kimrel", Source: "api-key", Metrics: []providers.MetricValue{}, Status: "unknown", - Error: "Kimi K2 API key unauthorized. Check KIMI_K2_API_KEY.", + Error: "Kimrel API key unauthorized. Check KIMREL_API_KEY.", }, nil } return providers.Snapshot{}, err @@ -323,7 +347,7 @@ func (Provider) Fetch(_ providers.FetchContext) (providers.Snapshot, error) { metrics = append(metrics, providers.MetricValue{ ID: "credits-balance", Label: "CREDITS", - Name: "Kimi K2 credits remaining", + Name: "Kimrel credits remaining", Value: math.Round(remainPct), NumericValue: &remainPct, NumericUnit: "percent", @@ -341,7 +365,7 @@ func (Provider) Fetch(_ providers.FetchContext) (providers.Snapshot, error) { metrics = append(metrics, providers.MetricValue{ ID: "credits-balance", Label: "CREDITS", - Name: "Kimi K2 credits", + Name: "Kimrel credits", Value: fmt.Sprintf("%d", int(math.Round(r))), NumericValue: &r, NumericUnit: "count", @@ -357,7 +381,7 @@ func (Provider) Fetch(_ providers.FetchContext) (providers.Snapshot, error) { return providers.Snapshot{ ProviderID: "kimi-k2", - ProviderName: "Kimi K2", + ProviderName: "Kimrel", Source: "api-key", Metrics: metrics, Status: "operational", @@ -374,7 +398,7 @@ func firstInContexts(body map[string]any, paths [][]string) (float64, bool) { return 0, false } -// init registers the Kimi K2 provider with the package registry. +// init registers the Kimrel provider with the package registry. func init() { providers.Register(Provider{}) } diff --git a/internal/providers/kimik2/kimik2_test.go b/internal/providers/kimrel/kimrel_test.go similarity index 99% rename from internal/providers/kimik2/kimik2_test.go rename to internal/providers/kimrel/kimrel_test.go index c63a7d1..425e75f 100644 --- a/internal/providers/kimik2/kimik2_test.go +++ b/internal/providers/kimrel/kimrel_test.go @@ -1,4 +1,4 @@ -package kimik2 +package kimrel import ( "encoding/json" diff --git a/internal/wsl/wsl.go b/internal/wsl/wsl.go new file mode 100644 index 0000000..9fb03c6 --- /dev/null +++ b/internal/wsl/wsl.go @@ -0,0 +1,28 @@ +// Package wsl discovers running WSL distributions on Windows so that +// providers which scan local CLI state files (~/.claude/projects, +// ~/.codex/sessions) can also surface usage from inside WSL distros as +// separate "machines." +// +// The whole feature is Windows-only by design — on non-Windows builds +// Sources returns nil, so callers can compose unconditionally. +package wsl + +// Source describes one WSL distribution surface that providers can scan. +// +// Native Windows is NOT represented here; callers continue to use +// os.UserHomeDir() for that. Sources returns ONLY the additional WSL +// distros so the call site reads as "scan native + each WSL source." +type Source struct { + // Key is the metric-ID-safe identifier for this distro (alnum + "_"). + // Distro names like "Ubuntu-22.04" become "Ubuntu_22_04" so they can + // be appended to existing metric IDs (e.g. "cost-today-wsl-Ubuntu_22_04") + // without breaking the codebase's hyphen-delimited convention. + Key string + // Label is the human-friendly distro name as reported by wsl.exe + // (e.g. "Ubuntu-22.04"). Used verbatim in PI dropdown labels. + Label string + // Home is the UNC path to the distro user's home directory + // (e.g. \\wsl.localhost\Debian\home\anthony). Suitable for direct + // filepath.Join with relative subpaths like ".claude/projects". + Home string +} diff --git a/internal/wsl/wsl_other.go b/internal/wsl/wsl_other.go new file mode 100644 index 0000000..0accd24 --- /dev/null +++ b/internal/wsl/wsl_other.go @@ -0,0 +1,8 @@ +//go:build !windows + +package wsl + +// Sources returns nil on non-Windows platforms. WSL only exists on +// Windows, so providers compiled for macOS/Linux see an empty source +// list and never try to spawn wsl.exe or read \\wsl.localhost paths. +func Sources() []Source { return nil } diff --git a/internal/wsl/wsl_windows.go b/internal/wsl/wsl_windows.go new file mode 100644 index 0000000..e8ab035 --- /dev/null +++ b/internal/wsl/wsl_windows.go @@ -0,0 +1,174 @@ +//go:build windows + +package wsl + +import ( + "context" + "encoding/binary" + "os" + "os/exec" + "path/filepath" + "sort" + "strings" + "sync" + "time" + "unicode/utf16" + + "github.com/anthonybaldwin/UsageButtons/internal/winutil" +) + +// sourcesCacheTTL bounds how often we shell out to wsl.exe and walk +// \\wsl.localhost. Distro state is stable on a human timescale; rescanning +// on every Fetch would slow down the cost-tile path needlessly. +const sourcesCacheTTL = 30 * time.Second + +var ( + sourcesMu sync.Mutex + sourcesCache []Source + sourcesCacheT time.Time +) + +// Sources returns the running, user-bearing WSL distributions on this +// machine. Stopped distros are deliberately omitted — accessing +// \\wsl.localhost\\... cold-starts the VM, which is far too +// expensive to do on a 5-minute scan loop. +// +// Returns nil (not error) if WSL is not installed, no distros are +// running, or any probe fails. Callers should treat nil as "no extra +// sources" rather than as an error condition. +func Sources() []Source { + sourcesMu.Lock() + defer sourcesMu.Unlock() + if sourcesCache != nil && time.Since(sourcesCacheT) < sourcesCacheTTL { + return sourcesCache + } + sourcesCache = discover() + sourcesCacheT = time.Now() + return sourcesCache +} + +// discover does the actual probe: enumerate running distros via wsl.exe, +// then resolve each distro's home directory under \\wsl.localhost. +func discover() []Source { + if _, err := exec.LookPath("wsl.exe"); err != nil { + return nil + } + + distros := listRunningDistros() + if len(distros) == 0 { + return nil + } + + var out []Source + for _, name := range distros { + // docker-desktop is a Docker-internal distro with no real user + // home; skip even if it's running. + if strings.EqualFold(name, "docker-desktop") || strings.EqualFold(name, "docker-desktop-data") { + continue + } + home := resolveHome(name) + if home == "" { + continue + } + out = append(out, Source{ + Key: sanitizeKey(name), + Label: name, + Home: home, + }) + } + return out +} + +// listRunningDistros runs `wsl.exe -l -q --running` and returns the +// distro names it prints. The --running flag is critical: bare `-l -q` +// would also list stopped distros, but querying their FS via +// \\wsl.localhost would silently boot them. +// +// wsl.exe writes UTF-16LE to stdout, with each distro on its own line. +// We decode and trim before returning. +func listRunningDistros() []string { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + cmd := exec.CommandContext(ctx, "wsl.exe", "-l", "-q", "--running") + winutil.HideConsoleWindow(cmd) + raw, err := cmd.Output() + if err != nil { + return nil + } + + text := decodeUTF16LE(raw) + var names []string + for _, line := range strings.Split(text, "\n") { + line = strings.TrimSpace(strings.TrimRight(line, "\r\x00")) + if line == "" { + continue + } + names = append(names, line) + } + sort.Strings(names) + return names +} + +// decodeUTF16LE converts UTF-16LE bytes (with optional BOM) into a Go +// string. wsl.exe always writes its list output in UTF-16LE on Windows, +// regardless of console code page. +func decodeUTF16LE(raw []byte) string { + if len(raw) >= 2 && raw[0] == 0xFF && raw[1] == 0xFE { + raw = raw[2:] + } + if len(raw)%2 != 0 { + raw = raw[:len(raw)-1] + } + u16 := make([]uint16, len(raw)/2) + for i := range u16 { + u16[i] = binary.LittleEndian.Uint16(raw[i*2:]) + } + return string(utf16.Decode(u16)) +} + +// resolveHome returns the UNC path to the default user's home directory +// inside the given distro, or "" if it can't be determined. +// +// Strategy: enumerate \\wsl.localhost\\home\ and pick the first +// (and almost always only) user directory. Shelling out via `wsl -d +// -u root -e cat /etc/passwd` would be more correct on +// multi-user distros but adds a per-distro process spawn for a case +// that doesn't exist in practice. +func resolveHome(distro string) string { + root := `\\wsl.localhost\` + distro + `\home` + entries, err := os.ReadDir(root) + if err != nil { + // Older Windows builds expose the legacy \\wsl$\... share instead + // of \\wsl.localhost\... — try that as a fallback before giving up. + root = `\\wsl$\` + distro + `\home` + entries, err = os.ReadDir(root) + if err != nil { + return "" + } + } + for _, e := range entries { + if e.IsDir() { + return filepath.Join(root, e.Name()) + } + } + return "" +} + +// sanitizeKey converts a distro name to a metric-ID-safe identifier by +// replacing every non-alphanumeric character with underscore. "Ubuntu-22.04" +// becomes "Ubuntu_22_04" so it can be appended to hyphenated metric IDs +// without ambiguity. +func sanitizeKey(name string) string { + var b strings.Builder + b.Grow(len(name)) + for _, r := range name { + switch { + case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9': + b.WriteRune(r) + default: + b.WriteByte('_') + } + } + return b.String() +} diff --git a/io.github.anthonybaldwin.UsageButtons.sdPlugin/assets/action-kimi-k2-key.svg b/io.github.anthonybaldwin.UsageButtons.sdPlugin/assets/action-kimi-k2-key.svg index 5f5cee6..e0dbeb4 100644 --- a/io.github.anthonybaldwin.UsageButtons.sdPlugin/assets/action-kimi-k2-key.svg +++ b/io.github.anthonybaldwin.UsageButtons.sdPlugin/assets/action-kimi-k2-key.svg @@ -1,6 +1,6 @@ - - - + + + diff --git a/io.github.anthonybaldwin.UsageButtons.sdPlugin/assets/action-kimi-k2.svg b/io.github.anthonybaldwin.UsageButtons.sdPlugin/assets/action-kimi-k2.svg index 0948f79..80e90bc 100644 --- a/io.github.anthonybaldwin.UsageButtons.sdPlugin/assets/action-kimi-k2.svg +++ b/io.github.anthonybaldwin.UsageButtons.sdPlugin/assets/action-kimi-k2.svg @@ -1,3 +1,3 @@ - - \ No newline at end of file + + diff --git a/io.github.anthonybaldwin.UsageButtons.sdPlugin/assets/action-moonshot-key.svg b/io.github.anthonybaldwin.UsageButtons.sdPlugin/assets/action-moonshot-key.svg index 4469264..6551bcc 100644 --- a/io.github.anthonybaldwin.UsageButtons.sdPlugin/assets/action-moonshot-key.svg +++ b/io.github.anthonybaldwin.UsageButtons.sdPlugin/assets/action-moonshot-key.svg @@ -1,6 +1,6 @@ - - - + + + diff --git a/io.github.anthonybaldwin.UsageButtons.sdPlugin/assets/action-moonshot.svg b/io.github.anthonybaldwin.UsageButtons.sdPlugin/assets/action-moonshot.svg index 7142857..206e769 100644 --- a/io.github.anthonybaldwin.UsageButtons.sdPlugin/assets/action-moonshot.svg +++ b/io.github.anthonybaldwin.UsageButtons.sdPlugin/assets/action-moonshot.svg @@ -1,5 +1,5 @@ - - + + diff --git a/io.github.anthonybaldwin.UsageButtons.sdPlugin/manifest.json b/io.github.anthonybaldwin.UsageButtons.sdPlugin/manifest.json index e594c58..ce07792 100644 --- a/io.github.anthonybaldwin.UsageButtons.sdPlugin/manifest.json +++ b/io.github.anthonybaldwin.UsageButtons.sdPlugin/manifest.json @@ -124,7 +124,7 @@ }, { "UUID": "io.github.anthonybaldwin.UsageButtons.claude", - "Name": "Claude Code", + "Name": "Claude", "Tooltip": "Claude usage stats — session %, weekly %, Sonnet %, Opus %, Design, Routines, extras, costs.", "Icon": "assets/action-claude", "PropertyInspectorPath": "ui/stat.html", @@ -301,7 +301,7 @@ { "UUID": "io.github.anthonybaldwin.UsageButtons.kimi", "Name": "Kimi", - "Tooltip": "Kimi usage stats — weekly coding quota and five-hour rate limit.", + "Tooltip": "Kimi usage — weekly coding quota and 5-hour rate limit. Helper extension preferred; OAuth from `kimi login` is the fallback.", "Icon": "assets/action-kimi", "PropertyInspectorPath": "ui/stat.html", "SupportedInMultiActions": false, @@ -316,8 +316,8 @@ }, { "UUID": "io.github.anthonybaldwin.UsageButtons.kimi-k2", - "Name": "Kimi K2", - "Tooltip": "Kimi K2 usage stats — credits remaining.", + "Name": "Kimrel", + "Tooltip": "Kimrel credits — third-party reseller of Kimi K2 (kimrel.com), not affiliated with Moonshot AI. For the official Moonshot dev API balance, use the Moonshot action instead.", "Icon": "assets/action-kimi-k2", "PropertyInspectorPath": "ui/stat.html", "SupportedInMultiActions": false, diff --git a/io.github.anthonybaldwin.UsageButtons.sdPlugin/ui/stat.html b/io.github.anthonybaldwin.UsageButtons.sdPlugin/ui/stat.html index 39bc0e7..254dca8 100644 --- a/io.github.anthonybaldwin.UsageButtons.sdPlugin/ui/stat.html +++ b/io.github.anthonybaldwin.UsageButtons.sdPlugin/ui/stat.html @@ -428,7 +428,7 @@

Admin keys only. Personal sk-ant- won't work. Fallback: ANTHROPIC_ADMIN_API_KEY.

- +

Admin keys only. Personal sk- won't work. Fallback: OPENAI_ADMIN_API_KEY.

+ - + @@ -1289,7 +1290,7 @@ zai: "#e85a6a", kimi: "#fe603c", minimax: "#fe603c", - "kimi-k2": "#4c00ff", + "kimi-k2": "#64748b", jetbrains: "#ff3399", kilo: "#f27027", kiro: "#ff9900", @@ -1330,7 +1331,7 @@ zai: "#211012", kimi: "#260d08", minimax: "#260d08", - "kimi-k2": "#0c0324", + "kimi-k2": "#1e293b", jetbrains: "#25051a", kilo: "#21130a", kiro: "#241405", @@ -1343,7 +1344,7 @@ // Display name shown on the per-provider tab. const PROVIDER_DISPLAY_NAMES = { - claude: "Claude Code", + claude: "Claude", anthropic: "Anthropic", openai: "OpenAI", codex: "Codex", @@ -1372,7 +1373,7 @@ zai: "z.ai", kimi: "Kimi", minimax: "MiniMax", - "kimi-k2": "Kimi K2", + "kimi-k2": "Kimrel", jetbrains: "JetBrains AI", kilo: "Kilo", kiro: "Kiro", @@ -2073,6 +2074,11 @@ setInterval(requestCookieStatus, 2500); requestProviderStatus(); setInterval(requestProviderStatus, 5000); + // One-shot: ask the plugin which WSL distros are running so the + // metric dropdown can offer per-distro local-cost entries + // (Windows-only feature; non-Windows builds reply with [] and + // the PI looks identical to today). + requestWSLDistros(); } function pushSettings() { @@ -2193,6 +2199,56 @@ })); } + // ---------- WSL distros ---------- + // Augments METRICS[claude]/METRICS[codex] with one extra + // "Cost today (WSL: …)" / "Cost last 30 days (WSL: …)" entry per + // running distro. Empty on non-Windows or when no distros are up, + // so the dropdown stays identical to before this feature shipped. + let wslDistrosKnown = false; + let wslDistrosCache = []; + + function requestWSLDistros() { + if (!websocket || websocket.readyState !== WebSocket.OPEN) return; + websocket.send(JSON.stringify({ + event: "sendToPlugin", + context: uuid, + payload: { action: "getWSLDistros" }, + })); + } + + function applyWSLDistros(distros) { + const list = Array.isArray(distros) ? distros : []; + // Idempotent: every PI session this is called exactly once on + // reply, so we don't need to strip prior wsl- entries. But guard + // against a future re-fetch by filtering them out first. + for (const provider of ["claude", "codex"]) { + const base = (METRICS[provider] || []).filter(([id]) => !/-wsl-/.test(id)); + const augmented = []; + for (const row of base) { + augmented.push(row); + const [id, label, kind] = row; + if (id === "cost-today" || id === "cost-30d") { + for (const d of list) { + const friendly = d.label || d.key; + const wslLabel = (id === "cost-today" ? "Cost today" : "Cost last 30 days") + + " (WSL: " + friendly + ")"; + augmented.push([id + "-wsl-" + d.key, wslLabel, kind]); + } + } + } + METRICS[provider] = augmented; + } + wslDistrosCache = list; + wslDistrosKnown = true; + // Repopulate the dropdown if the PI is currently looking at a + // provider we just augmented; the new options need to appear + // without requiring the user to reopen the inspector. + if (currentProviderId === "claude" || currentProviderId === "codex") { + const selected = document.getElementById("metric").value; + populateMetrics(currentProviderId, selected); + } + } + // ---------- Provider error banner ---------- function requestProviderStatus() { if (!websocket || websocket.readyState !== WebSocket.OPEN) return; @@ -2437,6 +2493,9 @@ case "providerStatus": setProviderError(payload); break; + case "wslDistros": + applyWSLDistros(payload.distros); + break; } } diff --git a/scripts/sync-lobe-icons.go b/scripts/sync-lobe-icons.go index c5685f4..c6b26ec 100644 --- a/scripts/sync-lobe-icons.go +++ b/scripts/sync-lobe-icons.go @@ -58,7 +58,8 @@ var mapping = []entry{ {ProviderID: "hermes-agent", LobeName: "hermesagent", Variant: "mono"}, {ProviderID: "kilo", LobeName: "kilocode", Variant: "mono"}, {ProviderID: "kimi", LobeName: "kimi", Variant: "mono"}, - {ProviderID: "kimi-k2", LobeName: "kimi", Variant: "mono"}, + // Kimrel (provider ID kimi-k2) is a third-party reseller — not allowed + // to use Kimi's Lobe icon. Custom mark lives in internal/icons/kimrel.go. {ProviderID: "minimax", LobeName: "minimax", Variant: "mono"}, {ProviderID: "mistral", LobeName: "mistral", Variant: "mono"}, {ProviderID: "moonshot", LobeName: "moonshot", Variant: "mono"}, @@ -92,7 +93,6 @@ type parsed struct { func main() { results := make(map[string]parsed, len(mapping)) - // dedupe fetches: kimi & kimi-k2 both pull kimi.svg type key struct{ name, variant string } cache := map[key]parsed{}