From 01605330b17fc9aff0121d8e264d717b094e52e1 Mon Sep 17 00:00:00 2001
From: Thomas Watson <thomas.watson@datadoghq.com>
Date: Mon, 29 Jun 2026 12:08:39 +0200
Subject: [PATCH] Fix Live Debugger returned arrow instrumentation

When a block-bodied function returned an anonymous expression-bodied arrow,
Live Debugger could queue the outer return wrapper at the same byte offset as
the inner arrow suffix. MagicString emitted the outer tail inside the inner
arrow instrumentation, producing invalid JavaScript.

Anchor the return-wrapper tail on the right side of the shared offset so the
inner arrow suffix closes first. Add regression coverage for returned arrows,
including JSX mapping, while asserting the output remains instrumented.
---
 .../live-debugger/src/transform/index.test.ts | 29 +++++++++++++++++++
 .../live-debugger/src/transform/index.ts      |  7 ++---
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/packages/plugins/live-debugger/src/transform/index.test.ts b/packages/plugins/live-debugger/src/transform/index.test.ts
index 13bf98558..8466e5ada 100644
--- a/packages/plugins/live-debugger/src/transform/index.test.ts
+++ b/packages/plugins/live-debugger/src/transform/index.test.ts
@@ -358,6 +358,34 @@ describe('transformCode', () => {
                 expectedInstrumentedCount: 1,
                 expectedTotalFunctions: 2,
             },
+            {
+                description: 'block-body arrow returning an expression-body arrow',
+                code: 'export const make = (cond) => { return () => cond ? a() : b(); };',
+                namedOnly: false,
+                expectedInstrumentedCount: 2,
+                expectedTotalFunctions: 2,
+            },
+            {
+                description: 'function declaration returning an expression-body arrow',
+                code: 'export function build(x) { const y = x + 1; return () => y > 0 ? pos(y) : neg(y); }',
+                namedOnly: false,
+                expectedInstrumentedCount: 2,
+                expectedTotalFunctions: 2,
+            },
+            {
+                description: 'block-body arrow returning an expression-body arrow with JSX mapping',
+                code: 'type Row = { id: string }; export const C = (rows: Row[]) => { const n = rows.length; return () => rows.map((r) => <li key={r.id}>{n}</li>); };',
+                namedOnly: false,
+                expectedInstrumentedCount: 3,
+                expectedTotalFunctions: 3,
+            },
+            {
+                description: 'block-body arrow returning a block-body arrow',
+                code: 'export const make = (cond) => { return () => { return cond ? a() : b(); }; };',
+                namedOnly: false,
+                expectedInstrumentedCount: 2,
+                expectedTotalFunctions: 2,
+            },
         ];
 
         test.each(nestedSyntaxCases)(
@@ -375,6 +403,7 @@ describe('transformCode', () => {
                 expect(result.instrumentedCount).toBe(expectedInstrumentedCount);
                 expect(result.totalFunctions).toBe(expectedTotalFunctions);
                 expect(probeMatches?.length).toBe(expectedInstrumentedCount);
+                expect(result.code).toContain('$dd_return');
                 expect(result.code).toContain('catch(e)');
             },
         );
diff --git a/packages/plugins/live-debugger/src/transform/index.ts b/packages/plugins/live-debugger/src/transform/index.ts
index 42d540728..df69f69ee 100644
--- a/packages/plugins/live-debugger/src/transform/index.ts
+++ b/packages/plugins/live-debugger/src/transform/index.ts
@@ -419,9 +419,8 @@ export function transformCode(options: TransformOptions): TransformResult {
         superCallRangesHandledByExpressionBodies,
     );
 
-    // Process inner (deeper) functions before outer ones so that MagicString
-    // appendLeft calls at shared positions (e.g. where an outer return wraps
-    // an inner arrow function) stack in the correct order.
+    // Process inner (deeper) functions before outer ones so same-side
+    // MagicString insertions at shared positions stack in the correct order.
     for (let i = targets.length - 1; i >= 0; i--) {
         injectInstrumentation(s, code, targets[i]);
     }
@@ -566,7 +565,7 @@ function injectInstrumentation(s: MagicStringType, code: string, target: Functio
                 const assignmentSuffix = ret.hasSequenceExpressionArgument ? ')' : '';
 
                 s.appendLeft(ret.argStart, assignmentPrefix);
-                s.appendLeft(
+                s.appendRight(
                     ret.argEnd,
                     `${assignmentSuffix}, ${probeVarName} ? $dd_return(${probeVarName}, ${rvVarName}, ${receiverArg}${returnCaptureArgs}) : ${rvVarName})`,
                 );