From f8f9d2a17f1fafab51abfd8fa7dacad71598afbc Mon Sep 17 00:00:00 2001 From: Aaron Perez Date: Thu, 19 Feb 2026 18:14:44 -0500 Subject: [PATCH] [SKY-7980] Patch Credential TOTP Over Webhook Logic (#4811) --- skyvern/forge/agent.py | 2 +- .../sdk/routes/streaming/notifications.py | 19 ++- skyvern/webeye/actions/parse_actions.py | 2 +- tests/unit/test_credential_totp_priority.py | 135 +++++++++++++++++- 4 files changed, 154 insertions(+), 4 deletions(-) diff --git a/skyvern/forge/agent.py b/skyvern/forge/agent.py index 649e5f8d..1a96430a 100644 --- a/skyvern/forge/agent.py +++ b/skyvern/forge/agent.py @@ -4529,7 +4529,7 @@ class ForgeAgent: # Try credential TOTP first (highest priority, doesn't need totp_url/totp_identifier) otp_value = try_generate_totp_from_credential(task.workflow_run_id) # Fall back to webhook/totp_identifier - if not otp_value and (task.totp_verification_url or task.totp_identifier): + 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) diff --git a/skyvern/forge/sdk/routes/streaming/notifications.py b/skyvern/forge/sdk/routes/streaming/notifications.py index bf713a73..c713c0fd 100644 --- a/skyvern/forge/sdk/routes/streaming/notifications.py +++ b/skyvern/forge/sdk/routes/streaming/notifications.py @@ -9,7 +9,7 @@ from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK from skyvern.config import settings from skyvern.forge import app from skyvern.forge.sdk.notification.factory import NotificationRegistryFactory -from skyvern.forge.sdk.routes.routers import base_router +from skyvern.forge.sdk.routes.routers import base_router, legacy_base_router from skyvern.forge.sdk.routes.streaming.auth import _auth as local_auth from skyvern.forge.sdk.routes.streaming.auth import auth as real_auth @@ -22,6 +22,23 @@ 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: auth = local_auth if settings.ENV == "local" else real_auth organization_id = await auth(apikey=apikey, token=token, websocket=websocket) diff --git a/skyvern/webeye/actions/parse_actions.py b/skyvern/webeye/actions/parse_actions.py index e53e7021..2129fb6c 100644 --- a/skyvern/webeye/actions/parse_actions.py +++ b/skyvern/webeye/actions/parse_actions.py @@ -916,7 +916,7 @@ async def generate_cua_fallback_actions( # Try credential TOTP first (highest priority, doesn't need totp_url/totp_identifier) otp_value = try_generate_totp_from_credential(task.workflow_run_id) # Fall back to webhook/totp_identifier - if not otp_value and (task.totp_verification_url or task.totp_identifier) and task.organization_id: + if not otp_value and task.organization_id: LOG.info( "Getting verification code for CUA", task_id=task.task_id, diff --git a/tests/unit/test_credential_totp_priority.py b/tests/unit/test_credential_totp_priority.py index caf0dabc..1dd9c7b9 100644 --- a/tests/unit/test_credential_totp_priority.py +++ b/tests/unit/test_credential_totp_priority.py @@ -5,9 +5,10 @@ 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 MagicMock, patch +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, try_generate_totp_from_credential @@ -165,3 +166,135 @@ class TestTryGenerateTotpFromCredential: 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