fix(dify-agent): end SSE streams on terminal events (#40042)

dify-agent/docs/dify-agent/guide/index.md

@@ -335,7 +335,10 @@ progress:
   `next_cursor` cursors.
 - `GET /runs/{run_id}/events/sse` replays and streams events over SSE. The SSE
   `id` is the event Redis Stream ID. `after` query cursors take precedence over
-  `Last-Event-ID` headers.
+  `Last-Event-ID` headers. The server closes the SSE response normally after
+  delivering a terminal event. Clients must stop reconnecting after consuming
+  that event. Both cursor forms remain exclusive resume cursors, so the server
+  does not resend a terminal event that the supplied cursor already excludes.
 
 Successful runs emit `run_started`, zero or more `pydantic_ai_event`, and
 `run_succeeded`. Failed runs end with `run_failed`, and accepted cancellations

dify-agent/src/dify_agent/client/_client.py

@@ -600,11 +600,16 @@ class Client:
         with an id, reconnects resume from that id using the ``after`` query
         parameter. HTTP 5xx stream responses are retried, but HTTP 4xx responses,
         DTO validation failures, and malformed SSE frames are not retried. By
-        default iteration stops after a succeeded, failed, or cancelled terminal event.
+        default, ``until_terminal=True`` returns immediately after yielding a
+        succeeded, failed, or cancelled terminal event. With
+        ``until_terminal=False``, iteration may consume the remainder of the current
+        response, but after observing a terminal event it will not reconnect when that
+        response ends normally or raises a reconnectable transport error.
         """
         _validate_stream_options(max_reconnects, reconnect_delay_seconds, timeout_seconds)
         cursor = after or "0-0"
         reconnect_attempts = 0
+        terminal_event_seen = False
         deadline = time.monotonic() + timeout_seconds if timeout_seconds is not None else None
         while True:
             _raise_if_stream_stopped(run_id, deadline=deadline, should_stop=should_stop)
@@ -617,10 +622,14 @@ class Client:
                 ):
                     if event.id is not None:
                         cursor = event.id
+                    if event.type in _TERMINAL_EVENT_TYPES:
+                        terminal_event_seen = True
                     yield event
-                    if until_terminal and event.type in _TERMINAL_EVENT_TYPES:
+                    if until_terminal and terminal_event_seen:
                         return
             except _ReconnectableStreamError as exc:
+                if terminal_event_seen:
+                    return
                 if not reconnect:
                     raise exc.error from exc
                 reconnect_attempts = _next_reconnect_attempt(
@@ -631,6 +640,8 @@ class Client:
                 _raise_if_stream_stopped(run_id, deadline=deadline, should_stop=should_stop)
                 await _sleep_async(_bounded_sleep_seconds(reconnect_delay_seconds, deadline))
                 continue
+            if terminal_event_seen:
+                return
             if not reconnect:
                 return
             reconnect_attempts = _next_reconnect_attempt(
@@ -657,6 +668,7 @@ class Client:
         _validate_stream_options(max_reconnects, reconnect_delay_seconds, timeout_seconds)
         cursor = after or "0-0"
         reconnect_attempts = 0
+        terminal_event_seen = False
         deadline = time.monotonic() + timeout_seconds if timeout_seconds is not None else None
         while True:
             _raise_if_stream_stopped(run_id, deadline=deadline, should_stop=should_stop)
@@ -669,10 +681,14 @@ class Client:
                 ):
                     if event.id is not None:
                         cursor = event.id
+                    if event.type in _TERMINAL_EVENT_TYPES:
+                        terminal_event_seen = True
                     yield event
-                    if until_terminal and event.type in _TERMINAL_EVENT_TYPES:
+                    if until_terminal and terminal_event_seen:
                         return
             except _ReconnectableStreamError as exc:
+                if terminal_event_seen:
+                    return
                 if not reconnect:
                     raise exc.error from exc
                 reconnect_attempts = _next_reconnect_attempt(
@@ -683,6 +699,8 @@ class Client:
                 _raise_if_stream_stopped(run_id, deadline=deadline, should_stop=should_stop)
                 _sleep_sync(_bounded_sleep_seconds(reconnect_delay_seconds, deadline))
                 continue
+            if terminal_event_seen:
+                return
             if not reconnect:
                 return
             reconnect_attempts = _next_reconnect_attempt(

dify-agent/src/dify_agent/storage/redis_run_store.py

@@ -26,6 +26,8 @@ from dify_agent.server.schemas import RunRecord, new_run_id
 from dify_agent.server.settings import DEFAULT_RUN_RETENTION_SECONDS
 from dify_agent.storage.redis_keys import run_events_key, run_record_key
 
+_TERMINAL_RUN_EVENT_TYPES = {"run_succeeded", "run_failed", "run_cancelled"}
+
 
 class RunNotFoundError(LookupError):
     """Raised when a requested run record does not exist."""
@@ -190,7 +192,7 @@ class RedisRunStore(RunEventSink):
         return RunEventsResponse(run_id=run_id, events=events, next_cursor=next_cursor)
 
     async def iter_events(self, run_id: str, *, after: str = "0-0") -> AsyncIterator[RunEvent]:
-        """Yield replayed and future events for SSE clients."""
+        """Yield replayed and future events through the first terminal event."""
         await self.get_run(run_id)
         cursor = after
         while True:
@@ -199,6 +201,8 @@ class RedisRunStore(RunEventSink):
                 if event.id is not None:
                     cursor = event.id
                 yield event
+                if event.type in _TERMINAL_RUN_EVENT_TYPES:
+                    return
             if not page.events:
                 break
         while True:
@@ -211,6 +215,8 @@ class RedisRunStore(RunEventSink):
                     if event.id is not None:
                         cursor = event.id
                     yield event
+                    if event.type in _TERMINAL_RUN_EVENT_TYPES:
+                        return
 
     @staticmethod
     def _decode_event(run_id: str, raw_id: object, fields: dict[object, object]) -> RunEvent:

dify-agent/tests/local/dify_agent/client/test_client.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import asyncio
 import json
-from collections.abc import Iterator
+from collections.abc import AsyncIterator, Iterator
 from datetime import UTC, datetime
 from typing import cast, override
 
@@ -125,6 +125,19 @@ class DisconnectingSyncStream(httpx.SyncByteStream):
         raise httpx.ReadError("stream disconnected")
 
 
+class DisconnectingAsyncStream(httpx.AsyncByteStream):
+    chunks: list[bytes]
+
+    def __init__(self, *chunks: str) -> None:
+        self.chunks = [chunk.encode() for chunk in chunks]
+
+    @override
+    async def __aiter__(self) -> AsyncIterator[bytes]:
+        for chunk in self.chunks:
+            yield chunk
+        raise httpx.ReadError("stream disconnected")
+
+
 def test_sse_decoder_accepts_function_tool_result_part_alias(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setattr(client_module, "_function_tool_result_payload_key_cache", "part")
     decoder = client_module._SSEDecoder()
@@ -624,6 +637,44 @@ def test_stream_events_stops_after_cancelled_terminal_event() -> None:
     assert calls == 1
 
 
+def test_stream_events_does_not_reconnect_after_terminal_when_until_terminal_is_false() -> None:
+    calls = 0
+
+    def handler(_request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        return httpx.Response(200, content=_event_frame(_run_succeeded_event()))
+
+    client = Client(
+        base_url="http://testserver",
+        sync_http_client=httpx.Client(transport=httpx.MockTransport(handler)),
+    )
+
+    events = list(client.stream_events_sync("run-1", until_terminal=False, reconnect_delay_seconds=0))
+
+    assert [event.type for event in events] == ["run_succeeded"]
+    assert calls == 1
+
+
+def test_stream_events_does_not_reconnect_after_terminal_transport_error() -> None:
+    calls = 0
+
+    def handler(_request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        return httpx.Response(200, stream=DisconnectingSyncStream(_event_frame(_run_succeeded_event())))
+
+    client = Client(
+        base_url="http://testserver",
+        sync_http_client=httpx.Client(transport=httpx.MockTransport(handler)),
+    )
+
+    events = list(client.stream_events_sync("run-1", until_terminal=False, reconnect_delay_seconds=0))
+
+    assert [event.type for event in events] == ["run_succeeded"]
+    assert calls == 1
+
+
 def test_stream_events_reconnects_from_latest_event_id() -> None:
     seen_after: list[str] = []
 
@@ -778,6 +829,101 @@ def test_async_stream_events_yields_terminal_event() -> None:
     asyncio.run(scenario())
 
 
+def test_async_stream_events_does_not_reconnect_after_terminal_when_until_terminal_is_false() -> None:
+    calls = 0
+
+    def handler(_request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        return httpx.Response(200, content=_event_frame(_run_succeeded_event()))
+
+    async def scenario() -> None:
+        http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
+        client = Client(base_url="http://testserver", async_http_client=http_client)
+
+        events = [event async for event in client.stream_events("run-1", until_terminal=False)]
+
+        assert [event.type for event in events] == ["run_succeeded"]
+        assert calls == 1
+        await http_client.aclose()
+
+    asyncio.run(scenario())
+
+
+def test_async_stream_events_does_not_reconnect_after_terminal_transport_error() -> None:
+    calls = 0
+
+    def handler(_request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        return httpx.Response(200, stream=DisconnectingAsyncStream(_event_frame(_run_succeeded_event())))
+
+    async def scenario() -> None:
+        http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
+        client = Client(base_url="http://testserver", async_http_client=http_client)
+
+        events = [event async for event in client.stream_events("run-1", until_terminal=False)]
+
+        assert [event.type for event in events] == ["run_succeeded"]
+        assert calls == 1
+        await http_client.aclose()
+
+    asyncio.run(scenario())
+
+
+def test_async_stream_events_reconnects_from_latest_event_after_transport_error() -> None:
+    seen_after: list[str] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        seen_after.append(request.url.params["after"])
+        if len(seen_after) == 1:
+            return httpx.Response(
+                200,
+                stream=DisconnectingAsyncStream(_event_frame(RunStartedEvent(id="1-0", run_id="run-1"))),
+            )
+        return httpx.Response(200, content=_event_frame(_run_succeeded_event()))
+
+    async def scenario() -> None:
+        http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
+        client = Client(base_url="http://testserver", async_http_client=http_client)
+
+        events = [event async for event in client.stream_events("run-1", reconnect_delay_seconds=0)]
+
+        assert seen_after == ["0-0", "1-0"]
+        assert [event.type for event in events] == ["run_started", "run_succeeded"]
+        await http_client.aclose()
+
+    asyncio.run(scenario())
+
+
+def test_async_stream_events_reconnects_after_eof_before_terminal() -> None:
+    calls = 0
+
+    def handler(_request: httpx.Request) -> httpx.Response:
+        nonlocal calls
+        calls += 1
+        return httpx.Response(200, content="")
+
+    async def scenario() -> None:
+        http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
+        client = Client(base_url="http://testserver", async_http_client=http_client)
+
+        with pytest.raises(DifyAgentStreamError, match="reconnect attempts exhausted"):
+            _ = [
+                event
+                async for event in client.stream_events(
+                    "run-1",
+                    max_reconnects=1,
+                    reconnect_delay_seconds=0,
+                )
+            ]
+
+        assert calls == 2
+        await http_client.aclose()
+
+    asyncio.run(scenario())
+
+
 def test_async_sse_parser_preserves_unicode_line_separators() -> None:
     error = "next-line:\x85line-separator:\u2028paragraph-separator:\u2029done"
     body = _event_frame(_run_failed_event(error))

dify-agent/tests/local/dify_agent/server/test_sse.py

@@ -3,6 +3,8 @@ import json
 from collections.abc import AsyncGenerator
 from typing import cast
 
+import pytest
+
 from dify_agent.protocol.schemas import RunFailedEvent, RunFailedEventData, RunStartedEvent
 from dify_agent.server.sse import format_sse_event, sse_event_stream
 
@@ -49,3 +51,16 @@ def test_sse_event_stream_emits_heartbeats_while_waiting() -> None:
         await stream.aclose()
 
     asyncio.run(scenario())
+
+
+def test_sse_event_stream_ends_after_finite_terminal_event_iterator() -> None:
+    async def scenario() -> None:
+        async def events():
+            yield RunFailedEvent(id="2-0", run_id="run-1", data=RunFailedEventData(error="model failed"))
+
+        stream = cast(AsyncGenerator[str, None], sse_event_stream(events(), heartbeat_interval_seconds=0.001))
+        assert (await anext(stream)).startswith("id: 2-0\nevent: run_failed")
+        with pytest.raises(StopAsyncIteration):
+            _ = await asyncio.wait_for(anext(stream), timeout=0.1)
+
+    asyncio.run(scenario())

dify-agent/tests/local/dify_agent/storage/test_redis_run_store.py

@@ -175,6 +175,25 @@ class FakeRedisPipeline:
         return list(self.results)
 
 
+def _terminal_event(
+    event_type: str,
+    run_id: str,
+) -> RunSucceededEvent | RunFailedEvent | RunCancelledEvent:
+    if event_type == "run_succeeded":
+        return RunSucceededEvent(
+            run_id=run_id,
+            data=RunSucceededEventData(
+                output="done",
+                session_snapshot=CompositorSessionSnapshot(layers=[]),
+            ),
+        )
+    if event_type == "run_failed":
+        return RunFailedEvent(run_id=run_id, data=RunFailedEventData(error="model failed"))
+    if event_type == "run_cancelled":
+        return RunCancelledEvent(run_id=run_id, data=RunCancelledEventData(reason="cancelled"))
+    raise AssertionError(f"unexpected terminal event type: {event_type}")
+
+
 def test_create_run_writes_running_record_without_job_queue_and_with_retention() -> None:
     redis = FakeRedis()
     store = RedisRunStore(redis, prefix="test")  # pyright: ignore[reportArgumentType]
@@ -504,3 +523,47 @@ def test_get_events_round_trips_run_succeeded_output_and_session_snapshot() -> N
     assert decoded.id == event_id
     assert decoded.data.output == output
     assert decoded.data.session_snapshot == session_snapshot
+
+
+@pytest.mark.parametrize("terminal_type", ["run_succeeded", "run_failed", "run_cancelled"])
+def test_iter_events_ends_after_replaying_terminal_event(terminal_type: str) -> None:
+    redis = FakeRedis()
+    store = RedisRunStore(redis, prefix="test")  # pyright: ignore[reportArgumentType]
+
+    async def scenario() -> list[str]:
+        record = await store.create_run()
+        _ = await store.append_event(RunStartedEvent(run_id=record.run_id))
+        _ = await store.finalize_run(_terminal_event(terminal_type, record.run_id))
+        redis.commands.clear()
+
+        async def collect_events() -> list[str]:
+            return [event.type async for event in store.iter_events(record.run_id)]
+
+        return await asyncio.wait_for(collect_events(), timeout=1)
+
+    event_types = asyncio.run(scenario())
+
+    assert event_types == ["run_started", terminal_type]
+    assert "xread" not in [command[0] for command in redis.commands]
+
+
+@pytest.mark.parametrize("terminal_type", ["run_succeeded", "run_failed", "run_cancelled"])
+def test_iter_events_ends_after_live_terminal_event(terminal_type: str) -> None:
+    redis = FakeRedis()
+    store = RedisRunStore(redis, prefix="test")  # pyright: ignore[reportArgumentType]
+
+    async def scenario() -> str:
+        record = await store.create_run()
+        events = store.iter_events(record.run_id)
+        next_event = asyncio.ensure_future(anext(events))
+        await asyncio.sleep(0)
+        assert not next_event.done()
+        assert "xread" in [command[0] for command in redis.commands]
+
+        _ = await store.finalize_run(_terminal_event(terminal_type, record.run_id))
+        event = await asyncio.wait_for(next_event, timeout=1)
+        with pytest.raises(StopAsyncIteration):
+            _ = await anext(events)
+        return event.type
+
+    assert asyncio.run(scenario()) == terminal_type