diff --git a/src/api/client.test.ts b/src/api/client.test.ts index 7985004..0753dc3 100644 --- a/src/api/client.test.ts +++ b/src/api/client.test.ts @@ -142,3 +142,48 @@ test("chatStream: unreachable host yields an error chunk instead of throwing", a globalThis.fetch = realFetch; } }); + +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: {}, + }; + const stats = KlaatAIClient.mergeUsagePlan( + new Headers({ "X-KlaatAI-Quota-Plan": "pro" }), + data, + ); + expect(stats.plan).toBe("pro"); +}); + +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, + total_tokens: 0, + total_cost_usd: 0, + by_tier: {}, + plan: "free", + }); + 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 5864614..86878ba 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,8 @@ 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 data = await res.json() as LifetimeUsageStats & { plan?: string }; + return KlaatAIClient.mergeUsagePlan(res.headers, data); } catch { return null; } @@ -835,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 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();