diff --git a/maestro-cli/src/main/java/maestro/cli/report/TestDebugReporter.kt b/maestro-cli/src/main/java/maestro/cli/report/TestDebugReporter.kt
index e16a065fc1..35f6352533 100644
--- a/maestro-cli/src/main/java/maestro/cli/report/TestDebugReporter.kt
+++ b/maestro-cli/src/main/java/maestro/cli/report/TestDebugReporter.kt
@@ -2,6 +2,8 @@ package maestro.cli.report
 
 import com.fasterxml.jackson.annotation.JsonInclude
 import com.fasterxml.jackson.annotation.JsonProperty
+import com.fasterxml.jackson.core.util.DefaultIndenter
+import com.fasterxml.jackson.core.util.DefaultPrettyPrinter
 import com.fasterxml.jackson.databind.JsonMappingException
 import com.fasterxml.jackson.databind.ObjectMapper
 import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
@@ -35,16 +37,17 @@ object TestDebugReporter {
 
     private val logger = LoggerFactory.getLogger(TestDebugReporter::class.java)
     private val mapper = jacksonObjectMapper()
-                .setSerializationInclusion(JsonInclude.Include.NON_NULL)
-                .setSerializationInclusion(JsonInclude.Include.NON_EMPTY)
-                .writerWithDefaultPrettyPrinter()
+        .setSerializationInclusion(JsonInclude.Include.NON_NULL)
+        .setSerializationInclusion(JsonInclude.Include.NON_EMPTY)
+        .writerWithDefaultPrettyPrinter()
 
     private var debugOutputPath: Path? = null
     private var debugOutputPathAsString: String? = null
     private var flattenDebugOutput: Boolean = false
 
     fun saveFlow(flowName: String, debugOutput: FlowDebugOutput, aiOutput: FlowAIOutput?, path: Path) {
-        // TODO(bartekpacia): Potentially accept a single "FlowPersistentOutput object
+        // TODO(bartekpacia): Potentially accept a single "FlowPersistentOutput" object
+        // TODO(bartekpacia: Build output incrementally, instead of single-shot on flow completion
 
         println("TestDebugReporter.saveFlow: saving flow metadata to $path")
 
@@ -75,18 +78,22 @@ object TestDebugReporter {
             it.screenshot.copyTo(file)
         }
 
-        // AI screenshots
-        aiOutput?.outputs?.forEach { output ->
-            val screenshotFilename = output.screenshotPath.name
-            val screenshotFile = File(path.absolutePathString(), screenshotFilename)
-            output.screenshotPath.copyTo(screenshotFile)
-        }
+        aiOutput?.run {
+            // Write AI screenshots. Paths need to be changed to the final ones.
+            val updatedOutputs = outputs.map { output ->
+                val screenshotFilename = output.screenshotPath.name
+                val screenshotFile = File(path.absolutePathString(), screenshotFilename)
+                output.screenshotPath.copyTo(screenshotFile)
+                output.copy(screenshotPath = screenshotFile)
+            }
+
+            outputs.clear()
+            outputs.addAll(updatedOutputs)
 
-        // AI JSON output
-        aiOutput?.let { flowAiOutput ->
+            // Write AI JSON output
             val filename = "ai-(${flowName.replace("/", "_")}).json"
             val file = File(path.absolutePathString(), filename)
-            mapper.writeValue(file, flowAiOutput)
+            mapper.writeValue(file, this)
         }
     }
 
@@ -140,8 +147,10 @@ object TestDebugReporter {
     fun getDebugOutputPath(): Path {
         if (debugOutputPath != null) return debugOutputPath as Path
 
-        val debugRootPath = if(debugOutputPathAsString != null) debugOutputPathAsString!! else System.getProperty("user.home")
-        val debugOutput = if(flattenDebugOutput) Paths.get(debugRootPath) else buildDefaultDebugOutputPath(debugRootPath)
+        val debugRootPath =
+            if (debugOutputPathAsString != null) debugOutputPathAsString!! else System.getProperty("user.home")
+        val debugOutput =
+            if (flattenDebugOutput) Paths.get(debugRootPath) else buildDefaultDebugOutputPath(debugRootPath)
 
         if (!debugOutput.exists()) {
             Files.createDirectories(debugOutput)
diff --git a/maestro-orchestra/src/main/java/maestro/orchestra/Orchestra.kt b/maestro-orchestra/src/main/java/maestro/orchestra/Orchestra.kt
index 9ebb2b453b..1dcff74381 100644
--- a/maestro-orchestra/src/main/java/maestro/orchestra/Orchestra.kt
+++ b/maestro-orchestra/src/main/java/maestro/orchestra/Orchestra.kt
@@ -362,12 +362,11 @@ class Orchestra(
 
         val imageData = Buffer()
         maestro.takeScreenshot(imageData, compressed = false)
-        val imageDataBytes = imageData.readByteArray()
 
         val defects = Prediction.findDefects(
             aiClient = ai,
             assertion = null,
-            screen = imageDataBytes,
+            screen = imageData.copy().readByteArray(),
             previousFalsePositives = listOf(), // TODO: take it from WorkspaceConfig (or MaestroConfig?)
         )