From 2f15234c451ddc031ef93cbbffcaeb3d06b27121 Mon Sep 17 00:00:00 2001 From: Ziming Wang Date: Thu, 3 Sep 2026 07:17:52 -0700 Subject: [PATCH 1/4] Collect model input after cadence pacing Signed-off-by: Ziming Wang --- flashdreams/flashdreams/api_v2/loop.py | 12 ++++++-- flashdreams/test_v2/test_session_runner.py | 35 ++++++++++++++++++++++ 2 files changed, 44 insertions(+), 3 deletions(-) diff --git a/flashdreams/flashdreams/api_v2/loop.py b/flashdreams/flashdreams/api_v2/loop.py index 491703b0b..93a691a0c 100644 --- a/flashdreams/flashdreams/api_v2/loop.py +++ b/flashdreams/flashdreams/api_v2/loop.py @@ -236,13 +236,19 @@ def _run_model_loop( while not self._shutdown_event.is_set() and ( max_steps is None or steps_run < max_steps ): - events, generation = event_buffer.read(reader_id) - step_index = self._begin_run(events, generation) - if step_index is None: + if ( + self.is_finished() + and event_buffer.generation == self._generation + and self._message_queue.empty() + ): break last_run_started = self._pace(last_run_started) if self._shutdown_event.is_set(): break + events, generation = event_buffer.read(reader_id) + step_index = self._begin_run(events, generation) + if step_index is None: + break step_started_at = time.monotonic() raw_result = self.step(step_index, events) step_elapsed_s = time.monotonic() - step_started_at diff --git a/flashdreams/test_v2/test_session_runner.py b/flashdreams/test_v2/test_session_runner.py index 4f87af25d..8e2214a08 100644 --- a/flashdreams/test_v2/test_session_runner.py +++ b/flashdreams/test_v2/test_session_runner.py @@ -535,6 +535,41 @@ def _lifecycle_event(event_type: type[UserInputEvent]) -> UserInputEvents: return UserInputEvents([event_type(timestamp=uint64(0))]) +def test_model_loop_collects_input_after_pacing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session = FakeSession(_session_desc(), CallLog()) + session.init() + event_buffer = EventBuffer() + event_buffer.register(0) + pace_calls = 0 + + def pace(last_run_started: float | None) -> float: + nonlocal pace_calls + del last_run_started + if pace_calls == 1: + event_buffer.append(_key_event()) + pace_calls += 1 + return float(pace_calls) + + monkeypatch.setattr(session.model_loop, "_pace", pace) + session.model_loop._run_model_loop( + event_buffer=event_buffer, + reader_id=0, + publish=lambda generation, results, elapsed: None, + max_steps=2, + ) + + assert session.observed_events[0].get_events() == [] + observed = session.observed_events[1].get_events() + assert len(observed) == 1 + assert isinstance(observed[0], KeyboardUserInputEvent) + assert (observed[0].key, observed[0].state) == ( + "a", + KeyboardInputState.PRESSED, + ) + + def test_run_session_presents_every_step_in_order() -> None: log = CallLog() session = FakeSession(_session_desc(), log) From 9c41be1b6ece2e5a67572697840b2bbb85a9e854 Mon Sep 17 00:00:00 2001 From: Ziming Wang Date: Thu, 3 Sep 2026 07:34:14 -0700 Subject: [PATCH 2/4] Make finished-loop shutdown atomic Signed-off-by: Ziming Wang --- flashdreams/flashdreams/api_v2/loop.py | 18 ++-- flashdreams/test_v2/test_session_runner.py | 98 ++++++++++++++++++++++ 2 files changed, 111 insertions(+), 5 deletions(-) diff --git a/flashdreams/flashdreams/api_v2/loop.py b/flashdreams/flashdreams/api_v2/loop.py index 93a691a0c..2a8593732 100644 --- a/flashdreams/flashdreams/api_v2/loop.py +++ b/flashdreams/flashdreams/api_v2/loop.py @@ -212,6 +212,18 @@ class IModelLoop(ILoop[StateT], ABC): bare :class:`StepResult` or ``None`` raises :class:`TypeError`. """ + def _claim_finished_shutdown(self, generation: int) -> bool: + """Atomically stop message acceptance for an idle finished loop.""" + with self._lifecycle_lock: + if ( + not self.is_finished() + or generation != self._generation + or not self._message_queue.empty() + ): + return False + self._accepting_messages = False + return True + @final def _run_model_loop( self, @@ -236,11 +248,7 @@ def _run_model_loop( while not self._shutdown_event.is_set() and ( max_steps is None or steps_run < max_steps ): - if ( - self.is_finished() - and event_buffer.generation == self._generation - and self._message_queue.empty() - ): + if self._claim_finished_shutdown(event_buffer.generation): break last_run_started = self._pace(last_run_started) if self._shutdown_event.is_set(): diff --git a/flashdreams/test_v2/test_session_runner.py b/flashdreams/test_v2/test_session_runner.py index 8e2214a08..5aeea6281 100644 --- a/flashdreams/test_v2/test_session_runner.py +++ b/flashdreams/test_v2/test_session_runner.py @@ -570,6 +570,104 @@ def pace(last_run_started: float | None) -> float: ) +def test_finished_model_loop_executes_an_accepted_message() -> None: + session = FiniteSession(_session_desc(), CallLog(), length=1, generated=1) + session.init() + event_buffer = EventBuffer() + event_buffer.register(0) + + invoke_async(session.model_loop, lambda state: state.reset()) + session.model_loop._run_model_loop( + event_buffer=event_buffer, + reader_id=0, + publish=lambda generation, results, elapsed: None, + max_steps=1, + ) + + assert session._generated == 1 + assert len(session.observed_events) == 1 + + +def test_finished_model_loop_drops_a_self_message_at_shutdown() -> None: + session = FiniteSession(_session_desc(), CallLog(), length=1, generated=1) + session.init() + event_buffer = EventBuffer() + event_buffer.register(0) + calls = 0 + + def self_message(state: FakeSession) -> None: + nonlocal calls + del state + calls += 1 + invoke_async(session.model_loop, self_message) + + invoke_async(session.model_loop, self_message) + session.model_loop._run_model_loop( + event_buffer=event_buffer, + reader_id=0, + publish=lambda generation, results, elapsed: None, + ) + + assert calls == 1 + assert session.model_loop._message_queue.empty() + assert session._failure_queue.empty() + + +def test_finished_model_loop_rejects_a_message_after_shutdown_is_claimed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session = FiniteSession(_session_desc(), CallLog(), length=1, generated=1) + session.init() + event_buffer = EventBuffer() + event_buffer.register(0) + queue_checked = threading.Event() + finish_check = threading.Event() + original_empty = session.model_loop._message_queue.empty + + def synchronized_empty() -> bool: + empty = original_empty() + queue_checked.set() + assert finish_check.wait(timeout=1.0) + return empty + + monkeypatch.setattr(session.model_loop._message_queue, "empty", synchronized_empty) + model_thread = threading.Thread( + target=session.model_loop._run_model_loop, + kwargs={ + "event_buffer": event_buffer, + "reader_id": 0, + "publish": lambda generation, results, elapsed: None, + }, + ) + model_thread.start() + assert queue_checked.wait(timeout=1.0) + + errors: list[BaseException] = [] + sender_attempted = threading.Event() + + def send_message() -> None: + sender_attempted.set() + try: + invoke_async(session.model_loop, lambda state: state.reset()) + except BaseException as error: + errors.append(error) + + sender_thread = threading.Thread(target=send_message) + sender_thread.start() + assert sender_attempted.wait(timeout=1.0) + finish_check.set() + model_thread.join(timeout=1.0) + sender_thread.join(timeout=1.0) + + assert not model_thread.is_alive() + assert not sender_thread.is_alive() + assert len(errors) == 1 + assert isinstance(errors[0], RuntimeError) + assert str(errors[0]) == "Loop is shutting down." + assert session._generated == 1 + assert session._failure_queue.empty() + + def test_run_session_presents_every_step_in_order() -> None: log = CallLog() session = FakeSession(_session_desc(), log) From 91f6f94c106fd2f2573ab05050b49bec9b213be3 Mon Sep 17 00:00:00 2001 From: Ziming Wang Date: Thu, 3 Sep 2026 07:48:02 -0700 Subject: [PATCH 3/4] Settle terminal messages before shutdown Signed-off-by: Ziming Wang --- flashdreams/flashdreams/api_v2/loop.py | 43 +++++++--- .../flashdreams/runtime_v2/event_buffer.py | 8 ++ flashdreams/test_v2/test_session_runner.py | 80 +++++++++++++++++-- 3 files changed, 112 insertions(+), 19 deletions(-) diff --git a/flashdreams/flashdreams/api_v2/loop.py b/flashdreams/flashdreams/api_v2/loop.py index 2a8593732..7cf34b143 100644 --- a/flashdreams/flashdreams/api_v2/loop.py +++ b/flashdreams/flashdreams/api_v2/loop.py @@ -184,6 +184,9 @@ def _run_message_batch(self) -> None: batch.append(self._message_queue.get_nowait()) except queue.Empty: break + self._run_messages(batch) + + def _run_messages(self, batch: list[_Message[StateT]]) -> None: for message in batch: result = message.operation(self.state) if result is not None: @@ -212,18 +215,30 @@ class IModelLoop(ILoop[StateT], ABC): bare :class:`StepResult` or ``None`` raises :class:`TypeError`. """ - def _claim_finished_shutdown(self, generation: int) -> bool: - """Atomically stop message acceptance for an idle finished loop.""" - with self._lifecycle_lock: - if ( - not self.is_finished() - or generation != self._generation - or not self._message_queue.empty() - ): - return False - self._accepting_messages = False + def _settle_finished_state(self, event_buffer: EventBuffer) -> bool: + """Process the terminal message batch and return whether to stop.""" + with event_buffer.hold_generation() as generation: + with self._lifecycle_lock: + if not self.is_finished() or generation != self._generation: + return False + self._accepting_messages = False + batch: list[_Message[StateT]] = [] + while True: + try: + batch.append(self._message_queue.get_nowait()) + except queue.Empty: + break + + self._run_messages(batch) + if self.is_finished(): return True + with self._lifecycle_lock: + if self._closed or self._shutdown_event.is_set(): + return True + self._accepting_messages = True + return False + @final def _run_model_loop( self, @@ -248,7 +263,7 @@ def _run_model_loop( while not self._shutdown_event.is_set() and ( max_steps is None or steps_run < max_steps ): - if self._claim_finished_shutdown(event_buffer.generation): + if self._settle_finished_state(event_buffer): break last_run_started = self._pace(last_run_started) if self._shutdown_event.is_set(): @@ -256,7 +271,11 @@ def _run_model_loop( events, generation = event_buffer.read(reader_id) step_index = self._begin_run(events, generation) if step_index is None: - break + if self._shutdown_event.is_set() or self._settle_finished_state( + event_buffer + ): + break + continue step_started_at = time.monotonic() raw_result = self.step(step_index, events) step_elapsed_s = time.monotonic() - step_started_at diff --git a/flashdreams/flashdreams/runtime_v2/event_buffer.py b/flashdreams/flashdreams/runtime_v2/event_buffer.py index 66265a567..cefc18fd9 100644 --- a/flashdreams/flashdreams/runtime_v2/event_buffer.py +++ b/flashdreams/flashdreams/runtime_v2/event_buffer.py @@ -4,6 +4,8 @@ """Input events shared by the model and UI loops.""" import threading +from collections.abc import Iterator +from contextlib import contextmanager from flashdreams.runtime_v2.user_input_event import ( ResetUserInputEvent, @@ -39,6 +41,12 @@ def generation(self) -> int: with self._lock: return self._generation + @contextmanager + def hold_generation(self) -> Iterator[int]: + """Yield the current generation while input appends are paused.""" + with self._lock: + yield self._generation + def register(self, reader_id: int) -> None: """Register a reader before input collection begins.""" with self._lock: diff --git a/flashdreams/test_v2/test_session_runner.py b/flashdreams/test_v2/test_session_runner.py index 5aeea6281..74fc67413 100644 --- a/flashdreams/test_v2/test_session_runner.py +++ b/flashdreams/test_v2/test_session_runner.py @@ -588,7 +588,7 @@ def test_finished_model_loop_executes_an_accepted_message() -> None: assert len(session.observed_events) == 1 -def test_finished_model_loop_drops_a_self_message_at_shutdown() -> None: +def test_finished_model_loop_rejects_a_self_message_at_shutdown_cutoff() -> None: session = FiniteSession(_session_desc(), CallLog(), length=1, generated=1) session.init() event_buffer = EventBuffer() @@ -599,7 +599,8 @@ def self_message(state: FakeSession) -> None: nonlocal calls del state calls += 1 - invoke_async(session.model_loop, self_message) + with pytest.raises(RuntimeError, match="Loop is shutting down"): + invoke_async(session.model_loop, self_message) invoke_async(session.model_loop, self_message) session.model_loop._run_model_loop( @@ -613,6 +614,70 @@ def self_message(state: FakeSession) -> None: assert session._failure_queue.empty() +def test_finished_model_loop_executes_a_message_accepted_after_the_main_batch() -> None: + session = FiniteSession(_session_desc(), CallLog(), length=1) + session.init() + event_buffer = EventBuffer() + event_buffer.register(0) + + def finish_and_reset(state: FakeSession) -> None: + assert isinstance(state, FiniteSession) + state._generated = 1 + invoke_async(session.model_loop, lambda owner: owner.reset()) + + invoke_async(session.model_loop, finish_and_reset) + session.model_loop._run_model_loop( + event_buffer=event_buffer, + reader_id=0, + publish=lambda generation, results, elapsed: None, + max_steps=1, + ) + + assert session._generated == 1 + assert len(session.observed_events) == 1 + assert session._failure_queue.empty() + + +def test_finished_model_loop_observes_a_reset_appended_after_its_read() -> None: + class ResetAfterReadEventBuffer(EventBuffer): + def __init__(self) -> None: + super().__init__() + self._reset_appended = False + + def read(self, reader_id: int) -> tuple[UserInputEvents, int]: + events, generation = super().read(reader_id) + if not self._reset_appended: + self._reset_appended = True + self.append(_lifecycle_event(ResetUserInputEvent)) + return events, generation + + session = FiniteSession(_session_desc(), CallLog(), length=1) + session.init() + event_buffer = ResetAfterReadEventBuffer() + event_buffer.register(0) + + def finish(state: FakeSession) -> None: + assert isinstance(state, FiniteSession) + state._generated = 1 + + invoke_async(session.model_loop, finish) + session.model_loop._run_model_loop( + event_buffer=event_buffer, + reader_id=0, + publish=lambda generation, results, elapsed: None, + max_steps=1, + ) + + assert session.model_loop._generation == 1 + assert session._generated == 1 + assert len(session.observed_events) == 1 + assert any( + isinstance(event, ResetUserInputEvent) + for event in session.observed_events[0].get_events() + ) + assert session._failure_queue.empty() + + def test_finished_model_loop_rejects_a_message_after_shutdown_is_claimed( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -622,15 +687,16 @@ def test_finished_model_loop_rejects_a_message_after_shutdown_is_claimed( event_buffer.register(0) queue_checked = threading.Event() finish_check = threading.Event() - original_empty = session.model_loop._message_queue.empty + original_get_nowait = session.model_loop._message_queue.get_nowait - def synchronized_empty() -> bool: - empty = original_empty() + def synchronized_get_nowait() -> object: queue_checked.set() assert finish_check.wait(timeout=1.0) - return empty + return original_get_nowait() - monkeypatch.setattr(session.model_loop._message_queue, "empty", synchronized_empty) + monkeypatch.setattr( + session.model_loop._message_queue, "get_nowait", synchronized_get_nowait + ) model_thread = threading.Thread( target=session.model_loop._run_model_loop, kwargs={ From f7913e96a24b7868d3a3570f445290870658b159 Mon Sep 17 00:00:00 2001 From: Ziming Wang Date: Thu, 3 Sep 2026 10:41:59 -0700 Subject: [PATCH 4/4] Narrow cadence fix to model input ordering Keep event collection immediately after model pacing. Restore the terminal lifecycle path for separate follow-up. Signed-off-by: Ziming Wang --- flashdreams/flashdreams/api_v2/loop.py | 35 +--- .../flashdreams/runtime_v2/event_buffer.py | 8 - flashdreams/test_v2/test_session_runner.py | 164 ------------------ 3 files changed, 1 insertion(+), 206 deletions(-) diff --git a/flashdreams/flashdreams/api_v2/loop.py b/flashdreams/flashdreams/api_v2/loop.py index 7cf34b143..f01c2ad68 100644 --- a/flashdreams/flashdreams/api_v2/loop.py +++ b/flashdreams/flashdreams/api_v2/loop.py @@ -184,9 +184,6 @@ def _run_message_batch(self) -> None: batch.append(self._message_queue.get_nowait()) except queue.Empty: break - self._run_messages(batch) - - def _run_messages(self, batch: list[_Message[StateT]]) -> None: for message in batch: result = message.operation(self.state) if result is not None: @@ -215,30 +212,6 @@ class IModelLoop(ILoop[StateT], ABC): bare :class:`StepResult` or ``None`` raises :class:`TypeError`. """ - def _settle_finished_state(self, event_buffer: EventBuffer) -> bool: - """Process the terminal message batch and return whether to stop.""" - with event_buffer.hold_generation() as generation: - with self._lifecycle_lock: - if not self.is_finished() or generation != self._generation: - return False - self._accepting_messages = False - batch: list[_Message[StateT]] = [] - while True: - try: - batch.append(self._message_queue.get_nowait()) - except queue.Empty: - break - - self._run_messages(batch) - if self.is_finished(): - return True - - with self._lifecycle_lock: - if self._closed or self._shutdown_event.is_set(): - return True - self._accepting_messages = True - return False - @final def _run_model_loop( self, @@ -263,19 +236,13 @@ def _run_model_loop( while not self._shutdown_event.is_set() and ( max_steps is None or steps_run < max_steps ): - if self._settle_finished_state(event_buffer): - break last_run_started = self._pace(last_run_started) if self._shutdown_event.is_set(): break events, generation = event_buffer.read(reader_id) step_index = self._begin_run(events, generation) if step_index is None: - if self._shutdown_event.is_set() or self._settle_finished_state( - event_buffer - ): - break - continue + break step_started_at = time.monotonic() raw_result = self.step(step_index, events) step_elapsed_s = time.monotonic() - step_started_at diff --git a/flashdreams/flashdreams/runtime_v2/event_buffer.py b/flashdreams/flashdreams/runtime_v2/event_buffer.py index cefc18fd9..66265a567 100644 --- a/flashdreams/flashdreams/runtime_v2/event_buffer.py +++ b/flashdreams/flashdreams/runtime_v2/event_buffer.py @@ -4,8 +4,6 @@ """Input events shared by the model and UI loops.""" import threading -from collections.abc import Iterator -from contextlib import contextmanager from flashdreams.runtime_v2.user_input_event import ( ResetUserInputEvent, @@ -41,12 +39,6 @@ def generation(self) -> int: with self._lock: return self._generation - @contextmanager - def hold_generation(self) -> Iterator[int]: - """Yield the current generation while input appends are paused.""" - with self._lock: - yield self._generation - def register(self, reader_id: int) -> None: """Register a reader before input collection begins.""" with self._lock: diff --git a/flashdreams/test_v2/test_session_runner.py b/flashdreams/test_v2/test_session_runner.py index 74fc67413..8e2214a08 100644 --- a/flashdreams/test_v2/test_session_runner.py +++ b/flashdreams/test_v2/test_session_runner.py @@ -570,170 +570,6 @@ def pace(last_run_started: float | None) -> float: ) -def test_finished_model_loop_executes_an_accepted_message() -> None: - session = FiniteSession(_session_desc(), CallLog(), length=1, generated=1) - session.init() - event_buffer = EventBuffer() - event_buffer.register(0) - - invoke_async(session.model_loop, lambda state: state.reset()) - session.model_loop._run_model_loop( - event_buffer=event_buffer, - reader_id=0, - publish=lambda generation, results, elapsed: None, - max_steps=1, - ) - - assert session._generated == 1 - assert len(session.observed_events) == 1 - - -def test_finished_model_loop_rejects_a_self_message_at_shutdown_cutoff() -> None: - session = FiniteSession(_session_desc(), CallLog(), length=1, generated=1) - session.init() - event_buffer = EventBuffer() - event_buffer.register(0) - calls = 0 - - def self_message(state: FakeSession) -> None: - nonlocal calls - del state - calls += 1 - with pytest.raises(RuntimeError, match="Loop is shutting down"): - invoke_async(session.model_loop, self_message) - - invoke_async(session.model_loop, self_message) - session.model_loop._run_model_loop( - event_buffer=event_buffer, - reader_id=0, - publish=lambda generation, results, elapsed: None, - ) - - assert calls == 1 - assert session.model_loop._message_queue.empty() - assert session._failure_queue.empty() - - -def test_finished_model_loop_executes_a_message_accepted_after_the_main_batch() -> None: - session = FiniteSession(_session_desc(), CallLog(), length=1) - session.init() - event_buffer = EventBuffer() - event_buffer.register(0) - - def finish_and_reset(state: FakeSession) -> None: - assert isinstance(state, FiniteSession) - state._generated = 1 - invoke_async(session.model_loop, lambda owner: owner.reset()) - - invoke_async(session.model_loop, finish_and_reset) - session.model_loop._run_model_loop( - event_buffer=event_buffer, - reader_id=0, - publish=lambda generation, results, elapsed: None, - max_steps=1, - ) - - assert session._generated == 1 - assert len(session.observed_events) == 1 - assert session._failure_queue.empty() - - -def test_finished_model_loop_observes_a_reset_appended_after_its_read() -> None: - class ResetAfterReadEventBuffer(EventBuffer): - def __init__(self) -> None: - super().__init__() - self._reset_appended = False - - def read(self, reader_id: int) -> tuple[UserInputEvents, int]: - events, generation = super().read(reader_id) - if not self._reset_appended: - self._reset_appended = True - self.append(_lifecycle_event(ResetUserInputEvent)) - return events, generation - - session = FiniteSession(_session_desc(), CallLog(), length=1) - session.init() - event_buffer = ResetAfterReadEventBuffer() - event_buffer.register(0) - - def finish(state: FakeSession) -> None: - assert isinstance(state, FiniteSession) - state._generated = 1 - - invoke_async(session.model_loop, finish) - session.model_loop._run_model_loop( - event_buffer=event_buffer, - reader_id=0, - publish=lambda generation, results, elapsed: None, - max_steps=1, - ) - - assert session.model_loop._generation == 1 - assert session._generated == 1 - assert len(session.observed_events) == 1 - assert any( - isinstance(event, ResetUserInputEvent) - for event in session.observed_events[0].get_events() - ) - assert session._failure_queue.empty() - - -def test_finished_model_loop_rejects_a_message_after_shutdown_is_claimed( - monkeypatch: pytest.MonkeyPatch, -) -> None: - session = FiniteSession(_session_desc(), CallLog(), length=1, generated=1) - session.init() - event_buffer = EventBuffer() - event_buffer.register(0) - queue_checked = threading.Event() - finish_check = threading.Event() - original_get_nowait = session.model_loop._message_queue.get_nowait - - def synchronized_get_nowait() -> object: - queue_checked.set() - assert finish_check.wait(timeout=1.0) - return original_get_nowait() - - monkeypatch.setattr( - session.model_loop._message_queue, "get_nowait", synchronized_get_nowait - ) - model_thread = threading.Thread( - target=session.model_loop._run_model_loop, - kwargs={ - "event_buffer": event_buffer, - "reader_id": 0, - "publish": lambda generation, results, elapsed: None, - }, - ) - model_thread.start() - assert queue_checked.wait(timeout=1.0) - - errors: list[BaseException] = [] - sender_attempted = threading.Event() - - def send_message() -> None: - sender_attempted.set() - try: - invoke_async(session.model_loop, lambda state: state.reset()) - except BaseException as error: - errors.append(error) - - sender_thread = threading.Thread(target=send_message) - sender_thread.start() - assert sender_attempted.wait(timeout=1.0) - finish_check.set() - model_thread.join(timeout=1.0) - sender_thread.join(timeout=1.0) - - assert not model_thread.is_alive() - assert not sender_thread.is_alive() - assert len(errors) == 1 - assert isinstance(errors[0], RuntimeError) - assert str(errors[0]) == "Loop is shutting down." - assert session._generated == 1 - assert session._failure_queue.empty() - - def test_run_session_presents_every_step_in_order() -> None: log = CallLog() session = FakeSession(_session_desc(), log)