Initial scenario tests infra (#4118)

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]
+            app.AZURE_CLIENT_FACTORY.create_from_client_secret(
-            client_id=settings.AZURE_CLIENT_ID,  # type: ignore[arg-type]
+                AzureClientSecretCredential(
-            client_secret=settings.AZURE_CLIENT_SECRET,  # type: ignore[arg-type]
+                    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

skyvern/forge/sdk/api/azure.py

@@ -1,115 +1,126 @@
-from typing import Self
+from typing import Protocol, 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.schemas.organizations import AzureClientSecretCredential
 
-LOG = structlog.get_logger()
 
+class AsyncAzureVaultClient(Protocol):
+    """Protocol defining the interface for Azure Vault clients.
 
-class AsyncAzureVaultClient:
+    This client provides methods to interact with Azure Key Vault for secret management.
-    def __init__(self, credential: ClientSecretCredential | DefaultAzureCredential) -> None:
+    """
-        self.credential = credential
 
     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)
+        """Retrieve a secret from Azure Key Vault.
-        try:
+
-            secret = await secret_client.get_secret(secret_name)
+        Args:
-            return secret.value
+            secret_name: The name of the secret to retrieve
-        except Exception as e:
+            vault_name: The name of the Azure Key Vault
-            LOG.exception("Failed to get secret from Azure Key Vault.", secret_name=secret_name, error=e)
+
-            return None
+        Returns:
-        finally:
+            The secret value as a string, or None if the secret doesn't exist or an error occurs
-            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)
+        """Create or update a secret in Azure Key Vault.
-        try:
+
-            secret = await secret_client.set_secret(secret_name, secret_value)
+        Args:
-            return secret.name
+            secret_name: The name of the secret to create or update
-        except Exception as e:
+            secret_value: The value to store
-            LOG.exception("Failed to create secret from Azure Key Vault.", secret_name=secret_name, error=e)
+            vault_name: The name of the Azure Key Vault
-            raise e
+
-        finally:
+        Returns:
-            await secret_client.close()
+            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)
+        """Delete a secret from Azure Key Vault.
-        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:
+        Args:
-        # Azure Key Vault URL format: https://<your-key-vault-name>.vault.azure.net
+            secret_name: The name of the secret to delete
-        # Assuming the secret_name is actually the Key Vault URL and the secret name
+            vault_name: The name of the Azure Key Vault
-        # 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.
+        Returns:
-        key_vault_url = f"https://{vault_name}.vault.azure.net"  # Placeholder, adjust as needed
+            The name of the deleted secret
-        return SecretClient(vault_url=key_vault_url, credential=self.credential)
+
+        Raises:
+            Exception: If the operation fails
+        """
+        ...
 
     async def close(self) -> None:
-        await self.credential.close()
+        """Close the client and release all resources."""
-
+        ...
-    @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)
 
 
-class AsyncAzureStorageClient:
+class AsyncAzureStorageClient(Protocol):
-    def __init__(self, storage_account_name: str, storage_account_key: str):
+    """Protocol defining the interface for Azure Storage clients.
-        self.blob_service_client = BlobServiceClient(
+
-            account_url=f"https://{storage_account_name}.blob.core.windows.net",
+    This client provides methods to interact with Azure Blob Storage for file operations.
-            credential=storage_account_key,
+    """
-        )
 
     async def upload_file_from_path(self, container_name: str, blob_name: str, file_path: str) -> None:
-        try:
+        """Upload a file from the local filesystem to Azure Blob Storage.
-            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:
+        Args:
-                await container_client.upload_blob(name=blob_name, data=data, overwrite=True)
+            container_name: The name of the Azure Blob container
-            LOG.info("File uploaded to Azure Blob Storage", container_name=container_name, blob_name=blob_name)
+            blob_name: The name to give the blob in storage
-        except Exception as e:
+            file_path: The local path to the file to upload
-            LOG.error(
+
-                "Failed to upload file to Azure Blob Storage",
+        Raises:
-                container_name=container_name,
+            Exception: If the upload fails
-                blob_name=blob_name,
+        """
-                error=e,
+        ...
-            )
-            raise e
 
     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
+        """
+        ...

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://<your-key-vault-name>.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)

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):
+    def __init__(self, client: AsyncAzureVaultClient, vault_name: str):
-        self._client = AsyncAzureVaultClient(
+        self._client = client
-            ClientSecretCredential(
-                tenant_id=tenant_id,
-                client_id=client_id,
-                client_secret=client_secret,
-            )
-        )
         self._vault_name = vault_name
 
     async def create_credential(self, organization_id: str, data: CreateCredentialRequest) -> Credential:

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_<uuid>`.
         # 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:

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,
                 )