pi-provider-quota 0.1.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/README.md ADDED
@@ -0,0 +1,92 @@
1
+ # pi-provider-quota
2
+
3
+ Track live quota usage for **Z.Ai**, **Kimi Code**, **Ollama Cloud**, and **DeepSeek** in [pi](https://pi.dev)'s status bar.
4
+
5
+ ## What it does
6
+
7
+ Polls provider APIs and Ollama's settings page every 30 seconds, then displays a color-coded quota bar:
8
+
9
+ ```
10
+ quota: zai pro · 5h: 🟡 62% · wk: 🟢 34%
11
+ quota: ollama pro · 5h: 🟢 18% · reset 2h 30m · wk: 🟢 5% · wk reset 3d 4h
12
+ quota: ds · ¥12.50 · ✅ Yes
13
+ ```
14
+
15
+ Color thresholds: 🟢 < 50%, 🟡 50–79%, 🔴 ≥ 80%.
16
+
17
+ ## Installation
18
+
19
+ ```bash
20
+ pi install git:github.com/ronnieops/pi-provider-quota
21
+ ```
22
+
23
+ ## Setup
24
+
25
+ Store your credentials in `~/.pi/agent/auth.json`:
26
+
27
+ ```json
28
+ {
29
+ "zai": { "key": "sk-..." },
30
+ "kimi-coding": { "key": "sk-..." },
31
+ "deepseek": { "key": "sk-..." },
32
+ "ollama-session": { "key": "<__Secure-session cookie value>", "aid": "<aid cookie value>" }
33
+ }
34
+ ```
35
+
36
+ Environment variable fallbacks are also supported (`ZAI_API_KEY`, `KIMI_CODING_API_KEY`, `DEEPSEEK_API_KEY`).
37
+
38
+ Select a tracked model in pi. Quota status appears in the status bar automatically.
39
+
40
+ ## Providers
41
+
42
+ | Provider | Live quota source | Fields displayed |
43
+ |----------|-------------------|------------------|
44
+ | **Z.Ai** | Dashboard API | 5h %, weekly %, reset times, plan tier, MCP model breakdown |
45
+ | **Kimi Code** | Billing API (web JWT) | 5h remaining, weekly remaining, weekly reset |
46
+ | **Ollama Cloud** | Settings page (session cookie) | Session %, weekly %, session reset, weekly reset, plan tier |
47
+ | **DeepSeek** | Balance API | Total balance, availability flag |
48
+
49
+ Ollama falls back to local GPU-time weighted turn counting when no session cookie is configured.
50
+
51
+ ## Commands
52
+
53
+ | Command | Description |
54
+ |---------|-------------|
55
+ | `/quota` | Detailed breakdown: live percentages, turn counts, reset times, cookie status, raw headers |
56
+ | `/quota-plan <tier>` | Override detected plan tier for the active provider |
57
+ | `/quota-json` | JSON snapshot of all quota state (for scripting) |
58
+
59
+ ## LLM Tool
60
+
61
+ The `provider_quota` tool is registered automatically. When an LLM calls it, it returns a JSON object with live percentages, plan tier, turn counts, reset times, and any fetch errors for the active provider. This lets the agent reason about remaining quota without manual `/quota` calls.
62
+
63
+ ## Security
64
+
65
+ - API keys and session cookies are stored in `~/.pi/agent/auth.json` only — never sent to any third party besides the provider's own API
66
+ - Credentials are used in HTTP headers only, never stored in state, never logged, never displayed
67
+ - HTML scraping uses regex captures that are numeric-only or fixed-alternation — no scraped text reaches output
68
+ - Error messages are sanitized — no response bodies, headers, or credential fragments in status output
69
+ - Session persistence strips credentials and raw headers before saving
70
+
71
+ ## Architecture
72
+
73
+ Single-file extension (`extensions/provider-quota.ts`, ~1100 lines) with:
74
+
75
+ - **Discriminated-union result types** (`OllamaQuotaResult`, `OllamaCookieResult`) — no `unknown` returns
76
+ - **Stale-data clearing** — all live fields nulled on fetch failure; plan tier preserved across transient errors
77
+ - **Session persistence** — quota state survives pi restarts (sanitized; credentials and raw headers stripped)
78
+
79
+ ## Development
80
+
81
+ ```bash
82
+ npm install
83
+ npx tsc --noEmit --target esnext --module esnext --moduleResolution bundler --strict --skipLibCheck --esModuleInterop extensions/provider-quota.ts
84
+ ```
85
+
86
+ ## License
87
+
88
+ MIT
89
+
90
+ ---
91
+
92
+ <p align="center">Proudly created with <a href="https://pi.dev">pi</a></p>
@@ -0,0 +1,1144 @@
1
+ /**
2
+ * pi-provider-quota
3
+ * Track Z.Ai GLM Coding Plan, Kimi Code, Ollama Cloud, and DeepSeek quota in pi's status bar.
4
+ */
5
+ import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
6
+ import { Type } from "typebox";
7
+
8
+ // =============================================================================
9
+ // Types
10
+ // =============================================================================
11
+
12
+ interface QuotaState {
13
+ provider: string;
14
+ model: string;
15
+ plan: string;
16
+
17
+ // Turn counting (all providers)
18
+ turnsUsed5h: number;
19
+ quotaUsed5h: number;
20
+ cycleStart5h: number | null;
21
+ turnsUsedWeekly: number;
22
+ weeklyStart: number | null;
23
+
24
+ // Live quota (Z.Ai, Kimi, DeepSeek - refreshed every 30s)
25
+ live5hPercentage: number | null;
26
+ liveWeeklyPercentage: number | null;
27
+ livePlan: string | null;
28
+ live5hResetAt: number | null;
29
+ liveWeeklyResetAt: number | null;
30
+ lastSuccessfulFetchAt: number | null;
31
+
32
+ // Kimi specific
33
+ liveKimi5hRemaining: number | null;
34
+ liveKimiWeeklyRemaining: number | null;
35
+ liveKimiWeeklyResetAt: number | null;
36
+ kimiJwtExpired: boolean;
37
+
38
+ // DeepSeek specific
39
+ liveDsAvailable: boolean | null;
40
+ liveDsTotalBalance: string | null;
41
+ liveDsCurrency: string | null;
42
+
43
+ // Ollama specific (GPU-time weighted + live quota)
44
+ ollamaGpuUnitsUsed: number;
45
+ liveOllama5hPercentage: number | null;
46
+ liveOllamaWeeklyPercentage: number | null;
47
+ liveOllamaResetAt: number | null;
48
+ liveOllamaWeeklyResetAt: number | null;
49
+
50
+ // MCP (Z.Ai)
51
+ liveMcpDetails: Array<{ modelCode: string; usage: number }>;
52
+
53
+ // Header-based quota tracking
54
+ headerPromptsLimit: number | null;
55
+ rawHeaders: Record<string, string>;
56
+
57
+ // Error handling
58
+ lastFetchError: string | null;
59
+ lastResponseAt: number | null;
60
+ }
61
+
62
+ interface TierData {
63
+ limit5h: number;
64
+ limitWeekly: number;
65
+ }
66
+
67
+ // =============================================================================
68
+ // Config
69
+ // =============================================================================
70
+
71
+ const ZAI_TIERS: Record<string, TierData> = {
72
+ lite: { limit5h: 80, limitWeekly: 400 },
73
+ pro: { limit5h: 400, limitWeekly: 2000 },
74
+ max: { limit5h: 1600, limitWeekly: 8000 },
75
+ };
76
+
77
+ const KIMI_TIERS: Record<string, TierData> = {
78
+ moderato: { limit5h: 100, limitWeekly: 100 },
79
+ allegretto: { limit5h: 100, limitWeekly: 600 },
80
+ vivace: { limit5h: 100, limitWeekly: 1200 },
81
+ };
82
+
83
+ const OLLAMA_TIERS: Record<string, TierData> = {
84
+ free: { limit5h: 500, limitWeekly: 3000 },
85
+ pro: { limit5h: 2400, limitWeekly: 13700 },
86
+ max: { limit5h: 12000, limitWeekly: 68500 },
87
+ };
88
+
89
+ // =============================================================================
90
+ // State
91
+ // =============================================================================
92
+
93
+ let state: QuotaState = createFreshState();
94
+
95
+ function createFreshState(): QuotaState {
96
+ return {
97
+ provider: "",
98
+ model: "",
99
+ plan: "",
100
+ turnsUsed5h: 0,
101
+ quotaUsed5h: 0,
102
+ cycleStart5h: null,
103
+ turnsUsedWeekly: 0,
104
+ weeklyStart: null,
105
+ live5hPercentage: null,
106
+ liveWeeklyPercentage: null,
107
+ livePlan: null,
108
+ live5hResetAt: null,
109
+ liveWeeklyResetAt: null,
110
+ lastSuccessfulFetchAt: null,
111
+ liveKimi5hRemaining: null,
112
+ liveKimiWeeklyRemaining: null,
113
+ liveKimiWeeklyResetAt: null,
114
+ kimiJwtExpired: false,
115
+ liveDsAvailable: null,
116
+ liveDsTotalBalance: null,
117
+ liveDsCurrency: null,
118
+ ollamaGpuUnitsUsed: 0,
119
+ liveOllama5hPercentage: null,
120
+ liveOllamaWeeklyPercentage: null,
121
+ liveOllamaResetAt: null,
122
+ liveOllamaWeeklyResetAt: null,
123
+ liveMcpDetails: [],
124
+ headerPromptsLimit: null,
125
+ rawHeaders: {},
126
+ lastFetchError: null,
127
+ lastResponseAt: null,
128
+ };
129
+ }
130
+
131
+ const ALLOWED_KEYS = new Set(Object.keys(createFreshState()));
132
+
133
+ function resetProviderData(): void {
134
+ state.live5hPercentage = null;
135
+ state.liveWeeklyPercentage = null;
136
+ state.livePlan = null;
137
+ state.live5hResetAt = null;
138
+ state.liveWeeklyResetAt = null;
139
+ state.lastSuccessfulFetchAt = null;
140
+ state.liveKimi5hRemaining = null;
141
+ state.liveKimiWeeklyRemaining = null;
142
+ state.liveKimiWeeklyResetAt = null;
143
+ state.kimiJwtExpired = false;
144
+ state.liveDsAvailable = null;
145
+ state.liveDsTotalBalance = null;
146
+ state.liveDsCurrency = null;
147
+ state.ollamaGpuUnitsUsed = 0;
148
+ state.liveOllama5hPercentage = null;
149
+ state.liveOllamaWeeklyPercentage = null;
150
+ state.liveOllamaResetAt = null;
151
+ state.liveOllamaWeeklyResetAt = null;
152
+ state.liveMcpDetails = [];
153
+ state.headerPromptsLimit = null;
154
+ state.rawHeaders = {};
155
+ state.lastFetchError = null;
156
+ }
157
+
158
+ function sanitizeState(data: unknown): Partial<QuotaState> {
159
+ if (!data || typeof data !== "object") return {};
160
+ return Object.fromEntries(
161
+ Object.entries(data as Record<string, unknown>).filter(
162
+ ([k]) => ALLOWED_KEYS.has(k),
163
+ ),
164
+ ) as Partial<QuotaState>;
165
+ }
166
+
167
+ // =============================================================================
168
+ // Helpers
169
+ // =============================================================================
170
+
171
+ function isZai(name: string): boolean {
172
+ const normalized = name.toLowerCase().replace(/[^a-z0-9._-]/g, "");
173
+ return ["zai", "glm"].some(p => normalized.includes(p));
174
+ }
175
+
176
+ function isKimi(name: string): boolean {
177
+ const normalized = name.toLowerCase().replace(/[^a-z0-9._-]/g, "");
178
+ return ["kimi", "k2p6", "kimi-for-coding", "moonshot"].some(p => normalized.includes(p));
179
+ }
180
+
181
+ function isOllama(name: string): boolean {
182
+ const normalized = name.toLowerCase().replace(/[^a-z0-9._-]/g, "");
183
+ return ["ollama", "ollama-cloud"].some(p => normalized.includes(p));
184
+ }
185
+
186
+ function isDeepseek(name: string): boolean {
187
+ const normalized = name.toLowerCase().replace(/[^a-z0-9._-]/g, "");
188
+ return ["deepseek"].some(p => normalized.includes(p));
189
+ }
190
+
191
+ function isTrackedProvider(name: string): boolean {
192
+ return isZai(name) || isKimi(name) || isOllama(name) || isDeepseek(name);
193
+ }
194
+
195
+ function getLimit(): number | null {
196
+ if (isDeepseek(state.provider)) return null;
197
+ if (state.headerPromptsLimit && !isOllama(state.provider)) {
198
+ return state.headerPromptsLimit;
199
+ }
200
+ const tiers = getProviderTiers();
201
+ const tier = tiers[state.plan];
202
+ return tier ? tier.limit5h : null;
203
+ }
204
+
205
+ function getWeeklyLimit(): number | null {
206
+ if (isDeepseek(state.provider)) return null;
207
+ const tiers = getProviderTiers();
208
+ const tier = tiers[state.plan];
209
+ return tier ? tier.limitWeekly : null;
210
+ }
211
+
212
+ function getProviderTiers(): Record<string, TierData> {
213
+ if (isZai(state.provider)) return ZAI_TIERS;
214
+ if (isKimi(state.provider)) return KIMI_TIERS;
215
+ if (isOllama(state.provider)) return OLLAMA_TIERS;
216
+ return {};
217
+ }
218
+
219
+ function ollamaGpuMultiplier(modelId: string): number {
220
+ const lower = modelId.toLowerCase();
221
+ if (lower.includes("405b")) return 50;
222
+ if (lower.includes("70b") || lower.includes("72b")) return 25;
223
+ if (lower.includes("34b") || lower.includes("30b")) return 10;
224
+ if (lower.includes("14b")) return 4;
225
+ if (lower.includes("8b") || lower.includes("7b")) return 2;
226
+ return 1;
227
+ }
228
+
229
+ function quotaMultiplier(modelId: string): number {
230
+ if (isOllama(state.provider)) {
231
+ return ollamaGpuMultiplier(modelId);
232
+ }
233
+ if (isZai(state.provider) && isPeakHours()) {
234
+ return 3;
235
+ }
236
+ return 1;
237
+ }
238
+
239
+ function isPeakHours(): boolean {
240
+ const now = new Date();
241
+ const utc8Hour = (now.getUTCHours() + 8) % 24;
242
+ return utc8Hour >= 14 && utc8Hour < 18;
243
+ }
244
+
245
+ function isCycleExpired(): boolean {
246
+ if (!state.cycleStart5h) return false;
247
+ return Date.now() - state.cycleStart5h >= 5 * 60 * 60 * 1000;
248
+ }
249
+
250
+ function maybeRotateCycle(): boolean {
251
+ if (!isCycleExpired()) return false;
252
+ state.turnsUsed5h = 0;
253
+ state.quotaUsed5h = 0;
254
+ state.cycleStart5h = Date.now();
255
+ state.headerPromptsLimit = null;
256
+ return true;
257
+ }
258
+
259
+ function maybeRotateWeekly(): void {
260
+ if (!state.weeklyStart) {
261
+ state.weeklyStart = Date.now();
262
+ return;
263
+ }
264
+ const weekMs = 7 * 24 * 60 * 60 * 1000;
265
+ if (Date.now() - state.weeklyStart >= weekMs) {
266
+ state.turnsUsedWeekly = 0;
267
+ state.weeklyStart = Date.now();
268
+ }
269
+ }
270
+
271
+ function relativeTime(targetMs: number | null): string {
272
+ if (!targetMs) return "—";
273
+ const diff = targetMs - Date.now();
274
+ if (diff <= 0) return "0m";
275
+ const hours = Math.floor(diff / (60 * 60 * 1000));
276
+ const mins = Math.floor((diff % (60 * 60 * 1000)) / (60 * 1000));
277
+ const days = Math.floor(hours / 24);
278
+ if (days > 0) return `${days}d ${hours % 24}h`;
279
+ if (hours > 0) return `${hours}h ${mins}m`;
280
+ return `${mins}m`;
281
+ }
282
+
283
+ function readApiKey(provider: string): string | null {
284
+ try {
285
+ const home = process.env.HOME || process.env.USERPROFILE || "";
286
+ const authPath = `${home}/.pi/agent/auth.json`;
287
+ const { existsSync, readFileSync } = require("node:fs");
288
+ if (!existsSync(authPath)) {
289
+ const envKey = process.env[`${provider.toUpperCase().replace(/[- ]/g, "_")}_API_KEY`];
290
+ return envKey || null;
291
+ }
292
+ const raw = readFileSync(authPath, "utf8");
293
+ const json = JSON.parse(raw);
294
+ const aliases: Record<string, string[]> = {
295
+ zai: ["zai", "zai-api", "glm"],
296
+ "kimi-coding": ["kimi-coding", "kimi", "moonshot"],
297
+ deepseek: ["deepseek", "deepseek-api"],
298
+ ollama: ["ollama", "ollama-cloud"],
299
+ };
300
+ const keys = aliases[provider] || [provider];
301
+ for (const key of keys) {
302
+ if (json[key]?.key) return json[key].key;
303
+ if (json[key]?.apiKey) return json[key].apiKey;
304
+ }
305
+ const envName = provider.toUpperCase().replace(/[- ]/g, "_") + "_API_KEY";
306
+ return process.env[envName] || null;
307
+ } catch {
308
+ return null;
309
+ }
310
+ }
311
+
312
+ function readKimiWebJwt(): string | null {
313
+ try {
314
+ const home = process.env.HOME || process.env.USERPROFILE || "";
315
+ const settingsPath = `${home}/.pi/agent/settings.json`;
316
+ const authPath = `${home}/.pi/agent/auth.json`;
317
+ const { existsSync, readFileSync } = require("node:fs");
318
+
319
+ // Check settings.json first (deprecated location, but supported)
320
+ if (existsSync(settingsPath)) {
321
+ const raw = readFileSync(settingsPath, "utf8");
322
+ const json = JSON.parse(raw);
323
+ if (typeof json.kimiWebJwt === "string" && json.kimiWebJwt.length > 0) {
324
+ return json.kimiWebJwt;
325
+ }
326
+ }
327
+
328
+ // Check auth.json
329
+ if (existsSync(authPath)) {
330
+ const raw = readFileSync(authPath, "utf8");
331
+ const json = JSON.parse(raw);
332
+ if (json["kimi-web"]?.key) return json["kimi-web"].key;
333
+ if (json["kimi-web-jwt"]?.key) return json["kimi-web-jwt"].key;
334
+ }
335
+
336
+ return null;
337
+ } catch {
338
+ return null;
339
+ }
340
+ }
341
+
342
+ function base64UrlDecode(str: string): string {
343
+ let s = str.replace(/-/g, "+").replace(/_/g, "/");
344
+ while (s.length % 4) s += "=";
345
+ return Buffer.from(s, "base64").toString("utf8");
346
+ }
347
+
348
+ function getJwtExpiry(jwt: string): number | null {
349
+ try {
350
+ const parts = jwt.split(".");
351
+ if (parts.length !== 3) return null;
352
+ const payload = JSON.parse(base64UrlDecode(parts[1]));
353
+ const exp = payload.exp;
354
+ if (typeof exp === "number" && exp > 0) return exp;
355
+ return null;
356
+ } catch {
357
+ return null;
358
+ }
359
+ }
360
+
361
+ function dsCurrencyPrefix(currency: string | null): string {
362
+ if (currency === "CNY") return "¥";
363
+ return "$";
364
+ }
365
+
366
+ // =============================================================================
367
+ // API Calls
368
+ // =============================================================================
369
+
370
+ async function fetchLiveQuota(apiKey: string): Promise<unknown> {
371
+ try {
372
+ const res = await fetch("https://z.ai/api/monitor/usage/quota/limit", {
373
+ headers: {
374
+ Authorization: `Bearer ${apiKey}`,
375
+ "Content-Type": "application/json",
376
+ },
377
+ signal: AbortSignal.timeout(8000),
378
+ });
379
+ if (!res.ok) return null;
380
+ return await res.json();
381
+ } catch {
382
+ return null;
383
+ }
384
+ }
385
+
386
+ function applyLiveData(data: unknown, apiState: QuotaState): boolean {
387
+ if (!data || typeof data !== "object") return false;
388
+ const d = data as Record<string, unknown>;
389
+ if (typeof d.level === "string") {
390
+ apiState.livePlan = d.level;
391
+ }
392
+ if (Array.isArray(d.limits)) {
393
+ let found5h = false;
394
+ let foundWeekly = false;
395
+ for (const limit of d.limits) {
396
+ if (!limit || typeof limit !== "object") continue;
397
+ const l = limit as Record<string, unknown>;
398
+ if (l.type === "TOKENS_LIMIT" && l.unit === 3 && !found5h) {
399
+ if (typeof l.percentage === "number") apiState.live5hPercentage = l.percentage;
400
+ if (typeof l.nextResetTime === "number") apiState.live5hResetAt = l.nextResetTime;
401
+ found5h = true;
402
+ }
403
+ if (l.type === "TOKENS_LIMIT" && l.unit === 6 && !foundWeekly) {
404
+ if (typeof l.percentage === "number") apiState.liveWeeklyPercentage = l.percentage;
405
+ if (typeof l.nextResetTime === "number") apiState.liveWeeklyResetAt = l.nextResetTime;
406
+ foundWeekly = true;
407
+ }
408
+ if (l.type === "TIME_LIMIT" && l.unit === 5) {
409
+ if (Array.isArray(l.usageDetails)) {
410
+ apiState.liveMcpDetails = l.usageDetails.map((item: unknown) => {
411
+ const i = item as Record<string, unknown>;
412
+ return {
413
+ modelCode: String(i.modelCode ?? ""),
414
+ usage: typeof i.usage === "number" ? i.usage : 0,
415
+ };
416
+ });
417
+ }
418
+ }
419
+ }
420
+ }
421
+ return true;
422
+ }
423
+
424
+ async function fetchKimiLiveQuota(webJwt: string): Promise<unknown> {
425
+ try {
426
+ const res = await fetch("https://www.kimi.com/apiv2/kimi.gateway.billing.v1.BillingService/GetUsages", {
427
+ method: "POST",
428
+ headers: {
429
+ Authorization: `Bearer ${webJwt}`,
430
+ "connect-protocol-version": "1",
431
+ "Content-Type": "application/json",
432
+ },
433
+ body: JSON.stringify({ scope: ["FEATURE_CODING"] }),
434
+ signal: AbortSignal.timeout(8000),
435
+ });
436
+ if (!res.ok) return null;
437
+ return await res.json();
438
+ } catch {
439
+ return null;
440
+ }
441
+ }
442
+
443
+ function applyKimiLiveData(data: unknown, apiState: QuotaState): boolean {
444
+ if (!data || typeof data !== "object") return false;
445
+ const d = data as Record<string, unknown>;
446
+
447
+ if (Array.isArray(d.usages)) {
448
+ const codingUsage = d.usages.find((u: unknown) => {
449
+ const usage = u as Record<string, unknown>;
450
+ return usage.scope === "FEATURE_CODING";
451
+ });
452
+
453
+ if (codingUsage && typeof codingUsage === "object") {
454
+ const cu = codingUsage as Record<string, unknown>;
455
+
456
+ // Weekly data
457
+ if (cu.detail && typeof cu.detail === "object") {
458
+ const detail = cu.detail as Record<string, unknown>;
459
+ const remaining = parseInt(String(detail.remaining ?? ""), 10);
460
+ if (!Number.isNaN(remaining)) {
461
+ apiState.liveKimiWeeklyRemaining = remaining;
462
+ }
463
+ const resetTime = detail.resetTime;
464
+ if (typeof resetTime === "string" && resetTime) {
465
+ const ts = new Date(resetTime).getTime();
466
+ if (!Number.isNaN(ts)) {
467
+ apiState.liveKimiWeeklyResetAt = ts;
468
+ }
469
+ }
470
+ }
471
+
472
+ // 5h rate limit
473
+ if (Array.isArray(cu.limits)) {
474
+ const windowLimit = cu.limits.find((l: unknown) => {
475
+ const limit = l as Record<string, unknown>;
476
+ return limit.window?.duration === 300;
477
+ });
478
+ if (windowLimit && typeof windowLimit === "object") {
479
+ const wl = windowLimit as Record<string, unknown>;
480
+ if (wl.detail && typeof wl.detail === "object") {
481
+ const wld = wl.detail as Record<string, unknown>;
482
+ const remaining = parseInt(String(wld.remaining ?? ""), 10);
483
+ if (!Number.isNaN(remaining)) {
484
+ apiState.liveKimi5hRemaining = remaining;
485
+ }
486
+ }
487
+ }
488
+ }
489
+ }
490
+ }
491
+
492
+ // Update plan from limit
493
+ if (typeof d.totalQuota === "object" && d.totalQuota) {
494
+ const tq = d.totalQuota as Record<string, unknown>;
495
+ const limit = parseInt(String(tq.limit ?? ""), 10);
496
+ if (!Number.isNaN(limit)) {
497
+ if (limit >= 1200) apiState.livePlan = "vivace";
498
+ else if (limit >= 600) apiState.livePlan = "allegretto";
499
+ else if (limit >= 100) apiState.livePlan = "moderato";
500
+ }
501
+ }
502
+
503
+ return true;
504
+ }
505
+
506
+ async function fetchDeepSeekBalance(apiKey: string): Promise<unknown> {
507
+ try {
508
+ const res = await fetch("https://api.deepseek.com/user/balance", {
509
+ headers: {
510
+ Authorization: `Bearer ${apiKey}`,
511
+ "Content-Type": "application/json",
512
+ },
513
+ signal: AbortSignal.timeout(8000),
514
+ });
515
+ if (!res.ok) return null;
516
+ return await res.json();
517
+ } catch {
518
+ return null;
519
+ }
520
+ }
521
+
522
+ function applyDeepSeekData(data: unknown, apiState: QuotaState): boolean {
523
+ if (!data || typeof data !== "object") return false;
524
+ const d = data as Record<string, unknown>;
525
+
526
+ if (typeof d.is_available === "boolean") {
527
+ apiState.liveDsAvailable = d.is_available;
528
+ }
529
+
530
+ if (Array.isArray(d.balance_infos) && d.balance_infos.length > 0) {
531
+ const info = d.balance_infos.find((i: unknown) => {
532
+ const item = i as Record<string, unknown>;
533
+ return item.currency === "USD" || item.currency === "CNY";
534
+ }) || d.balance_infos[0];
535
+
536
+ if (info && typeof info === "object") {
537
+ const i = info as Record<string, unknown>;
538
+ if (typeof i.total_balance === "string") {
539
+ apiState.liveDsTotalBalance = i.total_balance;
540
+ }
541
+ if (typeof i.currency === "string") {
542
+ apiState.liveDsCurrency = i.currency;
543
+ }
544
+ }
545
+ }
546
+
547
+ return true;
548
+ }
549
+
550
+ // =============================================================================
551
+ // Ollama Cloud Session Cookie & Live Quota
552
+ // =============================================================================
553
+
554
+ type OllamaCookieResult =
555
+ | { ok: true; session: string; aid: string | null }
556
+ | { ok: false; error: string };
557
+
558
+ function readOllamaSessionCookie(): OllamaCookieResult {
559
+ try {
560
+ const home = process.env.HOME || process.env.USERPROFILE || "";
561
+ const authPath = `${home}/.pi/agent/auth.json`;
562
+ const { existsSync, readFileSync } = require("node:fs");
563
+
564
+ if (!existsSync(authPath)) {
565
+ return { ok: false, error: "auth.json not found" };
566
+ }
567
+ let json: Record<string, unknown>;
568
+ try {
569
+ json = JSON.parse(readFileSync(authPath, "utf8"));
570
+ } catch {
571
+ return { ok: false, error: "auth.json is corrupted" };
572
+ }
573
+ const ollamaSession = json["ollama-session"] as Record<string, unknown> | undefined;
574
+ const ollamaGeneral = json["ollama"] as Record<string, unknown> | undefined;
575
+ if (ollamaSession?.key) {
576
+ return {
577
+ ok: true,
578
+ session: ollamaSession.key as string,
579
+ aid: (ollamaSession.aid as string | null) ?? null,
580
+ };
581
+ }
582
+ if (ollamaGeneral?.session) {
583
+ return {
584
+ ok: true,
585
+ session: ollamaGeneral.session as string,
586
+ aid: (ollamaGeneral.aid as string | null) ?? null,
587
+ };
588
+ }
589
+ return { ok: false, error: "No ollama-session key in auth.json" };
590
+ } catch (err) {
591
+ return { ok: false, error: err instanceof Error ? err.message : "Failed to read auth.json" };
592
+ }
593
+ }
594
+
595
+ type OllamaQuotaResult =
596
+ | { ok: true; sessionPercentage: number | null; weeklyPercentage: number | null; sessionResetAt: number | null; weeklyResetAt: number | null; plan: string | null }
597
+ | { ok: false; error: string };
598
+
599
+ async function fetchOllamaLiveQuota(sessionCookie: string, aidCookie?: string | null): Promise<OllamaQuotaResult> {
600
+ try {
601
+ let cookie = `__Secure-session=${sessionCookie}`;
602
+ if (aidCookie) cookie += `; aid=${aidCookie}`;
603
+
604
+ const res = await fetch("https://ollama.com/settings", {
605
+ headers: {
606
+ accept: "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
607
+ "accept-language": "en-US,en;q=0.9",
608
+ cookie,
609
+ "user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/146.0.0.0 Safari/537.36",
610
+ },
611
+ signal: AbortSignal.timeout(10000),
612
+ });
613
+
614
+ if (!res.ok) {
615
+ return { ok: false, error: `HTTP ${res.status}` };
616
+ }
617
+
618
+ const html = await res.text();
619
+
620
+ // Strip HTML tags to get clean text content — collapse to single-line prose
621
+ // so regexes can match naturally across inline elements like <span>54.1</span>%
622
+ const text = html
623
+ .replace(/<script[>\s][\s\S]*?<\/script>/g, " ")
624
+ .replace(/<style[>\s][\s\S]*?<\/style>/g, " ")
625
+ .replace(/<svg[\s\S]*?<\/svg>/g, " ")
626
+ .replace(/<[^>]+>/g, " ")
627
+ .replace(/\s+/g, " ")
628
+ .trim();
629
+
630
+ // Extract usage percentages and reset times from the collapsed text
631
+ const sessionMatch = text.match(/Session\s+usage[\s:]+(\d+\.?\d*)\s*%/i);
632
+ const weeklyMatch = text.match(/Weekly\s+usage[\s:]+(\d+\.?\d*)\s*%/i);
633
+ const sessionResetMatch = text.match(/Session.*?Resets?\s+in\s+(\d+)\s*(?:hours?|hrs?|h)/i);
634
+ const weeklyResetMatch = text.match(/Weekly.*?Resets?\s+in\s+(\d+)\s*(?:days?|d)/i);
635
+
636
+ if (!sessionMatch && !weeklyMatch) {
637
+ return { ok: false, error: "Quota data not found in page" };
638
+ }
639
+
640
+ const rawSession = sessionMatch ? parseFloat(sessionMatch[1]) : null;
641
+ const rawWeekly = weeklyMatch ? parseFloat(weeklyMatch[1]) : null;
642
+ if ((rawSession !== null && !Number.isFinite(rawSession)) || (rawWeekly !== null && !Number.isFinite(rawWeekly))) {
643
+ return { ok: false, error: "Invalid quota data" };
644
+ }
645
+ const sessionPct = rawSession;
646
+ const weeklyPct = rawWeekly;
647
+
648
+ // Compute reset timestamps from relative offsets
649
+ let sessionResetAt: number | null = null;
650
+ let weeklyResetAt: number | null = null;
651
+ if (sessionResetMatch) {
652
+ sessionResetAt = Date.now() + parseInt(sessionResetMatch[1], 10) * 3600 * 1000;
653
+ }
654
+ if (weeklyResetMatch) {
655
+ weeklyResetAt = Date.now() + parseInt(weeklyResetMatch[1], 10) * 86400 * 1000;
656
+ }
657
+
658
+ // Detect plan tier from the page (e.g. "Usage pro Cloud models...")
659
+ const planMatch = text.match(/\bUsage\s+(free|pro|max)\b/i);
660
+ const plan = planMatch ? planMatch[1].toLowerCase() : null;
661
+
662
+ return {
663
+ ok: true,
664
+ sessionPercentage: sessionPct,
665
+ weeklyPercentage: weeklyPct,
666
+ sessionResetAt,
667
+ weeklyResetAt,
668
+ plan,
669
+ };
670
+ } catch (err) {
671
+ return { ok: false, error: err instanceof Error ? err.message : "Network/timeout error" };
672
+ }
673
+ }
674
+
675
+ function applyOllamaLiveData(result: Extract<OllamaQuotaResult, { ok: true }>, apiState: QuotaState): void {
676
+ if (result.sessionPercentage !== null) apiState.liveOllama5hPercentage = result.sessionPercentage;
677
+ if (result.weeklyPercentage !== null) apiState.liveOllamaWeeklyPercentage = result.weeklyPercentage;
678
+ if (result.sessionResetAt !== null) apiState.liveOllamaResetAt = result.sessionResetAt;
679
+ if (result.weeklyResetAt !== null) apiState.liveOllamaWeeklyResetAt = result.weeklyResetAt;
680
+ if (result.plan !== null) apiState.livePlan = result.plan;
681
+ }
682
+
683
+ // =============================================================================
684
+ // Header Parsing
685
+ // =============================================================================
686
+
687
+ const QUOTA_HEADER_KEYS = new Set([
688
+ "x-usage-total-tokens",
689
+ "x-usage-prompt-tokens",
690
+ "x-usage-completion-tokens",
691
+ "x-quota-limit",
692
+ "x-quota-used",
693
+ "x-ratelimit-limit-tokens",
694
+ "x-ratelimit-remaining-tokens",
695
+ "x-ratelimit-reset",
696
+ ]);
697
+
698
+ function applyHeaderRules(headers: Record<string, string>, apiState: QuotaState): void {
699
+ for (const [key, value] of Object.entries(headers)) {
700
+ const normalized = key.toLowerCase();
701
+ if (QUOTA_HEADER_KEYS.has(normalized)) {
702
+ apiState.rawHeaders[normalized] = value;
703
+ }
704
+ }
705
+
706
+ if (apiState.rawHeaders["x-quota-limit"]) {
707
+ const limit = parseInt(apiState.rawHeaders["x-quota-limit"], 10);
708
+ if (!Number.isNaN(limit) && limit > 0) {
709
+ apiState.headerPromptsLimit = limit;
710
+ }
711
+ }
712
+ }
713
+
714
+ // =============================================================================
715
+ // Formatting
716
+ // =============================================================================
717
+
718
+ function progressBar(pct: number, width: number): string {
719
+ const filled = Math.round((pct / 100) * Math.max(0, Math.min(100, pct)) * width / 100);
720
+ return "█".repeat(Math.max(0, filled)) + "░".repeat(Math.max(0, width - filled));
721
+ }
722
+
723
+ function formatStatus(apiState: QuotaState): string {
724
+ if (!apiState.provider || !isTrackedProvider(apiState.provider)) {
725
+ return "";
726
+ }
727
+
728
+ const plan = apiState.plan || apiState.livePlan || "";
729
+ const planLabel = plan ? ` ${plan}` : "";
730
+
731
+ if (isDeepseek(apiState.provider)) {
732
+ const balance = apiState.liveDsTotalBalance ?? "—";
733
+ const prefix = dsCurrencyPrefix(apiState.liveDsCurrency);
734
+ const available = apiState.liveDsAvailable === false ? " · ❌ No" : "";
735
+ const error = apiState.lastFetchError ? ` · ⚠️ ${apiState.lastFetchError}` : "";
736
+ return `quota: ds${planLabel}${available} · ${prefix}${balance}${error}`;
737
+ }
738
+
739
+ if (isOllama(apiState.provider)) {
740
+ // Ollama: prefer live quota data if available, otherwise use turn counting
741
+ if (apiState.liveOllama5hPercentage !== null) {
742
+ // Use live data from Ollama settings page
743
+ const pct5h = apiState.liveOllama5hPercentage;
744
+ const pctWk = apiState.liveOllamaWeeklyPercentage ?? 0;
745
+ const reset5h = apiState.liveOllamaResetAt ? ` · reset ${relativeTime(apiState.liveOllamaResetAt)}` : "";
746
+ const resetWk = apiState.liveOllamaWeeklyResetAt ? ` · wk reset ${relativeTime(apiState.liveOllamaWeeklyResetAt)}` : "";
747
+
748
+ const color5h = pct5h >= 80 ? "🔴" : pct5h >= 50 ? "🟡" : "🟢";
749
+ const colorWk = pctWk >= 80 ? "🔴" : pctWk >= 50 ? "🟡" : "🟢";
750
+
751
+ return `quota: ollama${planLabel} · 5h: ${color5h} ${Math.round(pct5h)}%${reset5h} · wk: ${colorWk} ${Math.round(pctWk)}%${resetWk}`;
752
+ }
753
+
754
+ // Fall back to turn counting with GPU-time weighted units
755
+ const limit = getLimit();
756
+ const weeklyLimit = getWeeklyLimit();
757
+ const pct5h = limit ? Math.min(100, (apiState.ollamaGpuUnitsUsed / limit) * 100) : 0;
758
+ const pctWk = weeklyLimit ? Math.min(100, (apiState.turnsUsedWeekly / weeklyLimit) * 100) : 0;
759
+ const mult = ollamaGpuMultiplier(apiState.model);
760
+ const multLabel = mult > 1 ? ` · ${mult}×` : "";
761
+
762
+ const color5h = pct5h >= 80 ? "🔴" : pct5h >= 50 ? "🟡" : "🟢";
763
+ const colorWk = pctWk >= 80 ? "🔴" : pctWk >= 50 ? "🟡" : "🟢";
764
+
765
+ const resetLabel = apiState.cycleStart5h ? ` · reset ${relativeTime(apiState.cycleStart5h + 5 * 60 * 60 * 1000)}` : "";
766
+ const errorLabel = apiState.lastFetchError ? ` · ⚠️ ${apiState.lastFetchError}` : "";
767
+
768
+ return `quota: ollama${planLabel} · 5h: ${color5h} ${Math.round(pct5h)}% · wk: ${colorWk} ${Math.round(pctWk)}%${resetLabel}${multLabel}${errorLabel}`;
769
+ }
770
+
771
+ if (isKimi(apiState.provider)) {
772
+ // Kimi: show live data or turn counting
773
+ if (apiState.liveKimi5hRemaining !== null) {
774
+ const error = apiState.kimiJwtExpired ? " · ⚠️ JWT expired" : "";
775
+ return `quota: kimi${planLabel} · 5h: 🟢 ${apiState.liveKimi5hRemaining} remaining${error}`;
776
+ }
777
+
778
+ const limit = getLimit() ?? 100;
779
+ const pct = Math.min(100, (apiState.turnsUsed5h / limit) * 100);
780
+ const color = pct >= 80 ? "🔴" : pct >= 50 ? "🟡" : "🟢";
781
+ return `quota: kimi${planLabel} · 5h: ${color} ${Math.round(pct)}%`;
782
+ }
783
+
784
+ // Z.Ai: show live data or turn counting
785
+ if (apiState.live5hPercentage !== null) {
786
+ const pct5h = apiState.live5hPercentage;
787
+ const pctWk = apiState.liveWeeklyPercentage ?? 0;
788
+ const resetLabel = apiState.live5hResetAt ? ` · reset ${relativeTime(apiState.live5hResetAt)}` : "";
789
+ const color5h = pct5h >= 80 ? "🔴" : pct5h >= 50 ? "🟡" : "🟢";
790
+ const colorWk = pctWk >= 80 ? "🔴" : pctWk >= 50 ? "🟡" : "🟢";
791
+ return `quota: zai${planLabel} · 5h: ${color5h} ${Math.round(pct5h)}% · wk: ${colorWk} ${Math.round(pctWk)}%${resetLabel}`;
792
+ }
793
+
794
+ // Fallback to turn counting
795
+ const limit = getLimit();
796
+ const weeklyLimit = getWeeklyLimit();
797
+ const pct5h = limit ? Math.min(100, (apiState.quotaUsed5h / limit) * 100) : 0;
798
+ const pctWk = weeklyLimit ? Math.min(100, (apiState.turnsUsedWeekly / weeklyLimit) * 100) : 0;
799
+ const mult = quotaMultiplier(apiState.model);
800
+ const multLabel = mult > 1 ? ` · ${mult}×` : "";
801
+
802
+ const color5h = pct5h >= 80 ? "🔴" : pct5h >= 50 ? "🟡" : "🟢";
803
+ const colorWk = pctWk >= 80 ? "🔴" : pctWk >= 50 ? "🟡" : "🟢";
804
+
805
+ return `quota: zai${planLabel} · 5h: ${color5h} ${Math.round(pct5h)}% · wk: ${colorWk} ${Math.round(pctWk)}%${multLabel}`;
806
+ }
807
+
808
+ // =============================================================================
809
+ // Polling
810
+ // =============================================================================
811
+
812
+ let pollTimer: ReturnType<typeof setInterval> | null = null;
813
+ const POLL_INTERVAL_MS = 30_000; // 30 seconds
814
+
815
+ async function refreshLiveQuota(): Promise<void> {
816
+ if (!isTrackedProvider(state.provider)) return;
817
+
818
+ // For Ollama, try to fetch live quota from session cookie, otherwise use turn counting
819
+ if (isOllama(state.provider)) {
820
+ const cookieResult = readOllamaSessionCookie();
821
+ if (!cookieResult.ok) {
822
+ state.liveOllama5hPercentage = null;
823
+ state.liveOllamaWeeklyPercentage = null;
824
+ state.liveOllamaResetAt = null;
825
+ state.liveOllamaWeeklyResetAt = null;
826
+ state.lastFetchError = `No ollama-session cookie: ${cookieResult.error} (see ~/.pi/agent/auth.json)`;
827
+ refreshStatus();
828
+ return;
829
+ }
830
+ const { session, aid } = cookieResult;
831
+ const result = await fetchOllamaLiveQuota(session, aid);
832
+ if (result.ok) {
833
+ applyOllamaLiveData(result, state);
834
+ state.lastSuccessfulFetchAt = Date.now();
835
+ state.lastFetchError = null;
836
+ } else {
837
+ state.liveOllama5hPercentage = null;
838
+ state.liveOllamaWeeklyPercentage = null;
839
+ state.liveOllamaResetAt = null;
840
+ state.liveOllamaWeeklyResetAt = null;
841
+ state.lastFetchError = `Ollama fetch failed: ${result.error}`;
842
+ }
843
+ refreshStatus();
844
+ return;
845
+ }
846
+
847
+ if (isZai(state.provider)) {
848
+ const apiKey = readApiKey("zai");
849
+ if (apiKey) {
850
+ const data = await fetchLiveQuota(apiKey);
851
+ if (data) {
852
+ applyLiveData(data, state);
853
+ state.lastSuccessfulFetchAt = Date.now();
854
+ state.lastFetchError = null;
855
+ }
856
+ }
857
+ refreshStatus();
858
+ return;
859
+ }
860
+
861
+ if (isKimi(state.provider)) {
862
+ if (state.kimiJwtExpired) {
863
+ refreshStatus();
864
+ return;
865
+ }
866
+ const webJwt = readKimiWebJwt();
867
+ if (webJwt) {
868
+ const exp = getJwtExpiry(webJwt);
869
+ if (exp && exp * 1000 < Date.now()) {
870
+ state.kimiJwtExpired = true;
871
+ refreshStatus();
872
+ return;
873
+ }
874
+ const data = await fetchKimiLiveQuota(webJwt);
875
+ if (data) {
876
+ applyKimiLiveData(data, state);
877
+ state.lastSuccessfulFetchAt = Date.now();
878
+ state.lastFetchError = null;
879
+ }
880
+ }
881
+ refreshStatus();
882
+ return;
883
+ }
884
+
885
+ if (isDeepseek(state.provider)) {
886
+ const apiKey = readApiKey("deepseek");
887
+ if (apiKey) {
888
+ const data = await fetchDeepSeekBalance(apiKey);
889
+ if (data) {
890
+ applyDeepSeekData(data, state);
891
+ state.lastSuccessfulFetchAt = Date.now();
892
+ state.lastFetchError = null;
893
+ }
894
+ }
895
+ refreshStatus();
896
+ return;
897
+ }
898
+ }
899
+
900
+ function startPolling(): void {
901
+ if (pollTimer) return;
902
+ pollTimer = setInterval(async () => {
903
+ await refreshLiveQuota();
904
+ }, POLL_INTERVAL_MS);
905
+ if (pollTimer && "unref" in pollTimer) pollTimer.unref();
906
+ }
907
+
908
+ function stopPolling(): void {
909
+ if (pollTimer) {
910
+ clearInterval(pollTimer);
911
+ pollTimer = null;
912
+ }
913
+ }
914
+
915
+ function refreshStatus(): void {
916
+ const status = formatStatus(state);
917
+ if (status) {
918
+ pi.ui.setStatus("provider-quota", status);
919
+ } else {
920
+ pi.ui.setStatus("provider-quota", undefined);
921
+ }
922
+ }
923
+
924
+ // =============================================================================
925
+ // Default Export
926
+ // =============================================================================
927
+
928
+ export default function (pi: ExtensionAPI): void {
929
+ (globalThis as unknown as { pi: ExtensionAPI }).pi = pi;
930
+
931
+ // Session start - restore state
932
+ pi.on("session_start", async (_event, ctx) => {
933
+ for (const entry of ctx.sessionManager.getEntries()) {
934
+ if (entry.type === "custom" && entry.customType === "provider-quota") {
935
+ const safe = sanitizeState(entry.data);
936
+ state = { ...createFreshState(), ...safe };
937
+ break;
938
+ }
939
+ }
940
+
941
+ if (state.provider && isTrackedProvider(state.provider)) {
942
+ startPolling();
943
+ await refreshLiveQuota();
944
+ }
945
+ });
946
+
947
+ // Model select - switch provider
948
+ pi.on("model_select", async (event, ctx) => {
949
+ const { provider, model } = event;
950
+
951
+ if (!isTrackedProvider(provider)) {
952
+ resetProviderData();
953
+ state.provider = "";
954
+ state.model = "";
955
+ stopPolling();
956
+ refreshStatus();
957
+ return;
958
+ }
959
+
960
+ resetProviderData();
961
+ state.provider = provider;
962
+ state.model = model;
963
+ state.cycleStart5h = state.cycleStart5h ?? Date.now();
964
+ state.weeklyStart = state.weeklyStart ?? Date.now();
965
+
966
+ startPolling();
967
+ await refreshLiveQuota();
968
+ });
969
+
970
+ // Turn start - rotate cycles
971
+ pi.on("turn_start", async () => {
972
+ maybeRotateCycle();
973
+ maybeRotateWeekly();
974
+ refreshStatus();
975
+ });
976
+
977
+ // Turn end - count usage
978
+ pi.on("turn_end", async (event) => {
979
+ maybeRotateCycle();
980
+ maybeRotateWeekly();
981
+
982
+ const multiplier = quotaMultiplier(state.model);
983
+ state.turnsUsed5h++;
984
+ state.quotaUsed5h += multiplier;
985
+ state.turnsUsedWeekly++;
986
+
987
+ if (isOllama(state.provider)) {
988
+ state.ollamaGpuUnitsUsed += ollamaGpuMultiplier(state.model);
989
+ }
990
+
991
+ refreshStatus();
992
+ });
993
+
994
+ // After provider response - parse headers
995
+ pi.on("after_provider_response", async (event) => {
996
+ const headers = event.response?.headers;
997
+ if (headers && typeof headers === "object") {
998
+ const headerObj: Record<string, string> = {};
999
+ for (const [key, value] of Object.entries(headers)) {
1000
+ if (typeof value === "string") {
1001
+ headerObj[key] = value;
1002
+ } else if (Array.isArray(value)) {
1003
+ headerObj[key] = value.join(", ");
1004
+ }
1005
+ }
1006
+ applyHeaderRules(headerObj, state);
1007
+ }
1008
+ state.lastResponseAt = Date.now();
1009
+ refreshStatus();
1010
+ });
1011
+
1012
+ // Session shutdown - persist state
1013
+ pi.on("session_shutdown", async () => {
1014
+ stopPolling();
1015
+ const persist = { ...state, rawHeaders: {} };
1016
+ pi.appendEntry("provider-quota", persist);
1017
+ });
1018
+
1019
+ // Commands
1020
+ pi.registerCommand("quota", {
1021
+ description: "Show quota usage details",
1022
+ handler: async (_args, ctx) => {
1023
+ await refreshLiveQuota();
1024
+ const lines: string[] = [];
1025
+ lines.push(`Provider: ${state.provider || "—"}`);
1026
+ lines.push(`Model: ${state.model || "—"}`);
1027
+ lines.push(`Plan: ${state.plan || state.livePlan || "—"}`);
1028
+
1029
+ if (isZai(state.provider)) {
1030
+ lines.push(`Live 5h: ${state.live5hPercentage !== null ? `${Math.round(state.live5hPercentage)}%` : "—"}${state.live5hResetAt ? ` (resets ${relativeTime(state.live5hResetAt)})` : ""}`);
1031
+ lines.push(`Live wk: ${state.liveWeeklyPercentage !== null ? `${Math.round(state.liveWeeklyPercentage)}%` : "—"}`);
1032
+ if (state.liveMcpDetails.length > 0) {
1033
+ lines.push("MCP Usage:");
1034
+ for (const d of state.liveMcpDetails) {
1035
+ if (d.usage > 0) lines.push(` ${d.modelCode}: ${d.usage}`);
1036
+ }
1037
+ }
1038
+ }
1039
+
1040
+ if (isKimi(state.provider)) {
1041
+ lines.push(`5h remaining: ${state.liveKimi5hRemaining ?? "—"}`);
1042
+ lines.push(`Weekly remaining: ${state.liveKimiWeeklyRemaining ?? "—"}`);
1043
+ if (state.kimiJwtExpired) lines.push("⚠️ JWT expired");
1044
+ }
1045
+
1046
+ if (isOllama(state.provider)) {
1047
+ if (state.liveOllama5hPercentage !== null) {
1048
+ lines.push(`Live session: ${Math.round(state.liveOllama5hPercentage)}% used`);
1049
+ lines.push(`Live weekly: ${state.liveOllamaWeeklyPercentage !== null ? Math.round(state.liveOllamaWeeklyPercentage) + "% used" : "—"}`);
1050
+ if (state.liveOllamaResetAt) lines.push(`Session resets: ${relativeTime(state.liveOllamaResetAt)}`);
1051
+ if (state.liveOllamaWeeklyResetAt) lines.push(`Weekly resets: ${relativeTime(state.liveOllamaWeeklyResetAt)}`);
1052
+ } else {
1053
+ const cookieResult = readOllamaSessionCookie();
1054
+ lines.push(`GPU units (5h): ${state.ollamaGpuUnitsUsed}`);
1055
+ lines.push(`Weekly turns: ${state.turnsUsedWeekly}`);
1056
+ lines.push(`Session cookie: ${cookieResult.ok ? "✅ found" : "❌ " + cookieResult.error}`);
1057
+ }
1058
+ }
1059
+
1060
+ if (isDeepseek(state.provider)) {
1061
+ lines.push(`Balance: ${dsCurrencyPrefix(state.liveDsCurrency)}${state.liveDsTotalBalance ?? "—"}`);
1062
+ lines.push(`Available: ${state.liveDsAvailable === false ? "No" : state.liveDsAvailable === true ? "Yes" : "—"}`);
1063
+ }
1064
+
1065
+ lines.push(`Local 5h: ${state.quotaUsed5h} units (${state.turnsUsed5h} turns)`);
1066
+ lines.push(`Local wk: ${state.turnsUsedWeekly} turns`);
1067
+
1068
+ if (Object.keys(state.rawHeaders).length > 0) {
1069
+ lines.push("Headers:");
1070
+ for (const key of Object.keys(state.rawHeaders).sort()) {
1071
+ lines.push(` ${key}: ${state.rawHeaders[key]}`);
1072
+ }
1073
+ }
1074
+
1075
+ if (state.lastFetchError) {
1076
+ lines.push(`Error: ${state.lastFetchError}`);
1077
+ }
1078
+
1079
+ ctx.ui.notify(lines.join("\n"), "info");
1080
+ },
1081
+ });
1082
+
1083
+ pi.registerCommand("quota-plan", {
1084
+ description: "Set plan tier",
1085
+ getArgumentCompletions: () => {
1086
+ if (isZai(state.provider)) return ["lite", "pro", "max"];
1087
+ if (isKimi(state.provider)) return ["moderato", "allegretto", "vivace"];
1088
+ if (isOllama(state.provider)) return ["free", "pro", "max"];
1089
+ return [];
1090
+ },
1091
+ handler: async (args, ctx) => {
1092
+ if (isDeepseek(state.provider)) {
1093
+ ctx.ui.notify("DeepSeek uses pay-as-you-go billing with no plan tiers.", "info");
1094
+ return;
1095
+ }
1096
+ const plan = args.plan?.toLowerCase();
1097
+ if (!plan) {
1098
+ ctx.ui.notify(`Current plan: ${state.plan || "not set"}`, "info");
1099
+ return;
1100
+ }
1101
+ state.plan = plan;
1102
+ ctx.ui.notify(`Plan set to: ${plan}`, "info");
1103
+ refreshStatus();
1104
+ },
1105
+ });
1106
+
1107
+ // Tool for LLM
1108
+ pi.registerTool({
1109
+ name: "provider_quota",
1110
+ label: "Check Provider Quota",
1111
+ description: "Check current quota usage for the active provider",
1112
+ parameters: Type.Object({}),
1113
+ execute: async () => {
1114
+ await refreshLiveQuota();
1115
+ return {
1116
+ provider: state.provider,
1117
+ model: state.model,
1118
+ plan: state.plan || state.livePlan,
1119
+ live5hPercentage: state.live5hPercentage,
1120
+ liveWeeklyPercentage: state.liveWeeklyPercentage,
1121
+ livePlan: state.livePlan,
1122
+ turnsUsed5h: state.turnsUsed5h,
1123
+ quotaUsed5h: state.quotaUsed5h,
1124
+ turnsUsedWeekly: state.turnsUsedWeekly,
1125
+ kimi5hRemaining: state.liveKimi5hRemaining,
1126
+ kimiWeeklyRemaining: state.liveKimiWeeklyRemaining,
1127
+ kimiJwtExpired: state.kimiJwtExpired,
1128
+ deepseekBalance: state.liveDsTotalBalance ? `${dsCurrencyPrefix(state.liveDsCurrency)}${state.liveDsTotalBalance}` : null,
1129
+ deepseekAvailable: state.liveDsAvailable,
1130
+ ollamaGpuUnitsUsed: state.ollamaGpuUnitsUsed,
1131
+ ollamaLiveSessionPct: state.liveOllama5hPercentage,
1132
+ ollamaLiveWeeklyPct: state.liveOllamaWeeklyPercentage,
1133
+ ollamaSessionResetAt: state.liveOllamaResetAt,
1134
+ ollamaWeeklyResetAt: state.liveOllamaWeeklyResetAt,
1135
+ lastFetchError: state.lastFetchError,
1136
+ lastResponseAt: state.lastResponseAt,
1137
+ };
1138
+ },
1139
+ });
1140
+ }
1141
+
1142
+ // TypeScript needs this reference
1143
+ declare const pi: ExtensionAPI;
1144
+ (globalThis as unknown as { pi: ExtensionAPI }).pi = {} as ExtensionAPI;
package/package.json ADDED
@@ -0,0 +1,38 @@
1
+ {
2
+ "name": "pi-provider-quota",
3
+ "version": "0.1.0",
4
+ "description": "Track Z.Ai, Kimi Code, Ollama Cloud, and DeepSeek quota in pi's status bar",
5
+ "type": "module",
6
+ "main": "extensions/provider-quota.ts",
7
+ "license": "MIT",
8
+ "author": "",
9
+ "repository": {
10
+ "type": "git",
11
+ "url": "git+https://github.com/ronnieops/pi-provider-quota.git"
12
+ },
13
+ "keywords": [
14
+ "pi-package",
15
+ "pi",
16
+ "pi-coding-agent",
17
+ "quota",
18
+ "ollama",
19
+ "deepseek",
20
+ "kimi",
21
+ "zai",
22
+ "status-bar"
23
+ ],
24
+ "pi": {
25
+ "extensions": [
26
+ "./extensions/provider-quota.ts"
27
+ ]
28
+ },
29
+ "peerDependencies": {
30
+ "@mariozechner/pi-coding-agent": "*"
31
+ },
32
+ "dependencies": {
33
+ "typebox": "^1.1.24"
34
+ },
35
+ "devDependencies": {
36
+ "typescript": "^6.0.3"
37
+ }
38
+ }