fixes

Author: HardMax71

backend/app/db/repositories/sse_repository.py

@@ -15,9 +15,7 @@ def __init__(self, database: Database) -> None:
         self.mapper = SSEMapper()
 
     async def get_execution_status(self, execution_id: str) -> SSEExecutionStatusDomain | None:
-        doc = await self.executions_collection.find_one(
-            {"execution_id": execution_id}, {"status": 1, "_id": 0}
-        )
+        doc = await self.executions_collection.find_one({"execution_id": execution_id}, {"status": 1, "_id": 0})
         if not doc:
             return None
         return SSEExecutionStatusDomain(

backend/app/schemas_pydantic/execution.py

@@ -6,9 +6,9 @@
 
 from pydantic import BaseModel, ConfigDict, Field, model_validator
 
-from app.domain.enums.storage import ExecutionErrorType
 from app.domain.enums.events import EventType
 from app.domain.enums.execution import ExecutionStatus
+from app.domain.enums.storage import ExecutionErrorType
 from app.settings import get_settings
 
 
@@ -80,8 +80,9 @@ def validate_runtime_supported(self) -> "ExecutionRequest":  # noqa: D401
         if not (lang_info := runtimes.get(self.lang)):
             raise ValueError(f"Language '{self.lang}' not supported. Supported: {list(runtimes.keys())}")
         if self.lang_version not in lang_info.versions:
-            raise ValueError(f"Version '{self.lang_version}' not supported for {self.lang}. "
-                             f"Supported: {lang_info.versions}")
+            raise ValueError(
+                f"Version '{self.lang_version}' not supported for {self.lang}. Supported: {lang_info.versions}"
+            )
         return self
 
 

backend/app/schemas_pydantic/sse.py

@@ -9,10 +9,10 @@
 from app.schemas_pydantic.execution import ExecutionResult, ResourceUsage
 
 # Control event types sent by SSE (not from Kafka)
-SSEControlEventType = Literal['connected', 'heartbeat', 'shutdown', 'status', 'error']
+SSEControlEventType = Literal["connected", "heartbeat", "shutdown", "status", "error"]
 
 # Type variable for generic Redis message parsing
-T = TypeVar('T', bound=BaseModel)
+T = TypeVar("T", bound=BaseModel)
 
 
 class SSEExecutionEventData(BaseModel):
@@ -69,7 +69,7 @@ class RedisSSEMessage(BaseModel):
 
 
 # Control event types for notification SSE stream
-SSENotificationControlEventType = Literal['connected', 'heartbeat', 'notification']
+SSENotificationControlEventType = Literal["connected", "heartbeat", "notification"]
 
 
 class SSENotificationEventData(BaseModel):

backend/app/services/notification_service.py

@@ -34,9 +34,9 @@
     ExecutionTimeoutEvent,
 )
 from app.infrastructure.kafka.mappings import get_topic_for_event
+from app.schemas_pydantic.sse import RedisNotificationMessage
 from app.services.event_bus import EventBusManager
 from app.services.kafka_event_service import KafkaEventService
-from app.schemas_pydantic.sse import RedisNotificationMessage
 from app.services.sse.redis_bus import SSERedisBus
 from app.settings import Settings, get_settings
 

backend/app/services/sse/redis_bus.py

@@ -18,9 +18,9 @@ def __init__(self, pubsub: redis.client.PubSub, channel: str) -> None:
         self._pubsub = pubsub
         self._channel = channel
 
-    async def get(self, model: Type[T], timeout: float = 0.5) -> T | None:
+    async def get(self, model: Type[T]) -> T | None:
         """Get next typed message from the subscription."""
-        msg = await self._pubsub.get_message(ignore_subscribe_messages=True, timeout=timeout)
+        msg = await self._pubsub.get_message(ignore_subscribe_messages=True, timeout=0.5)
         if not msg or msg.get("type") != "message":
             return None
         try:
@@ -39,7 +39,7 @@ class SSERedisBus:
     """Redis-backed pub/sub bus for SSE event fan-out across workers."""
 
     def __init__(
-            self, redis_client: redis.Redis, exec_prefix: str = "sse:exec:", notif_prefix: str = "sse:notif:"
+        self, redis_client: redis.Redis, exec_prefix: str = "sse:exec:", notif_prefix: str = "sse:notif:"
     ) -> None:
         self._redis = redis_client
         self._exec_prefix = exec_prefix

backend/app/services/sse/sse_service.py

@@ -51,24 +51,28 @@ async def create_execution_stream(self, execution_id: str, user_id: str) -> Asyn
 
         shutdown_event = await self.shutdown_manager.register_connection(execution_id, connection_id)
         if shutdown_event is None:
-            yield self._format_sse_event(SSEExecutionEventData(
-                event_type="error",
-                execution_id=execution_id,
-                timestamp=datetime.now(timezone.utc).isoformat(),
-                error="Server is shutting down",
-            ))
+            yield self._format_sse_event(
+                SSEExecutionEventData(
+                    event_type="error",
+                    execution_id=execution_id,
+                    timestamp=datetime.now(timezone.utc).isoformat(),
+                    error="Server is shutting down",
+                )
+            )
             return
 
         subscription = None
         try:
             # Start opening subscription concurrently, then yield handshake
             sub_task = asyncio.create_task(self.sse_bus.open_subscription(execution_id))
-            yield self._format_sse_event(SSEExecutionEventData(
-                event_type="connected",
-                execution_id=execution_id,
-                timestamp=datetime.now(timezone.utc).isoformat(),
-                connection_id=connection_id,
-            ))
+            yield self._format_sse_event(
+                SSEExecutionEventData(
+                    event_type="connected",
+                    execution_id=execution_id,
+                    timestamp=datetime.now(timezone.utc).isoformat(),
+                    connection_id=connection_id,
+                )
+            )
 
             # Complete Redis subscription after handshake
             logger.info(f"Opening Redis subscription for execution {execution_id}")
@@ -77,12 +81,14 @@ async def create_execution_stream(self, execution_id: str, user_id: str) -> Asyn
 
             initial_status = await self.repository.get_execution_status(execution_id)
             if initial_status:
-                yield self._format_sse_event(SSEExecutionEventData(
-                    event_type="status",
-                    execution_id=initial_status.execution_id,
-                    timestamp=initial_status.timestamp,
-                    status=initial_status.status,
-                ))
+                yield self._format_sse_event(
+                    SSEExecutionEventData(
+                        event_type="status",
+                        execution_id=initial_status.execution_id,
+                        timestamp=initial_status.timestamp,
+                        status=initial_status.status,
+                    )
+                )
                 self.metrics.record_sse_message_sent("executions", "status")
 
             async for event_data in self._stream_events_redis(
@@ -109,26 +115,30 @@ async def _stream_events_redis(
         last_heartbeat = datetime.now(timezone.utc)
         while True:
             if shutdown_event.is_set():
-                yield self._format_sse_event(SSEExecutionEventData(
-                    event_type="shutdown",
-                    execution_id=execution_id,
-                    timestamp=datetime.now(timezone.utc).isoformat(),
-                    message="Server is shutting down",
-                    grace_period=30,
-                ))
+                yield self._format_sse_event(
+                    SSEExecutionEventData(
+                        event_type="shutdown",
+                        execution_id=execution_id,
+                        timestamp=datetime.now(timezone.utc).isoformat(),
+                        message="Server is shutting down",
+                        grace_period=30,
+                    )
+                )
                 break
 
             now = datetime.now(timezone.utc)
             if include_heartbeat and (now - last_heartbeat).total_seconds() >= self.heartbeat_interval:
-                yield self._format_sse_event(SSEExecutionEventData(
-                    event_type="heartbeat",
-                    execution_id=execution_id,
-                    timestamp=now.isoformat(),
-                    message="SSE connection active",
-                ))
+                yield self._format_sse_event(
+                    SSEExecutionEventData(
+                        event_type="heartbeat",
+                        execution_id=execution_id,
+                        timestamp=now.isoformat(),
+                        message="SSE connection active",
+                    )
+                )
                 last_heartbeat = now
 
-            msg: RedisSSEMessage | None = await subscription.get(RedisSSEMessage, timeout=0.5)
+            msg: RedisSSEMessage | None = await subscription.get(RedisSSEMessage)
             if not msg:
                 continue
 
@@ -145,36 +155,38 @@ async def _stream_events_redis(
                 # Ignore malformed messages
                 continue
 
-    async def _build_sse_event_from_redis(
-        self, execution_id: str, msg: RedisSSEMessage
-    ) -> SSEExecutionEventData:
+    async def _build_sse_event_from_redis(self, execution_id: str, msg: RedisSSEMessage) -> SSEExecutionEventData:
         """Build typed SSE event from Redis message."""
         result: ExecutionResult | None = None
         if msg.event_type == EventType.RESULT_STORED:
             exec_domain = await self.repository.get_execution(execution_id)
             if exec_domain:
                 result = ExecutionResult.model_validate(exec_domain)
 
-        return SSEExecutionEventData.model_validate({
-            **msg.data,
-            "event_type": msg.event_type,
-            "execution_id": execution_id,
-            "type": msg.event_type,
-            "result": result,
-        })
+        return SSEExecutionEventData.model_validate(
+            {
+                **msg.data,
+                "event_type": msg.event_type,
+                "execution_id": execution_id,
+                "type": msg.event_type,
+                "result": result,
+            }
+        )
 
     async def create_notification_stream(self, user_id: str) -> AsyncGenerator[Dict[str, Any], None]:
         subscription = None
 
         try:
             # Start opening subscription concurrently, then yield handshake
             sub_task = asyncio.create_task(self.sse_bus.open_notification_subscription(user_id))
-            yield self._format_notification_event(SSENotificationEventData(
-                event_type="connected",
-                user_id=user_id,
-                timestamp=datetime.now(timezone.utc).isoformat(),
-                message="Connected to notification stream",
-            ))
+            yield self._format_notification_event(
+                SSENotificationEventData(
+                    event_type="connected",
+                    user_id=user_id,
+                    timestamp=datetime.now(timezone.utc).isoformat(),
+                    message="Connected to notification stream",
+                )
+            )
 
             # Complete Redis subscription after handshake
             subscription = await sub_task
@@ -184,28 +196,32 @@ async def create_notification_stream(self, user_id: str) -> AsyncGenerator[Dict[
                 # Heartbeat
                 now = datetime.now(timezone.utc)
                 if (now - last_heartbeat).total_seconds() >= self.heartbeat_interval:
-                    yield self._format_notification_event(SSENotificationEventData(
-                        event_type="heartbeat",
-                        user_id=user_id,
-                        timestamp=now.isoformat(),
-                        message="Notification stream active",
-                    ))
+                    yield self._format_notification_event(
+                        SSENotificationEventData(
+                            event_type="heartbeat",
+                            user_id=user_id,
+                            timestamp=now.isoformat(),
+                            message="Notification stream active",
+                        )
+                    )
                     last_heartbeat = now
 
                 # Forward notification messages as SSE data
-                redis_msg = await subscription.get(RedisNotificationMessage, timeout=0.5)
+                redis_msg = await subscription.get(RedisNotificationMessage)
                 if redis_msg:
-                    yield self._format_notification_event(SSENotificationEventData(
-                        event_type="notification",
-                        notification_id=redis_msg.notification_id,
-                        severity=redis_msg.severity,
-                        status=redis_msg.status,
-                        tags=redis_msg.tags,
-                        subject=redis_msg.subject,
-                        body=redis_msg.body,
-                        action_url=redis_msg.action_url,
-                        created_at=redis_msg.created_at,
-                    ))
+                    yield self._format_notification_event(
+                        SSENotificationEventData(
+                            event_type="notification",
+                            notification_id=redis_msg.notification_id,
+                            severity=redis_msg.severity,
+                            status=redis_msg.status,
+                            tags=redis_msg.tags,
+                            subject=redis_msg.subject,
+                            body=redis_msg.body,
+                            action_url=redis_msg.action_url,
+                            created_at=redis_msg.created_at,
+                        )
+                    )
         finally:
             try:
                 if subscription is not None:

backend/tests/integration/notifications/test_notification_sse.py

@@ -4,6 +4,7 @@
 import pytest
 
 from app.domain.enums.notification import NotificationChannel, NotificationSeverity
+from app.schemas_pydantic.sse import RedisNotificationMessage
 from app.services.notification_service import NotificationService
 from app.services.sse.redis_bus import SSERedisBus
 from tests.helpers.eventually import eventually
@@ -34,13 +35,13 @@ async def test_in_app_notification_published_to_sse(scope) -> None:  # type: ign
     )
 
     # Receive published SSE payload
-    async def _recv():
-        m = await sub.get(timeout=0.5)
+    async def _recv() -> RedisNotificationMessage:
+        m = await sub.get(RedisNotificationMessage)
         assert m is not None
         return m
 
     msg = await eventually(_recv, timeout=5.0, interval=0.1)
     # Basic shape assertions
-    assert msg.get("subject") == "Hello"
-    assert msg.get("body") == "World"
-    assert "notification_id" in msg
+    assert msg.subject == "Hello"
+    assert msg.body == "World"
+    assert msg.notification_id

backend/tests/integration/services/sse/test_partitioned_event_router_integration.py

@@ -8,6 +8,7 @@
 from app.events.schema.schema_registry import SchemaRegistryManager
 from tests.helpers import make_execution_requested_event
 from app.infrastructure.kafka.events.pod import PodCreatedEvent
+from app.schemas_pydantic.sse import RedisSSEMessage
 from app.services.sse.kafka_redis_bridge import SSEKafkaRedisBridge
 from app.services.sse.redis_bus import SSERedisBus
 from app.settings import Settings
@@ -42,12 +43,12 @@ async def test_router_bridges_to_redis(redis_client) -> None:  # type: ignore[va
     await handler(ev)
 
     async def _recv():
-        m = await subscription.get(timeout=0.2)
+        m = await subscription.get(RedisSSEMessage)
         assert m is not None
         return m
 
     msg = await eventually(_recv, timeout=2.0, interval=0.05)
-    assert msg.get("event_type") == str(ev.event_type)
+    assert str(msg.event_type) == str(ev.event_type)
 
 
 @pytest.mark.asyncio

backend/tests/integration/test_sse_routes.py

@@ -6,7 +6,8 @@
 import pytest
 from httpx import AsyncClient
 
-from app.schemas_pydantic.sse import SSEHealthResponse
+from app.domain.enums.notification import NotificationSeverity, NotificationStatus
+from app.schemas_pydantic.sse import RedisNotificationMessage, SSEHealthResponse
 from app.infrastructure.kafka.events.pod import PodCreatedEvent
 from app.infrastructure.kafka.events.metadata import EventMetadata
 from app.services.sse.redis_bus import SSERedisBus
@@ -73,7 +74,17 @@ async def _connected() -> None:
         await eventually(_connected, timeout=2.0, interval=0.05)
 
         # Publish a notification
-        await bus.publish_notification(user_id, {"subject": "Hello", "body": "World", "event_type": "notification"})
+        notification = RedisNotificationMessage(
+            notification_id=f"notif-{uuid4().hex[:8]}",
+            severity=NotificationSeverity.MEDIUM,
+            status=NotificationStatus.PENDING,
+            tags=[],
+            subject="Hello",
+            body="World",
+            action_url="",
+            created_at="2024-01-01T00:00:00Z",
+        )
+        await bus.publish_notification(user_id, notification)
 
         # Wait for collection to complete
         try: