From 2262b6a564e09caebd2b79594954a5277c6cd037 Mon Sep 17 00:00:00 2001 From: SmartUp Developer Date: Mon, 6 Jul 2026 10:31:29 +0800 Subject: [PATCH] feat: proactive Bearer token refresh with unified fallback - Parse JWT 'exp' claim to track token expiry (token_expires_at) - Proactively refresh Bearer tokens 1 hour before expiry in ensure_authenticated() - Refresh response without expires_in now falls back to JWT exp for expiry derivation - Support custom refresh_path via bearer auth_config whitelist - Unify 401 fallback: all bearer types (including Sub2API) now go through _refresh_bearer_token() instead of separate old path - Remove dead _is_sub2api_bearer() / _refresh_sub2api_bearer_token() - Add 2 new tests: refresh_path override + JWT exp backfill Co-Authored-By: Claude Opus 4.8 --- backend/app/services/auth_config.py | 2 +- backend/app/services/upstream_client.py | 99 ++++++++-- backend/test_upstream_token_refresh.py | 245 +++++++++++++++++++++++- 3 files changed, 331 insertions(+), 15 deletions(-) diff --git a/backend/app/services/auth_config.py b/backend/app/services/auth_config.py index d6c2585..f5a0597 100644 --- a/backend/app/services/auth_config.py +++ b/backend/app/services/auth_config.py @@ -5,7 +5,7 @@ SECRET_KEYS = {"password", "token", "refresh_token", "access_token", "auth_token AUTH_CONFIG_ALLOWED_KEYS = { "none": set(), - "bearer": {"token", "refresh_token", "expires_in", "token_expires_at"}, + "bearer": {"token", "refresh_token", "expires_in", "token_expires_at", "refresh_path"}, "new_api_token": {"token", "user_id", "new_api_user", "provider"}, "nox_token": {"token", "user_id", "new_api_user", "provider"}, "api_key": {"key", "header"}, diff --git a/backend/app/services/upstream_client.py b/backend/app/services/upstream_client.py index 92294e0..b5d4464 100644 --- a/backend/app/services/upstream_client.py +++ b/backend/app/services/upstream_client.py @@ -1,6 +1,7 @@ """Upstream HTTP client — ported from monitor_ai98pro_group_rates.py.""" from __future__ import annotations +import base64 import hashlib import json import logging @@ -133,6 +134,27 @@ def _find_expires_in(value: Any) -> int | None: return None +def _parse_jwt_exp_timestamp(token: str) -> int | None: + """Extract the 'exp' claim (Unix seconds) from a JWT access token. + + Returns None if the token isn't a JWT or has no valid exp claim. + """ + if not token or token.count(".") != 2: + return None + try: + payload_b64 = token.split(".")[1] + padding = 4 - len(payload_b64) % 4 + if padding != 4: + payload_b64 += "=" * padding + claims = json.loads(base64.urlsafe_b64decode(payload_b64)) + exp = claims.get("exp") + if exp is not None: + return int(exp) + except Exception: + pass + return None + + def _clean_auth_header_value(value: Any, field_name: str) -> str: text = str(value or "").strip() if not text: @@ -542,14 +564,12 @@ 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, + token_expires_at: int | None = None, ) -> None: changed = False if token and self.auth_config.get("token") != token: @@ -566,6 +586,10 @@ class UpstreamClient: if self.auth_config.get("token_expires_at") != token_expires_at: self.auth_config["token_expires_at"] = token_expires_at changed = True + elif token_expires_at is not None: + 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)) @@ -668,16 +692,25 @@ class UpstreamClient: ) return True - def _refresh_sub2api_bearer_token(self) -> bool: - if not self._is_sub2api_bearer(): - return False + def _refresh_bearer_token(self) -> bool: + """Refresh a Bearer access token using its refresh_token. + + Default refresh path depends on api_prefix: + - api_prefix == "" → /api/user/refresh + - otherwise → /auth/refresh + Can be overridden via auth_config ``refresh_path``. + When the response has no ``expires_in``, falls back to parsing + the new access token's JWT ``exp`` claim. + """ refresh_token = str(self.auth_config.get("refresh_token") or "").strip() if not refresh_token: return False + default_path = "/api/user/refresh" if self.api_prefix == "" else "/auth/refresh" + refresh_path = self.auth_config.get("refresh_path") or default_path try: resp = self._do_request( "POST", - self._url("/auth/refresh"), + self._url(refresh_path), json={"refresh_token": refresh_token}, headers=self._headers(auth=False), cookies=self._cookies, @@ -691,13 +724,49 @@ class UpstreamClient: token = _find_token(payload) if not token: return False + + expires_in = _find_expires_in(payload) + token_expires_at = None + if expires_in is None: + token_expires_at = _parse_jwt_exp_timestamp(token) + self._remember_auth_tokens( token=token, refresh_token=_find_refresh_token(payload) or refresh_token, - expires_in=_find_expires_in(payload), + expires_in=expires_in, + token_expires_at=token_expires_at, ) return True + def _ensure_bearer_token_fresh(self) -> None: + """Proactively refresh Bearer token if it will expire within 1 hour. + + Writes missing ``token_expires_at`` from JWT ``exp`` claim if absent. + Does nothing if there is no ``refresh_token`` or no expiry info. + """ + refresh_token = str(self.auth_config.get("refresh_token") or "").strip() + if not refresh_token: + return + + # Try to derive token_expires_at from JWT exp if not already saved + if self.auth_config.get("token_expires_at") is None: + token = str(self.auth_config.get("token") or "").strip() + exp = _parse_jwt_exp_timestamp(token) + if exp is not None: + self._remember_auth_tokens(token_expires_at=exp) + + token_expires_at = self.auth_config.get("token_expires_at") + if token_expires_at is None: + return + try: + expires_at_int = int(token_expires_at) + except (TypeError, ValueError): + return + + # Refresh if within 1 hour of expiry + if int(time.time()) + 3600 >= expires_at_int: + self._refresh_bearer_token() + def _send_request( self, method: str, @@ -719,8 +788,8 @@ class UpstreamClient: and auth and allow_refresh ): - # sub2api bearer: 尝试 refresh - if self._is_sub2api_bearer() and self._refresh_sub2api_bearer_token(): + # bearer: 尝试 refresh（含 Sub2API，统一走 _refresh_bearer_token） + if self.auth_type == "bearer" and self._refresh_bearer_token(): resp = self._do_request( method, url, @@ -1050,12 +1119,18 @@ class UpstreamClient: def ensure_authenticated(self) -> None: """确保已认证，优先复用 token，过期后 refresh，失败再重新登录。 - 对于 login_password 类型： + 对于 bearer 类型： + - 有 refresh_token 时，在过期前 1 小时主动 refresh（避免上游在过期后拒绝刷新） + - 无 refresh_token 时不做任何事（静默失败，后续 API 调用会 401） + + 对于 login_password 类型（保持原有逻辑）： - 已有未过期 token：不登录，直接使用 - token 快过期或已过期且有 refresh_token：先 refresh - refresh 失败、无 token、无 refresh token：回退到现有登录流程 - 其他认证类型：调用 login() 保持原有行为 """ + if self.auth_type == "bearer": + self._ensure_bearer_token_fresh() + return if self.auth_type != "login_password": self.login() return diff --git a/backend/test_upstream_token_refresh.py b/backend/test_upstream_token_refresh.py index 4907db1..79b2210 100644 --- a/backend/test_upstream_token_refresh.py +++ b/backend/test_upstream_token_refresh.py @@ -1,5 +1,9 @@ from __future__ import annotations +import base64 +import json +import time + import httpx import pytest @@ -36,6 +40,17 @@ def _response(request: httpx.Request, status_code: int, payload=None) -> httpx.R return httpx.Response(status_code, json=payload, request=request) +def _jwt_token(exp_timestamp: int) -> str: + """Build a fake JWT with the given exp claim.""" + payload_b64 = base64.urlsafe_b64encode( + json.dumps({"exp": exp_timestamp}).encode() + ).rstrip(b"=").decode() + return f"header.{payload_b64}.sig" + + +# ── Existing Sub2API bearer tests (unchanged) ───────────────────────────── + + def test_sub2api_bearer_refreshes_on_401_and_retries(monkeypatch): updates: list[dict] = [] @@ -98,7 +113,51 @@ def test_sub2api_bearer_without_refresh_token_does_not_refresh(monkeypatch): ] -def test_non_sub2api_bearer_does_not_refresh(monkeypatch): +# ── Updated: non-Sub2API bearer now refreshes with refresh_token ────────── + + +def test_bearer_refreshes_on_401_with_default_refresh_path(monkeypatch): + """Non-sub2api bearer with refresh_token refreshes on 401 (new behaviour).""" + updates: list[dict] = [] + + def handler(request: httpx.Request, kwargs): + if request.url.path == "/groups/available": + auth = kwargs["headers"].get("Authorization") + if auth == "Bearer old-access": + return _response(request, 401, {"detail": "expired"}) + assert auth == "Bearer new-access" + return _response(request, 200, [{"id": "g1", "name": "G1"}]) + if request.url.path == "/api/user/refresh": + assert kwargs["json"] == {"refresh_token": "old-refresh"} + return _response(request, 200, { + "access_token": "new-access", + "refresh_token": "new-refresh", + "expires_in": 3600, + }) + raise AssertionError(f"unexpected: {request.method} {request.url}") + + fake = _install_fake_client(monkeypatch, handler) + client = UpstreamClient( + base_url="http://plain-bearer.local", + api_prefix="", + auth_type="bearer", + auth_config={"token": "old-access", "refresh_token": "old-refresh"}, + on_auth_config_update=updates.append, + ) + + groups = client.get_available_groups("/groups/available") + + assert groups == [{"id": "g1", "name": "G1"}] + assert [(c["method"], c["path"]) for c in fake.calls] == [ + ("GET", "/groups/available"), + ("POST", "/api/user/refresh"), + ("GET", "/groups/available"), + ] + assert updates[-1]["token"] == "new-access" + + +def test_bearer_without_refresh_token_still_fails_on_401(monkeypatch): + """Bearer without refresh_token does NOT retry on 401.""" def handler(request: httpx.Request, kwargs): assert request.url.path == "/groups/available" return _response(request, 401, {"detail": "expired"}) @@ -108,7 +167,7 @@ def test_non_sub2api_bearer_does_not_refresh(monkeypatch): base_url="http://plain-bearer.local", api_prefix="", auth_type="bearer", - auth_config={"token": "old-access", "refresh_token": "old-refresh"}, + auth_config={"token": "old-access"}, ) with pytest.raises(httpx.HTTPStatusError): @@ -119,6 +178,188 @@ def test_non_sub2api_bearer_does_not_refresh(monkeypatch): ] +# ── New: proactive refresh tests ───────────────────────────────────────── + + +def test_bearer_proactive_refresh_when_expiry_approaching(monkeypatch): + """ensure_authenticated proactively refreshes token expiring within 1 h.""" + updates: list[dict] = [] + + def handler(request: httpx.Request, kwargs): + if request.url.path == "/api/user/refresh": + assert kwargs["json"] == {"refresh_token": "the-refresh"} + return _response(request, 200, { + "access_token": "new-token", + "refresh_token": "new-refresh", + "expires_in": 7200, + }) + raise AssertionError(f"unexpected: {request.method} {request.url}") + + _install_fake_client(monkeypatch, handler) + client = UpstreamClient( + base_url="http://bearer.local", + api_prefix="", + auth_type="bearer", + auth_config={ + "token": "old-token", + "refresh_token": "the-refresh", + "token_expires_at": int(time.time()) + 1800, # 30 min left + }, + on_auth_config_update=updates.append, + ) + + client.ensure_authenticated() + + assert updates[-1]["token"] == "new-token" + assert updates[-1]["refresh_token"] == "new-refresh" + + +def test_bearer_proactive_skips_when_token_fresh(monkeypatch): + """ensure_authenticated does NOT refresh token that is far from expiry.""" + + def handler(request: httpx.Request, kwargs): + raise AssertionError(f"unexpected request: {request.method} {request.url}") + + _install_fake_client(monkeypatch, handler) + client = UpstreamClient( + base_url="http://bearer.local", + api_prefix="", + auth_type="bearer", + auth_config={ + "token": "still-fresh", + "refresh_token": "the-refresh", + "token_expires_at": int(time.time()) + 86400 * 7, # 7 days + }, + ) + + # Should not make any HTTP request + client.ensure_authenticated() + + +def test_bearer_proactive_parses_jwt_exp_when_missing(monkeypatch): + """Missing token_expires_at is derived from JWT exp claim.""" + future_exp = int(time.time()) + 86400 * 7 # 7 days → far away, no refresh + fake_jwt = _jwt_token(future_exp) + + def handler(request: httpx.Request, kwargs): + raise AssertionError(f"unexpected request: {request.method} {request.url}") + + _install_fake_client(monkeypatch, handler) + client = UpstreamClient( + base_url="http://bearer.local", + api_prefix="", + auth_type="bearer", + auth_config={"token": fake_jwt, "refresh_token": "the-refresh"}, + ) + + client.ensure_authenticated() + + assert client.auth_config.get("token_expires_at") == future_exp + + +def test_bearer_proactive_refresh_from_jwt_exp(monkeypatch): + """When token_expires_at is missing but JWT exp is within 1h, refresh.""" + updates: list[dict] = [] + + def handler(request: httpx.Request, kwargs): + if request.url.path == "/api/user/refresh": + return _response(request, 200, { + "access_token": "refreshed-token", + "expires_in": 7200, + }) + raise AssertionError(f"unexpected: {request.method} {request.url}") + + _install_fake_client(monkeypatch, handler) + near_exp = int(time.time()) + 600 # 10 min left + client = UpstreamClient( + base_url="http://bearer.local", + api_prefix="", + auth_type="bearer", + auth_config={"token": _jwt_token(near_exp), "refresh_token": "the-refresh"}, + on_auth_config_update=updates.append, + ) + + client.ensure_authenticated() + + assert updates[-1]["token"] == "refreshed-token" + + +def test_bearer_refresh_uses_custom_refresh_path(monkeypatch): + """refresh_path in auth_config overrides the default refresh endpoint.""" + updates: list[dict] = [] + + def handler(request: httpx.Request, kwargs): + if request.url.path == "/custom/refresh/endpoint": + assert kwargs["json"] == {"refresh_token": "my-refresh"} + return _response(request, 200, { + "access_token": "custom-refreshed", + "expires_in": 7200, + }) + raise AssertionError(f"unexpected: {request.method} {request.url}") + + _install_fake_client(monkeypatch, handler) + # Token is within 1h → triggers proactive refresh + near_exp = int(time.time()) + 1800 + client = UpstreamClient( + base_url="http://bearer.local", + api_prefix="", # default would be /api/user/refresh + auth_type="bearer", + auth_config={ + "token": "stale-token", + "refresh_token": "my-refresh", + "refresh_path": "/custom/refresh/endpoint", + "token_expires_at": near_exp, + }, + on_auth_config_update=updates.append, + ) + + client.ensure_authenticated() + + assert updates[-1]["token"] == "custom-refreshed" + + +def test_bearer_refresh_backfills_token_expires_at_from_jwt(monkeypatch): + """When refresh response has no expires_in, derive token_expires_at from JWT exp.""" + updates: list[dict] = [] + future_exp = int(time.time()) + 3600 # 1 h from now + new_access_token = _jwt_token(future_exp) + + def handler(request: httpx.Request, kwargs): + if request.url.path == "/api/user/refresh": + assert kwargs["json"] == {"refresh_token": "my-refresh"} + # Response deliberately omits expires_in + return _response(request, 200, { + "access_token": new_access_token, + "refresh_token": "new-refresh", + }) + raise AssertionError(f"unexpected: {request.method} {request.url}") + + _install_fake_client(monkeypatch, handler) + # Token is within 1h → triggers proactive refresh + near_exp = int(time.time()) + 1800 + client = UpstreamClient( + base_url="http://bearer.local", + api_prefix="", + auth_type="bearer", + auth_config={ + "token": "stale-token", + "refresh_token": "my-refresh", + "token_expires_at": near_exp, + }, + on_auth_config_update=updates.append, + ) + + client.ensure_authenticated() + + assert updates[-1]["token"] == new_access_token + assert updates[-1]["token_expires_at"] == future_exp + # expires_in was NOT in the response → should not appear in the update + assert "expires_in" not in updates[-1] + + +# ── Existing login_password tests (unchanged) ───────────────────────────── + + def test_login_password_saves_sub2api_refresh_token(monkeypatch): updates: list[dict] = []