From 778c3aeab920d74ea8e774be5f6811a2a50d3f81 Mon Sep 17 00:00:00 2001 From: Diego Rodrigues de Sa e Souza <8016841+diegosouzapw@users.noreply.github.com> Date: Sat, 11 Jul 2026 04:25:47 -0300 Subject: [PATCH] fix(kiro): probe IdC region during profileArn discovery, cross-region (recovers #6099) (#6840) * fix(kiro): route Amazon Q runtime by profileArn region for cross-region IdC Enterprise AWS IAM Identity Center accounts whose IdC instance lives outside the two Amazon Q Developer profile regions (us-east-1 / eu-central-1) - e.g. eu-north-1 (Stockholm), start URL https://d-XXXX.awsapps.com/start - showed no limits and returned 502 on every request. Root cause: the backend used the IdC/OIDC token region (providerSpecificData.region, e.g. eu-north-1) for every CodeWhisperer runtime call, hitting q.eu-north-1.amazonaws.com - a host that does not exist as a Q Developer runtime endpoint. Per AWS docs ("Supported Regions for the Q Developer console and Q Developer profile"), the Q Developer *profile* (which produces the profileArn and hosts generateAssistantResponse / GetUsageLimits / ListAvailableModels / ListAvailableProfiles) is only hosted in us-east-1 and eu-central-1, regardless of the IdC region; "data is stored in the Region where you create the Amazon Q Developer profile." Fix (new open-sse/services/kiroRegion.ts) decouples the two regions: - providerSpecificData.region stays the IdC/OIDC region, used ONLY for oidc.{region}.amazonaws.com token mint/refresh. - The runtime region is derived from the profileArn (resolveKiroRuntimeRegion): profileArn region -> a valid stored profile region -> us-east-1. A stored IdC region that is not a Q profile region (eu-north-1) is ignored for runtime. - Profile discovery (discoverKiroProfileArnAcrossRegions) probes the Q profile regions (EU IdC -> eu-central-1 first) with the cross-region SSO token instead of q.{idcRegion}. Wired into: executors/kiro.ts (generateAssistantResponse targets the profile region), services/usage/kiro.ts (getKiroUsage multi-region discovery + profileArn runtime region so Limits resolves), services/kiroModels.ts (ListAvailableModels), and src/lib/oauth/providers/kiro.ts (login-time postExchange profile discovery). Adds tests/unit/kiro-idc-cross-region.test.ts (15 cases). All Kiro suites pass (60 tests). * fix(kiro): probe the IdC region too during profileArn discovery (any IdC region) Make profile discovery general for an IdC in ANY of the ~30 IdC-supported AWS regions (us-west-2, ap-southeast-2, me-central-1, af-south-1, ...), not just eu-north-1. buildKiroProfileDiscoveryRegions now probes the two documented Q Developer profile regions FIRST (us-east-1 / eu-central-1, EU-first for EMEA IdC regions to cut latency), then appends the IdC/stored region itself as a forward-compatible fallback: if AWS ever co-locates the profile with the IdC or expands the profile-region list, a same-region probe still finds it. Probing a region with no profile simply returns nothing and we fall through. The profileArn's own region remains authoritative for every runtime call (resolveKiroRuntimeRegion), so a newly-issued ARN in any region is honored automatically. Adds ap-southeast-2 (APAC) cross-region coverage and updates the discovery-order tests. --------- Co-authored-by: artickc --- open-sse/executors/kiro.ts | 40 ++-- open-sse/services/kiroModels.ts | 21 +- open-sse/services/kiroRegion.ts | 185 ++++++++++++++++ open-sse/services/usage/kiro.ts | 89 ++++---- src/lib/oauth/providers/kiro.ts | 51 +---- tests/unit/kiro-idc-cross-region.test.ts | 267 +++++++++++++++++++++++ 6 files changed, 536 insertions(+), 117 deletions(-) create mode 100644 open-sse/services/kiroRegion.ts create mode 100644 tests/unit/kiro-idc-cross-region.test.ts diff --git a/open-sse/executors/kiro.ts b/open-sse/executors/kiro.ts index d5623aacdd..3dab64c395 100644 --- a/open-sse/executors/kiro.ts +++ b/open-sse/executors/kiro.ts @@ -19,6 +19,7 @@ import { type KiroThinkingState, } from "./kiroThinking.ts"; import { ByteQueue, TEXT_ENCODER, parseEventFrame } from "./kiro/eventstream.ts"; +import { kiroRuntimeHost, resolveKiroRuntimeRegion } from "../services/kiroRegion.ts"; type JsonRecord = Record; @@ -152,35 +153,28 @@ function ensureKiroUsage(state: KiroStreamState) { } /** - * Resolve the AWS region for a Kiro/CodeWhisperer connection. Enterprise AWS IAM Identity - * Center accounts are region-bound: the access token, the Q Developer profile ARN and the - * runtime endpoint must all match the region the IdC instance lives in (e.g. eu-central-1). - * A request signed for one region is rejected by another ("bearer token is invalid"), and a - * regional profileArn sent to us-east-1 fails with "Improperly formed request". Falls back to - * the region embedded in the profileArn, then us-east-1 (the AWS Builder ID default). + * Resolve the RUNTIME AWS region for a Kiro/CodeWhisperer connection. + * + * The runtime region is the region of the Amazon Q Developer profile (embedded in the + * profileArn — always us-east-1 or eu-central-1), NOT the IAM Identity Center / OIDC token + * region. An enterprise IdC instance may live in eu-north-1 (or any region), but the Q Developer + * profile that serves generateAssistantResponse only exists in us-east-1 / eu-central-1, so a + * runtime call must target the profileArn's region — routing to q.{idcRegion}.amazonaws.com + * (a host that does not exist) is what caused "no limits + 502 on every request". Delegates to + * the shared resolver (profileArn region → valid stored profile region → us-east-1). The IdC + * token region is used only for oidc.{region} token mint/refresh, elsewhere. */ export function resolveKiroRegion( credentials: { providerSpecificData?: unknown } | null | undefined ): string { - const psd = (credentials?.providerSpecificData || {}) as Record; - const region = typeof psd.region === "string" ? psd.region.trim().toLowerCase() : ""; - if (region) return region; - const arn = typeof psd.profileArn === "string" ? psd.profileArn.toLowerCase() : ""; - const match = arn.match(/^arn:aws:codewhisperer:([a-z0-9-]+):/); - return match ? match[1] : "us-east-1"; + return resolveKiroRuntimeRegion( + (credentials?.providerSpecificData || {}) as { region?: unknown; profileArn?: unknown } + ); } -/** - * CodeWhisperer/Amazon Q runtime host for a region. us-east-1 keeps the legacy - * codewhisperer.us-east-1 host (AWS Builder ID); other regions use the regional Amazon Q - * endpoint q.{region}.amazonaws.com — codewhisperer.{region}.amazonaws.com does not resolve - * for non-us-east-1 regions. - */ -export function kiroRuntimeHost(region: string): string { - return region === "us-east-1" - ? "https://codewhisperer.us-east-1.amazonaws.com" - : `https://q.${region}.amazonaws.com`; -} +// Re-exported from the shared region module so existing importers (and tests) that pull +// kiroRuntimeHost from this executor keep working. +export { kiroRuntimeHost }; /** * KiroExecutor - Executor for Kiro AI (AWS CodeWhisperer) diff --git a/open-sse/services/kiroModels.ts b/open-sse/services/kiroModels.ts index 89680655d4..717e8bac66 100644 --- a/open-sse/services/kiroModels.ts +++ b/open-sse/services/kiroModels.ts @@ -27,6 +27,8 @@ import { createHash } from "node:crypto"; import { v4 as uuidv4 } from "uuid"; +import { resolveKiroRuntimeRegion } from "./kiroRegion.ts"; + type RawRecord = Record; const KIRO_RUNTIME_SDK_VERSION = "1.0.0"; @@ -185,20 +187,15 @@ function expandKiroModels(data: unknown): KiroModel[] { } /** - * Derive the AWS region for a Kiro connection. Mirrors getKiroUsage: prefer the - * stored region, then the region embedded in the profileArn, else us-east-1. + * Derive the RUNTIME AWS region for a Kiro connection's model discovery. Delegates to the shared + * resolver: the profileArn region wins (that is where the Q Developer profile + ListAvailableModels + * live — us-east-1 / eu-central-1), then a valid stored profile region, else us-east-1. The IdC + * token region (e.g. eu-north-1) is deliberately not used as a runtime region. */ export function resolveKiroRegion(providerSpecificData: unknown): string { - const psd = asRecord(providerSpecificData); - const explicit = toNonEmptyString(psd.region); - if (explicit) return explicit.toLowerCase(); - - const profileArn = toNonEmptyString(psd.profileArn); - const fromArn = profileArn - ? profileArn.toLowerCase().match(/^arn:aws:codewhisperer:([a-z0-9-]+):/)?.[1] - : undefined; - - return fromArn || "us-east-1"; + return resolveKiroRuntimeRegion( + asRecord(providerSpecificData) as { region?: unknown; profileArn?: unknown } + ); } /** diff --git a/open-sse/services/kiroRegion.ts b/open-sse/services/kiroRegion.ts new file mode 100644 index 0000000000..e87495cf33 --- /dev/null +++ b/open-sse/services/kiroRegion.ts @@ -0,0 +1,185 @@ +/** + * Shared Amazon Q Developer (Kiro / AWS CodeWhisperer) region resolution. + * + * TWO DISTINCT REGIONS — verified against the AWS docs "Amazon Q Developer Pro Region support" + * ("Supported Regions for the Q Developer console and Q Developer profile"): + * + * • IdC / OIDC / token region — `providerSpecificData.region`. May be ANY of the ~30 IdC- + * supported AWS regions (us-east-1, us-west-2, ca-central-1, sa-east-1, eu-west-1/2/3, + * eu-central-1/2, eu-north-1, eu-south-1/2, ap-south-1/2, ap-east-1/2, ap-southeast-1..7, + * ap-northeast-1/2/3, me-central-1, me-south-1, af-south-1, il-central-1, …). Used ONLY for + * `oidc.{region}.amazonaws.com` token mint/refresh (see tokenRefresh.ts / oauth providers). + * • Q Developer PROFILE / RUNTIME region — where the `profileArn` lives and every CodeWhisperer + * runtime call is served (generateAssistantResponse, GetUsageLimits, ListAvailableModels, + * ListAvailableProfiles). AWS currently hosts the profile ONLY in us-east-1 and eu-central-1, + * REGARDLESS of the IdC region ("Regardless of the IAM Identity Center Region, data is stored + * in the Region where you create the Amazon Q Developer profile"). The AWS docs' own example: + * an IdC in us-west-1 → profile in us-east-1. + * + * Consequences enforced here: + * • The RUNTIME region is the region embedded in the `profileArn` (authoritative — whatever + * region AWS actually hosts the profile in), NOT the IdC region. Routing a runtime call to + * `q.{idcRegion}.amazonaws.com` for a non-profile IdC region (e.g. q.eu-north-1, which does + * not exist as a Q Developer runtime endpoint) is the root cause of the "Kiro IAM shows no + * limits + every request returns 502" failure. + * • profileArn discovery works for an IdC in ANY region: it probes the known profile regions + * (us-east-1 / eu-central-1) with the cross-region SSO token, AND the IdC's own region as a + * forward-compatible fallback (in case AWS ever co-locates or expands profile regions). The + * discovered ARN's region then drives every runtime call. + */ + +// Canonical AWS region shape — kept local (identical to AWS_REGION_PATTERN in +// src/lib/oauth/constants/oauth.ts) so this open-sse module has no cross-tree import just to +// validate a string. Guards against SSRF via region injection (GHSA-6mwv-4mrm-5p3m): the value +// is interpolated into upstream URLs. +export const AWS_REGION_PATTERN = /^[a-z]{2}-[a-z]+-\d{1,2}$/; + +/** + * Regions where the Amazon Q Developer *profile* is currently hosted (AWS docs: "Supported + * Regions for the Q Developer console and Q Developer profile"). These are the guaranteed + * discovery targets and the only regions trusted as a runtime fallback when no profileArn is + * known. The profileArn's own region is always honored above this list, so a future AWS + * profile-region expansion works automatically once an ARN is discovered. + */ +export const KIRO_PROFILE_REGIONS = ["us-east-1", "eu-central-1"] as const; + +/** + * CodeWhisperer / Amazon Q runtime host for a region. us-east-1 keeps the legacy + * codewhisperer.us-east-1 host (AWS Builder ID home region); other regions use the regional + * Amazon Q endpoint `q.{region}.amazonaws.com` — codewhisperer.{region}.amazonaws.com does not + * resolve for non-us-east-1 regions. + */ +export function kiroRuntimeHost(region: string): string { + return region === "us-east-1" + ? "https://codewhisperer.us-east-1.amazonaws.com" + : `https://q.${region}.amazonaws.com`; +} + +/** Extract the region from a CodeWhisperer profile ARN (`arn:aws:codewhisperer:{region}:...`). */ +export function regionFromKiroProfileArn(profileArn?: string | null): string | undefined { + if (typeof profileArn !== "string") return undefined; + return profileArn.toLowerCase().match(/^arn:aws:codewhisperer:([a-z0-9-]+):/)?.[1]; +} + +function normalizeRegion(region: unknown): string { + return typeof region === "string" ? region.trim().toLowerCase() : ""; +} + +/** + * Resolve the RUNTIME region for CodeWhisperer / Amazon Q calls. + * + * Priority: + * 1. The region embedded in the `profileArn` — authoritative, this is where the Q Developer + * profile (and thus the runtime) actually lives. + * 2. A stored region ONLY when it is a valid Q Developer profile region (us-east-1 / + * eu-central-1). A stored IdC region that is not a Q profile region (e.g. eu-north-1) is + * deliberately IGNORED for runtime — it is a token/OIDC region, not a runtime region. + * 3. us-east-1 (CodeWhisperer home region) as the final fallback. + */ +export function resolveKiroRuntimeRegion( + providerSpecificData: { region?: unknown; profileArn?: unknown } | null | undefined +): string { + const fromArn = regionFromKiroProfileArn( + typeof providerSpecificData?.profileArn === "string" + ? providerSpecificData.profileArn + : undefined + ); + if (fromArn) return fromArn; + + const stored = normalizeRegion(providerSpecificData?.region); + if (stored && (KIRO_PROFILE_REGIONS as readonly string[]).includes(stored)) return stored; + + return "us-east-1"; +} + +/** + * Build the ordered list of regions to probe for `ListAvailableProfiles`. + * + * The Amazon Q Developer profile (and thus every runtime endpoint) is currently hosted only in + * KIRO_PROFILE_REGIONS (us-east-1 / eu-central-1) regardless of the IdC region, so those are + * probed FIRST — EU-first when the IdC region is in EMEA (eu-, af-, me-, il- prefixes) to + * minimize latency. The IdC/stored region is then appended as a forward-compatible fallback: if + * AWS ever co-locates the profile with the IdC, or expands the profile-region list, a same-region + * probe still finds it. It is only appended when it is a valid AWS region distinct from the known + * profile regions; probing a region with no profile simply returns nothing and we fall through. + * This makes discovery work for an IdC in ANY region (us-west-2, ap-southeast-2, me-central-1, + * af-south-1, …), not just eu-north-1. + */ +export function buildKiroProfileDiscoveryRegions(storedRegion?: string | null): string[] { + const stored = normalizeRegion(storedRegion); + const preferEu = /^(eu|af|me|il)-/.test(stored); + const regions: string[] = preferEu + ? ["eu-central-1", "us-east-1"] + : ["us-east-1", "eu-central-1"]; + + if (stored && AWS_REGION_PATTERN.test(stored) && !regions.includes(stored)) { + regions.push(stored); + } + return regions; +} + +async function listKiroProfileArnForRegion( + accessToken: string, + region: string, + fetchImpl: typeof fetch +): Promise { + // Defensive: region comes from a hardcoded allowlist here, but validate before it is + // interpolated into the runtime host (SSRF guard, GHSA-6mwv-4mrm-5p3m). + if (!AWS_REGION_PATTERN.test(region)) return undefined; + try { + const response = await fetchImpl(`${kiroRuntimeHost(region)}/`, { + method: "POST", + headers: { + "Content-Type": "application/x-amz-json-1.0", + Accept: "application/json", + "x-amz-target": "AmazonCodeWhispererService.ListAvailableProfiles", + Authorization: `Bearer ${accessToken}`, + }, + body: JSON.stringify({ maxResults: 10 }), + // Never let a hung/region-mismatched profile lookup block login or the quota refresh. + signal: AbortSignal.timeout(10000), + }); + if (!response.ok) return undefined; + + const data = (await response.json()) as { profiles?: unknown }; + const profiles = Array.isArray(data?.profiles) ? data.profiles : []; + // Prefer a profile whose ARN region matches the region we queried; else take the first. + const matched = + profiles.find((profile: unknown) => { + const arn = (profile as { arn?: unknown })?.arn; + return typeof arn === "string" && regionFromKiroProfileArn(arn) === region; + }) || profiles[0]; + const arn = (matched as { arn?: unknown })?.arn; + return typeof arn === "string" && arn.length > 0 ? arn : undefined; + } catch { + return undefined; + } +} + +/** + * Discover a Kiro/CodeWhisperer profile ARN by probing the Q Developer profile regions + * (us-east-1 / eu-central-1) AND the IdC/stored region with the account's access token. The SSO + * bearer token minted from the IdC region works cross-region against the Q Developer profile's + * region (AWS's documented multi-region IdC ⇄ profile setup), so an IdC in ANY region resolves. + * Returns the first ARN found (its embedded region is the authoritative runtime region), or + * undefined when no profile is available (e.g. AWS Builder ID accounts, or an org/token with no + * Kiro entitlement). Best-effort: never throws. + */ +export async function discoverKiroProfileArnAcrossRegions( + accessToken: string | null | undefined, + storedRegion?: string | null, + fetchImpl?: typeof fetch +): Promise { + const token = typeof accessToken === "string" ? accessToken.trim() : ""; + if (!token) return undefined; + + // Resolve fetch at call time (not module-load) so callers/tests that swap globalThis.fetch + // are honored when no explicit implementation is injected. + const doFetch = fetchImpl ?? globalThis.fetch; + + for (const region of buildKiroProfileDiscoveryRegions(storedRegion)) { + const arn = await listKiroProfileArnForRegion(token, region, doFetch); + if (arn) return arn; + } + return undefined; +} diff --git a/open-sse/services/usage/kiro.ts b/open-sse/services/usage/kiro.ts index b96b5041ca..6ed5f72ed6 100644 --- a/open-sse/services/usage/kiro.ts +++ b/open-sse/services/usage/kiro.ts @@ -13,6 +13,11 @@ import { toRecord, toNumber } from "./scalars.ts"; import { type UsageQuota, parseResetTime } from "./quota.ts"; +import { + discoverKiroProfileArnAcrossRegions, + kiroRuntimeHost, + resolveKiroRuntimeRegion, +} from "../kiroRegion.ts"; import { isExternalIdpAuthMethod, KIRO_EXTERNAL_IDP_TOKEN_TYPE_HEADER, @@ -158,10 +163,20 @@ export async function discoverKiroProfileArn( } /** - * The three GetUsageLimits attempts (regional GET, CodeWhisperer POST, Q GET) tried in + * The three GetUsageLimits attempts (CodeWhisperer POST, regional GET, Q GET) tried in * order by getKiroUsage — extracted so the auth-method header variants (api_key * `tokentype`, external_idp `TokenType`) stay in one authHeaders object and the parent * function stays under the function-length gate. + * + * The POST variant (x-amz-json-1.0 + `x-amz-target: ...GetUsageLimits`) is tried FIRST: it + * is the canonical AWS JSON-RPC shape used everywhere else in the Kiro integration + * (discoverKiroProfileArn's ListAvailableProfiles call, kiroModels.ts fingerprinting) and, + * critically, it is the only variant whose URL is built from the RUNTIME-region-resolved + * `usageBaseUrl` (see resolveKiroRuntimeRegion in getKiroUsage) with the profileArn in the + * payload — required for cross-region IAM Identity Center accounts (#6099) where the profile + * lives in a different region than the IdC/token region. The two GET variants are + * best-effort fallbacks (added for #6587 API-key auth) for accounts/regions where the POST + * shape is rejected. */ function buildKiroUsageAttempts(opts: { authHeaders: Record; @@ -173,18 +188,6 @@ function buildKiroUsageAttempts(opts: { }): Array<{ name: string; run: () => Promise }> { const { authHeaders, usageParams, qParams, payload, usageBaseUrl, qBaseUrl } = opts; return [ - { - name: "codewhisperer-get", - run: () => - fetch(`${CODEWHISPERER_BASE_URL}/getUsageLimits?${usageParams.toString()}`, { - method: "GET", - headers: { - ...authHeaders, - "x-amz-user-agent": "aws-sdk-js/1.0.0 KiroIDE", - "user-agent": "aws-sdk-js/1.0.0 KiroIDE", - }, - }), - }, { name: "codewhisperer-post", run: () => @@ -198,6 +201,18 @@ function buildKiroUsageAttempts(opts: { body: JSON.stringify(payload), }), }, + { + name: "codewhisperer-get", + run: () => + fetch(`${CODEWHISPERER_BASE_URL}/getUsageLimits?${usageParams.toString()}`, { + method: "GET", + headers: { + ...authHeaders, + "x-amz-user-agent": "aws-sdk-js/1.0.0 KiroIDE", + "user-agent": "aws-sdk-js/1.0.0 KiroIDE", + }, + }), + }, { name: "q-get", run: () => @@ -209,27 +224,6 @@ function buildKiroUsageAttempts(opts: { ]; } -/** - * 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). - */ -function resolveKiroUsageEndpoints(providerSpecificData?: JsonRecord, profileArn?: string) { - const regionFromArn = profileArn - ? profileArn.toLowerCase().match(/^arn:aws:codewhisperer:([a-z0-9-]+):/)?.[1] - : undefined; - const region = - (typeof providerSpecificData?.region === "string" && - providerSpecificData.region.trim().toLowerCase()) || - regionFromArn || - "us-east-1"; - const usageBaseUrl = - region === "us-east-1" ? CODEWHISPERER_BASE_URL : `https://q.${region}.amazonaws.com`; - const qBaseUrl = `https://q.${region}.amazonaws.com`; - return { region, usageBaseUrl, qBaseUrl }; -} - /** * Base auth headers for the usage endpoints, per auth method: long-lived API keys add * `tokentype: API_KEY`; enterprise / Microsoft Entra (external_idp) org accounts require @@ -308,22 +302,31 @@ export async function getKiroUsage(accessToken?: string, providerSpecificData?: ? providerSpecificData.profileArn : undefined; - const { region, usageBaseUrl, qBaseUrl } = resolveKiroUsageEndpoints( - providerSpecificData, - profileArn - ); + const storedRegion = + typeof providerSpecificData?.region === "string" + ? providerSpecificData.region.trim().toLowerCase() + : undefined; - // 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. + // Enterprise IAM Identity Center logins and kiro-cli imports frequently don't persist a + // profileArn. Discover it by probing the Q Developer PROFILE regions (us-east-1 / eu-central-1) + // — NOT the IdC/token region. An IdC in eu-north-1 has no Q runtime host (q.eu-north-1 does not + // exist); its profile lives in eu-central-1 (or us-east-1) and the SSO token works cross-region + // against it. Without this, the quota card previously showed nothing ("no limits") for such + // accounts because the single-region lookup at q.{idcRegion} always failed. if (!profileArn && accessToken) { - profileArn = await discoverKiroProfileArn(accessToken, usageBaseUrl, region, authMethod); + profileArn = await discoverKiroProfileArnAcrossRegions(accessToken, storedRegion); } if (!profileArn && !isApiKey) { return { message: "Kiro connected. Profile ARN not available for quota tracking." }; } + // The RUNTIME region is the profileArn region (us-east-1 / eu-central-1), never the IdC token + // region. Route GetUsageLimits to that region's host so quota resolves for cross-region IdC. + const region = resolveKiroRuntimeRegion({ region: storedRegion, profileArn }); + const usageBaseUrl = region === "us-east-1" ? CODEWHISPERER_BASE_URL : kiroRuntimeHost(region); + const qBaseUrl = `https://q.${region}.amazonaws.com`; + const authHeaders = buildKiroAuthHeaders(accessToken, isApiKey, providerSpecificData); const usageParams = new URLSearchParams({ @@ -342,7 +345,7 @@ export async function getKiroUsage(accessToken?: string, providerSpecificData?: resourceType: "AGENTIC_REQUEST", }; -const attempts = buildKiroUsageAttempts({ + const attempts = buildKiroUsageAttempts({ authHeaders, usageParams, qParams, diff --git a/src/lib/oauth/providers/kiro.ts b/src/lib/oauth/providers/kiro.ts index 2efb7fc609..f90bbae0f0 100644 --- a/src/lib/oauth/providers/kiro.ts +++ b/src/lib/oauth/providers/kiro.ts @@ -1,4 +1,5 @@ import { KIRO_CONFIG, AWS_REGION_PATTERN, assertValidAwsRegion } from "../constants/oauth"; +import { discoverKiroProfileArnAcrossRegions } from "@omniroute/open-sse/services/kiroRegion.ts"; export const kiro = { config: KIRO_CONFIG, @@ -8,9 +9,7 @@ export const kiro = { const candidateRegion = regionMatch?.[1] || "us-east-1"; // Region is sourced from KIRO_CONFIG.tokenUrl (trusted constant) but defensively // re-validate before letting it influence later fetches (GHSA-6mwv-4mrm-5p3m). - const resolvedRegion = AWS_REGION_PATTERN.test(candidateRegion) - ? candidateRegion - : "us-east-1"; + const resolvedRegion = AWS_REGION_PATTERN.test(candidateRegion) ? candidateRegion : "us-east-1"; const registerPayload: { clientName: string; clientType: string; @@ -131,45 +130,19 @@ export const kiro = { }, // Enterprise IAM Identity Center accounts require a region-bound Q Developer profileArn on every // CodeWhisperer call; without it AWS returns 403 "User is not authorized to make this call". The - // device-code flow does not return one, so discover it here via ListAvailableProfiles against the - // same regional endpoint the token was issued for. Best-effort: AWS Builder ID accounts have no - // profile and this simply yields none; failures never block login. + // device-code flow does not return one, so discover it here via ListAvailableProfiles. + // + // The IdC/token region (`_region`, e.g. eu-north-1) is NOT where the Q Developer profile lives — + // AWS only hosts the profile (and its runtime) in us-east-1 / eu-central-1. So probe those + // profile regions with the freshly-minted SSO token (which works cross-region against the + // profile's home region), NOT q.{idcRegion} which does not resolve. Best-effort: AWS Builder ID + // accounts have no profile and this simply yields none; failures never block login. postExchange: async (tokenData) => { const accessToken = tokenData?.access_token; if (!accessToken) return null; - const region = String(tokenData?._region || "us-east-1").toLowerCase(); - // Defensive: tokenData._region came from upstream JSON or extraData - // and is interpolated into the runtime host below (GHSA-6mwv-4mrm-5p3m). - if (!AWS_REGION_PATTERN.test(region)) return null; - const runtimeHost = - region === "us-east-1" - ? "https://codewhisperer.us-east-1.amazonaws.com" - : `https://q.${region}.amazonaws.com`; - try { - const profRes = await fetch(`${runtimeHost}/`, { - method: "POST", - headers: { - "Content-Type": "application/x-amz-json-1.0", - Accept: "application/json", - "x-amz-target": "AmazonCodeWhispererService.ListAvailableProfiles", - Authorization: `Bearer ${accessToken}`, - }, - body: JSON.stringify({ maxResults: 10 }), - signal: AbortSignal.timeout(10000), - }); - if (!profRes.ok) return null; - const profData = await profRes.json(); - const profiles = Array.isArray(profData?.profiles) ? profData.profiles : []; - const matched = - profiles.find( - (p: { arn?: string }) => typeof p?.arn === "string" && p.arn.includes(`:${region}:`) - ) || profiles[0]; - const arn = matched && typeof matched.arn === "string" ? matched.arn : undefined; - return arn ? { profileArn: arn } : null; - } catch { - // Best-effort profile discovery — never block login on it. - return null; - } + const storedRegion = typeof tokenData?._region === "string" ? tokenData._region : undefined; + const arn = await discoverKiroProfileArnAcrossRegions(accessToken, storedRegion); + return arn ? { profileArn: arn } : null; }, mapTokens: (tokens, extra) => ({ accessToken: tokens.access_token, diff --git a/tests/unit/kiro-idc-cross-region.test.ts b/tests/unit/kiro-idc-cross-region.test.ts new file mode 100644 index 0000000000..6e34c44bb9 --- /dev/null +++ b/tests/unit/kiro-idc-cross-region.test.ts @@ -0,0 +1,267 @@ +/** + * Regression: Kiro enterprise IAM Identity Center accounts whose IdC instance lives OUTSIDE the + * two Amazon Q Developer profile regions (us-east-1 / eu-central-1) — e.g. eu-north-1 (Stockholm), + * start URL https://d-XXXX.awsapps.com/start. + * + * Root cause fixed here: the backend used the IdC/OIDC token region (eu-north-1) for every + * CodeWhisperer runtime call, hitting q.eu-north-1.amazonaws.com — a host that does not exist as a + * Q Developer runtime endpoint. Result: profileArn discovery failed (Limits showed nothing) and + * generateAssistantResponse failed (every request 502). AWS hosts the Q Developer PROFILE (and its + * runtime) only in us-east-1 / eu-central-1, regardless of the IdC region. + * + * The fix: the RUNTIME region is derived from the profileArn (us-east-1 / eu-central-1), and + * profileArn discovery probes those profile regions with the cross-region SSO token — never the + * IdC region. + */ + +import test from "node:test"; +import assert from "node:assert/strict"; + +import { + resolveKiroRuntimeRegion, + buildKiroProfileDiscoveryRegions, + discoverKiroProfileArnAcrossRegions, + kiroRuntimeHost, + regionFromKiroProfileArn, + KIRO_PROFILE_REGIONS, +} from "../../open-sse/services/kiroRegion.ts"; +import { resolveKiroRegion as resolveExecutorRegion } from "../../open-sse/executors/kiro.ts"; +import { kiro } from "@/lib/oauth/providers/kiro"; +import { __testing } from "@omniroute/open-sse/services/usage.ts"; + +const { getKiroUsage } = __testing; + +const EU_CENTRAL_ARN = "arn:aws:codewhisperer:eu-central-1:820374639727:profile/RX4VNUHGHGAQ"; + +test("KIRO_PROFILE_REGIONS is exactly us-east-1 and eu-central-1", () => { + assert.deepEqual([...KIRO_PROFILE_REGIONS], ["us-east-1", "eu-central-1"]); +}); + +test("regionFromKiroProfileArn extracts the region from a CodeWhisperer ARN", () => { + assert.equal(regionFromKiroProfileArn(EU_CENTRAL_ARN), "eu-central-1"); + assert.equal( + regionFromKiroProfileArn("arn:aws:codewhisperer:us-east-1:1:profile/X"), + "us-east-1" + ); + assert.equal(regionFromKiroProfileArn(undefined), undefined); + assert.equal(regionFromKiroProfileArn("not-an-arn"), undefined); +}); + +test("resolveKiroRuntimeRegion: profileArn region beats the IdC (eu-north-1) stored region", () => { + // The exact failing scenario: IdC token region eu-north-1, profile in eu-central-1. + assert.equal( + resolveKiroRuntimeRegion({ region: "eu-north-1", profileArn: EU_CENTRAL_ARN }), + "eu-central-1" + ); +}); + +test("resolveKiroRuntimeRegion: an IdC region that is not a Q profile region is ignored for runtime", () => { + // No profileArn yet, IdC region eu-north-1 → must NOT route to q.eu-north-1; fall back to us-east-1. + assert.equal(resolveKiroRuntimeRegion({ region: "eu-north-1" }), "us-east-1"); + assert.equal(resolveKiroRuntimeRegion({ region: "us-west-1" }), "us-east-1"); +}); + +test("resolveKiroRuntimeRegion: a valid stored profile region is honored, defaults to us-east-1", () => { + assert.equal(resolveKiroRuntimeRegion({ region: "eu-central-1" }), "eu-central-1"); + assert.equal(resolveKiroRuntimeRegion({ region: "us-east-1" }), "us-east-1"); + assert.equal(resolveKiroRuntimeRegion({}), "us-east-1"); + assert.equal(resolveKiroRuntimeRegion(null), "us-east-1"); +}); + +test("the executor's resolveKiroRegion routes an eu-north-1 IdC account to the profile region", () => { + assert.equal( + resolveExecutorRegion({ + providerSpecificData: { region: "eu-north-1", profileArn: EU_CENTRAL_ARN }, + }), + "eu-central-1" + ); + // generateAssistantResponse must therefore target the real Q host, not q.eu-north-1. + assert.equal(kiroRuntimeHost("eu-central-1"), "https://q.eu-central-1.amazonaws.com"); +}); + +test("buildKiroProfileDiscoveryRegions: EU IdC probes the profile regions first, then the IdC region", () => { + const regions = buildKiroProfileDiscoveryRegions("eu-north-1"); + assert.deepEqual(regions, ["eu-central-1", "us-east-1", "eu-north-1"]); + // The profile regions (fast path) are tried BEFORE the IdC-region fallback. + assert.ok(regions.indexOf("eu-central-1") < regions.indexOf("eu-north-1")); + assert.ok(regions.indexOf("us-east-1") < regions.indexOf("eu-north-1")); + // Another EMEA IdC region → still EU-first, IdC region appended as fallback. + assert.deepEqual(buildKiroProfileDiscoveryRegions("me-central-1"), [ + "eu-central-1", + "us-east-1", + "me-central-1", + ]); +}); + +test("buildKiroProfileDiscoveryRegions: non-EU IdC probes us-east-1 first, then the IdC region", () => { + assert.deepEqual(buildKiroProfileDiscoveryRegions("us-west-2"), [ + "us-east-1", + "eu-central-1", + "us-west-2", + ]); + assert.deepEqual(buildKiroProfileDiscoveryRegions("ap-southeast-2"), [ + "us-east-1", + "eu-central-1", + "ap-southeast-2", + ]); + // No stored region → just the two profile regions. + assert.deepEqual(buildKiroProfileDiscoveryRegions(undefined), ["us-east-1", "eu-central-1"]); +}); + +test("buildKiroProfileDiscoveryRegions: a stored profile region is probed first", () => { + assert.deepEqual(buildKiroProfileDiscoveryRegions("eu-central-1"), ["eu-central-1", "us-east-1"]); + assert.deepEqual(buildKiroProfileDiscoveryRegions("us-east-1"), ["us-east-1", "eu-central-1"]); +}); + +test("discoverKiroProfileArnAcrossRegions: eu-north-1 IdC finds the eu-central-1 profile, skips q.eu-north-1", async () => { + const requested: string[] = []; + const fetchImpl = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + // Simulate reality: q.eu-north-1 does not exist (network failure); eu-central-1 hosts the profile. + if (url.includes("eu-north-1")) throw new Error("ENOTFOUND q.eu-north-1.amazonaws.com"); + if (url.includes("eu-central-1")) { + return new Response(JSON.stringify({ profiles: [{ arn: EU_CENTRAL_ARN }] }), { status: 200 }); + } + // us-east-1 has no profile for this identity. + return new Response(JSON.stringify({ profiles: [] }), { status: 200 }); + }) as unknown as typeof fetch; + + const arn = await discoverKiroProfileArnAcrossRegions("sso-token", "eu-north-1", fetchImpl); + assert.equal(arn, EU_CENTRAL_ARN); + assert.ok( + requested.every((u) => !u.includes("eu-north-1")), + `must never probe q.eu-north-1, got: ${JSON.stringify(requested)}` + ); + assert.ok( + requested.some((u) => u.startsWith("https://q.eu-central-1.amazonaws.com/")), + "must probe the eu-central-1 Q Developer host" + ); +}); + +test("discoverKiroProfileArnAcrossRegions: a non-EU (ap-southeast-2) IdC resolves a us-east-1 profile", async () => { + // Proves the fix is general, not eu-north-1-specific: an APAC IdC's profile lives in a Q + // profile region (us-east-1 here) and is found via the cross-region SSO token. + const US_EAST_ARN = "arn:aws:codewhisperer:us-east-1:111111111111:profile/APAC"; + const requested: string[] = []; + const fetchImpl = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + if (url.includes("us-east-1")) { + return new Response(JSON.stringify({ profiles: [{ arn: US_EAST_ARN }] }), { status: 200 }); + } + return new Response(JSON.stringify({ profiles: [] }), { status: 200 }); + }) as unknown as typeof fetch; + + const arn = await discoverKiroProfileArnAcrossRegions("sso-token", "ap-southeast-2", fetchImpl); + assert.equal(arn, US_EAST_ARN); + assert.equal( + resolveKiroRuntimeRegion({ region: "ap-southeast-2", profileArn: arn }), + "us-east-1" + ); + // The us-east-1 profile region is probed before the ap-southeast-2 IdC-region fallback. + assert.ok(requested.some((u) => u.startsWith("https://codewhisperer.us-east-1.amazonaws.com/"))); +}); + +test("discoverKiroProfileArnAcrossRegions: no token / no profile yields undefined without throwing", async () => { + assert.equal(await discoverKiroProfileArnAcrossRegions("", "eu-north-1"), undefined); + const emptyFetch = (async () => + new Response(JSON.stringify({ profiles: [] }), { status: 200 })) as unknown as typeof fetch; + assert.equal( + await discoverKiroProfileArnAcrossRegions("tok", "eu-north-1", emptyFetch), + undefined + ); +}); + +test("kiro.postExchange (login) discovers the eu-central-1 profile for an eu-north-1 IdC token", async () => { + const originalFetch = global.fetch; + const requested: string[] = []; + global.fetch = (async (input: RequestInfo | URL) => { + const url = String(input); + requested.push(url); + if (url.includes("eu-north-1")) throw new Error("ENOTFOUND"); + if (url.includes("eu-central-1")) { + return new Response(JSON.stringify({ profiles: [{ arn: EU_CENTRAL_ARN }] }), { status: 200 }); + } + return new Response(JSON.stringify({ profiles: [] }), { status: 200 }); + }) as typeof fetch; + + try { + const extra = await kiro.postExchange({ access_token: "sso-token", _region: "eu-north-1" }); + assert.deepEqual(extra, { profileArn: EU_CENTRAL_ARN }); + assert.ok(requested.every((u) => !u.includes("eu-north-1"))); + } finally { + global.fetch = originalFetch; + } +}); + +test("kiro.mapTokens keeps region=eu-north-1 (for OIDC refresh) AND stores the eu-central-1 profileArn", () => { + const mapped = kiro.mapTokens( + { access_token: "at", refresh_token: "rt", expires_in: 3600, _region: "eu-north-1" }, + { profileArn: EU_CENTRAL_ARN } + ); + // region stays the IdC/OIDC region so token refresh hits oidc.eu-north-1.amazonaws.com … + assert.equal(mapped.providerSpecificData.region, "eu-north-1"); + // … while the profileArn carries the eu-central-1 runtime region for CodeWhisperer calls. + assert.equal(mapped.providerSpecificData.profileArn, EU_CENTRAL_ARN); +}); + +test("getKiroUsage: eu-north-1 IdC account resolves quota via the eu-central-1 profile", async () => { + const originalFetch = globalThis.fetch; + const requested: string[] = []; + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const url = String(input); + const target = String( + (init?.headers as Record | undefined)?.["x-amz-target"] || "" + ); + requested.push(`${target} ${url}`); + // q.eu-north-1 must never be contacted. + if (url.includes("eu-north-1")) throw new Error("ENOTFOUND q.eu-north-1"); + if (target.endsWith("ListAvailableProfiles")) { + if (url.includes("eu-central-1")) { + return new Response(JSON.stringify({ profiles: [{ arn: EU_CENTRAL_ARN }] }), { + status: 200, + }); + } + return new Response(JSON.stringify({ profiles: [] }), { status: 200 }); + } + // GetUsageLimits at the eu-central-1 host → real IAM CREDIT breakdown. + return new Response( + JSON.stringify({ + subscriptionInfo: { subscriptionTitle: "KIRO POWER" }, + usageBreakdownList: [ + { + resourceType: "CREDIT", + currentUsageWithPrecision: 12, + usageLimitWithPrecision: 1000, + }, + ], + }), + { status: 200, headers: { "Content-Type": "application/json" } } + ); + }) as typeof fetch; + + try { + // No persisted profileArn (the broken state), IdC region eu-north-1. + const result = (await getKiroUsage("sso-token", { + authMethod: "idc", + region: "eu-north-1", + })) as { + plan?: string; + quotas?: Record; + message?: string; + }; + assert.ok(result.quotas, `expected quotas, got: ${JSON.stringify(result)}`); + assert.equal(result.plan, "KIRO POWER"); + assert.equal(result.quotas!.credit.used, 12); + assert.equal(result.quotas!.credit.total, 1000); + // GetUsageLimits must have gone to the eu-central-1 runtime host, never q.eu-north-1. + assert.ok( + requested.some((r) => r.includes("GetUsageLimits") && r.includes("eu-central-1")), + `GetUsageLimits should hit eu-central-1, got: ${JSON.stringify(requested)}` + ); + assert.ok(requested.every((r) => !r.includes("eu-north-1"))); + } finally { + globalThis.fetch = originalFetch; + } +});