From 190f83ab7abac10ef7871d01be8fc087aa69df25 Mon Sep 17 00:00:00 2001 From: radu-mocanu Date: Wed, 29 Jul 2026 13:56:10 +0300 Subject: [PATCH] feat: enforce configurable tool timeouts --- pyproject.toml | 2 +- .../_utils/durable_interrupt/decorator.py | 79 +++++++++++++------ .../agent/tools/extraction_tool.py | 12 ++- .../internal_tools/batch_transform_tool.py | 5 +- .../tools/internal_tools/deeprag_tool.py | 5 +- .../agent/tools/process_tool.py | 2 +- src/uipath_langchain/agent/tools/tool_node.py | 10 ++- .../test_batch_transform_tool.py | 9 ++- .../tools/internal_tools/test_deeprag_tool.py | 8 +- tests/agent/tools/test_durable_interrupt.py | 73 +++++++++++++++++ tests/agent/tools/test_extraction_tool.py | 16 ++-- tests/agent/tools/test_process_tool.py | 10 ++- tests/agent/tools/test_tool_node.py | 34 ++++++++ uv.lock | 2 +- 14 files changed, 220 insertions(+), 47 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ac43b3b36..dc93a7151 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath-langchain" -version = "0.14.16" +version = "0.14.17" 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" diff --git a/src/uipath_langchain/_utils/durable_interrupt/decorator.py b/src/uipath_langchain/_utils/durable_interrupt/decorator.py index 1d304f11f..3a420b133 100644 --- a/src/uipath_langchain/_utils/durable_interrupt/decorator.py +++ b/src/uipath_langchain/_utils/durable_interrupt/decorator.py @@ -33,11 +33,13 @@ async def start_job(): import asyncio import contextvars import functools -from typing import Any, Callable, TypeVar +from datetime import UTC, datetime, timedelta +from typing import Any, Callable, TypeVar, overload from langgraph._internal._constants import CONFIG_KEY_SCRATCHPAD from langgraph.config import get_config from langgraph.types import interrupt +from uipath.platform.common import WaitUntil, assert_no_timeout from .skip_interrupt import SkipInterruptValue @@ -94,7 +96,34 @@ def _inject_resume(scratchpad: Any, value: Any) -> Any: return value -def durable_interrupt(fn: F) -> F: +def _interrupt_with_timeout(value: Any, timeout: int | None) -> Any: + """Interrupt with an optional timeout expressed in milliseconds.""" + if timeout is None or timeout <= 0: + return interrupt(value) + + resume_time = datetime.now(UTC) + timedelta(milliseconds=timeout) + return assert_no_timeout(interrupt([value, WaitUntil(resume_time=resume_time)])) + + +def _resume_interrupt(timeout: int | None) -> Any: + """Resume an interrupt and raise when its configured timer completed first.""" + result = interrupt(None) + if timeout is None or timeout <= 0: + return result + return assert_no_timeout(result) + + +@overload +def durable_interrupt(fn: F, *, timeout: int | None = None) -> F: ... + + +@overload +def durable_interrupt(*, timeout: int | None = None) -> Callable[[F], F]: ... + + +def durable_interrupt( + fn: F | None = None, *, timeout: int | None = None +) -> F | Callable[[F], F]: """Decorator that executes a side-effecting function exactly once and interrupts. On first execution the body runs and its return value is passed to @@ -110,9 +139,10 @@ def durable_interrupt(fn: F) -> F: decorator that enforces the pairing contract. Works correctly in both parent graphs and subgraphs. - Supports both sync and async functions:: + Supports both sync and async functions. Pass ``timeout`` in milliseconds + to resume on either the operation or a ``WaitUntil`` timer:: - @durable_interrupt + @durable_interrupt(timeout=60_000) async def create_task(): return await client.tasks.create_async(...) @@ -126,28 +156,33 @@ def create_task_sync(): result = create_task_sync() """ - if asyncio.iscoroutinefunction(fn): + def decorate(func: F) -> F: + if asyncio.iscoroutinefunction(func): - @functools.wraps(fn) - async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + @functools.wraps(func) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + scratchpad, idx = _next_durable_index() + if _is_resumed(scratchpad, idx): + return _resume_interrupt(timeout) + result = await func(*args, **kwargs) + if isinstance(result, SkipInterruptValue): + return _inject_resume(scratchpad, result.resume_value) + return _interrupt_with_timeout(result, timeout) + + return async_wrapper # type: ignore[return-value] + + @functools.wraps(func) + def sync_wrapper(*args: Any, **kwargs: Any) -> Any: scratchpad, idx = _next_durable_index() if _is_resumed(scratchpad, idx): - return interrupt(None) - result = await fn(*args, **kwargs) + return _resume_interrupt(timeout) + result = func(*args, **kwargs) if isinstance(result, SkipInterruptValue): return _inject_resume(scratchpad, result.resume_value) - return interrupt(result) - - return async_wrapper # type: ignore[return-value] + return _interrupt_with_timeout(result, timeout) - @functools.wraps(fn) - def sync_wrapper(*args: Any, **kwargs: Any) -> Any: - scratchpad, idx = _next_durable_index() - if _is_resumed(scratchpad, idx): - return interrupt(None) - result = fn(*args, **kwargs) - if isinstance(result, SkipInterruptValue): - return _inject_resume(scratchpad, result.resume_value) - return interrupt(result) + return sync_wrapper # type: ignore[return-value] - return sync_wrapper # type: ignore[return-value] + if fn is None: + return decorate + return decorate(fn) diff --git a/src/uipath_langchain/agent/tools/extraction_tool.py b/src/uipath_langchain/agent/tools/extraction_tool.py index dd6ed9fbf..33a344966 100644 --- a/src/uipath_langchain/agent/tools/extraction_tool.py +++ b/src/uipath_langchain/agent/tools/extraction_tool.py @@ -6,7 +6,7 @@ from langchain.tools import BaseTool from langchain_core.messages import ToolCall, ToolMessage from langchain_core.tools import StructuredTool -from langgraph.types import Command, interrupt +from langgraph.types import Command from pydantic import BaseModel, Field from uipath.agent.models.agent import AgentIxpExtractionResourceConfig from uipath.eval.mocks import mockable @@ -14,6 +14,7 @@ from uipath.platform.documents import ExtractionResponseIXP from uipath.platform.errors import EnrichedException +from uipath_langchain._utils.durable_interrupt import durable_interrupt from uipath_langchain.agent.react.job_attachments import raise_for_job_attachment_error from uipath_langchain.agent.react.types import AgentGraphState from uipath_langchain.agent.tools.tool_node import ( @@ -94,13 +95,16 @@ async def extraction_tool_fn(**kwargs: Any) -> ExtractionResponseIXP: attachment_id=attachment.id, ) raise - document_extraction_response = interrupt( - DocumentExtraction( + + @durable_interrupt(timeout=resource.settings.timeout) + async def extract_document() -> Any: + return DocumentExtraction( project_name=project_name, tag=version_tag, file_path=attachment_local_file_path, ) - ) + + document_extraction_response = await extract_document() return document_extraction_response diff --git a/src/uipath_langchain/agent/tools/internal_tools/batch_transform_tool.py b/src/uipath_langchain/agent/tools/internal_tools/batch_transform_tool.py index 3f828c91c..07454abd0 100644 --- a/src/uipath_langchain/agent/tools/internal_tools/batch_transform_tool.py +++ b/src/uipath_langchain/agent/tools/internal_tools/batch_transform_tool.py @@ -67,6 +67,7 @@ def create_batch_transform_tool( tool_name = sanitize_tool_name(resource.name) properties = resource.properties settings = properties.settings + operation_timeout = resource.settings.timeout if resource.settings else None # Extract settings query_setting = settings.query @@ -130,7 +131,7 @@ async def batch_transform_tool_fn(**kwargs: Any) -> dict[str, Any]: example_calls=[], # Examples cannot be provided for internal tools ) async def invoke_batch_transform(**_tool_kwargs: Any): - @durable_interrupt + @durable_interrupt(timeout=operation_timeout) async def create_ephemeral_index(): uipath = UiPath() ephemeral_index = ( @@ -150,7 +151,7 @@ async def create_ephemeral_index(): else: ephemeral_index = index_result - @durable_interrupt + @durable_interrupt(timeout=operation_timeout) async def create_batch_transform(): return CreateBatchTransform( name=f"task-{uuid.uuid4()}", diff --git a/src/uipath_langchain/agent/tools/internal_tools/deeprag_tool.py b/src/uipath_langchain/agent/tools/internal_tools/deeprag_tool.py index 4f369b08d..540cec672 100644 --- a/src/uipath_langchain/agent/tools/internal_tools/deeprag_tool.py +++ b/src/uipath_langchain/agent/tools/internal_tools/deeprag_tool.py @@ -65,6 +65,7 @@ def create_deeprag_tool( tool_name = sanitize_tool_name(resource.name) properties = resource.properties settings = properties.settings + operation_timeout = resource.settings.timeout if resource.settings else None # Extract settings query_setting = settings.query @@ -114,7 +115,7 @@ async def deeprag_tool_fn(**kwargs: Any) -> dict[str, Any]: example_calls=[], # Examples cannot be provided for internal tools ) async def invoke_deeprag(**_tool_kwargs: Any): - @durable_interrupt + @durable_interrupt(timeout=operation_timeout) async def create_ephemeral_index(): uipath = UiPath() ephemeral_index = ( @@ -134,7 +135,7 @@ async def create_ephemeral_index(): else: ephemeral_index = index_result - @durable_interrupt + @durable_interrupt(timeout=operation_timeout) async def create_deeprag(): return CreateDeepRag( name=f"task-{uuid.uuid4()}", diff --git a/src/uipath_langchain/agent/tools/process_tool.py b/src/uipath_langchain/agent/tools/process_tool.py index 52b158ad2..12bfb9cb0 100644 --- a/src/uipath_langchain/agent/tools/process_tool.py +++ b/src/uipath_langchain/agent/tools/process_tool.py @@ -75,7 +75,7 @@ async def invoke_process(**_tool_kwargs: Any): parent_span_id = _span_context.pop("parent_span_id", None) parent_operation_id = _bts_context.pop("parent_operation_id", None) - @durable_interrupt + @durable_interrupt(timeout=resource.settings.timeout) async def start_job(): client = UiPath() try: diff --git a/src/uipath_langchain/agent/tools/tool_node.py b/src/uipath_langchain/agent/tools/tool_node.py index f28a3b28f..68e1dcbf2 100644 --- a/src/uipath_langchain/agent/tools/tool_node.py +++ b/src/uipath_langchain/agent/tools/tool_node.py @@ -10,6 +10,7 @@ from langgraph.errors import GraphBubbleUp from langgraph.types import Command from pydantic import BaseModel +from uipath.platform.common import UiPathTimeoutError from uipath.platform.resume_triggers import is_no_content_marker from uipath.runtime.errors import UiPathErrorCategory @@ -251,15 +252,16 @@ def _wrap_tool_error_handling( ) -> RunnableCallable: """Wrap a tool node to catch errors and return them as ToolMessages, rather than failing the entire graph execution. - Catch and re-raise GraphBubbleUp, since LangGraph uses exceptions for interrupt control flow. - This is so we don't swallow expected interrupts as tool errors. + Catch and re-raise GraphBubbleUp, since LangGraph uses exceptions for interrupt + control flow. Also re-raise UiPathTimeoutError so a configured operation timeout + fails the agent instead of being returned to the model as a tool result. (https://langchain-ai.github.io/langgraph/concepts/human_in_the_loop/) """ def _func(state: AgentGraphState) -> OutputType: try: return tool_node.invoke(state) - except GraphBubbleUp: + except (GraphBubbleUp, UiPathTimeoutError): raise except Exception as e: result = _get_tool_error_result(e, state, tool_name) @@ -270,7 +272,7 @@ def _func(state: AgentGraphState) -> OutputType: async def _afunc(state: AgentGraphState) -> OutputType: try: return await tool_node.ainvoke(state) - except GraphBubbleUp: + except (GraphBubbleUp, UiPathTimeoutError): raise except Exception as e: result = _get_tool_error_result(e, state, tool_name) diff --git a/tests/agent/tools/internal_tools/test_batch_transform_tool.py b/tests/agent/tools/internal_tools/test_batch_transform_tool.py index 7bb3c11e0..5ac31e2c3 100644 --- a/tests/agent/tools/internal_tools/test_batch_transform_tool.py +++ b/tests/agent/tools/internal_tools/test_batch_transform_tool.py @@ -13,11 +13,13 @@ AgentInternalBatchTransformToolProperties, AgentInternalToolResourceConfig, AgentInternalToolType, + AgentToolSettings, BatchTransformFileExtension, BatchTransformFileExtensionSetting, BatchTransformWebSearchGrounding, BatchTransformWebSearchGroundingSetting, ) +from uipath.platform.common import WaitUntil from uipath.platform.context_grounding.context_grounding_index import ( ContextGroundingIndex, ) @@ -165,7 +167,8 @@ async def test_create_batch_transform_tool_static_query_index_ready( resource_config_static, mock_llm, ): - """Test Batch Transform tool with static query when index is immediately ready.""" + """Test Batch Transform uses a composite interrupt when timeout is configured.""" + resource_config_static.settings = AgentToolSettings(timeout=60_000) # Setup mocks mock_uipath = AsyncMock() mock_uipath_class.return_value = mock_uipath @@ -209,6 +212,10 @@ async def test_create_batch_transform_tool_static_query_index_ready( assert tool.coroutine is not None result = await tool.coroutine(attachment=mock_attachment) + interrupt_value = mock_interrupt.call_args.args[0] + assert isinstance(interrupt_value, list) + assert isinstance(interrupt_value[1], WaitUntil) + # Verify result contains attachment info assert result == { "result": { diff --git a/tests/agent/tools/internal_tools/test_deeprag_tool.py b/tests/agent/tools/internal_tools/test_deeprag_tool.py index bcfe2cca8..12e7f7462 100644 --- a/tests/agent/tools/internal_tools/test_deeprag_tool.py +++ b/tests/agent/tools/internal_tools/test_deeprag_tool.py @@ -12,11 +12,13 @@ AgentInternalDeepRagToolProperties, AgentInternalToolResourceConfig, AgentInternalToolType, + AgentToolSettings, CitationMode, DeepRagCitationModeSetting, DeepRagFileExtension, DeepRagFileExtensionSetting, ) +from uipath.platform.common import WaitUntil from uipath.platform.context_grounding.context_grounding_index import ( ContextGroundingIndex, ) @@ -136,7 +138,8 @@ async def test_create_deeprag_tool_static_query_index_ready( resource_config_static, mock_llm, ): - """Test DeepRAG tool with static query when index is immediately ready.""" + """Test DeepRAG uses a composite interrupt when timeout is configured.""" + resource_config_static.settings = AgentToolSettings(timeout=60_000) # Setup mocks mock_uipath = AsyncMock() mock_uipath_class.return_value = mock_uipath @@ -188,6 +191,9 @@ async def test_create_deeprag_tool_static_query_index_ready( # Only create_deeprag calls interrupt(); index was instant-resumed assert mock_interrupt.call_count == 1 + interrupt_value = mock_interrupt.call_args.args[0] + assert isinstance(interrupt_value, list) + assert isinstance(interrupt_value[1], WaitUntil) @patch( "uipath_langchain.agent.wrappers.job_attachment_wrapper.get_job_attachment_wrapper" diff --git a/tests/agent/tools/test_durable_interrupt.py b/tests/agent/tools/test_durable_interrupt.py index 6f6825561..75f994254 100644 --- a/tests/agent/tools/test_durable_interrupt.py +++ b/tests/agent/tools/test_durable_interrupt.py @@ -1,11 +1,14 @@ """Tests for the durable_interrupt decorator.""" from collections.abc import Generator +from datetime import UTC, datetime from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest from langgraph._internal._constants import CONFIG_KEY_SCRATCHPAD +from uipath.core.triggers import UIPATH_METADATA_KEY +from uipath.platform.common import UiPathTimeoutError, WaitUntil from uipath_langchain._utils.durable_interrupt import ( _durable_state, @@ -119,6 +122,76 @@ async def start_job() -> str: assert result == "resume-value" +class TestTimeout: + """Configured timeouts use a composite interrupt and reject timer resumes.""" + + @patch(PATCH_INTERRUPT) + @patch(PATCH_GET_CONFIG) + async def test_async_first_execution_adds_wait_until( + self, mock_get_config: MagicMock, mock_interrupt: MagicMock + ) -> None: + scratchpad = FakeScratchpad(resume=[]) + mock_get_config.return_value = _make_config(scratchpad) + mock_interrupt.side_effect = lambda value: value + before = datetime.now(UTC) + + @durable_interrupt(timeout=60_000) + async def start_job() -> dict[str, str]: + return {"wait": "job-123"} + + result = await start_job() + + interrupt_value = mock_interrupt.call_args.args[0] + assert interrupt_value[0] == {"wait": "job-123"} + assert isinstance(interrupt_value[1], WaitUntil) + assert interrupt_value[1].resume_time >= before + assert result is interrupt_value + + @patch(PATCH_INTERRUPT) + @patch(PATCH_GET_CONFIG) + def test_sync_first_execution_adds_wait_until( + self, mock_get_config: MagicMock, mock_interrupt: MagicMock + ) -> None: + scratchpad = FakeScratchpad(resume=[]) + mock_get_config.return_value = _make_config(scratchpad) + mock_interrupt.side_effect = lambda value: value + + @durable_interrupt(timeout=30_000) + def create_task() -> str: + return "task-123" + + result = create_task() + + interrupt_value = mock_interrupt.call_args.args[0] + assert interrupt_value[0] == "task-123" + assert isinstance(interrupt_value[1], WaitUntil) + assert result is interrupt_value + + @patch(PATCH_INTERRUPT) + @patch(PATCH_GET_CONFIG) + async def test_timeout_resume_raises( + self, mock_get_config: MagicMock, mock_interrupt: MagicMock + ) -> None: + scratchpad = FakeScratchpad(resume=["timer-result"]) + mock_get_config.return_value = _make_config(scratchpad) + timeout_result = { + UIPATH_METADATA_KEY: { + "triggerType": "Timer", + "triggerName": "Timer", + } + } + mock_interrupt.return_value = timeout_result + + @durable_interrupt(timeout=30_000) + async def start_job() -> str: + return "should-not-reach" + + with pytest.raises(UiPathTimeoutError) as exc_info: + await start_job() + + assert exc_info.value.value is timeout_result + + class TestSyncFirstExecution: """Sync first execution (no resume values): body runs and interrupt is called with result.""" diff --git a/tests/agent/tools/test_extraction_tool.py b/tests/agent/tools/test_extraction_tool.py index f6221fadb..b9083cb38 100644 --- a/tests/agent/tools/test_extraction_tool.py +++ b/tests/agent/tools/test_extraction_tool.py @@ -8,8 +8,10 @@ from uipath.agent.models.agent import ( AgentIxpExtractionResourceConfig, AgentIxpExtractionToolProperties, + AgentToolSettings, ) from uipath.platform.attachments import Attachment +from uipath.platform.common import WaitUntil from uipath.platform.documents import ExtractionResponseIXP from uipath.platform.errors import EnrichedException from uipath.runtime.errors import UiPathErrorCategory @@ -166,11 +168,12 @@ def extraction_resource(self): @pytest.mark.asyncio @patch("uipath.platform.UiPath") - @patch("uipath_langchain.agent.tools.extraction_tool.interrupt") + @patch("uipath_langchain._utils.durable_interrupt.decorator.interrupt") async def test_extraction_tool_downloads_attachment_and_calls_interrupt( self, mock_interrupt, mock_uipath_class, extraction_resource ): - """Test that extraction tool downloads attachment and calls interrupt with correct params.""" + """Test extraction adds WaitUntil when a timeout is configured.""" + extraction_resource.settings = AgentToolSettings(timeout=60_000) mock_client = MagicMock() mock_uipath_class.return_value = mock_client mock_client.attachments.download_async = AsyncMock( @@ -194,16 +197,19 @@ async def test_extraction_tool_downloads_attachment_and_calls_interrupt( ) assert mock_interrupt.called - interrupt_arg = mock_interrupt.call_args[0][0] + interrupt_value = mock_interrupt.call_args[0][0] + assert isinstance(interrupt_value, list) + interrupt_arg, wait_until_arg = interrupt_value assert interrupt_arg.project_name == "TestProject" assert interrupt_arg.tag == "v1.0" assert interrupt_arg.file_path == "/path/to/document.pdf" + assert isinstance(wait_until_arg, WaitUntil) assert result == {"extracted_data": {"field1": "value1"}} @pytest.mark.asyncio @patch("uipath.platform.UiPath") - @patch("uipath_langchain.agent.tools.extraction_tool.interrupt") + @patch("uipath_langchain._utils.durable_interrupt.decorator.interrupt") async def test_extraction_tool_with_different_version_tag( self, mock_interrupt, mock_uipath_class ): @@ -327,7 +333,7 @@ async def test_extraction_tool_missing_attachment_raises_system( @pytest.mark.asyncio @patch("uipath.platform.UiPath") - @patch("uipath_langchain.agent.tools.extraction_tool.interrupt") + @patch("uipath_langchain._utils.durable_interrupt.decorator.interrupt") async def test_extraction_tool_handles_alias_keyed_input( self, mock_interrupt, mock_uipath_class, extraction_resource ): diff --git a/tests/agent/tools/test_process_tool.py b/tests/agent/tools/test_process_tool.py index 4da3d92c2..8ac01bf2f 100644 --- a/tests/agent/tools/test_process_tool.py +++ b/tests/agent/tools/test_process_tool.py @@ -9,7 +9,7 @@ AgentProcessToolResourceConfig, AgentToolType, ) -from uipath.platform.common import WaitJob +from uipath.platform.common import WaitJob, WaitUntil from uipath.platform.orchestrator import Job from uipath_langchain.agent.tools.process_tool import create_process_tool @@ -188,7 +188,8 @@ async def test_invoke_calls_processes_invoke_async( async def test_invoke_interrupts_with_wait_job( self, mock_uipath_class, mock_interrupt, process_resource ): - """Test that after invoking, the tool interrupts with WaitJobRaw.""" + """Test that a configured timeout adds WaitUntil beside WaitJobRaw.""" + process_resource.settings.timeout = 30_000 mock_job = MagicMock(spec=Job) mock_job.key = "job-key-456" mock_job.folder_key = "folder-key-456" @@ -207,10 +208,13 @@ async def test_invoke_interrupts_with_wait_job( await tool.ainvoke({}) mock_interrupt.assert_called_once() - wait_job_arg = mock_interrupt.call_args[0][0] + interrupt_value = mock_interrupt.call_args[0][0] + assert isinstance(interrupt_value, list) + wait_job_arg, wait_until_arg = interrupt_value assert isinstance(wait_job_arg, WaitJob) assert wait_job_arg.job == mock_job assert wait_job_arg.process_folder_key == "folder-key-456" + assert isinstance(wait_until_arg, WaitUntil) @pytest.mark.asyncio @patch.dict(os.environ, {"UIPATH_FOLDER_PATH": "/Shared/DataFolder"}) diff --git a/tests/agent/tools/test_tool_node.py b/tests/agent/tools/test_tool_node.py index 1212e2be8..f1e19430c 100644 --- a/tests/agent/tools/test_tool_node.py +++ b/tests/agent/tools/test_tool_node.py @@ -10,6 +10,7 @@ from langchain_core.tools import BaseTool from langgraph.types import Command from pydantic import BaseModel +from uipath.platform.common import UiPathTimeoutError from uipath_langchain.agent.exceptions import ( AgentRuntimeError, @@ -68,6 +69,19 @@ async def _arun(self, input_text: str = "") -> str: raise ValueError(f"Async tool execution failed: {input_text}") +class MockTimeoutTool(BaseTool): + """Mock tool that always times out.""" + + name: str = "mock_timeout_tool" + description: str = "A mock tool that times out" + + def _run(self, input_text: str = "") -> str: + raise UiPathTimeoutError({"input": input_text}) + + async def _arun(self, input_text: str = "") -> str: + raise UiPathTimeoutError({"input": input_text}) + + class FilteredState(BaseModel): """Mock filtered state model for testing wrappers.""" @@ -442,6 +456,26 @@ async def test_wrap_tools_with_error_handling_captures_error(self): assert isinstance(tool_message.content, str) assert "Async tool execution failed: test input" in tool_message.content + async def test_wrap_tools_with_error_handling_reraises_timeout(self): + """A timeout fails the agent instead of becoming an error ToolMessage.""" + timeout_tool = MockTimeoutTool() + tool_call = { + "name": "mock_timeout_tool", + "args": {"input_text": "test input"}, + "id": "test_call_id", + } + state = AgentGraphState( + messages=[AIMessage(content="Using tool", tool_calls=[tool_call])] + ) + + node = UiPathToolNode(timeout_tool) + wrapped = wrap_tools_with_error_handling({"mock_timeout_tool": node})[ + "mock_timeout_tool" + ] + + with pytest.raises(UiPathTimeoutError, match="UiPath interrupt timed out"): + await wrapped.ainvoke(state) + class TestToolNodeConfirmation: """Tests for confirmation flow in UiPathToolNode._func / _afunc.""" diff --git a/uv.lock b/uv.lock index a90039156..eedd4ee6f 100644 --- a/uv.lock +++ b/uv.lock @@ -4498,7 +4498,7 @@ wheels = [ [[package]] name = "uipath-langchain" -version = "0.14.16" +version = "0.14.17" source = { editable = "." } dependencies = [ { name = "a2a-sdk" },