KTNB-794: Provide not only classpath, but the whole REPL state on ker…

…nel_info_reply

jupyter-lib/api/src/main/kotlin/org/jetbrains/kotlinx/jupyter/api/KernelRunMode.kt

@@ -3,7 +3,7 @@ package org.jetbrains.kotlinx.jupyter.api
 import org.jetbrains.kotlinx.jupyter.util.createDefaultDelegatingClassLoader
 
 /**
- * Represents settings that depend on the environment in which kernel is running
+ * Represents settings that depend on the environment in which the kernel is running
  */
 interface KernelRunMode {
     val name: String
@@ -22,6 +22,8 @@ interface KernelRunMode {
     val isRunInsideIntellijProcess: Boolean
 
     val streamSubstitutionType: StreamSubstitutionType
+
+    fun initializeSession(notebook: Notebook) = Unit
 }
 
 abstract class AbstractKernelRunMode(override val name: String) : KernelRunMode {
@@ -48,4 +50,8 @@ object EmbeddedKernelRunMode : AbstractKernelRunMode("Embedded") {
     override val isRunInsideIntellijProcess: Boolean get() = false
     override val streamSubstitutionType: StreamSubstitutionType
         get() = StreamSubstitutionType.BLOCKING
+
+    override fun initializeSession(notebook: Notebook) {
+        notebook.sessionOptions.serializeScriptData = true
+    }
 }

jupyter-lib/api/src/main/kotlin/org/jetbrains/kotlinx/jupyter/api/Notebook.kt

@@ -27,6 +27,19 @@ interface Notebook {
      */
     val variablesState: Map<String, VariableState>
 
+    /**
+     * Current session classpath
+     */
+    val currentClasspath: List<String>
+
+    /**
+     * Configuration options for the current notebook session.
+     *
+     * Provides various settings that control the behavior of the session, such as
+     * resolving sources, handling multi-platform projects, and serializing script data.
+     */
+    val sessionOptions: SessionOptions
+
     /**
      * Stores info about useful variables in a cell.
      * Key: cellId;
@@ -49,7 +62,7 @@ interface Notebook {
     fun getCell(id: Int): CodeCell
 
     /**
-     * Mapping allowing to get result by execution number
+     * Mapping allowing to get a result by execution number
      */
     @Throws(IndexOutOfBoundsException::class)
     fun getResult(id: Int): Any?
@@ -122,7 +135,7 @@ interface Notebook {
     val kernelRunMode: KernelRunMode
 
     /**
-     * Renderers processor gives an ability to render values and
+     * Renderers processor gives an ability to render values
      * and add new renderers
      */
     val renderersProcessor: RenderersProcessor

jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/messaging/IdeCompatibleMessageRequestProcessor.kt

@@ -220,7 +220,7 @@ open class IdeCompatibleMessageRequestProcessor(
                     ),
                 metadata =
                     Json.encodeToJsonElement(
-                        KernelInfoReplyMetadata(repl.currentClasspath.toList()),
+                        KernelInfoReplyMetadata(repl.currentSessionState),
                     ),
             ),
         )

jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/messaging/MessageTypes.kt

@@ -28,6 +28,7 @@ import kotlinx.serialization.serializer
 import org.jetbrains.kotlinx.jupyter.config.LanguageInfo
 import org.jetbrains.kotlinx.jupyter.exceptions.ReplException
 import org.jetbrains.kotlinx.jupyter.protocol.messageDataJson
+import org.jetbrains.kotlinx.jupyter.repl.EvaluatedSnippetMetadata
 import org.jetbrains.kotlinx.jupyter.util.EMPTY
 import org.jetbrains.kotlinx.jupyter.util.toUpperCaseAsciiOnly
 import java.util.concurrent.ConcurrentHashMap
@@ -363,7 +364,7 @@ class KernelInfoReply(
 
 @Serializable
 class KernelInfoReplyMetadata(
-    val classpath: List<String>,
+    val state: EvaluatedSnippetMetadata,
 )
 
 @Serializable

jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/ClasspathProvider.kt

@@ -0,0 +1,7 @@
+package org.jetbrains.kotlinx.jupyter.repl
+
+import org.jetbrains.kotlinx.jupyter.repl.result.Classpath
+
+fun interface ClasspathProvider {
+    fun provideClasspath(): Classpath
+}

jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/ReplForJupyter.kt

@@ -39,6 +39,8 @@ interface ReplForJupyter {
 
     val options: ReplOptions
 
+    val currentSessionState: EvaluatedSnippetMetadata
+
     val currentClasspath: Collection<String>
 
     val currentClassLoader: ClassLoader

jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/SharedReplContext.kt

@@ -4,6 +4,7 @@ import org.jetbrains.kotlinx.jupyter.api.ExecutionCallback
 import org.jetbrains.kotlinx.jupyter.api.ExtensionsProcessor
 import org.jetbrains.kotlinx.jupyter.api.KernelLoggerFactory
 import org.jetbrains.kotlinx.jupyter.api.Notebook
+import org.jetbrains.kotlinx.jupyter.api.SessionOptions
 import org.jetbrains.kotlinx.jupyter.api.ThrowableRenderersProcessor
 import org.jetbrains.kotlinx.jupyter.api.outputs.DisplayHandler
 import org.jetbrains.kotlinx.jupyter.codegen.ClassAnnotationsProcessor
@@ -43,4 +44,6 @@ data class SharedReplContext(
     val colorSchemeChangeCallbacksProcessor: ColorSchemeChangeCallbacksProcessor,
     val displayHandler: DisplayHandler,
     val inMemoryReplResultsHolder: InMemoryReplResultsHolder,
+    val sessionOptions: SessionOptions,
+    val currentClasspathProvider: ClasspathProvider,
 )

jupyter-lib/shared-compiler/src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/notebook/impl/NotebookImpl.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlinx.jupyter.api.KotlinKernelVersion
 import org.jetbrains.kotlinx.jupyter.api.LibraryLoader
 import org.jetbrains.kotlinx.jupyter.api.MimeTypedResult
 import org.jetbrains.kotlinx.jupyter.api.ResultsAccessor
+import org.jetbrains.kotlinx.jupyter.api.SessionOptions
 import org.jetbrains.kotlinx.jupyter.api.ThrowableRenderersProcessor
 import org.jetbrains.kotlinx.jupyter.api.VariableState
 import org.jetbrains.kotlinx.jupyter.api.libraries.ColorScheme
@@ -86,6 +87,14 @@ class NotebookImpl(
             ?: throw IllegalStateException("Evaluator is not initialized yet")
     }
 
+    override val currentClasspath: List<String>
+        get() = sharedReplContext?.currentClasspathProvider?.provideClasspath() ?: emptyList()
+
+    override val sessionOptions: SessionOptions get() {
+        return sharedReplContext?.sessionOptions
+            ?: throw IllegalStateException("Session options are not initialized yet")
+    }
+
     override val resultsAccessor = ResultsAccessor { getResult(it) }
 
     override fun getCell(id: Int): MutableCodeCell {

src/main/kotlin/org/jetbrains/kotlinx/jupyter/Ikotlin.kt

@@ -3,6 +3,7 @@ package org.jetbrains.kotlinx.jupyter
 import org.jetbrains.kotlinx.jupyter.api.EmbeddedKernelRunMode
 import org.jetbrains.kotlinx.jupyter.api.KernelLoggerFactory
 import org.jetbrains.kotlinx.jupyter.api.KernelRunMode
+import org.jetbrains.kotlinx.jupyter.api.StandaloneKernelRunMode
 import org.jetbrains.kotlinx.jupyter.api.libraries.JupyterConnection
 import org.jetbrains.kotlinx.jupyter.api.libraries.JupyterSocketType
 import org.jetbrains.kotlinx.jupyter.api.libraries.rawMessageCallback
@@ -97,7 +98,7 @@ fun main(vararg args: String) {
         val kernelConfig = kernelArgs.getConfig()
         startKernel(
             loggerFactory,
-            EmbeddedKernelRunMode,
+            StandaloneKernelRunMode,
             kernelConfig,
         )
     } catch (e: Exception) {

src/main/kotlin/org/jetbrains/kotlinx/jupyter/repl/impl/ReplForJupyterImpl.kt

@@ -85,6 +85,7 @@ import org.jetbrains.kotlinx.jupyter.messaging.comms.installCommHandler
 import org.jetbrains.kotlinx.jupyter.messaging.comms.requireUniqueTargets
 import org.jetbrains.kotlinx.jupyter.registerDefaultRenderers
 import org.jetbrains.kotlinx.jupyter.repl.BaseKernelHost
+import org.jetbrains.kotlinx.jupyter.repl.ClasspathProvider
 import org.jetbrains.kotlinx.jupyter.repl.CompletionResult
 import org.jetbrains.kotlinx.jupyter.repl.ContextUpdater
 import org.jetbrains.kotlinx.jupyter.repl.EvalData
@@ -116,6 +117,7 @@ import org.jetbrains.kotlinx.jupyter.repl.result.InternalMetadata
 import org.jetbrains.kotlinx.jupyter.repl.result.InternalMetadataImpl
 import org.jetbrains.kotlinx.jupyter.repl.result.InternalReplResult
 import org.jetbrains.kotlinx.jupyter.repl.result.SerializedCompiledScriptsData
+import org.jetbrains.kotlinx.jupyter.repl.result.buildScriptsData
 import java.io.Closeable
 import java.io.File
 import java.net.URLClassLoader
@@ -261,6 +263,30 @@ class ReplForJupyterImpl(
 
     override val currentClasspath = compilerConfiguration.classpath.map { it.canonicalPath }.toMutableSet()
     private val currentSources = mutableSetOf<String>()
+    private val evaluatedSnippetsMetadata = mutableListOf<InternalMetadata>()
+
+    private val allEvaluatedSnippetsMetadata: InternalMetadata get() {
+        val allCompiledData =
+            buildScriptsData {
+                for (metadata in evaluatedSnippetsMetadata) {
+                    addData(metadata.compiledData)
+                }
+            }
+        val allImports = evaluatedSnippetsMetadata.flatMap { it.newImports }
+
+        return InternalMetadataImpl(
+            allCompiledData,
+            allImports,
+        )
+    }
+
+    override val currentSessionState: EvaluatedSnippetMetadata get() {
+        return EvaluatedSnippetMetadata(
+            currentClasspath.toList(),
+            currentSources.toList(),
+            allEvaluatedSnippetsMetadata,
+        )
+    }
 
     private val evaluatorConfiguration =
         ScriptEvaluationConfiguration {
@@ -355,6 +381,8 @@ class ReplForJupyterImpl(
     private val afterCellExecutionsProcessor = AfterCellExecutionsProcessor(loggerFactory)
     private val shutdownExecutionsProcessor = ShutdownExecutionsProcessor(loggerFactory)
 
+    private val classpathProvider = ClasspathProvider { currentClasspath.toList() }
+
     override fun checkComplete(code: String) = jupyterCompiler.checkComplete(code)
 
     internal val sharedContext =
@@ -381,8 +409,11 @@ class ReplForJupyterImpl(
             colorSchemeChangeCallbacksProcessor,
             displayHandler,
             inMemoryReplResultsHolder,
+            sessionOptions,
+            classpathProvider,
         ).also {
             notebook.sharedReplContext = it
+            kernelRunMode.initializeSession(notebook)
             commHandlers.requireUniqueTargets()
             commHandlers.forEach { handler -> installCommHandler(handler) }
         }
@@ -582,6 +613,7 @@ class ReplForJupyterImpl(
         }
 
         val metadata = InternalMetadataImpl(compiledData, newImports)
+        evaluatedSnippetsMetadata.add(metadata)
         return if (throwable != null) {
             InternalReplResult.Error(throwable, metadata)
         } else {

src/test/kotlin/org/jetbrains/kotlinx/jupyter/test/TestUtil.kt

@@ -28,6 +28,7 @@ import org.jetbrains.kotlinx.jupyter.api.MimeTypes
 import org.jetbrains.kotlinx.jupyter.api.Notebook
 import org.jetbrains.kotlinx.jupyter.api.RenderersProcessor
 import org.jetbrains.kotlinx.jupyter.api.ResultsAccessor
+import org.jetbrains.kotlinx.jupyter.api.SessionOptions
 import org.jetbrains.kotlinx.jupyter.api.StandaloneKernelRunMode
 import org.jetbrains.kotlinx.jupyter.api.TextRenderersProcessor
 import org.jetbrains.kotlinx.jupyter.api.ThrowableRenderersProcessor
@@ -279,6 +280,11 @@ object NotebookMock : Notebook {
     override val variablesState = mutableMapOf<String, VariableStateImpl>()
     override val cellVariables = mapOf<Int, Set<String>>()
     override val resultsAccessor = ResultsAccessor { getResult(it) }
+    override val currentClasspath: List<String>
+        get() = notImplemented()
+
+    override val sessionOptions: SessionOptions
+        get() = notImplemented()
 
     override val loggerFactory: KernelLoggerFactory get() = DefaultKernelLoggerFactory