From 0ec2de9c59b28cef57b317f973b0c649eaa62ea3 Mon Sep 17 00:00:00 2001
From: Hew Li Yang <hewliyang@u.nus.edu>
Date: Tue, 25 Nov 2025 14:44:24 +0800
Subject: [PATCH] wip: add traffic_access_token if needed

---
 js/src/sandbox.ts                             | 19 +++++++++++++++----
 .../code_interpreter_async.py                 |  8 ++++++--
 .../code_interpreter_sync.py                  |  8 ++++++--
 3 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/js/src/sandbox.ts b/js/src/sandbox.ts
index 4272d846..0877e163 100644
--- a/js/src/sandbox.ts
+++ b/js/src/sandbox.ts
@@ -226,6 +226,9 @@ export class Sandbox extends BaseSandbox {
     if (this.envdAccessToken) {
       headers['X-Access-Token'] = this.envdAccessToken
     }
+    if (this.trafficAccessToken) {
+      headers['e2b-traffic-access-token'] = this.trafficAccessToken
+    }
 
     try {
       const res = await fetch(`${this.jupyterUrl}/execute`, {
@@ -295,13 +298,21 @@ export class Sandbox extends BaseSandbox {
    * @returns context object.
    */
   async createCodeContext(opts?: CreateCodeContextOpts): Promise<Context> {
+    const headers: Record<string, string> = {
+      'Content-Type': 'application/json',
+      ...this.connectionConfig.headers,
+    }
+    if (this.envdAccessToken) {
+      headers['X-Access-Token'] = this.envdAccessToken
+    }
+    if (this.trafficAccessToken) {
+      headers['e2b-traffic-access-token'] = this.trafficAccessToken
+    }
+
     try {
       const res = await fetch(`${this.jupyterUrl}/contexts`, {
         method: 'POST',
-        headers: {
-          'Content-Type': 'application/json',
-          ...this.connectionConfig.headers,
-        },
+        headers,
         body: JSON.stringify({
           language: opts?.language,
           cwd: opts?.cwd,
diff --git a/python/e2b_code_interpreter/code_interpreter_async.py b/python/e2b_code_interpreter/code_interpreter_async.py
index 789d6a70..7f2b5595 100644
--- a/python/e2b_code_interpreter/code_interpreter_async.py
+++ b/python/e2b_code_interpreter/code_interpreter_async.py
@@ -193,7 +193,9 @@ async def run_code(
 
         headers: Dict[str, str] = {}
         if self._envd_access_token:
-            headers = {"X-Access-Token": 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
 
         try:
             async with self._client.stream(
@@ -255,7 +257,9 @@ async def create_code_context(
 
         headers: Dict[str, str] = {}
         if self._envd_access_token:
-            headers = {"X-Access-Token": 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
 
         try:
             response = await self._client.post(
diff --git a/python/e2b_code_interpreter/code_interpreter_sync.py b/python/e2b_code_interpreter/code_interpreter_sync.py
index 6cf56c11..a44ad16c 100644
--- a/python/e2b_code_interpreter/code_interpreter_sync.py
+++ b/python/e2b_code_interpreter/code_interpreter_sync.py
@@ -190,7 +190,9 @@ def run_code(
 
         headers: Dict[str, str] = {}
         if self._envd_access_token:
-            headers = {"X-Access-Token": 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
 
         try:
             with self._client.stream(
@@ -252,7 +254,9 @@ def create_code_context(
 
         headers: Dict[str, str] = {}
         if self._envd_access_token:
-            headers = {"X-Access-Token": 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
 
         try:
             response = self._client.post(