diff --git a/python/packages/anthropic/agent_framework_anthropic/_bedrock_client.py b/python/packages/anthropic/agent_framework_anthropic/_bedrock_client.py index 8b7b6bd71b..6c4c32738c 100644 --- a/python/packages/anthropic/agent_framework_anthropic/_bedrock_client.py +++ b/python/packages/anthropic/agent_framework_anthropic/_bedrock_client.py @@ -6,13 +6,13 @@ from typing import Any, ClassVar, Generic, TypedDict from agent_framework import ( - AGENT_FRAMEWORK_USER_AGENT, ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer, FunctionInvocationConfiguration, FunctionInvocationLayer, ) from agent_framework._settings import SecretString, load_settings +from agent_framework._telemetry import get_user_agent from agent_framework.observability import ChatTelemetryLayer from anthropic import AsyncAnthropicBedrock @@ -94,7 +94,7 @@ def __init__( aws_profile=settings.get("aws_profile"), aws_session_token=session_token_secret.get_secret_value() if session_token_secret is not None else None, base_url=settings.get("anthropic_bedrock_base_url"), - default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT}, + default_headers={"User-Agent": get_user_agent()}, ) super().__init__( diff --git a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py index 6e5e5f14a6..53d6d6a65e 100644 --- a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py +++ b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py @@ -8,7 +8,6 @@ from typing import Any, ClassVar, Final, Generic, Literal, TypedDict from agent_framework import ( - AGENT_FRAMEWORK_USER_AGENT, Annotation, BaseChatClient, ChatAndFunctionMiddlewareTypes, @@ -28,6 +27,7 @@ tool, ) from agent_framework._settings import SecretString, load_settings +from agent_framework._telemetry import get_user_agent from agent_framework._tools import SHELL_TOOL_KIND_VALUE from agent_framework._types import _get_data_bytes_as_str # type: ignore from agent_framework.observability import ChatTelemetryLayer @@ -332,7 +332,7 @@ class MyOptions(AnthropicChatOptions, total=False): anthropic_client = AsyncAnthropic( api_key=api_key_secret.get_secret_value(), - default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT}, + default_headers={"User-Agent": get_user_agent()}, ) # Initialize parent @@ -604,7 +604,7 @@ def _prepare_options( run_options["betas"] = self._prepare_betas(options) # extra headers - run_options["extra_headers"] = {"User-Agent": AGENT_FRAMEWORK_USER_AGENT} + run_options["extra_headers"] = {"User-Agent": get_user_agent()} # Handle user option -> metadata.user_id (Anthropic uses metadata.user_id instead of user) if user := run_options.pop("user", None): diff --git a/python/packages/anthropic/agent_framework_anthropic/_foundry_client.py b/python/packages/anthropic/agent_framework_anthropic/_foundry_client.py index f5fe37296f..997aac0e44 100644 --- a/python/packages/anthropic/agent_framework_anthropic/_foundry_client.py +++ b/python/packages/anthropic/agent_framework_anthropic/_foundry_client.py @@ -6,13 +6,13 @@ from typing import Any, ClassVar, Generic, TypedDict from agent_framework import ( - AGENT_FRAMEWORK_USER_AGENT, ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer, FunctionInvocationConfiguration, FunctionInvocationLayer, ) from agent_framework._settings import SecretString, load_settings +from agent_framework._telemetry import get_user_agent from agent_framework.observability import ChatTelemetryLayer from anthropic import AsyncAnthropicFoundry @@ -91,14 +91,14 @@ def __init__( base_url=base_url_setting, api_key=api_key_value, azure_ad_token_provider=azure_ad_token_provider, - default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT}, + default_headers={"User-Agent": get_user_agent()}, ) else: anthropic_client = AsyncAnthropicFoundry( resource=resource_setting, api_key=api_key_value, azure_ad_token_provider=azure_ad_token_provider, - default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT}, + default_headers={"User-Agent": get_user_agent()}, ) super().__init__( diff --git a/python/packages/anthropic/agent_framework_anthropic/_vertex_client.py b/python/packages/anthropic/agent_framework_anthropic/_vertex_client.py index 73b16af4a7..7af22be0f0 100644 --- a/python/packages/anthropic/agent_framework_anthropic/_vertex_client.py +++ b/python/packages/anthropic/agent_framework_anthropic/_vertex_client.py @@ -6,13 +6,13 @@ from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypedDict from agent_framework import ( - AGENT_FRAMEWORK_USER_AGENT, ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer, FunctionInvocationConfiguration, FunctionInvocationLayer, ) from agent_framework._settings import load_settings +from agent_framework._telemetry import get_user_agent from agent_framework.observability import ChatTelemetryLayer from anthropic import NOT_GIVEN, AsyncAnthropicVertex @@ -89,7 +89,7 @@ def __init__( access_token=access_token, credentials=credentials, base_url=settings.get("anthropic_vertex_base_url"), - default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT}, + default_headers={"User-Agent": get_user_agent()}, ) super().__init__( diff --git a/python/packages/anthropic/tests/test_anthropic_provider_clients.py b/python/packages/anthropic/tests/test_anthropic_provider_clients.py index d076062fe6..0fb208b5e4 100644 --- a/python/packages/anthropic/tests/test_anthropic_provider_clients.py +++ b/python/packages/anthropic/tests/test_anthropic_provider_clients.py @@ -3,7 +3,8 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from agent_framework import AGENT_FRAMEWORK_USER_AGENT, ChatMiddlewareLayer, FunctionInvocationLayer +from agent_framework import ChatMiddlewareLayer, FunctionInvocationLayer +from agent_framework._telemetry import get_user_agent from agent_framework.observability import ChatTelemetryLayer from agent_framework_anthropic import ( @@ -61,7 +62,7 @@ def test_raw_anthropic_foundry_client_creates_sdk_client_from_settings(tmp_path) resource="test-resource", api_key="test-key", azure_ad_token_provider=None, - default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT}, + default_headers={"User-Agent": get_user_agent()}, ) @@ -85,7 +86,7 @@ def test_raw_anthropic_foundry_client_creates_sdk_client_from_base_url_settings( base_url="https://test-resource.services.ai.azure.com/anthropic/", api_key="test-key", azure_ad_token_provider=None, - default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT}, + default_headers={"User-Agent": get_user_agent()}, ) @@ -130,7 +131,7 @@ def test_raw_anthropic_bedrock_client_creates_sdk_client_from_arguments() -> Non aws_profile=None, aws_session_token=None, base_url=None, - default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT}, + default_headers={"User-Agent": get_user_agent()}, ) @@ -152,5 +153,5 @@ def test_raw_anthropic_vertex_client_creates_sdk_client_from_arguments() -> None access_token=None, credentials=None, base_url=None, - default_headers={"User-Agent": AGENT_FRAMEWORK_USER_AGENT}, + default_headers={"User-Agent": get_user_agent()}, ) diff --git a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py index de3c190e66..5a0b79f29d 100644 --- a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py +++ b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py @@ -14,7 +14,6 @@ from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict, overload from agent_framework import ( - AGENT_FRAMEWORK_USER_AGENT, AgentSession, Annotation, Content, @@ -25,6 +24,7 @@ SupportsGetEmbeddings, load_settings, ) +from agent_framework._telemetry import get_user_agent from agent_framework.exceptions import SettingNotFoundError from azure.core.credentials import AzureKeyCredential, TokenCredential from azure.core.credentials_async import AsyncTokenCredential @@ -535,7 +535,7 @@ def __init__( endpoint=self.endpoint, index_name=self.index_name, credential=self.credential, - user_agent=AGENT_FRAMEWORK_USER_AGENT, + user_agent=get_user_agent(), ) self._index_client: SearchIndexClient | None = None @@ -544,7 +544,7 @@ def __init__( self._index_client = SearchIndexClient( endpoint=self.endpoint, credential=self.credential, - user_agent=AGENT_FRAMEWORK_USER_AGENT, + user_agent=get_user_agent(), ) self._knowledge_base_initialized = False @@ -640,7 +640,7 @@ async def _auto_discover_vector_field(self) -> None: self._index_client = SearchIndexClient( endpoint=self.endpoint, credential=self.credential, - user_agent=AGENT_FRAMEWORK_USER_AGENT, + user_agent=get_user_agent(), ) if not self.index_name: logger.warning("Cannot auto-discover vector field: index_name is not set.") @@ -740,7 +740,7 @@ async def _ensure_knowledge_base(self) -> None: endpoint=self.endpoint, knowledge_base_name=knowledge_base_name, credential=self.credential, - user_agent=AGENT_FRAMEWORK_USER_AGENT, + user_agent=get_user_agent(), ) self._knowledge_base_initialized = True return @@ -802,7 +802,7 @@ async def _ensure_knowledge_base(self) -> None: endpoint=self.endpoint, knowledge_base_name=knowledge_base_name, credential=self.credential, - user_agent=AGENT_FRAMEWORK_USER_AGENT, + user_agent=get_user_agent(), ) async def _agentic_search(self, messages: list[Message]) -> list[Message]: diff --git a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py index 496d95d7c3..915eee432b 100644 --- a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py +++ b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py @@ -7,8 +7,8 @@ import logging from typing import Any, TypedDict -from agent_framework import AGENT_FRAMEWORK_USER_AGENT from agent_framework._settings import SecretString, load_settings +from agent_framework._telemetry import get_user_agent from agent_framework._workflows._checkpoint import CheckpointID, WorkflowCheckpoint from agent_framework._workflows._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value from agent_framework.exceptions import WorkflowCheckpointException @@ -194,7 +194,7 @@ def __init__( self._cosmos_client = CosmosClient( url=settings["endpoint"], # type: ignore[arg-type] credential=credential or settings["key"].get_secret_value(), # type: ignore[arg-type,union-attr] - user_agent_suffix=AGENT_FRAMEWORK_USER_AGENT, + user_agent_suffix=get_user_agent(), ) self._owns_client = True diff --git a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_history_provider.py b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_history_provider.py index d13f285249..62a83e6a0f 100644 --- a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_history_provider.py +++ b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_history_provider.py @@ -10,9 +10,10 @@ from collections.abc import Sequence from typing import Any, ClassVar, TypedDict -from agent_framework import AGENT_FRAMEWORK_USER_AGENT, Message +from agent_framework import Message from agent_framework._sessions import HistoryProvider from agent_framework._settings import SecretString, load_settings +from agent_framework._telemetry import get_user_agent from azure.core.credentials import TokenCredential from azure.core.credentials_async import AsyncTokenCredential from azure.cosmos import PartitionKey @@ -121,7 +122,7 @@ def __init__( self._cosmos_client = CosmosClient( url=settings["endpoint"], # type: ignore[arg-type] credential=credential or settings["key"].get_secret_value(), # type: ignore[arg-type,union-attr] - user_agent_suffix=AGENT_FRAMEWORK_USER_AGENT, + user_agent_suffix=get_user_agent(), ) self._owns_client = True diff --git a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py index 3606cdf26b..8e05591372 100644 --- a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py +++ b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py @@ -13,7 +13,6 @@ from uuid import uuid4 from agent_framework import ( - AGENT_FRAMEWORK_USER_AGENT, BaseChatClient, ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer, @@ -31,6 +30,7 @@ validate_tool_mode, ) from agent_framework._settings import SecretString, load_settings +from agent_framework._telemetry import get_user_agent from agent_framework.exceptions import ChatClientInvalidResponseException from agent_framework.observability import ChatTelemetryLayer from boto3.session import Session as Boto3Session @@ -299,7 +299,7 @@ class MyOptions(BedrockChatOptions, total=False): self._bedrock_client = session.client( "bedrock-runtime", region_name=region, - config=BotoConfig(user_agent_extra=AGENT_FRAMEWORK_USER_AGENT), + config=BotoConfig(user_agent_extra=get_user_agent()), ) super().__init__( diff --git a/python/packages/bedrock/agent_framework_bedrock/_embedding_client.py b/python/packages/bedrock/agent_framework_bedrock/_embedding_client.py index 99250b8248..52f5126e3d 100644 --- a/python/packages/bedrock/agent_framework_bedrock/_embedding_client.py +++ b/python/packages/bedrock/agent_framework_bedrock/_embedding_client.py @@ -11,7 +11,6 @@ from typing import Any, ClassVar, Generic, TypedDict from agent_framework import ( - AGENT_FRAMEWORK_USER_AGENT, BaseEmbeddingClient, Embedding, EmbeddingGenerationOptions, @@ -20,6 +19,7 @@ UsageDetails, load_settings, ) +from agent_framework._telemetry import get_user_agent from agent_framework.observability import EmbeddingTelemetryLayer from boto3.session import Session as Boto3Session from botocore.client import BaseClient @@ -140,7 +140,7 @@ def __init__( self._bedrock_client = boto3_session.client( "bedrock-runtime", region_name=region_name or resolved_region, - config=BotoConfig(user_agent_extra=AGENT_FRAMEWORK_USER_AGENT), + config=BotoConfig(user_agent_extra=get_user_agent()), ) self.model: str = settings["embedding_model"] # type: ignore[assignment] # pyright: ignore[reportTypedDictNotRequiredAccess] diff --git a/python/packages/core/agent_framework/_telemetry.py b/python/packages/core/agent_framework/_telemetry.py index a3b0f74146..f7ca2ce030 100644 --- a/python/packages/core/agent_framework/_telemetry.py +++ b/python/packages/core/agent_framework/_telemetry.py @@ -4,9 +4,6 @@ import logging import os -from collections.abc import Generator -from contextlib import contextmanager -from contextvars import ContextVar from typing import Any, Final from . import __version__ as version_info @@ -29,34 +26,73 @@ HTTP_USER_AGENT: Final[str] = "agent-framework-python" AGENT_FRAMEWORK_USER_AGENT = f"{HTTP_USER_AGENT}/{version_info}" # type: ignore[has-type] -_user_agent_prefixes: ContextVar[tuple[str, ...]] = ContextVar("_user_agent_prefixes", default=()) +# This environment variable is reserved by the Foundry hosting environment to +# indicate that the agent is running in a hosted environment. +_FOUNDRY_HOSTING_ENV_VAR = "FOUNDRY_HOSTING_ENVIRONMENT" +# This prefix is added to the user agent string when the agent is running in a hosted environment. +_HOSTED_USER_AGENT_PREFIX = "foundry-hosting" +_user_agent_prefixes: set[str] = set() +_hosted_env_detected: bool = False -@contextmanager -def user_agent_prefix(prefix: str) -> Generator[None]: - """Context manager that adds a prefix to the user agent string for the current scope. - This is useful for upstream layers that want to identify themselves in telemetry - for the duration of a request without permanently mutating global state. +def _add_user_agent_prefix(prefix: str) -> None: + """Permanently add a prefix to the user agent string. + + This is used by hosting layers to identify themselves in telemetry. + Once added, the prefix applies to all subsequent user agent strings. Args: prefix: The prefix to add (e.g. "foundry-hosting"). """ - current = _user_agent_prefixes.get() - token = _user_agent_prefixes.set((*current, prefix)) if prefix and prefix not in current else None + if prefix: + _user_agent_prefixes.add(prefix) + + +def _detect_hosted_environment() -> None: + """Detect if running in a hosted environment and add the user agent prefix. + + Checks the ``FOUNDRY_HOSTING_ENVIRONMENT`` env var first, then falls back + to checking whether the agent server SDK is installed (via + ``importlib.util.find_spec``) before importing it, to avoid unnecessary + import overhead for non-hosted scenarios. + """ + global _hosted_env_detected + if _hosted_env_detected: + return + _hosted_env_detected = True + + env_value = os.environ.get(_FOUNDRY_HOSTING_ENV_VAR) + if env_value is not None: + # Env var exists — trust its value and skip the fallback. + if env_value: + _add_user_agent_prefix(_HOSTED_USER_AGENT_PREFIX) + return + + # Env var not set — fall back to AgentConfig as a second layer of defense. + # Use find_spec to avoid the cost of a full import when the SDK is not installed. + import importlib.util + try: - yield - finally: - if token is not None: - _user_agent_prefixes.reset(token) + if importlib.util.find_spec("azure.ai.agentserver.core") is None: + return + except (ModuleNotFoundError, ValueError): + return + try: + from azure.ai.agentserver.core import AgentConfig # pyright: ignore[reportMissingImports] + + if AgentConfig.from_env().is_hosted: + _add_user_agent_prefix(_HOSTED_USER_AGENT_PREFIX) + except (ImportError, AttributeError): + pass -def _get_user_agent() -> str: - """Return the full user agent string including any context-scoped prefixes.""" - prefixes = _user_agent_prefixes.get() - if not prefixes: +def get_user_agent() -> str: + """Return the full user agent string including any registered prefixes.""" + _detect_hosted_environment() + if not _user_agent_prefixes: return AGENT_FRAMEWORK_USER_AGENT - return f"{'/'.join(prefixes)}/{AGENT_FRAMEWORK_USER_AGENT}" + return f"{'/'.join(sorted(_user_agent_prefixes))}/{AGENT_FRAMEWORK_USER_AGENT}" def prepend_agent_framework_to_user_agent(headers: dict[str, Any] | None = None) -> dict[str, Any]: @@ -89,7 +125,7 @@ def prepend_agent_framework_to_user_agent(headers: dict[str, Any] | None = None) """ if not IS_TELEMETRY_ENABLED: return headers or {} - user_agent = _get_user_agent() + user_agent = get_user_agent() if not headers: return {USER_AGENT_KEY: user_agent} headers[USER_AGENT_KEY] = f"{user_agent} {headers[USER_AGENT_KEY]}" if USER_AGENT_KEY in headers else user_agent diff --git a/python/packages/core/tests/core/test_telemetry.py b/python/packages/core/tests/core/test_telemetry.py index 1ba1df1dde..b0b01706ef 100644 --- a/python/packages/core/tests/core/test_telemetry.py +++ b/python/packages/core/tests/core/test_telemetry.py @@ -1,14 +1,21 @@ # Copyright (c) Microsoft. All rights reserved. -from unittest.mock import patch +import os +from unittest.mock import MagicMock, patch +import agent_framework._telemetry as _telemetry_mod from agent_framework import ( AGENT_FRAMEWORK_USER_AGENT, USER_AGENT_KEY, USER_AGENT_TELEMETRY_DISABLED_ENV_VAR, prepend_agent_framework_to_user_agent, ) -from agent_framework._telemetry import user_agent_prefix +from agent_framework._telemetry import ( + _FOUNDRY_HOSTING_ENV_VAR, + _HOSTED_USER_AGENT_PREFIX, + _add_user_agent_prefix, + _detect_hosted_environment, +) # region Test constants @@ -83,7 +90,7 @@ def test_prepend_to_empty_headers(): def test_prepend_to_empty_dict(): """Test prepending to empty headers dict.""" - headers = {} + headers: dict[str, str] = {} result = prepend_agent_framework_to_user_agent(headers) assert "User-Agent" in result @@ -99,54 +106,184 @@ def test_modifies_original_dict(): assert "User-Agent" in headers -# region Test user_agent_prefix context manager - +# region Test _add_user_agent_prefix -def test_user_agent_prefix_adds_prefix(): - """Test that the context manager adds a prefix within its scope.""" - with user_agent_prefix("test-host"): - result = prepend_agent_framework_to_user_agent() - assert result["User-Agent"].startswith("test-host/") - assert AGENT_FRAMEWORK_USER_AGENT in result["User-Agent"] - # Prefix is removed after exiting the context +def test_add_user_agent_prefix_adds_prefix(): + """Test that _add_user_agent_prefix permanently adds a prefix.""" + _telemetry_mod._user_agent_prefixes.clear() + _add_user_agent_prefix("test-host") result = prepend_agent_framework_to_user_agent() - assert result["User-Agent"] == AGENT_FRAMEWORK_USER_AGENT + assert result["User-Agent"].startswith("test-host/") + assert AGENT_FRAMEWORK_USER_AGENT in result["User-Agent"] + _telemetry_mod._user_agent_prefixes.clear() -def test_user_agent_prefix_ignores_duplicates(): - """Test that duplicate prefixes are not added within nested scopes.""" - with user_agent_prefix("test-host"), user_agent_prefix("test-host"): - result = prepend_agent_framework_to_user_agent() - assert result["User-Agent"].count("test-host") == 1 +def test_add_user_agent_prefix_ignores_duplicates(): + """Test that duplicate prefixes are not added.""" + _telemetry_mod._user_agent_prefixes.clear() + _add_user_agent_prefix("test-host") + _add_user_agent_prefix("test-host") + result = prepend_agent_framework_to_user_agent() + assert result["User-Agent"].count("test-host") == 1 + _telemetry_mod._user_agent_prefixes.clear() -def test_user_agent_prefix_ignores_empty(): +def test_add_user_agent_prefix_ignores_empty(): """Test that empty strings are not added as prefixes.""" - with user_agent_prefix(""): - result = prepend_agent_framework_to_user_agent() - assert result["User-Agent"] == AGENT_FRAMEWORK_USER_AGENT - - -def test_user_agent_prefix_restores_on_exit(): - """Test that prefixes are fully restored after the context manager exits.""" - with user_agent_prefix("test-host"): - pass + _telemetry_mod._user_agent_prefixes.clear() + _add_user_agent_prefix("") result = prepend_agent_framework_to_user_agent() assert result["User-Agent"] == AGENT_FRAMEWORK_USER_AGENT + _telemetry_mod._user_agent_prefixes.clear() -def test_user_agent_prefix_nesting(): - """Test that nested context managers compose prefixes correctly.""" - with user_agent_prefix("outer"): - with user_agent_prefix("inner"): - result = prepend_agent_framework_to_user_agent() - assert "outer" in result["User-Agent"] - assert "inner" in result["User-Agent"] - # Inner prefix removed - result = prepend_agent_framework_to_user_agent() - assert "outer" in result["User-Agent"] - assert "inner" not in result["User-Agent"] - # Both removed +def test_add_user_agent_prefix_multiple(): + """Test that multiple prefixes compose correctly.""" + _telemetry_mod._user_agent_prefixes.clear() + _add_user_agent_prefix("outer") + _add_user_agent_prefix("inner") result = prepend_agent_framework_to_user_agent() - assert result["User-Agent"] == AGENT_FRAMEWORK_USER_AGENT + assert "outer" in result["User-Agent"] + assert "inner" in result["User-Agent"] + _telemetry_mod._user_agent_prefixes.clear() + + +# region Test _detect_hosted_environment + + +def test_detect_hosted_env_var_truthy_adds_prefix(): + """Test that a truthy FOUNDRY_HOSTING_ENVIRONMENT env var adds the prefix.""" + _telemetry_mod._user_agent_prefixes.clear() + _telemetry_mod._hosted_env_detected = False + with patch.dict("os.environ", {_FOUNDRY_HOSTING_ENV_VAR: "production"}): + _detect_hosted_environment() + assert _HOSTED_USER_AGENT_PREFIX in _telemetry_mod._user_agent_prefixes + _telemetry_mod._user_agent_prefixes.clear() + _telemetry_mod._hosted_env_detected = False + + +def test_detect_hosted_env_var_empty_skips_prefix(): + """Test that an empty FOUNDRY_HOSTING_ENVIRONMENT env var does NOT add the prefix.""" + _telemetry_mod._user_agent_prefixes.clear() + _telemetry_mod._hosted_env_detected = False + with patch.dict("os.environ", {_FOUNDRY_HOSTING_ENV_VAR: ""}): + _detect_hosted_environment() + assert _HOSTED_USER_AGENT_PREFIX not in _telemetry_mod._user_agent_prefixes + _telemetry_mod._user_agent_prefixes.clear() + _telemetry_mod._hosted_env_detected = False + + +def test_detect_hosted_env_var_set_skips_agent_config_fallback(): + """Test that when the env var is set, AgentConfig is never consulted even if import would fail.""" + _telemetry_mod._user_agent_prefixes.clear() + _telemetry_mod._hosted_env_detected = False + import builtins + + real_import = builtins.__import__ + + def _block_agentconfig(name: str, *args, **kwargs): # type: ignore[no-untyped-def] + if "agentserver" in name: + raise AssertionError("AgentConfig should not be imported when env var is set") + return real_import(name, *args, **kwargs) + + with ( + patch.dict("os.environ", {_FOUNDRY_HOSTING_ENV_VAR: "prod"}), + patch("builtins.__import__", side_effect=_block_agentconfig), + ): + _detect_hosted_environment() + assert _HOSTED_USER_AGENT_PREFIX in _telemetry_mod._user_agent_prefixes + _telemetry_mod._user_agent_prefixes.clear() + _telemetry_mod._hosted_env_detected = False + + +def _mock_agent_config(*, is_hosted: bool) -> MagicMock: + """Create a mock azure.ai.agentserver.core module with AgentConfig.""" + mock_config = MagicMock() + mock_config.is_hosted = is_hosted + mock_module = MagicMock() + mock_module.AgentConfig.from_env.return_value = mock_config + return mock_module + + +def test_detect_hosted_fallback_agent_config_is_hosted(): + """Test that AgentConfig fallback adds the prefix when is_hosted is True.""" + _telemetry_mod._user_agent_prefixes.clear() + _telemetry_mod._hosted_env_detected = False + env = {k: v for k, v in os.environ.items() if k != _FOUNDRY_HOSTING_ENV_VAR} + mock_module = _mock_agent_config(is_hosted=True) + mock_spec = MagicMock() + with ( + patch.dict("os.environ", env, clear=True), + patch.dict("sys.modules", {"azure.ai.agentserver.core": mock_module}), + patch("importlib.util.find_spec", return_value=mock_spec), + ): + _detect_hosted_environment() + assert _HOSTED_USER_AGENT_PREFIX in _telemetry_mod._user_agent_prefixes + _telemetry_mod._user_agent_prefixes.clear() + _telemetry_mod._hosted_env_detected = False + + +def test_detect_hosted_fallback_agent_config_not_hosted(): + """Test that AgentConfig fallback does NOT add the prefix when is_hosted is False.""" + _telemetry_mod._user_agent_prefixes.clear() + _telemetry_mod._hosted_env_detected = False + mock_module = _mock_agent_config(is_hosted=False) + mock_spec = MagicMock() + env = {k: v for k, v in os.environ.items() if k != _FOUNDRY_HOSTING_ENV_VAR} + with ( + patch.dict("os.environ", env, clear=True), + patch.dict("sys.modules", {"azure.ai.agentserver.core": mock_module}), + patch("importlib.util.find_spec", return_value=mock_spec), + ): + _detect_hosted_environment() + assert _HOSTED_USER_AGENT_PREFIX not in _telemetry_mod._user_agent_prefixes + _telemetry_mod._user_agent_prefixes.clear() + _telemetry_mod._hosted_env_detected = False + + +def test_detect_hosted_fallback_import_error(): + """Test that ImportError from AgentConfig is silently handled.""" + _telemetry_mod._user_agent_prefixes.clear() + _telemetry_mod._hosted_env_detected = False + env = {k: v for k, v in os.environ.items() if k != _FOUNDRY_HOSTING_ENV_VAR} + with patch.dict("os.environ", env, clear=True): + # The real import may succeed or fail depending on the environment; + # force the ImportError path by making the import raise. + import builtins + + real_import = builtins.__import__ + + def _block_agentconfig(name: str, *args, **kwargs): # type: ignore[no-untyped-def] + if "agentserver" in name: + raise ImportError("mocked") + return real_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=_block_agentconfig): + _detect_hosted_environment() + assert _HOSTED_USER_AGENT_PREFIX not in _telemetry_mod._user_agent_prefixes + _telemetry_mod._user_agent_prefixes.clear() + _telemetry_mod._hosted_env_detected = False + + +# region Test module-level auto-detection + + +def test_lazy_detection_on_get_user_agent(): + """Test that get_user_agent() lazily detects the hosted environment. + + Since detection is deferred to the first ``get_user_agent()`` call, + this verifies the prefix is included without any explicit call to + ``_detect_hosted_environment()`` by consumer code. + """ + _telemetry_mod._user_agent_prefixes.clear() + _telemetry_mod._hosted_env_detected = False + with patch.dict("os.environ", {_FOUNDRY_HOSTING_ENV_VAR: "production"}): + user_agent = _telemetry_mod.get_user_agent() + + assert _HOSTED_USER_AGENT_PREFIX in _telemetry_mod._user_agent_prefixes + assert user_agent.startswith(f"{_HOSTED_USER_AGENT_PREFIX}/") + + # Clean up + _telemetry_mod._user_agent_prefixes.clear() + _telemetry_mod._hosted_env_detected = False diff --git a/python/packages/foundry/agent_framework_foundry/_agent.py b/python/packages/foundry/agent_framework_foundry/_agent.py index 0c7f93ba1f..2a0e4beb3f 100644 --- a/python/packages/foundry/agent_framework_foundry/_agent.py +++ b/python/packages/foundry/agent_framework_foundry/_agent.py @@ -15,7 +15,6 @@ from typing import TYPE_CHECKING, Any, ClassVar, Generic, cast from agent_framework import ( - AGENT_FRAMEWORK_USER_AGENT, AgentMiddlewareLayer, ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer, @@ -28,6 +27,7 @@ load_settings, ) from agent_framework._compaction import CompactionStrategy, TokenizerProtocol +from agent_framework._telemetry import get_user_agent from agent_framework.observability import AgentTelemetryLayer, ChatTelemetryLayer from agent_framework_openai._chat_client import OpenAIChatOptions, RawOpenAIChatClient from azure.ai.projects.aio import AIProjectClient @@ -190,7 +190,7 @@ def __init__( project_client_kwargs: dict[str, Any] = { "endpoint": resolved_endpoint, "credential": credential, - "user_agent": AGENT_FRAMEWORK_USER_AGENT, + "user_agent": get_user_agent(), } if allow_preview is not None: project_client_kwargs["allow_preview"] = allow_preview diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py index 735ba9fb57..d097fd73a1 100644 --- a/python/packages/foundry/agent_framework_foundry/_chat_client.py +++ b/python/packages/foundry/agent_framework_foundry/_chat_client.py @@ -8,7 +8,6 @@ from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal from agent_framework import ( - AGENT_FRAMEWORK_USER_AGENT, ChatMiddlewareLayer, Content, FunctionInvocationConfiguration, @@ -17,6 +16,7 @@ ) from agent_framework._compaction import CompactionStrategy, TokenizerProtocol from agent_framework._feature_stage import ExperimentalFeature, experimental +from agent_framework._telemetry import get_user_agent from agent_framework.observability import ChatTelemetryLayer from agent_framework_openai._chat_client import OpenAIChatOptions, RawOpenAIChatClient from azure.ai.projects.aio import AIProjectClient @@ -198,7 +198,7 @@ def __init__( project_client_kwargs: dict[str, Any] = { "endpoint": project_endpoint, "credential": credential, # type: ignore[arg-type] - "user_agent": AGENT_FRAMEWORK_USER_AGENT, + "user_agent": get_user_agent(), } if allow_preview is not None: project_client_kwargs["allow_preview"] = allow_preview diff --git a/python/packages/foundry/agent_framework_foundry/_memory_provider.py b/python/packages/foundry/agent_framework_foundry/_memory_provider.py index c49f950b6b..169da4fc85 100644 --- a/python/packages/foundry/agent_framework_foundry/_memory_provider.py +++ b/python/packages/foundry/agent_framework_foundry/_memory_provider.py @@ -14,13 +14,13 @@ from typing import TYPE_CHECKING, Any, ClassVar from agent_framework import ( - AGENT_FRAMEWORK_USER_AGENT, AgentSession, ContextProvider, Message, SessionContext, load_settings, ) +from agent_framework._telemetry import get_user_agent from azure.ai.projects.aio import AIProjectClient from azure.core.credentials import TokenCredential from azure.core.credentials_async import AsyncTokenCredential @@ -119,7 +119,7 @@ def __init__( project_client_kwargs: dict[str, Any] = { "endpoint": resolved_endpoint, "credential": credential, # type: ignore[arg-type] - "user_agent": AGENT_FRAMEWORK_USER_AGENT, + "user_agent": get_user_agent(), } if allow_preview is not None: project_client_kwargs["allow_preview"] = allow_preview diff --git a/python/packages/foundry/tests/foundry/test_foundry_chat_client.py b/python/packages/foundry/tests/foundry/test_foundry_chat_client.py index 68e7adc6fb..a7b57029ba 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_chat_client.py +++ b/python/packages/foundry/tests/foundry/test_foundry_chat_client.py @@ -12,7 +12,7 @@ import pytest from agent_framework import ChatResponse, Content, Message, SupportsChatGetResponse, tool -from agent_framework._telemetry import AGENT_FRAMEWORK_USER_AGENT +from agent_framework._telemetry import get_user_agent from agent_framework.exceptions import ChatClientException, ChatClientInvalidRequestException from agent_framework_openai import OpenAIContentFilterException from azure.ai.projects.models import MCPTool as FoundryMCPTool @@ -199,7 +199,7 @@ def test_init_with_project_endpoint_creates_project_client() -> None: assert factory.call_args.kwargs["endpoint"] == _TEST_FOUNDRY_PROJECT_ENDPOINT assert factory.call_args.kwargs["credential"] is credential assert factory.call_args.kwargs["allow_preview"] is True - assert factory.call_args.kwargs["user_agent"] == AGENT_FRAMEWORK_USER_AGENT + assert factory.call_args.kwargs["user_agent"] == get_user_agent() def test_init_with_empty_model_raises(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py b/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py index 2e25ba38d4..6377dfa602 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py +++ b/python/packages/foundry/tests/foundry/test_foundry_memory_provider.py @@ -7,8 +7,9 @@ from unittest.mock import AsyncMock, Mock, patch import pytest -from agent_framework import AGENT_FRAMEWORK_USER_AGENT, AgentResponse, Message +from agent_framework import AgentResponse, Message from agent_framework._sessions import AgentSession, SessionContext +from agent_framework._telemetry import get_user_agent from agent_framework_foundry._memory_provider import FoundryMemoryProvider @@ -94,7 +95,7 @@ def test_init_with_project_endpoint_and_credential(mock_project_client: AsyncMoc endpoint="https://test.project.endpoint", credential=mock_credential, allow_preview=True, - user_agent=AGENT_FRAMEWORK_USER_AGENT, + user_agent=get_user_agent(), ) diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py index b5bf98291b..05105ec768 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py @@ -1,7 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. from agent_framework import AgentSession, BaseAgent, SupportsAgentRun -from agent_framework._telemetry import user_agent_prefix from azure.ai.agentserver.invocations import InvocationAgentServerHost from starlette.requests import Request from starlette.responses import JSONResponse, Response, StreamingResponse @@ -11,8 +10,6 @@ class InvocationsHostServer(InvocationAgentServerHost): """An invocations server host for an agent.""" - USER_AGENT_PREFIX = "foundry-hosting" - def __init__( self, agent: BaseAgent, @@ -42,11 +39,6 @@ def __init__( async def _handle_invoke(self, request: Request) -> Response: """Invoke the agent with the given request.""" - with user_agent_prefix(self.USER_AGENT_PREFIX): - return await self._handle_invoke_inner(request) - - async def _handle_invoke_inner(self, request: Request) -> Response: - """Core invoke handler logic.""" data = await request.json() session_id: str = request.state.session_id diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index cdd4b34b4f..63bc730e78 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -20,7 +20,6 @@ SupportsAgentRun, WorkflowAgent, ) -from agent_framework._telemetry import user_agent_prefix from azure.ai.agentserver.responses import ( ResponseContext, ResponseEventStream, @@ -90,7 +89,6 @@ class ResponsesHostServer(ResponsesAgentServerHost): """A responses server host for an agent.""" - USER_AGENT_PREFIX = "foundry-hosting" # TODO(@taochen): Allow a different checkpoint storage that stores checkpoints externally CHECKPOINT_STORAGE_PATH = "/.checkpoints" @@ -150,37 +148,32 @@ def __init__( self._is_workflow_agent = True self._agent = agent - self.response_handler(self._handler) # pyright: ignore[reportUnknownMemberType] + self.response_handler(self._handle_response) # pyright: ignore[reportUnknownMemberType] @staticmethod def _is_streaming_request(request: CreateResponse) -> bool: """Check if the request is a streaming request.""" return request.stream is not None and request.stream is True - async def _handler( + def _handle_response( self, request: CreateResponse, context: ResponseContext, cancellation_signal: asyncio.Event, ) -> AsyncIterable[ResponseStreamEvent | dict[str, Any]]: """Handle the creation of a response.""" - with user_agent_prefix(self.USER_AGENT_PREFIX): - async for event in self._handle_inner(request, context, cancellation_signal): - yield event + if self._is_workflow_agent: + # Workflow agents are handled differently because they require checkpoint restoration + return self._handle_workflow_agent(request, context) + + return self._handle_regular_agent(request, context) - async def _handle_inner( + async def _handle_regular_agent( self, request: CreateResponse, context: ResponseContext, - cancellation_signal: asyncio.Event, ) -> AsyncIterable[ResponseStreamEvent | dict[str, Any]]: - """Core handler logic.""" - if self._is_workflow_agent: - # Workflow agents are handled differently because they require checkpoint restoration - async for event in self._handle_workflow_agent(request, context, cancellation_signal): - yield event - return - + """Handle the creation of a response for a regular (non-workflow) agent.""" input_text = await context.get_input_text() history = await context.get_history() messages: list[str | Content | Message] = [*_to_messages(history), input_text] @@ -243,7 +236,6 @@ async def _handle_workflow_agent( self, request: CreateResponse, context: ResponseContext, - cancellation_signal: asyncio.Event, ) -> AsyncIterable[ResponseStreamEvent | dict[str, Any]]: """Handle the creation of a response for a workflow agent. diff --git a/python/packages/gemini/agent_framework_gemini/_chat_client.py b/python/packages/gemini/agent_framework_gemini/_chat_client.py index b0fa52a676..59a4033fa2 100644 --- a/python/packages/gemini/agent_framework_gemini/_chat_client.py +++ b/python/packages/gemini/agent_framework_gemini/_chat_client.py @@ -10,7 +10,6 @@ from uuid import uuid4 from agent_framework import ( - AGENT_FRAMEWORK_USER_AGENT, BaseChatClient, ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer, @@ -28,6 +27,7 @@ validate_tool_mode, ) from agent_framework._settings import SecretString, load_settings +from agent_framework._telemetry import get_user_agent from agent_framework.observability import ChatTelemetryLayer from google import genai from google.auth.credentials import Credentials @@ -355,7 +355,7 @@ def __init__( ) client_kwargs: dict[str, Any] = { - "http_options": {"headers": {"x-goog-api-client": AGENT_FRAMEWORK_USER_AGENT}}, + "http_options": {"headers": {"x-goog-api-client": get_user_agent()}}, } if configured_vertexai is not None: client_kwargs["vertexai"] = configured_vertexai diff --git a/python/packages/purview/agent_framework_purview/_client.py b/python/packages/purview/agent_framework_purview/_client.py index af5c3f8224..43c8adce4e 100644 --- a/python/packages/purview/agent_framework_purview/_client.py +++ b/python/packages/purview/agent_framework_purview/_client.py @@ -11,7 +11,7 @@ from uuid import uuid4 import httpx -from agent_framework import AGENT_FRAMEWORK_USER_AGENT +from agent_framework._telemetry import get_user_agent from agent_framework.observability import get_tracer from azure.core.credentials import TokenCredential from azure.core.credentials_async import AsyncTokenCredential @@ -189,7 +189,7 @@ async def _post( payload = model.model_dump(by_alias=True, exclude_none=True, mode="json") request_headers = { "Authorization": f"Bearer {token}", - "User-Agent": AGENT_FRAMEWORK_USER_AGENT, + "User-Agent": get_user_agent(), "Content-Type": "application/json", } if correlation_id: