diff --git a/docs-site/src/content/docs/reference/adapters.md b/docs-site/src/content/docs/reference/adapters.md index f00d96edc08..6b6c73740e6 100644 --- a/docs-site/src/content/docs/reference/adapters.md +++ b/docs-site/src/content/docs/reference/adapters.md @@ -344,7 +344,9 @@ MiMo model Command Code serves. - Reports per-account usage. `AmazonCodeWhispererService.GetUsageLimits` on `https://management.{region}.kiro.dev/` returns the plan allowance, which becomes the monthly quota window for that account; a free-trial balance is reported as its own window. - The region comes from the account's profile ARN, then its stored API/SSO region. An + The region comes from the account's profile ARN, then its stored API/SSO region. An AWS + Builder ID account has no profile ARN of its own, so the probe sends the same Builder ID + service profile that generation requests use; that fixed ARN never selects the region. An unreadable or unrecognised response is reported as unknown rather than as zero usage, and an account whose overage is enabled is not treated as exhausted merely for passing its limit. The operation is undocumented by AWS, so treat the numbers as best-effort. diff --git a/src/providers/kiro-usage.ts b/src/providers/kiro-usage.ts index 46fa8279604..e5005d83fd3 100644 --- a/src/providers/kiro-usage.ts +++ b/src/providers/kiro-usage.ts @@ -12,6 +12,7 @@ * the pool needs those two answers to decide how long to cool a 429'd account. */ import { getValidAccessSnapshotForAccount } from "../oauth"; +import { resolveKiroRequestProfile } from "../oauth/kiro"; import type { ProviderQuota, ProviderQuotaWindow } from "./quota-types"; import { ACCOUNT_QUOTA_TTL_MS, @@ -46,6 +47,8 @@ export interface KiroUsageContext { accountId: string; access: string; profileArn?: string; + /** The ARN is the Builder ID service profile, not the account's; it must not pick the region. */ + builderIdFallback?: boolean; apiRegion?: string; ssoRegion?: string; } @@ -82,7 +85,10 @@ function safeRegion(value: string | undefined): string | undefined { * the SSO session that minted the token. */ function usageRegion(ctx: KiroUsageContext): string { - return safeRegion(ctx.profileArn?.split(":")[3]) + // The Builder ID service profile is Amazon's fixed us-east-1 ARN, not the account's own, so it + // must not pin the region, as kiro-constants.ts requires for the runtime path. + const arnRegion = ctx.builderIdFallback ? undefined : ctx.profileArn?.split(":")[3]; + return safeRegion(arnRegion) ?? safeRegion(ctx.apiRegion) ?? safeRegion(ctx.ssoRegion) ?? "us-east-1"; @@ -207,10 +213,19 @@ export async function fetchKiroUsageSnapshot(ctx: KiroUsageContext): Promise { const snapshot = await getValidAccessSnapshotForAccount("kiro", accountId); + // Builder ID accounts never get an account-scoped ARN, and GetUsageLimits rejects a missing one + // with 400 "Invalid profileArn". Ask the same resolver the runtime path uses, so the usage probe + // sends exactly the ARN a generation request would. An account object is always passed, so the + // accountless env/local-import fallbacks never apply to a pooled account. + const profile = resolveKiroRequestProfile({ + profileArn: snapshot.kiro?.profileArn, + authType: snapshot.kiro?.authType, + }); return { accountId, access: snapshot.accessToken, - ...(snapshot.kiro?.profileArn ? { profileArn: snapshot.kiro.profileArn } : {}), + ...(profile.profileArn ? { profileArn: profile.profileArn } : {}), + ...(profile.builderIdFallback ? { builderIdFallback: true } : {}), ...(snapshot.kiro?.apiRegion ? { apiRegion: snapshot.kiro.apiRegion } : {}), ...(snapshot.kiro?.ssoRegion ? { ssoRegion: snapshot.kiro.ssoRegion } : {}), }; diff --git a/tests/providers/kiro/kiro-account-quota.test.ts b/tests/providers/kiro/kiro-account-quota.test.ts index f5a869c20d8..ca4314e0c1c 100644 --- a/tests/providers/kiro/kiro-account-quota.test.ts +++ b/tests/providers/kiro/kiro-account-quota.test.ts @@ -178,3 +178,60 @@ describe("Kiro per-account quota", () => { expect(getKiroAccountExhaustion(key)?.exhausted).toBe(false); }); }); + +describe("Kiro Builder ID usage probe", () => { + const BUILDER_ID_ARN = "arn:aws:codewhisperer:us-east-1:638616132270:profile/AAAACCCCXXXX"; + + function captureProbe(): Array<{ host: string; arn: string | null; bodyArn: unknown }> { + const seen: Array<{ host: string; arn: string | null; bodyArn: unknown }> = []; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const request = new Request(input instanceof Request ? input : String(input), init); + const url = new URL(request.url); + const body = await request.json() as Record; + seen.push({ host: url.host, arn: url.searchParams.get("profileArn"), bodyArn: body.profileArn }); + // Builder ID accounts answer with a CREDIT-only breakdown (observed live), so this also pins + // the parser's fallback past AGENTIC_REQUEST. + return new Response(JSON.stringify({ + usageBreakdownList: [{ + resourceType: "CREDIT", + currentUsageWithPrecision: 25.21, + usageLimitWithPrecision: 5000, + unit: "CREDITS", + }], + overageConfiguration: { overageStatus: "DISABLED" }, + nextDateReset: Math.floor(Date.now() / 1000) + 3 * 24 * 3600, + }), { status: 200, headers: { "content-type": "application/json" } }); + }) as typeof fetch; + return seen; + } + + test("a Builder ID account without a stored ARN sends the service profile and reports usage", async () => { + // Builder ID credentials carry the OIDC client pair and never an account-scoped ARN. + await saveCredential("kiro", { + access: "token-bid", refresh: "refresh-bid", expires: Date.now() + 60 * 60_000, + accountId: "kiro-bid", email: "bid@example.com", + kiro: { apiRegion: "us-west-2", ssoRegion: "us-west-2", clientId: "client-id", clientSecret: "client-secret" }, + }); + const seen = captureProbe(); + const rows = await fetchProviderAccountQuotas("kiro"); + expect(seen).toHaveLength(1); + expect(seen[0]!.arn).toBe(BUILDER_ID_ARN); + expect(seen[0]!.bodyArn).toBe(BUILDER_ID_ARN); + // The fixed ARN is Amazon's us-east-1 profile; it must not pin the account's own region. + expect(seen[0]!.host).toBe("management.us-west-2.kiro.dev"); + expect(rows[0]!.quota?.monthlyPercent).toBeCloseTo(0.5042, 4); + }); + + test("a non-Builder-ID account without a stored ARN still sends none", async () => { + await saveCredential("kiro", { + access: "token-desk", refresh: "refresh-desk", expires: Date.now() + 60 * 60_000, + accountId: "kiro-desk", email: "desk@example.com", + kiro: { apiRegion: "us-east-1", ssoRegion: "us-east-1" }, + }); + const seen = captureProbe(); + await fetchProviderAccountQuotas("kiro"); + expect(seen).toHaveLength(1); + expect(seen[0]!.arn).toBeNull(); + expect(seen[0]!.bodyArn).toBeUndefined(); + }); +});