diff --git a/diktat-api/src/main/kotlin/com/saveourtool/diktat/DiktatRunner.kt b/diktat-api/src/main/kotlin/com/saveourtool/diktat/DiktatRunner.kt
index 8285dd58b4..b3f5d6181c 100644
--- a/diktat-api/src/main/kotlin/com/saveourtool/diktat/DiktatRunner.kt
+++ b/diktat-api/src/main/kotlin/com/saveourtool/diktat/DiktatRunner.kt
@@ -19,14 +19,12 @@ private typealias RunAction = (DiktatProcessor, DiktatProcessorListener) -> Unit
  * @property diktatBaseline
  * @property diktatBaselineGenerator
  * @property diktatReporter
- * @property diktatReporterCloser
  */
 data class DiktatRunner(
     val diktatProcessor: DiktatProcessor,
     val diktatBaseline: DiktatBaseline,
     private val diktatBaselineGenerator: DiktatProcessorListener,
     val diktatReporter: DiktatReporter,
-    private val diktatReporterCloser: DiktatProcessorListener,
 ) {
     private fun doRun(
         args: DiktatRunnerArguments,
@@ -38,7 +36,6 @@ data class DiktatRunner(
             DiktatProcessorListener(
                 args.loggingListener,
                 diktatReporter.skipKnownErrors(diktatBaseline),
-                diktatReporterCloser,
                 diktatBaselineGenerator,
                 errorCounter.countErrorsAsProcessorListener()
             ),
@@ -47,7 +44,7 @@ data class DiktatRunner(
     }
 
     /**
-     * Run `diktat fix` for all [files].
+     * Run `diktat fix` for all [DiktatRunnerArguments.files].
      *
      * @param args
      * @param fileUpdateNotifier notifier about updated files
@@ -74,7 +71,7 @@ data class DiktatRunner(
     }
 
     /**
-     * Run `diktat check` for all [files].
+     * Run `diktat check` for all [DiktatRunnerArguments.files].
      *
      * @param args
      * @return count of detected errors
diff --git a/diktat-api/src/main/kotlin/com/saveourtool/diktat/DiktatRunnerFactory.kt b/diktat-api/src/main/kotlin/com/saveourtool/diktat/DiktatRunnerFactory.kt
index a5fc58002c..81e0d3cd0c 100644
--- a/diktat-api/src/main/kotlin/com/saveourtool/diktat/DiktatRunnerFactory.kt
+++ b/diktat-api/src/main/kotlin/com/saveourtool/diktat/DiktatRunnerFactory.kt
@@ -3,7 +3,6 @@ package com.saveourtool.diktat
 import com.saveourtool.diktat.api.DiktatBaseline
 import com.saveourtool.diktat.api.DiktatBaselineFactory
 import com.saveourtool.diktat.api.DiktatProcessorListener
-import com.saveourtool.diktat.api.DiktatProcessorListener.Companion.closeAfterAllAsProcessorListener
 import com.saveourtool.diktat.api.DiktatReporter
 import com.saveourtool.diktat.api.DiktatReporterFactory
 import com.saveourtool.diktat.api.DiktatRuleConfigReader
@@ -32,7 +31,7 @@ class DiktatRunnerFactory(
         val diktatRuleSet = diktatRuleSetFactory(diktatRuleConfigs)
         val processor = diktatProcessorFactory(diktatRuleSet)
         val (baseline, baselineGenerator) = resolveBaseline(args.baselineFile, args.sourceRootDir)
-        val (reporter, closer) = resolveReporter(
+        val reporter = resolveReporter(
             args.reporterType, args.reporterOutput,
             args.colorNameInPlain, args.groupByFileInPlain,
             args.sourceRootDir
@@ -42,7 +41,6 @@ class DiktatRunnerFactory(
             diktatBaseline = baseline,
             diktatBaselineGenerator = baselineGenerator,
             diktatReporter = reporter,
-            diktatReporterCloser = closer,
         )
     }
 
@@ -65,14 +63,10 @@ class DiktatRunnerFactory(
         colorNameInPlain: String?,
         groupByFileInPlain: Boolean?,
         sourceRootDir: Path,
-    ): Pair<DiktatReporter, DiktatProcessorListener> {
-        val (outputStream, closeListener) = reporterOutput
-            ?.let { it to it.closeAfterAllAsProcessorListener() }
-            ?: run {
-                System.`out` to DiktatProcessorListener.empty
-            }
-        val actualReporter = if (reporterType == diktatReporterFactory.plainId) {
-            diktatReporterFactory.createPlain(outputStream, sourceRootDir, colorNameInPlain, groupByFileInPlain)
+    ): DiktatReporter {
+        val (outputStream, closeOutputStream) = reporterOutput?.let { it to true } ?: (System.`out` to false)
+        return if (reporterType == diktatReporterFactory.plainId) {
+            diktatReporterFactory.createPlain(outputStream, closeOutputStream, sourceRootDir, colorNameInPlain, groupByFileInPlain)
         } else {
             require(colorNameInPlain == null) {
                 "colorization is applicable only for plain reporter"
@@ -80,8 +74,7 @@ class DiktatRunnerFactory(
             require(groupByFileInPlain == null) {
                 "groupByFile is applicable only for plain reporter"
             }
-            diktatReporterFactory(reporterType, outputStream, sourceRootDir)
+            diktatReporterFactory(reporterType, outputStream, closeOutputStream, sourceRootDir)
         }
-        return actualReporter to closeListener
     }
 }
diff --git a/diktat-api/src/main/kotlin/com/saveourtool/diktat/api/DiktatProcessorListener.kt b/diktat-api/src/main/kotlin/com/saveourtool/diktat/api/DiktatProcessorListener.kt
index b8497b0090..5da4aa2404 100644
--- a/diktat-api/src/main/kotlin/com/saveourtool/diktat/api/DiktatProcessorListener.kt
+++ b/diktat-api/src/main/kotlin/com/saveourtool/diktat/api/DiktatProcessorListener.kt
@@ -1,6 +1,5 @@
 package com.saveourtool.diktat.api
 
-import java.io.OutputStream
 import java.nio.file.Path
 import java.util.concurrent.atomic.AtomicInteger
 
@@ -82,15 +81,5 @@ interface DiktatProcessorListener {
                 incrementAndGet()
             }
         }
-
-        /**
-         * @return An implementation of [DiktatProcessorListener] which closes [OutputStream] at the end
-         */
-        fun OutputStream.closeAfterAllAsProcessorListener(): DiktatProcessorListener = object : DiktatProcessorListener {
-            override fun afterAll() {
-                this@closeAfterAllAsProcessorListener.flush()
-                this@closeAfterAllAsProcessorListener.close()
-            }
-        }
     }
 }
diff --git a/diktat-api/src/main/kotlin/com/saveourtool/diktat/api/DiktatReporterFactory.kt b/diktat-api/src/main/kotlin/com/saveourtool/diktat/api/DiktatReporterFactory.kt
index 51438c9c8c..04794ab897 100644
--- a/diktat-api/src/main/kotlin/com/saveourtool/diktat/api/DiktatReporterFactory.kt
+++ b/diktat-api/src/main/kotlin/com/saveourtool/diktat/api/DiktatReporterFactory.kt
@@ -8,7 +8,7 @@ typealias DiktatReporter = DiktatProcessorListener
 /**
  * A factory to create [DiktatReporter]
  */
-interface DiktatReporterFactory : Function3<String, OutputStream, Path, DiktatReporter> {
+interface DiktatReporterFactory : Function4<String, OutputStream, Boolean, Path, DiktatReporter> {
     /**
      * Set of supported IDs
      */
@@ -27,17 +27,20 @@ interface DiktatReporterFactory : Function3<String, OutputStream, Path, DiktatRe
     /**
      * @param id ID of [DiktatReporter]
      * @param outputStream
+     * @param closeOutputStreamAfterAll close [outputStream] in [DiktatProcessorListener.afterAll]
      * @param sourceRootDir a dir to detect relative path for processing files
      * @return created [DiktatReporter]
      */
     override operator fun invoke(
         id: String,
         outputStream: OutputStream,
+        closeOutputStreamAfterAll: Boolean,
         sourceRootDir: Path,
     ): DiktatReporter
 
     /**
      * @param outputStream
+     * @param closeOutputStreamAfterAll close [outputStream] in [DiktatProcessorListener.afterAll]
      * @param sourceRootDir a dir to detect relative path for processing files
      * @param colorName name of color for colorful output, `null` means to disable colorization.
      * @param groupByFile
@@ -45,6 +48,7 @@ interface DiktatReporterFactory : Function3<String, OutputStream, Path, DiktatRe
      */
     fun createPlain(
         outputStream: OutputStream,
+        closeOutputStreamAfterAll: Boolean,
         sourceRootDir: Path,
         colorName: String? = null,
         groupByFile: Boolean? = null,
diff --git a/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/DiktatBaselineFactoryImpl.kt b/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/DiktatBaselineFactoryImpl.kt
index 75cfff3781..b71778964f 100644
--- a/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/DiktatBaselineFactoryImpl.kt
+++ b/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/DiktatBaselineFactoryImpl.kt
@@ -3,14 +3,12 @@ package com.saveourtool.diktat.ktlint
 import com.saveourtool.diktat.api.DiktatBaseline
 import com.saveourtool.diktat.api.DiktatBaselineFactory
 import com.saveourtool.diktat.api.DiktatProcessorListener
-import com.saveourtool.diktat.api.DiktatProcessorListener.Companion.closeAfterAllAsProcessorListener
 import com.saveourtool.diktat.ktlint.DiktatReporterImpl.Companion.wrap
 
 import com.pinterest.ktlint.cli.reporter.baseline.Baseline
-import com.pinterest.ktlint.cli.reporter.baseline.BaselineReporter
+import com.pinterest.ktlint.cli.reporter.baseline.BaselineReporterProvider
 import com.pinterest.ktlint.cli.reporter.baseline.loadBaseline
 
-import java.io.PrintStream
 import java.nio.file.Path
 
 import kotlin.io.path.absolutePathString
@@ -20,6 +18,8 @@ import kotlin.io.path.outputStream
  * A factory to create or generate [DiktatBaseline] using `KtLint`
  */
 class DiktatBaselineFactoryImpl : DiktatBaselineFactory {
+    private val baselineReporterProvider = BaselineReporterProvider()
+
     override fun tryToLoad(
         baselineFile: Path,
         sourceRootDir: Path,
@@ -34,11 +34,6 @@ class DiktatBaselineFactoryImpl : DiktatBaselineFactory {
             }
         }
 
-    override fun generator(baselineFile: Path, sourceRootDir: Path): DiktatProcessorListener {
-        val outputStream = baselineFile.outputStream()
-        return DiktatProcessorListener(
-            BaselineReporter(PrintStream(outputStream)).wrap(sourceRootDir),
-            outputStream.closeAfterAllAsProcessorListener()
-        )
-    }
+    override fun generator(baselineFile: Path, sourceRootDir: Path): DiktatProcessorListener =
+        baselineReporterProvider.get(baselineFile.outputStream(), closeOutAfterAll = true, emptyMap()).wrap(sourceRootDir)
 }
diff --git a/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/DiktatReporterFactoryImpl.kt b/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/DiktatReporterFactoryImpl.kt
index 2469fdbd51..40a1f0842a 100644
--- a/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/DiktatReporterFactoryImpl.kt
+++ b/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/DiktatReporterFactoryImpl.kt
@@ -10,7 +10,6 @@ import com.pinterest.ktlint.cli.reporter.plain.Color
 import com.pinterest.ktlint.cli.reporter.plain.PlainReporterProvider
 import com.pinterest.ktlint.cli.reporter.sarif.SarifReporterProvider
 import java.io.OutputStream
-import java.io.PrintStream
 import java.nio.file.Path
 import kotlin.io.path.pathString
 
@@ -44,6 +43,7 @@ class DiktatReporterFactoryImpl : DiktatReporterFactory {
     override fun invoke(
         id: String,
         outputStream: OutputStream,
+        closeOutputStreamAfterAll: Boolean,
         sourceRootDir: Path,
     ): DiktatReporter {
         val reporterProvider = reporterProviders[id] ?: throw IllegalArgumentException("Not supported reporter id by ${DiktatBaselineFactoryImpl::class.simpleName}")
@@ -55,11 +55,12 @@ class DiktatReporterFactoryImpl : DiktatReporterFactory {
         } else {
             emptyMap()
         }
-        return reporterProvider.get(outputStream.asPrintStream(), opt).wrap(sourceRootDir)
+        return reporterProvider.get(outputStream, closeOutputStreamAfterAll, opt).wrap(sourceRootDir)
     }
 
     override fun createPlain(
         outputStream: OutputStream,
+        closeOutputStreamAfterAll: Boolean,
         sourceRootDir: Path,
         colorName: String?,
         groupByFile: Boolean?,
@@ -74,10 +75,6 @@ class DiktatReporterFactoryImpl : DiktatReporterFactory {
             }
             groupByFile?.let { put("group_by_file", it) }
         }.mapValues { it.value.toString() }
-        return plainReporterProvider.get(outputStream.asPrintStream(), opt).wrap(sourceRootDir)
-    }
-
-    companion object {
-        private fun OutputStream.asPrintStream(): PrintStream = (this as? PrintStream) ?: PrintStream(this)
+        return plainReporterProvider.get(outputStream, closeOutputStreamAfterAll, opt).wrap(sourceRootDir)
     }
 }
diff --git a/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/DiktatReporterImpl.kt b/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/DiktatReporterImpl.kt
index e891266f19..dcca2964ef 100644
--- a/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/DiktatReporterImpl.kt
+++ b/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/DiktatReporterImpl.kt
@@ -2,6 +2,7 @@ package com.saveourtool.diktat.ktlint
 
 import com.saveourtool.diktat.api.DiktatError
 import com.saveourtool.diktat.api.DiktatReporter
+import com.saveourtool.diktat.ktlint.ReporterV2Wrapper.Companion.unwrapIfNeeded
 import com.pinterest.ktlint.cli.reporter.core.api.ReporterV2
 import java.nio.file.Path
 
@@ -32,7 +33,7 @@ class DiktatReporterImpl(
         /**
          * @return __KtLint__'s [ReporterV2]
          */
-        fun DiktatReporter.unwrap(): ReporterV2 = (this as? DiktatReporterImpl)?.ktLintReporter
+        fun DiktatReporter.unwrap(): ReporterV2 = (this as? DiktatReporterImpl)?.ktLintReporter?.unwrapIfNeeded()
             ?: error("Unsupported wrapper of ${DiktatReporter::class.java.simpleName}: ${this::class.java.canonicalName}")
     }
 }
diff --git a/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/KtLintUtils.kt b/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/KtLintUtils.kt
index b71753324f..82a61e855a 100644
--- a/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/KtLintUtils.kt
+++ b/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/KtLintUtils.kt
@@ -10,9 +10,14 @@ import com.saveourtool.diktat.api.DiktatRuleSet
 import com.saveourtool.diktat.common.config.rules.DIKTAT_RULE_SET_ID
 
 import com.pinterest.ktlint.cli.reporter.core.api.KtlintCliError
+import com.pinterest.ktlint.cli.reporter.core.api.ReporterProviderV2
+import com.pinterest.ktlint.cli.reporter.core.api.ReporterV2
 import com.pinterest.ktlint.rule.engine.api.LintError
 import com.pinterest.ktlint.rule.engine.core.api.RuleId
 import org.intellij.lang.annotations.Language
+import org.jetbrains.kotlin.utils.addToStdlib.applyIf
+import java.io.OutputStream
+import java.io.PrintStream
 
 import java.nio.file.Path
 
@@ -89,6 +94,20 @@ fun String.correctErrorDetail(canBeAutoCorrected: Boolean): String = if (canBeAu
  */
 fun Path.relativePathStringTo(sourceRootDir: Path): String = relativeTo(sourceRootDir).invariantSeparatorsPathString
 
+/**
+ * @param out [OutputStream] for [ReporterV2]
+ * @param closeOutAfterAll close [OutputStream] in [ReporterV2.afterAll]
+ * @param opt configuration for [ReporterV2]
+ * @return created [ReporterV2] which closes [out] in [ReporterV2.afterAll] if it's required
+ */
+fun <R : ReporterV2> ReporterProviderV2<R>.get(
+    out: OutputStream,
+    closeOutAfterAll: Boolean,
+    opt: Map<String, String>,
+): ReporterV2 = get(out.printStream(), opt).applyIf(closeOutAfterAll) {
+    closeAfterAll(out)
+}
+
 /**
  * Enables ignoring autocorrected errors when in "fix" mode (i.e. when
  * [com.pinterest.ktlint.core.KtLint.format] is invoked).
@@ -108,6 +127,16 @@ private fun DiktatCallback.ignoreCorrectedErrors(): DiktatCallback = DiktatCallb
     }
 }
 
+private fun OutputStream.printStream(): PrintStream = (this as? PrintStream) ?: PrintStream(this)
+
+private fun ReporterV2.closeAfterAll(outputStream: OutputStream): ReporterV2 = object : ReporterV2Wrapper(this@closeAfterAll) {
+    override fun afterAll() {
+        super.afterAll()
+        outputStream.flush()
+        outputStream.close()
+    }
+}
+
 /**
  * @param ruleSetSupplier
  * @param file
diff --git a/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/ReporterV2Wrapper.kt b/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/ReporterV2Wrapper.kt
new file mode 100644
index 0000000000..6669c794ce
--- /dev/null
+++ b/diktat-ktlint-engine/src/main/kotlin/com/saveourtool/diktat/ktlint/ReporterV2Wrapper.kt
@@ -0,0 +1,32 @@
+package com.saveourtool.diktat.ktlint
+
+import com.pinterest.ktlint.cli.reporter.core.api.KtlintCliError
+import com.pinterest.ktlint.cli.reporter.core.api.ReporterV2
+
+/**
+ * Wrapper for [ReporterV2]
+ *
+ * @param reporterV2
+ */
+open class ReporterV2Wrapper(private val reporterV2: ReporterV2) : ReporterV2 {
+    override fun beforeAll() = reporterV2.beforeAll()
+
+    override fun before(file: String) = reporterV2.before(file)
+
+    override fun onLintError(file: String, ktlintCliError: KtlintCliError) = reporterV2.onLintError(file, ktlintCliError)
+
+    override fun after(file: String) = reporterV2.after(file)
+
+    override fun afterAll() = reporterV2.afterAll()
+
+    companion object {
+        /**
+         * @return unwrapped [ReporterV2Wrapper] if it's required
+         */
+        fun ReporterV2.unwrapIfNeeded(): ReporterV2 = if (this is ReporterV2Wrapper) {
+            this.reporterV2
+        } else {
+            this
+        }
+    }
+}