diff --git a/skyvern/forge/sdk/workflow/context_manager.py b/skyvern/forge/sdk/workflow/context_manager.py index e288535a..cc3d59ea 100644 --- a/skyvern/forge/sdk/workflow/context_manager.py +++ b/skyvern/forge/sdk/workflow/context_manager.py @@ -296,7 +296,7 @@ class WorkflowRunContext: credential_item = await credential_service.get_credential_item(db_credential) credential = credential_item.credential - credential_totp_identifier = getattr(credential, "totp_identifier", None) + credential_totp_identifier = db_credential.totp_identifier or getattr(credential, "totp_identifier", None) if credential_totp_identifier: self.credential_totp_identifiers[parameter.key] = credential_totp_identifier diff --git a/skyvern/forge/sdk/workflow/models/block.py b/skyvern/forge/sdk/workflow/models/block.py index 7676b0f9..369bfadd 100644 --- a/skyvern/forge/sdk/workflow/models/block.py +++ b/skyvern/forge/sdk/workflow/models/block.py @@ -3493,6 +3493,13 @@ class TaskV2Block(Block): max_iterations: int = settings.MAX_ITERATIONS_PER_TASK_V2 max_steps: int = settings.MAX_STEPS_PER_TASK_V2 + def _resolve_totp_identifier(self, workflow_run_context: WorkflowRunContext) -> str | None: + if self.totp_identifier: + return self.totp_identifier + if workflow_run_context.credential_totp_identifiers: + return next(iter(workflow_run_context.credential_totp_identifiers.values()), None) + return None + def get_all_parameters( self, workflow_run_id: str, @@ -3535,7 +3542,7 @@ class TaskV2Block(Block): # Use the resolved values directly resolved_prompt = self.prompt resolved_url = self.url - resolved_totp_identifier = self.totp_identifier + resolved_totp_identifier = self._resolve_totp_identifier(workflow_run_context) resolved_totp_verification_url = self.totp_verification_url except Exception as e: