fix: use durable_interrupt for tool confirmation

Author: JoshParkSJ

src/uipath_langchain/agent/tools/durable_interrupt/__init__.py

@@ -2,16 +2,12 @@
 
 from .decorator import (
     _durable_state,
-    _interrupt_offset,
-    add_interrupt_offset,
     durable_interrupt,
 )
 from .skip_interrupt import SkipInterruptValue
 
 __all__ = [
-    "add_interrupt_offset",
     "durable_interrupt",
     "SkipInterruptValue",
     "_durable_state",
-    "_interrupt_offset",
 ]

src/uipath_langchain/agent/tools/durable_interrupt/decorator.py

@@ -48,15 +48,6 @@ async def start_job():
 _durable_state: contextvars.ContextVar[tuple[int, int] | None] = contextvars.ContextVar(
     "_durable_interrupt_state", default=None
 )
-# Number of interrupt() calls before the first @durable_interrupt
-_interrupt_offset: contextvars.ContextVar[int] = contextvars.ContextVar(
-    "_durable_interrupt_offset", default=0
-)
-
-
-def add_interrupt_offset(n: int = 1) -> None:
-    """Increment durable_interrupt's starting index offset by n"""
-    _interrupt_offset.set(_interrupt_offset.get(0) + n)
 
 
 def _next_durable_index() -> tuple[Any, int]:
@@ -76,8 +67,7 @@ def _next_durable_index() -> tuple[Any, int]:
     state = _durable_state.get()
 
     if state is None or state[0] != sp_id:
-        idx = _interrupt_offset.get(0)
-        _interrupt_offset.set(0)  # consume offset
+        idx = 0
     else:
         idx = state[1]
 

src/uipath_langchain/chat/hitl.py

@@ -7,7 +7,6 @@
 from langchain_core.messages.tool import ToolCall, ToolMessage
 from langchain_core.tools import BaseTool, InjectedToolCallId
 from langchain_core.tools import tool as langchain_tool
-from langgraph.types import interrupt
 from uipath.core.chat import (
     UiPathConversationToolCallConfirmationValue,
 )
@@ -119,21 +118,20 @@ def request_approval(
     if tool_call_schema is not None:
         input_schema = tool_call_schema.model_json_schema()
 
-    response = interrupt(
-        UiPathConversationToolCallConfirmationValue(
+    # Lazy import to avoid circular dependency:
+    # hitl -> agent.tools.durable_interrupt -> agent.tools -> tool_node -> hitl
+    from uipath_langchain.agent.tools.durable_interrupt import durable_interrupt
+
+    @durable_interrupt
+    def ask_confirmation():
+        return UiPathConversationToolCallConfirmationValue(
             tool_call_id=tool_call_id,
             tool_name=tool.name,
             input_schema=input_schema,
             input_value=tool_args,
         )
-    )
-    # Lazy import to avoid circular dependency:
-    # hitl -> agent.tools.durable_interrupt -> agent.tools -> tool_node -> hitl
-    from uipath_langchain.agent.tools.durable_interrupt import add_interrupt_offset
 
-    # Workaround for langgraph#6792 — remove when subgraph @task + interrupt()
-    # checkpoint caching is fixed upstream
-    add_interrupt_offset()
+    response = ask_confirmation()
 
     # The resume payload from CAS has shape:
     #   {"type": "uipath_cas_tool_call_confirmation",

src/uipath_langchain/runtime/messages.py

@@ -433,9 +433,7 @@ async def map_tool_message_to_events(
                 # Keep as string if not valid JSON
                 pass
 
-        events: list[UiPathConversationMessageEvent] = []
-
-        events.append(
+        events = [
             UiPathConversationMessageEvent(
                 message_id=message_id,
                 tool_call=UiPathConversationToolCallEvent(
@@ -447,7 +445,7 @@ async def map_tool_message_to_events(
                     ),
                 ),
             )
-        )
+        ]
 
         if is_last_tool_call:
             events.append(self.map_to_message_end_event(message_id))

tests/agent/tools/test_durable_interrupt.py

@@ -9,8 +9,6 @@
 
 from uipath_langchain.agent.tools.durable_interrupt import (
     _durable_state,
-    _interrupt_offset,
-    add_interrupt_offset,
     durable_interrupt,
 )
 
@@ -35,12 +33,10 @@ def _make_config(scratchpad: FakeScratchpad | None = None) -> dict[str, Any]:
 
 @pytest.fixture(autouse=True)
 def _reset_durable_state() -> Generator[None]:
-    """Reset per-node counter and offset between tests for isolation."""
-    state_token = _durable_state.set(None)
-    offset_token = _interrupt_offset.set(0)
+    """Reset per-node counter between tests for isolation."""
+    token = _durable_state.set(None)
     yield
-    _durable_state.reset(state_token)
-    _interrupt_offset.reset(offset_token)
+    _durable_state.reset(token)
 
 
 class TestAsyncFirstExecution:
@@ -515,161 +511,3 @@ def my_sync_fn() -> str:
             return "value"
 
         assert not asyncio.iscoroutinefunction(my_sync_fn)
-
-
-class TestInterruptOffsetWithPriorInterrupts:
-    """Offset accounts for prior interrupt() calls (e.g. HITL) in the same node."""
-
-    @patch(PATCH_INTERRUPT)
-    @patch(PATCH_GET_CONFIG)
-    def test_offset_skips_hitl_resume_slot(
-        self, mock_get_config: MagicMock, mock_interrupt: MagicMock
-    ) -> None:
-        """With resume=["hitl_value"] and offset=1, durable's idx=1, 1 < 1 is False → body runs."""
-        scratchpad = FakeScratchpad(resume=["hitl_value"])
-        mock_get_config.return_value = _make_config(scratchpad)
-        mock_interrupt.side_effect = lambda v: v
-
-        action = MagicMock(return_value="job-started")
-
-        @durable_interrupt
-        def start_job() -> str:
-            return action()
-
-        add_interrupt_offset()
-        start_job()
-
-        action.assert_called_once()
-        mock_interrupt.assert_called_once_with("job-started")
-
-    @patch(PATCH_INTERRUPT, return_value="durable-result")
-    @patch(PATCH_GET_CONFIG)
-    def test_offset_full_resume(
-        self, mock_get_config: MagicMock, mock_interrupt: MagicMock
-    ) -> None:
-        """With resume=["hitl_value", "durable_result"] and offset=1, durable reads index 1."""
-        scratchpad = FakeScratchpad(resume=["hitl_value", "durable_result"])
-        mock_get_config.return_value = _make_config(scratchpad)
-
-        action = MagicMock()
-
-        @durable_interrupt
-        def start_job() -> Any:
-            return action()
-
-        add_interrupt_offset()
-        result = start_job()
-
-        action.assert_not_called()  # body skipped — idx=1 < len(resume)=2
-        mock_interrupt.assert_called_once_with(None)
-        assert result == "durable-result"
-
-    @patch(PATCH_INTERRUPT, return_value="resumed")
-    @patch(PATCH_GET_CONFIG)
-    def test_offset_resets_on_scratchpad_change(
-        self, mock_get_config: MagicMock, mock_interrupt: MagicMock
-    ) -> None:
-        """After consuming offset on first scratchpad, a new scratchpad starts at 0."""
-        sp1 = FakeScratchpad(resume=["hitl", "durable"])
-        mock_get_config.return_value = _make_config(sp1)
-
-        @durable_interrupt
-        def task_a() -> str:
-            return "should-not-run"
-
-        add_interrupt_offset()
-        task_a()  # consumes offset, idx=1
-
-        # New scratchpad — offset should have been consumed/reset
-        sp2 = FakeScratchpad(resume=["val"])
-        mock_get_config.return_value = _make_config(sp2)
-
-        action = MagicMock()
-
-        @durable_interrupt
-        def task_b() -> Any:
-            return action()
-
-        task_b()  # idx should be 0, not 1
-
-        action.assert_not_called()  # idx=0 < len(resume)=1 → skipped
-
-    @patch(PATCH_INTERRUPT)
-    @patch(PATCH_GET_CONFIG)
-    def test_offset_with_multiple_durable_interrupts(
-        self, mock_get_config: MagicMock, mock_interrupt: MagicMock
-    ) -> None:
-        """HITL + two durable_interrupts: offset shifts both indices correctly."""
-        # resume[0] = HITL value, resume[1] = durable_a result
-        # durable_a at idx=1 → resumed, durable_b at idx=2 → body runs
-        scratchpad = FakeScratchpad(resume=["hitl_approval", "durable_a_result"])
-        mock_get_config.return_value = _make_config(scratchpad)
-        mock_interrupt.side_effect = lambda v: f"interrupt({v})"
-
-        action_a = MagicMock()
-        action_b = MagicMock(return_value="job-B")
-
-        @durable_interrupt
-        def durable_a() -> Any:
-            return action_a()
-
-        @durable_interrupt
-        def durable_b() -> str:
-            return action_b()
-
-        add_interrupt_offset()
-        result_a = durable_a()  # idx=1, 1 < 2 → skipped
-        result_b = durable_b()  # idx=2, 2 < 2 → False → body runs
-
-        action_a.assert_not_called()
-        action_b.assert_called_once()
-        assert result_a == "interrupt(None)"
-        assert result_b == "interrupt(job-B)"
-
-    @patch(PATCH_INTERRUPT)
-    @patch(PATCH_GET_CONFIG)
-    def test_multiple_offsets_accumulate(
-        self, mock_get_config: MagicMock, mock_interrupt: MagicMock
-    ) -> None:
-        """Two prior interrupt() calls → offset accumulates to 2."""
-        # resume[0] = first confirmation, resume[1] = second confirmation
-        # durable at idx=2 → body runs (2 < 2 is False)
-        scratchpad = FakeScratchpad(resume=["confirm_1", "confirm_2"])
-        mock_get_config.return_value = _make_config(scratchpad)
-        mock_interrupt.side_effect = lambda v: v
-
-        action = MagicMock(return_value="job-started")
-
-        @durable_interrupt
-        def start_job() -> str:
-            return action()
-
-        add_interrupt_offset()  # first confirmation
-        add_interrupt_offset()  # second confirmation
-        start_job()
-
-        action.assert_called_once()  # idx=2, 2 < 2 → False → body runs
-        mock_interrupt.assert_called_once_with("job-started")
-
-    @patch(PATCH_INTERRUPT, return_value="durable-result")
-    @patch(PATCH_GET_CONFIG)
-    def test_multiple_offsets_full_resume(
-        self, mock_get_config: MagicMock, mock_interrupt: MagicMock
-    ) -> None:
-        """Two prior interrupts fully resumed + durable resumed → body skipped."""
-        scratchpad = FakeScratchpad(resume=["confirm_1", "confirm_2", "durable-result"])
-        mock_get_config.return_value = _make_config(scratchpad)
-
-        action = MagicMock()
-
-        @durable_interrupt
-        def start_job() -> Any:
-            return action()
-
-        add_interrupt_offset()
-        add_interrupt_offset()
-        result = start_job()
-
-        action.assert_not_called()  # idx=2, 2 < 3 → True → skipped
-        mock_interrupt.assert_called_once_with(None)
-        assert result == "durable-result"

tests/agent/tools/test_tool_node.py

@@ -15,7 +15,6 @@
     AgentRuntimeError,
     AgentRuntimeErrorCode,
 )
-from uipath_langchain.agent.react.types import AgentGraphState
 from uipath_langchain.agent.tools.tool_node import (
     ToolWrapperMixin,
     UiPathToolNode,
@@ -73,9 +72,10 @@ class FilteredState(BaseModel):
     session_id: str = "test_session"
 
 
-class MockState(AgentGraphState):
+class MockState(BaseModel):
     """Mock state for testing."""
 
+    messages: list[Any] = []
     user_id: str = "test_user"
     session_id: str = "test_session"
 
@@ -316,7 +316,8 @@ def test_tool_error_propagates_when_handle_errors_false(self, mock_state):
         node = UiPathToolNode(failing_tool, handle_tool_errors=False)
 
         with pytest.raises(ValueError) as exc_info:
-            node._func(state)
+            node._func(state)  # type: ignore[arg-type]
+
         assert "Tool execution failed: test input" in str(exc_info.value)
 
     async def test_async_tool_error_propagates_when_handle_errors_false(self):
@@ -333,7 +334,8 @@ async def test_async_tool_error_propagates_when_handle_errors_false(self):
         node = UiPathToolNode(failing_tool, handle_tool_errors=False)
 
         with pytest.raises(ValueError) as exc_info:
-            await node._afunc(state)
+            await node._afunc(state)  # type: ignore[arg-type]
+
         assert "Async tool execution failed: test input" in str(exc_info.value)
 
     def test_tool_error_captured_when_handle_errors_true(self):
@@ -349,7 +351,8 @@ def test_tool_error_captured_when_handle_errors_true(self):
 
         node = UiPathToolNode(failing_tool, handle_tool_errors=True)
 
-        result = node._func(state)
+        result = node._func(state)  # type: ignore[arg-type]
+
         assert result is not None
         assert isinstance(result, dict)
         assert "messages" in result
@@ -375,7 +378,8 @@ async def test_async_tool_error_captured_when_handle_errors_true(self):
 
         node = UiPathToolNode(failing_tool, handle_tool_errors=True)
 
-        result = await node._afunc(state)
+        result = await node._afunc(state)  # type: ignore[arg-type]
+
         assert result is not None
         assert isinstance(result, dict)
         assert "messages" in result

tests/chat/test_hitl.py

@@ -146,39 +146,39 @@ def test_annotate_wraps_content_when_modified(self):
 class TestRequestApprovalTruthiness:
     """Tests for the truthiness fix in request_approval."""
 
-    @patch("uipath_langchain.chat.hitl.interrupt")
+    @patch("uipath_langchain.agent.tools.durable_interrupt.decorator.interrupt")
     def test_empty_dict_input_preserved(self, mock_interrupt):
         """Empty dict from user edits should not be replaced by original args."""
         mock_interrupt.return_value = {"value": {"approved": True, "input": {}}}
         tool = MockTool()
         result = request_approval({"query": "test", "tool_call_id": "c1"}, tool)
         assert result == {}
 
-    @patch("uipath_langchain.chat.hitl.interrupt")
+    @patch("uipath_langchain.agent.tools.durable_interrupt.decorator.interrupt")
     def test_empty_list_input_preserved(self, mock_interrupt):
         """Empty list from user edits should not be replaced by original args."""
         mock_interrupt.return_value = {"value": {"approved": True, "input": []}}
         tool = MockTool()
         result = request_approval({"query": "test", "tool_call_id": "c1"}, tool)
         assert result == []
 
-    @patch("uipath_langchain.chat.hitl.interrupt")
+    @patch("uipath_langchain.agent.tools.durable_interrupt.decorator.interrupt")
     def test_none_input_falls_back_to_original(self, mock_interrupt):
         """None input should fall back to original tool_args."""
         mock_interrupt.return_value = {"value": {"approved": True, "input": None}}
         tool = MockTool()
         result = request_approval({"query": "test", "tool_call_id": "c1"}, tool)
         assert result == {"query": "test"}
 
-    @patch("uipath_langchain.chat.hitl.interrupt")
+    @patch("uipath_langchain.agent.tools.durable_interrupt.decorator.interrupt")
     def test_missing_input_falls_back_to_original(self, mock_interrupt):
         """Missing input key should fall back to original tool_args."""
         mock_interrupt.return_value = {"value": {"approved": True}}
         tool = MockTool()
         result = request_approval({"query": "test", "tool_call_id": "c1"}, tool)
         assert result == {"query": "test"}
 
-    @patch("uipath_langchain.chat.hitl.interrupt")
+    @patch("uipath_langchain.agent.tools.durable_interrupt.decorator.interrupt")
     def test_rejected_returns_none(self, mock_interrupt):
         """Rejected approval returns None."""
         mock_interrupt.return_value = {"value": {"approved": False}}