diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_constants.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_constants.py index 21f1286390cd..5a766389be95 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_constants.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_constants.py @@ -19,6 +19,7 @@ class Constants: # Tracing APPLICATIONINSIGHTS_CONNECTION_STRING = "APPLICATIONINSIGHTS_CONNECTION_STRING" + APPLICATIONINSIGHTS_AUTH_MODE = "APPLICATIONINSIGHTS_AUTH_MODE" OTEL_EXPORTER_OTLP_ENDPOINT = "OTEL_EXPORTER_OTLP_ENDPOINT" FOUNDRY_AGENT365_TRACING_ENABLED = "FOUNDRY_AGENT365_TRACING_ENABLED" OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT = "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT" diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_tracing.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_tracing.py index 1266d751e4ec..9ef0bcb6eea1 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_tracing.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_tracing.py @@ -41,6 +41,7 @@ from opentelemetry import baggage as _otel_baggage, context as _otel_context, trace from . import _config +from ._constants import Constants _Content = Union[str, bytes, memoryview] @@ -246,6 +247,23 @@ def _setup_distro_export( kwargs["enable_azure_monitor"] = True kwargs["azure_monitor_connection_string"] = connection_string + # When Entra-based auth is requested, export to Azure Monitor using a + # system-assigned managed identity (no client id) rather than the + # connection string's instrumentation key alone. azure-identity is + # normally available transitively via microsoft-opentelemetry; if it + # is missing we log and continue rather than failing tracing setup. + auth_mode = os.environ.get(Constants.APPLICATIONINSIGHTS_AUTH_MODE, "") + if auth_mode.strip().lower() == "entra": + try: + from azure.identity import ManagedIdentityCredential + + kwargs["azure_monitor_exporter_credential"] = ManagedIdentityCredential() + except ImportError: + logger.warning( + "APPLICATIONINSIGHTS_AUTH_MODE=Entra requested but azure-identity is " + "not installed — continuing with connection-string authentication." + ) + # A365 tracing export — enabled only in hosted environments. if ( os.environ.get("FOUNDRY_HOSTING_ENVIRONMENT", "") diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/test_tracing.py b/sdk/agentserver/azure-ai-agentserver-core/tests/test_tracing.py index bacf0f4d6dea..cc52c6006151 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/test_tracing.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/test_tracing.py @@ -148,6 +148,78 @@ def test_distro_called_without_conn_str(self) -> None: assert kwargs["connection_string"] is None +# ------------------------------------------------------------------ # +# Entra-based Azure Monitor export credential +# ------------------------------------------------------------------ # + + +class TestEntraAuthMode: + """Verify _setup_distro_export wires a managed identity credential for Entra auth.""" + + def _run(self, env: dict) -> dict: + from azure.ai.agentserver.core import _tracing + with mock.patch("microsoft.opentelemetry.use_microsoft_opentelemetry") as mock_use, \ + mock.patch.dict(os.environ, env, clear=False): + _tracing._setup_distro_export( + resource=Resource.create({}), + span_processors=[], + log_record_processors=[], + connection_string="InstrumentationKey=00000000-0000-0000-0000-000000000000", + ) + mock_use.assert_called_once() + return mock_use.call_args[1] + + def test_entra_auth_mode_passes_managed_identity_credential(self) -> None: + sentinel = object() + with mock.patch( + "azure.identity.ManagedIdentityCredential", + return_value=sentinel, + ): + kwargs = self._run({"APPLICATIONINSIGHTS_AUTH_MODE": "Entra"}) + assert kwargs["enable_azure_monitor"] is True + assert kwargs["azure_monitor_exporter_credential"] is sentinel + + def test_entra_auth_mode_case_insensitive(self) -> None: + sentinel = object() + with mock.patch( + "azure.identity.ManagedIdentityCredential", + return_value=sentinel, + ): + kwargs = self._run({"APPLICATIONINSIGHTS_AUTH_MODE": "entra"}) + assert kwargs["azure_monitor_exporter_credential"] is sentinel + + def test_no_credential_when_auth_mode_not_entra(self) -> None: + env = {"APPLICATIONINSIGHTS_AUTH_MODE": ""} + with mock.patch.dict(os.environ, {}, clear=False): + os.environ.pop("APPLICATIONINSIGHTS_AUTH_MODE", None) + kwargs = self._run(env) + assert kwargs["enable_azure_monitor"] is True + assert "azure_monitor_exporter_credential" not in kwargs + + def test_entra_auth_continues_when_azure_identity_missing(self) -> None: + # If azure-identity cannot be imported, tracing setup must not fail: + # Azure Monitor stays enabled (connection-string auth) and no credential + # is passed. + import builtins + + real_import = builtins.__import__ + + def _fake_import(name, *args, **kwargs): + if name == "azure.identity": + raise ImportError("No module named 'azure.identity'") + return real_import(name, *args, **kwargs) + + with mock.patch("builtins.__import__", side_effect=_fake_import): + kwargs = self._run({"APPLICATIONINSIGHTS_AUTH_MODE": "Entra"}) + assert kwargs["enable_azure_monitor"] is True + assert "azure_monitor_exporter_credential" not in kwargs + + def test_managed_identity_credential_has_no_client_id(self) -> None: + with mock.patch("azure.identity.ManagedIdentityCredential") as mock_cred: + self._run({"APPLICATIONINSIGHTS_AUTH_MODE": "Entra"}) + mock_cred.assert_called_once_with() + + # ------------------------------------------------------------------ # # Constructor passes / skips connection string # ------------------------------------------------------------------ #