fix(ui): improve Kiro dashboard auth flow

This commit is contained in:
Tam Nhu Tran
2026-04-06 15:44:34 -04:00
parent 1434c3d1e0
commit 699b192137
4 changed files with 271 additions and 9 deletions
@@ -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<KiroAuthMethod>(DEFAULT_KIRO_AUTH_METHOD);
const [kiroIDCStartUrl, setKiroIDCStartUrl] = useState('');
const [kiroIDCRegion, setKiroIDCRegion] = useState('');
const [kiroIDCFlow, setKiroIDCFlow] = useState<KiroIDCFlow>(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<boolean> => {
@@ -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({
</SelectContent>
</Select>
<p className="text-xs text-muted-foreground">{kiroMethodOption.description}</p>
{isKiroSocial && (
<p className="text-xs text-muted-foreground">
If your browser does not return automatically after login, CCS can accept the
final
<span className="mx-1 rounded bg-muted px-1.5 py-0.5 font-mono text-[11px]">
kiro://...
</span>
callback URL in the next step.
</p>
)}
</div>
)}
{isKiroIdc && !showAuthUI && (
<div className="space-y-4 rounded-lg border border-border/60 bg-muted/20 p-4">
<div className="space-y-2">
<Label htmlFor="kiro-idc-start-url">IDC Start URL</Label>
<Input
id="kiro-idc-start-url"
value={kiroIDCStartUrl}
onChange={(e) => {
setKiroIDCStartUrl(e.target.value);
setLocalError(null);
}}
placeholder="https://d-xxx.awsapps.com/start"
disabled={isPending}
/>
<p className="text-xs text-muted-foreground">
Required for organization IAM Identity Center login.
</p>
</div>
<div className="space-y-2">
<Label htmlFor="kiro-idc-region">IDC Region</Label>
<Input
id="kiro-idc-region"
value={kiroIDCRegion}
onChange={(e) => {
setKiroIDCRegion(e.target.value);
setLocalError(null);
}}
placeholder="us-east-1"
disabled={isPending}
/>
<p className="text-xs text-muted-foreground">
Optional. Leave blank to use the upstream default region.
</p>
</div>
<div className="space-y-2">
<Label htmlFor="kiro-idc-flow">IDC Flow</Label>
<Select
value={kiroIDCFlow}
onValueChange={(value) => {
setKiroIDCFlow(value as KiroIDCFlow);
setLocalError(null);
}}
>
<SelectTrigger id="kiro-idc-flow">
<SelectValue placeholder="Select IDC flow" />
</SelectTrigger>
<SelectContent>
<SelectItem value="authcode">Authorization Code</SelectItem>
<SelectItem value="device">Device Code</SelectItem>
</SelectContent>
</Select>
<p className="text-xs text-muted-foreground">
Auth Code opens a browser and may need the final callback URL pasted back. Device
Code shows a verification code instead.
</p>
</div>
</div>
)}
@@ -438,7 +536,9 @@ export function AddAccountDialog({
<p className="text-xs text-muted-foreground mt-1">
{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')}
</p>
</div>
@@ -486,13 +586,19 @@ export function AddAccountDialog({
{/* Callback paste field */}
<div className="space-y-2">
<Label htmlFor="callback-url" className="text-xs">
{t('addAccountDialog.redirectPasteLabel')}
{isKiroSocial
? 'Browser did not return? Paste the final kiro:// callback URL:'
: t('addAccountDialog.redirectPasteLabel')}
</Label>
<Input
id="callback-url"
value={callbackUrl}
onChange={(e) => 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"
/>
<Button
@@ -516,7 +622,9 @@ export function AddAccountDialog({
{!authFlow.authUrl && !authFlow.isDeviceCodeFlow && (
<p className="text-xs text-center text-muted-foreground">
{t('addAccountDialog.preparingUrl')}
{isKiroSocial
? 'Preparing the Kiro sign-in URL. If it does not open automatically, it will appear here shortly.'
: t('addAccountDialog.preparingUrl')}
</p>
)}
</div>
+21
View File
@@ -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<AuthFlowState>(INITIAL_STATE);
const stateRef = useRef<AuthFlowState>(INITIAL_STATE);
const attemptIdRef = useRef(0);
const abortControllerRef = useRef<AbortController | null>(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);
+44 -1
View File
@@ -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';
}
@@ -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<string, unknown> | 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<string, unknown>;
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',