diff --git a/tests/tools/test_oneshot_completion_linger.py b/tests/tools/test_oneshot_completion_linger.py index e9c9fb5b54..b253105196 100644 --- a/tests/tools/test_oneshot_completion_linger.py +++ b/tests/tools/test_oneshot_completion_linger.py @@ -88,6 +88,24 @@ def test_already_exited_session_not_waited_on(registry): assert result["waited"] == [] +def test_session_mid_finish_is_still_waited_on(registry): + """``_move_to_finished`` moves a session out of ``_running`` BEFORE it enqueues the completion and + sets the event. A parent whose turn ends inside that window must still linger for the event, or + it drains nothing and exits without the follow-up turn (the CI-flaky quiet-notify resume).""" + s = _make_session(exited=True) # exited, moved, but completion not yet published + with registry._lock: + registry._finished[s.id] = s + + def _publish(): + time.sleep(0.3) + s._completion_event.set() + + threading.Thread(target=_publish, daemon=True).start() + result = registry.wait_for_pending_completions(timeout=30, poll_interval=0.1) + assert result["waited"] == [s.id] + assert result["completed"] == [s.id] + + def test_wait_returns_when_process_completes(registry): s = _make_session() with registry._lock: diff --git a/tools/process_registry.py b/tools/process_registry.py index eebb59e4e8..8300c745dc 100644 --- a/tools/process_registry.py +++ b/tools/process_registry.py @@ -1438,8 +1438,11 @@ class ProcessRegistry(ProcessCheckpointMixin): timeout = self._oneshot_completion_wait_seconds() result: dict = {"waited": [], "completed": [], "timed_out": []} with self._lock: + # `_finished` too: `_move_to_finished` pops a session from `_running` and enqueues its completion + # only after releasing handles and writing the checkpoint. A parent whose turn ends inside that + # window would otherwise see nothing pending, drain nothing and exit without the follow-up turn. pending = [ - s for s in self._running.values() + s for store in (self._running, self._finished) for s in store.values() if s.notify_on_complete and not s._completion_event.is_set() and (task_id is None or s.task_id == task_id) ] if not pending or timeout <= 0: