Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion docs-site/src/content/docs/reference/adapters.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
19 changes: 17 additions & 2 deletions src/providers/kiro-usage.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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";
Expand Down Expand Up @@ -207,10 +213,19 @@ export async function fetchKiroUsageSnapshot(ctx: KiroUsageContext): Promise<Kir
*/
export async function kiroUsageContextForAccount(accountId: string): Promise<KiroUsageContext> {
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 } : {}),
};
Expand Down
57 changes: 57 additions & 0 deletions tests/providers/kiro/kiro-account-quota.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<string, unknown>;
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();
});
});
Loading