From 913f77faca55df564ae4d4877eaa71881ca03116 Mon Sep 17 00:00:00 2001 From: Ankit Singhal Date: Tue, 10 Mar 2026 22:56:47 -0700 Subject: [PATCH 1/2] Extract incoming traceparent header for distributed trace propagation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The hosted agent server now extracts the W3C trace context (traceparent header) from incoming HTTP requests and uses it as the parent context for the root HostedAgents server span. This enables end-to-end distributed tracing where the caller's client span is the parent of the server span. When no traceparent header is present, behavior is unchanged — the server span becomes the root of a new trace. --- .../azure/ai/agentserver/core/server/base.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/server/base.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/server/base.py index 7a9f488227a7..acc93ce49dc3 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/server/base.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/server/base.py @@ -111,10 +111,17 @@ async def runs_endpoint(request): # Set up tracing context and span context = request.state.agent_run_context ctx = request_context.get() + # Extract W3C trace context from incoming request headers so the + # server span becomes a child of the caller's span when a + # traceparent header is present. + parent_ctx = TraceContextTextMapPropagator().extract( + carrier=dict(request.headers) + ) with self.tracer.start_as_current_span( name=f"HostedAgents-{context.response_id}", attributes=ctx, kind=trace.SpanKind.SERVER, + context=parent_ctx, ): try: logger.info("Start processing CreateResponse request.") From f4dccc7eb58b3742b5737b9aae30a40a2afc86a4 Mon Sep 17 00:00:00 2001 From: Ankit Singhal Date: Tue, 10 Mar 2026 23:18:31 -0700 Subject: [PATCH 2/2] Address review: use request.headers directly and add trace propagation tests - Use request.headers as carrier instead of dict(request.headers) since Starlette Headers is already a valid mapping type - Add unit tests for W3C trace context propagation: - Server span becomes child when traceparent header is present - Server span is root when no traceparent header is present - Validates lowercase header normalization works correctly --- .../azure/ai/agentserver/core/server/base.py | 2 +- .../tests/test_trace_context_propagation.py | 116 ++++++++++++++++++ 2 files changed, 117 insertions(+), 1 deletion(-) create mode 100644 sdk/agentserver/azure-ai-agentserver-core/tests/test_trace_context_propagation.py diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/server/base.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/server/base.py index acc93ce49dc3..1991a8803aa7 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/server/base.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/server/base.py @@ -115,7 +115,7 @@ async def runs_endpoint(request): # server span becomes a child of the caller's span when a # traceparent header is present. parent_ctx = TraceContextTextMapPropagator().extract( - carrier=dict(request.headers) + carrier=request.headers ) with self.tracer.start_as_current_span( name=f"HostedAgents-{context.response_id}", diff --git a/sdk/agentserver/azure-ai-agentserver-core/tests/test_trace_context_propagation.py b/sdk/agentserver/azure-ai-agentserver-core/tests/test_trace_context_propagation.py new file mode 100644 index 000000000000..eb296b8ca327 --- /dev/null +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/test_trace_context_propagation.py @@ -0,0 +1,116 @@ +"""Unit tests for W3C trace context propagation in the hosted agent server.""" + +import pytest +from opentelemetry import trace, context +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter, SpanExportResult +from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator + + +class ListSpanExporter(SpanExporter): + """A minimal span exporter that collects finished spans into a list.""" + + def __init__(self): + self.spans = [] + + def export(self, spans): + self.spans.extend(spans) + return SpanExportResult.SUCCESS + + def shutdown(self): + pass + + def clear(self): + self.spans.clear() + + +# Module-level provider (OTel only allows setting the global provider once) +_exporter = ListSpanExporter() +_provider = TracerProvider() +_provider.add_span_processor(SimpleSpanProcessor(_exporter)) +trace.set_tracer_provider(_provider) + + +@pytest.fixture(autouse=True) +def _clear_spans(): + """Clear collected spans before each test.""" + _exporter.clear() + yield + + +class TestTraceContextPropagation: + """Tests for W3C traceparent header extraction in runs_endpoint.""" + + def _make_traceparent(self): + """Create a traceparent header from a fresh span context.""" + tracer = trace.get_tracer("test-client") + with tracer.start_as_current_span("client-call") as span: + ctx = span.get_span_context() + carrier = {} + TraceContextTextMapPropagator().inject(carrier) + return carrier["traceparent"], ctx.trace_id, ctx.span_id + + def test_server_span_parents_to_incoming_traceparent(self): + """When a request carries a traceparent header, the server span must + become a child of that trace (same trace_id, parent_span_id matches).""" + traceparent, parent_trace_id, parent_span_id = self._make_traceparent() + + # Simulate what the server does: extract and start span with parent context + server_tracer = trace.get_tracer("azure.ai.agentserver") + incoming_headers = {"traceparent": traceparent} + extracted_ctx = TraceContextTextMapPropagator().extract(carrier=incoming_headers) + + with server_tracer.start_as_current_span( + name="HostedAgents-test", + kind=trace.SpanKind.SERVER, + context=extracted_ctx, + ) as server_span: + server_ctx = server_span.get_span_context() + + assert server_ctx.trace_id == parent_trace_id, ( + f"Server span trace_id {server_ctx.trace_id:#034x} should match " + f"parent trace_id {parent_trace_id:#034x}" + ) + + server_spans = [s for s in _exporter.spans if s.name == "HostedAgents-test"] + assert len(server_spans) == 1 + assert server_spans[0].parent.span_id == parent_span_id + + def test_server_span_is_root_without_traceparent(self): + """When no traceparent header is present, the server span must be a + root span (no parent, new trace_id).""" + server_tracer = trace.get_tracer("azure.ai.agentserver") + incoming_headers = {} + extracted_ctx = TraceContextTextMapPropagator().extract(carrier=incoming_headers) + + with server_tracer.start_as_current_span( + name="HostedAgents-no-parent", + kind=trace.SpanKind.SERVER, + context=extracted_ctx, + ): + pass + + server_spans = [s for s in _exporter.spans if s.name == "HostedAgents-no-parent"] + assert len(server_spans) == 1 + assert server_spans[0].parent is None, "Server span should be root when no traceparent is sent" + + def test_traceparent_header_case_insensitive(self): + """Starlette's request.headers is a case-insensitive mapping, so + mixed-case header keys are handled by the framework, not by our code. + This test verifies the propagator works with the lowercase key that + Starlette normalizes headers to.""" + traceparent, parent_trace_id, _ = self._make_traceparent() + + # Starlette normalizes all header keys to lowercase + server_tracer = trace.get_tracer("azure.ai.agentserver") + incoming_headers = {"traceparent": traceparent} + extracted_ctx = TraceContextTextMapPropagator().extract(carrier=incoming_headers) + + with server_tracer.start_as_current_span( + name="HostedAgents-case-test", + kind=trace.SpanKind.SERVER, + context=extracted_ctx, + ) as server_span: + server_ctx = server_span.get_span_context() + + assert server_ctx.trace_id == parent_trace_id