fix: fix time travel on fastapi implementation

Author: ranst91

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

@@ -106,7 +106,6 @@ async def _handle_stream_events(self, input: RunAgentInput) -> AsyncGenerator[st
             "thinking_process": None,
         }
 
-        messages = input.messages or []
         forwarded_props = input.forwarded_props
         node_name_input = forwarded_props.get('node_name', None) if forwarded_props else None
 
@@ -120,30 +119,13 @@ async def _handle_stream_events(self, input: RunAgentInput) -> AsyncGenerator[st
 
         agent_state = await self.graph.aget_state(config)
         self.active_run["mode"] = "continue" if thread_id and self.active_run.get("node_name") != "__end__" and self.active_run.get("node_name") else "start"
+
         prepared_stream_response = await self.prepare_stream(input=input, agent_state=agent_state, config=config)
 
         yield self._dispatch_event(
             RunStartedEvent(type=EventType.RUN_STARTED, thread_id=thread_id, run_id=self.active_run["id"])
         )
 
-        langchain_messages = agui_messages_to_langchain(messages)
-        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):
-            # Find the last user message by working backwards from the last message
-            last_user_message = None
-            for i in range(len(langchain_messages) - 1, -1, -1):
-                if isinstance(langchain_messages[i], HumanMessage):
-                    last_user_message = langchain_messages[i]
-                    break
-
-            if last_user_message:
-                prepared_stream_response = await self.prepare_regenerate_stream(
-                    input=input,
-                    message_checkpoint=last_user_message,
-                    config=config
-                )
-
         state = prepared_stream_response["state"]
         stream = prepared_stream_response["stream"]
         config = prepared_stream_response["config"]
@@ -274,14 +256,31 @@ async def prepare_stream(self, input: RunAgentInput, agent_state: State, config:
         thread_id = input.thread_id
 
         state_input["messages"] = agent_state.values.get("messages", [])
-        self.active_run["current_graph_state"] = agent_state.values
+        self.active_run["current_graph_state"] = agent_state.values.copy()
         langchain_messages = agui_messages_to_langchain(messages)
         state = self.langgraph_default_merge_state(state_input, langchain_messages, tools)
         self.active_run["current_graph_state"].update(state)
         config["configurable"]["thread_id"] = thread_id
         interrupts = agent_state.tasks[0].interrupts if agent_state.tasks and len(agent_state.tasks) > 0 else []
         has_active_interrupts = len(interrupts) > 0
         resume_input = forwarded_props.get('command', {}).get('resume', None)
+        self.active_run["schema_keys"] = self.get_schema_keys(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):
+            # Find the last user message by working backwards from the last message
+            last_user_message = None
+            for i in range(len(langchain_messages) - 1, -1, -1):
+                if isinstance(langchain_messages[i], HumanMessage):
+                    last_user_message = langchain_messages[i]
+                    break
+
+            if last_user_message:
+                return await self.prepare_regenerate_stream(
+                    input=input,
+                    message_checkpoint=last_user_message,
+                    config=config
+                )
 
         events_to_dispatch = []
         if has_active_interrupts and not resume_input:
@@ -312,8 +311,6 @@ async def prepare_stream(self, input: RunAgentInput, agent_state: State, config:
         if self.active_run["mode"] == "continue":
             await self.graph.aupdate_state(config, state, as_node=self.active_run.get("node_name"))
 
-        self.active_run["schema_keys"] = self.get_schema_keys(config)
-
         if resume_input:
             stream_input = Command(resume=resume_input)
         else:

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

@@ -236,6 +236,7 @@ export class LangGraphAgent extends AbstractAgent {
       values: { ...stateValuesDiff, messages: [...agentStateMessages, ...stateValuesDiff.messages] },
     }
     let stateValues = threadState.values;
+    this.activeRun!.schemaKeys = await this.getSchemaKeys();
 
     if ((agentState.values.messages ?? []).length > messages.filter((m) => m.role !== "system").length) {
       let lastUserMessage: LangGraphMessage | null = null;
@@ -268,8 +269,6 @@ export class LangGraphAgent extends AbstractAgent {
       });
     }
 
-    this.activeRun!.schemaKeys = await this.getSchemaKeys();
-
     const payloadInput = getStreamPayloadInput({
       mode,
       state: stateValues,