|
| 1 | +# Copyright (C) 2026 Tencent. All rights reserved. |
| 2 | +# tRPC-Agent-Python is licensed under Apache-2.0. |
| 3 | +from __future__ import annotations |
| 4 | + |
| 5 | +import json |
| 6 | + |
| 7 | +from trpc_agent_sdk.tools.safety._audit import AuditRecord |
| 8 | +from trpc_agent_sdk.tools.safety._audit import build_audit_record |
| 9 | +from trpc_agent_sdk.tools.safety._audit import emit_otel |
| 10 | +from trpc_agent_sdk.tools.safety._audit import record_safety_decision |
| 11 | +from trpc_agent_sdk.tools.safety._audit import write_audit |
| 12 | +from trpc_agent_sdk.tools.safety._types import Decision |
| 13 | +from trpc_agent_sdk.tools.safety._types import Finding |
| 14 | +from trpc_agent_sdk.tools.safety._types import RiskLevel |
| 15 | +from trpc_agent_sdk.tools.safety._types import SafetyReport |
| 16 | + |
| 17 | + |
| 18 | +def _report(decision=Decision.DENY, risk=RiskLevel.HIGH, findings=None, duration=7): |
| 19 | + return SafetyReport( |
| 20 | + decision=decision, |
| 21 | + risk_level=risk, |
| 22 | + findings=findings or [ |
| 23 | + Finding( |
| 24 | + rule_id="tool-fs-recursive-delete", |
| 25 | + risk_level=RiskLevel.HIGH, |
| 26 | + rule_decision=Decision.DENY, |
| 27 | + evidence="rm -rf /", |
| 28 | + recommendation="refuse", |
| 29 | + ) |
| 30 | + ], |
| 31 | + recommendation="blocked", |
| 32 | + scan_duration_ms=duration, |
| 33 | + sanitized=False, |
| 34 | + ) |
| 35 | + |
| 36 | + |
| 37 | +def test_build_audit_record_has_all_issue_required_fields(): |
| 38 | + rec = build_audit_record( |
| 39 | + _report(), tool_name="weather_tool", language="python", intercepted=True |
| 40 | + ) |
| 41 | + # Issue #90 mandatory audit fields: |
| 42 | + assert rec.tool_name == "weather_tool" # tool name |
| 43 | + assert rec.decision == "DENY" # decision |
| 44 | + assert rec.risk_level == "HIGH" # risk level |
| 45 | + assert rec.rule_ids == ["tool-fs-recursive-delete"] # rule id |
| 46 | + assert rec.scan_duration_ms == 7 # duration |
| 47 | + assert rec.sanitized is False # sanitized flag |
| 48 | + assert rec.intercepted is True # intercepted flag |
| 49 | + assert rec.timestamp # non-empty ISO timestamp |
| 50 | + |
| 51 | + |
| 52 | +def test_write_audit_appends_jsonl_lines(tmp_path): |
| 53 | + path = tmp_path / "audit.jsonl" |
| 54 | + rec = AuditRecord( |
| 55 | + timestamp="2026-07-11T00:00:00+00:00", |
| 56 | + tool_name="t", |
| 57 | + language="python", |
| 58 | + decision="DENY", |
| 59 | + risk_level="HIGH", |
| 60 | + rule_ids=["tool-fs-recursive-delete"], |
| 61 | + scan_duration_ms=1, |
| 62 | + sanitized=False, |
| 63 | + intercepted=True, |
| 64 | + ) |
| 65 | + write_audit(rec, str(path)) |
| 66 | + write_audit(rec, str(path)) |
| 67 | + lines = path.read_text(encoding="utf-8").strip().split("\n") |
| 68 | + assert len(lines) == 2 |
| 69 | + obj = json.loads(lines[0]) |
| 70 | + assert obj["tool_name"] == "t" |
| 71 | + assert obj["decision"] == "DENY" |
| 72 | + assert obj["intercepted"] is True |
| 73 | + |
| 74 | + |
| 75 | +def test_record_safety_decision_writes_and_returns(tmp_path): |
| 76 | + path = tmp_path / "a.jsonl" |
| 77 | + rec = record_safety_decision( |
| 78 | + _report(decision=Decision.ALLOW, risk=RiskLevel.NONE, findings=[], duration=2), |
| 79 | + tool_name="safe_tool", |
| 80 | + language="bash", |
| 81 | + intercepted=False, |
| 82 | + audit_path=str(path), |
| 83 | + ) |
| 84 | + assert rec.intercepted is False |
| 85 | + assert rec.decision == "ALLOW" |
| 86 | + assert path.exists() |
| 87 | + line = json.loads(path.read_text(encoding="utf-8").strip()) |
| 88 | + assert line["tool_name"] == "safe_tool" |
| 89 | + |
| 90 | + |
| 91 | +def test_emit_otel_is_safe_without_active_span(): |
| 92 | + # No active OTel span in unit tests -> must be a no-op, never raise. |
| 93 | + emit_otel(_report(), tool_name="t", intercepted=True) |
0 commit comments