dd-trace 5.102.0 → 5.104.0

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

@@ -3,6 +3,7 @@
 // Capture real timers at module load time, before any test can install fake timers.
 const realSetTimeout = setTimeout
 
+const { performance } = require('node:perf_hooks')
 const satisfies = require('../../../vendor/dist/semifies')
 
 const shimmer = require('../../datadog-shimmer')
@@ -12,9 +13,12 @@ const {
   PLAYWRIGHT_WORKER_TRACE_PAYLOAD_CODE,
   getIsFaultyEarlyFlakeDetection,
   DYNAMIC_NAME_RE,
+  getEfdRetryCount,
+  getMaxEfdRetryCount,
   recordAttemptToFixExecution,
   logAttemptToFixTestExecution,
   logTestOptimizationSummary,
+  getTestOptimizationRequestResults,
 } = require('../../dd-trace/src/plugins/util/test')
 const log = require('../../dd-trace/src/log')
 const {
@@ -68,6 +72,7 @@ let remainingTestsByFile = {}
 let isKnownTestsEnabled = false
 let isEarlyFlakeDetectionEnabled = false
 let earlyFlakeDetectionNumRetries = 0
+let earlyFlakeDetectionSlowTestRetries = {}
 let isEarlyFlakeDetectionFaulty = false
 let earlyFlakeDetectionFaultyThreshold = 0
 let isFlakyTestRetriesEnabled = false
@@ -83,10 +88,19 @@ let testsReportedInGenerateSummary = new Set()
 const newTestsWithDynamicNames = new Set()
 const attemptToFixExecutions = new Map()
 const loggedAttemptToFixTests = new Set()
+const efdManagedTestKeys = new Set()
+const efdRetryCountByTestKey = new Map()
+const efdRetryCountRequestsByTestKey = new Map()
+const efdRetryTestsById = new Map()
+const efdScheduledOriginalTestKeys = new Set()
+const efdStartedOriginalTestKeys = new Set()
+const efdSlowAbortedTests = new Set()
 let rootDir = ''
 let sessionProjects = []
 
 const MINIMUM_SUPPORTED_VERSION_RANGE_EFD = '>=1.38.0' // TODO: remove this once we drop support for v5
+const EFD_RETRY_COUNT_REQUEST = 'ddEfdRetryCountRequest'
+const EFD_RETRY_COUNT_RESPONSE = 'ddEfdRetryCountResponse'
 
 function isValidKnownTests (receivedKnownTests) {
   return !!receivedKnownTests.playwright
@@ -97,6 +111,222 @@ function getTestFullyQualifiedName (test) {
   return `${test._requireFile} ${fullname}`
 }
 
+/**
+ * @param {object} test
+ * @returns {string|undefined}
+ */
+function getTestProjectKey (test) {
+  const { _projectIndex, _projectId } = test
+  if (_projectIndex !== undefined) {
+    return `index:${_projectIndex}`
+  }
+  if (_projectId !== undefined) {
+    return `id:${_projectId}`
+  }
+
+  const projectSuite = getSuiteType(test, 'project')
+  const projectName = projectSuite?._fullProject?.project?.name ||
+    projectSuite?._fullProject?.name ||
+    projectSuite?.title
+  if (projectName) {
+    return `name:${projectName}`
+  }
+}
+
+/**
+ * @param {object} test
+ * @returns {number|undefined}
+ */
+function getTestEfdRepeatEachIndex (test) {
+  if (Object.hasOwn(test, '_ddEfdOriginalRepeatEachIndex')) {
+    return test._ddEfdOriginalRepeatEachIndex
+  }
+  return test.repeatEachIndex
+}
+
+/**
+ * @param {object} test
+ * @returns {string|undefined}
+ */
+function getTestRepeatEachKey (test) {
+  const repeatEachIndex = getTestEfdRepeatEachIndex(test)
+  if (repeatEachIndex !== undefined) {
+    return `repeat:${repeatEachIndex}`
+  }
+}
+
+/**
+ * @param {object} test
+ * @returns {string}
+ */
+function getTestEfdKey (test) {
+  const projectKey = getTestProjectKey(test)
+  const repeatEachKey = getTestRepeatEachKey(test)
+  const testFqn = getTestFullyQualifiedName(test)
+  return [projectKey, repeatEachKey, testFqn].filter(Boolean).join(' ')
+}
+
+function getConfiguredEfdRetryCount () {
+  if (!earlyFlakeDetectionSlowTestRetries || !Object.keys(earlyFlakeDetectionSlowTestRetries).length) {
+    return earlyFlakeDetectionNumRetries
+  }
+  return getMaxEfdRetryCount(earlyFlakeDetectionSlowTestRetries)
+}
+
+function markEfdManagedTest (test) {
+  test._ddIsEfdManagedTest = true
+  test._ddEfdSlowTestRetries = earlyFlakeDetectionSlowTestRetries
+  efdManagedTestKeys.add(getTestEfdKey(test))
+}
+
+function markEfdRetryTest (test, retryIndex, originalTest) {
+  test._ddIsEfdRetry = true
+  test._ddEfdRetryIndex = retryIndex
+  if (originalTest) {
+    test._ddEfdOriginalRepeatEachIndex = getTestEfdRepeatEachIndex(originalTest)
+  }
+}
+
+function registerEfdRetryTest (test) {
+  if (!test._ddIsEfdRetry) {
+    return
+  }
+
+  efdRetryTestsById.set(test.id, {
+    retryIndex: test._ddEfdRetryIndex,
+    testEfdKey: getTestEfdKey(test),
+  })
+}
+
+function getTestEfdSlowTestRetries (test) {
+  return test._ddEfdSlowTestRetries || earlyFlakeDetectionSlowTestRetries
+}
+
+function isTestEfdManaged (test) {
+  return !!test._ddIsEfdManagedTest || (
+    (test._ddIsNew || test._ddIsModified) &&
+    !test._ddIsAttemptToFix &&
+    isEarlyFlakeDetectionEnabled
+  )
+}
+
+function getFileSuiteRepeatEachIndex (fileSuite) {
+  const test = fileSuite.allTests()[0]
+  return test ? getTestEfdRepeatEachIndex(test) || 0 : 0
+}
+
+function getEfdRetryRepeatEachIndex (fileSuite, projectSuite, retryIndex, retryCount) {
+  const nativeRepeatEach = projectSuite._fullProject?.project?.repeatEach || 1
+  const originalRepeatEachIndex = getFileSuiteRepeatEachIndex(fileSuite)
+  return nativeRepeatEach + (originalRepeatEachIndex * retryCount) + retryIndex - 1
+}
+
+function getEfdRetryCountForTest (test) {
+  return efdRetryCountByTestKey.get(getTestEfdKey(test)) ?? getConfiguredEfdRetryCount()
+}
+
+function setEfdRetryCountForTest (test, retryCount) {
+  const testEfdKey = getTestEfdKey(test)
+  efdRetryCountByTestKey.set(testEfdKey, retryCount)
+
+  const requests = efdRetryCountRequestsByTestKey.get(testEfdKey)
+  if (requests) {
+    efdRetryCountRequestsByTestKey.delete(testEfdKey)
+    for (const resolveRequest of requests) {
+      resolveRequest(retryCount)
+    }
+  }
+}
+
+function sendEfdRetryCountToWorker (workerProcess, testId, retryIndex, retryCount) {
+  workerProcess.send({
+    type: EFD_RETRY_COUNT_RESPONSE,
+    testId,
+    isEfdRetry: retryIndex !== undefined,
+    retryIndex,
+    retryCount,
+  })
+}
+
+function sendEfdRetryCountToWorkerWhenAvailable (workerProcess, testId) {
+  const efdRetryTest = efdRetryTestsById.get(testId)
+  if (!efdRetryTest) {
+    sendEfdRetryCountToWorker(workerProcess, testId)
+    return
+  }
+
+  const { retryIndex, testEfdKey } = efdRetryTest
+
+  if (!testEfdKey || !efdManagedTestKeys.has(testEfdKey)) {
+    sendEfdRetryCountToWorker(workerProcess, testId)
+    return
+  }
+
+  const retryCount = efdRetryCountByTestKey.get(testEfdKey)
+  if (retryCount !== undefined) {
+    sendEfdRetryCountToWorker(workerProcess, testId, retryIndex, retryCount)
+    return
+  }
+
+  if (!efdStartedOriginalTestKeys.has(testEfdKey) && !efdScheduledOriginalTestKeys.has(testEfdKey)) {
+    sendEfdRetryCountToWorker(workerProcess, testId, retryIndex, 0)
+    return
+  }
+
+  if (!efdRetryCountRequestsByTestKey.has(testEfdKey)) {
+    efdRetryCountRequestsByTestKey.set(testEfdKey, [])
+  }
+  efdRetryCountRequestsByTestKey.get(testEfdKey).push((retryCount) => {
+    sendEfdRetryCountToWorker(workerProcess, testId, retryIndex, retryCount)
+  })
+}
+
+/**
+ * @param {object} test
+ * @returns {boolean}
+ */
+function shouldRequestEfdRetryCount (test) {
+  // The main process remains the source of truth. repeatEachIndex is only used as
+  // a cheap worker-side filter so first executions do not block on coordination.
+  return test._ddIsEfdRetry || test.repeatEachIndex > 0
+}
+
+function waitForEfdRetryCount (test) {
+  if (!process.send || !shouldRequestEfdRetryCount(test)) {
+    return Promise.resolve()
+  }
+
+  const testEfdKey = getTestEfdKey(test)
+  return new Promise(resolve => {
+    const messageHandler = (message) => {
+      if (message?.type === EFD_RETRY_COUNT_RESPONSE && message.testId === test.id) {
+        if (message.isEfdRetry) {
+          test._ddIsEfdRetry = true
+          test._ddEfdRetryIndex = message.retryIndex
+          test._ddEfdRetryCount = message.retryCount
+          efdRetryCountByTestKey.set(testEfdKey, message.retryCount)
+        }
+        process.removeListener('message', messageHandler)
+        resolve()
+      }
+    }
+
+    process.on('message', messageHandler)
+    process.send({
+      type: EFD_RETRY_COUNT_REQUEST,
+      testId: test.id,
+    })
+  })
+}
+
+function shouldSkipEfdRetry (test) {
+  if (!test._ddIsEfdRetry) {
+    return false
+  }
+  const retryCount = test._ddEfdRetryCount ?? efdRetryCountByTestKey.get(getTestEfdKey(test))
+  return retryCount !== undefined && test._ddEfdRetryIndex > retryCount
+}
+
 function getTestProperties (test) {
   const testName = getTestFullname(test)
   const testSuite = getTestSuitePath(test._requireFile, rootDir)
@@ -125,14 +355,17 @@ function getSuiteType (test, type) {
 }
 
 // Copy of Suite#_deepClone but with a function to filter tests
-function deepCloneSuite (suite, filterTest, tags = []) {
+function deepCloneSuite (suite, filterTest, tags = [], configureCopiedTest) {
   const copy = suite._clone()
   for (const entry of suite._entries) {
     if (entry.constructor.name === 'Suite') {
-      copy._addSuite(deepCloneSuite(entry, filterTest, tags))
+      copy._addSuite(deepCloneSuite(entry, filterTest, tags, configureCopiedTest))
     } else {
       if (filterTest(entry)) {
         const copiedTest = entry._clone()
+        if (configureCopiedTest) {
+          configureCopiedTest(copiedTest, entry)
+        }
         for (const tag of tags) {
           const resolvedTag = typeof tag === 'function' ? tag(entry) : tag
 
@@ -303,6 +536,7 @@ function getFinalStatus ({
   isAttemptToFix,
   hasFailedAllRetries,
   hasFailedAttemptToFixRetries,
+  hasPassedAnyEfdAttempt,
   testStatus,
 }) {
   if (!isFinalExecution) {
@@ -311,9 +545,12 @@ function getFinalStatus ({
   if (isDisabled || isQuarantined || testStatus === 'skip') {
     return 'skip'
   }
-  if (isAtrRetry || isEfdManagedTest) {
+  if (isAtrRetry) {
     return hasFailedAllRetries ? 'fail' : 'pass'
   }
+  if (isEfdManagedTest) {
+    return hasPassedAnyEfdAttempt ? 'pass' : 'fail'
+  }
   if (isAttemptToFix) {
     return hasFailedAttemptToFixRetries ? 'fail' : 'pass'
   }
@@ -350,6 +587,14 @@ function testBeginHandler (test, browserName, shouldCreateTestSpan) {
   if (_type === 'beforeAll' || _type === 'afterAll') {
     return
   }
+  if (shouldSkipEfdRetry(test)) {
+    test._ddShouldSkipEfdRetry = true
+    return
+  }
+  test._ddStartTime = performance.now()
+  if (isTestEfdManaged(test) && !test._ddIsEfdRetry) {
+    efdStartedOriginalTestKeys.add(getTestEfdKey(test))
+  }
   // this means that a skipped test is being handled
   if (!remainingTestsByFile[testSuiteAbsolutePath].length) {
     return
@@ -391,6 +636,45 @@ function testBeginHandler (test, browserName, shouldCreateTestSpan) {
   }
 }
 
+function finishTestSuiteIfDone (testSuiteAbsolutePath, projects) {
+  if (!shouldFinishTestSuite(testSuiteAbsolutePath)) {
+    return
+  }
+
+  const skippedTests = remainingTestsByFile[testSuiteAbsolutePath]
+    .filter(test => test.expectedStatus === 'skipped')
+
+  for (const test of skippedTests) {
+    const browserName = getBrowserNameFromProjects(projects, test)
+    testSkipCh.publish({
+      testName: getTestFullname(test),
+      testSuiteAbsolutePath,
+      testSourceFileAbsolutePath: test.location.file,
+      testSourceLine: test.location.line,
+      browserName,
+      isNew: test._ddIsNew,
+      isDisabled: test._ddIsDisabled,
+      isModified: test._ddIsModified,
+      isQuarantined: test._ddIsQuarantined,
+    })
+  }
+  remainingTestsByFile[testSuiteAbsolutePath] = []
+
+  const testStatuses = testSuiteToTestStatuses.get(testSuiteAbsolutePath)
+  let testSuiteStatus = 'pass'
+  if (testStatuses?.includes('fail')) {
+    testSuiteStatus = 'fail'
+  } else if (testStatuses?.every(status => status === 'skip')) {
+    testSuiteStatus = 'skip'
+  }
+
+  const suiteError = getTestSuiteError(testSuiteAbsolutePath)
+  const testSuiteCtx = testSuiteToCtx.get(testSuiteAbsolutePath)
+  if (testSuiteCtx) {
+    testSuiteFinishCh.publish({ status: testSuiteStatus, error: suiteError, ...testSuiteCtx.currentStore })
+  }
+}
+
 function testEndHandler ({
   test,
   annotations,
@@ -420,11 +704,21 @@ function testEndHandler ({
     return
   }
 
+  if (test._ddShouldSkipEfdRetry || shouldSkipEfdRetry(test)) {
+    test._ddShouldSkipEfdRetry = true
+    remainingTestsByFile[testSuiteAbsolutePath] = remainingTestsByFile[testSuiteAbsolutePath]
+      .filter(currentTest => currentTest !== test)
+    finishTestSuiteIfDone(testSuiteAbsolutePath, projects)
+    return
+  }
+
+  const isEfdManagedTest = isTestEfdManaged(test)
   const testFqn = getTestFullyQualifiedName(test)
-  const testStatuses = testsToTestStatuses.get(testFqn) || []
+  const testStatusKey = isEfdManagedTest ? getTestEfdKey(test) : testFqn
+  const testStatuses = testsToTestStatuses.get(testStatusKey) || []
 
   if (testStatuses.length === 0) {
-    testsToTestStatuses.set(testFqn, [testStatus])
+    testsToTestStatuses.set(testStatusKey, [testStatus])
     if (test._ddIsNew && DYNAMIC_NAME_RE.test(getTestFullname(test))) {
       newTestsWithDynamicNames.add(`${getTestSuitePath(test._requireFile, rootDir)} › ${getTestFullname(test)}`)
     }
@@ -432,6 +726,17 @@ function testEndHandler ({
     testStatuses.push(testStatus)
   }
 
+  const testEfdKey = getTestEfdKey(test)
+  if (isEfdManagedTest && !test._ddIsEfdRetry && !efdRetryCountByTestKey.has(testEfdKey)) {
+    const testResult = results.at(-1)
+    const duration = testResult?.duration > 0 ? testResult.duration : performance.now() - test._ddStartTime
+    const retryCount = getEfdRetryCount(duration, getTestEfdSlowTestRetries(test))
+    setEfdRetryCountForTest(test, retryCount)
+    if (retryCount === 0) {
+      efdSlowAbortedTests.add(testEfdKey)
+    }
+  }
+
   const testProperties = getTestProperties(test)
 
   if (testProperties.attemptToFix) {
@@ -460,9 +765,10 @@ function testEndHandler ({
   }
 
   // Check if all EFD retries failed
-  if (testStatuses.length === earlyFlakeDetectionNumRetries + 1 &&
+  const efdRetryCount = getEfdRetryCountForTest(test)
+  if (efdRetryCount > 0 && testStatuses.length === efdRetryCount + 1 &&
     (test._ddIsNew || test._ddIsModified) &&
-    test._ddIsEfdRetry &&
+    isEarlyFlakeDetectionEnabled &&
     testStatuses.every(status => status === 'fail')) {
     test._ddHasFailedAllRetries = true
   }
@@ -480,9 +786,6 @@ function testEndHandler ({
   if (shouldCreateTestSpan) {
     const testResult = results.at(-1)
     const testCtx = testToCtx.get(test)
-    const isEfdManagedTest = (test._ddIsNew || test._ddIsModified) &&
-      !test._ddIsAttemptToFix &&
-      isEarlyFlakeDetectionEnabled
     const isAtrRetry = testResult?.retry > 0 &&
       isFlakyTestRetriesEnabled &&
       !test._ddIsAttemptToFix &&
@@ -497,6 +800,7 @@ function testEndHandler ({
       isAttemptToFix: test._ddIsAttemptToFix,
       hasFailedAllRetries: test._ddHasFailedAllRetries,
       hasFailedAttemptToFixRetries: test._ddHasFailedAttemptToFixRetries,
+      hasPassedAnyEfdAttempt: testStatuses.includes('pass'),
       testStatus,
     })
 
@@ -520,6 +824,7 @@ function testEndHandler ({
         isAtrRetry,
         isModified: test._ddIsModified,
         finalStatus,
+        earlyFlakeAbortReason: efdSlowAbortedTests.has(testEfdKey) ? 'slow' : undefined,
         ...testCtx.currentStore,
       })
     }
@@ -540,38 +845,7 @@ function testEndHandler ({
       .filter(currentTest => currentTest !== test)
   }
 
-  if (shouldFinishTestSuite(testSuiteAbsolutePath)) {
-    const skippedTests = remainingTestsByFile[testSuiteAbsolutePath]
-      .filter(test => test.expectedStatus === 'skipped')
-
-    for (const test of skippedTests) {
-      const browserName = getBrowserNameFromProjects(projects, test)
-      testSkipCh.publish({
-        testName: getTestFullname(test),
-        testSuiteAbsolutePath,
-        testSourceFileAbsolutePath: test.location.file,
-        testSourceLine: test.location.line,
-        browserName,
-        isNew: test._ddIsNew,
-        isDisabled: test._ddIsDisabled,
-        isModified: test._ddIsModified,
-        isQuarantined: test._ddIsQuarantined,
-      })
-    }
-    remainingTestsByFile[testSuiteAbsolutePath] = []
-
-    const testStatuses = testSuiteToTestStatuses.get(testSuiteAbsolutePath)
-    let testSuiteStatus = 'pass'
-    if (testStatuses.includes('fail')) {
-      testSuiteStatus = 'fail'
-    } else if (testStatuses.every(status => status === 'skip')) {
-      testSuiteStatus = 'skip'
-    }
-
-    const suiteError = getTestSuiteError(testSuiteAbsolutePath)
-    const testSuiteCtx = testSuiteToCtx.get(testSuiteAbsolutePath)
-    testSuiteFinishCh.publish({ status: testSuiteStatus, error: suiteError, ...testSuiteCtx.currentStore })
-  }
+  finishTestSuiteIfDone(testSuiteAbsolutePath, projects)
 }
 
 function dispatcherRunWrapper (run) {
@@ -581,6 +855,39 @@ function dispatcherRunWrapper (run) {
   }
 }
 
+function deferEfdRetryGroups (testGroups) {
+  const groupsWithOriginalTests = []
+  const efdRetryOnlyGroups = []
+
+  for (const group of testGroups) {
+    const originalTests = []
+    const efdRetryTests = []
+
+    for (const test of group.tests) {
+      if (test._ddIsEfdRetry) {
+        efdRetryTests.push(test)
+      } else {
+        originalTests.push(test)
+        if (isTestEfdManaged(test)) {
+          efdScheduledOriginalTestKeys.add(getTestEfdKey(test))
+        }
+      }
+    }
+
+    if (efdRetryTests.length && originalTests.length) {
+      group.tests = [...originalTests, ...efdRetryTests]
+    }
+
+    if (originalTests.length) {
+      groupsWithOriginalTests.push(group)
+    } else {
+      efdRetryOnlyGroups.push(group)
+    }
+  }
+
+  return [...groupsWithOriginalTests, ...efdRetryOnlyGroups]
+}
+
 function dispatcherRunWrapperNew (run) {
   return function (testGroups) {
     // Filter out disabled tests from testGroups before they get scheduled,
@@ -593,12 +900,17 @@ function dispatcherRunWrapperNew (run) {
       testGroups = testGroups.filter(group => group.tests.length > 0)
     }
 
+    if (isEarlyFlakeDetectionEnabled) {
+      testGroups = deferEfdRetryGroups(testGroups)
+    }
+
     if (!this._allTests) {
       // Removed in https://github.com/microsoft/playwright/commit/1e52c37b254a441cccf332520f60225a5acc14c7
       // Not available from >=1.44.0
       this._ddAllTests = testGroups.flatMap(g => g.tests)
     }
     remainingTestsByFile = getTestsBySuiteFromTestGroups(testGroups)
+    arguments[0] = testGroups
     return run.apply(this, arguments)
   }
 }
@@ -638,7 +950,6 @@ function dispatcherHook (dispatcherExport) {
         )
       }
     })
-
     return worker
   })
   return dispatcherExport
@@ -690,14 +1001,11 @@ function dispatcherHookNew (dispatcherExport, runWrapper) {
       // above) and mark the execution final once the count reaches the expected total.
       // This mirrors how ATF finality is detected and centralizes the decision in the
       // main process, so workers only need to act on the _ddIsFinalExecution flag.
-      const isEfdManagedTest = (test._ddIsNew || test._ddIsModified) &&
-        !test._ddIsAttemptToFix &&
-        isEarlyFlakeDetectionEnabled
+      const isEfdManagedTest = isTestEfdManaged(test)
       let isFinalExecution
       if (isEfdManagedTest) {
-        const testFqn = getTestFullyQualifiedName(test)
-        const efdTestStatuses = testsToTestStatuses.get(testFqn) || []
-        isFinalExecution = efdTestStatuses.length === earlyFlakeDetectionNumRetries + 1
+        const efdTestStatuses = testsToTestStatuses.get(getTestEfdKey(test)) || []
+        isFinalExecution = efdTestStatuses.length === getEfdRetryCountForTest(test) + 1
       } else if (test._ddIsAttemptToFix) {
         isFinalExecution = !!(test._ddHasPassedAttemptToFixRetries || test._ddHasFailedAttemptToFixRetries)
       } else {
@@ -722,10 +1030,11 @@ function dispatcherHookNew (dispatcherExport, runWrapper) {
           _ddIsModified: test._ddIsModified,
           _ddIsFinalExecution: isFinalExecution,
           _ddIsEfdManagedTest: isEfdManagedTest,
+          _ddEarlyFlakeAbortReason: efdSlowAbortedTests.has(getTestEfdKey(test)) ? 'slow' : undefined,
+          _ddHasPassedAnyEfdAttempt: (testsToTestStatuses.get(getTestEfdKey(test)) || []).includes('pass'),
         },
       })
     })
-
     return worker
   })
   return dispatcherExport
@@ -751,6 +1060,7 @@ function runAllTestsWrapper (runAllTests, playwrightVersion) {
         isKnownTestsEnabled = libraryConfig.isKnownTestsEnabled
         isEarlyFlakeDetectionEnabled = libraryConfig.isEarlyFlakeDetectionEnabled
         earlyFlakeDetectionNumRetries = libraryConfig.earlyFlakeDetectionNumRetries
+        earlyFlakeDetectionSlowTestRetries = libraryConfig.earlyFlakeDetectionSlowTestRetries ?? {}
         earlyFlakeDetectionFaultyThreshold = libraryConfig.earlyFlakeDetectionFaultyThreshold
         isFlakyTestRetriesEnabled = libraryConfig.isFlakyTestRetriesEnabled
         flakyTestRetriesCount = libraryConfig.flakyTestRetriesCount
@@ -766,9 +1076,24 @@ function runAllTestsWrapper (runAllTests, playwrightVersion) {
       log.error('Playwright session start error', e)
     }
 
-    if (isKnownTestsEnabled && satisfies(playwrightVersion, MINIMUM_SUPPORTED_VERSION_RANGE_EFD)) {
+    const isTestOptimizationSupported = satisfies(playwrightVersion, MINIMUM_SUPPORTED_VERSION_RANGE_EFD)
+    const shouldGetKnownTests = isKnownTestsEnabled && isTestOptimizationSupported
+    const shouldGetTestManagementTests = isTestManagementTestsEnabled && isTestOptimizationSupported
+
+    const {
+      knownTestsResponse,
+      testManagementTestsResponse,
+    } = await getTestOptimizationRequestResults({
+      isKnownTestsEnabled: shouldGetKnownTests,
+      isTestManagementTestsEnabled: shouldGetTestManagementTests,
+      getKnownTests: () => getChannelPromise(knownTestsCh),
+      getTestManagementTests: () => getChannelPromise(testManagementTestsCh),
+    })
+
+    if (shouldGetKnownTests) {
       try {
-        const { err, knownTests: receivedKnownTests } = await getChannelPromise(knownTestsCh)
+        const { err, knownTests: receivedKnownTests } =
+          knownTestsResponse || await getChannelPromise(knownTestsCh)
         if (err) {
           isEarlyFlakeDetectionEnabled = false
           isKnownTestsEnabled = false
@@ -787,9 +1112,10 @@ function runAllTestsWrapper (runAllTests, playwrightVersion) {
       }
     }
 
-    if (isTestManagementTestsEnabled && satisfies(playwrightVersion, MINIMUM_SUPPORTED_VERSION_RANGE_EFD)) {
+    if (shouldGetTestManagementTests) {
       try {
-        const { err, testManagementTests: receivedTestManagementTests } = await getChannelPromise(testManagementTestsCh)
+        const { err, testManagementTests: receivedTestManagementTests } =
+          testManagementTestsResponse || await getChannelPromise(testManagementTestsCh)
         if (err) {
           isTestManagementTestsEnabled = false
         } else {
@@ -801,7 +1127,7 @@ function runAllTestsWrapper (runAllTests, playwrightVersion) {
       }
     }
 
-    if (isImpactedTestsEnabled && satisfies(playwrightVersion, MINIMUM_SUPPORTED_VERSION_RANGE_EFD)) {
+    if (isImpactedTestsEnabled && isTestOptimizationSupported) {
       try {
         const { err, modifiedFiles: receivedModifiedFiles } = await getChannelPromise(modifiedFilesCh)
         if (err) {
@@ -900,6 +1226,13 @@ function runAllTestsWrapper (runAllTests, playwrightVersion) {
     remainingTestsByFile = {}
     quarantinedButNotAttemptToFixFqns = new Set()
     testsReportedInGenerateSummary = new Set()
+    efdManagedTestKeys.clear()
+    efdRetryCountByTestKey.clear()
+    efdRetryCountRequestsByTestKey.clear()
+    efdRetryTestsById.clear()
+    efdScheduledOriginalTestKeys.clear()
+    efdStartedOriginalTestKeys.clear()
+    efdSlowAbortedTests.clear()
 
     // TODO: we can trick playwright into thinking the session passed by returning
     // 'passed' here. We might be able to use this for both EFD and Test Management tests.
@@ -1001,12 +1334,29 @@ addHook({
  * - we execute `applyRepeatEachIndex` for each of these cloned file suites
  * - we add the cloned file suites to the project suite
  */
-function applyRetriesToTests (fileSuitesWithTestsToRetry, filterTest, tagsToApply, numRetries) {
+function applyRetriesToTests (
+  fileSuitesWithTestsToRetry,
+  filterTest,
+  tagsToApply,
+  numRetries,
+  configureCopiedTest,
+  getRetryRepeatEachIndex
+) {
   for (const [fileSuite, projectSuite] of fileSuitesWithTestsToRetry.entries()) {
     for (let repeatEachIndex = 1; repeatEachIndex <= numRetries; repeatEachIndex++) {
-      const copyFileSuite = deepCloneSuite(fileSuite, filterTest, tagsToApply)
-      applyRepeatEachIndex(projectSuite._fullProject, copyFileSuite, repeatEachIndex + 1)
+      const copyFileSuite = deepCloneSuite(fileSuite, filterTest, tagsToApply, (copiedTest, originalTest) => {
+        if (configureCopiedTest) {
+          configureCopiedTest(copiedTest, originalTest, repeatEachIndex)
+        }
+      })
+      const retryRepeatEachIndex = getRetryRepeatEachIndex
+        ? getRetryRepeatEachIndex(fileSuite, projectSuite, repeatEachIndex, numRetries)
+        : repeatEachIndex + 1
+      applyRepeatEachIndex(projectSuite._fullProject, copyFileSuite, retryRepeatEachIndex)
       projectSuite._addSuite(copyFileSuite)
+      for (const copiedTest of copyFileSuite.allTests()) {
+        registerEfdRetryTest(copiedTest)
+      }
     }
   }
 }
@@ -1091,6 +1441,7 @@ addHook({
       for (const impactedTest of impactedTests) {
         impactedTest._ddIsModified = true
         if (isEarlyFlakeDetectionEnabled && impactedTest.expectedStatus !== 'skipped') {
+          markEfdManagedTest(impactedTest)
           const fileSuite = getSuiteType(impactedTest, 'file')
           if (!fileSuitesWithImpactedTestsToProjects.has(fileSuite)) {
             fileSuitesWithImpactedTestsToProjects.set(fileSuite, getSuiteType(impactedTest, 'project'))
@@ -1106,7 +1457,12 @@ addHook({
           '_ddIsEfdRetry',
           (test) => (isKnownTestsEnabled && isNewTest(test) ? '_ddIsNew' : null),
         ],
-        earlyFlakeDetectionNumRetries
+        getConfiguredEfdRetryCount(),
+        (copiedTest, originalTest, retryIndex) => {
+          markEfdRetryTest(copiedTest, retryIndex, originalTest)
+          markEfdManagedTest(copiedTest)
+        },
+        getEfdRetryRepeatEachIndex
       )
     }
 
@@ -1130,6 +1486,7 @@ addHook({
           if (isEarlyFlakeDetectionEnabled && newTest.expectedStatus !== 'skipped' && !newTest._ddIsModified) {
             // Prevent ATR or `--retries` from retrying new tests if EFD is enabled
             newTest.retries = 0
+            markEfdManagedTest(newTest)
             const fileSuite = getSuiteType(newTest, 'file')
             if (!fileSuitesWithNewTestsToProjects.has(fileSuite)) {
               fileSuitesWithNewTestsToProjects.set(fileSuite, getSuiteType(newTest, 'project'))
@@ -1141,7 +1498,12 @@ addHook({
           fileSuitesWithNewTestsToProjects,
           isNewTest,
           ['_ddIsNew', '_ddIsEfdRetry'],
-          earlyFlakeDetectionNumRetries
+          getConfiguredEfdRetryCount(),
+          (copiedTest, originalTest, retryIndex) => {
+            markEfdRetryTest(copiedTest, retryIndex, originalTest)
+            markEfdManagedTest(copiedTest)
+          },
+          getEfdRetryRepeatEachIndex
         )
       }
     }
@@ -1177,6 +1539,10 @@ addHook({
 
     // We add a new listener to `this.process`, which is represents the worker
     this.process.on('message', (message) => {
+      if (message?.type === EFD_RETRY_COUNT_REQUEST) {
+        sendEfdRetryCountToWorkerWhenAvailable(this.process, message.testId)
+        return
+      }
       // These messages are [code, payload]. The payload is test data
       if (Array.isArray(message) && message[0] === PLAYWRIGHT_WORKER_TRACE_PAYLOAD_CODE) {
         workerReportCh.publish(message[1])
@@ -1235,9 +1601,15 @@ addHook({
   const stepInfoByStepId = {}
 
   shimmer.wrap(workerPackage.WorkerMain.prototype, '_runTest', _runTest => async function (test) {
+    await waitForEfdRetryCount(test)
+    if (shouldSkipEfdRetry(test)) {
+      test._ddShouldSkipEfdRetry = true
+      test.expectedStatus = 'skipped'
+    }
     if (test.expectedStatus === 'skipped') {
       return _runTest.apply(this, arguments)
     }
+    test._ddStartTime = performance.now()
     steps = []
 
     const {
@@ -1320,6 +1692,21 @@ addHook({
     await res
 
     const { status, error, annotations, retry, testId } = testInfo
+    const testEfdKey = getTestEfdKey(test)
+    const isEfdManagedTest = isTestEfdManaged(test)
+    if (isEfdManagedTest && !test._ddIsEfdRetry && !efdRetryCountByTestKey.has(testEfdKey)) {
+      const duration = test.results?.at(-1)?.duration > 0
+        ? test.results.at(-1).duration
+        : performance.now() - test._ddStartTime
+      const retryCount = getEfdRetryCount(
+        duration,
+        getTestEfdSlowTestRetries(test)
+      )
+      setEfdRetryCountForTest(test, retryCount)
+      if (retryCount === 0) {
+        efdSlowAbortedTests.add(testEfdKey)
+      }
+    }
 
     // testInfo.errors could be better than "error",
     // which will only include timeout error (even though the test failed because of a different error)
@@ -1365,6 +1752,7 @@ addHook({
       isAttemptToFix: test._ddIsAttemptToFix,
       hasFailedAllRetries: test._ddHasFailedAllRetries,
       hasFailedAttemptToFixRetries: test._ddHasFailedAttemptToFixRetries,
+      hasPassedAnyEfdAttempt: test._ddHasPassedAnyEfdAttempt,
       testStatus: STATUS_TO_TEST_STATUS[status],
     })
 
@@ -1388,6 +1776,7 @@ addHook({
       isModified: test._ddIsModified,
       onDone,
       finalStatus,
+      earlyFlakeAbortReason: test._ddEarlyFlakeAbortReason,
       ...testCtx.currentStore,
     })