Merge pull request #17 from e2b-dev/move-start-cmd-to-envd

Move start cmd to envd

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.
  */
@@ -117,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
 
@@ -147,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)
@@ -166,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

@@ -2,6 +2,8 @@
 
 import logging
 import threading
+import uuid
+
 import requests
 
 from concurrent.futures import Future
@@ -91,16 +93,16 @@ def exec_cell(
             ws = ws_future.result(timeout=timeout)
         else:
             logger.debug(f"Creating new websocket connection to kernel {kernel_id}")
-            ws = self._connect_to_kernel_ws(kernel_id, timeout=timeout)
+            ws = self._connect_to_kernel_ws(kernel_id, None, 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
@@ -139,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
 
@@ -186,7 +190,7 @@ def restart_kernel(
         logger.debug(f"Restarted kernel {kernel_id}")
 
         threading.Thread(
-            target=self._connect_to_kernel_ws, args=(kernel_id, timeout)
+            target=self._connect_to_kernel_ws, args=(kernel_id, None, timeout)
         ).start()
 
     def shutdown_kernel(
@@ -243,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.
@@ -257,9 +261,12 @@ def _connect_to_kernel_ws(
         future = Future()
         self._connected_kernels[kernel_id] = future
 
+        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.")
 
@@ -279,12 +286,14 @@ def setup_default_kernel():
             kernel_id = self._sandbox.filesystem.read(
                 "/root/.jupyter/kernel_id", timeout=timeout
             )
+
             if kernel_id is None and not self._sandbox.is_open:
                 return
 
             kernel_id = kernel_id.strip()
+
             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]:
@@ -249,7 +251,7 @@ def serialize_results(results: List[Result]) -> List[Dict[str, str]]:
         serialized = []
         for result in results:
             serialized_dict = {key: result[key] for key in result.formats()}
-            serialized_dict['text'] = result.text
+            serialized_dict["text"] = result.text
             serialized.append(serialized_dict)
         return serialized
 

python/example.py

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

python/poetry.lock

@@ -14,7 +14,7 @@ python = "^3.8"
 
 pydantic = ">1, <3"
 websocket-client = "^1.7.0"
-e2b = ">=0.17.0"
+e2b = ">=0.17.1"
 
 [tool.poetry.group.dev.dependencies]
 black = "^24.3.0"

python/pyproject.toml

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