From fb3ffa46b6631bb79219a5c4caaab749ad2d6c49 Mon Sep 17 00:00:00 2001 From: Alessandro Pogliaghi Date: Thu, 30 Apr 2026 08:49:31 +0100 Subject: [PATCH 1/3] fix(cloud-agent): handle stream churn --- .../activities/relay_sandbox_events.py | 38 ++++++- .../tests/test_relay_sandbox_events.py | 98 +++++++++++++++++++ 2 files changed, 133 insertions(+), 3 deletions(-) diff --git a/products/tasks/backend/temporal/process_task/activities/relay_sandbox_events.py b/products/tasks/backend/temporal/process_task/activities/relay_sandbox_events.py index 1a73051801c1..f0d99183709d 100644 --- a/products/tasks/backend/temporal/process_task/activities/relay_sandbox_events.py +++ b/products/tasks/backend/temporal/process_task/activities/relay_sandbox_events.py @@ -256,7 +256,6 @@ async def _relay_loop( params=params, ) as event_source: event_source.response.raise_for_status() - reconnect_count = 0 # Reset on successful connection last_event_time[0] = time.monotonic() async for sse_event in event_source.aiter_sse(): @@ -274,6 +273,8 @@ async def _relay_loop( continue await redis_stream.write_event(event_data) + if not _is_keepalive_event(event_data): + reconnect_count = 0 last_event_time[0] = time.monotonic() if _is_end_of_turn(event_data): @@ -302,8 +303,17 @@ async def _relay_loop( return # SSE stream ended normally (sandbox closed connection) - await redis_stream.mark_complete() - return + if await _mark_complete_if_run_is_terminal(redis_stream, run_id): + logger.info("relay_sandbox_events_stopped_after_terminal_run", run_id=run_id) + return + + reconnect_count += 1 + logger.warning( + "relay_sandbox_events_stream_closed_before_terminal_run", + run_id=run_id, + reconnect_count=reconnect_count, + ) + await asyncio.sleep(min(reconnect_count * 2, 10)) except httpx.ReadTimeout: reconnect_count += 1 @@ -337,6 +347,24 @@ async def _relay_loop( pass +async def _mark_complete_if_run_is_terminal(redis_stream: TaskRunRedisStream, run_id: str) -> bool: + try: + task_run = await TaskRunModel.objects.only("status").aget(id=run_id) + except TaskRunModel.DoesNotExist: + await redis_stream.mark_error("Task run not found") + return True + + if task_run.status in ( + TaskRunModel.Status.COMPLETED, + TaskRunModel.Status.FAILED, + TaskRunModel.Status.CANCELLED, + ): + await redis_stream.mark_complete() + return True + + return False + + def _is_session_update(event_data: dict) -> bool: """Check if an event is a session/update notification (active agent processing).""" if event_data.get("type") != "notification": @@ -345,6 +373,10 @@ def _is_session_update(event_data: dict) -> bool: return notification.get("method") == "session/update" +def _is_keepalive_event(event_data: dict) -> bool: + return event_data.get("type") == "keepalive" + + _is_end_of_turn = is_turn_complete diff --git a/products/tasks/backend/temporal/process_task/activities/tests/test_relay_sandbox_events.py b/products/tasks/backend/temporal/process_task/activities/tests/test_relay_sandbox_events.py index 52fbea982512..3131034c42cf 100644 --- a/products/tasks/backend/temporal/process_task/activities/tests/test_relay_sandbox_events.py +++ b/products/tasks/backend/temporal/process_task/activities/tests/test_relay_sandbox_events.py @@ -17,6 +17,7 @@ TaskRunRedisStream, _is_end_of_turn, _is_session_update, + _mark_complete_if_run_is_terminal, _mark_error_unless_run_is_terminal, _relay_loop, relay_sandbox_events, @@ -333,6 +334,103 @@ async def aget(self, id: str) -> SimpleNamespace: redis_stream_mock.mark_complete.assert_awaited_once() redis_stream_mock.mark_error.assert_not_awaited() + async def test_terminal_run_marks_stream_complete_on_normal_stream_close( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + redis_stream_mock = SimpleNamespace(mark_complete=AsyncMock(), mark_error=AsyncMock()) + redis_stream = cast(TaskRunRedisStream, redis_stream_mock) + + class StubTaskRunQuerySet: + def only(self, *_fields: str) -> "StubTaskRunQuerySet": + return self + + async def aget(self, id: str) -> SimpleNamespace: + return SimpleNamespace(status="completed") + + monkeypatch.setattr( + relay_sandbox_events_module, + "TaskRunModel", + SimpleNamespace( + Status=SimpleNamespace(COMPLETED="completed", FAILED="failed", CANCELLED="cancelled"), + DoesNotExist=Exception, + objects=StubTaskRunQuerySet(), + ), + ) + + marked_complete = await _mark_complete_if_run_is_terminal(redis_stream, "run-id") + + assert marked_complete is True + redis_stream_mock.mark_complete.assert_awaited_once() + redis_stream_mock.mark_error.assert_not_awaited() + + async def test_in_progress_normal_stream_close_reconnects_until_exhausted( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + redis_stream = SimpleNamespace( + write_event=AsyncMock(), + mark_complete=AsyncMock(), + mark_error=AsyncMock(), + ) + sleep_mock = AsyncMock() + connect_attempts = 0 + + class EmptyEventSource: + response = SimpleNamespace(raise_for_status=lambda: None) + + async def __aenter__(self) -> "EmptyEventSource": + return self + + async def __aexit__(self, *_args: object) -> None: + return None + + async def aiter_sse(self): + if False: + yield SimpleNamespace(data="") + + def fake_connect_sse(*_args: object, **_kwargs: object) -> EmptyEventSource: + nonlocal connect_attempts + connect_attempts += 1 + return EmptyEventSource() + + async def fake_background_heartbeat(*_args: object, **_kwargs: object) -> None: + return None + + class StubTaskRunQuerySet: + def only(self, *_fields: str) -> "StubTaskRunQuerySet": + return self + + async def aget(self, id: str) -> SimpleNamespace: + return SimpleNamespace(status="in_progress") + + monkeypatch.setattr(relay_sandbox_events_module.httpx_sse, "aconnect_sse", fake_connect_sse) + monkeypatch.setattr(relay_sandbox_events_module.asyncio, "sleep", sleep_mock) + monkeypatch.setattr(relay_sandbox_events_module, "_background_heartbeat", fake_background_heartbeat) + monkeypatch.setattr( + relay_sandbox_events_module, + "TaskRunModel", + SimpleNamespace( + Status=SimpleNamespace(COMPLETED="completed", FAILED="failed", CANCELLED="cancelled"), + DoesNotExist=Exception, + objects=StubTaskRunQuerySet(), + ), + ) + + await _relay_loop( + events_url="https://sandbox.example/events", + headers={"Authorization": "Bearer token"}, + params={}, + redis_stream=cast(TaskRunRedisStream, redis_stream), + run_id="run-id", + task_id="task-id", + ) + + assert connect_attempts == relay_sandbox_events_module.MAX_RECONNECT_ATTEMPTS + 1 + redis_stream.write_event.assert_not_awaited() + redis_stream.mark_complete.assert_not_awaited() + redis_stream.mark_error.assert_awaited_once_with( + f"Lost connection to sandbox after {relay_sandbox_events_module.MAX_RECONNECT_ATTEMPTS} reconnection attempts" + ) + async def test_in_progress_run_marks_stream_error_on_relay_error(self, monkeypatch: pytest.MonkeyPatch) -> None: redis_stream_mock = SimpleNamespace(mark_complete=AsyncMock(), mark_error=AsyncMock()) redis_stream = cast(TaskRunRedisStream, redis_stream_mock) From 8a5307b3e79df4d7008d48e454536c47e3ee3274 Mon Sep 17 00:00:00 2001 From: Alessandro Pogliaghi Date: Thu, 30 Apr 2026 09:05:47 +0100 Subject: [PATCH 2/3] chore(cloud-agent): cover relay helper branches --- .../tests/test_relay_sandbox_events.py | 48 ++++++++++++++++++- 1 file changed, 46 insertions(+), 2 deletions(-) diff --git a/products/tasks/backend/temporal/process_task/activities/tests/test_relay_sandbox_events.py b/products/tasks/backend/temporal/process_task/activities/tests/test_relay_sandbox_events.py index 3131034c42cf..5bf6df7105c8 100644 --- a/products/tasks/backend/temporal/process_task/activities/tests/test_relay_sandbox_events.py +++ b/products/tasks/backend/temporal/process_task/activities/tests/test_relay_sandbox_events.py @@ -16,6 +16,7 @@ RelaySandboxEventsInput, TaskRunRedisStream, _is_end_of_turn, + _is_keepalive_event, _is_session_update, _mark_complete_if_run_is_terminal, _mark_error_unless_run_is_terminal, @@ -113,6 +114,18 @@ def test_is_session_update(self, _name: str, event_data: dict, expected: bool): assert _is_session_update(event_data) == expected +class TestIsKeepaliveEvent: + @parameterized.expand( + [ + ("keepalive", {"type": "keepalive"}, True), + ("notification", {"type": "notification"}, False), + ("missing_type", {}, False), + ], + ) + def test_is_keepalive_event(self, _name: str, event_data: dict, expected: bool) -> None: + assert _is_keepalive_event(event_data) == expected + + class TestAgentActiveReactivation: """Verify that agent_active is only re-activated by session/update events. @@ -363,6 +376,36 @@ async def aget(self, id: str) -> SimpleNamespace: redis_stream_mock.mark_complete.assert_awaited_once() redis_stream_mock.mark_error.assert_not_awaited() + async def test_missing_run_marks_stream_error_on_normal_stream_close(self, monkeypatch: pytest.MonkeyPatch) -> None: + redis_stream_mock = SimpleNamespace(mark_complete=AsyncMock(), mark_error=AsyncMock()) + redis_stream = cast(TaskRunRedisStream, redis_stream_mock) + + class DoesNotExist(Exception): + pass + + class StubTaskRunQuerySet: + def only(self, *_fields: str) -> "StubTaskRunQuerySet": + return self + + async def aget(self, id: str) -> SimpleNamespace: + raise DoesNotExist + + monkeypatch.setattr( + relay_sandbox_events_module, + "TaskRunModel", + SimpleNamespace( + Status=SimpleNamespace(COMPLETED="completed", FAILED="failed", CANCELLED="cancelled"), + DoesNotExist=DoesNotExist, + objects=StubTaskRunQuerySet(), + ), + ) + + marked_complete = await _mark_complete_if_run_is_terminal(redis_stream, "run-id") + + assert marked_complete is True + redis_stream_mock.mark_complete.assert_not_awaited() + redis_stream_mock.mark_error.assert_awaited_once_with("Task run not found") + async def test_in_progress_normal_stream_close_reconnects_until_exhausted( self, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -384,8 +427,9 @@ async def __aexit__(self, *_args: object) -> None: return None async def aiter_sse(self): - if False: - yield SimpleNamespace(data="") + events: list[SimpleNamespace] = [] + for event in events: + yield event def fake_connect_sse(*_args: object, **_kwargs: object) -> EmptyEventSource: nonlocal connect_attempts From a557cca7b93f22304d5e3b19ab327860a9d4d94a Mon Sep 17 00:00:00 2001 From: Alessandro Pogliaghi Date: Thu, 30 Apr 2026 10:30:21 +0100 Subject: [PATCH 3/3] fix(cloud-agent): complete streams on clean eof --- .../activities/relay_sandbox_events.py | 32 +------ .../tests/test_relay_sandbox_events.py | 89 ++----------------- 2 files changed, 8 insertions(+), 113 deletions(-) diff --git a/products/tasks/backend/temporal/process_task/activities/relay_sandbox_events.py b/products/tasks/backend/temporal/process_task/activities/relay_sandbox_events.py index f0d99183709d..a70ac34e55c2 100644 --- a/products/tasks/backend/temporal/process_task/activities/relay_sandbox_events.py +++ b/products/tasks/backend/temporal/process_task/activities/relay_sandbox_events.py @@ -303,17 +303,9 @@ async def _relay_loop( return # SSE stream ended normally (sandbox closed connection) - if await _mark_complete_if_run_is_terminal(redis_stream, run_id): - logger.info("relay_sandbox_events_stopped_after_terminal_run", run_id=run_id) - return - - reconnect_count += 1 - logger.warning( - "relay_sandbox_events_stream_closed_before_terminal_run", - run_id=run_id, - reconnect_count=reconnect_count, - ) - await asyncio.sleep(min(reconnect_count * 2, 10)) + await redis_stream.mark_complete() + logger.info("relay_sandbox_events_stream_closed", run_id=run_id) + return except httpx.ReadTimeout: reconnect_count += 1 @@ -347,24 +339,6 @@ async def _relay_loop( pass -async def _mark_complete_if_run_is_terminal(redis_stream: TaskRunRedisStream, run_id: str) -> bool: - try: - task_run = await TaskRunModel.objects.only("status").aget(id=run_id) - except TaskRunModel.DoesNotExist: - await redis_stream.mark_error("Task run not found") - return True - - if task_run.status in ( - TaskRunModel.Status.COMPLETED, - TaskRunModel.Status.FAILED, - TaskRunModel.Status.CANCELLED, - ): - await redis_stream.mark_complete() - return True - - return False - - def _is_session_update(event_data: dict) -> bool: """Check if an event is a session/update notification (active agent processing).""" if event_data.get("type") != "notification": diff --git a/products/tasks/backend/temporal/process_task/activities/tests/test_relay_sandbox_events.py b/products/tasks/backend/temporal/process_task/activities/tests/test_relay_sandbox_events.py index 5bf6df7105c8..7896b0a4dd50 100644 --- a/products/tasks/backend/temporal/process_task/activities/tests/test_relay_sandbox_events.py +++ b/products/tasks/backend/temporal/process_task/activities/tests/test_relay_sandbox_events.py @@ -18,7 +18,6 @@ _is_end_of_turn, _is_keepalive_event, _is_session_update, - _mark_complete_if_run_is_terminal, _mark_error_unless_run_is_terminal, _relay_loop, relay_sandbox_events, @@ -347,68 +346,7 @@ async def aget(self, id: str) -> SimpleNamespace: redis_stream_mock.mark_complete.assert_awaited_once() redis_stream_mock.mark_error.assert_not_awaited() - async def test_terminal_run_marks_stream_complete_on_normal_stream_close( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: - redis_stream_mock = SimpleNamespace(mark_complete=AsyncMock(), mark_error=AsyncMock()) - redis_stream = cast(TaskRunRedisStream, redis_stream_mock) - - class StubTaskRunQuerySet: - def only(self, *_fields: str) -> "StubTaskRunQuerySet": - return self - - async def aget(self, id: str) -> SimpleNamespace: - return SimpleNamespace(status="completed") - - monkeypatch.setattr( - relay_sandbox_events_module, - "TaskRunModel", - SimpleNamespace( - Status=SimpleNamespace(COMPLETED="completed", FAILED="failed", CANCELLED="cancelled"), - DoesNotExist=Exception, - objects=StubTaskRunQuerySet(), - ), - ) - - marked_complete = await _mark_complete_if_run_is_terminal(redis_stream, "run-id") - - assert marked_complete is True - redis_stream_mock.mark_complete.assert_awaited_once() - redis_stream_mock.mark_error.assert_not_awaited() - - async def test_missing_run_marks_stream_error_on_normal_stream_close(self, monkeypatch: pytest.MonkeyPatch) -> None: - redis_stream_mock = SimpleNamespace(mark_complete=AsyncMock(), mark_error=AsyncMock()) - redis_stream = cast(TaskRunRedisStream, redis_stream_mock) - - class DoesNotExist(Exception): - pass - - class StubTaskRunQuerySet: - def only(self, *_fields: str) -> "StubTaskRunQuerySet": - return self - - async def aget(self, id: str) -> SimpleNamespace: - raise DoesNotExist - - monkeypatch.setattr( - relay_sandbox_events_module, - "TaskRunModel", - SimpleNamespace( - Status=SimpleNamespace(COMPLETED="completed", FAILED="failed", CANCELLED="cancelled"), - DoesNotExist=DoesNotExist, - objects=StubTaskRunQuerySet(), - ), - ) - - marked_complete = await _mark_complete_if_run_is_terminal(redis_stream, "run-id") - - assert marked_complete is True - redis_stream_mock.mark_complete.assert_not_awaited() - redis_stream_mock.mark_error.assert_awaited_once_with("Task run not found") - - async def test_in_progress_normal_stream_close_reconnects_until_exhausted( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: + async def test_normal_stream_close_marks_stream_complete(self, monkeypatch: pytest.MonkeyPatch) -> None: redis_stream = SimpleNamespace( write_event=AsyncMock(), mark_complete=AsyncMock(), @@ -439,25 +377,9 @@ def fake_connect_sse(*_args: object, **_kwargs: object) -> EmptyEventSource: async def fake_background_heartbeat(*_args: object, **_kwargs: object) -> None: return None - class StubTaskRunQuerySet: - def only(self, *_fields: str) -> "StubTaskRunQuerySet": - return self - - async def aget(self, id: str) -> SimpleNamespace: - return SimpleNamespace(status="in_progress") - monkeypatch.setattr(relay_sandbox_events_module.httpx_sse, "aconnect_sse", fake_connect_sse) monkeypatch.setattr(relay_sandbox_events_module.asyncio, "sleep", sleep_mock) monkeypatch.setattr(relay_sandbox_events_module, "_background_heartbeat", fake_background_heartbeat) - monkeypatch.setattr( - relay_sandbox_events_module, - "TaskRunModel", - SimpleNamespace( - Status=SimpleNamespace(COMPLETED="completed", FAILED="failed", CANCELLED="cancelled"), - DoesNotExist=Exception, - objects=StubTaskRunQuerySet(), - ), - ) await _relay_loop( events_url="https://sandbox.example/events", @@ -468,12 +390,11 @@ async def aget(self, id: str) -> SimpleNamespace: task_id="task-id", ) - assert connect_attempts == relay_sandbox_events_module.MAX_RECONNECT_ATTEMPTS + 1 + assert connect_attempts == 1 + sleep_mock.assert_not_awaited() redis_stream.write_event.assert_not_awaited() - redis_stream.mark_complete.assert_not_awaited() - redis_stream.mark_error.assert_awaited_once_with( - f"Lost connection to sandbox after {relay_sandbox_events_module.MAX_RECONNECT_ATTEMPTS} reconnection attempts" - ) + redis_stream.mark_complete.assert_awaited_once() + redis_stream.mark_error.assert_not_awaited() async def test_in_progress_run_marks_stream_error_on_relay_error(self, monkeypatch: pytest.MonkeyPatch) -> None: redis_stream_mock = SimpleNamespace(mark_complete=AsyncMock(), mark_error=AsyncMock())