Add Jupyter workspace runtime
This commit is contained in:
wuyang6
+201
-2
@@ -1683,6 +1683,8 @@ def _display_path(path: Path, context: ToolExecutionContext) -> str:
|
||||
|
||||
|
||||
def _execution_cwd(context: ToolExecutionContext) -> Path:
|
||||
if context.jupyter_runtime is not None:
|
||||
return Path(context.jupyter_runtime.binding.workspace_cwd)
|
||||
if context.scratchpad_directory is not None:
|
||||
context.scratchpad_directory.mkdir(parents=True, exist_ok=True)
|
||||
return context.scratchpad_directory
|
||||
@@ -1781,6 +1783,12 @@ def _list_dir(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
|
||||
if not isinstance(raw_path, str):
|
||||
raise ToolExecutionError('path must be a string')
|
||||
max_entries = _coerce_int(arguments, 'max_entries', 200)
|
||||
if context.jupyter_runtime is not None:
|
||||
return context.jupyter_runtime.list_dir(
|
||||
raw_path,
|
||||
max_entries=max_entries,
|
||||
max_output_chars=context.max_output_chars,
|
||||
)
|
||||
target = _resolve_path(raw_path, context, allow_outside_root=True)
|
||||
if not target.exists():
|
||||
raise ToolExecutionError(f'Path not found: {raw_path}')
|
||||
@@ -1798,6 +1806,23 @@ def _list_dir(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
|
||||
|
||||
|
||||
def _read_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
|
||||
if context.jupyter_runtime is not None:
|
||||
start_line = arguments.get('start_line')
|
||||
end_line = arguments.get('end_line')
|
||||
if start_line is not None and (
|
||||
isinstance(start_line, bool) or not isinstance(start_line, int) or start_line < 1
|
||||
):
|
||||
raise ToolExecutionError('start_line must be an integer >= 1')
|
||||
if end_line is not None and (
|
||||
isinstance(end_line, bool) or not isinstance(end_line, int) or end_line < 1
|
||||
):
|
||||
raise ToolExecutionError('end_line must be an integer >= 1')
|
||||
return context.jupyter_runtime.read_text(
|
||||
_require_string(arguments, 'path'),
|
||||
start_line=start_line,
|
||||
end_line=end_line,
|
||||
max_output_chars=context.max_output_chars,
|
||||
)
|
||||
target = _resolve_path(
|
||||
_require_string(arguments, 'path'),
|
||||
context,
|
||||
@@ -1825,8 +1850,6 @@ 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 = _coerce_write_file_content(arguments)
|
||||
append = arguments.get('append', False)
|
||||
if not isinstance(append, bool):
|
||||
@@ -1834,6 +1857,25 @@ def _write_file(arguments: dict[str, Any], context: ToolExecutionContext) -> str
|
||||
newline_at_end = arguments.get('newline_at_end', False)
|
||||
if not isinstance(newline_at_end, bool):
|
||||
raise ToolExecutionError('newline_at_end must be a boolean')
|
||||
if context.jupyter_runtime is not None:
|
||||
message = context.jupyter_runtime.write_text(
|
||||
_require_string(arguments, 'path'),
|
||||
content,
|
||||
append=append,
|
||||
newline_at_end=newline_at_end,
|
||||
max_output_chars=context.max_output_chars,
|
||||
)
|
||||
return (
|
||||
message,
|
||||
{
|
||||
'action': 'remote_append_file' if append else 'remote_write_file',
|
||||
'path': _require_string(arguments, 'path'),
|
||||
'remote_workspace': context.jupyter_runtime.binding.workspace_cwd,
|
||||
'content_length': len(content),
|
||||
},
|
||||
)
|
||||
target = _resolve_path(_require_string(arguments, 'path'), context)
|
||||
_ensure_not_platform_code_write(target, context)
|
||||
if newline_at_end and content and not content.endswith('\n'):
|
||||
content += '\n'
|
||||
previous_text: str | None = None
|
||||
@@ -2068,6 +2110,34 @@ def _notebook_edit(arguments: dict[str, Any], context: ToolExecutionContext) ->
|
||||
|
||||
def _glob_search(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
|
||||
pattern = _require_string(arguments, 'pattern')
|
||||
if context.jupyter_runtime is not None:
|
||||
code = r'''
|
||||
from pathlib import Path
|
||||
import glob
|
||||
pattern = PATTERN
|
||||
workspace = WORKSPACE
|
||||
if pattern.startswith("/"):
|
||||
matches = sorted(glob.glob(pattern, recursive=True))
|
||||
else:
|
||||
matches = sorted(glob.glob(str(Path(workspace) / pattern), recursive=True))
|
||||
print("\n".join(matches) if matches else "(no matches)")
|
||||
'''
|
||||
script = (
|
||||
code.replace('PATTERN', repr(pattern))
|
||||
.replace('WORKSPACE', repr(context.jupyter_runtime.binding.workspace_cwd))
|
||||
)
|
||||
result = context.jupyter_runtime.run_python(
|
||||
code=script,
|
||||
script_path=None,
|
||||
args=[],
|
||||
stdin=None,
|
||||
timeout_seconds=20.0,
|
||||
max_output_chars=context.max_output_chars,
|
||||
cancel_event=context.cancel_event,
|
||||
)
|
||||
if result.exit_code != 0:
|
||||
raise ToolExecutionError(result.stdout.strip() or 'remote glob_search failed')
|
||||
return result.stdout.strip() or '(no matches)'
|
||||
matches = sorted(context.root.glob(pattern))
|
||||
if not matches:
|
||||
return '(no matches)'
|
||||
@@ -2093,6 +2163,53 @@ 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)
|
||||
if context.jupyter_runtime is not None:
|
||||
code = r'''
|
||||
from pathlib import Path
|
||||
import re
|
||||
pattern = PATTERN
|
||||
literal = LITERAL
|
||||
max_matches = MAX_MATCHES
|
||||
root = Path(ROOT)
|
||||
if not root.exists():
|
||||
raise SystemExit(f"Path not found: {root}")
|
||||
regex = re.compile(re.escape(pattern) if literal else pattern)
|
||||
hits = []
|
||||
files = root.rglob("*") if root.is_dir() else [root]
|
||||
for file_path in files:
|
||||
if not file_path.is_file():
|
||||
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):
|
||||
hits.append(f"{file_path}:{line_no}: {line[:500]}")
|
||||
if len(hits) >= max_matches:
|
||||
print("\n".join(hits + [f"... truncated at {max_matches} matches ..."]))
|
||||
raise SystemExit(0)
|
||||
print("\n".join(hits) if hits else "(no matches)")
|
||||
'''
|
||||
root = context.jupyter_runtime.resolve_workspace_path(raw_path)
|
||||
script = (
|
||||
code.replace('PATTERN', repr(pattern))
|
||||
.replace('LITERAL', repr(literal))
|
||||
.replace('MAX_MATCHES', str(max_matches))
|
||||
.replace('ROOT', repr(root))
|
||||
)
|
||||
result = context.jupyter_runtime.run_python(
|
||||
code=script,
|
||||
script_path=None,
|
||||
args=[],
|
||||
stdin=None,
|
||||
timeout_seconds=30.0,
|
||||
max_output_chars=context.max_output_chars,
|
||||
cancel_event=context.cancel_event,
|
||||
)
|
||||
if result.exit_code != 0:
|
||||
raise ToolExecutionError(result.stdout.strip() or 'remote grep_search failed')
|
||||
return result.stdout.strip() or '(no matches)'
|
||||
root = _resolve_path(raw_path, context, allow_outside_root=True)
|
||||
if not root.exists():
|
||||
raise ToolExecutionError(f'Path not found: {raw_path}')
|
||||
@@ -2180,6 +2297,46 @@ def _run_python_exec(arguments: dict[str, Any], context: ToolExecutionContext) -
|
||||
context.command_timeout_seconds,
|
||||
)
|
||||
max_output_chars = _coerce_int(arguments, 'max_output_chars', context.max_output_chars)
|
||||
if context.jupyter_runtime is not None:
|
||||
result = context.jupyter_runtime.run_python(
|
||||
code=code,
|
||||
script_path=(
|
||||
context.jupyter_runtime.resolve_workspace_path(script_path)
|
||||
if script_path
|
||||
else None
|
||||
),
|
||||
args=raw_args,
|
||||
stdin=stdin,
|
||||
timeout_seconds=timeout_seconds,
|
||||
max_output_chars=max_output_chars,
|
||||
cancel_event=context.cancel_event,
|
||||
)
|
||||
payload = [
|
||||
f'exit_code={result.exit_code}',
|
||||
'interpreter=remote:python3',
|
||||
f'remote_workspace={context.jupyter_runtime.binding.workspace_cwd}',
|
||||
'[stdout]',
|
||||
result.stdout.rstrip(),
|
||||
'[stderr]',
|
||||
result.stderr.rstrip(),
|
||||
]
|
||||
if result.timed_out:
|
||||
payload.insert(0, f'timed_out=true\ntimeout_seconds={timeout_seconds:g}')
|
||||
return (
|
||||
_truncate_output('\n'.join(payload).strip(), max_output_chars),
|
||||
{
|
||||
'action': 'remote_python_exec',
|
||||
'mode': 'code' if code else 'script',
|
||||
'interpreter': 'remote:python3',
|
||||
'remote_workspace': context.jupyter_runtime.binding.workspace_cwd,
|
||||
'timed_out': result.timed_out,
|
||||
'timeout_seconds': timeout_seconds,
|
||||
'exit_code': result.exit_code,
|
||||
'stdout_preview': _snapshot_text(result.stdout),
|
||||
'stderr_preview': _snapshot_text(result.stderr),
|
||||
'output_preview': _snapshot_text('\n'.join(payload).strip()),
|
||||
},
|
||||
)
|
||||
interpreter = _resolve_python_interpreter(context)
|
||||
if code:
|
||||
command = [interpreter, '-c', code]
|
||||
@@ -2484,6 +2641,34 @@ def _resolve_python_interpreter(context: ToolExecutionContext) -> str:
|
||||
def _run_bash(arguments: dict[str, Any], context: ToolExecutionContext) -> str:
|
||||
command = _require_string(arguments, 'command')
|
||||
_ensure_shell_allowed(command, context)
|
||||
if context.jupyter_runtime is not None:
|
||||
result = context.jupyter_runtime.run_command(
|
||||
command,
|
||||
timeout_seconds=context.command_timeout_seconds,
|
||||
max_output_chars=context.max_output_chars,
|
||||
cancel_event=context.cancel_event,
|
||||
)
|
||||
payload = [
|
||||
f'exit_code={result.exit_code}',
|
||||
f'remote_workspace={context.jupyter_runtime.binding.workspace_cwd}',
|
||||
'[stdout]',
|
||||
result.stdout.rstrip(),
|
||||
'[stderr]',
|
||||
result.stderr.rstrip(),
|
||||
]
|
||||
return (
|
||||
_truncate_output('\n'.join(payload).strip(), context.max_output_chars),
|
||||
{
|
||||
'action': 'remote_bash',
|
||||
'command': command,
|
||||
'exit_code': result.exit_code,
|
||||
'remote_workspace': context.jupyter_runtime.binding.workspace_cwd,
|
||||
'timed_out': result.timed_out,
|
||||
'stdout_preview': _snapshot_text(result.stdout),
|
||||
'stderr_preview': _snapshot_text(result.stderr),
|
||||
'output_preview': _snapshot_text('\n'.join(payload).strip()),
|
||||
},
|
||||
)
|
||||
completed = subprocess.run(
|
||||
command,
|
||||
shell=True,
|
||||
@@ -3784,6 +3969,20 @@ def _stream_bash(
|
||||
arguments: dict[str, Any],
|
||||
context: ToolExecutionContext,
|
||||
) -> Iterator[ToolStreamUpdate]:
|
||||
if context.jupyter_runtime is not None:
|
||||
result = _run_bash(arguments, context)
|
||||
content, metadata = result if isinstance(result, tuple) else (result, {})
|
||||
tool_result = ToolExecutionResult(
|
||||
name='bash',
|
||||
ok=True,
|
||||
content=content,
|
||||
metadata=metadata,
|
||||
)
|
||||
if content:
|
||||
yield from _stream_static_text_result(tool_result)
|
||||
else:
|
||||
yield ToolStreamUpdate(kind='result', result=tool_result)
|
||||
return
|
||||
try:
|
||||
command = _require_string(arguments, 'command')
|
||||
_ensure_shell_allowed(command, context)
|
||||
|
||||
Reference in New Issue
Block a user