add timeout to agent.receive and fix integ tests

Author: Murat Kaan Meral

src/strands/experimental/bidirectional_streaming/agent/agent.py

@@ -339,9 +339,11 @@ async def receive(self) -> AsyncIterable[BidiOutputEvent]:
         """
         while self.active:
             try:
-                event = await self._output_queue.get()
+                # Use a timeout to periodically check if we should stop
+                event = await asyncio.wait_for(self._output_queue.get(), timeout=0.5)
                 yield event
             except asyncio.TimeoutError:
+                # Timeout allows us to check self.active periodically
                 continue
 
     async def stop(self) -> None:

tests/strands/experimental/bidirectional_streaming/models/test_gemini_live.py

@@ -7,6 +7,8 @@
 - Event receiving and conversion
 """
 
+import base64
+import json
 import unittest.mock
 
 import pytest
@@ -16,8 +18,13 @@
 from strands.experimental.bidirectional_streaming.models.gemini_live import BidiGeminiLiveModel
 from strands.experimental.bidirectional_streaming.types.events import (
     BidiAudioInputEvent,
+    BidiAudioStreamEvent,
+    BidiConnectionCloseEvent,
+    BidiConnectionStartEvent,
     BidiImageInputEvent,
+    BidiInterruptionEvent,
     BidiTextInputEvent,
+    BidiTranscriptStreamEvent,
 )
 from strands.types._events import ToolResultEvent
 from strands.types.tools import ToolResult
@@ -198,7 +205,6 @@ async def test_send_all_content_types(mock_genai_client, model):
     assert content.parts[0].text == "Hello"
 
     # Test audio input (base64 encoded)
-    import base64
     audio_b64 = base64.b64encode(b"audio_bytes").decode('utf-8')
     audio_input = BidiAudioInputEvent(
         audio=audio_b64,
@@ -219,7 +225,6 @@ async def test_send_all_content_types(mock_genai_client, model):
     mock_live_session.send.assert_called_once()
 
     # Test tool result
-    from strands.types._events import ToolResultEvent
     tool_result: ToolResult = {
         "toolUseId": "tool-123",
         "status": "success",
@@ -255,11 +260,6 @@ async def test_send_edge_cases(mock_genai_client, model):
 @pytest.mark.asyncio
 async def test_receive_lifecycle_events(mock_genai_client, model, agenerator):
     """Test that receive() emits connection start and end events."""
-    from strands.experimental.bidirectional_streaming.types.events import (
-        BidiConnectionStartEvent,
-        BidiConnectionCloseEvent,
-    )
-    
     _, mock_live_session, _ = mock_genai_client
     mock_live_session.receive.return_value = agenerator([])
 
@@ -285,12 +285,6 @@ async def test_receive_lifecycle_events(mock_genai_client, model, agenerator):
 @pytest.mark.asyncio
 async def test_event_conversion(mock_genai_client, model):
     """Test conversion of all Gemini Live event types to standard format."""
-    from strands.experimental.bidirectional_streaming.types.events import (
-        BidiTranscriptStreamEvent,
-        BidiAudioStreamEvent,
-        BidiInterruptionEvent,
-    )
-    
     _, _, _ = mock_genai_client
     await model.start()
 
@@ -311,7 +305,6 @@ async def test_event_conversion(mock_genai_client, model):
     assert text_event.current_transcript == "Hello from Gemini"
 
     # Test audio output (base64 encoded)
-    import base64
     mock_audio = unittest.mock.Mock()
     mock_audio.text = None
     mock_audio.data = b"audio_data"

tests/strands/experimental/bidirectional_streaming/models/test_novasonic.py

@@ -15,6 +15,17 @@
 from strands.experimental.bidirectional_streaming.models.novasonic import (
     BidiNovaSonicModel,
 )
+from strands.experimental.bidirectional_streaming.types.events import (
+    BidiAudioInputEvent,
+    BidiAudioStreamEvent,
+    BidiImageInputEvent,
+    BidiInterruptionEvent,
+    BidiResponseStartEvent,
+    BidiTextInputEvent,
+    BidiTranscriptStreamEvent,
+    BidiUsageEvent,
+)
+from strands.types._events import ToolResultEvent
 from strands.types.tools import ToolResult
 
 
@@ -131,12 +142,6 @@ async def test_connection_edge_cases(nova_model, mock_client, mock_stream, model
 @pytest.mark.asyncio
 async def test_send_all_content_types(nova_model, mock_client, mock_stream):
     """Test sending all content types through unified send() method."""
-    from strands.experimental.bidirectional_streaming.types.events import (
-        BidiTextInputEvent,
-        BidiAudioInputEvent,
-    )
-    from strands.types._events import ToolResultEvent
-    
     with patch.object(nova_model, "_initialize_client", new_callable=AsyncMock):
         nova_model.client = mock_client
 
@@ -177,11 +182,6 @@ async def test_send_all_content_types(nova_model, mock_client, mock_stream):
 @pytest.mark.asyncio
 async def test_send_edge_cases(nova_model, mock_client, mock_stream, caplog):
     """Test send() edge cases and error handling."""
-    from strands.experimental.bidirectional_streaming.types.events import (
-        BidiTextInputEvent,
-        BidiImageInputEvent,
-    )
-    
     with patch.object(nova_model, "_initialize_client", new_callable=AsyncMock):
         nova_model.client = mock_client
 
@@ -191,7 +191,6 @@ async def test_send_edge_cases(nova_model, mock_client, mock_stream, caplog):
 
         # Test image content (not supported, base64 encoded, no encoding parameter)
         await nova_model.start()
-        import base64
         image_b64 = base64.b64encode(b"image data").decode('utf-8')
         image_event = BidiImageInputEvent(
             image=image_b64,
@@ -237,7 +236,6 @@ async def mock_wait_for(*args, **kwargs):
 async def test_event_conversion(nova_model):
     """Test conversion of all Nova Sonic event types to standard format."""
     # Test audio output (now returns BidiAudioStreamEvent)
-    from strands.experimental.bidirectional_streaming.types.events import BidiAudioStreamEvent
     audio_bytes = b"test audio data"
     audio_base64 = base64.b64encode(audio_bytes).decode("utf-8")
     nova_event = {"audioOutput": {"content": audio_base64}}
@@ -251,7 +249,6 @@ async def test_event_conversion(nova_model):
     assert result.get("sample_rate") == 24000
 
     # Test text output (now returns BidiTranscriptStreamEvent)
-    from strands.experimental.bidirectional_streaming.types.events import BidiTranscriptStreamEvent
     nova_event = {"textOutput": {"content": "Hello, world!", "role": "ASSISTANT"}}
     result = nova_model._convert_nova_event(nova_event)
     assert result is not None
@@ -282,7 +279,6 @@ async def test_event_conversion(nova_model):
     assert tool_use["input"] == tool_input
 
     # Test interruption (now returns BidiInterruptionEvent)
-    from strands.experimental.bidirectional_streaming.types.events import BidiInterruptionEvent
     nova_event = {"stopReason": "INTERRUPTED"}
     result = nova_model._convert_nova_event(nova_event)
     assert result is not None
@@ -291,7 +287,6 @@ async def test_event_conversion(nova_model):
     assert result.get("reason") == "user_speech"
 
     # Test usage metrics (now returns BidiUsageEvent)
-    from strands.experimental.bidirectional_streaming.types.events import BidiUsageEvent
     nova_event = {
         "usageEvent": {
             "totalTokens": 100,
@@ -315,7 +310,6 @@ async def test_event_conversion(nova_model):
     assert result.get("outputTokens") == 60
 
     # Test content start tracks role and emits BidiResponseStartEvent
-    from strands.experimental.bidirectional_streaming.types.events import BidiResponseStartEvent
     nova_event = {"contentStart": {"role": "USER"}}
     result = nova_model._convert_nova_event(nova_event)
     assert result is not None
@@ -349,16 +343,13 @@ async def test_audio_connection_lifecycle(nova_model, mock_client, mock_stream):
 @pytest.mark.asyncio
 async def test_silence_detection(nova_model, mock_client, mock_stream):
     """Test that silence detection automatically ends audio input."""
-    from strands.experimental.bidirectional_streaming.types.events import BidiAudioInputEvent
-    
     with patch.object(nova_model, "_initialize_client", new_callable=AsyncMock):
         nova_model.client = mock_client
         nova_model.silence_threshold = 0.1  # Short threshold for testing
 
         await nova_model.start()
 
         # Send audio to start connection (base64 encoded)
-        import base64
         audio_b64 = base64.b64encode(b"audio data").decode('utf-8')
         audio_event = BidiAudioInputEvent(
             audio=audio_b64,

tests/strands/experimental/bidirectional_streaming/models/test_openai_realtime.py

@@ -18,9 +18,13 @@
 from strands.experimental.bidirectional_streaming.models.openai import BidiOpenAIRealtimeModel
 from strands.experimental.bidirectional_streaming.types.events import (
     BidiAudioInputEvent,
+    BidiAudioStreamEvent,
     BidiImageInputEvent,
+    BidiInterruptionEvent,
     BidiTextInputEvent,
+    BidiTranscriptStreamEvent,
 )
+from strands.types._events import ToolResultEvent
 from strands.types.tools import ToolResult
 
 
@@ -222,8 +226,6 @@ async def async_connect(*args, **kwargs):
 @pytest.mark.asyncio
 async def test_send_all_content_types(mock_websockets_connect, model):
     """Test sending all content types through unified send() method."""
-    from strands.types._events import ToolResultEvent
-    
     _, mock_ws = mock_websockets_connect
     await model.start()
 
@@ -343,7 +345,6 @@ async def test_event_conversion(mock_websockets_connect, model):
     await model.start()
 
     # Test audio output (now returns list with BidiAudioStreamEvent)
-    from strands.experimental.bidirectional_streaming.types.events import BidiAudioStreamEvent
     audio_event = {
         "type": "response.output_audio.delta",
         "delta": base64.b64encode(b"audio_data").decode()
@@ -357,7 +358,6 @@ async def test_event_conversion(mock_websockets_connect, model):
     assert converted[0].get("format") == "pcm"
 
     # Test text output (now returns list with BidiTranscriptStreamEvent)
-    from strands.experimental.bidirectional_streaming.types.events import BidiTranscriptStreamEvent
     text_event = {
         "type": "response.output_text.delta",
         "delta": "Hello from OpenAI"
@@ -407,7 +407,6 @@ async def test_event_conversion(mock_websockets_connect, model):
     assert tool_use["input"]["expression"] == "2+2"
 
     # Test voice activity (now returns list with BidiInterruptionEvent for speech_started)
-    from strands.experimental.bidirectional_streaming.types.events import BidiInterruptionEvent
     speech_started = {
         "type": "input_audio_buffer.speech_started"
     }
@@ -465,7 +464,6 @@ def test_helper_methods(model):
     model._active = False
 
     # Test _create_text_event (now returns BidiTranscriptStreamEvent)
-    from strands.experimental.bidirectional_streaming.types.events import BidiTranscriptStreamEvent
     text_event = model._create_text_event("Hello", "user")
     assert isinstance(text_event, BidiTranscriptStreamEvent)
     assert text_event.get("type") == "bidi_transcript_stream"
@@ -476,7 +474,6 @@ def test_helper_methods(model):
     assert text_event.current_transcript == "Hello"
 
     # Test _create_voice_activity_event (now returns BidiInterruptionEvent for speech_started)
-    from strands.experimental.bidirectional_streaming.types.events import BidiInterruptionEvent
     voice_event = model._create_voice_activity_event("speech_started")
     assert isinstance(voice_event, BidiInterruptionEvent)
     assert voice_event.get("type") == "bidi_interruption"

tests_integ/bidirectional_streaming/context.py

@@ -5,6 +5,7 @@
 """
 
 import asyncio
+import base64
 import logging
 import time
 from typing import TYPE_CHECKING
@@ -81,12 +82,13 @@ async def __aenter__(self):
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         """Stop context manager, cleanup threads, and end agent session."""
-        await self.stop()
-        
-        # End agent session
+        # End agent session FIRST - this will cause receive() to exit cleanly
         if self.agent._agent_loop and self.agent._agent_loop.active:
             await self.agent.stop()
-            logger.debug("Agent session ended")
+            logger.debug("Agent session stopped")
+        
+        # Then stop the context threads
+        await self.stop()
 
         return False
 
@@ -254,8 +256,6 @@ def get_audio_outputs(self) -> list[bytes]:
         Returns:
             List of audio data bytes.
         """
-        import base64
-        
         # Drain queue first to get latest events
         events = self.get_events()
         audio_data = []
@@ -332,6 +332,7 @@ async def _input_thread(self):
 
         except asyncio.CancelledError:
             logger.debug("Input thread cancelled")
+            raise  # Re-raise to properly propagate cancellation
         except Exception as e:
             logger.error(f"Input thread error: {e}", exc_info=True)
         finally:
@@ -350,6 +351,7 @@ async def _event_collection_thread(self):
 
         except asyncio.CancelledError:
             logger.debug("Event collection thread cancelled")
+            raise  # Re-raise to properly propagate cancellation
         except Exception as e:
             logger.error(f"Event collection thread error: {e}")
 

tests_integ/bidirectional_streaming/generators/audio.py

@@ -4,6 +4,7 @@
 without requiring physical audio devices or pre-recorded files.
 """
 
+import base64
 import hashlib
 import logging
 from pathlib import Path
@@ -120,8 +121,6 @@ def create_audio_input_event(
         Returns:
             BidiAudioInputEvent dict ready for agent.send().
         """
-        import base64
-        
         # Convert bytes to base64 string for JSON compatibility
         audio_b64 = base64.b64encode(audio_data).decode('utf-8')