fix(realtime): send the access token on every join and push it on every session change

sdk-compliance.yaml

@@ -407,8 +407,8 @@ features:
 
   client.authentication_integration.third_party_auth: not_implemented
   client.authentication_integration.cross_client_token_sync:
-    status: partially_implemented
-    note: "Every auth event that carries a session sets the Authorization header used by auth, postgrest, storage and functions to that session's access token, and an event without a session resets it to the project key. AsyncClient also passes the token to realtime. The sync Client has no realtime client. SupabaseClient.with_token returns a client that shares the HTTP session and options and sends a given user JWT to every sub-client, for servers that act for one user per request. Like a supabase-js client with accessToken, it has no auth client: its auth raises."
+    status: implemented
+    note: "Every auth event that carries a session, including INITIAL_SESSION, sets the Authorization header used by auth, postgrest, storage and functions to that session's access token, and an event without a session resets it to the project key. AsyncClient also sends the token to the joined realtime channels right away, and every channel join carries it. The sync Client has no realtime client. SupabaseClient.with_token returns a client that shares the HTTP session and options and sends a given user JWT to every sub-client, for servers that act for one user per request. Like a supabase-js client with accessToken, it has no auth client: its auth raises."
     symbols:
       - SupabaseClient.with_token
   client.authentication_integration.oauth_flow_type:
@@ -1041,7 +1041,7 @@ features:
       - validate_heartbeat
   realtime.configuration.access_token_callback:
     status: implemented
-    note: "token_callback parameter on connect_once, awaited on connect and on each heartbeat. AsyncClient.realtime passes a callback that returns the access token of the current session, or the project key when there is none (see cross_client_token_sync)."
+    note: "token_callback parameter on connect_once and RealtimeClient, awaited on connect, on every channel join and on each heartbeat. Every join carries the token, and a token that changes while a join is pending is sent once the join completes. AsyncClient.realtime passes a callback that returns the access token of the current session, or the project key when there is none (see cross_client_token_sync)."
     symbols:
       - connect_once
   realtime.configuration.deferred_disconnect: not_implemented

src/realtime/README.md

@@ -78,7 +78,7 @@ asyncio.run(main())
 ### Notes:
 
 - `REALTIME_URL` is `http://localhost:54321/realtime/v1` when developing locally with the Supabase CLI and `https://<project_ref>.supabase.co/realtime/v1` when connecting to your Supabase project. The client switches to `ws`/`wss` and appends `/websocket` itself.
-- The token callback is awaited when connecting and on every heartbeat, so it can return a refreshed user access token. It must return a JWT whose claims contain `exp` and `role` (an existing database role), or the publishable key.
+- The token callback is awaited when connecting, on every channel join and on every heartbeat, so it can return a refreshed user access token. `client.set_auth()` awaits it right away and sends a new token to the joined channels. A token passed as `client.set_auth(token)` is used instead of the callback until `client.set_auth()` is called without one. It must return a JWT whose claims contain `exp` and `role` (an existing database role), or the publishable key.
 - `async with client.channel(...)` subscribes on entry and unsubscribes on exit. `subscribe()` raises if the server rejects the join.
 - `RealtimeClient(url, get_token)` creates a client without connecting. It connects when the first channel subscribes, shares that connection with every channel, and closes it when the last channel is unsubscribed or `aclose()` is called.
 - Channel name can be any `string`.

src/realtime/src/realtime/channel.py

@@ -228,6 +228,7 @@ async def subscribe(self) -> ReplyMessage:
         self._subscribe_attempt = attempt
         self._raise_if_topic_is_taken()
         await self.socket.connect()
+        await self.socket._refresh_token()
         self._raise_if_join_was_stopped(attempt)
         self._raise_if_topic_is_taken()
         self.socket.channels[self.topic] = self
@@ -252,9 +253,7 @@ async def subscribe(self) -> ReplyMessage:
             msg = await self.socket._send(message)
         except asyncio.CancelledError:
             self.state = ChannelStates.CLOSED
-            if self._subscribe_attempt is attempt:
-                self._subscribe_attempt = None
-                await self._leave_cancelled_join(join_ref, failure)
+            await self._stop_cancelled_join(attempt, join_ref, failure)
             raise
         except Exception:
             self.state = ChannelStates.CLOSED
@@ -268,6 +267,11 @@ async def subscribe(self) -> ReplyMessage:
                 f"error while subscribing to channel: {msg.payload.response!r}"
             )
         self.state = ChannelStates.JOINED
+        try:
+            await self.socket._send_token_if_changed(self, payload.access_token)
+        except asyncio.CancelledError:
+            await self._stop_cancelled_join(attempt, join_ref, failure)
+            raise
         return msg
 
     def _raise_if_join_was_stopped(self, attempt: object) -> None:
@@ -277,6 +281,18 @@ def _raise_if_join_was_stopped(self, attempt: object) -> None:
                 "or a closed connection before it completed"
             )
 
+    async def _stop_cancelled_join(
+        self,
+        attempt: object,
+        join_ref: str,
+        failure: asyncio.Future[Message] | None,
+    ) -> None:
+        if self._subscribe_attempt is not attempt:
+            return
+        self.state = ChannelStates.CLOSED
+        self._subscribe_attempt = None
+        await self._leave_cancelled_join(join_ref, failure)
+
     async def _leave_cancelled_join(
         self, join_ref: str, failure: asyncio.Future[Message] | None
     ) -> None:

src/realtime/src/realtime/client.py

@@ -159,7 +159,7 @@ class RealtimeClient:
     disconnects when the last channel is removed.
 
     :param url: The realtime URL of the project
-    :param token_callback: Returns the access token, awaited on connect and on every heartbeat
+    :param token_callback: Returns the access token, awaited on connect, on every channel join and on every heartbeat
     :param heartbeat_interval: Seconds between heartbeats, must be positive
     :param heartbeat_timeout: Seconds to wait for a heartbeat reply before the connection is considered lost, must be positive
     """
@@ -182,6 +182,8 @@ def __init__(
         self._listen_task: asyncio.Task[None] | None = None
         self._heartbeat_task: asyncio.Task[None] | None = None
         self._connect_lock = asyncio.Lock()
+        self._auth_lock = asyncio.Lock()
+        self._token_is_manual = False
 
         self.ref = 0
         self.channels: dict[str, RealtimeChannel] = {}
@@ -210,7 +212,7 @@ async def connect(self) -> None:
             websocket = await websockets.connect(
                 str(normalize_url(self.url, token)), logger=websockets_logger
             )
-            if token and self.last_token is not None:
+            if token and self.last_token is not None and not self._token_is_manual:
                 self.last_token = token
             failure: asyncio.Future[Message] = (
                 asyncio.get_running_loop().create_future()
@@ -318,7 +320,7 @@ async def _heartbeat(self, failure: asyncio.Future[Message]) -> None:
                     self._send(data), timeout=self.heartbeat_timeout
                 )
                 logger.info("heartbeat response: %r", res)
-                await self.set_auth()
+                await self._refresh_token()
                 await asyncio.sleep(self.heartbeat_interval)
         except asyncio.CancelledError:
             return
@@ -379,16 +381,18 @@ async def remove_all_channels(self) -> None:
 
     async def set_auth(self, token: str | None = None) -> None:
         """
-        Set the authentication token for the connection and update all joined channels.
+        Set the access token of the connection and send it to every joined channel.
 
-        This method updates the access token for the current connection and sends the new token
-        to all joined channels. This is useful for refreshing authentication or changing users.
+        Channels joined later send it in their join message. Tokens are sent
+        in the order of the calls.
 
         Args:
-            token (Optional[str]): The new authentication token. Can be None to remove authentication.
-
-        Returns:
-            None
+            token (Optional[str]): The new access token. The token callback no
+                longer replaces it on joins, heartbeats and reconnects, until
+                this method is called without a token. When None, the token
+                callback is awaited for the token and replaces it again from
+                then on. Nothing changes if there is no callback or it returns
+                no token.
 
         Raises:
             websockets.ConnectionClosed: If the connection is closed. The token is
@@ -397,27 +401,58 @@ async def set_auth(self, token: str | None = None) -> None:
             Exception: If the connection already failed, the exception that failed
                 it, such as a TimeoutError when a heartbeat timed out.
         """
-        if token is None:
-            if self.get_access_token is not None and (
-                new_token := await self.get_access_token()
-            ):
-                token = new_token
+        if token is None and self.get_access_token is None:
+            logger.warning("tried setting auth but no token callback was registered")
+            return
+        async with self._auth_lock:
+            self._token_is_manual = token is not None
+            if token is None:
+                await self._send_callback_token()
             else:
-                logger.warning(
-                    "tried setting auth but no token callback was registered"
-                )
-                return
+                await self._send_token(token)
+
+    async def _refresh_token(self) -> None:
+        if self.get_access_token is None:
+            return
+        async with self._auth_lock:
+            if not self._token_is_manual:
+                await self._send_callback_token()
+
+    async def _send_callback_token(self) -> None:
+        assert self.get_access_token is not None
+        token = await self.get_access_token()
+        if token:
+            await self._send_token(token)
+
+    async def _send_token(self, token: str) -> None:
         if token == self.last_token:
             logger.info("generated access token is still the same, skipping")
             return
-
         for channel in list(self.channels.values()):
             if channel.joined:
                 await channel._send_event_no_wait(
                     ChannelEvents.access_token, {"access_token": token}
                 )
         self.last_token = token
 
+    async def _send_token_if_changed(
+        self, channel: RealtimeChannel, joined_with: str | None
+    ) -> None:
+        async with self._auth_lock:
+            token = self.last_token
+            if token is None or token == joined_with or not channel.joined:
+                return
+            try:
+                await channel._send_event_no_wait(
+                    ChannelEvents.access_token, {"access_token": token}
+                )
+            except Exception as exception:
+                logger.warning(
+                    "could not send the new access token to %s: %r",
+                    channel.topic,
+                    exception,
+                )
+
     def _make_ref(self) -> str:
         self.ref += 1
         return f"{self.ref}"