fix(kiro): resolve quota for IAM Identity Center accounts missing a profileArn (#3722)

Integrated into release/v3.8.23

open-sse/services/usage.ts

@@ -2995,24 +2995,65 @@ export function buildKiroUsageResult(
   };
 }
 
+/**
+ * Discover a Kiro/CodeWhisperer profile ARN for an account that didn't persist one (common for
+ * AWS IAM Identity Center logins and kiro-cli imports). Calls ListAvailableProfiles on the
+ * region-matched endpoint and prefers a profile whose ARN is in the same region. Returns
+ * undefined when no profile is available (e.g. the org/token has no Kiro entitlement).
+ * Exported for testing.
+ */
+export async function discoverKiroProfileArn(
+  accessToken: string,
+  usageBaseUrl: string,
+  region: string
+): Promise<string | undefined> {
+  try {
+    const response = await fetch(usageBaseUrl, {
+      method: "POST",
+      headers: {
+        Authorization: `Bearer ${accessToken}`,
+        "Content-Type": "application/x-amz-json-1.0",
+        "x-amz-target": "AmazonCodeWhispererService.ListAvailableProfiles",
+        Accept: "application/json",
+      },
+      body: JSON.stringify({ maxResults: 10 }),
+      // Don't let a hung profile lookup block the usage/quota refresh indefinitely.
+      signal: AbortSignal.timeout(10000),
+    });
+    if (!response.ok) return undefined;
+
+    const data = toRecord(await response.json());
+    const profiles = Array.isArray(data.profiles) ? data.profiles : [];
+    const normalizedRegion = region.toLowerCase();
+    const matched =
+      profiles.find((profile: unknown) => {
+        const arn = toRecord(profile).arn;
+        return typeof arn === "string" && arn.toLowerCase().includes(`:${normalizedRegion}:`);
+      }) || profiles[0];
+    const arn = toRecord(matched).arn;
+    return typeof arn === "string" && arn.length > 0 ? arn : undefined;
+  } catch {
+    return undefined;
+  }
+}
+
 /**
  * Kiro (AWS CodeWhisperer) Usage
  */
 async function getKiroUsage(accessToken?: string, providerSpecificData?: JsonRecord) {
   try {
-    const profileArn = providerSpecificData?.profileArn;
-    if (!profileArn) {
-      return { message: "Kiro connected. Profile ARN not available for quota tracking." };
-    }
+    let profileArn =
+      typeof providerSpecificData?.profileArn === "string"
+        ? providerSpecificData.profileArn
+        : undefined;
 
     // Enterprise IAM Identity Center accounts are region-bound: the profileArn, token and
     // endpoint must all match the region. Derive the region from the stored region (preferred)
     // or the profileArn, then route to the regional Amazon Q endpoint (us-east-1 keeps the
     // legacy codewhisperer host; codewhisperer.{region} does not resolve for other regions).
-    const regionFromArn =
-      typeof profileArn === "string"
-        ? profileArn.toLowerCase().match(/^arn:aws:codewhisperer:([a-z0-9-]+):/)?.[1]
-        : undefined;
+    const regionFromArn = profileArn
+      ? profileArn.toLowerCase().match(/^arn:aws:codewhisperer:([a-z0-9-]+):/)?.[1]
+      : undefined;
     const region =
       (typeof providerSpecificData?.region === "string" &&
         providerSpecificData.region.trim().toLowerCase()) ||
@@ -3021,6 +3062,17 @@ async function getKiroUsage(accessToken?: string, providerSpecificData?: JsonRec
     const usageBaseUrl =
       region === "us-east-1" ? CODEWHISPERER_BASE_URL : `https://q.${region}.amazonaws.com`;
 
+    // IAM Identity Center logins and kiro-cli imports frequently don't persist a profileArn, which
+    // previously caused the quota card to show nothing ("0 used"). Discover it on demand from
+    // ListAvailableProfiles (region-matched) so usage still resolves for those accounts.
+    if (!profileArn && accessToken) {
+      profileArn = await discoverKiroProfileArn(accessToken, usageBaseUrl, region);
+    }
+
+    if (!profileArn) {
+      return { message: "Kiro connected. Profile ARN not available for quota tracking." };
+    }
+
     // Kiro uses AWS CodeWhisperer GetUsageLimits API
     const payload = {
       origin: "AI_EDITOR",

tests/unit/kiro-iam-profilearn-usage.test.ts

@@ -0,0 +1,100 @@
+import test from "node:test";
+import assert from "node:assert/strict";
+
+import {
+  buildKiroUsageResult,
+  discoverKiroProfileArn,
+} from "@omniroute/open-sse/services/usage.ts";
+
+// Real-world shape returned by GetUsageLimits for an AWS IAM Identity Center ("KIRO POWER")
+// account — the usage is reported under resourceType "CREDIT" (not "AGENTIC_REQUEST").
+const IAM_CREDIT_RESPONSE = {
+  daysUntilReset: 0,
+  nextDateReset: 1.782864e9,
+  subscriptionInfo: { subscriptionTitle: "KIRO POWER", type: "Q_DEVELOPER_STANDALONE_POWER" },
+  usageBreakdownList: [
+    {
+      currency: "USD",
+      currentUsage: 3670,
+      currentUsageWithPrecision: 3670.9,
+      displayName: "Credit",
+      resourceType: "CREDIT",
+      unit: "INVOCATIONS",
+      usageLimit: 10000,
+      usageLimitWithPrecision: 10000.0,
+    },
+  ],
+};
+
+test("buildKiroUsageResult parses the IAM CREDIT breakdown into non-zero usage", () => {
+  const result = buildKiroUsageResult(IAM_CREDIT_RESPONSE) as {
+    plan: string;
+    quotas: Record<string, { used: number; total: number; remaining: number }>;
+  };
+  assert.ok("quotas" in result, "must return quotas for a CREDIT breakdown");
+  assert.equal(result.plan, "KIRO POWER");
+  const credit = result.quotas.credit;
+  assert.ok(credit, "CREDIT resource should map to a 'credit' quota key");
+  assert.equal(credit.used, 3670.9);
+  assert.equal(credit.total, 10000);
+  assert.equal(credit.remaining, 10000 - 3670.9);
+});
+
+test("discoverKiroProfileArn prefers the region-matched profile ARN", async () => {
+  const originalFetch = globalThis.fetch;
+  globalThis.fetch = (async () =>
+    new Response(
+      JSON.stringify({
+        profiles: [
+          { arn: "arn:aws:codewhisperer:us-east-1:111111111111:profile/AAAA" },
+          { arn: "arn:aws:codewhisperer:eu-central-1:820374639727:profile/RX4VNUHGHGAQ" },
+        ],
+      }),
+      { status: 200, headers: { "Content-Type": "application/json" } }
+    )) as typeof fetch;
+  try {
+    const arn = await discoverKiroProfileArn(
+      "tok",
+      "https://q.eu-central-1.amazonaws.com",
+      "eu-central-1"
+    );
+    assert.equal(arn, "arn:aws:codewhisperer:eu-central-1:820374639727:profile/RX4VNUHGHGAQ");
+  } finally {
+    globalThis.fetch = originalFetch;
+  }
+});
+
+test("discoverKiroProfileArn falls back to the first profile when no region match", async () => {
+  const originalFetch = globalThis.fetch;
+  globalThis.fetch = (async () =>
+    new Response(
+      JSON.stringify({ profiles: [{ arn: "arn:aws:codewhisperer:us-east-1:1:profile/X" }] }),
+      { status: 200, headers: { "Content-Type": "application/json" } }
+    )) as typeof fetch;
+  try {
+    const arn = await discoverKiroProfileArn("tok", "https://q.eu-west-1.amazonaws.com", "eu-west-1");
+    assert.equal(arn, "arn:aws:codewhisperer:us-east-1:1:profile/X");
+  } finally {
+    globalThis.fetch = originalFetch;
+  }
+});
+
+test("discoverKiroProfileArn returns undefined for empty profiles or non-ok response", async () => {
+  const originalFetch = globalThis.fetch;
+  try {
+    globalThis.fetch = (async () =>
+      new Response(JSON.stringify({ profiles: [] }), { status: 200 })) as typeof fetch;
+    assert.equal(
+      await discoverKiroProfileArn("tok", "https://q.eu-central-1.amazonaws.com", "eu-central-1"),
+      undefined
+    );
+
+    globalThis.fetch = (async () => new Response("nope", { status: 403 })) as typeof fetch;
+    assert.equal(
+      await discoverKiroProfileArn("tok", "https://q.eu-central-1.amazonaws.com", "eu-central-1"),
+      undefined
+    );
+  } finally {
+    globalThis.fetch = originalFetch;
+  }
+});