adds wait for completion for audio upload

Author: seanreed1111

src/agent.py

@@ -160,19 +160,22 @@ def on_session_close(_event):
         logger.info(f"Session closing for room {room_name}, saving transcript...")
 
         async def cleanup():
-            # Stop egress recording - this triggers S3 upload of the audio file
+            # Stop egress recording and wait for S3 upload to complete
             if egress_manager is not None:
                 try:
-                    logger.info("Stopping egress recording...")
-                    stopped = await egress_manager.stop_recording()
-                    if stopped:
+                    logger.info("Stopping egress recording and waiting for upload...")
+                    file_info = await egress_manager.stop_recording()
+                    if file_info:
                         logger.info(
-                            f"Egress recording stopped for room {room_name}, "
-                            f"audio uploaded to s3://{S3_BUCKET}/{S3_PREFIX}/"
+                            f"Audio recording uploaded for room {room_name}: "
+                            f"location={file_info.location}, "
+                            f"filename={file_info.filename}, "
+                            f"size={file_info.size} bytes"
                         )
                     else:
                         logger.warning(
-                            f"Failed to stop egress recording for room {room_name}"
+                            f"No audio file info returned for room {room_name} "
+                            "(recording may have failed or no audio was captured)"
                         )
                 except Exception as e:
                     logger.error(f"Error stopping egress recording: {e}")

src/egress_manager.py

@@ -1,12 +1,24 @@
 """Egress manager for recording dual-channel audio to S3."""
 
+import asyncio
 import os
+from dataclasses import dataclass
 
 from livekit import api
 from livekit.protocol import egress as egress_proto
 from loguru import logger
 
 
+@dataclass
+class EgressFileInfo:
+    """Information about an uploaded egress file."""
+
+    filename: str
+    location: str
+    duration: int
+    size: int
+
+
 class EgressConfig:
     """Configuration for egress recordings."""
 
@@ -63,6 +75,7 @@ def __init__(self, config: EgressConfig):
         self.config = config
         self._api: api.LiveKitAPI | None = None
         self._egress_id: str | None = None
+        self._expected_filepath: str | None = None
 
     @property
     def livekit_api(self) -> api.LiveKitAPI:
@@ -120,6 +133,11 @@ async def start_dual_channel_recording(
             else:
                 filepath = f"{filepath_prefix}/{{room_name}}-{{time}}.ogg"
 
+            # Store expected filepath for logging
+            self._expected_filepath = filepath
+            expected_s3_path = f"s3://{self.config.s3_bucket}/{filepath}"
+            logger.info(f"Expected audio file path: {expected_s3_path}")
+
             file_output = egress_proto.EncodedFileOutput(
                 filepath=filepath,
                 s3=s3_upload,
@@ -138,35 +156,145 @@ async def start_dual_channel_recording(
 
             self._egress_id = info.egress_id
             logger.info(
-                f"Started dual-channel egress recording for room {room_name}, "
-                f"egress_id={self._egress_id}"
+                f"Started egress recording for room {room_name}, "
+                f"egress_id={self._egress_id}, "
+                f"expected_file={expected_s3_path}"
             )
             return self._egress_id
 
         except Exception as e:
             logger.error(f"Failed to start egress recording: {e}")
             return None
 
-    async def stop_recording(self) -> bool:
-        """Stop the active egress recording.
+    async def stop_recording(self) -> EgressFileInfo | None:
+        """Stop the active egress recording and wait for upload to complete.
 
         Returns:
-            True if stopped successfully or no active recording, False on error
+            EgressFileInfo with file details if successful, None on error
         """
         if not self._egress_id:
             logger.debug("No active egress to stop")
-            return True
+            return None
 
+        egress_id = self._egress_id
         try:
+            logger.info(f"Stopping egress recording, egress_id={egress_id}...")
             await self.livekit_api.egress.stop_egress(
-                egress_proto.StopEgressRequest(egress_id=self._egress_id)
+                egress_proto.StopEgressRequest(egress_id=egress_id)
+            )
+            logger.info(
+                f"Stop request sent for egress_id={egress_id}, waiting for upload to complete..."
             )
-            logger.info(f"Stopped egress recording, egress_id={self._egress_id}")
+
+            # Wait for the egress to complete and get file info
+            file_info = await self._wait_for_completion(egress_id)
             self._egress_id = None
-            return True
+            self._expected_filepath = None
+            return file_info
+
         except Exception as e:
             logger.error(f"Failed to stop egress recording: {e}")
-            return False
+            return None
+
+    async def _wait_for_completion(
+        self, egress_id: str, timeout: float = 60.0, poll_interval: float = 1.0
+    ) -> EgressFileInfo | None:
+        """Wait for egress to complete and return file info.
+
+        Args:
+            egress_id: The egress ID to wait for
+            timeout: Maximum time to wait in seconds
+            poll_interval: Time between status checks in seconds
+
+        Returns:
+            EgressFileInfo with file details if successful, None on error/timeout
+        """
+        start_time = asyncio.get_event_loop().time()
+        last_status = None
+
+        while True:
+            elapsed = asyncio.get_event_loop().time() - start_time
+            if elapsed > timeout:
+                logger.error(
+                    f"Timeout waiting for egress completion after {timeout}s, "
+                    f"egress_id={egress_id}, last_status={last_status}"
+                )
+                return None
+
+            try:
+                # List egress by ID to get current status
+                response = await self.livekit_api.egress.list_egress(
+                    egress_proto.ListEgressRequest(egress_id=egress_id)
+                )
+
+                if not response.items:
+                    logger.warning(f"Egress {egress_id} not found in list response")
+                    await asyncio.sleep(poll_interval)
+                    continue
+
+                egress_info = response.items[0]
+                status = egress_info.status
+                last_status = egress_proto.EgressStatus.Name(status)
+
+                logger.debug(
+                    f"Egress status: {last_status}, egress_id={egress_id}, "
+                    f"elapsed={elapsed:.1f}s"
+                )
+
+                # Check terminal states
+                if status == egress_proto.EgressStatus.EGRESS_COMPLETE:
+                    # Extract file info from results
+                    if egress_info.file_results:
+                        file_result = egress_info.file_results[0]
+                        file_info = EgressFileInfo(
+                            filename=file_result.filename,
+                            location=file_result.location,
+                            duration=file_result.duration,
+                            size=file_result.size,
+                        )
+                        logger.info(
+                            f"Egress completed successfully! "
+                            f"File uploaded to: {file_info.location}, "
+                            f"filename={file_info.filename}, "
+                            f"duration={file_info.duration}ns, "
+                            f"size={file_info.size} bytes"
+                        )
+                        return file_info
+                    else:
+                        logger.warning(
+                            f"Egress completed but no file_results found, "
+                            f"egress_id={egress_id}"
+                        )
+                        return None
+
+                elif status == egress_proto.EgressStatus.EGRESS_FAILED:
+                    error_msg = egress_info.error or "Unknown error"
+                    logger.error(f"Egress failed: {error_msg}, egress_id={egress_id}")
+                    return None
+
+                elif status == egress_proto.EgressStatus.EGRESS_ABORTED:
+                    logger.warning(f"Egress was aborted, egress_id={egress_id}")
+                    return None
+
+                elif status == egress_proto.EgressStatus.EGRESS_LIMIT_REACHED:
+                    logger.warning(f"Egress limit reached, egress_id={egress_id}")
+                    # Still try to get file info as partial recording may exist
+                    if egress_info.file_results:
+                        file_result = egress_info.file_results[0]
+                        return EgressFileInfo(
+                            filename=file_result.filename,
+                            location=file_result.location,
+                            duration=file_result.duration,
+                            size=file_result.size,
+                        )
+                    return None
+
+                # Still in progress, wait and poll again
+                await asyncio.sleep(poll_interval)
+
+            except Exception as e:
+                logger.error(f"Error polling egress status: {e}")
+                await asyncio.sleep(poll_interval)
 
     async def close(self) -> None:
         """Clean up resources."""

tests/test_egress_manager.py

@@ -3,8 +3,14 @@
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
+from livekit.protocol import egress as egress_proto
 
-from egress_manager import EgressConfig, EgressManager, create_default_egress_manager
+from egress_manager import (
+    EgressConfig,
+    EgressFileInfo,
+    EgressManager,
+    create_default_egress_manager,
+)
 
 
 class TestEgressConfig:
@@ -153,16 +159,35 @@ async def test_start_dual_channel_recording_failure(self):
 
     @pytest.mark.asyncio
     async def test_stop_recording_success(self):
-        """Test successful recording stop."""
+        """Test successful recording stop with file info."""
         config = EgressConfig(
             s3_bucket="test-bucket",
         )
         manager = EgressManager(config)
         manager._egress_id = "EG_TEST123456"
 
+        # Mock file result
+        mock_file_result = MagicMock()
+        mock_file_result.filename = "test-room-20251212-120000.ogg"
+        mock_file_result.location = (
+            "s3://test-bucket/audio/test-room-20251212-120000.ogg"
+        )
+        mock_file_result.duration = 60000000000  # 60 seconds in nanoseconds
+        mock_file_result.size = 1024000
+
+        # Mock egress info with COMPLETE status
+        mock_egress_info = MagicMock()
+        mock_egress_info.status = egress_proto.EgressStatus.EGRESS_COMPLETE
+        mock_egress_info.file_results = [mock_file_result]
+
+        # Mock list response
+        mock_list_response = MagicMock()
+        mock_list_response.items = [mock_egress_info]
+
         # Mock the LiveKit API
         mock_egress_service = AsyncMock()
         mock_egress_service.stop_egress = AsyncMock()
+        mock_egress_service.list_egress = AsyncMock(return_value=mock_list_response)
 
         mock_api = MagicMock()
         mock_api.egress = mock_egress_service
@@ -171,7 +196,12 @@ async def test_stop_recording_success(self):
 
         result = await manager.stop_recording()
 
-        assert result is True
+        assert result is not None
+        assert isinstance(result, EgressFileInfo)
+        assert result.filename == "test-room-20251212-120000.ogg"
+        assert result.location == "s3://test-bucket/audio/test-room-20251212-120000.ogg"
+        assert result.duration == 60000000000
+        assert result.size == 1024000
         assert manager.egress_id is None
         mock_egress_service.stop_egress.assert_called_once()
 
@@ -185,7 +215,7 @@ async def test_stop_recording_no_active_egress(self):
 
         result = await manager.stop_recording()
 
-        assert result is True
+        assert result is None  # No file info when no egress was active
 
     @pytest.mark.asyncio
     async def test_stop_recording_failure(self):
@@ -209,7 +239,7 @@ async def test_stop_recording_failure(self):
 
         result = await manager.stop_recording()
 
-        assert result is False
+        assert result is None  # Returns None on failure
 
     @pytest.mark.asyncio
     async def test_close(self):