From 4d24d568590db3d44d74fe39d3fbfc00c8384ddf Mon Sep 17 00:00:00 2001 From: Prajval Rasik Date: Sat, 4 Jul 2026 02:26:48 +0530 Subject: [PATCH] fix(langgraph): regenerate first assistant response --- .../langgraph/python/ag_ui_langgraph/agent.py | 63 ++++++++++++++++ .../test_prepare_stream_edited_message.py | 75 ++++++++++++++++++- 2 files changed, 137 insertions(+), 1 deletion(-) diff --git a/integrations/langgraph/python/ag_ui_langgraph/agent.py b/integrations/langgraph/python/ag_ui_langgraph/agent.py index ddc25a7c1d..0319c93204 100644 --- a/integrations/langgraph/python/ag_ui_langgraph/agent.py +++ b/integrations/langgraph/python/ag_ui_langgraph/agent.py @@ -560,6 +560,17 @@ async def prepare_stream(self, input: RunAgentInput, agent_state: State, config: config=config, ) + regenerate_placeholder = self._detect_regenerate_placeholder( + langchain_messages, + agent_state.values.get("messages", []), + ) + if regenerate_placeholder: + return await self.prepare_regenerate_stream( + input=input, + message_checkpoint=regenerate_placeholder, + config=config, + ) + non_system_messages = [msg for msg in langchain_messages if not isinstance(msg, SystemMessage)] if len(agent_state.values.get("messages", [])) > len(non_system_messages): # Only trigger time-travel regeneration if the incoming messages are NOT already @@ -1083,6 +1094,58 @@ def _detect_edited_human_message( return msg return None + @staticmethod + def _detect_regenerate_placeholder( + incoming_messages: List[BaseMessage], + checkpoint_messages: List[BaseMessage], + ) -> Optional[HumanMessage]: + """Return the checkpointed user message preceding an empty assistant placeholder. + + Regenerating the first assistant response can produce equal message + counts: checkpoint ``[Human, AI]`` and incoming ``[Human, empty AI]``. + The count-based heuristic below never sees that as time travel, so use + the explicit empty assistant placeholder as the regenerate signal. + """ + non_system_incoming = [ + msg for msg in incoming_messages if not isinstance(msg, SystemMessage) + ] + if len(non_system_incoming) < 2: + return None + + placeholder = non_system_incoming[-1] + if not isinstance(placeholder, AIMessage): + return None + if placeholder.content or getattr(placeholder, "tool_calls", None): + return None + + last_user_message: Optional[HumanMessage] = None + for candidate in reversed(non_system_incoming[:-1]): + if isinstance(candidate, HumanMessage): + last_user_message = candidate + break + if not last_user_message or not getattr(last_user_message, "id", None): + return None + + checkpoint_messages_by_id = { + getattr(message, "id", None): message + for message in checkpoint_messages + if getattr(message, "id", None) + } + checkpoint_user = checkpoint_messages_by_id.get(last_user_message.id) + if not isinstance(checkpoint_user, HumanMessage): + return None + + try: + user_index = checkpoint_messages.index(checkpoint_user) + except ValueError: + return None + + has_downstream_ai = any( + isinstance(message, AIMessage) + for message in checkpoint_messages[user_index + 1:] + ) + return last_user_message if has_downstream_ai else None + def _interrupts_to_agui(self, lg_interrupts) -> List[AGUIInterrupt]: """Map LangGraph task interrupts to AG-UI Interrupts. diff --git a/integrations/langgraph/python/tests/test_prepare_stream_edited_message.py b/integrations/langgraph/python/tests/test_prepare_stream_edited_message.py index 355a7f5ae7..a5c034a0c2 100644 --- a/integrations/langgraph/python/tests/test_prepare_stream_edited_message.py +++ b/integrations/langgraph/python/tests/test_prepare_stream_edited_message.py @@ -18,7 +18,7 @@ from langchain_core.messages import AIMessage, HumanMessage -from ag_ui.core import EventType, UserMessage +from ag_ui.core import AssistantMessage, EventType, UserMessage from tests._helpers import make_agent @@ -48,6 +48,7 @@ def _make_input(messages, thread_id="t1", forwarded_props=None): inp.tools = [] inp.context = [] inp.forwarded_props = forwarded_props or {} + inp.resume = None return inp @@ -113,6 +114,49 @@ def test_ignores_id_only_in_checkpoint(self): self.assertIsNone(agent._detect_edited_human_message(incoming, checkpoint)) +class TestDetectRegeneratePlaceholder(unittest.TestCase): + """Direct tests for empty-assistant regenerate detection.""" + + def test_detects_empty_assistant_placeholder_after_checkpointed_user(self): + agent = make_agent() + checkpoint = [ + HumanMessage(id="h1", content="hello"), + AIMessage(id="a1", content="hi"), + ] + incoming = [ + HumanMessage(id="h1", content="hello"), + AIMessage(id="", content=""), + ] + + result = agent._detect_regenerate_placeholder(incoming, checkpoint) + + self.assertIsNotNone(result) + self.assertEqual(result.id, "h1") + + def test_ignores_placeholder_without_downstream_assistant_response(self): + agent = make_agent() + checkpoint = [HumanMessage(id="h1", content="hello")] + incoming = [ + HumanMessage(id="h1", content="hello"), + AIMessage(id="", content=""), + ] + + self.assertIsNone(agent._detect_regenerate_placeholder(incoming, checkpoint)) + + def test_ignores_non_empty_assistant_message(self): + agent = make_agent() + checkpoint = [ + HumanMessage(id="h1", content="hello"), + AIMessage(id="a1", content="hi"), + ] + incoming = [ + HumanMessage(id="h1", content="hello"), + AIMessage(id="a1", content="hi"), + ] + + self.assertIsNone(agent._detect_regenerate_placeholder(incoming, checkpoint)) + + class TestPrepareStreamRoutesEditedMessage(unittest.IsolatedAsyncioTestCase): """Integration-level tests: ``prepare_stream`` must route a detected edit to ``prepare_regenerate_stream`` and skip the normal flow.""" @@ -143,6 +187,35 @@ async def test_edited_message_routes_to_regenerate(self): self.assertEqual(call_kwargs["message_checkpoint"].content, "What is 3+3?") self.assertEqual(result, {"stream": "regen"}) + async def test_empty_assistant_placeholder_routes_first_response_to_regenerate(self): + """Regenerating the first assistant response sends the same number of + messages as the checkpoint, so the count heuristic must not be the only + regenerate signal.""" + agent = make_agent() + agent.active_run = {"id": "run-1", "mode": "start"} + + checkpoint = [ + HumanMessage(id="h1", content="Hello"), + AIMessage(id="a1", content="Hi there"), + ] + state = _make_state(messages=checkpoint, tasks=[FakeTask()]) + + incoming = [ + UserMessage(id="h1", role="user", content="Hello"), + AssistantMessage(id="", role="assistant", content=""), + ] + inp = _make_input(messages=incoming) + + agent.prepare_regenerate_stream = AsyncMock(return_value={"stream": "regen"}) + config = {"configurable": {"thread_id": "t1"}} + + result = await agent.prepare_stream(inp, state, config) + + agent.prepare_regenerate_stream.assert_awaited_once() + call_kwargs = agent.prepare_regenerate_stream.await_args.kwargs + self.assertEqual(call_kwargs["message_checkpoint"].id, "h1") + self.assertEqual(result, {"stream": "regen"}) + async def test_unchanged_messages_do_not_regenerate(self): """Continuation (same id, same content) must NOT trigger regeneration — that path is reserved for true edits and rewinds."""