mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-11 22:48:38 +00:00
perf(sandbox): cache local-path masking patterns instead of recompiling per match (#3713)
mask_local_paths_in_output runs once per glob/grep match (tools.py:1540, :1625) and once per bash/ls output, and each call rebuilt every skills / ACP / user-data masking regex from scratch — Path.resolve() syscalls + re.escape + re.compile. A grep returning up to 100 matches recompiled the same patterns 100x. Compile the host->virtual patterns once per ordered source set via functools.lru_cache (keyed on the config-stable + per-thread (host, virtual) pairs, so it self-invalidates when they change) and collapse the three byte-identical replacer blocks into one applier. Behavior is unchanged: same patterns, same application order, same per-thread mappings. Verified behavior-preserving with an equivalence harness (original vs refactored logic over 54 inputs incl. slash-style variants, base-only matches, subpaths, overlapping mappings, missing skills/ACP, thread_data None) -> 0 mismatches. Added cache/consistency tests alongside the existing masking tests. Fixes #3712.
This commit is contained in:
parent
14d9bb87a4
commit
f7f2a500d1
@ -4,6 +4,7 @@ import posixpath
|
|||||||
import re
|
import re
|
||||||
import shlex
|
import shlex
|
||||||
from collections.abc import Callable
|
from collections.abc import Callable
|
||||||
|
from functools import lru_cache
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from langchain.tools import tool
|
from langchain.tools import tool
|
||||||
@ -554,76 +555,74 @@ def _thread_actual_to_virtual_mappings(thread_data: ThreadDataState) -> dict[str
|
|||||||
return {actual: virtual for virtual, actual in _thread_virtual_to_actual_mappings(thread_data).items()}
|
return {actual: virtual for virtual, actual in _thread_virtual_to_actual_mappings(thread_data).items()}
|
||||||
|
|
||||||
|
|
||||||
|
@lru_cache(maxsize=512)
|
||||||
|
def _compiled_mask_patterns(sources: tuple[tuple[str, str], ...]) -> tuple[tuple[re.Pattern[str], str, str], ...]:
|
||||||
|
"""Compile the host→virtual masking patterns once per source set.
|
||||||
|
|
||||||
|
``sources`` is an ordered tuple of ``(host_base, virtual_base)`` pairs
|
||||||
|
(skills, then ACP workspace, then per-thread user-data mappings sorted by
|
||||||
|
host-path length, longest first). The patterns derive only from
|
||||||
|
config-stable + per-thread inputs, so they're cached and reused instead of
|
||||||
|
being rebuilt — ``re.escape`` + ``re.compile`` + ``Path.resolve`` (a
|
||||||
|
syscall) — on every call. ``mask_local_paths_in_output`` runs once per
|
||||||
|
glob/grep match, so without this the same patterns are recompiled per
|
||||||
|
match.
|
||||||
|
"""
|
||||||
|
compiled: list[tuple[re.Pattern[str], str, str]] = []
|
||||||
|
for host_base, virtual_base in sources:
|
||||||
|
seen: set[str] = set()
|
||||||
|
# Same base set as ``_path_variants(raw) | _path_variants(resolved)``;
|
||||||
|
# ordered deterministically so the cached tuple is stable (variants of
|
||||||
|
# one host map to the same virtual and don't overlap after substitution,
|
||||||
|
# so order within a source is irrelevant to the result).
|
||||||
|
for root in (str(Path(host_base)), str(Path(host_base).resolve())):
|
||||||
|
for variant in sorted(_path_variants(root)):
|
||||||
|
if variant in seen:
|
||||||
|
continue
|
||||||
|
seen.add(variant)
|
||||||
|
escaped = re.escape(variant).replace(r"\\", r"[/\\]")
|
||||||
|
compiled.append((re.compile(escaped + r"(?:[/\\][^\s\"';&|<>()]*)?"), variant, virtual_base))
|
||||||
|
return tuple(compiled)
|
||||||
|
|
||||||
|
|
||||||
def mask_local_paths_in_output(output: str, thread_data: ThreadDataState | None) -> str:
|
def mask_local_paths_in_output(output: str, thread_data: ThreadDataState | None) -> str:
|
||||||
"""Mask host absolute paths from local sandbox output using virtual paths.
|
"""Mask host absolute paths from local sandbox output using virtual paths.
|
||||||
|
|
||||||
Handles user-data paths (per-thread), skills paths, and ACP workspace paths (global).
|
Handles user-data paths (per-thread), skills paths, and ACP workspace paths (global).
|
||||||
"""
|
"""
|
||||||
result = output
|
# Build the ordered (host_base, virtual_base) source list. Order is
|
||||||
|
# preserved from the original implementation: skills, then ACP workspace,
|
||||||
|
# then user-data mappings (longest host path first). Custom mount host
|
||||||
|
# paths are masked by LocalSandbox._reverse_resolve_paths_in_output().
|
||||||
|
sources: list[tuple[str, str]] = []
|
||||||
|
|
||||||
# Mask skills host paths
|
|
||||||
skills_host = _get_skills_host_path()
|
skills_host = _get_skills_host_path()
|
||||||
skills_container = _get_skills_container_path()
|
|
||||||
if skills_host:
|
if skills_host:
|
||||||
raw_base = str(Path(skills_host))
|
sources.append((skills_host, _get_skills_container_path()))
|
||||||
resolved_base = str(Path(skills_host).resolve())
|
|
||||||
for base in _path_variants(raw_base) | _path_variants(resolved_base):
|
|
||||||
escaped = re.escape(base).replace(r"\\", r"[/\\]")
|
|
||||||
pattern = re.compile(escaped + r"(?:[/\\][^\s\"';&|<>()]*)?")
|
|
||||||
|
|
||||||
def replace_skills(match: re.Match, _base: str = base) -> str:
|
acp_host = _get_acp_workspace_host_path(_extract_thread_id_from_thread_data(thread_data))
|
||||||
matched_path = match.group(0)
|
|
||||||
if matched_path == _base:
|
|
||||||
return skills_container
|
|
||||||
relative = matched_path[len(_base) :].lstrip("/\\")
|
|
||||||
return f"{skills_container}/{relative}" if relative else skills_container
|
|
||||||
|
|
||||||
result = pattern.sub(replace_skills, result)
|
|
||||||
|
|
||||||
# Mask ACP workspace host paths
|
|
||||||
_thread_id = _extract_thread_id_from_thread_data(thread_data)
|
|
||||||
acp_host = _get_acp_workspace_host_path(_thread_id)
|
|
||||||
if acp_host:
|
if acp_host:
|
||||||
raw_base = str(Path(acp_host))
|
sources.append((acp_host, _ACP_WORKSPACE_VIRTUAL_PATH))
|
||||||
resolved_base = str(Path(acp_host).resolve())
|
|
||||||
for base in _path_variants(raw_base) | _path_variants(resolved_base):
|
|
||||||
escaped = re.escape(base).replace(r"\\", r"[/\\]")
|
|
||||||
pattern = re.compile(escaped + r"(?:[/\\][^\s\"';&|<>()]*)?")
|
|
||||||
|
|
||||||
def replace_acp(match: re.Match, _base: str = base) -> str:
|
if thread_data is not None:
|
||||||
matched_path = match.group(0)
|
mappings = _thread_actual_to_virtual_mappings(thread_data)
|
||||||
if matched_path == _base:
|
for actual_base, virtual_base in sorted(mappings.items(), key=lambda item: len(item[0]), reverse=True):
|
||||||
return _ACP_WORKSPACE_VIRTUAL_PATH
|
sources.append((actual_base, virtual_base))
|
||||||
relative = matched_path[len(_base) :].lstrip("/\\")
|
|
||||||
return f"{_ACP_WORKSPACE_VIRTUAL_PATH}/{relative}" if relative else _ACP_WORKSPACE_VIRTUAL_PATH
|
|
||||||
|
|
||||||
result = pattern.sub(replace_acp, result)
|
if not sources:
|
||||||
|
return output
|
||||||
|
|
||||||
# Custom mount host paths are masked by LocalSandbox._reverse_resolve_paths_in_output()
|
result = output
|
||||||
|
for pattern, base, virtual in _compiled_mask_patterns(tuple(sources)):
|
||||||
|
|
||||||
# Mask user-data host paths
|
def replace_match(match: re.Match, _base: str = base, _virtual: str = virtual) -> str:
|
||||||
if thread_data is None:
|
matched_path = match.group(0)
|
||||||
return result
|
if matched_path == _base:
|
||||||
|
return _virtual
|
||||||
|
relative = matched_path[len(_base) :].lstrip("/\\")
|
||||||
|
return f"{_virtual}/{relative}" if relative else _virtual
|
||||||
|
|
||||||
mappings = _thread_actual_to_virtual_mappings(thread_data)
|
result = pattern.sub(replace_match, result)
|
||||||
if not mappings:
|
|
||||||
return result
|
|
||||||
|
|
||||||
for actual_base, virtual_base in sorted(mappings.items(), key=lambda item: len(item[0]), reverse=True):
|
|
||||||
raw_base = str(Path(actual_base))
|
|
||||||
resolved_base = str(Path(actual_base).resolve())
|
|
||||||
for base in _path_variants(raw_base) | _path_variants(resolved_base):
|
|
||||||
escaped_actual = re.escape(base).replace(r"\\", r"[/\\]")
|
|
||||||
pattern = re.compile(escaped_actual + r"(?:[/\\][^\s\"';&|<>()]*)?")
|
|
||||||
|
|
||||||
def replace_match(match: re.Match, _base: str = base, _virtual: str = virtual_base) -> str:
|
|
||||||
matched_path = match.group(0)
|
|
||||||
if matched_path == _base:
|
|
||||||
return _virtual
|
|
||||||
relative = matched_path[len(_base) :].lstrip("/\\")
|
|
||||||
return f"{_virtual}/{relative}" if relative else _virtual
|
|
||||||
|
|
||||||
result = pattern.sub(replace_match, result)
|
|
||||||
|
|
||||||
return result
|
return result
|
||||||
|
|
||||||
|
|||||||
@ -9,6 +9,7 @@ from deerflow.sandbox.exceptions import SandboxError
|
|||||||
from deerflow.sandbox.tools import (
|
from deerflow.sandbox.tools import (
|
||||||
VIRTUAL_PATH_PREFIX,
|
VIRTUAL_PATH_PREFIX,
|
||||||
_apply_cwd_prefix,
|
_apply_cwd_prefix,
|
||||||
|
_compiled_mask_patterns,
|
||||||
_get_custom_mount_for_path,
|
_get_custom_mount_for_path,
|
||||||
_get_custom_mounts,
|
_get_custom_mounts,
|
||||||
_is_acp_workspace_path,
|
_is_acp_workspace_path,
|
||||||
@ -114,6 +115,40 @@ def test_mask_local_paths_in_output_hides_skills_host_paths() -> None:
|
|||||||
assert "/mnt/skills/public/bootstrap/SKILL.md" in masked
|
assert "/mnt/skills/public/bootstrap/SKILL.md" in masked
|
||||||
|
|
||||||
|
|
||||||
|
def test_mask_local_paths_compiled_patterns_are_cached() -> None:
|
||||||
|
"""The compiled patterns for a given source set are built once and reused
|
||||||
|
(mask runs once per glob/grep match, so this avoids per-match recompiles)."""
|
||||||
|
sources = (("/tmp/deer-flow/threads/t1/user-data/workspace", "/mnt/user-data/workspace"),)
|
||||||
|
first = _compiled_mask_patterns(sources)
|
||||||
|
second = _compiled_mask_patterns(sources)
|
||||||
|
assert first is second # cache hit -> identical object, not rebuilt
|
||||||
|
|
||||||
|
|
||||||
|
def test_mask_local_paths_stable_across_repeated_and_batched_calls() -> None:
|
||||||
|
"""Masking is identical whether applied once or repeatedly (per-match path)."""
|
||||||
|
output = "a /tmp/deer-flow/threads/t1/user-data/workspace/x.txt and /tmp/deer-flow/threads/t1/user-data/outputs/y.log"
|
||||||
|
once = mask_local_paths_in_output(output, _THREAD_DATA)
|
||||||
|
twice = mask_local_paths_in_output(once, _THREAD_DATA)
|
||||||
|
assert "/tmp/deer-flow/threads/t1/user-data" not in once
|
||||||
|
assert "/mnt/user-data/workspace/x.txt" in once
|
||||||
|
assert "/mnt/user-data/outputs/y.log" in once
|
||||||
|
# Re-masking already-masked output leaves it unchanged (no host paths left).
|
||||||
|
assert twice == once
|
||||||
|
# Mapping outputs one-by-one matches masking each independently.
|
||||||
|
assert [mask_local_paths_in_output(o, _THREAD_DATA) for o in (output, output)] == [once, once]
|
||||||
|
|
||||||
|
|
||||||
|
def test_mask_local_paths_no_thread_data_still_masks_skills() -> None:
|
||||||
|
"""With thread_data=None, skills host paths are still masked (user-data skipped)."""
|
||||||
|
with (
|
||||||
|
patch("deerflow.sandbox.tools._get_skills_container_path", return_value="/mnt/skills"),
|
||||||
|
patch("deerflow.sandbox.tools._get_skills_host_path", return_value="/home/user/deer-flow/skills"),
|
||||||
|
):
|
||||||
|
masked = mask_local_paths_in_output("Reading: /home/user/deer-flow/skills/a/b.md", None)
|
||||||
|
assert "/home/user/deer-flow/skills" not in masked
|
||||||
|
assert "/mnt/skills/a/b.md" in masked
|
||||||
|
|
||||||
|
|
||||||
# ---------- _reject_path_traversal ----------
|
# ---------- _reject_path_traversal ----------
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user