diff --git a/ui/src/components/account/add-account-dialog.tsx b/ui/src/components/account/add-account-dialog.tsx index 49bb0448..646a862a 100644 --- a/ui/src/components/account/add-account-dialog.tsx +++ b/ui/src/components/account/add-account-dialog.tsx @@ -39,11 +39,15 @@ import { } from '@/components/account/antigravity-responsibility-constants'; import { DEFAULT_KIRO_AUTH_METHOD, + DEFAULT_KIRO_IDC_FLOW, + getKiroEffectiveFlowType, + getKiroEffectiveStartEndpoint, getKiroAuthMethodOption, + isKiroSocialAuthMethod, isDeviceCodeProvider, KIRO_AUTH_METHOD_OPTIONS, } from '@/lib/provider-config'; -import type { KiroAuthMethod } from '@/lib/provider-config'; +import type { KiroAuthMethod, KiroIDCFlow } from '@/lib/provider-config'; import { useTranslation } from 'react-i18next'; import { toast } from 'sonner'; @@ -81,6 +85,9 @@ export function AddAccountDialog({ const [powerUserModeEnabled, setPowerUserModeEnabled] = useState(false); const [powerUserModeLoading, setPowerUserModeLoading] = useState(false); const [kiroAuthMethod, setKiroAuthMethod] = useState(DEFAULT_KIRO_AUTH_METHOD); + const [kiroIDCStartUrl, setKiroIDCStartUrl] = useState(''); + const [kiroIDCRegion, setKiroIDCRegion] = useState(''); + const [kiroIDCFlow, setKiroIDCFlow] = useState(DEFAULT_KIRO_IDC_FLOW); const { t } = useTranslation(); const wasAuthenticatingRef = useRef(false); const powerUserModeRequestIdRef = useRef(0); @@ -97,9 +104,19 @@ export function AddAccountDialog({ const isGeminiRiskAcknowledged = normalizeRiskPhrase(riskAcknowledgementText) === RISK_ACK_PHRASE; const defaultDeviceCode = isDeviceCodeProvider(provider); const kiroMethodOption = getKiroAuthMethodOption(kiroAuthMethod); - const isDeviceCode = isKiro ? kiroMethodOption.flowType === 'device_code' : defaultDeviceCode; + const isKiroIdc = isKiro && kiroAuthMethod === 'idc'; + const isKiroSocial = isKiro && isKiroSocialAuthMethod(kiroAuthMethod); + const selectedKiroFlowType = isKiro + ? getKiroEffectiveFlowType(kiroAuthMethod, kiroIDCFlow) + : undefined; + const selectedKiroStartEndpoint = isKiro + ? getKiroEffectiveStartEndpoint(kiroAuthMethod) + : undefined; + const isDeviceCode = isKiro ? selectedKiroFlowType === 'device_code' : defaultDeviceCode; const isPending = authFlow.isAuthenticating || kiroImportMutation.isPending; const nicknameTrimmed = nickname.trim(); + const kiroIDCStartUrlTrimmed = kiroIDCStartUrl.trim(); + const kiroIDCRegionTrimmed = kiroIDCRegion.trim(); const errorMessage = localError || authFlow.error; const fetchPowerUserModeState = useCallback(async (): Promise => { @@ -168,6 +185,9 @@ export function AddAccountDialog({ setPowerUserModeEnabled(false); setPowerUserModeLoading(false); setKiroAuthMethod(DEFAULT_KIRO_AUTH_METHOD); + setKiroIDCStartUrl(''); + setKiroIDCRegion(''); + setKiroIDCFlow(DEFAULT_KIRO_IDC_FLOW); powerUserModeRequestIdRef.current += 1; powerUserModeLoadErrorShownRef.current = false; wasAuthenticatingRef.current = false; @@ -282,12 +302,19 @@ export function AddAccountDialog({ return; } setLocalError(null); + if (isKiroIdc && !kiroIDCStartUrlTrimmed) { + setLocalError('IDC Start URL is required for Kiro IAM Identity Center login.'); + return; + } wasAuthenticatingRef.current = true; authFlow.startAuth(provider, { nickname: nicknameTrimmed || undefined, kiroMethod: isKiro ? kiroAuthMethod : undefined, - flowType: isKiro ? kiroMethodOption.flowType : undefined, - startEndpoint: isKiro ? kiroMethodOption.startEndpoint : undefined, + kiroIDCStartUrl: isKiroIdc ? kiroIDCStartUrlTrimmed : undefined, + kiroIDCRegion: isKiroIdc && kiroIDCRegionTrimmed ? kiroIDCRegionTrimmed : undefined, + kiroIDCFlow: isKiroIdc ? kiroIDCFlow : undefined, + flowType: isKiro ? selectedKiroFlowType : undefined, + startEndpoint: isKiro ? selectedKiroStartEndpoint : undefined, riskAcknowledgement: requiresAgyResponsibilityFlow ? { version: ANTIGRAVITY_ACK_VERSION, @@ -399,6 +426,77 @@ export function AddAccountDialog({

{kiroMethodOption.description}

+ {isKiroSocial && ( +

+ If your browser does not return automatically after login, CCS can accept the + final + + kiro://... + + callback URL in the next step. +

+ )} + + )} + + {isKiroIdc && !showAuthUI && ( +
+
+ + { + setKiroIDCStartUrl(e.target.value); + setLocalError(null); + }} + placeholder="https://d-xxx.awsapps.com/start" + disabled={isPending} + /> +

+ Required for organization IAM Identity Center login. +

+
+ +
+ + { + setKiroIDCRegion(e.target.value); + setLocalError(null); + }} + placeholder="us-east-1" + disabled={isPending} + /> +

+ Optional. Leave blank to use the upstream default region. +

+
+ +
+ + +

+ Auth Code opens a browser and may need the final callback URL pasted back. Device + Code shows a verification code instead. +

+
)} @@ -438,7 +536,9 @@ export function AddAccountDialog({

{authFlow.isDeviceCodeFlow ? t('addAccountDialog.deviceCodeHint') - : t('addAccountDialog.browserHint')} + : isKiroSocial + ? 'Complete sign-in in your browser. If it does not return automatically, paste the final kiro:// callback URL below.' + : t('addAccountDialog.browserHint')}

@@ -486,13 +586,19 @@ export function AddAccountDialog({ {/* Callback paste field */}
setCallbackUrl(e.target.value)} - placeholder={t('addAccountDialog.callbackPlaceholder')} + placeholder={ + isKiroSocial + ? 'kiro://kiro.kiroAgent/authenticate-success?code=...&state=...' + : t('addAccountDialog.callbackPlaceholder') + } className="font-mono text-xs" />
diff --git a/ui/src/hooks/use-cliproxy-auth-flow.ts b/ui/src/hooks/use-cliproxy-auth-flow.ts index dcc57e09..1ac8199f 100644 --- a/ui/src/hooks/use-cliproxy-auth-flow.ts +++ b/ui/src/hooks/use-cliproxy-auth-flow.ts @@ -26,6 +26,9 @@ interface AuthFlowState { interface StartAuthOptions { nickname?: string; kiroMethod?: string; + kiroIDCStartUrl?: string; + kiroIDCRegion?: string; + kiroIDCFlow?: 'authcode' | 'device'; flowType?: 'authorization_code' | 'device_code'; startEndpoint?: 'start' | 'start-url'; riskAcknowledgement?: { @@ -70,6 +73,7 @@ const INITIAL_STATE: AuthFlowState = { export function useCliproxyAuthFlow() { const [state, setState] = useState(INITIAL_STATE); + const stateRef = useRef(INITIAL_STATE); const attemptIdRef = useRef(0); const abortControllerRef = useRef(null); @@ -102,6 +106,10 @@ export function useCliproxyAuthFlow() { }; }, [stopPolling]); + useEffect(() => { + stateRef.current = state; + }, [state]); + // Poll OAuth status const pollStatus = useCallback( async (provider: string, oauthState: string, attemptId: number) => { @@ -253,6 +261,9 @@ export function useCliproxyAuthFlow() { const payload = { nickname: options?.nickname, kiroMethod: options?.kiroMethod, + kiroIDCStartUrl: options?.kiroIDCStartUrl, + kiroIDCRegion: options?.kiroIDCRegion, + kiroIDCFlow: options?.kiroIDCFlow, riskAcknowledgement: options?.riskAcknowledgement, }; @@ -368,6 +379,16 @@ export function useCliproxyAuthFlow() { // Start polling for completion if (oauthState) { pollStartRef.current = Date.now(); + if (!authUrl) { + await pollStatus(provider, oauthState, attemptId); + if (!isActiveAttempt(attemptId)) { + return; + } + const currentState = stateRef.current; + if (!currentState.isAuthenticating || currentState.provider !== provider) { + return; + } + } pollIntervalRef.current = setInterval(() => { void pollStatus(provider, oauthState, attemptId); }, POLL_INTERVAL); diff --git a/ui/src/lib/provider-config.ts b/ui/src/lib/provider-config.ts index b1cbf0d2..516b85e8 100644 --- a/ui/src/lib/provider-config.ts +++ b/ui/src/lib/provider-config.ts @@ -234,8 +234,11 @@ export function getDeviceCodeProviderInstruction(provider: unknown): string { } /** Kiro auth methods exposed in CCS UI (aligned with CLIProxyAPIPlus support). */ -export const KIRO_AUTH_METHODS = ['aws', 'aws-authcode', 'google', 'github'] as const; +export const KIRO_AUTH_METHODS = ['aws', 'aws-authcode', 'google', 'github', 'idc'] as const; export type KiroAuthMethod = (typeof KIRO_AUTH_METHODS)[number]; +export const KIRO_IDC_FLOWS = ['authcode', 'device'] as const; +export type KiroIDCFlow = (typeof KIRO_IDC_FLOWS)[number]; +export const DEFAULT_KIRO_IDC_FLOW: KiroIDCFlow = 'authcode'; export type KiroFlowType = 'authorization_code' | 'device_code'; export type KiroStartEndpoint = 'start' | 'start-url'; @@ -280,6 +283,13 @@ export const KIRO_AUTH_METHOD_OPTIONS: readonly KiroAuthMethodOption[] = [ flowType: 'authorization_code', startEndpoint: 'start-url', }, + { + id: 'idc', + label: 'AWS Identity Center (IDC)', + description: 'Use your organization start URL with auth code or device flow.', + flowType: 'authorization_code', + startEndpoint: 'start', + }, ]; export function isKiroAuthMethod(value: string): value is KiroAuthMethod { @@ -292,7 +302,40 @@ export function normalizeKiroAuthMethod(value?: string): KiroAuthMethod { return isKiroAuthMethod(normalized) ? normalized : DEFAULT_KIRO_AUTH_METHOD; } +export function isKiroIDCFlow(value: string): value is KiroIDCFlow { + return KIRO_IDC_FLOWS.includes(value as KiroIDCFlow); +} + +export function normalizeKiroIDCFlow(value?: string): KiroIDCFlow { + if (!value) return DEFAULT_KIRO_IDC_FLOW; + const normalized = value.trim().toLowerCase(); + return isKiroIDCFlow(normalized) ? normalized : DEFAULT_KIRO_IDC_FLOW; +} + export function getKiroAuthMethodOption(method: KiroAuthMethod): KiroAuthMethodOption { const option = KIRO_AUTH_METHOD_OPTIONS.find((candidate) => candidate.id === method); return option || KIRO_AUTH_METHOD_OPTIONS[0]; } + +export function getKiroEffectiveFlowType( + method: KiroAuthMethod, + idcFlow: KiroIDCFlow = DEFAULT_KIRO_IDC_FLOW +): KiroFlowType { + if (method === 'aws') { + return 'device_code'; + } + + if (method === 'idc') { + return normalizeKiroIDCFlow(idcFlow) === 'device' ? 'device_code' : 'authorization_code'; + } + + return 'authorization_code'; +} + +export function getKiroEffectiveStartEndpoint(method: KiroAuthMethod): KiroStartEndpoint { + return method === 'google' || method === 'github' ? 'start-url' : 'start'; +} + +export function isKiroSocialAuthMethod(method: KiroAuthMethod): boolean { + return method === 'google' || method === 'github'; +} diff --git a/ui/tests/unit/hooks/use-cliproxy-auth-flow.test.tsx b/ui/tests/unit/hooks/use-cliproxy-auth-flow.test.tsx index d90c4593..b37b8533 100644 --- a/ui/tests/unit/hooks/use-cliproxy-auth-flow.test.tsx +++ b/ui/tests/unit/hooks/use-cliproxy-auth-flow.test.tsx @@ -221,6 +221,96 @@ describe('useCliproxyAuthFlow', () => { expect(toast.success).toHaveBeenCalledWith('codex authentication successful'); }); + it('promotes a state-first auth bootstrap into an immediate auth URL without waiting for the first interval', async () => { + let pollCount = 0; + + vi.stubGlobal( + 'fetch', + vi.fn((input: RequestInfo | URL) => { + const url = String(input); + + if (url.includes('/start-url')) { + return Promise.resolve( + createJsonResponse({ + success: true, + authUrl: null, + state: 'state-kiro-social', + }) + ); + } + + if (url.includes('/status?state=state-kiro-social')) { + pollCount += 1; + return Promise.resolve( + createJsonResponse({ + status: 'auth_url', + url: 'https://auth.example/kiro-social', + }) + ); + } + + return Promise.reject(new Error(`Unexpected fetch: ${url}`)); + }) + ); + + const { result } = renderHook(() => useCliproxyAuthFlow(), { wrapper }); + + await act(async () => { + await result.current.startAuth('kiro', { startEndpoint: 'start-url', kiroMethod: 'google' }); + }); + + expect(result.current.authUrl).toBe('https://auth.example/kiro-social'); + expect(result.current.oauthState).toBe('state-kiro-social'); + expect(result.current.isAuthenticating).toBe(true); + expect(pollCount).toBe(1); + }); + + it('forwards Kiro IDC options to the backend start endpoint payload', async () => { + let requestBody: Record | null = null; + + vi.stubGlobal( + 'fetch', + vi.fn(async (input: RequestInfo | URL, init?: RequestInit) => { + const url = String(input); + + if (url.includes('/start')) { + requestBody = JSON.parse(String(init?.body)) as Record; + return createJsonResponse({ + success: true, + account: { + id: 'kiro-idc-account', + provider: 'kiro', + }, + }); + } + + return Promise.reject(new Error(`Unexpected fetch: ${url}`)); + }) + ); + + const { result } = renderHook(() => useCliproxyAuthFlow(), { wrapper }); + + await act(async () => { + await result.current.startAuth('kiro', { + startEndpoint: 'start', + flowType: 'authorization_code', + kiroMethod: 'idc', + kiroIDCStartUrl: 'https://d-123.awsapps.com/start', + kiroIDCRegion: 'ca-central-1', + kiroIDCFlow: 'authcode', + }); + }); + + expect(requestBody).toEqual({ + nickname: undefined, + kiroMethod: 'idc', + kiroIDCStartUrl: 'https://d-123.awsapps.com/start', + kiroIDCRegion: 'ca-central-1', + kiroIDCFlow: 'authcode', + riskAcknowledgement: undefined, + }); + }); + it('treats callback responses without an account as failures', async () => { vi.stubGlobal( 'fetch',