Fix Exception Replay in Lambda

dd-java-agent/agent-debugger/src/main/java/com/datadog/debugger/exception/DefaultExceptionDebugger.java

@@ -156,6 +156,13 @@ private static void processSnapshotsAndSetTags(
       int[] mapping = createThrowableMapping(currentEx, t);
       StackTraceElement[] innerTrace = currentEx.getStackTrace();
       int currentIdx = innerTrace.length - snapshot.getStack().size();
+
+      if (currentIdx < 0) {
+        // This means the innerTrace was truncated by the underlying environment.
+        // This is known to happen in AWS Lambda, but may also happen elsewhere.
+        currentIdx = i;
+      }
+
       if (!sanityCheckSnapshotAssignment(snapshot, innerTrace, currentIdx)) {
         continue;
       }

dd-java-agent/agent-debugger/src/test/java/com/datadog/debugger/exception/DefaultExceptionDebuggerTest.java

@@ -30,6 +30,8 @@
 import datadog.trace.bootstrap.debugger.CapturedContext;
 import datadog.trace.bootstrap.debugger.CapturedStackFrame;
 import datadog.trace.bootstrap.debugger.MethodLocation;
+import datadog.trace.bootstrap.debugger.ProbeImplementation;
+import datadog.trace.bootstrap.debugger.ProbeLocation;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
 import java.io.BufferedReader;
 import java.io.IOException;
@@ -280,6 +282,77 @@ public void filteringOutErrors() {
     verify(manager, times(0)).isAlreadyInstrumented(any());
   }
 
+  @Test
+  public void lambdaTruncatedInnerTraceFallback() {
+    AgentSpan span = mock(AgentSpan.class);
+    doAnswer(this::recordTags).when(span).setTag(anyString(), anyString());
+    when(span.getTag(anyString())).thenAnswer(inv -> spanTags.get(inv.getArgument(0)));
+    when(span.getTags()).thenReturn(spanTags);
+
+    // Create an exception with a real truncated stack trace from Lambda
+    RuntimeException lambdaException =
+        new RuntimeException("lambda") {
+          @Override
+          public StackTraceElement[] getStackTrace() {
+            return new StackTraceElement[] {
+              new StackTraceElement("Main", "handleRequest", "Main.java", 11),
+              new StackTraceElement(
+                  "jdk.internal.reflect.DirectMethodHandleAccessor",
+                  "invoke",
+                  "Unknown Source",
+                  -1),
+              new StackTraceElement("java.lang.reflect.Method", "invoke", "Unknown Source", -1)
+            };
+          }
+        };
+
+    // Set up the snapshot with a longer stack to represent original data
+    List<CapturedStackFrame> snapshotStack = new ArrayList<>();
+    snapshotStack.add(
+        CapturedStackFrame.from(new StackTraceElement("Main", "handleRequest", "Main.java", 11)));
+    for (int i = 0; i < 5; i++) {
+      snapshotStack.add(
+          CapturedStackFrame.from(
+              new StackTraceElement("Lambda.Frame" + i, "method", "Lambda.java", 100 + i)));
+    }
+
+    // Mock snapshot
+    Snapshot mockSnapshot = mock(Snapshot.class);
+    when(mockSnapshot.getId()).thenReturn("test-snapshot-id");
+    when(mockSnapshot.getStack()).thenReturn(snapshotStack);
+    when(mockSnapshot.getChainedExceptionIdx()).thenReturn(0);
+
+    // Mock probe
+    ProbeLocation mockLocation = mock(ProbeLocation.class);
+    when(mockLocation.getType()).thenReturn("Main");
+    when(mockLocation.getMethod()).thenReturn("handleRequest");
+    ProbeImplementation mockProbe = mock(ProbeImplementation.class);
+    when(mockProbe.getLocation()).thenReturn(mockLocation);
+    when(mockSnapshot.getProbe()).thenReturn(mockProbe);
+
+    // Mock exception state
+    ExceptionProbeManager.ThrowableState state = mock(ExceptionProbeManager.ThrowableState.class);
+    when(state.getSnapshots()).thenReturn(singletonList(mockSnapshot));
+    when(state.getExceptionId()).thenReturn("test-exception-id");
+    when(state.isSnapshotSent()).thenReturn(false);
+
+    // Create mock manager that returns our state
+    ExceptionProbeManager mockManager = mock(ExceptionProbeManager.class);
+    when(mockManager.isAlreadyInstrumented(anyString())).thenReturn(true);
+    when(mockManager.getStateByThrowable(lambdaException)).thenReturn(state);
+
+    DefaultExceptionDebugger testDebugger =
+        new DefaultExceptionDebugger(mockManager, configurationUpdater, classNameFiltering, 100);
+
+    // Test
+    testDebugger.handleException(lambdaException, span);
+
+    // Verify
+    String tagName = String.format(SNAPSHOT_ID_TAG_FMT, 0);
+    assertTrue(spanTags.containsKey(tagName));
+    assertEquals("test-snapshot-id", spanTags.get(tagName));
+  }
+
   private Object recordTags(InvocationOnMock invocationOnMock) {
     Object[] args = invocationOnMock.getArguments();
     String key = (String) args[0];