test: raise tool safety coverage

Author: zzp1221

tests/tools/safety/test_tool_safety_guard.py

@@ -17,17 +17,28 @@
 from trpc_agent_sdk.tools import FunctionTool
 from trpc_agent_sdk.tools.safety import SafetyDecision
 from trpc_agent_sdk.tools.safety import SafetyGuardedCodeExecutor
+from trpc_agent_sdk.tools.safety import ToolSafetyBlockedError
+from trpc_agent_sdk.tools.safety import ToolSafetyFinding
 from trpc_agent_sdk.tools.safety import SafetyRiskLevel
 from trpc_agent_sdk.tools.safety import ToolSafetyAuditLogger
 from trpc_agent_sdk.tools.safety import ToolSafetyFilter
 from trpc_agent_sdk.tools.safety import ToolSafetyGuard
 from trpc_agent_sdk.tools.safety import ToolSafetyPolicy
+from trpc_agent_sdk.tools.safety import ToolSafetyReport
 from trpc_agent_sdk.tools.safety import ToolSafetyScanRequest
 from trpc_agent_sdk.tools.safety import ToolSafetyScanner
+from trpc_agent_sdk.tools.safety import apply_tool_safety_span_attributes
+from trpc_agent_sdk.tools.safety import extract_script_from_tool_args
+from trpc_agent_sdk.tools.safety import load_tool_safety_policy
 from trpc_agent_sdk.code_executors import CodeBlock
 from trpc_agent_sdk.code_executors import CodeExecutionInput
 from trpc_agent_sdk.code_executors import BaseCodeExecutor
 from trpc_agent_sdk.code_executors import create_code_execution_result
+from trpc_agent_sdk.filter import FilterResult
+from trpc_agent_sdk.tools.safety._guard import command_to_args
+from trpc_agent_sdk.tools.safety._guard import infer_language_from_key
+from trpc_agent_sdk.tools.safety._types import max_risk_level
+from trpc_agent_sdk.tools.safety._types import risk_level_value
 from scripts.tool_safety_check import main as tool_safety_check_main
 
 
@@ -236,6 +247,129 @@ def test_subprocess_output_capture_has_output_size_recommendation():
     assert "max_output_bytes" in finding.recommendation
 
 
+@pytest.mark.parametrize(
+    ("script", "language", "rule_id"),
+    [
+        ("print(", "python", "TSG-PY-SYNTAX"),
+        ("password = 'abcdefghijklmnopqrstuvwxyz'", "python", "TSG-SECRETS-LITERAL"),
+        ("for i in range(100000):\n    pass", "python", "TSG-RESOURCE-LARGE-LOOP"),
+        ("import pip", "python", "TSG-DEPENDENCY-INSTALL-API"),
+        ("from urllib import request", "python", "TSG-NETWORK-IMPORT"),
+        ("import os\nos.remove('/tmp/demo')", "python", "TSG-FILE-DANGEROUS-OP"),
+        ("open('/etc/passwd', 'w')", "python", "TSG-FILE-SYSTEM-WRITE"),
+        ("from pathlib import Path\nPath('/etc/app.conf').write_text('x')", "python", "TSG-FILE-SYSTEM-WRITE"),
+        ("from pathlib import Path\nPath('out.txt').write_text('" + ("x" * 120) + "')", "python", "TSG-RESOURCE-LARGE-WRITE"),
+        ("import time\ntime.sleep(99)", "python", "TSG-RESOURCE-LONG-SLEEP"),
+        ("import os\nos.fork()", "python", "TSG-RESOURCE-PROCESS-FANOUT"),
+        ("import asyncio\nasyncio.gather(" + ", ".join(f't{i}()' for i in range(40)) + ")", "python",
+         "TSG-RESOURCE-PARALLELISM"),
+        (":(){ :|:& };:", "bash", "TSG-RESOURCE-FORK-BOMB"),
+        ("while true; do echo x; done", "bash", "TSG-RESOURCE-INFINITE-LOOP"),
+        ("sleep 99", "bash", "TSG-RESOURCE-LONG-SLEEP"),
+        ("curl https://evil.example.net", "bash", "TSG-NETWORK-COMMAND"),
+        ("systemctl restart demo", "bash", "TSG-SHELL-DANGEROUS-COMMAND"),
+        ("cat ~/.ssh/id_rsa", "bash", "TSG-SECRETS-READ"),
+        ("dd if=/dev/zero of=big.bin bs=1M count=1024", "bash", "TSG-RESOURCE-LARGE-WRITE"),
+        ("echo $API_KEY", "bash", "TSG-SECRETS-SINK"),
+        ("tee /etc/app.conf", "bash", "TSG-FILE-SYSTEM-WRITE"),
+        ("echo x > ~/.ssh/config", "bash", "TSG-FILE-DENIED-PATH"),
+    ],
+)
+def test_scanner_additional_branches(script, language, rule_id):
+    tuned_policy = policy(max_loop_iterations=10, max_sleep_seconds=10, max_literal_write_bytes=10, max_parallel_tasks=5)
+    report = ToolSafetyScanner(tuned_policy).scan(ToolSafetyScanRequest(script=script, language=language))
+
+    assert any(finding.rule_id == rule_id for finding in report.findings)
+
+
+def test_scanner_unknown_language_and_env_secret_branches():
+    report = ToolSafetyScanner(policy()).scan(
+        ToolSafetyScanRequest(script="noop", language="ruby", env={"API_KEY": "secret-value"}))
+
+    assert report.decision == SafetyDecision.NEEDS_HUMAN_REVIEW
+    assert any(finding.rule_id == "TSG-LANG-UNKNOWN" for finding in report.findings)
+    assert any(finding.rule_id == "TSG-SECRETS-ENV" for finding in report.findings)
+    assert report.redacted is True
+
+
+def test_scanner_command_arg_prefixes_and_duration_helpers():
+    scripts = [
+        ("", "bash", ["sudo curl https://evil.example.net"], "TSG-NETWORK-COMMAND"),
+        ("", "bash", ["uv pip install requests"], "TSG-DEPENDENCY-INSTALL"),
+        ("", "bash", ["poetry add requests"], "TSG-DEPENDENCY-INSTALL"),
+        ("", "bash", ["parallel -j0 echo {}"], "TSG-RESOURCE-PARALLELISM"),
+        ("", "bash", ["xargs --max-procs=999"], "TSG-RESOURCE-PARALLELISM"),
+        ("", "bash", ["timeout -k 1s 10m python job.py"], "TSG-RESOURCE-LONG-TIMEOUT"),
+    ]
+    scanner = ToolSafetyScanner(policy(max_timeout_seconds=30, max_parallel_tasks=8))
+
+    for script, language, command_args, rule_id in scripts:
+        report = scanner.scan(ToolSafetyScanRequest(script=script, language=language, command_args=command_args))
+        assert any(finding.rule_id == rule_id for finding in report.findings), [
+            (finding.rule_id, finding.evidence) for finding in report.findings
+        ]
+
+
+def test_policy_loader_and_matching_branches(tmp_path):
+    policy_file = tmp_path / "policy.yaml"
+    policy_file.write_text(
+        """
+policy:
+  name: nested
+  allowed_domains:
+    - "*.example.com"
+  allowed_commands:
+    - custom-tool
+  denied_paths:
+    - ""
+    - .env
+  audit_log_path: audit.jsonl
+extra_flag: yes
+""",
+        encoding="utf-8",
+    )
+
+    loaded = ToolSafetyPolicy.load(policy_file)
+
+    assert loaded.name == "nested"
+    assert loaded.metadata == {"extra_flag": True}
+    assert loaded.audit_log_path == "audit.jsonl"
+    assert loaded.is_domain_allowed("api.example.com") is True
+    assert loaded.is_domain_allowed("example.com") is True
+    assert loaded.is_domain_allowed(None) is False
+    assert loaded.is_command_allowed("custom-tool --flag") is True
+    assert loaded.is_command_allowed(None) is False
+    assert loaded.is_denied_path(None) is False
+    assert loaded.is_system_write_path("/") is True
+    assert loaded.is_system_write_path(None) is False
+    assert load_tool_safety_policy(None).name == "default"
+
+
+def test_policy_load_rejects_non_mapping(tmp_path):
+    policy_file = tmp_path / "policy.yaml"
+    policy_file.write_text("- not-a-mapping\n", encoding="utf-8")
+
+    with pytest.raises(ValueError):
+        ToolSafetyPolicy.load(policy_file)
+
+
+def test_report_and_type_helpers_cover_empty_and_fallback_paths():
+    assert max_risk_level([]) == SafetyRiskLevel.NONE
+    assert risk_level_value("unknown") == 0
+
+    report = ToolSafetyReport(
+        decision="allow",
+        risk_level="none",
+        findings=[],
+        duration_ms=1.2345,
+        language="python",
+        scanned_at="2026-07-05T00:00:00+00:00",
+    )
+
+    assert report.is_allowed is True
+    assert json.loads(report.to_json(indent=None))["risk_count"] == 0
+
+
 def test_audit_event_contains_required_fields(tmp_path):
     report = ToolSafetyScanner(policy()).scan(
         ToolSafetyScanRequest(script="rm -rf /tmp/demo", language="bash", tool_metadata={"name": "Bash"}))
@@ -258,6 +392,61 @@ def test_report_includes_telemetry_attributes():
     assert report.telemetry_attributes["tool.safety.rule_id"] == "TSG-RESOURCE-INFINITE-LOOP"
 
 
+def test_telemetry_handles_list_values_and_span_errors(monkeypatch):
+    class Span:
+        def __init__(self):
+            self.attributes = {}
+
+        def set_attribute(self, key, value):
+            self.attributes[key] = value
+
+    span = Span()
+    report = ToolSafetyReport(
+        decision=SafetyDecision.ALLOW,
+        risk_level=SafetyRiskLevel.LOW,
+        findings=[],
+        duration_ms=0,
+        language="python",
+        scanned_at="2026-07-05T00:00:00+00:00",
+        telemetry_attributes={
+            "string": "value",
+            "list": ["a", "b"],
+            "object": {"x": 1},
+        },
+    )
+
+    monkeypatch.setattr("trpc_agent_sdk.tools.safety._telemetry.trace.get_current_span", lambda: span)
+    apply_tool_safety_span_attributes(report)
+    assert span.attributes["list"] == "a,b"
+    assert span.attributes["object"] == "{'x': 1}"
+
+    class RaisingSpan:
+        def set_attribute(self, key, value):
+            raise RuntimeError("boom")
+
+    monkeypatch.setattr("trpc_agent_sdk.tools.safety._telemetry.trace.get_current_span", lambda: RaisingSpan())
+    apply_tool_safety_span_attributes(report)
+
+
+@pytest.mark.asyncio
+async def test_guard_check_and_run_if_allowed_paths(tmp_path):
+    audit_path = tmp_path / "audit.jsonl"
+    guard = ToolSafetyGuard(policy=policy(audit_log_path=str(audit_path)), apply_telemetry=False)
+
+    report = guard.check(ToolSafetyScanRequest(script="print('ok')", language="python"))
+    assert report.is_allowed
+    assert audit_path.exists()
+
+    result = await guard.run_if_allowed(
+        ToolSafetyScanRequest(script="print('ok')", language="python"),
+        lambda: AsyncMock(return_value="ran")(),
+    )
+    assert result == "ran"
+
+    with pytest.raises(ToolSafetyBlockedError):
+        guard.check(ToolSafetyScanRequest(script="rm -rf /tmp/demo", language="bash"))
+
+
 @pytest.mark.asyncio
 async def test_tool_filter_blocks_before_handle():
     filter_ = ToolSafetyFilter(guard=ToolSafetyGuard(policy=policy()))
@@ -286,6 +475,20 @@ async def test_tool_filter_block_writes_audit_event(tmp_path):
     assert event["blocked"] is True
 
 
+@pytest.mark.asyncio
+async def test_tool_filter_policy_path_and_passthrough_branches(tmp_path):
+    policy_path = tmp_path / "policy.yaml"
+    policy_path.write_text("allowed_commands:\n  - echo\n", encoding="utf-8")
+    filter_ = ToolSafetyFilter(policy_path=policy_path)
+
+    handle_result = FilterResult(rsp={"ok": True})
+    result = await filter_.run(create_agent_context(), {"note": "not script"}, AsyncMock(return_value=handle_result))
+    assert result is handle_result
+
+    result = await filter_.run(create_agent_context(), "not mapping", AsyncMock(return_value={"plain": True}))
+    assert result.rsp == {"plain": True}
+
+
 @pytest.mark.asyncio
 async def test_tool_filter_allows_safe_command():
     filter_ = ToolSafetyFilter(guard=ToolSafetyGuard(policy=policy()))
@@ -318,6 +521,54 @@ async def execute_code(self, invocation_context, code_execution_input):
     assert delegate.called is False
 
 
+@pytest.mark.asyncio
+async def test_safety_guarded_code_executor_code_string_paths():
+    class DelegateExecutor(BaseCodeExecutor):
+        called: bool = False
+        work_dir: str = "/workspace"
+
+        async def execute_code(self, invocation_context, code_execution_input):
+            self.called = True
+            return create_code_execution_result(stdout="executed")
+
+    delegate = DelegateExecutor()
+    executor = SafetyGuardedCodeExecutor(delegate=delegate, guard=ToolSafetyGuard(policy=policy()))
+
+    blocked = await executor.execute_code(
+        invocation_context=AsyncMock(),
+        code_execution_input=CodeExecutionInput(code="while True:\n    pass"),
+    )
+    assert "Tool safety guard blocked" in blocked.output
+    assert delegate.called is False
+
+    allowed = await executor.execute_code(
+        invocation_context=AsyncMock(),
+        code_execution_input=CodeExecutionInput(code="print('ok')"),
+    )
+    assert "executed" in allowed.output
+    assert delegate.called is True
+
+
+@pytest.mark.parametrize(
+    ("args", "expected"),
+    [
+        ({"command": "echo 'unterminated", "cwd": "/tmp"}, ("echo 'unterminated", "bash", "/tmp", ["echo", "'unterminated"])),
+        ({"script": "#!/bin/bash\necho ok"}, ("#!/bin/bash\necho ok", "bash", None, [])),
+        ({"script": "set -e\necho ok"}, ("set -e\necho ok", "bash", None, [])),
+        ({"code": "print('ok')", "lang": "py"}, ("print('ok')", "python", None, [])),
+        ({"code_blocks": [{"code": "echo ok", "language": "sh"}]}, ("echo ok", "bash", None, [])),
+        ({"code_blocks": [{"language": "sh"}]}, None),
+    ],
+)
+def test_extract_script_from_tool_args_branches(args, expected):
+    assert extract_script_from_tool_args(args) == expected
+
+
+def test_infer_language_from_key_default():
+    assert infer_language_from_key("source", "anything") == "python"
+    assert command_to_args("echo ok") == ["echo", "ok"]
+
+
 def test_scans_500_line_script_under_one_second():
     script = "\n".join(f'print("line {i}")' for i in range(500))
     scanner = ToolSafetyScanner(policy())

trpc_agent_sdk/tools/safety/_policy.py

@@ -161,13 +161,15 @@ def is_system_write_path(self, value: str | None) -> bool:
             return False
         normalized = self._normalize_path_text(value)
         for pattern in self.system_write_paths:
-            normalized_pattern = self._normalize_path_text(str(pattern)).rstrip("/")
+            normalized_pattern = self._normalize_path_text(str(pattern))
             if not normalized_pattern:
                 continue
             if normalized_pattern == "/":
                 if normalized in {"/", "/*", "/."}:
                     return True
-            elif normalized == normalized_pattern or normalized.startswith(f"{normalized_pattern}/"):
+                continue
+            normalized_pattern = normalized_pattern.rstrip("/")
+            if normalized == normalized_pattern or normalized.startswith(f"{normalized_pattern}/"):
                 return True
         return False
 

trpc_agent_sdk/tools/safety/_scanner.py

@@ -1208,8 +1208,17 @@ def timeout_seconds_from_tokens(tokens: Sequence[str]) -> float | None:
     command = Path(tokens[0]).name.lower() if tokens else ""
     if command != "timeout":
         return None
-    for token in tokens[1:]:
+    index = 1
+    while index < len(tokens):
+        token = tokens[index]
+        if token in {"-k", "--kill-after"}:
+            index += 2
+            continue
+        if token.startswith("--kill-after="):
+            index += 1
+            continue
         if token.startswith("-"):
+            index += 1
             continue
         return duration_token_seconds(token)
     return None