Enable redirect-based OAuth flow

components/Chat/Chat.tsx

@@ -228,6 +228,7 @@
   const [interactionMessage, setInteractionMessage] = useState(null);
   const webSocketRef = useRef<WebSocket | null>(null);
   const webSocketConnectedRef = useRef(false);
+  const oauthPopupCancelledRef = useRef(false);
   const webSocketModeRef = useRef(
     sessionStorage.getItem('webSocketMode') === 'false' ? false : webSocketMode,
   );
@@ -252,7 +253,7 @@
    */
   const handleStopConversation = useCallback(() => {
     if (webSocketModeRef?.current) {
       console.log(
         'Stopping generation for user message:',
         activeUserMessageId.current,
       );
@@ -276,7 +277,7 @@
           controllerRef.current = new AbortController(); // Reset the controller
         }, 100);
       } catch (error) {
         console.log('error aborting - ', error);
       }
     }
   }, [webSocketModeRef, homeDispatch]);
@@ -291,7 +292,7 @@
     userResponse = '',
   }: any) => {
     if (!selectedConversation) {
       console.error(
         'Cannot send interaction response: no conversation selected',
       );
       return;
@@ -333,7 +334,7 @@
       homeDispatch({ field: 'messageIsStreaming', value: false });
       homeDispatch({ field: 'loading', value: false });
     }
   }, [selectedConversation?.id]);
 
   useEffect(() => {
     if (webSocketMode && !webSocketConnectedRef.current) {
@@ -349,7 +350,7 @@
         webSocketConnectedRef.current = false;
       }
     };
   }, [webSocketMode]);
 
   const connectWebSocket = async (retryCount = 0) => {
     const maxRetries = 3;
@@ -485,7 +486,7 @@
       };
 
       ws.onerror = (evt) => {
         console.error('[WebSocket] error:', evt);
         // Do not ws.close() here; let server/proxy drive closure
       };
 
@@ -495,7 +496,7 @@
       };
 
       ws.onclose = async (event) => {
         console.log(
           '[WebSocket] Connection closed. Code:',
           event.code,
           'Reason:',
@@ -545,6 +546,53 @@
     }
   }, [intermediateStepOverride]);
 
+  const persistOAuthPendingMessage = () => {
+    const conversation = selectedConversationRef.current;
+    if (!conversation) return;
+    const lastUserMessage = fetchLastMessage({ messages: conversation.messages, role: 'user' });
+    if (!lastUserMessage) return;
+    sessionStorage.setItem('oauth_pending_message', JSON.stringify(lastUserMessage));
+    sessionStorage.setItem('oauth_pending_conversation_id', conversation.id);
+  };
+
+  /**
+   * Handles OAuth consent flow by opening a popup window or navigating in the same tab
+   */
+  const handleOAuthConsent = (message: WebSocketInbound) => {
+    if (!isSystemInteractionMessage(message)) return false;
+
+    if (message.content?.input_type === 'oauth_consent') {
+      const oauthUrl = extractOAuthUrl(message);
+      if (oauthUrl) {
+        // Validate URL before opening
+        if (!isValidConsentPromptURL(oauthUrl)) {
+          console.error('OAuth URL validation failed in popup handler, refusing to open potentially malicious URL.');
+          toast.error('OAuth URL validation failed.');
+          return false;
+        }
+
+        const shouldUsePopup = !message.content?.use_redirect;
+        if (shouldUsePopup) {
+          const popup = window.open(
+            oauthUrl,
+            'oauth-popup',
+            'noopener,noreferrer'
+          );
+          const handleOAuthComplete = (event: MessageEvent) => {
+            if (popup && !popup.closed) popup.close();
+            window.removeEventListener('message', handleOAuthComplete);
+          };
+          window.addEventListener('message', handleOAuthComplete);
+        } else {
+          persistOAuthPendingMessage();
+          window.location.href = oauthUrl;
+        }
+      }
+      return true;
+    }
+    return false;
+  };
+
   /**
    * Updates refs immediately before React dispatch to prevent stale reads
    */
@@ -774,8 +822,23 @@
         if (oauthUrl) {
           // Validate URL before opening to prevent Open Redirect attacks
           if (isValidConsentPromptURL(oauthUrl)) {
-            // Open the validated OAuth URL in a new tab
-            window.open(oauthUrl, '_blank', 'noopener,noreferrer');
+            const shouldUsePopup = !message?.content?.use_redirect;
+            if (shouldUsePopup) {
+              if (oauthPopupCancelledRef.current) return;
+              // Open the validated OAuth URL in a new tab
+              const popup = window.open(oauthUrl, 'oauth-popup', 'noopener,noreferrer');
+              const handleOAuthComplete = (event: MessageEvent) => {
+                if (popup && !popup.closed) popup.close();
+                window.removeEventListener('message', handleOAuthComplete);
+                if (event.data?.type === 'AUTH_CANCELLED') {
+                  oauthPopupCancelledRef.current = true;
+                }
+              };
+              window.addEventListener('message', handleOAuthComplete);
+            } else {
+              persistOAuthPendingMessage();
+              window.location.href = oauthUrl;
+            }
           } else {
             console.error(
               'OAuth URL validation failed, refusing to open potentially malicious URL:',
@@ -863,6 +926,7 @@
   const handleSend = useCallback(
     async (message: Message, deleteCount = 0, _retry = false) => {
       message.id = uuidv4();
+      oauthPopupCancelledRef.current = false;
 
       // Set the active user message ID for WebSocket message tracking
       activeUserMessageId.current = message.id;
@@ -1488,6 +1552,63 @@
     [handleSend],
   );
 
+  // After returning from the OAuth provider, resubmit the message that triggered auth.
+  useEffect(() => {
+    const pendingMessageRaw = sessionStorage.getItem('oauth_pending_message');
+    const pendingConversationId = sessionStorage.getItem('oauth_pending_conversation_id');
+    if (!pendingMessageRaw || !pendingConversationId) return;
+    if (!selectedConversation || selectedConversation.id !== pendingConversationId) return;
+
+    // The success page runs at the NAT server origin, so sessionStorage is cross-origin and
+    // unavailable here. The flag is instead passed back as a URL query parameter.
+    const urlParams = new URLSearchParams(window.location.search);
+    const authCompleted = urlParams.get('oauth_auth_completed');
+    if (authCompleted) {
+      urlParams.delete('oauth_auth_completed');
+      const cleanUrl = window.location.pathname + (urlParams.toString() ? '?' + urlParams.toString() : '');
+      window.history.replaceState({}, '', cleanUrl);
+    }
+
+    sessionStorage.removeItem('oauth_pending_message');
+    sessionStorage.removeItem('oauth_pending_conversation_id');
+
+    // If the user pressed back without completing OAuth, show a cancellation message.
+    if (!authCompleted) {
+      const conversation = selectedConversationRef.current;
+      if (conversation) {
+        const messages = conversation.messages;
+        const lastMessage = messages.at(-1);
+        const updatedMessages = lastMessage?.role === 'assistant'
+          ? messages.map((m, idx) =>
+              idx === messages.length - 1 ? updateAssistantMessage(m, 'Authorization cancelled.') : m
+            )
+          : [...messages, createAssistantMessage(undefined, undefined, 'Authorization cancelled.')];
+        const updatedConversation = { ...conversation, messages: updatedMessages };
+        const updatedConversations = conversationsRef.current.map(c =>
+          c.id === updatedConversation.id ? updatedConversation : c
+        );
+        updateRefsAndDispatch(updatedConversations, updatedConversation, conversation);
+      }
+      return;
+    }
+
+    const resume = async () => {
+      let pendingMessage: Message;
+      try {
+        pendingMessage = JSON.parse(pendingMessageRaw);
+      } catch {
+        return;
+      }
+      // Ensure the WebSocket is connected before calling handleSend
+      if (webSocketModeRef.current && !webSocketConnectedRef.current) {
+        await connectWebSocket();
+      }
+      // Delete the user message + empty assistant placeholder appended during original send, then resubmit.
+      handleSend(pendingMessage, 2);
+    };
+    resume();
+  }, [selectedConversation?.id]);
+
   // Add a new effect to handle streaming state changes
   useEffect(() => {
     if (messageIsStreaming) {

pages/api/home/home.tsx

@@ -214,6 +214,11 @@ const Home = (_props: any) => {
       dispatch({ field: 'showChatbar', value: showChatbar === 'true' });
     }
 
+    const webSocketMode = sessionStorage.getItem('webSocketMode');
+    if (webSocketMode !== null) {
+      dispatch({ field: 'webSocketMode', value: webSocketMode === 'true' });
+    }
+
     const enableIntermediateSteps = sessionStorage.getItem(
       'enableIntermediateSteps',
     );

types/websocket.ts

@@ -44,6 +44,7 @@ export interface SystemInteractionMessage extends WebSocketMessageBase {
     text?: string;
     timeout?: number | null;
     error?: string | null;
+    use_redirect?: boolean;
   };
   thread_id?: string;
 }