dd-trace 6.0.0 → 6.2.0

package/packages/datadog-plugin-graphql/src/execute.js

@@ -5,12 +5,10 @@ const dc = require('dc-polyfill')
 const { storage } = require('../../datadog-core')
 const TracingPlugin = require('../../dd-trace/src/plugins/tracing')
 const GraphQLParsePlugin = require('./parse')
-const { extractErrorIntoSpanEvent, getSignature } = require('./utils')
+const { extractErrorIntoSpanEvent, getOperation, getSignature } = require('./utils')
 
 const legacyStorage = storage('legacy')
 
-const types = new Set(['query', 'mutation', 'subscription'])
-
 const iastResolveCh = dc.channel('apm:graphql:resolve:start')
 const resolverStartCh = dc.channel('datadog:graphql:resolver:start')
 const updateFieldCh = dc.channel('apm:graphql:resolve:updateField')
@@ -115,9 +113,11 @@ class GraphQLExecutePlugin extends TracingPlugin {
 
     ctx.collapse = this.config.collapse
 
+    const signature = getSignature(document, name, type, this.config.signature)
+
     const span = this.startSpan(this.operationName(), {
       service: this.config.service || this.serviceName(),
-      resource: getSignature(document, name, type, this.config.signature),
+      resource: signature,
       kind: this.constructor.kind,
       type: this.constructor.type,
       meta: {
@@ -241,7 +241,7 @@ class GraphQLExecutePlugin extends TracingPlugin {
     if (res?.errors?.length) {
       span.setTag('error', res.errors[0])
       for (const err of res.errors) {
-        extractErrorIntoSpanEvent(this._tracerConfig, span, err)
+        extractErrorIntoSpanEvent(this.config, span, err)
       }
     }
 
@@ -271,6 +271,9 @@ class GraphQLExecutePlugin extends TracingPlugin {
     const loc = this.config.source && document && fieldNode && fieldNode.loc
     const source = loc && document.slice(loc.start, loc.end)
 
+    // ctx form: startSpan sets field.currentStore = { ...activeStore, span }
+    // without entering it. Only the field's first resolver call runs in that
+    // store (isFirst check in wrapResolve); siblings use field.parentStore.
     const span = this.startSpan('graphql.resolve', {
       service: this.config.service,
       resource: `${fieldName}:${returnType}`,
@@ -283,7 +286,9 @@ class GraphQLExecutePlugin extends TracingPlugin {
         'graphql.field.type': baseTypeName,
         'graphql.source': source,
       },
-    }, false)
+    }, field)
+
+    field.span = span
 
     if (fieldNode && this.config.variables && fieldNode.arguments) {
       const variables = this.config.variables(variableValues)
@@ -388,15 +393,22 @@ function wrapResolve (resolve) {
         pathString,
         collapsedKey: collapsedKey ?? pathString,
         span: null,
+        // Set by startResolveSpan; currentStore is used by the first resolver
+        // call only, siblings use parentStore (see the isFirst check below).
+        parentStore: null,
+        currentStore: null,
       }
       rootCtx.fields.set(fieldKey, field)
     }
 
     // Collapsed siblings still publish updateField (master's contract: one
     // publish per resolver call, even when the span is collapsed) and route
-    // through callInAsyncScope so the abort signal stops them mid-flight.
+    // through callInAsyncScope so the abort signal stops them mid-flight. They
+    // run in the parent store, not field.currentStore: the first sibling's
+    // synchronous resolver already finished the shared graphql.resolve span, so
+    // re-entering its store would parent user spans to a closed span.
     if (!isFirst) {
-      return callInAsyncScope(resolve, this, arguments, rootCtx.abortController, (err) => {
+      return callInAsyncScope(resolve, this, arguments, rootCtx.abortController, field.parentStore, (err) => {
         if (updateFieldCh.hasSubscribers) {
           updateFieldCh.publish({ rootCtx, field, error: err, pathString: field.pathString })
         }
@@ -406,9 +418,8 @@ function wrapResolve (resolve) {
     const executeSpan = rootCtx.executeSpan
     const startTime = executeSpan._getTime()
     const span = rootCtx.plugin.startResolveSpan(field, rootCtx, executeSpan, startTime)
-    field.span = span
 
-    return callInAsyncScope(resolve, this, arguments, rootCtx.abortController, (err, res) => {
+    return callInAsyncScope(resolve, this, arguments, rootCtx.abortController, field.currentStore, (err, res) => {
       const endTime = executeSpan._getTime()
       rootCtx.plugin.finishResolveSpan(span, field, err, res, endTime || startTime)
       if (updateFieldCh.hasSubscribers) {
@@ -446,16 +457,16 @@ function wrapFieldType (field) {
   wrapFields(unwrapped)
 }
 
-function callInAsyncScope (fn, thisArg, args, abortController, cb) {
-  cb ??= () => {}
-
+// Runs the resolver inside `store`, including any code after an internal
+// `await`. A `.then()` the caller attaches afterward runs outside `store`.
+function callInAsyncScope (fn, thisArg, args, abortController, store, cb) {
   if (abortController?.signal.aborted) {
     cb(null, null)
     throw new AbortError('Aborted')
   }
 
   try {
-    const result = fn.apply(thisArg, args)
+    const result = legacyStorage.run(store, () => fn.apply(thisArg, args))
     if (typeof result?.then === 'function') {
       return result.then(
         res => { cb(null, res); return res },
@@ -658,17 +669,6 @@ function defaultFieldResolver (source, args, contextValue, info) {
   }
 }
 
-function getOperation (document, operationName) {
-  if (!document || !Array.isArray(document.definitions)) return
-
-  for (const definition of document.definitions) {
-    if (definition && types.has(definition.operation) &&
-        (!operationName || definition.name?.value === operationName)) {
-      return definition
-    }
-  }
-}
-
 function addVariableTags (config, span, variableValues) {
   if (!variableValues || !config.variables) return
 

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

@@ -6,6 +6,7 @@ const CompositePlugin = require('../../dd-trace/src/plugins/composite')
 const log = require('../../dd-trace/src/log')
 const GraphQLExecutePlugin = require('./execute')
 const GraphQLParsePlugin = require('./parse')
+const GraphQLRequestPlugin = require('./request')
 const GraphQLValidatePlugin = require('./validate')
 
 class GraphQLPlugin extends CompositePlugin {
@@ -14,6 +15,11 @@ class GraphQLPlugin extends CompositePlugin {
     return {
       execute: GraphQLExecutePlugin,
       parse: GraphQLParsePlugin,
+      // Top-level request span for drivers (mercurius) that funnel through a
+      // single entry point and parse/execute internally. graphql-js, apollo,
+      // and yoga produce no such channel, so the plugin simply never fires for
+      // them.
+      request: GraphQLRequestPlugin,
       validate: GraphQLValidatePlugin,
       // resolve plugin is absorbed into execute: per-field data is recorded
       // synchronously in wrapResolve, and all graphql.resolve spans are
@@ -30,6 +36,13 @@ class GraphQLPlugin extends CompositePlugin {
 }
 
 // config validator helpers
+//
+// `collapse`, `depth`, `variables`, and `errorExtensions` arrive pre-merged on
+// `config`: plugin_manager seeds the `DD_TRACE_GRAPHQL_*` env values (already
+// parsed to their declared type) as the base and a programmatic
+// `tracer.use('graphql', …)` value overrides them. So these helpers only shape
+// the merged value — coerce, validate, and turn `variables` into a filter — and
+// never read the environment themselves.
 
 function validateConfig (config) {
   const collapse = config.collapse === undefined || !!config.collapse
@@ -41,6 +54,7 @@ function validateConfig (config) {
     // v5 counted collapsed list indices toward `depth`, so the same query reached a
     // different depth depending on `collapse`. v6 counts selection-set depth only.
     countListIndices: DD_MAJOR < 6 && collapse,
+    errorExtensions: getErrorExtensions(config),
     hooks: getHooks(config),
   }
 }
@@ -58,13 +72,21 @@ function getVariablesFilter (config) {
   if (typeof config.variables === 'function') {
     return config.variables
   } else if (Array.isArray(config.variables)) {
-    return variables => pick(variables, config.variables)
+    return config.variables.length > 0 ? variables => pick(variables, config.variables) : null
   } else if (config.hasOwnProperty('variables')) {
     log.error('Expected `variables` to be an array or function.')
   }
   return null
 }
 
+function getErrorExtensions (config) {
+  if (Array.isArray(config.errorExtensions)) {
+    return config.errorExtensions.length > 0 ? config.errorExtensions : undefined
+  } else if (config.hasOwnProperty('errorExtensions')) {
+    log.error('Expected `errorExtensions` to be an array.')
+  }
+}
+
 const noop = () => {}
 const noopHooks = { execute: noop, parse: noop, validate: noop, resolve: undefined }
 

package/packages/datadog-plugin-graphql/src/request.js

@@ -0,0 +1,104 @@
+'use strict'
+
+const TracingPlugin = require('../../dd-trace/src/plugins/tracing')
+const { extractErrorIntoSpanEvent, getCachedRequestOperation } = require('./utils')
+
+// Top-level GraphQL request span for drivers that funnel every operation
+// through a single entry point but parse/validate/execute internally (mercurius
+// today). It parents the `graphql.parse`/`graphql.validate`/`graphql.execute`
+// (or JIT) sub-spans and carries the request text + operation name/type, which
+// otherwise have no home when the query is JIT-compiled and `graphql.execute`
+// never fires.
+//
+// The entry boundary only hands us the raw `source` (string or pre-parsed AST)
+// and `operationName`; the parsed document — and therefore the precise
+// operation signature — is only known once mercurius parses internally. On the
+// cold path the `validate` sub-plugin refines the resource/operation tags onto
+// this span via `ctx.currentStore.graphqlRequestSpan` once the document is
+// available, so we never re-parse on the hot path. On the JIT warm path no
+// sub-span fires, so we recover the same tags from the cache the cold path
+// populated, keyed by source + operationName.
+class GraphQLRequestPlugin extends TracingPlugin {
+  static id = 'graphql'
+  static operation = 'request'
+  static type = 'graphql'
+  static kind = 'server'
+  static prefix = 'tracing:orchestrion:mercurius:apm:graphql:request'
+
+  bindStart (ctx) {
+    // fastifyGraphQl(source, context, variables, operationName)
+    const source = ctx.arguments?.[0]
+    const operationName = ctx.arguments?.[3]
+
+    // `source` is the request text on the common path, but mercurius also
+    // accepts a pre-parsed document AST; only a string is the query text, and
+    // `graphql.source` carries only the text form.
+    const docSource = typeof source === 'string' ? source : undefined
+
+    // Warm (JIT-compiled) path: execute never fires, so recover the operation
+    // signature/type the cold path cached, keyed by source + operationName —
+    // by query text for a string, by document identity for a pre-parsed AST.
+    // Empty on the cold path — validate hasn't refined yet — where the request
+    // span is refined from the parsed document instead.
+    const cached = getCachedRequestOperation(source, operationName)
+
+    const span = this.startSpan(this.operationName({ id: 'request' }), {
+      service: this.config.service || this.serviceName(),
+      // The cached signature is the precise resource; otherwise provisional and
+      // refined by the validate sub-plugin once the document is parsed.
+      // `operationName` is the best name at the boundary; falls back to the
+      // operation signature once validate sees the document.
+      resource: cached?.signature || operationName || undefined,
+      kind: this.constructor.kind,
+      type: this.constructor.type,
+      meta: {
+        'graphql.operation.type': cached?.type,
+        'graphql.operation.name': cached?.name || operationName,
+        'graphql.source': this.config.source ? docSource : undefined,
+      },
+    }, ctx)
+
+    // Hand the span, the requested operation name, and the raw source to the
+    // validate sub-plugin running inside this store so it can refine the
+    // resource + operation tags from the parsed document (validate is the first
+    // boundary that has it) and cache them keyed by the source the request
+    // boundary saw. The raw source is the cache key — validate sees mercurius's
+    // internally parsed document, not the caller's source, and for a pre-parsed
+    // AST the two are different objects.
+    ctx.currentStore.graphqlRequestSpan = span
+    ctx.currentStore.graphqlRequestOperationName = operationName
+    ctx.currentStore.graphqlRequestSource = source
+
+    return ctx.currentStore
+  }
+
+  asyncEnd (ctx) {
+    /* istanbul ignore next: currentStore is populated for the request lifecycle; activeSpan is base-plugin fallback. */
+    const span = ctx?.currentStore?.span || this.activeSpan
+    /* istanbul ignore if: startSpan always populates currentStore for the request lifecycle. */
+    if (!span) return ctx.parentStore
+
+    const result = ctx.result
+    if (result?.errors?.length) {
+      span.setTag('error', result.errors[0])
+      for (const error of result.errors) {
+        extractErrorIntoSpanEvent(this._tracerConfig, span, error)
+      }
+    }
+
+    span.finish()
+
+    return ctx.parentStore
+  }
+
+  error (ctx) {
+    /* istanbul ignore next: currentStore is populated for request errors; activeSpan is base-plugin fallback. */
+    const span = ctx?.currentStore?.span || this.activeSpan
+    /* istanbul ignore else: errors are only routed after the request span has started. */
+    if (span && ctx?.error) {
+      span.setTag('error', ctx.error)
+    }
+  }
+}
+
+module.exports = GraphQLRequestPlugin

package/packages/datadog-plugin-graphql/src/utils.js

@@ -1,5 +1,186 @@
 'use strict'
 
+const { LRUCache } = require('../../../vendor/dist/lru-cache')
+
+/**
+ * @typedef {{ signature?: string, type?: string, name?: string }} RequestOperation
+ */
+
+const operationTypes = new Set(['query', 'mutation', 'subscription'])
+
+// Mercurius funnels every operation through `fastifyGraphQl`, but the parsed
+// document — and therefore the operation signature/type/name — is only known
+// once mercurius parses internally. The top-level request span opens before
+// that, and on the JIT warm path neither parse/validate nor execute fires, so
+// the span would otherwise be left with only the provisional resource. The cold
+// path caches the computed metadata; the request boundary reads it back on the
+// warm path. Bounded so a flood of distinct queries can't grow it without limit.
+//
+// The key is the operation name plus the raw query text, not the source alone:
+// mercurius keys its document LRU by source but compiles the JIT for a single
+// `operationName`, and the compiled query then serves that operation for every
+// later request that shares the source — regardless of the `operationName` those
+// requests ask for. A source-only key would hand a warm request for operation B
+// the metadata of whichever operation was cached last for that source (A),
+// mislabeling the span. Operation names cannot contain a newline, so it is a
+// safe separator that keeps the two parts from colliding.
+const requestOperationCache = new LRUCache({ max: 500 })
+
+// Mercurius also accepts a pre-parsed document AST as the source, which reaches
+// the request boundary as an object rather than query text — so there is no
+// string to key the LRU by. Mercurius keys its own document LRU by that source
+// object's identity, and the same object reaches the boundary on the warm path,
+// so a WeakMap keyed by the caller-owned document recovers the metadata without
+// mutating the document and releases with it. The value carries the requested
+// operation name so a JIT-only sibling selection is not handed another
+// operation's metadata (same reason the string cache keys by operation name).
+/** @type {WeakMap<object, Map<string | undefined, RequestOperation>>} */
+const documentOperationCache = new WeakMap()
+
+/**
+ * @param {string} source - The raw query text; the same key mercurius uses.
+ * @param {string | undefined} operationName - The requested operation name.
+ * @returns {string}
+ */
+function requestOperationKey (source, operationName) {
+  return `${operationName ?? ''}\n${source}`
+}
+
+/**
+ * @param {unknown} source - Query text on the common path; a pre-parsed
+ *   document AST otherwise. Any other shape (mercurius rejects it before
+ *   execute) has no cache entry and yields undefined.
+ * @param {string | undefined} operationName - The requested operation name.
+ * @returns {RequestOperation | undefined}
+ */
+function getCachedRequestOperation (source, operationName) {
+  if (typeof source === 'string') {
+    return requestOperationCache.get(requestOperationKey(source, operationName))
+  }
+  if (source === null || typeof source !== 'object') return
+  return documentOperationCache.get(source)?.get(operationName)
+}
+
+/**
+ * A string source keys the text LRU; a document AST keys the WeakMap. Any other
+ * shape has no usable key — mercurius rejects it before execute, so the warm
+ * path never reaches the request span for it either.
+ *
+ * @param {unknown} source
+ * @returns {source is string | object}
+ */
+function isCacheableSource (source) {
+  return typeof source === 'string' || (source !== null && typeof source === 'object')
+}
+
+/**
+ * Select the operation definition matching `operationName`, or the first one
+ * when no name is given (graphql/mercurius default selection).
+ *
+ * @param {import('graphql').DocumentNode | undefined} document
+ * @param {string | undefined} operationName
+ * @returns {import('graphql').OperationDefinitionNode | undefined}
+ */
+function getOperation (document, operationName) {
+  /* istanbul ignore if: validate/execute only call this with a parsed GraphQL document. */
+  if (!document || !Array.isArray(document.definitions)) return
+
+  for (const definition of document.definitions) {
+    if (operationTypes.has(definition?.operation) &&
+        (!operationName || definition.name?.value === operationName)) {
+      return definition
+    }
+  }
+}
+
+/**
+ * Refine the top-level graphql.request span (mercurius) from the parsed
+ * document and cache the metadata so the JIT warm path — where no sub-span
+ * fires — can recover the same tags at the request boundary.
+ *
+ * This runs at the first boundary that has the document (validate on the cold
+ * path, which also precedes a pre-execute validation failure). It is idempotent
+ * across the later execute boundary via the `ddRequestRefined` flag, and a
+ * no-op for graphql-js/apollo/yoga, which never open a request span.
+ *
+ * Every named operation in the document is cached, not just the selected one:
+ * a multi-operation document parses once, and a later request may select a
+ * sibling operation that mercurius then serves exclusively through its JIT path
+ * (no execute span), so its metadata has to be ready before that happens.
+ *
+ * @param {import('../../dd-trace/src/opentracing/span') | undefined} requestSpan
+ * @param {import('graphql').DocumentNode | undefined} document
+ * @param {unknown} requestSource - The raw source the request boundary saw:
+ *   query text on the common path, a pre-parsed document AST otherwise. The
+ *   cache is keyed by it, not by the parsed document, so the request boundary
+ *   recovers the metadata on the warm path from the same value mercurius keys
+ *   its own document LRU by. Any other shape has no usable key and is not
+ *   cached (the warm path never reaches this span for it either).
+ * @param {string | undefined} operationName - The requested operation name.
+ * @param {boolean} calculateSignature - The graphql plugin's `signature` config.
+ */
+function refineRequestSpan (requestSpan, document, requestSource, operationName, calculateSignature) {
+  /* istanbul ignore if: validate only refines after the request span and parsed document exist. */
+  if (!requestSpan || requestSpan.ddRequestRefined || !document) return
+  requestSpan.ddRequestRefined = true
+
+  const operation = getOperation(document, operationName)
+  const type = operation?.operation
+  const name = operation?.name?.value
+  const signature = getSignature(document, name, type, calculateSignature)
+
+  if (signature) requestSpan.setTag('resource.name', signature)
+  if (type) requestSpan.setTag('graphql.operation.type', type)
+  if (name) requestSpan.setTag('graphql.operation.name', name)
+
+  if (!isCacheableSource(requestSource)) return
+
+  // Cache the selected operation under the requested name (undefined selects
+  // the document's first operation, so it shares the entry with that name).
+  cacheRequestOperation(requestSource, operationName, { signature, type, name })
+
+  // Cache every named operation so a JIT-only sibling selection is labeled from
+  // this single parse instead of falling back to a bare operation name.
+  for (const definition of document.definitions) {
+    const definitionName = definition?.name?.value
+    if (definitionName === undefined || !operationTypes.has(definition.operation)) continue
+    if (definitionName === operationName) continue
+
+    cacheRequestOperation(requestSource, definitionName, {
+      signature: getSignature(document, definitionName, definition.operation, calculateSignature),
+      type: definition.operation,
+      name: definitionName,
+    })
+  }
+}
+
+/**
+ * @param {string | import('graphql').DocumentNode} source - Query text keys the
+ *   text LRU; a caller-owned document AST keys the WeakMap (never mutated).
+ * @param {string | undefined} operationName - The requested operation name.
+ * @param {RequestOperation} operation
+ */
+function cacheRequestOperation (source, operationName, operation) {
+  if (typeof source === 'string') {
+    requestOperationCache.set(requestOperationKey(source, operationName), operation)
+    return
+  }
+
+  let operations = documentOperationCache.get(source)
+  if (operations === undefined) {
+    operations = new Map()
+    documentOperationCache.set(source, operations)
+  }
+  operations.set(operationName, operation)
+}
+
+/**
+ * @param {{ errorExtensions?: string[] }} config Resolved plugin config; `errorExtensions` lists the
+ *   GraphQL error `extensions` keys to copy onto the span event.
+ * @param {import('../../dd-trace/src/opentracing/span')} span
+ * @param {{ name?: string, message?: string, stack?: string, locations?: Array<{ line: number, column: number }>,
+ *   path?: Array<string|number>, extensions?: Record<string, unknown> }} exc
+ */
 function extractErrorIntoSpanEvent (config, span, exc) {
   const attributes = {}
 
@@ -29,8 +210,8 @@ function extractErrorIntoSpanEvent (config, span, exc) {
     attributes.message = exc.message
   }
 
-  if (config.DD_TRACE_GRAPHQL_ERROR_EXTENSIONS) {
-    for (const ext of config.DD_TRACE_GRAPHQL_ERROR_EXTENSIONS) {
+  if (config.errorExtensions) {
+    for (const ext of config.errorExtensions) {
       if (exc.extensions?.[ext]) {
         const value = exc.extensions[ext]
 
@@ -77,5 +258,8 @@ function getSignature (document, operationName, operationType, calculate) {
 
 module.exports = {
   extractErrorIntoSpanEvent,
+  getCachedRequestOperation,
+  getOperation,
   getSignature,
+  refineRequestSpan,
 }

package/packages/datadog-plugin-graphql/src/validate.js

@@ -1,8 +1,11 @@
 'use strict'
 
+const { storage } = require('../../datadog-core')
 const TracingPlugin = require('../../dd-trace/src/plugins/tracing')
 const GraphQLParsePlugin = require('./parse')
-const { extractErrorIntoSpanEvent } = require('./utils')
+const { extractErrorIntoSpanEvent, refineRequestSpan } = require('./utils')
+
+const legacyStorage = storage('legacy')
 
 class GraphQLValidatePlugin extends TracingPlugin {
   static id = 'graphql'
@@ -15,6 +18,26 @@ class GraphQLValidatePlugin extends TracingPlugin {
     const docSource = document ? GraphQLParsePlugin.documentSources.get(document) : undefined
     const source = this.config.source && document && docSource
 
+    // Refine the top-level graphql.request span (mercurius) from the parsed
+    // document. validate is the first boundary that has it and precedes both
+    // execute and any pre-execute rejection (unknown field, GET mutation), so a
+    // failing request still ends up with a resource and operation tags. The
+    // request span, its operation name, and the raw source ride the active
+    // store the request boundary entered (validate's own `ctx.currentStore` is
+    // not populated yet). The cache is keyed by that raw source, not the parsed
+    // document — for a pre-parsed AST mercurius validates a structuredClone, so
+    // the document here is a different object from the one the boundary saw and
+    // recovers on the warm path. No-op for graphql-js/apollo/yoga, which never
+    // open a request span.
+    const requestStore = legacyStorage.getStore()
+    refineRequestSpan(
+      requestStore?.graphqlRequestSpan,
+      document,
+      requestStore?.graphqlRequestSource,
+      requestStore?.graphqlRequestOperationName,
+      this.config.signature
+    )
+
     this.startSpan('graphql.validate', {
       service: this.config.service,
       type: 'graphql',
@@ -38,7 +61,7 @@ class GraphQLValidatePlugin extends TracingPlugin {
     if (errors?.length) {
       span.setTag('error', errors[0])
       for (const err of errors) {
-        extractErrorIntoSpanEvent(this._tracerConfig, span, err)
+        extractErrorIntoSpanEvent(this.config, span, err)
       }
     }
 

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

@@ -61,15 +61,6 @@ const isJestWorker = !!getEnvironmentVariable('JEST_WORKER_ID')
 // https://github.com/facebook/jest/blob/d6ad15b0f88a05816c2fe034dd6900d28315d570/packages/jest-worker/src/types.ts#L38
 const CHILD_MESSAGE_END = 2
 
-function withTimeout (promise, timeoutMs) {
-  return new Promise(resolve => {
-    realSetTimeout(resolve, timeoutMs)
-
-    // Also resolve if the original promise resolves
-    promise.then(resolve)
-  })
-}
-
 class JestPlugin extends CiPlugin {
   static id = 'jest'
 
@@ -451,6 +442,7 @@ class JestPlugin extends CiPlugin {
       isAtrRetry,
       finalStatus,
       earlyFlakeAbortReason,
+      promises,
     }) => {
       span.setTag(TEST_STATUS, status)
       if (finalStatus) {
@@ -481,12 +473,33 @@ class JestPlugin extends CiPlugin {
         this.getTestTelemetryTags(span)
       )
 
-      span.finish()
-      finishAllTraceSpans(span)
-      this.activeTestSpan = null
+      const finish = () => {
+        span.finish()
+        finishAllTraceSpans(span)
+        this.activeTestSpan = null
+      }
+
+      if (finalStatus) {
+        if (promises?.hitBreakpointPromise) {
+          finish()
+          this.cancelDiBreakpointHitWait()
+          return
+        }
+        if (promises && this.diBreakpointHitPromise) {
+          promises.hitBreakpointPromise = this.waitForPreparedDiBreakpointHit().then(finish)
+          return
+        }
+        finish()
+        this.cancelDiBreakpointHitWait()
+        return
+      }
+      finish()
+      if (status === 'fail' && promises?.hitBreakpointPromise) {
+        this.prepareDiBreakpointHitWait()
+      }
     })
 
-    this.addSub('ci:jest:test:err', ({ span, error, shouldSetProbe, promises }) => {
+    this.addSub('ci:jest:test:err', ({ span, error, shouldSetProbe, shouldWaitForHitProbe, promises }) => {
       if (error && span) {
         span.setTag(TEST_STATUS, 'fail')
         span.setTag('error', getFormattedError(error, this.repositoryRoot))
@@ -494,8 +507,11 @@ class JestPlugin extends CiPlugin {
           const probeInformation = this.addDiProbe(error)
           if (probeInformation) {
             const { setProbePromise } = probeInformation
-            promises.isProbeReady = withTimeout(setProbePromise, 2000)
+            this.prepareDiBreakpointHitWait()
+            promises.isProbeReady = this.waitForDiOperation(setProbePromise)
           }
+        } else if (shouldWaitForHitProbe && this.di) {
+          promises.hitBreakpointPromise = this.waitForDiBreakpointHits()
         }
       }
     })