fix(langgraph): handle Command tool output (#639)

- fixes #431

Co-authored-by: Ran Shemtov <[email protected]>

Author: xdave

integrations/langgraph/python/ag_ui_langgraph/agent.py

@@ -6,10 +6,10 @@
 from langgraph.graph.state import CompiledStateGraph
 
 try:
-    from langchain.schema import BaseMessage, SystemMessage
+    from langchain.schema import BaseMessage, SystemMessage, ToolMessage
 except ImportError:
     # Langchain >= 1.0.0
-    from langchain_core.messages import BaseMessage, SystemMessage
+    from langchain_core.messages import BaseMessage, SystemMessage, ToolMessage
 
 from langchain_core.runnables import RunnableConfig, ensure_config
 from langchain_core.messages import HumanMessage
@@ -682,6 +682,51 @@ async def _handle_single_event(self, event: Any, state: State) -> AsyncGenerator
 
         elif event_type == LangGraphEventTypes.OnToolEnd:
             tool_call_output = event["data"]["output"]
+
+            if isinstance(tool_call_output, Command):
+                # Extract ToolMessages from Command.update
+                messages = tool_call_output.update.get('messages', [])
+                tool_messages = [m for m in messages if isinstance(m, ToolMessage)]
+
+                # Process each tool message
+                for tool_msg in tool_messages:
+                    if not self.active_run["has_function_streaming"]:
+                        yield self._dispatch_event(
+                            ToolCallStartEvent(
+                                type=EventType.TOOL_CALL_START,
+                                tool_call_id=tool_msg.tool_call_id,
+                                tool_call_name=tool_msg.name,
+                                parent_message_id=tool_msg.id,
+                                raw_event=event,
+                            )
+                        )
+                        yield self._dispatch_event(
+                            ToolCallArgsEvent(
+                                type=EventType.TOOL_CALL_ARGS,
+                                tool_call_id=tool_msg.tool_call_id,
+                                delta=json.dumps(event["data"].get("input", {})),
+                                raw_event=event
+                            )
+                        )
+                        yield self._dispatch_event(
+                            ToolCallEndEvent(
+                                type=EventType.TOOL_CALL_END,
+                                tool_call_id=tool_msg.tool_call_id,
+                                raw_event=event
+                            )
+                        )
+
+                    yield self._dispatch_event(
+                        ToolCallResultEvent(
+                            type=EventType.TOOL_CALL_RESULT,
+                            tool_call_id=tool_msg.tool_call_id,
+                            message_id=str(uuid.uuid4()),
+                            content=tool_msg.content,
+                            role="tool"
+                        )
+                    )
+                return
+
             if not self.active_run["has_function_streaming"]:
                 yield self._dispatch_event(
                     ToolCallStartEvent(