fix(code-interpreter): forward context authentication tokens

.changeset/secure-interpreter-contexts.md

@@ -0,0 +1,6 @@
+---
+"@e2b/code-interpreter": patch
+"@e2b/code-interpreter-python": patch
+---
+
+Forward sandbox authentication tokens for code context creation, listing, restart, and removal.

packages/code-interpreter-js/src/sandbox.ts

@@ -221,7 +221,7 @@
     }
 
     try {
       const res = await fetch(`${this.jupyterUrl}/execute`, {
         method: 'POST',
         headers,
         body: JSON.stringify({
@@ -298,6 +298,9 @@
       if (this.trafficAccessToken) {
         headers['E2B-Traffic-Access-Token'] = this.trafficAccessToken
       }
+      if (this.envdAccessToken) {
+        headers['X-Access-Token'] = this.envdAccessToken
+      }
 
       const res = await fetch(`${this.jupyterUrl}/contexts`, {
         method: 'POST',
@@ -340,6 +343,9 @@
       if (this.trafficAccessToken) {
         headers['E2B-Traffic-Access-Token'] = this.trafficAccessToken
       }
+      if (this.envdAccessToken) {
+        headers['X-Access-Token'] = this.envdAccessToken
+      }
 
       const res = await fetch(`${this.jupyterUrl}/contexts/${id}`, {
         method: 'DELETE',
@@ -375,6 +381,9 @@
       if (this.trafficAccessToken) {
         headers['E2B-Traffic-Access-Token'] = this.trafficAccessToken
       }
+      if (this.envdAccessToken) {
+        headers['X-Access-Token'] = this.envdAccessToken
+      }
 
       const res = await fetch(`${this.jupyterUrl}/contexts`, {
         method: 'GET',
@@ -415,6 +424,9 @@
       if (this.trafficAccessToken) {
         headers['E2B-Traffic-Access-Token'] = this.trafficAccessToken
       }
+      if (this.envdAccessToken) {
+        headers['X-Access-Token'] = this.envdAccessToken
+      }
 
       const res = await fetch(`${this.jupyterUrl}/contexts/${id}/restart`, {
         method: 'POST',

packages/code-interpreter-js/tests/contextAuth.test.ts

@@ -0,0 +1,46 @@
+import { afterEach, expect, test, vi } from 'vitest'
+
+import { Sandbox } from '../src'
+
+afterEach(() => vi.unstubAllGlobals())
+
+test.each(['sandbox-token', undefined])(
+  'context operations forward the sandbox token (%s)',
+  async (envdAccessToken) => {
+    const context = { id: 'context-id', language: 'python', cwd: '/home/user' }
+    const requests: Request[] = []
+    vi.stubGlobal('fetch', async (url: string, init: RequestInit) => {
+      const request = new Request(url, init)
+      requests.push(request)
+      if (request.headers.get('X-Access-Token') !== (envdAccessToken ?? null)) {
+        return new Response('Unauthorized', { status: 401 })
+      }
+      return Response.json(request.method === 'GET' ? [context] : context)
+    })
+    const SandboxClass = Sandbox as unknown as new (opts: object) => Sandbox
+    const sandbox = new SandboxClass({
+      sandboxId: 'sandbox-id',
+      envdVersion: '0.2.0',
+      envdAccessToken,
+      trafficAccessToken: 'traffic-token',
+      sandboxUrl: 'https://interpreter.example.test',
+    })
+
+    expect(await sandbox.createCodeContext()).toEqual(context)
+    expect(await sandbox.listCodeContexts()).toEqual([context])
+    await sandbox.restartCodeContext(context)
+    await sandbox.removeCodeContext(context)
+
+    expect(requests).toHaveLength(4)
+    for (const request of requests) {
+      expect(request.headers.get('X-Access-Token')).toBe(
+        envdAccessToken ?? null
+      )
+      expect(request.headers.get('E2B-Traffic-Access-Token')).toBe(
+        'traffic-token'
+      )
+      expect(request.headers.get('E2b-Sandbox-Id')).toBe('sandbox-id')
+      expect(request.headers.get('E2b-Sandbox-Port')).toBe('49999')
+    }
+  }
+)

packages/code-interpreter-python/e2b_code_interpreter/code_interpreter_async.py

@@ -290,6 +290,8 @@ async def create_code_context(
                 "E2b-Sandbox-Id": self.sandbox_id,
                 "E2b-Sandbox-Port": str(JUPYTER_PORT),
             }
+            if self._envd_access_token:
+                headers["X-Access-Token"] = self._envd_access_token
             if self.traffic_access_token:
                 headers["E2B-Traffic-Access-Token"] = self.traffic_access_token
 
@@ -331,6 +333,8 @@ async def remove_code_context(
                 "E2b-Sandbox-Id": self.sandbox_id,
                 "E2b-Sandbox-Port": str(JUPYTER_PORT),
             }
+            if self._envd_access_token:
+                headers["X-Access-Token"] = self._envd_access_token
             if self.traffic_access_token:
                 headers["E2B-Traffic-Access-Token"] = self.traffic_access_token
 
@@ -361,6 +365,8 @@ async def list_code_contexts(self) -> List[Context]:
                 "E2b-Sandbox-Id": self.sandbox_id,
                 "E2b-Sandbox-Port": str(JUPYTER_PORT),
             }
+            if self._envd_access_token:
+                headers["X-Access-Token"] = self._envd_access_token
             if self.traffic_access_token:
                 headers["E2B-Traffic-Access-Token"] = self.traffic_access_token
 
@@ -400,6 +406,8 @@ async def restart_code_context(
                 "E2b-Sandbox-Id": self.sandbox_id,
                 "E2b-Sandbox-Port": str(JUPYTER_PORT),
             }
+            if self._envd_access_token:
+                headers["X-Access-Token"] = self._envd_access_token
             if self.traffic_access_token:
                 headers["E2B-Traffic-Access-Token"] = self.traffic_access_token
 

packages/code-interpreter-python/e2b_code_interpreter/code_interpreter_sync.py

@@ -287,6 +287,8 @@ def create_code_context(
                 "E2b-Sandbox-Id": self.sandbox_id,
                 "E2b-Sandbox-Port": str(JUPYTER_PORT),
             }
+            if self._envd_access_token:
+                headers["X-Access-Token"] = self._envd_access_token
             if self.traffic_access_token:
                 headers["E2B-Traffic-Access-Token"] = self.traffic_access_token
 
@@ -328,6 +330,8 @@ def remove_code_context(
                 "E2b-Sandbox-Id": self.sandbox_id,
                 "E2b-Sandbox-Port": str(JUPYTER_PORT),
             }
+            if self._envd_access_token:
+                headers["X-Access-Token"] = self._envd_access_token
             if self.traffic_access_token:
                 headers["E2B-Traffic-Access-Token"] = self.traffic_access_token
 
@@ -358,6 +362,8 @@ def list_code_contexts(self) -> List[Context]:
                 "E2b-Sandbox-Id": self.sandbox_id,
                 "E2b-Sandbox-Port": str(JUPYTER_PORT),
             }
+            if self._envd_access_token:
+                headers["X-Access-Token"] = self._envd_access_token
             if self.traffic_access_token:
                 headers["E2B-Traffic-Access-Token"] = self.traffic_access_token
 
@@ -398,6 +404,8 @@ def restart_code_context(
                 "E2b-Sandbox-Id": self.sandbox_id,
                 "E2b-Sandbox-Port": str(JUPYTER_PORT),
             }
+            if self._envd_access_token:
+                headers["X-Access-Token"] = self._envd_access_token
             if self.traffic_access_token:
                 headers["E2B-Traffic-Access-Token"] = self.traffic_access_token
 

packages/code-interpreter-python/tests/test_context_auth.py

@@ -0,0 +1,62 @@
+import inspect
+
+import httpx
+import pytest
+
+from e2b.connection_config import ConnectionConfig
+from e2b_code_interpreter import AsyncSandbox, Sandbox
+from e2b_code_interpreter import code_interpreter_async, code_interpreter_sync
+
+
+@pytest.mark.parametrize(
+    "cls,module",
+    [(Sandbox, code_interpreter_sync), (AsyncSandbox, code_interpreter_async)],
+)
+@pytest.mark.parametrize("token", ["sandbox-token", None])
+async def test_context_operations_forward_auth(cls, module, token, monkeypatch):
+    context = {"id": "context-id", "language": "python", "cwd": "/home/user"}
+    requests = []
+
+    def handle(request):
+        requests.append(request)
+        if request.headers.get("X-Access-Token") != token:
+            return httpx.Response(401, text="Unauthorized")
+        return httpx.Response(
+            200, json=[context] if request.method == "GET" else context
+        )
+
+    monkeypatch.setattr(
+        module, "get_transport", lambda *args, **kwargs: httpx.MockTransport(handle)
+    )
+    sandbox = cls(
+        sandbox_id="sandbox-id",
+        sandbox_domain=None,
+        envd_version="0.2.0",
+        envd_access_token=token,
+        traffic_access_token="traffic-token",
+        connection_config=ConnectionConfig(
+            sandbox_url="https://interpreter.example.test"
+        ),
+    )
+
+    async def call(method, *args):
+        result = method(*args)
+        return await result if inspect.isawaitable(result) else result
+
+    try:
+        assert (await call(sandbox.create_code_context)).id == context["id"]
+        assert (await call(sandbox.list_code_contexts))[0].id == context["id"]
+        await call(sandbox.restart_code_context, context["id"])
+        await call(sandbox.remove_code_context, context["id"])
+    finally:
+        if cls is AsyncSandbox:
+            await sandbox._client.aclose()
+        else:
+            sandbox._client.close()
+
+    assert len(requests) == 4
+    for request in requests:
+        assert request.headers.get("X-Access-Token") == token
+        assert request.headers["E2B-Traffic-Access-Token"] == "traffic-token"
+        assert request.headers["E2b-Sandbox-Id"] == "sandbox-id"
+        assert request.headers["E2b-Sandbox-Port"] == "49999"