diff --git a/python/api/openai_auth_logout.py b/python/api/openai_auth_logout.py new file mode 100644 index 0000000000..7d8021702c --- /dev/null +++ b/python/api/openai_auth_logout.py @@ -0,0 +1,8 @@ +from python.helpers.api import ApiHandler, Request, Response +from python.helpers import openai_auth + + +class OpenaiAuthLogout(ApiHandler): + async def process(self, input: dict, request: Request) -> dict | Response: + openai_auth.clear_auth() + return {"ok": True} diff --git a/python/api/openai_auth_status.py b/python/api/openai_auth_status.py new file mode 100644 index 0000000000..d0d3f25459 --- /dev/null +++ b/python/api/openai_auth_status.py @@ -0,0 +1,7 @@ +from python.helpers.api import ApiHandler, Request, Response +from python.helpers import openai_auth + + +class OpenaiAuthStatus(ApiHandler): + async def process(self, input: dict, request: Request) -> dict | Response: + return openai_auth.get_auth_status() diff --git a/python/helpers/openai_auth.py b/python/helpers/openai_auth.py new file mode 100644 index 0000000000..0b96e1b84b --- /dev/null +++ b/python/helpers/openai_auth.py @@ -0,0 +1,358 @@ +import base64 +import hashlib +import json +import os +import secrets +import threading +import time +from http import HTTPStatus +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import TypedDict, Any +from urllib.parse import parse_qs, urlencode, urlparse + +import httpx + +from python.helpers import files +from python.helpers.print_style import PrintStyle + + +CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann" +AUTHORIZE_URL = "https://auth.openai.com/oauth/authorize" +TOKEN_URL = "https://auth.openai.com/oauth/token" +REDIRECT_URI = "http://localhost:1455/auth/callback" +SCOPE = "openid profile email offline_access" + +AUTH_FILE = "tmp/openai_auth.json" + + +class OpenAIAuth(TypedDict): + access_token: str + refresh_token: str + expires_at: int + account_id: str | None + + +class AuthorizationFlow(TypedDict): + state: str + verifier: str + url: str + + +_pending_lock = threading.Lock() +_pending_flow: AuthorizationFlow | None = None +_callback_server: ThreadingHTTPServer | None = None +_callback_thread: threading.Thread | None = None + + +def _base64url(data: bytes) -> str: + return base64.urlsafe_b64encode(data).decode().rstrip("=") + + +def _build_pkce_pair() -> tuple[str, str]: + verifier = secrets.token_urlsafe(64) + digest = hashlib.sha256(verifier.encode()).digest() + challenge = _base64url(digest) + return verifier, challenge + + +def _decode_jwt_payload(token: str) -> dict[str, Any] | None: + try: + parts = token.split(".") + if len(parts) != 3: + return None + payload = parts[1] + payload += "=" * (-len(payload) % 4) + decoded = base64.urlsafe_b64decode(payload.encode()).decode("utf-8") + return json.loads(decoded) + except Exception: + return None + + +def _extract_account_id(access_token: str) -> str | None: + payload = _decode_jwt_payload(access_token) or {} + auth_claim = payload.get("https://api.openai.com/auth") or {} + if isinstance(auth_claim, dict): + return auth_claim.get("chatgpt_account_id") + return None + + +def _load_auth_file() -> OpenAIAuth | None: + abs_path = files.get_abs_path(AUTH_FILE) + if not os.path.exists(abs_path): + return None + try: + content = files.read_file(AUTH_FILE) + data = json.loads(content) + if not isinstance(data, dict): + return None + access_token = data.get("access_token") + refresh_token = data.get("refresh_token") + expires_at = data.get("expires_at") + if not access_token or not refresh_token or not isinstance(expires_at, (int, float)): + return None + return OpenAIAuth( + access_token=access_token, + refresh_token=refresh_token, + expires_at=int(expires_at), + account_id=data.get("account_id"), + ) + except Exception: + return None + + +def _save_auth_file(auth: OpenAIAuth) -> None: + files.write_file(AUTH_FILE, json.dumps(auth, indent=2)) + + +def clear_auth() -> None: + abs_path = files.get_abs_path(AUTH_FILE) + if os.path.exists(abs_path): + os.remove(abs_path) + + +def get_auth_status() -> dict[str, Any]: + auth = _load_auth_file() + now = int(time.time() * 1000) + if not auth: + return { + "connected": False, + "expired": False, + "has_token": False, + "expires_at": None, + "account_id": None, + } + expired = auth["expires_at"] <= now + return { + "connected": not expired, + "expired": expired, + "has_token": True, + "expires_at": auth["expires_at"], + "account_id": auth.get("account_id"), + } + + +def create_authorization_flow() -> AuthorizationFlow: + verifier, challenge = _build_pkce_pair() + state = secrets.token_hex(16) + + params = { + "response_type": "code", + "client_id": CLIENT_ID, + "redirect_uri": REDIRECT_URI, + "scope": SCOPE, + "code_challenge": challenge, + "code_challenge_method": "S256", + "state": state, + "id_token_add_organizations": "true", + "codex_cli_simplified_flow": "true", + "originator": "codex_cli_rs", + } + url = f"{AUTHORIZE_URL}?{urlencode(params)}" + return AuthorizationFlow(state=state, verifier=verifier, url=url) + + +def set_pending_flow(flow: AuthorizationFlow) -> None: + with _pending_lock: + global _pending_flow + _pending_flow = flow + + +def _consume_pending_flow() -> AuthorizationFlow | None: + with _pending_lock: + global _pending_flow + flow = _pending_flow + _pending_flow = None + return flow + + +def _peek_pending_flow() -> AuthorizationFlow | None: + with _pending_lock: + return _pending_flow + + +def exchange_authorization_code(code: str, verifier: str) -> OpenAIAuth | None: + response = httpx.post( + TOKEN_URL, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + data={ + "grant_type": "authorization_code", + "client_id": CLIENT_ID, + "code": code, + "code_verifier": verifier, + "redirect_uri": REDIRECT_URI, + }, + timeout=15.0, + ) + if response.status_code != HTTPStatus.OK: + PrintStyle.error( + f"OpenAI OAuth token exchange failed: {response.status_code} {response.text}" + ) + return None + + payload = response.json() + access_token = payload.get("access_token") + refresh_token = payload.get("refresh_token") + expires_in = payload.get("expires_in") + if not access_token or not refresh_token or not isinstance(expires_in, (int, float)): + PrintStyle.error("OpenAI OAuth token response missing fields.") + return None + + account_id = _extract_account_id(access_token) + return OpenAIAuth( + access_token=access_token, + refresh_token=refresh_token, + expires_at=int(time.time() * 1000 + float(expires_in) * 1000), + account_id=account_id, + ) + + +SUCCESS_HTML = """ + + + + OpenAI Auth Complete + + + +
+

Authentication successful

+

You can close this tab and return to Agent Zero.

+
+ + +""" + + +ERROR_HTML = """ + + + + OpenAI Auth Failed + + + +
+

Authentication failed

+

Return to Agent Zero and try signing in again.

+
+ + +""" + + +class _CallbackHandler(BaseHTTPRequestHandler): + def log_message(self, format: str, *args: Any) -> None: + return + + def do_GET(self) -> None: + parsed = urlparse(self.path) + if parsed.path != "/auth/callback": + self._send_text(HTTPStatus.NOT_FOUND, "Not found") + return + + params = parse_qs(parsed.query) + code = (params.get("code") or [None])[0] + state = (params.get("state") or [None])[0] + + flow = _peek_pending_flow() + if not flow or not state or state != flow["state"]: + self._send_html(HTTPStatus.BAD_REQUEST, ERROR_HTML) + return + if not code: + self._send_html(HTTPStatus.BAD_REQUEST, ERROR_HTML) + return + + _consume_pending_flow() + auth = exchange_authorization_code(code, flow["verifier"]) + if not auth: + self._send_html(HTTPStatus.BAD_REQUEST, ERROR_HTML) + return + + _save_auth_file(auth) + PrintStyle().print("OpenAI OAuth completed. Tokens stored.") + self._send_html(HTTPStatus.OK, SUCCESS_HTML) + + def _send_text(self, status: HTTPStatus, text: str) -> None: + self.send_response(status) + self.send_header("Content-Type", "text/plain; charset=utf-8") + self.end_headers() + self.wfile.write(text.encode("utf-8")) + + def _send_html(self, status: HTTPStatus, html: str) -> None: + self.send_response(status) + self.send_header("Content-Type", "text/html; charset=utf-8") + self.end_headers() + self.wfile.write(html.encode("utf-8")) + + +def ensure_callback_server() -> None: + global _callback_server, _callback_thread + if _callback_server: + return + + try: + _callback_server = ThreadingHTTPServer(("127.0.0.1", 1455), _CallbackHandler) + except OSError as exc: + PrintStyle.error(f"OpenAI OAuth callback server unavailable: {exc}") + _callback_server = None + return + + _callback_thread = threading.Thread( + target=_callback_server.serve_forever, daemon=True + ) + _callback_thread.start() diff --git a/run_ui.py b/run_ui.py index 1691f69e74..a60c40c46a 100644 --- a/run_ui.py +++ b/run_ui.py @@ -17,7 +17,7 @@ from python.helpers.extract_tools import load_classes_from_folder from python.helpers.api import ApiHandler from python.helpers.print_style import PrintStyle -from python.helpers import login +from python.helpers import login, openai_auth # disable logging import logging @@ -167,6 +167,14 @@ async def logout_handler(): session.pop('authentication', None) return redirect(url_for('login_handler')) +@webapp.route("/auth/openai", methods=["GET"]) +@requires_auth +async def openai_auth_start(): + flow = openai_auth.create_authorization_flow() + openai_auth.set_pending_flow(flow) + openai_auth.ensure_callback_server() + return redirect(flow["url"]) + # handle default address, load index @webapp.route("/", methods=["GET"]) @requires_auth @@ -283,4 +291,4 @@ def init_a0(): if __name__ == "__main__": runtime.initialize() dotenv.load_dotenv() - run() \ No newline at end of file + run() diff --git a/webui/components/settings/external/external-settings.html b/webui/components/settings/external/external-settings.html index 63816f2c30..dda5df22fc 100644 --- a/webui/components/settings/external/external-settings.html +++ b/webui/components/settings/external/external-settings.html @@ -15,6 +15,12 @@ API Keys +
  • + + OpenAI Auth + OpenAI Auth + +
  • LiteLLM Global Settings @@ -57,6 +63,9 @@
    +
    + +
    diff --git a/webui/components/settings/external/openai_auth.html b/webui/components/settings/external/openai_auth.html new file mode 100644 index 0000000000..05d35b6568 --- /dev/null +++ b/webui/components/settings/external/openai_auth.html @@ -0,0 +1,130 @@ + + + OpenAI Auth + + + + + +
    + +
    + +