Limit PaperBench computer retries to startup

project/paperbench/paperbench/computer_utils.py

@@ -78,15 +78,19 @@ def before_sleep(state: RetryCallState) -> None:
             f"retrying due to '{exception}'"
         )
 
-    async for attempt in AsyncRetrying(
-        stop=stop_after_attempt(max_attempts),
-        retry=retry_if_exception_type(exception_types if exception_types else ()),
-        before_sleep=before_sleep,
-        reraise=True,
-    ):
-        with attempt:
-            async with computer_runtime.run(computer_config) as computer:
-                yield computer
+    async with AsyncExitStack() as stack:
+        computer: ComputerInterface | None = None
+        async for attempt in AsyncRetrying(
+            stop=stop_after_attempt(max_attempts),
+            retry=retry_if_exception_type(exception_types if exception_types else ()),
+            before_sleep=before_sleep,
+            reraise=True,
+        ):
+            with attempt:
+                computer = await stack.enter_async_context(computer_runtime.run(computer_config))
+
+        assert computer is not None
+        yield computer
 
 
 class ReleasableComputer(ComputerInterface):

project/paperbench/tests/unit/test_computer_retry.py

@@ -0,0 +1,56 @@
+from contextlib import asynccontextmanager
+
+import pytest
+
+from paperbench.computer_utils import start_computer_with_retry
+
+
+class FakeRuntime:
+    def __init__(self, startup_failures: int = 0) -> None:
+        self.startup_failures = startup_failures
+        self.enter_count = 0
+        self.exit_count = 0
+        self.computer = object()
+
+    @asynccontextmanager
+    async def run(self, _config):
+        self.enter_count += 1
+        if self.enter_count <= self.startup_failures:
+            raise ValueError("startup failed")
+        try:
+            yield self.computer
+        finally:
+            self.exit_count += 1
+
+
+@pytest.mark.asyncio
+async def test_start_computer_retries_only_startup_failures() -> None:
+    runtime = FakeRuntime(startup_failures=2)
+
+    async with start_computer_with_retry(
+        runtime,
+        object(),
+        exception_types=ValueError,
+        max_attempts=3,
+    ) as computer:
+        assert computer is runtime.computer
+
+    assert runtime.enter_count == 3
+    assert runtime.exit_count == 1
+
+
+@pytest.mark.asyncio
+async def test_start_computer_does_not_retry_caller_body() -> None:
+    runtime = FakeRuntime()
+
+    with pytest.raises(ValueError, match="caller failed"):
+        async with start_computer_with_retry(
+            runtime,
+            object(),
+            exception_types=ValueError,
+            max_attempts=3,
+        ):
+            raise ValueError("caller failed")
+
+    assert runtime.enter_count == 1
+    assert runtime.exit_count == 1