From 6f5de50eb32d6be0a1e6ab098cba9e94907c3354 Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Wed, 10 Jun 2026 16:58:51 -0300 Subject: [PATCH] fix(api/v2): buffer parallel messages in the AG-UI translator instead of dropping them Parallel components stream tokens for different message ids interleaved. The translator tracked a single open message: the first foreign token closed the open message and tombstoned its id, so every later event for it was dropped and its remaining text never reached the client. Tokens for a message that cannot take the wire now buffer until the open message genuinely ends (its add_message finalizer), then flush in arrival order; complete messages landing mid-stream buffer the same way instead of interleaving a second START. end/error drain all buffers before the terminal event. The wire still carries at most one open text message, so the stream stays AG-UI-conformant. --- .../base/langflow/api/v2/agui_translator.py | 149 ++++++++++++++---- .../tests/unit/api/v2/test_agui_translator.py | 127 +++++++++++---- 2 files changed, 217 insertions(+), 59 deletions(-) diff --git a/src/backend/base/langflow/api/v2/agui_translator.py b/src/backend/base/langflow/api/v2/agui_translator.py index 727e71e56e..0fa76a39f2 100644 --- a/src/backend/base/langflow/api/v2/agui_translator.py +++ b/src/backend/base/langflow/api/v2/agui_translator.py @@ -12,6 +12,7 @@ list. The translator is stateful: one instance per run. from __future__ import annotations import json +from dataclasses import dataclass, field from ag_ui.core import ( BaseEvent, @@ -37,6 +38,28 @@ from ag_ui.core import ( _CUSTOM_CONTENT_TYPES = frozenset({"json", "code", "media", "error"}) +@dataclass +class _BufferedMessage: + """A text message waiting for the wire while another message is open. + + AG-UI allows one open text message at a time, but parallel components + stream tokens for different message ids interleaved. Tokens for a message + that cannot open yet accumulate here; ``final_text`` is set when its + ``add_message`` finalizer arrives before it ever reached the wire. + """ + + chunks: list[str] = field(default_factory=list) + final_text: str | None = None + + @property + def complete(self) -> bool: + return self.final_text is not None + + @property + def text(self) -> str: + return self.final_text if self.final_text is not None else "".join(self.chunks) + + class AGUITranslator: """Translates Langflow ``EventManager`` events into AG-UI protocol events. @@ -47,11 +70,16 @@ class AGUITranslator: def __init__(self, run_id: str, thread_id: str) -> None: self.run_id = run_id self.thread_id = thread_id - # Id of the text message currently being streamed by ``token`` events, - # or ``None`` when no message is open. + # Id of the text message currently open on the wire, or ``None``. self._open_message_id: str | None = None - # Message ids already emitted as a complete (non-streamed) text message. + # Message ids whose TEXT_MESSAGE_END has been emitted; the protocol + # considers them closed, so nothing may reopen them. self._emitted_text_message_ids: set[str] = set() + # Messages from parallel components waiting for the wire, in arrival + # order (dict preserves insertion order). Invariant: non-empty only + # while another message holds the wire — closing the open message + # immediately promotes the next buffered one. + self._buffered_messages: dict[str, _BufferedMessage] = {} # Tool-call ids already emitted as TOOL_CALL_START / already resolved # with a TOOL_CALL_RESULT. ``add_message`` can re-fire with the same # (append-only) content_blocks, so emissions must be deduplicated. @@ -95,11 +123,11 @@ class AGUITranslator: # streamed message and must stay transparent, or the message would be # split into multiple START/END pairs reusing an already-ended id. if event_type == "end": - events = self._close_open_message() + events = self._drain_messages() events.append(RunFinishedEvent(run_id=self.run_id, thread_id=self.thread_id)) return events if event_type == "error": - events = self._close_open_message() + events = self._drain_messages() # The ``error`` payload varies by emission path: a full ErrorMessage # dump carries the reason in ``text``; the minimal path sends # ``{"error": str}``. @@ -111,11 +139,13 @@ class AGUITranslator: def _translate_token(self, data: dict) -> list[BaseEvent]: """Map a ``token`` event to text-message events. - The first token of a message opens it with ``TEXT_MESSAGE_START``; a - token for a different message id closes the previous one first. A - token whose id was already ended via ``add_message`` (or a prior - token boundary) is dropped: re-opening it would emit a second - ``TEXT_MESSAGE_START`` for an id the protocol considers closed. + The first token of a message opens it with ``TEXT_MESSAGE_START``. + While a message holds the wire, tokens for other message ids (parallel + components stream interleaved) buffer instead of closing it; the + buffered message flushes when the open one genuinely ends. A token + whose id was already ended via ``add_message`` is dropped: re-opening + it would emit a second ``TEXT_MESSAGE_START`` for an id the protocol + considers closed. """ message_id = str(data.get("id") or "") if not message_id: @@ -123,16 +153,22 @@ class AGUITranslator: # be correlated. Dropping the event is preferable to emitting a # malformed stream with empty message_ids. return [] - if message_id in self._emitted_text_message_ids and self._open_message_id != message_id: + if message_id in self._emitted_text_message_ids: return [] chunk = data.get("chunk", "") - events: list[BaseEvent] = [] - if self._open_message_id != message_id: - events.extend(self._close_open_message()) - events.append(TextMessageStartEvent(message_id=message_id, role="assistant")) + if self._open_message_id == message_id: + return [TextMessageContentEvent(message_id=message_id, delta=chunk)] + if self._open_message_id is None: self._open_message_id = message_id - events.append(TextMessageContentEvent(message_id=message_id, delta=chunk)) - return events + return [ + TextMessageStartEvent(message_id=message_id, role="assistant"), + TextMessageContentEvent(message_id=message_id, delta=chunk), + ] + # Another message holds the wire: buffer this one until it is free. + buffered = self._buffered_messages.setdefault(message_id, _BufferedMessage()) + if not buffered.complete: + buffered.chunks.append(chunk) + return [] def _translate_vertices_sorted(self, data: dict) -> list[BaseEvent]: """Map ``vertices_sorted`` to a ``STATE_SNAPSHOT`` of the node graph. @@ -198,21 +234,32 @@ class AGUITranslator: # Message text. if message_id and message_id == self._open_message_id: - # Finalizer of a token-streamed message: close it. The text was - # already streamed token by token, so it must not be re-emitted now - # or by any later add_message that re-fires for the same id. - self._emitted_text_message_ids.add(message_id) + # Finalizer of the wire-open streamed message: close it. The text + # was already streamed token by token, so it must not be re-emitted + # now or by any later add_message that re-fires for the same id. events.extend(self._close_open_message()) + elif message_id and message_id in self._buffered_messages: + # Finalizer for a message still waiting for the wire: record the + # authoritative full text; the trio is emitted on promotion. + self._buffered_messages[message_id].final_text = ( + data.get("text") or self._buffered_messages[message_id].text + ) else: text = data.get("text") or "" # Skip text-message lifecycle emission without a stable message_id; # tool-call events above are namespaced by block/content index so # they can still ride a missing id, but TEXT_MESSAGE_* cannot. if text and message_id and message_id not in self._emitted_text_message_ids: - self._emitted_text_message_ids.add(message_id) - events.append(TextMessageStartEvent(message_id=message_id, role="assistant")) - events.append(TextMessageContentEvent(message_id=message_id, delta=text)) - events.append(TextMessageEndEvent(message_id=message_id)) + if self._open_message_id is None: + self._emitted_text_message_ids.add(message_id) + events.append(TextMessageStartEvent(message_id=message_id, role="assistant")) + events.append(TextMessageContentEvent(message_id=message_id, delta=text)) + events.append(TextMessageEndEvent(message_id=message_id)) + else: + # A parallel component finished while another message holds + # the wire: buffer the complete message instead of opening + # a second text message mid-stream. + self._buffered_messages[message_id] = _BufferedMessage(final_text=text) return events def _translate_tool_use( @@ -286,17 +333,57 @@ class AGUITranslator: return {"op": "add", "path": f"/nodes/{node_id}", "value": {"status": status, "output": output}} def _close_open_message(self) -> list[BaseEvent]: - """Emit ``TEXT_MESSAGE_END`` for the open message, if any. + """Emit ``TEXT_MESSAGE_END`` for the open message and free the wire. The closed id is recorded in ``_emitted_text_message_ids`` so a later - token (e.g. an interleaved ``A, B, A`` sequence) cannot re-open it - and emit a second ``TEXT_MESSAGE_START`` for an id the protocol - already considers closed. + token cannot re-open it and emit a second ``TEXT_MESSAGE_START`` for + an id the protocol already considers closed. With the wire free, the + next buffered parallel message (if any) is promoted onto it. """ if self._open_message_id is None: return [] closed_id = self._open_message_id - end = TextMessageEndEvent(message_id=closed_id) self._emitted_text_message_ids.add(closed_id) self._open_message_id = None - return [end] + events: list[BaseEvent] = [TextMessageEndEvent(message_id=closed_id)] + events.extend(self._promote_next_buffered()) + return events + + def _promote_next_buffered(self) -> list[BaseEvent]: + """Move buffered parallel messages onto the freed wire, in arrival order. + + Already-complete messages emit their full START/CONTENT/END trio and the + promotion continues; the first still-streaming message replays its + buffered chunks, takes the wire, and stays open for its live tokens. + """ + events: list[BaseEvent] = [] + while self._buffered_messages and self._open_message_id is None: + message_id, buffered = next(iter(self._buffered_messages.items())) + del self._buffered_messages[message_id] + if message_id in self._emitted_text_message_ids: + continue + text = buffered.text + if buffered.complete: + self._emitted_text_message_ids.add(message_id) + if text: + events.append(TextMessageStartEvent(message_id=message_id, role="assistant")) + events.append(TextMessageContentEvent(message_id=message_id, delta=text)) + events.append(TextMessageEndEvent(message_id=message_id)) + continue + self._open_message_id = message_id + events.append(TextMessageStartEvent(message_id=message_id, role="assistant")) + if text: + events.append(TextMessageContentEvent(message_id=message_id, delta=text)) + return events + + def _drain_messages(self) -> list[BaseEvent]: + """Close the open message and flush every buffered one (run boundary). + + At ``end``/``error`` nothing else will free the wire, so buffered + parallel messages flush now — each promoted, emitted, and closed — + rather than being silently lost. + """ + events = self._close_open_message() + while self._open_message_id is not None: + events.extend(self._close_open_message()) + return events diff --git a/src/backend/tests/unit/api/v2/test_agui_translator.py b/src/backend/tests/unit/api/v2/test_agui_translator.py index 0797105bc7..995cea3d86 100644 --- a/src/backend/tests/unit/api/v2/test_agui_translator.py +++ b/src/backend/tests/unit/api/v2/test_agui_translator.py @@ -90,13 +90,31 @@ def test_token_sequence_emits_start_contents_then_end_on_boundary(): assert isinstance(ended[1], RunFinishedEvent) -def test_new_message_id_closes_previous_message_and_opens_new(): +def test_token_for_second_message_is_buffered_until_first_closes(): + """A token for a different message id must not close the streaming one. + + Parallel components stream tokens for different message ids interleaved. + Closing the open message on the first foreign token burned its id, so all + its later tokens were dropped. Instead the foreign message buffers until + the open one genuinely ends (its ``add_message`` finalizer), then flushes. + """ t = AGUITranslator(run_id="r1", thread_id="t1") t.start() t.translate("token", {"chunk": "a", "id": "m1"}) - out = t.translate("token", {"chunk": "b", "id": "m2"}) + buffered = t.translate("token", {"chunk": "b", "id": "m2"}) + # m2 buffers silently; m1 stays open. + assert buffered == [] + + # m1 keeps streaming: its tokens still flow. + more = t.translate("token", {"chunk": "a2", "id": "m1"}) + assert len(more) == 1 + assert isinstance(more[0], TextMessageContentEvent) + assert more[0].message_id == "m1" + + # m1's finalizer closes it and promotes m2 with its buffered content. + out = t.translate("add_message", {"id": "m1", "text": "aa2"}) assert isinstance(out[0], TextMessageEndEvent) assert out[0].message_id == "m1" assert isinstance(out[1], TextMessageStartEvent) @@ -166,32 +184,83 @@ def test_token_for_already_ended_message_id_is_dropped(): ) -def test_token_after_boundary_close_is_dropped(): - """An interleaved token sequence ``A, B, A`` must not re-open id A. +def test_parallel_interleaved_tokens_preserve_all_content(): + """Two components streaming in parallel must not lose either stream. - Switching from id A to id B closes A through ``_close_open_message``. A - later token for A used to slip past the dedup guard because - ``_close_open_message`` was the only finalizer that did not record the - closed id in ``_emitted_text_message_ids``. The translator must treat a - token-boundary close the same way it treats an ``add_message`` close. + Reproduces the reported parallel-components drop: with a single-slot + tracker, START(m2), CONTENT(m2), START(m3) burned m2, so CONTENT(m2) + was dropped and m2's remaining text never reached the client. With + buffering, m2 holds the wire until it genuinely ends; m3's tokens + buffer and flush afterwards. Nothing is dropped, and the wire carries + at most one open text message at a time. + """ + t = AGUITranslator(run_id="r1", thread_id="t1") + sequence = [ + ("token", {"chunk": "Result ", "id": "m2"}), + ("token", {"chunk": "Side ", "id": "m3"}), + ("token", {"chunk": "for X", "id": "m2"}), + ("token", {"chunk": "for Y", "id": "m3"}), + ("add_message", {"id": "m2", "text": "Result for X"}), + ("add_message", {"id": "m3", "text": "Side for Y"}), + ("end", {}), + ] + + out = _run_sequence(t, sequence) + + _assert_well_formed(out) + text_by_message: dict[str, str] = {} + for event in out: + if isinstance(event, TextMessageContentEvent): + text_by_message[event.message_id] = text_by_message.get(event.message_id, "") + event.delta + assert text_by_message["m2"] == "Result for X" + assert text_by_message["m3"] == "Side for Y" + + +def test_add_message_for_other_message_while_streaming_is_buffered(): + """A complete message landing mid-stream must not interleave its trio. + + A parallel non-streaming component can finish (add_message with full + text) while another component holds the wire with an open streamed + message. Emitting START/CONTENT/END for the finished one immediately + would open a second text message mid-stream; instead it buffers and + flushes when the open message closes. """ t = AGUITranslator(run_id="r1", thread_id="t1") t.start() - a1 = t.translate("token", {"chunk": "hi", "id": "m1"}) - assert any(isinstance(e, TextMessageStartEvent) and e.message_id == "m1" for e in a1) + t.translate("token", {"chunk": "streaming...", "id": "m1"}) + parallel_done = t.translate("add_message", {"id": "m9", "text": "finished early"}) + assert all( + not isinstance(e, (TextMessageStartEvent, TextMessageContentEvent, TextMessageEndEvent)) for e in parallel_done + ), f"buffered message leaked text events mid-stream: {parallel_done}" - # Switching to a new id closes m1 via _close_open_message. - b = t.translate("token", {"chunk": "yo", "id": "m2"}) - assert any(isinstance(e, TextMessageEndEvent) and e.message_id == "m1" for e in b) - assert any(isinstance(e, TextMessageStartEvent) and e.message_id == "m2" for e in b) + out = t.translate("add_message", {"id": "m1", "text": "streaming..."}) + assert [type(e) for e in out] == [ + TextMessageEndEvent, + TextMessageStartEvent, + TextMessageContentEvent, + TextMessageEndEvent, + ] + assert out[0].message_id == "m1" + assert out[1].message_id == "m9" + assert out[2].delta == "finished early" - # A late token for m1 must be dropped, not re-open the ended message. - late = t.translate("token", {"chunk": "again", "id": "m1"}) - assert all(not isinstance(e, TextMessageStartEvent) for e in late), ( - f"Token boundary did not mark m1 as ended; emitted: {late}" - ) - assert late == [], f"Expected no events for late token after boundary close; got {late}" + +def test_end_drains_buffered_messages_before_run_finished(): + """A run ending while messages are still buffered must flush them all.""" + t = AGUITranslator(run_id="r1", thread_id="t1") + t.start() + + t.translate("token", {"chunk": "open", "id": "m1"}) + t.translate("token", {"chunk": "waiting", "id": "m2"}) + out = t.translate("end", {}) + + assert isinstance(out[-1], RunFinishedEvent) + m2_content = [e for e in out if isinstance(e, TextMessageContentEvent) and e.message_id == "m2"] + assert len(m2_content) == 1 + assert m2_content[0].delta == "waiting" + ends = [e.message_id for e in out if isinstance(e, TextMessageEndEvent)] + assert ends == ["m1", "m2"] def test_vertices_sorted_emits_state_snapshot_of_all_nodes(): @@ -610,20 +679,22 @@ def _assert_well_formed(events: list) -> None: assert isinstance(events[0], RunStartedEvent) assert isinstance(events[-1], (RunFinishedEvent, RunErrorEvent)) - open_messages: set[str] = set() + open_message: str | None = None seen_messages: set[str] = set() for event in events: if isinstance(event, TextMessageStartEvent): - assert event.message_id not in open_messages, "text message started while already open" + # AG-UI's reference verifier allows at most one open text message; + # interleaved STARTs are rejected by conforming clients. + assert open_message is None, f"text message {event.message_id} started while {open_message} is open" assert event.message_id not in seen_messages, "text message id reused after it ended" - open_messages.add(event.message_id) + open_message = event.message_id seen_messages.add(event.message_id) elif isinstance(event, TextMessageContentEvent): - assert event.message_id in open_messages, "text content for a message that is not open" + assert event.message_id == open_message, "text content for a message that is not open" elif isinstance(event, TextMessageEndEvent): - assert event.message_id in open_messages, "text message ended without being open" - open_messages.discard(event.message_id) - assert not open_messages, "text messages left unclosed" + assert event.message_id == open_message, "text message ended without being open" + open_message = None + assert open_message is None, "text message left unclosed" started_tools: set[str] = set() ended_tools: set[str] = set()