From 34aefb731b00335fcfbee76fa23834388b79c21c Mon Sep 17 00:00:00 2001 From: Bohdan Triapitsyn Date: Thu, 30 Jul 2026 00:05:13 +0300 Subject: [PATCH] fix(quota): show all Z.ai usage windows --- packages/vscode/src/quotaProviders.test.ts | 28 ++++++++++ packages/vscode/src/quotaProviders.ts | 27 +++++++--- .../web/server/lib/quota/DOCUMENTATION.md | 1 + .../web/server/lib/quota/providers/zai.js | 24 +++++---- .../server/lib/quota/providers/zai.test.js | 54 +++++++++++++++++++ .../server/lib/quota/utils/transformers.js | 6 ++- 6 files changed, 122 insertions(+), 18 deletions(-) create mode 100644 packages/web/server/lib/quota/providers/zai.test.js diff --git a/packages/vscode/src/quotaProviders.test.ts b/packages/vscode/src/quotaProviders.test.ts index 00af15d3..9f7b466b 100644 --- a/packages/vscode/src/quotaProviders.test.ts +++ b/packages/vscode/src/quotaProviders.test.ts @@ -10,6 +10,7 @@ const AUTH = JSON.stringify({ openai: { access: 'test-token' }, crof: { key: 'test-token' }, neuralwatt: { key: 'test-token' }, + 'zai-coding-plan': { key: 'test-token' }, }); ((fs as unknown) as { existsSync: () => boolean }).existsSync = () => true; ((fs as unknown) as { readFileSync: () => string }).readFileSync = () => AUTH; @@ -138,6 +139,33 @@ describe('Codex quota provider (VS Code parity)', () => { }); }); +describe('Z.ai quota provider (VS Code parity)', () => { + test('surfaces 5-hour, weekly, and MCP quota windows', async () => { + stubFetchReturning(() => Promise.resolve(mockResponse({ + data: { + limits: [ + { type: 'TOKENS_LIMIT', unit: 3, number: 5, percentage: 0 }, + { type: 'TOKENS_LIMIT', unit: 6, number: 1, percentage: 100, nextResetTime: 1785659659993 }, + { type: 'TIME_LIMIT', unit: 5, number: 1, percentage: 0, nextResetTime: 1787128459979 }, + ], + }, + }))); + + const result = await fetchQuotaForProvider('zai-coding-plan'); + const windows = result.usage!.windows; + + assert.equal(result.ok, true); + assert.equal(windows['5h']!.usedPercent, 0); + assert.equal(windows['5h']!.windowSeconds, 5 * 60 * 60); + assert.equal(windows.weekly!.usedPercent, 100); + assert.equal(windows.weekly!.windowSeconds, 7 * 24 * 60 * 60); + assert.equal(windows.weekly!.resetAt, 1785659659993); + assert.equal(windows['MCP Tools']!.usedPercent, 0); + assert.equal(windows['MCP Tools']!.windowSeconds, 30 * 24 * 60 * 60); + assert.equal(windows['MCP Tools']!.resetAt, 1787128459979); + }); +}); + describe('NeuralWatt quota provider (VS Code parity)', () => { test('builds subscription window keyed by plan name (windowSeconds null)', async () => { stubFetchReturning(() => Promise.resolve(mockResponse(DOCUMENTED_SUBSCRIPTION_PAYLOAD))); diff --git a/packages/vscode/src/quotaProviders.ts b/packages/vscode/src/quotaProviders.ts index 759cb03e..807d8ace 100644 --- a/packages/vscode/src/quotaProviders.ts +++ b/packages/vscode/src/quotaProviders.ts @@ -222,7 +222,10 @@ const resolveGoogleWindow = (sourceId: GoogleAuthSource['sourceId'], resetAt: nu return { label: 'daily', seconds: GOOGLE_DAILY_WINDOW_SECONDS } as const; }; -const ZAI_TOKEN_WINDOW_SECONDS: Record = { 3: 3600 }; +const ZAI_TOKEN_WINDOW_SECONDS: Record = { + 3: 60 * 60, + 6: 7 * 24 * 60 * 60, +}; const readAuthFile = (): AuthFile => { if (!fs.existsSync(AUTH_FILE)) { @@ -1592,14 +1595,13 @@ const fetchZaiQuota = async (): Promise => { const payload = await response.json() as ZaiPayload; const limits = Array.isArray(payload?.data?.limits) ? payload.data.limits : []; - const tokensLimit = limits.find((limit: Record) => limit?.type === 'TOKENS_LIMIT'); - const windowSeconds = resolveWindowSeconds(tokensLimit as Record | undefined); - const windowLabel = resolveWindowLabel(windowSeconds); - const resetAt = tokensLimit?.nextResetTime ? normalizeTimestamp(tokensLimit.nextResetTime) : null; - const usedPercent = typeof tokensLimit?.percentage === 'number' ? tokensLimit.percentage : null; - const windows: Record = {}; - if (tokensLimit) { + for (const tokensLimit of limits.filter((limit) => limit?.type === 'TOKENS_LIMIT')) { + const windowSeconds = resolveWindowSeconds(tokensLimit as Record); + const windowLabel = resolveWindowLabel(windowSeconds); + const resetAt = tokensLimit.nextResetTime ? normalizeTimestamp(tokensLimit.nextResetTime) : null; + const usedPercent = typeof tokensLimit.percentage === 'number' ? tokensLimit.percentage : null; + windows[windowLabel] = toUsageWindow({ usedPercent, windowSeconds, @@ -1607,6 +1609,15 @@ const fetchZaiQuota = async (): Promise => { }); } + const mcpToolsTimeLimit = limits.find((limit) => limit?.type === 'TIME_LIMIT'); + if (mcpToolsTimeLimit) { + windows['MCP Tools'] = toUsageWindow({ + usedPercent: typeof mcpToolsTimeLimit.percentage === 'number' ? mcpToolsTimeLimit.percentage : null, + windowSeconds: 30 * 24 * 60 * 60, + resetAt: mcpToolsTimeLimit.nextResetTime ? normalizeTimestamp(mcpToolsTimeLimit.nextResetTime) : null, + }); + } + return buildResult({ providerId: 'zai-coding-plan', providerName: 'z.ai', diff --git a/packages/web/server/lib/quota/DOCUMENTATION.md b/packages/web/server/lib/quota/DOCUMENTATION.md index 8769d9e1..475ff2cf 100644 --- a/packages/web/server/lib/quota/DOCUMENTATION.md +++ b/packages/web/server/lib/quota/DOCUMENTATION.md @@ -74,3 +74,4 @@ In 2025/2026 MiniMax rebranded "Coding Plan" to "Token Plan" alongside the M3 mo - Keep provider IDs stable; clients use them directly. - Avoid adding alias-based dispatch in `fetchQuotaForProvider`; dispatch currently expects exact provider IDs. - Keep Google behavior changes isolated and review `providers/google/*` together. +- Z.ai Coding Plan exposes separate 5-hour and weekly `TOKENS_LIMIT` entries plus a monthly `TIME_LIMIT` for MCP tools; web and VS Code must preserve all three windows. diff --git a/packages/web/server/lib/quota/providers/zai.js b/packages/web/server/lib/quota/providers/zai.js index 980c51bd..0c526f99 100644 --- a/packages/web/server/lib/quota/providers/zai.js +++ b/packages/web/server/lib/quota/providers/zai.js @@ -4,8 +4,6 @@ import { normalizeAuthEntry, buildResult, toUsageWindow, - toNumber, - toTimestamp, resolveWindowSeconds, resolveWindowLabel, normalizeTimestamp @@ -57,14 +55,13 @@ export const fetchQuota = async () => { const payload = await response.json(); const limits = Array.isArray(payload?.data?.limits) ? payload.data.limits : []; - const tokensLimit = limits.find((limit) => limit?.type === 'TOKENS_LIMIT'); - const windowSeconds = resolveWindowSeconds(tokensLimit); - const windowLabel = resolveWindowLabel(windowSeconds); - const resetAt = tokensLimit?.nextResetTime ? normalizeTimestamp(tokensLimit.nextResetTime) : null; - const usedPercent = typeof tokensLimit?.percentage === 'number' ? tokensLimit.percentage : null; - const windows = {}; - if (tokensLimit) { + for (const tokensLimit of limits.filter((limit) => limit?.type === 'TOKENS_LIMIT')) { + const windowSeconds = resolveWindowSeconds(tokensLimit); + const windowLabel = resolveWindowLabel(windowSeconds); + const resetAt = tokensLimit?.nextResetTime ? normalizeTimestamp(tokensLimit.nextResetTime) : null; + const usedPercent = typeof tokensLimit?.percentage === 'number' ? tokensLimit.percentage : null; + windows[windowLabel] = toUsageWindow({ usedPercent, windowSeconds, @@ -72,6 +69,15 @@ export const fetchQuota = async () => { }); } + const mcpToolsTimeLimit = limits.find((limit) => limit?.type === 'TIME_LIMIT'); + if (mcpToolsTimeLimit) { + windows['MCP Tools'] = toUsageWindow({ + usedPercent: typeof mcpToolsTimeLimit.percentage === 'number' ? mcpToolsTimeLimit.percentage : null, + windowSeconds: 30 * 24 * 60 * 60, + resetAt: mcpToolsTimeLimit.nextResetTime ? normalizeTimestamp(mcpToolsTimeLimit.nextResetTime) : null + }); + } + return buildResult({ providerId, providerName, diff --git a/packages/web/server/lib/quota/providers/zai.test.js b/packages/web/server/lib/quota/providers/zai.test.js new file mode 100644 index 00000000..39d5dcf5 --- /dev/null +++ b/packages/web/server/lib/quota/providers/zai.test.js @@ -0,0 +1,54 @@ +import { afterEach, describe, expect, it, vi } from 'vitest'; + +vi.mock('../../opencode/auth.js', () => ({ + readAuthFile: () => ({ 'zai-coding-plan': { key: 'test-token' } }), +})); + +import { fetchQuota } from './zai.js'; + +afterEach(() => { + vi.unstubAllGlobals(); +}); + +const mockResponse = (body) => ({ + ok: true, + status: 200, + json: async () => body, +}); + +describe('Z.ai quota provider', () => { + it('surfaces 5-hour, weekly, and MCP quota windows', async () => { + vi.stubGlobal('fetch', vi.fn().mockResolvedValue(mockResponse({ + data: { + limits: [ + { type: 'TOKENS_LIMIT', unit: 3, number: 5, percentage: 0 }, + { type: 'TOKENS_LIMIT', unit: 6, number: 1, percentage: 100, nextResetTime: 1785659659993 }, + { type: 'TIME_LIMIT', unit: 5, number: 1, percentage: 0, nextResetTime: 1787128459979 }, + ], + }, + }))); + + const result = await fetchQuota(); + const windows = result.usage.windows; + + expect(result.ok).toBe(true); + expect(windows['5h']).toMatchObject({ + usedPercent: 0, + remainingPercent: 100, + windowSeconds: 5 * 60 * 60, + resetAt: null, + }); + expect(windows.weekly).toMatchObject({ + usedPercent: 100, + remainingPercent: 0, + windowSeconds: 7 * 24 * 60 * 60, + resetAt: 1785659659993, + }); + expect(windows['MCP Tools']).toMatchObject({ + usedPercent: 0, + remainingPercent: 100, + windowSeconds: 30 * 24 * 60 * 60, + resetAt: 1787128459979, + }); + }); +}); diff --git a/packages/web/server/lib/quota/utils/transformers.js b/packages/web/server/lib/quota/utils/transformers.js index 8c9ac101..5dcc0011 100644 --- a/packages/web/server/lib/quota/utils/transformers.js +++ b/packages/web/server/lib/quota/utils/transformers.js @@ -34,8 +34,12 @@ export const normalizeTimestamp = (value) => { return value < 1_000_000_000_000 ? value * 1000 : value; }; +const ZAI_TOKEN_WINDOW_SECONDS = { + 3: 60 * 60, + 6: 7 * 24 * 60 * 60 +}; + export const resolveWindowSeconds = (limit) => { - const ZAI_TOKEN_WINDOW_SECONDS = { 3: 3600 }; if (!limit || !limit.number) return null; const unitSeconds = ZAI_TOKEN_WINDOW_SECONDS[limit.unit]; if (!unitSeconds) return null;