dd-trace 5.34.0 → 5.36.0

package/packages/datadog-plugin-cypress/src/support.js

@@ -4,6 +4,10 @@ let isKnownTestsEnabled = false
 let knownTestsForSuite = []
 let suiteTests = []
 let earlyFlakeDetectionNumRetries = 0
+// We need to grab the original window as soon as possible,
+// in case the test changes the origin. If the test does change the origin,
+// any call to `cy.window()` will result in a cross origin error.
+let originalWindow
 
 // If the test is using multi domain with cy.origin, trying to access
 // window properties will result in a cross origin error.
@@ -61,6 +65,9 @@ beforeEach(function () {
       this.skip()
     }
   })
+  cy.window().then(win => {
+    originalWindow = win
+  })
 })
 
 before(function () {
@@ -78,39 +85,39 @@ before(function () {
 })
 
 after(() => {
-  cy.window().then(win => {
-    if (safeGetRum(win)) {
-      win.dispatchEvent(new Event('beforeunload'))
+  try {
+    if (safeGetRum(originalWindow)) {
+      originalWindow.dispatchEvent(new Event('beforeunload'))
     }
-  })
+  } catch (e) {
+    // ignore error. It's usually a multi origin issue.
+  }
 })
 
 
 afterEach(function () {
-  cy.window().then(win => {
-    const currentTest = Cypress.mocha.getRunner().suite.ctx.currentTest
-    const testInfo = {
-      testName: currentTest.fullTitle(),
-      testSuite: Cypress.mocha.getRootSuite().file,
-      testSuiteAbsolutePath: Cypress.spec && Cypress.spec.absolute,
-      state: currentTest.state,
-      error: currentTest.err,
-      isNew: currentTest._ddIsNew,
-      isEfdRetry: currentTest._ddIsEfdRetry
-    }
-    try {
-      testInfo.testSourceLine = Cypress.mocha.getRunner().currentRunnable.invocationDetails.line
-    } catch (e) {}
+  const currentTest = Cypress.mocha.getRunner().suite.ctx.currentTest
+  const testInfo = {
+    testName: currentTest.fullTitle(),
+    testSuite: Cypress.mocha.getRootSuite().file,
+    testSuiteAbsolutePath: Cypress.spec && Cypress.spec.absolute,
+    state: currentTest.state,
+    error: currentTest.err,
+    isNew: currentTest._ddIsNew,
+    isEfdRetry: currentTest._ddIsEfdRetry
+  }
+  try {
+    testInfo.testSourceLine = Cypress.mocha.getRunner().currentRunnable.invocationDetails.line
+  } catch (e) {}
 
-    if (safeGetRum(win)) {
-      testInfo.isRUMActive = true
-    }
-    let coverage
-    try {
-      coverage = win.__coverage__
-    } catch (e) {
-      // ignore error and continue
-    }
-    cy.task('dd:afterEach', { test: testInfo, coverage })
-  })
+  if (safeGetRum(originalWindow)) {
+    testInfo.isRUMActive = true
+  }
+  let coverage
+  try {
+    coverage = originalWindow.__coverage__
+  } catch (e) {
+    // ignore error and continue
+  }
+  cy.task('dd:afterEach', { test: testInfo, coverage })
 })

package/packages/datadog-plugin-dd-trace-api/src/index.js

@@ -0,0 +1,120 @@
+'use strict'
+
+const Plugin = require('../../dd-trace/src/plugins/plugin')
+const telemetryMetrics = require('../../dd-trace/src/telemetry/metrics')
+const apiMetrics = telemetryMetrics.manager.namespace('tracers')
+
+// api ==> here
+const objectMap = new WeakMap()
+
+const injectionEnabledTag =
+  `injection_enabled:${process.env.DD_INJECTION_ENABLED ? 'yes' : 'no'}`
+
+module.exports = class DdTraceApiPlugin extends Plugin {
+  static get id () {
+    return 'dd-trace-api'
+  }
+
+  constructor (...args) {
+    super(...args)
+
+    const tracer = this._tracer
+
+    this.addSub('datadog-api:v1:tracerinit', ({ proxy }) => {
+      const proxyVal = proxy()
+      objectMap.set(proxyVal, tracer)
+      objectMap.set(proxyVal.appsec, tracer.appsec)
+      objectMap.set(proxyVal.dogstatsd, tracer.dogstatsd)
+    })
+
+    const handleEvent = (name) => {
+      const counter = apiMetrics.count('dd_trace_api.called', [
+        `name:${name.replaceAll(':', '.')}`,
+        'api_version:v1',
+        injectionEnabledTag
+      ])
+
+      // For v1, APIs are 1:1 with their internal equivalents, so we can just
+      // call the internal method directly. That's what we do here unless we
+      // want to override. As the API evolves, this may change.
+      this.addSub(`datadog-api:v1:${name}`, ({ self, args, ret, proxy, revProxy }) => {
+        counter.inc()
+
+        if (name.includes(':')) {
+          name = name.split(':').pop()
+        }
+
+        if (objectMap.has(self)) {
+          self = objectMap.get(self)
+        }
+
+        for (let i = 0; i < args.length; i++) {
+          if (objectMap.has(args[i])) {
+            args[i] = objectMap.get(args[i])
+          }
+          if (typeof args[i] === 'function') {
+            const orig = args[i]
+            args[i] = (...fnArgs) => {
+              for (let j = 0; j < fnArgs.length; j++) {
+                if (revProxy && revProxy[j]) {
+                  const proxyVal = revProxy[j]()
+                  objectMap.set(proxyVal, fnArgs[j])
+                  fnArgs[j] = proxyVal
+                }
+              }
+              // TODO do we need to apply(this, ...) here?
+              return orig(...fnArgs)
+            }
+          }
+        }
+
+        try {
+          ret.value = self[name](...args)
+          if (proxy) {
+            const proxyVal = proxy()
+            objectMap.set(proxyVal, ret.value)
+            ret.value = proxyVal
+          } else if (ret.value && typeof ret.value === 'object') {
+            throw new TypeError(`Objects need proxies when returned via API (${name})`)
+          }
+        } catch (e) {
+          ret.error = e
+        }
+      })
+    }
+
+    // handleEvent('configure')
+    handleEvent('startSpan')
+    handleEvent('wrap')
+    handleEvent('trace')
+    handleEvent('inject')
+    handleEvent('extract')
+    handleEvent('getRumData')
+    handleEvent('profilerStarted')
+    handleEvent('context:toTraceId')
+    handleEvent('context:toSpanId')
+    handleEvent('context:toTraceparent')
+    handleEvent('span:context')
+    handleEvent('span:setTag')
+    handleEvent('span:addTags')
+    handleEvent('span:finish')
+    handleEvent('span:addLink')
+    handleEvent('scope')
+    handleEvent('scope:activate')
+    handleEvent('scope:active')
+    handleEvent('scope:bind')
+    handleEvent('appsec:blockRequest')
+    handleEvent('appsec:isUserBlocked')
+    handleEvent('appsec:setUser')
+    handleEvent('appsec:trackCustomEvent')
+    handleEvent('appsec:trackUserLoginFailureEvent')
+    handleEvent('appsec:trackUserLoginSuccessEvent')
+    handleEvent('dogstatsd:decrement')
+    handleEvent('dogstatsd:distribution')
+    handleEvent('dogstatsd:flush')
+    handleEvent('dogstatsd:gauge')
+    handleEvent('dogstatsd:histogram')
+    handleEvent('dogstatsd:increment')
+    handleEvent('use')
+  }
+}

package/packages/datadog-plugin-grpc/src/client.js

@@ -20,7 +20,7 @@ class GrpcClientPlugin extends ClientPlugin {
   }
 
   bindStart (message) {
-    const store = storage.getStore()
+    const store = storage('legacy').getStore()
     const { metadata, path, type } = message
     const metadataFilter = this.config.metadataFilter
     const method = getMethodMetadata(path, type)

package/packages/datadog-plugin-grpc/src/server.js

@@ -27,7 +27,7 @@ class GrpcServerPlugin extends ServerPlugin {
   }
 
   bindStart (message) {
-    const store = storage.getStore()
+    const store = storage('legacy').getStore()
     const { name, metadata, type } = message
     const metadataFilter = this.config.metadataFilter
     const childOf = extract(this.tracer, metadata)

package/packages/datadog-plugin-hapi/src/index.js

@@ -15,7 +15,7 @@ class HapiPlugin extends RouterPlugin {
     this._requestSpans = new WeakMap()
 
     this.addSub('apm:hapi:request:handle', ({ req }) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       const span = store && store.span
 
       this.setFramework(req, 'hapi', this.config)

package/packages/datadog-plugin-http/src/client.js

@@ -21,7 +21,7 @@ class HttpClientPlugin extends ClientPlugin {
 
   bindStart (message) {
     const { args, http = {} } = message
-    const store = storage.getStore()
+    const store = storage('legacy').getStore()
     const options = args.options
     const agent = options.agent || options._defaultAgent || http.globalAgent || {}
     const protocol = options.protocol || agent.protocol || 'http:'

package/packages/datadog-plugin-http/src/server.js

@@ -22,7 +22,7 @@ class HttpServerPlugin extends ServerPlugin {
   }
 
   start ({ req, res, abortController }) {
-    const store = storage.getStore()
+    const store = storage('legacy').getStore()
     const span = web.startSpan(
       this.tracer,
       {

package/packages/datadog-plugin-http2/src/client.js

@@ -36,7 +36,7 @@ class Http2ClientPlugin extends ClientPlugin {
     const uri = `${sessionDetails.protocol}//${sessionDetails.host}:${sessionDetails.port}${pathname}`
     const allowed = this.config.filter(uri)
 
-    const store = storage.getStore()
+    const store = storage('legacy').getStore()
     const childOf = store && allowed ? store.span : null
     const span = this.startSpan(this.operationName(), {
       childOf,
@@ -85,7 +85,7 @@ class Http2ClientPlugin extends ClientPlugin {
         return parentStore
     }
 
-    return storage.getStore()
+    return storage('legacy').getStore()
   }
 
   configure (config) {
@@ -98,7 +98,7 @@ class Http2ClientPlugin extends ClientPlugin {
     store.span.setTag(HTTP_STATUS_CODE, status)
 
     if (!this.config.validateStatus(status)) {
-      storage.run(store, () => this.addError())
+      storage('legacy').run(store, () => this.addError())
     }
 
     addHeaderTags(store.span, headers, HTTP_RESPONSE_HEADERS, this.config)

package/packages/datadog-plugin-http2/src/server.js

@@ -17,7 +17,7 @@ class Http2ServerPlugin extends ServerPlugin {
   }
 
   start ({ req, res }) {
-    const store = storage.getStore()
+    const store = storage('legacy').getStore()
     const span = web.startSpan(
       this.tracer,
       {

package/packages/datadog-plugin-jest/src/index.js

@@ -291,6 +291,7 @@ class JestPlugin extends CiPlugin {
       if (isJestWorker) {
         this.tracer._exporter.flush()
       }
+      this.removeAllDiProbes()
     })
 
     /**
@@ -317,15 +318,15 @@ class JestPlugin extends CiPlugin {
     })
 
     this.addSub('ci:jest:test:start', (test) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       const span = this.startTestSpan(test)
 
       this.enter(span, store)
       this.activeTestSpan = span
     })
 
-    this.addSub('ci:jest:test:finish', ({ status, testStartLine, promises, shouldRemoveProbe }) => {
-      const span = storage.getStore().span
+    this.addSub('ci:jest:test:finish', ({ status, testStartLine }) => {
+      const span = storage('legacy').getStore().span
       span.setTag(TEST_STATUS, status)
       if (testStartLine) {
         span.setTag(TEST_SOURCE_START, testStartLine)
@@ -346,15 +347,11 @@ class JestPlugin extends CiPlugin {
       span.finish()
       finishAllTraceSpans(span)
       this.activeTestSpan = null
-      if (shouldRemoveProbe && this.runningTestProbeId) {
-        promises.isProbeRemoved = withTimeout(this.removeDiProbe(this.runningTestProbeId), 2000)
-        this.runningTestProbeId = null
-      }
     })
 
     this.addSub('ci:jest:test:err', ({ error, shouldSetProbe, promises }) => {
       if (error) {
-        const store = storage.getStore()
+        const store = storage('legacy').getStore()
         if (store && store.span) {
           const span = store.span
           span.setTag(TEST_STATUS, 'fail')
@@ -362,9 +359,7 @@ class JestPlugin extends CiPlugin {
           if (shouldSetProbe) {
             const probeInformation = this.addDiProbe(error)
             if (probeInformation) {
-              const { probeId, setProbePromise, stackIndex } = probeInformation
-              this.runningTestProbeId = probeId
-              this.testErrorStackIndex = stackIndex
+              const { setProbePromise } = probeInformation
               promises.isProbeReady = withTimeout(setProbePromise, 2000)
             }
           }

package/packages/datadog-plugin-langchain/src/tracing.js

@@ -61,7 +61,7 @@ class LangChainTracingPlugin extends TracingPlugin {
       }
     }, false)
 
-    const store = storage.getStore() || {}
+    const store = storage('legacy').getStore() || {}
     ctx.currentStore = { ...store, span }
 
     return ctx.currentStore

package/packages/datadog-plugin-mariadb/src/index.js

@@ -13,12 +13,12 @@ class MariadbPlugin extends MySQLPlugin {
     super(...args)
 
     this.addSub(`apm:${this.component}:pool:skip`, () => {
-      skippedStore = storage.getStore()
-      storage.enterWith({ noop: true })
+      skippedStore = storage('legacy').getStore()
+      storage('legacy').enterWith({ noop: true })
     })
 
     this.addSub(`apm:${this.component}:pool:unskip`, () => {
-      storage.enterWith(skippedStore)
+      storage('legacy').enterWith(skippedStore)
       skippedStore = undefined
     })
   }

package/packages/datadog-plugin-mocha/src/index.js

@@ -155,13 +155,13 @@ class MochaPlugin extends CiPlugin {
       if (itrCorrelationId) {
         testSuiteSpan.setTag(ITR_CORRELATION_ID, itrCorrelationId)
       }
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       this.enter(testSuiteSpan, store)
       this._testSuites.set(testSuite, testSuiteSpan)
     })
 
     this.addSub('ci:mocha:test-suite:finish', (status) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       if (store && store.span) {
         const span = store.span
         // the test status of the suite may have been set in ci:mocha:test-suite:error already
@@ -174,7 +174,7 @@ class MochaPlugin extends CiPlugin {
     })
 
     this.addSub('ci:mocha:test-suite:error', (err) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       if (store && store.span) {
         const span = store.span
         span.setTag('error', err)
@@ -183,7 +183,7 @@ class MochaPlugin extends CiPlugin {
     })
 
     this.addSub('ci:mocha:test:start', (testInfo) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       const span = this.startTestSpan(testInfo)
 
       this.enter(span, store)
@@ -195,7 +195,7 @@ class MochaPlugin extends CiPlugin {
     })
 
     this.addSub('ci:mocha:test:finish', ({ status, hasBeenRetried, isLastRetry }) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       const span = store?.span
 
       if (span) {
@@ -219,15 +219,15 @@ class MochaPlugin extends CiPlugin {
         span.finish()
         finishAllTraceSpans(span)
         this.activeTestSpan = null
-        if (this.di && this.libraryConfig?.isDiEnabled && this.runningTestProbeId && isLastRetry) {
-          this.removeDiProbe(this.runningTestProbeId)
-          this.runningTestProbeId = null
+        if (this.di && this.libraryConfig?.isDiEnabled && this.runningTestProbe && isLastRetry) {
+          this.removeDiProbe(this.runningTestProbe)
+          this.runningTestProbe = null
         }
       }
     })
 
     this.addSub('ci:mocha:test:skip', (testInfo) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       // skipped through it.skip, so the span is not created yet
       // for this test
       if (!store) {
@@ -237,7 +237,7 @@ class MochaPlugin extends CiPlugin {
     })
 
     this.addSub('ci:mocha:test:error', (err) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       const span = store?.span
       if (err && span) {
         if (err.constructor.name === 'Pending' && !this.forbidPending) {
@@ -250,7 +250,7 @@ class MochaPlugin extends CiPlugin {
     })
 
     this.addSub('ci:mocha:test:retry', ({ isFirstAttempt, willBeRetried, err, test }) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       const span = store?.span
       if (span) {
         span.setTag(TEST_STATUS, 'fail')
@@ -275,8 +275,8 @@ class MochaPlugin extends CiPlugin {
         if (isFirstAttempt && willBeRetried && this.di && this.libraryConfig?.isDiEnabled) {
           const probeInformation = this.addDiProbe(err)
           if (probeInformation) {
-            const { probeId, stackIndex } = probeInformation
-            this.runningTestProbeId = probeId
+            const { file, line, stackIndex } = probeInformation
+            this.runningTestProbe = { file, line }
             this.testErrorStackIndex = stackIndex
             test._ddShouldWaitForHitProbe = true
             // TODO: we're not waiting for setProbePromise to be resolved, so there might be race conditions

package/packages/datadog-plugin-next/src/index.js

@@ -20,7 +20,7 @@ class NextPlugin extends ServerPlugin {
   }
 
   bindStart ({ req, res }) {
-    const store = storage.getStore()
+    const store = storage('legacy').getStore()
     const childOf = store ? store.span : store
     const span = this.tracer.startSpan(this.operationName(), {
       childOf,
@@ -43,7 +43,7 @@ class NextPlugin extends ServerPlugin {
 
   error ({ span, error }) {
     if (!span) {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       if (!store) return
 
       span = store.span
@@ -53,7 +53,7 @@ class NextPlugin extends ServerPlugin {
   }
 
   finish ({ req, res, nextRequest = {} }) {
-    const store = storage.getStore()
+    const store = storage('legacy').getStore()
 
     if (!store) return
 
@@ -85,7 +85,7 @@ class NextPlugin extends ServerPlugin {
   }
 
   pageLoad ({ page, isAppPath = false, isStatic = false }) {
-    const store = storage.getStore()
+    const store = storage('legacy').getStore()
 
     if (!store) return
 

package/packages/datadog-plugin-openai/src/tracing.js

@@ -60,7 +60,7 @@ class OpenAiTracingPlugin extends TracingPlugin {
   bindStart (ctx) {
     const { methodName, args, basePath, apiKey } = ctx
     const payload = normalizeRequestPayload(methodName, args)
-    const store = storage.getStore() || {}
+    const store = storage('legacy').getStore() || {}
 
     const span = this.startSpan('openai.request', {
       service: this.config.service,

package/packages/datadog-plugin-playwright/src/index.js

@@ -68,7 +68,7 @@ class PlaywrightPlugin extends CiPlugin {
     })
 
     this.addSub('ci:playwright:test-suite:start', (testSuiteAbsolutePath) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.rootDir)
       const testSourceFile = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
 
@@ -102,7 +102,7 @@ class PlaywrightPlugin extends CiPlugin {
     })
 
     this.addSub('ci:playwright:test-suite:finish', ({ status, error }) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       const span = store && store.span
       if (!span) return
       if (error) {
@@ -121,7 +121,7 @@ class PlaywrightPlugin extends CiPlugin {
     })
 
     this.addSub('ci:playwright:test:start', ({ testName, testSuiteAbsolutePath, testSourceLine, browserName }) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.rootDir)
       const testSourceFile = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
       const span = this.startTestSpan(testName, testSuite, testSourceFile, testSourceLine, browserName)
@@ -129,7 +129,7 @@ class PlaywrightPlugin extends CiPlugin {
       this.enter(span, store)
     })
     this.addSub('ci:playwright:test:finish', ({ testStatus, steps, error, extraTags, isNew, isEfdRetry, isRetry }) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       const span = store && store.span
       if (!span) return
 

package/packages/datadog-plugin-rhea/src/consumer.js

@@ -11,7 +11,7 @@ class RheaConsumerPlugin extends ConsumerPlugin {
     super(...args)
 
     this.addTraceSub('dispatch', ({ state }) => {
-      const span = storage.getStore().span
+      const span = storage('legacy').getStore().span
       span.setTag('amqp.delivery.state', state)
     })
   }

package/packages/datadog-plugin-router/src/index.js

@@ -29,7 +29,7 @@ class RouterPlugin extends WebPlugin {
         context.middleware.push(span)
       }
 
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       this._storeStack.push(store)
       this.enter(span, store)
 
@@ -94,7 +94,7 @@ class RouterPlugin extends WebPlugin {
   }
 
   _getStoreSpan () {
-    const store = storage.getStore()
+    const store = storage('legacy').getStore()
 
     return store && store.span
   }

package/packages/datadog-plugin-selenium/src/index.js

@@ -39,7 +39,7 @@ class SeleniumPlugin extends CiPlugin {
       browserVersion,
       isRumActive
     }) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       const span = store?.span
       if (!span) {
         return

package/packages/datadog-plugin-vitest/src/index.js

@@ -70,7 +70,7 @@ class VitestPlugin extends CiPlugin {
       isRetryReasonEfd
     }) => {
       const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
 
       const extraTags = {
         [TEST_SOURCE_FILE]: testSuite
@@ -102,7 +102,7 @@ class VitestPlugin extends CiPlugin {
     })
 
     this.addSub('ci:vitest:test:finish-time', ({ status, task }) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       const span = store?.span
 
       // we store the finish time to finish at a later hook
@@ -114,7 +114,7 @@ class VitestPlugin extends CiPlugin {
     })
 
     this.addSub('ci:vitest:test:pass', ({ task }) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       const span = store?.span
 
       if (span) {
@@ -128,15 +128,15 @@ class VitestPlugin extends CiPlugin {
     })
 
     this.addSub('ci:vitest:test:error', ({ duration, error, shouldSetProbe, promises }) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       const span = store?.span
 
       if (span) {
         if (shouldSetProbe && this.di) {
           const probeInformation = this.addDiProbe(error)
           if (probeInformation) {
-            const { probeId, stackIndex, setProbePromise } = probeInformation
-            this.runningTestProbeId = probeId
+            const { file, line, stackIndex, setProbePromise } = probeInformation
+            this.runningTestProbe = { file, line }
             this.testErrorStackIndex = stackIndex
             promises.setProbePromise = setProbePromise
           }
@@ -221,13 +221,13 @@ class VitestPlugin extends CiPlugin {
         }
       })
       this.telemetry.ciVisEvent(TELEMETRY_EVENT_CREATED, 'suite')
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       this.enter(testSuiteSpan, store)
       this.testSuiteSpan = testSuiteSpan
     })
 
     this.addSub('ci:vitest:test-suite:finish', ({ status, onFinish }) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       const span = store?.span
       if (span) {
         span.setTag(TEST_STATUS, status)
@@ -237,13 +237,13 @@ class VitestPlugin extends CiPlugin {
       this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'suite')
       // TODO: too frequent flush - find for method in worker to decrease frequency
       this.tracer._exporter.flush(onFinish)
-      if (this.runningTestProbeId) {
-        this.removeDiProbe(this.runningTestProbeId)
+      if (this.runningTestProbe) {
+        this.removeDiProbe(this.runningTestProbe)
       }
     })
 
     this.addSub('ci:vitest:test-suite:error', ({ error }) => {
-      const store = storage.getStore()
+      const store = storage('legacy').getStore()
       const span = store?.span
       if (span && error) {
         span.setTag('error', error)