[SKY-7980] Patch Credential TOTP Over Webhook Logic (#4811)

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