feat: Enhance error handling by adding trace_id to exceptions and formatting messages

Author: mfzzf

.gitignore

@@ -266,3 +266,4 @@ cython_debug/
 
 # Test results
 results.csv
+.ace-tool/

e2b_connect/client.py

@@ -61,9 +61,19 @@ def make_error_from_http_code(http_code: int):
 
 
 class ConnectException(Exception):
-    def __init__(self, status: Code, message: str):
+    def __init__(self, status: Code, message: str, trace_id: str = None):
         self.status = status
         self.message = message
+        self.trace_id = trace_id
+        super().__init__(self._format_message())
+
+    def _format_message(self) -> str:
+        if self.trace_id:
+            return f"{self.message} [X-Trace-ID: {self.trace_id}]"
+        return self.message
+
+    def __str__(self) -> str:
+        return self._format_message()
 
 
 envelope_header_length = 5
@@ -84,15 +94,28 @@ def decode_envelope_header(header):
 
 
 def error_for_response(http_resp: Response):
+    trace_id = None
+    # Try to get X-Trace-ID from headers
+    if hasattr(http_resp, 'headers'):
+        headers = http_resp.headers
+        if isinstance(headers, (list, tuple)):
+            # httpcore returns headers as list of tuples
+            for name, value in headers:
+                if name.lower() == b'x-trace-id':
+                    trace_id = value.decode('utf-8') if isinstance(value, bytes) else value
+                    break
+        elif hasattr(headers, 'get'):
+            trace_id = headers.get("X-Trace-ID") or headers.get("x-trace-id")
+
     try:
         error = json.loads(http_resp.content)
-        return make_error(error)
+        return make_error(error, trace_id)
     except (json.decoder.JSONDecodeError, KeyError):
         error = {"code": http_resp.status, "message": http_resp.content.decode("utf-8")}
-        return make_error(error)
+        return make_error(error, trace_id)
 
 
-def make_error(error):
+def make_error(error, trace_id: str = None):
     status = None
     try:
         code_value = error.get("code")
@@ -104,7 +127,7 @@ def make_error(error):
     except (KeyError, ValueError):
         status = Code.unknown
 
-    return ConnectException(status, error.get("message", ""))
+    return ConnectException(status, error.get("message", ""), trace_id)
 
 
 def _sync_retry(func, exc, retries):

ucloud_sandbox/api/__init__.py

@@ -45,23 +45,26 @@ def handle_api_exception(
     except json.JSONDecodeError:
         body = {}
 
+    # Extract X-Trace-ID from response headers
+    trace_id = e.headers.get("X-Trace-ID") or e.headers.get("x-trace-id")
+
     if e.status_code == 401:
         message = f"{e.status_code}: Unauthorized, please check your credentials."
         if body.get("message"):
             message += f" - {body['message']}"
-        return AuthenticationException(message)
+        return AuthenticationException(message, trace_id=trace_id)
 
     if e.status_code == 429:
         message = f"{e.status_code}: Rate limit exceeded, please try again later."
         if body.get("message"):
             message += f" - {body['message']}"
-        return RateLimitException(message)
+        return RateLimitException(message, trace_id=trace_id)
 
     if "message" in body:
         return default_exception_class(
-            f"{e.status_code}: {body['message']}"
+            f"{e.status_code}: {body['message']}", trace_id=trace_id
         ).with_traceback(stack_trace)
-    return default_exception_class(f"{e.status_code}: {e.content}").with_traceback(
+    return default_exception_class(f"{e.status_code}: {e.content}", trace_id=trace_id).with_traceback(
         stack_trace
     )
 

ucloud_sandbox/envd/api.py

@@ -29,31 +29,33 @@ def handle_envd_api_exception(res: httpx.Response):
         return
 
     res.read()
+    trace_id = res.headers.get("X-Trace-ID") or res.headers.get("x-trace-id")
 
-    return format_envd_api_exception(res.status_code, get_message(res))
+    return format_envd_api_exception(res.status_code, get_message(res), trace_id)
 
 
 async def ahandle_envd_api_exception(res: httpx.Response):
     if res.is_success:
         return
 
     await res.aread()
+    trace_id = res.headers.get("X-Trace-ID") or res.headers.get("x-trace-id")
 
-    return format_envd_api_exception(res.status_code, get_message(res))
+    return format_envd_api_exception(res.status_code, get_message(res), trace_id)
 
 
-def format_envd_api_exception(status_code: int, message: str):
+def format_envd_api_exception(status_code: int, message: str, trace_id: str = None):
     if status_code == 400:
-        return InvalidArgumentException(message)
+        return InvalidArgumentException(message, trace_id=trace_id)
     elif status_code == 401:
-        return AuthenticationException(message)
+        return AuthenticationException(message, trace_id=trace_id)
     elif status_code == 404:
-        return NotFoundException(message)
+        return NotFoundException(message, trace_id=trace_id)
     elif status_code == 429:
-        return SandboxException(f"{message}: The requests are being rate limited.")
+        return SandboxException(f"{message}: The requests are being rate limited.", trace_id=trace_id)
     elif status_code == 502:
-        return format_sandbox_timeout_exception(message)
+        return format_sandbox_timeout_exception(message, trace_id)
     elif status_code == 507:
-        return NotEnoughSpaceException(message)
+        return NotEnoughSpaceException(message, trace_id=trace_id)
     else:
-        return SandboxException(f"{status_code}: {message}")
+        return SandboxException(f"{status_code}: {message}", trace_id=trace_id)

ucloud_sandbox/envd/rpc.py

@@ -19,28 +19,32 @@
 
 def handle_rpc_exception(e: Exception):
     if isinstance(e, ConnectException):
+        trace_id = getattr(e, "trace_id", None)
         if e.status == Code.invalid_argument:
-            return InvalidArgumentException(e.message)
+            return InvalidArgumentException(e.message, trace_id=trace_id)
         elif e.status == Code.unauthenticated:
-            return AuthenticationException(e.message)
+            return AuthenticationException(e.message, trace_id=trace_id)
         elif e.status == Code.not_found:
-            return NotFoundException(e.message)
+            return NotFoundException(e.message, trace_id=trace_id)
         elif e.status == Code.unavailable:
-            return format_sandbox_timeout_exception(e.message)
+            return format_sandbox_timeout_exception(e.message, trace_id=trace_id)
         elif e.status == Code.resource_exhausted:
             return RateLimitException(
-                f"{e.message}: Rate limit exceeded, please try again later."
+                f"{e.message}: Rate limit exceeded, please try again later.",
+                trace_id=trace_id,
             )
         elif e.status == Code.canceled:
             return TimeoutException(
-                f"{e.message}: This error is likely due to exceeding 'request_timeout'. You can pass the request timeout value as an option when making the request."
+                f"{e.message}: This error is likely due to exceeding 'request_timeout'. You can pass the request timeout value as an option when making the request.",
+                trace_id=trace_id,
             )
         elif e.status == Code.deadline_exceeded:
             return TimeoutException(
-                f"{e.message}: This error is likely due to exceeding 'timeout' — the total time a long running request (like process or directory watch) can be active. It can be modified by passing 'timeout' when making the request. Use '0' to disable the timeout."
+                f"{e.message}: This error is likely due to exceeding 'timeout' — the total time a long running request (like process or directory watch) can be active. It can be modified by passing 'timeout' when making the request. Use '0' to disable the timeout.",
+                trace_id=trace_id,
             )
         else:
-            return SandboxException(f"{e.status}: {e.message}")
+            return SandboxException(f"{e.status}: {e.message}", trace_id=trace_id)
     else:
         return e
 

ucloud_sandbox/exceptions.py

@@ -1,6 +1,7 @@
-def format_sandbox_timeout_exception(message: str):
+def format_sandbox_timeout_exception(message: str, trace_id: str = None):
     return TimeoutException(
-        f"{message}: This error is likely due to sandbox timeout. You can modify the sandbox timeout by passing 'timeout' when starting the sandbox or calling '.set_timeout' on the sandbox with the desired timeout."
+        f"{message}: This error is likely due to sandbox timeout. You can modify the sandbox timeout by passing 'timeout' when starting the sandbox or calling '.set_timeout' on the sandbox with the desired timeout.",
+        trace_id=trace_id,
     )
 
 
@@ -23,7 +24,18 @@ class SandboxException(Exception):
     Raised when a general sandbox exception occurs.
     """
 
-    pass
+    def __init__(self, message: str = "", trace_id: str = None):
+        self.message = message
+        self.trace_id = trace_id
+        super().__init__(self._format_message())
+
+    def _format_message(self) -> str:
+        if self.trace_id:
+            return f"{self.message} [X-Trace-ID: {self.trace_id}]"
+        return self.message
+
+    def __str__(self) -> str:
+        return self._format_message()
 
 
 class TimeoutException(SandboxException):
@@ -68,7 +80,18 @@ class AuthenticationException(Exception):
     Raised when authentication fails.
     """
 
-    pass
+    def __init__(self, message: str = "", trace_id: str = None):
+        self.message = message
+        self.trace_id = trace_id
+        super().__init__(self._format_message())
+
+    def _format_message(self) -> str:
+        if self.trace_id:
+            return f"{self.message} [X-Trace-ID: {self.trace_id}]"
+        return self.message
+
+    def __str__(self) -> str:
+        return self._format_message()
 
 
 class TemplateException(SandboxException):