perf(chat): wait on turn completion notifications

Signed-off-by: duanjialing.777 <duanjialing.777@bytedance.com>

Author: Duang777

loopx/chat_runtime.py

@@ -1398,14 +1398,16 @@ def interrupt_turn(self, *, session_id: str, turn_id: str) -> dict[str, Any]:
 
     def wait_for_turn(self, *, session_id: str, turn_id: str, timeout_sec: float = 920.0) -> dict[str, Any]:
         deadline = time.monotonic() + timeout_sec
-        while time.monotonic() < deadline:
-            turn = self.store.load_turn(session_id, turn_id)
-            if turn is None:
+        while True:
+            if (turn := self.store.load_turn(session_id, turn_id)) is None:
                 raise KeyError("chat turn was not found")
             if turn.get("status") in TERMINAL_TURN_STATES:
                 return turn
-            time.sleep(0.02)
-        raise TimeoutError("chat turn wait timed out")
+            if (remaining := deadline - time.monotonic()) <= 0:
+                raise TimeoutError("chat turn wait timed out")
+            with self.lock:
+                done_event = self.turn_done_events.get((session_id, turn_id))
+            (done_event.wait if done_event else time.sleep)(remaining if done_event else min(0.02, remaining))
 
     def close_session(self, session_id: str) -> bool:
         with self._session_adapter_lock(session_id):

tests/test_chat_turn_wait.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+import threading
+from unittest.mock import Mock
+
+from loopx.chat_runtime import ChatRuntimeController
+
+
+def _runtime() -> ChatRuntimeController:
+    runtime = ChatRuntimeController.__new__(ChatRuntimeController)
+    runtime.store = Mock()  # type: ignore[assignment]
+    runtime.lock = threading.RLock()
+    runtime.turn_done_events = {}
+    return runtime
+
+
+def test_wait_for_turn_uses_managed_completion_event() -> None:
+    runtime = _runtime()
+    runtime.store.load_turn.side_effect = [{"status": "running"}, {"status": "completed"}]
+    completion = Mock()
+    runtime.turn_done_events[("session", "turn")] = completion  # type: ignore[assignment]
+
+    turn = runtime.wait_for_turn(session_id="session", turn_id="turn", timeout_sec=0.1)
+
+    assert turn["status"] == "completed"
+    completion.wait.assert_called_once()
+    assert runtime.store.load_turn.call_count == 2
+
+
+def test_wait_for_turn_performs_final_fallback_read_at_deadline() -> None:
+    runtime = _runtime()
+    runtime.store.load_turn.side_effect = [{"status": "running"}, {"status": "completed"}]
+
+    turn = runtime.wait_for_turn(session_id="session", turn_id="turn", timeout_sec=0.001)
+
+    assert turn["status"] == "completed"
+    assert runtime.store.load_turn.call_count == 2