72 lines
2.2 KiB
Python
72 lines
2.2 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",
|
|
"action_sequence",
|
|
"drag",
|
|
"pinch",
|
|
"touch_gesture",
|
|
"read_control",
|
|
"wait_until",
|
|
"scroll_to",
|
|
}
|
|
)
|
|
_ARGUMENT_LIMITS = {
|
|
**dict.fromkeys(_SMALL_ACTIONS, SMALL_ACTION_ARGUMENT_LIMIT),
|
|
"control_action": 512 * 1024,
|
|
}
|
|
|
|
|
|
class ToolArgumentOverflow(ValueError):
|
|
def __init__(self, name: str, characters: int):
|
|
limit = _ARGUMENT_LIMITS.get(name, SMALL_ACTION_ARGUMENT_LIMIT)
|
|
self.details = {
|
|
"reason": "tool_arguments_too_large",
|
|
"tool": name,
|
|
"argument_chars": characters,
|
|
"limit_chars": limit,
|
|
}
|
|
super().__init__(
|
|
f"收到的 {name} 参数流超出安全限制（{characters} 字符，"
|
|
f"保护上限 {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 _ARGUMENT_LIMITS and state.characters > _ARGUMENT_LIMITS[state.name]:
|
|
raise ToolArgumentOverflow(state.name, state.characters)
|