fix(auth): show live subscription plan in whoami

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");
+});

src/api/client.ts

@@ -109,6 +109,8 @@ export interface LifetimeUsageStats {
   total_tokens: number;
   total_cost_usd: number;
   by_tier: Record<string, TierUsage>;
+  /** 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<LifetimeUsageStats>;
+      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

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;
       }

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" };

src/auth/login.ts

@@ -28,19 +28,23 @@ export async function runWhoami(baseUrl: string, json = false): Promise<void> {
   }
   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();