wuyang/zk-data-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Enforce session workspace boundaries

Browse Source
This commit is contained in:
武阳
2026-05-08 17:09:07 +08:00
parent 5b14f55736
commit 0bbba2b936
13 changed files with 509 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/agent_tools.py
+213 -19
View File
@@ -1079,6 +1079,53 @@ def _snapshot_text(text: str, limit: int = 240) -> str:
return normalized[: limit - 3] + '...'
_GREP_DEFAULT_SKIPPED_DIRS = {
'.git',
'.hg',
'.mypy_cache',
'.next',
'.port_sessions',
'.pytest_cache',
'.ruff_cache',
'.svn',
'.venv',
'__pycache__',
'build',
'coverage',
'dist',
'node_modules',
'router_session_parquet',
'venv',
}
_GREP_BINARY_SUFFIXES = {
'.7z',
'.avif',
'.bin',
'.doc',
'.docx',
'.gz',
'.jpeg',
'.jpg',
'.parquet',
'.pdf',
'.png',
'.pyc',
'.snappy',
'.tar',
'.webp',
'.xls',
'.xlsx',
'.zip',
}
_GREP_MAX_LINE_CHARS = 800
_PLATFORM_READONLY_DIRS = {'src', 'backend', 'frontend', 'scripts'}
_PLATFORM_ROOT_MARKERS = (
Path('src') / 'agent_tools.py',
Path('backend') / 'api' / 'server.py',
Path('frontend') / 'app',
)
def _require_string(arguments: dict[str, Any], key: str) -> str:
value = arguments.get(key)
if not isinstance(value, str) or not value:
@@ -1543,19 +1590,90 @@ def _data_agent_session_output_root(context: ToolExecutionContext) -> Path:
return context.root / '.port_sessions' / 'data_agent_output'
def _resolve_path(raw_path: str, context: ToolExecutionContext, *, allow_missing: bool = True) -> Path:
def _resolve_path(
raw_path: str,
context: ToolExecutionContext,
*,
allow_missing: bool = True,
allow_outside_root: bool = False,
) -> Path:
expanded = Path(raw_path).expanduser()
candidate = expanded if expanded.is_absolute() else context.root / expanded
session_candidate = _session_logical_path(expanded, context)
candidate = (
expanded
if expanded.is_absolute()
else session_candidate
if session_candidate is not None
else context.root / expanded
)
resolved = candidate.resolve(strict=not allow_missing)
try:
resolved.relative_to(context.root)
except ValueError as exc:
if allow_outside_root:
return resolved
raise ToolExecutionError(
f'Path {raw_path!r} escapes the workspace root {context.root}'
) from exc
return resolved
def _session_logical_path(path: Path, context: ToolExecutionContext) -> Path | None:
"""把 output/scratchpad/input 这类逻辑路径路由到当前 session。"""
if path.is_absolute() or context.scratchpad_directory is None:
return None
if not path.parts:
return None
session_root = context.scratchpad_directory.parent
head, *tail = path.parts
tail_path = Path(*tail) if tail else Path()
if head in {'output', 'outputs'}:
return session_root / 'output' / tail_path
if head in {'scratchpad', 'scratch'}:
return context.scratchpad_directory / tail_path
if head in {'input', 'inputs'}:
return session_root / 'input' / tail_path
return None
def _display_path(path: Path, context: ToolExecutionContext) -> str:
resolved = path.resolve()
try:
return resolved.relative_to(context.root.resolve()).as_posix()
except ValueError:
return resolved.as_posix()
def _execution_cwd(context: ToolExecutionContext) -> Path:
if context.scratchpad_directory is not None:
context.scratchpad_directory.mkdir(parents=True, exist_ok=True)
return context.scratchpad_directory
return context.root
def _is_platform_app_root(root: Path) -> bool:
return all((root / marker).exists() for marker in _PLATFORM_ROOT_MARKERS)
def _is_platform_code_path(path: Path, context: ToolExecutionContext) -> bool:
if not _is_platform_app_root(context.root):
return False
try:
rel = path.resolve().relative_to(context.root.resolve())
except ValueError:
return False
return bool(rel.parts) and rel.parts[0] in _PLATFORM_READONLY_DIRS
def _ensure_not_platform_code_write(path: Path, context: ToolExecutionContext) -> None:
if _is_platform_code_path(path, context):
raise ToolPermissionError(
'Platform code is read-only in data-agent sessions. '
'Do not modify src, backend, frontend, or scripts from this agent.'
)
def _ensure_write_allowed(context: ToolExecutionContext) -> None:
if not context.permissions.allow_file_write:
raise ToolPermissionError(
@@ -1572,6 +1690,11 @@ def _ensure_process_execution_allowed(context: ToolExecutionContext, tool_label:
def _ensure_shell_allowed(command: str, context: ToolExecutionContext) -> None:
_ensure_process_execution_allowed(context, 'Shell commands')
if _looks_like_platform_code_shell_write(command, context):
raise ToolPermissionError(
'Shell command appears to write platform code. '
'Platform code directories are read-only in data-agent sessions.'
)
if context.permissions.allow_destructive_shell_commands:
return
destructive_patterns = [
@@ -1594,12 +1717,34 @@ def _ensure_shell_allowed(command: str, context: ToolExecutionContext) -> None:
)
def _looks_like_platform_code_shell_write(
command: str,
context: ToolExecutionContext,
) -> bool:
if not _is_platform_app_root(context.root):
return False
lowered = command.lower()
write_marker = re.search(
r'>|(^|[\s;&|])(tee|touch|python|python3|sed\s+-i|perl\s+-pi)\b',
lowered,
)
if write_marker is None:
return False
platform_fragments = [
f'{name}/' for name in _PLATFORM_READONLY_DIRS
] + [
f'{context.root.as_posix().lower()}/{name}/'
for name in _PLATFORM_READONLY_DIRS
]
return any(fragment in lowered for fragment in platform_fragments)
def _list_dir(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
raw_path = arguments.get('path', '.')
if not isinstance(raw_path, str):
raise ToolExecutionError('path must be a string')
max_entries = _coerce_int(arguments, 'max_entries', 200)
target = _resolve_path(raw_path, context)
target = _resolve_path(raw_path, context, allow_outside_root=True)
if not target.exists():
raise ToolExecutionError(f'Path not found: {raw_path}')
if not target.is_dir():
@@ -1608,7 +1753,7 @@ def _list_dir(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
lines: list[str] = []
for entry in entries[:max_entries]:
kind = 'dir' if entry.is_dir() else 'file'
rel = entry.relative_to(context.root)
rel = _display_path(entry, context)
lines.append(f'{kind}\t{rel}')
if len(entries) > max_entries:
lines.append(f'... truncated at {max_entries} entries ...')
@@ -1616,7 +1761,12 @@ def _list_dir(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
def _read_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
target = _resolve_path(_require_string(arguments, 'path'), context, allow_missing=False)
target = _resolve_path(
_require_string(arguments, 'path'),
context,
allow_missing=False,
allow_outside_root=True,
)
if not target.is_file():
raise ToolExecutionError(f'Path is not a file: {target}')
text = target.read_text(encoding='utf-8', errors='replace')
@@ -1639,6 +1789,7 @@ def _read_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
def _write_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
_ensure_write_allowed(context)
target = _resolve_path(_require_string(arguments, 'path'), context)
_ensure_not_platform_code_write(target, context)
content = arguments.get('content')
if not isinstance(content, str):
raise ToolExecutionError('content must be a string')
@@ -1649,13 +1800,13 @@ def _write_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str
previous_sha256 = hashlib.sha256(previous_text.encode('utf-8')).hexdigest()
target.parent.mkdir(parents=True, exist_ok=True)
target.write_text(content, encoding='utf-8')
rel = target.relative_to(context.root)
rel = _display_path(target, context)
new_sha256 = hashlib.sha256(content.encode('utf-8')).hexdigest()
return (
f'wrote {rel} ({len(content)} chars)',
{
'action': 'write_file',
'path': str(rel),
'path': rel,
'before_exists': previous_text is not None,
'before_sha256': previous_sha256,
'before_size': len(previous_text) if previous_text is not None else 0,
@@ -1675,6 +1826,7 @@ def _write_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str
def _edit_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
_ensure_write_allowed(context)
target = _resolve_path(_require_string(arguments, 'path'), context, allow_missing=False)
_ensure_not_platform_code_write(target, context)
if not target.is_file():
raise ToolExecutionError(f'Path is not a file: {target}')
old_text = arguments.get('old_text')
@@ -1697,14 +1849,14 @@ def _edit_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
before_sha256 = hashlib.sha256(current.encode('utf-8')).hexdigest()
updated = current.replace(old_text, new_text) if replace_all else current.replace(old_text, new_text, 1)
target.write_text(updated, encoding='utf-8')
rel = target.relative_to(context.root)
rel = _display_path(target, context)
replaced = occurrences if replace_all else 1
after_sha256 = hashlib.sha256(updated.encode('utf-8')).hexdigest()
return (
f'edited {rel}; replaced {replaced} occurrence(s)',
{
'action': 'edit_file',
'path': str(rel),
'path': rel,
'before_sha256': before_sha256,
'after_sha256': after_sha256,
'before_size': len(current),
@@ -1721,6 +1873,7 @@ def _edit_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
def _notebook_edit(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
_ensure_write_allowed(context)
target = _resolve_path(_require_string(arguments, 'path'), context, allow_missing=False)
_ensure_not_platform_code_write(target, context)
if target.suffix != '.ipynb':
raise ToolExecutionError('notebook_edit requires a .ipynb target')
if not target.is_file():
@@ -1781,12 +1934,12 @@ def _notebook_edit(arguments: dict[str, Any], context: ToolExecutionContext) ->
updated = json.dumps(notebook, ensure_ascii=True, indent=1) + '\n'
target.write_text(updated, encoding='utf-8')
after_sha256 = hashlib.sha256(updated.encode('utf-8')).hexdigest()
rel = target.relative_to(context.root)
rel = _display_path(target, context)
return (
f'updated notebook cell {cell_index} in {rel}',
{
'action': 'notebook_edit',
'path': str(rel),
'path': rel,
'cell_index': cell_index,
'cell_type': cell['cell_type'],
'before_sha256': before_sha256,
@@ -1824,7 +1977,7 @@ def _grep_search(arguments: dict[str, Any], context: ToolExecutionContext) -> st
if not isinstance(literal, bool):
raise ToolExecutionError('literal must be a boolean')
max_matches = _coerce_int(arguments, 'max_matches', 100)
root = _resolve_path(raw_path, context)
root = _resolve_path(raw_path, context, allow_outside_root=True)
if not root.exists():
raise ToolExecutionError(f'Path not found: {raw_path}')
try:
@@ -1833,20 +1986,61 @@ def _grep_search(arguments: dict[str, Any], context: ToolExecutionContext) -> st
raise ToolExecutionError(f'Invalid regex pattern: {exc}') from exc
hits: list[str] = []
file_iter = root.rglob('*') if root.is_dir() else [root]
root_parts = _relative_parts(root, context)
for file_path in file_iter:
if not file_path.is_file():
continue
if _grep_should_skip_path(file_path, context, root_parts):
continue
try:
text = file_path.read_text(encoding='utf-8', errors='replace')
except OSError:
continue
for line_no, line in enumerate(text.splitlines(), start=1):
if regex.search(line):
rel = file_path.relative_to(context.root)
hits.append(f'{rel}:{line_no}: {line}')
rel = _display_path(file_path, context)
hits.append(f'{rel}:{line_no}: {_truncate_grep_line(line)}')
if len(hits) >= max_matches:
return '\n'.join(hits + [f'... truncated at {max_matches} matches ...'])
return '\n'.join(hits) if hits else '(no matches)'
return _truncate_output(
'\n'.join(
hits
+ [f'... truncated at {max_matches} matches ...']
),
context.max_output_chars,
)
if not hits:
return '(no matches)'
return _truncate_output('\n'.join(hits), context.max_output_chars)
def _relative_parts(path: Path, context: ToolExecutionContext) -> tuple[str, ...]:
try:
return path.resolve().relative_to(context.root.resolve()).parts
except ValueError:
return ()
def _grep_should_skip_path(
file_path: Path,
context: ToolExecutionContext,
root_parts: tuple[str, ...],
) -> bool:
rel_parts = _relative_parts(file_path, context)
if file_path.suffix.lower() in _GREP_BINARY_SUFFIXES:
return True
explicit_roots = set(root_parts)
path_parts = rel_parts[:-1] if rel_parts else file_path.parts[:-1]
for part in path_parts:
if part in _GREP_DEFAULT_SKIPPED_DIRS and part not in explicit_roots:
return True
return False
def _truncate_grep_line(line: str) -> str:
if len(line) <= _GREP_MAX_LINE_CHARS:
return line
omitted = len(line) - _GREP_MAX_LINE_CHARS
return f'{line[:_GREP_MAX_LINE_CHARS]}... [line truncated, {omitted} chars omitted]'
def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
@@ -1887,7 +2081,7 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) -
try:
process = subprocess.Popen(
command,
cwd=context.root,
cwd=_execution_cwd(context),
stdin=subprocess.PIPE if stdin else None,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
@@ -2045,7 +2239,7 @@ def _run_python_package(arguments: dict[str, Any], context: ToolExecutionContext
try:
completed = subprocess.run(
command,
cwd=context.root,
cwd=_execution_cwd(context),
capture_output=True,
text=True,
timeout=timeout_seconds,
@@ -2130,7 +2324,7 @@ def _run_bash(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
command,
shell=True,
executable='/bin/bash',
cwd=context.root,
cwd=_execution_cwd(context),
capture_output=True,
text=True,
timeout=context.command_timeout_seconds,