diff --git a/src/cliproxy/auth-session-manager.ts b/src/cliproxy/auth-session-manager.ts new file mode 100644 index 00000000..ce718be9 --- /dev/null +++ b/src/cliproxy/auth-session-manager.ts @@ -0,0 +1,119 @@ +/** + * Auth Session Manager + * + * Tracks active OAuth sessions and provides cancellation capability. + * Used to properly terminate in-progress OAuth flows from UI. + */ + +import { EventEmitter } from 'events'; +import { ChildProcess } from 'child_process'; + +export interface ActiveAuthSession { + sessionId: string; + provider: string; + startedAt: number; + process?: ChildProcess; +} + +export const authSessionEvents = new EventEmitter(); + +const activeSessions = new Map(); + +/** + * Register an active OAuth session + */ +export function registerAuthSession( + sessionId: string, + provider: string, + process?: ChildProcess +): void { + activeSessions.set(sessionId, { + sessionId, + provider, + startedAt: Date.now(), + process, + }); + authSessionEvents.emit('session:started', sessionId, provider); +} + +/** + * Update session with process reference (if registered before spawn) + */ +export function attachProcessToSession(sessionId: string, process: ChildProcess): void { + const session = activeSessions.get(sessionId); + if (session) { + session.process = process; + } +} + +/** + * Unregister an auth session (on completion or cancellation) + */ +export function unregisterAuthSession(sessionId: string): void { + activeSessions.delete(sessionId); + authSessionEvents.emit('session:ended', sessionId); +} + +/** + * Cancel an active OAuth session + * Returns true if session was found and killed + */ +export function cancelAuthSession(sessionId: string): boolean { + const session = activeSessions.get(sessionId); + if (!session) { + return false; + } + + // Kill the process if attached + if (session.process && !session.process.killed) { + session.process.kill('SIGTERM'); + } + + activeSessions.delete(sessionId); + authSessionEvents.emit('session:cancelled', sessionId); + return true; +} + +/** + * Get active session by session ID + */ +export function getActiveSession(sessionId: string): ActiveAuthSession | null { + return activeSessions.get(sessionId) || null; +} + +/** + * Get active session for a provider (most recent) + */ +export function getActiveSessionForProvider(provider: string): ActiveAuthSession | null { + for (const session of activeSessions.values()) { + if (session.provider === provider) { + return session; + } + } + return null; +} + +/** + * Check if there's an active session for provider + */ +export function hasActiveSession(provider: string): boolean { + return getActiveSessionForProvider(provider) !== null; +} + +/** + * Cancel all sessions for a provider + */ +export function cancelAllSessionsForProvider(provider: string): number { + let count = 0; + for (const [sessionId, session] of activeSessions.entries()) { + if (session.provider === provider) { + if (session.process && !session.process.killed) { + session.process.kill('SIGTERM'); + } + activeSessions.delete(sessionId); + authSessionEvents.emit('session:cancelled', sessionId); + count++; + } + } + return count; +} diff --git a/src/cliproxy/auth/oauth-process.ts b/src/cliproxy/auth/oauth-process.ts index 309df84d..c09a9536 100644 --- a/src/cliproxy/auth/oauth-process.ts +++ b/src/cliproxy/auth/oauth-process.ts @@ -25,6 +25,12 @@ import { getTimeoutTroubleshooting, showStep } from './environment-detector'; import { isAuthenticated, registerAccountFromToken } from './token-manager'; import { deviceCodeEvents, type DeviceCodePrompt } from '../device-code-handler'; import { OAUTH_FLOW_TYPES } from '../../management'; +import { + registerAuthSession, + attachProcessToSession, + unregisterAuthSession, + authSessionEvents, +} from '../auth-session-manager'; /** Options for OAuth process execution */ export interface OAuthProcessOptions { @@ -326,6 +332,19 @@ export function executeOAuthProcess(options: OAuthProcessOptions): Promise { + if (cancelledSessionId === state.sessionId && authProcess && !authProcess.killed) { + log('Session cancelled externally'); + authProcess.kill('SIGTERM'); + } + }; + authSessionEvents.on('session:cancelled', handleCancel); + const startTime = Date.now(); authProcess.stdout?.on('data', async (data: Buffer) => { @@ -377,6 +396,8 @@ export function executeOAuthProcess(options: OAuthProcessOptions): Promise { + const { provider } = req.params; + + // Validate provider + if (!validProviders.includes(provider as CLIProxyProvider)) { + res.status(400).json({ error: `Invalid provider: ${provider}` }); + return; + } + + // Check if there's an active session + if (!hasActiveSession(provider)) { + res.status(404).json({ error: 'No active authentication session for this provider' }); + return; + } + + // Cancel all sessions for this provider + const cancelledCount = cancelAllSessionsForProvider(provider); + + res.json({ + success: true, + cancelled: cancelledCount, + provider, + }); +}); + /** * GET /api/cliproxy/auth/project-selection/:sessionId - Get pending project selection prompt * Returns project list for user to select from during OAuth flow diff --git a/ui/src/components/account/add-account-dialog.tsx b/ui/src/components/account/add-account-dialog.tsx index c8276ace..0290bda3 100644 --- a/ui/src/components/account/add-account-dialog.tsx +++ b/ui/src/components/account/add-account-dialog.tsx @@ -17,7 +17,7 @@ import { Button } from '@/components/ui/button'; import { Input } from '@/components/ui/input'; import { Label } from '@/components/ui/label'; import { Loader2, ExternalLink, User, Download } from 'lucide-react'; -import { useStartAuth, useKiroImport } from '@/hooks/use-cliproxy'; +import { useStartAuth, useKiroImport, useCancelAuth } from '@/hooks/use-cliproxy'; import { applyDefaultPreset } from '@/lib/preset-utils'; import { toast } from 'sonner'; @@ -40,10 +40,19 @@ export function AddAccountDialog({ const [nickname, setNickname] = useState(''); const startAuthMutation = useStartAuth(); const kiroImportMutation = useKiroImport(); + const cancelAuthMutation = useCancelAuth(); const isKiro = provider === 'kiro'; const isPending = startAuthMutation.isPending || kiroImportMutation.isPending; + const handleCancel = () => { + if (isPending) { + cancelAuthMutation.mutate(provider); + } + setNickname(''); + onClose(); + }; + const handleStartAuth = () => { startAuthMutation.mutate( { provider, nickname: nickname.trim() || undefined }, @@ -83,9 +92,8 @@ export function AddAccountDialog({ }; const handleOpenChange = (isOpen: boolean) => { - if (!isOpen && !isPending) { - setNickname(''); - onClose(); + if (!isOpen) { + handleCancel(); } }; @@ -121,7 +129,7 @@ export function AddAccountDialog({
- {isKiro && ( diff --git a/ui/src/components/setup/wizard/index.tsx b/ui/src/components/setup/wizard/index.tsx index cea897c4..29e98e76 100644 --- a/ui/src/components/setup/wizard/index.tsx +++ b/ui/src/components/setup/wizard/index.tsx @@ -15,7 +15,12 @@ import { DialogDescription, } from '@/components/ui/dialog'; import { Sparkles } from 'lucide-react'; -import { useCliproxyAuth, useCreateVariant, useStartAuth } from '@/hooks/use-cliproxy'; +import { + useCliproxyAuth, + useCreateVariant, + useStartAuth, + useCancelAuth, +} from '@/hooks/use-cliproxy'; import type { AuthStatus, OAuthAccount } from '@/lib/api-client'; import { applyDefaultPreset } from '@/lib/preset-utils'; import { usePrivacy } from '@/contexts/privacy-context'; @@ -42,6 +47,7 @@ export function QuickSetupWizard({ open, onClose }: QuickSetupWizardProps) { const { data: authData, refetch } = useCliproxyAuth(); const createMutation = useCreateVariant(); const startAuthMutation = useStartAuth(); + const cancelAuthMutation = useCancelAuth(); const { privacyMode } = usePrivacy(); // Get auth status for selected provider @@ -146,6 +152,10 @@ export function QuickSetupWizard({ open, onClose }: QuickSetupWizardProps) { // Prevent accidental close when user has made progress const handleOpenChange = (isOpen: boolean) => { if (!isOpen) { + // Cancel any in-progress auth when closing + if (startAuthMutation.isPending && selectedProvider) { + cancelAuthMutation.mutate(selectedProvider); + } if (step === 'success' || step === 'provider') { onClose(); return; diff --git a/ui/src/hooks/use-cliproxy-auth-flow.ts b/ui/src/hooks/use-cliproxy-auth-flow.ts index 5bc452b5..a91a05c3 100644 --- a/ui/src/hooks/use-cliproxy-auth-flow.ts +++ b/ui/src/hooks/use-cliproxy-auth-flow.ts @@ -6,6 +6,7 @@ import { useState, useCallback, useRef, useEffect, useMemo } from 'react'; import { useQueryClient } from '@tanstack/react-query'; import { toast } from 'sonner'; +import { api } from '@/lib/api-client'; interface AuthFlowState { provider: string | null; @@ -82,9 +83,16 @@ export function useCliproxyAuthFlow() { ); const cancelAuth = useCallback(() => { + const currentProvider = state.provider; abortControllerRef.current?.abort(); setState({ provider: null, isAuthenticating: false, error: null }); - }, []); + // Also cancel on backend + if (currentProvider) { + api.cliproxy.auth.cancel(currentProvider).catch(() => { + // Ignore errors - session may have already completed + }); + } + }, [state.provider]); return useMemo( () => ({ diff --git a/ui/src/hooks/use-cliproxy.ts b/ui/src/hooks/use-cliproxy.ts index 7dd7811b..d4b6ec9c 100644 --- a/ui/src/hooks/use-cliproxy.ts +++ b/ui/src/hooks/use-cliproxy.ts @@ -136,6 +136,16 @@ export function useStartAuth() { }); } +// Cancel OAuth flow hook +export function useCancelAuth() { + return useMutation({ + mutationFn: (provider: string) => api.cliproxy.auth.cancel(provider), + onError: (error: Error) => { + toast.error(error.message); + }, + }); +} + // Kiro IDE import hook (alternative auth path when OAuth callback fails) export function useKiroImport() { const queryClient = useQueryClient(); diff --git a/ui/src/lib/api-client.ts b/ui/src/lib/api-client.ts index 0a003cb4..60e3ec02 100644 --- a/ui/src/lib/api-client.ts +++ b/ui/src/lib/api-client.ts @@ -357,6 +357,12 @@ export const api = { method: 'POST', body: JSON.stringify({ nickname }), }), + /** Cancel in-progress OAuth flow */ + cancel: (provider: string) => + request<{ success: boolean; cancelled: number; provider: string }>( + `/cliproxy/auth/${provider}/cancel`, + { method: 'POST' } + ), /** Import Kiro token from Kiro IDE (Kiro only) */ kiroImport: () => request<{ success: boolean; account: OAuthAccount | null; error?: string }>(