Fix GPU sandbox hardware OAuth failure (#246)

* Fix GPU sandbox hardware OAuth failure

Co-authored-by: OpenAI Codex <codex@openai.com>

* Address sandbox sleep-time review

Co-authored-by: OpenAI Codex <codex@openai.com>

---------

Co-authored-by: OpenAI Codex <codex@openai.com>

agent/tools/sandbox_client.py

@@ -65,7 +65,6 @@ MAX_TIMEOUT = 1200
 WAIT_TIMEOUT = 600
 WAIT_INTERVAL = 5
 API_WAIT_TIMEOUT = 180
-HARDWARE_REQUEST_TIMEOUT = 60
 CPU_BASIC_HARDWARE = "cpu-basic"
 
 
@@ -78,58 +77,6 @@ def _is_transient_space_visibility_error(error: Exception) -> bool:
     return "Repository Not Found" in message or "404 Client Error" in message
 
 
-def _is_transient_space_management_error(error: Exception) -> bool:
-    """Return True when a just-created private Space is not manageable yet."""
-    response = getattr(error, "response", None)
-    if getattr(response, "status_code", None) in {401, 404}:
-        return True
-    message = str(error)
-    return (
-        "Repository Not Found" in message
-        or "401 Client Error" in message
-        or "404 Client Error" in message
-    )
-
-
-def _request_space_hardware_with_retry(
-    api: HfApi,
-    space_id: str,
-    *,
-    hardware: str,
-    sleep_time: int | None,
-    log: Callable[[str], object],
-    check_cancel: Callable[[], object],
-) -> None:
-    """Request hardware, retrying while Hub permissions propagate for a new Space."""
-    deadline = time.time() + HARDWARE_REQUEST_TIMEOUT
-    attempt = 0
-    while True:
-        check_cancel()
-        try:
-            api.request_space_hardware(
-                space_id,
-                hardware=hardware,
-                sleep_time=sleep_time,
-            )
-            return
-        except Exception as e:
-            if not _is_transient_space_management_error(e):
-                raise
-
-            remaining = deadline - time.time()
-            if remaining <= 0:
-                raise
-
-            attempt += 1
-            status_code = getattr(getattr(e, "response", None), "status_code", None)
-            status = f"HTTP {status_code}" if status_code else type(e).__name__
-            log(
-                f"  Hardware request not accepted yet ({status}); "
-                f"retrying ({attempt})..."
-            )
-            time.sleep(min(WAIT_INTERVAL, remaining))
-
-
 _DOCKERFILE = """\
 FROM ghcr.io/astral-sh/uv:python3.12-bookworm-slim
 
@@ -679,24 +626,21 @@ class Sandbox:
 
         _check_cancel()
 
-        # ``duplicate_space`` already receives the target hardware. The extra
-        # /hardware call is useful for paid tiers, but hosted OAuth tokens can
-        # 401 on that endpoint for a fresh private Space even after duplication
-        # succeeds. Avoid the redundant call for default CPU sandboxes when no
-        # auto-sleep timer is requested; with sleep_time set, the hardware
-        # endpoint is still needed to configure auto-sleep.
-        if hardware == CPU_BASIC_HARDWARE and sleep_time is None:
-            _log(f"Using duplicated Space hardware: {hardware}")
-        else:
-            _request_space_hardware_with_retry(
-                api,
-                space_id,
-                hardware=hardware,
-                sleep_time=sleep_time,
-                log=_log,
-                check_cancel=_check_cancel,
-            )
-            _log(f"Requested hardware: {hardware}")
+        # ``duplicate_space`` sends hardware and sleepTimeSeconds in the
+        # initial create request. Avoid a second /hardware call: deployed HF
+        # OAuth tokens can 401 on that endpoint for a just-created private
+        # Space even though duplication itself succeeded. We rely on the
+        # duplicate endpoint to honor sleepTimeSeconds for upgraded hardware;
+        # cpu-basic auto-sleep is fixed by the Hub.
+        _log(f"Using duplicated Space hardware: {hardware}")
+        if sleep_time is not None:
+            if hardware == CPU_BASIC_HARDWARE:
+                _log(
+                    f"Requested duplicated Space sleep time: {sleep_time}s "
+                    "(cpu-basic auto-sleep is fixed by the Hub)"
+                )
+            else:
+                _log(f"Using duplicated Space sleep time: {sleep_time}s")
 
         # Inject secrets BEFORE uploading server files (which triggers rebuild).
         # Secrets added after a Space is running aren't available until restart,

tests/unit/test_sandbox_private_spaces.py

@@ -3,8 +3,6 @@ import threading
 import time
 from types import SimpleNamespace
 
-import pytest
-
 from agent.core import telemetry
 from agent.tools import sandbox_client, sandbox_tool
 from agent.tools.sandbox_client import Sandbox
@@ -17,6 +15,7 @@ def _fail_metadata_update(*args, **kwargs):
 
 def test_sandbox_client_defaults_to_private_spaces(monkeypatch):
     duplicate_kwargs = {}
+    logs: list[str] = []
     requested_hardware = []
 
     class FakeApi:
@@ -44,11 +43,12 @@ def test_sandbox_client_defaults_to_private_spaces(monkeypatch):
     )
     monkeypatch.setattr(Sandbox, "_wait_for_api", lambda self, *args, **kwargs: None)
 
-    Sandbox.create(owner="alice", token="hf-token", log=lambda msg: None)
+    Sandbox.create(owner="alice", token="hf-token", log=logs.append)
 
     assert duplicate_kwargs["private"] is True
     assert duplicate_kwargs["hardware"] == "cpu-basic"
     assert requested_hardware == []
+    assert not any("sleep time" in log for log in logs)
 
 
 def test_sandbox_client_retries_transient_runtime_404(monkeypatch):
@@ -98,32 +98,20 @@ def test_sandbox_client_retries_transient_runtime_404(monkeypatch):
     assert runtime_calls == 2
 
 
-def test_sandbox_client_retries_transient_hardware_401(monkeypatch):
-    hardware_calls = 0
+def test_sandbox_client_configures_gpu_at_duplication(monkeypatch):
+    duplicate_kwargs = {}
     logs: list[str] = []
-
-    class FakeResponse:
-        status_code = 401
-
-    class FakeHardware401(Exception):
-        response = FakeResponse()
-
-        def __str__(self):
-            return "401 Client Error: Repository Not Found"
+    requested_hardware = []
 
     class FakeApi:
         def __init__(self, token=None):
             self.token = token
 
         def duplicate_space(self, **kwargs):
-            pass
+            duplicate_kwargs.update(kwargs)
 
         def request_space_hardware(self, space_id, hardware, sleep_time=None):
-            nonlocal hardware_calls
-            hardware_calls += 1
-            if hardware_calls == 1:
-                raise FakeHardware401()
-            return SimpleNamespace(stage="BUILDING", hardware=None)
+            requested_hardware.append((space_id, hardware, sleep_time))
 
         def add_space_secret(self, *args, **kwargs):
             pass
@@ -144,58 +132,62 @@ def test_sandbox_client_retries_transient_hardware_401(monkeypatch):
         owner="alice",
         token="hf-token",
         hardware="t4-small",
+        sleep_time=2700,
         log=logs.append,
     )
 
     assert sandbox.space_id.startswith("alice/sandbox-")
-    assert hardware_calls == 2
-    assert any("Hardware request not accepted yet (HTTP 401)" in log for log in logs)
+    assert duplicate_kwargs["hardware"] == "t4-small"
+    assert duplicate_kwargs["sleep_time"] == 2700
+    assert requested_hardware == []
+    assert "Using duplicated Space hardware: t4-small" in logs
+    assert "Using duplicated Space sleep time: 2700s" in logs
 
 
-def test_sandbox_hardware_retry_reraises_after_timeout(monkeypatch):
-    calls = 0
+def test_sandbox_client_logs_cpu_sleep_time_as_hub_fixed(monkeypatch):
+    duplicate_kwargs = {}
     logs: list[str] = []
-    sleeps: list[float] = []
+    requested_hardware = []
 
-    class FakeResponse:
-        status_code = 401
+    class FakeApi:
+        def __init__(self, token=None):
+            self.token = token
 
-    class FakeHardware401(Exception):
-        response = FakeResponse()
+        def duplicate_space(self, **kwargs):
+            duplicate_kwargs.update(kwargs)
 
-        def __str__(self):
-            return "401 Client Error: Repository Not Found"
+        def request_space_hardware(self, space_id, hardware, sleep_time=None):
+            requested_hardware.append((space_id, hardware, sleep_time))
 
-    first_error = FakeHardware401("first")
-    timeout_error = FakeHardware401("timeout")
+        def add_space_secret(self, *args, **kwargs):
+            pass
 
-    class FakeApi:
-        def request_space_hardware(self, space_id, hardware, sleep_time=None):
-            nonlocal calls
-            calls += 1
-            if calls == 1:
-                raise first_error
-            raise timeout_error
-
-    timestamps = iter([100.0, 100.0, 161.0])
-
-    monkeypatch.setattr(sandbox_client.time, "time", lambda: next(timestamps))
-    monkeypatch.setattr(sandbox_client.time, "sleep", sleeps.append)
-
-    with pytest.raises(FakeHardware401) as excinfo:
-        sandbox_client._request_space_hardware_with_retry(
-            FakeApi(),
-            "alice/sandbox-12345678",
-            hardware="cpu-basic",
-            sleep_time=None,
-            log=logs.append,
-            check_cancel=lambda: None,
-        )
+        def get_space_runtime(self, space_id):
+            return SimpleNamespace(stage="RUNNING", hardware="cpu-basic")
+
+    monkeypatch.setattr(sandbox_client, "HfApi", FakeApi)
+    monkeypatch.setattr(
+        Sandbox,
+        "_setup_server",
+        staticmethod(lambda *args, **kwargs: None),
+    )
+    monkeypatch.setattr(Sandbox, "_wait_for_api", lambda self, *args, **kwargs: None)
+
+    Sandbox.create(
+        owner="alice",
+        token="hf-token",
+        sleep_time=2700,
+        log=logs.append,
+    )
 
-    assert excinfo.value is timeout_error
-    assert calls == 2
-    assert sleeps == [sandbox_client.WAIT_INTERVAL]
-    assert len(logs) == 1
+    assert duplicate_kwargs["hardware"] == "cpu-basic"
+    assert duplicate_kwargs["sleep_time"] == 2700
+    assert requested_hardware == []
+    assert "Using duplicated Space hardware: cpu-basic" in logs
+    assert (
+        "Requested duplicated Space sleep time: 2700s "
+        "(cpu-basic auto-sleep is fixed by the Hub)"
+    ) in logs
 
 
 def test_sandbox_tool_forces_private_spaces(monkeypatch):