diff --git a/src/handlers/session.ts b/src/handlers/session.ts index 5d1b4d8..2cbf8cf 100644 --- a/src/handlers/session.ts +++ b/src/handlers/session.ts @@ -73,6 +73,9 @@ export async function newSession( if (!sid) throw new Error("zcode create returned no sessionId"); server.sessionMap.set(sid, sid); + // Only freshly-created sessions are eligible for auto-title on first + // end_turn; resumed/loaded sessions already have a title and must keep it. + server.titleEligibleSessions.add(sid); log(`session/new → ${sid}`); // Sync to the App's tasks-index.sqlite so the App UI shows this session. @@ -342,10 +345,16 @@ export async function prompt( turn, ); - // Session title: set once on the first end_turn of this session. The title - // is the first prompt text (truncated). Subsequent turns never overwrite it - // (set-once gate), and the App's title_overridden flag always wins. - if (result.stopReason === "end_turn" && !server.sessionTitles.has(params.sessionId)) { + // Session title: set once on the first end_turn, but ONLY for freshly + // created sessions. Resumed/loaded sessions already carry a title from + // their history and must not be overwritten by the first post-load message. + // sessionTitles enforces set-once within a session; titleEligibleSessions + // gates which sessions are titled at all. + if ( + result.stopReason === "end_turn" && + server.titleEligibleSessions.has(params.sessionId) && + !server.sessionTitles.has(params.sessionId) + ) { const title = text.slice(0, 80); server.sessionTitles.set(params.sessionId, title); const { updateSessionTitle } = await import("../tasks-index.js"); diff --git a/src/handlers/slash.ts b/src/handlers/slash.ts index 5658ab2..0ef772c 100644 --- a/src/handlers/slash.ts +++ b/src/handlers/slash.ts @@ -6,6 +6,9 @@ * feedback `agent_message_chunk`, and return `end_turn` — never reaching the * normal turn loop. Unknown `/x` falls through to the model (extensibility). * + * `/quota` is the exception: it does not call ZCode at all — it queries the + * GLM Coding Plan usage API directly and renders the result. + * * Returns the PromptResponse when intercepted, or null to let the caller run a * normal turn. */ @@ -16,6 +19,7 @@ import type * as acp from "@agentclientprotocol/sdk"; import { RequestError } from "@agentclientprotocol/sdk"; import { applyModelSwitch } from "../config/runtime-model.js"; import { emitConfigOptionUpdate } from "../config/options.js"; +import { formatQuota, queryQuota } from "../quota/index.js"; import { CONFIG_DISPATCH, warn } from "../utils.js"; import type { ZcodeAcpServer } from "../server.js"; import { sendTextChunk } from "./io.js"; @@ -48,6 +52,13 @@ export async function handleSlashCommand( try { switch (cmd) { + case "quota": { + // Does not touch ZCode — queries the GLM usage API directly. Always + // returns a status line (success card or an error fallback), so this + // never throws into the catch below under normal conditions. + const result = await queryQuota(); + return ok(formatQuota(result)); + } case "compact": { const result = (await compact(server, { sessionId: acpSid }, cx)) as { __lockTimeout?: boolean; diff --git a/src/quota/cache.ts b/src/quota/cache.ts new file mode 100644 index 0000000..96074b5 --- /dev/null +++ b/src/quota/cache.ts @@ -0,0 +1,39 @@ +/** + * In-memory TTL cache for quota results. + * + * `/quota` is user-triggered, so the cache only guards against rapid repeats + * (e.g. double-entering the command). The TTL is intentionally short — 10s — + * so the displayed figures stay fresh while still debouncing bursts. There is + * no file persistence and no per-key sharding: a single process-wide slot. + */ + +import type { QuotaResult } from "./types.js"; + +/** How long a cached result is served before re-querying. */ +const TTL_MS = 10_000; + +let slot: { result: QuotaResult; at: number } | null = null; + +/** Inject a fake clock — only tests should need this. */ +let now: () => number = () => Date.now(); + +/** Return the cached result if still fresh, else `null`. */ +export function getCached(): QuotaResult | null { + if (slot && now() - slot.at < TTL_MS) return slot.result; + return null; +} + +/** Store a fresh result, stamping it with the current time. */ +export function setCached(result: QuotaResult): void { + slot = { result, at: now() }; +} + +/** Clear the cache (test helper). */ +export function clearCache(): void { + slot = null; +} + +/** Override the clock (test-only). Pass `undefined` to restore real time. */ +export function setClock(fn?: () => number): void { + now = fn ?? (() => Date.now()); +} diff --git a/src/quota/client.ts b/src/quota/client.ts new file mode 100644 index 0000000..76e8191 --- /dev/null +++ b/src/quota/client.ts @@ -0,0 +1,73 @@ +/** + * GLM quota API HTTP client. + * + * Queries `bigmodel.cn` (CN) or `api.z.ai` (intl) for the current account's + * usage limits. Credentials come from the active ZCode provider in + * `~/.zcode/v2/config.json` via {@link loadZcodeCredentials} — the same apiKey + * the backend already uses for model calls. + */ + +import { loadZcodeCredentials } from "../backend/credentials.js"; + +/** Quota endpoint path, appended to the chosen host. */ +const QUOTA_PATH = "/api/monitor/usage/quota/limit"; + +/** Request timeout (ms). The endpoint is fast; keep this tight. */ +const TIMEOUT_MS = 8000; + +/** Hosts for the CN and intl deployments. */ +const HOST_CN = "https://open.bigmodel.cn"; +const HOST_INTL = "https://api.z.ai"; + +/** Raw response from the quota endpoint, pre-parsing. */ +export interface QuotaResponse { + status: number; + json: unknown; + text: string; +} + +/** + * Pick the quota host from a provider `baseURL`. + * + * `api.z.ai` → intl; anything else (including `open.bigmodel.cn`, + * `bigmodel.cn`, empty) → CN. This mirrors how the backend routes model + * traffic. + */ +export function resolveQuotaHost(baseURL: string): string { + return baseURL.includes("api.z.ai") ? HOST_INTL : HOST_CN; +} + +/** + * Fetch the quota envelope from the GLM API. + * + * @throws if the active provider has no apiKey in config, or on network/timeout + * errors. The caller maps these to `unavailable`. + */ +export async function fetchQuotaResponse( + fetchImpl: typeof globalThis.fetch = globalThis.fetch, +): Promise { + const { ANTHROPIC_API_KEY, ZCODE_BASE_URL } = loadZcodeCredentials(); + if (!ANTHROPIC_API_KEY) { + throw new Error("no apiKey in ZCode config — cannot query quota"); + } + + const url = resolveQuotaHost(ZCODE_BASE_URL ?? "") + QUOTA_PATH; + const resp = await fetchImpl(url, { + method: "GET", + headers: { + Accept: "application/json, text/plain, */*", + Authorization: `Bearer ${ANTHROPIC_API_KEY}`, + }, + signal: AbortSignal.timeout(TIMEOUT_MS), + }); + + const text = await resp.text(); + let json: unknown = null; + try { + json = text ? JSON.parse(text) : null; + } catch { + // Non-JSON body — leave json null; the parser treats it as unavailable. + } + + return { status: resp.status, json, text }; +} diff --git a/src/quota/format.ts b/src/quota/format.ts new file mode 100644 index 0000000..4dfa6d3 --- /dev/null +++ b/src/quota/format.ts @@ -0,0 +1,126 @@ +/** + * Quota result formatting — renders a {@link QuotaResult} as a multi-line + * plain-text card for the `agent_message_chunk` feedback. + * + * The progress bar uses two block characters at {@link BAR_WIDTH}-cell width + * for a clean, precise fill (each cell = 10%): + * █ (U+2588, full) — used portion + * ░ (U+2591, light) — remaining portion + */ + +import type { QuotaItem, QuotaResult } from "./types.js"; + +/** Progress-bar cell count. Compact to fit chat-frame width without wrapping. */ +const BAR_WIDTH = 10; + +/** Filled / empty cell characters. */ +const CHAR_FULL = "█"; +const CHAR_EMPTY = "░"; + +/** Pad a percent number to 2 chars (right-aligned). */ +function padPercent(n: number): string { + return String(n).padStart(2); +} + +/** + * Render a {@link BAR_WIDTH}-cell progress bar for a used-percent in [0, 100]. + * + * 0% → all empty; 100% → all full; otherwise the used portion is `█` and the + * rest is `░`, rounded to the nearest cell (each cell = 10%). + */ +export function renderBar(usedPercent: number): string { + const clamped = Math.max(0, Math.min(100, usedPercent)); + const filled = Math.round((clamped / 100) * BAR_WIDTH); + const empty = BAR_WIDTH - filled; + return CHAR_FULL.repeat(filled) + CHAR_EMPTY.repeat(empty); +} + +/** Format a reset timestamp (ms) as a local `MM-DD HH:MM` string, or `null`. */ +function formatResetTime(nextResetTime?: number): string | null { + if (nextResetTime === undefined || !Number.isFinite(nextResetTime)) return null; + const d = new Date(nextResetTime); + if (Number.isNaN(d.getTime())) return null; + const mm = String(d.getMonth() + 1).padStart(2, "0"); + const dd = String(d.getDate()).padStart(2, "0"); + const hh = String(d.getHours()).padStart(2, "0"); + const min = String(d.getMinutes()).padStart(2, "0"); + return `${mm}-${dd} ${hh}:${min}`; +} + +/** Capitalise the first letter of a plan level (e.g. "pro" → "Pro"). */ +function capitalise(s: string): string { + if (!s) return ""; + return s.charAt(0).toUpperCase() + s.slice(1); +} + +/** + * Build the trailing annotation for an item: reset time first (so all items + * align), then `(used/total)` last (only some limits carry absolute counters). + */ +function formatTrailing(item: QuotaItem): string { + const parts: string[] = []; + + const reset = formatResetTime(item.nextResetTime); + if (reset) parts.push(reset); + + // Absolute counters — only some limit kinds carry them (MCP/TIME does, + // legacy TOKENS_LIMIT often does not). Placed last so the reset times of + // counter-less items (e.g. 5h) stay left-aligned with counter-bearing ones. + if ( + typeof item.usedCount === "number" && + typeof item.totalCount === "number" && + Number.isFinite(item.usedCount) && + Number.isFinite(item.totalCount) + ) { + parts.push(`(${item.usedCount}/${item.totalCount})`); + } + + return parts.length > 0 ? ` · ${parts.join(" · ")}` : ""; +} + +/** Right-pad a model code for aligned detail sub-lines. */ +const DETAIL_LABEL_WIDTH = 14; + +/** Render one quota item line (+ indented detail sub-lines if present). */ +function formatItem(item: QuotaItem): string[] { + const lines: string[] = []; + const bar = renderBar(item.usedPercent); + lines.push( + `${item.label.padEnd(5)} ${bar} ${padPercent(item.usedPercent)}%${formatTrailing(item)}`, + ); + + if (item.detail && item.detail.length > 0) { + const last = item.detail.length - 1; + item.detail.forEach((d, i) => { + const branch = i === last ? "└" : "├"; + const name = d.modelCode.padEnd(DETAIL_LABEL_WIDTH); + lines.push(` ${branch} ${name}${d.usage}`); + }); + } + return lines; +} + +/** Error/fallback messages keyed by `kind`. */ +const STATUS_MESSAGES: Record, string> = { + auth_error: "🔒 Quota auth expired — re-login in the ZCode app", + rate_limited: "⏳ Quota service busy, try again shortly", + unavailable: "⚠ Quota info unavailable", +}; + +/** + * Render a {@link QuotaResult} as a multi-line plain-text card. + * + * Success → header line + divider + one progress-bar line per item (plus + * indented per-model detail where present). Non-success kinds → a single + * explanatory line. + */ +export function formatQuota(result: QuotaResult): string { + if (result.kind !== "success") { + return STATUS_MESSAGES[result.kind]; + } + + const header = `GLM Coding Plan${result.level ? ` · ${capitalise(result.level)}` : ""}`; + const divider = "─".repeat(50); + const body = result.items.flatMap(formatItem); + return [header, divider, ...body].join("\n"); +} diff --git a/src/quota/index.ts b/src/quota/index.ts new file mode 100644 index 0000000..0c05937 --- /dev/null +++ b/src/quota/index.ts @@ -0,0 +1,47 @@ +/** + * Quota orchestration — the single entry point used by the `/quota` slash + * command. + * + * Flow: cache check → fetch → parse → cache write. Any thrown error (missing + * apiKey, network/timeout) degrades to `unavailable` rather than propagating, + * so the command always produces a user-visible message. + */ + +import { getCached, setCached } from "./cache.js"; +import { fetchQuotaResponse } from "./client.js"; +import { log } from "../utils.js"; +import { parseQuotaEnvelope } from "./parse.js"; +import type { QuotaResult } from "./types.js"; + +/** + * Query the GLM quota API and return a normalised {@link QuotaResult}. + * + * Serves a cached result when fresh (< 10s). On fetch failure the result is + * `unavailable`; on a successful fetch the parsed envelope is cached and + * returned. + */ +export async function queryQuota(): Promise { + const cached = getCached(); + if (cached) { + log("quota: serving cached result"); + return cached; + } + + let result: QuotaResult; + try { + const resp = await fetchQuotaResponse(); + result = parseQuotaEnvelope(resp); + } catch (e) { + log(`quota: fetch failed (${e instanceof Error ? e.message : String(e)})`); + result = { kind: "unavailable" }; + } + + setCached(result); + return result; +} + +// Re-exports for consumers (slash handler + tests). +export { formatQuota } from "./format.js"; +export { parseLimit, parseQuotaEnvelope } from "./parse.js"; +export { renderBar } from "./format.js"; +export type { QuotaItem, QuotaResult, RawLimit } from "./types.js"; diff --git a/src/quota/parse.ts b/src/quota/parse.ts new file mode 100644 index 0000000..7e6cdf6 --- /dev/null +++ b/src/quota/parse.ts @@ -0,0 +1,236 @@ +/** + * Quota response parsing — turns the raw GLM `limits[]` array into a flat list + * of normalised {@link QuotaItem}s. + * + * Design goals (vs. the reference project): + * - **No hardcoded type whitelist.** Unknown `type` values still render, + * labelled by their raw type string, so future windows surface + * automatically. + * - **Percentage fallback chain** that tolerates inconsistent field sets + * across limit kinds (some carry absolute counters, some only a legacy + * `percentage`). + * - **Labels derived from `type` + `number`**, without magic-number coupling. + */ + +import type { QuotaItem, QuotaResult, RawLimit } from "./types.js"; + +/** Coerce an unknown value to a finite number, or `null`. */ +function asFiniteNumber(value: unknown): number | null { + const n = typeof value === "number" ? value : Number(value); + return Number.isFinite(n) ? n : null; +} + +/** Clamp a percentage to the valid [0, 100] range, or `null` if not finite. */ +function clampPercent(value: unknown): number | null { + const n = asFiniteNumber(value); + if (n === null) return null; + return Math.max(0, Math.min(100, Math.round(n))); +} + +/** + * Compute `{ leftPercent, usedPercent }` from a single raw limit, tolerating + * several field combinations (preferred → fallback): + * 1. `remaining` + `currentValue` (absolute counters) — most precise; + * 2. `currentValue` / `usage` — used over total; + * 3. `percentage` (legacy, treated as *used* percent) — last resort. + * Returns `null` when no percentage can be derived. + */ +function computePercentages(limit: RawLimit): { + leftPercent: number; + usedPercent: number; +} | null { + const usage = asFiniteNumber(limit.usage); + const remaining = asFiniteNumber(limit.remaining); + const currentValue = asFiniteNumber(limit.currentValue); + + // Path 1: absolute remaining/used counters → derive total from parts. + const totalFromParts = + remaining !== null && currentValue !== null ? remaining + currentValue : null; + const total = totalFromParts !== null && totalFromParts > 0 ? totalFromParts : usage; + + if (total !== null && total > 0) { + if (remaining !== null && remaining >= 0 && remaining <= total) { + const leftPercent = clampPercent((remaining / total) * 100); + if (leftPercent !== null) { + return { leftPercent, usedPercent: 100 - leftPercent }; + } + } + if (currentValue !== null && currentValue >= 0 && currentValue <= total) { + const usedPercent = clampPercent((currentValue / total) * 100); + if (usedPercent !== null) { + return { leftPercent: 100 - usedPercent, usedPercent }; + } + } + } + + // Path 3: legacy `percentage` (treated as used percent — GLM convention). + const usedPercent = clampPercent(limit.percentage); + if (usedPercent === null) return null; + return { leftPercent: 100 - usedPercent, usedPercent }; +} + +/** Map a `type` to a coarse category used for label/key derivation. */ +function categoryOf(type: string): "token" | "mcp" | "other" { + switch (type) { + case "TOKENS_LIMIT": + return "token"; + case "MCP_LIMIT": + case "TIME_LIMIT": + return "mcp"; + default: + return "other"; + } +} + +/** + * Derive a stable key and human-readable label for a limit. + * + * Known windows get friendly labels; anything else falls back to the raw + * `type` (or `"Quota"`) so it still renders. + */ +function deriveLabel( + limit: RawLimit, + category: "token" | "mcp" | "other", +): { + key: string; + label: string; +} { + const number = limit.number; + if (category === "token") { + // GLM convention: number===5 → 5-hour window; number===7 → weekly. + if (number === 5) return { key: "token_5h", label: "5h" }; + if (number === 7) return { key: "token_week", label: "Week" }; + if (number !== undefined) return { key: `token_${number}`, label: `${number}` }; + return { key: "token", label: "Token" }; + } + if (category === "mcp") { + return { key: "mcp", label: "MCP" }; + } + // Unknown type — surface it verbatim rather than hide it. + const fallback = limit.type || "Quota"; + return { key: fallback.toLowerCase(), label: fallback }; +} + +/** + * Parse one raw limit into a {@link QuotaItem}, or `null` if no percentage + * can be computed (the limit is then dropped from display). + */ +export function parseLimit(limit: RawLimit): QuotaItem | null { + const percent = computePercentages(limit); + if (!percent) return null; + + const category = categoryOf(limit.type); + const { key, label } = deriveLabel(limit, category); + + const item: QuotaItem = { + key, + label, + usedPercent: percent.usedPercent, + leftPercent: percent.leftPercent, + }; + + // Absolute counters — only some limit kinds carry them (MCP/TIME_LIMIT do, + // legacy TOKENS_LIMIT often does not). Surface them when present so the + // formatter can show "(used / total)". + const usedCount = asFiniteNumber(limit.currentValue); + if (usedCount !== null) item.usedCount = usedCount; + const totalCount = asFiniteNumber(limit.usage); + if (totalCount !== null) item.totalCount = totalCount; + + const nextResetTime = asFiniteNumber(limit.nextResetTime); + if (nextResetTime !== null) item.nextResetTime = nextResetTime; + + if (Array.isArray(limit.usageDetails) && limit.usageDetails.length > 0) { + item.detail = limit.usageDetails; + } + + return item; +} + +/** Auth-failure detection from the business-layer `msg` (zh/en keywords). */ +function isAuthFailureMessage(msg: unknown): boolean { + return typeof msg === "string" && /authorization|auth|token|鉴权|授权|未登录/i.test(msg); +} + +/** Rate-limit detection from the business-layer `msg` (zh/en keywords). */ +function isRateLimitedMessage(msg: unknown): boolean { + return ( + typeof msg === "string" && + /rate\s*limit|too many requests|too frequent|frequency|限流|频率|过于频繁|稍后再试/i.test(msg) + ); +} + +/** Shape of the successful API envelope we expect: `{ success, data: { level, limits } }`. */ +interface QuotaEnvelope { + success?: boolean; + code?: number; + msg?: string; + data?: { level?: string; limits?: RawLimit[] }; +} + +/** + * Parse the full API response (already classified by the fetch layer as a + * `response` with a status + body) into a {@link QuotaResult}. + * + * State machine: 429/rate-limit text → `rate_limited`; auth codes/text → + * `auth_error`; success with zero parseable items → `unavailable`; otherwise + * the items are returned in original order. + */ +export function parseQuotaEnvelope(envelope: { + status: number; + json: unknown; + text: string; +}): QuotaResult { + // HTTP-level rate limit. + if (envelope.status === 429 || isRateLimitedMessage(envelope.text)) { + return { kind: "rate_limited" }; + } + + const payload = envelope.json as QuotaEnvelope | null; + if (!payload || typeof payload !== "object") { + return { kind: "unavailable" }; + } + + if (payload.success !== true) { + if (payload.code === 1001 || payload.code === 401 || isAuthFailureMessage(payload.msg)) { + return { kind: "auth_error" }; + } + if (isRateLimitedMessage(payload.msg)) { + return { kind: "rate_limited" }; + } + return { kind: "unavailable" }; + } + + const limits = Array.isArray(payload.data?.limits) ? payload.data!.limits! : []; + const items = limits + .map(parseLimit) + .filter((x): x is QuotaItem => x !== null) + .sort(byDisplayOrder); + if (items.length === 0) { + return { kind: "unavailable" }; + } + + return { + kind: "success", + level: typeof payload.data?.level === "string" ? payload.data.level : "", + items, + }; +} + +/** + * Display ordering: token windows first (5h before Week before other token + * windows), then MCP, then anything else — by stable key. Keeps the card + * layout predictable regardless of the API's array order. + */ +function byDisplayOrder(a: QuotaItem, b: QuotaItem): number { + return rankOf(a.key) - rankOf(b.key) || a.key.localeCompare(b.key); +} + +/** Stable rank for a quota key — lower renders first. */ +function rankOf(key: string): number { + if (key === "token_5h") return 0; + if (key === "token_week") return 1; + if (key.startsWith("token_")) return 2; // other token windows + if (key === "mcp") return 3; + return 4; // unknown types last +} diff --git a/src/quota/types.ts b/src/quota/types.ts new file mode 100644 index 0000000..dbc158d --- /dev/null +++ b/src/quota/types.ts @@ -0,0 +1,62 @@ +/** + * Type definitions for the GLM Coding Plan quota feature. + * + * The `/quota` slash command queries `bigmodel.cn` / `api.z.ai` for the + * current account's usage limits and renders them as a multi-line card. + * These types model the raw API response and the normalised shape consumed + * by the formatter. + */ + +/** A single raw limit entry from the GLM quota API (`data.limits[]`). */ +export interface RawLimit { + /** Limit family — "TOKENS_LIMIT" | "TIME_LIMIT" | "MCP_LIMIT" | . */ + type: string; + /** Window-count sub-identifier; meaning varies by backend (unused for classification). */ + unit?: number; + /** Window span; GLM uses 5 for the 5-hour window, 7 for the weekly window. */ + number?: number; + /** Total allowance for the window (absolute counter). */ + usage?: number; + /** Absolute remaining counter (preferred for percentage). */ + remaining?: number; + /** Absolute used counter (preferred for percentage alongside `usage`). */ + currentValue?: number; + /** Legacy percentage field — historically represents *used* percent. */ + percentage?: number; + /** Reset timestamp in epoch milliseconds. */ + nextResetTime?: number; + /** Per-model usage breakdown (present on MCP/TIME limits). */ + usageDetails?: { modelCode: string; usage: number }[]; +} + +/** + * Normalised quota item — one rendered line. + * + * Every raw limit that yields a valid percentage becomes a `QuotaItem`; + * limits with no computable percentage are silently dropped. + */ +export interface QuotaItem { + /** Stable key derived from type+number, e.g. "token_5h" / "mcp". */ + key: string; + /** Human-readable label for the line, e.g. "5h" / "MCP" / "Week". */ + label: string; + /** Used percentage, clamped to [0, 100]. */ + usedPercent: number; + /** Remaining percentage = 100 - usedPercent. */ + leftPercent: number; + /** Absolute used count (only some limits carry this, e.g. MCP/TIME). */ + usedCount?: number; + /** Absolute total allowance (only some limits carry this). */ + totalCount?: number; + /** Reset timestamp in epoch milliseconds, if present. */ + nextResetTime?: number; + /** Per-model usage breakdown, if present (rendered as indented sub-lines). */ + detail?: { modelCode: string; usage: number }[]; +} + +/** Top-level parsed result — a 4-state sum type. */ +export type QuotaResult = + | { kind: "success"; level: string; items: QuotaItem[] } + | { kind: "auth_error" } + | { kind: "rate_limited" } + | { kind: "unavailable" }; diff --git a/src/server.ts b/src/server.ts index 2c0ad99..2d178d8 100644 --- a/src/server.ts +++ b/src/server.ts @@ -49,6 +49,13 @@ export class ZcodeAcpServer { clientCapabilities: ClientCapabilities = {}; /** Session titles already set, to enforce set-once (acp_sid → title). */ readonly sessionTitles = new Map(); + /** + * Sessions eligible for auto-title on first end_turn. Only `session/new` + * populates this — resumed/loaded sessions already carry a title, so their + * first post-load message must NOT overwrite it. (sessionTitles alone can't + * distinguish "freshly created" from "resumed but not yet titled in-process".) + */ + readonly titleEligibleSessions = new Set(); /** Last mode id advertised to the client (acp_sid → modeId), for change detection. */ readonly lastMode = new Map(); /** Per-session ProjectionDiffers (persists across turns). */ diff --git a/src/utils.ts b/src/utils.ts index 65af528..c2d2151 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -58,6 +58,7 @@ export const SLASH_COMMANDS = [ description: "Set the reasoning effort", input: { hint: "max|high|nothink" }, }, + { name: "quota", description: "Show remaining usage quota (5h / weekly / MCP)" }, ] as const; /** Static metadata for the configOptions selects (model/mode/thought). */ diff --git a/tests/quota.test.ts b/tests/quota.test.ts new file mode 100644 index 0000000..67e99f4 --- /dev/null +++ b/tests/quota.test.ts @@ -0,0 +1,332 @@ +/** + * Tests for the quota feature: parsing (percentage fallbacks, label derivation, + * error states), formatting (progress-bar rendering, detail sub-lines), the + * in-memory TTL cache, and the queryQuota orchestration (cache hit / fetch + * failure → unavailable). + * + * Parser/formatter tests are pure-function. The orchestration test injects a + * fake fetch + fake clock to avoid real network calls. + */ + +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +// Mock the credentials loader so queryQuota orchestration tests don't depend +// on a real ~/.zcode/v2/config.json (absent in CI → "no apiKey" → unavailable). +// resolveQuotaHost is a pure function unaffected by this mock. +vi.mock("../src/backend/credentials.js", () => ({ + loadZcodeCredentials: () => ({ + ANTHROPIC_API_KEY: "test-key", + ZCODE_BASE_URL: "https://open.bigmodel.cn", + }), +})); + +import { clearCache, getCached, setCached, setClock } from "../src/quota/cache.js"; +import { formatQuota, renderBar } from "../src/quota/format.js"; +import { parseLimit, parseQuotaEnvelope } from "../src/quota/parse.js"; +import type { QuotaResult } from "../src/quota/types.js"; +import { resolveQuotaHost } from "../src/quota/client.js"; + +// Real-shape fixture mirroring the live API response (CN, Pro plan). +const PRO_FIXTURE = { + success: true, + code: 200, + msg: "操作成功", + data: { + level: "pro", + limits: [ + { + type: "TIME_LIMIT", + unit: 5, + number: 1, + usage: 1000, + currentValue: 237, + remaining: 763, + percentage: 24, + nextResetTime: 1784166659961, + usageDetails: [ + { modelCode: "search-prime", usage: 169 }, + { modelCode: "web-reader", usage: 68 }, + ], + }, + { + type: "TOKENS_LIMIT", + unit: 3, + number: 5, + percentage: 5, + nextResetTime: 1783436462284, + }, + ], + }, +} as const; + +describe("parseLimit", () => { + it("prefers remaining+currentValue over percentage (most precise)", () => { + const item = parseLimit({ + type: "TIME_LIMIT", + remaining: 763, + currentValue: 237, + percentage: 90, // misleading legacy field — must be ignored + }); + expect(item?.usedPercent).toBe(24); // 237 / (763+237) + expect(item?.leftPercent).toBe(76); + }); + + it("falls back to currentValue / usage when no remaining", () => { + const item = parseLimit({ + type: "TOKENS_LIMIT", + number: 5, + currentValue: 50, + usage: 200, + }); + expect(item?.usedPercent).toBe(25); + }); + + it("falls back to legacy percentage field (treated as used)", () => { + const item = parseLimit({ type: "TOKENS_LIMIT", number: 5, percentage: 5 }); + expect(item?.usedPercent).toBe(5); + expect(item?.leftPercent).toBe(95); + }); + + it("returns null when no percentage can be computed", () => { + expect(parseLimit({ type: "TOKENS_LIMIT", number: 5 })).toBeNull(); + expect(parseLimit({ type: "UNKNOWN" })).toBeNull(); + }); + + it("derives labels: 5h / Week / MCP / unknown", () => { + expect(parseLimit({ type: "TOKENS_LIMIT", number: 5, percentage: 10 })?.label).toBe("5h"); + expect(parseLimit({ type: "TOKENS_LIMIT", number: 7, percentage: 10 })?.label).toBe("Week"); + expect(parseLimit({ type: "TIME_LIMIT", percentage: 10 })?.label).toBe("MCP"); + expect(parseLimit({ type: "MCP_LIMIT", percentage: 10 })?.label).toBe("MCP"); + expect(parseLimit({ type: "FUTURE_LIMIT", percentage: 10 })?.label).toBe("FUTURE_LIMIT"); + }); + + it("carries nextResetTime and detail when present", () => { + const item = parseLimit({ + type: "TIME_LIMIT", + percentage: 24, + nextResetTime: 1784166659961, + usageDetails: [{ modelCode: "search-prime", usage: 169 }], + }); + expect(item?.nextResetTime).toBe(1784166659961); + expect(item?.detail).toEqual([{ modelCode: "search-prime", usage: 169 }]); + }); + + it("clamps out-of-range percentages to [0, 100]", () => { + expect(parseLimit({ type: "X", percentage: 150 })?.usedPercent).toBe(100); + expect(parseLimit({ type: "X", percentage: -5 })?.usedPercent).toBe(0); + }); +}); + +describe("parseQuotaEnvelope", () => { + it("parses the Pro fixture into 2 items with correct kinds", () => { + const result = parseQuotaEnvelope({ + status: 200, + json: PRO_FIXTURE, + text: JSON.stringify(PRO_FIXTURE), + }); + expect(result.kind).toBe("success"); + if (result.kind !== "success") return; + expect(result.level).toBe("pro"); + expect(result.items).toHaveLength(2); + // Both items must have valid percentages. + for (const item of result.items) { + expect(item.usedPercent).toBeGreaterThanOrEqual(0); + expect(item.usedPercent).toBeLessThanOrEqual(100); + } + }); + + it("HTTP 429 → rate_limited", () => { + const result = parseQuotaEnvelope({ status: 429, json: null, text: "" }); + expect(result.kind).toBe("rate_limited"); + }); + + it("business code 1001 / 401 → auth_error", () => { + expect( + parseQuotaEnvelope({ status: 200, json: { success: false, code: 1001 }, text: "" }).kind, + ).toBe("auth_error"); + expect( + parseQuotaEnvelope({ status: 200, json: { success: false, code: 401 }, text: "" }).kind, + ).toBe("auth_error"); + }); + + it("non-JSON body → unavailable", () => { + expect(parseQuotaEnvelope({ status: 200, json: null, text: "oops" }).kind).toBe("unavailable"); + }); + + it("success but zero parseable limits → unavailable", () => { + const result = parseQuotaEnvelope({ + status: 200, + json: { success: true, data: { level: "pro", limits: [{ type: "X" }] } }, + text: "", + }); + expect(result.kind).toBe("unavailable"); + }); +}); + +describe("renderBar", () => { + it("0% → all empty cells", () => { + expect(renderBar(0)).toBe("░".repeat(10)); + }); + + it("100% → all full cells", () => { + expect(renderBar(100)).toBe("█".repeat(10)); + }); + + it("50% → 5 full + 5 empty", () => { + expect(renderBar(50)).toBe("█".repeat(5) + "░".repeat(5)); + }); + + it("rounds to the nearest cell (each cell = 10%)", () => { + // 24% of 10 = 2.4 → rounds to 2 full. + expect(renderBar(24)).toBe("██" + "░".repeat(8)); + // 5% of 10 = 0.5 → rounds to 1 full. + expect(renderBar(5)).toBe("█" + "░".repeat(9)); + // 21% of 10 = 2.1 → rounds to 2 full. + expect(renderBar(21)).toBe("██" + "░".repeat(8)); + }); + + it("clamps inputs outside [0, 100]", () => { + expect(renderBar(150)).toBe("█".repeat(10)); + expect(renderBar(-10)).toBe("░".repeat(10)); + }); +}); + +describe("formatQuota", () => { + it("renders a success card with header, divider, bars, and detail", () => { + const result: QuotaResult = { + kind: "success", + level: "pro", + items: [ + { + key: "token_5h", + label: "5h", + usedPercent: 5, + leftPercent: 95, + nextResetTime: 1783436462284, + }, + { + key: "mcp", + label: "MCP", + usedPercent: 24, + leftPercent: 76, + usedCount: 237, + totalCount: 1000, + nextResetTime: 1784166659961, + detail: [ + { modelCode: "search-prime", usage: 169 }, + { modelCode: "web-reader", usage: 68 }, + ], + }, + ], + }; + const out = formatQuota(result); + const lines = out.split("\n"); + expect(lines[0]).toBe("GLM Coding Plan · Pro"); + expect(lines[1]).toMatch(/^─+$/); + // 5h line: percent + reset time, no "resets" word, no absolute counts. + expect(lines[2]).toContain("5h"); + expect(lines[2]).toContain("5%"); + expect(lines[2]).not.toContain("resets"); + expect(lines[2]).not.toMatch(/\(\d+\/\d+\)/); + // MCP line: percent + absolute counts + reset time. + expect(lines[3]).toContain("MCP"); + expect(lines[3]).toContain("24%"); + expect(lines[3]).toContain("(237/1000)"); + expect(lines[3]).not.toContain("resets"); + // Detail branches (now padded model codes). + expect(lines[4]).toMatch(/├ search-prime\s+\d+/); + expect(lines[5]).toMatch(/└ web-reader\s+\d+/); + }); + + it("omits absolute counts when the limit carries no counters (5h)", () => { + const out = formatQuota({ + kind: "success", + level: "pro", + items: [{ key: "token_5h", label: "5h", usedPercent: 18, leftPercent: 82 }], + }); + expect(out).not.toMatch(/\(\d+\/\d+\)/); + }); + + it("renders auth_error / rate_limited / unavailable fallbacks", () => { + expect(formatQuota({ kind: "auth_error" })).toMatch(/auth expired/i); + expect(formatQuota({ kind: "rate_limited" })).toMatch(/busy/i); + expect(formatQuota({ kind: "unavailable" })).toMatch(/unavailable/i); + }); +}); + +describe("cache", () => { + beforeEach(() => { + clearCache(); + setClock(() => 1000); + }); + afterEach(() => { + clearCache(); + setClock(undefined); + }); + + it("serves a cached result within the TTL window", () => { + const r: QuotaResult = { kind: "unavailable" }; + setCached(r); + setClock(() => 1000 + 9_999); // 9.999s later — still fresh + expect(getCached()).toBe(r); + }); + + it("returns null once the TTL expires", () => { + setCached({ kind: "unavailable" }); + setClock(() => 1000 + 10_001); // 10.001s later — expired + expect(getCached()).toBeNull(); + }); +}); + +describe("resolveQuotaHost", () => { + it("routes api.z.ai → intl, everything else → CN", () => { + expect(resolveQuotaHost("https://api.z.ai/api/anthropic")).toBe("https://api.z.ai"); + expect(resolveQuotaHost("https://open.bigmodel.cn/api/anthropic")).toBe( + "https://open.bigmodel.cn", + ); + expect(resolveQuotaHost("")).toBe("https://open.bigmodel.cn"); + }); +}); + +describe("queryQuota orchestration", () => { + // queryQuota imports client.ts which calls loadZcodeCredentials at call + // time, so we mock the module's fetchQuotaResponse via a spy on global fetch. + let fetchSpy: ReturnType; + + beforeEach(() => { + clearCache(); + setClock(() => 5000); + fetchSpy = vi + .spyOn(globalThis, "fetch") + .mockResolvedValue(new Response(JSON.stringify(PRO_FIXTURE), { status: 200 })); + }); + afterEach(() => { + fetchSpy.mockRestore(); + clearCache(); + setClock(undefined); + }); + + it("returns parsed success on a 200 response", async () => { + const { queryQuota } = await import("../src/quota/index.js"); + const result = await queryQuota(); + expect(result.kind).toBe("success"); + if (result.kind === "success") { + expect(result.items.length).toBeGreaterThanOrEqual(1); + } + }); + + it("degrades to unavailable on fetch throw", async () => { + fetchSpy.mockRejectedValue(new Error("network down")); + const { queryQuota } = await import("../src/quota/index.js"); + const result = await queryQuota(); + expect(result.kind).toBe("unavailable"); + }); + + it("serves cached result without re-fetching within TTL", async () => { + const { queryQuota } = await import("../src/quota/index.js"); + await queryQuota(); + expect(fetchSpy).toHaveBeenCalledTimes(1); + await queryQuota(); // cached — no new fetch + expect(fetchSpy).toHaveBeenCalledTimes(1); + }); +});