fix: align tool safety CI behavior

Author: Spock12138

examples/tool_safety_guard/run_safety_scan.py

@@ -10,24 +10,18 @@
 import sys
 from pathlib import Path
 
-try:
-    from trpc_agent_sdk.tools.safety._cli_helpers import FIXTURE_GENERATED_AT
-    from trpc_agent_sdk.tools.safety._cli_helpers import format_mismatches
-    from trpc_agent_sdk.tools.safety._cli_helpers import load_policy
-    from trpc_agent_sdk.tools.safety._cli_helpers import load_samples
-    from trpc_agent_sdk.tools.safety._cli_helpers import scan_samples
-    from trpc_agent_sdk.tools.safety._cli_helpers import write_audit_log
-    from trpc_agent_sdk.tools.safety._cli_helpers import write_json_report
-except ModuleNotFoundError:  # pragma: no cover - exercised by direct script execution in fresh checkouts.
-    repo_root = Path(__file__).resolve().parents[2]
-    sys.path.insert(0, str(repo_root))
-    from trpc_agent_sdk.tools.safety._cli_helpers import FIXTURE_GENERATED_AT
-    from trpc_agent_sdk.tools.safety._cli_helpers import format_mismatches
-    from trpc_agent_sdk.tools.safety._cli_helpers import load_policy
-    from trpc_agent_sdk.tools.safety._cli_helpers import load_samples
-    from trpc_agent_sdk.tools.safety._cli_helpers import scan_samples
-    from trpc_agent_sdk.tools.safety._cli_helpers import write_audit_log
-    from trpc_agent_sdk.tools.safety._cli_helpers import write_json_report
+REPO_ROOT = Path(__file__).resolve().parents[2]
+repo_root = str(REPO_ROOT)
+sys.path[:] = [path for path in sys.path if path != repo_root]
+sys.path.insert(0, repo_root)
+
+from trpc_agent_sdk.tools.safety._cli_helpers import FIXTURE_GENERATED_AT
+from trpc_agent_sdk.tools.safety._cli_helpers import format_mismatches
+from trpc_agent_sdk.tools.safety._cli_helpers import load_policy
+from trpc_agent_sdk.tools.safety._cli_helpers import load_samples
+from trpc_agent_sdk.tools.safety._cli_helpers import scan_samples
+from trpc_agent_sdk.tools.safety._cli_helpers import write_audit_log
+from trpc_agent_sdk.tools.safety._cli_helpers import write_json_report
 
 EXAMPLE_DIR = Path(__file__).resolve().parent
 POLICY_PATH = EXAMPLE_DIR / "tool_safety_policy.yaml"

tests/tools/safety/test_code_executor_wrapper.py

@@ -8,7 +8,6 @@
 
 import asyncio
 import logging
-from typing import Any
 from unittest.mock import Mock
 from unittest.mock import patch
 
@@ -53,6 +52,12 @@ async def execute_code(
         return self.result
 
 
+def _enable_caplog_logger(name: str) -> None:
+    target_logger = logging.getLogger(name)
+    target_logger.disabled = False
+    target_logger.propagate = True
+
+
 class RecordingAuditLogger:
     def __init__(self):
         self.events: list[SafetyAuditEvent] = []
@@ -452,6 +457,7 @@ def test_fail_closed_blocks_and_logs_without_exception_detail(self, caplog):
             scanner=RaisingScanner(policy),
             audit_logger=audit_logger,
         )
+        _enable_caplog_logger("trpc_agent_sdk.tools.safety._code_executor")
         caplog.set_level(logging.WARNING, logger="trpc_agent_sdk.tools.safety._code_executor")
 
         with patch("trpc_agent_sdk.tools.safety._code_executor.set_safety_span_attributes") as mock_span:
@@ -475,6 +481,7 @@ def test_fail_open_delegates_and_skips_audit_span(self, caplog):
             scanner=RaisingScanner(policy),
             audit_logger=audit_logger,
         )
+        _enable_caplog_logger("trpc_agent_sdk.tools.safety._code_executor")
         caplog.set_level(logging.WARNING, logger="trpc_agent_sdk.tools.safety._code_executor")
 
         with patch("trpc_agent_sdk.tools.safety._code_executor.set_safety_span_attributes") as mock_span:

tests/tools/safety/test_filter_and_audit.py

@@ -15,6 +15,7 @@
 from trpc_agent_sdk.context import new_agent_context
 from trpc_agent_sdk.filter import BaseFilter
 from trpc_agent_sdk.filter import get_tool_filter
+from trpc_agent_sdk.filter import register_tool_filter
 from trpc_agent_sdk.filter import run_filters
 from trpc_agent_sdk.tools._context_var import reset_tool_var
 from trpc_agent_sdk.tools._context_var import set_tool_var
@@ -42,6 +43,17 @@ def emit(self, event: SafetyAuditEvent) -> None:
         self.events.append(event)
 
 
+def _ensure_safety_filter_registered() -> None:
+    if get_tool_filter("tool_safety_guard") is None:
+        register_tool_filter("tool_safety_guard")(ToolSafetyFilter)
+
+
+def _enable_caplog_logger(name: str) -> None:
+    target_logger = logging.getLogger(name)
+    target_logger.disabled = False
+    target_logger.propagate = True
+
+
 class StaticScanner:
     def __init__(
         self,
@@ -190,6 +202,8 @@ def test_review_blocks_by_default_and_can_be_nonblocking(self):
         assert audit_logger.events[0].blocked is False
 
     def test_registry_name_uses_default_policy_instance(self):
+        _ensure_safety_filter_registered()
+
         registered = get_tool_filter("tool_safety_guard")
 
         assert isinstance(registered, ToolSafetyFilter)
@@ -257,6 +271,7 @@ def test_shell_tools_default_to_shell_language(self):
     def test_fail_closed_blocks_and_logs_without_exception_detail(self, caplog):
         policy = SafetyPolicy(fail_closed=True)
         filter_ = ToolSafetyFilter(scanner=RaisingScanner(policy), audit_logger=RecordingAuditLogger())
+        _enable_caplog_logger("trpc_agent_sdk.tools.safety._filter")
         caplog.set_level(logging.WARNING, logger="trpc_agent_sdk.tools.safety._filter")
 
         result, _, calls = _run_filter(filter_, {"command": "echo ok"})
@@ -272,6 +287,7 @@ def test_fail_open_allows_and_logs_without_exception_detail(self, caplog):
         policy = SafetyPolicy(fail_closed=False)
         audit_logger = RecordingAuditLogger()
         filter_ = ToolSafetyFilter(scanner=RaisingScanner(policy), audit_logger=audit_logger)
+        _enable_caplog_logger("trpc_agent_sdk.tools.safety._filter")
         caplog.set_level(logging.WARNING, logger="trpc_agent_sdk.tools.safety._filter")
 
         with patch("trpc_agent_sdk.tools.safety._filter.set_safety_span_attributes") as mock_span:

trpc_agent_sdk/tools/safety/_audit.py

@@ -23,17 +23,16 @@ def _ordered_rule_ids(report: SafetyReport) -> list[str]:
 
 
 def build_safety_audit_event(
-        report: SafetyReport,
-        *,
-        tool_name: str = "",
-        cwd: str = "",
-        function_call_id: str = "",
-        agent_name: str = "",
+    report: SafetyReport,
+    *,
+    tool_name: str = "",
+    cwd: str = "",
+    function_call_id: str = "",
+    agent_name: str = "",
 ) -> SafetyAuditEvent:
     """Build an audit-safe summary event from a safety report."""
 
-    resolved_tool_name = tool_name or str(
-        report.metadata.get("target_tool", "") or "")
+    resolved_tool_name = tool_name or str(report.metadata.get("target_tool", "") or "")
     return SafetyAuditEvent(
         tool_name=resolved_tool_name,
         decision=report.decision,
@@ -54,6 +53,7 @@ def build_safety_audit_event(
 
 class SafetyAuditLogger:
     """Optional JSONL writer for safety audit events."""
+
     def __init__(self, path: str | Path | None = None, enabled: bool = True):
         self.path = Path(path) if path is not None else None
         self.enabled = enabled
@@ -73,13 +73,11 @@ def emit(self, event: SafetyAuditEvent) -> None:
             logger.warning("Failed to write safety audit event: %s", ex)
 
 
-def set_safety_span_attributes(report: SafetyReport, *,
-                               tool_name: str = "") -> None:
+def set_safety_span_attributes(report: SafetyReport, *, tool_name: str = "") -> None:
     """Best-effort write of safety summary fields onto the current OTel span."""
 
     rule_ids = _ordered_rule_ids(report)
-    resolved_tool_name = tool_name or str(
-        report.metadata.get("target_tool", "") or "")
+    resolved_tool_name = tool_name or str(report.metadata.get("target_tool", "") or "")
     attributes: dict[str, str | bool | int | float] = {
         "tool.safety.decision": report.decision.value,
         "tool.safety.risk_level": report.risk_level.value,