vikarti.anatra/Skyvern
1
0
Fork 0
mirror of https://github.com/Skyvern-AI/skyvern.git synced 2026-07-10 00:19:39 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 5

[SKY-6] Backend 2FA Cleanup (#4826)

Browse source
This commit is contained in:
Aaron Perez 2026-02-20 17:06:08 -05:00 • committed by GitHub
parent b56d724ed8
commit aa2eb8d4e8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 355 additions and 130 deletions
Download patch file Download diff file Expand all files Collapse all files

2
skyvern/forge/sdk/db/models.py
View file

@ -116,7 +116,6 @@ class TaskModel(Base):
model = Column(JSON, nullable=True)
browser_address = Column(String, nullable=True)
download_timeout = Column(Numeric, nullable=True)
# 2FA verification code waiting state fields
waiting_for_verification_code = Column(Boolean, nullable=False, default=False, server_default=sqlalchemy.false())
verification_code_identifier = Column(String, nullable=True)
verification_code_polling_started_at = Column(DateTime, nullable=True)
@ -354,7 +353,6 @@ class WorkflowRunModel(Base):
debug_session_id: Column = Column(String, nullable=True)
ai_fallback = Column(Boolean, nullable=True)
code_gen = Column(Boolean, nullable=True)
# 2FA verification code waiting state fields
waiting_for_verification_code = Column(Boolean, nullable=False, default=False, server_default=sqlalchemy.false())
verification_code_identifier = Column(String, nullable=True)
verification_code_polling_started_at = Column(DateTime, nullable=True)

17
skyvern/forge/sdk/routes/streaming/auth.py
View file

@ -73,7 +73,8 @@ async def auth(apikey: str | None, token: str | None, websocket: WebSocket) -> s
# NOTE(jdo:streaming-local-dev): use this instead of the above `auth`
async def _auth(apikey: str | None, token: str | None, websocket: WebSocket) -> str | None:
"""
Dummy auth for local testing.
Local dev auth: extracts org_id from API key without strict validation.
Falls back to o_temp123 if no key provided.
"""
try:
@ -82,4 +83,18 @@ async def _auth(apikey: str | None, token: str | None, websocket: WebSocket) ->
LOG.info("WebSocket connection closed cleanly.")
return None
# Try to extract real org_id from the API key
if apikey:
try:
from jose import jwt
from skyvern.config import settings
payload = jwt.decode(apikey, settings.SECRET_KEY, algorithms=["HS256"])
org_id = payload.get("sub")
if org_id:
return org_id
except Exception:
LOG.warning("Local auth: failed to decode API key, falling back to o_temp123")
return "o_temp123"

31
skyvern/forge/sdk/routes/streaming/notifications.py
View file

@ -54,15 +54,22 @@ async def _notification_stream_handler(
# Send initial state: all currently active verification requests
active_requests = await app.DATABASE.get_active_verification_requests(organization_id)
for req in active_requests:
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"),
}
)
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()
@ -100,6 +107,12 @@ async def _notification_stream_handler(
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()

283
skyvern/services/otp_service.py
View file

@ -1,4 +1,5 @@
import asyncio
from dataclasses import dataclass
from datetime import datetime, timedelta
import pyotp
@ -18,6 +19,21 @@ from skyvern.forge.sdk.schemas.totp_codes import OTPType
LOG = structlog.get_logger()
@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):
value: str = Field(..., description="The value of the OTP code.")
type: OTPType | None = Field(None, description="The type of the OTP code.")
@ -108,80 +124,37 @@ async def poll_otp_value(
totp_verification_url: str | None = None,
totp_identifier: str | None = None,
) -> OTPValue | None:
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",
ctx = OTPPollContext(
organization_id=organization_id,
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,
)
# Set the waiting state in the database when polling starts
identifier_for_ui = totp_identifier
if workflow_run_id:
try:
await app.DATABASE.update_workflow_run(
workflow_run_id=workflow_run_id,
waiting_for_verification_code=True,
verification_code_identifier=identifier_for_ui,
verification_code_polling_started_at=start_datetime,
)
LOG.info(
"Set 2FA waiting state for workflow run",
workflow_run_id=workflow_run_id,
verification_code_identifier=identifier_for_ui,
)
try:
NotificationRegistryFactory.get_registry().publish(
organization_id,
{
"type": "verification_code_required",
"workflow_run_id": workflow_run_id,
"task_id": task_id,
"identifier": identifier_for_ui,
"polling_started_at": start_datetime.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 task_id:
try:
await app.DATABASE.update_task_2fa_state(
task_id=task_id,
organization_id=organization_id,
waiting_for_verification_code=True,
verification_code_identifier=identifier_for_ui,
verification_code_polling_started_at=start_datetime,
)
LOG.info(
"Set 2FA waiting state for task",
task_id=task_id,
verification_code_identifier=identifier_for_ui,
)
try:
NotificationRegistryFactory.get_registry().publish(
organization_id,
{
"type": "verification_code_required",
"task_id": task_id,
"identifier": identifier_for_ui,
"polling_started_at": start_datetime.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)
await _set_waiting_state(ctx, start_datetime)
try:
while True:
@ -190,80 +163,154 @@ async def poll_otp_value(
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,
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 totp_verification_url:
if ctx.totp_verification_url:
assert org_api_key is not None
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,
ctx.organization_id,
ctx.totp_verification_url,
org_api_key,
task_id=ctx.task_id,
workflow_run_id=ctx.workflow_run_id,
)
elif totp_identifier:
elif ctx.totp_identifier:
otp_value = await _get_otp_value_from_db(
organization_id,
totp_identifier,
task_id=task_id,
workflow_id=workflow_id,
workflow_run_id=workflow_run_id,
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(
organization_id,
task_id=task_id,
workflow_run_id=workflow_run_id,
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(
organization_id,
task_id=task_id,
workflow_run_id=workflow_run_id,
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:
# Clear the waiting state when polling completes (success, timeout, or error)
if workflow_run_id:
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,
)
LOG.info(
"Set 2FA waiting state for workflow run",
workflow_run_id=ctx.workflow_run_id,
verification_code_identifier=identifier_for_ui,
)
try:
await app.DATABASE.update_workflow_run(
workflow_run_id=workflow_run_id,
waiting_for_verification_code=False,
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(),
},
)
LOG.info("Cleared 2FA waiting state for workflow run", workflow_run_id=workflow_run_id)
try:
NotificationRegistryFactory.get_registry().publish(
organization_id,
{"type": "verification_code_resolved", "workflow_run_id": workflow_run_id, "task_id": 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 task_id:
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,
)
LOG.info(
"Set 2FA waiting state for task",
task_id=ctx.task_id,
verification_code_identifier=identifier_for_ui,
)
try:
await app.DATABASE.update_task_2fa_state(
task_id=task_id,
organization_id=organization_id,
waiting_for_verification_code=False,
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(),
},
)
LOG.info("Cleared 2FA waiting state for task", task_id=task_id)
try:
NotificationRegistryFactory.get_registry().publish(
organization_id,
{"type": "verification_code_resolved", "task_id": 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)
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)
async def _get_otp_value_from_url(

152
tests/unit_tests/test_otp_service_organization_context.py Normal file
View file

@ -0,0 +1,152 @@
"""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()