Extract CredentialVaultService (#3669)

skyvern/forge/app.py

@@ -14,6 +14,8 @@ from skyvern.forge.sdk.cache.factory import CacheFactory
 from skyvern.forge.sdk.db.client import AgentDB
 from skyvern.forge.sdk.experimentation.providers import BaseExperimentationProvider, NoOpExperimentationProvider
 from skyvern.forge.sdk.schemas.organizations import Organization
+from skyvern.forge.sdk.services.credential.bitwarden_credential_service import BitwardenCredentialVaultService
+from skyvern.forge.sdk.services.credential.credential_vault_service import CredentialVaultService
 from skyvern.forge.sdk.settings_manager import SettingsManager
 from skyvern.forge.sdk.trace import TraceManager
 from skyvern.forge.sdk.trace.lmnr import LaminarTrace
@@ -94,6 +96,8 @@ WORKFLOW_CONTEXT_MANAGER = WorkflowContextManager()
 WORKFLOW_SERVICE = WorkflowService()
 AGENT_FUNCTION = AgentFunction()
 PERSISTENT_SESSIONS_MANAGER = PersistentSessionsManager(database=DATABASE)
+CREDENTIAL_VAULT_SERVICE: CredentialVaultService = BitwardenCredentialVaultService()
+
 scrape_exclude: ScrapeExcludeFunc | None = None
 authentication_function: Callable[[str], Awaitable[Organization]] | None = None
 authenticate_user_function: Callable[[str], Awaitable[str | None]] | None = None

skyvern/forge/sdk/routes/credentials.py

@@ -155,37 +155,12 @@ async def create_credential(
     ),
     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,
-        totp_type=data.credential.totp_type if hasattr(data.credential, "totp_type") else "none",
+    credential = await app.CREDENTIAL_VAULT_SERVICE.create_credential(
+        organization_id=current_org.organization_id, data=data
     )
 
     # Early resyncing the Bitwarden vault
-    background_tasks.add_task(fetch_credential_item_background, item_id)
+    background_tasks.add_task(fetch_credential_item_background, credential.item_id)
 
     if data.credential_type == CredentialType.PASSWORD:
         credential_response = PasswordCredentialResponse(
@@ -209,6 +184,8 @@ async def create_credential(
             credential_type=data.credential_type,
             name=data.name,
         )
+    else:
+        raise HTTPException(status_code=400, detail=f"Unsupported credential type: {data.credential_type}")
 
 
 @legacy_base_router.delete("/credentials/{credential_id}")
@@ -238,20 +215,13 @@ async def delete_credential(
     ),
     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)
+    await app.CREDENTIAL_VAULT_SERVICE.delete_credential(credential)
 
     return None
 
@@ -283,45 +253,7 @@ async def get_credential(
     ),
     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,
-            totp_type=credential.totp_type,
-        )
-        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")
+    return await app.CREDENTIAL_VAULT_SERVICE.get_credential(current_org.organization_id, credential_id)
 
 
 @legacy_base_router.get("/credentials")
@@ -359,47 +291,7 @@ async def get_credentials(
         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,
-                totp_type=credential.totp_type,
-            )
-            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
+    return await app.CREDENTIAL_VAULT_SERVICE.get_credentials(current_org.organization_id, page, page_size)
 
 
 @base_router.get(

skyvern/forge/sdk/services/credential/bitwarden_credential_service.py

@@ -0,0 +1,150 @@
+import structlog
+from fastapi import HTTPException
+
+from skyvern.forge import app
+from skyvern.forge.sdk.schemas.credentials import (
+    CreateCredentialRequest,
+    Credential,
+    CredentialItem,
+    CredentialResponse,
+    CredentialType,
+    CreditCardCredentialResponse,
+    PasswordCredentialResponse,
+)
+from skyvern.forge.sdk.services.bitwarden import BitwardenService
+from skyvern.forge.sdk.services.credential.credential_vault_service import CredentialVaultService
+
+LOG = structlog.get_logger()
+
+
+class BitwardenCredentialVaultService(CredentialVaultService):
+    async def create_credential(self, organization_id: str, data: CreateCredentialRequest) -> Credential:
+        org_collection = await app.DATABASE.get_organization_bitwarden_collection(organization_id)
+
+        if not org_collection:
+            LOG.info(
+                "There is no collection for the organization. Creating new collection.",
+                organization_id=organization_id,
+            )
+            collection_id = await BitwardenService.create_collection(
+                name=organization_id,
+            )
+            org_collection = await app.DATABASE.create_organization_bitwarden_collection(
+                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=organization_id,
+            item_id=item_id,
+            name=data.name,
+            credential_type=data.credential_type,
+            totp_type=data.credential.totp_type if hasattr(data.credential, "totp_type") else "none",
+        )
+
+        return credential
+
+    async def delete_credential(
+        self,
+        credential: Credential,
+    ) -> None:
+        organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection(
+            credential.organization_id
+        )
+        if not organization_bitwarden_collection:
+            raise HTTPException(status_code=404, detail="Credential account not found. It might have been deleted.")
+
+        await app.DATABASE.delete_credential(credential.credential_id, credential.organization_id)
+        await BitwardenService.delete_credential_item(credential.item_id)
+
+    async def get_credential(self, organization_id: str, credential_id: str) -> CredentialResponse:
+        organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection(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=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,
+                totp_type=credential.totp_type,
+            )
+            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")
+
+    async def get_credentials(self, organization_id: str, page: int, page_size: int) -> list[CredentialResponse]:
+        organization_bitwarden_collection = await app.DATABASE.get_organization_bitwarden_collection(organization_id)
+        if not organization_bitwarden_collection:
+            return []
+
+        credentials = await app.DATABASE.get_credentials(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:
+                LOG.warning(
+                    "Credential item not found in vault",
+                    credential_id=credential.credential_id,
+                    item_id=credential.item_id,
+                )
+                continue
+            if item.credential_type == CredentialType.PASSWORD:
+                credential_response = PasswordCredentialResponse(
+                    username=item.credential.username,
+                    totp_type=credential.totp_type,
+                )
+                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
+
+    async def get_credential_item(self, db_credential: Credential) -> CredentialItem:
+        return await BitwardenService.get_credential_item(db_credential.item_id)

skyvern/forge/sdk/services/credential/credential_vault_service.py

@@ -0,0 +1,36 @@
+from abc import ABC, abstractmethod
+
+from skyvern.forge.sdk.schemas.credentials import (
+    CreateCredentialRequest,
+    Credential,
+    CredentialItem,
+    CredentialResponse,
+)
+
+
+class CredentialVaultService(ABC):
+    """Abstract interface for credential vault services.
+
+    This interface defines the contract for storing and retrieving credentials
+    from different vault providers (e.g., Bitwarden, OnePassword, AWS Secrets Manager).
+    """
+
+    @abstractmethod
+    async def create_credential(self, organization_id: str, data: CreateCredentialRequest) -> Credential:
+        """Create a new credential in the vault and database."""
+
+    @abstractmethod
+    async def delete_credential(self, credential: Credential) -> None:
+        """Delete a credential from the vault and database."""
+
+    @abstractmethod
+    async def get_credential(self, organization_id: str, credential_id: str) -> CredentialResponse:
+        """Retrieve a credential with masked sensitive data."""
+
+    @abstractmethod
+    async def get_credentials(self, organization_id: str, page: int, page_size: int) -> list[CredentialResponse]:
+        """Retrieve all credentials for an organization with pagination."""
+
+    @abstractmethod
+    async def get_credential_item(self, db_credential: Credential) -> CredentialItem:
+        """Retrieve the full credential data from the vault."""

skyvern/forge/sdk/workflow/context_manager.py

@@ -295,15 +295,14 @@ class WorkflowRunContext:
             if db_credential is None:
                 raise CredentialParameterNotFoundError(credential_id)
 
-            bitwarden_credential = await BitwardenService.get_credential_item(db_credential.item_id)
-
-            credential_item = bitwarden_credential.credential
+            credential_item = await app.CREDENTIAL_VAULT_SERVICE.get_credential_item(db_credential)
+            credential = credential_item.credential
 
             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)",
             }
-            credential_dict = credential_item.model_dump()
+            credential_dict = credential.model_dump()
             for key, value in credential_dict.items():
                 if value is None:
                     continue
@@ -312,12 +311,12 @@ class WorkflowRunContext:
                 self.secrets[secret_id] = value
                 self.values[parameter.key][key] = secret_id
 
-            if isinstance(credential_item, PasswordCredential) and credential_item.totp is not None:
+            if isinstance(credential, PasswordCredential) and credential.totp is not None:
                 random_secret_id = self.generate_random_secret_id()
                 totp_secret_id = f"{random_secret_id}_totp"
                 self.secrets[totp_secret_id] = BitwardenConstants.TOTP
                 totp_secret_value = self.totp_secret_value_key(totp_secret_id)
-                self.secrets[totp_secret_value] = parse_totp_secret(credential_item.totp)
+                self.secrets[totp_secret_value] = parse_totp_secret(credential.totp)
                 self.values[parameter.key]["totp"] = totp_secret_id
         except Exception as e:
             LOG.error(f"Failed to get credential from database: {credential_id}. Error: {e}")
@@ -345,27 +344,26 @@ class WorkflowRunContext:
         if db_credential is None:
             raise CredentialParameterNotFoundError(credential_id)
 
-        bitwarden_credential = await BitwardenService.get_credential_item(db_credential.item_id)
-
-        credential_item = bitwarden_credential.credential
+        credential_item = await app.CREDENTIAL_VAULT_SERVICE.get_credential_item(db_credential)
+        credential = credential_item.credential
 
         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)",
         }
-        credential_dict = credential_item.model_dump()
+        credential_dict = credential.model_dump()
         for key, value in credential_dict.items():
             random_secret_id = self.generate_random_secret_id()
             secret_id = f"{random_secret_id}_{key}"
             self.secrets[secret_id] = value
             self.values[parameter.key][key] = secret_id
 
-        if isinstance(credential_item, PasswordCredential) and credential_item.totp is not None:
+        if isinstance(credential, PasswordCredential) and credential.totp is not None:
             random_secret_id = self.generate_random_secret_id()
             totp_secret_id = f"{random_secret_id}_totp"
             self.secrets[totp_secret_id] = BitwardenConstants.TOTP
             totp_secret_value = self.totp_secret_value_key(totp_secret_id)
-            self.secrets[totp_secret_value] = parse_totp_secret(credential_item.totp)
+            self.secrets[totp_secret_value] = parse_totp_secret(credential.totp)
             self.values[parameter.key]["totp"] = totp_secret_id
 
     async def register_aws_secret_parameter_value(