diff --git a/.env.example b/.env.example index 030c40271..dd5d44599 100644 --- a/.env.example +++ b/.env.example @@ -143,10 +143,4 @@ SKYVERN_AUTH_BITWARDEN_CLIENT_SECRET=your-client-secret-here # BITWARDEN_TIMEOUT_SECONDS=60 # Shared Redis URL used by any service that needs Redis (pub/sub, cache, etc.) -# REDIS_URL=redis://localhost:6379/0 - -# Notification registry type: "local" (default, in-process) or "redis" (multi-pod) -# NOTIFICATION_REGISTRY_TYPE=local - -# Optional: override Redis URL specifically for notifications (falls back to REDIS_URL) -# NOTIFICATION_REDIS_URL= \ No newline at end of file +# REDIS_URL=redis://localhost:6379/0 \ No newline at end of file diff --git a/skyvern/config.py b/skyvern/config.py index 6be624f11..0df1910fe 100644 --- a/skyvern/config.py +++ b/skyvern/config.py @@ -101,10 +101,6 @@ class Settings(BaseSettings): # Shared Redis URL (used by any service that needs Redis) REDIS_URL: str = "redis://localhost:6379/0" - # Notification registry settings ("local" or "redis") - NOTIFICATION_REGISTRY_TYPE: str = "local" - NOTIFICATION_REDIS_URL: str | None = None # Deprecated: falls back to REDIS_URL - # S3/AWS settings AWS_REGION: str = "us-east-1" MAX_UPLOAD_FILE_SIZE: int = 10 * 1024 * 1024 # 10 MB diff --git a/skyvern/core/script_generations/real_skyvern_page_ai.py b/skyvern/core/script_generations/real_skyvern_page_ai.py index de2a41bb8..c1493ca45 100644 --- a/skyvern/core/script_generations/real_skyvern_page_ai.py +++ b/skyvern/core/script_generations/real_skyvern_page_ai.py @@ -20,7 +20,7 @@ from skyvern.forge.sdk.core import skyvern_context from skyvern.forge.sdk.schemas.totp_codes import OTPType from skyvern.schemas.workflows import BlockStatus from skyvern.services import script_service -from skyvern.services.otp_service import poll_otp_value, try_generate_totp_from_credential +from skyvern.services.otp_service import poll_otp_value from skyvern.utils.prompt_engine import load_prompt_with_elements from skyvern.webeye.actions import handler_utils from skyvern.webeye.actions.actions import ( @@ -257,10 +257,7 @@ class RealSkyvernPageAi(SkyvernPageAi): if value and isinstance(data, dict) and "value" not in data: data["value"] = value - # Try credential TOTP first (highest priority, doesn't need totp_url/totp_identifier) - otp_value = try_generate_totp_from_credential(workflow_run_id) - # Fall back to webhook/totp_identifier - if not otp_value and (totp_identifier or totp_url) and context and organization_id and task_id: + if (totp_identifier or totp_url) and context and organization_id and task_id: if totp_identifier: totp_identifier = _render_template_with_label(totp_identifier, label=self.current_label) if totp_url: diff --git a/skyvern/core/script_generations/script_skyvern_page.py b/skyvern/core/script_generations/script_skyvern_page.py index b97125543..6d5516cba 100644 --- a/skyvern/core/script_generations/script_skyvern_page.py +++ b/skyvern/core/script_generations/script_skyvern_page.py @@ -27,7 +27,7 @@ from skyvern.forge.sdk.api.files import ( from skyvern.forge.sdk.artifact.models import ArtifactType from skyvern.forge.sdk.core import skyvern_context from skyvern.schemas.steps import AgentStepOutput -from skyvern.services.otp_service import poll_otp_value, try_generate_totp_from_credential +from skyvern.services.otp_service import poll_otp_value from skyvern.utils.url_validators import prepend_scheme_and_validate_url from skyvern.webeye.actions.action_types import ActionType from skyvern.webeye.actions.actions import ( @@ -626,21 +626,16 @@ class ScriptSkyvernPage(SkyvernPage): if is_totp_value: value = generate_totp_value(context.workflow_run_id, original_value) elif (totp_identifier or totp_url) and organization_id: - # Try credential TOTP first (higher priority than webhook/totp_identifier) - credential_totp = try_generate_totp_from_credential(workflow_run_id) - if credential_totp: - value = credential_totp.value - else: - totp_value = await poll_otp_value( - organization_id=organization_id, - task_id=task_id, - workflow_run_id=workflow_run_id, - totp_verification_url=totp_url, - totp_identifier=totp_identifier, - ) - if totp_value: - # use the totp verification code - value = totp_value.value + totp_value = await poll_otp_value( + organization_id=organization_id, + task_id=task_id, + workflow_run_id=workflow_run_id, + totp_verification_url=totp_url, + totp_identifier=totp_identifier, + ) + if totp_value: + # use the totp verification code + value = totp_value.value return value diff --git a/skyvern/forge/agent.py b/skyvern/forge/agent.py index e16df2696..ee79717be 100644 --- a/skyvern/forge/agent.py +++ b/skyvern/forge/agent.py @@ -105,11 +105,7 @@ from skyvern.schemas.steps import AgentStepOutput from skyvern.services import run_service, service_utils from skyvern.services.action_service import get_action_history from skyvern.services.error_detection_service import detect_user_defined_errors_for_task -from skyvern.services.otp_service import ( - extract_totp_from_navigation_inputs, - poll_otp_value, - try_generate_totp_from_credential, -) +from skyvern.services.otp_service import poll_otp_value from skyvern.utils.image_resizer import Resolution from skyvern.utils.prompt_engine import MaxStepsReasonResponse, load_prompt_with_elements from skyvern.webeye.actions.action_types import ActionType @@ -2591,7 +2587,7 @@ class ForgeAgent: step, browser_state, scraped_page, - verification_code_check=True, + verification_code_check=bool(task.totp_verification_url or task.totp_identifier), expire_verification_code=True, ) @@ -4526,6 +4522,9 @@ class ForgeAgent: if not task.organization_id: return json_response, [] + if not task.totp_verification_url and not task.totp_identifier: + return json_response, [] + should_verify_by_magic_link = json_response.get("should_verify_by_magic_link") place_to_enter_verification_code = json_response.get("place_to_enter_verification_code") should_enter_verification_code = json_response.get("should_enter_verification_code") @@ -4546,10 +4545,8 @@ class ForgeAgent: return json_response, actions if should_verify_by_magic_link: - # Magic links still require TOTP config (need a source to poll the link from) - if task.totp_verification_url or task.totp_identifier: - actions = await self.handle_potential_magic_link(task, step, scraped_page, browser_state, json_response) - return json_response, actions + actions = await self.handle_potential_magic_link(task, step, scraped_page, browser_state, json_response) + return json_response, actions return json_response, [] @@ -4606,40 +4603,31 @@ class ForgeAgent: ) -> dict[str, Any]: place_to_enter_verification_code = json_response.get("place_to_enter_verification_code") should_enter_verification_code = json_response.get("should_enter_verification_code") - if place_to_enter_verification_code and should_enter_verification_code and task.organization_id: + if ( + place_to_enter_verification_code + and should_enter_verification_code + and (task.totp_verification_url or task.totp_identifier) + and task.organization_id + ): LOG.info("Need verification code") - # 1. Check navigation payload first for inline OTP. - otp_value = extract_totp_from_navigation_inputs(task.navigation_payload) - if otp_value: - # Code was already in the payload the LLM saw, so its actions are already correct. - # No need to re-prompt, generate from credentials, or poll. - return json_response - - # 2. Then try to generate TOTP from credential if payload has no OTP. - otp_value = try_generate_totp_from_credential(task.workflow_run_id) - # 3. Lastly, poll for OTP if organization has config and no OTP was found yet. - if not otp_value: - workflow_id = workflow_permanent_id = None - if task.workflow_run_id: - workflow_run = await app.DATABASE.get_workflow_run(task.workflow_run_id) - if workflow_run: - workflow_id = workflow_run.workflow_id - workflow_permanent_id = workflow_run.workflow_permanent_id - otp_value = await poll_otp_value( - organization_id=task.organization_id, - task_id=task.task_id, - workflow_id=workflow_id, - workflow_run_id=task.workflow_run_id, - workflow_permanent_id=workflow_permanent_id, - totp_verification_url=task.totp_verification_url, - totp_identifier=task.totp_identifier, - ) - + workflow_id = workflow_permanent_id = None + if task.workflow_run_id: + workflow_run = await app.DATABASE.get_workflow_run(task.workflow_run_id) + if workflow_run: + workflow_id = workflow_run.workflow_id + workflow_permanent_id = workflow_run.workflow_permanent_id + otp_value = await poll_otp_value( + organization_id=task.organization_id, + task_id=task.task_id, + workflow_id=workflow_id, + workflow_run_id=task.workflow_run_id, + workflow_permanent_id=workflow_permanent_id, + totp_verification_url=task.totp_verification_url, + totp_identifier=task.totp_identifier, + ) if not otp_value or otp_value.get_otp_type() != OTPType.TOTP: return json_response - # Store the code so _build_navigation_payload injects it, then re-prompt - # the LLM so it generates actions that type the real code. current_context = skyvern_context.ensure_context() current_context.totp_codes[task.task_id] = otp_value.value diff --git a/skyvern/forge/api_app.py b/skyvern/forge/api_app.py index 9b9dce16d..7396ab778 100644 --- a/skyvern/forge/api_app.py +++ b/skyvern/forge/api_app.py @@ -118,20 +118,6 @@ async def lifespan(fastapi_app: FastAPI) -> AsyncGenerator[None, Any]: # Stop cleanup scheduler await stop_cleanup_scheduler() - # Close notification registry (e.g. cancel Redis listener tasks) - from skyvern.forge.sdk.notification.factory import NotificationRegistryFactory - - registry = NotificationRegistryFactory.get_registry() - if hasattr(registry, "close"): - await registry.close() - - # Close shared Redis client (after registry so listener tasks drain first) - from skyvern.forge.sdk.redis.factory import RedisClientFactory - - redis_client = RedisClientFactory.get_client() - if redis_client is not None: - await redis_client.close() - # Close all persistent browser sessions from skyvern.webeye.default_persistent_sessions_manager import DefaultPersistentSessionsManager diff --git a/skyvern/forge/forge_app.py b/skyvern/forge/forge_app.py index 78b591220..20fdeac59 100644 --- a/skyvern/forge/forge_app.py +++ b/skyvern/forge/forge_app.py @@ -113,16 +113,6 @@ def create_forge_app() -> ForgeApp: app.STORAGE = StorageFactory.get_storage() app.CACHE = CacheFactory.get_cache() - if settings.NOTIFICATION_REGISTRY_TYPE == "redis" and settings.NOTIFICATION_REDIS_URL: - from redis.asyncio import from_url as redis_from_url - - from skyvern.forge.sdk.notification.factory import NotificationRegistryFactory - from skyvern.forge.sdk.notification.redis import RedisNotificationRegistry - from skyvern.forge.sdk.redis.factory import RedisClientFactory - - redis_client = redis_from_url(settings.NOTIFICATION_REDIS_URL, decode_responses=True) - RedisClientFactory.set_client(redis_client) - NotificationRegistryFactory.set_registry(RedisNotificationRegistry(redis_client)) app.ARTIFACT_MANAGER = ArtifactManager() app.BROWSER_MANAGER = RealBrowserManager() app.EXPERIMENTATION_PROVIDER = NoOpExperimentationProvider() diff --git a/skyvern/forge/sdk/db/utils.py b/skyvern/forge/sdk/db/utils.py index b3326a5eb..23d1c99f6 100644 --- a/skyvern/forge/sdk/db/utils.py +++ b/skyvern/forge/sdk/db/utils.py @@ -211,9 +211,6 @@ def convert_to_task(task_obj: TaskModel, debug_enabled: bool = False, workflow_p browser_session_id=task_obj.browser_session_id, browser_address=task_obj.browser_address, download_timeout=task_obj.download_timeout, - waiting_for_verification_code=task_obj.waiting_for_verification_code or False, - verification_code_identifier=task_obj.verification_code_identifier, - verification_code_polling_started_at=task_obj.verification_code_polling_started_at, ) return task @@ -429,9 +426,6 @@ def convert_to_workflow_run( run_with=workflow_run_model.run_with, code_gen=workflow_run_model.code_gen, ai_fallback=workflow_run_model.ai_fallback, - waiting_for_verification_code=workflow_run_model.waiting_for_verification_code or False, - verification_code_identifier=workflow_run_model.verification_code_identifier, - verification_code_polling_started_at=workflow_run_model.verification_code_polling_started_at, ) diff --git a/skyvern/forge/sdk/notification/base.py b/skyvern/forge/sdk/notification/base.py deleted file mode 100644 index 66ef84807..000000000 --- a/skyvern/forge/sdk/notification/base.py +++ /dev/null @@ -1,21 +0,0 @@ -"""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: ... diff --git a/skyvern/forge/sdk/notification/factory.py b/skyvern/forge/sdk/notification/factory.py deleted file mode 100644 index 85e7bfcfe..000000000 --- a/skyvern/forge/sdk/notification/factory.py +++ /dev/null @@ -1,14 +0,0 @@ -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 diff --git a/skyvern/forge/sdk/notification/local.py b/skyvern/forge/sdk/notification/local.py deleted file mode 100644 index 9ed1d7f11..000000000 --- a/skyvern/forge/sdk/notification/local.py +++ /dev/null @@ -1,45 +0,0 @@ -"""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, - ) diff --git a/skyvern/forge/sdk/notification/redis.py b/skyvern/forge/sdk/notification/redis.py deleted file mode 100644 index fe1e9b844..000000000 --- a/skyvern/forge/sdk/notification/redis.py +++ /dev/null @@ -1,55 +0,0 @@ -"""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) diff --git a/skyvern/forge/sdk/redis/__init__.py b/skyvern/forge/sdk/redis/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/skyvern/forge/sdk/redis/factory.py b/skyvern/forge/sdk/redis/factory.py deleted file mode 100644 index 31a69612a..000000000 --- a/skyvern/forge/sdk/redis/factory.py +++ /dev/null @@ -1,21 +0,0 @@ -from __future__ import annotations - -from redis.asyncio import Redis - - -class RedisClientFactory: - """Singleton factory for a shared async Redis client. - - Follows the same static set/get pattern as ``CacheFactory``. - Defaults to ``None`` (no Redis in local/OSS mode). - """ - - __client: Redis | None = None - - @staticmethod - def set_client(client: Redis) -> None: - RedisClientFactory.__client = client - - @staticmethod - def get_client() -> Redis | None: - return RedisClientFactory.__client diff --git a/skyvern/forge/sdk/redis/pubsub.py b/skyvern/forge/sdk/redis/pubsub.py deleted file mode 100644 index fed4f79fc..000000000 --- a/skyvern/forge/sdk/redis/pubsub.py +++ /dev/null @@ -1,130 +0,0 @@ -"""Generic Redis pub/sub layer. - -Extracted from ``RedisNotificationRegistry`` so that any feature -(notifications, events, cache invalidation, etc.) can reuse the same -pattern with its own channel prefix. -""" - -import asyncio -import json -from collections import defaultdict - -import structlog -from redis.asyncio import Redis - -LOG = structlog.get_logger() - - -class RedisPubSub: - """Fan-out pub/sub backed by Redis. One Redis PubSub channel per key.""" - - def __init__(self, redis_client: Redis, channel_prefix: str) -> None: - self._redis = redis_client - self._channel_prefix = channel_prefix - self._subscribers: dict[str, list[asyncio.Queue[dict]]] = defaultdict(list) - # One listener task per key channel - self._listener_tasks: dict[str, asyncio.Task[None]] = {} - - # ------------------------------------------------------------------ - # Public interface - # ------------------------------------------------------------------ - - def subscribe(self, key: str) -> asyncio.Queue[dict]: - queue: asyncio.Queue[dict] = asyncio.Queue() - self._subscribers[key].append(queue) - - # Spin up a Redis listener if this is the first local subscriber - if key not in self._listener_tasks: - task = asyncio.get_running_loop().create_task(self._listen(key)) - self._listener_tasks[key] = task - - LOG.info("PubSub subscriber added", key=key, channel_prefix=self._channel_prefix) - return queue - - def unsubscribe(self, key: str, queue: asyncio.Queue[dict]) -> None: - queues = self._subscribers.get(key) - if queues: - try: - queues.remove(queue) - except ValueError: - pass - if not queues: - del self._subscribers[key] - self._cancel_listener(key) - LOG.info("PubSub subscriber removed", key=key, channel_prefix=self._channel_prefix) - - def publish(self, key: str, message: dict) -> None: - """Fire-and-forget Redis PUBLISH.""" - try: - loop = asyncio.get_running_loop() - loop.create_task(self._publish_to_redis(key, message)) - except RuntimeError: - LOG.warning( - "No running event loop; cannot publish via Redis", - key=key, - channel_prefix=self._channel_prefix, - ) - - async def close(self) -> None: - """Cancel all listener tasks and clear state. Call on shutdown.""" - for key in list(self._listener_tasks): - self._cancel_listener(key) - self._subscribers.clear() - LOG.info("RedisPubSub closed", channel_prefix=self._channel_prefix) - - # ------------------------------------------------------------------ - # Internal helpers - # ------------------------------------------------------------------ - - async def _publish_to_redis(self, key: str, message: dict) -> None: - channel = f"{self._channel_prefix}{key}" - try: - await self._redis.publish(channel, json.dumps(message)) - except Exception: - LOG.exception("Failed to publish to Redis", key=key, channel_prefix=self._channel_prefix) - - async def _listen(self, key: str) -> None: - """Subscribe to a Redis channel and fan out messages locally.""" - channel = f"{self._channel_prefix}{key}" - pubsub = self._redis.pubsub() - try: - await pubsub.subscribe(channel) - LOG.info("Redis listener started", channel=channel) - async for raw_message in pubsub.listen(): - if raw_message["type"] != "message": - continue - try: - data = json.loads(raw_message["data"]) - except (json.JSONDecodeError, TypeError): - LOG.warning("Invalid JSON on Redis channel", channel=channel) - continue - self._dispatch_local(key, data) - except asyncio.CancelledError: - LOG.info("Redis listener cancelled", channel=channel) - raise - except Exception: - LOG.exception("Redis listener error", channel=channel) - finally: - try: - await pubsub.unsubscribe(channel) - await pubsub.close() - except Exception: - LOG.warning("Error closing Redis pubsub", channel=channel) - - def _dispatch_local(self, key: str, message: dict) -> None: - """Fan out a message to all local asyncio queues for this key.""" - queues = self._subscribers.get(key, []) - for queue in queues: - try: - queue.put_nowait(message) - except asyncio.QueueFull: - LOG.warning( - "Queue full, dropping message", - key=key, - channel_prefix=self._channel_prefix, - ) - - def _cancel_listener(self, key: str) -> None: - task = self._listener_tasks.pop(key, None) - if task and not task.done(): - task.cancel() diff --git a/skyvern/forge/sdk/routes/__init__.py b/skyvern/forge/sdk/routes/__init__.py index 42e038d05..a24204fca 100644 --- a/skyvern/forge/sdk/routes/__init__.py +++ b/skyvern/forge/sdk/routes/__init__.py @@ -11,6 +11,5 @@ from skyvern.forge.sdk.routes import sdk # noqa: F401 from skyvern.forge.sdk.routes import webhooks # noqa: F401 from skyvern.forge.sdk.routes import workflow_copilot # noqa: F401 from skyvern.forge.sdk.routes.streaming import messages # noqa: F401 -from skyvern.forge.sdk.routes.streaming import notifications # noqa: F401 from skyvern.forge.sdk.routes.streaming import screenshot # noqa: F401 from skyvern.forge.sdk.routes.streaming import vnc # noqa: F401 diff --git a/skyvern/forge/sdk/routes/credentials.py b/skyvern/forge/sdk/routes/credentials.py index 940c85f4d..f2b31ba3c 100644 --- a/skyvern/forge/sdk/routes/credentials.py +++ b/skyvern/forge/sdk/routes/credentials.py @@ -157,15 +157,10 @@ async def send_totp_code( task = await app.DATABASE.get_task(data.task_id, curr_org.organization_id) if not task: raise HTTPException(status_code=400, detail=f"Invalid task id: {data.task_id}") - workflow_id_for_storage: str | None = None if data.workflow_id: - if data.workflow_id.startswith("wpid_"): - workflow = await app.DATABASE.get_workflow_by_permanent_id(data.workflow_id, curr_org.organization_id) - else: - workflow = await app.DATABASE.get_workflow(data.workflow_id, curr_org.organization_id) + workflow = await app.DATABASE.get_workflow(data.workflow_id, curr_org.organization_id) if not workflow: raise HTTPException(status_code=400, detail=f"Invalid workflow id: {data.workflow_id}") - workflow_id_for_storage = workflow.workflow_id if data.workflow_run_id: workflow_run = await app.DATABASE.get_workflow_run(data.workflow_run_id, curr_org.organization_id) if not workflow_run: @@ -193,7 +188,7 @@ async def send_totp_code( content=data.content, code=otp_value.value, task_id=data.task_id, - workflow_id=workflow_id_for_storage, + workflow_id=data.workflow_id, workflow_run_id=data.workflow_run_id, source=data.source, expired_at=data.expired_at, diff --git a/skyvern/forge/sdk/routes/streaming/notifications.py b/skyvern/forge/sdk/routes/streaming/notifications.py deleted file mode 100644 index 2a00c899d..000000000 --- a/skyvern/forge/sdk/routes/streaming/notifications.py +++ /dev/null @@ -1,128 +0,0 @@ -"""WebSocket endpoint for streaming global 2FA verification code notifications.""" - -import asyncio - -import structlog -from fastapi import WebSocket, WebSocketDisconnect -from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK - -from skyvern.forge import app -from skyvern.forge.sdk.notification.factory import NotificationRegistryFactory -from skyvern.forge.sdk.routes.routers import base_router, legacy_base_router -from skyvern.forge.sdk.routes.streaming.auth import auth - -LOG = structlog.get_logger() -HEARTBEAT_INTERVAL = 60 - - -@base_router.websocket("/stream/notifications") -async def notification_stream( - websocket: WebSocket, - apikey: str | None = None, - token: str | None = None, -) -> None: - return await _notification_stream_handler(websocket=websocket, apikey=apikey, token=token) - - -@legacy_base_router.websocket("/stream/notifications") -async def notification_stream_legacy( - websocket: WebSocket, - apikey: str | None = None, - token: str | None = None, -) -> None: - return await _notification_stream_handler(websocket=websocket, apikey=apikey, token=token) - - -async def _notification_stream_handler( - websocket: WebSocket, - apikey: str | None = None, - token: str | None = None, -) -> None: - organization_id = await auth(apikey=apikey, token=token, websocket=websocket) - if not organization_id: - LOG.info("Notifications: Authentication failed") - return - - LOG.info("Notifications: Started streaming", organization_id=organization_id) - registry = NotificationRegistryFactory.get_registry() - queue = registry.subscribe(organization_id) - - try: - # Send initial state: all currently active verification requests - active_requests = await app.DATABASE.get_active_verification_requests(organization_id) - for req in active_requests: - try: - await websocket.send_json( - { - "type": "verification_code_required", - "task_id": req.get("task_id"), - "workflow_run_id": req.get("workflow_run_id"), - "identifier": req.get("verification_code_identifier"), - "polling_started_at": req.get("verification_code_polling_started_at"), - } - ) - except (WebSocketDisconnect, ConnectionClosedOK, ConnectionClosedError, RuntimeError): - LOG.info( - "Notifications: Client disconnected during initial state send", - organization_id=organization_id, - ) - return - - # Watch for client disconnect while streaming events - disconnect_event = asyncio.Event() - - async def _watch_disconnect() -> None: - try: - while True: - await websocket.receive() - except (WebSocketDisconnect, ConnectionClosedOK, ConnectionClosedError): - disconnect_event.set() - - watcher = asyncio.create_task(_watch_disconnect()) - try: - while not disconnect_event.is_set(): - queue_task = asyncio.ensure_future(asyncio.wait_for(queue.get(), timeout=HEARTBEAT_INTERVAL)) - disconnect_wait = asyncio.ensure_future(disconnect_event.wait()) - done, pending = await asyncio.wait({queue_task, disconnect_wait}, return_when=asyncio.FIRST_COMPLETED) - for p in pending: - p.cancel() - - if disconnect_event.is_set(): - return - - try: - message = queue_task.result() - await websocket.send_json(message) - except TimeoutError: - try: - await websocket.send_json({"type": "heartbeat"}) - except Exception: - LOG.info( - "Notifications: Client unreachable during heartbeat. Closing.", - organization_id=organization_id, - ) - return - except asyncio.CancelledError: - return - except (WebSocketDisconnect, ConnectionClosedOK, ConnectionClosedError, RuntimeError): - LOG.info( - "Notifications: Client disconnected during send", - organization_id=organization_id, - ) - return - finally: - watcher.cancel() - - except WebSocketDisconnect: - LOG.info("Notifications: WebSocket disconnected", organization_id=organization_id) - except ConnectionClosedOK: - LOG.info("Notifications: ConnectionClosedOK", organization_id=organization_id) - except ConnectionClosedError: - LOG.warning( - "Notifications: ConnectionClosedError (client likely disconnected)", organization_id=organization_id - ) - except Exception: - LOG.warning("Notifications: Error while streaming", organization_id=organization_id, exc_info=True) - finally: - registry.unsubscribe(organization_id, queue) - LOG.info("Notifications: Connection closed", organization_id=organization_id) diff --git a/skyvern/forge/sdk/schemas/tasks.py b/skyvern/forge/sdk/schemas/tasks.py index 139891a49..a1a5918af 100644 --- a/skyvern/forge/sdk/schemas/tasks.py +++ b/skyvern/forge/sdk/schemas/tasks.py @@ -288,10 +288,6 @@ class Task(TaskBase): queued_at: datetime | None = None started_at: datetime | None = None finished_at: datetime | None = None - # 2FA verification code waiting state fields - waiting_for_verification_code: bool = False - verification_code_identifier: str | None = None - verification_code_polling_started_at: datetime | None = None @property def llm_key(self) -> str | None: @@ -369,9 +365,6 @@ class Task(TaskBase): max_screenshot_scrolls=self.max_screenshot_scrolls, step_count=step_count, browser_session_id=self.browser_session_id, - waiting_for_verification_code=self.waiting_for_verification_code, - verification_code_identifier=self.verification_code_identifier, - verification_code_polling_started_at=self.verification_code_polling_started_at, ) @@ -399,10 +392,6 @@ class TaskResponse(BaseModel): max_screenshot_scrolls: int | None = None step_count: int | None = None browser_session_id: str | None = None - # 2FA verification code waiting state fields - waiting_for_verification_code: bool = False - verification_code_identifier: str | None = None - verification_code_polling_started_at: datetime | None = None class TaskOutput(BaseModel): diff --git a/skyvern/forge/sdk/workflow/models/workflow.py b/skyvern/forge/sdk/workflow/models/workflow.py index bcc1d5349..41fc3dd87 100644 --- a/skyvern/forge/sdk/workflow/models/workflow.py +++ b/skyvern/forge/sdk/workflow/models/workflow.py @@ -168,10 +168,6 @@ class WorkflowRun(BaseModel): sequential_key: str | None = None ai_fallback: bool | None = None code_gen: bool | None = None - # 2FA verification code waiting state fields - waiting_for_verification_code: bool = False - verification_code_identifier: str | None = None - verification_code_polling_started_at: datetime | None = None queued_at: datetime | None = None started_at: datetime | None = None @@ -235,10 +231,6 @@ class WorkflowRunResponseBase(BaseModel): browser_address: str | None = None script_run: ScriptRunResponse | None = None errors: list[dict[str, Any]] | None = None - # 2FA verification code waiting state fields - waiting_for_verification_code: bool = False - verification_code_identifier: str | None = None - verification_code_polling_started_at: datetime | None = None class WorkflowRunWithWorkflowResponse(WorkflowRunResponseBase): diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py index 28c47c459..b121a80c6 100644 --- a/skyvern/forge/sdk/workflow/service.py +++ b/skyvern/forge/sdk/workflow/service.py @@ -3789,10 +3789,6 @@ class WorkflowService: browser_address=workflow_run.browser_address, script_run=workflow_run.script_run, errors=errors, - # 2FA verification code waiting state fields - waiting_for_verification_code=workflow_run.waiting_for_verification_code, - verification_code_identifier=workflow_run.verification_code_identifier, - verification_code_polling_started_at=workflow_run.verification_code_polling_started_at, ) async def clean_up_workflow( diff --git a/skyvern/services/otp_service.py b/skyvern/services/otp_service.py index 2c63600b1..fcd283994 100644 --- a/skyvern/services/otp_service.py +++ b/skyvern/services/otp_service.py @@ -1,9 +1,6 @@ import asyncio -import re -from dataclasses import dataclass from datetime import datetime, timedelta -import pyotp import structlog from pydantic import BaseModel, Field @@ -14,29 +11,9 @@ from skyvern.forge.prompts import prompt_engine from skyvern.forge.sdk.core.aiohttp_helper import aiohttp_post from skyvern.forge.sdk.core.security import generate_skyvern_webhook_signature from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType -from skyvern.forge.sdk.notification.factory import NotificationRegistryFactory from skyvern.forge.sdk.schemas.totp_codes import OTPType LOG = structlog.get_logger() -_MFA_PARAMETER_KEY_HINTS = ("mfa", "otp", "verification") -_NON_ALNUM_PATTERN = re.compile(r"[^a-z0-9]") - -MFANavigationPayload = dict | list | str | None - - -@dataclass(slots=True) -class OTPPollContext: - organization_id: str - task_id: str | None = None - workflow_id: str | None = None - workflow_run_id: str | None = None - workflow_permanent_id: str | None = None - totp_verification_url: str | None = None - totp_identifier: str | None = None - - @property - def needs_manual_input(self) -> bool: - return not self.totp_verification_url class OTPValue(BaseModel): @@ -59,62 +36,6 @@ class OTPResultParsedByLLM(BaseModel): otp_value: str | None = Field(None, description="The OTP value.") -def _is_mfa_like_parameter_key(key: object) -> bool: - """Return True when a payload key appears to represent an MFA/OTP parameter.""" - normalized_key = _NON_ALNUM_PATTERN.sub("", str(key).lower()) - return any(hint in normalized_key for hint in _MFA_PARAMETER_KEY_HINTS) - - -def extract_totp_from_navigation_inputs(navigation_payload: MFANavigationPayload) -> OTPValue | None: - """Extract TOTP from runtime navigation inputs. - - Runtime inline OTP extraction is intentionally payload-only. - """ - # Early return if payload is not a traversable container - if not isinstance(navigation_payload, (dict, list)): - return None - - traversal_stack: list[dict | list | str] = [navigation_payload] - visited_container_ids: set[int] = set() - - while traversal_stack: - current_item = traversal_stack.pop() - - # Only save OTP if a string - if isinstance(current_item, str): - return OTPValue(value=current_item, type=OTPType.TOTP) - - # Check if we've already visited this container to avoid cycles - current_id = id(current_item) - if current_id in visited_container_ids: - continue - visited_container_ids.add(current_id) - - # For lists, add all nested containers to the stack for further traversal - if isinstance(current_item, list): - for item in reversed(current_item): - if isinstance(item, (dict, list)): - traversal_stack.append(item) - continue - - # For dictionaries, process key-value pairs - for key, value in reversed(list(current_item.items())): - # Add nested containers (dicts/lists) to the stack regardless of key name - if isinstance(value, (dict, list)): - traversal_stack.append(value) - # Only process string values if the key suggests it's MFA-related - if not _is_mfa_like_parameter_key(key): - continue - if not isinstance(value, str): - continue - # If the value is a non-empty string under an MFA-like key, add it for evaluation - candidate_value = value.strip() - if candidate_value: - traversal_stack.append(candidate_value) - # No OTP code found after traversing the entire payload - return None - - async def parse_otp_login( content: str, organization_id: str, @@ -135,47 +56,6 @@ async def parse_otp_login( return None -def try_generate_totp_from_credential(workflow_run_id: str | None) -> OTPValue | None: - """Try to generate a TOTP code from a credential secret stored in the workflow run context. - - Scans workflow_run_context.values for credential entries with a "totp" key - (e.g. Bitwarden, 1Password, Azure Key Vault credentials) and generates a - TOTP code using pyotp. This should be checked BEFORE poll_otp_value so that - credential-based TOTP takes priority over webhook (totp_url) and totp_identifier. - """ - if not workflow_run_id: - return None - - workflow_run_context = app.WORKFLOW_CONTEXT_MANAGER.get_workflow_run_context(workflow_run_id) - if not workflow_run_context: - return None - - for key, value in workflow_run_context.values.items(): - if isinstance(value, dict) and "totp" in value: - totp_secret_id = value.get("totp") - if not totp_secret_id or not isinstance(totp_secret_id, str): - continue - totp_secret_key = workflow_run_context.totp_secret_value_key(totp_secret_id) - totp_secret = workflow_run_context.get_original_secret_value_or_none(totp_secret_key) - if totp_secret: - try: - code = pyotp.TOTP(totp_secret).now() - LOG.info( - "Generated TOTP from credential secret", - workflow_run_id=workflow_run_id, - credential_key=key, - ) - return OTPValue(value=code, type=OTPType.TOTP) - except Exception: - LOG.warning( - "Failed to generate TOTP from credential secret", - workflow_run_id=workflow_run_id, - credential_key=key, - exc_info=True, - ) - return None - - async def poll_otp_value( organization_id: str, task_id: str | None = None, @@ -185,194 +65,53 @@ async def poll_otp_value( totp_verification_url: str | None = None, totp_identifier: str | None = None, ) -> OTPValue | None: - ctx = OTPPollContext( - organization_id=organization_id, + timeout = timedelta(minutes=settings.VERIFICATION_CODE_POLLING_TIMEOUT_MINS) + start_datetime = datetime.utcnow() + timeout_datetime = start_datetime + timeout + org_token = await app.DATABASE.get_valid_org_auth_token(organization_id, OrganizationAuthTokenType.api.value) + if not org_token: + LOG.error("Failed to get organization token when trying to get otp value") + return None + LOG.info( + "Polling otp value", task_id=task_id, - workflow_id=workflow_id, workflow_run_id=workflow_run_id, workflow_permanent_id=workflow_permanent_id, totp_verification_url=totp_verification_url, totp_identifier=totp_identifier, ) - timeout = timedelta(minutes=settings.VERIFICATION_CODE_POLLING_TIMEOUT_MINS) - start_datetime = datetime.utcnow() - timeout_datetime = start_datetime + timeout - org_api_key: str | None = None - if ctx.totp_verification_url: - org_token = await app.DATABASE.get_valid_org_auth_token( - ctx.organization_id, OrganizationAuthTokenType.api.value - ) - if not org_token: - LOG.error("Failed to get organization token when trying to get otp value") - return None - org_api_key = org_token.token - LOG.info( - "Polling otp value", - task_id=ctx.task_id, - workflow_run_id=ctx.workflow_run_id, - workflow_permanent_id=ctx.workflow_permanent_id, - totp_verification_url=ctx.totp_verification_url, - totp_identifier=ctx.totp_identifier, - ) - - await _set_waiting_state(ctx, start_datetime) - - try: - while True: - await asyncio.sleep(10) - # check timeout - if datetime.utcnow() > timeout_datetime: - LOG.warning("Polling otp value timed out") - raise NoTOTPVerificationCodeFound( - task_id=ctx.task_id, - workflow_run_id=ctx.workflow_run_id, - workflow_id=ctx.workflow_permanent_id, - totp_verification_url=ctx.totp_verification_url, - totp_identifier=ctx.totp_identifier, - ) - otp_value: OTPValue | None = None - if ctx.totp_verification_url: - assert org_api_key is not None - otp_value = await _get_otp_value_from_url( - ctx.organization_id, - ctx.totp_verification_url, - org_api_key, - task_id=ctx.task_id, - workflow_run_id=ctx.workflow_run_id, - workflow_permanent_id=ctx.workflow_permanent_id, - ) - elif ctx.totp_identifier: - otp_value = await _get_otp_value_from_db( - ctx.organization_id, - ctx.totp_identifier, - task_id=ctx.task_id, - workflow_id=ctx.workflow_id, - workflow_run_id=ctx.workflow_run_id, - ) - if not otp_value: - otp_value = await _get_otp_value_by_run( - ctx.organization_id, - task_id=ctx.task_id, - workflow_run_id=ctx.workflow_run_id, - ) - else: - # No pre-configured TOTP — poll for manually submitted codes by run context - otp_value = await _get_otp_value_by_run( - ctx.organization_id, - task_id=ctx.task_id, - workflow_run_id=ctx.workflow_run_id, - ) - if otp_value: - LOG.info("Got otp value", otp_value=otp_value) - return otp_value - finally: - await _clear_waiting_state(ctx) - - -async def _set_waiting_state(ctx: OTPPollContext, started_at: datetime) -> None: - if not ctx.needs_manual_input: - return - - identifier_for_ui = ctx.totp_identifier - if ctx.workflow_run_id: - try: - await app.DATABASE.update_workflow_run( - workflow_run_id=ctx.workflow_run_id, - waiting_for_verification_code=True, - verification_code_identifier=identifier_for_ui, - verification_code_polling_started_at=started_at, + while True: + await asyncio.sleep(10) + # check timeout + if datetime.utcnow() > timeout_datetime: + LOG.warning("Polling otp value timed out") + raise NoTOTPVerificationCodeFound( + task_id=task_id, + workflow_run_id=workflow_run_id, + workflow_id=workflow_permanent_id, + totp_verification_url=totp_verification_url, + totp_identifier=totp_identifier, ) - LOG.info( - "Set 2FA waiting state for workflow run", - workflow_run_id=ctx.workflow_run_id, - verification_code_identifier=identifier_for_ui, + otp_value: OTPValue | None = None + if totp_verification_url: + otp_value = await _get_otp_value_from_url( + organization_id, + totp_verification_url, + org_token.token, + task_id=task_id, + workflow_run_id=workflow_run_id, ) - try: - NotificationRegistryFactory.get_registry().publish( - ctx.organization_id, - { - "type": "verification_code_required", - "workflow_run_id": ctx.workflow_run_id, - "task_id": ctx.task_id, - "identifier": identifier_for_ui, - "polling_started_at": started_at.isoformat(), - }, - ) - except Exception: - LOG.warning("Failed to publish 2FA required notification for workflow run", exc_info=True) - except Exception: - LOG.warning("Failed to set 2FA waiting state for workflow run", exc_info=True) - elif ctx.task_id: - try: - await app.DATABASE.update_task_2fa_state( - task_id=ctx.task_id, - organization_id=ctx.organization_id, - waiting_for_verification_code=True, - verification_code_identifier=identifier_for_ui, - verification_code_polling_started_at=started_at, + elif totp_identifier: + otp_value = await _get_otp_value_from_db( + organization_id, + totp_identifier, + task_id=task_id, + workflow_id=workflow_permanent_id, + workflow_run_id=workflow_run_id, ) - LOG.info( - "Set 2FA waiting state for task", - task_id=ctx.task_id, - verification_code_identifier=identifier_for_ui, - ) - try: - NotificationRegistryFactory.get_registry().publish( - ctx.organization_id, - { - "type": "verification_code_required", - "task_id": ctx.task_id, - "identifier": identifier_for_ui, - "polling_started_at": started_at.isoformat(), - }, - ) - except Exception: - LOG.warning("Failed to publish 2FA required notification for task", exc_info=True) - except Exception: - LOG.warning("Failed to set 2FA waiting state for task", exc_info=True) - - -async def _clear_waiting_state(ctx: OTPPollContext) -> None: - if not ctx.needs_manual_input: - return - - if ctx.workflow_run_id: - try: - await app.DATABASE.update_workflow_run( - workflow_run_id=ctx.workflow_run_id, - waiting_for_verification_code=False, - ) - LOG.info("Cleared 2FA waiting state for workflow run", workflow_run_id=ctx.workflow_run_id) - try: - NotificationRegistryFactory.get_registry().publish( - ctx.organization_id, - { - "type": "verification_code_resolved", - "workflow_run_id": ctx.workflow_run_id, - "task_id": ctx.task_id, - }, - ) - except Exception: - LOG.warning("Failed to publish 2FA resolved notification for workflow run", exc_info=True) - except Exception: - LOG.warning("Failed to clear 2FA waiting state for workflow run", exc_info=True) - elif ctx.task_id: - try: - await app.DATABASE.update_task_2fa_state( - task_id=ctx.task_id, - organization_id=ctx.organization_id, - waiting_for_verification_code=False, - ) - LOG.info("Cleared 2FA waiting state for task", task_id=ctx.task_id) - try: - NotificationRegistryFactory.get_registry().publish( - ctx.organization_id, - {"type": "verification_code_resolved", "task_id": ctx.task_id}, - ) - except Exception: - LOG.warning("Failed to publish 2FA resolved notification for task", exc_info=True) - except Exception: - LOG.warning("Failed to clear 2FA waiting state for task", exc_info=True) + if otp_value: + LOG.info("Got otp value", otp_value=otp_value) + return otp_value async def _get_otp_value_from_url( @@ -436,28 +175,6 @@ async def _get_otp_value_from_url( return otp_value -async def _get_otp_value_by_run( - organization_id: str, - task_id: str | None = None, - workflow_run_id: str | None = None, -) -> OTPValue | None: - """Look up OTP codes by task_id/workflow_run_id when no totp_identifier is configured. - - Used for the manual 2FA input flow where users submit codes through the UI - without pre-configured TOTP credentials. - """ - codes = await app.DATABASE.get_otp_codes_by_run( - organization_id=organization_id, - task_id=task_id, - workflow_run_id=workflow_run_id, - limit=1, - ) - if codes: - code = codes[0] - return OTPValue(value=code.code, type=code.otp_type) - return None - - async def _get_otp_value_from_db( organization_id: str, totp_identifier: str, diff --git a/skyvern/webeye/actions/parse_actions.py b/skyvern/webeye/actions/parse_actions.py index 574f8d955..40785d9bd 100644 --- a/skyvern/webeye/actions/parse_actions.py +++ b/skyvern/webeye/actions/parse_actions.py @@ -14,11 +14,7 @@ from skyvern.forge.sdk.core import skyvern_context from skyvern.forge.sdk.models import Step from skyvern.forge.sdk.schemas.tasks import Task from skyvern.forge.sdk.schemas.totp_codes import OTPType -from skyvern.services.otp_service import ( - extract_totp_from_navigation_inputs, - poll_otp_value, - try_generate_totp_from_credential, -) +from skyvern.services.otp_service import poll_otp_value from skyvern.utils.image_resizer import Resolution, scale_coordinates from skyvern.webeye.actions.action_types import ActionType from skyvern.webeye.actions.actions import ( @@ -917,13 +913,7 @@ async def generate_cua_fallback_actions( ) elif skyvern_action_type == "get_verification_code": - # 1. Check navigation payload first for inline OTP. - otp_value = extract_totp_from_navigation_inputs(task.navigation_payload) - # 2. Then try to generate TOTP from credential if payload has no OTP. - if not otp_value: - otp_value = try_generate_totp_from_credential(task.workflow_run_id) - # 3. Lastly, poll for OTP if organization has config and no OTP was found yet. - if not otp_value and task.organization_id: + if (task.totp_verification_url or task.totp_identifier) and task.organization_id: LOG.info( "Getting verification code for CUA", task_id=task.task_id, @@ -940,23 +930,29 @@ async def generate_cua_fallback_actions( totp_verification_url=task.totp_verification_url, totp_identifier=task.totp_identifier, ) + if not otp_value or otp_value.get_otp_type() != OTPType.TOTP: + raise NoTOTPVerificationCodeFound() + verification_code = otp_value.value + reasoning = reasoning or f"Received verification code: {verification_code}" + action = VerificationCodeAction( + verification_code=verification_code, + reasoning=reasoning, + intention=reasoning, + ) except NoTOTPVerificationCodeFound: reasoning_suffix = "No verification code found" reasoning = f"{reasoning}. {reasoning_suffix}" if reasoning else reasoning_suffix + action = TerminateAction( + reasoning=reasoning, + intention=reasoning, + ) except FailedToGetTOTPVerificationCode as e: reasoning_suffix = f"Failed to get verification code. Reason: {e.reason}" reasoning = f"{reasoning}. {reasoning_suffix}" if reasoning else reasoning_suffix - - # OTP value obtained - use it as verification code - if otp_value and otp_value.get_otp_type() == OTPType.TOTP: - verification_code = otp_value.value - reasoning = reasoning or f"Received verification code: {verification_code}" - action = VerificationCodeAction( - verification_code=verification_code, - reasoning=reasoning, - intention=reasoning, - ) - # Terminate the task since OTP is necessary but wasn't found/generated + action = TerminateAction( + reasoning=reasoning, + intention=reasoning, + ) else: action = TerminateAction( reasoning=reasoning, diff --git a/tests/unit/test_credential_totp_priority.py b/tests/unit/test_credential_totp_priority.py deleted file mode 100644 index d67ba8dd7..000000000 --- a/tests/unit/test_credential_totp_priority.py +++ /dev/null @@ -1,395 +0,0 @@ -"""Tests for credential TOTP priority over webhook (totp_url) and totp_identifier. - -Verifies that try_generate_totp_from_credential() correctly generates TOTP codes -from credential secrets stored in workflow run context, and that callers check -credential TOTP before falling back to poll_otp_value. -""" - -from unittest.mock import AsyncMock, MagicMock, patch - -import pyotp -import pytest - -from skyvern.forge.sdk.schemas.totp_codes import OTPType -from skyvern.services.otp_service import ( - OTPValue, - _get_otp_value_from_url, - poll_otp_value, - try_generate_totp_from_credential, -) - -# A valid base32 TOTP secret for testing -TEST_TOTP_SECRET = "JBSWY3DPEHPK3PXP" - - -def _make_workflow_run_context( - values: dict | None = None, - secrets: dict | None = None, -) -> MagicMock: - """Create a mock WorkflowRunContext with the given values and secrets.""" - ctx = MagicMock() - ctx.values = values or {} - ctx.secrets = secrets or {} - - def totp_secret_value_key(totp_secret_id: str) -> str: - return f"{totp_secret_id}_value" - - ctx.totp_secret_value_key = totp_secret_value_key - - def get_original_secret_value_or_none(secret_key: str) -> str | None: - return ctx.secrets.get(secret_key) - - ctx.get_original_secret_value_or_none = get_original_secret_value_or_none - return ctx - - -class TestTryGenerateTotpFromCredential: - """Tests for the try_generate_totp_from_credential helper.""" - - def test_returns_none_when_workflow_run_id_is_none(self) -> None: - result = try_generate_totp_from_credential(None) - assert result is None - - def test_returns_none_when_no_workflow_run_context(self) -> None: - with patch("skyvern.services.otp_service.app") as mock_app: - mock_app.WORKFLOW_CONTEXT_MANAGER.get_workflow_run_context.return_value = None - result = try_generate_totp_from_credential("wfr_123") - assert result is None - - def test_returns_none_when_no_credential_values(self) -> None: - ctx = _make_workflow_run_context(values={"some_param": "plain_string"}) - with patch("skyvern.services.otp_service.app") as mock_app: - mock_app.WORKFLOW_CONTEXT_MANAGER.get_workflow_run_context.return_value = ctx - result = try_generate_totp_from_credential("wfr_123") - assert result is None - - def test_returns_none_when_dict_value_has_no_totp_key(self) -> None: - ctx = _make_workflow_run_context( - values={"cred_param": {"username": "user", "password": "pass"}}, - ) - with patch("skyvern.services.otp_service.app") as mock_app: - mock_app.WORKFLOW_CONTEXT_MANAGER.get_workflow_run_context.return_value = ctx - result = try_generate_totp_from_credential("wfr_123") - assert result is None - - def test_returns_none_when_totp_secret_id_is_empty(self) -> None: - ctx = _make_workflow_run_context( - values={"cred_param": {"totp": ""}}, - ) - with patch("skyvern.services.otp_service.app") as mock_app: - mock_app.WORKFLOW_CONTEXT_MANAGER.get_workflow_run_context.return_value = ctx - result = try_generate_totp_from_credential("wfr_123") - assert result is None - - def test_returns_none_when_totp_secret_not_in_secrets(self) -> None: - """When the secret ID doesn't resolve to an actual secret value.""" - ctx = _make_workflow_run_context( - values={"cred_param": {"totp": "secret_id_123"}}, - secrets={}, # no secret stored - ) - with patch("skyvern.services.otp_service.app") as mock_app: - mock_app.WORKFLOW_CONTEXT_MANAGER.get_workflow_run_context.return_value = ctx - result = try_generate_totp_from_credential("wfr_123") - assert result is None - - def test_generates_totp_from_credential_secret(self) -> None: - """Happy path: credential with valid TOTP secret generates a code.""" - ctx = _make_workflow_run_context( - values={"cred_param": {"username": "user", "password": "pass", "totp": "totp_ref_1"}}, - secrets={"totp_ref_1_value": TEST_TOTP_SECRET}, - ) - with patch("skyvern.services.otp_service.app") as mock_app: - mock_app.WORKFLOW_CONTEXT_MANAGER.get_workflow_run_context.return_value = ctx - result = try_generate_totp_from_credential("wfr_123") - - assert result is not None - assert isinstance(result, OTPValue) - assert result.type == OTPType.TOTP - # Verify the code matches what pyotp would generate - expected_code = pyotp.TOTP(TEST_TOTP_SECRET).now() - assert result.value == expected_code - - def test_returns_first_matching_credential(self) -> None: - """When multiple credentials have TOTP, returns the first one found.""" - ctx = _make_workflow_run_context( - values={ - "cred_a": {"totp": "ref_a"}, - "cred_b": {"totp": "ref_b"}, - }, - secrets={ - "ref_a_value": TEST_TOTP_SECRET, - "ref_b_value": "ORSXG5DJNZTQ====", - }, - ) - with patch("skyvern.services.otp_service.app") as mock_app: - mock_app.WORKFLOW_CONTEXT_MANAGER.get_workflow_run_context.return_value = ctx - result = try_generate_totp_from_credential("wfr_123") - - assert result is not None - assert result.value == pyotp.TOTP(TEST_TOTP_SECRET).now() - - def test_skips_invalid_secret_and_continues(self) -> None: - """If one credential has an invalid TOTP secret, skip it and try the next.""" - ctx = _make_workflow_run_context( - values={ - "cred_bad": {"totp": "ref_bad"}, - "cred_good": {"totp": "ref_good"}, - }, - secrets={ - "ref_bad_value": "NOT_A_VALID_BASE32!!!", - "ref_good_value": TEST_TOTP_SECRET, - }, - ) - with patch("skyvern.services.otp_service.app") as mock_app: - mock_app.WORKFLOW_CONTEXT_MANAGER.get_workflow_run_context.return_value = ctx - result = try_generate_totp_from_credential("wfr_123") - - assert result is not None - assert result.value == pyotp.TOTP(TEST_TOTP_SECRET).now() - - def test_skips_non_string_totp_id(self) -> None: - """If the totp value is not a string (e.g., int or None), skip it.""" - ctx = _make_workflow_run_context( - values={"cred_param": {"totp": 12345}}, - ) - with patch("skyvern.services.otp_service.app") as mock_app: - mock_app.WORKFLOW_CONTEXT_MANAGER.get_workflow_run_context.return_value = ctx - result = try_generate_totp_from_credential("wfr_123") - assert result is None - - def test_skips_non_dict_values(self) -> None: - """Non-dict values in the context should be ignored.""" - ctx = _make_workflow_run_context( - values={ - "string_param": "hello", - "int_param": 42, - "list_param": [1, 2, 3], - "none_param": None, - }, - ) - with patch("skyvern.services.otp_service.app") as mock_app: - mock_app.WORKFLOW_CONTEXT_MANAGER.get_workflow_run_context.return_value = ctx - result = try_generate_totp_from_credential("wfr_123") - assert result is None - - -class TestPollOtpCalledWithoutTotpConfig: - """Tests that poll_otp_value is called even when totp_verification_url and - totp_identifier are both None — i.e. the manual 2FA code submission path.""" - - @pytest.mark.asyncio - async def test_poll_otp_called_when_no_totp_config_cua_fallback(self) -> None: - """When credential TOTP returns None and no totp_verification_url or - totp_identifier is set, poll_otp_value should still be called via - generate_cua_fallback_actions so the manual code submission flow works.""" - mock_task = MagicMock() - mock_task.totp_verification_url = None - mock_task.totp_identifier = None - mock_task.organization_id = "org_123" - mock_task.task_id = "task_456" - mock_task.workflow_run_id = None - mock_task.navigation_goal = "test goal" - - mock_step = MagicMock() - mock_step.step_id = "step_789" - mock_step.order = 0 - - mock_otp_value = MagicMock(spec=OTPValue) - mock_otp_value.get_otp_type.return_value = OTPType.TOTP - mock_otp_value.value = "123456" - - with ( - patch( - "skyvern.webeye.actions.parse_actions.try_generate_totp_from_credential", - return_value=None, - ), - patch( - "skyvern.webeye.actions.parse_actions.poll_otp_value", - new_callable=AsyncMock, - return_value=mock_otp_value, - ) as mock_poll, - patch("skyvern.webeye.actions.parse_actions.app") as mock_app, - patch("skyvern.webeye.actions.parse_actions.prompt_engine") as mock_prompt_engine, - ): - mock_prompt_engine.load_prompt.return_value = "test prompt" - # LLM returns get_verification_code action - mock_app.LLM_API_HANDLER = AsyncMock( - return_value={"action": "get_verification_code", "useful_information": "Need 2FA code"}, - ) - - from skyvern.webeye.actions.parse_actions import generate_cua_fallback_actions - - actions = await generate_cua_fallback_actions( - task=mock_task, - step=mock_step, - assistant_message="Enter verification code", - reasoning="Need 2FA code", - ) - - # poll_otp_value should have been called even though - # totp_verification_url and totp_identifier are both None - mock_poll.assert_called_once_with( - organization_id="org_123", - task_id="task_456", - workflow_run_id=None, - totp_verification_url=None, - totp_identifier=None, - ) - - # Verify we got a VerificationCodeAction back - assert len(actions) == 1 - from skyvern.webeye.actions.actions import VerificationCodeAction - - assert isinstance(actions[0], VerificationCodeAction) - assert actions[0].verification_code == "123456" - - @pytest.mark.asyncio - async def test_poll_otp_called_when_no_totp_config_agent(self) -> None: - """When credential TOTP returns None and no totp_verification_url or - totp_identifier is set, poll_otp_value should still be called via the - agent's handle_potential_verification_code so the manual code submission - flow works.""" - # Return None from poll_otp_value so we hit the early return at line 4548 - # (no valid OTP) — this avoids needing to mock the deeper context/LLM calls - with ( - patch( - "skyvern.forge.agent.try_generate_totp_from_credential", - return_value=None, - ), - patch( - "skyvern.forge.agent.poll_otp_value", - new_callable=AsyncMock, - return_value=None, - ) as mock_poll, - patch("skyvern.forge.agent.app") as mock_app, - ): - mock_app.DATABASE.get_workflow_run = AsyncMock(return_value=None) - - from skyvern.forge.agent import ForgeAgent - - agent = ForgeAgent.__new__(ForgeAgent) - - mock_task = MagicMock() - mock_task.totp_verification_url = None - mock_task.totp_identifier = None - mock_task.organization_id = "org_123" - mock_task.task_id = "task_456" - mock_task.workflow_run_id = None - - json_response = { - "place_to_enter_verification_code": True, - "should_enter_verification_code": True, - } - - result = await agent.handle_potential_verification_code( - task=mock_task, - step=MagicMock(), - scraped_page=MagicMock(), - browser_state=MagicMock(), - json_response=json_response, - ) - - # poll_otp_value should have been called even though - # totp_verification_url and totp_identifier are both None - mock_poll.assert_called_once_with( - organization_id="org_123", - task_id="task_456", - workflow_id=None, - workflow_run_id=None, - workflow_permanent_id=None, - totp_verification_url=None, - totp_identifier=None, - ) - - # When poll_otp_value returns None, the method returns json_response unchanged - assert result == json_response - - -class TestWorkflowPermanentIdPassedToWebhook: - """Tests that workflow_permanent_id is included in the TOTP webhook POST body.""" - - @pytest.mark.asyncio - async def test_get_otp_value_from_url_includes_workflow_permanent_id(self) -> None: - """_get_otp_value_from_url should include workflow_permanent_id in the request payload.""" - with ( - patch("skyvern.services.otp_service.generate_skyvern_webhook_signature") as mock_sign, - patch("skyvern.services.otp_service.aiohttp_post", new_callable=AsyncMock) as mock_post, - ): - mock_sign.return_value = MagicMock(signed_payload="signed", headers={"x-sig": "abc"}) - mock_post.return_value = {"verification_code": "123456"} - - result = await _get_otp_value_from_url( - organization_id="org_1", - url="https://example.com/totp", - api_key="key_1", - task_id="tsk_1", - workflow_run_id="wr_1", - workflow_permanent_id="wpid_1", - ) - - # Verify the payload includes workflow_permanent_id - mock_sign.assert_called_once() - payload = mock_sign.call_args.kwargs["payload"] - assert payload == { - "task_id": "tsk_1", - "workflow_run_id": "wr_1", - "workflow_permanent_id": "wpid_1", - } - assert result is not None - assert result.value == "123456" - - @pytest.mark.asyncio - async def test_get_otp_value_from_url_omits_workflow_permanent_id_when_none(self) -> None: - """When workflow_permanent_id is None, it should not appear in the payload.""" - with ( - patch("skyvern.services.otp_service.generate_skyvern_webhook_signature") as mock_sign, - patch("skyvern.services.otp_service.aiohttp_post", new_callable=AsyncMock) as mock_post, - ): - mock_sign.return_value = MagicMock(signed_payload="signed", headers={"x-sig": "abc"}) - mock_post.return_value = {"verification_code": "123456"} - - await _get_otp_value_from_url( - organization_id="org_1", - url="https://example.com/totp", - api_key="key_1", - task_id="tsk_1", - workflow_run_id="wr_1", - ) - - payload = mock_sign.call_args.kwargs["payload"] - assert "workflow_permanent_id" not in payload - - @pytest.mark.asyncio - async def test_poll_otp_passes_workflow_permanent_id_to_url_handler(self) -> None: - """poll_otp_value should forward workflow_permanent_id to _get_otp_value_from_url.""" - with ( - patch("skyvern.services.otp_service.app") as mock_app, - patch("skyvern.services.otp_service._set_waiting_state", new_callable=AsyncMock), - patch("skyvern.services.otp_service._clear_waiting_state", new_callable=AsyncMock), - patch("skyvern.services.otp_service._get_otp_value_from_url", new_callable=AsyncMock) as mock_url_handler, - patch("skyvern.services.otp_service.asyncio.sleep", new_callable=AsyncMock), - ): - mock_org_token = MagicMock() - mock_org_token.token = "api_key_123" - mock_app.DATABASE.get_valid_org_auth_token = AsyncMock(return_value=mock_org_token) - - mock_url_handler.return_value = OTPValue(value="123456", type=OTPType.TOTP) - - result = await poll_otp_value( - organization_id="org_1", - task_id="tsk_1", - workflow_run_id="wr_1", - workflow_permanent_id="wpid_1", - totp_verification_url="https://example.com/totp", - ) - - mock_url_handler.assert_called_once_with( - "org_1", - "https://example.com/totp", - "api_key_123", - task_id="tsk_1", - workflow_run_id="wr_1", - workflow_permanent_id="wpid_1", - ) - assert result is not None - assert result.value == "123456" diff --git a/tests/unit_tests/test_notification_registry.py b/tests/unit_tests/test_notification_registry.py deleted file mode 100644 index 1654740ac..000000000 --- a/tests/unit_tests/test_notification_registry.py +++ /dev/null @@ -1,106 +0,0 @@ -"""Tests for NotificationRegistry pub/sub and get_active_verification_requests (SKY-6).""" - -import pytest - -from skyvern.forge.sdk.db.agent_db import AgentDB -from skyvern.forge.sdk.notification.base import BaseNotificationRegistry -from skyvern.forge.sdk.notification.factory import NotificationRegistryFactory -from skyvern.forge.sdk.notification.local import LocalNotificationRegistry - -# === Task 1: NotificationRegistry subscribe / publish / unsubscribe === - - -@pytest.mark.asyncio -async def test_subscribe_and_publish(): - """Published messages should be received by subscribers.""" - registry = LocalNotificationRegistry() - queue = registry.subscribe("org_1") - - registry.publish("org_1", {"type": "verification_code_required", "task_id": "tsk_1"}) - msg = queue.get_nowait() - assert msg["type"] == "verification_code_required" - assert msg["task_id"] == "tsk_1" - - -@pytest.mark.asyncio -async def test_multiple_subscribers(): - """All subscribers for an org should receive the same message.""" - registry = LocalNotificationRegistry() - q1 = registry.subscribe("org_1") - q2 = registry.subscribe("org_1") - - registry.publish("org_1", {"type": "verification_code_required"}) - assert not q1.empty() - assert not q2.empty() - assert q1.get_nowait() == q2.get_nowait() - - -@pytest.mark.asyncio -async def test_publish_wrong_org_does_not_leak(): - """Messages for org_A should not appear in org_B's queue.""" - registry = LocalNotificationRegistry() - q_a = registry.subscribe("org_a") - q_b = registry.subscribe("org_b") - - registry.publish("org_a", {"type": "test"}) - assert not q_a.empty() - assert q_b.empty() - - -@pytest.mark.asyncio -async def test_unsubscribe(): - """After unsubscribe, the queue should no longer receive messages.""" - registry = LocalNotificationRegistry() - queue = registry.subscribe("org_1") - - registry.unsubscribe("org_1", queue) - registry.publish("org_1", {"type": "test"}) - assert queue.empty() - - -@pytest.mark.asyncio -async def test_unsubscribe_idempotent(): - """Unsubscribing a queue that's already removed should not raise.""" - registry = LocalNotificationRegistry() - queue = registry.subscribe("org_1") - registry.unsubscribe("org_1", queue) - registry.unsubscribe("org_1", queue) # should not raise - - -# === Task: BaseNotificationRegistry ABC === - - -def test_base_notification_registry_cannot_be_instantiated(): - """ABC should not be directly instantiable.""" - with pytest.raises(TypeError): - BaseNotificationRegistry() - - -# === Task: NotificationRegistryFactory === - - -@pytest.mark.asyncio -async def test_factory_returns_local_by_default(): - """Factory should return a LocalNotificationRegistry by default.""" - registry = NotificationRegistryFactory.get_registry() - assert isinstance(registry, LocalNotificationRegistry) - - -@pytest.mark.asyncio -async def test_factory_set_and_get(): - """Factory should allow swapping the registry implementation.""" - original = NotificationRegistryFactory.get_registry() - try: - custom = LocalNotificationRegistry() - NotificationRegistryFactory.set_registry(custom) - assert NotificationRegistryFactory.get_registry() is custom - finally: - NotificationRegistryFactory.set_registry(original) - - -# === Task 2: get_active_verification_requests DB method === - - -def test_get_active_verification_requests_method_exists(): - """AgentDB should have get_active_verification_requests method.""" - assert hasattr(AgentDB, "get_active_verification_requests") diff --git a/tests/unit_tests/test_otp_no_config.py b/tests/unit_tests/test_otp_no_config.py deleted file mode 100644 index 930ed19e2..000000000 --- a/tests/unit_tests/test_otp_no_config.py +++ /dev/null @@ -1,736 +0,0 @@ -"""Tests for manual 2FA input without pre-configured TOTP credentials (SKY-6).""" - -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -from skyvern.constants import SPECIAL_FIELD_VERIFICATION_CODE -from skyvern.forge.agent import ForgeAgent -from skyvern.forge.sdk.db.agent_db import AgentDB -from skyvern.forge.sdk.notification.local import LocalNotificationRegistry -from skyvern.forge.sdk.routes.credentials import send_totp_code -from skyvern.forge.sdk.schemas.totp_codes import OTPType, TOTPCodeCreate -from skyvern.schemas.runs import RunEngine -from skyvern.services.otp_service import ( - OTPValue, - _get_otp_value_by_run, - extract_totp_from_navigation_inputs, - poll_otp_value, -) - - -@pytest.mark.asyncio -async def test_get_otp_codes_by_run_exists(): - """get_otp_codes_by_run should exist on AgentDB.""" - assert hasattr(AgentDB, "get_otp_codes_by_run"), "AgentDB missing get_otp_codes_by_run method" - - -@pytest.mark.asyncio -async def test_get_otp_codes_by_run_returns_empty_without_identifiers(): - """get_otp_codes_by_run should return [] when neither task_id nor workflow_run_id is given.""" - db = AgentDB.__new__(AgentDB) - result = await db.get_otp_codes_by_run( - organization_id="org_1", - ) - assert result == [] - - -# === Task 2: _get_otp_value_by_run OTP service function === - - -@pytest.mark.asyncio -async def test_get_otp_value_by_run_returns_code(): - """_get_otp_value_by_run should find OTP codes by task_id.""" - mock_code = MagicMock() - mock_code.code = "123456" - mock_code.otp_type = "totp" - - mock_db = AsyncMock() - mock_db.get_otp_codes_by_run.return_value = [mock_code] - - mock_app = MagicMock() - mock_app.DATABASE = mock_db - - with patch("skyvern.services.otp_service.app", new=mock_app): - result = await _get_otp_value_by_run( - organization_id="org_1", - task_id="tsk_1", - ) - assert result is not None - assert result.value == "123456" - - -@pytest.mark.asyncio -async def test_get_otp_value_by_run_returns_none_when_no_codes(): - """_get_otp_value_by_run should return None when no codes found.""" - mock_db = AsyncMock() - mock_db.get_otp_codes_by_run.return_value = [] - - mock_app = MagicMock() - mock_app.DATABASE = mock_db - - with patch("skyvern.services.otp_service.app", new=mock_app): - result = await _get_otp_value_by_run( - organization_id="org_1", - task_id="tsk_1", - ) - assert result is None - - -def test_extract_totp_from_navigation_inputs_prefers_payload_code_over_goal_text(): - """Payload alias code should be used even if goal text contains another code.""" - otp_value = extract_totp_from_navigation_inputs( - {"mfaChoice": "520265"}, - ) - - assert otp_value is not None - assert otp_value.value == "520265" - assert otp_value.get_otp_type() == OTPType.TOTP - - -def test_extract_totp_from_navigation_inputs_ignores_navigation_goal_when_payload_missing(): - """Goal text alone should not produce inline OTP in payload-only mode.""" - otp_value = extract_totp_from_navigation_inputs( - None, - ) - - assert otp_value is None - - -def test_extract_totp_from_navigation_inputs_ignores_goal_input_action_when_payload_missing(): - """Input-action goal text should be ignored in payload-only mode.""" - otp_value = extract_totp_from_navigation_inputs( - {}, - ) - - assert otp_value is None - - -def test_extract_totp_from_navigation_inputs_no_code_anywhere(): - """No code in payload or goal should return None.""" - otp_value = extract_totp_from_navigation_inputs( - None, - ) - - assert otp_value is None - - -# === Task 3: poll_otp_value without identifier === - - -@pytest.mark.asyncio -async def test_poll_otp_value_without_identifier_uses_run_lookup(): - """poll_otp_value should use _get_otp_value_by_run when no identifier/URL provided.""" - mock_code = MagicMock() - mock_code.code = "123456" - mock_code.otp_type = "totp" - - mock_db = AsyncMock() - mock_db.get_valid_org_auth_token.return_value = MagicMock(token="tok") - mock_db.get_otp_codes_by_run.return_value = [mock_code] - mock_db.update_task_2fa_state = AsyncMock() - - mock_app = MagicMock() - mock_app.DATABASE = mock_db - - with ( - patch("skyvern.services.otp_service.app", new=mock_app), - patch("skyvern.services.otp_service.asyncio.sleep", new_callable=AsyncMock), - ): - result = await poll_otp_value( - organization_id="org_1", - task_id="tsk_1", - ) - assert result is not None - assert result.value == "123456" - - -# === Task 6: Integration test — handle_potential_OTP_actions without TOTP config === - - -@pytest.mark.asyncio -async def test_handle_potential_OTP_actions_without_totp_config(): - """When LLM detects 2FA but no TOTP config exists, should still enter verification flow.""" - agent = ForgeAgent.__new__(ForgeAgent) - - task = MagicMock() - task.organization_id = "org_1" - task.totp_verification_url = None - task.totp_identifier = None - task.task_id = "tsk_1" - task.workflow_run_id = None - - step = MagicMock() - scraped_page = MagicMock() - browser_state = MagicMock() - - json_response = { - "should_enter_verification_code": True, - "place_to_enter_verification_code": "input#otp-code", - "actions": [], - } - - with patch.object(agent, "handle_potential_verification_code", new_callable=AsyncMock) as mock_handler: - mock_handler.return_value = {"actions": []} - with patch("skyvern.forge.agent.parse_actions", return_value=[]): - result_json, result_actions = await agent.handle_potential_OTP_actions( - task, step, scraped_page, browser_state, json_response - ) - mock_handler.assert_called_once() - - -@pytest.mark.asyncio -async def test_handle_potential_verification_code_uses_navigation_payload_and_skips_poll(): - """When payload includes MFA code, should consume it and skip poll_otp_value.""" - from skyvern.forge import agent as forge_agent_module - - agent = ForgeAgent.__new__(ForgeAgent) - - task = MagicMock() - task.organization_id = "org_1" - task.totp_verification_url = None - task.totp_identifier = None - task.task_id = "tsk_payload_code" - task.workflow_run_id = None - task.navigation_payload = {"mfaChoice": "520265"} - task.llm_key = None - - step = MagicMock() - step.step_id = "step_1" - step.order = 0 - - scraped_page = MagicMock() - scraped_page.screenshots = [] - browser_state = MagicMock() - json_response = { - "should_enter_verification_code": True, - "place_to_enter_verification_code": "input#otp-code", - "actions": [ - {"action_type": "INPUT_TEXT", "id": "AAAb", "reasoning": "Enter MFA code", "text": "520265"}, - ], - } - - original_app_inst = object.__getattribute__(forge_agent_module.app, "_inst") - object.__setattr__(forge_agent_module.app, "_inst", MagicMock(LLM_API_HANDLER=AsyncMock())) - try: - with ( - patch("skyvern.forge.agent.try_generate_totp_from_credential") as mock_credential_totp, - patch("skyvern.forge.agent.poll_otp_value", new_callable=AsyncMock) as mock_poll, - patch("skyvern.forge.agent.skyvern_context") as mock_skyvern_context, - patch("skyvern.forge.agent.service_utils.is_cua_task", new_callable=AsyncMock, return_value=False), - patch( - "skyvern.forge.agent.LLMAPIHandlerFactory.get_override_llm_api_handler", - return_value=AsyncMock(return_value={"actions": []}), - ), - patch.object( - agent, - "_build_extract_action_prompt", - new_callable=AsyncMock, - return_value=("prompt", False, "extract-actions"), - ), - ): - mock_context = MagicMock() - mock_context.totp_codes = {} - mock_skyvern_context.ensure_context.return_value = mock_context - mock_skyvern_context.current.return_value = mock_context - - await agent.handle_potential_verification_code( - task=task, - step=step, - scraped_page=scraped_page, - browser_state=browser_state, - json_response=json_response, - ) - finally: - object.__setattr__(forge_agent_module.app, "_inst", original_app_inst) - - mock_credential_totp.assert_not_called() - mock_poll.assert_not_called() - # Early return path: code already visible to LLM in payload, no context storage or re-prompt needed - - -@pytest.mark.asyncio -async def test_handle_potential_OTP_actions_skips_magic_link_without_totp_config(): - """Magic links should still require TOTP config.""" - agent = ForgeAgent.__new__(ForgeAgent) - - task = MagicMock() - task.organization_id = "org_1" - task.totp_verification_url = None - task.totp_identifier = None - - step = MagicMock() - scraped_page = MagicMock() - browser_state = MagicMock() - - json_response = { - "should_verify_by_magic_link": True, - } - - with patch.object(agent, "handle_potential_magic_link", new_callable=AsyncMock) as mock_handler: - result_json, result_actions = await agent.handle_potential_OTP_actions( - task, step, scraped_page, browser_state, json_response - ) - mock_handler.assert_not_called() - assert result_actions == [] - - -# === Task 7: verification_code_check always True in LLM prompt === - - -@pytest.mark.asyncio -async def test_verification_code_check_always_true_without_totp_config(): - """_build_extract_action_prompt should receive verification_code_check=True even when task has no TOTP config.""" - agent = ForgeAgent.__new__(ForgeAgent) - agent.async_operation_pool = MagicMock() - - task = MagicMock() - task.totp_verification_url = None - task.totp_identifier = None - task.task_id = "tsk_1" - task.workflow_run_id = None - task.organization_id = "org_1" - task.url = "https://example.com" - - step = MagicMock() - step.step_id = "step_1" - step.order = 0 - step.retry_index = 0 - - scraped_page = MagicMock() - scraped_page.elements = [] - - browser_state = MagicMock() - - with ( - patch("skyvern.forge.agent.skyvern_context") as mock_ctx, - patch.object(agent, "_scrape_with_type", new_callable=AsyncMock, return_value=scraped_page), - patch.object( - agent, - "_build_extract_action_prompt", - new_callable=AsyncMock, - return_value=("prompt", False, "extract_action"), - ) as mock_build, - ): - mock_ctx.current.return_value = None - await agent.build_and_record_step_prompt( - task, step, browser_state, RunEngine.skyvern_v1, persist_artifacts=False - ) - mock_build.assert_called_once() - _, kwargs = mock_build.call_args - assert kwargs["verification_code_check"] is True - - -@pytest.mark.asyncio -async def test_verification_code_check_always_true_with_totp_config(): - """_build_extract_action_prompt should receive verification_code_check=True when task HAS TOTP config (unchanged).""" - agent = ForgeAgent.__new__(ForgeAgent) - agent.async_operation_pool = MagicMock() - - task = MagicMock() - task.totp_verification_url = "https://otp.example.com" - task.totp_identifier = "user@example.com" - task.task_id = "tsk_2" - task.workflow_run_id = None - task.organization_id = "org_1" - task.url = "https://example.com" - - step = MagicMock() - step.step_id = "step_1" - step.order = 0 - step.retry_index = 0 - - scraped_page = MagicMock() - scraped_page.elements = [] - - browser_state = MagicMock() - - with ( - patch("skyvern.forge.agent.skyvern_context") as mock_ctx, - patch.object(agent, "_scrape_with_type", new_callable=AsyncMock, return_value=scraped_page), - patch.object( - agent, - "_build_extract_action_prompt", - new_callable=AsyncMock, - return_value=("prompt", False, "extract_action"), - ) as mock_build, - ): - mock_ctx.current.return_value = None - await agent.build_and_record_step_prompt( - task, step, browser_state, RunEngine.skyvern_v1, persist_artifacts=False - ) - mock_build.assert_called_once() - _, kwargs = mock_build.call_args - assert kwargs["verification_code_check"] is True - - -# === Fix: poll_otp_value should pass workflow_id, not workflow_permanent_id === - - -@pytest.mark.asyncio -async def test_poll_otp_value_passes_workflow_id_not_permanent_id(): - """poll_otp_value should pass workflow_id (w_* format) to _get_otp_value_from_db, not workflow_permanent_id.""" - mock_db = AsyncMock() - mock_db.get_valid_org_auth_token.return_value = MagicMock(token="tok") - mock_db.update_workflow_run = AsyncMock() - - mock_app = MagicMock() - mock_app.DATABASE = mock_db - - with ( - patch("skyvern.services.otp_service.app", new=mock_app), - patch("skyvern.services.otp_service.asyncio.sleep", new_callable=AsyncMock), - patch( - "skyvern.services.otp_service._get_otp_value_from_db", - new_callable=AsyncMock, - return_value=OTPValue(value="654321", type="totp"), - ) as mock_get_from_db, - ): - result = await poll_otp_value( - organization_id="org_1", - workflow_id="w_123", - workflow_run_id="wr_789", - workflow_permanent_id="wpid_456", - totp_identifier="user@example.com", - ) - assert result is not None - assert result.value == "654321" - mock_get_from_db.assert_called_once_with( - "org_1", - "user@example.com", - task_id=None, - workflow_id="w_123", - workflow_run_id="wr_789", - ) - - -# === Fix: send_totp_code should resolve wpid_* to w_* before storage === - - -@pytest.mark.asyncio -async def test_send_totp_code_resolves_wpid_to_workflow_id(): - """send_totp_code should resolve wpid_* to w_* before storing in DB.""" - mock_workflow = MagicMock() - mock_workflow.workflow_id = "w_abc123" - - mock_totp_code = MagicMock() - - mock_db = AsyncMock() - mock_db.get_workflow_by_permanent_id = AsyncMock(return_value=mock_workflow) - mock_db.create_otp_code = AsyncMock(return_value=mock_totp_code) - - mock_app = MagicMock() - mock_app.DATABASE = mock_db - - data = TOTPCodeCreate( - totp_identifier="user@example.com", - content="123456", - workflow_id="wpid_xyz789", - ) - curr_org = MagicMock() - curr_org.organization_id = "org_1" - - with patch("skyvern.forge.sdk.routes.credentials.app", new=mock_app): - await send_totp_code(data=data, curr_org=curr_org) - - mock_db.create_otp_code.assert_called_once() - call_kwargs = mock_db.create_otp_code.call_args[1] - assert call_kwargs["workflow_id"] == "w_abc123", f"Expected w_abc123 but got {call_kwargs['workflow_id']}" - - -@pytest.mark.asyncio -async def test_send_totp_code_w_format_passes_through(): - """send_totp_code should resolve and store w_* format workflow_id correctly.""" - mock_workflow = MagicMock() - mock_workflow.workflow_id = "w_abc123" - - mock_totp_code = MagicMock() - - mock_db = AsyncMock() - mock_db.get_workflow = AsyncMock(return_value=mock_workflow) - mock_db.create_otp_code = AsyncMock(return_value=mock_totp_code) - - mock_app = MagicMock() - mock_app.DATABASE = mock_db - - data = TOTPCodeCreate( - totp_identifier="user@example.com", - content="123456", - workflow_id="w_abc123", - ) - curr_org = MagicMock() - curr_org.organization_id = "org_1" - - with patch("skyvern.forge.sdk.routes.credentials.app", new=mock_app): - await send_totp_code(data=data, curr_org=curr_org) - - call_kwargs = mock_db.create_otp_code.call_args[1] - assert call_kwargs["workflow_id"] == "w_abc123" - - -@pytest.mark.asyncio -async def test_send_totp_code_none_workflow_id(): - """send_totp_code should pass None workflow_id when not provided.""" - mock_totp_code = MagicMock() - - mock_db = AsyncMock() - mock_db.create_otp_code = AsyncMock(return_value=mock_totp_code) - - mock_app = MagicMock() - mock_app.DATABASE = mock_db - - data = TOTPCodeCreate( - totp_identifier="user@example.com", - content="123456", - ) - curr_org = MagicMock() - curr_org.organization_id = "org_1" - - with patch("skyvern.forge.sdk.routes.credentials.app", new=mock_app): - await send_totp_code(data=data, curr_org=curr_org) - - call_kwargs = mock_db.create_otp_code.call_args[1] - assert call_kwargs["workflow_id"] is None - - -# === Fix: _build_navigation_payload should inject code without TOTP config === - - -def test_build_navigation_payload_injects_code_without_totp_config(): - """_build_navigation_payload should inject SPECIAL_FIELD_VERIFICATION_CODE even when - task has no totp_verification_url or totp_identifier (manual 2FA flow).""" - agent = ForgeAgent.__new__(ForgeAgent) - - task = MagicMock() - task.totp_verification_url = None - task.totp_identifier = None - task.task_id = "tsk_manual_2fa" - task.workflow_run_id = "wr_123" - task.navigation_payload = {"username": "user@example.com"} - - mock_context = MagicMock() - mock_context.totp_codes = {"tsk_manual_2fa": "123456"} - mock_context.has_magic_link_page.return_value = False - - with patch("skyvern.forge.agent.skyvern_context") as mock_skyvern_ctx: - mock_skyvern_ctx.ensure_context.return_value = mock_context - result = agent._build_navigation_payload(task) - - assert isinstance(result, dict) - assert SPECIAL_FIELD_VERIFICATION_CODE in result - assert result[SPECIAL_FIELD_VERIFICATION_CODE] == "123456" - # Original payload preserved - assert result["username"] == "user@example.com" - - -def test_build_navigation_payload_injects_code_when_payload_is_none(): - """_build_navigation_payload should create a dict with the code when payload is None.""" - agent = ForgeAgent.__new__(ForgeAgent) - - task = MagicMock() - task.totp_verification_url = None - task.totp_identifier = None - task.task_id = "tsk_manual_2fa" - task.workflow_run_id = "wr_123" - task.navigation_payload = None - - mock_context = MagicMock() - mock_context.totp_codes = {"tsk_manual_2fa": "999999"} - mock_context.has_magic_link_page.return_value = False - - with patch("skyvern.forge.agent.skyvern_context") as mock_skyvern_ctx: - mock_skyvern_ctx.ensure_context.return_value = mock_context - result = agent._build_navigation_payload(task) - - assert isinstance(result, dict) - assert result[SPECIAL_FIELD_VERIFICATION_CODE] == "999999" - - -def test_build_navigation_payload_no_code_no_injection(): - """_build_navigation_payload should NOT inject anything when no code in context.""" - agent = ForgeAgent.__new__(ForgeAgent) - - task = MagicMock() - task.totp_verification_url = None - task.totp_identifier = None - task.task_id = "tsk_no_code" - task.workflow_run_id = "wr_456" - task.navigation_payload = {"field": "value"} - - mock_context = MagicMock() - mock_context.totp_codes = {} # No code in context - mock_context.has_magic_link_page.return_value = False - - with patch("skyvern.forge.agent.skyvern_context") as mock_skyvern_ctx: - mock_skyvern_ctx.ensure_context.return_value = mock_context - result = agent._build_navigation_payload(task) - - assert isinstance(result, dict) - assert SPECIAL_FIELD_VERIFICATION_CODE not in result - assert result["field"] == "value" - - -# === Task: poll_otp_value publishes 2FA events to notification registry === - - -@pytest.mark.asyncio -async def test_poll_otp_value_publishes_required_event_for_task(): - """poll_otp_value should publish verification_code_required when task waiting state is set.""" - mock_code = MagicMock() - mock_code.code = "123456" - mock_code.otp_type = "totp" - - mock_db = AsyncMock() - mock_db.get_valid_org_auth_token.return_value = MagicMock(token="tok") - mock_db.get_otp_codes_by_run.return_value = [mock_code] - mock_db.update_task_2fa_state = AsyncMock() - - mock_app = MagicMock() - mock_app.DATABASE = mock_db - - registry = LocalNotificationRegistry() - queue = registry.subscribe("org_1") - - with ( - patch("skyvern.services.otp_service.app", new=mock_app), - patch("skyvern.services.otp_service.asyncio.sleep", new_callable=AsyncMock), - patch( - "skyvern.forge.sdk.notification.factory.NotificationRegistryFactory._NotificationRegistryFactory__registry", - new=registry, - ), - ): - await poll_otp_value(organization_id="org_1", task_id="tsk_1") - - # Should have received required + resolved messages - messages = [] - while not queue.empty(): - messages.append(queue.get_nowait()) - - types = [m["type"] for m in messages] - assert "verification_code_required" in types - assert "verification_code_resolved" in types - - required = next(m for m in messages if m["type"] == "verification_code_required") - assert required["task_id"] == "tsk_1" - - resolved = next(m for m in messages if m["type"] == "verification_code_resolved") - assert resolved["task_id"] == "tsk_1" - - -@pytest.mark.asyncio -async def test_reprompt_with_verification_code_calls_llm(): - """When poll_otp_value returns a TOTP code, handle_potential_verification_code should - store the code in context, re-prompt the LLM with verification_code_check=False, and - return the LLM's new response.""" - from skyvern.forge import agent as forge_agent_module - - agent = ForgeAgent.__new__(ForgeAgent) - - task = MagicMock() - task.organization_id = "org_1" - task.totp_verification_url = None - task.totp_identifier = None - task.task_id = "tsk_reprompt" - task.workflow_run_id = None - task.navigation_payload = None - task.llm_key = None - - step = MagicMock() - step.step_id = "step_1" - step.order = 0 - - scraped_page = MagicMock() - scraped_page.screenshots = [] - browser_state = MagicMock() - json_response = { - "should_enter_verification_code": True, - "place_to_enter_verification_code": "input#otp-code", - "actions": [], - } - - llm_response = {"actions": [{"action_type": "input_text", "text": "999888"}]} - - original_app_inst = object.__getattribute__(forge_agent_module.app, "_inst") - object.__setattr__(forge_agent_module.app, "_inst", MagicMock(LLM_API_HANDLER=AsyncMock())) - try: - with ( - patch("skyvern.forge.agent.extract_totp_from_navigation_inputs", return_value=None), - patch("skyvern.forge.agent.try_generate_totp_from_credential", return_value=None), - patch( - "skyvern.forge.agent.poll_otp_value", - new_callable=AsyncMock, - return_value=OTPValue(value="999888", type=OTPType.TOTP), - ), - patch("skyvern.forge.agent.skyvern_context") as mock_skyvern_context, - patch("skyvern.forge.agent.service_utils.is_cua_task", new_callable=AsyncMock, return_value=False), - patch( - "skyvern.forge.agent.LLMAPIHandlerFactory.get_override_llm_api_handler", - return_value=AsyncMock(return_value=llm_response), - ), - patch.object( - agent, - "_build_extract_action_prompt", - new_callable=AsyncMock, - return_value=("prompt", False, "extract-actions"), - ) as mock_build, - ): - mock_context = MagicMock() - mock_context.totp_codes = {} - mock_skyvern_context.ensure_context.return_value = mock_context - mock_skyvern_context.current.return_value = mock_context - - result = await agent.handle_potential_verification_code( - task=task, - step=step, - scraped_page=scraped_page, - browser_state=browser_state, - json_response=json_response, - ) - finally: - object.__setattr__(forge_agent_module.app, "_inst", original_app_inst) - - mock_build.assert_called_once() - _, kwargs = mock_build.call_args - assert kwargs["verification_code_check"] is False - assert mock_context.totp_codes["tsk_reprompt"] == "999888" - assert result == llm_response - - -@pytest.mark.asyncio -async def test_poll_otp_value_publishes_required_event_for_workflow_run(): - """poll_otp_value should publish verification_code_required when workflow run waiting state is set.""" - mock_code = MagicMock() - mock_code.code = "654321" - mock_code.otp_type = "totp" - - mock_db = AsyncMock() - mock_db.get_valid_org_auth_token.return_value = MagicMock(token="tok") - mock_db.update_workflow_run = AsyncMock() - mock_db.get_otp_codes_by_run.return_value = [mock_code] - - mock_app = MagicMock() - mock_app.DATABASE = mock_db - - registry = LocalNotificationRegistry() - queue = registry.subscribe("org_1") - - with ( - patch("skyvern.services.otp_service.app", new=mock_app), - patch("skyvern.services.otp_service.asyncio.sleep", new_callable=AsyncMock), - patch( - "skyvern.forge.sdk.notification.factory.NotificationRegistryFactory._NotificationRegistryFactory__registry", - new=registry, - ), - ): - await poll_otp_value(organization_id="org_1", workflow_run_id="wr_1") - - messages = [] - while not queue.empty(): - messages.append(queue.get_nowait()) - - types = [m["type"] for m in messages] - assert "verification_code_required" in types - assert "verification_code_resolved" in types - - required = next(m for m in messages if m["type"] == "verification_code_required") - assert required["workflow_run_id"] == "wr_1" diff --git a/tests/unit_tests/test_otp_service_organization_context.py b/tests/unit_tests/test_otp_service_organization_context.py deleted file mode 100644 index 250280975..000000000 --- a/tests/unit_tests/test_otp_service_organization_context.py +++ /dev/null @@ -1,152 +0,0 @@ -"""Tests for poll_otp_value organization token usage by OTP context.""" - -from datetime import datetime -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -from skyvern.forge.sdk.db.enums import OrganizationAuthTokenType -from skyvern.forge.sdk.notification.local import LocalNotificationRegistry -from skyvern.services import otp_service -from skyvern.services.otp_service import OTPValue, poll_otp_value - - -@pytest.mark.asyncio -async def test_poll_otp_value_without_totp_url_does_not_require_org_token() -> None: - """poll_otp_value should not depend on org token when no totp_verification_url is set.""" - expected_otp = OTPValue(value="123456") - - mock_db = AsyncMock() - mock_db.get_valid_org_auth_token.return_value = None - mock_db.update_task_2fa_state = AsyncMock() - - mock_app = MagicMock() - mock_app.DATABASE = mock_db - - with ( - patch("skyvern.services.otp_service.app", new=mock_app), - patch( - "skyvern.services.otp_service._get_otp_value_by_run", - new_callable=AsyncMock, - return_value=expected_otp, - ) as mock_get_otp_by_run, - patch("skyvern.services.otp_service.asyncio.sleep", new_callable=AsyncMock), - ): - result = await poll_otp_value( - organization_id="org_1", - task_id="task_1", - ) - - assert result == expected_otp - mock_get_otp_by_run.assert_awaited_once_with( - "org_1", - task_id="task_1", - workflow_run_id=None, - ) - mock_db.get_valid_org_auth_token.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_poll_otp_value_with_totp_url_still_checks_org_token() -> None: - """poll_otp_value should continue checking org token when totp_verification_url is configured.""" - mock_db = AsyncMock() - mock_db.get_valid_org_auth_token.return_value = None - - mock_app = MagicMock() - mock_app.DATABASE = mock_db - - with ( - patch("skyvern.services.otp_service.app", new=mock_app), - patch("skyvern.services.otp_service._get_otp_value_from_url", new_callable=AsyncMock) as mock_from_url, - ): - result = await poll_otp_value( - organization_id="org_1", - task_id="task_1", - totp_verification_url="https://otp.example.com", - ) - - assert result is None - mock_db.get_valid_org_auth_token.assert_awaited_once_with( - "org_1", - OrganizationAuthTokenType.api.value, - ) - mock_from_url.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_set_waiting_state_updates_task_and_publishes_required_event() -> None: - """_set_waiting_state should write task waiting state and publish required event to org channel.""" - started_at = datetime(2026, 1, 2, 3, 4, 5) - - mock_db = AsyncMock() - mock_db.update_task_2fa_state = AsyncMock() - - mock_app = MagicMock() - mock_app.DATABASE = mock_db - - registry = LocalNotificationRegistry() - queue = registry.subscribe("org_1") - - ctx = otp_service.OTPPollContext( - organization_id="org_1", - task_id="tsk_1", - ) - - with ( - patch("skyvern.services.otp_service.app", new=mock_app), - patch( - "skyvern.forge.sdk.notification.factory.NotificationRegistryFactory._NotificationRegistryFactory__registry", - new=registry, - ), - ): - await otp_service._set_waiting_state(ctx, started_at) - - mock_db.update_task_2fa_state.assert_awaited_once() - update_kwargs = mock_db.update_task_2fa_state.await_args.kwargs - assert update_kwargs["organization_id"] == "org_1" - assert update_kwargs["task_id"] == "tsk_1" - assert update_kwargs["waiting_for_verification_code"] is True - assert update_kwargs["verification_code_polling_started_at"] == started_at - - message = queue.get_nowait() - assert message["type"] == "verification_code_required" - assert message["task_id"] == "tsk_1" - assert queue.empty() - - -@pytest.mark.asyncio -async def test_clear_waiting_state_updates_task_and_publishes_resolved_event() -> None: - """_clear_waiting_state should clear task waiting state and publish resolved event to org channel.""" - mock_db = AsyncMock() - mock_db.update_task_2fa_state = AsyncMock() - - mock_app = MagicMock() - mock_app.DATABASE = mock_db - - registry = LocalNotificationRegistry() - queue = registry.subscribe("org_1") - - ctx = otp_service.OTPPollContext( - organization_id="org_1", - task_id="tsk_1", - ) - - with ( - patch("skyvern.services.otp_service.app", new=mock_app), - patch( - "skyvern.forge.sdk.notification.factory.NotificationRegistryFactory._NotificationRegistryFactory__registry", - new=registry, - ), - ): - await otp_service._clear_waiting_state(ctx) - - mock_db.update_task_2fa_state.assert_awaited_once() - update_kwargs = mock_db.update_task_2fa_state.await_args.kwargs - assert update_kwargs["organization_id"] == "org_1" - assert update_kwargs["task_id"] == "tsk_1" - assert update_kwargs["waiting_for_verification_code"] is False - - message = queue.get_nowait() - assert message["type"] == "verification_code_resolved" - assert message["task_id"] == "tsk_1" - assert queue.empty() diff --git a/tests/unit_tests/test_redis_client_factory.py b/tests/unit_tests/test_redis_client_factory.py deleted file mode 100644 index 36e31b44c..000000000 --- a/tests/unit_tests/test_redis_client_factory.py +++ /dev/null @@ -1,22 +0,0 @@ -"""Tests for RedisClientFactory.""" - -from unittest.mock import MagicMock - -from skyvern.forge.sdk.redis.factory import RedisClientFactory - - -def test_default_is_none(): - """Factory returns None when no client has been set.""" - # Reset to default state - RedisClientFactory.set_client(None) # type: ignore[arg-type] - assert RedisClientFactory.get_client() is None - - -def test_set_and_get(): - """Round-trip: set_client then get_client returns the same object.""" - mock_client = MagicMock() - RedisClientFactory.set_client(mock_client) - assert RedisClientFactory.get_client() is mock_client - - # Cleanup - RedisClientFactory.set_client(None) # type: ignore[arg-type] diff --git a/tests/unit_tests/test_redis_notification_registry.py b/tests/unit_tests/test_redis_notification_registry.py deleted file mode 100644 index 7370476d6..000000000 --- a/tests/unit_tests/test_redis_notification_registry.py +++ /dev/null @@ -1,237 +0,0 @@ -"""Tests for RedisNotificationRegistry (SKY-6). - -All tests use a mock Redis client — no real Redis instance required. -""" - -import asyncio -import json -from unittest.mock import AsyncMock, MagicMock - -import pytest - -from skyvern.forge.sdk.notification.redis import RedisNotificationRegistry - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _make_mock_redis() -> MagicMock: - """Return a mock redis.asyncio.Redis client.""" - redis = MagicMock() - redis.publish = AsyncMock() - redis.pubsub = MagicMock() - return redis - - -def _make_mock_pubsub(messages: list[dict] | None = None, *, block: bool = False) -> MagicMock: - """Return a mock PubSub that yields *messages* from ``listen()``. - - Each entry in *messages* should look like: - {"type": "message", "data": '{"key": "val"}'} - - If *block* is True the async generator will hang forever after - exhausting *messages*, which keeps the listener task alive so that - cancellation semantics can be tested. - """ - pubsub = MagicMock() - pubsub.subscribe = AsyncMock() - pubsub.unsubscribe = AsyncMock() - pubsub.close = AsyncMock() - - async def _listen(): - for msg in messages or []: - yield msg - if block: - # Keep the listener alive until cancelled - await asyncio.Event().wait() - - pubsub.listen = _listen - return pubsub - - -# --------------------------------------------------------------------------- -# Tests: subscribe -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_subscribe_creates_queue_and_starts_listener(): - """subscribe() should return a queue and start a background listener task.""" - redis = _make_mock_redis() - pubsub = _make_mock_pubsub() - redis.pubsub.return_value = pubsub - - registry = RedisNotificationRegistry(redis) - queue = registry.subscribe("org_1") - - assert isinstance(queue, asyncio.Queue) - assert "org_1" in registry._listener_tasks - task = registry._listener_tasks["org_1"] - assert isinstance(task, asyncio.Task) - - # Cleanup - await registry.close() - - -@pytest.mark.asyncio -async def test_subscribe_reuses_listener_for_same_org(): - """A second subscribe for the same org should NOT create a new listener task.""" - redis = _make_mock_redis() - pubsub = _make_mock_pubsub() - redis.pubsub.return_value = pubsub - - registry = RedisNotificationRegistry(redis) - registry.subscribe("org_1") - first_task = registry._listener_tasks["org_1"] - - registry.subscribe("org_1") - assert registry._listener_tasks["org_1"] is first_task - - await registry.close() - - -# --------------------------------------------------------------------------- -# Tests: unsubscribe -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_unsubscribe_cancels_listener_when_last_subscriber_leaves(): - """When the last subscriber unsubscribes, the listener task should be cancelled.""" - redis = _make_mock_redis() - pubsub = _make_mock_pubsub(block=True) - redis.pubsub.return_value = pubsub - - registry = RedisNotificationRegistry(redis) - queue = registry.subscribe("org_1") - - # Let the listener task start running - await asyncio.sleep(0) - - task = registry._listener_tasks["org_1"] - - registry.unsubscribe("org_1", queue) - assert "org_1" not in registry._listener_tasks - - # Wait for the task to fully complete after cancellation - await asyncio.gather(task, return_exceptions=True) - assert task.cancelled() - - await registry.close() - - -@pytest.mark.asyncio -async def test_unsubscribe_keeps_listener_when_subscribers_remain(): - """If other subscribers remain, the listener task should stay alive.""" - redis = _make_mock_redis() - pubsub = _make_mock_pubsub() - redis.pubsub.return_value = pubsub - - registry = RedisNotificationRegistry(redis) - q1 = registry.subscribe("org_1") - registry.subscribe("org_1") # second subscriber - - registry.unsubscribe("org_1", q1) - assert "org_1" in registry._listener_tasks - - await registry.close() - - -# --------------------------------------------------------------------------- -# Tests: publish -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_publish_calls_redis_publish(): - """publish() should fire-and-forget a Redis PUBLISH.""" - redis = _make_mock_redis() - registry = RedisNotificationRegistry(redis) - - registry.publish("org_1", {"type": "verification_code_required"}) - - # Allow the fire-and-forget task to execute - await asyncio.sleep(0) - - redis.publish.assert_awaited_once_with( - "skyvern:notifications:org_1", - json.dumps({"type": "verification_code_required"}), - ) - - await registry.close() - - -# --------------------------------------------------------------------------- -# Tests: _dispatch_local -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_dispatch_local_fans_out_to_all_queues(): - """_dispatch_local should put the message into every local queue for the org.""" - redis = _make_mock_redis() - pubsub = _make_mock_pubsub() - redis.pubsub.return_value = pubsub - - registry = RedisNotificationRegistry(redis) - q1 = registry.subscribe("org_1") - q2 = registry.subscribe("org_1") - - msg = {"type": "test", "value": 42} - registry._dispatch_local("org_1", msg) - - assert q1.get_nowait() == msg - assert q2.get_nowait() == msg - - await registry.close() - - -@pytest.mark.asyncio -async def test_dispatch_local_does_not_leak_across_orgs(): - """Messages dispatched for org_a should not appear in org_b queues.""" - redis = _make_mock_redis() - pubsub = _make_mock_pubsub() - redis.pubsub.return_value = pubsub - - registry = RedisNotificationRegistry(redis) - q_a = registry.subscribe("org_a") - q_b = registry.subscribe("org_b") - - registry._dispatch_local("org_a", {"type": "test"}) - assert not q_a.empty() - assert q_b.empty() - - await registry.close() - - -# --------------------------------------------------------------------------- -# Tests: close -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_close_cancels_all_listeners_and_clears_state(): - """close() should cancel every listener task and empty subscriber maps.""" - redis = _make_mock_redis() - pubsub = _make_mock_pubsub(block=True) - redis.pubsub.return_value = pubsub - - registry = RedisNotificationRegistry(redis) - registry.subscribe("org_1") - registry.subscribe("org_2") - - # Let the listener tasks start running - await asyncio.sleep(0) - - task_1 = registry._listener_tasks["org_1"] - task_2 = registry._listener_tasks["org_2"] - - await registry.close() - - # Wait for the tasks to fully complete after cancellation - await asyncio.gather(task_1, task_2, return_exceptions=True) - assert task_1.cancelled() - assert task_2.cancelled() - assert len(registry._listener_tasks) == 0 - assert len(registry._subscribers) == 0 diff --git a/tests/unit_tests/test_redis_pubsub.py b/tests/unit_tests/test_redis_pubsub.py deleted file mode 100644 index 97dae81e2..000000000 --- a/tests/unit_tests/test_redis_pubsub.py +++ /dev/null @@ -1,262 +0,0 @@ -"""Tests for RedisPubSub (generic pub/sub layer). - -All tests use a mock Redis client — no real Redis instance required. -""" - -import asyncio -import json -from unittest.mock import AsyncMock, MagicMock - -import pytest - -from skyvern.forge.sdk.redis.pubsub import RedisPubSub - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _make_mock_redis() -> MagicMock: - """Return a mock redis.asyncio.Redis client.""" - redis = MagicMock() - redis.publish = AsyncMock() - redis.pubsub = MagicMock() - return redis - - -def _make_mock_pubsub(messages: list[dict] | None = None, *, block: bool = False) -> MagicMock: - """Return a mock PubSub that yields *messages* from ``listen()``. - - Each entry in *messages* should look like: - {"type": "message", "data": '{"key": "val"}'} - - If *block* is True the async generator will hang forever after - exhausting *messages*, which keeps the listener task alive so that - cancellation semantics can be tested. - """ - pubsub = MagicMock() - pubsub.subscribe = AsyncMock() - pubsub.unsubscribe = AsyncMock() - pubsub.close = AsyncMock() - - async def _listen(): - for msg in messages or []: - yield msg - if block: - await asyncio.Event().wait() - - pubsub.listen = _listen - return pubsub - - -PREFIX = "skyvern:test:" - - -# --------------------------------------------------------------------------- -# Tests: subscribe -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_subscribe_creates_queue_and_starts_listener(): - """subscribe() should return a queue and start a background listener task.""" - redis = _make_mock_redis() - pubsub = _make_mock_pubsub() - redis.pubsub.return_value = pubsub - - ps = RedisPubSub(redis, channel_prefix=PREFIX) - queue = ps.subscribe("key_1") - - assert isinstance(queue, asyncio.Queue) - assert "key_1" in ps._listener_tasks - task = ps._listener_tasks["key_1"] - assert isinstance(task, asyncio.Task) - - await ps.close() - - -@pytest.mark.asyncio -async def test_subscribe_reuses_listener_for_same_key(): - """A second subscribe for the same key should NOT create a new listener task.""" - redis = _make_mock_redis() - pubsub = _make_mock_pubsub() - redis.pubsub.return_value = pubsub - - ps = RedisPubSub(redis, channel_prefix=PREFIX) - ps.subscribe("key_1") - first_task = ps._listener_tasks["key_1"] - - ps.subscribe("key_1") - assert ps._listener_tasks["key_1"] is first_task - - await ps.close() - - -# --------------------------------------------------------------------------- -# Tests: unsubscribe -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_unsubscribe_cancels_listener_when_last_subscriber_leaves(): - """When the last subscriber unsubscribes, the listener task should be cancelled.""" - redis = _make_mock_redis() - pubsub = _make_mock_pubsub(block=True) - redis.pubsub.return_value = pubsub - - ps = RedisPubSub(redis, channel_prefix=PREFIX) - queue = ps.subscribe("key_1") - - await asyncio.sleep(0) - - task = ps._listener_tasks["key_1"] - - ps.unsubscribe("key_1", queue) - assert "key_1" not in ps._listener_tasks - - await asyncio.gather(task, return_exceptions=True) - assert task.cancelled() - - await ps.close() - - -@pytest.mark.asyncio -async def test_unsubscribe_keeps_listener_when_subscribers_remain(): - """If other subscribers remain, the listener task should stay alive.""" - redis = _make_mock_redis() - pubsub = _make_mock_pubsub() - redis.pubsub.return_value = pubsub - - ps = RedisPubSub(redis, channel_prefix=PREFIX) - q1 = ps.subscribe("key_1") - ps.subscribe("key_1") - - ps.unsubscribe("key_1", q1) - assert "key_1" in ps._listener_tasks - - await ps.close() - - -# --------------------------------------------------------------------------- -# Tests: publish -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_publish_calls_redis_publish(): - """publish() should fire-and-forget a Redis PUBLISH with prefixed channel.""" - redis = _make_mock_redis() - ps = RedisPubSub(redis, channel_prefix=PREFIX) - - ps.publish("key_1", {"type": "event"}) - - await asyncio.sleep(0) - - redis.publish.assert_awaited_once_with( - f"{PREFIX}key_1", - json.dumps({"type": "event"}), - ) - - await ps.close() - - -# --------------------------------------------------------------------------- -# Tests: _dispatch_local -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_dispatch_local_fans_out_to_all_queues(): - """_dispatch_local should put the message into every local queue for the key.""" - redis = _make_mock_redis() - pubsub = _make_mock_pubsub() - redis.pubsub.return_value = pubsub - - ps = RedisPubSub(redis, channel_prefix=PREFIX) - q1 = ps.subscribe("key_1") - q2 = ps.subscribe("key_1") - - msg = {"type": "test", "value": 42} - ps._dispatch_local("key_1", msg) - - assert q1.get_nowait() == msg - assert q2.get_nowait() == msg - - await ps.close() - - -@pytest.mark.asyncio -async def test_dispatch_local_does_not_leak_across_keys(): - """Messages dispatched for key_a should not appear in key_b queues.""" - redis = _make_mock_redis() - pubsub = _make_mock_pubsub() - redis.pubsub.return_value = pubsub - - ps = RedisPubSub(redis, channel_prefix=PREFIX) - q_a = ps.subscribe("key_a") - q_b = ps.subscribe("key_b") - - ps._dispatch_local("key_a", {"type": "test"}) - assert not q_a.empty() - assert q_b.empty() - - await ps.close() - - -# --------------------------------------------------------------------------- -# Tests: close -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_close_cancels_all_listeners_and_clears_state(): - """close() should cancel every listener task and empty subscriber maps.""" - redis = _make_mock_redis() - pubsub = _make_mock_pubsub(block=True) - redis.pubsub.return_value = pubsub - - ps = RedisPubSub(redis, channel_prefix=PREFIX) - ps.subscribe("key_1") - ps.subscribe("key_2") - - await asyncio.sleep(0) - - task_1 = ps._listener_tasks["key_1"] - task_2 = ps._listener_tasks["key_2"] - - await ps.close() - - await asyncio.gather(task_1, task_2, return_exceptions=True) - assert task_1.cancelled() - assert task_2.cancelled() - assert len(ps._listener_tasks) == 0 - assert len(ps._subscribers) == 0 - - -# --------------------------------------------------------------------------- -# Tests: prefix isolation -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_different_prefixes_do_not_interfere(): - """Two RedisPubSub instances with different prefixes use separate channels.""" - redis = _make_mock_redis() - - ps_a = RedisPubSub(redis, channel_prefix="prefix_a:") - ps_b = RedisPubSub(redis, channel_prefix="prefix_b:") - - ps_a.publish("key_1", {"from": "a"}) - ps_b.publish("key_1", {"from": "b"}) - - await asyncio.sleep(0) - - calls = redis.publish.await_args_list - assert len(calls) == 2 - - channels = {call.args[0] for call in calls} - assert "prefix_a:key_1" in channels - assert "prefix_b:key_1" in channels - - await ps_a.close() - await ps_b.close()