fix(devboxes): clean up create timeout

src/runloop_api_client/resources/devboxes/devboxes.py

@@ -77,7 +77,7 @@
     AsyncDiskSnapshotsCursorIDPage,
 )
 from ..._exceptions import RunloopError, APIStatusError, APIConnectionError
-from ...lib.polling import PollingConfig, poll_until
+from ...lib.polling import PollingConfig, PollingTimeout, poll_until
 from ..._base_client import AsyncPaginator, make_request_options
 from .disk_snapshots import (
     DiskSnapshotsResource,
@@ -471,6 +471,7 @@ def create_and_await_running(
         mounts: Optional[Iterable[Mount]] | Omit = omit,
         name: Optional[str] | Omit = omit,
         polling_config: PollingConfig | None = None,
+        shutdown_on_timeout: bool = True,
         secrets: Optional[Dict[str, str]] | Omit = omit,
         snapshot_id: Optional[str] | Omit = omit,
         tunnel: Optional[devbox_create_params.Tunnel] | Omit = omit,
@@ -489,9 +490,11 @@ def create_and_await_running(
         Args:
             create_args: Arguments to pass to the `create` method. See the `create` method for detailed documentation.
             request_args: Optional request arguments including polling configuration and additional request options
+            shutdown_on_timeout: Shutdown the created devbox if waiting for running state times out.
 
         Returns:
-            The devbox in running state
+            The devbox in running state, or the created devbox if waiting times out and
+            shutdown_on_timeout is False.
 
         Raises:
             PollingTimeout: If polling times out before devbox is running
@@ -521,10 +524,16 @@ def create_and_await_running(
             idempotency_key=idempotency_key,
         )
 
-        return self.await_running(
-            devbox.id,
-            polling_config=polling_config,
-        )
+        try:
+            return self.await_running(
+                devbox.id,
+                polling_config=polling_config,
+            )
+        except PollingTimeout:
+            if not shutdown_on_timeout:
+                return devbox
+            self.shutdown(devbox.id)
+            raise
 
     def list(
         self,
@@ -2041,6 +2050,7 @@ async def create_and_await_running(
         mounts: Optional[Iterable[Mount]] | Omit = omit,
         name: Optional[str] | Omit = omit,
         polling_config: PollingConfig | None = None,
+        shutdown_on_timeout: bool = True,
         secrets: Optional[Dict[str, str]] | Omit = omit,
         snapshot_id: Optional[str] | Omit = omit,
         tunnel: Optional[devbox_create_params.Tunnel] | Omit = omit,
@@ -2059,9 +2069,11 @@ async def create_and_await_running(
         Args:
             See the `create` method for detailed documentation.
             polling_config: Optional polling configuration
+            shutdown_on_timeout: Shutdown the created devbox if waiting for running state times out.
 
         Returns:
-            The devbox in running state
+            The devbox in running state, or the created devbox if waiting times out and
+            shutdown_on_timeout is False.
 
         Raises:
             PollingTimeout: If polling times out before devbox is running
@@ -2092,10 +2104,16 @@ async def create_and_await_running(
             idempotency_key=idempotency_key,
         )
 
-        return await self.await_running(
-            devbox.id,
-            polling_config=polling_config,
-        )
+        try:
+            return await self.await_running(
+                devbox.id,
+                polling_config=polling_config,
+            )
+        except PollingTimeout:
+            if not shutdown_on_timeout:
+                return devbox
+            await self.shutdown(devbox.id)
+            raise
 
     async def await_running(
         self,

src/runloop_api_client/sdk/_types.py

@@ -92,11 +92,16 @@ class LongPollingRequestOptions(LongRequestOptions, PollingRequestOptions):  # t
     pass
 
 
-class SDKDevboxCreateParams(DevboxCreateParams, LongPollingRequestOptions):
+class DevboxTimeoutCleanupOptions(TypedDict, total=False):
+    shutdown_on_timeout: bool
+    """Shutdown the created devbox when waiting for it to run times out. Defaults to true."""
+
+
+class SDKDevboxCreateParams(DevboxCreateParams, LongPollingRequestOptions, DevboxTimeoutCleanupOptions):
     pass
 
 
-class SDKDevboxCreateFromImageParams(DevboxBaseCreateParams, LongPollingRequestOptions):
+class SDKDevboxCreateFromImageParams(DevboxBaseCreateParams, LongPollingRequestOptions, DevboxTimeoutCleanupOptions):
     pass
 
 

tests/api_resources/test_devboxes.py

@@ -4,7 +4,7 @@
 
 import os
 from typing import Any, cast
-from unittest.mock import Mock, patch
+from unittest.mock import Mock, AsyncMock, patch
 
 import httpx
 import pytest
@@ -1546,6 +1546,49 @@ def test_method_create_and_await_running_await_failure(self, client: Runloop) ->
                         name="test",
                     )
 
+    @parametrize
+    def test_method_create_and_await_running_timeout_shuts_down(self, client: Runloop) -> None:
+        devbox = DevboxView(
+            id="test_id",
+            status="provisioning",
+            capabilities=[],
+            create_time_ms=1234567890,
+            launch_parameters=LaunchParameters(resource_size_request="X_SMALL"),
+            metadata={},
+            state_transitions=[],
+        )
+        timeout = PollingTimeout("Timed out", devbox)
+
+        with patch.object(client.devboxes, "create", return_value=devbox):
+            with patch.object(client.devboxes, "await_running", side_effect=timeout):
+                with patch.object(client.devboxes, "shutdown") as mock_shutdown:
+                    with pytest.raises(PollingTimeout) as exc_info:
+                        client.devboxes.create_and_await_running()
+
+        assert exc_info.value is timeout
+        mock_shutdown.assert_called_once_with("test_id")
+
+    @parametrize
+    def test_method_create_and_await_running_timeout_returns_devbox_when_configured(self, client: Runloop) -> None:
+        devbox = DevboxView(
+            id="test_id",
+            status="provisioning",
+            capabilities=[],
+            create_time_ms=1234567890,
+            launch_parameters=LaunchParameters(resource_size_request="X_SMALL"),
+            metadata={},
+            state_transitions=[],
+        )
+        timeout = PollingTimeout("Timed out", devbox)
+
+        with patch.object(client.devboxes, "create", return_value=devbox):
+            with patch.object(client.devboxes, "await_running", side_effect=timeout):
+                with patch.object(client.devboxes, "shutdown") as mock_shutdown:
+                    result = client.devboxes.create_and_await_running(shutdown_on_timeout=False)
+
+        assert result is devbox
+        mock_shutdown.assert_not_called()
+
     @parametrize
     def test_method_await_suspended_success(self, client: Runloop) -> None:
         """Test await_suspended with successful polling to suspended state"""
@@ -1744,6 +1787,51 @@ async def test_method_create(self, async_client: AsyncRunloop) -> None:
         devbox = await async_client.devboxes.create()
         assert_matches_type(DevboxView, devbox, path=["response"])
 
+    @parametrize
+    async def test_method_create_and_await_running_timeout_shuts_down(self, async_client: AsyncRunloop) -> None:
+        devbox = DevboxView(
+            id="test_id",
+            status="provisioning",
+            capabilities=[],
+            create_time_ms=1234567890,
+            launch_parameters=LaunchParameters(resource_size_request="X_SMALL"),
+            metadata={},
+            state_transitions=[],
+        )
+        timeout = PollingTimeout("Timed out", devbox)
+
+        with patch.object(async_client.devboxes, "create", AsyncMock(return_value=devbox)):
+            with patch.object(async_client.devboxes, "await_running", AsyncMock(side_effect=timeout)):
+                with patch.object(async_client.devboxes, "shutdown", AsyncMock()) as mock_shutdown:
+                    with pytest.raises(PollingTimeout) as exc_info:
+                        await async_client.devboxes.create_and_await_running()
+
+        assert exc_info.value is timeout
+        mock_shutdown.assert_awaited_once_with("test_id")
+
+    @parametrize
+    async def test_method_create_and_await_running_timeout_returns_devbox_when_configured(
+        self, async_client: AsyncRunloop
+    ) -> None:
+        devbox = DevboxView(
+            id="test_id",
+            status="provisioning",
+            capabilities=[],
+            create_time_ms=1234567890,
+            launch_parameters=LaunchParameters(resource_size_request="X_SMALL"),
+            metadata={},
+            state_transitions=[],
+        )
+        timeout = PollingTimeout("Timed out", devbox)
+
+        with patch.object(async_client.devboxes, "create", AsyncMock(return_value=devbox)):
+            with patch.object(async_client.devboxes, "await_running", AsyncMock(side_effect=timeout)):
+                with patch.object(async_client.devboxes, "shutdown", AsyncMock()) as mock_shutdown:
+                    result = await async_client.devboxes.create_and_await_running(shutdown_on_timeout=False)
+
+        assert result is devbox
+        mock_shutdown.assert_not_awaited()
+
     @parametrize
     async def test_method_create_with_all_params(self, async_client: AsyncRunloop) -> None:
         devbox = await async_client.devboxes.create(

tests/sdk/test_async_ops.py

@@ -65,11 +65,13 @@ async def test_create(self, mock_async_client: AsyncMock, devbox_view: MockDevbo
             name="test-devbox",
             metadata={"key": "value"},
             polling_config=PollingConfig(timeout_seconds=60.0),
+            shutdown_on_timeout=False,
         )
 
         assert isinstance(devbox, AsyncDevbox)
         assert devbox.id == "dbx_123"
         mock_async_client.devboxes.create_and_await_running.assert_awaited_once()
+        assert mock_async_client.devboxes.create_and_await_running.call_args.kwargs["shutdown_on_timeout"] is False
 
     @pytest.mark.asyncio
     async def test_create_from_blueprint_id(self, mock_async_client: AsyncMock, devbox_view: MockDevboxView) -> None:

tests/sdk/test_ops.py

@@ -64,11 +64,13 @@ def test_create(self, mock_client: Mock, devbox_view: MockDevboxView) -> None:
             name="test-devbox",
             metadata={"key": "value"},
             polling_config=PollingConfig(timeout_seconds=60.0),
+            shutdown_on_timeout=False,
         )
 
         assert isinstance(devbox, Devbox)
         assert devbox.id == "dbx_123"
         mock_client.devboxes.create_and_await_running.assert_called_once()
+        assert mock_client.devboxes.create_and_await_running.call_args.kwargs["shutdown_on_timeout"] is False
 
     def test_create_from_blueprint_id(self, mock_client: Mock, devbox_view: MockDevboxView) -> None:
         """Test create_from_blueprint_id method."""