PY-76049 [Jupyter] Fix document intervals sync and custom folding update

GitOrigin-RevId: 2a22e541554e5e158fdb4db9dfd32fd386e7f684
2026-03-22 06:50:54 +07:00 · 2024-09-25 08:20:48 +02:00
parent 0273ef891d
commit 8fd29d3516
13 changed files with 203 additions and 166 deletions
--- a/notebooks/visualization/src/com/intellij/notebooks/visualization/NotebookCellInlayManager.kt
+++ b/notebooks/visualization/src/com/intellij/notebooks/visualization/NotebookCellInlayManager.kt
@@ -1,25 +1,6 @@
 package com.intellij.notebooks.visualization

 import com.intellij.ide.ui.LafManagerListener
-import com.intellij.openapi.Disposable
-import com.intellij.openapi.application.ApplicationManager
-import com.intellij.openapi.application.runInEdt
-import com.intellij.openapi.diagnostic.thisLogger
-import com.intellij.openapi.editor.Editor
-import com.intellij.openapi.editor.FoldRegion
-import com.intellij.openapi.editor.colors.EditorColorsListener
-import com.intellij.openapi.editor.colors.EditorColorsManager
-import com.intellij.openapi.editor.event.CaretEvent
-import com.intellij.openapi.editor.event.CaretListener
-import com.intellij.openapi.editor.ex.EditorEx
-import com.intellij.openapi.editor.ex.FoldingListener
-import com.intellij.openapi.editor.impl.EditorImpl
-import com.intellij.openapi.util.Disposer
-import com.intellij.openapi.util.Key
-import com.intellij.openapi.util.registry.Registry
-import com.intellij.util.EventDispatcher
-import com.intellij.util.SmartList
-import com.intellij.util.concurrency.ThreadingAssertions
 import com.intellij.notebooks.ui.editor.actions.command.mode.NOTEBOOK_EDITOR_MODE
 import com.intellij.notebooks.ui.editor.actions.command.mode.NotebookEditorMode
 import com.intellij.notebooks.ui.editor.actions.command.mode.NotebookEditorModeListener
@@ -29,7 +10,28 @@ import com.intellij.notebooks.visualization.ui.*
 import com.intellij.notebooks.visualization.ui.EditorCellEventListener.*
 import com.intellij.notebooks.visualization.ui.EditorCellViewEventListener.CellViewCreated
 import com.intellij.notebooks.visualization.ui.EditorCellViewEventListener.CellViewRemoved
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.runInEdt
+import com.intellij.openapi.diagnostic.thisLogger
+import com.intellij.openapi.editor.Document
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.editor.FoldRegion
+import com.intellij.openapi.editor.colors.EditorColorsListener
+import com.intellij.openapi.editor.colors.EditorColorsManager
+import com.intellij.openapi.editor.event.BulkAwareDocumentListener
+import com.intellij.openapi.editor.event.CaretEvent
+import com.intellij.openapi.editor.event.CaretListener
+import com.intellij.openapi.editor.ex.EditorEx
+import com.intellij.openapi.editor.ex.FoldingListener
 import com.intellij.openapi.editor.ex.FoldingModelEx
+import com.intellij.openapi.editor.impl.EditorImpl
+import com.intellij.openapi.util.Disposer
+import com.intellij.openapi.util.Key
+import com.intellij.openapi.util.registry.Registry
+import com.intellij.util.EventDispatcher
+import com.intellij.util.SmartList
+import com.intellij.util.concurrency.ThreadingAssertions

 class NotebookCellInlayManager private constructor(
  val editor: EditorImpl,
@@ -46,6 +48,8 @@ class NotebookCellInlayManager private constructor(

  val cells: List<EditorCell> get() = _cells.toList()

+  val views = mutableMapOf<EditorCell, EditorCellView>()
+
  /**
   * Listens for inlay changes (called after all inlays are updated). Feel free to convert it to the EP if you need another listener
   */
@@ -61,6 +65,14 @@ class NotebookCellInlayManager private constructor(

  private var updateCtx: UpdateContext? = null

+  /*
+   EditorImpl sets `myDocumentChangeInProgress` attribute to true during document update processing, that prevents correct update
+   of custom folding regions.When this flag is set, folding updates will be postponed until the editor finishes its work.
+   */
+  private var editorIsProcessingDocument = false
+
+  private var postponedUpdates = mutableListOf<UpdateContext>()
+
  fun <T> update(force: Boolean = false, block: (updateCtx: UpdateContext) -> T): T {
    val ctx = updateCtx
    return if (ctx != null) {
@@ -74,7 +86,11 @@ class NotebookCellInlayManager private constructor(
        val r = keepScrollingPositionWhile(editor) {
          val r = block(newCtx)
          updateCtx = null
-          newCtx.applyUpdates(editor)
+          if (editorIsProcessingDocument) {
+            postponedUpdates.add(newCtx)
+          } else {
+            newCtx.applyUpdates(editor)
+          }
          r
        }
        inlaysChanged()
@@ -146,8 +162,6 @@ class NotebookCellInlayManager private constructor(

    editor.putUserData(CELL_INLAY_MANAGER_KEY, this)

-    handleRefreshedDocument()
-
    val connection = ApplicationManager.getApplication().messageBus.connect(editor.disposable)
    connection.subscribe(EditorColorsManager.TOPIC, EditorColorsListener {
      updateAll()
@@ -170,6 +184,77 @@ class NotebookCellInlayManager private constructor(
    setupSelectionUI()

    ApplicationManager.getApplication().messageBus.connect(this).subscribe(NOTEBOOK_EDITOR_MODE, this)
+
+    cellEventListeners.addListener(object : EditorCellEventListener {
+      override fun onEditorCellEvents(events: List<EditorCellEvent>) {
+        updateUI(events)
+      }
+    })
+
+    editor.document.addDocumentListener(object : BulkAwareDocumentListener.Simple {
+      override fun beforeDocumentChange(document: Document) {
+        editorIsProcessingDocument = true
+      }
+
+      override fun afterDocumentChange(document: Document) {
+        editorIsProcessingDocument = false
+        postponedUpdates.forEach {
+          it.applyUpdates(editor)
+        }
+        postponedUpdates.clear()
+      }
+    })
+
+    handleRefreshedDocument()
+  }
+
+  private fun updateUI(events: List<EditorCellEvent>) {
+    update { ctx ->
+      for (event in events) {
+        when (event) {
+          is CellCreated -> {
+            val cell = event.cell
+            cell.visible.afterChange { visible ->
+              if (visible) {
+                createCellViewIfNecessary(cell, ctx)
+              }
+              else {
+                disposeCellView(cell)
+              }
+            }
+            createCellViewIfNecessary(cell, ctx)
+          }
+          is CellRemoved -> {
+            disposeCellView(event.cell)
+          }
+        }
+      }
+    }
+  }
+
+
+  private fun createCellViewIfNecessary(cell: EditorCell, ctx: UpdateContext) {
+    if (views[cell] == null) {
+      createCellView(cell, ctx)
+    }
+  }
+
+  private fun createCellView(
+    cell: EditorCell,
+    ctx: UpdateContext,
+  ) {
+    val view = EditorCellView(editor, notebookCellLines, cell, this)
+    Disposer.register(cell, view)
+    view.updateCellFolding(ctx)
+    views[cell] = view
+    fireCellViewCreated(view)
+  }
+
+  private fun disposeCellView(cell: EditorCell) {
+    views.remove(cell)?.let { view ->
+      fireCellViewRemoved(view)
+      Disposer.dispose(view)
+    }
  }

  private fun setupSelectionUI() {
@@ -184,9 +269,9 @@ class NotebookCellInlayManager private constructor(
    val selectionModel = editor.cellSelectionModel ?: error("The selection model is supposed to be installed")
    val selectedCells = selectionModel.selectedCells.map { it.ordinal }
    for (cell in cells) {
-      cell.selected = cell.intervalPointer.get()?.ordinal in selectedCells
+      cell.selected.set(cell.intervalPointer.get()?.ordinal in selectedCells)

-      if (cell.selected) {
+      if (cell.selected.get()) {
        editor.project?.messageBus?.syncPublisher(JupyterCellSelectionNotifier.TOPIC)?.cellSelected(cell.interval, editor)
      }
    }
@@ -221,12 +306,12 @@ class NotebookCellInlayManager private constructor(
        update { ctx ->
          changedRegions.forEach { region ->
            editorCells(region).forEach {
-              it.visible = region.isExpanded
+              it.visible.set(region.isExpanded)
            }
          }
          removedRegions.forEach { region ->
            editorCells(region).forEach {
-              it.visible = true
+              it.visible.set(true)
            }
          }
        }
@@ -253,16 +338,9 @@ class NotebookCellInlayManager private constructor(
      }.toMutableList()
    }
    cellEventListeners.multicaster.onEditorCellEvents(_cells.map { CellCreated(it) })
-    update {
-      _cells.forEach {
-        it.initView()
-      }
-    }
  }

-  private fun createCell(interval: NotebookIntervalPointer) = EditorCell(editor, this, interval) { cell ->
-    EditorCellView(editor, notebookCellLines, cell, this).also { Disposer.register(cell, it) }
-  }.also {
+  private fun createCell(interval: NotebookIntervalPointer) = EditorCell(editor, this, interval).also {
    cellExtensionFactories.forEach { factory ->
      factory.onCellCreated(it)
    }
@@ -285,7 +363,7 @@ class NotebookCellInlayManager private constructor(
      shouldCheckInlayOffsets: Boolean,
      inputFactories: List<NotebookCellInlayController.InputFactory> = listOf(),
      cellExtensionFactories: List<CellExtensionFactory> = listOf(),
-    ) : NotebookCellInlayManager {
+    ): NotebookCellInlayManager {
      EditorEmbeddedComponentContainer(editor as EditorEx)
      val notebookCellInlayManager = NotebookCellInlayManager(
        editor,
@@ -294,8 +372,8 @@ class NotebookCellInlayManager private constructor(
        cellExtensionFactories
      ).also { Disposer.register(editor.disposable, it) }
      editor.putUserData(isFoldingEnabledKey, Registry.`is`("jupyter.editor.folding.cells"))
-      NotebookIntervalPointerFactory.get(editor).changeListeners.addListener(notebookCellInlayManager, editor.disposable)
      notebookCellInlayManager.initialize()
+      NotebookIntervalPointerFactory.get(editor).changeListeners.addListener(notebookCellInlayManager, editor.disposable)
      return notebookCellInlayManager
    }

@@ -315,7 +393,6 @@ class NotebookCellInlayManager private constructor(
        is NotebookIntervalPointersEvent.OnEdited -> {
          val cell = _cells[change.intervalAfter.ordinal]
          cell.updateInput()
-          events.add(CellUpdated(cell))
        }
        is NotebookIntervalPointersEvent.OnInserted -> {
          change.subsequentPointers.forEach {
@@ -344,9 +421,6 @@ class NotebookCellInlayManager private constructor(
      fixInlaysOffsetsAfterNewCellInsert(change, ctx)
    }
    cellEventListeners.multicaster.onEditorCellEvents(events)
-
-    events.filterIsInstance<CellCreated>().forEach { it.cell.initView() }
-
    checkInlayOffsets()
  }

--- a/notebooks/visualization/src/com/intellij/notebooks/visualization/NotebookIntervalPointerImpl.kt
+++ b/notebooks/visualization/src/com/intellij/notebooks/visualization/NotebookIntervalPointerImpl.kt
@@ -52,7 +52,6 @@ class UndoableActionListener : BulkAwareDocumentListener.Simple {
        override fun redo() {}
      }
      registerUndoableAction(action)
-      document.notebookIntervalPointerFactory.onUpdated(NotebookIntervalPointersEvent(eventChanges))
    }
  }
 }
@@ -183,6 +182,8 @@ class NotebookIntervalPointerFactoryImpl(

    //Postpone UndoableAction registration until DocumentUndoProvider registers its own action.
    document.putUserData(POSTPONED_CHANGES_KEY, PostponedChanges(eventChanges, shiftChanges))
+
+    onUpdated(NotebookIntervalPointersEvent(eventChanges))
  }

  private fun applyPostponedChanges(event: NotebookCellLinesEvent) {
@@ -190,6 +191,7 @@ class NotebookIntervalPointerFactoryImpl(
    ThreadingAssertions.assertWriteAccess()
    updatePointersByChanges(postponedChanges.eventChanges)
    updatePointersByChanges(postponedChanges.shiftChanges)
+    onUpdated(NotebookIntervalPointersEvent(postponedChanges.eventChanges))
  }

  private fun updatePointersByChanges(changes: List<Change>) {
--- a/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/ControllerEditorCellViewComponent.kt
+++ b/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/ControllerEditorCellViewComponent.kt
@@ -14,15 +14,22 @@ class ControllerEditorCellViewComponent(
  internal val controller: NotebookCellInlayController,
  private val editor: Editor,
  private val cell: EditorCell,
-) : EditorCellViewComponent(), HasGutterIcon {
+) : EditorCellViewComponent() {

  private var foldedRegion: FoldRegion? = null

-  override fun updateGutterIcons(gutterAction: AnAction?) {
+  private fun updateGutterIcons(gutterAction: AnAction?) {
    val inlay = controller.inlay
    inlay.update()
  }

+  init {
+    cell.gutterAction.afterChange(this) { action ->
+      updateGutterIcons(action)
+    }
+    updateGutterIcons(cell.gutterAction.get())
+  }
+
  override fun dispose() {
    super.dispose()
    Disposer.dispose(controller.inlay)
--- a/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/CustomFoldingEditorCellViewComponent.kt
+++ b/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/CustomFoldingEditorCellViewComponent.kt
@@ -25,7 +25,7 @@ class CustomFoldingEditorCellViewComponent(
  internal val component: JComponent,
  private val editor: EditorEx,
  private val cell: EditorCell,
-) : EditorCellViewComponent(), HasGutterIcon {
+) : EditorCellViewComponent() {

  private var foldingRegion: CustomFoldRegion? = null

@@ -50,26 +50,36 @@ class CustomFoldingEditorCellViewComponent(
    return component
  }

-  override fun updateGutterIcons(gutterAction: AnAction?) {
-    gutterActionRenderer = gutterAction?.let { ActionToGutterRendererAdapter(it) }
-    foldingRegion?.update()
+  private fun updateGutterIcons(gutterAction: AnAction?) {
+    cell.manager.update { ctx ->
+      gutterActionRenderer = gutterAction?.let { ActionToGutterRendererAdapter(it) }
+      ctx.addFoldingOperation { modelEx ->
+        foldingRegion?.update()
+      }
+    }
  }

-  override fun dispose() {
-    super.dispose()
-    disposeFolding()
+  init {
+    cell.gutterAction.afterChange(this) { action ->
+      updateGutterIcons(action)
+    }
+    updateGutterIcons(cell.gutterAction.get())
  }

-  private fun disposeFolding() = cell.manager.update { ctx ->
-    foldingRegion?.let { region ->
-      ctx.addFoldingOperation {
+  override fun dispose() = cell.manager.update { ctx ->
+    disposeFolding(ctx)
+  }
+
+  private fun disposeFolding(ctx: UpdateContext) {
+    ctx.addFoldingOperation {
+      foldingRegion?.let { region ->
        if (region.isValid == true) {
          editor.foldingModel.removeFoldRegion(region)
        }
      }
+      foldingRegion = null
    }
    editor.componentContainer.remove(mainComponent)
-    foldingRegion = null
  }

  override fun calculateBounds(): Rectangle {
--- a/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/EditorCell.kt
+++ b/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/EditorCell.kt
@@ -13,6 +13,7 @@ import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.editor.ex.EditorEx
 import com.intellij.openapi.editor.impl.EditorImpl
+import com.intellij.openapi.observable.properties.AtomicBooleanProperty
 import com.intellij.openapi.observable.properties.AtomicProperty
 import com.intellij.openapi.util.*
 import java.time.ZonedDateTime
@@ -24,7 +25,6 @@ class EditorCell(
  private val editor: EditorEx,
  val manager: NotebookCellInlayManager,
  var intervalPointer: NotebookIntervalPointer,
-  private val viewFactory: (EditorCell) -> EditorCellView,
 ) : Disposable, UserDataHolder by UserDataHolderBase() {

  val source = AtomicProperty<String>(getSource())
@@ -42,68 +42,20 @@ class EditorCell(

  val interval get() = intervalPointer.get() ?: error("Invalid interval")

-  var view: EditorCellView? = null
+  val view: EditorCellView?
+    get() = manager.views[this]

-  var visible: Boolean = true
-    set(value) {
-      if (field == value) return
-      field = value
-      manager.update<Unit> { ctx ->
-        if (!value) {
-          view?.let {
-            disposeView(it)
-          }
-        }
-        else {
-          if (view == null) {
-            view = createView()
-          }
-        }
-      }
-    }
+  var visible = AtomicBooleanProperty(true)

  init {
    CELL_EXTENSION_CONTAINER_KEY.set(this, mutableMapOf())
  }

-  fun initView() {
-    view = createView()
-  }
+  val selected = AtomicBooleanProperty(false)

-  private fun createView(): EditorCellView = manager.update { ctx ->
-    val view = viewFactory(this).also { Disposer.register(this, it) }
-    gutterAction?.let { view.setGutterAction(it) }
-    view.updateExecutionStatus(executionCount, progressStatus, executionStartTime, executionEndTime)
-    view.selected = selected
-    manager.fireCellViewCreated(view)
-    view.updateCellFolding(ctx)
-    view
-  }
+  val gutterAction = AtomicProperty<AnAction?>(null)

-  private fun disposeView(it: EditorCellView) {
-    Disposer.dispose(it)
-    view = null
-    manager.fireCellViewRemoved(it)
-  }
-
-  var selected: Boolean = false
-    set(value) {
-      if (field != value) {
-        field = value
-        view?.selected = value
-      }
-    }
-
-  var gutterAction: AnAction? = null
-    private set
-
-  private var executionCount: Int? = null
-
-  private var progressStatus: ProgressStatus? = null
-
-  private var executionStartTime: ZonedDateTime? = null
-
-  private var executionEndTime: ZonedDateTime? = null
+  val executionStatus = AtomicProperty<ExecutionStatus>(ExecutionStatus())

  val mode = AtomicProperty<NotebookEditorMode>(NotebookEditorMode.COMMAND)

@@ -111,7 +63,6 @@ class EditorCell(

  override fun dispose() {
    cleanupExtensions()
-    view?.let { disposeView(it) }
  }

  private fun cleanupExtensions() {
@@ -139,8 +90,7 @@ class EditorCell(
  }

  fun setGutterAction(action: AnAction?) {
-    gutterAction = action
-    view?.setGutterAction(action)
+    gutterAction.set(action)
  }

  inline fun <reified T : NotebookCellInlayController> getController(): T? {
@@ -174,10 +124,10 @@ class EditorCell(

  private fun getOutputs(): List<NotebookOutputDataKey> =
    NotebookOutputDataKeyExtractor.EP_NAME.extensionList.asSequence()
-       .mapNotNull { it.extract(editor as EditorImpl, interval) }
-       .firstOrNull()
-       ?.takeIf { it.isNotEmpty() }
-     ?: emptyList()
+      .mapNotNull { it.extract(editor as EditorImpl, interval) }
+      .firstOrNull()
+      ?.takeIf { it.isNotEmpty() }
+    ?: emptyList()

  private fun updateOutputs(newOutputs: List<NotebookOutputDataKey>) = runInEdt {
    outputs.set(newOutputs)
@@ -186,22 +136,18 @@ class EditorCell(
  fun onExecutionEvent(event: ExecutionEvent) {
    when (event) {
      is ExecutionEvent.ExecutionStarted -> {
-        executionStartTime = event.startTime
-        progressStatus = event.status
+        executionStatus.set(executionStatus.get().copy(status = event.status, startTime = event.startTime))
      }
      is ExecutionEvent.ExecutionStopped -> {
-        executionEndTime = event.endTime
-        progressStatus = event.status
-        executionCount = event.executionCount
+        executionStatus.set(executionStatus.get().copy(status = event.status, endTime = event.endTime, count = event.executionCount))
      }
      is ExecutionEvent.ExecutionSubmitted -> {
-        progressStatus = event.status
+        executionStatus.set(executionStatus.get().copy(status = event.status))
      }
      is ExecutionEvent.ExecutionReset -> {
-        progressStatus = event.status
+        executionStatus.set(executionStatus.get().copy(status = event.status))
      }
    }
-    view?.updateExecutionStatus(executionCount, progressStatus, executionStartTime, executionEndTime)
  }

  fun switchToEditMode() = runInEdt {
@@ -236,4 +182,11 @@ class EditorCell(
  private fun forEachExtension(action: (EditorCellExtension) -> Unit) {
    CELL_EXTENSION_CONTAINER_KEY.get(this)?.values?.forEach { action(it) }
  }
+
+  data class ExecutionStatus(
+    val status: ProgressStatus? = null,
+    val count: Int? = null,
+    val startTime: ZonedDateTime? = null,
+    val endTime: ZonedDateTime? = null
+  )
 }
--- a/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/EditorCellEventListener.kt
+++ b/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/EditorCellEventListener.kt
@@ -8,5 +8,4 @@ interface EditorCellEventListener : EventListener {
  sealed interface EditorCellEvent
  data class CellCreated(val cell: EditorCell) : EditorCellEvent
  data class CellRemoved(val cell: EditorCell) : EditorCellEvent
-  data class CellUpdated(val cell: EditorCell) : EditorCellEvent
 }
--- a/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/EditorCellInput.kt
+++ b/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/EditorCellInput.kt
@@ -4,7 +4,6 @@ import com.intellij.notebooks.ui.visualization.NotebookEditorAppearanceUtils.isO
 import com.intellij.notebooks.ui.visualization.notebookAppearance
 import com.intellij.notebooks.visualization.NotebookCellInlayController
 import com.intellij.notebooks.visualization.NotebookCellLines
-import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.editor.impl.EditorImpl
 import java.awt.Rectangle

@@ -25,8 +24,6 @@ class EditorCellInput(

  val folding = EditorCellFoldingBar(editor, ::getFoldingBounds) { toggleFolding() }

-  private var gutterAction: AnAction? = null
-
  var folded = false
    private set

@@ -65,16 +62,6 @@ class EditorCellInput(

  fun update() {
    updateInput()
-    updateGutterIcons()
-  }
-
-  private fun updateGutterIcons() {
-    (component as? HasGutterIcon)?.updateGutterIcons(gutterAction)
-  }
-
-  fun setGutterAction(action: AnAction?) {
-    gutterAction = action
-    updateGutterIcons()
  }

  override fun calculateBounds(): Rectangle {
--- a/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/EditorCellSelectionModel.kt
+++ b/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/EditorCellSelectionModel.kt
@@ -21,7 +21,7 @@ class EditorCellSelectionModel(manager: NotebookCellInlayManager) {

  private fun removeCell(selectedCell: EditorCell) {
    _selection.remove(selectedCell)
-    selectedCell.selected = false
+    selectedCell.selected.set(false)
  }

  fun replaceSelection(cells: Collection<EditorCell>) {
@@ -29,10 +29,10 @@ class EditorCellSelectionModel(manager: NotebookCellInlayManager) {
    val toRemove = _selection - selectionSet
    val toAdd = selectionSet - _selection
    toRemove.forEach {
-      it.selected = false
+      it.selected.set(false)
    }
    toAdd.forEach {
-      it.selected = true
+      it.selected.set(true)
    }
    _selection.removeAll(toRemove)
    _selection.addAll(toAdd)
--- a/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/EditorCellView.kt
+++ b/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/EditorCellView.kt
@@ -8,7 +8,6 @@ import com.intellij.notebooks.visualization.*
 import com.intellij.notebooks.visualization.NotebookCellInlayController.InputFactory
 import com.intellij.notebooks.visualization.context.NotebookDataContext
 import com.intellij.openapi.Disposable
-import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.actionSystem.DataProvider
 import com.intellij.openapi.actionSystem.PlatformCoreDataKeys
 import com.intellij.openapi.actionSystem.PlatformDataKeys
@@ -90,6 +89,15 @@ class EditorCellView(
    cell.source.afterChange(this) {
      updateInput()
    }
+    cell.selected.afterChange(this) { selected ->
+      this.selected = selected
+    }
+    this.selected = cell.selected.get()
+    cell.executionStatus.afterChange(this) { execution ->
+      updateExecutionStatus(execution.count, execution.status, execution.startTime, execution.endTime)
+    }
+    val executionStatus = cell.executionStatus.get()
+    updateExecutionStatus(executionStatus.count, executionStatus.status, executionStatus.startTime, executionStatus.endTime)
    recreateControllers()
    updateSelection(false)
    updateOutputs()
@@ -211,10 +219,6 @@ class EditorCellView(
    outputs?.onViewportChange()
  }

-  fun setGutterAction(action: AnAction?) {
-    input.setGutterAction(action)
-  }
-
  fun mouseExited() {
    mouseOver = false
    updateFolding()
@@ -336,7 +340,7 @@ class EditorCellView(
    return Rectangle(0, inputBounds.y, editor.contentSize.width, height)
  }

-  fun updateExecutionStatus(executionCount: Int?, progressStatus: ProgressStatus?, startTime: ZonedDateTime?, endTime: ZonedDateTime?) {
+  private fun updateExecutionStatus(executionCount: Int?, progressStatus: ProgressStatus?, startTime: ZonedDateTime?, endTime: ZonedDateTime?) {
    _controllers.filterIsInstance<CellExecutionStatusView>().firstOrNull()
      ?.updateExecutionStatus(executionCount, progressStatus, startTime, endTime)
    input.runCellButton?.updateGutterAction(progressStatus)
--- a/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/EditorCellViewComponent.kt
+++ b/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/EditorCellViewComponent.kt
@@ -6,15 +6,19 @@ import com.intellij.openapi.editor.Inlay
 import com.intellij.openapi.util.Disposer
 import com.intellij.notebooks.visualization.UpdateContext
 import java.awt.Rectangle
+import java.util.Collections

 abstract class EditorCellViewComponent : Disposable {
  protected var parent: EditorCellViewComponent? = null

-  private val children = mutableListOf<EditorCellViewComponent>()
+  private val _children = mutableListOf<EditorCellViewComponent>()
+
+  val children: List<EditorCellViewComponent>
+    get() = Collections.unmodifiableList(_children)

  /* Add automatically registers child disposable. */
  fun add(child: EditorCellViewComponent) {
-    children.add(child)
+    _children.add(child)
    child.parent = this
    Disposer.register(this, child)
  }
@@ -22,14 +26,14 @@ abstract class EditorCellViewComponent : Disposable {
  /* Chile disposable will be automatically disposed. */
  fun remove(child: EditorCellViewComponent) {
    Disposer.dispose(child)
-    children.remove(child)
+    _children.remove(child)
    child.parent = null
  }

  override fun dispose() = Unit

  fun onViewportChange() {
-    children.forEach { it.onViewportChange() }
+    _children.forEach { it.onViewportChange() }
    doViewportChange()
  }

@@ -38,13 +42,13 @@ abstract class EditorCellViewComponent : Disposable {
  abstract fun calculateBounds(): Rectangle

  open fun updateCellFolding(updateContext: UpdateContext) {
-    children.forEach {
+    _children.forEach {
      it.updateCellFolding(updateContext)
    }
  }

  fun getInlays(): Sequence<Inlay<*>> {
-    return doGetInlays() + children.asSequence().flatMap { it.getInlays() }
+    return doGetInlays() + _children.asSequence().flatMap { it.getInlays() }
  }

  open fun doGetInlays(): Sequence<Inlay<*>> {
--- a/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/HasGutterIcon.kt
+++ b/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/HasGutterIcon.kt
@@ -1,7 +0,0 @@
-package com.intellij.notebooks.visualization.ui
-
-import com.intellij.openapi.actionSystem.AnAction
-
-interface HasGutterIcon {
-  fun updateGutterIcons(gutterAction: AnAction?)
-}
--- a/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/InputComponent.kt
+++ b/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/InputComponent.kt
@@ -2,7 +2,7 @@ package com.intellij.notebooks.visualization.ui

 import com.intellij.notebooks.visualization.UpdateContext

-interface InputComponent: HasGutterIcon {
+interface InputComponent {
  fun updateInput(ctx: UpdateContext) {}
  fun updateFolding(ctx: UpdateContext, folded: Boolean)
  fun requestCaret()
--- a/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/TextEditorCellViewComponent.kt
+++ b/notebooks/visualization/src/com/intellij/notebooks/visualization/ui/TextEditorCellViewComponent.kt
@@ -18,6 +18,7 @@ import com.intellij.notebooks.ui.editor.actions.command.mode.NotebookEditorMode
 import com.intellij.notebooks.ui.editor.actions.command.mode.setMode
 import com.intellij.notebooks.visualization.NotebookCellLines
 import com.intellij.notebooks.visualization.UpdateContext
+import com.intellij.openapi.application.runInEdt
 import java.awt.Dimension
 import java.awt.Rectangle
 import java.awt.event.MouseAdapter
@@ -27,7 +28,7 @@ import kotlin.text.lines
 class TextEditorCellViewComponent(
  private val editor: EditorEx,
  private val cell: EditorCell,
-) : EditorCellViewComponent(), HasGutterIcon, InputComponent {
+) : EditorCellViewComponent(), InputComponent {

  private var highlighters: List<RangeHighlighter>? = null

@@ -52,13 +53,17 @@ class TextEditorCellViewComponent(

  init {
    editor.contentComponent.addMouseListener(mouseListener)
+    cell.gutterAction.afterChange(this) { action ->
+      updateGutterIcons(action)
+    }
+    updateGutterIcons(cell.gutterAction.get())
  }

-  override fun updateGutterIcons(gutterAction: AnAction?) {
+  private fun updateGutterIcons(gutterAction: AnAction?) = runInEdt {
    disposeExistingHighlighter()
    if (gutterAction != null) {
      val markupModel = editor.markupModel
-      val interval = safeInterval ?: return
+      val interval = safeInterval ?: return@runInEdt
      val startOffset = editor.document.getLineStartOffset(interval.lines.first)
      val endOffset = editor.document.getLineEndOffset(interval.lines.last)
      val highlighter = markupModel.addRangeHighlighter(
@@ -73,8 +78,7 @@ class TextEditorCellViewComponent(
    }
  }

-  override fun dispose() {
-    super.dispose()
+  override fun dispose() = cell.manager.update { ctx ->
    disposeExistingHighlighter()
    editor.contentComponent.removeMouseListener(mouseListener)
  }