Add execute and await completion function (#643)

* Add execute and await completion

* Fix lint

Author: alb-rl

pyproject.toml

@@ -9,18 +9,18 @@ authors = [
 ]
 dependencies = [
     "httpx>=0.23.0, <1",
-    "pydantic>=1.9.0, <3",
+    "pydantic>=1.9.0, <2.11",
     "typing-extensions>=4.10, <5",
     "anyio>=3.5.0, <5",
     "distro>=1.7.0, <2",
     "sniffio",
+    "uuid-utils>=0.11.0",
 ]
 
-requires-python = ">= 3.8"
+requires-python = ">= 3.9"
 classifiers = [
   "Typing :: Typed",
   "Intended Audience :: Developers",
-  "Programming Language :: Python :: 3.8",
   "Programming Language :: Python :: 3.9",
   "Programming Language :: Python :: 3.10",
   "Programming Language :: Python :: 3.11",
@@ -59,6 +59,7 @@ dev-dependencies = [
     "importlib-metadata>=6.7.0",
     "rich>=13.7.1",
     "pytest-xdist>=3.6.1",
+    "uuid-utils>=0.11.0"
 ]
 
 [tool.rye.scripts]
@@ -86,6 +87,7 @@ typecheck = { chain = [
   "typecheck:pyright",
   "typecheck:mypy"
 ]}
+
 "typecheck:pyright" = "pyright"
 "typecheck:verify-types" = "pyright --verifytypes runloop_api_client --ignoreexternal"
 "typecheck:mypy" = "mypy ."
@@ -130,8 +132,8 @@ path = "README.md"
 
 [[tool.hatch.metadata.hooks.fancy-pypi-readme.substitutions]]
 # replace relative links with absolute links
-pattern = '\[(.+?)\]\(((?!https?://)\S+?)\)'
-replacement = '[\1](https://github.com/runloopai/api-client-python/tree/main/\g<2>)'
+pattern = '\\[(.+?)\\]\\(((?!https?://)\\S+?)\\)'
+replacement = '[\\1](https://github.com/runloopai/api-client-python/tree/main/\\g<2>)'
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]
@@ -153,7 +155,7 @@ markers = [
 # there are a couple of flags that are still disabled by
 # default in strict mode as they are experimental and niche.
 typeCheckingMode = "strict"
-pythonVersion = "3.8"
+pythonVersion = "3.9"
 
 exclude = [
     "_dev",
@@ -223,7 +225,7 @@ ignore_missing_imports = true
 [tool.ruff]
 line-length = 120
 output-format = "grouped"
-target-version = "py38"
+target-version = "py39"
 
 [tool.ruff.format]
 docstring-code-format = true

requirements-dev.lock

@@ -97,10 +97,10 @@ pygments==2.18.0
 pyright==1.1.399
 pytest==8.3.3
     # via pytest-asyncio
+    # via pytest-timeout
     # via pytest-xdist
-pytest-timeout==2.3.1
-    # via runloop-api-client (dev)
 pytest-asyncio==0.24.0
+pytest-timeout==2.3.1
 pytest-xdist==3.7.0
 python-dateutil==2.8.2
     # via time-machine
@@ -128,6 +128,8 @@ typing-extensions==4.12.2
     # via pydantic-core
     # via pyright
     # via runloop-api-client
+uuid-utils==0.11.0
+    # via runloop-api-client
 virtualenv==20.24.5
     # via nox
 yarl==1.20.0

requirements.lock

@@ -68,5 +68,7 @@ typing-extensions==4.12.2
     # via pydantic
     # via pydantic-core
     # via runloop-api-client
+uuid-utils==0.11.0
+    # via runloop-api-client
 yarl==1.20.0
     # via aiohttp

src/runloop_api_client/_models.py

@@ -179,7 +179,7 @@ def __str__(self) -> str:
     @classmethod
     @override
     def construct(  # pyright: ignore[reportIncompatibleMethodOverride]
-        __cls: Type[ModelT],
+        __cls: Type[ModelT], # type: ignore
         _fields_set: set[str] | None = None,
         **values: object,
     ) -> ModelT:

src/runloop_api_client/resources/devboxes/devboxes.py

@@ -6,6 +6,7 @@
 from typing_extensions import Literal
 
 import httpx
+from uuid_utils import uuid7
 
 from .lsp import (
     LspResource,
@@ -787,6 +788,59 @@ def execute(
             cast_to=DevboxAsyncExecutionDetailView,
         )
 
+    def execute_and_await_completion(
+        self,
+        devbox_id: str,
+        *,
+        command: str,
+        shell_name: Optional[str] | NotGiven = NOT_GIVEN,
+        polling_config: PollingConfig | None = None,
+        # The following are forwarded to the initial execute request
+        extra_headers: Headers | None = None,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
+        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        idempotency_key: str | None = None,
+    ) -> DevboxAsyncExecutionDetailView:
+        """
+        Execute a command and wait for it to complete with optimal latency for long running commands.
+
+        This method launches an execution with a generated command_id and first attempts to
+        return the result within the initial request's timeout. If the execution is not yet
+        complete, it switches to using wait_for_command to minimize latency while waiting.
+        """
+        command_id = str(uuid7())
+        execution = self.execute(
+            devbox_id,
+            command=command,
+            command_id=command_id,
+            shell_name=shell_name,
+            extra_headers=extra_headers,
+            extra_query=extra_query,
+            extra_body=extra_body,
+            timeout=timeout,
+            idempotency_key=idempotency_key,
+        )
+        if execution.status == "completed":
+            return execution
+
+        def handle_timeout_error(error: Exception) -> DevboxAsyncExecutionDetailView:
+            if isinstance(error, APITimeoutError) or (
+                isinstance(error, APIStatusError) and error.response.status_code == 408
+            ):
+                return execution
+            raise error
+
+        def is_done(result: DevboxAsyncExecutionDetailView) -> bool:
+            return result.status == "completed"
+
+        return poll_until(
+            lambda: self.wait_for_command(execution.execution_id, devbox_id=devbox_id, statuses=["completed"]),
+            is_done,
+            polling_config,
+            handle_timeout_error,
+        )
+
     def execute_async(
         self,
         id: str,
@@ -2175,6 +2229,60 @@ async def execute(
             cast_to=DevboxAsyncExecutionDetailView,
         )
 
+    async def execute_and_await_completion(
+        self,
+        devbox_id: str,
+        *,
+        command: str,
+        shell_name: Optional[str] | NotGiven = NOT_GIVEN,
+        polling_config: PollingConfig | None = None,
+        # The following are forwarded to the initial execute request
+        extra_headers: Headers | None = None,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
+        timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+        idempotency_key: str | None = None,
+    ) -> DevboxAsyncExecutionDetailView:
+        """
+        Execute a command and wait for it to complete with optimal latency for long running commands.
+
+        This method launches an execution with a generated command_id and first attempts to
+        return the result within the initial request's timeout. If the execution is not yet
+        complete, it switches to using wait_for_command to minimize latency while waiting.
+        """
+
+        command_id = str(uuid7())
+        execution = await self.execute(
+            devbox_id,
+            command=command,
+            command_id=command_id,
+            shell_name=shell_name,
+            extra_headers=extra_headers,
+            extra_query=extra_query,
+            extra_body=extra_body,
+            timeout=timeout,
+            idempotency_key=idempotency_key,
+        )
+        if execution.status == "completed":
+            return execution
+
+        def handle_timeout_error(error: Exception) -> DevboxAsyncExecutionDetailView:
+            if isinstance(error, APITimeoutError) or (
+                isinstance(error, APIStatusError) and error.response.status_code == 408
+            ):
+                return execution
+            raise error
+
+        def is_done(result: DevboxAsyncExecutionDetailView) -> bool:
+            return result.status == "completed"
+
+        return await async_poll_until(
+            lambda: self.wait_for_command(execution.execution_id, devbox_id=devbox_id, statuses=["completed"]),
+            is_done,
+            polling_config,
+            handle_timeout_error,
+        )
+
     async def execute_async(
         self,
         id: str,

tests/smoketests/test_executions.py

@@ -61,3 +61,37 @@ def test_tail_stdout_logs() -> None:
         if received:
             break
     assert isinstance(received, str)
+
+
+@pytest.mark.timeout(30)
+def test_execute_and_await_completion() -> None:
+    assert _devbox_id
+    completed = client.devboxes.execute_and_await_completion(
+        _devbox_id,
+        command="echo hello && sleep 1",
+        polling_config=PollingConfig(max_attempts=120, interval_seconds=2.0, timeout_seconds=10 * 60),
+    )
+    assert completed.status == "completed"
+
+
+@pytest.mark.timeout(90)
+def test_execute_and_await_completion_long_running() -> None:
+    assert _devbox_id
+    completed = client.devboxes.execute_and_await_completion(
+        _devbox_id,
+        command="echo hello && sleep 70",
+        polling_config=PollingConfig(max_attempts=120, interval_seconds=2.0),
+    )
+    assert completed.status == "completed"
+
+
+# TODO: Uncomment this test when we fix timeouts for polling
+# @pytest.mark.timeout(30)
+# def test_execute_and_await_completion_timeout() -> None:
+#     assert _devbox_id
+#     with pytest.raises(PollingTimeout):
+#         client.devboxes.execute_and_await_completion(
+#             devbox_id=_devbox_id,
+#             command="echo hello && sleep 10",
+#             polling_config=PollingConfig(max_attempts=1, interval_seconds=2.0, timeout_seconds=3),
+#         )