From f376b334fa9acb4e38546b69d183929b098c5a85 Mon Sep 17 00:00:00 2001 From: Phototonic <35551999+Phototonic@users.noreply.github.com> Date: Sat, 12 Sep 2026 13:13:45 +1000 Subject: [PATCH] feat(zai): add Coding Plan usage provider --- README.md | 45 +- docs/spec.md | 96 +++- fixtures/zai/README.md | 15 + fixtures/zai/credit_limit.json | 29 ++ fixtures/zai/mcp_only.json | 19 + fixtures/zai/token_mcp.json | 39 ++ i18n/en/yapcap.ftl | 9 + i18n/pl/yapcap.ftl | 9 + resources/providers/zai-reversed.svg | 4 + resources/providers/zai.svg | 4 + src/account_storage/mod.rs | 1 + src/account_storage/tests.rs | 18 + src/app/login/flows/mod.rs | 2 + src/app/login/flows/zai.rs | 180 +++++++ .../login_flow_tests/key_authentication.rs | 10 +- .../key_authentication_cases.rs | 62 +++ src/app/login/login_flow_tests/mod.rs | 1 + src/app/login/login_flow_tests/support.rs | 2 + src/app/login/login_flow_tests/zai.rs | 35 ++ src/app/login/mod.rs | 28 +- src/app/mod.rs | 9 + src/app/popup_view.rs | 2 + src/app/popup_view/detail.rs | 42 ++ src/app/popup_view/settings/accounts.rs | 6 +- .../settings/accounts/login_controls.rs | 84 +++- src/app/provider_actions.rs | 13 +- src/app/provider_assets.rs | 6 + src/app/session.rs | 4 + src/app/tests.rs | 21 +- src/config.rs | 97 ++++ src/config/watch_update.rs | 43 +- src/demo_env.rs | 159 ++++++- src/detection.rs | 102 +++- src/error.rs | 83 ++++ src/model.rs | 27 +- src/providers/adapters.rs | 3 + src/providers/adapters/zai_adapter.rs | 190 ++++++++ src/providers/interface.rs | 4 +- src/providers/mod.rs | 1 + src/providers/opencode_auth.rs | 112 ++++- src/providers/registry/tests.rs | 39 +- src/providers/zai/account.rs | 81 ++++ src/providers/zai/login.rs | 178 +++++++ src/providers/zai/mod.rs | 450 ++++++++++++++++++ src/providers/zai/opencode.rs | 67 +++ src/providers/zai/quota.rs | 343 +++++++++++++ src/providers/zai/storage.rs | 103 ++++ src/runtime.rs | 29 ++ 48 files changed, 2846 insertions(+), 60 deletions(-) create mode 100644 fixtures/zai/README.md create mode 100644 fixtures/zai/credit_limit.json create mode 100644 fixtures/zai/mcp_only.json create mode 100644 fixtures/zai/token_mcp.json create mode 100644 resources/providers/zai-reversed.svg create mode 100644 resources/providers/zai.svg create mode 100644 src/app/login/flows/zai.rs create mode 100644 src/app/login/login_flow_tests/zai.rs create mode 100644 src/providers/adapters/zai_adapter.rs create mode 100644 src/providers/zai/account.rs create mode 100644 src/providers/zai/login.rs create mode 100644 src/providers/zai/mod.rs create mode 100644 src/providers/zai/opencode.rs create mode 100644 src/providers/zai/quota.rs create mode 100644 src/providers/zai/storage.rs diff --git a/README.md b/README.md index 4444a6a..4e67bdc 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ # YapCap -**A native COSMIC panel applet that tracks AI coding quota for Codex, Claude Code, Cursor, Antigravity, Gemini, Minimax, GitHub Copilot, Kimi for Coding, and OpenCode Go.** +**A native COSMIC panel applet that tracks AI coding quota for Codex, Claude Code, Cursor, Antigravity, Gemini, GitHub Copilot, Minimax, Z.AI Coding Plan, Kimi for Coding, and OpenCode Go.** YapCap panel applet @@ -28,15 +28,16 @@ YapCap lives in your COSMIC panel and shows how much of your AI coding quota you - **Cursor** — Auto/Composer and API usage - **Antigravity** — grouped Gemini and Claude/GPT model quota (5h + weekly) - **Gemini** — Pro / Flash / Lite quota bars (OAuth accounts only) - - **Minimax** — API key usage tracking - **GitHub Copilot** — Free chat/completions or paid premium interactions + - **Minimax** — API key usage tracking + - **Z.AI Coding Plan** — global personal quota with 5-hour, weekly, and optional MCP windows - **Kimi for Coding** — API key usage tracking with weekly and rate-limit windows - **OpenCode Go** — API key usage tracking with 5-hour, weekly, and monthly windows - 👥 **Multi-account support** — add, switch, and remove accounts per provider. The popup pages through stored accounts one at a time, while the panel remains fixed-width for the active account. - 🔎 **Automatic discovery** — detected providers appear automatically, provider availability updates live, and an empty setup points directly to Settings. Gemini remains opt-in and must be enabled manually. -- 🔐 **In-app login** — guided browser login for Codex, Claude, Antigravity, Gemini, and Copilot; API-key forms for Minimax, Kimi, and OpenCode Go; Cursor scans the local IDE state. -- 🔑 **OpenCode integration** — compatible keys can optionally prefill Minimax, Kimi, and OpenCode Go forms; Codex and Copilot offer explicit OAuth imports. Credentials are copied only after confirmation and are never synchronized with OpenCode. -- ✅ **Active badge** — marks the account currently in use by the host tool or environment for Codex, Claude, Cursor, Gemini, Minimax, Kimi, and OpenCode Go. +- 🔐 **In-app login** — guided browser login for Codex, Claude, Antigravity, Gemini, and Copilot; API-key forms for Minimax, Z.AI, Kimi, and OpenCode Go; Cursor scans the local IDE state. +- 🔑 **OpenCode integration** — compatible keys can optionally prefill Minimax, Z.AI, Kimi, and OpenCode Go forms; Codex and Copilot offer explicit OAuth imports. Credentials are copied only after confirmation and are never synchronized with OpenCode. +- ✅ **Active badge** — marks the account currently in use by the host tool or environment for Codex, Claude, Cursor, Gemini, Minimax, Kimi, and OpenCode Go. Z.AI has no host Active badge. - ⚙️ **Configurable panel** — logo+bars, bars only, logo+%, or %-only; used/left toggle; relative or absolute reset times. ## Screenshots @@ -180,11 +181,39 @@ just install Each provider supports multiple accounts. Manage them from the popup under **Settings → [Provider]**. -- **Add account** — triggers the provider's own login flow: Codex browser OAuth, native Claude OAuth in the browser, Antigravity and Gemini browser OAuth, GitHub Copilot browser device flow, Minimax, Kimi, or OpenCode Go API-key entry, or Cursor IDE account scanning, without leaving YapCap. +- **Add account** — triggers the provider's own login flow: Codex browser OAuth, native Claude OAuth in the browser, Antigravity and Gemini browser OAuth, GitHub Copilot browser device flow, Minimax, Z.AI, Kimi, or OpenCode Go API-key entry, or Cursor IDE account scanning, without leaving YapCap. - **Switch account** — tap any account row to make it active; the panel and popup update immediately. - **Remove account** — deletes only YapCap's copy of the credentials. Provider accounts and host app configs are never touched. -Codex, Claude, Cursor, Antigravity, and Gemini keep at most one account per provider identity. Copilot keeps at most one account per GitHub numeric user id and displays the current GitHub username. Minimax, Kimi, and OpenCode Go use unique user-provided labels and reject duplicate API keys. +Codex, Claude, Cursor, Antigravity, and Gemini keep at most one account per provider identity. Copilot keeps at most one account per GitHub numeric user id and displays the current GitHub username. Minimax, Z.AI, Kimi, and OpenCode Go use unique user-provided labels and reject duplicate API keys. + +### Z.AI Coding Plan + +Add a Z.AI Coding Plan account from **Settings → Z.AI → Add account** and enter +the API key manually. YapCap stores only non-secret account metadata in COSMIC +configuration; the key belongs to the YapCap-managed account directory under +`/yapcap/zai-accounts//api_key.txt`. + +When adding or reauthenticating, YapCap may prefill a usable typed API key from +OpenCode's local `~/.local/share/opencode/auth.json`, checking `zai-coding-plan` +before `zai`. This is a +one-time read-only prefill: YapCap never writes to or synchronizes with +OpenCode, and it has no environment-key fallback. A key is stored locally after +save without remote validation. + +Automatic Z.AI detection is content-aware: a usable typed API entry at either +OpenCode key is required; a bare auth file, malformed entry, or non-API entry is +not treated as detected. + +YapCap reads the global personal Coding Plan quota from the fixed endpoint +`https://api.z.ai/api/monitor/usage/quota/limit`. The primary request uses +`Bearer ` and only a 401 receives one raw-key retry; 5xx responses are not +retried with raw authorization. Usage is shown as **5 Hour**, **Weekly**, and +optional **MCP** windows in that order. MCP has its own label and no assumed +calendar duration. An MCP-only response is valid, but the UI reports Coding +Plan usage as unavailable rather than fabricating token or weekly usage. OAuth, +alternate endpoints or hosts, region/team/promotional scopes, and fixed monthly +duration are not supported. ## Panel styles @@ -225,7 +254,7 @@ YapCap stores provider credentials under YapCap-owned account storage and calls | --- | --- | | `~/.config/cosmic/io.github.TopiCsarno.YapCap/v600/` | Settings (provider toggles, accounts, display options) | | `~/.cache/yapcap/snapshots.json` | Former cached usage state; current builds leave it on disk but do not load it | -| `~/.local/state/yapcap/-accounts/` | Managed credential copies (`` is one of `codex`, `claude`, `cursor`, `antigravity`, `gemini`, `minimax`, `copilot`, `kimi`, `opencode-go`) | +| `~/.local/state/yapcap/-accounts/` | Managed credential copies (`` is one of `codex`, `claude`, `cursor`, `antigravity`, `gemini`, `copilot`, `minimax`, `zai`, `kimi`, `opencode-go`) | | `~/.local/state/yapcap/logs/yapcap.log` | Log output | **Flatpak** (`io.github.TopiCsarno.YapCap`): YapCap account state and logs live only under `~/.var/app/io.github.TopiCsarno.YapCap/data/yapcap/`. Old Flatpak snapshot caches under `~/.var/app/io.github.TopiCsarno.YapCap/cache/yapcap/` may remain on disk but are no longer active runtime state. The manifest mounts host `~/.config/cosmic` read-write for COSMIC app settings (not `xdg-config/cosmic`, for compatibility with Flatpak path resolution). diff --git a/docs/spec.md b/docs/spec.md index 233a006..151f165 100644 --- a/docs/spec.md +++ b/docs/spec.md @@ -8,7 +8,7 @@ read_when: # YapCap — COSMIC Panel Applet Architecture -**Status:** As-built v0.6.0 · **Last updated:** 2026-09-03 +**Status:** As-built v0.6.0 · **Last updated:** 2026-09-12 ## Document Metadata @@ -18,7 +18,7 @@ read_when: | Target desktop | COSMIC | | Target language | Rust (edition 2024) | | Target runtime | libcosmic applet runtime | -| Providers | Codex, Claude Code, Cursor, Antigravity, Gemini, GitHub Copilot, Minimax, Kimi for Coding, OpenCode Go | +| Providers | Codex, Claude Code, Cursor, Antigravity, Gemini, GitHub Copilot, Minimax, Z.AI Coding Plan, Kimi for Coding, OpenCode Go | ## Document Map @@ -26,7 +26,7 @@ read_when: | --- | --- | | 1. Product Definition | 1.1 Scope and Non-Goals
1.2 Supported Sources | | 2. Architecture | 2.1 System Context
2.2 Crate Layout
2.3 Runtime and Message Flow
2.4 Multi-Process Applet Model | -| 3. Providers | 3.1 Codex
3.2 Claude
3.3 Cursor
3.4 Copilot
3.5 Gemini
3.6 Minimax
3.7 Kimi
3.8 OpenCode Go
3.9 Antigravity | +| 3. Providers | 3.1 Codex
3.2 Claude
3.3 Cursor
3.4 Copilot
3.5 Gemini
3.6 Minimax
3.7 Z.AI Coding Plan
3.8 Kimi
3.9 OpenCode Go
3.10 Antigravity | | 4. Auth and Config | 4.1 OAuth Credential Files
4.2 Cursor Token Source
4.3 Configuration | | 5. Data Model | 5.1 UsageSnapshot
5.2 ProviderRuntimeState and Health
5.3 Stale/Fresh Rules | | 6. Persistence, Logging, Paths | | @@ -39,7 +39,7 @@ read_when: ### 1.1 Scope and Non-Goals -- YapCap is a native Linux COSMIC panel applet that shows local usage state for Codex, Claude Code, Cursor, Antigravity, Gemini, GitHub Copilot, Minimax, Kimi for Coding, and OpenCode Go. +- YapCap is a native Linux COSMIC panel applet that shows local usage state for Codex, Claude Code, Cursor, Antigravity, Gemini, GitHub Copilot, Minimax, Z.AI Coding Plan, Kimi for Coding, and OpenCode Go. - Ships only on COSMIC. No GNOME, KDE, tray, or generic indicator paths exist. - Reads locally available credentials and caches. No user account, no cloud sync, no telemetry. - Out of scope: additional providers, historical charts, notifications, plugin architecture, doctor command, secret vault, alternative DEs. @@ -55,17 +55,18 @@ read_when: | Gemini | Active Gemini account resolved from YapCap-owned `gemini-accounts//` (`metadata.json`, `tokens.json`, optional `snapshot.json`) | OAuth refresh-token grant against `oauth2.googleapis.com/token` before expiry or once after a `loadCodeAssist` / `retrieveUserQuota` 401 | | Copilot | Active GitHub Copilot account resolved from YapCap-owned `copilot-accounts//` (`metadata.json`, `tokens.json`) | None; token is long-lived and re-auth is user-driven after revocation | | Minimax | Selected YapCap-managed API-key account under `/yapcap/minimax-accounts//` | `MINIMAX_API_KEY` identifies an environment-backed active account | +| Z.AI Coding Plan | Selected YapCap-managed API-key account under `/yapcap/zai-accounts//api_key.txt` | One-time content-aware prefill from OpenCode's `zai-coding-plan`, then `zai`, entries during add or reauthentication; no runtime fallback | | Kimi | Active YapCap-managed Kimi account resolved from `/yapcap/kimi-accounts//api_key.txt` | `KIMI_API_KEY` when the managed account has no stored key | | OpenCode Go | Active YapCap-managed OpenCode Go account whose private key matches the `opencode-go` API credential in OpenCode's `~/.local/share/opencode/auth.json` | `OPENCODE_API_KEY` or legacy `OPENCODE_GO_API_KEY` when the managed account has no stored key | -Claude, Codex, Cursor, Antigravity, Gemini, Minimax, Copilot, Kimi, and OpenCode Go all use YapCap-managed account storage. There +Claude, Codex, Cursor, Antigravity, Gemini, Copilot, Minimax, Z.AI, Kimi, and OpenCode Go all use YapCap-managed account storage. There is no web-cookie path for Claude and no forced-source environment variable. Gemini supports only Google OAuth accounts; gemini-cli API-key and Vertex AI configurations are out of scope. Minimax uses API key authentication without host -CLI integration. Kimi and OpenCode Go optionally prefill an API key once for add +CLI integration. Z.AI, Kimi, and OpenCode Go optionally prefill an API key once for add or reauthentication from OpenCode's local auth file; OpenCode Go also compares that file's `opencode-go` key with YapCap's private account keys for the Active -badge. Minimax does the same. OpenCode Go's auth file is watched for changes, +badge. Minimax does the same. Z.AI has no host Active badge. OpenCode Go's auth file is watched for changes, but YapCap never writes to or synchronizes with OpenCode during refresh. Codex and Copilot offer explicit compatible OAuth imports, including restoration of an existing managed account. These are one-time copies: YapCap never writes @@ -85,6 +86,7 @@ flowchart LR Panel --> Antigravity[Antigravity module] Panel --> Gemini[Gemini module] Panel --> Minimax[Minimax module] + Panel --> Zai[Z.AI module] Panel --> Copilot[Copilot module] Panel --> Kimi[Kimi module] Panel --> OpenCodeGo[OpenCode Go module] @@ -95,6 +97,7 @@ flowchart LR Antigravity --> AntigravityAPI[cloudcode-pa.googleapis.com] Gemini --> GeminiAPI[cloudcode-pa.googleapis.com] Minimax --> MinimaxAPI[www.minimax.io] + Zai --> ZaiAPI[api.z.ai] Copilot --> GitHubCopilot[api.github.com/copilot_internal/user] Kimi --> KimiAPI[api.kimi.com/coding/v1/usages] OpenCodeGo --> OpenCodeGoAPI[opencode.ai/zen/go/v1/usage] @@ -120,16 +123,17 @@ Library modules (`src/`, also usable from tests): | --- | --- | | `runtime` | `refresh_one(provider)`, `refresh_provider(...)`, `load_initial_state`, `persist_state`. Startup state is reconciled from shared runtime config, not from `snapshots.json`. | | `providers::registry` | Provider-facing interface used by runtime and UI code. It exposes provider capabilities, account discovery, prepared account-settings facts, account deletion, account status refresh, and usage fetch through provider adapters. | -| `providers::adapters` | Provider adapter implementations for Codex, Claude, Cursor, Antigravity, Gemini, Minimax, Copilot, Kimi, and OpenCode Go. Each adapter maps the shared provider interface onto provider-specific account and fetch modules. | +| `providers::adapters` | Provider adapter implementations for Codex, Claude, Cursor, Antigravity, Gemini, Copilot, Minimax, Z.AI, Kimi, and OpenCode Go. Each adapter maps the shared provider interface onto provider-specific account and fetch modules. | | `providers::interface` | Shared provider adapter trait, capability flags, account descriptors, prepared account-settings facts, account handles, and async future alias. | | `providers::codex` | Codex managed login, YapCap-owned account listing, OAuth usage fetch, and refresh-on-401/403 under `src/providers/codex/`. | | `providers::claude` | Managed native OAuth login and YapCap-owned account listing under `src/providers/claude/`, OAuth usage fetch, token refresh against Anthropic’s OAuth token endpoint (no Claude CLI), and read-only host `~/.claude.json` matching for `system_active_account_id`. | | `providers::cursor` | Cursor web API via YapCap-owned tokens scanned from Cursor IDE's local SQLite state. | | `providers::copilot` | GitHub device-flow login, id-based YapCap-owned account listing/dedupe, single-call usage fetch, and Free/paid Copilot schema parsing under `src/providers/copilot/`. | | `providers::minimax` | API key-based authentication, YapCap-owned account storage, usage quota tracking, and Minimax API integration under `src/providers/minimax/`. | +| `providers::zai` | Z.AI Coding Plan API-key login, content-aware OpenCode prefill/detection, private key storage, fixed global quota fetch, and token/MCP usage-window parsing under `src/providers/zai/`. | | `providers::kimi` | Kimi for Coding API-key account management, OpenCode key prefill, private key storage, usage fetch, and Kimi usage-window parsing under `src/providers/kimi/`. | | `providers::opencode_go` | OpenCode Go API-key account management, OpenCode key prefill and active-account matching, private key storage, usage fetch, and 5-hour/weekly/monthly window parsing under `src/providers/opencode_go/`. | -| `account_storage` | Shared explicit-account storage foundation for provider migrations. It writes account metadata, provider tokens, and per-account cached snapshots as separate JSON files under opaque YapCap-owned account directories. All write entry points (`create_account`, `replace_account`, `save_metadata`, `save_tokens`, `save_snapshot`) create the full account directory chain on demand, so callers do not need to pre-create the provider account root. It also exposes lower-level `create_private_dir`/`set_private_file_permissions`/`write_json`/`read_json` primitives (owner-only `0o700` dirs, `0o600` files) that providers with a bespoke on-disk schema — Copilot, Minimax, and Kimi — use directly instead of hand-rolling their own permission-setting code. | +| `account_storage` | Shared explicit-account storage foundation for provider migrations. It writes account metadata, provider tokens, and per-account cached snapshots as separate JSON files under opaque YapCap-owned account directories. All write entry points (`create_account`, `replace_account`, `save_metadata`, `save_tokens`, `save_snapshot`) create the full account directory chain on demand, so callers do not need to pre-create the provider account root. It also exposes lower-level `create_private_dir`/`set_private_file_permissions`/`write_json`/`read_json` primitives (owner-only `0o700` dirs, `0o600` files) that providers with a bespoke on-disk schema — Copilot, Minimax, Z.AI, and Kimi — use directly instead of hand-rolling their own permission-setting code. | | `auth` | Parses JWT identity claims used by Codex OAuth compatibility paths. | | `config` | COSMIC config entry, provider toggles, provider account preferences, and the shared app ID constant used by all COSMIC config entries. | | `shared_state` | Versioned COSMIC-backed shared runtime and shared control entries. Shared runtime wraps the app runtime payload with generation and write timestamp metadata. Shared control stores per-provider explicit refresh requests with request metadata. | @@ -243,7 +247,7 @@ sequenceDiagram request; only the refresh owner executes provider refresh and publishes shared runtime status, usage snapshots, account health, auth state, refresh errors, or runtime cleanup. -- A successful login (or Minimax or Kimi API-key save) clears the login state +- A successful login (or Minimax, Z.AI, or Kimi API-key save) clears the login state immediately: the account controls return to the normal add-account state with no confirmation message or dismiss step. Failed login preparation keeps showing the error with `Add another` / `Dismiss` controls. A failed Kimi API-key save @@ -1022,7 +1026,54 @@ Error classification (`MinimaxError`): - **Transient:** `RateLimited { retry_after_secs }`, network errors, timeouts. - **No usage data:** invalid or missing quota response preserves prior snapshot. -### 3.7 Kimi +### 3.7 Z.AI Coding Plan + +Z.AI Coding Plan uses API-key authentication and YapCap-managed accounts for the +global personal Coding Plan scope. + +- Each managed account stores non-secret metadata in `zai_managed_accounts` and + stores the API key separately at + `/yapcap/zai-accounts//api_key.txt`. The account directory is + owner-only (`0o700`) and the key file is owner-readable and owner-writable only + (`0o600`). The key is held transiently in memory in the add or reauthentication + form, masked by default with a visibility toggle. It is never written to general + COSMIC configuration or logs. +- The account label is required; duplicate labels and API keys are rejected. A + new account is selected after a successful save and requests a provider refresh. + Reauthentication preserves the existing id, label, and creation time while + replacing the stored key and authentication metadata. Saving performs local + normalization and storage only; it does not remotely validate the key. +- The add and reauthentication forms accept manual key entry. They may also + prefill once from OpenCode's local `~/.local/share/opencode/auth.json`, checking + the `zai-coding-plan` API entry first and the `zai` API entry second. Only a + usable typed API entry with a non-empty key is eligible; a bare auth file, + malformed entry, or non-API entry does not count as detection. Editing a + prefilled key clears its imported provenance. YapCap never writes to or + synchronizes with OpenCode, and Z.AI has no environment-key fallback or host + Active badge. + +Usage fetch: + +- Each refresh sends `GET https://api.z.ai/api/monitor/usage/quota/limit` with + `Accept: application/json` and `Authorization: Bearer `. This is the + fixed global endpoint used by the implementation, not a claim that Z.AI + exposes a stable public contract for it. +- A first HTTP 401 gets exactly one retry with the raw key in the authorization + header. A final 401/403 requires login. HTTP 429 is transient and preserves a + numeric `Retry-After` value when present; HTTP 5xx responses are transient and + are not retried with raw authorization. Other invalid or unsuccessful responses + preserve any prior successful snapshot. +- The parser accepts the fixture-backed legacy token and MCP response forms and + emits recognized windows in this order: **5 Hour**, **Weekly**, then optional + **MCP**. MCP is separately labelled and never receives an assumed calendar + duration; reset or duration metadata is retained only when supplied by the + response. An MCP-only response is valid and displays its MCP usage, while the + UI reports Coding Plan usage as unavailable rather than inventing token or + weekly zeroes. +- Region-specific, team, promotional, alternate-host, alternate-endpoint, OAuth, + and monthly fixed-duration variants are out of scope. + +### 3.8 Kimi Kimi for Coding uses API-key authentication and YapCap-managed accounts. @@ -1081,7 +1132,7 @@ Usage windows: seven-day duration, so both Weekly and rate-limit windows support the shared popup pace meter whenever their reset timestamp is present. -### 3.8 OpenCode Go +### 3.9 OpenCode Go OpenCode Go uses API-key authentication and YapCap-managed accounts. @@ -1113,7 +1164,7 @@ Usage fetch: - OpenCode credentials are read-only integration input; YapCap never writes to or synchronizes with OpenCode. -### 3.9 Antigravity +### 3.10 Antigravity Antigravity is a second Google Code Assist provider modeled on Gemini (§3.5): YapCap runs its own Google OAuth login and stores its own tokens; it never @@ -1311,6 +1362,7 @@ antigravity_enablement = "auto" gemini_enablement = "disabled" copilot_enablement = "auto" minimax_enablement = "auto" +zai_enablement = "auto" kimi_enablement = "auto" opencode_go_enablement = "auto" selected_codex_account_ids = [] @@ -1327,6 +1379,8 @@ selected_copilot_account_ids = [] copilot_managed_accounts = [] selected_minimax_account_ids = [] minimax_managed_accounts = [] +selected_zai_account_ids = [] +zai_managed_accounts = [] selected_kimi_account_ids = [] kimi_managed_accounts = [] selected_opencode_go_account_ids = [] @@ -1378,7 +1432,7 @@ log_level = "info" account directory is derived from the id and current runtime paths so native and Flatpak installs do not persist each other's absolute state directories. There is at most one managed Copilot account per GitHub numeric user id. -- Minimax, Kimi, and OpenCode Go each retain list-shaped selected-account ids +- Minimax, Z.AI, Kimi, and OpenCode Go each retain list-shaped selected-account ids with the same single-active-account semantics. Their managed-account entries store non-secret labels and timestamps; API keys remain in private provider account directories. New accounts reject duplicate labels and API keys. @@ -1456,7 +1510,7 @@ drive the status line and headline percentage. ```rust struct UsageSnapshot { - provider: ProviderId, // Codex | Claude | Cursor | Antigravity | Gemini | Copilot | Minimax | Kimi | OpenCodeGo + provider: ProviderId, // Codex | Claude | Cursor | Antigravity | Gemini | Copilot | Minimax | Zai | Kimi | OpenCodeGo source: String, // "OAuth" | "API Key" | "Managed Account" | ... updated_at: DateTime, headline: UsageHeadline, // index into windows for the panel badge @@ -1540,7 +1594,7 @@ struct ProviderAccountRuntimeState { - `reset_at` is present and `≤ now` (elapsed), or - `used_percent ≤ 0` and the window is in its **fresh fraction** — `now - (reset_at - window_seconds) < window_seconds / 20` (the first 5 % of the window since it last reset). When `window_seconds` or `reset_at` are missing, the fresh-fraction check degrades to "used_percent ≤ 0" so providers like Claude that can omit `resets_at` after a reset still surface the label. -Otherwise it formats `reset_at` per `ResetTimeFormat`. The rule is provider-agnostic and applies uniformly to every `UsageWindow` rendered in the popup (Codex Session/Weekly, Claude Session/Weekly plus per-model scoped windows such as Sonnet/Opus/Cowork/Fable, Cursor Total/Auto+Composer/API, Antigravity grouped Five Hour/Weekly, Gemini Pro/Flash/Lite, Copilot Free Chat/Completions, Copilot Paid Credits/Premium, Minimax Token, Kimi Weekly/Rate Limit, and OpenCode Go 5 Hour/Weekly/Monthly). +Otherwise it formats `reset_at` per `ResetTimeFormat`. The rule is provider-agnostic and applies uniformly to every `UsageWindow` rendered in the popup (Codex Session/Weekly, Claude Session/Weekly plus per-model scoped windows such as Sonnet/Opus/Cowork/Fable, Cursor Total/Auto+Composer/API, Antigravity grouped Five Hour/Weekly, Gemini Pro/Flash/Lite, Copilot Free Chat/Completions, Copilot Paid Credits/Premium, Minimax Token, Z.AI 5 Hour/Weekly/MCP, Kimi Weekly/Rate Limit, and OpenCode Go 5 Hour/Weekly/Monthly). ## 6. Persistence, Logging, Paths @@ -1555,7 +1609,7 @@ All paths come from `config::paths()`. - Managed accounts and logs: under the XDG state root (typically `~/.local/state/yapcap/`), including `codex-accounts/`, `claude-accounts/`, `cursor-accounts/`, `antigravity-accounts/`, `gemini-accounts/`, - `copilot-accounts/`, `minimax-accounts/`, `kimi-accounts/`, and + `copilot-accounts/`, `minimax-accounts/`, `zai-accounts/`, `kimi-accounts/`, and `opencode-go-accounts/` **Flatpak** (`FLATPAK_ID` set): YapCap-owned cache and state **only** under the per-app tree on the host filesystem: @@ -1565,7 +1619,7 @@ All paths come from `config::paths()`. - Managed accounts and logs: `~/.var/app//data/yapcap/`, including `codex-accounts/`, `claude-accounts/`, `cursor-accounts/`, `antigravity-accounts/`, `gemini-accounts/`, `copilot-accounts/`, - `minimax-accounts/`, `kimi-accounts/`, and `opencode-go-accounts/` + `minimax-accounts/`, `zai-accounts/`, `kimi-accounts/`, and `opencode-go-accounts/` Flatpak does **not** read or write the native install’s `~/.local/state/yapcap/` or `~/.cache/yapcap/` for YapCap data. The `~` in the `.var` paths is the passwd home directory (`pw_dir`), not `dirs::home_dir()` / `$HOME`, so locations stay correct when the sandbox overrides `HOME`. @@ -1677,7 +1731,7 @@ owns provider detail cards and `app::popup_view::settings::*` owns the settings - Navigation row: - provider detail: an icon-only, six-provider viewport. With two to five enabled providers, the icons are centered horizontally within the viewport. With exactly one enabled provider the navigation row is hidden entirely. Previous/next controls shift the viewport by one additional enabled provider without wrapping or showing usage bars; their disabled end states make the viewport bounds explicit. The selected provider has a soft neutral fill and a thick foreground underline, while the full row shares a divider baseline. Selecting an icon keeps the existing persisted selection and refresh behavior. - secondary routes: Settings, Manage providers, Manage accounts, and About do not show a navigation row and return to provider detail with Back. -- Providers render in a single fixed order (`ProviderId::ALL`) everywhere they are listed: Codex, Claude, Cursor, Antigravity, Gemini, Copilot, Minimax, Kimi, OpenCode Go. +- Providers render in a single fixed order (`ProviderId::ALL`) everywhere they are listed: Codex, Claude, Cursor, Antigravity, Gemini, Copilot, Minimax, Z.AI Coding Plan, Kimi, OpenCode Go. - Provider and settings tabs and selected account rows use a soft accent fill and accent border. Global settings segmented option groups are softly merged neutral component surfaces; their selected option has a slightly deeper neutral fill plus accent-colored text and a checkmark. Settings section wrappers around titles and bodies stay visually neutral (layout only). - Popup child surfaces follow the active COSMIC applet transparency state. Neutral component containers and custom neutral buttons use a 40% alpha overlay in transparent mode, including provider Account and Usage cards, while retaining their normal component fill when transparency is disabled. Accent and status overlays retain their semantic colors and interaction states. - Body panel (scrollable where content can grow): shows either selected provider details, global Settings, Manage providers, provider-scoped Manage accounts, or About. Manage providers lists every provider in one rounded component card with horizontal row dividers and trailing enable switches. Settings retains the Refresh interval, panel-icon, reset-time, and usage-amount controls. About centers the YapCap logo and identity, groups project/developer/license links into full-width link rows, and shows checking/error/update state; an available update uses a destructive callout that links to its release and drives the header notification dot. When no provider tabs are available, the provider route suppresses the navigation row and shows a centered YapCap/provider-logo hero with “No providers set up yet”, guidance to manage providers, and a suggested action. @@ -1687,7 +1741,7 @@ owns provider detail cards and `app::popup_view::settings::*` owns the settings - Global Settings contains app-wide settings such as Autorefresh segmented interval buttons, panel icon style preview buttons, reset time format, and usage amount format. Each selectable option shows a tooltip explaining its effect on hover. If the startup update check fails, YapCap keeps retrying in the background with exponential backoff and shows the latest detailed failure plus the next retry delay in About. Error state also shows a manual "Check again" action. - When an update is available, a small red notification dot appears on the header About action. - Debug builds can force the About update-available state with `YAPCAP_DEBUG_UPDATE_AVAILABLE`. Values `1`, `true`, `yes`, and empty string use `v9.9.9`; any other value is treated as the release version. Debug builds can also simulate offline HTTP with `YAPCAP_DEBUG_OFFLINE`; values `0`, `false`, `no`, and `off` disable it, while any other present value enables it. - - `YAPCAP_DEMO` (debug only; inert in release) seeds a screenshot-oriented synthetic config plus `AppState`: all nine providers are enabled with `provider_visibility_mode = user_managed`; **Codex** gets two managed demo accounts, with the Pro account selected; **Claude** gets two managed demo accounts, with the Pro account selected and synthetic **extra usage**; **Cursor** gets one managed demo account; **Gemini** gets one Pro-tier managed demo account; **Minimax** gets one managed demo account; **Copilot** gets two managed demo accounts, with the Free account selected and the Pro+ account carrying a **Credits** window, dollar cost card (`$28.00 / $70.00`), and `+42 over plan`; **Antigravity** gets two managed demo accounts, with the Pro account selected; and **Kimi** and **OpenCode Go** each get one API-key demo account. Multiple managed demo accounts remain available through the account pager, but only one account is selected per provider. Every synthetic provider usage window includes pace timing so demo behavior matches production. Display settings otherwise follow defaults (panel icon style, reset time format, usage format, autorefresh interval); the default startup `Task` batch is skipped; provider refresh becomes a no-op; shared-runtime writes are skipped; and demo data is re-applied after config reconciliation. + - `YAPCAP_DEMO` (debug only; inert in release) seeds a screenshot-oriented synthetic config plus `AppState`: all ten providers are enabled with `provider_visibility_mode = user_managed`; **Codex** gets two managed demo accounts, with the Pro account selected; **Claude** gets two managed demo accounts, with the Pro account selected and synthetic **extra usage**; **Cursor** gets one managed demo account; **Gemini** gets one Pro-tier managed demo account; **Minimax** gets one managed demo account; **Z.AI** gets one managed Coding Plan demo account; **Copilot** gets two managed demo accounts, with the Free account selected and the Pro+ account carrying a **Credits** window, dollar cost card (`$28.00 / $70.00`), and `+42 over plan`; **Antigravity** gets two managed demo accounts, with the Pro account selected; and **Kimi** and **OpenCode Go** each get one API-key demo account. Multiple managed demo accounts remain available through the account pager, but only one account is selected per provider. Every synthetic provider usage window includes pace timing so demo behavior matches production. Display settings otherwise follow defaults (panel icon style, reset time format, usage format, autorefresh interval); the default startup `Task` batch is skipped; provider refresh becomes a no-op; shared-runtime writes are skipped; and demo data is re-applied after config reconciliation. - `YAPCAP_DEMO` and `YAPCAP_DEBUG_UPDATE_AVAILABLE` are independent debug toggles and can be combined; `just run-demo-update` launches the synthetic demo with the forced update state. - Provider account cards list currently valid account sources as separate selector rows with a selected outline/checkmark, a row press to make an account active, and account action icons. Long account labels are truncated in-row and reveal the full label on hover. With no accounts, the provider settings page shows a tighter centered `No accounts` card with `Add account`; providers with a currently available OpenCode import also show `Import from OpenCode` in that card. With accounts present, account-creation controls are grouped in a separate bordered container; each action is a full-width row with a trailing arrow matching the provider-detail Manage accounts action. Providers with an OpenCode import path show that action as a second stacked row. Codex add-account login opens the browser from the Settings flow and stores the result in YapCap-owned account storage. Codex account rows show the same login-required warning badge and row highlight as other providers when `auth_state = ActionRequired` (for example after refresh token failure). Claude add-account opens the native OAuth browser flow from Settings, shows the same browser account/private-window hint as Copilot, and asks the user to paste the returned authentication code; malformed pasted input is rejected with plain-language guidance to paste the authentication code (no internal format jargon). Claude account rows use email-derived labels and show login-required, error, or stale badges when account state needs attention. Claude accounts with `auth_state = ActionRequired` show a per-account re-authenticate action (refresh icon) in Settings alongside the delete action; clicking it starts a targeted OAuth flow that must complete with the same email — a different email is rejected with an error and the existing account is left unchanged; success immediately triggers a usage refresh. Generic Claude add-account keeps duplicate-by-email upsert behavior. Cursor add-account scans Cursor IDE's local SQLite state database and imports the currently logged-in Cursor account tokens into YapCap-owned storage. Cursor accounts that need user action show a `Re-auth needed` badge plus a per-account refresh action in Settings, and the provider status text tells the user to log into that account in Cursor and rescan. Cursor `Active` reflects the account currently used by Cursor IDE and can appear alongside `Re-auth needed` when YapCap's copied session needs a fresh scan. Antigravity and Gemini add-account open Google OAuth browser flows and store only YapCap-owned account storage. Copilot add-account starts GitHub device flow, shows the shared browser account/private-window hint near the Settings control, displays the user code and `Open Browser` fallback while polling, and stores accounts by GitHub numeric user id. Minimax, Kimi, and OpenCode Go use API-key forms and may prefill once from OpenCode. Copilot and Antigravity account rows never show an Active badge. Codex, Claude, Cursor, Antigravity, Gemini, Copilot, Minimax, Kimi, and OpenCode Go account removal deletes only YapCap-owned account homes/config dirs/profile roots. Cursor accounts are always managed and displayed with the email address as the account label. Copilot accounts are displayed with the GitHub login label. If no accounts remain for a provider, the provider detail shows an empty state pointing the user to Settings. - Footer: Back on secondary routes. @@ -1746,7 +1800,7 @@ Most user-visible strings in `src/app/popup_view.rs`, `src/app/popup_view/detail ## 10. Testing - `cargo test` runs unit and integration tests covering: config defaults and legacy-field compatibility, usage display formatting, app-state helpers, model status/headline helpers, provider normalization and error paths, managed-account storage and selection, OAuth and API-key login flows, refresh/backoff state machines, update checks, debug/demo behavior, provider adapters, and app-level state transitions. -- No tests hit real provider APIs. Fixtures under `fixtures/{antigravity,claude,codex,copilot,cursor,gemini}/` are redacted probe captures or handcrafted JSON. Copilot uses device-code, OAuth-token, GitHub identity, and `copilot_internal/user` captures; Cursor uses usage, identity, and OAuth-token captures; Gemini and Antigravity use OAuth, Code Assist, and quota captures. Minimax and Kimi response shapes are covered by provider-local test payloads rather than fixture directories. +- No tests hit real provider APIs. Fixtures under `fixtures/{antigravity,claude,codex,copilot,cursor,gemini,zai}/` are redacted probe captures or handcrafted JSON. Copilot uses device-code, OAuth-token, GitHub identity, and `copilot_internal/user` captures; Cursor uses usage, identity, and OAuth-token captures; Gemini and Antigravity use OAuth, Code Assist, and quota captures; Z.AI uses fixture-backed token and MCP response shapes. Minimax and Kimi response shapes are covered by provider-local test payloads rather than fixture directories. - `cargo clippy` and `cargo fmt --check` are expected clean on main. - Tests must never read or write the developer's real COSMIC config or state. `config::cosmic_config_context()` is the only way to open a `cosmic_config::Config`; under `cfg(test)` it resolves to a per-test-thread temporary root via `Config::with_custom_path`, so isolation does not depend on a test remembering a guard. `test_support::test_env()` additionally points `XDG_CONFIG_HOME`, `XDG_STATE_HOME`, and `XDG_CACHE_HOME` at a temporary root and clears `FLATPAK_ID`, which is what isolates `config::paths()` for tests that touch account storage. Opening `cosmic_config::Config::new` directly from app code reintroduces the hazard: a test that constructs an `AppModel` with a default config and reaches `write_config` will overwrite the real account registry. - Manual QA should cover: install via `just install`, each provider's auth refresh flow, transient provider failures showing "Stale" not "Error", stale shared-runtime display on cold-start, settings persistence across restarts, multi-process two-display sync and owner takeover, update-check UI states, and dark/light theme icon variants. diff --git a/fixtures/zai/README.md b/fixtures/zai/README.md new file mode 100644 index 0000000..e4a1f0a --- /dev/null +++ b/fixtures/zai/README.md @@ -0,0 +1,15 @@ +# Z.AI quota fixtures + +These sanitized responses exercise the fixed `GET https://api.z.ai/api/monitor/usage/quota/limit` +endpoint without credentials or personal data. + +The parser sends `Authorization: Bearer ` first and permits one retry with the raw +authorization value only after HTTP 401. `CREDIT_LIMIT` and `TOKENS_LIMIT` rows with `(unit: 3, +number: 5)` are the five-hour slot; `(unit: 6, number: 1)` is the weekly slot. A fixture-backed +`TIME_LIMIT` `(unit: 5, number: 1)` is rendered as MCP. The MCP fixture intentionally omits a +fixed duration because a reset instant does not establish a calendar-month span. + +`percentage` is preferred over counters. When it is unavailable, `currentValue / usage` or +`(usage - remaining) / usage` supplies the used percentage only when the values are finite and +the usage limit is positive. Reset timestamps are epoch milliseconds. The responses contain only +synthetic values and are not raw captures. diff --git a/fixtures/zai/credit_limit.json b/fixtures/zai/credit_limit.json new file mode 100644 index 0000000..19fbbb1 --- /dev/null +++ b/fixtures/zai/credit_limit.json @@ -0,0 +1,29 @@ +{ + "code": 200, + "success": true, + "data": { + "level": "pro", + "limits": [ + { + "type": "CREDIT_LIMIT", + "unit": 6, + "number": 1, + "usage": 100000, + "currentValue": 100001, + "remaining": 0, + "percentage": 100, + "nextResetTime": 1790000000000 + }, + { + "type": "CREDIT_LIMIT", + "unit": 3, + "number": 5, + "usage": 100000, + "currentValue": 0, + "remaining": 100000, + "percentage": 0, + "nextResetTime": 1789000000000 + } + ] + } +} diff --git a/fixtures/zai/mcp_only.json b/fixtures/zai/mcp_only.json new file mode 100644 index 0000000..6a08511 --- /dev/null +++ b/fixtures/zai/mcp_only.json @@ -0,0 +1,19 @@ +{ + "code": 200, + "success": true, + "data": { + "level": "lite", + "limits": [ + { + "type": "TIME_LIMIT", + "unit": 5, + "number": 1, + "usage": 1000, + "currentValue": 250, + "remaining": 750, + "percentage": 25, + "nextResetTime": 1791000000000 + } + ] + } +} diff --git a/fixtures/zai/token_mcp.json b/fixtures/zai/token_mcp.json new file mode 100644 index 0000000..7a91a4d --- /dev/null +++ b/fixtures/zai/token_mcp.json @@ -0,0 +1,39 @@ +{ + "code": 200, + "success": true, + "data": { + "level": "max", + "limits": [ + { + "type": "TIME_LIMIT", + "unit": 5, + "number": 1, + "usage": 1000, + "currentValue": 300, + "remaining": 700, + "percentage": 30, + "nextResetTime": 1791000000000 + }, + { + "type": "TOKENS_LIMIT", + "unit": 6, + "number": 1, + "usage": 100000, + "currentValue": 25000, + "remaining": 75000, + "percentage": 25, + "nextResetTime": 1790000000000 + }, + { + "type": "TOKENS_LIMIT", + "unit": 3, + "number": 5, + "usage": 100000, + "currentValue": 10000, + "remaining": 90000, + "percentage": 10, + "nextResetTime": 1789000000000 + } + ] + } +} diff --git a/i18n/en/yapcap.ftl b/i18n/en/yapcap.ftl index 0cd2a4c..a1f3bc4 100644 --- a/i18n/en/yapcap.ftl +++ b/i18n/en/yapcap.ftl @@ -78,6 +78,13 @@ gemini-account-reauth-tooltip = Re-authenticate this Gemini account minimax-accounts-title = Minimax Accounts minimax-account-select-required = Select a Minimax account before refreshing minimax-account-reauth-tooltip = Re-authenticate this Minimax account +zai-accounts-title = Z.AI Coding Plan Accounts +zai-account-select-required = Select a Z.AI Coding Plan account before refreshing +zai-account-reauth-tooltip = Re-authenticate this Z.AI Coding Plan account +zai-login-editing = Enter your Z.AI Coding Plan API key. +zai-login-failed = Failed to save Z.AI Coding Plan account. +zai-api-key-placeholder = API Key +zai-api-key-imported-from-opencode = API key pre-filled from OpenCode. You can edit or clear it. kimi-accounts-title = Kimi Accounts kimi-account-select-required = Select a Kimi account before refreshing kimi-account-reauth-tooltip = Re-authenticate this Kimi account @@ -180,6 +187,8 @@ extra-usage-disabled = Disabled credits-label = Credits credits-available = { $balance } available copilot-overage-over-plan = +{ $count } over plan +zai-coding-plan-unavailable-title = Coding Plan usage unavailable +zai-coding-plan-unavailable-detail = Coding Plan quotas were not reported by Z.AI. ## Provider status badges badge-disabled = Disabled diff --git a/i18n/pl/yapcap.ftl b/i18n/pl/yapcap.ftl index 9bf35bb..c2d9a61 100644 --- a/i18n/pl/yapcap.ftl +++ b/i18n/pl/yapcap.ftl @@ -78,6 +78,13 @@ gemini-account-reauth-tooltip = Uwierzytelnij ponownie to konto Gemini minimax-accounts-title = Konta Minimax minimax-account-select-required = Wybierz konto Minimax przed odświeżeniem minimax-account-reauth-tooltip = Uwierzytelnij ponownie to konto Minimax +zai-accounts-title = Konta Z.AI Coding Plan +zai-account-select-required = Wybierz konto Z.AI Coding Plan przed odświeżeniem +zai-account-reauth-tooltip = Uwierzytelnij ponownie to konto Z.AI Coding Plan +zai-login-editing = Wprowadź swój klucz API Z.AI Coding Plan. +zai-login-failed = Nie udało się zapisać konta Z.AI Coding Plan. +zai-api-key-placeholder = Klucz API +zai-api-key-imported-from-opencode = Klucz API został wstępnie uzupełniony z OpenCode. Możesz go zmienić lub usunąć. kimi-accounts-title = Konta Kimi kimi-account-select-required = Wybierz konto Kimi przed odświeżeniem kimi-account-reauth-tooltip = Uwierzytelnij ponownie to konto Kimi @@ -173,6 +180,8 @@ extra-usage-disabled = Zablokowane credits-label = Kredyty credits-available = { $balance } dostępne copilot-overage-over-plan = +{ $count } ponad plan +zai-coding-plan-unavailable-title = Dane o zużyciu Coding Plan są niedostępne +zai-coding-plan-unavailable-detail = Z.AI zwróciło tylko dane o zużyciu narzędzi MCP. Z.AI nie zwróciło danych o limitach Coding Plan. ## Provider status badges badge-disabled = Zablokowane diff --git a/resources/providers/zai-reversed.svg b/resources/providers/zai-reversed.svg new file mode 100644 index 0000000..963db05 --- /dev/null +++ b/resources/providers/zai-reversed.svg @@ -0,0 +1,4 @@ + + Z.AI Coding Plan + + diff --git a/resources/providers/zai.svg b/resources/providers/zai.svg new file mode 100644 index 0000000..530c690 --- /dev/null +++ b/resources/providers/zai.svg @@ -0,0 +1,4 @@ + + Z.AI Coding Plan + + diff --git a/src/account_storage/mod.rs b/src/account_storage/mod.rs index 5571671..2afe07c 100644 --- a/src/account_storage/mod.rs +++ b/src/account_storage/mod.rs @@ -355,6 +355,7 @@ impl ProviderAccountStorage { ProviderId::Gemini => "gemini", ProviderId::Copilot => "copilot", ProviderId::Minimax => "minimax", + ProviderId::Zai => "zai", ProviderId::Kimi => "kimi", ProviderId::Antigravity => "antigravity", ProviderId::OpenCodeGo => "opencode_go", diff --git a/src/account_storage/tests.rs b/src/account_storage/tests.rs index b3d906a..a1a4085 100644 --- a/src/account_storage/tests.rs +++ b/src/account_storage/tests.rs @@ -201,6 +201,24 @@ fn create_account_creates_missing_provider_root() { assert!(stored.account_dir.join(TOKENS_FILE).exists()); } +#[test] +fn creates_zai_account_ids_with_zai_prefix() { + let storage = ProviderAccountStorage::new(test_dir("zai-prefix")); + let stored = storage + .create_account(NewProviderAccount { + provider: ProviderId::Zai, + email: String::new(), + provider_account_id: None, + organization_id: None, + organization_name: None, + tokens: tokens(), + snapshot: None, + }) + .unwrap(); + + assert!(stored.account_ref.account_id.starts_with("zai-")); +} + #[cfg(unix)] #[test] fn managed_account_storage_repairs_directory_and_file_permissions() { diff --git a/src/app/login/flows/mod.rs b/src/app/login/flows/mod.rs index 10a4306..a7dc974 100644 --- a/src/app/login/flows/mod.rs +++ b/src/app/login/flows/mod.rs @@ -7,6 +7,7 @@ mod grok; mod kimi; mod minimax; mod opencode_go; +mod zai; pub(crate) use antigravity::AntigravityLoginFlow; pub(crate) use claude::ClaudeLoginFlow; @@ -17,3 +18,4 @@ pub(crate) use grok::GrokLoginFlow; pub(crate) use kimi::KimiLoginFlow; pub(crate) use minimax::MinimaxLoginFlow; pub(crate) use opencode_go::OpenCodeGoLoginFlow; +pub(crate) use zai::ZaiLoginFlow; diff --git a/src/app/login/flows/zai.rs b/src/app/login/flows/zai.rs new file mode 100644 index 0000000..65ef6e7 --- /dev/null +++ b/src/app/login/flows/zai.rs @@ -0,0 +1,180 @@ +use super::super::{LoginEventKind, LoginFlow, apply_login_success_checked}; +use crate::account_selection::select_account_after_login; +use crate::app::{AppModel, Config, Handle, Message, ProviderId, Task}; +use crate::providers::zai; +use crate::providers::zai::login::{ZaiLoginEvent, ZaiLoginState, ZaiLoginStatus}; + +pub(crate) struct ZaiLoginFlow; + +impl LoginFlow for ZaiLoginFlow { + type State = ZaiLoginState; + type Event = ZaiLoginEvent; + const PROVIDER: ProviderId = ProviderId::Zai; + + fn state(app: &AppModel) -> &Option { + &app.zai_login + } + + fn state_mut(app: &mut AppModel) -> &mut Option { + &mut app.zai_login + } + + fn handle_mut(app: &mut AppModel) -> &mut Option { + &mut app.zai_login_handle + } + + fn is_running(state: &Self::State) -> bool { + state.status == ZaiLoginStatus::Editing + } + + fn log_id(state: &Self::State) -> &str { + &state.account_id + } + + fn status_debug(state: &Self::State) -> String { + format!("{:?}", state.status) + } + + fn account_exists(config: &Config, account_id: &str) -> bool { + config + .zai_managed_accounts + .iter() + .any(|account| account.id == account_id) + } + + fn failed_state(error: String) -> Self::State { + ZaiLoginState::failed(error) + } + + fn prepare(config: Config) -> Result<(Self::State, cosmic::iced::Task), String> { + let _ = config; + Ok((zai::login::prepare(), cosmic::iced::Task::none())) + } + + fn prepare_for_reauth( + config: Config, + account_id: &str, + ) -> Result<(Self::State, cosmic::iced::Task), String> { + Ok(( + zai::login::prepare_for_reauth(config, account_id)?, + cosmic::iced::Task::none(), + )) + } + + fn wrap_event(event: Self::Event) -> Message { + Message::LoginEvent(ProviderId::Zai, Box::new(LoginEventKind::Zai(event))) + } + + fn on_event(app: &mut AppModel, event: Self::Event) -> Task { + match event { + ZaiLoginEvent::ApiKeyChanged(api_key) => { + if let Some(login) = app.zai_login.as_mut() { + login.update_api_key(api_key); + } + Task::none() + } + ZaiLoginEvent::ApiKeyVisibilityToggled => { + if let Some(login) = app.zai_login.as_mut() { + login.toggle_api_key_visibility(); + } + Task::none() + } + ZaiLoginEvent::LabelChanged(label) => { + if let Some(login) = app.zai_login.as_mut() { + login.update_label(label); + } + Task::none() + } + ZaiLoginEvent::Saved => { + let Some(login) = app.zai_login.as_mut() else { + return Task::none(); + }; + let flow_id = login.account_id.clone(); + let is_reauthentication = app + .config + .zai_managed_accounts + .iter() + .any(|account| account.id == flow_id); + match zai::login::save(&app.config, login) { + Ok(managed_account) => { + let account_id = managed_account.id.clone(); + let selected_account_id = account_id.clone(); + let result = apply_login_success_checked( + app, + ProviderId::Zai, + &flow_id, + account_id.clone(), + move |config| { + zai::account::apply_login_account(config, managed_account); + select_account_after_login( + config, + ProviderId::Zai, + selected_account_id, + ); + }, + ); + match result { + Ok(task) => { + app.zai_login_handle = None; + app.zai_login = None; + task + } + Err(()) => { + cleanup_new_account_storage(&account_id, is_reauthentication); + if let Some(login) = app.zai_login.as_mut() { + login.status = ZaiLoginStatus::Editing; + login.error = Some( + "Failed to save Z.AI account configuration".to_string(), + ); + } + Task::none() + } + } + } + Err(_) => Task::none(), + } + } + } + } +} + +fn cleanup_new_account_storage(account_id: &str, is_reauthentication: bool) { + if is_reauthentication { + return; + } + if let Err(error) = zai::storage::delete_account(account_id) { + tracing::error!( + account_id, + error = %error, + "failed to clean up new Z.AI account storage" + ); + } +} + +#[cfg(test)] +mod tests { + use super::cleanup_new_account_storage; + use crate::providers::zai::storage; + + #[test] + fn config_failure_removes_new_account_storage() { + let _env = crate::test_support::test_env(); + let account_id = "zai-config-failure-new"; + storage::write_api_key(account_id, "new-key").unwrap(); + + cleanup_new_account_storage(account_id, false); + + assert!(storage::load_api_key(account_id).is_err()); + } + + #[test] + fn config_failure_does_not_replace_reauthentication_candidate() { + let _env = crate::test_support::test_env(); + let account_id = "zai-config-failure-existing"; + storage::write_api_key(account_id, "candidate-key").unwrap(); + + cleanup_new_account_storage(account_id, true); + + assert_eq!(storage::load_api_key(account_id).unwrap(), "candidate-key"); + } +} diff --git a/src/app/login/login_flow_tests/key_authentication.rs b/src/app/login/login_flow_tests/key_authentication.rs index 61cdbd1..8f562dc 100644 --- a/src/app/login/login_flow_tests/key_authentication.rs +++ b/src/app/login/login_flow_tests/key_authentication.rs @@ -1,5 +1,5 @@ use super::key_authentication_cases::{ - KeyAuthenticationCase, KimiCase, MinimaxCase, OpenCodeGoCase, + KeyAuthenticationCase, KimiCase, MinimaxCase, OpenCodeGoCase, ZaiCase, }; use super::support::{isolated_xdg, test_app}; use crate::app::AppModel; @@ -248,6 +248,7 @@ fn key_authentication_rejects_empty_keys_for_all_api_key_providers() { assert_empty_key_is_editable::("common-empty-kimi"); assert_empty_key_is_editable::("common-empty-minimax"); assert_empty_key_is_editable::("common-empty-opencode-go"); + assert_empty_key_is_editable::("common-empty-zai"); } #[test] @@ -255,6 +256,7 @@ fn key_authentication_shares_masking_and_visibility_for_all_api_key_providers() assert_masking_and_visibility_are_shared::("common-visibility-kimi"); assert_masking_and_visibility_are_shared::("common-visibility-minimax"); assert_masking_and_visibility_are_shared::("common-visibility-opencode-go"); + assert_masking_and_visibility_are_shared::("common-visibility-zai"); } #[test] @@ -264,6 +266,7 @@ fn key_authentication_shares_import_provenance_for_all_api_key_providers() { assert_imported_provenance_is_cleared_by_input::( "common-provenance-opencode-go", ); + assert_imported_provenance_is_cleared_by_input::("common-provenance-zai"); } #[cfg(unix)] @@ -272,6 +275,7 @@ fn key_authentication_preserves_editable_forms_after_save_failure_for_all_provid assert_save_failure_is_editable::("common-failure-kimi"); assert_save_failure_is_editable::("common-failure-minimax"); assert_save_failure_is_editable::("common-failure-opencode-go"); + assert_save_failure_is_editable::("common-failure-zai"); } #[test] @@ -279,6 +283,7 @@ fn key_authentication_success_selects_and_requests_refresh_for_all_providers() { assert_save_selects_and_refreshes::("common-success-kimi"); assert_save_selects_and_refreshes::("common-success-minimax"); assert_save_selects_and_refreshes::("common-success-opencode-go"); + assert_save_selects_and_refreshes::("common-success-zai"); } #[test] @@ -286,6 +291,7 @@ fn key_authentication_reauthentication_preserves_identity_for_all_providers() { assert_reauthentication_preserves_identity::("common-reauth-kimi"); assert_reauthentication_preserves_identity::("common-reauth-minimax"); assert_reauthentication_preserves_identity::("common-reauth-opencode-go"); + assert_reauthentication_preserves_identity::("common-reauth-zai"); } #[test] @@ -293,6 +299,7 @@ fn key_authentication_rejects_empty_account_names_for_all_providers() { assert_rejects_invalid_account_details::("common-empty-name-kimi"); assert_rejects_invalid_account_details::("common-empty-name-minimax"); assert_rejects_invalid_account_details::("common-empty-name-opencode-go"); + assert_rejects_invalid_account_details::("common-empty-name-zai"); } #[test] @@ -300,4 +307,5 @@ fn key_authentication_rejects_duplicate_names_and_keys_for_all_providers() { assert_rejects_duplicate_account_details::("common-duplicate-kimi"); assert_rejects_duplicate_account_details::("common-duplicate-minimax"); assert_rejects_duplicate_account_details::("common-duplicate-opencode-go"); + assert_rejects_duplicate_account_details::("common-duplicate-zai"); } diff --git a/src/app/login/login_flow_tests/key_authentication_cases.rs b/src/app/login/login_flow_tests/key_authentication_cases.rs index 8d0c89d..b37735e 100644 --- a/src/app/login/login_flow_tests/key_authentication_cases.rs +++ b/src/app/login/login_flow_tests/key_authentication_cases.rs @@ -2,12 +2,14 @@ use crate::app::AppModel; use crate::app::login::LoginFlow; use crate::config::{ Config, ManagedKimiAccountConfig, ManagedMinimaxAccountConfig, ManagedOpenCodeGoAccountConfig, + ManagedZaiAccountConfig, }; use crate::key_authentication::KeyAuthenticationState; use crate::model::ProviderId; use crate::providers::kimi::storage as kimi_storage; use crate::providers::minimax::storage as minimax_storage; use crate::providers::opencode_go::storage as opencode_go_storage; +use crate::providers::zai::storage as zai_storage; use chrono::{DateTime, Utc}; use std::path::{Path, PathBuf}; @@ -186,6 +188,57 @@ impl KeyAuthenticationCase for OpenCodeGoCase { } } +pub(super) struct ZaiCase; + +impl KeyAuthenticationCase for ZaiCase { + type Flow = super::super::ZaiLoginFlow; + + const PROVIDER: ProviderId = ProviderId::Zai; + + fn state(app: &AppModel) -> Option<&KeyAuthenticationState> { + app.zai_login.as_ref() + } + + fn opencode_auth(key: &str) -> String { + format!(r#"{{"zai-coding-plan":{{"type":"api","key":"{key}"}}}}"#) + } + + fn storage_root(root: &Path) -> PathBuf { + root.join("yapcap/zai-accounts") + } + + fn account_facts(config: &Config, account_id: &str) -> Option { + config + .zai_managed_accounts + .iter() + .find(|account| account.id == account_id) + .map(AccountFacts::from_zai) + } + + fn set_account(config: &mut Config, account_id: &str, label: &str, created_at: DateTime) { + config.zai_managed_accounts.push(ManagedZaiAccountConfig { + id: account_id.to_string(), + label: label.to_string(), + api_key_source: "stored".to_string(), + created_at, + updated_at: created_at, + last_authenticated_at: Some(created_at), + }); + } + + fn load_api_key(account_id: &str) -> Result { + zai_storage::load_api_key(account_id) + } + + fn write_api_key(account_id: &str, api_key: &str) -> Result<(), String> { + zai_storage::write_api_key(account_id, api_key) + } + + fn save_error_prefix() -> &'static str { + "Failed to save Z.AI API key:" + } +} + #[derive(Debug)] pub(super) struct AccountFacts { pub label: String, @@ -221,4 +274,13 @@ impl AccountFacts { last_authenticated_at: account.last_authenticated_at, } } + + fn from_zai(account: &ManagedZaiAccountConfig) -> Self { + Self { + label: account.label.clone(), + created_at: account.created_at, + updated_at: account.updated_at, + last_authenticated_at: account.last_authenticated_at, + } + } } diff --git a/src/app/login/login_flow_tests/mod.rs b/src/app/login/login_flow_tests/mod.rs index 780d19f..c2f2415 100644 --- a/src/app/login/login_flow_tests/mod.rs +++ b/src/app/login/login_flow_tests/mod.rs @@ -12,3 +12,4 @@ mod key_authentication_cases; mod kimi; mod minimax; mod opencode_go; +mod zai; diff --git a/src/app/login/login_flow_tests/support.rs b/src/app/login/login_flow_tests/support.rs index 9056ec8..37794fb 100644 --- a/src/app/login/login_flow_tests/support.rs +++ b/src/app/login/login_flow_tests/support.rs @@ -51,6 +51,8 @@ pub(super) fn test_app() -> AppModel { opencode_go_login_handle: None, grok_login: None, grok_login_handle: None, + zai_login: None, + zai_login_handle: None, } } diff --git a/src/app/login/login_flow_tests/zai.rs b/src/app/login/login_flow_tests/zai.rs new file mode 100644 index 0000000..838bc83 --- /dev/null +++ b/src/app/login/login_flow_tests/zai.rs @@ -0,0 +1,35 @@ +use super::support::{isolated_xdg, test_app}; +use crate::app::login::{ZaiLoginFlow, start_login}; +use crate::config::Config; +use crate::providers::zai::login::prepare_for_reauth; +use std::fs; + +#[test] +fn zai_login_uses_the_zai_opencode_credential_for_prefill() { + let (mut env, root) = isolated_xdg("zai-opencode-prefill"); + fs::create_dir_all(&root).unwrap(); + let auth_path = root.join("auth.json"); + fs::write( + &auth_path, + r#"{"zai-coding-plan":{"type":"api","key":"fake-zai-key"}}"#, + ) + .unwrap(); + env.set("YAPCAP_OPENCODE_AUTH_PATH", &auth_path); + let mut app = test_app(); + + let _ = start_login::(&mut app); + + let login = app.zai_login.as_ref().unwrap(); + assert_eq!(login.api_key, "fake-zai-key"); + assert!(login.api_key_from_opencode); +} + +#[test] +fn zai_reauth_reports_a_zai_specific_missing_account_error() { + let (_env, _root) = isolated_xdg("zai-missing-account"); + + assert_eq!( + prepare_for_reauth(Config::default(), "missing").unwrap_err(), + "Z.AI account not found" + ); +} diff --git a/src/app/login/mod.rs b/src/app/login/mod.rs index 9fc685d..1220952 100644 --- a/src/app/login/mod.rs +++ b/src/app/login/mod.rs @@ -3,7 +3,7 @@ mod legacy; pub(crate) use flows::{ AntigravityLoginFlow, ClaudeLoginFlow, CodexLoginFlow, CopilotLoginFlow, GeminiLoginFlow, - GrokLoginFlow, KimiLoginFlow, MinimaxLoginFlow, OpenCodeGoLoginFlow, + GrokLoginFlow, KimiLoginFlow, MinimaxLoginFlow, OpenCodeGoLoginFlow, ZaiLoginFlow, }; use super::{ @@ -136,13 +136,36 @@ fn apply_login_success( account_id, "login flow succeeded" ); - app.write_config(|new_config| apply(new_config)); + let _ = app.write_config(|new_config| apply(new_config)); runtime::reconcile_provider(&app.config, &app.detection, &mut app.state, provider); runtime::mark_account_reauthenticated(&mut app.state, provider, &account_id); app.sync_panel_suggested_bounds(); app.request_provider_refresh(provider, RefreshRequestReason::AccountAction) } +pub(super) fn apply_login_success_checked( + app: &mut AppModel, + provider: ProviderId, + flow_id: &str, + account_id: String, + apply: impl FnOnce(&mut Config), +) -> Result, ()> { + if !app.write_config(|new_config| apply(new_config)) { + return Err(()); + } + tracing::info!( + process_id = %app.process_info.id, + provider = provider.label(), + flow_id, + account_id, + "login flow succeeded" + ); + runtime::reconcile_provider(&app.config, &app.detection, &mut app.state, provider); + runtime::mark_account_reauthenticated(&mut app.state, provider, &account_id); + app.sync_panel_suggested_bounds(); + Ok(app.request_provider_refresh(provider, RefreshRequestReason::AccountAction)) +} + fn log_login_failed(process_id: &str, provider: ProviderId, flow_id: &str, error: &str) { tracing::info!( process_id, @@ -222,6 +245,7 @@ pub(crate) enum LoginEventKind { Antigravity(AntigravityLoginEvent), OpenCodeGo(OpenCodeGoLoginEvent), Grok(GrokLoginEvent), + Zai(crate::providers::zai::login::ZaiLoginEvent), } #[cfg(test)] diff --git a/src/app/mod.rs b/src/app/mod.rs index fa47a5a..b1deecb 100644 --- a/src/app/mod.rs +++ b/src/app/mod.rs @@ -54,6 +54,7 @@ use crate::providers::kimi::{ use crate::providers::minimax::{self, MinimaxLoginEvent, MinimaxLoginState}; use crate::providers::opencode_go::login::{OpenCodeGoLoginEvent, OpenCodeGoLoginState}; use crate::providers::registry; +use crate::providers::zai::ZaiLoginState; use crate::refresh_owner::{ self, ProcessInfo, RefreshOwner, RefreshOwnerAttempt, RefreshOwnerWaiter, }; @@ -135,6 +136,8 @@ pub struct AppModel { opencode_go_login_handle: Option, pub grok_login: Option, pub grok_login_handle: Option, + zai_login: Option, + zai_login_handle: Option, } impl Drop for AppModel { @@ -318,6 +321,8 @@ impl cosmic::Application for AppModel { opencode_go_login_handle: None, grok_login: None, grok_login_handle: None, + zai_login: None, + zai_login_handle: None, }; tracing::info!( pid = app.process_info.pid, @@ -401,6 +406,7 @@ impl cosmic::Application for AppModel { antigravity: self.antigravity_login.as_ref(), opencode_go: self.opencode_go_login.as_ref(), grok: self.grok_login.as_ref(), + zai: self.zai_login.as_ref(), }, popup_view::DetailSelection { provider: self.selected_provider, @@ -602,6 +608,9 @@ impl AppModel { (ProviderId::Grok, login::LoginEventKind::Grok(event)) => { login::GrokLoginFlow::on_event(self, event) } + (ProviderId::Zai, login::LoginEventKind::Zai(event)) => { + login::ZaiLoginFlow::on_event(self, event) + } _ => Task::none(), }); } diff --git a/src/app/popup_view.rs b/src/app/popup_view.rs index af8fe7a..a52b703 100644 --- a/src/app/popup_view.rs +++ b/src/app/popup_view.rs @@ -27,6 +27,7 @@ use crate::providers::cursor::CursorScanState; use crate::providers::gemini::{GeminiLoginState, GeminiLoginStatus}; use crate::providers::kimi::login::KimiLoginState; use crate::providers::minimax::MinimaxLoginState; +use crate::providers::zai::login::ZaiLoginState; use crate::updates::UpdateStatus; use crate::usage_display; use cosmic::Element; @@ -57,6 +58,7 @@ pub struct ProviderLoginStates<'a> { pub antigravity: Option<&'a AntigravityLoginState>, pub opencode_go: Option<&'a crate::providers::opencode_go::login::OpenCodeGoLoginState>, pub grok: Option<&'a crate::providers::grok::GrokLoginState>, + pub zai: Option<&'a ZaiLoginState>, } #[derive(Clone, Copy)] diff --git a/src/app/popup_view/detail.rs b/src/app/popup_view/detail.rs index 50367ee..34ae379 100644 --- a/src/app/popup_view/detail.rs +++ b/src/app/popup_view/detail.rs @@ -477,6 +477,14 @@ fn window_sections<'a>( config: &'a Config, ) -> Vec> { let mut items = Vec::new(); + if zai_coding_plan_absent(snapshot.provider, &snapshot.windows) { + items.push(info_block( + fl!("zai-coding-plan-unavailable-title"), + fl!("zai-coding-plan-unavailable-detail"), + None, + None, + )); + } let mut windows = snapshot.windows.iter().peekable(); while let Some(window) = windows.next() { let Some(group) = window.group.as_deref() else { @@ -496,6 +504,12 @@ fn window_sections<'a>( items } +fn zai_coding_plan_absent(provider: ProviderId, windows: &[UsageWindow]) -> bool { + provider == ProviderId::Zai + && !windows.is_empty() + && windows.iter().all(|window| window.label == "MCP") +} + fn usage_group_card( group: &str, sections: Vec>, @@ -855,6 +869,34 @@ mod tests { use super::*; use crate::model::{AccountSelectionStatus, AuthState, ProviderHealth}; + #[test] + fn zai_coding_plan_notice_only_for_mcp_only_usage() { + let mut windows = vec![UsageWindow { + label: "MCP".to_string(), + used_percent: 35.0, + reset_at: None, + window_seconds: None, + reset_description: None, + group: None, + }]; + assert!(zai_coding_plan_absent(ProviderId::Zai, &windows)); + assert_eq!( + fl!("zai-coding-plan-unavailable-detail"), + "Coding Plan quotas were not reported by Z.AI." + ); + assert!(!zai_coding_plan_absent(ProviderId::Minimax, &windows)); + assert!(!zai_coding_plan_absent(ProviderId::Zai, &[])); + + let mut coding_plan = windows[0].clone(); + coding_plan.label = "Weekly".to_string(); + windows.push(coding_plan); + assert!(!zai_coding_plan_absent(ProviderId::Zai, &windows)); + windows[1].label = "5 Hour".to_string(); + assert!(!zai_coding_plan_absent(ProviderId::Zai, &windows)); + windows.remove(0); + assert!(!zai_coding_plan_absent(ProviderId::Zai, &windows)); + } + #[test] fn action_required_account_reports_reauth_message() { let provider = ProviderRuntimeState { diff --git a/src/app/popup_view/settings/accounts.rs b/src/app/popup_view/settings/accounts.rs index c4e4eb3..e0d0402 100644 --- a/src/app/popup_view/settings/accounts.rs +++ b/src/app/popup_view/settings/accounts.rs @@ -6,7 +6,7 @@ use self::empty::empty_accounts_state; use self::login_controls::{ antigravity_login_controls, claude_login_controls, codex_login_controls, copilot_login_controls, cursor_scan_controls, gemini_login_controls, grok_login_controls, - kimi_login_controls, minimax_login_controls, opencode_go_login_controls, + kimi_login_controls, minimax_login_controls, opencode_go_login_controls, zai_login_controls, }; use self::rows::{ AccountRowPosition, account_action_container, account_selector_list, account_settings_row, @@ -109,6 +109,10 @@ pub(super) fn provider_settings_view<'a>( minimax_login_controls(logins.minimax, enabled), logins.minimax.is_some(), ), + ProviderLoginKind::Zai => ( + zai_login_controls(logins.zai, enabled), + logins.zai.is_some(), + ), ProviderLoginKind::Kimi => ( kimi_login_controls(logins.kimi, enabled), logins.kimi.is_some(), diff --git a/src/app/popup_view/settings/accounts/login_controls.rs b/src/app/popup_view/settings/accounts/login_controls.rs index ab4e279..89249ec 100644 --- a/src/app/popup_view/settings/accounts/login_controls.rs +++ b/src/app/popup_view/settings/accounts/login_controls.rs @@ -4,18 +4,25 @@ use super::super::super::{ Element, GeminiLoginState, GeminiLoginStatus, Length, Message, account_add_button, account_import_button, fl, row, widget, }; -use crate::app::login::{KimiLoginFlow, LoginFlow, MinimaxLoginFlow, OpenCodeGoLoginFlow}; +use crate::app::login::{ + KimiLoginFlow, LoginFlow, MinimaxLoginFlow, OpenCodeGoLoginFlow, ZaiLoginFlow, +}; use crate::providers::grok::{GrokLoginState, GrokLoginStatus}; use crate::providers::kimi::login::{KimiLoginEvent, KimiLoginState, KimiLoginStatus}; use crate::providers::minimax::{MinimaxLoginEvent, MinimaxLoginState, MinimaxLoginStatus}; use crate::providers::opencode_go::login::{ OpenCodeGoLoginEvent, OpenCodeGoLoginState, OpenCodeGoLoginStatus, }; +use crate::providers::zai::{ZaiLoginEvent, ZaiLoginState, ZaiLoginStatus}; fn minimax_login_message(event: MinimaxLoginEvent) -> Message { MinimaxLoginFlow::wrap_event(event) } +fn zai_login_message(event: ZaiLoginEvent) -> Message { + ZaiLoginFlow::wrap_event(event) +} + fn kimi_login_message(event: KimiLoginEvent) -> Message { KimiLoginFlow::wrap_event(event) } @@ -563,6 +570,81 @@ fn minimax_login_status(login: &MinimaxLoginState) -> String { } } +pub(super) fn zai_login_controls( + login: Option<&ZaiLoginState>, + enabled: bool, +) -> Element<'_, Message> { + let Some(login) = login else { + return account_add_button( + fl!("account-add"), + enabled.then_some(Message::StartLogin(crate::model::ProviderId::Zai)), + ); + }; + let status = match login.status { + ZaiLoginStatus::Editing => fl!("zai-login-editing"), + ZaiLoginStatus::Failed => fl!("zai-login-failed"), + }; + let mut content = cosmic::iced::widget::column![widget::text(status).size(13)] + .spacing(10) + .width(Length::Fill); + if let Some(error) = &login.error { + content = content.push(widget::text(error).size(13)); + } + if login.status == ZaiLoginStatus::Editing { + content = content.push(zai_login_fields(login)).push( + row![ + widget::button::standard(fl!("account-add")) + .on_press_maybe(enabled.then_some(zai_login_message(ZaiLoginEvent::Saved))), + widget::button::text(fl!("account-cancel")).on_press_maybe( + enabled.then_some(Message::CancelLogin(crate::model::ProviderId::Zai)) + ), + ] + .spacing(8), + ); + } else { + content = content.push( + row![ + widget::button::text(fl!("account-add-another")).on_press_maybe( + enabled.then_some(Message::StartLogin(crate::model::ProviderId::Zai)) + ), + widget::button::text(fl!("account-dismiss")).on_press_maybe( + enabled.then_some(Message::CancelLogin(crate::model::ProviderId::Zai)) + ), + ] + .spacing(8), + ); + } + content.into() +} + +fn zai_login_fields(login: &ZaiLoginState) -> Element<'_, Message> { + let mut fields = cosmic::iced::widget::column![ + widget::text(fl!("zai-api-key-placeholder")).size(12), + widget::text_input::secure_input( + fl!("zai-api-key-placeholder"), + &login.api_key, + Some(zai_login_message(ZaiLoginEvent::ApiKeyVisibilityToggled)), + !login.api_key_visible, + ) + .on_input(|api_key| zai_login_message(ZaiLoginEvent::ApiKeyChanged(api_key))) + .on_submit(|_| zai_login_message(ZaiLoginEvent::Saved)) + .width(Length::Fill), + ] + .spacing(10) + .width(Length::Fill); + if login.api_key_from_opencode { + fields = fields.push(widget::text(fl!("zai-api-key-imported-from-opencode")).size(12)); + } + fields + .push(widget::text(fl!("account-label")).size(12)) + .push( + widget::text_input(fl!("account-label"), &login.label) + .on_input(|label| zai_login_message(ZaiLoginEvent::LabelChanged(label))) + .width(Length::Fill), + ) + .into() +} + pub(super) fn kimi_login_controls( login: Option<&KimiLoginState>, enabled: bool, diff --git a/src/app/provider_actions.rs b/src/app/provider_actions.rs index 15a3391..d9d5cfd 100644 --- a/src/app/provider_actions.rs +++ b/src/app/provider_actions.rs @@ -264,15 +264,15 @@ impl AppModel { ))) } - pub(super) fn write_config(&mut self, f: impl FnOnce(&mut Config)) { + pub(super) fn write_config(&mut self, f: impl FnOnce(&mut Config)) -> bool { let mut new_config = self.config.clone(); f(&mut new_config); if new_config == self.config { - return; + return true; } if demo_env::is_active() { self.config = new_config; - return; + return true; } let ctx = match crate::config::cosmic_config_context( ::APP_ID, @@ -287,7 +287,7 @@ impl AppModel { error = ?error, "failed to open config for writing" ); - return; + return false; } }; if let Err(error) = @@ -300,9 +300,10 @@ impl AppModel { error = ?error, "failed to write config" ); - return; + return false; } self.config = new_config; + true } pub(super) fn set_provider_enabled( @@ -574,6 +575,7 @@ pub(super) fn popup_route_label(route: PopupRoute) -> &'static str { ProviderId::Gemini => "manage_accounts_gemini", ProviderId::Copilot => "manage_accounts_copilot", ProviderId::Minimax => "manage_accounts_minimax", + ProviderId::Zai => "manage_accounts_zai", ProviderId::Kimi => "manage_accounts_kimi", ProviderId::Antigravity => "manage_accounts_antigravity", ProviderId::OpenCodeGo => "manage_accounts_opencode_go", @@ -616,6 +618,7 @@ fn managed_account_count(config: &Config) -> usize { + config.gemini_managed_accounts.len() + config.copilot_managed_accounts.len() + config.minimax_managed_accounts.len() + + config.zai_managed_accounts.len() + config.kimi_managed_accounts.len() + config.antigravity_managed_accounts.len() + config.opencode_go_managed_accounts.len() diff --git a/src/app/provider_assets.rs b/src/app/provider_assets.rs index 38f47b2..d3e0a68 100644 --- a/src/app/provider_assets.rs +++ b/src/app/provider_assets.rs @@ -41,6 +41,12 @@ pub fn provider_icon_handle(provider: ProviderId, variant: ProviderIconVariant) (ProviderId::Minimax, ProviderIconVariant::Reversed) => { include_bytes!("../../resources/providers/minimax-reversed.svg") } + (ProviderId::Zai, ProviderIconVariant::Default) => { + include_bytes!("../../resources/providers/zai.svg") + } + (ProviderId::Zai, ProviderIconVariant::Reversed) => { + include_bytes!("../../resources/providers/zai-reversed.svg") + } (ProviderId::Kimi, ProviderIconVariant::Default) => { include_bytes!("../../resources/providers/kimi.svg") } diff --git a/src/app/session.rs b/src/app/session.rs index 9ef8221..44cbf04 100644 --- a/src/app/session.rs +++ b/src/app/session.rs @@ -41,6 +41,7 @@ pub(super) fn start_login(app: &mut AppModel, provider: ProviderId) -> Task login::start_login::(app), ProviderId::Copilot => login::start_login::(app), ProviderId::Minimax => login::start_login::(app), + ProviderId::Zai => login::start_login::(app), ProviderId::Kimi => login::start_login::(app), ProviderId::Antigravity => login::start_login::(app), ProviderId::OpenCodeGo => login::start_login::(app), @@ -93,6 +94,7 @@ pub(super) fn cancel_login(app: &mut AppModel, provider: ProviderId) { ProviderId::Gemini => login::cancel_login::(app), ProviderId::Copilot => login::cancel_login::(app), ProviderId::Minimax => login::cancel_login::(app), + ProviderId::Zai => login::cancel_login::(app), ProviderId::Kimi => login::cancel_login::(app), ProviderId::Antigravity => login::cancel_login::(app), ProviderId::OpenCodeGo => login::cancel_login::(app), @@ -112,6 +114,7 @@ pub(super) fn reauthenticate( ProviderId::Gemini => login::reauthenticate::(app, account_id), ProviderId::Copilot => login::reauthenticate::(app, account_id), ProviderId::Minimax => login::reauthenticate::(app, account_id), + ProviderId::Zai => login::reauthenticate::(app, account_id), ProviderId::Kimi => login::reauthenticate::(app, account_id), ProviderId::Antigravity => { login::reauthenticate::(app, account_id) @@ -141,6 +144,7 @@ pub(super) fn sync_metadata_after_refresh(app: &mut AppModel, provider: Provider ProviderId::Gemini | ProviderId::Copilot | ProviderId::Minimax + | ProviderId::Zai | ProviderId::Kimi | ProviderId::Antigravity | ProviderId::OpenCodeGo diff --git a/src/app/tests.rs b/src/app/tests.rs index 71ea35f..c2a5dc3 100644 --- a/src/app/tests.rs +++ b/src/app/tests.rs @@ -18,7 +18,7 @@ use crate::account_storage::{NewProviderAccount, ProviderAccountStorage, Provide use crate::config::{ ManagedClaudeAccountConfig, ManagedCodexAccountConfig, ManagedCopilotAccountConfig, ManagedCursorAccountConfig, ManagedGeminiAccountConfig, ManagedKimiAccountConfig, - ManagedMinimaxAccountConfig, + ManagedMinimaxAccountConfig, ManagedZaiAccountConfig, }; use crate::model::{ AccountSelectionStatus, ProviderAccountRuntimeState, ProviderIdentity, ProviderRuntimeState, @@ -1150,6 +1150,8 @@ pub(super) fn test_app(refresh_owner: Option) -> AppModel { opencode_go_login_handle: None, grok_login: None, grok_login_handle: None, + zai_login: None, + zai_login_handle: None, } } @@ -1288,6 +1290,17 @@ fn kimi_account(id: &str) -> ManagedKimiAccountConfig { } } +fn zai_account(id: &str) -> ManagedZaiAccountConfig { + ManagedZaiAccountConfig { + id: id.to_string(), + label: id.to_string(), + api_key_source: "env:ZAI_API_KEY".to_string(), + created_at: Utc::now(), + updated_at: Utc::now(), + last_authenticated_at: None, + } +} + fn antigravity_account(id: &str) -> crate::config::ManagedAntigravityAccountConfig { crate::config::ManagedAntigravityAccountConfig { id: id.to_string(), @@ -1453,6 +1466,12 @@ fn delete_account_requests_refresh_for_all_providers() { app.config.selected_minimax_account_ids = vec![keep_id.to_string()]; "remove".to_string() } + ProviderId::Zai => { + app.config.zai_managed_accounts.push(zai_account(keep_id)); + app.config.zai_managed_accounts.push(zai_account("remove")); + app.config.selected_zai_account_ids = vec![keep_id.to_string()]; + "remove".to_string() + } ProviderId::Kimi => { app.config.kimi_managed_accounts.push(kimi_account(keep_id)); app.config diff --git a/src/config.rs b/src/config.rs index fe50b5c..7c77e5b 100644 --- a/src/config.rs +++ b/src/config.rs @@ -37,6 +37,8 @@ pub struct Config { #[serde(default)] pub minimax_enablement: ProviderEnablement, #[serde(default)] + pub zai_enablement: ProviderEnablement, + #[serde(default)] pub kimi_enablement: ProviderEnablement, #[serde(default)] pub antigravity_enablement: ProviderEnablement, @@ -63,6 +65,10 @@ pub struct Config { #[serde(default)] pub minimax_managed_accounts: Vec, #[serde(default)] + pub selected_zai_account_ids: Vec, + #[serde(default)] + pub zai_managed_accounts: Vec, + #[serde(default)] pub selected_kimi_account_ids: Vec, #[serde(default)] pub kimi_managed_accounts: Vec, @@ -96,6 +102,7 @@ impl Default for Config { gemini_enablement: default_gemini_enablement(), copilot_enablement: ProviderEnablement::Auto, minimax_enablement: ProviderEnablement::Auto, + zai_enablement: ProviderEnablement::Auto, kimi_enablement: ProviderEnablement::Auto, antigravity_enablement: ProviderEnablement::Auto, opencode_go_enablement: ProviderEnablement::Auto, @@ -112,6 +119,8 @@ impl Default for Config { copilot_managed_accounts: Vec::new(), selected_minimax_account_ids: Vec::new(), minimax_managed_accounts: Vec::new(), + selected_zai_account_ids: Vec::new(), + zai_managed_accounts: Vec::new(), selected_kimi_account_ids: Vec::new(), kimi_managed_accounts: Vec::new(), selected_antigravity_account_ids: Vec::new(), @@ -147,6 +156,7 @@ impl Config { ProviderId::Gemini => self.gemini_enablement, ProviderId::Copilot => self.copilot_enablement, ProviderId::Minimax => self.minimax_enablement, + ProviderId::Zai => self.zai_enablement, ProviderId::Kimi => self.kimi_enablement, ProviderId::Antigravity => self.antigravity_enablement, ProviderId::OpenCodeGo => self.opencode_go_enablement, @@ -163,6 +173,7 @@ impl Config { ProviderId::Gemini => &self.selected_gemini_account_ids, ProviderId::Copilot => &self.selected_copilot_account_ids, ProviderId::Minimax => &self.selected_minimax_account_ids, + ProviderId::Zai => &self.selected_zai_account_ids, ProviderId::Kimi => &self.selected_kimi_account_ids, ProviderId::Antigravity => &self.selected_antigravity_account_ids, ProviderId::OpenCodeGo => &self.selected_opencode_go_account_ids, @@ -178,6 +189,7 @@ impl Config { ProviderId::Gemini => &mut self.selected_gemini_account_ids, ProviderId::Copilot => &mut self.selected_copilot_account_ids, ProviderId::Minimax => &mut self.selected_minimax_account_ids, + ProviderId::Zai => &mut self.selected_zai_account_ids, ProviderId::Kimi => &mut self.selected_kimi_account_ids, ProviderId::Antigravity => &mut self.selected_antigravity_account_ids, ProviderId::OpenCodeGo => &mut self.selected_opencode_go_account_ids, @@ -232,6 +244,7 @@ fn provider_enabled_key(provider: ProviderId) -> &'static str { ProviderId::Gemini => "gemini_enabled", ProviderId::Copilot => "copilot_enabled", ProviderId::Minimax => "minimax_enabled", + ProviderId::Zai => "zai_enabled", ProviderId::Kimi => "kimi_enabled", ProviderId::Antigravity => "antigravity_enabled", ProviderId::OpenCodeGo => "opencode_go_enabled", @@ -247,6 +260,7 @@ fn provider_enablement_key(provider: ProviderId) -> &'static str { ProviderId::Gemini => "gemini_enablement", ProviderId::Copilot => "copilot_enablement", ProviderId::Minimax => "minimax_enablement", + ProviderId::Zai => "zai_enablement", ProviderId::Kimi => "kimi_enablement", ProviderId::Antigravity => "antigravity_enablement", ProviderId::OpenCodeGo => "opencode_go_enablement", @@ -262,6 +276,7 @@ fn provider_enablement_mut(config: &mut Config, provider: ProviderId) -> &mut Pr ProviderId::Gemini => &mut config.gemini_enablement, ProviderId::Copilot => &mut config.copilot_enablement, ProviderId::Minimax => &mut config.minimax_enablement, + ProviderId::Zai => &mut config.zai_enablement, ProviderId::Kimi => &mut config.kimi_enablement, ProviderId::Antigravity => &mut config.antigravity_enablement, ProviderId::OpenCodeGo => &mut config.opencode_go_enablement, @@ -406,6 +421,16 @@ pub struct ManagedMinimaxAccountConfig { pub last_authenticated_at: Option>, } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct ManagedZaiAccountConfig { + pub id: String, + pub label: String, + pub api_key_source: String, + pub created_at: DateTime, + pub updated_at: DateTime, + pub last_authenticated_at: Option>, +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct ManagedKimiAccountConfig { pub id: String, @@ -466,6 +491,7 @@ pub struct AppPaths { pub gemini_accounts_dir: PathBuf, pub copilot_accounts_dir: PathBuf, pub minimax_accounts_dir: PathBuf, + pub zai_accounts_dir: PathBuf, pub kimi_accounts_dir: PathBuf, pub antigravity_accounts_dir: PathBuf, pub opencode_go_accounts_dir: PathBuf, @@ -509,6 +535,7 @@ pub fn write_changed_config_entries( gemini_enablement, copilot_enablement, minimax_enablement, + zai_enablement, kimi_enablement, antigravity_enablement, opencode_go_enablement, @@ -524,6 +551,8 @@ pub fn write_changed_config_entries( copilot_managed_accounts, selected_minimax_account_ids, minimax_managed_accounts, + selected_zai_account_ids, + zai_managed_accounts, selected_kimi_account_ids, kimi_managed_accounts, selected_antigravity_account_ids, @@ -556,6 +585,7 @@ pub fn write_changed_config_entries( set_changed!(gemini_enablement); set_changed!(copilot_enablement); set_changed!(minimax_enablement); + set_changed!(zai_enablement); set_changed!(kimi_enablement); set_changed!(antigravity_enablement); set_changed!(opencode_go_enablement); @@ -572,6 +602,8 @@ pub fn write_changed_config_entries( set_changed!(copilot_managed_accounts); set_changed!(selected_minimax_account_ids); set_changed!(minimax_managed_accounts); + set_changed!(selected_zai_account_ids); + set_changed!(zai_managed_accounts); set_changed!(selected_kimi_account_ids); set_changed!(kimi_managed_accounts); set_changed!(selected_antigravity_account_ids); @@ -712,6 +744,7 @@ pub fn paths() -> AppPaths { let gemini_accounts_dir = state_dir.join("gemini-accounts"); let copilot_accounts_dir = state_dir.join("copilot-accounts"); let minimax_accounts_dir = state_dir.join("minimax-accounts"); + let zai_accounts_dir = state_dir.join("zai-accounts"); let kimi_accounts_dir = state_dir.join("kimi-accounts"); let antigravity_accounts_dir = state_dir.join("antigravity-accounts"); let opencode_go_accounts_dir = state_dir.join("opencode-go-accounts"); @@ -726,6 +759,7 @@ pub fn paths() -> AppPaths { gemini_accounts_dir, copilot_accounts_dir, minimax_accounts_dir, + zai_accounts_dir, kimi_accounts_dir, antigravity_accounts_dir, opencode_go_accounts_dir, @@ -809,6 +843,17 @@ mod tests { new.refresh_interval_seconds = old.refresh_interval_seconds + 60; new.kimi_enablement = ProviderEnablement::Enabled; new.selected_kimi_account_ids = vec!["kimi-test".to_string()]; + new.zai_enablement = ProviderEnablement::Enabled; + new.selected_zai_account_ids = vec!["zai-test".to_string()]; + let now = Utc::now(); + new.zai_managed_accounts = vec![ManagedZaiAccountConfig { + id: "zai-test".to_string(), + label: "Z.AI Coding Plan".to_string(), + api_key_source: "stored".to_string(), + created_at: now, + updated_at: now, + last_authenticated_at: None, + }]; write_changed_config_entries(&ctx, &old, &new).unwrap(); @@ -836,6 +881,8 @@ mod tests { assert_eq!(config.gemini_enablement, ProviderEnablement::Enabled); assert!(config.set_provider_enabled(ProviderId::Gemini, false)); assert_eq!(config.gemini_enablement, ProviderEnablement::Disabled); + assert!(config.set_provider_enabled(ProviderId::Zai, true)); + assert_eq!(config.zai_enablement, ProviderEnablement::Enabled); } #[test] @@ -852,6 +899,9 @@ mod tests { assert_eq!(config.selected_provider, ProviderId::Codex); assert_eq!(config.codex_enablement, ProviderEnablement::Auto); assert_eq!(config.gemini_enablement, ProviderEnablement::Disabled); + assert_eq!(config.zai_enablement, ProviderEnablement::Auto); + assert!(config.selected_zai_account_ids.is_empty()); + assert!(config.zai_managed_accounts.is_empty()); } #[test] @@ -864,6 +914,21 @@ mod tests { assert_eq!(config.gemini_enablement, ProviderEnablement::Disabled); } + #[test] + fn missing_zai_fields_default_for_existing_config() { + let mut value = serde_json::to_value(Config::default()).unwrap(); + let object = value.as_object_mut().unwrap(); + object.remove("zai_enablement"); + object.remove("selected_zai_account_ids"); + object.remove("zai_managed_accounts"); + + let config: Config = serde_json::from_value(value).unwrap(); + + assert_eq!(config.zai_enablement, ProviderEnablement::Auto); + assert!(config.selected_zai_account_ids.is_empty()); + assert!(config.zai_managed_accounts.is_empty()); + } + #[test] fn config_schema_version_marks_fresh_patch_boundary() { let config = Config::default(); @@ -874,6 +939,7 @@ mod tests { assert!(config.gemini_managed_accounts.is_empty()); assert!(config.copilot_managed_accounts.is_empty()); assert!(config.minimax_managed_accounts.is_empty()); + assert!(config.zai_managed_accounts.is_empty()); assert!(config.kimi_managed_accounts.is_empty()); assert!(config.antigravity_managed_accounts.is_empty()); assert!(config.grok_managed_accounts.is_empty()); @@ -1085,6 +1151,37 @@ mod tests { ); } + #[test] + fn zai_accounts_dir_is_configured_under_state_root() { + let p = paths(); + assert!( + p.zai_accounts_dir + .ends_with(std::path::Path::new("yapcap/zai-accounts")), + "unexpected zai_accounts_dir: {}", + p.zai_accounts_dir.display() + ); + } + + #[test] + fn zai_managed_account_config_roundtrips_without_api_key() { + let now = Utc::now(); + let account = ManagedZaiAccountConfig { + id: "zai-test-1".to_string(), + label: "Z.AI Coding Plan".to_string(), + api_key_source: "stored".to_string(), + created_at: now, + updated_at: now, + last_authenticated_at: Some(now), + }; + let value = serde_json::to_value(&account).unwrap(); + + assert!(value.get("api_key").is_none()); + assert_eq!( + serde_json::from_value::(value).unwrap(), + account + ); + } + #[test] fn grok_managed_account_config_roundtrips() { let id = "grok-test-1"; diff --git a/src/config/watch_update.rs b/src/config/watch_update.rs index 6cbf826..cd01929 100644 --- a/src/config/watch_update.rs +++ b/src/config/watch_update.rs @@ -33,6 +33,7 @@ impl Config { "gemini_enablement" => self.gemini_enablement = update.gemini_enablement, "copilot_enablement" => self.copilot_enablement = update.copilot_enablement, "minimax_enablement" => self.minimax_enablement = update.minimax_enablement, + "zai_enablement" => self.zai_enablement = update.zai_enablement, "kimi_enablement" => self.kimi_enablement = update.kimi_enablement, "antigravity_enablement" => { self.antigravity_enablement = update.antigravity_enablement; @@ -87,6 +88,12 @@ impl Config { "minimax_managed_accounts" => { self.minimax_managed_accounts = update.minimax_managed_accounts.clone(); } + "selected_zai_account_ids" => { + self.selected_zai_account_ids = update.selected_zai_account_ids.clone(); + } + "zai_managed_accounts" => { + self.zai_managed_accounts = update.zai_managed_accounts.clone(); + } "selected_kimi_account_ids" => { self.selected_kimi_account_ids = update.selected_kimi_account_ids.clone(); } @@ -123,7 +130,7 @@ mod tests { use super::*; use crate::config::{ ManagedGrokAccountConfig, ManagedKimiAccountConfig, ManagedMinimaxAccountConfig, - ManagedOpenCodeGoAccountConfig, ProviderEnablement, + ManagedOpenCodeGoAccountConfig, ManagedZaiAccountConfig, ProviderEnablement, }; use chrono::Utc; @@ -193,6 +200,40 @@ mod tests { ); } + #[test] + fn applies_zai_watcher_keys_without_replacing_unrelated_configuration() { + let mut config = Config::default(); + let now = Utc::now(); + let update = Config { + codex_enablement: ProviderEnablement::Disabled, + zai_enablement: ProviderEnablement::Enabled, + selected_zai_account_ids: vec!["zai-1".to_string()], + zai_managed_accounts: vec![ManagedZaiAccountConfig { + id: "zai-1".to_string(), + label: "Z.AI Coding Plan".to_string(), + api_key_source: "stored".to_string(), + created_at: now, + updated_at: now, + last_authenticated_at: None, + }], + ..Config::default() + }; + + config.apply_watcher_update( + update, + &[ + "zai_enablement", + "selected_zai_account_ids", + "zai_managed_accounts", + ], + ); + + assert_eq!(config.zai_enablement, ProviderEnablement::Enabled); + assert_eq!(config.selected_zai_account_ids, ["zai-1"]); + assert_eq!(config.zai_managed_accounts[0].id, "zai-1"); + assert_eq!(config.codex_enablement, ProviderEnablement::Auto); + } + #[test] fn applies_opencode_go_watcher_keys_without_replacing_unrelated_configuration() { let mut config = Config::default(); diff --git a/src/demo_env.rs b/src/demo_env.rs index 9a9e170..2efeaf1 100644 --- a/src/demo_env.rs +++ b/src/demo_env.rs @@ -4,7 +4,8 @@ use crate::config::{ Config, ManagedAntigravityAccountConfig, ManagedClaudeAccountConfig, ManagedCodexAccountConfig, ManagedCopilotAccountConfig, ManagedCursorAccountConfig, ManagedGeminiAccountConfig, ManagedGrokAccountConfig, ManagedKimiAccountConfig, ManagedMinimaxAccountConfig, - ManagedOpenCodeGoAccountConfig, ProviderEnablement, ProviderVisibilityMode, paths, + ManagedOpenCodeGoAccountConfig, ManagedZaiAccountConfig, ProviderEnablement, + ProviderVisibilityMode, paths, }; use crate::model::{ AccountSelectionStatus, AppState, AuthState, ExtraUsageState, ProviderAccountRuntimeState, @@ -31,6 +32,7 @@ const ANTIGRAVITY_PRIMARY_ID: &str = "yapcap-demo:antigravity-primary"; const ANTIGRAVITY_FREE_ID: &str = "yapcap-demo:antigravity-free"; const OPENCODE_GO_ID: &str = "yapcap-demo:opencode-go"; const GROK_PRIMARY_ID: &str = "yapcap-demo:grok-primary"; +const ZAI_PRIMARY_ID: &str = "yapcap-demo:zai-coding-plan"; fn env_truthy() -> bool { std::env::var(DEMO_ENV).is_ok_and(|value| { @@ -68,6 +70,7 @@ pub fn apply_config(config: &mut Config) { config.antigravity_enablement = ProviderEnablement::Enabled; config.opencode_go_enablement = ProviderEnablement::Enabled; config.grok_enablement = ProviderEnablement::Enabled; + config.zai_enablement = ProviderEnablement::Enabled; config.codex_managed_accounts = demo_codex_accounts(); config.claude_managed_accounts = demo_claude_accounts(); @@ -79,6 +82,7 @@ pub fn apply_config(config: &mut Config) { config.antigravity_managed_accounts = demo_antigravity_accounts(); config.opencode_go_managed_accounts = demo_opencode_go_accounts(); config.grok_managed_accounts = demo_grok_accounts(); + config.zai_managed_accounts = demo_zai_accounts(); config.provider_visibility_mode = ProviderVisibilityMode::UserManaged; @@ -92,6 +96,7 @@ pub fn apply_config(config: &mut Config) { config.selected_antigravity_account_ids = vec![ANTIGRAVITY_PRIMARY_ID.to_string()]; config.selected_opencode_go_account_ids = vec![OPENCODE_GO_ID.to_string()]; config.selected_grok_account_ids = vec![GROK_PRIMARY_ID.to_string()]; + config.selected_zai_account_ids = vec![ZAI_PRIMARY_ID.to_string()]; } pub fn strip_leaked_state(config: &mut Config) -> bool { @@ -119,6 +124,8 @@ pub fn strip_leaked_state(config: &mut Config) -> bool { }); changed |= strip_ids(&mut config.selected_grok_account_ids); changed |= retain_len_changed(&mut config.grok_managed_accounts, |account| &account.id); + changed |= strip_ids(&mut config.selected_zai_account_ids); + changed |= retain_len_changed(&mut config.zai_managed_accounts, |account| &account.id); changed } @@ -188,6 +195,7 @@ fn demo_system_active_account_id(provider: ProviderId) -> Option { ProviderId::Antigravity => return None, ProviderId::OpenCodeGo => return None, ProviderId::Grok => return None, + ProviderId::Zai => return None, }; Some(id.to_string()) } @@ -204,6 +212,7 @@ fn demo_source(provider: ProviderId) -> String { ProviderId::Kimi => "API Key".to_string(), ProviderId::Antigravity => "OAuth".to_string(), ProviderId::OpenCodeGo => "API Key".to_string(), + ProviderId::Zai => "API Key".to_string(), } } @@ -386,6 +395,18 @@ fn demo_runtime_accounts(provider: ProviderId) -> Vec vec![demo_account( + provider, + DemoAccount { + account_id: ZAI_PRIMARY_ID, + label: "Z.AI Coding Plan", + last_success_at: now - Duration::minutes(2), + health: ProviderHealth::Ok, + auth_state: AuthState::Ready, + error: None, + snapshot: snapshot_zai_primary(), + }, + )], } } @@ -812,6 +833,18 @@ fn demo_grok_accounts() -> Vec { }] } +fn demo_zai_accounts() -> Vec { + let now = demo_timestamp(); + vec![ManagedZaiAccountConfig { + id: ZAI_PRIMARY_ID.to_string(), + label: "Z.AI Coding Plan".to_string(), + api_key_source: "demo".to_string(), + created_at: now, + updated_at: now, + last_authenticated_at: Some(now), + }] +} + fn snapshot_minimax_primary() -> UsageSnapshot { let now = Utc::now(); let interval_reset = now + Duration::hours(3); @@ -888,6 +921,53 @@ fn snapshot_kimi_primary() -> UsageSnapshot { } } +fn snapshot_zai_primary() -> UsageSnapshot { + let now = Utc::now(); + let five_hour_reset = now + Duration::hours(3); + let weekly_reset = now + Duration::days(4); + let mcp_reset = now + Duration::hours(1); + UsageSnapshot { + provider: ProviderId::Zai, + source: "API Key".to_string(), + updated_at: now, + headline: UsageHeadline(0), + windows: vec![ + UsageWindow { + label: "5 Hour".to_string(), + used_percent: 34.0, + reset_at: Some(five_hour_reset), + window_seconds: Some(5 * 60 * 60), + reset_description: Some(five_hour_reset.to_rfc3339()), + group: None, + }, + UsageWindow { + label: "Weekly".to_string(), + used_percent: 57.0, + reset_at: Some(weekly_reset), + window_seconds: Some(7 * 24 * 60 * 60), + reset_description: Some(weekly_reset.to_rfc3339()), + group: None, + }, + UsageWindow { + label: "MCP".to_string(), + used_percent: 18.0, + reset_at: Some(mcp_reset), + window_seconds: None, + reset_description: Some(mcp_reset.to_rfc3339()), + group: None, + }, + ], + provider_cost: None, + extra_usage: None, + identity: ProviderIdentity { + email: None, + account_id: None, + plan: Some("Coding Plan".to_string()), + display_name: Some("Z.AI Coding Plan".to_string()), + }, + } +} + fn snapshot_opencode_go() -> UsageSnapshot { let now = Utc::now(); let five_hour_reset = now + Duration::hours(3); @@ -1239,15 +1319,18 @@ mod tests { snapshot_minimax_primary(), snapshot_kimi_primary(), snapshot_opencode_go(), + snapshot_zai_primary(), ] { assert!(!snapshot.windows.is_empty()); for window in &snapshot.windows { - assert!( - crate::usage_display::pace(window, snapshot.updated_at).is_some(), - "{} {} demo window must support the standard pace meter", - snapshot.provider.label(), - window.label - ); + if window.window_seconds.is_some() { + assert!( + crate::usage_display::pace(window, snapshot.updated_at).is_some(), + "{} {} demo window must support the standard pace meter", + snapshot.provider.label(), + window.label + ); + } } } } @@ -1281,6 +1364,7 @@ mod tests { assert_eq!(config.kimi_managed_accounts.len(), 1); assert_eq!(config.antigravity_managed_accounts.len(), 2); assert_eq!(config.opencode_go_managed_accounts.len(), 1); + assert_eq!(config.zai_managed_accounts.len(), 1); assert_eq!(config.selected_codex_account_ids.len(), 1); assert_eq!(config.selected_claude_account_ids.len(), 1); assert_eq!(config.selected_cursor_account_ids.len(), 1); @@ -1290,6 +1374,7 @@ mod tests { assert_eq!(config.selected_kimi_account_ids.len(), 1); assert_eq!(config.selected_antigravity_account_ids.len(), 1); assert_eq!(config.selected_opencode_go_account_ids.len(), 1); + assert_eq!(config.selected_zai_account_ids.len(), 1); for provider in ProviderId::ALL { assert_eq!( config.provider_enablement(provider), @@ -1318,6 +1403,7 @@ mod tests { selected_kimi_account_ids: vec!["real-kimi".to_string()], selected_antigravity_account_ids: vec!["real-antigravity".to_string()], selected_opencode_go_account_ids: vec!["real-opencode-go".to_string()], + selected_zai_account_ids: vec!["real-zai".to_string()], ..Config::default() }; apply_config(&mut config); @@ -1375,6 +1461,7 @@ mod tests { .map(|a| &a.id) .collect(), ProviderId::Grok => config.grok_managed_accounts.iter().map(|a| &a.id).collect(), + ProviderId::Zai => config.zai_managed_accounts.iter().map(|a| &a.id).collect(), }; for id in selected { assert!( @@ -1687,6 +1774,64 @@ mod tests { assert_eq!(labels, vec!["5 Hour", "Weekly", "Monthly"]); } + #[test] + fn zai_demo_seeds_one_account_with_coding_plan_windows() { + let _guard = test_support::env_lock(); + unsafe { + std::env::set_var(DEMO_ENV, "1"); + } + let mut config = Config::default(); + apply_config(&mut config); + let mut state = AppState::empty(); + apply(&config, &mut state); + unsafe { + std::env::remove_var(DEMO_ENV); + } + + assert_eq!( + config.selected_zai_account_ids, + vec![ZAI_PRIMARY_ID.to_string()] + ); + let account = state + .provider_accounts + .iter() + .find(|account| { + account.provider == ProviderId::Zai && account.account_id == ZAI_PRIMARY_ID + }) + .expect("Z.AI demo account"); + assert_eq!(account.label, "Z.AI Coding Plan"); + assert_eq!(account.source_label.as_deref(), Some("API Key")); + let snapshot = account.snapshot.as_ref().expect("Z.AI demo snapshot"); + assert_eq!(snapshot.provider, ProviderId::Zai); + assert_eq!(snapshot.source, "API Key"); + assert_eq!(snapshot.identity.plan.as_deref(), Some("Coding Plan")); + let labels: Vec<&str> = snapshot + .windows + .iter() + .map(|window| window.label.as_str()) + .collect(); + assert_eq!(labels, vec!["5 Hour", "Weekly", "MCP"]); + assert_eq!(snapshot.windows[0].window_seconds, Some(5 * 60 * 60)); + assert_eq!(snapshot.windows[1].window_seconds, Some(7 * 24 * 60 * 60)); + assert_eq!(snapshot.windows[2].window_seconds, None); + } + + #[test] + fn strip_leaked_state_removes_zai_demo_ids() { + let mut config = Config { + selected_zai_account_ids: vec![ZAI_PRIMARY_ID.to_string(), "real-zai".to_string()], + zai_managed_accounts: demo_zai_accounts(), + ..Config::default() + }; + + assert!(strip_leaked_state(&mut config)); + assert_eq!( + config.selected_zai_account_ids, + vec!["real-zai".to_string()] + ); + assert!(config.zai_managed_accounts.is_empty()); + } + #[test] fn codex_demo_seeds_active_pro_and_free_accounts() { let accounts = demo_runtime_accounts(ProviderId::Codex); diff --git a/src/detection.rs b/src/detection.rs index a0e097d..b6a7255 100644 --- a/src/detection.rs +++ b/src/detection.rs @@ -41,6 +41,7 @@ fn markers(provider: ProviderId) -> &'static [Marker] { const KIMI: [Marker; 0] = []; const OPENCODE_GO: [Marker; 1] = [file(".local/share/opencode/auth.json")]; const GROK: [Marker; 2] = [dir(".grok"), file(".grok/auth.json")]; + const ZAI: [Marker; 0] = []; match provider { ProviderId::Codex => &CODEX, ProviderId::Claude => &CLAUDE, @@ -52,6 +53,7 @@ fn markers(provider: ProviderId) -> &'static [Marker] { ProviderId::Kimi => &KIMI, ProviderId::OpenCodeGo => &OPENCODE_GO, ProviderId::Grok => &GROK, + ProviderId::Zai => &ZAI, } } @@ -96,13 +98,28 @@ fn provider_index(provider: ProviderId) -> usize { pub fn detect(home: &Path) -> DetectionSnapshot { let mut snapshot = DetectionSnapshot::default(); for provider in ProviderId::ALL { - snapshot.detected[provider_index(provider)] = markers(provider) - .iter() - .any(|marker| marker.exists_in(home)); + snapshot.detected[provider_index(provider)] = if provider == ProviderId::Zai { + detect_zai(home) + } else { + markers(provider) + .iter() + .any(|marker| marker.exists_in(home)) + }; } snapshot } +fn detect_zai(home: &Path) -> bool { + let default_path = || home.join(".local/share/opencode/auth.json"); + let path = + if std::env::var_os(crate::providers::opencode_auth::OPENCODE_AUTH_PATH_ENV).is_some() { + crate::providers::opencode_auth::auth_path().unwrap_or_else(default_path) + } else { + default_path() + }; + crate::providers::zai::opencode::has_usable_api_key_at(&path) +} + #[must_use] pub fn startup_snapshot(home: Option) -> DetectionSnapshot { home.map_or_else(DetectionSnapshot::default, |home| detect(&home)) @@ -315,4 +332,83 @@ mod tests { let snapshot = detect(home.path()); assert!(snapshot.detected(ProviderId::Grok)); } + + fn detect_zai_source(contents: Option<&str>) -> bool { + let home = home(); + let path = home.path().join(".local/share/opencode/auth.json"); + if let Some(contents) = contents { + fs::create_dir_all(path.parent().expect("auth path has a parent")).unwrap(); + fs::write(&path, contents).unwrap(); + } + let mut env = crate::test_support::test_env(); + env.remove(crate::providers::opencode_auth::OPENCODE_AUTH_CONTENT_ENV); + env.set( + crate::providers::opencode_auth::OPENCODE_AUTH_PATH_ENV, + &path, + ); + detect(home.path()).detected(ProviderId::Zai) + } + + #[test] + fn zai_detects_primary_coding_plan_credential() { + assert!(detect_zai_source(Some( + r#"{"zai-coding-plan":{"type":"api","key":"primary"},"zai":{"type":"api","key":"alias"}}"# + ))); + } + + #[test] + fn zai_detection_uses_alias_when_primary_is_not_usable() { + assert!(detect_zai_source(Some( + r#"{"zai-coding-plan":{"type":"oauth","refresh":"r","access":"a","expires":1},"zai":{"type":"api","key":"alias"}}"# + ))); + } + + #[test] + fn zai_detection_requires_an_ordered_provider_entry() { + assert!(!detect_zai_source(Some( + r#"{"other":{"type":"api","key":"unrelated"}}"# + ))); + } + + #[test] + fn zai_detection_rejects_non_api_blank_malformed_and_missing_sources() { + for source in [ + Some(r#"{"zai":{"type":"oauth","refresh":"r","access":"a","expires":1}}"#), + Some(r#"{"zai":{"type":"wellknown","key":"key","token":"token"}}"#), + Some(r#"{"zai":{"type":"future","value":"unknown"}}"#), + Some(r#"{"zai":{"type":"api","key":" "}}"#), + Some("not-json"), + None, + ] { + assert!(!detect_zai_source(source)); + } + } + + #[test] + fn zai_detection_does_not_use_a_bare_auth_file_marker() { + assert!(!detect_zai_source(Some("{}"))); + } + + #[test] + fn detection_snapshot_stores_only_the_zai_detection_fact() { + let secret = "key-that-must-not-enter-the-snapshot"; + let home = home(); + let path = home.path().join(".local/share/opencode/auth.json"); + fs::create_dir_all(path.parent().expect("auth path has a parent")).unwrap(); + fs::write( + &path, + format!(r#"{{"zai":{{"type":"api","key":"{secret}"}}}}"#), + ) + .unwrap(); + let mut env = crate::test_support::test_env(); + env.remove(crate::providers::opencode_auth::OPENCODE_AUTH_CONTENT_ENV); + env.set( + crate::providers::opencode_auth::OPENCODE_AUTH_PATH_ENV, + &path, + ); + + let snapshot = detect(home.path()); + assert!(snapshot.detected(ProviderId::Zai)); + assert!(!format!("{snapshot:?}").contains(secret)); + } } diff --git a/src/error.rs b/src/error.rs index ab2f5ca..62c3302 100644 --- a/src/error.rs +++ b/src/error.rs @@ -55,6 +55,12 @@ impl From for AppError { } } +impl From for AppError { + fn from(value: ZaiError) -> Self { + Self::Provider(ProviderError::Zai(value)) + } +} + impl From for AppError { fn from(value: KimiError) -> Self { Self::Provider(ProviderError::Kimi(value)) @@ -122,6 +128,7 @@ impl AppError { Self::Provider(ProviderError::Gemini(e)) => e.rate_limit_retry_after_secs(), Self::Provider(ProviderError::Copilot(e)) => e.rate_limit_retry_after_secs(), Self::Provider(ProviderError::Minimax(e)) => e.rate_limit_retry_after_secs(), + Self::Provider(ProviderError::Zai(e)) => e.rate_limit_retry_after_secs(), Self::Provider(ProviderError::Kimi(e)) => e.rate_limit_retry_after_secs(), Self::Provider(ProviderError::Antigravity(e)) => e.rate_limit_retry_after_secs(), Self::Provider(ProviderError::OpenCodeGo(e)) => e.rate_limit_retry_after_secs(), @@ -158,6 +165,8 @@ pub enum ProviderError { #[error(transparent)] Minimax(#[from] MinimaxError), #[error(transparent)] + Zai(#[from] ZaiError), + #[error(transparent)] Kimi(#[from] KimiError), #[error(transparent)] Antigravity(#[from] AntigravityError), @@ -177,6 +186,7 @@ impl ProviderError { Self::Gemini(error) => error.is_network_unavailable(), Self::Copilot(error) => error.is_network_unavailable(), Self::Minimax(error) => error.is_network_unavailable(), + Self::Zai(error) => error.is_network_unavailable(), Self::Kimi(error) => error.is_network_unavailable(), Self::Antigravity(error) => error.is_network_unavailable(), Self::OpenCodeGo(error) => error.is_network_unavailable(), @@ -193,6 +203,7 @@ impl ProviderError { Self::Gemini(error) => error.requires_user_action(), Self::Copilot(error) => error.requires_user_action(), Self::Minimax(error) => error.requires_user_action(), + Self::Zai(error) => error.requires_user_action(), Self::Kimi(error) => error.requires_user_action(), Self::Antigravity(error) => error.requires_user_action(), Self::OpenCodeGo(error) => error.requires_user_action(), @@ -209,6 +220,7 @@ impl ProviderError { Self::Gemini(error) => error.is_transient(), Self::Copilot(error) => error.is_transient(), Self::Minimax(error) => error.is_transient(), + Self::Zai(error) => error.is_transient(), Self::Kimi(error) => error.is_transient(), Self::Antigravity(error) => error.is_transient(), Self::OpenCodeGo(error) => error.is_transient(), @@ -646,6 +658,56 @@ impl MinimaxError { } } +#[derive(Debug, Error)] +pub enum ZaiError { + #[error("Z.AI login required")] + LoginRequired, + #[error("Z.AI API key is invalid")] + InvalidApiKey, + #[error("Z.AI usage request failed")] + UsageRequest(#[source] reqwest::Error), + #[error("Z.AI usage endpoint returned HTTP {status}")] + UsageHttp { status: u16 }, + #[error("failed to decode Z.AI usage response")] + DecodeUsage(#[source] serde_json::Error), + #[error("Z.AI usage response envelope is invalid")] + InvalidEnvelope, + #[error("Z.AI response had no valid usage windows")] + NoUsageData, + #[error("Rate limited by Z.AI — will retry automatically")] + RateLimited { retry_after_secs: Option }, +} + +impl ZaiError { + #[must_use] + pub fn is_network_unavailable(&self) -> bool { + matches!(self, Self::UsageRequest(source) if request_could_not_reach_network(source)) + } + + #[must_use] + pub fn requires_user_action(&self) -> bool { + matches!(self, Self::LoginRequired | Self::InvalidApiKey) + } + + #[must_use] + pub fn rate_limit_retry_after_secs(&self) -> Option { + match self { + Self::RateLimited { retry_after_secs } => *retry_after_secs, + _ => None, + } + } + + #[must_use] + pub fn is_transient(&self) -> bool { + match self { + Self::RateLimited { .. } => true, + Self::UsageRequest(source) => request_could_not_reach_network(source), + Self::UsageHttp { status } => *status >= 500, + _ => false, + } + } +} + #[derive(Debug, Error)] pub enum KimiError { #[error("Kimi login required")] @@ -1085,4 +1147,25 @@ mod tests { assert!(AppError::from(GrokError::CredentialsMissing).requires_user_action()); assert!(AppError::from(GrokError::RefreshUnavailable).requires_user_action()); } + + #[test] + fn zai_auth_failures_require_user_action() { + assert!(AppError::from(ZaiError::LoginRequired).requires_user_action()); + assert!(AppError::from(ZaiError::InvalidApiKey).requires_user_action()); + assert!(!AppError::from(ZaiError::LoginRequired).is_transient()); + } + + #[test] + fn zai_rate_limit_and_server_errors_are_classified() { + let rate_limited = AppError::from(ZaiError::RateLimited { + retry_after_secs: Some(42), + }); + assert_eq!(rate_limited.rate_limit_retry_after_secs(), Some(42)); + assert!(rate_limited.is_transient()); + assert!(!rate_limited.requires_user_action()); + + let server_error = AppError::from(ZaiError::UsageHttp { status: 503 }); + assert!(server_error.is_transient()); + assert!(!server_error.requires_user_action()); + } } diff --git a/src/model.rs b/src/model.rs index 7018274..9f8630f 100644 --- a/src/model.rs +++ b/src/model.rs @@ -16,13 +16,14 @@ pub enum ProviderId { Gemini, Copilot, Minimax, + Zai, Kimi, OpenCodeGo, Grok, } impl ProviderId { - pub const ALL: [Self; 10] = [ + pub const ALL: [Self; 11] = [ Self::Codex, Self::Claude, Self::Cursor, @@ -30,6 +31,7 @@ impl ProviderId { Self::Gemini, Self::Copilot, Self::Minimax, + Self::Zai, Self::Kimi, Self::OpenCodeGo, Self::Grok, @@ -45,6 +47,7 @@ impl ProviderId { Self::Antigravity => "Antigravity", Self::Copilot => "Copilot", Self::Minimax => "Minimax", + Self::Zai => "Z.AI Coding Plan", Self::Kimi => "Kimi", Self::OpenCodeGo => "OpenCode Go", Self::Grok => "Grok", @@ -364,6 +367,28 @@ mod tests { assert!(ProviderId::ALL.contains(&ProviderId::Grok)); } + #[test] + fn zai_provider_id_has_serialization_label_and_canonical_order() { + assert_eq!(serde_json::to_string(&ProviderId::Zai).unwrap(), "\"zai\""); + assert_eq!(ProviderId::Zai.label(), "Z.AI Coding Plan"); + assert_eq!( + ProviderId::ALL[ProviderId::ALL + .iter() + .position(|provider| *provider == ProviderId::Minimax) + .unwrap() + + 1], + ProviderId::Zai + ); + assert_eq!( + ProviderId::ALL[ProviderId::ALL + .iter() + .position(|provider| *provider == ProviderId::Zai) + .unwrap() + + 1], + ProviderId::Kimi + ); + } + fn window(label: &str) -> UsageWindow { UsageWindow { label: label.to_string(), diff --git a/src/providers/adapters.rs b/src/providers/adapters.rs index f181c75..aeaaffd 100644 --- a/src/providers/adapters.rs +++ b/src/providers/adapters.rs @@ -10,6 +10,7 @@ mod grok_adapter; mod kimi_adapter; mod minimax_adapter; mod opencode_go_adapter; +mod zai_adapter; use crate::account_storage::ProviderAccountStorage; use crate::config::{Config, host_user_home_dir, paths}; @@ -28,6 +29,7 @@ pub(super) fn adapter(provider: ProviderId) -> &'static dyn ProviderAdapter { ProviderId::Copilot => &COPILOT_ADAPTER, ProviderId::Kimi => &KIMI_ADAPTER, ProviderId::Minimax => &MINIMAX_ADAPTER, + ProviderId::Zai => &ZAI_ADAPTER, ProviderId::Antigravity => &ANTIGRAVITY_ADAPTER, ProviderId::OpenCodeGo => &OPENCODE_GO_ADAPTER, ProviderId::Grok => &GROK_ADAPTER, @@ -41,6 +43,7 @@ static GEMINI_ADAPTER: gemini_adapter::GeminiAdapter = gemini_adapter::GeminiAda static COPILOT_ADAPTER: copilot_adapter::CopilotAdapter = copilot_adapter::CopilotAdapter; static KIMI_ADAPTER: kimi_adapter::KimiAdapter = kimi_adapter::KimiAdapter; static MINIMAX_ADAPTER: minimax_adapter::MinimaxAdapter = minimax_adapter::MinimaxAdapter; +static ZAI_ADAPTER: zai_adapter::ZaiAdapter = zai_adapter::ZaiAdapter; static ANTIGRAVITY_ADAPTER: antigravity_adapter::AntigravityAdapter = antigravity_adapter::AntigravityAdapter; static OPENCODE_GO_ADAPTER: opencode_go_adapter::OpenCodeGoAdapter = diff --git a/src/providers/adapters/zai_adapter.rs b/src/providers/adapters/zai_adapter.rs new file mode 100644 index 0000000..8fc05d6 --- /dev/null +++ b/src/providers/adapters/zai_adapter.rs @@ -0,0 +1,190 @@ +// SPDX-License-Identifier: MPL-2.0 + +use super::reconcile_provider_account_descriptors; +use crate::account_storage::ProviderAccountStorage; +use crate::config::{Config, paths}; +use crate::error::AppError; +use crate::model::{AppState, ProviderId, UsageSnapshot}; +use crate::providers::interface::{ + BoxFuture, ProviderAccountAction, ProviderAccountDescriptor, ProviderAccountHandle, + ProviderAdapter, ProviderCapabilities, ProviderLoginKind, +}; +use crate::providers::zai; + +pub(super) struct ZaiAdapter; + +impl ProviderAdapter for ZaiAdapter { + fn id(&self) -> ProviderId { + ProviderId::Zai + } + + fn login_kind(&self) -> ProviderLoginKind { + ProviderLoginKind::Zai + } + + fn capabilities(&self) -> ProviderCapabilities { + ProviderCapabilities { + supports_background_status_refresh: false, + requires_auth_prompt_on_auth_failure: false, + } + } + + fn selection_required_message(&self) -> Option { + Some(crate::fl!("badge-select-required")) + } + + fn discover_accounts(&self, config: &Config) -> Vec { + zai::discover_accounts(config) + .into_iter() + .filter_map(|account| { + config + .zai_managed_accounts + .iter() + .find(|managed| managed.id == account.id) + .cloned() + .map(|managed| ProviderAccountDescriptor { + provider: self.id(), + account_id: account.id, + label: account.label, + actions: vec![ + ProviderAccountAction::Delete, + ProviderAccountAction::Reauthenticate, + ], + handle: ProviderAccountHandle::Zai(managed), + }) + }) + .collect() + } + + fn sync_managed_accounts(&self, config: &mut Config) -> bool { + zai::sync_managed_accounts(config) + } + + fn delete_account(&self, account_id: &str, config: &mut Config) -> bool { + if !config + .zai_managed_accounts + .iter() + .any(|account| account.id == account_id) + { + return false; + } + if ProviderAccountStorage::new(paths().zai_accounts_dir) + .delete_account(account_id) + .is_err() + { + return false; + } + config + .zai_managed_accounts + .retain(|account| account.id != account_id); + config + .selected_zai_account_ids + .retain(|id| id != account_id); + true + } + + fn reconcile_provider_accounts(&self, config: &Config, state: &mut AppState) { + let accounts = self.discover_accounts(config); + reconcile_provider_account_descriptors(self.id(), config, state, &accounts); + if let Some(provider_state) = state.provider_mut(self.id()) { + provider_state.system_active_account_id = None; + } + } + + fn fetch_account<'a>( + &self, + handle: &'a ProviderAccountHandle, + client: &'a reqwest::Client, + ) -> BoxFuture<'a, crate::error::Result> { + let provider = self.id(); + Box::pin(async move { + match handle { + ProviderAccountHandle::Zai(account) => { + zai::fetch(client, account).await.map_err(AppError::from) + } + _ => Err(AppError::InvalidAccountHandle { provider }), + } + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::{ManagedZaiAccountConfig, paths}; + use chrono::Utc; + use tempfile::tempdir; + + fn account(id: &str) -> ManagedZaiAccountConfig { + let now = Utc::now(); + ManagedZaiAccountConfig { + id: id.to_string(), + label: "Z.AI account".to_string(), + api_key_source: "stored".to_string(), + created_at: now, + updated_at: now, + last_authenticated_at: None, + } + } + + #[test] + fn exposes_managed_key_account_contract() { + let adapter = ZaiAdapter; + assert_eq!(adapter.id(), ProviderId::Zai); + assert_eq!(adapter.login_kind(), ProviderLoginKind::Zai); + assert_eq!( + adapter.capabilities(), + ProviderCapabilities { + supports_background_status_refresh: false, + requires_auth_prompt_on_auth_failure: false, + } + ); + assert!(!adapter.supports_opencode_import()); + assert_eq!(adapter.system_active_account_id(&Config::default()), None); + } + + #[test] + fn discovers_zai_accounts_with_only_delete_and_reauthenticate() { + let adapter = ZaiAdapter; + let config = Config { + zai_managed_accounts: vec![account("zai-1")], + ..Config::default() + }; + + let descriptors = adapter.discover_accounts(&config); + + assert_eq!(descriptors.len(), 1); + assert_eq!(descriptors[0].provider, ProviderId::Zai); + assert_eq!( + descriptors[0].actions, + vec![ + ProviderAccountAction::Delete, + ProviderAccountAction::Reauthenticate + ] + ); + assert!(matches!( + descriptors[0].handle, + ProviderAccountHandle::Zai(_) + )); + } + + #[test] + fn deleting_account_removes_storage_config_and_selection() { + let mut env = crate::test_support::test_env(); + let root = tempdir().unwrap(); + env.set("XDG_STATE_HOME", root.path()); + let id = "zai-delete"; + let storage = ProviderAccountStorage::new(paths().zai_accounts_dir.clone()); + storage.write_text_file(id, "api_key.txt", "key").unwrap(); + let mut config = Config { + zai_managed_accounts: vec![account(id)], + selected_zai_account_ids: vec![id.to_string()], + ..Config::default() + }; + + assert!(ZaiAdapter.delete_account(id, &mut config)); + assert!(config.zai_managed_accounts.is_empty()); + assert!(config.selected_zai_account_ids.is_empty()); + assert!(!paths().zai_accounts_dir.join(id).exists()); + } +} diff --git a/src/providers/interface.rs b/src/providers/interface.rs index 0f33ee4..af630b5 100644 --- a/src/providers/interface.rs +++ b/src/providers/interface.rs @@ -4,7 +4,7 @@ use crate::config::{ Config, ManagedAntigravityAccountConfig, ManagedClaudeAccountConfig, ManagedCodexAccountConfig, ManagedCopilotAccountConfig, ManagedCursorAccountConfig, ManagedGeminiAccountConfig, ManagedGrokAccountConfig, ManagedKimiAccountConfig, ManagedMinimaxAccountConfig, - ManagedOpenCodeGoAccountConfig, + ManagedOpenCodeGoAccountConfig, ManagedZaiAccountConfig, }; use crate::error::AppError; use crate::model::{AppState, AuthState, ProviderAccountRuntimeState, ProviderId, UsageSnapshot}; @@ -67,6 +67,7 @@ pub enum ProviderLoginKind { Gemini, Copilot, Minimax, + Zai, Kimi, Antigravity, OpenCodeGo, @@ -126,6 +127,7 @@ pub enum ProviderAccountHandle { Gemini(ManagedGeminiAccountConfig), Copilot(ManagedCopilotAccountConfig), Minimax(ManagedMinimaxAccountConfig), + Zai(ManagedZaiAccountConfig), Kimi(ManagedKimiAccountConfig), Antigravity(ManagedAntigravityAccountConfig), OpenCodeGo(ManagedOpenCodeGoAccountConfig), diff --git a/src/providers/mod.rs b/src/providers/mod.rs index ca940f1..453bf51 100644 --- a/src/providers/mod.rs +++ b/src/providers/mod.rs @@ -15,3 +15,4 @@ pub mod minimax; pub mod opencode_auth; pub mod opencode_go; pub mod registry; +pub mod zai; diff --git a/src/providers/opencode_auth.rs b/src/providers/opencode_auth.rs index ada3b1c..73f21db 100644 --- a/src/providers/opencode_auth.rs +++ b/src/providers/opencode_auth.rs @@ -4,7 +4,7 @@ use crate::config::host_user_home_dir; use serde::Deserialize; use serde_json::Value; use std::fs; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; pub const OPENCODE_AUTH_PATH_ENV: &str = "YAPCAP_OPENCODE_AUTH_PATH"; pub const OPENCODE_AUTH_CONTENT_ENV: &str = "OPENCODE_AUTH_CONTENT"; @@ -72,10 +72,43 @@ pub fn auth_path() -> Option { pub fn discover(provider_id: &str) -> Option { let auth = auth_json()?; + discover_from_auth(&auth, provider_id) +} + +pub fn discover_api_key(provider_ids: &[&str]) -> Option { + let auth = auth_json()?; + discover_api_key_from_auth(&auth, provider_ids) +} + +pub fn has_api_key_at_path(path: &Path, provider_ids: &[&str]) -> bool { + let auth = fs::read_to_string(path) + .ok() + .and_then(|body| serde_json::from_str(&body).ok()); + auth.as_ref() + .and_then(|auth| discover_api_key_from_auth(auth, provider_ids)) + .is_some() +} + +fn discover_from_auth(auth: &Value, provider_id: &str) -> Option { let entry = auth.get(provider_id)?.clone(); let credential: RawCredential = serde_json::from_value(entry).ok()?; + let credential = into_credential(credential); + + credential.is_valid().then_some(credential) +} + +fn discover_api_key_from_auth(auth: &Value, provider_ids: &[&str]) -> Option { + provider_ids.iter().find_map(|provider_id| { + let credential = discover_from_auth(auth, provider_id)?; + match credential { + OpenCodeCredential::Api { key } => usable_api_key(&key), + OpenCodeCredential::OAuth { .. } | OpenCodeCredential::WellKnown { .. } => None, + } + }) +} - let credential = match credential { +fn into_credential(credential: RawCredential) -> OpenCodeCredential { + match credential { RawCredential::Api { key } => OpenCodeCredential::Api { key }, RawCredential::OAuth { access, @@ -91,12 +124,15 @@ pub fn discover(provider_id: &str) -> Option { enterprise_url, }, RawCredential::WellKnown { key, token } => OpenCodeCredential::WellKnown { key, token }, - }; + } +} - if credential.is_valid() { - Some(credential) - } else { +fn usable_api_key(key: &str) -> Option { + let key = key.trim(); + if key.is_empty() || key.chars().any(char::is_control) { None + } else { + Some(key.to_string()) } } @@ -358,4 +394,68 @@ mod tests { Some(OpenCodeCredential::Api { key }) if key == "file-key" )); } + + #[test] + fn discovers_api_key_in_requested_order() { + let temp = tempdir().unwrap(); + let path = temp.path().join("auth.json"); + fs::write( + &path, + r#"{"zai-coding-plan":{"type":"api","key":"primary"},"zai":{"type":"api","key":"alias"}}"#, + ) + .unwrap(); + let mut env = crate::test_support::test_env(); + env.remove(OPENCODE_AUTH_CONTENT_ENV); + env.set(OPENCODE_AUTH_PATH_ENV, &path); + + assert_eq!( + discover_api_key(&["zai-coding-plan", "zai"]).as_deref(), + Some("primary") + ); + } + + #[test] + fn api_key_discovery_uses_alias_when_primary_is_not_usable() { + let temp = tempdir().unwrap(); + let path = temp.path().join("auth.json"); + fs::write( + &path, + r#"{"zai-coding-plan":{"type":"oauth","refresh":"r","access":"a","expires":1},"zai":{"type":"api","key":" alias "}}"#, + ) + .unwrap(); + let mut env = crate::test_support::test_env(); + env.remove(OPENCODE_AUTH_CONTENT_ENV); + env.set(OPENCODE_AUTH_PATH_ENV, &path); + + assert_eq!( + discover_api_key(&["zai-coding-plan", "zai"]).as_deref(), + Some("alias") + ); + } + + #[test] + fn api_key_presence_at_path_returns_only_usable_api_credential_fact() { + let temp = tempdir().unwrap(); + let path = temp.path().join("auth.json"); + fs::write(&path, r#"{"zai":{"type":"api","key":"key"}}"#).unwrap(); + + assert!(has_api_key_at_path(&path, &["zai-coding-plan", "zai"])); + assert!(!has_api_key_at_path(&path, &["other"])); + } + + #[test] + fn api_key_presence_at_path_rejects_empty_control_and_non_api_credentials() { + let temp = tempdir().unwrap(); + let path = temp.path().join("auth.json"); + + for body in [ + r#"{"zai":{"type":"api","key":" "}}"#, + r#"{"zai":{"type":"api","key":"bad\nkey"}}"#, + r#"{"zai":{"type":"oauth","refresh":"r","access":"a","expires":1}}"#, + "not-json", + ] { + fs::write(&path, body).unwrap(); + assert!(!has_api_key_at_path(&path, &["zai"])); + } + } } diff --git a/src/providers/registry/tests.rs b/src/providers/registry/tests.rs index bdcb038..487556e 100644 --- a/src/providers/registry/tests.rs +++ b/src/providers/registry/tests.rs @@ -32,6 +32,13 @@ fn providers_expose_expected_capabilities() { requires_auth_prompt_on_auth_failure: false, } ); + assert_eq!( + capabilities(ProviderId::Zai), + ProviderCapabilities { + supports_background_status_refresh: false, + requires_auth_prompt_on_auth_failure: false, + } + ); assert_eq!( capabilities(ProviderId::OpenCodeGo), ProviderCapabilities { @@ -77,6 +84,20 @@ fn grok_provider_registered_and_discovers_accounts() { assert_eq!(login_kind(ProviderId::Grok), ProviderLoginKind::Grok); } +#[test] +fn zai_provider_registered_as_managed_key_provider() { + assert_eq!(login_kind(ProviderId::Zai), ProviderLoginKind::Zai); + assert!(!supports_opencode_import(ProviderId::Zai)); + assert_eq!( + selection_required_message(ProviderId::Zai), + Some("Select".to_string()) + ); + assert_eq!( + system_active_account_id(ProviderId::Zai, &Config::default()), + None + ); +} + #[test] fn cursor_supports_background_status_refresh() { assert!(supports_background_status_refresh(ProviderId::Cursor)); @@ -325,6 +346,7 @@ fn host_aware_providers_resolve_system_active_account_id() { (ProviderId::Cursor, true), (ProviderId::Copilot, false), (ProviderId::Minimax, true), + (ProviderId::Zai, false), (ProviderId::Kimi, true), (ProviderId::OpenCodeGo, true), (ProviderId::Grok, true), @@ -406,7 +428,7 @@ fn every_provider_descriptor_declares_supported_account_actions() { ManagedAntigravityAccountConfig, ManagedClaudeAccountConfig, ManagedCodexAccountConfig, ManagedCopilotAccountConfig, ManagedCursorAccountConfig, ManagedGeminiAccountConfig, ManagedGrokAccountConfig, ManagedKimiAccountConfig, ManagedMinimaxAccountConfig, - ManagedOpenCodeGoAccountConfig, paths, + ManagedOpenCodeGoAccountConfig, ManagedZaiAccountConfig, paths, }; use crate::providers::opencode_auth::{OPENCODE_AUTH_CONTENT_ENV, OPENCODE_AUTH_PATH_ENV}; use std::path::PathBuf; @@ -531,6 +553,14 @@ fn every_provider_descriptor_declares_supported_account_actions() { updated_at: now, last_authenticated_at: None, }], + zai_managed_accounts: vec![ManagedZaiAccountConfig { + id: "zai-1".to_string(), + label: "Z.AI account".to_string(), + api_key_source: "stored".to_string(), + created_at: now, + updated_at: now, + last_authenticated_at: None, + }], kimi_managed_accounts: vec![ManagedKimiAccountConfig { id: "kimi-1".to_string(), label: "Kimi account".to_string(), @@ -613,6 +643,13 @@ fn every_provider_descriptor_declares_supported_account_actions() { ProviderAccountAction::Reauthenticate, ], ), + ( + ProviderId::Zai, + vec![ + ProviderAccountAction::Delete, + ProviderAccountAction::Reauthenticate, + ], + ), ( ProviderId::Kimi, vec![ diff --git a/src/providers/zai/account.rs b/src/providers/zai/account.rs new file mode 100644 index 0000000..f09cc8d --- /dev/null +++ b/src/providers/zai/account.rs @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: MPL-2.0 + +use crate::account_storage::validated_account_dir; +use crate::config::{Config, ManagedZaiAccountConfig, paths}; +use std::path::PathBuf; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ZaiAccount { + pub id: String, + pub label: String, + pub config_dir: PathBuf, +} + +pub fn discover_accounts(config: &Config) -> Vec { + config + .zai_managed_accounts + .iter() + .filter_map(|managed| { + let config_dir = validated_account_dir(&paths().zai_accounts_dir, &managed.id).ok()?; + Some(ZaiAccount { + id: managed.id.clone(), + label: managed.label.clone(), + config_dir, + }) + }) + .collect() +} + +pub fn apply_login_account(config: &mut Config, account: ManagedZaiAccountConfig) { + let account_id = account.id.clone(); + config + .zai_managed_accounts + .retain(|existing| existing.id != account_id); + config.zai_managed_accounts.push(account); +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::Utc; + + fn account(id: &str) -> ManagedZaiAccountConfig { + let now = Utc::now(); + ManagedZaiAccountConfig { + id: id.to_string(), + label: id.to_string(), + api_key_source: "stored".to_string(), + created_at: now, + updated_at: now, + last_authenticated_at: Some(now), + } + } + + #[test] + fn applies_login_by_replacing_the_same_account_id() { + let mut config = Config { + zai_managed_accounts: vec![account("zai-1")], + ..Config::default() + }; + let mut replacement = account("zai-1"); + replacement.label = "Replacement".to_string(); + + apply_login_account(&mut config, replacement); + + assert_eq!(config.zai_managed_accounts.len(), 1); + assert_eq!(config.zai_managed_accounts[0].label, "Replacement"); + } + + #[test] + fn discovery_rejects_path_escaping_account_ids() { + let config = Config { + zai_managed_accounts: vec![account("zai-1"), account("../outside")], + ..Config::default() + }; + + let discovered = discover_accounts(&config); + + assert_eq!(discovered.len(), 1); + assert_eq!(discovered[0].id, "zai-1"); + } +} diff --git a/src/providers/zai/login.rs b/src/providers/zai/login.rs new file mode 100644 index 0000000..221b43d --- /dev/null +++ b/src/providers/zai/login.rs @@ -0,0 +1,178 @@ +// SPDX-License-Identifier: MPL-2.0 + +use crate::config::{Config, ManagedZaiAccountConfig}; +use crate::key_authentication::{ + KeyAuthenticationAdapter, KeyAuthenticationTarget, prepare as prepare_key_authentication, + prepare_for_reauth as prepare_key_authentication_for_reauth, +}; +use crate::providers::zai::{opencode, storage}; +use chrono::{DateTime, Utc}; + +pub use crate::key_authentication::{ + KeyAuthenticationEvent as ZaiLoginEvent, KeyAuthenticationState as ZaiLoginState, + KeyAuthenticationStatus as ZaiLoginStatus, +}; + +pub(crate) struct ZaiKeyAuthentication; + +impl KeyAuthenticationAdapter for ZaiKeyAuthentication { + type Account = ManagedZaiAccountConfig; + + fn new_account_id() -> String { + format!( + "zai-{}", + Utc::now().timestamp_nanos_opt().unwrap_or_default() + ) + } + + fn find_account(config: &Config, account_id: &str) -> Result { + config + .zai_managed_accounts + .iter() + .find(|account| account.id == account_id) + .map(|account| KeyAuthenticationTarget { + id: account.id.clone(), + label: account.label.clone(), + created_at: account.created_at, + }) + .ok_or_else(|| "Z.AI account not found".to_string()) + } + + fn discover_api_key() -> Option { + opencode::discover_api_key() + } + + fn empty_key_error() -> String { + "API key is required".to_string() + } + + fn validate( + config: &Config, + account_id: &str, + label: &str, + api_key: &str, + ) -> Result<(), String> { + let api_key = storage::normalize_api_key(api_key)?; + if label.trim().is_empty() { + return Err("Account name is required".to_string()); + } + if config + .zai_managed_accounts + .iter() + .any(|account| account.id != account_id && account.label == label) + { + return Err("An account with this name already exists".to_string()); + } + if config.zai_managed_accounts.iter().any(|account| { + account.id != account_id + && storage::load_api_key(&account.id) + .ok() + .and_then(|stored_key| storage::normalize_api_key(&stored_key).ok()) + .is_some_and(|stored_key| stored_key == api_key) + }) { + return Err("An account with this API key already exists".to_string()); + } + Ok(()) + } + + fn build_account( + account_id: String, + label: String, + created_at: DateTime, + authenticated_at: DateTime, + ) -> Self::Account { + ManagedZaiAccountConfig { + id: account_id, + label, + api_key_source: "stored".to_string(), + created_at, + updated_at: authenticated_at, + last_authenticated_at: Some(authenticated_at), + } + } + + fn persist(account: &Self::Account, api_key: &str) -> Result<(), String> { + storage::write_api_key(&account.id, api_key) + .map_err(|error| format!("Failed to save Z.AI API key: {error}")) + } +} + +pub fn prepare() -> ZaiLoginState { + prepare_key_authentication::() +} + +pub fn prepare_for_reauth(config: Config, account_id: &str) -> Result { + prepare_key_authentication_for_reauth::(&config, account_id) +} + +pub(crate) fn save( + config: &Config, + state: &mut ZaiLoginState, +) -> Result { + state.save::(config) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::paths; + use std::fs; + use tempfile::tempdir; + + #[test] + fn prepares_with_ordered_opencode_key() { + let temp = tempdir().unwrap(); + let path = temp.path().join("auth.json"); + fs::write( + &path, + r#"{"zai-coding-plan":{"type":"api","key":"primary"},"zai":{"type":"api","key":"alias"}}"#, + ) + .unwrap(); + let mut env = crate::test_support::test_env(); + env.set( + crate::providers::opencode_auth::OPENCODE_AUTH_PATH_ENV, + &path, + ); + env.remove(crate::providers::opencode_auth::OPENCODE_AUTH_CONTENT_ENV); + + let state = prepare(); + + assert_eq!(state.api_key, "primary"); + assert!(state.api_key_from_opencode); + } + + #[test] + fn save_normalizes_key_and_keeps_metadata_secret_free() { + let _env = crate::test_support::test_env(); + let mut state = ZaiLoginState::new("zai-test".to_string()); + state.update_label("Test account".to_string()); + state.update_api_key(" key-value ".to_string()); + let config = Config::default(); + + let account = state.save::(&config).unwrap(); + + assert_eq!(storage::load_api_key(&account.id).unwrap(), "key-value"); + assert_eq!(account.api_key_source, "stored"); + assert!( + serde_json::to_string(&account) + .unwrap() + .find("key-value") + .is_none() + ); + let _ = fs::remove_dir_all(paths().zai_accounts_dir.join(account.id)); + } + + #[test] + fn save_rejects_control_characters_without_persisting() { + let _env = crate::test_support::test_env(); + let mut state = ZaiLoginState::new("zai-test-control".to_string()); + state.update_label("Test account".to_string()); + state.update_api_key("valid\nkey".to_string()); + + let result = state.save::(&Config::default()); + + assert!(result.is_err()); + assert!(state.error.is_some()); + assert!(storage::load_api_key("zai-test-control").is_err()); + } +} diff --git a/src/providers/zai/mod.rs b/src/providers/zai/mod.rs new file mode 100644 index 0000000..646f5ad --- /dev/null +++ b/src/providers/zai/mod.rs @@ -0,0 +1,450 @@ +// SPDX-License-Identifier: MPL-2.0 + +pub mod account; +pub mod login; +pub mod opencode; +mod quota; +pub mod storage; + +use crate::config::{Config, ManagedZaiAccountConfig}; +use crate::error::ZaiError; +use crate::model::UsageSnapshot; +use crate::providers::zai::storage::normalize_api_key; + +pub use account::discover_accounts; +pub use login::{ZaiLoginEvent, ZaiLoginState, ZaiLoginStatus}; +pub use quota::parse; + +const ZAI_API_URL: &str = "https://api.z.ai/api/monitor/usage/quota/limit"; + +pub fn sync_managed_accounts(config: &mut Config) -> bool { + let original_len = config.zai_managed_accounts.len(); + config + .zai_managed_accounts + .retain(|account| matches!(account.api_key_source.as_str(), "stored" | "demo")); + config.zai_managed_accounts.len() != original_len +} + +pub async fn fetch( + client: &reqwest::Client, + account: &ManagedZaiAccountConfig, +) -> Result { + fetch_with_endpoint(client, account, ZAI_API_URL).await +} + +pub(crate) async fn fetch_with_endpoint( + _client: &reqwest::Client, + account: &ManagedZaiAccountConfig, + endpoint: &str, +) -> Result { + let client = crate::runtime::http_client_without_redirects(); + let stored_key = storage::load_api_key(&account.id).map_err(|_| ZaiError::LoginRequired)?; + let api_key = normalize_api_key(&stored_key).map_err(|error| { + if error == "API key is required" { + ZaiError::LoginRequired + } else { + ZaiError::InvalidApiKey + } + })?; + + let mut response = send_request(&client, endpoint, &api_key, true).await?; + if response.status() == reqwest::StatusCode::UNAUTHORIZED { + response = send_request(&client, endpoint, &api_key, false).await?; + } + + classify_response(response).await +} + +async fn send_request( + client: &reqwest::Client, + endpoint: &str, + api_key: &str, + bearer: bool, +) -> Result { + let authorization = if bearer { + format!("Bearer {api_key}") + } else { + api_key.to_string() + }; + let authorization = reqwest::header::HeaderValue::from_str(&authorization) + .map_err(|_| ZaiError::InvalidApiKey)?; + + client + .get(endpoint) + .header(reqwest::header::ACCEPT, "application/json") + .header(reqwest::header::AUTHORIZATION, authorization) + .send() + .await + .map_err(ZaiError::UsageRequest) +} + +async fn classify_response(response: reqwest::Response) -> Result { + let status = response.status(); + if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN { + return Err(ZaiError::LoginRequired); + } + if status == reqwest::StatusCode::TOO_MANY_REQUESTS { + let retry_after_secs = response + .headers() + .get(reqwest::header::RETRY_AFTER) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse().ok()); + return Err(ZaiError::RateLimited { retry_after_secs }); + } + if status.is_server_error() || !status.is_success() { + return Err(ZaiError::UsageHttp { + status: status.as_u16(), + }); + } + + let body = response.text().await.map_err(ZaiError::UsageRequest)?; + parse(&body, chrono::Utc::now()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::providers::zai::storage::write_api_key; + use std::time::{Duration, SystemTime, UNIX_EPOCH}; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; + use tokio::time::timeout; + + const SERVER_TIMEOUT: Duration = Duration::from_secs(5); + + fn account_id(suffix: &str) -> String { + let millis = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis(); + format!("zai-test-{millis}-{suffix}") + } + + fn account(id: &str) -> ManagedZaiAccountConfig { + let now = chrono::Utc::now(); + ManagedZaiAccountConfig { + id: id.to_string(), + label: "Z.AI".to_string(), + api_key_source: "stored".to_string(), + created_at: now, + updated_at: now, + last_authenticated_at: Some(now), + } + } + + async fn server(responses: Vec) -> (String, tokio::task::JoinHandle>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let task = tokio::spawn(async move { + let mut headers = Vec::new(); + for response in responses { + let (mut stream, _) = timeout(SERVER_TIMEOUT, listener.accept()) + .await + .expect("test server accept timed out") + .expect("test server accept failed"); + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + loop { + let read = timeout(SERVER_TIMEOUT, stream.read(&mut buffer)) + .await + .expect("test server read timed out") + .expect("test server read failed"); + if read == 0 { + break; + } + request.extend_from_slice(&buffer[..read]); + if request.windows(4).any(|window| window == b"\r\n\r\n") { + break; + } + } + let request = String::from_utf8_lossy(&request); + headers.push( + request + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("authorization") + .then(|| value.trim().to_string()) + }) + .unwrap_or_default() + .to_string(), + ); + timeout(SERVER_TIMEOUT, stream.write_all(response.as_bytes())) + .await + .expect("test server write timed out") + .expect("test server write failed"); + timeout(SERVER_TIMEOUT, stream.shutdown()) + .await + .expect("test server shutdown timed out") + .expect("test server shutdown failed"); + } + headers + }); + (endpoint, task) + } + + async fn redirect_server() -> (String, tokio::task::JoinHandle>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let redirect_endpoint = format!("{endpoint}/redirected"); + let task = tokio::spawn(async move { + let mut headers = Vec::new(); + let (mut stream, _) = timeout(SERVER_TIMEOUT, listener.accept()) + .await + .expect("redirect test server accept timed out") + .expect("redirect test server accept failed"); + let mut request = Vec::new(); + let mut buffer = [0_u8; 1024]; + loop { + let read = timeout(SERVER_TIMEOUT, stream.read(&mut buffer)) + .await + .expect("redirect test server read timed out") + .expect("redirect test server read failed"); + if read == 0 { + break; + } + request.extend_from_slice(&buffer[..read]); + if request.windows(4).any(|window| window == b"\r\n\r\n") { + break; + } + } + let request = String::from_utf8_lossy(&request); + headers.push( + request + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("authorization") + .then(|| value.trim().to_string()) + }) + .unwrap_or_default() + .to_string(), + ); + let redirect = response(302, &format!("Location: {redirect_endpoint}\r\n"), ""); + timeout(SERVER_TIMEOUT, stream.write_all(redirect.as_bytes())) + .await + .expect("redirect test server write timed out") + .expect("redirect test server write failed"); + timeout(SERVER_TIMEOUT, stream.shutdown()) + .await + .expect("redirect test server shutdown timed out") + .expect("redirect test server shutdown failed"); + + if let Ok(Ok((mut redirected, _))) = + timeout(Duration::from_millis(250), listener.accept()).await + { + let mut redirected_request = Vec::new(); + loop { + let read = timeout(SERVER_TIMEOUT, redirected.read(&mut buffer)) + .await + .expect("redirected request read timed out") + .expect("redirected request read failed"); + if read == 0 { + break; + } + redirected_request.extend_from_slice(&buffer[..read]); + if redirected_request + .windows(4) + .any(|window| window == b"\r\n\r\n") + { + break; + } + } + let redirected_request = String::from_utf8_lossy(&redirected_request); + headers.push( + redirected_request + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("authorization") + .then(|| value.trim().to_string()) + }) + .unwrap_or_default() + .to_string(), + ); + let body = r#"{"code":200,"success":true,"data":{"limits":[]}}"#; + let success = response(200, "Content-Type: application/json\r\n", body); + let _ = redirected.write_all(success.as_bytes()).await; + let _ = redirected.shutdown().await; + } + headers + }); + (endpoint, task) + } + + fn response(status: u16, headers: &str, body: &str) -> String { + format!( + "HTTP/1.1 {status} Test\r\nConnection: close\r\nContent-Length: {}\r\n{headers}\r\n{body}", + body.len(), + headers = headers, + body = body, + ) + } + + #[test] + fn sync_keeps_only_stored_and_demo_sources() { + let now = chrono::Utc::now(); + let mut config = Config { + zai_managed_accounts: vec![ + ManagedZaiAccountConfig { + id: "stored".to_string(), + label: "Stored".to_string(), + api_key_source: "stored".to_string(), + created_at: now, + updated_at: now, + last_authenticated_at: None, + }, + ManagedZaiAccountConfig { + id: "demo".to_string(), + label: "Demo".to_string(), + api_key_source: "demo".to_string(), + created_at: now, + updated_at: now, + last_authenticated_at: None, + }, + ManagedZaiAccountConfig { + id: "other".to_string(), + label: "Other".to_string(), + api_key_source: "env:ZAI_API_KEY".to_string(), + created_at: now, + updated_at: now, + last_authenticated_at: None, + }, + ], + ..Config::default() + }; + + assert!(sync_managed_accounts(&mut config)); + assert_eq!(config.zai_managed_accounts.len(), 2); + } + + #[tokio::test] + async fn bearer_success_sends_fixed_auth_shape() { + let _env = crate::test_support::test_env(); + let id = account_id("bearer"); + write_api_key(&id, "key-value").unwrap(); + let body = r#"{"code":200,"success":true,"data":{"limits":[{"type":"TIME_LIMIT","unit":5,"number":1,"percentage":1}]}}"#; + let (endpoint, task) = server(vec![response( + 200, + "Content-Type: application/json\r\n", + body, + )]) + .await; + + let result = fetch_with_endpoint(&reqwest::Client::new(), &account(&id), &endpoint).await; + let headers = task.await.unwrap(); + + assert!(result.is_ok()); + assert_eq!(headers, ["Bearer key-value"]); + let _ = crate::account_storage::ProviderAccountStorage::new( + &crate::config::paths().zai_accounts_dir, + ) + .delete_account(&id); + } + + #[tokio::test] + async fn retries_once_with_raw_authorization_after_bearer_401() { + let _env = crate::test_support::test_env(); + let id = account_id("raw-retry"); + write_api_key(&id, "key-value").unwrap(); + let body = r#"{"code":200,"success":true,"data":{"limits":[{"type":"TIME_LIMIT","unit":5,"number":1,"percentage":1}]}}"#; + let (endpoint, task) = server(vec![ + response(401, "", ""), + response(200, "Content-Type: application/json\r\n", body), + ]) + .await; + + let result = fetch_with_endpoint(&reqwest::Client::new(), &account(&id), &endpoint).await; + let headers = task.await.unwrap(); + + assert!(result.is_ok()); + assert_eq!(headers, ["Bearer key-value", "key-value"]); + let _ = crate::account_storage::ProviderAccountStorage::new( + &crate::config::paths().zai_accounts_dir, + ) + .delete_account(&id); + } + + #[tokio::test] + async fn final_401_is_login_required_without_a_third_request() { + let _env = crate::test_support::test_env(); + let id = account_id("auth-failure"); + write_api_key(&id, "key-value").unwrap(); + let (endpoint, task) = server(vec![response(401, "", ""), response(401, "", "")]).await; + + let result = fetch_with_endpoint(&reqwest::Client::new(), &account(&id), &endpoint).await; + let headers = task.await.unwrap(); + + assert!(matches!(result, Err(ZaiError::LoginRequired))); + assert_eq!(headers, ["Bearer key-value", "key-value"]); + let _ = crate::account_storage::ProviderAccountStorage::new( + &crate::config::paths().zai_accounts_dir, + ) + .delete_account(&id); + } + + #[tokio::test] + async fn rate_limit_preserves_numeric_retry_after() { + let _env = crate::test_support::test_env(); + let id = account_id("rate-limit"); + write_api_key(&id, "key-value").unwrap(); + let (endpoint, task) = server(vec![response(429, "Retry-After: 17\r\n", "")]).await; + + let result = fetch_with_endpoint(&reqwest::Client::new(), &account(&id), &endpoint).await; + let headers = task.await.unwrap(); + + assert!(matches!( + result, + Err(ZaiError::RateLimited { + retry_after_secs: Some(17) + }) + )); + assert_eq!(headers, ["Bearer key-value"]); + let _ = crate::account_storage::ProviderAccountStorage::new( + &crate::config::paths().zai_accounts_dir, + ) + .delete_account(&id); + } + + #[tokio::test] + async fn server_errors_are_transient_without_retrying_raw_auth() { + let _env = crate::test_support::test_env(); + let id = account_id("server-error"); + write_api_key(&id, "key-value").unwrap(); + let (endpoint, task) = server(vec![response(503, "", "")]).await; + + let result = fetch_with_endpoint(&reqwest::Client::new(), &account(&id), &endpoint).await; + let headers = task.await.unwrap(); + + assert!(matches!(&result, Err(ZaiError::UsageHttp { status: 503 }))); + assert_eq!(headers, ["Bearer key-value"]); + assert!( + result + .as_ref() + .err() + .is_some_and(|error| error.is_transient()) + ); + let _ = crate::account_storage::ProviderAccountStorage::new( + &crate::config::paths().zai_accounts_dir, + ) + .delete_account(&id); + } + + #[tokio::test] + async fn redirects_are_rejected_without_forwarding_credentials() { + let _env = crate::test_support::test_env(); + let id = account_id("redirect"); + write_api_key(&id, "key-value").unwrap(); + let (endpoint, task) = redirect_server().await; + + let result = fetch_with_endpoint(&reqwest::Client::new(), &account(&id), &endpoint).await; + let headers = task.await.unwrap(); + + assert!(matches!(result, Err(ZaiError::UsageHttp { status: 302 }))); + assert_eq!(headers, ["Bearer key-value"]); + let _ = crate::account_storage::ProviderAccountStorage::new( + &crate::config::paths().zai_accounts_dir, + ) + .delete_account(&id); + } +} diff --git a/src/providers/zai/opencode.rs b/src/providers/zai/opencode.rs new file mode 100644 index 0000000..6bbf2b9 --- /dev/null +++ b/src/providers/zai/opencode.rs @@ -0,0 +1,67 @@ +// SPDX-License-Identifier: MPL-2.0 + +use std::path::Path; + +const PROVIDER_IDS: [&str; 2] = ["zai-coding-plan", "zai"]; + +pub fn discover_api_key() -> Option { + crate::providers::opencode_auth::discover_api_key(&PROVIDER_IDS) +} + +pub fn has_usable_api_key_at(path: &Path) -> bool { + crate::providers::opencode_auth::has_api_key_at_path(path, &PROVIDER_IDS) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + use tempfile::tempdir; + + #[test] + fn discovers_primary_key_before_alias() { + let temp = tempdir().unwrap(); + let path = temp.path().join("auth.json"); + fs::write( + &path, + r#"{"zai-coding-plan":{"type":"api","key":"primary"},"zai":{"type":"api","key":"alias"}}"#, + ) + .unwrap(); + let mut env = crate::test_support::test_env(); + env.remove(crate::providers::opencode_auth::OPENCODE_AUTH_CONTENT_ENV); + env.set( + crate::providers::opencode_auth::OPENCODE_AUTH_PATH_ENV, + &path, + ); + + assert_eq!(discover_api_key().as_deref(), Some("primary")); + } + + #[test] + fn falls_back_to_alias_when_primary_is_not_an_api_key() { + let temp = tempdir().unwrap(); + let path = temp.path().join("auth.json"); + fs::write( + &path, + r#"{"zai-coding-plan":{"type":"oauth","refresh":"r","access":"a","expires":1},"zai":{"type":"api","key":" alias "}}"#, + ) + .unwrap(); + let mut env = crate::test_support::test_env(); + env.remove(crate::providers::opencode_auth::OPENCODE_AUTH_CONTENT_ENV); + env.set( + crate::providers::opencode_auth::OPENCODE_AUTH_PATH_ENV, + &path, + ); + + assert_eq!(discover_api_key().as_deref(), Some("alias")); + } + + #[test] + fn path_predicate_does_not_expose_the_key() { + let temp = tempdir().unwrap(); + let path = temp.path().join("auth.json"); + fs::write(&path, r#"{"zai":{"type":"api","key":"key"}}"#).unwrap(); + + assert!(has_usable_api_key_at(&path)); + } +} diff --git a/src/providers/zai/quota.rs b/src/providers/zai/quota.rs new file mode 100644 index 0000000..ac66e32 --- /dev/null +++ b/src/providers/zai/quota.rs @@ -0,0 +1,343 @@ +// SPDX-License-Identifier: MPL-2.0 + +use crate::error::ZaiError; +use crate::model::{ProviderId, ProviderIdentity, UsageHeadline, UsageSnapshot, UsageWindow}; +use chrono::{DateTime, Utc}; +use serde::Deserialize; +use serde_json::Value; + +#[derive(Debug, Deserialize)] +struct ZaiQuotaResponse { + code: Option, + success: Option, + data: Option, +} + +#[derive(Debug, Deserialize)] +struct ZaiQuotaData { + #[serde(default, alias = "limit")] + limits: Option>, + level: Option, +} + +#[derive(Debug, Deserialize)] +struct ZaiLimit { + #[serde(rename = "type")] + limit_type: Option, + unit: Option, + number: Option, + usage: Option, + #[serde(rename = "currentValue")] + current_value: Option, + remaining: Option, + percentage: Option, + #[serde(rename = "nextResetTime")] + next_reset_time: Option, + #[serde(rename = "startTime")] + start_time: Option, + #[serde(rename = "endTime")] + end_time: Option, +} + +#[derive(Clone, Copy)] +enum ZaiWindowSlot { + FiveHour, + Weekly, + Mcp, +} + +impl ZaiWindowSlot { + const fn index(self) -> usize { + match self { + Self::FiveHour => 0, + Self::Weekly => 1, + Self::Mcp => 2, + } + } + + const fn label(self) -> &'static str { + match self { + Self::FiveHour => "5 Hour", + Self::Weekly => "Weekly", + Self::Mcp => "MCP", + } + } + + const fn window_seconds(self) -> Option { + match self { + Self::FiveHour => Some(5 * 3600), + Self::Weekly => Some(7 * 24 * 3600), + Self::Mcp => None, + } + } + + fn reset_description(self) -> Option { + match self { + Self::FiveHour => Some("Resets every 5 hours".to_string()), + Self::Weekly => Some("Resets weekly".to_string()), + Self::Mcp => None, + } + } +} + +pub fn parse(body: &str, updated_at: DateTime) -> Result { + let response: ZaiQuotaResponse = serde_json::from_str(body).map_err(ZaiError::DecodeUsage)?; + if response.code != Some(200) || response.success != Some(true) { + return Err(ZaiError::InvalidEnvelope); + } + let data = response.data.ok_or(ZaiError::InvalidEnvelope)?; + let mut slots: [Option; 3] = std::array::from_fn(|_| None); + + for limit in data.limits.unwrap_or_default() { + let Some(slot) = classify(&limit) else { + continue; + }; + let Some(used_percent) = used_percent(&limit) else { + continue; + }; + let window = UsageWindow { + label: slot.label().to_string(), + used_percent, + reset_at: reset_at(&limit), + window_seconds: slot + .window_seconds() + .or_else(|| measured_window_seconds(&limit)), + reset_description: slot.reset_description(), + group: None, + }; + let index = slot.index(); + if slots[index].is_none() { + slots[index] = Some(window); + } + } + + let windows: Vec = slots.into_iter().flatten().collect(); + if windows.is_empty() { + return Err(ZaiError::NoUsageData); + } + + Ok(UsageSnapshot { + provider: ProviderId::Zai, + source: "API Key".to_string(), + updated_at, + headline: UsageHeadline(0), + windows, + provider_cost: None, + extra_usage: None, + identity: ProviderIdentity { + email: None, + account_id: None, + plan: normalized_plan(data.level.as_deref()), + display_name: None, + }, + }) +} + +fn classify(limit: &ZaiLimit) -> Option { + let limit_type = limit.limit_type.as_deref()?; + let unit = integer(limit.unit.as_ref())?; + let number = integer(limit.number.as_ref())?; + match (limit_type, unit, number) { + ("CREDIT_LIMIT" | "TOKENS_LIMIT", 3, 5) => Some(ZaiWindowSlot::FiveHour), + ("CREDIT_LIMIT" | "TOKENS_LIMIT", 6, 1) => Some(ZaiWindowSlot::Weekly), + ("TIME_LIMIT", 5, 1) => Some(ZaiWindowSlot::Mcp), + _ => None, + } +} + +fn used_percent(limit: &ZaiLimit) -> Option { + if let Some(percentage) = finite_number(limit.percentage.as_ref()) { + return Some((percentage as f32).clamp(0.0, 100.0)); + } + let usage = finite_number(limit.usage.as_ref()).filter(|usage| *usage > 0.0)?; + let numerator = finite_number(limit.current_value.as_ref()) + .or_else(|| finite_number(limit.remaining.as_ref()).map(|remaining| usage - remaining))?; + let percent = numerator / usage * 100.0; + percent + .is_finite() + .then_some((percent as f32).clamp(0.0, 100.0)) +} + +fn reset_at(limit: &ZaiLimit) -> Option> { + integer(limit.next_reset_time.as_ref()).and_then(DateTime::from_timestamp_millis) +} + +fn measured_window_seconds(limit: &ZaiLimit) -> Option { + let start = integer(limit.start_time.as_ref())?; + let end = integer(limit.end_time.as_ref())?; + let millis = end.checked_sub(start)?; + (millis > 0) + .then_some(millis / 1000) + .filter(|seconds| *seconds > 0) +} + +fn integer(value: Option<&Value>) -> Option { + match value? { + Value::Number(value) => value.as_i64(), + Value::String(value) => value.trim().parse().ok(), + _ => None, + } +} + +fn finite_number(value: Option<&Value>) -> Option { + let number = match value? { + Value::Number(value) => value.as_f64()?, + Value::String(value) => value.trim().parse().ok()?, + _ => return None, + }; + number.is_finite().then_some(number) +} + +fn normalized_plan(level: Option<&str>) -> Option { + let level = level?.trim(); + if level.is_empty() { + return None; + } + Some(match level.to_ascii_lowercase().as_str() { + "lite" => "Lite".to_string(), + "pro" => "Pro".to_string(), + "max" => "Max".to_string(), + _ => level.to_string(), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + const UPDATED_AT: &str = "2026-08-04T06:21:48Z"; + + fn updated_at() -> DateTime { + DateTime::parse_from_rfc3339(UPDATED_AT) + .unwrap() + .with_timezone(&Utc) + } + + #[test] + fn parses_live_credit_fixture_in_canonical_order() { + let snapshot = parse( + include_str!("../../../fixtures/zai/credit_limit.json"), + updated_at(), + ) + .unwrap(); + + assert_eq!(snapshot.provider, ProviderId::Zai); + assert_eq!(snapshot.source, "API Key"); + assert_eq!(snapshot.identity.plan.as_deref(), Some("Pro")); + assert_eq!(snapshot.windows.len(), 2); + assert_eq!(snapshot.windows[0].label, "5 Hour"); + assert_eq!(snapshot.windows[0].used_percent, 0.0); + assert_eq!(snapshot.windows[0].window_seconds, Some(5 * 3600)); + assert_eq!(snapshot.windows[1].label, "Weekly"); + assert_eq!(snapshot.windows[1].used_percent, 100.0); + assert_eq!(snapshot.windows[1].window_seconds, Some(7 * 24 * 3600)); + assert!(snapshot.windows[1].reset_at.is_some()); + } + + #[test] + fn parses_token_and_mcp_fixture_without_inventing_mcp_duration() { + let snapshot = parse( + include_str!("../../../fixtures/zai/token_mcp.json"), + updated_at(), + ) + .unwrap(); + + assert_eq!(snapshot.windows.len(), 3); + assert_eq!( + snapshot + .windows + .iter() + .map(|window| window.label.as_str()) + .collect::>(), + ["5 Hour", "Weekly", "MCP"] + ); + assert_eq!(snapshot.windows[2].window_seconds, None); + assert_eq!(snapshot.identity.plan.as_deref(), Some("Max")); + } + + #[test] + fn parses_mcp_only_fixture_as_successful_usage() { + let snapshot = parse( + include_str!("../../../fixtures/zai/mcp_only.json"), + updated_at(), + ) + .unwrap(); + + assert_eq!(snapshot.windows.len(), 1); + assert_eq!(snapshot.windows[0].label, "MCP"); + assert_eq!(snapshot.windows[0].used_percent, 25.0); + assert_eq!(snapshot.windows[0].window_seconds, None); + } + + #[test] + fn input_order_does_not_change_window_order() { + let body = r#"{ + "code":200, + "success":true, + "data":{"level":"custom","limits":[ + {"type":"TIME_LIMIT","unit":5,"number":1,"percentage":20}, + {"type":"CREDIT_LIMIT","unit":6,"number":1,"percentage":40}, + {"type":"CREDIT_LIMIT","unit":3,"number":5,"percentage":10} + ]} + }"#; + + let snapshot = parse(body, updated_at()).unwrap(); + + assert_eq!( + snapshot + .windows + .iter() + .map(|window| window.label.as_str()) + .collect::>(), + ["5 Hour", "Weekly", "MCP"] + ); + assert_eq!(snapshot.identity.plan.as_deref(), Some("custom")); + } + + #[test] + fn derives_and_clamps_usage_percent_without_creating_invalid_rows() { + let body = r#"{ + "code":200, + "success":true, + "data":{"limits":[ + {"type":"CREDIT_LIMIT","unit":3,"number":5,"usage":100,"currentValue":150}, + {"type":"CREDIT_LIMIT","unit":6,"number":1,"usage":100,"remaining":25}, + {"type":"TIME_LIMIT","unit":5,"number":1,"usage":0,"remaining":0}, + {"type":"CREDIT_LIMIT","unit":99,"number":1,"percentage":0} + ]} + }"#; + + let snapshot = parse(body, updated_at()).unwrap(); + + assert_eq!(snapshot.windows.len(), 2); + assert_eq!(snapshot.windows[0].used_percent, 100.0); + assert_eq!(snapshot.windows[1].used_percent, 75.0); + } + + #[test] + fn rejects_unsuccessful_empty_and_invalid_envelopes() { + for body in [ + r#"{"code":500,"success":false,"data":{"limits":[]}}"#, + r#"{"code":200,"success":true}"#, + r#"{"code":200,"success":true,"data":{"limits":[]}}"#, + r#"{"code":200,"success":true,"data":{"limits":[{"type":"UNKNOWN","unit":3,"number":5,"percentage":0}]}}"#, + ] { + assert!(parse(body, updated_at()).is_err()); + } + } + + #[test] + fn accepts_reported_mcp_start_and_end_span() { + let body = r#"{ + "code":200, + "success":true, + "data":{"limits":[ + {"type":"TIME_LIMIT","unit":5,"number":1,"percentage":20,"startTime":1000,"endTime":86401000} + ]} + }"#; + + let snapshot = parse(body, updated_at()).unwrap(); + + assert_eq!(snapshot.windows[0].window_seconds, Some(86400)); + } +} diff --git a/src/providers/zai/storage.rs b/src/providers/zai/storage.rs new file mode 100644 index 0000000..165788a --- /dev/null +++ b/src/providers/zai/storage.rs @@ -0,0 +1,103 @@ +// SPDX-License-Identifier: MPL-2.0 + +use crate::account_storage::ProviderAccountStorage; +use std::path::Path; + +pub const API_KEY_FILE: &str = "api_key.txt"; + +pub fn write_api_key(account_id: &str, api_key: &str) -> Result<(), String> { + write_api_key_at( + &crate::config::paths().zai_accounts_dir, + account_id, + api_key, + ) +} + +pub fn load_api_key(account_id: &str) -> Result { + load_api_key_at(&crate::config::paths().zai_accounts_dir, account_id) +} + +pub(crate) fn delete_account(account_id: &str) -> Result<(), String> { + ProviderAccountStorage::new(&crate::config::paths().zai_accounts_dir) + .delete_account(account_id) + .map(|_| ()) + .map_err(|error| error.to_string()) +} + +pub(crate) fn write_api_key_at(root: &Path, account_id: &str, api_key: &str) -> Result<(), String> { + let api_key = normalize_api_key(api_key)?; + ProviderAccountStorage::new(root) + .write_text_file(account_id, API_KEY_FILE, &api_key) + .map_err(|error| error.to_string()) +} + +pub(crate) fn load_api_key_at(root: &Path, account_id: &str) -> Result { + ProviderAccountStorage::new(root) + .read_text_file(account_id, API_KEY_FILE) + .map_err(|error| error.to_string()) +} + +pub(crate) fn normalize_api_key(api_key: &str) -> Result { + let api_key = api_key.trim(); + if api_key.is_empty() { + return Err("API key is required".to_string()); + } + if api_key.chars().any(char::is_control) { + return Err("API key contains invalid characters".to_string()); + } + Ok(api_key.to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + use tempfile::tempdir; + + #[test] + fn writes_normalized_key_under_managed_root() { + let temp = tempdir().unwrap(); + let root = temp.path().join("zai-accounts"); + + write_api_key_at(&root, "zai-1", " test-key ").unwrap(); + + assert_eq!(load_api_key_at(&root, "zai-1").unwrap(), "test-key"); + assert!(root.join("zai-1").join(API_KEY_FILE).exists()); + } + + #[test] + fn rejects_empty_and_control_character_keys() { + let temp = tempdir().unwrap(); + let root = temp.path().join("zai-accounts"); + + assert!(write_api_key_at(&root, "zai-1", " \t ").is_err()); + assert!(write_api_key_at(&root, "zai-1", "bad\nkey").is_err()); + } + + #[cfg(unix)] + #[test] + fn rejects_symlinked_account_and_api_key_paths() { + use std::os::unix::fs::symlink; + + let temp = tempdir().unwrap(); + let root = temp.path().join("zai-accounts"); + let outside = temp.path().join("outside"); + fs::create_dir_all(&outside).unwrap(); + let target = outside.join(API_KEY_FILE); + fs::write(&target, "unchanged").unwrap(); + fs::create_dir_all(&root).unwrap(); + symlink(&outside, root.join("zai-1")).unwrap(); + + assert!(write_api_key_at(&root, "zai-1", "replacement").is_err()); + assert!(load_api_key_at(&root, "zai-1").is_err()); + assert_eq!(fs::read_to_string(&target).unwrap(), "unchanged"); + + fs::remove_file(root.join("zai-1")).unwrap(); + fs::create_dir(root.join("zai-1")).unwrap(); + symlink(&target, root.join("zai-1").join(API_KEY_FILE)).unwrap(); + + assert!(write_api_key_at(&root, "zai-1", "replacement").is_err()); + assert!(load_api_key_at(&root, "zai-1").is_err()); + assert_eq!(fs::read_to_string(&target).unwrap(), "unchanged"); + } +} diff --git a/src/runtime.rs b/src/runtime.rs index c1e67a7..d0d80bc 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -10,6 +10,7 @@ use crate::model::{ use crate::providers; use crate::shared_state::{SharedRuntimeState, SharedStateWriter}; use chrono::Utc; +use std::sync::OnceLock; use std::time::Duration; pub(crate) const HTTP_TIMEOUT: Duration = Duration::from_secs(20); @@ -74,6 +75,34 @@ pub fn http_client() -> reqwest::Client { }) } +pub(crate) fn http_client_without_redirects() -> reqwest::Client { + static CLIENT: OnceLock = OnceLock::new(); + + CLIENT + .get_or_init(|| { + let builder = reqwest::Client::builder() + .timeout(HTTP_TIMEOUT) + .connect_timeout(HTTP_CONNECT_TIMEOUT) + .redirect(reqwest::redirect::Policy::none()); + #[cfg(debug_assertions)] + let builder = apply_debug_offline_proxy(builder); + + builder.build().unwrap_or_else(|error| { + tracing::warn!( + error = %error, + "failed to build no-redirect HTTP client" + ); + reqwest::Client::builder() + .timeout(HTTP_TIMEOUT) + .connect_timeout(HTTP_CONNECT_TIMEOUT) + .redirect(reqwest::redirect::Policy::none()) + .build() + .expect("no-redirect HTTP client must build") + }) + }) + .clone() +} + #[cfg(debug_assertions)] fn apply_debug_offline_proxy(builder: reqwest::ClientBuilder) -> reqwest::ClientBuilder { if !std::env::var(DEBUG_OFFLINE_ENV).is_ok_and(|value| debug_env_value_enabled(&value)) {