official totp code endpoint under the credentials route (#2357)

skyvern/exceptions.py

@@ -660,6 +660,7 @@ class NoTOTPVerificationCodeFound(SkyvernHTTPException):
         self,
         task_id: str | None = None,
         workflow_run_id: str | None = None,
+        workflow_id: str | None = None,
         totp_verification_url: str | None = None,
         totp_identifier: str | None = None,
     ) -> None:
@@ -668,6 +669,8 @@ class NoTOTPVerificationCodeFound(SkyvernHTTPException):
             msg += f" task_id={task_id}"
         if workflow_run_id:
             msg += f" workflow_run_id={workflow_run_id}"
+        if workflow_id:
+            msg += f" workflow_id={workflow_id}"
         if totp_verification_url:
             msg += f" totp_verification_url={totp_verification_url}"
         if totp_identifier:

skyvern/forge/agent.py

@@ -2631,10 +2631,18 @@ class ForgeAgent:
             and task.organization_id
         ):
             LOG.info("Need verification code", step_id=step.step_id)
+            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
             verification_code = await poll_verification_code(
                 task.task_id,
                 task.organization_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,
             )

skyvern/forge/sdk/routes/__init__.py

@@ -1,4 +1,4 @@
 from skyvern.forge.sdk.routes import agent_protocol  # noqa: F401
 from skyvern.forge.sdk.routes import browser_sessions  # noqa: F401
+from skyvern.forge.sdk.routes import credentials  # noqa: F401
 from skyvern.forge.sdk.routes import streaming  # noqa: F401
-from skyvern.forge.sdk.routes import totp  # noqa: F401

skyvern/forge/sdk/routes/agent_protocol.py

@@ -6,19 +6,7 @@ from typing import Annotated, Any
 
 import structlog
 import yaml
-from fastapi import (
+from fastapi import BackgroundTasks, Depends, Header, HTTPException, Path, Query, Request, Response, UploadFile, status
-    BackgroundTasks,
-    Body,
-    Depends,
-    Header,
-    HTTPException,
-    Path,
-    Query,
-    Request,
-    Response,
-    UploadFile,
-    status,
-)
 from fastapi.responses import ORJSONResponse
 
 from skyvern import analytics
@@ -37,13 +25,6 @@ from skyvern.forge.sdk.executor.factory import AsyncExecutorFactory
 from skyvern.forge.sdk.models import Step
 from skyvern.forge.sdk.routes.routers import base_router, legacy_base_router, legacy_v2_router
 from skyvern.forge.sdk.schemas.ai_suggestions import AISuggestionBase, AISuggestionRequest
-from skyvern.forge.sdk.schemas.credentials import (
-    CreateCredentialRequest,
-    CredentialResponse,
-    CredentialType,
-    CreditCardCredentialResponse,
-    PasswordCredentialResponse,
-)
 from skyvern.forge.sdk.schemas.organizations import (
     GetOrganizationAPIKeysResponse,
     GetOrganizationsResponse,
@@ -63,7 +44,6 @@ from skyvern.forge.sdk.schemas.tasks import (
 )
 from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunTimeline
 from skyvern.forge.sdk.services import org_auth_service
-from skyvern.forge.sdk.services.bitwarden import BitwardenService
 from skyvern.forge.sdk.workflow.exceptions import (
     FailedToCreateWorkflow,
     FailedToUpdateWorkflow,
@@ -1706,253 +1686,3 @@ async def cancel_run(
     analytics.capture("skyvern-oss-agent-cancel-run")
 
     await run_service.cancel_run(run_id, organization_id=current_org.organization_id, api_key=x_api_key)
-
-
-@legacy_base_router.get("/credentials")
-@legacy_base_router.get("/credentials/", include_in_schema=False)
-@base_router.get(
-    "/credentials",
-    response_model=list[CredentialResponse],
-    summary="Get all credentials",
-    description="Retrieves a paginated list of credentials for the current organization",
-    tags=["Credentials"],
-    openapi_extra={
-        "x-fern-sdk-group-name": "credentials",
-        "x-fern-sdk-method-name": "get_credentials",
-    },
-)
-async def get_credentials(
-    current_org: Organization = Depends(org_auth_service.get_current_org),
-    page: int = Query(
-        1,
-        ge=1,
-        description="Page number for pagination",
-        example=1,
-        openapi_extra={"x-fern-sdk-parameter-name": "page"},
-    ),
-    page_size: int = Query(
-        10,
-        ge=1,
-        description="Number of items per page",
-        example=10,
-        openapi_extra={"x-fern-sdk-parameter-name": "page_size"},
-    ),
-) -> list[CredentialResponse]:
-    organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection(
-        current_org.organization_id
-    )
-    if not organization_bitwarden_collection:
-        return []
-
-    credentials = await app.DATABASE.get_credentials(current_org.organization_id, page=page, page_size=page_size)
-    items = await BitwardenService.get_collection_items(organization_bitwarden_collection.collection_id)
-
-    response_items = []
-    for credential in credentials:
-        item = next((item for item in items if item.item_id == credential.item_id), None)
-        if not item:
-            continue
-        if item.credential_type == CredentialType.PASSWORD:
-            credential_response = PasswordCredentialResponse(username=item.credential.username)
-            response_items.append(
-                CredentialResponse(
-                    credential=credential_response,
-                    credential_id=credential.credential_id,
-                    credential_type=item.credential_type,
-                    name=item.name,
-                )
-            )
-        elif item.credential_type == CredentialType.CREDIT_CARD:
-            credential_response = CreditCardCredentialResponse(
-                last_four=item.credential.card_number[-4:],
-                brand=item.credential.card_brand,
-            )
-            response_items.append(
-                CredentialResponse(
-                    credential=credential_response,
-                    credential_id=credential.credential_id,
-                    credential_type=item.credential_type,
-                    name=item.name,
-                )
-            )
-    return response_items
-
-
-@legacy_base_router.get("/credentials/{credential_id}")
-@legacy_base_router.get("/credentials/{credential_id}/", include_in_schema=False)
-@base_router.get(
-    "/credentials/{credential_id}",
-    response_model=CredentialResponse,
-    summary="Get credential by ID",
-    description="Retrieves a specific credential by its ID",
-    tags=["Credentials"],
-    openapi_extra={
-        "x-fern-sdk-group-name": "credentials",
-        "x-fern-sdk-method-name": "get_credential",
-    },
-)
-async def get_credential(
-    credential_id: str = Path(
-        ...,
-        description="The unique identifier of the credential",
-        example="cred_1234567890",
-        openapi_extra={"x-fern-sdk-parameter-name": "credential_id"},
-    ),
-    current_org: Organization = Depends(org_auth_service.get_current_org),
-) -> CredentialResponse:
-    organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection(
-        current_org.organization_id
-    )
-    if not organization_bitwarden_collection:
-        raise HTTPException(status_code=404, detail="Credential account not found. It might have been deleted.")
-
-    credential = await app.DATABASE.get_credential(
-        credential_id=credential_id, organization_id=current_org.organization_id
-    )
-    if not credential:
-        raise HTTPException(status_code=404, detail="Credential not found")
-
-    credential_item = await BitwardenService.get_credential_item(credential.item_id)
-    if not credential_item:
-        raise HTTPException(status_code=404, detail="Credential not found")
-
-    if credential_item.credential_type == CredentialType.PASSWORD:
-        credential_response = PasswordCredentialResponse(
-            username=credential_item.credential.username,
-        )
-        return CredentialResponse(
-            credential=credential_response,
-            credential_id=credential.credential_id,
-            credential_type=credential_item.credential_type,
-            name=credential_item.name,
-        )
-    if credential_item.credential_type == CredentialType.CREDIT_CARD:
-        credential_response = CreditCardCredentialResponse(
-            last_four=credential_item.credential.card_number[-4:],
-            brand=credential_item.credential.card_brand,
-        )
-        return CredentialResponse(
-            credential=credential_response,
-            credential_id=credential.credential_id,
-            credential_type=credential_item.credential_type,
-            name=credential_item.name,
-        )
-    raise HTTPException(status_code=400, detail="Invalid credential type")
-
-
-@legacy_base_router.delete("/credentials/{credential_id}")
-@legacy_base_router.delete("/credentials/{credential_id}/", include_in_schema=False)
-@base_router.post(
-    "/credentials/{credential_id}/delete",
-    status_code=204,
-    summary="Delete credential",
-    description="Deletes a specific credential by its ID",
-    tags=["Credentials"],
-    openapi_extra={
-        "x-fern-sdk-group-name": "credentials",
-        "x-fern-sdk-method-name": "delete_credential",
-    },
-)
-async def delete_credential(
-    credential_id: str = Path(
-        ...,
-        description="The unique identifier of the credential to delete",
-        example="cred_1234567890",
-        openapi_extra={"x-fern-sdk-parameter-name": "credential_id"},
-    ),
-    current_org: Organization = Depends(org_auth_service.get_current_org),
-) -> None:
-    organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection(
-        current_org.organization_id
-    )
-    if not organization_bitwarden_collection:
-        raise HTTPException(status_code=404, detail="Credential account not found. It might have been deleted.")
-
-    credential = await app.DATABASE.get_credential(
-        credential_id=credential_id, organization_id=current_org.organization_id
-    )
-    if not credential:
-        raise HTTPException(status_code=404, detail=f"Credential not found, credential_id={credential_id}")
-
-    await app.DATABASE.delete_credential(credential.credential_id, current_org.organization_id)
-    await BitwardenService.delete_credential_item(credential.item_id)
-
-    return None
-
-
-@legacy_base_router.post("/credentials")
-@legacy_base_router.post("/credentials/", include_in_schema=False)
-@base_router.post(
-    "/credentials",
-    response_model=CredentialResponse,
-    status_code=201,
-    summary="Create credential",
-    description="Creates a new credential for the current organization",
-    tags=["Credentials"],
-    openapi_extra={
-        "x-fern-sdk-group-name": "credentials",
-        "x-fern-sdk-method-name": "create_credential",
-    },
-)
-async def create_credential(
-    data: CreateCredentialRequest = Body(
-        ...,
-        description="The credential data to create",
-        example={
-            "name": "My Credential",
-            "credential_type": "PASSWORD",
-            "credential": {"username": "user@example.com", "password": "securepassword123", "totp": "JBSWY3DPEHPK3PXP"},
-        },
-        openapi_extra={"x-fern-sdk-parameter-name": "data"},
-    ),
-    current_org: Organization = Depends(org_auth_service.get_current_org),
-) -> CredentialResponse:
-    org_collection = await app.DATABASE.get_organization_bitwarden_collection(current_org.organization_id)
-
-    if not org_collection:
-        LOG.info(
-            "There is no collection for the organization. Creating new collection.",
-            organization_id=current_org.organization_id,
-        )
-        collection_id = await BitwardenService.create_collection(
-            name=current_org.organization_id,
-        )
-        org_collection = await app.DATABASE.create_organization_bitwarden_collection(
-            current_org.organization_id,
-            collection_id,
-        )
-
-    item_id = await BitwardenService.create_credential_item(
-        collection_id=org_collection.collection_id,
-        name=data.name,
-        credential=data.credential,
-    )
-
-    credential = await app.DATABASE.create_credential(
-        organization_id=current_org.organization_id,
-        item_id=item_id,
-        name=data.name,
-        credential_type=data.credential_type,
-    )
-
-    if data.credential_type == CredentialType.PASSWORD:
-        credential_response = PasswordCredentialResponse(
-            username=data.credential.username,
-        )
-        return CredentialResponse(
-            credential=credential_response,
-            credential_id=credential.credential_id,
-            credential_type=data.credential_type,
-            name=data.name,
-        )
-    elif data.credential_type == CredentialType.CREDIT_CARD:
-        credential_response = CreditCardCredentialResponse(
-            last_four=data.credential.card_number[-4:],
-            brand=data.credential.card_brand,
-        )
-        return CredentialResponse(
-            credential=credential_response,
-            credential_id=credential.credential_id,
-            credential_type=data.credential_type,
-            name=data.name,
-        )

skyvern/forge/sdk/routes/credentials.py

@@ -0,0 +1,321 @@
+import structlog
+from fastapi import Body, Depends, HTTPException, Path, Query
+
+from skyvern.forge import app
+from skyvern.forge.prompts import prompt_engine
+from skyvern.forge.sdk.routes.routers import base_router, legacy_base_router
+from skyvern.forge.sdk.schemas.credentials import (
+    CreateCredentialRequest,
+    CredentialResponse,
+    CredentialType,
+    CreditCardCredentialResponse,
+    PasswordCredentialResponse,
+)
+from skyvern.forge.sdk.schemas.organizations import Organization
+from skyvern.forge.sdk.schemas.totp_codes import TOTPCode, TOTPCodeCreate
+from skyvern.forge.sdk.services import org_auth_service
+from skyvern.forge.sdk.services.bitwarden import BitwardenService
+
+LOG = structlog.get_logger()
+
+
+@legacy_base_router.post(
+    "/totp",
+    tags=["agent"],
+    openapi_extra={
+        "x-fern-sdk-group-name": "agent",
+        "x-fern-sdk-method-name": "send_totp_code",
+    },
+)
+@legacy_base_router.post("/totp/", include_in_schema=False)
+@base_router.post(
+    "/credentials/totp",
+    response_model=TOTPCode,
+    summary="Send TOTP code",
+    description="Send a TOTP code to the user",
+    tags=["Credentials"],
+    openapi_extra={
+        "x-fern-sdk-group-name": "credentials",
+        "x-fern-sdk-method-name": "send_totp_code",
+    },
+)
+async def send_totp_code(
+    data: TOTPCodeCreate, curr_org: Organization = Depends(org_auth_service.get_current_org)
+) -> TOTPCode:
+    LOG.info(
+        "Saving TOTP code",
+        organization_id=curr_org.organization_id,
+        totp_identifier=data.totp_identifier,
+        task_id=data.task_id,
+        workflow_id=data.workflow_id,
+    )
+    code = await parse_totp_code(data.content)
+    if not code:
+        raise HTTPException(status_code=400, detail="Failed to parse totp code")
+    return await app.DATABASE.create_totp_code(
+        organization_id=curr_org.organization_id,
+        totp_identifier=data.totp_identifier,
+        content=data.content,
+        code=code,
+        task_id=data.task_id,
+        workflow_id=data.workflow_id,
+        workflow_run_id=data.workflow_run_id,
+        source=data.source,
+        expired_at=data.expired_at,
+    )
+
+
+async def parse_totp_code(content: str) -> str | None:
+    prompt = prompt_engine.load_prompt("parse-verification-code", content=content)
+    code_resp = await app.SECONDARY_LLM_API_HANDLER(prompt=prompt, prompt_name="parse-verification-code")
+    return code_resp.get("code", None)
+
+
+@legacy_base_router.get("/credentials")
+@legacy_base_router.get("/credentials/", include_in_schema=False)
+@base_router.get(
+    "/credentials",
+    response_model=list[CredentialResponse],
+    summary="Get all credentials",
+    description="Retrieves a paginated list of credentials for the current organization",
+    tags=["Credentials"],
+    openapi_extra={
+        "x-fern-sdk-group-name": "credentials",
+        "x-fern-sdk-method-name": "get_credentials",
+    },
+)
+async def get_credentials(
+    current_org: Organization = Depends(org_auth_service.get_current_org),
+    page: int = Query(
+        1,
+        ge=1,
+        description="Page number for pagination",
+        example=1,
+        openapi_extra={"x-fern-sdk-parameter-name": "page"},
+    ),
+    page_size: int = Query(
+        10,
+        ge=1,
+        description="Number of items per page",
+        example=10,
+        openapi_extra={"x-fern-sdk-parameter-name": "page_size"},
+    ),
+) -> list[CredentialResponse]:
+    organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection(
+        current_org.organization_id
+    )
+    if not organization_bitwarden_collection:
+        return []
+
+    credentials = await app.DATABASE.get_credentials(current_org.organization_id, page=page, page_size=page_size)
+    items = await BitwardenService.get_collection_items(organization_bitwarden_collection.collection_id)
+
+    response_items = []
+    for credential in credentials:
+        item = next((item for item in items if item.item_id == credential.item_id), None)
+        if not item:
+            continue
+        if item.credential_type == CredentialType.PASSWORD:
+            credential_response = PasswordCredentialResponse(username=item.credential.username)
+            response_items.append(
+                CredentialResponse(
+                    credential=credential_response,
+                    credential_id=credential.credential_id,
+                    credential_type=item.credential_type,
+                    name=item.name,
+                )
+            )
+        elif item.credential_type == CredentialType.CREDIT_CARD:
+            credential_response = CreditCardCredentialResponse(
+                last_four=item.credential.card_number[-4:],
+                brand=item.credential.card_brand,
+            )
+            response_items.append(
+                CredentialResponse(
+                    credential=credential_response,
+                    credential_id=credential.credential_id,
+                    credential_type=item.credential_type,
+                    name=item.name,
+                )
+            )
+    return response_items
+
+
+@legacy_base_router.get("/credentials/{credential_id}")
+@legacy_base_router.get("/credentials/{credential_id}/", include_in_schema=False)
+@base_router.get(
+    "/credentials/{credential_id}",
+    response_model=CredentialResponse,
+    summary="Get credential by ID",
+    description="Retrieves a specific credential by its ID",
+    tags=["Credentials"],
+    openapi_extra={
+        "x-fern-sdk-group-name": "credentials",
+        "x-fern-sdk-method-name": "get_credential",
+    },
+)
+async def get_credential(
+    credential_id: str = Path(
+        ...,
+        description="The unique identifier of the credential",
+        example="cred_1234567890",
+        openapi_extra={"x-fern-sdk-parameter-name": "credential_id"},
+    ),
+    current_org: Organization = Depends(org_auth_service.get_current_org),
+) -> CredentialResponse:
+    organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection(
+        current_org.organization_id
+    )
+    if not organization_bitwarden_collection:
+        raise HTTPException(status_code=404, detail="Credential account not found. It might have been deleted.")
+
+    credential = await app.DATABASE.get_credential(
+        credential_id=credential_id, organization_id=current_org.organization_id
+    )
+    if not credential:
+        raise HTTPException(status_code=404, detail="Credential not found")
+
+    credential_item = await BitwardenService.get_credential_item(credential.item_id)
+    if not credential_item:
+        raise HTTPException(status_code=404, detail="Credential not found")
+
+    if credential_item.credential_type == CredentialType.PASSWORD:
+        credential_response = PasswordCredentialResponse(
+            username=credential_item.credential.username,
+        )
+        return CredentialResponse(
+            credential=credential_response,
+            credential_id=credential.credential_id,
+            credential_type=credential_item.credential_type,
+            name=credential_item.name,
+        )
+    if credential_item.credential_type == CredentialType.CREDIT_CARD:
+        credential_response = CreditCardCredentialResponse(
+            last_four=credential_item.credential.card_number[-4:],
+            brand=credential_item.credential.card_brand,
+        )
+        return CredentialResponse(
+            credential=credential_response,
+            credential_id=credential.credential_id,
+            credential_type=credential_item.credential_type,
+            name=credential_item.name,
+        )
+    raise HTTPException(status_code=400, detail="Invalid credential type")
+
+
+@legacy_base_router.delete("/credentials/{credential_id}")
+@legacy_base_router.delete("/credentials/{credential_id}/", include_in_schema=False)
+@base_router.post(
+    "/credentials/{credential_id}/delete",
+    status_code=204,
+    summary="Delete credential",
+    description="Deletes a specific credential by its ID",
+    tags=["Credentials"],
+    openapi_extra={
+        "x-fern-sdk-group-name": "credentials",
+        "x-fern-sdk-method-name": "delete_credential",
+    },
+)
+async def delete_credential(
+    credential_id: str = Path(
+        ...,
+        description="The unique identifier of the credential to delete",
+        example="cred_1234567890",
+        openapi_extra={"x-fern-sdk-parameter-name": "credential_id"},
+    ),
+    current_org: Organization = Depends(org_auth_service.get_current_org),
+) -> None:
+    organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection(
+        current_org.organization_id
+    )
+    if not organization_bitwarden_collection:
+        raise HTTPException(status_code=404, detail="Credential account not found. It might have been deleted.")
+
+    credential = await app.DATABASE.get_credential(
+        credential_id=credential_id, organization_id=current_org.organization_id
+    )
+    if not credential:
+        raise HTTPException(status_code=404, detail=f"Credential not found, credential_id={credential_id}")
+
+    await app.DATABASE.delete_credential(credential.credential_id, current_org.organization_id)
+    await BitwardenService.delete_credential_item(credential.item_id)
+
+    return None
+
+
+@legacy_base_router.post("/credentials")
+@legacy_base_router.post("/credentials/", include_in_schema=False)
+@base_router.post(
+    "/credentials",
+    response_model=CredentialResponse,
+    status_code=201,
+    summary="Create credential",
+    description="Creates a new credential for the current organization",
+    tags=["Credentials"],
+    openapi_extra={
+        "x-fern-sdk-group-name": "credentials",
+        "x-fern-sdk-method-name": "create_credential",
+    },
+)
+async def create_credential(
+    data: CreateCredentialRequest = Body(
+        ...,
+        description="The credential data to create",
+        example={
+            "name": "My Credential",
+            "credential_type": "PASSWORD",
+            "credential": {"username": "user@example.com", "password": "securepassword123", "totp": "JBSWY3DPEHPK3PXP"},
+        },
+        openapi_extra={"x-fern-sdk-parameter-name": "data"},
+    ),
+    current_org: Organization = Depends(org_auth_service.get_current_org),
+) -> CredentialResponse:
+    org_collection = await app.DATABASE.get_organization_bitwarden_collection(current_org.organization_id)
+
+    if not org_collection:
+        LOG.info(
+            "There is no collection for the organization. Creating new collection.",
+            organization_id=current_org.organization_id,
+        )
+        collection_id = await BitwardenService.create_collection(
+            name=current_org.organization_id,
+        )
+        org_collection = await app.DATABASE.create_organization_bitwarden_collection(
+            current_org.organization_id,
+            collection_id,
+        )
+
+    item_id = await BitwardenService.create_credential_item(
+        collection_id=org_collection.collection_id,
+        name=data.name,
+        credential=data.credential,
+    )
+
+    credential = await app.DATABASE.create_credential(
+        organization_id=current_org.organization_id,
+        item_id=item_id,
+        name=data.name,
+        credential_type=data.credential_type,
+    )
+
+    if data.credential_type == CredentialType.PASSWORD:
+        credential_response = PasswordCredentialResponse(
+            username=data.credential.username,
+        )
+        return CredentialResponse(
+            credential=credential_response,
+            credential_id=credential.credential_id,
+            credential_type=data.credential_type,
+            name=data.name,
+        )
+    elif data.credential_type == CredentialType.CREDIT_CARD:
+        credential_response = CreditCardCredentialResponse(
+            last_four=data.credential.card_number[-4:],
+            brand=data.credential.card_brand,
+        )
+        return CredentialResponse(
+            credential=credential_response,
+            credential_id=credential.credential_id,
+            credential_type=data.credential_type,
+            name=data.name,
+        )

skyvern/forge/sdk/routes/totp.py

@@ -1,53 +0,0 @@
-import structlog
-from fastapi import Depends, HTTPException
-
-from skyvern.forge import app
-from skyvern.forge.prompts import prompt_engine
-from skyvern.forge.sdk.routes.routers import legacy_base_router
-from skyvern.forge.sdk.schemas.organizations import Organization
-from skyvern.forge.sdk.schemas.totp_codes import TOTPCode, TOTPCodeCreate
-from skyvern.forge.sdk.services import org_auth_service
-
-LOG = structlog.get_logger()
-
-
-@legacy_base_router.post(
-    "/totp",
-    tags=["agent"],
-    openapi_extra={
-        "x-fern-sdk-group-name": "agent",
-        "x-fern-sdk-method-name": "send_totp_code",
-    },
-)
-@legacy_base_router.post("/totp/", include_in_schema=False)
-async def send_totp_code(
-    data: TOTPCodeCreate, curr_org: Organization = Depends(org_auth_service.get_current_org)
-) -> TOTPCode:
-    LOG.info(
-        "Saving TOTP code",
-        data=data,
-        organization_id=curr_org.organization_id,
-        totp_identifier=data.totp_identifier,
-        task_id=data.task_id,
-        workflow_id=data.workflow_id,
-    )
-    code = await parse_totp_code(data.content)
-    if not code:
-        raise HTTPException(status_code=400, detail="Failed to parse totp code")
-    return await app.DATABASE.create_totp_code(
-        organization_id=curr_org.organization_id,
-        totp_identifier=data.totp_identifier,
-        content=data.content,
-        code=code,
-        task_id=data.task_id,
-        workflow_id=data.workflow_id,
-        workflow_run_id=data.workflow_run_id,
-        source=data.source,
-        expired_at=data.expired_at,
-    )
-
-
-async def parse_totp_code(content: str) -> str | None:
-    prompt = prompt_engine.load_prompt("parse-verification-code", content=content)
-    code_resp = await app.SECONDARY_LLM_API_HANDLER(prompt=prompt, prompt_name="parse-verification-code")
-    return code_resp.get("code", None)

skyvern/forge/sdk/schemas/totp_codes.py

@@ -1,6 +1,6 @@
 from datetime import datetime
 
-from pydantic import BaseModel, ConfigDict, field_validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 from skyvern.forge.sdk.utils.sanitization import sanitize_postgres_text
 
@@ -8,19 +8,55 @@ from skyvern.forge.sdk.utils.sanitization import sanitize_postgres_text
 class TOTPCodeBase(BaseModel):
     model_config = ConfigDict(from_attributes=True)
 
-    totp_identifier: str | None = None
+    totp_identifier: str | None = Field(
-    task_id: str | None = None
+        default=None,
-    workflow_id: str | None = None
+        description="The identifier of the TOTP code. It can be the email address, phone number, or the identifier of the user.",
-    workflow_run_id: str | None = None
+        examples=["john.doe@example.com", "4155555555", "user_123"],
-    source: str | None = None
+    )
-    content: str | None = None
+    task_id: str | None = Field(
+        default=None,
+        description="The task_id the totp code is for. It can be the task_id of the task that the TOTP code is for.",
+        examples=["task_123456"],
+    )
+    workflow_id: str | None = Field(
+        default=None,
+        description="The workflow ID the TOTP code is for. It can be the workflow ID of the workflow that the TOTP code is for.",
+        examples=["wpid_123456"],
+    )
+    workflow_run_id: str | None = Field(
+        default=None,
+        description="The workflow run id that the TOTP code is for. It can be the workflow run id of the workflow run that the TOTP code is for.",
+        examples=["wr_123456"],
+    )
+    source: str | None = Field(
+        default=None,
+        description="An optional field. The source of the TOTP code. e.g. email, sms, etc.",
+        examples=["email", "sms", "app"],
+    )
+    content: str | None = Field(
+        default=None,
+        description="The content of the TOTP code. It can be the email content that contains the TOTP code, or the sms message that contains the TOTP code. Skyvern will automatically extract the TOTP code from the content.",
+        examples=["Hello, your verification code is 123456"],
+    )
 
-    expired_at: datetime | None = None
+    expired_at: datetime | None = Field(
+        default=None,
+        description="The timestamp when the TOTP code expires",
+        examples=["2025-01-01T00:00:00Z"],
+    )
 
 
 class TOTPCodeCreate(TOTPCodeBase):
-    totp_identifier: str
+    totp_identifier: str = Field(
-    content: str
+        ...,
+        description="The identifier of the TOTP code. It can be the email address, phone number, or the identifier of the user.",
+        examples=["john.doe@example.com", "4155555555", "user_123"],
+    )
+    content: str = Field(
+        ...,
+        description="The content of the TOTP code. It can be the email content that contains the TOTP code, or the sms message that contains the TOTP code. Skyvern will automatically extract the TOTP code from the content.",
+        examples=["Hello, your verification code is 123456"],
+    )
 
     @field_validator("content")
     @classmethod
@@ -30,8 +66,8 @@ class TOTPCodeCreate(TOTPCodeBase):
 
 
 class TOTPCode(TOTPCodeCreate):
-    totp_code_id: str
+    totp_code_id: str = Field(..., description="The skyvern ID of the TOTP code.")
-    code: str
+    code: str = Field(..., description="The TOTP code extracted from the content.")
-    organization_id: str
+    organization_id: str = Field(..., description="The ID of the organization that the TOTP code is for.")
-    created_at: datetime
+    created_at: datetime = Field(..., description="The timestamp when the TOTP code was created.")
-    modified_at: datetime
+    modified_at: datetime = Field(..., description="The timestamp when the TOTP code was modified.")

skyvern/webeye/actions/handler.py

@@ -3319,10 +3319,11 @@ async def poll_verification_code(
     while True:
         # check timeout
         if datetime.utcnow() > timeout_datetime:
-            LOG.warning("Polling verification code timed out", workflow_id=workflow_id)
+            LOG.warning("Polling verification code 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,
             )
@@ -3339,7 +3340,7 @@ async def poll_verification_code(
                 task_id,
                 organization_id,
                 totp_identifier,
-                workflow_id=workflow_id,
+                workflow_id=workflow_permanent_id,
                 workflow_run_id=workflow_run_id,
             )
         if verification_code: