From 83c77fb0bfb5a84d21c6bef9211ade44049aa7af Mon Sep 17 00:00:00 2001 From: Michael YC JO Date: Thu, 20 Aug 2026 18:29:08 +0900 Subject: [PATCH] fix(oauth): keep the Kiro profileArn on IAM Identity Center logins (#10725) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merged via merge-train (release/v3.8.50, batch1 2026-08-20) — static gates (typecheck/file-size/complexity/cognitive/changelog) green on the combined tree; test:unit reds observed in the boarded run were verified pre-existing on the pure release tip (unrelated flake), not caused by this PR. Thanks for the contribution! --- src/shared/components/OAuthModal.tsx | 6 +- .../kiro-idc-profilearn-extradata.test.ts | 75 +++++++++++++++++++ 2 files changed, 80 insertions(+), 1 deletion(-) create mode 100644 tests/unit/kiro-idc-profilearn-extradata.test.ts diff --git a/src/shared/components/OAuthModal.tsx b/src/shared/components/OAuthModal.tsx index 5cc190c5f8..2a885f7573 100644 --- a/src/shared/components/OAuthModal.tsx +++ b/src/shared/components/OAuthModal.tsx @@ -430,13 +430,17 @@ export default function OAuthModal({ const verifyUrl = data.verification_uri_complete || data.verification_uri; if (typeof verifyUrl === "string" && verifyUrl) window.open(verifyUrl, "oauth_verify"); - // Start polling - pass extraData for Kiro (contains _clientId, _clientSecret) + // Start polling - pass extraData for Kiro (contains _clientId, _clientSecret). + // _authMethod must be forwarded too: pollToken falls back to "builder-id" without it, + // which makes postExchange skip the Q Developer profile lookup. An IdC connection then + // gets persisted with no profileArn and every usage call returns 403. const extraData = provider === "kiro" || provider === "amazon-q" ? { _clientId: data._clientId, _clientSecret: data._clientSecret, _region: data._region, + _authMethod: data._authMethod, } : provider === "ghe-copilot" && gheUrl.trim() ? { gheUrl: gheUrl.trim() } diff --git a/tests/unit/kiro-idc-profilearn-extradata.test.ts b/tests/unit/kiro-idc-profilearn-extradata.test.ts new file mode 100644 index 0000000000..79f1e604f1 --- /dev/null +++ b/tests/unit/kiro-idc-profilearn-extradata.test.ts @@ -0,0 +1,75 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { readFileSync } from "node:fs"; +import { fileURLToPath } from "node:url"; +import { dirname, join } from "node:path"; + +// Source-level contract tests for the IdC (IAM Identity Center) Kiro login path. +// +// A Kiro connection added through "Your organization (AWS IAM Identity Center)" must end up +// with a `profileArn` in providerSpecificData. Without it every CodeWhisperer usage call is +// made unbound to the Q Developer profile and AWS answers 403 "User is not authorized to make +// this call", so the quota card shows "Kiro quota API rejected the current token" while chat +// keeps working. +// +// The value is only discovered when `_authMethod` survives the device-code round trip: +// +// requestDeviceCode -> _authMethod: "idc" (startUrl present) +// OAuthModal -> forwards extraData to the poll request <- the link that broke +// pollToken -> extraData?._authMethod || "builder-id" +// postExchange -> returns early for "builder-id" (no profile lookup) +// mapTokens -> persists authMethod + profileArn +// +// The modal used to copy only _clientId/_clientSecret/_region, so an IdC login was persisted as +// a Builder ID one and the profile lookup never ran. + +const here = dirname(fileURLToPath(import.meta.url)); +const modalSource = readFileSync(join(here, "../../src/shared/components/OAuthModal.tsx"), "utf8"); +const kiroProviderSource = readFileSync( + join(here, "../../src/lib/oauth/providers/kiro.ts"), + "utf8" +); + +test("OAuthModal forwards _authMethod to the poll request for Kiro device code", () => { + const extraDataBlock = modalSource.match( + /provider === "kiro" \|\| provider === "amazon-q"\s*\?\s*\{[\s\S]*?\}/ + ); + assert.ok(extraDataBlock, "Kiro extraData block not found in OAuthModal"); + + const block = extraDataBlock[0]; + assert.match(block, /_clientId: data\._clientId/); + assert.match(block, /_clientSecret: data\._clientSecret/); + assert.match(block, /_region: data\._region/); + // The regression: dropping this line silently downgrades an IdC login to Builder ID. + assert.match( + block, + /_authMethod: data\._authMethod/, + "OAuthModal must forward _authMethod so pollToken does not fall back to builder-id" + ); +}); + +test("requestDeviceCode marks IdC logins so the profile lookup can run", () => { + assert.match( + kiroProviderSource, + /_authMethod: config\.skipIssuerUrlForRegistration \? "idc" : "builder-id"/ + ); +}); + +test("pollToken keeps the device-code _authMethod when the caller forwards it", () => { + assert.match(kiroProviderSource, /_authMethod: extraData\?\._authMethod \|\| "builder-id"/); +}); + +test("postExchange skips the profile lookup only for Builder ID logins", () => { + assert.match(kiroProviderSource, /if \(tokenData\?\._authMethod === "builder-id"\) return null;/); + assert.match( + kiroProviderSource, + /discoverKiroProfileArnAcrossRegions\(accessToken, storedRegion\)/ + ); +}); + +test("mapTokens persists profileArn when the exchange discovered one", () => { + assert.match( + kiroProviderSource, + /\.\.\.\(extra\?\.profileArn \? \{ profileArn: extra\.profileArn \} : \{\}\)/ + ); +});