mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
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
committed by
GitHub
Unverified
GitHub
parent
b084d0461d
commit
0989e68d1c
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user