From e12fb5351b1160bb7df3e6145da368dd31315112 Mon Sep 17 00:00:00 2001 From: NOXX - Commiter Date: Fri, 12 Jun 2026 15:45:37 +0300 Subject: [PATCH] fix(kiro): resolve quota for IAM Identity Center accounts missing a profileArn (#3722) Integrated into release/v3.8.23 --- open-sse/services/usage.ts | 68 +++++++++++-- tests/unit/kiro-iam-profilearn-usage.test.ts | 100 +++++++++++++++++++ 2 files changed, 160 insertions(+), 8 deletions(-) create mode 100644 tests/unit/kiro-iam-profilearn-usage.test.ts diff --git a/open-sse/services/usage.ts b/open-sse/services/usage.ts index d33cbcea4a..0493f7e65d 100644 --- a/open-sse/services/usage.ts +++ b/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 { + 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", diff --git a/tests/unit/kiro-iam-profilearn-usage.test.ts b/tests/unit/kiro-iam-profilearn-usage.test.ts new file mode 100644 index 0000000000..7c98eb9c7f --- /dev/null +++ b/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; + }; + 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; + } +});