fix(langgraph): regenerate first assistant response

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.
 

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."""