diff --git a/agent/tool_dispatch_helpers.py b/agent/tool_dispatch_helpers.py index 54bedab1ae..fd56cb0477 100644 --- a/agent/tool_dispatch_helpers.py +++ b/agent/tool_dispatch_helpers.py @@ -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):