update

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

Author: raj-subhankar

ui/desktop/src/components/ProviderGuard.tsx

@@ -3,7 +3,7 @@ import { useNavigate } from 'react-router-dom';
 import { useConfig } from './ConfigContext';
 import { SetupModal } from './SetupModal';
 import { startOpenRouterSetup } from '../utils/openRouterSetup';
-import { startTetrateSetup } from '../utils/tetrateSetup';
+import { cancelTetrateSetup, startTetrateSetup } from '../utils/tetrateSetup';
 import WelcomeGooseLogo from './WelcomeGooseLogo';
 import { toastService } from '../toasts';
 import { OllamaSetup } from './OllamaSetup';
@@ -36,7 +36,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);
 
@@ -65,15 +67,38 @@ export default function ProviderGuard({ didSelectProvider, children }: ProviderG
     show: boolean;
     title: string;
     message: string;
+    showProgress?: boolean;
     showRetry: boolean;
+    closeLabel?: string;
     autoClose?: number;
   } | 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 {
@@ -86,6 +111,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({
@@ -97,6 +127,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`;
@@ -177,7 +220,7 @@ export default function ProviderGuard({ didSelectProvider, children }: ProviderG
     if (setupType === 'openrouter') {
       setOpenRouterSetupState(null);
     } else {
-      setTetrateSetupState(null);
+      void handleCancelTetrateSetup();
     }
   };
 
@@ -284,8 +327,12 @@ export default function ProviderGuard({ didSelectProvider, children }: ProviderG
                 </div>
 
                 <div
-                  onClick={handleTetrateSetup}
-                  className="w-full p-4 sm:p-6 bg-transparent border border-background-hover rounded-xl hover:border-text-muted 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">
@@ -408,9 +455,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

@@ -49,7 +49,7 @@ import {
 import { UPDATES_ENABLED } from './updates';
 import './utils/recipeHash';
 import { Client, createClient, createConfig } from './api/client';
-import { handleTetrateCallbackUrl, runTetrateAuthFlow } from './tetrateAuth';
+import { cancelTetrateAuthFlow, handleTetrateCallbackUrl, runTetrateAuthFlow } from './tetrateAuth';
 import { GooseApp } from './api';
 import installExtension, { REACT_DEVELOPER_TOOLS } from 'electron-devtools-installer';
 
@@ -1321,6 +1321,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) => {
   try {

ui/desktop/src/preload.ts

@@ -124,6 +124,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 }>;
@@ -251,6 +252,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;
+  }
+}