fix(tools): one-shot linger waits for a completion that is mid-publish

`ProcessRegistry._move_to_finished` pops the session out of `_running`, then
saves the receipt, releases handles and writes the checkpoint, and only THEN
enqueues the completion and sets `_completion_event`. A quiet one-shot parent
whose turn ends inside that window called `wait_for_pending_completions`,
found nothing in `_running`, drained an empty queue and exited without the
follow-up turn. That is the CI flake in
tests/tools/test_completed_process_results.py::test_headless_terminal_result_survives_cli_exit
(`follow_ups == []`), which also hit unrelated branches.

Consider `_finished` sessions whose event is not yet set as pending too.

Repro: a 6s sleep before the enqueue plus a 2s delay before the parent's first
wait fails the E2E 2/2 on main and passes 2/2 with this change.
This commit is contained in:
teknium1
2026-09-14 18:47:07 -07:00
committed by Teknium
parent dfc28b61a0
commit f5a457ad5b
2 changed files with 22 additions and 1 deletions

View File

@@ -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:

View File

@@ -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: