livekit: Preserve tokens on channel rejoin (#59388)

- **collab: Don't revoke LiveKit token when rejoining the same channel**
- **call: Bound LiveKit reconnect attempts and decouple from collab
rejoin**

# Objective

- Describe the objective or issue this PR addresses.
- If you're fixing a specific issue, use "Fixes #X" for each issue as
[described in the GitHub
docs](https://docs.github.com/en/issues/tracking-your-work-with-issues/using-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword).

## Solution

- Describe the solution used to achieve the objective above.

## Testing

- Did you test these changes? If so, how?
- Are there any parts that need more testing?
- How can other people (reviewers) test your changes? Is there anything
specific they need to know?
- If relevant, what platforms did you test these changes on, and are
there any important ones you can't test?

## Self-Review Checklist:

- [ ] I've reviewed my own diff for quality, security, and reliability
- [ ] Unsafe blocks (if any) have justifying comments
- [ ] The content adheres to Zed's UI standards
([UX/UI](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)
and
[icon](https://github.com/zed-industries/zed/blob/main/crates/icons/README.md)
guidelines)
- [ ] Tests cover the new/changed behavior
- [ ] Performance impact has been considered and is acceptable

## Showcase

> This section is optional. If this PR does not include a visual change
or does not add a new user-facing feature, you can delete this section.

- Help others understand the result of this PR by showcasing your
awesome work!
- If this PR includes a visual change, consider adding a screenshot,
GIF, or video
- A before/after comparison is very useful for changes to existing
features!

While a showcase should aim to be brief and digestible, you can use a
toggleable section to save space on longer showcases:

<details>
  <summary>Click to view showcase</summary>

My super cool demos here

</details>

---

Release Notes:

- N/A or Added/Fixed/Improved ...

crates/call/src/call_impl/room.rs

@@ -1773,14 +1773,26 @@ impl Room {
     }
 }
 
+/// The number of times we attempt to establish the LiveKit connection before
+/// giving up. Retries cover transient failures (e.g. a network blip); a
+/// genuinely bad token won't be fixed by retrying, so the bound keeps us from
+/// hammering LiveKit indefinitely.
+const LIVEKIT_CONNECT_ATTEMPTS: usize = 5;
+
 fn spawn_room_connection(
     livekit_connection_info: Option<proto::LiveKitConnectionInfo>,
     cx: &mut Context<Room>,
 ) -> Option<Task<()>> {
-    let mut connection_info = livekit_connection_info?;
+    let connection_info = livekit_connection_info?;
     Some(cx.spawn(async move |this, cx| {
+        // Retry the LiveKit connection with backoff, but deliberately do NOT
+        // ask collab for a fresh token here: a refreshed token is only needed
+        // when the collab connection itself is re-established, which is handled
+        // separately in `maintain_connection` via `ensure_livekit_connection`.
+        // Coupling this loop to `rejoin` previously turned an isolated LiveKit
+        // failure into a channel-wide `RejoinRoom` storm.
         let mut backoff = Duration::from_secs(1);
-        loop {
+        for attempt in 1..=LIVEKIT_CONNECT_ATTEMPTS {
             match livekit::Room::connect(
                 connection_info.server_url.clone(),
                 connection_info.token.clone(),
@@ -1839,40 +1851,21 @@ fn spawn_room_connection(
                     return;
                 }
                 Err(error) => {
-                    log::error!("failed to connect to LiveKit room: {error:#}");
+                    log::error!(
+                        "failed to connect to LiveKit room (attempt {attempt}/{LIVEKIT_CONNECT_ATTEMPTS}): {error:#}"
+                    );
                 }
             }
 
-            cx.background_executor().timer(backoff).await;
-            backoff = (backoff * 2).min(Duration::from_secs(30));
-
-            // The token we were given may no longer be valid: LiveKit revokes
-            // it when this user's stale connection is cleaned up around the
-            // time the token is issued (e.g. when rejoining a channel right
-            // after a crash). Rejoin the room to obtain fresh connection info
-            // before trying again.
-            let Ok(rejoin) = this.update(cx, |this, cx| this.rejoin(cx)) else {
-                return;
-            };
-            match rejoin.await {
-                Ok(Some(new_connection_info)) => {
-                    log::info!(
-                        "refreshed LiveKit connection info (token changed: {})",
-                        new_connection_info.token != connection_info.token
-                    );
-                    connection_info = new_connection_info;
-                }
-                Ok(None) => {
-                    log::warn!(
-                        "server returned no fresh LiveKit connection info; \
-                         retrying with the existing token"
-                    );
-                }
-                Err(error) => {
-                    log::error!("failed to refresh LiveKit connection info: {error:#}");
-                }
+            if attempt < LIVEKIT_CONNECT_ATTEMPTS {
+                cx.background_executor().timer(backoff).await;
+                backoff = (backoff * 2).min(Duration::from_secs(10));
             }
         }
+
+        log::error!(
+            "giving up on LiveKit room connection after {LIVEKIT_CONNECT_ATTEMPTS} attempts"
+        );
     }))
 }
 

crates/collab/src/rpc.rs

@@ -1339,7 +1339,7 @@ async fn connection_lost(
         _ = executor.sleep(RECONNECT_TIMEOUT).fuse() => {
 
             log::info!("connection lost, removing all resources for user:{}, connection:{:?}", session.user_id(), session.connection_id);
-            leave_room_for_session(&session, session.connection_id).await.trace_err();
+            leave_room_for_session(&session, session.connection_id, None).await.trace_err();
             leave_channel_buffers_for_session(&session)
                 .await
                 .trace_err();
@@ -1702,7 +1702,7 @@ async fn leave_room(
     response: Response<proto::LeaveRoom>,
     session: MessageContext,
 ) -> Result<()> {
-    leave_room_for_session(&session, session.connection_id).await?;
+    leave_room_for_session(&session, session.connection_id, None).await?;
     response.send(proto::Ack {})?;
     Ok(())
 }
@@ -3437,7 +3437,7 @@ async fn join_channel_internal(
                 "cleaning up stale connection",
             );
             drop(db);
-            leave_room_for_session(&session, connection).await?;
+            leave_room_for_session(&session, connection, Some(channel_id)).await?;
             db = session.db().await;
         }
 
@@ -4067,7 +4067,11 @@ async fn update_user_contacts(user_id: UserId, session: &Session) -> Result<()>
     Ok(())
 }
 
-async fn leave_room_for_session(session: &Session, connection_id: ConnectionId) -> Result<()> {
+async fn leave_room_for_session(
+    session: &Session,
+    connection_id: ConnectionId,
+    rejoining_channel_id: Option<ChannelId>,
+) -> Result<()> {
     let mut contacts_to_update = HashSet::default();
 
     let room_id;
@@ -4076,6 +4080,7 @@ async fn leave_room_for_session(session: &Session, connection_id: ConnectionId)
     let delete_livekit_room;
     let room;
     let channel;
+    let left_channel_id;
 
     if let Some(mut left_room) = session.db().await.leave_room(connection_id).await? {
         contacts_to_update.insert(session.user_id());
@@ -4090,12 +4095,23 @@ async fn leave_room_for_session(session: &Session, connection_id: ConnectionId)
         delete_livekit_room = left_room.deleted;
         room = mem::take(&mut left_room.room);
         channel = mem::take(&mut left_room.channel);
+        left_channel_id = channel.as_ref().map(|channel| channel.id);
 
         room_updated(&room, &session.peer);
     } else {
         return Ok(());
     }
 
+    // When this cleanup is part of rejoining the same channel, the user is
+    // immediately re-entering this LiveKit room with the same identity, which
+    // LiveKit handles by evicting the prior session. Calling
+    // `remove_participant` here is therefore redundant, and worse: on LiveKit
+    // Cloud it revokes the identity's tokens (including the freshly issued one
+    // for the rejoin), leaving the user connected to collab but unable to join
+    // audio. Skip the removal in that case.
+    let skip_livekit_removal =
+        rejoining_channel_id.is_some() && left_channel_id == rejoining_channel_id;
+
     if let Some(channel) = channel {
         channel_updated(
             &channel,
@@ -4128,10 +4144,12 @@ async fn leave_room_for_session(session: &Session, connection_id: ConnectionId)
     }
 
     if let Some(live_kit) = session.app_state.livekit_client.as_ref() {
-        live_kit
-            .remove_participant(livekit_room.clone(), session.user_id().to_string())
-            .await
-            .trace_err();
+        if !skip_livekit_removal {
+            live_kit
+                .remove_participant(livekit_room.clone(), session.user_id().to_string())
+                .await
+                .trace_err();
+        }
 
         if delete_livekit_room {
             live_kit.delete_room(livekit_room).await.trace_err();

crates/collab/tests/integration/channel_tests.rs

@@ -367,69 +367,73 @@ async fn test_joining_channel_ancestor_member(
 }
 
 #[gpui::test]
-async fn test_channel_call_recovers_from_livekit_connect_failure(
+async fn test_rejoining_channel_does_not_remove_livekit_participant(
     executor: BackgroundExecutor,
-    cx_a: &mut TestAppContext,
+    cx_a1: &mut TestAppContext,
+    cx_a2: &mut TestAppContext,
 ) {
     let mut server = TestServer::start(executor.clone()).await;
-    let client_a = server.create_client(cx_a, "user_a").await;
-    client_a.initialize_channel_store(cx_a);
-
-    let channel_id = server
-        .make_channel("zed", None, (&client_a, cx_a), &mut [])
-        .await;
-
-    // Simulate LiveKit rejecting the token issued for this join. In
-    // production this happens when the collab server's stale connection
-    // cleanup (`leave_room_for_session`) calls `remove_participant` for this
-    // user's identity just before issuing a new token for the same identity
-    // and room (e.g. rejoining a channel right after an abrupt disconnect or
-    // restart): LiveKit Cloud revokes the freshly issued token and the
-    // client's initial connection fails with "401 invalid token: revoked".
-    let identity = client_a.user_id().unwrap().to_string();
+    // Two connections for the same user. The second one joining the same
+    // channel reproduces the production scenario where the server runs
+    // stale-connection cleanup for the first connection (as it does after an
+    // abrupt disconnect/crash followed by a rejoin within RECONNECT_TIMEOUT).
+    // Mirror LiveKit Cloud: removing a participant revokes their tokens. This
+    // is what turns the redundant `remove_participant` during stale cleanup
+    // into a user-visible failure.
     server
         .test_livekit_server
-        .set_token_revoked(&identity, true);
+        .set_revoke_tokens_on_removal(true);
 
-    let active_call_a = cx_a.read(ActiveCall::global);
-    active_call_a
-        .update(cx_a, |active_call, cx| {
+    let client_a1 = server.create_client(cx_a1, "user_a").await;
+    let client_a2 = server.create_client(cx_a2, "user_a").await;
+    client_a1.initialize_channel_store(cx_a1);
+
+    let channel_id = server
+        .make_channel("zed", None, (&client_a1, cx_a1), &mut [])
+        .await;
+
+    let active_call_a1 = cx_a1.read(ActiveCall::global);
+    active_call_a1
+        .update(cx_a1, |active_call, cx| {
             active_call.join_channel(channel_id, cx)
         })
         .await
         .unwrap();
     executor.run_until_parked();
 
-    // The collab server considers user A a participant in the call...
-    let room_a =
-        cx_a.read(|cx| active_call_a.read_with(cx, |call, _| call.room().unwrap().clone()));
-    cx_a.read(|cx| {
-        client_a.channel_store().read_with(cx, |channels, _| {
-            assert_participants_eq(
-                channels.channel_participants(channel_id),
-                &[client_a.user_id().unwrap()],
-            );
+    // The second connection joins the same channel, triggering stale-connection
+    // cleanup of the first connection on the server.
+    let active_call_a2 = cx_a2.read(ActiveCall::global);
+    active_call_a2
+        .update(cx_a2, |active_call, cx| {
+            active_call.join_channel(channel_id, cx)
         })
-    });
-    // ...but the LiveKit connection failed, so they have no audio.
-    cx_a.read(|cx| room_a.read_with(cx, |room, cx| assert!(!room.is_connected(cx))));
-
-    // Once LiveKit accepts the user's tokens again, the client should
-    // re-establish its LiveKit connection instead of silently staying in the
-    // call without audio.
-    server
-        .test_livekit_server
-        .set_token_revoked(&identity, false);
-    executor.advance_clock(RECONNECT_TIMEOUT);
+        .await
+        .unwrap();
     executor.run_until_parked();
-    cx_a.read(|cx| {
-        room_a.read_with(cx, |room, cx| {
+
+    // The user-visible symptom: with the bug, stale cleanup revokes the
+    // rejoining user's token and they end up in the call with no audio.
+    let room_a2 =
+        cx_a2.read(|cx| active_call_a2.read_with(cx, |call, _| call.room().unwrap().clone()));
+    cx_a2.read(|cx| {
+        room_a2.read_with(cx, |room, cx| {
             assert!(
                 room.is_connected(cx),
-                "client should have re-established its LiveKit connection"
+                "rejoining the same channel should not break audio"
             )
         })
     });
+
+    // The mechanism: that cleanup must NOT have removed the user from the
+    // LiveKit room, since they are immediately rejoining it.
+    let identity = client_a2.user_id().unwrap().to_string();
+    assert!(
+        !server
+            .test_livekit_server
+            .participant_was_removed(&identity),
+        "rejoining the same channel should not remove the LiveKit participant"
+    );
 }
 
 #[gpui::test]

crates/livekit_client/src/test.rs

@@ -57,6 +57,13 @@ pub struct TestServer {
     pub secret_key: String,
     rooms: Mutex<HashMap<String, TestServerRoom>>,
     revoked_identities: Mutex<HashSet<String>>,
+    /// Log of `(room, identity)` pairs passed to `remove_participant`, for
+    /// tests that need to assert whether a participant was forcibly removed.
+    removed_participants: Mutex<Vec<(String, String)>>,
+    /// When set, `remove_participant` revokes the identity's tokens, mirroring
+    /// LiveKit Cloud (whose docs state removal "immediately revokes their
+    /// access token"). Off by default so unrelated tests are unaffected.
+    revoke_tokens_on_removal: AtomicBool,
     executor: BackgroundExecutor,
 }
 
@@ -75,6 +82,8 @@ impl TestServer {
                 secret_key,
                 rooms: Default::default(),
                 revoked_identities: Default::default(),
+                removed_participants: Default::default(),
+                revoke_tokens_on_removal: AtomicBool::new(false),
                 executor,
             });
             e.insert(server.clone());
@@ -106,16 +115,18 @@ impl TestServer {
         }
     }
 
-    /// Simulates LiveKit Cloud revoking an identity's tokens, as happens when
-    /// `remove_participant` is called for that identity (e.g. by the collab
-    /// server's stale connection cleanup) around the time a token is issued.
-    pub fn set_token_revoked(&self, identity: &str, revoked: bool) {
-        let mut revoked_identities = self.revoked_identities.lock();
-        if revoked {
-            revoked_identities.insert(identity.to_string());
-        } else {
-            revoked_identities.remove(identity);
-        }
+    /// Makes `remove_participant` revoke the removed identity's tokens, as
+    /// LiveKit Cloud does. Used to reproduce the revoked-token failure.
+    pub fn set_revoke_tokens_on_removal(&self, revoke: bool) {
+        self.revoke_tokens_on_removal.store(revoke, SeqCst);
+    }
+
+    /// Whether `remove_participant` was ever called for the given identity.
+    pub fn participant_was_removed(&self, identity: &str) -> bool {
+        self.removed_participants
+            .lock()
+            .iter()
+            .any(|(_, removed_identity)| removed_identity == identity)
     }
 
     pub async fn create_room(&self, room: String) -> Result<()> {
@@ -266,6 +277,14 @@ impl TestServer {
     ) -> Result<()> {
         self.simulate_random_delay().await;
 
+        self.removed_participants
+            .lock()
+            .push((room_name.clone(), identity.0.clone()));
+
+        if self.revoke_tokens_on_removal.load(SeqCst) {
+            self.revoked_identities.lock().insert(identity.0.clone());
+        }
+
         let mut server_rooms = self.rooms.lock();
         let room = server_rooms
             .get_mut(&room_name)