Python: Capture context provider instructions in agent telemetry (microsoft#6515)

* Fix agent instructions telemetry

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Simplify agent instructions telemetry guard

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix observability mypy cast

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: eavanvalkenburg

python/packages/core/agent_framework/observability.py

@@ -1488,14 +1488,21 @@ def get_response(
         )
 
         if stream:
+            agent_span = trace.get_current_span()
             span = _start_streaming_span(attributes, OtelAttr.REQUEST_MODEL)
 
             if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and messages and span.is_recording():
+                system_instructions = _get_instructions_from_options(opts)
+                _capture_current_agent_system_instructions(
+                    agent_span,
+                    span,
+                    system_instructions,
+                )
                 _capture_messages(
                     span=span,
                     provider_name=provider_name,
                     messages=messages,
-                    system_instructions=opts.get("instructions"),
+                    system_instructions=system_instructions,
                 )
 
             span_state = {"closed": False}
@@ -1585,13 +1592,20 @@ async def _finalize_stream() -> None:
             return wrapped_stream
 
         async def _get_response() -> ChatResponse:
+            agent_span = trace.get_current_span()
             with _get_span(attributes=attributes, span_name_attribute=OtelAttr.REQUEST_MODEL) as span:
                 if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and messages and span.is_recording():
+                    system_instructions = _get_instructions_from_options(opts)
+                    _capture_current_agent_system_instructions(
+                        agent_span,
+                        span,
+                        system_instructions,
+                    )
                     _capture_messages(
                         span=span,
                         provider_name=provider_name,
                         messages=messages,
-                        system_instructions=opts.get("instructions"),
+                        system_instructions=system_instructions,
                     )
                 start_time_stamp = perf_counter()
                 try:
@@ -1761,7 +1775,6 @@ def _trace_agent_invocation(
             inner_response_telemetry_captured_fields
         )
         inner_accumulated_usage_token = INNER_ACCUMULATED_USAGE.set({})
-
         if stream:
             span = _start_streaming_span(attributes, OtelAttr.AGENT_NAME)
 
@@ -1859,38 +1872,44 @@ async def _finalize_stream() -> None:
         async def _run() -> AgentResponse[Any]:
             try:
                 with _get_span(attributes=attributes, span_name_attribute=OtelAttr.AGENT_NAME) as span:
-                    if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and messages and span.is_recording():
-                        _capture_messages(
-                            span=span,
-                            provider_name=provider_name,
-                            messages=messages,
-                            system_instructions=_get_instructions_from_options(dict(merged_options)),
-                        )
-                    start_time_stamp = perf_counter()
                     try:
-                        response: AgentResponse[Any] = await execute()
-                    except Exception as exception:
-                        capture_exception(span=span, exception=exception, timestamp=time_ns())
-                        raise
-                    duration = perf_counter() - start_time_stamp
-                    if response:
-                        response_attributes = _get_response_attributes(
-                            attributes,
-                            response,
-                            capture_response_id=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,
-                        )
-                        _apply_accumulated_usage(response_attributes, inner_response_telemetry_captured_fields)
-                        _capture_response(span=span, attributes=response_attributes, duration=duration)
-                        if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and response.messages and span.is_recording():
+                        if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED and messages and span.is_recording():
                             _capture_messages(
                                 span=span,
                                 provider_name=provider_name,
-                                messages=response.messages,
-                                output=True,
+                                messages=messages,
+                                system_instructions=_get_instructions_from_options(dict(merged_options)),
+                            )
+                        start_time_stamp = perf_counter()
+                        response: AgentResponse[Any] = await execute()
+                        duration = perf_counter() - start_time_stamp
+                        if response:
+                            response_attributes = _get_response_attributes(
+                                attributes,
+                                response,
+                                capture_response_id=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
+                                ),
                             )
-                    return response  # type: ignore[return-value,no-any-return]
+                            _apply_accumulated_usage(response_attributes, inner_response_telemetry_captured_fields)
+                            _capture_response(span=span, attributes=response_attributes, duration=duration)
+                            if (
+                                OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED
+                                and response.messages
+                                and span.is_recording()
+                            ):
+                                _capture_messages(
+                                    span=span,
+                                    provider_name=provider_name,
+                                    messages=response.messages,
+                                    output=True,
+                                )
+                        return response  # type: ignore[return-value,no-any-return]
+                    except Exception as exception:
+                        capture_exception(span=span, exception=exception, timestamp=time_ns())
+                        raise
             finally:
                 INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.reset(inner_response_telemetry_captured_fields_token)
                 INNER_ACCUMULATED_USAGE.reset(inner_accumulated_usage_token)
@@ -2263,6 +2282,83 @@ def capture_exception(span: trace.Span, exception: Exception, timestamp: int | N
     span.set_status(status=trace.StatusCode.ERROR, description=repr(exception))
 
 
+def _capture_system_instructions(span: trace.Span, system_instructions: str | list[str] | None) -> None:
+    """Capture system instructions on a span."""
+    if not system_instructions:
+        return
+    otel_sys_instructions = [
+        {"type": "text", "content": instruction} for instruction in _normalize_instructions(system_instructions)
+    ]
+    span.set_attribute(OtelAttr.SYSTEM_INSTRUCTIONS, json.dumps(otel_sys_instructions, ensure_ascii=False))
+
+
+def _capture_current_agent_system_instructions(
+    agent_span: trace.Span,
+    chat_span: trace.Span,
+    system_instructions: str | list[str] | None,
+) -> None:
+    """Capture final chat instructions on the current agent span when the chat span belongs to it."""
+    if not system_instructions or not agent_span.is_recording():
+        return
+
+    agent_attributes_obj = getattr(agent_span, "attributes", None)
+    if not isinstance(agent_attributes_obj, Mapping):
+        return
+    agent_attributes = cast(Mapping[str, Any], agent_attributes_obj)
+    if agent_attributes.get(OtelAttr.OPERATION.value) != OtelAttr.AGENT_INVOKE_OPERATION:
+        return
+
+    if not _instructions_preserve_existing_agent_instructions(agent_attributes, system_instructions):
+        return
+
+    chat_parent = getattr(chat_span, "parent", None)
+    agent_context = agent_span.get_span_context()
+    if (
+        chat_parent is None
+        or chat_parent.span_id != agent_context.span_id
+        or chat_parent.trace_id != agent_context.trace_id
+    ):
+        return
+
+    _capture_system_instructions(agent_span, system_instructions)
+
+
+def _normalize_instructions(system_instructions: str | list[str]) -> list[str]:
+    """Normalize system instructions to telemetry text items."""
+    return system_instructions if isinstance(system_instructions, list) else [system_instructions]
+
+
+def _instructions_preserve_existing_agent_instructions(
+    agent_attributes: Mapping[str, Any],
+    system_instructions: str | list[str],
+) -> bool:
+    """Return True when chat instructions preserve the agent span's existing instructions."""
+    existing = agent_attributes.get(OtelAttr.SYSTEM_INSTRUCTIONS)
+    if not isinstance(existing, str):
+        return True
+
+    try:
+        existing_items_obj = json.loads(existing)
+    except json.JSONDecodeError:
+        return False
+
+    if not isinstance(existing_items_obj, list):
+        return False
+    existing_items = cast(list[object], existing_items_obj)
+
+    existing_contents: list[str] = []
+    for item in existing_items:
+        if not isinstance(item, Mapping):
+            continue
+        content = cast(Mapping[str, Any], item).get("content")
+        if isinstance(content, str):
+            existing_contents.append(content)
+
+    existing_text = "\n".join(existing_contents)
+    new_text = "\n".join(_normalize_instructions(system_instructions))
+    return new_text == existing_text or new_text.startswith(f"{existing_text}\n")
+
+
 def _capture_messages(
     span: trace.Span,
     provider_name: str,
@@ -2294,11 +2390,7 @@ def _capture_messages(
     span.set_attribute(
         OtelAttr.OUTPUT_MESSAGES if output else OtelAttr.INPUT_MESSAGES, json.dumps(otel_messages, ensure_ascii=False)
     )
-    if system_instructions:
-        if not isinstance(system_instructions, list):
-            system_instructions = [system_instructions]
-        otel_sys_instructions = [{"type": "text", "content": instruction} for instruction in system_instructions]
-        span.set_attribute(OtelAttr.SYSTEM_INSTRUCTIONS, json.dumps(otel_sys_instructions, ensure_ascii=False))
+    _capture_system_instructions(span, system_instructions)
 
 
 def _to_otel_message(message: Message) -> dict[str, Any]:

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

@@ -17,6 +17,7 @@
     ChatResponse,
     ChatResponseUpdate,
     Content,
+    ContextProvider,
     Message,
     RawAgent,
     ResponseStream,
@@ -3647,6 +3648,124 @@ async def test_agent_streaming_instructions_merged_from_default_and_options(
     assert "Stream override." in system_instructions[0]["content"]
 
 
+@pytest.mark.parametrize("enable_sensitive_data", [True], indirect=True)
+@pytest.mark.parametrize("stream", [False, True])
+async def test_agent_instructions_include_context_provider_extensions(
+    mock_chat_client,
+    span_exporter: InMemorySpanExporter,
+    enable_sensitive_data,
+    stream: bool,
+) -> None:
+    """Agent span instructions include instructions added by context providers."""
+    import json
+
+    class UserMemoryProvider(ContextProvider):
+        def __init__(self) -> None:
+            super().__init__(source_id="user-memory")
+
+        async def before_run(
+            self,
+            *,
+            agent: Any,
+            session: Any,
+            context: Any,
+            state: dict[str, Any],
+        ) -> None:
+            context.extend_instructions(self.source_id, "The user's name is Alice.")
+
+    agent = Agent(
+        client=mock_chat_client(),
+        name="memory_agent",
+        instructions="You are a friendly assistant.",
+        context_providers=[UserMemoryProvider()],
+    )
+
+    span_exporter.clear()
+    if stream:
+        result_stream = agent.run("Hello", stream=True)
+        async for _ in result_stream:
+            pass
+        await result_stream.get_final_response()
+    else:
+        await agent.run("Hello")
+
+    spans = span_exporter.get_finished_spans()
+    agent_spans = [
+        span for span in spans if span.attributes.get(OtelAttr.OPERATION.value) == OtelAttr.AGENT_INVOKE_OPERATION
+    ]
+    assert len(agent_spans) == 1
+
+    system_instructions = json.loads(agent_spans[0].attributes[OtelAttr.SYSTEM_INSTRUCTIONS])
+    contents = [item["content"] for item in system_instructions]
+    assert any("You are a friendly assistant." in content for content in contents)
+    assert any("The user's name is Alice." in content for content in contents)
+
+
+@pytest.mark.parametrize("enable_sensitive_data", [True], indirect=True)
+async def test_agent_instructions_not_overwritten_by_unrelated_nested_chat(
+    mock_chat_client,
+    span_exporter: InMemorySpanExporter,
+    enable_sensitive_data,
+) -> None:
+    """Unrelated nested chat calls must not overwrite agent span instructions."""
+    import json
+
+    class NestedChatProvider(ContextProvider):
+        def __init__(self, nested_client: BaseChatClient[Any]) -> None:
+            super().__init__(source_id="nested-chat")
+            self.nested_client = nested_client
+
+        async def before_run(
+            self,
+            *,
+            agent: Any,
+            session: Any,
+            context: Any,
+            state: dict[str, Any],
+        ) -> None:
+            context.extend_instructions(self.source_id, "Context-provided instructions.")
+
+        async def after_run(
+            self,
+            *,
+            agent: Any,
+            session: Any,
+            context: Any,
+            state: dict[str, Any],
+        ) -> None:
+            await self.nested_client.get_response(
+                messages=[Message(role="user", contents=["Nested request"])],
+                options={"model": "NestedModel", "instructions": "Unrelated nested instructions."},
+                client_kwargs={"session": session},
+            )
+
+    agent = Agent(
+        client=mock_chat_client(),
+        name="guarded_agent",
+        instructions="Base agent instructions.",
+        context_providers=[NestedChatProvider(mock_chat_client())],
+    )
+
+    span_exporter.clear()
+    await agent.run("Hello")
+
+    spans = span_exporter.get_finished_spans()
+    agent_spans = [
+        span for span in spans if span.attributes.get(OtelAttr.OPERATION.value) == OtelAttr.AGENT_INVOKE_OPERATION
+    ]
+    assert len(agent_spans) == 1
+    chat_spans = [
+        span for span in spans if span.attributes.get(OtelAttr.OPERATION.value) == OtelAttr.CHAT_COMPLETION_OPERATION
+    ]
+    assert len(chat_spans) == 2
+
+    system_instructions = json.loads(agent_spans[0].attributes[OtelAttr.SYSTEM_INSTRUCTIONS])
+    contents = [item["content"] for item in system_instructions]
+    assert any("Base agent instructions." in content for content in contents)
+    assert any("Context-provided instructions." in content for content in contents)
+    assert all("Unrelated nested instructions." not in content for content in contents)
+
+
 @pytest.mark.parametrize("enable_sensitive_data", [True], indirect=True)
 async def test_agent_no_instructions_in_default_or_options(
     mock_chat_agent, span_exporter: InMemorySpanExporter, enable_sensitive_data