feat: process tools tracing and interrupt (#546)

Author: cotovanu-cristian

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.5.39"
+version = "0.5.40"
 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/tools/process_tool.py

@@ -5,10 +5,12 @@
 from langchain.tools import BaseTool
 from langchain_core.messages import ToolCall
 from langchain_core.tools import StructuredTool
+from langgraph.func import task
 from langgraph.types import interrupt
 from uipath.agent.models.agent import AgentProcessToolResourceConfig
 from uipath.eval.mocks import mockable
-from uipath.platform.common import InvokeProcess
+from uipath.platform import UiPath
+from uipath.platform.common import WaitJob
 
 from uipath_langchain.agent.react.job_attachments import get_job_attachments
 from uipath_langchain.agent.react.jsonschema_pydantic_converter import create_model
@@ -36,25 +38,37 @@ def create_process_tool(resource: AgentProcessToolResourceConfig) -> StructuredT
     input_model: Any = create_model(resource.input_schema)
     output_model: Any = create_model(resource.output_schema)
 
-    @mockable(
-        name=resource.name,
-        description=resource.description,
-        input_schema=input_model.model_json_schema(),
-        output_schema=output_model.model_json_schema(),
-        example_calls=resource.properties.example_calls,
-    )
+    _span_context: dict[str, Any] = {}
+
     async def process_tool_fn(**kwargs: Any):
         attachments = get_job_attachments(input_model, kwargs)
         input_arguments = input_model.model_validate(kwargs).model_dump(mode="json")
-        return interrupt(
-            InvokeProcess(
-                name=process_name,
-                input_arguments=input_arguments,
-                process_folder_path=folder_path,
-                process_folder_key=None,
-                attachments=attachments,
-            )
+
+        @mockable(
+            name=tool_name.lower(),
+            description=resource.description,
+            input_schema=input_model.model_json_schema(),
+            output_schema=output_model.model_json_schema(),
+            example_calls=resource.properties.example_calls,
         )
+        async def invoke_process():
+            parent_span_id = _span_context.pop("parent_span_id", None)
+
+            @task
+            async def start_job():
+                client = UiPath()
+                return await client.processes.invoke_async(
+                    name=process_name,
+                    input_arguments=input_arguments,
+                    folder_path=folder_path,
+                    attachments=attachments,
+                    parent_span_id=parent_span_id,
+                )
+
+            job = await start_job()
+            return interrupt(WaitJob(job=job, process_folder_key=job.folder_key))
+
+        return await invoke_process()
 
     job_attachment_wrapper = get_job_attachment_wrapper(output_type=output_model)
 
@@ -73,11 +87,12 @@ async def process_tool_wrapper(
         coroutine=process_tool_fn,
         output_type=output_model,
         metadata={
-            "tool_type": "process",
+            "tool_type": resource.type.lower(),
             "display_name": process_name,
             "folder_path": folder_path,
             "args_schema": input_model,
             "output_schema": output_model,
+            "_span_context": _span_context,
         },
         argument_properties=resource.argument_properties,
     )

tests/agent/tools/test_process_tool.py

@@ -1,33 +1,77 @@
-"""Tests for process_tool.py metadata."""
+"""Tests for process_tool.py."""
+
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from uipath.agent.models.agent import (
     AgentProcessToolProperties,
     AgentProcessToolResourceConfig,
     AgentToolType,
 )
+from uipath.platform.common import WaitJob
+from uipath.platform.orchestrator import Job
 
 from uipath_langchain.agent.tools.process_tool import create_process_tool
 
 
+def _noop_task(fn):
+    """No-op replacement for @task so it works outside Pregel context."""
+    return fn
+
+
+@pytest.fixture(autouse=True)
+def _patch_lg_task():
+    """Patch @task decorator to no-op since unit tests run outside Pregel context."""
+    with patch("uipath_langchain.agent.tools.process_tool.task", _noop_task):
+        yield
+
+
+@pytest.fixture
+def process_resource():
+    """Create a minimal process tool resource config."""
+    return AgentProcessToolResourceConfig(
+        type=AgentToolType.PROCESS,
+        name="test_process",
+        description="Test process description",
+        input_schema={"type": "object", "properties": {}},
+        output_schema={"type": "object", "properties": {}},
+        properties=AgentProcessToolProperties(
+            process_name="MyProcess",
+            folder_path="/Shared/MyFolder",
+        ),
+    )
+
+
+@pytest.fixture
+def process_resource_with_inputs():
+    """Create a process tool resource config with input arguments."""
+    return AgentProcessToolResourceConfig(
+        type=AgentToolType.PROCESS,
+        name="data_processor",
+        description="Process data with arguments",
+        input_schema={
+            "type": "object",
+            "properties": {
+                "name": {"type": "string"},
+                "count": {"type": "integer"},
+            },
+        },
+        output_schema={
+            "type": "object",
+            "properties": {
+                "result": {"type": "string"},
+            },
+        },
+        properties=AgentProcessToolProperties(
+            process_name="DataProcessor",
+            folder_path="/Shared/DataFolder",
+        ),
+    )
+
+
 class TestProcessToolMetadata:
     """Test that process tool has correct metadata for observability."""
 
-    @pytest.fixture
-    def process_resource(self):
-        """Create a minimal process tool resource config."""
-        return AgentProcessToolResourceConfig(
-            type=AgentToolType.PROCESS,
-            name="test_process",
-            description="Test process description",
-            input_schema={"type": "object", "properties": {}},
-            output_schema={"type": "object", "properties": {}},
-            properties=AgentProcessToolProperties(
-                process_name="MyProcess",
-                folder_path="/Shared/MyFolder",
-            ),
-        )
-
     def test_process_tool_has_metadata(self, process_resource):
         """Test that process tool has metadata dict."""
         tool = create_process_tool(process_resource)
@@ -36,7 +80,7 @@ def test_process_tool_has_metadata(self, process_resource):
         assert isinstance(tool.metadata, dict)
 
     def test_process_tool_metadata_has_tool_type(self, process_resource):
-        """Test that metadata contains tool_type for span detection."""
+        """Test that metadata contains tool_type derived from resource type."""
         tool = create_process_tool(process_resource)
         assert tool.metadata is not None
         assert tool.metadata["tool_type"] == "process"
@@ -52,3 +96,207 @@ def test_process_tool_metadata_has_folder_path(self, process_resource):
         tool = create_process_tool(process_resource)
         assert tool.metadata is not None
         assert tool.metadata["folder_path"] == "/Shared/MyFolder"
+
+    def test_process_tool_metadata_has_span_context(self, process_resource):
+        """Test that metadata contains _span_context dict for tracing."""
+        tool = create_process_tool(process_resource)
+        assert tool.metadata is not None
+        assert "_span_context" in tool.metadata
+        assert isinstance(tool.metadata["_span_context"], dict)
+
+    def test_process_tool_metadata_tool_type_uses_resource_type(self):
+        """Test that tool_type is derived from resource.type.lower()."""
+        resource = AgentProcessToolResourceConfig(
+            type=AgentToolType.PROCESS,
+            name="test_process",
+            description="Test",
+            input_schema={"type": "object", "properties": {}},
+            output_schema={"type": "object", "properties": {}},
+            properties=AgentProcessToolProperties(
+                process_name="MyProcess",
+                folder_path="/Shared",
+            ),
+        )
+        tool = create_process_tool(resource)
+        assert tool.metadata is not None
+        assert tool.metadata["tool_type"] == resource.type.lower()
+
+
+class TestProcessToolInvocation:
+    """Test process tool invocation behavior: invoke then interrupt."""
+
+    @pytest.mark.asyncio
+    @patch("uipath_langchain.agent.tools.process_tool.interrupt")
+    @patch("uipath_langchain.agent.tools.process_tool.UiPath")
+    async def test_invoke_calls_processes_invoke_async(
+        self, mock_uipath_class, mock_interrupt, process_resource
+    ):
+        """Test that invoking the tool calls client.processes.invoke_async."""
+        mock_job = MagicMock(spec=Job)
+        mock_job.folder_key = "folder-key-123"
+
+        mock_client = MagicMock()
+        mock_client.processes.invoke_async = AsyncMock(return_value=mock_job)
+        mock_uipath_class.return_value = mock_client
+
+        mock_interrupt.return_value = {"output": "result"}
+
+        tool = create_process_tool(process_resource)
+        await tool.ainvoke({})
+
+        mock_client.processes.invoke_async.assert_called_once_with(
+            name="MyProcess",
+            input_arguments={},
+            folder_path="/Shared/MyFolder",
+            attachments=[],
+            parent_span_id=None,
+        )
+
+    @pytest.mark.asyncio
+    @patch("uipath_langchain.agent.tools.process_tool.interrupt")
+    @patch("uipath_langchain.agent.tools.process_tool.UiPath")
+    async def test_invoke_interrupts_with_wait_job(
+        self, mock_uipath_class, mock_interrupt, process_resource
+    ):
+        """Test that after invoking, the tool interrupts with WaitJob."""
+        mock_job = MagicMock(spec=Job)
+        mock_job.folder_key = "folder-key-456"
+
+        mock_client = MagicMock()
+        mock_client.processes.invoke_async = AsyncMock(return_value=mock_job)
+        mock_uipath_class.return_value = mock_client
+
+        mock_interrupt.return_value = {"output": "done"}
+
+        tool = create_process_tool(process_resource)
+        await tool.ainvoke({})
+
+        mock_interrupt.assert_called_once()
+        wait_job_arg = mock_interrupt.call_args[0][0]
+        assert isinstance(wait_job_arg, WaitJob)
+        assert wait_job_arg.job == mock_job
+        assert wait_job_arg.process_folder_key == "folder-key-456"
+
+    @pytest.mark.asyncio
+    @patch("uipath_langchain.agent.tools.process_tool.interrupt")
+    @patch("uipath_langchain.agent.tools.process_tool.UiPath")
+    async def test_invoke_passes_input_arguments(
+        self, mock_uipath_class, mock_interrupt, process_resource_with_inputs
+    ):
+        """Test that input arguments are correctly passed to invoke_async."""
+        mock_job = MagicMock(spec=Job)
+        mock_job.folder_key = "folder-key"
+
+        mock_client = MagicMock()
+        mock_client.processes.invoke_async = AsyncMock(return_value=mock_job)
+        mock_uipath_class.return_value = mock_client
+
+        mock_interrupt.return_value = {"result": "processed"}
+
+        tool = create_process_tool(process_resource_with_inputs)
+        await tool.ainvoke({"name": "test-data", "count": 42})
+
+        call_kwargs = mock_client.processes.invoke_async.call_args[1]
+        assert call_kwargs["input_arguments"] == {"name": "test-data", "count": 42}
+        assert call_kwargs["name"] == "DataProcessor"
+        assert call_kwargs["folder_path"] == "/Shared/DataFolder"
+
+    @pytest.mark.asyncio
+    @patch("uipath_langchain.agent.tools.process_tool.interrupt")
+    @patch("uipath_langchain.agent.tools.process_tool.UiPath")
+    async def test_invoke_returns_interrupt_value(
+        self, mock_uipath_class, mock_interrupt, process_resource
+    ):
+        """Test that the tool returns the value from interrupt()."""
+        mock_job = MagicMock(spec=Job)
+        mock_job.folder_key = "folder-key"
+
+        mock_client = MagicMock()
+        mock_client.processes.invoke_async = AsyncMock(return_value=mock_job)
+        mock_uipath_class.return_value = mock_client
+
+        mock_interrupt.return_value = {"output_arg": "value123"}
+
+        tool = create_process_tool(process_resource)
+        result = await tool.ainvoke({})
+
+        assert result == {"output_arg": "value123"}
+
+
+class TestProcessToolSpanContext:
+    """Test that _span_context is properly wired for tracing."""
+
+    @pytest.mark.asyncio
+    @patch("uipath_langchain.agent.tools.process_tool.interrupt")
+    @patch("uipath_langchain.agent.tools.process_tool.UiPath")
+    async def test_span_context_parent_span_id_passed_to_invoke(
+        self, mock_uipath_class, mock_interrupt, process_resource
+    ):
+        """Test that parent_span_id from _span_context is forwarded to invoke_async."""
+        mock_job = MagicMock(spec=Job)
+        mock_job.folder_key = "folder-key"
+
+        mock_client = MagicMock()
+        mock_client.processes.invoke_async = AsyncMock(return_value=mock_job)
+        mock_uipath_class.return_value = mock_client
+
+        mock_interrupt.return_value = {}
+
+        tool = create_process_tool(process_resource)
+        assert tool.metadata is not None
+
+        # Simulate tracing setting parent_span_id via the shared _span_context
+        tool.metadata["_span_context"]["parent_span_id"] = "span-abc-123"
+
+        await tool.ainvoke({})
+
+        call_kwargs = mock_client.processes.invoke_async.call_args[1]
+        assert call_kwargs["parent_span_id"] == "span-abc-123"
+
+    @pytest.mark.asyncio
+    @patch("uipath_langchain.agent.tools.process_tool.interrupt")
+    @patch("uipath_langchain.agent.tools.process_tool.UiPath")
+    async def test_span_context_consumed_after_invoke(
+        self, mock_uipath_class, mock_interrupt, process_resource
+    ):
+        """Test that parent_span_id is popped (consumed) from _span_context after use."""
+        mock_job = MagicMock(spec=Job)
+        mock_job.folder_key = "folder-key"
+
+        mock_client = MagicMock()
+        mock_client.processes.invoke_async = AsyncMock(return_value=mock_job)
+        mock_uipath_class.return_value = mock_client
+
+        mock_interrupt.return_value = {}
+
+        tool = create_process_tool(process_resource)
+        assert tool.metadata is not None
+        tool.metadata["_span_context"]["parent_span_id"] = "span-xyz"
+
+        await tool.ainvoke({})
+
+        # parent_span_id should be consumed (popped) after invocation
+        assert "parent_span_id" not in tool.metadata["_span_context"]
+
+    @pytest.mark.asyncio
+    @patch("uipath_langchain.agent.tools.process_tool.interrupt")
+    @patch("uipath_langchain.agent.tools.process_tool.UiPath")
+    async def test_span_context_defaults_to_none_when_empty(
+        self, mock_uipath_class, mock_interrupt, process_resource
+    ):
+        """Test that parent_span_id defaults to None when _span_context is empty."""
+        mock_job = MagicMock(spec=Job)
+        mock_job.folder_key = "folder-key"
+
+        mock_client = MagicMock()
+        mock_client.processes.invoke_async = AsyncMock(return_value=mock_job)
+        mock_uipath_class.return_value = mock_client
+
+        mock_interrupt.return_value = {}
+
+        tool = create_process_tool(process_resource)
+        # Don't set any parent_span_id
+        await tool.ainvoke({})
+
+        call_kwargs = mock_client.processes.invoke_async.call_args[1]
+        assert call_kwargs["parent_span_id"] is None