From 75f486b7a221e0c988f07b7e192df7690b9c044a Mon Sep 17 00:00:00 2001 From: decolua Date: Fri, 6 Mar 2026 00:21:27 +0700 Subject: [PATCH] Added profile ARN handling in OAuth provider mapping and improved polling logic in OAuth modal for better user experience. --- open-sse/services/usage.js | 74 ++++++++++++++++++++--- package.json | 2 +- src/lib/oauth/providers.js | 23 ++++--- src/shared/components/KiroOAuthWrapper.js | 2 + src/shared/components/OAuthModal.js | 21 +++++++ 5 files changed, 105 insertions(+), 17 deletions(-) diff --git a/open-sse/services/usage.js b/open-sse/services/usage.js index ec7ec54..733768a 100644 --- a/open-sse/services/usage.js +++ b/open-sse/services/usage.js @@ -470,13 +470,12 @@ async function getCodexUsage(accessToken) { * Kiro (AWS CodeWhisperer) Usage */ async function getKiroUsage(accessToken, providerSpecificData) { - try { - const profileArn = providerSpecificData?.profileArn; - if (!profileArn) { - return { message: "Kiro connected. Profile ARN not available for quota tracking." }; - } + // Default profileArn fallback + const DEFAULT_PROFILE_ARN = "arn:aws:codewhisperer:us-east-1:638616132270:profile/AAAACCCCXXXX"; + const profileArn = providerSpecificData?.profileArn || DEFAULT_PROFILE_ARN; - // Kiro uses AWS CodeWhisperer GetUsageLimits API + try { + // Try old API first (POST method) const payload = { origin: "AI_EDITOR", profileArn: profileArn, @@ -550,7 +549,68 @@ async function getKiroUsage(accessToken, providerSpecificData) { quotas: quotaInfo, }; } catch (error) { - throw new Error(`Failed to fetch Kiro usage: ${error.message}`); + // 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", + headers: { + "Authorization": `Bearer ${accessToken}`, + "Accept": "application/json", + }, + }); + + if (!fallbackResponse.ok) { + throw new Error(`Fallback API error (${fallbackResponse.status})`); + } + + 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}`); + } } } diff --git a/package.json b/package.json index ae3508e..d2547dc 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "9router-app", - "version": "0.3.31", + "version": "0.3.32", "description": "9Router web dashboard", "private": true, "scripts": { diff --git a/src/lib/oauth/providers.js b/src/lib/oauth/providers.js index 9878089..ebd579e 100644 --- a/src/lib/oauth/providers.js +++ b/src/lib/oauth/providers.js @@ -660,6 +660,7 @@ const PROVIDERS = { access_token: data.accessToken, refresh_token: data.refreshToken, expires_in: data.expiresIn, + profile_arn: data?.profileArn || null, // Store client credentials for refresh _clientId: extraData?._clientId, _clientSecret: extraData?._clientSecret, @@ -675,15 +676,19 @@ const PROVIDERS = { }, }; }, - mapTokens: (tokens) => ({ - accessToken: tokens.access_token, - refreshToken: tokens.refresh_token, - expiresIn: tokens.expires_in, - providerSpecificData: { - clientId: tokens._clientId, - clientSecret: tokens._clientSecret, - }, - }), + mapTokens: (tokens) => { + const mapped = { + accessToken: tokens.access_token, + refreshToken: tokens.refresh_token, + expiresIn: tokens.expires_in, + providerSpecificData: { + profileArn: tokens?.profile_arn || null, + clientId: tokens._clientId, + clientSecret: tokens._clientSecret, + }, + }; + return mapped; + }, }, cursor: { diff --git a/src/shared/components/KiroOAuthWrapper.js b/src/shared/components/KiroOAuthWrapper.js index be53ccd..93e45e3 100644 --- a/src/shared/components/KiroOAuthWrapper.js +++ b/src/shared/components/KiroOAuthWrapper.js @@ -43,12 +43,14 @@ export default function KiroOAuthWrapper({ isOpen, providerInfo, onSuccess, onCl setAuthMethod(null); setSocialProvider(null); onSuccess?.(); + onClose?.(); // Close modal after success }; const handleDeviceSuccess = () => { setAuthMethod(null); setIdcConfig(null); onSuccess?.(); + onClose?.(); // Close modal after success }; // Show method selection first diff --git a/src/shared/components/OAuthModal.js b/src/shared/components/OAuthModal.js index c78f9f3..c25c501 100644 --- a/src/shared/components/OAuthModal.js +++ b/src/shared/components/OAuthModal.js @@ -19,6 +19,7 @@ export default function OAuthModal({ isOpen, provider, providerInfo, onSuccess, const [deviceData, setDeviceData] = useState(null); const [polling, setPolling] = useState(false); const popupRef = useRef(null); + const pollingAbortRef = useRef(false); const { copied, copy } = useCopyToClipboard(); // State for client-only values to avoid hydration mismatch @@ -66,12 +67,27 @@ export default function OAuthModal({ isOpen, provider, providerInfo, onSuccess, // Poll for device code token const startPolling = useCallback(async (deviceCode, codeVerifier, interval, extraData) => { + pollingAbortRef.current = false; setPolling(true); const maxAttempts = 60; for (let i = 0; i < maxAttempts; i++) { + // Check if polling should be aborted + if (pollingAbortRef.current) { + console.log("[OAuthModal] Polling aborted"); + setPolling(false); + return; + } + await new Promise((r) => setTimeout(r, interval * 1000)); + // Check again after sleep + if (pollingAbortRef.current) { + console.log("[OAuthModal] Polling aborted after sleep"); + setPolling(false); + return; + } + try { const res = await fetch(`/api/oauth/${provider}/poll`, { method: "POST", @@ -82,6 +98,7 @@ export default function OAuthModal({ isOpen, provider, providerInfo, onSuccess, const data = await res.json(); if (data.success) { + pollingAbortRef.current = true; // Stop polling immediately setStep("success"); setPolling(false); onSuccess?.(); @@ -182,7 +199,11 @@ export default function OAuthModal({ isOpen, provider, providerInfo, onSuccess, setIsDeviceCode(false); setDeviceData(null); setPolling(false); + pollingAbortRef.current = false; startOAuthFlow(); + } else if (!isOpen) { + // Abort polling when modal closes + pollingAbortRef.current = true; } }, [isOpen, provider, startOAuthFlow]);