feat: support upstream token refresh

backend/app/services/upstream_client.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 
 import json
 import time
-from typing import Any, Optional
+from typing import Any, Callable, Optional
 from urllib.parse import urljoin
 import httpx
 
@@ -32,6 +32,38 @@ def _find_token(value: Any) -> str:
     return ""
 
 
+def _find_refresh_token(value: Any) -> str:
+    if isinstance(value, dict):
+        for key in ("refresh_token", "refreshToken"):
+            candidate = value.get(key)
+            if isinstance(candidate, str) and candidate:
+                return candidate
+        for key in ("data", "result", "user", "session"):
+            tok = _find_refresh_token(value.get(key))
+            if tok:
+                return tok
+    return ""
+
+
+def _find_expires_in(value: Any) -> int | None:
+    if isinstance(value, dict):
+        for key in ("expires_in", "expiresIn"):
+            candidate = value.get(key)
+            if candidate is None:
+                continue
+            try:
+                expires_in = int(float(candidate))
+            except (TypeError, ValueError):
+                continue
+            if expires_in > 0:
+                return expires_in
+        for key in ("data", "result", "user", "session"):
+            expires_in = _find_expires_in(value.get(key))
+            if expires_in is not None:
+                return expires_in
+    return None
+
+
 def _clean_auth_header_value(value: Any, field_name: str) -> str:
     text = str(value or "").strip()
     if not text:
@@ -289,12 +321,14 @@ class UpstreamClient:
         auth_type: str,
         auth_config: dict[str, Any],
         timeout: float = 30.0,
+        on_auth_config_update: Callable[[dict[str, Any]], None] | None = None,
     ) -> None:
         self.base_url = base_url.rstrip("/")
         self.api_prefix = api_prefix.strip("/")
         self.auth_type = auth_type
         self.auth_config = auth_config
         self.timeout = timeout
+        self.on_auth_config_update = on_auth_config_update
         self._token: str = ""
         self._cookies: dict[str, str] = {}
         self._new_api_user: str = ""
@@ -321,6 +355,7 @@ class UpstreamClient:
                 bool(self.auth_config.get("new_api_user"))
                 or login_path == "/api/user/login"
                 or self.auth_type == "cookie"
+                or self.auth_type == "new_api_token"
                 or self.auth_type == "nox_token"
             )
         )
@@ -349,6 +384,13 @@ class UpstreamClient:
                 headers["Authorization"] = f"Bearer {token}"
             if user_id:
                 headers["Nox-Api-User"] = user_id
+        elif self.auth_type == "new_api_token":
+            token = _clean_auth_header_value(self.auth_config.get("token", ""), "New-API access token")
+            user_id = self._user_header_value()
+            if token:
+                headers["Authorization"] = f"Bearer {token}"
+            if user_id:
+                headers["New-Api-User"] = user_id
         elif self.auth_type == "api_key":
             key = _clean_auth_header_value(self.auth_config.get("key", ""), "API key")
             header = self.auth_config.get("header", "Authorization")
@@ -371,28 +413,117 @@ class UpstreamClient:
             headers["Nox-Api-User"] = self._new_api_user
         return headers
 
+    def _is_sub2api_bearer(self) -> bool:
+        return self.auth_type == "bearer" and self.api_prefix == "api/v1"
+
+    def _remember_auth_tokens(
+        self,
+        token: str = "",
+        refresh_token: str = "",
+        expires_in: int | None = None,
+    ) -> None:
+        changed = False
+        if token and self.auth_config.get("token") != token:
+            self.auth_config["token"] = token
+            changed = True
+        if refresh_token and self.auth_config.get("refresh_token") != refresh_token:
+            self.auth_config["refresh_token"] = refresh_token
+            changed = True
+        if expires_in is not None:
+            if self.auth_config.get("expires_in") != expires_in:
+                self.auth_config["expires_in"] = expires_in
+                changed = True
+            token_expires_at = int(time.time()) + expires_in
+            if self.auth_config.get("token_expires_at") != token_expires_at:
+                self.auth_config["token_expires_at"] = token_expires_at
+                changed = True
+        if changed and self.on_auth_config_update:
+            self.on_auth_config_update(dict(self.auth_config))
+
+    def _refresh_sub2api_bearer_token(self) -> bool:
+        if not self._is_sub2api_bearer():
+            return False
+        refresh_token = str(self.auth_config.get("refresh_token") or "").strip()
+        if not refresh_token:
+            return False
+        try:
+            resp = self._client.request(
+                "POST",
+                self._url("/auth/refresh"),
+                json={"refresh_token": refresh_token},
+                headers=self._headers(auth=False),
+                cookies=self._cookies,
+            )
+            self._cookies.update(dict(resp.cookies))
+            resp.raise_for_status()
+            payload = resp.json() if resp.content else {}
+        except Exception:
+            return False
+
+        token = _find_token(payload)
+        if not token:
+            return False
+        self._remember_auth_tokens(
+            token=token,
+            refresh_token=_find_refresh_token(payload) or refresh_token,
+            expires_in=_find_expires_in(payload),
+        )
+        return True
+
+    def _send_request(
+        self,
+        method: str,
+        url: str,
+        auth: bool = True,
+        allow_refresh: bool = True,
+        **kwargs: Any,
+    ) -> httpx.Response:
+        resp = self._client.request(
+            method,
+            url,
+            headers=self._headers(auth),
+            cookies=self._cookies,
+            **kwargs,
+        )
+        self._cookies.update(dict(resp.cookies))
+        if (
+            getattr(resp, "status_code", None) == 401
+            and auth
+            and allow_refresh
+            and self._is_sub2api_bearer()
+            and self._refresh_sub2api_bearer_token()
+        ):
+            resp = self._client.request(
+                method,
+                url,
+                headers=self._headers(auth),
+                cookies=self._cookies,
+                **kwargs,
+            )
+            self._cookies.update(dict(resp.cookies))
+        return resp
+
     def _request(self, method: str, path: str, body: Any = None, auth: bool = True) -> Any:
-        if auth and self.auth_type == "cookie" and "user/self" in path and not self.auth_config.get("new_api_user"):
+        if auth and self.auth_type == "cookie" and "user/self" in path and not self._user_header_value():
             raise UpstreamError("New-API user endpoint requires New-Api-User; re-extract the session cookie after login and save the upstream")
+        if auth and self.auth_type == "new_api_token" and "user/self" in path and not self._user_header_value():
+            raise UpstreamError("New-API endpoint requires New-Api-User; please fill user id and retry")
         if auth and self.auth_type == "nox_token" and "user/self" in path and not self._user_header_value():
             raise UpstreamError("Nox-API endpoint requires Nox-Api-User; please fill user id and retry")
         url = self._url(path)
         if body is not None:
-            resp = self._client.request(
+            resp = self._send_request(
                 method,
                 url,
                 json=body,
-                headers=self._headers(auth),
-                cookies=self._cookies,
+                auth=auth,
             )
         else:
-            resp = self._client.request(
+            resp = self._send_request(
                 method,
                 url,
-                headers=self._headers(auth),
-                cookies=self._cookies,
+                auth=auth,
             )
-        self._cookies.update(dict(resp.cookies))
         resp.raise_for_status()
         ct = resp.headers.get("content-type", "")
         if not resp.content:
@@ -449,14 +580,11 @@ class UpstreamClient:
         return items, meta
 
     def _request_new_api_token_list(self, path: str, params: dict[str, Any]) -> tuple[list[dict[str, Any]], dict[str, Any]]:
-        resp = self._client.request(
+        resp = self._send_request(
             "GET",
             self._url(path),
             params=params,
-            headers=self._headers(),
-            cookies=self._cookies,
         )
-        self._cookies.update(dict(resp.cookies))
         resp.raise_for_status()
         data = resp.json()
         self._ensure_api_success(data, "list New-API tokens")
@@ -585,14 +713,28 @@ class UpstreamClient:
             return
         email = self.auth_config.get("email", "")
         password = self.auth_config.get("password", "")
-        login_path = self.auth_config.get("login_path", "/auth/login")
-        username_field = self.auth_config.get("username_field", "email")
+        default_login_path = "/api/user/login" if self.api_prefix == "" else "/auth/login"
+        login_path = self.auth_config.get("login_path") or default_login_path
+        default_username_field = "username" if login_path == "/api/user/login" else "email"
+        username_field = self.auth_config.get("username_field") or default_username_field
         if not email or not password:
             raise UpstreamError("login_password auth requires email and password in auth_config")
         resp = self._request("POST", login_path, {username_field: email, "password": password}, auth=False)
         token = _find_token(resp)
         if token:
             self._token = token
+            self._new_api_user = (
+                self.auth_config.get("new_api_user", "")
+                or self.auth_config.get("user_id", "")
+                or _find_user_id(resp)
+            )
+            refresh_token = _find_refresh_token(resp)
+            if self.api_prefix == "api/v1" or refresh_token:
+                self._remember_auth_tokens(
+                    token=token,
+                    refresh_token=refresh_token,
+                    expires_in=_find_expires_in(resp),
+                )
             return
         if self._cookies:
             self._new_api_user = self.auth_config.get("new_api_user", "") or _find_user_id(resp)
@@ -655,12 +797,10 @@ class UpstreamClient:
         if status:
             params["status"] = status
         url = self._url(endpoint)
-        resp = self._client.request(
+        resp = self._send_request(
             "GET",
             url,
             params=params if params else None,
-            headers=self._headers(),
-            cookies=self._cookies,
         )
         resp.raise_for_status()
         data = resp.json()