Files
dsh_shenxian/dsh-server-docs/scripts/bash-output-guard.py
T
admin 03c8363960 feat(cost): 省积分机制 —— 会话预算告警 + Bash 大输出门禁
本会话实测:553 轮 × 平均 42 万 token = 2.33 亿 input,output 仅占 0.35%(286:1)。
根因 = **对话历史是 append-only**:工具输出(`function_call_result`)一旦生成就**永久留在 messages 里、
每轮全量重发**,模型无权删自己的历史 ⇒ **唯一能"去掉"的时机 = 输出被生成之前**
(平台 `contextWindow=100 万` / 压缩阈值 90% 只是放大器,不是根因)。

- `scripts/stop-dialog-guard.py`:
  ① 修 `-S -E` ⇒ stdin 回退 cp936 ⇒ 含中文 payload 解析失败导致的**静默空转**(改走 buffer 显式 UTF-8,
     读写都加固;入口即留痕;根治"日志缺失时无法区分『没被调用』与『被静默 return』")
  ② 新增 `session_budget()`:读转录 `usage.input_tokens` ⇒ 当前体量 + 累计工具调用 + **上一轮增量**
  ③ 阈值 `BUDGET_TOKENS=120000` / `BUDGET_TOOLS=80` ⇒ 超预算经 `UserPromptSubmit` 注入告警
- `scripts/bash-output-guard.py`(新):`PreToolUse(Bash)` 拦"几乎必然巨大"的 6 条读命令
  (`cat *.log/jsonl` 无管道 / `grep -r` 无管道 / `ls -laR` / `find /` / `journalctl` 无 `-n` / `dmesg` 无 `head`),
  deny 文案**必给等价限流写法**。
  **设计(用户质疑"全拦也有问题")**:⛔ **不全拦** —— 误拦挡住正事比漏拦更贵 ⇒
  默认只拦高置信度、**拿不准一律放行**、`rg`/`git log`/`du·tree` 仅 `hard` 模式拦、
  可 `off` 急停(`<工作区>/.workbuddy/bash-guard-mode`)、任何异常 **fail-open**。
  实测 8/8:soft 拦 cat/ls-lR、放行 rg/git-log/grep+head;hard 全拦;off 全放;deny JSON 正确输出。
- (`settings.json` 的 `PreToolUse` 已加 `"matcher": "Bash"` 条目 —— **该文件不在仓库内**,
  且 hooks 是**启动时快照** ⇒ **Bash 门禁需完全重启才生效**;脚本内容本身每次现读、改完即生效。)

⚠️ 记录两个已踩的坑:Python `%` 格式串里的**裸 `%` 必须写 `%%`**(否则 `TypeError` ⇒ 在 `except: pass` 里
**静默失效**,表现为"日志有 DENY 但 stdout 空"= 看起来拦了其实没生效)⇒ **hook 必须"写日志 + emit"双动作**。
2026-09-15 21:17:08 +08:00

174 lines
7.8 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""bash-output-guard —— `PreToolUse(Bash)`:在**命令执行前**拦下"会灌爆上下文"的读命令。
为什么需要(2026-09-15 实测)
────────────────────────────
对话历史是 **append-only**:一次工具调用在转录里落两条记录 —— `function_call`(命令)
+ `function_call_result`(**输出正文**);**输出一旦生成就永久留在 messages 里,且每轮全量重发**。
实测某会话 **553 轮 × 平均 42 万 token = 2.33 亿 input**,而 output 仅占 **0.35%**。
⇒ 想真正"把大输出从上下文里去掉",**唯一的时机就是它被生成之前**(模型无权删自己的历史)。
边界(重要)
──────────
* **只拦"读"类高风险命令**,且**必须给出等价限流写法**(教而非堵)。
* **判不准就放行**(fail-open)—— 本机不是沙箱,但这类命令本身不破坏数据,误拦的代价只是麻烦。
* 急停:env `DSH_OUTPUT_GUARD_OFF=1`,或新建 `<工作区>/.workbuddy/bash-guard.disabled`。
* 命中才写一行日志 `<工作区>/.workbuddy/bash-guard.log`(用于调误报,上限 300 行)。
安装(`settings.json` 的 hooks 段 · **新增一条** · ⚠️ **需完全重启才生效**)
────────────────────────────────────────────────────────────────────
"PreToolUse": [ …,
{ "matcher": "Bash",
"hooks": [{ "type": "command",
"command": "\"<python>\" \"<此脚本>\"", "timeout": 10 }] } ]
⛔ **别给本脚本加 `-E`**:`-E` 会屏蔽 `PYTHONUTF8`/`PYTHONIOENCODING` ⇒ stdin 回退 cp936 ⇒
含中文的 payload 解析失败且**静默 fail-open**(同目录 `stop-dialog-guard.py` 已因此"白排查一天")。
"""
import io
import json
import os
import re
import sys
import time
LOG_REL = os.path.join('.workbuddy', 'bash-guard.log')
SCOPE = 'aliyun-dsh-server'
# 兜底工作区:脚本位于 <工作区>/dsh-server-docs/scripts/ ⇒ 上溯三级
WS_FALLBACK = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# (是否 core, 正则, 名称, 更省的写法)
# ⚠️ 设计原则:**不能全拦** —— 误拦挡住正事 ⇒ 我会来回试/绕路 ⇒ 反而更贵;漏拦只是多花点 token。
# core=True = 几乎必然巨大、误拦率≈0 ⇒ **默认就拦**
# core=False = 经常确实需要全量 ⇒ **仅 hard 模式拦**(默认放行)
RULES = [
(True, r'\bcat\s+[^|>;]*\.(log|jsonl|json|ndjson|csv|txt)\b',
'`cat` 大文本文件', '改用 `head -30 文件` / `sed -n \'1,30p\' 文件` / `wc -l 文件`'),
(True, r'\bgrep\b[^|;]*-[a-zA-Z]*[rR]',
'递归 grep', '改用 `grep -rn … | head -30`,或 `grep -rc …`(只要计数)'),
(True, r'\bls\b[^|;]*-[a-zA-Z]*[lR]',
'`ls -l/-R` 全量列表', '改用 `ls -la 目录 | head -20`,或 `ls 目录 | wc -l`'),
(True, r'\bfind\s+(/[A-Za-z]|[A-Za-z]:)',
'`find` 从盘符/根起全树扫', '改用 `find 具体目录 -maxdepth 3 … | head -20`'),
(True, r'\bjournalctl\b(?!.*(-n\s*\d|head))',
'`journalctl` 无行数限制', '改用 `journalctl -u X -n 50 --no-pager`'),
(True, r'\bdmesg\b(?!.*head)',
'`dmesg` 无行数限制', '改用 `dmesg | tail -30`'),
(False, r'(^|[|;&]\s*)rg\b(?!.*\|)',
'递归 rg 无管道限流', '改用 `rg … | head -30`,或 `rg -c …`'),
(False, r'\bgit\s+(log|diff)\b(?!.*(-n\s*\d|--max-count|head))',
'`git log/diff` 无行数限制', '改用 `git log --oneline -5` / `git diff --stat`'),
(False, r'\b(du|tree)\b[^|;]*\s(/|[A-Za-z]:)',
'`du/tree` 从根起', '改用 `du -sh 具体目录` / `tree -L 2 目录 | head -40`'),
]
# 命令里已有限流 ⇒ 放行(只按"最外层"有就好了)
SAFE = re.compile(r'\|\s*(head|tail|wc|grep\s+-c|cut|sed\s+-n|awk|uniq|sort\s+-u)\b')
def _read_stdin():
"""⚠️ 必须走 buffer 显式 UTF-8(`-E` 下 sys.stdin 是 cp936)。"""
try:
return sys.stdin.buffer.read().decode('utf-8', 'replace')
except Exception:
try:
return sys.stdin.read()
except Exception:
return ''
def _emit(obj):
data = json.dumps(obj, ensure_ascii=False).encode('utf-8')
try:
sys.stdout.buffer.write(data)
sys.stdout.buffer.flush()
except Exception:
try:
sys.stdout.write(data.decode('utf-8', 'replace'))
sys.stdout.flush()
except Exception:
pass
def log(root, detail):
try:
p = os.path.join(root, LOG_REL)
os.makedirs(os.path.dirname(p), exist_ok=True)
with io.open(p, 'a', encoding='utf-8', newline='\n') as f:
f.write('%s\t%s\n' % (time.strftime('%Y-%m-%d %H:%M:%S'), detail))
lines = io.open(p, encoding='utf-8').read().split('\n')
if len(lines) > 300:
io.open(p, 'w', encoding='utf-8', newline='\n').write('\n'.join(lines[-150:]))
except Exception:
pass
def reason_for(cmd, hard=False):
for core, pat, why, fix in RULES:
if not core and not hard: # loose 条只在 hard 模式生效
continue
if re.search(pat, cmd):
return why, fix
return None, None
def main():
raw = _read_stdin()
payload = None
if raw.strip():
try:
payload = json.loads(raw)
except ValueError:
payload = None
root = (os.environ.get('CODEBUDDY_PROJECT_DIR') or os.environ.get('DSH_WORKSPACE')
or (isinstance(payload, dict) and payload.get('cwd')) or WS_FALLBACK)
tp = str((payload or {}).get('transcript_path') or '') if isinstance(payload, dict) else ''
if isinstance(payload, dict):
# 入口即留痕(只记 event,低频;便于判"有没有被调用")
log(str(root), 'entry|event=%s|in_scope=%s' % (payload.get('hook_event_name') or '(parse-fail)', SCOPE in tp))
if not isinstance(payload, dict):
return
if (payload.get('hook_event_name') or '') != 'PreToolUse':
return
if (payload.get('tool_name') or '') != 'Bash':
return
if os.environ.get('DSH_OUTPUT_GUARD_OFF'):
return
try:
if os.path.exists(os.path.join(root, '.workbuddy', 'bash-guard.disabled')):
return
except Exception:
pass
cmd = ((payload.get('tool_input') or {}).get('command')) or ''
if not cmd or SAFE.search(cmd):
return
try:
md = io.open(os.path.join(root, '.workbuddy', 'bash-guard-mode'),
encoding='utf-8').read().lower()
except OSError:
md = ''
if 'off' in md:
return
why, fix = reason_for(cmd, hard=('hard' in md))
if not why:
return
log(str(root), 'DENY|%s|%s' % (why, cmd.replace('\n', ' ')[:110]))
_emit({'hookSpecificOutput': {
'hookEventName': 'PreToolUse',
'permissionDecision': 'deny',
'permissionDecisionReason': (
'💰 拦下:**%s** —— 这类命令的输出会**永久留在会话上下文里、每轮全量重发**'
'**永久留在会话上下文里、每轮全量重发**』(实测某会话 553 轮 × 平均 42 万 token = 2.33 亿 input,'
'output 仅占 0.35%%)。\n'
'✅ 换成限流写法再发:%s\n'
'ℹ️ 若确实需要全量:**先落盘再只读关键行**(`… > /tmp/x.txt 2>&1` 然后 `sed -n \'1,40p\' /tmp/x.txt`);'
'急停用 env `DSH_OUTPUT_GUARD_OFF=1` 或新建 `<工作区>/.workbuddy/bash-guard.disabled`。'
% (why, fix))}})
if __name__ == '__main__':
try:
main()
except Exception:
pass # fail-open:本钩子只为省积分,绝不因自身异常阻断正常工作
sys.exit(0)