From 7100b7e0045518bd33c3bd832d47b28a391b8cec Mon Sep 17 00:00:00 2001 From: Marc Kelechava Date: Mon, 1 Dec 2025 16:19:37 -0800 Subject: [PATCH] store totp_identifier to credentials with fallback for login runs (#4154) --- skyvern/forge/sdk/db/client.py | 2 ++ skyvern/forge/sdk/routes/run_blocks.py | 9 ++++-- skyvern/forge/sdk/schemas/credentials.py | 15 ++++++++++ .../credential/credential_vault_service.py | 2 ++ skyvern/forge/sdk/workflow/context_manager.py | 8 +++++ skyvern/forge/sdk/workflow/models/block.py | 29 ++++++++++--------- 6 files changed, 49 insertions(+), 16 deletions(-) diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index dd1bc3c8..4f023738 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -4515,6 +4515,7 @@ class AgentDB: totp_type: str, card_last4: str | None, card_brand: str | None, + totp_identifier: str | None = None, ) -> Credential: async with self.Session() as session: credential = CredentialModel( @@ -4525,6 +4526,7 @@ class AgentDB: credential_type=credential_type, username=username, totp_type=totp_type, + totp_identifier=totp_identifier, card_last4=card_last4, card_brand=card_brand, ) diff --git a/skyvern/forge/sdk/routes/run_blocks.py b/skyvern/forge/sdk/routes/run_blocks.py index a1ba180a..c2a94af8 100644 --- a/skyvern/forge/sdk/routes/run_blocks.py +++ b/skyvern/forge/sdk/routes/run_blocks.py @@ -94,12 +94,15 @@ async def login( label = "login" yaml_parameters = [] parameter_key = "credential" + resolved_totp_identifier = login_request.totp_identifier if login_request.credential_type == CredentialType.skyvern: if not login_request.credential_id: raise HTTPException(status_code=400, detail="credential_id is required to login with Skyvern credential") credential = await app.DATABASE.get_credential(login_request.credential_id, organization.organization_id) if not credential: raise HTTPException(status_code=404, detail=f"Credential {login_request.credential_id} not found") + if not resolved_totp_identifier: + resolved_totp_identifier = credential.totp_identifier yaml_parameters = [ WorkflowParameterYAML( @@ -169,7 +172,7 @@ async def login( max_steps_per_run=10, parameter_keys=[parameter_key], totp_verification_url=totp_verification_url, - totp_identifier=login_request.totp_identifier, + totp_identifier=resolved_totp_identifier, ) yaml_blocks = [login_block_yaml] workflow_definition_yaml = WorkflowDefinitionYAML( @@ -198,7 +201,7 @@ async def login( legacy_workflow_request = WorkflowRequestBody( proxy_location=login_request.proxy_location, webhook_callback_url=webhook_url, - totp_identifier=login_request.totp_identifier, + totp_identifier=resolved_totp_identifier, totp_verification_url=totp_verification_url, browser_session_id=login_request.browser_session_id, browser_profile_id=login_request.browser_profile_id, @@ -235,7 +238,7 @@ async def login( proxy_location=login_request.proxy_location, webhook_url=webhook_url, totp_url=totp_verification_url, - totp_identifier=login_request.totp_identifier, + totp_identifier=resolved_totp_identifier, browser_session_id=login_request.browser_session_id, browser_profile_id=login_request.browser_profile_id, max_screenshot_scrolls=login_request.max_screenshot_scrolling_times, diff --git a/skyvern/forge/sdk/schemas/credentials.py b/skyvern/forge/sdk/schemas/credentials.py index 4690881b..0036ffe6 100644 --- a/skyvern/forge/sdk/schemas/credentials.py +++ b/skyvern/forge/sdk/schemas/credentials.py @@ -34,6 +34,11 @@ class PasswordCredentialResponse(BaseModel): description="Type of 2FA method used for this credential", examples=[TotpType.AUTHENTICATOR], ) + totp_identifier: str | None = Field( + default=None, + description="Identifier (email or phone number) used to fetch TOTP codes", + examples=["user@example.com", "+14155550123"], + ) class CreditCardCredentialResponse(BaseModel): @@ -58,6 +63,11 @@ class PasswordCredential(BaseModel): description="Type of 2FA method used for this credential", examples=[TotpType.AUTHENTICATOR], ) + totp_identifier: str | None = Field( + default=None, + description="Identifier (email or phone number) used to fetch TOTP codes", + examples=["user@example.com", "+14155550123"], + ) class NonEmptyPasswordCredential(PasswordCredential): @@ -155,6 +165,11 @@ class Credential(BaseModel): description="Type of 2FA method used for this credential", examples=[TotpType.AUTHENTICATOR], ) + totp_identifier: str | None = Field( + default=None, + description="Identifier (email or phone number) used to fetch TOTP codes", + examples=["user@example.com", "+14155550123"], + ) card_last4: str | None = Field(..., description="For credit_card credentials: the last four digits of the card") card_brand: str | None = Field(..., description="For credit_card credentials: the card brand") diff --git a/skyvern/forge/sdk/services/credential/credential_vault_service.py b/skyvern/forge/sdk/services/credential/credential_vault_service.py index 3e70134d..8d5ee55b 100644 --- a/skyvern/forge/sdk/services/credential/credential_vault_service.py +++ b/skyvern/forge/sdk/services/credential/credential_vault_service.py @@ -51,6 +51,7 @@ class CredentialVaultService(ABC): credential_type=data.credential_type, username=data.credential.username, totp_type=data.credential.totp_type, + totp_identifier=data.credential.totp_identifier, card_last4=None, card_brand=None, ) @@ -65,6 +66,7 @@ class CredentialVaultService(ABC): totp_type="none", card_last4=data.credential.card_number[-4:], card_brand=data.credential.card_brand, + totp_identifier=None, ) else: raise Exception(f"Unsupported credential type: {data.credential_type}") diff --git a/skyvern/forge/sdk/workflow/context_manager.py b/skyvern/forge/sdk/workflow/context_manager.py index 413eee08..e288535a 100644 --- a/skyvern/forge/sdk/workflow/context_manager.py +++ b/skyvern/forge/sdk/workflow/context_manager.py @@ -173,6 +173,7 @@ class WorkflowRunContext: self._aws_client = aws_client self.organization_id: str | None = None self.include_secrets_in_templates: bool = False + self.credential_totp_identifiers: dict[str, str] = {} def get_parameter(self, key: str) -> Parameter: return self.parameters[key] @@ -295,6 +296,10 @@ class WorkflowRunContext: credential_item = await credential_service.get_credential_item(db_credential) credential = credential_item.credential + credential_totp_identifier = getattr(credential, "totp_identifier", None) + if credential_totp_identifier: + self.credential_totp_identifiers[parameter.key] = credential_totp_identifier + self.parameters[parameter.key] = parameter self.values[parameter.key] = { "context": "These values are placeholders. When you type this in, the real value gets inserted (For security reasons)", @@ -319,6 +324,9 @@ class WorkflowRunContext: self.secrets[totp_secret_value] = parse_totp_secret(credential.totp) self.values[parameter.key]["totp"] = totp_secret_id + def get_credential_totp_identifier(self, parameter_key: str) -> str | None: + return self.credential_totp_identifiers.get(parameter_key) + async def register_secret_workflow_parameter_value( self, parameter: WorkflowParameter, diff --git a/skyvern/forge/sdk/workflow/models/block.py b/skyvern/forge/sdk/workflow/models/block.py index 102c71ac..5602f5da 100644 --- a/skyvern/forge/sdk/workflow/models/block.py +++ b/skyvern/forge/sdk/workflow/models/block.py @@ -588,19 +588,22 @@ class BaseTaskBlock(Block): ) self.url = task_url_parameter_value - if ( - self.totp_identifier - and workflow_run_context.has_parameter(self.totp_identifier) - and workflow_run_context.has_value(self.totp_identifier) - ): - totp_identifier_parameter_value = workflow_run_context.get_value(self.totp_identifier) - if totp_identifier_parameter_value: - LOG.info( - "TOTP identifier is parameterized, using parameter value", - totp_identifier_parameter_value=totp_identifier_parameter_value, - totp_identifier_parameter_key=self.totp_identifier, - ) - self.totp_identifier = totp_identifier_parameter_value + if self.totp_identifier: + if workflow_run_context.has_parameter(self.totp_identifier) and workflow_run_context.has_value( + self.totp_identifier + ): + totp_identifier_parameter_value = workflow_run_context.get_value(self.totp_identifier) + if totp_identifier_parameter_value: + self.totp_identifier = totp_identifier_parameter_value + else: + for parameter in self.get_all_parameters(workflow_run_id): + parameter_key = getattr(parameter, "key", None) + if not parameter_key: + continue + credential_totp_identifier = workflow_run_context.get_credential_totp_identifier(parameter_key) + if credential_totp_identifier: + self.totp_identifier = credential_totp_identifier + break if self.download_suffix and workflow_run_context.has_parameter(self.download_suffix): download_suffix_parameter_value = workflow_run_context.get_value(self.download_suffix)