fix(media): preserve audio across route reconnects (#245)

Keep persistent relay streams and LiveKit SFU publications bound to their current audio senders across reconnects.

src/client/sfu/publisher.ts

@@ -712,18 +712,77 @@ export class SfuPublisher {
       if (!this.owns(room, generation) || this.state !== "active") {
         return false;
       }
-      const audio = this.audio;
-      const profile = this.desiredProfile;
       const sdk = this.sdk;
+      if (!sdk) {
+        return false;
+      }
+
+      const previousVideo = this.video;
+      const previousAudio = this.audio;
+      const previousVideoSender =
+        previousVideo?.publication.videoTrack?.sender ?? null;
+      const video = previousVideo
+        ? currentPublishedTrack(
+            room,
+            sdk.Track.Source.ScreenShare,
+            previousVideo.rawTrack,
+          )
+        : null;
+      const videoSender = video?.publication.videoTrack?.sender ?? null;
+      if (!video || !videoSender) {
+        this.failureStage = "transport";
+        await this.failClosed(room, generation);
+        return false;
+      }
+      const audio = previousAudio
+        ? currentPublishedTrack(
+            room,
+            sdk.Track.Source.ScreenShareAudio,
+            previousAudio.rawTrack,
+          )
+        : null;
+      const audioSender = audio?.publication.audioTrack?.sender ?? null;
+      if (previousAudio && (!audio || !audioSender)) {
+        this.failureStage = "transport";
+        await this.failClosed(room, generation);
+        return false;
+      }
+      this.stopStats();
+      this.video = video;
+      this.audio = audio;
+      if (
+        video.publication !== previousVideo?.publication ||
+        videoSender !== previousVideoSender
+      ) {
+        this.senderParameters = null;
+        this.videoQualityWarning = null;
+      }
+      this.startStats(video);
+
+      const profile = this.desiredProfile;
       const profileRevision = this.profileRevision;
-      if (!audio || !profile || !sdk) {
+      if (!audio || !profile) {
+        this.audioSenderParameters = null;
+        this.audioParametersSender = null;
+        this.audioQualityWarning = null;
         return true;
       }
       retainPublishedAudioOptions(audio, profile, sdk);
 
-      let audioSender: RTCRtpSender | null = null;
+      if (!audioSender) {
+        this.failureStage = "transport";
+        await this.failClosed(room, generation);
+        return false;
+      }
+      if (
+        audio.publication !== previousAudio?.publication ||
+        this.audioParametersSender !== audioSender
+      ) {
+        this.audioSenderParameters = null;
+        this.audioParametersSender = null;
+        this.audioQualityWarning = null;
+      }
       try {
-        audioSender = publishedAudioSender(audio);
         const ownedAudioSender = audioSender;
         const configured = await configurePublishedAudio(
           audio,
@@ -1107,6 +1166,23 @@ function publishedAudioSender(published: PublishedTrack): RTCRtpSender {
   return sender;
 }
 
+function currentPublishedTrack(
+  room: Room,
+  source: Track.Source,
+  expectedRawTrack: MediaStreamTrack,
+): PublishedTrack | null {
+  const publication = room.localParticipant.getTrackPublication(source);
+  const rawTrack = publication?.track?.mediaStreamTrack;
+  if (
+    !publication ||
+    !rawTrack ||
+    rawTrack !== expectedRawTrack
+  ) {
+    return null;
+  }
+  return { publication, rawTrack };
+}
+
 function retainPublishedAudioOptions(
   published: PublishedTrack,
   profile: QualityProfile,

src/client/webrtc/viewer-relay.ts

@@ -172,10 +172,10 @@ export class ViewerRelay {
     if (this.disposed) {
       return;
     }
-    const replacePrepared = this.stream !== null && this.stream !== stream;
+    const preparedPeer = this.preparedChild?.peer ?? null;
     this.stream = stream;
     this.syncQueue = this.syncQueue
-      .then(() => this.syncStream(stream, replacePrepared))
+      .then(() => this.syncStream(stream, preparedPeer))
       .catch(() => undefined);
   }
 
@@ -318,14 +318,17 @@ export class ViewerRelay {
 
   private async syncStream(
     stream: MediaStream,
-    replacePrepared: boolean,
+    preparedPeer: HostPeer | null,
   ): Promise<void> {
     if (this.disposed || this.stream !== stream) {
       return;
     }
-    const prepared = replacePrepared ? this.preparedChild : null;
-    if (prepared && !(await this.replacePeerStream(prepared.peer, stream))) {
-      this.failPreparedChild(prepared.peer);
+    if (
+      preparedPeer &&
+      this.preparedChild?.peer === preparedPeer &&
+      !(await this.replacePeerStream(preparedPeer, stream))
+    ) {
+      this.failPreparedChild(preparedPeer);
     }
     if (this.disposed || this.stream !== stream) {
       return;

tests/host-peer-tracks.test.ts

@@ -1581,6 +1581,46 @@ describe("ViewerRelay downstream ownership", () => {
     relay.dispose();
   });
 
+  it("syncs same-stream audio changes into the prepared connection", async () => {
+    const relay = new ViewerRelay(
+      { iceServers: [] },
+      QUALITY_PROFILES["720p30"],
+      { sendSignal: () => true },
+    );
+    const video = createTrack("video", "persistent-video");
+    let audio: MediaStreamTrack | null = null;
+    const persistentStream = {
+      getTracks: () => (audio ? [video, audio] : [video]),
+      getVideoTracks: () => [video],
+      getAudioTracks: () => (audio ? [audio] : []),
+    } as unknown as MediaStream;
+
+    relay.setStream(persistentStream);
+    expect(
+      relay.prepareChild(
+        7,
+        routeCandidate(7, "prepared-child"),
+        ["prepared-child"],
+      ),
+    ).toBe(true);
+    const preparedConnection = FakePeerConnection.latest!;
+    expect(preparedConnection.senders[1]?.track).toBeNull();
+
+    audio = createTrack("audio", "late-audio");
+    relay.setStream(persistentStream);
+    await vi.waitFor(() =>
+      expect(preparedConnection.senders[1]?.track).toBe(audio),
+    );
+
+    audio = null;
+    relay.setStream(persistentStream);
+    await vi.waitFor(() =>
+      expect(preparedConnection.senders[1]?.track).toBeNull(),
+    );
+    expect(FakePeerConnection.latest).toBe(preparedConnection);
+    relay.dispose();
+  });
+
   it("admits provisional children until endpoint cap three is full", async () => {
     const relay = new ViewerRelay(
       { iceServers: [] },

tests/sfu-client.test.ts

@@ -92,6 +92,10 @@ const livekit = vi.hoisted(() => {
       this.sender = new FakeSender(track);
     }
 
+    get mediaStreamTrack(): MediaStreamTrack {
+      return this.currentTrack;
+    }
+
     readonly replaceTrack = vi.fn(
       async (nextTrack: MediaStreamTrack): Promise<void> => {
         this.currentTrack = nextTrack;
@@ -169,6 +173,24 @@ const livekit = vi.hoisted(() => {
       },
     );
 
+    getTrackPublication(source: string) {
+      return this.publications.find(
+        (publication) => publication.options.source === source,
+      );
+    }
+
+    async republishForReconnect(): Promise<void> {
+      const previous = [...this.publications];
+      this.publications.length = 0;
+      for (const publication of previous) {
+        delete publication.videoTrack;
+        delete publication.audioTrack;
+        await this.publishTrack(publication.rawTrack, {
+          ...publication.options,
+        });
+      }
+    }
+
     republishedOptions: Array<Record<string, unknown>> = [];
     readonly republishAllTracks = vi.fn(async (): Promise<void> => {
       this.republishedOptions = this.publications.map((publication) => ({
@@ -1016,6 +1038,123 @@ describe("SfuPublisher", () => {
     expect(room.localParticipant.republishAllTracks).not.toHaveBeenCalled();
   });
 
+  it("rebinds republished tracks after a full reconnect", async () => {
+    const publisher = new SfuPublisher();
+    await publisher.connect(connection);
+    await publisher.activate(
+      stream(track("video", "video-1"), track("audio", "audio-1")),
+      { ...qualityProfile, screenAudioQuality: "music" },
+    );
+    await expect(
+      publisher.updateProfile({
+        ...qualityProfile,
+        screenAudioQuality: "very-high",
+      }),
+    ).resolves.toBe(true);
+    const room = livekit.state.rooms[0];
+    const oldVideoPublication = room.localParticipant.publications[0];
+    const oldAudioPublication = room.localParticipant.publications[1];
+
+    await room.localParticipant.republishForReconnect();
+    const videoPublication = room.localParticipant.getTrackPublication(
+      Track.Source.ScreenShare,
+    )!;
+    const audioPublication = room.localParticipant.getTrackPublication(
+      Track.Source.ScreenShareAudio,
+    )!;
+    expect(videoPublication).not.toBe(oldVideoPublication);
+    expect(audioPublication).not.toBe(oldAudioPublication);
+
+    room.emit(RoomEvent.Reconnected);
+
+    await vi.waitFor(() =>
+      expect(audioPublication.track.sender.setParameters).toHaveBeenCalledOnce(),
+    );
+    expect(audioPublication.track.sender.parameters.encodings[0]?.maxBitrate).toBe(
+      256_000,
+    );
+    expect(publisher.getAudioSenderParameters()?.appliedMaxBitrate).toBe(
+      256_000,
+    );
+    expect(publisher.getSenderParameters()).toBeNull();
+    expect(publisher.getQualityWarning()).toBeNull();
+  });
+
+  it("retains video sender readback across a signal-only reconnect", async () => {
+    const publisher = new SfuPublisher();
+    await publisher.connect(connection);
+    await publisher.activate(
+      stream(track("video", "video-1"), track("audio", "audio-1")),
+      { ...qualityProfile, screenAudioQuality: "music" },
+    );
+    const room = livekit.state.rooms[0];
+    const videoSender = room.localParticipant.publications[0].track.sender;
+    const audioSender = room.localParticipant.publications[1].track.sender;
+    videoSender.setParameters.mockImplementationOnce(async (parameters) => {
+      videoSender.parameters = {
+        ...parameters,
+        encodings: parameters.encodings.map((encoding) => ({
+          ...encoding,
+          maxBitrate: 2_000_000,
+        })),
+      };
+    });
+    await expect(
+      publisher.updateProfile({
+        ...qualityProfile,
+        maxBitrate: 5_000_000,
+        degradationPreference: "balanced",
+        screenAudioQuality: "music",
+      }),
+    ).resolves.toBe(true);
+    const senderParameters = publisher.getSenderParameters();
+    const qualityWarning = publisher.getQualityWarning();
+    expect(senderParameters).not.toBeNull();
+    expect(qualityWarning).not.toBeNull();
+    const audioUpdates = audioSender.setParameters.mock.calls.length;
+
+    room.emit(RoomEvent.Reconnected);
+
+    await vi.waitFor(() =>
+      expect(audioSender.setParameters).toHaveBeenCalledTimes(audioUpdates + 1),
+    );
+    expect(publisher.getSenderParameters()).toBe(senderParameters);
+    expect(publisher.getQualityWarning()).toBe(qualityWarning);
+    expect(room.disconnect).not.toHaveBeenCalled();
+  });
+
+  it.each(["publication", "sender"] as const)(
+    "fails the SFU route when republished audio loses its %s",
+    async (missing) => {
+      const disconnected = vi.fn();
+      const publisher = new SfuPublisher({ onDisconnected: disconnected });
+      await publisher.connect(connection);
+      await publisher.activate(
+        stream(track("video", "video-1"), track("audio", "audio-1")),
+        { ...qualityProfile, screenAudioQuality: "music" },
+      );
+      const room = livekit.state.rooms[0];
+      await room.localParticipant.republishForReconnect();
+      const audioPublication = room.localParticipant.getTrackPublication(
+        Track.Source.ScreenShareAudio,
+      )!;
+      if (missing === "publication") {
+        room.localParticipant.publications.splice(
+          room.localParticipant.publications.indexOf(audioPublication),
+          1,
+        );
+      } else {
+        delete audioPublication.audioTrack;
+      }
+
+      room.emit(RoomEvent.Reconnected);
+
+      await vi.waitFor(() => expect(room.disconnect).toHaveBeenCalledWith(false));
+      expect(publisher.getFailureStage()).toBe("transport");
+      expect(disconnected).toHaveBeenCalledOnce();
+    },
+  );
+
   it("retains the active profile for LiveKit track restart and republish", async () => {
     const publisher = new SfuPublisher();
     await publisher.connect(connection);