53 lines
1.9 KiB
Python
53 lines
1.9 KiB
Python
"""Reject runaway argument streams for small, built-in device actions.
|
|
|
|
This is a transport safety bound, not a model output-token budget. Free-text
|
|
tools (input_text, save_note, reports and plugin calls) are deliberately excluded.
|
|
Only lengths and tool names are retained; argument bodies may contain user data.
|
|
"""
|
|
|
|
from dataclasses import dataclass
|
|
|
|
|
|
SMALL_ACTION_ARGUMENT_LIMIT = 16 * 1024
|
|
_SMALL_ACTIONS = frozenset(
|
|
{"click", "long_press", "swipe", "press_key", "manage_app", "wait_for_delay"}
|
|
)
|
|
|
|
|
|
class ToolArgumentOverflow(ValueError):
|
|
def __init__(self, name: str, characters: int):
|
|
self.details = {
|
|
"reason": "tool_arguments_too_large",
|
|
"tool": name,
|
|
"argument_chars": characters,
|
|
"limit_chars": SMALL_ACTION_ARGUMENT_LIMIT,
|
|
}
|
|
super().__init__(
|
|
f"收到的 {name} 参数流超出安全限制（{characters} 字符，"
|
|
f"保护上限 {SMALL_ACTION_ARGUMENT_LIMIT}）；已关闭输出流，未执行本次动作。"
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class _ToolStream:
|
|
name: str | None = None
|
|
characters: int = 0
|
|
|
|
|
|
class ToolArgumentGuard:
|
|
def __init__(self):
|
|
self.calls: dict[int | str, _ToolStream] = {}
|
|
|
|
def observe(self, chunk) -> None:
|
|
for tool in getattr(chunk, "tool_call_chunks", None) or []:
|
|
index = tool.get("index")
|
|
key = index if index is not None else tool.get("id") or "unknown"
|
|
state = self.calls.setdefault(key, _ToolStream())
|
|
if isinstance(tool.get("name"), str) and tool["name"]:
|
|
state.name = tool["name"]
|
|
arguments = tool.get("args")
|
|
if isinstance(arguments, str):
|
|
state.characters += len(arguments)
|
|
if state.name in _SMALL_ACTIONS and state.characters > SMALL_ACTION_ARGUMENT_LIMIT:
|
|
raise ToolArgumentOverflow(state.name, state.characters)
|