fix: address critical security and concurrency issues

- Replace global session state with per-session registry in window_message.py
- Add session_id parameter to audio hooks for proper stream identification
- Add missing /api/window-message/receive endpoint for frontend integration
- Fix postMessage security by avoiding wildcard origins where possible
- Update all tests to reflect new hook signatures

Addresses critical findings from gemini-code-assist review.

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-authored-by: Mervin Praison <[email protected]>

src/frontend/src/hooks/useWindowMessage.ts

@@ -98,7 +98,11 @@ export function useWindowMessage(options: UseWindowMessageOptions = {}): UseWind
           
           if (window.parent && target) {
             if (target === "parent") {
-              window.parent.postMessage(messageData, "*");
+              // Use specific origin when possible, avoid wildcard
+              const origin = targetOrigin !== "*" && targetOrigin !== "parent" 
+                ? targetOrigin 
+                : window.location.origin;
+              window.parent.postMessage(messageData, origin);
             } else {
               window.parent.postMessage(messageData, target);
             }
@@ -169,11 +173,11 @@ export function useWindowMessage(options: UseWindowMessageOptions = {}): UseWind
     }
     
     try {
-      if (targetOrigin === "parent") {
-        window.parent.postMessage(data, "*");
-      } else {
-        window.parent.postMessage(data, targetOrigin);
-      }
+      // Use specific origin when possible, avoid wildcard
+      const origin = targetOrigin !== "*" && targetOrigin !== "parent" 
+        ? targetOrigin 
+        : window.location.origin;
+      window.parent.postMessage(data, origin);
     } catch (error) {
       console.error('Failed to send message to parent:', error);
     }

src/praisonaiui/features/audio.py

@@ -160,7 +160,7 @@ class AudioFeature(BaseFeatureProtocol):
             _audio_stats["active_sessions"] += 1
             
             # Trigger start hooks
-            await self._trigger_start_hooks()
+            await self._trigger_start_hooks(session_id)
             
             await websocket.send_json({"status": "ready"})
             
@@ -182,7 +182,7 @@ class AudioFeature(BaseFeatureProtocol):
                             try:
                                 control_msg = json.loads(message["text"])
                                 if control_msg.get("type") == "end":
-                                    await self._trigger_end_hooks()
+                                    await self._trigger_end_hooks(session_id)
                                     break
                             except json.JSONDecodeError:
                                 _log.warning("Invalid JSON in audio control message")
@@ -202,7 +202,8 @@ class AudioFeature(BaseFeatureProtocol):
                 _audio_stats["active_sessions"] = max(0, _audio_stats["active_sessions"] - 1)
             
             # Trigger end hooks if not already triggered
-            await self._trigger_end_hooks()
+            if session_id:
+                await self._trigger_end_hooks(session_id)
             
             if not websocket.client_state.disconnected:
                 await websocket.close()
@@ -219,35 +220,35 @@ class AudioFeature(BaseFeatureProtocol):
             _audio_stats["total_bytes"] += len(pcm_data)
         
         # Trigger chunk hooks
-        await self._trigger_chunk_hooks(pcm_data, sample_rate)
+        await self._trigger_chunk_hooks(session_id, pcm_data, sample_rate)
 
-    async def _trigger_start_hooks(self) -> None:
+    async def _trigger_start_hooks(self, session_id: str) -> None:
         """Execute all audio start hooks."""
         for hook in _audio_start_hooks:
             try:
                 _log.debug(f"Executing audio start hook: {getattr(hook, '__name__', str(hook))}")
-                result = hook()
+                result = hook(session_id)
                 if asyncio.iscoroutine(result):
                     await result
             except Exception as e:
                 _log.error(f"Audio start hook failed: {getattr(hook, '__name__', str(hook))}: {e}")
 
-    async def _trigger_chunk_hooks(self, pcm_data: bytes, sample_rate: int) -> None:
+    async def _trigger_chunk_hooks(self, session_id: str, pcm_data: bytes, sample_rate: int) -> None:
         """Execute all audio chunk hooks."""
         for hook in _audio_chunk_hooks:
             try:
-                result = hook(pcm_data, sample_rate)
+                result = hook(session_id, pcm_data, sample_rate)
                 if asyncio.iscoroutine(result):
                     await result
             except Exception as e:
                 _log.error(f"Audio chunk hook failed: {getattr(hook, '__name__', str(hook))}: {e}")
 
-    async def _trigger_end_hooks(self) -> None:
+    async def _trigger_end_hooks(self, session_id: str) -> None:
         """Execute all audio end hooks."""
         for hook in _audio_end_hooks:
             try:
                 _log.debug(f"Executing audio end hook: {getattr(hook, '__name__', str(hook))}")
-                result = hook()
+                result = hook(session_id)
                 if asyncio.iscoroutine(result):
                     await result
             except Exception as e:
@@ -300,7 +301,7 @@ def register_audio_start_hook(func: Callable) -> Callable:
     """Register an audio start hook.
     
     Args:
-        func: Function to call when audio recording starts
+        func: Function to call when audio recording starts (receives session_id)
         
     Returns:
         The original function (for use as decorator)
@@ -315,7 +316,7 @@ def register_audio_chunk_hook(func: Callable) -> Callable:
     """Register an audio chunk hook.
     
     Args:
-        func: Function to call for each audio chunk (pcm_data, sample_rate)
+        func: Function to call for each audio chunk (session_id, pcm_data, sample_rate)
         
     Returns:
         The original function (for use as decorator)
@@ -330,7 +331,7 @@ def register_audio_end_hook(func: Callable) -> Callable:
     """Register an audio end hook.
     
     Args:
-        func: Function to call when audio recording ends
+        func: Function to call when audio recording ends (receives session_id)
         
     Returns:
         The original function (for use as decorator)
@@ -367,7 +368,7 @@ def on_audio_start(func: Callable) -> Callable:
     Example::
     
         @aiui.on_audio_start
-        async def start():
+        async def start(session_id: str):
             await aiui.Message(content="🎙 Listening...").send()
     """
     return register_audio_start_hook(func)
@@ -381,7 +382,7 @@ def on_audio_chunk(func: Callable) -> Callable:
     Example::
     
         @aiui.on_audio_chunk
-        async def chunk(pcm: bytes, sample_rate: int):
+        async def chunk(session_id: str, pcm: bytes, sample_rate: int):
             # Forward to Whisper server / Deepgram / local vosk
             await stt_buffer.append(pcm)
     """
@@ -396,7 +397,7 @@ def on_audio_end(func: Callable) -> Callable:
     Example::
     
         @aiui.on_audio_end
-        async def end():
+        async def end(session_id: str):
             transcript = await stt_buffer.finalise()
             await handler(aiui.InboundMessage(content=transcript, source="voice"))
     """

src/praisonaiui/features/window_message.py

@@ -24,8 +24,8 @@ _log = logging.getLogger(__name__)
 _window_message_hooks: Dict[str, List[Callable]] = {}
 _message_log: List[Dict[str, Any]] = []
 
-# Current session context for sending messages
-_current_session_context: Optional[Dict[str, Any]] = None
+# Session contexts registry for sending messages
+_session_contexts: Dict[str, Dict[str, Any]] = {}
 
 
 class WindowMessageFeature(BaseFeatureProtocol):
@@ -45,6 +45,7 @@ class WindowMessageFeature(BaseFeatureProtocol):
     def routes(self) -> List[Route]:
         return [
             Route("/api/window-message", self._send_message, methods=["POST"]),
+            Route("/api/window-message/receive", self._receive_message, methods=["POST"]),
             Route("/api/window-message/hooks", self._list_hooks, methods=["GET"]),
             Route("/api/window-message/log", self._message_log, methods=["GET"]),
             Route("/sse/window-message", self._sse_stream, methods=["GET"]),
@@ -96,19 +97,35 @@ class WindowMessageFeature(BaseFeatureProtocol):
         
         _message_log.append(message)
         
-        # If there's an active session context, also deliver via SSE
-        if _current_session_context and "sse_queue" in _current_session_context:
-            try:
-                await _current_session_context["sse_queue"].put({
-                    "type": "window.message",
-                    "data": data,
-                    "target": target,
-                })
-            except Exception as e:
-                _log.warning(f"Failed to send message via SSE: {e}")
+        # Send to all active session contexts
+        for session_id, context in _session_contexts.items():
+            if "sse_queue" in context:
+                try:
+                    await context["sse_queue"].put({
+                        "type": "window.message",
+                        "data": data,
+                        "target": target,
+                    })
+                except Exception as e:
+                    _log.warning(f"Failed to send message via SSE to session {session_id}: {e}")
         
         return JSONResponse({"status": "sent", "target": target})
 
+    async def _receive_message(self, request: Request) -> JSONResponse:
+        """API endpoint to receive window messages from frontend."""
+        try:
+            body = await request.json()
+        except Exception:
+            return JSONResponse({"error": "Invalid JSON"}, status_code=400)
+        
+        data = body.get("data", {})
+        source = body.get("source")
+        
+        # Handle the incoming message
+        await handle_window_message(data, source)
+        
+        return JSONResponse({"status": "received"})
+
     async def _list_hooks(self, request: Request) -> JSONResponse:
         """List all registered window message hooks."""
         hooks_info = {}
@@ -141,9 +158,12 @@ class WindowMessageFeature(BaseFeatureProtocol):
             # Create a queue for this connection
             message_queue = asyncio.Queue()
             
+            # Generate a unique session ID for this connection
+            import uuid
+            session_id = str(uuid.uuid4())
+            
             # Store session context for message sending
-            global _current_session_context
-            _current_session_context = {"sse_queue": message_queue}
+            _session_contexts[session_id] = {"sse_queue": message_queue}
             
             try:
                 while True:
@@ -157,7 +177,9 @@ class WindowMessageFeature(BaseFeatureProtocol):
             except asyncio.CancelledError:
                 pass
             finally:
-                _current_session_context = None
+                # Clean up session context
+                if session_id in _session_contexts:
+                    del _session_contexts[session_id]
         
         return StreamingResponse(
             event_stream(),
@@ -276,24 +298,25 @@ async def send_window_message(data: Dict[str, Any], target: str = "parent") -> N
     
     _message_log.append(message)
     
-    # Send via current SSE connection if available
-    if _current_session_context and "sse_queue" in _current_session_context:
-        try:
-            await _current_session_context["sse_queue"].put({
-                "type": "window.message",
-                "data": data,
-                "target": target,
-            })
-        except Exception as e:
-            _log.warning(f"Failed to send window message via SSE: {e}")
+    # Send to all active SSE connections
+    for session_id, context in _session_contexts.items():
+        if "sse_queue" in context:
+            try:
+                await context["sse_queue"].put({
+                    "type": "window.message",
+                    "data": data,
+                    "target": target,
+                })
+            except Exception as e:
+                _log.warning(f"Failed to send window message via SSE to session {session_id}: {e}")
 
 
 def reset_window_message_state() -> None:
     """Reset window message state for testing."""
-    global _window_message_hooks, _message_log, _current_session_context
+    global _window_message_hooks, _message_log, _session_contexts
     _window_message_hooks.clear()
     _message_log.clear()
-    _current_session_context = None
+    _session_contexts.clear()
 
 
 # Public decorators

tests/unit/test_audio_hooks.py

@@ -46,15 +46,15 @@ class TestAudioFeature:
         
         # Add some hooks
         @on_audio_start
-        async def start_hook():
+        async def start_hook(session_id):
             pass
         
         @on_audio_chunk
-        async def chunk_hook(pcm, sample_rate):
+        async def chunk_hook(session_id, pcm, sample_rate):
             pass
         
         @on_audio_end
-        async def end_hook():
+        async def end_hook(session_id):
             pass
         
         health = await feature.health()
@@ -108,15 +108,15 @@ class TestAudioHookRegistration:
         """Test decorator syntax for audio hooks."""
         
         @on_audio_start
-        def start():
+        def start(session_id):
             pass
         
         @on_audio_chunk
-        def chunk(pcm, sample_rate):
+        def chunk(session_id, pcm, sample_rate):
             pass
         
         @on_audio_end
-        def end():
+        def end(session_id):
             pass
         
         assert len(_audio_start_hooks) == 1
@@ -155,20 +155,20 @@ class TestAudioHookExecution:
         execution_order = []
         
         @on_audio_start
-        def hook1():
+        def hook1(session_id):
             execution_order.append(1)
         
         @on_audio_start
-        async def hook2():
+        async def hook2(session_id):
             execution_order.append(2)
         
         @on_audio_start
-        def hook3():
+        def hook3(session_id):
             execution_order.append(3)
         
         # Create feature instance and trigger hooks
         feature = AudioFeature()
-        await feature._trigger_start_hooks()
+        await feature._trigger_start_hooks("test_session")
         
         assert execution_order == [1, 2, 3]
     
@@ -178,11 +178,11 @@ class TestAudioHookExecution:
         received_chunks = []
         
         @on_audio_chunk
-        def hook1(pcm_data, sample_rate):
+        def hook1(session_id, pcm_data, sample_rate):
             received_chunks.append(("hook1", len(pcm_data), sample_rate))
         
         @on_audio_chunk
-        async def hook2(pcm_data, sample_rate):
+        async def hook2(session_id, pcm_data, sample_rate):
             received_chunks.append(("hook2", len(pcm_data), sample_rate))
         
         # Create feature instance and trigger hooks
@@ -190,7 +190,7 @@ class TestAudioHookExecution:
         test_data = b"test_pcm_data"
         test_rate = 16000
         
-        await feature._trigger_chunk_hooks(test_data, test_rate)
+        await feature._trigger_chunk_hooks("test_session", test_data, test_rate)
         
         assert len(received_chunks) == 2
         assert ("hook1", len(test_data), test_rate) in received_chunks
@@ -202,20 +202,20 @@ class TestAudioHookExecution:
         execution_order = []
         
         @on_audio_end
-        def hook1():
+        def hook1(session_id):
             execution_order.append(1)
         
         @on_audio_end
-        async def hook2():
+        async def hook2(session_id):
             execution_order.append(2)
         
         @on_audio_end
-        def hook3():
+        def hook3(session_id):
             execution_order.append(3)
         
         # Create feature instance and trigger hooks
         feature = AudioFeature()
-        await feature._trigger_end_hooks()
+        await feature._trigger_end_hooks("test_session")
         
         assert execution_order == [1, 2, 3]
     
@@ -225,21 +225,21 @@ class TestAudioHookExecution:
         execution_order = []
         
         @on_audio_start
-        def good_hook1():
+        def good_hook1(session_id):
             execution_order.append(1)
         
         @on_audio_start
-        def failing_hook():
+        def failing_hook(session_id):
             execution_order.append("fail")
             raise RuntimeError("Hook failed")
         
         @on_audio_start
-        def good_hook2():
+        def good_hook2(session_id):
             execution_order.append(2)
         
         # Should not raise exception
         feature = AudioFeature()
-        await feature._trigger_start_hooks()
+        await feature._trigger_start_hooks("test_session")
         
         # All hooks should have been attempted
         assert execution_order == [1, "fail", 2]
@@ -368,7 +368,7 @@ class TestAudioStats:
         _audio_sessions["test"] = {}
         
         @on_audio_start
-        def hook():
+        def hook(session_id):
             pass
         
         assert len(_audio_start_hooks) == 1
@@ -398,18 +398,18 @@ class TestAudioIntegration:
         stt_buffer = {"chunks": [], "finalized": False, "transcript": ""}
         
         @on_audio_start
-        async def start():
+        async def start(session_id):
             # Reset buffer for new session
             stt_buffer["chunks"].clear()
             stt_buffer["finalized"] = False
         
         @on_audio_chunk
-        async def chunk(pcm_data, sample_rate):
+        async def chunk(session_id, pcm_data, sample_rate):
             # Accumulate audio chunks
             stt_buffer["chunks"].append(pcm_data)
         
         @on_audio_end
-        async def end():
+        async def end(session_id):
             # Simulate STT processing
             if stt_buffer["chunks"]:
                 stt_buffer["transcript"] = f"Processed {len(stt_buffer['chunks'])} chunks"
@@ -418,20 +418,20 @@ class TestAudioIntegration:
         feature = AudioFeature()
         
         # Simulate audio session
-        await feature._trigger_start_hooks()
+        await feature._trigger_start_hooks("test_session")
         assert len(stt_buffer["chunks"]) == 0
         assert not stt_buffer["finalized"]
         
         # Send audio chunks
         chunks = [b"chunk1", b"chunk2", b"chunk3"]
         for chunk in chunks:
-            await feature._trigger_chunk_hooks(chunk, 16000)
+            await feature._trigger_chunk_hooks("test_session", chunk, 16000)
         
         assert len(stt_buffer["chunks"]) == len(chunks)
         assert stt_buffer["chunks"] == chunks
         
         # End session
-        await feature._trigger_end_hooks()
+        await feature._trigger_end_hooks("test_session")
         assert stt_buffer["finalized"]
         assert stt_buffer["transcript"] == "Processed 3 chunks"
     
@@ -445,35 +445,35 @@ class TestAudioIntegration:
         }
         
         @on_audio_start
-        async def start_all():
+        async def start_all(session_id):
             for provider in providers.values():
                 provider["active"] = True
                 provider["chunks"].clear()
         
         @on_audio_chunk
-        async def chunk_to_whisper(pcm_data, sample_rate):
+        async def chunk_to_whisper(session_id, pcm_data, sample_rate):
             if providers["whisper"]["active"]:
                 providers["whisper"]["chunks"].append(pcm_data)
         
         @on_audio_chunk
-        async def chunk_to_deepgram(pcm_data, sample_rate):
+        async def chunk_to_deepgram(session_id, pcm_data, sample_rate):
             if providers["deepgram"]["active"]:
                 providers["deepgram"]["chunks"].append(pcm_data)
         
         @on_audio_chunk
-        async def chunk_to_vosk(pcm_data, sample_rate):
+        async def chunk_to_vosk(session_id, pcm_data, sample_rate):
             if providers["vosk"]["active"]:
                 providers["vosk"]["chunks"].append(pcm_data)
         
         @on_audio_end
-        async def end_all():
+        async def end_all(session_id):
             for provider in providers.values():
                 provider["active"] = False
         
         feature = AudioFeature()
         
         # Start session
-        await feature._trigger_start_hooks()
+        await feature._trigger_start_hooks("test_session")
         
         # All providers should be active
         for provider in providers.values():
@@ -482,7 +482,7 @@ class TestAudioIntegration:
         
         # Send audio chunk
         test_chunk = b"test_audio_data"
-        await feature._trigger_chunk_hooks(test_chunk, 16000)
+        await feature._trigger_chunk_hooks("test_session", test_chunk, 16000)
         
         # All providers should have received the chunk
         for provider in providers.values():
@@ -490,7 +490,7 @@ class TestAudioIntegration:
             assert provider["chunks"][0] == test_chunk
         
         # End session
-        await feature._trigger_end_hooks()
+        await feature._trigger_end_hooks("test_session")
         
         # All providers should be deactivated
         for provider in providers.values():
@@ -502,7 +502,7 @@ class TestAudioIntegration:
         received_chunks = []
         
         @on_audio_chunk
-        async def ordered_processor(pcm_data, sample_rate):
+        async def ordered_processor(session_id, pcm_data, sample_rate):
             # Decode chunk number from data
             chunk_num = int(pcm_data.decode())
             received_chunks.append(chunk_num)
@@ -512,7 +512,7 @@ class TestAudioIntegration:
         # Send chunks in order
         for i in range(10):
             chunk_data = str(i).encode()
-            await feature._trigger_chunk_hooks(chunk_data, 16000)
+            await feature._trigger_chunk_hooks("test_session", chunk_data, 16000)
         
         # Chunks should be processed in order
         assert received_chunks == list(range(10))
@@ -523,7 +523,7 @@ class TestAudioIntegration:
         session_data = {}
         
         @on_audio_chunk
-        async def track_by_session(pcm_data, sample_rate):
+        async def track_by_session(session_id, pcm_data, sample_rate):
             # Extract session ID from chunk data (simulated)
             session_id = pcm_data.decode().split(":")[0]
             if session_id not in session_data:
@@ -539,7 +539,7 @@ class TestAudioIntegration:
         for session_id in sessions:
             for chunk_num in range(chunks_per_session):
                 chunk_data = f"{session_id}:chunk{chunk_num}".encode()
-                await feature._trigger_chunk_hooks(chunk_data, 16000)
+                await feature._trigger_chunk_hooks("test_session", chunk_data, 16000)
         
         # Each session should have received its chunks
         assert len(session_data) == len(sessions)

tests/unit/test_window_message.py

@@ -15,7 +15,7 @@ from praisonaiui.features.window_message import (
     reset_window_message_state,
     _window_message_hooks,
     _message_log,
-    _current_session_context,
+    _session_contexts,
 )
 
 
@@ -32,7 +32,7 @@ class TestWindowMessageFeature:
         
         assert feature.name == "window_message"
         assert feature.description == "Browser window.postMessage communication"
-        assert len(feature.routes()) == 4
+        assert len(feature.routes()) == 5  # Added receive endpoint
         assert len(feature.cli_commands()) == 1
     
     @pytest.mark.asyncio
@@ -263,10 +263,9 @@ class TestWindowMessageSending:
         # Mock session context with SSE queue
         sse_queue = asyncio.Queue()
         
-        # Set global session context
-        global _current_session_context
-        original_context = _current_session_context
-        _current_session_context = {"sse_queue": sse_queue}
+        # Add a session context for testing
+        test_session_id = "test_session"
+        _session_contexts[test_session_id] = {"sse_queue": sse_queue}
         
         try:
             await send_window_message({"type": "test"}, "parent")
@@ -278,7 +277,9 @@ class TestWindowMessageSending:
             assert message["data"]["type"] == "test"
             assert message["target"] == "parent"
         finally:
-            _current_session_context = original_context
+            # Clean up test session
+            if test_session_id in _session_contexts:
+                del _session_contexts[test_session_id]
 
 
 class TestWindowMessageSecurity: