fix: stabilize tool safety CI checks

Author: Spock12138

scripts/tool_safety_check.py

@@ -12,24 +12,18 @@
 import sys
 from pathlib import Path
 
-try:
-    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_file
-    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[1]
-    sys.path.insert(0, str(repo_root))
-    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_file
-    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[1]
+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 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_file
+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
 
 
 def parse_args(argv: list[str] | None = None) -> argparse.Namespace:

trpc_agent_sdk/tools/safety/_audit.py

@@ -23,16 +23,17 @@ 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,
@@ -53,7 +54,6 @@ 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,11 +73,13 @@ 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,

trpc_agent_sdk/tools/safety/_cli_helpers.py

@@ -35,7 +35,8 @@ def load_samples(path: str | Path) -> list[dict[str, Any]]:
     sample_path = Path(path)
     loaded = yaml.safe_load(sample_path.read_text(encoding="utf-8"))
     if not isinstance(loaded, list):
-        raise ValueError("samples YAML must contain a list of sample mappings.")
+        raise ValueError(
+            "samples YAML must contain a list of sample mappings.")
 
     samples: list[dict[str, Any]] = []
     seen_ids: set[str] = set()
@@ -44,23 +45,32 @@ def load_samples(path: str | Path) -> list[dict[str, Any]]:
             raise ValueError(f"sample at index {index} must be a mapping.")
         sample_id = _string_field(item, "id")
         if not sample_id:
-            raise ValueError(f"sample at index {index} must have a non-empty id.")
+            raise ValueError(
+                f"sample at index {index} must have a non-empty id.")
         if sample_id in seen_ids:
             raise ValueError(f"duplicate sample id: {sample_id}")
         seen_ids.add(sample_id)
 
         if not _string_field(item, "description"):
-            raise ValueError(f"sample {sample_id} must have a non-empty description.")
+            raise ValueError(
+                f"sample {sample_id} must have a non-empty description.")
         if not _string_field(item, "language"):
-            raise ValueError(f"sample {sample_id} must have a non-empty language.")
-        if not (_string_field(item, "content") or _string_field(item, "command")):
-            raise ValueError(f"sample {sample_id} must define content or command.")
+            raise ValueError(
+                f"sample {sample_id} must have a non-empty language.")
+        if not (_string_field(item, "content")
+                or _string_field(item, "command")):
+            raise ValueError(
+                f"sample {sample_id} must define content or command.")
         if not _string_field(item, "expected_decision"):
-            raise ValueError(f"sample {sample_id} must define expected_decision.")
+            raise ValueError(
+                f"sample {sample_id} must define expected_decision.")
 
         expected_rules = item.get("expected_rules", [])
-        if not isinstance(expected_rules, list) or not all(isinstance(rule, str) for rule in expected_rules):
-            raise ValueError(f"sample {sample_id} expected_rules must be a list of strings.")
+        if not isinstance(expected_rules, list) or not all(
+                isinstance(rule, str) for rule in expected_rules):
+            raise ValueError(
+                f"sample {sample_id} expected_rules must be a list of strings."
+            )
 
         samples.append(dict(item))
     return samples
@@ -77,11 +87,11 @@ def load_policy(path: str | Path | None) -> SafetyPolicy:
 
 
 def scan_samples(
-    samples: list[dict[str, Any]],
-    policy: SafetyPolicy,
-    *,
-    generated_at: str = FIXTURE_GENERATED_AT,
-    stable_elapsed_ms: float | None = None,
+        samples: list[dict[str, Any]],
+        policy: SafetyPolicy,
+        *,
+        generated_at: str = FIXTURE_GENERATED_AT,
+        stable_elapsed_ms: float | None = None,
 ) -> tuple[dict[str, Any], list[SafetyAuditEvent], list[str]]:
     """Scan sample mappings and return aggregate report, audit events, and mismatches."""
 
@@ -93,10 +103,14 @@ def scan_samples(
     for sample in samples:
         report = scanner.scan(build_target_from_sample(sample))
         if stable_elapsed_ms is not None:
-            report = report.model_copy(update={"elapsed_ms": stable_elapsed_ms})
+            report = report.model_copy(
+                update={"elapsed_ms": stable_elapsed_ms})
         result = build_sample_result(sample, report)
         results.append(result)
-        audit_events.append(build_audit_event(result["sample_id"], report, generated_at=generated_at))
+        audit_events.append(
+            build_audit_event(result["sample_id"],
+                              report,
+                              generated_at=generated_at))
         if not result["match"]:
             mismatches.append(_mismatch_message(result))
 
@@ -108,11 +122,11 @@ def scan_samples(
 
 
 def scan_file(
-    path: str | Path,
-    policy: SafetyPolicy,
-    *,
-    language: str = "unknown",
-    generated_at: str = FIXTURE_GENERATED_AT,
+        path: str | Path,
+        policy: SafetyPolicy,
+        *,
+        language: str = "unknown",
+        generated_at: str = FIXTURE_GENERATED_AT,
 ) -> tuple[dict[str, Any], list[SafetyAuditEvent]]:
     """Scan one local file and return an aggregate report with one result."""
 
@@ -122,7 +136,10 @@ def scan_file(
         content=content,
         language=parse_language(language),
         tool_name=TOOL_NAME,
-        tool_metadata={"source": "file", "file_name": file_path.name},
+        tool_metadata={
+            "source": "file",
+            "file_name": file_path.name
+        },
     )
     report = SafetyScanner(policy).scan(target)
     sample = {
@@ -132,7 +149,9 @@ def scan_file(
         "expected_rules": [finding.rule_id for finding in report.findings],
     }
     result = build_sample_result(sample, report)
-    event = build_audit_event(result["sample_id"], report, generated_at=generated_at)
+    event = build_audit_event(result["sample_id"],
+                              report,
+                              generated_at=generated_at)
     aggregate = build_aggregate_report(
         policy_name=policy.name,
         results=[result],
@@ -147,14 +166,18 @@ def build_target_from_sample(sample: dict[str, Any]) -> ScanTarget:
     env_keys = sample.get("env_keys", [])
     if env_keys is None:
         env_keys = []
-    if not isinstance(env_keys, list) or not all(isinstance(key, str) for key in env_keys):
-        raise ValueError(f"sample {sample.get('id', '<unknown>')} env_keys must be a list of strings.")
+    if not isinstance(env_keys, list) or not all(
+            isinstance(key, str) for key in env_keys):
+        raise ValueError(
+            f"sample {sample.get('id', '<unknown>')} env_keys must be a list of strings."
+        )
 
     return ScanTarget(
         content=_string_field(sample, "content"),
         command=_string_field(sample, "command"),
         language=parse_language(_string_field(sample, "language")),
-        env={key: "" for key in env_keys},
+        env={key: ""
+             for key in env_keys},
         tool_name=TOOL_NAME,
         tool_metadata={"sample_id": _string_field(sample, "id")},
     )
@@ -173,41 +196,53 @@ def parse_language(value: str) -> ScriptLanguage:
     return ScriptLanguage.UNKNOWN
 
 
-def build_sample_result(sample: dict[str, Any], report: SafetyReport) -> dict[str, Any]:
+def build_sample_result(sample: dict[str, Any],
+                        report: SafetyReport) -> dict[str, Any]:
     """Build the report entry for one sample without including full source text."""
 
     expected_decision = _string_field(sample, "expected_decision")
     expected_rules = list(sample.get("expected_rules") or [])
     actual_rules = [finding.rule_id for finding in report.findings]
     return {
-        "sample_id": _string_field(sample, "id"),
-        "description": _string_field(sample, "description"),
-        "expected_decision": expected_decision,
-        "expected_rules": expected_rules,
-        "match": expected_decision == report.decision.value and _is_subset(expected_rules, actual_rules),
-        "report": report.model_dump(mode="json"),
+        "sample_id":
+        _string_field(sample, "id"),
+        "description":
+        _string_field(sample, "description"),
+        "expected_decision":
+        expected_decision,
+        "expected_rules":
+        expected_rules,
+        "match":
+        expected_decision == report.decision.value
+        and _is_subset(expected_rules, actual_rules),
+        "report":
+        report.model_dump(mode="json"),
     }
 
 
 def build_aggregate_report(
-    *,
-    policy_name: str,
-    results: list[dict[str, Any]],
-    generated_at: str = FIXTURE_GENERATED_AT,
+        *,
+        policy_name: str,
+        results: list[dict[str, Any]],
+        generated_at: str = FIXTURE_GENERATED_AT,
 ) -> dict[str, Any]:
     """Build the stable public JSON report shape."""
 
-    decisions = Counter(str(result["report"]["decision"]) for result in results)
+    decisions = Counter(
+        str(result["report"]["decision"]) for result in results)
     return {
         "policy_name": policy_name,
         "generated_at": generated_at,
         "sample_count": len(results),
-        "decision_summary": {decision: decisions.get(decision, 0) for decision in DECISION_VALUES},
+        "decision_summary":
+        {decision: decisions.get(decision, 0)
+         for decision in DECISION_VALUES},
         "results": results,
     }
 
 
-def build_audit_event(sample_id: str, report: SafetyReport, *, generated_at: str) -> SafetyAuditEvent:
+def build_audit_event(sample_id: str, report: SafetyReport, *,
+                      generated_at: str) -> SafetyAuditEvent:
     """Build a stable audit event for sample/file scans."""
 
     event = build_safety_audit_event(
@@ -223,10 +258,8 @@ def write_json_report(report: dict[str, Any], path: str | Path) -> None:
 
     report_path = Path(path)
     report_path.parent.mkdir(parents=True, exist_ok=True)
-    report_path.write_text(
-        json.dumps(report, ensure_ascii=False, indent=2, sort_keys=False) + "\n",
-        encoding="utf-8",
-    )
+    payload = json.dumps(report, ensure_ascii=False, indent=2, sort_keys=False)
+    report_path.write_text(f"{payload}\n", encoding="utf-8")
 
 
 def write_audit_log(events: list[SafetyAuditEvent], path: str | Path) -> None:
@@ -258,7 +291,9 @@ def _is_subset(expected_rules: list[str], actual_rules: list[str]) -> bool:
 
 def _mismatch_message(result: dict[str, Any]) -> str:
     report = result["report"]
-    actual_rules = [finding["rule_id"] for finding in report.get("findings", [])]
+    actual_rules = [
+        finding["rule_id"] for finding in report.get("findings", [])
+    ]
     return (
         f"{result['sample_id']}: expected decision={result['expected_decision']} "
         f"rules={result['expected_rules']}, got decision={report['decision']} rules={actual_rules}"