Timeout and retry mechanism for Bitwarden CLI (#670)

skyvern/config.py

@@ -62,6 +62,12 @@ class Settings(BaseSettings):
     # Workflow constant parameters
     WORKFLOW_DOWNLOAD_DIRECTORY_PARAMETER_KEY: str = "SKYVERN_DOWNLOAD_DIRECTORY"
 
+    #####################
+    # Bitwarden Configs #
+    #####################
+    BITWARDEN_TIMEOUT_SECONDS: int = 60
+    BITWARDEN_MAX_RETRIES: int = 1
+
     #####################
     # LLM Configuration #
     #####################

skyvern/exceptions.py

@@ -278,6 +278,11 @@ class BitwardenLogoutError(BitwardenBaseError):
         super().__init__(f"Error logging out of Bitwarden: {message}")
 
 
+class BitwardenSyncError(BitwardenBaseError):
+    def __init__(self, message: str) -> None:
+        super().__init__(f"Error syncing Bitwarden: {message}")
+
+
 class UnknownElementTreeFormat(SkyvernException):
     def __init__(self, fmt: str) -> None:
         super().__init__(f"Unknown element tree format {fmt}")

skyvern/forge/sdk/services/bitwarden.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 import os
 import re
@@ -7,10 +8,12 @@ from enum import StrEnum
 import structlog
 import tldextract
 
+from skyvern.config import settings
 from skyvern.exceptions import (
     BitwardenListItemsError,
     BitwardenLoginError,
     BitwardenLogoutError,
+    BitwardenSyncError,
     BitwardenTOTPError,
     BitwardenUnlockError,
 )
@@ -68,7 +71,47 @@ class BitwardenService:
         return None
 
     @staticmethod
-    def get_secret_value_from_url(
+    async def get_secret_value_from_url(
+        client_id: str,
+        client_secret: str,
+        master_password: str,
+        url: str,
+        collection_id: str | None = None,
+        remaining_retries: int = settings.BITWARDEN_MAX_RETRIES,
+        fail_reasons: list[str] = [],
+    ) -> dict[str, str]:
+        """
+        Get the secret value from the Bitwarden CLI.
+        """
+        try:
+            async with asyncio.timeout(settings.BITWARDEN_TIMEOUT_SECONDS):
+                return await BitwardenService._get_secret_value_from_url(
+                    client_id=client_id,
+                    client_secret=client_secret,
+                    master_password=master_password,
+                    url=url,
+                    collection_id=collection_id,
+                )
+        except Exception as e:
+            if remaining_retries <= 0:
+                raise BitwardenListItemsError(
+                    f"Bitwarden CLI failed after all retry attempts. Fail reasons: {fail_reasons}"
+                )
+
+            remaining_retries -= 1
+            LOG.info("Retrying to get secret value from Bitwarden", remaining_retries=remaining_retries)
+            return await BitwardenService.get_secret_value_from_url(
+                client_id=client_id,
+                client_secret=client_secret,
+                master_password=master_password,
+                url=url,
+                collection_id=collection_id,
+                remaining_retries=remaining_retries,
+                fail_reasons=fail_reasons + [f"{type(e).__name__}: {str(e)}"],
+            )
+
+    @staticmethod
+    async def _get_secret_value_from_url(
         client_id: str,
         client_secret: str,
         master_password: str,
@@ -80,6 +123,7 @@ class BitwardenService:
         """
         try:
             BitwardenService.login(client_id, client_secret)
+            BitwardenService.sync()
             session_key = BitwardenService.unlock(master_password)
 
             # Extract the domain from the URL and search for items in Bitwarden with that domain
@@ -151,7 +195,50 @@ class BitwardenService:
             BitwardenService.logout()
 
     @staticmethod
-    def get_sensitive_information_from_identity(
+    async def get_sensitive_information_from_identity(
+        client_id: str,
+        client_secret: str,
+        master_password: str,
+        collection_id: str,
+        identity_key: str,
+        identity_fields: list[str],
+        remaining_retries: int = settings.BITWARDEN_MAX_RETRIES,
+        fail_reasons: list[str] = [],
+    ) -> dict[str, str]:
+        """
+        Get the secret value from the Bitwarden CLI.
+        """
+        try:
+            async with asyncio.timeout(settings.BITWARDEN_TIMEOUT_SECONDS):
+                return await BitwardenService._get_sensitive_information_from_identity(
+                    client_id=client_id,
+                    client_secret=client_secret,
+                    master_password=master_password,
+                    collection_id=collection_id,
+                    identity_key=identity_key,
+                    identity_fields=identity_fields,
+                )
+        except Exception as e:
+            if remaining_retries <= 0:
+                raise BitwardenListItemsError(
+                    f"Bitwarden CLI failed after all retry attempts. Fail reasons: {fail_reasons}"
+                )
+
+            remaining_retries -= 1
+            LOG.info("Retrying to get sensitive information from Bitwarden", remaining_retries=remaining_retries)
+            return await BitwardenService.get_sensitive_information_from_identity(
+                client_id=client_id,
+                client_secret=client_secret,
+                master_password=master_password,
+                collection_id=collection_id,
+                identity_key=identity_key,
+                identity_fields=identity_fields,
+                remaining_retries=remaining_retries,
+                fail_reasons=fail_reasons + [f"{type(e).__name__}: {str(e)}"],
+            )
+
+    @staticmethod
+    async def _get_sensitive_information_from_identity(
         client_id: str,
         client_secret: str,
         master_password: str,
@@ -164,6 +251,7 @@ class BitwardenService:
         """
         try:
             BitwardenService.login(client_id, client_secret)
+            BitwardenService.sync()
             session_key = BitwardenService.unlock(master_password)
 
             # Step 3: Retrieve the items
@@ -267,6 +355,18 @@ class BitwardenService:
 
         return session_key
 
+    @staticmethod
+    def sync() -> None:
+        """
+        Sync the Bitwarden CLI.
+        """
+        sync_command = ["bw", "sync"]
+        LOG.info("Bitwarden CLI sync started")
+        sync_result = BitwardenService.run_command(sync_command)
+        LOG.info("Bitwarden CLI sync completed")
+        if sync_result.stderr:
+            raise BitwardenSyncError(sync_result.stderr)
+
     @staticmethod
     def logout() -> None:
         """
@@ -274,5 +374,5 @@ class BitwardenService:
         """
         logout_command = ["bw", "logout"]
         logout_result = BitwardenService.run_command(logout_command)
-        if logout_result.stderr:
+        if logout_result.stderr and "You are not logged in." not in logout_result.stderr:
             raise BitwardenLogoutError(logout_result.stderr)

skyvern/forge/sdk/workflow/context_manager.py

@@ -97,11 +97,11 @@ class WorkflowRunContext:
             return self.secrets.get(secret_id_or_value)
         return None
 
-    def get_secrets_from_password_manager(self) -> dict[str, Any]:
+    async def get_secrets_from_password_manager(self) -> dict[str, Any]:
         """
         Get the secrets from the password manager. The secrets dict will contain the actual secret values.
         """
-        secret_credentials = BitwardenService.get_secret_value_from_url(
+        secret_credentials = await BitwardenService.get_secret_value_from_url(
             url=self.secrets[BitwardenConstants.URL],
             client_secret=self.secrets[BitwardenConstants.CLIENT_SECRET],
             client_id=self.secrets[BitwardenConstants.CLIENT_ID],
@@ -161,7 +161,7 @@ class WorkflowRunContext:
                     collection_id = parameter.bitwarden_collection_id
 
             try:
-                secret_credentials = BitwardenService.get_secret_value_from_url(
+                secret_credentials = await BitwardenService.get_secret_value_from_url(
                     client_id,
                     client_secret,
                     master_password,
@@ -218,7 +218,7 @@ class WorkflowRunContext:
                 collection_id = self.values[parameter.bitwarden_collection_id]
 
             try:
-                sensitive_values = BitwardenService.get_sensitive_information_from_identity(
+                sensitive_values = await BitwardenService.get_sensitive_information_from_identity(
                     client_id,
                     client_secret,
                     master_password,

skyvern/webeye/actions/handler.py

@@ -706,7 +706,7 @@ async def get_actual_value_of_parameter_if_secret(task: Task, parameter: str) ->
     secret_value = workflow_run_context.get_original_secret_value_or_none(parameter)
 
     if secret_value == BitwardenConstants.TOTP:
-        secrets = workflow_run_context.get_secrets_from_password_manager()
+        secrets = await workflow_run_context.get_secrets_from_password_manager()
         secret_value = secrets[BitwardenConstants.TOTP]
     return secret_value if secret_value is not None else parameter