From 2d36e85e9e27170558a608a7ca1231f5c74bf6f3 Mon Sep 17 00:00:00 2001 From: "mxiamxia@gmail.com" Date: Mon, 30 Jun 2025 13:50:24 -0700 Subject: [PATCH] Improve resource attributes handling --- .../distro/_aws_attribute_keys.py | 1 + .../distro/aws_opentelemetry_configurator.py | 14 ++++-- .../distro/aws_opentelemetry_distro.py | 4 +- .../test_aws_opentelementry_configurator.py | 50 +++++++++++++++++++ 4 files changed, 64 insertions(+), 5 deletions(-) diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py index 23ba661af..41a01d662 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py @@ -33,3 +33,4 @@ AWS_LAMBDA_FUNCTION_NAME: str = "aws.lambda.function.name" AWS_LAMBDA_RESOURCEMAPPING_ID: str = "aws.lambda.resource_mapping.id" AWS_LAMBDA_FUNCTION_ARN: str = "aws.lambda.function.arn" +AWS_AI_AGENT_TYPE: str = "aws.ai.agent.type" diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py index ff486285e..761852f39 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_configurator.py @@ -9,7 +9,7 @@ from importlib_metadata import version from typing_extensions import override -from amazon.opentelemetry.distro._aws_attribute_keys import AWS_LOCAL_SERVICE +from amazon.opentelemetry.distro._aws_attribute_keys import AWS_AI_AGENT_TYPE, AWS_LOCAL_SERVICE from amazon.opentelemetry.distro._aws_resource_attribute_configurator import get_service_attribute from amazon.opentelemetry.distro._utils import is_agent_observability_enabled, is_installed from amazon.opentelemetry.distro.always_record_sampler import AlwaysRecordSampler @@ -174,7 +174,7 @@ def _initialize_components(): AwsEksResourceDetector(), AwsEcsResourceDetector(), ] - if not _is_lambda_environment() + if not (_is_lambda_environment() or is_agent_observability_enabled()) else [] ) @@ -555,7 +555,15 @@ def _customize_resource(resource: Resource) -> Resource: if is_unknown: _logger.debug("No valid service name found") - return resource.merge(Resource.create({AWS_LOCAL_SERVICE: service_name})) + custom_attributes = {AWS_LOCAL_SERVICE: service_name} + + if is_agent_observability_enabled(): + # Add aws.ai.agent.type if it doesn't exist in the resource + if resource and resource.attributes.get(AWS_AI_AGENT_TYPE) is None: + # Set a default agent type for AI agent observability + custom_attributes[AWS_AI_AGENT_TYPE] = "default" + + return resource.merge(Resource.create(custom_attributes)) def _is_application_signals_enabled(): diff --git a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_distro.py b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_distro.py index cf8109780..9fde5e248 100644 --- a/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_distro.py +++ b/aws-opentelemetry-distro/src/amazon/opentelemetry/distro/aws_opentelemetry_distro.py @@ -70,8 +70,6 @@ def _configure(self, **kwargs): os.environ.setdefault(OTEL_EXPORTER_OTLP_PROTOCOL, "http/protobuf") - super(AwsOpenTelemetryDistro, self)._configure() - os.environ.setdefault(OTEL_PROPAGATORS, "xray,tracecontext,b3,b3multi") os.environ.setdefault(OTEL_PYTHON_ID_GENERATOR, "xray") os.environ.setdefault( @@ -117,5 +115,7 @@ def _configure(self, **kwargs): # Disable AWS Application Signals by default os.environ.setdefault(APPLICATION_SIGNALS_ENABLED_CONFIG, "false") + super(AwsOpenTelemetryDistro, self)._configure() + if kwargs.get("apply_patches", True): apply_instrumentation_patches() diff --git a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py index d37568a97..795e1ce4f 100644 --- a/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py +++ b/aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_opentelementry_configurator.py @@ -10,6 +10,7 @@ from requests import Session +from amazon.opentelemetry.distro._aws_attribute_keys import AWS_AI_AGENT_TYPE, AWS_LOCAL_SERVICE from amazon.opentelemetry.distro.always_record_sampler import AlwaysRecordSampler from amazon.opentelemetry.distro.attribute_propagating_span_processor import AttributePropagatingSpanProcessor from amazon.opentelemetry.distro.aws_batch_unsampled_span_processor import BatchUnsampledSpanProcessor @@ -27,6 +28,7 @@ _custom_import_sampler, _customize_logs_exporter, _customize_metric_exporters, + _customize_resource, _customize_sampler, _customize_span_exporter, _customize_span_processors, @@ -67,6 +69,7 @@ from opentelemetry.sdk.trace import Span, SpanProcessor, Tracer, TracerProvider from opentelemetry.sdk.trace.export import SpanExporter from opentelemetry.sdk.trace.sampling import DEFAULT_ON, Sampler +from opentelemetry.semconv.resource import ResourceAttributes from opentelemetry.trace import get_tracer_provider @@ -1091,6 +1094,53 @@ def test_customize_metric_exporters_with_emf(self): self.assertEqual(len(metric_readers), 1) self.assertIsInstance(metric_readers[0], PeriodicExportingMetricReader) + @patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.is_agent_observability_enabled") + @patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.get_service_attribute") + def test_customize_resource_without_agent_observability(self, mock_get_service_attribute, mock_is_agent_enabled): + """Test _customize_resource when agent observability is disabled""" + mock_is_agent_enabled.return_value = False + mock_get_service_attribute.return_value = ("test-service", False) + + resource = Resource.create({ResourceAttributes.SERVICE_NAME: "test-service"}) + result = _customize_resource(resource) + + # Should only have AWS_LOCAL_SERVICE added + self.assertEqual(result.attributes[AWS_LOCAL_SERVICE], "test-service") + self.assertNotIn(AWS_AI_AGENT_TYPE, result.attributes) + + @patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.is_agent_observability_enabled") + @patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.get_service_attribute") + def test_customize_resource_with_agent_observability_default( + self, mock_get_service_attribute, mock_is_agent_enabled + ): + """Test _customize_resource when agent observability is enabled with default agent type""" + mock_is_agent_enabled.return_value = True + mock_get_service_attribute.return_value = ("test-service", False) + + resource = Resource.create({ResourceAttributes.SERVICE_NAME: "test-service"}) + result = _customize_resource(resource) + + # Should have both AWS_LOCAL_SERVICE and AWS_AI_AGENT_TYPE with default value + self.assertEqual(result.attributes[AWS_LOCAL_SERVICE], "test-service") + self.assertEqual(result.attributes[AWS_AI_AGENT_TYPE], "default") + + @patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.is_agent_observability_enabled") + @patch("amazon.opentelemetry.distro.aws_opentelemetry_configurator.get_service_attribute") + def test_customize_resource_with_existing_agent_type(self, mock_get_service_attribute, mock_is_agent_enabled): + """Test _customize_resource when agent type already exists in resource""" + mock_is_agent_enabled.return_value = True + mock_get_service_attribute.return_value = ("test-service", False) + + # Create resource with existing agent type + resource = Resource.create( + {ResourceAttributes.SERVICE_NAME: "test-service", AWS_AI_AGENT_TYPE: "existing-agent"} + ) + result = _customize_resource(resource) + + # Should preserve existing agent type and not override it + self.assertEqual(result.attributes[AWS_LOCAL_SERVICE], "test-service") + self.assertEqual(result.attributes[AWS_AI_AGENT_TYPE], "existing-agent") + def validate_distro_environ(): tc: TestCase = TestCase()