Fix watch subscription timeout handling

python/src/docket/cli/_support.py

@@ -24,28 +24,51 @@ async def iterate_with_timeout(
     """Iterate over an async iterator with timeout, ensuring proper cleanup.
 
     Wraps an async iterator to add timeout support and guaranteed cleanup.
-    On timeout, yields None to allow the caller to handle polling fallback.
+    On timeout, yields None without cancelling the pending read.
 
     Args:
-        iterator: An async generator (must have __anext__ and aclose methods)
-        timeout: Timeout in seconds for each iteration
+        iterator: An async iterator with __anext__ and aclose methods.
+        timeout: Timeout in seconds for each iteration.
 
     Yields:
-        Items from the iterator, or None if timeout expires
+        Items from the iterator, or None if the timeout expires.
     """
     import asyncio
 
+    pending: asyncio.Task[dict[str, Any]] | None = None
+
     try:
         while True:
-            try:
-                yield await asyncio.wait_for(iterator.__anext__(), timeout=timeout)
-            except asyncio.TimeoutError:
-                # Yield None to signal timeout, allowing caller to handle polling
+            if pending is None:
+                pending = asyncio.create_task(iterator.__anext__())
+
+            done, _ = await asyncio.wait({pending}, timeout=timeout)
+
+            if not done:
                 yield None
+                continue
+
+            task = pending
+            pending = None
+
+            try:
+                item = task.result()
             except StopAsyncIteration:
                 break
+
+            yield item
     finally:
-        await iterator.aclose()
+        try:
+            if pending is not None:
+                if not pending.done():
+                    pending.cancel()
+
+                try:
+                    await pending
+                except (asyncio.CancelledError, StopAsyncIteration):
+                    pass
+        finally:
+            await iterator.aclose()
 
 
 class LogLevel(str, enum.Enum):

python/tests/cli/test_iterate_with_timeout.py

@@ -99,3 +99,60 @@ async def aclose(self) -> None:
             pass  # pragma: no cover - exception always raised before this
 
     assert close_called, "aclose() should be called even on exception"
+
+
+async def test_iterate_with_timeout_preserves_iterator_after_timeout():
+    """A timeout must not cancel the pending read from the iterator."""
+
+    async def delayed_iterator():
+        await asyncio.sleep(0.3)
+        yield {"value": "completed"}
+
+    results: list[dict[str, Any] | None] = []
+
+    async for item in iterate_with_timeout(
+        delayed_iterator(), timeout=0.1
+    ):
+        results.append(item)
+        if item is not None:
+            break
+
+    assert results[-1] == {"value": "completed"}
+    assert None in results[:-1]
+
+
+async def test_iterate_with_timeout_closes_iterator_when_pending_read_fails():
+    """Close the iterator even if the pending read raises an exception."""
+
+    class FailingIterator:
+        def __init__(self):
+            self.fail_read = asyncio.Event()
+            self.read_finished = asyncio.Event()
+            self.close_called = False
+
+        def __aiter__(self):
+            return self
+
+        async def __anext__(self):
+            await self.fail_read.wait()
+            try:
+                raise RuntimeError("read failed")
+            finally:
+                self.read_finished.set()
+
+        async def aclose(self):
+            self.close_called = True
+
+    iterator = FailingIterator()
+    wrapped = iterate_with_timeout(iterator, timeout=0.01)
+
+    assert await wrapped.__anext__() is None
+
+    iterator.fail_read.set()
+    await iterator.read_finished.wait()
+    await asyncio.sleep(0)
+
+    with pytest.raises(RuntimeError, match="read failed"):
+        await wrapped.aclose()
+
+    assert iterator.close_called