fix(langgraph): include the payload with the original user message that is going to be regenerated

Author: ranst91

typescript-sdk/integrations/langgraph/python/ag_ui_langgraph/agent.py

@@ -685,7 +685,12 @@ async def get_checkpoint_before_message(self, message_id: str, thread_id: str):
                     empty_snapshot = snapshot
                     empty_snapshot.values["messages"] = []
                     return empty_snapshot
-                return history_list[idx - 1]  # return one snapshot *before* the one that includes the message
+
+                snapshot_values_without_messages = snapshot.values.copy()
+                del snapshot_values_without_messages["messages"]
+                checkpoint = history_list[idx - 1] # use one snapshot *before* the one that includes the message
+                checkpoint.values = {**checkpoint.values, **snapshot_values_without_messages}
+                return checkpoint
 
         raise ValueError("Message ID not found in history")
 

typescript-sdk/integrations/langgraph/src/agent.ts

@@ -992,7 +992,9 @@ export class LangGraphAgent extends AbstractAgent {
 
     const targetStateIndex = reversed.indexOf(targetState);
 
-    return reversed[targetStateIndex - 1] ?? { ...targetState, values: {} };
+    const { messages, ...targetStateValuesWithoutMessages } = targetState.values as State
+    const selectedCheckpoint = reversed[targetStateIndex - 1] ?? { ...targetState, values: {} }
+    return { ...selectedCheckpoint, values: { ...selectedCheckpoint.values, ...targetStateValuesWithoutMessages } };
   }
 }