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 b0ed26bbeda1..9e5310d7ff6f 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 @@ -246,6 +246,15 @@ 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. + auth_mode = os.environ.get("APPLICATIONINSIGHTS_AUTH_MODE", "") + if auth_mode.strip().lower() == "entra": + credential = _create_managed_identity_credential() + if credential is not None: + kwargs["azure_monitor_exporter_credential"] = credential + # A365 tracing export — enabled only in hosted environments. if ( os.environ.get("FOUNDRY_HOSTING_ENVIRONMENT", "") @@ -259,6 +268,28 @@ def _setup_distro_export( use_microsoft_opentelemetry(**kwargs) +def _create_managed_identity_credential() -> Optional[Any]: + """Create a system-assigned :class:`ManagedIdentityCredential`. + + Instantiated without a ``client_id`` so the platform-assigned (system) + managed identity is used for Entra-based Azure Monitor export. Returns + ``None`` when ``azure-identity`` is not installed so tracing setup can + continue without export credentials. + + :return: A managed identity credential, or ``None`` if unavailable. + :rtype: Any or None + """ + try: + from azure.identity import ManagedIdentityCredential + except ImportError: + logger.warning( + "APPLICATIONINSIGHTS_AUTH_MODE=Entra requires azure-identity, " + "which is not installed — Azure Monitor export credential disabled." + ) + return None + return ManagedIdentityCredential() + + # ====================================================================== # Public API: span operations # ====================================================================== 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..40a3a35258a3 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,60 @@ 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.ai.agentserver.core._tracing._create_managed_identity_credential", + 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.ai.agentserver.core._tracing._create_managed_identity_credential", + 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 "azure_monitor_exporter_credential" not in kwargs + + def test_managed_identity_credential_has_no_client_id(self) -> None: + from azure.ai.agentserver.core import _tracing + with mock.patch("azure.identity.ManagedIdentityCredential") as mock_cred: + _tracing._create_managed_identity_credential() + mock_cred.assert_called_once_with() + + # ------------------------------------------------------------------ # # Constructor passes / skips connection string # ------------------------------------------------------------------ #