[SKY-6] Backend: Enable 2FA code detection without TOTP credentials (#4786)

skyvern/forge/sdk/notification/base.py

@@ -0,0 +1,21 @@
+"""Abstract base for notification registries."""
+
+import asyncio
+from abc import ABC, abstractmethod
+
+
+class BaseNotificationRegistry(ABC):
+    """Abstract pub/sub registry scoped by organization.
+
+    Implementations must fan-out: a single publish call delivers the
+    message to every active subscriber for that organization.
+    """
+
+    @abstractmethod
+    def subscribe(self, organization_id: str) -> asyncio.Queue[dict]: ...
+
+    @abstractmethod
+    def unsubscribe(self, organization_id: str, queue: asyncio.Queue[dict]) -> None: ...
+
+    @abstractmethod
+    def publish(self, organization_id: str, message: dict) -> None: ...

skyvern/forge/sdk/notification/factory.py

@@ -0,0 +1,14 @@
+from skyvern.forge.sdk.notification.base import BaseNotificationRegistry
+from skyvern.forge.sdk.notification.local import LocalNotificationRegistry
+
+
+class NotificationRegistryFactory:
+    __registry: BaseNotificationRegistry = LocalNotificationRegistry()
+
+    @staticmethod
+    def set_registry(registry: BaseNotificationRegistry) -> None:
+        NotificationRegistryFactory.__registry = registry
+
+    @staticmethod
+    def get_registry() -> BaseNotificationRegistry:
+        return NotificationRegistryFactory.__registry

skyvern/forge/sdk/notification/local.py

@@ -0,0 +1,45 @@
+"""In-process notification registry using asyncio queues (single-pod only)."""
+
+import asyncio
+from collections import defaultdict
+
+import structlog
+
+from skyvern.forge.sdk.notification.base import BaseNotificationRegistry
+
+LOG = structlog.get_logger()
+
+
+class LocalNotificationRegistry(BaseNotificationRegistry):
+    """In-process fan-out pub/sub using asyncio queues. Single-pod only."""
+
+    def __init__(self) -> None:
+        self._subscribers: dict[str, list[asyncio.Queue[dict]]] = defaultdict(list)
+
+    def subscribe(self, organization_id: str) -> asyncio.Queue[dict]:
+        queue: asyncio.Queue[dict] = asyncio.Queue()
+        self._subscribers[organization_id].append(queue)
+        LOG.info("Notification subscriber added", organization_id=organization_id)
+        return queue
+
+    def unsubscribe(self, organization_id: str, queue: asyncio.Queue[dict]) -> None:
+        queues = self._subscribers.get(organization_id)
+        if queues:
+            try:
+                queues.remove(queue)
+            except ValueError:
+                pass
+            if not queues:
+                del self._subscribers[organization_id]
+        LOG.info("Notification subscriber removed", organization_id=organization_id)
+
+    def publish(self, organization_id: str, message: dict) -> None:
+        queues = self._subscribers.get(organization_id, [])
+        for queue in queues:
+            try:
+                queue.put_nowait(message)
+            except asyncio.QueueFull:
+                LOG.warning(
+                    "Notification queue full, dropping message",
+                    organization_id=organization_id,
+                )

skyvern/forge/sdk/notification/redis.py

@@ -0,0 +1,55 @@
+"""Redis-backed notification registry for multi-pod deployments.
+
+Thin adapter around :class:`RedisPubSub` — all Redis pub/sub logic
+lives in the generic layer; this class maps the ``organization_id``
+domain concept onto generic string keys.
+"""
+
+import asyncio
+
+from redis.asyncio import Redis
+
+from skyvern.forge.sdk.notification.base import BaseNotificationRegistry
+from skyvern.forge.sdk.redis.pubsub import RedisPubSub
+
+
+class RedisNotificationRegistry(BaseNotificationRegistry):
+    """Fan-out pub/sub backed by Redis.  One Redis PubSub channel per org."""
+
+    def __init__(self, redis_client: Redis) -> None:
+        self._pubsub = RedisPubSub(redis_client, channel_prefix="skyvern:notifications:")
+
+    # ------------------------------------------------------------------
+    # Property accessors (used by existing tests)
+    # ------------------------------------------------------------------
+
+    @property
+    def _listener_tasks(self) -> dict[str, asyncio.Task[None]]:
+        return self._pubsub._listener_tasks
+
+    @property
+    def _subscribers(self) -> dict[str, list[asyncio.Queue[dict]]]:
+        return self._pubsub._subscribers
+
+    # ------------------------------------------------------------------
+    # Public interface
+    # ------------------------------------------------------------------
+
+    def subscribe(self, organization_id: str) -> asyncio.Queue[dict]:
+        return self._pubsub.subscribe(organization_id)
+
+    def unsubscribe(self, organization_id: str, queue: asyncio.Queue[dict]) -> None:
+        self._pubsub.unsubscribe(organization_id, queue)
+
+    def publish(self, organization_id: str, message: dict) -> None:
+        self._pubsub.publish(organization_id, message)
+
+    async def close(self) -> None:
+        await self._pubsub.close()
+
+    # ------------------------------------------------------------------
+    # Internal helper (exposed for tests)
+    # ------------------------------------------------------------------
+
+    def _dispatch_local(self, organization_id: str, message: dict) -> None:
+        self._pubsub._dispatch_local(organization_id, message)