diff --git a/CHANGELOG.md b/CHANGELOG.md index a1f06a6..62654b6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,11 @@ the fuller account of each version, including verification notes. ## [Unreleased] +### 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 diff --git a/src/comfy_low/transport.py b/src/comfy_low/transport.py index 1e550c8..fa82991 100644 --- a/src/comfy_low/transport.py +++ b/src/comfy_low/transport.py @@ -53,7 +53,7 @@ import platform import secrets import sys -from collections.abc import AsyncIterator, Iterator, Mapping +from collections.abc import AsyncGenerator, AsyncIterator, Generator, Iterator, Mapping from contextlib import asynccontextmanager, contextmanager from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -1330,7 +1330,9 @@ def get_job(self, job_id_or_url: str, *, timeout: Any = _UNSET) -> Job: resp = self.raw_request("GET", path, timeout=timeout) return _job(self._p.parse_or_raise(resp, (200,))) - def get_job_events(self, job_id_or_url: str, *, timeout: Any = _UNSET) -> Iterator[RawEvent]: + def get_job_events( + self, job_id_or_url: str, *, timeout: Any = _UNSET + ) -> Generator[RawEvent, None, None]: """GET /api/v2/jobs/{id}/events — raw live SSE iterator (escape hatch). No reconnection here; a single connection's frames. ``comfy_sdk`` adds the @@ -1851,7 +1853,7 @@ async def get_job(self, job_id_or_url: str, *, timeout: Any = _UNSET) -> Job: async def get_job_events( self, job_id_or_url: str, *, timeout: Any = _UNSET - ) -> AsyncIterator[RawEvent]: + ) -> AsyncGenerator[RawEvent, None]: # Read-idle timeout by default (see the sync get_job_events); a stalled # connection raises so comfy_sdk falls back to poll/reconnect. path = job_id_or_url if _looks_like_path(job_id_or_url) else f"/jobs/{job_id_or_url}/events" diff --git a/src/comfy_sdk/jobs.py b/src/comfy_sdk/jobs.py index e0d37de..832b431 100644 --- a/src/comfy_sdk/jobs.py +++ b/src/comfy_sdk/jobs.py @@ -12,6 +12,7 @@ import time from collections.abc import AsyncIterator, Iterator +from contextlib import aclosing, closing from dataclasses import dataclass, field from datetime import datetime from typing import Any, Literal @@ -262,15 +263,16 @@ 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 + with closing(self._low.get_job_events(events_url)) as raw_events: + 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 except ApiError as exc: if exc.http_status == 501: return # surface has no SSE — poll paths remain authoritative @@ -379,15 +381,16 @@ 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 + async with aclosing(self._low.get_job_events(events_url)) as raw_events: + 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 except ApiError as exc: if exc.http_status == 501: return # surface has no SSE — poll paths remain authoritative diff --git a/tests/test_event_stream_cleanup.py b/tests/test_event_stream_cleanup.py new file mode 100644 index 0000000..63f0341 --- /dev/null +++ b/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()