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..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 @@ -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): @@ -303,7 +304,8 @@ async def _relay_loop( # SSE stream ended normally (sandbox closed connection) await redis_stream.mark_complete() - return + logger.info("relay_sandbox_events_stream_closed", run_id=run_id) + return except httpx.ReadTimeout: reconnect_count += 1 @@ -345,6 +347,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..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 @@ -16,6 +16,7 @@ RelaySandboxEventsInput, TaskRunRedisStream, _is_end_of_turn, + _is_keepalive_event, _is_session_update, _mark_error_unless_run_is_terminal, _relay_loop, @@ -112,6 +113,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. @@ -333,6 +346,56 @@ 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_normal_stream_close_marks_stream_complete(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): + events: list[SimpleNamespace] = [] + for event in events: + yield event + + 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 + + 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) + + 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 == 1 + sleep_mock.assert_not_awaited() + redis_stream.write_event.assert_not_awaited() + 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()) redis_stream = cast(TaskRunRedisStream, redis_stream_mock)