From d82c8f02ae07ee8713fff061043a23ebd1aba3d9 Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 14:28:54 -0400 Subject: [PATCH 01/19] Add per-account OpenAI ChatGPT OAuth adapter --- .../adapter_processor_v2/adapter_processor.py | 25 +- backend/adapter_processor_v2/serializers.py | 15 +- backend/adapter_processor_v2/views.py | 107 ++++- backend/connector_auth_v2/openai_oauth.py | 345 ++++++++++++++ backend/connector_auth_v2/urls.py | 4 + backend/connector_auth_v2/views.py | 46 ++ .../input-output/configure-ds/ConfigureDs.jsx | 30 +- .../components/oauth-ds/oauth-ds/OAuthDs.jsx | 133 +++++- .../oauth-ds/openai/OpenAIOAuthButton.css | 17 + .../oauth-ds/openai/OpenAIOAuthButton.jsx | 56 +++ frontend/src/helpers/GetStaticData.js | 1 + .../platform_service/controller/platform.py | 54 ++- .../src/unstract/sdk1/adapters/adapterkit.py | 4 + .../sdk1/src/unstract/sdk1/adapters/base1.py | 6 + .../unstract/sdk1/adapters/llm1/__init__.py | 2 + .../sdk1/adapters/llm1/openai_oauth.py | 106 +++++ .../adapters/llm1/static/openai-oauth.json | 93 ++++ .../sdk1/src/unstract/sdk1/auth/__init__.py | 25 + .../src/unstract/sdk1/auth/openai_oauth.py | 236 ++++++++++ unstract/sdk1/src/unstract/sdk1/llm.py | 437 +++++++++++++++--- unstract/sdk1/tests/test_openai_oauth.py | 209 +++++++++ 21 files changed, 1864 insertions(+), 87 deletions(-) create mode 100644 backend/connector_auth_v2/openai_oauth.py create mode 100644 frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.css create mode 100644 frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx create mode 100644 unstract/sdk1/src/unstract/sdk1/adapters/llm1/openai_oauth.py create mode 100644 unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json create mode 100644 unstract/sdk1/src/unstract/sdk1/auth/__init__.py create mode 100644 unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py create mode 100644 unstract/sdk1/tests/test_openai_oauth.py diff --git a/backend/adapter_processor_v2/adapter_processor.py b/backend/adapter_processor_v2/adapter_processor.py index a9c7f40e32..3db703096c 100644 --- a/backend/adapter_processor_v2/adapter_processor.py +++ b/backend/adapter_processor_v2/adapter_processor.py @@ -8,6 +8,13 @@ from django.core.exceptions import ObjectDoesNotExist from platform_settings_v2.platform_auth_service import PlatformAuthenticationService from tenant_account_v2.organization_member_service import OrganizationMemberService +from unstract.sdk1.adapters.adapterkit import Adapterkit +from unstract.sdk1.adapters.base import Adapter +from unstract.sdk1.adapters.x2text.constants import X2TextConstants +from unstract.sdk1.constants import AdapterTypes +from unstract.sdk1.embedding import EmbeddingCompat +from unstract.sdk1.exceptions import SdkError +from unstract.sdk1.llm import LLM from adapter_processor_v2.constants import AdapterKeys, AllowedDomains from adapter_processor_v2.exceptions import ( @@ -16,13 +23,6 @@ InValidAdapterId, TestAdapterError, ) -from unstract.sdk1.adapters.adapterkit import Adapterkit -from unstract.sdk1.adapters.base import Adapter -from unstract.sdk1.adapters.x2text.constants import X2TextConstants -from unstract.sdk1.constants import AdapterTypes -from unstract.sdk1.embedding import EmbeddingCompat -from unstract.sdk1.exceptions import SdkError -from unstract.sdk1.llm import LLM from .models import AdapterInstance, UserDefaultAdapter @@ -46,6 +46,9 @@ def get_json_schema(adapter_id: str) -> dict[str, Any]: schema_details[AdapterKeys.JSON_SCHEMA] = json.loads( updated_adapters[0].get(AdapterKeys.JSON_SCHEMA) ) + for key in ("oauth", "oauth_provider", "python_social_auth_backend"): + if key in updated_adapters[0]: + schema_details[key] = updated_adapters[0][key] else: logger.error(f"Invalid adapter Id : {adapter_id} while fetching JSON Schema") raise InValidAdapterId() @@ -68,8 +71,7 @@ def get_all_supported_adapters(user_email: str, type: str) -> list[dict[Any, Any if not is_special_user and adapter_id.startswith("noOp"): continue - supported_adapters.append( - { + adapter_details = { AdapterKeys.ID: adapter_id, AdapterKeys.NAME: each_adapter.get(AdapterKeys.NAME), AdapterKeys.DESCRIPTION: each_adapter.get(AdapterKeys.DESCRIPTION), @@ -77,7 +79,10 @@ def get_all_supported_adapters(user_email: str, type: str) -> list[dict[Any, Any AdapterKeys.ADAPTER_TYPE: each_adapter.get(AdapterKeys.ADAPTER_TYPE), AdapterKeys.DOC_URL: each_adapter.get(AdapterKeys.DOC_URL, ""), } - ) + for key in ("oauth", "oauth_provider", "python_social_auth_backend"): + if key in each_adapter: + adapter_details[key] = each_adapter[key] + supported_adapters.append(adapter_details) return supported_adapters @staticmethod diff --git a/backend/adapter_processor_v2/serializers.py b/backend/adapter_processor_v2/serializers.py index 8209dd75c7..123be31f45 100644 --- a/backend/adapter_processor_v2/serializers.py +++ b/backend/adapter_processor_v2/serializers.py @@ -2,6 +2,9 @@ from typing import Any from account_v2.serializer import UserSerializer +from backend.constants import FieldLengthConstants as FLC +from backend.serializers import AuditSerializer +from connector_auth_v2.openai_oauth import redact_openai_oauth_metadata from cryptography.fernet import Fernet from django.conf import settings from rest_framework import serializers @@ -10,14 +13,13 @@ serialize_group_refs, serialize_owner_refs, ) +from unstract.sdk1.auth.openai_oauth import is_openai_oauth_adapter +from unstract.sdk1.constants import AdapterTypes +from unstract.sdk1.constants import Common as common from utils.input_sanitizer import validate_name_field, validate_no_html_tags from adapter_processor_v2.adapter_processor import AdapterProcessor from adapter_processor_v2.constants import AdapterKeys -from backend.constants import FieldLengthConstants as FLC -from backend.serializers import AuditSerializer -from unstract.sdk1.constants import AdapterTypes -from unstract.sdk1.constants import Common as common from .models import AdapterInstance, UserDefaultAdapter @@ -87,6 +89,11 @@ def to_representation(self, instance: AdapterInstance) -> dict[str, str]: rep.pop(AdapterKeys.ADAPTER_METADATA_B) adapter_metadata = instance.metadata + if is_openai_oauth_adapter(instance.adapter_id): + # OAuth tokens and account identity stay encrypted in the database + # and are never returned to the browser after the initial login. + adapter_metadata = redact_openai_oauth_metadata(adapter_metadata) + # Hide unstract_key when use_platform_provided_unstract_key is True if ( adapter_metadata.get("use_platform_provided_unstract_key") is True diff --git a/backend/adapter_processor_v2/views.py b/backend/adapter_processor_v2/views.py index f2aef82b0c..c69b0e0fdb 100644 --- a/backend/adapter_processor_v2/views.py +++ b/backend/adapter_processor_v2/views.py @@ -3,6 +3,7 @@ from typing import Any from account_v2.models import User +from connector_auth_v2.openai_oauth import OpenAIOAuthService from django.db import IntegrityError from django.db.models import ProtectedError, QuerySet from django.http import HttpRequest @@ -19,7 +20,7 @@ from plugins import get_plugin from rest_framework import status from rest_framework.decorators import action -from rest_framework.exceptions import PermissionDenied +from rest_framework.exceptions import PermissionDenied, ValidationError from rest_framework.request import Request from rest_framework.response import Response from rest_framework.serializers import ModelSerializer @@ -27,6 +28,11 @@ from rest_framework.viewsets import GenericViewSet, ModelViewSet from tenant_account_v2.organization_member_service import OrganizationMemberService from tool_instance_v2.models import ToolInstance +from unstract.sdk1.auth.openai_oauth import ( + OPENAI_OAUTH_PRIVATE_FIELDS, + OpenAIOAuthError, + is_openai_oauth_adapter, +) from utils.filtering import FilterHelper from utils.pagination import OptionalPagination from utils.user_context import UserContext @@ -60,6 +66,63 @@ logger = logging.getLogger(__name__) +def _prepare_openai_oauth_payload( + request: Request, + payload: Any, + existing_metadata: dict[str, Any] | None = None, +) -> tuple[Any, str | None]: + """Inject credentials from one owned OAuth login into an adapter payload.""" + adapter_id = payload.get(AdapterKeys.ADAPTER_ID) + if not is_openai_oauth_adapter(adapter_id): + return payload, None + + submitted_metadata = payload.get(AdapterKeys.ADAPTER_METADATA) + if submitted_metadata is None and existing_metadata is not None: + adapter_metadata = dict(existing_metadata) + elif isinstance(submitted_metadata, dict): + adapter_metadata = dict(submitted_metadata) + else: + adapter_metadata = {} + + # Tokens/account identity are always sourced from the server-side login + # session or the already-encrypted row. Never trust values sent in JSON. + for key in OPENAI_OAUTH_PRIVATE_FIELDS: + adapter_metadata.pop(key, None) + adapter_metadata.pop("oauth_authenticated", None) + adapter_metadata.pop("oauth_account_label", None) + + oauth_key = request.query_params.get("oauth-key") + if oauth_key: + try: + credentials = OpenAIOAuthService.credentials_for_request(oauth_key, request) + except OpenAIOAuthError as exc: + raise ValidationError({"oauth-key": str(exc)}) from exc + adapter_metadata.update(credentials) + elif existing_metadata is not None: + for key in OPENAI_OAUTH_PRIVATE_FIELDS: + if key in existing_metadata: + adapter_metadata[key] = existing_metadata[key] + else: + raise ValidationError( + {"oauth-key": "OpenAI OAuth authentication is required for this adapter."} + ) + + payload[AdapterKeys.ADAPTER_METADATA] = adapter_metadata + return payload, oauth_key + + +def _consume_openai_oauth_key(cache_key: str | None, request: Request) -> None: + """Best-effort cleanup after credentials are durably stored.""" + if not cache_key: + return + try: + OpenAIOAuthService.consume(cache_key, request) + except OpenAIOAuthError: + # The adapter row is already the durable credential store. A cache + # expiry/race must not turn a successful create/update into a 500. + logger.warning("Could not consume OpenAI OAuth hand-off session") + + class DefaultAdapterViewSet(ModelViewSet): versioning_class = URLPathVersioning serializer_class = DefaultAdapterSerializer @@ -127,10 +190,13 @@ def get_adapter_schema( def test(self, request: Request) -> Response: """Tests the connector against the credentials passed.""" - serializer: AdapterInstanceSerializer = self.get_serializer(data=request.data) + payload = request.data.copy() + payload, _ = _prepare_openai_oauth_payload(request, payload) + serializer: AdapterInstanceSerializer = self.get_serializer(data=payload) serializer.is_valid(raise_exception=True) adapter_id = serializer.validated_data.get(AdapterKeys.ADAPTER_ID) adapter_metadata = serializer.validated_data.get(AdapterKeys.ADAPTER_METADATA) + adapter_metadata = dict(adapter_metadata or {}) adapter_metadata[AdapterKeys.ADAPTER_TYPE] = serializer.validated_data.get( AdapterKeys.ADAPTER_TYPE ) @@ -237,10 +303,12 @@ def _enforce_llm_creation_restriction(request: Any, adapter_type: str) -> None: ) def create(self, request: Any) -> Response: - serializer = self.get_serializer(data=request.data) + payload = request.data.copy() + payload, oauth_key = _prepare_openai_oauth_payload(request, payload) + serializer = self.get_serializer(data=payload) use_platform_unstract_key = False - adapter_metadata = request.data.get(AdapterKeys.ADAPTER_METADATA) + adapter_metadata = payload.get(AdapterKeys.ADAPTER_METADATA) if adapter_metadata and adapter_metadata.get( AdapterKeys.PLATFORM_PROVIDED_UNSTRACT_KEY, False ): @@ -312,6 +380,9 @@ def create(self, request: Any) -> Response: user_default_adapter.save() + # The encrypted adapter row is now the durable credential store. + _consume_openai_oauth_key(oauth_key, request) + except IntegrityError: raise DuplicateAdapterNameError( name=serializer.validated_data.get(AdapterKeys.ADAPTER_NAME) @@ -410,7 +481,10 @@ def partial_update( adapter = self.get_object() before = self.snapshot_share_axes(adapter) - response = super().partial_update(request, *args, **kwargs) + if is_openai_oauth_adapter(adapter.adapter_id): + response = self._update_openai_oauth(request, adapter, partial=True) + else: + response = super().partial_update(request, *args, **kwargs) if response.status_code == 200 and notification_plugin: self._notify_shared_users(adapter, before, request.data, request.user) return response @@ -555,6 +629,13 @@ def list_of_shared_users(self, request: HttpRequest, pk: Any = None) -> Response def update( self, request: Request, *args: tuple[Any], **kwargs: dict[str, Any] ) -> Response: + # OAuth adapters carry their credentials in the encrypted metadata row; + # inject a newly completed login or preserve the existing account when + # a metadata-only edit does not include a new login session. + adapter = self.get_object() + if is_openai_oauth_adapter(adapter.adapter_id): + return self._update_openai_oauth(request, adapter, partial=False) + # Check if adapter metadata is being updated and contains the platform key flag use_platform_unstract_key = False adapter_metadata = request.data.get(AdapterKeys.ADAPTER_METADATA) @@ -565,9 +646,6 @@ def update( use_platform_unstract_key = True logger.error(f"Platform key flag detected: {use_platform_unstract_key}") - # Get the adapter instance for update - adapter = self.get_object() - if use_platform_unstract_key: logger.error("Processing adapter with platform key") serializer = self.get_serializer(adapter, data=request.data, partial=True) @@ -597,6 +675,19 @@ def update( # For non-platform-key cases, use the default update behavior return super().update(request, *args, **kwargs) + def _update_openai_oauth( + self, request: Request, adapter: AdapterInstance, *, partial: bool + ) -> Response: + payload = request.data.copy() + payload, oauth_key = _prepare_openai_oauth_payload( + request, payload, existing_metadata=adapter.metadata + ) + serializer = self.get_serializer(adapter, data=payload, partial=partial) + serializer.is_valid(raise_exception=True) + serializer.save() + _consume_openai_oauth_key(oauth_key, request) + return Response(serializer.data) + @action(detail=True, methods=["get"]) def adapter_info(self, request: HttpRequest, pk: uuid) -> Response: adapter = self.get_object() diff --git a/backend/connector_auth_v2/openai_oauth.py b/backend/connector_auth_v2/openai_oauth.py new file mode 100644 index 0000000000..68f851e364 --- /dev/null +++ b/backend/connector_auth_v2/openai_oauth.py @@ -0,0 +1,345 @@ +"""Server-side device login and credential hand-off for OpenAI OAuth.""" + +from __future__ import annotations + +import json +import os +import time +import uuid +from datetime import UTC, datetime +from typing import Any + +import httpx +from cryptography.fernet import Fernet, InvalidToken +from django.conf import settings +from django.core.cache import cache +from rest_framework.request import Request +from unstract.sdk1.auth.openai_oauth import ( + OPENAI_OAUTH_CLIENT_ID, + OPENAI_OAUTH_DEVICE_REDIRECT_URI, + OPENAI_OAUTH_DEVICE_TOKEN_URL, + OPENAI_OAUTH_DEVICE_USERCODE_URL, + OPENAI_OAUTH_DEVICE_VERIFICATION_URL, + OPENAI_OAUTH_PRIVATE_FIELDS, + OPENAI_OAUTH_TOKEN_URL, + OpenAIOAuthError, + extract_account_id, + extract_email, + refresh_openai_oauth_metadata, +) +from utils.user_session import UserSessionUtils + +_CACHE_PREFIX = "openai-oauth:" +_DEFAULT_STATE_TTL_SECONDS = 900 +_STATE_TTL_SECONDS = int( + os.environ.get( + "OPENAI_OAUTH_STATE_TTL_SECONDS", str(_DEFAULT_STATE_TTL_SECONDS) + ) +) + + +class OpenAIOAuthSessionError(OpenAIOAuthError): + """Raised for an invalid, expired, or unauthorized browser login session.""" + + +def _safe_response_json(response: httpx.Response, operation: str) -> dict[str, Any]: + if not 200 <= response.status_code < 300: + raise OpenAIOAuthError( + f"OpenAI OAuth {operation} failed with status {response.status_code}" + ) + try: + payload = response.json() + except (ValueError, json.JSONDecodeError) as exc: + raise OpenAIOAuthError( + f"OpenAI OAuth {operation} returned an invalid response" + ) from exc + if not isinstance(payload, dict): + raise OpenAIOAuthError( + f"OpenAI OAuth {operation} returned an invalid response" + ) + return payload + + +def _post_json(url: str, *, json_body: dict[str, Any], operation: str) -> dict[str, Any]: + try: + response = httpx.post(url, json=json_body, timeout=15.0) + except httpx.HTTPError as exc: + raise OpenAIOAuthError( + f"OpenAI OAuth {operation} could not reach the authorization server" + ) from exc + return _safe_response_json(response, operation) + + +def _post_form( + url: str, *, form_data: dict[str, str], operation: str +) -> dict[str, Any]: + try: + response = httpx.post( + url, + data=form_data, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + timeout=15.0, + ) + except httpx.HTTPError as exc: + raise OpenAIOAuthError( + f"OpenAI OAuth {operation} could not reach the authorization server" + ) from exc + return _safe_response_json(response, operation) + + +def _as_expiry_seconds(value: object) -> float | None: + if value is None: + return None + if isinstance(value, (int, float)): + return float(value) + if isinstance(value, str): + try: + return float(value) + except ValueError: + try: + return datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp() + except ValueError: + return None + return None + + +def _account_label(account_id: str, email: str | None) -> str: + if email: + return email + suffix = account_id[-8:] if len(account_id) > 8 else account_id + return f"OpenAI account ({suffix})" + + +class OpenAIOAuthService: + """Own one short-lived OAuth login session per browser/account.""" + + @staticmethod + def _identity(request: Request) -> tuple[str, str]: + user_id = getattr(request.user, "pk", None) or getattr(request.user, "id", None) + organization_id = UserSessionUtils.get_organization_id(request) + if user_id is None or not organization_id: + raise OpenAIOAuthSessionError( + "An authenticated organization session is required for OpenAI OAuth" + ) + return str(user_id), str(organization_id) + + @staticmethod + def _encrypt(credentials: dict[str, Any]) -> str: + fernet = Fernet(str(settings.ENCRYPTION_KEY).encode("utf-8")) + return fernet.encrypt(json.dumps(credentials).encode("utf-8")).decode("utf-8") + + @staticmethod + def _decrypt(value: str) -> dict[str, Any]: + try: + fernet = Fernet(str(settings.ENCRYPTION_KEY).encode("utf-8")) + credentials = json.loads(fernet.decrypt(value.encode("utf-8"))) + except (InvalidToken, ValueError, json.JSONDecodeError) as exc: + raise OpenAIOAuthSessionError( + "OpenAI OAuth login session is no longer valid" + ) from exc + if not isinstance(credentials, dict): + raise OpenAIOAuthSessionError( + "OpenAI OAuth login session is no longer valid" + ) + return credentials + + @classmethod + def _save_state(cls, cache_key: str, state: dict[str, Any]) -> None: + expiry = _as_expiry_seconds(state.get("expires_at")) + ttl = _STATE_TTL_SECONDS + if expiry is not None: + ttl = min(ttl, max(30, int(expiry - time.time()))) + cache.set(cache_key, state, max(ttl, 30)) + + @classmethod + def _get_owned_state(cls, cache_key: str, request: Request) -> dict[str, Any]: + if not cache_key or not cache_key.startswith(_CACHE_PREFIX): + raise OpenAIOAuthSessionError("OpenAI OAuth login session is invalid") + user_id, organization_id = cls._identity(request) + state = cache.get(cache_key) + if not isinstance(state, dict): + raise OpenAIOAuthSessionError( + "OpenAI OAuth login session was not found or has expired" + ) + if state.get("owner_id") != user_id or state.get("organization_id") != organization_id: + raise OpenAIOAuthSessionError( + "OpenAI OAuth login session was not found or is not owned by this user" + ) + expiry = _as_expiry_seconds(state.get("expires_at")) + if expiry is not None and expiry <= time.time(): + cache.delete(cache_key) + raise OpenAIOAuthSessionError("OpenAI OAuth login session has expired") + return state + + @classmethod + def start(cls, request: Request) -> dict[str, Any]: + user_id, organization_id = cls._identity(request) + payload = _post_json( + OPENAI_OAUTH_DEVICE_USERCODE_URL, + json_body={"client_id": OPENAI_OAUTH_CLIENT_ID}, + operation="device login start", + ) + device_auth_id = payload.get("device_auth_id") + user_code = payload.get("user_code") or payload.get("usercode") + if not isinstance(device_auth_id, str) or not device_auth_id: + raise OpenAIOAuthError("OpenAI OAuth device login returned no device id") + if not isinstance(user_code, str) or not user_code: + raise OpenAIOAuthError("OpenAI OAuth device login returned no user code") + + try: + interval = max(1, int(payload.get("interval", 5))) + except (TypeError, ValueError): + interval = 5 + expires_at = payload.get("expires_at") + if not isinstance(expires_at, str): + expires_at = datetime.fromtimestamp( + time.time() + _STATE_TTL_SECONDS, tz=UTC + ).isoformat() + + cache_key = f"{_CACHE_PREFIX}{uuid.uuid4().hex}" + cls._save_state( + cache_key, + { + "status": "pending", + "owner_id": user_id, + "organization_id": organization_id, + "device_auth_id": device_auth_id, + "user_code": user_code, + "poll_interval": interval, + "expires_at": expires_at, + }, + ) + return { + "cache_key": cache_key, + "verification_url": OPENAI_OAUTH_DEVICE_VERIFICATION_URL, + "user_code": user_code, + "expires_at": expires_at, + "poll_interval": interval, + } + + @classmethod + def _success_response(cls, state: dict[str, Any]) -> dict[str, Any]: + return { + "status": "success", + "account_label": state.get("account_label", "OpenAI account"), + } + + @classmethod + def poll(cls, request: Request, cache_key: str) -> dict[str, Any]: + state = cls._get_owned_state(cache_key, request) + if state.get("status") == "success": + return cls._success_response(state) + + try: + response = httpx.post( + OPENAI_OAUTH_DEVICE_TOKEN_URL, + json={ + "device_auth_id": state["device_auth_id"], + "user_code": state["user_code"], + }, + timeout=15.0, + ) + except httpx.HTTPError as exc: + raise OpenAIOAuthError( + "OpenAI OAuth device login could not reach the authorization server" + ) from exc + + # The official device endpoint uses 403/404 while the user has not + # finished entering the code. Keep the browser poll alive for those + # expected responses (and for rate limiting). + if response.status_code in {403, 404, 429}: + return { + "status": "pending", + "poll_interval": state.get("poll_interval", 5), + } + device_result = _safe_response_json(response, "device login poll") + authorization_code = device_result.get("authorization_code") + code_verifier = device_result.get("code_verifier") + if not isinstance(authorization_code, str) or not isinstance( + code_verifier, str + ): + raise OpenAIOAuthError( + "OpenAI OAuth device login returned incomplete authorization data" + ) + + token_data = _post_form( + OPENAI_OAUTH_TOKEN_URL, + form_data={ + "grant_type": "authorization_code", + "code": authorization_code, + "redirect_uri": OPENAI_OAUTH_DEVICE_REDIRECT_URI, + "client_id": OPENAI_OAUTH_CLIENT_ID, + "code_verifier": code_verifier, + }, + operation="token exchange", + ) + access_token = token_data.get("access_token") + refresh_token = token_data.get("refresh_token") + id_token = token_data.get("id_token") + if not all(isinstance(value, str) and value for value in (access_token, refresh_token)): + raise OpenAIOAuthError( + "OpenAI OAuth token exchange returned incomplete credentials" + ) + + account_id = extract_account_id(id_token, access_token) + if not account_id: + raise OpenAIOAuthError( + "OpenAI OAuth token exchange returned no ChatGPT account" + ) + email = extract_email(id_token, access_token) + try: + expires_in = float(token_data.get("expires_in", 3600)) + except (TypeError, ValueError): + expires_in = 3600 + credentials = { + "oauth_access_token": access_token, + "oauth_refresh_token": refresh_token, + "oauth_id_token": id_token, + "oauth_account_id": account_id, + "oauth_account_email": email, + "oauth_expires_at": time.time() + expires_in, + "oauth_authenticated": True, + } + state.update( + { + "status": "success", + "credentials": cls._encrypt(credentials), + "account_label": _account_label(account_id, email), + } + ) + cls._save_state(cache_key, state) + return cls._success_response(state) + + @classmethod + def credentials_for_request( + cls, cache_key: str, request: Request + ) -> dict[str, Any]: + state = cls._get_owned_state(cache_key, request) + if state.get("status") != "success" or not state.get("credentials"): + raise OpenAIOAuthSessionError("Complete OpenAI OAuth authentication first") + credentials = cls._decrypt(state["credentials"]) + refreshed = refresh_openai_oauth_metadata(credentials) + if refreshed != credentials: + state["credentials"] = cls._encrypt(refreshed) + cls._save_state(cache_key, state) + return refreshed + + @classmethod + def consume(cls, cache_key: str, request: Request) -> None: + # Validate ownership before deleting so a leaked cache key cannot be + # used to consume another user's pending login. + cls._get_owned_state(cache_key, request) + cache.delete(cache_key) + + +def redact_openai_oauth_metadata(metadata: dict[str, Any]) -> dict[str, Any]: + """Return the non-secret account state safe for API responses.""" + redacted = { + key: value for key, value in metadata.items() if key not in OPENAI_OAUTH_PRIVATE_FIELDS + } + redacted["oauth_authenticated"] = bool(metadata.get("oauth_access_token")) + account_id = metadata.get("oauth_account_id") + email = metadata.get("oauth_account_email") + if isinstance(account_id, str): + redacted["oauth_account_label"] = _account_label(account_id, email) + return redacted diff --git a/backend/connector_auth_v2/urls.py b/backend/connector_auth_v2/urls.py index 55337ad20f..b4532ee3d9 100644 --- a/backend/connector_auth_v2/urls.py +++ b/backend/connector_auth_v2/urls.py @@ -8,6 +8,8 @@ "get": "cache_key", } ) +openai_oauth_start = ConnectorAuthViewSet.as_view({"post": "openai_start"}) +openai_oauth_poll = ConnectorAuthViewSet.as_view({"get": "openai_poll"}) urlpatterns = format_suffix_patterns( [ @@ -17,5 +19,7 @@ connector_auth_cache, name="connector-cache", ), + path("oauth/openai/start/", openai_oauth_start, name="openai-oauth-start"), + path("oauth/openai/poll/", openai_oauth_poll, name="openai-oauth-poll"), ] ) diff --git a/backend/connector_auth_v2/views.py b/backend/connector_auth_v2/views.py index 4d515563b6..a17bd99228 100644 --- a/backend/connector_auth_v2/views.py +++ b/backend/connector_auth_v2/views.py @@ -6,10 +6,15 @@ from rest_framework.request import Request from rest_framework.response import Response from rest_framework.versioning import URLPathVersioning +from unstract.sdk1.auth.openai_oauth import OpenAIOAuthError from utils.user_session import UserSessionUtils from connector_auth_v2.constants import SocialAuthConstants from connector_auth_v2.exceptions import KeyNotConfigured +from connector_auth_v2.openai_oauth import ( + OpenAIOAuthService, + OpenAIOAuthSessionError, +) logger = logging.getLogger(__name__) @@ -47,3 +52,44 @@ def cache_key( status=status.HTTP_200_OK, data={"cache_key": f"{cache_key}"}, ) + + @staticmethod + def _require_authenticated(request: Request) -> Response | None: + if not getattr(request.user, "is_authenticated", False): + return Response( + {"message": "Authentication is required for OpenAI OAuth."}, + status=status.HTTP_401_UNAUTHORIZED, + ) + return None + + def openai_start(self, request: Request) -> Response: + """Start a server-side OpenAI device-code login.""" + if unauthorized := self._require_authenticated(request): + return unauthorized + try: + return Response( + OpenAIOAuthService.start(request), status=status.HTTP_200_OK + ) + except OpenAIOAuthSessionError as exc: + return Response({"message": str(exc)}, status=status.HTTP_400_BAD_REQUEST) + except OpenAIOAuthError as exc: + logger.warning("OpenAI OAuth device login start failed: %s", exc) + return Response( + {"message": str(exc)}, status=status.HTTP_502_BAD_GATEWAY + ) + + def openai_poll(self, request: Request) -> Response: + """Poll one OpenAI device-code login and exchange it when complete.""" + if unauthorized := self._require_authenticated(request): + return unauthorized + cache_key = request.query_params.get("oauth-key") + try: + result = OpenAIOAuthService.poll(request, cache_key or "") + return Response(result, status=status.HTTP_200_OK) + except OpenAIOAuthSessionError as exc: + return Response({"message": str(exc)}, status=status.HTTP_400_BAD_REQUEST) + except OpenAIOAuthError as exc: + logger.warning("OpenAI OAuth device login poll failed: %s", exc) + return Response( + {"message": str(exc)}, status=status.HTTP_502_BAD_GATEWAY + ) diff --git a/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx b/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx index 5bb0d6f632..e26c88a8ad 100644 --- a/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx +++ b/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx @@ -56,7 +56,10 @@ function ConfigureDs({ // Determine if this is a new or existing connector const hasOAuthCredentials = - metadata && (metadata.access_token || (metadata.provider && metadata.uid)); + metadata && + (metadata.access_token || + (metadata.provider && metadata.uid) || + metadata.oauth_authenticated); const isExistingConnector = Boolean(editItemId || hasOAuthCredentials); // Determine if OAuth authentication method is selected @@ -176,7 +179,11 @@ function ConfigureDs({ (status !== "success" || !cacheKey?.length) ) { const providerName = - oAuthProvider === "google-oauth2" ? "Google" : "OAuth provider"; + oAuthProvider === "google-oauth2" + ? "Google" + : oAuthProvider === "openai-oauth" + ? "OpenAI" + : "OAuth provider"; setAlertDetails({ type: "error", content: `OAuth authentication required. Please sign in with ${providerName} first.`, @@ -211,10 +218,14 @@ function ConfigureDs({ } if (oAuthProvider?.length > 0 && isOAuthMethodSelected()) { - body["connector_metadata"] = { - ...body["connector_metadata"], - ...{ "oauth-key": cacheKey }, - }; + if (isConnector) { + body["connector_metadata"] = { + ...body["connector_metadata"], + ...{ "oauth-key": cacheKey }, + }; + } else { + url = `${url}?oauth-key=${encodeURIComponent(cacheKey)}`; + } } const requestOptions = { @@ -311,7 +322,11 @@ function ConfigureDs({ url = `${url}${editItemId}/`; } - if (oAuthProvider?.length > 0 && isOAuthMethodSelected()) { + if ( + oAuthProvider?.length > 0 && + isOAuthMethodSelected() && + cacheKey?.length + ) { const encodedCacheKey = encodeURIComponent(cacheKey); url = url + `?oauth-key=${encodedCacheKey}`; } @@ -342,6 +357,7 @@ function ConfigureDs({ if (oAuthProvider?.length > 0 && isOAuthMethodSelected()) { localStorage.removeItem(oauthCacheKey); localStorage.removeItem(oauthStatusKey); + localStorage.removeItem(`oauth-device-${selectedSourceId}`); } setOpen(false); diff --git a/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx b/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx index cc40c50bbf..9019e5d17c 100644 --- a/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx +++ b/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx @@ -1,5 +1,5 @@ import PropTypes from "prop-types"; -import { useEffect, useState } from "react"; +import { useCallback, useEffect, useState } from "react"; import { Typography } from "@/components/ui/shims/antd-typography"; import { getBaseUrl, O_AUTH_PROVIDERS } from "../../../helpers/GetStaticData"; @@ -8,6 +8,7 @@ import { useExceptionHandler } from "../../../hooks/useExceptionHandler.jsx"; import { useAlertStore } from "../../../store/alert-store"; import GoogleOAuthButton from "../google/GoogleOAuthButton.jsx"; import MicrosoftOAuthButton from "../microsoft/MicrosoftOAuthButton.jsx"; +import OpenAIOAuthButton from "../openai/OpenAIOAuthButton.jsx"; function OAuthDs({ oAuthProvider, @@ -36,6 +37,9 @@ function OAuthDs({ if (oAuthProvider === O_AUTH_PROVIDERS.GOOGLE) { return "Authenticate with Google"; } + if (oAuthProvider === O_AUTH_PROVIDERS.OPENAI) { + return "Sign in with OpenAI"; + } return "Authenticate"; }; @@ -45,6 +49,22 @@ function OAuthDs({ // Initialize from connector-specific status return localStorage.getItem(oauthStatusKey); }); + const [loginCacheKey, setLoginCacheKey] = useState(() => + localStorage.getItem(oauthCacheKey), + ); + const [deviceLogin, setDeviceLogin] = useState(() => { + try { + return JSON.parse(localStorage.getItem(`oauth-device-${selectedSourceId}`)); + } catch { + return null; + } + }); + + const updateOAuthStatus = useCallback((newStatus) => { + setOAuthStatus(newStatus); + setStatus(newStatus); + localStorage.setItem(oauthStatusKey, newStatus); + }, [oauthStatusKey, setStatus]); useEffect(() => { const handleStorageChange = () => { @@ -60,15 +80,29 @@ function OAuthDs({ // Load persisted cache key if available const persistedCacheKey = localStorage.getItem(oauthCacheKey); + setLoginCacheKey(persistedCacheKey || null); if (persistedCacheKey) { setCacheKey(persistedCacheKey); + } else { + setCacheKey(""); } // Set initial status from connector-specific status const connectorStatus = localStorage.getItem(oauthStatusKey); - if (connectorStatus) { - setStatus(connectorStatus); - setOAuthStatus(connectorStatus); + setStatus(connectorStatus || ""); + setOAuthStatus(connectorStatus || ""); + + const persistedDeviceLogin = localStorage.getItem( + `oauth-device-${selectedSourceId}`, + ); + if (persistedDeviceLogin) { + try { + setDeviceLogin(JSON.parse(persistedDeviceLogin)); + } catch { + localStorage.removeItem(`oauth-device-${selectedSourceId}`); + } + } else { + setDeviceLogin(null); } return () => { @@ -77,8 +111,85 @@ function OAuthDs({ }; }, [selectedSourceId, oauthCacheKey, oauthStatusKey, setCacheKey, setStatus]); + useEffect(() => { + if ( + oAuthProvider !== O_AUTH_PROVIDERS.OPENAI || + oauthStatus !== "pending" || + !loginCacheKey + ) { + return undefined; + } + + let isActive = true; + const pollLogin = async () => { + try { + const response = await axiosPrivate({ + method: "GET", + url: `/api/v1/oauth/openai/poll/?oauth-key=${encodeURIComponent(loginCacheKey)}`, + }); + if (!isActive) { + return; + } + const result = response?.data || {}; + if (result.status === "success") { + setDeviceLogin((current) => ({ + ...(current || {}), + account_label: result.account_label, + })); + updateOAuthStatus("success"); + } + } catch (err) { + if (!isActive) { + return; + } + const message = + err?.response?.data?.message || "OpenAI authentication failed"; + updateOAuthStatus("error"); + setAlertDetails(handleException(err, message)); + } + }; + + const initialPoll = setTimeout(pollLogin, 1000); + const pollInterval = setInterval(pollLogin, 5000); + return () => { + isActive = false; + clearTimeout(initialPoll); + clearInterval(pollInterval); + }; + }, [ + axiosPrivate, + handleException, + loginCacheKey, + oauthStatus, + oAuthProvider, + setAlertDetails, + updateOAuthStatus, + ]); + const handleOAuth = async () => { try { + if (oAuthProvider === O_AUTH_PROVIDERS.OPENAI) { + const response = await axiosPrivate({ + method: "POST", + url: "/api/v1/oauth/openai/start/", + }); + const loginDetails = response?.data || {}; + const newCacheKey = loginDetails.cache_key; + if (!newCacheKey) { + throw new Error("OpenAI OAuth did not return a login session"); + } + setLoginCacheKey(newCacheKey); + setCacheKey(newCacheKey); + localStorage.setItem(oauthCacheKey, newCacheKey); + setDeviceLogin(loginDetails); + localStorage.setItem( + `oauth-device-${selectedSourceId}`, + JSON.stringify(loginDetails), + ); + updateOAuthStatus("pending"); + return; + } + // Store connector context in sessionStorage for OAuth callback (survives window.open) sessionStorage.setItem("oauth-current-connector", selectedSourceId); @@ -136,6 +247,20 @@ function OAuthDs({ ); } + if (O_AUTH_PROVIDERS.OPENAI === oAuthProvider) { + return ( + + ); + } + return Provider not available.; } diff --git a/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.css b/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.css new file mode 100644 index 0000000000..2aa3289628 --- /dev/null +++ b/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.css @@ -0,0 +1,17 @@ +.openai-oauth-layout { + display: flex; + flex-direction: column; + gap: 8px; + margin-bottom: 20px; +} + +.openai-oauth-instructions { + margin: 0; + font-size: 12px; + line-height: 1.5; +} + +.openai-oauth-account { + display: block; + font-size: 12px; +} diff --git a/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx b/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx new file mode 100644 index 0000000000..565a62146c --- /dev/null +++ b/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx @@ -0,0 +1,56 @@ +import PropTypes from "prop-types"; +import { Button } from "@/components/ui/shims/antd-button"; +import { Typography } from "@/components/ui/shims/antd-typography"; + +import "./OpenAIOAuthButton.css"; + +const OpenAIOAuthButton = ({ + handleOAuth, + status, + buttonText = "Sign in with OpenAI", + disabled = false, + verificationUrl, + userCode, + accountLabel, +}) => { + const buttonLabel = status === "success" ? "Authenticated" : buttonText; + + return ( +
+ + {status === "pending" && verificationUrl && userCode && ( + + Open{" "} + + OpenAI device login + {" "} + and enter code {userCode}. + + )} + {status === "success" && accountLabel && ( + + {accountLabel} + + )} +
+ ); +}; + +OpenAIOAuthButton.propTypes = { + handleOAuth: PropTypes.func.isRequired, + status: PropTypes.string, + buttonText: PropTypes.string, + disabled: PropTypes.bool, + verificationUrl: PropTypes.string, + userCode: PropTypes.string, + accountLabel: PropTypes.string, +}; + +export default OpenAIOAuthButton; diff --git a/frontend/src/helpers/GetStaticData.js b/frontend/src/helpers/GetStaticData.js index 9e067896c0..35abc833e3 100644 --- a/frontend/src/helpers/GetStaticData.js +++ b/frontend/src/helpers/GetStaticData.js @@ -85,6 +85,7 @@ const formatBytes = (bytes, decimals = 1) => { const O_AUTH_PROVIDERS = { GOOGLE: "google-oauth2", MICROSOFT: "azuread-tenant-oauth2", + OPENAI: "openai-oauth", }; const CONNECTOR_TYPE_MAP = { diff --git a/platform-service/src/unstract/platform_service/controller/platform.py b/platform-service/src/unstract/platform_service/controller/platform.py index fbc8cdd8bb..dacc06a959 100644 --- a/platform-service/src/unstract/platform_service/controller/platform.py +++ b/platform-service/src/unstract/platform_service/controller/platform.py @@ -7,7 +7,6 @@ from cryptography.fernet import Fernet, InvalidToken from flask import Blueprint, Request, jsonify, make_response, request from flask import current_app as app - from unstract.core.flask import PluginManager from unstract.core.flask.exceptions import APIError from unstract.platform_service.constants import DBTable @@ -17,6 +16,11 @@ AdapterInstanceRequestHelper, ) from unstract.platform_service.helper.prompt_studio import PromptStudioRequestHelper +from unstract.sdk1.auth.openai_oauth import ( + OpenAIOAuthRefreshError, + is_openai_oauth_adapter, + refresh_openai_oauth_metadata, +) platform_bp = Blueprint("platform", __name__) @@ -374,10 +378,56 @@ def adapter_instance() -> Any: f: Fernet = Fernet(Env.ENCRYPTION_KEY.encode("utf-8")) - data_dict["adapter_metadata"] = json.loads( + adapter_metadata = json.loads( f.decrypt(bytes(data_dict.pop("adapter_metadata_b")).decode("utf-8")) ) + if is_openai_oauth_adapter(data_dict.get("adapter_id")): + try: + refreshed_metadata = refresh_openai_oauth_metadata(adapter_metadata) + except OpenAIOAuthRefreshError as exc: + app.logger.warning( + "OpenAI OAuth credentials could not be refreshed for adapter " + "%s: %s", + adapter_instance_id, + exc, + ) + raise APIError( + message=( + "OpenAI OAuth credentials have expired. Reauthenticate " + "this adapter in the platform settings." + ), + code=401, + ) from exc + + if refreshed_metadata != adapter_metadata: + refreshed_ciphertext = f.encrypt( + json.dumps(refreshed_metadata).encode("utf-8") + ) + try: + update_query = ( + f'UPDATE "{Env.DB_SCHEMA}".{DBTable.ADAPTER_INSTANCE} ' + "SET adapter_metadata_b=%s " + "WHERE id=%s AND organization_id=%s" + ) + with safe_cursor( + update_query, + (refreshed_ciphertext, adapter_instance_id, organization_uid), + ): + pass + except Exception: + # The current workflow can use the refreshed copy. A later + # request may refresh again if the durable write lost a race. + app.logger.warning( + "Could not persist refreshed OpenAI OAuth credentials for " + "adapter %s", + adapter_instance_id, + exc_info=True, + ) + adapter_metadata = refreshed_metadata + + data_dict["adapter_metadata"] = adapter_metadata + return jsonify(data_dict) except InvalidToken: msg = ( diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/adapterkit.py b/unstract/sdk1/src/unstract/sdk1/adapters/adapterkit.py index 330ad2db92..daf88f6681 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/adapterkit.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/adapterkit.py @@ -5,6 +5,7 @@ from typing import Any from singleton_decorator import singleton + from unstract.sdk1.adapters import AdapterDict from unstract.sdk1.adapters.base import Adapter from unstract.sdk1.adapters.constants import Common @@ -71,6 +72,8 @@ def get_adapters_list(self) -> list[dict[str, "Any"]]: json_schema = m.get_json_schema() desc = m.get_description() icon = m.get_icon() + get_auth_metadata = getattr(m, "get_auth_metadata", None) + auth_metadata = get_auth_metadata() if get_auth_metadata else {} adapters.append( { "id": _id, @@ -81,6 +84,7 @@ def get_adapters_list(self) -> list[dict[str, "Any"]]: "adapter_type": adapter_type, "json_schema": json_schema, "doc_url": m.get_doc_url(), + **auth_metadata, } ) return adapters diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py index 716587a23e..d2174ef16b 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/base1.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/base1.py @@ -12,6 +12,7 @@ from typing import Any from pydantic import BaseModel, Field, model_validator + from unstract.sdk1.adapters.constants import AdapterDocs, Common from unstract.sdk1.adapters.enums import AdapterTypes @@ -255,6 +256,11 @@ def get_icon() -> str: def get_doc_url(cls) -> str: return AdapterDocs.type_index_url(cls.get_adapter_type()) + @classmethod + def get_auth_metadata(cls) -> dict[str, "Any"]: + """Return optional authentication metadata for the UI/API catalog.""" + return {} + @classmethod def get_json_schema(cls) -> str: schema_path = ( diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/__init__.py b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/__init__.py index a3a03c7da3..f3424265f1 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/__init__.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/__init__.py @@ -11,6 +11,7 @@ from unstract.sdk1.adapters.llm1.ollama import OllamaLLMAdapter from unstract.sdk1.adapters.llm1.openai import OpenAILLMAdapter from unstract.sdk1.adapters.llm1.openai_compatible import OpenAICompatibleLLMAdapter +from unstract.sdk1.adapters.llm1.openai_oauth import OpenAIOAuthLLMAdapter from unstract.sdk1.adapters.llm1.openrouter import OpenRouterLLMAdapter from unstract.sdk1.adapters.llm1.vertexai import VertexAILLMAdapter @@ -29,6 +30,7 @@ "OllamaLLMAdapter", "OpenAILLMAdapter", "OpenAICompatibleLLMAdapter", + "OpenAIOAuthLLMAdapter", "OpenRouterLLMAdapter", "VertexAILLMAdapter", ] diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/openai_oauth.py b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/openai_oauth.py new file mode 100644 index 0000000000..c9aea521f6 --- /dev/null +++ b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/openai_oauth.py @@ -0,0 +1,106 @@ +"""OpenAI ChatGPT OAuth LLM adapter. + +This adapter deliberately has a separate id and provider from the API-key +OpenAI adapter. OAuth credentials are injected by the authenticated web/API +flow and are never part of the configuration form. +""" + +from typing import Any + +from unstract.sdk1.adapters.base1 import BaseAdapter, BaseChatCompletionParameters +from unstract.sdk1.adapters.enums import AdapterTypes +from unstract.sdk1.auth.openai_oauth import ( + OPENAI_OAUTH_CHATGPT_API_BASE, + OPENAI_OAUTH_PROVIDER, +) + + +class OpenAIOAuthLLMParameters(BaseChatCompletionParameters): + """Validated per-adapter credentials for the ChatGPT Responses API.""" + + oauth_access_token: str + oauth_refresh_token: str + oauth_id_token: str | None = None + oauth_account_id: str + oauth_account_email: str | None = None + oauth_expires_at: float | int | None = None + # This is fixed by the adapter. It is not exposed in the JSON schema. + api_base: str = OPENAI_OAUTH_CHATGPT_API_BASE + reasoning_effort: str | None = None + + @staticmethod + def validate(adapter_metadata: dict[str, Any]) -> dict[str, Any]: + metadata = dict(adapter_metadata) + model = str(metadata.get("model", "")).strip() + for prefix in ("openai/", "chatgpt/"): + if model.startswith(prefix): + model = model[len(prefix) :] + break + for field in ("oauth_access_token", "oauth_refresh_token", "oauth_account_id"): + if not isinstance(metadata.get(field), str) or not metadata[field].strip(): + raise ValueError(f"Missing required OpenAI OAuth field: {field}") + if not model: + raise ValueError("Missing required OpenAI OAuth field: model") + metadata["model"] = model + metadata["api_base"] = OPENAI_OAUTH_CHATGPT_API_BASE + + validated = OpenAIOAuthLLMParameters(**metadata).model_dump() + # The LLM wrapper uses this for usage/cost bookkeeping; it is not sent + # as a provider parameter. + validated["cost_model"] = model + return validated + + @staticmethod + def validate_model(adapter_metadata: dict[str, Any]) -> str: + model = str(adapter_metadata.get("model", "")).strip() + return model.removeprefix("openai/").removeprefix("chatgpt/") + + +class OpenAIOAuthLLMAdapter(OpenAIOAuthLLMParameters, BaseAdapter): + """OpenAI LLM backed by a user-authorized ChatGPT account.""" + + @staticmethod + def get_id() -> str: + return "openai-oauth|a5ce9b7d-5f8a-4c6d-8a65-6e1f626b2b6e" + + @staticmethod + def get_metadata() -> dict[str, Any]: + return { + "name": "OpenAI (OAuth)", + "version": "1.0.0", + "adapter": OpenAIOAuthLLMAdapter, + "description": "OpenAI LLM adapter authenticated with ChatGPT OAuth", + "is_active": True, + } + + @classmethod + def get_auth_metadata(cls) -> dict[str, Any]: + return { + "oauth": True, + "oauth_provider": "openai", + "python_social_auth_backend": OPENAI_OAUTH_PROVIDER, + } + + @staticmethod + def get_name() -> str: + return "OpenAI (OAuth)" + + @staticmethod + def get_description() -> str: + return "OpenAI LLM adapter authenticated with ChatGPT OAuth" + + @staticmethod + def get_provider() -> str: + return OPENAI_OAUTH_PROVIDER + + @staticmethod + def get_icon() -> str: + return "/icons/adapter-icons/OpenAI.png" + + @staticmethod + def get_doc_url() -> str: + return "https://developers.openai.com/codex/auth/" + + @staticmethod + def get_adapter_type() -> AdapterTypes: + return AdapterTypes.LLM diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json new file mode 100644 index 0000000000..beebec0e43 --- /dev/null +++ b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json @@ -0,0 +1,93 @@ +{ + "title": "OpenAI OAuth LLM", + "type": "object", + "required": [ + "adapter_name", + "model" + ], + "properties": { + "adapter_name": { + "type": "string", + "title": "Name", + "default": "", + "description": "Provide a unique name for this adapter instance. Example: openai-chatgpt-personal" + }, + "model": { + "type": "string", + "title": "Model", + "default": "gpt-5-codex", + "description": "The Codex-capable model to use through your ChatGPT account." + }, + "max_tokens": { + "type": "number", + "minimum": 1, + "multipleOf": 1, + "title": "Maximum Output Tokens", + "default": 4096, + "description": "Maximum number of output tokens to request." + }, + "max_retries": { + "type": "number", + "minimum": 0, + "multipleOf": 1, + "title": "Max Retries", + "default": 5, + "description": "The maximum number of times to retry a transient request." + }, + "timeout": { + "type": "number", + "minimum": 0, + "multipleOf": 1, + "title": "Timeout", + "default": 900, + "description": "Timeout in seconds." + }, + "enable_reasoning": { + "type": "boolean", + "title": "Enable Reasoning", + "default": false, + "description": "Allow the model to apply extra reasoning for complex tasks." + } + }, + "allOf": [ + { + "if": { + "properties": { + "enable_reasoning": { + "const": true + } + } + }, + "then": { + "properties": { + "reasoning_effort": { + "type": "string", + "enum": [ + "low", + "medium", + "high" + ], + "default": "medium", + "title": "Reasoning Effort", + "description": "Sets the reasoning strength when reasoning is enabled." + } + }, + "required": [ + "reasoning_effort" + ] + } + }, + { + "if": { + "properties": { + "enable_reasoning": { + "const": false + } + } + }, + "then": { + "properties": {} + } + } + ] +} diff --git a/unstract/sdk1/src/unstract/sdk1/auth/__init__.py b/unstract/sdk1/src/unstract/sdk1/auth/__init__.py new file mode 100644 index 0000000000..14d4447960 --- /dev/null +++ b/unstract/sdk1/src/unstract/sdk1/auth/__init__.py @@ -0,0 +1,25 @@ +"""Authentication helpers shared by SDK adapters and runtime services.""" + +from unstract.sdk1.auth.openai_oauth import ( + OPENAI_OAUTH_ADAPTER_PREFIX, + OPENAI_OAUTH_CHATGPT_API_BASE, + OPENAI_OAUTH_PROVIDER, + OpenAIOAuthError, + OpenAIOAuthRefreshError, + extract_account_id, + extract_email, + is_openai_oauth_adapter, + refresh_openai_oauth_metadata, +) + +__all__ = [ + "OPENAI_OAUTH_ADAPTER_PREFIX", + "OPENAI_OAUTH_CHATGPT_API_BASE", + "OPENAI_OAUTH_PROVIDER", + "OpenAIOAuthError", + "OpenAIOAuthRefreshError", + "extract_account_id", + "extract_email", + "is_openai_oauth_adapter", + "refresh_openai_oauth_metadata", +] diff --git a/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py b/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py new file mode 100644 index 0000000000..b2ea9deb13 --- /dev/null +++ b/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py @@ -0,0 +1,236 @@ +"""Small, provider-specific helpers for OpenAI ChatGPT OAuth credentials. + +The browser/device login belongs to the web application. This module stays +free of Django and Flask so the platform service can refresh credentials for an +individual adapter instance without keeping a process-global account. +""" + +from __future__ import annotations + +import base64 +import json +import os +import time +from collections.abc import Mapping +from typing import Any + +import httpx + +OPENAI_OAUTH_PROVIDER = "openai-oauth" +OPENAI_OAUTH_ADAPTER_PREFIX = f"{OPENAI_OAUTH_PROVIDER}|" +OPENAI_OAUTH_AUTH_BASE = os.environ.get( + "OPENAI_OAUTH_AUTH_BASE", "https://auth.openai.com" +).rstrip("/") +OPENAI_OAUTH_CLIENT_ID = os.environ.get( + "OPENAI_OAUTH_CLIENT_ID", "app_EMoamEEZ73f0CkXaXp7hrann" +) +OPENAI_OAUTH_CHATGPT_API_BASE = "https://chatgpt.com/backend-api/codex" +OPENAI_OAUTH_DEVICE_USERCODE_URL = ( + f"{OPENAI_OAUTH_AUTH_BASE}/api/accounts/deviceauth/usercode" +) +OPENAI_OAUTH_DEVICE_TOKEN_URL = ( + f"{OPENAI_OAUTH_AUTH_BASE}/api/accounts/deviceauth/token" +) +OPENAI_OAUTH_TOKEN_URL = f"{OPENAI_OAUTH_AUTH_BASE}/oauth/token" +OPENAI_OAUTH_DEVICE_VERIFICATION_URL = f"{OPENAI_OAUTH_AUTH_BASE}/codex/device" +OPENAI_OAUTH_DEVICE_REDIRECT_URI = f"{OPENAI_OAUTH_AUTH_BASE}/deviceauth/callback" + +# OAuth metadata names are intentionally namespaced. It prevents a normal +# OpenAI API-key adapter from accidentally being treated as a ChatGPT OAuth +# adapter, and makes redaction easy at the API boundary. +OPENAI_OAUTH_SECRET_FIELDS = frozenset( + { + "oauth_access_token", + "oauth_refresh_token", + "oauth_id_token", + } +) +OPENAI_OAUTH_PRIVATE_FIELDS = frozenset( + { + *OPENAI_OAUTH_SECRET_FIELDS, + "oauth_account_id", + "oauth_account_email", + "oauth_expires_at", + } +) + + +class OpenAIOAuthError(RuntimeError): + """Base error for OpenAI OAuth credential handling.""" + + +class OpenAIOAuthRefreshError(OpenAIOAuthError): + """Raised when an expired access token cannot be refreshed.""" + + +def is_openai_oauth_adapter(adapter_id: str | None) -> bool: + """Return whether ``adapter_id`` identifies the OAuth-backed adapter.""" + return bool(adapter_id and adapter_id.startswith(OPENAI_OAUTH_ADAPTER_PREFIX)) + + +def _decode_jwt_claims(token: str | None) -> dict[str, Any]: + """Decode the unverified JWT payload for routing metadata only. + + The authorization server has already issued the token and the provider + validates it on every model request. We only use claims to find the + account label and expiry; this function is not an authentication check. + """ + if not token or not isinstance(token, str): + return {} + try: + parts = token.split(".") + if len(parts) != 3: + return {} + payload = parts[1] + "=" * (-len(parts[1]) % 4) + decoded = base64.urlsafe_b64decode(payload.encode("ascii")) + claims = json.loads(decoded.decode("utf-8")) + return claims if isinstance(claims, dict) else {} + except (ValueError, UnicodeDecodeError, json.JSONDecodeError): + return {} + + +def _nested_claim(claims: Mapping[str, Any], key: str) -> object | None: + """Read a claim from the common OpenAI auth namespace as a fallback.""" + for namespace in ( + "https://api.openai.com/auth", + "https://auth.openai.com/auth", + "auth", + ): + auth_claims = claims.get(namespace) + if isinstance(auth_claims, Mapping) and auth_claims.get(key): + return auth_claims[key] + return None + + +def extract_account_id(*tokens: str | None) -> str | None: + """Extract the ChatGPT account/workspace id from issued JWT claims.""" + for token in tokens: + claims = _decode_jwt_claims(token) + account_id = claims.get("chatgpt_account_id") or claims.get("account_id") + if not account_id: + account_id = _nested_claim(claims, "chatgpt_account_id") + if isinstance(account_id, str) and account_id.strip(): + return account_id.strip() + return None + + +def extract_email(*tokens: str | None) -> str | None: + """Extract an optional account email used only as a friendly label.""" + for token in tokens: + claims = _decode_jwt_claims(token) + email = claims.get("email") + if not email and isinstance(claims.get("profile"), Mapping): + email = claims["profile"].get("email") + if not email: + email = _nested_claim(claims, "email") + if isinstance(email, str) and email.strip(): + return email.strip() + return None + + +def _extract_expiry(*tokens: str | None) -> float | None: + for token in tokens: + expiry = _decode_jwt_claims(token).get("exp") + try: + if expiry is not None: + return float(expiry) + except (TypeError, ValueError): + continue + return None + + +def _response_error(response: httpx.Response, operation: str) -> OpenAIOAuthRefreshError: + """Build a safe refresh error without including response/token contents.""" + return OpenAIOAuthRefreshError( + f"OpenAI OAuth {operation} failed with status {response.status_code}" + ) + + +def refresh_openai_oauth_metadata( + metadata: Mapping[str, Any], + *, + force: bool = False, + now: float | None = None, +) -> dict[str, Any]: + """Refresh one adapter's OAuth metadata when its access token is expiring. + + The returned dictionary is a copy. The old refresh token is retained when + the token endpoint does not rotate it, which is permitted by OAuth 2.0. + Callers decide where the refreshed copy is persisted; no account is stored + in module-level state. + """ + refreshed = dict(metadata) + access_token = refreshed.get("oauth_access_token") + refresh_token = refreshed.get("oauth_refresh_token") + account_id = refreshed.get("oauth_account_id") or extract_account_id( + refreshed.get("oauth_id_token"), access_token + ) + if not access_token or not refresh_token or not account_id: + raise OpenAIOAuthRefreshError( + "OpenAI OAuth metadata is missing the credentials required for refresh" + ) + + current_time = time.time() if now is None else now + try: + expires_at = float(refreshed.get("oauth_expires_at")) + except (TypeError, ValueError): + expires_at = _extract_expiry(access_token, refreshed.get("oauth_id_token")) + + # Refresh slightly before expiry so a long-running request does not start + # with a token that expires while it is in flight. + if not force and expires_at is not None and expires_at > current_time + 60: + return refreshed + + try: + response = httpx.post( + OPENAI_OAUTH_TOKEN_URL, + json={ + "client_id": OPENAI_OAUTH_CLIENT_ID, + "grant_type": "refresh_token", + "refresh_token": refresh_token, + }, + timeout=15.0, + ) + except httpx.HTTPError as exc: + raise OpenAIOAuthRefreshError( + "OpenAI OAuth token refresh could not reach the authorization server" + ) from exc + + if not 200 <= response.status_code < 300: + raise _response_error(response, "token refresh") + + try: + token_data = response.json() + except (ValueError, json.JSONDecodeError) as exc: + raise OpenAIOAuthRefreshError( + "OpenAI OAuth token refresh returned an invalid response" + ) from exc + + new_access_token = token_data.get("access_token") + if not isinstance(new_access_token, str) or not new_access_token: + raise OpenAIOAuthRefreshError( + "OpenAI OAuth token refresh returned no access token" + ) + + new_id_token = token_data.get("id_token") or refreshed.get("oauth_id_token") + new_refresh_token = token_data.get("refresh_token") or refresh_token + expires_in = token_data.get("expires_in") + try: + new_expires_at = current_time + float(expires_in) + except (TypeError, ValueError): + new_expires_at = _extract_expiry(new_access_token, new_id_token) + + refreshed.update( + { + "oauth_access_token": new_access_token, + "oauth_refresh_token": new_refresh_token, + "oauth_id_token": new_id_token, + "oauth_account_id": extract_account_id(new_id_token, new_access_token) + or account_id, + "oauth_account_email": extract_email(new_id_token, new_access_token) + or refreshed.get("oauth_account_email"), + "oauth_expires_at": new_expires_at, + "oauth_authenticated": True, + } + ) + return refreshed diff --git a/unstract/sdk1/src/unstract/sdk1/llm.py b/unstract/sdk1/src/unstract/sdk1/llm.py index e780685180..e3f1aa9ba4 100644 --- a/unstract/sdk1/src/unstract/sdk1/llm.py +++ b/unstract/sdk1/src/unstract/sdk1/llm.py @@ -1,6 +1,8 @@ +import json import logging import os import re +import uuid from collections.abc import Callable, Generator, Mapping, Sequence from dataclasses import dataclass, field from enum import Enum @@ -11,8 +13,13 @@ # from litellm import get_supported_openai_params from litellm import get_max_tokens + from unstract.sdk1.adapters.constants import Common from unstract.sdk1.adapters.llm1 import adapters +from unstract.sdk1.auth.openai_oauth import ( + OPENAI_OAUTH_CHATGPT_API_BASE, + is_openai_oauth_adapter, +) from unstract.sdk1.constants import Common as SdkCommon from unstract.sdk1.constants import ToolEnv from unstract.sdk1.exceptions import LLMError, SdkError, strip_litellm_prefix @@ -496,6 +503,272 @@ def _build_messages( {"role": "user", "content": user_content}, ] + def _uses_openai_oauth(self) -> bool: + """Whether this LLM uses the per-account ChatGPT OAuth adapter.""" + return is_openai_oauth_adapter(self._adapter_id) + + @staticmethod + def _response_value( + response: object, key: str, default: object | None = None + ) -> object | None: + """Read a field from a LiteLLM dict, model, or response event.""" + if isinstance(response, Mapping): + return response.get(key, default) + try: + value = response[key] # type: ignore[index] + except (KeyError, TypeError, AttributeError, IndexError): + value = getattr(response, key, default) + return value if value is not None else default + + @staticmethod + def _responses_block(block: object) -> dict[str, object]: + """Convert one chat content block to a Responses input block.""" + if not isinstance(block, Mapping): + return {"type": "input_text", "text": str(block)} + block_type = block.get("type") + if block_type in ("text", "input_text"): + return {"type": "input_text", "text": block.get("text", "")} + if block_type == "image_url": + image_url = block.get("image_url") + detail = None + if isinstance(image_url, Mapping): + detail = image_url.get("detail") + image_url = image_url.get("url") + image_block: dict[str, object] = { + "type": "input_image", + "image_url": image_url, + } + if detail: + image_block["detail"] = detail + return image_block + if isinstance(block_type, str) and block_type.startswith("input_"): + # Preserve already-converted Responses blocks for callers that use + # complete_vision() directly. + return dict(block) + return {"type": "input_text", "text": json.dumps(dict(block))} + + @classmethod + def _responses_content(cls, content: object) -> list[dict[str, object]]: + """Convert OpenAI chat content blocks to Responses input blocks.""" + if isinstance(content, str): + return [{"type": "input_text", "text": content}] + if not isinstance(content, list): + if content is None: + return [] + return [{"type": "input_text", "text": str(content)}] + return [cls._responses_block(block) for block in content] + + @classmethod + def _responses_input( + cls, messages: list[dict[str, object]] + ) -> tuple[str, list[dict[str, object]]]: + instructions: list[str] = [] + response_input: list[dict[str, object]] = [] + for message in messages: + role = str(message.get("role", "user")) + content = message.get("content") + if role == "system": + text_blocks = cls._responses_content(content) + instructions.extend( + str(block.get("text", "")) + for block in text_blocks + if block.get("type") == "input_text" + ) + continue + if role not in {"user", "assistant"}: + role = "user" + response_input.append( + { + "type": "message", + "role": role, + "content": cls._responses_content(content), + } + ) + return "\n".join(text for text in instructions if text), response_input + + @staticmethod + def _responses_tools(tools: object) -> object: + """Translate chat function tools to the Responses tool shape.""" + if not isinstance(tools, list): + return tools + converted: list[object] = [] + for tool in tools: + if not isinstance(tool, Mapping) or tool.get("type") != "function": + converted.append(tool) + continue + function = tool.get("function") + if not isinstance(function, Mapping): + converted.append(tool) + continue + converted.append( + { + "type": "function", + "name": function.get("name"), + "description": function.get("description"), + "parameters": function.get("parameters"), + } + ) + return converted + + def _build_openai_oauth_responses_kwargs( + self, + messages: list[dict[str, object]], + completion_kwargs: dict[str, object], + *, + stream: bool, + ) -> dict[str, object]: + """Build a Responses API call with credentials for one adapter only.""" + values = dict(completion_kwargs) + access_token = str(values.pop("oauth_access_token")) + account_id = str(values.pop("oauth_account_id")) + values.pop("oauth_refresh_token", None) + values.pop("oauth_id_token", None) + values.pop("oauth_account_email", None) + values.pop("oauth_expires_at", None) + values.pop("oauth_authenticated", None) + + model = str(values.pop("model")) + api_base = str(values.pop("api_base", OPENAI_OAUTH_CHATGPT_API_BASE)) + max_tokens = values.pop("max_tokens", None) + values.pop("temperature", None) + values.pop("n", None) + values.pop("api_version", None) + values.pop("enable_reasoning", None) + values.pop("max_retries", None) + values.pop("num_retries", None) + values.pop("cost_model", None) + values.pop("context_window", None) + + if max_tokens is not None: + values["max_output_tokens"] = max_tokens + reasoning_effort = values.pop("reasoning_effort", None) + # The ChatGPT Codex endpoint expects this field on Responses requests + # so encrypted reasoning can be returned and reused by the account. + values["include"] = ["reasoning.encrypted_content"] + if reasoning_effort: + values["reasoning"] = {"effort": reasoning_effort} + if "tools" in values: + values["tools"] = self._responses_tools(values["tools"]) + + instructions, response_input = self._responses_input(messages) + values.update( + { + "model": model, + "input": response_input, + "custom_llm_provider": "openai", + # LiteLLM's OpenAI handler uses this key to select its client; + # the explicit Authorization header below is account-specific. + "api_key": access_token, + "api_base": api_base, + "stream": stream, + "store": False, + } + ) + if instructions: + values["instructions"] = instructions + + headers = dict(values.pop("extra_headers", {}) or {}) + headers.update( + { + "Authorization": f"Bearer {access_token}", + "ChatGPT-Account-Id": account_id, + "Content-Type": "application/json", + "Accept": "text/event-stream" if stream else "application/json", + "originator": "unstract", + "session-id": str(uuid.uuid4()), + } + ) + values["extra_headers"] = headers + return values + + @classmethod + def _responses_output_text(cls, response: object) -> str | None: + output_text = cls._response_value(response, "output_text") + if isinstance(output_text, str): + return output_text + output = cls._response_value(response, "output") + if not isinstance(output, list): + return None + text_parts: list[str] = [] + for item in output: + if cls._response_value(item, "type") != "message": + continue + content = cls._response_value(item, "content") + if not isinstance(content, list): + continue + for block in content: + if cls._response_value(block, "type") in { + "output_text", + "text", + }: + text = cls._response_value(block, "text") + if isinstance(text, str): + text_parts.append(text) + return "".join(text_parts) or None + + @classmethod + def _responses_usage(cls, usage: object) -> dict[str, int]: + if usage is None: + return {} + + def integer(key: str, fallback: str) -> int: + value = cls._response_value(usage, key) + if value is None: + value = cls._response_value(usage, fallback, 0) + try: + return int(value or 0) + except (TypeError, ValueError): + return 0 + + return { + "prompt_tokens": integer("input_tokens", "prompt_tokens"), + "completion_tokens": integer("output_tokens", "completion_tokens"), + "total_tokens": integer("total_tokens", "total_tokens"), + } + + def _stream_openai_oauth( + self, + messages: list[dict[str, object]], + completion_kwargs: dict[str, object], + callback_manager: object | None, + max_retries: int, + ) -> Generator[LLMResponseCompat, None, None]: + """Yield text events from one account's Responses API stream.""" + response_kwargs = self._build_openai_oauth_responses_kwargs( + messages, completion_kwargs, stream=True + ) + for event in iter_with_retry( + lambda: litellm.responses(**response_kwargs), + max_retries=max_retries, + retry_predicate=is_retryable_litellm_error, + description=self._get_adapter_info(), + ): + event_type = self._response_value(event, "type") + if event_type == "response.completed": + completed_response = self._response_value(event, "response") + usage = self._responses_usage( + self._response_value(completed_response, "usage") + ) + if usage: + self._record_usage( + self._cost_model or self.kwargs["model"], + messages, + usage, + "stream_complete", + response=completed_response, + ) + continue + if event_type != "response.output_text.delta": + continue + text = self._response_value(event, "delta", "") + if not isinstance(text, str) or not text: + continue + if callback_manager and hasattr(callback_manager, "on_stream"): + callback_manager.on_stream(text) + stream_response = LLMResponseCompat(text) + stream_response.delta = text + yield stream_response + @capture_metrics def complete( self, @@ -541,20 +814,37 @@ def complete( max_retries = pop_litellm_retry_kwargs( completion_kwargs, self._get_adapter_info() ) - response: dict[str, object] = call_with_retry( - lambda: litellm.completion(messages=messages, **completion_kwargs), - max_retries=max_retries, - retry_predicate=is_retryable_litellm_error, - description=self._get_adapter_info(), - ) - - response_text = response["choices"][0]["message"]["content"] - finish_reason = response["choices"][0].get("finish_reason") + if self._uses_openai_oauth(): + response = call_with_retry( + lambda: litellm.responses( + **self._build_openai_oauth_responses_kwargs( + messages, completion_kwargs, stream=False + ) + ), + max_retries=max_retries, + retry_predicate=is_retryable_litellm_error, + description=self._get_adapter_info(), + ) + response_text = self._responses_output_text(response) + finish_reason = None + usage = self._responses_usage( + self._response_value(response, "usage") + ) + else: + response = call_with_retry( + lambda: litellm.completion(messages=messages, **completion_kwargs), + max_retries=max_retries, + retry_predicate=is_retryable_litellm_error, + description=self._get_adapter_info(), + ) + response_text = response["choices"][0]["message"]["content"] + finish_reason = response["choices"][0].get("finish_reason") + usage = response.get("usage") self._record_usage( self._cost_model or self.kwargs["model"], messages, - response.get("usage"), + usage, "complete", response=response, ) @@ -658,18 +948,38 @@ def complete_vision( completion_kwargs.pop("enable_prompt_caching", None) completion_kwargs.pop("context_window", None) - response: dict[str, object] = litellm.completion( - messages=messages, - **completion_kwargs, - ) - - response_text = response["choices"][0]["message"]["content"] - finish_reason = response["choices"][0].get("finish_reason") + if self._uses_openai_oauth(): + max_retries = pop_litellm_retry_kwargs( + completion_kwargs, self._get_adapter_info() + ) + response = call_with_retry( + lambda: litellm.responses( + **self._build_openai_oauth_responses_kwargs( + messages, completion_kwargs, stream=False + ) + ), + max_retries=max_retries, + retry_predicate=is_retryable_litellm_error, + description=self._get_adapter_info(), + ) + response_text = self._responses_output_text(response) + finish_reason = None + usage = self._responses_usage( + self._response_value(response, "usage") + ) + else: + response = litellm.completion( + messages=messages, + **completion_kwargs, + ) + response_text = response["choices"][0]["message"]["content"] + finish_reason = response["choices"][0].get("finish_reason") + usage = response.get("usage") self._record_usage( self._cost_model or self.kwargs["model"], messages, - response.get("usage"), + usage, "complete_vision", response=response, ) @@ -732,32 +1042,37 @@ def stream_complete( completion_kwargs, self._get_adapter_info() ) has_yielded_content = False - for chunk in iter_with_retry( - lambda: litellm.completion( - messages=messages, - stream=True, - stream_options={"include_usage": True}, - **completion_kwargs, - ), - max_retries=max_retries, - retry_predicate=is_retryable_litellm_error, - description=self._get_adapter_info(), - ): - if chunk.get("usage"): - self._record_usage( - self._cost_model or self.kwargs["model"], - messages, - chunk.get("usage"), - "stream_complete", - response=chunk, - ) - - response = self._process_stream_chunk( - chunk, callback_manager, has_yielded_content + if self._uses_openai_oauth(): + yield from self._stream_openai_oauth( + messages, completion_kwargs, callback_manager, max_retries ) - if response is not None: - has_yielded_content = True - yield response + else: + for chunk in iter_with_retry( + lambda: litellm.completion( + messages=messages, + stream=True, + stream_options={"include_usage": True}, + **completion_kwargs, + ), + max_retries=max_retries, + retry_predicate=is_retryable_litellm_error, + description=self._get_adapter_info(), + ): + if chunk.get("usage"): + self._record_usage( + self._cost_model or self.kwargs["model"], + messages, + chunk.get("usage"), + "stream_complete", + response=chunk, + ) + + response = self._process_stream_chunk( + chunk, callback_manager, has_yielded_content + ) + if response is not None: + has_yielded_content = True + yield response except LLMError: # Already wrapped LLMError, re-raise as is @@ -812,19 +1127,37 @@ async def acomplete( max_retries = pop_litellm_retry_kwargs( completion_kwargs, self._get_adapter_info() ) - response = await acall_with_retry( - lambda: litellm.acompletion(messages=messages, **completion_kwargs), - max_retries=max_retries, - retry_predicate=is_retryable_litellm_error, - description=self._get_adapter_info(), - ) - response_text = response["choices"][0]["message"]["content"] - finish_reason = response["choices"][0].get("finish_reason") + if self._uses_openai_oauth(): + response = await acall_with_retry( + lambda: litellm.aresponses( + **self._build_openai_oauth_responses_kwargs( + messages, completion_kwargs, stream=False + ) + ), + max_retries=max_retries, + retry_predicate=is_retryable_litellm_error, + description=self._get_adapter_info(), + ) + response_text = self._responses_output_text(response) + finish_reason = None + usage = self._responses_usage( + self._response_value(response, "usage") + ) + else: + response = await acall_with_retry( + lambda: litellm.acompletion(messages=messages, **completion_kwargs), + max_retries=max_retries, + retry_predicate=is_retryable_litellm_error, + description=self._get_adapter_info(), + ) + response_text = response["choices"][0]["message"]["content"] + finish_reason = response["choices"][0].get("finish_reason") + usage = response.get("usage") self._record_usage( self._cost_model or self.kwargs["model"], messages, - response.get("usage"), + usage, "acomplete", response=response, ) diff --git a/unstract/sdk1/tests/test_openai_oauth.py b/unstract/sdk1/tests/test_openai_oauth.py new file mode 100644 index 0000000000..94a4953a21 --- /dev/null +++ b/unstract/sdk1/tests/test_openai_oauth.py @@ -0,0 +1,209 @@ +"""Tests for the per-account OpenAI ChatGPT OAuth adapter.""" + +from __future__ import annotations + +import asyncio +import base64 +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from unstract.sdk1.adapters.constants import Common +from unstract.sdk1.adapters.llm1 import adapters +from unstract.sdk1.adapters.llm1.openai_oauth import ( + OpenAIOAuthLLMAdapter, + OpenAIOAuthLLMParameters, +) +from unstract.sdk1.auth.openai_oauth import ( + OPENAI_OAUTH_CHATGPT_API_BASE, + OPENAI_OAUTH_TOKEN_URL, + extract_account_id, + extract_email, + refresh_openai_oauth_metadata, +) +from unstract.sdk1.llm import LLM + + +def _jwt(payload: dict[str, object]) -> str: + encoded = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode().rstrip("=") + return f"header.{encoded}.signature" + + +def _metadata(access_token: str, account_id: str) -> dict[str, object]: + return { + "model": "gpt-5-codex", + "oauth_access_token": access_token, + "oauth_refresh_token": f"refresh-{account_id}", + "oauth_id_token": "id-token", + "oauth_account_id": account_id, + "oauth_account_email": f"{account_id}@example.test", + "oauth_expires_at": 4_000_000_000, + } + + +def test_openai_oauth_adapter_is_registered_with_auth_metadata() -> None: + adapter_id = OpenAIOAuthLLMAdapter.get_id() + assert adapters[adapter_id][Common.MODULE] is OpenAIOAuthLLMAdapter + + registry_metadata = adapters[adapter_id][Common.METADATA] + assert registry_metadata[Common.ADAPTER] is OpenAIOAuthLLMAdapter + assert OpenAIOAuthLLMAdapter.get_auth_metadata() == { + "oauth": True, + "oauth_provider": "openai", + "python_social_auth_backend": "openai-oauth", + } + schema = json.loads(OpenAIOAuthLLMAdapter.get_json_schema()) + assert schema["properties"]["model"]["default"] == "gpt-5-codex" + + +def test_openai_oauth_parameters_normalize_model_and_fix_endpoint() -> None: + metadata = _metadata("access", "account") + metadata["model"] = "openai/gpt-5-codex" + + validated = OpenAIOAuthLLMParameters.validate(metadata) + + assert validated["model"] == "gpt-5-codex" + assert validated["api_base"] == OPENAI_OAUTH_CHATGPT_API_BASE + assert validated["cost_model"] == "gpt-5-codex" + assert metadata["model"] == "openai/gpt-5-codex" + + +@pytest.mark.parametrize( + "field", + ["oauth_access_token", "oauth_refresh_token", "oauth_account_id"], +) +def test_openai_oauth_parameters_require_account_credentials(field: str) -> None: + metadata = _metadata("access", "account") + metadata[field] = " " + + with pytest.raises(ValueError, match=field): + OpenAIOAuthLLMParameters.validate(metadata) + + +def test_openai_oauth_claim_helpers_read_account_and_email() -> None: + token = _jwt( + { + "chatgpt_account_id": "workspace-123", + "email": "person@example.test", + } + ) + + assert extract_account_id(token) == "workspace-123" + assert extract_email(token) == "person@example.test" + + +def test_refresh_openai_oauth_metadata_is_scoped_to_one_account() -> None: + response = MagicMock(status_code=200) + response.json.return_value = {"access_token": "new-access", "expires_in": 3600} + metadata = _metadata("old-access", "workspace-123") + metadata["oauth_expires_at"] = 0 + + with patch( + "unstract.sdk1.auth.openai_oauth.httpx.post", return_value=response + ) as post: + refreshed = refresh_openai_oauth_metadata(metadata, now=100, force=True) + + assert refreshed["oauth_access_token"] == "new-access" + assert refreshed["oauth_refresh_token"] == "refresh-workspace-123" + assert refreshed["oauth_account_id"] == "workspace-123" + post.assert_called_once_with( + OPENAI_OAUTH_TOKEN_URL, + json={ + "client_id": "app_EMoamEEZ73f0CkXaXp7hrann", + "grant_type": "refresh_token", + "refresh_token": "refresh-workspace-123", + }, + timeout=15.0, + ) + + +def test_llm_sends_each_account_token_and_workspace_header() -> None: + first = LLM( + adapter_id=OpenAIOAuthLLMAdapter.get_id(), + adapter_metadata=_metadata("token-a", "account-a"), + ) + second = LLM( + adapter_id=OpenAIOAuthLLMAdapter.get_id(), + adapter_metadata=_metadata("token-b", "account-b"), + ) + + with ( + patch( + "unstract.sdk1.llm.litellm.responses", + side_effect=[ + {"output_text": "first", "usage": {}}, + {"output_text": "second", "usage": {}}, + ], + ) as responses, + patch.object(first, "_record_usage"), + patch.object(second, "_record_usage"), + ): + assert first.complete("hello")["response"].text == "first" + assert second.complete("hello")["response"].text == "second" + + first_request = responses.call_args_list[0].kwargs + second_request = responses.call_args_list[1].kwargs + assert first_request["api_key"] == "token-a" + assert second_request["api_key"] == "token-b" + assert first_request["extra_headers"]["ChatGPT-Account-Id"] == "account-a" + assert second_request["extra_headers"]["ChatGPT-Account-Id"] == "account-b" + assert first_request["extra_headers"]["Authorization"] == "Bearer token-a" + assert second_request["extra_headers"]["Authorization"] == "Bearer token-b" + assert first_request["extra_headers"]["session-id"] + assert first_request["include"] == ["reasoning.encrypted_content"] + assert "oauth_refresh_token" not in first_request + assert first_request["api_base"] == OPENAI_OAUTH_CHATGPT_API_BASE + + +def test_llm_streams_responses_text_and_records_completion_usage() -> None: + llm = LLM( + adapter_id=OpenAIOAuthLLMAdapter.get_id(), + adapter_metadata=_metadata("token", "account"), + ) + events = iter( + [ + {"type": "response.output_text.delta", "delta": "hello"}, + { + "type": "response.completed", + "response": { + "id": "response-1", + "usage": {"input_tokens": 4, "output_tokens": 2, "total_tokens": 6}, + }, + }, + ] + ) + + with ( + patch("unstract.sdk1.llm.litellm.responses", return_value=events), + patch.object(llm, "_record_usage") as record_usage, + ): + chunks = list(llm.stream_complete("hello")) + + assert [chunk.text for chunk in chunks] == ["hello"] + record_usage.assert_called_once() + assert record_usage.call_args.args[2] == { + "prompt_tokens": 4, + "completion_tokens": 2, + "total_tokens": 6, + } + + +def test_llm_async_completion_uses_responses_api() -> None: + llm = LLM( + adapter_id=OpenAIOAuthLLMAdapter.get_id(), + adapter_metadata=_metadata("token", "account"), + ) + response = {"output_text": "async response", "usage": {}} + + with ( + patch( + "unstract.sdk1.llm.litellm.aresponses", + new=AsyncMock(return_value=response), + ) as responses, + patch.object(llm, "_record_usage"), + ): + result = asyncio.run(llm.acomplete("hello")) + + assert result["response"].text == "async response" + assert responses.await_args.kwargs["extra_headers"]["ChatGPT-Account-Id"] == "account" From 9f9aecc70fbb3d2f087ec7e73fce53dd9ee94b88 Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 15:31:13 -0400 Subject: [PATCH 02/19] Fix CSRF header for OpenAI OAuth start --- frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx b/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx index 9019e5d17c..4ca71c0b02 100644 --- a/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx +++ b/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx @@ -1,5 +1,6 @@ import PropTypes from "prop-types"; import { useCallback, useEffect, useState } from "react"; +import Cookies from "js-cookie"; import { Typography } from "@/components/ui/shims/antd-typography"; import { getBaseUrl, O_AUTH_PROVIDERS } from "../../../helpers/GetStaticData"; @@ -172,6 +173,9 @@ function OAuthDs({ const response = await axiosPrivate({ method: "POST", url: "/api/v1/oauth/openai/start/", + headers: { + "X-CSRFToken": Cookies.get("csrftoken"), + }, }); const loginDetails = response?.data || {}; const newCacheKey = loginDetails.cache_key; From 64a86ade448cc9c20394352a0395354b8b86daab Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 15:55:09 -0400 Subject: [PATCH 03/19] Allow reopening pending OpenAI device login --- .../src/components/oauth-ds/oauth-ds/OAuthDs.jsx | 9 +++++++++ .../oauth-ds/openai/OpenAIOAuthButton.jsx | 13 ++++++++++--- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx b/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx index 4ca71c0b02..f857edac3c 100644 --- a/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx +++ b/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx @@ -170,6 +170,15 @@ function OAuthDs({ const handleOAuth = async () => { try { if (oAuthProvider === O_AUTH_PROVIDERS.OPENAI) { + if (oauthStatus === "pending" && deviceLogin?.verification_url) { + window.open( + deviceLogin.verification_url, + "_blank", + "toolbar=yes,scrollbars=yes,resizable=yes,top=200,left=500,width=500,height=600", + ); + return; + } + const response = await axiosPrivate({ method: "POST", url: "/api/v1/oauth/openai/start/", diff --git a/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx b/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx index 565a62146c..6a23f73ea5 100644 --- a/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx +++ b/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx @@ -13,20 +13,27 @@ const OpenAIOAuthButton = ({ userCode, accountLabel, }) => { - const buttonLabel = status === "success" ? "Authenticated" : buttonText; + const isPending = status === "pending"; + const buttonLabel = + status === "success" + ? "Authenticated" + : isPending && verificationUrl + ? "Open OpenAI device login" + : buttonText; return (
{status === "pending" && verificationUrl && userCode && ( + Waiting for authorization. {" "} Open{" "} OpenAI device login From 504cdd4a654a803726e2e9537d9802c1e9a8cde0 Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 16:06:55 -0400 Subject: [PATCH 04/19] Make OpenAI OAuth model a dropdown --- .../adapters/llm1/static/openai-oauth.json | 18 +++++++++++++++++- unstract/sdk1/tests/test_openai_oauth.py | 13 +++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json index beebec0e43..a038ccb523 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json +++ b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json @@ -15,8 +15,24 @@ "model": { "type": "string", "title": "Model", + "enum": [ + "gpt-5.6-sol", + "gpt-5.6-terra", + "gpt-5.6-luna", + "gpt-5.6", + "gpt-5.3-codex", + "gpt-5-codex" + ], + "enumNames": [ + "GPT-5.6 Sol (flagship)", + "GPT-5.6 Terra (balanced)", + "GPT-5.6 Luna (efficient)", + "GPT-5.6 (default alias)", + "GPT-5.3 Codex", + "GPT-5 Codex (legacy)" + ], "default": "gpt-5-codex", - "description": "The Codex-capable model to use through your ChatGPT account." + "description": "Choose the Codex-capable model to use through your ChatGPT account." }, "max_tokens": { "type": "number", diff --git a/unstract/sdk1/tests/test_openai_oauth.py b/unstract/sdk1/tests/test_openai_oauth.py index 94a4953a21..fb627a9330 100644 --- a/unstract/sdk1/tests/test_openai_oauth.py +++ b/unstract/sdk1/tests/test_openai_oauth.py @@ -55,6 +55,19 @@ def test_openai_oauth_adapter_is_registered_with_auth_metadata() -> None: } schema = json.loads(OpenAIOAuthLLMAdapter.get_json_schema()) assert schema["properties"]["model"]["default"] == "gpt-5-codex" + assert schema["properties"]["model"]["enum"] == [ + "gpt-5.6-sol", + "gpt-5.6-terra", + "gpt-5.6-luna", + "gpt-5.6", + "gpt-5.3-codex", + "gpt-5-codex", + ] + assert schema["allOf"][0]["then"]["properties"]["reasoning_effort"]["enum"] == [ + "low", + "medium", + "high", + ] def test_openai_oauth_parameters_normalize_model_and_fix_endpoint() -> None: From 13d254298cb6a8fc8e24d46ed473f75d77b4c84e Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 16:20:09 -0400 Subject: [PATCH 05/19] Expose full OpenAI reasoning effort range --- .../sdk1/adapters/llm1/static/openai-oauth.json | 15 +++++++++++++-- unstract/sdk1/tests/test_openai_oauth.py | 3 +++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json index a038ccb523..ac9bc4bac5 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json +++ b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json @@ -79,13 +79,24 @@ "reasoning_effort": { "type": "string", "enum": [ + "none", "low", "medium", - "high" + "high", + "xhigh", + "max" + ], + "enumNames": [ + "None", + "Low", + "Medium", + "High", + "Extra high", + "Maximum" ], "default": "medium", "title": "Reasoning Effort", - "description": "Sets the reasoning strength when reasoning is enabled." + "description": "Sets the reasoning strength when reasoning is enabled. Use None to minimize reasoning overhead." } }, "required": [ diff --git a/unstract/sdk1/tests/test_openai_oauth.py b/unstract/sdk1/tests/test_openai_oauth.py index fb627a9330..00109a5d8e 100644 --- a/unstract/sdk1/tests/test_openai_oauth.py +++ b/unstract/sdk1/tests/test_openai_oauth.py @@ -64,9 +64,12 @@ def test_openai_oauth_adapter_is_registered_with_auth_metadata() -> None: "gpt-5-codex", ] assert schema["allOf"][0]["then"]["properties"]["reasoning_effort"]["enum"] == [ + "none", "low", "medium", "high", + "xhigh", + "max", ] From 33f359a66e4d0d0ba7c972d1adfa5d7732422482 Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 16:29:15 -0400 Subject: [PATCH 06/19] Complete OpenAI OAuth model picker --- .../sdk1/adapters/llm1/static/openai-oauth.json | 16 ++++++++++++++-- unstract/sdk1/tests/test_openai_oauth.py | 11 ++++++++++- 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json index ac9bc4bac5..fc20a50ec5 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json +++ b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json @@ -16,23 +16,35 @@ "type": "string", "title": "Model", "enum": [ + "gpt-6-astra", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna", "gpt-5.6", + "gpt-5.3-codex-spark", + "gpt-5.5", + "gpt-5.4", + "gpt-5.4-mini", + "gpt-5.2", "gpt-5.3-codex", "gpt-5-codex" ], "enumNames": [ + "GPT-6 Astra (preview)", "GPT-5.6 Sol (flagship)", "GPT-5.6 Terra (balanced)", "GPT-5.6 Luna (efficient)", "GPT-5.6 (default alias)", + "GPT-5.3 Codex Spark (preview)", + "GPT-5.5 (legacy)", + "GPT-5.4 (legacy)", + "GPT-5.4 Mini (legacy)", + "GPT-5.2 (deprecated)", "GPT-5.3 Codex", "GPT-5 Codex (legacy)" ], - "default": "gpt-5-codex", - "description": "Choose the Codex-capable model to use through your ChatGPT account." + "default": "gpt-5.6", + "description": "Choose a model available to your ChatGPT/Codex account. Availability varies by account, plan, and rollout." }, "max_tokens": { "type": "number", diff --git a/unstract/sdk1/tests/test_openai_oauth.py b/unstract/sdk1/tests/test_openai_oauth.py index 00109a5d8e..1fb4fdc1fd 100644 --- a/unstract/sdk1/tests/test_openai_oauth.py +++ b/unstract/sdk1/tests/test_openai_oauth.py @@ -54,15 +54,24 @@ def test_openai_oauth_adapter_is_registered_with_auth_metadata() -> None: "python_social_auth_backend": "openai-oauth", } schema = json.loads(OpenAIOAuthLLMAdapter.get_json_schema()) - assert schema["properties"]["model"]["default"] == "gpt-5-codex" + assert schema["properties"]["model"]["default"] == "gpt-5.6" assert schema["properties"]["model"]["enum"] == [ + "gpt-6-astra", "gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna", "gpt-5.6", + "gpt-5.3-codex-spark", + "gpt-5.5", + "gpt-5.4", + "gpt-5.4-mini", + "gpt-5.2", "gpt-5.3-codex", "gpt-5-codex", ] + assert len(schema["properties"]["model"]["enum"]) == len( + schema["properties"]["model"]["enumNames"] + ) assert schema["allOf"][0]["then"]["properties"]["reasoning_effort"]["enum"] == [ "none", "low", From 258cf014f96c983e43683c26319e610f7dea4ee5 Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 16:53:03 -0400 Subject: [PATCH 07/19] Discover OpenAI OAuth models per account --- backend/connector_auth_v2/openai_oauth.py | 15 + backend/connector_auth_v2/urls.py | 2 + backend/connector_auth_v2/views.py | 61 +++- .../input-output/configure-ds/ConfigureDs.jsx | 73 ++++- .../components/oauth-ds/oauth-ds/OAuthDs.jsx | 74 +++++ .../adapters/llm1/static/openai-oauth.json | 50 +-- .../sdk1/src/unstract/sdk1/auth/__init__.py | 10 + .../src/unstract/sdk1/auth/openai_oauth.py | 304 +++++++++++++++++- unstract/sdk1/tests/test_openai_oauth.py | 120 +++++-- 9 files changed, 636 insertions(+), 73 deletions(-) diff --git a/backend/connector_auth_v2/openai_oauth.py b/backend/connector_auth_v2/openai_oauth.py index 68f851e364..fafc97c34c 100644 --- a/backend/connector_auth_v2/openai_oauth.py +++ b/backend/connector_auth_v2/openai_oauth.py @@ -23,8 +23,10 @@ OPENAI_OAUTH_PRIVATE_FIELDS, OPENAI_OAUTH_TOKEN_URL, OpenAIOAuthError, + build_openai_oauth_json_schema, extract_account_id, extract_email, + fetch_openai_oauth_model_catalog, refresh_openai_oauth_metadata, ) from utils.user_session import UserSessionUtils @@ -324,6 +326,19 @@ def credentials_for_request( cls._save_state(cache_key, state) return refreshed + @staticmethod + def dynamic_model_schema( + credentials: dict[str, Any], *, current_model: str | None = None + ) -> tuple[dict[str, Any], dict[str, Any]]: + """Return a live account-specific form schema and refreshed credentials.""" + refreshed = refresh_openai_oauth_metadata(credentials) + catalog = fetch_openai_oauth_model_catalog(refreshed) + schema = build_openai_oauth_json_schema( + catalog, + current_model=current_model, + ) + return schema, refreshed + @classmethod def consume(cls, cache_key: str, request: Request) -> None: # Validate ownership before deleting so a leaked cache key cannot be diff --git a/backend/connector_auth_v2/urls.py b/backend/connector_auth_v2/urls.py index b4532ee3d9..050f268a47 100644 --- a/backend/connector_auth_v2/urls.py +++ b/backend/connector_auth_v2/urls.py @@ -10,6 +10,7 @@ ) openai_oauth_start = ConnectorAuthViewSet.as_view({"post": "openai_start"}) openai_oauth_poll = ConnectorAuthViewSet.as_view({"get": "openai_poll"}) +openai_oauth_models = ConnectorAuthViewSet.as_view({"get": "openai_models"}) urlpatterns = format_suffix_patterns( [ @@ -21,5 +22,6 @@ ), path("oauth/openai/start/", openai_oauth_start, name="openai-oauth-start"), path("oauth/openai/poll/", openai_oauth_poll, name="openai-oauth-poll"), + path("oauth/openai/models/", openai_oauth_models, name="openai-oauth-models"), ] ) diff --git a/backend/connector_auth_v2/views.py b/backend/connector_auth_v2/views.py index a17bd99228..a03c646dfb 100644 --- a/backend/connector_auth_v2/views.py +++ b/backend/connector_auth_v2/views.py @@ -1,12 +1,16 @@ import logging import uuid +from adapter_processor_v2.models import AdapterInstance from django.conf import settings from rest_framework import status, viewsets from rest_framework.request import Request from rest_framework.response import Response from rest_framework.versioning import URLPathVersioning -from unstract.sdk1.auth.openai_oauth import OpenAIOAuthError +from unstract.sdk1.auth.openai_oauth import ( + OpenAIOAuthError, + is_openai_oauth_adapter, +) from utils.user_session import UserSessionUtils from connector_auth_v2.constants import SocialAuthConstants @@ -93,3 +97,58 @@ def openai_poll(self, request: Request) -> Response: return Response( {"message": str(exc)}, status=status.HTTP_502_BAD_GATEWAY ) + + def openai_models(self, request: Request) -> Response: + """Return the live model/reasoning schema for one OAuth account.""" + if unauthorized := self._require_authenticated(request): + return unauthorized + + oauth_key = request.query_params.get("oauth-key") + adapter_instance_id = request.query_params.get("adapter-instance-id") + current_model = request.query_params.get("model") + + try: + if oauth_key: + credentials = OpenAIOAuthService.credentials_for_request( + oauth_key, request + ) + elif adapter_instance_id: + try: + adapter_uuid = uuid.UUID(adapter_instance_id) + except (AttributeError, TypeError, ValueError) as exc: + raise OpenAIOAuthSessionError( + "OpenAI OAuth adapter was not found or is not accessible" + ) from exc + adapter = ( + AdapterInstance.objects.for_user(request.user) + .filter(pk=adapter_uuid) + .first() + ) + if adapter is None or not is_openai_oauth_adapter(adapter.adapter_id): + raise OpenAIOAuthSessionError( + "OpenAI OAuth adapter was not found or is not accessible" + ) + credentials = adapter.metadata + if not isinstance(credentials, dict): + raise OpenAIOAuthSessionError( + "OpenAI OAuth adapter credentials are invalid" + ) + if not current_model and isinstance(credentials.get("model"), str): + current_model = credentials["model"] + else: + raise OpenAIOAuthSessionError( + "An OpenAI OAuth login session or adapter is required" + ) + + schema, _ = OpenAIOAuthService.dynamic_model_schema( + credentials, + current_model=current_model, + ) + return Response({"json_schema": schema}, status=status.HTTP_200_OK) + except OpenAIOAuthSessionError as exc: + return Response({"message": str(exc)}, status=status.HTTP_400_BAD_REQUEST) + except OpenAIOAuthError as exc: + logger.warning("OpenAI OAuth model discovery failed: %s", exc) + return Response( + {"message": str(exc)}, status=status.HTTP_502_BAD_GATEWAY + ) diff --git a/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx b/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx index e26c88a8ad..6740897c22 100644 --- a/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx +++ b/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx @@ -1,6 +1,6 @@ import { Info } from "lucide-react"; import PropTypes from "prop-types"; -import { createRef, useEffect, useState } from "react"; +import { createRef, useCallback, useEffect, useState } from "react"; import { Col, Row } from "@/components/ui/shims/antd-layout"; import { Popover } from "@/components/ui/shims/antd-overlays"; @@ -36,6 +36,7 @@ function ConfigureDs({ const [isTcSuccessful, setIsTcSuccessful] = useState(false); const [isTcLoading, setIsTcLoading] = useState(false); const [isSubmitApiLoading, setIsSubmitApiLoading] = useState(false); + const [oauthSpec, setOAuthSpec] = useState(null); const [cacheKey, setCacheKey] = useState(""); const [status, setStatus] = useState(""); @@ -54,6 +55,72 @@ function ConfigureDs({ const oauthCacheKey = `oauth-cachekey-${selectedSourceId}`; const oauthStatusKey = `oauth-status-${selectedSourceId}`; + const findReasoningSchema = useCallback((schema, model) => { + if (!schema || !model || !Array.isArray(schema.allOf)) { + return undefined; + } + const condition = schema.allOf.find( + (candidate) => + candidate?.if?.properties?.enable_reasoning?.const === true && + candidate?.if?.properties?.model?.const === model, + ); + return condition?.then?.properties?.reasoning_effort; + }, []); + + const handleOpenAIModelSchema = useCallback( + (dynamicSchema) => { + const modelSchema = dynamicSchema?.properties?.model; + if (!Array.isArray(modelSchema?.enum) || modelSchema.enum.length === 0) { + return; + } + + setOAuthSpec(dynamicSchema); + setFormData((currentFormData) => { + const current = currentFormData || {}; + const next = { ...current }; + if (!modelSchema.enum.includes(next.model)) { + next.model = modelSchema.default || modelSchema.enum[0]; + } + + if (next.enable_reasoning) { + const reasoningSchema = findReasoningSchema(dynamicSchema, next.model); + if ( + Array.isArray(reasoningSchema?.enum) && + reasoningSchema.enum.length > 0 && + !reasoningSchema.enum.includes(next.reasoning_effort) + ) { + next.reasoning_effort = + reasoningSchema.default || reasoningSchema.enum[0]; + } + } + return next; + }); + }, + [findReasoningSchema, setFormData], + ); + + useEffect(() => { + setOAuthSpec(null); + }, [spec]); + + useEffect(() => { + if (!oauthSpec || !formData?.enable_reasoning || !formData?.model) { + return; + } + const reasoningSchema = findReasoningSchema(oauthSpec, formData.model); + if ( + !Array.isArray(reasoningSchema?.enum) || + reasoningSchema.enum.length === 0 || + reasoningSchema.enum.includes(formData.reasoning_effort) + ) { + return; + } + setFormData((currentFormData) => ({ + ...(currentFormData || {}), + reasoning_effort: reasoningSchema.default || reasoningSchema.enum[0], + })); + }, [findReasoningSchema, formData, oauthSpec, setFormData]); + // Determine if this is a new or existing connector const hasOAuthCredentials = metadata && @@ -405,7 +472,7 @@ function ConfigureDs({
)} )} diff --git a/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx b/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx index f857edac3c..fe12df1af3 100644 --- a/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx +++ b/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx @@ -17,6 +17,8 @@ function OAuthDs({ setStatus, selectedSourceId, isExistingConnector, + adapterInstanceId, + onModelsLoaded, disabled = false, }) { const axiosPrivate = useAxiosPrivate(); @@ -60,6 +62,7 @@ function OAuthDs({ return null; } }); + const [activeLoginCacheKey, setActiveLoginCacheKey] = useState(null); const updateOAuthStatus = useCallback((newStatus) => { setOAuthStatus(newStatus); @@ -67,6 +70,30 @@ function OAuthDs({ localStorage.setItem(oauthStatusKey, newStatus); }, [oauthStatusKey, setStatus]); + const loadOpenAIModelSchema = useCallback( + async (oauthKey = "") => { + const params = new URLSearchParams(); + if (oauthKey) { + params.set("oauth-key", oauthKey); + } else if (adapterInstanceId) { + params.set("adapter-instance-id", adapterInstanceId); + } else { + return; + } + + const response = await axiosPrivate({ + method: "GET", + url: `/api/v1/oauth/openai/models/?${params.toString()}`, + }); + const dynamicSchema = response?.data?.json_schema; + if (!dynamicSchema) { + throw new Error("OpenAI OAuth returned no model schema"); + } + onModelsLoaded?.(dynamicSchema); + }, + [adapterInstanceId, axiosPrivate, onModelsLoaded], + ); + useEffect(() => { const handleStorageChange = () => { // Listen for changes to our specific connector only @@ -112,6 +139,50 @@ function OAuthDs({ }; }, [selectedSourceId, oauthCacheKey, oauthStatusKey, setCacheKey, setStatus]); + useEffect(() => { + if (oAuthProvider !== O_AUTH_PROVIDERS.OPENAI || !onModelsLoaded) { + return undefined; + } + + // Existing adapters use their server-side encrypted credentials. A new + // OAuth login takes precedence only after that login has completed, so a + // stale localStorage cache key cannot select another account by accident. + const sessionKey = + oauthStatus === "success" + ? activeLoginCacheKey || (!adapterInstanceId ? loginCacheKey : "") + : ""; + const source = sessionKey || (adapterInstanceId ? adapterInstanceId : ""); + if (!source) { + return undefined; + } + + let isActive = true; + loadOpenAIModelSchema(sessionKey) + .catch((err) => { + if (!isActive) { + return; + } + const message = + err?.response?.data?.message || + "Could not load models available to this OpenAI account"; + setAlertDetails(handleException(err, message)); + }); + + return () => { + isActive = false; + }; + }, [ + activeLoginCacheKey, + adapterInstanceId, + handleException, + loadOpenAIModelSchema, + loginCacheKey, + oauthStatus, + oAuthProvider, + onModelsLoaded, + setAlertDetails, + ]); + useEffect(() => { if ( oAuthProvider !== O_AUTH_PROVIDERS.OPENAI || @@ -192,6 +263,7 @@ function OAuthDs({ throw new Error("OpenAI OAuth did not return a login session"); } setLoginCacheKey(newCacheKey); + setActiveLoginCacheKey(newCacheKey); setCacheKey(newCacheKey); localStorage.setItem(oauthCacheKey, newCacheKey); setDeviceLogin(loginDetails); @@ -283,6 +355,8 @@ OAuthDs.propTypes = { setStatus: PropTypes.func, selectedSourceId: PropTypes.string.isRequired, isExistingConnector: PropTypes.bool, + adapterInstanceId: PropTypes.string, + onModelsLoaded: PropTypes.func, disabled: PropTypes.bool, }; diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json index fc20a50ec5..5e15ff513d 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json +++ b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json @@ -15,36 +15,7 @@ "model": { "type": "string", "title": "Model", - "enum": [ - "gpt-6-astra", - "gpt-5.6-sol", - "gpt-5.6-terra", - "gpt-5.6-luna", - "gpt-5.6", - "gpt-5.3-codex-spark", - "gpt-5.5", - "gpt-5.4", - "gpt-5.4-mini", - "gpt-5.2", - "gpt-5.3-codex", - "gpt-5-codex" - ], - "enumNames": [ - "GPT-6 Astra (preview)", - "GPT-5.6 Sol (flagship)", - "GPT-5.6 Terra (balanced)", - "GPT-5.6 Luna (efficient)", - "GPT-5.6 (default alias)", - "GPT-5.3 Codex Spark (preview)", - "GPT-5.5 (legacy)", - "GPT-5.4 (legacy)", - "GPT-5.4 Mini (legacy)", - "GPT-5.2 (deprecated)", - "GPT-5.3 Codex", - "GPT-5 Codex (legacy)" - ], - "default": "gpt-5.6", - "description": "Choose a model available to your ChatGPT/Codex account. Availability varies by account, plan, and rollout." + "description": "Sign in with OpenAI to load the models available to this ChatGPT/Codex account." }, "max_tokens": { "type": "number", @@ -90,25 +61,8 @@ "properties": { "reasoning_effort": { "type": "string", - "enum": [ - "none", - "low", - "medium", - "high", - "xhigh", - "max" - ], - "enumNames": [ - "None", - "Low", - "Medium", - "High", - "Extra high", - "Maximum" - ], - "default": "medium", "title": "Reasoning Effort", - "description": "Sets the reasoning strength when reasoning is enabled. Use None to minimize reasoning overhead." + "description": "Sign in with OpenAI to load reasoning levels supported by the selected model." } }, "required": [ diff --git a/unstract/sdk1/src/unstract/sdk1/auth/__init__.py b/unstract/sdk1/src/unstract/sdk1/auth/__init__.py index 14d4447960..5f8177fa52 100644 --- a/unstract/sdk1/src/unstract/sdk1/auth/__init__.py +++ b/unstract/sdk1/src/unstract/sdk1/auth/__init__.py @@ -3,11 +3,16 @@ from unstract.sdk1.auth.openai_oauth import ( OPENAI_OAUTH_ADAPTER_PREFIX, OPENAI_OAUTH_CHATGPT_API_BASE, + OPENAI_OAUTH_CODEX_CLIENT_VERSION, + OPENAI_OAUTH_MODELS_URL, OPENAI_OAUTH_PROVIDER, OpenAIOAuthError, + OpenAIOAuthModelCatalogError, OpenAIOAuthRefreshError, + build_openai_oauth_json_schema, extract_account_id, extract_email, + fetch_openai_oauth_model_catalog, is_openai_oauth_adapter, refresh_openai_oauth_metadata, ) @@ -15,11 +20,16 @@ __all__ = [ "OPENAI_OAUTH_ADAPTER_PREFIX", "OPENAI_OAUTH_CHATGPT_API_BASE", + "OPENAI_OAUTH_CODEX_CLIENT_VERSION", + "OPENAI_OAUTH_MODELS_URL", "OPENAI_OAUTH_PROVIDER", "OpenAIOAuthError", + "OpenAIOAuthModelCatalogError", "OpenAIOAuthRefreshError", + "build_openai_oauth_json_schema", "extract_account_id", "extract_email", + "fetch_openai_oauth_model_catalog", "is_openai_oauth_adapter", "refresh_openai_oauth_metadata", ] diff --git a/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py b/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py index b2ea9deb13..993028fb91 100644 --- a/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py +++ b/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py @@ -11,7 +11,7 @@ import json import os import time -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import Any import httpx @@ -25,6 +25,10 @@ "OPENAI_OAUTH_CLIENT_ID", "app_EMoamEEZ73f0CkXaXp7hrann" ) OPENAI_OAUTH_CHATGPT_API_BASE = "https://chatgpt.com/backend-api/codex" +OPENAI_OAUTH_MODELS_URL = f"{OPENAI_OAUTH_CHATGPT_API_BASE}/models" +OPENAI_OAUTH_CODEX_CLIENT_VERSION = os.environ.get( + "OPENAI_OAUTH_CODEX_CLIENT_VERSION", "0.149.0" +) OPENAI_OAUTH_DEVICE_USERCODE_URL = ( f"{OPENAI_OAUTH_AUTH_BASE}/api/accounts/deviceauth/usercode" ) @@ -63,6 +67,10 @@ class OpenAIOAuthRefreshError(OpenAIOAuthError): """Raised when an expired access token cannot be refreshed.""" +class OpenAIOAuthModelCatalogError(OpenAIOAuthError): + """Raised when the account-specific Codex model catalog cannot be read.""" + + def is_openai_oauth_adapter(adapter_id: str | None) -> bool: """Return whether ``adapter_id`` identifies the OAuth-backed adapter.""" return bool(adapter_id and adapter_id.startswith(OPENAI_OAUTH_ADAPTER_PREFIX)) @@ -146,6 +154,300 @@ def _response_error(response: httpx.Response, operation: str) -> OpenAIOAuthRefr ) +def _reasoning_levels(value: object) -> list[dict[str, str]]: + """Normalize the reasoning options returned by the Codex catalog.""" + if not isinstance(value, list): + return [] + + levels: list[dict[str, str]] = [] + seen: set[str] = set() + for item in value: + if not isinstance(item, Mapping): + continue + effort = item.get("effort") or item.get("reasoning_effort") + if not isinstance(effort, str) or not effort.strip(): + continue + effort = effort.strip() + if effort in seen: + continue + description = item.get("description") + if not isinstance(description, str) or not description.strip(): + description = effort.replace("_", " ").replace("-", " ").title() + levels.append({"effort": effort, "description": description.strip()}) + seen.add(effort) + return levels + + +def _normalize_catalog_model( + item: Mapping[str, Any], index: int +) -> dict[str, Any] | None: + slug = item.get("slug") + if not isinstance(slug, str) or not slug.strip(): + return None + slug = slug.strip() + + visibility = item.get("visibility") + if isinstance(visibility, str) and visibility.lower() in { + "hidden", + "none", + "unlisted", + }: + return None + if item.get("supported_in_api") is False: + return None + + display_name = item.get("display_name") + if not isinstance(display_name, str) or not display_name.strip(): + display_name = slug + description = item.get("description") + if not isinstance(description, str): + description = "" + priority = item.get("priority") + try: + normalized_priority = int(priority) + except (TypeError, ValueError): + normalized_priority = 2**31 - 1 + + return { + "slug": slug, + "display_name": display_name.strip(), + "description": description.strip(), + "default_reasoning_level": item.get("default_reasoning_level"), + "supported_reasoning_levels": _reasoning_levels( + item.get("supported_reasoning_levels") + ), + "priority": normalized_priority, + "_catalog_index": index, + "is_deprecated": bool(item.get("upgrade")), + } + + +def _normalize_catalog_models(raw_models: list[object]) -> list[dict[str, Any]]: + models: list[dict[str, Any]] = [] + seen_slugs: set[str] = set() + for index, item in enumerate(raw_models): + if not isinstance(item, Mapping): + continue + model = _normalize_catalog_model(item, index) + if model is None or model["slug"] in seen_slugs: + continue + models.append(model) + seen_slugs.add(model["slug"]) + + models.sort(key=lambda item: (item["priority"], item["_catalog_index"])) + for item in models: + item.pop("_catalog_index", None) + return models + + +def _catalog_models_from_response(response: httpx.Response) -> list[object]: + if not 200 <= response.status_code < 300: + raise OpenAIOAuthModelCatalogError( + f"OpenAI OAuth model discovery failed with status {response.status_code}" + ) + + try: + payload = response.json() + except (ValueError, json.JSONDecodeError) as exc: + raise OpenAIOAuthModelCatalogError( + "OpenAI OAuth model discovery returned an invalid response" + ) from exc + if not isinstance(payload, Mapping) or not isinstance(payload.get("models"), list): + raise OpenAIOAuthModelCatalogError( + "OpenAI OAuth model discovery returned no model catalog" + ) + return payload["models"] + + +def fetch_openai_oauth_model_catalog( + metadata: Mapping[str, Any], + *, + client_version: str | None = None, +) -> list[dict[str, Any]]: + """Fetch the visible model catalog for one authenticated ChatGPT account. + + The Codex endpoint is account- and plan-aware. No model ids or reasoning + levels are defined here; the response is normalized only enough for the + configuration UI to consume it safely. + """ + access_token = metadata.get("oauth_access_token") + account_id = metadata.get("oauth_account_id") + if not isinstance(access_token, str) or not access_token.strip(): + raise OpenAIOAuthModelCatalogError( + "OpenAI OAuth model discovery requires an access token" + ) + if not isinstance(account_id, str) or not account_id.strip(): + raise OpenAIOAuthModelCatalogError( + "OpenAI OAuth model discovery requires a ChatGPT account" + ) + + try: + response = httpx.get( + OPENAI_OAUTH_MODELS_URL, + params={ + "client_version": client_version or OPENAI_OAUTH_CODEX_CLIENT_VERSION + }, + headers={ + "Authorization": f"Bearer {access_token}", + "ChatGPT-Account-ID": account_id, + "Accept": "application/json", + # Keep the catalog aligned with the originator used for model + # requests by the Unstract adapter. + "originator": "unstract", + }, + timeout=15.0, + ) + except httpx.HTTPError as exc: + raise OpenAIOAuthModelCatalogError( + "OpenAI OAuth model discovery could not reach ChatGPT" + ) from exc + + models = _normalize_catalog_models(_catalog_models_from_response(response)) + + if not models: + raise OpenAIOAuthModelCatalogError( + "OpenAI OAuth model discovery returned no available models" + ) + return models + + +def _valid_catalog_models( + model_catalog: Sequence[Mapping[str, Any]], +) -> list[dict[str, Any]]: + return [ + dict(item) + for item in model_catalog + if isinstance(item, Mapping) + and isinstance(item.get("slug"), str) + and item["slug"].strip() + ] + + +def _model_labels(models: Sequence[Mapping[str, Any]]) -> list[str]: + labels: list[str] = [] + for model in models: + slug = str(model["slug"]).strip() + label = str(model.get("display_name") or slug).strip() + if model.get("is_deprecated") and "deprecated" not in label.lower(): + label = f"{label} (deprecated)" + labels.append(label) + return labels + + +def _selected_catalog_model(slugs: Sequence[str], current_model: str | None) -> str: + selected_model = current_model.strip() if isinstance(current_model, str) else "" + for prefix in ("openai/", "chatgpt/"): + if selected_model.startswith(prefix): + selected_model = selected_model[len(prefix) :] + break + return selected_model if selected_model in slugs else slugs[0] + + +def _reasoning_schema_for_model( + model: Mapping[str, Any], +) -> dict[str, Any] | None: + levels_value = model.get("supported_reasoning_levels") + if not isinstance(levels_value, list): + return None + levels = [ + dict(level) + for level in levels_value + if isinstance(level, Mapping) + and isinstance(level.get("effort"), str) + and level["effort"].strip() + ] + if not levels: + return None + + efforts = [str(level["effort"]).strip() for level in levels] + effort_labels = [ + str(level.get("description") or effort).strip() + for level, effort in zip(levels, efforts, strict=True) + ] + default_effort = model.get("default_reasoning_level") + if default_effort not in efforts: + default_effort = efforts[0] + return { + "type": "string", + "enum": efforts, + "enumNames": effort_labels, + "default": default_effort, + "title": "Reasoning Effort", + "description": "Reasoning levels reported for this model by ChatGPT.", + } + + +def _add_reasoning_conditions( + all_of: list[dict[str, Any]], models: Sequence[Mapping[str, Any]] +) -> None: + for model in models: + slug = str(model["slug"]).strip() + reasoning_schema = _reasoning_schema_for_model(model) + if reasoning_schema is None: + continue + all_of.append( + { + "if": { + "required": ["enable_reasoning", "model"], + "properties": { + "enable_reasoning": {"const": True}, + "model": {"const": slug}, + }, + }, + "then": { + "properties": {"reasoning_effort": reasoning_schema}, + "required": ["reasoning_effort"], + }, + } + ) + + +def build_openai_oauth_json_schema( + model_catalog: Sequence[Mapping[str, Any]], + *, + current_model: str | None = None, + base_schema: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """Build a form schema from one account's live Codex model catalog. + + ``base_schema`` is injectable for tests. In production the adapter's + schema is loaded lazily to avoid an import cycle with this auth module. + """ + if base_schema is None: + from unstract.sdk1.adapters.llm1.openai_oauth import OpenAIOAuthLLMAdapter + + base_schema = json.loads(OpenAIOAuthLLMAdapter.get_json_schema()) + + schema = json.loads(json.dumps(base_schema)) + models = _valid_catalog_models(model_catalog) + if not models: + raise OpenAIOAuthModelCatalogError( + "Cannot build an OpenAI OAuth form without available models" + ) + + model_property = schema.setdefault("properties", {}).setdefault("model", {}) + slugs = [str(item["slug"]).strip() for item in models] + labels = _model_labels(models) + selected_model = _selected_catalog_model(slugs, current_model) + + model_property["enum"] = slugs + model_property["enumNames"] = labels + model_property["default"] = selected_model + model_property["description"] = ( + "Models reported as available by this ChatGPT/Codex account." + ) + + all_of = schema.setdefault("allOf", []) + if not isinstance(all_of, list): + all_of = [] + schema["allOf"] = all_of + + _add_reasoning_conditions(all_of, models) + + schema["x-openai-oauth-model-source"] = "chatgpt-account" + return schema + + def refresh_openai_oauth_metadata( metadata: Mapping[str, Any], *, diff --git a/unstract/sdk1/tests/test_openai_oauth.py b/unstract/sdk1/tests/test_openai_oauth.py index 1fb4fdc1fd..f2e5f173bc 100644 --- a/unstract/sdk1/tests/test_openai_oauth.py +++ b/unstract/sdk1/tests/test_openai_oauth.py @@ -17,9 +17,13 @@ ) from unstract.sdk1.auth.openai_oauth import ( OPENAI_OAUTH_CHATGPT_API_BASE, + OPENAI_OAUTH_CODEX_CLIENT_VERSION, + OPENAI_OAUTH_MODELS_URL, OPENAI_OAUTH_TOKEN_URL, + build_openai_oauth_json_schema, extract_account_id, extract_email, + fetch_openai_oauth_model_catalog, refresh_openai_oauth_metadata, ) from unstract.sdk1.llm import LLM @@ -54,31 +58,105 @@ def test_openai_oauth_adapter_is_registered_with_auth_metadata() -> None: "python_social_auth_backend": "openai-oauth", } schema = json.loads(OpenAIOAuthLLMAdapter.get_json_schema()) - assert schema["properties"]["model"]["default"] == "gpt-5.6" - assert schema["properties"]["model"]["enum"] == [ - "gpt-6-astra", - "gpt-5.6-sol", - "gpt-5.6-terra", - "gpt-5.6-luna", - "gpt-5.6", - "gpt-5.3-codex-spark", - "gpt-5.5", - "gpt-5.4", - "gpt-5.4-mini", - "gpt-5.2", - "gpt-5.3-codex", - "gpt-5-codex", + assert "enum" not in schema["properties"]["model"] + assert "enum" not in schema["allOf"][0]["then"]["properties"]["reasoning_effort"] + + +def test_openai_oauth_model_catalog_is_normalized_per_account() -> None: + response = MagicMock(status_code=200) + response.json.return_value = { + "models": [ + { + "slug": "hidden-model", + "display_name": "Hidden", + "visibility": "hidden", + }, + { + "slug": "gpt-account-fast", + "display_name": "Account Fast", + "description": "Fast account model", + "priority": 2, + "default_reasoning_level": "low", + "supported_reasoning_levels": [ + {"effort": "low", "description": "Light"}, + {"effort": "high", "description": "Deep"}, + ], + }, + { + "slug": "gpt-account-best", + "display_name": "Account Best", + "priority": 1, + "supported_in_api": True, + "supported_reasoning_levels": [{"effort": "max"}], + }, + {"slug": "not-supported", "supported_in_api": False}, + ] + } + metadata = _metadata("account-token", "workspace-123") + + with patch("unstract.sdk1.auth.openai_oauth.httpx.get", return_value=response) as get: + catalog = fetch_openai_oauth_model_catalog(metadata) + + assert [model["slug"] for model in catalog] == [ + "gpt-account-best", + "gpt-account-fast", + ] + assert catalog[1]["supported_reasoning_levels"] == [ + {"effort": "low", "description": "Light"}, + {"effort": "high", "description": "Deep"}, ] - assert len(schema["properties"]["model"]["enum"]) == len( - schema["properties"]["model"]["enumNames"] + get.assert_called_once_with( + OPENAI_OAUTH_MODELS_URL, + params={"client_version": OPENAI_OAUTH_CODEX_CLIENT_VERSION}, + headers={ + "Authorization": "Bearer account-token", + "ChatGPT-Account-ID": "workspace-123", + "Accept": "application/json", + "originator": "unstract", + }, + timeout=15.0, ) - assert schema["allOf"][0]["then"]["properties"]["reasoning_effort"]["enum"] == [ - "none", + + +def test_openai_oauth_dynamic_schema_uses_model_specific_reasoning_levels() -> None: + base_schema = json.loads(OpenAIOAuthLLMAdapter.get_json_schema()) + schema = build_openai_oauth_json_schema( + [ + { + "slug": "account-model-a", + "display_name": "Model A", + "default_reasoning_level": "high", + "supported_reasoning_levels": [ + {"effort": "low", "description": "Low"}, + {"effort": "high", "description": "High"}, + ], + }, + { + "slug": "account-model-b", + "display_name": "Model B", + "supported_reasoning_levels": [ + {"effort": "max", "description": "Maximum"}, + ], + }, + ], + current_model="account-model-b", + base_schema=base_schema, + ) + + model = schema["properties"]["model"] + assert model["enum"] == ["account-model-a", "account-model-b"] + assert model["enumNames"] == ["Model A", "Model B"] + assert model["default"] == "account-model-b" + model_conditions = schema["allOf"][2:] + assert model_conditions[0]["if"]["properties"]["model"] == { + "const": "account-model-a" + } + assert model_conditions[0]["then"]["properties"]["reasoning_effort"]["enum"] == [ "low", - "medium", "high", - "xhigh", - "max", + ] + assert model_conditions[1]["then"]["properties"]["reasoning_effort"]["enum"] == [ + "max" ] From 93185785c9f9856e240763e381940c1ba3947da0 Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 17:01:57 -0400 Subject: [PATCH 08/19] Keep OAuth form fields dynamic dropdowns --- .../unstract/sdk1/adapters/llm1/static/openai-oauth.json | 4 ++++ unstract/sdk1/tests/test_openai_oauth.py | 7 +++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json index 5e15ff513d..f03410d88c 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json +++ b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json @@ -15,6 +15,8 @@ "model": { "type": "string", "title": "Model", + "enum": [], + "enumNames": [], "description": "Sign in with OpenAI to load the models available to this ChatGPT/Codex account." }, "max_tokens": { @@ -62,6 +64,8 @@ "reasoning_effort": { "type": "string", "title": "Reasoning Effort", + "enum": [], + "enumNames": [], "description": "Sign in with OpenAI to load reasoning levels supported by the selected model." } }, diff --git a/unstract/sdk1/tests/test_openai_oauth.py b/unstract/sdk1/tests/test_openai_oauth.py index f2e5f173bc..7ef5eb85a3 100644 --- a/unstract/sdk1/tests/test_openai_oauth.py +++ b/unstract/sdk1/tests/test_openai_oauth.py @@ -58,8 +58,11 @@ def test_openai_oauth_adapter_is_registered_with_auth_metadata() -> None: "python_social_auth_backend": "openai-oauth", } schema = json.loads(OpenAIOAuthLLMAdapter.get_json_schema()) - assert "enum" not in schema["properties"]["model"] - assert "enum" not in schema["allOf"][0]["then"]["properties"]["reasoning_effort"] + assert schema["properties"]["model"]["enum"] == [] + assert ( + schema["allOf"][0]["then"]["properties"]["reasoning_effort"]["enum"] + == [] + ) def test_openai_oauth_model_catalog_is_normalized_per_account() -> None: From cae221022e8ec5442669dfefce41be02ad686ce4 Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 17:13:41 -0400 Subject: [PATCH 09/19] Render account reasoning choices after OAuth --- .../src/unstract/sdk1/auth/openai_oauth.py | 25 +++++++++++++++++++ unstract/sdk1/tests/test_openai_oauth.py | 2 +- 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py b/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py index 993028fb91..991bb48d72 100644 --- a/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py +++ b/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py @@ -402,6 +402,27 @@ def _add_reasoning_conditions( ) +def _remove_empty_reasoning_placeholder(all_of: list[dict[str, Any]]) -> None: + """Remove the pre-auth reasoning rule before adding live model rules.""" + all_of[:] = [ + condition + for condition in all_of + if not ( + isinstance(condition, Mapping) + and isinstance(condition.get("then"), Mapping) + and isinstance(condition["then"].get("properties"), Mapping) + and isinstance( + condition["then"]["properties"].get("reasoning_effort"), + Mapping, + ) + and condition["then"]["properties"]["reasoning_effort"].get( + "enum" + ) + == [] + ) + ] + + def build_openai_oauth_json_schema( model_catalog: Sequence[Mapping[str, Any]], *, @@ -442,6 +463,10 @@ def build_openai_oauth_json_schema( all_of = [] schema["allOf"] = all_of + # The static adapter schema keeps an empty reasoning field so RJSF can + # render the form before login. Once the account catalog is available, + # that placeholder must not remain as an additional empty constraint. + _remove_empty_reasoning_placeholder(all_of) _add_reasoning_conditions(all_of, models) schema["x-openai-oauth-model-source"] = "chatgpt-account" diff --git a/unstract/sdk1/tests/test_openai_oauth.py b/unstract/sdk1/tests/test_openai_oauth.py index 7ef5eb85a3..8a34889bea 100644 --- a/unstract/sdk1/tests/test_openai_oauth.py +++ b/unstract/sdk1/tests/test_openai_oauth.py @@ -150,7 +150,7 @@ def test_openai_oauth_dynamic_schema_uses_model_specific_reasoning_levels() -> N assert model["enum"] == ["account-model-a", "account-model-b"] assert model["enumNames"] == ["Model A", "Model B"] assert model["default"] == "account-model-b" - model_conditions = schema["allOf"][2:] + model_conditions = schema["allOf"][1:] assert model_conditions[0]["if"]["properties"]["model"] == { "const": "account-model-a" } From c5cd1444a88009444a2eb9a1dd17d983c18bb329 Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 17:36:14 -0400 Subject: [PATCH 10/19] Persist OpenAI OAuth adapter sessions --- backend/adapter_processor_v2/views.py | 32 ++++++++++++++++++- .../input-output/configure-ds/ConfigureDs.jsx | 22 +++++++------ .../components/oauth-ds/oauth-ds/OAuthDs.jsx | 12 ++++++- .../oauth-ds/openai/OpenAIOAuthButton.jsx | 2 +- frontend/src/hooks/useExceptionHandler.jsx | 5 +-- .../src/unstract/sdk1/auth/openai_oauth.py | 28 ++++++---------- unstract/sdk1/tests/test_openai_oauth.py | 15 +++++++++ 7 files changed, 83 insertions(+), 33 deletions(-) diff --git a/backend/adapter_processor_v2/views.py b/backend/adapter_processor_v2/views.py index c69b0e0fdb..409f9e55f0 100644 --- a/backend/adapter_processor_v2/views.py +++ b/backend/adapter_processor_v2/views.py @@ -111,6 +111,31 @@ def _prepare_openai_oauth_payload( return payload, oauth_key +def _saved_openai_oauth_metadata(request: Request) -> dict[str, Any] | None: + """Load one owned adapter's encrypted OAuth metadata for a test request.""" + adapter_instance_id = request.query_params.get("adapter-instance-id") + if not adapter_instance_id: + return None + try: + adapter_uuid = uuid.UUID(adapter_instance_id) + except (AttributeError, TypeError, ValueError) as exc: + raise ValidationError( + {"adapter-instance-id": "OpenAI OAuth adapter was not found"} + ) from exc + + adapter = ( + AdapterInstance.objects.for_user(request.user) + .filter(pk=adapter_uuid) + .first() + ) + if adapter is None or not is_openai_oauth_adapter(adapter.adapter_id): + raise ValidationError( + {"adapter-instance-id": "OpenAI OAuth adapter was not found"} + ) + metadata = adapter.metadata + return metadata if isinstance(metadata, dict) else None + + def _consume_openai_oauth_key(cache_key: str | None, request: Request) -> None: """Best-effort cleanup after credentials are durably stored.""" if not cache_key: @@ -191,7 +216,12 @@ def get_adapter_schema( def test(self, request: Request) -> Response: """Tests the connector against the credentials passed.""" payload = request.data.copy() - payload, _ = _prepare_openai_oauth_payload(request, payload) + existing_metadata = None + if is_openai_oauth_adapter(payload.get(AdapterKeys.ADAPTER_ID)): + existing_metadata = _saved_openai_oauth_metadata(request) + payload, _ = _prepare_openai_oauth_payload( + request, payload, existing_metadata=existing_metadata + ) serializer: AdapterInstanceSerializer = self.get_serializer(data=payload) serializer.is_valid(raise_exception=True) adapter_id = serializer.validated_data.get(AdapterKeys.ADAPTER_ID) diff --git a/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx b/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx index 6740897c22..222901ca65 100644 --- a/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx +++ b/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx @@ -128,6 +128,9 @@ function ConfigureDs({ (metadata.provider && metadata.uid) || metadata.oauth_authenticated); const isExistingConnector = Boolean(editItemId || hasOAuthCredentials); + const hasPersistedOAuthCredentials = Boolean( + !isConnector && isExistingConnector && hasOAuthCredentials, + ); // Determine if OAuth authentication method is selected const isOAuthMethodSelected = () => { @@ -227,14 +230,6 @@ function ConfigureDs({ } }, [selectedSourceId, oAuthProvider, oauthStatusKey, oauthCacheKey]); - // Cleanup OAuth localStorage when component unmounts (modal close) - useEffect(() => { - return () => { - localStorage.removeItem(oauthCacheKey); - localStorage.removeItem(oauthStatusKey); - }; - }, [oauthCacheKey, oauthStatusKey]); - const handleTestConnection = (updatedFormData) => { // Check if there any error in form proceed to test connection only there is no error. if (formRef && !formRef.current?.validateForm()) { @@ -243,7 +238,10 @@ function ConfigureDs({ if ( oAuthProvider?.length && isOAuthMethodSelected() && - (status !== "success" || !cacheKey?.length) + !( + (status === "success" && cacheKey?.length) || + hasPersistedOAuthCredentials + ) ) { const providerName = oAuthProvider === "google-oauth2" @@ -290,8 +288,10 @@ function ConfigureDs({ ...body["connector_metadata"], ...{ "oauth-key": cacheKey }, }; - } else { + } else if (cacheKey?.length) { url = `${url}?oauth-key=${encodeURIComponent(cacheKey)}`; + } else if (editItemId && hasPersistedOAuthCredentials) { + url = `${url}?adapter-instance-id=${encodeURIComponent(editItemId)}`; } } @@ -487,6 +487,8 @@ function ConfigureDs({ setStatus={handleSetStatus} selectedSourceId={selectedSourceId} isExistingConnector={isExistingConnector} + hasOAuthCredentials={hasPersistedOAuthCredentials} + oauthAccountLabel={metadata?.oauth_account_label} adapterInstanceId={editItemId} onModelsLoaded={handleOpenAIModelSchema} /> diff --git a/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx b/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx index fe12df1af3..5a7aac4767 100644 --- a/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx +++ b/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx @@ -17,6 +17,8 @@ function OAuthDs({ setStatus, selectedSourceId, isExistingConnector, + hasOAuthCredentials = false, + oauthAccountLabel, adapterInstanceId, onModelsLoaded, disabled = false, @@ -31,6 +33,12 @@ function OAuthDs({ // Determine button text based on connector state and provider const getButtonText = () => { + if ( + oAuthProvider === O_AUTH_PROVIDERS.OPENAI && + hasOAuthCredentials + ) { + return "Authenticated"; + } if (isExistingConnector) { return "Reauthenticate"; } @@ -341,7 +349,7 @@ function OAuthDs({ disabled={disabled} verificationUrl={deviceLogin?.verification_url} userCode={deviceLogin?.user_code} - accountLabel={deviceLogin?.account_label} + accountLabel={deviceLogin?.account_label || oauthAccountLabel} /> ); } @@ -355,6 +363,8 @@ OAuthDs.propTypes = { setStatus: PropTypes.func, selectedSourceId: PropTypes.string.isRequired, isExistingConnector: PropTypes.bool, + hasOAuthCredentials: PropTypes.bool, + oauthAccountLabel: PropTypes.string, adapterInstanceId: PropTypes.string, onModelsLoaded: PropTypes.func, disabled: PropTypes.bool, diff --git a/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx b/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx index 6a23f73ea5..eff11a5285 100644 --- a/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx +++ b/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx @@ -41,7 +41,7 @@ const OpenAIOAuthButton = ({ and enter code {userCode}. )} - {status === "success" && accountLabel && ( + {accountLabel && ( {accountLabel} diff --git a/frontend/src/hooks/useExceptionHandler.jsx b/frontend/src/hooks/useExceptionHandler.jsx index eef7dbfc7b..227573317d 100644 --- a/frontend/src/hooks/useExceptionHandler.jsx +++ b/frontend/src/hooks/useExceptionHandler.jsx @@ -1,3 +1,4 @@ +import { useCallback } from "react"; import { useNavigate } from "react-router-dom"; import { getRequestIdFromError } from "../helpers/requestId"; @@ -5,7 +6,7 @@ import { getRequestIdFromError } from "../helpers/requestId"; const useExceptionHandler = () => { const navigate = useNavigate(); - const handleException = ( + const handleException = useCallback(( err, errMessage = "Something went wrong", setBackendErrors = undefined, @@ -99,7 +100,7 @@ const useExceptionHandler = () => { } else { return alert(errMessage); } - }; + }, [navigate]); return handleException; }; diff --git a/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py b/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py index 991bb48d72..d40cd3de14 100644 --- a/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py +++ b/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py @@ -162,15 +162,19 @@ def _reasoning_levels(value: object) -> list[dict[str, str]]: levels: list[dict[str, str]] = [] seen: set[str] = set() for item in value: - if not isinstance(item, Mapping): + if isinstance(item, str): + effort = item + description = item + elif isinstance(item, Mapping): + effort = item.get("effort") or item.get("reasoning_effort") + description = item.get("description") + else: continue - effort = item.get("effort") or item.get("reasoning_effort") if not isinstance(effort, str) or not effort.strip(): continue effort = effort.strip() if effort in seen: continue - description = item.get("description") if not isinstance(description, str) or not description.strip(): description = effort.replace("_", " ").replace("-", " ").title() levels.append({"effort": effort, "description": description.strip()}) @@ -346,24 +350,12 @@ def _selected_catalog_model(slugs: Sequence[str], current_model: str | None) -> def _reasoning_schema_for_model( model: Mapping[str, Any], ) -> dict[str, Any] | None: - levels_value = model.get("supported_reasoning_levels") - if not isinstance(levels_value, list): - return None - levels = [ - dict(level) - for level in levels_value - if isinstance(level, Mapping) - and isinstance(level.get("effort"), str) - and level["effort"].strip() - ] + levels = _reasoning_levels(model.get("supported_reasoning_levels")) if not levels: return None - efforts = [str(level["effort"]).strip() for level in levels] - effort_labels = [ - str(level.get("description") or effort).strip() - for level, effort in zip(levels, efforts, strict=True) - ] + efforts = [level["effort"] for level in levels] + effort_labels = [level["description"] for level in levels] default_effort = model.get("default_reasoning_level") if default_effort not in efforts: default_effort = efforts[0] diff --git a/unstract/sdk1/tests/test_openai_oauth.py b/unstract/sdk1/tests/test_openai_oauth.py index 8a34889bea..34c9061da6 100644 --- a/unstract/sdk1/tests/test_openai_oauth.py +++ b/unstract/sdk1/tests/test_openai_oauth.py @@ -163,6 +163,21 @@ def test_openai_oauth_dynamic_schema_uses_model_specific_reasoning_levels() -> N ] +def test_openai_oauth_dynamic_schema_accepts_string_reasoning_levels() -> None: + schema = build_openai_oauth_json_schema( + [ + { + "slug": "account-model", + "supported_reasoning_levels": ["low", "medium", "high"], + } + ], + base_schema=json.loads(OpenAIOAuthLLMAdapter.get_json_schema()), + ) + + reasoning = schema["allOf"][1]["then"]["properties"]["reasoning_effort"] + assert reasoning["enum"] == ["low", "medium", "high"] + + def test_openai_oauth_parameters_normalize_model_and_fix_endpoint() -> None: metadata = _metadata("access", "account") metadata["model"] = "openai/gpt-5-codex" From 3f06f6044294be83c49cec92f60277c12bb181f8 Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 17:55:36 -0400 Subject: [PATCH 11/19] Make OpenAI device login recoverable --- .../input-output/configure-ds/ConfigureDs.jsx | 8 +- .../components/oauth-ds/oauth-ds/OAuthDs.jsx | 92 ++++++++++++++++--- .../oauth-ds/openai/OpenAIOAuthButton.jsx | 15 ++- 3 files changed, 95 insertions(+), 20 deletions(-) diff --git a/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx b/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx index 222901ca65..bc0f4e4f6a 100644 --- a/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx +++ b/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx @@ -52,8 +52,10 @@ function ConfigureDs({ } = usePostHogEvents(); const { getUrl } = useRequestUrl(); - const oauthCacheKey = `oauth-cachekey-${selectedSourceId}`; - const oauthStatusKey = `oauth-status-${selectedSourceId}`; + const oauthStateScope = editItemId || selectedSourceId; + const oauthCacheKey = `oauth-cachekey-${oauthStateScope}`; + const oauthStatusKey = `oauth-status-${oauthStateScope}`; + const oauthDeviceKey = `oauth-device-${oauthStateScope}`; const findReasoningSchema = useCallback((schema, model) => { if (!schema || !model || !Array.isArray(schema.allOf)) { @@ -424,7 +426,7 @@ function ConfigureDs({ if (oAuthProvider?.length > 0 && isOAuthMethodSelected()) { localStorage.removeItem(oauthCacheKey); localStorage.removeItem(oauthStatusKey); - localStorage.removeItem(`oauth-device-${selectedSourceId}`); + localStorage.removeItem(oauthDeviceKey); } setOpen(false); diff --git a/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx b/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx index 5a7aac4767..559a916b8a 100644 --- a/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx +++ b/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx @@ -27,9 +27,13 @@ function OAuthDs({ const { setAlertDetails } = useAlertStore(); const handleException = useExceptionHandler(); - // Simple OAuth storage keys per connector - const oauthCacheKey = `oauth-cachekey-${selectedSourceId}`; - const oauthStatusKey = `oauth-status-${selectedSourceId}`; + // Keep transient OAuth hand-off state isolated per saved adapter. New + // adapters use the provider id until they are saved; saved adapters use + // their instance id so multiple ChatGPT accounts never share browser state. + const oauthStateScope = adapterInstanceId || selectedSourceId; + const oauthCacheKey = `oauth-cachekey-${oauthStateScope}`; + const oauthStatusKey = `oauth-status-${oauthStateScope}`; + const oauthDeviceKey = `oauth-device-${oauthStateScope}`; // Determine button text based on connector state and provider const getButtonText = () => { @@ -57,20 +61,24 @@ function OAuthDs({ const buttonText = getButtonText(); const [oauthStatus, setOAuthStatus] = useState(() => { - // Initialize from connector-specific status - return localStorage.getItem(oauthStatusKey); + // A durable OAuth adapter is authenticated even when there is no + // browser-side hand-off session left to restore. + return hasOAuthCredentials + ? "success" + : localStorage.getItem(oauthStatusKey); }); const [loginCacheKey, setLoginCacheKey] = useState(() => localStorage.getItem(oauthCacheKey), ); const [deviceLogin, setDeviceLogin] = useState(() => { try { - return JSON.parse(localStorage.getItem(`oauth-device-${selectedSourceId}`)); + return JSON.parse(localStorage.getItem(oauthDeviceKey)); } catch { return null; } }); const [activeLoginCacheKey, setActiveLoginCacheKey] = useState(null); + const [isStarting, setIsStarting] = useState(false); const updateOAuthStatus = useCallback((newStatus) => { setOAuthStatus(newStatus); @@ -128,14 +136,12 @@ function OAuthDs({ setStatus(connectorStatus || ""); setOAuthStatus(connectorStatus || ""); - const persistedDeviceLogin = localStorage.getItem( - `oauth-device-${selectedSourceId}`, - ); + const persistedDeviceLogin = localStorage.getItem(oauthDeviceKey); if (persistedDeviceLogin) { try { setDeviceLogin(JSON.parse(persistedDeviceLogin)); } catch { - localStorage.removeItem(`oauth-device-${selectedSourceId}`); + localStorage.removeItem(oauthDeviceKey); } } else { setDeviceLogin(null); @@ -145,7 +151,38 @@ function OAuthDs({ window.removeEventListener("storage", handleStorageChange); // Don't clear localStorage on unmount to persist across tab switches }; - }, [selectedSourceId, oauthCacheKey, oauthStatusKey, setCacheKey, setStatus]); + }, [ + oauthCacheKey, + oauthDeviceKey, + oauthStatusKey, + selectedSourceId, + setCacheKey, + setStatus, + ]); + + useEffect(() => { + // Do not let an old pending browser state mask credentials already saved + // on an existing adapter. A live reauthentication session is allowed to + // remain pending and can still replace those credentials after testing. + if ( + oAuthProvider !== O_AUTH_PROVIDERS.OPENAI || + !hasOAuthCredentials || + loginCacheKey || + activeLoginCacheKey || + oauthStatus === "success" + ) { + return; + } + setOAuthStatus("success"); + setStatus("success"); + }, [ + activeLoginCacheKey, + hasOAuthCredentials, + loginCacheKey, + oAuthProvider, + oauthStatus, + setStatus, + ]); useEffect(() => { if (oAuthProvider !== O_AUTH_PROVIDERS.OPENAI || !onModelsLoaded) { @@ -247,6 +284,7 @@ function OAuthDs({ ]); const handleOAuth = async () => { + let loginWindow; try { if (oAuthProvider === O_AUTH_PROVIDERS.OPENAI) { if (oauthStatus === "pending" && deviceLogin?.verification_url) { @@ -258,6 +296,15 @@ function OAuthDs({ return; } + setIsStarting(true); + // Open a user-initiated window before awaiting the API call. Browsers + // may block a window opened only after the device-code request returns; + // the visible link below remains the fallback when that happens. + loginWindow = window.open( + "about:blank", + "_blank", + "toolbar=yes,scrollbars=yes,resizable=yes,top=200,left=500,width=500,height=600", + ); const response = await axiosPrivate({ method: "POST", url: "/api/v1/oauth/openai/start/", @@ -276,10 +323,19 @@ function OAuthDs({ localStorage.setItem(oauthCacheKey, newCacheKey); setDeviceLogin(loginDetails); localStorage.setItem( - `oauth-device-${selectedSourceId}`, + oauthDeviceKey, JSON.stringify(loginDetails), ); updateOAuthStatus("pending"); + if (loginWindow && loginDetails.verification_url) { + loginWindow.location.href = loginDetails.verification_url; + } else if (loginDetails.verification_url) { + window.open( + loginDetails.verification_url, + "_blank", + "toolbar=yes,scrollbars=yes,resizable=yes,top=200,left=500,width=500,height=600", + ); + } return; } @@ -310,7 +366,18 @@ function OAuthDs({ "toolbar=yes,scrollbars=yes,resizable=yes,top=200,left=500,width=500,height=600", ); } catch (err) { + if (loginWindow && !loginWindow.closed) { + loginWindow.close(); + } + if (oAuthProvider === O_AUTH_PROVIDERS.OPENAI) { + setIsStarting(false); + updateOAuthStatus("error"); + } setAlertDetails(handleException(err)); + } finally { + if (oAuthProvider === O_AUTH_PROVIDERS.OPENAI) { + setIsStarting(false); + } } }; @@ -350,6 +417,7 @@ function OAuthDs({ verificationUrl={deviceLogin?.verification_url} userCode={deviceLogin?.user_code} accountLabel={deviceLogin?.account_label || oauthAccountLabel} + isStarting={isStarting} /> ); } diff --git a/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx b/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx index eff11a5285..3f494b1a4e 100644 --- a/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx +++ b/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx @@ -12,22 +12,26 @@ const OpenAIOAuthButton = ({ verificationUrl, userCode, accountLabel, + isStarting = false, }) => { const isPending = status === "pending"; - const buttonLabel = - status === "success" + const buttonLabel = isStarting + ? "Starting OpenAI sign-in" + : status === "success" ? "Authenticated" : isPending && verificationUrl ? "Open OpenAI device login" - : buttonText; + : isPending + ? "Retry OpenAI sign-in" + : buttonText; return (
@@ -58,6 +62,7 @@ OpenAIOAuthButton.propTypes = { verificationUrl: PropTypes.string, userCode: PropTypes.string, accountLabel: PropTypes.string, + isStarting: PropTypes.bool, }; export default OpenAIOAuthButton; From 6bc35fa94ac42bd58c348b1f038b14dd81d73fdf Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 19:55:33 -0400 Subject: [PATCH 12/19] Fix persisted OAuth reasoning control --- .../input-output/configure-ds/ConfigureDs.jsx | 44 ++++++----- .../configure-ds/openai-oauth-form-schema.js | 78 +++++++++++++++++++ .../openai-oauth-form-schema.test.js | 77 ++++++++++++++++++ .../rjsf-form-layout/RjsfFormLayout.jsx | 9 +++ 4 files changed, 190 insertions(+), 18 deletions(-) create mode 100644 frontend/src/components/input-output/configure-ds/openai-oauth-form-schema.js create mode 100644 frontend/src/components/input-output/configure-ds/openai-oauth-form-schema.test.js diff --git a/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx b/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx index bc0f4e4f6a..7417f8ba2c 100644 --- a/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx +++ b/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx @@ -1,6 +1,6 @@ import { Info } from "lucide-react"; import PropTypes from "prop-types"; -import { createRef, useCallback, useEffect, useState } from "react"; +import { createRef, useCallback, useEffect, useMemo, useState } from "react"; import { Col, Row } from "@/components/ui/shims/antd-layout"; import { Popover } from "@/components/ui/shims/antd-overlays"; @@ -13,6 +13,10 @@ import { useAlertStore } from "../../../store/alert-store"; import { useSessionStore } from "../../../store/session-store"; import { OAuthDs } from "../../oauth-ds/oauth-ds/OAuthDs.jsx"; import { CustomButton } from "../../widgets/custom-button/CustomButton.jsx"; +import { + getReasoningSchemaForModel, + materializeReasoningProperty, +} from "./openai-oauth-form-schema.js"; import "./ConfigureDs.css"; function ConfigureDs({ @@ -57,18 +61,6 @@ function ConfigureDs({ const oauthStatusKey = `oauth-status-${oauthStateScope}`; const oauthDeviceKey = `oauth-device-${oauthStateScope}`; - const findReasoningSchema = useCallback((schema, model) => { - if (!schema || !model || !Array.isArray(schema.allOf)) { - return undefined; - } - const condition = schema.allOf.find( - (candidate) => - candidate?.if?.properties?.enable_reasoning?.const === true && - candidate?.if?.properties?.model?.const === model, - ); - return condition?.then?.properties?.reasoning_effort; - }, []); - const handleOpenAIModelSchema = useCallback( (dynamicSchema) => { const modelSchema = dynamicSchema?.properties?.model; @@ -85,7 +77,10 @@ function ConfigureDs({ } if (next.enable_reasoning) { - const reasoningSchema = findReasoningSchema(dynamicSchema, next.model); + const reasoningSchema = getReasoningSchemaForModel( + dynamicSchema, + next.model, + ); if ( Array.isArray(reasoningSchema?.enum) && reasoningSchema.enum.length > 0 && @@ -98,7 +93,17 @@ function ConfigureDs({ return next; }); }, - [findReasoningSchema, setFormData], + [setFormData], + ); + + const renderSchema = useMemo( + () => + materializeReasoningProperty( + oauthSpec || spec, + formData?.model, + formData?.enable_reasoning, + ), + [formData?.enable_reasoning, formData?.model, oauthSpec, spec], ); useEffect(() => { @@ -109,7 +114,10 @@ function ConfigureDs({ if (!oauthSpec || !formData?.enable_reasoning || !formData?.model) { return; } - const reasoningSchema = findReasoningSchema(oauthSpec, formData.model); + const reasoningSchema = getReasoningSchemaForModel( + oauthSpec, + formData.model, + ); if ( !Array.isArray(reasoningSchema?.enum) || reasoningSchema.enum.length === 0 || @@ -121,7 +129,7 @@ function ConfigureDs({ ...(currentFormData || {}), reasoning_effort: reasoningSchema.default || reasoningSchema.enum[0], })); - }, [findReasoningSchema, formData, oauthSpec, setFormData]); + }, [formData, oauthSpec, setFormData]); // Determine if this is a new or existing connector const hasOAuthCredentials = @@ -474,7 +482,7 @@ function ConfigureDs({
)} = 0 ? modelValues[labelIndex] : model; +} + +function getReasoningSchemaForModel(schema, model) { + if (!schema || !model || !Array.isArray(schema.allOf)) { + return undefined; + } + + const selectedModel = canonicalModelValue(schema, model); + const condition = schema.allOf.find( + (candidate) => + candidate?.if?.properties?.enable_reasoning?.const === true && + candidate?.if?.properties?.model?.const === selectedModel, + ); + return condition?.then?.properties?.reasoning_effort; +} + +/** + * RJSF can retain a resolved allOf branch when an account schema replaces the + * pre-auth schema while the checkbox is already enabled. Put the selected + * account's live field in the visible properties as well; the allOf branch + * remains in place for validation and model-specific requiredness. + */ +function materializeReasoningProperty(schema, model, enabled) { + if ( + !schema || + schema["x-openai-oauth-model-source"] !== OPENAI_OAUTH_MODEL_SOURCE || + !enabled + ) { + return schema; + } + + const reasoningSchema = getReasoningSchemaForModel(schema, model); + if ( + !reasoningSchema || + !Array.isArray(reasoningSchema.enum) || + reasoningSchema.enum.length === 0 + ) { + return schema; + } + + const sourceProperties = schema.properties || {}; + const properties = {}; + let inserted = false; + Object.entries(sourceProperties).forEach(([name, property]) => { + properties[name] = property; + if (name === "enable_reasoning") { + properties.reasoning_effort = reasoningSchema; + inserted = true; + } + }); + if (!inserted) { + properties.reasoning_effort = reasoningSchema; + } + + return { ...schema, properties }; +} + +export { + getReasoningSchemaForModel, + materializeReasoningProperty, + OPENAI_OAUTH_MODEL_SOURCE, +}; diff --git a/frontend/src/components/input-output/configure-ds/openai-oauth-form-schema.test.js b/frontend/src/components/input-output/configure-ds/openai-oauth-form-schema.test.js new file mode 100644 index 0000000000..0683aca979 --- /dev/null +++ b/frontend/src/components/input-output/configure-ds/openai-oauth-form-schema.test.js @@ -0,0 +1,77 @@ +import { describe, expect, it } from "vitest"; + +import { + getReasoningSchemaForModel, + materializeReasoningProperty, +} from "./openai-oauth-form-schema.js"; + +const schema = { + "x-openai-oauth-model-source": "chatgpt-account", + properties: { + model: { + enum: ["gpt-5.6-luna", "gpt-5.6-sol"], + enumNames: ["GPT-5.6-Luna", "GPT-5.6-Sol"], + }, + enable_reasoning: { type: "boolean" }, + }, + allOf: [ + { + if: { + properties: { + enable_reasoning: { const: true }, + model: { const: "gpt-5.6-luna" }, + }, + }, + then: { + properties: { + reasoning_effort: { + type: "string", + enum: ["low", "medium", "high", "xhigh", "max"], + }, + }, + }, + }, + ], +}; + +describe("OpenAI OAuth form schema", () => { + it("resolves reasoning by live model slug or display label", () => { + expect(getReasoningSchemaForModel(schema, "gpt-5.6-luna").enum).toEqual([ + "low", + "medium", + "high", + "xhigh", + "max", + ]); + expect(getReasoningSchemaForModel(schema, "GPT-5.6-Luna").enum).toEqual([ + "low", + "medium", + "high", + "xhigh", + "max", + ]); + }); + + it("materializes the selected account options after OAuth schema load", () => { + const rendered = materializeReasoningProperty(schema, "gpt-5.6-luna", true); + + expect(Object.keys(rendered.properties)).toEqual([ + "model", + "enable_reasoning", + "reasoning_effort", + ]); + expect(rendered.properties.reasoning_effort.enum).toEqual([ + "low", + "medium", + "high", + "xhigh", + "max", + ]); + }); + + it("does not add the field while reasoning is disabled", () => { + expect(materializeReasoningProperty(schema, "gpt-5.6-luna", false)).toBe( + schema, + ); + }); +}); diff --git a/frontend/src/layouts/rjsf-form-layout/RjsfFormLayout.jsx b/frontend/src/layouts/rjsf-form-layout/RjsfFormLayout.jsx index 3855513500..548cf5a08b 100644 --- a/frontend/src/layouts/rjsf-form-layout/RjsfFormLayout.jsx +++ b/frontend/src/layouts/rjsf-form-layout/RjsfFormLayout.jsx @@ -276,6 +276,14 @@ function RjsfFormLayout({ }; }, [formSchema]); + // An OAuth schema arrives after the form has already mounted. Include the + // materialized reasoning field in the key so RJSF does not retain the + // pre-auth resolved allOf branch when reasoning was already enabled. + const formInstanceKey = + schema?.["x-openai-oauth-model-source"] === "chatgpt-account" + ? `openai-oauth-${Boolean(schema?.properties?.reasoning_effort)}` + : undefined; + return ( <> {isLoading ? ( @@ -289,6 +297,7 @@ function RjsfFormLayout({ /> )}
Date: Thu, 3 Sep 2026 20:11:22 -0400 Subject: [PATCH 13/19] Support streaming-only OpenAI OAuth responses --- unstract/sdk1/src/unstract/sdk1/llm.py | 151 +++++++++++++++++------ unstract/sdk1/tests/test_openai_oauth.py | 38 +++++- 2 files changed, 146 insertions(+), 43 deletions(-) diff --git a/unstract/sdk1/src/unstract/sdk1/llm.py b/unstract/sdk1/src/unstract/sdk1/llm.py index e3f1aa9ba4..f32a89901a 100644 --- a/unstract/sdk1/src/unstract/sdk1/llm.py +++ b/unstract/sdk1/src/unstract/sdk1/llm.py @@ -726,6 +726,104 @@ def integer(key: str, fallback: str) -> int: "total_tokens": integer("total_tokens", "total_tokens"), } + def _collect_openai_oauth_response( + self, + messages: list[dict[str, object]], + completion_kwargs: dict[str, object], + max_retries: int, + ) -> tuple[str | None, object | None, dict[str, int]]: + """Consume a streaming-only Responses API call as one completion. + + The ChatGPT/Codex endpoint rejects ``stream=False``. Keep the public + ``complete()`` contract by consuming the stream internally and + returning the assembled text, completed response, and usage. + """ + response_kwargs = self._build_openai_oauth_responses_kwargs( + messages, completion_kwargs, stream=True + ) + text_parts: list[str] = [] + completed_response: object | None = None + last_event: object | None = None + + for event in iter_with_retry( + lambda: litellm.responses(**response_kwargs), + max_retries=max_retries, + retry_predicate=is_retryable_litellm_error, + description=self._get_adapter_info(), + ): + last_event = event + event_type = self._response_value(event, "type") + if event_type == "response.output_text.delta": + text = self._response_value(event, "delta", "") + if isinstance(text, str): + text_parts.append(text) + elif event_type == "response.completed": + response = self._response_value(event, "response") + if response is not None: + completed_response = response + + response = ( + completed_response if completed_response is not None else last_event + ) + response_text = "".join(text_parts) or self._responses_output_text(response) + usage_source = ( + completed_response if completed_response is not None else response + ) + usage = self._responses_usage(self._response_value(usage_source, "usage")) + return response_text, response, usage + + async def _acollect_openai_oauth_response( + self, + messages: list[dict[str, object]], + completion_kwargs: dict[str, object], + max_retries: int, + ) -> tuple[str | None, object | None, dict[str, int]]: + """Async counterpart to :meth:`_collect_openai_oauth_response`.""" + response_kwargs = self._build_openai_oauth_responses_kwargs( + messages, completion_kwargs, stream=True + ) + + async def consume_stream() -> tuple[ + str | None, object | None, dict[str, int] + ]: + text_parts: list[str] = [] + completed_response: object | None = None + last_event: object | None = None + stream = await litellm.aresponses(**response_kwargs) + + async for event in stream: + last_event = event + event_type = self._response_value(event, "type") + if event_type == "response.output_text.delta": + text = self._response_value(event, "delta", "") + if isinstance(text, str): + text_parts.append(text) + elif event_type == "response.completed": + response = self._response_value(event, "response") + if response is not None: + completed_response = response + + response = ( + completed_response if completed_response is not None else last_event + ) + response_text = "".join(text_parts) or self._responses_output_text( + response + ) + usage_source = ( + completed_response if completed_response is not None else response + ) + usage = self._responses_usage( + self._response_value(usage_source, "usage") + ) + return response_text, response, usage + + return await acall_with_retry( + consume_stream, + max_retries=max_retries, + retry_predicate=is_retryable_litellm_error, + description=self._get_adapter_info(), + ) + def _stream_openai_oauth( self, messages: list[dict[str, object]], @@ -815,21 +913,12 @@ def complete( completion_kwargs, self._get_adapter_info() ) if self._uses_openai_oauth(): - response = call_with_retry( - lambda: litellm.responses( - **self._build_openai_oauth_responses_kwargs( - messages, completion_kwargs, stream=False - ) - ), - max_retries=max_retries, - retry_predicate=is_retryable_litellm_error, - description=self._get_adapter_info(), + response_text, response, usage = self._collect_openai_oauth_response( + messages, + completion_kwargs, + max_retries, ) - response_text = self._responses_output_text(response) finish_reason = None - usage = self._responses_usage( - self._response_value(response, "usage") - ) else: response = call_with_retry( lambda: litellm.completion(messages=messages, **completion_kwargs), @@ -952,21 +1041,12 @@ def complete_vision( max_retries = pop_litellm_retry_kwargs( completion_kwargs, self._get_adapter_info() ) - response = call_with_retry( - lambda: litellm.responses( - **self._build_openai_oauth_responses_kwargs( - messages, completion_kwargs, stream=False - ) - ), - max_retries=max_retries, - retry_predicate=is_retryable_litellm_error, - description=self._get_adapter_info(), + response_text, response, usage = self._collect_openai_oauth_response( + messages, + completion_kwargs, + max_retries, ) - response_text = self._responses_output_text(response) finish_reason = None - usage = self._responses_usage( - self._response_value(response, "usage") - ) else: response = litellm.completion( messages=messages, @@ -1128,21 +1208,14 @@ async def acomplete( completion_kwargs, self._get_adapter_info() ) if self._uses_openai_oauth(): - response = await acall_with_retry( - lambda: litellm.aresponses( - **self._build_openai_oauth_responses_kwargs( - messages, completion_kwargs, stream=False - ) - ), - max_retries=max_retries, - retry_predicate=is_retryable_litellm_error, - description=self._get_adapter_info(), + response_text, response, usage = await ( + self._acollect_openai_oauth_response( + messages, + completion_kwargs, + max_retries, + ) ) - response_text = self._responses_output_text(response) finish_reason = None - usage = self._responses_usage( - self._response_value(response, "usage") - ) else: response = await acall_with_retry( lambda: litellm.acompletion(messages=messages, **completion_kwargs), diff --git a/unstract/sdk1/tests/test_openai_oauth.py b/unstract/sdk1/tests/test_openai_oauth.py index 34c9061da6..472d93dfb9 100644 --- a/unstract/sdk1/tests/test_openai_oauth.py +++ b/unstract/sdk1/tests/test_openai_oauth.py @@ -46,6 +46,22 @@ def _metadata(access_token: str, account_id: str) -> dict[str, object]: } +def _response_events(text: str) -> object: + return iter( + [ + {"type": "response.output_text.delta", "delta": text}, + { + "type": "response.completed", + "response": { + "id": f"response-{text}", + "output_text": text, + "usage": {}, + }, + }, + ] + ) + + def test_openai_oauth_adapter_is_registered_with_auth_metadata() -> None: adapter_id = OpenAIOAuthLLMAdapter.get_id() assert adapters[adapter_id][Common.MODULE] is OpenAIOAuthLLMAdapter @@ -253,8 +269,8 @@ def test_llm_sends_each_account_token_and_workspace_header() -> None: patch( "unstract.sdk1.llm.litellm.responses", side_effect=[ - {"output_text": "first", "usage": {}}, - {"output_text": "second", "usage": {}}, + _response_events("first"), + _response_events("second"), ], ) as responses, patch.object(first, "_record_usage"), @@ -275,6 +291,8 @@ def test_llm_sends_each_account_token_and_workspace_header() -> None: assert first_request["include"] == ["reasoning.encrypted_content"] assert "oauth_refresh_token" not in first_request assert first_request["api_base"] == OPENAI_OAUTH_CHATGPT_API_BASE + assert first_request["stream"] is True + assert second_request["stream"] is True def test_llm_streams_responses_text_and_records_completion_usage() -> None: @@ -315,12 +333,23 @@ def test_llm_async_completion_uses_responses_api() -> None: adapter_id=OpenAIOAuthLLMAdapter.get_id(), adapter_metadata=_metadata("token", "account"), ) - response = {"output_text": "async response", "usage": {}} + async def response_events(): + yield { + "type": "response.output_text.delta", + "delta": "async response", + } + yield { + "type": "response.completed", + "response": { + "output_text": "async response", + "usage": {}, + }, + } with ( patch( "unstract.sdk1.llm.litellm.aresponses", - new=AsyncMock(return_value=response), + new=AsyncMock(return_value=response_events()), ) as responses, patch.object(llm, "_record_usage"), ): @@ -328,3 +357,4 @@ def test_llm_async_completion_uses_responses_api() -> None: assert result["response"].text == "async response" assert responses.await_args.kwargs["extra_headers"]["ChatGPT-Account-Id"] == "account" + assert responses.await_args.kwargs["stream"] is True From 55b35b448bbd5d42d3eb80baece2a20e6f7ef30f Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 20:21:55 -0400 Subject: [PATCH 14/19] Show reasoning effort names in OAuth form --- unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py | 6 ++++-- unstract/sdk1/tests/test_openai_oauth.py | 4 ++++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py b/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py index d40cd3de14..76e07e08f5 100644 --- a/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py +++ b/unstract/sdk1/src/unstract/sdk1/auth/openai_oauth.py @@ -355,14 +355,16 @@ def _reasoning_schema_for_model( return None efforts = [level["effort"] for level in levels] - effort_labels = [level["description"] for level in levels] default_effort = model.get("default_reasoning_level") if default_effort not in efforts: default_effort = efforts[0] return { "type": "string", "enum": efforts, - "enumNames": effort_labels, + # RJSF renders enumNames in the select. Keep the stable effort names + # visible; the catalog descriptions are explanatory metadata, not the + # value users need to choose or persist. + "enumNames": efforts, "default": default_effort, "title": "Reasoning Effort", "description": "Reasoning levels reported for this model by ChatGPT.", diff --git a/unstract/sdk1/tests/test_openai_oauth.py b/unstract/sdk1/tests/test_openai_oauth.py index 472d93dfb9..3f5add7afa 100644 --- a/unstract/sdk1/tests/test_openai_oauth.py +++ b/unstract/sdk1/tests/test_openai_oauth.py @@ -174,6 +174,9 @@ def test_openai_oauth_dynamic_schema_uses_model_specific_reasoning_levels() -> N "low", "high", ] + assert model_conditions[0]["then"]["properties"]["reasoning_effort"][ + "enumNames" + ] == ["low", "high"] assert model_conditions[1]["then"]["properties"]["reasoning_effort"]["enum"] == [ "max" ] @@ -192,6 +195,7 @@ def test_openai_oauth_dynamic_schema_accepts_string_reasoning_levels() -> None: reasoning = schema["allOf"][1]["then"]["properties"]["reasoning_effort"] assert reasoning["enum"] == ["low", "medium", "high"] + assert reasoning["enumNames"] == ["low", "medium", "high"] def test_openai_oauth_parameters_normalize_model_and_fix_endpoint() -> None: From 7b959218f3182a4075e111538006cfa7ac80417e Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 20:22:51 -0400 Subject: [PATCH 15/19] Do not forward unsupported Codex output cap --- .../adapters/llm1/static/openai-oauth.json | 2 +- unstract/sdk1/src/unstract/sdk1/llm.py | 8 +++++--- unstract/sdk1/tests/test_openai_oauth.py | 18 ++++++++++++++++++ 3 files changed, 24 insertions(+), 4 deletions(-) diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json index f03410d88c..c8938bccf6 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json +++ b/unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/openai-oauth.json @@ -25,7 +25,7 @@ "multipleOf": 1, "title": "Maximum Output Tokens", "default": 4096, - "description": "Maximum number of output tokens to request." + "description": "Compatibility setting retained by Unstract; the ChatGPT/Codex subscription endpoint uses its provider-managed output budget." }, "max_retries": { "type": "number", diff --git a/unstract/sdk1/src/unstract/sdk1/llm.py b/unstract/sdk1/src/unstract/sdk1/llm.py index f32a89901a..a21561b072 100644 --- a/unstract/sdk1/src/unstract/sdk1/llm.py +++ b/unstract/sdk1/src/unstract/sdk1/llm.py @@ -629,7 +629,11 @@ def _build_openai_oauth_responses_kwargs( model = str(values.pop("model")) api_base = str(values.pop("api_base", OPENAI_OAUTH_CHATGPT_API_BASE)) - max_tokens = values.pop("max_tokens", None) + # The ChatGPT/Codex subscription endpoint rejects both the legacy + # max_tokens option and the Responses API max_output_tokens option. + # Keep the schema value for adapter compatibility, but do not put it + # on the provider request. + values.pop("max_tokens", None) values.pop("temperature", None) values.pop("n", None) values.pop("api_version", None) @@ -639,8 +643,6 @@ def _build_openai_oauth_responses_kwargs( values.pop("cost_model", None) values.pop("context_window", None) - if max_tokens is not None: - values["max_output_tokens"] = max_tokens reasoning_effort = values.pop("reasoning_effort", None) # The ChatGPT Codex endpoint expects this field on Responses requests # so encrypted reasoning can be returned and reused by the account. diff --git a/unstract/sdk1/tests/test_openai_oauth.py b/unstract/sdk1/tests/test_openai_oauth.py index 3f5add7afa..bfe1bc10f2 100644 --- a/unstract/sdk1/tests/test_openai_oauth.py +++ b/unstract/sdk1/tests/test_openai_oauth.py @@ -210,6 +210,22 @@ def test_openai_oauth_parameters_normalize_model_and_fix_endpoint() -> None: assert metadata["model"] == "openai/gpt-5-codex" +def test_openai_oauth_does_not_forward_unsupported_output_token_parameters() -> None: + llm = LLM( + adapter_id=OpenAIOAuthLLMAdapter.get_id(), + adapter_metadata=_metadata("access", "account"), + ) + + request = llm._build_openai_oauth_responses_kwargs( + [], + {**_metadata("access", "account"), "max_tokens": 4096}, + stream=True, + ) + + assert "max_tokens" not in request + assert "max_output_tokens" not in request + + @pytest.mark.parametrize( "field", ["oauth_access_token", "oauth_refresh_token", "oauth_account_id"], @@ -297,6 +313,8 @@ def test_llm_sends_each_account_token_and_workspace_header() -> None: assert first_request["api_base"] == OPENAI_OAUTH_CHATGPT_API_BASE assert first_request["stream"] is True assert second_request["stream"] is True + assert "max_tokens" not in first_request + assert "max_output_tokens" not in first_request def test_llm_streams_responses_text_and_records_completion_usage() -> None: From eee9ab387043642188553a167ec4a9df4348bdbf Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 20:55:36 -0400 Subject: [PATCH 16/19] Recover from expired OpenAI OAuth handoff --- .../components/oauth-ds/oauth-ds/OAuthDs.jsx | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx b/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx index 559a916b8a..ca3abf424d 100644 --- a/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx +++ b/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx @@ -210,6 +210,27 @@ function OAuthDs({ const message = err?.response?.data?.message || "Could not load models available to this OpenAI account"; + + // The hand-off key is intentionally short-lived and is also removed + // after credentials are saved. Do not keep treating an expired key + // from localStorage as an authenticated new login. Saved adapters can + // fall back to their durable credentials on the next effect pass; + // unsaved adapters are returned to the sign-in state. + const isExpiredLoginSession = + sessionKey && + typeof message === "string" && + message.toLowerCase().includes("openai oauth login session"); + if (isExpiredLoginSession) { + setLoginCacheKey(null); + setActiveLoginCacheKey(null); + setDeviceLogin(null); + setCacheKey(""); + setStatus(""); + setOAuthStatus(""); + localStorage.removeItem(oauthCacheKey); + localStorage.removeItem(oauthStatusKey); + localStorage.removeItem(oauthDeviceKey); + } setAlertDetails(handleException(err, message)); }); @@ -226,6 +247,8 @@ function OAuthDs({ oAuthProvider, onModelsLoaded, setAlertDetails, + setCacheKey, + setStatus, ]); useEffect(() => { From 557927d9869efe3c9954957584c803906bf337c2 Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 21:11:07 -0400 Subject: [PATCH 17/19] Persist OpenAI OAuth form drafts --- .../add-source-modal/AddSourceModal.jsx | 28 ++++++++- .../input-output/add-source/AddSource.jsx | 57 ++++++++++++++++++- .../input-output/configure-ds/ConfigureDs.css | 7 +++ .../input-output/configure-ds/ConfigureDs.jsx | 23 +++++--- frontend/src/components/onboard/OnBoard.jsx | 49 ++++++++++++---- frontend/src/components/onboard/onBoard.css | 23 ++++++++ 6 files changed, 162 insertions(+), 25 deletions(-) diff --git a/frontend/src/components/input-output/add-source-modal/AddSourceModal.jsx b/frontend/src/components/input-output/add-source-modal/AddSourceModal.jsx index 398026fa9d..cf7d9faaf9 100644 --- a/frontend/src/components/input-output/add-source-modal/AddSourceModal.jsx +++ b/frontend/src/components/input-output/add-source-modal/AddSourceModal.jsx @@ -1,6 +1,6 @@ import { ArrowLeft } from "lucide-react"; import PropTypes from "prop-types"; -import { useEffect, useState } from "react"; +import { useCallback, useEffect, useState } from "react"; import { Button } from "@/components/ui/shims/antd-button"; import { Modal } from "@/components/ui/shims/antd-overlays"; @@ -28,6 +28,7 @@ function AddSourceModal({ const [titles, setTitles] = useState({}); const [selectedSourceName, setSelectedSourceName] = useState(""); const [selectedDocUrl, setSelectedDocUrl] = useState(""); + const [draftsBySourceId, setDraftsBySourceId] = useState({}); const { setAlertDetails } = useAlertStore(); const axiosPrivate = useAxiosPrivate(); const handleException = useExceptionHandler(); @@ -35,6 +36,21 @@ function AddSourceModal({ const [isLoading, setIsLoading] = useState(false); const [sourcesList, setSourcesList] = useState([]); + const updateSourceDraft = useCallback((sourceId, draft) => { + setDraftsBySourceId((currentDrafts) => ({ + ...currentDrafts, + [sourceId]: draft, + })); + }, []); + + const clearSourceDraft = useCallback((sourceId) => { + setDraftsBySourceId((currentDrafts) => { + const nextDrafts = { ...currentDrafts }; + delete nextDrafts[sourceId]; + return nextDrafts; + }); + }, []); + const disabledIdsByType = { EMBEDDING: ["huggingface|90ec9ec2-1768-4d69-8fb1-c88b95de5e5a"], LLM: ["replicate|2715ce84-05af-4ab4-b8e9-67ac3211b81e"], @@ -58,13 +74,16 @@ function AddSourceModal({ useEffect(() => { if (!open) { - setTimeout(() => { + const resetTimeout = setTimeout(() => { // A delay added in order to avoid glitch in the UI when the modal is closed. setSelectedSourceId(null); setEditItemId(null); // Clear metadata to prevent stale data when adding a new connector setMetadata({}); + setDraftsBySourceId({}); }, 500); + + return () => clearTimeout(resetTimeout); } getListOfSources(); @@ -185,6 +204,7 @@ function AddSourceModal({ onCancel={() => { setOpen(false); setMetadata(null); + setDraftsBySourceId({}); }} maskClosable={false} title={modalTitle} @@ -197,6 +217,7 @@ function AddSourceModal({ > {selectedSourceId ? ( ) : isLoading ? ( diff --git a/frontend/src/components/input-output/add-source/AddSource.jsx b/frontend/src/components/input-output/add-source/AddSource.jsx index 60b1585f81..3d287f7dc0 100644 --- a/frontend/src/components/input-output/add-source/AddSource.jsx +++ b/frontend/src/components/input-output/add-source/AddSource.jsx @@ -1,7 +1,7 @@ import { getDefaultFormState } from "@rjsf/utils"; import validator from "@rjsf/validator-ajv8"; import PropTypes from "prop-types"; -import { useEffect, useMemo, useState } from "react"; +import { useEffect, useMemo, useRef, useState } from "react"; import { Typography } from "@/components/ui/shims/antd-typography"; import { useAxiosPrivate } from "../../../hooks/useAxiosPrivate"; @@ -36,6 +36,24 @@ try { // Ignore if not available } +const OPENAI_OAUTH_SOURCE_PREFIX = "openai-oauth|"; +const OPENAI_OAUTH_DRAFT_FIELDS = new Set([ + "adapter_name", + "model", + "max_tokens", + "max_retries", + "timeout", + "enable_reasoning", + "reasoning_effort", +]); + +const getOpenAIOAuthDraft = (formData) => + Object.fromEntries( + Object.entries(formData || {}).filter(([fieldName]) => + OPENAI_OAUTH_DRAFT_FIELDS.has(fieldName), + ), + ); + function AddSource({ selectedSourceId, selectedSourceName, @@ -46,9 +64,17 @@ function AddSource({ editItemId, metadata, selectedDocUrl, + draftFormData, + onDraftFormDataChange, + onClearDraft, }) { const [spec, setSpec] = useState({}); const [formData, setFormData] = useState({}); + const initialDraftFormData = useRef( + selectedSourceId.startsWith(OPENAI_OAUTH_SOURCE_PREFIX) + ? getOpenAIOAuthDraft(draftFormData) + : undefined, + ); const [isLoading, setIsLoading] = useState(false); const [oAuthProvider, setOAuthProvider] = useState(""); const { setAlertDetails } = useAlertStore(); @@ -131,6 +157,11 @@ function AddSource({ if (metadata && Object.keys(metadata).length > 0) { setFormData(metadata); + } else if ( + initialDraftFormData.current && + Object.keys(initialDraftFormData.current).length > 0 + ) { + setFormData(initialDraftFormData.current); } else { const defaults = getDefaultFormState( validator, @@ -162,6 +193,25 @@ function AddSource({ } }, [metadata]); + useEffect(() => { + if ( + !selectedSourceId.startsWith(OPENAI_OAUTH_SOURCE_PREFIX) || + editItemId?.length || + (metadata && Object.keys(metadata).length > 0) + ) { + return; + } + + onDraftFormDataChange?.(selectedSourceId, getOpenAIOAuthDraft(formData)); + }, [editItemId, formData, metadata, onDraftFormDataChange, selectedSourceId]); + + const handleAddNewItem = (row, isEdit) => { + if (!isEdit) { + onClearDraft?.(selectedSourceId); + } + addNewItem?.(row, isEdit); + }; + if (selectedSourceId.includes("pcs|")) { return ( @@ -183,7 +233,7 @@ function AddSource({ oAuthProvider={oAuthProvider} selectedSourceId={selectedSourceId} isLoading={isLoading} - addNewItem={addNewItem} + addNewItem={handleAddNewItem} type={type} editItemId={editItemId} isConnector={isConnector} @@ -204,6 +254,9 @@ AddSource.propTypes = { editItemId: PropTypes.string, metadata: PropTypes.object, selectedDocUrl: PropTypes.string, + draftFormData: PropTypes.object, + onDraftFormDataChange: PropTypes.func, + onClearDraft: PropTypes.func, }; export { AddSource }; diff --git a/frontend/src/components/input-output/configure-ds/ConfigureDs.css b/frontend/src/components/input-output/configure-ds/ConfigureDs.css index 5d44bc3fef..24498cec5a 100644 --- a/frontend/src/components/input-output/configure-ds/ConfigureDs.css +++ b/frontend/src/components/input-output/configure-ds/ConfigureDs.css @@ -36,3 +36,10 @@ color: var(--primary); cursor: pointer; } + +.config-submit-hint { + color: var(--success); + font-size: 12px; + margin-top: 10px; + text-align: center; +} diff --git a/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx b/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx index 7417f8ba2c..526636b92f 100644 --- a/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx +++ b/frontend/src/components/input-output/configure-ds/ConfigureDs.jsx @@ -202,11 +202,9 @@ function ConfigureDs({ }, [formData]); useEffect(() => { - if (!metadata) { - setFormData({}); - return; + if (metadata && Object.keys(metadata).length > 0) { + setFormData(metadata); } - setFormData(metadata); }, [selectedSourceId, metadata, setFormData]); // Clear OAuth state when switching to a different connector @@ -449,11 +447,13 @@ function ConfigureDs({ const updateSession = (type) => { const adapterType = type.toLowerCase(); - const adaptersList = sessionDetails?.adapters; - if (adaptersList && !adaptersList.includes(adapterType)) { - adaptersList.push(adapterType); - const adaptersListInSession = { adapters: adaptersList }; - updateSessionDetails(adaptersListInSession); + const adaptersList = Array.isArray(sessionDetails?.adapters) + ? sessionDetails.adapters + : []; + if (!adaptersList.includes(adapterType)) { + updateSessionDetails({ + adapters: [...adaptersList, adapterType], + }); } }; @@ -503,6 +503,11 @@ function ConfigureDs({ onModelsLoaded={handleOpenAIModelSchema} /> )} + {isTcSuccessful && ( +
+ Connection tested successfully. Click Submit to save this adapter. +
+ )} + (Array.isArray(adapterTypes) ? adapterTypes : []) + .filter((adapterType) => typeof adapterType === "string") + .map((adapterType) => adapterType.toLowerCase()); + function OnBoard() { const navigate = useNavigate(); const { sessionDetails } = useSessionStore(); @@ -26,12 +31,19 @@ function OnBoard() { const [editItemId, setEditItemId] = useState(null); const [type, setType] = useState(null); const homePageUrl = `/${orgName}/${homePagePath}`; - const [adaptersList, setAdaptersList] = useState(adapters || []); + const [adaptersList, setAdaptersList] = useState(() => + normalizeAdapterTypes(adapters), + ); + + useEffect(() => { + setAdaptersList(normalizeAdapterTypes(adapters)); + }, [adapters]); + useEffect(() => { if (onboardCompleted(adaptersList)) { navigate(homePageUrl); } - }, [adaptersList]); + }, [adaptersList, homePageUrl, navigate]); const steps = [ { @@ -73,9 +85,17 @@ function OnBoard() { setOpenAddSourcesModal(true); }; - const addNewItem = (row, isEdit) => { - const newAdapter = row?.adapter_type.toLowerCase(); - setAdaptersList([...adaptersList, newAdapter]); + const addNewItem = (row) => { + const newAdapter = row?.adapter_type?.toLowerCase(); + if (!newAdapter) { + return; + } + + setAdaptersList((currentAdapters) => + currentAdapters.includes(newAdapter) + ? currentAdapters + : [...currentAdapters, newAdapter], + ); }; return ( @@ -114,17 +134,22 @@ function OnBoard() { {adaptersList?.includes(step.type) ? ( -
+
Configured
) : ( - +
+ + Not configured + + +
)} diff --git a/frontend/src/components/onboard/onBoard.css b/frontend/src/components/onboard/onBoard.css index 57cd3b5651..79407cb47b 100644 --- a/frontend/src/components/onboard/onBoard.css +++ b/frontend/src/components/onboard/onBoard.css @@ -132,6 +132,29 @@ vertical-align: middle; } +.configured-status, +.unconfigured-status { + display: flex; + flex-direction: column; + align-items: center; + gap: 6px; +} + +.configured-status { + color: var(--success); + font-weight: 600; +} + +.configured-status .configured-text { + margin-left: 0; +} + +.not-configured-text { + color: #6b7280; + font-size: 0.85rem; + font-weight: 600; +} + @media screen and (max-height: 900px) { .ant-space .ant-space-item .ant-card { padding: 6px; From 7f4c851e4f5081b2d7e8d1528b234e394c6df8ec Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 23:21:21 -0400 Subject: [PATCH 18/19] Persist OpenAI OAuth accounts across sessions --- .../0002_openai_oauth_credential.py | 68 ++++++++ backend/connector_auth_v2/models.py | 41 +++++ backend/connector_auth_v2/openai_oauth.py | 147 +++++++++++++++--- .../connector_auth_v2/test_openai_oauth.py | 82 ++++++++++ backend/connector_auth_v2/urls.py | 6 + backend/connector_auth_v2/views.py | 29 ++-- .../components/oauth-ds/oauth-ds/OAuthDs.jsx | 126 +++++++++++++-- .../oauth-ds/openai/OpenAIOAuthButton.jsx | 10 +- 8 files changed, 463 insertions(+), 46 deletions(-) create mode 100644 backend/connector_auth_v2/migrations/0002_openai_oauth_credential.py create mode 100644 backend/connector_auth_v2/test_openai_oauth.py diff --git a/backend/connector_auth_v2/migrations/0002_openai_oauth_credential.py b/backend/connector_auth_v2/migrations/0002_openai_oauth_credential.py new file mode 100644 index 0000000000..1c1c0526c3 --- /dev/null +++ b/backend/connector_auth_v2/migrations/0002_openai_oauth_credential.py @@ -0,0 +1,68 @@ +import uuid + +import django.db.models.deletion +from django.conf import settings +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + migrations.swappable_dependency(settings.AUTH_USER_MODEL), + ("connector_auth_v2", "0001_initial"), + ] + + operations = [ + migrations.CreateModel( + name="OpenAIOAuthCredential", + fields=[ + ( + "id", + models.UUIDField( + default=uuid.uuid4, + editable=False, + primary_key=True, + serialize=False, + ), + ), + ( + "organization_id", + models.CharField(max_length=64), + ), + ("account_id", models.CharField(max_length=255)), + ( + "account_label", + models.CharField(blank=True, default="", max_length=255), + ), + ("encrypted_credentials", models.TextField()), + ("created_at", models.DateTimeField(auto_now_add=True)), + ("modified_at", models.DateTimeField(auto_now=True)), + ( + "user", + models.ForeignKey( + on_delete=django.db.models.deletion.CASCADE, + related_name="openai_oauth_credentials", + to=settings.AUTH_USER_MODEL, + ), + ), + ], + options={ + "verbose_name": "OpenAI OAuth credential", + "verbose_name_plural": "OpenAI OAuth credentials", + "db_table": "openai_oauth_credential", + }, + ), + migrations.AddConstraint( + model_name="openaioauthcredential", + constraint=models.UniqueConstraint( + fields=("user", "organization_id", "account_id"), + name="unique_openai_oauth_user_org_account", + ), + ), + migrations.AddIndex( + model_name="openaioauthcredential", + index=models.Index( + fields=("user", "organization_id", "-modified_at"), + name="openai_oauth_user_org_mod_idx", + ), + ), + ] diff --git a/backend/connector_auth_v2/models.py b/backend/connector_auth_v2/models.py index a92630c41a..af7aa00d02 100644 --- a/backend/connector_auth_v2/models.py +++ b/backend/connector_auth_v2/models.py @@ -3,6 +3,7 @@ from typing import Any from account_v2.models import User +from backend.constants import FieldLengthConstants as FieldLength from django.db import models from django.db.models.query import QuerySet from rest_framework.request import Request @@ -157,3 +158,43 @@ class Meta: class ConnectorDjangoStorage(DjangoStorage): user = ConnectorAuth + + +class OpenAIOAuthCredential(models.Model): + """Encrypted OpenAI OAuth credentials retained for the signed-in user. + + The browser hand-off in Redis is intentionally short-lived. This model is + the durable server-side account record that lets a user reopen the OAuth + form without repeating device login. Secrets are encrypted with the same + Fernet key used for adapter metadata and are never serialized to clients. + """ + + id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False) + user = models.ForeignKey( + User, + related_name="openai_oauth_credentials", + on_delete=models.CASCADE, + ) + organization_id = models.CharField(max_length=FieldLength.ORG_NAME_SIZE) + account_id = models.CharField(max_length=255) + account_label = models.CharField(max_length=255, blank=True, default="") + encrypted_credentials = models.TextField() + created_at = models.DateTimeField(auto_now_add=True) + modified_at = models.DateTimeField(auto_now=True) + + class Meta: + db_table = "openai_oauth_credential" + verbose_name = "OpenAI OAuth credential" + verbose_name_plural = "OpenAI OAuth credentials" + constraints = [ + models.UniqueConstraint( + fields=["user", "organization_id", "account_id"], + name="unique_openai_oauth_user_org_account", + ), + ] + indexes = [ + models.Index( + fields=["user", "organization_id", "-modified_at"], + name="openai_oauth_user_org_mod_idx", + ), + ] diff --git a/backend/connector_auth_v2/openai_oauth.py b/backend/connector_auth_v2/openai_oauth.py index fafc97c34c..1e112214e1 100644 --- a/backend/connector_auth_v2/openai_oauth.py +++ b/backend/connector_auth_v2/openai_oauth.py @@ -31,12 +31,12 @@ ) from utils.user_session import UserSessionUtils +from connector_auth_v2.models import OpenAIOAuthCredential + _CACHE_PREFIX = "openai-oauth:" _DEFAULT_STATE_TTL_SECONDS = 900 _STATE_TTL_SECONDS = int( - os.environ.get( - "OPENAI_OAUTH_STATE_TTL_SECONDS", str(_DEFAULT_STATE_TTL_SECONDS) - ) + os.environ.get("OPENAI_OAUTH_STATE_TTL_SECONDS", str(_DEFAULT_STATE_TTL_SECONDS)) ) @@ -56,9 +56,7 @@ def _safe_response_json(response: httpx.Response, operation: str) -> dict[str, A f"OpenAI OAuth {operation} returned an invalid response" ) from exc if not isinstance(payload, dict): - raise OpenAIOAuthError( - f"OpenAI OAuth {operation} returned an invalid response" - ) + raise OpenAIOAuthError(f"OpenAI OAuth {operation} returned an invalid response") return payload @@ -72,9 +70,7 @@ def _post_json(url: str, *, json_body: dict[str, Any], operation: str) -> dict[s return _safe_response_json(response, operation) -def _post_form( - url: str, *, form_data: dict[str, str], operation: str -) -> dict[str, Any]: +def _post_form(url: str, *, form_data: dict[str, str], operation: str) -> dict[str, Any]: try: response = httpx.post( url, @@ -140,10 +136,61 @@ def _decrypt(value: str) -> dict[str, Any]: "OpenAI OAuth login session is no longer valid" ) from exc if not isinstance(credentials, dict): + raise OpenAIOAuthSessionError("OpenAI OAuth login session is no longer valid") + return credentials + + @classmethod + def _persist_credentials( + cls, request: Request, credentials: dict[str, Any] + ) -> OpenAIOAuthCredential: + """Upsert encrypted credentials for the current user and organization.""" + _, organization_id = cls._identity(request) + account_id = credentials.get("oauth_account_id") + if not isinstance(account_id, str) or not account_id: raise OpenAIOAuthSessionError( - "OpenAI OAuth login session is no longer valid" + "OpenAI OAuth credentials returned no ChatGPT account" ) - return credentials + + return OpenAIOAuthCredential.objects.update_or_create( + user_id=request.user.pk, + organization_id=organization_id, + account_id=account_id, + defaults={ + "account_label": _account_label( + account_id, credentials.get("oauth_account_email") + ), + "encrypted_credentials": cls._encrypt(credentials), + }, + )[0] + + @classmethod + def _load_persisted_credentials( + cls, request: Request, credential: OpenAIOAuthCredential + ) -> dict[str, Any]: + """Decrypt and refresh one durable credential record when necessary.""" + credentials = cls._decrypt(credential.encrypted_credentials) + refreshed = refresh_openai_oauth_metadata(credentials) + account_id = refreshed.get("oauth_account_id") + if not isinstance(account_id, str) or not account_id: + raise OpenAIOAuthSessionError( + "OpenAI OAuth credentials returned no ChatGPT account" + ) + + if refreshed != credentials or credential.account_id != account_id: + credential.account_id = account_id + credential.account_label = _account_label( + account_id, refreshed.get("oauth_account_email") + ) + credential.encrypted_credentials = cls._encrypt(refreshed) + credential.save( + update_fields=[ + "account_id", + "account_label", + "encrypted_credentials", + "modified_at", + ] + ) + return refreshed @classmethod def _save_state(cls, cache_key: str, state: dict[str, Any]) -> None: @@ -163,7 +210,10 @@ def _get_owned_state(cls, cache_key: str, request: Request) -> dict[str, Any]: raise OpenAIOAuthSessionError( "OpenAI OAuth login session was not found or has expired" ) - if state.get("owner_id") != user_id or state.get("organization_id") != organization_id: + if ( + state.get("owner_id") != user_id + or state.get("organization_id") != organization_id + ): raise OpenAIOAuthSessionError( "OpenAI OAuth login session was not found or is not owned by this user" ) @@ -221,15 +271,48 @@ def start(cls, request: Request) -> dict[str, Any]: @classmethod def _success_response(cls, state: dict[str, Any]) -> dict[str, Any]: - return { + response = { "status": "success", "account_label": state.get("account_label", "OpenAI account"), } + if state.get("restored"): + response["restored"] = True + return response + + @classmethod + def _create_success_handoff( + cls, request: Request, credentials: dict[str, Any], *, restored: bool = False + ) -> dict[str, Any]: + """Create a fresh short-lived hand-off for a durable account record.""" + user_id, organization_id = cls._identity(request) + account_id = credentials.get("oauth_account_id") + if not isinstance(account_id, str) or not account_id: + raise OpenAIOAuthSessionError( + "OpenAI OAuth credentials returned no ChatGPT account" + ) + cache_key = f"{_CACHE_PREFIX}{uuid.uuid4().hex}" + state = { + "status": "success", + "owner_id": user_id, + "organization_id": organization_id, + "credentials": cls._encrypt(credentials), + "account_label": _account_label( + account_id, credentials.get("oauth_account_email") + ), + "restored": restored, + } + cls._save_state(cache_key, state) + return {"cache_key": cache_key, **cls._success_response(state)} @classmethod def poll(cls, request: Request, cache_key: str) -> dict[str, Any]: state = cls._get_owned_state(cache_key, request) if state.get("status") == "success": + # A hand-off created by an older backend may have completed before + # durable persistence was introduced. Backfill it on first use. + if state.get("credentials"): + credentials = cls._decrypt(state["credentials"]) + cls._persist_credentials(request, credentials) return cls._success_response(state) try: @@ -257,9 +340,7 @@ def poll(cls, request: Request, cache_key: str) -> dict[str, Any]: device_result = _safe_response_json(response, "device login poll") authorization_code = device_result.get("authorization_code") code_verifier = device_result.get("code_verifier") - if not isinstance(authorization_code, str) or not isinstance( - code_verifier, str - ): + if not isinstance(authorization_code, str) or not isinstance(code_verifier, str): raise OpenAIOAuthError( "OpenAI OAuth device login returned incomplete authorization data" ) @@ -278,7 +359,9 @@ def poll(cls, request: Request, cache_key: str) -> dict[str, Any]: access_token = token_data.get("access_token") refresh_token = token_data.get("refresh_token") id_token = token_data.get("id_token") - if not all(isinstance(value, str) and value for value in (access_token, refresh_token)): + if not all( + isinstance(value, str) and value for value in (access_token, refresh_token) + ): raise OpenAIOAuthError( "OpenAI OAuth token exchange returned incomplete credentials" ) @@ -302,6 +385,9 @@ def poll(cls, request: Request, cache_key: str) -> dict[str, Any]: "oauth_expires_at": time.time() + expires_in, "oauth_authenticated": True, } + # Persist immediately after authorization succeeds. The Redis state + # below remains only a short-lived, browser-scoped hand-off. + cls._persist_credentials(request, credentials) state.update( { "status": "success", @@ -313,9 +399,7 @@ def poll(cls, request: Request, cache_key: str) -> dict[str, Any]: return cls._success_response(state) @classmethod - def credentials_for_request( - cls, cache_key: str, request: Request - ) -> dict[str, Any]: + def credentials_for_request(cls, cache_key: str, request: Request) -> dict[str, Any]: state = cls._get_owned_state(cache_key, request) if state.get("status") != "success" or not state.get("credentials"): raise OpenAIOAuthSessionError("Complete OpenAI OAuth authentication first") @@ -324,8 +408,27 @@ def credentials_for_request( if refreshed != credentials: state["credentials"] = cls._encrypt(refreshed) cls._save_state(cache_key, state) + cls._persist_credentials(request, refreshed) return refreshed + @classmethod + def restore(cls, request: Request) -> dict[str, Any] | None: + """Restore the most recently used durable OpenAI account, if any.""" + user_id, organization_id = cls._identity(request) + credential = ( + OpenAIOAuthCredential.objects.filter( + user_id=user_id, + organization_id=organization_id, + ) + .order_by("-modified_at") + .first() + ) + if credential is None: + return None + + credentials = cls._load_persisted_credentials(request, credential) + return cls._create_success_handoff(request, credentials, restored=True) + @staticmethod def dynamic_model_schema( credentials: dict[str, Any], *, current_model: str | None = None @@ -350,7 +453,9 @@ def consume(cls, cache_key: str, request: Request) -> None: def redact_openai_oauth_metadata(metadata: dict[str, Any]) -> dict[str, Any]: """Return the non-secret account state safe for API responses.""" redacted = { - key: value for key, value in metadata.items() if key not in OPENAI_OAUTH_PRIVATE_FIELDS + key: value + for key, value in metadata.items() + if key not in OPENAI_OAUTH_PRIVATE_FIELDS } redacted["oauth_authenticated"] = bool(metadata.get("oauth_access_token")) account_id = metadata.get("oauth_account_id") diff --git a/backend/connector_auth_v2/test_openai_oauth.py b/backend/connector_auth_v2/test_openai_oauth.py new file mode 100644 index 0000000000..f2e8fd2f28 --- /dev/null +++ b/backend/connector_auth_v2/test_openai_oauth.py @@ -0,0 +1,82 @@ +from types import SimpleNamespace +from unittest.mock import patch + +from account_v2.models import User +from cryptography.fernet import Fernet +from django.core.cache import cache +from django.test import TestCase, override_settings + +from connector_auth_v2.models import OpenAIOAuthCredential +from connector_auth_v2.openai_oauth import OpenAIOAuthService + +TEST_ENCRYPTION_KEY = Fernet.generate_key().decode("utf-8") +TEST_ORGANIZATION_ID = "oauth-test-org" +TEST_CACHES = { + "default": { + "BACKEND": "django.core.cache.backends.locmem.LocMemCache", + "LOCATION": "openai-oauth-tests", + } +} + + +@override_settings(ENCRYPTION_KEY=TEST_ENCRYPTION_KEY, CACHES=TEST_CACHES) +class OpenAIOAuthCredentialTests(TestCase): + def setUp(self) -> None: + self.user = User.objects.create_user( + username="oauth@example.com", + email="oauth@example.com", + password="not-used-in-this-test", + ) + self.request = SimpleNamespace(user=self.user) + cache.clear() + + def tearDown(self) -> None: + cache.clear() + + def test_authenticated_account_is_encrypted_and_restorable(self) -> None: + credentials = { + "oauth_access_token": "access-token", + "oauth_refresh_token": "refresh-token", + "oauth_id_token": "id-token", + "oauth_account_id": "account-123", + "oauth_account_email": "oauth@example.com", + "oauth_expires_at": 4_000_000_000, + "oauth_authenticated": True, + } + + with ( + patch.object( + OpenAIOAuthService, + "_identity", + return_value=(str(self.user.pk), TEST_ORGANIZATION_ID), + ), + patch( + "connector_auth_v2.openai_oauth.refresh_openai_oauth_metadata", + side_effect=lambda metadata: dict(metadata), + ), + ): + OpenAIOAuthService._persist_credentials(self.request, credentials) + + stored = OpenAIOAuthCredential.objects.get(user=self.user) + assert "access-token" not in stored.encrypted_credentials + assert stored.account_label == "oauth@example.com" + + restored = OpenAIOAuthService.restore(self.request) + + assert restored is not None + assert restored["status"] == "success" + assert restored["restored"] is True + assert restored["cache_key"].startswith("openai-oauth:") + + state = cache.get(restored["cache_key"]) + assert state["status"] == "success" + assert state["owner_id"] == str(self.user.pk) + assert OpenAIOAuthService._decrypt(state["credentials"]) == credentials + + def test_restore_returns_no_account_before_first_login(self) -> None: + with patch.object( + OpenAIOAuthService, + "_identity", + return_value=(str(self.user.pk), TEST_ORGANIZATION_ID), + ): + assert OpenAIOAuthService.restore(self.request) is None diff --git a/backend/connector_auth_v2/urls.py b/backend/connector_auth_v2/urls.py index 050f268a47..6f6ccbf074 100644 --- a/backend/connector_auth_v2/urls.py +++ b/backend/connector_auth_v2/urls.py @@ -10,6 +10,7 @@ ) openai_oauth_start = ConnectorAuthViewSet.as_view({"post": "openai_start"}) openai_oauth_poll = ConnectorAuthViewSet.as_view({"get": "openai_poll"}) +openai_oauth_restore = ConnectorAuthViewSet.as_view({"get": "openai_restore"}) openai_oauth_models = ConnectorAuthViewSet.as_view({"get": "openai_models"}) urlpatterns = format_suffix_patterns( @@ -22,6 +23,11 @@ ), path("oauth/openai/start/", openai_oauth_start, name="openai-oauth-start"), path("oauth/openai/poll/", openai_oauth_poll, name="openai-oauth-poll"), + path( + "oauth/openai/restore/", + openai_oauth_restore, + name="openai-oauth-restore", + ), path("oauth/openai/models/", openai_oauth_models, name="openai-oauth-models"), ] ) diff --git a/backend/connector_auth_v2/views.py b/backend/connector_auth_v2/views.py index a03c646dfb..9291916ec3 100644 --- a/backend/connector_auth_v2/views.py +++ b/backend/connector_auth_v2/views.py @@ -71,16 +71,12 @@ def openai_start(self, request: Request) -> Response: if unauthorized := self._require_authenticated(request): return unauthorized try: - return Response( - OpenAIOAuthService.start(request), status=status.HTTP_200_OK - ) + return Response(OpenAIOAuthService.start(request), status=status.HTTP_200_OK) except OpenAIOAuthSessionError as exc: return Response({"message": str(exc)}, status=status.HTTP_400_BAD_REQUEST) except OpenAIOAuthError as exc: logger.warning("OpenAI OAuth device login start failed: %s", exc) - return Response( - {"message": str(exc)}, status=status.HTTP_502_BAD_GATEWAY - ) + return Response({"message": str(exc)}, status=status.HTTP_502_BAD_GATEWAY) def openai_poll(self, request: Request) -> Response: """Poll one OpenAI device-code login and exchange it when complete.""" @@ -94,9 +90,20 @@ def openai_poll(self, request: Request) -> Response: return Response({"message": str(exc)}, status=status.HTTP_400_BAD_REQUEST) except OpenAIOAuthError as exc: logger.warning("OpenAI OAuth device login poll failed: %s", exc) - return Response( - {"message": str(exc)}, status=status.HTTP_502_BAD_GATEWAY - ) + return Response({"message": str(exc)}, status=status.HTTP_502_BAD_GATEWAY) + + def openai_restore(self, request: Request) -> Response: + """Restore the user's most recently authenticated OpenAI account.""" + if unauthorized := self._require_authenticated(request): + return unauthorized + try: + result = OpenAIOAuthService.restore(request) + return Response(result or {"status": "none"}, status=status.HTTP_200_OK) + except OpenAIOAuthSessionError as exc: + return Response({"message": str(exc)}, status=status.HTTP_400_BAD_REQUEST) + except OpenAIOAuthError as exc: + logger.warning("OpenAI OAuth account restore failed: %s", exc) + return Response({"message": str(exc)}, status=status.HTTP_502_BAD_GATEWAY) def openai_models(self, request: Request) -> Response: """Return the live model/reasoning schema for one OAuth account.""" @@ -149,6 +156,4 @@ def openai_models(self, request: Request) -> Response: return Response({"message": str(exc)}, status=status.HTTP_400_BAD_REQUEST) except OpenAIOAuthError as exc: logger.warning("OpenAI OAuth model discovery failed: %s", exc) - return Response( - {"message": str(exc)}, status=status.HTTP_502_BAD_GATEWAY - ) + return Response({"message": str(exc)}, status=status.HTTP_502_BAD_GATEWAY) diff --git a/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx b/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx index ca3abf424d..27a701e38b 100644 --- a/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx +++ b/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx @@ -1,5 +1,5 @@ import PropTypes from "prop-types"; -import { useCallback, useEffect, useState } from "react"; +import { useCallback, useEffect, useRef, useState } from "react"; import Cookies from "js-cookie"; import { Typography } from "@/components/ui/shims/antd-typography"; @@ -79,6 +79,8 @@ function OAuthDs({ }); const [activeLoginCacheKey, setActiveLoginCacheKey] = useState(null); const [isStarting, setIsStarting] = useState(false); + const [isRestoring, setIsRestoring] = useState(false); + const restoreScopeRef = useRef(null); const updateOAuthStatus = useCallback((newStatus) => { setOAuthStatus(newStatus); @@ -86,6 +88,24 @@ function OAuthDs({ localStorage.setItem(oauthStatusKey, newStatus); }, [oauthStatusKey, setStatus]); + const clearOAuthHandoff = useCallback(() => { + setLoginCacheKey(null); + setActiveLoginCacheKey(null); + setDeviceLogin(null); + setCacheKey(""); + setStatus(""); + setOAuthStatus(""); + localStorage.removeItem(oauthCacheKey); + localStorage.removeItem(oauthStatusKey); + localStorage.removeItem(oauthDeviceKey); + }, [ + oauthCacheKey, + oauthDeviceKey, + oauthStatusKey, + setCacheKey, + setStatus, + ]); + const loadOpenAIModelSchema = useCallback( async (oauthKey = "") => { const params = new URLSearchParams(); @@ -160,6 +180,87 @@ function OAuthDs({ setStatus, ]); + useEffect(() => { + if ( + oAuthProvider !== O_AUTH_PROVIDERS.OPENAI || + adapterInstanceId || + hasOAuthCredentials || + oauthStatus === "pending" || + activeLoginCacheKey || + restoreScopeRef.current === oauthStateScope + ) { + return undefined; + } + + restoreScopeRef.current = oauthStateScope; + let isActive = true; + setIsRestoring(true); + + axiosPrivate({ + method: "GET", + url: "/api/v1/oauth/openai/restore/", + }) + .then((response) => { + if (!isActive) { + return; + } + const result = response?.data || {}; + if (result.status !== "success" || !result.cache_key) { + clearOAuthHandoff(); + return; + } + + const restoredCacheKey = result.cache_key; + const restoredLogin = { + ...result, + status: "success", + }; + setLoginCacheKey(restoredCacheKey); + setActiveLoginCacheKey(restoredCacheKey); + setCacheKey(restoredCacheKey); + setDeviceLogin(restoredLogin); + localStorage.setItem(oauthCacheKey, restoredCacheKey); + localStorage.setItem(oauthDeviceKey, JSON.stringify(restoredLogin)); + updateOAuthStatus("success"); + }) + .catch((err) => { + if (!isActive) { + return; + } + clearOAuthHandoff(); + setAlertDetails( + handleException( + err, + "Could not restore your saved OpenAI authentication", + ), + ); + }) + .finally(() => { + if (isActive) { + setIsRestoring(false); + } + }); + + return () => { + isActive = false; + }; + }, [ + activeLoginCacheKey, + adapterInstanceId, + axiosPrivate, + clearOAuthHandoff, + handleException, + hasOAuthCredentials, + oauthCacheKey, + oauthDeviceKey, + oauthStatus, + oAuthProvider, + oauthStateScope, + setAlertDetails, + setCacheKey, + updateOAuthStatus, + ]); + useEffect(() => { // Do not let an old pending browser state mask credentials already saved // on an existing adapter. A live reauthentication session is allowed to @@ -221,15 +322,8 @@ function OAuthDs({ typeof message === "string" && message.toLowerCase().includes("openai oauth login session"); if (isExpiredLoginSession) { - setLoginCacheKey(null); - setActiveLoginCacheKey(null); - setDeviceLogin(null); - setCacheKey(""); - setStatus(""); - setOAuthStatus(""); - localStorage.removeItem(oauthCacheKey); - localStorage.removeItem(oauthStatusKey); - localStorage.removeItem(oauthDeviceKey); + clearOAuthHandoff(); + return; } setAlertDetails(handleException(err, message)); }); @@ -241,6 +335,7 @@ function OAuthDs({ activeLoginCacheKey, adapterInstanceId, handleException, + clearOAuthHandoff, loadOpenAIModelSchema, loginCacheKey, oauthStatus, @@ -284,6 +379,15 @@ function OAuthDs({ } const message = err?.response?.data?.message || "OpenAI authentication failed"; + const isExpiredLoginSession = + typeof message === "string" && + message.toLowerCase().includes("openai oauth login session"); + if (isExpiredLoginSession) { + // A Redis hand-off is disposable. Clear the stale browser key so a + // durable server-side account can be restored on the next render. + clearOAuthHandoff(); + return; + } updateOAuthStatus("error"); setAlertDetails(handleException(err, message)); } @@ -298,6 +402,7 @@ function OAuthDs({ }; }, [ axiosPrivate, + clearOAuthHandoff, handleException, loginCacheKey, oauthStatus, @@ -441,6 +546,7 @@ function OAuthDs({ userCode={deviceLogin?.user_code} accountLabel={deviceLogin?.account_label || oauthAccountLabel} isStarting={isStarting} + isRestoring={isRestoring} /> ); } diff --git a/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx b/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx index 3f494b1a4e..8b731bfefe 100644 --- a/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx +++ b/frontend/src/components/oauth-ds/openai/OpenAIOAuthButton.jsx @@ -13,9 +13,12 @@ const OpenAIOAuthButton = ({ userCode, accountLabel, isStarting = false, + isRestoring = false, }) => { const isPending = status === "pending"; - const buttonLabel = isStarting + const buttonLabel = isRestoring + ? "Restoring OpenAI authentication" + : isStarting ? "Starting OpenAI sign-in" : status === "success" ? "Authenticated" @@ -30,8 +33,8 @@ const OpenAIOAuthButton = ({ @@ -63,6 +66,7 @@ OpenAIOAuthButton.propTypes = { userCode: PropTypes.string, accountLabel: PropTypes.string, isStarting: PropTypes.bool, + isRestoring: PropTypes.bool, }; export default OpenAIOAuthButton; From 1d82aa38169402b4a00293213cfdfb2727c969b9 Mon Sep 17 00:00:00 2001 From: Timothy Wayne Gregg <5861166+CompleteDotTech@users.noreply.github.com> Date: Thu, 3 Sep 2026 23:28:34 -0400 Subject: [PATCH 19/19] Avoid stale OAuth restore race --- .../src/components/oauth-ds/oauth-ds/OAuthDs.jsx | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx b/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx index 27a701e38b..22a750a86c 100644 --- a/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx +++ b/frontend/src/components/oauth-ds/oauth-ds/OAuthDs.jsx @@ -290,6 +290,17 @@ function OAuthDs({ return undefined; } + // A new adapter may have a stale success/cache key in localStorage. Wait + // for the server-side restore request before trying that disposable key; + // otherwise an expired model request can cancel the restore race. + if ( + !adapterInstanceId && + !activeLoginCacheKey && + (isRestoring || restoreScopeRef.current === oauthStateScope) + ) { + return undefined; + } + // Existing adapters use their server-side encrypted credentials. A new // OAuth login takes precedence only after that login has completed, so a // stale localStorage cache key cannot select another account by accident. @@ -336,11 +347,13 @@ function OAuthDs({ adapterInstanceId, handleException, clearOAuthHandoff, + isRestoring, loadOpenAIModelSchema, loginCacheKey, oauthStatus, oAuthProvider, onModelsLoaded, + oauthStateScope, setAlertDetails, setCacheKey, setStatus,