added context methods to the SDK + tests (#179)

* added context methods to the SDK + tests

* lint

* format

* remove body return on delete context

Author: mishushakov

.changeset/open-lamps-drop.md

@@ -0,0 +1,6 @@
+---
+'@e2b/code-interpreter-python': minor
+'@e2b/code-interpreter': minor
+---
+
+added context methods to the sdk

js/src/sandbox.ts

@@ -320,4 +320,96 @@ export class Sandbox extends BaseSandbox {
       throw formatRequestTimeoutError(error)
     }
   }
+
+  /**
+   * Removes a context.
+   *
+   * @param context context to remove.
+   *
+   * @returns void.
+   */
+  async removeCodeContext(context: Context | string): Promise<void> {
+    try {
+      const id = typeof context === 'string' ? context : context.id
+      const res = await fetch(`${this.jupyterUrl}/contexts/${id}`, {
+        method: 'DELETE',
+        headers: {
+          'Content-Type': 'application/json',
+          ...this.connectionConfig.headers,
+        },
+        keepalive: true,
+        signal: this.connectionConfig.getSignal(
+          this.connectionConfig.requestTimeoutMs
+        ),
+      })
+
+      const error = await extractError(res)
+      if (error) {
+        throw error
+      }
+    } catch (error) {
+      throw formatRequestTimeoutError(error)
+    }
+  }
+
+  /**
+   * List all contexts.
+   *
+   * @returns list of contexts.
+   */
+  async listCodeContexts(): Promise<Context[]> {
+    try {
+      const res = await fetch(`${this.jupyterUrl}/contexts`, {
+        method: 'GET',
+        headers: {
+          'Content-Type': 'application/json',
+          ...this.connectionConfig.headers,
+        },
+        keepalive: true,
+        signal: this.connectionConfig.getSignal(
+          this.connectionConfig.requestTimeoutMs
+        ),
+      })
+
+      const error = await extractError(res)
+      if (error) {
+        throw error
+      }
+
+      return await res.json()
+    } catch (error) {
+      throw formatRequestTimeoutError(error)
+    }
+  }
+
+  /**
+   * Restart a context.
+   *
+   * @param context context to restart.
+   *
+   * @returns void.
+   */
+  async restartCodeContext(context: Context | string): Promise<void> {
+    try {
+      const id = typeof context === 'string' ? context : context.id
+      const res = await fetch(`${this.jupyterUrl}/contexts/${id}/restart`, {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+          ...this.connectionConfig.headers,
+        },
+        keepalive: true,
+        signal: this.connectionConfig.getSignal(
+          this.connectionConfig.requestTimeoutMs
+        ),
+      })
+
+      const error = await extractError(res)
+      if (error) {
+        throw error
+      }
+    } catch (error) {
+      throw formatRequestTimeoutError(error)
+    }
+  }
 }

js/tests/contexts.test.ts

@@ -0,0 +1,63 @@
+import { expect } from 'vitest'
+
+import { sandboxTest } from './setup'
+
+sandboxTest('create context with no options', async ({ sandbox }) => {
+  const context = await sandbox.createCodeContext()
+
+  const contexts = await sandbox.listCodeContexts()
+  const lastContext = contexts[contexts.length - 1]
+
+  expect(lastContext.id).toBe(context.id)
+  expect(lastContext.language).toBe(context.language)
+  expect(lastContext.cwd).toBe(context.cwd)
+})
+
+sandboxTest('create context with options', async ({ sandbox }) => {
+  const context = await sandbox.createCodeContext({
+    language: 'python',
+    cwd: '/root',
+  })
+
+  const contexts = await sandbox.listCodeContexts()
+  const lastContext = contexts[contexts.length - 1]
+
+  expect(lastContext.id).toBe(context.id)
+  expect(lastContext.language).toBe(context.language)
+  expect(lastContext.cwd).toBe(context.cwd)
+})
+
+sandboxTest('remove context', async ({ sandbox }) => {
+  const context = await sandbox.createCodeContext()
+
+  await sandbox.removeCodeContext(context.id)
+  const contexts = await sandbox.listCodeContexts()
+
+  expect(contexts.map((context) => context.id)).not.toContain(context.id)
+})
+
+sandboxTest('list contexts', async ({ sandbox }) => {
+  const contexts = await sandbox.listCodeContexts()
+
+  // default contexts should include python and javascript
+  expect(contexts.map((context) => context.language)).toContain('python')
+  expect(contexts.map((context) => context.language)).toContain('javascript')
+})
+
+sandboxTest('restart context', async ({ sandbox }) => {
+  const context = await sandbox.createCodeContext()
+
+  // set a variable in the context
+  await sandbox.runCode('x = 1', { context: context })
+
+  // restart the context
+  await sandbox.restartCodeContext(context.id)
+
+  // check that the variable no longer exists
+  const execution = await sandbox.runCode('x', { context: context })
+
+  // check for an NameError with message "name 'x' is not defined"
+  expect(execution.error).toBeDefined()
+  expect(execution.error?.name).toBe('NameError')
+  expect(execution.error?.value).toBe("name 'x' is not defined")
+})

python/e2b_code_interpreter/code_interpreter_async.py

@@ -1,7 +1,7 @@
 import logging
 import httpx
 
-from typing import Optional, Dict, overload, Union, Literal
+from typing import Optional, Dict, overload, Union, Literal, List
 from httpx import AsyncClient
 
 from e2b import (
@@ -273,3 +273,89 @@ async def create_code_context(
             return Context.from_json(data)
         except httpx.TimeoutException:
             raise format_request_timeout_error()
+
+    async def remove_code_context(
+        self,
+        context: Union[Context, str],
+    ) -> None:
+        """
+        Removes a context.
+
+        :param context: Context to remove. Can be a Context object or a context ID string.
+
+        :return: None
+        """
+        context_id = context.id if isinstance(context, Context) else context
+
+        headers: Dict[str, str] = {}
+        if self._envd_access_token:
+            headers = {"X-Access-Token": self._envd_access_token}
+
+        try:
+            response = await self._client.delete(
+                f"{self._jupyter_url}/contexts/{context_id}",
+                headers=headers,
+                timeout=self.connection_config.request_timeout,
+            )
+
+            err = await aextract_exception(response)
+            if err:
+                raise err
+        except httpx.TimeoutException:
+            raise format_request_timeout_error()
+
+    async def list_code_contexts(self) -> List[Context]:
+        """
+        List all contexts.
+
+        :return: List of contexts.
+        """
+        headers: Dict[str, str] = {}
+        if self._envd_access_token:
+            headers = {"X-Access-Token": self._envd_access_token}
+
+        try:
+            response = await self._client.get(
+                f"{self._jupyter_url}/contexts",
+                headers=headers,
+                timeout=self.connection_config.request_timeout,
+            )
+
+            err = await aextract_exception(response)
+            if err:
+                raise err
+
+            data = response.json()
+            return [Context.from_json(context_data) for context_data in data]
+        except httpx.TimeoutException:
+            raise format_request_timeout_error()
+
+    async def restart_code_context(
+        self,
+        context: Union[Context, str],
+    ) -> None:
+        """
+        Restart a context.
+
+        :param context: Context to restart. Can be a Context object or a context ID string.
+
+        :return: None
+        """
+        context_id = context.id if isinstance(context, Context) else context
+
+        headers: Dict[str, str] = {}
+        if self._envd_access_token:
+            headers = {"X-Access-Token": self._envd_access_token}
+
+        try:
+            response = await self._client.post(
+                f"{self._jupyter_url}/contexts/{context_id}/restart",
+                headers=headers,
+                timeout=self.connection_config.request_timeout,
+            )
+
+            err = await aextract_exception(response)
+            if err:
+                raise err
+        except httpx.TimeoutException:
+            raise format_request_timeout_error()

python/e2b_code_interpreter/code_interpreter_sync.py

@@ -1,7 +1,7 @@
 import logging
 import httpx
 
-from typing import Optional, Dict, overload, Literal, Union
+from typing import Optional, Dict, overload, Literal, Union, List
 from httpx import Client
 from e2b import Sandbox as BaseSandbox, InvalidArgumentException
 
@@ -270,3 +270,89 @@ def create_code_context(
             return Context.from_json(data)
         except httpx.TimeoutException:
             raise format_request_timeout_error()
+
+    def remove_code_context(
+        self,
+        context: Union[Context, str],
+    ) -> None:
+        """
+        Removes a context.
+
+        :param context: Context to remove. Can be a Context object or a context ID string.
+
+        :return: None
+        """
+        context_id = context.id if isinstance(context, Context) else context
+
+        headers: Dict[str, str] = {}
+        if self._envd_access_token:
+            headers = {"X-Access-Token": self._envd_access_token}
+
+        try:
+            response = self._client.delete(
+                f"{self._jupyter_url}/contexts/{context_id}",
+                headers=headers,
+                timeout=self.connection_config.request_timeout,
+            )
+
+            err = extract_exception(response)
+            if err:
+                raise err
+        except httpx.TimeoutException:
+            raise format_request_timeout_error()
+
+    def list_code_contexts(self) -> List[Context]:
+        """
+        List all contexts.
+
+        :return: List of contexts.
+        """
+        headers: Dict[str, str] = {}
+        if self._envd_access_token:
+            headers = {"X-Access-Token": self._envd_access_token}
+
+        try:
+            response = self._client.get(
+                f"{self._jupyter_url}/contexts",
+                headers=headers,
+                timeout=self.connection_config.request_timeout,
+            )
+
+            err = extract_exception(response)
+            if err:
+                raise err
+
+            data = response.json()
+            return [Context.from_json(context_data) for context_data in data]
+        except httpx.TimeoutException:
+            raise format_request_timeout_error()
+
+    def restart_code_context(
+        self,
+        context: Union[Context, str],
+    ) -> None:
+        """
+        Restart a context.
+
+        :param context: Context to restart. Can be a Context object or a context ID string.
+
+        :return: None
+        """
+        context_id = context.id if isinstance(context, Context) else context
+
+        headers: Dict[str, str] = {}
+        if self._envd_access_token:
+            headers = {"X-Access-Token": self._envd_access_token}
+
+        try:
+            response = self._client.post(
+                f"{self._jupyter_url}/contexts/{context_id}/restart",
+                headers=headers,
+                timeout=self.connection_config.request_timeout,
+            )
+
+            err = extract_exception(response)
+            if err:
+                raise err
+        except httpx.TimeoutException:
+            raise format_request_timeout_error()