From 6de70c61216d4e13268cbfb1015a4fd802d3ee44 Mon Sep 17 00:00:00 2001 From: syf2211 Date: Wed, 26 Aug 2026 00:12:34 +0000 Subject: [PATCH 1/2] fix(auth): show live subscription plan in whoami Fetch plan from authenticated /v1/me/usage (quota headers or JSON body) instead of relying on an unpopulated credentials field. Also capture plan from the OAuth callback when the web redirect includes it. Fixes #24 --- src/api/client.test.ts | 40 ++++++++++++++++++++++++++++++++++++++++ src/api/client.ts | 7 ++++++- src/auth/browser.ts | 1 + src/auth/login.test.ts | 24 ++++++++++++++++++++++++ src/auth/login.ts | 10 +++++++--- 5 files changed, 78 insertions(+), 4 deletions(-) diff --git a/src/api/client.test.ts b/src/api/client.test.ts index 7985004..1953e10 100644 --- a/src/api/client.test.ts +++ b/src/api/client.test.ts @@ -142,3 +142,43 @@ test("chatStream: unreachable host yields an error chunk instead of throwing", a globalThis.fetch = realFetch; } }); + +test("getUsageStats: reads plan from quota headers", async () => { + const realFetch = globalThis.fetch; + globalThis.fetch = (async () => new Response(JSON.stringify({ + total_requests: 3, + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + total_cost_usd: 0.01, + by_tier: {}, + }), { + status: 200, + headers: { "X-KlaatAI-Quota-Plan": "pro" }, + })) as typeof fetch; + try { + const stats = await new KlaatAIClient({ apiKey: "test" }).getUsageStats(); + expect(stats?.plan).toBe("pro"); + } finally { + globalThis.fetch = realFetch; + } +}); + +test("getUsageStats: reads plan from JSON body when headers omit it", async () => { + const realFetch = globalThis.fetch; + globalThis.fetch = (async () => new Response(JSON.stringify({ + total_requests: 1, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + total_cost_usd: 0, + by_tier: {}, + plan: "free", + }), { status: 200 })) as typeof fetch; + try { + const stats = await new KlaatAIClient({ apiKey: "test" }).getUsageStats(); + expect(stats?.plan).toBe("free"); + } finally { + globalThis.fetch = realFetch; + } +}); diff --git a/src/api/client.ts b/src/api/client.ts index 5864614..0f5329f 100644 --- a/src/api/client.ts +++ b/src/api/client.ts @@ -109,6 +109,8 @@ export interface LifetimeUsageStats { total_tokens: number; total_cost_usd: number; by_tier: Record; + /** Subscription plan when the backend exposes it on /v1/me/usage. */ + plan?: string; } export interface ClientOptions { @@ -532,7 +534,10 @@ export class KlaatAIClient { res = await fetch(`${this.baseUrl}/v1/me/usage`, { headers: this.headers() }); } if (!res.ok) return null; - return res.json() as Promise; + const quota = KlaatAIClient.parseQuotaHeaders(res.headers); + const data = await res.json() as LifetimeUsageStats & { plan?: string }; + const plan = quota?.plan ?? data.plan; + return plan ? { ...data, plan } : data; } catch { return null; } diff --git a/src/auth/browser.ts b/src/auth/browser.ts index 148f1df..62eddc9 100644 --- a/src/auth/browser.ts +++ b/src/auth/browser.ts @@ -205,6 +205,7 @@ export async function startOAuthBrowserAuth( expiresAt: Math.floor(Date.now() / 1000) + expiresIn, userId: url.searchParams.get("user_id") ?? undefined, email: url.searchParams.get("email") ?? undefined, + plan: url.searchParams.get("plan") ?? undefined, } : null); return; } diff --git a/src/auth/login.test.ts b/src/auth/login.test.ts index b1463cd..88f57e6 100644 --- a/src/auth/login.test.ts +++ b/src/auth/login.test.ts @@ -19,6 +19,7 @@ type Creds = { accessToken?: string | null; email?: string | null; plan?: string let mockToken: string | null = null; let mockCreds: Creds = {}; let pingImpl: () => Promise<{ status: string }>; +let usageImpl: () => Promise<{ plan?: string } | null>; mock.module("../auth/credentials.js", () => ({ ...realCredentials, @@ -35,6 +36,7 @@ mock.module("../api/client.js", () => ({ KlaatAIClient: class extends realClient.KlaatAIClient { constructor() { super({ baseUrl: "http://mock.invalid" }); } override async ping(): Promise<{ status: string }> { return pingImpl(); } + override async getUsageStats() { return usageImpl(); } }, })); @@ -53,6 +55,7 @@ let origErr: typeof console.error; beforeEach(() => { mockToken = null; mockCreds = {}; + usageImpl = async () => null; stdout = []; stderr = []; origLog = console.log; @@ -80,6 +83,7 @@ describe("runWhoami (json: false)", () => { mockToken = "jwt-abc"; mockCreds = { accessToken: "jwt-abc", email: "demo@klaatai.com", plan: "pro" }; pingImpl = async () => ({ status: "ok" }); + usageImpl = async () => ({ plan: "pro", total_requests: 1, prompt_tokens: 0, completion_tokens: 0, total_tokens: 0, total_cost_usd: 0, by_tier: {} }); await runWhoami("http://127.0.0.1:8765", false); const blob = jsonOut(); expect(blob).toContain("demo@klaatai.com"); @@ -89,6 +93,15 @@ describe("runWhoami (json: false)", () => { expect(errOut()).toBe(""); }); + test("authenticated + live usage plan shown when credentials omit plan", async () => { + mockToken = "jwt-abc"; + mockCreds = { accessToken: "jwt-abc", email: "demo@klaatai.com" }; + pingImpl = async () => ({ status: "ok" }); + usageImpl = async () => ({ plan: "pro", total_requests: 1, prompt_tokens: 0, completion_tokens: 0, total_tokens: 0, total_cost_usd: 0, by_tier: {} }); + await runWhoami("http://127.0.0.1:8765", false); + expect(jsonOut()).toContain("pro"); + }); + test("authenticated but API unreachable → writes error to stderr", async () => { mockToken = "jwt-abc"; mockCreds = { accessToken: "jwt-abc", email: "x@y.z", plan: "free" }; @@ -114,6 +127,7 @@ describe("runWhoami (json: true)", () => { mockToken = "jwt-abc"; mockCreds = { accessToken: "jwt-abc", email: "demo@klaatai.com", plan: "pro" }; pingImpl = async () => ({ status: "ok" }); + usageImpl = async () => ({ plan: "pro", total_requests: 1, prompt_tokens: 0, completion_tokens: 0, total_tokens: 0, total_cost_usd: 0, by_tier: {} }); await runWhoami("http://127.0.0.1:8765", true); const obj = parse(jsonOut()); expect(obj.signedIn).toBe(true); @@ -123,6 +137,16 @@ describe("runWhoami (json: true)", () => { expect(errOut()).toBe(""); }); + test("authenticated + live usage plan overrides stale stored plan", async () => { + mockToken = "jwt-abc"; + mockCreds = { accessToken: "jwt-abc", email: "demo@klaatai.com", plan: "free" }; + pingImpl = async () => ({ status: "ok" }); + usageImpl = async () => ({ plan: "pro", total_requests: 1, prompt_tokens: 0, completion_tokens: 0, total_tokens: 0, total_cost_usd: 0, by_tier: {} }); + await runWhoami("http://127.0.0.1:8765", true); + const obj = parse(jsonOut()); + expect(obj.plan).toBe("pro"); + }); + test("authenticated + backend offline → online:false schema on stderr", async () => { mockToken = "jwt-abc"; mockCreds = { accessToken: "jwt-abc", email: "x@y.z", plan: "free" }; diff --git a/src/auth/login.ts b/src/auth/login.ts index 5ff4753..3155709 100644 --- a/src/auth/login.ts +++ b/src/auth/login.ts @@ -28,19 +28,23 @@ export async function runWhoami(baseUrl: string, json = false): Promise { } const client = new KlaatAIClient({ apiKey: token, baseUrl }); try { - const info = await client.ping(); + const [info, usage] = await Promise.all([ + client.ping(), + client.getUsageStats(), + ]); const creds = loadCredentials(); + const plan = usage?.plan ?? creds.plan ?? null; if (json) { console.log(JSON.stringify({ signedIn: true, email: creds.email ?? null, - plan: creds.plan ?? null, + plan, backend: info.status === "ok" ? "online" : "offline", }, null, 2)); } else { console.log(); if (creds.email) console.log(chalk.bold(" Account: ") + creds.email); - if (creds.plan) console.log(chalk.bold(" Plan: ") + creds.plan); + if (plan) console.log(chalk.bold(" Plan: ") + plan); console.log(chalk.bold(" Session: ") + "subscription (JWT)"); console.log(chalk.bold(" Backend: ") + (info.status === "ok" ? chalk.green("Online") : chalk.red("Offline"))); console.log(); From 40f7dcccbdf6ce1e37d6f8e26c4bc0b3b644cdfc Mon Sep 17 00:00:00 2001 From: syf2211 Date: Wed, 26 Aug 2026 00:14:31 +0000 Subject: [PATCH 2/2] test(client): avoid fetch mocks that leak across test files Extract mergeUsagePlan static helper and unit-test it directly so login.test.ts module mocks do not break getUsageStats coverage in CI. --- src/api/client.test.ts | 51 +++++++++++++++++++++++------------------- src/api/client.ts | 16 ++++++++++--- 2 files changed, 41 insertions(+), 26 deletions(-) diff --git a/src/api/client.test.ts b/src/api/client.test.ts index 1953e10..0753dc3 100644 --- a/src/api/client.test.ts +++ b/src/api/client.test.ts @@ -143,30 +143,24 @@ test("chatStream: unreachable host yields an error chunk instead of throwing", a } }); -test("getUsageStats: reads plan from quota headers", async () => { - const realFetch = globalThis.fetch; - globalThis.fetch = (async () => new Response(JSON.stringify({ +test("mergeUsagePlan: reads plan from quota headers", () => { + const data = { total_requests: 3, prompt_tokens: 10, completion_tokens: 5, total_tokens: 15, total_cost_usd: 0.01, by_tier: {}, - }), { - status: 200, - headers: { "X-KlaatAI-Quota-Plan": "pro" }, - })) as typeof fetch; - try { - const stats = await new KlaatAIClient({ apiKey: "test" }).getUsageStats(); - expect(stats?.plan).toBe("pro"); - } finally { - globalThis.fetch = realFetch; - } + }; + const stats = KlaatAIClient.mergeUsagePlan( + new Headers({ "X-KlaatAI-Quota-Plan": "pro" }), + data, + ); + expect(stats.plan).toBe("pro"); }); -test("getUsageStats: reads plan from JSON body when headers omit it", async () => { - const realFetch = globalThis.fetch; - globalThis.fetch = (async () => new Response(JSON.stringify({ +test("mergeUsagePlan: reads plan from JSON body when headers omit it", () => { + const stats = KlaatAIClient.mergeUsagePlan(new Headers(), { total_requests: 1, prompt_tokens: 0, completion_tokens: 0, @@ -174,11 +168,22 @@ test("getUsageStats: reads plan from JSON body when headers omit it", async () = total_cost_usd: 0, by_tier: {}, plan: "free", - }), { status: 200 })) as typeof fetch; - try { - const stats = await new KlaatAIClient({ apiKey: "test" }).getUsageStats(); - expect(stats?.plan).toBe("free"); - } finally { - globalThis.fetch = realFetch; - } + }); + expect(stats.plan).toBe("free"); +}); + +test("mergeUsagePlan: header plan wins over JSON body", () => { + const stats = KlaatAIClient.mergeUsagePlan( + new Headers({ "X-KlaatAI-Quota-Plan": "pro" }), + { + total_requests: 1, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + total_cost_usd: 0, + by_tier: {}, + plan: "free", + }, + ); + expect(stats.plan).toBe("pro"); }); diff --git a/src/api/client.ts b/src/api/client.ts index 0f5329f..86878ba 100644 --- a/src/api/client.ts +++ b/src/api/client.ts @@ -534,10 +534,8 @@ export class KlaatAIClient { res = await fetch(`${this.baseUrl}/v1/me/usage`, { headers: this.headers() }); } if (!res.ok) return null; - const quota = KlaatAIClient.parseQuotaHeaders(res.headers); const data = await res.json() as LifetimeUsageStats & { plan?: string }; - const plan = quota?.plan ?? data.plan; - return plan ? { ...data, plan } : data; + return KlaatAIClient.mergeUsagePlan(res.headers, data); } catch { return null; } @@ -840,6 +838,18 @@ export class KlaatAIClient { return null; } + /** + * Attach subscription plan from quota headers or JSON body to usage stats. + * Header wins when both are present (matches chat completion behavior). + */ + static mergeUsagePlan( + h: Headers, + data: LifetimeUsageStats & { plan?: string }, + ): LifetimeUsageStats { + const plan = KlaatAIClient.parseQuotaHeaders(h)?.plan ?? data.plan; + return plan ? { ...data, plan } : data; + } + /** * Parse the E1 weighted-unit quota + tier from response headers. Returns null * when none are present (older server / non-subscription auth). Tolerant of a