* 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
@@ -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." };
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user