From 888483a2f80300bacb40d7b5567b15e8fbb8e167 Mon Sep 17 00:00:00 2001 From: zyssyz123 <916125788@qq.com> Date: Tue, 2 Jun 2026 15:20:07 +0800 Subject: [PATCH] fix: user token (#36930) --- api/services/account_service.py | 40 +++++++++- .../services/test_account_service.py | 74 ++++++++++++++++++- 2 files changed, 109 insertions(+), 5 deletions(-) diff --git a/api/services/account_service.py b/api/services/account_service.py index 6705bdc4e6..89a0a5d8e2 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -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) diff --git a/api/tests/unit_tests/services/test_account_service.py b/api/tests/unit_tests/services/test_account_service.py index e09102a788..f18e5a61b8 100644 --- a/api/tests/unit_tests/services/test_account_service.py +++ b/api/tests/unit_tests/services/test_account_service.py @@ -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: """