From fd36871d6076f7f827f41acba6a3eba666129730 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Mon, 20 Apr 2026 18:34:30 -0700 Subject: [PATCH] User agent scoped --- .../core/agent_framework/_telemetry.py | 32 +++++++---- .../core/tests/core/test_telemetry.py | 54 +++++++++++++++++++ .../_invocations.py | 12 +++-- .../_responses.py | 16 ++++-- 4 files changed, 95 insertions(+), 19 deletions(-) diff --git a/python/packages/core/agent_framework/_telemetry.py b/python/packages/core/agent_framework/_telemetry.py index 8e59704b8a..6351b30bbb 100644 --- a/python/packages/core/agent_framework/_telemetry.py +++ b/python/packages/core/agent_framework/_telemetry.py @@ -4,6 +4,9 @@ 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 @@ -26,27 +29,34 @@ 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: list[str] = [] +_user_agent_prefixes: ContextVar[tuple[str, ...]] = ContextVar("_user_agent_prefixes", default=()) -def append_to_user_agent(prefix: str) -> None: - """Prepend a prefix to the agent framework user agent string. +@contextmanager +def user_agent_prefix(prefix: str) -> Generator[None, None, None]: + """Context manager that adds a prefix to the user agent string for the current scope. - This is useful for hosting layers that want to identify themselves in telemetry. - Duplicate prefixes are ignored. + This is useful for upstream layers that want to identify themselves in telemetry + for the duration of a request without permanently mutating global state. Args: - prefix: The prefix to prepend (e.g. "foundry-hosting"). + prefix: The prefix to add (e.g. "foundry-hosting"). """ - if prefix and prefix not in _user_agent_prefixes: - _user_agent_prefixes.append(prefix) + current = _user_agent_prefixes.get() + token = _user_agent_prefixes.set((*current, prefix)) if prefix and prefix not in current else None + try: + yield + finally: + if token is not None: + _user_agent_prefixes.reset(token) def _get_user_agent() -> str: - """Return the full user agent string including any prepended prefixes.""" - if not _user_agent_prefixes: + """Return the full user agent string including any context-scoped prefixes.""" + prefixes = _user_agent_prefixes.get() + if not prefixes: return AGENT_FRAMEWORK_USER_AGENT - return f"{'/'.join(_user_agent_prefixes)}/{AGENT_FRAMEWORK_USER_AGENT}" + return f"{'/'.join(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 f2d9392b2d..1ba1df1dde 100644 --- a/python/packages/core/tests/core/test_telemetry.py +++ b/python/packages/core/tests/core/test_telemetry.py @@ -8,6 +8,7 @@ from agent_framework import ( USER_AGENT_TELEMETRY_DISABLED_ENV_VAR, prepend_agent_framework_to_user_agent, ) +from agent_framework._telemetry import user_agent_prefix # region Test constants @@ -96,3 +97,56 @@ def test_modifies_original_dict(): assert result is headers # Same object assert "User-Agent" in headers + + +# region Test user_agent_prefix context manager + + +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 + result = prepend_agent_framework_to_user_agent() + assert result["User-Agent"] == AGENT_FRAMEWORK_USER_AGENT + + +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_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 + result = prepend_agent_framework_to_user_agent() + assert result["User-Agent"] == AGENT_FRAMEWORK_USER_AGENT + + +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 + result = prepend_agent_framework_to_user_agent() + assert result["User-Agent"] == AGENT_FRAMEWORK_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 df5b8c5293..b5bf98291b 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py @@ -1,11 +1,11 @@ # Copyright (c) Microsoft. All rights reserved. from agent_framework import AgentSession, BaseAgent, SupportsAgentRun -from agent_framework._telemetry import append_to_user_agent +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 -from typing_extensions import Any, AsyncGenerator, Optional +from typing_extensions import Any, AsyncGenerator class InvocationsHostServer(InvocationAgentServerHost): @@ -17,7 +17,7 @@ class InvocationsHostServer(InvocationAgentServerHost): self, agent: BaseAgent, *, - openapi_spec: Optional[dict[str, Any]] = None, + openapi_spec: dict[str, Any] | None = None, **kwargs: Any, ) -> None: """Initialize an InvocationsHostServer. @@ -36,13 +36,17 @@ class InvocationsHostServer(InvocationAgentServerHost): if not isinstance(agent, SupportsAgentRun): raise TypeError("Agent must support the SupportsAgentRun interface") - append_to_user_agent(self.USER_AGENT_PREFIX) self._agent = agent self._sessions: dict[str, AgentSession] = {} self.invoke_handler(self._handle_invoke) # pyright: ignore[reportUnknownMemberType] 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 936c05858e..45c67b8137 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -19,7 +19,7 @@ from agent_framework import ( SupportsAgentRun, WorkflowAgent, ) -from agent_framework._telemetry import append_to_user_agent +from agent_framework._telemetry import user_agent_prefix from azure.ai.agentserver.responses import ( ResponseContext, ResponseEventStream, @@ -151,9 +151,6 @@ class ResponsesHostServer(ResponsesAgentServerHost): self._agent = agent self.response_handler(self._handler) # pyright: ignore[reportUnknownMemberType] - # Append the user agent prefix for telemetry purposes - append_to_user_agent(self.USER_AGENT_PREFIX) - @staticmethod def _is_streaming_request(request: CreateResponse) -> bool: """Check if the request is a streaming request.""" @@ -166,6 +163,17 @@ class ResponsesHostServer(ResponsesAgentServerHost): 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 + + async def _handle_inner( + 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):