From d97516d728603aba0ca8a1b9f8526dec045109c2 Mon Sep 17 00:00:00 2001 From: Justin Tahara <105671973+justin-tahara@users.noreply.github.com> Date: Wed, 23 Sep 2026 18:13:04 +0000 Subject: [PATCH] fix(craft): stop event pumps when a stream ends (#14984) --- .../features/build/session/streaming.py | 5 ++ .../craft/test_session_manager_merger.py | 51 +++++++++++++++++++ 2 files changed, 56 insertions(+) diff --git a/backend/onyx/server/features/build/session/streaming.py b/backend/onyx/server/features/build/session/streaming.py index 7b837cbe24..ef37f5ff81 100644 --- a/backend/onyx/server/features/build/session/streaming.py +++ b/backend/onyx/server/features/build/session/streaming.py @@ -319,10 +319,15 @@ def merge_events_with_announces( ) try: for evt in event_iter: + if stop.is_set(): + break output.put(evt) except Exception as e: output.put(e) finally: + # Close on this thread: it owns the generator, and close() runs the + # subscription teardown (bus unsubscribe, client close). + event_iter.close() output.put(done_sentinel) def drive_announces() -> None: diff --git a/backend/tests/unit/onyx/server/features/craft/test_session_manager_merger.py b/backend/tests/unit/onyx/server/features/craft/test_session_manager_merger.py index e8c4e1bedc..206277b4e3 100644 --- a/backend/tests/unit/onyx/server/features/craft/test_session_manager_merger.py +++ b/backend/tests/unit/onyx/server/features/craft/test_session_manager_merger.py @@ -9,6 +9,7 @@ import json import threading import time from collections.abc import Generator +from queue import Empty from typing import Any from unittest.mock import MagicMock from uuid import UUID, uuid4 @@ -20,6 +21,8 @@ from onyx.server.features.build.packets import ( ApprovalRequestedPacket, ConnectAppRequestPacket, ) +from onyx.server.features.build.sandbox.opencode.event_bus import PodEventBus +from onyx.server.features.build.sandbox.sse import SSEKeepalive from onyx.server.features.build.session import streaming as streaming_mod @@ -325,3 +328,51 @@ def test_pop_announcement_exception_is_swallowed( assert out == ["still-ok"] assert len(calls) >= 2 + + +def test_events_pump_stops_and_unsubscribes_when_consumer_exits( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An attach stream ends on its first terminal chunk. The events pump must + then exit and the bus subscription must be removed.""" + _stub_get_cache_backend(monkeypatch) + _stub_pop_announcement(monkeypatch, _always_none) + monkeypatch.setattr(streaming_mod.connect_app, "pop_announcement", _always_none) + + bus = PodEventBus("http://sandbox.invalid", auth=None) + monkeypatch.setattr(bus, "_ensure_reader_started", lambda: None) + opencode_session_id = "ses_test" + session_id = uuid4() + + def subscription() -> Generator[Any, None, None]: + # Mirrors subscribe_to_opencode_session: never ends on its own and + # yields a keepalive between bus polls. + sub = bus.subscribe(opencode_session_id) + try: + while True: + try: + yield sub.queue.get(timeout=0.05) + except Empty: + yield SSEKeepalive() + finally: + bus.unsubscribe(sub) + + merged = streaming_mod.merge_events_with_announces( + subscription(), session_id=session_id, tenant_id="public" + ) + assert isinstance(next(merged), SSEKeepalive) + assert len(bus._subscribers.get(opencode_session_id, [])) == 1 + merged.close() + + pump_name = f"events-pump-{session_id}" + deadline = time.monotonic() + 2.0 + while time.monotonic() < deadline: + pump_alive = any( + t.name == pump_name and t.is_alive() for t in threading.enumerate() + ) + if not pump_alive and opencode_session_id not in bus._subscribers: + break + time.sleep(0.05) + + assert not any(t.name == pump_name for t in threading.enumerate()) + assert opencode_session_id not in bus._subscribers