fix(tokens): claim redis socket tokens atomically

[anxkhn]

All Submissions:

• Have you followed the guidelines stated in CONTRIBUTING.md file?
• Have you checked to ensure there aren't any other open Pull Requests for the desired changed?

Type of change

• Bug fix (non-breaking change which fixes an issue)

Description

RedisTokenManager.link_token_to_sid previously checked whether a token existed
before storing its socket record in a separate Redis command. Concurrent workers
could both observe an absent token and then associate it with different sockets,
allowing two connections to share the same client state.

This change claims each token with an atomic Redis SET using NX and the existing
expiration. A worker that loses the claim retries with a new UUID, while Redis
errors continue to use the local token-manager fallback.

The regression test coordinates two managers against shared Redis state and verifies
that only one keeps the requested token. It fails against the previous implementation
and passes with the atomic claim.

Testing

• uv run pytest tests/units/utils/test_token_manager.py -q (28 passed, 4 skipped)
• uv run pytest tests/units --cov --no-cov-on-fail --cov-report= (6625 passed,
  17 skipped, 78.04% coverage)
• uv run ruff check .
• uv run ruff format . --check
• uv run pyright reflex tests (0 errors)
• uv run towncrier check --config pyproject.toml --compare-with origin/main

Changes To Core Features:

• Have you added an explanation of what your changes do and why you'd like us to include them?
• Have you written new tests for your core changes, as applicable?
• Have you successfully ran tests with your changes locally?

[codspeed-hq]

Merging this PR will not alter performance

✅ 26 untouched benchmarks
⏩ 8 skipped benchmarks¹

---

_{Comparing anxkhn:fix/atomic-redis-token-claim (553fbf1) with main (65a2889)}

Footnotes

1. 8 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩

[greptile-apps]

Greptile Summary

This PR fixes a TOCTOU race in RedisTokenManager.link_token_to_sid by collapsing the old check-then-set pattern (EXISTS → SET) into a single atomic SET ... NX call, so only one worker can ever claim a given token and concurrent workers that lose the claim transparently retry with a fresh UUID.

• reflex/utils/token_manager.py: The while not await redis.set(..., nx=True) loop atomically claims the token; losers generate a new UUID and retry. The exception path falls back to the local manager correctly, returning whichever token (original or retried) was in flight at the time of failure.
• tests/units/utils/test_token_manager.py: Unit tests are updated for the new SET NX API and a new integration-style test (test_link_token_to_sid_claims_token_atomically) verifies the no-double-claim invariant, though it retains unused scaffolding (exists function, both_checked event) from the old approach that should be removed.

Confidence Score: 4/5

The production fix is correct and safe to merge; the core atomic-claim logic works as intended and error fallback paths are preserved.

The implementation change is straightforward and well-targeted. The retry loop has no iteration cap, which is benign in practice with UUID4 but could cause a hang under unusual Redis proxy behaviour. The new concurrency test carries dead coordination code from the old design, making the test intent misleading without affecting correctness or coverage.

tests/units/utils/test_token_manager.py — the new atomicity test contains unused async exists coordination code that should be cleaned up.

Important Files Changed

Filename	Overview
reflex/utils/token_manager.py	Replaces the non-atomic exists+set pattern with a single SET NX call; the logic is correct and the fallback path is preserved, though the retry loop has no iteration cap.
tests/units/utils/test_token_manager.py	Tests updated to reflect new SET NX API; the new concurrency test contains dead coordination code (exists function, both_checked event) carried over from the old implementation.
news/6740.bugfix.md	Changelog entry accurately describes the fix.

_{Reviews (1): Last reviewed commit: "fix(tokens): claim redis socket tokens a..." | Re-trigger Greptile}

[greptile-apps]

Dead code from old implementation left in concurrency test

The exists async function, both_checked event, and exists_calls counter are never exercised because the new link_token_to_sid no longer calls redis.exists at all. mock_redis.exists.side_effect = exists is also set but never triggered. These are artefacts from the previous check-then-set design. As written, the test only exercises the sequential path through set_value (a synchronous side effect with no await), which means the event loop does not interleave the two tasks — the test passes, but for a different reason than the comment implies. Removing the dead coordination code and documenting that the test validates the NX-wins-first-caller contract would make the intent clearer.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[greptile-apps]

Unbounded retry loop on persistent Redis False returns

The while not await self.redis.set(..., nx=True) loop has no iteration cap. UUID4 collision is practically impossible, so a legitimate livelock is essentially inconceivable. However, if Redis misbehaves and returns falsy (e.g. due to a proxy that strips the NX response or a mock returning 0 instead of raising), the loop will spin indefinitely without any timeout or circuit-breaker. A guard such as a maximum retry count with an emergency Exception raise would prevent a silent hang in such edge cases.