From 360b633b9cc2fbdbd780d80cda597188e795ce01 Mon Sep 17 00:00:00 2001 From: Neehar Duvvuri Date: Thu, 2 Apr 2026 11:21:27 -0400 Subject: [PATCH] Remove leaf_customer_span_id baggage re-parenting hack from tracing Use standard OTel traceparent propagation instead of overriding the parent span ID via the baggage header's leaf_customer_span_id key. - Remove _parse_baggage_key() and _override_parent_span_id() helpers - Remove baggage_header parameter from span/start_span/_extract_context - Simplify _prepare_request_span_args to return 3-tuple (no baggage) - Remove associated tests (TestParseBaggageKey, test_baggage_leaf_customer_span_id) - Preserve set_baggage/detach_baggage OTel context API (unrelated) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../azure/ai/agentserver/core/_tracing.py | 141 ++---------------- .../tests/test_tracing.py | 39 ----- .../tests/test_tracing.py | 36 ----- 3 files changed, 13 insertions(+), 203 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_tracing.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_tracing.py index 3e79af4ed93c..56fa18652f3e 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_tracing.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_tracing.py @@ -36,8 +36,6 @@ #: W3C Trace Context header names used for distributed trace propagation. _W3C_HEADERS = ("traceparent", "tracestate") -#: Baggage key whose value overrides the parent span ID. -_LEAF_CUSTOMER_SPAN_ID = "leaf_customer_span_id" # ------------------------------------------------------------------ # GenAI semantic convention attribute keys @@ -161,35 +159,21 @@ def __init__( def _extract_context( self, carrier: Optional[dict[str, str]], - baggage_header: Optional[str] = None, ) -> Any: """Extract parent trace context from a W3C carrier dict. - When a ``baggage`` header is provided and contains a - ``leaf_customer_span_id`` key, the parent span ID is overridden - so that the server's root span is parented under the leaf customer - span rather than the span referenced in the ``traceparent`` header. + Uses the standard ``traceparent`` / ``tracestate`` headers via + the OpenTelemetry :class:`TraceContextTextMapPropagator`. :param carrier: W3C trace-context headers or None. :type carrier: Optional[dict[str, str]] - :param baggage_header: Raw ``baggage`` header value or None. - :type baggage_header: Optional[str] :return: The extracted OTel context, or None. :rtype: Any """ if not carrier or self._propagator is None: return None - ctx = self._propagator.extract(carrier=carrier) - - if not baggage_header: - return ctx - - leaf_span_id = _parse_baggage_key(baggage_header, _LEAF_CUSTOMER_SPAN_ID) - if not leaf_span_id: - return ctx - - return _override_parent_span_id(ctx, leaf_span_id) + return self._propagator.extract(carrier=carrier) @staticmethod def _setup_azure_monitor(connection_string: str, resource: Any, trace_provider: Any) -> None: @@ -303,7 +287,6 @@ def span( name: str, attributes: Optional[dict[str, str]] = None, carrier: Optional[dict[str, str]] = None, - baggage_header: Optional[str] = None, ) -> Iterator[Any]: """Create a traced span if tracing is enabled, otherwise no-op. @@ -317,9 +300,6 @@ def span( :type attributes: Optional[dict[str, str]] :param carrier: Incoming HTTP headers for W3C trace-context propagation. :type carrier: Optional[dict[str, str]] - :param baggage_header: Raw ``baggage`` header value for - ``leaf_customer_span_id`` extraction. - :type baggage_header: Optional[str] :return: Context manager that yields the OTel span or *None*. :rtype: Iterator[Any] """ @@ -327,7 +307,7 @@ def span( yield None return - ctx = self._extract_context(carrier, baggage_header) + ctx = self._extract_context(carrier) with self._tracer.start_as_current_span( name=name, @@ -342,7 +322,6 @@ def start_span( name: str, attributes: Optional[dict[str, str]] = None, carrier: Optional[dict[str, str]] = None, - baggage_header: Optional[str] = None, ) -> Any: """Start a span without a context manager. @@ -356,16 +335,13 @@ def start_span( :type attributes: Optional[dict[str, str]] :param carrier: Incoming HTTP headers for W3C trace-context propagation. :type carrier: Optional[dict[str, str]] - :param baggage_header: Raw ``baggage`` header value for - ``leaf_customer_span_id`` extraction. - :type baggage_header: Optional[str] :return: The OTel span, or *None* when tracing is disabled. :rtype: Any """ if not self._enabled or self._tracer is None: return None - ctx = self._extract_context(carrier, baggage_header) + ctx = self._extract_context(carrier) return self._tracer.start_span( name=name, @@ -385,7 +361,7 @@ def _prepare_request_span_args( span_operation: str, operation_name: Optional[str] = None, session_id: str = "", - ) -> tuple[str, dict[str, str], dict[str, str], Optional[str]]: + ) -> tuple[str, dict[str, str], dict[str, str]]: """Extract headers and build span arguments for a request. Shared pipeline used by :meth:`start_request_span` and @@ -403,16 +379,15 @@ def _prepare_request_span_args( :param session_id: Session ID from the ``agent_session_id`` query parameter. Defaults to ``""`` (no session). :type session_id: str - :return: ``(name, attributes, carrier, baggage)`` ready for + :return: ``(name, attributes, carrier)`` ready for :meth:`span` or :meth:`start_span`. - :rtype: tuple[str, dict[str, str], dict[str, str], Optional[str]] + :rtype: tuple[str, dict[str, str], dict[str, str]] """ carrier = _extract_w3c_carrier(headers) - baggage = headers.get("baggage") span_attrs = self.build_span_attrs( invocation_id, session_id, operation_name=operation_name ) - return self.span_name(span_operation), span_attrs, carrier, baggage + return self.span_name(span_operation), span_attrs, carrier def start_request_span( self, @@ -447,11 +422,11 @@ def start_request_span( :return: The OTel span, or *None* when tracing is disabled. :rtype: Any """ - name, attrs, carrier, baggage = self._prepare_request_span_args( + name, attrs, carrier = self._prepare_request_span_args( headers, invocation_id, span_operation, operation_name, session_id=session_id, ) - return self.start_span(name, attributes=attrs, carrier=carrier, baggage_header=baggage) + return self.start_span(name, attributes=attrs, carrier=carrier) @contextmanager def request_span( @@ -486,11 +461,11 @@ def request_span( :return: Context manager that yields the OTel span or *None*. :rtype: Iterator[Any] """ - name, attrs, carrier, baggage = self._prepare_request_span_args( + name, attrs, carrier = self._prepare_request_span_args( headers, invocation_id, span_operation, operation_name, session_id=session_id, ) - with self.span(name, attributes=attrs, carrier=carrier, baggage_header=baggage) as otel_span: + with self.span(name, attributes=attrs, carrier=carrier) as otel_span: yield otel_span # ------------------------------------------------------------------ @@ -902,93 +877,3 @@ def _extract_w3c_carrier(headers: Mapping[str, str]) -> dict[str, str]: """ result: dict[str, str] = {k: v for k in _W3C_HEADERS if (v := headers.get(k)) is not None} return result - - -def _parse_baggage_key(baggage: str, key: str) -> Optional[str]: - """Parse a single key from a W3C Baggage header value. - - The `W3C Baggage`_ format is a comma-separated list of - ``key=value`` pairs with optional properties after a ``;``. - - Example:: - - leaf_customer_span_id=abc123,other=val - - .. _W3C Baggage: https://www.w3.org/TR/baggage/ - - :param baggage: The raw header value. - :type baggage: str - :param key: The baggage key to look up. - :type key: str - :return: The value for *key*, or *None* if not found. - :rtype: Optional[str] - """ - for member in baggage.split(","): - member = member.strip() - if not member: - continue - # Split on first '=' only; value may contain '=' - kv_part = member.split(";", 1)[0] # strip optional properties - eq_idx = kv_part.find("=") - if eq_idx < 0: - continue - k = kv_part[:eq_idx].strip() - v = kv_part[eq_idx + 1:].strip() - if k == key: - return v - return None - - -def _override_parent_span_id(ctx: Any, hex_span_id: str) -> Any: - """Create a new context with the same trace ID but a different parent span ID. - - Constructs a :class:`~opentelemetry.trace.SpanContext` with the trace ID - taken from the existing context and the span ID replaced by - *hex_span_id*. The resulting context can be used as the ``context`` - argument to ``start_span`` / ``start_as_current_span``. - - Returns the original *ctx* unchanged if *hex_span_id* is invalid or - ``opentelemetry-api`` is not installed. - - Per invocation-protocol-spec. - - :param ctx: An OTel context produced by ``TraceContextTextMapPropagator.extract()``. - :type ctx: Any - :param hex_span_id: 16-character lower-case hex string representing the - desired parent span ID. - :type hex_span_id: str - :return: A context with the overridden parent span ID, or the original. - :rtype: Any - """ - if not _HAS_OTEL: - return ctx - - # A valid OTel span ID is exactly 16 hex characters (8 bytes). - if len(hex_span_id) != 16: - logger.warning("Invalid leaf_customer_span_id length in baggage: %r (expected 16 hex chars)", hex_span_id) - return ctx - - try: - new_span_id = int(hex_span_id, 16) - except (ValueError, TypeError): - logger.warning("Invalid leaf_customer_span_id in baggage: %r", hex_span_id) - return ctx - - if new_span_id == 0: - return ctx - - # Grab the trace ID from the current parent span in ctx. - current_span = trace.get_current_span(ctx) - current_ctx = current_span.get_span_context() - if current_ctx is None or not current_ctx.is_valid: - return ctx - - custom_span_ctx = trace.SpanContext( - trace_id=current_ctx.trace_id, - span_id=new_span_id, - is_remote=True, - trace_flags=current_ctx.trace_flags, - trace_state=current_ctx.trace_state, - ) - custom_parent = trace.NonRecordingSpan(custom_span_ctx) - return trace.set_span_in_context(custom_parent, ctx) diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/test_tracing.py b/sdk/agentserver/azure-ai-agentserver-core/tests/test_tracing.py index bcacefe59bda..1c1b7348a7d4 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/test_tracing.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/test_tracing.py @@ -13,7 +13,6 @@ resolve_appinsights_connection_string, ) from azure.ai.agentserver.core._constants import Constants -from azure.ai.agentserver.core._tracing import _parse_baggage_key # ------------------------------------------------------------------ # @@ -184,41 +183,3 @@ def test_agent_version_default_empty(self) -> None: env.pop(Constants.FOUNDRY_AGENT_VERSION, None) with mock.patch.dict(os.environ, env, clear=True): assert resolve_agent_version() == "" - - -# ------------------------------------------------------------------ # -# Baggage parsing (unit tests for _parse_baggage_key) -# ------------------------------------------------------------------ # - - -class TestParseBaggageKey: - """Unit tests for _parse_baggage_key().""" - - def test_single_key(self) -> None: - assert _parse_baggage_key("leaf_customer_span_id=abc123", "leaf_customer_span_id") == "abc123" - - def test_multiple_keys(self) -> None: - baggage = "key1=val1,leaf_customer_span_id=def456,key2=val2" - assert _parse_baggage_key(baggage, "leaf_customer_span_id") == "def456" - - def test_key_not_found(self) -> None: - assert _parse_baggage_key("key1=val1,key2=val2", "leaf_customer_span_id") is None - - def test_empty_baggage(self) -> None: - assert _parse_baggage_key("", "leaf_customer_span_id") is None - - def test_key_with_properties(self) -> None: - baggage = "leaf_customer_span_id=abc123;prop1=x" - assert _parse_baggage_key(baggage, "leaf_customer_span_id") == "abc123" - - def test_whitespace_handling(self) -> None: - baggage = " leaf_customer_span_id = abc123 , other = val " - assert _parse_baggage_key(baggage, "leaf_customer_span_id") == "abc123" - - def test_value_with_equals(self) -> None: - baggage = "leaf_customer_span_id=abc=123" - assert _parse_baggage_key(baggage, "leaf_customer_span_id") == "abc=123" - - def test_no_equals_in_member(self) -> None: - baggage = "malformed_entry,leaf_customer_span_id=good" - assert _parse_baggage_key(baggage, "leaf_customer_span_id") == "good" diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_tracing.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_tracing.py index a5d8a78ff9a8..49a1219c4ff0 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_tracing.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_tracing.py @@ -428,42 +428,6 @@ async def test_namespaced_invocation_id_attribute(): assert attrs.get("azure.ai.agentserver.invocations.invocation_id") == inv_id -# --------------------------------------------------------------------------- -# Baggage tests -# --------------------------------------------------------------------------- - -@pytest.mark.asyncio -async def test_baggage_leaf_customer_span_id(): - """Baggage leaf_customer_span_id overrides parent span ID.""" - server = _make_tracing_server() - transport = ASGITransport(app=server.app) - - trace_id_hex = uuid.uuid4().hex - original_span_id = uuid.uuid4().hex[:16] - leaf_span_id = uuid.uuid4().hex[:16] - traceparent = f"00-{trace_id_hex}-{original_span_id}-01" - baggage = f"leaf_customer_span_id={leaf_span_id}" - - async with AsyncClient(transport=transport, base_url="http://testserver") as client: - await client.post( - "/invocations", - content=b"test", - headers={ - "traceparent": traceparent, - "baggage": baggage, - }, - ) - - spans = _get_spans() - invoke_spans = [s for s in spans if "invoke_agent" in s.name] - assert len(invoke_spans) >= 1 - span = invoke_spans[0] - # The parent span ID should be overridden to leaf_span_id - if span.parent is not None: - actual_parent_span_id = format(span.parent.span_id, "016x") - assert actual_parent_span_id == leaf_span_id - - # --------------------------------------------------------------------------- # Agent name/version in span names # ---------------------------------------------------------------------------