feat: add CancellationToken support to polling infrastructure

- Create CancellationToken class with sync/async event support
- Update poll_until() and async_poll_until() with cancellation_token parameter
- Add cancellable sleep using threading.Event.wait() and asyncio.wait_for()
- Update PollingRequestOptions TypedDict with cancellation_token field
- Propagate cancellation_token through Blueprint and ScenarioRun polling methods

Part of porting TypeScript PR #765 features to Python SDK.

Author: tode-rl

src/runloop_api_client/lib/cancellation.py

@@ -0,0 +1,103 @@
+"""Cancellation support for polling operations."""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+from typing import TYPE_CHECKING
+
+from .._exceptions import RunloopError
+
+if TYPE_CHECKING:
+    pass
+
+__all__ = ["PollingCancelled", "CancellationToken"]
+
+
+class PollingCancelled(RunloopError):
+    """Exception raised when a polling operation is cancelled."""
+
+    pass
+
+
+class CancellationToken:
+    """Thread-safe cancellation token for polling operations.
+
+    Similar to JavaScript's AbortSignal. Works in both sync and async contexts.
+
+    Example (sync):
+        >>> token = CancellationToken()
+        >>> # In another thread:
+        >>> token.cancel()
+        >>> # In polling code:
+        >>> token.raise_if_cancelled()  # Raises PollingCancelled
+
+    Example (async):
+        >>> token = CancellationToken()
+        >>> # In another task:
+        >>> token.cancel()
+        >>> # In async polling code:
+        >>> await asyncio.wait_for(token.async_event.wait(), timeout=1.0)
+    """
+
+    def __init__(self) -> None:
+        """Create a new cancellation token."""
+        self._cancelled = False
+        self._sync_event = threading.Event()
+        self._async_event: asyncio.Event | None = None
+        self._lock = threading.Lock()
+
+    def cancel(self) -> None:
+        """Mark this token as cancelled.
+
+        Thread-safe and can be called multiple times. Sets both sync and async events.
+        """
+        with self._lock:
+            if self._cancelled:
+                return
+            self._cancelled = True
+            self._sync_event.set()
+            if self._async_event is not None:
+                self._async_event.set()
+
+    def is_cancelled(self) -> bool:
+        """Check if this token has been cancelled.
+
+        Returns:
+            True if cancel() has been called, False otherwise.
+        """
+        return self._cancelled
+
+    def raise_if_cancelled(self) -> None:
+        """Raise PollingCancelled if this token has been cancelled.
+
+        Raises:
+            PollingCancelled: If cancel() has been called.
+        """
+        if self._cancelled:
+            raise PollingCancelled("Polling operation was cancelled")
+
+    @property
+    def sync_event(self) -> threading.Event:
+        """Get the synchronous event for cancellation checking.
+
+        Returns:
+            threading.Event that is set when cancel() is called.
+        """
+        return self._sync_event
+
+    @property
+    def async_event(self) -> asyncio.Event:
+        """Get the asynchronous event for cancellation checking.
+
+        Lazily creates the async event on first access. If cancel() was already called,
+        the event will be set immediately.
+
+        Returns:
+            asyncio.Event that is set when cancel() is called.
+        """
+        if self._async_event is None:
+            self._async_event = asyncio.Event()
+            if self._cancelled:
+                self._async_event.set()
+        return self._async_event

src/runloop_api_client/lib/polling.py

@@ -2,6 +2,8 @@
 from typing import Any, TypeVar, Callable, Optional
 from dataclasses import dataclass
 
+from .cancellation import CancellationToken
+
 T = TypeVar("T")
 
 
@@ -27,6 +29,7 @@ def poll_until(
     is_terminal: Callable[[T], bool],
     config: Optional[PollingConfig] = None,
     on_error: Optional[Callable[[Exception], T]] = None,
+    cancellation_token: Optional[CancellationToken] = None,
 ) -> T:
     """
     Poll until a condition is met or timeout/max attempts are reached.
@@ -37,12 +40,14 @@ def poll_until(
         config: Optional polling configuration
         on_error: Optional error handler that can return a value to continue polling
                  or re-raise the exception to stop polling
+        cancellation_token: Optional token to cancel the polling operation
 
     Returns:
         The final state of the polled object
 
     Raises:
         PollingTimeout: When max attempts or timeout is reached
+        PollingCancelled: If cancellation_token.cancel() is called
     """
     if config is None:
         config = PollingConfig()
@@ -52,6 +57,10 @@ def poll_until(
     last_result = None
 
     while True:
+        # Check for cancellation before each iteration
+        if cancellation_token is not None:
+            cancellation_token.raise_if_cancelled()
+
         try:
             last_result = retriever()
         except Exception as e:
@@ -72,4 +81,9 @@ def poll_until(
             if elapsed >= config.timeout_seconds:
                 raise PollingTimeout(f"Exceeded timeout of {config.timeout_seconds} seconds", last_result)
 
-        time.sleep(config.interval_seconds)
+        # Cancellable sleep
+        if cancellation_token is not None:
+            if cancellation_token.sync_event.wait(timeout=config.interval_seconds):
+                cancellation_token.raise_if_cancelled()
+        else:
+            time.sleep(config.interval_seconds)

src/runloop_api_client/lib/polling_async.py

@@ -3,6 +3,7 @@
 from typing import Union, TypeVar, Callable, Optional, Awaitable
 
 from .polling import PollingConfig, PollingTimeout
+from .cancellation import CancellationToken
 
 T = TypeVar("T")
 
@@ -12,6 +13,7 @@ async def async_poll_until(
     is_terminal: Callable[[T], bool],
     config: Optional[PollingConfig] = None,
     on_error: Optional[Callable[[Exception], T]] = None,
+    cancellation_token: Optional[CancellationToken] = None,
 ) -> T:
     """
     Poll until a condition is met or timeout/max attempts are reached.
@@ -22,12 +24,14 @@ async def async_poll_until(
         config: Optional polling configuration
         on_error: Optional error handler that can return a value to continue polling
                  or re-raise the exception to stop polling
+        cancellation_token: Optional token to cancel the polling operation
 
     Returns:
         The final state of the polled object
 
     Raises:
         PollingTimeout: When max attempts or timeout is reached
+        PollingCancelled: If cancellation_token.cancel() is called
     """
     if config is None:
         config = PollingConfig()
@@ -37,6 +41,10 @@ async def async_poll_until(
     last_result: Union[T, None] = None
 
     while True:
+        # Check for cancellation before each iteration
+        if cancellation_token is not None:
+            cancellation_token.raise_if_cancelled()
+
         try:
             last_result = await retriever()
         except Exception as e:
@@ -57,4 +65,15 @@ async def async_poll_until(
             if elapsed >= config.timeout_seconds:
                 raise PollingTimeout(f"Exceeded timeout of {config.timeout_seconds} seconds", last_result)
 
-        await asyncio.sleep(config.interval_seconds)
+        # Cancellable async sleep
+        if cancellation_token is not None:
+            try:
+                await asyncio.wait_for(
+                    cancellation_token.async_event.wait(),
+                    timeout=config.interval_seconds,
+                )
+                cancellation_token.raise_if_cancelled()
+            except asyncio.TimeoutError:
+                pass  # Normal sleep completion
+        else:
+            await asyncio.sleep(config.interval_seconds)

src/runloop_api_client/resources/blueprints.py

@@ -29,6 +29,7 @@
 from ..lib.polling import PollingConfig, poll_until
 from .._base_client import AsyncPaginator, make_request_options
 from ..lib.polling_async import async_poll_until
+from ..lib.cancellation import CancellationToken
 from .._utils._validation import ValidationNotification
 from ..types.blueprint_view import BlueprintView
 from ..types.blueprint_preview_view import BlueprintPreviewView
@@ -41,6 +42,7 @@
 # Type for request arguments that combine polling config with additional request options
 class BlueprintRequestArgs(TypedDict, total=False):
     polling_config: PollingConfig | None
+    cancellation_token: CancellationToken | None
     # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
     # The extra values given here take precedence over values defined on the client or passed to this method.
     extra_headers: Headers | None
@@ -280,6 +282,7 @@ def await_build_complete(
         id: str,
         *,
         polling_config: PollingConfig | None = None,
+        cancellation_token: CancellationToken | None = None,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
@@ -292,6 +295,7 @@ def await_build_complete(
         Args:
             id: The ID of the blueprint to wait for
             polling_config: Optional polling configuration
+            cancellation_token: Token to cancel the wait operation
             extra_headers: Send extra headers
             extra_query: Add additional query parameters to the request
             extra_body: Add additional JSON properties to the request
@@ -302,6 +306,7 @@ def await_build_complete(
 
         Raises:
             PollingTimeout: If polling times out before blueprint is built
+            PollingCancelled: If cancellation_token.cancel() is called
             RunloopError: If blueprint enters a non-built terminal state
         """
 
@@ -313,7 +318,12 @@ def retrieve_blueprint() -> BlueprintView:
         def is_done_building(blueprint: BlueprintView) -> bool:
             return blueprint.status not in ["queued", "building", "provisioning"]
 
-        blueprint = poll_until(retrieve_blueprint, is_done_building, polling_config)
+        blueprint = poll_until(
+            retrieve_blueprint,
+            is_done_building,
+            polling_config,
+            cancellation_token=cancellation_token,
+        )
 
         if blueprint.status != "build_complete":
             raise RunloopError(f"Blueprint entered non-built terminal state: {blueprint.status}")
@@ -338,6 +348,7 @@ def create_and_await_build_complete(
         services: Optional[Iterable[blueprint_create_params.Service]] | Omit = omit,
         system_setup_commands: Optional[SequenceNotStr[str]] | Omit = omit,
         polling_config: PollingConfig | None = None,
+        cancellation_token: CancellationToken | None = None,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
@@ -353,12 +364,14 @@ def create_and_await_build_complete(
         Args:
             See the `create` method for detailed documentation.
             polling_config: Optional polling configuration
+            cancellation_token: Token to cancel the wait operation
 
         Returns:
             The built blueprint
 
         Raises:
             PollingTimeout: If polling times out before blueprint is built
+            PollingCancelled: If cancellation_token.cancel() is called
             RunloopError: If blueprint enters a non-built terminal state
         """
         # Pass all create_args to the underlying create method
@@ -387,6 +400,7 @@ def create_and_await_build_complete(
         return self.await_build_complete(
             blueprint.id,
             polling_config=polling_config,
+            cancellation_token=cancellation_token,
             extra_headers=extra_headers,
             extra_query=extra_query,
             extra_body=extra_body,
@@ -960,6 +974,7 @@ async def await_build_complete(
         id: str,
         *,
         polling_config: PollingConfig | None = None,
+        cancellation_token: CancellationToken | None = None,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
@@ -972,6 +987,7 @@ async def await_build_complete(
         Args:
             id: The ID of the blueprint to wait for
             polling_config: Optional polling configuration
+            cancellation_token: Token to cancel the wait operation
             extra_headers: Send extra headers
             extra_query: Add additional query parameters to the request
             extra_body: Add additional JSON properties to the request
@@ -982,6 +998,7 @@ async def await_build_complete(
 
         Raises:
             PollingTimeout: If polling times out before blueprint is built
+            PollingCancelled: If cancellation_token.cancel() is called
             RunloopError: If blueprint enters a non-built terminal state
         """
 
@@ -993,7 +1010,12 @@ async def retrieve_blueprint() -> BlueprintView:
         def is_done_building(blueprint: BlueprintView) -> bool:
             return blueprint.status not in ["queued", "building", "provisioning"]
 
-        blueprint = await async_poll_until(retrieve_blueprint, is_done_building, polling_config)
+        blueprint = await async_poll_until(
+            retrieve_blueprint,
+            is_done_building,
+            polling_config,
+            cancellation_token=cancellation_token,
+        )
 
         if blueprint.status != "build_complete":
             raise RunloopError(f"Blueprint entered non-built terminal state: {blueprint.status}")
@@ -1018,6 +1040,7 @@ async def create_and_await_build_complete(
         services: Optional[Iterable[blueprint_create_params.Service]] | Omit = omit,
         system_setup_commands: Optional[SequenceNotStr[str]] | Omit = omit,
         polling_config: PollingConfig | None = None,
+        cancellation_token: CancellationToken | None = None,
         # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
         # The extra values given here take precedence over values defined on the client or passed to this method.
         extra_headers: Headers | None = None,
@@ -1033,12 +1056,14 @@ async def create_and_await_build_complete(
         Args:
             See the `create` method for detailed documentation.
             polling_config: Optional polling configuration
+            cancellation_token: Token to cancel the wait operation
 
         Returns:
             The built blueprint
 
         Raises:
             PollingTimeout: If polling times out before blueprint is built
+            PollingCancelled: If cancellation_token.cancel() is called
             RunloopError: If blueprint enters a non-built terminal state
         """
         # Pass all create_args to the underlying create method
@@ -1067,6 +1092,7 @@ async def create_and_await_build_complete(
         return await self.await_build_complete(
             blueprint.id,
             polling_config=polling_config,
+            cancellation_token=cancellation_token,
             extra_headers=extra_headers,
             extra_query=extra_query,
             extra_body=extra_body,