User agent scoped

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]:

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

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
 

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):