Standardize tool error capture with ToolErrorData across framework handlers

Signed-off-by: Eric Evans <194135482+ericevans-nv@users.noreply.github.com>

Author: ericevans-nv

packages/nvidia_nat_adk/src/nat/plugins/adk/callback_handler.py

@@ -25,6 +25,7 @@
 from nat.data_models.intermediate_step import IntermediateStepPayload
 from nat.data_models.intermediate_step import IntermediateStepType
 from nat.data_models.intermediate_step import StreamEventData
+from nat.data_models.intermediate_step import ToolErrorData
 from nat.data_models.intermediate_step import TraceMetadata
 from nat.data_models.intermediate_step import UsageInfo
 from nat.data_models.profiler_callback import BaseProfilerCallback
@@ -200,8 +201,29 @@ async def wrapped_tool_use(base_tool_instance, *args, **kwargs) -> Any:
 
                 return result
 
-            except Exception as _e:
-                logger.exception("BaseTool error occured")
+            except Exception as e:
+                logger.error("BaseTool error: %s", e)
+                kwargs_args = (kwargs.get("args", {}) if isinstance(kwargs.get("args"), dict) else {})
+                tool_error: ToolErrorData = ToolErrorData(
+                    content=f"{type(e).__name__}: {e!s}",
+                    error_type=type(e).__name__,
+                    error_message=str(e),
+                )
+                self.step_manager.push_intermediate_step(
+                    IntermediateStepPayload(
+                        event_type=IntermediateStepType.TOOL_END,
+                        span_event_timestamp=time.time(),
+                        framework=LLMFrameworkEnum.ADK,
+                        name=tool_name,
+                        data=StreamEventData(
+                            input={
+                                "args": args, "kwargs": dict(kwargs_args)
+                            },
+                            output=tool_error,
+                        ),
+                        usage_info=UsageInfo(token_usage=TokenUsageBaseModel()),
+                        UUID=step_uuid,
+                    ))
                 raise
 
         return wrapped_tool_use

packages/nvidia_nat_adk/tests/test_adk_callback_handler.py

@@ -23,6 +23,7 @@
 
 from nat.data_models.intermediate_step import IntermediateStepType
 from nat.data_models.intermediate_step import LLMFrameworkEnum
+from nat.data_models.intermediate_step import ToolErrorData
 from nat.data_models.profiler_callback import BaseProfilerCallback
 from nat.plugins.adk.callback_handler import ADKProfilerHandler
 
@@ -219,30 +220,34 @@ async def test_tool_use_monkey_patch_functionality(handler, mock_context):
 
 @pytest.mark.asyncio
 async def test_tool_use_monkey_patch_with_exception(handler, mock_context):
-    """Test tool use monkey patch handles exceptions properly."""
-    # Create a mock tool instance
+    """When a tool raises an exception, TOOL_END event contains ToolErrorData with parsed error details."""
     mock_tool_instance = MagicMock()
-    mock_tool_instance.name = "test_tool"
+    mock_tool_instance.name = "lookup_tool"
 
-    # Create mock original function that raises an exception
-    mock_original_func = AsyncMock(side_effect=Exception("Tool error"))
+    mock_original_func = AsyncMock(side_effect=ValueError("Column 'revenue' not found"))
     handler._original_tool_call = mock_original_func
 
-    # Get the wrapped function
     wrapped_func = handler._tool_use_monkey_patch()
 
-    # Test that exception is re-raised
-    with pytest.raises(Exception, match="Tool error"):
+    with pytest.raises(ValueError, match="Column 'revenue' not found"):
         await wrapped_func(mock_tool_instance, "arg1")
 
-    # Verify original function was called
     mock_original_func.assert_called_once()
 
-    # Verify start event was still pushed
-    assert mock_context.push_intermediate_step.call_count >= 1
+    assert mock_context.push_intermediate_step.call_count == 2
     start_call = mock_context.push_intermediate_step.call_args_list[0][0][0]
     assert start_call.event_type == IntermediateStepType.TOOL_START
 
+    end_call = mock_context.push_intermediate_step.call_args_list[1][0][0]
+    assert end_call.event_type == IntermediateStepType.TOOL_END
+    assert end_call.name == "lookup_tool"
+    assert isinstance(end_call.data.output, ToolErrorData)
+
+    error_data: ToolErrorData = end_call.data.output
+    assert error_data.content == "ValueError: Column 'revenue' not found"
+    assert error_data.error_type == "ValueError"
+    assert error_data.error_message == "Column 'revenue' not found"
+
 
 @pytest.mark.asyncio
 async def test_tool_use_monkey_patch_tool_name_error(handler, mock_context):

packages/nvidia_nat_autogen/src/nat/plugins/autogen/callback_handler.py

@@ -45,6 +45,7 @@
 from nat.data_models.intermediate_step import IntermediateStepPayload
 from nat.data_models.intermediate_step import IntermediateStepType
 from nat.data_models.intermediate_step import StreamEventData
+from nat.data_models.intermediate_step import ToolErrorData
 from nat.data_models.intermediate_step import TraceMetadata
 from nat.data_models.intermediate_step import UsageInfo
 from nat.data_models.profiler_callback import BaseProfilerCallback
@@ -557,15 +558,18 @@ async def wrapped_tool_call(*args: Any, **kwargs: Any) -> Any:
             except Exception:
                 logger.debug("Error getting tool name")
 
-            # Extract tool input
-            tool_input = ""
+            # Extract tool input from args
+            # run_json signature: (self, args: Mapping[str, Any], cancellation_token, call_id=None)
+            # args[0] = self (tool instance)
+            # args[1] = args (the tool arguments as a Mapping)
+            tool_input: dict[str, Any] = {}
             try:
                 if len(args) > 1:
-                    call_data = args[1]
-                    if hasattr(call_data, "kwargs"):
-                        tool_input = str(call_data.kwargs)
-                    elif isinstance(call_data, dict):
-                        tool_input = str(call_data.get("kwargs", {}))
+                    tool_args = args[1]
+                    if isinstance(tool_args, dict):
+                        tool_input = dict(tool_args)
+                    elif hasattr(tool_args, "items"):
+                        tool_input = dict(tool_args)
             except Exception:
                 logger.debug("Error extracting tool input")
 
@@ -590,14 +594,18 @@ async def wrapped_tool_call(*args: Any, **kwargs: Any) -> Any:
                 output = await original_func(*args, **kwargs)
             except Exception as e:
                 logger.error("Tool execution failed: %s", e)
+                tool_error: ToolErrorData = ToolErrorData(
+                    content=f"{type(e).__name__}: {e!s}",
+                    error_type=type(e).__name__,
+                    error_message=str(e),
+                )
                 handler.step_manager.push_intermediate_step(
                     IntermediateStepPayload(
                         event_type=IntermediateStepType.TOOL_END,
                         span_event_timestamp=time.time(),
                         framework=LLMFrameworkEnum.AUTOGEN,
                         name=tool_name,
-                        data=StreamEventData(input=tool_input, output=str(e)),
-                        metadata=TraceMetadata(error=str(e)),
+                        data=StreamEventData(input=tool_input, output=tool_error),
                         usage_info=UsageInfo(token_usage=TokenUsageBaseModel()),
                         UUID=start_uuid,
                     ))

packages/nvidia_nat_autogen/tests/test_callback_handler_autogen.py

@@ -22,6 +22,7 @@
 
 import pytest
 
+from nat.data_models.intermediate_step import ToolErrorData
 from nat.plugins.autogen.callback_handler import AutoGenProfilerHandler
 from nat.plugins.autogen.callback_handler import ClientPatchInfo
 from nat.plugins.autogen.callback_handler import PatchedClients
@@ -561,29 +562,56 @@ async def test_tool_wrapper_handles_dict_input(self, mock_get):
 
     @patch('nat.plugins.autogen.callback_handler.Context.get')
     async def test_tool_wrapper_handles_exception(self, mock_get):
-        """Test tool wrapper handles tool execution errors."""
+        """When a tool raises an exception, TOOL_END event contains ToolErrorData with parsed error details."""
         mock_context = Mock()
         mock_step_manager = Mock()
         mock_context.intermediate_step_manager = mock_step_manager
         mock_get.return_value = mock_context
 
         handler = AutoGenProfilerHandler()
 
-        original_func = AsyncMock(side_effect=ValueError("Tool failed"))
+        original_func = AsyncMock(side_effect=ValueError("Column 'revenue' not found"))
         wrapped = handler._create_tool_wrapper(original_func)
 
         tool = Mock()
-        tool.name = "failing_tool"
+        tool.name = "lookup_tool"
         call_data = Mock()
         call_data.kwargs = {}
 
-        with pytest.raises(ValueError, match="Tool failed"):
+        with pytest.raises(ValueError, match="Column 'revenue' not found"):
             await wrapped(tool, call_data)
 
-        # Should have START and error END
         assert mock_step_manager.push_intermediate_step.call_count == 2
         error_call = mock_step_manager.push_intermediate_step.call_args_list[1][0][0]
-        assert "Tool failed" in error_call.data.output
+        assert error_call.name == "lookup_tool"
+        assert isinstance(error_call.data.output, ToolErrorData)
+
+        error_data: ToolErrorData = error_call.data.output
+        assert error_data.content == "ValueError: Column 'revenue' not found"
+        assert error_data.error_type == "ValueError"
+        assert error_data.error_message == "Column 'revenue' not found"
+
+    @patch('nat.plugins.autogen.callback_handler.Context.get')
+    async def test_tool_wrapper_extracts_input_from_run_json_args(self, mock_get):
+        """Test tool wrapper extracts input from run_json signature: (self, args: Mapping, ...)."""
+        mock_context = Mock()
+        mock_step_manager = Mock()
+        mock_context.intermediate_step_manager = mock_step_manager
+        mock_get.return_value = mock_context
+
+        handler = AutoGenProfilerHandler()
+
+        original_func = AsyncMock(return_value="result")
+        wrapped = handler._create_tool_wrapper(original_func)
+
+        tool = Mock()
+        tool.name = "failing_lookup"
+        tool_args = {"query": "revenue"}
+
+        await wrapped(tool, tool_args, Mock())
+
+        start_event = mock_step_manager.push_intermediate_step.call_args_list[0][0][0]
+        assert start_event.data.input == {"query": "revenue"}
 
 
 class TestIntegration:

packages/nvidia_nat_core/src/nat/data_models/intermediate_step.py

@@ -68,6 +68,14 @@ class IntermediateStepState(StrEnum):
     END = "END"
 
 
+class ToolErrorData(BaseModel):
+    """ToolErrorData is a data model that represents the output field in a TOOL_END event when an error occurs."""
+
+    content: str = Field(description="Full error string, e.g. 'ValueError: Column not found'")
+    error_type: str = Field(description="Exception type, e.g. 'ValueError'")
+    error_message: str = Field(description="Error message without type, e.g. 'Column not found'")
+
+
 class StreamEventData(BaseModel):
     """
     StreamEventData is a data model that represents the data field in an streaming event.

packages/nvidia_nat_core/src/nat/utils/atif_converter.py

@@ -46,6 +46,7 @@
 from nat.data_models.intermediate_step import IntermediateStepCategory
 from nat.data_models.intermediate_step import IntermediateStepState
 from nat.data_models.intermediate_step import IntermediateStepType
+from nat.data_models.intermediate_step import ToolErrorData
 from nat.data_models.intermediate_step import TraceMetadata
 
 logger = logging.getLogger(__name__)
@@ -101,29 +102,6 @@ def _safe_str(value: Any) -> str:
     return str(value)
 
 
-def _extract_tool_error(output: Any) -> dict[str, str] | None:
-    """Extract error metadata from a tool output for ``step.extra["tool_errors"]``."""
-    # TODO: return a model instead of a plain dict once ATIF spec adds error support
-    status: str | None = getattr(output, "status", None) or (output.get("status") if isinstance(output, dict) else None)
-    if status != "error":
-        return None
-    content: str = (getattr(output, "content", None) or (output.get("content") if isinstance(output, dict) else None)
-                    or _safe_str(output))
-    error_type: str = "Unknown"
-    error_message: str = content
-    if ":" in content:
-        candidate: str = content.split(":", 1)[0].strip()
-        if candidate.isidentifier():
-            error_type = candidate
-            error_message = content.split(":", 1)[1].strip()
-    return {
-        "error": content,
-        "error_type": error_type,
-        "error_message": error_message,
-        "status": "error",
-    }
-
-
 def _extract_user_input(value: Any) -> str:
     """Extract the user-facing input text from a workflow start payload.
 
@@ -366,11 +344,15 @@ def _flush_pending() -> None:
                 call_id = f"call_{ist.UUID}"
                 tc = ATIFToolCall(tool_call_id=call_id, function_name=tool_name, arguments=tool_input)
                 obs = ATIFObservationResult(source_call_id=call_id, content=tool_output)
-                tool_error: dict[str, str] | None = _extract_tool_error(raw_output)
 
-                if tool_error is not None:
-                    tool_error["tool"] = tool_name
-                extra: dict[str, Any] | None = ({"tool_errors": [tool_error]} if tool_error else None)
+                tool_error: dict[str, str] | None = None
+                if isinstance(raw_output, ToolErrorData):
+                    tool_error = {
+                        "tool": tool_name,
+                        "error": raw_output.content,
+                        "error_type": raw_output.error_type,
+                        "error_message": raw_output.error_message,
+                    }
 
                 if pending is not None:
                     pending.tool_calls.append(tc)
@@ -379,7 +361,7 @@ def _flush_pending() -> None:
                         pending.extra.setdefault("tool_errors", []).append(tool_error)
                     pending.tool_ancestry.append(_atif_ancestry_from_ist(ist))
                 else:
-                    extra = _atif_step_extra_model_from_ist(ist).model_dump(exclude_none=True)
+                    extra: dict[str, Any] = _atif_step_extra_model_from_ist(ist).model_dump(exclude_none=True)
                     if tool_error:
                         extra.setdefault("tool_errors", []).append(tool_error)
                     atif_steps.append(
@@ -552,9 +534,16 @@ def push(self, ist: IntermediateStep) -> ATIFStep | None:
             call_id = f"call_{ist.UUID}"
             tc = ATIFToolCall(tool_call_id=call_id, function_name=tool_name, arguments=tool_input)
             obs = ATIFObservationResult(source_call_id=call_id, content=tool_output)
-            tool_error: dict[str, str] | None = _extract_tool_error(raw_output)
-            if tool_error is not None:
-                tool_error["tool"] = tool_name
+
+            tool_error: dict[str, str] | None = None
+            if isinstance(raw_output, ToolErrorData):
+                tool_error = {
+                    "tool": tool_name,
+                    "error": raw_output.content,
+                    "error_type": raw_output.error_type,
+                    "error_message": raw_output.error_message,
+                }
+
             if self._pending is not None:
                 self._pending.tool_calls.append(tc)
                 self._pending.observations.append(obs)
@@ -563,7 +552,7 @@ def push(self, ist: IntermediateStep) -> ATIFStep | None:
                 self._pending.tool_ancestry.append(_atif_ancestry_from_ist(ist))
                 return None
 
-            extra = _atif_step_extra_model_from_ist(ist).model_dump(exclude_none=True)
+            extra: dict[str, Any] = _atif_step_extra_model_from_ist(ist).model_dump(exclude_none=True)
             if tool_error:
                 extra.setdefault("tool_errors", []).append(tool_error)
             orphan_step = ATIFStep(