From f2370991ece0a898685019e6a22028033bb91e53 Mon Sep 17 00:00:00 2001 From: artickc Date: Fri, 3 Jul 2026 16:22:06 +0300 Subject: [PATCH] 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). --- open-sse/executors/kiro.ts | 40 ++-- open-sse/services/kiroModels.ts | 21 +-- open-sse/services/kiroRegion.ts | 177 ++++++++++++++++++ open-sse/services/usage/kiro.ts | 39 ++-- src/lib/oauth/providers/kiro.ts | 51 ++--- tests/unit/kiro-idc-cross-region.test.ts | 225 +++++++++++++++++++++++ 6 files changed, 461 insertions(+), 92 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 7def296e39..6ee29f33e6 100644 --- a/open-sse/executors/kiro.ts +++ b/open-sse/executors/kiro.ts @@ -14,6 +14,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; @@ -147,35 +148,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 24bc13d810..0871b0fb69 100644 --- a/open-sse/services/kiroModels.ts +++ b/open-sse/services/kiroModels.ts @@ -23,6 +23,8 @@ * never breaks when the account is offline / unauthenticated / token-expired. */ +import { resolveKiroRuntimeRegion } from "./kiroRegion.ts"; + type RawRecord = Record; function asRecord(value: unknown): RawRecord { @@ -76,20 +78,15 @@ export function parseKiroModels(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..8f1eee3f61 --- /dev/null +++ b/open-sse/services/kiroRegion.ts @@ -0,0 +1,177 @@ +/** + * Shared Amazon Q Developer (Kiro / AWS CodeWhisperer) region resolution. + * + * CRITICAL AWS constraint — verified against the AWS docs + * "Amazon Q Developer Pro Region support" → "Supported Regions for the Q Developer console and + * Q Developer profile": + * + * • An IAM Identity Center instance (and therefore the SSO OIDC token endpoint + * `oidc.{region}.amazonaws.com`) may live in MANY regions, e.g. eu-north-1 (Stockholm), + * us-west-1, ap-southeast-2, ... + * • The Amazon Q Developer PROFILE — which produces the `profileArn` and hosts EVERY + * CodeWhisperer runtime call (generateAssistantResponse, GetUsageLimits, ListAvailableModels, + * ListAvailableProfiles) — is only hosted in **us-east-1** and **eu-central-1**. + * • "Regardless of the IAM Identity Center Region, data is stored in the Region where you create + * the Amazon Q Developer profile." → the IdC region and the runtime region are frequently + * DIFFERENT (e.g. IdC in eu-north-1 → profile in eu-central-1). + * + * Consequences enforced here: + * • `providerSpecificData.region` is the IdC/OIDC/token region. It must ONLY be used for + * `oidc.{region}.amazonaws.com` token mint/refresh (see tokenRefresh.ts / oauth providers). + * • The RUNTIME region is the region embedded in the `profileArn` (us-east-1 / eu-central-1), + * NOT the IdC region. Routing a runtime call to `q.eu-north-1.amazonaws.com` — a host that + * 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 for enterprise IdC accounts + * whose IdC lives outside us-east-1 / eu-central-1. + */ + +// 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* (identity-aware / IdC Pro tier) can exist, and + * therefore the only regions that host a Q Developer runtime endpoint for profile-bound calls. + * Order is the default probe order (us-east-1 is CodeWhisperer's home region). + */ +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 Q Developer profile regions to probe for `ListAvailableProfiles`. + * The IdC/token region is only useful here as a hint for geographic proximity — the actual + * profile always lives in one of KIRO_PROFILE_REGIONS, so those are the only regions probed. + */ +export function buildKiroProfileDiscoveryRegions(storedRegion?: string | null): string[] { + const regions: string[] = []; + const stored = normalizeRegion(storedRegion); + + // If the IdC/token region happens to be a Q profile region itself, probe it first. + if (stored && (KIRO_PROFILE_REGIONS as readonly string[]).includes(stored)) { + regions.push(stored); + } + + // Otherwise order the two known profile regions by rough geographic proximity to the IdC + // region so an EU IdC (e.g. eu-north-1) hits eu-central-1 first. + const preferEu = /^(eu|af|me|il)-/.test(stored); + const ordered = preferEu ? ["eu-central-1", "us-east-1"] : ["us-east-1", "eu-central-1"]; + for (const r of ordered) { + if (!regions.includes(r)) regions.push(r); + } + 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) 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). 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 9b85579103..1f8d083c90 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"; type JsonRecord = Record; @@ -155,32 +160,30 @@ export async function getKiroUsage(accessToken?: string, providerSpecificData?: ? 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 = 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 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); + profileArn = await discoverKiroProfileArnAcrossRegions(accessToken, storedRegion); } if (!profileArn) { 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); + // Kiro uses AWS CodeWhisperer GetUsageLimits API const payload = { origin: "AI_EDITOR", 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..fd8a375fec --- /dev/null +++ b/tests/unit/kiro-idc-cross-region.test.ts @@ -0,0 +1,225 @@ +/** + * 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 eu-central-1 first, never eu-north-1", () => { + const regions = buildKiroProfileDiscoveryRegions("eu-north-1"); + assert.deepEqual(regions, ["eu-central-1", "us-east-1"]); + assert.ok(!regions.includes("eu-north-1"), "must not probe the nonexistent q.eu-north-1 host"); +}); + +test("buildKiroProfileDiscoveryRegions: non-EU IdC probes us-east-1 first", () => { + assert.deepEqual(buildKiroProfileDiscoveryRegions("us-west-2"), ["us-east-1", "eu-central-1"]); + 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: 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; + } +});