diff --git a/skyvern/exceptions.py b/skyvern/exceptions.py index 80c0a9ac..d61e0988 100644 --- a/skyvern/exceptions.py +++ b/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: diff --git a/skyvern/forge/agent.py b/skyvern/forge/agent.py index 39fd990d..37e476cf 100644 --- a/skyvern/forge/agent.py +++ b/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, ) diff --git a/skyvern/forge/sdk/routes/__init__.py b/skyvern/forge/sdk/routes/__init__.py index d80f63ed..b463444d 100644 --- a/skyvern/forge/sdk/routes/__init__.py +++ b/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 diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py index c336f801..5255c654 100644 --- a/skyvern/forge/sdk/routes/agent_protocol.py +++ b/skyvern/forge/sdk/routes/agent_protocol.py @@ -6,19 +6,7 @@ from typing import Annotated, Any import structlog import yaml -from fastapi import ( - BackgroundTasks, - Body, - Depends, - Header, - HTTPException, - Path, - Query, - Request, - Response, - UploadFile, - status, -) +from fastapi import BackgroundTasks, 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, - ) diff --git a/skyvern/forge/sdk/routes/credentials.py b/skyvern/forge/sdk/routes/credentials.py new file mode 100644 index 00000000..b06ef422 --- /dev/null +++ b/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, + ) diff --git a/skyvern/forge/sdk/routes/totp.py b/skyvern/forge/sdk/routes/totp.py deleted file mode 100644 index c7c3fe83..00000000 --- a/skyvern/forge/sdk/routes/totp.py +++ /dev/null @@ -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) diff --git a/skyvern/forge/sdk/schemas/totp_codes.py b/skyvern/forge/sdk/schemas/totp_codes.py index 7b058d65..1c7bee77 100644 --- a/skyvern/forge/sdk/schemas/totp_codes.py +++ b/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 - task_id: str | None = None - workflow_id: str | None = None - workflow_run_id: str | None = None - source: str | None = None - content: str | None = None + totp_identifier: str | None = Field( + default=None, + 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"], + ) + 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 - content: str + totp_identifier: str = Field( + ..., + 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 - code: str - organization_id: str - created_at: datetime - modified_at: datetime + totp_code_id: str = Field(..., description="The skyvern ID of the TOTP code.") + code: str = Field(..., description="The TOTP code extracted from the content.") + organization_id: str = Field(..., description="The ID of the organization that the TOTP code is for.") + created_at: datetime = Field(..., description="The timestamp when the TOTP code was created.") + modified_at: datetime = Field(..., description="The timestamp when the TOTP code was modified.") diff --git a/skyvern/webeye/actions/handler.py b/skyvern/webeye/actions/handler.py index 436860d4..e7b0bc53 100644 --- a/skyvern/webeye/actions/handler.py +++ b/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: