From 6a339b3131ae43413c3b9df33fa300d029ee44ae Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Tue, 25 Nov 2025 14:56:10 +0100
Subject: [PATCH 1/9] added context methods to the SDK + tests

---
 .changeset/open-lamps-drop.md                 |  6 ++
 js/src/messaging.ts                           |  5 +-
 js/src/sandbox.ts                             | 90 +++++++++++++++++++
 js/tests/contexts.test.ts                     | 40 +++++++++
 .../code_interpreter_async.py                 | 88 +++++++++++++++++-
 .../code_interpreter_sync.py                  | 88 +++++++++++++++++-
 python/tests/async/test_async_contexts.py     | 39 ++++++++
 python/tests/sync/test_contexts.py            | 39 ++++++++
 8 files changed, 389 insertions(+), 6 deletions(-)
 create mode 100644 .changeset/open-lamps-drop.md
 create mode 100644 js/tests/contexts.test.ts
 create mode 100644 python/tests/async/test_async_contexts.py
 create mode 100644 python/tests/sync/test_contexts.py

diff --git a/.changeset/open-lamps-drop.md b/.changeset/open-lamps-drop.md
new file mode 100644
index 00000000..2fc5b72a
--- /dev/null
+++ b/.changeset/open-lamps-drop.md
@@ -0,0 +1,6 @@
+---
+'@e2b/code-interpreter-python': minor
+'@e2b/code-interpreter': minor
+---
+
+added context methods to the sdk
diff --git a/js/src/messaging.ts b/js/src/messaging.ts
index 5f1e4e2c..11b1a981 100644
--- a/js/src/messaging.ts
+++ b/js/src/messaging.ts
@@ -145,10 +145,7 @@ export class Result {
 
   readonly raw: RawData
 
-  constructor(
-    rawData: RawData,
-    public readonly isMainResult: boolean
-  ) {
+  constructor(rawData: RawData, public readonly isMainResult: boolean) {
     const data = { ...rawData }
     delete data['type']
     delete data['is_main_result']
diff --git a/js/src/sandbox.ts b/js/src/sandbox.ts
index 4272d846..d985264a 100644
--- a/js/src/sandbox.ts
+++ b/js/src/sandbox.ts
@@ -320,4 +320,94 @@ 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
+      }
+
+      return await res.json()
+    } 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,
+        },
+      })
+
+      const error = await extractError(res)
+      if (error) {
+        throw error
+      }
+    } catch (error) {
+      throw formatRequestTimeoutError(error)
+    }
+  }
 }
diff --git a/js/tests/contexts.test.ts b/js/tests/contexts.test.ts
new file mode 100644
index 00000000..bb36c9da
--- /dev/null
+++ b/js/tests/contexts.test.ts
@@ -0,0 +1,40 @@
+import { expect } from 'vitest'
+
+import { isDebug, sandboxTest } from './setup'
+
+sandboxTest('create context with no options', async ({ sandbox }) => {
+  const context = await sandbox.createCodeContext()
+
+  expect(context.id).toBeDefined()
+  expect(context.language).toBe('python')
+  expect(context.cwd).toBe('/home/user')
+})
+
+sandboxTest('create context with options', async ({ sandbox }) => {
+  const context = await sandbox.createCodeContext({
+    language: 'python',
+    cwd: '/home/user/test',
+  })
+
+  expect(context.id).toBeDefined()
+  expect(context.language).toBe('python')
+  expect(context.cwd).toBe('/home/user/test')
+})
+
+sandboxTest('remove context', async ({ sandbox }) => {
+  const context = await sandbox.createCodeContext()
+
+  await sandbox.removeCodeContext(context.id)
+})
+
+sandboxTest('list contexts', async ({ sandbox }) => {
+  const contexts = await sandbox.listCodeContexts()
+
+  expect(contexts.length).toBeGreaterThan(0)
+})
+
+sandboxTest('restart context', async ({ sandbox }) => {
+  const context = await sandbox.createCodeContext()
+
+  await sandbox.restartCodeContext(context.id)
+})
diff --git a/python/e2b_code_interpreter/code_interpreter_async.py b/python/e2b_code_interpreter/code_interpreter_async.py
index 789d6a70..98f59dd0 100644
--- a/python/e2b_code_interpreter/code_interpreter_async.py
+++ b/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()
diff --git a/python/e2b_code_interpreter/code_interpreter_sync.py b/python/e2b_code_interpreter/code_interpreter_sync.py
index 6cf56c11..67492398 100644
--- a/python/e2b_code_interpreter/code_interpreter_sync.py
+++ b/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()
diff --git a/python/tests/async/test_async_contexts.py b/python/tests/async/test_async_contexts.py
new file mode 100644
index 00000000..4ae6b0fe
--- /dev/null
+++ b/python/tests/async/test_async_contexts.py
@@ -0,0 +1,39 @@
+from e2b_code_interpreter.code_interpreter_async import AsyncSandbox
+
+
+async def test_create_context_with_no_options(async_sandbox: AsyncSandbox):
+    context = await async_sandbox.create_code_context()
+
+    assert context.id is not None
+    assert context.language == "python"
+    assert context.cwd == "/home/user"
+
+
+async def test_create_context_with_options(async_sandbox: AsyncSandbox):
+    context = await async_sandbox.create_code_context(
+        language="python",
+        cwd="/home/user/test",
+    )
+
+    assert context.id is not None
+    assert context.language == "python"
+    assert context.cwd == "/home/user/test"
+
+
+async def test_remove_context(async_sandbox: AsyncSandbox):
+    context = await async_sandbox.create_code_context()
+
+    await async_sandbox.remove_code_context(context.id)
+
+
+async def test_list_contexts(async_sandbox: AsyncSandbox):
+    contexts = await async_sandbox.list_code_contexts()
+
+    assert len(contexts) > 0
+
+
+async def test_restart_context(async_sandbox: AsyncSandbox):
+    context = await async_sandbox.create_code_context()
+
+    await async_sandbox.restart_code_context(context.id)
+
diff --git a/python/tests/sync/test_contexts.py b/python/tests/sync/test_contexts.py
new file mode 100644
index 00000000..8731e409
--- /dev/null
+++ b/python/tests/sync/test_contexts.py
@@ -0,0 +1,39 @@
+from e2b_code_interpreter.code_interpreter_sync import Sandbox
+
+
+def test_create_context_with_no_options(sandbox: Sandbox):
+    context = sandbox.create_code_context()
+
+    assert context.id is not None
+    assert context.language == "python"
+    assert context.cwd == "/home/user"
+
+
+def test_create_context_with_options(sandbox: Sandbox):
+    context = sandbox.create_code_context(
+        language="python",
+        cwd="/home/user/test",
+    )
+
+    assert context.id is not None
+    assert context.language == "python"
+    assert context.cwd == "/home/user/test"
+
+
+def test_remove_context(sandbox: Sandbox):
+    context = sandbox.create_code_context()
+
+    sandbox.remove_code_context(context.id)
+
+
+def test_list_contexts(sandbox: Sandbox):
+    contexts = sandbox.list_code_contexts()
+
+    assert len(contexts) > 0
+
+
+def test_restart_context(sandbox: Sandbox):
+    context = sandbox.create_code_context()
+
+    sandbox.restart_code_context(context.id)
+

From 870581c44c39d896482a06349e93cbbebb4dc041 Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Tue, 25 Nov 2025 14:57:29 +0100
Subject: [PATCH 2/9] lint

---
 js/tests/contexts.test.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/js/tests/contexts.test.ts b/js/tests/contexts.test.ts
index bb36c9da..736d21d1 100644
--- a/js/tests/contexts.test.ts
+++ b/js/tests/contexts.test.ts
@@ -1,6 +1,6 @@
 import { expect } from 'vitest'
 
-import { isDebug, sandboxTest } from './setup'
+import { sandboxTest } from './setup'
 
 sandboxTest('create context with no options', async ({ sandbox }) => {
   const context = await sandbox.createCodeContext()

From 8f2b1c94051f6a6239947f3bb7069be43b2e4b18 Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Tue, 25 Nov 2025 15:01:08 +0100
Subject: [PATCH 3/9] format

---
 js/src/messaging.ts                       | 5 ++++-
 python/tests/async/test_async_contexts.py | 1 -
 python/tests/sync/test_contexts.py        | 1 -
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/js/src/messaging.ts b/js/src/messaging.ts
index 11b1a981..5f1e4e2c 100644
--- a/js/src/messaging.ts
+++ b/js/src/messaging.ts
@@ -145,7 +145,10 @@ export class Result {
 
   readonly raw: RawData
 
-  constructor(rawData: RawData, public readonly isMainResult: boolean) {
+  constructor(
+    rawData: RawData,
+    public readonly isMainResult: boolean
+  ) {
     const data = { ...rawData }
     delete data['type']
     delete data['is_main_result']
diff --git a/python/tests/async/test_async_contexts.py b/python/tests/async/test_async_contexts.py
index 4ae6b0fe..b51d75eb 100644
--- a/python/tests/async/test_async_contexts.py
+++ b/python/tests/async/test_async_contexts.py
@@ -36,4 +36,3 @@ async def test_restart_context(async_sandbox: AsyncSandbox):
     context = await async_sandbox.create_code_context()
 
     await async_sandbox.restart_code_context(context.id)
-
diff --git a/python/tests/sync/test_contexts.py b/python/tests/sync/test_contexts.py
index 8731e409..a97269c9 100644
--- a/python/tests/sync/test_contexts.py
+++ b/python/tests/sync/test_contexts.py
@@ -36,4 +36,3 @@ def test_restart_context(sandbox: Sandbox):
     context = sandbox.create_code_context()
 
     sandbox.restart_code_context(context.id)
-

From 699cb0516d3b0feac84e5040d89d4ad7d0d1fb81 Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Tue, 25 Nov 2025 15:02:01 +0100
Subject: [PATCH 4/9] remove body return on delete context

---
 js/src/sandbox.ts | 2 --
 1 file changed, 2 deletions(-)

diff --git a/js/src/sandbox.ts b/js/src/sandbox.ts
index d985264a..b093d537 100644
--- a/js/src/sandbox.ts
+++ b/js/src/sandbox.ts
@@ -347,8 +347,6 @@ export class Sandbox extends BaseSandbox {
       if (error) {
         throw error
       }
-
-      return await res.json()
     } catch (error) {
       throw formatRequestTimeoutError(error)
     }

From 325cebcb2cc3a15ffbcb8ab3da0ce8fa9f5d5e7a Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Tue, 25 Nov 2025 19:17:17 +0100
Subject: [PATCH 5/9] updated tests

---
 js/tests/contexts.test.ts                 | 27 ++++++++++++++++-------
 python/tests/async/test_async_contexts.py | 23 +++++++++++++------
 python/tests/sync/test_contexts.py        | 23 +++++++++++++------
 3 files changed, 51 insertions(+), 22 deletions(-)

diff --git a/js/tests/contexts.test.ts b/js/tests/contexts.test.ts
index 736d21d1..abcdd513 100644
--- a/js/tests/contexts.test.ts
+++ b/js/tests/contexts.test.ts
@@ -5,32 +5,43 @@ import { sandboxTest } from './setup'
 sandboxTest('create context with no options', async ({ sandbox }) => {
   const context = await sandbox.createCodeContext()
 
-  expect(context.id).toBeDefined()
-  expect(context.language).toBe('python')
-  expect(context.cwd).toBe('/home/user')
+  const contexts = await sandbox.listCodeContexts()
+  const lastContext = contexts[contexts.length - 1]
+
+  expect(lastContext.id).toBeDefined()
+  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: '/home/user/test',
+    cwd: '/root',
   })
 
-  expect(context.id).toBeDefined()
-  expect(context.language).toBe('python')
-  expect(context.cwd).toBe('/home/user/test')
+  const contexts = await sandbox.listCodeContexts()
+  const lastContext = contexts[contexts.length - 1]
+
+  expect(lastContext.id).toBeDefined()
+  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()
 
-  expect(contexts.length).toBeGreaterThan(0)
+  // 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 }) => {
diff --git a/python/tests/async/test_async_contexts.py b/python/tests/async/test_async_contexts.py
index b51d75eb..c8ee1332 100644
--- a/python/tests/async/test_async_contexts.py
+++ b/python/tests/async/test_async_contexts.py
@@ -4,9 +4,12 @@
 async def test_create_context_with_no_options(async_sandbox: AsyncSandbox):
     context = await async_sandbox.create_code_context()
 
-    assert context.id is not None
-    assert context.language == "python"
-    assert context.cwd == "/home/user"
+    contexts = await async_sandbox.list_code_contexts()
+    last_context = contexts[-1]
+
+    assert last_context.id is not None
+    assert last_context.language == context.language
+    assert last_context.cwd == context.cwd
 
 
 async def test_create_context_with_options(async_sandbox: AsyncSandbox):
@@ -15,9 +18,12 @@ async def test_create_context_with_options(async_sandbox: AsyncSandbox):
         cwd="/home/user/test",
     )
 
-    assert context.id is not None
-    assert context.language == "python"
-    assert context.cwd == "/home/user/test"
+    contexts = await async_sandbox.list_code_contexts()
+    last_context = contexts[-1]
+
+    assert last_context.id is not None
+    assert last_context.language == context.language
+    assert last_context.cwd == context.cwd
 
 
 async def test_remove_context(async_sandbox: AsyncSandbox):
@@ -29,7 +35,10 @@ async def test_remove_context(async_sandbox: AsyncSandbox):
 async def test_list_contexts(async_sandbox: AsyncSandbox):
     contexts = await async_sandbox.list_code_contexts()
 
-    assert len(contexts) > 0
+    # default contexts should include python and javascript
+    languages = [context.language for context in contexts]
+    assert "python" in languages
+    assert "javascript" in languages
 
 
 async def test_restart_context(async_sandbox: AsyncSandbox):
diff --git a/python/tests/sync/test_contexts.py b/python/tests/sync/test_contexts.py
index a97269c9..4a282ddc 100644
--- a/python/tests/sync/test_contexts.py
+++ b/python/tests/sync/test_contexts.py
@@ -4,9 +4,12 @@
 def test_create_context_with_no_options(sandbox: Sandbox):
     context = sandbox.create_code_context()
 
-    assert context.id is not None
-    assert context.language == "python"
-    assert context.cwd == "/home/user"
+    contexts = sandbox.list_code_contexts()
+    last_context = contexts[-1]
+
+    assert last_context.id is not None
+    assert last_context.language == context.language
+    assert last_context.cwd == context.cwd
 
 
 def test_create_context_with_options(sandbox: Sandbox):
@@ -15,9 +18,12 @@ def test_create_context_with_options(sandbox: Sandbox):
         cwd="/home/user/test",
     )
 
-    assert context.id is not None
-    assert context.language == "python"
-    assert context.cwd == "/home/user/test"
+    contexts = sandbox.list_code_contexts()
+    last_context = contexts[-1]
+
+    assert last_context.id is not None
+    assert last_context.language == context.language
+    assert last_context.cwd == context.cwd
 
 
 def test_remove_context(sandbox: Sandbox):
@@ -29,7 +35,10 @@ def test_remove_context(sandbox: Sandbox):
 def test_list_contexts(sandbox: Sandbox):
     contexts = sandbox.list_code_contexts()
 
-    assert len(contexts) > 0
+    # default contexts should include python and javascript
+    languages = [context.language for context in contexts]
+    assert "python" in languages
+    assert "javascript" in languages
 
 
 def test_restart_context(sandbox: Sandbox):

From c6eeca7b224a02d2ad7a134439bf64c23d091564 Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Tue, 25 Nov 2025 19:25:51 +0100
Subject: [PATCH 6/9] updated tests

---
 js/src/sandbox.ts                         |  4 ++++
 js/tests/contexts.test.ts                 | 12 ++++++++++++
 python/tests/async/test_async_contexts.py | 15 +++++++++++++++
 python/tests/sync/test_contexts.py        | 15 +++++++++++++++
 4 files changed, 46 insertions(+)

diff --git a/js/src/sandbox.ts b/js/src/sandbox.ts
index b093d537..0a87d20f 100644
--- a/js/src/sandbox.ts
+++ b/js/src/sandbox.ts
@@ -398,6 +398,10 @@ export class Sandbox extends BaseSandbox {
           'Content-Type': 'application/json',
           ...this.connectionConfig.headers,
         },
+        keepalive: true,
+        signal: this.connectionConfig.getSignal(
+          this.connectionConfig.requestTimeoutMs
+        ),
       })
 
       const error = await extractError(res)
diff --git a/js/tests/contexts.test.ts b/js/tests/contexts.test.ts
index abcdd513..63697bff 100644
--- a/js/tests/contexts.test.ts
+++ b/js/tests/contexts.test.ts
@@ -47,5 +47,17 @@ sandboxTest('list contexts', async ({ sandbox }) => {
 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")
 })
diff --git a/python/tests/async/test_async_contexts.py b/python/tests/async/test_async_contexts.py
index c8ee1332..1a85b1d8 100644
--- a/python/tests/async/test_async_contexts.py
+++ b/python/tests/async/test_async_contexts.py
@@ -31,6 +31,9 @@ async def test_remove_context(async_sandbox: AsyncSandbox):
 
     await async_sandbox.remove_code_context(context.id)
 
+    contexts = await async_sandbox.list_code_contexts()
+    assert context.id not in [ctx.id for ctx in contexts]
+
 
 async def test_list_contexts(async_sandbox: AsyncSandbox):
     contexts = await async_sandbox.list_code_contexts()
@@ -44,4 +47,16 @@ async def test_list_contexts(async_sandbox: AsyncSandbox):
 async def test_restart_context(async_sandbox: AsyncSandbox):
     context = await async_sandbox.create_code_context()
 
+    # set a variable in the context
+    await async_sandbox.run_code("x = 1", context=context)
+
+    # restart the context
     await async_sandbox.restart_code_context(context.id)
+
+    # check that the variable no longer exists
+    execution = await async_sandbox.run_code("x", context=context)
+
+    # check for a NameError with message "name 'x' is not defined"
+    assert execution.error is not None
+    assert execution.error.name == "NameError"
+    assert execution.error.value == "name 'x' is not defined"
diff --git a/python/tests/sync/test_contexts.py b/python/tests/sync/test_contexts.py
index 4a282ddc..e854abdd 100644
--- a/python/tests/sync/test_contexts.py
+++ b/python/tests/sync/test_contexts.py
@@ -31,6 +31,9 @@ def test_remove_context(sandbox: Sandbox):
 
     sandbox.remove_code_context(context.id)
 
+    contexts = sandbox.list_code_contexts()
+    assert context.id not in [ctx.id for ctx in contexts]
+
 
 def test_list_contexts(sandbox: Sandbox):
     contexts = sandbox.list_code_contexts()
@@ -44,4 +47,16 @@ def test_list_contexts(sandbox: Sandbox):
 def test_restart_context(sandbox: Sandbox):
     context = sandbox.create_code_context()
 
+    # set a variable in the context
+    sandbox.run_code("x = 1", context=context)
+
+    # restart the context
     sandbox.restart_code_context(context.id)
+
+    # check that the variable no longer exists
+    execution = sandbox.run_code("x", context=context)
+
+    # check for a NameError with message "name 'x' is not defined"
+    assert execution.error is not None
+    assert execution.error.name == "NameError"
+    assert execution.error.value == "name 'x' is not defined"

From dc14e0fef9f8a9aba8949b69fc24672fdb84880b Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Tue, 25 Nov 2025 19:50:04 +0100
Subject: [PATCH 7/9] added 1s delay after context create before list

---
 js/tests/contexts.test.ts                 | 6 ++++++
 python/tests/async/test_async_contexts.py | 9 ++++++++-
 python/tests/sync/test_contexts.py        | 9 ++++++++-
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/js/tests/contexts.test.ts b/js/tests/contexts.test.ts
index 63697bff..60037159 100644
--- a/js/tests/contexts.test.ts
+++ b/js/tests/contexts.test.ts
@@ -5,6 +5,9 @@ import { sandboxTest } from './setup'
 sandboxTest('create context with no options', async ({ sandbox }) => {
   const context = await sandbox.createCodeContext()
 
+  // wait 1 second for the context to be created
+  await new Promise((resolve) => setTimeout(resolve, 1000))
+
   const contexts = await sandbox.listCodeContexts()
   const lastContext = contexts[contexts.length - 1]
 
@@ -19,6 +22,9 @@ sandboxTest('create context with options', async ({ sandbox }) => {
     cwd: '/root',
   })
 
+  // wait 1 second for the context to be created
+  await new Promise((resolve) => setTimeout(resolve, 1000))
+
   const contexts = await sandbox.listCodeContexts()
   const lastContext = contexts[contexts.length - 1]
 
diff --git a/python/tests/async/test_async_contexts.py b/python/tests/async/test_async_contexts.py
index 1a85b1d8..4dc984b2 100644
--- a/python/tests/async/test_async_contexts.py
+++ b/python/tests/async/test_async_contexts.py
@@ -1,9 +1,13 @@
 from e2b_code_interpreter.code_interpreter_async import AsyncSandbox
+import asyncio
 
 
 async def test_create_context_with_no_options(async_sandbox: AsyncSandbox):
     context = await async_sandbox.create_code_context()
 
+    # wait 1 second for the context to be created
+    await asyncio.sleep(1)
+
     contexts = await async_sandbox.list_code_contexts()
     last_context = contexts[-1]
 
@@ -15,9 +19,12 @@ async def test_create_context_with_no_options(async_sandbox: AsyncSandbox):
 async def test_create_context_with_options(async_sandbox: AsyncSandbox):
     context = await async_sandbox.create_code_context(
         language="python",
-        cwd="/home/user/test",
+        cwd="/root",
     )
 
+    # wait 1 second for the context to be created
+    await asyncio.sleep(1)
+
     contexts = await async_sandbox.list_code_contexts()
     last_context = contexts[-1]
 
diff --git a/python/tests/sync/test_contexts.py b/python/tests/sync/test_contexts.py
index e854abdd..2b155685 100644
--- a/python/tests/sync/test_contexts.py
+++ b/python/tests/sync/test_contexts.py
@@ -1,9 +1,13 @@
 from e2b_code_interpreter.code_interpreter_sync import Sandbox
+import time
 
 
 def test_create_context_with_no_options(sandbox: Sandbox):
     context = sandbox.create_code_context()
 
+    # wait 1 second for the context to be created
+    time.sleep(1)
+
     contexts = sandbox.list_code_contexts()
     last_context = contexts[-1]
 
@@ -15,9 +19,12 @@ def test_create_context_with_no_options(sandbox: Sandbox):
 def test_create_context_with_options(sandbox: Sandbox):
     context = sandbox.create_code_context(
         language="python",
-        cwd="/home/user/test",
+        cwd="/root",
     )
 
+    # wait 1 second for the context to be created
+    time.sleep(1)
+
     contexts = sandbox.list_code_contexts()
     last_context = contexts[-1]
 

From 7306a5a79b9915c6fd5790e679c1cb36529d20f1 Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Tue, 25 Nov 2025 19:51:17 +0100
Subject: [PATCH 8/9] verify context id

---
 js/tests/contexts.test.ts                 | 4 ++--
 python/tests/async/test_async_contexts.py | 4 ++--
 python/tests/sync/test_contexts.py        | 4 ++--
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/js/tests/contexts.test.ts b/js/tests/contexts.test.ts
index 60037159..d74114e6 100644
--- a/js/tests/contexts.test.ts
+++ b/js/tests/contexts.test.ts
@@ -11,7 +11,7 @@ sandboxTest('create context with no options', async ({ sandbox }) => {
   const contexts = await sandbox.listCodeContexts()
   const lastContext = contexts[contexts.length - 1]
 
-  expect(lastContext.id).toBeDefined()
+  expect(lastContext.id).toBe(context.id)
   expect(lastContext.language).toBe(context.language)
   expect(lastContext.cwd).toBe(context.cwd)
 })
@@ -28,7 +28,7 @@ sandboxTest('create context with options', async ({ sandbox }) => {
   const contexts = await sandbox.listCodeContexts()
   const lastContext = contexts[contexts.length - 1]
 
-  expect(lastContext.id).toBeDefined()
+  expect(lastContext.id).toBe(context.id)
   expect(lastContext.language).toBe(context.language)
   expect(lastContext.cwd).toBe(context.cwd)
 })
diff --git a/python/tests/async/test_async_contexts.py b/python/tests/async/test_async_contexts.py
index 4dc984b2..28462147 100644
--- a/python/tests/async/test_async_contexts.py
+++ b/python/tests/async/test_async_contexts.py
@@ -11,7 +11,7 @@ async def test_create_context_with_no_options(async_sandbox: AsyncSandbox):
     contexts = await async_sandbox.list_code_contexts()
     last_context = contexts[-1]
 
-    assert last_context.id is not None
+    assert last_context.id == context.id
     assert last_context.language == context.language
     assert last_context.cwd == context.cwd
 
@@ -28,7 +28,7 @@ async def test_create_context_with_options(async_sandbox: AsyncSandbox):
     contexts = await async_sandbox.list_code_contexts()
     last_context = contexts[-1]
 
-    assert last_context.id is not None
+    assert last_context.id == context.id
     assert last_context.language == context.language
     assert last_context.cwd == context.cwd
 
diff --git a/python/tests/sync/test_contexts.py b/python/tests/sync/test_contexts.py
index 2b155685..d130569f 100644
--- a/python/tests/sync/test_contexts.py
+++ b/python/tests/sync/test_contexts.py
@@ -11,7 +11,7 @@ def test_create_context_with_no_options(sandbox: Sandbox):
     contexts = sandbox.list_code_contexts()
     last_context = contexts[-1]
 
-    assert last_context.id is not None
+    assert last_context.id == context.id
     assert last_context.language == context.language
     assert last_context.cwd == context.cwd
 
@@ -28,7 +28,7 @@ def test_create_context_with_options(sandbox: Sandbox):
     contexts = sandbox.list_code_contexts()
     last_context = contexts[-1]
 
-    assert last_context.id is not None
+    assert last_context.id == context.id
     assert last_context.language == context.language
     assert last_context.cwd == context.cwd
 

From 51b0d17984aff95c0eda799ab9cb1f15ccb05358 Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Tue, 25 Nov 2025 20:32:30 +0100
Subject: [PATCH 9/9] update /contexts method

---
 js/tests/contexts.test.ts                 |  6 ------
 python/tests/async/test_async_contexts.py |  7 -------
 python/tests/sync/test_contexts.py        |  7 -------
 template/server/main.py                   | 19 +++++++++----------
 4 files changed, 9 insertions(+), 30 deletions(-)

diff --git a/js/tests/contexts.test.ts b/js/tests/contexts.test.ts
index d74114e6..7dda2090 100644
--- a/js/tests/contexts.test.ts
+++ b/js/tests/contexts.test.ts
@@ -5,9 +5,6 @@ import { sandboxTest } from './setup'
 sandboxTest('create context with no options', async ({ sandbox }) => {
   const context = await sandbox.createCodeContext()
 
-  // wait 1 second for the context to be created
-  await new Promise((resolve) => setTimeout(resolve, 1000))
-
   const contexts = await sandbox.listCodeContexts()
   const lastContext = contexts[contexts.length - 1]
 
@@ -22,9 +19,6 @@ sandboxTest('create context with options', async ({ sandbox }) => {
     cwd: '/root',
   })
 
-  // wait 1 second for the context to be created
-  await new Promise((resolve) => setTimeout(resolve, 1000))
-
   const contexts = await sandbox.listCodeContexts()
   const lastContext = contexts[contexts.length - 1]
 
diff --git a/python/tests/async/test_async_contexts.py b/python/tests/async/test_async_contexts.py
index 28462147..6be07734 100644
--- a/python/tests/async/test_async_contexts.py
+++ b/python/tests/async/test_async_contexts.py
@@ -1,13 +1,9 @@
 from e2b_code_interpreter.code_interpreter_async import AsyncSandbox
-import asyncio
 
 
 async def test_create_context_with_no_options(async_sandbox: AsyncSandbox):
     context = await async_sandbox.create_code_context()
 
-    # wait 1 second for the context to be created
-    await asyncio.sleep(1)
-
     contexts = await async_sandbox.list_code_contexts()
     last_context = contexts[-1]
 
@@ -22,9 +18,6 @@ async def test_create_context_with_options(async_sandbox: AsyncSandbox):
         cwd="/root",
     )
 
-    # wait 1 second for the context to be created
-    await asyncio.sleep(1)
-
     contexts = await async_sandbox.list_code_contexts()
     last_context = contexts[-1]
 
diff --git a/python/tests/sync/test_contexts.py b/python/tests/sync/test_contexts.py
index d130569f..a7cbd884 100644
--- a/python/tests/sync/test_contexts.py
+++ b/python/tests/sync/test_contexts.py
@@ -1,13 +1,9 @@
 from e2b_code_interpreter.code_interpreter_sync import Sandbox
-import time
 
 
 def test_create_context_with_no_options(sandbox: Sandbox):
     context = sandbox.create_code_context()
 
-    # wait 1 second for the context to be created
-    time.sleep(1)
-
     contexts = sandbox.list_code_contexts()
     last_context = contexts[-1]
 
@@ -22,9 +18,6 @@ def test_create_context_with_options(sandbox: Sandbox):
         cwd="/root",
     )
 
-    # wait 1 second for the context to be created
-    time.sleep(1)
-
     contexts = sandbox.list_code_contexts()
     last_context = contexts[-1]
 
diff --git a/template/server/main.py b/template/server/main.py
index ea89a9d8..1f296926 100644
--- a/template/server/main.py
+++ b/template/server/main.py
@@ -2,7 +2,7 @@
 import sys
 import httpx
 
-from typing import Dict, Union, Literal, Set
+from typing import Dict, Union, Literal, List
 
 from contextlib import asynccontextmanager
 from fastapi import FastAPI, Request
@@ -133,19 +133,18 @@ async def post_contexts(request: CreateContext) -> Context:
 
 
 @app.get("/contexts")
-async def get_contexts() -> Set[Context]:
+async def get_contexts() -> List[Context]:
     logger.info("Listing contexts")
 
-    context_ids = websockets.keys()
-
-    return set(
+    return [
         Context(
-            id=websockets[context_id].context_id,
-            language=websockets[context_id].language,
-            cwd=websockets[context_id].cwd,
+            id=ws.context_id,
+            language=ws.language,
+            cwd=ws.cwd,
         )
-        for context_id in context_ids
-    )
+        for key, ws in websockets.items()
+        if key != "default"
+    ]
 
 
 @app.post("/contexts/{context_id}/restart")