mirror of
https://github.com/langgenius/dify.git
synced 2026-06-07 16:13:59 +08:00
fix: user token (#36930)
This commit is contained in:
parent
7056985f72
commit
888483a2f8
@ -119,6 +119,8 @@ class TokenPair(BaseModel):
|
||||
REFRESH_TOKEN_PREFIX = "refresh_token:"
|
||||
ACCOUNT_REFRESH_TOKEN_PREFIX = "account_refresh_token:"
|
||||
REFRESH_TOKEN_EXPIRY = timedelta(days=dify_config.REFRESH_TOKEN_EXPIRE_DAYS)
|
||||
ACCOUNT_LAST_ACTIVE_REFRESH_PREFIX = "account_last_active_refresh:"
|
||||
ACCOUNT_LAST_ACTIVE_REFRESH_INTERVAL = timedelta(minutes=10)
|
||||
|
||||
|
||||
class AccountService:
|
||||
@ -152,6 +154,40 @@ class AccountService:
|
||||
def _get_account_refresh_token_key(account_id: str) -> str:
|
||||
return f"{ACCOUNT_REFRESH_TOKEN_PREFIX}{account_id}"
|
||||
|
||||
@staticmethod
|
||||
def _get_account_last_active_refresh_key(account_id: str) -> str:
|
||||
return f"{ACCOUNT_LAST_ACTIVE_REFRESH_PREFIX}{account_id}"
|
||||
|
||||
@staticmethod
|
||||
@redis_fallback(default_return=True)
|
||||
def _should_refresh_account_last_active(account_id: str) -> bool:
|
||||
return bool(
|
||||
redis_client.set(
|
||||
AccountService._get_account_last_active_refresh_key(account_id),
|
||||
1,
|
||||
ex=int(ACCOUNT_LAST_ACTIVE_REFRESH_INTERVAL.total_seconds()),
|
||||
nx=True,
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _refresh_account_last_active(account: Account) -> None:
|
||||
now = naive_utc_now()
|
||||
refresh_before = now - ACCOUNT_LAST_ACTIVE_REFRESH_INTERVAL
|
||||
|
||||
if account.last_active_at >= refresh_before:
|
||||
return
|
||||
|
||||
if not AccountService._should_refresh_account_last_active(account.id):
|
||||
return
|
||||
|
||||
db.session.execute(
|
||||
update(Account)
|
||||
.where(Account.id == account.id, Account.last_active_at < refresh_before)
|
||||
.values(last_active_at=now, updated_at=func.current_timestamp())
|
||||
)
|
||||
db.session.commit()
|
||||
|
||||
@staticmethod
|
||||
def _store_refresh_token(refresh_token: str, account_id: str):
|
||||
redis_client.setex(AccountService._get_refresh_token_key(refresh_token), REFRESH_TOKEN_EXPIRY, account_id)
|
||||
@ -229,9 +265,7 @@ class AccountService:
|
||||
available_ta.current = True
|
||||
db.session.commit()
|
||||
|
||||
if naive_utc_now() - account.last_active_at > timedelta(minutes=10):
|
||||
account.last_active_at = naive_utc_now()
|
||||
db.session.commit()
|
||||
AccountService._refresh_account_last_active(account)
|
||||
# NOTE: make sure account is accessible outside of a db session
|
||||
# This ensures that it will work correctly after upgrading to Flask version 3.1.2
|
||||
db.session.refresh(account)
|
||||
|
||||
@ -436,7 +436,10 @@ class TestAccountService:
|
||||
mock_db_dependencies["db"].session.scalar.return_value = mock_tenant_join
|
||||
|
||||
# Mock datetime
|
||||
with patch("services.account_service.datetime") as mock_datetime:
|
||||
with (
|
||||
patch("services.account_service.datetime") as mock_datetime,
|
||||
patch.object(AccountService, "_refresh_account_last_active") as mock_refresh_last_active,
|
||||
):
|
||||
mock_now = datetime.now()
|
||||
mock_datetime.now.return_value = mock_now
|
||||
mock_datetime.UTC = "UTC"
|
||||
@ -447,6 +450,7 @@ class TestAccountService:
|
||||
# Verify results
|
||||
assert result == mock_account
|
||||
assert mock_account.set_tenant_id.called
|
||||
mock_refresh_last_active.assert_called_once_with(mock_account)
|
||||
|
||||
def test_load_user_not_found(self, mock_db_dependencies):
|
||||
"""Test user loading when user does not exist."""
|
||||
@ -483,7 +487,10 @@ class TestAccountService:
|
||||
mock_db_dependencies["db"].session.scalar.side_effect = [None, mock_available_tenant]
|
||||
|
||||
# Mock datetime
|
||||
with patch("services.account_service.datetime") as mock_datetime:
|
||||
with (
|
||||
patch("services.account_service.datetime") as mock_datetime,
|
||||
patch.object(AccountService, "_refresh_account_last_active") as mock_refresh_last_active,
|
||||
):
|
||||
mock_now = datetime.now()
|
||||
mock_datetime.now.return_value = mock_now
|
||||
mock_datetime.UTC = "UTC"
|
||||
@ -495,6 +502,7 @@ class TestAccountService:
|
||||
assert result == mock_account
|
||||
assert mock_available_tenant.current is True
|
||||
self._assert_database_operations_called(mock_db_dependencies["db"])
|
||||
mock_refresh_last_active.assert_called_once_with(mock_account)
|
||||
|
||||
def test_load_user_no_tenants(self, mock_db_dependencies):
|
||||
"""Test user loading when user has no tenants at all."""
|
||||
@ -517,6 +525,68 @@ class TestAccountService:
|
||||
# Verify results
|
||||
assert result is None
|
||||
|
||||
def test_refresh_account_last_active_uses_redis_gate_and_conditional_update(self, mock_db_dependencies):
|
||||
"""Test last-active refresh is gated in Redis and conditionally written to DB."""
|
||||
mock_account = TestAccountAssociatedDataFactory.create_account_mock()
|
||||
now = datetime(2026, 6, 2, 2, 45, 49)
|
||||
mock_account.last_active_at = now - timedelta(minutes=15)
|
||||
|
||||
with (
|
||||
patch("services.account_service.naive_utc_now", return_value=now),
|
||||
patch("services.account_service.redis_client") as mock_redis_client,
|
||||
):
|
||||
mock_redis_client.set.return_value = True
|
||||
|
||||
AccountService._refresh_account_last_active(mock_account)
|
||||
|
||||
mock_redis_client.set.assert_called_once_with(
|
||||
"account_last_active_refresh:user-123",
|
||||
1,
|
||||
ex=600,
|
||||
nx=True,
|
||||
)
|
||||
mock_db_dependencies["db"].session.execute.assert_called_once()
|
||||
mock_db_dependencies["db"].session.commit.assert_called_once()
|
||||
|
||||
def test_refresh_account_last_active_skips_db_when_redis_gate_exists(self, mock_db_dependencies):
|
||||
"""Test concurrent refresh attempts do not enqueue duplicate DB updates."""
|
||||
mock_account = TestAccountAssociatedDataFactory.create_account_mock()
|
||||
now = datetime(2026, 6, 2, 2, 45, 49)
|
||||
mock_account.last_active_at = now - timedelta(minutes=15)
|
||||
|
||||
with (
|
||||
patch("services.account_service.naive_utc_now", return_value=now),
|
||||
patch("services.account_service.redis_client") as mock_redis_client,
|
||||
):
|
||||
mock_redis_client.set.return_value = None
|
||||
|
||||
AccountService._refresh_account_last_active(mock_account)
|
||||
|
||||
mock_redis_client.set.assert_called_once_with(
|
||||
"account_last_active_refresh:user-123",
|
||||
1,
|
||||
ex=600,
|
||||
nx=True,
|
||||
)
|
||||
mock_db_dependencies["db"].session.execute.assert_not_called()
|
||||
mock_db_dependencies["db"].session.commit.assert_not_called()
|
||||
|
||||
def test_refresh_account_last_active_skips_recent_account(self, mock_db_dependencies):
|
||||
"""Test recent activity does not touch Redis or DB."""
|
||||
mock_account = TestAccountAssociatedDataFactory.create_account_mock()
|
||||
now = datetime(2026, 6, 2, 2, 45, 49)
|
||||
mock_account.last_active_at = now - timedelta(minutes=5)
|
||||
|
||||
with (
|
||||
patch("services.account_service.naive_utc_now", return_value=now),
|
||||
patch("services.account_service.redis_client") as mock_redis_client,
|
||||
):
|
||||
AccountService._refresh_account_last_active(mock_account)
|
||||
|
||||
mock_redis_client.set.assert_not_called()
|
||||
mock_db_dependencies["db"].session.execute.assert_not_called()
|
||||
mock_db_dependencies["db"].session.commit.assert_not_called()
|
||||
|
||||
|
||||
class TestTenantService:
|
||||
"""
|
||||
|
||||
Loading…
Reference in New Issue
Block a user