[Bug]: OTel ignores the invocation-local X-Ray header on Lambda Managed Instances

import java.lang.reflect.Proxy;
import java.time.Instant;
import java.util.*;
import com.amazonaws.services.lambda.runtime.Context;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.opentelemetry.sdk.trace.SdkTracerProvider;
import io.opentelemetry.sdk.trace.samplers.Sampler;
import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter;
import software.amazon.awssdk.utilslite.SdkInternalThreadLocal;
import software.amazon.lambda.durable.otel.*;
import software.amazon.lambda.durable.plugin.*;

public class LmiTraceCarrierProbe {
    static final String HEADER = "Root=1-6955b900-123456789012345678901234;Parent=1234567890123456;Sampled=0";
    static final String EXPECTED_TRACE = "6955b900123456789012345678901234";
    static final String PROPERTY = "com.amazonaws.xray.traceHeader";
    static final String KEY = "AWS_LAMBDA_X_TRACE_ID";

    public static void main(String[] args) throws Exception {
        if (System.getenv("_X_AMZN_TRACE_ID") != null) throw new AssertionError("Probe requires no environment header");
        Context runtimeContext = (Context) Proxy.newProxyInstance(Context.class.getClassLoader(), new Class<?>[]{Context.class},
            (proxy, method, values) -> {
                if (method.getName().equals("getXrayTraceId")) return HEADER;
                throw new UnsupportedOperationException(method.getName());
            });
        var results = new ArrayList<Map<String,Object>>();
        for (boolean ordinary : new boolean[]{false,true}) {
            System.clearProperty(PROPERTY);
            SdkInternalThreadLocal.remove(KEY);
            if (ordinary) System.setProperty(PROPERTY, HEADER);
            else SdkInternalThreadLocal.put(KEY, HEADER);
            var extracted = new XRayContextExtractor().extract();
            if (ordinary && (extracted == null || !EXPECTED_TRACE.equals(extracted.traceId()) || extracted.sampling() != ExtractedContext.Sampling.NOT_SAMPLED))
                throw new AssertionError("Ordinary carrier failed");
            if (!ordinary && (extracted != null || !HEADER.equals(runtimeContext.getXrayTraceId()) || !HEADER.equals(SdkInternalThreadLocal.get(KEY))))
                throw new AssertionError("LMI fixture or expected extractor result differs");
            for (boolean executionView : new boolean[]{true,false}) {
                var exporter = InMemorySpanExporter.create();
                var builder = SdkTracerProvider.builder().setSampler(Sampler.alwaysOn())
                    .addSpanProcessor(SimpleSpanProcessor.create(exporter));
                var config = OtelPluginConfig.builder().enableMdc(false).build();
                DurableExecutionPlugin plugin = executionView ? new ExecutionOtelPlugin(builder,config) : new InvocationOtelPlugin(builder,config);
                var arn="arn:aws:lambda:us-west-2:123456789012:function:probe:1/durable-execution/probe/id";
                plugin.onInvocationStart(new InvocationInfo("request",arn,true,Instant.parse("2026-10-02T00:00:00Z")));
                plugin.onInvocationEnd(new InvocationEndInfo("request",arn,true,InvocationStatus.SUCCEEDED,null));
                var spans = exporter.getFinishedSpanItems().stream().map(s -> Map.of(
                    "name",s.getName(),"traceId",s.getTraceId(),"sampled",s.getSpanContext().isSampled())).toList();
                if (ordinary && !spans.isEmpty()) throw new AssertionError("Upstream Sampled=0 must suppress exports");
                if (!ordinary && (spans.size()!=2 || exporter.getFinishedSpanItems().stream().anyMatch(s -> EXPECTED_TRACE.equals(s.getTraceId()) || !s.getSpanContext().isSampled())))
                    throw new AssertionError("Unexpected LMI-path observation");
                var record = new LinkedHashMap<String,Object>();
                record.put("carrier",ordinary?"ordinary Lambda system property":"LMI context plus runtime ThreadLocal");
                record.put("plugin",plugin.getClass().getSimpleName());
                record.put("inputSampling","NOT_SAMPLED");
                record.put("extractor",extracted==null?null:Map.of("traceId",extracted.traceId(),"sampling",extracted.sampling().name()));
                record.put("spans",spans);
                results.add(record);
            }
        }
        System.clearProperty(PROPERTY);
        SdkInternalThreadLocal.remove(KEY);
        System.out.println(new ObjectMapper().writeValueAsString(Map.of("sdkAndPluginVersion","2.2.1","results",results,"checksPassed",true)));
    }
}