Buffer early WebRTC ICE candidates

streaming_app.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import copy
 import json
 import multiprocessing
@@ -8,6 +9,7 @@ import tarfile
 import threading
 import time
 import uuid
+from collections import defaultdict
 from pathlib import Path
 
 import gradio as gr
@@ -16,6 +18,7 @@ import sentencepiece
 import spaces
 import torch
 from fastrtc import StreamHandler, WebRTC, get_cloudflare_turn_credentials_async
+from gradio.components.base import server
 from huggingface_hub import hf_hub_download
 
 
@@ -224,6 +227,50 @@ async def get_turn_configuration():
     return configuration
 
 
+class ReliableWebRTC(WebRTC):
+    """Serialize FastRTC signaling and replay ICE candidates that beat the SDP offer."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._signaling_lock = None
+        self._pending_ice = defaultdict(list)
+
+    @server
+    async def offer(self, body):
+        if self._signaling_lock is None:
+            self._signaling_lock = asyncio.Lock()
+
+        async with self._signaling_lock:
+            webrtc_id = body.get("webrtc_id")
+            is_candidate = body.get("type") == "ice-candidate" and "candidate" in body
+
+            if is_candidate and webrtc_id not in self.pcs:
+                pending = self._pending_ice[webrtc_id]
+                if len(pending) < 32:
+                    pending.append(body)
+                print(f"Buffered early ICE candidate for pending connection: {webrtc_id}")
+                return {"status": "success"}
+
+            response = await self.handle_offer(
+                body,
+                self.set_additional_outputs(webrtc_id),
+            )
+
+            if not is_candidate and webrtc_id in self.pcs:
+                pending = self._pending_ice.pop(webrtc_id, [])
+                for candidate in pending:
+                    await self.handle_offer(
+                        candidate,
+                        self.set_additional_outputs(webrtc_id),
+                    )
+                print(
+                    f"Registered WebRTC offer and replayed {len(pending)} early ICE candidates: "
+                    f"{webrtc_id}"
+                )
+
+            return response
+
+
 def new_bridge(oauth_token: gr.OAuthToken | None):
     remember_turn_token(oauth_token)
     input_queue = _manager.Queue(INPUT_QUEUE_FRAMES)
@@ -284,6 +331,8 @@ class PersonaPlexQueueHandler(StreamHandler):
             input_sample_rate=OUTPUT_SAMPLE_RATE,
             fps=WEBRTC_FRAME_RATE,
         )
+        self._logged_input = False
+        self._logged_output = False
 
     def _bridge(self):
         # FastRTC prepends the WebRTC component value to additional inputs.
@@ -301,6 +350,9 @@ class PersonaPlexQueueHandler(StreamHandler):
             return
         sample_rate, array = frame
         audio = np.asarray(array, dtype=np.int16).reshape(-1)
+        if not self._logged_input:
+            print("WebRTC audio bridge received its first input frame.")
+            self._logged_input = True
         if bounded_put(input_queue, (int(sample_rate), audio)):
             counters["input_drops"] = int(counters.get("input_drops", 0)) + 1
 
@@ -315,6 +367,9 @@ class PersonaPlexQueueHandler(StreamHandler):
             audio = output_queue.get_nowait()
         except queue.Empty:
             return None
+        if not self._logged_output:
+            print("WebRTC audio bridge emitted its first output frame.")
+            self._logged_output = True
         return OUTPUT_SAMPLE_RATE, np.asarray(audio, dtype=np.int16).reshape(1, -1)
 
     def copy(self):
@@ -407,6 +462,7 @@ def run_streaming_session(
         drain_proxy_queue(input_queue)
         ready_seconds = round(time.perf_counter() - session_started, 3)
         ready_event.set()
+        print("PersonaPlex GPU session is ready and waiting for WebRTC audio.")
         live_started = time.perf_counter()
         status = (
             f"**Engine ready in {ready_seconds:.1f}s.** Start the WebRTC microphone and speak naturally. "
@@ -561,7 +617,7 @@ with gr.Blocks(title="RoleForge Streaming Spike", theme=gr.themes.Soft(), css=CS
         with gr.Column(scale=2):
             gr.Markdown("### 🎙️ Live Stage")
             status = gr.Markdown("Sign in, choose the cue, then start the live engine.")
-            webrtc = WebRTC(
+            webrtc = ReliableWebRTC(
                 label="2. Connect microphone after Engine ready",
                 modality="audio",
                 mode="send-receive",