feat(webui): support remote Codex OAuth login (#5174)

nanobot/providers/openai_codex_oauth.py

@@ -0,0 +1,272 @@
+"""WebUI adapter around oauth-cli-kit's interactive Codex login."""
+
+# oauth-cli-kit does not publish type stubs.
+# pyright: reportMissingTypeStubs=false
+
+from __future__ import annotations
+
+import hmac
+import queue
+import re
+import threading
+import time
+from concurrent.futures import Future
+from contextlib import suppress
+from urllib.parse import parse_qs, urlsplit
+
+from oauth_cli_kit import login_oauth_interactive
+from oauth_cli_kit.models import OAuthToken
+from oauth_cli_kit.providers import OPENAI_CODEX_PROVIDER
+
+_AUTHORIZATION_URL_TIMEOUT_S = 5.0
+_CALLBACK = urlsplit(OPENAI_CODEX_PROVIDER.redirect_uri)
+_CALLBACK_HOSTS = {"localhost", "127.0.0.1", "::1"}
+_TOKEN_EXCHANGE_STATUS = re.compile(r"Token exchange failed:\s*(\d{3})\b")
+
+
+class OpenAICodexOAuthError(RuntimeError):
+    """An actionable Codex OAuth failure that contains no credential material."""
+
+
+class OpenAICodexOAuthInputError(OpenAICodexOAuthError):
+    """A recoverable error in a callback URL pasted by the user."""
+
+
+class OpenAICodexOAuthLoginFlow:
+    """Expose oauth-cli-kit's blocking prompt as a two-stage WebUI flow."""
+
+    def __init__(
+        self,
+        *,
+        proxy: str | None,
+        timeout_s: float,
+        open_browser: bool,
+    ) -> None:
+        self.authorization_url = ""
+        self._expected_state = ""
+        self._proxy = proxy
+        self._open_browser = open_browser
+        self._expires_at = time.monotonic() + timeout_s
+        self._callback_input: queue.Queue[str] = queue.Queue(maxsize=1)
+        self._result: Future[OAuthToken] = Future()
+        self._ready = threading.Event()
+        self._submission_lock = threading.Lock()
+        self._submitted = False
+        self._thread = threading.Thread(
+            target=self._run,
+            name="nanobot-openai-codex-oauth",
+            daemon=True,
+        )
+
+    @property
+    def expired(self) -> bool:
+        return time.monotonic() >= self._expires_at
+
+    @property
+    def remaining_seconds(self) -> int:
+        return max(0, int(self._expires_at - time.monotonic()))
+
+    def start(self) -> OpenAICodexOAuthLoginFlow:
+        self._thread.start()
+        wait_s = min(
+            _AUTHORIZATION_URL_TIMEOUT_S,
+            max(0.0, self._expires_at - time.monotonic()),
+        )
+        if not self._ready.wait(wait_s):
+            error = OpenAICodexOAuthError(
+                "OpenAI Codex sign-in could not create an authorization URL."
+            )
+            self._fail(error)
+            raise error
+        if self._result.done():
+            self._result.result()
+        if self.authorization_url:
+            return self
+        error = OpenAICodexOAuthError(
+            "OpenAI Codex sign-in returned no authorization URL."
+        )
+        self._fail(error)
+        raise error
+
+    def complete(self, callback_url: str | None = None) -> OAuthToken | None:
+        """Submit a full callback URL, or return ``None`` while waiting for one."""
+        if self._result.done():
+            return self._result.result()
+        if self.expired:
+            error = OpenAICodexOAuthError(
+                "OpenAI Codex sign-in expired. Start a new sign-in flow."
+            )
+            self._fail(error)
+            raise error
+        if callback_url is None:
+            return None
+
+        callback_state, authorization_failed = _validate_callback_url(callback_url)
+        if not hmac.compare_digest(callback_state, self._expected_state):
+            raise OpenAICodexOAuthInputError(
+                "The callback URL does not belong to this sign-in flow. Copy the latest URL."
+            )
+        if authorization_failed:
+            error = OpenAICodexOAuthError(
+                "OpenAI Codex sign-in was not completed by the authorization server."
+            )
+            self._fail(error)
+            raise error
+
+        with self._submission_lock:
+            if self._submitted:
+                return None
+            self._submitted = True
+            try:
+                self._callback_input.put_nowait(callback_url.strip())
+            except queue.Full:
+                return None
+        return self._result.result() if self._result.done() else None
+
+    def cancel(self) -> None:
+        """Unblock an abandoned interactive login."""
+        self._fail(OpenAICodexOAuthError("OpenAI Codex sign-in was cancelled."))
+        if threading.current_thread() is not self._thread:
+            self._thread.join(timeout=0.5)
+
+    def _run(self) -> None:
+        try:
+            token = login_oauth_interactive(
+                print_fn=self._capture_output,
+                prompt_fn=self._prompt_for_callback,
+                provider=OPENAI_CODEX_PROVIDER,
+                proxy=self._proxy,
+                open_browser=self._open_browser,
+            )
+        except Exception as exc:
+            with suppress(Exception):
+                self._result.set_exception(_safe_login_error(exc))
+        else:
+            with suppress(Exception):
+                self._result.set_result(token)
+        finally:
+            self._ready.set()
+
+    def _capture_output(self, message: str) -> None:
+        raw = str(message)
+        start = raw.find(OPENAI_CODEX_PROVIDER.authorize_url)
+        if start < 0:
+            return
+        candidate = raw[start:].split(maxsplit=1)[0]
+        state = _first(parse_qs(urlsplit(candidate).query), "state")
+        if not state:
+            return
+        self.authorization_url = candidate
+        self._expected_state = state
+        self._ready.set()
+
+    def _prompt_for_callback(self, _prompt: str) -> str:
+        remaining = max(0.0, self._expires_at - time.monotonic())
+        try:
+            value = self._callback_input.get(timeout=remaining)
+        except queue.Empty as exc:
+            raise OpenAICodexOAuthError(
+                "OpenAI Codex sign-in expired. Start a new sign-in flow."
+            ) from exc
+        if not value:
+            error = self._result.exception() if self._result.done() else None
+            if error is not None:
+                raise error
+            raise OpenAICodexOAuthError("OpenAI Codex sign-in was cancelled.")
+        return value
+
+    def _fail(self, error: OpenAICodexOAuthError) -> None:
+        try:
+            self._result.set_exception(error)
+        except Exception:
+            pass
+        else:
+            with suppress(queue.Full):
+                self._callback_input.put_nowait("")
+        self._ready.set()
+
+
+def start_openai_codex_oauth_login(
+    *,
+    proxy: str | None = None,
+    timeout_s: float = 600,
+    open_browser: bool = True,
+) -> OpenAICodexOAuthLoginFlow:
+    """Start a non-blocking wrapper around oauth-cli-kit's Codex login."""
+    return OpenAICodexOAuthLoginFlow(
+        proxy=proxy,
+        timeout_s=timeout_s,
+        open_browser=open_browser,
+    ).start()
+
+
+def complete_openai_codex_oauth_login(
+    flow: OpenAICodexOAuthLoginFlow,
+    callback_url: str | None = None,
+) -> OAuthToken | None:
+    """Complete a pending Codex login from a full callback URL."""
+    return flow.complete(callback_url)
+
+
+def _validate_callback_url(raw: str) -> tuple[str, bool]:
+    value = raw.strip()
+    if not value:
+        raise OpenAICodexOAuthInputError("Paste the full callback URL from your browser.")
+    try:
+        parsed = urlsplit(value)
+        port = parsed.port
+    except ValueError as exc:
+        raise OpenAICodexOAuthInputError(
+            "The callback URL is invalid. Copy the full URL from your browser's address bar."
+        ) from exc
+    if (
+        parsed.scheme != _CALLBACK.scheme
+        or parsed.hostname not in _CALLBACK_HOSTS
+        or port != _CALLBACK.port
+        or parsed.path != _CALLBACK.path
+        or parsed.username is not None
+        or parsed.password is not None
+    ):
+        raise OpenAICodexOAuthInputError(
+            f"Paste the full callback URL from your browser ({OPENAI_CODEX_PROVIDER.redirect_uri}?...)."
+        )
+    params = parse_qs(parsed.query)
+    code = _first(params, "code")
+    state = _first(params, "state")
+    error = _first(params, "error")
+    if not state:
+        raise OpenAICodexOAuthInputError(
+            "The callback URL is missing OAuth state. Copy the entire browser address."
+        )
+    if not code and not error:
+        raise OpenAICodexOAuthInputError(
+            "The callback URL has no authorization result. Finish signing in, then copy it again."
+        )
+    return state, error is not None
+
+
+def _safe_login_error(exc: Exception) -> OpenAICodexOAuthError:
+    if isinstance(exc, OpenAICodexOAuthError):
+        return exc
+    message = str(exc).strip()
+    if message == "State validation failed.":
+        return OpenAICodexOAuthError(
+            "OpenAI Codex sign-in failed because the OAuth state did not match."
+        )
+    if message == "Authorization code not found.":
+        return OpenAICodexOAuthError(
+            "OpenAI Codex sign-in returned no authorization code."
+        )
+    status = _TOKEN_EXCHANGE_STATUS.search(message)
+    if status:
+        return OpenAICodexOAuthError(
+            f"OpenAI Codex OAuth token exchange failed with HTTP {status.group(1)}."
+        )
+    return OpenAICodexOAuthError(
+        f"OpenAI Codex sign-in failed ({type(exc).__name__})."
+    )
+
+
+def _first(params: dict[str, list[str]], key: str) -> str | None:
+    values = params.get(key)
+    return values[0] if values else None

nanobot/webui/settings_api.py

@@ -131,10 +131,10 @@ _IMAGE_GENERATION_ASPECT_RATIOS = {
 }
 _CONTEXT_WINDOW_TOKEN_OPTIONS = {65_536, 200_000, 262_144, 500_000, 1_048_576}
 _OAUTH_PROXY_PROVIDERS = {"openai_codex", "xai_grok"}
-_XAI_WEBUI_OAUTH_TIMEOUT_S = 600
-_XAI_WEBUI_OAUTH_MAX_FLOWS = 8
-_xai_webui_oauth_flows: dict[str, Any] = {}
-_xai_webui_oauth_flows_lock = threading.Lock()
+_WEBUI_OAUTH_TIMEOUT_S = 600
+_WEBUI_OAUTH_MAX_FLOWS = 8
+_webui_oauth_flows: dict[str, tuple[str, Any]] = {}
+_webui_oauth_flows_lock = threading.Lock()
 _MODEL_CONFIGURATION_SLUG_RE = re.compile(r"[^a-z0-9_-]+")
 _ENV_REF_RE = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)\}")
 
@@ -1810,7 +1810,7 @@ def login_oauth_provider(query: QueryParams) -> dict[str, Any]:
 
     if spec.name == "openai_codex":
         try:
-            from oauth_cli_kit import get_token, login_oauth_interactive
+            from nanobot.providers.openai_codex_oauth import start_openai_codex_oauth_login
         except ImportError:
             raise WebUISettingsError(
                 "oauth_cli_kit not installed. Run: pip install oauth-cli-kit", status=500
@@ -1820,19 +1820,30 @@ def login_oauth_provider(query: QueryParams) -> dict[str, Any]:
             proxy = resolve_config_env_vars(load_config()).providers.openai_codex.proxy or None
         except ValueError as e:
             raise WebUISettingsError(str(e), status=400) from e
-        token = None
-        with suppress(Exception):
-            token = get_token(proxy=proxy)
-        if not (token and token.access):
-            messages: list[str] = []
-            token = login_oauth_interactive(
-                print_fn=lambda message: messages.append(str(message)),
-                prompt_fn=lambda _prompt: "",
+        remote_browser_value = _query_first(query, "remote_browser")
+        remote_browser = (
+            _parse_bool(remote_browser_value, "remote_browser")
+            if remote_browser_value is not None
+            else False
+        )
+        try:
+            flow = start_openai_codex_oauth_login(
                 proxy=proxy,
+                timeout_s=_WEBUI_OAUTH_TIMEOUT_S,
+                open_browser=not remote_browser,
             )
-        if not (token and token.access):
-            raise WebUISettingsError("OAuth login failed", status=401)
-        return settings_payload()
+        except Exception as e:
+            raise WebUISettingsError(f"OpenAI Codex OAuth login failed: {e}", status=502) from e
+        flow_id = secrets.token_urlsafe(24)
+        _register_webui_oauth_flow(spec.name, flow_id, flow)
+        return {
+            "status": "authorization_required",
+            "provider": spec.name,
+            "flow_id": flow_id,
+            "authorization_url": flow.authorization_url,
+            "expires_in": flow.remaining_seconds,
+            "completion_input": "callback_url",
+        }
 
     if spec.name == "github_copilot":
         try:
@@ -1862,18 +1873,19 @@ def login_oauth_provider(query: QueryParams) -> dict[str, Any]:
         try:
             flow = start_xai_oauth_login(
                 proxy=proxy,
-                timeout_s=_XAI_WEBUI_OAUTH_TIMEOUT_S,
+                timeout_s=_WEBUI_OAUTH_TIMEOUT_S,
             )
         except Exception as e:
             raise WebUISettingsError(f"xAI OAuth login failed: {e}", status=502) from e
         flow_id = secrets.token_urlsafe(24)
-        _register_xai_webui_oauth_flow(flow_id, flow)
+        _register_webui_oauth_flow(spec.name, flow_id, flow)
         return {
             "status": "authorization_required",
             "provider": spec.name,
             "flow_id": flow_id,
             "authorization_url": flow.authorization_url,
             "expires_in": flow.remaining_seconds,
+            "completion_input": "authorization_code",
         }
 
     raise WebUISettingsError("OAuth login is not supported for this provider")
@@ -1881,34 +1893,47 @@ def login_oauth_provider(query: QueryParams) -> dict[str, Any]:
 
 def complete_oauth_provider(
     query: QueryParams,
-    authorization_code: str | None = None,
+    authorization_response: str | None = None,
 ) -> dict[str, Any]:
     provider_name = (_query_first(query, "provider") or "").strip()
     flow_id = (_query_first(query, "flow_id") or "").strip()
     spec = find_by_name(provider_name)
-    if spec is None or spec.name != "xai_grok":
+    if spec is None or spec.name not in {"openai_codex", "xai_grok"}:
         raise WebUISettingsError("OAuth completion is not supported for this provider")
     if not flow_id:
         raise WebUISettingsError("flow_id is required")
 
-    flow = _get_xai_webui_oauth_flow(flow_id)
+    flow = _get_webui_oauth_flow(spec.name, flow_id)
     if flow is None:
-        raise WebUISettingsError("xAI sign-in expired. Start again.", status=410)
-
-    from nanobot.providers.xai_oauth import complete_xai_oauth_login
+        raise WebUISettingsError(f"{spec.label} sign-in expired. Start again.", status=410)
 
     try:
-        token = complete_xai_oauth_login(flow, authorization_code)
+        if spec.name == "openai_codex":
+            from nanobot.providers.openai_codex_oauth import (
+                OpenAICodexOAuthInputError,
+                complete_openai_codex_oauth_login,
+            )
+
+            try:
+                token = complete_openai_codex_oauth_login(flow, authorization_response)
+            except OpenAICodexOAuthInputError as e:
+                raise WebUISettingsError(str(e), status=400) from e
+        else:
+            from nanobot.providers.xai_oauth import complete_xai_oauth_login
+
+            token = complete_xai_oauth_login(flow, authorization_response)
+    except WebUISettingsError:
+        raise
     except Exception as e:
-        _remove_xai_webui_oauth_flow(flow_id, flow)
-        raise WebUISettingsError(f"xAI OAuth login failed: {e}", status=502) from e
+        _remove_webui_oauth_flow(spec.name, flow_id, flow)
+        raise WebUISettingsError(f"{spec.label} OAuth login failed: {e}", status=502) from e
     if token is None:
         return {
             "status": "pending",
             "provider": spec.name,
             "flow_id": flow_id,
         }
-    _remove_xai_webui_oauth_flow(flow_id, flow, cancel=False)
+    _remove_webui_oauth_flow(spec.name, flow_id, flow, cancel=False)
     if not token.access:
         raise WebUISettingsError("OAuth login failed", status=401)
     return settings_payload()
@@ -1930,6 +1955,7 @@ def logout_oauth_provider(query: QueryParams) -> dict[str, Any]:
             raise WebUISettingsError(
                 "oauth_cli_kit not installed. Run: pip install oauth-cli-kit", status=500
             ) from None
+        _clear_webui_oauth_flows(spec.name)
         token_path = FileTokenStorage(token_filename=OPENAI_CODEX_PROVIDER.token_filename).get_token_path()
     elif spec.name == "github_copilot":
         try:
@@ -1942,7 +1968,7 @@ def logout_oauth_provider(query: QueryParams) -> dict[str, Any]:
     elif spec.name == "xai_grok":
         from nanobot.providers.xai_oauth import logout_xai_oauth
 
-        _clear_xai_webui_oauth_flows()
+        _clear_webui_oauth_flows(spec.name)
         logout_xai_oauth()
         return settings_payload()
     else:
@@ -1954,47 +1980,60 @@ def logout_oauth_provider(query: QueryParams) -> dict[str, Any]:
     return settings_payload()
 
 
-def _register_xai_webui_oauth_flow(flow_id: str, flow: Any) -> None:
+def _register_webui_oauth_flow(provider_name: str, flow_id: str, flow: Any) -> None:
     discarded: list[Any] = []
-    with _xai_webui_oauth_flows_lock:
-        for existing_id, existing in list(_xai_webui_oauth_flows.items()):
+    with _webui_oauth_flows_lock:
+        for existing_id, (_provider_name, existing) in list(_webui_oauth_flows.items()):
             if existing.expired:
-                discarded.append(_xai_webui_oauth_flows.pop(existing_id))
-        while len(_xai_webui_oauth_flows) >= _XAI_WEBUI_OAUTH_MAX_FLOWS:
-            oldest_id = next(iter(_xai_webui_oauth_flows))
-            discarded.append(_xai_webui_oauth_flows.pop(oldest_id))
-        _xai_webui_oauth_flows[flow_id] = flow
+                discarded.append(_webui_oauth_flows.pop(existing_id)[1])
+        while len(_webui_oauth_flows) >= _WEBUI_OAUTH_MAX_FLOWS:
+            oldest_id = next(iter(_webui_oauth_flows))
+            discarded.append(_webui_oauth_flows.pop(oldest_id)[1])
+        _webui_oauth_flows[flow_id] = (provider_name, flow)
     for existing in discarded:
         existing.cancel()
 
 
-def _get_xai_webui_oauth_flow(flow_id: str) -> Any | None:
-    with _xai_webui_oauth_flows_lock:
-        flow = _xai_webui_oauth_flows.get(flow_id)
-        if flow is None or not flow.expired:
+def _get_webui_oauth_flow(provider_name: str, flow_id: str) -> Any | None:
+    with _webui_oauth_flows_lock:
+        registered = _webui_oauth_flows.get(flow_id)
+        if registered is None or registered[0] != provider_name:
+            return None
+        flow = registered[1]
+        if not flow.expired:
             return flow
-        _xai_webui_oauth_flows.pop(flow_id, None)
+        _webui_oauth_flows.pop(flow_id, None)
     flow.cancel()
     return None
 
 
-def _remove_xai_webui_oauth_flow(
+def _remove_webui_oauth_flow(
+    provider_name: str,
     flow_id: str,
     flow: Any,
     *,
     cancel: bool = True,
 ) -> None:
-    with _xai_webui_oauth_flows_lock:
-        if _xai_webui_oauth_flows.get(flow_id) is flow:
-            _xai_webui_oauth_flows.pop(flow_id)
+    with _webui_oauth_flows_lock:
+        registered = _webui_oauth_flows.get(flow_id)
+        if (
+            registered is not None
+            and registered[0] == provider_name
+            and registered[1] is flow
+        ):
+            _webui_oauth_flows.pop(flow_id)
     if cancel:
         flow.cancel()
 
 
-def _clear_xai_webui_oauth_flows() -> None:
-    with _xai_webui_oauth_flows_lock:
-        flows = list(_xai_webui_oauth_flows.values())
-        _xai_webui_oauth_flows.clear()
+def _clear_webui_oauth_flows(provider_name: str) -> None:
+    with _webui_oauth_flows_lock:
+        flow_ids = [
+            flow_id
+            for flow_id, (registered_provider, _flow) in _webui_oauth_flows.items()
+            if registered_provider == provider_name
+        ]
+        flows = [_webui_oauth_flows.pop(flow_id)[1] for flow_id in flow_ids]
     for flow in flows:
         flow.cancel()
 

nanobot/webui/settings_routes.py

@@ -85,7 +85,8 @@ _CHANNEL_VALUES_HEADER_MAX_BYTES = 64 * 1024
 _API_SERVICE_VALUES_HEADER = "X-Nanobot-API-Service-Values"
 _API_SERVICE_VALUES_HEADER_MAX_BYTES = 8 * 1024
 _OAUTH_CODE_HEADER = "X-Nanobot-OAuth-Code"
-_OAUTH_CODE_HEADER_MAX_BYTES = 8 * 1024
+_OAUTH_CALLBACK_HEADER = "X-Nanobot-OAuth-Callback"
+_OAUTH_RESPONSE_HEADER_MAX_BYTES = 8 * 1024
 
 _SKIP_FIELD = object()
 _CHANNEL_CONNECT_ACTIONS = frozenset({"start", "poll", "cancel"})
@@ -471,16 +472,22 @@ class WebUISettingsRouter:
             if action == "login":
                 payload = await asyncio.to_thread(login_oauth_provider, query)
             elif action == "complete":
-                authorization_code = case_insensitive_header(
+                authorization_response = case_insensitive_header(
+                    request.headers,
+                    _OAUTH_CALLBACK_HEADER,
+                ) or case_insensitive_header(
                     request.headers,
                     _OAUTH_CODE_HEADER,
                 )
-                if len(authorization_code.encode("utf-8")) > _OAUTH_CODE_HEADER_MAX_BYTES:
-                    raise WebUISettingsError("OAuth authorization code is too large")
+                if (
+                    len(authorization_response.encode("utf-8"))
+                    > _OAUTH_RESPONSE_HEADER_MAX_BYTES
+                ):
+                    raise WebUISettingsError("OAuth authorization response is too large")
                 payload = await asyncio.to_thread(
                     complete_oauth_provider,
                     query,
-                    authorization_code or None,
+                    authorization_response or None,
                 )
             else:
                 payload = await asyncio.to_thread(logout_oauth_provider, query)