feat: use WaitJobRaw for process_tool (#635)

Author: GabrielVasilescu04

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.7.14"
+version = "0.7.15"
 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

@@ -1,5 +1,6 @@
 """Process tool creation for UiPath process execution."""
 
+import json
 from typing import Any
 
 from langchain.tools import BaseTool
@@ -8,7 +9,8 @@
 from uipath.agent.models.agent import AgentProcessToolResourceConfig, AgentToolType
 from uipath.eval.mocks import mockable
 from uipath.platform import UiPath
-from uipath.platform.common import WaitJob
+from uipath.platform.common import WaitJobRaw
+from uipath.platform.orchestrator import JobState
 
 from uipath_langchain.agent.react.job_attachments import get_job_attachments
 from uipath_langchain.agent.react.jsonschema_pydantic_converter import create_model
@@ -75,9 +77,22 @@ async def start_job():
                     )
                     _bts_context[bts_key] = str(job.key)
 
-                return WaitJob(job=job, process_folder_key=job.folder_key)
+                return WaitJobRaw(job=job, process_folder_key=job.folder_key)
 
-            return await start_job()
+            job = await start_job()
+
+            if (job.state or "").lower() == JobState.FAULTED:
+                error_info = str(job.info or "Unknown error")
+                return f"{error_info}"
+
+            client = UiPath()
+            output_str = await client.jobs.extract_output_async(job)
+            if output_str:
+                try:
+                    return json.loads(output_str)
+                except (json.JSONDecodeError, TypeError):
+                    return output_str
+            return output_str
 
         return await invoke_process(**kwargs)
 

tests/agent/tools/test_process_tool.py

@@ -124,11 +124,15 @@ async def test_invoke_calls_processes_invoke_async(
         mock_job.key = "job-key-123"
         mock_job.folder_key = "folder-key-123"
 
+        mock_resumed_job = MagicMock(spec=Job)
+        mock_resumed_job.state = "successful"
+
         mock_client = MagicMock()
         mock_client.processes.invoke_async = AsyncMock(return_value=mock_job)
+        mock_client.jobs.extract_output_async = AsyncMock(return_value=None)
         mock_uipath_class.return_value = mock_client
 
-        mock_interrupt.return_value = {"output": "result"}
+        mock_interrupt.return_value = mock_resumed_job
 
         tool = create_process_tool(process_resource)
         await tool.ainvoke({})
@@ -148,16 +152,20 @@ 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 WaitJob."""
+        """Test that after invoking, the tool interrupts with WaitJobRaw."""
         mock_job = MagicMock(spec=Job)
         mock_job.key = "job-key-456"
         mock_job.folder_key = "folder-key-456"
 
+        mock_resumed_job = MagicMock(spec=Job)
+        mock_resumed_job.state = "successful"
+
         mock_client = MagicMock()
         mock_client.processes.invoke_async = AsyncMock(return_value=mock_job)
+        mock_client.jobs.extract_output_async = AsyncMock(return_value=None)
         mock_uipath_class.return_value = mock_client
 
-        mock_interrupt.return_value = {"output": "done"}
+        mock_interrupt.return_value = mock_resumed_job
 
         tool = create_process_tool(process_resource)
         await tool.ainvoke({})
@@ -179,11 +187,15 @@ async def test_invoke_passes_input_arguments(
         mock_job.key = "job-key"
         mock_job.folder_key = "folder-key"
 
+        mock_resumed_job = MagicMock(spec=Job)
+        mock_resumed_job.state = "successful"
+
         mock_client = MagicMock()
         mock_client.processes.invoke_async = AsyncMock(return_value=mock_job)
+        mock_client.jobs.extract_output_async = AsyncMock(return_value=None)
         mock_uipath_class.return_value = mock_client
 
-        mock_interrupt.return_value = {"result": "processed"}
+        mock_interrupt.return_value = mock_resumed_job
 
         tool = create_process_tool(process_resource_with_inputs)
         await tool.ainvoke({"name": "test-data", "count": 42})
@@ -196,25 +208,59 @@ async def test_invoke_passes_input_arguments(
     @pytest.mark.asyncio
     @patch("uipath_langchain.agent.tools.durable_interrupt.decorator.interrupt")
     @patch("uipath_langchain.agent.tools.process_tool.UiPath")
-    async def test_invoke_returns_interrupt_value(
+    async def test_invoke_returns_output_from_extract(
         self, mock_uipath_class, mock_interrupt, process_resource
     ):
-        """Test that the tool returns the value from interrupt()."""
+        """Test that the tool returns the extracted job output on success."""
         mock_job = MagicMock(spec=Job)
         mock_job.key = "job-key"
         mock_job.folder_key = "folder-key"
 
+        mock_resumed_job = MagicMock(spec=Job)
+        mock_resumed_job.state = "successful"
+
         mock_client = MagicMock()
         mock_client.processes.invoke_async = AsyncMock(return_value=mock_job)
+        mock_client.jobs.extract_output_async = AsyncMock(
+            return_value='{"output_arg": "value123"}'
+        )
         mock_uipath_class.return_value = mock_client
 
-        mock_interrupt.return_value = {"output_arg": "value123"}
+        mock_interrupt.return_value = mock_resumed_job
 
         tool = create_process_tool(process_resource)
         result = await tool.ainvoke({})
 
         assert result == {"output_arg": "value123"}
 
+    @pytest.mark.asyncio
+    @patch("uipath_langchain.agent.tools.durable_interrupt.decorator.interrupt")
+    @patch("uipath_langchain.agent.tools.process_tool.UiPath")
+    async def test_invoke_returns_error_message_on_faulted_job(
+        self, mock_uipath_class, mock_interrupt, process_resource
+    ):
+        """Test that the tool returns an error message string when the job is faulted."""
+        mock_job = MagicMock(spec=Job)
+        mock_job.key = "job-key"
+        mock_job.folder_key = "folder-key"
+
+        mock_resumed_job = MagicMock(spec=Job)
+        mock_resumed_job.state = "faulted"
+        mock_resumed_job.job_error = None
+        mock_resumed_job.info = "Something went wrong in the workflow"
+
+        mock_client = MagicMock()
+        mock_client.processes.invoke_async = AsyncMock(return_value=mock_job)
+        mock_uipath_class.return_value = mock_client
+
+        mock_interrupt.return_value = mock_resumed_job
+
+        tool = create_process_tool(process_resource)
+        result = await tool.ainvoke({})
+
+        assert isinstance(result, str)
+        assert "Something went wrong in the workflow" in result
+
 
 class TestProcessToolSpanContext:
     """Test that _span_context is properly wired for tracing."""
@@ -230,11 +276,15 @@ async def test_span_context_parent_span_id_passed_to_invoke(
         mock_job.key = "job-key"
         mock_job.folder_key = "folder-key"
 
+        mock_resumed_job = MagicMock(spec=Job)
+        mock_resumed_job.state = "successful"
+
         mock_client = MagicMock()
         mock_client.processes.invoke_async = AsyncMock(return_value=mock_job)
+        mock_client.jobs.extract_output_async = AsyncMock(return_value=None)
         mock_uipath_class.return_value = mock_client
 
-        mock_interrupt.return_value = {}
+        mock_interrupt.return_value = mock_resumed_job
 
         tool = create_process_tool(process_resource)
         assert tool.metadata is not None
@@ -258,11 +308,15 @@ async def test_span_context_consumed_after_invoke(
         mock_job.key = "job-key"
         mock_job.folder_key = "folder-key"
 
+        mock_resumed_job = MagicMock(spec=Job)
+        mock_resumed_job.state = "successful"
+
         mock_client = MagicMock()
         mock_client.processes.invoke_async = AsyncMock(return_value=mock_job)
+        mock_client.jobs.extract_output_async = AsyncMock(return_value=None)
         mock_uipath_class.return_value = mock_client
 
-        mock_interrupt.return_value = {}
+        mock_interrupt.return_value = mock_resumed_job
 
         tool = create_process_tool(process_resource)
         assert tool.metadata is not None
@@ -284,11 +338,15 @@ async def test_span_context_defaults_to_none_when_empty(
         mock_job.key = "job-key"
         mock_job.folder_key = "folder-key"
 
+        mock_resumed_job = MagicMock(spec=Job)
+        mock_resumed_job.state = "successful"
+
         mock_client = MagicMock()
         mock_client.processes.invoke_async = AsyncMock(return_value=mock_job)
+        mock_client.jobs.extract_output_async = AsyncMock(return_value=None)
         mock_uipath_class.return_value = mock_client
 
-        mock_interrupt.return_value = {}
+        mock_interrupt.return_value = mock_resumed_job
 
         tool = create_process_tool(process_resource)
         # Don't set any parent_span_id