fix(cloud-agent): handle stream churn

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
 
 

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)