Add Jupyter workspace runtime
This commit is contained in:
wuyang6
@@ -0,0 +1,648 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import re
|
||||
import shlex
|
||||
import ssl
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from html import unescape
|
||||
from pathlib import PurePosixPath
|
||||
from typing import Any
|
||||
from urllib.parse import quote, urlparse, urlunparse
|
||||
|
||||
try: # 依赖在部署环境安装；这里延迟报错，方便无网络环境做静态检查。
|
||||
import requests
|
||||
except ImportError: # pragma: no cover
|
||||
requests = None # type: ignore[assignment]
|
||||
|
||||
try:
|
||||
import websocket
|
||||
except ImportError: # pragma: no cover
|
||||
websocket = None # type: ignore[assignment]
|
||||
|
||||
|
||||
DEFAULT_JUPYTER_WORKSPACE_ROOT = '/root/zk_agent_workspaces'
|
||||
|
||||
|
||||
class JupyterRuntimeError(RuntimeError):
|
||||
"""Raised when a remote Jupyter runtime cannot complete a request."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class JupyterWorkspaceBinding:
|
||||
account_id: str
|
||||
session_id: str
|
||||
base_url: str
|
||||
workspace_root: str
|
||||
workspace_cwd: str
|
||||
workspace_api_path: str
|
||||
jupyter_tree_url: str
|
||||
created_at: float
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
'account_id': self.account_id,
|
||||
'session_id': self.session_id,
|
||||
'base_url': self.base_url,
|
||||
'workspace_root': self.workspace_root,
|
||||
'workspace_cwd': self.workspace_cwd,
|
||||
'workspace_api_path': self.workspace_api_path,
|
||||
'jupyter_tree_url': self.jupyter_tree_url,
|
||||
'created_at': self.created_at,
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RemoteCommandResult:
|
||||
exit_code: int
|
||||
stdout: str
|
||||
stderr: str = ''
|
||||
timed_out: bool = False
|
||||
|
||||
|
||||
class JupyterRuntimeSession:
|
||||
"""A small terminal-backed execution bridge for one account/session pair."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
binding: JupyterWorkspaceBinding,
|
||||
http_session: Any,
|
||||
xsrf_token: str,
|
||||
) -> None:
|
||||
self.binding = binding
|
||||
self._session = http_session
|
||||
self._xsrf_token = xsrf_token
|
||||
self._terminal_name: str | None = None
|
||||
self._lock = threading.RLock()
|
||||
|
||||
@classmethod
|
||||
def connect(
|
||||
cls,
|
||||
*,
|
||||
account_id: str,
|
||||
session_id: str,
|
||||
base_url: str,
|
||||
password: str,
|
||||
workspace_root: str = DEFAULT_JUPYTER_WORKSPACE_ROOT,
|
||||
timeout_seconds: float = 20.0,
|
||||
) -> 'JupyterRuntimeSession':
|
||||
if requests is None:
|
||||
raise JupyterRuntimeError(
|
||||
'Missing dependency: requests. Please install project dependencies.'
|
||||
)
|
||||
normalized_base_url = normalize_jupyter_base_url(base_url)
|
||||
http_session = requests.Session()
|
||||
xsrf_token = _login_jupyter(
|
||||
http_session,
|
||||
normalized_base_url,
|
||||
password,
|
||||
timeout_seconds=timeout_seconds,
|
||||
)
|
||||
workspace_root = normalize_posix_path(workspace_root)
|
||||
workspace_cwd = f'{workspace_root}/{sanitize_remote_path_part(session_id)}'
|
||||
workspace_api_path = api_path_for_absolute_path(workspace_cwd)
|
||||
binding = JupyterWorkspaceBinding(
|
||||
account_id=account_id,
|
||||
session_id=session_id,
|
||||
base_url=normalized_base_url,
|
||||
workspace_root=workspace_root,
|
||||
workspace_cwd=workspace_cwd,
|
||||
workspace_api_path=workspace_api_path,
|
||||
jupyter_tree_url=(
|
||||
f'{normalized_base_url}/lab/tree/{quote(workspace_api_path)}'
|
||||
),
|
||||
created_at=time.time(),
|
||||
)
|
||||
runtime = cls(
|
||||
binding=binding,
|
||||
http_session=http_session,
|
||||
xsrf_token=xsrf_token,
|
||||
)
|
||||
runtime.bootstrap_workspace(timeout_seconds=timeout_seconds)
|
||||
return runtime
|
||||
|
||||
def bootstrap_workspace(self, *, timeout_seconds: float = 20.0) -> None:
|
||||
result = self.run_command(
|
||||
(
|
||||
'mkdir -p '
|
||||
f'{shlex.quote(self.binding.workspace_cwd)}/output '
|
||||
f'{shlex.quote(self.binding.workspace_cwd)}/scratchpad'
|
||||
),
|
||||
timeout_seconds=timeout_seconds,
|
||||
max_output_chars=4000,
|
||||
)
|
||||
if result.exit_code != 0:
|
||||
raise JupyterRuntimeError(
|
||||
result.stdout.strip() or 'Unable to bootstrap Jupyter workspace.'
|
||||
)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
payload = self.binding.to_dict()
|
||||
payload['connected'] = True
|
||||
return payload
|
||||
|
||||
def render_context(self) -> str:
|
||||
return '\n'.join(
|
||||
[
|
||||
'[远端 Jupyter 工作区]',
|
||||
'当前 session 已切换到远端 Jupyter 运行时。',
|
||||
f'- 默认工作目录：{self.binding.workspace_cwd}',
|
||||
f'- 输出文件优先写入：{self.binding.workspace_cwd}/output',
|
||||
'- 线上挂载目录（例如 /mnt/...）可以通过绝对路径读取和分析。',
|
||||
'- 除非用户明确要求，不要修改平台代码目录；数据产物应留在当前 session 工作区。',
|
||||
f'- Jupyter 工作区入口：{self.binding.jupyter_tree_url}',
|
||||
]
|
||||
)
|
||||
|
||||
def run_command(
|
||||
self,
|
||||
command: str,
|
||||
*,
|
||||
timeout_seconds: float,
|
||||
max_output_chars: int,
|
||||
cancel_event: Any | None = None,
|
||||
) -> RemoteCommandResult:
|
||||
if websocket is None:
|
||||
raise JupyterRuntimeError(
|
||||
'Missing dependency: websocket-client. Please install project dependencies.'
|
||||
)
|
||||
if not command.strip():
|
||||
raise JupyterRuntimeError('Remote command is empty.')
|
||||
marker = f'__ZK_AGENT_EXIT_{uuid.uuid4().hex}__'
|
||||
workspace = shlex.quote(self.binding.workspace_cwd)
|
||||
wrapped = (
|
||||
f'mkdir -p {workspace} && '
|
||||
f'cd {workspace} && '
|
||||
f'bash -lc {shlex.quote(command)}; '
|
||||
f'printf "\\n{marker}:%s\\n" "$?"'
|
||||
)
|
||||
with self._lock:
|
||||
name = self._ensure_terminal()
|
||||
ws_url = self._terminal_websocket_url(name)
|
||||
ws = websocket.create_connection(
|
||||
ws_url,
|
||||
timeout=min(timeout_seconds, 30.0),
|
||||
cookie=self._cookie_header(),
|
||||
origin=self.binding.base_url,
|
||||
sslopt={'cert_reqs': ssl.CERT_NONE},
|
||||
header=[
|
||||
f'X-XSRFToken: {self._xsrf_token}',
|
||||
f'Referer: {self.binding.base_url}/lab',
|
||||
],
|
||||
)
|
||||
try:
|
||||
self._drain_initial_messages(ws, max_wait_seconds=1.0)
|
||||
ws.send(json.dumps(['stdin', 'stty -echo\n']))
|
||||
self._drain_initial_messages(ws, max_wait_seconds=0.3)
|
||||
ws.send(json.dumps(['stdin', f'{wrapped}\n']))
|
||||
return self._collect_until_marker(
|
||||
ws,
|
||||
marker=marker,
|
||||
timeout_seconds=timeout_seconds,
|
||||
max_output_chars=max_output_chars,
|
||||
cancel_event=cancel_event,
|
||||
)
|
||||
finally:
|
||||
try:
|
||||
ws.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def run_python(
|
||||
self,
|
||||
*,
|
||||
code: str | None,
|
||||
script_path: str | None,
|
||||
args: list[str],
|
||||
stdin: str | None,
|
||||
timeout_seconds: float,
|
||||
max_output_chars: int,
|
||||
cancel_event: Any | None = None,
|
||||
) -> RemoteCommandResult:
|
||||
if bool(code) == bool(script_path):
|
||||
raise JupyterRuntimeError('code and script_path must specify exactly one value')
|
||||
temp_dir = f'{self.binding.workspace_cwd}/scratchpad/.python_exec'
|
||||
command_parts = [f'mkdir -p {shlex.quote(temp_dir)}']
|
||||
if code is not None:
|
||||
script_path = f'{temp_dir}/snippet_{uuid.uuid4().hex}.py'
|
||||
command_parts.append(
|
||||
self._remote_write_text_command(script_path, code)
|
||||
)
|
||||
assert script_path is not None
|
||||
quoted_args = ' '.join(shlex.quote(arg) for arg in args)
|
||||
python_command = f'python3 {shlex.quote(script_path)}'
|
||||
if quoted_args:
|
||||
python_command = f'{python_command} {quoted_args}'
|
||||
if stdin is not None:
|
||||
stdin_path = f'{temp_dir}/stdin_{uuid.uuid4().hex}.txt'
|
||||
command_parts.append(
|
||||
self._remote_write_text_command(stdin_path, stdin)
|
||||
)
|
||||
python_command = f'{python_command} < {shlex.quote(stdin_path)}'
|
||||
command_parts.append(python_command)
|
||||
return self.run_command(
|
||||
' && '.join(command_parts),
|
||||
timeout_seconds=timeout_seconds,
|
||||
max_output_chars=max_output_chars,
|
||||
cancel_event=cancel_event,
|
||||
)
|
||||
|
||||
def list_dir(self, path: str, *, max_entries: int, max_output_chars: int) -> str:
|
||||
target = self.resolve_workspace_path(path)
|
||||
code = r'''
|
||||
import json
|
||||
from pathlib import Path
|
||||
target = Path(PATH)
|
||||
if not target.exists():
|
||||
raise SystemExit(f"Path not found: {target}")
|
||||
if not target.is_dir():
|
||||
raise SystemExit(f"Path is not a directory: {target}")
|
||||
entries = []
|
||||
for item in sorted(target.iterdir(), key=lambda p: (not p.is_dir(), p.name.lower()))[:MAX_ENTRIES]:
|
||||
entries.append(("dir" if item.is_dir() else "file") + "\t" + str(item))
|
||||
print("\n".join(entries) if entries else "(empty directory)")
|
||||
'''
|
||||
script = (
|
||||
code.replace('PATH', python_literal(target))
|
||||
.replace('MAX_ENTRIES', str(max_entries))
|
||||
)
|
||||
result = self.run_python(
|
||||
code=script,
|
||||
script_path=None,
|
||||
args=[],
|
||||
stdin=None,
|
||||
timeout_seconds=20.0,
|
||||
max_output_chars=max_output_chars,
|
||||
)
|
||||
if result.exit_code != 0:
|
||||
raise JupyterRuntimeError(result.stdout.strip() or 'remote list_dir failed')
|
||||
return result.stdout.strip()
|
||||
|
||||
def read_text(
|
||||
self,
|
||||
path: str,
|
||||
*,
|
||||
start_line: int | None,
|
||||
end_line: int | None,
|
||||
max_output_chars: int,
|
||||
) -> str:
|
||||
target = self.resolve_workspace_path(path)
|
||||
code = r'''
|
||||
from pathlib import Path
|
||||
target = Path(PATH)
|
||||
if not target.is_file():
|
||||
raise SystemExit(f"Path is not a file: {target}")
|
||||
text = target.read_text(encoding="utf-8", errors="replace")
|
||||
start_line = START_LINE
|
||||
end_line = END_LINE
|
||||
if start_line is None and end_line is None:
|
||||
print(text, end="")
|
||||
else:
|
||||
lines = text.splitlines()
|
||||
start_idx = max((start_line or 1) - 1, 0)
|
||||
end_idx = end_line or len(lines)
|
||||
print("\n".join(f"{start_idx + idx + 1}: {line}" for idx, line in enumerate(lines[start_idx:end_idx])), end="")
|
||||
'''
|
||||
script = (
|
||||
code.replace('PATH', python_literal(target))
|
||||
.replace('START_LINE', python_literal(start_line))
|
||||
.replace('END_LINE', python_literal(end_line))
|
||||
)
|
||||
result = self.run_python(
|
||||
code=script,
|
||||
script_path=None,
|
||||
args=[],
|
||||
stdin=None,
|
||||
timeout_seconds=20.0,
|
||||
max_output_chars=max_output_chars,
|
||||
)
|
||||
if result.exit_code != 0:
|
||||
raise JupyterRuntimeError(result.stdout.strip() or 'remote read_file failed')
|
||||
return result.stdout[:max_output_chars]
|
||||
|
||||
def write_text(
|
||||
self,
|
||||
path: str,
|
||||
content: str,
|
||||
*,
|
||||
append: bool,
|
||||
newline_at_end: bool,
|
||||
max_output_chars: int,
|
||||
) -> str:
|
||||
target = self.resolve_workspace_path(path)
|
||||
if newline_at_end and content and not content.endswith('\n'):
|
||||
content += '\n'
|
||||
mode = 'a' if append else 'w'
|
||||
code = r'''
|
||||
from pathlib import Path
|
||||
target = Path(PATH)
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
content = CONTENT
|
||||
with target.open(MODE, encoding="utf-8") as fh:
|
||||
fh.write(content)
|
||||
print(f"{'appended' if MODE == 'a' else 'wrote'} {target} ({len(content)} chars)")
|
||||
'''
|
||||
script = (
|
||||
code.replace('PATH', python_literal(target))
|
||||
.replace('CONTENT', python_literal(content))
|
||||
.replace('MODE', python_literal(mode))
|
||||
)
|
||||
result = self.run_python(
|
||||
code=script,
|
||||
script_path=None,
|
||||
args=[],
|
||||
stdin=None,
|
||||
timeout_seconds=20.0,
|
||||
max_output_chars=max_output_chars,
|
||||
)
|
||||
if result.exit_code != 0:
|
||||
raise JupyterRuntimeError(result.stdout.strip() or 'remote write_file failed')
|
||||
return result.stdout.strip()
|
||||
|
||||
def resolve_workspace_path(self, raw_path: str) -> str:
|
||||
value = raw_path.strip() or '.'
|
||||
if value.startswith('/'):
|
||||
return normalize_posix_path(value)
|
||||
path = PurePosixPath(value)
|
||||
if path.parts and path.parts[0] in {'output', 'outputs'}:
|
||||
tail = PurePosixPath(*path.parts[1:]) if len(path.parts) > 1 else PurePosixPath()
|
||||
return normalize_posix_path(f'{self.binding.workspace_cwd}/output/{tail}')
|
||||
if path.parts and path.parts[0] in {'scratchpad', 'scratch'}:
|
||||
tail = PurePosixPath(*path.parts[1:]) if len(path.parts) > 1 else PurePosixPath()
|
||||
return normalize_posix_path(f'{self.binding.workspace_cwd}/scratchpad/{tail}')
|
||||
return normalize_posix_path(f'{self.binding.workspace_cwd}/{value}')
|
||||
|
||||
def _ensure_terminal(self) -> str:
|
||||
if self._terminal_name:
|
||||
return self._terminal_name
|
||||
response = self._session.post(
|
||||
f'{self.binding.base_url}/api/terminals',
|
||||
headers={'X-XSRFToken': self._xsrf_token},
|
||||
timeout=20,
|
||||
)
|
||||
if response.status_code >= 400:
|
||||
raise JupyterRuntimeError(
|
||||
f'Unable to create Jupyter terminal: HTTP {response.status_code} {response.text[:500]}'
|
||||
)
|
||||
payload = response.json()
|
||||
name = str(payload.get('name') or '').strip()
|
||||
if not name:
|
||||
raise JupyterRuntimeError('Jupyter terminal response did not contain a terminal name.')
|
||||
self._terminal_name = name
|
||||
return name
|
||||
|
||||
def _terminal_websocket_url(self, name: str) -> str:
|
||||
parsed = urlparse(self.binding.base_url)
|
||||
scheme = 'wss' if parsed.scheme == 'https' else 'ws'
|
||||
base_path = parsed.path.rstrip('/')
|
||||
path = f'{base_path}/terminals/websocket/{quote(name)}'
|
||||
return urlunparse((scheme, parsed.netloc, path, '', '', ''))
|
||||
|
||||
def _cookie_header(self) -> str:
|
||||
return '; '.join(
|
||||
f'{cookie.name}={cookie.value}'
|
||||
for cookie in self._session.cookies
|
||||
)
|
||||
|
||||
def _drain_initial_messages(self, ws: Any, *, max_wait_seconds: float) -> None:
|
||||
deadline = time.monotonic() + max_wait_seconds
|
||||
old_timeout = ws.gettimeout()
|
||||
try:
|
||||
while time.monotonic() < deadline:
|
||||
ws.settimeout(max(0.05, deadline - time.monotonic()))
|
||||
try:
|
||||
ws.recv()
|
||||
except Exception:
|
||||
break
|
||||
finally:
|
||||
ws.settimeout(old_timeout)
|
||||
|
||||
def _collect_until_marker(
|
||||
self,
|
||||
ws: Any,
|
||||
*,
|
||||
marker: str,
|
||||
timeout_seconds: float,
|
||||
max_output_chars: int,
|
||||
cancel_event: Any | None,
|
||||
) -> RemoteCommandResult:
|
||||
deadline = time.monotonic() + timeout_seconds
|
||||
transcript: list[str] = []
|
||||
exit_code: int | None = None
|
||||
marker_pattern = re.compile(rf'{re.escape(marker)}:(\d+)')
|
||||
while time.monotonic() < deadline:
|
||||
if cancel_event is not None and cancel_event.is_set():
|
||||
try:
|
||||
ws.send(json.dumps(['stdin', '\x03']))
|
||||
except Exception:
|
||||
pass
|
||||
raise JupyterRuntimeError('Run cancelled by user')
|
||||
ws.settimeout(max(0.1, min(1.0, deadline - time.monotonic())))
|
||||
try:
|
||||
raw = ws.recv()
|
||||
except Exception:
|
||||
continue
|
||||
for chunk in parse_terminal_payload(raw):
|
||||
match = marker_pattern.search(chunk)
|
||||
if match:
|
||||
exit_code = int(match.group(1))
|
||||
chunk = marker_pattern.sub('', chunk)
|
||||
if chunk:
|
||||
transcript.append(chunk)
|
||||
break
|
||||
transcript.append(chunk)
|
||||
if exit_code is not None:
|
||||
break
|
||||
if exit_code is None:
|
||||
try:
|
||||
ws.send(json.dumps(['stdin', '\x03']))
|
||||
except Exception:
|
||||
pass
|
||||
return RemoteCommandResult(
|
||||
exit_code=124,
|
||||
stdout=truncate_text(''.join(transcript), max_output_chars),
|
||||
timed_out=True,
|
||||
)
|
||||
return RemoteCommandResult(
|
||||
exit_code=exit_code,
|
||||
stdout=truncate_text(clean_terminal_output(''.join(transcript)), max_output_chars),
|
||||
)
|
||||
|
||||
def _remote_write_text_command(self, path: str, content: str) -> str:
|
||||
encoded = base64.b64encode(content.encode('utf-8')).decode('ascii')
|
||||
script = '; '.join(
|
||||
[
|
||||
'import base64',
|
||||
'from pathlib import Path',
|
||||
f'path = Path({json.dumps(path)})',
|
||||
'path.parent.mkdir(parents=True, exist_ok=True)',
|
||||
f'path.write_bytes(base64.b64decode({json.dumps(encoded)}))',
|
||||
]
|
||||
)
|
||||
return f'python3 -c {shlex.quote(script)}'
|
||||
|
||||
|
||||
class JupyterRuntimeManager:
|
||||
def __init__(self) -> None:
|
||||
self._lock = threading.RLock()
|
||||
self._sessions: dict[tuple[str, str], JupyterRuntimeSession] = {}
|
||||
|
||||
def bind_session(
|
||||
self,
|
||||
*,
|
||||
account_id: str,
|
||||
session_id: str,
|
||||
base_url: str,
|
||||
password: str,
|
||||
workspace_root: str = DEFAULT_JUPYTER_WORKSPACE_ROOT,
|
||||
) -> JupyterRuntimeSession:
|
||||
runtime = JupyterRuntimeSession.connect(
|
||||
account_id=account_id,
|
||||
session_id=session_id,
|
||||
base_url=base_url,
|
||||
password=password,
|
||||
workspace_root=workspace_root,
|
||||
)
|
||||
with self._lock:
|
||||
self._sessions[(account_id, session_id)] = runtime
|
||||
return runtime
|
||||
|
||||
def get(self, account_id: str, session_id: str) -> JupyterRuntimeSession | None:
|
||||
with self._lock:
|
||||
return self._sessions.get((account_id, session_id))
|
||||
|
||||
def session_payload(self, account_id: str, session_id: str) -> dict[str, Any]:
|
||||
runtime = self.get(account_id, session_id)
|
||||
if runtime is None:
|
||||
return {
|
||||
'connected': False,
|
||||
'account_id': account_id,
|
||||
'session_id': session_id,
|
||||
}
|
||||
return runtime.to_dict()
|
||||
|
||||
|
||||
def normalize_jupyter_base_url(value: str) -> str:
|
||||
raw = value.strip()
|
||||
if not raw:
|
||||
raise JupyterRuntimeError('Jupyter URL is empty.')
|
||||
parsed = urlparse(raw)
|
||||
if not parsed.scheme:
|
||||
parsed = urlparse(f'https://{raw}')
|
||||
if parsed.scheme not in {'http', 'https'} or not parsed.netloc:
|
||||
raise JupyterRuntimeError(f'Invalid Jupyter URL: {value}')
|
||||
path = parsed.path.rstrip('/')
|
||||
for marker in ('/lab', '/tree', '/notebooks', '/login'):
|
||||
index = path.find(marker)
|
||||
if index >= 0:
|
||||
path = path[:index]
|
||||
break
|
||||
return urlunparse((parsed.scheme, parsed.netloc, path.rstrip('/'), '', '', '')).rstrip('/')
|
||||
|
||||
|
||||
def normalize_posix_path(value: str) -> str:
|
||||
path = PurePosixPath(value)
|
||||
if not str(path).startswith('/'):
|
||||
path = PurePosixPath('/') / path
|
||||
normalized = str(path)
|
||||
normalized = re.sub(r'/+', '/', normalized)
|
||||
return normalized.rstrip('/') or '/'
|
||||
|
||||
|
||||
def sanitize_remote_path_part(value: str) -> str:
|
||||
return re.sub(r'[^A-Za-z0-9_.-]+', '_', value).strip('._') or uuid.uuid4().hex
|
||||
|
||||
|
||||
def api_path_for_absolute_path(path: str) -> str:
|
||||
normalized = normalize_posix_path(path)
|
||||
if normalized == '/root':
|
||||
return ''
|
||||
if normalized.startswith('/root/'):
|
||||
return normalized[len('/root/') :]
|
||||
return normalized.lstrip('/')
|
||||
|
||||
|
||||
def parse_terminal_payload(raw: str) -> list[str]:
|
||||
try:
|
||||
payload = json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
return [raw]
|
||||
if isinstance(payload, list) and len(payload) >= 2:
|
||||
kind, content = payload[0], payload[1]
|
||||
if kind in {'stdout', 'stderr'} and isinstance(content, str):
|
||||
return [content]
|
||||
return []
|
||||
|
||||
|
||||
def clean_terminal_output(text: str) -> str:
|
||||
# 去掉少量 ANSI 控制符，避免输出里带终端颜色和光标控制。
|
||||
cleaned = re.sub(r'\x1b\[[0-?]*[ -/]*[@-~]', '', text)
|
||||
lines = [
|
||||
line.rstrip('\r')
|
||||
for line in cleaned.splitlines()
|
||||
if line.strip() not in {'#', '>'} and not re.fullmatch(r'(>\s*)+', line.strip())
|
||||
]
|
||||
return '\n'.join(lines).strip('\r\n')
|
||||
|
||||
|
||||
def truncate_text(text: str, limit: int) -> str:
|
||||
if limit <= 0 or len(text) <= limit:
|
||||
return text
|
||||
omitted = len(text) - limit
|
||||
return f'{text[:limit]}\n... [truncated, {omitted} chars omitted]'
|
||||
|
||||
|
||||
def python_literal(value: object) -> str:
|
||||
return repr(value)
|
||||
|
||||
|
||||
def _login_jupyter(
|
||||
http_session: Any,
|
||||
base_url: str,
|
||||
password: str,
|
||||
*,
|
||||
timeout_seconds: float,
|
||||
) -> str:
|
||||
login_url = f'{base_url}/login?next=%2Flab'
|
||||
response = http_session.get(login_url, timeout=timeout_seconds)
|
||||
if response.status_code >= 400:
|
||||
raise JupyterRuntimeError(
|
||||
f'Unable to open Jupyter login page: HTTP {response.status_code}'
|
||||
)
|
||||
xsrf_token = _extract_xsrf_token(response.text) or http_session.cookies.get('_xsrf')
|
||||
if not xsrf_token:
|
||||
raise JupyterRuntimeError('Unable to find Jupyter _xsrf token on login page.')
|
||||
post = http_session.post(
|
||||
f'{base_url}/login',
|
||||
data={'_xsrf': xsrf_token, 'password': password},
|
||||
headers={'Referer': login_url},
|
||||
timeout=timeout_seconds,
|
||||
allow_redirects=True,
|
||||
)
|
||||
if post.status_code >= 400:
|
||||
raise JupyterRuntimeError(
|
||||
f'Jupyter login failed: HTTP {post.status_code} {post.text[:500]}'
|
||||
)
|
||||
probe = http_session.get(f'{base_url}/api/contents', timeout=timeout_seconds)
|
||||
if probe.status_code >= 400:
|
||||
raise JupyterRuntimeError(
|
||||
f'Jupyter contents API probe failed after login: HTTP {probe.status_code} {probe.text[:500]}'
|
||||
)
|
||||
return str(xsrf_token)
|
||||
|
||||
|
||||
def _extract_xsrf_token(html: str) -> str | None:
|
||||
patterns = [
|
||||
r'name=["\']_xsrf["\']\s+value=["\']([^"\']+)["\']',
|
||||
r'value=["\']([^"\']+)["\']\s+name=["\']_xsrf["\']',
|
||||
]
|
||||
for pattern in patterns:
|
||||
match = re.search(pattern, html)
|
||||
if match:
|
||||
return unescape(match.group(1))
|
||||
return None
|
||||
Reference in New Issue
Block a user