fix(agno): propagate run identity to child spans (#239)

Author: sipercai

instrumentation-loongsuite/loongsuite-instrumentation-agno/CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Changed
+
+- Propagate Agno run user and session identity to model and tool spans while
+  preferring ENTRY baggage identity over Agno framework-provided values.
+
 ### Added
 
 - Capture `gen_ai.skill.*` attributes on Agno skill tool spans such as

instrumentation-loongsuite/loongsuite-instrumentation-agno/src/opentelemetry/instrumentation/agno/_wrapper.py

@@ -28,6 +28,8 @@
     create_agent_invocation,
     create_llm_invocation,
     create_tool_invocation,
+    reset_current_agno_run_identity,
+    set_current_agno_run_identity,
     update_agent_invocation_from_events,
     update_agent_invocation_from_response,
     update_llm_invocation_from_response,
@@ -113,6 +115,7 @@ def _run(
         self._handler.start_invoke_agent(
             invocation, context=otel_context.get_current()
         )
+        identity_token = set_current_agno_run_identity(invocation.attributes)
         try:
             response = wrapped(*args, **kwargs)
             update_agent_invocation_from_response(invocation, response)
@@ -125,6 +128,8 @@ def _run(
                 _error(exc),
             )
             raise
+        finally:
+            reset_current_agno_run_identity(identity_token)
 
     def _run_stream(
         self,
@@ -144,6 +149,21 @@ def generator() -> Iterator[Any]:
             self._handler.start_invoke_agent(
                 invocation, context=parent_context
             )
+            identity_token = None
+
+            def enter_identity() -> None:
+                nonlocal identity_token
+                identity_token = set_current_agno_run_identity(
+                    invocation.attributes
+                )
+
+            def exit_identity() -> None:
+                nonlocal identity_token
+                if identity_token is not None:
+                    reset_current_agno_run_identity(identity_token)
+                    identity_token = None
+
+            enter_identity()
             try:
                 stream = wrapped(*args, **kwargs)
                 for event in stream:
@@ -152,26 +172,37 @@ def generator() -> Iterator[Any]:
                             timeit.default_timer()
                         )
                     events.append(event)
-                    yield event
+                    # Do not expose run-local identity to caller code while
+                    # the stream chunk is yielded outside this wrapper.
+                    exit_identity()
+                    try:
+                        yield event
+                    finally:
+                        enter_identity()
                 update_agent_invocation_from_events(invocation, events)
                 _finish_invocation(self._handler.stop_invoke_agent, invocation)
                 finalized = True
             except BaseException as exc:
                 error = exc
                 raise
             finally:
-                if not finalized:
-                    if error is None or _is_stream_close(error):
-                        update_agent_invocation_from_events(invocation, events)
-                        _finish_invocation(
-                            self._handler.stop_invoke_agent, invocation
-                        )
-                    else:
-                        _finish_invocation(
-                            self._handler.fail_invoke_agent,
-                            invocation,
-                            _error(error),
-                        )
+                try:
+                    if not finalized:
+                        if error is None or _is_stream_close(error):
+                            update_agent_invocation_from_events(
+                                invocation, events
+                            )
+                            _finish_invocation(
+                                self._handler.stop_invoke_agent, invocation
+                            )
+                        else:
+                            _finish_invocation(
+                                self._handler.fail_invoke_agent,
+                                invocation,
+                                _error(error),
+                            )
+                finally:
+                    exit_identity()
 
         return generator()
 
@@ -214,6 +245,7 @@ async def _arun(
     ) -> Any:
         invocation = create_agent_invocation(instance, arguments)
         self._handler.start_invoke_agent(invocation, context=parent_context)
+        identity_token = set_current_agno_run_identity(invocation.attributes)
         try:
             response = wrapped(*args, **kwargs)
             if inspect.isawaitable(response):
@@ -228,6 +260,8 @@ async def _arun(
                 _error(exc),
             )
             raise
+        finally:
+            reset_current_agno_run_identity(identity_token)
 
     async def _arun_stream(
         self,
@@ -244,6 +278,21 @@ async def _arun_stream(
         finalized = False
         error = None
         self._handler.start_invoke_agent(invocation, context=parent_context)
+        identity_token = None
+
+        def enter_identity() -> None:
+            nonlocal identity_token
+            identity_token = set_current_agno_run_identity(
+                invocation.attributes
+            )
+
+        def exit_identity() -> None:
+            nonlocal identity_token
+            if identity_token is not None:
+                reset_current_agno_run_identity(identity_token)
+                identity_token = None
+
+        enter_identity()
         try:
             stream = wrapped(*args, **kwargs)
             if inspect.isawaitable(stream):
@@ -252,26 +301,35 @@ async def _arun_stream(
                 if invocation.monotonic_first_token_s is None:
                     invocation.monotonic_first_token_s = timeit.default_timer()
                 events.append(event)
-                yield event
+                # Do not expose run-local identity to caller code while
+                # the stream chunk is yielded outside this wrapper.
+                exit_identity()
+                try:
+                    yield event
+                finally:
+                    enter_identity()
             update_agent_invocation_from_events(invocation, events)
             _finish_invocation(self._handler.stop_invoke_agent, invocation)
             finalized = True
         except BaseException as exc:
             error = exc
             raise
         finally:
-            if not finalized:
-                if error is None or _is_stream_close(error):
-                    update_agent_invocation_from_events(invocation, events)
-                    _finish_invocation(
-                        self._handler.stop_invoke_agent, invocation
-                    )
-                else:
-                    _finish_invocation(
-                        self._handler.fail_invoke_agent,
-                        invocation,
-                        _error(error),
-                    )
+            try:
+                if not finalized:
+                    if error is None or _is_stream_close(error):
+                        update_agent_invocation_from_events(invocation, events)
+                        _finish_invocation(
+                            self._handler.stop_invoke_agent, invocation
+                        )
+                    else:
+                        _finish_invocation(
+                            self._handler.fail_invoke_agent,
+                            invocation,
+                            _error(error),
+                        )
+            finally:
+                exit_identity()
 
 
 class AgnoFunctionCallWrapper:

instrumentation-loongsuite/loongsuite-instrumentation-agno/src/opentelemetry/instrumentation/agno/utils.py

@@ -15,9 +15,11 @@
 from __future__ import annotations
 
 import json
+from contextvars import ContextVar, Token
 from dataclasses import asdict, is_dataclass
 from typing import Any, Dict, Mapping, Optional, Sequence
 
+from opentelemetry import baggage
 from opentelemetry.util.genai.extended_semconv.gen_ai_extended_attributes import (
     GEN_AI_SESSION_ID,
     GEN_AI_USER_ID,
@@ -44,6 +46,106 @@
     "get_skill_reference",
     "get_skill_script",
 }
+_CURRENT_AGNO_RUN_IDENTITY: ContextVar[dict[str, str] | None] = ContextVar(
+    "agno_current_run_identity",
+    default=None,
+)
+
+
+def _identity_value(value: Any) -> str | None:
+    if value is None:
+        return None
+    text = str(value).strip()
+    return text or None
+
+
+def _first_identity_value(*values: Any) -> str | None:
+    for value in values:
+        identity = _identity_value(value)
+        if identity is not None:
+            return identity
+    return None
+
+
+def _current_baggage_value(key: str) -> str | None:
+    try:
+        value = baggage.get_baggage(key)
+    except Exception:
+        return None
+    return _identity_value(value)
+
+
+def entry_baggage_identity_attributes() -> dict[str, str]:
+    attributes: dict[str, str] = {}
+    session_id = _current_baggage_value(GEN_AI_SESSION_ID)
+    user_id = _current_baggage_value(GEN_AI_USER_ID)
+    if session_id:
+        attributes[GEN_AI_SESSION_ID] = session_id
+    if user_id:
+        attributes[GEN_AI_USER_ID] = user_id
+    return attributes
+
+
+def _framework_identity_attributes(
+    agent: Any, arguments: Mapping[str, Any]
+) -> dict[str, str]:
+    attributes: dict[str, str] = {}
+    user_id = _first_identity_value(
+        arguments.get("user_id"),
+        getattr(agent, "user_id", None),
+    )
+    session_id = _first_identity_value(
+        arguments.get("session_id"),
+        getattr(agent, "session_id", None),
+    )
+    if user_id:
+        attributes[GEN_AI_USER_ID] = user_id
+    if session_id:
+        attributes[GEN_AI_SESSION_ID] = session_id
+    return attributes
+
+
+def agno_run_identity_attributes(
+    agent: Any, arguments: Mapping[str, Any]
+) -> dict[str, str]:
+    attributes = _framework_identity_attributes(agent, arguments)
+    attributes.update(entry_baggage_identity_attributes())
+    return attributes
+
+
+def set_current_agno_run_identity(
+    attributes: Mapping[str, Any],
+) -> Token[dict[str, str] | None]:
+    identity = {
+        key: value
+        for key in (GEN_AI_SESSION_ID, GEN_AI_USER_ID)
+        if (value := _identity_value(attributes.get(key))) is not None
+    }
+    return _CURRENT_AGNO_RUN_IDENTITY.set(identity or None)
+
+
+def reset_current_agno_run_identity(
+    token: Token[dict[str, str] | None],
+) -> None:
+    _CURRENT_AGNO_RUN_IDENTITY.reset(token)
+
+
+def _current_identity_attributes() -> dict[str, str]:
+    attributes: dict[str, str] = {}
+    run_identity = _CURRENT_AGNO_RUN_IDENTITY.get()
+    if run_identity:
+        attributes.update(run_identity)
+    # Re-read baggage so ENTRY identity stays highest priority for child spans
+    # even if the active OTel context changes while an Agno run is in progress.
+    attributes.update(entry_baggage_identity_attributes())
+    return attributes
+
+
+def _apply_current_identity(invocation: Any) -> str | None:
+    attributes = _current_identity_attributes()
+    for key, value in attributes.items():
+        invocation.attributes[key] = value
+    return attributes.get(GEN_AI_SESSION_ID)
 
 
 def _json_default(value: Any) -> Any:
@@ -421,15 +523,8 @@ def create_agent_invocation(
 ) -> InvokeAgentInvocation:
     model = getattr(agent, "model", None)
     input_value = arguments.get("input")
-    user_id = arguments.get("user_id") or getattr(agent, "user_id", None)
-    session_id = arguments.get("session_id") or getattr(
-        agent, "session_id", None
-    )
-    attributes: dict[str, Any] = {}
-    if user_id:
-        attributes[GEN_AI_USER_ID] = str(user_id)
-    if session_id:
-        attributes[GEN_AI_SESSION_ID] = str(session_id)
+    attributes: dict[str, Any] = agno_run_identity_attributes(agent, arguments)
+    session_id = attributes.get(GEN_AI_SESSION_ID)
 
     invocation = InvokeAgentInvocation(
         provider=_PROVIDER,
@@ -568,6 +663,9 @@ def create_llm_invocation(
         input_messages=convert_model_messages(arguments.get("messages")),
         tool_definitions=convert_tool_definitions(arguments.get("tools")),
     )
+    session_id = _apply_current_identity(invocation)
+    if session_id and not getattr(invocation, "conversation_id", None):
+        invocation.conversation_id = session_id
 
     for name in (
         "temperature",
@@ -657,6 +755,7 @@ def create_tool_invocation(function_call: Any) -> ExecuteToolInvocation:
         tool_type="function",
         tool_call_arguments=getattr(function_call, "arguments", None),
     )
+    _apply_current_identity(invocation)
     _apply_agno_skill_tool_metadata(invocation, invocation.tool_call_arguments)
     return invocation