fix: enhance connection state checks (#2137)

* fix: enhance connection state checks

* feat: support prefetch

---------

Co-authored-by: hanhandi <1540984562@qq.com>

Author: diyuyi-agora

ai_agents/agents/ten_packages/extension/bytedance_llm_based_asr/extension.py

@@ -305,7 +305,9 @@ def is_connected(self) -> bool:
         #     ):
         #         return True  # Still consider connected during finalize grace period
 
-        return self.connected and self.client is not None
+        return (
+            self.connected and self.client is not None and self.client.connected
+        )
 
     @override
     async def send_audio(
@@ -351,6 +353,8 @@ async def send_audio(
             return True
 
         except Exception as e:
+            if self.stopped:
+                return False
             self.ten_env.log(LogLevel.ERROR, f"Error sending audio: {e}")
             await self._handle_error(e)
             return False
@@ -466,16 +470,32 @@ async def _handle_reconnect(self) -> None:
         finally:
             self._reconnecting = False
 
+    def _result_level_asr_info_fields(
+        self, result: ASRResponse
+    ) -> dict[str, Any]:
+        """Fields from vendor `result` object that belong in metadata.asr_info.
+
+        Only includes keys when present on the vendor payload (e.g. prefetch).
+        """
+        rd = result.result
+        if not rd or not isinstance(rd, dict):
+            return {}
+        if "prefetch" not in rd:
+            return {}
+        return {"prefetch": rd["prefetch"]}
+
     def _build_metadata_with_asr_info(
         self,
         base_metadata: dict[str, Any] | None = None,
         additional_fields: dict[str, Any] | None = None,
+        result_level_fields: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
         """Build metadata according to protocol: session_id at root, others in asr_info.
 
         Args:
             base_metadata: Base metadata dict (defaults to self.metadata if None)
             additional_fields: Additional fields to add to asr_info
+            result_level_fields: Vendor result-level fields (e.g. prefetch) for asr_info
 
         Returns:
             Metadata dict with structure: {"session_id": "...", "asr_info": {...}}
@@ -501,6 +521,10 @@ def _build_metadata_with_asr_info(
         if additional_fields:
             asr_info.update(additional_fields)
 
+        # Vendor result-level fields (e.g. prefetch under result.{prefetch})
+        if result_level_fields:
+            asr_info.update(result_level_fields)
+
         # Build final metadata structure
         metadata: dict[str, Any] = {}
         if session_id is not None:
@@ -510,7 +534,9 @@ def _build_metadata_with_asr_info(
         return metadata
 
     def _extract_final_result_metadata(
-        self, utterance: Utterance
+        self,
+        utterance: Utterance,
+        result_level_fields: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
         """Extract metadata from utterance additions.
 
@@ -524,11 +550,14 @@ def _extract_final_result_metadata(
             additional_fields = utterance.additions
 
         return self._build_metadata_with_asr_info(
-            additional_fields=additional_fields
+            additional_fields=additional_fields,
+            result_level_fields=result_level_fields,
         )
 
     def _extract_non_final_result_metadata(
-        self, utterance: Utterance
+        self,
+        utterance: Utterance,
+        result_level_fields: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
         """Extract metadata from utterance additions for non-final results.
 
@@ -550,7 +579,8 @@ def _extract_non_final_result_metadata(
                 additional_fields["source"] = additions["source"]
 
         return self._build_metadata_with_asr_info(
-            additional_fields=additional_fields
+            additional_fields=additional_fields,
+            result_level_fields=result_level_fields,
         )
 
     def _calculate_utterance_start_ms(
@@ -719,6 +749,8 @@ async def _on_asr_result(self, result: ASRResponse) -> None:
                 category=LOG_CATEGORY_VENDOR,
             )
 
+            result_level_asr_info = self._result_level_asr_info_fields(result)
+
             # Process utterances: send definite=true individually,
             # and concatenate adjacent definite=false utterances together
             if not result.utterances:
@@ -728,7 +760,9 @@ async def _on_asr_result(self, result: ASRResponse) -> None:
                     result.start_ms
                 )
                 # Build metadata according to protocol: session_id at root, others in asr_info
-                metadata = self._build_metadata_with_asr_info()
+                metadata = self._build_metadata_with_asr_info(
+                    result_level_fields=result_level_asr_info
+                )
                 await self._send_asr_result_from_text(
                     text=result.text,
                     is_final=False,
@@ -767,11 +801,13 @@ async def _on_asr_result(self, result: ASRResponse) -> None:
                     if is_final:
                         has_final_result = True
                         metadata = self._extract_final_result_metadata(
-                            utterance
+                            utterance,
+                            result_level_fields=result_level_asr_info,
                         )
                     else:
                         metadata = self._extract_non_final_result_metadata(
-                            utterance
+                            utterance,
+                            result_level_fields=result_level_asr_info,
                         )
 
                     await self._send_asr_result_from_text(
@@ -816,10 +852,14 @@ async def _on_asr_result(self, result: ASRResponse) -> None:
                             "start_time": first.start_time,
                             "duration_ms": last.end_time - first.start_time,
                             "metadata": (
-                                self._extract_final_result_metadata(last)
+                                self._extract_final_result_metadata(
+                                    last,
+                                    result_level_fields=result_level_asr_info,
+                                )
                                 if is_final
                                 else self._extract_non_final_result_metadata(
-                                    last
+                                    last,
+                                    result_level_fields=result_level_asr_info,
                                 )
                             ),
                             "utterance": last,  # Keep reference for timestamp tracking

ai_agents/agents/ten_packages/extension/bytedance_llm_based_asr/manifest.json

@@ -1,7 +1,7 @@
 {
   "type": "extension",
   "name": "bytedance_llm_based_asr",
-  "version": "0.3.18",
+  "version": "0.3.20",
   "dependencies": [
     {
       "type": "system",

ai_agents/agents/ten_packages/extension/bytedance_llm_based_asr/tests/test_connection_state_sync.py

@@ -0,0 +1,63 @@
+import os
+import sys
+import types
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+
+extension_dir = os.path.join(os.path.dirname(__file__), "..")
+sys.path.insert(0, extension_dir)
+
+package = types.ModuleType("bytedance_llm_based_asr")
+package.__path__ = [extension_dir]
+sys.modules["bytedance_llm_based_asr"] = package
+
+from bytedance_llm_based_asr import config as config_module
+
+sys.modules["bytedance_llm_based_asr.config"] = config_module
+
+from bytedance_llm_based_asr import extension as extension_module
+
+sys.modules["bytedance_llm_based_asr.extension"] = extension_module
+
+from bytedance_llm_based_asr.extension import BytedanceASRLLMExtension
+
+
+@pytest.fixture
+def mock_ten_env():
+    env = AsyncMock()
+    env.log = MagicMock()
+    env.log_debug = MagicMock()
+    env.log_info = MagicMock()
+    env.log_warn = MagicMock()
+    env.log_error = MagicMock()
+    return env
+
+
+@pytest.fixture
+def extension(mock_ten_env):
+    ext = BytedanceASRLLMExtension("test_extension")
+    ext.ten_env = mock_ten_env
+    ext.connected = True
+    ext.client = MagicMock()
+    ext.client.connected = False
+    ext.client.send_audio = AsyncMock()
+    return ext
+
+
+@pytest.mark.asyncio
+async def test_handle_audio_frame_buffers_when_client_state_is_disconnected(
+    extension, mock_ten_env
+):
+    frame = MagicMock()
+    frame.get_buf.return_value = b"\x00\x01"
+
+    await extension._handle_audio_frame(mock_ten_env, frame)
+
+    assert extension.is_connected() is False
+    assert extension.buffered_frames.qsize() == 1
+    extension.client.send_audio.assert_not_awaited()
+    mock_ten_env.log_debug.assert_called_once_with(
+        "send_frame: service not connected."
+    )

ai_agents/agents/ten_packages/extension/bytedance_llm_based_asr/tests/test_stop_phase_error_suppression.py

@@ -0,0 +1,104 @@
+import os
+import sys
+import types
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+
+extension_dir = os.path.join(os.path.dirname(__file__), "..")
+sys.path.insert(0, extension_dir)
+
+package = types.ModuleType("bytedance_llm_based_asr")
+package.__path__ = [extension_dir]
+sys.modules["bytedance_llm_based_asr"] = package
+
+from bytedance_llm_based_asr import config as config_module
+
+sys.modules["bytedance_llm_based_asr.config"] = config_module
+
+from bytedance_llm_based_asr import extension as extension_module
+
+sys.modules["bytedance_llm_based_asr.extension"] = extension_module
+
+from bytedance_llm_based_asr.extension import BytedanceASRLLMExtension
+
+
+@pytest.fixture
+def mock_ten_env():
+    env = AsyncMock()
+    env.log = MagicMock()
+    env.log_debug = MagicMock()
+    env.log_info = MagicMock()
+    env.log_warn = MagicMock()
+    env.log_error = MagicMock()
+    return env
+
+
+@pytest.fixture
+def mock_frame():
+    frame = MagicMock()
+    frame.lock_buf.return_value = b"\x00\x01"
+    return frame
+
+
+@pytest.fixture
+def extension(mock_ten_env):
+    ext = BytedanceASRLLMExtension("test_extension")
+    ext.ten_env = mock_ten_env
+    ext.connected = True
+    ext.client = MagicMock()
+    ext.client.send_audio = AsyncMock()
+    ext.send_asr_error = AsyncMock()
+    return ext
+
+
+@pytest.mark.asyncio
+async def test_send_audio_suppresses_stop_phase_disconnect_error(
+    extension, mock_ten_env, mock_frame
+):
+    extension.stopped = True
+    extension.client.send_audio.side_effect = RuntimeError(
+        "Not connected to ASR service"
+    )
+
+    result = await extension.send_audio(mock_frame, "session-1")
+
+    assert result is False
+    extension.send_asr_error.assert_not_awaited()
+    mock_ten_env.log.assert_not_called()
+    mock_ten_env.log_debug.assert_not_called()
+    mock_frame.unlock_buf.assert_called_once_with(b"\x00\x01")
+
+
+@pytest.mark.asyncio
+async def test_send_audio_reports_runtime_disconnect_error(
+    extension, mock_ten_env, mock_frame
+):
+    extension.stopped = False
+    extension.client.send_audio.side_effect = RuntimeError(
+        "Not connected to ASR service"
+    )
+
+    result = await extension.send_audio(mock_frame, "session-1")
+
+    assert result is False
+    extension.send_asr_error.assert_awaited_once()
+    mock_ten_env.log.assert_called_once()
+    mock_frame.unlock_buf.assert_called_once_with(b"\x00\x01")
+
+
+@pytest.mark.asyncio
+async def test_send_audio_does_not_suppress_other_stop_phase_errors(
+    extension, mock_ten_env, mock_frame
+):
+    extension.stopped = True
+    extension.client.send_audio.side_effect = ValueError("unexpected failure")
+
+    result = await extension.send_audio(mock_frame, "session-1")
+
+    assert result is False
+    extension.send_asr_error.assert_not_awaited()
+    mock_ten_env.log.assert_not_called()
+    mock_ten_env.log_debug.assert_not_called()
+    mock_frame.unlock_buf.assert_called_once_with(b"\x00\x01")