[PR #32195] fix: add distributed lock and fix session persistence for OAuth token refresh #33605

Open
opened 2026-02-21 20:53:35 -05:00 by yindo · 0 comments
Owner

Original Pull Request: https://github.com/langgenius/dify/pull/32195

State: open
Merged: No


Summary

Fixes two critical bugs in the OAuth token refresh logic that cause:

  1. Refreshed credentials never persist to the databasedb.session.commit() after a successful OAuth refresh does not persist in Celery worker and gevent greenlet contexts because Flask-SQLAlchemy's scoped db.session is not reliable in these contexts. Every subsequent request re-reads the stale expires_at and re-triggers the refresh, leading to hundreds of unnecessary OAuth calls per day.

  2. No distributed lock on concurrent refresh — Multiple concurrent workers (e.g., GraphWorker-0 and GraphWorker-1 in workflow execution) can simultaneously detect expired credentials and call the OAuth provider to refresh. With providers that use refresh token rotation (e.g., QuickBooks), this leads to invalid_grant errors and permanent failure requiring manual re-authorization.

Together, these bugs create a cascading failure: the token is refreshed hundreds of times but never saved, until the original refresh token is eventually invalidated by the OAuth provider.

Changes

api/core/tools/tool_manager.py

  • Add non-blocking Redis distributed lock around the OAuth refresh block
  • Replace db.session.refresh() / db.session.commit() with explicit Session(db.engine) and session.get() / session.commit() to ensure persistence in all execution contexts
  • Add double-check after lock acquisition (another request may have already refreshed)
  • Add exponential backoff polling for requests that don't acquire the lock

api/services/datasource_provider_service.py

  • Add the same non-blocking Redis distributed lock + double-check pattern
  • This file already uses explicit Session(db.engine), so only the lock was needed

New test files

  • api/tests/unit_tests/core/tools/test_tool_manager_oauth_refresh_lock.py — 6 tests covering lock acquired + refresh, double-check skip, lock not acquired + polling, lock release on exception, non-expired skip, never-expires skip
  • api/tests/unit_tests/services/test_datasource_provider_oauth_refresh_lock.py — 6 matching tests for datasource provider

Evidence from production

Before fix (24 hours):

  • 693 OAuth refresh HTTP calls (expected: ~24)
  • Database updated_at == created_at — credentials never persisted
  • Concurrent refresh race condition caused permanent invalid_grant failure

After fix (5.5 hours):

  • 3 OAuth refresh calls (1 manual test + 2 automatic hourly refreshes)
  • Database updated_at correctly updated on each refresh
  • Zero errors, zero concurrent conflicts
  • All workflows executing successfully

Related issue

Closes #32174

Test plan

  • 6 unit tests for tool_manager.py OAuth refresh lock — all pass
  • 6 unit tests for datasource_provider_service.py OAuth refresh lock — all pass
  • 78 existing tools tests — zero regressions
  • Production deployment verified: token refresh persists, lock prevents concurrent refresh
**Original Pull Request:** https://github.com/langgenius/dify/pull/32195 **State:** open **Merged:** No --- ## Summary Fixes two critical bugs in the OAuth token refresh logic that cause: 1. **Refreshed credentials never persist to the database** — `db.session.commit()` after a successful OAuth refresh does not persist in Celery worker and gevent greenlet contexts because Flask-SQLAlchemy's scoped `db.session` is not reliable in these contexts. Every subsequent request re-reads the stale `expires_at` and re-triggers the refresh, leading to hundreds of unnecessary OAuth calls per day. 2. **No distributed lock on concurrent refresh** — Multiple concurrent workers (e.g., `GraphWorker-0` and `GraphWorker-1` in workflow execution) can simultaneously detect expired credentials and call the OAuth provider to refresh. With providers that use refresh token rotation (e.g., QuickBooks), this leads to `invalid_grant` errors and permanent failure requiring manual re-authorization. Together, these bugs create a cascading failure: the token is refreshed hundreds of times but never saved, until the original refresh token is eventually invalidated by the OAuth provider. ## Changes ### `api/core/tools/tool_manager.py` - Add non-blocking Redis distributed lock around the OAuth refresh block - Replace `db.session.refresh()` / `db.session.commit()` with explicit `Session(db.engine)` and `session.get()` / `session.commit()` to ensure persistence in all execution contexts - Add double-check after lock acquisition (another request may have already refreshed) - Add exponential backoff polling for requests that don't acquire the lock ### `api/services/datasource_provider_service.py` - Add the same non-blocking Redis distributed lock + double-check pattern - This file already uses explicit `Session(db.engine)`, so only the lock was needed ### New test files - `api/tests/unit_tests/core/tools/test_tool_manager_oauth_refresh_lock.py` — 6 tests covering lock acquired + refresh, double-check skip, lock not acquired + polling, lock release on exception, non-expired skip, never-expires skip - `api/tests/unit_tests/services/test_datasource_provider_oauth_refresh_lock.py` — 6 matching tests for datasource provider ## Evidence from production Before fix (24 hours): - **693** OAuth refresh HTTP calls (expected: ~24) - Database `updated_at == created_at` — credentials never persisted - Concurrent refresh race condition caused permanent `invalid_grant` failure After fix (5.5 hours): - **3** OAuth refresh calls (1 manual test + 2 automatic hourly refreshes) - Database `updated_at` correctly updated on each refresh - Zero errors, zero concurrent conflicts - All workflows executing successfully ## Related issue Closes #32174 ## Test plan - [x] 6 unit tests for `tool_manager.py` OAuth refresh lock — all pass - [x] 6 unit tests for `datasource_provider_service.py` OAuth refresh lock — all pass - [x] 78 existing tools tests — zero regressions - [x] Production deployment verified: token refresh persists, lock prevents concurrent refresh
yindo added the pull-request label 2026-02-21 20:53:35 -05:00
Sign in to join this conversation.
1 Participants
Notifications
Due Date
No due date set.
Dependencies

No dependencies set.

Reference: langgenius/dify#33605