update

Signed-off-by: raj-subhankar <subhankar.rj@gmail.com>

Author: raj-subhankar

ui/desktop/src/components/ProviderGuard.tsx

@@ -5,6 +5,7 @@ import { SetupModal } from './SetupModal';
 import { startOpenRouterSetup } from '../utils/openRouterSetup';
 import { startTetrateSetup } from '../utils/tetrateSetup';
 import { startChatGptCodexSetup } from '../utils/chatgptCodexSetup';
+import { cancelTetrateSetup } from '../utils/tetrateSetup';
 import WelcomeGooseLogo from './WelcomeGooseLogo';
 import { toastService } from '../toasts';
 import { OllamaSetup } from './OllamaSetup';
@@ -37,7 +38,9 @@ export default function ProviderGuard({ didSelectProvider, children }: ProviderG
   const [userInActiveSetup, setUserInActiveSetup] = useState(false);
   const [showSwitchModelModal, setShowSwitchModelModal] = useState(false);
   const [switchModelProvider, setSwitchModelProvider] = useState<string | null>(null);
+  const [isTetrateSetupInProgress, setIsTetrateSetupInProgress] = useState(false);
   const onboardingTracked = useRef(false);
+  const tetrateSetupRunId = useRef(0);
   const scrollContainerRef = useRef<HTMLDivElement>(null);
   const [showScrollIndicator, setShowScrollIndicator] = useState(true);
 
@@ -66,7 +69,9 @@ export default function ProviderGuard({ didSelectProvider, children }: ProviderG
     show: boolean;
     title: string;
     message: string;
+    showProgress?: boolean;
     showRetry: boolean;
+    closeLabel?: string;
     autoClose?: number;
   } | null>(null);
 
@@ -79,10 +84,31 @@ export default function ProviderGuard({ didSelectProvider, children }: ProviderG
   } | null>(null);
 
   const handleTetrateSetup = async () => {
+    if (isTetrateSetupInProgress) {
+      return;
+    }
+
+    const runId = ++tetrateSetupRunId.current;
+    setIsTetrateSetupInProgress(true);
+    setTetrateSetupState({
+      show: true,
+      title: 'Complete setup in your browser',
+      message: 'After finishing sign-in in the browser, return to Goose. You can cancel setup anytime.',
+      showProgress: true,
+      showRetry: false,
+      closeLabel: 'Cancel Setup',
+    });
     trackOnboardingProviderSelected('tetrate');
     try {
       const result = await startTetrateSetup();
+      if (runId !== tetrateSetupRunId.current) {
+        return;
+      }
+
+      setIsTetrateSetupInProgress(false);
+
       if (result.success) {
+        setTetrateSetupState(null);
         setSwitchModelProvider('tetrate');
         setShowSwitchModelModal(true);
       } else {
@@ -95,6 +121,11 @@ export default function ProviderGuard({ didSelectProvider, children }: ProviderG
         });
       }
     } catch (error) {
+      if (runId !== tetrateSetupRunId.current) {
+        return;
+      }
+
+      setIsTetrateSetupInProgress(false);
       console.error('Tetrate setup error:', error);
       trackOnboardingSetupFailed('tetrate', 'unexpected_error');
       setTetrateSetupState({
@@ -134,6 +165,19 @@ export default function ProviderGuard({ didSelectProvider, children }: ProviderG
     }
   };
 
+  const handleCancelTetrateSetup = async () => {
+    if (!isTetrateSetupInProgress) {
+      setTetrateSetupState(null);
+      return;
+    }
+
+    tetrateSetupRunId.current += 1;
+    setIsTetrateSetupInProgress(false);
+    setTetrateSetupState(null);
+    trackOnboardingAbandoned('tetrate_setup');
+    await cancelTetrateSetup();
+  };
+
   const handleApiKeySuccess = async (provider: string, _model: string, apiKey: string) => {
     trackOnboardingProviderSelected('api_key');
     const keyName = `${provider.toUpperCase()}_API_KEY`;
@@ -217,7 +261,7 @@ export default function ProviderGuard({ didSelectProvider, children }: ProviderG
     if (setupType === 'openrouter') {
       setOpenRouterSetupState(null);
     } else if (setupType === 'tetrate') {
-      setTetrateSetupState(null);
+      void handleCancelTetrateSetup();
     } else {
       setChatgptCodexSetupState(null);
     }
@@ -366,8 +410,12 @@ export default function ProviderGuard({ didSelectProvider, children }: ProviderG
                 </div>
 
                 <div
-                  onClick={handleTetrateSetup}
-                  className="w-full p-4 sm:p-6 bg-transparent border rounded-xl transition-all duration-200 cursor-pointer group"
+                  onClick={isTetrateSetupInProgress ? undefined : handleTetrateSetup}
+                  className={`w-full p-4 sm:p-6 bg-transparent border border-background-hover rounded-xl transition-all duration-200 group ${
+                    isTetrateSetupInProgress
+                      ? 'cursor-not-allowed opacity-70'
+                      : 'cursor-pointer hover:border-text-muted'
+                  }`}
                 >
                   <div className="flex items-start justify-between mb-3">
                     <div className="flex items-center gap-2">
@@ -490,9 +538,11 @@ export default function ProviderGuard({ didSelectProvider, children }: ProviderG
           <SetupModal
             title={tetrateSetupState.title}
             message={tetrateSetupState.message}
+            showProgress={tetrateSetupState.showProgress}
             showRetry={tetrateSetupState.showRetry}
             onRetry={() => handleRetrySetup('tetrate')}
             onClose={() => closeSetupModal('tetrate')}
+            closeLabel={tetrateSetupState.closeLabel}
             autoClose={tetrateSetupState.autoClose}
           />
         )}

ui/desktop/src/components/SetupModal.tsx

@@ -9,6 +9,7 @@ interface SetupModalProps {
   onRetry?: () => void;
   autoClose?: number;
   onClose?: () => void;
+  closeLabel?: string;
 }
 
 export function SetupModal({
@@ -19,6 +20,7 @@ export function SetupModal({
   onRetry,
   autoClose,
   onClose,
+  closeLabel,
 }: SetupModalProps) {
   useEffect(() => {
     if (autoClose && onClose) {
@@ -45,7 +47,7 @@ export function SetupModal({
         {onClose && (
           <div className="mb-4">
             <Button onClick={onClose} className="w-full">
-              Close
+              {closeLabel || 'Close'}
             </Button>
             <br />
           </div>

ui/desktop/src/main.ts

@@ -48,6 +48,7 @@ import { Client, createClient, createConfig } from './api/client';
 import { GooseApp } from './api';
 import {
   TETRATE_AUTH_CLEANUP_INTERVAL_MS,
+  cancelTetrateAuthFlow,
   cleanupExpiredTetrateAuthFlows,
   handleTetrateCallbackUrl,
   runTetrateAuthFlow,
@@ -1266,6 +1267,10 @@ ipcMain.handle('tetrate-auth-start', async (event) => {
   return runTetrateAuthFlow(client);
 });
 
+ipcMain.handle('tetrate-auth-cancel', () => {
+  return cancelTetrateAuthFlow();
+});
+
 // Handle menu bar icon visibility
 ipcMain.handle('set-menu-bar-icon', async (_event, show: boolean) => {
   updateSettings((s) => {

ui/desktop/src/preload.ts

@@ -115,6 +115,7 @@ type ElectronAPI = {
   // Function to serve temp images
   getTempImage: (filePath: string) => Promise<string | null>;
   startTetrateAuth: () => Promise<{ success: boolean; message: string }>;
+  cancelTetrateAuth: () => Promise<boolean>;
   // Update-related functions
   getVersion: () => string;
   checkForUpdates: () => Promise<{ updateInfo: unknown; error: string | null }>;
@@ -238,6 +239,9 @@ const electronAPI: ElectronAPI = {
   startTetrateAuth: (): Promise<{ success: boolean; message: string }> => {
     return ipcRenderer.invoke('tetrate-auth-start');
   },
+  cancelTetrateAuth: (): Promise<boolean> => {
+    return ipcRenderer.invoke('tetrate-auth-cancel');
+  },
   getVersion: (): string => {
     return config.GOOSE_VERSION || ipcRenderer.sendSync('get-app-version') || '';
   },

ui/desktop/src/tetrateAuth.test.ts

@@ -19,7 +19,12 @@ vi.mock('./api', () => ({
 import { shell } from 'electron';
 import type { Client } from './api/client';
 import { verifyTetrateSetup } from './api';
-import { __test, handleTetrateCallbackUrl, runTetrateAuthFlow } from './tetrateAuth';
+import {
+  __test,
+  cancelTetrateAuthFlow,
+  handleTetrateCallbackUrl,
+  runTetrateAuthFlow,
+} from './tetrateAuth';
 
 describe('tetrateAuth', () => {
   afterEach(() => {
@@ -124,6 +129,21 @@ describe('tetrateAuth', () => {
     await expect(waitPromise).rejects.toThrow('access_denied');
   });
 
+  it('uses error_description when the auth callback includes one', async () => {
+    const { flowId, authUrl } = __test.createTetrateAuthFlow();
+    const callbackUrl = new URL(authUrl).searchParams.get('callback');
+    expect(callbackUrl).toBeTruthy();
+
+    const errorUrl = new URL(callbackUrl as string);
+    errorUrl.searchParams.set('error', 'access_denied');
+    errorUrl.searchParams.set('error_description', 'User denied authorization');
+
+    const waitPromise = __test.waitForTetrateCallback(flowId);
+    expect(handleTetrateCallbackUrl(errorUrl.toString())).toBe(true);
+
+    await expect(waitPromise).rejects.toThrow('User denied authorization');
+  });
+
   it('rejects immediately when the callback has no code and no error', async () => {
     const { flowId, authUrl } = __test.createTetrateAuthFlow();
     const callbackUrl = new URL(authUrl).searchParams.get('callback');
@@ -136,6 +156,19 @@ describe('tetrateAuth', () => {
     await expect(waitPromise).rejects.toThrow('Authentication failed');
   });
 
+  it('supports canceling an active auth flow', async () => {
+    const openExternalMock = vi.mocked(shell.openExternal);
+    openExternalMock.mockResolvedValue();
+
+    const flowPromise = runTetrateAuthFlow({} as Client);
+    expect(cancelTetrateAuthFlow()).toBe(true);
+
+    await expect(flowPromise).resolves.toEqual({
+      success: false,
+      message: 'Authentication canceled by user',
+    });
+  });
+
   it('runs the full auth flow and verifies the code', async () => {
     const verifyMock = vi.mocked(verifyTetrateSetup);
     const request = new globalThis.Request('http://localhost/test');

ui/desktop/src/tetrateAuth.ts

@@ -30,13 +30,16 @@ type TetrateCallbackMatch = {
   state: string;
   code?: string;
   error?: string;
+  errorDescription?: string;
 };
 
 const TETRATE_AUTH_URL = 'https://router.tetrate.ai/auth';
-const TETRATE_AUTH_TTL_MS = 10 * 60 * 1000;
+const TETRATE_AUTH_TTL_MS = 2 * 60 * 1000;
 const TETRATE_AUTH_CALLBACK_SCHEME = 'goose';
 
 const tetrateAuthFlows = new Map<string, TetrateAuthFlow>();
+const completedTetrateAuthFlowErrors = new Map<string, string>();
+let activeTetrateFlowId: string | null = null;
 
 function createPkcePair(): { codeVerifier: string; codeChallenge: string } {
   const codeVerifier = crypto.randomBytes(96).toString('base64url');
@@ -115,8 +118,10 @@ function matchTetrateCallbackUrl(callbackUrl: string): TetrateCallbackMatch | nu
   const result: TetrateCallbackMatch = { flowId, state };
   const code = parsedUrl.searchParams.get('code');
   const error = parsedUrl.searchParams.get('error');
+  const errorDescription = parsedUrl.searchParams.get('error_description');
   if (code) result.code = code;
   if (error) result.error = error;
+  if (errorDescription) result.errorDescription = errorDescription;
 
   return result;
 }
@@ -141,6 +146,9 @@ function expireTetrateAuthFlow(flowId: string, message: string): void {
   }
 
   log.warn('Tetrate auth flow expired:', { flowId, reason: message });
+  if (!flow.reject) {
+    completedTetrateAuthFlowErrors.set(flowId, message);
+  }
   flow.reject?.(new Error(message));
   cleanupTetrateAuthFlow(flowId);
 }
@@ -194,7 +202,8 @@ export function handleTetrateCallbackUrl(
   }
 
   if (match.error) {
-    expireTetrateAuthFlow(match.flowId, `Authentication denied: ${match.error}`);
+    const errorMessage = match.errorDescription || match.error;
+    expireTetrateAuthFlow(match.flowId, `Authentication denied: ${errorMessage}`);
     return true;
   }
 
@@ -211,6 +220,11 @@ export function handleTetrateCallbackUrl(
 function waitForTetrateCallback(flowId: string): Promise<string> {
   const flow = tetrateAuthFlows.get(flowId);
   if (!flow) {
+    const completedFlowError = completedTetrateAuthFlowErrors.get(flowId);
+    if (completedFlowError) {
+      completedTetrateAuthFlowErrors.delete(flowId);
+      return Promise.reject(new Error(completedFlowError));
+    }
     return Promise.reject(new Error('Authentication expired'));
   }
 
@@ -234,6 +248,16 @@ async function startTetrateAuthSession(flowId: string, authUrl: string): Promise
   return waitForTetrateCallback(flowId);
 }
 
+export function cancelTetrateAuthFlow(message = 'Authentication canceled by user'): boolean {
+  if (!activeTetrateFlowId) {
+    return false;
+  }
+  const flowId = activeTetrateFlowId;
+  activeTetrateFlowId = null;
+  expireTetrateAuthFlow(flowId, message);
+  return true;
+}
+
 function getTetrateAuthErrorMessage(error: unknown): string {
   if (error instanceof Error && error.message) {
     return error.message;
@@ -251,7 +275,15 @@ function getTetrateAuthErrorMessage(error: unknown): string {
 }
 
 export async function runTetrateAuthFlow(client: Client): Promise<TetrateSetupResponse> {
+  if (activeTetrateFlowId) {
+    return {
+      success: false,
+      message: 'Authentication already in progress',
+    };
+  }
+
   const { flowId, authUrl } = createTetrateAuthFlow();
+  activeTetrateFlowId = flowId;
 
   try {
     const callbackUrl = await startTetrateAuthSession(flowId, authUrl);
@@ -294,6 +326,10 @@ export async function runTetrateAuthFlow(client: Client): Promise<TetrateSetupRe
       success: false,
       message: getTetrateAuthErrorMessage(error),
     };
+  } finally {
+    if (activeTetrateFlowId === flowId) {
+      activeTetrateFlowId = null;
+    }
   }
 }
 
@@ -311,6 +347,8 @@ export const __test = {
       }
     }
     tetrateAuthFlows.clear();
+    completedTetrateAuthFlowErrors.clear();
+    activeTetrateFlowId = null;
   },
   waitForTetrateCallback,
 };

ui/desktop/src/utils/tetrateSetup.ts

@@ -11,3 +11,11 @@ export async function startTetrateSetup(): Promise<{
     };
   }
 }
+
+export async function cancelTetrateSetup(): Promise<boolean> {
+  try {
+    return await window.electron.cancelTetrateAuth();
+  } catch {
+    return false;
+  }
+}