dd-trace 5.116.0 → 5.118.0

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

@@ -46,6 +46,10 @@ const {
   addCoverageBackfillUntestedFiles,
   getCoverageBackfillFiles,
 } = require('./jest/coverage-backfill')
+const {
+  getChannelPromise,
+  publishWithCompletion,
+} = require('./helpers/channel')
 const { addHook, channel } = require('./helpers/instrument')
 
 const testSessionStartCh = channel('ci:jest:session:start')
@@ -1971,6 +1975,9 @@ function shouldFinishBailTestSession (globalConfig, results) {
   return !!globalConfig?.bail && getNumBailFailures(results) >= globalConfig.bail
 }
 
+/**
+ * @param {Record<string, unknown> & { onDone?: () => void }} payload
+ */
 async function waitForTestSessionFinish (payload) {
   if (!testSessionFinishCh.hasSubscribers || hasFinishedTestSession) return
 
@@ -1978,8 +1985,9 @@ async function waitForTestSessionFinish (payload) {
 
   let timeoutId
 
+  let onDone
   const flushPromise = new Promise((resolve) => {
-    payload.onDone = () => {
+    onDone = () => {
       clearTimeout(timeoutId)
       resolve()
     }
@@ -1992,7 +2000,7 @@ async function waitForTestSessionFinish (payload) {
     timeoutId.unref?.()
   })
 
-  testSessionFinishCh.publish(payload)
+  publishWithCompletion(testSessionFinishCh, payload, onDone)
 
   const waitingResult = await Promise.race([flushPromise, timeoutPromise])
 
@@ -2157,12 +2165,6 @@ function getWrappedScheduleTests (scheduleTests, frameworkVersion) {
   }
 }
 
-function getChannelPromise (channelToPublishTo, payload = {}) {
-  return new Promise(resolve => {
-    channelToPublishTo.publish({ ...payload, onDone: resolve })
-  })
-}
-
 function searchSourceWrapper (searchSourcePackage, frameworkVersion) {
   const SearchSource = searchSourcePackage.default ?? searchSourcePackage
 
@@ -2517,9 +2519,7 @@ function getCliWrapper (isNewJestVersion) {
 
       if (codeCoverageReportCh.hasSubscribers) {
         const rootDir = result.globalConfig?.rootDir || process.cwd()
-        await new Promise((resolve) => {
-          codeCoverageReportCh.publish({ rootDir, onDone: resolve })
-        })
+        await getChannelPromise(codeCoverageReportCh, { rootDir })
       }
 
       logSessionSummary(ignoredFailuresSummary, getAttemptToFixExecutionsFromJestResults(result))
@@ -2537,6 +2537,10 @@ function shouldWaitForTestSuiteFinish (environment) {
   return isJestWorker && environment.globalConfig?.workerIdleMemoryLimit !== undefined
 }
 
+/**
+ * @param {Record<string, unknown>} payload
+ * @param {boolean} waitForFinish
+ */
 function publishTestSuiteFinish (payload, waitForFinish) {
   if (!testSuiteFinishCh.hasSubscribers) return
 
@@ -2545,12 +2549,9 @@ function publishTestSuiteFinish (payload, waitForFinish) {
     return
   }
 
-  return new Promise(resolve => {
-    testSuiteFinishCh.publish({
-      ...payload,
-      waitForFinish,
-      onDone: resolve,
-    })
+  return getChannelPromise(testSuiteFinishCh, {
+    ...payload,
+    waitForFinish,
   })
 }
 

package/packages/datadog-instrumentations/src/mocha/main.js

@@ -3,6 +3,11 @@
 const { createCoverageMap } = require('../../../../vendor/dist/istanbul-lib-coverage')
 const satisfies = require('../../../../vendor/dist/semifies')
 const { DD_MAJOR } = require('../../../../version')
+const {
+  getRunStoresPromise,
+  publishWithCompletion,
+  runStoresWithCompletion,
+} = require('../helpers/channel')
 const { addHook, channel } = require('../helpers/instrument')
 const shimmer = require('../../../datadog-shimmer')
 const { isMarkedAsUnskippable } = require('../../../datadog-plugin-jest/src/util')
@@ -240,6 +245,41 @@ function resetSuiteSkippingRunState () {
   writeCoverageBackfillToCache({})
 }
 
+/**
+ * @param {((failures: number) => void) | undefined} callback
+ * @returns {{ onRunDone: (failures: number) => void, onFlushDone: () => void }}
+ */
+function getRunCompletionCallbacks (callback) {
+  let failures
+  let hasRunFinished = false
+  let hasFlushFinished = false
+  let hasCompleted = false
+  const onDone = callback || (() => {})
+
+  const completeIfReady = () => {
+    if (hasCompleted || !hasRunFinished || !hasFlushFinished) return
+
+    hasCompleted = true
+    onDone(failures)
+  }
+
+  return {
+    onRunDone: (runFailures) => {
+      if (hasRunFinished) return
+
+      failures = runFailures
+      hasRunFinished = true
+      completeIfReady()
+    },
+    onFlushDone: () => {
+      if (hasFlushFinished) return
+
+      hasFlushFinished = true
+      completeIfReady()
+    },
+  }
+}
+
 function getOnStartHandler (frameworkVersion) {
   return function () {
     const processArgv = process.argv.slice(2).join(' ')
@@ -251,7 +291,12 @@ function getOnStartHandler (frameworkVersion) {
   }
 }
 
-function getOnEndHandler (isParallel) {
+/**
+ * @param {boolean} isParallel
+ * @param {() => void} onDone
+ * @returns {() => void}
+ */
+function getOnEndHandler (isParallel, onDone) {
   return function () {
     let status = 'pass'
     let error
@@ -337,7 +382,7 @@ function getOnEndHandler (isParallel) {
       global.__coverage__ = fromCoverageMapToCoverage(originalCoverageMap)
     }
 
-    testSessionFinishCh.publish({
+    publishWithCompletion(testSessionFinishCh, {
       status,
       isSuitesSkipped,
       testCodeCoverageLinesTotal,
@@ -350,19 +395,13 @@ function getOnEndHandler (isParallel) {
       isEarlyFlakeDetectionFaulty: config.isEarlyFlakeDetectionFaulty,
       isTestManagementEnabled: config.isTestManagementTestsEnabled,
       isParallel,
-    })
+    }, onDone)
 
     logTestOptimizationSummary({ attemptToFixExecutions, newTestsWithDynamicNames })
     loggedAttemptToFixTests.clear()
   }
 }
 
-function getRunStoresPromise (channelToPublishTo, ctx) {
-  return new Promise(resolve => {
-    channelToPublishTo.runStores({ ...ctx, onDone: resolve }, () => {})
-  })
-}
-
 function applyKnownTestsResponse ({ err, knownTests }) {
   if (err) {
     config.knownTests = []
@@ -395,13 +434,14 @@ function getExecutionConfiguration (runner, isParallel, frameworkVersion, onFini
   let skippableSuitesResponse
   resetSuiteSkippingRunState()
 
-  const onReceivedSkippableSuites = ({
-    err,
-    skippableSuites,
-    itrCorrelationId: responseItrCorrelationId,
-    skippableSuitesCoverage: responseSkippableSuitesCoverage,
-  }) => {
-    if (err) {
+  const onReceivedSkippableSuites = (response) => {
+    const {
+      err,
+      skippableSuites,
+      itrCorrelationId: responseItrCorrelationId,
+      skippableSuitesCoverage: responseSkippableSuitesCoverage,
+    } = response || {}
+    if (!response || err) {
       suitesToSkip = []
       skippableSuitesCoverage = {}
     } else {
@@ -443,12 +483,12 @@ function getExecutionConfiguration (runner, isParallel, frameworkVersion, onFini
       return
     }
 
-    ctx.onDone = onReceivedSkippableSuites
-    skippableSuitesCh.runStores(ctx, () => {})
+    runStoresWithCompletion(skippableSuitesCh, ctx, onReceivedSkippableSuites)
   }
 
-  const onReceivedImpactedTests = ({ err, modifiedFiles: receivedModifiedFiles }) => {
-    if (err) {
+  const onReceivedImpactedTests = (response) => {
+    const { err, modifiedFiles: receivedModifiedFiles } = response || {}
+    if (!response || err) {
       config.modifiedFiles = []
       config.isImpactedTestsEnabled = false
     } else {
@@ -465,8 +505,7 @@ function getExecutionConfiguration (runner, isParallel, frameworkVersion, onFini
 
   const continueAfterTestRequests = () => {
     if (config.isImpactedTestsEnabled) {
-      ctx.onDone = onReceivedImpactedTests
-      modifiedFilesCh.runStores(ctx, () => {})
+      runStoresWithCompletion(modifiedFilesCh, ctx, onReceivedImpactedTests)
     } else if (config.isSuitesSkippingEnabled) {
       requestSkippableSuites()
     } else {
@@ -476,8 +515,14 @@ function getExecutionConfiguration (runner, isParallel, frameworkVersion, onFini
     }
   }
 
-  const onReceivedConfiguration = ({ err, isTestDynamicInstrumentationEnabled, libraryConfig, repositoryRoot }) => {
-    if (err || !skippableSuitesCh.hasSubscribers || !knownTestsCh.hasSubscribers) {
+  const onReceivedConfiguration = (response) => {
+    const {
+      err,
+      isTestDynamicInstrumentationEnabled,
+      libraryConfig,
+      repositoryRoot,
+    } = response || {}
+    if (!response || err || !skippableSuitesCh.hasSubscribers || !knownTestsCh.hasSubscribers) {
       return mochaGlobalRunCh.runStores(ctx, () => {
         onFinishRequest()
       })
@@ -526,9 +571,7 @@ function getExecutionConfiguration (runner, isParallel, frameworkVersion, onFini
     })
   }
 
-  ctx.onDone = onReceivedConfiguration
-
-  libraryConfigurationCh.runStores(ctx, () => {})
+  runStoresWithCompletion(libraryConfigurationCh, ctx, onReceivedConfiguration)
 }
 
 // In this hook we delay the execution with options.delay to grab library configuration,
@@ -641,6 +684,9 @@ addHook({
       return run.apply(this, args)
     }
 
+    const { onRunDone, onFlushDone } = getRunCompletionCallbacks(args[0])
+    args[0] = onRunDone
+
     const { suitesByTestFile, numSuitesByTestFile } = getSuitesByTestFile(this.suite)
     // Root-level tests (direct children of root, no describe wrapper) keyed by file.
     // Populated during the root 'suite' event so the normal finish path can include them
@@ -752,7 +798,7 @@ addHook({
       finishRootSuiteForFile(test.file)
     }
 
-    const onEnd = getOnEndHandler(false)
+    const onEnd = getOnEndHandler(false, onFlushDone)
 
     this.once('start', getOnStartHandler(frameworkVersion))
 
@@ -1041,8 +1087,11 @@ addHook({
       return run.apply(this, arguments)
     }
 
+    const { onRunDone, onFlushDone } = getRunCompletionCallbacks(cb)
+    arguments[0] = onRunDone
+
     this.once('start', getOnStartHandler(frameworkVersion))
-    this.once('end', getOnEndHandler(true))
+    this.once('end', getOnEndHandler(true, onFlushDone))
 
     // Populate unskippable suites before config is fetched (matches serial mode at Mocha.prototype.run)
     for (const filePath of files) {
@@ -1082,7 +1131,7 @@ addHook({
         skippedSuites = skippedFiles
         skippedSuitesCoverage = getSkippedSuitesCoverageForRun()
         writeCoverageBackfillToCache(skippedSuitesCoverage, getCoverageRootDir())
-        run.apply(this, [cb, { files: filteredFiles }])
+        run.apply(this, [onRunDone, { files: filteredFiles }])
       } else {
         run.apply(this, arguments)
       }

package/packages/datadog-instrumentations/src/openai-agents.js

@@ -0,0 +1,159 @@
+'use strict'
+
+const { channel, tracingChannel } = require('dc-polyfill')
+const shimmer = require('../../datadog-shimmer')
+const { addHook, getHooks } = require('./helpers/instrument')
+
+// `WeakSet` keyed by module exports — replaces the underscored
+// `mod._datadogPatched` flag while keeping dedupe semantics. Mods are kept
+// alive by `require.cache` anyway, so this doesn't add lifetime to anything.
+const patchedMods = new WeakSet()
+const modelBaseURLs = new WeakMap()
+const chatCompletionsModelConstructorCh =
+  tracingChannel('orchestrion:@openai/agents-openai:OpenAIChatCompletionsModel_constructor')
+
+/**
+ * Capture the client base URL before the SDK stores the client in a #private field.
+ *
+ * @param {{ arguments?: [{ baseURL?: string }], self?: object }} ctx
+ */
+function captureChatCompletionsModelBaseURL (ctx) {
+  const baseURL = ctx.arguments?.[0]?.baseURL
+  if (ctx.self && typeof baseURL === 'string') modelBaseURLs.set(ctx.self, baseURL)
+}
+
+chatCompletionsModelConstructorCh.end.subscribe(captureChatCompletionsModelBaseURL)
+
+for (const hook of getHooks('@openai/agents-openai')) {
+  addHook(hook, moduleExports => moduleExports)
+}
+
+// Plugin subscribes to this and registers its TracingProcessor when
+// `@openai/agents` loads. Publishing from here keeps this file free of
+// any cross-package import from the plugin.
+const agentsCoreLoadedCh = channel('apm:openai-agents:agents-core:loaded')
+
+// Plugin uses addBind on this channel so that legacyStorage.run(store, fn) wraps
+// the model call — including async iterator advancement for streaming responses.
+// This ensures the active dd-trace span is visible to the openai plugin when it
+// creates its openai.request span, correctly parenting it under the agent span.
+const modelStartCh = channel('apm:openai-agents:model:start')
+
+// Tool.invoke runs inside agents-core's public function-span context. Bind the
+// corresponding dd-trace tool span around that public invocation boundary so
+// spans created by user tool code inherit it.
+const toolStartCh = channel('apm:openai-agents:tool:start')
+
+// Reference to the loaded @openai/agents module, captured in the first hook
+// so that wrapResponseMethod can call getCurrentSpan() without an additional
+// require (and without triggering n/no-missing-require on agents-core internals).
+let agentsMod
+
+// @openai/agents >=0.8.0 moved addTraceProcessor / getCurrentSpan out of the
+// top-level re-exports.  The new public surface uses getGlobalTraceProvider():
+//   provider.registerProcessor(processor)  (replaces addTraceProcessor)
+//   provider.getCurrentSpan()              (replaces getCurrentSpan)
+// Both APIs are tried so this file works across the full supported version range.
+// The plugin subscriber (index.js) handles processor registration via the channel.
+function getCurrentSpan () {
+  if (typeof agentsMod?.getCurrentSpan === 'function') {
+    return agentsMod.getCurrentSpan()
+  }
+  if (typeof agentsMod?.getGlobalTraceProvider === 'function') {
+    return agentsMod.getGlobalTraceProvider().getCurrentSpan()
+  }
+}
+
+function getCurrentSpanId () {
+  return getCurrentSpan()?.spanId
+}
+
+addHook({ name: '@openai/agents', versions: ['>=0.7.0'] }, (mod) => {
+  if (patchedMods.has(mod)) return mod
+  if (typeof mod?.addTraceProcessor !== 'function' && typeof mod?.getGlobalTraceProvider !== 'function') return mod
+  patchedMods.add(mod)
+  agentsMod = mod
+  if (typeof mod.tool === 'function') {
+    shimmer.wrap(mod, 'tool', wrapToolFactory, { replaceGetter: true })
+  }
+  agentsCoreLoadedCh.publish({ mod })
+  return mod
+})
+
+function wrapToolFactory (original) {
+  return function (...args) {
+    const tool = original.apply(this, args)
+    if (typeof tool?.invoke === 'function') {
+      shimmer.wrap(tool, 'invoke', wrapToolInvoke)
+    }
+    return tool
+  }
+}
+
+function wrapToolInvoke (original) {
+  return function (...args) {
+    const agentsCoreSpan = getCurrentSpan()
+    return toolStartCh.runStores({ agentsCoreSpan }, () => original.apply(this, args))
+  }
+}
+
+function wrapResponseMethod (original) {
+  return function (...args) {
+    const agentsCoreSpanId = getCurrentSpanId()
+    const baseURL = getClientBaseURL(this)
+    return modelStartCh.runStores({ agentsCoreSpanId, baseURL }, () => original.apply(this, args))
+  }
+}
+
+function wrapStreamedResponseMethod (original) {
+  return function (...args) {
+    const agentsCoreSpanId = getCurrentSpanId()
+    const baseURL = getClientBaseURL(this)
+    const context = { agentsCoreSpanId, baseURL }
+    const iterator = modelStartCh.runStores(context, () => original.apply(this, args))
+    return wrapAsyncIterator(iterator, context)
+  }
+}
+
+function getClientBaseURL (model) {
+  return model?.client?.baseURL ?? model?._client?.baseURL ?? modelBaseURLs.get(model)
+}
+
+function wrapAsyncIterator (iterator, context) {
+  if (!iterator || typeof iterator !== 'object') return iterator
+
+  return {
+    next () {
+      return modelStartCh.runStores(context, () => iterator.next.apply(iterator, arguments))
+    },
+    throw () {
+      if (typeof iterator.throw !== 'function') return Promise.reject(arguments[0])
+      return modelStartCh.runStores(context, () => iterator.throw.apply(iterator, arguments))
+    },
+    return () {
+      if (typeof iterator.return !== 'function') return Promise.resolve({ done: true, value: arguments[0] })
+      return modelStartCh.runStores(context, () => iterator.return.apply(iterator, arguments))
+    },
+    [Symbol.asyncIterator] () {
+      return this
+    },
+  }
+}
+
+addHook({ name: '@openai/agents-openai', versions: ['>=0.7.0'] }, (mod) => {
+  if (patchedMods.has(mod)) return mod
+  const responseProto = mod?.OpenAIResponsesModel?.prototype
+  const chatCompletionsProto = mod?.OpenAIChatCompletionsModel?.prototype
+  if (!responseProto && !chatCompletionsProto) return mod
+
+  patchedMods.add(mod)
+  for (const proto of [responseProto, chatCompletionsProto]) {
+    if (typeof proto?.getResponse === 'function') {
+      shimmer.wrap(proto, 'getResponse', wrapResponseMethod)
+    }
+    if (typeof proto?.getStreamedResponse === 'function') {
+      shimmer.wrap(proto, 'getStreamedResponse', wrapStreamedResponseMethod)
+    }
+  }
+  return mod
+})

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

@@ -29,6 +29,7 @@ const {
   RUM_TEST_EXECUTION_ID_COOKIE_NAME: RUM_COOKIE_NAME,
 } = require('../../dd-trace/src/ci-visibility/rum')
 const { DD_MAJOR } = require('../../../version')
+const { getChannelPromise } = require('./helpers/channel')
 const { addHook, channel, tracingChannel } = require('./helpers/instrument')
 
 const testStartCh = channel('ci:playwright:test:start')
@@ -614,12 +615,6 @@ function getTestByTestId (dispatcher, testId) {
   }
 }
 
-function getChannelPromise (channelToPublishTo, params) {
-  return new Promise(resolve => {
-    channelToPublishTo.publish({ onDone: resolve, ...params })
-  })
-}
-
 // Inspired by https://github.com/microsoft/playwright/blob/2b77ed4d7aafa85a600caa0b0d101b72c8437eeb/packages/playwright/src/reporters/base.ts#L293
 // We can't use test.outcome() directly because it's set on follow up handlers:
 // our `testEndHandler` is called before the outcome is set.
@@ -1206,8 +1201,6 @@ function dispatcherHookNew (dispatcherExport, runWrapper) {
 function runAllTestsWrapper (runAllTests, playwrightVersion) {
   // Config parameter is only available from >=1.55.0
   return async function (config) {
-    let onDone
-
     rootDir = getRootDir(this, config)
     const projects = getProjectsFromRunner(this, config)
     const isFailureScreenshotEnabled = isFailureScreenshotCaptureEnabled(projects)
@@ -1379,17 +1372,12 @@ function runAllTestsWrapper (runAllTests, playwrightVersion) {
     logTestOptimizationSummary({ attemptToFixExecutions, newTestsWithDynamicNames })
     loggedAttemptToFixTests.clear()
 
-    const flushWait = new Promise(resolve => {
-      onDone = resolve
-    })
-    testSessionFinishCh.publish({
+    await getChannelPromise(testSessionFinishCh, {
       status: preventedToFail ? 'pass' : STATUS_TO_TEST_STATUS[sessionStatus],
       isEarlyFlakeDetectionEnabled,
       isEarlyFlakeDetectionFaulty,
       isTestManagementTestsEnabled,
-      onDone,
     })
-    await flushWait
 
     startedSuites = []
     remainingTestsByFile = {}
@@ -2112,12 +2100,6 @@ function instrumentWorkerMainMethods (workerMain) {
       annotationTags = parseAnnotations(annotations)
     }
 
-    let onDone
-
-    const flushPromise = new Promise(resolve => {
-      onDone = resolve
-    })
-
     // Wait for the properties to be received, but do not block the worker forever if IPC fails.
     const ddPropertiesTimeoutPromise = new Promise(resolve => {
       const ddPropertiesTimeout = realSetTimeout(() => {
@@ -2143,7 +2125,7 @@ function instrumentWorkerMainMethods (workerMain) {
       testStatus: STATUS_TO_TEST_STATUS[status],
     })
 
-    testFinishCh.publish({
+    await getChannelPromise(testFinishCh, {
       testStatus: STATUS_TO_TEST_STATUS[status],
       steps: steps.filter(step => step.testId === testId),
       error,
@@ -2161,14 +2143,11 @@ function instrumentWorkerMainMethods (workerMain) {
       hasFailedAttemptToFixRetries: test._ddHasFailedAttemptToFixRetries,
       isAtrRetry: test._ddIsAtrRetry,
       isModified: test._ddIsModified,
-      onDone,
       finalStatus,
       earlyFlakeAbortReason: test._ddEarlyFlakeAbortReason,
       ...testCtx.currentStore,
     })
 
-    await flushPromise
-
     return res
   })
 

package/packages/datadog-instrumentations/src/vitest-main-no-worker-init.js

@@ -419,7 +419,7 @@ function createMainProcessReporter (reporterState) {
     testSuiteFinishCh.publish({
       status: getDatadogStatus(testSuiteResult),
       deferFlush: true,
-      onFinish: noop,
+      onDone: noop,
       ...testSuiteCtx.currentStore,
     })
     testSuiteContexts.delete(testModuleId)

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

@@ -16,6 +16,7 @@ const {
   getTestSuitePath,
   isModifiedTest,
 } = require('../../dd-trace/src/plugins/util/test')
+const { getChannelPromise } = require('./helpers/channel')
 const { addHook } = require('./helpers/instrument')
 const noWorkerInit = require('./vitest-main-no-worker-init')
 const {
@@ -38,7 +39,6 @@ const {
   workerReportLogsCh,
   codeCoverageReportCh,
   findExportByName,
-  getChannelPromise,
   getTypeTasks,
   getWorkspaceProject,
   setProvidedContext,
@@ -433,7 +433,8 @@ async function runMainProcessSetup (ctx, frameworkVersion, testSpecifications, s
       err,
       libraryConfig,
       requestErrorTags: receivedRequestErrorTags = {},
-    } = await getChannelPromise(libraryConfigurationCh, frameworkVersion, {
+    } = await getChannelPromise(libraryConfigurationCh, {
+      frameworkVersion,
       isVitestNoWorkerInitActive: shouldInstallNoWorkerInit,
     })
     requestErrorTags = receivedRequestErrorTags
@@ -450,7 +451,7 @@ async function runMainProcessSetup (ctx, frameworkVersion, testSpecifications, s
   resetMainProcessProvidedContext(ctx)
 
   if (testSessionConfigurationCh.hasSubscribers) {
-    testSessionConfiguration = await getChannelPromise(testSessionConfigurationCh, frameworkVersion)
+    testSessionConfiguration = await getChannelPromise(testSessionConfigurationCh, { frameworkVersion }) || {}
     const {
       testSessionId,
       testModuleId,
@@ -494,7 +495,7 @@ async function runMainProcessSetup (ctx, frameworkVersion, testSpecifications, s
 
   if (isKnownTestsEnabled) {
     const currentKnownTestsResponse = knownTestsResponse || await getChannelPromise(knownTestsCh)
-    if (currentKnownTestsResponse.err) {
+    if (!currentKnownTestsResponse || currentKnownTestsResponse.err) {
       isEarlyFlakeDetectionEnabled = false
     } else {
       knownTests = currentKnownTestsResponse.knownTests
@@ -538,12 +539,13 @@ async function runMainProcessSetup (ctx, frameworkVersion, testSpecifications, s
   }
 
   if (isTestManagementTestsEnabled) {
-    const { err, testManagementTests: receivedTestManagementTests } =
+    const testManagementResponse =
       testManagementTestsResponse || await getChannelPromise(testManagementTestsCh)
-    if (err) {
+    if (!testManagementResponse || testManagementResponse.err) {
       isTestManagementTestsEnabled = false
       log.error('Could not get test management tests.')
     } else {
+      const { testManagementTests: receivedTestManagementTests } = testManagementResponse
       testManagementTests = receivedTestManagementTests
       testManagementTestsBySuite = getTestManagementTestsBySuite(receivedTestManagementTests)
       shouldSendTestProperties = true
@@ -558,8 +560,7 @@ async function runMainProcessSetup (ctx, frameworkVersion, testSpecifications, s
 
   if (isImpactedTestsEnabled) {
     const modifiedFilesResponse = await getChannelPromise(modifiedFilesCh)
-    const { err } = modifiedFilesResponse
-    if (err) {
+    if (!modifiedFilesResponse || modifiedFilesResponse.err) {
       log.error('Could not get modified tests.')
     } else {
       modifiedFiles = modifiedFilesResponse.modifiedFiles
@@ -812,18 +813,13 @@ function getFinishWrapper (exitOrClose) {
       return exitOrClose.apply(this, arguments)
     }
 
-    let onFinish
-
-    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}.`)
     }
 
-    testSessionFinishCh.publish({
+    const flushPromise = getChannelPromise(testSessionFinishCh, {
       status: getSessionStatus(this.state),
       testCodeCoverageLinesTotal,
       error,
@@ -833,7 +829,6 @@ function getFinishWrapper (exitOrClose) {
       requestErrorTags,
       vitestPool,
       isVitestNoWorkerInitActive,
-      onFinish,
     })
 
     logTestOptimizationSummary({ attemptToFixExecutions, newTestsWithDynamicNames })
@@ -842,9 +837,7 @@ function getFinishWrapper (exitOrClose) {
 
     // If coverage was generated, publish coverage report channel for upload
     if (coverageRootDir && codeCoverageReportCh.hasSubscribers) {
-      await new Promise((resolve) => {
-        codeCoverageReportCh.publish({ rootDir: coverageRootDir, onDone: resolve })
-      })
+      await getChannelPromise(codeCoverageReportCh, { rootDir: coverageRootDir })
     }
 
     return exitOrClose.apply(this, arguments)
@@ -1213,16 +1206,11 @@ async function reportTypecheckFile (file, sessionConfiguration, frameworkVersion
     testSuiteErrorCh.runStores(testSuiteCtx, () => {})
   }
 
-  let onFinish
-  const onFinishPromise = new Promise(resolve => {
-    onFinish = resolve
-  })
-  testSuiteFinishCh.publish({
+  await getChannelPromise(testSuiteFinishCh, {
+    frameworkVersion,
     status: getTypecheckTaskStatus(file),
-    onFinish,
     ...testSuiteCtx.currentStore,
   })
-  await onFinishPromise
 }
 
 async function reportTypecheckResults (result, frameworkVersion, ctx, typechecker) {
@@ -1234,7 +1222,7 @@ async function reportTypecheckResults (result, frameworkVersion, ctx, typechecke
   }
   const providedContext = getMainProcessProvidedContext(ctx)
   const sessionConfiguration = testSessionConfigurationCh.hasSubscribers
-    ? await getChannelPromise(testSessionConfigurationCh, frameworkVersion)
+    ? await getChannelPromise(testSessionConfigurationCh, { frameworkVersion }) || {}
     : {}
 
   await Promise.all(result.files.map(file => reportTypecheckFile(