refactor(agent/tool_dispatch_helpers): flatten bridge peel / overlap / risk helpers; compact docstrings
This commit is contained in:
@@ -1,12 +1,10 @@
|
||||
"""Tool-dispatch helpers — parallelism gating, multimodal envelopes, mutation tracking.
|
||||
|
||||
Stateless module-level utilities extracted from ``run_agent.py``, which
|
||||
re-exports each name so existing ``from run_agent import ...`` imports keep
|
||||
working. Groups: batch-parallelism planner (path-overlap admission; V4A patch
|
||||
scope comes from patch-body headers, not a decoy ``path=``), multimodal
|
||||
``{"_multimodal": True, "content": [...], "text_summary": ...}`` envelope
|
||||
helpers, per-turn file-mutation verifier inputs, trajectory normalisation, and
|
||||
the tool-result message constructor with its untrusted-content wrapping.
|
||||
Stateless utilities extracted from ``run_agent.py`` (which re-exports each name): the
|
||||
batch-parallelism planner (path-overlap admission; V4A patch scope comes from patch-body
|
||||
headers, not a decoy ``path=``), multimodal ``{"_multimodal": True, "content": [...],
|
||||
"text_summary": ...}`` envelope helpers, file-mutation verifier inputs, trajectory
|
||||
normalisation, and the tool-result message constructor with untrusted-content wrapping.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -45,15 +43,13 @@ _PARALLEL_SAFE_TOOLS = frozenset({
|
||||
"web_search",
|
||||
})
|
||||
|
||||
# Filesystem tools admitted by path overlap. Readers may share a subtree; a
|
||||
# writer conflicts with ANY overlapping reservation. This keeps a batched
|
||||
# read_file/search_files from observing pre-mutation state when the model
|
||||
# batches it alongside the patch/write_file it depends on.
|
||||
# Filesystem tools admitted by path overlap: readers may share a subtree, a writer conflicts
|
||||
# with ANY overlapping reservation (so a batched read never observes pre-mutation state).
|
||||
_PATH_SCOPED_READERS = frozenset({"read_file", "search_files"})
|
||||
_PATH_SCOPED_WRITERS = frozenset({"write_file", "patch"})
|
||||
_PATH_SCOPED_TOOLS = _PATH_SCOPED_READERS | _PATH_SCOPED_WRITERS
|
||||
|
||||
# Patterns that indicate a terminal command may modify/delete files.
|
||||
# Terminal commands that may modify/delete files.
|
||||
_DESTRUCTIVE_PATTERNS = re.compile(
|
||||
r"""(?:^|\s|&&|\|\||;|`)(?:
|
||||
rm\s|rmdir\s|
|
||||
@@ -90,23 +86,18 @@ _PARALLEL_SAFE_BRIDGE_LOOKUPS = frozenset({"tool_search", "tool_describe"})
|
||||
|
||||
|
||||
def _peel_bridge_call(tool_name: str, function_args: dict) -> tuple[str, dict]:
|
||||
"""Resolve a ``tool_call`` bridge invocation to ``(underlying_name, underlying_args)``.
|
||||
|
||||
With tool search active the model emits the literal name ``tool_call`` for
|
||||
every deferred tool, so admission must be decided on the underlying tool
|
||||
(as the executors' unwrap does). An unparseable bridge call is returned
|
||||
unchanged: it stays a sequential barrier and fails at dispatch as before.
|
||||
"""
|
||||
"""Resolve a ``tool_call`` bridge invocation to ``(underlying_name, underlying_args)`` so
|
||||
admission is decided on the real tool (as the executors' unwrap does). An unparseable
|
||||
bridge call is returned unchanged: it stays a sequential barrier and fails at dispatch."""
|
||||
try:
|
||||
from tools.tool_search import TOOL_CALL_NAME, resolve_underlying_call
|
||||
if tool_name != TOOL_CALL_NAME:
|
||||
return tool_name, function_args
|
||||
underlying, underlying_args, err = resolve_underlying_call(function_args)
|
||||
if err is not None or not underlying:
|
||||
return tool_name, function_args
|
||||
return underlying, underlying_args
|
||||
if tool_name == TOOL_CALL_NAME:
|
||||
underlying, underlying_args, err = resolve_underlying_call(function_args)
|
||||
if err is None and underlying:
|
||||
return underlying, underlying_args
|
||||
except Exception:
|
||||
return tool_name, function_args
|
||||
pass
|
||||
return tool_name, function_args
|
||||
|
||||
|
||||
def _batch_admission(tool_call, execution_cwd: Optional[Path]) -> tuple[str, List[Path], bool] | None:
|
||||
@@ -121,16 +112,11 @@ def _batch_admission(tool_call, execution_cwd: Optional[Path]) -> tuple[str, Lis
|
||||
_raw = tool_call.function.arguments
|
||||
logging.debug(
|
||||
"Could not parse args for %s — treating as sequential barrier; raw=%s",
|
||||
tool_name,
|
||||
_raw[:200] if isinstance(_raw, str) else repr(_raw)[:200],
|
||||
tool_name, _raw[:200] if isinstance(_raw, str) else repr(_raw)[:200],
|
||||
)
|
||||
return None
|
||||
if not isinstance(function_args, dict):
|
||||
logging.debug(
|
||||
"Non-dict args for %s (%s) — treating as sequential barrier",
|
||||
tool_name,
|
||||
type(function_args).__name__,
|
||||
)
|
||||
logging.debug("Non-dict args for %s (%s) — treating as sequential barrier", tool_name, type(function_args).__name__)
|
||||
return None
|
||||
|
||||
name, args = _peel_bridge_call(tool_name, function_args)
|
||||
@@ -147,17 +133,13 @@ def _batch_admission(tool_call, execution_cwd: Optional[Path]) -> tuple[str, Lis
|
||||
def _plan_tool_batch_segments(tool_calls, *, execution_cwd: Optional[Path] = None) -> List[tuple]:
|
||||
"""Split a tool-call batch into ordered ``("parallel"|"sequential", calls)`` segments.
|
||||
|
||||
Segments preserve the model's call order exactly — a later call never
|
||||
crosses an earlier barrier — so result ordering and side-effect boundaries
|
||||
match fully-sequential execution. Barriers: ``_NEVER_PARALLEL_TOOLS``,
|
||||
unparseable/non-dict args, and anything not parallel-safe (built-in list,
|
||||
bridge lookups, opted-in MCP tools). Path-scoped tools join a run only when
|
||||
their paths don't conflict with the run's reservations: reader↔reader
|
||||
overlap commutes and stays parallel; any overlap involving a writer closes
|
||||
the run so the call starts a NEW run after the conflicting one lands.
|
||||
``search_files`` reserves its root (default ``.``) as a reader. Parallel
|
||||
runs shorter than two calls demote to sequential (the sequential executor
|
||||
owns the richer inline dispatch); adjacent sequential segments merge.
|
||||
Call order is preserved exactly (a later call never crosses an earlier barrier), so
|
||||
result order and side-effect boundaries match fully-sequential execution. Barriers:
|
||||
``_NEVER_PARALLEL_TOOLS``, unparseable/non-dict args, anything not parallel-safe.
|
||||
Path-scoped tools join a run only if they don't conflict with its reservations:
|
||||
reader↔reader overlap stays parallel; any overlap involving a writer closes the run so
|
||||
the call starts a NEW run after the conflicting one lands. Runs shorter than two calls
|
||||
demote to sequential (it owns the richer inline dispatch); adjacent sequential merge.
|
||||
"""
|
||||
segments: List[tuple] = []
|
||||
current: list = []
|
||||
@@ -207,9 +189,8 @@ def _should_parallelize_tool_batch(tool_calls) -> bool:
|
||||
|
||||
|
||||
def _canonical_path(raw_path: str, execution_cwd: Optional[Path] = None) -> Path:
|
||||
"""Canonical, OS-aware path for overlap detection (realpath for symlinks on
|
||||
existing components, normcase for case-insensitive platforms); relative
|
||||
paths resolve against *execution_cwd* or ``Path.cwd()``."""
|
||||
"""Canonical, OS-aware path for overlap detection (realpath + normcase); relative paths
|
||||
resolve against *execution_cwd* or ``Path.cwd()``."""
|
||||
expanded = Path(raw_path).expanduser()
|
||||
base = execution_cwd if execution_cwd is not None else Path.cwd()
|
||||
candidate = expanded if expanded.is_absolute() else base / expanded
|
||||
@@ -221,13 +202,9 @@ def _extract_parallel_scope_paths(
|
||||
function_args: dict,
|
||||
execution_cwd: Optional[Path] = None,
|
||||
) -> List[Path]:
|
||||
"""Every canonical path this call reserves for overlap checks.
|
||||
|
||||
*execution_cwd* should be the cwd the tool will actually use (may differ
|
||||
from the process cwd on WSL / sandboxed backends). For V4A ``patch`` the
|
||||
scope comes from patch-body file headers. An empty result means the scope
|
||||
is unknown and the planner must treat the call as a sequential barrier.
|
||||
"""
|
||||
"""Every canonical path this call reserves for overlap checks. *execution_cwd* is the cwd
|
||||
the tool will actually use (may differ from the process cwd on WSL / sandboxed backends);
|
||||
V4A ``patch`` scope comes from patch-body headers. Empty = unknown scope = barrier."""
|
||||
if tool_name not in _PATH_SCOPED_TOOLS:
|
||||
return []
|
||||
|
||||
@@ -239,9 +216,7 @@ def _extract_parallel_scope_paths(
|
||||
if isinstance(raw_path, str) and raw_path.strip():
|
||||
raw_paths.append(raw_path)
|
||||
elif tool_name == "search_files":
|
||||
# search_files defaults its root to the cwd; reserve that rather than
|
||||
# demoting every bare search to a barrier.
|
||||
raw_paths.append(".")
|
||||
raw_paths.append(".") # its default root; reserving it beats demoting every bare search
|
||||
|
||||
# dict.fromkeys dedupes while preserving first-seen order.
|
||||
return list(dict.fromkeys(
|
||||
@@ -261,23 +236,18 @@ def _extract_parallel_scope_path(
|
||||
|
||||
|
||||
def _paths_overlap(left: Path, right: Path) -> bool:
|
||||
"""True when two already-canonical paths may refer to the same subtree."""
|
||||
left_parts = left.parts
|
||||
right_parts = right.parts
|
||||
"""True when two already-canonical paths may refer to the same subtree (an empty path
|
||||
overlaps nothing)."""
|
||||
left_parts, right_parts = left.parts, right.parts
|
||||
if not left_parts or not right_parts:
|
||||
# Empty paths are guarded upstream; only two non-empty equal prefixes overlap.
|
||||
return bool(left_parts) == bool(right_parts) and bool(left_parts)
|
||||
return False
|
||||
common_len = min(len(left_parts), len(right_parts))
|
||||
return left_parts[:common_len] == right_parts[:common_len]
|
||||
|
||||
|
||||
def _is_multimodal_tool_result(value: Any) -> bool:
|
||||
"""True for the multimodal envelope: dict with ``_multimodal=True`` and a ``content`` list."""
|
||||
return (
|
||||
isinstance(value, dict)
|
||||
and value.get("_multimodal") is True
|
||||
and isinstance(value.get("content"), list)
|
||||
)
|
||||
return isinstance(value, dict) and value.get("_multimodal") is True and isinstance(value.get("content"), list)
|
||||
|
||||
|
||||
def _is_text_part(p: Any) -> bool:
|
||||
@@ -322,12 +292,8 @@ _V4A_MOVE_HEADER = re.compile(r'^\*\*\*\s*Move\s+File:\s*(.+?)\s*->\s*(.+)$', re
|
||||
|
||||
|
||||
def _extract_file_mutation_targets(tool_name: str, args: Dict[str, Any]) -> List[str]:
|
||||
"""File paths a ``write_file`` / ``patch`` call targets.
|
||||
|
||||
Replace mode uses ``args["path"]``; V4A patch mode parses the
|
||||
``*** Update/Add/Delete/Move File:`` headers so each file in a multi-file
|
||||
patch is tracked separately.
|
||||
"""
|
||||
"""File paths a ``write_file`` / ``patch`` call targets: ``args["path"]`` in replace mode,
|
||||
every ``*** Update/Add/Delete/Move File:`` header in V4A patch mode."""
|
||||
if tool_name not in _FILE_MUTATING_TOOLS:
|
||||
return []
|
||||
mode = "replace" if tool_name == "write_file" else (args.get("mode") or "replace")
|
||||
@@ -389,8 +355,8 @@ def _extract_error_preview(result: Any, max_len: int = 180) -> str:
|
||||
|
||||
|
||||
def _trajectory_normalize_msg(msg: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Shallow copy with image blobs stripped for trajectory saving: multimodal
|
||||
results become their text summary, image parts become ``[screenshot]``."""
|
||||
"""Shallow copy for trajectory saving: multimodal results become their text summary,
|
||||
image parts become ``[screenshot]``."""
|
||||
if not isinstance(msg, dict):
|
||||
return msg
|
||||
content = msg.get("content")
|
||||
@@ -420,19 +386,14 @@ def make_tool_result_message(
|
||||
*,
|
||||
effect_disposition: str | None = None,
|
||||
) -> dict:
|
||||
"""Build a tool-result message with the OpenAI ``name`` field (wire format)
|
||||
and the internal ``tool_name`` field (session DB).
|
||||
|
||||
Content from high-risk tools (web_extract, web_search, browser_*, mcp_*) is
|
||||
wrapped in untrusted-data delimiters — the architectural defense against
|
||||
indirect prompt injection; see ``_maybe_wrap_untrusted``.
|
||||
"""Build a tool-result message: OpenAI ``name`` (wire format) plus internal ``tool_name``
|
||||
(session DB). High-risk tool content (web_extract, web_search, browser_*, mcp_*) is
|
||||
wrapped in untrusted-data delimiters — the defense against indirect prompt injection.
|
||||
"""
|
||||
# Replay-recovery callers bypass the executor's canonical-id helper, so normalize here too.
|
||||
tool_call_id = _normalize_tool_call_id(tool_call_id)
|
||||
|
||||
# Order matters: detect elision on the RAW content and append the notice
|
||||
# first, THEN wrap, so the notice sits inside the untrusted block next to
|
||||
# the data it describes — once, at construction time (cache-safe).
|
||||
# Elision notice is appended to the RAW content first, THEN wrapped, so it sits inside
|
||||
# the untrusted block next to the data it describes — once, at construction (cache-safe).
|
||||
wrapped = _maybe_wrap_untrusted(name, _maybe_append_elision_notice(name, content))
|
||||
message = stamp_message_timestamp({
|
||||
"role": "tool",
|
||||
@@ -453,14 +414,12 @@ def make_tool_result_message(
|
||||
return message
|
||||
|
||||
|
||||
# Tools whose results carry attacker-controllable content. Short outputs
|
||||
# (under 32 chars) skip wrapping: the overhead outweighs any injection risk.
|
||||
# Tools whose results carry attacker-controllable content; outputs under 32 chars skip wrapping.
|
||||
_UNTRUSTED_TOOL_NAMES = frozenset({"web_extract", "web_search"})
|
||||
_UNTRUSTED_TOOL_PREFIXES = ("browser_", "mcp_")
|
||||
_UNTRUSTED_WRAP_MIN_CHARS = 32
|
||||
|
||||
# Case-insensitive so attacker content can't forge or prematurely close the
|
||||
# boundary with a differently-cased tag the model would still read as one.
|
||||
# Case-insensitive so a differently-cased tag can't forge or prematurely close the boundary.
|
||||
_DELIMITER_TOKEN_RE = re.compile(r"untrusted_tool_result", re.IGNORECASE)
|
||||
|
||||
|
||||
@@ -472,21 +431,16 @@ def _is_text_item(item: Any) -> bool:
|
||||
return _is_text_part(item) and isinstance(item.get("text"), str)
|
||||
|
||||
|
||||
# --- Upstream-elision detection ---
|
||||
# Some MCP servers elide data SERVER-SIDE and mark it inside the payload
|
||||
# ('...13 more items', '"has_more": true', 'saved to sandbox', 'data_preview'
|
||||
# envelopes). The result looks structurally complete, so models treat the
|
||||
# visible slice as the whole dataset. Conservative, explicit markers only —
|
||||
# not a generic truncation heuristic. One notice is appended at construction
|
||||
# time, never mutated later (prompt-cache safe).
|
||||
# Some MCP servers elide data SERVER-SIDE and mark it inside a structurally complete payload,
|
||||
# so models treat the visible slice as the whole dataset. Conservative explicit markers only —
|
||||
# not a generic truncation heuristic; the notice is appended once at construction (cache-safe).
|
||||
_UPSTREAM_ELISION_PATTERNS = (
|
||||
re.compile(r"\.\.\.\s*\d+\s+more\s+items?", re.IGNORECASE),
|
||||
re.compile(r'"has_more"\s*:\s*true', re.IGNORECASE),
|
||||
re.compile(r"saved to sandbox", re.IGNORECASE),
|
||||
re.compile(r"data_preview", re.IGNORECASE),
|
||||
)
|
||||
# Tiny results can't hide an elided enumeration; markers for the sizes that
|
||||
# matter (20-50K) are always inside the first 64KB.
|
||||
# Tiny results can't hide an elided enumeration; markers for the sizes that matter sit in the first 64KB.
|
||||
_ELISION_SCAN_MIN_CHARS = 1_000
|
||||
_ELISION_SCAN_MAX_CHARS = 65_536
|
||||
|
||||
@@ -513,49 +467,39 @@ def _maybe_append_elision_notice(name: str, content: Any) -> Any:
|
||||
|
||||
|
||||
def _tool_output_risk_metadata(name: str, content: Any) -> Optional[Dict[str, Any]]:
|
||||
"""Internal-only advisory classification of attacker-controlled output.
|
||||
|
||||
Records deterministic finding ids, never blocks or redacts, and omits the scanned text.
|
||||
"""
|
||||
"""Internal-only advisory classification of attacker-controlled output: deterministic
|
||||
finding ids, never blocks or redacts, omits the scanned text."""
|
||||
if not _is_untrusted_tool(name):
|
||||
return None
|
||||
if isinstance(content, str):
|
||||
text_parts = [content]
|
||||
elif isinstance(content, list):
|
||||
text_parts = [item["text"] for item in content if _is_text_item(item)]
|
||||
if not text_parts:
|
||||
return None
|
||||
else:
|
||||
return None
|
||||
if not text_parts:
|
||||
return None
|
||||
|
||||
findings: List[str] = []
|
||||
for text in text_parts:
|
||||
for finding in scan_for_threats(text, scope="context"):
|
||||
if finding not in findings:
|
||||
findings.append(finding)
|
||||
return {
|
||||
"risk": "high" if findings else "low",
|
||||
"findings": findings,
|
||||
"redacted": False,
|
||||
}
|
||||
return {"risk": "high" if findings else "low", "findings": findings, "redacted": False}
|
||||
|
||||
|
||||
def _neutralize_delimiters(content: str) -> str:
|
||||
"""Defang embedded ``untrusted_tool_result`` tokens so poisoned content
|
||||
can't close the trust boundary early (hyphens keep it readable but non-matching)."""
|
||||
"""Defang embedded ``untrusted_tool_result`` tokens so poisoned content can't close the
|
||||
trust boundary early (hyphens keep it readable but non-matching)."""
|
||||
return _DELIMITER_TOKEN_RE.sub("untrusted-tool-result", content)
|
||||
|
||||
|
||||
def _maybe_wrap_untrusted(name: str, content: Any) -> Any:
|
||||
"""Wrap high-risk tool content in untrusted-data delimiters.
|
||||
|
||||
Strings are neutralized and wrapped in exactly one block; text parts of a
|
||||
multimodal list are wrapped individually (non-text parts preserved, outer
|
||||
list rebuilt — compare by value, not ``is``). Unchanged when the tool is
|
||||
not high-risk, the content is neither str nor list, or a string is too
|
||||
short. There is deliberately no "already wrapped" fast-path: it would be
|
||||
attacker-forgeable, so harmless re-wrapping is the safe choice.
|
||||
"""
|
||||
"""Wrap high-risk tool content in untrusted-data delimiters: strings are neutralized and
|
||||
wrapped in exactly one block; text parts of a multimodal list are wrapped individually
|
||||
(outer list rebuilt — compare by value, not ``is``). Unchanged for non-high-risk tools,
|
||||
non-str/list content, or short strings. Deliberately no "already wrapped" fast-path:
|
||||
it would be attacker-forgeable, so harmless re-wrapping is the safe choice."""
|
||||
if not _is_untrusted_tool(name):
|
||||
return content
|
||||
if isinstance(content, str):
|
||||
|
||||
Reference in New Issue
Block a user