|
6 | 6 | """Wrapper helpers showing how to attach the safety filter to existing tools |
7 | 7 | and code executors without modifying SDK internals. |
8 | 8 |
|
9 | | -Two shapes are provided: |
| 9 | +Four shapes are provided: |
10 | 10 | - wrap_tool: returns a new tool with the safety filter prepended. |
11 | | -- wrap_executor: returns a new BaseCodeExecutor subclass whose execute_code |
| 11 | +- SafeCodeExecutor: returns a new BaseCodeExecutor subclass whose execute_code |
12 | 12 | runs the scanner before delegating to the wrapped executor. |
| 13 | +- safety_wrapper: decorator that scans a function's script argument before |
| 14 | + the function body runs (sync or async). |
| 15 | +- SafetyReviewedSkillRunner: wraps a skill runner callable with pre-execution |
| 16 | + safety scanning, covering the Skill execution path. |
13 | 17 | """ |
14 | 18 | from __future__ import annotations |
15 | 19 |
|
| 20 | +import asyncio |
| 21 | +import functools |
16 | 22 | from typing import Any |
17 | 23 | from typing import Optional |
18 | 24 |
|
| 25 | +from .audit import AuditLogger |
19 | 26 | from .policy import PolicyConfig |
20 | 27 | from .scanner import SafetyScanner |
21 | 28 | from .tool_filter import ToolSafetyFilter |
@@ -59,3 +66,150 @@ async def execute_code(self, invocation_context, input_data): |
59 | 66 | return await inner.execute_code(invocation_context, input_data) |
60 | 67 |
|
61 | 68 | return _SafeCodeExecutor() |
| 69 | + |
| 70 | + |
| 71 | +# --------------------------------------------------------------------------- |
| 72 | +# Exception |
| 73 | +# --------------------------------------------------------------------------- |
| 74 | + |
| 75 | + |
| 76 | +class SafetyDeniedError(RuntimeError): |
| 77 | + """Raised when a safety wrapper blocks a script (decision == DENY).""" |
| 78 | + |
| 79 | + def __init__(self, report): |
| 80 | + self.report = report |
| 81 | + rule_ids = report.rule_ids if report.rule_ids else ["unknown"] |
| 82 | + super().__init__(f"script denied by rule(s) {rule_ids}") |
| 83 | + |
| 84 | + |
| 85 | +# --------------------------------------------------------------------------- |
| 86 | +# Decorator |
| 87 | +# --------------------------------------------------------------------------- |
| 88 | + |
| 89 | + |
| 90 | +def safety_wrapper(tool_name="unknown", *, script_arg="script", |
| 91 | + policy=None, audit_path=None, raise_on_deny=True): |
| 92 | + """Decorator: scan the *script_arg* of a function before it runs. |
| 93 | +
|
| 94 | + Works on both sync and async functions. When the scan decision is DENY |
| 95 | + and *raise_on_deny* is True, raises :class:`SafetyDeniedError`. |
| 96 | +
|
| 97 | + The decorated function's keyword argument named *script_arg* (or a key |
| 98 | + inside a dict positional argument) is scanned before the body runs. |
| 99 | +
|
| 100 | + Example:: |
| 101 | +
|
| 102 | + @safety_wrapper(tool_name="runner", script_arg="code") |
| 103 | + async def execute(*, tool_context, args): |
| 104 | + ... |
| 105 | +
|
| 106 | + Args: |
| 107 | + tool_name: Name used in audit/report. |
| 108 | + script_arg: Name of the kwarg (or key in a dict positional arg) |
| 109 | + that holds the script text. |
| 110 | + policy: Optional PolicyConfig; uses defaults when None. |
| 111 | + audit_path: Optional path for JSONL audit log. |
| 112 | + raise_on_deny: Raise SafetyDeniedError on DENY (default True). |
| 113 | + """ |
| 114 | + if policy is None: |
| 115 | + policy = PolicyConfig() |
| 116 | + _scanner = SafetyScanner(policy=policy) |
| 117 | + _audit = AuditLogger(audit_path) |
| 118 | + |
| 119 | + def _extract_script(args, kwargs): |
| 120 | + script = kwargs.get(script_arg) |
| 121 | + if script is None: |
| 122 | + for arg in args: |
| 123 | + if isinstance(arg, dict) and script_arg in arg: |
| 124 | + return arg[script_arg] |
| 125 | + return script |
| 126 | + |
| 127 | + def _guard(args, kwargs): |
| 128 | + script = _extract_script(args, kwargs) |
| 129 | + if not script or not isinstance(script, str): |
| 130 | + return |
| 131 | + report = _scanner.scan(ScanInput(script=script, tool_name=tool_name)) |
| 132 | + _audit.log(report, intercepted=report.blocked) |
| 133 | + if report.decision == Decision.DENY and raise_on_deny: |
| 134 | + raise SafetyDeniedError(report) |
| 135 | + |
| 136 | + def decorator(func): |
| 137 | + @functools.wraps(func) |
| 138 | + async def async_wrapper(*args, **kwargs): |
| 139 | + _guard(args, kwargs) |
| 140 | + return await func(*args, **kwargs) |
| 141 | + |
| 142 | + @functools.wraps(func) |
| 143 | + def sync_wrapper(*args, **kwargs): |
| 144 | + _guard(args, kwargs) |
| 145 | + return func(*args, **kwargs) |
| 146 | + |
| 147 | + return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper |
| 148 | + |
| 149 | + return decorator |
| 150 | + |
| 151 | + |
| 152 | +# --------------------------------------------------------------------------- |
| 153 | +# Skill runner wrapper |
| 154 | +# --------------------------------------------------------------------------- |
| 155 | + |
| 156 | + |
| 157 | +class SafetyReviewedSkillRunner: |
| 158 | + """Wrap a skill runner callable with pre-execution safety scanning. |
| 159 | +
|
| 160 | + Covers the Skill execution path: works with any callable that accepts |
| 161 | + ``(tool_context, args)`` or has a ``run_async(tool_context=, args=)`` |
| 162 | + method. When the scan decision is DENY (or NEEDS_HUMAN_REVIEW when |
| 163 | + *block_review* is True), execution is skipped and a blocked-response |
| 164 | + dict is returned. |
| 165 | +
|
| 166 | + Example:: |
| 167 | +
|
| 168 | + safe_skill = SafetyReviewedSkillRunner( |
| 169 | + my_skill_runner, policy=policy, tool_name="skill_run", |
| 170 | + ) |
| 171 | + result = await safe_skill.run(tool_context, args) |
| 172 | + """ |
| 173 | + |
| 174 | + def __init__(self, runner, policy, *, audit_path=None, |
| 175 | + block_review=False, tool_name="skill_run"): |
| 176 | + self._runner = runner |
| 177 | + self._scanner = SafetyScanner(policy=policy) |
| 178 | + self._audit = AuditLogger(audit_path) |
| 179 | + self._block_review = block_review |
| 180 | + self._tool_name = tool_name |
| 181 | + |
| 182 | + async def run(self, tool_context, args): |
| 183 | + """Scan skill args and delegate to the wrapped runner when allowed.""" |
| 184 | + script = self._extract_script(args) |
| 185 | + if script: |
| 186 | + report = self._scanner.scan( |
| 187 | + ScanInput(script=script, tool_name=self._tool_name) |
| 188 | + ) |
| 189 | + self._audit.log(report, intercepted=report.blocked) |
| 190 | + if report.decision == Decision.DENY: |
| 191 | + return {"success": False, "error": "SKILL_BLOCKED", |
| 192 | + "safety": report.to_dict()} |
| 193 | + if report.decision == Decision.NEEDS_HUMAN_REVIEW and self._block_review: |
| 194 | + return {"success": False, "error": "SKILL_NEEDS_REVIEW", |
| 195 | + "safety": report.to_dict()} |
| 196 | + |
| 197 | + # Delegate to the wrapped runner. |
| 198 | + if hasattr(self._runner, "run_async"): |
| 199 | + result = self._runner.run_async(tool_context=tool_context, args=args) |
| 200 | + else: |
| 201 | + result = self._runner(tool_context, args) |
| 202 | + if hasattr(result, "__await__"): |
| 203 | + return await result |
| 204 | + return result |
| 205 | + |
| 206 | + @staticmethod |
| 207 | + def _extract_script(args): |
| 208 | + """Pull script content from common skill arg shapes.""" |
| 209 | + if not isinstance(args, dict): |
| 210 | + return None |
| 211 | + for key in ("script", "code", "command", "cmd"): |
| 212 | + val = args.get(key) |
| 213 | + if isinstance(val, str): |
| 214 | + return val |
| 215 | + return None |
0 commit comments