feat: Expose TokenManager with lifecycle APIs for session management

[timon0305]

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

• New feature (non-breaking change which adds functionality)

New Feature Submission:

• Does your submission pass the tests?
• Have you linted your code locally prior to submission?

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?

---

Description

Exposes rx.get_token_manager() with new lifecycle APIs for handling client/session lifecycle within background events, as requested in #5669.

New APIs on TokenManager:

Method	Description
session_is_connected(sid)	Async iterator yielding client token until the session disconnects
token_is_connected(client_token)	Async iterator yielding session ID until the token disconnects
when_session_disconnects(sid)	Returns asyncio.Event set on disconnect
when_token_disconnects(client_token)	Returns asyncio.Event set on disconnect
when_token_connects(client_token)	Returns asyncio.Event set on connect

Public accessor:

• rx.get_token_manager() — returns the active TokenManager or raises RuntimeError

Files changed (4):

• reflex/utils/token_manager.py — Added lifecycle methods, notification hooks, get_token_manager(), and _TokenNotConnectedError
• reflex/app.py — Wired _notify_connect() / _notify_disconnect() into EventNamespace.on_connect / on_disconnect
• reflex/__init__.py — Exposed get_token_manager via lazy loader mapping
• tests/units/utils/test_token_manager.py — Added 19 new tests covering all lifecycle APIs

Example usage:

@rx.event(background=True)
async def thing_checker(self):
    for _task_active in rx.get_token_manager().session_is_connected(
        sid=self.router_data.session.session_id
    ):
        async with self:
            if await self._check_thing():
                await self._update_thing()
        await asyncio.sleep(CHECK_INTERVAL)

Test results:

• 3567 tests passed (3548 existing + 19 new), 0 failures

closes #5669

[greptile-apps]

Greptile Overview

Greptile Summary

This PR exposes rx.get_token_manager() and adds lifecycle primitives on TokenManager to let background tasks watch for token/session connect and disconnect events. It wires connect/disconnect notifications into EventNamespace.on_connect / on_disconnect, exports the accessor via reflex.__init__ lazy loading, and adds unit tests covering the new APIs.

Key integration point is EventNamespace (Socket.IO namespace) calling _notify_connect() after a token is linked to a sid, and _notify_disconnect() before cleanup removes the token/sid mappings.

Confidence Score: 3/5

• This PR is close to merge-ready but has a correctness/performance issue in the new lifecycle iterators that should be fixed first.
• The new lifecycle APIs are well-scoped and tests cover the surface area, but session_is_connected() / token_is_connected() are implemented as tight async-generator loops without awaits, which will cause CPU spin and event-loop starvation under typical async for usage. Fixing those iterator semantics would significantly reduce merge risk.
• reflex/utils/token_manager.py (lifecycle iterator implementations)

Important Files Changed

Filename	Overview
reflex/utils/token_manager.py	Adds lifecycle connect/disconnect watcher APIs and get_token_manager(), but session_is_connected()/token_is_connected() are implemented as tight async-generator loops with no awaits, causing CPU spin for typical async for usage.
reflex/app.py	Wires token manager lifecycle notifications into EventNamespace.on_connect/on_disconnect; change is localized and consistent with existing token linking/cleanup flow.
reflex/init.py	Exposes get_token_manager via the lazy loader mapping; straightforward export change.
tests/units/utils/test_token_manager.py	Adds unit tests for lifecycle APIs; coverage is good, but current tests also implicitly accept the tight-loop iterator behavior.

Sequence Diagram

sequenceDiagram
    participant Client
    participant SIO as EventNamespace
    participant TM as TokenManager

    Client->>SIO: websocket connect (sid, token in query)
    alt token provided
        SIO->>TM: link_token_to_sid(token, sid)
        SIO->>TM: _notify_connect(actual_token, sid)
    else no token
        SIO-->>Client: warn (no token)
    end

    Note over Client,SIO: During session lifetime, app may run background events

    Client->>SIO: websocket disconnect (sid)
    SIO->>TM: _notify_disconnect(disconnect_token, sid)
    SIO->>TM: disconnect_token(disconnect_token, sid) (async task)
    TM-->>TM: remove token<->sid mappings

Loading

[greptile-apps]

_{4 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[masenf]

this is good, but lets remove the on token connect methods for now, since they won't work reliably across instances