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. ↩

2. No successful run was found on main (127e4d8) during the generation of this report, so 65a2889 was used instead as the comparison base. There might be some changes unrelated to this pull request in this report. ↩

[greptile-apps]

Greptile Summary

This PR replaces the previous check-then-set (EXISTS → SET) pattern in RedisTokenManager.link_token_to_sid with a single atomic SET NX, eliminating the TOCTOU window where two concurrent workers could both observe an absent key and then each bind it to a different socket.

• Core fix (token_manager.py): The while loop retries with a fresh UUID whenever SET NX returns falsy (key already owned), and falls back to super().link_token_to_sid() on any Redis exception, preserving the original degraded-mode behaviour. The return fallback_token or new_token expression correctly handles all four combinations of fallback_token/new_token being set or None.
• Tests: Existing tests are updated to mock SET (with nx=True) rather than EXISTS. A new concurrency test (test_link_token_to_sid_claims_token_atomically) validates the NX-wins-first-caller contract over shared in-process state.

Confidence Score: 5/5

The change is narrowly scoped to the token-claiming path, the atomic SET NX approach is the standard Redis solution to this class of race condition, and the fallback to local-only storage on Redis errors preserves the pre-existing degraded-mode contract.

The implementation handles all code paths correctly: successful atomic claim, key-already-owned retry, and Redis exception fallback. Local-dict writes are always consistent with what was claimed. No data-loss or incorrect-state scenario was identified in the changed code.

tests/units/utils/test_token_manager.py — the dead coordination scaffolding in the concurrency test (noted in a prior thread) is the one area worth cleaning up before merge.

Important Files Changed

Filename	Overview
reflex/utils/token_manager.py	Replaces the old two-step check-then-set with an atomic Redis SET NX, correctly eliminating the TOCTOU race. The retry loop, fallback on exception, and local-dict synchronisation all appear logically sound.
tests/units/utils/test_token_manager.py	Tests updated to match the new NX-based contract. The new concurrency test verifies mutual exclusion over shared state; the dead coordination scaffolding (exists, both_checked, exists_calls) left from the prior design is noted in a previous thread.
news/6740.bugfix.md	Changelog entry accurately describes the fix.

_{Reviews (2): 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.

[Alek99]

Hi appreciate the contributions but can you prioritize Github issue also please address the AI review before marking it out of draft