Skip to content

fix(tokens): claim redis socket tokens atomically #6771

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
anxkhn wants to merge 1 commit into
base: main
Choose a base branch
from
+67 −36
Open

fix(tokens): claim redis socket tokens atomically #6771

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions news/6740.bugfix.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Redis-backed socket token claims are now atomic across workers, preventing simultaneous connections from sharing client state.
37 changes: 14 additions & 23 deletions reflex/utils/token_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -309,36 +309,27 @@ async def link_token_to_sid(self, token: str, sid: str) -> str | None:
# Make sure the update subscriber is running
self._ensure_socket_record_task()

# Check Redis for cross-worker duplicates
redis_key = self._get_redis_key(token)
socket_record = SocketRecord(instance_id=self.instance_id, sid=sid)
socket_record_data = pickle.dumps(socket_record)
new_token = None

try:
token_exists_in_redis = await self.redis.exists(redis_key)
while not await self.redis.set(
self._get_redis_key(token),
socket_record_data,
ex=self.token_expiration,
nx=True,
):
token = new_token = _get_new_token()
Comment on lines +317 to +323

@greptile-apps greptile-apps Bot Jul 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

except Exception as e:
console.error(f"Redis error checking token existence: {e}")
return await super().link_token_to_sid(token, sid)

new_token = None
if token_exists_in_redis:
# Duplicate exists somewhere - generate new token
token = new_token = _get_new_token()
redis_key = self._get_redis_key(new_token)
console.error(f"Redis error claiming token: {e}")
fallback_token = await super().link_token_to_sid(token, sid)
return fallback_token or new_token

# Store in local dicts
socket_record = self.token_to_socket[token] = SocketRecord(
instance_id=self.instance_id, sid=sid
)
self.token_to_socket[token] = socket_record
self.sid_to_token[sid] = token

# Store in Redis if possible
try:
await self.redis.set(
redis_key,
pickle.dumps(socket_record),
ex=self.token_expiration,
)
except Exception as e:
console.error(f"Redis error storing token: {e}")
# Return the new token if one was generated
return new_token

Expand Down
65 changes: 52 additions & 13 deletions tests/units/utils/test_token_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -291,18 +291,16 @@ async def test_link_token_to_sid_normal_case(self, manager, mock_redis):
mock_redis: Mock Redis client fixture.
"""
token, sid = "token1", "sid1"
mock_redis.exists.return_value = False
mock_redis.set.return_value = True

result = await manager.link_token_to_sid(token, sid)

assert result is None
mock_redis.exists.assert_called_once_with(
f"token_manager_socket_record_{token}"
)
mock_redis.set.assert_called_once_with(
f"token_manager_socket_record_{token}",
pickle.dumps(SocketRecord(instance_id=manager.instance_id, sid=sid)),
ex=3600,
nx=True,
)
assert manager.token_to_socket[token].sid == sid
assert manager.sid_to_token[sid] == token
Expand All @@ -324,7 +322,6 @@ async def test_link_token_to_sid_reconnection_skips_redis(
result = await manager.link_token_to_sid(token, sid)

assert result is None
mock_redis.exists.assert_not_called()
mock_redis.set.assert_not_called()

async def test_link_token_to_sid_duplicate_detected(self, manager, mock_redis):
Expand All @@ -335,21 +332,20 @@ async def test_link_token_to_sid_duplicate_detected(self, manager, mock_redis):
mock_redis: Mock Redis client fixture.
"""
token, sid = "token1", "sid1"
mock_redis.exists.return_value = True
mock_redis.set.side_effect = [False, True]

result = await manager.link_token_to_sid(token, sid)

assert result is not None
assert result != token
assert len(result) == 36 # UUID4 length

mock_redis.exists.assert_called_once_with(
f"token_manager_socket_record_{token}"
)
mock_redis.set.assert_called_once_with(
assert mock_redis.set.await_count == 2
mock_redis.set.assert_awaited_with(
f"token_manager_socket_record_{result}",
pickle.dumps(SocketRecord(instance_id=manager.instance_id, sid=sid)),
ex=3600,
nx=True,
)
assert manager.token_to_sid[result] == sid
assert manager.sid_to_token[sid] == result
Expand All @@ -362,7 +358,7 @@ async def test_link_token_to_sid_redis_error_fallback(self, manager, mock_redis)
mock_redis: Mock Redis client fixture.
"""
token, sid = "token1", "sid1"
mock_redis.exists.side_effect = Exception("Redis connection error")
mock_redis.set.side_effect = Exception("Redis connection error")

with patch.object(
LocalTokenManager, "link_token_to_sid", new_callable=AsyncMock
Expand All @@ -384,7 +380,6 @@ async def test_link_token_to_sid_redis_set_error_continues(
mock_redis: Mock Redis client fixture.
"""
token, sid = "token1", "sid1"
mock_redis.exists.return_value = False
mock_redis.set.side_effect = Exception("Redis set error")

result = await manager.link_token_to_sid(token, sid)
Expand All @@ -393,6 +388,50 @@ async def test_link_token_to_sid_redis_set_error_continues(
assert manager.token_to_sid[token] == sid
assert manager.sid_to_token[sid] == token

async def test_link_token_to_sid_claims_token_atomically(self, mock_redis):
"""Test concurrent managers cannot both claim the same token.

Args:
mock_redis: Mock Redis client fixture.
"""
redis_values = {}
exists_calls = 0
both_checked = asyncio.Event()

async def exists(key):
nonlocal exists_calls
result = key in redis_values
exists_calls += 1
if exists_calls == 2:
both_checked.set()
await both_checked.wait()
return result

def set_value(key, value, *, ex, nx=False):
if nx and key in redis_values:
return False
redis_values[key] = value
return True

mock_redis.exists.side_effect = exists
Comment on lines +397 to +416

@greptile-apps greptile-apps Bot Jul 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 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!

mock_redis.set.side_effect = set_value
with patch("reflex_base.config.get_config") as mock_get_config:
mock_get_config.return_value.redis_token_expiration = 3600
managers = [RedisTokenManager(mock_redis), RedisTokenManager(mock_redis)]

with patch.object(RedisTokenManager, "_ensure_socket_record_task"):
results = await asyncio.gather(
managers[0].link_token_to_sid("token1", "sid1"),
managers[1].link_token_to_sid("token1", "sid2"),
)

assert sum(result is None for result in results) == 1
claimed_tokens = {
manager.sid_to_token[sid]
for manager, sid in zip(managers, ("sid1", "sid2"), strict=True)
}
assert len(claimed_tokens) == 2

async def test_disconnect_token_owned_locally(self, manager, mock_redis):
"""Test disconnect cleans up both Redis and local mappings when owned locally.

Expand Down Expand Up @@ -465,7 +504,7 @@ async def test_various_redis_errors_handled_gracefully(
redis_error: Exception to test error handling.
"""
token, sid = "token1", "sid1"
mock_redis.exists.side_effect = redis_error
mock_redis.set.side_effect = redis_error

with patch.object(
LocalTokenManager, "link_token_to_sid", new_callable=AsyncMock
Expand Down
Loading