fix: preserve AIMessage id in replace_tool_calls (#688)

Co-authored-by: Valentina Bojan <valentina.bojan@uipath.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: 3 people

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.8.19"
+version = "0.8.20"
 description = "Python SDK that enables developers to build and deploy LangGraph agents to the UiPath Cloud Platform"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

src/uipath_langchain/agent/messages/message_utils.py

@@ -37,4 +37,5 @@ def replace_tool_calls(message: AIMessage, tool_calls: list[ToolCall]) -> AIMess
         content_blocks=content_blocks,
         tool_calls=tool_calls,
         response_metadata=response_metadata,
+        id=message.id,
     )

tests/agent/messages/test_message_utils.py

@@ -1,14 +1,20 @@
 """Tests for agent/messages/message_utils.py module."""
 
-from langchain.messages import AIMessage, ToolCall
+from typing import Any, Union
+
+from langchain.messages import AIMessage, HumanMessage, ToolCall
+from langchain_core.messages import BaseMessage
 from langchain_core.messages.content import (
     ContentBlock,
     create_text_block,
     create_tool_call,
 )
+from langgraph.graph.message import add_messages
 
 from uipath_langchain.agent.messages.message_utils import replace_tool_calls
 
+MessageItem = Union[BaseMessage, list[str], tuple[str, str], str, dict[str, Any]]
+
 
 class TestReplaceToolCalls:
     """Test cases for replace_tool_calls function."""
@@ -179,6 +185,74 @@ def test_replace_tool_calls_no_original_metadata(self):
         assert len(tool_call_blocks) == 1
         assert tool_call_blocks[0]["name"] == "new_tool"
 
+    def test_replace_tool_calls_preserves_message_id(self):
+        """Test that the original message id is preserved after replacement."""
+        original_tool_calls = [ToolCall(name="old_tool", args={}, id="old_id")]
+        original_content_blocks: list[ContentBlock] = [
+            create_text_block("Test message"),
+            create_tool_call(name="old_tool", args={}, id="old_id"),
+        ]
+        original_message = AIMessage(
+            content_blocks=original_content_blocks,
+            tool_calls=original_tool_calls,
+            id="msg-original-id",
+        )
+
+        new_tool_calls = [ToolCall(name="new_tool", args={}, id="new_id")]
+
+        result = replace_tool_calls(original_message, new_tool_calls)
+
+        assert result.id == "msg-original-id"
+
+    def test_replace_tool_calls_updated_args_visible_via_add_messages(self):
+        """Test that updated tool call args are visible after add_messages processes them.
+
+        Reproduces the HITL bug: when a human reviews and updates activity input
+        during an escalation, the activity must execute with the reviewed args.
+        Without id preservation, add_messages appends a duplicate AIMessage
+        instead of replacing the original, causing the tool to run with stale args.
+        """
+        original_tool_calls = [
+            ToolCall(name="my_activity", args={"input": "original_value"}, id="call_1")
+        ]
+        original_ai_message = AIMessage(
+            content_blocks=[
+                create_text_block("I will invoke the activity"),
+                create_tool_call(
+                    name="my_activity", args={"input": "original_value"}, id="call_1"
+                ),
+            ],
+            tool_calls=original_tool_calls,
+            id="msg-from-llm",
+        )
+
+        messages: list[MessageItem] = [
+            HumanMessage(content="do something", id="msg-human"),
+            original_ai_message,
+        ]
+
+        # Simulate HITL review: human changes the input
+        reviewed_tool_calls = [
+            ToolCall(name="my_activity", args={"input": "reviewed_value"}, id="call_1")
+        ]
+        updated_ai_message = replace_tool_calls(
+            original_ai_message, reviewed_tool_calls
+        )
+
+        # Simulate what Command(update={"messages": [updated_ai_message]}) does
+        new_messages: list[MessageItem] = [updated_ai_message]
+        result_messages = add_messages(messages, new_messages)
+
+        # There must be exactly one AIMessage — not a duplicate
+        ai_messages = [m for m in result_messages if isinstance(m, AIMessage)]
+        assert len(ai_messages) == 1, (
+            f"Expected 1 AIMessage but got {len(ai_messages)}; "
+            "add_messages appended instead of replacing (id mismatch)"
+        )
+
+        # The surviving AIMessage must carry the reviewed args
+        assert ai_messages[0].tool_calls[0]["args"] == {"input": "reviewed_value"}
+
     def test_replace_tool_calls_content_blocks(self):
         """Test that non-tool content blocks are preserved."""
         original_tool_calls = [ToolCall(name="old_tool", args={}, id="old_id")]