diff --git a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/actions/ScratchAction.kt b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/actions/ScratchAction.kt
index 8abc7fc24f8..0870e284e8a 100644
--- a/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/actions/ScratchAction.kt
+++ b/idea/idea-jvm/src/org/jetbrains/kotlin/idea/scratch/actions/ScratchAction.kt
@@ -7,12 +7,18 @@ package org.jetbrains.kotlin.idea.scratch.actions
 
 import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.actionSystem.AnActionEvent
+import com.intellij.openapi.actionSystem.CommonDataKeys
+import com.intellij.openapi.fileEditor.impl.text.TextEditorProvider
+import org.jetbrains.kotlin.idea.scratch.getScratchPanel
 import org.jetbrains.kotlin.idea.scratch.getScratchPanelFromSelectedEditor
 import javax.swing.Icon
 
 abstract class ScratchAction(message: String, icon: Icon) : AnAction(message, message, icon) {
     override fun update(e: AnActionEvent) {
-        val scratchPanel = e.project?.let { getScratchPanelFromSelectedEditor(it) }
+        val scratchPanel = e.getData(CommonDataKeys.EDITOR)
+            ?.let { TextEditorProvider.getInstance().getTextEditor(it).getScratchPanel() }
+            ?: e.project?.let { getScratchPanelFromSelectedEditor(it) }
+
         e.presentation.isVisible = scratchPanel != null
     }
 }
\ No newline at end of file