Python: Show FoundryAgent client spans in Foundry traces

python/packages/core/agent_framework/observability.py

@@ -124,6 +124,9 @@
 )
 INNER_RESPONSE_ID_CAPTURED_FIELD: Final[str] = "response_id"
 INNER_USAGE_CAPTURED_FIELD: Final[str] = "usage"
+INNER_CAPTURED_RESPONSE_ID: Final[contextvars.ContextVar[str | None]] = contextvars.ContextVar(
+    "inner_captured_response_id", default=None
+)
 
 # Tracks accumulated token usage from all inner chat completion spans within an agent invoke.
 INNER_ACCUMULATED_USAGE: Final[contextvars.ContextVar[UsageDetails | None]] = contextvars.ContextVar(
@@ -2290,6 +2293,14 @@ def __init__(
         self.token_usage_histogram = _get_token_usage_histogram()
         self.duration_histogram = _get_duration_histogram()
 
+    def _get_additional_otel_agent_attributes(self) -> Mapping[str, Any]:
+        """Return provider-specific attributes emitted on agent spans."""
+        return {}
+
+    def _should_capture_agent_response_id(self) -> bool:
+        """Return whether the agent span must retain an inner response ID."""
+        return False
+
     def _trace_agent_invocation(
         self,
         *,
@@ -2333,6 +2344,7 @@ def _trace_agent_invocation(
             all_options=dict(merged_options),
             **merged_client_kwargs,
         )
+        attributes.update(self._get_additional_otel_agent_attributes())
 
         if stream:
             # Do NOT set the inner-telemetry context vars here: this synchronous run() body executes
@@ -2344,6 +2356,7 @@ def _trace_agent_invocation(
             # below), so set and reset both happen in the consumer's context.
             inner_response_telemetry_captured_fields: set[str] = set()
             inner_response_telemetry_captured_fields_token: contextvars.Token[set[str] | None] | None = None
+            inner_captured_response_id_token: contextvars.Token[str | None] | None = None
             inner_accumulated_usage_token: contextvars.Token[UsageDetails | None] | None = None
             # Agent Framework's agents run in-process (the actual network call happens on a nested
             # chat span), so invoke_agent spans use the default INTERNAL kind.
@@ -2411,10 +2424,14 @@ async def _finalize_stream() -> None:
                     response_attributes = _get_response_attributes(
                         attributes,
                         response,
-                        capture_response_id=INNER_RESPONSE_ID_CAPTURED_FIELD
-                        not in inner_response_telemetry_captured_fields,
+                        capture_response_id=(
+                            self._should_capture_agent_response_id()
+                            or INNER_RESPONSE_ID_CAPTURED_FIELD not in inner_response_telemetry_captured_fields
+                        ),
                         capture_usage=INNER_USAGE_CAPTURED_FIELD not in inner_response_telemetry_captured_fields,
                     )
+                    if self._should_capture_agent_response_id():
+                        _apply_captured_response_id(response_attributes)
                     _apply_accumulated_usage(response_attributes, inner_response_telemetry_captured_fields)
                     _capture_response(span=span, attributes=response_attributes, duration=duration)
                     if (
@@ -2436,6 +2453,8 @@ async def _finalize_stream() -> None:
                     # pull-context factory below set the tokens in — so the reset is cross-context safe.
                     if inner_response_telemetry_captured_fields_token is not None:
                         INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.reset(inner_response_telemetry_captured_fields_token)
+                    if inner_captured_response_id_token is not None:
+                        INNER_CAPTURED_RESPONSE_ID.reset(inner_captured_response_id_token)
                     if inner_accumulated_usage_token is not None:
                         INNER_ACCUMULATED_USAGE.reset(inner_accumulated_usage_token)
                     _close_span()
@@ -2447,11 +2466,14 @@ def _inner_telemetry_pull_context() -> contextlib.AbstractContextManager[Any]:
                 # avoiding the cross-context Token reset failure. Setting happens before the
                 # underlying iterator is pulled, so inner chat completion spans created during the
                 # pull can still accumulate usage / mark captured fields.
-                nonlocal inner_response_telemetry_captured_fields_token, inner_accumulated_usage_token
+                nonlocal inner_response_telemetry_captured_fields_token
+                nonlocal inner_captured_response_id_token
+                nonlocal inner_accumulated_usage_token
                 if inner_response_telemetry_captured_fields_token is None:
                     inner_response_telemetry_captured_fields_token = INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.set(
                         inner_response_telemetry_captured_fields
                     )
+                    inner_captured_response_id_token = INNER_CAPTURED_RESPONSE_ID.set(None)
                     inner_accumulated_usage_token = INNER_ACCUMULATED_USAGE.set({})
                 return _activate_span(span)
 
@@ -2483,6 +2505,7 @@ async def _run() -> AgentResponse[Any]:
             inner_response_telemetry_captured_fields_token = INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.set(
                 inner_response_telemetry_captured_fields
             )
+            inner_captured_response_id_token = INNER_CAPTURED_RESPONSE_ID.set(None)
             inner_accumulated_usage_token = INNER_ACCUMULATED_USAGE.set({})
             try:
                 with _get_span(attributes=attributes, span_name_attribute=OtelAttr.AGENT_NAME) as span:
@@ -2500,12 +2523,16 @@ async def _run() -> AgentResponse[Any]:
                             response_attributes = _get_response_attributes(
                                 attributes,
                                 response,
-                                capture_response_id=INNER_RESPONSE_ID_CAPTURED_FIELD
-                                not in inner_response_telemetry_captured_fields,
+                                capture_response_id=(
+                                    self._should_capture_agent_response_id()
+                                    or INNER_RESPONSE_ID_CAPTURED_FIELD not in inner_response_telemetry_captured_fields
+                                ),
                                 capture_usage=(
                                     INNER_USAGE_CAPTURED_FIELD not in inner_response_telemetry_captured_fields
                                 ),
                             )
+                            if self._should_capture_agent_response_id():
+                                _apply_captured_response_id(response_attributes)
                             _apply_accumulated_usage(
                                 response_attributes,
                                 inner_response_telemetry_captured_fields,
@@ -2531,6 +2558,7 @@ async def _run() -> AgentResponse[Any]:
                         raise
             finally:
                 INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.reset(inner_response_telemetry_captured_fields_token)
+                INNER_CAPTURED_RESPONSE_ID.reset(inner_captured_response_id_token)
                 INNER_ACCUMULATED_USAGE.reset(inner_accumulated_usage_token)
 
         return _run()
@@ -3439,6 +3467,7 @@ def _mark_inner_response_telemetry_captured(
         return
     if response.response_id:
         captured_fields.add(INNER_RESPONSE_ID_CAPTURED_FIELD)
+        INNER_CAPTURED_RESPONSE_ID.set(response.response_id)
     if response.usage_details:
         captured_fields.add(INNER_USAGE_CAPTURED_FIELD)
         accumulated = INNER_ACCUMULATED_USAGE.get()
@@ -3448,6 +3477,12 @@ def _mark_inner_response_telemetry_captured(
             INNER_ACCUMULATED_USAGE.set(add_usage_details(accumulated, response.usage_details))
 
 
+def _apply_captured_response_id(attributes: dict[str, Any]) -> None:
+    """Apply the inner chat response ID to an agent span when the provider requires it."""
+    if response_id := INNER_CAPTURED_RESPONSE_ID.get():
+        attributes.setdefault(OtelAttr.RESPONSE_ID, response_id)
+
+
 def _apply_accumulated_usage(attributes: dict[str, Any], captured_fields: set[str]) -> None:
     """Apply accumulated usage from inner chat spans to the invoke_agent span attributes."""
     if INNER_USAGE_CAPTURED_FIELD not in captured_fields:

python/packages/core/tests/core/test_observability.py

@@ -4847,10 +4847,13 @@ async def _get() -> ChatResponse:
 
 
 @pytest.mark.parametrize("stream", [False, True])
-async def test_agent_and_chat_spans_do_not_duplicate_response_telemetry(
-    span_exporter: InMemorySpanExporter, stream: bool
+@pytest.mark.parametrize("retain_agent_response_id", [False, True])
+async def test_agent_provider_hooks_control_response_telemetry(
+    span_exporter: InMemorySpanExporter,
+    stream: bool,
+    retain_agent_response_id: bool,
 ):
-    """The inner chat span owns response-id; usage is aggregated on the agent span."""
+    """Provider hooks can add root attributes and retain an inner response ID when required."""
 
     class NestedTelemetryChatClient(ChatTelemetryLayer, BaseChatClient[Any]):
         def service_url(self):
@@ -4890,7 +4893,14 @@ async def _get() -> ChatResponse:
 
             return _get()
 
-    agent = Agent(
+    class ProviderTelemetryAgent(Agent):
+        def _get_additional_otel_agent_attributes(self) -> Mapping[str, Any]:
+            return {"test.provider.attribute": "provider-value"}
+
+        def _should_capture_agent_response_id(self) -> bool:
+            return retain_agent_response_id
+
+    agent = ProviderTelemetryAgent(
         client=NestedTelemetryChatClient(),  # ty: ignore[invalid-argument-type]
         id="nested_agent_id",
         name="nested_agent",
@@ -4921,7 +4931,11 @@ async def _get() -> ChatResponse:
     assert chat_span.attributes[OtelAttr.INPUT_TOKENS] == 11  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[not-subscriptable]
     assert chat_span.attributes[OtelAttr.OUTPUT_TOKENS] == 22  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[not-subscriptable]
 
-    assert OtelAttr.RESPONSE_ID not in agent_span.attributes  # type: ignore[operator]  # pyrefly: ignore[not-iterable]  # ty: ignore[unsupported-operator]
+    assert agent_span.attributes["test.provider.attribute"] == "provider-value"  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[not-subscriptable]
+    if retain_agent_response_id:
+        assert agent_span.attributes[OtelAttr.RESPONSE_ID] == "nested_resp_123"  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[not-subscriptable]
+    else:
+        assert OtelAttr.RESPONSE_ID not in agent_span.attributes  # type: ignore[operator]  # pyrefly: ignore[not-iterable]  # ty: ignore[unsupported-operator]
     # The agent span carries the aggregated usage from all inner chat completions
     assert agent_span.attributes[OtelAttr.INPUT_TOKENS] == 11  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[not-subscriptable]
     assert agent_span.attributes[OtelAttr.OUTPUT_TOKENS] == 22  # type: ignore[index]  # pyrefly: ignore[unsupported-operation]  # ty: ignore[not-subscriptable]

python/packages/foundry/agent_framework_foundry/_agent.py

@@ -36,6 +36,7 @@
 from agent_framework.observability import AgentTelemetryLayer, ChatTelemetryLayer
 from agent_framework_openai._chat_client import OpenAIChatOptions, RawOpenAIChatClient
 from azure.ai.projects.aio import AIProjectClient
+from azure.ai.projects.models import ConnectionType
 from azure.core.credentials import TokenCredential
 from azure.core.credentials_async import AsyncTokenCredential
 
@@ -95,6 +96,7 @@ class FoundryAgentSettings(TypedDict, total=False):
 
 
 FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY = "foundry_hosted_agent_session_id"
+_FOUNDRY_PROJECT_ARM_ID_ATTRIBUTE = "microsoft.foundry.project.id"
 
 
 class FoundryAgentOptions(OpenAIChatOptions, total=False):
@@ -750,6 +752,7 @@ def __init__(
             client_kwargs["function_invocation_configuration"] = function_invocation_configuration
 
         client = actual_client_type(**client_kwargs)
+        self._foundry_project_arm_id: str | None = None
 
         super().__init__(
             client=client,  # type: ignore[arg-type]
@@ -841,6 +844,22 @@ def _update_session_from_chat_response_update(
         if session is not None and isinstance(agent_session_id, str) and agent_session_id:
             session.state[FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY] = agent_session_id
 
+    async def _get_foundry_project_arm_id(self) -> str:
+        """Get the Foundry project ARM ID from its Application Insights connection."""
+        client = cast(RawFoundryAgentChatClient, self.client)
+        # AIProjectClient does not expose the project ARM ID directly. Derive it from the
+        # project-scoped connection until https://github.com/Azure/azure-sdk-for-python/issues/48825 is addressed.
+        connections = client.project_client.connections.list(connection_type=ConnectionType.APPLICATION_INSIGHTS)
+        async for connection in connections:
+            connection_suffix = f"/connections/{connection.name}"
+            if not connection.id.lower().endswith(connection_suffix.lower()):
+                raise ValueError(
+                    f"The Foundry Application Insights connection ID has an unexpected format: {connection.id!r}."
+                )
+            return connection.id[: -len(connection_suffix)]
+
+        raise ValueError("The Foundry project does not have an Application Insights connection.")
+
     async def configure_azure_monitor(
         self,
         enable_sensitive_data: bool = False,
@@ -858,6 +877,8 @@ async def configure_azure_monitor(
 
         Raises:
             ImportError: If azure-monitor-opentelemetry-exporter is not installed.
+            ValueError: If the Application Insights connection does not contain the expected
+                project-scoped ARM resource ID.
         """
         from agent_framework.observability import (
             OBSERVABILITY_SETTINGS,
@@ -897,6 +918,8 @@ async def configure_azure_monitor(
                 "Install it with: pip install azure-monitor-opentelemetry"
             ) from exc
 
+        self._foundry_project_arm_id = await self._get_foundry_project_arm_id()
+
         if "resource" not in kwargs:
             kwargs["resource"] = create_resource()
 
@@ -951,6 +974,18 @@ class FoundryAgent(  # type: ignore[misc]
             )
     """
 
+    @override
+    def _get_additional_otel_agent_attributes(self) -> Mapping[str, Any]:
+        """Return Foundry attributes required to discover the agent trace."""
+        if self._foundry_project_arm_id:
+            return {_FOUNDRY_PROJECT_ARM_ID_ATTRIBUTE: self._foundry_project_arm_id}
+        return {}
+
+    @override
+    def _should_capture_agent_response_id(self) -> bool:
+        """Keep the response ID on the client agent span for Foundry trace discovery."""
+        return True
+
     def __init__(
         self,
         *,

python/packages/foundry/tests/foundry/test_foundry_agent.py

@@ -947,6 +947,43 @@ def test_raw_foundry_agent_init_creates_client() -> None:
 
     assert agent.client is not None
     assert cast(Any, agent.client).agent_name == "test-agent"
+    assert agent.name == "test-agent"
+
+
+async def test_get_foundry_project_arm_id_from_application_insights_connection() -> None:
+    """Test that project attribution uses the public project-scoped connection ID."""
+
+    project_arm_id = (
+        "/subscriptions/test-sub/resourceGroups/test-rg/providers/"
+        "Microsoft.CognitiveServices/accounts/test-account/projects/test-project"
+    )
+    project_client = MagicMock()
+
+    async def connections():
+        yield SimpleNamespace(id=f"{project_arm_id}/connections/appinsights", name="appinsights")
+
+    project_client.connections.list.return_value = connections()
+    agent = RawFoundryAgent(project_client=project_client, agent_name="test-agent")
+
+    assert await agent._get_foundry_project_arm_id() == project_arm_id
+    project_client.connections.list.assert_called_once_with(
+        connection_type=projects_models.ConnectionType.APPLICATION_INSIGHTS
+    )
+
+
+async def test_get_foundry_project_arm_id_rejects_unexpected_connection_id() -> None:
+    """Test that malformed connection metadata does not silently disable portal attribution."""
+
+    project_client = MagicMock()
+
+    async def connections():
+        yield SimpleNamespace(id="appinsights", name="appinsights")
+
+    project_client.connections.list.return_value = connections()
+    agent = RawFoundryAgent(project_client=project_client, agent_name="test-agent")
+
+    with pytest.raises(ValueError, match="unexpected format"):
+        await agent._get_foundry_project_arm_id()
 
 
 def test_raw_foundry_agent_init_passes_default_headers_to_client() -> None:
@@ -1243,6 +1280,8 @@ def test_foundry_agent_init() -> None:
 
     assert agent.client is not None
     assert cast(Any, agent.client).agent_name == "test-agent"
+    assert agent.name == "test-agent"
+    assert agent._should_capture_agent_response_id()
 
 
 def test_foundry_agent_init_with_middleware() -> None:
@@ -1278,6 +1317,10 @@ async def test_foundry_agent_configure_azure_monitor() -> None:
     mock_views = MagicMock(return_value=[])
     mock_resource = MagicMock()
     mock_enable = MagicMock()
+    project_arm_id = (
+        "/subscriptions/test-sub/resourceGroups/test-rg/providers/"
+        "Microsoft.CognitiveServices/accounts/test-account/projects/test-project"
+    )
 
     with (
         patch.dict(
@@ -1287,15 +1330,24 @@ async def test_foundry_agent_configure_azure_monitor() -> None:
         patch("agent_framework.observability.create_metric_views", mock_views),
         patch("agent_framework.observability.create_resource", return_value=mock_resource),
         patch("agent_framework.observability.enable_instrumentation", mock_enable),
+        patch(
+            "agent_framework_foundry._agent.RawFoundryAgent._get_foundry_project_arm_id",
+            new_callable=AsyncMock,
+            return_value=project_arm_id,
+        ) as mock_get_project_arm_id,
     ):
         await agent.configure_azure_monitor(enable_sensitive_data=True)
 
     mock_project.telemetry.get_application_insights_connection_string.assert_called_once()
+    mock_get_project_arm_id.assert_awaited_once_with()
     call_kwargs = mock_configure.call_args.kwargs
     assert call_kwargs["connection_string"] == "InstrumentationKey=test-key;IngestionEndpoint=https://test.endpoint"
     assert call_kwargs["views"] == []
     assert call_kwargs["resource"] is mock_resource
     mock_enable.assert_called_once_with(enable_sensitive_data=True)
+    assert agent._get_additional_otel_agent_attributes() == {
+        "microsoft.foundry.project.id": project_arm_id,
+    }
 
 
 async def test_foundry_agent_configure_azure_monitor_resource_not_found() -> None: