PY-80672 Jupyter: Rewrite AI to self managed controllers

Signed-off-by: Nikita.Ashihmin <nikita.ashihmin@jetbrains.com>
(cherry picked from commit 5c53d822123b4f955b6d88c6fffd342fe8a0d956)

GitOrigin-RevId: 5ff60c807c27441575d7e757c299779f68ca6aad

notebooks/visualization/resources/intellij.notebooks.visualization.xml

@@ -29,6 +29,11 @@
                     interface="com.intellij.notebooks.visualization.controllers.selfUpdate.SelfManagedControllerFactory"
                     dynamic="true"/>
 
+    <extensionPoint qualifiedName="org.jetbrains.plugins.notebooks.editorNotebookEndInlayProvider"
+                    interface="com.intellij.notebooks.visualization.ui.endInlay.EditorNotebookEndInlayProvider"
+                    dynamic="true"/>
+
+
     <extensionPoint qualifiedName="org.jetbrains.plugins.notebooks.inputFactory"
                     interface="com.intellij.notebooks.visualization.NotebookCellInlayController$InputFactory"
                     dynamic="true"/>

notebooks/visualization/src/com/intellij/notebooks/visualization/NotebookBelowLastCellPanel.kt

@@ -3,17 +3,17 @@ package com.intellij.notebooks.visualization
 
 import com.intellij.notebooks.ui.visualization.NotebookEditorAppearanceUtils.isOrdinaryNotebookEditor
 import com.intellij.notebooks.ui.visualization.NotebookUtil.notebookAppearance
-import com.intellij.notebooks.visualization.ui.cellsDnD.DropHighlightableCellPanel
+import com.intellij.notebooks.visualization.ui.cellsDnD.DropHighlightable
 import com.intellij.notebooks.visualization.ui.jupyterToolbars.JupyterAddNewCellToolbar
 import com.intellij.openapi.actionSystem.ActionGroup
 import com.intellij.openapi.actionSystem.ActionManager
 import com.intellij.openapi.editor.impl.EditorImpl
 import com.intellij.util.ui.JBEmptyBorder
-import java.awt.Graphics2D
 import org.intellij.lang.annotations.Language
 import java.awt.FlowLayout
-import javax.swing.JPanel
 import java.awt.Graphics
+import java.awt.Graphics2D
+import javax.swing.JPanel
 
 /**
  * Basically, this panel consists only of
@@ -21,8 +21,8 @@ import java.awt.Graphics
  * * a highlightable border to show drop destination
  */
 class NotebookBelowLastCellPanel(
-  val editor: EditorImpl
-) : JPanel(FlowLayout(FlowLayout.CENTER)), DropHighlightableCellPanel {
+  val editor: EditorImpl,
+) : JPanel(FlowLayout(FlowLayout.CENTER)), DropHighlightable {
 
   private var isHighlighted = false
 

notebooks/visualization/src/com/intellij/notebooks/visualization/NotebookCellInlayManager.kt

@@ -2,20 +2,18 @@ package com.intellij.notebooks.visualization
 
 import com.intellij.ide.ui.LafManagerListener
 import com.intellij.notebooks.ui.bind
-import com.intellij.notebooks.ui.visualization.NotebookUtil.notebookAppearance
-import com.intellij.notebooks.visualization.context.NotebookDataContext.NOTEBOOK_CELL_LINES_INTERVAL
 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.notebooks.visualization.ui.endInlay.EditorNotebookEndInlay
+import com.intellij.notebooks.visualization.ui.endInlay.EditorNotebookEndInlayProvider
 import com.intellij.openapi.Disposable
-import com.intellij.openapi.actionSystem.UiDataProvider
 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.Inlay
 import com.intellij.openapi.editor.colors.EditorColorsListener
 import com.intellij.openapi.editor.colors.EditorColorsManager
 import com.intellij.openapi.editor.event.CaretEvent
@@ -35,7 +33,7 @@ import java.awt.Point
 class NotebookCellInlayManager private constructor(
   val editor: EditorImpl,
   private val shouldCheckInlayOffsets: Boolean,
-  private val notebook: EditorNotebook,
+  val notebook: EditorNotebook,
 ) : Disposable, NotebookIntervalPointerFactory.ChangeListener {
 
   private val notebookCellLines = NotebookCellLines.get(editor)
@@ -44,10 +42,9 @@ class NotebookCellInlayManager private constructor(
 
   val cells: List<EditorCell> get() = notebook.cells
 
-  internal val views: MutableMap<EditorCell, EditorCellView> = mutableMapOf()
+  val endNotebookInlays: List<EditorNotebookEndInlay> = EditorNotebookEndInlayProvider.create(this)
 
-  val belowLastCellPanel: NotebookBelowLastCellPanel = NotebookBelowLastCellPanel(editor)
-  private var belowLastCellInlay: Inlay<*>? = null
+  internal val views: MutableMap<EditorCell, EditorCellView> = mutableMapOf()
 
   private val cellViewEventListeners = EventDispatcher.create(EditorCellViewEventListener::class.java)
 
@@ -128,7 +125,7 @@ class NotebookCellInlayManager private constructor(
 
     setupFoldingListener()
     setupSelectionUI()
-    addBelowLastCellInlay()
+
 
     notebook.addCellEventsListener(this, object : EditorCellEventListener {
       override fun onEditorCellEvents(events: List<EditorCellEvent>) {
@@ -216,22 +213,6 @@ class NotebookCellInlayManager private constructor(
     }
   }
 
-  private fun addBelowLastCellInlay() {  // PY-77218
-    belowLastCellInlay = editor.addComponentInlay(
-      UiDataProvider.wrapComponent(belowLastCellPanel) { sink ->
-        sink[NOTEBOOK_CELL_LINES_INTERVAL] = editor.notebook?.cells?.lastOrNull()?.interval
-      },
-      isRelatedToPrecedingText = true,
-      showAbove = false,
-      priority = editor.notebookAppearance.jupyterBelowLastCellInlayPriority,
-      offset = editor.document.getLineEndOffset((editor.document.lineCount - 1).coerceAtLeast(0))
-    )
-  }
-
-  fun removeBelowLastCellInlay() {
-    belowLastCellInlay?.let { Disposer.dispose(it) }
-    belowLastCellInlay = null
-  }
 
   private fun setupFoldingListener() {
     val foldingModel = editor.foldingModel

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/EditorCellView.kt

@@ -7,7 +7,7 @@ import com.intellij.notebooks.visualization.*
 import com.intellij.notebooks.visualization.NotebookCellInlayController.InputFactory
 import com.intellij.notebooks.visualization.controllers.selfUpdate.SelfManagedCellController
 import com.intellij.notebooks.visualization.controllers.selfUpdate.SelfManagedControllerFactory
-import com.intellij.notebooks.visualization.ui.cellsDnD.DropHighlightableCellPanel
+import com.intellij.notebooks.visualization.ui.cellsDnD.DropHighlightable
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.diagnostic.thisLogger
@@ -59,9 +59,6 @@ class EditorCellView(
   // We are storing last lines range for highlighters to prevent highlighters unnecessary recreation on the same lines.
   private var lastHighLightersLines: IntRange? = null
 
-  //We do not use it
-  var disableActions: Boolean = false
-
   var isUnderDiff: Boolean
     get() = cell.isUnderDiff.get()
     set(value) = cell.isUnderDiff.set(value)
@@ -322,11 +319,11 @@ class EditorCellView(
   }
 
   fun addDropHighlightIfApplicable() {
-    selfManagedControllers.filterIsInstance<DropHighlightableCellPanel>().firstOrNull()?.addDropHighlight()
+    selfManagedControllers.filterIsInstance<DropHighlightable>().firstOrNull()?.addDropHighlight()
   }
 
   fun removeDropHighlightIfPresent() {
-    selfManagedControllers.filterIsInstance<DropHighlightableCellPanel>().firstOrNull()?.removeDropHighlight()
+    selfManagedControllers.filterIsInstance<DropHighlightable>().firstOrNull()?.removeDropHighlight()
   }
 
 

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/cellsDnD/DropHighlightable.kt

@@ -1,7 +1,7 @@
 // Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package com.intellij.notebooks.visualization.ui.cellsDnD
 
-interface DropHighlightableCellPanel {
+interface DropHighlightable {
   fun addDropHighlight()
   fun removeDropHighlight()
 }

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/cellsDnD/EditorCellDragAssistant.kt

@@ -197,9 +197,14 @@ class EditorCellDragAssistant(
     }
   }
 
-  private fun addHighlightAfterLastCell() = inlayManager?.belowLastCellPanel?.addDropHighlight()
+  private fun addHighlightAfterLastCell() = inlayManager?.endNotebookInlays?.filterIsInstance<DropHighlightable>()?.forEach {
+    it.addDropHighlight()
+  }
+
+  private fun removeHighlightAfterLastCell() = inlayManager?.endNotebookInlays?.filterIsInstance<DropHighlightable>()?.forEach {
+    it.removeDropHighlight()
+  }
 
-  private fun removeHighlightAfterLastCell() = inlayManager?.belowLastCellPanel?.removeDropHighlight()
 
   private fun deleteDropIndicatorForTargetCell(cell: EditorCell) = try {
     cell.view?.removeDropHighlightIfPresent()

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/endInlay/EditorNotebookEndInlay.kt

@@ -0,0 +1,6 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.notebooks.visualization.ui.endInlay
+
+import com.intellij.openapi.Disposable
+
+interface EditorNotebookEndInlay : Disposable.Default

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/endInlay/EditorNotebookEndInlayProvider.kt

@@ -0,0 +1,22 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.notebooks.visualization.ui.endInlay
+
+import com.intellij.notebooks.visualization.NotebookCellInlayManager
+import com.intellij.openapi.extensions.ExtensionPointName
+import com.intellij.openapi.util.Disposer
+
+interface EditorNotebookEndInlayProvider {
+  fun create(inlayManager: NotebookCellInlayManager): EditorNotebookEndInlay?
+
+  companion object {
+    private val EP = ExtensionPointName.create<EditorNotebookEndInlayProvider>("org.jetbrains.plugins.notebooks.editorNotebookEndInlayProvider")
+
+    fun create(inlayManager: NotebookCellInlayManager): List<EditorNotebookEndInlay> {
+      return EP.extensionsIfPointIsRegistered.mapNotNull {
+        val endInlay = it.create(inlayManager) ?: return@mapNotNull null
+        Disposer.register(inlayManager, endInlay)
+        endInlay
+      }
+    }
+  }
+}

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/endInlay/addToolbar/EditorNotebookEndAddToolbar.kt

@@ -0,0 +1,41 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.notebooks.visualization.ui.endInlay.addToolbar
+
+import com.intellij.notebooks.ui.visualization.NotebookUtil.notebookAppearance
+import com.intellij.notebooks.visualization.NotebookBelowLastCellPanel
+import com.intellij.notebooks.visualization.NotebookCellInlayManager
+import com.intellij.notebooks.visualization.context.NotebookDataContext.NOTEBOOK_CELL_LINES_INTERVAL
+import com.intellij.notebooks.visualization.ui.addComponentInlay
+import com.intellij.notebooks.visualization.ui.cellsDnD.DropHighlightable
+import com.intellij.notebooks.visualization.ui.endInlay.EditorNotebookEndInlay
+import com.intellij.openapi.actionSystem.UiDataProvider
+import com.intellij.openapi.util.Disposer
+
+class EditorNotebookEndAddToolbar(val inlayManager: NotebookCellInlayManager) : EditorNotebookEndInlay, DropHighlightable {
+  private val editor = inlayManager.editor
+
+  private val belowLastCellPanel: NotebookBelowLastCellPanel = NotebookBelowLastCellPanel(editor)
+
+  init {
+    // PY-77218
+    editor.addComponentInlay(
+      UiDataProvider.wrapComponent(belowLastCellPanel) { sink ->
+        sink[NOTEBOOK_CELL_LINES_INTERVAL] = inlayManager.notebook.cells.lastOrNull()?.interval
+      },
+      isRelatedToPrecedingText = true,
+      showAbove = false,
+      priority = editor.notebookAppearance.jupyterBelowLastCellInlayPriority,
+      offset = editor.document.getLineEndOffset((editor.document.lineCount - 1).coerceAtLeast(0))
+    ).also {
+      Disposer.register(this, it)
+    }
+  }
+
+  override fun addDropHighlight() {
+    belowLastCellPanel.addDropHighlight()
+  }
+
+  override fun removeDropHighlight() {
+    belowLastCellPanel.removeDropHighlight()
+  }
+}

notebooks/visualization/src/com/intellij/notebooks/visualization/ui/endInlay/addToolbar/EditorNotebookEndAddToolbarProvider.kt

@@ -0,0 +1,15 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.notebooks.visualization.ui.endInlay.addToolbar
+
+import com.intellij.notebooks.ui.visualization.NotebookEditorAppearanceUtils.isOrdinaryNotebookEditor
+import com.intellij.notebooks.visualization.NotebookCellInlayManager
+import com.intellij.notebooks.visualization.ui.endInlay.EditorNotebookEndInlay
+import com.intellij.notebooks.visualization.ui.endInlay.EditorNotebookEndInlayProvider
+
+class EditorNotebookEndAddToolbarProvider : EditorNotebookEndInlayProvider {
+  override fun create(inlayManager: NotebookCellInlayManager): EditorNotebookEndInlay? {
+    if (!inlayManager.editor.isOrdinaryNotebookEditor())
+      return null
+    return EditorNotebookEndAddToolbar(inlayManager)
+  }
+}