218 lines
9.1 KiB
Python
218 lines
9.1 KiB
Python
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""**UserPromptSubmit 多守卫合并入口**（2026-10-06 用户报「三个 hook 同时 10s 超时」）。
|
|||
|
|
|
|||
|
|
## 为什么要合并（实测根因，⛔ 不是"脚本慢"）
|
|||
|
|
|
|||
|
|
单一 `UserPromptSubmit` 上串行挂了 **7 条 hook**，每条都是**一个独立的 Python 进程**
|
|||
|
|
（`python.exe <某 hook>.py`）⇒ 宿主要**冷启 7 次解释器**，每次 0.3–1 s
|
|||
|
|
（Windows 上首次 import 开销尤其大，`datetime`/`json`/`re` 全都要重新加载）。
|
|||
|
|
|
|||
|
|
实测（本机，同一份真实 190 MB transcript 负载）：
|
|||
|
|
|
|||
|
|
| 跑法 | 耗时 |
|
|||
|
|
|---|---|
|
|||
|
|
| 任一条 guard 单跑 | 0.23–0.36 s |
|
|||
|
|
| 7 条**并发**（多进程） | 0.84 s（墙钟） |
|
|||
|
|
| ⚠️ 宿主里**串行** 7 条 | 累计轻松顶破 10 s ⇒ **4 条 `timeout=10` 的一起报超时** |
|
|||
|
|
|
|||
|
|
⇒ 结论：**单条都不慢，死在"7 次冷启动 + 串行排队"**。
|
|||
|
|
|
|||
|
|
## 本入口做什么
|
|||
|
|
|
|||
|
|
把**同一事件、同一 stdin 契约**的 4 个 guard 合进**一个进程**：
|
|||
|
|
`reply-style-guard` / `stop-dialog-guard` / `skill-load-guard` / `session-log-guard`
|
|||
|
|
|
|||
|
|
一次读 stdin ⇒ 依次在**本进程内**调各自的 `main()` ⇒ 合并 stdout ⇒ 一次输出。
|
|||
|
|
**省掉 3 次 Python 冷启动**（4 条 → 1 条）。
|
|||
|
|
|
|||
|
|
## 契约（必须与各 guard 保持一致）
|
|||
|
|
|
|||
|
|
- 各自都从 `sys.stdin.buffer.read()` 读**同一份** payload（事件名都是 `UserPromptSubmit`）
|
|||
|
|
- 各自都往 `sys.stdout` 写**一个 JSON**（可能空 ⇒ 表示放行）
|
|||
|
|
- 合并规则：把各自的 JSON 合并 ——
|
|||
|
|
· `hookSpecificOutput.additionalContext` **字符串拼接**（多个守卫都要注入时不能互相顶掉）
|
|||
|
|
· `decision` / `permissionDecision` 等**谁拦谁说了算**（先出现的非空值优先，并记来源）
|
|||
|
|
⚠️ ⛔ 不吞任何一条的"阻止"决定 —— 那是安全语义，合并只能加信息、不能减判定。
|
|||
|
|
|
|||
|
|
## 硬纪律（沿用本包同族）
|
|||
|
|
|
|||
|
|
1. 🔴 **输出走 `sys.stdout.buffer.write(bytes)`**（cp936 下文本模式会炸，见 `_env`/`pitfalls`）。
|
|||
|
|
2. 🔴 **fail-open**：本入口自己不判任何业务规则；任一 guard 抛异常 ⇒ 只记 stderr，
|
|||
|
|
其余 guard 照跑（⛔ 不许一个坏 guard 把别的都拖没）。
|
|||
|
|
3. 🔴 **不改各 guard 源码**：用 `runpy` 按路径加载、临时接管 stdout 取回它们的 JSON
|
|||
|
|
⇒ 单跑与合并跑**判据同源**（⛔ 不复制一份规则，避免"一条规则两份实现必漂"）。
|
|||
|
|
4. ⚠️ 保持**单跑入口不变**（各 `*-guard.py` 仍可直接被宿主调用）——本文件是**可选的加速入口**。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import io
|
|||
|
|
import json
|
|||
|
|
import os
|
|||
|
|
import runpy
|
|||
|
|
import sys
|
|||
|
|
import time
|
|||
|
|
|
|||
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|||
|
|
|
|||
|
|
# 🔴 合并的守卫清单（⛔ 顺序＝执行顺序；都是 fail-open 的"补充信息/提醒"类）
|
|||
|
|
# ⚠️ **不含** `supervise-ensure-hook` / `wb-result-hook` / `decision_bridge` ——
|
|||
|
|
# 它们有副作用（投递/常驻确保）或属别处接线，**合并会改变语义**，一律不碰。
|
|||
|
|
GUARDS = (
|
|||
|
|
"reply-style-guard.py",
|
|||
|
|
"stop-dialog-guard.py",
|
|||
|
|
"skill-load-guard.py",
|
|||
|
|
"session-log-guard.py",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _log(msg: str) -> None:
|
|||
|
|
"""⛔ 绝不写 stdout（那是协议通道）；失败只上 stderr 且吞异常。"""
|
|||
|
|
try:
|
|||
|
|
sys.stderr.write("[prompt-guards] %s\n" % msg)
|
|||
|
|
sys.stderr.flush()
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
|
|||
|
|
|
|||
|
|
class _Tee(io.TextIOBase):
|
|||
|
|
"""临时接管 `sys.stdout`：把 guard 写出的 JSON 收下来（⛔ 不落终端）。
|
|||
|
|
|
|||
|
|
⚠️🔴 **踩过的坑（2026-10-06 实测）**：多数 guard 走 `sys.stdout.buffer.write(bytes)`。
|
|||
|
|
最初把 `self.buffer = self` ⇒ 那次 `write` **撞上 `TextIOBase.write`**，
|
|||
|
|
于是 bytes 被 `str()` 成 `b'{"hookSpecificOutput": ...}'` ⇒ **JSON 解析必失败**
|
|||
|
|
（现象：合并入口报「输出了非 JSON（6228 字节）」）。
|
|||
|
|
✅ 正解：`buffer` 用一个**独立的 `io.BytesIO`**；文本写与字节写**各收各的**，
|
|||
|
|
`text()` 时按"谁有内容取谁"取回（正常只会有一路有内容）。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
self._raw = io.BytesIO() # 给 `sys.stdout.buffer.write(bytes)` 用
|
|||
|
|
self.buffer = self._raw
|
|||
|
|
self._txt = [] # 给 `sys.stdout.write(str)` 用（少数 guard 会走）
|
|||
|
|
|
|||
|
|
# TextIOBase 接口
|
|||
|
|
def write(self, s): # noqa: D102
|
|||
|
|
try:
|
|||
|
|
self._txt.append(str(s))
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
return len(s or "")
|
|||
|
|
def flush(self): # noqa: D102
|
|||
|
|
pass
|
|||
|
|
def isatty(self): # noqa: D102
|
|||
|
|
return False
|
|||
|
|
def writable(self): # noqa: D102
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
def text(self) -> str:
|
|||
|
|
"""取回收集到的内容（字节路优先；两路都空 ⇒ 空串）。"""
|
|||
|
|
b = self._raw.getvalue()
|
|||
|
|
if b:
|
|||
|
|
return b.decode("utf-8", "replace")
|
|||
|
|
return "".join(self._txt)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _run_one(name: str, payload_bytes: bytes):
|
|||
|
|
"""在同进程内跑一个 guard，返回 (json_or_None, seconds)。
|
|||
|
|
|
|||
|
|
⚠️ guard 读的是 **stdin**；本入口已把 stdin 读过一次 ⇒ 这里给它一个"假 stdin"，
|
|||
|
|
内容是**同一份 payload 字节**（保证各 guard 看到的输入完全一致）。
|
|||
|
|
"""
|
|||
|
|
path = os.path.join(HERE, name)
|
|||
|
|
if not os.path.isfile(path):
|
|||
|
|
_log("跳过（不存在）：%s" % name)
|
|||
|
|
return None, 0.0
|
|||
|
|
|
|||
|
|
old_out, old_in = sys.stdout, sys.stdin
|
|||
|
|
tee = _Tee()
|
|||
|
|
# ⚠️ `TextIOWrapper.buffer` 是**只读属性**（实测 `AttributeError: readonly attribute`）
|
|||
|
|
# ⇒ 不能后赋值；正解＝**直接构造一个带 BytesIO 的 TextIOWrapper**，
|
|||
|
|
# 它的 `.buffer` 天然就是那个 BytesIO（guard 走 `sys.stdin.buffer.read()`）。
|
|||
|
|
fake_in = io.TextIOWrapper(io.BytesIO(payload_bytes), encoding="utf-8")
|
|||
|
|
t0 = time.time()
|
|||
|
|
try:
|
|||
|
|
sys.stdout = tee
|
|||
|
|
sys.stdin = fake_in
|
|||
|
|
runpy.run_path(path, run_name="__main__")
|
|||
|
|
except SystemExit:
|
|||
|
|
pass # guard 自己 exit ⇒ 正常
|
|||
|
|
except Exception as e: # noqa: BLE001
|
|||
|
|
_log("%s 抛异常（已忽略，fail-open）：%r" % (name, e))
|
|||
|
|
finally:
|
|||
|
|
sys.stdout, sys.stdin = old_out, old_in
|
|||
|
|
el = time.time() - t0
|
|||
|
|
txt = tee.text().strip()
|
|||
|
|
if not txt:
|
|||
|
|
return None, el
|
|||
|
|
try:
|
|||
|
|
return json.loads(txt), el
|
|||
|
|
except Exception:
|
|||
|
|
# 非 JSON（不合契约）⇒ 留痕但不当成阻止
|
|||
|
|
_log("%s 输出了非 JSON（%d 字节，已忽略）" % (name, len(txt)))
|
|||
|
|
return None, el
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _merge(dst: dict, src: dict, who: str) -> None:
|
|||
|
|
"""把 `src` 合并进 `dst`（⛔ 只加不减；冲突时保留先到者并记来源）。"""
|
|||
|
|
if not isinstance(src, dict):
|
|||
|
|
return
|
|||
|
|
for k, v in src.items():
|
|||
|
|
if k == "hookSpecificOutput" and isinstance(v, dict):
|
|||
|
|
cur = dst.setdefault("hookSpecificOutput", {})
|
|||
|
|
if not isinstance(cur, dict):
|
|||
|
|
continue
|
|||
|
|
for k2, v2 in v.items():
|
|||
|
|
if k2 == "additionalContext" and isinstance(v2, str):
|
|||
|
|
# 🔴 多守卫都要注入 ⇒ **拼接**（⛔ 别互相顶掉）
|
|||
|
|
old = cur.get("additionalContext") or ""
|
|||
|
|
cur[k2] = (old + ("\n\n" if old else "") + v2) if old else v2
|
|||
|
|
else:
|
|||
|
|
if k2 not in cur:
|
|||
|
|
cur[k2] = v2
|
|||
|
|
elif cur[k2] != v2:
|
|||
|
|
_log("字段冲突 %s.%s：保留先到者，丢弃 %s 的值" % (who, k2, who))
|
|||
|
|
else:
|
|||
|
|
if k not in dst:
|
|||
|
|
dst[k] = v
|
|||
|
|
# ⚠️ 同名字段（如 `decision`）冲突 ⇒ **保留先到者**；因为均为 fail-open 类，
|
|||
|
|
# 先到者通常来自更"硬"的守卫（格式/收尾）；真要强制多判定应拆开接线。
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main() -> None:
|
|||
|
|
try:
|
|||
|
|
raw = sys.stdin.buffer.read()
|
|||
|
|
except Exception:
|
|||
|
|
raw = b""
|
|||
|
|
if not raw:
|
|||
|
|
return # 空 payload ⇒ 静默放行（同各 guard）
|
|||
|
|
|
|||
|
|
out = {}
|
|||
|
|
timings = []
|
|||
|
|
for name in GUARDS:
|
|||
|
|
obj, el = _run_one(name, raw)
|
|||
|
|
timings.append("%s=%.3fs" % (name.split("-")[0], el))
|
|||
|
|
if obj:
|
|||
|
|
_merge(out, obj, name)
|
|||
|
|
|
|||
|
|
if not out:
|
|||
|
|
return # 全都没话说 ⇒ 零输出（⛔ 不写空 JSON）
|
|||
|
|
|
|||
|
|
# 🔴 出口：必须走 buffer bytes（cp936 下文本模式炸）
|
|||
|
|
try:
|
|||
|
|
out.setdefault("hookSpecificOutput", {})["hookEventName"] = "UserPromptSubmit"
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
try:
|
|||
|
|
sys.stdout.buffer.write(json.dumps(out, ensure_ascii=False).encode("utf-8"))
|
|||
|
|
sys.stdout.buffer.flush()
|
|||
|
|
except Exception as e: # noqa: BLE001
|
|||
|
|
_log("写 stdout 失败：%r" % (e,))
|
|||
|
|
_log("合并完成：%s" % " ".join(timings))
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
try:
|
|||
|
|
main()
|
|||
|
|
except Exception as e: # noqa: BLE001
|
|||
|
|
_log("FATAL %r" % (e,)) # fail-open：异常也不阻断提交
|