fix: capture stack trace before async boundary

src/core/ProtobufCommunicator.ts

@@ -170,7 +170,7 @@ export class ProtobufCommunicator {
       requestId,
       tags: {},
       outboundSpan: cleanSpan,
-      stackTrace: this.getStackTrace(),
+      stackTrace: cleanSpan?.stackTrace,
     });
 
     const sdkMessage: SDKMessage = SDKMessage.create({
@@ -231,7 +231,7 @@ export class ProtobufCommunicator {
       requestId,
       tags: {},
       outboundSpan: cleanSpan,
-      stackTrace: this.getStackTrace(),
+      stackTrace: cleanSpan?.stackTrace,
     });
 
     const sdkMessage = SDKMessage.create({

src/core/types.ts

@@ -99,6 +99,7 @@ export type CleanSpanData = {
   };
   // sdk-specific
   isUsed?: boolean;
+  stackTrace?: string;
 };
 
 export type MockRequestData = {
@@ -111,6 +112,7 @@ export type MockRequestData = {
   submoduleName: string;
   inputValue: unknown;
   kind: SpanKind;
+  stackTrace?: string;
 };
 
 export type MetadataObject = {

src/instrumentation/core/utils/index.ts

@@ -2,3 +2,4 @@ export * from "./modeUtils";
 export * from "./selfExclusion";
 export * from "./modeUtils";
 export * from "./shimmerUtils";
+export * from "./stackTraceUtils";

src/instrumentation/core/utils/mockResponseUtils.ts

@@ -58,6 +58,7 @@ function convertMockRequestDataToCleanSpanData(
       code: StatusCode.OK,
       message: "OK",
     },
+    stackTrace: mockRequestData.stackTrace,
   };
 }
 

src/instrumentation/core/utils/stackTraceUtils.ts

@@ -0,0 +1,46 @@
+/**
+ * Helper functions for capturing stack traces in replay mode
+ *
+ * TODO: Consider using a structured format for stack frames:
+ * 
+ * {
+ *   "frames": [
+ *     {
+ *       "fileName": "file.js",
+ *       "lineNumber": 10,
+ *       "columnNumber": 20,
+ *       "functionName": "functionName"
+ *     }
+ *   ]
+ * }
+ *
+ * This would allow for more efficient matching and filtering of stack frames.
+ * It would also allow for more accurate stack trace reconstruction in replay mode.
+ */
+
+/**
+ * 
+ * @param excludeClassNames - Class names to exclude from the stack trace
+ * @returns The stack trace as a string
+ */
+export function captureStackTrace(excludeClassNames: string[] = []): string {
+    const originalStackTraceLimit = Error.stackTraceLimit; // Default is 10
+
+    Error.stackTraceLimit = 100;
+    const s = new Error().stack || "";
+    Error.stackTraceLimit = originalStackTraceLimit;
+
+    const defaultExcludes = [
+    "drift-node-sdk/src/instrumentation",
+    "drift-node-sdk/src/core",
+    "node_modules/@use-tusk",
+    ];
+
+    const allExcludes = [...defaultExcludes, ...excludeClassNames];
+
+    return s
+    .split("\n")
+    .slice(2) // Skip "Error" and capture method lines
+    .filter((l) => !allExcludes.some(exclude => l.includes(exclude)))
+    .join("\n");
+}

src/instrumentation/libraries/fetch/Instrumentation.ts

@@ -4,7 +4,7 @@ import { SpanInfo, SpanUtils } from "../../../core/tracing/SpanUtils";
 import { SpanKind, SpanStatusCode } from "@opentelemetry/api";
 import { TuskDriftCore, TuskDriftMode } from "../../../core/TuskDrift";
 import { getDecodedType, httpBodyEncoder } from "../http/utils";
-import { isTuskDriftIngestionUrl, TUSK_SKIP_HEADER } from "../../core/utils";
+import { captureStackTrace, isTuskDriftIngestionUrl, TUSK_SKIP_HEADER } from "../../core/utils";
 import { findMockResponseAsync } from "../../core/utils/mockResponseUtils";
 import { handleReplayMode, handleRecordMode } from "../../core/utils/modeUtils";
 import { FetchInputValue, FetchOutputValue, FetchInstrumentationConfig } from "./types";
@@ -60,7 +60,8 @@ export class FetchInstrumentation extends TdInstrumentationBase {
       input: string | URL | Request,
       init?: RequestInit,
     ): Promise<Response> {
-      return self._handleFetchRequest(input, init);
+      const stackTrace = captureStackTrace(["FetchInstrumentation"]);
+      return self._handleFetchRequest(input, init, stackTrace);
     } as typeof globalThis.fetch;
 
     logger.debug("Global fetch patching complete");
@@ -69,6 +70,7 @@ export class FetchInstrumentation extends TdInstrumentationBase {
   private async _handleFetchRequest(
     input: string | URL | Request,
     init?: RequestInit,
+    stackTrace?: string,
   ): Promise<Response> {
     // Parse request details
     const url = typeof input === "string" ? input : input instanceof URL ? input.href : input.url;
@@ -118,7 +120,7 @@ export class FetchInstrumentation extends TdInstrumentationBase {
               isPreAppStart: false,
             },
             (spanInfo) => {
-              return this._handleReplayFetch(inputValue, spanInfo);
+              return this._handleReplayFetch(inputValue, spanInfo, stackTrace);
             },
           );
         },
@@ -242,6 +244,7 @@ export class FetchInstrumentation extends TdInstrumentationBase {
   private async _handleReplayFetch(
     inputValue: FetchInputValue,
     spanInfo: SpanInfo,
+    stackTrace?: string,
   ): Promise<Response> {
     const mockData = await findMockResponseAsync({
       mockRequestData: {
@@ -254,6 +257,7 @@ export class FetchInstrumentation extends TdInstrumentationBase {
         submoduleName: inputValue.method,
         inputValue,
         kind: SpanKind.CLIENT,
+        stackTrace,
       },
       tuskDrift: this.tuskDrift,
       inputValueSchemaMerges: {

src/instrumentation/libraries/grpc/Instrumentation.ts

@@ -4,7 +4,7 @@ import { TdInstrumentationNodeModuleFile } from "../../core/baseClasses/TdInstru
 import { SpanUtils, SpanInfo } from "../../../core/tracing/SpanUtils";
 import { SpanKind, SpanStatusCode } from "@opentelemetry/api";
 import { TuskDriftCore, TuskDriftMode } from "../../../core/TuskDrift";
-import { wrap } from "../../core/utils";
+import { captureStackTrace, wrap } from "../../core/utils";
 import { findMockResponseAsync } from "../../core/utils/mockResponseUtils";
 import { handleRecordMode, handleReplayMode } from "../../core/utils/modeUtils";
 import { PackageType } from "@use-tusk/drift-schemas/core/span";
@@ -303,6 +303,8 @@ export class GrpcInstrumentation extends TdInstrumentationBase {
 
         // Handle replay mode
         if (self.mode === TuskDriftMode.REPLAY) {
+          const stackTrace = captureStackTrace(["GrpcInstrumentation"]);
+
           return handleReplayMode({
             replayModeHandler: () => {
               return SpanUtils.createAndExecuteSpan(
@@ -324,6 +326,7 @@ export class GrpcInstrumentation extends TdInstrumentationBase {
                     inputValue,
                     callback,
                     MetadataConstructor,
+                    stackTrace,
                   );
                 },
               );
@@ -525,6 +528,7 @@ export class GrpcInstrumentation extends TdInstrumentationBase {
     inputValue: GrpcClientInputValue,
     callback: Function,
     MetadataConstructor: any,
+    stackTrace?: string,
   ): Promise<any> {
     logger.debug(`[GrpcInstrumentation] Replaying gRPC unary request`);
 
@@ -539,6 +543,7 @@ export class GrpcInstrumentation extends TdInstrumentationBase {
         instrumentationName: this.INSTRUMENTATION_NAME,
         submoduleName: "client",
         kind: SpanKind.CLIENT,
+        stackTrace,
       },
       tuskDrift: this.tuskDrift,
     });

src/instrumentation/libraries/http/HttpReplayHooks.ts

@@ -31,12 +31,14 @@ export class HttpReplayHooks {
     protocol,
     args,
     spanInfo,
+    stackTrace,
   }: {
     method: string;
     requestOptions: RequestOptions;
     protocol: "http" | "https";
     args: any[];
     spanInfo: SpanInfo;
+    stackTrace?: string;
   }): TdMockClientRequest | undefined {
     logger.debug(
       `[HttpReplayHooks] Handling outbound ${protocol.toUpperCase()} ${method} request in replay mode`,
@@ -66,7 +68,7 @@ export class HttpReplayHooks {
     };
 
     // Create and return the mock client request
-    const mockRequest = new TdMockClientRequest(mockOptions, spanInfo, callback);
+    const mockRequest = new TdMockClientRequest(mockOptions, spanInfo, callback, stackTrace);
 
     // For GET/HEAD requests, automatically call .end() to match native http.get() behavior
     // Native http.get() and https.get() are convenience methods that automatically call .end()

src/instrumentation/libraries/http/Instrumentation.ts

@@ -29,6 +29,7 @@ import {
   handleRecordMode,
   handleReplayMode,
   isTuskDriftIngestionUrl,
+  captureStackTrace,
 } from "../../core/utils";
 import { PackageType, StatusCode } from "@use-tusk/drift-schemas/core/span";
 import {
@@ -1047,6 +1048,8 @@ export class HttpInstrumentation extends TdInstrumentationBase {
         );
 
         if (self.mode === TuskDriftMode.REPLAY) {
+          const stackTrace = captureStackTrace(["HttpInstrumentation"]);
+
           return handleReplayMode({
             replayModeHandler: () => {
               // Build input value object for replay mode
@@ -1082,6 +1085,7 @@ export class HttpInstrumentation extends TdInstrumentationBase {
                     protocol: requestProtocol,
                     args,
                     spanInfo,
+                    stackTrace,
                   });
                 },
               );

src/instrumentation/libraries/http/mocks/TdMockClientRequest.ts

@@ -42,6 +42,7 @@ export class TdMockClientRequest extends EventEmitter {
   public finished: boolean = false;
   private tuskDrift: TuskDriftCore;
   private spanInfo: SpanInfo;
+  private stackTrace?: string;
 
   private requestBodyBuffers: Buffer[] = [];
   private playbackStarted: boolean = false;
@@ -52,12 +53,14 @@ export class TdMockClientRequest extends EventEmitter {
     options: TdMockClientRequestOptions,
     spanInfo: SpanInfo,
     callback?: (res: IncomingMessage) => void,
+    stackTrace?: string,
   ) {
     super();
     TdMockClientRequest._setupPrototype();
     this.tuskDrift = TuskDriftCore.getInstance();
 
     this.spanInfo = spanInfo;
+    this.stackTrace = stackTrace;
 
     if (!options || Object.keys(options).length === 0) {
       throw new Error(
@@ -302,6 +305,7 @@ export class TdMockClientRequest extends EventEmitter {
           submoduleName: rawInputValue.method,
           inputValue,
           kind: SpanKind.CLIENT,
+          stackTrace: this.stackTrace,
         },
         tuskDrift: this.tuskDrift,
         inputValueSchemaMerges: {

src/instrumentation/libraries/ioredis/Instrumentation.ts

@@ -20,6 +20,7 @@ import {
 import { convertValueToJsonable } from "./utils";
 import { PackageType } from "@use-tusk/drift-schemas/core/span";
 import { logger } from "../../../core/utils/logger";
+import { captureStackTrace } from "src/instrumentation/core/utils";
 
 const SUPPORTED_VERSIONS = [">=4.11.0 <5", "5.*"];
 
@@ -175,6 +176,8 @@ export class IORedisInstrumentation extends TdInstrumentationBase {
 
         // Handle replay mode
         if (self.mode === TuskDriftMode.REPLAY) {
+          const stackTrace = captureStackTrace(["IORedisInstrumentation"]);
+
           return handleReplayMode({
             replayModeHandler: () => {
               return SpanUtils.createAndExecuteSpan(
@@ -191,7 +194,7 @@ export class IORedisInstrumentation extends TdInstrumentationBase {
                   isPreAppStart: false,
                 },
                 (spanInfo) => {
-                  return self._handleReplaySendCommand(spanInfo, cmd, inputValue, commandName);
+                  return self._handleReplaySendCommand(spanInfo, cmd, inputValue, commandName, stackTrace);
                 },
               );
             },
@@ -449,6 +452,7 @@ export class IORedisInstrumentation extends TdInstrumentationBase {
     cmd: IORedisCommand,
     inputValue: IORedisInputValue,
     commandName: string,
+    stackTrace?: string,
   ): Promise<any> {
     logger.debug(`[IORedisInstrumentation] Replaying IORedis command ${cmd.name}`);
 
@@ -462,6 +466,7 @@ export class IORedisInstrumentation extends TdInstrumentationBase {
         instrumentationName: this.INSTRUMENTATION_NAME,
         submoduleName: cmd.name,
         kind: SpanKind.CLIENT,
+        stackTrace,
       },
       tuskDrift: this.tuskDrift,
     });

src/instrumentation/libraries/jsonwebtoken/Instrumentation.ts

@@ -3,7 +3,7 @@ import { TdInstrumentationNodeModule } from "../../core/baseClasses/TdInstrument
 import { SpanUtils, SpanInfo } from "../../../core/tracing/SpanUtils";
 import { SpanKind, SpanStatusCode } from "@opentelemetry/api";
 import { TuskDriftCore, TuskDriftMode } from "../../../core/TuskDrift";
-import { wrap } from "../../core/utils";
+import { captureStackTrace, wrap } from "../../core/utils";
 import { findMockResponseAsync } from "../../core/utils/mockResponseUtils";
 import { handleRecordMode, handleReplayMode } from "../../core/utils/modeUtils";
 import {
@@ -116,6 +116,8 @@ export class JsonwebtokenInstrumentation extends TdInstrumentationBase {
 
         // Handle replay mode (only if app is ready)
         if (self.mode === TuskDriftMode.REPLAY) {
+          const stackTrace = captureStackTrace(["JsonwebtokenInstrumentation"]);
+
           return handleReplayMode({
             replayModeHandler: () => {
               // Create span in replay mode
@@ -132,7 +134,7 @@ export class JsonwebtokenInstrumentation extends TdInstrumentationBase {
                   isPreAppStart: false,
                 },
                 (spanInfo) => {
-                  return self.handleReplayVerify(verifyConfig, inputValue, spanInfo);
+                  return self.handleReplayVerify(verifyConfig, inputValue, spanInfo, stackTrace);
                 },
               );
             },
@@ -204,6 +206,8 @@ export class JsonwebtokenInstrumentation extends TdInstrumentationBase {
 
         // Handle replay mode (only if app is ready)
         if (self.mode === TuskDriftMode.REPLAY) {
+          const stackTrace = captureStackTrace(["JsonwebtokenInstrumentation"]);
+
           return handleReplayMode({
             replayModeHandler: () => {
               // Create span in replay mode
@@ -220,7 +224,7 @@ export class JsonwebtokenInstrumentation extends TdInstrumentationBase {
                   isPreAppStart: false,
                 },
                 (spanInfo) => {
-                  return self.handleReplaySign(signConfig, inputValue, spanInfo);
+                  return self.handleReplaySign(signConfig, inputValue, spanInfo, stackTrace);
                 },
               );
             },
@@ -500,6 +504,7 @@ export class JsonwebtokenInstrumentation extends TdInstrumentationBase {
     verifyConfig: VerifyQueryConfig,
     inputValue: JwtVerifyInputValue,
     spanInfo: SpanInfo,
+    stackTrace?: string,
   ): Promise<any> {
     logger.debug(`[JsonwebtokenInstrumentation] Replaying JWT verify`);
 
@@ -514,6 +519,7 @@ export class JsonwebtokenInstrumentation extends TdInstrumentationBase {
         submoduleName: "verify",
         inputValue: inputValue,
         kind: SpanKind.CLIENT,
+        stackTrace,
       },
       tuskDrift: this.tuskDrift,
     });
@@ -578,6 +584,7 @@ export class JsonwebtokenInstrumentation extends TdInstrumentationBase {
     signConfig: SignQueryConfig,
     inputValue: JwtSignInputValue,
     spanInfo: SpanInfo,
+    stackTrace?: string,
   ): Promise<any> {
     logger.debug(`[JsonwebtokenInstrumentation] Replaying JWT sign`);
 
@@ -592,6 +599,7 @@ export class JsonwebtokenInstrumentation extends TdInstrumentationBase {
         submoduleName: "sign",
         inputValue: inputValue,
         kind: SpanKind.CLIENT,
+        stackTrace,
       },
       tuskDrift: this.tuskDrift,
     });