User agent scoped

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