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 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
parent
5a20e792fe
commit
2262b6a564
@@ -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"},
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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] = []
|
||||
|
||||
|
||||
Reference in New Issue
Block a user