add cleanup

Author: giancarloromeo

packages/celery-library/src/celery_library/backends/redis.py

@@ -1,3 +1,4 @@
+import asyncio
 import contextlib
 import logging
 from collections.abc import AsyncIterator
@@ -186,11 +187,15 @@ async def consume_task_events(
     ) -> AsyncIterator[tuple[TaskEventID, TaskEvent]]:
         stream_key = _build_stream_key(task_id)
         while True:
-            messages = await self._redis_client_sdk.redis.xread(
-                {stream_key: last_id or _CELERY_TASK_STREAM_DEFAULT_ID},
-                block=_CELERY_TASK_STREAM_BLOCK_TIMEOUT,
-                count=_CELERY_TASK_STREAM_COUNT,
-            )
+            try:
+                messages = await self._redis_client_sdk.redis.xread(
+                    {stream_key: last_id or _CELERY_TASK_STREAM_DEFAULT_ID},
+                    block=_CELERY_TASK_STREAM_BLOCK_TIMEOUT,
+                    count=_CELERY_TASK_STREAM_COUNT,
+                )
+            except asyncio.CancelledError:
+                break
+
             if not messages:
                 continue
             for _, events in messages:
@@ -204,6 +209,7 @@ async def consume_task_events(
                             raw_event
                         )
                         yield msg_id, event
+                        last_id = msg_id
                     except ValidationError:
                         continue
 

services/web/server/src/simcore_service_webserver/tasks/_controller/_rest.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 from typing import Final
 
@@ -193,18 +194,23 @@ async def get_async_job_stream(request: web.Request) -> web.Response:
     _path_params = parse_request_path_parameters_as(TaskPathParams, request)
 
     async def event_generator():
-        async for event_id, event in get_task_manager(request.app).consume_task_events(
-            owner_metadata=OwnerMetadata.model_validate(
-                WebServerOwnerMetadata(
-                    user_id=_req_ctx.user_id,
-                    product_name=_req_ctx.product_name,
-                ).model_dump()
-            ),
-            task_uuid=_path_params.task_id,
-            last_id=_header_params.last_event_id,
-        ):
-            yield SSEEvent(
-                id=event_id, event=event.type, data=[json_dumps(event.data)]
-            ).serialize()
+        try:
+            async for event_id, event in get_task_manager(
+                request.app
+            ).consume_task_events(
+                owner_metadata=OwnerMetadata.model_validate(
+                    WebServerOwnerMetadata(
+                        user_id=_req_ctx.user_id,
+                        product_name=_req_ctx.product_name,
+                    ).model_dump()
+                ),
+                task_uuid=_path_params.task_id,
+                last_id=_header_params.last_event_id,
+            ):
+                yield SSEEvent(
+                    id=event_id, event=event.type, data=[json_dumps(event.data)]
+                ).serialize()
+        except asyncio.CancelledError:
+            return
 
     return create_event_stream_response(event_generator=event_generator)