From d6bb3de2c56a8693029d645d15b42f84836ed0fe Mon Sep 17 00:00:00 2001 From: Stanislav Novosad Date: Tue, 14 Oct 2025 20:25:22 -0600 Subject: [PATCH] Delete Azure Vault secrets async (#3716) --- skyvern/forge/api_app.py | 13 +----- skyvern/forge/sdk/api/azure.py | 2 +- skyvern/forge/sdk/routes/credentials.py | 4 ++ .../azure_credential_vault_service.py | 43 +++++++++++++++---- .../credential/credential_vault_service.py | 6 +++ 5 files changed, 48 insertions(+), 20 deletions(-) diff --git a/skyvern/forge/api_app.py b/skyvern/forge/api_app.py index 68593be6..2354300f 100644 --- a/skyvern/forge/api_app.py +++ b/skyvern/forge/api_app.py @@ -1,7 +1,6 @@ import uuid -from contextlib import asynccontextmanager from datetime import datetime -from typing import Any, AsyncGenerator, Awaitable, Callable +from typing import Awaitable, Callable import structlog from fastapi import FastAPI, Response, status @@ -51,20 +50,12 @@ def custom_openapi() -> dict: return app.openapi_schema -@asynccontextmanager -async def lifespan(_: FastAPI) -> AsyncGenerator[None, Any]: - """Lifespan context manager for FastAPI app startup and shutdown.""" - LOG.info("Server started") - yield - LOG.info("Server shutting down") - - def get_agent_app() -> FastAPI: """ Start the agent server. """ - app = FastAPI(lifespan=lifespan) + app = FastAPI() # Add CORS middleware app.add_middleware( diff --git a/skyvern/forge/sdk/api/azure.py b/skyvern/forge/sdk/api/azure.py index b0121f59..f639dd2b 100644 --- a/skyvern/forge/sdk/api/azure.py +++ b/skyvern/forge/sdk/api/azure.py @@ -33,7 +33,7 @@ class AsyncAzureVaultClient: finally: await secret_client.close() - async def create_secret(self, secret_name: str, secret_value: str, vault_name: str) -> str: + async def create_or_update_secret(self, secret_name: str, secret_value: str, vault_name: str) -> str: secret_client = await self._get_secret_client(vault_name) try: secret = await secret_client.set_secret(secret_name, secret_value) diff --git a/skyvern/forge/sdk/routes/credentials.py b/skyvern/forge/sdk/routes/credentials.py index c4ea1d6a..5550ab2c 100644 --- a/skyvern/forge/sdk/routes/credentials.py +++ b/skyvern/forge/sdk/routes/credentials.py @@ -205,6 +205,7 @@ async def create_credential( include_in_schema=False, ) async def delete_credential( + background_tasks: BackgroundTasks, credential_id: str = Path( ..., description="The unique identifier of the credential to delete", @@ -226,6 +227,9 @@ async def delete_credential( await credential_service.delete_credential(credential) + # Schedule background cleanup if the service implements it + background_tasks.add_task(credential_service.post_delete_credential_item, credential.item_id) + return None diff --git a/skyvern/forge/sdk/services/credential/azure_credential_vault_service.py b/skyvern/forge/sdk/services/credential/azure_credential_vault_service.py index 866ec234..949d1457 100644 --- a/skyvern/forge/sdk/services/credential/azure_credential_vault_service.py +++ b/skyvern/forge/sdk/services/credential/azure_credential_vault_service.py @@ -1,6 +1,7 @@ import uuid from typing import Annotated, Literal, Union +import structlog from azure.identity.aio import ClientSecretCredential from fastapi import HTTPException from pydantic import BaseModel, Field, TypeAdapter @@ -21,6 +22,8 @@ from skyvern.forge.sdk.schemas.credentials import ( ) from skyvern.forge.sdk.services.credential.credential_vault_service import CredentialVaultService +LOG = structlog.get_logger() + class AzureCredentialVaultService(CredentialVaultService): class _PasswordCredentialDataImage(BaseModel): @@ -72,7 +75,37 @@ class AzureCredentialVaultService(CredentialVaultService): credential: Credential, ) -> None: await app.DATABASE.delete_credential(credential.credential_id, credential.organization_id) - await self.delete_credential_item(credential.item_id) + # Deleting takes several seconds, so we empty the value and delete async so customers do not have to wait + await self._client.create_or_update_secret( + vault_name=self._vault_name, + secret_name=credential.item_id, + secret_value="", + ) + + async def post_delete_credential_item(self, item_id: str) -> None: + """ + Background task to delete the credential item from Azure Key Vault. + This allows the API to respond quickly while the deletion happens asynchronously. + """ + try: + LOG.info( + "Deleting credential item from Azure Key Vault in background", + item_id=item_id, + vault_name=self._vault_name, + ) + await self._client.delete_secret(secret_name=item_id, vault_name=self._vault_name) + LOG.info( + "Successfully deleted credential item from Azure Key Vault", + item_id=item_id, + vault_name=self._vault_name, + ) + except Exception as e: + LOG.exception( + "Failed to delete credential item from Azure Key Vault in background", + item_id=item_id, + vault_name=self._vault_name, + error=str(e), + ) async def get_credential(self, organization_id: str, credential_id: str) -> CredentialResponse: credential = await app.DATABASE.get_credential(credential_id=credential_id, organization_id=organization_id) @@ -85,12 +118,6 @@ class AzureCredentialVaultService(CredentialVaultService): credentials = await app.DATABASE.get_credentials(organization_id, page=page, page_size=page_size) return [_convert_to_response(credential) for credential in credentials] - async def delete_credential_item(self, item_id: str) -> None: - await self._client.delete_secret( - vault_name=self._vault_name, - secret_name=item_id, - ) - async def get_credential_item(self, db_credential: Credential) -> CredentialItem: secret_json_str = await self._client.get_secret(secret_name=db_credential.item_id, vault_name=self._vault_name) if secret_json_str is None: @@ -154,7 +181,7 @@ class AzureCredentialVaultService(CredentialVaultService): secret_name = f"{organization_id}-{uuid.uuid4()}".replace("_", "") secret_value = data.model_dump_json(exclude_none=True) - return await self._client.create_secret( + return await self._client.create_or_update_secret( vault_name=self._vault_name, secret_name=secret_name, secret_value=secret_value, diff --git a/skyvern/forge/sdk/services/credential/credential_vault_service.py b/skyvern/forge/sdk/services/credential/credential_vault_service.py index f39e2b28..93c16b19 100644 --- a/skyvern/forge/sdk/services/credential/credential_vault_service.py +++ b/skyvern/forge/sdk/services/credential/credential_vault_service.py @@ -26,6 +26,12 @@ class CredentialVaultService(ABC): async def delete_credential(self, credential: Credential) -> None: """Delete a credential from the vault and database.""" + async def post_delete_credential_item(self, item_id: str) -> None: + """ + Optional hook for scheduling background cleanup tasks after credential deletion. + Default implementation does nothing. Override in subclasses as needed. + """ + @abstractmethod async def get_credential(self, organization_id: str, credential_id: str) -> CredentialResponse: """Retrieve a credential with masked sensitive data."""