Make it backward compatible + add execution count

Author: jakubno

js/src/code-interpreter.ts

@@ -1,16 +1,11 @@
 import { ProcessMessage, Sandbox, SandboxOpts } from 'e2b'
 import { Result, JupyterKernelWebSocket, Execution } from './messaging'
-import { createDeferredPromise } from './utils'
+import { createDeferredPromise, id } from './utils'
 
 interface Kernels {
   [kernelID: string]: JupyterKernelWebSocket
 }
 
-export interface CreateKernelProps {
-  path: string
-  kernelName?: string
-}
-
 /**
  * E2B code interpreter sandbox extension.
  */
@@ -101,11 +96,9 @@ export class JupyterExtension {
     resolve: (value: string) => void,
     opts?: { timeout?: number }
   ) {
-    const sessionInfo = (
-      await this.sandbox.filesystem.read('/root/.jupyter/.session_info', opts)
-    )
-    const parsedSessionInfo = JSON.parse(sessionInfo)
-    const kernelID = parsedSessionInfo.kernel.id
+    const kernelID = (
+      await this.sandbox.filesystem.read('/root/.jupyter/kernel_id', opts)
+    ).trim()
     await this.connectToKernelWS(kernelID)
     resolve(kernelID)
   }
@@ -119,13 +112,16 @@ export class JupyterExtension {
    * the retrieval of the necessary WebSocket URL from the kernel's information.
    *
    * @param kernelID The unique identifier of the kernel to connect to.
+   * @param sessionID The unique identifier of the session to connect to.
    * @throws {Error} Throws an error if the connection to the kernel's WebSocket cannot be established.
    */
-  private async connectToKernelWS(kernelID: string) {
+  private async connectToKernelWS(kernelID: string, sessionID?: string) {
     const url = `${this.sandbox.getProtocol('ws')}://${this.sandbox.getHostname(
       8888
     )}/api/kernels/${kernelID}/channels`
-    const ws = new JupyterKernelWebSocket(url)
+
+    sessionID = sessionID || id(16)
+    const ws = new JupyterKernelWebSocket(url, sessionID)
     await ws.connect()
     this.connectedKernels[kernelID] = ws
 
@@ -149,15 +145,15 @@ export class JupyterExtension {
     cwd: string = '/home/user',
     kernelName?: string
   ): Promise<string> {
-    const data: CreateKernelProps = { path: cwd }
+    const data = { path: cwd, kernel: {name: "python3"}, type: "notebook", name: id(16) }
     if (kernelName) {
-      data.kernelName = kernelName
+      data.kernel.name = kernelName
     }
 
     const response = await fetch(
       `${this.sandbox.getProtocol()}://${this.sandbox.getHostname(
         8888
-      )}/api/kernels`,
+      )}/api/sessions`,
       {
         method: 'POST',
         body: JSON.stringify(data)
@@ -168,8 +164,10 @@ export class JupyterExtension {
       throw new Error(`Failed to create kernel: ${response.statusText}`)
     }
 
-    const kernelID = (await response.json()).id
-    await this.connectToKernelWS(kernelID)
+    const sessionInfo = await response.json()
+    const kernelID = sessionInfo.kernel.id
+    await this.connectToKernelWS(kernelID, sessionInfo.id)
+
     return kernelID
   }
 

js/src/messaging.ts

@@ -226,7 +226,11 @@ export class Execution {
     /**
      * An Error object if an error occurred, null otherwise.
      */
-    public error?: ExecutionError
+    public error?: ExecutionError,
+    /**
+     * Execution count of the cell.
+     */
+    public executionCount?: number
   ) { }
 
   /**
@@ -305,7 +309,7 @@ export class JupyterKernelWebSocket {
    * Does not start WebSocket connection!
    * You need to call connect() method first.
    */
-  constructor(private readonly url: string) { }
+  constructor(private readonly url: string, private readonly sessionID: string) { }
 
   // public
   /**
@@ -401,6 +405,7 @@ export class JupyterKernelWebSocket {
         }
       } else if (message.msg_type == 'execute_input') {
         cell.inputAccepted = true
+        cell.execution.executionCount = message.content.execution_count
       } else {
         console.warn('[UNHANDLED MESSAGE TYPE]:', message.msg_type)
       }
@@ -486,12 +491,11 @@ export class JupyterKernelWebSocket {
    * @param code Code to be executed.
    */
   private sendExecuteRequest(msg_id: string, code: string) {
-    const session = id(16)
     return {
       header: {
         msg_id: msg_id,
         username: 'e2b',
-        session: session,
+        session: this.sessionID,
         msg_type: 'execute_request',
         version: '5.3'
       },
@@ -500,7 +504,7 @@ export class JupyterKernelWebSocket {
       content: {
         code: code,
         silent: false,
-        store_history: false,
+        store_history: true,
         user_expressions: {},
         allow_stdin: false
       }

js/tests/executionCount.test.ts

@@ -0,0 +1,15 @@
+import { CodeInterpreter } from '../src'
+
+import { expect, test } from 'vitest'
+
+test('execution count', async () => {
+  const sandbox = await CodeInterpreter.create()
+
+  await sandbox.notebook.execCell('!pwd')
+  const result = await sandbox.notebook.execCell('!pwd')
+
+
+  await sandbox.close()
+
+  expect(result.executionCount).toEqual(2)
+})

python/e2b_code_interpreter/main.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
-import json
 import logging
 import threading
+import uuid
+
 import requests
 
 from concurrent.futures import Future
@@ -94,14 +95,14 @@ def exec_cell(
             logger.debug(f"Creating new websocket connection to kernel {kernel_id}")
             ws = self._connect_to_kernel_ws(kernel_id, timeout=timeout)
 
-        session_id = ws.send_execution_message(code, on_stdout, on_stderr, on_result)
+        message_id = ws.send_execution_message(code, on_stdout, on_stderr, on_result)
         logger.debug(
-            f"Sent execution message to kernel {kernel_id}, session_id: {session_id}"
+            f"Sent execution message to kernel {kernel_id}, message_id: {message_id}"
         )
 
-        result = ws.get_result(session_id, timeout=timeout)
+        result = ws.get_result(message_id, timeout=timeout)
         logger.debug(
-            f"Received result from kernel {kernel_id}, session_id: {session_id}, result: {result}"
+            f"Received result from kernel {kernel_id}, message_id: {message_id}, result: {result}"
         )
 
         return result
@@ -140,24 +141,26 @@ def create_kernel(
         :param timeout: Timeout for the kernel creation request.
         :return: Kernel id of the created kernel
         """
-        data = {"path": cwd}
+        data = {"path": cwd, "kernel": {"name": "python3"}, "type": "notebook", "name": str(uuid.uuid4())}
         if kernel_name:
-            data["kernel_name"] = kernel_name
+            data["kernel"]['name'] = kernel_name
         logger.debug(f"Creating kernel with data: {data}")
 
         response = requests.post(
-            f"{self._sandbox.get_protocol()}://{self._sandbox.get_hostname(8888)}/api/kernels",
+            f"{self._sandbox.get_protocol()}://{self._sandbox.get_hostname(8888)}/api/sessions",
             json=data,
             timeout=timeout,
         )
         if not response.ok:
             raise KernelException(f"Failed to create kernel: {response.text}")
 
-        kernel_id = response.json()["id"]
+        session_data = response.json()
+        kernel_id = session_data["kernel"]["id"]
+        session_id = session_data["id"]
         logger.debug(f"Created kernel {kernel_id}")
 
         threading.Thread(
-            target=self._connect_to_kernel_ws, args=(kernel_id, timeout)
+            target=self._connect_to_kernel_ws, args=(kernel_id, session_id, timeout)
         ).start()
         return kernel_id
 
@@ -244,7 +247,7 @@ def close(self):
             ws.result().close()
 
     def _connect_to_kernel_ws(
-        self, kernel_id: str, timeout: Optional[float] = TIMEOUT
+        self, kernel_id: str, session_id: Optional[str], timeout: Optional[float] = TIMEOUT
     ) -> JupyterKernelWebSocket:
         """
         Establishes a WebSocket connection to a specified Jupyter kernel.
@@ -258,9 +261,13 @@ def _connect_to_kernel_ws(
         future = Future()
         self._connected_kernels[kernel_id] = future
 
+        print(session_id)
+        session_id = session_id or str(uuid.uuid4())
         ws = JupyterKernelWebSocket(
             url=f"{self._sandbox.get_protocol('ws')}://{self._sandbox.get_hostname(8888)}/api/kernels/{kernel_id}/channels",
+            session_id=session_id
         )
+
         ws.connect(timeout=timeout)
         logger.debug(f"Connected to kernel's ({kernel_id}) websocket.")
 
@@ -277,16 +284,13 @@ def _start_connecting_to_default_kernel(
         logger.debug("Starting to connect to the default kernel")
 
         def setup_default_kernel():
-            session_info = self._sandbox.filesystem.read(
-                "/root/.jupyter/.session_info", timeout=timeout
+            kernel_id = self._sandbox.filesystem.read(
+                "/root/.jupyter/kernel_id", timeout=timeout
             )
-            if session_info is None and not self._sandbox.is_open:
-                return
+            kernel_id = kernel_id.strip()
 
-            session_info = json.loads(session_info)
-            kernel_id = session_info["kernel"]["id"]
             logger.debug(f"Default kernel id: {kernel_id}")
-            self._connect_to_kernel_ws(kernel_id, timeout=timeout)
+            self._connect_to_kernel_ws(kernel_id, None, timeout=timeout)
             self._kernel_id_set.set_result(kernel_id)
 
         threading.Thread(target=setup_default_kernel).start()

python/e2b_code_interpreter/messaging.py

@@ -25,6 +25,7 @@ class CellExecution:
     """
 
     input_accepted: bool = False
+
     on_stdout: Optional[Callable[[ProcessMessage], Any]] = None
     on_stderr: Optional[Callable[[ProcessMessage], Any]] = None
     on_result: Optional[Callable[[Result], Any]] = None
@@ -44,8 +45,9 @@ def __init__(
 
 class JupyterKernelWebSocket:
 
-    def __init__(self, url: str):
+    def __init__(self, url: str, session_id: str):
         self.url = url
+        self.session_id = session_id
         self._cells: Dict[str, CellExecution] = {}
         self._waiting_for_replies: Dict[str, DeferredFuture] = {}
         self._queue_in = Queue()
@@ -101,14 +103,13 @@ def connect(self, timeout: float = TIMEOUT):
 
         logger.debug("WebSocket started")
 
-    @staticmethod
-    def _get_execute_request(msg_id: str, code: str) -> str:
+    def _get_execute_request(self, msg_id: str, code: str) -> str:
         return json.dumps(
             {
                 "header": {
                     "msg_id": msg_id,
                     "username": "e2b",
-                    "session": str(uuid.uuid4()),
+                    "session": self.session_id,
                     "msg_type": "execute_request",
                     "version": "5.3",
                 },
@@ -117,7 +118,7 @@ def _get_execute_request(msg_id: str, code: str) -> str:
                 "content": {
                     "code": code,
                     "silent": False,
-                    "store_history": False,
+                    "store_history": True,
                     "user_expressions": {},
                     "allow_stdin": False,
                 },
@@ -237,6 +238,7 @@ def _receive_message(self, data: dict):
 
         elif data["msg_type"] == "execute_input":
             logger.debug(f"Input accepted for {parent_msg_ig}")
+            cell.partial_result.execution_count = data["content"]["execution_count"]
             cell.input_accepted = True
         else:
             logger.warning(f"[UNHANDLED MESSAGE TYPE]: {data['msg_type']}")

python/e2b_code_interpreter/models.py

@@ -222,6 +222,8 @@ class Config:
     "Logs printed to stdout and stderr during execution."
     error: Optional[Error] = None
     "Error object if an error occurred, None otherwise."
+    execution_count: Optional[int] = None
+    "Execution count of the cell."
 
     @property
     def text(self) -> Optional[str]:

python/example.py

@@ -24,6 +24,7 @@
 """
 
 with CodeInterpreter() as sandbox:
+    print(sandbox.id)
     execution = sandbox.notebook.exec_cell(code)
 
 print(execution.results[0].formats())

python/tests/test_execution_count.py

@@ -0,0 +1,12 @@
+from e2b_code_interpreter.main import CodeInterpreter
+
+
+def test_bash():
+    with CodeInterpreter() as sandbox:
+        result = sandbox.notebook.exec_cell("echo 'E2B is awesome!'")
+        print(result.execution_count)
+        result = sandbox.notebook.exec_cell("!pwd")
+        print(result.execution_count)
+        result = sandbox.notebook.exec_cell("!pwd")
+        print(result.execution_count)
+

template/start-up.sh

@@ -23,6 +23,8 @@ function start_jupyter_server() {
 	echo "Kernel created"
 
 	sudo mkdir -p /root/.jupyter
+  kernel_id=$(echo "${response}" | jq -r '.kernel.id')
+	sudo echo "${kernel_id}" | sudo tee /root/.jupyter/kernel_id >/dev/null
 	sudo echo "${response}" | sudo tee /root/.jupyter/.session_info >/dev/null
 	echo "Jupyter Server started"
 }