init

Dockerfile

@@ -12,4 +12,5 @@ COPY --chown=user ./requirements.txt requirements.txt
 RUN pip install --no-cache-dir --upgrade -r requirements.txt
 
 COPY --chown=user . /app
+
 CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

main.py

@@ -1,7 +1,7 @@
 # main.py
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
 from fastapi.middleware.cors import CORSMiddleware
-from aiortc import RTCPeerConnection, RTCSessionDescription, RTCConfiguration, RTCIceServer
+from aiortc import RTCPeerConnection, RTCSessionDescription, RTCConfiguration, RTCIceServer, RTCIceCandidate
 import json
 import re
 
@@ -15,9 +15,11 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-peer_connections = {}
+# NOTE: The peer_connections dictionary is less critical with WebSockets 
+# because the WebSocket lifecycle can manage the PeerConnection.
+peer_connections = {} 
 
-# ✅ Add STUN + TURN servers here
+# ✅ STUN + TURN servers remain for the peer-to-peer connection
 ice_servers = [
     RTCIceServer(urls="stun:stun.l.google.com:19302"),
     RTCIceServer(
@@ -29,16 +31,17 @@ ice_servers = [
 
 rtc_config = RTCConfiguration(iceServers=ice_servers)
 
+# 🚀 NEW: WebSocket endpoint for Signaling
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    print("🔌 WebSocket accepted for signaling")
 
-@app.post("/signal")
-async def signal(request: Request):
-    body = await request.json()
-    offer = body["offer"]
-
-    # ✅ Peer connection with TURN support
+    # 1. Create a new Peer Connection for this session
     pc = RTCPeerConnection(rtc_config)
-    peer_connections[id(pc)] = pc  # prevent GC
+    peer_connections[id(pc)] = pc  # Keep reference to prevent GC
 
+    # 2. Add handlers for DataChannel messages (RPC logic remains)
     @pc.on("datachannel")
     def on_datachannel(channel):
         print("✅ Data channel opened")
@@ -46,42 +49,84 @@ async def signal(request: Request):
         @channel.on("message")
         def on_message(message):
             print("📩 Received:", message)
-
-            rpc = json.loads(message)
-            result = handle_rpc(rpc)
-
-            response = json.dumps({
-                "jsonrpc": "2.0",
-                "id": rpc["id"],
-                "result": result,
+            try:
+                rpc = json.loads(message)
+                result = handle_rpc(rpc)
+
+                response = json.dumps({
+                    "jsonrpc": "2.0",
+                    "id": rpc["id"],
+                    "result": result,
+                })
+                print("📤 Sending:", response)
+                channel.send(response)
+            except json.JSONDecodeError:
+                print(f"Error decoding RPC message: {message}")
+            
+
+    # 3. Handle ICE Candidates and send them over the WebSocket
+    @pc.on("icecandidate")
+    async def on_icecandidate(candidate):
+        if candidate:
+            print(f"📤 Sending ICE candidate: {candidate.sdpMid}")
+            await websocket.send_json({
+                "type": "candidate",
+                # aiortc candidates need to be converted to a dictionary for safe JSON transfer
+                "candidate": {
+                    "candidate": candidate.candidate,
+                    "sdpMid": candidate.sdpMid,
+                    "sdpMLineIndex": candidate.sdpMLineIndex,
+                    "usernameFragment": candidate.usernameFragment,
+                }
             })
 
-            print("📤 Sending:", response)
-            channel.send(response)
-
-    # Apply remote SDP
-    await pc.setRemoteDescription(
-        RTCSessionDescription(
-            sdp=offer["sdp"],
-            type=offer["type"]
-        )
-    )
-
-    # Create & return answer
-    answer = await pc.createAnswer()
-    await pc.setLocalDescription(answer)
-
-    print("✅ Answer generated and sent")
-
-    return {
-        "answer": {
-            "sdp": pc.localDescription.sdp,
-            "type": pc.localDescription.type
-        }
-    }
-
-
-# ✅ Your RPC logic stays the same
+    try:
+        # 4. Main loop to receive signaling messages (Offer, Answer, Candidates)
+        while True:
+            data = await websocket.receive_json()
+            message_type = data.get("type")
+
+            if message_type == "offer":
+                # Handle initial Offer
+                offer = RTCSessionDescription(sdp=data["sdp"], type="offer")
+                await pc.setRemoteDescription(offer)
+                
+                # Create and send Answer
+                answer = await pc.createAnswer()
+                await pc.setLocalDescription(answer)
+                
+                print("✅ Sending Answer")
+                await websocket.send_json({
+                    "type": "answer",
+                    "sdp": pc.localDescription.sdp
+                })
+
+            elif message_type == "candidate":
+                # Handle incoming ICE candidates
+                try:
+                    candidate_init = data["candidate"]
+                    candidate = RTCIceCandidate(
+                        candidate=candidate_init["candidate"],
+                        sdpMid=candidate_init["sdpMid"],
+                        sdpMLineIndex=candidate_init["sdpMLineIndex"],
+                        usernameFragment=candidate_init.get("usernameFragment"),
+                    )
+                    await pc.addIceCandidate(candidate)
+                    print(f"📥 Added remote ICE candidate: {candidate.sdpMid}")
+                except Exception as e:
+                    print(f"Error adding ICE candidate: {e}")
+
+
+    except WebSocketDisconnect:
+        print("🚫 WebSocket disconnected. Closing PeerConnection.")
+    finally:
+        # Clean up
+        await pc.close()
+        if id(pc) in peer_connections:
+            del peer_connections[id(pc)]
+
+
+# ✅ Your RPC logic remains the same
 def handle_rpc(rpc):
     method = rpc.get("method")
     params = rpc.get("params", {})
@@ -122,4 +167,4 @@ def handle_rpc(rpc):
         return "Unknown method"
 
 
-print("✅ FastAPI WebRTC backend with TURN is ready")
+print("✅ FastAPI WebRTC backend with WebSocket signaling is ready")