From 2ea6a16479a6d782fa7b119d1c18b1e45c151cfc Mon Sep 17 00:00:00 2001
From: "duanjialing.777" <duanjialing.777@bytedance.com>
Date: Tue, 15 Sep 2026 14:35:49 +0800
Subject: [PATCH] perf(chat): wait on turn completion notifications

Signed-off-by: duanjialing.777 <duanjialing.777@bytedance.com>
---
 loopx/chat_runtime.py        | 12 +++++++-----
 tests/test_chat_turn_wait.py | 37 ++++++++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 5 deletions(-)
 create mode 100644 tests/test_chat_turn_wait.py

diff --git a/loopx/chat_runtime.py b/loopx/chat_runtime.py
index 756643ff48..fa17ad829f 100644
--- a/loopx/chat_runtime.py
+++ b/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):
diff --git a/tests/test_chat_turn_wait.py b/tests/test_chat_turn_wait.py
new file mode 100644
index 0000000000..b527a03316
--- /dev/null
+++ b/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