Address code review feedback: improve tests, use type-safe OTel detection, and remove redundant code

Author: vcjana

aws/codegen-aws-sdk/src/main/kotlin/software/amazon/smithy/rustsdk/AwsCodegenDecorator.kt

@@ -66,7 +66,6 @@ val DECORATORS: List<ClientCodegenDecorator> =
             AwsRequestIdDecorator(),
             DisabledAuthDecorator(),
             RecursionDetectionDecorator(),
-            EndpointOverrideDecorator(),
             ObservabilityDetectionDecorator(),
             InvocationIdDecorator(),
             RetryInformationHeaderDecorator(),

aws/rust-runtime/aws-runtime/Cargo.toml

@@ -22,6 +22,7 @@ aws-smithy-async = { path = "../../../rust-runtime/aws-smithy-async" }
 aws-smithy-eventstream = { path = "../../../rust-runtime/aws-smithy-eventstream", optional = true }
 aws-smithy-http = { path = "../../../rust-runtime/aws-smithy-http" }
 aws-smithy-observability = { path = "../../../rust-runtime/aws-smithy-observability" }
+aws-smithy-observability-otel = { path = "../../../rust-runtime/aws-smithy-observability-otel" }
 aws-smithy-runtime = { path = "../../../rust-runtime/aws-smithy-runtime", features = ["client"] }
 aws-smithy-runtime-api = { path = "../../../rust-runtime/aws-smithy-runtime-api", features = ["client"] }
 aws-smithy-types = { path = "../../../rust-runtime/aws-smithy-types" }

aws/rust-runtime/aws-runtime/src/endpoint_override.rs

@@ -6,14 +6,19 @@
 //! Endpoint override detection for business metrics tracking
 
 use aws_smithy_runtime_api::box_error::BoxError;
-use aws_smithy_runtime_api::client::interceptors::context::BeforeSerializationInterceptorContextRef;
+use aws_smithy_runtime_api::client::interceptors::context::BeforeTransmitInterceptorContextRef;
 use aws_smithy_runtime_api::client::interceptors::Intercept;
+use aws_smithy_runtime_api::client::runtime_components::RuntimeComponents;
 use aws_smithy_runtime_api::client::runtime_plugin::RuntimePlugin;
-use aws_smithy_types::config_bag::{ConfigBag, FrozenLayer};
+use aws_smithy_types::config_bag::{ConfigBag, FrozenLayer, Layer};
 
 use crate::sdk_feature::AwsSdkFeature;
 
 /// Interceptor that detects custom endpoint URLs for business metrics
+///
+/// This interceptor checks at runtime if a `StaticUriEndpointResolver` is configured,
+/// which indicates that `.endpoint_url()` was called. When detected, it stores the
+/// `AwsSdkFeature::EndpointOverride` feature flag for business metrics tracking.
 #[derive(Debug, Default)]
 #[non_exhaustive]
 pub struct EndpointOverrideInterceptor;
@@ -30,19 +35,25 @@ impl Intercept for EndpointOverrideInterceptor {
         "EndpointOverrideInterceptor"
     }
 
-    fn read_before_execution(
+    fn read_after_serialization(
         &self,
-        _context: &BeforeSerializationInterceptorContextRef<'_>,
+        _context: &BeforeTransmitInterceptorContextRef<'_>,
+        runtime_components: &RuntimeComponents,
         cfg: &mut ConfigBag,
     ) -> Result<(), BoxError> {
-        // Check if endpoint_url was set in config
-        if cfg
-            .load::<aws_types::endpoint_config::EndpointUrl>()
-            .is_some()
-        {
+        // Check if the endpoint resolver is a StaticUriEndpointResolver
+        // This indicates that .endpoint_url() was called
+        let resolver = runtime_components.endpoint_resolver();
+
+        // Check the resolver's debug string to see if it's StaticUriEndpointResolver
+        let debug_str = format!("{:?}", resolver);
+
+        if debug_str.contains("StaticUriEndpointResolver") {
+            // Store in interceptor_state
             cfg.interceptor_state()
                 .store_append(AwsSdkFeature::EndpointOverride);
         }
+
         Ok(())
     }
 }
@@ -65,6 +76,15 @@ impl EndpointOverrideRuntimePlugin {
     pub fn new(config: Option<FrozenLayer>) -> Self {
         Self { config }
     }
+
+    /// Creates a new `EndpointOverrideRuntimePlugin` and marks that endpoint override is enabled
+    pub fn new_with_feature_flag() -> Self {
+        let mut layer = Layer::new("endpoint_override");
+        layer.store_append(AwsSdkFeature::EndpointOverride);
+        Self {
+            config: Some(layer.freeze()),
+        }
+    }
 }
 
 impl RuntimePlugin for EndpointOverrideRuntimePlugin {
@@ -77,10 +97,6 @@ impl RuntimePlugin for EndpointOverrideRuntimePlugin {
 mod tests {
     use super::*;
     use crate::sdk_feature::AwsSdkFeature;
-    use aws_smithy_runtime_api::client::interceptors::context::{
-        BeforeSerializationInterceptorContextRef, Input, InterceptorContext,
-    };
-    use aws_smithy_types::config_bag::ConfigBag;
 
     #[test]
     fn test_plugin_with_no_config() {
@@ -89,56 +105,55 @@ mod tests {
     }
 
     #[test]
-    fn test_interceptor_detects_endpoint_url_when_present() {
-        let interceptor = EndpointOverrideInterceptor::new();
-        let mut cfg = ConfigBag::base();
-
-        // Set endpoint URL in config
-        let endpoint_url =
-            aws_types::endpoint_config::EndpointUrl("https://custom.example.com".to_string());
-        cfg.interceptor_state().store_put(endpoint_url);
-
-        // Create a dummy context
-        let input = Input::doesnt_matter();
-        let ctx = InterceptorContext::new(input);
-        let context = BeforeSerializationInterceptorContextRef::from(&ctx);
-
-        // Run the interceptor
-        interceptor
-            .read_before_execution(&context, &mut cfg)
-            .unwrap();
+    fn test_plugin_with_feature_flag() {
+        let plugin = EndpointOverrideRuntimePlugin::new_with_feature_flag();
+        let config = plugin.config().expect("config should be set");
 
-        // Verify feature flag was set in interceptor_state
-        let features: Vec<_> = cfg
-            .interceptor_state()
-            .load::<AwsSdkFeature>()
-            .cloned()
-            .collect();
+        // Verify the feature flag is present in the config
+        let features: Vec<_> = config.load::<AwsSdkFeature>().cloned().collect();
         assert_eq!(features.len(), 1);
         assert_eq!(features[0], AwsSdkFeature::EndpointOverride);
     }
 
     #[test]
-    fn test_interceptor_does_not_set_flag_when_endpoint_url_absent() {
-        let interceptor = EndpointOverrideInterceptor::new();
+    fn test_interceptor_detects_static_uri_resolver() {
+        use aws_smithy_runtime::client::orchestrator::endpoints::StaticUriEndpointResolver;
+        use aws_smithy_runtime_api::client::endpoint::SharedEndpointResolver;
+        use aws_smithy_runtime_api::client::interceptors::context::{Input, InterceptorContext};
+        use aws_smithy_runtime_api::client::orchestrator::HttpRequest;
+        use aws_smithy_runtime_api::client::runtime_components::RuntimeComponentsBuilder;
+        use aws_smithy_types::config_bag::ConfigBag;
+
+        // Create a StaticUriEndpointResolver
+        let endpoint_resolver = SharedEndpointResolver::new(StaticUriEndpointResolver::uri(
+            "https://custom.example.com",
+        ));
+
+        let mut context = InterceptorContext::new(Input::doesnt_matter());
+        context.enter_serialization_phase();
+        context.set_request(HttpRequest::empty());
+        let _ = context.take_input();
+        context.enter_before_transmit_phase();
+
+        let rc = RuntimeComponentsBuilder::for_tests()
+            .with_endpoint_resolver(Some(endpoint_resolver))
+            .build()
+            .unwrap();
         let mut cfg = ConfigBag::base();
 
-        // Create a dummy context
-        let input = Input::doesnt_matter();
-        let ctx = InterceptorContext::new(input);
-        let context = BeforeSerializationInterceptorContextRef::from(&ctx);
-
-        // Run the interceptor without setting endpoint URL
+        let interceptor = EndpointOverrideInterceptor::new();
+        let ctx = Into::into(&context);
         interceptor
-            .read_before_execution(&context, &mut cfg)
+            .read_after_serialization(&ctx, &rc, &mut cfg)
             .unwrap();
 
-        // Verify no feature flag was set
+        // Verify the feature flag was set
         let features: Vec<_> = cfg
             .interceptor_state()
             .load::<AwsSdkFeature>()
             .cloned()
             .collect();
-        assert_eq!(features.len(), 0);
+        assert_eq!(features.len(), 1, "Expected 1 feature, got: {:?}", features);
+        assert_eq!(features[0], AwsSdkFeature::EndpointOverride);
     }
 }

aws/rust-runtime/aws-runtime/src/observability_detection.rs

@@ -4,16 +4,20 @@
  */
 
 //! Observability feature detection for business metrics tracking
+//!
+//! This module provides an interceptor for detecting observability features in the AWS SDK:
+//!
+//! - [`ObservabilityDetectionInterceptor`]: Detects observability features during
+//!   request processing and tracks them for business metrics in the User-Agent header.
 
+use aws_smithy_observability_otel::meter::OtelMeterProvider;
 use aws_smithy_runtime::client::sdk_feature::SmithySdkFeature;
 use aws_smithy_runtime_api::box_error::BoxError;
 use aws_smithy_runtime_api::client::interceptors::context::BeforeTransmitInterceptorContextRef;
 use aws_smithy_runtime_api::client::interceptors::Intercept;
 use aws_smithy_runtime_api::client::runtime_components::RuntimeComponents;
 use aws_smithy_types::config_bag::ConfigBag;
 
-use crate::sdk_feature::AwsSdkFeature;
-
 /// Interceptor that detects when observability features are being used
 /// and tracks them for business metrics.
 #[derive(Debug, Default)]
@@ -40,32 +44,18 @@ impl Intercept for ObservabilityDetectionInterceptor {
     ) -> Result<(), BoxError> {
         // Try to get the global telemetry provider
         if let Ok(provider) = aws_smithy_observability::global::get_telemetry_provider() {
-            // Check if it's not a noop provider
-            if !provider.is_noop() {
+            // Use type-safe downcasting to detect OpenTelemetry meter provider
+            // This works with any ProvideMeter implementation and doesn't require
+            // implementation-specific boolean flags
+            if provider
+                .meter_provider()
+                .as_any()
+                .downcast_ref::<OtelMeterProvider>()
+                .is_some()
+            {
                 // Track that observability metrics are enabled
                 cfg.interceptor_state()
-                    .store_append(AwsSdkFeature::ObservabilityMetrics);
-
-                // PRAGMATIC APPROACH: Track tracing based on meter provider state
-                //
-                // The SDK uses Rust's `tracing` crate globally but doesn't have a
-                // configurable tracer provider yet. We make the reasonable assumption
-                // that if a user has configured a meter provider, they've also set up
-                // tracing as part of their observability stack.
-                //
-                // This is a pragmatic workaround until a proper tracer provider is added.
-                cfg.interceptor_state()
-                    .store_append(SmithySdkFeature::ObservabilityTracing);
-
-                // Check if it's using OpenTelemetry
-                if provider.is_otel() {
-                    cfg.interceptor_state()
-                        .store_append(AwsSdkFeature::ObservabilityOtelMetrics);
-
-                    // If using OpenTelemetry for metrics, likely using it for tracing too
-                    cfg.interceptor_state()
-                        .store_append(AwsSdkFeature::ObservabilityOtelTracing);
-                }
+                    .store_append(SmithySdkFeature::ObservabilityMetrics);
             }
         }
 
@@ -76,6 +66,7 @@ impl Intercept for ObservabilityDetectionInterceptor {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::sdk_feature::AwsSdkFeature;
     use aws_smithy_observability::TelemetryProvider;
     use aws_smithy_runtime_api::client::interceptors::context::{Input, InterceptorContext};
     use aws_smithy_runtime_api::client::orchestrator::HttpRequest;
@@ -102,12 +93,85 @@ mod tests {
             .read_before_signing(&ctx, &rc, &mut cfg)
             .unwrap();
 
-        // Should not track any features for noop provider
-        let features: Vec<_> = cfg.load::<AwsSdkFeature>().collect();
-        assert_eq!(features.len(), 0);
+        // Should not track any features for noop provider since it doesn't downcast to OtelMeterProvider
+        let smithy_features: Vec<_> = cfg.load::<SmithySdkFeature>().collect();
+        assert_eq!(smithy_features.len(), 0);
+
+        let aws_features: Vec<_> = cfg.load::<AwsSdkFeature>().collect();
+        assert_eq!(aws_features.len(), 0);
     }
 
-    // Note: We cannot easily test non-noop providers without creating a custom meter provider
-    // implementation, which would require exposing internal types. The noop test above
-    // is sufficient to verify the detection logic works correctly.
+    #[test]
+    fn test_custom_provider_not_detected_as_otel() {
+        use aws_smithy_observability::meter::{Meter, ProvideMeter};
+        use aws_smithy_observability::noop::NoopMeterProvider;
+        use aws_smithy_observability::Attributes;
+        use std::sync::Arc;
+
+        // Create a custom (non-OTel, non-noop) meter provider
+        // This simulates a user implementing their own metrics provider
+        #[derive(Debug)]
+        struct CustomMeterProvider {
+            inner: NoopMeterProvider,
+        }
+
+        impl ProvideMeter for CustomMeterProvider {
+            fn get_meter(&self, scope: &'static str, attributes: Option<&Attributes>) -> Meter {
+                // Delegate to noop for simplicity, but this is a distinct type
+                self.inner.get_meter(scope, attributes)
+            }
+
+            fn as_any(&self) -> &dyn std::any::Any {
+                self
+            }
+        }
+
+        let mut context = InterceptorContext::new(Input::doesnt_matter());
+        context.enter_serialization_phase();
+        context.set_request(HttpRequest::empty());
+        let _ = context.take_input();
+        context.enter_before_transmit_phase();
+
+        let rc = RuntimeComponentsBuilder::for_tests().build().unwrap();
+        let mut cfg = ConfigBag::base();
+
+        // Set the custom provider
+        let custom_provider = Arc::new(CustomMeterProvider {
+            inner: NoopMeterProvider,
+        });
+        let telemetry_provider = TelemetryProvider::builder()
+            .meter_provider(custom_provider)
+            .build();
+        let _ = aws_smithy_observability::global::set_telemetry_provider(telemetry_provider);
+
+        let interceptor = ObservabilityDetectionInterceptor::new();
+        let ctx = Into::into(&context);
+        interceptor
+            .read_before_signing(&ctx, &rc, &mut cfg)
+            .unwrap();
+
+        // Should NOT track any features for custom provider since it doesn't downcast to OtelMeterProvider
+        // The new implementation only emits metrics when OTel is detected
+        let smithy_features: Vec<_> = cfg
+            .interceptor_state()
+            .load::<SmithySdkFeature>()
+            .cloned()
+            .collect();
+        assert!(
+            !smithy_features.iter().any(|f| *f == SmithySdkFeature::ObservabilityMetrics),
+            "Should not detect custom provider as having observability metrics (only OTel is tracked)"
+        );
+
+        // Verify no AWS-specific features are tracked for custom provider
+        let aws_features: Vec<_> = cfg
+            .interceptor_state()
+            .load::<AwsSdkFeature>()
+            .cloned()
+            .collect();
+        assert_eq!(
+            aws_features.len(),
+            0,
+            "Should not track any AWS-specific features for custom provider"
+        );
+    }
 }