fix: multiple control flow tools same llm call (#673)

Author: cotovanu-cristian

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.8.10"
+version = "0.8.11"
 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/react/llm_node.py

@@ -10,6 +10,7 @@
 )
 from langchain_core.tools import BaseTool, StructuredTool
 from pydantic import BaseModel
+from uipath.agent.react import RAISE_ERROR_TOOL
 from uipath.runtime.errors import UiPathErrorCategory
 
 from uipath_langchain.agent.tools.static_args import (
@@ -30,11 +31,24 @@
 def _filter_control_flow_tool_calls(
     tool_calls: list[ToolCall],
 ) -> list[ToolCall]:
-    """Remove control flow tools when multiple tool calls exist."""
+    """Remove control flow tool calls only when regular tool calls exist alongside them.
+
+    When only control flow tool calls are present and raise_error is among them,
+    keep only the first raise_error (takes precedence over end_execution).
+    """
     if len(tool_calls) <= 1:
         return tool_calls
 
-    return [tc for tc in tool_calls if tc.get("name") not in FLOW_CONTROL_TOOLS]
+    non_control_flow_tool_calls = [
+        tc for tc in tool_calls if tc.get("name") not in FLOW_CONTROL_TOOLS
+    ]
+    if not non_control_flow_tool_calls:
+        raise_error_calls = [
+            tc for tc in tool_calls if tc.get("name") == RAISE_ERROR_TOOL.name
+        ]
+        return raise_error_calls[:1] if raise_error_calls else tool_calls
+
+    return non_control_flow_tool_calls
 
 
 StateT = TypeVar("StateT", bound=AgentGraphState)

tests/agent/react/test_llm_node.py

@@ -263,3 +263,82 @@ async def test_multiple_flow_control_calls_all_filtered(self):
         ]
         assert len(tool_call_blocks) == 1
         assert tool_call_blocks[0]["name"] == "regular_tool"
+
+    @pytest.mark.asyncio
+    async def test_end_execution_and_raise_error_keeps_only_raise_error(self):
+        """When only end_execution and raise_error are called, keep only raise_error."""
+        mock_response = AIMessage(
+            content_blocks=[
+                create_tool_call(
+                    name=END_EXECUTION_TOOL.name,
+                    args={"result": "done"},
+                    id="call_1",
+                ),
+                create_tool_call(
+                    name=RAISE_ERROR_TOOL.name,
+                    args={"message": "conflict"},
+                    id="call_2",
+                ),
+            ],
+            tool_calls=[
+                {
+                    "name": END_EXECUTION_TOOL.name,
+                    "args": {"result": "done"},
+                    "id": "call_1",
+                },
+                {
+                    "name": RAISE_ERROR_TOOL.name,
+                    "args": {"message": "conflict"},
+                    "id": "call_2",
+                },
+            ],
+        )
+        self.mock_model.ainvoke = AsyncMock(return_value=mock_response)
+
+        llm_node = create_llm_node(self.mock_model, [self.regular_tool])
+
+        result = await llm_node(self.test_state)
+
+        response_message = result["messages"][0]
+        assert len(response_message.tool_calls) == 1
+        assert response_message.tool_calls[0]["name"] == RAISE_ERROR_TOOL.name
+
+    @pytest.mark.asyncio
+    async def test_multiple_raise_error_calls_keeps_only_first(self):
+        """When raise_error is called multiple times, keep only the first."""
+        mock_response = AIMessage(
+            content_blocks=[
+                create_tool_call(
+                    name=RAISE_ERROR_TOOL.name,
+                    args={"message": "first error"},
+                    id="call_1",
+                ),
+                create_tool_call(
+                    name=RAISE_ERROR_TOOL.name,
+                    args={"message": "second error"},
+                    id="call_2",
+                ),
+            ],
+            tool_calls=[
+                {
+                    "name": RAISE_ERROR_TOOL.name,
+                    "args": {"message": "first error"},
+                    "id": "call_1",
+                },
+                {
+                    "name": RAISE_ERROR_TOOL.name,
+                    "args": {"message": "second error"},
+                    "id": "call_2",
+                },
+            ],
+        )
+        self.mock_model.ainvoke = AsyncMock(return_value=mock_response)
+
+        llm_node = create_llm_node(self.mock_model, [self.regular_tool])
+
+        result = await llm_node(self.test_state)
+
+        response_message = result["messages"][0]
+        assert len(response_message.tool_calls) == 1
+        assert response_message.tool_calls[0]["name"] == RAISE_ERROR_TOOL.name
+        assert response_message.tool_calls[0]["args"]["message"] == "first error"