From 03f7dc86d345fa96ec0759a439f04a8989351e70 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Thu, 23 Apr 2026 13:49:27 -0700 Subject: [PATCH] Fix hosting user agent missing --- .../core/agent_framework/_telemetry.py | 61 ++++-- .../core/tests/core/test_telemetry.py | 180 ++++++++++++++---- .../_invocations.py | 8 - .../_responses.py | 26 +-- 4 files changed, 190 insertions(+), 85 deletions(-) diff --git a/python/packages/core/agent_framework/_telemetry.py b/python/packages/core/agent_framework/_telemetry.py index a3b0f74146..2b91b3938c 100644 --- a/python/packages/core/agent_framework/_telemetry.py +++ b/python/packages/core/agent_framework/_telemetry.py @@ -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,60 @@ 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() -@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 importing ``AgentConfig`` from the agent server SDK as a second layer. + """ + 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. try: - yield - finally: - if token is not None: - _user_agent_prefixes.reset(token) + 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 + + +# Run detection of hosted environment on import +_detect_hosted_environment() def _get_user_agent() -> str: - """Return the full user agent string including any context-scoped prefixes.""" - prefixes = _user_agent_prefixes.get() - if not prefixes: + """Return the full user agent string including any registered prefixes.""" + 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]: diff --git a/python/packages/core/tests/core/test_telemetry.py b/python/packages/core/tests/core/test_telemetry.py index 1ba1df1dde..b27f407f0c 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,145 @@ 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() + 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() + + +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() + 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() + + +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() + 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() + + +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() + env = {k: v for k, v in os.environ.items() if k != _FOUNDRY_HOSTING_ENV_VAR} + mock_module = _mock_agent_config(is_hosted=True) + with ( + patch.dict("os.environ", env, clear=True), + patch.dict("sys.modules", {"azure.ai.agentserver.core": mock_module}), + ): + _detect_hosted_environment() + assert _HOSTED_USER_AGENT_PREFIX in _telemetry_mod._user_agent_prefixes + _telemetry_mod._user_agent_prefixes.clear() + + +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() + mock_module = _mock_agent_config(is_hosted=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), + patch.dict("sys.modules", {"azure.ai.agentserver.core": mock_module}), + ): + _detect_hosted_environment() + assert _HOSTED_USER_AGENT_PREFIX not in _telemetry_mod._user_agent_prefixes + _telemetry_mod._user_agent_prefixes.clear() + + +def test_detect_hosted_fallback_import_error(): + """Test that ImportError from AgentConfig is silently handled.""" + _telemetry_mod._user_agent_prefixes.clear() + 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() 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 @@ 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 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..605f279411 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 @@ 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( + async 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.