From 50ff63f839a92f6c39ae3349dffc9012633c02a8 Mon Sep 17 00:00:00 2001
From: Ray Chen <ray@raychen.io>
Date: Wed, 18 Mar 2026 01:05:10 +0800
Subject: [PATCH] Fix async method tracing

---
 packages/container/src/helpers.test.ts | 29 ++++++++++++++++++++++++++
 packages/container/src/helpers.ts      | 17 +++++++++------
 2 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/packages/container/src/helpers.test.ts b/packages/container/src/helpers.test.ts
index dd8f141..2fb6fe5 100644
--- a/packages/container/src/helpers.test.ts
+++ b/packages/container/src/helpers.test.ts
@@ -219,4 +219,33 @@ describe("wrapWithTracing", () => {
     instance.childMethod();
     expect(trace.mock.calls[0][0]).toBe("Child.childMethod");
   });
+
+  it("passes an async callback to trace for async methods", async () => {
+    class WithAsync extends StopClass {
+      async fetchData() {
+        return "data";
+      }
+      syncMethod() {
+        return "sync";
+      }
+    }
+    const callbacks: Array<{ name: string; isAsync: boolean }> = [];
+    const trace = vi.fn((name: string, fn: () => unknown) => {
+      callbacks.push({
+        name,
+        isAsync: fn.constructor.name === "AsyncFunction",
+      });
+      return fn();
+    });
+    const instance = new WithAsync();
+    wrapWithTracing(instance, trace, StopClass.prototype);
+
+    instance.syncMethod();
+    await instance.fetchData();
+
+    const syncCall = callbacks.find((c) => c.name === "WithAsync.syncMethod");
+    const asyncCall = callbacks.find((c) => c.name === "WithAsync.fetchData");
+    expect(syncCall?.isAsync).toBe(false);
+    expect(asyncCall?.isAsync).toBe(true);
+  });
 });
diff --git a/packages/container/src/helpers.ts b/packages/container/src/helpers.ts
index 2c00092..a42d467 100644
--- a/packages/container/src/helpers.ts
+++ b/packages/container/src/helpers.ts
@@ -66,13 +66,18 @@ export function wrapWithTracing(
 
       const original = desc.value as (...args: unknown[]) => unknown;
       const spanName = `${className}.${key}`;
+      const isAsync = original.constructor.name === "AsyncFunction";
 
-      (instance as Record<string, unknown>)[key] = function (
-        this: unknown,
-        ...args: unknown[]
-      ) {
-        return trace(spanName, () => original.apply(this, args));
-      };
+      // For async methods, pass an async callback to trace() so tracing
+      // libraries (e.g., DD APM) that detect async functions keep the
+      // span open until the promise settles.
+      (instance as Record<string, unknown>)[key] = isAsync
+        ? function (this: unknown, ...args: unknown[]) {
+            return trace(spanName, async () => original.apply(this, args));
+          }
+        : function (this: unknown, ...args: unknown[]) {
+            return trace(spanName, () => original.apply(this, args));
+          };
     }
     proto = Object.getPrototypeOf(proto);
   }