* feat(kiro): wire aws identity center device flow into provider oauth (#587)

* feat(kiro): wire aws identity center device flow into provider oauth

Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com>
This commit is contained in:
Dang Dinh Quan
2026-04-15 11:44:46 +07:00
committed by GitHub
parent c3a2bd01b7
commit b1288c5064
7 changed files with 243 additions and 164 deletions

View File

@@ -808,7 +808,7 @@ Models:
```bash
Dashboard → Connect Kiro
→ AWS Builder ID or Google/GitHub
→ AWS Builder ID, AWS IAM Identity Center, Google, GitHub
→ Unlimited usage
Models:
@@ -1208,4 +1208,3 @@ MIT License - see [LICENSE](LICENSE) for details.
<div align="center">
<sub>Built with ❤️ for developers who code 24/7</sub>
</div>

View File

@@ -532,149 +532,155 @@ async function getCodexUsage(accessToken) {
/**
* Kiro (AWS CodeWhisperer) Usage
*/
function parseKiroQuotaData(data) {
const usageList = data.usageBreakdownList || [];
const quotaInfo = {};
const resetAt = parseResetTime(data.nextDateReset || data.resetDate);
usageList.forEach((breakdown) => {
const resourceType = breakdown.resourceType?.toLowerCase() || "unknown";
const used = breakdown.currentUsageWithPrecision || 0;
const total = breakdown.usageLimitWithPrecision || 0;
quotaInfo[resourceType] = {
used,
total,
remaining: total - used,
resetAt,
unlimited: false,
};
// Add free trial if available
if (breakdown.freeTrialInfo) {
const freeUsed = breakdown.freeTrialInfo.currentUsageWithPrecision || 0;
const freeTotal = breakdown.freeTrialInfo.usageLimitWithPrecision || 0;
quotaInfo[`${resourceType}_freetrial`] = {
used: freeUsed,
total: freeTotal,
remaining: freeTotal - freeUsed,
resetAt: parseResetTime(breakdown.freeTrialInfo.freeTrialExpiry || resetAt),
unlimited: false,
};
}
});
return {
plan: data.subscriptionInfo?.subscriptionTitle || "Kiro",
quotas: quotaInfo,
};
}
async function getKiroUsage(accessToken, providerSpecificData) {
// Default profileArn fallback
const DEFAULT_PROFILE_ARN = "arn:aws:codewhisperer:us-east-1:638616132270:profile/AAAACCCCXXXX";
const profileArn = providerSpecificData?.profileArn || DEFAULT_PROFILE_ARN;
const authMethod = providerSpecificData?.authMethod || "builder-id";
try {
// Try old API first (POST method)
const payload = {
origin: "AI_EDITOR",
profileArn: profileArn,
resourceType: "AGENTIC_REQUEST",
};
const getUsageParams = new URLSearchParams({
isEmailRequired: "true",
origin: "AI_EDITOR",
resourceType: "AGENTIC_REQUEST",
});
const response = await fetch("https://codewhisperer.us-east-1.amazonaws.com", {
method: "POST",
headers: {
"Authorization": `Bearer ${accessToken}`,
"Content-Type": "application/x-amz-json-1.0",
"x-amz-target": "AmazonCodeWhispererService.GetUsageLimits",
"Accept": "application/json",
},
body: JSON.stringify(payload),
});
if (!response.ok) {
const errorText = await response.text();
// Handle authentication errors gracefully
if (response.status === 403 || response.status === 401) {
return {
message: "Kiro quota API authentication expired. Chat may still work.",
quotas: {}
};
}
throw new Error(`Kiro API error (${response.status}): ${errorText}`);
}
const data = await response.json();
// Parse usage data from usageBreakdownList
const usageList = data.usageBreakdownList || [];
const quotaInfo = {};
// Parse reset time - supports multiple formats (nextDateReset, resetDate, etc.)
const resetAt = parseResetTime(data.nextDateReset || data.resetDate);
usageList.forEach((breakdown) => {
const resourceType = breakdown.resourceType?.toLowerCase() || "unknown";
const used = breakdown.currentUsageWithPrecision || 0;
const total = breakdown.usageLimitWithPrecision || 0;
quotaInfo[resourceType] = {
used,
total,
remaining: total - used,
resetAt,
unlimited: false,
};
// Add free trial if available
if (breakdown.freeTrialInfo) {
const freeUsed = breakdown.freeTrialInfo.currentUsageWithPrecision || 0;
const freeTotal = breakdown.freeTrialInfo.usageLimitWithPrecision || 0;
quotaInfo[`${resourceType}_freetrial`] = {
used: freeUsed,
total: freeTotal,
remaining: freeTotal - freeUsed,
resetAt,
unlimited: false,
};
}
});
return {
plan: data.subscriptionInfo?.subscriptionTitle || "Kiro",
quotas: quotaInfo,
};
} catch (error) {
// Fallback to new API (GET method)
try {
const params = new URLSearchParams({
origin: "AI_EDITOR",
profileArn: profileArn,
resourceType: "AGENTIC_REQUEST",
});
const fallbackResponse = await fetch(`https://q.us-east-1.amazonaws.com/getUsageLimits?${params}`, {
method: "GET",
// For compatibility, try multiple known Kiro usage endpoints
const attempts = [
{
name: "codewhisperer-get",
run: async () => fetch(
`https://codewhisperer.us-east-1.amazonaws.com/getUsageLimits?${getUsageParams.toString()}`,
{
method: "GET",
headers: {
"Authorization": `Bearer ${accessToken}`,
"Accept": "application/json",
"x-amz-user-agent": "aws-sdk-js/1.0.0 KiroIDE",
"user-agent": "aws-sdk-js/1.0.0 KiroIDE",
},
},
),
},
{
name: "codewhisperer-post",
run: async () => fetch("https://codewhisperer.us-east-1.amazonaws.com", {
method: "POST",
headers: {
"Authorization": `Bearer ${accessToken}`,
"Content-Type": "application/x-amz-json-1.0",
"x-amz-target": "AmazonCodeWhispererService.GetUsageLimits",
"Accept": "application/json",
},
});
body: JSON.stringify({
origin: "AI_EDITOR",
profileArn,
resourceType: "AGENTIC_REQUEST",
}),
}),
},
{
name: "q-get",
run: async () => {
const params = new URLSearchParams({
origin: "AI_EDITOR",
profileArn,
resourceType: "AGENTIC_REQUEST",
});
return fetch(`https://q.us-east-1.amazonaws.com/getUsageLimits?${params}`, {
method: "GET",
headers: {
"Authorization": `Bearer ${accessToken}`,
"Accept": "application/json",
},
});
},
},
];
if (!fallbackResponse.ok) {
throw new Error(`Fallback API error (${fallbackResponse.status})`);
let sawAuthError = false;
const errors = [];
for (const attempt of attempts) {
try {
const response = await attempt.run();
if (!response.ok) {
const errorText = await response.text().catch(() => "");
if (response.status === 401 || response.status === 403) {
sawAuthError = true;
}
errors.push(`${attempt.name}:${response.status}${errorText ? `:${errorText}` : ""}`);
continue;
}
const fallbackData = await fallbackResponse.json();
// Parse new API response structure
const usageList = fallbackData.usageBreakdownList || [];
const quotaInfo = {};
const resetAt = parseResetTime(fallbackData.nextDateReset || fallbackData.resetDate);
usageList.forEach((breakdown) => {
const resourceType = breakdown.resourceType?.toLowerCase() || "unknown";
const used = breakdown.currentUsageWithPrecision || 0;
const total = breakdown.usageLimitWithPrecision || 0;
quotaInfo[resourceType] = {
used,
total,
remaining: total - used,
resetAt,
unlimited: false,
};
// Add free trial if available
if (breakdown.freeTrialInfo) {
const freeUsed = breakdown.freeTrialInfo.currentUsageWithPrecision || 0;
const freeTotal = breakdown.freeTrialInfo.usageLimitWithPrecision || 0;
quotaInfo[`${resourceType}_freetrial`] = {
used: freeUsed,
total: freeTotal,
remaining: freeTotal - freeUsed,
resetAt: parseResetTime(breakdown.freeTrialInfo.freeTrialExpiry),
unlimited: false,
};
}
});
return {
plan: fallbackData.subscriptionInfo?.subscriptionTitle || "Kiro",
quotas: quotaInfo,
};
} catch (fallbackError) {
throw new Error(`Failed to fetch Kiro usage: ${error.message} | Fallback: ${fallbackError.message}`);
const data = await response.json();
return parseKiroQuotaData(data);
} catch (error) {
errors.push(`${attempt.name}:${error.message}`);
}
}
if (sawAuthError && authMethod === "idc") {
return {
message: "Kiro quota API is unavailable for the current AWS IAM Identity Center session. Chat may still work. If this persists after renewing your session, reconnect Kiro.",
quotas: {},
};
}
if (sawAuthError) {
return {
message: "Kiro quota API rejected the current token. Chat may still work.",
quotas: {},
};
}
const fallbackMessage =
errors.length > 0
? `Unable to fetch Kiro usage right now. (${errors[errors.length - 1]})`
: "Unable to fetch Kiro usage right now.";
return {
message: fallbackMessage,
quotas: {},
};
}
/**
@@ -705,4 +711,3 @@ async function getIflowUsage(accessToken) {
return { message: "Unable to fetch iFlow usage." };
}
}

View File

@@ -83,11 +83,6 @@ export default function QuotaTable({ quotas = [], compact = false }) {
return (
<div className="overflow-x-auto">
<table className="w-full table-fixed text-left">
<colgroup>
<col className="w-[30%]" /> {/* Model Name */}
<col className="w-[45%]" /> {/* Limit Progress */}
<col className="w-[25%]" /> {/* Reset Time */}
</colgroup>
<tbody>
{quotas.map((quota, index) => {
const remaining = quota.remainingPercentage !== undefined
@@ -104,7 +99,7 @@ export default function QuotaTable({ quotas = [], compact = false }) {
className="border-b border-black/5 dark:border-white/5 hover:bg-black/[0.02] dark:hover:bg-white/[0.02] transition-colors"
>
{/* Model Name with Status Emoji */}
<td className={cellPad}>
<td className={`${cellPad} w-[30%]`}>
<div className="flex items-center gap-1.5 min-w-0">
<span className="text-[10px] shrink-0">{colors.emoji}</span>
<span className={`${nameText} font-medium text-text-primary truncate`}>
@@ -114,7 +109,7 @@ export default function QuotaTable({ quotas = [], compact = false }) {
</td>
{/* Limit (Progress + Numbers) */}
<td className={cellPad}>
<td className={`${cellPad} w-[45%]`}>
<div className={compact ? "space-y-1" : "space-y-1.5"}>
{/* Progress bar - always show with border for visibility */}
<div className={`${compact ? "h-1" : "h-1.5"} rounded-full overflow-hidden border ${colors.bgLight} ${
@@ -139,7 +134,7 @@ export default function QuotaTable({ quotas = [], compact = false }) {
</td>
{/* Reset Time */}
<td className={cellPad}>
<td className={`${cellPad} w-[25%]`}>
{countdown !== "-" || resetDisplay ? (
<div className="space-y-0.5">
{countdown !== "-" && (

View File

@@ -58,15 +58,25 @@ export async function GET(request, { params }) {
}
const authData = generateAuthData(provider, null);
const startUrl = searchParams.get("start_url");
const region = searchParams.get("region");
const authMethod = searchParams.get("auth_method");
const deviceOptions = provider === "kiro"
? {
...(startUrl ? { startUrl } : {}),
...(region ? { region } : {}),
...(authMethod ? { authMethod } : {}),
}
: undefined;
// Providers that don't use PKCE for device code
const noPkceDeviceProviders = ["github", "kiro", "kimi-coding", "kilocode", "codebuddy"];
let deviceData;
if (noPkceDeviceProviders.includes(provider)) {
deviceData = await requestDeviceCode(provider);
deviceData = await requestDeviceCode(provider, undefined, deviceOptions);
} else {
// Qwen and other PKCE providers
deviceData = await requestDeviceCode(provider, authData.codeChallenge);
deviceData = await requestDeviceCode(provider, authData.codeChallenge, deviceOptions);
}
return NextResponse.json({

View File

@@ -25,6 +25,28 @@ import {
CODEBUDDY_CONFIG,
} from "./constants/oauth";
const BASE64_BLOCK_SIZE = 4;
/**
* Decode JWT access token and extract a stable account identifier for display/upsert.
* @param {string} accessToken
* @returns {string|undefined}
*/
function extractEmailFromAccessToken(accessToken) {
try {
if (!accessToken || typeof accessToken !== "string") return undefined;
const parts = accessToken.split(".");
if (parts.length !== 3) return undefined;
const base64 = parts[1].replace(/-/g, "+").replace(/_/g, "/");
const missingPadding = (BASE64_BLOCK_SIZE - (base64.length % BASE64_BLOCK_SIZE)) % BASE64_BLOCK_SIZE;
const padded = base64 + "=".repeat(missingPadding);
const payload = JSON.parse(Buffer.from(padded, "base64").toString("utf8"));
return payload.email || payload.preferred_username || payload.sub || undefined;
} catch {
return undefined;
}
}
// Provider configurations
const PROVIDERS = {
claude: {
@@ -652,9 +674,17 @@ const PROVIDERS = {
config: KIRO_CONFIG,
flowType: "device_code",
// Kiro uses AWS SSO OIDC - requires client registration first
requestDeviceCode: async (config) => {
requestDeviceCode: async (config, codeChallenge, options = {}) => {
const trimmedRegion = typeof options.region === "string" ? options.region.trim() : "";
const region = trimmedRegion || "us-east-1";
const trimmedStartUrl = typeof options.startUrl === "string" ? options.startUrl.trim() : "";
const startUrl = trimmedStartUrl || config.startUrl;
const authMethod = options.authMethod === "idc" ? "idc" : "builder-id";
const registerClientUrl = `https://oidc.${region}.amazonaws.com/client/register`;
const deviceAuthUrl = `https://oidc.${region}.amazonaws.com/device_authorization`;
// Step 1: Register client with AWS SSO OIDC
const registerRes = await fetch(config.registerClientUrl, {
const registerRes = await fetch(registerClientUrl, {
method: "POST",
headers: {
"Content-Type": "application/json",
@@ -677,7 +707,7 @@ const PROVIDERS = {
const clientInfo = await registerRes.json();
// Step 2: Request device authorization
const deviceRes = await fetch(config.deviceAuthUrl, {
const deviceRes = await fetch(deviceAuthUrl, {
method: "POST",
headers: {
"Content-Type": "application/json",
@@ -686,7 +716,7 @@ const PROVIDERS = {
body: JSON.stringify({
clientId: clientInfo.clientId,
clientSecret: clientInfo.clientSecret,
startUrl: config.startUrl,
startUrl,
}),
});
@@ -708,10 +738,15 @@ const PROVIDERS = {
// Store client credentials for token exchange
_clientId: clientInfo.clientId,
_clientSecret: clientInfo.clientSecret,
_region: region,
_authMethod: authMethod,
_startUrl: startUrl,
};
},
pollToken: async (config, deviceCode, codeVerifier, extraData) => {
const response = await fetch(config.tokenUrl, {
const region = extraData?._region || "us-east-1";
const tokenUrl = `https://oidc.${region}.amazonaws.com/token`;
const response = await fetch(tokenUrl, {
method: "POST",
headers: {
"Content-Type": "application/json",
@@ -745,6 +780,9 @@ const PROVIDERS = {
// Store client credentials for refresh
_clientId: extraData?._clientId,
_clientSecret: extraData?._clientSecret,
_region: extraData?._region,
_authMethod: extraData?._authMethod,
_startUrl: extraData?._startUrl,
},
};
}
@@ -758,14 +796,19 @@ const PROVIDERS = {
};
},
mapTokens: (tokens) => {
const email = extractEmailFromAccessToken(tokens.access_token);
const mapped = {
accessToken: tokens.access_token,
refreshToken: tokens.refresh_token,
expiresIn: tokens.expires_in,
email,
providerSpecificData: {
profileArn: tokens?.profile_arn || null,
clientId: tokens._clientId,
clientSecret: tokens._clientSecret,
region: tokens._region || "us-east-1",
authMethod: tokens._authMethod || "builder-id",
startUrl: tokens._startUrl || KIRO_CONFIG.startUrl,
},
};
return mapped;
@@ -1158,12 +1201,12 @@ export async function exchangeTokens(providerName, code, redirectUri, codeVerifi
/**
* Request device code (for device_code flow)
*/
export async function requestDeviceCode(providerName, codeChallenge) {
export async function requestDeviceCode(providerName, codeChallenge, options) {
const provider = getProvider(providerName);
if (provider.flowType !== "device_code") {
throw new Error(`Provider ${providerName} does not support device code flow`);
}
return await provider.requestDeviceCode(provider.config, codeChallenge);
return await provider.requestDeviceCode(provider.config, codeChallenge, options || {});
}
/**
@@ -1213,4 +1256,3 @@ export async function pollForToken(providerName, deviceCode, codeVerifier, extra
return { success: false, error: result.data.error, errorDescription: result.data.error_description };
}

View File

@@ -126,10 +126,10 @@ export default function KiroAuthModal({ isOpen, onMethodSelect, onClose }) {
</div>
</button>
{/* AWS IAM Identity Center (IDC) - HIDDEN */}
{/* AWS IAM Identity Center (IDC) */}
<button
onClick={() => handleMethodSelect("idc")}
className="hidden w-full p-4 text-left border border-border rounded-lg hover:bg-sidebar transition-colors"
className="w-full p-4 text-left border border-border rounded-lg hover:bg-sidebar transition-colors"
>
<div className="flex items-start gap-3">
<span className="material-symbols-outlined text-primary mt-0.5">business</span>

View File

@@ -10,7 +10,7 @@ import { useCopyToClipboard } from "@/shared/hooks/useCopyToClipboard";
* - Localhost: Auto callback via popup message
* - Remote: Manual paste callback URL
*/
export default function OAuthModal({ isOpen, provider, providerInfo, onSuccess, onClose, oauthMeta }) {
export default function OAuthModal({ isOpen, provider, providerInfo, onSuccess, onClose, oauthMeta, idcConfig }) {
const [step, setStep] = useState("waiting"); // waiting | input | success | error
const [authData, setAuthData] = useState(null);
const [callbackUrl, setCallbackUrl] = useState("");
@@ -138,18 +138,30 @@ export default function OAuthModal({ isOpen, provider, providerInfo, onSuccess,
setIsDeviceCode(true);
setStep("waiting");
const res = await fetch(`/api/oauth/${provider}/device-code`);
const deviceCodeUrl = new URL(`/api/oauth/${provider}/device-code`, window.location.origin);
if (provider === "kiro" && idcConfig?.startUrl) {
deviceCodeUrl.searchParams.set("start_url", idcConfig.startUrl);
if (idcConfig.region) {
deviceCodeUrl.searchParams.set("region", idcConfig.region);
}
deviceCodeUrl.searchParams.set("auth_method", "idc");
}
const res = await fetch(deviceCodeUrl.toString());
const data = await res.json();
if (!res.ok) throw new Error(data.error);
setDeviceData(data);
// Open verification URL
const verifyUrl = data.verification_uri_complete || data.verification_uri;
if (verifyUrl) window.open(verifyUrl, "_blank");
// Pass extraData for Kiro (contains _clientId, _clientSecret)
const extraData = provider === "kiro" ? { _clientId: data._clientId, _clientSecret: data._clientSecret } : null;
const extraData = provider === "kiro"
? {
_clientId: data._clientId,
_clientSecret: data._clientSecret,
_region: data._region,
_authMethod: data._authMethod,
_startUrl: data._startUrl,
}
: null;
startPolling(data.device_code, data.codeVerifier, data.interval || 5, extraData);
return;
}
@@ -209,7 +221,7 @@ export default function OAuthModal({ isOpen, provider, providerInfo, onSuccess,
setError(err.message);
setStep("error");
}
}, [provider, isLocalhost, startPolling]);
}, [provider, isLocalhost, startPolling, oauthMeta, idcConfig]);
// Reset state and start OAuth when modal opens
useEffect(() => {
@@ -345,6 +357,7 @@ export default function OAuthModal({ isOpen, provider, providerInfo, onSuccess,
}, [onClose, provider]);
if (!provider || !providerInfo) return null;
const deviceLoginUrl = deviceData?.verification_uri_complete || deviceData?.verification_uri || "";
return (
<Modal isOpen={isOpen} title={`Connect ${providerInfo.name}`} onClose={handleClose} size="lg">
@@ -372,18 +385,28 @@ export default function OAuthModal({ isOpen, provider, providerInfo, onSuccess,
<>
<div className="text-center py-4">
<p className="text-sm text-text-muted mb-4">
Visit the URL below and enter the code:
Visit the login URL below and authorize:
</p>
<div className="bg-sidebar p-4 rounded-lg mb-4">
<p className="text-xs text-text-muted mb-1">Verification URL</p>
<p className="text-xs text-text-muted mb-1">Login URL</p>
<div className="flex items-center gap-2">
<code className="flex-1 text-sm break-all">{deviceData.verification_uri}</code>
<code className="flex-1 text-sm break-all">{deviceLoginUrl}</code>
<Button
size="sm"
variant="ghost"
icon={copied === "verify_url" ? "check" : "content_copy"}
onClick={() => copy(deviceData.verification_uri, "verify_url")}
icon={copied === "login_url" ? "check" : "content_copy"}
onClick={() => copy(deviceLoginUrl, "login_url")}
disabled={!deviceLoginUrl}
/>
<Button
size="sm"
variant="ghost"
icon="open_in_new"
onClick={() => window.open(deviceLoginUrl, "_blank", "noopener,noreferrer")}
disabled={!deviceLoginUrl}
>
Open
</Button>
</div>
</div>
<div className="bg-primary/10 p-4 rounded-lg">
@@ -494,4 +517,9 @@ OAuthModal.propTypes = {
onClose: PropTypes.func.isRequired,
/** Extra metadata passed to /authorize and /exchange (e.g. gitlab clientId/baseUrl) */
oauthMeta: PropTypes.object,
/** Optional Kiro IDC config for AWS IAM Identity Center device flow */
idcConfig: PropTypes.shape({
startUrl: PropTypes.string,
region: PropTypes.string,
}),
};