From 6d1f257ad803227f115df3d7364ff1db15df81e0 Mon Sep 17 00:00:00 2001 From: zfoong Date: Sat, 21 Mar 2026 09:57:34 +0900 Subject: [PATCH] Fix browser interface freezing issue when oauth failed --- agent_core/__init__.py | 2 + agent_core/core/credentials/__init__.py | 3 +- agent_core/core/credentials/oauth_server.py | 186 +++++++++++++----- app/credentials/handlers.py | 20 +- app/ui_layer/adapters/browser_adapter.py | 61 +++++- .../src/pages/Settings/SettingsPage.tsx | 56 ++++-- 6 files changed, 258 insertions(+), 70 deletions(-) diff --git a/agent_core/__init__.py b/agent_core/__init__.py index ee6a3fd6..d0757090 100644 --- a/agent_core/__init__.py +++ b/agent_core/__init__.py @@ -75,6 +75,7 @@ get_credentials, has_embedded_credentials, run_oauth_flow, + run_oauth_flow_async, ) from agent_core.core.config import ( ConfigRegistry, @@ -312,6 +313,7 @@ "get_credentials", "has_embedded_credentials", "run_oauth_flow", + "run_oauth_flow_async", # Config "ConfigRegistry", "get_workspace_root", diff --git a/agent_core/core/credentials/__init__.py b/agent_core/core/credentials/__init__.py index 39200ffc..055a6c77 100644 --- a/agent_core/core/credentials/__init__.py +++ b/agent_core/core/credentials/__init__.py @@ -8,7 +8,7 @@ encode_credential, generate_credentials_block, ) -from agent_core.core.credentials.oauth_server import run_oauth_flow +from agent_core.core.credentials.oauth_server import run_oauth_flow, run_oauth_flow_async __all__ = [ "get_credential", @@ -17,4 +17,5 @@ "encode_credential", "generate_credentials_block", "run_oauth_flow", + "run_oauth_flow_async", ] diff --git a/agent_core/core/credentials/oauth_server.py b/agent_core/core/credentials/oauth_server.py index ac9f4770..9d8a701f 100644 --- a/agent_core/core/credentials/oauth_server.py +++ b/agent_core/core/credentials/oauth_server.py @@ -16,8 +16,12 @@ # HTTPS (for Slack and other providers requiring https redirect URIs) code, error = run_oauth_flow("https://slack.com/oauth/...", use_https=True) + + # Async version with cancellation support (recommended for UI contexts) + code, error = await run_oauth_flow_async("https://provider.com/oauth/...") """ +import asyncio import ipaddress import logging import os @@ -29,7 +33,7 @@ from datetime import datetime, timedelta, timezone from http.server import HTTPServer, BaseHTTPRequestHandler from urllib.parse import urlparse, parse_qs -from typing import Optional, Tuple +from typing import Any, Dict, Optional, Tuple logger = logging.getLogger(__name__) @@ -104,58 +108,78 @@ def _cleanup_files(*paths: str) -> None: pass -class _OAuthCallbackHandler(BaseHTTPRequestHandler): - """Handler for OAuth callback requests.""" - - code: Optional[str] = None - state: Optional[str] = None - error: Optional[str] = None - - def do_GET(self): - """Handle GET request from OAuth callback.""" - params = parse_qs(urlparse(self.path).query) - _OAuthCallbackHandler.code = params.get("code", [None])[0] - _OAuthCallbackHandler.state = params.get("state", [None])[0] - _OAuthCallbackHandler.error = params.get("error", [None])[0] +def _make_callback_handler(result_holder: Dict[str, Any]): + """ + Create a callback handler class that stores results in the provided dict. - self.send_response(200) - self.send_header("Content-Type", "text/html") - self.end_headers() - if _OAuthCallbackHandler.code: - self.wfile.write( - b"

Authorization successful!

You can close this tab.

" - ) - else: - self.wfile.write( - f"

Failed

{_OAuthCallbackHandler.error}

".encode() - ) + This avoids class-level state that would be shared across OAuth flows. + """ + class _OAuthCallbackHandler(BaseHTTPRequestHandler): + """Handler for OAuth callback requests.""" + + def do_GET(self): + """Handle GET request from OAuth callback.""" + params = parse_qs(urlparse(self.path).query) + result_holder["code"] = params.get("code", [None])[0] + result_holder["state"] = params.get("state", [None])[0] + result_holder["error"] = params.get("error", [None])[0] + + self.send_response(200) + self.send_header("Content-Type", "text/html") + self.end_headers() + if result_holder["code"]: + self.wfile.write( + b"

Authorization successful!

You can close this tab.

" + ) + else: + self.wfile.write( + f"

Failed

{result_holder['error']}

".encode() + ) + + def log_message(self, format, *args): + """Suppress default HTTP server logging.""" + pass - def log_message(self, format, *args): - """Suppress default HTTP server logging.""" - pass + return _OAuthCallbackHandler -def _serve_until_code(server: HTTPServer, deadline: float) -> None: +def _serve_until_code( + server: HTTPServer, + deadline: float, + result_holder: Dict[str, Any], + cancel_event: Optional[threading.Event] = None, +) -> None: """ - Handle requests in a loop until we capture the OAuth code/error or timeout. + Handle requests in a loop until we capture the OAuth code/error, timeout, or cancelled. A single handle_request() can be consumed by TLS handshake failures, favicon requests, browser pre-connects, etc. Looping ensures the server stays alive for the actual callback. """ while time.time() < deadline: - remaining = max(0.5, deadline - time.time()) - server.timeout = min(remaining, 2.0) + # Check for cancellation + if cancel_event and cancel_event.is_set(): + logger.debug("[OAUTH] Cancellation requested, stopping server") + break + + remaining = max(0.1, deadline - time.time()) + # Use shorter timeout (0.5s) for responsive cancellation checking + server.timeout = min(remaining, 0.5) try: server.handle_request() except Exception as e: logger.debug(f"[OAUTH] handle_request error (will retry): {e}") - if _OAuthCallbackHandler.code or _OAuthCallbackHandler.error: + + if result_holder.get("code") or result_holder.get("error"): break def run_oauth_flow( - auth_url: str, port: int = 8765, timeout: int = 120, use_https: bool = False + auth_url: str, + port: int = 8765, + timeout: int = 120, + use_https: bool = False, + cancel_event: Optional[threading.Event] = None, ) -> Tuple[Optional[str], Optional[str]]: """ Open browser for OAuth, wait for callback. @@ -167,17 +191,27 @@ def run_oauth_flow( use_https: If True, serve HTTPS with a self-signed cert. Required for providers like Slack that reject http:// redirect URIs. Default False (plain HTTP — works with Google, Notion, etc.). + cancel_event: Optional threading.Event to signal cancellation. + When set, the OAuth flow will stop and return a cancellation error. Returns: Tuple of (code, error_message): - On success: (authorization_code, None) - On failure: (None, error_message) """ - _OAuthCallbackHandler.code = None - _OAuthCallbackHandler.state = None - _OAuthCallbackHandler.error = None + # Check for early cancellation + if cancel_event and cancel_event.is_set(): + return None, "OAuth cancelled" - server = HTTPServer(("127.0.0.1", port), _OAuthCallbackHandler) + # Use instance-level result holder instead of class-level state + result_holder: Dict[str, Any] = {"code": None, "state": None, "error": None} + handler_class = _make_callback_handler(result_holder) + + try: + server = HTTPServer(("127.0.0.1", port), handler_class) + except OSError as e: + # Port already in use + return None, f"Failed to start OAuth server: {e}" if use_https: cert_path = key_path = None @@ -198,21 +232,85 @@ def run_oauth_flow( deadline = time.time() + timeout thread = threading.Thread( - target=_serve_until_code, args=(server, deadline), daemon=True + target=_serve_until_code, + args=(server, deadline, result_holder, cancel_event), + daemon=True ) thread.start() + # Check cancellation before opening browser + if cancel_event and cancel_event.is_set(): + server.server_close() + return None, "OAuth cancelled" + try: webbrowser.open(auth_url) except Exception: server.server_close() return None, f"Could not open browser. Visit manually:\n{auth_url}" - thread.join(timeout=timeout) + # Wait for thread with periodic cancellation checks + while thread.is_alive(): + thread.join(timeout=0.5) + if cancel_event and cancel_event.is_set(): + logger.debug("[OAUTH] Cancellation detected during wait") + break + server.server_close() - if _OAuthCallbackHandler.error: - return None, _OAuthCallbackHandler.error - if _OAuthCallbackHandler.code: - return _OAuthCallbackHandler.code, None + # Check cancellation first + if cancel_event and cancel_event.is_set(): + return None, "OAuth cancelled" + + if result_holder.get("error"): + return None, result_holder["error"] + if result_holder.get("code"): + return result_holder["code"], None return None, "OAuth timed out." + + +async def run_oauth_flow_async( + auth_url: str, + port: int = 8765, + timeout: int = 120, + use_https: bool = False, +) -> Tuple[Optional[str], Optional[str]]: + """ + Async version of run_oauth_flow with proper cancellation support. + + This function runs the OAuth flow in a thread executor and properly handles + asyncio task cancellation by signaling the OAuth server to stop. + + Args: + auth_url: The full OAuth authorization URL to open. + port: Local port for callback server (default: 8765). + timeout: Seconds to wait for callback (default: 120). + use_https: If True, serve HTTPS with a self-signed cert. + + Returns: + Tuple of (code, error_message): + - On success: (authorization_code, None) + - On failure: (None, error_message) + + Raises: + asyncio.CancelledError: If the task is cancelled (after signaling OAuth to stop) + """ + cancel_event = threading.Event() + loop = asyncio.get_event_loop() + + def run_flow(): + return run_oauth_flow( + auth_url=auth_url, + port=port, + timeout=timeout, + use_https=use_https, + cancel_event=cancel_event, + ) + + try: + return await loop.run_in_executor(None, run_flow) + except asyncio.CancelledError: + # Signal the OAuth server to stop + cancel_event.set() + logger.debug("[OAUTH] Async task cancelled, signaled OAuth server to stop") + raise diff --git a/app/credentials/handlers.py b/app/credentials/handlers.py index c8eceb31..1924910f 100644 --- a/app/credentials/handlers.py +++ b/app/credentials/handlers.py @@ -76,8 +76,8 @@ async def login(self, args): "code_challenge": code_challenge, "code_challenge_method": "S256", } - from agent_core import run_oauth_flow - code, error = run_oauth_flow(f"https://accounts.google.com/o/oauth2/v2/auth?{urlencode(params)}") + from agent_core import run_oauth_flow_async + code, error = await run_oauth_flow_async(f"https://accounts.google.com/o/oauth2/v2/auth?{urlencode(params)}") if error: return False, f"Google OAuth failed: {error}" token_data = { @@ -141,8 +141,8 @@ async def invite(self, args): scopes = "chat:write,channels:read,channels:history,groups:read,groups:history,users:read,files:write,im:read,im:write,im:history" params = {"client_id": SLACK_SHARED_CLIENT_ID, "scope": scopes, "redirect_uri": REDIRECT_URI_HTTPS, "state": secrets.token_urlsafe(32)} - from agent_core import run_oauth_flow - code, error = run_oauth_flow(f"https://slack.com/oauth/v2/authorize?{urlencode(params)}", use_https=True) + from agent_core import run_oauth_flow_async + code, error = await run_oauth_flow_async(f"https://slack.com/oauth/v2/authorize?{urlencode(params)}", use_https=True) if error: return False, f"Slack OAuth failed: {error}" import aiohttp @@ -206,8 +206,8 @@ async def invite(self, args): return False, "CraftOS Notion integration not configured. Set NOTION_SHARED_CLIENT_ID and NOTION_SHARED_CLIENT_SECRET env vars.\nAlternatively, use /notion login with your own integration token." params = {"client_id": NOTION_SHARED_CLIENT_ID, "redirect_uri": REDIRECT_URI, "response_type": "code", "owner": "user", "state": secrets.token_urlsafe(32)} - from agent_core import run_oauth_flow - code, error = run_oauth_flow(f"https://api.notion.com/v1/oauth/authorize?{urlencode(params)}") + from agent_core import run_oauth_flow_async + code, error = await run_oauth_flow_async(f"https://api.notion.com/v1/oauth/authorize?{urlencode(params)}") if error: return False, f"Notion OAuth failed: {error}" import aiohttp @@ -264,8 +264,8 @@ async def login(self, args): return False, "Not configured. Set LINKEDIN_CLIENT_ID and LINKEDIN_CLIENT_SECRET env vars." params = {"response_type": "code", "client_id": LINKEDIN_CLIENT_ID, "redirect_uri": REDIRECT_URI, "scope": "openid profile email w_member_social", "state": secrets.token_urlsafe(32)} - from agent_core import run_oauth_flow - code, error = run_oauth_flow(f"https://www.linkedin.com/oauth/v2/authorization?{urlencode(params)}") + from agent_core import run_oauth_flow_async + code, error = await run_oauth_flow_async(f"https://www.linkedin.com/oauth/v2/authorization?{urlencode(params)}") if error: return False, f"LinkedIn OAuth failed: {error}" import aiohttp @@ -818,8 +818,8 @@ async def login(self, args): "code_challenge": code_challenge, "code_challenge_method": "S256", } - from agent_core import run_oauth_flow - code, error = run_oauth_flow( + from agent_core import run_oauth_flow_async + code, error = await run_oauth_flow_async( f"https://login.microsoftonline.com/common/oauth2/v2.0/authorize?{urlencode(params)}" ) if error: diff --git a/app/ui_layer/adapters/browser_adapter.py b/app/ui_layer/adapters/browser_adapter.py index c659a4d1..4d1536fa 100644 --- a/app/ui_layer/adapters/browser_adapter.py +++ b/app/ui_layer/adapters/browser_adapter.py @@ -648,6 +648,9 @@ def __init__( self._metrics_collector = MetricsCollector(controller.agent) self._metrics_task: Optional[asyncio.Task] = None + # Track active OAuth tasks for cancellation support + self._oauth_tasks: Dict[str, asyncio.Task] = {} + @property def theme_adapter(self) -> ThemeAdapter: return self._theme_adapter @@ -1165,6 +1168,10 @@ async def _handle_ws_message(self, data: Dict[str, Any]) -> None: integration_id = data.get("id", "") await self._handle_integration_connect_interactive(integration_id) + elif msg_type == "integration_connect_cancel": + integration_id = data.get("id", "") + await self._handle_integration_connect_cancel(integration_id) + elif msg_type == "integration_disconnect": integration_id = data.get("id", "") account_id = data.get("account_id") @@ -2989,7 +2996,17 @@ async def _handle_integration_connect_token( }) async def _handle_integration_connect_oauth(self, integration_id: str) -> None: - """Start OAuth flow for an integration.""" + """Start OAuth flow for an integration (non-blocking).""" + # Cancel any existing OAuth task for this integration + if integration_id in self._oauth_tasks: + self._oauth_tasks[integration_id].cancel() + + # Run OAuth in background task so WebSocket message loop stays responsive + task = asyncio.create_task(self._run_oauth_flow(integration_id)) + self._oauth_tasks[integration_id] = task + + async def _run_oauth_flow(self, integration_id: str) -> None: + """Execute OAuth flow and broadcast result (runs as background task).""" try: success, message = await connect_integration_oauth(integration_id) await self._broadcast({ @@ -3003,6 +3020,16 @@ async def _handle_integration_connect_oauth(self, integration_id: str) -> None: # Refresh the list on success (listener is started by connect_integration_oauth) if success: await self._handle_integration_list() + except asyncio.CancelledError: + # OAuth was cancelled by user closing the modal + await self._broadcast({ + "type": "integration_connect_result", + "data": { + "success": False, + "message": "OAuth cancelled", + "id": integration_id, + }, + }) except Exception as e: await self._broadcast({ "type": "integration_connect_result", @@ -3012,9 +3039,21 @@ async def _handle_integration_connect_oauth(self, integration_id: str) -> None: "id": integration_id, }, }) + finally: + self._oauth_tasks.pop(integration_id, None) async def _handle_integration_connect_interactive(self, integration_id: str) -> None: - """Connect an integration using interactive flow (e.g. Telegram QR login).""" + """Connect an integration using interactive flow (non-blocking).""" + # Cancel any existing interactive task for this integration + if integration_id in self._oauth_tasks: + self._oauth_tasks[integration_id].cancel() + + # Run interactive flow in background task so WebSocket message loop stays responsive + task = asyncio.create_task(self._run_interactive_flow(integration_id)) + self._oauth_tasks[integration_id] = task + + async def _run_interactive_flow(self, integration_id: str) -> None: + """Execute interactive flow and broadcast result (runs as background task).""" try: success, message = await connect_integration_interactive(integration_id) await self._broadcast({ @@ -3028,6 +3067,16 @@ async def _handle_integration_connect_interactive(self, integration_id: str) -> # Refresh the list on success (listener is started by connect_integration_interactive) if success: await self._handle_integration_list() + except asyncio.CancelledError: + # Interactive flow was cancelled by user closing the modal + await self._broadcast({ + "type": "integration_connect_result", + "data": { + "success": False, + "message": "Connection cancelled", + "id": integration_id, + }, + }) except Exception as e: await self._broadcast({ "type": "integration_connect_result", @@ -3037,6 +3086,14 @@ async def _handle_integration_connect_interactive(self, integration_id: str) -> "id": integration_id, }, }) + finally: + self._oauth_tasks.pop(integration_id, None) + + async def _handle_integration_connect_cancel(self, integration_id: str) -> None: + """Cancel an in-progress OAuth/interactive flow.""" + if integration_id in self._oauth_tasks: + self._oauth_tasks[integration_id].cancel() + # Result will be broadcast by the cancelled task's CancelledError handler async def _handle_integration_disconnect( self, integration_id: str, account_id: Optional[str] = None diff --git a/app/ui_layer/browser/frontend/src/pages/Settings/SettingsPage.tsx b/app/ui_layer/browser/frontend/src/pages/Settings/SettingsPage.tsx index 51ca70ca..2ece1dd9 100644 --- a/app/ui_layer/browser/frontend/src/pages/Settings/SettingsPage.tsx +++ b/app/ui_layer/browser/frontend/src/pages/Settings/SettingsPage.tsx @@ -3324,7 +3324,18 @@ function IntegrationsSettings() { setCredentials({}) setConnectError('') } else { - setConnectError(d.error || d.message || 'Connection failed') + const errorMsg = d.error || d.message || 'Connection failed' + // Don't show error for user-initiated cancellation (modal already closed) + if (errorMsg.toLowerCase().includes('cancelled')) { + // Silent - user already closed the modal + return + } + // Show toast for timeout so user knows what happened + if (errorMsg.toLowerCase().includes('timed out')) { + showToast('error', 'OAuth timed out. Please try again.') + } + // Show error in modal if still open + setConnectError(errorMsg) } }), onMessage('integration_disconnect_result', (data: unknown) => { @@ -3473,6 +3484,25 @@ function IntegrationsSettings() { setShowConnectModal(false) } + const handleCloseConnectModal = () => { + // Cancel any in-progress OAuth/interactive flow for this integration + if (isConnecting && selectedIntegration) { + send('integration_connect_cancel', { id: selectedIntegration.id }) + } + + // Handle WhatsApp-specific cleanup (has its own polling mechanism) + if (selectedIntegration?.id === 'whatsapp' && whatsappStatus !== 'idle') { + handleCancelWhatsApp() + return // handleCancelWhatsApp already closes the modal + } + + // Reset state + setIsConnecting(false) + setShowConnectModal(false) + setConnectError('') + setCredentials({}) + } + const handleOpenManage = (integration: Integration) => { send('integration_info', { id: integration.id }) } @@ -3644,11 +3674,11 @@ function IntegrationsSettings() { {/* Connect Modal */} {showConnectModal && selectedIntegration && ( -
setShowConnectModal(false)}> +
e.stopPropagation()}>

Connect {selectedIntegration.name}

-
@@ -3669,10 +3699,10 @@ function IntegrationsSettings() { disabled={isConnecting} > {isConnecting ? ( - <> + Connecting... - + ) : ( <>Sign in with {selectedIntegration.name} )} @@ -3707,10 +3737,10 @@ function IntegrationsSettings() { disabled={isConnecting} > {isConnecting ? ( - <> + Connecting... - + ) : ( 'Connect' )} @@ -3745,10 +3775,10 @@ function IntegrationsSettings() { disabled={isConnecting} > {isConnecting ? ( - <> + Connecting... - + ) : ( 'Connect with Token' )} @@ -3792,10 +3822,10 @@ function IntegrationsSettings() { disabled={isConnecting} > {isConnecting ? ( - <> + Connecting... - + ) : ( 'Connect Bot' )} @@ -3810,10 +3840,10 @@ function IntegrationsSettings() { disabled={isConnecting} > {isConnecting ? ( - <> + Waiting for QR scan... - + ) : ( 'Connect User Account (QR Code)' )}