From 0ba8780f9b463faff93d3be32375c61e7ac29d69 Mon Sep 17 00:00:00 2001 From: Ragini-Microsoft Date: Tue, 29 Jul 2025 15:15:59 +0530 Subject: [PATCH] replacing DefaultAzureCredential with ManagedIdentityCredential --- infra/main.bicep | 5 ++ src/backend/.env.sample | 1 + src/backend/app_config.py | 25 +----- src/backend/config_kernel.py | 3 +- src/backend/context/cosmos_memory_kernel.py | 4 +- src/backend/helpers/azure_credential_utils.py | 41 ++++++++++ .../helpers/test_azure_credential_utils.py | 78 +++++++++++++++++++ src/backend/tests/test_config.py | 8 -- src/backend/utils_kernel.py | 8 +- src/frontend/.env.sample | 1 + 10 files changed, 139 insertions(+), 35 deletions(-) create mode 100644 src/backend/helpers/azure_credential_utils.py create mode 100644 src/backend/tests/helpers/test_azure_credential_utils.py diff --git a/infra/main.bicep b/infra/main.bicep index 3dfd2db02..8ee54772d 100644 --- a/infra/main.bicep +++ b/infra/main.bicep @@ -1034,6 +1034,10 @@ module containerApp 'br/public:avm/res/app/container-app:0.14.2' = if (container name: 'AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME' value: aiFoundryAiServicesModelDeployment.name } + { + name: 'APP_ENV' + value: 'Prod' + } ] } ] @@ -1087,6 +1091,7 @@ module webSite 'br/public:avm/res/web/site:0.15.1' = if (webSiteEnabled) { WEBSITES_CONTAINER_START_TIME_LIMIT: '1800' // 30 minutes, adjust as needed BACKEND_API_URL: 'https://${containerApp.outputs.fqdn}' AUTH_ENABLED: 'false' + APP_ENV: 'Prod' } } } diff --git a/src/backend/.env.sample b/src/backend/.env.sample index 2a651df39..ab1c41369 100644 --- a/src/backend/.env.sample +++ b/src/backend/.env.sample @@ -16,6 +16,7 @@ AZURE_AI_MODEL_DEPLOYMENT_NAME=gpt-4o APPLICATIONINSIGHTS_CONNECTION_STRING= AZURE_AI_AGENT_MODEL_DEPLOYMENT_NAME=gpt-4o AZURE_AI_AGENT_ENDPOINT= +APP_ENV="dev" BACKEND_API_URL=http://localhost:8000 FRONTEND_SITE_NAME=http://127.0.0.1:3000 \ No newline at end of file diff --git a/src/backend/app_config.py b/src/backend/app_config.py index d4b1a9e9a..17771fed2 100644 --- a/src/backend/app_config.py +++ b/src/backend/app_config.py @@ -5,7 +5,7 @@ from azure.ai.projects.aio import AIProjectClient from azure.cosmos.aio import CosmosClient -from azure.identity import DefaultAzureCredential +from helpers.azure_credential_utils import get_azure_credential from dotenv import load_dotenv from semantic_kernel.kernel import Kernel @@ -106,23 +106,6 @@ def _get_bool(self, name: str) -> bool: """ return name in os.environ and os.environ[name].lower() in ["true", "1"] - def get_azure_credentials(self): - """Get Azure credentials using DefaultAzureCredential. - - Returns: - DefaultAzureCredential instance for Azure authentication - """ - # Cache the credentials object - if self._azure_credentials is not None: - return self._azure_credentials - - try: - self._azure_credentials = DefaultAzureCredential() - return self._azure_credentials - except Exception as exc: - logging.warning("Failed to create DefaultAzureCredential: %s", exc) - return None - def get_cosmos_database_client(self): """Get a Cosmos DB client for the configured database. @@ -132,7 +115,7 @@ def get_cosmos_database_client(self): try: if self._cosmos_client is None: self._cosmos_client = CosmosClient( - self.COSMOSDB_ENDPOINT, credential=self.get_azure_credentials() + self.COSMOSDB_ENDPOINT, credential=get_azure_credential() ) if self._cosmos_database is None: @@ -169,10 +152,10 @@ def get_ai_project_client(self): return self._ai_project_client try: - credential = self.get_azure_credentials() + credential = get_azure_credential() if credential is None: raise RuntimeError( - "Unable to acquire Azure credentials; ensure DefaultAzureCredential is configured" + "Unable to acquire Azure credentials; ensure Managed Identity is configured" ) endpoint = self.AZURE_AI_AGENT_ENDPOINT diff --git a/src/backend/config_kernel.py b/src/backend/config_kernel.py index 80d0738af..598a88dc5 100644 --- a/src/backend/config_kernel.py +++ b/src/backend/config_kernel.py @@ -1,5 +1,6 @@ # Import AppConfig from app_config from app_config import config +from helpers.azure_credential_utils import get_azure_credential # This file is left as a lightweight wrapper around AppConfig for backward compatibility @@ -31,7 +32,7 @@ class Config: @staticmethod def GetAzureCredentials(): """Get Azure credentials using the AppConfig implementation.""" - return config.get_azure_credentials() + return get_azure_credential() @staticmethod def GetCosmosDatabaseClient(): diff --git a/src/backend/context/cosmos_memory_kernel.py b/src/backend/context/cosmos_memory_kernel.py index 02e732706..d547979da 100644 --- a/src/backend/context/cosmos_memory_kernel.py +++ b/src/backend/context/cosmos_memory_kernel.py @@ -10,7 +10,7 @@ from azure.cosmos.partition_key import PartitionKey from azure.cosmos.aio import CosmosClient -from azure.identity import DefaultAzureCredential +from helpers.azure_credential_utils import get_azure_credential from semantic_kernel.memory.memory_record import MemoryRecord from semantic_kernel.memory.memory_store_base import MemoryStoreBase from semantic_kernel.contents import ChatMessageContent, ChatHistory, AuthorRole @@ -73,7 +73,7 @@ async def initialize(self): if not self._database: # Create Cosmos client cosmos_client = CosmosClient( - self._cosmos_endpoint, credential=DefaultAzureCredential() + self._cosmos_endpoint, credential=get_azure_credential() ) self._database = cosmos_client.get_database_client( self._cosmos_database diff --git a/src/backend/helpers/azure_credential_utils.py b/src/backend/helpers/azure_credential_utils.py new file mode 100644 index 000000000..646efb444 --- /dev/null +++ b/src/backend/helpers/azure_credential_utils.py @@ -0,0 +1,41 @@ +import os +from azure.identity import ManagedIdentityCredential, DefaultAzureCredential +from azure.identity.aio import ManagedIdentityCredential as AioManagedIdentityCredential, DefaultAzureCredential as AioDefaultAzureCredential + + +async def get_azure_credential_async(client_id=None): + """ + Returns an Azure credential asynchronously based on the application environment. + + If the environment is 'dev', it uses AioDefaultAzureCredential. + Otherwise, it uses AioManagedIdentityCredential. + + Args: + client_id (str, optional): The client ID for the Managed Identity Credential. + + Returns: + Credential object: Either AioDefaultAzureCredential or AioManagedIdentityCredential. + """ + if os.getenv("APP_ENV", "prod").lower() == 'dev': + return AioDefaultAzureCredential() # CodeQL [SM05139] Okay use of DefaultAzureCredential as it is only used in development + else: + return AioManagedIdentityCredential(client_id=client_id) + + +def get_azure_credential(client_id=None): + """ + Returns an Azure credential based on the application environment. + + If the environment is 'dev', it uses DefaultAzureCredential. + Otherwise, it uses ManagedIdentityCredential. + + Args: + client_id (str, optional): The client ID for the Managed Identity Credential. + + Returns: + Credential object: Either DefaultAzureCredential or ManagedIdentityCredential. + """ + if os.getenv("APP_ENV", "prod").lower() == 'dev': + return DefaultAzureCredential() # CodeQL [SM05139] Okay use of DefaultAzureCredential as it is only used in development + else: + return ManagedIdentityCredential(client_id=client_id) diff --git a/src/backend/tests/helpers/test_azure_credential_utils.py b/src/backend/tests/helpers/test_azure_credential_utils.py new file mode 100644 index 000000000..fd98527f5 --- /dev/null +++ b/src/backend/tests/helpers/test_azure_credential_utils.py @@ -0,0 +1,78 @@ +import pytest +import sys +import os +from unittest.mock import patch, MagicMock + +# Ensure src/backend is on the Python path for imports +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))) + +import helpers.azure_credential_utils as azure_credential_utils + +# Synchronous tests + +@patch("helpers.azure_credential_utils.os.getenv") +@patch("helpers.azure_credential_utils.DefaultAzureCredential") +@patch("helpers.azure_credential_utils.ManagedIdentityCredential") +def test_get_azure_credential_dev_env(mock_managed_identity_credential, mock_default_azure_credential, mock_getenv): + """Test get_azure_credential in dev environment.""" + mock_getenv.return_value = "dev" + mock_default_credential = MagicMock() + mock_default_azure_credential.return_value = mock_default_credential + + credential = azure_credential_utils.get_azure_credential() + + mock_getenv.assert_called_once_with("APP_ENV", "prod") + mock_default_azure_credential.assert_called_once() + mock_managed_identity_credential.assert_not_called() + assert credential == mock_default_credential + +@patch("helpers.azure_credential_utils.os.getenv") +@patch("helpers.azure_credential_utils.DefaultAzureCredential") +@patch("helpers.azure_credential_utils.ManagedIdentityCredential") +def test_get_azure_credential_non_dev_env(mock_managed_identity_credential, mock_default_azure_credential, mock_getenv): + """Test get_azure_credential in non-dev environment.""" + mock_getenv.return_value = "prod" + mock_managed_credential = MagicMock() + mock_managed_identity_credential.return_value = mock_managed_credential + credential = azure_credential_utils.get_azure_credential(client_id="test-client-id") + + mock_getenv.assert_called_once_with("APP_ENV", "prod") + mock_managed_identity_credential.assert_called_once_with(client_id="test-client-id") + mock_default_azure_credential.assert_not_called() + assert credential == mock_managed_credential + +# Asynchronous tests + +@pytest.mark.asyncio +@patch("helpers.azure_credential_utils.os.getenv") +@patch("helpers.azure_credential_utils.AioDefaultAzureCredential") +@patch("helpers.azure_credential_utils.AioManagedIdentityCredential") +async def test_get_azure_credential_async_dev_env(mock_aio_managed_identity_credential, mock_aio_default_azure_credential, mock_getenv): + """Test get_azure_credential_async in dev environment.""" + mock_getenv.return_value = "dev" + mock_aio_default_credential = MagicMock() + mock_aio_default_azure_credential.return_value = mock_aio_default_credential + + credential = await azure_credential_utils.get_azure_credential_async() + + mock_getenv.assert_called_once_with("APP_ENV", "prod") + mock_aio_default_azure_credential.assert_called_once() + mock_aio_managed_identity_credential.assert_not_called() + assert credential == mock_aio_default_credential + +@pytest.mark.asyncio +@patch("helpers.azure_credential_utils.os.getenv") +@patch("helpers.azure_credential_utils.AioDefaultAzureCredential") +@patch("helpers.azure_credential_utils.AioManagedIdentityCredential") +async def test_get_azure_credential_async_non_dev_env(mock_aio_managed_identity_credential, mock_aio_default_azure_credential, mock_getenv): + """Test get_azure_credential_async in non-dev environment.""" + mock_getenv.return_value = "prod" + mock_aio_managed_credential = MagicMock() + mock_aio_managed_identity_credential.return_value = mock_aio_managed_credential + + credential = await azure_credential_utils.get_azure_credential_async(client_id="test-client-id") + + mock_getenv.assert_called_once_with("APP_ENV", "prod") + mock_aio_managed_identity_credential.assert_called_once_with(client_id="test-client-id") + mock_aio_default_azure_credential.assert_not_called() + assert credential == mock_aio_managed_credential \ No newline at end of file diff --git a/src/backend/tests/test_config.py b/src/backend/tests/test_config.py index 3c4b0efe2..07ff0d0b4 100644 --- a/src/backend/tests/test_config.py +++ b/src/backend/tests/test_config.py @@ -52,11 +52,3 @@ def test_get_bool_config(): assert GetBoolConfig("FEATURE_ENABLED") is True with patch.dict("os.environ", {"FEATURE_ENABLED": "0"}): assert GetBoolConfig("FEATURE_ENABLED") is False - - -@patch("config.DefaultAzureCredential") -def test_get_azure_credentials_with_env_vars(mock_default_cred): - """Test Config.GetAzureCredentials with explicit credentials.""" - with patch.dict(os.environ, MOCK_ENV_VARS): - creds = Config.GetAzureCredentials() - assert creds is not None diff --git a/src/backend/utils_kernel.py b/src/backend/utils_kernel.py index a95dc52e3..b57e75ad8 100644 --- a/src/backend/utils_kernel.py +++ b/src/backend/utils_kernel.py @@ -11,9 +11,11 @@ # Import AppConfig from app_config from app_config import config -from azure.identity import DefaultAzureCredential from context.cosmos_memory_kernel import CosmosMemoryContext +# Import the credential utility +from helpers.azure_credential_utils import get_azure_credential + # Import agent factory and the new AppConfig from kernel_agents.agent_factory import AgentFactory from kernel_agents.group_chat_manager import GroupChatManager @@ -169,8 +171,8 @@ async def rai_success(description: str) -> bool: True if it passes, False otherwise """ try: - # Use DefaultAzureCredential for authentication to Azure OpenAI - credential = DefaultAzureCredential() + # Use managed identity for authentication to Azure OpenAI + credential = get_azure_credential() access_token = credential.get_token( "https://cognitiveservices.azure.com/.default" ).token diff --git a/src/frontend/.env.sample b/src/frontend/.env.sample index 3f56e3400..0817d28e2 100644 --- a/src/frontend/.env.sample +++ b/src/frontend/.env.sample @@ -2,6 +2,7 @@ API_URL=http://localhost:8000 ENABLE_AUTH=false +APP_ENV="dev" # VITE_APP_MSAL_AUTH_CLIENTID="" # VITE_APP_MSAL_AUTH_AUTHORITY="" # VITE_APP_MSAL_REDIRECT_URL="/"