dd-trace 5.102.0 → 5.104.0

package/packages/datadog-instrumentations/src/redis.js

@@ -13,6 +13,7 @@ const errorCh = channel('apm:redis:command:error')
 let createClientUrl
 let createClientName
 const instanceInfo = new WeakMap()
+const connectionInfoCache = new WeakMap()
 
 function wrapAddCommand (addCommand) {
   return function (command) {
@@ -20,10 +21,7 @@ function wrapAddCommand (addCommand) {
       return addCommand.apply(this, arguments)
     }
 
-    const name = command[0]
-    const args = command.slice(1)
-
-    const ctx = getStartCtx(this, name, args)
+    const ctx = getStartCtx(this, command[0], command, 1)
     return startCh.runStores(ctx, () => {
       const res = addCommand.apply(this, arguments)
 
@@ -136,22 +134,33 @@ addHook({ name: 'redis', versions: ['>=0.12 <2.6'] }, redis => {
   return redis
 })
 
-function getStartCtx (client, command, args) {
-  const { url, connectionName } = instanceInfo.get(client) || {}
+function getStartCtx (client, command, args, argsStartIndex) {
+  let cached = connectionInfoCache.get(client)
+  if (cached === undefined) {
+    const info = instanceInfo.get(client)
+    cached = {
+      connectionOptions:
+        client.connection_options || client.connection_option || client.connectionOption || info?.url,
+      connectionName: info?.connectionName,
+    }
+    connectionInfoCache.set(client, cached)
+  }
 
   return {
     db: client.selected_db,
     command,
     args,
-    connectionOptions: client.connection_options || client.connection_option || client.connectionOption || url,
-    connectionName,
+    argsStartIndex,
+    connectionOptions: cached.connectionOptions,
+    connectionName: cached.connectionName,
   }
 }
 
 function wrapCallback (finishCh, errorCh, ctx, callback) {
-  return shimmer.wrapFunction(callback, callback => function (err) {
+  if (typeof callback !== 'function') return callback
+  return function (err) {
     return finish(finishCh, errorCh, ctx, err, callback, this, arguments)
-  })
+  }
 }
 
 function finish (finishCh, errorCh, ctx, error, callback, thisArg, args) {

package/packages/datadog-instrumentations/src/router.js

@@ -112,8 +112,10 @@ function createWrapRouterMethod (name, compile) {
     }
   }
 
-  function wrapNext (req, next) {
-    return shimmer.wrapFunction(next, next => function (error) {
+  function wrapNext (req, originalNext) {
+    // Per layer dispatch, N per request. `shimmer.wrapCallback` preserves
+    // only `name` + `length`; see its JSDoc for the full contract.
+    return shimmer.wrapCallback(originalNext, next => function (error) {
       if (error && error !== 'route' && error !== 'router') {
         errorChannel.publish({ req, error })
       }

package/packages/datadog-instrumentations/src/vitest.js

@@ -19,6 +19,7 @@ const {
   collectTestOptimizationSummariesFromTraces,
   logAttemptToFixTestExecution,
   logTestOptimizationSummary,
+  getTestOptimizationRequestResults,
 } = require('../../dd-trace/src/plugins/util/test')
 const { addHook, channel } = require('./helpers/instrument')
 
@@ -43,6 +44,7 @@ const testSuiteErrorCh = channel('ci:vitest:test-suite:error')
 // test session hooks
 const testSessionStartCh = channel('ci:vitest:session:start')
 const testSessionFinishCh = channel('ci:vitest:session:finish')
+const testSessionConfigurationCh = channel('ci:vitest:session:configuration')
 const libraryConfigurationCh = channel('ci:vitest:library-configuration')
 const knownTestsCh = channel('ci:vitest:known-tests')
 const isEarlyFlakeDetectionFaultyCh = channel('ci:vitest:is-early-flake-detection-faulty')
@@ -74,6 +76,9 @@ let isRetryReasonEfd = false
 let isRetryReasonAttemptToFix = false
 const switchedStatuses = new WeakSet()
 const workerProcesses = new WeakSet()
+const mainProcessSetupPromises = new WeakMap()
+const coverageWrappedProviders = new WeakSet()
+const finishWrappedContexts = new WeakSet()
 let isFlakyTestRetriesEnabled = false
 let flakyTestRetriesCount = 0
 let isEarlyFlakeDetectionEnabled = false
@@ -134,6 +139,9 @@ function getProvidedContext () {
       _ddFlakyTestRetriesCount: flakyTestRetriesCount,
       _ddIsImpactedTestsEnabled: isImpactedTestsEnabled,
       _ddModifiedFiles: modifiedFiles,
+      _ddTestSessionId: testSessionId,
+      _ddTestModuleId: testModuleId,
+      _ddTestCommand: testCommand,
     } = globalThis.__vitest_worker__.providedContext
 
     return {
@@ -150,6 +158,9 @@ function getProvidedContext () {
       flakyTestRetriesCount: flakyTestRetriesCount ?? 0,
       isImpactedTestsEnabled,
       modifiedFiles,
+      testSessionId,
+      testModuleId,
+      testCommand,
     }
   } catch {
     log.error('Vitest workers could not parse provided context, so some features will not work.')
@@ -167,6 +178,9 @@ function getProvidedContext () {
       flakyTestRetriesCount: 0,
       isImpactedTestsEnabled: false,
       modifiedFiles: {},
+      testSessionId: undefined,
+      testModuleId: undefined,
+      testCommand: undefined,
     }
   }
 }
@@ -217,6 +231,10 @@ function getTestRunnerExport (testPackage) {
   return findExportByName(testPackage, 'VitestTestRunner') || findExportByName(testPackage, 'TestRunner')
 }
 
+function getVitestExport (vitestPackage) {
+  return findExportByName(vitestPackage, 'Vitest')
+}
+
 function getForksPoolWorkerExport (vitestPackage) {
   return findExportByName(vitestPackage, 'ForksPoolWorker')
 }
@@ -329,170 +347,223 @@ function wrapBeforeEachCleanupResult (task, result) {
   return result
 }
 
-function getSortWrapper (sort, frameworkVersion) {
-  return async function () {
-    if (!testSessionFinishCh.hasSubscribers) {
-      return sort.apply(this, arguments)
-    }
-    // There isn't any other async function that we seem to be able to hook into
-    // So we will use the sort from BaseSequencer. This means that a custom sequencer
-    // will not work. This will be a known limitation.
+function getWorkspaceProject (ctx) {
+  return ctx.getCoreWorkspaceProject
+    ? ctx.getCoreWorkspaceProject()
+    : ctx.getRootProject()
+}
+
+function setProvidedContext (ctx, values, warningMessage) {
+  try {
+    Object.assign(getWorkspaceProject(ctx)._provided, values)
+  } catch {
+    log.warn(warningMessage)
+  }
+}
+
+function getTestFilepathsFromSpecifications (testSpecifications) {
+  if (!Array.isArray(testSpecifications) || !testSpecifications.length) {
+    return
+  }
+
+  return testSpecifications.map(testSpecification => {
+    const testFile = Array.isArray(testSpecification) ? testSpecification[1] : testSpecification
+    return testFile?.moduleId || testFile?.filepath || testFile
+  })
+}
+
+function getTestFilepaths (ctx, testSpecifications) {
+  const testFilepaths = getTestFilepathsFromSpecifications(testSpecifications)
+  if (testFilepaths) {
+    return testFilepaths
+  }
+
+  const getFilePaths = ctx.getTestFilepaths || ctx._globTestFilepaths
+  return getFilePaths.call(ctx)
+}
+
+function wrapCoverageProvider (ctx) {
+  const { coverageProvider } = ctx
+  if (!coverageProvider?.generateCoverage || coverageWrappedProviders.has(coverageProvider)) {
+    return
+  }
+  coverageWrappedProviders.add(coverageProvider)
+
+  // Capture coverage root directory from config (default is 'coverage' in cwd)
+  try {
+    const coverageConfig = ctx.config?.coverage
+    const reportsDirectory = coverageConfig?.reportsDirectory || 'coverage'
+    const rootDir = ctx.config?.root || process.cwd()
+    coverageRootDir = path.isAbsolute(reportsDirectory) ? reportsDirectory : path.join(rootDir, reportsDirectory)
+  } catch {
+    // Fallback to cwd if we can't get config
+    coverageRootDir = process.cwd()
+  }
+
+  shimmer.wrap(coverageProvider, 'generateCoverage', generateCoverage => async function () {
+    const totalCodeCoverage = await generateCoverage.apply(this, arguments)
 
     try {
-      const { err, libraryConfig } = await getChannelPromise(libraryConfigurationCh, frameworkVersion)
-      if (!err) {
-        isFlakyTestRetriesEnabled = libraryConfig.isFlakyTestRetriesEnabled
-        flakyTestRetriesCount = libraryConfig.flakyTestRetriesCount
-        isEarlyFlakeDetectionEnabled = libraryConfig.isEarlyFlakeDetectionEnabled
-        earlyFlakeDetectionNumRetries = libraryConfig.earlyFlakeDetectionNumRetries
-        earlyFlakeDetectionSlowTestRetries = libraryConfig.earlyFlakeDetectionSlowTestRetries ?? {}
-        isDiEnabled = libraryConfig.isDiEnabled
-        isKnownTestsEnabled = libraryConfig.isKnownTestsEnabled
-        isTestManagementTestsEnabled = libraryConfig.isTestManagementEnabled
-        testManagementAttemptToFixRetries = libraryConfig.testManagementAttemptToFixRetries
-        isImpactedTestsEnabled = libraryConfig.isImpactedTestsEnabled
-      }
+      testCodeCoverageLinesTotal = totalCodeCoverage.getCoverageSummary().lines.pct
     } catch {
-      isFlakyTestRetriesEnabled = false
-      isEarlyFlakeDetectionEnabled = false
-      isDiEnabled = false
-      isKnownTestsEnabled = false
-      isImpactedTestsEnabled = false
+      // ignore errors
     }
+    return totalCodeCoverage
+  })
+}
 
-    if (isFlakyTestRetriesEnabled && !this.ctx.config.retry && flakyTestRetriesCount > 0) {
-      this.ctx.config.retry = flakyTestRetriesCount
-      try {
-        const workspaceProject = this.ctx.getCoreWorkspaceProject
-          ? this.ctx.getCoreWorkspaceProject()
-          : this.ctx.getRootProject()
-        workspaceProject._provided._ddIsFlakyTestRetriesEnabled = isFlakyTestRetriesEnabled
-        workspaceProject._provided._ddFlakyTestRetriesCount = flakyTestRetriesCount
-      } catch {
-        log.warn('Could not send library configuration to workers.')
-      }
-    }
+function wrapSessionFinish (ctx) {
+  if (finishWrappedContexts.has(ctx)) {
+    return
+  }
+  finishWrappedContexts.add(ctx)
 
-    if (isKnownTestsEnabled) {
-      const knownTestsResponse = await getChannelPromise(knownTestsCh)
-      if (knownTestsResponse.err) {
-        isEarlyFlakeDetectionEnabled = false
-      } else {
-        const knownTests = knownTestsResponse.knownTests
-        const getFilePaths = this.ctx.getTestFilepaths || this.ctx._globTestFilepaths
-
-        const testFilepaths = await getFilePaths.call(this.ctx)
-
-        if (isValidKnownTests(knownTests)) {
-          isEarlyFlakeDetectionFaultyCh.publish({
-            knownTests: knownTests.vitest,
-            testFilepaths,
-            onDone: (isFaulty) => {
-              isEarlyFlakeDetectionFaulty = isFaulty
-            },
-          })
-          if (isEarlyFlakeDetectionFaulty) {
-            isEarlyFlakeDetectionEnabled = false
-            log.warn('New test detection is disabled because the number of new tests is too high.')
-          } else {
-            // TODO: use this to pass session and module IDs to the worker, instead of polluting process.env
-            // Note: setting this.ctx.config.provide directly does not work because it's cached
-            try {
-              const workspaceProject = this.ctx.getCoreWorkspaceProject
-                ? this.ctx.getCoreWorkspaceProject()
-                : this.ctx.getRootProject()
-              workspaceProject._provided._ddIsKnownTestsEnabled = isKnownTestsEnabled
-              workspaceProject._provided._ddKnownTests = knownTests
-              workspaceProject._provided._ddIsEarlyFlakeDetectionEnabled = isEarlyFlakeDetectionEnabled
-              workspaceProject._provided._ddEarlyFlakeDetectionNumRetries =
-                getConfiguredEfdRetryCount(earlyFlakeDetectionSlowTestRetries, earlyFlakeDetectionNumRetries)
-              workspaceProject._provided._ddEarlyFlakeDetectionSlowTestRetries = earlyFlakeDetectionSlowTestRetries
-            } catch {
-              log.warn('Could not send known tests to workers so Early Flake Detection will not work.')
-            }
-          }
-        } else {
-          isEarlyFlakeDetectionFaulty = true
-          isEarlyFlakeDetectionEnabled = false
-        }
-      }
-    }
+  shimmer.wrap(ctx, 'exit', getFinishWrapper)
+  shimmer.wrap(ctx, 'close', getFinishWrapper)
+}
 
-    if (isDiEnabled) {
-      try {
-        const workspaceProject = this.ctx.getCoreWorkspaceProject
-          ? this.ctx.getCoreWorkspaceProject()
-          : this.ctx.getRootProject()
-        workspaceProject._provided._ddIsDiEnabled = isDiEnabled
-      } catch {
-        log.warn('Could not send Dynamic Instrumentation configuration to workers.')
-      }
-    }
+async function runMainProcessSetup (ctx, frameworkVersion, testSpecifications) {
+  if (!testSessionFinishCh.hasSubscribers) {
+    return
+  }
 
-    if (isTestManagementTestsEnabled) {
-      const { err, testManagementTests: receivedTestManagementTests } = await getChannelPromise(testManagementTestsCh)
-      if (err) {
-        isTestManagementTestsEnabled = false
-        log.error('Could not get test management tests.')
-      } else {
-        const testManagementTests = receivedTestManagementTests
-        try {
-          const workspaceProject = this.ctx.getCoreWorkspaceProject
-            ? this.ctx.getCoreWorkspaceProject()
-            : this.ctx.getRootProject()
-          workspaceProject._provided._ddIsTestManagementTestsEnabled = isTestManagementTestsEnabled
-          workspaceProject._provided._ddTestManagementAttemptToFixRetries = testManagementAttemptToFixRetries
-          workspaceProject._provided._ddTestManagementTests = testManagementTests
-        } catch {
-          log.warn('Could not send test management tests to workers so Test Management will not work.')
-        }
-      }
+  try {
+    const { err, libraryConfig } = await getChannelPromise(libraryConfigurationCh, frameworkVersion)
+    if (!err) {
+      isFlakyTestRetriesEnabled = libraryConfig.isFlakyTestRetriesEnabled
+      flakyTestRetriesCount = libraryConfig.flakyTestRetriesCount
+      isEarlyFlakeDetectionEnabled = libraryConfig.isEarlyFlakeDetectionEnabled
+      earlyFlakeDetectionNumRetries = libraryConfig.earlyFlakeDetectionNumRetries
+      earlyFlakeDetectionSlowTestRetries = libraryConfig.earlyFlakeDetectionSlowTestRetries ?? {}
+      isDiEnabled = libraryConfig.isDiEnabled
+      isKnownTestsEnabled = libraryConfig.isKnownTestsEnabled
+      isTestManagementTestsEnabled = libraryConfig.isTestManagementEnabled
+      testManagementAttemptToFixRetries = libraryConfig.testManagementAttemptToFixRetries
+      isImpactedTestsEnabled = libraryConfig.isImpactedTestsEnabled
     }
+  } catch {
+    isFlakyTestRetriesEnabled = false
+    isEarlyFlakeDetectionEnabled = false
+    isDiEnabled = false
+    isKnownTestsEnabled = false
+    isImpactedTestsEnabled = false
+  }
 
-    if (isImpactedTestsEnabled) {
-      const { err, modifiedFiles } = await getChannelPromise(modifiedFilesCh)
-      if (err) {
-        log.error('Could not get modified tests.')
-      } else {
-        try {
-          const workspaceProject = this.ctx.getCoreWorkspaceProject
-            ? this.ctx.getCoreWorkspaceProject()
-            : this.ctx.getRootProject()
-          workspaceProject._provided._ddIsImpactedTestsEnabled = isImpactedTestsEnabled
-          workspaceProject._provided._ddModifiedFiles = modifiedFiles
-        } catch {
-          log.warn('Could not send modified tests to workers so Impacted Tests will not work.')
+  if (testSessionConfigurationCh.hasSubscribers) {
+    const { testSessionId, testModuleId, testCommand } = await getChannelPromise(
+      testSessionConfigurationCh,
+      frameworkVersion
+    )
+    setProvidedContext(ctx, {
+      _ddTestSessionId: testSessionId,
+      _ddTestModuleId: testModuleId,
+      _ddTestCommand: testCommand,
+    }, 'Could not send test session configuration to workers.')
+  }
+
+  const {
+    knownTestsResponse,
+    testManagementTestsResponse,
+  } = await getTestOptimizationRequestResults({
+    isKnownTestsEnabled,
+    isTestManagementTestsEnabled,
+    getKnownTests: () => getChannelPromise(knownTestsCh),
+    getTestManagementTests: () => getChannelPromise(testManagementTestsCh),
+  })
+
+  if (isFlakyTestRetriesEnabled && !ctx.config.retry && flakyTestRetriesCount > 0) {
+    ctx.config.retry = flakyTestRetriesCount
+    setProvidedContext(ctx, {
+      _ddIsFlakyTestRetriesEnabled: isFlakyTestRetriesEnabled,
+      _ddFlakyTestRetriesCount: flakyTestRetriesCount,
+    }, 'Could not send library configuration to workers.')
+  }
+
+  if (isKnownTestsEnabled) {
+    const currentKnownTestsResponse = knownTestsResponse || await getChannelPromise(knownTestsCh)
+    if (currentKnownTestsResponse.err) {
+      isEarlyFlakeDetectionEnabled = false
+    } else {
+      const knownTests = currentKnownTestsResponse.knownTests
+      const testFilepaths = await getTestFilepaths(ctx, testSpecifications)
+
+      if (isValidKnownTests(knownTests)) {
+        isEarlyFlakeDetectionFaultyCh.publish({
+          knownTests: knownTests.vitest,
+          testFilepaths,
+          onDone: (isFaulty) => {
+            isEarlyFlakeDetectionFaulty = isFaulty
+          },
+        })
+        if (isEarlyFlakeDetectionFaulty) {
+          isEarlyFlakeDetectionEnabled = false
+          log.warn('New test detection is disabled because the number of new tests is too high.')
+        } else {
+          setProvidedContext(ctx, {
+            _ddIsKnownTestsEnabled: isKnownTestsEnabled,
+            _ddKnownTests: knownTests,
+            _ddIsEarlyFlakeDetectionEnabled: isEarlyFlakeDetectionEnabled,
+            _ddEarlyFlakeDetectionNumRetries:
+              getConfiguredEfdRetryCount(earlyFlakeDetectionSlowTestRetries, earlyFlakeDetectionNumRetries),
+            _ddEarlyFlakeDetectionSlowTestRetries: earlyFlakeDetectionSlowTestRetries,
+          }, 'Could not send known tests to workers so Early Flake Detection will not work.')
         }
+      } else {
+        isEarlyFlakeDetectionFaulty = true
+        isEarlyFlakeDetectionEnabled = false
       }
     }
+  }
 
-    if (this.ctx.coverageProvider?.generateCoverage) {
-      // Capture coverage root directory from config (default is 'coverage' in cwd)
-      try {
-        const coverageConfig = this.ctx.config?.coverage
-        const reportsDirectory = coverageConfig?.reportsDirectory || 'coverage'
-        const rootDir = this.ctx.config?.root || process.cwd()
-        coverageRootDir = path.isAbsolute(reportsDirectory) ? reportsDirectory : path.join(rootDir, reportsDirectory)
-      } catch {
-        // Fallback to cwd if we can't get config
-        coverageRootDir = process.cwd()
-      }
+  if (isDiEnabled) {
+    setProvidedContext(ctx, {
+      _ddIsDiEnabled: isDiEnabled,
+    }, 'Could not send Dynamic Instrumentation configuration to workers.')
+  }
 
-      shimmer.wrap(this.ctx.coverageProvider, 'generateCoverage', generateCoverage => async function () {
-        const totalCodeCoverage = await generateCoverage.apply(this, arguments)
+  if (isTestManagementTestsEnabled) {
+    const { err, testManagementTests: receivedTestManagementTests } =
+      testManagementTestsResponse || await getChannelPromise(testManagementTestsCh)
+    if (err) {
+      isTestManagementTestsEnabled = false
+      log.error('Could not get test management tests.')
+    } else {
+      setProvidedContext(ctx, {
+        _ddIsTestManagementTestsEnabled: isTestManagementTestsEnabled,
+        _ddTestManagementAttemptToFixRetries: testManagementAttemptToFixRetries,
+        _ddTestManagementTests: receivedTestManagementTests,
+      }, 'Could not send test management tests to workers so Test Management will not work.')
+    }
+  }
 
-        try {
-          testCodeCoverageLinesTotal = totalCodeCoverage.getCoverageSummary().lines.pct
-        } catch {
-          // ignore errors
-        }
-        return totalCodeCoverage
-      })
+  if (isImpactedTestsEnabled) {
+    const { err, modifiedFiles } = await getChannelPromise(modifiedFilesCh)
+    if (err) {
+      log.error('Could not get modified tests.')
+    } else {
+      setProvidedContext(ctx, {
+        _ddIsImpactedTestsEnabled: isImpactedTestsEnabled,
+        _ddModifiedFiles: modifiedFiles,
+      }, 'Could not send modified tests to workers so Impacted Tests will not work.')
     }
+  }
 
-    shimmer.wrap(this.ctx, 'exit', getFinishWrapper)
-    shimmer.wrap(this.ctx, 'close', getFinishWrapper)
+  wrapCoverageProvider(ctx)
+  wrapSessionFinish(ctx)
+}
 
+function ensureMainProcessSetup (ctx, frameworkVersion, testSpecifications) {
+  let setupPromise = mainProcessSetupPromises.get(ctx)
+  if (!setupPromise) {
+    setupPromise = runMainProcessSetup(ctx, frameworkVersion, testSpecifications)
+    mainProcessSetupPromises.set(ctx, setupPromise)
+  }
+  return setupPromise
+}
+
+function getSortWrapper (sort, frameworkVersion) {
+  return async function () {
+    await ensureMainProcessSetup(this.ctx, frameworkVersion, arguments[0])
     return sort.apply(this, arguments)
   }
 }
@@ -504,6 +575,11 @@ function getFinishWrapper (exitOrClose) {
       return exitOrClose.apply(this, arguments)
     }
     isClosed = true
+
+    if (!testSessionFinishCh.hasSubscribers) {
+      return exitOrClose.apply(this, arguments)
+    }
+
     let onFinish
 
     const flushPromise = new Promise(resolve => {
@@ -555,6 +631,17 @@ function getCliOrStartVitestWrapper (frameworkVersion) {
   }
 }
 
+function wrapVitestRunFiles (Vitest, frameworkVersion) {
+  if (!Vitest?.prototype?.runFiles) {
+    return
+  }
+
+  shimmer.wrap(Vitest.prototype, 'runFiles', runFiles => async function (testSpecifications) {
+    await ensureMainProcessSetup(this, frameworkVersion, testSpecifications)
+    return runFiles.apply(this, arguments)
+  })
+}
+
 function getCreateCliWrapper (vitestPackage, frameworkVersion) {
   const createCliExport = findExportByName(vitestPackage, 'createCLI')
   if (!createCliExport) {
@@ -658,6 +745,11 @@ function getStartVitestWrapper (cliApiPackage, frameworkVersion) {
   const startVitestExport = findExportByName(cliApiPackage, 'startVitest')
   shimmer.wrap(cliApiPackage, startVitestExport.key, getCliOrStartVitestWrapper(frameworkVersion))
 
+  const vitest = getVitestExport(cliApiPackage)
+  if (vitest) {
+    wrapVitestRunFiles(vitest.value, frameworkVersion)
+  }
+
   const forksPoolWorker = getForksPoolWorkerExport(cliApiPackage)
   if (forksPoolWorker) {
     // function is async
@@ -921,7 +1013,6 @@ function wrapVitestTestRunner (VitestTestRunner) {
       // Here we finish the earlier iteration,
       // as long as it's not the _last_ iteration (which will be finished normally)
 
-      // TODO: check test duration (not to repeat if it's too slow)
       const ctx = taskToCtx.get(task)
       if (ctx) {
         if (lastExecutionStatus === 'fail') {
@@ -1269,8 +1360,15 @@ addHook({
     }
     // From >=3.0.1, the first arguments changes from a string to an object containing the filepath
     const testSuiteAbsolutePath = testPaths[0]?.filepath || testPaths[0]
-
-    const testSuiteCtx = { testSuiteAbsolutePath, frameworkVersion }
+    const providedContext = getProvidedContext()
+
+    const testSuiteCtx = {
+      testSuiteAbsolutePath,
+      frameworkVersion,
+      testSessionId: providedContext.testSessionId,
+      testModuleId: providedContext.testModuleId,
+      testCommand: providedContext.testCommand,
+    }
     testSuiteStartCh.runStores(testSuiteCtx, () => {})
     const startTestsResponse = await startTests.apply(this, arguments)
 
@@ -1288,7 +1386,6 @@ addHook({
       // We have to trick vitest into thinking that the test has passed
       // but we want to report it as failed if it did fail
       const isSwitchedStatus = switchedStatuses.has(task)
-      const providedContext = getProvidedContext()
 
       if (result) {
         const { state, duration, errors } = result

package/packages/datadog-plugin-apollo/src/gateway/request.js

@@ -2,8 +2,7 @@
 
 const { storage } = require('../../../datadog-core')
 const ApolloBasePlugin = require('../../../dd-trace/src/plugins/apollo')
-
-let tools
+const { getSignature } = require('../../../datadog-plugin-graphql/src/utils')
 
 const OPERATION_DEFINITION = 'OperationDefinition'
 const FRAGMENT_DEFINITION = 'FragmentDefinition'
@@ -100,23 +99,4 @@ function buildOperationContext (schema, operationDocument, operationName) {
   }
 }
 
-function getSignature (document, operationName, operationType, calculate) {
-  if (calculate !== false && tools !== false) {
-    try {
-      try {
-        tools = tools || require('../../../datadog-plugin-graphql/src/tools')
-      } catch (e) {
-        tools = false
-        throw e
-      }
-
-      return tools.defaultEngineReportingSignature(document, operationName)
-    } catch {
-      // safety net
-    }
-  }
-
-  return [operationType, operationName].filter(Boolean).join(' ')
-}
-
 module.exports = ApolloGatewayRequestPlugin