feat: enforce configurable tool timeouts

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"

src/uipath_langchain/_utils/durable_interrupt/decorator.py

@@ -33,11 +33,13 @@
 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 @@
     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 @@
     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 @@
         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)

src/uipath_langchain/agent/tools/extraction_tool.py

@@ -6,14 +6,15 @@
 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
 from uipath.platform.common import DocumentExtraction
 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
 

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()}",

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()}",

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:

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)

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": {

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"