diff --git a/instrumentation-loongsuite/loongsuite-instrumentation-microsoft-agent-framework/CHANGELOG.md b/instrumentation-loongsuite/loongsuite-instrumentation-microsoft-agent-framework/CHANGELOG.md index f1820e343..e04229104 100644 --- a/instrumentation-loongsuite/loongsuite-instrumentation-microsoft-agent-framework/CHANGELOG.md +++ b/instrumentation-loongsuite/loongsuite-instrumentation-microsoft-agent-framework/CHANGELOG.md @@ -10,6 +10,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed - Avoid enriching non-MAF spans that use overlapping GenAI operation names. +- Preserve parent-child context for legacy Microsoft Agent Framework streaming + agent spans. ## Version 0.7.0 (2026-07-03) diff --git a/instrumentation-loongsuite/loongsuite-instrumentation-microsoft-agent-framework/src/opentelemetry/instrumentation/microsoft_agent_framework/util_genai_bridge.py b/instrumentation-loongsuite/loongsuite-instrumentation-microsoft-agent-framework/src/opentelemetry/instrumentation/microsoft_agent_framework/util_genai_bridge.py index c45493cde..c5c08e287 100644 --- a/instrumentation-loongsuite/loongsuite-instrumentation-microsoft-agent-framework/src/opentelemetry/instrumentation/microsoft_agent_framework/util_genai_bridge.py +++ b/instrumentation-loongsuite/loongsuite-instrumentation-microsoft-agent-framework/src/opentelemetry/instrumentation/microsoft_agent_framework/util_genai_bridge.py @@ -27,7 +27,10 @@ import inspect import json import logging +import sys import timeit +import weakref +from time import perf_counter, time_ns from typing import Any, AsyncGenerator, Callable, Generator, Mapping, Optional from opentelemetry import trace as otel_trace @@ -102,6 +105,8 @@ _original_create_mcp_client_span: Any = None _original_tools_get_function_span: Any = None _original_mcp_create_mcp_client_span: Any = None +_original_agent_trace_invocation: Any = None +_legacy_response_stream_originals: dict[str, Any] = {} _FINALIZED_ATTR = "_loongsuite_util_genai_finalized" _END_WRAPPED_ATTR = "_loongsuite_util_genai_end_wrapped" @@ -129,6 +134,7 @@ def apply_util_genai_bridge() -> None: global _original_get_span global _original_start_streaming_span global _original_tools_get_function_span + global _original_agent_trace_invocation if _applied: return @@ -157,6 +163,11 @@ def apply_util_genai_bridge() -> None: _original_create_mcp_client_span = getattr( observability, "create_mcp_client_span", None ) + _original_agent_trace_invocation = getattr( + getattr(observability, "AgentTelemetryLayer", None), + "_trace_agent_invocation", + None, + ) wrapped_get_span = ( _wrap_get_span(_original_get_span) @@ -192,6 +203,16 @@ def apply_util_genai_bridge() -> None: if wrapped_create_mcp_client_span is not None: observability.create_mcp_client_span = wrapped_create_mcp_client_span # type: ignore[attr-defined] + legacy_response_stream_patched = _patch_legacy_response_stream() + if legacy_response_stream_patched and _original_agent_trace_invocation: + agent_cls = getattr(observability, "AgentTelemetryLayer", None) + if agent_cls is not None: + agent_cls._trace_agent_invocation = ( + _wrap_legacy_agent_trace_invocation( # type: ignore[attr-defined] + _original_agent_trace_invocation + ) + ) + try: import agent_framework._tools as tools_mod # type: ignore @@ -228,6 +249,7 @@ def revert_util_genai_bridge() -> None: global _original_get_span global _original_start_streaming_span global _original_tools_get_function_span + global _original_agent_trace_invocation if not _applied: return @@ -248,6 +270,14 @@ def revert_util_genai_bridge() -> None: observability.create_mcp_client_span = ( _original_create_mcp_client_span # type: ignore[attr-defined] ) + agent_cls = getattr(observability, "AgentTelemetryLayer", None) + if ( + agent_cls is not None + and _original_agent_trace_invocation is not None + ): + agent_cls._trace_agent_invocation = ( # type: ignore[attr-defined] + _original_agent_trace_invocation + ) except ImportError: pass try: @@ -275,6 +305,301 @@ def revert_util_genai_bridge() -> None: _original_create_mcp_client_span = None _original_tools_get_function_span = None _original_mcp_create_mcp_client_span = None + _original_agent_trace_invocation = None + _restore_legacy_response_stream() + + +def _patch_legacy_response_stream() -> bool: + """Add per-pull context support to older MAF ``ResponseStream``. + + MAF 1.0 streaming spans were created detached and the stream type had no + pull-context hook, so child spans created while resolving/iterating the + stream became roots. MAF 1.10 added ``with_pull_context_manager``; this + compatibility shim backports only that context propagation surface. + """ + global _legacy_response_stream_originals + + if _legacy_response_stream_originals: + return True + try: + from agent_framework._types import ResponseStream # type: ignore + except ImportError: + return False + if hasattr(ResponseStream, "with_pull_context_manager"): + return False + + originals = { + "__init__": ResponseStream.__init__, + "_get_stream": ResponseStream._get_stream, + "__anext__": ResponseStream.__anext__, + "_run_cleanup_hooks": ResponseStream._run_cleanup_hooks, + } + _legacy_response_stream_originals = originals + + original_init = originals["__init__"] + original_get_stream = originals["_get_stream"] + original_anext = originals["__anext__"] + original_run_cleanup_hooks = originals["_run_cleanup_hooks"] + + def _init(self: Any, *args: Any, **kwargs: Any) -> None: + original_init(self, *args, **kwargs) + self._pull_context_manager_factories = [] + self._stream_error = None + + async def _get_stream(self: Any) -> Any: + if getattr(self, "_stream", None) is not None: + # The stream is already resolved; child spans are created while + # pulling updates, which is covered by the patched ``__anext__``. + return await original_get_stream(self) + with contextlib.ExitStack() as stack: + for factory in getattr( + self, "_pull_context_manager_factories", () + ): + stack.enter_context(factory()) + return await original_get_stream(self) + + async def _anext(self: Any) -> Any: + with contextlib.ExitStack() as stack: + for factory in getattr( + self, "_pull_context_manager_factories", () + ): + stack.enter_context(factory()) + return await original_anext(self) + + async def _run_cleanup_hooks(self: Any) -> Any: + _, exc, _ = sys.exc_info() + if exc is not None and not isinstance(exc, StopAsyncIteration): + self._stream_error = exc + try: + with contextlib.ExitStack() as stack: + for factory in getattr( + self, "_pull_context_manager_factories", () + ): + stack.enter_context(factory()) + return await original_run_cleanup_hooks(self) + finally: + if exc is not None: + self._stream_error = None + + def _with_pull_context_manager(self: Any, cm_factory: Callable[[], Any]): + self._pull_context_manager_factories.append(cm_factory) + return self + + ResponseStream.__init__ = _init # type: ignore[assignment] + ResponseStream._get_stream = _get_stream # type: ignore[assignment] + ResponseStream.__anext__ = _anext # type: ignore[assignment] + ResponseStream._run_cleanup_hooks = _run_cleanup_hooks # type: ignore[assignment] + ResponseStream.with_pull_context_manager = _with_pull_context_manager # type: ignore[attr-defined] + return True + + +def _restore_legacy_response_stream() -> None: + global _legacy_response_stream_originals + if not _legacy_response_stream_originals: + return + try: + from agent_framework._types import ResponseStream # type: ignore + except ImportError: + _legacy_response_stream_originals = {} + return + ResponseStream.__init__ = _legacy_response_stream_originals["__init__"] # type: ignore[assignment] + ResponseStream._get_stream = _legacy_response_stream_originals[ + "_get_stream" + ] # type: ignore[assignment] + ResponseStream.__anext__ = _legacy_response_stream_originals["__anext__"] # type: ignore[assignment] + ResponseStream._run_cleanup_hooks = _legacy_response_stream_originals[ # type: ignore[assignment] + "_run_cleanup_hooks" + ] + try: + delattr(ResponseStream, "with_pull_context_manager") + except AttributeError: + pass + _legacy_response_stream_originals = {} + + +def _wrap_legacy_agent_trace_invocation( + original: Callable[..., Any], +) -> Callable[..., Any]: + def _trace_agent_invocation(self: Any, *args: Any, **kwargs: Any) -> Any: + if args or not kwargs.get("stream"): + return original(self, *args, **kwargs) + return _legacy_streaming_agent_invocation(self, **kwargs) + + return _trace_agent_invocation + + +def _legacy_streaming_agent_invocation(self: Any, **kwargs: Any) -> Any: + import agent_framework.observability as observability # type: ignore + from agent_framework._types import ResponseStream # type: ignore + + execute = kwargs["execute"] + messages = kwargs.get("messages") + session = kwargs.get("session") + merged_options = kwargs.get("merged_options") or {} + client_kwargs = kwargs.get("client_kwargs") + if not observability.OBSERVABILITY_SETTINGS.ENABLED: + return execute() + + provider_name = str(getattr(self, "otel_provider_name", "unknown")) + merged_client_kwargs = ( + dict(client_kwargs) if client_kwargs is not None else {} + ) + OtelAttr = observability.OtelAttr + attributes = observability._get_span_attributes( + operation_name=OtelAttr.AGENT_INVOKE_OPERATION, + provider_name=provider_name, + agent_id=getattr(self, "id", "unknown"), + agent_name=getattr(self, "name", None) + or getattr(self, "id", "unknown"), + agent_description=getattr(self, "description", None), + thread_id=session.service_session_id if session else None, + all_options=dict(merged_options), + **merged_client_kwargs, + ) + + operation = attributes.get(OtelAttr.OPERATION, "operation") + span_name = attributes.get(OtelAttr.AGENT_NAME, "unknown") + span = observability.get_tracer().start_span( + f"{operation} {span_name}", + kind=otel_trace.SpanKind.INTERNAL, + attributes=attributes, + ) + _mark_maf_live_span(span) + _wrap_span_end(span) + + if ( + observability.OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED + and messages + and span.is_recording() + ): + observability._capture_messages( + span=span, + provider_name=provider_name, + messages=messages, + system_instructions=observability._get_instructions_from_options( + dict(merged_options) + ), + ) + + span_state = {"closed": False} + duration_state: dict[str, float] = {} + start_time = perf_counter() + inner_response_telemetry_captured_fields: set[str] = set() + inner_accumulated_usage: dict[str, Any] = {} + + def _close_span() -> None: + if span_state["closed"]: + return + span_state["closed"] = True + span.end() + + def _record_duration() -> None: + duration_state["duration"] = perf_counter() - start_time + + try: + with _activate_live_span(span): + run_result = execute() + if isinstance(run_result, ResponseStream): + result_stream = run_result + elif inspect.isawaitable(run_result): + result_stream = ResponseStream.from_awaitable(run_result) + else: + raise RuntimeError( + "Streaming telemetry requires a ResponseStream result." + ) + except Exception as exception: + observability.capture_exception( + span=span, exception=exception, timestamp=time_ns() + ) + _close_span() + raise + + async def _finalize_stream() -> None: + try: + stream_error = getattr(result_stream, "_stream_error", None) + if stream_error is not None: + observability.capture_exception( + span=span, exception=stream_error, timestamp=time_ns() + ) + return + response = await result_stream.get_final_response() + response_attributes = observability._get_response_attributes( + attributes, + response, + capture_response_id=( + observability.INNER_RESPONSE_ID_CAPTURED_FIELD + not in inner_response_telemetry_captured_fields + ), + capture_usage=( + observability.INNER_USAGE_CAPTURED_FIELD + not in inner_response_telemetry_captured_fields + ), + ) + observability._apply_accumulated_usage( + response_attributes, + inner_response_telemetry_captured_fields, + ) + observability._capture_response( + span=span, + attributes=response_attributes, + duration=duration_state.get("duration"), + ) + if ( + observability.OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED + and getattr(response, "messages", None) + and span.is_recording() + ): + observability._capture_messages( + span=span, + provider_name=provider_name, + messages=response.messages, + output=True, + ) + except Exception as exception: + observability.capture_exception( + span=span, exception=exception, timestamp=time_ns() + ) + finally: + _close_span() + + @contextlib.contextmanager + def _inner_telemetry_pull_context() -> Any: + fields_token = ( + observability.INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.set( + inner_response_telemetry_captured_fields + ) + ) + usage_token = observability.INNER_ACCUMULATED_USAGE.set( + inner_accumulated_usage + ) + try: + with _activate_live_span(span): + yield + finally: + observability.INNER_ACCUMULATED_USAGE.reset(usage_token) + observability.INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS.reset( + fields_token + ) + + wrapped_stream = ( + result_stream.with_cleanup_hook(_record_duration) + .with_cleanup_hook(_finalize_stream) + .with_pull_context_manager(_inner_telemetry_pull_context) + ) + try: + weakref.finalize(wrapped_stream, _close_span) + except TypeError: + logger.debug("MAF ResponseStream is not weak-referenceable") + return wrapped_stream + + +def _activate_live_span(span: OtelSpan) -> Any: + return otel_trace.use_span( + span=span, + end_on_exit=False, + record_exception=False, + set_status_on_exception=False, + ) def _wrap_get_span(original: Callable[..., Any]) -> Callable[..., Any]: diff --git a/instrumentation-loongsuite/loongsuite-instrumentation-microsoft-agent-framework/tests/test_util_genai_bridge.py b/instrumentation-loongsuite/loongsuite-instrumentation-microsoft-agent-framework/tests/test_util_genai_bridge.py index 1d18e3730..9891d4e96 100644 --- a/instrumentation-loongsuite/loongsuite-instrumentation-microsoft-agent-framework/tests/test_util_genai_bridge.py +++ b/instrumentation-loongsuite/loongsuite-instrumentation-microsoft-agent-framework/tests/test_util_genai_bridge.py @@ -583,3 +583,260 @@ async def _run(): assert ( getattr(span, util_genai_bridge._STREAM_FIRST_TOKEN_ATTR) is not None ) + + +def test_legacy_streaming_agent_context_parents_inner_spans(monkeypatch): + """MAF 1.0 created streaming agent spans after ``execute()``. + + The returned stream then produced inner chat/tool spans without the agent + span as current context. The legacy bridge patch must create the agent span + first and activate it while the stream resolves and pulls updates. + """ + + tp = TracerProvider() + exporter = InMemorySpanExporter() + tp.add_span_processor(SimpleSpanProcessor(exporter)) + tracer = tp.get_tracer("fake-legacy-maf") + + class _OtelAttr: + OPERATION = GEN_AI_OPERATION_NAME + AGENT_NAME = "gen_ai.agent.name" + AGENT_ID = "gen_ai.agent.id" + AGENT_INVOKE_OPERATION = GenAIOperation.INVOKE_AGENT + + class _Settings: + ENABLED = True + SENSITIVE_DATA_ENABLED = True + + class _ResponseStream: + def __init__(self, stream, *, finalizer=None): + self._stream_source = stream + self._stream = None + self._iterator = None + self._updates = [] + self._finalizer = finalizer + self._finalized = False + self._final_result = None + self._cleanup_hooks = [] + self._cleanup_run = False + self._transform_hooks = [] + self._result_hooks = [] + self._map_update = None + self._consumed = False + self._wrap_inner = False + + @classmethod + def from_awaitable(cls, awaitable): + return cls(awaitable) + + async def _get_stream(self): + if self._stream is None: + if hasattr(self._stream_source, "__aiter__"): + self._stream = self._stream_source + else: + self._stream = await self._stream_source + return self._stream + + def __aiter__(self): + return self + + async def __anext__(self): + if self._iterator is None: + stream = await self._get_stream() + self._iterator = stream.__aiter__() + try: + update = await self._iterator.__anext__() + except StopAsyncIteration: + self._consumed = True + await self._run_cleanup_hooks() + await self.get_final_response() + raise + self._updates.append(update) + return update + + def __await__(self): + async def _wrap(): + await self._get_stream() + return self + + return _wrap().__await__() + + async def get_final_response(self): + if not self._finalized and not self._consumed: + async for _ in self: + pass + if not self._finalized: + if self._finalizer is None: + self._final_result = types.SimpleNamespace( + messages=[], usage_details=None + ) + else: + result = self._finalizer(self._updates) + if asyncio.iscoroutine(result): + result = await result + self._final_result = result + self._finalized = True + return self._final_result + + def with_cleanup_hook(self, hook): + self._cleanup_hooks.append(hook) + return self + + async def _run_cleanup_hooks(self): + if self._cleanup_run: + return + self._cleanup_run = True + for hook in self._cleanup_hooks: + result = hook() + if asyncio.iscoroutine(result): + await result + + @contextlib.contextmanager + def _get_span(attributes, span_name_attribute): + span = tracer.start_span( + f"{attributes[GEN_AI_OPERATION_NAME]} {attributes[span_name_attribute]}" + ) + span.set_attributes(attributes) + with trace.use_span(span, end_on_exit=True) as current_span: + yield current_span + + @contextlib.contextmanager + def get_function_span(attributes): + span = tracer.start_span( + f"{attributes[GEN_AI_OPERATION_NAME]} {attributes['gen_ai.tool.name']}" + ) + span.set_attributes(attributes) + with trace.use_span(span, end_on_exit=True) as current_span: + yield current_span + + class _AgentTelemetryLayer: + id = "agent-1" + name = "legacy-agent" + description = None + otel_provider_name = "microsoft.agent_framework" + + def _trace_agent_invocation(self, **kwargs): + return kwargs["execute"]() + + def _get_span_attributes(**kwargs): + attrs = { + GEN_AI_OPERATION_NAME: kwargs["operation_name"], + GEN_AI_PROVIDER_NAME: kwargs["provider_name"], + } + if kwargs.get("agent_name"): + attrs["gen_ai.agent.name"] = kwargs["agent_name"] + if kwargs.get("agent_id"): + attrs["gen_ai.agent.id"] = kwargs["agent_id"] + return attrs + + def _capture_response(span, attributes, **_kwargs): + span.set_attributes(attributes) + + def _capture_messages(span, messages, output=False, **_kwargs): + attr = "gen_ai.output.messages" if output else "gen_ai.input.messages" + role = "assistant" if output else "user" + span.set_attribute( + attr, + json.dumps( + [ + { + "role": role, + "parts": [{"type": "text", "content": str(message)}], + } + for message in messages + ] + ), + ) + + obs_mod = types.ModuleType("agent_framework.observability") + obs_mod.OtelAttr = _OtelAttr + obs_mod.OBSERVABILITY_SETTINGS = _Settings + obs_mod.INNER_RESPONSE_TELEMETRY_CAPTURED_FIELDS = __import__( + "contextvars" + ).ContextVar("fields") + obs_mod.INNER_ACCUMULATED_USAGE = __import__("contextvars").ContextVar( + "usage" + ) + obs_mod.INNER_RESPONSE_ID_CAPTURED_FIELD = "response_id" + obs_mod.INNER_USAGE_CAPTURED_FIELD = "usage" + obs_mod.AgentTelemetryLayer = _AgentTelemetryLayer + obs_mod._get_span = _get_span + obs_mod.get_function_span = get_function_span + obs_mod.get_tracer = lambda: tracer + obs_mod._get_span_attributes = _get_span_attributes + obs_mod._get_instructions_from_options = lambda _options: None + obs_mod._capture_messages = _capture_messages + obs_mod._get_response_attributes = lambda _attrs, _response, **_kwargs: {} + obs_mod._apply_accumulated_usage = lambda _attrs, _fields: None + obs_mod._capture_response = _capture_response + obs_mod.capture_exception = lambda span, exception, timestamp: ( + span.record_exception(exception) + ) + + af_mod = types.ModuleType("agent_framework") + af_mod.observability = obs_mod + types_mod = types.ModuleType("agent_framework._types") + types_mod.ResponseStream = _ResponseStream + tools_mod = types.ModuleType("agent_framework._tools") + tools_mod.get_function_span = get_function_span + mcp_mod = types.ModuleType("agent_framework._mcp") + + monkeypatch.setitem(sys.modules, "agent_framework", af_mod) + monkeypatch.setitem(sys.modules, "agent_framework.observability", obs_mod) + monkeypatch.setitem(sys.modules, "agent_framework._types", types_mod) + monkeypatch.setitem(sys.modules, "agent_framework._tools", tools_mod) + monkeypatch.setitem(sys.modules, "agent_framework._mcp", mcp_mod) + + original_anext = _ResponseStream.__anext__ + original_run_cleanup_hooks = _ResponseStream._run_cleanup_hooks + + util_genai_bridge.revert_util_genai_bridge() + util_genai_bridge.apply_util_genai_bridge() + try: + assert hasattr(_ResponseStream, "with_pull_context_manager") + patched_anext = _ResponseStream.__anext__ + util_genai_bridge.apply_util_genai_bridge() + assert _ResponseStream.__anext__ is patched_anext + + async def _updates(): + with tracer.start_as_current_span("chat qwen-plus"): + yield "delta" + + def _execute(): + return _ResponseStream( + _updates(), + finalizer=lambda _updates: types.SimpleNamespace( + messages=["done"], usage_details=None + ), + ) + + stream = _AgentTelemetryLayer()._trace_agent_invocation( + messages=["hello"], + session=None, + merged_options={}, + client_kwargs=None, + stream=True, + execute=_execute, + ) + + async def _consume(): + async for _ in stream: + pass + + asyncio.run(_consume()) + finally: + util_genai_bridge.revert_util_genai_bridge() + + spans = exporter.get_finished_spans() + by_name = {span.name: span for span in spans} + agent = by_name["invoke_agent legacy-agent"] + chat = by_name["chat qwen-plus"] + assert chat.parent.span_id == agent.context.span_id + assert chat.context.trace_id == agent.context.trace_id + input_messages = json.loads(agent.attributes["gen_ai.input.messages"]) + assert input_messages[0]["parts"][0]["content"] == "hello" + output_messages = json.loads(agent.attributes["gen_ai.output.messages"]) + assert output_messages[0]["parts"][0]["content"] == "done" + assert not hasattr(_ResponseStream, "with_pull_context_manager") + assert _ResponseStream.__anext__ is original_anext + assert _ResponseStream._run_cleanup_hooks is original_run_cleanup_hooks