fix: refactor

Author: giancarloromeo

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

@@ -9,6 +9,7 @@
 from servicelib.celery.models import (
     Task,
     TaskEvent,
+    TaskEventID,
     TaskFilter,
     TaskID,
     TaskInfoStore,
@@ -161,7 +162,7 @@ async def publish_task_event(self, task_id: TaskID, event: TaskEvent) -> None:
 
     async def consume_task_events(
         self, task_id: TaskID, last_id: str | None = None
-    ) -> AsyncIterator[TaskEvent]:
+    ) -> AsyncIterator[tuple[TaskEventID, TaskEvent]]:
         stream_key = _build_stream_key(task_id)
         while True:
             messages = await self._redis_client_sdk.redis.xread(
@@ -183,8 +184,7 @@ async def consume_task_events(
                         event: TaskEvent = TypeAdapter(TaskEvent).validate_json(
                             raw_event
                         )
-                        event.event_id = msg_id
-                        yield event
+                        yield msg_id, event
                     except ValidationError:
                         continue
 

packages/service-library/src/servicelib/celery/models.py

@@ -127,16 +127,15 @@ class TaskMetadata(BaseModel):
     queue: TasksQueue = TasksQueue.DEFAULT
 
 
-class BaseTaskEvent(BaseModel):
-    event_id: str | None = None
+TaskEventID: TypeAlias = str
 
 
-class TaskDataEvent(BaseTaskEvent):
+class TaskDataEvent(BaseModel):
     type: Literal["data"] = "data"
     data: Any
 
 
-class TaskStatusEvent(BaseTaskEvent):
+class TaskStatusEvent(BaseModel):
     type: Literal["status"] = "status"
     data: Literal["done", "error"]
 
@@ -217,7 +216,7 @@ def consume_task_events(
         self,
         task_id: TaskID,
         last_id: str | None = None,
-    ) -> AsyncIterator[TaskEvent]: ...
+    ) -> AsyncIterator[tuple[TaskEventID, TaskEvent]]: ...
 
 
 class TaskStatus(BaseModel):

packages/service-library/src/servicelib/celery/task_manager.py

@@ -6,6 +6,7 @@
 from ..celery.models import (
     Task,
     TaskEvent,
+    TaskEventID,
     TaskFilter,
     TaskID,
     TaskMetadata,
@@ -53,4 +54,4 @@ def consume_task_events(
         task_filter: TaskFilter,
         task_uuid: TaskUUID,
         last_id: str | None = None,
-    ) -> AsyncIterator[TaskEvent]: ...
+    ) -> AsyncIterator[tuple[TaskEventID, TaskEvent]]: ...

packages/service-library/src/servicelib/sse.py

@@ -0,0 +1,30 @@
+from typing import Annotated
+
+from pydantic import BaseModel, BeforeValidator
+
+
+def _normalize_data(v: str | list[str]) -> list[str]:
+    if isinstance(v, str):
+        lines = v.splitlines()
+        return lines if lines else [""]
+    return v
+
+
+class SSEEvent(BaseModel):
+    id: str | None = None
+    event: str | None = None
+    data: Annotated[str | list[str], BeforeValidator(_normalize_data)]
+    retry: int | None = None
+
+    def serialize(self) -> bytes:
+        lines = []
+        if self.id is not None:
+            lines.append(f"id: {self.id}")
+        if self.event is not None:
+            lines.append(f"event: {self.event}")
+        lines.extend(f"data: {line}" for line in self.data)
+        if self.retry is not None:
+            lines.append(f"retry: {self.retry}")
+
+        payload = "\n".join(lines) + "\n\n"
+        return payload.encode("utf-8")

services/web/server/src/simcore_service_webserver/storage/_rest.py

@@ -52,7 +52,7 @@
     create_data_response,
     create_event_stream_response,
 )
-from servicelib.celery.models import TaskEvent, TaskFilter
+from servicelib.celery.models import TaskFilter
 from servicelib.common_headers import X_FORWARDED_PROTO
 from servicelib.rabbitmq.rpc_interfaces.storage.paths import (
     compute_path_size as remote_compute_path_size,
@@ -66,6 +66,7 @@
 )
 from servicelib.request_keys import RQT_USERID_KEY
 from servicelib.rest_responses import unwrap_envelope
+from servicelib.sse import SSEEvent
 from simcore_service_webserver.celery import get_task_manager
 from simcore_service_webserver.storage._rest_schemas import StreamHeaders
 from yarl import URL
@@ -583,14 +584,6 @@ class _PathParams(BaseModel):
     )
 
 
-def _format_sse(event: TaskEvent, event_name) -> bytes:
-    sse = ""
-    if event_name:
-        sse += f"event: {event_name}\n"
-    sse += f"data: {event.model_dump_json()}\n\n"
-    return sse.encode("utf-8")
-
-
 @routes.get(
     _storage_locations_prefix + "/{location_id}/search/{job_id}/stream",
     name="stream_search",
@@ -614,12 +607,14 @@ class _PathParams(BaseModel):
     )
 
     async def event_generator():
-        async for event in task_manager.consume_task_events(
+        async for event_id, event in task_manager.consume_task_events(
             task_filter=TaskFilter.model_validate(task_filter.model_dump()),
             task_uuid=path_params.job_id,
             last_id=header_params.last_event_id,
         ):
-            yield _format_sse(event, event_name=event.type)
+            yield SSEEvent(
+                id=event_id, event=event.type, data=event.model_dump_json()
+            ).serialize()
             if event.type == "status" and getattr(event, "data", None) in (
                 "done",
                 "error",