From 05c73860131fb734e3577bbf4256a5f288fdfd79 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=9B=90=E7=B2=92=20Yanli?= Date: Thu, 6 Aug 2026 18:53:10 +0800 Subject: [PATCH] fix(dify-agent): end SSE streams on terminal events (#40042) --- dify-agent/docs/dify-agent/guide/index.md | 5 +- dify-agent/src/dify_agent/client/_client.py | 24 ++- .../src/dify_agent/storage/redis_run_store.py | 8 +- .../local/dify_agent/client/test_client.py | 148 +++++++++++++++++- .../tests/local/dify_agent/server/test_sse.py | 15 ++ .../storage/test_redis_run_store.py | 63 ++++++++ 6 files changed, 257 insertions(+), 6 deletions(-) diff --git a/dify-agent/docs/dify-agent/guide/index.md b/dify-agent/docs/dify-agent/guide/index.md index 14b86650c51..2774d4438f6 100644 --- a/dify-agent/docs/dify-agent/guide/index.md +++ b/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 diff --git a/dify-agent/src/dify_agent/client/_client.py b/dify-agent/src/dify_agent/client/_client.py index a584555ed34..d72eb975c82 100644 --- a/dify-agent/src/dify_agent/client/_client.py +++ b/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( diff --git a/dify-agent/src/dify_agent/storage/redis_run_store.py b/dify-agent/src/dify_agent/storage/redis_run_store.py index 9afe26a40ea..140ea5a7f3d 100644 --- a/dify-agent/src/dify_agent/storage/redis_run_store.py +++ b/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: diff --git a/dify-agent/tests/local/dify_agent/client/test_client.py b/dify-agent/tests/local/dify_agent/client/test_client.py index 527dc2828b0..d87c192d783 100644 --- a/dify-agent/tests/local/dify_agent/client/test_client.py +++ b/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)) diff --git a/dify-agent/tests/local/dify_agent/server/test_sse.py b/dify-agent/tests/local/dify_agent/server/test_sse.py index 8f146ad396c..274725ce3cb 100644 --- a/dify-agent/tests/local/dify_agent/server/test_sse.py +++ b/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()) diff --git a/dify-agent/tests/local/dify_agent/storage/test_redis_run_store.py b/dify-agent/tests/local/dify_agent/storage/test_redis_run_store.py index eb56e6ffc87..ce839f80f33 100644 --- a/dify-agent/tests/local/dify_agent/storage/test_redis_run_store.py +++ b/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