From 5b530cab52a4ae070b2e9c844dc4c91808a0ca1b Mon Sep 17 00:00:00 2001 From: Stanislav Novosad Date: Wed, 26 Nov 2025 19:26:48 -0700 Subject: [PATCH] Initial scenario tests infra (#4118) --- skyvern/forge/forge_app.py | 16 +- skyvern/forge/sdk/api/azure.py | 187 +++++++++--------- skyvern/forge/sdk/api/real_azure.py | 127 ++++++++++++ .../azure_credential_vault_service.py | 11 +- skyvern/forge/sdk/workflow/context_manager.py | 7 +- skyvern/forge/sdk/workflow/models/block.py | 3 +- 6 files changed, 245 insertions(+), 106 deletions(-) create mode 100644 skyvern/forge/sdk/api/real_azure.py diff --git a/skyvern/forge/forge_app.py b/skyvern/forge/forge_app.py index 6a4b3a27..8384e1d5 100644 --- a/skyvern/forge/forge_app.py +++ b/skyvern/forge/forge_app.py @@ -9,8 +9,10 @@ from openai import AsyncAzureOpenAI, AsyncOpenAI from skyvern.config import Settings from skyvern.forge.agent import ForgeAgent from skyvern.forge.agent_functions import AgentFunction +from skyvern.forge.sdk.api.azure import AzureClientFactory from skyvern.forge.sdk.api.llm.api_handler_factory import LLMAPIHandlerFactory from skyvern.forge.sdk.api.llm.models import LLMAPIHandler +from skyvern.forge.sdk.api.real_azure import RealAzureClientFactory from skyvern.forge.sdk.artifact.manager import ArtifactManager from skyvern.forge.sdk.artifact.storage.base import BaseStorage from skyvern.forge.sdk.artifact.storage.factory import StorageFactory @@ -20,7 +22,7 @@ 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.credentials import CredentialVaultType -from skyvern.forge.sdk.schemas.organizations import Organization +from skyvern.forge.sdk.schemas.organizations import AzureClientSecretCredential, Organization from skyvern.forge.sdk.services.credential.azure_credential_vault_service import AzureCredentialVaultService from skyvern.forge.sdk.services.credential.bitwarden_credential_service import BitwardenCredentialVaultService from skyvern.forge.sdk.services.credential.credential_vault_service import CredentialVaultService @@ -46,6 +48,7 @@ class ForgeApp: OPENAI_CLIENT: AsyncOpenAI | AsyncAzureOpenAI ANTHROPIC_CLIENT: AsyncAnthropic | AsyncAnthropicBedrock UI_TARS_CLIENT: AsyncOpenAI | None + AZURE_CLIENT_FACTORY: AzureClientFactory SECONDARY_LLM_API_HANDLER: LLMAPIHandler SELECT_AGENT_LLM_API_HANDLER: LLMAPIHandler NORMAL_SELECT_AGENT_LLM_API_HANDLER: LLMAPIHandler @@ -166,12 +169,17 @@ def create_forge_app() -> ForgeApp: app.AGENT_FUNCTION = AgentFunction() app.PERSISTENT_SESSIONS_MANAGER = PersistentSessionsManager(database=app.DATABASE) + app.AZURE_CLIENT_FACTORY = RealAzureClientFactory() app.BITWARDEN_CREDENTIAL_VAULT_SERVICE = BitwardenCredentialVaultService() app.AZURE_CREDENTIAL_VAULT_SERVICE = ( AzureCredentialVaultService( - tenant_id=settings.AZURE_TENANT_ID, # type: ignore[arg-type] - client_id=settings.AZURE_CLIENT_ID, # type: ignore[arg-type] - client_secret=settings.AZURE_CLIENT_SECRET, # type: ignore[arg-type] + app.AZURE_CLIENT_FACTORY.create_from_client_secret( + AzureClientSecretCredential( + tenant_id=settings.AZURE_TENANT_ID, # type: ignore + client_id=settings.AZURE_CLIENT_ID, # type: ignore + client_secret=settings.AZURE_CLIENT_SECRET, # type: ignore + ) + ), vault_name=settings.AZURE_CREDENTIAL_VAULT, # type: ignore[arg-type] ) if settings.AZURE_CREDENTIAL_VAULT diff --git a/skyvern/forge/sdk/api/azure.py b/skyvern/forge/sdk/api/azure.py index f639dd2b..af07cfb4 100644 --- a/skyvern/forge/sdk/api/azure.py +++ b/skyvern/forge/sdk/api/azure.py @@ -1,115 +1,126 @@ -from typing import Self - -import structlog -from azure.identity.aio import ClientSecretCredential, DefaultAzureCredential -from azure.keyvault.secrets.aio import SecretClient -from azure.storage.blob.aio import BlobServiceClient +from typing import Protocol, Self from skyvern.forge.sdk.schemas.organizations import AzureClientSecretCredential -LOG = structlog.get_logger() +class AsyncAzureVaultClient(Protocol): + """Protocol defining the interface for Azure Vault clients. -class AsyncAzureVaultClient: - def __init__(self, credential: ClientSecretCredential | DefaultAzureCredential) -> None: - self.credential = credential + This client provides methods to interact with Azure Key Vault for secret management. + """ async def __aenter__(self) -> Self: - return self + """Enter async context manager.""" + ... async def __aexit__( self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: object ) -> None: - await self.credential.close() + """Exit async context manager and cleanup resources.""" + ... async def get_secret(self, secret_name: str, vault_name: str) -> str | None: - secret_client = await self._get_secret_client(vault_name) - try: - secret = await secret_client.get_secret(secret_name) - return secret.value - except Exception as e: - LOG.exception("Failed to get secret from Azure Key Vault.", secret_name=secret_name, error=e) - return None - finally: - await secret_client.close() + """Retrieve a secret from Azure Key Vault. + + Args: + secret_name: The name of the secret to retrieve + vault_name: The name of the Azure Key Vault + + Returns: + The secret value as a string, or None if the secret doesn't exist or an error occurs + """ + ... 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) - return secret.name - except Exception as e: - LOG.exception("Failed to create secret from Azure Key Vault.", secret_name=secret_name, error=e) - raise e - finally: - await secret_client.close() + """Create or update a secret in Azure Key Vault. + + Args: + secret_name: The name of the secret to create or update + secret_value: The value to store + vault_name: The name of the Azure Key Vault + + Returns: + The name of the created/updated secret + + Raises: + Exception: If the operation fails + """ + ... async def delete_secret(self, secret_name: str, vault_name: str) -> str: - secret_client = await self._get_secret_client(vault_name) - try: - secret = await secret_client.delete_secret(secret_name) - return secret.name - except Exception as e: - LOG.exception("Failed to delete secret from Azure Key Vault.", secret_name=secret_name, error=e) - raise e - finally: - await secret_client.close() + """Delete a secret from Azure Key Vault. - async def _get_secret_client(self, vault_name: str) -> SecretClient: - # Azure Key Vault URL format: https://.vault.azure.net - # Assuming the secret_name is actually the Key Vault URL and the secret name - # This needs to be clarified or passed as separate parameters - # For now, let's assume secret_name is the actual secret name and Key Vault URL is in settings. - key_vault_url = f"https://{vault_name}.vault.azure.net" # Placeholder, adjust as needed - return SecretClient(vault_url=key_vault_url, credential=self.credential) + Args: + secret_name: The name of the secret to delete + vault_name: The name of the Azure Key Vault + + Returns: + The name of the deleted secret + + Raises: + Exception: If the operation fails + """ + ... async def close(self) -> None: - await self.credential.close() - - @classmethod - def create_default(cls) -> "AsyncAzureVaultClient": - return cls(DefaultAzureCredential()) - - @classmethod - def create_from_client_secret( - cls, - credential: AzureClientSecretCredential, - ) -> "AsyncAzureVaultClient": - cred = ClientSecretCredential( - tenant_id=credential.tenant_id, - client_id=credential.client_id, - client_secret=credential.client_secret, - ) - return cls(cred) + """Close the client and release all resources.""" + ... -class AsyncAzureStorageClient: - def __init__(self, storage_account_name: str, storage_account_key: str): - self.blob_service_client = BlobServiceClient( - account_url=f"https://{storage_account_name}.blob.core.windows.net", - credential=storage_account_key, - ) +class AsyncAzureStorageClient(Protocol): + """Protocol defining the interface for Azure Storage clients. + + This client provides methods to interact with Azure Blob Storage for file operations. + """ async def upload_file_from_path(self, container_name: str, blob_name: str, file_path: str) -> None: - try: - container_client = self.blob_service_client.get_container_client(container_name) - # Create the container if it doesn't exist - try: - await container_client.create_container() - except Exception as e: - LOG.info("Azure container already exists or failed to create", container_name=container_name, error=e) + """Upload a file from the local filesystem to Azure Blob Storage. - with open(file_path, "rb") as data: - await container_client.upload_blob(name=blob_name, data=data, overwrite=True) - LOG.info("File uploaded to Azure Blob Storage", container_name=container_name, blob_name=blob_name) - except Exception as e: - LOG.error( - "Failed to upload file to Azure Blob Storage", - container_name=container_name, - blob_name=blob_name, - error=e, - ) - raise e + Args: + container_name: The name of the Azure Blob container + blob_name: The name to give the blob in storage + file_path: The local path to the file to upload + + Raises: + Exception: If the upload fails + """ + ... async def close(self) -> None: - await self.blob_service_client.close() + """Close the storage client and release resources.""" + ... + + +class AzureClientFactory(Protocol): + """Protocol defining the interface for creating Azure Vault and Storage clients.""" + + def create_default(self) -> "AsyncAzureVaultClient": + """Create an Azure Vault client using default credentials. + + Returns: + An AsyncAzureVaultClient instance using DefaultAzureCredential + """ + ... + + def create_from_client_secret(self, credential: AzureClientSecretCredential) -> "AsyncAzureVaultClient": + """Create an Azure Vault client using client secret credentials. + + Args: + credential: Azure client secret credentials containing tenant_id, client_id, and client_secret + + Returns: + An AsyncAzureVaultClient instance + """ + ... + + def create_storage_client(self, storage_account_name: str, storage_account_key: str) -> "AsyncAzureStorageClient": + """Create an Azure Storage client with the provided credentials. + + Args: + storage_account_name: The name of the Azure storage account + storage_account_key: The access key for the storage account + + Returns: + An AsyncAzureStorageClient instance + """ + ... diff --git a/skyvern/forge/sdk/api/real_azure.py b/skyvern/forge/sdk/api/real_azure.py new file mode 100644 index 00000000..ad7b8aab --- /dev/null +++ b/skyvern/forge/sdk/api/real_azure.py @@ -0,0 +1,127 @@ +"""Real implementations of Azure clients (Vault and Storage) and their factories.""" + +from typing import Self + +import structlog +from azure.identity.aio import ClientSecretCredential, DefaultAzureCredential +from azure.keyvault.secrets.aio import SecretClient +from azure.storage.blob.aio import BlobServiceClient + +from skyvern.forge.sdk.api.azure import AsyncAzureStorageClient, AsyncAzureVaultClient, AzureClientFactory +from skyvern.forge.sdk.schemas.organizations import AzureClientSecretCredential + +LOG = structlog.get_logger() + + +class RealAsyncAzureVaultClient(AsyncAzureVaultClient): + """Real implementation of Azure Vault client using Azure SDK.""" + + def __init__(self, credential: ClientSecretCredential | DefaultAzureCredential) -> None: + self.credential = credential + + async def __aenter__(self) -> Self: + return self + + async def __aexit__( + self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: object + ) -> None: + await self.credential.close() + + async def get_secret(self, secret_name: str, vault_name: str) -> str | None: + secret_client = await self._get_secret_client(vault_name) + try: + secret = await secret_client.get_secret(secret_name) + return secret.value + except Exception as e: + LOG.exception("Failed to get secret from Azure Key Vault.", secret_name=secret_name, error=e) + return None + finally: + await secret_client.close() + + 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) + return secret.name + except Exception as e: + LOG.exception("Failed to create secret from Azure Key Vault.", secret_name=secret_name, error=e) + raise e + finally: + await secret_client.close() + + async def delete_secret(self, secret_name: str, vault_name: str) -> str: + secret_client = await self._get_secret_client(vault_name) + try: + secret = await secret_client.delete_secret(secret_name) + return secret.name + except Exception as e: + LOG.exception("Failed to delete secret from Azure Key Vault.", secret_name=secret_name, error=e) + raise e + finally: + await secret_client.close() + + async def _get_secret_client(self, vault_name: str) -> SecretClient: + # Azure Key Vault URL format: https://.vault.azure.net + # Assuming the secret_name is actually the Key Vault URL and the secret name + # This needs to be clarified or passed as separate parameters + # For now, let's assume secret_name is the actual secret name and Key Vault URL is in settings. + key_vault_url = f"https://{vault_name}.vault.azure.net" # Placeholder, adjust as needed + return SecretClient(vault_url=key_vault_url, credential=self.credential) + + async def close(self) -> None: + await self.credential.close() + + +class RealAsyncAzureStorageClient(AsyncAzureStorageClient): + """Real implementation of Azure Storage client using Azure SDK.""" + + def __init__(self, storage_account_name: str, storage_account_key: str): + self.blob_service_client = BlobServiceClient( + account_url=f"https://{storage_account_name}.blob.core.windows.net", + credential=storage_account_key, + ) + + async def upload_file_from_path(self, container_name: str, blob_name: str, file_path: str) -> None: + try: + container_client = self.blob_service_client.get_container_client(container_name) + # Create the container if it doesn't exist + try: + await container_client.create_container() + except Exception as e: + LOG.info("Azure container already exists or failed to create", container_name=container_name, error=e) + + with open(file_path, "rb") as data: + await container_client.upload_blob(name=blob_name, data=data, overwrite=True) + LOG.info("File uploaded to Azure Blob Storage", container_name=container_name, blob_name=blob_name) + except Exception as e: + LOG.error( + "Failed to upload file to Azure Blob Storage", + container_name=container_name, + blob_name=blob_name, + error=e, + ) + raise e + + async def close(self) -> None: + await self.blob_service_client.close() + + +class RealAzureClientFactory(AzureClientFactory): + """Factory for creating real Azure Vault and Storage clients.""" + + def create_default(self) -> AsyncAzureVaultClient: + """Create an Azure Vault client using DefaultAzureCredential.""" + return RealAsyncAzureVaultClient(DefaultAzureCredential()) + + def create_from_client_secret(self, credential: AzureClientSecretCredential) -> AsyncAzureVaultClient: + """Create an Azure Vault client using client secret credentials.""" + cred = ClientSecretCredential( + tenant_id=credential.tenant_id, + client_id=credential.client_id, + client_secret=credential.client_secret, + ) + return RealAsyncAzureVaultClient(cred) + + def create_storage_client(self, storage_account_name: str, storage_account_key: str) -> AsyncAzureStorageClient: + """Create an Azure Storage client with the provided credentials.""" + return RealAsyncAzureStorageClient(storage_account_name, storage_account_key) 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 629d8df3..2b46abf7 100644 --- a/skyvern/forge/sdk/services/credential/azure_credential_vault_service.py +++ b/skyvern/forge/sdk/services/credential/azure_credential_vault_service.py @@ -2,7 +2,6 @@ import uuid from typing import Annotated, Literal, Union import structlog -from azure.identity.aio import ClientSecretCredential from pydantic import BaseModel, Field, TypeAdapter from skyvern.forge import app @@ -41,14 +40,8 @@ class AzureCredentialVaultService(CredentialVaultService): Union[_PasswordCredentialDataImage, _CreditCardCredentialDataImage], Field(discriminator="type") ] - def __init__(self, tenant_id: str, client_id: str, client_secret: str, vault_name: str): - self._client = AsyncAzureVaultClient( - ClientSecretCredential( - tenant_id=tenant_id, - client_id=client_id, - client_secret=client_secret, - ) - ) + def __init__(self, client: AsyncAzureVaultClient, vault_name: str): + self._client = client self._vault_name = vault_name async def create_credential(self, organization_id: str, data: CreateCredentialRequest) -> Credential: diff --git a/skyvern/forge/sdk/workflow/context_manager.py b/skyvern/forge/sdk/workflow/context_manager.py index 66495204..413eee08 100644 --- a/skyvern/forge/sdk/workflow/context_manager.py +++ b/skyvern/forge/sdk/workflow/context_manager.py @@ -387,7 +387,8 @@ class WorkflowRunContext: # If the parameter is an Azure secret, fetch the secret value and store it in the secrets dict # The value of the parameter will be the random secret id with format `secret_`. # We'll replace the random secret id with the actual secret value when we need to use it. - async with AsyncAzureVaultClient.create_default() as azure_vault_client: + azure_vault_client = app.AZURE_CLIENT_FACTORY.create_default() + async with azure_vault_client: secret_value = await azure_vault_client.get_secret(parameter.azure_key, vault_name) if secret_value is not None: random_secret_id = self.generate_random_secret_id() @@ -989,10 +990,10 @@ class WorkflowRunContext: organization.organization_id, OrganizationAuthTokenType.azure_client_secret_credential.value ) if org_auth_token: - azure_vault_client = AsyncAzureVaultClient.create_from_client_secret(org_auth_token.credential) + azure_vault_client = app.AZURE_CLIENT_FACTORY.create_from_client_secret(org_auth_token.credential) else: # Use the DefaultAzureCredential if not configured on organization level - azure_vault_client = AsyncAzureVaultClient.create_default() + azure_vault_client = app.AZURE_CLIENT_FACTORY.create_default() return azure_vault_client def _add_secret_parameter_value(self, parameter: Parameter, key: str, value: str) -> None: diff --git a/skyvern/forge/sdk/workflow/models/block.py b/skyvern/forge/sdk/workflow/models/block.py index 620d5660..d290e4c8 100644 --- a/skyvern/forge/sdk/workflow/models/block.py +++ b/skyvern/forge/sdk/workflow/models/block.py @@ -49,7 +49,6 @@ from skyvern.forge import app from skyvern.forge.prompts import prompt_engine from skyvern.forge.sdk.api import email from skyvern.forge.sdk.api.aws import AsyncAWSClient -from skyvern.forge.sdk.api.azure import AsyncAzureStorageClient from skyvern.forge.sdk.api.files import ( calculate_sha256_for_file, create_named_temporary_file, @@ -2322,7 +2321,7 @@ class FileUploadBlock(Block): if actual_azure_storage_account_name is None or actual_azure_storage_account_key is None: raise AzureConfigurationError("Azure Storage is not configured") - azure_client = AsyncAzureStorageClient( + azure_client = app.AZURE_CLIENT_FACTORY.create_storage_client( storage_account_name=actual_azure_storage_account_name, storage_account_key=actual_azure_storage_account_key, )