Revive token socket_record after expiration

[masenf]

If the socket_record expires in redis, but the instance managing it is still alive, then re-link the token to the sid.

For del and evicted, do not re-link the token as this is an explicit disconnect.

[codspeed-hq]

CodSpeed Performance Report

Merging #5977 will not alter performance

_{Comparing masenf/token-lifetime (f8e9110) with main (3617ee4)}

Summary

✅ 8 untouched

[greptile-apps]

Greptile Overview

Greptile Summary

This PR implements token revival logic when Redis records expire while the managing instance is still alive. The change also fixes a critical bug in reflex/utils/token_manager.py:243 where the condition was checking socket_record.instance_id != self.instance_id (not equal), which would only clean up tokens owned by other instances rather than the current instance.

Key Changes:

• Made _handle_socket_record_del async and added an expired parameter
• Fixed instance ID comparison from != to == on line 243
• Added token revival logic that calls link_token_to_sid when expired=True
• Distinguished between expired events (which trigger revival) and del/evicted events (which don't)

Issues Found:

• Race condition in reflex/utils/token_manager.py:242-247 where the token is popped from token_to_socket before revival, creating a window where the socket record is missing

Confidence Score: 3/5

• This PR has a critical bug fix but introduces a race condition
• Score reflects the important bug fix (instance_id comparison) and useful feature (token revival), but is lowered due to the race condition where the token is popped before being re-inserted during revival
• reflex/utils/token_manager.py needs attention for the race condition in the token revival logic

Important Files Changed

File Analysis

Filename	Score	Overview
reflex/utils/token_manager.py	3/5	Adds token revival after expiration, changes instance_id check logic, and makes _handle_socket_record_del async; one race condition identified

Sequence Diagram

sequenceDiagram
    participant Redis
    participant Subscriber as Redis Keyspace Subscriber
    participant TokenManager as RedisTokenManager
    participant LocalDict as token_to_socket dict

    Note over Redis: Token expires in Redis
    Redis->>Subscriber: keyspace event: "expired"
    Subscriber->>TokenManager: _handle_socket_record_del(token, expired=True)
    TokenManager->>LocalDict: pop(token)
    LocalDict-->>TokenManager: socket_record
    Note over TokenManager: Check instance_id == self.instance_id
    alt expired=True
        TokenManager->>TokenManager: link_token_to_sid(token, sid)
        TokenManager->>Redis: set(token_key, socket_record, ex=...)
        TokenManager->>LocalDict: token_to_socket[token] = socket_record
        Note over TokenManager: Token revived!
    end
    TokenManager->>LocalDict: sid_to_token.pop(sid)
    
    Note over Redis: Explicit delete
    Redis->>Subscriber: keyspace event: "del"
    Subscriber->>TokenManager: _handle_socket_record_del(token, expired=False)
    TokenManager->>LocalDict: pop(token)
    LocalDict-->>TokenManager: socket_record
    Note over TokenManager: Check instance_id == self.instance_id
    alt expired=False
        Note over TokenManager: No revival - token stays deleted
    end
    TokenManager->>LocalDict: sid_to_token.pop(sid)

Loading

[greptile-apps]

_{1 file reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

logic: the token has been popped from token_to_socket on line 242, but then re-inserted during link_token_to_sid on line 247 - this creates a race where the socket_record is temporarily missing

Suggested change

	socket_record := self.token_to_socket.pop(token, None)
	) is not None and socket_record.instance_id == self.instance_id:
	if expired:
	# Keep the record alive as long as this process is alive and not deleted.
	await self.link_token_to_sid(token, socket_record.sid)
	self.sid_to_token.pop(socket_record.sid, None)

Prompt To Fix With AI

This is a comment left during a code review.
Path: reflex/utils/token_manager.py
Line: 242:242

Comment:
**logic:** the token has been popped from `token_to_socket` on line 242, but then re-inserted during `link_token_to_sid` on line 247 - this creates a race where the socket_record is temporarily missing

```suggestion
        ) is not None and socket_record.instance_id == self.instance_id:
            if expired:
                # Keep the record alive as long as this process is alive and not deleted.
                await self.link_token_to_sid(token, socket_record.sid)
            self.sid_to_token.pop(socket_record.sid, None)
```

How can I resolve this? If you propose a fix, please make it concise.