fix: close owned job event streams

CHANGELOG.md

@@ -18,6 +18,11 @@ the fuller account of each version, including verification notes.
   deliberately not `RequestNotFound`, which would say the id was never valid. Running the model
   again means submitting a new request. Before this, the bucket arrived as a bare `RouterError`.
 
+### Fixed
+
+- Job event iterators close their underlying HTTP stream when iteration completes or the
+  iterator is explicitly closed, for both synchronous and asynchronous clients.
+
 ## [0.5.0] - 2026-10-10
 
 ### Added

src/comfy_sdk/jobs.py

@@ -262,15 +262,21 @@ def events(self) -> Iterator[Event]:
         while True:
             terminal_seen = False
             try:
-                for raw in self._low.get_job_events(events_url):
-                    ev = event_from_raw(raw, self._bind_output)
-                    if ev is None:
-                        continue
-                    if isinstance(ev, StatusChange) and _core.is_terminal(ev.status):
-                        terminal_seen = True
+                raw_events = self._low.get_job_events(events_url)
+                try:
+                    for raw in raw_events:
+                        ev = event_from_raw(raw, self._bind_output)
+                        if ev is None:
+                            continue
+                        if isinstance(ev, StatusChange) and _core.is_terminal(ev.status):
+                            terminal_seen = True
+                            yield ev
+                            return
                         yield ev
-                        return
-                    yield ev
+                finally:
+                    close = getattr(raw_events, "close", None)
+                    if callable(close):
+                        close()
             except ApiError as exc:
                 if exc.http_status == 501:
                     return  # surface has no SSE — poll paths remain authoritative
@@ -379,15 +385,21 @@ async def events(self) -> AsyncIterator[Event]:
         while True:
             terminal_seen = False
             try:
-                async for raw in self._low.get_job_events(events_url):
-                    ev = event_from_raw(raw, self._bind_output)
-                    if ev is None:
-                        continue
-                    if isinstance(ev, StatusChange) and _core.is_terminal(ev.status):
-                        terminal_seen = True
+                raw_events = self._low.get_job_events(events_url)
+                try:
+                    async for raw in raw_events:
+                        ev = event_from_raw(raw, self._bind_output)
+                        if ev is None:
+                            continue
+                        if isinstance(ev, StatusChange) and _core.is_terminal(ev.status):
+                            terminal_seen = True
+                            yield ev
+                            return
                         yield ev
-                        return
-                    yield ev
+                finally:
+                    aclose = getattr(raw_events, "aclose", None)
+                    if callable(aclose):
+                        await aclose()
             except ApiError as exc:
                 if exc.http_status == 501:
                     return  # surface has no SSE — poll paths remain authoritative

tests/test_event_iterator_overrides.py

@@ -0,0 +1,144 @@
+"""Iterator-only low-level overrides stay valid when high-level events close."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Iterator
+from typing import Any
+
+import pytest
+
+from comfy_low.models import Job as LowJob
+from comfy_low.sse import RawEvent
+from comfy_low.transport import AsyncComfyLow, ComfyLow
+from comfy_sdk import AsyncJob, Job, Progress, StatusChange
+
+
+def _model(status: str = "running") -> LowJob:
+    return LowJob.model_validate(
+        {
+            "id": "iterator-job",
+            "status": status,
+            "created_at": "2026-01-01T00:00:00Z",
+            "started_at": None,
+            "completed_at": None,
+            "expires_at": "2026-01-02T00:00:00Z",
+            "queue_position": None,
+            "progress": None,
+            "outputs": [],
+            "error": None,
+            "urls": {"self": "/jobs/iterator-job", "events": "/events", "cancel": "/cancel"},
+        }
+    )
+
+
+class _Events(Iterator[RawEvent]):
+    def __init__(self, frames: list[RawEvent]) -> None:
+        self.frames = iter(frames)
+        self.reads = 0
+
+    def __next__(self) -> RawEvent:
+        self.reads += 1
+        return next(self.frames)
+
+
+class _AsyncEvents(AsyncIterator[RawEvent]):
+    def __init__(self, frames: list[RawEvent]) -> None:
+        self.frames = iter(frames)
+        self.reads = 0
+
+    async def __anext__(self) -> RawEvent:
+        self.reads += 1
+        try:
+            return next(self.frames)
+        except StopIteration:
+            raise StopAsyncIteration from None
+
+
+class _IteratorLow(ComfyLow):
+    def __init__(self, events: Iterator[RawEvent]) -> None:
+        super().__init__("http://iterator.invalid")
+        self.events = events
+        self.stream_calls = 0
+        self.poll_calls = 0
+
+    def get_job_events(self, job_id_or_url: str, *, timeout: Any = None) -> Iterator[RawEvent]:
+        self.stream_calls += 1
+        return self.events
+
+    def get_job(self, job_id_or_url: str, *, timeout: Any = None) -> LowJob:
+        self.poll_calls += 1
+        return _model("succeeded")
+
+
+class _IteratorAsyncLow(AsyncComfyLow):
+    def __init__(self, events: AsyncIterator[RawEvent]) -> None:
+        super().__init__("http://iterator.invalid")
+        self.events = events
+        self.stream_calls = 0
+        self.poll_calls = 0
+
+    def get_job_events(self, job_id_or_url: str, *, timeout: Any = None) -> AsyncIterator[RawEvent]:
+        self.stream_calls += 1
+        return self.events
+
+    async def get_job(self, job_id_or_url: str, *, timeout: Any = None) -> LowJob:
+        self.poll_calls += 1
+        return _model("succeeded")
+
+
+def _frames(exit_mode: str) -> list[RawEvent]:
+    if exit_mode == "empty":
+        return []
+    if exit_mode == "close":
+        return [RawEvent(event="progress", data={"value": 0.5})]
+    return [RawEvent(event="status", data={"status": "succeeded"})]
+
+
+@pytest.mark.parametrize("exit_mode", ["terminal", "close", "empty"])
+@pytest.mark.parametrize("cleanup", ["absent", "noncallable", "callable"])
+def test_events_accept_iterator_overrides(exit_mode, cleanup):
+    raw = _Events(_frames(exit_mode))
+    closed = []
+    if cleanup != "absent":
+        raw.close = (lambda: closed.append(True)) if cleanup == "callable" else None
+    with _IteratorLow(raw) as low:
+        events = Job(low, _model()).events()
+        try:
+            if exit_mode == "close":
+                assert next(events) == Progress(value=0.5)
+                events.close()
+            else:
+                assert list(events) == [StatusChange(status="succeeded")]
+            assert low.stream_calls == 1
+            assert low.poll_calls == (1 if exit_mode == "empty" else 0)
+            assert raw.reads == 1
+            assert closed == ([True] if cleanup == "callable" else [])
+        finally:
+            events.close()
+
+
+@pytest.mark.parametrize("exit_mode", ["terminal", "close", "empty"])
+@pytest.mark.parametrize("cleanup", ["absent", "noncallable", "callable"])
+async def test_events_accept_async_iterator_overrides(exit_mode, cleanup):
+    raw = _AsyncEvents(_frames(exit_mode))
+    closed = []
+
+    async def close():
+        closed.append(True)
+
+    if cleanup != "absent":
+        raw.aclose = close if cleanup == "callable" else None
+    async with _IteratorAsyncLow(raw) as low:
+        events = AsyncJob(low, _model()).events()
+        try:
+            if exit_mode == "close":
+                assert await anext(events) == Progress(value=0.5)
+                await events.aclose()
+            else:
+                assert [event async for event in events] == [StatusChange(status="succeeded")]
+            assert low.stream_calls == 1
+            assert low.poll_calls == (1 if exit_mode == "empty" else 0)
+            assert raw.reads == 1
+            assert closed == ([True] if cleanup == "callable" else [])
+        finally:
+            await events.aclose()

tests/test_event_stream_cleanup.py

@@ -0,0 +1,115 @@
+"""High-level event iteration owns each low-level HTTP stream."""
+
+from __future__ import annotations
+
+import asyncio
+from contextlib import asynccontextmanager, contextmanager
+
+import pytest
+
+from comfy_sdk import AsyncComfy, Comfy, StatusChange
+
+
+@pytest.mark.parametrize("exit_mode", ["terminal", "close"])
+def test_events_close_upstream(server, monkeypatch, exit_mode):
+    with Comfy() as client:
+        job = client.submit(
+            client.workflows.from_json({"3": {"class_type": "KSampler", "inputs": {}}})
+        )
+        low = job._low
+        original_events = low.get_job_events
+        original_open = low.open
+        retained = []
+        responses = []
+        exited = []
+
+        def retain_events(*args, **kwargs):
+            stream = original_events(*args, **kwargs)
+            retained.append(stream)
+            return stream
+
+        @contextmanager
+        def observe_open(*args, **kwargs):
+            try:
+                with original_open(*args, **kwargs) as response:
+                    responses.append(response)
+                    yield response
+            finally:
+                exited.append(True)
+
+        monkeypatch.setattr(low, "get_job_events", retain_events)
+        monkeypatch.setattr(low, "open", observe_open)
+        stream = job.events()
+        try:
+            if exit_mode == "terminal":
+                seen = list(stream)
+                assert isinstance(seen[-1], StatusChange)
+                assert seen[-1].status == "succeeded"
+            else:
+                next(stream)
+                stream.close()
+            assert len(retained) == 1
+            assert exited == [True]
+            assert responses[0].is_closed
+        finally:
+            stream.close()
+            for upstream in retained:
+                upstream.close()
+
+
+@pytest.mark.parametrize("exit_mode", ["terminal", "close", "cancel"])
+async def test_async_events_close_upstream(server, monkeypatch, exit_mode):
+    if exit_mode == "cancel":
+        server.state.sse_mode = "stall"
+        server.state.stall_seconds = 0.2
+    async with AsyncComfy() as client:
+        job = await client.submit(
+            client.workflows.from_json({"3": {"class_type": "KSampler", "inputs": {}}})
+        )
+        low = job._low
+        original_events = low.get_job_events
+        original_open = low.open
+        retained = []
+        responses = []
+        exited = []
+
+        def retain_events(*args, **kwargs):
+            stream = original_events(*args, **kwargs)
+            retained.append(stream)
+            return stream
+
+        @asynccontextmanager
+        async def observe_open(*args, **kwargs):
+            try:
+                async with original_open(*args, **kwargs) as response:
+                    responses.append(response)
+                    yield response
+            finally:
+                exited.append(True)
+
+        monkeypatch.setattr(low, "get_job_events", retain_events)
+        monkeypatch.setattr(low, "open", observe_open)
+        stream = job.events()
+        try:
+            if exit_mode == "terminal":
+                seen = [event async for event in stream]
+                assert isinstance(seen[-1], StatusChange)
+                assert seen[-1].status == "succeeded"
+            elif exit_mode == "close":
+                await anext(stream)
+                await stream.aclose()
+            else:
+                await anext(stream)
+                await anext(stream)
+                pending = asyncio.create_task(anext(stream))
+                await asyncio.sleep(0)
+                pending.cancel()
+                with pytest.raises(asyncio.CancelledError):
+                    await pending
+            assert len(retained) == 1
+            assert exited == [True]
+            assert responses[0].is_closed
+        finally:
+            await stream.aclose()
+            for upstream in retained:
+                await upstream.aclose()