diff --git a/lib/plugin/aiTrace.js b/lib/plugin/aiTrace.js
index e0b114a34..8129ee9c9 100644
--- a/lib/plugin/aiTrace.js
+++ b/lib/plugin/aiTrace.js
@@ -127,7 +127,7 @@ export default function (config) {
     firstFailedStepSaved = false
   })
 
-  event.dispatcher.on(event.step.after, async step => {
+  event.dispatcher.on(event.step.after, step => {
     if (!currentTest) return
     if (step.status === 'failed') {
       testFailed = true
@@ -140,13 +140,10 @@ export default function (config) {
       output.debug(`aiTrace: Skipping failed step "${step.toString()}" - already handled by step.failed event`)
       return
     }
-    recorder.add(`save artifacts for step ${step.toString()}`, async () => {
-      try {
-        await persistStep(step)
-      } catch (err) {
-        output.debug(`aiTrace: Error saving step: ${err.message}`)
-      }
-    }, true)
+    const stepPersistPromise = persistStep(step).catch(err => {
+      output.debug(`aiTrace: Error saving step: ${err.message}`)
+    })
+    recorder.add(`wait aiTrace step persistence: ${step.toString()}`, () => stepPersistPromise, true)
   })
 
   event.dispatcher.on(event.step.failed, async step => {
@@ -284,7 +281,22 @@ export default function (config) {
         output.debug(`aiTrace: Browser unavailable, partial artifact capture: ${err.message}`)
       }
 
-      if (!step.artifacts?.screenshot) {
+      if (step.artifacts?.screenshot) {
+        const screenshotPath = path.isAbsolute(step.artifacts.screenshot)
+          ? step.artifacts.screenshot
+          : path.resolve(dir, step.artifacts.screenshot)
+        const screenshotFile = path.basename(screenshotPath)
+        stepData.artifacts.screenshot = screenshotFile
+        step.artifacts.screenshot = screenshotPath
+
+        if (!fs.existsSync(screenshotPath)) {
+          try {
+            await helper.saveScreenshot(screenshotPath, config.fullPageScreenshots)
+          } catch (err) {
+            output.debug(`aiTrace: Could not save screenshot: ${err.message}`)
+          }
+        }
+      } else {
         try {
           const screenshotFile = `${stepPrefix}_screenshot.png`
           const screenshotPath = path.join(dir, screenshotFile)