address error types

Author: Dingway98

hyperbrowser/client/managers/async_manager/sandboxes/sandbox_transport.py

@@ -11,7 +11,7 @@
     parse_json_response,
     resolve_runtime_transport_target,
 )
-from ...sandboxes.shared import _build_query_path
+from ...sandboxes.shared import _build_query_path, _is_replayable_http_content
 
 
 class RuntimeTransport:
@@ -167,6 +167,12 @@ async def _request(
 
         if response.status_code == 401 and allow_refresh:
             await response.aclose()
+            if not _is_replayable_http_content(content):
+                return ensure_response_ok(
+                    response,
+                    "runtime",
+                    "Runtime token expired during a streaming request; retry with a fresh stream.",
+                )
             refreshed = await self._resolve_connection(True)
             retry = await self._send(
                 refreshed,

hyperbrowser/client/managers/sandboxes/image_build.py

@@ -8,14 +8,22 @@
 import uuid
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, List, Optional, Sequence
+from typing import Dict, FrozenSet, List, Optional, Sequence
 
 import httpx
 
-from ....models.sandbox import SandboxImageBuildUpload, SandboxImageInit
+from ....models.sandbox import (
+    SandboxImageBuildInputFormat,
+    SandboxImageBuildStatus,
+    SandboxImageBuildUpload,
+    SandboxImageInit,
+)
 
-IMAGE_BUILD_INPUT_FORMAT = "rootfs_export_tar_gz"
+IMAGE_BUILD_INPUT_FORMAT: SandboxImageBuildInputFormat = "rootfs_export_tar_gz"
 IMAGE_BUILD_SOURCE_PLATFORM = "linux/amd64"
+TERMINAL_IMAGE_BUILD_STATUSES: FrozenSet[SandboxImageBuildStatus] = frozenset(
+    {"completed", "failed", "canceled", "cancelled"}
+)
 _IMAGE_INIT_ENV_KEY_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
 _RESERVED_IMAGE_INIT_ENV_KEYS = {
     "SANDBOX_ENABLED",
@@ -35,7 +43,7 @@ class DockerImageBuildArtifact:
     path: str
     sha256_hex: str
     size_bytes: int
-    input_format: str
+    input_format: SandboxImageBuildInputFormat
     source_platform: str
     image_config_user: str
     image_init: Optional[SandboxImageInit]
@@ -161,8 +169,8 @@ def make_temp_docker_tag(prefix: str = "hyperbrowser-sdk-build") -> str:
     return f"{prefix}:{uuid.uuid4().hex}"
 
 
-def is_terminal_image_build_status(status: str) -> bool:
-    return status.strip().lower() in {"completed", "failed", "canceled", "cancelled"}
+def is_terminal_image_build_status(status: SandboxImageBuildStatus) -> bool:
+    return status in TERMINAL_IMAGE_BUILD_STATUSES
 
 
 def _ensure_docker_image_source_platform(docker_image: str, platform: str) -> None:
@@ -232,6 +240,7 @@ def _package_docker_container(
     hasher = hashlib.sha256()
     writer = _HashingCountingWriter(tmp, hasher)
     stderr_file = tempfile.TemporaryFile()
+    process = None
     try:
         process = subprocess.Popen(
             ["docker", "export", container_id],
@@ -272,11 +281,27 @@ def _package_docker_container(
             pass
         raise
     finally:
+        if process is not None:
+            _cleanup_docker_export_process(process)
         tmp.close()
         stderr_file.close()
         _remove_docker_container(container_id)
 
 
+def _cleanup_docker_export_process(process) -> None:
+    if process.stdout is not None:
+        process.stdout.close()
+    if process.poll() is None:
+        process.terminate()
+        try:
+            process.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+            process.kill()
+            process.wait()
+        return
+    process.wait()
+
+
 def _derive_auto_image_init(config: Dict[str, object]) -> Optional[SandboxImageInit]:
     env = _derive_auto_image_env(_list_string_config(config.get("Env")))
     args = _derive_auto_startup_args(

hyperbrowser/client/managers/sandboxes/shared.py

@@ -1,4 +1,5 @@
 import base64
+import collections.abc
 import posixpath
 import re
 from datetime import datetime, timedelta, timezone
@@ -28,6 +29,24 @@ def _copy_model(model):
     return model.model_copy(deep=True)
 
 
+def _is_replayable_http_content(content) -> bool:
+    if content is None:
+        return True
+    if isinstance(content, (bytes, bytearray, memoryview, str)):
+        return True
+    if hasattr(content, "read"):
+        return False
+    if hasattr(content, "__aiter__"):
+        return False
+    if isinstance(content, collections.abc.Iterator):
+        return False
+    try:
+        iter(content)
+    except TypeError:
+        return True
+    return True
+
+
 def _quote_shell_token(token: str) -> str:
     if token == "":
         return "''"

hyperbrowser/client/managers/sync_manager/sandboxes/sandbox_transport.py

@@ -11,7 +11,7 @@
     parse_json_response,
     resolve_runtime_transport_target,
 )
-from ...sandboxes.shared import _build_query_path
+from ...sandboxes.shared import _build_query_path, _is_replayable_http_content
 
 
 class RuntimeTransport:
@@ -165,6 +165,12 @@ def _request(
 
         if response.status_code == 401 and allow_refresh:
             response.close()
+            if not _is_replayable_http_content(content):
+                return ensure_response_ok(
+                    response,
+                    "runtime",
+                    "Runtime token expired during a streaming request; retry with a fresh stream.",
+                )
             refreshed = self._resolve_connection(True)
             retry = self._send(
                 refreshed,

hyperbrowser/models/sandbox.py

@@ -397,7 +397,7 @@ class SandboxImageBuild(SandboxBaseModel):
     namespace: Optional[str] = None
     image_name: str = Field(alias="imageName")
     image_id: Optional[str] = Field(default=None, alias="imageId")
-    status: str
+    status: SandboxImageBuildStatus
     input_bucket: Optional[str] = Field(default=None, alias="inputBucket")
     input_key: Optional[str] = Field(default=None, alias="inputKey")
     input_sha256: Optional[str] = Field(default=None, alias="inputSha256")

tests/test_sandbox_image_build_helpers.py

@@ -3,6 +3,7 @@
 
 import httpx
 import pytest
+from pydantic import ValidationError
 
 import hyperbrowser.client.managers.async_manager.sandbox as async_sandbox_module
 import hyperbrowser.client.managers.sync_manager.sandbox as sync_sandbox_module
@@ -11,7 +12,27 @@
     SandboxManager as AsyncSandboxManager,
 )
 from hyperbrowser.client.managers.sync_manager.sandbox import SandboxManager
-from hyperbrowser.models import SandboxImageBuildUpload
+from hyperbrowser.models import SandboxImageBuild, SandboxImageBuildUpload
+
+
+def _image_build(
+    status: str,
+    *,
+    error_code: str = "",
+    error_message: str = "",
+) -> SandboxImageBuild:
+    return SandboxImageBuild(
+        id="build-123",
+        imageName="custom",
+        status=status,
+        errorCode=error_code,
+        errorMessage=error_message,
+    )
+
+
+def test_image_build_model_rejects_invalid_status_casing():
+    with pytest.raises(ValidationError):
+        _image_build(" Completed ")
 
 
 def test_build_docker_image_from_dockerfile_targets_linux_amd64(monkeypatch, tmp_path):
@@ -61,6 +82,71 @@ def fake_run(args, check, stdout, stderr, text):
         image_build.package_docker_image("local/app:latest")
 
 
+def test_package_docker_container_reaps_export_process_on_read_failure(
+    monkeypatch,
+    tmp_path,
+):
+    removed = []
+
+    class BrokenStdout:
+        def __init__(self):
+            self.closed = False
+
+        def read(self, size):
+            raise RuntimeError("read failed")
+
+        def close(self):
+            self.closed = True
+
+    class FakeProcess:
+        def __init__(self):
+            self.stdout = BrokenStdout()
+            self.terminated = False
+            self.killed = False
+            self.waits = 0
+            self.return_code = None
+
+        def poll(self):
+            return self.return_code
+
+        def terminate(self):
+            self.terminated = True
+            self.return_code = -15
+
+        def kill(self):
+            self.killed = True
+            self.return_code = -9
+
+        def wait(self, timeout=None):
+            self.waits += 1
+            if self.return_code is None:
+                self.return_code = 0
+            return self.return_code
+
+    fake_process = FakeProcess()
+
+    def fake_popen(args, stdout, stderr):
+        return fake_process
+
+    monkeypatch.setattr(image_build.subprocess, "Popen", fake_popen)
+    monkeypatch.setattr(image_build, "_remove_docker_container", removed.append)
+
+    with pytest.raises(RuntimeError, match="read failed"):
+        image_build._package_docker_container(
+            "local/app:latest",
+            "container-123",
+            {},
+            platform="linux/amd64",
+            temp_dir=str(tmp_path),
+        )
+
+    assert fake_process.stdout.closed is True
+    assert fake_process.terminated is True
+    assert fake_process.killed is False
+    assert fake_process.waits == 1
+    assert removed == ["container-123"]
+
+
 def test_upload_image_build_artifact_streams_file_with_content_length(
     monkeypatch,
     tmp_path,
@@ -140,6 +226,48 @@ def fake_build(**kwargs):
     assert removed == ["temp:tag"]
 
 
+def test_sync_wait_for_image_build_returns_completed_status(monkeypatch):
+    manager = SandboxManager(
+        SimpleNamespace(
+            timeout=30,
+            config=SimpleNamespace(runtime_proxy_override=None),
+        )
+    )
+    monkeypatch.setattr(
+        manager,
+        "get_image_build",
+        lambda build_id: _image_build("completed"),
+    )
+
+    build = manager.wait_for_image_build("build-123", poll_interval=0, timeout=1)
+
+    assert build.status == "completed"
+
+
+def test_sync_wait_for_image_build_raises_detailed_failed_status(monkeypatch):
+    manager = SandboxManager(
+        SimpleNamespace(
+            timeout=30,
+            config=SimpleNamespace(runtime_proxy_override=None),
+        )
+    )
+    monkeypatch.setattr(
+        manager,
+        "get_image_build",
+        lambda build_id: _image_build(
+            "failed",
+            error_code="E_BAD",
+            error_message="backend failure",
+        ),
+    )
+
+    with pytest.raises(
+        RuntimeError,
+        match=r"image build failed \[E_BAD\]: backend failure",
+    ):
+        manager.wait_for_image_build("build-123", poll_interval=0, timeout=1)
+
+
 @pytest.mark.anyio
 async def test_async_dockerfile_image_build_cleans_temp_tag_on_build_failure(
     monkeypatch,
@@ -175,3 +303,55 @@ def fake_build(**kwargs):
         )
 
     assert removed == ["temp:tag"]
+
+
+@pytest.mark.anyio
+async def test_async_wait_for_image_build_returns_completed_status(monkeypatch):
+    manager = AsyncSandboxManager(
+        SimpleNamespace(
+            timeout=30,
+            config=SimpleNamespace(runtime_proxy_override=None),
+        )
+    )
+
+    async def fake_get_image_build(build_id):
+        return _image_build("completed")
+
+    monkeypatch.setattr(manager, "get_image_build", fake_get_image_build)
+
+    build = await manager.wait_for_image_build(
+        "build-123",
+        poll_interval=0,
+        timeout=1,
+    )
+
+    assert build.status == "completed"
+
+
+@pytest.mark.anyio
+async def test_async_wait_for_image_build_raises_detailed_failed_status(monkeypatch):
+    manager = AsyncSandboxManager(
+        SimpleNamespace(
+            timeout=30,
+            config=SimpleNamespace(runtime_proxy_override=None),
+        )
+    )
+
+    async def fake_get_image_build(build_id):
+        return _image_build(
+            "failed",
+            error_code="E_BAD",
+            error_message="backend failure",
+        )
+
+    monkeypatch.setattr(manager, "get_image_build", fake_get_image_build)
+
+    with pytest.raises(
+        RuntimeError,
+        match=r"image build failed \[E_BAD\]: backend failure",
+    ):
+        await manager.wait_for_image_build(
+            "build-123",
+            poll_interval=0,
+            timeout=1,
+        )