fix useWebRTC hook to handle events correctly

src/hooks/useWebRTC.ts

@@ -1,133 +1,140 @@
-import { useState } from "react";
+import { useState, useRef, useCallback, useEffect } from "react";
 
 import { ICreateConversationResponse } from "../interfaces/conversation";
-import { createConversation, createConnection } from "../services/api/chatService";
+import { createConnection, createConversation } from "../services/api/chatService";
 import { conversationWebSocket } from "../services/websockets/conversation";
 
-export function useWebRTC() {
-  const [peerConnection, setPeerConnection] = useState<RTCPeerConnection | null>(null);
-  const [socket, setSocket] = useState<WebSocket | null>(null);
-  const [transcript, setTranscript] = useState<string>("");
+export const useWebRTC = () => {
   const [isConnected, setIsConnected] = useState(false);
+  const [transcript, setTranscript] = useState("");
   const [remoteStream, setRemoteStream] = useState<MediaStream | null>(null);
+  const [callStatus, setCallStatus] = useState("");
+  const peerConnectionRef = useRef<RTCPeerConnection | null>(null);
+  const dataChannelRef = useRef<RTCDataChannel | null>(null);
+  const socketRef = useRef<WebSocket | null>(null);
   const [error, setError] = useState<string | null>(null);
+  const mediaStreamRef = useRef<MediaStream | null>(null);
+  const isConnectedRef = useRef(isConnected);
 
-  async function startCall() {
-    try {
-      setError(null);
-      const sessionResponse = await createConversation({ modality: "voice" });
-      const sessionData: ICreateConversationResponse = await sessionResponse.data;
-      const conversationId = sessionData.conversation_id;
+  useEffect(() => {
+    isConnectedRef.current = isConnected;
+  }, [isConnected]);
 
-      const pc = new RTCPeerConnection();
-      setPeerConnection(pc);
+  const endCall = useCallback(() => {
+    if (peerConnectionRef.current) {
+      peerConnectionRef.current.close();
+      peerConnectionRef.current = null;
+    }
+    if (socketRef.current) {
+      socketRef.current.close();
+      socketRef.current = null;
+    }
+    if (mediaStreamRef.current) {
+      mediaStreamRef.current.getTracks().forEach((track) => track.stop());
+      mediaStreamRef.current = null;
+    }
+    if (dataChannelRef.current) {
+      dataChannelRef.current.close();
+      dataChannelRef.current = null;
+    }
+    setIsConnected(false);
+    setCallStatus("Disconnected");
+    setRemoteStream(null);
+    setTranscript("");
+  }, []);
 
-      pc.ontrack = (event) => {
-        const newStream = new MediaStream();
-        newStream.addTrack(event.track);
-        setRemoteStream(newStream);
+  const startCall = useCallback(async () => {
+    try {
+      const peerConnection = new RTCPeerConnection();
+      peerConnectionRef.current = peerConnection;
+      peerConnection.onconnectionstatechange = () => {
+        switch (peerConnection.connectionState) {
+          case "connected":
+            setIsConnected(true);
+            setCallStatus("Connected");
+            break;
+          case "disconnected":
+          case "failed":
+          case "closed":
+            endCall();
+            break;
+          default:
+            break;
+        }
+      };
+      peerConnection.ontrack = (event) => {
+        setRemoteStream(event.streams[0]);
       };
 
-      pc.onconnectionstatechange = () => {
-        if (pc.connectionState === "connected") {
-          setIsConnected(true);
+      const mediaStream = await navigator.mediaDevices.getUserMedia({ audio: true });
+      mediaStreamRef.current = mediaStream;
+      peerConnection.addTrack(mediaStream.getTracks()[0], mediaStream);
+      const dataChannel = peerConnection.createDataChannel("response");
+      dataChannelRef.current = dataChannel;
+      dataChannel.onmessage = (event) => {
+        if (socketRef.current?.readyState === WebSocket.OPEN) {
+          socketRef.current.send(event.data);
         }
       };
 
-      const stream = await navigator.mediaDevices.getUserMedia({ audio: true, video: false });
-      stream.getTracks().forEach((track) => {
-        pc.addTrack(track, stream);
-      });
-
-      const offer = await pc.createOffer();
-      await pc.setLocalDescription(offer);
+      const sessionResponse = await createConversation({ modality: "voice" });
+      const sessionData: ICreateConversationResponse = await sessionResponse.data;
+      const conversationId = sessionData.conversation_id;
 
+      const offer = await peerConnection.createOffer();
+      await peerConnection.setLocalDescription(offer);
       const webrtcResponse = await createConnection(conversationId, {
         conversation_id: conversationId,
         offer: { sdp: offer.sdp, type: offer.type },
       });
 
-      const response = await webrtcResponse.data;
-      const ephemeralKey = response.ephemeral_key;
-
-      const ws = conversationWebSocket({
+      const responseData = await webrtcResponse.data;
+      const ephemeralKey = responseData.ephemeral_key;
+      const socket = conversationWebSocket({
         conversationId,
         modality: "voice",
       });
 
-      setSocket(ws);
-
-      ws.onopen = () => {
-        const auth_token = localStorage.getItem("auth_token");
-        const token = `Bearer ${auth_token}`;
-        ws.send(
+      socketRef.current = socket;
+      socket.onopen = () => {
+        socket.send(
           JSON.stringify({
             type: "headers",
-            headers: {
-              Authorization: `Bearer ${ephemeralKey}`,
-              Token: token,
-            },
+            headers: { Authorization: `Bearer ${ephemeralKey}` },
           })
         );
       };
 
-      ws.onmessage = (event) => {
+      socket.onmessage = (event) => {
         try {
           const data = JSON.parse(event.data);
           if (data.transcript) {
-            setTranscript(data.transcript);
+            setTranscript((prev) => `${prev}\n${data.transcript}`);
+          } else if (dataChannelRef.current?.readyState === "open") {
+            dataChannelRef.current.send(JSON.stringify(data));
           }
         } catch {
-          setError("Error processing message from server");
-        }
-      };
-
-      pc.onicecandidate = (event) => {
-        if (event.candidate && ws.readyState === WebSocket.OPEN) {
-          ws.send(
-            JSON.stringify({
-              type: "ice-candidate",
-              candidate: event.candidate,
-            })
-          );
+          setError("Error handling message");
         }
       };
 
-      ws.onclose = (_event) => {
-        if (isConnected) {
-          endCall();
-        }
+      socket.onclose = (_event) => {
+        if (isConnectedRef.current) endCall();
       };
-
-      await pc.setRemoteDescription(new RTCSessionDescription(response.answer));
-
-      const dataChannel = pc.createDataChannel("chat");
-      dataChannel.onmessage = (event) => {
-        if (ws && ws.readyState === WebSocket.OPEN) {
-          ws.send(event.data);
-        }
-      };
-    } catch (error) {
-      setError(error instanceof Error ? error.message : "Failed to start call");
+      await peerConnection.setRemoteDescription(new RTCSessionDescription(responseData.answer));
+    } catch {
+      setError("Call setup failed");
       endCall();
-      throw error;
-    }
-  }
-
-  function endCall() {
-    if (peerConnection) {
-      peerConnection.close();
-      setPeerConnection(null);
-    }
-
-    if (socket) {
-      socket.close();
-      setSocket(null);
     }
-
-    setIsConnected(false);
-    setRemoteStream(null);
-  }
-
-  return { startCall, endCall, transcript, isConnected, remoteStream, error };
-}
+  }, [endCall]);
+
+  return {
+    startCall,
+    endCall,
+    isConnected,
+    transcript,
+    remoteStream,
+    callStatus,
+    error,
+  };
+};