From b74c2f9b2accccb94f53671e51c9ac514850554a Mon Sep 17 00:00:00 2001
From: yumoqing <yumoqing@gmail.com>
Date: Wed, 11 Sep 2024 14:53:53 +0800
Subject: [PATCH] bugfix

---
 rtcllm/aav.py | 28 +++++++++++++++++++++-------
 rtcllm/rtc.py |  4 ++--
 2 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/rtcllm/aav.py b/rtcllm/aav.py
index 43c4cb1..9d03c80 100644
--- a/rtcllm/aav.py
+++ b/rtcllm/aav.py
@@ -1,4 +1,5 @@
 import random
+from traceback import print_exc
 from aiortc.contrib.media import MediaBlackhole, MediaPlayer, MediaRecorder, MediaRelay
 from aiortc import MediaStreamTrack, VideoStreamTrack, AudioStreamTrack
 
@@ -43,11 +44,18 @@ class MyAudioStreamTrack(AudioStreamTrack):
 		print(f'{self.kind=}, {self.__class__.__name__}, {dir(self)}')
 
 	async def recv(self):
-		print(f'============{self.__class__.__name__}, {self.source.duration=}, {self.source.time=}')
+		try:
+			return await self._recv()
+		except MediaStreamError:
+			self.source_reload()
+			return await self._recv()
+		except Exception as e:
+			print_exc()
+			raise e
+
+	async def _recv(self):
 		if self.source is None:
 			return None
-		if self.track.readyState != 'live':
-			self.source_reload()
 		f = await self.track.recv()
 		return f
 
@@ -73,11 +81,18 @@ class MyVideoStreamTrack(VideoStreamTrack):
 		print(f'{self.kind=}, {self.__class__.__name__}, {dir(self)}')
 
 	async def recv(self):
-		print(f'============{self.__class__.__name__}, {self.source.duration=}, {self.source.time=}')
+		try:
+			return await self._recv()
+		except MediaStreamError:
+			self.source_reload()
+			return self._recv()
+		except Exception as e:
+			print_exc()
+			raise e
+
+	async def _recv(self):
 		if self.source is None:
 			return None
-		if self.track.readyState != 'live':
-			self.source_reload()
 		f = await self.track.recv()
 		return f
 
@@ -107,7 +122,6 @@ class LoopingVideoTrack(VideoStreamTrack):
 
 	async def recv(self):
 		frame = await self.player.video.recv()
-		print(f'{self.__class__.__name__}, {self.source.duration=}, {self.source.time=}')
 		if self.player.video.readyState != 'live':
 			self.player = MyMediaPlayer(self.player._file_path)
 			frame = await self.player.video.recv()
diff --git a/rtcllm/rtc.py b/rtcllm/rtc.py
index 0b7c3e3..095673d 100644
--- a/rtcllm/rtc.py
+++ b/rtcllm/rtc.py
@@ -127,8 +127,8 @@ class RTCLLM:
 			'pc':pc
 		})
 		pc.addTrack(MyAudioStreamTrack(player))
-		pc.addTrack(LoopingVideoTrack('./1.mp4'))
-		# pc.addTrack(MyVideoStreamTrack(player))
+		pc.addTrack(MyVideoStreamTrack(player))
+		# pc.addTrack(LoopingVideoTrack('./1.mp4'))
 		# pc.addTrack(player.video)
 		await self.ws_send(json.dumps({'type':'callAccepted', 'to':data['from']}))