feat(escalation): [2/2] send the escalation span id with the Action Center task [ACTN-11909]

src/uipath_langchain/agent/tools/escalation_tool.py

@@ -225,6 +225,7 @@ async def create_task_for_channel(
     data: dict[str, Any],
     recipient: TaskRecipient | None,
     folder_path: str | None,
+    task_source_metadata: dict[str, Any] | None = None,
 ) -> Task:
     """Create the human task backing an escalation channel."""
     if isinstance(channel, AgentQuickFormEscalationChannel):
@@ -241,6 +242,7 @@ async def create_task_for_channel(
             labels=channel.labels,
             is_actionable_message_enabled=channel.properties.is_actionable_message_enabled,
             actionable_message_metadata=channel.properties.actionable_message_meta_data,
+            task_source_metadata=task_source_metadata,
         )
     return await client.tasks.create_async(
         title=title,
@@ -252,6 +254,7 @@ async def create_task_for_channel(
         labels=channel.labels,
         is_actionable_message_enabled=channel.properties.is_actionable_message_enabled,
         actionable_message_metadata=channel.properties.actionable_message_meta_data,
+        task_source_metadata=task_source_metadata,
     )
 
 
@@ -361,13 +364,17 @@ async def escalate(**_tool_kwargs: Any):
             @durable_interrupt
             async def create_escalation_task():
                 client = UiPath()
+                parent_span_id = _span_context.get("parent_span_id")
                 created_task = await create_task_for_channel(
                     client,
                     channel,
                     title=task_title,
                     data=serialized_data,
                     recipient=recipient,
                     folder_path=folder_path,
+                    task_source_metadata=(
+                        {"ParentSpanId": parent_span_id} if parent_span_id else None
+                    ),
                 )
 
                 if created_task.id is not None:

tests/agent/tools/test_escalation_tool.py

@@ -262,6 +262,36 @@ async def test_escalation_tool_with_string_task_title(
         # Verify create_async was called with the static title
         create_call = mock_client.tasks.create_async.call_args
         assert create_call[1]["title"] == "Static Task Title"
+        # No escalation span yet, so no ParentSpanId.
+        assert create_call[1]["task_source_metadata"] is None
+
+    @pytest.mark.asyncio
+    @patch("uipath_langchain.agent.tools.escalation_tool.UiPath")
+    @patch("uipath_langchain._utils.durable_interrupt.decorator.interrupt")
+    async def test_escalation_task_carries_the_escalation_span_id(
+        self, mock_interrupt, mock_uipath_class, escalation_resource
+    ):
+        mock_client = MagicMock()
+        mock_client.tasks.create_async = AsyncMock(return_value=_make_mock_task())
+        mock_uipath_class.return_value = mock_client
+
+        mock_result = MagicMock()
+        mock_result.action = None
+        mock_result.data = {}
+        mock_result.is_deleted = False
+        mock_interrupt.return_value = mock_result
+
+        tool = create_escalation_tool(escalation_resource)
+        # The tool instrumentor sets this to the escalation span when the tool starts.
+        assert tool.metadata is not None
+        tool.metadata["_span_context"]["parent_span_id"] = "a3c27222815e80ca"
+        call = ToolCall(args={}, id="test-call", name=tool.name)
+        await tool.awrapper(tool, call, {})  # type: ignore[attr-defined]
+
+        create_call = mock_client.tasks.create_async.call_args
+        assert create_call[1]["task_source_metadata"] == {
+            "ParentSpanId": "a3c27222815e80ca"
+        }
 
     @pytest.mark.asyncio
     @patch("uipath_langchain.agent.tools.escalation_tool.UiPath")
@@ -1264,6 +1294,33 @@ async def test_dispatches_to_create_quickform_async(
         assert "app_name" not in kwargs
         assert "app_folder_path" not in kwargs
 
+    @pytest.mark.asyncio
+    @patch("uipath_langchain.agent.tools.escalation_tool.UiPath")
+    @patch("uipath_langchain._utils.durable_interrupt.decorator.interrupt")
+    async def test_quick_form_task_carries_the_escalation_span_id(
+        self, mock_interrupt, mock_uipath_class, quick_form_resource
+    ):
+        mock_client = MagicMock()
+        mock_client.tasks.create_quickform_async = AsyncMock(
+            return_value=_make_mock_task(id=779)
+        )
+        mock_uipath_class.return_value = mock_client
+
+        mock_result = MagicMock()
+        mock_result.action = "approve"
+        mock_result.data = {}
+        mock_result.is_deleted = False
+        mock_interrupt.return_value = mock_result
+
+        tool = create_escalation_tool(quick_form_resource)
+        assert tool.metadata is not None
+        tool.metadata["_span_context"]["parent_span_id"] = "a3c27222815e80ca"
+        call = ToolCall(args={}, id="test-call", name=tool.name)
+        await tool.awrapper(tool, call, {})  # type: ignore[attr-defined]
+
+        kwargs = mock_client.tasks.create_quickform_async.call_args[1]
+        assert kwargs["task_source_metadata"] == {"ParentSpanId": "a3c27222815e80ca"}
+
     @pytest.mark.asyncio
     @patch("uipath_langchain.agent.tools.escalation_tool.UiPath")
     @patch("uipath_langchain._utils.durable_interrupt.decorator.interrupt")