diff --git a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ScratchFileAutoRunner.kt b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ScratchFileAutoRunner.kt
index 8ededf92b24..93cd29645c5 100644
--- a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ScratchFileAutoRunner.kt
+++ b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ScratchFileAutoRunner.kt
@@ -19,6 +19,7 @@ import com.intellij.util.Alarm
 import org.jetbrains.kotlin.idea.scratch.actions.RunScratchAction
 import org.jetbrains.kotlin.idea.scratch.actions.RunScratchFromHereAction
 import org.jetbrains.kotlin.idea.scratch.actions.ScratchCompilationSupport
+import org.jetbrains.kotlin.idea.scratch.ui.findScratchFileEditorWithPreview
 
 class ScratchFileAutoRunner(private val project: Project) : DocumentListener {
     companion object {
@@ -72,6 +73,7 @@ class ScratchFileAutoRunner(private val project: Project) : DocumentListener {
     }
 
     private fun getScratchFile(file: VirtualFile, project: Project): ScratchFile? {
-        return getScratchFileFromEditorSelectedForFile(FileEditorManager.getInstance(project), file)
+        val editor = FileEditorManager.getInstance(project).getSelectedEditor(file) as? TextEditor
+        return editor?.findScratchFileEditorWithPreview()?.scratchFile
     }
 }
\ No newline at end of file
diff --git a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/scratchUtils.kt b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/scratchUtils.kt
index 78b51e8b1af..9830a3ca17d 100644
--- a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/scratchUtils.kt
+++ b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/scratchUtils.kt
@@ -13,7 +13,9 @@ import com.intellij.openapi.fileEditor.TextEditor
 import com.intellij.openapi.fileEditor.impl.text.TextEditorProvider
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
+import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.idea.actions.KOTLIN_WORKSHEET_EXTENSION
+import org.jetbrains.kotlin.idea.scratch.ui.KtScratchFileEditorWithPreview
 import org.jetbrains.kotlin.idea.scratch.ui.findScratchFileEditorWithPreview
 
 internal val LOG = Logger.getInstance("#org.jetbrains.kotlin.idea.scratch")
@@ -27,9 +29,10 @@ val VirtualFile.isKotlinWorksheet: Boolean
 val VirtualFile.isKotlinScratch: Boolean
     get() = ScratchFileService.getInstance().getRootType(this) is ScratchRootType
 
-fun getScratchFileFromEditorSelectedForFile(fileManager: FileEditorManager, virtualFile: VirtualFile): ScratchFile? {
+@TestOnly
+fun getScratchEditorForSelectedFile(fileManager: FileEditorManager, virtualFile: VirtualFile): KtScratchFileEditorWithPreview? {
     val editor = fileManager.getSelectedEditor(virtualFile) as? TextEditor ?: return null
-    return editor.getScratchFile()
+    return editor.findScratchFileEditorWithPreview()
 }
 
 fun getScratchFileFromSelectedEditor(project: Project): ScratchFile? {
diff --git a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ui/KtScratchFileEditorProvider.kt b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ui/KtScratchFileEditorProvider.kt
index e4ec3f997f5..3cfc9c05170 100644
--- a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ui/KtScratchFileEditorProvider.kt
+++ b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/ui/KtScratchFileEditorProvider.kt
@@ -23,6 +23,7 @@ import com.intellij.openapi.util.Key
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.pom.Navigatable
 import com.intellij.psi.PsiManager
+import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.idea.scratch.*
 import org.jetbrains.kotlin.idea.scratch.output.InlayScratchOutputHandler
 import org.jetbrains.kotlin.idea.scratch.output.ScratchOutputHandlerAdapter
@@ -93,6 +94,11 @@ class KtScratchFileEditorWithPreview private constructor(
         inlayOutputHandler.clear(scratchFile)
     }
 
+    @TestOnly
+    fun setPreviewEnabled(isPreviewEnabled: Boolean) {
+        layout = if (isPreviewEnabled) Layout.SHOW_EDITOR_AND_PREVIEW else Layout.SHOW_EDITOR
+    }
+
     companion object {
         fun create(scratchFile: ScratchFile): KtScratchFileEditorWithPreview {
             val textEditorProvider = TextEditorProvider.getInstance()
diff --git a/idea/tests/org/jetbrains/kotlin/idea/scratch/AbstractScratchRunActionTest.kt b/idea/tests/org/jetbrains/kotlin/idea/scratch/AbstractScratchRunActionTest.kt
index f56477d8772..42a30211bb2 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/scratch/AbstractScratchRunActionTest.kt
+++ b/idea/tests/org/jetbrains/kotlin/idea/scratch/AbstractScratchRunActionTest.kt
@@ -30,6 +30,7 @@ import org.jetbrains.kotlin.idea.scratch.actions.RunScratchAction
 import org.jetbrains.kotlin.idea.scratch.actions.ScratchCompilationSupport
 import org.jetbrains.kotlin.idea.scratch.output.InlayScratchFileRenderer
 import org.jetbrains.kotlin.idea.scratch.output.getInlays
+import org.jetbrains.kotlin.idea.scratch.ui.KtScratchFileEditorWithPreview
 import org.jetbrains.kotlin.idea.test.KotlinLightProjectDescriptor
 import org.jetbrains.kotlin.idea.test.KotlinWithJdkAndRuntimeLightProjectDescriptor
 import org.jetbrains.kotlin.idea.test.PluginTestCaseBase
@@ -175,12 +176,12 @@ abstract class AbstractScratchRunActionTest : FileEditorManagerTestCase() {
 
         ScriptDependenciesManager.updateScriptDependenciesSynchronously(scratchVirtualFile, project)
 
-        val scratchFile = getScratchFileFromEditorSelectedForFile(myManager, myFixture.file.virtualFile)
+        val scratchFileEditor = getScratchEditorForSelectedFile(myManager, myFixture.file.virtualFile)
             ?: error("Couldn't find scratch file")
 
-        configureOptions(scratchFile, text, myFixture.module)
+        configureOptions(scratchFileEditor, text, myFixture.module)
 
-        return scratchFile
+        return scratchFileEditor.scratchFile
     }
 
     protected fun configureWorksheetByText(name: String, text: String): ScratchFile {
@@ -188,14 +189,14 @@ abstract class AbstractScratchRunActionTest : FileEditorManagerTestCase() {
 
         ScriptDependenciesManager.updateScriptDependenciesSynchronously(worksheetFile, project)
 
-        val scratchFile = getScratchFileFromEditorSelectedForFile(myManager, myFixture.file.virtualFile)
+        val scratchFileEditor = getScratchEditorForSelectedFile(myManager, myFixture.file.virtualFile)
             ?: error("Couldn't find scratch panel")
 
         // We want to check that correct module is selected automatically,
         // that's why we set `module` to null so it wouldn't be changed
-        configureOptions(scratchFile, text, null)
+        configureOptions(scratchFileEditor, text, null)
 
-        return scratchFile
+        return scratchFileEditor.scratchFile
     }
 
 
@@ -227,7 +228,7 @@ abstract class AbstractScratchRunActionTest : FileEditorManagerTestCase() {
 
     protected fun stopReplProcess() {
         if (myFixture.file != null) {
-            val scratchFile = getScratchFileFromEditorSelectedForFile(myManager, myFixture.file.virtualFile)
+            val scratchFile = getScratchEditorForSelectedFile(myManager, myFixture.file.virtualFile)?.scratchFile
                 ?: error("Couldn't find scratch panel")
             scratchFile.replScratchExecutor?.stopAndWait()
         }
@@ -303,10 +304,12 @@ abstract class AbstractScratchRunActionTest : FileEditorManagerTestCase() {
         }
 
         fun configureOptions(
-            scratchFile: ScratchFile,
+            scratchFileEditor: KtScratchFileEditorWithPreview,
             fileText: String,
             module: Module?
         ) {
+            val scratchFile = scratchFileEditor.scratchFile
+
             if (InTextDirectivesUtils.getPrefixedBoolean(fileText, "// INTERACTIVE_MODE: ") != true) {
                 scratchFile.saveOptions { copy(isInteractiveMode = false) }
             }
@@ -318,6 +321,9 @@ abstract class AbstractScratchRunActionTest : FileEditorManagerTestCase() {
             if (module != null && !InTextDirectivesUtils.isDirectiveDefined(fileText, "// NO_MODULE")) {
                 scratchFile.setModule(module)
             }
+
+            val isPreviewEnabled = InTextDirectivesUtils.getPrefixedBoolean(fileText, "// PREVIEW_ENABLED: ") == true
+            scratchFileEditor.setPreviewEnabled(isPreviewEnabled)
         }
 
     }
diff --git a/idea/tests/org/jetbrains/kotlin/idea/scratch/ScratchLineMarkersTest.kt b/idea/tests/org/jetbrains/kotlin/idea/scratch/ScratchLineMarkersTest.kt
index a692b9b9bac..4b77a2c9f66 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/scratch/ScratchLineMarkersTest.kt
+++ b/idea/tests/org/jetbrains/kotlin/idea/scratch/ScratchLineMarkersTest.kt
@@ -39,10 +39,10 @@ abstract class AbstractScratchLineMarkersTest : FileEditorManagerTestCase() {
 
         ScriptDependenciesManager.updateScriptDependenciesSynchronously(scratchVirtualFile, project)
 
-        val scratchFile = getScratchFileFromEditorSelectedForFile(FileEditorManager.getInstance(project), myFixture.file.virtualFile)
+        val scratchFileEditor = getScratchEditorForSelectedFile(FileEditorManager.getInstance(project), myFixture.file.virtualFile)
             ?: error("Couldn't find scratch panel")
 
-        configureOptions(scratchFile, fileText, null)
+        configureOptions(scratchFileEditor, fileText, null)
 
         val project = myFixture.project
         val document = myFixture.editor.document