Python: Fix user agent prefix (#5455)

* Fix hosting user agent missing

* Fix other providers

* Add more tests

* comments

* Fix tests
This commit is contained in:
Tao Chen
2026-04-23 16:40:38 -07:00
committed by GitHub
Unverified
parent b084d0461d
commit 0989e68d1c
21 changed files with 290 additions and 130 deletions
@@ -6,13 +6,13 @@ from collections.abc import Sequence
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 @@ class RawAnthropicBedrockClient(RawAnthropicClient[AnthropicOptionsT], Generic[A
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__(
@@ -8,7 +8,6 @@ from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, Sequenc
from typing import Any, ClassVar, Final, Generic, Literal, TypedDict
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
Annotation,
BaseChatClient,
ChatAndFunctionMiddlewareTypes,
@@ -28,6 +27,7 @@ from agent_framework import (
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 RawAnthropicClient(
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 @@ class RawAnthropicClient(
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):
@@ -6,13 +6,13 @@ from collections.abc import Awaitable, Callable, Sequence
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 @@ class RawAnthropicFoundryClient(RawAnthropicClient[AnthropicOptionsT], Generic[A
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__(
@@ -6,13 +6,13 @@ from collections.abc import Sequence
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 @@ class RawAnthropicVertexClient(RawAnthropicClient[AnthropicOptionsT], Generic[An
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__(
@@ -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()},
)
@@ -14,7 +14,6 @@ from collections.abc import Awaitable, Callable
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 @@ from agent_framework import (
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 @@ class AzureAISearchContextProvider(ContextProvider):
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 @@ class AzureAISearchContextProvider(ContextProvider):
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 @@ class AzureAISearchContextProvider(ContextProvider):
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 @@ class AzureAISearchContextProvider(ContextProvider):
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 @@ class AzureAISearchContextProvider(ContextProvider):
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]:
@@ -7,8 +7,8 @@ from __future__ import annotations
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 @@ class CosmosCheckpointStorage:
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
@@ -10,9 +10,10 @@ import uuid
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 @@ class CosmosHistoryProvider(HistoryProvider):
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
@@ -13,7 +13,6 @@ from typing import Any, ClassVar, Generic, Literal, TypedDict
from uuid import uuid4
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
BaseChatClient,
ChatAndFunctionMiddlewareTypes,
ChatMiddlewareLayer,
@@ -31,6 +30,7 @@ from agent_framework import (
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 BedrockChatClient(
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__(
@@ -11,7 +11,6 @@ from collections.abc import Sequence
from typing import Any, ClassVar, Generic, TypedDict
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
BaseEmbeddingClient,
Embedding,
EmbeddingGenerationOptions,
@@ -20,6 +19,7 @@ from agent_framework import (
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 @@ class RawBedrockEmbeddingClient(
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]
@@ -4,9 +4,6 @@ from __future__ import annotations
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 @@ USER_AGENT_KEY: Final[str] = "User-Agent"
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.
def _add_user_agent_prefix(prefix: str) -> None:
"""Permanently add a prefix to the user agent string.
This is useful for upstream layers that want to identify themselves in telemetry
for the duration of a request without permanently mutating global state.
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
+178 -41
View File
@@ -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
@@ -15,7 +15,6 @@ from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequen
from typing import TYPE_CHECKING, Any, ClassVar, Generic, cast
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
AgentMiddlewareLayer,
ChatAndFunctionMiddlewareTypes,
ChatMiddlewareLayer,
@@ -28,6 +27,7 @@ from agent_framework import (
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 @@ class RawFoundryAgentChatClient( # type: ignore[misc]
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
@@ -8,7 +8,6 @@ from collections.abc import Awaitable, Callable, Mapping, Sequence
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 import (
)
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 @@ class RawFoundryChatClient( # type: ignore[misc]
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
@@ -14,13 +14,13 @@ from contextlib import AbstractAsyncContextManager
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 @@ class FoundryMemoryProvider(ContextProvider):
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
@@ -12,7 +12,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
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:
@@ -7,8 +7,9 @@ import os
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(),
)
@@ -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 @@ from typing_extensions import Any, AsyncGenerator
class InvocationsHostServer(InvocationAgentServerHost):
"""An invocations server host for an agent."""
USER_AGENT_PREFIX = "foundry-hosting"
def __init__(
self,
agent: BaseAgent,
@@ -42,11 +39,6 @@ class InvocationsHostServer(InvocationAgentServerHost):
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
@@ -20,7 +20,6 @@ from agent_framework import (
SupportsAgentRun,
WorkflowAgent,
)
from agent_framework._telemetry import user_agent_prefix
from azure.ai.agentserver.responses import (
ResponseContext,
ResponseEventStream,
@@ -90,7 +89,6 @@ logger = logging.getLogger(__name__)
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 @@ class ResponsesHostServer(ResponsesAgentServerHost):
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)
async def _handle_inner(
return self._handle_regular_agent(request, context)
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 @@ class ResponsesHostServer(ResponsesAgentServerHost):
self,
request: CreateResponse,
context: ResponseContext,
cancellation_signal: asyncio.Event,
) -> AsyncIterable[ResponseStreamEvent | dict[str, Any]]:
"""Handle the creation of a response for a workflow agent.
@@ -10,7 +10,6 @@ from typing import Any, ClassVar, Generic, cast
from uuid import uuid4
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
BaseChatClient,
ChatAndFunctionMiddlewareTypes,
ChatMiddlewareLayer,
@@ -28,6 +27,7 @@ from agent_framework import (
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 @@ class RawGeminiChatClient(
)
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
@@ -11,7 +11,7 @@ from typing import Any, Literal, TypeVar, Union, overload
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 @@ class PurviewClient:
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: