diff --git a/README.md b/README.md index 86d38eb6..533e3cbc 100644 --- a/README.md +++ b/README.md @@ -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.
Built with ❤️ for developers who code 24/7
- diff --git a/open-sse/services/usage.js b/open-sse/services/usage.js index bf0c28ab..1430c5aa 100644 --- a/open-sse/services/usage.js +++ b/open-sse/services/usage.js @@ -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." }; } } - diff --git a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/QuotaTable.js b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/QuotaTable.js index 8fc2ecf7..e34da0de 100644 --- a/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/QuotaTable.js +++ b/src/app/(dashboard)/dashboard/usage/components/ProviderLimits/QuotaTable.js @@ -83,11 +83,6 @@ export default function QuotaTable({ quotas = [], compact = false }) { return (
- - {/* Model Name */} - {/* Limit Progress */} - {/* Reset Time */} - {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 */} - {/* Limit (Progress + Numbers) */} -
+
{colors.emoji} @@ -114,7 +109,7 @@ export default function QuotaTable({ quotas = [], compact = false }) {
+
{/* Progress bar - always show with border for visibility */}
{/* Reset Time */} -
+ {countdown !== "-" || resetDisplay ? (
{countdown !== "-" && ( diff --git a/src/app/api/oauth/[provider]/[action]/route.js b/src/app/api/oauth/[provider]/[action]/route.js index 3e496221..771e1fb1 100644 --- a/src/app/api/oauth/[provider]/[action]/route.js +++ b/src/app/api/oauth/[provider]/[action]/route.js @@ -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({ diff --git a/src/lib/oauth/providers.js b/src/lib/oauth/providers.js index f004fa64..ffb7e2a1 100644 --- a/src/lib/oauth/providers.js +++ b/src/lib/oauth/providers.js @@ -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 }; } - diff --git a/src/shared/components/KiroAuthModal.js b/src/shared/components/KiroAuthModal.js index 0a9133df..6f567c6d 100644 --- a/src/shared/components/KiroAuthModal.js +++ b/src/shared/components/KiroAuthModal.js @@ -126,10 +126,10 @@ export default function KiroAuthModal({ isOpen, onMethodSelect, onClose }) {
- {/* AWS IAM Identity Center (IDC) - HIDDEN */} + {/* AWS IAM Identity Center (IDC) */}
@@ -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, + }), };