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 f1d8dd3db86f..493f58794776 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 @@ -345,6 +345,15 @@ def resolve_project_id() -> str: return os.environ.get(_ENV_FOUNDRY_PROJECT_ARM_ID, "") +def resolve_session_id() -> str: + """Resolve the default session ID from the ``FOUNDRY_AGENT_SESSION_ID`` environment variable. + + :return: The default session ID, or an empty string if not set. + :rtype: str + """ + return os.environ.get(_ENV_FOUNDRY_AGENT_SESSION_ID, "") + + def resolve_sse_keepalive_interval(interval: Optional[int] = None) -> int: """Resolve the SSE keep-alive interval from argument, env var, or default. 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 f9ab0ff60750..b0ed26bbeda1 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 @@ -146,10 +146,16 @@ def configure_observability( logging.getLogger(_noisy).setLevel(logging.WARNING) # Tracing and OTel export - _configure_tracing(connection_string=connection_string, enable_sensitive_data=enable_sensitive_data) + _configure_tracing( + connection_string=connection_string, + enable_sensitive_data=enable_sensitive_data, + ) -def _configure_tracing(connection_string: Optional[str] = None, enable_sensitive_data: bool = False) -> 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`. @@ -182,7 +188,13 @@ def _configure_tracing(connection_string: Optional[str] = None, enable_sensitive agent_tenant_id=agent_tenant_id, ), ] - log_record_processors = [_BaggageLogRecordProcessor()] # type: ignore[list-item] + log_record_processors = [ # type: ignore[list-item] + _BaggageLogRecordProcessor( + agent_name=agent_name, + agent_version=agent_version, + session_id=_config.resolve_session_id() or None, + ), + ] try: _setup_distro_export( @@ -521,6 +533,17 @@ class _BaggageLogRecordProcessor: for end-to-end correlation. """ + def __init__( + self, + *, + agent_name: Optional[str] = None, + agent_version: Optional[str] = None, + session_id: Optional[str] = None, + ) -> None: + self.agent_name = agent_name + self.agent_version = agent_version + self.session_id = session_id + def on_emit(self, log_data: Any) -> None: # pylint: disable=unused-argument """Copy baggage entries into the log record's attributes. @@ -528,11 +551,26 @@ def on_emit(self, log_data: Any) -> None: # pylint: disable=unused-argument :type log_data: any """ try: + if not hasattr(log_data, "log_record") or not log_data.log_record: + return + + attrs = log_data.log_record.attributes # type: ignore[assignment] + ctx = _otel_context.get_current() entries = _otel_baggage.get_all(context=ctx) - if entries and hasattr(log_data, 'log_record') and log_data.log_record: + if entries: for key, value in entries.items(): - log_data.log_record.attributes[key] = value # type: ignore[index] + attrs[key] = value # type: ignore[index] + + if self.agent_name and _ATTR_GEN_AI_AGENT_NAME not in attrs: + attrs[_ATTR_GEN_AI_AGENT_NAME] = self.agent_name + if self.agent_version and _ATTR_GEN_AI_AGENT_VERSION not in attrs: + attrs[_ATTR_GEN_AI_AGENT_VERSION] = self.agent_version + + bag_session = _otel_baggage.get_baggage(_BAGGAGE_SESSION_ID, context=ctx) + resolved_session = bag_session or self.session_id + if resolved_session and _ATTR_SESSION_ID not in attrs: + attrs[_ATTR_SESSION_ID] = resolved_session except Exception: # pylint: disable=broad-except pass 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 600b587afab6..bacf0f4d6dea 100644 --- a/sdk/agentserver/azure-ai-agentserver-core/tests/test_tracing.py +++ b/sdk/agentserver/azure-ai-agentserver-core/tests/test_tracing.py @@ -15,8 +15,12 @@ resolve_agent_name, resolve_agent_version, resolve_appinsights_connection_string, + resolve_session_id, +) +from azure.ai.agentserver.core._tracing import ( + _BaggageLogRecordProcessor, + _FoundryEnrichmentSpanProcessor, ) -from azure.ai.agentserver.core._tracing import _FoundryEnrichmentSpanProcessor class _CollectorExporter(SpanExporter): @@ -408,7 +412,7 @@ def test_invocation_id_propagates_to_child_spans(self) -> None: class TestAgentIdentityResolution: - """Tests for resolve_agent_name() and resolve_agent_version().""" + """Tests for agent identity/session resolution helpers.""" def test_agent_name_from_env(self) -> None: with mock.patch.dict(os.environ, {"FOUNDRY_AGENT_NAME": "my-agent"}): @@ -430,5 +434,78 @@ def test_agent_version_default_empty(self) -> None: with mock.patch.dict(os.environ, env, clear=True): assert resolve_agent_version() == "" + def test_session_id_from_env(self) -> None: + with mock.patch.dict(os.environ, {"FOUNDRY_AGENT_SESSION_ID": "session-1"}): + assert resolve_session_id() == "session-1" + + def test_session_id_default_empty(self) -> None: + env = os.environ.copy() + env.pop("FOUNDRY_AGENT_SESSION_ID", None) + with mock.patch.dict(os.environ, env, clear=True): + assert resolve_session_id() == "" + + +class _FakeLogRecord: + def __init__(self, attributes): + self.attributes = attributes + + +class _FakeLogData: + def __init__(self, attributes): + self.log_record = _FakeLogRecord(attributes) + + +class TestBaggageLogRecordProcessor: + def test_adds_agent_and_fallback_session_attributes(self) -> None: + proc = _BaggageLogRecordProcessor( + agent_name="agent-a", + agent_version="1.2.3", + session_id="session-fallback-1", + ) + log_data = _FakeLogData({}) + + proc.on_emit(log_data) + attrs = log_data.log_record.attributes + assert attrs["gen_ai.agent.name"] == "agent-a" + assert attrs["gen_ai.agent.version"] == "1.2.3" + assert attrs["microsoft.session.id"] == "session-fallback-1" + def test_prefers_baggage_session_id_over_fallback(self) -> None: + proc = _BaggageLogRecordProcessor( + agent_name="agent-a", + agent_version="1.2.3", + session_id="session-fallback-1", + ) + log_data = _FakeLogData({}) + + ctx = _otel_baggage.set_baggage( + "azure.ai.agentserver.session_id", "session-from-baggage", + ) + token = _otel_context.attach(ctx) + try: + proc.on_emit(log_data) + finally: + _otel_context.detach(token) + + attrs = log_data.log_record.attributes + assert attrs["microsoft.session.id"] == "session-from-baggage" + + def test_does_not_overwrite_existing_log_attributes(self) -> None: + proc = _BaggageLogRecordProcessor( + agent_name="agent-a", + agent_version="1.2.3", + session_id="session-fallback-1", + ) + attrs = { + "gen_ai.agent.name": "existing-name", + "gen_ai.agent.version": "0.0.1", + "microsoft.session.id": "existing-session", + } + log_data = _FakeLogData(attrs) + + proc.on_emit(log_data) + + assert attrs["gen_ai.agent.name"] == "existing-name" + assert attrs["gen_ai.agent.version"] == "0.0.1" + assert attrs["microsoft.session.id"] == "existing-session"