From 02688205639a64622abe96f56bc843d307376fea Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Tue, 18 Jun 2024 12:12:31 +0200
Subject: [PATCH 01/26] add vitest support wip

---
 .../src/helpers/hooks.js                      |   3 +
 .../datadog-instrumentations/src/vitest.js    | 174 ++++++++++++++++++
 packages/datadog-plugin-vitest/src/index.js   |  85 +++++++++
 packages/dd-trace/src/plugins/index.js        |   3 +
 4 files changed, 265 insertions(+)
 create mode 100644 packages/datadog-instrumentations/src/vitest.js
 create mode 100644 packages/datadog-plugin-vitest/src/index.js

diff --git a/packages/datadog-instrumentations/src/helpers/hooks.js b/packages/datadog-instrumentations/src/helpers/hooks.js
index 0723ceabd84..61f514f1232 100644
--- a/packages/datadog-instrumentations/src/helpers/hooks.js
+++ b/packages/datadog-instrumentations/src/helpers/hooks.js
@@ -71,6 +71,9 @@ module.exports = {
   'microgateway-core': () => require('../microgateway-core'),
   mocha: () => require('../mocha'),
   'mocha-each': () => require('../mocha'),
+  vitest: () => require('../vitest'),
+  '@vitest': () => require('../vitest'),
+  '@vitest/runner': () => require('../vitest'),
   workerpool: () => require('../mocha'),
   moleculer: () => require('../moleculer'),
   mongodb: () => require('../mongodb'),
diff --git a/packages/datadog-instrumentations/src/vitest.js b/packages/datadog-instrumentations/src/vitest.js
new file mode 100644
index 00000000000..e79ea866fb1
--- /dev/null
+++ b/packages/datadog-instrumentations/src/vitest.js
@@ -0,0 +1,174 @@
+const { addHook, channel, AsyncResource } = require('./helpers/instrument')
+const { isMainThread } = require('node:worker_threads')
+const shimmer = require('../../datadog-shimmer')
+
+// Needs modifying node_modules/.pnpm/import-in-the-middle@1.7.3/node_modules/import-in-the-middle/index.js
+// the specifiers thing I don't understand
+
+// needs creating a fake file node_modules/vitest/dist/@vitest/spy -- why????
+
+const testStartCh = channel('ci:vitest:test:start')
+const testFinishCh = channel('ci:vitest:test:finish')
+
+const onAfterRunFilesCh = channel('ci:vitest:run-files')
+
+const testSuiteStartCh = channel('ci:vitest:test-suite:start')
+const testSuiteFinishCh = channel('ci:vitest:test-suite:finish')
+
+const testSessionStartCh = channel('ci:vitest:session:start')
+const testSessionFinishCh = channel('ci:vitest:session:finish')
+
+const taskToAsync = new WeakMap()
+
+const sessionAsyncResource = new AsyncResource('bound-anonymous-fn')
+
+// can I access VitestRunner???
+
+const isVitestTestRunner = (vitestPackage) => {
+  return vitestPackage.VitestTestRunner
+}
+
+const isReporter = (vitestPackage) => {
+  // maybe we can use this as session start? maybe onCollected or similar
+  return vitestPackage.D?.name === 'DefaultReporter'
+}
+
+const isVitestPlugin = (vitestPackage) => {
+  return vitestPackage.B?.name === 'BaseSequencer' // maybe we can use #sort or something to start session
+}
+
+// --- do I need to specify a file? That's problematic because they contain a hash
+addHook({
+  name: 'vitest',
+  versions: ['>=0.0.0']
+}, (vitestPackage, frameworkVersion) => {
+  // debugger
+  if (isVitestPlugin(vitestPackage)) {
+    // const oldVitestPlugin = vitestPackage.V
+    // vitestPackage.V = async function () {
+    //   debugger
+    //   return oldVitestPlugin.apply(this, arguments)
+    // }
+    // const oldCreateVitest = vitestPackage.c
+    // vitestPackage.c = async function () {
+    //   debugger
+    //   return oldCreateVitest.apply(this, arguments)
+    // }
+
+    // TODO: use shimmer
+    const oldStartVitest = vitestPackage.s
+    vitestPackage.s = async function () {
+      // this has actually worked!!!
+      debugger
+      sessionAsyncResource.runInAsyncScope(() => {
+        // TODO: change command to proper command
+        testSessionStartCh.publish({ command: 'vitest run', frameworkVersion })
+      })
+      const res = await oldStartVitest.apply(this, arguments)
+
+      sessionAsyncResource.runInAsyncScope(() => {
+        // TODO: get proper status
+        testSessionFinishCh.publish('pass')
+      })
+      debugger
+
+      return res
+    }
+  }
+
+  // maybe not needed
+  // if (isReporter(vitestPackage)) {
+  //   //  maybe not used
+  //   shimmer.wrap(vitestPackage.D.prototype, 'onPathsCollected', onPathsCollected => async function (test) {
+  //     debugger
+  //     return onPathsCollected.apply(this, arguments)
+  //   })
+  // }
+
+  if (isVitestTestRunner(vitestPackage)) {
+    // test start
+    shimmer.wrap(vitestPackage.VitestTestRunner.prototype, 'onBeforeTryTask', onBeforeTryTask => async function (task) {
+      const asyncResource = new AsyncResource('bound-anonymous-fn')
+      taskToAsync.set(task, asyncResource)
+
+      asyncResource.runInAsyncScope(() => {
+        testStartCh.publish({ testName: task.name, testSuiteAbsolutePath: task.suite.file.filepath })
+      })
+      return onBeforeTryTask.apply(this, arguments)
+    })
+
+    shimmer.wrap(vitestPackage.VitestTestRunner.prototype, 'onAfterTryTask', onAfterTryTask => async function (task) {
+      const res = await onAfterTryTask.apply(this, arguments)
+
+      const asyncResource = taskToAsync.get(task)
+
+      asyncResource.runInAsyncScope(() => {
+        // TODO: if no error has been found, it's a pass.
+        // See logic in packages/runner/src/run.ts in vitest after calling onAfterTryTask
+        testFinishCh.publish('pass')
+      })
+      return res
+    })
+
+    // TODO: probably no need to flush so often - maybe we can use the test suite finish
+    shimmer.wrap(vitestPackage.VitestTestRunner.prototype, 'onAfterRunFiles', onAfterRunFiles => async function () {
+      let onFinish = null
+      const onFinishPromise = new Promise(resolve => {
+        onFinish = resolve
+      })
+
+      onAfterRunFilesCh.publish(onFinish)
+
+      await onFinishPromise
+
+      return await onAfterRunFiles.apply(this, arguments)
+    })
+  }
+  return vitestPackage
+})
+
+addHook({
+  name: '@vitest/runner',
+  versions: ['>=0.0.0']
+}, vitestPackage => {
+  debugger
+  if (vitestPackage.VitestRunner) {
+    console.log('VitestRunner!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')
+  }
+  // TODO: it does not seem to reach this point --- but it does using iitm directly --- why?
+  // it reaches on workers only
+  if (vitestPackage.startTests) {
+    shimmer.wrap(vitestPackage, 'startTests', startTests => async function (testPath, vitestTestRunner) {
+      const asyncResource = new AsyncResource('bound-anonymous-fn')
+      asyncResource.runInAsyncScope(() => {
+        testSuiteStartCh.publish(testPath[0])
+      })
+      const startTestsResponse = await startTests.apply(this, arguments)
+      asyncResource.runInAsyncScope(() => {
+        testSuiteFinishCh.publish(startTestsResponse[0].result.state)
+      })
+      return startTestsResponse
+    })
+
+    // const oldStartTests = vitestPackage.startTests
+    // vitestPackage.startTests = async function (testPath, vitestTestRunner) {
+    //   console.log('start suite', testPath[0])
+    //   const res = await oldStartTests.apply(this, arguments)
+    //   console.log('test suite finished', testPath[0])
+    //   console.log('result', res[0].result.state)
+    //   // console.log('tasks', res[0].tasks)
+    //   return res
+    // }
+  }
+
+  return vitestPackage
+})
+
+addHook({
+  name: '@vitest',
+  versions: ['>=0.0.0']
+}, vitestPackage => {
+  debugger
+  console.log('@vitest', { isMainThread })
+  return vitestPackage
+})
diff --git a/packages/datadog-plugin-vitest/src/index.js b/packages/datadog-plugin-vitest/src/index.js
new file mode 100644
index 00000000000..25c72952a10
--- /dev/null
+++ b/packages/datadog-plugin-vitest/src/index.js
@@ -0,0 +1,85 @@
+const CiPlugin = require('../../dd-trace/src/plugins/ci_plugin')
+const { storage } = require('../../datadog-core')
+
+const {
+  TEST_STATUS,
+  finishAllTraceSpans,
+  getTestSuitePath,
+  getTestSuiteCommonTags
+} = require('../../dd-trace/src/plugins/util/test')
+const { COMPONENT } = require('../../dd-trace/src/constants')
+
+class VitestPlugin extends CiPlugin {
+  static get id () {
+    return 'vitest'
+  }
+
+  constructor (...args) {
+    super(...args)
+    this.addSub('ci:vitest:test:start', ({ testName, testSuiteAbsolutePath }) => {
+      const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
+      const store = storage.getStore()
+      const span = this.startTestSpan(testName, testSuite, this.testSuiteSpan)
+
+      this.enter(span, store)
+    })
+
+    this.addSub('ci:vitest:test:finish', (status) => {
+      const store = storage.getStore()
+      const span = store?.span
+
+      if (span) {
+        span.setTag(TEST_STATUS, status)
+
+        span.finish()
+        finishAllTraceSpans(span)
+      }
+    })
+
+    this.addSub('ci:vitest:test-suite:start', (testSuiteAbsolutePath) => {
+      const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
+      const testSuiteMetadata = getTestSuiteCommonTags(
+        this.command,
+        this.frameworkVersion,
+        testSuite,
+        'vitest'
+      )
+      const testSuiteSpan = this.tracer.startSpan('mocha.test_suite', {
+        childOf: this.testModuleSpan,
+        tags: {
+          [COMPONENT]: this.constructor.id,
+          ...this.testEnvironmentMetadata,
+          ...testSuiteMetadata
+        }
+      })
+      const store = storage.getStore()
+      this.enter(testSuiteSpan, store)
+      this.testSuiteSpan = testSuiteSpan
+    })
+
+    this.addSub('ci:vitest:test-suite:finish', status => {
+      const store = storage.getStore()
+      if (store && store.span) {
+        const span = store.span
+        span.setTag(TEST_STATUS, status)
+        span.finish()
+      }
+    })
+
+    // TODO: do we need to flush?
+    this.addSub('ci:vitest:session:finish', (status) => {
+      console.log('session finish!', status)
+      this.testSessionSpan.setTag(TEST_STATUS, status)
+      this.testModuleSpan.setTag(TEST_STATUS, status)
+      this.testModuleSpan.finish()
+      this.testSessionSpan.finish()
+      finishAllTraceSpans(this.testSessionSpan)
+    })
+
+    this.addSub('ci:vitest:run-files', onFinish => {
+      this.tracer._exporter.flush(onFinish)
+    })
+  }
+}
+
+module.exports = VitestPlugin
diff --git a/packages/dd-trace/src/plugins/index.js b/packages/dd-trace/src/plugins/index.js
index 0b98cd9c076..076e67104b8 100644
--- a/packages/dd-trace/src/plugins/index.js
+++ b/packages/dd-trace/src/plugins/index.js
@@ -54,6 +54,9 @@ module.exports = {
   get 'microgateway-core' () { return require('../../../datadog-plugin-microgateway-core/src') },
   get mocha () { return require('../../../datadog-plugin-mocha/src') },
   get 'mocha-each' () { return require('../../../datadog-plugin-mocha/src') },
+  get vitest () { return require('../../../datadog-plugin-vitest/src') },
+  get '@vitest' () { return require('../../../datadog-plugin-vitest/src') },
+  get '@vitest/runner' () { return require('../../../datadog-plugin-vitest/src') },
   get workerpool () { return require('../../../datadog-plugin-mocha/src') },
   get moleculer () { return require('../../../datadog-plugin-moleculer/src') },
   get mongodb () { return require('../../../datadog-plugin-mongodb-core/src') },

From 4b03ee1f1c34b210940e2a05e2f1868249137765 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Tue, 18 Jun 2024 14:48:33 +0200
Subject: [PATCH 02/26] better vitest support

---
 .../datadog-instrumentations/src/vitest.js    | 116 +++++++++++++++---
 packages/datadog-plugin-vitest/src/index.js   |  23 ++--
 packages/dd-trace/src/plugins/ci_plugin.js    |   4 +
 3 files changed, 117 insertions(+), 26 deletions(-)

diff --git a/packages/datadog-instrumentations/src/vitest.js b/packages/datadog-instrumentations/src/vitest.js
index e79ea866fb1..8b537d0a336 100644
--- a/packages/datadog-instrumentations/src/vitest.js
+++ b/packages/datadog-instrumentations/src/vitest.js
@@ -111,18 +111,18 @@ addHook({
     })
 
     // TODO: probably no need to flush so often - maybe we can use the test suite finish
-    shimmer.wrap(vitestPackage.VitestTestRunner.prototype, 'onAfterRunFiles', onAfterRunFiles => async function () {
-      let onFinish = null
-      const onFinishPromise = new Promise(resolve => {
-        onFinish = resolve
-      })
+    // shimmer.wrap(vitestPackage.VitestTestRunner.prototype, 'onAfterRunFiles', onAfterRunFiles => async function () {
+    //   let onFinish = null
+    //   const onFinishPromise = new Promise(resolve => {
+    //     onFinish = resolve
+    //   })
 
-      onAfterRunFilesCh.publish(onFinish)
+    //   onAfterRunFilesCh.publish(onFinish)
 
-      await onFinishPromise
+    //   await onFinishPromise
 
-      return await onAfterRunFiles.apply(this, arguments)
-    })
+    //   return await onAfterRunFiles.apply(this, arguments)
+    // })
   }
   return vitestPackage
 })
@@ -144,9 +144,20 @@ addHook({
         testSuiteStartCh.publish(testPath[0])
       })
       const startTestsResponse = await startTests.apply(this, arguments)
+
+      let onFinish = null
+      const onFinishPromise = new Promise(resolve => {
+        onFinish = resolve
+      })
+
       asyncResource.runInAsyncScope(() => {
-        testSuiteFinishCh.publish(startTestsResponse[0].result.state)
+        testSuiteFinishCh.publish({ status: startTestsResponse[0].result.state, onFinish })
       })
+
+      console.log('waiting flush')
+      await onFinishPromise
+      console.log('flushed')
+
       return startTestsResponse
     })
 
@@ -164,11 +175,80 @@ addHook({
   return vitestPackage
 })
 
-addHook({
-  name: '@vitest',
-  versions: ['>=0.0.0']
-}, vitestPackage => {
-  debugger
-  console.log('@vitest', { isMainThread })
-  return vitestPackage
-})
+let privateSet
+
+const isPrivateTools = (tinypoolPackage) => {
+  return tinypoolPackage.__privateGet
+}
+
+const isTinyPoolClass = (tinypoolPackage) => {
+  return tinypoolPackage.name === 'Tinypool'
+}
+
+const isThreadPool = (value) => {
+  return value.constructor.name === 'ThreadPool'
+}
+
+let threadPool
+
+// MAYBE NOT NEEDED IF I CAN PASS SESSION AND MODULE ID TO THE WORKERS
+// addHook({
+//   name: 'tinypool',
+//   versions: ['>=0.0.0']
+// }, tinypoolPackage => {
+//   debugger
+//   if (isPrivateTools(tinypoolPackage)) {
+//     // maybe get private set??
+//     const oldPrivateSet = tinypoolPackage.__privateSet
+
+//     tinypoolPackage.__privateSet = function (obj, member, value, setter) {
+//       if (isThreadPool(value)) {
+//         threadPool = value
+//         for (const worker of threadPool.workers.readyItems.values()) {
+//           debugger
+//           const oldWorkerMessage = worker.onMessage
+//           worker.onMessage = function (message) {
+//             debugger
+//             // we have to stop our message from being handled
+//             if (message.type === 'ci:vitest:worker:ready') {
+//               return
+//             }
+//             return oldWorkerMessage.apply(this, arguments)
+//           }
+//           const oldMessageHandler = worker.worker.thread._events.message
+
+//           worker.worker.thread.removeListener('message', oldMessageHandler)
+
+//           worker.worker.thread.on('message', function (message) {
+//             debugger
+//             // we have to stop our message from being handled
+//             if (message.type === 'ci:vitest:worker:ready') {
+//               return
+//             }
+//             return oldMessageHandler.apply(this, arguments)
+//           })
+
+//           // debugger
+//           // worker.worker.thread._events.message = function (message) {
+//           //   debugger
+//           //   // we have to stop our message from being handled
+//           //   if (message.type === 'ci:vitest:worker:ready') {
+//           //     return
+//           //   }
+//           //   return oldMessageHandler.apply(this, arguments)
+//           // }
+//         }
+//       }
+//       return oldPrivateSet.apply(this, arguments)
+//     }
+//   }
+//   if (isTinyPoolClass(tinypoolPackage)) {
+//     shimmer.wrap(tinypoolPackage.prototype, 'run', run => async function (task) {
+//       debugger
+//       const res = await run.apply(this, arguments)
+//       debugger
+//       return res
+//     })
+//   }
+//   return tinypoolPackage
+// })
diff --git a/packages/datadog-plugin-vitest/src/index.js b/packages/datadog-plugin-vitest/src/index.js
index 25c72952a10..7398d4cb307 100644
--- a/packages/datadog-plugin-vitest/src/index.js
+++ b/packages/datadog-plugin-vitest/src/index.js
@@ -37,6 +37,11 @@ class VitestPlugin extends CiPlugin {
     })
 
     this.addSub('ci:vitest:test-suite:start', (testSuiteAbsolutePath) => {
+      const testSessionSpanContext = this.tracer.extract('text_map', {
+        'x-datadog-trace-id': process.env.DD_CIVISIBILITY_TEST_SESSION_ID,
+        'x-datadog-parent-id': process.env.DD_CIVISIBILITY_TEST_MODULE_ID
+      })
+
       const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
       const testSuiteMetadata = getTestSuiteCommonTags(
         this.command,
@@ -44,8 +49,8 @@ class VitestPlugin extends CiPlugin {
         testSuite,
         'vitest'
       )
-      const testSuiteSpan = this.tracer.startSpan('mocha.test_suite', {
-        childOf: this.testModuleSpan,
+      const testSuiteSpan = this.tracer.startSpan('vitest.test_suite', {
+        childOf: testSessionSpanContext,
         tags: {
           [COMPONENT]: this.constructor.id,
           ...this.testEnvironmentMetadata,
@@ -57,18 +62,19 @@ class VitestPlugin extends CiPlugin {
       this.testSuiteSpan = testSuiteSpan
     })
 
-    this.addSub('ci:vitest:test-suite:finish', status => {
+    this.addSub('ci:vitest:test-suite:finish', ({ status, onFinish }) => {
       const store = storage.getStore()
       if (store && store.span) {
         const span = store.span
         span.setTag(TEST_STATUS, status)
         span.finish()
+        finishAllTraceSpans(span)
       }
+      this.tracer._exporter.flush(onFinish)
     })
 
-    // TODO: do we need to flush?
+    // TODO: do we need to flush? - probably not because it's just two spans in the main process
     this.addSub('ci:vitest:session:finish', (status) => {
-      console.log('session finish!', status)
       this.testSessionSpan.setTag(TEST_STATUS, status)
       this.testModuleSpan.setTag(TEST_STATUS, status)
       this.testModuleSpan.finish()
@@ -76,9 +82,10 @@ class VitestPlugin extends CiPlugin {
       finishAllTraceSpans(this.testSessionSpan)
     })
 
-    this.addSub('ci:vitest:run-files', onFinish => {
-      this.tracer._exporter.flush(onFinish)
-    })
+    // this.addSub('ci:vitest:run-files', onFinish => {
+    //   console.log('flushing in run files')
+    //   this.tracer._exporter.flush(onFinish)
+    // })
   }
 }
 
diff --git a/packages/dd-trace/src/plugins/ci_plugin.js b/packages/dd-trace/src/plugins/ci_plugin.js
index 310cb2dc940..a93eb6906a7 100644
--- a/packages/dd-trace/src/plugins/ci_plugin.js
+++ b/packages/dd-trace/src/plugins/ci_plugin.js
@@ -91,6 +91,10 @@ module.exports = class CiPlugin extends Plugin {
           ...testModuleSpanMetadata
         }
       })
+      // only for vitest
+      process.env.DD_CIVISIBILITY_TEST_SESSION_ID = this.testSessionSpan.context().toTraceId()
+      process.env.DD_CIVISIBILITY_TEST_MODULE_ID = this.testModuleSpan.context().toSpanId()
+
       this.telemetry.ciVisEvent(TELEMETRY_EVENT_CREATED, 'module')
     })
 

From 689d9aa99b1740f0d8b352fe968ea2b4d200dadf Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Tue, 18 Jun 2024 17:22:44 +0200
Subject: [PATCH 03/26] wip

---
 ci/init.js                                             | 10 ++++++++++
 ext/exporters.d.ts                                     |  5 +++--
 ext/exporters.js                                       |  3 ++-
 packages/datadog-instrumentations/src/helpers/hook.js  |  5 +++--
 packages/datadog-instrumentations/src/helpers/hooks.js |  1 +
 packages/dd-trace/src/exporter.js                      |  1 +
 packages/dd-trace/src/iitm.js                          |  3 +++
 packages/dd-trace/src/plugins/index.js                 |  1 +
 8 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/ci/init.js b/ci/init.js
index b54e29abd4d..3d3e345aaaa 100644
--- a/ci/init.js
+++ b/ci/init.js
@@ -1,10 +1,14 @@
 /* eslint-disable no-console */
 const tracer = require('../packages/dd-trace')
 const { isTrue } = require('../packages/dd-trace/src/util')
+const { isMainThread, parentPort } = require('node:worker_threads')
 
 const isJestWorker = !!process.env.JEST_WORKER_ID
 const isCucumberWorker = !!process.env.CUCUMBER_WORKER_ID
 const isMochaWorker = !!process.env.MOCHA_WORKER_ID
+// eslint-disable-next-line
+// https://github.com/vitest-dev/vitest/blob/f969fb0f9f0247a7daa2afee8f70de25ea5e123f/packages/vitest/src/node/pool.ts#L110-L111
+const isVitestWorker = !isMainThread && process.env.VITEST === 'true' && process.env.TEST === 'true'
 
 const options = {
   startupLogs: false,
@@ -51,6 +55,12 @@ if (isMochaWorker) {
   }
 }
 
+if (isVitestWorker) {
+  // if (parentPort?.postMessage) {
+  //   parentPort.postMessage({ type: 'ci:vitest:worker:ready' })
+  // }
+}
+
 if (shouldInit) {
   tracer.init(options)
   tracer.use('fs', false)
diff --git a/ext/exporters.d.ts b/ext/exporters.d.ts
index 07bc2cd29e3..901025d1aa4 100644
--- a/ext/exporters.d.ts
+++ b/ext/exporters.d.ts
@@ -4,8 +4,9 @@ declare const exporters: {
   DATADOG: 'datadog',
   AGENT_PROXY: 'agent_proxy',
   JEST_WORKER: 'jest_worker',
-  CUCUMBER_WORKER: 'cucumber_worker'
-  MOCHA_WORKER: 'mocha_worker'
+  CUCUMBER_WORKER: 'cucumber_worker',
+  MOCHA_WORKER: 'mocha_worker',
+  VITEST_WORKER: 'vitest_worker'
 }
 
 export = exporters
diff --git a/ext/exporters.js b/ext/exporters.js
index 770116c3152..43ebb33603e 100644
--- a/ext/exporters.js
+++ b/ext/exporters.js
@@ -6,5 +6,6 @@ module.exports = {
   AGENT_PROXY: 'agent_proxy',
   JEST_WORKER: 'jest_worker',
   CUCUMBER_WORKER: 'cucumber_worker',
-  MOCHA_WORKER: 'mocha_worker'
+  MOCHA_WORKER: 'mocha_worker',
+  VITEST_WORKER: 'vitest_worker'
 }
diff --git a/packages/datadog-instrumentations/src/helpers/hook.js b/packages/datadog-instrumentations/src/helpers/hook.js
index 7bec453187a..31d53cee135 100644
--- a/packages/datadog-instrumentations/src/helpers/hook.js
+++ b/packages/datadog-instrumentations/src/helpers/hook.js
@@ -20,7 +20,8 @@ function Hook (modules, onrequire) {
     const parts = [moduleBaseDir, moduleName].filter(v => v)
     const filename = path.join(...parts)
 
-    if (this._patched[filename]) return moduleExports
+    // this does not seem to work with IITM: the same file seems to import different stuff
+    // if (this._patched[filename]) return moduleExports
 
     this._patched[filename] = true
 
@@ -30,7 +31,7 @@ function Hook (modules, onrequire) {
   this._ritmHook = ritm(modules, {}, safeHook)
   this._iitmHook = iitm(modules, {}, (moduleExports, moduleName, moduleBaseDir) => {
     // TODO: Move this logic to import-in-the-middle and only do it for CommonJS
-    // modules and not ESM. In the meantime, all the modules we instrument are
+    // modules and not ESM. In the meantime, all thpe modules we instrument are
     // CommonJS modules for which the default export is always moved to
     // `default` anyway.
     if (moduleExports && moduleExports.default) {
diff --git a/packages/datadog-instrumentations/src/helpers/hooks.js b/packages/datadog-instrumentations/src/helpers/hooks.js
index 61f514f1232..32501d3fc1d 100644
--- a/packages/datadog-instrumentations/src/helpers/hooks.js
+++ b/packages/datadog-instrumentations/src/helpers/hooks.js
@@ -74,6 +74,7 @@ module.exports = {
   vitest: () => require('../vitest'),
   '@vitest': () => require('../vitest'),
   '@vitest/runner': () => require('../vitest'),
+  tinypool: () => require('../vitest'),
   workerpool: () => require('../mocha'),
   moleculer: () => require('../moleculer'),
   mongodb: () => require('../mongodb'),
diff --git a/packages/dd-trace/src/exporter.js b/packages/dd-trace/src/exporter.js
index 02d50c3b57e..a657a3a60c8 100644
--- a/packages/dd-trace/src/exporter.js
+++ b/packages/dd-trace/src/exporter.js
@@ -20,6 +20,7 @@ module.exports = name => {
     case exporters.JEST_WORKER:
     case exporters.CUCUMBER_WORKER:
     case exporters.MOCHA_WORKER:
+    case exporters.VITEST_WORKER:
       return require('./ci-visibility/exporters/test-worker')
     default:
       return inAWSLambda && !usingLambdaExtension ? require('./exporters/log') : require('./exporters/agent')
diff --git a/packages/dd-trace/src/iitm.js b/packages/dd-trace/src/iitm.js
index 86a8d4dcecd..e9317f20f2e 100644
--- a/packages/dd-trace/src/iitm.js
+++ b/packages/dd-trace/src/iitm.js
@@ -8,6 +8,9 @@ const dc = require('dc-polyfill')
 if (semver.satisfies(process.versions.node, '>=14.13.1')) {
   const moduleLoadStartChannel = dc.channel('dd-trace:moduleLoadStart')
   addHook((name, namespace) => {
+    // if (name.includes('/vitest/') || name.includes('@vitest')) {
+    //   debugger
+    // }
     if (moduleLoadStartChannel.hasSubscribers) {
       moduleLoadStartChannel.publish({
         filename: name,
diff --git a/packages/dd-trace/src/plugins/index.js b/packages/dd-trace/src/plugins/index.js
index 076e67104b8..13adbe3b649 100644
--- a/packages/dd-trace/src/plugins/index.js
+++ b/packages/dd-trace/src/plugins/index.js
@@ -57,6 +57,7 @@ module.exports = {
   get vitest () { return require('../../../datadog-plugin-vitest/src') },
   get '@vitest' () { return require('../../../datadog-plugin-vitest/src') },
   get '@vitest/runner' () { return require('../../../datadog-plugin-vitest/src') },
+  get tinypool () { return require('../../../datadog-plugin-vitest/src') },
   get workerpool () { return require('../../../datadog-plugin-mocha/src') },
   get moleculer () { return require('../../../datadog-plugin-moleculer/src') },
   get mongodb () { return require('../../../datadog-plugin-mongodb-core/src') },

From 6512002ca0a1a65fbf6aa17c364663fe4b8437a5 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Wed, 19 Jun 2024 11:30:20 +0200
Subject: [PATCH 04/26] clean up

---
 ci/init.js                                    |  10 --
 ext/exporters.d.ts                            |   3 +-
 ext/exporters.js                              |   3 +-
 .../src/helpers/hook.js                       |   2 +-
 .../src/helpers/hooks.js                      |   6 +-
 .../datadog-instrumentations/src/vitest.js    | 159 +-----------------
 6 files changed, 13 insertions(+), 170 deletions(-)

diff --git a/ci/init.js b/ci/init.js
index 3d3e345aaaa..b54e29abd4d 100644
--- a/ci/init.js
+++ b/ci/init.js
@@ -1,14 +1,10 @@
 /* eslint-disable no-console */
 const tracer = require('../packages/dd-trace')
 const { isTrue } = require('../packages/dd-trace/src/util')
-const { isMainThread, parentPort } = require('node:worker_threads')
 
 const isJestWorker = !!process.env.JEST_WORKER_ID
 const isCucumberWorker = !!process.env.CUCUMBER_WORKER_ID
 const isMochaWorker = !!process.env.MOCHA_WORKER_ID
-// eslint-disable-next-line
-// https://github.com/vitest-dev/vitest/blob/f969fb0f9f0247a7daa2afee8f70de25ea5e123f/packages/vitest/src/node/pool.ts#L110-L111
-const isVitestWorker = !isMainThread && process.env.VITEST === 'true' && process.env.TEST === 'true'
 
 const options = {
   startupLogs: false,
@@ -55,12 +51,6 @@ if (isMochaWorker) {
   }
 }
 
-if (isVitestWorker) {
-  // if (parentPort?.postMessage) {
-  //   parentPort.postMessage({ type: 'ci:vitest:worker:ready' })
-  // }
-}
-
 if (shouldInit) {
   tracer.init(options)
   tracer.use('fs', false)
diff --git a/ext/exporters.d.ts b/ext/exporters.d.ts
index 901025d1aa4..2f462dd93e7 100644
--- a/ext/exporters.d.ts
+++ b/ext/exporters.d.ts
@@ -5,8 +5,7 @@ declare const exporters: {
   AGENT_PROXY: 'agent_proxy',
   JEST_WORKER: 'jest_worker',
   CUCUMBER_WORKER: 'cucumber_worker',
-  MOCHA_WORKER: 'mocha_worker',
-  VITEST_WORKER: 'vitest_worker'
+  MOCHA_WORKER: 'mocha_worker'
 }
 
 export = exporters
diff --git a/ext/exporters.js b/ext/exporters.js
index 43ebb33603e..770116c3152 100644
--- a/ext/exporters.js
+++ b/ext/exporters.js
@@ -6,6 +6,5 @@ module.exports = {
   AGENT_PROXY: 'agent_proxy',
   JEST_WORKER: 'jest_worker',
   CUCUMBER_WORKER: 'cucumber_worker',
-  MOCHA_WORKER: 'mocha_worker',
-  VITEST_WORKER: 'vitest_worker'
+  MOCHA_WORKER: 'mocha_worker'
 }
diff --git a/packages/datadog-instrumentations/src/helpers/hook.js b/packages/datadog-instrumentations/src/helpers/hook.js
index 31d53cee135..23046e00184 100644
--- a/packages/datadog-instrumentations/src/helpers/hook.js
+++ b/packages/datadog-instrumentations/src/helpers/hook.js
@@ -31,7 +31,7 @@ function Hook (modules, onrequire) {
   this._ritmHook = ritm(modules, {}, safeHook)
   this._iitmHook = iitm(modules, {}, (moduleExports, moduleName, moduleBaseDir) => {
     // TODO: Move this logic to import-in-the-middle and only do it for CommonJS
-    // modules and not ESM. In the meantime, all thpe modules we instrument are
+    // modules and not ESM. In the meantime, all the modules we instrument are
     // CommonJS modules for which the default export is always moved to
     // `default` anyway.
     if (moduleExports && moduleExports.default) {
diff --git a/packages/datadog-instrumentations/src/helpers/hooks.js b/packages/datadog-instrumentations/src/helpers/hooks.js
index 32501d3fc1d..c06e1101ecd 100644
--- a/packages/datadog-instrumentations/src/helpers/hooks.js
+++ b/packages/datadog-instrumentations/src/helpers/hooks.js
@@ -23,6 +23,7 @@ module.exports = {
   '@opentelemetry/sdk-trace-node': () => require('../otel-sdk-trace'),
   '@redis/client': () => require('../redis'),
   '@smithy/smithy-client': () => require('../aws-sdk'),
+  '@vitest/runner': () => require('../vitest'),
   aerospike: () => require('../aerospike'),
   amqp10: () => require('../amqp10'),
   amqplib: () => require('../amqplib'),
@@ -71,10 +72,6 @@ module.exports = {
   'microgateway-core': () => require('../microgateway-core'),
   mocha: () => require('../mocha'),
   'mocha-each': () => require('../mocha'),
-  vitest: () => require('../vitest'),
-  '@vitest': () => require('../vitest'),
-  '@vitest/runner': () => require('../vitest'),
-  tinypool: () => require('../vitest'),
   workerpool: () => require('../mocha'),
   moleculer: () => require('../moleculer'),
   mongodb: () => require('../mongodb'),
@@ -114,6 +111,7 @@ module.exports = {
   sharedb: () => require('../sharedb'),
   tedious: () => require('../tedious'),
   undici: () => require('../undici'),
+  vitest: () => require('../vitest'),
   when: () => require('../when'),
   winston: () => require('../winston')
 }
diff --git a/packages/datadog-instrumentations/src/vitest.js b/packages/datadog-instrumentations/src/vitest.js
index 8b537d0a336..b4e7b0c15c8 100644
--- a/packages/datadog-instrumentations/src/vitest.js
+++ b/packages/datadog-instrumentations/src/vitest.js
@@ -1,5 +1,4 @@
 const { addHook, channel, AsyncResource } = require('./helpers/instrument')
-const { isMainThread } = require('node:worker_threads')
 const shimmer = require('../../datadog-shimmer')
 
 // Needs modifying node_modules/.pnpm/import-in-the-middle@1.7.3/node_modules/import-in-the-middle/index.js
@@ -10,8 +9,6 @@ const shimmer = require('../../datadog-shimmer')
 const testStartCh = channel('ci:vitest:test:start')
 const testFinishCh = channel('ci:vitest:test:finish')
 
-const onAfterRunFilesCh = channel('ci:vitest:run-files')
-
 const testSuiteStartCh = channel('ci:vitest:test-suite:start')
 const testSuiteFinishCh = channel('ci:vitest:test-suite:finish')
 
@@ -22,19 +19,12 @@ const taskToAsync = new WeakMap()
 
 const sessionAsyncResource = new AsyncResource('bound-anonymous-fn')
 
-// can I access VitestRunner???
-
 const isVitestTestRunner = (vitestPackage) => {
   return vitestPackage.VitestTestRunner
 }
 
-const isReporter = (vitestPackage) => {
-  // maybe we can use this as session start? maybe onCollected or similar
-  return vitestPackage.D?.name === 'DefaultReporter'
-}
-
 const isVitestPlugin = (vitestPackage) => {
-  return vitestPackage.B?.name === 'BaseSequencer' // maybe we can use #sort or something to start session
+  return vitestPackage.B?.name === 'BaseSequencer'
 }
 
 // --- do I need to specify a file? That's problematic because they contain a hash
@@ -42,49 +32,23 @@ addHook({
   name: 'vitest',
   versions: ['>=0.0.0']
 }, (vitestPackage, frameworkVersion) => {
-  // debugger
   if (isVitestPlugin(vitestPackage)) {
-    // const oldVitestPlugin = vitestPackage.V
-    // vitestPackage.V = async function () {
-    //   debugger
-    //   return oldVitestPlugin.apply(this, arguments)
-    // }
-    // const oldCreateVitest = vitestPackage.c
-    // vitestPackage.c = async function () {
-    //   debugger
-    //   return oldCreateVitest.apply(this, arguments)
-    // }
-
-    // TODO: use shimmer
-    const oldStartVitest = vitestPackage.s
-    vitestPackage.s = async function () {
-      // this has actually worked!!!
-      debugger
+    // TODO: will the "s" (minified version) be the same in future versions?
+    shimmer.wrap(vitestPackage, 's', startVitest => async function () {
       sessionAsyncResource.runInAsyncScope(() => {
         // TODO: change command to proper command
         testSessionStartCh.publish({ command: 'vitest run', frameworkVersion })
       })
-      const res = await oldStartVitest.apply(this, arguments)
+      const res = await startVitest.apply(this, arguments)
 
       sessionAsyncResource.runInAsyncScope(() => {
         // TODO: get proper status
         testSessionFinishCh.publish('pass')
       })
-      debugger
-
       return res
-    }
+    })
   }
 
-  // maybe not needed
-  // if (isReporter(vitestPackage)) {
-  //   //  maybe not used
-  //   shimmer.wrap(vitestPackage.D.prototype, 'onPathsCollected', onPathsCollected => async function (test) {
-  //     debugger
-  //     return onPathsCollected.apply(this, arguments)
-  //   })
-  // }
-
   if (isVitestTestRunner(vitestPackage)) {
     // test start
     shimmer.wrap(vitestPackage.VitestTestRunner.prototype, 'onBeforeTryTask', onBeforeTryTask => async function (task) {
@@ -97,6 +61,7 @@ addHook({
       return onBeforeTryTask.apply(this, arguments)
     })
 
+    // test finish
     shimmer.wrap(vitestPackage.VitestTestRunner.prototype, 'onAfterTryTask', onAfterTryTask => async function (task) {
       const res = await onAfterTryTask.apply(this, arguments)
 
@@ -109,34 +74,16 @@ addHook({
       })
       return res
     })
-
-    // TODO: probably no need to flush so often - maybe we can use the test suite finish
-    // shimmer.wrap(vitestPackage.VitestTestRunner.prototype, 'onAfterRunFiles', onAfterRunFiles => async function () {
-    //   let onFinish = null
-    //   const onFinishPromise = new Promise(resolve => {
-    //     onFinish = resolve
-    //   })
-
-    //   onAfterRunFilesCh.publish(onFinish)
-
-    //   await onFinishPromise
-
-    //   return await onAfterRunFiles.apply(this, arguments)
-    // })
   }
   return vitestPackage
 })
 
+// test suite start and finish
+// only relevant for workers
 addHook({
   name: '@vitest/runner',
   versions: ['>=0.0.0']
 }, vitestPackage => {
-  debugger
-  if (vitestPackage.VitestRunner) {
-    console.log('VitestRunner!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')
-  }
-  // TODO: it does not seem to reach this point --- but it does using iitm directly --- why?
-  // it reaches on workers only
   if (vitestPackage.startTests) {
     shimmer.wrap(vitestPackage, 'startTests', startTests => async function (testPath, vitestTestRunner) {
       const asyncResource = new AsyncResource('bound-anonymous-fn')
@@ -154,101 +101,11 @@ addHook({
         testSuiteFinishCh.publish({ status: startTestsResponse[0].result.state, onFinish })
       })
 
-      console.log('waiting flush')
       await onFinishPromise
-      console.log('flushed')
 
       return startTestsResponse
     })
-
-    // const oldStartTests = vitestPackage.startTests
-    // vitestPackage.startTests = async function (testPath, vitestTestRunner) {
-    //   console.log('start suite', testPath[0])
-    //   const res = await oldStartTests.apply(this, arguments)
-    //   console.log('test suite finished', testPath[0])
-    //   console.log('result', res[0].result.state)
-    //   // console.log('tasks', res[0].tasks)
-    //   return res
-    // }
   }
 
   return vitestPackage
 })
-
-let privateSet
-
-const isPrivateTools = (tinypoolPackage) => {
-  return tinypoolPackage.__privateGet
-}
-
-const isTinyPoolClass = (tinypoolPackage) => {
-  return tinypoolPackage.name === 'Tinypool'
-}
-
-const isThreadPool = (value) => {
-  return value.constructor.name === 'ThreadPool'
-}
-
-let threadPool
-
-// MAYBE NOT NEEDED IF I CAN PASS SESSION AND MODULE ID TO THE WORKERS
-// addHook({
-//   name: 'tinypool',
-//   versions: ['>=0.0.0']
-// }, tinypoolPackage => {
-//   debugger
-//   if (isPrivateTools(tinypoolPackage)) {
-//     // maybe get private set??
-//     const oldPrivateSet = tinypoolPackage.__privateSet
-
-//     tinypoolPackage.__privateSet = function (obj, member, value, setter) {
-//       if (isThreadPool(value)) {
-//         threadPool = value
-//         for (const worker of threadPool.workers.readyItems.values()) {
-//           debugger
-//           const oldWorkerMessage = worker.onMessage
-//           worker.onMessage = function (message) {
-//             debugger
-//             // we have to stop our message from being handled
-//             if (message.type === 'ci:vitest:worker:ready') {
-//               return
-//             }
-//             return oldWorkerMessage.apply(this, arguments)
-//           }
-//           const oldMessageHandler = worker.worker.thread._events.message
-
-//           worker.worker.thread.removeListener('message', oldMessageHandler)
-
-//           worker.worker.thread.on('message', function (message) {
-//             debugger
-//             // we have to stop our message from being handled
-//             if (message.type === 'ci:vitest:worker:ready') {
-//               return
-//             }
-//             return oldMessageHandler.apply(this, arguments)
-//           })
-
-//           // debugger
-//           // worker.worker.thread._events.message = function (message) {
-//           //   debugger
-//           //   // we have to stop our message from being handled
-//           //   if (message.type === 'ci:vitest:worker:ready') {
-//           //     return
-//           //   }
-//           //   return oldMessageHandler.apply(this, arguments)
-//           // }
-//         }
-//       }
-//       return oldPrivateSet.apply(this, arguments)
-//     }
-//   }
-//   if (isTinyPoolClass(tinypoolPackage)) {
-//     shimmer.wrap(tinypoolPackage.prototype, 'run', run => async function (task) {
-//       debugger
-//       const res = await run.apply(this, arguments)
-//       debugger
-//       return res
-//     })
-//   }
-//   return tinypoolPackage
-// })

From d9537638fe36dc15caa81c2843358fdd350bfd21 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Wed, 19 Jun 2024 11:40:27 +0200
Subject: [PATCH 05/26] clean up

---
 packages/datadog-plugin-vitest/src/index.js | 5 -----
 packages/dd-trace/src/exporter.js           | 1 -
 packages/dd-trace/src/iitm.js               | 3 ---
 packages/dd-trace/src/plugins/ci_plugin.js  | 7 +++++--
 packages/dd-trace/src/plugins/index.js      | 4 +---
 5 files changed, 6 insertions(+), 14 deletions(-)

diff --git a/packages/datadog-plugin-vitest/src/index.js b/packages/datadog-plugin-vitest/src/index.js
index 7398d4cb307..1eab2cf1f69 100644
--- a/packages/datadog-plugin-vitest/src/index.js
+++ b/packages/datadog-plugin-vitest/src/index.js
@@ -81,11 +81,6 @@ class VitestPlugin extends CiPlugin {
       this.testSessionSpan.finish()
       finishAllTraceSpans(this.testSessionSpan)
     })
-
-    // this.addSub('ci:vitest:run-files', onFinish => {
-    //   console.log('flushing in run files')
-    //   this.tracer._exporter.flush(onFinish)
-    // })
   }
 }
 
diff --git a/packages/dd-trace/src/exporter.js b/packages/dd-trace/src/exporter.js
index a657a3a60c8..02d50c3b57e 100644
--- a/packages/dd-trace/src/exporter.js
+++ b/packages/dd-trace/src/exporter.js
@@ -20,7 +20,6 @@ module.exports = name => {
     case exporters.JEST_WORKER:
     case exporters.CUCUMBER_WORKER:
     case exporters.MOCHA_WORKER:
-    case exporters.VITEST_WORKER:
       return require('./ci-visibility/exporters/test-worker')
     default:
       return inAWSLambda && !usingLambdaExtension ? require('./exporters/log') : require('./exporters/agent')
diff --git a/packages/dd-trace/src/iitm.js b/packages/dd-trace/src/iitm.js
index e9317f20f2e..86a8d4dcecd 100644
--- a/packages/dd-trace/src/iitm.js
+++ b/packages/dd-trace/src/iitm.js
@@ -8,9 +8,6 @@ const dc = require('dc-polyfill')
 if (semver.satisfies(process.versions.node, '>=14.13.1')) {
   const moduleLoadStartChannel = dc.channel('dd-trace:moduleLoadStart')
   addHook((name, namespace) => {
-    // if (name.includes('/vitest/') || name.includes('@vitest')) {
-    //   debugger
-    // }
     if (moduleLoadStartChannel.hasSubscribers) {
       moduleLoadStartChannel.publish({
         filename: name,
diff --git a/packages/dd-trace/src/plugins/ci_plugin.js b/packages/dd-trace/src/plugins/ci_plugin.js
index a93eb6906a7..4daeb02a4bf 100644
--- a/packages/dd-trace/src/plugins/ci_plugin.js
+++ b/packages/dd-trace/src/plugins/ci_plugin.js
@@ -92,8 +92,11 @@ module.exports = class CiPlugin extends Plugin {
         }
       })
       // only for vitest
-      process.env.DD_CIVISIBILITY_TEST_SESSION_ID = this.testSessionSpan.context().toTraceId()
-      process.env.DD_CIVISIBILITY_TEST_MODULE_ID = this.testModuleSpan.context().toSpanId()
+      // These are added for the worker threads to use
+      if (this.constructor.id === 'vitest') {
+        process.env.DD_CIVISIBILITY_TEST_SESSION_ID = this.testSessionSpan.context().toTraceId()
+        process.env.DD_CIVISIBILITY_TEST_MODULE_ID = this.testModuleSpan.context().toSpanId()
+      }
 
       this.telemetry.ciVisEvent(TELEMETRY_EVENT_CREATED, 'module')
     })
diff --git a/packages/dd-trace/src/plugins/index.js b/packages/dd-trace/src/plugins/index.js
index 13adbe3b649..fd9288afcc4 100644
--- a/packages/dd-trace/src/plugins/index.js
+++ b/packages/dd-trace/src/plugins/index.js
@@ -18,6 +18,7 @@ module.exports = {
   get '@opensearch-project/opensearch' () { return require('../../../datadog-plugin-opensearch/src') },
   get '@redis/client' () { return require('../../../datadog-plugin-redis/src') },
   get '@smithy/smithy-client' () { return require('../../../datadog-plugin-aws-sdk/src') },
+  get '@vitest/runner' () { return require('../../../datadog-plugin-vitest/src') },
   get aerospike () { return require('../../../datadog-plugin-aerospike/src') },
   get amqp10 () { return require('../../../datadog-plugin-amqp10/src') },
   get amqplib () { return require('../../../datadog-plugin-amqplib/src') },
@@ -55,9 +56,6 @@ module.exports = {
   get mocha () { return require('../../../datadog-plugin-mocha/src') },
   get 'mocha-each' () { return require('../../../datadog-plugin-mocha/src') },
   get vitest () { return require('../../../datadog-plugin-vitest/src') },
-  get '@vitest' () { return require('../../../datadog-plugin-vitest/src') },
-  get '@vitest/runner' () { return require('../../../datadog-plugin-vitest/src') },
-  get tinypool () { return require('../../../datadog-plugin-vitest/src') },
   get workerpool () { return require('../../../datadog-plugin-mocha/src') },
   get moleculer () { return require('../../../datadog-plugin-moleculer/src') },
   get mongodb () { return require('../../../datadog-plugin-mongodb-core/src') },

From 0d4f4acf91f4cdd81116ed6ab7596ca5800db681 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Wed, 19 Jun 2024 14:49:16 +0200
Subject: [PATCH 06/26] add e2e tests

---
 .../ci-visibility/vitest-tests/sum.mjs        |   3 +
 .../vitest-tests/test-visibility-test-2.mjs   |   8 ++
 .../vitest-tests/test-visibility-test.mjs     |   8 ++
 integration-tests/vitest.config.mjs           |   9 ++
 integration-tests/vitest/vitest.spec.js       | 102 ++++++++++++++++++
 package.json                                  |   6 +-
 .../datadog-instrumentations/src/vitest.js    |  19 +++-
 packages/datadog-plugin-vitest/src/index.js   |   6 +-
 yarn.lock                                     |   7 +-
 9 files changed, 159 insertions(+), 9 deletions(-)
 create mode 100644 integration-tests/ci-visibility/vitest-tests/sum.mjs
 create mode 100644 integration-tests/ci-visibility/vitest-tests/test-visibility-test-2.mjs
 create mode 100644 integration-tests/ci-visibility/vitest-tests/test-visibility-test.mjs
 create mode 100644 integration-tests/vitest.config.mjs
 create mode 100644 integration-tests/vitest/vitest.spec.js

diff --git a/integration-tests/ci-visibility/vitest-tests/sum.mjs b/integration-tests/ci-visibility/vitest-tests/sum.mjs
new file mode 100644
index 00000000000..f1c6520acbd
--- /dev/null
+++ b/integration-tests/ci-visibility/vitest-tests/sum.mjs
@@ -0,0 +1,3 @@
+export function sum (a, b) {
+  return a + b
+}
diff --git a/integration-tests/ci-visibility/vitest-tests/test-visibility-test-2.mjs b/integration-tests/ci-visibility/vitest-tests/test-visibility-test-2.mjs
new file mode 100644
index 00000000000..9fd618318c0
--- /dev/null
+++ b/integration-tests/ci-visibility/vitest-tests/test-visibility-test-2.mjs
@@ -0,0 +1,8 @@
+import { describe, test, expect } from 'vitest'
+import { sum } from './sum'
+
+describe('test visibility 2', () => {
+  test('can report tests 2', () => {
+    expect(sum(1, 2)).to.equal(3)
+  })
+})
diff --git a/integration-tests/ci-visibility/vitest-tests/test-visibility-test.mjs b/integration-tests/ci-visibility/vitest-tests/test-visibility-test.mjs
new file mode 100644
index 00000000000..c3e23fed22f
--- /dev/null
+++ b/integration-tests/ci-visibility/vitest-tests/test-visibility-test.mjs
@@ -0,0 +1,8 @@
+import { describe, test, expect } from 'vitest'
+import { sum } from './sum'
+
+describe('test visibility', () => {
+  test('can report tests', () => {
+    expect(sum(1, 2)).to.equal(3)
+  })
+})
diff --git a/integration-tests/vitest.config.mjs b/integration-tests/vitest.config.mjs
new file mode 100644
index 00000000000..f04d63785fd
--- /dev/null
+++ b/integration-tests/vitest.config.mjs
@@ -0,0 +1,9 @@
+import { defineConfig } from 'vite'
+
+export default defineConfig({
+  test: {
+    include: [
+      'ci-visibility/vitest-tests/test-visibility*'
+    ]
+  }
+})
diff --git a/integration-tests/vitest/vitest.spec.js b/integration-tests/vitest/vitest.spec.js
new file mode 100644
index 00000000000..f5d85e4abf1
--- /dev/null
+++ b/integration-tests/vitest/vitest.spec.js
@@ -0,0 +1,102 @@
+'use strict'
+
+const { exec } = require('child_process')
+
+const getPort = require('get-port')
+const { assert } = require('chai')
+
+const {
+  createSandbox,
+  getCiVisAgentlessConfig
+} = require('../helpers')
+const { FakeCiVisIntake } = require('../ci-visibility-intake')
+const {
+  TEST_STATUS,
+  TEST_TYPE
+} = require('../../packages/dd-trace/src/plugins/util/test')
+
+// should probably work with 1.6.0
+// tested with 1.1.0
+const versions = ['1.1.0'] // only one I've tested so far
+
+versions.forEach((version) => {
+  describe(`vitest@${version}`, () => {
+    let sandbox, cwd, receiver, childProcess
+
+    before(async function () {
+      sandbox = await createSandbox([`vitest@${version}`], true)
+      // debugger
+      cwd = sandbox.folder
+    })
+
+    after(async () => {
+      await sandbox.remove()
+    })
+
+    beforeEach(async function () {
+      const port = await getPort()
+      receiver = await new FakeCiVisIntake(port).start()
+    })
+
+    afterEach(async () => {
+      childProcess.kill()
+      await receiver.stop()
+    })
+
+    it('can run and report tests', (done) => {
+      receiver.gatherPayloadsMaxTimeout(({ url }) => url === '/api/v2/citestcycle', payloads => {
+        debugger
+        const events = payloads.flatMap(({ payload }) => payload.events)
+
+        const testSessionEvent = events.find(event => event.type === 'test_session_end')
+        const testModuleEvent = events.find(event => event.type === 'test_module_end')
+        const testSuiteEvents = events.filter(event => event.type === 'test_suite_end')
+        const testEvents = events.filter(event => event.type === 'test')
+
+        assert.include(testSessionEvent.content.resource, 'test_session.vitest run')
+        assert.equal(testSessionEvent.content.meta[TEST_STATUS], 'pass')
+        assert.include(testModuleEvent.content.resource, 'test_module.vitest run')
+        assert.equal(testModuleEvent.content.meta[TEST_STATUS], 'pass')
+        assert.equal(testSessionEvent.content.meta[TEST_TYPE], 'test')
+        assert.equal(testModuleEvent.content.meta[TEST_TYPE], 'test')
+
+        assert.includeMembers(testSuiteEvents.map(suite => suite.content.resource), [
+          'test_suite.ci-visibility/vitest-tests/test-visibility-test.mjs',
+          'test_suite.ci-visibility/vitest-tests/test-visibility-test-2.mjs'
+        ])
+
+        // TODO: just check pass
+        assert.includeMembers(testSuiteEvents.map(suite => suite.content.meta[TEST_STATUS]), [
+          'pass',
+        ])
+
+
+        assert.includeMembers(testEvents.map(test => test.content.resource), [
+          'ci-visibility/vitest-tests/test-visibility-test.mjs.can report tests',
+          'ci-visibility/vitest-tests/test-visibility-test-2.mjs.can report tests 2',
+        ])
+
+        // TODO: just check pass
+        assert.includeMembers(testEvents.map(test => test.content.meta[TEST_STATUS]), [
+          'pass',
+        ])
+      }).then(() => done()).catch(done)
+
+      childProcess = exec(
+        './node_modules/.bin/vitest run',
+        {
+          cwd,
+          env: {
+            ...getCiVisAgentlessConfig(receiver.port),
+            // maybe only in node@20
+            NODE_OPTIONS: '--import dd-trace/register.js -r dd-trace/ci/init', // ESM requires more stuff
+          },
+          stdio: 'pipe'
+        }
+      )
+
+      childProcess.stdout.pipe(process.stdout)
+      childProcess.stderr.pipe(process.stderr)
+    })
+  })
+})
diff --git a/package.json b/package.json
index 905841d2ada..1fb7ee26b11 100644
--- a/package.json
+++ b/package.json
@@ -39,6 +39,7 @@
     "test:integration:selenium": "mocha --colors --timeout 30000 -r \"packages/dd-trace/test/setup/core.js\" \"integration-tests/selenium/*.spec.js\"",
     "test:integration:profiler": "mocha --colors --timeout 90000 -r \"packages/dd-trace/test/setup/core.js\" \"integration-tests/profiler/*.spec.js\"",
     "test:integration:serverless": "mocha --colors --timeout 30000 -r \"packages/dd-trace/test/setup/core.js\" \"integration-tests/serverless/*.spec.js\"",
+    "test:integration:vitest": "mocha --colors --timeout 30000 -r \"packages/dd-trace/test/setup/core.js\" \"integration-tests/vitest/*.spec.js\"",
     "test:integration:plugins": "mocha --colors --exit -r \"packages/dd-trace/test/setup/mocha.js\" \"packages/datadog-plugin-@($(echo $PLUGINS))/test/integration-test/**/*.spec.js\"",
     "test:unit:plugins": "mocha --colors --exit -r \"packages/dd-trace/test/setup/mocha.js\" \"packages/datadog-instrumentations/test/@($(echo $PLUGINS)).spec.js\" \"packages/datadog-plugin-@($(echo $PLUGINS))/test/**/*.spec.js\" --exclude \"packages/datadog-plugin-@($(echo $PLUGINS))/test/integration-test/**/*.spec.js\"",
     "test:shimmer": "mocha --colors 'packages/datadog-shimmer/test/**/*.spec.js'",
@@ -81,7 +82,7 @@
     "crypto-randomuuid": "^1.0.0",
     "dc-polyfill": "^0.1.4",
     "ignore": "^5.2.4",
-    "import-in-the-middle": "^1.7.4",
+    "import-in-the-middle": "https://github.com/DataDog/import-in-the-middle#dbaca128f6eb517d71bcf7e6b6b45f21a07128a4",
     "int64-buffer": "^0.1.9",
     "istanbul-lib-coverage": "3.2.0",
     "jest-docblock": "^29.7.0",
@@ -137,5 +138,8 @@
     "sinon-chai": "^3.7.0",
     "tap": "^16.3.7",
     "tape": "^5.6.5"
+  },
+  "volta": {
+    "node": "20.14.0"
   }
 }
diff --git a/packages/datadog-instrumentations/src/vitest.js b/packages/datadog-instrumentations/src/vitest.js
index b4e7b0c15c8..eb7e7449a5e 100644
--- a/packages/datadog-instrumentations/src/vitest.js
+++ b/packages/datadog-instrumentations/src/vitest.js
@@ -28,13 +28,24 @@ const isVitestPlugin = (vitestPackage) => {
 }
 
 // --- do I need to specify a file? That's problematic because they contain a hash
+
+// vitestPackage.s only works for 1.1.0 - it doesn't for 1.6.0
 addHook({
   name: 'vitest',
-  versions: ['>=0.0.0']
+  versions: ['>=1.1.0 <1.6.0']
 }, (vitestPackage, frameworkVersion) => {
+  debugger
   if (isVitestPlugin(vitestPackage)) {
+    //
+    console.log('is S still s???????', vitestPackage)
     // TODO: will the "s" (minified version) be the same in future versions?
     shimmer.wrap(vitestPackage, 's', startVitest => async function () {
+      let onFinish
+
+      const flushPromise = new Promise(resolve => {
+        onFinish = resolve
+      })
+
       sessionAsyncResource.runInAsyncScope(() => {
         // TODO: change command to proper command
         testSessionStartCh.publish({ command: 'vitest run', frameworkVersion })
@@ -43,8 +54,11 @@ addHook({
 
       sessionAsyncResource.runInAsyncScope(() => {
         // TODO: get proper status
-        testSessionFinishCh.publish('pass')
+        testSessionFinishCh.publish({ status: 'pass', onFinish })
       })
+
+      await flushPromise
+
       return res
     })
   }
@@ -55,6 +69,7 @@ addHook({
       const asyncResource = new AsyncResource('bound-anonymous-fn')
       taskToAsync.set(task, asyncResource)
 
+      // TODO: task.name is just the name of the test() block: include parent describes
       asyncResource.runInAsyncScope(() => {
         testStartCh.publish({ testName: task.name, testSuiteAbsolutePath: task.suite.file.filepath })
       })
diff --git a/packages/datadog-plugin-vitest/src/index.js b/packages/datadog-plugin-vitest/src/index.js
index 1eab2cf1f69..dcce39b1cce 100644
--- a/packages/datadog-plugin-vitest/src/index.js
+++ b/packages/datadog-plugin-vitest/src/index.js
@@ -70,16 +70,18 @@ class VitestPlugin extends CiPlugin {
         span.finish()
         finishAllTraceSpans(span)
       }
+      // TODO: too frequent - how to decrease frequency?
       this.tracer._exporter.flush(onFinish)
     })
 
-    // TODO: do we need to flush? - probably not because it's just two spans in the main process
-    this.addSub('ci:vitest:session:finish', (status) => {
+    this.addSub('ci:vitest:session:finish', ({ status, onFinish }) => {
       this.testSessionSpan.setTag(TEST_STATUS, status)
       this.testModuleSpan.setTag(TEST_STATUS, status)
       this.testModuleSpan.finish()
       this.testSessionSpan.finish()
       finishAllTraceSpans(this.testSessionSpan)
+      // apparently this is indeed needed
+      this.tracer._exporter.flush(onFinish)
     })
   }
 }
diff --git a/yarn.lock b/yarn.lock
index fcbb9cd6c3f..28a7abc232c 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -2824,10 +2824,9 @@ import-fresh@^3.0.0, import-fresh@^3.2.1:
     parent-module "^1.0.0"
     resolve-from "^4.0.0"
 
-import-in-the-middle@^1.7.4:
-  version "1.7.4"
-  resolved "https://registry.yarnpkg.com/import-in-the-middle/-/import-in-the-middle-1.7.4.tgz#508da6e91cfa84f210dcdb6c0a91ab0c9e8b3ebc"
-  integrity sha512-Lk+qzWmiQuRPPulGQeK5qq0v32k2bHnWrRPFgqyvhw7Kkov5L6MOLOIU3pcWeujc9W4q54Cp3Q2WV16eQkc7Bg==
+"import-in-the-middle@https://github.com/DataDog/import-in-the-middle#dbaca128f6eb517d71bcf7e6b6b45f21a07128a4":
+  version "1.8.1"
+  resolved "https://github.com/DataDog/import-in-the-middle#dbaca128f6eb517d71bcf7e6b6b45f21a07128a4"
   dependencies:
     acorn "^8.8.2"
     acorn-import-attributes "^1.9.5"

From 17de8593f6715f1379ddba779b90c821348cd63b Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Wed, 19 Jun 2024 17:40:45 +0200
Subject: [PATCH 07/26] support 1.6.0 too

---
 .../datadog-instrumentations/src/vitest.js    | 107 ++++++++++++++----
 1 file changed, 82 insertions(+), 25 deletions(-)

diff --git a/packages/datadog-instrumentations/src/vitest.js b/packages/datadog-instrumentations/src/vitest.js
index eb7e7449a5e..54b891e8496 100644
--- a/packages/datadog-instrumentations/src/vitest.js
+++ b/packages/datadog-instrumentations/src/vitest.js
@@ -1,6 +1,5 @@
 const { addHook, channel, AsyncResource } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
-
 // Needs modifying node_modules/.pnpm/import-in-the-middle@1.7.3/node_modules/import-in-the-middle/index.js
 // the specifiers thing I don't understand
 
@@ -23,45 +22,103 @@ const isVitestTestRunner = (vitestPackage) => {
   return vitestPackage.VitestTestRunner
 }
 
-const isVitestPlugin = (vitestPackage) => {
-  return vitestPackage.B?.name === 'BaseSequencer'
+// maybe not valid???
+// const isVitestPlugin = (vitestPackage) => {
+//   return vitestPackage.B?.name === 'BaseSequencer'
+// }
+
+// from 1.6.0 at least, BaseSequencer is not exported at the same point as startVitest
+// const isVitestPlugin = (vitestPackage) => {
+//   return vitestPackage.s?.name === 'startVitest'
+// }
+
+const isCac = (vitestPackage) => {
+  return vitestPackage.c?.name === 'createCLI'
+}
+
+const isReporterPackage = (vitestPackage) => {
+  // return vitestPackage.a?.name === 'BasicReporter' // it doesn't even export, so can't use BaseReport
+  return vitestPackage.B?.name === 'BaseSequencer' // this is probably used everywhere, so maybe it has info
 }
 
 // --- do I need to specify a file? That's problematic because they contain a hash
 
+// maybe use getEnvPackageName as "start" of session?
+
 // vitestPackage.s only works for 1.1.0 - it doesn't for 1.6.0
 addHook({
   name: 'vitest',
-  versions: ['>=1.1.0 <1.6.0']
+  versions: ['>=1.1.0']
 }, (vitestPackage, frameworkVersion) => {
-  debugger
-  if (isVitestPlugin(vitestPackage)) {
-    //
-    console.log('is S still s???????', vitestPackage)
-    // TODO: will the "s" (minified version) be the same in future versions?
-    shimmer.wrap(vitestPackage, 's', startVitest => async function () {
-      let onFinish
-
-      const flushPromise = new Promise(resolve => {
-        onFinish = resolve
-      })
-
+  if (isCac(vitestPackage)) {
+    const oldCreateCli = vitestPackage.c
+    // could be start session
+    // TODO: change to shimmer
+    vitestPackage.c = function () {
       sessionAsyncResource.runInAsyncScope(() => {
         // TODO: change command to proper command
         testSessionStartCh.publish({ command: 'vitest run', frameworkVersion })
       })
-      const res = await startVitest.apply(this, arguments)
-
-      sessionAsyncResource.runInAsyncScope(() => {
-        // TODO: get proper status
-        testSessionFinishCh.publish({ status: 'pass', onFinish })
+      return oldCreateCli.apply(this, arguments)
+    }
+  }
+  if (isReporterPackage(vitestPackage)) {
+    // this hack seems to work for 1.6.0
+    shimmer.wrap(vitestPackage.B.prototype, 'sort', sort => async function () {
+      // maybe this could also be start, instead of having two different places - easier
+
+      // BY THIS TIME IT'S TOO LATE!!!!! -> changing process.env does not propagate to the worker threads
+      // sessionAsyncResource.runInAsyncScope(() => {
+      //   // TODO: change command to proper command
+      //   testSessionStartCh.publish({ command: 'vitest run', frameworkVersion })
+      // })
+
+      shimmer.wrap(this.ctx, 'exit', exit => async function () {
+        let onFinish
+
+        const flushPromise = new Promise(resolve => {
+          onFinish = resolve
+        })
+        sessionAsyncResource.runInAsyncScope(() => {
+          // TODO: get proper status
+          testSessionFinishCh.publish({ status: 'pass', onFinish })
+        })
+
+        await flushPromise
+
+        return exit.apply(this, arguments)
       })
 
-      await flushPromise
-
-      return res
+      return sort.apply(this, arguments)
     })
   }
+  // if (isVitestPlugin(vitestPackage)) {
+  //   // this DOES NOT WORK for 1.6.0
+  //   // TODO: will the "s" (minified version) be the same in future versions?
+  //   shimmer.wrap(vitestPackage, 's', startVitest => async function () {
+  //     let onFinish
+
+  //     const flushPromise = new Promise(resolve => {
+  //       onFinish = resolve
+  //     })
+
+  //     sessionAsyncResource.runInAsyncScope(() => {
+  //       // TODO: change command to proper command
+  //       testSessionStartCh.publish({ command: 'vitest run', frameworkVersion })
+  //     })
+  //     const res = await startVitest.apply(this, arguments)
+
+  //     console.log('SHOULD NOT LOG THIS!!!!!!!!!!')
+  //     sessionAsyncResource.runInAsyncScope(() => {
+  //       // TODO: get proper status
+  //       testSessionFinishCh.publish({ status: 'pass', onFinish })
+  //     })
+
+  //     await flushPromise
+
+  //     return res
+  //   })
+  // }
 
   if (isVitestTestRunner(vitestPackage)) {
     // test start
@@ -100,7 +157,7 @@ addHook({
   versions: ['>=0.0.0']
 }, vitestPackage => {
   if (vitestPackage.startTests) {
-    shimmer.wrap(vitestPackage, 'startTests', startTests => async function (testPath, vitestTestRunner) {
+    shimmer.wrap(vitestPackage, 'startTests', startTests => async function (testPath) {
       const asyncResource = new AsyncResource('bound-anonymous-fn')
       asyncResource.runInAsyncScope(() => {
         testSuiteStartCh.publish(testPath[0])

From b1ce709eb31ecab448826e584f047017f30453d7 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Thu, 20 Jun 2024 16:03:50 +0200
Subject: [PATCH 08/26] clean up and add error handling

---
 .../test-visibility-failed-suite.mjs          |  20 ++
 .../test-visibility-passed-suite.mjs          |  20 ++
 .../vitest-tests/test-visibility-test-2.mjs   |   8 -
 .../vitest-tests/test-visibility-test.mjs     |   8 -
 integration-tests/vitest/vitest.spec.js       |  52 ++++--
 .../datadog-instrumentations/src/vitest.js    | 172 ++++++++++--------
 packages/datadog-plugin-vitest/src/index.js   |  38 +++-
 7 files changed, 201 insertions(+), 117 deletions(-)
 create mode 100644 integration-tests/ci-visibility/vitest-tests/test-visibility-failed-suite.mjs
 create mode 100644 integration-tests/ci-visibility/vitest-tests/test-visibility-passed-suite.mjs
 delete mode 100644 integration-tests/ci-visibility/vitest-tests/test-visibility-test-2.mjs
 delete mode 100644 integration-tests/ci-visibility/vitest-tests/test-visibility-test.mjs

diff --git a/integration-tests/ci-visibility/vitest-tests/test-visibility-failed-suite.mjs b/integration-tests/ci-visibility/vitest-tests/test-visibility-failed-suite.mjs
new file mode 100644
index 00000000000..75551a80cf4
--- /dev/null
+++ b/integration-tests/ci-visibility/vitest-tests/test-visibility-failed-suite.mjs
@@ -0,0 +1,20 @@
+import { describe, test, expect } from 'vitest'
+import { sum } from './sum'
+
+describe('context', () => {
+  test('can report failed test', () => {
+    expect(sum(1, 2)).to.equal(4)
+  })
+  test('can report more', () => {
+    expect(sum(1, 2)).to.equal(3)
+  })
+})
+
+describe('other context', () => {
+  test('can report passed test', () => {
+    expect(sum(1, 2)).to.equal(3)
+  })
+  test('can report more', () => {
+    expect(sum(1, 2)).to.equal(3)
+  })
+})
diff --git a/integration-tests/ci-visibility/vitest-tests/test-visibility-passed-suite.mjs b/integration-tests/ci-visibility/vitest-tests/test-visibility-passed-suite.mjs
new file mode 100644
index 00000000000..3e351a8e4f8
--- /dev/null
+++ b/integration-tests/ci-visibility/vitest-tests/test-visibility-passed-suite.mjs
@@ -0,0 +1,20 @@
+import { describe, test, expect } from 'vitest'
+import { sum } from './sum'
+
+describe('context', () => {
+  test('can report passed test', () => {
+    expect(sum(1, 2)).to.equal(3)
+  })
+  test('can report more', () => {
+    expect(sum(1, 2)).to.equal(3)
+  })
+})
+
+describe('other context', () => {
+  test('can report passed test', () => {
+    expect(sum(1, 2)).to.equal(3)
+  })
+  test('can report more', () => {
+    expect(sum(1, 2)).to.equal(3)
+  })
+})
diff --git a/integration-tests/ci-visibility/vitest-tests/test-visibility-test-2.mjs b/integration-tests/ci-visibility/vitest-tests/test-visibility-test-2.mjs
deleted file mode 100644
index 9fd618318c0..00000000000
--- a/integration-tests/ci-visibility/vitest-tests/test-visibility-test-2.mjs
+++ /dev/null
@@ -1,8 +0,0 @@
-import { describe, test, expect } from 'vitest'
-import { sum } from './sum'
-
-describe('test visibility 2', () => {
-  test('can report tests 2', () => {
-    expect(sum(1, 2)).to.equal(3)
-  })
-})
diff --git a/integration-tests/ci-visibility/vitest-tests/test-visibility-test.mjs b/integration-tests/ci-visibility/vitest-tests/test-visibility-test.mjs
deleted file mode 100644
index c3e23fed22f..00000000000
--- a/integration-tests/ci-visibility/vitest-tests/test-visibility-test.mjs
+++ /dev/null
@@ -1,8 +0,0 @@
-import { describe, test, expect } from 'vitest'
-import { sum } from './sum'
-
-describe('test visibility', () => {
-  test('can report tests', () => {
-    expect(sum(1, 2)).to.equal(3)
-  })
-})
diff --git a/integration-tests/vitest/vitest.spec.js b/integration-tests/vitest/vitest.spec.js
index f5d85e4abf1..6afb44de920 100644
--- a/integration-tests/vitest/vitest.spec.js
+++ b/integration-tests/vitest/vitest.spec.js
@@ -15,9 +15,8 @@ const {
   TEST_TYPE
 } = require('../../packages/dd-trace/src/plugins/util/test')
 
-// should probably work with 1.6.0
-// tested with 1.1.0
-const versions = ['1.1.0'] // only one I've tested so far
+// tested with 1.6.0
+const versions = ['latest']
 
 versions.forEach((version) => {
   describe(`vitest@${version}`, () => {
@@ -54,33 +53,50 @@ versions.forEach((version) => {
         const testEvents = events.filter(event => event.type === 'test')
 
         assert.include(testSessionEvent.content.resource, 'test_session.vitest run')
-        assert.equal(testSessionEvent.content.meta[TEST_STATUS], 'pass')
+        assert.equal(testSessionEvent.content.meta[TEST_STATUS], 'fail')
         assert.include(testModuleEvent.content.resource, 'test_module.vitest run')
-        assert.equal(testModuleEvent.content.meta[TEST_STATUS], 'pass')
+        assert.equal(testModuleEvent.content.meta[TEST_STATUS], 'fail')
         assert.equal(testSessionEvent.content.meta[TEST_TYPE], 'test')
         assert.equal(testModuleEvent.content.meta[TEST_TYPE], 'test')
 
-        assert.includeMembers(testSuiteEvents.map(suite => suite.content.resource), [
-          'test_suite.ci-visibility/vitest-tests/test-visibility-test.mjs',
-          'test_suite.ci-visibility/vitest-tests/test-visibility-test-2.mjs'
-        ])
+        const passedSuite = testSuiteEvents.find(
+          suite => suite.content.resource === 'test_suite.ci-visibility/vitest-tests/test-visibility-passed-suite.mjs'
+        )
+        assert.equal(passedSuite.content.meta[TEST_STATUS], 'pass')
 
-        // TODO: just check pass
-        assert.includeMembers(testSuiteEvents.map(suite => suite.content.meta[TEST_STATUS]), [
-          'pass',
-        ])
+        const failedSuite = testSuiteEvents.find(
+          suite => suite.content.resource === 'test_suite.ci-visibility/vitest-tests/test-visibility-failed-suite.mjs'
+        )
+        assert.equal(failedSuite.content.meta[TEST_STATUS], 'fail')
+
+        const failedTest = testSuiteEvents.find(
+          ({ content: { resource } }) =>
+            resource === 'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs.can report failed test'
+        )
+
+        assert.equal(failedTest.content.meta[TEST_STATUS], 'fail')
 
+        const passedTests = testEvents.filter(testEvent => testEvent.content.meta[TEST_STATUS] === 'pass')
 
-        assert.includeMembers(testEvents.map(test => test.content.resource), [
-          'ci-visibility/vitest-tests/test-visibility-test.mjs.can report tests',
-          'ci-visibility/vitest-tests/test-visibility-test-2.mjs.can report tests 2',
+        assert.includeMembers(passedTests.map(test => test.content.resource), [
+          'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs.can report more',
+          'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs.can report passed test',
+          'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.can report passed test',
+          'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.can report more'
         ])
 
         // TODO: just check pass
         assert.includeMembers(testEvents.map(test => test.content.meta[TEST_STATUS]), [
           'pass',
+          'pass',
+          'pass',
+          'pass',
+          'pass',
+          'pass',
+          'pass',
+          'fail'
         ])
-      }).then(() => done()).catch(done)
+      }, 25000).then(() => done()).catch(done)
 
       childProcess = exec(
         './node_modules/.bin/vitest run',
@@ -89,7 +105,7 @@ versions.forEach((version) => {
           env: {
             ...getCiVisAgentlessConfig(receiver.port),
             // maybe only in node@20
-            NODE_OPTIONS: '--import dd-trace/register.js -r dd-trace/ci/init', // ESM requires more stuff
+            NODE_OPTIONS: '--import dd-trace/register.js -r dd-trace/ci/init' // ESM requires more stuff
           },
           stdio: 'pipe'
         }
diff --git a/packages/datadog-instrumentations/src/vitest.js b/packages/datadog-instrumentations/src/vitest.js
index 54b891e8496..d5971fceaad 100644
--- a/packages/datadog-instrumentations/src/vitest.js
+++ b/packages/datadog-instrumentations/src/vitest.js
@@ -1,15 +1,18 @@
 const { addHook, channel, AsyncResource } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
-// Needs modifying node_modules/.pnpm/import-in-the-middle@1.7.3/node_modules/import-in-the-middle/index.js
-// the specifiers thing I don't understand
-
-// needs creating a fake file node_modules/vitest/dist/@vitest/spy -- why????
+// Needs removing specifiers check in
+// node_modules/.pnpm/import-in-the-middle@1.7.3/node_modules/import-in-the-middle/index.js
+// TODO: why?
 
+// needs creating a fake file node_modules/vitest/dist/@vitest/spy
+// TODO: why?
 const testStartCh = channel('ci:vitest:test:start')
 const testFinishCh = channel('ci:vitest:test:finish')
+const testErrorCh = channel('ci:vitest:test:error')
 
 const testSuiteStartCh = channel('ci:vitest:test-suite:start')
 const testSuiteFinishCh = channel('ci:vitest:test-suite:finish')
+const testSuiteErrorCh = channel('ci:vitest:test-suite:error')
 
 const testSessionStartCh = channel('ci:vitest:session:start')
 const testSessionFinishCh = channel('ci:vitest:session:finish')
@@ -22,57 +25,71 @@ const isVitestTestRunner = (vitestPackage) => {
   return vitestPackage.VitestTestRunner
 }
 
-// maybe not valid???
-// const isVitestPlugin = (vitestPackage) => {
-//   return vitestPackage.B?.name === 'BaseSequencer'
-// }
+function isCac (vitestPackage) {
+  return vitestPackage.c?.name === 'createCLI'
+}
 
-// from 1.6.0 at least, BaseSequencer is not exported at the same point as startVitest
-// const isVitestPlugin = (vitestPackage) => {
-//   return vitestPackage.s?.name === 'startVitest'
-// }
+function isReporterPackage (vitestPackage) {
+  return vitestPackage.B?.name === 'BaseSequencer'
+}
 
-const isCac = (vitestPackage) => {
-  return vitestPackage.c?.name === 'createCLI'
+function getSessionStatus (state) {
+  if (state.getCountOfFailedTests() > 0) {
+    return 'fail'
+  }
+  if (state.pathsSet.size === 0) {
+    return 'skip'
+  }
+  return 'pass'
 }
 
-const isReporterPackage = (vitestPackage) => {
-  // return vitestPackage.a?.name === 'BasicReporter' // it doesn't even export, so can't use BaseReport
-  return vitestPackage.B?.name === 'BaseSequencer' // this is probably used everywhere, so maybe it has info
+// eslint-disable-next-line
+// From https://github.com/vitest-dev/vitest/blob/51c04e2f44d91322b334f8ccbcdb368facc3f8ec/packages/runner/src/run.ts#L243-L250
+function getVitestTestStatus (test, retryCount) {
+  if (test.result.state !== 'fail') {
+    if (!test.repeats) {
+      return 'pass'
+    } else if (test.repeats && (test.retry ?? 0) === retryCount) {
+      return 'pass'
+    }
+  }
+  return 'fail'
 }
 
-// --- do I need to specify a file? That's problematic because they contain a hash
+function getTestTasks (fileTasks) {
+  const testTasks = []
 
-// maybe use getEnvPackageName as "start" of session?
+  function getTasks (tasks) {
+    for (const task of tasks) {
+      if (task.type === 'test') {
+        testTasks.push(task)
+      } else {
+        getTasks(task.tasks)
+      }
+    }
+  }
+
+  getTasks(fileTasks)
+
+  return testTasks
+}
 
-// vitestPackage.s only works for 1.1.0 - it doesn't for 1.6.0
+// Can't specify file because compiled vitest includes hashes in their files
 addHook({
   name: 'vitest',
-  versions: ['>=1.1.0']
+  versions: ['>=1.6.0']
 }, (vitestPackage, frameworkVersion) => {
   if (isCac(vitestPackage)) {
-    const oldCreateCli = vitestPackage.c
-    // could be start session
-    // TODO: change to shimmer
-    vitestPackage.c = function () {
+    shimmer.wrap(vitestPackage, 'c', oldCreateCli => function () {
       sessionAsyncResource.runInAsyncScope(() => {
-        // TODO: change command to proper command
-        testSessionStartCh.publish({ command: 'vitest run', frameworkVersion })
+        const processArgv = process.argv.slice(2).join(' ')
+        testSessionStartCh.publish({ command: `vitest ${processArgv}`, frameworkVersion })
       })
       return oldCreateCli.apply(this, arguments)
-    }
+    })
   }
   if (isReporterPackage(vitestPackage)) {
-    // this hack seems to work for 1.6.0
     shimmer.wrap(vitestPackage.B.prototype, 'sort', sort => async function () {
-      // maybe this could also be start, instead of having two different places - easier
-
-      // BY THIS TIME IT'S TOO LATE!!!!! -> changing process.env does not propagate to the worker threads
-      // sessionAsyncResource.runInAsyncScope(() => {
-      //   // TODO: change command to proper command
-      //   testSessionStartCh.publish({ command: 'vitest run', frameworkVersion })
-      // })
-
       shimmer.wrap(this.ctx, 'exit', exit => async function () {
         let onFinish
 
@@ -80,8 +97,7 @@ addHook({
           onFinish = resolve
         })
         sessionAsyncResource.runInAsyncScope(() => {
-          // TODO: get proper status
-          testSessionFinishCh.publish({ status: 'pass', onFinish })
+          testSessionFinishCh.publish({ status: getSessionStatus(this.state), onFinish })
         })
 
         await flushPromise
@@ -92,37 +108,12 @@ addHook({
       return sort.apply(this, arguments)
     })
   }
-  // if (isVitestPlugin(vitestPackage)) {
-  //   // this DOES NOT WORK for 1.6.0
-  //   // TODO: will the "s" (minified version) be the same in future versions?
-  //   shimmer.wrap(vitestPackage, 's', startVitest => async function () {
-  //     let onFinish
-
-  //     const flushPromise = new Promise(resolve => {
-  //       onFinish = resolve
-  //     })
-
-  //     sessionAsyncResource.runInAsyncScope(() => {
-  //       // TODO: change command to proper command
-  //       testSessionStartCh.publish({ command: 'vitest run', frameworkVersion })
-  //     })
-  //     const res = await startVitest.apply(this, arguments)
-
-  //     console.log('SHOULD NOT LOG THIS!!!!!!!!!!')
-  //     sessionAsyncResource.runInAsyncScope(() => {
-  //       // TODO: get proper status
-  //       testSessionFinishCh.publish({ status: 'pass', onFinish })
-  //     })
-
-  //     await flushPromise
-
-  //     return res
-  //   })
-  // }
 
   if (isVitestTestRunner(vitestPackage)) {
+    const { VitestTestRunner } = vitestPackage
     // test start
-    shimmer.wrap(vitestPackage.VitestTestRunner.prototype, 'onBeforeTryTask', onBeforeTryTask => async function (task) {
+    // TODO: add test for beforeEach / afterEach and before/after hooks: they're likely tasks too
+    shimmer.wrap(VitestTestRunner.prototype, 'onBeforeTryTask', onBeforeTryTask => async function (task) {
       const asyncResource = new AsyncResource('bound-anonymous-fn')
       taskToAsync.set(task, asyncResource)
 
@@ -134,18 +125,19 @@ addHook({
     })
 
     // test finish
-    shimmer.wrap(vitestPackage.VitestTestRunner.prototype, 'onAfterTryTask', onAfterTryTask => async function (task) {
-      const res = await onAfterTryTask.apply(this, arguments)
+    // TODO: add test for beforeEach / afterEach and before/after hooks: they're likely tasks too
+    shimmer.wrap(VitestTestRunner.prototype, 'onAfterTryTask', onAfterTryTask =>
+      async function (task, { retry: retryCount }) {
+        const result = await onAfterTryTask.apply(this, arguments)
 
-      const asyncResource = taskToAsync.get(task)
+        const testStatus = getVitestTestStatus(task, retryCount)
+        const asyncResource = taskToAsync.get(task)
 
-      asyncResource.runInAsyncScope(() => {
-        // TODO: if no error has been found, it's a pass.
-        // See logic in packages/runner/src/run.ts in vitest after calling onAfterTryTask
-        testFinishCh.publish('pass')
+        asyncResource.runInAsyncScope(() => {
+          testFinishCh.publish(testStatus)
+        })
+        return result
       })
-      return res
-    })
   }
   return vitestPackage
 })
@@ -154,12 +146,12 @@ addHook({
 // only relevant for workers
 addHook({
   name: '@vitest/runner',
-  versions: ['>=0.0.0']
+  versions: ['>=1.6.0']
 }, vitestPackage => {
   if (vitestPackage.startTests) {
     shimmer.wrap(vitestPackage, 'startTests', startTests => async function (testPath) {
-      const asyncResource = new AsyncResource('bound-anonymous-fn')
-      asyncResource.runInAsyncScope(() => {
+      const testSuiteAsyncResource = new AsyncResource('bound-anonymous-fn')
+      testSuiteAsyncResource.runInAsyncScope(() => {
         testSuiteStartCh.publish(testPath[0])
       })
       const startTestsResponse = await startTests.apply(this, arguments)
@@ -169,10 +161,30 @@ addHook({
         onFinish = resolve
       })
 
-      asyncResource.runInAsyncScope(() => {
+      const testTasks = getTestTasks(startTestsResponse[0].tasks)
+      // we don't use getVitestTestStatus here because every hook call has finished, so it's already set
+      // unlike on onAfterTryTask
+      const failedTests = testTasks.filter(task => task.result.state === 'fail')
+
+      failedTests.forEach(failedTask => {
+        const testAsyncResource = taskToAsync.get(failedTask)
+        const { result: { duration, errors } } = failedTask
+        testAsyncResource.runInAsyncScope(() => {
+          // we need to manually finish them here because they won't be finished by onAfterTryTask
+          testErrorCh.publish({ duration, errors })
+        })
+        if (errors.length) {
+          testSuiteAsyncResource.runInAsyncScope(() => {
+            testSuiteErrorCh.publish({ errors })
+          })
+        }
+      })
+
+      testSuiteAsyncResource.runInAsyncScope(() => {
         testSuiteFinishCh.publish({ status: startTestsResponse[0].result.state, onFinish })
       })
 
+      // TODO: fix too frequent flushes
       await onFinishPromise
 
       return startTestsResponse
diff --git a/packages/datadog-plugin-vitest/src/index.js b/packages/datadog-plugin-vitest/src/index.js
index dcce39b1cce..72e54322b7d 100644
--- a/packages/datadog-plugin-vitest/src/index.js
+++ b/packages/datadog-plugin-vitest/src/index.js
@@ -9,6 +9,11 @@ const {
 } = require('../../dd-trace/src/plugins/util/test')
 const { COMPONENT } = require('../../dd-trace/src/constants')
 
+// Milliseconds that we subtract from the error test duration
+// so that they do not overlap with the following test
+// This is because there's some loss of resolution.
+const MILLISECONDS_TO_SUBTRACT_FROM_FAILED_TEST_DURATION = 5
+
 class VitestPlugin extends CiPlugin {
   static get id () {
     return 'vitest'
@@ -36,6 +41,23 @@ class VitestPlugin extends CiPlugin {
       }
     })
 
+    this.addSub('ci:vitest:test:error', ({ duration, errors }) => {
+      const store = storage.getStore()
+      const span = store?.span
+
+      if (span) {
+        span.setTag(TEST_STATUS, 'fail')
+
+        if (errors.length) {
+          // TODO: what to do with multiple errors?
+          const [error] = errors
+          span.setTag('error', error)
+        }
+        span.finish(span._startTime + duration - MILLISECONDS_TO_SUBTRACT_FROM_FAILED_TEST_DURATION) // milliseconds
+        finishAllTraceSpans(span)
+      }
+    })
+
     this.addSub('ci:vitest:test-suite:start', (testSuiteAbsolutePath) => {
       const testSessionSpanContext = this.tracer.extract('text_map', {
         'x-datadog-trace-id': process.env.DD_CIVISIBILITY_TEST_SESSION_ID,
@@ -64,16 +86,26 @@ class VitestPlugin extends CiPlugin {
 
     this.addSub('ci:vitest:test-suite:finish', ({ status, onFinish }) => {
       const store = storage.getStore()
-      if (store && store.span) {
-        const span = store.span
+      const span = store?.span
+      if (span) {
         span.setTag(TEST_STATUS, status)
         span.finish()
         finishAllTraceSpans(span)
       }
-      // TODO: too frequent - how to decrease frequency?
+      // TODO: too frequent flush - find for method in worker to decrease frequency
       this.tracer._exporter.flush(onFinish)
     })
 
+    this.addSub('ci:vitest:test-suite:error', ({ errors }) => {
+      const store = storage.getStore()
+      const span = store?.span
+      if (span) {
+        // TODO: what to do with multiple errors?
+        const [error] = errors
+        span.setTag('error', error)
+      }
+    })
+
     this.addSub('ci:vitest:session:finish', ({ status, onFinish }) => {
       this.testSessionSpan.setTag(TEST_STATUS, status)
       this.testModuleSpan.setTag(TEST_STATUS, status)

From 81d8a1cd3df3f7481109b9671babfe08884269c1 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Thu, 20 Jun 2024 16:19:00 +0200
Subject: [PATCH 09/26] add session error message

---
 packages/datadog-instrumentations/src/vitest.js | 12 +++++++++++-
 packages/datadog-plugin-vitest/src/index.js     |  7 +++++--
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/packages/datadog-instrumentations/src/vitest.js b/packages/datadog-instrumentations/src/vitest.js
index d5971fceaad..b58f368d452 100644
--- a/packages/datadog-instrumentations/src/vitest.js
+++ b/packages/datadog-instrumentations/src/vitest.js
@@ -96,8 +96,18 @@ addHook({
         const flushPromise = new Promise(resolve => {
           onFinish = resolve
         })
+        const failedSuites = this.state.getFailedFilepaths()
+        let error
+        if (failedSuites.length) {
+          error = new Error(`${failedSuites.length} test suites failed.`)
+        }
+
         sessionAsyncResource.runInAsyncScope(() => {
-          testSessionFinishCh.publish({ status: getSessionStatus(this.state), onFinish })
+          testSessionFinishCh.publish({
+            status: getSessionStatus(this.state),
+            onFinish,
+            error
+          })
         })
 
         await flushPromise
diff --git a/packages/datadog-plugin-vitest/src/index.js b/packages/datadog-plugin-vitest/src/index.js
index 72e54322b7d..d19a3e25db5 100644
--- a/packages/datadog-plugin-vitest/src/index.js
+++ b/packages/datadog-plugin-vitest/src/index.js
@@ -106,13 +106,16 @@ class VitestPlugin extends CiPlugin {
       }
     })
 
-    this.addSub('ci:vitest:session:finish', ({ status, onFinish }) => {
+    this.addSub('ci:vitest:session:finish', ({ status, onFinish, error }) => {
       this.testSessionSpan.setTag(TEST_STATUS, status)
       this.testModuleSpan.setTag(TEST_STATUS, status)
+      if (error) {
+        this.testModuleSpan.setTag('error', error)
+        this.testSessionSpan.setTag('error', error)
+      }
       this.testModuleSpan.finish()
       this.testSessionSpan.finish()
       finishAllTraceSpans(this.testSessionSpan)
-      // apparently this is indeed needed
       this.tracer._exporter.flush(onFinish)
     })
   }

From 8c49dbd9d0bc9feb820464c507de9d85a7a743b1 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Thu, 20 Jun 2024 16:27:55 +0200
Subject: [PATCH 10/26] add ci job

---
 .github/workflows/project.yml                  | 17 +++++++++++++++++
 integration-tests/vitest/vitest.spec.js        | 18 ++++++++----------
 .../datadog-instrumentations/src/vitest.js     |  2 +-
 3 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/.github/workflows/project.yml b/.github/workflows/project.yml
index 258d827bdf2..58249badccf 100644
--- a/.github/workflows/project.yml
+++ b/.github/workflows/project.yml
@@ -96,6 +96,23 @@ jobs:
           CYPRESS_VERSION: ${{ matrix.cypress-version }}
           NODE_OPTIONS: '-r ./ci/init'
 
+  integration-vitest:
+    runs-on: ubuntu-latest
+    env:
+      DD_SERVICE: dd-trace-js-integration-tests
+      DD_CIVISIBILITY_AGENTLESS_ENABLED: 1
+      DD_API_KEY: ${{ secrets.DD_API_KEY_CI_APP }}
+    steps:
+      - uses: actions/checkout@v4
+      - uses: ./.github/actions/node/setup
+      - run: yarn install
+      - uses: actions/setup-node@v3
+        with:
+          node-version: 20
+      - run: yarn test:integration:vitest
+        env:
+          NODE_OPTIONS: '-r ./ci/init'
+
   lint:
     runs-on: ubuntu-latest
     steps:
diff --git a/integration-tests/vitest/vitest.spec.js b/integration-tests/vitest/vitest.spec.js
index 6afb44de920..2c0b733ecc7 100644
--- a/integration-tests/vitest/vitest.spec.js
+++ b/integration-tests/vitest/vitest.spec.js
@@ -24,7 +24,6 @@ versions.forEach((version) => {
 
     before(async function () {
       sandbox = await createSandbox([`vitest@${version}`], true)
-      // debugger
       cwd = sandbox.folder
     })
 
@@ -44,7 +43,6 @@ versions.forEach((version) => {
 
     it('can run and report tests', (done) => {
       receiver.gatherPayloadsMaxTimeout(({ url }) => url === '/api/v2/citestcycle', payloads => {
-        debugger
         const events = payloads.flatMap(({ payload }) => payload.events)
 
         const testSessionEvent = events.find(event => event.type === 'test_session_end')
@@ -69,7 +67,7 @@ versions.forEach((version) => {
         )
         assert.equal(failedSuite.content.meta[TEST_STATUS], 'fail')
 
-        const failedTest = testSuiteEvents.find(
+        const failedTest = testEvents.find(
           ({ content: { resource } }) =>
             resource === 'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs.can report failed test'
         )
@@ -79,13 +77,15 @@ versions.forEach((version) => {
         const passedTests = testEvents.filter(testEvent => testEvent.content.meta[TEST_STATUS] === 'pass')
 
         assert.includeMembers(passedTests.map(test => test.content.resource), [
+          'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.can report passed test',
+          'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.can report more',
+          'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.can report passed test',
+          'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.can report more',
           'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs.can report more',
           'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs.can report passed test',
-          'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.can report passed test',
-          'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.can report more'
+          'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs.can report more'
         ])
 
-        // TODO: just check pass
         assert.includeMembers(testEvents.map(test => test.content.meta[TEST_STATUS]), [
           'pass',
           'pass',
@@ -96,7 +96,8 @@ versions.forEach((version) => {
           'pass',
           'fail'
         ])
-      }, 25000).then(() => done()).catch(done)
+        // TODO: check error messages
+      }).then(() => done()).catch(done)
 
       childProcess = exec(
         './node_modules/.bin/vitest run',
@@ -110,9 +111,6 @@ versions.forEach((version) => {
           stdio: 'pipe'
         }
       )
-
-      childProcess.stdout.pipe(process.stdout)
-      childProcess.stderr.pipe(process.stderr)
     })
   })
 })
diff --git a/packages/datadog-instrumentations/src/vitest.js b/packages/datadog-instrumentations/src/vitest.js
index b58f368d452..53cea8d9dd4 100644
--- a/packages/datadog-instrumentations/src/vitest.js
+++ b/packages/datadog-instrumentations/src/vitest.js
@@ -99,7 +99,7 @@ addHook({
         const failedSuites = this.state.getFailedFilepaths()
         let error
         if (failedSuites.length) {
-          error = new Error(`${failedSuites.length} test suites failed.`)
+          error = new Error(`Test suites failed: ${failedSuites.length}.`)
         }
 
         sessionAsyncResource.runInAsyncScope(() => {

From b2a163e33a9f0a1ff6009ab40c7b3983da108d0b Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Thu, 20 Jun 2024 16:30:03 +0200
Subject: [PATCH 11/26] remove volta

---
 package.json | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/package.json b/package.json
index 1fb7ee26b11..7cca672d918 100644
--- a/package.json
+++ b/package.json
@@ -138,8 +138,5 @@
     "sinon-chai": "^3.7.0",
     "tap": "^16.3.7",
     "tape": "^5.6.5"
-  },
-  "volta": {
-    "node": "20.14.0"
   }
 }

From 420e0ddcc99a6a11d54130b0620d8ddf4260de72 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Fri, 21 Jun 2024 13:42:00 +0200
Subject: [PATCH 12/26] support for before after all hooks

---
 .../datadog-instrumentations/src/vitest.js    | 71 ++++++++++++++-----
 packages/datadog-plugin-vitest/src/index.js   | 13 ++--
 2 files changed, 59 insertions(+), 25 deletions(-)

diff --git a/packages/datadog-instrumentations/src/vitest.js b/packages/datadog-instrumentations/src/vitest.js
index 53cea8d9dd4..82c679f9e5c 100644
--- a/packages/datadog-instrumentations/src/vitest.js
+++ b/packages/datadog-instrumentations/src/vitest.js
@@ -56,14 +56,14 @@ function getVitestTestStatus (test, retryCount) {
   return 'fail'
 }
 
-function getTestTasks (fileTasks) {
-  const testTasks = []
+function getTypeTasks (fileTasks, type = 'test') {
+  const typeTasks = []
 
   function getTasks (tasks) {
     for (const task of tasks) {
-      if (task.type === 'test') {
-        testTasks.push(task)
-      } else {
+      if (task.type === type) {
+        typeTasks.push(task)
+      } else if (task.tasks) {
         getTasks(task.tasks)
       }
     }
@@ -71,7 +71,21 @@ function getTestTasks (fileTasks) {
 
   getTasks(fileTasks)
 
-  return testTasks
+  return typeTasks
+}
+
+function getTestName (task) {
+  let testName = ''
+  let currentTask = task
+
+  while (currentTask.suite) {
+    if (currentTask.name) {
+      testName = `${currentTask.name} ${testName}`
+    }
+    currentTask = currentTask.suite
+  }
+
+  return testName
 }
 
 // Can't specify file because compiled vitest includes hashes in their files
@@ -127,9 +141,8 @@ addHook({
       const asyncResource = new AsyncResource('bound-anonymous-fn')
       taskToAsync.set(task, asyncResource)
 
-      // TODO: task.name is just the name of the test() block: include parent describes
       asyncResource.runInAsyncScope(() => {
-        testStartCh.publish({ testName: task.name, testSuiteAbsolutePath: task.suite.file.filepath })
+        testStartCh.publish({ testName: getTestName(task), testSuiteAbsolutePath: task.suite.file.filepath })
       })
       return onBeforeTryTask.apply(this, arguments)
     })
@@ -160,6 +173,8 @@ addHook({
 }, vitestPackage => {
   if (vitestPackage.startTests) {
     shimmer.wrap(vitestPackage, 'startTests', startTests => async function (testPath) {
+      let testSuiteError = null
+
       const testSuiteAsyncResource = new AsyncResource('bound-anonymous-fn')
       testSuiteAsyncResource.runInAsyncScope(() => {
         testSuiteStartCh.publish(testPath[0])
@@ -171,27 +186,49 @@ addHook({
         onFinish = resolve
       })
 
-      const testTasks = getTestTasks(startTestsResponse[0].tasks)
+      const testTasks = getTypeTasks(startTestsResponse[0].tasks)
       // we don't use getVitestTestStatus here because every hook call has finished, so it's already set
       // unlike on onAfterTryTask
-      const failedTests = testTasks.filter(task => task.result.state === 'fail')
+      const failedTests = testTasks.filter(task => task.result?.state === 'fail')
 
+      // Errored tests do not go through onAfterTryTask, so we need to handle them here
       failedTests.forEach(failedTask => {
         const testAsyncResource = taskToAsync.get(failedTask)
         const { result: { duration, errors } } = failedTask
+        let testError
+
+        if (errors?.length) {
+          testError = errors[0]
+        }
+
         testAsyncResource.runInAsyncScope(() => {
-          // we need to manually finish them here because they won't be finished by onAfterTryTask
-          testErrorCh.publish({ duration, errors })
+          testErrorCh.publish({ duration, error: testError })
         })
-        if (errors.length) {
-          testSuiteAsyncResource.runInAsyncScope(() => {
-            testSuiteErrorCh.publish({ errors })
-          })
+        if (errors?.length) {
+          testSuiteError = testError // we store the error to bubble it up to the suite
         }
       })
 
+      const testSuiteResult = startTestsResponse[0].result
+
+      if (testSuiteResult.errors?.length) { // Errors from root level hooks
+        testSuiteError = testSuiteResult.errors[0]
+      } else if (testSuiteResult.state === 'fail') { // Errors from `describe` level hooks
+        const suiteTasks = getTypeTasks(startTestsResponse[0].tasks, 'suite')
+        const failedSuites = suiteTasks.filter(task => task.result?.state === 'fail')
+        if (failedSuites.length && failedSuites[0].result?.errors?.length) {
+          testSuiteError = failedSuites[0].result.errors[0]
+        }
+      }
+
+      if (testSuiteError) {
+        testSuiteAsyncResource.runInAsyncScope(() => {
+          testSuiteErrorCh.publish({ error: testSuiteError })
+        })
+      }
+
       testSuiteAsyncResource.runInAsyncScope(() => {
-        testSuiteFinishCh.publish({ status: startTestsResponse[0].result.state, onFinish })
+        testSuiteFinishCh.publish({ status: testSuiteResult.state, onFinish })
       })
 
       // TODO: fix too frequent flushes
diff --git a/packages/datadog-plugin-vitest/src/index.js b/packages/datadog-plugin-vitest/src/index.js
index d19a3e25db5..ce42aa405f5 100644
--- a/packages/datadog-plugin-vitest/src/index.js
+++ b/packages/datadog-plugin-vitest/src/index.js
@@ -41,16 +41,14 @@ class VitestPlugin extends CiPlugin {
       }
     })
 
-    this.addSub('ci:vitest:test:error', ({ duration, errors }) => {
+    this.addSub('ci:vitest:test:error', ({ duration, error }) => {
       const store = storage.getStore()
       const span = store?.span
 
       if (span) {
         span.setTag(TEST_STATUS, 'fail')
 
-        if (errors.length) {
-          // TODO: what to do with multiple errors?
-          const [error] = errors
+        if (error) {
           span.setTag('error', error)
         }
         span.finish(span._startTime + duration - MILLISECONDS_TO_SUBTRACT_FROM_FAILED_TEST_DURATION) // milliseconds
@@ -96,13 +94,12 @@ class VitestPlugin extends CiPlugin {
       this.tracer._exporter.flush(onFinish)
     })
 
-    this.addSub('ci:vitest:test-suite:error', ({ errors }) => {
+    this.addSub('ci:vitest:test-suite:error', ({ error }) => {
       const store = storage.getStore()
       const span = store?.span
-      if (span) {
-        // TODO: what to do with multiple errors?
-        const [error] = errors
+      if (span && error) {
         span.setTag('error', error)
+        span.setTag(TEST_STATUS, 'fail')
       }
     })
 

From e5ae80f3cd666f8cf7e165bea5855a0df78ea23d Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Fri, 21 Jun 2024 14:51:20 +0200
Subject: [PATCH 13/26] support beforeeach and aftereach

---
 .../datadog-instrumentations/src/vitest.js    | 49 +++++++++++--------
 packages/datadog-plugin-vitest/src/index.js   | 18 ++++++-
 2 files changed, 44 insertions(+), 23 deletions(-)

diff --git a/packages/datadog-instrumentations/src/vitest.js b/packages/datadog-instrumentations/src/vitest.js
index 82c679f9e5c..3275ff96451 100644
--- a/packages/datadog-instrumentations/src/vitest.js
+++ b/packages/datadog-instrumentations/src/vitest.js
@@ -7,7 +7,8 @@ const shimmer = require('../../datadog-shimmer')
 // needs creating a fake file node_modules/vitest/dist/@vitest/spy
 // TODO: why?
 const testStartCh = channel('ci:vitest:test:start')
-const testFinishCh = channel('ci:vitest:test:finish')
+const testFinishTimeCh = channel('ci:vitest:test:finish-time')
+const testPassCh = channel('ci:vitest:test:pass')
 const testErrorCh = channel('ci:vitest:test:error')
 
 const testSuiteStartCh = channel('ci:vitest:test-suite:start')
@@ -153,11 +154,12 @@ addHook({
       async function (task, { retry: retryCount }) {
         const result = await onAfterTryTask.apply(this, arguments)
 
-        const testStatus = getVitestTestStatus(task, retryCount)
+        const status = getVitestTestStatus(task, retryCount)
         const asyncResource = taskToAsync.get(task)
 
+        // We don't finish here because the test might fail in a later hook
         asyncResource.runInAsyncScope(() => {
-          testFinishCh.publish(testStatus)
+          testFinishTimeCh.publish({ status, task })
         })
         return result
       })
@@ -187,25 +189,30 @@ addHook({
       })
 
       const testTasks = getTypeTasks(startTestsResponse[0].tasks)
-      // we don't use getVitestTestStatus here because every hook call has finished, so it's already set
-      // unlike on onAfterTryTask
-      const failedTests = testTasks.filter(task => task.result?.state === 'fail')
-
-      // Errored tests do not go through onAfterTryTask, so we need to handle them here
-      failedTests.forEach(failedTask => {
-        const testAsyncResource = taskToAsync.get(failedTask)
-        const { result: { duration, errors } } = failedTask
-        let testError
-
-        if (errors?.length) {
-          testError = errors[0]
-        }
 
-        testAsyncResource.runInAsyncScope(() => {
-          testErrorCh.publish({ duration, error: testError })
-        })
-        if (errors?.length) {
-          testSuiteError = testError // we store the error to bubble it up to the suite
+      testTasks.forEach(task => {
+        const testAsyncResource = taskToAsync.get(task)
+        const { result: { state, duration, errors } } = task
+
+        // what about 'skip'??
+        if (state === 'pass') {
+          testAsyncResource.runInAsyncScope(() => {
+            testPassCh.publish({ task })
+          })
+        } else if (state === 'fail') {
+          // If it's failing, we have no accurate finish time, so we have to use `duration`
+          let testError
+
+          if (errors?.length) {
+            testError = errors[0]
+          }
+
+          testAsyncResource.runInAsyncScope(() => {
+            testErrorCh.publish({ duration, error: testError })
+          })
+          if (errors?.length) {
+            testSuiteError = testError // we store the error to bubble it up to the suite
+          }
         }
       })
 
diff --git a/packages/datadog-plugin-vitest/src/index.js b/packages/datadog-plugin-vitest/src/index.js
index ce42aa405f5..bcec8d51793 100644
--- a/packages/datadog-plugin-vitest/src/index.js
+++ b/packages/datadog-plugin-vitest/src/index.js
@@ -21,6 +21,9 @@ class VitestPlugin extends CiPlugin {
 
   constructor (...args) {
     super(...args)
+
+    this.taskToFinishTime = new WeakMap()
+
     this.addSub('ci:vitest:test:start', ({ testName, testSuiteAbsolutePath }) => {
       const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
       const store = storage.getStore()
@@ -29,14 +32,25 @@ class VitestPlugin extends CiPlugin {
       this.enter(span, store)
     })
 
-    this.addSub('ci:vitest:test:finish', (status) => {
+    // If there's a hook error, this is called AND THEN test:error - which will not work
+    this.addSub('ci:vitest:test:finish-time', ({ status, task }) => {
       const store = storage.getStore()
       const span = store?.span
 
+      // we store the finish time
       if (span) {
         span.setTag(TEST_STATUS, status)
+        this.taskToFinishTime.set(task, span._getTime())
+      }
+    })
 
-        span.finish()
+    this.addSub('ci:vitest:test:pass', ({ task }) => {
+      const store = storage.getStore()
+      const span = store?.span
+
+      if (span) {
+        span.setTag(TEST_STATUS, 'pass')
+        span.finish(this.taskToFinishTime.get(task))
         finishAllTraceSpans(span)
       }
     })

From 5fd20c78dbb80bb3d3e2144e77408ebf0b332318 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Fri, 21 Jun 2024 15:43:54 +0200
Subject: [PATCH 14/26] add test source file

---
 packages/datadog-plugin-vitest/src/index.js | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/packages/datadog-plugin-vitest/src/index.js b/packages/datadog-plugin-vitest/src/index.js
index bcec8d51793..9d314320c65 100644
--- a/packages/datadog-plugin-vitest/src/index.js
+++ b/packages/datadog-plugin-vitest/src/index.js
@@ -5,7 +5,8 @@ const {
   TEST_STATUS,
   finishAllTraceSpans,
   getTestSuitePath,
-  getTestSuiteCommonTags
+  getTestSuiteCommonTags,
+  TEST_SOURCE_FILE
 } = require('../../dd-trace/src/plugins/util/test')
 const { COMPONENT } = require('../../dd-trace/src/constants')
 
@@ -27,7 +28,14 @@ class VitestPlugin extends CiPlugin {
     this.addSub('ci:vitest:test:start', ({ testName, testSuiteAbsolutePath }) => {
       const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
       const store = storage.getStore()
-      const span = this.startTestSpan(testName, testSuite, this.testSuiteSpan)
+      const span = this.startTestSpan(
+        testName,
+        testSuite,
+        this.testSuiteSpan,
+        {
+          [TEST_SOURCE_FILE]: testSuite
+        }
+      )
 
       this.enter(span, store)
     })

From 7edbaf66711e333bd21e8ea7b313762f07bf7ec5 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Fri, 21 Jun 2024 15:54:17 +0200
Subject: [PATCH 15/26] better tests

---
 .../test-visibility-failed-hooks.mjs          | 26 ++++++++
 .../test-visibility-failed-suite.mjs          | 15 ++++-
 integration-tests/vitest/vitest.spec.js       | 64 +++++++++++--------
 3 files changed, 75 insertions(+), 30 deletions(-)
 create mode 100644 integration-tests/ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs

diff --git a/integration-tests/ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs b/integration-tests/ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs
new file mode 100644
index 00000000000..a97f95e0df1
--- /dev/null
+++ b/integration-tests/ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs
@@ -0,0 +1,26 @@
+import { describe, test, expect, beforeEach, afterEach } from 'vitest'
+import { sum } from './sum'
+
+describe('context', () => {
+  beforeEach(() => {
+    throw new Error('failed before each')
+  })
+  test('can report failed test', () => {
+    expect(sum(1, 2)).to.equal(4)
+  })
+  test('can report more', () => {
+    expect(sum(1, 2)).to.equal(3)
+  })
+})
+
+describe('other context', () => {
+  afterEach(() => {
+    throw new Error('failed after each')
+  })
+  test('can report passed test', () => {
+    expect(sum(1, 2)).to.equal(3)
+  })
+  test('can report more', () => {
+    expect(sum(1, 2)).to.equal(3)
+  })
+})
diff --git a/integration-tests/ci-visibility/vitest-tests/test-visibility-failed-suite.mjs b/integration-tests/ci-visibility/vitest-tests/test-visibility-failed-suite.mjs
index 75551a80cf4..f2df345a87f 100644
--- a/integration-tests/ci-visibility/vitest-tests/test-visibility-failed-suite.mjs
+++ b/integration-tests/ci-visibility/vitest-tests/test-visibility-failed-suite.mjs
@@ -1,16 +1,25 @@
-import { describe, test, expect } from 'vitest'
+import { describe, test, expect, beforeEach, afterEach } from 'vitest'
 import { sum } from './sum'
 
-describe('context', () => {
+let preparedValue = 1
+
+describe('test-visibility-failed-suite-first-describe', () => {
+  beforeEach(() => {
+    preparedValue = 2
+  })
   test('can report failed test', () => {
     expect(sum(1, 2)).to.equal(4)
   })
   test('can report more', () => {
     expect(sum(1, 2)).to.equal(3)
+    expect(preparedValue).to.equal(2)
   })
 })
 
-describe('other context', () => {
+describe('test-visibility-failed-suite-second-describe', () => {
+  afterEach(() => {
+    preparedValue = 1
+  })
   test('can report passed test', () => {
     expect(sum(1, 2)).to.equal(3)
   })
diff --git a/integration-tests/vitest/vitest.spec.js b/integration-tests/vitest/vitest.spec.js
index 2c0b733ecc7..df2efbc208c 100644
--- a/integration-tests/vitest/vitest.spec.js
+++ b/integration-tests/vitest/vitest.spec.js
@@ -67,35 +67,45 @@ versions.forEach((version) => {
         )
         assert.equal(failedSuite.content.meta[TEST_STATUS], 'fail')
 
-        const failedTest = testEvents.find(
-          ({ content: { resource } }) =>
-            resource === 'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs.can report failed test'
+        const failedSuiteHooks = testSuiteEvents.find(
+          suite => suite.content.resource === 'test_suite.ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs'
+        )
+        assert.equal(failedSuiteHooks.content.meta[TEST_STATUS], 'fail')
+
+        assert.includeMembers(testEvents.map(test => test.content.resource),
+          [
+            'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs' +
+            '.test-visibility-failed-suite-first-describe can report failed test ',
+            'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs' +
+            '.test-visibility-failed-suite-first-describe can report more ',
+            'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs' +
+            '.test-visibility-failed-suite-second-describe can report passed test ',
+            'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs' +
+            '.test-visibility-failed-suite-second-describe can report more ',
+            'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.context can report passed test ',
+            'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.context can report more ',
+            'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.other context can report passed test ',
+            'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.other context can report more ',
+            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.context can report failed test ',
+            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.context can report more ',
+            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.other context can report passed test ',
+            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.other context can report more '
+          ]
         )
 
-        assert.equal(failedTest.content.meta[TEST_STATUS], 'fail')
-
-        const passedTests = testEvents.filter(testEvent => testEvent.content.meta[TEST_STATUS] === 'pass')
-
-        assert.includeMembers(passedTests.map(test => test.content.resource), [
-          'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.can report passed test',
-          'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.can report more',
-          'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.can report passed test',
-          'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.can report more',
-          'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs.can report more',
-          'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs.can report passed test',
-          'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs.can report more'
-        ])
-
-        assert.includeMembers(testEvents.map(test => test.content.meta[TEST_STATUS]), [
-          'pass',
-          'pass',
-          'pass',
-          'pass',
-          'pass',
-          'pass',
-          'pass',
-          'fail'
-        ])
+        const failedTests = testEvents.filter(test => test.content.meta[TEST_STATUS] === 'fail')
+
+        assert.includeMembers(
+          failedTests.map(test => test.content.resource),
+          [
+            'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs' +
+            '.test-visibility-failed-suite-first-describe can report failed test ',
+            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.context can report failed test ',
+            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.context can report more ',
+            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.other context can report passed test ',
+            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.other context can report more '
+          ]
+        )
         // TODO: check error messages
       }).then(() => done()).catch(done)
 

From 482db61e9cc7e1ad7b8b21486db51c1a0dd983bf Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Mon, 24 Jun 2024 11:37:18 +0200
Subject: [PATCH 16/26] ts definitions

---
 docs/test.ts | 2 ++
 index.d.ts   | 9 ++++++++-
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/docs/test.ts b/docs/test.ts
index 380bfd6bba9..9b11d344841 100644
--- a/docs/test.ts
+++ b/docs/test.ts
@@ -363,6 +363,8 @@ tracer.use('sharedb');
 tracer.use('sharedb', sharedbOptions);
 tracer.use('tedious');
 tracer.use('undici');
+tracer.use('vitest');
+tracer.use('vitest', { service: 'vitest-service' });
 tracer.use('winston');
 
 tracer.use('express', false)
diff --git a/index.d.ts b/index.d.ts
index d0b634f5dd8..7d5846f8af2 100644
--- a/index.d.ts
+++ b/index.d.ts
@@ -198,6 +198,7 @@ interface Plugins {
   "sharedb": tracer.plugins.sharedb;
   "tedious": tracer.plugins.tedious;
   "undici": tracer.plugins.undici;
+  "vitest": tracer.plugins.vitest;
   "winston": tracer.plugins.winston;
 }
 
@@ -1556,7 +1557,7 @@ declare namespace tracer {
 
     /**
      * This plugin automatically instruments the
-     * [jest](https://github.com/facebook/jest) module.
+     * [jest](https://github.com/jestjs/jest) module.
      */
     interface jest extends Integration {}
 
@@ -1839,6 +1840,12 @@ declare namespace tracer {
      */
     interface undici extends HttpClient {}
 
+    /**
+     * This plugin automatically instruments the
+     * [vitest](https://github.com/vitest-dev/vitest) module.
+     */
+    interface vitest extends Integration {}
+
     /**
      * This plugin patches the [winston](https://github.com/winstonjs/winston)
      * to automatically inject trace identifiers in log records when the

From ba5088c27aed5a0ba6798d10abf4d1791fc8aa15 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Mon, 24 Jun 2024 14:04:04 +0200
Subject: [PATCH 17/26] add support for skipped tests

---
 .../test-visibility-passed-suite.mjs          | 12 ++++
 integration-tests/vitest/vitest.spec.js       | 52 ++++++++++------
 .../datadog-instrumentations/src/vitest.js    | 59 ++++++++++---------
 packages/datadog-plugin-vitest/src/index.js   | 13 ++++
 4 files changed, 91 insertions(+), 45 deletions(-)

diff --git a/integration-tests/ci-visibility/vitest-tests/test-visibility-passed-suite.mjs b/integration-tests/ci-visibility/vitest-tests/test-visibility-passed-suite.mjs
index 3e351a8e4f8..c2cf93431d8 100644
--- a/integration-tests/ci-visibility/vitest-tests/test-visibility-passed-suite.mjs
+++ b/integration-tests/ci-visibility/vitest-tests/test-visibility-passed-suite.mjs
@@ -17,4 +17,16 @@ describe('other context', () => {
   test('can report more', () => {
     expect(sum(1, 2)).to.equal(3)
   })
+  test.skip('can skip', () => {
+    expect(sum(1, 2)).to.equal(3)
+  })
+  test.todo('can todo', () => {
+    expect(sum(1, 2)).to.equal(3)
+  })
+  // eslint-disable-next-line
+  test('can programmatic skip', (context) => {
+    // eslint-disable-next-line
+    context.skip()
+    expect(sum(1, 2)).to.equal(3)
+  })
 })
diff --git a/integration-tests/vitest/vitest.spec.js b/integration-tests/vitest/vitest.spec.js
index df2efbc208c..88c2d50c383 100644
--- a/integration-tests/vitest/vitest.spec.js
+++ b/integration-tests/vitest/vitest.spec.js
@@ -75,21 +75,23 @@ versions.forEach((version) => {
         assert.includeMembers(testEvents.map(test => test.content.resource),
           [
             'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs' +
-            '.test-visibility-failed-suite-first-describe can report failed test ',
+            '.test-visibility-failed-suite-first-describe can report failed test',
             'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs' +
-            '.test-visibility-failed-suite-first-describe can report more ',
+            '.test-visibility-failed-suite-first-describe can report more',
             'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs' +
-            '.test-visibility-failed-suite-second-describe can report passed test ',
+            '.test-visibility-failed-suite-second-describe can report passed test',
             'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs' +
-            '.test-visibility-failed-suite-second-describe can report more ',
-            'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.context can report passed test ',
-            'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.context can report more ',
-            'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.other context can report passed test ',
-            'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.other context can report more ',
-            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.context can report failed test ',
-            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.context can report more ',
-            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.other context can report passed test ',
-            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.other context can report more '
+            '.test-visibility-failed-suite-second-describe can report more',
+            'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.context can report passed test',
+            'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.context can report more',
+            'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.other context can report passed test',
+            'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.other context can report more',
+            'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.other context can skip',
+            'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.other context can todo',
+            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.context can report failed test',
+            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.context can report more',
+            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.other context can report passed test',
+            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.other context can report more'
           ]
         )
 
@@ -99,11 +101,22 @@ versions.forEach((version) => {
           failedTests.map(test => test.content.resource),
           [
             'ci-visibility/vitest-tests/test-visibility-failed-suite.mjs' +
-            '.test-visibility-failed-suite-first-describe can report failed test ',
-            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.context can report failed test ',
-            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.context can report more ',
-            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.other context can report passed test ',
-            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.other context can report more '
+            '.test-visibility-failed-suite-first-describe can report failed test',
+            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.context can report failed test',
+            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.context can report more',
+            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.other context can report passed test',
+            'ci-visibility/vitest-tests/test-visibility-failed-hooks.mjs.other context can report more'
+          ]
+        )
+
+        const skippedTests = testEvents.filter(test => test.content.meta[TEST_STATUS] === 'skip')
+
+        assert.includeMembers(
+          skippedTests.map(test => test.content.resource),
+          [
+            'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.other context can skip',
+            'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.other context can todo',
+            'ci-visibility/vitest-tests/test-visibility-passed-suite.mjs.other context can programmatic skip'
           ]
         )
         // TODO: check error messages
@@ -116,11 +129,14 @@ versions.forEach((version) => {
           env: {
             ...getCiVisAgentlessConfig(receiver.port),
             // maybe only in node@20
-            NODE_OPTIONS: '--import dd-trace/register.js -r dd-trace/ci/init' // ESM requires more stuff
+            NODE_OPTIONS: '--import dd-trace/register.js -r dd-trace/ci/init' // ESM requires more flags
           },
           stdio: 'pipe'
         }
       )
+
+      childProcess.stderr.pipe(process.stderr)
+      childProcess.stdout.pipe(process.stdout)
     })
   })
 })
diff --git a/packages/datadog-instrumentations/src/vitest.js b/packages/datadog-instrumentations/src/vitest.js
index 3275ff96451..65a8f967ae5 100644
--- a/packages/datadog-instrumentations/src/vitest.js
+++ b/packages/datadog-instrumentations/src/vitest.js
@@ -10,6 +10,7 @@ const testStartCh = channel('ci:vitest:test:start')
 const testFinishTimeCh = channel('ci:vitest:test:finish-time')
 const testPassCh = channel('ci:vitest:test:pass')
 const testErrorCh = channel('ci:vitest:test:error')
+const testSkipCh = channel('ci:vitest:test:skip')
 
 const testSuiteStartCh = channel('ci:vitest:test-suite:start')
 const testSuiteFinishCh = channel('ci:vitest:test-suite:finish')
@@ -76,10 +77,10 @@ function getTypeTasks (fileTasks, type = 'test') {
 }
 
 function getTestName (task) {
-  let testName = ''
-  let currentTask = task
+  let testName = task.name
+  let currentTask = task.suite
 
-  while (currentTask.suite) {
+  while (currentTask) {
     if (currentTask.name) {
       testName = `${currentTask.name} ${testName}`
     }
@@ -136,8 +137,7 @@ addHook({
 
   if (isVitestTestRunner(vitestPackage)) {
     const { VitestTestRunner } = vitestPackage
-    // test start
-    // TODO: add test for beforeEach / afterEach and before/after hooks: they're likely tasks too
+    // test start (only tests that are not marked as skip or todo)
     shimmer.wrap(VitestTestRunner.prototype, 'onBeforeTryTask', onBeforeTryTask => async function (task) {
       const asyncResource = new AsyncResource('bound-anonymous-fn')
       taskToAsync.set(task, asyncResource)
@@ -148,8 +148,7 @@ addHook({
       return onBeforeTryTask.apply(this, arguments)
     })
 
-    // test finish
-    // TODO: add test for beforeEach / afterEach and before/after hooks: they're likely tasks too
+    // test finish (only passed tests)
     shimmer.wrap(VitestTestRunner.prototype, 'onAfterTryTask', onAfterTryTask =>
       async function (task, { retry: retryCount }) {
         const result = await onAfterTryTask.apply(this, arguments)
@@ -192,27 +191,33 @@ addHook({
 
       testTasks.forEach(task => {
         const testAsyncResource = taskToAsync.get(task)
-        const { result: { state, duration, errors } } = task
-
-        // what about 'skip'??
-        if (state === 'pass') {
-          testAsyncResource.runInAsyncScope(() => {
-            testPassCh.publish({ task })
-          })
-        } else if (state === 'fail') {
-          // If it's failing, we have no accurate finish time, so we have to use `duration`
-          let testError
-
-          if (errors?.length) {
-            testError = errors[0]
-          }
-
-          testAsyncResource.runInAsyncScope(() => {
-            testErrorCh.publish({ duration, error: testError })
-          })
-          if (errors?.length) {
-            testSuiteError = testError // we store the error to bubble it up to the suite
+        const { result } = task
+
+        if (result) {
+          const { state, duration, errors } = result
+          if (state === 'skip') { // programmatic skip
+            testSkipCh.publish({ testName: getTestName(task), testSuiteAbsolutePath: task.suite.file.filepath })
+          } else if (state === 'pass') {
+            testAsyncResource.runInAsyncScope(() => {
+              testPassCh.publish({ task })
+            })
+          } else if (state === 'fail') {
+            // If it's failing, we have no accurate finish time, so we have to use `duration`
+            let testError
+
+            if (errors?.length) {
+              testError = errors[0]
+            }
+
+            testAsyncResource.runInAsyncScope(() => {
+              testErrorCh.publish({ duration, error: testError })
+            })
+            if (errors?.length) {
+              testSuiteError = testError // we store the error to bubble it up to the suite
+            }
           }
+        } else { // test.skip or test.todo
+          testSkipCh.publish({ testName: getTestName(task), testSuiteAbsolutePath: task.suite.file.filepath })
         }
       })
 
diff --git a/packages/datadog-plugin-vitest/src/index.js b/packages/datadog-plugin-vitest/src/index.js
index 9d314320c65..4bbd70090e1 100644
--- a/packages/datadog-plugin-vitest/src/index.js
+++ b/packages/datadog-plugin-vitest/src/index.js
@@ -78,6 +78,19 @@ class VitestPlugin extends CiPlugin {
       }
     })
 
+    this.addSub('ci:vitest:test:skip', ({ testName, testSuiteAbsolutePath }) => {
+      const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
+      this.startTestSpan(
+        testName,
+        testSuite,
+        this.testSuiteSpan,
+        {
+          [TEST_SOURCE_FILE]: testSuite,
+          [TEST_STATUS]: 'skip'
+        }
+      ).finish()
+    })
+
     this.addSub('ci:vitest:test-suite:start', (testSuiteAbsolutePath) => {
       const testSessionSpanContext = this.tracer.extract('text_map', {
         'x-datadog-trace-id': process.env.DD_CIVISIBILITY_TEST_SESSION_ID,

From 5c570a9bdfa1c6b478ac6b19cc0bb054d9a9f7f2 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Mon, 24 Jun 2024 14:09:10 +0200
Subject: [PATCH 18/26] update comment

---
 packages/datadog-plugin-vitest/src/index.js | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/packages/datadog-plugin-vitest/src/index.js b/packages/datadog-plugin-vitest/src/index.js
index 4bbd70090e1..c47467528e6 100644
--- a/packages/datadog-plugin-vitest/src/index.js
+++ b/packages/datadog-plugin-vitest/src/index.js
@@ -40,12 +40,12 @@ class VitestPlugin extends CiPlugin {
       this.enter(span, store)
     })
 
-    // If there's a hook error, this is called AND THEN test:error - which will not work
     this.addSub('ci:vitest:test:finish-time', ({ status, task }) => {
       const store = storage.getStore()
       const span = store?.span
 
-      // we store the finish time
+      // we store the finish time to finish at a later hook
+      // this is because the test might fail at a `afterEach` hook
       if (span) {
         span.setTag(TEST_STATUS, status)
         this.taskToFinishTime.set(task, span._getTime())

From 1ddeee81ccadef5094ff2fbfa95a35ec0467daf3 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Tue, 25 Jun 2024 12:58:55 +0200
Subject: [PATCH 19/26] do not use iitm version and add wildcard logic

---
 integration-tests/vitest/vitest.spec.js       |   3 -
 package.json                                  |   4 +-
 .../src/helpers/hook.js                       |   5 +-
 .../src/helpers/register.js                   |   9 +-
 .../datadog-instrumentations/src/vitest.js    | 273 +++++++++---------
 yarn.lock                                     |   5 +-
 6 files changed, 155 insertions(+), 144 deletions(-)

diff --git a/integration-tests/vitest/vitest.spec.js b/integration-tests/vitest/vitest.spec.js
index 88c2d50c383..3065e52b202 100644
--- a/integration-tests/vitest/vitest.spec.js
+++ b/integration-tests/vitest/vitest.spec.js
@@ -134,9 +134,6 @@ versions.forEach((version) => {
           stdio: 'pipe'
         }
       )
-
-      childProcess.stderr.pipe(process.stderr)
-      childProcess.stdout.pipe(process.stdout)
     })
   })
 })
diff --git a/package.json b/package.json
index 90d40e849a4..484a9fb1190 100644
--- a/package.json
+++ b/package.json
@@ -37,9 +37,9 @@
     "test:integration:cypress": "mocha --colors --timeout 30000 -r \"packages/dd-trace/test/setup/core.js\" \"integration-tests/cypress/*.spec.js\"",
     "test:integration:playwright": "mocha --colors --timeout 30000 -r \"packages/dd-trace/test/setup/core.js\" \"integration-tests/playwright/*.spec.js\"",
     "test:integration:selenium": "mocha --colors --timeout 30000 -r \"packages/dd-trace/test/setup/core.js\" \"integration-tests/selenium/*.spec.js\"",
+    "test:integration:vitest": "mocha --colors --timeout 30000 -r \"packages/dd-trace/test/setup/core.js\" \"integration-tests/vitest/*.spec.js\"",
     "test:integration:profiler": "mocha --colors --timeout 180000 -r \"packages/dd-trace/test/setup/core.js\" \"integration-tests/profiler/*.spec.js\"",
     "test:integration:serverless": "mocha --colors --timeout 30000 -r \"packages/dd-trace/test/setup/core.js\" \"integration-tests/serverless/*.spec.js\"",
-    "test:integration:vitest": "mocha --colors --timeout 30000 -r \"packages/dd-trace/test/setup/core.js\" \"integration-tests/vitest/*.spec.js\"",
     "test:integration:plugins": "mocha --colors --exit -r \"packages/dd-trace/test/setup/mocha.js\" \"packages/datadog-plugin-@($(echo $PLUGINS))/test/integration-test/**/*.spec.js\"",
     "test:unit:plugins": "mocha --colors --exit -r \"packages/dd-trace/test/setup/mocha.js\" \"packages/datadog-instrumentations/test/@($(echo $PLUGINS)).spec.js\" \"packages/datadog-plugin-@($(echo $PLUGINS))/test/**/*.spec.js\" --exclude \"packages/datadog-plugin-@($(echo $PLUGINS))/test/integration-test/**/*.spec.js\"",
     "test:shimmer": "mocha --colors 'packages/datadog-shimmer/test/**/*.spec.js'",
@@ -82,7 +82,7 @@
     "crypto-randomuuid": "^1.0.0",
     "dc-polyfill": "^0.1.4",
     "ignore": "^5.2.4",
-    "import-in-the-middle": "https://github.com/DataDog/import-in-the-middle#dbaca128f6eb517d71bcf7e6b6b45f21a07128a4",
+    "import-in-the-middle": "^1.8.1",
     "int64-buffer": "^0.1.9",
     "istanbul-lib-coverage": "3.2.0",
     "jest-docblock": "^29.7.0",
diff --git a/packages/datadog-instrumentations/src/helpers/hook.js b/packages/datadog-instrumentations/src/helpers/hook.js
index 23046e00184..b56e2794ed0 100644
--- a/packages/datadog-instrumentations/src/helpers/hook.js
+++ b/packages/datadog-instrumentations/src/helpers/hook.js
@@ -20,8 +20,7 @@ function Hook (modules, onrequire) {
     const parts = [moduleBaseDir, moduleName].filter(v => v)
     const filename = path.join(...parts)
 
-    // this does not seem to work with IITM: the same file seems to import different stuff
-    // if (this._patched[filename]) return moduleExports
+    if (this._patched[filename]) return moduleExports
 
     this._patched[filename] = true
 
@@ -29,7 +28,7 @@ function Hook (modules, onrequire) {
   }
 
   this._ritmHook = ritm(modules, {}, safeHook)
-  this._iitmHook = iitm(modules, {}, (moduleExports, moduleName, moduleBaseDir) => {
+  this._iitmHook = iitm(modules, { internals: true }, (moduleExports, moduleName, moduleBaseDir) => {
     // TODO: Move this logic to import-in-the-middle and only do it for CommonJS
     // modules and not ESM. In the meantime, all the modules we instrument are
     // CommonJS modules for which the default export is always moved to
diff --git a/packages/datadog-instrumentations/src/helpers/register.js b/packages/datadog-instrumentations/src/helpers/register.js
index eba90d6a980..f8d1d2e0493 100644
--- a/packages/datadog-instrumentations/src/helpers/register.js
+++ b/packages/datadog-instrumentations/src/helpers/register.js
@@ -35,6 +35,8 @@ if (DD_TRACE_DEBUG && DD_TRACE_DEBUG.toLowerCase() !== 'false') {
   setImmediate(checkRequireCache.checkForPotentialConflicts)
 }
 
+const isWildCard = (name) => name.includes('*')
+
 // TODO: make this more efficient
 for (const packageName of names) {
   if (disabledInstrumentations.has(packageName)) continue
@@ -59,7 +61,12 @@ for (const packageName of names) {
         hook[HOOK_SYMBOL] = new WeakMap()
       }
 
-      if (moduleName === fullFilename) {
+      // Some libraries include a hash in their filenames when installed,
+      // so our instrumentation has to include a * to match them for more than a single version.
+      const matchesFile = moduleName === fullFilename ||
+        (isWildCard(fullFilename) && moduleName.includes(fullFilename.replace('*', '')))
+
+      if (matchesFile) {
         const version = moduleVersion || getVersion(moduleBaseDir)
 
         if (matchVersion(version, versions)) {
diff --git a/packages/datadog-instrumentations/src/vitest.js b/packages/datadog-instrumentations/src/vitest.js
index 65a8f967ae5..ad8e8c6f247 100644
--- a/packages/datadog-instrumentations/src/vitest.js
+++ b/packages/datadog-instrumentations/src/vitest.js
@@ -1,21 +1,19 @@
 const { addHook, channel, AsyncResource } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
-// Needs removing specifiers check in
-// node_modules/.pnpm/import-in-the-middle@1.7.3/node_modules/import-in-the-middle/index.js
-// TODO: why?
 
-// needs creating a fake file node_modules/vitest/dist/@vitest/spy
-// TODO: why?
+// test hooks
 const testStartCh = channel('ci:vitest:test:start')
 const testFinishTimeCh = channel('ci:vitest:test:finish-time')
 const testPassCh = channel('ci:vitest:test:pass')
 const testErrorCh = channel('ci:vitest:test:error')
 const testSkipCh = channel('ci:vitest:test:skip')
 
+// test suite hooks
 const testSuiteStartCh = channel('ci:vitest:test-suite:start')
 const testSuiteFinishCh = channel('ci:vitest:test-suite:finish')
 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')
 
@@ -23,14 +21,6 @@ const taskToAsync = new WeakMap()
 
 const sessionAsyncResource = new AsyncResource('bound-anonymous-fn')
 
-const isVitestTestRunner = (vitestPackage) => {
-  return vitestPackage.VitestTestRunner
-}
-
-function isCac (vitestPackage) {
-  return vitestPackage.c?.name === 'createCLI'
-}
-
 function isReporterPackage (vitestPackage) {
   return vitestPackage.B?.name === 'BaseSequencer'
 }
@@ -90,79 +80,97 @@ function getTestName (task) {
   return testName
 }
 
-// Can't specify file because compiled vitest includes hashes in their files
 addHook({
   name: 'vitest',
-  versions: ['>=1.6.0']
-}, (vitestPackage, frameworkVersion) => {
-  if (isCac(vitestPackage)) {
-    shimmer.wrap(vitestPackage, 'c', oldCreateCli => function () {
-      sessionAsyncResource.runInAsyncScope(() => {
-        const processArgv = process.argv.slice(2).join(' ')
-        testSessionStartCh.publish({ command: `vitest ${processArgv}`, frameworkVersion })
-      })
-      return oldCreateCli.apply(this, arguments)
+  versions: ['>=1.6.0'],
+  file: 'dist/runners.js'
+}, (vitestPackage) => {
+  const { VitestTestRunner } = vitestPackage
+  // test start (only tests that are not marked as skip or todo)
+  shimmer.wrap(VitestTestRunner.prototype, 'onBeforeTryTask', onBeforeTryTask => async function (task) {
+    const asyncResource = new AsyncResource('bound-anonymous-fn')
+    taskToAsync.set(task, asyncResource)
+
+    asyncResource.runInAsyncScope(() => {
+      testStartCh.publish({ testName: getTestName(task), testSuiteAbsolutePath: task.suite.file.filepath })
     })
-  }
-  if (isReporterPackage(vitestPackage)) {
-    shimmer.wrap(vitestPackage.B.prototype, 'sort', sort => async function () {
-      shimmer.wrap(this.ctx, 'exit', exit => async function () {
-        let onFinish
+    return onBeforeTryTask.apply(this, arguments)
+  })
 
-        const flushPromise = new Promise(resolve => {
-          onFinish = resolve
-        })
-        const failedSuites = this.state.getFailedFilepaths()
-        let error
-        if (failedSuites.length) {
-          error = new Error(`Test suites failed: ${failedSuites.length}.`)
-        }
-
-        sessionAsyncResource.runInAsyncScope(() => {
-          testSessionFinishCh.publish({
-            status: getSessionStatus(this.state),
-            onFinish,
-            error
-          })
-        })
+  // test finish (only passed tests)
+  shimmer.wrap(VitestTestRunner.prototype, 'onAfterTryTask', onAfterTryTask =>
+    async function (task, { retry: retryCount }) {
+      const result = await onAfterTryTask.apply(this, arguments)
 
-        await flushPromise
+      const status = getVitestTestStatus(task, retryCount)
+      const asyncResource = taskToAsync.get(task)
 
-        return exit.apply(this, arguments)
+      // We don't finish here because the test might fail in a later hook
+      asyncResource.runInAsyncScope(() => {
+        testFinishTimeCh.publish({ status, task })
       })
-
-      return sort.apply(this, arguments)
+      return result
     })
+
+  return vitestPackage
+})
+
+addHook({
+  name: 'vitest',
+  versions: ['>=1.6.0'],
+  file: 'dist/vendor/index.*'
+}, (vitestPackage) => {
+  // there are multiple index* files so we have to check the exported values
+  if (!isReporterPackage(vitestPackage)) {
+    return vitestPackage
   }
+  shimmer.wrap(vitestPackage.B.prototype, 'sort', sort => async function () {
+    shimmer.wrap(this.ctx, 'exit', exit => async function () {
+      let onFinish
 
-  if (isVitestTestRunner(vitestPackage)) {
-    const { VitestTestRunner } = vitestPackage
-    // test start (only tests that are not marked as skip or todo)
-    shimmer.wrap(VitestTestRunner.prototype, 'onBeforeTryTask', onBeforeTryTask => async function (task) {
-      const asyncResource = new AsyncResource('bound-anonymous-fn')
-      taskToAsync.set(task, asyncResource)
+      const flushPromise = new Promise(resolve => {
+        onFinish = resolve
+      })
+      const failedSuites = this.state.getFailedFilepaths()
+      let error
+      if (failedSuites.length) {
+        error = new Error(`Test suites failed: ${failedSuites.length}.`)
+      }
 
-      asyncResource.runInAsyncScope(() => {
-        testStartCh.publish({ testName: getTestName(task), testSuiteAbsolutePath: task.suite.file.filepath })
+      sessionAsyncResource.runInAsyncScope(() => {
+        testSessionFinishCh.publish({
+          status: getSessionStatus(this.state),
+          onFinish,
+          error
+        })
       })
-      return onBeforeTryTask.apply(this, arguments)
+
+      await flushPromise
+
+      return exit.apply(this, arguments)
     })
 
-    // test finish (only passed tests)
-    shimmer.wrap(VitestTestRunner.prototype, 'onAfterTryTask', onAfterTryTask =>
-      async function (task, { retry: retryCount }) {
-        const result = await onAfterTryTask.apply(this, arguments)
+    return sort.apply(this, arguments)
+  })
 
-        const status = getVitestTestStatus(task, retryCount)
-        const asyncResource = taskToAsync.get(task)
+  return vitestPackage
+})
+
+
+// Can't specify file because compiled vitest includes hashes in their files
+addHook({
+  name: 'vitest',
+  versions: ['>=1.6.0'],
+  file: 'dist/vendor/cac.*'
+}, (vitestPackage, frameworkVersion) => {
+  shimmer.wrap(vitestPackage, 'c', oldCreateCli => function () {
+    sessionAsyncResource.runInAsyncScope(() => {
+      const processArgv = process.argv.slice(2).join(' ')
+      testSessionStartCh.publish({ command: `vitest ${processArgv}`, frameworkVersion })
+    })
+    return oldCreateCli.apply(this, arguments)
+  })
 
-        // We don't finish here because the test might fail in a later hook
-        asyncResource.runInAsyncScope(() => {
-          testFinishTimeCh.publish({ status, task })
-        })
-        return result
-      })
-  }
   return vitestPackage
 })
 
@@ -170,85 +178,84 @@ addHook({
 // only relevant for workers
 addHook({
   name: '@vitest/runner',
-  versions: ['>=1.6.0']
+  versions: ['>=1.6.0'],
+  file: 'dist/index.js'
 }, vitestPackage => {
-  if (vitestPackage.startTests) {
-    shimmer.wrap(vitestPackage, 'startTests', startTests => async function (testPath) {
-      let testSuiteError = null
+  shimmer.wrap(vitestPackage, 'startTests', startTests => async function (testPath) {
+    let testSuiteError = null
 
-      const testSuiteAsyncResource = new AsyncResource('bound-anonymous-fn')
-      testSuiteAsyncResource.runInAsyncScope(() => {
-        testSuiteStartCh.publish(testPath[0])
-      })
-      const startTestsResponse = await startTests.apply(this, arguments)
+    const testSuiteAsyncResource = new AsyncResource('bound-anonymous-fn')
+    testSuiteAsyncResource.runInAsyncScope(() => {
+      testSuiteStartCh.publish(testPath[0])
+    })
+    const startTestsResponse = await startTests.apply(this, arguments)
 
-      let onFinish = null
-      const onFinishPromise = new Promise(resolve => {
-        onFinish = resolve
-      })
+    let onFinish = null
+    const onFinishPromise = new Promise(resolve => {
+      onFinish = resolve
+    })
 
-      const testTasks = getTypeTasks(startTestsResponse[0].tasks)
-
-      testTasks.forEach(task => {
-        const testAsyncResource = taskToAsync.get(task)
-        const { result } = task
-
-        if (result) {
-          const { state, duration, errors } = result
-          if (state === 'skip') { // programmatic skip
-            testSkipCh.publish({ testName: getTestName(task), testSuiteAbsolutePath: task.suite.file.filepath })
-          } else if (state === 'pass') {
-            testAsyncResource.runInAsyncScope(() => {
-              testPassCh.publish({ task })
-            })
-          } else if (state === 'fail') {
-            // If it's failing, we have no accurate finish time, so we have to use `duration`
-            let testError
-
-            if (errors?.length) {
-              testError = errors[0]
-            }
-
-            testAsyncResource.runInAsyncScope(() => {
-              testErrorCh.publish({ duration, error: testError })
-            })
-            if (errors?.length) {
-              testSuiteError = testError // we store the error to bubble it up to the suite
-            }
-          }
-        } else { // test.skip or test.todo
+    const testTasks = getTypeTasks(startTestsResponse[0].tasks)
+
+    testTasks.forEach(task => {
+      const testAsyncResource = taskToAsync.get(task)
+      const { result } = task
+
+      if (result) {
+        const { state, duration, errors } = result
+        if (state === 'skip') { // programmatic skip
           testSkipCh.publish({ testName: getTestName(task), testSuiteAbsolutePath: task.suite.file.filepath })
-        }
-      })
+        } else if (state === 'pass') {
+          testAsyncResource.runInAsyncScope(() => {
+            testPassCh.publish({ task })
+          })
+        } else if (state === 'fail') {
+          // If it's failing, we have no accurate finish time, so we have to use `duration`
+          let testError
 
-      const testSuiteResult = startTestsResponse[0].result
+          if (errors?.length) {
+            testError = errors[0]
+          }
 
-      if (testSuiteResult.errors?.length) { // Errors from root level hooks
-        testSuiteError = testSuiteResult.errors[0]
-      } else if (testSuiteResult.state === 'fail') { // Errors from `describe` level hooks
-        const suiteTasks = getTypeTasks(startTestsResponse[0].tasks, 'suite')
-        const failedSuites = suiteTasks.filter(task => task.result?.state === 'fail')
-        if (failedSuites.length && failedSuites[0].result?.errors?.length) {
-          testSuiteError = failedSuites[0].result.errors[0]
+          testAsyncResource.runInAsyncScope(() => {
+            testErrorCh.publish({ duration, error: testError })
+          })
+          if (errors?.length) {
+            testSuiteError = testError // we store the error to bubble it up to the suite
+          }
         }
+      } else { // test.skip or test.todo
+        testSkipCh.publish({ testName: getTestName(task), testSuiteAbsolutePath: task.suite.file.filepath })
       }
+    })
 
-      if (testSuiteError) {
-        testSuiteAsyncResource.runInAsyncScope(() => {
-          testSuiteErrorCh.publish({ error: testSuiteError })
-        })
+    const testSuiteResult = startTestsResponse[0].result
+
+    if (testSuiteResult.errors?.length) { // Errors from root level hooks
+      testSuiteError = testSuiteResult.errors[0]
+    } else if (testSuiteResult.state === 'fail') { // Errors from `describe` level hooks
+      const suiteTasks = getTypeTasks(startTestsResponse[0].tasks, 'suite')
+      const failedSuites = suiteTasks.filter(task => task.result?.state === 'fail')
+      if (failedSuites.length && failedSuites[0].result?.errors?.length) {
+        testSuiteError = failedSuites[0].result.errors[0]
       }
+    }
 
+    if (testSuiteError) {
       testSuiteAsyncResource.runInAsyncScope(() => {
-        testSuiteFinishCh.publish({ status: testSuiteResult.state, onFinish })
+        testSuiteErrorCh.publish({ error: testSuiteError })
       })
+    }
 
-      // TODO: fix too frequent flushes
-      await onFinishPromise
-
-      return startTestsResponse
+    testSuiteAsyncResource.runInAsyncScope(() => {
+      testSuiteFinishCh.publish({ status: testSuiteResult.state, onFinish })
     })
-  }
+
+    // TODO: fix too frequent flushes
+    await onFinishPromise
+
+    return startTestsResponse
+  })
 
   return vitestPackage
 })
diff --git a/yarn.lock b/yarn.lock
index 28a7abc232c..0c9111dd1b2 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -2824,9 +2824,10 @@ import-fresh@^3.0.0, import-fresh@^3.2.1:
     parent-module "^1.0.0"
     resolve-from "^4.0.0"
 
-"import-in-the-middle@https://github.com/DataDog/import-in-the-middle#dbaca128f6eb517d71bcf7e6b6b45f21a07128a4":
+import-in-the-middle@^1.8.1:
   version "1.8.1"
-  resolved "https://github.com/DataDog/import-in-the-middle#dbaca128f6eb517d71bcf7e6b6b45f21a07128a4"
+  resolved "https://registry.yarnpkg.com/import-in-the-middle/-/import-in-the-middle-1.8.1.tgz#8b51c2cc631b64e53e958d7048d2d9463ce628f8"
+  integrity sha512-yhRwoHtiLGvmSozNOALgjRPFI6uYsds60EoMqqnXyyv+JOIW/BrrLejuTGBt+bq0T5tLzOHrN0T7xYTm4Qt/ng==
   dependencies:
     acorn "^8.8.2"
     acorn-import-attributes "^1.9.5"

From aa20aae8deba514b84572a1d49624a39fb8167bc Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Tue, 25 Jun 2024 15:20:00 +0200
Subject: [PATCH 20/26] static list of esm plugins

---
 packages/datadog-instrumentations/src/helpers/hook.js  |  6 +++---
 .../datadog-instrumentations/src/helpers/register.js   | 10 +++++++++-
 packages/datadog-instrumentations/src/vitest.js        |  1 -
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/packages/datadog-instrumentations/src/helpers/hook.js b/packages/datadog-instrumentations/src/helpers/hook.js
index b56e2794ed0..0f9ebf4b22e 100644
--- a/packages/datadog-instrumentations/src/helpers/hook.js
+++ b/packages/datadog-instrumentations/src/helpers/hook.js
@@ -11,8 +11,8 @@ const ritm = require('../../../dd-trace/src/ritm')
  * @param {string[]} modules list of modules to hook into
  * @param {Function} onrequire callback to be executed upon encountering module
  */
-function Hook (modules, onrequire) {
-  if (!(this instanceof Hook)) return new Hook(modules, onrequire)
+function Hook (modules, hookOptions, onrequire) {
+  if (!(this instanceof Hook)) return new Hook(modules, hookOptions, onrequire)
 
   this._patched = Object.create(null)
 
@@ -28,7 +28,7 @@ function Hook (modules, onrequire) {
   }
 
   this._ritmHook = ritm(modules, {}, safeHook)
-  this._iitmHook = iitm(modules, { internals: true }, (moduleExports, moduleName, moduleBaseDir) => {
+  this._iitmHook = iitm(modules, hookOptions, (moduleExports, moduleName, moduleBaseDir) => {
     // TODO: Move this logic to import-in-the-middle and only do it for CommonJS
     // modules and not ESM. In the meantime, all the modules we instrument are
     // CommonJS modules for which the default export is always moved to
diff --git a/packages/datadog-instrumentations/src/helpers/register.js b/packages/datadog-instrumentations/src/helpers/register.js
index f8d1d2e0493..b6bbad55294 100644
--- a/packages/datadog-instrumentations/src/helpers/register.js
+++ b/packages/datadog-instrumentations/src/helpers/register.js
@@ -21,6 +21,8 @@ const disabledInstrumentations = new Set(
   DD_TRACE_DISABLED_INSTRUMENTATIONS ? DD_TRACE_DISABLED_INSTRUMENTATIONS.split(',') : []
 )
 
+const esmFirstLibraries = new Set(['vitest', '@vitest/runner'])
+
 const loadChannel = channel('dd-trace:instrumentation:load')
 
 // Globals
@@ -41,7 +43,13 @@ const isWildCard = (name) => name.includes('*')
 for (const packageName of names) {
   if (disabledInstrumentations.has(packageName)) continue
 
-  Hook([packageName], (moduleExports, moduleName, moduleBaseDir, moduleVersion) => {
+  const hookOptions = {}
+
+  if (esmFirstLibraries.has(packageName)) {
+    hookOptions.internals = true
+  }
+
+  Hook([packageName], hookOptions, (moduleExports, moduleName, moduleBaseDir, moduleVersion) => {
     moduleName = moduleName.replace(pathSepExpr, '/')
 
     // This executes the integration file thus adding its entries to `instrumentations`
diff --git a/packages/datadog-instrumentations/src/vitest.js b/packages/datadog-instrumentations/src/vitest.js
index ad8e8c6f247..b9870f31aa7 100644
--- a/packages/datadog-instrumentations/src/vitest.js
+++ b/packages/datadog-instrumentations/src/vitest.js
@@ -156,7 +156,6 @@ addHook({
   return vitestPackage
 })
 
-
 // Can't specify file because compiled vitest includes hashes in their files
 addHook({
   name: 'vitest',

From 5551e6334b837d120ea7d94819a57b45b93f2ae5 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Tue, 25 Jun 2024 16:46:30 +0200
Subject: [PATCH 21/26] add safeguard hook

---
 packages/datadog-instrumentations/src/helpers/hook.js | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/packages/datadog-instrumentations/src/helpers/hook.js b/packages/datadog-instrumentations/src/helpers/hook.js
index 0f9ebf4b22e..0177744ea1c 100644
--- a/packages/datadog-instrumentations/src/helpers/hook.js
+++ b/packages/datadog-instrumentations/src/helpers/hook.js
@@ -14,6 +14,11 @@ const ritm = require('../../../dd-trace/src/ritm')
 function Hook (modules, hookOptions, onrequire) {
   if (!(this instanceof Hook)) return new Hook(modules, hookOptions, onrequire)
 
+  if (typeof hookOptions === 'function') {
+    onrequire = hookOptions
+    hookOptions = {}
+  }
+
   this._patched = Object.create(null)
 
   const safeHook = (moduleExports, moduleName, moduleBaseDir, moduleVersion) => {

From c15e63c8722484209e39479d282c9b081d617275 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Tue, 25 Jun 2024 16:54:56 +0200
Subject: [PATCH 22/26] better regex logic

---
 .../src/helpers/instrument.js                  |  7 ++++---
 .../src/helpers/register.js                    | 18 +++++++++++++-----
 .../datadog-instrumentations/src/vitest.js     |  4 ++--
 3 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/packages/datadog-instrumentations/src/helpers/instrument.js b/packages/datadog-instrumentations/src/helpers/instrument.js
index 0889f1e5402..20657335044 100644
--- a/packages/datadog-instrumentations/src/helpers/instrument.js
+++ b/packages/datadog-instrumentations/src/helpers/instrument.js
@@ -17,10 +17,11 @@ exports.channel = function (name) {
 /**
  * @param {string} args.name module name
  * @param {string[]} args.versions array of semver range strings
- * @param {string} args.file path to file within package to instrument?
+ * @param {string} args.file path to file within package to instrument
+ * @param {string} args.filePattern pattern to match files within package to instrument
  * @param Function hook
  */
-exports.addHook = function addHook ({ name, versions, file }, hook) {
+exports.addHook = function addHook ({ name, versions, file, filePattern }, hook) {
   if (typeof name === 'string') {
     name = [name]
   }
@@ -29,7 +30,7 @@ exports.addHook = function addHook ({ name, versions, file }, hook) {
     if (!instrumentations[val]) {
       instrumentations[val] = []
     }
-    instrumentations[val].push({ name: val, versions, file, hook })
+    instrumentations[val].push({ name: val, versions, file, filePattern, hook })
   }
 }
 
diff --git a/packages/datadog-instrumentations/src/helpers/register.js b/packages/datadog-instrumentations/src/helpers/register.js
index b6bbad55294..66757dfb2cc 100644
--- a/packages/datadog-instrumentations/src/helpers/register.js
+++ b/packages/datadog-instrumentations/src/helpers/register.js
@@ -59,8 +59,12 @@ for (const packageName of names) {
       return moduleExports
     }
 
-    for (const { name, file, versions, hook } of instrumentations[packageName]) {
+    for (const { name, file, versions, hook, filePattern } of instrumentations[packageName]) {
+      let fullFilePattern = filePattern
       const fullFilename = filename(name, file)
+      if (fullFilePattern) {
+        fullFilePattern = filename(name, fullFilePattern)
+      }
 
       // Create a WeakMap associated with the hook function so that patches on the same moduleExport only happens once
       // for example by instrumenting both dns and node:dns double the spans would be created
@@ -68,11 +72,15 @@ for (const packageName of names) {
       if (!hook[HOOK_SYMBOL]) {
         hook[HOOK_SYMBOL] = new WeakMap()
       }
+      let matchesFile = false
+
+      matchesFile = moduleName === fullFilename
 
-      // Some libraries include a hash in their filenames when installed,
-      // so our instrumentation has to include a * to match them for more than a single version.
-      const matchesFile = moduleName === fullFilename ||
-        (isWildCard(fullFilename) && moduleName.includes(fullFilename.replace('*', '')))
+      if (fullFilePattern) {
+        // Some libraries include a hash in their filenames when installed,
+        // so our instrumentation has to include a '.*' to match them for more than a single version.
+        matchesFile = matchesFile || new RegExp(fullFilePattern).test(moduleName)
+      }
 
       if (matchesFile) {
         const version = moduleVersion || getVersion(moduleBaseDir)
diff --git a/packages/datadog-instrumentations/src/vitest.js b/packages/datadog-instrumentations/src/vitest.js
index b9870f31aa7..39545808c72 100644
--- a/packages/datadog-instrumentations/src/vitest.js
+++ b/packages/datadog-instrumentations/src/vitest.js
@@ -118,7 +118,7 @@ addHook({
 addHook({
   name: 'vitest',
   versions: ['>=1.6.0'],
-  file: 'dist/vendor/index.*'
+  filePattern: 'dist/vendor/index.*'
 }, (vitestPackage) => {
   // there are multiple index* files so we have to check the exported values
   if (!isReporterPackage(vitestPackage)) {
@@ -160,7 +160,7 @@ addHook({
 addHook({
   name: 'vitest',
   versions: ['>=1.6.0'],
-  file: 'dist/vendor/cac.*'
+  filePattern: 'dist/vendor/cac.*'
 }, (vitestPackage, frameworkVersion) => {
   shimmer.wrap(vitestPackage, 'c', oldCreateCli => function () {
     sessionAsyncResource.runInAsyncScope(() => {

From ee2b032716ada75cd5060c9e98c331e3257e0f31 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Tue, 25 Jun 2024 16:58:53 +0200
Subject: [PATCH 23/26] remove unused var

---
 packages/datadog-instrumentations/src/helpers/register.js | 2 --
 1 file changed, 2 deletions(-)

diff --git a/packages/datadog-instrumentations/src/helpers/register.js b/packages/datadog-instrumentations/src/helpers/register.js
index 66757dfb2cc..f3b7f8d8973 100644
--- a/packages/datadog-instrumentations/src/helpers/register.js
+++ b/packages/datadog-instrumentations/src/helpers/register.js
@@ -37,8 +37,6 @@ if (DD_TRACE_DEBUG && DD_TRACE_DEBUG.toLowerCase() !== 'false') {
   setImmediate(checkRequireCache.checkForPotentialConflicts)
 }
 
-const isWildCard = (name) => name.includes('*')
-
 // TODO: make this more efficient
 for (const packageName of names) {
   if (disabledInstrumentations.has(packageName)) continue

From 37f33e7ed025df0168e42889e6f6e0dc17b66998 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Fri, 28 Jun 2024 10:55:43 +0200
Subject: [PATCH 24/26] move esm first logic to hooks

---
 packages/datadog-esbuild/index.js                     |  6 +++++-
 .../datadog-instrumentations/src/helpers/hooks.js     |  4 ++--
 .../datadog-instrumentations/src/helpers/register.js  | 11 ++++++-----
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/packages/datadog-esbuild/index.js b/packages/datadog-esbuild/index.js
index 95a0e8ddd16..ce263799023 100644
--- a/packages/datadog-esbuild/index.js
+++ b/packages/datadog-esbuild/index.js
@@ -7,7 +7,11 @@ const hooks = require('../datadog-instrumentations/src/helpers/hooks.js')
 const extractPackageAndModulePath = require('../datadog-instrumentations/src/utils/src/extract-package-and-module-path')
 
 for (const hook of Object.values(hooks)) {
-  hook()
+  if (typeof hook === 'object') {
+    hook.fn()
+  } else {
+    hook()
+  }
 }
 
 const modulesOfInterest = new Set()
diff --git a/packages/datadog-instrumentations/src/helpers/hooks.js b/packages/datadog-instrumentations/src/helpers/hooks.js
index c06e1101ecd..94f3318fb62 100644
--- a/packages/datadog-instrumentations/src/helpers/hooks.js
+++ b/packages/datadog-instrumentations/src/helpers/hooks.js
@@ -23,7 +23,7 @@ module.exports = {
   '@opentelemetry/sdk-trace-node': () => require('../otel-sdk-trace'),
   '@redis/client': () => require('../redis'),
   '@smithy/smithy-client': () => require('../aws-sdk'),
-  '@vitest/runner': () => require('../vitest'),
+  '@vitest/runner': { esmFirst: true, fn: () => require('../vitest') },
   aerospike: () => require('../aerospike'),
   amqp10: () => require('../amqp10'),
   amqplib: () => require('../amqplib'),
@@ -111,7 +111,7 @@ module.exports = {
   sharedb: () => require('../sharedb'),
   tedious: () => require('../tedious'),
   undici: () => require('../undici'),
-  vitest: () => require('../vitest'),
+  vitest: { esmFirst: true, fn: () => require('../vitest') },
   when: () => require('../when'),
   winston: () => require('../winston')
 }
diff --git a/packages/datadog-instrumentations/src/helpers/register.js b/packages/datadog-instrumentations/src/helpers/register.js
index 7e70454d2ad..e45a0c0cd14 100644
--- a/packages/datadog-instrumentations/src/helpers/register.js
+++ b/packages/datadog-instrumentations/src/helpers/register.js
@@ -22,8 +22,6 @@ const disabledInstrumentations = new Set(
   DD_TRACE_DISABLED_INSTRUMENTATIONS ? DD_TRACE_DISABLED_INSTRUMENTATIONS.split(',') : []
 )
 
-const esmFirstLibraries = new Set(['vitest', '@vitest/runner'])
-
 const loadChannel = channel('dd-trace:instrumentation:load')
 
 // Globals
@@ -46,15 +44,18 @@ for (const packageName of names) {
 
   const hookOptions = {}
 
-  if (esmFirstLibraries.has(packageName)) {
-    hookOptions.internals = true
+  let hook = hooks[packageName]
+
+  if (typeof hook === 'object') {
+    hookOptions.internals = hook.esmFirst
+    hook = hook.fn
   }
 
   Hook([packageName], hookOptions, (moduleExports, moduleName, moduleBaseDir, moduleVersion) => {
     moduleName = moduleName.replace(pathSepExpr, '/')
 
     // This executes the integration file thus adding its entries to `instrumentations`
-    hooks[packageName]()
+    hook()
 
     if (!instrumentations[packageName]) {
       return moduleExports

From b29cfdb18b753d98e0b083798f20f42f815560b8 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Fri, 28 Jun 2024 11:53:45 +0200
Subject: [PATCH 25/26] remove dependency on getport

---
 integration-tests/vitest/vitest.spec.js | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/integration-tests/vitest/vitest.spec.js b/integration-tests/vitest/vitest.spec.js
index 3065e52b202..4a3151c2b72 100644
--- a/integration-tests/vitest/vitest.spec.js
+++ b/integration-tests/vitest/vitest.spec.js
@@ -2,7 +2,6 @@
 
 const { exec } = require('child_process')
 
-const getPort = require('get-port')
 const { assert } = require('chai')
 
 const {
@@ -32,8 +31,7 @@ versions.forEach((version) => {
     })
 
     beforeEach(async function () {
-      const port = await getPort()
-      receiver = await new FakeCiVisIntake(port).start()
+      receiver = await new FakeCiVisIntake().start()
     })
 
     afterEach(async () => {

From 6bee6563e5ce854a3256f6826e2af26a4261b678 Mon Sep 17 00:00:00 2001
From: Juan Fernandez <juan.fernandezdealba@datadoghq.com>
Date: Fri, 28 Jun 2024 12:30:54 +0200
Subject: [PATCH 26/26] adding checks for hasSubscriber and asyncresource

---
 .../datadog-instrumentations/src/vitest.js    | 42 ++++++++++++++-----
 1 file changed, 32 insertions(+), 10 deletions(-)

diff --git a/packages/datadog-instrumentations/src/vitest.js b/packages/datadog-instrumentations/src/vitest.js
index 39545808c72..42f32b1ac42 100644
--- a/packages/datadog-instrumentations/src/vitest.js
+++ b/packages/datadog-instrumentations/src/vitest.js
@@ -88,6 +88,9 @@ addHook({
   const { VitestTestRunner } = vitestPackage
   // test start (only tests that are not marked as skip or todo)
   shimmer.wrap(VitestTestRunner.prototype, 'onBeforeTryTask', onBeforeTryTask => async function (task) {
+    if (!testStartCh.hasSubscribers) {
+      return onBeforeTryTask.apply(this, arguments)
+    }
     const asyncResource = new AsyncResource('bound-anonymous-fn')
     taskToAsync.set(task, asyncResource)
 
@@ -100,15 +103,21 @@ addHook({
   // test finish (only passed tests)
   shimmer.wrap(VitestTestRunner.prototype, 'onAfterTryTask', onAfterTryTask =>
     async function (task, { retry: retryCount }) {
+      if (!testFinishTimeCh.hasSubscribers) {
+        return onAfterTryTask.apply(this, arguments)
+      }
       const result = await onAfterTryTask.apply(this, arguments)
 
       const status = getVitestTestStatus(task, retryCount)
       const asyncResource = taskToAsync.get(task)
 
-      // We don't finish here because the test might fail in a later hook
-      asyncResource.runInAsyncScope(() => {
-        testFinishTimeCh.publish({ status, task })
-      })
+      if (asyncResource) {
+        // We don't finish here because the test might fail in a later hook
+        asyncResource.runInAsyncScope(() => {
+          testFinishTimeCh.publish({ status, task })
+        })
+      }
+
       return result
     })
 
@@ -125,6 +134,9 @@ addHook({
     return vitestPackage
   }
   shimmer.wrap(vitestPackage.B.prototype, 'sort', sort => async function () {
+    if (!testSessionFinishCh.hasSubscribers) {
+      return sort.apply(this, arguments)
+    }
     shimmer.wrap(this.ctx, 'exit', exit => async function () {
       let onFinish
 
@@ -163,6 +175,9 @@ addHook({
   filePattern: 'dist/vendor/cac.*'
 }, (vitestPackage, frameworkVersion) => {
   shimmer.wrap(vitestPackage, 'c', oldCreateCli => function () {
+    if (!testSessionStartCh.hasSubscribers) {
+      return oldCreateCli.apply(this, arguments)
+    }
     sessionAsyncResource.runInAsyncScope(() => {
       const processArgv = process.argv.slice(2).join(' ')
       testSessionStartCh.publish({ command: `vitest ${processArgv}`, frameworkVersion })
@@ -182,6 +197,9 @@ addHook({
 }, vitestPackage => {
   shimmer.wrap(vitestPackage, 'startTests', startTests => async function (testPath) {
     let testSuiteError = null
+    if (!testSuiteStartCh.hasSubscribers) {
+      return startTests.apply(this, arguments)
+    }
 
     const testSuiteAsyncResource = new AsyncResource('bound-anonymous-fn')
     testSuiteAsyncResource.runInAsyncScope(() => {
@@ -205,9 +223,11 @@ addHook({
         if (state === 'skip') { // programmatic skip
           testSkipCh.publish({ testName: getTestName(task), testSuiteAbsolutePath: task.suite.file.filepath })
         } else if (state === 'pass') {
-          testAsyncResource.runInAsyncScope(() => {
-            testPassCh.publish({ task })
-          })
+          if (testAsyncResource) {
+            testAsyncResource.runInAsyncScope(() => {
+              testPassCh.publish({ task })
+            })
+          }
         } else if (state === 'fail') {
           // If it's failing, we have no accurate finish time, so we have to use `duration`
           let testError
@@ -216,9 +236,11 @@ addHook({
             testError = errors[0]
           }
 
-          testAsyncResource.runInAsyncScope(() => {
-            testErrorCh.publish({ duration, error: testError })
-          })
+          if (testAsyncResource) {
+            testAsyncResource.runInAsyncScope(() => {
+              testErrorCh.publish({ duration, error: testError })
+            })
+          }
           if (errors?.length) {
             testSuiteError = testError // we store the error to bubble it up to the suite
           }