From ad92d5f3b5f96ba28fb82a2bce929ea8914c55dd Mon Sep 17 00:00:00 2001
From: Kayvan Zahiri <123409429+Kayvan-Zahiri@users.noreply.github.com>
Date: Sat, 17 Oct 2026 21:07:49 -0700
Subject: [PATCH] fix(polling): stop synchronous waits on cancelled tasks

---
 src/runwayml/lib/polling.py     |  4 +-
 tests/test_task_cancellation.py | 75 +++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_task_cancellation.py

diff --git a/src/runwayml/lib/polling.py b/src/runwayml/lib/polling.py
index 4c1f7377..2b55a9cc 100644
--- a/src/runwayml/lib/polling.py
+++ b/src/runwayml/lib/polling.py
@@ -93,7 +93,7 @@ def wait_for_task_output(self, timeout: Union[float, None] = 60 * 10) -> TaskRet
                 task_details = client.tasks.retrieve(self.id)
                 if task_details.status == "SUCCEEDED":
                     return task_details
-                if task_details.status == "FAILED":
+                if task_details.status == "FAILED" or task_details.status == "CANCELLED":
                     raise TaskFailedError(task_details)
                 if timeout is not None and time.time() - start_time > timeout:
                     raise TaskTimeoutError(task_details)
@@ -209,7 +209,7 @@ def wait_for_task_output(
             task_details = client.tasks.retrieve(self.id)  # type: ignore[attr-defined]
             if task_details.status == "SUCCEEDED":
                 return task_details
-            if task_details.status == "FAILED":
+            if task_details.status == "FAILED" or task_details.status == "CANCELLED":
                 raise TaskFailedError(task_details)
             if timeout is not None and time.time() - start_time > timeout:
                 raise TaskTimeoutError(task_details)
diff --git a/tests/test_task_cancellation.py b/tests/test_task_cancellation.py
new file mode 100644
index 00000000..74dd7ba0
--- /dev/null
+++ b/tests/test_task_cancellation.py
@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from typing import Any
+
+import anyio
+import httpx
+import pytest
+
+from runwayml import RunwayML, AsyncRunwayML, TaskFailedError
+
+
+@pytest.mark.parametrize("asynchronous", [False, True])
+@pytest.mark.parametrize("entry_point", ["create", "retrieve"])
+@pytest.mark.parametrize("status", ["CANCELLED", "FAILED", "SUCCEEDED"])
+async def test_public_task_wait_handles_terminal_status(
+    monkeypatch: pytest.MonkeyPatch, asynchronous: bool, entry_point: str, status: str
+) -> None:
+    monkeypatch.setattr("runwayml.lib.polling.time.sleep", lambda _delay: None)
+
+    async def no_sleep(_delay: float) -> None:
+        pass
+
+    monkeypatch.setattr(anyio, "sleep", no_sleep)
+    gets = 0
+    expected_gets = 2 if entry_point == "retrieve" else 1
+
+    def respond(request: httpx.Request) -> httpx.Response:
+        nonlocal gets
+        if request.method == "POST":
+            return httpx.Response(200, json={"id": "task-example"})
+        gets += 1
+        assert gets <= expected_gets, "polled again after a terminal task response"
+        payload: dict[str, Any] = {
+            "id": "task-example",
+            "createdAt": "2026-01-01T00:00:00Z",
+            "status": status,
+            "cost": {"credits": 0},
+        }
+        if status == "FAILED":
+            payload["failure"] = "test failure"
+        if status == "SUCCEEDED":
+            payload["output"] = ["https://example.com/output.png"]
+        return httpx.Response(200, json=payload)
+
+    if asynchronous:
+        async with AsyncRunwayML(
+            api_key="test", max_retries=0, http_client=httpx.AsyncClient(transport=httpx.MockTransport(respond))
+        ) as client:
+            response = (
+                await client.generate.image.create(config_id="test", input={"prompt_text": "A test image"})
+                if entry_point == "create"
+                else await client.tasks.retrieve("task-example")
+            )
+            if status == "SUCCEEDED":
+                assert (await response.wait_for_task_output()).status == status
+            else:
+                with pytest.raises(TaskFailedError) as error:
+                    await response.wait_for_task_output()
+                assert error.value.task_details.status == status
+    else:
+        with RunwayML(
+            api_key="test", max_retries=0, http_client=httpx.Client(transport=httpx.MockTransport(respond))
+        ) as sync_client:
+            sync_response = (
+                sync_client.generate.image.create(config_id="test", input={"prompt_text": "A test image"})
+                if entry_point == "create"
+                else sync_client.tasks.retrieve("task-example")
+            )
+            if status == "SUCCEEDED":
+                assert sync_response.wait_for_task_output().status == status
+            else:
+                with pytest.raises(TaskFailedError) as sync_error:
+                    sync_response.wait_for_task_output()
+                assert sync_error.value.task_details.status == status
+    assert gets == expected_gets