diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py index 0785f01e36ba..44bacb0b86a9 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py @@ -188,11 +188,13 @@ def __init__( # Observability (logging + tracing) -------------------------------- _conn_str = applicationinsights_connection_string or self.config.appinsights_connection_string + _sensitive_data = os.environ.get("FOUNDRY_ENABLE_SENSITIVE_DATA", "true").lower() not in ("false", "0") if configure_observability is not None: try: configure_observability( connection_string=_conn_str, log_level=log_level, + enable_sensitive_data=_sensitive_data, ) except ValueError: raise # invalid log_level etc. โ€” user should fix their config diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_config.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_config.py index e22bc1ff1cf6..95111f467b91 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_config.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_config.py @@ -24,6 +24,9 @@ _ENV_FOUNDRY_AGENT_NAME = "FOUNDRY_AGENT_NAME" _ENV_FOUNDRY_AGENT_VERSION = "FOUNDRY_AGENT_VERSION" +_ENV_FOUNDRY_AGENT_INSTANCE_CLIENT_ID = "FOUNDRY_AGENT_INSTANCE_CLIENT_ID" +_ENV_FOUNDRY_AGENT_BLUEPRINT_CLIENT_ID = "FOUNDRY_AGENT_BLUEPRINT_CLIENT_ID" +_ENV_FOUNDRY_AGENT_TENANT_ID = "FOUNDRY_AGENT_TENANT_ID" _ENV_FOUNDRY_HOSTING_ENVIRONMENT = "FOUNDRY_HOSTING_ENVIRONMENT" _ENV_FOUNDRY_PROJECT_ENDPOINT = "FOUNDRY_PROJECT_ENDPOINT" _ENV_FOUNDRY_PROJECT_ARM_ID = "FOUNDRY_PROJECT_ARM_ID" @@ -283,6 +286,46 @@ def resolve_agent_version() -> str: return os.environ.get(_ENV_FOUNDRY_AGENT_VERSION, "") +def resolve_agent_id() -> str: + """Resolve the agent ID. + + Resolution order: + 1. ``FOUNDRY_AGENT_INSTANCE_CLIENT_ID`` environment variable. + 2. ``:`` if both are set. + 3. ```` if only name is set. + 4. Empty string if nothing is available. + + :return: The resolved agent ID, or an empty string if not determinable. + :rtype: str + """ + agent_id = os.environ.get(_ENV_FOUNDRY_AGENT_INSTANCE_CLIENT_ID, "") + if agent_id: + return agent_id + agent_name = os.environ.get(_ENV_FOUNDRY_AGENT_NAME, "") + agent_version = os.environ.get(_ENV_FOUNDRY_AGENT_VERSION, "") + if agent_name and agent_version: + return f"{agent_name}:{agent_version}" + return agent_name + + +def resolve_agent_blueprint_id() -> str: + """Resolve the agent blueprint client ID from the ``FOUNDRY_AGENT_BLUEPRINT_CLIENT_ID`` environment variable. + + :return: The agent blueprint client ID, or an empty string if not set. + :rtype: str + """ + return os.environ.get(_ENV_FOUNDRY_AGENT_BLUEPRINT_CLIENT_ID, "") + + +def resolve_agent_tenant_id() -> str: + """Resolve the agent tenant ID from the ``FOUNDRY_AGENT_TENANT_ID`` environment variable. + + :return: The agent tenant ID, or an empty string if not set. + :rtype: str + """ + return os.environ.get(_ENV_FOUNDRY_AGENT_TENANT_ID, "") + + def resolve_project_id() -> str: """Resolve the Foundry project ARM resource ID from the ``FOUNDRY_PROJECT_ARM_ID`` environment variable. diff --git a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_constants.py b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_constants.py index 74b7c0708931..8042b75f21cc 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_constants.py +++ b/sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_constants.py @@ -19,6 +19,8 @@ class Constants: # Tracing APPLICATIONINSIGHTS_CONNECTION_STRING = "APPLICATIONINSIGHTS_CONNECTION_STRING" OTEL_EXPORTER_OTLP_ENDPOINT = "OTEL_EXPORTER_OTLP_ENDPOINT" + FOUNDRY_AGENT365_TRACING_ENABLED = "FOUNDRY_AGENT365_TRACING_ENABLED" + FOUNDRY_ENABLE_SENSITIVE_DATA = "FOUNDRY_ENABLE_SENSITIVE_DATA" # SSE keep-alive SSE_KEEPALIVE_INTERVAL = "SSE_KEEPALIVE_INTERVAL" 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 faf5d23d7aaf..ac4b2e30a350 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 @@ -53,6 +53,9 @@ _ATTR_GEN_AI_SYSTEM = "gen_ai.system" _ATTR_GEN_AI_PROVIDER_NAME = "gen_ai.provider.name" _ATTR_GEN_AI_AGENT_ID = "gen_ai.agent.id" +_ATTR_GEN_AI_AGENT_BLUEPRINT_ID = "gen_ai.agent.blueprint.id" +_ATTR_GEN_AI_AGENT_TENANT_ID = "microsoft.tenant.id" +_ATTR_FOUNDRY_AGENT_TYPE = "microsoft.foundry.agent.type" _ATTR_GEN_AI_AGENT_NAME = "gen_ai.agent.name" _ATTR_GEN_AI_AGENT_VERSION = "gen_ai.agent.version" _ATTR_GEN_AI_RESPONSE_ID = "gen_ai.response.id" @@ -95,6 +98,7 @@ def configure_observability( *, connection_string: Optional[str] = None, log_level: Optional[str] = None, + enable_sensitive_data: bool = False, ) -> None: """Default observability setup: console logging + tracing/OTel export. @@ -111,6 +115,10 @@ def configure_observability( :paramtype connection_string: str or None :keyword log_level: Log level name (e.g. ``"INFO"``, ``"DEBUG"``). :paramtype log_level: str or None + :keyword enable_sensitive_data: Enable sensitive data recording + (prompts, tool arguments, results) for Agent Framework SDK + instrumentation. Defaults to False. + :paramtype enable_sensitive_data: bool """ # Console logging on the root logger so user logs are also visible. resolved_level = _config.resolve_log_level(log_level) @@ -135,10 +143,10 @@ def configure_observability( logging.getLogger("azure.core.pipeline.policies.http_logging_policy").setLevel(logging.WARNING) # Tracing and OTel export - _configure_tracing(connection_string=connection_string) + _configure_tracing(connection_string=connection_string, enable_sensitive_data=enable_sensitive_data) -def _configure_tracing(connection_string: Optional[str] = None) -> None: +def _configure_tracing(connection_string: Optional[str] = None, enable_sensitive_data: bool = False) -> None: """Configure OpenTelemetry exporters via the microsoft-opentelemetry distro. Internal helper called by :func:`configure_observability`. @@ -146,6 +154,9 @@ def _configure_tracing(connection_string: Optional[str] = None) -> None: :param connection_string: Application Insights connection string. When provided, traces and logs are exported to Azure Monitor. :type connection_string: str or None + :param enable_sensitive_data: Enable sensitive data recording for + Agent Framework SDK instrumentation. + :type enable_sensitive_data: bool """ resource = _create_resource() if resource is None: @@ -156,18 +167,17 @@ def _configure_tracing(connection_string: Optional[str] = None) -> None: agent_name = _config.resolve_agent_name() or None agent_version = _config.resolve_agent_version() or None project_id = _config.resolve_project_id() or None - - if agent_name and agent_version: - agent_id = f"{agent_name}:{agent_version}" - elif agent_name: - agent_id = agent_name - else: - agent_id = None + agent_id = _config.resolve_agent_id() or None + agent_blueprint_id = _config.resolve_agent_blueprint_id() or None + agent_tenant_id = _config.resolve_agent_tenant_id() or None span_processors = [ _FoundryEnrichmentSpanProcessor( agent_name=agent_name, agent_version=agent_version, agent_id=agent_id, project_id=project_id, + agent_blueprint_id=agent_blueprint_id, + agent_tenant_id=agent_tenant_id, + agent_type="hosted" if os.environ.get("FOUNDRY_HOSTING_ENVIRONMENT", "") else None, ), ] log_record_processors = [_BaggageLogRecordProcessor()] # type: ignore[list-item] @@ -178,6 +188,7 @@ def _configure_tracing(connection_string: Optional[str] = None) -> None: span_processors=span_processors, log_record_processors=log_record_processors, connection_string=connection_string, + enable_sensitive_data=enable_sensitive_data, ) logger.info("Tracing configured successfully via microsoft-opentelemetry distro.") except ImportError: @@ -192,6 +203,7 @@ def _setup_distro_export( span_processors: list[Any], log_record_processors: list[Any], connection_string: Optional[str] = None, + enable_sensitive_data: bool = False, ) -> None: """Delegate to microsoft-opentelemetry distro for exporter configuration. @@ -202,6 +214,8 @@ def _setup_distro_export( :keyword span_processors: Span processors to register. :keyword log_record_processors: Log record processors to register. :keyword connection_string: Application Insights connection string. + :keyword enable_sensitive_data: Enable sensitive data recording for + Agent Framework SDK instrumentation. """ from microsoft.opentelemetry import use_microsoft_opentelemetry @@ -209,6 +223,7 @@ def _setup_distro_export( "resource": resource, "span_processors": span_processors, "log_record_processors": log_record_processors, + "enable_sensitive_data": enable_sensitive_data, } # Azure Monitor export is off by default in the distro โ€” enable it @@ -217,6 +232,16 @@ def _setup_distro_export( kwargs["enable_azure_monitor"] = True kwargs["azure_monitor_connection_string"] = connection_string + # A365 tracing export โ€” enabled only in hosted environments. + if ( + os.environ.get("FOUNDRY_HOSTING_ENVIRONMENT", "") + and os.environ.get("FOUNDRY_AGENT365_TRACING_ENABLED", "").lower() in ("true", "1") + ): + kwargs["enable_a365"] = True + kwargs["a365_use_s2s_endpoint"] = True + kwargs["a365_enable_observability_exporter"] = True + kwargs["a365_observability_scope_override"] = "api://9b975845-388f-4429-889e-eab1ef63949c/.default" + use_microsoft_opentelemetry(**kwargs) @@ -460,15 +485,22 @@ class _FoundryEnrichmentSpanProcessor: def __init__( self, + *, agent_name: Optional[str] = None, agent_version: Optional[str] = None, agent_id: Optional[str] = None, project_id: Optional[str] = None, + agent_blueprint_id: Optional[str] = None, + agent_tenant_id: Optional[str] = None, + agent_type: Optional[str] = None, ) -> None: self.agent_name = agent_name self.agent_version = agent_version self.agent_id = agent_id self.project_id = project_id + self.agent_blueprint_id = agent_blueprint_id + self.agent_tenant_id = agent_tenant_id + self.agent_type = agent_type def on_start(self, span: Any, parent_context: Any = None) -> None: if self.project_id: @@ -504,6 +536,12 @@ def _on_ending(self, span: Any) -> None: attrs[_ATTR_GEN_AI_AGENT_VERSION] = self.agent_version if self.agent_id: attrs[_ATTR_GEN_AI_AGENT_ID] = self.agent_id + if self.agent_blueprint_id: + attrs[_ATTR_GEN_AI_AGENT_BLUEPRINT_ID] = self.agent_blueprint_id + if self.agent_tenant_id: + attrs[_ATTR_GEN_AI_AGENT_TENANT_ID] = self.agent_tenant_id + if self.agent_type and attrs.get(_ATTR_GEN_AI_OPERATION_NAME) == "invoke_agent": + attrs[_ATTR_FOUNDRY_AGENT_TYPE] = self.agent_type except Exception: # pylint: disable=broad-exception-caught logger.debug("Failed to enrich span attributes in _on_ending", exc_info=True) 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 2b3531b552d1..1a753cb50c1d 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/test_tracing.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/test_tracing.py @@ -74,6 +74,7 @@ def test_observability_receives_constructor_connection_string(self) -> None: mock_configure.assert_called_once_with( connection_string="InstrumentationKey=ctor", log_level=None, + enable_sensitive_data=True, ) def test_observability_disabled_when_none(self) -> None: @@ -160,6 +161,7 @@ def test_constructor_passes_connection_string(self) -> None: mock_configure.assert_called_once_with( connection_string="InstrumentationKey=ctor", log_level=None, + enable_sensitive_data=True, ) @@ -366,4 +368,66 @@ def test_agent_version_default_empty(self) -> None: assert resolve_agent_version() == "" +# ------------------------------------------------------------------ # +# agent_type attribute scoping +# ------------------------------------------------------------------ # + + +class TestAgentTypeAttribute: + """microsoft.foundry.agent.type is only set on invoke_agent spans.""" + + @staticmethod + def _create_provider(proc): + collector = _CollectorExporter() + provider = TracerProvider() + provider.add_span_processor(proc) + provider.add_span_processor(SimpleSpanProcessor(collector)) + return provider, collector + + def test_agent_type_set_on_invoke_agent_span(self) -> None: + """agent_type is written when gen_ai.operation.name == invoke_agent.""" + proc = _FoundryEnrichmentSpanProcessor( + agent_name="a", agent_version="1", agent_id="a:1", + agent_type="hosted", + ) + provider, collector = self._create_provider(proc) + tracer = provider.get_tracer("test") + + with tracer.start_as_current_span("invoke_agent") as span: + span.set_attribute("gen_ai.operation.name", "invoke_agent") + + attrs = dict(collector.spans[0].attributes) + assert attrs["microsoft.foundry.agent.type"] == "hosted" + + def test_agent_type_not_set_on_other_spans(self) -> None: + """agent_type must NOT appear on spans without invoke_agent operation.""" + proc = _FoundryEnrichmentSpanProcessor( + agent_name="a", agent_version="1", agent_id="a:1", + agent_type="hosted", + ) + provider, collector = self._create_provider(proc) + tracer = provider.get_tracer("test") + + with tracer.start_as_current_span("some_other_span") as span: + span.set_attribute("gen_ai.operation.name", "chat") + + attrs = dict(collector.spans[0].attributes) + assert "microsoft.foundry.agent.type" not in attrs + + def test_agent_type_none_skipped(self) -> None: + """When agent_type is None, attribute is never set even on invoke_agent.""" + proc = _FoundryEnrichmentSpanProcessor( + agent_name="a", agent_version="1", agent_id="a:1", + agent_type=None, + ) + provider, collector = self._create_provider(proc) + tracer = provider.get_tracer("test") + + with tracer.start_as_current_span("invoke_agent") as span: + span.set_attribute("gen_ai.operation.name", "invoke_agent") + + attrs = dict(collector.spans[0].attributes) + assert "microsoft.foundry.agent.type" not in attrs + + diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation.py index bf3120974fa0..fb1d52a36d3d 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation.py @@ -16,6 +16,7 @@ from typing import Any, Optional from opentelemetry import baggage as _otel_baggage, context as _otel_context +from opentelemetry.baggage.propagation import W3CBaggagePropagator from starlette.requests import Request from starlette.responses import JSONResponse, Response, StreamingResponse from starlette.routing import Route @@ -367,7 +368,14 @@ async def _create_invocation_endpoint(self, request: Request) -> Response: # Propagate invocation/session IDs as W3C baggage so downstream # services receive them automatically via the baggage header. + # Extract incoming baggage from request headers (only baggage, not traceparent) + # to preserve parent-child span relationships while inheriting caller's baggage entries. + _incoming_baggage_ctx = W3CBaggagePropagator().extract( + carrier={"baggage": request.headers.get("baggage", "")} + ) ctx = _otel_context.get_current() + for _bkey, _bval in _otel_baggage.get_all(context=_incoming_baggage_ctx).items(): + ctx = _otel_baggage.set_baggage(_bkey, _bval, context=ctx) ctx = _otel_baggage.set_baggage( "azure.ai.agentserver.invocation_id", invocation_id, context=ctx, ) diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/cspell.json b/sdk/agentserver/azure-ai-agentserver-invocations/cspell.json index 5858cd8e195b..e2180fd922d2 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/cspell.json +++ b/sdk/agentserver/azure-ai-agentserver-invocations/cspell.json @@ -4,6 +4,8 @@ "appinsights", "ASGI", "autouse", + "bkey", + "bval", "caplog", "genai", "hypercorn", 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 082ad23549ed..d7c4eef2985d 100644 --- a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_tracing.py +++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_tracing.py @@ -452,6 +452,82 @@ def test_agent_name_only_in_span_name(): assert "solo-agent" in invoke_spans[0].name +# --------------------------------------------------------------------------- +# Incoming W3C baggage propagation +# --------------------------------------------------------------------------- + +def test_incoming_baggage_merged_into_context(): + """Incoming W3C baggage header entries are merged into OTel context.""" + from opentelemetry import baggage as _otel_baggage, context as _otel_context + from opentelemetry.sdk.trace import SpanProcessor + + captured_baggage = {} + + class BaggageCaptureProcessor(SpanProcessor): + """Captures baggage visible when span starts.""" + def on_start(self, span, parent_context=None): + ctx = parent_context or _otel_context.get_current() + captured_baggage.update(_otel_baggage.get_all(context=ctx)) + + # Add our capture processor to the module provider + _MODULE_PROVIDER.add_span_processor(BaggageCaptureProcessor()) + + server = _make_tracing_server() + client = TestClient(server) + client.post( + "/invocations", + content=b"test", + headers={"baggage": "user.id=test-user-123,custom.key=custom-value"}, + ) + + # Incoming baggage entries should be present + assert captured_baggage.get("user.id") == "test-user-123" + assert captured_baggage.get("custom.key") == "custom-value" + + +def test_incoming_baggage_does_not_break_span_parenting(): + """Incoming baggage header does not break parent-child span relationships.""" + server = _make_tracing_server() + + # Create a traceparent to verify parenting is preserved + trace_id_hex = uuid.uuid4().hex + span_id_hex = uuid.uuid4().hex[:16] + traceparent = f"00-{trace_id_hex}-{span_id_hex}-01" + + client = TestClient(server) + client.post( + "/invocations", + content=b"test", + headers={ + "traceparent": traceparent, + "baggage": "user.id=test-user-456", + }, + ) + + 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 span should still have the same trace ID (parent-child preserved) + actual_trace_id = format(span.context.trace_id, "032x") + assert actual_trace_id == trace_id_hex + # And the parent span ID should match the traceparent + actual_parent_id = format(span.parent.span_id, "016x") + assert actual_parent_id == span_id_hex + + +def test_incoming_baggage_empty_header(): + """Empty baggage header does not cause errors.""" + server = _make_tracing_server() + client = TestClient(server) + resp = client.post( + "/invocations", + content=b"test", + headers={"baggage": ""}, + ) + assert resp.status_code == 200 + + # --------------------------------------------------------------------------- # Project endpoint attribute # --------------------------------------------------------------------------- diff --git a/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_endpoint_handler.py b/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_endpoint_handler.py index 85dcc182c35b..09ae02bc094c 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_endpoint_handler.py +++ b/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_endpoint_handler.py @@ -18,14 +18,13 @@ from opentelemetry import baggage as _otel_baggage from opentelemetry import context as _otel_context +from opentelemetry.baggage.propagation import W3CBaggagePropagator from starlette.requests import Request from starlette.responses import JSONResponse, Response, StreamingResponse from azure.ai.agentserver.core import ( # pylint: disable=import-error,no-name-in-module - detach_context, end_span, flush_spans, - set_current_span, trace_stream, ) from azure.ai.agentserver.responses.models._generated import ( @@ -416,15 +415,19 @@ def _wrap_streaming_response( # Inner wrap: trace_stream ends the span when iteration completes. traced = trace_stream(response.body_iterator, otel_span) - # Outer wrap: re-attach span as current context during streaming - # so child spans are correctly parented. + # Outer wrap: re-attach the full context (span + baggage) during streaming + # so child spans are correctly parented and baggage is visible to processors. + # We capture the context now (while baggage is still attached) rather than + # relying on get_current() later when the iterator actually runs. + _captured_ctx = _otel_context.get_current() + async def _iter_with_context(): # type: ignore[return] - token = set_current_span(otel_span) + token = _otel_context.attach(_captured_ctx) try: async for chunk in traced: yield chunk finally: - detach_context(token) + _otel_context.detach(token) response.body_iterator = _iter_with_context() return response @@ -716,7 +719,16 @@ async def handle_create(self, request: Request) -> Response: # pylint: disable= self._safe_set_attrs(otel_span, build_create_otel_attrs(ctx, request_id=request_id, project_id=_project_id)) # Set W3C baggage per spec ยง7.3 + # Extract incoming baggage from request headers (only baggage, not traceparent) + # to preserve parent-child span relationships while inheriting caller's baggage entries. + _incoming_baggage_ctx = W3CBaggagePropagator().extract( + carrier={"baggage": request.headers.get("baggage", "")} + ) bag_ctx = _otel_context.get_current() + # Merge incoming baggage entries (e.g. user.id) onto current context + for _bkey, _bval in _otel_baggage.get_all(context=_incoming_baggage_ctx).items(): + bag_ctx = _otel_baggage.set_baggage(_bkey, _bval, context=bag_ctx) + bag_ctx = _otel_baggage.set_baggage("azure.ai.agentserver.response_id", response_id, context=bag_ctx) bag_ctx = _otel_baggage.set_baggage( "azure.ai.agentserver.conversation_id", ctx.conversation_id or "", context=bag_ctx diff --git a/sdk/agentserver/azure-ai-agentserver-responses/cspell.json b/sdk/agentserver/azure-ai-agentserver-responses/cspell.json index 173bf9281425..69f59055e4b8 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/cspell.json +++ b/sdk/agentserver/azure-ai-agentserver-responses/cspell.json @@ -21,7 +21,9 @@ "JVBE", "hdrs", "myproj", - "myhost" + "myhost", + "bkey", + "bval" ], "ignorePaths": [ "*.csv", diff --git a/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_tracing.py b/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_tracing.py index 7a05437d118f..1e5799a7fac2 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_tracing.py +++ b/sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_tracing.py @@ -10,6 +10,7 @@ from typing import Any +import pytest from starlette.testclient import TestClient from azure.ai.agentserver.responses import ResponsesAgentServerHost, ResponsesServerOptions @@ -215,3 +216,103 @@ def test_tracing__span_tags_omit_request_id_when_header_absent() -> None: ) assert "request.id" not in hook.spans[0].tags + + +# --------------------------------------------------------------------------- +# Incoming W3C baggage propagation +# --------------------------------------------------------------------------- + + +def test_tracing__incoming_baggage_merged_into_context() -> None: + """Incoming W3C baggage header entries are merged into OTel context.""" + try: + from opentelemetry import baggage as _otel_baggage, context as _otel_context, trace + from opentelemetry.sdk.trace import TracerProvider as SdkTracerProvider, SpanProcessor + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + except ImportError: + pytest.skip("opentelemetry SDK not installed") + + captured_baggage: dict = {} + + class BaggageCaptureProcessor(SpanProcessor): + """Captures baggage visible when span starts.""" + def on_start(self, span, parent_context=None): + ctx = parent_context or _otel_context.get_current() + captured_baggage.update(_otel_baggage.get_all(context=ctx)) + + # Get or create a provider with our capture processor + existing = trace.get_tracer_provider() + if hasattr(existing, "add_span_processor"): + existing.add_span_processor(BaggageCaptureProcessor()) + else: + provider = SdkTracerProvider() + provider.add_span_processor(BaggageCaptureProcessor()) + trace.set_tracer_provider(provider) + + client = _build_client() + client.post( + "/responses", + json={"model": "gpt-4o-mini", "input": "hi", "stream": False}, + headers={"baggage": "user.id=test-user-789,custom.key=custom-value"}, + ) + + # Incoming baggage entries should be present + assert captured_baggage.get("user.id") == "test-user-789" + assert captured_baggage.get("custom.key") == "custom-value" + + +def test_tracing__incoming_baggage_does_not_break_span_parenting() -> None: + """Incoming baggage header does not break parent-child span relationships.""" + try: + from opentelemetry import trace + from opentelemetry.sdk.trace import TracerProvider as SdkTracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + except ImportError: + pytest.skip("opentelemetry SDK not installed") + + import uuid + + exporter = InMemorySpanExporter() + existing = trace.get_tracer_provider() + if hasattr(existing, "add_span_processor"): + existing.add_span_processor(SimpleSpanProcessor(exporter)) + else: + provider = SdkTracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + trace.set_tracer_provider(provider) + + trace_id_hex = uuid.uuid4().hex + span_id_hex = uuid.uuid4().hex[:16] + traceparent = f"00-{trace_id_hex}-{span_id_hex}-01" + + client = _build_client() + client.post( + "/responses", + json={"model": "gpt-4o-mini", "input": "hi", "stream": False}, + headers={ + "traceparent": traceparent, + "baggage": "user.id=test-user-parenting", + }, + ) + + spans = exporter.get_finished_spans() + # Find the invoke_agent span + matching_spans = [s for s in spans if "invoke_agent" in s.name] + assert len(matching_spans) >= 1 + span = matching_spans[0] + # The span should have the same trace ID (parent-child preserved) + actual_trace_id = format(span.context.trace_id, "032x") + assert actual_trace_id == trace_id_hex + + +def test_tracing__incoming_baggage_empty_header_no_error() -> None: + """Empty baggage header does not cause errors.""" + client = _build_client() + resp = client.post( + "/responses", + json={"model": "gpt-4o-mini", "input": "hi", "stream": False}, + headers={"baggage": ""}, + ) + assert resp.status_code == 200