From ba82d8d0b976a95b613a990692f7338eee13bfb8 Mon Sep 17 00:00:00 2001 From: antoniociccia Date: Thu, 16 Jul 2026 12:48:41 +0200 Subject: [PATCH] fix(agents): gate tool-span I/O behind DATAPIZZA_TRACE_CLIENT_IO Tool spans exported tool arguments and results unconditionally, so sensitive data handled by tools (file contents, DB rows, PII) leaked to trace exporters by default. Client I/O is already gated behind DATAPIZZA_TRACE_CLIENT_IO in core/clients/client.py; tool spans were simply not covered. Gate the two set_attribute calls in _execute_tool and _a_execute_tool behind the same flag (default off). The console log_panel output is intentionally left unchanged. Adds sync+async regression tests and documents the flag's tool coverage. --- datapizza-ai-core/datapizza/agents/runner.py | 33 +++- .../agents/tests/test_tool_span_privacy.py | 170 ++++++++++++++++++ docs/Guides/Monitoring/tracing.md | 4 + 3 files changed, 199 insertions(+), 8 deletions(-) create mode 100644 datapizza-ai-core/datapizza/agents/tests/test_tool_span_privacy.py diff --git a/datapizza-ai-core/datapizza/agents/runner.py b/datapizza-ai-core/datapizza/agents/runner.py index 70bc9ed4..20c73bd4 100644 --- a/datapizza-ai-core/datapizza/agents/runner.py +++ b/datapizza-ai-core/datapizza/agents/runner.py @@ -1,4 +1,5 @@ import inspect +import os from collections.abc import AsyncGenerator, Generator from dataclasses import dataclass from typing import TYPE_CHECKING, Literal @@ -21,6 +22,16 @@ from .agent import Agent, Plan, StepResult +def _trace_tool_io_enabled() -> bool: + """Whether tool arguments/results may be exported to trace spans. + + Mirrors the client I/O gate so tool spans do not leak arguments and + results (which can carry file contents, DB rows or PII) to trace + exporters unless the operator explicitly opts in. + """ + return os.getenv("DATAPIZZA_TRACE_CLIENT_IO", "false").lower() == "true" + + @dataclass class HandoffRequest: target: str @@ -1090,16 +1101,19 @@ def _execute_tool( self, function_call: FunctionCallBlock, agent: "Agent" ) -> FunctionCallResultBlock: with tool_span(f"Tool {function_call.tool.name}") as current_tool_span: - current_tool_span.set_attribute( - "tool_arguments", str(function_call.arguments) - ) + trace_io = _trace_tool_io_enabled() + if trace_io: + current_tool_span.set_attribute( + "tool_arguments", str(function_call.arguments) + ) result = function_call.tool(**function_call.arguments) if inspect.iscoroutine(result): result = AsyncExecutor.get_instance().run(result) if result: - current_tool_span.set_attribute("tool_result", result) + if trace_io: + current_tool_span.set_attribute("tool_result", result) agent._logger.log_panel( result, title=f"TOOL {function_call.tool.name.upper()} RESULT", @@ -1115,16 +1129,19 @@ async def _a_execute_tool( self, function_call: FunctionCallBlock, agent: "Agent" ) -> FunctionCallResultBlock: with tool_span(f"Tool {function_call.tool.name}") as current_tool_span: - current_tool_span.set_attribute( - "tool_arguments", str(function_call.arguments) - ) + trace_io = _trace_tool_io_enabled() + if trace_io: + current_tool_span.set_attribute( + "tool_arguments", str(function_call.arguments) + ) result = function_call.tool(**function_call.arguments) if inspect.iscoroutine(result): result = await result if result: - current_tool_span.set_attribute("tool_result", result) + if trace_io: + current_tool_span.set_attribute("tool_result", result) agent._logger.log_panel( result, title=f"TOOL {function_call.tool.name.upper()} RESULT", diff --git a/datapizza-ai-core/datapizza/agents/tests/test_tool_span_privacy.py b/datapizza-ai-core/datapizza/agents/tests/test_tool_span_privacy.py new file mode 100644 index 00000000..90202a37 --- /dev/null +++ b/datapizza-ai-core/datapizza/agents/tests/test_tool_span_privacy.py @@ -0,0 +1,170 @@ +import asyncio +from contextlib import contextmanager +from unittest.mock import MagicMock, patch + +from datapizza.agents.agent import Agent +from datapizza.agents.runner import AgentRunner +from datapizza.clients import MockClient +from datapizza.tools import tool +from datapizza.type import FunctionCallBlock + +TRACE_ENV = "DATAPIZZA_TRACE_CLIENT_IO" + + +@tool +def reveal(secret: str) -> str: + return f"SENSITIVE:{secret}" + + +@tool +def empty_tool() -> str: + return "" + + +def _make_call(): + return FunctionCallBlock( + id="call-1", + arguments={"secret": "hunter2"}, + name=reveal.name, + tool=reveal, + ) + + +def _capture_span(): + span = MagicMock() + + @contextmanager + def fake_tool_span(name=None): + yield span + + return span, fake_tool_span + + +def _attribute_keys(span): + return [call.args[0] for call in span.set_attribute.call_args_list] + + +def test_tool_span_omits_io_by_default(monkeypatch): + monkeypatch.delenv(TRACE_ENV, raising=False) + span, fake_tool_span = _capture_span() + agent = Agent(name="t", client=MockClient(), tools=[reveal]) + + with patch("datapizza.agents.runner.tool_span", fake_tool_span): + AgentRunner()._execute_tool(_make_call(), agent) + + keys = _attribute_keys(span) + assert "tool_arguments" not in keys + assert "tool_result" not in keys + + +def test_tool_span_includes_io_when_enabled(monkeypatch): + monkeypatch.setenv(TRACE_ENV, "true") + span, fake_tool_span = _capture_span() + agent = Agent(name="t", client=MockClient(), tools=[reveal]) + + with patch("datapizza.agents.runner.tool_span", fake_tool_span): + AgentRunner()._execute_tool(_make_call(), agent) + + keys = _attribute_keys(span) + assert "tool_arguments" in keys + assert "tool_result" in keys + + +def test_a_tool_span_omits_io_by_default(monkeypatch): + monkeypatch.delenv(TRACE_ENV, raising=False) + span, fake_tool_span = _capture_span() + agent = Agent(name="t", client=MockClient(), tools=[reveal]) + + with patch("datapizza.agents.runner.tool_span", fake_tool_span): + asyncio.run(AgentRunner()._a_execute_tool(_make_call(), agent)) + + keys = _attribute_keys(span) + assert "tool_arguments" not in keys + assert "tool_result" not in keys + + +def test_a_tool_span_includes_io_when_enabled(monkeypatch): + monkeypatch.setenv(TRACE_ENV, "true") + span, fake_tool_span = _capture_span() + agent = Agent(name="t", client=MockClient(), tools=[reveal]) + + with patch("datapizza.agents.runner.tool_span", fake_tool_span): + asyncio.run(AgentRunner()._a_execute_tool(_make_call(), agent)) + + keys = _attribute_keys(span) + assert "tool_arguments" in keys + assert "tool_result" in keys + + +def test_tool_span_includes_io_with_uppercase_env(monkeypatch): + # The docs instruct DATAPIZZA_TRACE_CLIENT_IO=TRUE (uppercase); the gate must honor it. + monkeypatch.setenv(TRACE_ENV, "TRUE") + span, fake_tool_span = _capture_span() + agent = Agent(name="t", client=MockClient(), tools=[reveal]) + + with patch("datapizza.agents.runner.tool_span", fake_tool_span): + AgentRunner()._execute_tool(_make_call(), agent) + + keys = _attribute_keys(span) + assert "tool_arguments" in keys + assert "tool_result" in keys + + +def test_tool_span_omits_io_for_non_true_value(monkeypatch): + # Only the literal true/TRUE enables (parity with the client I/O gate); "1" must not. + monkeypatch.setenv(TRACE_ENV, "1") + span, fake_tool_span = _capture_span() + agent = Agent(name="t", client=MockClient(), tools=[reveal]) + + with patch("datapizza.agents.runner.tool_span", fake_tool_span): + AgentRunner()._execute_tool(_make_call(), agent) + + keys = _attribute_keys(span) + assert "tool_arguments" not in keys + assert "tool_result" not in keys + + +def test_falsy_tool_result_never_sets_result_attribute(monkeypatch): + # Even with tracing enabled, a falsy result must not produce a tool_result attribute. + monkeypatch.setenv(TRACE_ENV, "true") + span, fake_tool_span = _capture_span() + agent = Agent(name="t", client=MockClient(), tools=[empty_tool]) + call = FunctionCallBlock( + id="call-empty", arguments={}, name=empty_tool.name, tool=empty_tool + ) + + with patch("datapizza.agents.runner.tool_span", fake_tool_span): + AgentRunner()._execute_tool(call, agent) + + keys = _attribute_keys(span) + assert "tool_arguments" in keys + assert "tool_result" not in keys + + +def test_return_value_unchanged_regardless_of_flag(monkeypatch): + _, fake_tool_span = _capture_span() + agent = Agent(name="t", client=MockClient(), tools=[reveal]) + + with patch("datapizza.agents.runner.tool_span", fake_tool_span): + monkeypatch.delenv(TRACE_ENV, raising=False) + off = AgentRunner()._execute_tool(_make_call(), agent) + monkeypatch.setenv(TRACE_ENV, "true") + on = AgentRunner()._execute_tool(_make_call(), agent) + + assert off.result == "SENSITIVE:hunter2" + assert on.result == off.result + + +def test_console_log_panel_called_even_when_tracing_disabled(monkeypatch): + # The console channel is intentionally independent of the trace gate: it must + # still fire while the span omits the sensitive attributes. + monkeypatch.delenv(TRACE_ENV, raising=False) + span, fake_tool_span = _capture_span() + agent = Agent(name="t", client=MockClient(), tools=[reveal]) + agent._logger = MagicMock() + + with patch("datapizza.agents.runner.tool_span", fake_tool_span): + AgentRunner()._execute_tool(_make_call(), agent) + + assert agent._logger.log_panel.called + assert "tool_result" not in _attribute_keys(span) diff --git a/docs/Guides/Monitoring/tracing.md b/docs/Guides/Monitoring/tracing.md index b84314aa..37883ebd 100644 --- a/docs/Guides/Monitoring/tracing.md +++ b/docs/Guides/Monitoring/tracing.md @@ -43,6 +43,10 @@ If you want to log the input/output and the memory passed to client invoke you s default is `FALSE` +The same flag also gates tool spans: tool arguments and results (which may contain +file contents, database rows or other sensitive data) are attached to trace spans +only when `DATAPIZZA_TRACE_CLIENT_IO=TRUE`. With the default they are omitted. + ## Manual Span Creation