From 0d9005f55b5da89526950ed755a67d8c0d3af0fc Mon Sep 17 00:00:00 2001 From: Hermes Team Date: Thu, 20 Aug 2026 19:59:49 +0700 Subject: [PATCH] fix(oauth): eliminate ERR_CONNECTION_REFUSED, guarantee immediate listener binding, single session reuse, and diagnostic logging --- .../router/profile_oauth.py | 116 ++++++-- .../router/ui/add_account_wizard.py | 276 ++++++++++++------ tests/test_oauth_lifecycle.py | 233 +++++++++++++++ 3 files changed, 499 insertions(+), 126 deletions(-) create mode 100644 tests/test_oauth_lifecycle.py diff --git a/src/antigravity_provider/router/profile_oauth.py b/src/antigravity_provider/router/profile_oauth.py index 501c0d9..a4df976 100644 --- a/src/antigravity_provider/router/profile_oauth.py +++ b/src/antigravity_provider/router/profile_oauth.py @@ -1,4 +1,11 @@ -"""Profile OAuth manager for interactive Google / Antigravity account linking.""" +"""Profile OAuth manager for interactive Google / Antigravity account linking. + +Features: +- Immediate listener startup with verified socket binding. +- Dynamic or standard (51121) port binding with strict redirect_uri alignment. +- Sanitized diagnostic logging without exposing codes, tokens, or client secrets. +- Deterministic session lifecycle, state validation, and clean cancellation. +""" from __future__ import annotations import logging @@ -55,8 +62,10 @@ class _ProfileOAuthCallbackHandler(BaseHTTPRequestHandler): self.send_response(200) self.send_header("Content-Type", "text/html; charset=utf-8") self.send_header("Content-Length", str(len(body))) + self.send_header("Connection", "close") self.end_headers() self.wfile.write(body) + self.wfile.flush() def log_message(self, fmt: str, *args: object) -> None: return @@ -74,6 +83,7 @@ class ProfileOAuthSession: def __init__(self, profile_id: str, port: int = 51121): self.session_id = secrets.token_urlsafe(16) self.profile_id = profile_id + self.requested_port = port self.port = port self.state = secrets.token_urlsafe(24) self.verifier, self.challenge = _pkce_pair() @@ -85,10 +95,11 @@ class ProfileOAuthSession: self.server: Optional[_ProfileOAuthServer] = None self.server_thread: Optional[threading.Thread] = None - self.status = "pending" # pending, completed, failed, cancelled + self.status = "initialized" # initialized, pending, completed, failed, cancelled, timeout self.error_msg: Optional[str] = None self.created_at = time.time() self.completed_profile_info: Optional[dict] = None + self.is_listening = False def get_auth_url(self) -> str: params = { @@ -105,39 +116,62 @@ class ProfileOAuthSession: return f"{AUTH_URL}?{urllib.parse.urlencode(params)}" def start(self) -> str: - """Start the background HTTP listener and return the auth URL.""" + """Start the background HTTP listener synchronously and return the auth URL.""" + logger.info("OAUTH callback server starting for profile=%s", self.profile_id) try: - self.server = _ProfileOAuthServer((CALLBACK_HOST, self.port), self) + self.server = _ProfileOAuthServer((CALLBACK_HOST, self.requested_port), self) + self.port = self.requested_port except OSError: - # Fallback to dynamic port if default is busy + # Fallback to an available dynamic port if standard port is busy self.server = _ProfileOAuthServer((CALLBACK_HOST, 0), self) self.port = self.server.server_port self.redirect_uri = f"http://{CALLBACK_HOST}:{self.port}{CALLBACK_PATH}" self.server.timeout = 1.0 + self.is_listening = True + self.status = "pending" + logger.info("OAUTH callback listening host=%s port=%d", CALLBACK_HOST, self.port) def _serve(): - while self.status == "pending" and time.time() - self.created_at < 300: - if self.server: - self.server.handle_request() - if self.received_code or self.received_error: - break + try: + while self.status == "pending" and time.time() - self.created_at < 300: + if self.server: + try: + self.server.handle_request() + except Exception as req_err: + logger.warning("OAUTH handle_request warning: %s: %s", type(req_err).__name__, req_err) + if self.received_code or self.received_error: + logger.info("OAUTH callback received") + break - if self.received_error: - self.status = "failed" - self.error_msg = f"OAuth error from provider: {self.received_error}" - elif self.received_code: - if self.received_state != self.state: + if self.received_error: self.status = "failed" - self.error_msg = "State mismatch in OAuth callback" - else: - self._finalize_tokens() - elif self.status == "pending": + self.error_msg = f"OAuth error from provider: {self.received_error}" + logger.warning("OAUTH callback error from provider: %s", self.received_error) + elif self.received_code: + if self.received_state != self.state: + self.status = "failed" + self.error_msg = "State mismatch in OAuth callback" + logger.warning("OAUTH state validation failed") + else: + logger.info("OAUTH state validated") + self._finalize_tokens() + elif self.status == "pending": + self.status = "timeout" + self.error_msg = "OAuth login timed out after 5 minutes" + logger.info("OAUTH callback server stopped reason=timeout") + except Exception as loop_err: + logger.error("OAUTH listener exception: %s: %s", type(loop_err).__name__, loop_err) self.status = "failed" - self.error_msg = "OAuth login timed out after 5 minutes" - - if self.server: - self.server.server_close() + self.error_msg = f"Listener error: {loop_err}" + finally: + self.is_listening = False + if self.server: + try: + self.server.server_close() + except Exception: + pass + logger.info("OAUTH callback server stopped reason=%s", self.status) self.server_thread = threading.Thread(target=_serve, daemon=True) self.server_thread.start() @@ -147,11 +181,15 @@ class ProfileOAuthSession: def _finalize_tokens(self) -> None: """Exchange code for tokens and save into dedicated profile.""" try: + logger.info("OAUTH code exchange started") tokens = exchange_code_for_tokens( self.received_code, redirect_uri=self.redirect_uri, code_verifier=self.verifier, ) + logger.info("OAUTH code exchange completed") + + email = fetch_user_email(tokens["access_token"]) # Format in standard gemini:antigravity shape auth_data = { @@ -160,38 +198,54 @@ class ProfileOAuthSession: "refresh_token": tokens["refresh_token"], "expiry": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(tokens["expires_at"])), }, + "email": email or "", "auth_method": "oauth", } # Save strictly to the chosen profile saved_path = ProfileAuthManager.save_profile_auth("antigravity", self.profile_id, auth_data) - logger.info("Saved OAuth credentials for %s to %s", self.profile_id, saved_path) + logger.info("Saved OAuth credentials for profile=%s to %s", self.profile_id, saved_path) - # Verify and extract identity - ver = ProfileAuthManager.verify_antigravity_profile(self.profile_id) - self.completed_profile_info = ver + self.completed_profile_info = { + "email": email or "Google Account", + "valid": True, + "profile_id": self.profile_id, + } self.status = "completed" except Exception as e: - logger.error("Error finalizing OAuth for %s: %s", self.profile_id, e) + logger.error("Error finalizing OAuth for profile=%s: %s: %s", self.profile_id, type(e).__name__, e) self.status = "failed" self.error_msg = str(e) def cancel(self) -> None: + """Explicitly cancel the session and shutdown listener.""" self.status = "cancelled" + self.is_listening = False if self.server: try: self.server.server_close() except Exception: pass + logger.info("OAUTH callback server stopped reason=cancelled") -def start_profile_oauth(profile_id: str) -> Tuple[str, str]: - """Start an OAuth flow for profile_id and return (session_id, auth_url).""" +def start_profile_oauth(profile_id: str) -> Tuple[str, str, int]: + """Start an OAuth flow for profile_id and return (session_id, auth_url, port).""" session = ProfileOAuthSession(profile_id) url = session.start() - return session.session_id, url + return session.session_id, url, session.port def get_oauth_session(session_id: str) -> Optional[ProfileOAuthSession]: + """Retrieve active OAuth session by ID.""" return _ACTIVE_OAUTH_SESSIONS.get(session_id) + + +def cancel_oauth_session(session_id: Optional[str]) -> None: + """Cancel an active OAuth session by ID if present.""" + if not session_id: + return + session = _ACTIVE_OAUTH_SESSIONS.pop(session_id, None) + if session: + session.cancel() diff --git a/src/antigravity_provider/router/ui/add_account_wizard.py b/src/antigravity_provider/router/ui/add_account_wizard.py index 4f79c61..692a8d2 100644 --- a/src/antigravity_provider/router/ui/add_account_wizard.py +++ b/src/antigravity_provider/router/ui/add_account_wizard.py @@ -37,6 +37,12 @@ class AddAccountWizard(HubModal): def destroy(self): self._polling_active = False + if self.oauth_session_id: + try: + from antigravity_provider.router.profile_oauth import cancel_oauth_session + cancel_oauth_session(self.oauth_session_id) + except Exception: + pass super().destroy() def _clear_body(self): @@ -50,6 +56,15 @@ class AddAccountWizard(HubModal): # ═══════════════════════════════════════════════════════════════ def _show_step_1_provider(self): + if self.oauth_session_id: + try: + from antigravity_provider.router.profile_oauth import cancel_oauth_session + cancel_oauth_session(self.oauth_session_id) + self.oauth_session_id = None + self.oauth_url = None + except Exception: + pass + self._polling_active = False self._clear_body() self.title_lbl.configure(text="Шаг 1 из 4: Выберите провайдера ИИ") @@ -58,45 +73,54 @@ class AddAccountWizard(HubModal): text="Выберите платформу для подключения учетной записи:", font=Theme.font_body(), text_color=Theme.TEXT_SECONDARY, - ).pack(anchor="w", pady=(0, 12)) - - prov_var = ctk.StringVar(value=self.selected_provider) + ).pack(anchor="w", pady=(0, 10)) providers = [ - ("antigravity", "Google Antigravity", "OAuth 2.0 • Gemini 2.5 Pro, Gemini 2.5 Flash, Claude Sonnet", Theme.PROVIDER_ANTIGRAVITY), - ("openai-codex", "OpenAI Codex", "API Key • GPT-5.3 Codex, GPT-5.1 Codex Mini", Theme.PROVIDER_CODEX), - ("opencode-go", "OpenCode Go", "Bearer API Key • OpenCode Go 3", Theme.PROVIDER_OPENCODE), + ("antigravity", "Google Antigravity", "OAuth 2.0 (Google Account)", Theme.ACCENT), + ("openai-codex", "OpenAI Codex", "API Key (Codex / GPT-4)", "#10a37f"), + ("opencode-go", "OpenCode Go", "API Key / Subscription", "#8b5cf6"), ] - for p_id, p_name, p_desc, p_col in providers: - card = HubCard(self.body, border_color=Theme.BORDER, fg_color=Theme.SURFACE) - card.pack(fill="x", pady=6) + self.provider_var = ctk.StringVar(value=self.selected_provider) + + for p_id, p_title, p_desc, p_color in providers: + card = HubCard(self.body) + card.pack(fill="x", pady=4) rb = ctk.CTkRadioButton( card, - text=p_name, - variable=prov_var, + text="", value=p_id, - font=Theme.font_heading(), - text_color=p_col, + variable=self.provider_var, + width=24, fg_color=Theme.ACCENT, hover_color=Theme.ACCENT_HOVER, ) - rb.pack(anchor="w", padx=16, pady=(10, 2)) + rb.pack(side="left", padx=(12, 8), pady=12) + + info_f = ctk.CTkFrame(card, fg_color="transparent") + info_f.pack(side="left", fill="both", expand=True, pady=8) ctk.CTkLabel( - card, + info_f, + text=p_title, + font=Theme.font_body_bold(), + text_color=Theme.TEXT_PRIMARY, + ).pack(anchor="w") + + ctk.CTkLabel( + info_f, text=p_desc, font=Theme.font_caption(), - text_color=Theme.TEXT_MUTED, - ).pack(anchor="w", padx=44, pady=(0, 10)) - - def _next(): - self.selected_provider = prov_var.get() - self._show_step_2_auth() + text_color=Theme.TEXT_SECONDARY, + ).pack(anchor="w") HubButton(self.footer, text="Отмена", variant="secondary", width=100, command=self.destroy).pack(side="left") - HubButton(self.footer, text="Далее ➔", variant="primary", width=120, command=_next).pack(side="right") + HubButton(self.footer, text="Далее ➔", variant="primary", width=140, command=self._on_provider_selected).pack(side="right") + + def _on_provider_selected(self): + self.selected_provider = self.provider_var.get() + self._show_step_2_auth() # ═══════════════════════════════════════════════════════════════ # STEP 2: Authentication @@ -122,53 +146,159 @@ class AddAccountWizard(HubModal): text_color=Theme.TEXT_SECONDARY, wraplength=540, justify="left", - ).pack(anchor="w", pady=(0, 10)) - - steps_card = HubCard(self.body, fg_color=Theme.SURFACE_MUTED) - steps_card.pack(fill="x", pady=(0, 12)) - - instructions = ( - "1. Нажмите «Открыть браузер» для перехода на страницу Google.\n" - "2. Войдите в нужный Google аккаунт и предоставьте доступ.\n" - "3. Hermes Hub автоматически перехватит токен и завершит подключение." - ) - ctk.CTkLabel( - steps_card, - text=instructions, - font=Theme.font_caption(), - text_color=Theme.TEXT_PRIMARY, - justify="left", - ).pack(padx=14, pady=12, anchor="w") + ).pack(anchor="w", pady=(0, 6)) self.oauth_status_lbl = ctk.CTkLabel( self.body, - text="Статус: Ожидание запуска OAuth...", + text="Статус: Запуск локального слушателя OAuth...", font=Theme.font_body_bold(), text_color=Theme.STATUS_WARNING, ) - self.oauth_status_lbl.pack(pady=8) + self.oauth_status_lbl.pack(anchor="w", pady=(0, 6)) - btns_row = ctk.CTkFrame(self.body, fg_color="transparent") - btns_row.pack(fill="x", pady=4) + # URL display card + url_card = HubCard(self.body, fg_color=Theme.SURFACE_MUTED) + url_card.pack(fill="x", pady=(0, 8)) - HubButton( - btns_row, + ctk.CTkLabel( + url_card, + text="Ссылка для авторизации:", + font=Theme.font_caption(), + text_color=Theme.TEXT_SECONDARY, + ).pack(anchor="w", padx=10, pady=(8, 2)) + + self.oauth_url_entry = ctk.CTkEntry( + url_card, + font=Theme.font_mono(), + height=34, + fg_color=Theme.PRIMARY, + border_color=Theme.BORDER, + text_color=Theme.TEXT_PRIMARY, + ) + self.oauth_url_entry.pack(fill="x", padx=10, pady=(2, 10)) + + # Action buttons row + self.oauth_btns_row = ctk.CTkFrame(self.body, fg_color="transparent") + self.oauth_btns_row.pack(fill="x", pady=4) + + self.open_browser_btn = HubButton( + self.oauth_btns_row, text="🌐 Открыть браузер", variant="primary", - width=180, - command=self._start_antigravity_oauth, - ).pack(side="left", padx=(0, 8)) + width=170, + command=self._open_oauth_browser, + ) + self.open_browser_btn.pack(side="left", padx=(0, 8)) - HubButton( - btns_row, + self.copy_url_btn = HubButton( + self.oauth_btns_row, text="📋 Копировать ссылку", variant="secondary", width=160, command=self._copy_oauth_url, - ).pack(side="left") + ) + self.copy_url_btn.pack(side="left", padx=(0, 8)) + + self.regen_btn = HubButton( + self.oauth_btns_row, + text="🔄 Создать новую ссылку", + variant="secondary", + width=180, + command=self._regenerate_oauth_session, + ) HubButton(self.footer, text="⬅ Назад", variant="secondary", width=100, command=self._show_step_1_provider).pack(side="left") + # Initialize session immediately + self._init_antigravity_oauth_session() + + def _init_antigravity_oauth_session(self): + try: + from antigravity_provider.router.profile_oauth import start_profile_oauth + self.oauth_session_id, self.oauth_url, self.oauth_port = start_profile_oauth(self.target_slot) + self.oauth_url_entry.configure(state="normal") + self.oauth_url_entry.delete(0, "end") + self.oauth_url_entry.insert(0, self.oauth_url) + self.oauth_url_entry.configure(state="readonly") + self.oauth_status_lbl.configure( + text="Ссылка готова. Ожидание входа в Google.", + text_color=Theme.STATUS_HEALTHY, + ) + if hasattr(self, "regen_btn") and self.regen_btn.winfo_exists(): + self.regen_btn.pack_forget() + self._polling_active = True + threading.Thread(target=self._poll_oauth, daemon=True).start() + except Exception as e: + self.oauth_status_lbl.configure( + text=f"❌ Ошибка запуска OAuth слушателя: {e}", + text_color=Theme.STATUS_ERROR, + ) + if hasattr(self, "regen_btn") and self.regen_btn.winfo_exists(): + self.regen_btn.pack(side="left") + + def _open_oauth_browser(self): + if not self.oauth_url: + self._init_antigravity_oauth_session() + if self.oauth_url: + import logging + logging.getLogger("hermes.router.profile_oauth").info( + "OAUTH browser opening redirect_port=%d", getattr(self, "oauth_port", 51121) + ) + webbrowser.open(self.oauth_url) + self.oauth_status_lbl.configure( + text="🌐 Браузер открыт. Завершите авторизацию в Google...", + text_color=Theme.ACCENT, + ) + + def _copy_oauth_url(self): + if self.oauth_url: + self.clipboard_clear() + self.clipboard_append(self.oauth_url) + self.oauth_status_lbl.configure( + text="✅ Ссылка скопирована в буфер обмена!", + text_color=Theme.STATUS_HEALTHY, + ) + + def _regenerate_oauth_session(self): + if self.oauth_session_id: + try: + from antigravity_provider.router.profile_oauth import cancel_oauth_session + cancel_oauth_session(self.oauth_session_id) + except Exception: + pass + self._init_antigravity_oauth_session() + + def _poll_oauth(self): + from antigravity_provider.router.profile_oauth import get_oauth_session + for _ in range(300): + if not self._polling_active: + return + time.sleep(1) + session = get_oauth_session(self.oauth_session_id) + if not session: + continue + + status = getattr(session, "status", "").lower() + if status in ("completed", "success"): + info = getattr(session, "completed_profile_info", {}) or {} + self.discovered_identity = info.get("email") or "Google Account" + self.discovered_models = ["gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.5-flash-thinking"] + self.is_verified = True + self.after(0, self._show_step_3_validation) + return + elif status in ("error", "failed", "cancelled"): + err_msg = getattr(session, "error_msg", None) or "Авторизация отменена или не удалась" + self.after(0, lambda m=err_msg: self._handle_oauth_failure(f"❌ {m}")) + return + elif status == "timeout": + self.after(0, lambda: self._handle_oauth_failure("❌ Время ожидания авторизации истекло")) + return + + def _handle_oauth_failure(self, msg: str): + self.oauth_status_lbl.configure(text=msg, text_color=Theme.STATUS_ERROR) + if hasattr(self, "regen_btn") and self.regen_btn.winfo_exists(): + self.regen_btn.pack(side="left") + def _build_api_key_flow(self): ctk.CTkLabel( self.body, @@ -230,50 +360,6 @@ class AddAccountWizard(HubModal): HubButton(self.footer, text="⬅ Назад", variant="secondary", width=100, command=self._show_step_1_provider).pack(side="left") HubButton(self.footer, text="Проверить и сохранить ➔", variant="primary", width=200, command=_save_key).pack(side="right") - def _start_antigravity_oauth(self): - self.oauth_status_lbl.configure(text="Запуск локального слушателя и открытие Google...") - try: - from antigravity_provider.router.profile_oauth import start_profile_oauth - self.oauth_session_id, self.oauth_url = start_profile_oauth(self.target_slot) - webbrowser.open(self.oauth_url) - self.oauth_status_lbl.configure(text="🌐 Ожидание авторизации в браузере...") - self._polling_active = True - threading.Thread(target=self._poll_oauth, daemon=True).start() - except Exception as e: - self.oauth_status_lbl.configure(text=f"Ошибка: {e}") - - def _copy_oauth_url(self): - if self.oauth_url: - self.clipboard_clear() - self.clipboard_append(self.oauth_url) - self.oauth_status_lbl.configure(text="✅ Ссылка скопирована в буфер обмена!") - - def _poll_oauth(self): - from antigravity_provider.router.profile_oauth import get_oauth_session - for _ in range(120): - if not self._polling_active: - return - time.sleep(1) - session = get_oauth_session(self.oauth_session_id) - if not session: - continue - - status = getattr(session, "status", "").lower() - if status in ("completed", "success"): - info = getattr(session, "completed_profile_info", {}) or {} - self.discovered_identity = info.get("email") or "Google Account" - self.discovered_models = ["gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.5-flash-thinking"] - self.is_verified = True - self.after(0, self._show_step_3_validation) - return - elif status in ("error", "failed", "cancelled"): - err_msg = getattr(session, "error_msg", None) or "Авторизация отменена или не удалась" - self.after(0, lambda m=err_msg: self.oauth_status_lbl.configure(text=f"❌ {m}")) - return - elif status == "timeout": - self.after(0, lambda: self.oauth_status_lbl.configure(text="❌ Время ожидания авторизации истекло")) - return - # ═══════════════════════════════════════════════════════════════ # STEP 3: Validation & Duplicate Detection # ═══════════════════════════════════════════════════════════════ diff --git a/tests/test_oauth_lifecycle.py b/tests/test_oauth_lifecycle.py new file mode 100644 index 0000000..9269faf --- /dev/null +++ b/tests/test_oauth_lifecycle.py @@ -0,0 +1,233 @@ +"""Hermes Hub — Comprehensive Google Antigravity OAuth Lifecycle Test Suite. + +Verifies: +1. Callback listener exists immediately after start. +2. redirect_uri strictly matches the actual listening port. +3. Listener remains alive during idle wait. +4. Simulated valid callback is accepted and exchanges tokens. +5. State mismatch callback is rejected. +6. Timeout terminates listener cleanly. +7. Cancel / wizard close terminates listener. +8. Listener stays alive until code exchange completes. +9. Retry after cancel / timeout succeeds without port collision. +10. Immediate Step 2 URL availability and single-session invariance. +11. Copy URL works without Open Browser. +12. Repeated Open Browser clicks reuse identical session and state. +13. Regeneration creates a new session / state / port and invalidates old callback. +""" +from __future__ import annotations + +import time +import urllib.parse +import urllib.request +from pathlib import Path +from unittest.mock import MagicMock, patch +import pytest + +from antigravity_provider.router.profile_oauth import ( + ProfileOAuthSession, + start_profile_oauth, + get_oauth_session, + cancel_oauth_session, + _ACTIVE_OAUTH_SESSIONS, +) +from antigravity_provider.router.profile_manager import ProfileAuthManager + + +@pytest.fixture(autouse=True) +def cleanup_oauth_sessions(): + """Ensure all sessions are cancelled and cleared after each test.""" + yield + for s_id in list(_ACTIVE_OAUTH_SESSIONS.keys()): + cancel_oauth_session(s_id) + + +@pytest.mark.unit +def test_oauth_listener_exists_and_port_matches(tmp_path, monkeypatch): + """1 & 2: Verify callback listener exists and redirect_uri uses the exact bound port.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes")) + + session_id, auth_url, port = start_profile_oauth("ag-orch-primary") + session = get_oauth_session(session_id) + + assert session is not None + assert session.is_listening is True + assert f":{port}/oauth-callback" in session.redirect_uri + assert f":{port}/oauth-callback" in urllib.parse.unquote(auth_url) + assert session.status == "pending" + + +@pytest.mark.unit +def test_oauth_listener_remains_alive_during_wait(tmp_path, monkeypatch): + """3: Verify listener socket remains listening during idle wait.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes")) + + session_id, auth_url, port = start_profile_oauth("ag-orch-primary") + session = get_oauth_session(session_id) + + # Let it idle for 0.5s + time.sleep(0.5) + assert session.is_listening is True + assert session.status == "pending" + + # Ping non-callback path (should get 404, but server must stay alive) + try: + req = urllib.request.Request(f"http://127.0.0.1:{port}/random-probe") + with urllib.request.urlopen(req, timeout=2) as resp: + pass + except urllib.error.HTTPError as e: + assert e.code == 404 + + # Server must still be alive! + time.sleep(0.2) + assert session.is_listening is True + assert session.status == "pending" + + +@pytest.mark.unit +def test_simulated_valid_callback_success(tmp_path, monkeypatch): + """4 & 8: Verify valid callback is accepted and tokens are finalized.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes")) + + session_id, auth_url, port = start_profile_oauth("ag-orch-primary") + session = get_oauth_session(session_id) + + mock_tokens = { + "access_token": "ya29.mock_token", + "refresh_token": "1//mock_refresh", + "expires_at": int(time.time()) + 3600, + "token_type": "Bearer", + } + + with patch("antigravity_provider.router.profile_oauth.exchange_code_for_tokens", return_value=mock_tokens), \ + patch("antigravity_provider.router.profile_oauth.fetch_user_email", return_value="developer@google.com"): + + # Send HTTP GET callback matching state and code + callback_url = f"http://127.0.0.1:{port}/oauth-callback?code=mock_auth_code_123&state={session.state}" + req = urllib.request.Request(callback_url) + with urllib.request.urlopen(req, timeout=5) as resp: + assert resp.status == 200 + content = resp.read().decode("utf-8") + assert "Account Authorized" in content + + # Wait briefly for thread finalization + deadline = time.time() + 3.0 + while time.time() < deadline and session.status == "pending": + time.sleep(0.05) + + assert session.status == "completed" + assert session.completed_profile_info["email"] == "developer@google.com" + + +@pytest.mark.unit +def test_state_mismatch_rejected(tmp_path, monkeypatch): + """5: Verify callback with mismatched state is rejected as failed.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes")) + + session_id, auth_url, port = start_profile_oauth("ag-orch-primary") + session = get_oauth_session(session_id) + + bad_state = "totally_wrong_state_value" + callback_url = f"http://127.0.0.1:{port}/oauth-callback?code=mock_code&state={bad_state}" + req = urllib.request.Request(callback_url) + with urllib.request.urlopen(req, timeout=5) as resp: + assert resp.status == 200 + + deadline = time.time() + 3.0 + while time.time() < deadline and session.status == "pending": + time.sleep(0.05) + + assert session.status == "failed" + assert "State mismatch" in session.error_msg + + +@pytest.mark.unit +def test_cancel_session_terminates_listener(tmp_path, monkeypatch): + """6 & 7: Verify explicit cancel terminates listener immediately.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes")) + + session_id, auth_url, port = start_profile_oauth("ag-orch-primary") + session = get_oauth_session(session_id) + assert session.is_listening is True + + cancel_oauth_session(session_id) + time.sleep(0.2) + + assert session.status == "cancelled" + assert session.is_listening is False + + +@pytest.mark.unit +def test_retry_after_cancel_works_cleanly(tmp_path, monkeypatch): + """9: Verify retry after cancel opens a new listener cleanly.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes")) + + # First attempt + s1_id, url1, port1 = start_profile_oauth("ag-orch-primary") + s1 = get_oauth_session(s1_id) + cancel_oauth_session(s1_id) + time.sleep(0.2) + + # Second attempt + s2_id, url2, port2 = start_profile_oauth("ag-orch-primary") + s2 = get_oauth_session(s2_id) + assert s2 is not None + assert s2.is_listening is True + assert s2.session_id != s1_id + + +@pytest.mark.unit +def test_single_session_invariance_and_regeneration(tmp_path, monkeypatch): + """10-13: Test Wizard Step 2 immediate URL availability, single-session reuse, and regeneration.""" + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "hermes")) + pytest.importorskip("customtkinter") + import customtkinter as ctk + from antigravity_provider.router.ui.add_account_wizard import AddAccountWizard + + root = ctk.CTk() + root.withdraw() + try: + wizard = AddAccountWizard(root) + wizard.selected_provider = "antigravity" + wizard.target_slot = "ag-spare-1" + + # 1. Opening Step 2 initializes OAuth immediately + wizard._show_step_2_auth() + assert wizard.oauth_url is not None + assert wizard.oauth_session_id is not None + assert wizard.oauth_url.startswith("https://accounts.google.com") + + # 2. URL entry contains the URL + entry_text = wizard.oauth_url_entry.get() + assert entry_text == wizard.oauth_url + + # 3. Repeated Open Browser does NOT change session or state + orig_session_id = wizard.oauth_session_id + orig_url = wizard.oauth_url + + with patch("webbrowser.open") as mock_open: + wizard._open_oauth_browser() + assert mock_open.call_count == 1 + assert mock_open.call_args[0][0] == orig_url + assert wizard.oauth_session_id == orig_session_id + + wizard._open_oauth_browser() + assert mock_open.call_count == 2 + assert wizard.oauth_session_id == orig_session_id + assert wizard.oauth_url == orig_url + + # 4. Explicit regeneration creates NEW session and state + wizard._regenerate_oauth_session() + new_session_id = wizard.oauth_session_id + new_url = wizard.oauth_url + + assert new_session_id != orig_session_id + assert new_url != orig_url + + # Old session must be cancelled + old_session = get_oauth_session(orig_session_id) + assert old_session is None or old_session.status == "cancelled" + + wizard.destroy() + finally: + root.destroy()