This fix ensures that WebSocket message listeners are properly cleaned up, preventing memory leaks when users leave sessions or navigate away from the component.

jam-ui/src/components/client/JKSessionScreen.js

@@ -1,5 +1,5 @@
 // jam-ui/src/components/client/JKSessionScreen.js
-import React, { useEffect, useContext, useState, memo, useMemo } from 'react'
+import React, { useEffect, useContext, useState, memo, useMemo, useCallback } from 'react'
 import { useParams, useHistory } from 'react-router-dom';
 
 //import useJamServer, { ConnectionStatus } from '../../hooks/useJamServer'
@@ -76,7 +76,8 @@ const JKSessionScreen = () => {
     lastError,
     jamClient,
     server,
-    registerMessageCallback } = useJamServerContext();
+    registerMessageCallback,
+    unregisterMessageCallback } = useJamServerContext();
   const { currentSession, setCurrentSession, currentSessionIdRef, setCurrentSessionId, inSession } = useCurrentSessionContext();
   const { globalObject, metronomeState, closeMetronome, resetMetronome } = useGlobalContext();
   const { getCurrentRecordingState, reset: resetRecordingState, currentlyRecording } = useRecordingHelpers();
@@ -139,6 +140,17 @@ const JKSessionScreen = () => {
   const [selectedJamTrack, setSelectedJamTrack] = useState(null);
   const [jamTrackStems, setJamTrackStems] = useState([]);
 
+  // Store references to registered message callbacks for cleanup
+  const [registeredCallbacks, setRegisteredCallbacks] = useState([]);
+
+  // Function to unregister message callbacks
+  const unregisterMessageCallbacks = useCallback(() => {
+    registeredCallbacks.forEach(({ type, callback }) => {
+      unregisterMessageCallback(type, callback);
+    });
+    setRegisteredCallbacks([]);
+  }, [registeredCallbacks, unregisterMessageCallback]);
+
   useEffect(() => {
     if (!isConnected || !jamClient) return;
     console.debug("JKSessionScreen: -DEBUG- isConnected changed to true");
@@ -326,10 +338,20 @@ const JKSessionScreen = () => {
 
         //@refreshCurrentSession(true);
 
-        registerMessageCallback(MessageType.SESSION_JOIN, trackChanges);
-        registerMessageCallback(MessageType.SESSION_DEPART, trackChanges);
-        registerMessageCallback(MessageType.TRACKS_CHANGED, trackChanges);
-        registerMessageCallback(MessageType.HEARTBEAT_ACK, trackChanges);
+        // Register message callbacks and store references for cleanup
+        const callbacksToRegister = [
+          { type: MessageType.SESSION_JOIN, callback: trackChanges },
+          { type: MessageType.SESSION_DEPART, callback: trackChanges },
+          { type: MessageType.TRACKS_CHANGED, callback: trackChanges },
+          { type: MessageType.HEARTBEAT_ACK, callback: trackChanges }
+        ];
+
+        callbacksToRegister.forEach(({ type, callback }) => {
+          registerMessageCallback(type, callback);
+        });
+
+        // Store registered callbacks for cleanup
+        setRegisteredCallbacks(callbacksToRegister);
 
         //TODO: revist the logic in following commented section
         //if (document) { $(document).trigger(EVENTS.SESSION_STARTED, { session: { id: this.currentSessionId, lesson_session: response.lesson_session } }); }
@@ -604,6 +626,9 @@ const JKSessionScreen = () => {
     try {
       setLeaveLoading(true);
 
+      // Unregister message callbacks before leaving
+      unregisterMessageCallbacks();
+
       // Close metronome if open before leaving
       if (metronomeState.isOpen) {
         console.log('Closing metronome before leaving session');
@@ -636,16 +661,19 @@ const JKSessionScreen = () => {
     }
   };
 
-  // Cleanup metronome state when leaving session
+  // Cleanup metronome state and message callbacks when component unmounts
   useEffect(() => {
     return () => {
+      // Unregister message callbacks
+      unregisterMessageCallbacks();
+
       // Reset metronome state when component unmounts (session ends)
       if (metronomeState.isOpen) {
         console.log('Resetting metronome state on session cleanup');
         resetMetronome();
       }
     };
-  }, [metronomeState.isOpen, resetMetronome]);
+  }, [metronomeState.isOpen, resetMetronome, unregisterMessageCallbacks]);
 
   // Check if user can use video (subscription/permission check)
   const canVideo = () => {

jam-ui/src/context/JamServerContext.js

@@ -10,7 +10,8 @@ export const JamServerProvider = ({ children }) => {
     lastError,
     jamClient,
     server,
-    registerMessageCallback } = useJamServer(process.env.REACT_APP_WEBSOCKET_GATEWAY_URL);
+    registerMessageCallback,
+    unregisterMessageCallback } = useJamServer(process.env.REACT_APP_WEBSOCKET_GATEWAY_URL);
 
   return (
     <JamServerContext.Provider value={{
@@ -22,6 +23,7 @@ export const JamServerProvider = ({ children }) => {
       reconnectAttempts,
       lastError,
       registerMessageCallback,
+      unregisterMessageCallback,
     }}
     >
       {children}